Skip to content

[None][feat] Update FMHA cubins for head_dim 80 #13808

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yuxianq merged 3 commits into from
May 13, 2026
Merged

[None][feat] Update FMHA cubins for head_dim 80 #13808

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6e026bb
[None][feat] Support headDimV=80 for trtllm-gen FMHA kernels on Black…
yuxianq Apr 13, 2026
544ecb4
[None][chore] Update trtllm-gen FMHA metadata version
yuxianq May 11, 2026
94d5bc4
Merge branch 'main' into trtllm-gen-head-80
yuxianq May 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,10 +49,10 @@ QkvLayout AttentionInputLayoutToQkvLayout(AttentionInputLayout layout)

FmhaDispatcher::FmhaDispatcher(MHARunnerFixedParams fixedParams)
: mFixedParams(fixedParams)
// TRTLLM-GEN only supports power of 2 head sizes.
// TRTLLM-GEN only supports power of 2 head sizes (and 80 with padding).
// The exception will fall back to fmha v2.
// Please update fmha_v2/setup.py if you want to add more supported head sizes.
, mUseTllmGen(tensorrt_llm::common::isSM100Family() && fixedParams.headSize != 80 && fixedParams.headSize != 72)
, mUseTllmGen(tensorrt_llm::common::isSM100Family() && fixedParams.headSize != 72)
{
if (mUseTllmGen)
{
Expand Down
2 changes: 1 addition & 1 deletion cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/cubin/kernelMetaInfo.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ namespace kernels
{
// clang-format off

#define TLLM_GEN_VERSION "3977eaf5-dirty"
#define TLLM_GEN_VERSION "09fbb2db-dirty"
#ifndef EXCLUDE_SM_100
extern unsigned char const FmhaSm100aKernel_QE4m3KvE2m1OE4m3H128PagedKvCausalP32VarSeqQ128Kv128PersistentContext_cubin[];
extern unsigned char const FmhaSm100aKernel_QE4m3KvE2m1OE4m3H128PagedKvCausalP32VarSeqQ128Kv128StaticContext_cubin[];
Expand Down
2 changes: 1 addition & 1 deletion tests/unittest/trt/attention/test_gpt_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -316,7 +316,7 @@ def load_test_cases():
[2], # batch_size
[165], # in_len
[2, 8, 32], # num_q_heads
[32, 64, 96, 128, 160], # head_size
[32, 64, 80, 96, 128, 160], # head_size
[2], # num_kv_heads
[False],
[1, 2, 4], # beam_width
Expand Down
4 changes: 2 additions & 2 deletions tests/unittest/utils/util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -157,10 +157,10 @@ def skip_fp8_pre_ada(use_fp8):


def skip_blackwell_for_fmha_tests(context_fmha_type, head_size):
if (isSM100Family()) and (head_size not in [32, 64, 128] and
if (isSM100Family()) and (head_size not in [32, 64, 80, 128] and
context_fmha_type != ContextFMHAType.disabled):
pytest.skip(
"Context FMHA only supports head sizes [32, 64, 128] currently on blackwell."
"Context FMHA only supports head sizes [32, 64, 80, 128] currently on blackwell."
)


Expand Down
Loading