[amukkara]

Description

Use tiled fmha_v2 kernel for FP8 (E4M3) context attention with head_dim 256 on SM120/SM121. Targets nvidia/Qwen3.6-35B-A3B-NVFP4 which uses FP8 KV cache.

• cpp/kernels/fmha_v2/setup.py: add a tiled 64x128 QMMA kernel for head_dim 256, alongside the existing non-tiled 64x32 kernel.
• fmhaRunner.cpp: restrict the forced non-tiled FP8 selection to SM89, so SM120/SM121 use the existing head-size heuristic, which selects tiled kernels for head_dim >= 256.

Performance

fmha_v2 harness on RTX PRO 6K BSE, FP8 Q/K/V, head_dim 256, tiled vs non-tiled kernel.
The tiled kernel is faster on all 35 shapes tested: 2.03x - 3.52x, geomean 3.01x.

sweep	range	speedup
KV length (b=1, 16 Q / 4 KV heads)	65 -> 32k	2.15x -> 3.32x
batch size	4 -> 64 at KV 256 / 1k / 4k	2.03x -> 3.24x
head config	8/2, 32/8, 32/32	2.93x -> 3.23x
chunked prefill	query chunk 128 / 512, KV 2k / 8k	3.38x -> 3.43x
packed QKV, padding mask, FP8 output	KV 1k / 4k	2.99x -> 3.52x

Test Coverage

The tiled kernel passes the fmha_v2 harness reference check on SM120 (-epsilon 0.2) for paged KV and packed QKV; causal, padding, sliding-window and custom masks; ALiBi; GQA; chunked prefill; and FP8 and BF16 output.

e2e accuracy verified by tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestQwen3_6_35B_A3B::test_nvfp4_w4a16

Dev Engineer Review

setup.py adds an SM120-only tiled 64x128 QMMA configuration for FP8 E4M3 with head_dim 256. fmhaRunner.cpp limits forced non-tiled FP8 selection to SM89, allowing SM120/SM121 to use the existing tiling heuristic. The reported benchmark and accuracy results are author-reported; no test execution results were provided.

QA Engineer Review

No test changes.

Per-File QA Perspective

• cpp/kernels/fmha_v2/setup.py: Adds a tiled-kernel configuration for SM120 and head_dim 256. QA should verify kernel selection and output correctness for supported shapes.
• cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp: Changes FP8 E4M3 kernel selection for SM120/SM121. QA should verify that the tiling heuristic selects the expected kernel and preserves output correctness.

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: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: ce39a9e1-0b74-468e-af8e-a89765dc5855

📥 Commits

Reviewing files that changed from the base of the PR and between 54f2dbc and 886388c.

📒 Files selected for processing (1)

• cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

Walkthrough

The FMHA kernel enumerator adds an SM120 tiled configuration for head size 256. The runner routes SM120/SM121 E4M3 cases through tiled-kernel selection. The setup file updates its copyright range through 2026.

Changes

SM120 FMHA kernel selection

Layer / File(s)	Summary
Tiled kernel configuration and selection cpp/kernels/fmha_v2/setup.py, cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp	The enumerator adds an SM120 tiled configuration for head size 256 with Q/KV loop steps 64/128. The runner limits the E4M3 non-tiled-kernel branch to SM89.

Priority: ➖ Normal

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

Change: Feature

Suggested reviewers: brnguyen2

Merge Risk: ⚪ Minimal · up to 88638

No specific user-impacting issue was established in the changed kernel-selection path; unsupported configurations retain the unfused fallback.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: enabling tiled FP8 context FMHA for head_dim 256 on SM120. It uses the repository’s ticket and type format.
Description check	✅ Passed	The description explains the change and its purpose, reports performance results, identifies test coverage, and completes the PR checklist. It provides the key information required by the template.
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.

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Autopilot is currently an internal CodeRabbit preview.

---

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

[amukkara]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #76442 [ run ] triggered by Bot. Commit: 886388c Link to invocation

[tensorrt-cicd]

PR_Github #76442 [ run ] completed with state FAILURE. Commit: 886388c
/LLM/main/L0_MergeRequest_PR pipeline #63032 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

[amukkara]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #76489 [ run ] triggered by Bot. Commit: 886388c Link to invocation

[tensorrt-cicd]

PR_Github #76489 [ run ] completed with state SUCCESS. Commit: 886388c
/LLM/main/L0_MergeRequest_PR pipeline #63069 completed with status: 'SUCCESS'

CI Report

Link to invocation

[farazkh80]

LGTM, thanks!

[github-actions]

Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements.