Skip to content

[https://nvbugs/6600098][test] Stabilize KV cache V2 scheduler tests #17823

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 4 commits into from
Aug 20, 2026
Merged

[https://nvbugs/6600098][test] Stabilize KV cache V2 scheduler tests #17823

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
480d9da
[https://nvbugs/6600098][test] Cap KV cache in scheduler comparison t…
yizhang-nv Aug 17, 2026
639ff47
[https://nvbugs/6600098][test] Unwaive remaining KV cache scheduler t…
yizhang-nv Aug 17, 2026
17958b6
[https://nvbugs/6600098][test] Isolate KV cache session reuse per test
yizhang-nv Aug 17, 2026
d4f9f98
[https://nvbugs/6600098][test] Isolate KV cache comparisons with priv…
yizhang-nv Aug 19, 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
12 changes: 12 additions & 0 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 @@ -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.
Expand Down Expand Up @@ -223,13 +228,20 @@ 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"

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,
Expand Down
6 changes: 0 additions & 6 deletions tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading