[https://nvbugs/6800150][test] Unwaive the KV budget split tests and fix their stale mocks

[moraxu]

Dev Engineer Review

The change removes 13 waivers for KV budget split tests and adds explicit KvCacheCreator mock attributes. The new values match the test configurations. No production code or assertions changed. The scope is limited to test setup and waiver removal.

QA Engineer Review

test_kv_cache_budget_split.py and tests/integration/test_lists/waives.txt changed. The waiver file removes entries for TestExternalDrafterKvDtype and 12 TestSplitGpuBudgetForDraft cases. The test file adds stubs for cache manager version, context-parallel helix state, sleep configuration, beam width, transceiver configuration, and sparse attention configuration. The reported focused run passed 60 tests; broader CI results are reported in the PR objectives, not independently verified here. Coverage verdict: sufficient for the reported test scope.

Per-File QA Perspective

• tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py: Covers KV cache budget splitting and external-drafter dtype behavior. The supplied context does not establish whether this file appears in test-db/ or qa/.
• tests/integration/test_lists/waives.txt: Removes 13 waivers for the KV budget split tests, enabling those tests to run. No other test-list changes are reported.

Description

The tests in tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py build KvCacheCreator with object.__new__ and stub only the attributes the cost path read when each test was written. Later changes that made the cost path read another attribute broke some of them. The tests are marked cpu_only, so they run only on the CPU stages. That is where they failed, and where they were waived under nvbugs/6800223, and they kept breaking while waived. Each row below was verified by running the test file at the listed commit and at its parent:

Commit	New read on the cost path	Tests broken
#19018 (09-17)	_cache_transceiver_config in _get_draft_cache_cost	test_target_cost_uses_derived_layer_type_windows (nvbugs/6800150, fixed in #19028)
#19040 (09-21)	Added TestExternalDrafterKvDtype without that stub, while all of unittest/_torch/executor was waived	TestExternalDrafterKvDtype (2 tests)
#19273 (09-22)	mapping.has_cp_helix() in _get_draft_cache_cost. A bare Mock answers it truthily, so the slope gets multiplied by Mock.cp_size	test_native_full_target_cost_preserves_draft_capacity (3), test_one_model_draft_cost_uses_derived_kv_config (2)
#19183 (10-09)	_is_kv_cache_manager_v2 in _bound_unsplit_v2_draft_by_tokens	test_returns_none_when_draft_kv_zero
#17495 (10-12)	_llm_args.sleep_config in _validate_or_fallback_kv_cache_manager_v2	None new: two tests that already failed now fail at an earlier line

#19183 and #17495 are only on main. On release/1.3 the failures come from #19040 and #19273.

This PR stubs each missing attribute with the value KvCacheCreator.__init__ sets for the configuration each test models:

• has_cp_helix() returns False for the non-CP mappings.
• _llm_args.sleep_config is None, _max_beam_width is 1 and _cache_transceiver_config is None.
• _is_kv_cache_manager_v2 matches the stubbed manager class.
• The draft config has sparse_attention_config = None.

release/1.3 does not read _llm_args.sleep_config or _is_kv_cache_manager_v2 on these paths yet. Those two stubs are kept: KvCacheCreator.__init__ sets both attributes on this branch too, and they keep the change correct for main, where both reads exist.

It then removes the 13 waive lines filed under nvbugs/6800223. No assertions change, and no production code changes. Of the 20 waived test items, 7 fail on release/1.3 and 13 already pass (on main, 8 fail and 12 pass).

nvbugs/6800150 and nvbugs/6800223 track the same CI failure. The waives were filed under 6800223.

Test Coverage

• test_kv_cache_budget_split.py with -m cpu_only and no waives file: 60 passed. On release/1.3 without this change: 7 failed, 53 passed.
• unittest/_torch/executor run as the CPU stage runs it (-m cpu_only --test-prefix=unittest --waives-file=tests/integration/test_lists/waives.txt, CUDA hidden), in release/1.3's CI devel image (PyTorch 26.09) with a release/1.3 CI wheel built on that image: 1614 passed, 15 skipped. Compared with release/1.3 without this change, the only difference is the 20 unwaived items, which move from skipped to passed. The 15 skips are test_pcg_fx_fallback_policy_is_model_specific, waived under nvbugs/6831653.
• This PR first targeted main, where the same change passed the full pre-merge pipeline (PR_Github #77598).
• In CI these tests run on CPU-Generic-x86-1 and CPU-Generic-arm-1, through the unittest/_torch/executor entry in l0_cpu.yml.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 9c0d20d7-1aa3-4494-8918-c78a550b005c

📥 Commits

Reviewing files that changed from the base of the PR and between b7c7cff and 30d90c7.

📒 Files selected for processing: 2

• tests/integration/test_lists/waives.txt
• tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py

💤 Files with no reviewable changes: 1

• tests/integration/test_lists/waives.txt

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Walkthrough

KV-cache budget test fixtures now initialize additional mock state and configuration. The skip list no longer contains the external-drafter dtype test or 12 split-GPU-budget draft tests.

Changes

KV-cache Budget Tests

Layer / File(s)	Summary
Update fixtures and enable tests tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py, tests/integration/test_lists/waives.txt	Test fixtures now set manager-version selection, CP-helix state, sleep configuration, sparse-attention configuration, beam width, and cache-transceiver configuration. The skip list no longer includes the external-drafter dtype test or 12 split-GPU-budget draft tests.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Other

Merge Risk: ⚪ Minimal · up to 30d90

This change re-enables the KV-cache budget tests with updated mocks. No concrete merge-blocking behavior change is identified; it is mergeable subject to normal checks.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 40.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 5 functions across 1 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly identifies the main change: unwaiving the KV budget split tests and fixing their stale mocks. It follows the required NVBugs and type format.
Description check		The description explains the issue, root cause, solution, affected tests, test coverage, CI results, and checklist status. It is complete and directly related to the changes.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[moraxu]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #77598 [ run ] triggered by Bot. Commit: b7c7cff Link to invocation

[tensorrt-cicd]

PR_Github #77598 [ run ] completed with state SUCCESS. Commit: b7c7cff
/LLM/main/L0_MergeRequest_PR pipeline #63972 completed with status: 'SUCCESS'

CI Report

Link to invocation

[moraxu]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #77921 [ run ] triggered by Bot. Commit: 30d90c7 Link to invocation

[tensorrt-cicd]

PR_Github #77921 [ run ] completed with state SUCCESS. Commit: 30d90c7
/LLM/release-1.3/L0_MergeRequest_PR pipeline #226 completed with status: 'SUCCESS'

CI Report

Link to invocation

[nvpohanh]

[by Claude Code] @yizhang-nv Could you review this PR? Thanks!

[moraxu]

Already fixed on main by #20399, closing