diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index 762ed7b54164..6bbeaa7944cd 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -760,6 +760,24 @@ size_t AttentionOp::getFmhaMultiCtasKvScratchSize() const noexcept return partialStatsSize + partialOSize; } +size_t AttentionOp::contextMlaWorkspaceBytesPerToken(int32_t numAttnHeads, int32_t qkRopeHeadDim, int32_t qkNopeHeadDim, + int32_t vHeadDim, bool fp8ContextMla, bool separateQAndKvInput, bool sparseMla) noexcept +{ + // Only the fp8 context-MLA separate-Q/KV path stages total_kv_len-scaled K/V dequant buffers. + // Sparse MLA reads K/V directly from the paged KV cache (no staging), so its per-token cost is 0. + if (!fp8ContextMla || !separateQAndKvInput || sparseMla) + { + return 0; + } + // Mirror getWorkspaceSizeForContext's dim layout for the non-sparse fp8 branch: + // total_k_dim_all_heads = numAttnHeads * (qk_rope_head_dim + qk_nope_head_dim) + // total_v_dim_all_heads = numAttnHeads * v_head_dim + // The buffers are fp8 (1 byte/element), so bytes/token == element count. + int const dimKPerHead = qkRopeHeadDim + qkNopeHeadDim; + int const dimVPerHead = vHeadDim; + return static_cast(numAttnHeads) * static_cast(dimKPerHead + dimVPerHead); +} + size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int32_t max_num_seq, int32_t input_seq_length, int32_t cross_kv_length, int32_t max_num_tokens, int32_t total_kv_len) const noexcept { @@ -843,10 +861,17 @@ size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int3 { // Use total_kv_len when available (KV cache reuse causes total_kv_len >> max_num_tokens). // enqueueContext sizes these buffers by total_kv_len, so workspace must match. + // NOTE: the per-token cost of these two buffers (total_k_dim_all_heads + total_v_dim_all_heads) is + // the single source of truth exposed via contextMlaWorkspaceBytesPerToken() for the KV-cache + // estimator's workspace reserve. Keep the two in sync if this dim layout changes. size_t const kv_buf_tokens = std::max( static_cast(total_kv_len), static_cast(mChunkPrefillBufferBatchSize) * max_num_tokens); fp8_k_buf_size = kv_buf_tokens * static_cast(total_k_dim_all_heads); fp8_v_buf_size = kv_buf_tokens * static_cast(total_v_dim_all_heads); + TLLM_CHECK(static_cast(total_k_dim_all_heads + total_v_dim_all_heads) + == contextMlaWorkspaceBytesPerToken(mNumAttnHeads, mMLAParams.qk_rope_head_dim, + mMLAParams.qk_nope_head_dim, mMLAParams.v_head_dim, mFP8ContextMLA, + /*separateQAndKvInput=*/true, useSparseMLA())); } } else if (useSageAttnSeparateQkv) diff --git a/cpp/tensorrt_llm/common/attentionOp.h b/cpp/tensorrt_llm/common/attentionOp.h index 438489348577..e71b2e4d238c 100644 --- a/cpp/tensorrt_llm/common/attentionOp.h +++ b/cpp/tensorrt_llm/common/attentionOp.h @@ -60,6 +60,13 @@ class AttentionOp [[nodiscard]] size_t getWorkspaceSizeForContext(tensorrt_llm::DataType type, int32_t nbReq, int32_t max_input_length, int32_t cross_kv_length = 0, int32_t max_num_tokens = 0, int32_t total_kv_len = 0) const noexcept; + // Per-token byte cost of the context-MLA K/V dequant staging buffers, whose size scales with the summed + // attended KV length (`total_kv_len`). Only the fp8 context-MLA separate-Q/KV path stages these buffers; + // every other path (incl. sparse MLA, which reads K/V straight from the paged cache) returns 0. Single + // source of truth shared by getWorkspaceSizeForContext (runtime sizing) and the KV-cache estimator, so + // the two cannot drift. + [[nodiscard]] static size_t contextMlaWorkspaceBytesPerToken(int32_t numAttnHeads, int32_t qkRopeHeadDim, + int32_t qkNopeHeadDim, int32_t vHeadDim, bool fp8ContextMla, bool separateQAndKvInput, bool sparseMla) noexcept; // total_num_seq is the sum of beam_width for multiple requests [[nodiscard]] size_t getWorkspaceSizeForGeneration(tensorrt_llm::DataType type, int32_t total_num_seq, int32_t max_attention_window_size, int32_t max_num_tokens, int32_t max_blocks_per_sequence) const noexcept; diff --git a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp index dd1ff0db5410..2ac9cc671277 100644 --- a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -187,6 +188,14 @@ void initBindings(nb::module_& m) [](int cp_size) { return tensorrt_llm::kernels::computeHelixWorkspaceSizePerRank(cp_size); }, nb::arg("cp_size"), "Get helix all-to-all workspace size per rank in bytes"); + m.def("get_context_mla_workspace_bytes_per_token", + &tensorrt_llm::common::op::AttentionOp::contextMlaWorkspaceBytesPerToken, nb::arg("num_attn_heads"), + nb::arg("qk_rope_head_dim"), nb::arg("qk_nope_head_dim"), nb::arg("v_head_dim"), nb::arg("fp8_context_mla"), + nb::arg("separate_q_and_kv_input"), nb::arg("sparse_mla"), + "Per-token byte cost of the context-MLA K/V dequant staging buffers (scales with summed attended KV " + "length). Returns 0 outside the fp8 context-MLA separate-Q/KV path. Used by the KV-cache estimator to " + "reserve workspace headroom before sizing the KV pool."); + m.def("compute_flash_mla_metadata", &tensorrt_llm::computeFlashMlaMetadata, nb::arg("seqlens_k"), nb::arg("tile_scheduler_metadata"), nb::arg("num_splits"), nb::arg("batch_size"), nb::arg("s_q"), nb::arg("num_q_heads"), nb::arg("num_kv_heads"), nb::arg("head_size_v"), diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 40ab2e3a64ed..6d08a7eaa6df 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -299,6 +299,106 @@ 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. + """ + 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, + )) + + +def get_mla_context_workspace_kv_len_cap(kv_cache_config, max_batch_size, + max_num_tokens, max_seq_len, + enable_chunked_prefill): + """Max summed attended-KV length (tokens) per forward step the fp8 context-MLA workspace is reserved + -- and the scheduler admits -- for, or ``None`` when no reservation is needed. + + Returns ``None`` unless KV-cache reuse can grow the workspace past the floor the profiling forward + measures: with block reuse off the summed attended KV is bounded by ``max_num_tokens`` (already + profiled), and with chunked prefill each attention launch is independently bounded by its chunk buffer. + In both cases reserving would only double-count and needlessly shrink the KV pool, so no cap is returned. + + Otherwise the default (no override) is the never-stall worst case ``min(max_batch_size, max_num_tokens) + * max_seq_len``: at most that many context requests run in a step, each attending at most ``max_seq_len`` + KV, so reserving for it never defers a request. An explicit ``fp8_context_mla_kv_len_cap`` override + reserves less workspace (freeing KV pool) and lets the scheduler defer over-cap requests; it is floored + at ``max_seq_len`` (one request must always fit) and capped at the worst case. + """ + if not kv_cache_config.enable_block_reuse or enable_chunked_prefill: + return None + worst_case = min(max_batch_size, max_num_tokens) * max_seq_len + override = kv_cache_config.fp8_context_mla_kv_len_cap + if override is None: + return worst_case + return min(max(int(override), max_seq_len), worst_case) + + +def get_mla_context_workspace_reserve(budget_bytes, k_bytes_per_token, + w_bytes_per_token, kv_len_cap): + """Bytes to reserve for the fp8 context-MLA workspace, and the token admission cap that reserve covers. + + Reserve ``w * kv_len_cap`` (the worst-case summed attended KV), clamped to the per-token split + ``budget * w / (k + w)`` so a memory-constrained node shares the budget at a common token count rather + than starving the KV pool. The admission cap is ``reserve / w == min(kv_len_cap, budget / (k + w))`` + tokens; the scheduler admits at most that much summed attended KV, so the fp8 dequant staging buffer + this reserve covers stays within it. This accounts for the fp8 staging term only -- the separate BF16 + full-gather buffers on the reuse path are not yet charged here (tracked as a follow-up), so this bounds + but does not by itself guarantee the reuse-path peak. Returns ``(reserve_bytes, cap_tokens)``, or + ``(0, None)`` when any input is non-positive. + """ + if not (budget_bytes > 0 and k_bytes_per_token > 0 and w_bytes_per_token > 0 + and kv_len_cap and kv_len_cap > 0): + return 0, None + reserve = min( + w_bytes_per_token * kv_len_cap, budget_bytes * w_bytes_per_token / + (k_bytes_per_token + w_bytes_per_token)) + return reserve, int(reserve / w_bytes_per_token) + + def is_vswa_enabled(kv_cache_config): max_attention_window = kv_cache_config.max_attention_window return max_attention_window is not None and len( @@ -421,6 +521,10 @@ def __init__( KVCacheManagerV2) self._draft_config = draft_config self._skip_est = skip_est + # Admission cap (tokens of summed context attended-KV) that the fp8 context-MLA workspace reservation + # covers, computed in configure_kv_cache_capacity and carried to the KV manager so the scheduler + # reads it directly instead of re-deriving it from pool layout. None until reserved (or w == 0). + self._fp8_ctx_mla_kv_len_cap = None self._maybe_enable_fabric_memory_for_python_transceiver() def _maybe_enable_fabric_memory_for_python_transceiver(self) -> None: @@ -1044,6 +1148,41 @@ 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 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( + 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, + self._max_seq_len, self._llm_args.enable_chunked_prefill) + if w_bytes_per_token > 0 and kv_len_cap: + budget_before = kv_cache_max_memory + workspace_reserve, self._fp8_ctx_mla_kv_len_cap = ( + get_mla_context_workspace_reserve( + budget_before, + self._get_kv_size_per_token().slope, w_bytes_per_token, + kv_len_cap)) + if workspace_reserve > 0: + kv_cache_max_memory = int(budget_before - workspace_reserve) + logger.info( + f"Reserving {workspace_reserve / (GB):.2f} GiB for the fp8 context-MLA attention " + f"workspace (w={w_bytes_per_token} B/token, admitting up to " + f"{self._fp8_ctx_mla_kv_len_cap} tokens of summed attended KV): KV cache budget " + f"{budget_before / (GB):.2f} -> {kv_cache_max_memory / (GB):.2f} GiB." + ) + # NOTE: # For KVCacheManager, KvCacheCreator currently controls capacity using two parameters in KVCacheConfig: # • max_tokens @@ -1729,6 +1868,13 @@ def build_managers(self, estimating_kv_cache, kv_cache_config_override=self_kv_cache_config) + # Carry the fp8 context-MLA workspace admission cap (computed in configure_kv_cache_capacity) onto + # the real KV manager so the scheduler reads it directly instead of re-deriving from pool layout. + # The estimation build reserves nothing and runs throwaway fresh-prefill dummies, so leave the + # attribute unset there (PyExecutor._get_ctx_mla_kv_len_cap does not cap during warmup). + if not estimating_kv_cache and kv_cache_manager is not None: + kv_cache_manager.fp8_ctx_mla_kv_len_cap = self._fp8_ctx_mla_kv_len_cap + if (not estimating_kv_cache and self._kv_connector_manager is not None and self._draft_model_engine is not None): raise NotImplementedError( diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 34ebd340a5ee..635cdf0e137b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5268,6 +5268,63 @@ def _waiting_requests(self, context_requests: list[LlmRequest], self.batch_wait_iters_count = 0 return context_requests + def _get_ctx_mla_kv_len_cap(self): + """Cap on the summed context attended-KV length (total_kv_len) per forward step, cached. + + The KV-cache estimator is the single decision point: it reserves the fp8 context-MLA workspace only + when KV-cache reuse can grow it past the profiled floor, and carries the exact token cap that reserve + covers onto the KV manager as `fp8_ctx_mla_kv_len_cap` (`min(L_cap, budget/(k+w))`). The scheduler + 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. + """ + cap = getattr(self, "_ctx_mla_kv_len_cap", "unset") + if cap != "unset": + return cap + if getattr(self, "is_warmup", False): + # 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 + 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 + return self._ctx_mla_kv_len_cap + + @staticmethod + def _context_attended_kv_len(ctx_req) -> int: + """Attended KV length this context request contributes to total_kv_len this step. + + Cached prefix (post-reuse begin position) plus this chunk's new tokens, clamped to the prompt length. + V1 leaves `context_current_position` at 0 with the reuse credit in `estimated_reusable_tokens` (first + chunk only); V2 has already advanced `context_current_position`. + """ + begin = ctx_req.context_current_position + if ctx_req.is_first_context_chunk: + begin = max(begin, ctx_req.estimated_reusable_tokens) + attended = begin + ctx_req.context_chunk_size + return min(attended, ctx_req.orig_prompt_len) + + def _cap_context_by_total_kv_len(self, context_requests): + """Trim scheduled context requests so their summed attended KV length stays within the fp8 + context-MLA workspace reservation (KV-cache reuse can push total_kv_len far past `max_num_tokens`). + The first request is always kept: it attends at most `max_seq_len` and at most its pool, both covered + by the cap, so one request always fits and forward progress is guaranteed. Deferred requests stay + active and retry next iteration, mirroring `_waiting_requests`. + """ + cap = self._get_ctx_mla_kv_len_cap() + if cap is None or len(context_requests) <= 1: + return context_requests + cumulative = 0 + for i, ctx_req in enumerate(context_requests): + cumulative += self._context_attended_kv_len(ctx_req) + if i > 0 and cumulative > cap: + logger.debug( + f"Deferring {len(context_requests) - i} context request(s): summed attended " + f"KV length {cumulative} would exceed the fp8 context-MLA workspace cap {cap}." + ) + return context_requests[:i] + return context_requests + @nvtx_range("_schedule") def _schedule(self): if hasattr(self.kv_cache_manager, "prepare_expect_snapshot_points"): @@ -5304,6 +5361,11 @@ def _schedule(self): scheduled_context_requests) num_fitting = len(scheduled_context_requests) + # Cap summed context attended-KV length so the fp8 context-MLA attention workspace stays within the + # headroom the estimator reserved for it (no-op for non-fp8-MLA models). + scheduled_context_requests = self._cap_context_by_total_kv_len( + scheduled_context_requests) + scheduled_requests = ScheduledRequests() scheduled_requests.encoder_requests = scheduler_output.encoder_requests scheduled_requests.reset_context_requests(scheduled_context_requests) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 68d9fdbeb207..55de55dfa4d9 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3819,6 +3819,20 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): "Set to 0 to disable prefetch. Only effective with KV cache manager v2 and block reuse enabled." ) + # This is a pure python field, not a pybind field. It is only for the Pytorch backend. + fp8_context_mla_kv_len_cap: Optional[int] = Field( + default=None, + status="prototype", + description= + "Override, in tokens, for the max summed attended-KV length (total_kv_len) per forward step that " + "the fp8 context-MLA attention workspace is reserved and scheduled for. Only affects fp8 " + "context-MLA models (e.g. DeepSeek / Kimi with an fp8 KV cache). None (default) reserves for the " + "never-stall worst case min(max_batch_size, max_num_tokens) * max_seq_len. A smaller value reserves " + "less workspace (freeing KV cache) and defers context requests whose summed attended KV would " + "exceed it; it is floored at max_seq_len and capped at the worst case. Safe at any value -- the " + "scheduler enforces it -- trading prefill batching under heavy reuse for KV cache capacity." + ) + # This is a pure python field, not a pybind field. It is only for the Pytorch backend. pool_ratio: Optional[List[float]] = Field( default=None, diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index f20e02169d62..fd61153bcdd2 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -656,6 +656,13 @@ "kind": "value", "path": "kv_cache_config.event_buffer_max_size" }, + { + "allowed_values": [], + "annotation": "Optional[int]", + "converter": "", + "kind": "value", + "path": "kv_cache_config.fp8_context_mla_kv_len_cap" + }, { "allowed_values": [], "annotation": "Optional[float]", diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index c2af496c1350..edba06398bb3 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -76,7 +76,6 @@ accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-ep4-cutl accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-tp4-cutlass-auto] SKIP (https://nvbugs/5596343) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_chunked_prefill[cutlass-auto] SKIP (https://nvbugs/5596343) accuracy/test_llm_api_pytorch.py::TestKanana_Instruct::test_auto_dtype SKIP (https://nvbugs/6209806) -accuracy/test_llm_api_pytorch.py::TestKimiK2::test_nvfp4[4gpus] SKIP (https://nvbugs/6368562) accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[pp4-attn_backend=TRTLLM-torch_compile=False] SKIP (https://nvbugs/6490043) accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[tp2pp2-attn_backend=FLASHINFER-torch_compile=True] SKIP (https://nvbugs/6422337) accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[tp4-attn_backend=TRTLLM-torch_compile=False] SKIP (https://nvbugs/5616182) diff --git a/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py b/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py index 4782ead014de..06e0231a1a0b 100644 --- a/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py +++ b/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py @@ -190,6 +190,7 @@ def _make_creator( creator._execution_stream = None creator._draft_config = None creator._skip_est = True + creator._fp8_ctx_mla_kv_len_cap = None return creator diff --git a/tests/unittest/_torch/executor/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/test_kv_cache_estimation.py index 6a30bf974d68..cd0264eb276c 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/test_kv_cache_estimation.py @@ -644,6 +644,12 @@ 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. + patch( + "tensorrt_llm._torch.pyexecutor._util.get_mla_context_workspace_bytes_per_token", + return_value=0, + ), ): assert creator.try_prepare_estimation() assert kv_cache_config.max_tokens == estimation_max_tokens diff --git a/tests/unittest/_torch/executor/test_mla_workspace_reserve.py b/tests/unittest/_torch/executor/test_mla_workspace_reserve.py new file mode 100644 index 000000000000..1cf9aa2e5706 --- /dev/null +++ b/tests/unittest/_torch/executor/test_mla_workspace_reserve.py @@ -0,0 +1,316 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for the fp8 context-MLA attention-workspace reservation and its admission cap. + +The estimator reserves KV-cache headroom for the total_kv_len-scaled context-MLA workspace and carries the +exact token cap that reserve covers onto the KV manager; the scheduler reads that cap directly (rather than +re-deriving it from pool layout) and trims context requests whose summed attended KV would exceed it. These +tests cover the pure-Python pieces: the reservation gate (reuse-on / chunked-off), the workspace reserve/cap +math, the L_cap derivation, the per-request attended-KV computation, the admission trim, the carried-cap +read (V1/V2 layout-independent), and the non-MLA no-op gate. +""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from tensorrt_llm._torch.pyexecutor import _util +from tensorrt_llm._torch.pyexecutor._util import ( + get_mla_context_workspace_bytes_per_token, + get_mla_context_workspace_kv_len_cap, + get_mla_context_workspace_reserve, +) +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + +def _ctx_req( + context_current_position, + is_first_context_chunk, + estimated_reusable_tokens, + context_chunk_size, + orig_prompt_len, +): + return SimpleNamespace( + context_current_position=context_current_position, + is_first_context_chunk=is_first_context_chunk, + estimated_reusable_tokens=estimated_reusable_tokens, + context_chunk_size=context_chunk_size, + orig_prompt_len=orig_prompt_len, + ) + + +@pytest.mark.parametrize( + "req,expected", + [ + # V1: reuse not yet applied at schedule time; credit is in estimated_reusable_tokens. + (_ctx_req(0, True, 100, 50, 150), 150), + # V2: context_current_position already advanced past the reused prefix. + (_ctx_req(100, True, 100, 50, 150), 150), + # Fresh prefill, no reuse. + (_ctx_req(0, True, 0, 80, 80), 80), + # Middle (non-first) chunk: reuse credit does not apply, position already advanced. + (_ctx_req(200, False, 0, 100, 500), 300), + # Reuse estimate above the prompt length is clamped to the prompt length. + (_ctx_req(0, True, 999, 50, 150), 150), + ], +) +def test_context_attended_kv_len(req, expected): + assert PyExecutor._context_attended_kv_len(req) == expected + + +def _make_executor(cap): + # Bypass __init__; the trim only reads the cached cap and the two helper methods. + exe = object.__new__(PyExecutor) + exe._ctx_mla_kv_len_cap = cap # pre-set so the cap getter skips the C++ binding + return exe + + +def test_cap_trims_tail_by_total_kv_len(): + exe = _make_executor(cap=120) + reqs = [_ctx_req(0, True, 0, 50, 50) for _ in range(3)] # cumulative 50, 100, 150 + kept = exe._cap_context_by_total_kv_len(reqs) + assert len(kept) == 2 # the third would push cumulative to 150 > 120 + + +def test_cap_keeps_all_when_within_budget(): + exe = _make_executor(cap=200) + reqs = [_ctx_req(0, True, 0, 50, 50) for _ in range(3)] + assert len(exe._cap_context_by_total_kv_len(reqs)) == 3 + + +def test_cap_always_keeps_first_even_if_it_alone_exceeds(): + exe = _make_executor(cap=120) + reqs = [_ctx_req(0, True, 0, 200, 200), _ctx_req(0, True, 0, 50, 50)] + kept = exe._cap_context_by_total_kv_len(reqs) + assert len(kept) == 1 # first kept despite exceeding cap (forward-progress guard) + + +def test_cap_single_request_never_trimmed(): + exe = _make_executor(cap=1) + reqs = [_ctx_req(0, True, 0, 500, 500)] + assert exe._cap_context_by_total_kv_len(reqs) == reqs + + +def test_no_cap_returns_untouched(): + exe = _make_executor(cap=None) # non-fp8-MLA model: no reservation, no trimming + reqs = [_ctx_req(0, True, 0, 50, 50) for _ in range(5)] + assert exe._cap_context_by_total_kv_len(reqs) is reqs + + +@pytest.mark.parametrize( + "override,expected", + [ + # Default (no override): never-stall worst case min(bs, num_tokens) * max_seq_len + # = min(64, 8192) * 4096. + (None, 64 * 4096), + # In-range override passes through. + (100_000, 100_000), + # Below the max_seq_len floor is raised to max_seq_len (one request must always fit). + (1_000, 4096), + # Above the worst case is capped at the worst case (a larger value only over-reserves). + (10**12, 64 * 4096), + ], +) +def test_kv_len_cap_default_floor_and_ceiling(override, expected): + # Reuse on + chunked off: a cap is reserved (default worst case, or the clamped override). + cfg = SimpleNamespace(enable_block_reuse=True, fp8_context_mla_kv_len_cap=override) + assert ( + get_mla_context_workspace_kv_len_cap( + cfg, + max_batch_size=64, + max_num_tokens=8192, + max_seq_len=4096, + enable_chunked_prefill=False, + ) + == expected + ) + + +# ---- reservation gate: only when reuse can grow the workspace past the profiled floor ---- + + +@pytest.mark.parametrize( + "enable_block_reuse,enable_chunked_prefill", + [ + # Reuse off: summed attended KV bounded by max_num_tokens (already profiled) -> no reservation. + (False, False), + # Chunked prefill: each attention launch bounded by its own chunk buffer -> no reservation. + (True, True), + (False, True), + ], +) +def test_kv_len_cap_none_when_reuse_cannot_grow_workspace( + enable_block_reuse, enable_chunked_prefill +): + # None cap -> configure_kv_cache_capacity reserves nothing and applies no admission cap for these + # unaffected configs, even with an explicit override present. + cfg = SimpleNamespace(enable_block_reuse=enable_block_reuse, fp8_context_mla_kv_len_cap=999_999) + assert ( + get_mla_context_workspace_kv_len_cap( + cfg, + max_batch_size=64, + max_num_tokens=8192, + max_seq_len=4096, + enable_chunked_prefill=enable_chunked_prefill, + ) + is None + ) + + +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 + + +def _fp8_mla_model_config(sparse_algorithm): + # Duck-typed ModelConfig: MLA (kv_lora_rank + qk_rope_head_dim), fp8 KV cache, DeepSeek-V3 head dims. + sparse_cfg = ( + SimpleNamespace(algorithm=sparse_algorithm) if sparse_algorithm is not None else None + ) + return SimpleNamespace( + pretrained_config=SimpleNamespace( + kv_lora_rank=512, + qk_rope_head_dim=64, + qk_nope_head_dim=128, + v_head_dim=128, + num_attention_heads=128, + ), + quant_config=SimpleNamespace(quant_mode=SimpleNamespace(has_fp8_kv_cache=lambda: True)), + sparse_attention_config=sparse_cfg, + ) + + +@pytest.mark.parametrize( + "sparse_algorithm,sm,short_seq_mha,expect_zero", + [ + # Absorption path, TRTLLM-gen SM, fallback off -> K/V read from the paged cache, nothing staged. + ("dsa", 100, "0", True), + ("deepseek_v4", 103, "0", True), + # Short-seq MHA fallback sends short contexts down the dense path -> buffer is staged. + ("dsa", 100, "1024", False), + # Skip-softmax is a sparse config but passes no sparse indices to C++, so MLA stays dense. + ("skip_softmax", 100, "0", False), + # mUseTllmGen is `sm >= 100 && sm != 120`, so SM90 / SM120 are dense whatever the config says. + ("dsa", 90, "0", False), + ("dsa", 120, "0", False), + # No sparse config at all. + (None, 100, "0", False), + ], +) +def test_workspace_bytes_zero_only_for_absorption_mode_sparse_mla( + monkeypatch, sparse_algorithm, sm, short_seq_mha, expect_zero +): + # 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.setenv("TRTLLM_MLA_SHORT_SEQ_MHA_THRESHOLD", short_seq_mha) + w = get_mla_context_workspace_bytes_per_token( + _fp8_mla_model_config(sparse_algorithm), + SimpleNamespace(enable_attention_dp=False, tp_size=8), + ) + assert (w == 0) if expect_zero else (w > 0) + + +# ---- workspace reserve / admission-cap math (estimator side) ---- + + +@pytest.mark.parametrize( + "budget,k,w,kv_len_cap,expected_cap", + [ + # Worst case fits the budget: reserve = w * kv_len_cap, so cap = kv_len_cap. + (10**12, 1000, 100, 50_000, 50_000), + # Memory-constrained: w * kv_len_cap exceeds the per-token split, so cap = budget / (k + w). + # budget/(k+w) = 1_100_000 / 1100 = 1000. + (1_100_000, 1000, 100, 10**9, 1000), + ], +) +def test_workspace_reserve_cap(budget, k, w, kv_len_cap, expected_cap): + reserve, cap = get_mla_context_workspace_reserve(budget, k, w, kv_len_cap) + assert cap == expected_cap + assert cap == int(reserve / w) # the cap is exactly what the reserve covers + assert reserve <= w * kv_len_cap # never reserves beyond the worst case + assert cap <= kv_len_cap # never admits more than L_cap + + +@pytest.mark.parametrize( + "budget,k,w,kv_len_cap", + [ + (0, 1000, 100, 50_000), # no budget + (10**12, 0, 100, 50_000), # k <= 0 + (10**12, 1000, 0, 50_000), # w <= 0 + (10**12, 1000, 100, None), # no L_cap + (10**12, 1000, 100, 0), # L_cap == 0 + ], +) +def test_workspace_reserve_zero_for_bad_inputs(budget, k, w, kv_len_cap): + # Any non-positive input -> reserve nothing and carry no cap (no admission cap applied). + assert get_mla_context_workspace_reserve(budget, k, w, kv_len_cap) == (0, None) + + +# ---- 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. + # 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 + return exe + + +@pytest.mark.parametrize( + "manager", + [ + # V1-style manager: unified-pool block count present but must NOT be consulted anymore. + SimpleNamespace( + fp8_ctx_mla_kv_len_cap=262144, blocks_in_primary_pool=4096, tokens_per_block=64 + ), + # V2-style manager: page-index upper bound overstates capacity; must NOT be consulted. + SimpleNamespace( + fp8_ctx_mla_kv_len_cap=262144, blocks_in_primary_pool=10**9, tokens_per_block=64 + ), + ], +) +def test_ctx_cap_reads_carried_value_ignoring_pool_layout(manager): + # Both managers carry the same estimator cap; the (very different) pool layout is irrelevant. + exe = _executor_with_manager(manager) + assert exe._get_ctx_mla_kv_len_cap() == 262144 + + +@pytest.mark.parametrize( + "manager", + [ + # Estimator made no reservation -- non-fp8-MLA, reuse off / chunked prefill, or estimation skipped. + SimpleNamespace(fp8_ctx_mla_kv_len_cap=None), + # Attribute absent (manager not built by the estimator): treated the same as None, no admission cap. + SimpleNamespace(), + ], +) +def test_ctx_cap_none_when_no_reservation(manager): + # A carried None (or absent) cap means no headroom to enforce -> admission disabled, no crash. + exe = _executor_with_manager(manager) + assert exe._get_ctx_mla_kv_len_cap() is None + + +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. + exe = _executor_with_manager(SimpleNamespace(fp8_ctx_mla_kv_len_cap=262144), is_warmup=True) + assert exe._get_ctx_mla_kv_len_cap() is None