[https://nvbugs/6368562][fix] Reserve fp8 context-MLA attention workspace in KV cache estimation - #16399
Conversation
|
/bot run |
|
PR_Github #59337 [ run ] triggered by Bot. Commit: |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
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
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: |
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.
9f0571a to
e5ae24a
Compare
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
cpp/tensorrt_llm/common/attentionOp.h (1)
63-69: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the new public API with Doxygen.
This newly added public interface uses regular
//comments. Convert the block to Doxygen comments so generated API documentation includes the estimator contract.Proposed fix
- // Per-token byte cost of the context-MLA K/V dequant staging buffers, whose size scales with the summed - // attended KV length (`total_kv_len`). Only the fp8 context-MLA separate-Q/KV path stages these buffers; - // every other path (incl. sparse MLA, which reads K/V straight from the paged cache) returns 0. Single - // source of truth shared by getWorkspaceSizeForContext (runtime sizing) and the KV-cache estimator, so - // the two cannot drift. + /** + * `@brief` Returns the per-token byte cost of FP8 context-MLA K/V staging buffers. + * + * The cost scales with attended KV length. Returns zero outside the FP8 + * context-MLA separate-Q/KV path, including sparse MLA. + */As per coding guidelines, new C++ interfaces must use Doxygen comments.
🤖 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 `@cpp/tensorrt_llm/common/attentionOp.h` around lines 63 - 69, Convert the comment immediately preceding contextMlaWorkspaceBytesPerToken into a Doxygen comment while preserving its existing description of buffer sizing, applicable paths, zero-return behavior, and shared usage by getWorkspaceSizeForContext and the KV-cache estimator.Source: Coding guidelines
🤖 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 `@cpp/tensorrt_llm/common/attentionOp.h`:
- Around line 63-69: Convert the comment immediately preceding
contextMlaWorkspaceBytesPerToken into a Doxygen comment while preserving its
existing description of buffer sizing, applicable paths, zero-return behavior,
and shared usage by getWorkspaceSizeForContext and the KV-cache estimator.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 1ef9297a-d156-4e56-97ce-a860eafcfc01
📒 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 (8)
- tensorrt_llm/usage/llm_args_golden_manifest.json
- tests/unittest/_torch/executor/test_kv_cache_estimation.py
- tensorrt_llm/llmapi/llm_args.py
- tests/unittest/_torch/executor/test_dual_pool_kv_cache.py
- cpp/tensorrt_llm/nanobind/thop/bindings.cpp
- cpp/tensorrt_llm/common/attentionOp.cpp
- tensorrt_llm/_torch/pyexecutor/py_executor.py
- tensorrt_llm/_torch/pyexecutor/_util.py
|
PR_Github #62402 [ run ] triggered by Bot. Commit: |
|
PR_Github #62402 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
…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, which the runtime skips only where AttentionOp::useSparseMLA() holds: a DSA / DeepSeek-V4 model (the only algorithms lowering to the absorption path that reads K/V from the paged cache) on an SM using TRTLLM-gen (sm >= 100 && sm != 120), with the short-seq MHA fallback off. Gate on that predicate rather than on a sparse configuration merely being present: skip-softmax passes no sparse indices to C++ and its ignore-list can exclude a layer, and the fallback (TRTLLM_MLA_SHORT_SEQ_MHA_THRESHOLD > 0) routes short sequences back through the dense path. The workspace is shared across layers, so one dense layer forces the reserve. - 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>
e5ae24a to
ff57ee0
Compare
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
🧹 Nitpick comments (2)
cpp/tensorrt_llm/common/attentionOp.h (2)
513-513: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winKeep the new member private or use public-member naming.
mForcePrepareSpecDecTreeMaskis declared beforeprivate:and is therefore publicly mutable despite using themprefix reserved for private members. Move it into the private section; public methods can still access it.As per coding guidelines,
m-prefixed members should be private class members.🤖 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 `@cpp/tensorrt_llm/common/attentionOp.h` at line 513, Move mForcePrepareSpecDecTreeMask into the class’s private section, preserving its existing type and default value; keep public APIs unchanged.Source: Coding guidelines
63-69: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the new public API with Doxygen.
contextMlaWorkspaceBytesPerTokenis a new public interface but is documented only with ordinary//comments. Use///or/** ... */documentation with a concise@briefso generated API documentation includes this contract.As per coding guidelines, new interfaces must use Doxygen comments.
🤖 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 `@cpp/tensorrt_llm/common/attentionOp.h` around lines 63 - 69, Convert the existing comment immediately preceding contextMlaWorkspaceBytesPerToken into Doxygen syntax and add a concise `@brief` describing its per-token workspace-byte calculation and scope. Preserve the existing contract details while ensuring generated API documentation associates them with contextMlaWorkspaceBytesPerToken.Source: Coding guidelines
🤖 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 `@cpp/tensorrt_llm/common/attentionOp.h`:
- Line 513: Move mForcePrepareSpecDecTreeMask into the class’s private section,
preserving its existing type and default value; keep public APIs unchanged.
- Around line 63-69: Convert the existing comment immediately preceding
contextMlaWorkspaceBytesPerToken into Doxygen syntax and add a concise `@brief`
describing its per-token workspace-byte calculation and scope. Preserve the
existing contract details while ensuring generated API documentation associates
them with contextMlaWorkspaceBytesPerToken.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2ec9e0a8-eadd-49f7-a14f-d78b35456501
📒 Files selected for processing (3)
cpp/tensorrt_llm/common/attentionOp.cppcpp/tensorrt_llm/common/attentionOp.hcpp/tensorrt_llm/nanobind/thop/bindings.cpp
🚧 Files skipped from review as they are similar to previous changes (2)
- cpp/tensorrt_llm/nanobind/thop/bindings.cpp
- cpp/tensorrt_llm/common/attentionOp.cpp
|
PR_Github #62512 [ run ] triggered by Bot. Commit: |
|
PR_Github #62512 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62697 [ run ] triggered by Bot. Commit: |
|
PR_Github #62697 [ run ] completed with state
|
|
/bot skip --comment "Single GPU pipelines have passed. Sufficient to prove the calculation is conservative and works. If multi-GPU pipeline tests are really failures from this MR we will further follow-up to resolve this." |
|
PR_Github #63357 [ skip ] triggered by Bot. Commit: |
|
PR_Github #63357 [ skip ] completed with state |
…end contract The fp8 context-MLA workspace reservation (nvbugs/6368562, NVIDIA#16399) was threaded imperatively through the KV-cache estimator, keyed on a model-config check specific to MLA rather than on the backend that allocates the buffer. Two reviewers flagged this on NVIDIA#16399: the reserve fires off a model-level MLA check, so a model running a backend that never stages the buffer is still charged for it -- shrinking the KV pool for a workspace it will not allocate. Lift the accounting into a declared contract on the attention backend: - AttentionBackend.runtime_workspace_bytes_per_token(model_config, mapping) returns the per-token bytes to reserve for a workspace the backend stages whose size scales with a runtime quantity the profiling forward does not drive to its serving maximum. Default 0 -- correct for every backend but fp8 context-MLA. - TrtllmAttention declares the fp8 context-MLA K/V dequant workspace, still sized by the single C++ source of truth (contextMlaWorkspaceBytesPerToken) and keeping NVIDIA#16399's runtime-matched sparse gate (dsa/deepseek_v4 on SM 100/103 with the short-seq MHA fallback off). - The estimator resolves the declaration through the model's selected backend via get_attention_workspace_bytes_per_token(). The reserve/cap math is unchanged; what changes is that a non-TRTLLM backend now correctly reserves nothing. - Document the contract in ATTENTION_DEVELOPER_GUIDE.md (required reading) so a new backend inherits the accounting instead of the OOM. The contract is deliberately a scalar per-token rate, not a typed driver/reservation abstraction: there is one driving quantity today (total_kv_len) and the scheduler's cap is specific to it, so a richer type would be unused scaffolding. A backend with a different driver introduces it then, alongside the enforcement it needs. Also carries two follow-through fixes from NVIDIA#16399 review threads that were resolved without a code change, both in the cap reader this contract feeds: - A carried cap of exactly 0 was collapsed to None ("no cap") by a truthiness check, inverting admission control for the tightest-budget case it exists to protect. Compare against None instead. - is_warmup is a real property on PyExecutor, so the defensive getattr is unnecessary. Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com>
15 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(...), returning non-zero only for the fp8 Context-MLA separate-Q/KV non-sparse path and validated viaTLLM_CHECKagainst derived K/V sizing inAttentionOp::getWorkspaceSizeForContext(...).thop.get_context_mla_workspace_bytes_per_token(...)(used for kv-cap math to prevent runtime vs estimator drift).kv_cache_manager.fp8_ctx_mla_kv_len_cap) so scheduling consumes the pre-reserved cap directly.KvCacheConfig.fp8_context_mla_kv_len_capto override the derived cap (leaving prior worst-case behavior unchanged when unset).accuracy/test_llm_api_pytorch.py::TestKimiK2::test_nvfp4[4gpus](nvbugs/6368562), re-enabling the integration coverage for the scenario impacted by admission-cap enforcement.tests/unittest/_torch/executor/test_mla_workspace_reserve.pywith extensive unit coverage for attended-KV length computation, reuse/chunking effects, cap trimming rules, estimator reserve/cap calculations, and executor behavior when caps are carried on the KV manager vs absent.tests/unittest/_torch/executor/test_dual_pool_kv_cache.pyandtests/unittest/_torch/executor/test_kv_cache_estimation.pyto isolate/neutralize the MLA reserve path where needed.CI follow-up
49973completed successfully for commit9f0571a.QA Engineer Review
Test-list change (test-list only)
tests/integration/test_lists/waives.txtSKIP accuracy/test_llm_api_pytorch.py::TestKimiK2::test_nvfp4[4gpus](nvbugs/6368562).Test-code changes (outside test-list files)
tests/unittest/_torch/executor/test_mla_workspace_reserve.pytest_context_attended_kv_lentest_cap_trims_tail_by_total_kv_lentest_cap_keeps_all_when_within_budgettest_cap_always_keeps_first_even_if_it_alone_exceedstest_cap_single_request_never_trimmedtest_no_cap_returns_untouchedtest_kv_len_cap_default_floor_and_ceilingtest_kv_len_cap_none_when_reuse_cannot_grow_workspacetest_workspace_bytes_zero_for_non_mla_modeltest_workspace_bytes_zero_only_for_absorption_mode_sparse_mlatest_workspace_reserve_captest_workspace_reserve_zero_for_bad_inputstest_ctx_cap_reads_carried_value_ignoring_pool_layouttest_ctx_cap_none_when_no_reservationtest_ctx_cap_no_cap_during_warmuptests/unittest/_torch/executor/test_dual_pool_kv_cache.py_make_creatorto initializecreator._fp8_ctx_mla_kv_len_cap = None.tests/unittest/_torch/executor/test_kv_cache_estimation.pytest_estimation_temporarily_uses_inferred_pool_sizingto patchget_mla_context_workspace_bytes_per_tokento return0.Coverage vs test-db/qa
tests/integration/test_lists/.test_nvfp4[4gpus]).