[None][fix] Pass max_num_tokens to MLA KV manager and keep warmup-feasibility floors

[qiaoxj07]

Description

_create_kv_cache_manager's is_mla branch omits max_num_tokens from the manager constructor call, while the generic (non-MLA) branch passes it. As a result DeepseekV4CacheManager falls back to _max_num_tokens=None and _build_cache_config sizes the context extra quota by the full max_tokens estimate instead of the runtime chunk size, inflating the KV-estimation temporary quota (23.69 GiB instead of 8.14 GiB per rank on DeepSeek-V4-Pro, ISL≈990K, MTP3, dep4/ADP) — and with it the profiling peak and its OOM risk.

Forwarding max_num_tokens (one line in _util.py) fixes that, but on its own it trades the profiling-peak risk for a deterministic post-estimation OOM: every DSV4-Pro agentx disagg GEN config failed with it applied. The tight quota exposed two latent sizing bugs:

1. pool_ratio disabled the sizing constraints entirely (two independent sites: _build_cache_config only built constraints when pool_ratio is None, and StorageManager dropped constraints from min_slots when initial_pool_ratio was provided). The estimation quota was then split purely by ratio: the ageless COMPRESS pool group got 3,084 blocks < the 3,872 blocks one full-length request needs, the estimation manager clamped max_seq_len 991048→789502, and CUDA-graph warmup silently skipped every batch size ≥ 2. The profiling peak missed ~5 GiB of capture working set, the KV estimate rose from the calibrated 36.09 GiB to 40.70 GiB, and the final executor OOM'd during the advanced-sampling bs=32 graph capture with only 5.75 GiB free after KV-pool creation.
2. _KVCache.resume's utilization gate rejects any pool group above max_util_for_resume, so a warmup batch that exactly fills its floor fails its last resume (observed: a 32-request warmup batch failing at request 32 with a 32-slot pool group at 31/32 = 97% utilization).

Fix

Treat constraints as feasibility floors everywhere:

• DeepseekV4CacheManager._build_cache_config always builds its constraints (pool_ratio only replaces the typical_step-derived ratio), and Constraint 1 now describes the CUDA-graph warmup batch the engine will actually run: one max_seq_len decode plus minimal decodes up to the engine's largest configured graph batch (model_engine._max_cuda_graph_batch_size, plumbed through _create_kv_cache_manager). With max(cuda_graph_batch_sizes) ≪ max_batch_size the floor sizes to the real graph batch, and with CUDA graphs disabled (0) the constraint is dropped entirely — no other warmup creates a max_seq_len-long request, so even the single long request would floor (and, on tight quotas, bump) the allocation for a warmup that will not happen. The floor never forces pool allocation for a warmup that won't run.
• StorageManager derives min_slots from constraints even when initial_pool_ratio is provided, and scales the floors by 1/max_util_for_resume so the last request of a constraint batch still passes resume's utilization gate. The scaling only applies in the binding range (0, 1): the llmapi field permits max_util_for_resume=0, where only the first resume of an empty pool (utilization 0) passes and every later one is rejected — no finite floor helps and construction must not divide by zero; values >= 1 never trigger the gate (utilization cannot exceed 1). Neither degenerate value scales the floors.
• model_engine warns instead of silently skipping CUDA-graph warmup batch sizes and general warmup shapes, since a skip during estimation makes the profiling peak unrepresentative.

Behavior change note (V2 cache manager owners)

initial_pool_ratio no longer overrides constraints completely: the explicit ratio still decides how quota beyond the floors is split, but pool groups are clamped up to the (utilization-scaled) constraint minimums. A ratio split that cannot hold the declared warmup batches produced a manager that could not run them — that behavior was the bug. The existing test test_initial_pool_ratio_overrides_typical_step_and_constraints encoded the old semantics and is split into test_initial_pool_ratio_overrides_typical_step + test_initial_pool_ratio_respects_constraint_floors accordingly. When a tier quota is below the floor bytes, construction rounds the quota up (pre-existing _compute_slot_count_for_level behavior); the estimation accounting stays correct because it adds back the actually-allocated pool bytes.

Verification

A/B on the previously always-failing config (DeepSeek-V4-Pro agentx disagg gen-only, GB300, dep4/ADP, MTP3, c32, max_batch_size=32, max_num_tokens=128, pool_ratio=[0.3,0.4,0.3]), GEN rank 0:

	main (no fix)	max_num_tokens only	this PR
