diff --git a/tests/integration/defs/kv_cache/test_kv_cache_v2_scheduler.py b/tests/integration/defs/kv_cache/test_kv_cache_v2_scheduler.py index b24ebdd1dcf6..e248c029c3d3 100644 --- a/tests/integration/defs/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/integration/defs/kv_cache/test_kv_cache_v2_scheduler.py @@ -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).""" @@ -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) @@ -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, ) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 880951138867..2ae317366330 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -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)