From 907577df7fd9b549aa4119e8c9e2ce7724d3603b Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Sun, 16 Aug 2026 22:20:43 -0700 Subject: [PATCH 1/5] [TRTLLM-15033][feat] Upstream Kimi K3 MLA decode backend selection to main Port the Kimi K3 MLA decode backend-selection feature from feat/kimi_k3 (PRs #17320 and #17363, TRTLLM-15001) onto main: K3's absorbed MLA generation runs on the FlashInfer CuTe-DSL kernel for BF16 KV cache and falls back to trtllm-gen for FP8 KV cache, mixed context/generation batches, and multi-token generation. - flashinfer_trtllm_gen.py: validate and store the requested MLA backend, size and stage the CuTe-DSL workspace (page table + sequence lengths staged once per step, keyed for CUDA-graph safety), and route the MLA decode call through the per-batch effective backend. - trtllm.py: TrtllmAttention gains the flashinfer_mla_backend selector and the mla_backend_policy per-batch override hook; the metadata carries the CuTe-DSL staging key with per-step resets. - utils.py / mla.py: thread flashinfer_mla_backend from MLA.__init__ into backend construction. - kimi_k3_mla_attention.py: select K3's generation backend (TLLM_K3_MLA_GEN_BACKEND, default cute-dsl; FP8 KV forces trtllm-gen) and install K3's per-batch fallback policy on the mqa backend. Fold in the four TRTLLM-15033 review follow-ups from #17320: - rename the backend-policy token argument to num_gen_tokens and pass the generation-token count from both call sites, making the previously accidental caller agreement an explicit contract; - log the FP8-KV override once per process instead of once per layer; - validate TLLM_K3_MLA_GEN_BACKEND at read time with an error naming the env var, and document the selector and K3 behavior in ATTENTION_DEVELOPER_GUIDE.md; - reject flashinfer_mla_backend on non-TRTLLM attention backends in create_attention instead of failing with a raw TypeError. Signed-off-by: Brian Nguyen --- .../fmha/flashinfer_trtllm_gen.py | 221 +++++++++++++++++- .../_torch/attention_backend/trtllm.py | 29 ++- .../_torch/attention_backend/utils.py | 13 +- .../modules/ATTENTION_DEVELOPER_GUIDE.md | 18 ++ .../kimi_k3_mla/kimi_k3_mla_attention.py | 71 ++++++ tensorrt_llm/_torch/modules/mla.py | 7 +- .../_torch/attention/test_fmha_page_index.py | 94 ++++++++ .../modules/test_kimi_k3_mla_backend.py | 91 ++++++++ 8 files changed, 530 insertions(+), 14 deletions(-) create mode 100644 tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py diff --git a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py index 49330a61e31d..4566c15d7fe6 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py @@ -16,9 +16,9 @@ """ FlashInfer TRTLLM-Gen FMHA -This module implements attention computation using flashinfer's trtllm-gen kernels. -It provides a TRT-LLM attention FMHA library for trtllm-gen kernels -(Blackwell architecture: SM100/SM103). Enable or disable it through +This module implements attention computation using flashinfer's trtllm-gen kernels, +with an optional CuTeDSL kernel for MLA generation. It provides a TRT-LLM attention +FMHA library for Blackwell architecture (SM100/SM103). Enable or disable it through ``TLLM_FMHA_LIBS``. Architecture: @@ -207,6 +207,18 @@ def _cached_build( _install_flashinfer_mla_decode_tuning_config_cache() +_SUPPORTED_MLA_BACKENDS = {"cute-dsl", "trtllm-gen"} + + +def _get_mla_backend(backend: str) -> str: + backend = backend.strip().lower() + if backend not in _SUPPORTED_MLA_BACKENDS: + raise ValueError( + f"flashinfer_mla_backend must be one of {_SUPPORTED_MLA_BACKENDS}, got {backend!r}." + ) + return backend + + _MULTI_CTAS_KV_COUNTER_ALIGNMENT = 8 @@ -230,6 +242,84 @@ def _get_bmm1_scale_log2(bmm1_scale: torch.Tensor) -> torch.Tensor: return bmm1_scale.narrow(0, 1, 1) +@lru_cache(maxsize=128) +def _get_cute_dsl_mla_workspace_size( + max_batch_size: int, + q_len: int, + num_heads: int, + kv_lora_rank: int, + multi_processor_count: int, +) -> int: + from flashinfer.cute_dsl.attention.monolithic.mla_decode import _get_split_kv_and_workspace_size + + return max( + _get_split_kv_and_workspace_size( + batch_size, + q_len, + num_heads, + kv_lora_rank, + multi_processor_count, + )[1] + for batch_size in range(1, max_batch_size + 1) + ) + + +def _get_cute_dsl_mla_buffer_layout( + batch_size: int, + padded_num_pages: int, +) -> Tuple[int, int, int]: + buffer_alignment_bytes = 32 + page_table_bytes = batch_size * padded_num_pages * torch.int32.itemsize + sequence_lengths_offset = ( + math.ceil(page_table_bytes / buffer_alignment_bytes) * buffer_alignment_bytes + ) + kernel_workspace_offset = ( + math.ceil( + (sequence_lengths_offset + batch_size * torch.int32.itemsize) / buffer_alignment_bytes + ) + * buffer_alignment_bytes + ) + return page_table_bytes, sequence_lengths_offset, kernel_workspace_offset + + +def _prepare_cute_dsl_mla_buffers( + workspace: torch.Tensor, + block_tables: torch.Tensor, + sequence_lengths: torch.Tensor, + padded_num_pages: int, + skip_copy: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if block_tables.size(-1) > padded_num_pages: + raise RuntimeError("CuTeDSL MLA page table exceeds its pre-allocated shape.") + batch_size = block_tables.size(0) + if sequence_lengths.numel() < batch_size: + raise RuntimeError("CuTeDSL MLA sequence lengths are smaller than the batch size.") + + page_table_bytes, sequence_lengths_offset, kernel_workspace_offset = ( + _get_cute_dsl_mla_buffer_layout(batch_size, padded_num_pages) + ) + workspace_bytes = workspace.view(torch.uint8).flatten() + page_table_storage = ( + workspace_bytes[:page_table_bytes].view(torch.int32).view(batch_size, padded_num_pages) + ) + if not skip_copy: + page_table_storage.zero_() + page_table_storage[:, : block_tables.size(-1)].copy_(block_tables[:, 0, :]) + + sequence_lengths_storage = workspace_bytes[ + sequence_lengths_offset:kernel_workspace_offset + ].view(torch.int32) + aligned_sequence_lengths = sequence_lengths_storage[:batch_size] + if not skip_copy: + aligned_sequence_lengths.copy_(sequence_lengths.flatten()[:batch_size]) + + kernel_workspace_bytes = (workspace_bytes.numel() - kernel_workspace_offset) // 4 * 4 + kernel_workspace = workspace_bytes[ + kernel_workspace_offset : kernel_workspace_offset + kernel_workspace_bytes + ].view(-1, 4) + return kernel_workspace, page_table_storage, aligned_sequence_lengths + + @lru_cache(maxsize=128) def _get_context_workspace_layout( dtype: torch.dtype, @@ -358,7 +448,7 @@ def _get_workspace_size( class FlashInferTrtllmGenFmha(PhasedFmha): """ - An attention backend using pure trtllm-gen kernels from flashinfer. + An attention backend using FlashInfer trtllm-gen and optional MLA CuTeDSL kernels. """ # Default KV layout for flashinfer @@ -415,6 +505,13 @@ class FlashInferTrtllmGenFmha(PhasedFmha): def __init__(self, attn: "TrtllmAttention"): super().__init__(attn) self._layout = self.DEFAULT_KV_LAYOUT + requested_mla_backend = _get_mla_backend(attn.flashinfer_mla_backend) + if requested_mla_backend == "cute-dsl" and attn.is_mla_enable and attn.has_fp8_kv_cache: + raise ValueError( + "flashinfer_mla_backend='cute-dsl' does not support FP8 KV cache device scales; " + "use 'trtllm-gen' instead." + ) + self._mla_backend = requested_mla_backend # Read once so the hot path is not sensitive to later environment changes. self._enable_pdl = get_env_enable_pdl() @@ -840,12 +937,39 @@ def prepare_workspace( fp8_context_fmha=fp8_context_fmha, ) + effective_mla_backend = self._get_effective_mla_backend(metadata, num_gen_tokens) + if is_gen_only and attn.is_mla_enable and effective_mla_backend == "cute-dsl": + if metadata.kv_cache_manager is None: + raise RuntimeError("CuTeDSL MLA requires a paged KV cache manager.") + max_batch_size = metadata.max_num_sequences or metadata.max_num_requests + max_num_pages = metadata.kv_cache_manager.max_blocks_per_seq + tokens_per_block = metadata.tokens_per_block or 64 + pages_per_superblock = 128 // tokens_per_block + padded_num_pages = ( + math.ceil(max_num_pages / pages_per_superblock) * pages_per_superblock + ) + _, _, buffer_metadata_size = _get_cute_dsl_mla_buffer_layout( + max_batch_size, padded_num_pages + ) + assert self._multi_processor_count is not None + cute_dsl_workspace_size = _get_cute_dsl_mla_workspace_size( + max_batch_size, + max(1, attn.predicted_tokens_per_seq), + attn.num_heads, + int(attn.kv_lora_rank or 0), + self._multi_processor_count, + ) + required_workspace_size = max( + required_workspace_size, buffer_metadata_size + cute_dsl_workspace_size + ) + required_workspace_size = math.ceil(required_workspace_size / 4) * 4 + current_workspace_size = workspace.numel() * workspace.element_size() if current_workspace_size < required_workspace_size: if metadata.is_cuda_graph and torch.cuda.is_current_stream_capturing(): raise RuntimeError( "Attention CUDA graph workspace is smaller than the " - "required size for trtllm-gen." + f"required size for {self._mla_backend}." ) required_workspace_numel = math.ceil(required_workspace_size / workspace.element_size()) workspace.resize_((required_workspace_numel,)) @@ -856,6 +980,26 @@ def _get_multi_ctas_kv_counter_buffer(self) -> torch.Tensor: raise RuntimeError("The trtllm-gen multi-CTA KV counter buffer is not initialized.") return counter_buffer + def _get_effective_mla_backend( + self, meta: "TrtllmAttentionMetadata", num_gen_tokens: int + ) -> str: + """Resolve the MLA decode backend for the current scheduler batch. + + ``attn.mla_backend_policy`` is a neutral per-batch override hook. + A model whose backend choice depends on the batch composition (e.g. + Kimi K3's MLA module) installs a policy on the attention instance it + owns; the policy receives the statically configured backend, the batch + metadata, and the batch's generation-token count, and returns the + backend to use for this batch. Without a policy (the default for every + model) this returns ``self._mla_backend`` unchanged, so every call + site behaves exactly like the plain ``self._mla_backend`` checks it + replaced. + """ + policy = self.attn.mla_backend_policy + if policy is None: + return self._mla_backend + return policy(self._mla_backend, meta, num_gen_tokens) + @staticmethod def _compute_window_left( cyclic_attention_window_size: int, @@ -1170,6 +1314,9 @@ def run_mla_generation( attn = params.attn meta = params.meta fwd = params.fwd + # ``params.num_tokens`` is the generation-token count here: phased.py + # dispatches MLA generation with the generation slice only. + effective_mla_backend = self._get_effective_mla_backend(meta, params.num_tokens) if 0 < params.cyclic_attention_window_size < params.max_past_kv_length: raise NotImplementedError( "Sliding-window attention is not supported by MLA decode path." @@ -1239,7 +1386,52 @@ def run_mla_generation( ) bmm1_scale = 1.0 / (attn.q_scaling * math.sqrt(qk_nope_head_dim + qk_rope_head_dim)) bmm2_scale = 1.0 - workspace_buffer = params.workspace.view(-1, 4) + if effective_mla_backend == "cute-dsl": + pages_per_superblock = 128 // params.tokens_per_block + padded_num_pages = ( + math.ceil(block_tables.size(-1) / pages_per_superblock) * pages_per_superblock + ) + # Every MLA layer of one forward step shares this metadata, its + # workspace, and (for a single paged pool) identical block tables + # and sequence lengths — so the staged copies are byte-identical + # across layers. Stage once per step: the first generation-only + # call copies, subsequent layers with a matching key skip the 3 + # copy kernels. ``prepare()`` / ``update_for_spec_dec`` reset + # the key each step so eager forwards always re-stage; under CUDA + # graphs the first layer's captured copies replay once per step. + # The capture flag is part of the key: CUDA-graph capture is + # preceded by warmup forwards on the same metadata without an + # intervening prepare(), and the capture pass MUST re-record the + # staging copies (a skip would freeze stale page tables into the + # graph). + staging_key = ( + torch.cuda.is_current_stream_capturing(), + params.workspace.data_ptr(), + block_tables.data_ptr(), + tuple(block_tables.shape), + params.sequence_lengths.data_ptr(), + params.seq_offset, + batch_beam, + padded_num_pages, + ) + skip_staging_copy = ( + meta.num_contexts == 0 + and getattr(meta, "_cute_dsl_mla_staging_key", None) == staging_key + ) + workspace_buffer, block_tables, sequence_lengths = _prepare_cute_dsl_mla_buffers( + params.workspace, + block_tables, + params.sequence_lengths, + padded_num_pages, + skip_copy=skip_staging_copy, + ) + if meta.num_contexts == 0: + meta._cute_dsl_mla_staging_key = staging_key + uses_shared_paged_kv_idx = True + else: + sequence_lengths = params.sequence_lengths + workspace_buffer = params.workspace.view(-1, 4) + uses_shared_paged_kv_idx = self.USE_SHARED_PAGED_KV_IDX flashinfer.mla.trtllm_batch_decode_with_kv_cache_mla( query, # query @@ -1249,7 +1441,7 @@ def run_mla_generation( kv_lora_rank, # kv_lora_rank qk_rope_head_dim, # qk_rope_head_dim block_tables, # block_tables - params.sequence_lengths, # seq_lens + sequence_lengths, # seq_lens params.max_past_kv_length, # max_seq_len 0, # sparse_mla_top_k params.context_buf.view(batch_beam, q_len_per_req, attn.num_heads, kv_lora_rank), # out @@ -1258,8 +1450,15 @@ def run_mla_generation( fwd.attention_sinks, # sinks None, # skip_softmax_threshold_scale_factor self._enable_pdl, # enable_pdl - "trtllm-gen", # backend - True, # is_var_seq - self.USE_SHARED_PAGED_KV_IDX, # uses_shared_paged_kv_idx - multi_ctas_kv_counter_buffer=self._get_multi_ctas_kv_counter_buffer(), + backend=effective_mla_backend, + is_var_seq=True, + uses_shared_paged_kv_idx=uses_shared_paged_kv_idx, + cute_dsl_impl="monolithic", + # flashinfer rejects the counter buffer unless the trtllm-gen + # runner is selected; the cute-dsl MLA path must pass None. + multi_ctas_kv_counter_buffer=( + self._get_multi_ctas_kv_counter_buffer() + if effective_mla_backend != "cute-dsl" + else None + ), ) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 70f8394ecf7b..762b5a15ce0f 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -18,7 +18,7 @@ import os import weakref from dataclasses import dataclass, field -from typing import TYPE_CHECKING, List, Optional, Tuple +from typing import TYPE_CHECKING, Callable, List, Optional, Tuple import torch @@ -177,6 +177,16 @@ def effective_beam_width(self) -> int: init=False, repr=False) + # Per-forward-pass staging key for the CuTeDSL MLA generation workspace + # (page table + sequence lengths). All MLA layers of one step stage + # byte-identical data into the shared workspace, so the first layer + # copies and later layers skip. Reset whenever kv lens can change so + # eager forwards always re-stage (under CUDA graphs the first layer's + # captured copies replay once per step). + _cute_dsl_mla_staging_key: Optional[tuple] = field(default=None, + init=False, + repr=False) + use_paged_context_fmha: bool = field(init=False, default=False, repr=False) # FMHA prologue buffers for the MLA generation path, hoisted out of the per-layer @@ -551,6 +561,9 @@ def _invalidate_mla_scheduler_buffers(self) -> None: # buffers below must be rebuilt before the next MLA layer reads them. self._mla_scheduler_buffers_valid = False self._mla_ctx_cu_seqlens_valid = False + # The staged CuTe DSL page table and sequence lengths are derived from + # the same per-iteration scheduler state. + self._cute_dsl_mla_staging_key = None def update_helix_param( self, @@ -1414,6 +1427,7 @@ def __init__( skip_create_weights_in_init: bool = False, attention_chunk_size: Optional[int] = None, sparse_params: Optional[SparseParams] = None, + flashinfer_mla_backend: str = "trtllm-gen", **kwargs, ): """ @@ -1429,10 +1443,23 @@ def __init__( If None, positional embedding should be applied by the model before calling the backend. Otherwise, the backend is in-charge of applying positional embedding and may cache K without embedding it first. mla_params (MLAParams): Optional parameters for MLA. If None, MLA is not enabled. + flashinfer_mla_backend (str): FlashInfer MLA generation backend selected for + this attention instance. """ super().__init__(layer_idx, num_heads, head_dim, num_kv_heads, quant_config, **kwargs) self.sparse_params = sparse_params + self.flashinfer_mla_backend = flashinfer_mla_backend + # Per-batch MLA decode backend override hook. Maps (statically + # configured backend, batch metadata, generation-token count) to the + # backend name for the current batch. None (the default) keeps the + # static ``flashinfer_mla_backend`` selection unchanged. Model code + # that needs batch-dependent selection (e.g. Kimi K3's MLA module) + # installs a policy on the attention instances it owns; it lives on + # the backend object rather than the FMHA lib instances because + # ``create_fmha_libs`` may recreate those after model construction. + self.mla_backend_policy: Optional[Callable[ + [str, "TrtllmAttentionMetadata", int], str]] = None self.is_mla_enable = mla_params is not None self.mla_params = mla_params or MLAParams() diff --git a/tensorrt_llm/_torch/attention_backend/utils.py b/tensorrt_llm/_torch/attention_backend/utils.py index 7bf3c6c305e7..2a48c18ec54c 100644 --- a/tensorrt_llm/_torch/attention_backend/utils.py +++ b/tensorrt_llm/_torch/attention_backend/utils.py @@ -64,7 +64,8 @@ def create_attention( sparse_params: Optional[SparseParams] = None, dtype: Optional[torch.dtype] = None, aux_stream: Optional[torch.cuda.Stream] = None, -): + flashinfer_mla_backend: Optional[str] = None, +) -> AttentionBackend: if attention_chunk_size is not None and backend_name.upper() != "TRTLLM": raise ValueError( f"Backend {backend_name} does not support chunked attention.") @@ -99,6 +100,16 @@ def create_attention( aux_stream=aux_stream, sparse_params=sparse_params, ) + if flashinfer_mla_backend is not None: + # Only TrtllmAttention understands this selector. Raise instead of + # silently dropping it: a model that configured a specific MLA + # generation kernel must not run on another backend's default. + if not issubclass(attn_cls, TrtllmAttention): + raise ValueError( + f"flashinfer_mla_backend={flashinfer_mla_backend!r} is only " + "supported by the TRTLLM attention backend, but backend " + f"{backend_name} resolves to {attn_cls.__name__}.") + kwargs["flashinfer_mla_backend"] = flashinfer_mla_backend return attn_cls( layer_idx, diff --git a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md index 844a58fe1974..b0088bec1867 100644 --- a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md @@ -311,6 +311,24 @@ or `TLLM_FMHA_LIBS=-cute_dsl_mla,-msa_sparse_gqa,-flashinfer_trtllm_gen` to forc path. Each FMHA library exposes `is_available()` for module/static environment checks and `is_supported()` for per-forward request checks. +The `TrtllmAttention` constructor's `flashinfer_mla_backend` argument selects +the FlashInfer MLA generation kernel inside `FlashInferTrtllmGenFmha` for that +attention instance. It accepts `trtllm-gen` (default) or `cute-dsl`; the +latter uses the monolithic CuTeDSL decode implementation. Selecting `cute-dsl` +for an MLA layer using FP8 KV cache raises an exception because the current +CuTeDSL kernel does not accept the device-resident BMM scale tensors produced +for FP8 KV. `TrtllmAttention.mla_backend_policy` is an optional per-batch +override hook: model code may install a callable +`(static_backend, metadata, num_gen_tokens) -> backend` on an attention +instance to adjust the selection to the batch composition. + +Kimi K3 defaults its absorbed-generation MLA backend to `cute-dsl` for BF16 KV +cache (override with `TLLM_K3_MLA_GEN_BACKEND=trtllm-gen`; other values are +rejected at model build). FP8 KV cache forces `trtllm-gen`. K3 also installs a +per-batch policy that falls back to `trtllm-gen` for mixed +context/generation batches and multi-token generation (speculative +verification), keeping `cute-dsl` for plain one-token-per-request decode. + The FMHA package is split by role: - `fmha/interface.py` defines the `Fmha` runtime contract. diff --git a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py index f6aa9128eac8..c598fbb3931b 100644 --- a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py +++ b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py @@ -10,17 +10,83 @@ from __future__ import annotations +import os from typing import Optional import torch from ....functional import PositionEmbeddingType +from ....logger import logger +from ....models.modeling_utils import QuantConfig from ...attention_backend import AttentionMetadata, TrtllmAttention from ...attention_backend.interface import PositionalEmbeddingParams, RopeParams from ...model_config import ModelConfig from ..linear import Linear, TensorParallelMode from ..mla import MLA +_KIMI_K3_MLA_GEN_BACKEND_ENV = "TLLM_K3_MLA_GEN_BACKEND" +_KIMI_K3_MLA_GEN_BACKENDS = ("cute-dsl", "trtllm-gen") + + +def _select_mla_generation_backend(quant_config: Optional[QuantConfig]) -> str: + """Select K3's absorbed-generation MLA backend. + + K3 was tuned with the FlashInfer CuTe-DSL backend for BF16 KV cache. + FP8 KV cache carries device scales that CuTe-DSL does not support, so + retain the TRTLLM-Gen fallback used by the pre-refactor implementation. + """ + backend = os.environ.get(_KIMI_K3_MLA_GEN_BACKEND_ENV, "cute-dsl") + # Validate here, where the env var is read: an invalid value would + # otherwise surface only deep inside attention-backend construction, + # with an error that never names the knob that caused it. + if backend not in _KIMI_K3_MLA_GEN_BACKENDS: + raise ValueError( + f"{_KIMI_K3_MLA_GEN_BACKEND_ENV}={backend!r} is invalid; " + f"expected one of {list(_KIMI_K3_MLA_GEN_BACKENDS)}." + ) + has_fp8_kv_cache = bool( + quant_config is not None and quant_config.layer_quant_mode.has_fp8_kv_cache() + ) + if has_fp8_kv_cache and backend != "trtllm-gen": + # info_once: this runs once per MLA layer (~60x at startup for + # FP8-KV) and the decision is identical for every layer. + logger.info_once( + "Kimi K3 MLA: FP8 KV cache requires the trtllm-gen MLA " + f"generation backend; overriding '{backend}' -> 'trtllm-gen'.", + key="kimi_k3_mla_gen_backend_fp8_override", + ) + return "trtllm-gen" + return backend + + +def _kimi_k3_mla_decode_backend_policy( + requested_backend: str, + metadata, + num_gen_tokens: int, +) -> str: + """Per-batch MLA decode backend selection for Kimi K3. + + Installed as ``mla_backend_policy`` on K3's generation attention backend + (see :class:`KimiK3MLAAttention`); the general attention code applies no + such policy on its own. + + CuTe-DSL reuses one staged page table across MLA layers for a + generation-only, one-token-per-request batch. A mixed context/generation + batch cannot use that reuse key, so selecting CuTe-DSL would repeat the + staging copies in every MLA layer and regress time to first token. The + CuTe-DSL kernel itself accepts multi-token queries, but K3's decode + tuning covers only the one-token-per-request regime, so speculative + multi-token verification also falls back. Keep the requested CuTe-DSL + fast path for plain decode and use TRTLLM-Gen for mixed batches and + multi-token generation. + """ + is_single_token_generation = num_gen_tokens == metadata.num_generations + if requested_backend == "cute-dsl" and ( + metadata.num_contexts > 0 or not is_single_token_generation + ): + return "trtllm-gen" + return requested_backend + def _meta_safe_cast_dtype(module, dtype): """``module.to(dtype=dtype)`` that also works under ``MetaInitMode``. @@ -183,6 +249,7 @@ def __init__( reduce_output=False, fuse_qkv_a_proj=False, rms_norm_eps=rms_norm_eps, + flashinfer_mla_backend=_select_mla_generation_backend(model_config.get_quant_config()), ) # K3 calls forward_impl() directly to insert its output gate before # the base row-parallel o_proj. The original executor metadata remains @@ -215,6 +282,10 @@ def __init__( assert isinstance(self.mqa, TrtllmAttention) _install_identity_rope_table(self.mha) _install_identity_rope_table(self.mqa) + # Only the absorbed-generation backend (mqa) requests CuTe-DSL, so + # only it needs K3's per-batch fallback policy; mha keeps the + # default trtllm-gen selection. + self.mqa.mla_backend_policy = _kimi_k3_mla_decode_backend_policy self.rotary_emb = None self.apply_rotary_emb = False diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py index 646509bb25da..d008dfd63b36 100644 --- a/tensorrt_llm/_torch/modules/mla.py +++ b/tensorrt_llm/_torch/modules/mla.py @@ -229,7 +229,8 @@ def __init__( o_lora_rank: int = 1024, fuse_qkv_a_proj: bool = True, rms_norm_eps: Optional[float] = None, - ): + flashinfer_mla_backend: Optional[str] = None, + ) -> None: """ Initialize the MLA module. @@ -260,6 +261,9 @@ def __init__( rms_norm_eps (Optional[float]): Override the RMSNorm epsilon from the pretrained config. If neither source provides a value (e.g. config.pretrained_config is None), falls back to 1e-6. + flashinfer_mla_backend (Optional[str]): Generation backend for the + FlashInfer/TRTLLM-Gen MLA dispatcher. ``None`` preserves the + attention backend default. """ super().__init__() self.layer_idx = layer_idx @@ -557,6 +561,7 @@ def yarn_get_mscale(scale=1, mscale=1): dtype=dtype, aux_stream=mqa_aux_stream, rope_append=True, + flashinfer_mla_backend=flashinfer_mla_backend, ) if self.mqa is None: raise RuntimeError("MLA requires a non-null MQA attention backend") diff --git a/tests/unittest/_torch/attention/test_fmha_page_index.py b/tests/unittest/_torch/attention/test_fmha_page_index.py index 9ffabeb9ede1..062a0b573b1a 100644 --- a/tests/unittest/_torch/attention/test_fmha_page_index.py +++ b/tests/unittest/_torch/attention/test_fmha_page_index.py @@ -12,6 +12,22 @@ ) +class _AttentionStub: + def __init__( + self, + *, + is_mla_enable: bool, + has_fp8_kv_cache: bool, + flashinfer_mla_backend: str = "trtllm-gen", + ) -> None: + self.is_mla_enable = is_mla_enable + self.has_fp8_kv_cache = has_fp8_kv_cache + self.flashinfer_mla_backend = flashinfer_mla_backend + self.kv_lora_rank = 512 if is_mla_enable else None + self.head_dim = 576 + self.v_head_dim = 512 if is_mla_enable else None + + def _get_total_num_blocks(manager: SimpleNamespace, kv_factor: int = 2) -> int: fmha = object.__new__(FlashInferTrtllmGenFmha) fmha.kv_factor = kv_factor @@ -96,3 +112,81 @@ def check_counter_size_args( forward_args=SimpleNamespace(), workspace=SimpleNamespace(), ) + + +def test_flashinfer_cute_dsl_mla_backend_rejects_fp8_kv_cache() -> None: + attn = _AttentionStub( + is_mla_enable=True, + has_fp8_kv_cache=True, + flashinfer_mla_backend="cute-dsl", + ) + + with pytest.raises(ValueError, match="does not support FP8 KV cache"): + FlashInferTrtllmGenFmha(attn) + + +def test_flashinfer_mla_backend_rejects_unknown_backend() -> None: + attn = _AttentionStub( + is_mla_enable=True, + has_fp8_kv_cache=False, + flashinfer_mla_backend="cutedsl", + ) + + with pytest.raises(ValueError, match="flashinfer_mla_backend must be one of"): + FlashInferTrtllmGenFmha(attn) + + +def _make_fmha(requested_backend: str, mla_backend_policy) -> FlashInferTrtllmGenFmha: + fmha = object.__new__(FlashInferTrtllmGenFmha) + fmha._mla_backend = requested_backend + # ``Fmha.attn`` is a read-only property that dereferences ``_attn_ref`` + # (normally a weakref to the owning TrtllmAttention). SimpleNamespace is + # not weak-referenceable, so stand in with a closure of the same shape. + attn = SimpleNamespace(mla_backend_policy=mla_backend_policy) + fmha._attn_ref = lambda: attn + return fmha + + +@pytest.mark.parametrize("requested_backend", ["cute-dsl", "trtllm-gen"]) +@pytest.mark.parametrize( + ("num_contexts", "num_generations", "num_gen_tokens"), + [ + (0, 4, 4), # generation-only, one token per request + (1, 3, 3), # mixed context/generation batch + (0, 4, 8), # multi-token generation (speculative verification) + ], +) +def test_flashinfer_mla_backend_default_matches_static_selection( + requested_backend: str, + num_contexts: int, + num_generations: int, + num_gen_tokens: int, +) -> None: + """Without an installed policy the static backend is used for every batch + composition, matching the behavior before the policy hook existed.""" + fmha = _make_fmha(requested_backend, mla_backend_policy=None) + + assert ( + fmha._get_effective_mla_backend( + SimpleNamespace( + num_contexts=num_contexts, + num_generations=num_generations, + ), + num_gen_tokens, + ) + == requested_backend + ) + + +def test_flashinfer_mla_backend_policy_hook_is_consulted() -> None: + calls = [] + + def policy(requested_backend: str, meta, num_gen_tokens: int) -> str: + calls.append((requested_backend, meta, num_gen_tokens)) + return "trtllm-gen" + + fmha = _make_fmha("cute-dsl", mla_backend_policy=policy) + meta = SimpleNamespace(num_contexts=0, num_generations=4) + + assert fmha._get_effective_mla_backend(meta, 4) == "trtllm-gen" + assert calls == [("cute-dsl", meta, 4)] diff --git a/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py b/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py new file mode 100644 index 000000000000..0ca4240ab783 --- /dev/null +++ b/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from typing import Optional + +import pytest + +from tensorrt_llm._torch.modules.kimi_k3_mla.kimi_k3_mla_attention import ( + _KIMI_K3_MLA_GEN_BACKEND_ENV, + _kimi_k3_mla_decode_backend_policy, + _select_mla_generation_backend, +) +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + + +@pytest.mark.parametrize( + ("configured_backend", "expected_backend"), + [(None, "cute-dsl"), ("trtllm-gen", "trtllm-gen")], +) +def test_select_kimi_k3_mla_generation_backend( + monkeypatch: pytest.MonkeyPatch, + configured_backend: Optional[str], + expected_backend: str, +) -> None: + if configured_backend is None: + monkeypatch.delenv(_KIMI_K3_MLA_GEN_BACKEND_ENV, raising=False) + else: + monkeypatch.setenv(_KIMI_K3_MLA_GEN_BACKEND_ENV, configured_backend) + + assert _select_mla_generation_backend(None) == expected_backend + + +@pytest.mark.parametrize("invalid_backend", ["cutedsl", "", "CUTE-DSL "]) +def test_select_kimi_k3_mla_generation_backend_rejects_invalid_env( + monkeypatch: pytest.MonkeyPatch, + invalid_backend: str, +) -> None: + """An invalid env value must fail at read time with a message naming the + knob, not propagate until attention-backend construction.""" + monkeypatch.setenv(_KIMI_K3_MLA_GEN_BACKEND_ENV, invalid_backend) + + with pytest.raises(ValueError, match=_KIMI_K3_MLA_GEN_BACKEND_ENV): + _select_mla_generation_backend(None) + + +def test_select_kimi_k3_mla_generation_backend_uses_trtllm_gen_for_fp8_kv_cache( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv(_KIMI_K3_MLA_GEN_BACKEND_ENV, "cute-dsl") + quant_config = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) + + assert _select_mla_generation_backend(quant_config) == "trtllm-gen" + + +@pytest.mark.parametrize( + ( + "requested_backend", + "num_contexts", + "num_generations", + "num_gen_tokens", + "expected_backend", + ), + [ + ("cute-dsl", 0, 4, 4, "cute-dsl"), + ("cute-dsl", 1, 3, 3, "trtllm-gen"), + ("cute-dsl", 0, 4, 8, "trtllm-gen"), + ("trtllm-gen", 1, 3, 3, "trtllm-gen"), + ], +) +def test_kimi_k3_mla_decode_backend_policy_falls_back_for_mixed_batches( + requested_backend: str, + num_contexts: int, + num_generations: int, + num_gen_tokens: int, + expected_backend: str, +) -> None: + """K3's per-batch policy keeps CuTe-DSL only for generation-only, + one-token-per-request batches; mixed or multi-token batches fall back + to trtllm-gen.""" + assert ( + _kimi_k3_mla_decode_backend_policy( + requested_backend, + SimpleNamespace( + num_contexts=num_contexts, + num_generations=num_generations, + ), + num_gen_tokens, + ) + == expected_backend + ) From 311f1c766e84da8f2d7fc5e42ca2de83f2acca08 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Wed, 19 Aug 2026 16:13:32 -0700 Subject: [PATCH 2/5] [TRTLLM-15033][fix] address Kimi K3 MLA review feedback Signed-off-by: Simeng Liu --- .../attention_backend/fmha/cute_dsl_mla.py | 19 ++++-- .../fmha/flashinfer_trtllm_gen.py | 23 ++++--- .../attention_backend/fmha/interface.py | 48 ++++++++++++- .../_torch/attention_backend/trtllm.py | 25 ++++--- .../modules/ATTENTION_DEVELOPER_GUIDE.md | 25 ++++--- .../kimi_k3_mla/kimi_k3_mla_attention.py | 36 ++++++---- .../_torch/attention/test_fmha_page_index.py | 67 +++++++++++++++++-- .../modules/test_kimi_k3_mla_backend.py | 18 ++--- 8 files changed, 201 insertions(+), 60 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py index 83341117796b..c4c305e7255d 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py @@ -44,6 +44,14 @@ class CuteDslMlaFmha(PhasedFmha): @classmethod def is_available(cls, attn: "TrtllmAttention") -> bool: + if attn.flashinfer_mla_backend is not None: + logger.debug( + "Standalone CuTe DSL MLA FMHA is unavailable: an explicit " + "flashinfer_mla_backend delegates MLA generation to " + "FlashInferTrtllmGenFmha." + ) + return False + if not IS_CUTLASS_DSL_AVAILABLE: logger.debug("CuTe DSL MLA FMHA is unavailable: nvidia-cutlass-dsl is not installed.") return False @@ -242,7 +250,9 @@ def _is_perf_favorable( (16, 2): 64, (16, 4): 32, (16, 8): 16, - (96, 1): 1, # trtllm-gen may fail under num_head == 96 + # Keep H=96 off TRTLLM-Gen: its heuristic may select a 64-head + # Q tile, which does not divide 96 and produces an invalid config. + (96, 1): 1, (128, 1): 64, (128, 2): 32, } @@ -271,9 +281,10 @@ def _is_supported_with_reason( ) -> tuple[bool, str]: if fwd.attention_input_type != AttentionInputType.generation_only: return False, "CuTe DSL MLA FMHA only supports generation-only attention." - # It is to disable mix-batch(context request + generation request) for now. - # Trtllm-gen may fail under num_head == 96 so it is exempted from this check. - # TODO: Eliminate high host overhead of cutedsl mla to enable mix-batch. + # Disable mixed context/generation batches until the CuTe DSL host + # overhead is reduced. H=96 is exempt: TRTLLM-Gen may select a + # 64-head Q tile, which does not divide 96 and produces an invalid + # configuration after K3's 96-to-128 head padding was removed. if (meta.num_contexts != 0 and attn.num_heads != 96) or meta.num_generations <= 0: return False, "CuTe DSL MLA FMHA only supports decode-only batches." if meta.helix_position_offsets is not None: diff --git a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py index 4566c15d7fe6..736dea282e06 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py @@ -58,6 +58,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantMode +from .interface import _CuteDslMlaStagingKey from .phased import FmhaParams, PhasedFmha if TYPE_CHECKING: @@ -502,10 +503,10 @@ class FlashInferTrtllmGenFmha(PhasedFmha): (576, 512), } - def __init__(self, attn: "TrtllmAttention"): + def __init__(self, attn: "TrtllmAttention") -> None: super().__init__(attn) self._layout = self.DEFAULT_KV_LAYOUT - requested_mla_backend = _get_mla_backend(attn.flashinfer_mla_backend) + requested_mla_backend = _get_mla_backend(attn.flashinfer_mla_backend or "trtllm-gen") if requested_mla_backend == "cute-dsl" and attn.is_mla_enable and attn.has_fp8_kv_cache: raise ValueError( "flashinfer_mla_backend='cute-dsl' does not support FP8 KV cache device scales; " @@ -1404,15 +1405,15 @@ def run_mla_generation( # intervening prepare(), and the capture pass MUST re-record the # staging copies (a skip would freeze stale page tables into the # graph). - staging_key = ( - torch.cuda.is_current_stream_capturing(), - params.workspace.data_ptr(), - block_tables.data_ptr(), - tuple(block_tables.shape), - params.sequence_lengths.data_ptr(), - params.seq_offset, - batch_beam, - padded_num_pages, + staging_key = _CuteDslMlaStagingKey( + is_capturing=torch.cuda.is_current_stream_capturing(), + workspace_ptr=params.workspace.data_ptr(), + block_tables_ptr=block_tables.data_ptr(), + block_tables_shape=tuple(block_tables.shape), + sequence_lengths_ptr=params.sequence_lengths.data_ptr(), + sequence_lengths_offset=params.seq_offset, + batch_beam=batch_beam, + padded_num_pages=padded_num_pages, ) skip_staging_copy = ( meta.num_contexts == 0 diff --git a/tensorrt_llm/_torch/attention_backend/fmha/interface.py b/tensorrt_llm/_torch/attention_backend/fmha/interface.py index f14f26964b25..a74fd4f36c58 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/interface.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/interface.py @@ -15,7 +15,7 @@ import weakref from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, NamedTuple, Optional, Protocol import torch @@ -28,6 +28,52 @@ ) +class _CuteDslMlaStagingKey(NamedTuple): + """Identifies CuTe-DSL MLA inputs staged into a shared workspace. + + Attributes: + is_capturing: Whether the staging occurred during CUDA graph capture. + workspace_ptr: Address of the shared staging workspace. + block_tables_ptr: Address of the source block tables. + block_tables_shape: Shape of the source block tables. + sequence_lengths_ptr: Address of the source sequence lengths. + sequence_lengths_offset: Offset applied to the source sequence lengths. + batch_beam: Number of generation sequences, including beam expansion. + padded_num_pages: Page-table width after CuTe-DSL alignment padding. + """ + + is_capturing: bool + workspace_ptr: int + block_tables_ptr: int + block_tables_shape: tuple[int, ...] + sequence_lengths_ptr: int + sequence_lengths_offset: int + batch_beam: int + padded_num_pages: int + + +class MlaBackendPolicy(Protocol): + """Selects the MLA generation backend for one scheduler batch.""" + + def __call__( + self, + requested_backend: str, + metadata: "TrtllmAttentionMetadata", + num_gen_tokens: int, + ) -> str: + """Return the backend to use for the supplied batch composition. + + Args: + requested_backend: Backend selected by the attention instance. + metadata: Runtime metadata for the current scheduler batch. + num_gen_tokens: Number of generation tokens in the batch. + + Returns: + Backend name to use for MLA generation in this batch. + """ + ... + + class Fmha(ABC): """Common runtime contract for TRT-LLM attention FMHA libraries.""" diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 762b5a15ce0f..bac23d041166 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -18,7 +18,7 @@ import os import weakref from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Callable, List, Optional, Tuple +from typing import TYPE_CHECKING, List, Optional, Tuple import torch @@ -31,6 +31,8 @@ from tensorrt_llm._torch.attention_backend.fmha import ( Fmha, get_enabled_fmha_lib_classes) +from tensorrt_llm._torch.attention_backend.fmha.interface import ( + MlaBackendPolicy, _CuteDslMlaStagingKey) from tensorrt_llm._utils import get_sm_version, maybe_pin_memory, prefer_pinned from tensorrt_llm.bindings.internal import thop from tensorrt_llm.functional import AttentionMaskType @@ -183,9 +185,11 @@ def effective_beam_width(self) -> int: # copies and later layers skip. Reset whenever kv lens can change so # eager forwards always re-stage (under CUDA graphs the first layer's # captured copies replay once per step). - _cute_dsl_mla_staging_key: Optional[tuple] = field(default=None, - init=False, - repr=False) + _cute_dsl_mla_staging_key: Optional[_CuteDslMlaStagingKey] = field( + default=None, + init=False, + repr=False, + ) use_paged_context_fmha: bool = field(init=False, default=False, repr=False) @@ -1427,9 +1431,9 @@ def __init__( skip_create_weights_in_init: bool = False, attention_chunk_size: Optional[int] = None, sparse_params: Optional[SparseParams] = None, - flashinfer_mla_backend: str = "trtllm-gen", + flashinfer_mla_backend: Optional[str] = None, **kwargs, - ): + ) -> None: """ Initialize the backend. Args: @@ -1443,8 +1447,10 @@ def __init__( If None, positional embedding should be applied by the model before calling the backend. Otherwise, the backend is in-charge of applying positional embedding and may cache K without embedding it first. mla_params (MLAParams): Optional parameters for MLA. If None, MLA is not enabled. - flashinfer_mla_backend (str): FlashInfer MLA generation backend selected for - this attention instance. + flashinfer_mla_backend (Optional[str]): FlashInfer MLA generation backend + selected for this attention instance. + None preserves the ordered FMHA-library + dispatch. """ super().__init__(layer_idx, num_heads, head_dim, num_kv_heads, quant_config, **kwargs) @@ -1458,8 +1464,7 @@ def __init__( # installs a policy on the attention instances it owns; it lives on # the backend object rather than the FMHA lib instances because # ``create_fmha_libs`` may recreate those after model construction. - self.mla_backend_policy: Optional[Callable[ - [str, "TrtllmAttentionMetadata", int], str]] = None + self.mla_backend_policy: Optional[MlaBackendPolicy] = None self.is_mla_enable = mla_params is not None self.mla_params = mla_params or MLAParams() diff --git a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md index b0088bec1867..062bf1fb312d 100644 --- a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md @@ -311,14 +311,19 @@ or `TLLM_FMHA_LIBS=-cute_dsl_mla,-msa_sparse_gqa,-flashinfer_trtllm_gen` to forc path. Each FMHA library exposes `is_available()` for module/static environment checks and `is_supported()` for per-forward request checks. -The `TrtllmAttention` constructor's `flashinfer_mla_backend` argument selects -the FlashInfer MLA generation kernel inside `FlashInferTrtllmGenFmha` for that -attention instance. It accepts `trtllm-gen` (default) or `cute-dsl`; the -latter uses the monolithic CuTeDSL decode implementation. Selecting `cute-dsl` -for an MLA layer using FP8 KV cache raises an exception because the current -CuTeDSL kernel does not accept the device-resident BMM scale tensors produced -for FP8 KV. `TrtllmAttention.mla_backend_policy` is an optional per-batch -override hook: model code may install a callable +The `TrtllmAttention` constructor's optional `flashinfer_mla_backend` argument +explicitly selects the MLA generation kernel inside +`FlashInferTrtllmGenFmha` for that attention instance. It accepts +`trtllm-gen` or `cute-dsl`; the latter uses the monolithic CuTeDSL decode +implementation. When the argument is `None`, the ordered FMHA-library +dispatch is preserved and FlashInfer uses `trtllm-gen` if reached. When it is +set, the standalone `CuteDslMlaFmha` defers to the explicit FlashInfer +selection. Selecting `cute-dsl` for an MLA layer using FP8 KV cache raises an +exception because the current CuTeDSL kernel does not accept the +device-resident BMM scale tensors produced for FP8 KV. + +`TrtllmAttention.mla_backend_policy` is an optional per-batch override hook: +model code may install a callable `(static_backend, metadata, num_gen_tokens) -> backend` on an attention instance to adjust the selection to the batch composition. @@ -327,7 +332,9 @@ cache (override with `TLLM_K3_MLA_GEN_BACKEND=trtllm-gen`; other values are rejected at model build). FP8 KV cache forces `trtllm-gen`. K3 also installs a per-batch policy that falls back to `trtllm-gen` for mixed context/generation batches and multi-token generation (speculative -verification), keeping `cute-dsl` for plain one-token-per-request decode. +verification), keeping `cute-dsl` for plain one-token-per-request decode. A +mixed H=96 batch remains on `cute-dsl`: TRTLLM-Gen may select a 64-head Q tile, +which does not divide 96 after K3's head padding removal. The FMHA package is split by role: diff --git a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py index c598fbb3931b..db0d816c42df 100644 --- a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py +++ b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py @@ -11,6 +11,7 @@ from __future__ import annotations import os +from functools import partial from typing import Optional import torch @@ -18,7 +19,7 @@ from ....functional import PositionEmbeddingType from ....logger import logger from ....models.modeling_utils import QuantConfig -from ...attention_backend import AttentionMetadata, TrtllmAttention +from ...attention_backend import AttentionMetadata, TrtllmAttention, TrtllmAttentionMetadata from ...attention_backend.interface import PositionalEmbeddingParams, RopeParams from ...model_config import ModelConfig from ..linear import Linear, TensorParallelMode @@ -61,8 +62,10 @@ def _select_mla_generation_backend(quant_config: Optional[QuantConfig]) -> str: def _kimi_k3_mla_decode_backend_policy( requested_backend: str, - metadata, + metadata: TrtllmAttentionMetadata, num_gen_tokens: int, + *, + num_heads: int, ) -> str: """Per-batch MLA decode backend selection for Kimi K3. @@ -71,18 +74,22 @@ def _kimi_k3_mla_decode_backend_policy( such policy on its own. CuTe-DSL reuses one staged page table across MLA layers for a - generation-only, one-token-per-request batch. A mixed context/generation - batch cannot use that reuse key, so selecting CuTe-DSL would repeat the - staging copies in every MLA layer and regress time to first token. The - CuTe-DSL kernel itself accepts multi-token queries, but K3's decode - tuning covers only the one-token-per-request regime, so speculative - multi-token verification also falls back. Keep the requested CuTe-DSL - fast path for plain decode and use TRTLLM-Gen for mixed batches and - multi-token generation. + generation-only, one-token-per-request batch. Other mixed batches repeat + the staging copies in every MLA layer and regress time to first token, so + they fall back to TRTLLM-Gen. The H=96 path is the correctness exception: + TRTLLM-Gen may select a 64-head Q tile, which does not divide 96 and + produces an invalid configuration after K3's head padding was removed. + + The CuTe-DSL kernel itself accepts multi-token queries, but K3's decode + tuning covers only the one-token-per-request regime, so generation-only + speculative verification also falls back. """ is_single_token_generation = num_gen_tokens == metadata.num_generations - if requested_backend == "cute-dsl" and ( - metadata.num_contexts > 0 or not is_single_token_generation + requires_cute_dsl_for_mixed_batch = metadata.num_contexts > 0 and num_heads == 96 + if ( + requested_backend == "cute-dsl" + and not requires_cute_dsl_for_mixed_batch + and (metadata.num_contexts > 0 or not is_single_token_generation) ): return "trtllm-gen" return requested_backend @@ -285,7 +292,10 @@ def __init__( # Only the absorbed-generation backend (mqa) requests CuTe-DSL, so # only it needs K3's per-batch fallback policy; mha keeps the # default trtllm-gen selection. - self.mqa.mla_backend_policy = _kimi_k3_mla_decode_backend_policy + self.mqa.mla_backend_policy = partial( + _kimi_k3_mla_decode_backend_policy, + num_heads=self.mqa.num_heads, + ) self.rotary_emb = None self.apply_rotary_emb = False diff --git a/tests/unittest/_torch/attention/test_fmha_page_index.py b/tests/unittest/_torch/attention/test_fmha_page_index.py index 062a0b573b1a..3f18a0b30e84 100644 --- a/tests/unittest/_torch/attention/test_fmha_page_index.py +++ b/tests/unittest/_torch/attention/test_fmha_page_index.py @@ -1,15 +1,20 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from collections.abc import Callable from types import SimpleNamespace +from typing import TypeAlias import pytest import torch +from tensorrt_llm._torch.attention_backend.fmha.cute_dsl_mla import CuteDslMlaFmha from tensorrt_llm._torch.attention_backend.fmha.flashinfer_trtllm_gen import ( FlashInferTrtllmGenFmha, _get_multi_ctas_kv_counter_size, ) +from tensorrt_llm._torch.attention_backend.fmha.interface import _CuteDslMlaStagingKey +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata class _AttentionStub: @@ -18,7 +23,7 @@ def __init__( *, is_mla_enable: bool, has_fp8_kv_cache: bool, - flashinfer_mla_backend: str = "trtllm-gen", + flashinfer_mla_backend: str | None = None, ) -> None: self.is_mla_enable = is_mla_enable self.has_fp8_kv_cache = has_fp8_kv_cache @@ -28,6 +33,9 @@ def __init__( self.v_head_dim = 512 if is_mla_enable else None +_MlaBackendPolicy: TypeAlias = Callable[[str, SimpleNamespace, int], str] + + def _get_total_num_blocks(manager: SimpleNamespace, kv_factor: int = 2) -> int: fmha = object.__new__(FlashInferTrtllmGenFmha) fmha.kv_factor = kv_factor @@ -125,6 +133,50 @@ def test_flashinfer_cute_dsl_mla_backend_rejects_fp8_kv_cache() -> None: FlashInferTrtllmGenFmha(attn) +@pytest.mark.parametrize("configured_backend", ["cute-dsl", "trtllm-gen"]) +def test_standalone_cute_dsl_mla_defers_to_explicit_flashinfer_backend( + configured_backend: str, +) -> None: + attn = _AttentionStub( + is_mla_enable=True, + has_fp8_kv_cache=False, + flashinfer_mla_backend=configured_backend, + ) + + assert not CuteDslMlaFmha.is_available(attn) + + +def test_flashinfer_mla_backend_defaults_to_trtllm_gen() -> None: + attn = _AttentionStub( + is_mla_enable=True, + has_fp8_kv_cache=False, + ) + + assert FlashInferTrtllmGenFmha(attn)._mla_backend == "trtllm-gen" + + +def test_mla_scheduler_invalidation_resets_cute_dsl_staging_key() -> None: + metadata = object.__new__(TrtllmAttentionMetadata) + metadata._mla_scheduler_buffers_valid = True + metadata._mla_ctx_cu_seqlens_valid = True + metadata._cute_dsl_mla_staging_key = _CuteDslMlaStagingKey( + is_capturing=True, + workspace_ptr=1, + block_tables_ptr=2, + block_tables_shape=(3, 4), + sequence_lengths_ptr=5, + sequence_lengths_offset=6, + batch_beam=7, + padded_num_pages=8, + ) + + metadata._invalidate_mla_scheduler_buffers() + + assert not metadata._mla_scheduler_buffers_valid + assert not metadata._mla_ctx_cu_seqlens_valid + assert metadata._cute_dsl_mla_staging_key is None + + def test_flashinfer_mla_backend_rejects_unknown_backend() -> None: attn = _AttentionStub( is_mla_enable=True, @@ -136,7 +188,10 @@ def test_flashinfer_mla_backend_rejects_unknown_backend() -> None: FlashInferTrtllmGenFmha(attn) -def _make_fmha(requested_backend: str, mla_backend_policy) -> FlashInferTrtllmGenFmha: +def _make_fmha( + requested_backend: str, + mla_backend_policy: _MlaBackendPolicy | None, +) -> FlashInferTrtllmGenFmha: fmha = object.__new__(FlashInferTrtllmGenFmha) fmha._mla_backend = requested_backend # ``Fmha.attn`` is a read-only property that dereferences ``_attn_ref`` @@ -179,9 +234,13 @@ def test_flashinfer_mla_backend_default_matches_static_selection( def test_flashinfer_mla_backend_policy_hook_is_consulted() -> None: - calls = [] + calls: list[tuple[str, SimpleNamespace, int]] = [] - def policy(requested_backend: str, meta, num_gen_tokens: int) -> str: + def policy( + requested_backend: str, + meta: SimpleNamespace, + num_gen_tokens: int, + ) -> str: calls.append((requested_backend, meta, num_gen_tokens)) return "trtllm-gen" diff --git a/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py b/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py index 0ca4240ab783..937daeabb2f9 100644 --- a/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py +++ b/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py @@ -59,25 +59,26 @@ def test_select_kimi_k3_mla_generation_backend_uses_trtllm_gen_for_fp8_kv_cache( "num_contexts", "num_generations", "num_gen_tokens", + "num_heads", "expected_backend", ), [ - ("cute-dsl", 0, 4, 4, "cute-dsl"), - ("cute-dsl", 1, 3, 3, "trtllm-gen"), - ("cute-dsl", 0, 4, 8, "trtllm-gen"), - ("trtllm-gen", 1, 3, 3, "trtllm-gen"), + ("cute-dsl", 0, 4, 4, 96, "cute-dsl"), + ("cute-dsl", 1, 3, 3, 12, "trtllm-gen"), + ("cute-dsl", 1, 3, 3, 96, "cute-dsl"), + ("cute-dsl", 0, 4, 8, 96, "trtllm-gen"), + ("trtllm-gen", 1, 3, 3, 96, "trtllm-gen"), ], ) -def test_kimi_k3_mla_decode_backend_policy_falls_back_for_mixed_batches( +def test_kimi_k3_mla_decode_backend_policy_by_batch_shape( requested_backend: str, num_contexts: int, num_generations: int, num_gen_tokens: int, + num_heads: int, expected_backend: str, ) -> None: - """K3's per-batch policy keeps CuTe-DSL only for generation-only, - one-token-per-request batches; mixed or multi-token batches fall back - to trtllm-gen.""" + """K3 falls back outside plain decode except for unsafe H=96 mixed batches.""" assert ( _kimi_k3_mla_decode_backend_policy( requested_backend, @@ -86,6 +87,7 @@ def test_kimi_k3_mla_decode_backend_policy_falls_back_for_mixed_batches( num_generations=num_generations, ), num_gen_tokens, + num_heads=num_heads, ) == expected_backend ) From cbd5171b7eea56710c9675cf238e83404a2229dd Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Thu, 20 Aug 2026 00:16:21 +0000 Subject: [PATCH 3/5] Address trivial review comments Signed-off-by: Brian Nguyen --- .../_torch/attention_backend/fmha/flashinfer_trtllm_gen.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py index 736dea282e06..14840b595565 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py @@ -300,6 +300,11 @@ def _prepare_cute_dsl_mla_buffers( _get_cute_dsl_mla_buffer_layout(batch_size, padded_num_pages) ) workspace_bytes = workspace.view(torch.uint8).flatten() + if workspace_bytes.numel() <= kernel_workspace_offset: + raise RuntimeError( + f"CuTeDSL MLA workspace has {workspace_bytes.numel()} bytes; " + f"staging metadata alone needs {kernel_workspace_offset} bytes." + ) page_table_storage = ( workspace_bytes[:page_table_bytes].view(torch.int32).view(batch_size, padded_num_pages) ) @@ -970,7 +975,7 @@ def prepare_workspace( if metadata.is_cuda_graph and torch.cuda.is_current_stream_capturing(): raise RuntimeError( "Attention CUDA graph workspace is smaller than the " - f"required size for {self._mla_backend}." + f"required size for {effective_mla_backend}." ) required_workspace_numel = math.ceil(required_workspace_size / workspace.element_size()) workspace.resize_((required_workspace_numel,)) From 3efdf8c0fb6a1291344935c78bb6f2427f12f19b Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Thu, 20 Aug 2026 06:07:06 +0000 Subject: [PATCH 4/5] Address trivial review comments Signed-off-by: Brian Nguyen --- tests/integration/defs/sysinfo/get_sysinfo.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/integration/defs/sysinfo/get_sysinfo.py b/tests/integration/defs/sysinfo/get_sysinfo.py index b637b44e6112..a421d1ed4ecc 100644 --- a/tests/integration/defs/sysinfo/get_sysinfo.py +++ b/tests/integration/defs/sysinfo/get_sysinfo.py @@ -111,7 +111,13 @@ def is_power(): # TODO(#17993): cherry-picked from PR #17993 to unblock this PR's CI (distro->'na' # empty-render bug). Drop this and take main's version when resolving the rebase # conflict after #17993 lands. -def get_linux_distribution(): +def get_linux_distribution() -> tuple[str, str, str]: + """Return Linux distribution ID, version, and codename. + + Returns: + A tuple containing the distribution ID, version, and codename. + Returns ``("na", "na", "na")`` when metadata is unavailable. + """ try: import distro return (distro.id(), distro.version(), distro.codename()) From a992880ef2ac373bdbbf965612b063e52ea2e058 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Fri, 21 Aug 2026 13:03:56 -0700 Subject: [PATCH 5/5] [TRTLLM-15033][fix] address remaining review feedback Signed-off-by: Simeng Liu --- .../attention_backend/sparse/dsa/metadata.py | 4 +++- tests/integration/defs/sysinfo/get_sysinfo.py | 2 +- .../sparse/dsa/test_req_idx_per_token.py | 24 ++++++++++++++++++- 3 files changed, 27 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 2854c71b6f96..06ba0bc12f56 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -336,7 +336,7 @@ def warmup_cute_dsl_radix_topk(self, next_n: int) -> None: num_sms=self.num_sms, ) - def on_update_kv_lens(self): + def on_update_kv_lens(self) -> None: # After changing the kv_lens/kv_lens_cuda, we may need to update other metadatas. # Especially for the changes in the _preprocess_inputs() of model_engine.py. # @@ -346,6 +346,8 @@ def on_update_kv_lens(self): # slot_mapping_* buffers also depend on these effective cached lengths. If we do not # refresh slot mappings here, indexer K-cache updates can be written with stale offsets. + super().on_update_kv_lens() + # _preprocess_inputs() also uses this as a general hook to "invalidate per-forward-pass # caches so they are recomputed (and captured) on every _forward_step". Invalidate the # pool_view cache here so it is recomputed on the next diff --git a/tests/integration/defs/sysinfo/get_sysinfo.py b/tests/integration/defs/sysinfo/get_sysinfo.py index a421d1ed4ecc..e3659b6f1c13 100644 --- a/tests/integration/defs/sysinfo/get_sysinfo.py +++ b/tests/integration/defs/sysinfo/get_sysinfo.py @@ -140,7 +140,7 @@ def get_linux_distribution() -> tuple[str, str, str]: except OSError: logger.error( f"Cannot determine the Linux distribution ({distro_reason}; " - "/etc/os-release also unreadable); reporting ('na', 'na', 'na'). " + "os-release metadata unavailable); reporting ('na', 'na', 'na'). " "Test-db conditions matching linux_distribution_name (e.g. ubuntu*) " "will select ZERO tests and the rendered test list will be empty.") return ("na", "na", "na") diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py b/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py index 58448da45733..6b5b43c27787 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py @@ -20,6 +20,7 @@ 2. on_update_kv_lens() must rebuild the map after the MTP draft loop rewrites seq_lens, or every draft token is misattributed to request 0 (https://nvbugs/6513132, https://nvbugs/6513093). +3. on_update_kv_lens() must preserve the base metadata invalidation contract. """ from unittest.mock import Mock @@ -64,7 +65,28 @@ def test_matches_host_repeat_interleave(seq_lens, device): assert result.to(torch.int32).tolist() == _host_reference(seq_lens).tolist() -def test_on_update_kv_lens_rebuilds_stale_map(): +def test_on_update_kv_lens_invalidates_base_mla_state() -> None: + md = object.__new__(DSAtrtllmAttentionMetadata) + md.enable_flash_mla = False + md._mla_scheduler_buffers_valid = True + md._mla_ctx_cu_seqlens_valid = True + md._cute_dsl_mla_staging_key = object() + md._invalidate_pool_view_cache = Mock() + md._num_tokens = 0 + md._num_generations = 0 + md.kv_cache_manager = None + md.kv_lens_cuda = torch.empty(0, dtype=torch.int32) + md._compute_kv_lens_row_reorder = Mock() + md.prepare_dense_topk_indices = Mock() + + md.on_update_kv_lens() + + assert not md._mla_scheduler_buffers_valid + assert not md._mla_ctx_cu_seqlens_valid + assert md._cute_dsl_mla_staging_key is None + + +def test_on_update_kv_lens_rebuilds_stale_map() -> None: """on_update_kv_lens() must replace prepare()'s stale map (fails pre-fix).""" if not torch.cuda.is_available(): pytest.skip("CUDA not available")