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

tensorrt_llm/_torch/attention_backend/interface.py

@@ -13,6 +13,7 @@
 from typing_extensions import Self
 
 if TYPE_CHECKING:
+    from ..model_config import ModelConfig
     from ..speculative.interface import SpecMetadata
     from ..speculative.spec_tree_manager import SpecTreeManager
 
@@ -1070,6 +1071,21 @@ def support_mla(cls) -> bool:
     def support_multi_item_scoring(cls) -> bool:
         return False
 
+    @classmethod
+    def runtime_workspace_bytes_per_token(cls, model_config: "ModelConfig",
+                                          mapping: Mapping) -> int:
+        """Per-token bytes to reserve for a workspace this backend stages whose size scales with a
+        runtime quantity the KV-cache estimator does not drive to its serving maximum while profiling
+        (e.g. ``total_kv_len``, inflated by KV-cache reuse). Default ``0`` -- correct for every backend
+        except fp8 context-MLA today.
+
+        A non-zero rate is reserved from the KV budget by the estimator, and the scheduler caps the
+        driving sum at what that reserve covers, so the declared buffer stays within its reservation
+        (buffers the backend does not declare are still unaccounted for). Keep the rate identical to the
+        runtime allocation's per-token cost. See ``ATTENTION_DEVELOPER_GUIDE.md`` §2.3.
+        """
+        return 0
+
     def create_output(self, q: torch.Tensor, **kwargs) -> List[torch.Tensor]:
         """
         Create the output tensors for the attention operation.

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -23,6 +23,9 @@
 import torch
 
 if TYPE_CHECKING:
+    from tensorrt_llm.mapping import Mapping
+
+    from ..model_config import ModelConfig
     from ..speculative.interface import SpecMetadata
     from ..speculative.spec_tree_manager import SpecTreeManager
 
@@ -34,6 +37,7 @@
 from tensorrt_llm.math_utils import ceil_div
 from tensorrt_llm.models.modeling_utils import QuantConfig
 
+from ..pyexecutor.config_utils import is_mla
 from ..utils import (compute_swizzled_sf_shape, get_global_attrs,
                      get_model_extra_attrs)
 from .interface import (AttentionBackend, AttentionForwardArgs,
@@ -1364,6 +1368,60 @@ def update_quant_config(self, new_quant_config: Optional[QuantConfig]):
             )
         self.create_fmha_libs()
 
+    @classmethod
+    def runtime_workspace_bytes_per_token(cls, model_config: "ModelConfig",
+                                          mapping: "Mapping") -> int:
+        """fp8 context-MLA stages a K/V dequant workspace sized by the summed attended KV length
+        (``total_kv_len``) across the context requests in a forward step, not by ``max_num_tokens`` -- so
+        KV-cache reuse can push it far past the profiling floor. This buffer is shared across attention
+        layers. ``0`` for non-MLA / non-fp8-KV / absorption-mode sparse MLA (which reads K/V straight from
+        the paged cache and stages no dequant buffer).
+
+        The per-token cost is the single source of truth in C++
+        (``AttentionOp::contextMlaWorkspaceBytesPerToken``, exposed via nanobind), so it cannot drift
+        from the runtime allocation.
+        """
+        config = model_config.pretrained_config
+        if not is_mla(config):
+            return 0
+        quant_config = model_config.quant_config
+        fp8_context_mla = (quant_config is not None
+                           and quant_config.quant_mode.has_fp8_kv_cache()
+                           and get_sm_version() in (90, 100, 103, 120))
+        if not fp8_context_mla:
+            return 0
+        # Attention-DP runs the full head set per rank; otherwise heads shard across TP (mirror
+        # mNumAttnHeads).
+        attn_tp = 1 if mapping.enable_attention_dp else mapping.tp_size
+        num_attn_heads = config.num_attention_heads // attn_tp
+        # The buffer is skipped only where AttentionOp::useSparseMLA() holds, which needs all three of:
+        #   * DSA / DeepSeek-V4 -- only these lower to the absorption path that reads K/V from the paged
+        #     cache. Skip-softmax passes no sparse indices to C++, and its ignore-list can exclude a layer,
+        #     so those layers still run dense MLA. The workspace is shared, so one dense layer forces the
+        #     reserve.
+        #   * SM 100 / 103 -- mUseTllmGen is `sm >= 100 && sm != 120`.
+        #   * short-seq MHA fallback off -- it routes short contexts back through the dense path.
+        # Match the runtime predicate, not just "a sparse config exists": over-reserving costs KV pool,
+        # under-reserving OOMs mid-forward.
+        sparse_algorithm = getattr(model_config.sparse_attention_config,
+                                   "algorithm", None)
+        sparse_mla = (sparse_algorithm in ("dsa", "deepseek_v4")
+                      and get_sm_version() in (100, 103))
+        short_seq_mha_enabled = int(
+            os.environ.get("TRTLLM_MLA_SHORT_SEQ_MHA_THRESHOLD", "0")) > 0
+        stages_no_buffer = sparse_mla and not short_seq_mha_enabled
+        return int(
+            thop.get_context_mla_workspace_bytes_per_token(
+                num_attn_heads=num_attn_heads,
+                qk_rope_head_dim=config.qk_rope_head_dim,
+                qk_nope_head_dim=config.qk_nope_head_dim,
+                v_head_dim=config.v_head_dim,
+                fp8_context_mla=fp8_context_mla,
+                # Paged context MLA always uses separate Q/KV input; the term is otherwise gated to 0.
+                separate_q_and_kv_input=True,
+                sparse_mla=stages_no_buffer,
+            ))
+
     def get_local_layer_idx(self, metadata: TrtllmAttentionMetadata) -> int:
         if self.local_layer_idx is not None:
             return self.local_layer_idx

tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md

@@ -204,6 +204,8 @@ The core contract is:
   - `support_fused_rope()`
   - `support_fused_qkv()`
   - `support_mla()`
+- `runtime_workspace_bytes_per_token(model_config, mapping)` — the memory-accounting
+  contract (default `0`); see below
 
 `**kwargs` is only a temporary compatibility path. It is merged into
 `AttentionForwardArgs`, rejects unknown fields, and must not be mixed with
@@ -212,6 +214,18 @@ an explicit `forward_args`.
 Those capability hooks are coarse checks. They do not prove that every
 required operator or sparse path already exists.
 
+**Workspace memory-accounting contract.** The KV-cache estimator profiles peak
+memory against an empty cache and hands the rest to the KV pool, so a workspace
+sized by a runtime quantity the profiling forward never drives to its serving
+maximum is under-reserved and can OOM mid-forward. If a backend stages such a
+buffer, declare its per-token cost via
+`runtime_workspace_bytes_per_token(model_config, mapping)` (default `0`): the
+estimator reserves it from the KV budget and the scheduler caps the driving sum.
+Keep the declared cost identical to the runtime allocation's (single source of
+truth). The one instance today is the fp8 context-MLA K/V dequant workspace,
+sized by summed attended KV length (`total_kv_len`) — which KV-cache reuse
+decouples from `max_num_tokens` (`TrtllmAttention.runtime_workspace_bytes_per_token`).
+
 ### 2.4 Capability reference
 
 Check each backend's capability hooks (`support_fused_rope()`,
@@ -323,6 +337,9 @@ agree on latent-cache layout, paged-KV read/write paths, and cached/chunked
 context behavior. Read `mla.py` and the relevant
 backend code for the current implementation details.
 
+fp8 context-MLA also stages a K/V dequant workspace sized by summed attended KV
+length; it is declared through the workspace memory-accounting contract (§2.3).
+
 #### 3.2.4 Sparse side-cache semantics
 
 Sparse backends may add side caches beyond the main KV cache. Some sparse
@@ -376,6 +393,11 @@ as the current blocker.
   reused, whether chunked prefill or speculative decoding matters, and whether
   sparse side caches are required.
 
+- **Workspace memory accounting**
+  Whether the backend stages a workspace sized by a runtime quantity the KV-cache
+  profiler does not max out (e.g. `total_kv_len` under reuse) — if so, declare it
+  via `runtime_workspace_bytes_per_token` (§2.3).
+
 ### 4.3 Default bring-up order
 
 Start with `TRTLLM` when the new attention fits or only needs limited changes.

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -321,56 +321,24 @@ def bytes_for_tokens(self, tokens: int) -> int:
         return self.slope * tokens + self.intercept
 
 
-def get_mla_context_workspace_bytes_per_token(model_config, mapping) -> int:
-    """Per-token byte cost of the fp8 context-MLA K/V dequant workspace.
-
-    This buffer is shared across attention layers and scales with the summed attended KV length
-    (``total_kv_len``) of the step's context requests, which KV-cache reuse can grow far past the floor the
-    profiling forward measures against an empty cache. The per-token size is the single source of truth in
-    C++ (``AttentionOp::contextMlaWorkspaceBytesPerToken``, exposed via nanobind), so the estimator's
-    reserve cannot drift from the runtime allocation. Returns 0 for non-MLA / non-fp8-KV / absorption-mode
-    sparse MLA (which reads K/V straight from the paged cache). A non-zero result drives both the reserve
-    and the scheduler's admission cap.
+def get_attention_workspace_bytes_per_token(model_config, mapping) -> int:
+    """Per-token workspace headroom the model's selected attention backend declares.
+
+    The KV-cache profiling forward under-measures any attention workspace sized by a runtime quantity it
+    does not drive to its serving maximum (e.g. ``total_kv_len``, inflated by KV reuse). Backends declare
+    such a buffer via ``AttentionBackend.runtime_workspace_bytes_per_token``; this resolves the model's
+    backend and returns its rate. A backend that stages no such buffer inherits the default 0, so no
+    workspace is reserved and no admission cap is installed for it. See ``ATTENTION_DEVELOPER_GUIDE.md``
+    §2.3.
     """
