[https://nvbugs/6566891][fix] Use FlashInfer FA2 for Gemma4 on SM120 and SM121

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.

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(