estimation temp quota	23.69 GiB	8.14 GiB	8.14 GiB (8.22 allocated with floors)
estimation-phase graph-warmup coverage	bs=32→1, greedy+advanced	bs=1 only	bs=32→1, greedy+advanced
profiling peak	260.21 GiB	239.54 GiB (unrepresentative)	244.74 GiB
estimated max KV memory	36.09 GiB	40.70 GiB	36.09 GiB
result	done	OOM in bs=32 advanced-sampling capture	done (benchmark runs to completion)

The same arithmetic was verified for the tep8 shape (max_batch_size=128, mtp0): estimation warmup previously capped at bs=64, now fits bs=128 at the reduced 19.46 GiB quota.

A second end-to-end run with the final code reproduced the result (estimate 35.94 GiB, run-to-run peak jitter ~0.2 GiB) and exercised the graphs-disabled path on the CTX worker (cuda_graph_config: null): its generation-graph constraint is dropped, and the (pre-existing) capacity-limited skip of the CTX 128-request general-warmup shape is now visible via the new warning instead of silent.

Unit red/green inside the container:

• test_initial_pool_ratio_respects_constraint_floors fails on the old storage manager (pool group pinned at 20 slots vs the required scaled floor of 132) and passes with the fix; the full kv_cache_manager_v2 suite (75 tests) passes unchanged.
• test_deepseek_v4_pool_ratio_overrides_typical_step_but_keeps_constraints and test_deepseek_v4_constraints_cover_cuda_graph_warmup_batch fail on stock code and pass with the fix.

Note on non-DSV4 MLA models (e.g. DeepSeek-V3/R1 with the V1 manager): they now receive the real max_num_tokens instead of the constructor default 8192, which only affects the C++ manager's chunk_size = min(max_num_tokens, max_seq_len) — the same value every non-MLA model already gets.

Related: #16236 (sibling stale-_indexer_compress_ratio fix from the same DSV4 port; independent of this change).

Test Coverage

• tests/unittest/_torch/executor/test_kv_cache_estimation.py::test_mla_branch_forwards_max_num_tokens_to_manager
• tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py::test_deepseek_v4_pool_ratio_overrides_typical_step_but_keeps_constraints
• tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py::test_deepseek_v4_constraints_cover_cuda_graph_warmup_batch
• tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py::test_deepseek_v4_warmup_floor_shrinks_to_actual_graph_batch
• tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py::test_deepseek_v4_no_warmup_floor_when_graphs_disabled
• tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py::TestInitRatioConfig::test_small_constraint_floors_do_not_inflate_quota
• tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py::TestInitRatioConfig::test_initial_pool_ratio_respects_constraint_floors (red/green against the old storage manager)
• tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py::TestInitRatioConfig::test_initial_pool_ratio_overrides_typical_step
• tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py::TestInitRatioConfig::test_constraint_floors_tolerate_degenerate_max_util_for_resume

PR Checklist

• PR title follows the [None][fix] format
• Commit is signed off (DCO)
• Pre-commit hooks pass on the changed files
• Regression tests added and red/green verified
• End-to-end verified on the previously failing GB300 benchmark config

[qiaoxj07]

/bot run

[coderabbitai]

📝 Walkthrough

Walkthrough

The MLA KV cache initialization path now forwards max_num_tokens to the cache manager. A regression test verifies that the configured value reaches the manager constructor.

Changes

MLA KV Cache Forwarding

Layer / File(s)	Summary
Forward maximum token budget tensorrt_llm/_torch/pyexecutor/_util.py, tests/unittest/_torch/executor/test_kv_cache_estimation.py	The MLA cache manager receives max_num_tokens, and a regression test verifies the forwarded constructor argument.

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

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
Title check	✅ Passed	The title is specific, concise, and uses the required [None][fix] format while reflecting the main change.
Description check	✅ Passed	The PR includes Description, Test Coverage, and PR Checklist sections with substantive details, matching the template.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

🧹 Nitpick comments (1)

tests/unittest/_torch/executor/test_kv_cache_estimation.py (1)

487-506: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add return annotations to the new test functions.

The regression coverage is sufficient for this forwarding bug, but both test_mla_branch_forwards_max_num_tokens_to_manager and _RecordingManager.__init__ should declare -> None.

Proposed fix

