diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index 24aa619b8e34..a9aa54d6029c 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -29,6 +29,7 @@ from tensorrt_llm._torch.modules.fused_moe.interface import MoEWeightLoadingMode from tensorrt_llm._torch.modules.fused_moe.routing import BaseMoeRoutingMethod from tensorrt_llm._torch.modules.qk_norm_attention import QKNormRoPEAttention +from tensorrt_llm._utils import is_sm_100f from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -328,14 +329,10 @@ def __init__( # the rotate split, matching HF's rotate_half(head_dim//2) pairing. self.rotary_emb.head_dim = layer_head_dim - # Use trtllm-gen for ALL layers. trtllm-gen has pre-compiled cubins - # for both H256+SWA and H512 across all supported dtypes. - # For FP8 KV cache (NVFP4), Q is also cast to FP8 in the FlashInfer - # backend so that QkvE4m3OBfloat16 context cubins can be used - # (context cubins require same Q/KV dtype; decode cubins support - # mixed dtypes natively). Uniform backend avoids workspace - # corruption between different wrapper types under CUDA graphs. - self.attn.flashinfer_backend = "trtllm-gen" + # 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" # KV shared layers: use target layer's index for KV cache access # so the attention backend reads from the target layer's cache slot. diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index 5913dd3a518d..5eff2c644555 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -45,6 +45,7 @@ Gemma4TextModel, Gemma4TextScaledWordEmbedding, ) +from tensorrt_llm._utils import is_sm_100f from tensorrt_llm.mapping import Mapping if TYPE_CHECKING: @@ -2410,8 +2411,9 @@ def test_attn_backend_dispatches_to_flashinfer(self): "FLASHINFER must dispatch to FlashInferAttention", ) - def test_all_layers_use_trtllm_gen(self): - """All Gemma4 layers use trtllm-gen backend uniformly. + @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. trtllm-gen has pre-compiled cubins for H256+H512, both BF16 and FP8 dtypes. For FP8 KV cache (NVFP4), the FlashInfer backend @@ -2424,13 +2426,34 @@ def test_all_layers_use_trtllm_gen(self): model_config = ModelConfig(pretrained_config=config) for i in range(config.num_hidden_layers): - attn = Gemma4Attention(model_config, i) + attn = Gemma4Attention( + model_config, + i, + is_sliding=config.layer_types[i] == "sliding_attention", + ) self.assertEqual( attn.attn.flashinfer_backend, "trtllm-gen", f"Layer {i} should use trtllm-gen", ) + @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.""" + config_dict = deepcopy(GEMMA4_SMALL_CONFIG) + config = Gemma4TextConfig(**config_dict) + model_config = ModelConfig(pretrained_config=config) + + for layer_idx in range(config.num_hidden_layers): + attn = Gemma4Attention( + model_config, + layer_idx=layer_idx, + is_sliding=config.layer_types[layer_idx] == "sliding_attention", + ) + self.assertEqual(attn.attn.flashinfer_backend, "fa2") + class TestGemma4CUDAGraph(unittest.TestCase): """Tests for Gemma4 attention with CUDA graph capture/replay.""" @@ -3388,8 +3411,7 @@ def _run_cuda_graph_real_headdim(self, config_dict, label=""): layers = [] for info in layers_info: kwargs = {} - # head_dim>256 needs trtllm-gen (fa2 JIT doesn't support it) - if info["head_dim"] > 256: + if info["head_dim"] > 256 and is_sm_100f(): kwargs["flashinfer_backend"] = "trtllm-gen" layers.append( FlashInferAttention(