From 31eb8010334d335a8a8fed00ce98c9767ccc36f4 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Wed, 12 Aug 2026 06:27:37 -0700 Subject: [PATCH 1/2] [NVBUG-6566891][fix] Use FA2 for Gemma4 on SM120 and SM121 Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_gemma4.py | 14 +++---- .../_torch/modeling/test_modeling_gemma4.py | 37 +++++++++++++++++-- 2 files changed, 39 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index 24aa619b8e34..b94559a593dc 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 get_sm_version from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -328,14 +329,11 @@ 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 does not provide SM120/SM121 kernels. FlashInfer FA2 + # supports Gemma4 H256/H512, sliding-window attention, and FP8 KV cache + # on these architectures. Keep the existing trtllm-gen routing on all + # other architectures. + self.attn.flashinfer_backend = "fa2" if get_sm_version() in (120, 121) else "trtllm-gen" # 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..8db5bd6ea7d1 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -2410,8 +2410,11 @@ 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.get_sm_version", return_value=100 + ) + def test_all_layers_use_trtllm_gen_on_sm100(self, _mock_get_sm_version): + """All Gemma4 layers use trtllm-gen uniformly on SM100. trtllm-gen has pre-compiled cubins for H256+H512, both BF16 and FP8 dtypes. For FP8 KV cache (NVFP4), the FlashInfer backend @@ -2424,13 +2427,39 @@ 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", ) + def test_sm12x_layers_use_fa2(self): + """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 sm_version in (120, 121): + with ( + self.subTest(sm_version=sm_version), + unittest.mock.patch( + "tensorrt_llm._torch.models.modeling_gemma4.get_sm_version", + return_value=sm_version, + ), + ): + 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,7 +3417,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) + # On B200, head_dim>256 needs trtllm-gen. if info["head_dim"] > 256: kwargs["flashinfer_backend"] = "trtllm-gen" layers.append( From 0d132a823449f1df45d7d4ff5bf8ea5221f25462 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Thu, 13 Aug 2026 02:44:21 -0700 Subject: [PATCH 2/2] [NVBUG-6566891][fix] Select Gemma4 attention by kernel capability Gate TRTLLM-Gen on datacenter Blackwell capability and use FA2 elsewhere. Align Gemma4 backend selection and CUDA Graph coverage with the same capability check. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_gemma4.py | 11 +++--- .../_torch/modeling/test_modeling_gemma4.py | 39 ++++++++----------- 2 files changed, 21 insertions(+), 29 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index b94559a593dc..a9aa54d6029c 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -29,7 +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 get_sm_version +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 @@ -329,11 +329,10 @@ def __init__( # the rotate split, matching HF's rotate_half(head_dim//2) pairing. self.rotary_emb.head_dim = layer_head_dim - # trtllm-gen does not provide SM120/SM121 kernels. FlashInfer FA2 - # supports Gemma4 H256/H512, sliding-window attention, and FP8 KV cache - # on these architectures. Keep the existing trtllm-gen routing on all - # other architectures. - self.attn.flashinfer_backend = "fa2" if get_sm_version() in (120, 121) else "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 8db5bd6ea7d1..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,11 +2411,9 @@ def test_attn_backend_dispatches_to_flashinfer(self): "FLASHINFER must dispatch to FlashInferAttention", ) - @unittest.mock.patch( - "tensorrt_llm._torch.models.modeling_gemma4.get_sm_version", return_value=100 - ) - def test_all_layers_use_trtllm_gen_on_sm100(self, _mock_get_sm_version): - """All Gemma4 layers use trtllm-gen uniformly on SM100. + @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 @@ -2438,27 +2437,22 @@ def test_all_layers_use_trtllm_gen_on_sm100(self, _mock_get_sm_version): f"Layer {i} should use trtllm-gen", ) - def test_sm12x_layers_use_fa2(self): + @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 sm_version in (120, 121): - with ( - self.subTest(sm_version=sm_version), - unittest.mock.patch( - "tensorrt_llm._torch.models.modeling_gemma4.get_sm_version", - return_value=sm_version, - ), - ): - 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") + 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): @@ -3417,8 +3411,7 @@ def _run_cuda_graph_real_headdim(self, config_dict, label=""): layers = [] for info in layers_info: kwargs = {} - # On B200, head_dim>256 needs trtllm-gen. - if info["head_dim"] > 256: + if info["head_dim"] > 256 and is_sm_100f(): kwargs["flashinfer_backend"] = "trtllm-gen" layers.append( FlashInferAttention(