diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h index 3a757c3263d9..81fbe2d7b3ae 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h @@ -209,6 +209,8 @@ struct QKVPreprocessingParams bool separate_q_kv_output{false}; bool quantized_fp8_output{false}; bool generation_phase{false}; + // The input contains Q only and the existing KV cache must not be updated. + bool q_only_input{false}; int multi_processor_count{0}; int rotary_vision_start{0}; int rotary_vision_length{0}; @@ -308,6 +310,7 @@ struct QKVPreprocessingParams ss << "separate_q_kv_output: " << std::boolalpha << separate_q_kv_output << std::endl; ss << "quantized_fp8_output: " << quantized_fp8_output << std::endl; ss << "generation_phase: " << generation_phase << std::endl; + ss << "q_only_input: " << q_only_input << std::endl; ss << "multi_processor_count: " << multi_processor_count << std::endl; return ss.str(); diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h index 22e17b438d58..a158567d8765 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h @@ -405,7 +405,7 @@ __global__ void applyBiasRopeUpdateKVCache(QKVPreprocessingParams(global_token_idx) * hidden_size + hidden_idx; - auto const src_k_idx = static_cast(global_token_idx) * hidden_size + src_k_offset + hidden_idx_kv; - auto const src_v_idx = static_cast(global_token_idx) * hidden_size + src_v_offset + hidden_idx_kv; - - VecType q, k, v, q_pair, k_pair; + auto const src_q_idx = static_cast(global_token_idx) * input_hidden_size + hidden_idx; + auto const src_k_idx + = static_cast(global_token_idx) * input_hidden_size + src_k_offset + hidden_idx_kv; + auto const src_v_idx + = static_cast(global_token_idx) * input_hidden_size + src_v_offset + hidden_idx_kv; + + VecType q, q_pair; + VecType k{}, v{}, k_pair{}; // key without position embedding - VecType k_wo_pos; + VecType k_wo_pos{}; // load q,k,v and add bias if (valid_head_dim_idx) { q = *reinterpret_cast(¶ms.qkv_input[src_q_idx]); - k = *reinterpret_cast(¶ms.qkv_input[src_k_idx]); - v = *reinterpret_cast(¶ms.qkv_input[src_v_idx]); q_pair = *reinterpret_cast(¶ms.qkv_input[src_q_idx + rotated_head_dim_offset]); - k_pair = *reinterpret_cast(¶ms.qkv_input[src_k_idx + rotated_head_dim_offset]); + if (!params.q_only_input) + { + k = *reinterpret_cast(¶ms.qkv_input[src_k_idx]); + v = *reinterpret_cast(¶ms.qkv_input[src_v_idx]); + k_pair = *reinterpret_cast(¶ms.qkv_input[src_k_idx + rotated_head_dim_offset]); + } if constexpr (ADD_BIAS) { auto const q_bias = *reinterpret_cast(¶ms.qkv_bias[hidden_idx]); - auto const k_bias - = *reinterpret_cast(¶ms.qkv_bias[hidden_idx_kv + src_k_offset]); - auto const v_bias - = *reinterpret_cast(¶ms.qkv_bias[hidden_idx_kv + src_v_offset]); auto const q_pair_bias = *reinterpret_cast(¶ms.qkv_bias[hidden_idx + rotated_head_dim_offset]); - auto const k_pair_bias = *reinterpret_cast( - ¶ms.qkv_bias[hidden_idx_kv + src_k_offset + rotated_head_dim_offset]); q = mmha::add(q, q_bias); - k = mmha::add(k, k_bias); - v = mmha::add(v, v_bias); q_pair = mmha::add(q_pair, q_pair_bias); - k_pair = mmha::add(k_pair, k_pair_bias); + if (!params.q_only_input) + { + auto const k_bias + = *reinterpret_cast(¶ms.qkv_bias[hidden_idx_kv + src_k_offset]); + auto const v_bias + = *reinterpret_cast(¶ms.qkv_bias[hidden_idx_kv + src_v_offset]); + auto const k_pair_bias = *reinterpret_cast( + ¶ms.qkv_bias[hidden_idx_kv + src_k_offset + rotated_head_dim_offset]); + k = mmha::add(k, k_bias); + v = mmha::add(v, v_bias); + k_pair = mmha::add(k_pair, k_pair_bias); + } } k_wo_pos = k; } @@ -550,7 +559,7 @@ __global__ void applyBiasRopeUpdateKVCache(QKVPreprocessingParams(bounded_global_token_idx) * params.hidden_size + hidden_idx; + auto const src_q_idx = static_cast(bounded_global_token_idx) * input_hidden_size + hidden_idx; auto const src_k_idx - = static_cast(bounded_global_token_idx) * params.hidden_size + src_k_offset + hidden_idx_kv; + = static_cast(bounded_global_token_idx) * input_hidden_size + src_k_offset + hidden_idx_kv; auto const src_v_idx - = static_cast(bounded_global_token_idx) * params.hidden_size + src_v_offset + hidden_idx_kv; + = static_cast(bounded_global_token_idx) * input_hidden_size + src_v_offset + hidden_idx_kv; auto q = *reinterpret_cast(¶ms.qkv_input[src_q_idx]); - auto k = *reinterpret_cast(¶ms.qkv_input[src_k_idx]); - auto v = *reinterpret_cast(¶ms.qkv_input[src_v_idx]); + VecT k{}; + VecT v{}; [[maybe_unused]] auto q_pair = *reinterpret_cast(¶ms.qkv_input[src_q_idx + rotated_head_dim_offset]); - [[maybe_unused]] auto k_pair - = *reinterpret_cast(¶ms.qkv_input[src_k_idx + rotated_head_dim_offset]); + [[maybe_unused]] VecT k_pair{}; + if (!params.q_only_input) + { + k = *reinterpret_cast(¶ms.qkv_input[src_k_idx]); + v = *reinterpret_cast(¶ms.qkv_input[src_v_idx]); + k_pair = *reinterpret_cast(¶ms.qkv_input[src_k_idx + rotated_head_dim_offset]); + } // Bias should have been fused with QKV projection, but we keep the logic here for unit tests. if constexpr (ADD_BIAS) { auto const q_bias = *reinterpret_cast(¶ms.qkv_bias[hidden_idx]); - auto const k_bias = *reinterpret_cast(¶ms.qkv_bias[hidden_idx_kv + src_k_offset]); - auto const v_bias = *reinterpret_cast(¶ms.qkv_bias[hidden_idx_kv + src_v_offset]); auto const q_pair_bias = *reinterpret_cast(¶ms.qkv_bias[hidden_idx + rotated_head_dim_offset]); - auto const k_pair_bias = *reinterpret_cast( - ¶ms.qkv_bias[hidden_idx_kv + src_k_offset + rotated_head_dim_offset]); q = mmha::add(q, q_bias); - k = mmha::add(k, k_bias); - v = mmha::add(v, v_bias); q_pair = mmha::add(q_pair, q_pair_bias); - k_pair = mmha::add(k_pair, k_pair_bias); + if (!params.q_only_input) + { + auto const k_bias = *reinterpret_cast(¶ms.qkv_bias[hidden_idx_kv + src_k_offset]); + auto const v_bias = *reinterpret_cast(¶ms.qkv_bias[hidden_idx_kv + src_v_offset]); + auto const k_pair_bias = *reinterpret_cast( + ¶ms.qkv_bias[hidden_idx_kv + src_k_offset + rotated_head_dim_offset]); + k = mmha::add(k, k_bias); + v = mmha::add(v, v_bias); + k_pair = mmha::add(k_pair, k_pair_bias); + } } // Cos/sin cache. @@ -967,7 +987,7 @@ __global__ void applyBiasRopeUpdateKVCacheV2(QKVPreprocessingParams(params, stream); break; case 224: kernelV2DispatchHeadSize<224, 224, T, TCache, KVCacheBuffer>(params, stream); break; case 256: kernelV2DispatchHeadSize<256, 256, T, TCache, KVCacheBuffer>(params, stream); break; + case 512: kernelV2DispatchHeadSize<512, 512, T, TCache, KVCacheBuffer>(params, stream); break; case 576: kernelV2DispatchHeadSize<576, 576, T, TCache, KVCacheBuffer>(params, stream); break; default: // Fall back to v1 kernel. diff --git a/cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp b/cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp index bc80d8446445..45866df9fa30 100644 --- a/cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp +++ b/cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp @@ -80,6 +80,12 @@ cudaStream_t currentStreamFor(at::Tensor const& tensor) return at::cuda::getCurrentCUDAStream(tensor.get_device()).stream(); } +bool isQOnlyInput(at::Tensor const& qkvInput, int64_t const numHeads, int64_t const headSize) +{ + TORCH_CHECK(qkvInput.dim() > 0, "QKV input must have at least one dimension."); + return qkvInput.size(-1) == numHeads * headSize; +} + struct WorkspaceAccessor { uint8_t* base{}; @@ -276,7 +282,8 @@ trtllmGenContextPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, tor TORCH_CHECK(kv_cache_block_offsets.has_value(), "kv_cache_block_offsets is required."); TORCH_CHECK(!cross_attention || !is_mla_enable, "trtllm-gen cross attention does not support MLA."); - bool const separateQKvOutput = paged_context_fmha || fp8_context_fmha || cross_attention; + bool const qOnlyInput = !cross_attention && isQOnlyInput(qkv_input, num_heads, head_size); + bool const separateQKvOutput = paged_context_fmha || fp8_context_fmha || cross_attention || qOnlyInput; auto const qkvScalarType = qkv_input.scalar_type(); auto const qkvElementSize = static_cast(qkv_input.element_size()); auto const quantMode = tensorrt_llm::common::QuantMode(static_cast(kv_cache_quant_mode)); @@ -396,6 +403,7 @@ trtllmGenContextPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, tor qkvParams.separate_q_kv_output = separateQKvOutput; qkvParams.quantized_fp8_output = fp8_context_fmha; qkvParams.generation_phase = false; + qkvParams.q_only_input = qOnlyInput; qkvParams.multi_processor_count = static_cast(multi_processor_count); qkvParams.rotary_vision_start = 0; qkvParams.rotary_vision_length = 0; @@ -470,6 +478,10 @@ void trtllmGenContextPostprocess(torch::Tensor qkv_input, torch::Tensor workspac int64_t const attention_chunk_size, int64_t const multi_processor_count) { (void) mask_type; + if (isQOnlyInput(qkv_input, num_heads, head_size)) + { + return; + } auto const qkvScalarType = qkv_input.scalar_type(); auto const qkvElementSize = static_cast(qkv_input.element_size()); auto const quantMode = tensorrt_llm::common::QuantMode(static_cast(kv_cache_quant_mode)); @@ -596,6 +608,7 @@ trtllmGenGenerationPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, (void) bmm2_scale; bool const isMultiTokenGen = spec_decoding_generation_lengths.has_value() && predicted_tokens_per_seq > 1; + bool const qOnlyInput = !cross_attention && isQOnlyInput(qkv_input, num_heads, head_size); TORCH_CHECK( !cross_attention || !isMultiTokenGen, "trtllm-gen cross attention does not support multi-token generation."); auto const qkvScalarType = qkv_input.scalar_type(); @@ -728,6 +741,7 @@ trtllmGenGenerationPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, qkvParams.separate_q_kv_output = true; qkvParams.quantized_fp8_output = fp8_context_fmha; qkvParams.generation_phase = true; + qkvParams.q_only_input = qOnlyInput; qkvParams.multi_processor_count = static_cast(multi_processor_count); qkvParams.rotary_vision_start = 0; qkvParams.rotary_vision_length = 0; diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index d49ab084a128..93ce6dd8eece 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -650,6 +650,20 @@ def get_draft_metadata( "views.") return self._draft_metadata + def get_shared_kv_draft_metadata( + self, + num_accepted_tokens: torch.Tensor, + num_contexts: int, + ) -> "FlashInferAttentionMetadata": + """Return the shared-target-KV draft view at the accepted prefix.""" + draft_metadata = self.get_draft_metadata() + draft_metadata.update_shared_kv_draft_lengths( + self, + num_accepted_tokens, + num_contexts, + ) + return draft_metadata + def _sync_draft_view(self, target: "FlashInferAttentionMetadata") -> None: """Refresh a shared- or separate-KV draft metadata view.""" if not (self._is_shared_kv_draft_view diff --git a/tensorrt_llm/_torch/attention_backend/fmha/__init__.py b/tensorrt_llm/_torch/attention_backend/fmha/__init__.py index b4970d93468e..5e9079ec584a 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/__init__.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/__init__.py @@ -13,24 +13,29 @@ # See the License for the specific language governing permissions and # limitations under the License. +from .combined import CombinedFmha from .cute_dsl_mla import CuteDslMlaFmha from .fallback import FallbackFmha from .flashinfer_trtllm_gen import FlashInferTrtllmGenFmha -from .interface import Fmha +from .interface import Fmha, FmhaPhase from .msa_sparse_gqa import MsaSparseGqaFmha from .phased import FmhaParams, PhasedFmha from .registry import DEFAULT_FMHA_LIBS, FMHA_LIBS, FmhaCls, get_enabled_fmha_lib_classes +from .triton_custom_mask import TritonCustomMaskFmha __all__ = [ "DEFAULT_FMHA_LIBS", "FMHA_LIBS", + "CombinedFmha", "CuteDslMlaFmha", "FallbackFmha", "FlashInferTrtllmGenFmha", "Fmha", "FmhaCls", "FmhaParams", + "FmhaPhase", "MsaSparseGqaFmha", "PhasedFmha", + "TritonCustomMaskFmha", "get_enabled_fmha_lib_classes", ] diff --git a/tensorrt_llm/_torch/attention_backend/fmha/combined.py b/tensorrt_llm/_torch/attention_backend/fmha/combined.py new file mode 100644 index 000000000000..be68d6b5a225 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/fmha/combined.py @@ -0,0 +1,89 @@ +# 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 .phased import FmhaParams, PhasedFmha + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention_backend.trtllm import ( + TrtllmAttention, + TrtllmAttentionMetadata, + ) + + +class CombinedFmha(PhasedFmha): + """Combine context and generation implementations for mixed batches.""" + + def __init__(self, attn: "TrtllmAttention"): + super().__init__(attn) + self._context_impl: Optional[PhasedFmha] = None + self._generation_impl: Optional[PhasedFmha] = None + + def set_fmha_impls( + self, + context_impl: PhasedFmha, + generation_impl: PhasedFmha, + ) -> None: + self._context_impl = context_impl + self._generation_impl = generation_impl + + def _get_context_impl(self) -> PhasedFmha: + if self._context_impl is None: + raise RuntimeError("CombinedFmha context implementation is not configured.") + return self._context_impl + + def _get_generation_impl(self) -> PhasedFmha: + if self._generation_impl is None: + raise RuntimeError("CombinedFmha generation implementation is not configured.") + return self._generation_impl + + def prepare_workspace( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + workspace: torch.Tensor, + ) -> None: + context_impl = self._get_context_impl() + generation_impl = self._get_generation_impl() + context_impl.prepare_workspace( + q, + k, + v, + metadata, + forward_args, + workspace, + ) + generation_impl.prepare_workspace( + q, + k, + v, + metadata, + forward_args, + workspace, + ) + + def run_context(self, params: FmhaParams) -> None: + self._get_context_impl().run_context(params) + + def run_generation(self, params: FmhaParams) -> None: + self._get_generation_impl().run_generation(params) diff --git a/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py index 83341117796b..534259a34f21 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py @@ -28,6 +28,7 @@ from tensorrt_llm._utils import get_sm_version from tensorrt_llm.logger import logger +from .interface import FmhaPhase from .phased import FmhaParams, PhasedFmha if TYPE_CHECKING: @@ -193,7 +194,11 @@ def is_supported( v: Optional[torch.Tensor], metadata: "TrtllmAttentionMetadata", forward_args: AttentionForwardArgs, + *, + phase: Optional[FmhaPhase] = None, ) -> bool: + if phase not in (None, FmhaPhase.GENERATION): + return False supported, reason = self._is_supported_with_reason( q, self.attn, diff --git a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py index 70e0d1feea69..233f007687dc 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py @@ -17,10 +17,13 @@ import torch -from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs +from tensorrt_llm._torch.attention_backend.interface import ( + AttentionForwardArgs, + CustomAttentionMask, +) from tensorrt_llm.bindings.internal import thop -from .interface import Fmha +from .interface import Fmha, FmhaPhase if TYPE_CHECKING: from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata @@ -49,6 +52,32 @@ class FallbackFmha(Fmha): """Fallback FMHA implementation using the fused TRT-LLM thop attention op.""" + def is_supported( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + *, + phase: Optional[FmhaPhase] = None, + ) -> bool: + del q, phase + if forward_args.attention_mask == CustomAttentionMask.CUSTOM: + return False + + # The fused THOP path cannot read an existing self-attention KV cache + # from a Q-only input. This form is used by external shared-KV draft + # layers and must be handled by a phased generation implementation. + is_q_only_self_attention = ( + not metadata.is_cross + and k is None + and v is None + and not forward_args.is_fused_qkv + and not forward_args.update_kv_cache + ) + return not is_q_only_self_attention + def forward( self, q: torch.Tensor, diff --git a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py index 49330a61e31d..be94f93cc3ae 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py @@ -58,7 +58,9 @@ from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantMode +from .interface import FmhaPhase from .phased import FmhaParams, PhasedFmha +from .trtllm_gen_utils import get_trtllm_gen_context_workspace_size if TYPE_CHECKING: from tensorrt_llm._torch.attention_backend.trtllm import ( @@ -230,52 +232,6 @@ def _get_bmm1_scale_log2(bmm1_scale: torch.Tensor) -> torch.Tensor: return bmm1_scale.narrow(0, 1, 1) -@lru_cache(maxsize=128) -def _get_context_workspace_layout( - dtype: torch.dtype, - batch_size: int, - num_tokens: int, - num_heads: int, - head_size: int, - rotary_embedding_dim: int, - fp8_context_fmha: bool, -) -> dict[str, int]: - return thop.get_trtllm_gen_context_workspace_layout( - dtype, - batch_size, - num_tokens, - num_heads, - head_size, - rotary_embedding_dim, - True, - fp8_context_fmha, - ) - - -@lru_cache(maxsize=128) -def _get_context_workspace_size( - dtype: torch.dtype, - max_num_seq: int, - max_num_tokens: int, - num_heads: int, - head_size: int, - rotary_embedding_dim: int, - fp8_context_fmha: bool, -) -> int: - if max_num_tokens == 0: - return 0 - layout = _get_context_workspace_layout( - dtype, - max_num_seq, - max_num_tokens, - num_heads, - head_size, - rotary_embedding_dim, - fp8_context_fmha, - ) - return int(layout["total_size"]) - - @lru_cache(maxsize=128) def _get_generation_workspace_layout( dtype: torch.dtype, @@ -335,7 +291,7 @@ def _get_workspace_size( rotary_embedding_dim: int, fp8_context_fmha: bool, ) -> int: - context_size = _get_context_workspace_size( + context_size = get_trtllm_gen_context_workspace_size( dtype, max_num_requests, num_tokens, @@ -422,18 +378,6 @@ def __init__(self, attn: "TrtllmAttention"): self._multi_processor_count: Optional[int] = None self._multi_ctas_kv_counter_buffer: Optional[torch.Tensor] = None - def _get_total_num_blocks(self, meta: "TrtllmAttentionMetadata") -> int: - kv_cache_manager = meta.kv_cache_manager - if kv_cache_manager is not None: - get_page_index_upper_bound = getattr( - getattr(kv_cache_manager, "impl", None), "get_page_index_upper_bound", None - ) - # KVCacheManagerV2 exposes this implementation-only API and reports an - # already-flattened page-index bound, unlike the legacy logical block count. - if get_page_index_upper_bound is not None: - return int(kv_cache_manager.blocks_in_primary_pool) - return super()._get_total_num_blocks(meta) - @classmethod def is_available(cls, attn: "TrtllmAttention") -> bool: if not IS_FLASHINFER_AVAILABLE: @@ -576,6 +520,8 @@ def is_supported( v: Optional[torch.Tensor], metadata: "TrtllmAttentionMetadata", forward_args: AttentionForwardArgs, + *, + phase: Optional[FmhaPhase] = None, ) -> bool: supported, reason = self._is_supported_with_reason( q, @@ -584,9 +530,10 @@ def is_supported( self.attn, metadata, forward_args, + phase=phase, ) if not supported: - logger.debug(f"FlashInfer trtllm-gen fmha library does not support {reason}") + logger.debug(f"FlashInfer trtllm-gen fmha library does not support request: {reason}") return supported def _is_supported_with_reason( @@ -597,8 +544,22 @@ def _is_supported_with_reason( attn: "TrtllmAttention", meta: "TrtllmAttentionMetadata", fwd: AttentionForwardArgs, + *, + phase: Optional[FmhaPhase] = None, ) -> Tuple[bool, str]: is_mla_enable = attn.is_mla_enable + if phase is None: + has_context_phase = fwd.attention_input_type != AttentionInputType.generation_only + has_generation_phase = fwd.attention_input_type != AttentionInputType.context_only + elif phase == FmhaPhase.CONTEXT: + has_context_phase = True + has_generation_phase = False + elif phase == FmhaPhase.GENERATION: + has_context_phase = False + has_generation_phase = True + else: + return False, f"invalid FMHA phase: {phase}." + sparse_params = attn.sparse_params has_skip_softmax = sparse_params is not None and sparse_params.algorithm == "skip_softmax" has_sparse_attention = sparse_params is not None and not has_skip_softmax @@ -627,8 +588,8 @@ def _is_supported_with_reason( return False, "relative attention bias." if meta.use_spec_decoding and meta.is_spec_dec_tree: return False, "spec-dec tree/custom masks." - if is_mla_enable and fwd.attention_input_type != AttentionInputType.generation_only: - return False, "MLA with non-generation-only attention." + if is_mla_enable and has_context_phase: + return False, "MLA context attention." if meta.kv_cache_block_offsets is None: return False, "non-paged KV cache; paged KV cache is required." @@ -648,9 +609,24 @@ def _is_supported_with_reason( if tokens_per_block is None: tokens_per_block = 0 - attn_input_type = fwd.attention_input_type - has_context_phase = attn_input_type != AttentionInputType.generation_only - has_generation_phase = attn_input_type != AttentionInputType.context_only + if not is_mla_enable and not meta.is_cross: + q_hidden_size = attn.num_heads * attn.head_dim + qkv_hidden_size = q_hidden_size + 2 * attn.num_kv_heads * attn.head_dim + has_fused_qkv = ( + fwd.is_fused_qkv and k is None and v is None and q.size(-1) == qkv_hidden_size + ) + has_q_only = ( + not fwd.is_fused_qkv + and not fwd.update_kv_cache + and k is None + and v is None + and q.size(-1) == q_hidden_size + ) + if not has_fused_qkv and not has_q_only: + return False, "self attention requires fused QKV or Q-only cached-KV input." + if has_q_only and q.dtype == torch.float8_e4m3fn: + return False, "Q-only cached-KV preprocessing supports FP16 or BF16 input." + q_dtype = q.dtype o_dtype = output.dtype @@ -763,20 +739,20 @@ 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 get_fp8_context_fmha( + def _use_fp8_context_fmha( self, - q: torch.Tensor, output: torch.Tensor, - metadata: "TrtllmAttentionMetadata", - forward_args: AttentionForwardArgs, - is_gen_only: bool, + attention_input_type: AttentionInputType, ) -> bool: 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) + or ( + kv_cache_quant_mode.has_fp8_kv_cache() + and attention_input_type != AttentionInputType.generation_only + ) ) def prepare_workspace( @@ -824,7 +800,7 @@ def prepare_workspace( 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) + fp8_context_fmha = self._use_fp8_context_fmha(output, attention_input_type) workspace_max_tokens = max(num_tokens, metadata.max_context_length) workspace_max_gen_tokens = max(num_gen_tokens, metadata.max_num_requests) @@ -899,6 +875,8 @@ def run_context( rope_params = attn.rope_params bmm1_scale_static = self._get_bmm1_scale(attn) attention_chunk_size = self._get_attention_chunk_size(attn) + output = fwd.output + fp8_context_fmha = self._use_fp8_context_fmha(output, fwd.attention_input_type) ( q_processed, kv_pool, @@ -948,7 +926,7 @@ def run_context( bmm1_scale_static, # bmm1_scale 1.0, # bmm2_scale attention_chunk_size, # attention_chunk_size - params.fp8_context_fmha, # fp8_context_fmha + fp8_context_fmha, # fp8_context_fmha meta.use_paged_context_fmha, # paged_context_fmha attn.is_mla_enable, # is_mla_enable self._multi_processor_count, # multi_processor_count @@ -962,7 +940,7 @@ def run_context( has_fp4_kv = QuantMode(attn.quant_mode).has_fp4_kv_cache() if has_fp4_kv and kv_scale_pool is None: raise RuntimeError("trtllm-gen FP4 KV cache requires KV scale pool.") - if has_fp4_kv or params.fp8_context_fmha: + if has_fp4_kv or fp8_context_fmha: q_processed = ( q_processed.view(torch.uint8) .flatten()[: params.num_tokens * attn.num_heads * attn.head_dim] @@ -970,9 +948,9 @@ def run_context( .view(params.num_tokens, attn.num_heads, attn.head_dim) ) ctx_bmm1_scale = bmm1_scale_static - if params.fp8_context_fmha and bmm1_scale is not None: + if fp8_context_fmha and bmm1_scale is not None: ctx_bmm1_scale = bmm1_scale.narrow(0, 0, 1) - ctx_bmm2_scale = bmm2_scale if params.fp8_context_fmha and bmm2_scale is not None else 1.0 + ctx_bmm2_scale = bmm2_scale if fp8_context_fmha and bmm2_scale is not None else 1.0 causal = ( False if params.is_cross @@ -1038,7 +1016,7 @@ def run_context( rope_params.max_positions, # rotary_embedding_max_positions attn.position_embedding_type, # position_embedding_type bmm1_scale_static, # bmm1_scale - params.fp8_context_fmha, # fp8_context_fmha + fp8_context_fmha, # fp8_context_fmha meta.use_paged_context_fmha, # paged_context_fmha attn.is_mla_enable, # is_mla_enable attention_chunk_size, # attention_chunk_size @@ -1055,6 +1033,8 @@ def run_generation( rope_params = attn.rope_params bmm1_scale_static = self._get_bmm1_scale(attn) attention_chunk_size = self._get_attention_chunk_size(attn) + output = fwd.output + fp8_context_fmha = self._use_fp8_context_fmha(output, fwd.attention_input_type) batch_beam = params.num_requests * meta.beam_width ( q_processed, @@ -1105,7 +1085,7 @@ def run_generation( attn.position_embedding_type, # position_embedding_type bmm1_scale_static, # bmm1_scale 1.0, # bmm2_scale - params.fp8_context_fmha, # fp8_context_fmha + fp8_context_fmha, # fp8_context_fmha attn.predicted_tokens_per_seq, # predicted_tokens_per_seq attention_chunk_size, # attention_chunk_size self._multi_processor_count, # multi_processor_count @@ -1122,7 +1102,7 @@ def run_generation( has_fp4_kv = QuantMode(attn.quant_mode).has_fp4_kv_cache() if has_fp4_kv and kv_scale_pool is None: raise RuntimeError("trtllm-gen FP4 KV cache requires KV scale pool.") - if has_fp4_kv or params.fp8_context_fmha: + if has_fp4_kv or fp8_context_fmha: q_processed = ( q_processed.view(torch.uint8) .flatten()[: params.num_tokens * attn.num_heads * attn.head_dim] @@ -1130,9 +1110,9 @@ def run_generation( .view(params.num_tokens, attn.num_heads, attn.head_dim) ) gen_bmm1_scale = ( - bmm1_scale if params.fp8_context_fmha and bmm1_scale is not None else bmm1_scale_static + bmm1_scale if fp8_context_fmha and bmm1_scale is not None else bmm1_scale_static ) - gen_bmm2_scale = bmm2_scale if params.fp8_context_fmha and bmm2_scale is not None else 1.0 + gen_bmm2_scale = bmm2_scale if fp8_context_fmha and bmm2_scale is not None else 1.0 flashinfer.decode.trtllm_batch_decode_with_kv_cache( query=q_processed, diff --git a/tensorrt_llm/_torch/attention_backend/fmha/interface.py b/tensorrt_llm/_torch/attention_backend/fmha/interface.py index f14f26964b25..f2209d013c8f 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/interface.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/interface.py @@ -15,6 +15,7 @@ import weakref from abc import ABC, abstractmethod +from enum import Enum from typing import TYPE_CHECKING, Optional import torch @@ -28,6 +29,13 @@ ) +class FmhaPhase(str, Enum): + """Attention phase checked by a phased FMHA library.""" + + CONTEXT = "context" + GENERATION = "generation" + + class Fmha(ABC): """Common runtime contract for TRT-LLM attention FMHA libraries.""" @@ -52,7 +60,10 @@ def is_supported( v: Optional[torch.Tensor], metadata: "TrtllmAttentionMetadata", forward_args: AttentionForwardArgs, + *, + phase: Optional[FmhaPhase] = None, ) -> bool: + """Return whether this library supports the request or requested phase.""" return True @abstractmethod diff --git a/tensorrt_llm/_torch/attention_backend/fmha/phased.py b/tensorrt_llm/_torch/attention_backend/fmha/phased.py index 3752403ecc89..7656839fcaa5 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/phased.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/phased.py @@ -18,7 +18,12 @@ import torch -from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs, AttentionInputType +from tensorrt_llm._torch.attention_backend.interface import ( + AttentionForwardArgs, + AttentionInputType, + CustomAttentionMask, + PredefinedAttentionMask, +) from .interface import Fmha @@ -37,6 +42,8 @@ class FmhaParams: workspace: torch.Tensor attention_input: Optional[torch.Tensor] = None qkv_input: Optional[torch.Tensor] = None + key_input: Optional[torch.Tensor] = None + value_input: Optional[torch.Tensor] = None context_buf: Optional[torch.Tensor] = None sequence_lengths: Optional[torch.Tensor] = None context_lengths: Optional[torch.Tensor] = None @@ -47,7 +54,6 @@ class FmhaParams: 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 @@ -83,6 +89,16 @@ def _get_total_num_blocks( if kv_cache_manager is None: return 0 + get_page_index_upper_bound = getattr( + getattr(kv_cache_manager, "impl", None), + "get_page_index_upper_bound", + None, + ) + if get_page_index_upper_bound is not None: + # KVCacheManagerV2 exposes an already-flattened page-index bound, + # unlike the legacy logical block count. + return int(kv_cache_manager.blocks_in_primary_pool) + 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) @@ -94,16 +110,6 @@ def _get_total_num_blocks( 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, @@ -146,7 +152,6 @@ def forward( 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, @@ -182,7 +187,6 @@ def forward( 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), is_cross=metadata.is_cross, @@ -205,6 +209,12 @@ def forward( params.attention_input = q[token_offset : token_offset + num_ctx_tokens] params.qkv_input = params.attention_input + params.key_input = ( + k[token_offset : token_offset + num_ctx_tokens] if k is not None else None + ) + params.value_input = ( + v[token_offset : token_offset + num_ctx_tokens] if v is not None else None + ) 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:] @@ -242,6 +252,12 @@ def forward( params.attention_input = q[token_offset : token_offset + num_gen_tokens] params.qkv_input = params.attention_input + params.key_input = ( + k[token_offset : token_offset + num_gen_tokens] if k is not None else None + ) + params.value_input = ( + v[token_offset : token_offset + num_gen_tokens] if v is not None else None + ) 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 @@ -254,6 +270,13 @@ def forward( if attn.is_mla_enable: self.run_mla_generation(params) else: + # The custom mask covers only the context portion of a mixed batch. + if ( + not metadata.is_cross + and params.fwd.attention_mask == CustomAttentionMask.CUSTOM + ): + params.fwd.attention_mask = PredefinedAttentionMask.CAUSAL + params.fwd.attention_mask_data = None self.run_generation(params) def run_context(self, params: FmhaParams) -> None: diff --git a/tensorrt_llm/_torch/attention_backend/fmha/registry.py b/tensorrt_llm/_torch/attention_backend/fmha/registry.py index dec41c668eff..fda09290239b 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/registry.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/registry.py @@ -20,6 +20,7 @@ from .fallback import FallbackFmha from .flashinfer_trtllm_gen import FlashInferTrtllmGenFmha from .interface import Fmha +from .triton_custom_mask import TritonCustomMaskFmha FmhaCls: TypeAlias = type[Fmha] @@ -34,6 +35,7 @@ def init_fmha_libs() -> dict[str, "FmhaCls"]: from .msa_sparse_gqa import MsaSparseGqaFmha return { + "triton_custom_mask": TritonCustomMaskFmha, "cute_dsl_mla": CuteDslMlaFmha, "msa_sparse_gqa": MsaSparseGqaFmha, "flashinfer_trtllm_gen": FlashInferTrtllmGenFmha, diff --git a/tensorrt_llm/_torch/attention_backend/fmha/triton_custom_mask.py b/tensorrt_llm/_torch/attention_backend/fmha/triton_custom_mask.py new file mode 100644 index 000000000000..b685f5989889 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/fmha/triton_custom_mask.py @@ -0,0 +1,322 @@ +# 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. + +"""Triton custom-mask context attention for the TRT-LLM backend.""" + +import math +from functools import lru_cache +from typing import TYPE_CHECKING, Optional + +import torch + +from tensorrt_llm._torch.attention_backend.interface import ( + AttentionForwardArgs, + CustomAttentionMask, +) +from tensorrt_llm.bindings.internal import thop +from tensorrt_llm.functional import AttentionMaskType, PositionEmbeddingType +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantMode + +from .interface import FmhaPhase +from .phased import FmhaParams, PhasedFmha +from .trtllm_gen_utils import get_trtllm_gen_context_workspace_size + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention_backend.trtllm import ( + TrtllmAttention, + TrtllmAttentionMetadata, + ) + + +class TritonCustomMaskFmha(PhasedFmha): + """Run custom-mask context attention with Triton.""" + + SUPPORTED_DTYPES = {torch.float16, torch.bfloat16} + + def __init__(self, attn: "TrtllmAttention"): + super().__init__(attn) + self._multi_processor_count: Optional[int] = None + + @classmethod + def is_available(cls, attn: "TrtllmAttention") -> bool: + required_ops = ( + "get_trtllm_gen_context_workspace_layout", + "trtllm_gen_context_preprocess", + ) + missing_ops = [op for op in required_ops if not hasattr(thop, op)] + if missing_ops: + logger.debug( + "Triton custom-mask FMHA is unavailable: missing fused " + f"nanobind ops: {', '.join(missing_ops)}." + ) + return False + if attn.num_heads <= 0 or attn.num_kv_heads <= 0: + return False + if attn.num_heads % attn.num_kv_heads != 0: + return False + return True + + def is_supported( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + *, + phase: FmhaPhase, + ) -> bool: + supported, reason = self._check_support_with_reason( + q, + k, + v, + metadata, + forward_args, + phase=phase, + ) + if not supported: + logger.debug(f"Triton custom-mask FMHA does not support request: {reason}") + return supported + + def _check_support_with_reason( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + *, + phase: FmhaPhase, + ) -> tuple[bool, str]: + if phase != FmhaPhase.CONTEXT: + return False, "Only context attention is supported." + if forward_args.attention_mask != CustomAttentionMask.CUSTOM: + return False, "Request does not use a custom attention mask." + if forward_args.attention_mask_data is None: + return False, "Custom attention requires attention_mask_data." + if metadata.num_contexts <= 0: + return False, "Custom attention requires at least one context request." + if metadata.is_cross: + return False, "Custom-mask cross attention is not supported." + if metadata.kv_cache_block_offsets is None: + return False, "Custom-mask TRT-LLM attention requires paged KV cache." + if self.attn.is_mla_enable: + return False, "Custom-mask MLA is not supported." + if not forward_args.is_fused_qkv or k is not None or v is not None: + return False, "Custom-mask TRT-LLM attention requires fused QKV input." + if q.dtype not in self.SUPPORTED_DTYPES: + return False, f"Input dtype {q.dtype} is not supported." + output = forward_args.output + if output is None or output.dtype != q.dtype: + return False, "Custom-mask Triton attention requires output to match the input dtype." + if forward_args.output_sf is not None: + return False, "Custom-mask Triton attention does not support NVFP4 output." + if QuantMode(self.attn.quant_mode).has_kv_cache_quant(): + return False, "Custom-mask TRT-LLM attention does not support quantized KV cache." + if forward_args.attention_sinks is not None: + return False, "Custom-mask Triton attention does not support attention sinks." + if ( + getattr(self.attn, "head_dim", 0) > 256 + and PositionEmbeddingType(self.attn.position_embedding_type).is_rope() + ): + return ( + False, + "Custom-mask head dimensions above 256 require RoPE to be applied before the backend.", + ) + + return True, "" + + @staticmethod + @lru_cache(maxsize=None) + def _get_multi_processor_count_for_device(device_index: int) -> int: + return torch.cuda.get_device_properties(device_index).multi_processor_count + + @classmethod + def _get_multi_processor_count(cls, device: torch.device) -> int: + device = torch.device(device) + if device.type != "cuda": + raise RuntimeError("Triton custom-mask FMHA requires CUDA tensors.") + device_index = device.index + if device_index is None: + device_index = torch.cuda.current_device() + return cls._get_multi_processor_count_for_device(device_index) + + @staticmethod + def _get_bmm1_scale(attn: "TrtllmAttention") -> float: + return 1.0 / (math.sqrt(attn.head_dim) * attn.q_scaling) + + @staticmethod + def _get_attention_chunk_size(attn: "TrtllmAttention") -> int: + return attn.attention_chunk_size if attn.attention_chunk_size is not None else 0 + + def prepare_workspace( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + workspace: torch.Tensor, + ) -> None: + if self._multi_processor_count is None: + self._multi_processor_count = self._get_multi_processor_count(q.device) + + attn = self.attn + required_workspace_size = get_trtllm_gen_context_workspace_size( + dtype=q.dtype, + max_num_seq=metadata.max_num_requests, + max_num_tokens=max(q.size(0), metadata.max_context_length), + num_heads=attn.num_heads, + head_size=attn.head_dim, + rotary_embedding_dim=attn.rope_dim, + fp8_context_fmha=False, + ) + current_workspace_size = workspace.numel() * workspace.element_size() + if current_workspace_size < required_workspace_size: + if metadata.is_cuda_graph and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "Attention CUDA graph workspace is smaller than the required size " + "for Triton custom-mask FMHA." + ) + required_workspace_numel = math.ceil(required_workspace_size / workspace.element_size()) + workspace.resize_((required_workspace_numel,)) + + def run_context(self, params: FmhaParams) -> None: + """Preprocess fused QKV, then run the Triton custom-mask context kernel.""" + attn = params.attn + meta = params.meta + fwd = params.fwd + rope_params = attn.rope_params + bmm1_scale = self._get_bmm1_scale(attn) + attention_chunk_size = self._get_attention_chunk_size(attn) + + ( + q_processed, + kv_pool, + block_tables, + _, + _, + _, + _, + cu_q_seqlens, + _, + _, + _, + window_left, + ) = thop.trtllm_gen_context_preprocess( + params.qkv_input, + params.workspace, + params.sequence_lengths, + params.context_lengths, + meta.kv_cache_block_offsets, + meta.host_kv_cache_pool_pointers, + meta.host_kv_cache_pool_mapping, + fwd.kv_scale_orig_quant, + fwd.kv_scale_quant_orig, + fwd.out_scale, + attn.rotary_inv_freq, + attn.rotary_cos_sin, + fwd.mrope_rotary_cos_sin, + attn.local_layer_idx, + attn.num_heads, + attn.num_kv_heads, + attn.head_dim, + params.tokens_per_block, + int(AttentionMaskType.causal), + attn.quant_mode, + params.max_attention_window_size, + params.cyclic_attention_window_size, + params.num_tokens, + params.batch_size, + params.input_seq_length, + params.max_past_kv_length, + rope_params.dim, + rope_params.theta, + int(rope_params.scale_type), + rope_params.scale, + rope_params.max_positions, + attn.position_embedding_type, + bmm1_scale, + 1.0, + attention_chunk_size, + False, # fp8_context_fmha + False, # paged_context_fmha; keep processed QKV packed in-place + False, # is_mla_enable + self._multi_processor_count, + params.total_num_blocks, + params.kv_factor, + True, # need_build_kv_cache_metadata + None, # cross_kv + False, # cross_attention + ) + + if kv_pool is None or block_tables is None: + raise RuntimeError("Custom-mask TRT-LLM attention requires paged KV metadata.") + if params.qkv_input is None or params.context_buf is None: + raise RuntimeError( + "Custom-mask TRT-LLM attention requires context QKV and output buffers." + ) + if params.sequence_lengths is None or params.context_lengths is None: + raise RuntimeError("Custom-mask TRT-LLM attention requires context sequence lengths.") + + q_size = attn.num_heads * attn.head_dim + kv_size = attn.num_kv_heads * attn.head_dim + k_processed = params.qkv_input.narrow(1, q_size, kv_size).view( + params.num_tokens, + attn.num_kv_heads, + attn.head_dim, + ) + v_processed = params.qkv_input.narrow(1, q_size + kv_size, kv_size).view( + params.num_tokens, + attn.num_kv_heads, + attn.head_dim, + ) + + block_tables = block_tables[: params.batch_size] + blocks_per_sequence = block_tables.size(-1) + page_table_indptr = ( + torch.arange( + params.batch_size + 1, + dtype=torch.int32, + device=block_tables.device, + ) + * blocks_per_sequence + ) + prefix_lens = ( + params.sequence_lengths[: params.batch_size] + - params.context_lengths[: params.batch_size] + ) + + from ..triton_prefill import triton_prefill_with_custom_mask + + triton_prefill_with_custom_mask( + q=q_processed, + k=k_processed, + v=v_processed, + output=params.context_buf, + qo_indptr=cu_q_seqlens, + kv_cache=None, + prefix_lens=prefix_lens, + page_table_indptr=page_table_indptr, + page_table_indices=block_tables[:, 0, :].reshape(-1), + page_size=params.tokens_per_block, + custom_mask=fwd.attention_mask_data.flatten().contiguous(), + sm_scale=bmm1_scale, + window_left=window_left, + k_cache=kv_pool, + v_cache=kv_pool, + v_page_table_indices=block_tables[:, 1, :].reshape(-1), + ) diff --git a/tensorrt_llm/_torch/attention_backend/fmha/trtllm_gen_utils.py b/tensorrt_llm/_torch/attention_backend/fmha/trtllm_gen_utils.py new file mode 100644 index 000000000000..73458519d12c --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/fmha/trtllm_gen_utils.py @@ -0,0 +1,46 @@ +# 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 functools import lru_cache + +import torch + +from tensorrt_llm.bindings.internal import thop + + +@lru_cache(maxsize=128) +def get_trtllm_gen_context_workspace_size( + dtype: torch.dtype, + max_num_seq: int, + max_num_tokens: int, + num_heads: int, + head_size: int, + rotary_embedding_dim: int, + fp8_context_fmha: bool, +) -> int: + """Return the fused context-preprocessing workspace size in bytes.""" + if max_num_tokens == 0: + return 0 + layout = thop.get_trtllm_gen_context_workspace_layout( + dtype, + max_num_seq, + max_num_tokens, + num_heads, + head_size, + rotary_embedding_dim, + True, + fp8_context_fmha, + ) + return int(layout["total_size"]) diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index 42708a4f33ac..2a5ed81ef42e 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -341,6 +341,21 @@ def prepare_encoder_only(self) -> None: """ self.prepare() + def get_shared_kv_draft_metadata( + self, + num_accepted_tokens: torch.Tensor, + num_contexts: int, + ) -> "AttentionMetadata": + """Return metadata for an external draft model sharing target KV. + + Backends supporting external shared-KV drafting override this method. + The returned view must expose only the accepted target prefix and must + not mutate the target metadata. + """ + raise NotImplementedError( + f"{type(self).__name__} does not support external shared-KV " + "draft attention.") + def _prepare_mamba_metadata(self): if self.mamba_metadata is False: return diff --git a/tensorrt_llm/_torch/attention_backend/triton_prefill.py b/tensorrt_llm/_torch/attention_backend/triton_prefill.py index f481d91a864a..d696e47889f0 100644 --- a/tensorrt_llm/_torch/attention_backend/triton_prefill.py +++ b/tensorrt_llm/_torch/attention_backend/triton_prefill.py @@ -112,8 +112,9 @@ def _fwd_kernel( # Sequence boundaries qo_indptr, # [batch+1] cumulative extend token counts prefix_indptr, # [batch+1] cumulative prefix TOKEN counts - # Page table for prefix - page_table, # [total_pages] physical page IDs + # Page tables for prefix + k_page_table, # [total_pages] physical K page IDs + v_page_table, # [total_pages] physical V page IDs page_table_indptr, # [batch+1] cumulative page counts per seq # Mask mask_ptr, @@ -240,8 +241,13 @@ def _fwd_kernel( token_positions = start_n + offs_n page_local_idx = token_positions // PAGE_SIZE token_in_page = token_positions % PAGE_SIZE - page_ids = tl.load( - page_table + cur_seq_page_start + page_local_idx, + k_page_ids = tl.load( + k_page_table + cur_seq_page_start + page_local_idx, + mask=mask_n, + other=-1, + ) + v_page_ids = tl.load( + v_page_table + cur_seq_page_start + page_local_idx, mask=mask_n, other=-1, ) @@ -250,19 +256,21 @@ def _fwd_kernel( # with mask=True still issues the memory access, so a -1 page_id # would OOB into K_Buffer/V_Buffer (NaN propagation or illegal # address fault). Filter at load and at qk merge. - valid_page = page_ids >= 0 - safe_page_ids = tl.where(valid_page, page_ids, 0) + valid_page = (k_page_ids >= 0) & (v_page_ids >= 0) + safe_k_page_ids = tl.where(valid_page, k_page_ids, 0) + safe_v_page_ids = tl.where(valid_page, v_page_ids, 0) # Promote to int64 before multiplying by stride: with KV pools at # production scale (e.g. shape [N, 2, 8, 32, 256] => stride 131072 # elements/page), page_id * stride overflows int32 once page_id # exceeds ~16K and silently wraps to a negative offset, IMA-ing # K_Buffer/V_Buffer. - safe_page_ids = safe_page_ids.to(tl.int64) + safe_k_page_ids = safe_k_page_ids.to(tl.int64) + safe_v_page_ids = safe_v_page_ids.to(tl.int64) final_mask &= valid_page[None, :] # Load K from paged buffer (transposed for dot product) offs_buf_k = ( - safe_page_ids[None, :] * stride_buf_kpage + safe_k_page_ids[None, :] * stride_buf_kpage + cur_kv_head * stride_buf_kh + token_in_page[None, :] * stride_buf_ktoken + offs_d[:, None] * stride_buf_kdim @@ -291,7 +299,7 @@ def _fwd_kernel( # Load V from paged buffer offs_buf_v = ( - safe_page_ids[:, None] * stride_buf_vpage + safe_v_page_ids[:, None] * stride_buf_vpage + cur_kv_head * stride_buf_vh + token_in_page[:, None] * stride_buf_vtoken + offs_dv[None, :] * stride_buf_vdim @@ -411,7 +419,8 @@ def _extend_attention_fwd( v_buffer: torch.Tensor, qo_indptr: torch.Tensor, prefix_indptr: torch.Tensor, - page_table: torch.Tensor, + k_page_table: torch.Tensor, + v_page_table: torch.Tensor, page_table_indptr: torch.Tensor, custom_mask: Optional[torch.Tensor], mask_indptr: Optional[torch.Tensor], @@ -434,7 +443,8 @@ def _extend_attention_fwd( v_buffer: V view of paged cache [num_pages, num_kv_heads, page_size, head_dim] qo_indptr: [batch+1] extend token boundaries prefix_indptr: [batch+1] cumulative prefix token counts - page_table: [total_pages] physical page IDs + k_page_table: [total_pages] physical K page IDs + v_page_table: [total_pages] physical V page IDs page_table_indptr: [batch+1] cumulative page counts per seq custom_mask: Flattened bool mask or None mask_indptr: [batch+1] cumulative mask element counts, or None @@ -467,7 +477,8 @@ def _extend_attention_fwd( v_buffer, qo_indptr, prefix_indptr, - page_table, + k_page_table, + v_page_table, page_table_indptr, custom_mask, mask_indptr, @@ -526,6 +537,9 @@ def triton_prefill_with_custom_mask( custom_mask: Optional[torch.Tensor], sm_scale: float, window_left: int = -1, + k_cache: Optional[torch.Tensor] = None, + v_cache: Optional[torch.Tensor] = None, + v_page_table_indices: Optional[torch.Tensor] = None, ) -> None: """Triton prefill attention with optional custom mask and paged KV cache. @@ -558,32 +572,27 @@ def triton_prefill_with_custom_mask( A query at position p attends to keys in [p - window_left, p]. Use (attention_window_size - 1) since TRTLLM window size is exclusive while flashinfer/triton convention is inclusive. + k_cache: Optional K-only paged cache + [num_pages, num_kv_heads, page_size, head_dim]. Use with + ``v_cache`` for backends that store K and V in separate pages. + v_cache: Optional V-only paged cache with the same shape as + ``k_cache``. + v_page_table_indices: Optional physical V page IDs. Defaults to + ``page_table_indices`` for caches whose K and V share page IDs. """ device = q.device num_contexts = qo_indptr.shape[0] - 1 extend_lens = qo_indptr[1:] - qo_indptr[:-1] - # Build causal mask when custom_mask is not provided total_kv_lens = prefix_lens + extend_lens + is_causal = custom_mask is None if custom_mask is None: - # Generate causal mask: each query at position i attends to - # positions [0..prefix_len+i] (all prefix + causal extend) - mask_parts = [] - for seq_idx in range(num_contexts): - ext = int(extend_lens[seq_idx].item()) - pre = int(prefix_lens[seq_idx].item()) - total = pre + ext - # Row i of extend can see columns [0..pre+i] - rows = torch.arange(ext, device=device).unsqueeze(1) - cols = torch.arange(total, device=device).unsqueeze(0) - seq_mask = cols <= (rows + pre) - mask_parts.append(seq_mask.reshape(-1)) - custom_mask = torch.cat(mask_parts) - - # Compute mask_indptr: each seq's mask is [extend_len, prefix_len + extend_len] - mask_sizes = extend_lens * total_kv_lens - mask_indptr = torch.zeros(num_contexts + 1, dtype=torch.int32, device=device) - mask_indptr[1:] = torch.cumsum(mask_sizes, dim=0) + mask_indptr = torch.empty(1, dtype=torch.int32, device=device) + else: + # Each sequence's custom mask is [extend_len, prefix_len + extend_len]. + mask_sizes = extend_lens * total_kv_lens + mask_indptr = torch.zeros(num_contexts + 1, dtype=torch.int32, device=device) + mask_indptr[1:] = torch.cumsum(mask_sizes, dim=0) # Compute prefix_indptr: cumulative prefix token counts prefix_indptr = torch.zeros(num_contexts + 1, dtype=torch.int32, device=device) @@ -595,8 +604,20 @@ def triton_prefill_with_custom_mask( # can run tl.dot in BF16. Only the accessed pages are copied. compute_dtype = q.dtype if kv_cache is not None: + if k_cache is not None or v_cache is not None: + raise ValueError("Specify either kv_cache or k_cache/v_cache, not both.") k_buffer = kv_cache.select(1, 0) # [num_pages, num_kv_heads, page_size, head_dim] v_buffer = kv_cache.select(1, 1) + elif k_cache is not None and v_cache is not None: + k_buffer = k_cache + v_buffer = v_cache + elif k_cache is not None or v_cache is not None: + raise ValueError("k_cache and v_cache must be specified together.") + else: + k_buffer = None + v_buffer = None + + if k_buffer is not None and v_buffer is not None: if k_buffer.dtype != compute_dtype: num_prefix_pages = int(page_table_indptr[-1].item()) if num_prefix_pages > 0: @@ -611,6 +632,8 @@ def triton_prefill_with_custom_mask( v_buffer = torch.empty(1, num_kv_heads, 1, head_dim, dtype=v.dtype, device=device) max_len_extend = int(extend_lens.max().item()) if num_contexts > 0 else 0 + if v_page_table_indices is None: + v_page_table_indices = page_table_indices _extend_attention_fwd( q_extend=q, @@ -621,14 +644,15 @@ def triton_prefill_with_custom_mask( v_buffer=v_buffer, qo_indptr=qo_indptr, prefix_indptr=prefix_indptr, - page_table=page_table_indices, + k_page_table=page_table_indices, + v_page_table=v_page_table_indices, page_table_indptr=page_table_indptr, custom_mask=custom_mask, mask_indptr=mask_indptr, max_len_extend=max_len_extend, sm_scale=sm_scale, page_size=page_size, - is_causal=False, + is_causal=is_causal, logit_cap=0.0, skip_prefix_custom_mask=False, window_left=window_left, diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 70f8394ecf7b..71c68cebe3ae 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import functools import math import os @@ -30,7 +31,7 @@ from ..speculative.spec_tree_manager import SpecTreeManager from tensorrt_llm._torch.attention_backend.fmha import ( - Fmha, get_enabled_fmha_lib_classes) + CombinedFmha, Fmha, FmhaPhase, PhasedFmha, 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 @@ -82,6 +83,22 @@ class TrtllmAttentionMetadata(AttentionMetadata): workspace: Optional[torch.Tensor] = None cuda_graph_workspace: Optional[torch.Tensor] = None + _draft_metadata: Optional["TrtllmAttentionMetadata"] = field(init=False, + default=None, + repr=False) + _draft_kv_runtime_lens: Optional[torch.Tensor] = field(init=False, + default=None, + repr=False) + _shared_kv_draft_seq_lens: Optional[torch.Tensor] = field(init=False, + default=None, + repr=False) + _shared_kv_draft_seq_lens_cuda: Optional[torch.Tensor] = field(init=False, + default=None, + repr=False) + _shared_kv_draft_request_types: Optional[torch.Tensor] = field(init=False, + default=None, + repr=False) + # TrtllmAttention needs to know the beam width to access to the cache indirection buffer, # when beam search is enabled. beam_width: int = 1 @@ -273,6 +290,205 @@ def __post_init__(self) -> None: ) if self.runtime_features is not None else False self._post_init_with_buffers(self.cuda_graph_buffers) + def create_cuda_graph_metadata( + self, + max_batch_size: int, + sub_cross_metadata: bool = False, + max_draft_tokens: int = 0, + buffers=None, + encode_only: bool = False) -> "TrtllmAttentionMetadata": + if self.is_cuda_graph: + return self + metadata = super().create_cuda_graph_metadata( + max_batch_size, + sub_cross_metadata, + max_draft_tokens, + buffers, + encode_only, + ) + # A graph metadata object owns its lazy draft view and all of that + # view's fixed-address buffers. Do not inherit the eager view through + # AttentionMetadata's shallow copy. + metadata._draft_metadata = None + return metadata + + def _validate_shared_kv_draft_metadata(self, + num_accepted_tokens: torch.Tensor, + num_contexts: int) -> None: + if self.is_cross: + raise ValueError( + "TRTLLM external shared-KV drafting does not support cross attention." + ) + if self.kv_cache_manager is None: + raise ValueError( + "TRTLLM external shared-KV drafting requires a paged KV cache.") + if getattr(self.kv_cache_manager, "kv_factor", 2) != 2: + raise ValueError( + "TRTLLM external shared-KV drafting does not support MLA.") + if self.sparse_metadata_params is not None or self.num_sparse_topk > 0: + raise ValueError( + "TRTLLM external shared-KV drafting does not support sparse attention." + ) + if self.beam_width != 1: + raise ValueError( + "TRTLLM external shared-KV drafting does not support beam search." + ) + if self.draft_kv_cache_manager is not None: + raise ValueError( + "TRTLLM external shared-KV drafting cannot use a separate draft KV cache." + ) + if self.kv_cache_block_offsets is None: + raise ValueError( + "TRTLLM external shared-KV drafting requires KV-cache block offsets." + ) + if num_contexts != self.num_contexts: + raise ValueError( + "num_contexts must match the target attention metadata: " + f"got {num_contexts}, expected {self.num_contexts}.") + if num_accepted_tokens.ndim != 1 or num_accepted_tokens.shape[ + 0] < self.num_seqs: + raise ValueError( + "num_accepted_tokens must be a 1-D tensor covering every sequence." + ) + if self.kv_lens_cuda_runtime is None or self.kv_lens_runtime is None: + raise RuntimeError( + "TRTLLM target attention metadata must be prepared before " + "requesting a shared-KV draft view.") + if num_accepted_tokens.device != self.kv_lens_cuda_runtime.device: + raise ValueError( + "num_accepted_tokens and target KV lengths must be on the same device." + ) + + def _create_shared_kv_draft_metadata(self) -> "TrtllmAttentionMetadata": + draft_metadata = copy.copy(self) + draft_metadata._draft_metadata = None + draft_metadata.cross = None + draft_metadata.cuda_graph_buffers = None + draft_metadata._saved_tensors = {} + + max_num_sequences = self.max_num_sequences or self.max_num_requests + draft_metadata._shared_kv_draft_seq_lens = torch.ones( + (max_num_sequences, ), + dtype=torch.int, + device="cpu", + pin_memory=prefer_pinned(), + ) + draft_metadata._shared_kv_draft_seq_lens_cuda = torch.ones( + (max_num_sequences, ), + dtype=self.kv_lens_cuda.dtype, + device=self.kv_lens_cuda.device, + ) + draft_metadata._shared_kv_draft_request_types = torch.ones( + (max_num_sequences, ), + dtype=self.host_request_types.dtype, + device="cpu", + pin_memory=prefer_pinned(), + ) + draft_metadata._draft_kv_runtime_lens = torch.empty( + (max_num_sequences, ), + dtype=self.kv_lens_cuda.dtype, + device=self.kv_lens_cuda.device, + ) + draft_metadata.host_total_kv_lens = torch.empty_like( + self.host_total_kv_lens) + + # The assistant issues Q-only generation requests. It must not inherit + # speculative-tree state from the target forward. + draft_metadata.is_spec_decoding_enabled = False + draft_metadata.use_spec_decoding = False + draft_metadata.is_spec_dec_tree = False + draft_metadata.is_spec_dec_dynamic_tree = False + draft_metadata.force_prepare_spec_dec_tree_mask = False + draft_metadata.max_total_draft_tokens = None + draft_metadata.spec_decoding_position_offsets = None + draft_metadata.spec_decoding_position_offsets_cpp = None + draft_metadata.spec_decoding_packed_mask = None + draft_metadata.spec_decoding_generation_lengths = None + draft_metadata.spec_decoding_bl_tree_mask_offset = None + draft_metadata.spec_decoding_bl_tree_mask = None + draft_metadata.spec_bl_tree_first_sparse_mask_offset_kv = None + draft_metadata.position_offsets_stride = 0 + draft_metadata.draft_kv_cache_manager = None + draft_metadata.draft_kv_cache_block_offsets = None + draft_metadata.padded_num_tokens = None + return draft_metadata + + def _sync_shared_kv_draft_metadata( + self, target: "TrtllmAttentionMetadata") -> None: + """Refresh this cached draft view from prepared target metadata.""" + if (self._shared_kv_draft_seq_lens is None + or self._shared_kv_draft_seq_lens_cuda is None + or self._shared_kv_draft_request_types is None + or self._draft_kv_runtime_lens is None): + raise RuntimeError("Shared-KV draft buffers are not initialized.") + + num_seqs = target.num_seqs + if num_seqs > self._shared_kv_draft_seq_lens.shape[0]: + raise ValueError( + f"Draft batch size {num_seqs} exceeds the allocated capacity " + f"{self._shared_kv_draft_seq_lens.shape[0]}.") + + self._seq_lens = self._shared_kv_draft_seq_lens[:num_seqs] + self._seq_lens_cuda = self._shared_kv_draft_seq_lens_cuda[:num_seqs] + self._seq_lens_kv = None + self._seq_lens_kv_cuda = None + self._num_contexts = 0 + self._num_generations = num_seqs + self._num_ctx_tokens = 0 + self._num_tokens = num_seqs + + self.kv_cache_manager = target.kv_cache_manager + self.kv_cache_params = target.kv_cache_params + self.kv_cache_block_offsets = target.kv_cache_block_offsets + self.host_kv_cache_block_offsets = target.host_kv_cache_block_offsets + self.workspace = target.workspace + self.cuda_graph_workspace = target.cuda_graph_workspace + self.request_ids = target.request_ids + self.prompt_lens = target.prompt_lens + self.all_rank_num_tokens = target.all_rank_num_tokens + + # Until accepted lengths are published, the complete target lengths + # are a safe initialized value for the view. The generic getter below + # replaces generation rows with the exact accepted prefix. + self._draft_kv_runtime_lens[:num_seqs].copy_( + target.kv_lens_cuda_runtime[:num_seqs]) + self._bind_runtime_views( + kv_lens_cuda=self._draft_kv_runtime_lens[:num_seqs], + kv_lens=target.kv_lens_runtime[:num_seqs], + prompt_lens_cuda=target.prompt_lens_cuda_runtime[:num_seqs], + prompt_lens_cpu=target.prompt_lens_cpu_runtime[:num_seqs], + host_request_types=self._shared_kv_draft_request_types[:num_seqs], + ) + self.host_total_kv_lens[0] = 0 + self.host_total_kv_lens[1] = target.host_total_kv_lens.sum() + + def get_shared_kv_draft_metadata( + self, + num_accepted_tokens: torch.Tensor, + num_contexts: int, + ) -> "TrtllmAttentionMetadata": + """Return a Q-only generation view over the accepted target KV.""" + self._validate_shared_kv_draft_metadata(num_accepted_tokens, + num_contexts) + if self._draft_metadata is None: + self._draft_metadata = self._create_shared_kv_draft_metadata() + + draft_metadata = self._draft_metadata + draft_metadata._sync_shared_kv_draft_metadata(self) + num_seqs = self.num_seqs + if num_contexts > 0: + draft_metadata._draft_kv_runtime_lens[:num_contexts].copy_( + self.kv_lens_cuda_runtime[:num_contexts]) + generation_slice = slice(num_contexts, num_seqs) + torch.sub( + self.kv_lens_cuda_runtime[generation_slice], + self.seq_lens_kv_cuda[generation_slice], + out=draft_metadata._draft_kv_runtime_lens[generation_slice], + ) + draft_metadata._draft_kv_runtime_lens[generation_slice].add_( + num_accepted_tokens[generation_slice]) + return draft_metadata + def update_position_offsets_for_cpp(self, query_len: int) -> None: """Refresh the C++ view of spec-dec position offsets.""" offsets = self.spec_decoding_position_offsets @@ -786,6 +1002,8 @@ def prepare(self) -> None: prompt_lens_cpu=self.prompt_lens_cpu[:self.num_seqs], host_request_types=self.host_request_types[:self.num_seqs], ) + if self._draft_metadata is not None: + self._draft_metadata._sync_shared_kv_draft_metadata(self) def prepare_encoder_decoder_from_precomputed_lengths( self, prompt_lens: torch.Tensor, kv_lens: torch.Tensor, @@ -1487,6 +1705,9 @@ def __init__( self.local_layer_idx: Optional[int] = None self.fmha_libs: List[Fmha] = [] + self.phased_fmha_libs: List[PhasedFmha] = [] + self.non_phased_fmha_libs: List[Fmha] = [] + self.combined_fmha: Optional[CombinedFmha] = None if not skip_create_weights_in_init: self.update_quant_config(self.quant_config) @@ -1747,6 +1968,100 @@ def create_fmha_libs(self) -> None: if fmha_cls.is_available(self): self.fmha_libs.append(fmha_cls(self)) + self.phased_fmha_libs = [ + fmha for fmha in self.fmha_libs if isinstance(fmha, PhasedFmha) + ] + self.non_phased_fmha_libs = [ + fmha for fmha in self.fmha_libs if not isinstance(fmha, PhasedFmha) + ] + self.combined_fmha = (CombinedFmha(self) if self.phased_fmha_libs + and not self.is_mla_enable else None) + + def _select_non_mla_fmha( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: TrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> Optional[Fmha]: + has_context = metadata.num_contexts > 0 + has_generation = metadata.num_generations > 0 + if not has_context and not has_generation: + return None + + context_fmha = None + generation_fmha = None + for fmha in self.fmha_libs: + if isinstance(fmha, PhasedFmha): + if (has_context and context_fmha is None and fmha.is_supported( + q, + k, + v, + metadata, + forward_args, + phase=FmhaPhase.CONTEXT, + )): + context_fmha = fmha + if (has_generation and generation_fmha is None + and fmha.is_supported( + q, + k, + v, + metadata, + forward_args, + phase=FmhaPhase.GENERATION, + )): + generation_fmha = fmha + + if (has_context and context_fmha is None) or ( + has_generation and generation_fmha is None): + continue + if context_fmha is None: + return generation_fmha + if generation_fmha is None or context_fmha is generation_fmha: + return context_fmha + + self.combined_fmha.set_fmha_impls( + context_fmha, + generation_fmha, + ) + return self.combined_fmha + if fmha.is_supported(q, k, v, metadata, forward_args): + return fmha + return None + + def _select_mla_fmha( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: TrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> Optional[Fmha]: + if forward_args.attention_input_type == AttentionInputType.context_only: + phase = FmhaPhase.CONTEXT + elif forward_args.attention_input_type == AttentionInputType.generation_only: + phase = FmhaPhase.GENERATION + else: + return None + + for fmha in self.fmha_libs: + if isinstance(fmha, PhasedFmha): + supported = fmha.is_supported( + q, + k, + v, + metadata, + forward_args, + phase=phase, + ) + else: + supported = fmha.is_supported(q, k, v, metadata, forward_args) + if supported: + return fmha + return None + def forward( self, q: torch.Tensor, @@ -1796,15 +2111,24 @@ def forward( 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_q_only = False + if not self.is_mla_enable and not metadata.is_cross and k is None and v is None: + q_hidden_size = self.num_heads * self.head_dim + qkv_hidden_size = q_hidden_size + 2 * self.num_kv_heads * self.head_dim + has_q_only = q.size(-1) == q_hidden_size + forward_args.is_fused_qkv = q.size(-1) == qkv_hidden_size + forward_args.update_kv_cache = not has_q_only + else: + 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 + assert (has_fused_qkv or has_unfused_kv or has_q_only + or uses_cached_cross_kv) # `quant_scale_qkv` only makes sense paired with `quant_q_buffer`: the # C++ op interprets the buffer as the destination of a pre-quantized # FP8 Q for the DSv4 fused norm+RoPE path. `quant_q_buffer` alone is @@ -2045,13 +2369,22 @@ def forward( 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 + fmha: Optional[Fmha] = None + if self.is_mla_enable: + fmha = self._select_mla_fmha(q, k, v, metadata, forward_args) else: + fmha = self._select_non_mla_fmha( + q, + k, + v, + metadata, + forward_args, + ) + + if fmha is None: raise RuntimeError( "No TRT-LLM attention FMHA library supports this request.") + fmha.forward(q, k, v, metadata, forward_args) if self.print_skip_softmax_stat: total_blocks, skipped_blocks = self.skip_softmax_stat diff --git a/tensorrt_llm/_torch/models/modeling_cohere2.py b/tensorrt_llm/_torch/models/modeling_cohere2.py index f228b75e9374..c2a538791343 100644 --- a/tensorrt_llm/_torch/models/modeling_cohere2.py +++ b/tensorrt_llm/_torch/models/modeling_cohere2.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. + from typing import Optional import torch @@ -6,7 +21,11 @@ from transformers import Cohere2Config from transformers.activations import ACT2FN -from tensorrt_llm._torch.attention_backend import AttentionMetadata, FlashInferAttentionMetadata +from tensorrt_llm._torch.attention_backend import ( + AttentionMetadata, + FlashInferAttentionMetadata, + TrtllmAttentionMetadata, +) from tensorrt_llm._torch.attention_backend.interface import ( AttentionMask, CustomAttentionMask, @@ -114,9 +133,10 @@ def forward( **kwargs, ): if attention_mask_data is not None: - assert isinstance(attn_metadata, FlashInferAttentionMetadata), ( - "Only FlashInfer backend supports custom attention mask currently." - ) + assert isinstance( + attn_metadata, + (FlashInferAttentionMetadata, TrtllmAttentionMetadata), + ), "Only FlashInfer and TRTLLM backends support custom attention masks." assert attention_mask == CustomAttentionMask.CUSTOM return super().forward( position_ids=position_ids, diff --git a/tensorrt_llm/_torch/models/modeling_gemma3.py b/tensorrt_llm/_torch/models/modeling_gemma3.py index 628d1939797a..9de6a9f23c1e 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma3.py +++ b/tensorrt_llm/_torch/models/modeling_gemma3.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 math from typing import Any, Dict, Literal, Optional, Tuple @@ -11,7 +26,8 @@ from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType from tensorrt_llm.mapping import Mapping -from ..attention_backend import AttentionMetadata, FlashInferAttentionMetadata +from ..attention_backend import (AttentionMetadata, FlashInferAttentionMetadata, + TrtllmAttentionMetadata) from ..attention_backend.interface import (AttentionMask, CustomAttentionMask, PositionalEmbeddingParams, PredefinedAttentionMask, RopeParams) @@ -117,8 +133,9 @@ def forward( if attention_mask_data is not None: assert isinstance( - attn_metadata, FlashInferAttentionMetadata - ), "Only FlashInfer backend supports custom attention mask currently." + attn_metadata, + (FlashInferAttentionMetadata, TrtllmAttentionMetadata), + ), "Only FlashInfer and TRTLLM backends support custom attention masks." assert attention_mask == CustomAttentionMask.CUSTOM return super().forward(position_ids=position_ids, hidden_states=hidden_states, @@ -379,14 +396,15 @@ def get_context_mask( # ASSUMPTIONS: # 1) Chunked prefill is disabled to avoid chunking image tokens as they need bidirectional attention. # 2) KV cache reuse is disabled to avoid partially matched image tokens (entire image must be reused to get things correct). - def get_flashinfer_attention_mask( + def get_attention_mask( self, image_token_mask: torch.BoolTensor, attn_metadata: AttentionMetadata, effective_sliding_window: Optional[int] = None) -> torch.Tensor: """ - This is specifically needed for context phase requests. Currently, we don't create custom mask for generation requests because FlashInfer backend - doesn't use it anyway and there's nothing special we need to do for generation requests. + This is specifically needed for context phase requests. Generation + requests use the regular causal attention semantics and do not need a + custom mask. - This function will only be called for a batch when there's at least one context request in the batch with image tokens. - In context phase, each sample's input_ids may have a mix of image tokens and text tokens where tokens corresponding to an image appear as a contiguous blob. Example: torch.IntTensor([2, 3, 4, 5, img_idx, img_idx, img_idx, ..., img_idx, 100]) @@ -402,26 +420,31 @@ def get_flashinfer_attention_mask( A flattened boolean mask of shape (sum(q_len[i] * k_len[i] for i in range(batch_size)). """ - assert isinstance( - attn_metadata, FlashInferAttentionMetadata - ), "Only FlashInfer backend supports custom mask currently." num_contexts = attn_metadata.num_contexts assert num_contexts > 0, "There should be at least one context request in the batch for custom mask." - qo_indptr = attn_metadata.qo_indptr[:num_contexts + 1] - cached_token_lens = attn_metadata.cached_token_lens[:num_contexts] - assert (cached_token_lens == 0).all( - ), "cached_token_lens should be 0 for context requests since chunked prefill and kv cache reuse must be disabled." + context_lens = attn_metadata.seq_lens[:num_contexts].tolist() + cached_token_lens = ( + attn_metadata.kv_cache_params. + num_cached_tokens_per_seq[:num_contexts] + if attn_metadata.kv_cache_params is not None and + attn_metadata.kv_cache_params.num_cached_tokens_per_seq is not None + else [0] * num_contexts) + assert not any(cached_token_lens), ( + "cached_token_lens should be 0 for context requests since chunked prefill " + "and kv cache reuse must be disabled.") # Create masks for context requests. context_mask_list = [] - for i in range(num_contexts): + token_offset = 0 + for context_len in context_lens: mask_i = self.get_context_mask( - image_token_mask=image_token_mask[qo_indptr[i]:qo_indptr[i + - 1]], + image_token_mask=image_token_mask[token_offset:token_offset + + context_len], effective_sliding_window=effective_sliding_window, ) context_mask_list.append(mask_i.flatten()) + token_offset += context_len return torch.cat(context_mask_list, dim=0).contiguous() @inference_mode_unless_compiling @@ -439,12 +462,12 @@ def forward( local_attention_mask_data = None global_attention_mask_data = None if image_token_mask is not None: - global_attention_mask_data = self.get_flashinfer_attention_mask( + global_attention_mask_data = self.get_attention_mask( image_token_mask=image_token_mask, attn_metadata=attn_metadata, effective_sliding_window=None, ) - local_attention_mask_data = self.get_flashinfer_attention_mask( + local_attention_mask_data = self.get_attention_mask( image_token_mask=image_token_mask, attn_metadata=attn_metadata, effective_sliding_window=self.config.sliding_window, diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index 1834217b57bc..8e86fc297c17 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -34,7 +34,11 @@ from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping -from ..attention_backend import AttentionMetadata, FlashInferAttentionMetadata +from ..attention_backend import ( + AttentionMetadata, + FlashInferAttentionMetadata, + TrtllmAttentionMetadata, +) from ..attention_backend.interface import ( AttentionMask, CustomAttentionMask, @@ -283,6 +287,14 @@ def __init__( dense_bias=False, config=model_config, q_scaling=q_scaling, + # Full-attention layers use proportional RoPE whose active + # frequencies are paired across the full head. Apply it at the + # module layer because fused preprocessing cannot represent that + # pairing with only the logical rotary dimension. + # KV-shared layers read a frozen cache prefix, so their explicit + # position_ids—not the readable KV length—must select the query's + # RoPE position. + rope_fusion=is_sliding and not is_kv_shared, ) # Restore original config head_dim @@ -331,10 +343,11 @@ def __init__( # the rotate split, matching HF's rotate_half(head_dim//2) pairing. self.rotary_emb.head_dim = layer_head_dim - # trtllm-gen FMHA kernels are available only on datacenter Blackwell. - # Use FlashInfer FA2 on other architectures, including SM120/SM121; - # multimodal custom masks then use FlashInfer's native mask planning. - self.attn.flashinfer_backend = "trtllm-gen" if is_sm_100f() else "fa2" + # When FlashInfer is explicitly selected, use trtllm-gen only on + # datacenter Blackwell. The default TRTLLM backend selects phased FMHA + # directly and does not expose this FlashInfer option. + if hasattr(self.attn, "flashinfer_backend"): + self.attn.flashinfer_backend = "trtllm-gen" if is_sm_100f() else "fa2" # KV shared layers: use target layer's index for KV cache access # so the attention backend reads from the target layer's cache slot. @@ -499,9 +512,10 @@ def forward( **kwargs, ) -> torch.Tensor: if attention_mask_data is not None: - assert isinstance(attn_metadata, FlashInferAttentionMetadata), ( - "Only FlashInfer backend supports custom attention mask currently." - ) + assert isinstance( + attn_metadata, + (FlashInferAttentionMetadata, TrtllmAttentionMetadata), + ), "Only FlashInfer and TRTLLM backends support custom attention masks." assert attention_mask == CustomAttentionMask.CUSTOM # Custom-mask (multimodal) prefill uses the Triton prefill fallback, # which consumes BF16 q/k/v — keep the unfused prep for those calls. @@ -1254,12 +1268,19 @@ def __init__( def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: """Gemma4-specific defaults. - FlashInfer backend is required for hybrid attention (per-layer - head_dim 256/512 with VSWA), trtllm-gen cubin dispatch, and - bidirectional attention masks for multimodal tokens. + The TRTLLM attention backend uses trtllm-gen for the regular attention + phases and a Triton context phase for bidirectional multimodal masks. + External shared-KV MTP keeps FlashInfer as its default because the + production FP8-KV multimodal path is not yet supported by TRTLLM's + custom-mask context implementation. """ + speculative_config = getattr(llm_args, "speculative_config", None) + spec_dec_mode = getattr(speculative_config, "spec_dec_mode", None) + uses_external_shared_kv = ( + spec_dec_mode is not None and spec_dec_mode.is_mtp_eagle_one_model() + ) return { - "attn_backend": "FLASHINFER", + "attn_backend": "FLASHINFER" if uses_external_shared_kv else "TRTLLM", } @classmethod @@ -1351,16 +1372,13 @@ def get_context_mask( return causal_mask - def get_flashinfer_attention_mask( + def get_attention_mask( self, mm_token_type_ids: torch.Tensor, attn_metadata: AttentionMetadata, effective_sliding_window: Optional[int] = None, ) -> torch.Tensor: - """Build FlashInfer custom mask for context requests.""" - assert isinstance(attn_metadata, FlashInferAttentionMetadata), ( - "Only FlashInfer backend supports custom mask currently." - ) + """Build a custom attention mask for context requests.""" num_contexts = attn_metadata.num_contexts assert num_contexts > 0 @@ -1449,7 +1467,7 @@ def forward( # full-attention layers retain their normal causal mask. use_bidir = getattr(self.config, "use_bidirectional_attention", None) if mm_token_type_ids is not None and use_bidir == "vision": - local_attention_mask_data = self.get_flashinfer_attention_mask( + local_attention_mask_data = self.get_attention_mask( mm_token_type_ids=mm_token_type_ids, attn_metadata=attn_metadata, effective_sliding_window=self.config.sliding_window, @@ -1608,6 +1626,17 @@ def __init__(self, model_config: ModelConfig): hidden_size=assistant_config.hidden_size, vocab_size=assistant_config.vocab_size, ) + # The assistant is invoked with draft-specific attention metadata. + # Keep its attention calls eager under torch.compile / breakable CUDA + # graphs so they consume that explicit metadata instead of the target + # model metadata registered in the custom attention op. + for layer in self.model.layers: + attention = layer.self_attn + if attention.register_to_config: + text_model_config.extra_attrs.get("attn_layers", {}).pop( + attention.layer_idx_str, None + ) + attention.register_to_config = False self.pre_projection = Linear( 2 * assistant_config.backbone_hidden_size, assistant_config.hidden_size, diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index a667130b0312..2588c34ee05a 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -562,9 +562,7 @@ class Gemma4MultimodalModelBase(MultimodalModelMixin, PreTrainedModel): @classmethod def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: """Gemma4-specific defaults — see Gemma4ForCausalLM.get_model_defaults.""" - return { - "attn_backend": "FLASHINFER", - } + return Gemma4ForCausalLM.get_model_defaults(llm_args) @classmethod def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: @@ -855,11 +853,11 @@ def get_sub_model_config( ) pretrained_config = getattr(model_config.pretrained_config, name) quant_config = model_config.quant_config if name == "text_config" else None - preferred_backend = "FLASHINFER" if name == "text_config" else "TRTLLM" + attn_backend = model_config.attn_backend if name == "text_config" else "TRTLLM" sub_config: ModelConfig = dataclasses.replace( model_config, pretrained_config=pretrained_config, - attn_backend=preferred_backend, + attn_backend=attn_backend, quant_config=quant_config, ) if ( diff --git a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md index 844a58fe1974..9d63d46152a7 100644 --- a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md @@ -300,31 +300,41 @@ The main differences across backends: #### 3.2.2 `TRTLLM` internal FMHA libraries `TrtllmAttention` dispatches attention through an ordered list of internal FMHA -libraries. `CuteDslMlaFmha` integrates Blackwell CuTe DSL MLA decode kernels, -`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. +libraries. `TritonCustomMaskFmha` provides Triton context attention for custom +masks, `CuteDslMlaFmha` integrates Blackwell CuTe DSL MLA decode kernels, +`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 -`cute_dsl_mla,msa_sparse_gqa,flashinfer_trtllm_gen,fallback`; use `TLLM_FMHA_LIBS=fallback` -or `TLLM_FMHA_LIBS=-cute_dsl_mla,-msa_sparse_gqa,-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. +`triton_custom_mask,cute_dsl_mla,msa_sparse_gqa,flashinfer_trtllm_gen,fallback`; +use `TLLM_FMHA_LIBS=fallback` to force the fallback path. Each FMHA library +exposes `is_available()` for module/static environment checks and +`is_supported()` for per-forward request checks. +For mixed non-MLA batches, the dispatcher checks each active phase independently +with `is_supported(..., phase=...)`; a phased library accepts only phases backed +by its corresponding `run_*()` entry point. 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/phased.py` defines `PhasedFmha`, shared phase splitting, and the + context/generation and MHA/MLA entry points. +- `fmha/combined.py` composes different context and generation implementations + for non-MLA mixed batches. +- `fmha/triton_custom_mask.py` implements the Triton custom-mask context phase. + Custom-mask data applies to context requests; for mixed batches, + `TrtllmAttention` can pair it with a later causal-generation provider through + `CombinedFmha`. - `fmha/cute_dsl_mla.py` implements the CuTe DSL MLA decode FMHA library. - `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. +Use `PhasedFmha` for libraries that implement one or more phase-specific entry +points. Use `Fmha` directly for libraries that already own the full request +shape. #### 3.2.3 MLA cached-context semantics diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index d497862e0c8e..c4d6590ab6ba 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -13,7 +13,6 @@ from tensorrt_llm.mapping import Mapping from ..attention_backend import AttentionMetadata -from ..attention_backend.flashinfer import FlashInferAttentionMetadata from ..model_config import ModelConfig from ..pyexecutor.llm_request import LlmRequest from ..pyexecutor.mamba_cache_manager import MambaHybridCacheManager @@ -838,11 +837,6 @@ def _forward_external_shared_target_kv_draft_loop( batch_size, ): """Draft with an external Q-only assistant over accepted target KV.""" - if not isinstance(attn_metadata, FlashInferAttentionMetadata): - raise TypeError( - "External shared-target-KV MTP requires FlashInfer attention " - "metadata.") - ( draft_input_ids, recurrent_hidden_states, @@ -857,10 +851,8 @@ def _forward_external_shared_target_kv_draft_loop( batch_indices=spec_metadata.batch_indices_cuda[:batch_size], ) - draft_metadata = attn_metadata.get_draft_metadata() - draft_metadata.update_shared_kv_draft_lengths(attn_metadata, - num_accepted_tokens, - num_contexts) + draft_metadata = attn_metadata.get_shared_kv_draft_metadata( + num_accepted_tokens, num_contexts) draft_metadata.use_spec_decoding = False draft_metadata.padded_num_tokens = None draft_metadata.all_rank_num_tokens = ( diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 941f02e0f27a..988aafaa685f 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -1274,7 +1274,7 @@ def maybe_compute_mm_embed_cumsum( # `multimodal_hashing_process` already setdefault()ed the cumsum (without # this flag) leaves scheduler_v2 reading mm_bidirectional_blocks=None and # silently skipping align — which then trips - # get_flashinfer_attention_mask's `cached_token_lens == 0` assert on + # get_attention_mask's `cached_token_lens == 0` assert on # 26B/31B once an image block is split across chunks. Gemma4 sets this # per-instance from text_config.use_bidirectional_attention. mm_data.setdefault( diff --git a/tests/unittest/_torch/attention/test_combined_fmha.py b/tests/unittest/_torch/attention/test_combined_fmha.py new file mode 100644 index 000000000000..2fd95c4447ba --- /dev/null +++ b/tests/unittest/_torch/attention/test_combined_fmha.py @@ -0,0 +1,303 @@ +# 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 types import SimpleNamespace + +import torch + +from tensorrt_llm._torch.attention_backend.fmha.combined import CombinedFmha +from tensorrt_llm._torch.attention_backend.fmha.flashinfer_trtllm_gen import FlashInferTrtllmGenFmha +from tensorrt_llm._torch.attention_backend.fmha.interface import Fmha, FmhaPhase +from tensorrt_llm._torch.attention_backend.fmha.phased import FmhaParams, PhasedFmha +from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs, AttentionInputType +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention +from tensorrt_llm.quantization.mode import QuantMode + + +class _TestAttention: + def __init__(self) -> None: + self.is_mla_enable = False + self.kv_lora_rank = None + self.v_head_dim = None + self.head_dim = 4 + self.num_heads = 1 + self.num_kv_heads = 1 + self.predicted_tokens_per_seq = 1 + + +class _TestPhasedFmha(PhasedFmha): + def __init__( + self, + attn: _TestAttention, + supported_phases: set[FmhaPhase], + name: str, + events: list[tuple], + workspace_size: int = 0, + ) -> None: + super().__init__(attn) + self._supported_phases = supported_phases + self._name = name + self._events = events + self._workspace_size = workspace_size + + def is_supported( + self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + metadata: object, + forward_args: AttentionForwardArgs, + *, + phase: FmhaPhase | None = None, + ) -> bool: + self._events.append(("support", self._name, phase)) + return phase in self._supported_phases + + def prepare_workspace( + self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + metadata: object, + forward_args: AttentionForwardArgs, + workspace: torch.Tensor, + ) -> None: + self._events.append(("prepare", self._name)) + if workspace.numel() < self._workspace_size: + workspace.resize_(self._workspace_size) + + def run_context(self, params: FmhaParams) -> None: + self._events.append(("run", self._name, FmhaPhase.CONTEXT, params.num_tokens)) + + def run_generation(self, params: FmhaParams) -> None: + self._events.append(("run", self._name, FmhaPhase.GENERATION, params.num_tokens)) + + +class _TestFmha(Fmha): + def __init__(self, attn: _TestAttention, name: str, events: list[tuple]) -> None: + super().__init__(attn) + self._name = name + self._events = events + + def is_supported( + self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + metadata: object, + forward_args: AttentionForwardArgs, + *, + phase: FmhaPhase | None = None, + ) -> bool: + self._events.append(("support", self._name, phase)) + return True + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + metadata: object, + forward_args: AttentionForwardArgs, + ) -> None: + self._events.append(("forward", self._name)) + + +def test_select_non_mla_fmha_combines_supported_phases() -> None: + events: list[tuple] = [] + attn = _TestAttention() + context_fmha = _TestPhasedFmha(attn, {FmhaPhase.CONTEXT}, "context", events) + generation_fmha = _TestPhasedFmha(attn, {FmhaPhase.GENERATION}, "generation", events) + combined_fmha = CombinedFmha(attn) + backend = SimpleNamespace( + fmha_libs=[context_fmha, generation_fmha], + phased_fmha_libs=[context_fmha, generation_fmha], + combined_fmha=combined_fmha, + ) + + selected = TrtllmAttention._select_non_mla_fmha( + backend, + torch.empty((2, 4)), + None, + None, + SimpleNamespace(num_contexts=1, num_generations=1), + AttentionForwardArgs(attention_input_type=AttentionInputType.mixed), + ) + + assert selected is combined_fmha + assert events == [ + ("support", "context", FmhaPhase.CONTEXT), + ("support", "context", FmhaPhase.GENERATION), + ("support", "generation", FmhaPhase.GENERATION), + ] + + +def test_select_non_mla_fmha_checks_followup_support() -> None: + events: list[tuple] = [] + attn = _TestAttention() + context_fmha = _TestPhasedFmha(attn, {FmhaPhase.CONTEXT}, "context", events) + unsupported_followup = _TestPhasedFmha(attn, set(), "followup", events) + backend = SimpleNamespace( + fmha_libs=[context_fmha, unsupported_followup], + phased_fmha_libs=[context_fmha, unsupported_followup], + combined_fmha=CombinedFmha(attn), + ) + + selected = TrtllmAttention._select_non_mla_fmha( + backend, + torch.empty((2, 4)), + None, + None, + SimpleNamespace(num_contexts=1, num_generations=1), + AttentionForwardArgs(attention_input_type=AttentionInputType.mixed), + ) + + assert selected is None + assert events[-1] == ("support", "followup", FmhaPhase.GENERATION) + + +def test_select_non_mla_fmha_reuses_one_implementation() -> None: + events: list[tuple] = [] + attn = _TestAttention() + fmha = _TestPhasedFmha( + attn, + {FmhaPhase.CONTEXT, FmhaPhase.GENERATION}, + "both", + events, + ) + backend = SimpleNamespace( + fmha_libs=[fmha], + phased_fmha_libs=[fmha], + combined_fmha=CombinedFmha(attn), + ) + + selected = TrtllmAttention._select_non_mla_fmha( + backend, + torch.empty((2, 4)), + None, + None, + SimpleNamespace(num_contexts=1, num_generations=1), + AttentionForwardArgs(attention_input_type=AttentionInputType.mixed), + ) + + assert selected is fmha + + +def test_select_non_mla_fmha_preserves_registry_order() -> None: + events: list[tuple] = [] + attn = _TestAttention() + non_phased_fmha = _TestFmha(attn, "non-phased", events) + phased_fmha = _TestPhasedFmha( + attn, + {FmhaPhase.CONTEXT, FmhaPhase.GENERATION}, + "phased", + events, + ) + backend = SimpleNamespace( + fmha_libs=[non_phased_fmha, phased_fmha], + phased_fmha_libs=[phased_fmha], + combined_fmha=CombinedFmha(attn), + ) + + selected = TrtllmAttention._select_non_mla_fmha( + backend, + torch.empty((2, 4)), + None, + None, + SimpleNamespace(num_contexts=1, num_generations=1), + AttentionForwardArgs(attention_input_type=AttentionInputType.mixed), + ) + + assert selected is non_phased_fmha + assert events == [("support", "non-phased", None)] + + +def test_combined_fmha_delegates_phases_and_prepares_max_workspace() -> None: + events: list[tuple] = [] + attn = _TestAttention() + context_fmha = _TestPhasedFmha( + attn, + {FmhaPhase.CONTEXT}, + "context", + events, + workspace_size=8, + ) + generation_fmha = _TestPhasedFmha( + attn, + {FmhaPhase.GENERATION}, + "generation", + events, + workspace_size=4, + ) + combined_fmha = CombinedFmha(attn) + combined_fmha.set_fmha_impls(context_fmha, generation_fmha) + metadata = SimpleNamespace( + kv_cache_block_offsets=object(), + effective_workspace=torch.empty(0, dtype=torch.uint8), + num_contexts=1, + num_ctx_tokens=2, + num_generations=1, + kv_lens_cuda_runtime=torch.tensor([2, 5], dtype=torch.int32), + kv_lens_runtime=torch.tensor([2, 5], dtype=torch.int32), + prompt_lens_cuda_runtime=torch.tensor([2, 1], dtype=torch.int32), + prompt_lens_cpu_runtime=torch.tensor([2, 1], dtype=torch.int32), + beam_width=1, + cache_indirection=None, + tokens_per_block=32, + kv_cache_manager=None, + is_cross=False, + is_spec_decoding_enabled=False, + ) + forward_args = AttentionForwardArgs( + output=torch.empty((3, 4)), + attention_input_type=AttentionInputType.mixed, + attention_window_size=8, + ) + + combined_fmha.forward(torch.empty((3, 4)), None, None, metadata, forward_args) + + assert events == [ + ("prepare", "context"), + ("prepare", "generation"), + ("run", "context", FmhaPhase.CONTEXT, 2), + ("run", "generation", FmhaPhase.GENERATION, 1), + ] + assert metadata.effective_workspace.numel() == 8 + + +def test_combined_fmha_uses_flattened_v2_page_bound() -> None: + attn = _TestAttention() + combined_fmha = CombinedFmha(attn) + kv_cache_manager = SimpleNamespace( + impl=SimpleNamespace(get_page_index_upper_bound=lambda: 23), + blocks_in_primary_pool=23, + num_local_layers=4, + ) + + assert ( + combined_fmha._get_total_num_blocks(SimpleNamespace(kv_cache_manager=kv_cache_manager)) + == 23 + ) + + +def test_flashinfer_fp8_mode_remains_implementation_local() -> None: + attn = _TestAttention() + attn.quant_mode = int(QuantMode.from_description(use_fp8_kv_cache=True)) + fmha = FlashInferTrtllmGenFmha(attn) + output = torch.empty(1, dtype=torch.bfloat16) + + assert fmha._use_fp8_context_fmha(output, AttentionInputType.context_only) + assert fmha._use_fp8_context_fmha(output, AttentionInputType.mixed) + assert not fmha._use_fp8_context_fmha(output, AttentionInputType.generation_only) diff --git a/tests/unittest/_torch/attention/test_fallback_fmha.py b/tests/unittest/_torch/attention/test_fallback_fmha.py new file mode 100644 index 000000000000..30cb5da2afa9 --- /dev/null +++ b/tests/unittest/_torch/attention/test_fallback_fmha.py @@ -0,0 +1,83 @@ +# 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 types import SimpleNamespace + +import torch + +from tensorrt_llm._torch.attention_backend.fmha.fallback import FallbackFmha +from tensorrt_llm._torch.attention_backend.interface import ( + AttentionForwardArgs, + CustomAttentionMask, +) + + +class _TestAttention: + pass + + +def test_fallback_rejects_q_only_self_attention() -> None: + fmha = FallbackFmha(_TestAttention()) + q = torch.empty((1, 64)) + metadata = SimpleNamespace(is_cross=False) + forward_args = AttentionForwardArgs(is_fused_qkv=False, update_kv_cache=False) + + assert not fmha.is_supported(q, None, None, metadata, forward_args) + + +def test_fallback_accepts_other_qkv_forms() -> None: + fmha = FallbackFmha(_TestAttention()) + q = torch.empty((1, 64)) + + fused_qkv_args = AttentionForwardArgs(is_fused_qkv=True, update_kv_cache=True) + assert fmha.is_supported( + q, + None, + None, + SimpleNamespace(is_cross=False), + fused_qkv_args, + ) + + unfused_qkv_args = AttentionForwardArgs(is_fused_qkv=False, update_kv_cache=True) + kv = torch.empty((1, 32)) + assert fmha.is_supported( + q, + kv, + kv, + SimpleNamespace(is_cross=False), + unfused_qkv_args, + ) + + cached_cross_args = AttentionForwardArgs(is_fused_qkv=False, update_kv_cache=False) + assert fmha.is_supported( + q, + None, + None, + SimpleNamespace(is_cross=True), + cached_cross_args, + ) + + +def test_fallback_still_rejects_custom_mask() -> None: + fmha = FallbackFmha(_TestAttention()) + qkv = torch.empty((1, 64)) + metadata = SimpleNamespace(is_cross=False) + forward_args = AttentionForwardArgs( + attention_mask=CustomAttentionMask.CUSTOM, + is_fused_qkv=True, + update_kv_cache=True, + ) + + assert not fmha.is_supported(qkv, None, None, metadata, forward_args) diff --git a/tests/unittest/_torch/attention/test_trtllm_shared_kv_metadata.py b/tests/unittest/_torch/attention/test_trtllm_shared_kv_metadata.py new file mode 100644 index 000000000000..833db25c8056 --- /dev/null +++ b/tests/unittest/_torch/attention/test_trtllm_shared_kv_metadata.py @@ -0,0 +1,464 @@ +# 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 copy +import math +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch +from kv_cache_utils import fill_kv_cache_logical + +import tensorrt_llm +from tensorrt_llm._torch.attention_backend.fmha.flashinfer_trtllm_gen import FlashInferTrtllmGenFmha +from tensorrt_llm._torch.attention_backend.fmha.interface import FmhaPhase +from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs, AttentionMetadata +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention, TrtllmAttentionMetadata +from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.executor import KvCacheConfig +from tensorrt_llm.mapping import Mapping + + +class _SelectorAttention: + def __init__(self, head_dim: int) -> None: + self.is_mla_enable = False + self.kv_lora_rank = None + self.v_head_dim = None + self.head_dim = head_dim + self.num_heads = 8 + self.num_kv_heads = 1 + self.predicted_tokens_per_seq = 1 + self.sparse_params = None + self.position_embedding_type = 0 + self.quant_mode = 0 + + +def _make_target_metadata() -> TrtllmAttentionMetadata: + metadata = object.__new__(TrtllmAttentionMetadata) + metadata.max_num_requests = 4 + metadata.max_num_sequences = 4 + metadata.max_num_tokens = 16 + metadata._seq_lens = torch.tensor([4, 2, 3], dtype=torch.int32) + metadata._seq_lens_cuda = metadata._seq_lens.clone() + metadata._seq_lens_kv = None + metadata._seq_lens_kv_cuda = None + metadata._num_contexts = 1 + metadata._num_generations = 2 + metadata._num_ctx_tokens = 4 + metadata._num_tokens = 9 + metadata.cross = None + + metadata.kv_cache_manager = SimpleNamespace(kv_factor=2) + metadata.draft_kv_cache_manager = None + metadata.kv_cache_params = object() + metadata.kv_cache_block_offsets = torch.arange(24).view(1, 3, 2, 4) + metadata.host_kv_cache_block_offsets = object() + metadata.draft_kv_cache_block_offsets = None + metadata.workspace = torch.empty(7, dtype=torch.int8) + metadata.cuda_graph_workspace = torch.empty(11, dtype=torch.int8) + metadata.kv_lens_cuda = torch.empty(4, dtype=torch.int32) + metadata.kv_lens_cuda_runtime = torch.tensor([10, 20, 30], dtype=torch.int32) + metadata.kv_lens_runtime = torch.tensor([10, 20, 30], dtype=torch.int32) + metadata.prompt_lens_cuda_runtime = torch.tensor([4, 8, 9], dtype=torch.int32) + metadata.prompt_lens_cpu_runtime = torch.tensor([4, 8, 9], dtype=torch.int32) + metadata.host_request_types = torch.tensor([0, 1, 1, -1], dtype=torch.int32) + metadata.host_request_types_runtime = metadata.host_request_types[:3] + metadata.host_total_kv_lens = torch.tensor([10, 50], dtype=torch.int32) + + metadata.request_ids = [101, 102, 103] + metadata.prompt_lens = [4, 8, 9] + metadata.all_rank_num_tokens = [9] + metadata.mapping = object() + metadata.sparse_metadata_params = None + metadata.num_sparse_topk = 0 + metadata.beam_width = 1 + metadata.padded_num_tokens = 9 + metadata.cuda_graph_buffers = None + metadata._saved_tensors = {"target": torch.tensor(1)} + + metadata.is_spec_decoding_enabled = True + metadata.use_spec_decoding = True + metadata.is_spec_dec_tree = True + metadata.is_spec_dec_dynamic_tree = True + metadata.force_prepare_spec_dec_tree_mask = True + metadata.max_total_draft_tokens = 4 + metadata.spec_decoding_position_offsets = torch.tensor([0]) + metadata.spec_decoding_position_offsets_cpp = torch.tensor([[0]]) + metadata.spec_decoding_packed_mask = torch.tensor([1]) + metadata.spec_decoding_generation_lengths = torch.tensor([1]) + metadata.spec_decoding_bl_tree_mask_offset = torch.tensor([0]) + metadata.spec_decoding_bl_tree_mask = torch.tensor([1]) + metadata.spec_bl_tree_first_sparse_mask_offset_kv = torch.tensor([0]) + metadata.position_offsets_stride = 1 + + metadata._draft_metadata = None + metadata._draft_kv_runtime_lens = None + metadata._shared_kv_draft_seq_lens = None + metadata._shared_kv_draft_seq_lens_cuda = None + metadata._shared_kv_draft_request_types = None + return metadata + + +def test_shared_kv_draft_view_uses_accepted_prefix_without_target_mutation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + "tensorrt_llm._torch.attention_backend.trtllm.prefer_pinned", + lambda: False, + ) + target = _make_target_metadata() + target_kv_lens = target.kv_lens_cuda_runtime.clone() + target_seq_lens = target.seq_lens_cuda.clone() + target_request_types = target.host_request_types.clone() + + draft = target.get_shared_kv_draft_metadata( + torch.tensor([4, 1, 2], dtype=torch.int32), + num_contexts=1, + ) + + torch.testing.assert_close( + draft.kv_lens_cuda_runtime, + torch.tensor([10, 19, 29], dtype=torch.int32), + ) + torch.testing.assert_close(draft.seq_lens, torch.ones(3, dtype=torch.int32)) + torch.testing.assert_close(draft.seq_lens_cuda, torch.ones(3, dtype=torch.int32)) + torch.testing.assert_close(draft.host_request_types_runtime, torch.ones(3, dtype=torch.int32)) + assert draft.num_contexts == 0 + assert draft.num_generations == 3 + assert draft.num_ctx_tokens == 0 + assert draft.num_tokens == 3 + assert not draft.is_cross + + assert draft.kv_cache_manager is target.kv_cache_manager + assert draft.kv_cache_params is target.kv_cache_params + assert draft.kv_cache_block_offsets is target.kv_cache_block_offsets + assert draft.host_kv_cache_block_offsets is target.host_kv_cache_block_offsets + assert draft.workspace is target.workspace + assert draft.cuda_graph_workspace is target.cuda_graph_workspace + assert draft.mapping is target.mapping + assert draft.kv_lens_runtime.data_ptr() == target.kv_lens_runtime.data_ptr() + assert not draft.is_spec_decoding_enabled + assert not draft.use_spec_decoding + assert draft.spec_decoding_packed_mask is None + assert draft.padded_num_tokens is None + + torch.testing.assert_close(target.kv_lens_cuda_runtime, target_kv_lens) + torch.testing.assert_close(target.seq_lens_cuda, target_seq_lens) + torch.testing.assert_close(target.host_request_types, target_request_types) + assert target.num_contexts == 1 + assert target.use_spec_decoding + assert target._saved_tensors + + +def test_shared_kv_draft_view_reuses_stable_buffers(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "tensorrt_llm._torch.attention_backend.trtllm.prefer_pinned", + lambda: False, + ) + target = _make_target_metadata() + draft = target.get_shared_kv_draft_metadata( + torch.tensor([4, 1, 2], dtype=torch.int32), + num_contexts=1, + ) + storage_pointers = ( + draft._draft_kv_runtime_lens.data_ptr(), + draft._shared_kv_draft_seq_lens.data_ptr(), + draft._shared_kv_draft_seq_lens_cuda.data_ptr(), + draft._shared_kv_draft_request_types.data_ptr(), + ) + + target.kv_lens_cuda_runtime.copy_(torch.tensor([11, 22, 33])) + target.kv_lens_runtime.copy_(torch.tensor([11, 22, 33])) + second = target.get_shared_kv_draft_metadata( + torch.tensor([4, 2, 1], dtype=torch.int32), + num_contexts=1, + ) + + assert second is draft + assert storage_pointers == ( + draft._draft_kv_runtime_lens.data_ptr(), + draft._shared_kv_draft_seq_lens.data_ptr(), + draft._shared_kv_draft_seq_lens_cuda.data_ptr(), + draft._shared_kv_draft_request_types.data_ptr(), + ) + torch.testing.assert_close( + draft.kv_lens_cuda_runtime, + torch.tensor([11, 22, 31], dtype=torch.int32), + ) + torch.testing.assert_close( + target.kv_lens_cuda_runtime, + torch.tensor([11, 22, 33], dtype=torch.int32), + ) + + +@pytest.mark.parametrize( + ("field", "value", "error"), + [ + ("beam_width", 2, "beam search"), + ("draft_kv_cache_manager", object(), "separate draft KV cache"), + ("sparse_metadata_params", object(), "sparse attention"), + ], +) +def test_shared_kv_draft_view_rejects_unsupported_modes( + field: str, + value: object, + error: str, +) -> None: + target = _make_target_metadata() + setattr(target, field, value) + + with pytest.raises(ValueError, match=error): + target.get_shared_kv_draft_metadata(torch.ones(3, dtype=torch.int32), num_contexts=1) + + +def test_shared_kv_draft_view_rejects_mla_and_cross_attention() -> None: + target = _make_target_metadata() + target.kv_cache_manager.kv_factor = 1 + with pytest.raises(ValueError, match="MLA"): + target.get_shared_kv_draft_metadata(torch.ones(3, dtype=torch.int32), num_contexts=1) + + target = _make_target_metadata() + target._seq_lens_kv = target._seq_lens.clone() + target._seq_lens_kv_cuda = target._seq_lens_cuda.clone() + with pytest.raises(ValueError, match="cross attention"): + target.get_shared_kv_draft_metadata(torch.ones(3, dtype=torch.int32), num_contexts=1) + + +def test_cuda_graph_metadata_does_not_inherit_eager_draft_view() -> None: + target = _make_target_metadata() + target._draft_metadata = object() + graph_metadata = copy.copy(target) + + with mock.patch.object( + AttentionMetadata, + "create_cuda_graph_metadata", + return_value=graph_metadata, + ): + result = target.create_cuda_graph_metadata(max_batch_size=4) + + assert result is graph_metadata + assert result._draft_metadata is None + assert target._draft_metadata is not None + + +def test_cuda_graph_metadata_preserves_its_draft_view() -> None: + graph_metadata = _make_target_metadata() + graph_metadata.is_cuda_graph = True + draft_metadata = object() + graph_metadata._draft_metadata = draft_metadata + + result = graph_metadata.create_cuda_graph_metadata(max_batch_size=4) + + assert result is graph_metadata + assert result._draft_metadata is draft_metadata + + +@pytest.mark.parametrize("head_dim", [256, 512]) +def test_q_only_generation_selects_flashinfer_trtllm_gen(head_dim: int) -> None: + attention = _SelectorAttention(head_dim) + fmha = FlashInferTrtllmGenFmha(attention) + metadata = SimpleNamespace( + num_contexts=0, + num_generations=2, + helix_position_offsets=None, + num_sparse_topk=0, + use_spec_decoding=False, + is_spec_dec_tree=False, + kv_cache_block_offsets=object(), + kv_cache_manager=SimpleNamespace( + dtype=DataType.BF16, + blocks_in_primary_pool=8, + num_local_layers=1, + impl=None, + ), + is_cross=False, + tokens_per_block=32, + beam_width=1, + ) + q = torch.empty((2, attention.num_heads * head_dim), dtype=torch.bfloat16) + forward_args = AttentionForwardArgs( + output=torch.empty_like(q), + is_fused_qkv=False, + update_kv_cache=False, + ) + + assert fmha.is_supported( + q, + None, + None, + metadata, + forward_args, + phase=FmhaPhase.GENERATION, + ) + backend = SimpleNamespace(fmha_libs=[fmha], combined_fmha=None) + selected = TrtllmAttention._select_non_mla_fmha( + backend, + q, + None, + None, + metadata, + forward_args, + ) + + assert selected is fmha + + +def _shared_kv_reference( + q: torch.Tensor, + keys: list[torch.Tensor], + values: list[torch.Tensor], + kv_lens: list[int], +) -> torch.Tensor: + num_heads = q.shape[1] // keys[0].shape[-1] + num_kv_heads = keys[0].shape[1] + repeats = num_heads // num_kv_heads + outputs = [] + for row, kv_len in enumerate(kv_lens): + q_row = q[row].view(num_heads, -1).float() + k_row = keys[row][:kv_len].repeat_interleave(repeats, dim=1).float() + v_row = values[row][:kv_len].repeat_interleave(repeats, dim=1).float() + scores = torch.einsum("hd,lhd->hl", q_row, k_row) / math.sqrt(q_row.shape[-1]) + probs = torch.softmax(scores, dim=-1) + outputs.append(torch.einsum("hl,lhd->hd", probs, v_row)) + return torch.stack(outputs).to(q.dtype).view_as(q) + + +@pytest.mark.parametrize("head_dim", [256, 512]) +def test_q_only_shared_kv_forward_preserves_target_cache(head_dim: int) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the shared-KV forward test") + if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): + pytest.skip("FlashInfer trtllm-gen requires SM100 or SM103") + if not IS_FLASHINFER_AVAILABLE: + pytest.skip("FlashInfer is required for the trtllm-gen forward test") + + num_heads = 8 + num_kv_heads = 1 + target_q_lens = [3, 3] + cached_lens = [5, 7] + target_kv_lens = [ + cached + query for cached, query in zip(cached_lens, target_q_lens, strict=True) + ] + accepted_lens = [2, 1] + draft_kv_lens = [ + target - query + accepted + for target, query, accepted in zip( + target_kv_lens, + target_q_lens, + accepted_lens, + strict=True, + ) + ] + request_ids = [0, 1] + mapping = Mapping(world_size=1, tp_size=1, rank=0) + manager = KVCacheManager( + KvCacheConfig(max_tokens=64), + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, + num_layers=1, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=32, + max_seq_len=32, + max_batch_size=len(request_ids), + mapping=mapping, + dtype=DataType.BF16, + ) + try: + manager.add_dummy_requests(request_ids, target_kv_lens) + generator = torch.Generator(device="cuda").manual_seed(2026 + head_dim) + keys = [ + torch.randn( + length, + num_kv_heads, + head_dim, + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + for length in target_kv_lens + ] + values = [ + torch.randn( + length, + num_kv_heads, + head_dim, + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + for length in target_kv_lens + ] + fill_kv_cache_logical( + manager, + 0, + request_ids, + keys, + values, + kv_layout="HND", + ) + + target_metadata = TrtllmAttentionMetadata( + num_contexts=0, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=cached_lens, + ), + seq_lens=torch.tensor(target_q_lens, dtype=torch.int32), + max_num_requests=len(request_ids), + max_num_tokens=16, + kv_cache_manager=manager, + request_ids=request_ids, + prompt_lens=target_kv_lens, + kv_layout="HND", + ) + target_metadata.prepare() + draft_metadata = target_metadata.get_shared_kv_draft_metadata( + torch.tensor(accepted_lens, dtype=torch.int32, device="cuda"), + num_contexts=0, + ) + + attention = TrtllmAttention( + layer_idx=0, + num_heads=num_heads, + head_dim=head_dim, + num_kv_heads=num_kv_heads, + ) + assert FlashInferTrtllmGenFmha.is_available(attention) + fmha = FlashInferTrtllmGenFmha(attention) + attention.fmha_libs = [fmha] + attention.phased_fmha_libs = [fmha] + attention.non_phased_fmha_libs = [] + + q = torch.randn( + len(request_ids), + num_heads * head_dim, + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + expected = _shared_kv_reference(q, keys, values, draft_kv_lens) + cache = manager.get_buffers(0, kv_layout="HND") + cache_before = cache.clone() + + actual = attention.forward(q, None, None, draft_metadata) + + torch.testing.assert_close(actual, expected, atol=3e-2, rtol=3e-3) + torch.testing.assert_close(cache, cache_before, atol=0, rtol=0) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py index 351e297627df..7e56c3b46462 100644 --- a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py +++ b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py @@ -25,7 +25,7 @@ Tier 2 (LLM_MODELS_ROOT gated, real tokenizer/processor): - ``TestGemma4InputProcessor`` — image/audio/video pre-processing. - - ``TestGemma4MultimodalE2E`` — full LLM+vision generate parity via the + - ``TestGemma4MultimodalE2E`` — full LLM+vision context/decode execution via the shared ``TestModelingMultimodal`` skeleton (image / multi-image modalities). Audio-side tests (``TestGemma4InputProcessor.test_audio_*``) are intentionally @@ -812,7 +812,7 @@ def test_instantiation_with_vision(self): mc = ModelConfig( pretrained_config=config, mapping=Mapping(world_size=1, tp_size=1, rank=0), - attn_backend="FLASHINFER", + attn_backend="TRTLLM", ) model = Gemma4ForConditionalGeneration(mc) @@ -825,6 +825,32 @@ def test_instantiation_with_vision(self): # TRT-LLM class, not ``transformers.AutoModel`` output. self.assertIsInstance(model.vision_tower, Gemma4VisionModel) + def test_text_submodel_preserves_explicit_attention_backend(self): + """The text submodel honors explicit backend overrides.""" + text_config = Gemma4TextConfig(**SMALL_TEXT_CONFIG) + vision_config = Gemma4VisionConfig(**SMALL_VISION_CONFIG) + config = Gemma4Config( + text_config=text_config, + vision_config=vision_config, + audio_config=None, + ) + + for backend in ("FLASHINFER", "TRTLLM"): + model_config = ModelConfig( + pretrained_config=config, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + attn_backend=backend, + ) + text_model_config = Gemma4ForConditionalGeneration.get_sub_model_config( + model_config, "text_config" + ) + vision_model_config = Gemma4ForConditionalGeneration.get_sub_model_config( + model_config, "vision_config" + ) + + self.assertEqual(text_model_config.attn_backend, backend) + self.assertEqual(vision_model_config.attn_backend, "TRTLLM") + def test_encoder_cache_reuses_image_embedding_across_requests(self): """Persistent cache reuse applies to the shared dense/MoE Gemma4 wrapper.""" self.assertTrue(issubclass(Gemma4ForConditionalGeneration, MultimodalModelMixin)) @@ -1089,7 +1115,7 @@ def test_instantiation_without_vision(self): mc = ModelConfig( pretrained_config=config, mapping=Mapping(world_size=1, tp_size=1, rank=0), - attn_backend="FLASHINFER", + attn_backend="TRTLLM", ) model = Gemma4ForConditionalGeneration(mc) @@ -1105,9 +1131,18 @@ def test_instantiation_without_vision(self): def _get_model_path(): llm_models_root = os.environ.get("LLM_MODELS_ROOT") - if llm_models_root: - return os.path.join(llm_models_root, "gemma4/gemma-4-26B-A4B-it") - return None + if not llm_models_root: + return None + + model_subdirs = ( + "gemma/gemma-4-26B-A4B-it", + "gemma4/gemma-4-26B-A4B-it", + ) + for model_subdir in model_subdirs: + model_path = os.path.join(llm_models_root, model_subdir) + if os.path.isfile(os.path.join(model_path, "config.json")): + return model_path + return os.path.join(llm_models_root, model_subdirs[0]) MODEL_26B_PATH = _get_model_path() @@ -1489,11 +1524,7 @@ def _video_token_id(self): # follow-up work. -_GEMMA4_E2E_PATH = ( - os.path.join(os.environ["LLM_MODELS_ROOT"], "gemma4/gemma-4-26B-A4B-it") - if os.environ.get("LLM_MODELS_ROOT") - else None -) +_GEMMA4_E2E_PATH = _get_model_path() def _gemma4_e2e_model_available() -> bool: @@ -1517,7 +1548,7 @@ def _gemma4_e2e_model_available() -> bool: # Reduced-layer text+vision config for E2E. Real ``_name_or_path`` so the # input processor + tokenizer pick up the genuine Gemma4 processor files; # attention head_dim follows the existing dummy-config recipe in -# ``test_gemma4_e2e_dummy.py`` (128 / 256 for FlashInfer compatibility). +# ``test_gemma4_e2e_dummy.py`` (128 / 256 for TRTLLM FMHA compatibility). GEMMA4_E2E_CONFIG = { "architectures": ["Gemma4ForConditionalGeneration"], "model_type": "gemma4", @@ -1526,7 +1557,8 @@ def _gemma4_e2e_model_available() -> bool: # HF Gemma4 (transformers 5.5.3) does not advertise SDPA support; pin # to eager so PreTrainedModel.__init__ does not raise on dispatch check. "_attn_implementation": "eager", - "image_token_id": 262145, + # Match the real processor output used by this E2E. + "image_token_id": 258880, "audio_token_id": 262273, "video_token_id": 262146, "boi_token_index": 255999, @@ -1555,6 +1587,8 @@ def _gemma4_e2e_model_available() -> bool: "hidden_size_per_layer_input": 0, "use_double_wide_mlp": False, "final_logit_softcapping": None, + # Exercise the custom mask for image-token blocks. + "use_bidirectional_attention": "vision", "torch_dtype": "bfloat16", "tie_word_embeddings": True, "attention_bias": False, @@ -1588,15 +1622,17 @@ def _gemma4_e2e_model_available() -> bool: not _gemma4_e2e_model_available(), reason=( "LLM_MODELS_ROOT not set or Gemma4 model not available — " - "E2E generate parity requires real tokenizer/processor files." + "E2E generation requires real tokenizer/processor files." ), ) class TestGemma4MultimodalE2E(unittest.TestCase): - """Full LLM + vision parity via the shared multimodal skeleton. + """Full LLM + vision execution via the shared multimodal skeleton. Implemented as a thin wrapper around ``TestModelingMultimodal`` so we - reuse the standardized context+generation phase, KV cache manager - setup, HF input loading, and tolerance/compare helpers. + reuse the standardized context+generation phases and KV cache setup. + Component tests above cover HF numerical parity; this test targets the + fused TRTLLM runtime path because full-vocabulary random-weight BF16 + parity is numerically unstable across the fused and eager implementations. """ # The actual abstract base + skeleton live in @@ -1630,6 +1666,29 @@ def get_model_type(self) -> str: def get_model_config_class(self) -> Type: return Gemma4Config + @property + def skip_hf_inference(self) -> bool: + return True + + def get_kv_cache_manager( + self, + dtype, + config, + tokens_per_block, + max_seq_len, + batch_size, + num_blocks, + ): + del dtype, max_seq_len + from test_modeling_gemma4 import _build_gemma4_kv_cache_manager + + return _build_gemma4_kv_cache_manager( + config.text_config, + num_blocks=num_blocks, + tokens_per_block=tokens_per_block, + batch_size=batch_size, + ) + def get_scenarios(self) -> List[MultimodalScenario]: # Single sanity-image scenario for the in-PR test list. # CUDA-graph / chunked-prefill / kv-cache-reuse permutations diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma3.py b/tests/unittest/_torch/modeling/test_modeling_gemma3.py index f252f4cbd2d6..af120f1db3fb 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma3.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma3.py @@ -403,7 +403,7 @@ def test_gemma3_allclose_to_hf(self, scenario: Scenario) -> None: kv_cache_manager.shutdown() - def test_gemma3_flashinfer_mask(self): + def test_gemma3_attention_mask(self): config_dict = deepcopy(GEMMA3_1B_CONFIG) gemma3_config = Gemma3TextConfig.from_dict(config_dict) @@ -461,7 +461,7 @@ def test_gemma3_flashinfer_mask(self): # First sample has 2 image tokens, second sample has 2 image tokens, third sample has none. image_token_mask = torch.tensor( [True, True, False, True, True, True, False, False], device=device) - causal_mask = gemma3.get_flashinfer_attention_mask( + causal_mask = gemma3.get_attention_mask( image_token_mask=image_token_mask, attn_metadata=attn_metadata) # Causal mask for context request 1. ctx_request_1_mask = torch.tensor( diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index 5eff2c644555..1984042fc251 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -617,11 +617,56 @@ def test_ordered_embedding_combines_vocab_parallel_shards(self): torch.testing.assert_close(actual, expected) - def test_assistant_uses_target_kv_sources(self): + @torch.no_grad() + def test_kv_shared_sliding_rope_uses_explicit_position_ids(self): assistant = Gemma4AssistantForCausalLM(_make_assistant_model_config()) - self.assertEqual(len(assistant.model.layers), 4) - self.assertTrue(all(layer.is_kv_shared_layer for layer in assistant.model.layers)) + attention = next(layer.self_attn for layer in assistant.model.layers if layer.is_sliding) + self.assertTrue(attention.is_kv_shared) + self.assertFalse(attention.rope_fusion) + self.assertIsNotNone(attention.rotary_emb) + + rotary_emb = attention.rotary_emb + device = rotary_emb.rotary_cos_sin.device + attention.q_norm.to(device) + query = torch.arange( + 1, + attention.q_size + 1, + dtype=attention.q_norm.weight.dtype, + device=device, + ).unsqueeze(0) + prefix_length = 7 + + def reference_rope(position: int) -> torch.Tensor: + normalized = attention.q_norm(query.reshape(-1, attention.head_dim)).reshape( + -1, attention.q_size + ) + cos_sin = rotary_emb.rotary_cos_sin[position : position + 1] + cos = cos_sin[:, 0, :].to(normalized.dtype).unsqueeze(0) + sin = cos_sin[:, 1, :].to(normalized.dtype).unsqueeze(0) + normalized = normalized.view(1, 1, -1, attention.head_dim).transpose(1, 2) + rotated = rotary_emb.apply_rotary_pos_emb( + normalized, + cos, + sin, + is_neox=rotary_emb.is_neox, + ) + return rotated.transpose(1, 2).contiguous().view(1, -1) + + actual, k, v = attention.apply_rope( + query, + None, + None, + torch.tensor([[prefix_length]], device=device), + ) + expected = reference_rope(prefix_length) + previous_position = reference_rope(prefix_length - 1) + + self.assertIsNone(k) + self.assertIsNone(v) + torch.testing.assert_close(actual, expected) + self.assertFalse(torch.allclose(actual, previous_position, rtol=1e-3, atol=1e-3)) + def test_assistant_uses_target_kv_sources(self): target_config = { **GEMMA4_SMALL_CONFIG, "layer_types": [ @@ -635,6 +680,18 @@ def test_assistant_uses_target_kv_sources(self): "num_kv_shared_layers": 2, } target = Gemma4ForCausalLM(_make_model_config(target_config)) + target_attn_layers = target.model.model_config.extra_attrs["attn_layers"].copy() + + assistant = Gemma4AssistantForCausalLM(_make_assistant_model_config()) + self.assertEqual(len(assistant.model.layers), 4) + self.assertTrue(all(layer.is_kv_shared_layer for layer in assistant.model.layers)) + self.assertTrue(all(not layer.self_attn.rope_fusion for layer in assistant.model.layers)) + self.assertTrue( + all(not layer.self_attn.register_to_config for layer in assistant.model.layers) + ) + self.assertEqual(assistant.model.model_config.extra_attrs["attn_layers"], {}) + self.assertEqual(target.model.model_config.extra_attrs["attn_layers"], target_attn_layers) + assistant.load_weights_from_target_model(target) expected_source_layers = [2, 2, 2, 3] @@ -972,7 +1029,7 @@ def _assert_most_elems_close(self, actual, ref, atol=0.5, rtol=0.5, max_failed_f f"{failed * 100:.2f}% of elements differ (max {max_failed_frac * 100}%)", ) - def _make_hf_and_trt_models(self, config_dict=None): + def _make_hf_and_trt_models(self, config_dict=None, attn_backend="FLASHINFER"): """Create paired HF and TRT-LLM models with shared weights.""" from transformers import Gemma4ForCausalLM as HFGemma4 @@ -986,7 +1043,7 @@ def _make_hf_and_trt_models(self, config_dict=None): device = torch.device("cuda") hf_model = HFGemma4(config).to(dtype).to(device).eval() - model_config = ModelConfig(pretrained_config=config, attn_backend="FLASHINFER") + model_config = ModelConfig(pretrained_config=config, attn_backend=attn_backend) trt_model = Gemma4ForCausalLM(model_config).to(dtype).to(device) wm = Gemma4HfWeightMapper() @@ -1095,6 +1152,7 @@ def _run_full_model_comparison( atol=0.5, rtol=0.5, max_failed_frac=0.01, + attn_backend="FLASHINFER", ): """Run context + generation comparison for a given config.""" from transformers.cache_utils import DynamicCache @@ -1103,13 +1161,15 @@ def _run_full_model_comparison( from tensorrt_llm._torch.metadata import KVCacheParams torch.random.manual_seed(42) - hf, trt, gemma4_config = self._make_hf_and_trt_models(config_dict) + hf, trt, gemma4_config = self._make_hf_and_trt_models( + config_dict, attn_backend=attn_backend + ) if gemma4_config.enable_moe_block: self._stabilize_moe_routing(hf, trt, gemma4_config) hf_cache = DynamicCache() device = torch.device("cuda") - backend = "FLASHINFER" + backend = attn_backend # Set up KV cache kv_cache_manager = self._get_kv_cache_manager(gemma4_config) @@ -1210,6 +1270,25 @@ def test_hybrid_headdim_config(self): """Hybrid head_dim: sliding=64, full=128 (per-layer KV cache).""" self._run_full_model_comparison(deepcopy(GEMMA4_HYBRID_HEADDIM_CONFIG)) + @torch.no_grad() + def test_hybrid_headdim_config_trtllm(self): + """TRTLLM backend parity for hybrid sliding/full head dimensions.""" + self._run_full_model_comparison( + deepcopy(GEMMA4_HYBRID_HEADDIM_CONFIG), + attn_backend="TRTLLM", + ) + + @torch.no_grad() + def test_kv_sharing_real_headdim_config_trtllm(self): + """TRTLLM backend reuses cached KV for Q-only H256/H512 layers.""" + config_dict = deepcopy(GEMMA4_E2B_REAL_DIMS_CONFIG) + config_dict["num_hidden_layers"] = 12 + config_dict["num_kv_shared_layers"] = 4 + self._run_full_model_comparison( + config_dict, + attn_backend="TRTLLM", + ) + @torch.no_grad() def test_diff_kv_heads_config(self): """Different num_kv_heads per layer type — tests V2 pool grouping.""" @@ -1831,6 +1910,10 @@ def test_kv_sharing_instantiation(self): if i >= first_shared: self.assertTrue(layer.is_kv_shared_layer, f"Layer {i} should be KV shared") self.assertTrue(layer.self_attn.is_kv_shared, f"Layer {i} attn should be KV shared") + self.assertFalse( + layer.self_attn.rope_fusion, + f"Layer {i} should apply RoPE from explicit position_ids", + ) else: self.assertFalse(layer.is_kv_shared_layer, f"Layer {i} should NOT be KV shared") @@ -2356,20 +2439,12 @@ def test_gemma4_routing_matches_hf_reference(self): class TestGemma4ModelDefaults(unittest.TestCase): """Tests for Gemma4 model defaults (get_model_defaults).""" - def test_causal_lm_requires_flashinfer_backend(self): - """Gemma4ForCausalLM must default to FLASHINFER attention backend. - - The TRTLLM backend does not support: - - FlashInfer VSWA per-pool page management for hybrid head_dim - - trtllm-gen cubin dispatch for head_dim=512 layers - - Custom attention masks for bidirectional multimodal tokens - Without this default, models crash with: - 'TrtllmAttentionMetadata' object has no attribute 'kv_layout' - """ + def test_causal_lm_requires_trtllm_backend(self): + """Gemma4ForCausalLM must default to the TRTLLM attention backend.""" defaults = Gemma4ForCausalLM.get_model_defaults(None) self.assertIn("attn_backend", defaults, "get_model_defaults must set attn_backend") self.assertEqual( - defaults["attn_backend"], "FLASHINFER", "Gemma4 requires FLASHINFER (exact uppercase)" + defaults["attn_backend"], "TRTLLM", "Gemma4 requires TRTLLM (exact uppercase)" ) def test_causal_lm_does_not_disable_cuda_graphs(self): @@ -2377,13 +2452,13 @@ def test_causal_lm_does_not_disable_cuda_graphs(self): defaults = Gemma4ForCausalLM.get_model_defaults(None) self.assertNotIn("cuda_graph_config", defaults) - def test_conditional_gen_requires_flashinfer_backend(self): - """Gemma4ForConditionalGeneration must also default to FLASHINFER.""" + def test_conditional_gen_requires_trtllm_backend(self): + """Gemma4ForConditionalGeneration must also default to TRTLLM.""" from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4ForConditionalGeneration defaults = Gemma4ForConditionalGeneration.get_model_defaults(None) self.assertIn("attn_backend", defaults) - self.assertEqual(defaults["attn_backend"], "FLASHINFER") + self.assertEqual(defaults["attn_backend"], "TRTLLM") def test_conditional_gen_does_not_disable_cuda_graphs(self): """Gemma4ForConditionalGeneration must not disable CUDA graphs.""" @@ -2392,28 +2467,55 @@ def test_conditional_gen_does_not_disable_cuda_graphs(self): defaults = Gemma4ForConditionalGeneration.get_model_defaults(None) self.assertNotIn("cuda_graph_config", defaults) - def test_attn_backend_dispatches_to_flashinfer(self): - """Verify the exact string 'FLASHINFER' dispatches correctly. + def test_external_shared_kv_mtp_defaults_to_flashinfer(self): + """Keep the production FP8 custom-mask MTP path on FlashInfer.""" + from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4ForConditionalGeneration + + spec_dec_mode = SimpleNamespace(is_mtp_eagle_one_model=lambda: True) + llm_args = SimpleNamespace(speculative_config=SimpleNamespace(spec_dec_mode=spec_dec_mode)) - get_attention_backend uses exact case-sensitive match. 'FlashInfer' - or 'flashinfer' would silently fall back to TrtllmAttention, which - causes 'TrtllmAttentionMetadata has no attribute kv_layout' crashes. - """ + self.assertEqual( + Gemma4ForCausalLM.get_model_defaults(llm_args)["attn_backend"], "FLASHINFER" + ) + self.assertEqual( + Gemma4ForConditionalGeneration.get_model_defaults(llm_args)["attn_backend"], + "FLASHINFER", + ) + + def test_attn_backend_dispatches_to_trtllm(self): + """Verify the Gemma4 default dispatches to TrtllmAttention.""" from tensorrt_llm._torch.attention_backend.utils import get_attention_backend defaults = Gemma4ForCausalLM.get_model_defaults(None) backend_cls = get_attention_backend(defaults["attn_backend"]) - # Must be FlashInferAttention, not TrtllmAttention self.assertEqual( backend_cls.__name__, - "FlashInferAttention", - "FLASHINFER must dispatch to FlashInferAttention", + "TrtllmAttention", + "TRTLLM must dispatch to TrtllmAttention", ) + def test_all_layers_use_trtllm_backend(self): + """All Gemma4 layers use the default TRTLLM attention backend.""" + config_dict = deepcopy(GEMMA4_SMALL_CONFIG) + config = Gemma4TextConfig(**config_dict) + model_config = ModelConfig(pretrained_config=config) + + for i in range(config.num_hidden_layers): + attn = Gemma4Attention( + model_config, + i, + is_sliding=config.layer_types[i] == "sliding_attention", + ) + self.assertEqual( + attn.attn.__class__.__name__, + "TrtllmAttention", + f"Layer {i} should use the TRTLLM backend", + ) + @unittest.mock.patch("tensorrt_llm._torch.models.modeling_gemma4.is_sm_100f", return_value=True) - def test_all_layers_use_trtllm_gen_on_sm100f(self, _mock_is_sm_100f): - """All Gemma4 layers use trtllm-gen uniformly on datacenter Blackwell. + def test_all_flashinfer_layers_use_trtllm_gen_on_sm100f(self, _mock_is_sm_100f): + """All explicit FlashInfer layers use trtllm-gen on datacenter Blackwell. trtllm-gen has pre-compiled cubins for H256+H512, both BF16 and FP8 dtypes. For FP8 KV cache (NVFP4), the FlashInfer backend @@ -2423,7 +2525,7 @@ def test_all_layers_use_trtllm_gen_on_sm100f(self, _mock_is_sm_100f): """ config_dict = deepcopy(GEMMA4_SMALL_CONFIG) config = Gemma4TextConfig(**config_dict) - model_config = ModelConfig(pretrained_config=config) + model_config = ModelConfig(pretrained_config=config, attn_backend="FLASHINFER") for i in range(config.num_hidden_layers): attn = Gemma4Attention( @@ -2440,11 +2542,11 @@ def test_all_layers_use_trtllm_gen_on_sm100f(self, _mock_is_sm_100f): @unittest.mock.patch( "tensorrt_llm._torch.models.modeling_gemma4.is_sm_100f", return_value=False ) - def test_non_sm100f_layers_use_fa2(self, _mock_is_sm_100f): - """Gemma4 uses FlashInfer FA2 where trtllm-gen kernels are unavailable.""" + def test_non_sm100f_flashinfer_layers_use_fa2(self, _mock_is_sm_100f): + """Explicit FlashInfer layers use FA2 where trtllm-gen is unavailable.""" config_dict = deepcopy(GEMMA4_SMALL_CONFIG) config = Gemma4TextConfig(**config_dict) - model_config = ModelConfig(pretrained_config=config) + model_config = ModelConfig(pretrained_config=config, attn_backend="FLASHINFER") for layer_idx in range(config.num_hidden_layers): attn = Gemma4Attention( diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_external_shared_kv_metadata.py b/tests/unittest/_torch/speculative/hw_agnostic/test_external_shared_kv_metadata.py new file mode 100644 index 000000000000..49913d98bcae --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_external_shared_kv_metadata.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import torch + +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker + + +def test_external_shared_kv_uses_backend_metadata_contract() -> None: + """The worker forwards the backend-provided draft metadata view.""" + num_accepted_tokens = torch.tensor([1, 2], dtype=torch.int32) + draft_metadata = SimpleNamespace( + use_spec_decoding=True, + padded_num_tokens=4, + all_rank_num_tokens=None, + ) + + class FakeAttentionMetadata: + def __init__(self) -> None: + self.seq_lens_cuda = torch.tensor([1, 1], dtype=torch.int32) + self.calls = [] + + def get_shared_kv_draft_metadata(self, accepted, num_contexts): + self.calls.append((accepted, num_contexts)) + return draft_metadata + + class FakeDraftModel: + def __init__(self) -> None: + self.received_metadata = [] + + def forward_draft_step(self, **kwargs): + self.received_metadata.append(kwargs["attn_metadata"]) + return torch.zeros((2, 1)), kwargs["recurrent_hidden_states"] + + attn_metadata = FakeAttentionMetadata() + draft_model = FakeDraftModel() + all_rank_num_tokens = torch.tensor([2], dtype=torch.int32) + spec_metadata = SimpleNamespace( + runtime_draft_len=1, + batch_indices_cuda=torch.arange(2), + subseq_all_rank_num_tokens=all_rank_num_tokens, + ) + worker = SimpleNamespace( + guided_decoder=None, + sa_enhancer=None, + _prepare_shared_kv_draft_inputs=lambda **kwargs: ( + torch.tensor([10, 20]), + torch.zeros((2, 4)), + torch.tensor([[1, 2]]), + ), + sample_draft_tokens=lambda *args, **kwargs: torch.tensor([11, 21]), + ) + + result = Eagle3OneModelWorker._forward_external_shared_target_kv_draft_loop( + worker, + position_ids=torch.tensor([[0, 1]]), + hidden_states=torch.zeros((2, 4)), + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + accepted_tokens=torch.tensor([[10], [20]]), + num_accepted_tokens=num_accepted_tokens, + num_contexts=1, + batch_size=2, + ) + + assert len(attn_metadata.calls) == 1 + assert attn_metadata.calls[0][0] is num_accepted_tokens + assert attn_metadata.calls[0][1] == 1 + assert draft_model.received_metadata == [draft_metadata] + assert draft_metadata.use_spec_decoding is False + assert draft_metadata.padded_num_tokens is None + assert draft_metadata.all_rank_num_tokens is all_rank_num_tokens + torch.testing.assert_close(result, torch.tensor([[11], [21]]))