[#15791][fix] Enable SM120 skip-softmax FMHA support probe

[mihai-chiorean]

Fixes #15791

Summary

• Teach checkIfKernelExist() to recognize the guarded SM120/SM121 BF16 causal PACKED_QKV bridge path.
• Keep the predicate compile-time gated by TLLM_ENABLE_SKIP_SOFTMAX_SM120 and limited to the already-linked handwritten SM120 bridge; this does not add a new kernel.
• Centralize the bridge eligibility predicate so support detection and dispatch stay in sync.

Test Coverage

• git diff --check
• git diff --cached --check
• Spark validation on validate-blockers-811cc0732 @ df113db388332ea665efe9133fb0194abb18e3cd: FMHA missing warnings 11 -> 0; unfused fallback warnings 11 -> 0; max attention workspace ~13.4 GB -> ~100 MB; batch 1 TPOT 18.37 -> 17.09 ms; batch 8 TPOT 33.06 -> 32.55 ms; output coherent.
• Build log: /tmp/sm120_fmha_build_20260629_222414.log; before/after logs: /tmp/sm120_fmha_before_batch_20260629_222045.log, /tmp/sm120_fmha_after_batch_20260629_225519.log.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds a compile-time-guarded helper isSupportedBySkipSoftmaxSm120Bridge that checks whether a MHARunnerFixedParams configuration is supported by the sm_120/sm_121 skip-softmax launch bridge, and updates checkIfKernelExist to return true immediately when that helper reports support, bypassing the legacy kernel-hash lookup.

Skip-softmax SM120 bridge support

Layer / File(s)	Summary
Bridge helper and checkIfKernelExist integration cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_v2.cpp	Defines isSupportedBySkipSoftmaxSm120Bridge (behind TLLM_ENABLE_SKIP_SOFTMAX_SM120) checking SM value, BF16 data type, causal mask, packed QKV, head sizes 128/256, and several disabled feature flags; integrates it at the top of checkIfKernelExist to short-circuit the legacy hash-based lookup.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
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.
Description check	✅ Passed	The PR description covers the issue, solution, and test coverage, with only minor template drift from using Summary instead of Description.
Title check	✅ Passed	The title accurately reflects the new SM120 skip-softmax FMHA support probe added by the patch.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

🧹 Nitpick comments (1)

cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_v2.cpp (1)

56-72: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Keep the SM120 bridge eligibility matrix in one place.

This helper is now a second copy of the bridge routing predicate in run(). The current PR fixes one drift between support detection and dispatch, but the next capability tweak can desynchronize them again. Please route both paths through one shared predicate and keep only the runtime-only checks (flash_attention, interleaved, softmax-stats pointer) outside it.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_v2.cpp`
around lines 56 - 72, The SM120 bridge eligibility logic is duplicated between
isSupportedBySkipSoftmaxSm120Bridge and the routing in run(), which can drift
over time. Refactor both paths to use one shared predicate for the capability
matrix, and leave only runtime-only checks such as flash_attention, interleaved,
and the softmax-stats pointer in the dispatch path. Keep the shared eligibility
logic centralized in the helper and have run() call it directly instead of
re-encoding the same conditions.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In
`@cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_v2.cpp`:
- Around line 56-72: The SM120 bridge eligibility logic is duplicated between
isSupportedBySkipSoftmaxSm120Bridge and the routing in run(), which can drift
over time. Refactor both paths to use one shared predicate for the capability
matrix, and leave only runtime-only checks such as flash_attention, interleaved,
and the softmax-stats pointer in the dispatch path. Keep the shared eligibility
logic centralized in the helper and have run() call it directly instead of
re-encoding the same conditions.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: d87717dc-aee5-411f-a161-7ec217bf6d46

📥 Commits

Reviewing files that changed from the base of the PR and between 19b674e and 7bab590.

📒 Files selected for processing (1)

• cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_v2.cpp

[waynehacking8]

Datapoint from an RTX PRO 6000 (SM120, x86_64, CUDA 13.0): I built current main (e4aba85) bare-metal with --cuda_architectures "120-real" (needs #15934 to get past a ptxas failure in moeAlltoAllKernels.cu) and drove the real FusedMultiHeadAttentionXMMAKernelV2::checkIfKernelExist through a small driver linked against the built libtensorrt_llm.so, compiled with the TU's exact define set. BF16 causal PACKED_QKV hd 128/256 already returns true for both sm=120 and sm=121 on stock main - this build carries the generated sm120 fmha_v2 BF16 cubins, which satisfy the probe (note KernelExistPredicate is a bitmask-subset match, not equality), so I could not reproduce the probe-miss this PR fixes.

Your validation was on validate-blockers-811cc0732 @ df113db on Spark, so the miss likely depends on which cubins that build carried. Might be worth re-checking whether the missing-FMHA warnings still occur on current main on Spark before merging - if they do, the delta between our builds (cubin set or arch list) would itself be informative. I did not run a serving-level A/B; this is probe-level only.

[pengbowang-nv]

Independent SM120 A/B result on public main:

Environment

• NVIDIA RTX PRO 6000 Blackwell Server Edition, SM120
• x86_64, CUDA 13.2
• Baseline: main@29f4552bfc
• Build: scripts/build_wheel.py -a 120-real
• PR [None][fix] Exclude SM12x from bf16x2 FP8 cvt specialization in moe_a2a #15934 applied symmetrically only to avoid the unrelated SM12x ptxas failure

Results

Check	Baseline	PR #15759
checkIfKernelExist() — SM120/121, BF16 causal PACKED_QKV, hd 128/256, FP32 acc on/off	8/8 true	8/8 true
test_skip_softmax_sm120.py	14 passed	14 passed
Missing-FMHA warnings	0	0
Unfused fallback warnings	0	0
Maximum workspace reported by the test	33,559,296 bytes	33,559,296 bytes

The probe was a real C++ driver linked against the locally built libtensorrt_llm.so. Extracted probe output was identical before and after the PR.

The stock SM120 build already generates BF16 causal PACKED_QKV metadata for head dimensions 128 and 256. Since KernelExistPredicate performs a bitmask-subset match rather than exact hash equality, those entries satisfy the support probe without the new bridge short-circuit.

Recommendation: close PR #15759 as unnecessary on current main, and close Issue #15791 as not reproducible on a standard SM120 build.

[waynehacking8]

Careful with my 07-04 datapoint - I don't think it supports closing this. On sm_120 the bf16 causal PACKED_QKV path is a hand-written bridge (run_skip_softmax_bf16_d128_causal_sm120 / _d256_, declared in fused_multihead_attention_v2.cpp and defined in skip_softmax_sm120/fused_multihead_flash_attention_ws_sm120.cu), not generated cubin metadata, and checkIfKernelExist builds its query with the flash/warp-spec/tiled bits zeroed and then matches with (stored & query) == query - so a true can come from an unrelated superset entry rather than from a kernel that actually covers the config.

That also makes the A/B non-discriminating: #15759 short-circuits true for exactly the configs your baseline already returns true for, so 8/8 -> 8/8 and an unchanged workspace is what you'd see whether or not the bug is real. And the ~32 MiB workspace on both sides suggests the failing precondition just never occurred on that build - the reporter saw ~13.4 GB and 11 missing-FMHA warnings on a Spark build, so I'd reproduce there before calling it not reproducible.

[mihai-chiorean]

Thanks for the investigation. I reran the validation on clean builds with #15759 absent.

Results:

• SM121 DGX Spark: Qwen3.6-35B-A3B-NVFP4 served coherently with TRTLLM attention, MTP, CUDA graphs, and overlap enabled. No missing-FMHA or unfused-MHA warnings; maximum attention workspace was 95.59 MiB.
• SM120 RTX PRO 6000: both no-MTP and MTP serving completed coherently. The focused SM120 test suite passed 14/14. Workspace remained approximately 95 MiB, with no 13.4 GB allocation or missing-FMHA failure.

The earlier subset-hash probe was not sufficient by itself, but these clean production-path runs exercise the actual serving path and do not reproduce the reported failure.

I’m therefore closing this PR as unnecessary on current main.