Skip to content

[https://nvbugs/6640875][test] Stabilize LoRA KV scheduler tests and re-enable seven cases #18949

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yizhang-nv merged 3 commits into from
Sep 9, 2026
Merged

[https://nvbugs/6640875][test] Stabilize LoRA KV scheduler tests and re-enable seven cases #18949

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
40cb8e2
[https://nvbugs/6428002][test] Re-enable KV scheduler sibling tests
yizhang-nv Sep 9, 2026
349ff17
[https://nvbugs/6600098][test] Unwaive remaining Llama scheduler cases
yizhang-nv Sep 9, 2026
315b498
[https://nvbugs/6640875][test] Isolate LoRA scheduler comparisons and…
yizhang-nv Sep 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 20 additions & 3 deletions tests/integration/defs/kv_cache/test_kv_cache_v2_scheduler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -364,6 +364,12 @@ def test_eviction_overlap(self):
# ===========================================================================
# LoRA tests on llama-7b-hf
# ===========================================================================
# 8192 tokens use ~4 GiB for llama-7b-hf, enough for these functional tests
# without scaling the KV allocation with the GPU's available memory.
_LORA_KV_CACHE_MAX_TOKENS = 8192


@pytest.mark.private_mpi_session
@pytest.mark.skip_less_device_memory(40000)
class TestKVCacheV2LoRA:
"""LoRA tests for V2 scheduler using llama-7b-hf (1 GPU, >=40GB)."""
Expand Down Expand Up @@ -393,7 +399,10 @@ def _run_v1_v2_lora(
if sampling_params is None:
sampling_params = SamplingParams(max_tokens=32, temperature=0.0)
if kv_extra is None:
kv_extra = {"free_gpu_memory_fraction": 0.4}
kv_extra = {
"max_tokens": _LORA_KV_CACHE_MAX_TOKENS,
"free_gpu_memory_fraction": 0.4,
}
lora_request = executor_request.LoRARequest("lora-0", 0, self.LORA_DIR)

kv_v1 = KvCacheConfig(use_kv_cache_manager_v2=False, **kv_extra)
Expand Down Expand Up @@ -462,12 +471,20 @@ def _run_multi_adapter(kv_config, **extra_llm_kwargs):
return out_lora, out_base

outs_v1 = _run_multi_adapter(
KvCacheConfig(use_kv_cache_manager_v2=False, free_gpu_memory_fraction=0.4),
KvCacheConfig(
use_kv_cache_manager_v2=False,
max_tokens=_LORA_KV_CACHE_MAX_TOKENS,
free_gpu_memory_fraction=0.4,
),
)
gc.collect()
torch.cuda.empty_cache()
outs_v2 = _run_multi_adapter(
KvCacheConfig(use_kv_cache_manager_v2=True, free_gpu_memory_fraction=0.4),
KvCacheConfig(
use_kv_cache_manager_v2=True,
max_tokens=_LORA_KV_CACHE_MAX_TOKENS,
free_gpu_memory_fraction=0.4,
),
scheduler_config=_V2_SCHEDULER_CONFIG,
)

Expand Down
7 changes: 0 additions & 7 deletions tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -281,13 +281,6 @@ full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::Tes
full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v1_kv_cache-True-True-cutlass-fp8] SKIP (https://nvbugs/6626640)
full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4[latency_moe_cutlass] SKIP (https://nvbugs/6470267)
full:sm100/unittest/bindings SKIP (Disable for Blackwell)
kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2Llama::test_chunked_prefill SKIP (https://nvbugs/6428002)
kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2Llama::test_eviction[cuda_graph] SKIP (https://nvbugs/6600098)
kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2Llama::test_eviction_with_block_reuse SKIP (https://nvbugs/6462303)
kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2Llama::test_overlap_scheduler[non_overlap] SKIP (https://nvbugs/6600098)
kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2Llama::test_overlap_scheduler[overlap] SKIP (https://nvbugs/6600098)
kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2Llama::test_token_budget_limited SKIP (https://nvbugs/6600098)
kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2LoRA::test_lora_chunked_prefill SKIP (https://nvbugs/6640875)
llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-greedy-tp2-bart-large-cnn] SKIP (https://nvbugs/6463812)
llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[cudagraph] SKIP (https://nvbugs/6683837)
llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[eager] SKIP (https://nvbugs/6683837)
Expand Down
Loading