-def test_mla_branch_forwards_max_num_tokens_to_manager():
+def test_mla_branch_forwards_max_num_tokens_to_manager() -> None:
...
-        def __init__(self, *args, **kwargs):
+        def __init__(self, *args: object, **kwargs: object) -> None:

As per coding guidelines, Python functions must be annotated with return types. As per path instructions, coverage is sufficient for this targeted regression and requires no follow-up test file.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/executor/test_kv_cache_estimation.py` around lines 487
- 506, Add -> None return annotations to the new test function
test_mla_branch_forwards_max_num_tokens_to_manager and the nested
_RecordingManager.__init__ method, without changing the existing regression
coverage.

Sources: Coding guidelines, Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tests/unittest/_torch/executor/test_kv_cache_estimation.py`:
- Around line 487-506: Add -> None return annotations to the new test function
test_mla_branch_forwards_max_num_tokens_to_manager and the nested
_RecordingManager.__init__ method, without changing the existing regression
coverage.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: afa6da09-a36e-4b12-bbac-68f76630a539

📥 Commits

Reviewing files that changed from the base of the PR and between 8bcec87 and d4f8edf.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/pyexecutor/_util.py
• tests/unittest/_torch/executor/test_kv_cache_estimation.py

[github-actions]

👎 Promotion blocked, new vulnerability found

Vulnerability report

Component	Vulnerability	Description	Severity
pytorch	CVE-2025-3000	A vulnerability classified as critical has been found in PyTorch 2.6.0. This affects the function torch.jit.script. The manipulation leads to memory corruption. It is possible to launch the attack on the local host. The exploit has been disclosed to the public and may be used.	MEDIUM

[qiaoxj07]

/bot run --disable-fail-fast

[qiaoxj07]

/bot run

[tensorrt-cicd]

PR_Github #58748 [ run ] triggered by Bot. Commit: d61da26 Link to invocation

[tensorrt-cicd]

PR_Github #58749 [ run ] triggered by Bot. Commit: d61da26 Link to invocation

[tensorrt-cicd]

PR_Github #58748 [ run ] completed with state ABORTED. Commit: d61da26

Link to invocation

[tensorrt-cicd]

PR_Github #58749 [ run ] completed with state SUCCESS. Commit: d61da26
/LLM/main/L0_MergeRequest_PR pipeline #47331 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

CI Agent Failure Analysis

Link to invocation

[longlee0622]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58759 [ run ] triggered by Bot. Commit: d61da26 Link to invocation

[qiaoxj07]

/bot run

[tensorrt-cicd]

PR_Github #58766 [ run ] triggered by Bot. Commit: 6d1bb26 Link to invocation

[tensorrt-cicd]

PR_Github #58759 [ run ] completed with state ABORTED. Commit: d61da26

Link to invocation

[tensorrt-cicd]

PR_Github #58773 [ run ] triggered by Bot. Commit: cbe98e4 Link to invocation

[tensorrt-cicd]

PR_Github #58766 [ run ] completed with state ABORTED. Commit: 6d1bb26

Link to invocation

[tensorrt-cicd]

PR_Github #58773 [ run ] completed with state SUCCESS. Commit: cbe98e4
/LLM/main/L0_MergeRequest_PR pipeline #47355 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

CI Agent Failure Analysis

Link to invocation

[longlee0622]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58776 [ run ] triggered by Bot. Commit: cbe98e4 Link to invocation

[tensorrt-cicd]

PR_Github #58776 [ run ] completed with state SUCCESS. Commit: cbe98e4
/LLM/main/L0_MergeRequest_PR pipeline #47358 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

CI Agent Failure Analysis

Link to invocation

[qiaoxj07]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58806 [ run ] triggered by Bot. Commit: cbe98e4 Link to invocation

[tensorrt-cicd]

PR_Github #58806 [ run ] completed with state DISABLED
Pipeline is freezed and top-1 instance is under maintenance.

Link to invocation

[chzblych]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58833 [ run ] triggered by Bot. Commit: cbe98e4 Link to invocation

[tensorrt-cicd]

PR_Github #58833 [ run ] completed with state SUCCESS. Commit: cbe98e4
/LLM/main/L0_MergeRequest_PR pipeline #47386 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

CI Agent Failure Analysis

Link to invocation

[jiaganc]

Superseded by #16311, which replaces this PR with the scoped KV cache estimation fix.