From c3f3fcb3693219260f263fecfc70461e296bb698 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 10 Jun 2026 06:52:25 +0000 Subject: [PATCH 1/8] [TRTLLM-12807][feat] Refactor TRTLLM attention FMHA libraries Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/attention_backend/fmha/__init__.py | 32 ++ .../_torch/attention_backend/fmha/fallback.py | 187 +++++++++ .../flashinfer_trtllm_gen.py} | 381 ++++++------------ .../attention_backend/fmha/interface.py | 59 +++ .../_torch/attention_backend/fmha/phased.py | 269 +++++++++++++ .../_torch/attention_backend/fmha/registry.py | 81 ++++ .../modules/ATTENTION_DEVELOPER_GUIDE.md | 33 +- .../test_attention_op_sync.py | 37 +- 8 files changed, 801 insertions(+), 278 deletions(-) create mode 100644 tensorrt_llm/_torch/attention_backend/fmha/__init__.py create mode 100644 tensorrt_llm/_torch/attention_backend/fmha/fallback.py rename tensorrt_llm/_torch/attention_backend/{trtllm_gen.py => fmha/flashinfer_trtllm_gen.py} (80%) create mode 100644 tensorrt_llm/_torch/attention_backend/fmha/interface.py create mode 100644 tensorrt_llm/_torch/attention_backend/fmha/phased.py create mode 100644 tensorrt_llm/_torch/attention_backend/fmha/registry.py diff --git a/tensorrt_llm/_torch/attention_backend/fmha/__init__.py b/tensorrt_llm/_torch/attention_backend/fmha/__init__.py new file mode 100644 index 000000000000..1c3981abcf91 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/fmha/__init__.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .fallback import FallbackFmha +from .flashinfer_trtllm_gen import FlashInferTrtllmGenFmha +from .interface import Fmha +from .phased import FmhaParams, PhasedFmha +from .registry import DEFAULT_FMHA_LIBS, FMHA_LIBS, FmhaCls, get_enabled_fmha_lib_classes + +__all__ = [ + "DEFAULT_FMHA_LIBS", + "FMHA_LIBS", + "FallbackFmha", + "FlashInferTrtllmGenFmha", + "Fmha", + "FmhaCls", + "FmhaParams", + "PhasedFmha", + "get_enabled_fmha_lib_classes", +] diff --git a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py new file mode 100644 index 000000000000..384f4b8e59e5 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py @@ -0,0 +1,187 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import TYPE_CHECKING, Optional + +import torch + +from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs +from tensorrt_llm._torch.attention_backend.sparse.skip_softmax import ( + SkipSoftmaxKernelParams, + SkipSoftmaxParams, +) +from tensorrt_llm.bindings.internal import thop + +from .interface import Fmha + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata + + +# ``AttentionForwardArgs`` fields that this backend does not consume. +# Sync test (test_attention_op_sync.py) requires every other field to map to a +# kwarg name, a @property on the dataclass, or a field that some @property +# transitively reads; entries here are exempt. +_THOP_EXCLUDED_FIELDS: frozenset = frozenset( + { + "topk_indices", # DSA-only + "attention_mask_data", # custom-mask code path + "out_scale_sf", # promoted into ``out_scale`` in ``TrtllmAttention._run`` for NVFP4 path + } +) + +# ``thop.attention`` kwargs hard-wired to a literal at the call site (no +# rich object owns them). Sync test enforces both the kwarg name and the +# literal value. +_THOP_LITERALS: dict = {} + + +class FallbackFmha(Fmha): + """Fallback FMHA implementation using the fused TRT-LLM thop attention op.""" + + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> None: + attn = self.attn + sparse_params = attn.sparse_params + skip_softmax_kernel_params = ( + sparse_params.scheduler.get_kernel_params(timestep=forward_args.timestep) + if isinstance(sparse_params, SkipSoftmaxParams) + else SkipSoftmaxKernelParams() + ) + + # Every kwarg sources from ``attn`` / ``metadata`` / ``forward_args`` + # (with ``forward_args.sparse_prediction`` for sparse-attn inputs), + # ``skip_softmax_kernel_params``, or a literal allowlisted in + # ``_THOP_LITERALS``. ``test_attention_op_sync.py`` enforces this + # statically. + thop.attention( + q=q, + k=k, + v=v, + output=forward_args.output, + output_sf=forward_args.output_sf, + workspace_=metadata.effective_workspace, + # --- Per-step batch state (TrtllmAttentionMetadata) --- + sequence_length=metadata.kv_lens_cuda_runtime, + host_past_key_value_lengths=metadata.kv_lens_runtime, + host_total_kv_lens=metadata.host_total_kv_lens, + context_lengths=metadata.prompt_lens_cuda_runtime, + host_context_lengths=metadata.prompt_lens_cpu_runtime, + host_request_types=metadata.host_request_types_runtime, + max_context_q_len_override=metadata.max_context_q_len_override, + kv_cache_block_offsets=metadata.kv_cache_block_offsets, + host_kv_cache_pool_pointers=metadata.host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping=metadata.host_kv_cache_pool_mapping, + cache_indirection=metadata.cache_indirection, + block_ids_per_seq=metadata.block_ids_per_seq, + tokens_per_block=metadata.tokens_per_block, + max_num_requests=metadata.max_num_requests, + beam_width=metadata.effective_beam_width, + use_paged_context_fmha=metadata.use_paged_context_fmha, + helix_position_offsets=metadata.helix_position_offsets, + helix_is_inactive_rank=metadata.helix_is_inactive_rank, + is_spec_decoding_enabled=metadata.is_spec_decoding_enabled, + use_spec_decoding=metadata.use_spec_decoding, + is_spec_dec_tree=metadata.is_spec_dec_tree, + spec_decoding_generation_lengths=metadata.spec_decoding_generation_lengths, + spec_decoding_position_offsets_for_cpp=metadata.spec_decoding_position_offsets_for_cpp, + spec_decoding_packed_mask=metadata.spec_decoding_packed_mask, + spec_decoding_bl_tree_mask_offset=metadata.spec_decoding_bl_tree_mask_offset, + spec_decoding_bl_tree_mask=metadata.spec_decoding_bl_tree_mask, + spec_decoding_target_max_draft_tokens=metadata.max_total_draft_tokens, + spec_bl_tree_first_sparse_mask_offset_kv=metadata.spec_bl_tree_first_sparse_mask_offset_kv, + num_sparse_topk=metadata.num_sparse_topk, + flash_mla_tile_scheduler_metadata=metadata.flash_mla_tile_scheduler_metadata, + flash_mla_num_splits=metadata.flash_mla_num_splits, + num_contexts=metadata.num_contexts, + num_ctx_tokens=metadata.num_ctx_tokens, + max_context_length=metadata.max_context_length, + max_seq_len=metadata.max_seq_len, + trtllm_gen_jit_warmup=metadata.trtllm_gen_jit_warmup, + is_cross=metadata.is_cross, + # --- Per-call (AttentionForwardArgs) --- + out_scale=forward_args.out_scale, + kv_scale_orig_quant=forward_args.kv_scale_orig_quant, + kv_scale_quant_orig=forward_args.kv_scale_quant_orig, + latent_cache=forward_args.latent_cache, + q_pe=forward_args.q_pe, + attention_sinks=forward_args.attention_sinks, + mask_type=forward_args.mask_type, + attention_input_type=int(forward_args.attention_input_type), + attention_window_size=forward_args.attention_window_size, + chunked_prefill_buffer_batch_size=forward_args.chunked_prefill_buffer_batch_size, + mrope_rotary_cos_sin=forward_args.mrope_rotary_cos_sin, + mrope_position_deltas=forward_args.mrope_position_deltas, + softmax_stats_tensor=forward_args.softmax_stats_tensor, + cu_q_seqlens=forward_args.cu_q_seqlens, + cu_kv_seqlens=forward_args.cu_kv_seqlens, + fmha_scheduler_counter=forward_args.fmha_scheduler_counter, + mla_bmm1_scale=forward_args.mla_bmm1_scale, + mla_bmm2_scale=forward_args.mla_bmm2_scale, + quant_q_buffer=forward_args.quant_q_buffer, + sage_attn_num_elts_per_blk_q=forward_args.sage_attn_num_elts_per_blk_q, + sage_attn_num_elts_per_blk_k=forward_args.sage_attn_num_elts_per_blk_k, + sage_attn_num_elts_per_blk_v=forward_args.sage_attn_num_elts_per_blk_v, + sage_attn_qk_int8=forward_args.sage_attn_qk_int8, + is_fused_qkv=forward_args.is_fused_qkv, + update_kv_cache=forward_args.update_kv_cache, + cross_kv=forward_args.cross_kv, + relative_attention_bias=forward_args.relative_attention_bias, + relative_attention_max_distance=forward_args.relative_attention_max_distance, + # --- Module config (TrtllmAttention) --- + rotary_inv_freq=attn.rotary_inv_freq, + rotary_cos_sin=attn.rotary_cos_sin, + predicted_tokens_per_seq=attn.predicted_tokens_per_seq, + local_layer_idx=attn.local_layer_idx, + num_heads=attn.num_heads, + num_kv_heads=attn.num_kv_heads, + head_size=attn.head_dim, + quant_mode=attn.quant_mode, + q_scaling=attn.q_scaling, + position_embedding_type=attn.position_embedding_type, + rope_dim=attn.rope_dim, + rope_base=attn.rope_base, + rope_scale_type=attn.rope_scale_type, + rope_scale=attn.rope_scale, + rope_short_m_scale=attn.rope_short_m_scale, + rope_long_m_scale=attn.rope_long_m_scale, + rope_max_positions=attn.rope_max_positions, + rope_original_max_positions=attn.rope_original_max_positions, + is_mla_enable=attn.is_mla_enable, + q_lora_rank=attn.q_lora_rank, + kv_lora_rank=attn.kv_lora_rank, + qk_nope_head_dim=attn.qk_nope_head_dim, + qk_rope_head_dim=attn.qk_rope_head_dim, + v_head_dim=attn.v_head_dim, + rope_append=attn.rope_append, + attention_chunk_size=attn.attention_chunk_size, + skip_softmax_threshold_scale_factor_prefill=skip_softmax_kernel_params.threshold_scale_factor_prefill, + skip_softmax_threshold_scale_factor_decode=skip_softmax_kernel_params.threshold_scale_factor_decode, + skip_softmax_stat=attn.skip_softmax_stat, + # --- Sparse-specific (AttentionForwardArgs.sparse_prediction) --- + sparse_kv_indices=forward_args.sparse_prediction.sparse_kv_indices, + sparse_kv_offsets=forward_args.sparse_prediction.sparse_kv_offsets, + sparse_attn_indices=forward_args.sparse_prediction.sparse_attn_indices, + sparse_attn_offsets=forward_args.sparse_prediction.sparse_attn_offsets, + sparse_attn_indices_block_size=forward_args.sparse_prediction.sparse_attn_indices_block_size, + sparse_mla_topk_lens=forward_args.sparse_prediction.sparse_mla_topk_lens, + compressed_kv_cache_pool_ptr=forward_args.sparse_prediction.compressed_kv_cache_pool_ptr, + ) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py similarity index 80% rename from tensorrt_llm/_torch/attention_backend/trtllm_gen.py rename to tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py index aca899c5ed1a..dc4d382e0db0 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py @@ -14,11 +14,12 @@ # limitations under the License. """ -TrtLLM-Gen Attention Backend +FlashInfer TRTLLM-Gen FMHA This module implements attention computation using flashinfer's trtllm-gen kernels. -It provides a drop-in replacement for thop.attention() with support for trtllm-gen -kernel only (Blackwell architecture: SM100/SM103). Enabled via TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION=1. +It provides a TRT-LLM attention FMHA library for trtllm-gen kernels +(Blackwell architecture: SM100/SM103). Enable or disable it through +``TLLM_FMHA_LIBS``. Architecture: - QKV preprocessing & RoPE: C++ kernels via tensorrt_llm.bindings.internal.thop, @@ -27,20 +28,17 @@ the paged KV cache fields carried by FmhaParams. Entry points: - FlashInferTrtllmGenAttention.is_supported() - Check if trtllm-gen can handle the given config. - FlashInferTrtllmGenAttention.forward() - Main attention method. + FlashInferTrtllmGenFmha.is_available() - Check if this FMHA library can be instantiated. + FlashInferTrtllmGenFmha.is_supported() - Check if trtllm-gen can handle the given request. + FlashInferTrtllmGenFmha.forward() - Main attention method. Example: - backend = FlashInferTrtllmGenAttention(attention_layer=...) - supported, reason = backend.is_supported(q, k, v, attn=..., meta=..., fwd=...) - if supported: - backend.forward(q, k, v, attn=..., meta=..., fwd=...) - else: - Fallback to thop.attention() + fmha = FlashInferTrtllmGenFmha(attn=...) + if fmha.is_supported(q, k, v, metadata, forward_args): + fmha.forward(q, k, v, metadata, forward_args) """ import math -from dataclasses import dataclass from functools import lru_cache from typing import TYPE_CHECKING, List, Optional, Tuple @@ -56,8 +54,11 @@ from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal import thop from tensorrt_llm.functional import AttentionMaskType +from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantMode +from .phased import FmhaParams, PhasedFmha + if TYPE_CHECKING: from tensorrt_llm._torch.attention_backend.trtllm import ( TrtllmAttention, @@ -350,36 +351,7 @@ def _get_workspace_size( return max(context_size, generation_size) -@dataclass(slots=True) -class FmhaParams: - attn: "TrtllmAttention" - meta: "TrtllmAttentionMetadata" - fwd: AttentionForwardArgs - workspace: torch.Tensor - attention_input: Optional[torch.Tensor] = None - qkv_input: Optional[torch.Tensor] = None - context_buf: Optional[torch.Tensor] = None - sequence_lengths: Optional[torch.Tensor] = None - context_lengths: Optional[torch.Tensor] = None - input_seq_length: int = 0 - max_past_kv_length: int = 0 - max_attention_window_size: int = 0 - cyclic_attention_window_size: int = 0 - num_tokens: int = 0 - seq_offset: int = 0 - tokens_per_block: int = 64 - fp8_context_fmha: bool = False - kv_factor: int = 0 - total_num_blocks: int = 0 - # Context-only fields - batch_size: int = 0 - # Generation-only fields - num_requests: int = 0 - spec_decoding_generation_lengths: Optional[torch.Tensor] = None - spec_decoding_position_offsets: Optional[torch.Tensor] = None - - -class FlashInferTrtllmGenAttention: +class FlashInferTrtllmGenFmha(PhasedFmha): """ An attention backend using pure trtllm-gen kernels from flashinfer. """ @@ -387,6 +359,7 @@ class FlashInferTrtllmGenAttention: # Default KV layout for flashinfer # HND = [max_num_pages, kv_factor, num_kv_heads, page_size, head_dim] DEFAULT_KV_LAYOUT = "HND" + REQUIRES_PAGED_KV = True # Keep shared paged indices disabled to match the current TensorRT-LLM # block-table layout used by the fused preprocessing path. USE_SHARED_PAGED_KV_IDX = False @@ -436,22 +409,63 @@ class FlashInferTrtllmGenAttention: (576, 512, 32), } - def __init__( - self, - attention_layer: "TrtllmAttention", - ): + def __init__(self, attn: "TrtllmAttention"): + super().__init__(attn) self._layout = self.DEFAULT_KV_LAYOUT # Read once so the hot path is not sensitive to later environment changes. self._enable_pdl = get_env_enable_pdl() - missing_ops = self._missing_fused_nanobind_ops() - if missing_ops: - raise RuntimeError( - f"trtllm-gen requires fused nanobind ops, missing: {', '.join(missing_ops)}." - ) # Lazily set on the first forward() call from the query device. self._multi_processor_count: Optional[int] = None + @classmethod + def is_available(cls, attn: "TrtllmAttention") -> bool: + if not IS_FLASHINFER_AVAILABLE: + logger.debug("FlashInfer TRTLLM-Gen FMHA is unavailable: flashinfer is not installed.") + return False + + missing_ops = cls._missing_fused_nanobind_ops() + if missing_ops: + logger.debug( + "FlashInfer TRTLLM-Gen FMHA is unavailable: missing fused " + f"nanobind ops: {', '.join(missing_ops)}." + ) + return False + + sm = get_sm_version() + if not is_sm_100f(sm): + logger.debug( + f"FlashInfer TRTLLM-Gen FMHA is unavailable: requires SM100 or SM103, got SM{sm}." + ) + return False + + has_skip_softmax = ( + attn.skip_softmax_threshold_scale_factor_prefill is not None + or attn.skip_softmax_threshold_scale_factor_decode is not None + ) + if has_skip_softmax: + logger.debug( + "FlashInfer TRTLLM-Gen FMHA is unavailable: skip-softmax attention is enabled." + ) + return False + + if attn.num_heads <= 0 or attn.num_kv_heads <= 0: + logger.debug( + "FlashInfer TRTLLM-Gen FMHA is unavailable: " + f"num_heads={attn.num_heads}, num_kv_heads={attn.num_kv_heads}." + ) + return False + + if attn.num_heads % attn.num_kv_heads != 0: + logger.debug( + "FlashInfer TRTLLM-Gen FMHA is unavailable: " + f"num_heads ({attn.num_heads}) must be divisible by " + f"num_kv_heads ({attn.num_kv_heads})." + ) + return False + + return True + @property def layout(self) -> str: """KV cache layout.""" @@ -483,45 +497,13 @@ def _get_kv_scale_params( return kv_scale_orig_quant, kv_scale_quant_orig @staticmethod - def _get_kv_cache_dtype_and_total_blocks( + def _get_kv_cache_dtype( meta: "TrtllmAttentionMetadata", - is_mla_enable: bool, - ) -> Tuple[Optional[DataType], int]: - kv_cache_dtype = None - total_num_blocks = 0 + ) -> Optional[DataType]: kv_cache_manager = meta.kv_cache_manager if kv_cache_manager is not None: - kv_cache_dtype = kv_cache_manager.dtype - kv_factor = 1 if is_mla_enable else 2 - blocks_in_primary_pool = getattr(kv_cache_manager, "blocks_in_primary_pool", None) - if blocks_in_primary_pool is None: - blocks_per_window = getattr(kv_cache_manager, "blocks_per_window", None) - if blocks_per_window: - blocks_in_primary_pool = max( - int(primary) for primary, _ in blocks_per_window.values() - ) - if blocks_in_primary_pool is not None: - total_num_blocks = ( - int(blocks_in_primary_pool) * kv_cache_manager.num_local_layers * kv_factor - ) - return kv_cache_dtype, total_num_blocks - - @staticmethod - def _get_kv_factor(attn: "TrtllmAttention") -> int: - return 1 if attn.is_mla_enable else 2 - - @staticmethod - def _get_generation_out_head_size(attn: "TrtllmAttention") -> int: - kv_lora_rank = attn.kv_lora_rank or 0 - if attn.is_mla_enable and kv_lora_rank: - return kv_lora_rank - return attn.head_dim - - @staticmethod - def _get_context_out_head_size(attn: "TrtllmAttention") -> int: - if attn.is_mla_enable and attn.v_head_dim: - return attn.v_head_dim - return attn.head_dim + return kv_cache_manager.dtype + return None @staticmethod def _get_bmm1_scale(attn: "TrtllmAttention") -> float: @@ -584,6 +566,26 @@ def _check_mla_generation_support( return True, "" def is_supported( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> bool: + supported, reason = self._is_supported_with_reason( + q, + k, + v, + self.attn, + metadata, + forward_args, + ) + if not supported: + logger.debug(f"FlashInfer TRTLLM-Gen FMHA does not support request: {reason}") + return supported + + def _is_supported_with_reason( self, q: torch.Tensor, k: Optional[torch.Tensor], @@ -625,8 +627,6 @@ def is_supported( if is_mla_enable and fwd.attention_input_type != AttentionInputType.generation_only: return False, "trtllm-gen MLA supports generation-only attention." - if not IS_FLASHINFER_AVAILABLE: - return False, "flashinfer package is not installed." if meta.kv_cache_block_offsets is None: return False, "trtllm-gen requires paged KV cache." output = fwd.output @@ -643,16 +643,12 @@ def is_supported( q_dtype = q.dtype o_dtype = output.dtype - sm = get_sm_version() - if not is_sm_100f(sm): - return False, (f"trtllm-gen requires SM100 or SM103 (Blackwell). Current: SM{sm}.") - if q_dtype not in self.SUPPORTED_INPUT_DTYPES: return False, ( f"Input dtype {q_dtype} not supported. Supported: FP16, BF16, FP8 (E4M3)." ) - kv_cache_dtype, _ = self._get_kv_cache_dtype_and_total_blocks(meta, is_mla_enable) + kv_cache_dtype = self._get_kv_cache_dtype(meta) if kv_cache_dtype is None: kv_cache_dtype = torch_dtype_to_binding(q_dtype) @@ -673,15 +669,6 @@ def is_supported( if o_dtype not in self.SUPPORTED_OUT_DTYPES: return False, f"Output dtype {o_dtype} not supported. Supported: FP16, BF16, FP8." - assert attn.num_heads > 0, "num_heads must be positive." - assert attn.num_kv_heads > 0, "num_kv_heads must be positive." - if attn.num_heads % attn.num_kv_heads != 0: - return ( - False, - f"num_heads ({attn.num_heads}) must be divisible by " - f"num_kv_heads ({attn.num_kv_heads}).", - ) - has_alibi = attn.position_embedding_type in (4, 5) check_context_phase = has_context_phase and not is_mla_enable if check_context_phase: @@ -772,80 +759,64 @@ def _get_multi_processor_count(self, device: torch.device) -> int: device_index = torch.cuda.current_device() return self._get_multi_processor_count_for_device(device_index) - def forward( + def get_fp8_context_fmha( + self, + q: torch.Tensor, + output: torch.Tensor, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + is_gen_only: bool, + ) -> bool: + del q, metadata, forward_args + kv_cache_quant_mode = QuantMode(self.attn.quant_mode) + return ( + output.dtype == torch.float8_e4m3fn + or output.dtype == torch.uint8 + or kv_cache_quant_mode.has_fp4_kv_cache() + or (kv_cache_quant_mode.has_fp8_kv_cache() and not is_gen_only) + ) + + def prepare_workspace( self, q: torch.Tensor, k: Optional[torch.Tensor], v: Optional[torch.Tensor], - attn: "TrtllmAttention", - meta: "TrtllmAttentionMetadata", - fwd: AttentionForwardArgs, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + workspace: torch.Tensor, ) -> None: - output = fwd.output - if output is None: - raise RuntimeError("trtllm-gen attention requires output.") - if meta.kv_cache_block_offsets is None: - raise RuntimeError("trtllm-gen attention requires paged KV cache.") - - workspace = meta.effective_workspace - if workspace is None: - workspace = torch.empty((0,), device=q.device, dtype=torch.int8) - + del k, v + attn = self.attn # Lazily cache the SM count from the first query tensor's device. if self._multi_processor_count is None: self._multi_processor_count = self._get_multi_processor_count(q.device) - num_heads = attn.num_heads - num_kv_heads = attn.num_kv_heads - head_size = attn.head_dim - quant_mode = attn.quant_mode - is_mla_enable = attn.is_mla_enable - tokens_per_block = meta.tokens_per_block - max_num_requests = meta.max_num_requests - max_context_length = meta.max_context_length - attention_window_size = fwd.attention_window_size - beam_width = meta.beam_width - num_tokens = q.size(0) - attn_input_type = fwd.attention_input_type - is_gen_only = attn_input_type == AttentionInputType.generation_only - is_fp8_out = output.dtype == torch.float8_e4m3fn - is_fp4_out = output.dtype == torch.uint8 - kv_cache_quant_mode = QuantMode(quant_mode) - fp8_context_fmha = ( - is_fp8_out - or is_fp4_out - or kv_cache_quant_mode.has_fp4_kv_cache() - or (kv_cache_quant_mode.has_fp8_kv_cache() and not is_gen_only) - ) - - num_contexts = meta.num_contexts - num_ctx_tokens = meta.num_ctx_tokens - num_generations = meta.num_generations - num_gen_tokens = num_tokens if is_gen_only else num_tokens - num_ctx_tokens - if num_gen_tokens < 0: - raise RuntimeError( - f"Invalid trtllm-gen attention token counts: num_tokens={num_tokens}, " - f"num_ctx_tokens={num_ctx_tokens}, attention_input_type={attn_input_type}." - ) + attention_input_type = forward_args.attention_input_type + is_gen_only = attention_input_type == AttentionInputType.generation_only + num_gen_tokens = num_tokens if is_gen_only else num_tokens - metadata.num_ctx_tokens + output = forward_args.output + if output is None: + raise RuntimeError(f"{type(self).__name__} requires output.") + fp8_context_fmha = self.get_fp8_context_fmha(q, output, metadata, forward_args, is_gen_only) - workspace_max_tokens = max(num_tokens, max_context_length) - workspace_max_gen_tokens = max(num_gen_tokens, max_num_requests) + workspace_max_tokens = max(num_tokens, metadata.max_context_length) + workspace_max_gen_tokens = max(num_gen_tokens, metadata.max_num_requests) required_workspace_size = _get_workspace_size( dtype=q.dtype, num_tokens=workspace_max_tokens, num_gen_tokens=workspace_max_gen_tokens, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - head_size=head_size, - max_num_requests=max_num_requests, + num_heads=attn.num_heads, + num_kv_heads=attn.num_kv_heads, + head_size=attn.head_dim, + max_num_requests=metadata.max_num_requests, rotary_embedding_dim=attn.rope_dim, fp8_context_fmha=fp8_context_fmha, ) current_workspace_size = workspace.numel() * workspace.element_size() if current_workspace_size < required_workspace_size: - if meta.is_cuda_graph and torch.cuda.is_current_stream_capturing(): + 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." @@ -853,104 +824,6 @@ def forward( required_workspace_numel = math.ceil(required_workspace_size / workspace.element_size()) workspace.resize_((required_workspace_numel,)) - out_head_size = ( - self._get_generation_out_head_size(attn) - if is_gen_only - else self._get_context_out_head_size(attn) - ) - out_tensor = output.view(num_tokens, num_heads, out_head_size) - - cache_indirection = meta.cache_indirection - max_attn_window_size = ( - attention_window_size - if beam_width == 1 - else ( - cache_indirection.size(2) - if cache_indirection is not None - else attention_window_size - ) - ) - cyclic_attn_window_size = attention_window_size - tokens_per_block = tokens_per_block if tokens_per_block is not None else 64 - _, total_num_blocks = self._get_kv_cache_dtype_and_total_blocks(meta, is_mla_enable) - - params = FmhaParams( - attn=attn, - meta=meta, - fwd=fwd, - workspace=workspace, - max_attention_window_size=max_attn_window_size, - cyclic_attention_window_size=cyclic_attn_window_size, - tokens_per_block=tokens_per_block, - fp8_context_fmha=fp8_context_fmha, - kv_factor=self._get_kv_factor(attn), - total_num_blocks=total_num_blocks, - ) - - sequence_length = meta.kv_lens_cuda_runtime - host_past_key_value_lengths = meta.kv_lens_runtime - - if num_contexts > 0 and attn_input_type != AttentionInputType.generation_only: - seq_offset = 0 - token_offset = 0 - num_seqs = num_contexts - - context_lengths = meta.prompt_lens_cuda_runtime - host_context_lengths = meta.prompt_lens_cpu_runtime - max_context_q_len = int(host_context_lengths[seq_offset : seq_offset + num_seqs].max()) - max_past_kv_len = int( - host_past_key_value_lengths[seq_offset : seq_offset + num_seqs].max() - ) - - params.attention_input = q[token_offset : token_offset + num_ctx_tokens] - params.qkv_input = params.attention_input - params.context_buf = out_tensor[token_offset : token_offset + num_ctx_tokens] - params.sequence_lengths = sequence_length[seq_offset:] - params.context_lengths = context_lengths[seq_offset:] - params.max_past_kv_length = max_past_kv_len - params.num_tokens = num_ctx_tokens - params.seq_offset = seq_offset - params.input_seq_length = max_context_q_len - params.batch_size = num_seqs - self.run_context(params) - - if num_generations > 0 and attn_input_type != AttentionInputType.context_only: - seq_offset = num_contexts - token_offset = 0 if is_gen_only else num_ctx_tokens - num_seqs = num_generations - - max_past_kv_len = int( - host_past_key_value_lengths[seq_offset : seq_offset + num_seqs].max() - ) - input_seq_length = num_gen_tokens // num_seqs if num_seqs > 0 else 1 - - predicted_tokens_per_seq = attn.predicted_tokens_per_seq - spec_gen_lengths = None - spec_pos_offsets = None - if meta.is_spec_decoding_enabled and predicted_tokens_per_seq > 1: - spec_gen_lengths = meta.spec_decoding_generation_lengths - position_offsets_for_cpp = meta.spec_decoding_position_offsets_for_cpp - if position_offsets_for_cpp is not None and position_offsets_for_cpp.dim() == 1: - position_offsets_for_cpp = position_offsets_for_cpp.view(max_num_requests, -1) - spec_pos_offsets = position_offsets_for_cpp - - params.attention_input = q[token_offset : token_offset + num_gen_tokens] - params.qkv_input = params.attention_input - params.context_buf = out_tensor[token_offset : token_offset + num_gen_tokens] - params.sequence_lengths = sequence_length[seq_offset:] - params.max_past_kv_length = max_past_kv_len - params.num_tokens = num_gen_tokens - params.seq_offset = seq_offset - params.input_seq_length = input_seq_length - params.num_requests = num_seqs // beam_width - params.spec_decoding_generation_lengths = spec_gen_lengths - params.spec_decoding_position_offsets = spec_pos_offsets - if is_mla_enable: - self.run_mla_generation(params) - else: - self.run_generation(params) - return - @staticmethod def _compute_window_left( cyclic_attention_window_size: int, diff --git a/tensorrt_llm/_torch/attention_backend/fmha/interface.py b/tensorrt_llm/_torch/attention_backend/fmha/interface.py new file mode 100644 index 000000000000..0e40e753d2da --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/fmha/interface.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Optional + +import torch + +from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention_backend.trtllm import ( + TrtllmAttention, + TrtllmAttentionMetadata, + ) + + +class Fmha(ABC): + """Common runtime contract for TRT-LLM attention FMHA libraries.""" + + def __init__(self, attn: "TrtllmAttention"): + self.attn = attn + + @classmethod + def is_available(cls, attn: "TrtllmAttention") -> bool: + return True + + def is_supported( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> bool: + return True + + @abstractmethod + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> None: + raise NotImplementedError diff --git a/tensorrt_llm/_torch/attention_backend/fmha/phased.py b/tensorrt_llm/_torch/attention_backend/fmha/phased.py new file mode 100644 index 000000000000..43a7cb2aed0a --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/fmha/phased.py @@ -0,0 +1,269 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional, cast + +import torch + +from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs, AttentionInputType + +from .interface import Fmha + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention_backend.trtllm import ( + TrtllmAttention, + TrtllmAttentionMetadata, + ) + + +@dataclass(slots=True) +class FmhaParams: + attn: "TrtllmAttention" + meta: "TrtllmAttentionMetadata" + fwd: AttentionForwardArgs + workspace: torch.Tensor + attention_input: Optional[torch.Tensor] = None + qkv_input: Optional[torch.Tensor] = None + context_buf: Optional[torch.Tensor] = None + sequence_lengths: Optional[torch.Tensor] = None + context_lengths: Optional[torch.Tensor] = None + input_seq_length: int = 0 + max_past_kv_length: int = 0 + max_attention_window_size: int = 0 + cyclic_attention_window_size: int = 0 + num_tokens: int = 0 + seq_offset: int = 0 + tokens_per_block: int = 64 + fp8_context_fmha: bool = False + kv_factor: int = 0 + total_num_blocks: int = 0 + # Context-only fields + batch_size: int = 0 + # Generation-only fields + num_requests: int = 0 + spec_decoding_generation_lengths: Optional[torch.Tensor] = None + spec_decoding_position_offsets: Optional[torch.Tensor] = None + + +class PhasedFmha(Fmha): + """FMHA base that dispatches mixed requests by phase and attention family.""" + + REQUIRES_PAGED_KV = False + + def __init__(self, attn: "TrtllmAttention"): + super().__init__(attn) + self.kv_factor = 1 if attn.is_mla_enable else 2 + kv_lora_rank = attn.kv_lora_rank or 0 + self.generation_out_head_size = ( + kv_lora_rank if attn.is_mla_enable and kv_lora_rank else attn.head_dim + ) + self.context_out_head_size = ( + attn.v_head_dim if attn.is_mla_enable and attn.v_head_dim else attn.head_dim + ) + + def _get_total_num_blocks( + self, + meta: "TrtllmAttentionMetadata", + ) -> int: + kv_cache_manager = meta.kv_cache_manager + if kv_cache_manager is None: + return 0 + + blocks_in_primary_pool = getattr(kv_cache_manager, "blocks_in_primary_pool", None) + if blocks_in_primary_pool is None: + blocks_per_window = getattr(kv_cache_manager, "blocks_per_window", None) + if blocks_per_window: + blocks_in_primary_pool = max( + int(primary) for primary, _ in blocks_per_window.values() + ) + if blocks_in_primary_pool is None: + return 0 + return int(blocks_in_primary_pool) * kv_cache_manager.num_local_layers * self.kv_factor + + def get_fp8_context_fmha( + self, + q: torch.Tensor, + output: torch.Tensor, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + is_gen_only: bool, + ) -> bool: + return False + + def prepare_workspace( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + workspace: torch.Tensor, + ) -> None: + pass + + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> None: + attn = self.attn + output = forward_args.output + if output is None: + raise RuntimeError(f"{type(self).__name__} requires output.") + if self.REQUIRES_PAGED_KV and metadata.kv_cache_block_offsets is None: + raise RuntimeError(f"{type(self).__name__} requires paged KV cache.") + + workspace = cast(torch.Tensor, metadata.effective_workspace) + + num_tokens = q.size(0) + attention_input_type = forward_args.attention_input_type + is_gen_only = attention_input_type == AttentionInputType.generation_only + + num_contexts = metadata.num_contexts + num_ctx_tokens = metadata.num_ctx_tokens + num_generations = metadata.num_generations + num_gen_tokens = num_tokens if is_gen_only else num_tokens - num_ctx_tokens + if num_gen_tokens < 0: + raise RuntimeError( + f"Invalid FMHA token counts: num_tokens={num_tokens}, " + f"num_ctx_tokens={num_ctx_tokens}, attention_input_type={attention_input_type}." + ) + + fp8_context_fmha = self.get_fp8_context_fmha(q, output, metadata, forward_args, is_gen_only) + self.prepare_workspace( + q, + k, + v, + metadata, + forward_args, + workspace, + ) + + out_head_size = self.generation_out_head_size if is_gen_only else self.context_out_head_size + out_tensor = output.view(num_tokens, attn.num_heads, out_head_size) + + attention_window_size = forward_args.attention_window_size + cache_indirection = metadata.cache_indirection + max_attention_window_size = ( + attention_window_size + if metadata.beam_width == 1 + else ( + cache_indirection.size(2) + if cache_indirection is not None + else attention_window_size + ) + ) + tokens_per_block = ( + metadata.tokens_per_block if metadata.tokens_per_block is not None else 64 + ) + + params = FmhaParams( + attn=attn, + meta=metadata, + fwd=forward_args, + workspace=workspace, + max_attention_window_size=max_attention_window_size, + cyclic_attention_window_size=attention_window_size, + tokens_per_block=tokens_per_block, + fp8_context_fmha=fp8_context_fmha, + kv_factor=self.kv_factor, + total_num_blocks=self._get_total_num_blocks(metadata), + ) + + sequence_length = metadata.kv_lens_cuda_runtime + host_past_key_value_lengths = metadata.kv_lens_runtime + + if num_contexts > 0 and attention_input_type != AttentionInputType.generation_only: + seq_offset = 0 + token_offset = 0 + num_seqs = num_contexts + + context_lengths = metadata.prompt_lens_cuda_runtime + host_context_lengths = metadata.prompt_lens_cpu_runtime + max_context_q_len = int(host_context_lengths[seq_offset : seq_offset + num_seqs].max()) + max_past_kv_len = int( + host_past_key_value_lengths[seq_offset : seq_offset + num_seqs].max() + ) + + params.attention_input = q[token_offset : token_offset + num_ctx_tokens] + params.qkv_input = params.attention_input + params.context_buf = out_tensor[token_offset : token_offset + num_ctx_tokens] + params.sequence_lengths = sequence_length[seq_offset:] + params.context_lengths = context_lengths[seq_offset:] + params.max_past_kv_length = max_past_kv_len + params.num_tokens = num_ctx_tokens + params.seq_offset = seq_offset + params.input_seq_length = max_context_q_len + params.batch_size = num_seqs + if attn.is_mla_enable: + self.run_mla_context(params) + else: + self.run_context(params) + + if num_generations > 0 and attention_input_type != AttentionInputType.context_only: + seq_offset = num_contexts + token_offset = 0 if is_gen_only else num_ctx_tokens + num_seqs = num_generations + + max_past_kv_len = int( + host_past_key_value_lengths[seq_offset : seq_offset + num_seqs].max() + ) + input_seq_length = num_gen_tokens // num_seqs if num_seqs > 0 else 1 + + predicted_tokens_per_seq = attn.predicted_tokens_per_seq + spec_gen_lengths = None + spec_pos_offsets = None + if metadata.is_spec_decoding_enabled and predicted_tokens_per_seq > 1: + spec_gen_lengths = metadata.spec_decoding_generation_lengths + position_offsets_for_cpp = metadata.spec_decoding_position_offsets_for_cpp + if position_offsets_for_cpp is not None and position_offsets_for_cpp.dim() == 1: + position_offsets_for_cpp = position_offsets_for_cpp.view( + metadata.max_num_requests, -1 + ) + spec_pos_offsets = position_offsets_for_cpp + + params.attention_input = q[token_offset : token_offset + num_gen_tokens] + params.qkv_input = params.attention_input + params.context_buf = out_tensor[token_offset : token_offset + num_gen_tokens] + params.sequence_lengths = sequence_length[seq_offset:] + params.max_past_kv_length = max_past_kv_len + params.num_tokens = num_gen_tokens + params.seq_offset = seq_offset + params.input_seq_length = input_seq_length + params.num_requests = num_seqs // metadata.beam_width + params.spec_decoding_generation_lengths = spec_gen_lengths + params.spec_decoding_position_offsets = spec_pos_offsets + if attn.is_mla_enable: + self.run_mla_generation(params) + else: + self.run_generation(params) + + def run_context(self, params: FmhaParams) -> None: + raise NotImplementedError(f"{type(self).__name__} does not support context attention.") + + def run_generation(self, params: FmhaParams) -> None: + raise NotImplementedError(f"{type(self).__name__} does not support generation attention.") + + def run_mla_context(self, params: FmhaParams) -> None: + raise NotImplementedError(f"{type(self).__name__} does not support MLA context attention.") + + def run_mla_generation(self, params: FmhaParams) -> None: + raise NotImplementedError( + f"{type(self).__name__} does not support MLA generation attention." + ) diff --git a/tensorrt_llm/_torch/attention_backend/fmha/registry.py b/tensorrt_llm/_torch/attention_backend/fmha/registry.py new file mode 100644 index 000000000000..97467657e9b5 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/fmha/registry.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from typing import TypeAlias + +from .fallback import FallbackFmha +from .flashinfer_trtllm_gen import FlashInferTrtllmGenFmha +from .interface import Fmha + +FmhaCls: TypeAlias = type[Fmha] + +FMHA_LIBS: dict[str, FmhaCls] = { + "flashinfer_trtllm_gen": FlashInferTrtllmGenFmha, + "fallback": FallbackFmha, +} +DEFAULT_FMHA_LIBS: tuple[str, ...] = tuple(FMHA_LIBS) + + +def _parse_fmha_libs_env() -> tuple[str, ...]: + value = os.environ.get("TLLM_FMHA_LIBS") + if value is None or not value.strip(): + return DEFAULT_FMHA_LIBS + + tokens = [token.strip() for token in value.split(",") if token.strip()] + if not tokens: + return DEFAULT_FMHA_LIBS + + has_delta_token = any(token[0] in "+-" for token in tokens) + if has_delta_token and not all(token[0] in "+-" for token in tokens): + raise ValueError( + "TLLM_FMHA_LIBS must use either an exact comma-separated list " + "or only +/- delta entries." + ) + + if has_delta_token: + names = list(DEFAULT_FMHA_LIBS) + for token in tokens: + sign = token[0] + name = token[1:].strip() + if not name: + raise ValueError(f"Invalid empty FMHA library entry in {value!r}.") + if name not in FMHA_LIBS: + raise ValueError(f"Unknown FMHA library {name!r} in TLLM_FMHA_LIBS.") + if sign == "+" and name not in names: + names.append(name) + elif sign == "-" and name in names: + names.remove(name) + else: + names = [] + for name in tokens: + if name not in FMHA_LIBS: + raise ValueError(f"Unknown FMHA library {name!r} in TLLM_FMHA_LIBS.") + if name not in names: + names.append(name) + + return tuple(names) + + +def get_enabled_fmha_lib_classes() -> list[FmhaCls]: + return [FMHA_LIBS[name] for name in _parse_fmha_libs_env()] + + +__all__ = [ + "DEFAULT_FMHA_LIBS", + "FMHA_LIBS", + "FmhaCls", + "get_enabled_fmha_lib_classes", +] diff --git a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md index a30849da0fa9..4e98b8d688a4 100644 --- a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md @@ -243,13 +243,32 @@ The main differences across backends: | `VANILLA` | Python-side | Python-side slicing | | `FLASHINFER` | Python-side (explicit append) | Page-table metadata | -#### 3.2.2 `TRTLLM` internal `trtllm_gen` path +#### 3.2.2 `TRTLLM` internal FMHA libraries -`trtllm_gen.py` integrates trtllm-gen kernels from FlashInfer into the -`TRTLLM` backend. It is not a separate backend. It is an internal fast path -disabled by default (`TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION`). It bridges the -TRTLLM block-offset format into the page-table shape expected by those kernels. -If it does not apply, `TrtllmAttention` stays on its regular runtime path. +`TrtllmAttention` dispatches attention through an ordered list of internal FMHA +libraries. `FlashInferTrtllmGenFmha` integrates trtllm-gen kernels from +FlashInfer into the `TRTLLM` backend, and `FallbackFmha` calls the regular +`thop.attention` runtime path. These are not separate attention backends. + +`TLLM_FMHA_LIBS` controls the ordered list. Unset means +`flashinfer_trtllm_gen,fallback`; use `TLLM_FMHA_LIBS=fallback` or +`TLLM_FMHA_LIBS=-flashinfer_trtllm_gen` to force the fallback path. Each FMHA +library exposes `is_available()` for module/static environment checks and +`is_supported()` for per-forward request checks. + +The FMHA package is split by role: + +- `fmha/interface.py` defines the `Fmha` runtime contract. +- `fmha/phased.py` defines `PhasedFmha`, which handles mixed context/generation + requests and dispatches them to phase-specific hooks. +- `fmha/flashinfer_trtllm_gen.py` implements the FlashInfer trtllm-gen FMHA + library. +- `fmha/fallback.py` implements the regular `thop.attention` fallback library. +- `fmha/registry.py` owns `TLLM_FMHA_LIBS` parsing and library ordering. + +Use `PhasedFmha` for libraries that need separate context/generation or MHA/MLA +entry points. Use `Fmha` directly for libraries that already own the full +request shape. #### 3.2.3 MLA cached-context semantics @@ -338,7 +357,7 @@ Working rules: | `tensorrt_llm/_torch/attention_backend/interface.py` | Backend contract, base metadata, capability hooks | | `tensorrt_llm/_torch/attention_backend/utils.py` | Backend and sparse-backend selection | | `tensorrt_llm/_torch/attention_backend/trtllm.py` | TRTLLM backend and metadata | -| `tensorrt_llm/_torch/attention_backend/trtllm_gen.py` | Internal dense fast path | +| `tensorrt_llm/_torch/attention_backend/fmha/` | Internal TRTLLM FMHA libraries | | `tensorrt_llm/_torch/attention_backend/vanilla.py` | Torch fallback backend and metadata | | `tensorrt_llm/_torch/attention_backend/flashinfer.py` | FlashInfer backend and metadata | | `tensorrt_llm/_torch/attention_backend/sparse/` | DSA, Rocket sparse backends, metadata, cache managers | diff --git a/tests/unittest/_torch/attention_backend/test_attention_op_sync.py b/tests/unittest/_torch/attention_backend/test_attention_op_sync.py index 284cc46bc909..0a91900c172a 100644 --- a/tests/unittest/_torch/attention_backend/test_attention_op_sync.py +++ b/tests/unittest/_torch/attention_backend/test_attention_op_sync.py @@ -11,8 +11,8 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Static sync test for the inline ``thop.attention(...)`` call in -``TrtllmAttention._run``. +"""Static sync test for the fallback ``thop.attention(...)`` call in +``FallbackFmha.forward``. That call is the single explicit-kwarg call site for the C++ ``thop.attention`` binding. This test parses both the call site (Python AST) and the C++ @@ -20,7 +20,7 @@ 1. Every C++ parameter name appears at the call site (and nothing extra). 2. Every call-site kwarg sourced as ``root.attr[.attr...]`` resolves on - exactly one of ``self`` / ``metadata`` / ``forward_args``, and its + exactly one of ``attn`` / ``metadata`` / ``forward_args``, and its declared C++ type matches the source attribute's Python type at a coarse-category level (tensor / int / bool / float / list-of-X). 3. Every dataclass field reachable from ``AttentionForwardArgs`` (including @@ -42,27 +42,28 @@ import typing from dataclasses import fields -from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs -from tensorrt_llm._torch.attention_backend.sparse.skip_softmax import SkipSoftmaxKernelParams -from tensorrt_llm._torch.attention_backend.trtllm import ( +from tensorrt_llm._torch.attention_backend.fmha.fallback import ( _THOP_EXCLUDED_FIELDS, _THOP_LITERALS, - TrtllmAttention, - TrtllmAttentionMetadata, + FallbackFmha, ) +from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs +from tensorrt_llm._torch.attention_backend.sparse.skip_softmax import SkipSoftmaxKernelParams +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention, TrtllmAttentionMetadata # Roots used as the LHS of attribute chains at the call site. Match the -# parameter names inside ``TrtllmAttention._run``. +# names inside ``FallbackFmha.forward``. _SOURCE_CLASSES = { - "self": TrtllmAttention, + "attn": TrtllmAttention, "metadata": TrtllmAttentionMetadata, "forward_args": AttentionForwardArgs, "skip_softmax_kernel_params": SkipSoftmaxKernelParams, } _THOP_KWARG_SOURCE_ALIASES: dict[str, tuple[str, tuple[str, ...]]] = { + "beam_width": ("metadata", ("effective_beam_width",)), "context_lengths": ("metadata", ("prompt_lens_cuda_runtime",)), - "head_size": ("self", ("head_dim",)), + "head_size": ("attn", ("head_dim",)), "host_context_lengths": ("metadata", ("prompt_lens_cpu_runtime",)), "host_past_key_value_lengths": ("metadata", ("kv_lens_runtime",)), "host_request_types": ("metadata", ("host_request_types_runtime",)), @@ -199,8 +200,9 @@ def _binding_types() -> dict[str, str]: def _parse_thop_attention_call() -> ast.Call: - """Locate the single ``thop.attention(...)`` call inside ``_run``.""" - src = textwrap.dedent(inspect.getsource(TrtllmAttention._run)) + """Locate the single ``thop.attention(...)`` call inside + ``FallbackFmha.forward``.""" + src = textwrap.dedent(inspect.getsource(FallbackFmha.forward)) tree = ast.parse(src) for node in ast.walk(tree): if ( @@ -211,7 +213,7 @@ def _parse_thop_attention_call() -> ast.Call: and node.func.value.id == "thop" ): return node - raise AssertionError("Could not find thop.attention(...) call in TrtllmAttention._run") + raise AssertionError("Could not find thop.attention(...) call in FallbackFmha.forward") def _attribute_path(node: ast.AST) -> tuple[str, tuple[str, ...]] | None: @@ -504,8 +506,9 @@ def _self_attrs_in_property(prop: property) -> set[str]: def _collect_chains(root: str) -> set[tuple[str, ...]]: - """All attribute paths in ``_run`` that start with ``Name(root).``.""" - src = textwrap.dedent(inspect.getsource(TrtllmAttention._run)) + """All attribute paths in ``FallbackFmha.forward`` that start with + ``Name(root).``.""" + src = textwrap.dedent(inspect.getsource(FallbackFmha.forward)) chains: set[tuple[str, ...]] = set() for node in ast.walk(ast.parse(src)): if not isinstance(node, ast.Attribute): @@ -568,7 +571,7 @@ def test_every_forward_args_field_is_consumed(): def test_no_unexpected_other_kwargs(): """The only call-site kwargs that aren't ``source.attr`` chains or - allowlisted literals are the ``_run`` positional parameters.""" + allowlisted literals are the ``FallbackFmha.forward`` parameters.""" _, _, other_kwargs = _classify_kwargs() expected = {"q", "k", "v"} unexpected = other_kwargs - expected From 3b62ab5cedc65249de74ec784435ce266d513bc0 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 10 Jun 2026 10:49:32 +0000 Subject: [PATCH 2/8] [TRTLLM-12807][chore] Address FMHA review comments Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/attention_backend/fmha/flashinfer_trtllm_gen.py | 3 --- tensorrt_llm/_torch/attention_backend/fmha/phased.py | 4 ++-- 2 files changed, 2 insertions(+), 5 deletions(-) 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 dc4d382e0db0..70f6eccc29bb 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py @@ -359,7 +359,6 @@ class FlashInferTrtllmGenFmha(PhasedFmha): # Default KV layout for flashinfer # HND = [max_num_pages, kv_factor, num_kv_heads, page_size, head_dim] DEFAULT_KV_LAYOUT = "HND" - REQUIRES_PAGED_KV = True # Keep shared paged indices disabled to match the current TensorRT-LLM # block-table layout used by the fused preprocessing path. USE_SHARED_PAGED_KV_IDX = False @@ -767,7 +766,6 @@ def get_fp8_context_fmha( forward_args: AttentionForwardArgs, is_gen_only: bool, ) -> bool: - del q, metadata, forward_args kv_cache_quant_mode = QuantMode(self.attn.quant_mode) return ( output.dtype == torch.float8_e4m3fn @@ -785,7 +783,6 @@ def prepare_workspace( forward_args: AttentionForwardArgs, workspace: torch.Tensor, ) -> None: - del k, v attn = self.attn # Lazily cache the SM count from the first query tensor's device. if self._multi_processor_count is None: diff --git a/tensorrt_llm/_torch/attention_backend/fmha/phased.py b/tensorrt_llm/_torch/attention_backend/fmha/phased.py index 43a7cb2aed0a..aced12909cc9 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/phased.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/phased.py @@ -59,9 +59,9 @@ class FmhaParams: class PhasedFmha(Fmha): - """FMHA base that dispatches mixed requests by phase and attention family.""" + """FMHA helper for paged-KV libraries that split work by request phase.""" - REQUIRES_PAGED_KV = False + REQUIRES_PAGED_KV = True def __init__(self, attn: "TrtllmAttention"): super().__init__(attn) From e2b1ece7c02fa7e27aaf94dba932df88106654d5 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Thu, 11 Jun 2026 07:47:21 +0000 Subject: [PATCH 3/8] [TRTLLM-12807][fix] Resolve FMHA rebase conflicts Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/fmha/fallback.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py index 384f4b8e59e5..73c5778379d8 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py @@ -38,7 +38,7 @@ { "topk_indices", # DSA-only "attention_mask_data", # custom-mask code path - "out_scale_sf", # promoted into ``out_scale`` in ``TrtllmAttention._run`` for NVFP4 path + "out_scale_sf", # promoted into ``out_scale`` in ``TrtllmAttention.forward`` for NVFP4 path } ) From c16824c666013585c1e81d7324dc39ffa248878f Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Thu, 11 Jun 2026 09:59:07 +0000 Subject: [PATCH 4/8] [TRTLLM-12807][fix] Store FMHA attention owner weakly Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/attention_backend/fmha/interface.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention_backend/fmha/interface.py b/tensorrt_llm/_torch/attention_backend/fmha/interface.py index 0e40e753d2da..f14f26964b25 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/interface.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/interface.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import weakref from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Optional @@ -31,7 +32,14 @@ class Fmha(ABC): """Common runtime contract for TRT-LLM attention FMHA libraries.""" def __init__(self, attn: "TrtllmAttention"): - self.attn = attn + self._attn_ref: weakref.ReferenceType["TrtllmAttention"] = weakref.ref(attn) + + @property + def attn(self) -> "TrtllmAttention": + attn = self._attn_ref() + if attn is None: + raise RuntimeError("The owning TrtllmAttention instance has been garbage collected.") + return attn @classmethod def is_available(cls, attn: "TrtllmAttention") -> bool: From 888fc48c3f04de133437c76b9ebf2203f2cc2dc4 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Mon, 15 Jun 2026 05:07:57 +0000 Subject: [PATCH 5/8] [TRTLLM-12807][fix] Restore FMHA dispatch after rebase Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/attention_backend/trtllm.py | 428 +++++------------- 1 file changed, 119 insertions(+), 309 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 7ada30a7c708..47f1831722b4 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -11,7 +11,8 @@ from ..speculative.interface import SpecMetadata from ..speculative.spec_tree_manager import SpecTreeManager -from tensorrt_llm._torch.attention_backend import trtllm_gen +from tensorrt_llm._torch.attention_backend.fmha import ( + Fmha, get_enabled_fmha_lib_classes) 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 @@ -25,27 +26,7 @@ PredefinedAttentionMask, RopeParams, SparsePrediction, merge_attention_forward_args) from .sparse.params import SparseParams -from .sparse.skip_softmax import SkipSoftmaxKernelParams, SkipSoftmaxParams - -# Enable TRTLLM-Gen attention backend by default. Set -# TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION=0 to force the thop.attention path. -_TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION = (os.environ.get( - "TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION", "1") == "1") - -# ``AttentionForwardArgs`` fields that this backend does not consume. -# Sync test (test_attention_op_sync.py) requires every other field to map to a -# kwarg name, a @property on the dataclass, or a field that some @property -# transitively reads; entries here are exempt. -_THOP_EXCLUDED_FIELDS: frozenset = frozenset({ - "topk_indices", # DSA-only - "attention_mask_data", # custom-mask code path - "out_scale_sf", # promoted into ``out_scale`` in ``_run`` for NVFP4 path -}) - -# ``thop.attention`` kwargs hard-wired to a literal at the call site (no -# rich object owns them). Sync test enforces both the kwarg name and the -# literal value. -_THOP_LITERALS: dict = {} +from .sparse.skip_softmax import SkipSoftmaxParams @functools.cache @@ -124,6 +105,7 @@ class TrtllmAttentionMetadata(AttentionMetadata): spec_decoding_bl_tree_mask_offset: Optional[torch.Tensor] = None spec_decoding_bl_tree_mask: Optional[torch.Tensor] = None spec_bl_tree_first_sparse_mask_offset_kv: Optional[torch.Tensor] = None + max_total_draft_tokens: Optional[int] = None # TRTLLM-Gen FMHA JIT warmup controls. trtllm_gen_jit_warmup: bool = False @@ -1213,6 +1195,7 @@ def __init__( self.kv_scale_orig_quant = 1.0 / self.kv_cache_scaling_factor self.local_layer_idx: Optional[int] = None + self.fmha_libs: List[Fmha] = [] if not skip_create_weights_in_init: self.update_quant_config(self.quant_config) @@ -1235,6 +1218,7 @@ def update_quant_config(self, new_quant_config: Optional[QuantConfig]): self.has_nvfp4 = self.quant_config.layer_quant_mode.has_nvfp4() self.has_w4a8_nvfp4_fp8 = self.quant_config.layer_quant_mode.has_w4a8_nvfp4_fp8( ) + self.create_fmha_libs() def get_local_layer_idx(self, metadata: TrtllmAttentionMetadata) -> int: if self.local_layer_idx is not None: @@ -1417,23 +1401,114 @@ def rope_max_positions(self) -> int: def rope_original_max_positions(self) -> int: return self.rope_params.original_max_positions - def _get_trtllm_gen_backend( - self) -> trtllm_gen.FlashInferTrtllmGenAttention: - backend = getattr(self, "_trtllm_gen_backend", None) - if backend is None: - backend = trtllm_gen.FlashInferTrtllmGenAttention( - attention_layer=self, ) - self._trtllm_gen_backend = backend - return backend - def _run( + def create_fmha_libs(self) -> None: + self.fmha_libs = [] + for fmha_cls in get_enabled_fmha_lib_classes(): + if fmha_cls.is_available(self): + self.fmha_libs.append(fmha_cls(self)) + + + def forward( self, q: torch.Tensor, k: Optional[torch.Tensor], v: Optional[torch.Tensor], metadata: TrtllmAttentionMetadata, - forward_args: AttentionForwardArgs, - ) -> None: + forward_args: Optional[AttentionForwardArgs] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Execute the TRTLLM attention backend.""" + forward_args = merge_attention_forward_args(forward_args, kwargs) + assert isinstance( + metadata, + TrtllmAttentionMetadata, + ) + # Cross-attention uses the THOP path; the trtllm-gen backend API does + # not carry encoder K/V tensors yet. + + if forward_args.multi_item_part_lens is not None: + raise ValueError( + "TRT-LLM Attention does not support multi-item scoring") + + # SM90 forces ``use_paged_context_fmha`` on for correctness + # (https://nvbugs/5624818). + if get_sm_version() == 90: + metadata.use_paged_context_fmha = True + + # Sparse mqa/gqa attention uses generation kernel which reads Q from qPtr (separate buffer). + # Force paged context FMHA so QKV preprocessing writes Q to q_buf_2_. + if (self.sparse_params is not None and getattr( + self.sparse_params, 'algorithm', None) == 'mqa_gqa'): + metadata.use_paged_context_fmha = True + + if self.is_mla_enable: + # Context MLA uses separate qkv instead of paged_context_fmha + metadata.use_paged_context_fmha = False + + if forward_args.output is None: + is_gen_only = (forward_args.attention_input_type == + AttentionInputType.generation_only) + outputs = self.create_output( + q, + is_quantize_output=forward_args.out_scale is not None, + metadata=metadata, + attention_mask=forward_args.attention_mask, + use_paged_context_fmha=metadata.use_paged_context_fmha, + is_mla_enable=self.is_mla_enable, + is_gen_only=is_gen_only, + ) + forward_args.output = outputs[0] + forward_args.output_sf = outputs[1] if len(outputs) == 2 else None + + forward_args.is_fused_qkv = not metadata.is_cross and k is None + forward_args.update_kv_cache = not metadata.is_cross or k is not None + has_fused_qkv = forward_args.is_fused_qkv and k is None and v is None + has_unfused_kv = (not forward_args.is_fused_qkv and k is not None + and v is not None) + uses_cached_cross_kv = (metadata.is_cross + and not forward_args.update_kv_cache + and k is None and v is None) + assert has_fused_qkv or has_unfused_kv or uses_cached_cross_kv + if forward_args.cu_q_seqlens is None: + forward_args.cu_q_seqlens = metadata.cu_q_seqlens + if forward_args.cu_kv_seqlens is None: + forward_args.cu_kv_seqlens = metadata.cu_kv_seqlens + + # RocketKV and DSA predict which blocks to keep, so build their sparse + # index tensors here. Skip-softmax needs no prediction. + sparse_params = self.sparse_params + if (sparse_params is not None + and not isinstance(sparse_params, SkipSoftmaxParams)): + kv_idx, kv_off = self.sparse_kv_predict(q, k, metadata, + forward_args) + at_idx, at_off = self.sparse_attn_predict(q, k, metadata, + forward_args) + forward_args.sparse_prediction = SparsePrediction( + sparse_kv_indices=kv_idx, + sparse_kv_offsets=kv_off, + sparse_attn_indices=at_idx, + sparse_attn_offsets=at_off, + sparse_attn_indices_block_size=sparse_params.indices_block_size, + ) + + # Compute FlashMLA tile-scheduler metadata once per forward pass. + # The flag is reset in prepare_flash_mla() and update_for_spec_dec() to trigger + # recomputation when cache_seq_lens change. The metadata must always match the + # compacted generation sub-batch, which is also the layout used by block_ids_per_seq. + if (metadata.enable_flash_mla and forward_args.attention_input_type + != AttentionInputType.context_only + and metadata.num_generations > 0 + and not metadata._flash_mla_metadata_valid): + self._compute_flash_mla_metadata(metadata) + metadata._flash_mla_metadata_valid = True + + # Blackwell first_sparse: refresh at layer 0 before kernel launch. + if self.get_local_layer_idx(metadata) == 0 and ( + metadata.spec_bl_tree_first_sparse_mask_offset_kv is not None + and metadata._seq_lens_cuda is not None): + metadata.update_blackwell_first_sparse_mask_offset() + if metadata.is_cross: if k is not None and v is not None: k_flat = k.contiguous().view(k.shape[0], -1) @@ -1501,8 +1576,8 @@ def _run( self._ensure_rope_table_size(metadata.max_seq_len) - # Prime ``self.local_layer_idx`` so the ``thop.attention`` kwarg - # below reads a populated int rather than the ``None`` placeholder. + # Prime ``self.local_layer_idx`` so FMHA implementations read a + # populated int rather than the ``None`` placeholder. # The call is a fast cache hit after the first forward. self.local_layer_idx = self.get_local_layer_idx(metadata) if metadata.spec_decoding_bl_tree_mask is not None and self.local_layer_idx == 0: @@ -1551,179 +1626,16 @@ def _run( assert metadata.kv_cache_manager is None assert metadata.num_contexts == metadata.num_seqs - use_trtllm_gen = False - if _TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION: - trtllm_gen_backend = self._get_trtllm_gen_backend() - use_trtllm_gen = trtllm_gen_backend.is_supported( - q, - k, - v, - attn=self, - meta=metadata, - fwd=forward_args, - )[0] - - if use_trtllm_gen: - trtllm_gen_backend.forward( - q, - k, - v, - attn=self, - meta=metadata, - fwd=forward_args, - ) + if not self.fmha_libs: + self.create_fmha_libs() + + for fmha in self.fmha_libs: + if fmha.is_supported(q, k, v, metadata, forward_args): + fmha.forward(q, k, v, metadata, forward_args) + break else: - sparse_params = self.sparse_params - skip_softmax_kernel_params = ( - sparse_params.scheduler.get_kernel_params( - timestep=forward_args.timestep) if isinstance( - sparse_params, - SkipSoftmaxParams) else SkipSoftmaxKernelParams()) - - # Every kwarg sources from ``self`` / ``metadata`` / - # ``forward_args`` (with ``forward_args.sparse_prediction`` for - # sparse-attn inputs), ``skip_softmax_kernel_params``, or a literal - # allowlisted in ``_THOP_LITERALS``. ``test_attention_op_sync.py`` - # enforces this statically. - thop.attention( - q=q, - k=k, - v=v, - output=forward_args.output, - output_sf=forward_args.output_sf, - workspace_=metadata.effective_workspace, - - # --- Per-step batch state (TrtllmAttentionMetadata) --- - sequence_length=metadata.kv_lens_cuda_runtime, - host_past_key_value_lengths=metadata.kv_lens_runtime, - host_total_kv_lens=metadata.host_total_kv_lens, - context_lengths=metadata.prompt_lens_cuda_runtime, - host_context_lengths=metadata.prompt_lens_cpu_runtime, - host_request_types=metadata.host_request_types_runtime, - max_context_q_len_override=metadata.max_context_q_len_override, - kv_cache_block_offsets=metadata.kv_cache_block_offsets, - host_kv_cache_pool_pointers=metadata. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=metadata.host_kv_cache_pool_mapping, - cache_indirection=metadata.cache_indirection, - block_ids_per_seq=metadata.block_ids_per_seq, - tokens_per_block=metadata.tokens_per_block, - max_num_requests=metadata.max_num_requests, - beam_width=metadata.effective_beam_width, - use_paged_context_fmha=metadata.use_paged_context_fmha, - helix_position_offsets=metadata.helix_position_offsets, - helix_is_inactive_rank=metadata.helix_is_inactive_rank, - is_spec_decoding_enabled=metadata.is_spec_decoding_enabled, - use_spec_decoding=metadata.use_spec_decoding, - is_spec_dec_tree=metadata.is_spec_dec_tree, - spec_decoding_generation_lengths=metadata. - spec_decoding_generation_lengths, - spec_decoding_position_offsets_for_cpp=metadata. - spec_decoding_position_offsets_for_cpp, - spec_decoding_packed_mask=metadata.spec_decoding_packed_mask, - spec_decoding_bl_tree_mask_offset=metadata. - spec_decoding_bl_tree_mask_offset, - spec_decoding_bl_tree_mask=metadata.spec_decoding_bl_tree_mask, - spec_decoding_target_max_draft_tokens=metadata. - max_total_draft_tokens, - spec_bl_tree_first_sparse_mask_offset_kv=metadata. - spec_bl_tree_first_sparse_mask_offset_kv, - num_sparse_topk=metadata.num_sparse_topk, - flash_mla_tile_scheduler_metadata=metadata. - flash_mla_tile_scheduler_metadata, - flash_mla_num_splits=metadata.flash_mla_num_splits, - num_contexts=metadata.num_contexts, - num_ctx_tokens=metadata.num_ctx_tokens, - max_context_length=metadata.max_context_length, - max_seq_len=metadata.max_seq_len, - trtllm_gen_jit_warmup=metadata.trtllm_gen_jit_warmup, - is_cross=metadata.is_cross, - - # --- Per-call (AttentionForwardArgs) --- - out_scale=forward_args.out_scale, - kv_scale_orig_quant=forward_args.kv_scale_orig_quant, - kv_scale_quant_orig=forward_args.kv_scale_quant_orig, - latent_cache=forward_args.latent_cache, - q_pe=forward_args.q_pe, - attention_sinks=forward_args.attention_sinks, - mask_type=forward_args.mask_type, - attention_input_type=int(forward_args.attention_input_type), - attention_window_size=forward_args.attention_window_size, - chunked_prefill_buffer_batch_size=forward_args. - chunked_prefill_buffer_batch_size, - mrope_rotary_cos_sin=forward_args.mrope_rotary_cos_sin, - mrope_position_deltas=forward_args.mrope_position_deltas, - softmax_stats_tensor=forward_args.softmax_stats_tensor, - cu_q_seqlens=forward_args.cu_q_seqlens, - cu_kv_seqlens=forward_args.cu_kv_seqlens, - fmha_scheduler_counter=forward_args.fmha_scheduler_counter, - mla_bmm1_scale=forward_args.mla_bmm1_scale, - mla_bmm2_scale=forward_args.mla_bmm2_scale, - quant_q_buffer=forward_args.quant_q_buffer, - sage_attn_num_elts_per_blk_q=forward_args. - sage_attn_num_elts_per_blk_q, - sage_attn_num_elts_per_blk_k=forward_args. - sage_attn_num_elts_per_blk_k, - sage_attn_num_elts_per_blk_v=forward_args. - sage_attn_num_elts_per_blk_v, - sage_attn_qk_int8=forward_args.sage_attn_qk_int8, - is_fused_qkv=forward_args.is_fused_qkv, - update_kv_cache=forward_args.update_kv_cache, - cross_kv=forward_args.cross_kv, - relative_attention_bias=forward_args.relative_attention_bias, - relative_attention_max_distance=forward_args. - relative_attention_max_distance, - position_embedding_type=self.position_embedding_type, - # --- Module config (TrtllmAttention) --- - rotary_inv_freq=self.rotary_inv_freq, - rotary_cos_sin=self.rotary_cos_sin, - predicted_tokens_per_seq=self.predicted_tokens_per_seq, - local_layer_idx=self.local_layer_idx, - num_heads=self.num_heads, - num_kv_heads=self.num_kv_heads, - head_size=self.head_dim, - quant_mode=self.quant_mode, - q_scaling=self.q_scaling, - rope_dim=self.rope_dim, - rope_base=self.rope_base, - rope_scale_type=self.rope_scale_type, - rope_scale=self.rope_scale, - rope_short_m_scale=self.rope_short_m_scale, - rope_long_m_scale=self.rope_long_m_scale, - rope_max_positions=self.rope_max_positions, - rope_original_max_positions=self.rope_original_max_positions, - is_mla_enable=self.is_mla_enable, - q_lora_rank=self.q_lora_rank, - kv_lora_rank=self.kv_lora_rank, - qk_nope_head_dim=self.qk_nope_head_dim, - qk_rope_head_dim=self.qk_rope_head_dim, - v_head_dim=self.v_head_dim, - rope_append=self.rope_append, - attention_chunk_size=self.attention_chunk_size, - skip_softmax_threshold_scale_factor_prefill= - skip_softmax_kernel_params.threshold_scale_factor_prefill, - skip_softmax_threshold_scale_factor_decode= - skip_softmax_kernel_params.threshold_scale_factor_decode, - skip_softmax_stat=self.skip_softmax_stat, - - # --- Sparse-specific (AttentionForwardArgs.sparse_prediction) --- - sparse_kv_indices=forward_args.sparse_prediction. - sparse_kv_indices, - sparse_kv_offsets=forward_args.sparse_prediction. - sparse_kv_offsets, - sparse_attn_indices=forward_args.sparse_prediction. - sparse_attn_indices, - sparse_attn_offsets=forward_args.sparse_prediction. - sparse_attn_offsets, - sparse_attn_indices_block_size=forward_args.sparse_prediction. - sparse_attn_indices_block_size, - sparse_mla_topk_lens=forward_args.sparse_prediction. - sparse_mla_topk_lens, - compressed_kv_cache_pool_ptr=forward_args.sparse_prediction. - compressed_kv_cache_pool_ptr, - - # --- Literals intentionally in _THOP_LITERALS --- - ) + raise RuntimeError( + "No TRT-LLM attention FMHA library supports this request.") if self.print_skip_softmax_stat: total_blocks, skipped_blocks = self.skip_softmax_stat @@ -1732,108 +1644,6 @@ def _run( f"SKIP_SOFTMAX_STAT: layer{self.layer_idx}: {skipped_blocks} / {total_blocks}" f" = {skipped_blocks / total_blocks * 100: .2f}%") - def forward( - self, - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - metadata: TrtllmAttentionMetadata, - forward_args: Optional[AttentionForwardArgs] = None, - **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - """Execute the TRTLLM attention backend.""" - forward_args = merge_attention_forward_args(forward_args, kwargs) - assert isinstance( - metadata, - TrtllmAttentionMetadata, - ) - # Cross-attention uses the THOP path; the trtllm-gen backend API does - # not carry encoder K/V tensors yet. - - if forward_args.multi_item_part_lens is not None: - raise ValueError( - "TRT-LLM Attention does not support multi-item scoring") - - # SM90 forces ``use_paged_context_fmha`` on for correctness - # (https://nvbugs/5624818). - if get_sm_version() == 90: - metadata.use_paged_context_fmha = True - - # Sparse mqa/gqa attention uses generation kernel which reads Q from qPtr (separate buffer). - # Force paged context FMHA so QKV preprocessing writes Q to q_buf_2_. - if (self.sparse_params is not None and getattr( - self.sparse_params, 'algorithm', None) == 'mqa_gqa'): - metadata.use_paged_context_fmha = True - - if self.is_mla_enable: - # Context MLA uses separate qkv instead of paged_context_fmha - metadata.use_paged_context_fmha = False - - if forward_args.output is None: - is_gen_only = (forward_args.attention_input_type == - AttentionInputType.generation_only) - outputs = self.create_output( - q, - is_quantize_output=forward_args.out_scale is not None, - metadata=metadata, - attention_mask=forward_args.attention_mask, - use_paged_context_fmha=metadata.use_paged_context_fmha, - is_mla_enable=self.is_mla_enable, - is_gen_only=is_gen_only, - ) - forward_args.output = outputs[0] - forward_args.output_sf = outputs[1] if len(outputs) == 2 else None - - forward_args.is_fused_qkv = not metadata.is_cross and k is None - forward_args.update_kv_cache = not metadata.is_cross or k is not None - has_fused_qkv = forward_args.is_fused_qkv and k is None and v is None - has_unfused_kv = (not forward_args.is_fused_qkv and k is not None - and v is not None) - uses_cached_cross_kv = (metadata.is_cross - and not forward_args.update_kv_cache - and k is None and v is None) - assert has_fused_qkv or has_unfused_kv or uses_cached_cross_kv - if forward_args.cu_q_seqlens is None: - forward_args.cu_q_seqlens = metadata.cu_q_seqlens - if forward_args.cu_kv_seqlens is None: - forward_args.cu_kv_seqlens = metadata.cu_kv_seqlens - - # RocketKV and DSA predict which blocks to keep, so build their sparse - # index tensors here. Skip-softmax needs no prediction. - sparse_params = self.sparse_params - if (sparse_params is not None - and not isinstance(sparse_params, SkipSoftmaxParams)): - kv_idx, kv_off = self.sparse_kv_predict(q, k, metadata, - forward_args) - at_idx, at_off = self.sparse_attn_predict(q, k, metadata, - forward_args) - forward_args.sparse_prediction = SparsePrediction( - sparse_kv_indices=kv_idx, - sparse_kv_offsets=kv_off, - sparse_attn_indices=at_idx, - sparse_attn_offsets=at_off, - sparse_attn_indices_block_size=sparse_params.indices_block_size, - ) - - # Compute FlashMLA tile-scheduler metadata once per forward pass. - # The flag is reset in prepare_flash_mla() and update_for_spec_dec() to trigger - # recomputation when cache_seq_lens change. The metadata must always match the - # compacted generation sub-batch, which is also the layout used by block_ids_per_seq. - if (metadata.enable_flash_mla and forward_args.attention_input_type - != AttentionInputType.context_only - and metadata.num_generations > 0 - and not metadata._flash_mla_metadata_valid): - self._compute_flash_mla_metadata(metadata) - metadata._flash_mla_metadata_valid = True - - # Blackwell first_sparse: refresh at layer 0 before kernel launch. - if self.get_local_layer_idx(metadata) == 0 and ( - metadata.spec_bl_tree_first_sparse_mask_offset_kv is not None - and metadata._seq_lens_cuda is not None): - metadata.update_blackwell_first_sparse_mask_offset() - - self._run(q, k, v, metadata, forward_args) - if forward_args.output_sf is None: return forward_args.output else: From 04083799d5f9f5b29792969f8f3cafcdb717bbe8 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 17 Jun 2026 03:16:11 +0000 Subject: [PATCH 6/8] [TRTLLM-12807][fix] Clean up FMHA rebase metadata Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/trtllm.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 47f1831722b4..b79b3731de5e 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import functools import math import os @@ -105,7 +120,6 @@ class TrtllmAttentionMetadata(AttentionMetadata): spec_decoding_bl_tree_mask_offset: Optional[torch.Tensor] = None spec_decoding_bl_tree_mask: Optional[torch.Tensor] = None spec_bl_tree_first_sparse_mask_offset_kv: Optional[torch.Tensor] = None - max_total_draft_tokens: Optional[int] = None # TRTLLM-Gen FMHA JIT warmup controls. trtllm_gen_jit_warmup: bool = False From 93cc65326bad7bc05337f0e143a51cf8441809da Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 17 Jun 2026 03:55:31 +0000 Subject: [PATCH 7/8] [TRTLLM-12807][fix] Resolve FMHA rebase fallout Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/trtllm.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index b79b3731de5e..93e1a2bfe4c4 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -1415,14 +1415,12 @@ def rope_max_positions(self) -> int: def rope_original_max_positions(self) -> int: return self.rope_params.original_max_positions - def create_fmha_libs(self) -> None: self.fmha_libs = [] for fmha_cls in get_enabled_fmha_lib_classes(): if fmha_cls.is_available(self): self.fmha_libs.append(fmha_cls(self)) - def forward( self, q: torch.Tensor, From 79d1e09e6a981259afd9f24a6f6a41071d700afc Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 17 Jun 2026 08:50:05 +0000 Subject: [PATCH 8/8] [TRTLLM-12807][fix] Fix FMHA CI regressions Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/attention_backend/fmha/flashinfer_trtllm_gen.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) 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 70f6eccc29bb..d880335a8b75 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py @@ -50,6 +50,7 @@ import flashinfer from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs, AttentionInputType +from tensorrt_llm._torch.attention_backend.sparse.skip_softmax import SkipSoftmaxParams from tensorrt_llm._utils import get_sm_version, is_sm_100f, torch_dtype_to_binding from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal import thop @@ -438,10 +439,7 @@ def is_available(cls, attn: "TrtllmAttention") -> bool: ) return False - has_skip_softmax = ( - attn.skip_softmax_threshold_scale_factor_prefill is not None - or attn.skip_softmax_threshold_scale_factor_decode is not None - ) + has_skip_softmax = isinstance(attn.sparse_params, SkipSoftmaxParams) if has_skip_softmax: logger.debug( "FlashInfer TRTLLM-Gen FMHA is unavailable: skip-softmax attention is enabled."