-    from tensorrt_llm.bindings.internal import thop
-    config = model_config.pretrained_config
-    if not is_mla(config):
-        return 0
-    quant_config = model_config.quant_config
-    fp8_context_mla = (quant_config is not None
-                       and quant_config.quant_mode.has_fp8_kv_cache()
-                       and get_sm_version() in (90, 100, 103, 120))
-    if not fp8_context_mla:
-        return 0
-    # Attention-DP runs the full head set per rank; otherwise heads shard across TP (mirror mNumAttnHeads).
-    attn_tp = 1 if mapping.enable_attention_dp else mapping.tp_size
-    num_attn_heads = config.num_attention_heads // attn_tp
-    # The buffer is skipped only where AttentionOp::useSparseMLA() holds, which needs all three of:
-    #   * DSA / DeepSeek-V4 -- only these lower to the absorption path that reads K/V from the paged cache.
-    #     Skip-softmax passes no sparse indices to C++, and its ignore-list can exclude a layer, so those
-    #     layers still run dense MLA. The workspace is shared, so one dense layer forces the reserve.
-    #   * SM 100 / 103 -- mUseTllmGen is `sm >= 100 && sm != 120`.
-    #   * short-seq MHA fallback off -- it routes short contexts back through the dense path.
-    # Match the runtime predicate, not just "a sparse config exists": over-reserving costs KV pool,
-    # under-reserving OOMs mid-forward.
-    sparse_algorithm = getattr(model_config.sparse_attention_config,
-                               "algorithm", None)
-    sparse_mla = (sparse_algorithm in ("dsa", "deepseek_v4")
-                  and get_sm_version() in (100, 103))
-    short_seq_mha_enabled = int(
-        os.environ.get("TRTLLM_MLA_SHORT_SEQ_MHA_THRESHOLD", "0")) > 0
-    stages_no_buffer = sparse_mla and not short_seq_mha_enabled
-    return int(
-        thop.get_context_mla_workspace_bytes_per_token(
-            num_attn_heads=num_attn_heads,
-            qk_rope_head_dim=config.qk_rope_head_dim,
-            qk_nope_head_dim=config.qk_nope_head_dim,
-            v_head_dim=config.v_head_dim,
-            fp8_context_mla=fp8_context_mla,
-            # Paged context MLA always uses separate Q/KV input; the term is otherwise gated to 0.
-            separate_q_and_kv_input=True,
-            sparse_mla=stages_no_buffer,
-        ))
+    from ..attention_backend.utils import get_attention_backend
+
+    # Resolved without ``sparse_params``: those are per-layer, while this workspace is one buffer shared
+    # across every attention layer, so the declaration is a whole-model question. The sparse backends all
+    # derive from the dense class and inherit its declaration, which reads the sparse gate off model_config.
+    return get_attention_backend(
+        model_config.attn_backend).runtime_workspace_bytes_per_token(
+            model_config, mapping)
 
 
 def get_mla_context_workspace_kv_len_cap(kv_cache_config, max_batch_size,
@@ -1171,21 +1139,21 @@ def configure_kv_cache_capacity(self,
         self._kv_cache_config.pool_ratio = self._pool_ratio_in
         self._kv_cache_config.avg_seq_len = self._avg_seq_len_in
 
-        # Reserve headroom for the fp8 context-MLA attention workspace, which the profiling forward
-        # under-measures (fresh-prefill dummies never exercise KV reuse). This is only needed when KV-cache
-        # reuse can push summed attended KV past the profiled floor: get_mla_context_workspace_kv_len_cap
-        # returns None (no reservation) with reuse off -- the workspace is then bounded by max_num_tokens --
-        # or with chunked prefill -- each attention launch is then bounded by its chunk buffer -- since
-        # reserving in those cases would double-count and needlessly shrink the KV pool (up to ~37% for
-        # Kimi-K2 attention-DP). When it does apply,
+        # Reserve headroom for the attention workspace the selected backend declares (today: fp8
+        # context-MLA), which the profiling forward under-measures (fresh-prefill dummies never exercise
+        # KV reuse). This is only needed when KV-cache reuse can push summed attended KV past the profiled
+        # floor: get_mla_context_workspace_kv_len_cap returns None (no reservation) with reuse off -- the
+        # workspace is then bounded by max_num_tokens -- or with chunked prefill -- each attention launch is
+        # then bounded by its chunk buffer -- since reserving in those cases would double-count and
+        # needlessly shrink the KV pool (up to ~37% for Kimi-K2 attention-DP). When it does apply,
         # reserve w * L_cap bytes -- covering the worst-case summed attended KV the scheduler admits
         # (get_mla_context_workspace_kv_len_cap) -- but clamp it 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 reserve
         # covers exactly reserve/w tokens of summed attended KV; that count is carried to the KV manager as
         # the scheduler's admission cap so it never re-derives the cap from pool layout (which V2
         # overstates). No cap or w == 0 -> no-op.
-        w_bytes_per_token = get_mla_context_workspace_bytes_per_token(
+        w_bytes_per_token = get_attention_workspace_bytes_per_token(
             self._model_engine.model.model_config, self._mapping)
         kv_len_cap = get_mla_context_workspace_kv_len_cap(
             self._kv_cache_config, self._max_batch_size, self._max_num_tokens,

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -5460,19 +5460,59 @@ def _get_ctx_mla_kv_len_cap(self):
         reads that decision directly rather than re-deriving it from pool layout, which V2 overstates
         (`blocks_in_primary_pool` forwards `get_page_index_upper_bound`, not the available-page count).
         A carried value of None -- non-fp8-MLA model, no reservation needed (reuse off / chunked prefill),
-        or estimation skipped -- means no admission cap is applied.
+        or estimation skipped -- means no admission cap is applied. A carried 0 is a real cap (a budget so
+        tight the reserve covers under one token), not "no cap", so it must not be collapsed into None.
         """
         cap = getattr(self, "_ctx_mla_kv_len_cap", "unset")
         if cap != "unset":
             return cap
-        if getattr(self, "is_warmup", False):
+        if self.is_warmup:
             # Estimation/warmup runs fresh-prefill dummies against a throwaway manager that reserved
             # nothing; don't cap them (and don't cache -- real serving recomputes from the real manager).
             return None
+        # getattr: managers not built by the estimator (and test doubles) never carry the attribute.
         carried = getattr(self.kv_cache_manager, "fp8_ctx_mla_kv_len_cap", None)
-        self._ctx_mla_kv_len_cap = int(carried) if carried else None
+        self._ctx_mla_kv_len_cap = int(carried) if carried is not None else None
+        self._warn_if_ctx_mla_kv_len_cap_degenerate()
         return self._ctx_mla_kv_len_cap
 
+    def _warn_if_ctx_mla_kv_len_cap_degenerate(self) -> None:
+        """Surface an fp8 context-MLA admission cap too tight to batch context requests on.
+
+        `_cap_context_by_total_kv_len` logs its deferrals at debug level, so a deployment whose budget
+        lands here sees context throughput collapse with nothing in the log at default level. Called from
+        the cached resolution in `_get_ctx_mla_kv_len_cap`, so it fires at most once per executor.
+
+        Only the memory budget can produce a cap this small: the reserve is clamped to
+        `budget * w / (k + w)` (`get_mla_context_workspace_reserve`) and the cap is that reserve divided
+        by `w`, while the `fp8_context_mla_kv_len_cap` override is floored at `max_seq_len`
+        (`get_mla_context_workspace_kv_len_cap`). So the override is never the cause, and raising it is
+        never the fix -- point at the budget instead.
+        """
+        cap = self._ctx_mla_kv_len_cap
+        if cap is None:
+            return
+        if cap == 0:
+            consequence = (
+                "every forward step will schedule exactly one context request -- the "
+                "forward-progress request kept unconditionally -- whatever its length"
+            )
+        elif self.max_seq_len and cap < self.max_seq_len:
+            consequence = (
+                f"a single request attending max_seq_len={self.max_seq_len} already exceeds it, so "
+                "context requests near that length will be scheduled one per forward step"
+            )
+        else:
+            return
+        logger.warning(
+            f"fp8 context-MLA admission cap resolved to {cap} token(s) of summed attended KV: "
+            f"{consequence}. Prefill batching is degraded, not incorrect -- generation is unaffected "
+            "and requests still complete. The cap is the workspace reserve the KV-cache estimator "
+            "could afford divided by its per-token cost, so the KV cache memory budget is the binding "
+            "constraint: raise free_gpu_memory_fraction or max_gpu_total_bytes. Raising "
+            "KvCacheConfig.fp8_context_mla_kv_len_cap will not help (it is already floored at "
+            "max_seq_len).")
+
     @staticmethod
     def _context_attended_kv_len(ctx_req) -> int:
         """Attended KV length this context request contributes to total_kv_len this step.

tests/unittest/_torch/executor/test_kv_cache_estimation.py

@@ -644,10 +644,11 @@ def test_estimation_temporarily_uses_inferred_pool_sizing() -> None:
         patch.object(torch.cuda, "memory_stats", return_value={"allocated_bytes.all.current": 128}),
         patch.object(torch.cuda, "empty_cache"),
         patch.object(torch.cuda, "reset_peak_memory_stats"),
-        # This test exercises inferred pool sizing, not the fp8 context-MLA workspace reserve; the mock
-        # model_config would otherwise walk into the MLA byte-cost path. Neutralize it so no reserve applies.
+        # This test exercises inferred pool sizing, not the backend workspace reserve; the mock
+        # model_config would otherwise walk into backend resolution and the MLA byte-cost path.
+        # Neutralize it so no reserve applies.
         patch(
-            "tensorrt_llm._torch.pyexecutor._util.get_mla_context_workspace_bytes_per_token",
+            "tensorrt_llm._torch.pyexecutor._util.get_attention_workspace_bytes_per_token",
             return_value=0,
         ),
     ):