[None][feat] Support MLA generation in TrtllmGen attention backend

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -29,9 +29,9 @@
                         RopeParams)
 from .trtllm_gen import trtllm_gen_attention
 
-# Enable TRTLLM-Gen attention backend via environment variable.
-_TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION = os.environ.get(
-    "TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION", "0") == "1"
+# Enable TRTLLM-Gen attention backend via environment variable (default: on).
+_TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION = (os.environ.get(
+    "TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION", "1") == "1")
 
 
 @dataclass(kw_only=True, init=False)

tensorrt_llm/_torch/attention_backend/trtllm_gen.py

@@ -159,10 +159,8 @@ def is_supported(
         has_sparse_attn = sparse_attn_indices is not None and sparse_attn_indices.numel() > 0
         if has_sparse_kv or has_sparse_attn:
             return False, "Sparse attention is not supported by trtllm-gen backend."
-        if is_mla_enable:
-            return False, "MLA is not supported by trtllm-gen backend."
-        if not is_fused_qkv:
-            return False, "Only fused QKV is supported by trtllm-gen backend."
+        if is_mla_enable and not is_fused_qkv:
+            return False, "MLA context (separate Q/K/V) falls back to thop."
         if not update_kv_cache:
             return False, "KV cache update cannot be disabled now."
         if cross_attention:
@@ -194,7 +192,6 @@ def is_supported(
             )
 
         o_dtype = out_dtype if out_dtype is not None else q_dtype
-        heads_ratio = num_heads // num_kv_heads if num_kv_heads > 0 else 1
 
         if phase in ("context", "both"):
             if head_size in cls.UNSUPPORTED_HEAD_SIZES_CONTEXT:
@@ -231,12 +228,11 @@ def is_supported(
                     False,
                     f"[Generation] tokens_per_block ({tokens_per_block}) must be >= {cls.MIN_TOKENS_PER_BLOCK}.",
                 )
-            if heads_ratio > cls.MAX_HEADS_RATIO_GENERATION:
-                max_ratio = cls.MAX_HEADS_RATIO_GENERATION
+            heads_ratio = num_heads // num_kv_heads
+            if not is_mla_enable and heads_ratio > cls.MAX_HEADS_RATIO_GENERATION:
                 return (
                     False,
-                    f"[Generation] num_heads/num_kv_heads ratio "
-                    f"({heads_ratio}) must be <= {max_ratio}.",
+                    f"[Generation] heads ratio ({heads_ratio}) exceeds maximum ({cls.MAX_HEADS_RATIO_GENERATION}).",
                 )
             if has_alibi:
                 return False, "[Generation] ALiBi is not supported."
@@ -261,7 +257,7 @@ def is_supported(
         return True, ""
 
 
-@dataclass
+@dataclass(slots=True)
 class ContextWorkspaceBuffers:
     """
     Workspace buffers for context phase.
@@ -292,7 +288,7 @@ class ContextWorkspaceBuffers:
     fmha_bmm2_scale: Optional[torch.Tensor] = None
 
 
-@dataclass
+@dataclass(slots=True)
 class GenerationWorkspaceBuffers:
     """
     Workspace buffers for generation phase.
@@ -855,6 +851,14 @@ class EnqueueParams:
     paged_context_fmha: bool = False
     attention_sinks: Optional[torch.Tensor] = None
     is_mla_enable: bool = False
+    # MLA parameters
+    kv_lora_rank: int = 0
+    qk_nope_head_dim: int = 0
+    qk_rope_head_dim: int = 0
+    v_head_dim: int = 0
+    q_scaling: float = 1.0
+    latent_cache: Optional[torch.Tensor] = None
+    num_layers: int = 0
 
 
 @dataclass
@@ -1361,6 +1365,77 @@ def run_generation(self, params: EnqueueGenerationParams):
                 q_len_per_req=params.input_seq_length,
             )
 
+    def run_mla_generation(self, params: EnqueueGenerationParams) -> None:
+        """MLA generation decode using flashinfer MLA kernel."""
+        if 0 < params.cyclic_attention_window_size < params.max_past_kv_length:
+            raise NotImplementedError(
+                "Sliding-window attention is not supported by MLA decode path."
+            )
+        if params.attention_chunk_size != 0:
+            raise NotImplementedError("Chunked-attention is not supported by MLA decode path.")
+
+        batch_beam = params.num_requests * params.beam_width
+        block_tables = self._build_block_tables(
+            params.kv_cache_block_offsets,
+            params.host_kv_cache_pool_mapping,
+            params.layer_idx,
+            params.global_layer_idx,
+            params.seq_offset,
+            batch_beam,
+        )
+
+        pages_per_superblock = 128 // params.tokens_per_block
+        if pages_per_superblock > 1:
+            num_blocks = block_tables.size(-1)
+            remainder = num_blocks % pages_per_superblock
+            if remainder != 0:
+                pad = pages_per_superblock - remainder
+                block_tables = torch.nn.functional.pad(block_tables, (0, pad), value=0)
+
+        mla_head_dim_qk = params.kv_lora_rank + params.qk_rope_head_dim
+        q_len_per_req = params.num_tokens // batch_beam if batch_beam > 0 else 1
+
+        query = params.qkv_input.view(batch_beam, q_len_per_req, params.num_heads, mla_head_dim_qk)
+
+        kv_cache = self._kv_cache_manager.get_buffers(
+            params.global_layer_idx, kv_layout=DEFAULT_KV_LAYOUT
+        )
+        if kv_cache.ndim == 5:
+            kv_cache = kv_cache.squeeze(2)
+
+        bmm1_scale = 1.0 / (
+            params.q_scaling * math.sqrt(params.qk_nope_head_dim + params.qk_rope_head_dim)
+        )
+        bmm2_scale = 1.0
+
+        # context_buf: [B*q_len, H, kv_lora_rank]
+        # flashinfer MLA decode out check only accepts [B, H, D] (3D),
+        # but kernel outputs [B, q_len, H, D] when q_len > 1.
+        # For q_len=1: pass context_buf directly (in-place, zero copy).
+        # For q_len>1: let flashinfer allocate, then copy back.
+        # Once https://github.com/flashinfer-ai/flashinfer/issues/2856 is fixed,
+        # we can remove the conditional and always pass params.context_buf.
+        out_buf = params.context_buf if q_len_per_req == 1 else None
+
+        mla_out = flashinfer.mla.trtllm_batch_decode_with_kv_cache_mla(
+            query=query,
+            kv_cache=kv_cache,
+            workspace_buffer=params.workspace.view(-1, 4),
+            qk_nope_head_dim=params.qk_nope_head_dim,
+            kv_lora_rank=params.kv_lora_rank,
+            qk_rope_head_dim=params.qk_rope_head_dim,
+            block_tables=block_tables,
+            seq_lens=params.sequence_lengths,
+            max_seq_len=params.max_past_kv_length,
+            out=out_buf,
+            bmm1_scale=bmm1_scale,
+            bmm2_scale=bmm2_scale,
+            sinks=params.attention_sinks,
+        )
+
+        if q_len_per_req > 1:
+            params.context_buf.copy_(mla_out.reshape_as(params.context_buf))
+
 
 def _parse_request_types(host_request_types: torch.Tensor) -> Tuple[int, int]:
     """
@@ -1693,8 +1768,18 @@ def trtllm_gen_attention(
 
     num_contexts, num_generations = _parse_request_types(host_request_types)
 
-    num_ctx_tokens = int(host_context_lengths[:num_contexts].sum()) if num_contexts > 0 else 0
-    num_gen_tokens = num_tokens - num_ctx_tokens
+    is_gen_only = attn_input_type == AttentionInputType.generation_only
+    is_ctx_only = attn_input_type == AttentionInputType.context_only
+
+    if is_gen_only:
+        num_ctx_tokens = 0
+        num_gen_tokens = num_tokens
+    elif is_ctx_only:
+        num_ctx_tokens = num_tokens
+        num_gen_tokens = 0
+    else:
+        num_ctx_tokens = int(host_context_lengths[:num_contexts].sum()) if num_contexts > 0 else 0
+        num_gen_tokens = num_tokens - num_ctx_tokens
 
     # Prepare Workspace
     # Use upper-bound token counts for workspace sizing to avoid repeated
@@ -1732,7 +1817,13 @@ def trtllm_gen_attention(
             workspace.resize_(required_workspace_size)
             workspace.zero_()
 
-    out_tensor = output.view(num_tokens, num_heads, head_size)
+    if is_mla_enable and is_gen_only and kv_lora_rank:
+        out_head_size = kv_lora_rank
+    elif is_mla_enable and v_head_dim:
+        out_head_size = v_head_dim
+    else:
+        out_head_size = head_size
+    out_tensor = output.view(num_tokens, num_heads, out_head_size)
 
     max_attn_window_size = (
         attention_window_size
@@ -1741,8 +1832,6 @@ def trtllm_gen_attention(
     )
     cyclic_attn_window_size = attention_window_size
 
-    is_gen_only = attn_input_type == AttentionInputType.generation_only
-
     common_params = dict(
         workspace=workspace,
         max_attention_window_size=max_attn_window_size,
@@ -1784,6 +1873,15 @@ def trtllm_gen_attention(
         paged_context_fmha=use_paged_context_fmha,
         attention_sinks=attention_sinks,
         is_mla_enable=is_mla_enable,
+        kv_lora_rank=kv_lora_rank or 0,
+        qk_nope_head_dim=qk_nope_head_dim or 0,
+        qk_rope_head_dim=qk_rope_head_dim or 0,
+        v_head_dim=v_head_dim or 0,
+        q_scaling=q_scaling,
+        latent_cache=latent_cache,
+        num_layers=host_kv_cache_pool_mapping.size(0)
+        if host_kv_cache_pool_mapping is not None
+        else 0,
     )
 
     # Context Phase
@@ -1853,6 +1951,9 @@ def trtllm_gen_attention(
             spec_decoding_position_offsets=spec_pos_offsets,
             spec_decoding_packed_mask=spec_packed_mask,
         )
-        backend.run_generation(gen_params)
+        if is_mla_enable:
+            backend.run_mla_generation(gen_params)
+        else:
+            backend.run_generation(gen_params)
 
     logger.debug(f"trtllm_gen_attention stops at layer {layer_idx}")