From 0f5c162d1d28e68613fdc4333f3e3973496e7ef9 Mon Sep 17 00:00:00 2001 From: Yueh-Ting Chen Date: Wed, 15 Jul 2026 15:58:43 +0800 Subject: [PATCH] [None][fix] Declare attention runtime-workspace bytes/token as a backend contract The fp8 context-MLA workspace reservation (nvbugs/6368562, #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 #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 #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 #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 --- .../_torch/attention_backend/interface.py | 16 ++++ .../_torch/attention_backend/trtllm.py | 58 ++++++++++++ .../modules/ATTENTION_DEVELOPER_GUIDE.md | 22 +++++ tensorrt_llm/_torch/pyexecutor/_util.py | 82 ++++++----------- tensorrt_llm/_torch/pyexecutor/py_executor.py | 46 +++++++++- .../executor/test_kv_cache_estimation.py | 7 +- .../executor/test_mla_workspace_reserve.py | 89 ++++++++++++++++--- 7 files changed, 246 insertions(+), 74 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index f94fbfca3e94..7a3977b910c6 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/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. diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index dd756d561849..38cbd312691b 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/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 diff --git a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md index 10cc9a9ae872..405c83cf6a99 100644 --- a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md +++ b/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. diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index c4fb115111fe..1eb79abd12c4 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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,13 +1139,13 @@ 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 @@ -1185,7 +1153,7 @@ def configure_kv_cache_capacity(self, # 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, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 62b7c6488ee9..de5281131aec 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/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. diff --git a/tests/unittest/_torch/executor/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/test_kv_cache_estimation.py index 5c75c9186bd7..7a24eb5d019c 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_estimation.py +++ b/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, ), ): diff --git a/tests/unittest/_torch/executor/test_mla_workspace_reserve.py b/tests/unittest/_torch/executor/test_mla_workspace_reserve.py index 1cf9aa2e5706..5a9fe1726f11 100644 --- a/tests/unittest/_torch/executor/test_mla_workspace_reserve.py +++ b/tests/unittest/_torch/executor/test_mla_workspace_reserve.py @@ -23,13 +23,14 @@ """ from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import Mock, patch import pytest -from tensorrt_llm._torch.pyexecutor import _util +from tensorrt_llm._torch.attention_backend import trtllm as trtllm_backend +from tensorrt_llm._torch.pyexecutor import py_executor as py_executor_module from tensorrt_llm._torch.pyexecutor._util import ( - get_mla_context_workspace_bytes_per_token, + get_attention_workspace_bytes_per_token, get_mla_context_workspace_kv_len_cap, get_mla_context_workspace_reserve, ) @@ -171,9 +172,25 @@ def test_kv_len_cap_none_when_reuse_cannot_grow_workspace( def test_workspace_bytes_zero_for_non_mla_model(): - # No kv_lora_rank on the config -> not MLA -> 0 (early return, no binding call needed). - model_config = SimpleNamespace(pretrained_config=SimpleNamespace(), quant_config=None) - assert get_mla_context_workspace_bytes_per_token(model_config, Mock()) == 0 + # No kv_lora_rank on the config -> the TRTLLM backend declares no reservation -> 0 (early return, no + # binding call needed). Exercises the full resolve-backend path, not just the TRTLLM classmethod. + model_config = SimpleNamespace( + pretrained_config=SimpleNamespace(), quant_config=None, attn_backend="TRTLLM" + ) + assert get_attention_workspace_bytes_per_token(model_config, Mock()) == 0 + + +def test_workspace_bytes_zero_for_backend_without_declaration(): + # A backend that stages no runtime-scaled workspace inherits the default 0 from AttentionBackend, so + # the estimator reserves nothing and the scheduler applies no admission cap for it -- even for an MLA + # model that the TRTLLM backend would charge for. VANILLA is used because it always resolves (FLASHINFER + # silently falls back to TRTLLM when flashinfer is not installed). + model_config = SimpleNamespace( + pretrained_config=SimpleNamespace(kv_lora_rank=512, qk_rope_head_dim=64), + quant_config=None, + attn_backend="VANILLA", + ) + assert get_attention_workspace_bytes_per_token(model_config, Mock()) == 0 def _fp8_mla_model_config(sparse_algorithm): @@ -191,6 +208,7 @@ def _fp8_mla_model_config(sparse_algorithm): ), quant_config=SimpleNamespace(quant_mode=SimpleNamespace(has_fp8_kv_cache=lambda: True)), sparse_attention_config=sparse_cfg, + attn_backend="TRTLLM", ) @@ -216,9 +234,9 @@ def test_workspace_bytes_zero_only_for_absorption_mode_sparse_mla( ): # Reporting w == 0 for a config that still stages the fp8 K/V buffers reserves nothing and installs no # admission cap -- exactly the mid-forward OOM this reservation exists to prevent. - monkeypatch.setattr(_util, "get_sm_version", lambda: sm) + monkeypatch.setattr(trtllm_backend, "get_sm_version", lambda: sm) monkeypatch.setenv("TRTLLM_MLA_SHORT_SEQ_MHA_THRESHOLD", short_seq_mha) - w = get_mla_context_workspace_bytes_per_token( + w = get_attention_workspace_bytes_per_token( _fp8_mla_model_config(sparse_algorithm), SimpleNamespace(enable_attention_dp=False, tp_size=8), ) @@ -264,14 +282,16 @@ def test_workspace_reserve_zero_for_bad_inputs(budget, k, w, kv_len_cap): # ---- carried admission-cap read (executor side, V1/V2) ---- -def _executor_with_manager(manager, is_warmup=False): - # Bypass __init__; the cap getter only reads is_warmup and the cap carried on kv_cache_manager. The - # estimator is the single decision point -- the scheduler never re-derives w or the pool layout. +def _executor_with_manager(manager, is_warmup=False, max_seq_len=4096): + # Bypass __init__; the cap getter reads is_warmup, the cap carried on kv_cache_manager, and max_seq_len + # (only to decide whether the resolved cap is too tight to batch on -- see the degenerate-cap warning). + # The estimator is the single decision point -- the scheduler never re-derives w or the pool layout. # Set the _is_warmup backing field directly -- the is_warmup property setter also propagates into # model_engine, which this bare (no-__init__) executor doesn't have. exe = object.__new__(PyExecutor) exe._is_warmup = is_warmup exe.kv_cache_manager = manager + exe.max_seq_len = max_seq_len return exe @@ -309,6 +329,53 @@ def test_ctx_cap_none_when_no_reservation(manager): assert exe._get_ctx_mla_kv_len_cap() is None +def test_ctx_cap_zero_is_a_cap_not_no_cap(): + # A budget so tight the reserve covers under one token carries 0. Collapsing that into None would + # disable admission control for exactly the case it is most needed -- it must stay a real cap. + exe = _executor_with_manager(SimpleNamespace(fp8_ctx_mla_kv_len_cap=0)) + assert exe._get_ctx_mla_kv_len_cap() == 0 + # Enforced as a cap: everything past the first request (the forward-progress guard) is deferred. + reqs = [_ctx_req(0, True, 0, 50, 50) for _ in range(3)] + assert len(exe._cap_context_by_total_kv_len(reqs)) == 1 + + +@pytest.mark.parametrize( + "cap,max_seq_len,should_warn", + [ + # Reserve covers under one token: one context request per step, forever. + (0, 4096, True), + # Below max_seq_len: a single max-length request already exceeds the cap, so requests near that + # length serialize. Only the memory budget can produce this -- the user override is floored at + # max_seq_len -- so the warning must fire here too, not just at 0. + (4095, 4096, True), + # Exactly max_seq_len: one max-length request still fits; batching degrades gracefully. + (4096, 4096, False), + (262144, 4096, False), + # No reservation at all -- admission disabled, nothing to warn about. + (None, 4096, False), + # max_seq_len unknown: only the unambiguous 0 case warns, and comparing against None must not raise. + (4095, None, False), + (0, None, True), + ], +) +def test_ctx_cap_degenerate_warns_once(cap, max_seq_len, should_warn): + # A cap this tight collapses context batching, and the per-iteration deferral log is debug-level, so + # the resolution point must say so once at warning level or the throughput loss is invisible. + exe = _executor_with_manager( + SimpleNamespace(fp8_ctx_mla_kv_len_cap=cap), max_seq_len=max_seq_len + ) + with patch.object(py_executor_module, "logger") as mock_logger: + assert exe._get_ctx_mla_kv_len_cap() == cap + # Resolution is cached, so re-reading must not re-warn. + exe._get_ctx_mla_kv_len_cap() + assert mock_logger.warning.call_count == (1 if should_warn else 0) + if should_warn: + # Name the budget as the lever; the override is floored at max_seq_len and cannot fix this. + message = mock_logger.warning.call_args[0][0] + assert "free_gpu_memory_fraction" in message + assert str(cap) in message + + def test_ctx_cap_no_cap_during_warmup(): # Estimation/warmup runs fresh-prefill dummies against a throwaway manager that reserved nothing: don't # cap even if a cap were carried, and short-circuit before reading the manager.