Skip to content

[None][fix] Declare attention runtime-workspace bytes/token as a backend contract - #16432

Merged
eopXD merged 1 commit into
NVIDIA:mainfrom
eopXD:attention-workspace-reservation-contract
Aug 17, 2026
Merged

[None][fix] Declare attention runtime-workspace bytes/token as a backend contract#16432
eopXD merged 1 commit into
NVIDIA:mainfrom
eopXD:attention-workspace-reservation-contract

Conversation

@eopXD

@eopXD eopXD commented Jul 15, 2026

Copy link
Copy Markdown
Collaborator

Description

Follow-up to #16399 (fp8 context-MLA attention-workspace reservation, nvbugs/6368562), now rebased
onto main with #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:

  • @SimengLiu-nv: "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."
  • @QiJune: "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."

#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 from max_num_tokens by KV-cache reuse) is under-reserved and can OOM. A
future 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 (default 0) — a
    backend declares the per-token cost of any workspace it stages whose size scales with such a runtime
    quantity.
  • TrtllmAttention declares the fp8 context-MLA workspace, still sized by the single C++ source of
    truth (AttentionOp::contextMlaWorkspaceBytesPerToken, via nanobind), and keeping [https://nvbugs/6368562][fix] Reserve fp8 context-MLA attention workspace in KV cache estimation #16399's
    runtime-matched sparse gate verbatim (dsa/deepseek_v4 on SM 100/103 with the short-seq MHA fallback
    off — not "a sparse config exists").
  • The estimator resolves the declaration through the model's selected backend via
    get_attention_workspace_bytes_per_token(). The reserve/cap math is unchanged.
  • Documented as a contract in 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 main charges it the MLA rate off a model-config check and shrinks the KV pool for a
buffer 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 the
scheduler'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 here
is 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 here
rather than left dangling:

  • A carried cap of exactly 0 was collapsed to None ("no cap") by a truthiness check
    (int(carried) if carried else None), inverting admission control for precisely the tightest-budget
    case the reservation exists to protect. Now compares against None. Flagged by CodeRabbit; new test
    test_ctx_cap_zero_is_a_cap_not_no_cap asserts 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 on PyExecutor
    (py_executor.py:1212), so the defensive getattr is unnecessary. Flagged (non-blocking) by
    @pengbowang-nv. The remaining getattr on kv_cache_manager is kept deliberately and now carries a
    comment saying why: managers not built by the estimator never carry the attribute.

Known follow-ups (not in this PR)

  • The BF16 full_k / full_kv full-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_estimation disables estimation for context parallelism and the VANILLA backend without
    setting _skip_est, so configure_kv_cache_capacity never runs and no cap is installed — fail-open
    rather 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. The
    non-MLA test now exercises the full resolve-backend path (get_attention_backend → backend classmethod
    0); a new test_workspace_bytes_zero_for_backend_without_declaration covers a backend that
    inherits the default 0 for a model the TRTLLM backend would charge for (uses VANILLA, which
    always resolves — FLASHINFER silently falls back to TRTLLM when flashinfer is absent). Plus
    test_ctx_cap_zero_is_a_cap_not_no_cap for the carried-zero fix. All of [https://nvbugs/6368562][fix] Reserve fp8 context-MLA attention workspace in KV cache estimation #16399's existing coverage is
    preserved.
  • tests/unittest/_torch/executor/test_kv_cache_estimation.py — patch target renamed to the resolver.

PR Checklist

  • PR description clearly explains what and why.
  • PR follows TRT-LLM coding guidelines; pre-commit run locally against the PR diff range (green).
  • Test cases provided for the new code paths.
  • No public API changes (the backend method is internal).
  • No new dependencies.

GitHub Bot Help

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

Dev Engineer Review

  • Added AttentionBackend.runtime_workspace_bytes_per_token(model_config, mapping).
  • Kept the default workspace reservation at 0.
  • Added the TRTLLM implementation for fp8 context-MLA configurations.
  • Routed KV-cache estimation through the selected attention backend.
  • Preserved TRTLLM workspace reservation and scheduler-cap behavior.
  • Preserved an explicit admission cap of 0.
  • Added a warning for zero or degenerate caps.
  • Documented the backend contract in ATTENTION_DEVELOPER_GUIDE.md.
  • No configuration or test-list files changed.
  • The implementation is behavior-preserving for supported TRTLLM fp8 context-MLA paths.
  • Non-TRTLLM MLA backends no longer reserve TRTLLM-specific workspace.
  • Review follow-up should confirm that operators can observe the throughput impact of a zero cap.

QA Engineer Review

  • Modified tests/unittest/_torch/executor/test_mla_workspace_reserve.py.
  • Modified tests/unittest/_torch/executor/test_kv_cache_estimation.py.
  • Added coverage for backend workspace resolution.
  • Added coverage for non-MLA models and backends without workspace declarations.
  • Added coverage for workspace reservation.
  • Added coverage for zero-cap enforcement and one-time warnings.
  • Updated KV-cache estimation mocking to use get_attention_workspace_bytes_per_token.
  • No corresponding tests/integration/test_lists/ entries were modified.
  • Verdict: needs follow-up because CI coverage-list data is unavailable.

@coderabbitai

coderabbitai Bot commented Jul 15, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Walkthrough

The change adds a backend workspace-accounting hook, implements FP8 context-MLA sizing in TrtllmAttention, and uses the generic value for KV-cache estimation. It also preserves zero admission caps and adds warnings and tests for degenerate context limits.

Changes

Attention workspace accounting

Layer / File(s) Summary
Backend workspace contract and FP8 context-MLA estimator
tensorrt_llm/_torch/attention_backend/interface.py, tensorrt_llm/_torch/attention_backend/trtllm.py, tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md
AttentionBackend defines a default-zero runtime_workspace_bytes_per_token hook. TrtllmAttention calculates FP8 context-MLA workspace for supported configurations and delegates sizing to THOP. The guide documents the contract.
Generic KV-cache workspace reservation
tensorrt_llm/_torch/pyexecutor/_util.py, tests/unittest/_torch/executor/test_kv_cache_estimation.py
KV-cache estimation resolves workspace requirements through the selected backend. Tests replace the MLA-specific mock with the generic estimator.
Context admission cap handling
tensorrt_llm/_torch/pyexecutor/py_executor.py, tests/unittest/_torch/executor/test_mla_workspace_reserve.py
Zero remains an active cap. PyExecutor warns once for degenerate caps. Tests cover backend resolution, sparse MLA, zero caps, and warning behavior.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Suggested labels: api-compatible

Suggested reviewers: bowenfu, schetlur-nv, vallis-neria

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 -->
Loading

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (3)
cpp/tensorrt_llm/common/attentionOp.h (1)

61-67: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Document the new public API with Doxygen.

AttentionOp::contextMlaWorkspaceBytesPerToken is 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 win

Coverage gap: cap-derivation and KV-budget-split logic are untested.

Coverage of _context_attended_kv_len and _cap_context_by_total_kv_len is 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_executor helper here bypasses it entirely by pre-setting _ctx_mla_kv_len_cap, so the actual blocks_in_primary_pool * tokens_per_block computation and the w > 0 gating are never exercised.
  • tensorrt_llm/_torch/pyexecutor/_util.py::KvCacheCreator.configure_kv_cache_capacity — the new cap = budget / (k + w) reservation-split branch has no unit coverage.

Both would need mocking (kv_cache_manager attributes / 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 win

Reuse the shared MLA predicate here.
tensorrt_llm._torch.pyexecutor.config_utils.is_mla() already checks both kv_lora_rank and qk_rope_head_dim; this guard only checks kv_lora_rank, so a malformed config can drift past the check and hit the later config.qk_rope_head_dim access. 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

📥 Commits

Reviewing files that changed from the base of the PR and between 0846183 and e18a63b.

📒 Files selected for processing (10)
  • cpp/tensorrt_llm/common/attentionOp.cpp
  • cpp/tensorrt_llm/common/attentionOp.h
  • cpp/tensorrt_llm/nanobind/thop/bindings.cpp
  • tensorrt_llm/_torch/attention_backend/interface.py
  • tensorrt_llm/_torch/attention_backend/trtllm.py
  • tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md
  • tensorrt_llm/_torch/pyexecutor/_util.py
  • tensorrt_llm/_torch/pyexecutor/py_executor.py
  • tests/integration/test_lists/waives.txt
  • tests/unittest/_torch/executor/test_mla_workspace_reserve.py
💤 Files with no reviewable changes (1)
  • tests/integration/test_lists/waives.txt

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py
@eopXD
eopXD marked this pull request as draft July 15, 2026 14:36
@mikeiovine
mikeiovine requested review from mikeiovine and removed request for schetlur-nv July 20, 2026 15:30
@eopXD
eopXD force-pushed the attention-workspace-reservation-contract branch from e18a63b to 5b0ff6b Compare August 3, 2026 06:06
@eopXD eopXD changed the title [None][chore] Declare attention runtime-workspace bytes/token as a backend contract [None][fix] Declare attention runtime-workspace bytes/token as a backend contract Aug 3, 2026
@eopXD
eopXD marked this pull request as ready for review August 3, 2026 06:11
@eopXD
eopXD requested a review from a team as a code owner August 3, 2026 06:11
@eopXD
eopXD requested review from VALLIS-NERIA and lowsfer August 3, 2026 06:11
@eopXD

eopXD commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63982 [ run ] triggered by Bot. Commit: 48bab81 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63987 [ run ] triggered by Bot. Commit: 48bab81 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github/16432-48bab81 #63982 was force-killed by a newer pipeline run.
L0 job information not available (job may not have been triggered yet).

Link to superseding invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63987 [ run ] completed with state SUCCESS. Commit: 48bab81
/LLM/main/L0_MergeRequest_PR pipeline #51922 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ 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

Link to invocation

@pengbowang-nv pengbowang-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@SimengLiu-nv SimengLiu-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approve for KVCM.

@brnguyen2 brnguyen2 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py
Comment thread tensorrt_llm/_torch/attention_backend/interface.py Outdated
Comment thread tensorrt_llm/_torch/attention_backend/trtllm.py
…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>
@eopXD
eopXD force-pushed the attention-workspace-reservation-contract branch from 48bab81 to 0f5c162 Compare August 10, 2026 08:36
@coderabbitai

coderabbitai Bot commented Aug 10, 2026

Copy link
Copy Markdown
Contributor

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.

@eopXD

eopXD commented Aug 10, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (1)
tensorrt_llm/_torch/attention_backend/trtllm.py (1)

1371-1423: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Use a Google-style contract docstring.

runtime_workspace_bytes_per_token is a public AttentionBackend hook. Document model_config, mapping, and the zero-return cases in Args and Returns sections. 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

📥 Commits

Reviewing files that changed from the base of the PR and between f13a0be and 0f5c162.

📒 Files selected for processing (7)
  • tensorrt_llm/_torch/attention_backend/interface.py
  • tensorrt_llm/_torch/attention_backend/trtllm.py
  • tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md
  • tensorrt_llm/_torch/pyexecutor/_util.py
  • tensorrt_llm/_torch/pyexecutor/py_executor.py
  • tests/unittest/_torch/executor/test_kv_cache_estimation.py
  • tests/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

Comment thread tests/unittest/_torch/executor/test_kv_cache_estimation.py
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #65000 [ run ] triggered by Bot. Commit: 0f5c162 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #65000 [ run ] completed with state FAILURE. Commit: 0f5c162
/LLM/main/L0_MergeRequest_PR pipeline #52815 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

@eopXD

eopXD commented Aug 11, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #65224 [ run ] triggered by Bot. Commit: 0f5c162 Link to invocation

@eopXD

eopXD commented Aug 12, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #65510 [ run ] triggered by Bot. Commit: 0f5c162 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #65224 [ run ] completed with state ABORTED. Commit: 0f5c162

Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #65510 [ run ] completed with state FAILURE. Commit: 0f5c162
/LLM/main/L0_MergeRequest_PR pipeline #53249 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

@eopXD

eopXD commented Aug 13, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #65821 [ run ] triggered by Bot. Commit: 0f5c162 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #65821 [ run ] completed with state SUCCESS. Commit: 0f5c162
/LLM/main/L0_MergeRequest_PR pipeline #53521 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

@eopXD
eopXD merged commit 3d311e3 into NVIDIA:main Aug 17, 2026
13 checks passed
yihwang-nv pushed a commit to yihwang-nv/TensorRT-LLM that referenced this pull request Aug 18, 2026
…end contract (NVIDIA#16432)

Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

10 participants