[TRTLLM-15033][feat] Upstream Kimi K3 MLA decode backend selection to main

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:

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:
@@ -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:
@@ -207,6 +208,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 +243,89 @@ 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()
+    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)
+    )
+    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 +454,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
@@ -412,9 +508,16 @@ 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 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; "
+                "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 +943,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 {effective_mla_backend}."
                 )
             required_workspace_numel = math.ceil(required_workspace_size / workspace.element_size())
             workspace.resize_((required_workspace_numel,))
@@ -856,6 +986,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 +1320,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 +1392,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 = _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
+                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 +1447,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 +1456,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
+            ),
         )