[https://nvbugs/6368562][fix] Reserve fp8 context-MLA attention workspace in KV cache estimation - #16399
[https://nvbugs/6368562][fix] Reserve fp8 context-MLA attention workspace in KV cache estimation#16399eopXD wants to merge 1 commit into
Conversation
|
/bot run |
|
PR_Github #59337 [ run ] triggered by Bot. Commit: |
WalkthroughAdds a shared FP8 Context-MLA workspace estimator, exposes it to Python, reserves workspace in KV-cache capacity calculations, and caps scheduled context requests by attended KV length. Tests cover reuse accounting, admission trimming, configuration, and non-MLA behavior. ChangesFP8 Context-MLA Workspace Reservation
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant KvCacheCreator
participant thop
participant AttentionOp
PyExecutor->>KvCacheCreator: configure KV-cache capacity
KvCacheCreator->>thop: estimate Context-MLA workspace bytes/token
thop->>AttentionOp: contextMlaWorkspaceBytesPerToken
AttentionOp-->>thop: workspace bytes/token
thop-->>KvCacheCreator: workspace bytes/token
KvCacheCreator-->>PyExecutor: carry attended-KV cap
PyExecutor->>PyExecutor: trim scheduled context requests
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tensorrt_llm/_torch/pyexecutor/_util.py (1)
222-243: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReuse
config_utils.is_mla()instead of re-deriving the MLA check.
is_mla = getattr(config, "kv_lora_rank", None) is not Noneonly checkskv_lora_rank, whiletensorrt_llm/_torch/pyexecutor/config_utils.py::is_mla()requires bothkv_lora_rankandqk_rope_head_dim. Functionally equivalent today but duplicated logic risks silent drift between the two checks.♻️ Proposed refactor
- config = model_config.pretrained_config - is_mla = getattr(config, "kv_lora_rank", None) is not None - if not is_mla: + from tensorrt_llm._torch.pyexecutor.config_utils import is_mla + config = model_config.pretrained_config + if not is_mla(config): return 0🤖 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 `@tensorrt_llm/_torch/pyexecutor/_util.py` around lines 222 - 243, Update get_mla_context_workspace_bytes_per_token to reuse config_utils.is_mla(config) for MLA detection instead of checking only config.kv_lora_rank. Preserve the existing early return for non-MLA models and use the shared helper so both kv_lora_rank and qk_rope_head_dim are validated consistently.tests/unittest/_torch/executor/test_mla_workspace_reserve.py (1)
1-110: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTest coverage: add cases for
_get_ctx_mla_kv_len_capderivation andget_mla_context_workspace_bytes_per_token's fp8/sparse branches.Current tests are sufficient for the pure-trim/pure-attended-length logic (
_context_attended_kv_len,_cap_context_by_total_kv_len), but_get_ctx_mla_kv_len_capis only exercised via a pre-set_ctx_mla_kv_len_capbypass, andget_mla_context_workspace_bytes_per_tokenis only tested on its non-MLA early-return. Consider adding, in this file:
- a test for
_get_ctx_mla_kv_len_capderivingblocks * tokens_per_blockfrom a mockedkv_cache_manager(and thew == 0→Nonebranch),- a test for
get_mla_context_workspace_bytes_per_tokenwithfp8_context_mla=True/sparse_mla=Trueby mockingtensorrt_llm.bindings.internal.thop.As per path instructions, coverage of the pure-Python reuse/trim logic is sufficient, but these two gaps are worth a follow-up since they directly gate whether the workspace reservation is ever activated.
🤖 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_mla_workspace_reserve.py` around lines 1 - 110, Extend this test module with coverage for PyExecutor._get_ctx_mla_kv_len_cap, verifying it derives blocks multiplied by tokens_per_block from a mocked kv_cache_manager and returns None when the workspace value is zero. Add a get_mla_context_workspace_bytes_per_token test for fp8_context_mla=True and sparse_mla=True, mocking tensorrt_llm.bindings.internal.thop and asserting the resulting workspace calculation.Source: 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.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 5108-5133: Invalidate the cached value set by
_get_ctx_mla_kv_len_cap whenever _maybe_rebalance_kv_pools invokes
mgr.impl.adjust() and changes the primary-pool capacity. Clear or recompute
_ctx_mla_kv_len_cap after the rebalance so subsequent scheduling reads the
updated blocks_in_primary_pool and tokens_per_block values.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 222-243: Update get_mla_context_workspace_bytes_per_token to reuse
config_utils.is_mla(config) for MLA detection instead of checking only
config.kv_lora_rank. Preserve the existing early return for non-MLA models and
use the shared helper so both kv_lora_rank and qk_rope_head_dim are validated
consistently.
In `@tests/unittest/_torch/executor/test_mla_workspace_reserve.py`:
- Around line 1-110: Extend this test module with coverage for
PyExecutor._get_ctx_mla_kv_len_cap, verifying it derives blocks multiplied by
tokens_per_block from a mocked kv_cache_manager and returns None when the
workspace value is zero. Add a get_mla_context_workspace_bytes_per_token test
for fp8_context_mla=True and sparse_mla=True, mocking
tensorrt_llm.bindings.internal.thop and asserting the resulting workspace
calculation.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4db55ce4-09dd-4367-b5cd-5d6976e4c3e4
📒 Files selected for processing (7)
cpp/tensorrt_llm/common/attentionOp.cppcpp/tensorrt_llm/common/attentionOp.hcpp/tensorrt_llm/nanobind/thop/bindings.cpptensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/py_executor.pytests/integration/test_lists/waives.txttests/unittest/_torch/executor/test_mla_workspace_reserve.py
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
|
PR_Github #59337 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59381 [ run ] triggered by Bot. Commit: |
|
PR_Github #59381 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59458 [ run ] triggered by Bot. Commit: |
| kv_len_cap = get_mla_context_workspace_kv_len_cap( | ||
| self._kv_cache_config, self._max_batch_size, self._max_num_tokens, | ||
| self._max_seq_len) | ||
| if w_bytes_per_token > 0 and kv_len_cap: |
There was a problem hiding this comment.
This branch currently reserves workspace whenever w_bytes_per_token > 0, including configurations where runtime workspace demand cannot exceed the profiled peak.
With block reuse disabled on the non-chunked path, summed attended KV is bounded by the processed-token budget, so profiling already covers the FP8 staging requirement. With chunked prefill, each AttentionOp launch is independently bounded by the chunk buffer. Reserving the workspace again therefore double-counts it.
This can cause a substantial KV-cache capacity regression for unaffected configurations. For Kimi-K2 with attention-DP, k ≈ 35,136 B/token and w = 20,480 B/token, so the clamped split can reduce KV-cache capacity by roughly 37% even when reuse-driven workspace growth is impossible.
Please apply the reserve and admission cap only when block reuse is enabled, chunked prefill is disabled, and the selected backend can reach the dense TRTLLM full-gather path. Please also add no-op coverage for the unaffected configurations.
There was a problem hiding this comment.
Good point as pengbo also mentioned this.
Addressed in the latest diff. Please check.
1d6539d to
3ba6487
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #60620 [ run ] triggered by Bot. Commit: |
|
PR_Github #60620 [ run ] completed with state
|
3ba6487 to
1eb4bd8
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #60855 [ run ] triggered by Bot. Commit: |
|
PR_Github #60855 [ run ] completed with state
|
…pace in KV cache estimation TestKimiK2::test_nvfp4[4gpus] (NVFP4, TP4 + attention-DP, block reuse) OOMs mid-forward in the MLA context attention. The fp8 context-MLA K/V dequant workspace is one buffer shared across attention layers whose size scales with the summed attended KV length (total_kv_len) of the step's context requests. The KV-cache estimator profiles fresh-prefill dummies against an empty cache, so total_kv_len there sits near max_num_tokens and the workspace is at its floor; with block reuse at serving time total_kv_len decouples from max_num_tokens and the workspace grows past the floor, but the estimator has already handed that headroom to the KV pool. The under-reservation was latent until NVIDIA#14852 sized the workspace by total_kv_len. Reserve for the workspace during estimation instead of lowering free_gpu_memory_fraction (a user co-tenancy knob): - Only reserve when block reuse is enabled and chunked prefill is disabled -- the sole conditions under which total_kv_len can exceed the profiled floor. Without reuse the workspace is bounded by max_num_tokens (already profiled); with chunked prefill each attention launch is independently bounded by its own chunk buffer. Reserving in those configs would double-count and needlessly shrink the KV pool (up to ~37% for Kimi-K2 attention-DP). No-op there and for non-fp8-MLA models. - Reserve w * L_cap bytes, where w is the per-token workspace cost and L_cap is the never-stall worst-case summed attended KV per step, min(max_batch_size, max_num_tokens) * max_seq_len. Clamp the reserve to the per-token split budget * w / (k + w) so a memory-constrained node shares the budget at a common token count instead of starving the pool; equivalently the pool keeps max((budget - w*L_cap)/k, budget/(k + w)) tokens. - The estimator carries the exact cap it reserved for (min(L_cap, budget/(k+w))) onto the KV manager; the scheduler reads it directly and trims context requests whose summed attended total_kv_len would exceed it, always keeping one request as a forward-progress guard. It does not re-derive the cap from pool layout, which KV-cache-manager V2 overstates (blocks_in_primary_pool forwards get_page_index_upper_bound, not the available-page count). A carried cap of None (no reservation) applies no admission cap. - w counts the fp8 K/V dequant staging buffer. A sparse-MLA model normally stages nothing, but with the short-seq MHA fallback enabled (TRTLLM_MLA_SHORT_SEQ_MHA_THRESHOLD > 0) it routes short sequences through the dense context path that does stage the buffer, so reserve for that case too. - KvCacheConfig.fp8_context_mla_kv_len_cap (prototype) overrides L_cap to trade reserved workspace for KV pool; the scheduler enforces it. w is a single source of truth in C++ (AttentionOp::contextMlaWorkspaceBytesPerToken, guarded against getWorkspaceSizeForContext by a TLLM_CHECK) exposed via nanobind, so the reserve cannot drift from the runtime allocation. This accounts for the fp8 staging term only; the separate BF16 full-gather buffers on the reuse path are a follow-up, so the reserve bounds but does not by itself eliminate reuse-driven OOM. Re-enable TestKimiK2::test_nvfp4[4gpus] (remove the nvbugs/6368562 waive). Co-Authored-By: Yueh-Ting Chen <yueh.ting.chen@gmail.com> Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com>
1eb4bd8 to
9f0571a
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #60949 [ run ] triggered by Bot. Commit: |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/unittest/_torch/executor/test_mla_workspace_reserve.py (1)
228-259: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd coverage for a carried cap of
0.Related to the
carried else Nonetruthiness issue flagged inpy_executor.py's_get_ctx_mla_kv_len_cap: none of themanagerfixtures here usefp8_ctx_mla_kv_len_cap=0, so the case where a legitimately-tiny reservation collapses to "no cap" is untested.🧪 Proposed test addition
def test_ctx_cap_zero_is_not_treated_as_no_cap(): # A carried cap of exactly 0 means "admit almost nothing", not "unlimited". exe = _executor_with_manager(SimpleNamespace(fp8_ctx_mla_kv_len_cap=0)) assert exe._get_ctx_mla_kv_len_cap() == 0🤖 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_mla_workspace_reserve.py` around lines 228 - 259, Add coverage for a manager carrying fp8_ctx_mla_kv_len_cap equal to 0, verifying _get_ctx_mla_kv_len_cap() returns 0 rather than None. Add a focused test alongside test_ctx_cap_none_when_no_reservation using SimpleNamespace(fp8_ctx_mla_kv_len_cap=0), preserving the distinction between zero and an absent or None cap.tensorrt_llm/_torch/pyexecutor/py_executor.py (1)
5267-5287: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winDeferred context requests keep any KV growth already applied during scheduling — same tradeoff as
_waiting_requests, but undocumented here.
context_requests[:i]drops the tail without calling_revert_ctx_allocfor the dropped requests, even thoughscheduler.schedule_request()may have already grown their KV cache capacity for this chunk (V2). This mirrors the already-documented, accepted tradeoff for_waiting_requests("still occupy KV cache and may reduce the batch size available for generation requests"), but that rationale isn't restated here, and this cap can trigger far more often (any reuse-heavy step) than the waiting-queue path.🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 5267 - 5287, The deferral path in _cap_context_by_total_kv_len must explicitly document that dropped tail requests may retain KV-cache capacity grown by scheduler.schedule_request(), matching the accepted _waiting_requests tradeoff. Update the method docstring to state that deferred requests remain active, retain any already-applied KV growth, and may reduce capacity for generation requests; preserve the existing slicing behavior.
🤖 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.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 5242-5251: Update the carried-cap assignment in the
`_ctx_mla_kv_len_cap` lookup to distinguish `None` from a legitimate zero value:
convert any non-None `fp8_ctx_mla_kv_len_cap` to int, while preserving None only
when no cap was carried. Ensure a carried value of 0 remains 0 and therefore
enforces the intended near-zero admission cap.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 5267-5287: The deferral path in _cap_context_by_total_kv_len must
explicitly document that dropped tail requests may retain KV-cache capacity
grown by scheduler.schedule_request(), matching the accepted _waiting_requests
tradeoff. Update the method docstring to state that deferred requests remain
active, retain any already-applied KV growth, and may reduce capacity for
generation requests; preserve the existing slicing behavior.
In `@tests/unittest/_torch/executor/test_mla_workspace_reserve.py`:
- Around line 228-259: Add coverage for a manager carrying
fp8_ctx_mla_kv_len_cap equal to 0, verifying _get_ctx_mla_kv_len_cap() returns 0
rather than None. Add a focused test alongside
test_ctx_cap_none_when_no_reservation using
SimpleNamespace(fp8_ctx_mla_kv_len_cap=0), preserving the distinction between
zero and an absent or None cap.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4906e7a2-c551-4133-9439-9538d6daea1a
📒 Files selected for processing (11)
cpp/tensorrt_llm/common/attentionOp.cppcpp/tensorrt_llm/common/attentionOp.hcpp/tensorrt_llm/nanobind/thop/bindings.cpptensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/integration/test_lists/waives.txttests/unittest/_torch/executor/test_dual_pool_kv_cache.pytests/unittest/_torch/executor/test_kv_cache_estimation.pytests/unittest/_torch/executor/test_mla_workspace_reserve.py
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
🚧 Files skipped from review as they are similar to previous changes (3)
- cpp/tensorrt_llm/common/attentionOp.h
- cpp/tensorrt_llm/nanobind/thop/bindings.cpp
- cpp/tensorrt_llm/common/attentionOp.cpp
| cap = getattr(self, "_ctx_mla_kv_len_cap", "unset") | ||
| if cap != "unset": | ||
| return cap | ||
| if getattr(self, "is_warmup", False): | ||
| # Estimation/warmup runs fresh-prefill dummies against a throwaway manager that reserved | ||
| # nothing; don't cap them (and don't cache -- real serving recomputes from the real manager). | ||
| return None | ||
| carried = getattr(self.kv_cache_manager, "fp8_ctx_mla_kv_len_cap", None) | ||
| self._ctx_mla_kv_len_cap = int(carried) if carried else None | ||
| return self._ctx_mla_kv_len_cap |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
carried == 0 is silently treated as "no cap" instead of "admit almost nothing".
self._ctx_mla_kv_len_cap = int(carried) if carried else None uses truthiness, so a carried cap of exactly 0 — which get_mla_context_workspace_reserve can legitimately return as int(reserve / w_bytes_per_token) when the reservation is tiny relative to w_bytes_per_token (a very memory-constrained node) — collapses to None, meaning "no admission cap" (unlimited) rather than the intended "cap admission to ~0 tokens". This inverts the safety property for exactly the tight-budget case this reservation is meant to protect.
🛡️ Proposed fix
- carried = getattr(self.kv_cache_manager, "fp8_ctx_mla_kv_len_cap", None)
- self._ctx_mla_kv_len_cap = int(carried) if carried else None
+ carried = getattr(self.kv_cache_manager, "fp8_ctx_mla_kv_len_cap", None)
+ self._ctx_mla_kv_len_cap = int(carried) if carried is not None else None📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| cap = getattr(self, "_ctx_mla_kv_len_cap", "unset") | |
| if cap != "unset": | |
| return cap | |
| if getattr(self, "is_warmup", False): | |
| # Estimation/warmup runs fresh-prefill dummies against a throwaway manager that reserved | |
| # nothing; don't cap them (and don't cache -- real serving recomputes from the real manager). | |
| return None | |
| carried = getattr(self.kv_cache_manager, "fp8_ctx_mla_kv_len_cap", None) | |
| self._ctx_mla_kv_len_cap = int(carried) if carried else None | |
| return self._ctx_mla_kv_len_cap | |
| cap = getattr(self, "_ctx_mla_kv_len_cap", "unset") | |
| if cap != "unset": | |
| return cap | |
| if getattr(self, "is_warmup", False): | |
| # Estimation/warmup runs fresh-prefill dummies against a throwaway manager that reserved | |
| # nothing; don't cap them (and don't cache -- real serving recomputes from the real manager). | |
| return None | |
| carried = getattr(self.kv_cache_manager, "fp8_ctx_mla_kv_len_cap", None) | |
| self._ctx_mla_kv_len_cap = int(carried) if carried is not None else None | |
| return self._ctx_mla_kv_len_cap |
🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 5242 - 5251,
Update the carried-cap assignment in the `_ctx_mla_kv_len_cap` lookup to
distinguish `None` from a legitimate zero value: convert any non-None
`fp8_ctx_mla_kv_len_cap` to int, while preserving None only when no cap was
carried. Ensure a carried value of 0 remains 0 and therefore enforces the
intended near-zero admission cap.
|
PR_Github #60949 [ run ] completed with state
|
pengbowang-nv
left a comment
pengbowang-nv
left a comment
There was a problem hiding this comment.
Attention Part LGTM
| # Estimation/warmup runs fresh-prefill dummies against a throwaway manager that reserved | ||
| # nothing; don't cap them (and don't cache -- real serving recomputes from the real manager). | ||
| return None | ||
| carried = getattr(self.kv_cache_manager, "fp8_ctx_mla_kv_len_cap", None) |
There was a problem hiding this comment.
Non blocking: no need getattr here also
|
/bot run --disable-fail-fast |
|
PR_Github #61023 [ run ] triggered by Bot. Commit: |
|
PR_Github #61023 [ run ] completed with state
|
|
[by Codex] @SimengLiu-nv Could you please review this PR? Thank you! |
|
/bot run --disable-fail-fast |
|
PR_Github #61772 [ run ] triggered by Bot. Commit: |
|
PR_Github #61772 [ run ] completed with state |
SimengLiu-nv
left a comment
SimengLiu-nv
left a comment
There was a problem hiding this comment.
Approve given the two comments will be addressed.
| return self.slope * tokens + self.intercept | ||
|
|
||
|
|
||
| def get_mla_context_workspace_bytes_per_token(model_config, mapping) -> int: |
There was a problem hiding this comment.
Question: does this function only apply to TRTLLM MLA or other sources as well like flashinfer? If only to TRTLLM MLA, it would be necessary to check the attention backend.
| # sequences through the dense context path that does stage the fp8 buffer. So the buffer is truly never | ||
| # staged only for sparse MLA with the fallback off; reserve for it otherwise (conservative -- may | ||
| # over-reserve for models that never take the fallback, but never under-reserves). | ||
| sparse_mla = model_config.sparse_attention_config is not None |
There was a problem hiding this comment.
[CODEX][P1] Do not use global sparse-config presence as the staging-free MLA workspace gate
The estimator currently assumes that any model with a sparse-attention configuration never needs the dense FP8 context-MLA K/V staging buffers:
sparse_mla = model_config.sparse_attention_config is not None
short_seq_mha_enabled = ...
stages_no_buffer = sparse_mla and not short_seq_mha_enabledstages_no_buffer is then passed to contextMlaWorkspaceBytesPerToken(). When it is true, the C++ helper returns zero, so KV-cache capacity estimation reserves no workspace and the executor installs no total_kv_len admission cap.
The problem is that sparse_attention_config is not None is only a model-level configuration check. It is not equivalent to the runtime predicate that determines whether the K/V staging buffers are allocated.
Runtime behavior is layer-specific
Each MLA layer independently converts the global sparse configuration into runtime parameters:
sparse_params = sparse_attn_cfg.to_sparse_params(
pretrained_config=config.pretrained_config,
layer_idx=self.layer_idx,
)Depending on the algorithm and layer, this can produce DSA parameters, DeepSeek-V4 parameters, skip-softmax parameters, or None for a layer excluded by the sparse configuration.
The C++ attention operator only enables staging-free sparse MLA when sparse top-k indices are actually supplied:
if (num_sparse_topk_value > 0
&& sparse_attn_indices.has_value()
&& sparse_attn_indices.value().numel() > 0)
{
op->mUseSparseAttention = true;
}AttentionOp::useSparseMLA() additionally requires TRTLLM-gen and MLA to be active. Only when that complete runtime predicate is true does getWorkspaceSizeForContext() set the FP8 K/V staging sizes to zero.
Otherwise, it allocates:
fp8_k_buf_size = kv_buf_tokens * total_k_dim_all_heads;
fp8_v_buf_size = kv_buf_tokens * total_v_dim_all_heads;where kv_buf_tokens grows with total_kv_len.
Skip-softmax is a concrete counterexample
Skip-softmax is represented by SparseAttentionConfig, but it does not use the sparse-MLA absorption path.
TrtllmAttention explicitly excludes SkipSoftmaxParams from sparse-index generation:
if sparse_params is not None and not isinstance(
sparse_params, SkipSoftmaxParams):
# Generate sparse KV and attention indices.For skip-softmax, the backend only passes threshold parameters to the attention kernel. It does not provide num_sparse_topk and sparse attention indices, so C++ leaves mUseSparseAttention false. Consequently, useSparseMLA() is false and the dense FP8 K/V staging buffers are still allocated.
The estimator nevertheless sees the global SkipSoftmaxAttentionConfig, sets stages_no_buffer=True, and reports w=0.
Resulting failure sequence
A reachable problematic configuration is:
- MLA model using the TRTLLM backend.
- FP8 KV cache enabled.
SkipSoftmaxAttentionConfigpresent.- Block reuse enabled.
- Chunked prefill disabled.
- Shared prefixes cause
total_kv_lento grow beyond themax_num_tokensfloor exercised during profiling.
For this configuration, the estimator reserves no staging workspace and installs no admission cap. At runtime, the dense MLA path still resizes the FP8 K/V buffers according to the larger total_kv_len. The resize can therefore consume memory already assigned to the KV pool and reproduce the mid-forward OOM this PR is intended to mitigate.
A layer-specific exclusion has the same problem: the global configuration remains non-None, but an excluded layer resolves sparse_params=None and uses dense MLA. Because the attention workspace is shared across layers, one reachable dense layer is enough to require the reservation.
The TRTLLM_MLA_SHORT_SEQ_MHA_THRESHOLD special case only handles one DSA-specific dense fallback. It does not address skip-softmax, layer exclusions, or another sparse configuration that does not guarantee useSparseMLA().
Suggested fix
Please derive the staging-free decision from the actual per-layer/runtime capability rather than global configuration presence.
A conservative implementation should return a nonzero workspace cost whenever any MLA layer can reach dense context attention. Only configurations that guarantee the sparse-MLA absorption path for every applicable layer, with all dense fallbacks disabled, should return zero.
Suggested regression coverage:
- DSA absorption with no dense fallback:
w == 0 - DSA with short-sequence MHA fallback enabled:
w > 0 - Skip-softmax MLA:
w > 0 - Sparse configuration excluded on one MLA layer:
w > 0
I identified this through static inspection rather than a runtime reproduction. The exact OOM requires sufficient prefix reuse and memory pressure, but the estimator and runtime predicates are directly inconsistent.
11 participants
Description
TestKimiK2::test_nvfp4[4gpus](Kimi-K2-Thinking NVFP4, TP4 + attention-DP, block reuse) OOMsmid-forward inside the MLA context attention (
mla_custom_op→thop.attentionworkspace resize).Root cause. The fp8 context-MLA K/V dequant workspace is a single buffer (reused across attention
layers) whose size scales with the summed attended KV length (
total_kv_len) of the context requests ina forward step. The KV-cache memory estimator profiles with fresh-prefill dummy requests against an empty
KV cache, so during profiling
total_kv_lenis pinned nearmax_num_tokensand this workspace sits atits floor. With block reuse at serving time (e.g. MMLU's shared few-shot prefixes)
total_kv_lendecouplesfrom
max_num_tokensand the workspace grows far past the profiled floor — but the estimator has alreadyhanded that headroom to the KV pool, so the workspace has nowhere to grow. This latent under-reservation
was exposed once #14852 correctly sized the workspace by
total_kv_len(before that it was under-sized,causing an OOB instead).
Fix (estimation, not the memory fraction —
free_gpu_memory_fractionis a user co-tenancy contract).Reserve KV-cache headroom for the workspace up front, bounded so it never over-reserves:
which
total_kv_lencan exceed the profiled floor. Without reuse, summed attended KV is bounded bymax_num_tokens, exactly what the profiling forward exercises; with chunked prefill, each attentionlaunch is independently bounded by its own chunk buffer. Reserving in those configurations would
double-count and needlessly shrink the KV pool (up to ~37% for Kimi-K2 with attention-DP). No-op there,
and for non-fp8-MLA models — addresses @QiJune's and @pengbowang-nv's blocking review.
w * L_capbytes, wherewis the per-token workspace cost andL_capis the never-stallworst-case summed attended KV per step,
min(max_batch_size, max_num_tokens) * max_seq_len.budget * w / (k + w)so a memory-constrained node — wherereserving the full worst case would starve the pool — shares the budget at a common token count instead.
Equivalently, the pool keeps
max((budget − w·L_cap)/k, budget/(k + w))tokens (per @pengbowang-nv'sreview: the worst case has an upper cap, so most deployments reserve only a small fixed amount rather
than a fixed proportion of the budget).
min(L_cap, budget/(k+w)), i.e.reserve/w) onto the KV manager; the scheduler reads it directly and trims context requests whose summedattended
total_kv_lenwould exceed it, always keeping at least one request as a forward-progress guard.It does not re-derive the cap from pool layout, which KV-cache-manager V2 overstates
(
blocks_in_primary_poolforwardsget_page_index_upper_bound, not the available-page count) —addresses @QiJune's blocking review. A carried cap of
None(nothing reserved) simply applies noadmission cap.
wcounts the fp8 K/V dequant staging buffer. A sparse-MLA model normally stages nothing, but with theshort-seq MHA fallback enabled (
TRTLLM_MLA_SHORT_SEQ_MHA_THRESHOLD > 0) it routes short sequencesthrough the dense context path that does stage the buffer, so we reserve for that reachable case too
(conservative) — addresses @QiJune's sparse-gate review.
KvCacheConfig.fp8_context_mla_kv_len_cap(prototype) overridesL_capto trade reserved workspace forKV pool; safe at any value because the scheduler enforces it (floored at
max_seq_len, capped at theworst case).
The per-token cost
wis a single source of truth in C++(
AttentionOp::contextMlaWorkspaceBytesPerToken, guarded againstgetWorkspaceSizeForContextby aTLLM_CHECK) and exposed to the Python estimator via a nanobind binding, so the reserve cannot drift fromthe runtime allocation.
Scope / limitations. This is a targeted, monotonic mitigation of the fp8 context-MLA staging
workspace.
forward_context_with_cached_kv()also retains additional BF16 reuse-scaled buffers (full_k/full_kv) that this reservation does not yet account for, so it reduces — but does not by itself eliminate— reuse-driven mid-forward OOM in every configuration. Full-path accounting and a high-fanout shared-prefix
memory test are tracked as follow-ups.
Changes:
cpp/.../common/attentionOp.{h,cpp}:contextMlaWorkspaceBytesPerToken()helper (single source of truth).cpp/.../nanobind/thop/bindings.cpp:get_context_mla_workspace_bytes_per_tokenbinding._torch/pyexecutor/_util.py: fold the reuse/chunked-prefill reservation gate intoget_mla_context_workspace_kv_len_cap()(returnsNonewhen no reservation is needed) and addget_mla_context_workspace_reserve(); reserve the workspace and carry the admission cap onto the KVmanager in
configure_kv_cache_capacity/build_managers._torch/pyexecutor/py_executor.py: read the carried cap and trim summed context attended-KV in_schedule(no admission cap when nothing was reserved).llmapi/llm_args.py:KvCacheConfig.fp8_context_mla_kv_len_capprototype override.TestKimiK2::test_nvfp4[4gpus](remove the nvbugs/6368562 waive).Test Coverage
tests/unittest/_torch/executor/test_mla_workspace_reserve.py(new): unit coverage for the reservationgate (reserve only when block reuse is on and chunked prefill off; no-op otherwise), the reserve/cap math
(
get_mla_context_workspace_reserve, both the worst-case-fits and memory-constrained branches), theL_capderivation (default worst case, override floor/ceiling), the per-request attended-KV computation(V1/V2 reuse timing, chunk clamp), the admission trim (tail trim, keep-all, first-request-always-kept
guard, no-cap no-op), the carried-cap read (V1/V2 layout-independent; no cap when the carried value
is absent or
None), and the non-MLA gate.accuracy/test_llm_api_pytorch.py::TestKimiK2::test_nvfp4[4gpus]: the original repro, re-enabled here.PR Checklist
KvCacheConfig.fp8_context_mla_kv_len_capis a new nested-config field — regeneratetensorrt_llm/usage/llm_args_golden_manifest.jsonand obtain telemetry/privacy CODEOWNER approval.GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Dev Engineer Review
AttentionOp::contextMlaWorkspaceBytesPerToken) and validates its computed fp8 context-MLA K/V buffer sizing in runtime workspace calculation to keep runtime workspace sizing and KV-cache estimation consistent.get_context_mla_workspace_bytes_per_token).fp8_context_mla_kv_len_cap(prototype, defaultNone) to override the derived cap (in tokens), while preserving prior behavior when not applicable (non-MLA / non-FP8 context-MLA / sparse-MLA / chunked-prefill paths).accuracy/test_llm_api_pytorch.py::TestKimiK2::test_nvfp4[4gpus](effectively re-enabling the test).CI follow-up
49127, associated with commit1eb4bd8.QA Engineer Review
Test-list change (test-list only):
tests/integration/test_lists/waives.txt: removedSKIP accuracy/test_llm_api_pytorch.py::TestKimiK2::test_nvfp4[4gpus].Test-code changes (outside test-list files):
tests/unittest/_torch/executor/test_mla_workspace_reserve.pyPyExecutor._context_attended_kv_len,PyExecutor._cap_context_by_total_kv_len, workspace reserve/cap estimation helpers (includingNonebehavior when reuse/chunking conditions prevent reserve growth), and executor-side cap carrying into the KV manager.tests/unittest/_torch/executor/test_dual_pool_kv_cache.pycreator._fp8_ctx_mla_kv_len_cap = Nonefor the additional attribute.tests/unittest/_torch/executor/test_kv_cache_estimation.pyget_mla_context_workspace_bytes_per_tokento return0for the specific estimation scenario.tests/integration/test_lists/.