From d4136c0bf2e177b6ec9334f5b916b58dfafedb0b Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Thu, 19 Mar 2026 21:20:22 -0700 Subject: [PATCH 1/3] [None][feat] Support MLA in TrtllmGen attention backend Signed-off-by: Yihan Wang --- .../_torch/attention_backend/trtllm.py | 5 +- .../_torch/attention_backend/trtllm_gen.py | 142 +++++++++++++++--- 2 files changed, 124 insertions(+), 23 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index d67eb907283e..5c918cf5639a 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -30,8 +30,9 @@ 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" +# _TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION = os.environ.get( +# "TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION", "0") == "1" +_TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION = True @dataclass(kw_only=True, init=False) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/trtllm_gen.py index a11d8f51fe19..a4ca9573d2f2 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py +++ b/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,13 +228,6 @@ 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 - return ( - False, - f"[Generation] num_heads/num_kv_heads ratio " - f"({heads_ratio}) must be <= {max_ratio}.", - ) if has_alibi: return False, "[Generation] ALiBi is not supported." if (q_dtype, kv_cache_dtype, o_dtype) not in cls.SUPPORTED_DTYPE_COMBOS_GENERATION: @@ -261,7 +251,7 @@ def is_supported( return True, "" -@dataclass +@dataclass(slots=True) class ContextWorkspaceBuffers: """ Workspace buffers for context phase. @@ -292,7 +282,7 @@ class ContextWorkspaceBuffers: fmha_bmm2_scale: Optional[torch.Tensor] = None -@dataclass +@dataclass(slots=True) class GenerationWorkspaceBuffers: """ Workspace buffers for generation phase. @@ -855,12 +845,24 @@ 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 class EnqueueContextParams(EnqueueParams): batch_size: int = 0 mrope_rotary_cos_sin: Optional[torch.Tensor] = None + # MLA context: separate K, V inputs (non-absorption mode) + k_input: Optional[torch.Tensor] = None + v_input: Optional[torch.Tensor] = None + absorption_mode: bool = False @dataclass @@ -1361,6 +1363,72 @@ def run_generation(self, params: EnqueueGenerationParams): q_len_per_req=params.input_seq_length, ) + def run_mla_generation(self, params: EnqueueGenerationParams): + """MLA generation decode using flashinfer MLA kernel.""" + import math + + 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 +1761,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 +1810,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 +1825,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 +1866,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 @@ -1808,8 +1899,14 @@ def trtllm_gen_attention( input_seq_length=max_context_q_len, batch_size=num_seqs, mrope_rotary_cos_sin=mrope_rotary_cos_sin, + k_input=k[token_offset : token_offset + num_ctx_tokens] if k is not None else None, + v_input=v[token_offset : token_offset + num_ctx_tokens] if v is not None else None, + absorption_mode=False, ) - backend.run_context(ctx_params) + if is_mla_enable and not is_fused_qkv: + backend.run_mla_context(ctx_params) + else: + backend.run_context(ctx_params) # Generation Phase if num_generations > 0 and attn_input_type != AttentionInputType.context_only: @@ -1853,6 +1950,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}") From c055a39df228c1d2e1ce601fa7da9e7b0d3d4c6e Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Wed, 1 Apr 2026 04:03:49 -0700 Subject: [PATCH 2/3] Address review comments Signed-off-by: Yihan Wang --- .../_torch/attention_backend/trtllm.py | 7 +++-- .../_torch/attention_backend/trtllm_gen.py | 27 ++++++++++--------- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 5c918cf5639a..ad2c4d9a6315 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -29,10 +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" -_TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION = True +# 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) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/trtllm_gen.py index a4ca9573d2f2..2ed6abe0b05a 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm_gen.py @@ -228,6 +228,12 @@ def is_supported( False, f"[Generation] tokens_per_block ({tokens_per_block}) must be >= {cls.MIN_TOKENS_PER_BLOCK}.", ) + heads_ratio = num_heads // num_kv_heads + if not is_mla_enable and heads_ratio > cls.MAX_HEADS_RATIO_GENERATION: + return ( + False, + f"[Generation] heads ratio ({heads_ratio}) exceeds maximum ({cls.MAX_HEADS_RATIO_GENERATION}).", + ) if has_alibi: return False, "[Generation] ALiBi is not supported." if (q_dtype, kv_cache_dtype, o_dtype) not in cls.SUPPORTED_DTYPE_COMBOS_GENERATION: @@ -859,10 +865,6 @@ class EnqueueParams: class EnqueueContextParams(EnqueueParams): batch_size: int = 0 mrope_rotary_cos_sin: Optional[torch.Tensor] = None - # MLA context: separate K, V inputs (non-absorption mode) - k_input: Optional[torch.Tensor] = None - v_input: Optional[torch.Tensor] = None - absorption_mode: bool = False @dataclass @@ -1363,9 +1365,14 @@ def run_generation(self, params: EnqueueGenerationParams): q_len_per_req=params.input_seq_length, ) - def run_mla_generation(self, params: EnqueueGenerationParams): + def run_mla_generation(self, params: EnqueueGenerationParams) -> None: """MLA generation decode using flashinfer MLA kernel.""" - import math + 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( @@ -1899,14 +1906,8 @@ def trtllm_gen_attention( input_seq_length=max_context_q_len, batch_size=num_seqs, mrope_rotary_cos_sin=mrope_rotary_cos_sin, - k_input=k[token_offset : token_offset + num_ctx_tokens] if k is not None else None, - v_input=v[token_offset : token_offset + num_ctx_tokens] if v is not None else None, - absorption_mode=False, ) - if is_mla_enable and not is_fused_qkv: - backend.run_mla_context(ctx_params) - else: - backend.run_context(ctx_params) + backend.run_context(ctx_params) # Generation Phase if num_generations > 0 and attn_input_type != AttentionInputType.context_only: From c5541d153b2b69a5dd682237dbcc7aa4d4985464 Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Thu, 2 Apr 2026 19:29:31 -0700 Subject: [PATCH 3/3] Disable TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION and run thop.attention path Signed-off-by: Yihan Wang --- tensorrt_llm/_torch/attention_backend/trtllm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index ad2c4d9a6315..7554ead2401a 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/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 (default: on). +# Enable TRTLLM-Gen attention backend via environment variable (default: off). _TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION = (os.environ.get( - "TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION", "1") == "1") + "TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION", "0") == "1") @dataclass(kw_only=True, init=False)