NVIDIA / NVIDIA/TensorRT-LLM

SM120 skip-softmax FMHA support probe misses bridge path

Open
#15,791 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Customized kernels
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.