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

[lfr-0531]

Dev Engineer Review

• Gemma4 attention now uses trtllm-gen only on SM100f.
• Gemma4 uses FlashInfer FA2 on SM120, SM121, and other non-SM100f architectures.
• The routing applies to both sliding-window and full-attention layers.
• Tests cover backend selection and CUDA Graph behavior for both paths.
• No public API or dependency changes were introduced.
• No configuration files or test-list files changed.
• Validation passed, including pre-commit, SM120 H256/H512 checks, and a short NVFP4 generation with FP8 KV cache.

Verdict: sufficient.

QA Engineer Review

Modified test functions:

• test_all_layers_use_trtllm_gen was renamed to test_all_layers_use_trtllm_gen_on_sm100f.
• test_non_sm100f_layers_use_fa2 was added.
• CUDA Graph backend-selection coverage was updated for SM100f and non-SM100f architectures.

These tests are not listed in tests/integration/test_lists/ because no test-list files changed. The unit tests provide direct coverage for the updated backend-selection behavior.

Verdict: sufficient.

Description

Gemma4 currently forces every FlashInferAttention layer to use the TRTLLM-Gen inner backend. TRTLLM-Gen does not provide the required attention kernels for SM120/SM121, so this path reaches an unsupported-architecture assertion on DGX Spark (SM121).

This change uses FlashInfer FA2 for Gemma4 on SM120 and SM121 while preserving the existing TRTLLM-Gen path on all other architectures. It covers both the H256 sliding-window layers and H512 full-attention layers used by Gemma4. The change is limited to backend selection and does not introduce public API changes or new dependencies.

Related issue: https://nvbugspro.nvidia.com/bug/6566891

Related PR: #16214

Test Coverage

• Added backend-selection regression coverage for:
  • SM100: all Gemma4 layers continue to use TRTLLM-Gen.
  • SM120 and SM121: all sliding-window and full-attention layers use FA2.
• Built and installed TensorRT-LLM on an NVIDIA RTX PRO 6000 Blackwell Workstation Edition (SM120) with CMAKE_CUDA_ARCHITECTURES=120-real.
• Verified that the imported Python modules and compiled bindings resolved to the modified worktree.
• Ran python3 -m pytest -q tests/unittest/_torch/modeling/test_modeling_gemma4.py::TestGemma4ModelDefaults: 7 passed, 2 subtests passed.
• Loaded the real nvidia-Gemma-4-26B-A4B-NVFP4 configuration and verified that H256 sliding-window and H512 full-attention layers both selected FA2 on SM120.
• Completed a short end-to-end generation with the real NVFP4 checkpoint, FP8 KV cache, and max_seq_len=max_num_tokens=6272. FlashInfer JIT compiled an SM120f FA2 kernel, and the run completed with E2E_OK.
• Ran pre-commit run --all-files successfully.
• SM121 backend selection is covered by the regression test; SM121 hardware was not available for this validation.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 5d130827-52c5-456c-a3e7-a54d4e065a6d

📥 Commits

Reviewing files that changed from the base of the PR and between 31eb801 and 0d132a8.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/models/modeling_gemma4.py
• tests/unittest/_torch/modeling/test_modeling_gemma4.py

---

Walkthrough

Gemma4 now selects trtllm-gen on SM100f devices and FlashInfer FA2 on other architectures. Tests cover both dispatch paths and update CUDA graph backend selection.

Changes

Gemma4 backend dispatch

Layer / File(s)	Summary
Architecture-specific backend selection tensorrt_llm/_torch/models/modeling_gemma4.py	The model uses trtllm-gen on SM100f and FlashInfer FA2 on other architectures.
Backend selection validation tests/unittest/_torch/modeling/test_modeling_gemma4.py	Tests verify both backend paths and select trtllm-gen for large head dimensions only on SM100f devices.

Estimated code review effort: 2 (Simple) | ~10 minutes

Mergeability Score: 🔵 Low · up to 0d132

The change narrowly routes Gemma4 attention through FA2 on SM120/SM121 while preserving existing behavior elsewhere. It is mergeable with owner awareness that multimodal image-prefill using custom masks on the affected architectures still needs explicit correctness validation; no concrete failure is currently demonstrated.

Suggested reviewers: kaiyux, 2ez4bz, hudayday

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the fix and the affected Gemma4 architectures, SM120 and SM121.
Description check	✅ Passed	The description explains the issue, solution, test coverage, validation results, and checklist status.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65582 [ run ] triggered by Bot. Commit: 31eb801 Link to invocation

[brnguyen2]

Approving — the comments below are optional touch-ups, not blockers.

The fix lands in the right place and the mocked-get_sm_version tests are a reasonable way to pin behavior without SM12x hardware. Two things worth resolving before merge:

• The reported hardware is SM121, and validation was on SM120. The two arches share the trtllm-gen gap, so the fix should carry over, but a Spark confirmation before merge would close the loop.
• test_modeling_gemma4.py only appears in l0_b200.yml, so CI exercises the trtllm-gen branch exclusively; the fa2 branch is covered only by attribute assertions with a mocked SM version. Nothing here actually runs an fa2 Gemma4 kernel in CI. Worth saying explicitly in the PR description which fa2 paths were hand-validated (text generation) and which were not (multimodal prefill), since the entry point in the linked bug is the VLM one.

[tensorrt-cicd]

PR_Github #65582 [ run ] completed with state SUCCESS. Commit: 31eb801
/LLM/main/L0_MergeRequest_PR pipeline #53315 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65731 [ run ] triggered by Bot. Commit: 31eb801 Link to invocation

[tensorrt-cicd]

PR_Github #65731 [ run ] completed with state SUCCESS. Commit: 31eb801
/LLM/main/L0_MergeRequest_PR pipeline #53447 completed with status: 'SUCCESS'

CI Report

Link to invocation

[Hudayday]

LGTM!

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65883 [ run ] triggered by Bot. Commit: 0d132a8 Link to invocation

[tensorrt-cicd]

PR_Github #65883 [ run ] completed with state SUCCESS. Commit: 0d132a8
/LLM/main/L0_MergeRequest_PR pipeline #53576 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65939 [ run ] triggered by Bot. Commit: 0d132a8 Link to invocation

[tensorrt-cicd]

PR_Github #65939 [ run ] completed with state SUCCESS. Commit: 0d132a8
/LLM/main/L0_MergeRequest_PR pipeline #53628 completed with status: 'SUCCESS'

CI Report

Link to invocation