SM120 skip-softmax FMHA support probe misses bridge path
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 14.7k
- Forks
- 2.8k
- Avg merge
- 2d 23h
- Merged PRs (30d)
- 489
Description
Symptom
On SM120/SM121, supported BF16 causal PACKED_QKV context FMHA shapes can be dispatched by the hand-written skip-softmax bridge, but the build-time support probe rejects them.
Observed with Qwen3.6 MTP validation: the model falls back to unfused MHA, emits repeated missing-FMHA warnings, and allocates a much larger attention workspace.
Representative warning from the rejected support-probe path:
FMHA kernels are not found with these parameters:
Validation evidence from the tested fix:
FMHA missing warnings: 11 -> 0
unfused fallback warnings: 11 -> 0
max attention workspace: ~13.4 GB -> ~100 MB
batch 1 TPOT: 18.37 ms -> 17.09 ms
batch 8 TPOT: 33.06 ms -> 32.55 ms
Root Cause
The SM120/SM121 skip-softmax BF16 causal PACKED_QKV path is provided by a hand-written bridge in the skip-softmax translation unit, not by generated cubin metadata. FusedMultiHeadAttentionXMMAKernelV2::checkIfKernelExist() only consulted the generated kernel metadata hash, so it returned false for bridge-supported shapes even though run() could launch the bridge.
Tested Fix Summary
Teach checkIfKernelExist() to recognize the guarded SM120/SM121 BF16 causal PACKED_QKV bridge eligibility matrix, without adding a new kernel. Keep the predicate compile-time gated by TLLM_ENABLE_SKIP_SOFTMAX_SM120 and centralize the eligibility matrix so build-time support detection and runtime dispatch do not drift.
Fix under test: https://github.com/NVIDIA/TensorRT-LLM/pull/15759
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by inspecting FusedMultiHeadAttentionXMMAKernelV2::checkIfKernelExist() and the corresponding run() bridge dispatch described in the issue. Compare the guarded SM120/SM121 BF16 causal PACKED_QKV eligibility matrix with the existing generated-metadata check. Done means eligible bridge shapes pass support probing without adding a kernel and the reported validation warnings and fallback disappear.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- ai-infra-agents, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100