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 fd35e7f0e8b3..b24ebdd1dcf6 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 @@ -96,6 +96,11 @@ # V2 scheduler requires MAX_UTILIZATION policy _V2_SCHEDULER_CONFIG = SchedulerConfig(capacity_scheduler_policy="MAX_UTILIZATION") +# These functional tests construct V1 and V2 LLMs back-to-back. Cap each KV pool +# at ~256 MiB instead of using the 90% default to reduce allocator pressure +# during teardown (Llama-3.2-1B uses ~32 KiB of KV cache per token). +_LLAMA_KV_CACHE_MAX_TOKENS = 8192 + # --------------------------------------------------------------------------- # Eviction test parameters. # Goal: CUDA graph warmup passes AND runtime triggers scheduler eviction. @@ -223,6 +228,9 @@ def _run_eviction_test( # =========================================================================== # Functional tests on Llama-3.2-1B # =========================================================================== +# Each comparison constructs stateful V1 and V2 engines back-to-back. Keep +# their MPI workers private so engine state cannot leak across comparisons. +@pytest.mark.private_mpi_session class TestKVCacheV2Llama: """Functional tests for V2 scheduler using Llama-3.2-1B (1 GPU).""" @@ -230,6 +238,10 @@ class TestKVCacheV2Llama: def _compare(self, prompts, max_tokens=32, kv_extra=None, **llm_kwargs): """Run V1 vs V2 with greedy sampling; assert outputs match.""" + kv_extra = { + "max_tokens": _LLAMA_KV_CACHE_MAX_TOKENS, + **(kv_extra or {}), + } return _run_v1_v2_compare( self.MODEL_PATH, prompts, diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index ea9594dd4901..2aea8e3a85b6 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -321,12 +321,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::TestQwen3_5_4B::test_dflash SKIP (https://nvbugs/6273850) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_fp8 SKIP (https://nvbugs/6273850) 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) 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_mixtral_moe_routed_expert_fp8_multi_lora_varying_ranks[cudagraph] SKIP (https://nvbugs/6463829) llmapi/test_llm_api_pytorch_moe_lora.py::test_mixtral_moe_routed_expert_fp8_multi_lora_varying_ranks[eager] SKIP (https://nvbugs/6463829)