From 480d9da452a9c29962289e97d943241bb2dade4f Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Mon, 17 Aug 2026 04:00:12 -0700 Subject: [PATCH 1/4] [https://nvbugs/6600098][test] Cap KV cache in scheduler comparison tests Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../defs/kv_cache/test_kv_cache_v2_scheduler.py | 9 +++++++++ tests/integration/test_lists/waives.txt | 4 ---- 2 files changed, 9 insertions(+), 4 deletions(-) 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..e59fa1b92848 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. @@ -230,6 +235,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..e0257ea6ace3 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -322,11 +322,7 @@ 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_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) From 639ff47af76e2d20cc19cc9546bf3ee84af15d8c Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Mon, 17 Aug 2026 05:52:52 -0700 Subject: [PATCH 2/4] [https://nvbugs/6600098][test] Unwaive remaining KV cache scheduler tests Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index e0257ea6ace3..2aea8e3a85b6 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -321,8 +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_with_block_reuse SKIP (https://nvbugs/6462303) 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) From 17958b6806f5b234c619f39afe8e7638e35464f7 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Mon, 17 Aug 2026 10:05:02 -0700 Subject: [PATCH 3/4] [https://nvbugs/6600098][test] Isolate KV cache session reuse per test Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../defs/kv_cache/test_kv_cache_v2_scheduler.py | 11 +++++++++++ tests/test_common/session_prefetcher.py | 7 +++++++ tests/unittest/llmapi/test_session_prefetcher.py | 14 ++++++++++++++ 3 files changed, 32 insertions(+) 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 e59fa1b92848..ce76e2309fd8 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 @@ -22,6 +22,8 @@ import pytest import torch +from test_common.session_prefetcher import PREFETCHER +from test_common.session_reuse import REUSE from tensorrt_llm import LLM from tensorrt_llm._torch.peft.lora.config import LoraConfig @@ -233,6 +235,15 @@ class TestKVCacheV2Llama: MODEL_PATH = f"{llm_models_root()}/llama-3.2-models/Llama-3.2-1B" + @pytest.fixture(autouse=True) + def _test_scoped_session_reuse(self, monkeypatch): + """Keep V1 -> V2 reuse, but never carry an MPI pool across tests.""" + monkeypatch.setenv("TRTLLM_TEST_PREFETCH_SESSION", "0") + PREFETCHER.discard_shadow() + REUSE.drain() + yield + REUSE.drain() + def _compare(self, prompts, max_tokens=32, kv_extra=None, **llm_kwargs): """Run V1 vs V2 with greedy sampling; assert outputs match.""" kv_extra = { diff --git a/tests/test_common/session_prefetcher.py b/tests/test_common/session_prefetcher.py index 188a6aa70cc3..7c759031630f 100644 --- a/tests/test_common/session_prefetcher.py +++ b/tests/test_common/session_prefetcher.py @@ -513,6 +513,13 @@ def take(self, spec: int) -> object | None: ).start() return None + def discard_shadow(self) -> None: + """Shut down an unused shadow without disabling future prefetching.""" + built = self._drain() + if built is not None: + self.stats["pools_discarded_test_boundary"] += 1 + built.session.shutdown() + def _make_factory(self, real_cls): """A drop-in for ``MpiPoolSession`` that consumes and re-arms the shadow.""" diff --git a/tests/unittest/llmapi/test_session_prefetcher.py b/tests/unittest/llmapi/test_session_prefetcher.py index 53928664c123..80b4fdb60c27 100644 --- a/tests/unittest/llmapi/test_session_prefetcher.py +++ b/tests/unittest/llmapi/test_session_prefetcher.py @@ -300,6 +300,20 @@ def test_factory_hit_hands_over_shadow(prefetcher): assert factory(4) is pool # prefetched pool handed over +def test_discard_shadow_keeps_prefetcher_usable(prefetcher): + pool = _FakePool(4) + _arm(prefetcher, pool, spec=4) + + prefetcher.discard_shadow() + + assert pool.shut + assert prefetcher._built is None + assert not prefetcher._disposed + prefetcher.schedule_shadow(2) + prefetcher._thread.join(timeout=10) + assert prefetcher.built == [2] + + def test_take_spec_mismatch_returns_none(prefetcher): pool = _FakePool(4) _arm(prefetcher, pool, spec=4) From d4f9f980b7ac1bad8d966cd9c5b3e57c23fc9acd Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:33:36 -0700 Subject: [PATCH 4/4] [https://nvbugs/6600098][test] Isolate KV cache comparisons with private sessions Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../defs/kv_cache/test_kv_cache_v2_scheduler.py | 14 +++----------- tests/test_common/session_prefetcher.py | 7 ------- tests/unittest/llmapi/test_session_prefetcher.py | 14 -------------- 3 files changed, 3 insertions(+), 32 deletions(-) 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 ce76e2309fd8..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 @@ -22,8 +22,6 @@ import pytest import torch -from test_common.session_prefetcher import PREFETCHER -from test_common.session_reuse import REUSE from tensorrt_llm import LLM from tensorrt_llm._torch.peft.lora.config import LoraConfig @@ -230,20 +228,14 @@ 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).""" MODEL_PATH = f"{llm_models_root()}/llama-3.2-models/Llama-3.2-1B" - @pytest.fixture(autouse=True) - def _test_scoped_session_reuse(self, monkeypatch): - """Keep V1 -> V2 reuse, but never carry an MPI pool across tests.""" - monkeypatch.setenv("TRTLLM_TEST_PREFETCH_SESSION", "0") - PREFETCHER.discard_shadow() - REUSE.drain() - yield - REUSE.drain() - def _compare(self, prompts, max_tokens=32, kv_extra=None, **llm_kwargs): """Run V1 vs V2 with greedy sampling; assert outputs match.""" kv_extra = { diff --git a/tests/test_common/session_prefetcher.py b/tests/test_common/session_prefetcher.py index 7c759031630f..188a6aa70cc3 100644 --- a/tests/test_common/session_prefetcher.py +++ b/tests/test_common/session_prefetcher.py @@ -513,13 +513,6 @@ def take(self, spec: int) -> object | None: ).start() return None - def discard_shadow(self) -> None: - """Shut down an unused shadow without disabling future prefetching.""" - built = self._drain() - if built is not None: - self.stats["pools_discarded_test_boundary"] += 1 - built.session.shutdown() - def _make_factory(self, real_cls): """A drop-in for ``MpiPoolSession`` that consumes and re-arms the shadow.""" diff --git a/tests/unittest/llmapi/test_session_prefetcher.py b/tests/unittest/llmapi/test_session_prefetcher.py index 80b4fdb60c27..53928664c123 100644 --- a/tests/unittest/llmapi/test_session_prefetcher.py +++ b/tests/unittest/llmapi/test_session_prefetcher.py @@ -300,20 +300,6 @@ def test_factory_hit_hands_over_shadow(prefetcher): assert factory(4) is pool # prefetched pool handed over -def test_discard_shadow_keeps_prefetcher_usable(prefetcher): - pool = _FakePool(4) - _arm(prefetcher, pool, spec=4) - - prefetcher.discard_shadow() - - assert pool.shut - assert prefetcher._built is None - assert not prefetcher._disposed - prefetcher.schedule_shadow(2) - prefetcher._thread.join(timeout=10) - assert prefetcher.built == [2] - - def test_take_spec_mismatch_returns_none(prefetcher): pool = _FakePool(4) _arm(prefetcher, pool, spec=4)