Skip to content

SM120 skip-softmax FMHA support probe misses bridge path #15791

Description

@mihai-chiorean

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: #15759

Metadata

Metadata

Assignees

No one assigned

    Labels

    Customized kernels<NV>Specialized/modified CUDA kernels in TRTLLM for LLM ops, beyond standard TRT. Dev & perf.

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions