[None][chore] Declare attention runtime-workspace bytes/token as a backend contract - #16432
Draft
eopXD wants to merge 2 commits into
Draft
[None][chore] Declare attention runtime-workspace bytes/token as a backend contract#16432eopXD wants to merge 2 commits into
eopXD wants to merge 2 commits into
Conversation
…pace in KV cache estimation The fp8 context-MLA K/V dequant workspace scales with the summed attended KV length (total_kv_len) of the context requests in a forward step. The KV-cache memory estimator profiles with fresh-prefill dummy requests against an empty cache, so total_kv_len there is pinned near max_num_tokens and this workspace sits at its floor. With block reuse at serving time total_kv_len decouples from max_num_tokens and the workspace can grow far past the profiled floor, but the estimator has already handed that headroom to the KV pool, causing an OOM mid-forward (TestKimiK2::test_nvfp4[4gpus], exposed once NVIDIA#14852 sized the workspace correctly by total_kv_len). Fix the estimation rather than the memory fraction (a user co-tenancy knob): - Split the KV budget between the pool (k bytes/token, all layers) and the workspace (w bytes/token, one shared layer buffer) at a common token count, so max_tokens = budget / (k + w) and the reserved workspace covers exactly max_tokens tokens of attended KV. - Enforce it at admission: the scheduler trims scheduled context requests so their summed attended KV length stays within the pool token capacity, always keeping at least one request as a forward-progress guard. No-op for non-fp8-MLA models. The per-token workspace cost is a single source of truth in C++ (AttentionOp::contextMlaWorkspaceBytesPerToken), exposed via nanobind, so the estimator's reserve can never drift from the runtime allocation. Remove the waive for TestKimiK2::test_nvfp4[4gpus] (nvbugs/6368562) to re-enable the test. Co-Authored-By: Yueh-Ting Chen <yueh.ting.chen@gmail.com> Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com>
…ckend contract The fp8 context-MLA workspace reservation (nvbugs/6368562) was threaded imperatively through the KV-cache estimator and the scheduler, keyed on a model-config check specific to MLA. Any future attention backend that stages a workspace sized by a runtime quantity the KV-cache profiler under-measures would have to rediscover and re-thread the same estimator + scheduler + cost-rate pieces, and silently re-introduce the same mid-forward OOM if it missed one. 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 except fp8 context-MLA. - TrtllmAttention declares the fp8 context-MLA K/V dequant workspace, still sized by the single C++ source of truth (contextMlaWorkspaceBytesPerToken). - The estimator and scheduler consume the declaration generically via get_attention_workspace_bytes_per_token(); their reservation/cap logic is unchanged. The active fp8-MLA path is behavior-identical; the reservation is now tied to the backend that actually allocates the buffer. - 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 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. Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com>
eopXD
requested review from
BowenFu,
PerkzZheng,
cascade812,
mzweilz,
niukuo,
nv-xtf,
reasonsolo,
ruodil,
schetlur-nv,
tongyuantongyu and
yunruis
July 15, 2026 14:18
Contributor
📝 WalkthroughWalkthroughAdds a shared FP8 context-MLA workspace estimator, exposes it to Python backends, reserves corresponding KV-cache headroom, and caps scheduled context requests by attended KV length. Tests cover reuse accounting, trimming behavior, disabled caps, and non-MLA configurations. ChangesFP8 context-MLA workspace accounting
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant AttentionBackend
participant AttentionOp
participant KvCacheCreator
participant PyExecutor
AttentionBackend->>AttentionOp: obtain FP8 context-MLA bytes per token
AttentionOp-->>AttentionBackend: return K/V staging cost
KvCacheCreator->>AttentionBackend: resolve workspace reservation
AttentionBackend-->>KvCacheCreator: return bytes per token
KvCacheCreator->>PyExecutor: establish KV capacity and attended-KV cap
PyExecutor-->>PyExecutor: trim context requests exceeding the cap
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}
Contributor
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
Comment on lines
+5108
to
+5133
| def _get_ctx_mla_kv_len_cap(self): | ||
| """Cap on the summed context attended-KV length (total_kv_len) per forward step. | ||
|
|
||
| Equals the KV pool's primary-pool token capacity, which is exactly what the estimator reserved the | ||
| fp8 context-MLA attention workspace for (`max_tokens = budget / (k + w)`). Returns None (no cap) for | ||
| non-fp8-MLA models — those reserved no workspace, so this per-forward constraint must not alter their | ||
| scheduling. Computed once and cached. | ||
| """ | ||
| cap = getattr(self, "_ctx_mla_kv_len_cap", "unset") | ||
| if cap != "unset": | ||
| return cap | ||
| # Lazy import: _util imports py_executor at module scope, so a top-level import here is circular. | ||
| from ._util import get_attention_workspace_bytes_per_token | ||
| cap = None | ||
| w = get_attention_workspace_bytes_per_token( | ||
| self.model_engine.model.model_config, self.dist.mapping) | ||
| if w > 0: | ||
| blocks = getattr(self.kv_cache_manager, "blocks_in_primary_pool", | ||
| None) | ||
| tokens_per_block = getattr(self.kv_cache_manager, | ||
| "tokens_per_block", None) | ||
| if blocks and tokens_per_block: | ||
| cap = int(blocks) * int(tokens_per_block) | ||
| self._ctx_mla_kv_len_cap = cap | ||
| return cap | ||
|
|
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the relevant methods and nearby context.
python3 - <<'PY'
from pathlib import Path
path = Path("tensorrt_llm/_torch/pyexecutor/py_executor.py")
lines = path.read_text().splitlines()
targets = ["def _get_ctx_mla_kv_len_cap", "def _maybe_rebalance_kv_pools", "def _context_attended_kv_len", "def _cap_context_by_total_kv_len", "def _schedule"]
for t in targets:
for i, line in enumerate(lines, 1):
if t in line:
start = max(1, i - 20)
end = min(len(lines), i + 80)
print(f"\n=== {t} @ line {i} ===")
for j in range(start, end + 1):
print(f"{j:5d}: {lines[j-1]}")
break
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 30267
Invalidate the cached context-MLA KV cap after KV pool rebalance. _get_ctx_mla_kv_len_cap() caches blocks_in_primary_pool * tokens_per_block once, but _maybe_rebalance_kv_pools() can change the primary pool via mgr.impl.adjust(). Reset _ctx_mla_kv_len_cap after a successful adjust so scheduling tracks the current capacity.
🤖 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 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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Follow-up to #16399 (fp8 context-MLA attention-workspace reservation, nvbugs/6368562).
Motivation. #16399 fixes a mid-forward OOM by reserving KV-cache headroom for the fp8
context-MLA attention workspace and capping summed attended KV length in the scheduler. That fix is
correct, but the accounting is threaded imperatively through the estimator
(
_util.configure_kv_cache_capacity) and the scheduler (py_executor._cap_context_by_total_kv_len),keyed on an MLA-specific model-config check. The failure mode it patches 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. A future backend would have torediscover and re-thread the same estimator + scheduler + 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).get_attention_workspace_bytes_per_token(); their reservation/cap logic is unchanged. The activefp8-MLA path is behavior-identical; the reservation is now tied to the backend that actually
allocates the buffer.
ATTENTION_DEVELOPER_GUIDE.md(required reading) — §2.3, §3.2.3, §4.2.The contract is deliberately a scalar per-token rate, not a typed driver/reservation abstraction:
there is one driving quantity today and the scheduler's cap is specific to it, so a richer type would
be unused scaffolding. A backend with a different driving quantity introduces it then, alongside the
enforcement it needs.
Stacking note
Stacked on #16399 (fork branch
nvbugs/6368562-mla-workspace-reserve). Until that merges, the diffhere also includes #16399's changes; will rebase onto
mainonce #16399 lands. Please merge #16399first.
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→ backendclassmethod →
0). Behavior-preserving refactor; no new runtime path.PR Checklist
pre-commitrun locally (isort/yapf/ruff/ruff-format green).GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Tests