[None][fix] Declare attention runtime-workspace bytes/token as a backend contract - #16432
Conversation
WalkthroughThe change adds a backend workspace-accounting hook, implements FP8 context-MLA sizing in ChangesAttention workspace accounting
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested labels: Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant KVCacheEstimator
participant TrtllmAttention
participant THOP
PyExecutor->>KVCacheEstimator: estimate KV-cache capacity
KVCacheEstimator->>TrtllmAttention: request workspace bytes per token
TrtllmAttention->>THOP: estimate FP8 context-MLA workspace
THOP-->>TrtllmAttention: return workspace requirement
TrtllmAttention-->>KVCacheEstimator: return workspace bytes per token
KVCacheEstimator-->>PyExecutor: return KV capacity and context cap
``
<!-- walkthrough_end -->
<!-- pre_merge_checks_walkthrough_start -->
<details>
<summary>🚥 Pre-merge checks | ✅ 4 | ❌ 1</summary>
### ❌ Failed checks (1 warning)
| Check name | Status | Explanation | Resolution |
| :----------------: | :--------- | :------------------------------------------------------------------------------------ | :--------------------------------------------------------------------------------- |
| Docstring Coverage | ⚠️ Warning | Docstring coverage is 37.93% which is insufficient. The required threshold is 80.00%. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
<details>
<summary>✅ Passed checks (4 passed)</summary>
| Check name | Status | Explanation |
| :------------------------: | :------- | :------------------------------------------------------------------------------------------------------------------------ |
| Title check | ✅ Passed | The title follows the required format and clearly summarizes the primary backend workspace-contract change. |
| Description check | ✅ Passed | The description explains the motivation, implementation, tests, behavior changes, known follow-ups, and checklist status. |
| 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. |
</details>
</details>
<!-- pre_merge_checks_walkthrough_end -->
<!-- finishing_touch_checkbox_start -->
<details>
<summary>✨ Finishing Touches</summary>
<details>
<summary>🧪 Generate unit tests (beta)</summary>
- [ ] <!-- {"checkboxId": "f47ac10b-58cc-4372-a567-0e02b2c3d479", "radioGroupId": "utg-output-choice-group-unknown_comment_id"} --> Create PR with unit tests
</details>
</details>
<!-- finishing_touch_checkbox_end -->
<!-- tips_start -->
---
<sub>Comment `@coderabbitai help` to get the list of available commands.</sub>
<!-- tips_end -->
|
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
cpp/tensorrt_llm/common/attentionOp.h (1)
61-67: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the new public API with Doxygen.
AttentionOp::contextMlaWorkspaceBytesPerTokenis a new public interface, but its declaration uses ordinary//comments and does not document its parameters or return value. Use a Doxygen comment here.As per coding guidelines, use Doxygen comments for new interfaces.
🤖 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 61 - 67, Replace the ordinary comment immediately preceding AttentionOp::contextMlaWorkspaceBytesPerToken with a Doxygen comment that documents the method’s purpose, every parameter, and its returned byte count. Preserve the existing sizing behavior and shared-source-of-truth description.Source: Coding guidelines
tests/unittest/_torch/executor/test_mla_workspace_reserve.py (1)
67-71: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCoverage gap: cap-derivation and KV-budget-split logic are untested.
Coverage of
_context_attended_kv_lenand_cap_context_by_total_kv_lenis solid, but two related pieces of new behavior have no test in this file:
tensorrt_llm/_torch/pyexecutor/py_executor.py::PyExecutor._get_ctx_mla_kv_len_cap— the_make_executorhelper here bypasses it entirely by pre-setting_ctx_mla_kv_len_cap, so the actualblocks_in_primary_pool * tokens_per_blockcomputation and thew > 0gating are never exercised.tensorrt_llm/_torch/pyexecutor/_util.py::KvCacheCreator.configure_kv_cache_capacity— the newcap = budget / (k + w)reservation-split branch has no unit coverage.Both would need mocking (
kv_cache_managerattributes /get_attention_workspace_bytes_per_token) similar to the pattern already used for_make_executor, so this is a reasonable, low-effort follow-up rather than a blocker — flagging for completeness. As per path instructions, calling out coverage gaps with concrete file names for QA follow-up.Also applies to: 106-112
🤖 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 67 - 71, Add focused tests for the uncovered cap-derivation and KV-budget-split branches. Extend the `_make_executor`-related tests to exercise `PyExecutor._get_ctx_mla_kv_len_cap`, including `blocks_in_primary_pool * tokens_per_block` and the `w > 0` gating, using mocked `kv_cache_manager` attributes; add `KvCacheCreator.configure_kv_cache_capacity` coverage for the `cap = budget / (k + w)` reservation split with a mocked `get_attention_workspace_bytes_per_token`.Source: Path instructions
tensorrt_llm/_torch/attention_backend/trtllm.py (1)
1296-1299: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winReuse the shared MLA predicate here.
tensorrt_llm._torch.pyexecutor.config_utils.is_mla()already checks bothkv_lora_rankandqk_rope_head_dim; this guard only checkskv_lora_rank, so a malformed config can drift past the check and hit the laterconfig.qk_rope_head_dimaccess. Importing the shared helper here is safe.🤖 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/attention_backend/trtllm.py` around lines 1296 - 1299, Update the MLA guard in the relevant attention backend method to use the shared config_utils.is_mla() predicate instead of checking kv_lora_rank directly, and import that helper. Preserve the existing early return of 0 for non-MLA configurations while ensuring both required MLA fields are validated before later qk_rope_head_dim access.
🤖 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: Reset the cached _ctx_mla_kv_len_cap in
_maybe_rebalance_kv_pools immediately after a successful mgr.impl.adjust() so
the next _get_ctx_mla_kv_len_cap() recomputes blocks_in_primary_pool *
tokens_per_block using the rebalanced KV pool. Do not invalidate the cache when
adjust fails or is not performed.
---
Nitpick comments:
In `@cpp/tensorrt_llm/common/attentionOp.h`:
- Around line 61-67: Replace the ordinary comment immediately preceding
AttentionOp::contextMlaWorkspaceBytesPerToken with a Doxygen comment that
documents the method’s purpose, every parameter, and its returned byte count.
Preserve the existing sizing behavior and shared-source-of-truth description.
In `@tensorrt_llm/_torch/attention_backend/trtllm.py`:
- Around line 1296-1299: Update the MLA guard in the relevant attention backend
method to use the shared config_utils.is_mla() predicate instead of checking
kv_lora_rank directly, and import that helper. Preserve the existing early
return of 0 for non-MLA configurations while ensuring both required MLA fields
are validated before later qk_rope_head_dim access.
In `@tests/unittest/_torch/executor/test_mla_workspace_reserve.py`:
- Around line 67-71: Add focused tests for the uncovered cap-derivation and
KV-budget-split branches. Extend the `_make_executor`-related tests to exercise
`PyExecutor._get_ctx_mla_kv_len_cap`, including `blocks_in_primary_pool *
tokens_per_block` and the `w > 0` gating, using mocked `kv_cache_manager`
attributes; add `KvCacheCreator.configure_kv_cache_capacity` coverage for the
`cap = budget / (k + w)` reservation split with a mocked
`get_attention_workspace_bytes_per_token`.
🪄 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: 4157f3a6-9cad-4fa5-892c-67e8a22b83bc
📒 Files selected for processing (10)
cpp/tensorrt_llm/common/attentionOp.cppcpp/tensorrt_llm/common/attentionOp.hcpp/tensorrt_llm/nanobind/thop/bindings.cpptensorrt_llm/_torch/attention_backend/interface.pytensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.mdtensorrt_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
e18a63b to
5b0ff6b
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #63982 [ run ] triggered by Bot. Commit: |
|
PR_Github #63987 [ run ] triggered by Bot. Commit: |
|
PR_Github/16432-48bab81 #63982 was force-killed by a newer pipeline run. |
|
PR_Github #63987 [ run ] completed with state
|
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
The contract belongs on the backend rather than on a model-config predicate — this lands cleanly, and I confirmed the estimator's resolution matches runtime: mla.py builds attention through create_attention(config.attn_backend, ...) (which asserts support_mla()), and the sparse TRTLLM variants all subclass TrtllmAttention, so dropping sparse_params at resolution time can't change the answer. The gating logic moved verbatim; no drift.
Two things worth resolving before merge, both noted inline: the degenerate cap == 0 case now silently serializes context scheduling, and the new interface hook is the only unannotated method among its neighbours.
Title/tracking: this is tagged [None][fix] but it is the completion of a fix for nvbugs/6368562 and lands a real behavior change (non-TRTLLM backends stop being charged the MLA rate). Tag it with that NVBug or a JIRA so the change is traceable from the bug — the description already explains the link, but the title is what tooling reads.
One gap the description doesn't cover: it says the estimator is the contract's only consumer, but the py_executor.py is_warmup and 0-cap changes are consumer-side behavior changes on the scheduler path. They're described further down as "two follow-ups", so this is only a framing nit — the 0-cap one in particular is a user-visible scheduling change and deserves to be called out as such, not filed under review follow-ups.
…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. Addressing review on this PR: - Warn once when the resolved admission cap is degenerate. Since kv_len_cap is always at least max_seq_len (the default is a multiple of it, an override is floored at it), a cap below max_seq_len means the KV budget cannot fund max_seq_len tokens of KV plus workspace -- the pool cannot hold one max-length sequence. Context requests then schedule one per forward step, and the only trace was a debug-level deferral log. Name the memory budget as the lever: the fp8_context_mla_kv_len_cap override is floored at max_seq_len, so it can neither cause nor fix this state. - Annotate the new backend hook's parameters (ModelConfig, Mapping) on both the base declaration and the TrtllmAttention override, matching the annotation convention of the surrounding files so third-party implementers know what they are handed. Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com>
48bab81 to
0f5c162
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tensorrt_llm/_torch/attention_backend/trtllm.py (1)
1371-1423: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse a Google-style contract docstring.
runtime_workspace_bytes_per_tokenis a publicAttentionBackendhook. Documentmodel_config,mapping, and the zero-return cases inArgsandReturnssections. This makes the backend contract clear for future implementations.🤖 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/attention_backend/trtllm.py` around lines 1371 - 1423, Update the docstring of AttentionBackend.runtime_workspace_bytes_per_token to use Google-style Args and Returns sections. Document the model_config and mapping parameters, and specify in Returns when the method returns zero, including non-MLA, non-FP8-KV, and absorption-mode sparse MLA cases.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.
Inline comments:
In `@tests/unittest/_torch/executor/test_kv_cache_estimation.py`:
- Around line 647-653: Complete test-list coverage for the modified and added
tests: run the relevant pytest tests/unittest/ suite, then update the applicable
CI and QA manifests so
tests/unittest/_torch/executor/test_kv_cache_estimation.py and
tests/unittest/_torch/executor/test_mla_workspace_reserve.py are included. Apply
this coverage-list update at the anchor site
tests/unittest/_torch/executor/test_kv_cache_estimation.py:647-653 and sibling
sites tests/unittest/_torch/executor/test_mla_workspace_reserve.py:175-193,
:237-243, and :332-378; each test file must be represented in the relevant
test-db and qa lists.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/trtllm.py`:
- Around line 1371-1423: Update the docstring of
AttentionBackend.runtime_workspace_bytes_per_token to use Google-style Args and
Returns sections. Document the model_config and mapping parameters, and specify
in Returns when the method returns zero, including non-MLA, non-FP8-KV, and
absorption-mode sparse MLA cases.
🪄 Autofix
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: a2a5073c-3026-45be-bbdb-8d70f3ffb629
📒 Files selected for processing (7)
tensorrt_llm/_torch/attention_backend/interface.pytensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.mdtensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/py_executor.pytests/unittest/_torch/executor/test_kv_cache_estimation.pytests/unittest/_torch/executor/test_mla_workspace_reserve.py
🚧 Files skipped from review as they are similar to previous changes (4)
- tensorrt_llm/_torch/attention_backend/interface.py
- tensorrt_llm/_torch/pyexecutor/py_executor.py
- tensorrt_llm/_torch/pyexecutor/_util.py
- tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md
|
PR_Github #65000 [ run ] triggered by Bot. Commit: |
|
PR_Github #65000 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #65224 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #65510 [ run ] triggered by Bot. Commit: |
|
PR_Github #65224 [ run ] completed with state |
|
PR_Github #65510 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #65821 [ run ] triggered by Bot. Commit: |
|
PR_Github #65821 [ run ] completed with state |
…end contract (NVIDIA#16432) Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com>
Description
Follow-up to #16399 (fp8 context-MLA attention-workspace reservation, nvbugs/6368562), now rebased
onto
mainwith #16399 merged.Motivation. #16399 reserves KV-cache headroom for the fp8 context-MLA attention workspace and caps
summed attended KV length in the scheduler. The accounting works, but it is keyed on a model-config
check (is this MLA + fp8 KV?) rather than on the backend that actually allocates the buffer. Two
reviewers flagged this on #16399:
If only to TRTLLM MLA, it would be necessary to check the attention backend."
is disabled, and the selected backend can reach the dense TRTLLM full-gather path."
#16399 landed the first two conditions of QiJune's list; this PR lands the third. It was deferred by
agreement, not dropped.
The underlying failure mode is structural, not MLA-specific: the KV-cache estimator profiles peak memory
against an empty cache and hands the rest to the KV pool, so any attention backend that stages a
workspace sized by a runtime quantity the profiling forward does not drive to its serving maximum (here
total_kv_len, decoupled frommax_num_tokensby KV-cache reuse) is under-reserved and can OOM. Afuture backend would have to rediscover and re-thread the same estimator + cost-rate pieces, and
silently re-introduce the same OOM if it missed one.
Change. Lift the accounting into a declared contract on the attention backend, so future backends
inherit the accounting instead of the OOM:
AttentionBackend.runtime_workspace_bytes_per_token(model_config, mapping) -> int(default0) — abackend declares the per-token cost of any workspace it stages whose size scales with such a runtime
quantity.
TrtllmAttentiondeclares the fp8 context-MLA workspace, still sized by the single C++ source oftruth (
AttentionOp::contextMlaWorkspaceBytesPerToken, via nanobind), and keeping [https://nvbugs/6368562][fix] Reserve fp8 context-MLA attention workspace in KV cache estimation #16399'sruntime-matched sparse gate verbatim (dsa/deepseek_v4 on SM 100/103 with the short-seq MHA fallback
off — not "a sparse config exists").
get_attention_workspace_bytes_per_token(). The reserve/cap math is unchanged.ATTENTION_DEVELOPER_GUIDE.md(required reading) — §2.3, §3.2.3, §4.2.This is not a pure refactor. A model that resolves to a non-TRTLLM backend now correctly reserves
nothing, where
maincharges it the MLA rate off a model-config check and shrinks the KV pool for abuffer that backend never allocates. That is the behavior change the two reviews asked for. The active
fp8-MLA-on-TRTLLM path — the nvbugs/6368562 repro — is unchanged.
Scope of the contract. Deliberately a scalar per-token rate, not a typed driver/reservation
abstraction. Only the rate is generalized; the driving quantity (
total_kv_len), the reservation gate(
get_mla_context_workspace_kv_len_cap) and the cap plumbing (kv_cache_manager.fp8_ctx_mla_kv_len_cap,KvCacheConfig.fp8_context_mla_kv_len_cap) remain MLA-named, because there is one driver today and thescheduler's cap is specific to it. A backend with a different driving quantity introduces it then,
alongside the enforcement it needs — a richer type now would be unused scaffolding.
Note that after #16399's review redesign (carrying the admission cap from the estimator onto the KV
manager instead of re-deriving it from pool layout), the scheduler no longer reads the per-token rate at
all. The contract therefore has exactly one consumer: the estimator.
py_executor.py's only delta hereis the two review follow-ups below.
Additionally: two follow-ups from #16399 review threads
Both threads were marked resolved on #16399 without a code change landing. Both are in
PyExecutor._get_ctx_mla_kv_len_cap, the cap reader this contract feeds, so they are folded in hererather than left dangling:
0was collapsed toNone("no cap") by a truthiness check(
int(carried) if carried else None), inverting admission control for precisely the tightest-budgetcase the reservation exists to protect. Now compares against
None. Flagged by CodeRabbit; new testtest_ctx_cap_zero_is_a_cap_not_no_capasserts both the read and that the trim still enforces it(keeping the first request as the forward-progress guard).
getattr(self, "is_warmup", False)→self.is_warmup— it is a real property onPyExecutor(
py_executor.py:1212), so the defensivegetattris unnecessary. Flagged (non-blocking) by@pengbowang-nv. The remaining
getattronkv_cache_manageris kept deliberately and now carries acomment saying why: managers not built by the estimator never carry the attribute.
Known follow-ups (not in this PR)
full_k/full_kvfull-gather buffers on the reuse path are still unaccounted for(@QiJune's [https://nvbugs/6368562][fix] Reserve fp8 context-MLA attention workspace in KV cache estimation #16399 thread, deferred by agreement), along with the high-fanout shared-prefix memory test
he asked to be tracked.
try_prepare_estimationdisables estimation for context parallelism and theVANILLAbackend withoutsetting
_skip_est, soconfigure_kv_cache_capacitynever runs and no cap is installed — fail-openrather than fail-closed. Narrow (neither path reaches fp8 context-MLA today) but worth closing.
Test Coverage
tests/unittest/_torch/executor/test_mla_workspace_reserve.py— retargeted to the new resolver. Thenon-MLA test now exercises the full resolve-backend path (
get_attention_backend→ backend classmethod→
0); a newtest_workspace_bytes_zero_for_backend_without_declarationcovers a backend thatinherits the default
0for a model the TRTLLM backend would charge for (usesVANILLA, whichalways resolves —
FLASHINFERsilently falls back to TRTLLM when flashinfer is absent). Plustest_ctx_cap_zero_is_a_cap_not_no_capfor the carried-zero fix. All of [https://nvbugs/6368562][fix] Reserve fp8 context-MLA attention workspace in KV cache estimation #16399's existing coverage ispreserved.
tests/unittest/_torch/executor/test_kv_cache_estimation.py— patch target renamed to the resolver.PR Checklist
pre-commitrun locally against the PR diff range (green).GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Dev Engineer Review
AttentionBackend.runtime_workspace_bytes_per_token(model_config, mapping).0.0.ATTENTION_DEVELOPER_GUIDE.md.QA Engineer Review
tests/unittest/_torch/executor/test_mla_workspace_reserve.py.tests/unittest/_torch/executor/test_kv_cache_estimation.py.get_attention_workspace_bytes_per_token.tests/integration/test_lists/entries were modified.