[https://nvbugs/6035425][fix] Fix KV cache host splitting logic#14373

mikeiovine · 2026-05-20T20:29:37Z

Description

The fix in #13130 only applied to the V2 KV cache manager.

Make it so that the KV cache host budget is always split proportionally across the draft and target. I did not mess with the GPU side splitting at all.

Note that V1 (non-VSWA) does not use the same budget splitting logic for the GPU side memory. It instead does it "implicitly" using a single max_num_tokens value. I did not want to mess with it since this is a sensitive part of the system, so the behavior there is unchanged (a simpler way to fix this is to simply remove the VSWA restriction on the call to _split_kv_cache_budget_for_draft, but the changes existing behavior).

Test Coverage

New tests added and manually verified that host budget is not doubled.

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.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Summary by CodeRabbit

Refactor
- Optimized KV cache memory management for speculative decoding. Improved internal budget allocation enables more efficient GPU and host memory distribution, providing better resource utilization and enhanced support for speculative decoding operations across various hardware configurations.

mikeiovine · 2026-05-20T20:30:05Z

/bot run

coderabbitai · 2026-05-20T20:33:01Z

📝 Walkthrough

Walkthrough

This PR refactors KV cache budget splitting for speculative decoding draft caches by introducing separate GPU and host budget-split helpers in KvCacheCreator, updating build_managers() to conditionally apply splits, and expanding test coverage to validate proportional budget allocation and integration scenarios.

Changes

KV Cache Budget Splitting for Speculative Decoding

Layer / File(s)	Summary
Cache cost computation and combined budget split foundation `tensorrt_llm/_torch/pyexecutor/_util.py`	Introduces new helpers: `_compute_cache_cost_per_manager()` computes per-manager cache costs for target vs draft layers; `_split_budget_by_slope()` splits a combined byte budget into target/draft shares using affine cost components (intercept fixed cost, slope variable cost).
GPU and host budget splitting refactor in build_managers `tensorrt_llm/_torch/pyexecutor/_util.py`	Replaces monolithic draft-budget split with separate GPU and host splitting helpers. `_compute_draft_kv_cache_size_fraction()` computes draft's fractional share; `_split_kv_cache_gpu_budget_for_draft()` clones a `KvCacheConfig` with adjusted GPU budget for draft; `_should_split_gpu_budget_for_draft()` determines when GPU-budget splitting is required. Updates `build_managers()` to apply splits only for non-estimation, non-V2 two-model cases. Updates V1 two-model path to temporarily override target config budgets with draft splits during draft manager creation, then restore original values.
Test documentation and GPU budget split test coverage `tests/unittest/_torch/executor/test_kv_cache_budget_split.py`	Updates module-level docstring to focus on KV cache budget splitting between managers. Simplifies `_make_creator()` helper docstring. Adds `TestSplitKvCacheGpuBudgetForDraft` test class validating proportional GPU budget splitting and `None` behavior when GPU budget is absent or per-token costs are equal.
Host budget split tests and GPU+host integration `tests/unittest/_torch/executor/test_kv_cache_budget_split.py`	Expands host budget splitting tests to validate `_split_host_cache_budget_for_draft()` proportional allocation, behavior without GPU budget (slope-ratio based), and regression checks ensuring target+draft budgets sum to original (not doubled). Adds chaining test where GPU split result is merged with host split, validating both GPU and host budgets. Updates "no host cache" and "zero host cache" tests to call new host-split helper.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 55.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title '[https://nvbugs/6035425][fix] Fix KV cache host splitting logic' clearly summarizes the main change: fixing KV cache host budget splitting, which aligns with the refactoring in the changeset.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The PR description explains the issue and solution clearly, references the prior fix (PR `#13130`), specifies what was and wasn't changed, and includes test coverage details.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

tensorrt-cicd · 2026-05-20T20:36:03Z

PR_Github #49484 [ run ] triggered by Bot. Commit: 1e3a9eb Link to invocation

tensorrt-cicd · 2026-05-20T22:46:29Z

PR_Github #49484 [ run ] completed with state SUCCESS. Commit: 1e3a9eb
/LLM/main/L0_MergeRequest_PR pipeline #39125 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Uh oh!

Conversation

mikeiovine commented May 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

GitHub Bot Help

Summary by CodeRabbit

Uh oh!

mikeiovine commented May 20, 2026

Uh oh!

coderabbitai Bot commented May 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

❌ Failed checks (1 warning)

Uh oh!

tensorrt-cicd commented May 20, 2026

Uh oh!

tensorrt-cicd commented May 20, 2026

Uh oh!

Uh oh!

mikeiovine commented May 21, 2026

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

mikeiovine commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

mikeiovine commented May 26, 2026

Uh oh!

tensorrt-cicd commented May 26, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

mikeiovine commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

mikeiovine commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

mikeiovine commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

mikeiovine commented May 29, 2026

Uh oh!

tensorrt-cicd commented May 29, 2026

Uh oh!

mikeiovine commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 11, 2026

Uh oh!

mikeiovine commented Jun 11, 2026

Uh oh!

tensorrt-cicd commented Jun 11, 2026

Uh oh!

tensorrt-cicd commented Jun 11, 2026

Uh oh!

mikeiovine commented May 20, 2026 •

edited

Loading

coderabbitai Bot commented May 20, 2026 •

edited

Loading