Skip to content

[Perf][Qwen4Exp] Add SM121 TP=1 skinny-GEMM plans #59753

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
benchislett merged 4 commits into from
Oct 2, 2026
+136 −0
Merged

[Perf][Qwen4Exp] Add SM121 TP=1 skinny-GEMM plans #59753

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
31e8b5a
[Perf][Qwen4Exp] Add SM121 TP=1 LM-head skinny-GEMM plans
stecasta Oct 2, 2026
6a1e7e3
[Perf][Qwen4Exp] Extend SM121 TP=1 plans to the dense BF16 projections
stecasta Oct 2, 2026
c8fddd2
Merge branch 'main' into feat/qwen4exp-sm121-lmhead-plans
benchislett Oct 2, 2026
e5fa725
Merge branch 'main' into feat/qwen4exp-sm121-lmhead-plans
mgoin Oct 2, 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
30 changes: 30 additions & 0 deletions tests/kernels/test_bf16_skinny_gemm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,12 @@
for num_tokens, config in plans.items()
]

QWEN4_EXP_SM121_CASES = [
(n, k, num_tokens, config)
for (n, k), plans in qwen4_exp_gemm.QWEN4_EXP_SM121_GEMM_PLANS.items()
for num_tokens, config in plans.items()
]

EXPECTED_CUTE_CONFIGS = {
(3072, 7168, 1): (224, 3, 4, 8),
(3072, 7168, 2): (128, 3, 2, 8),
Expand Down Expand Up @@ -725,6 +731,7 @@ def test_qwen4_exp_hopper_plans_are_valid() -> None:
((10, 3), qwen4_exp_gemm.QWEN4_EXP_GEMM_PLANS),
((10, 0), qwen4_exp_gemm.QWEN4_EXP_SM100_GEMM_PLANS),
((9, 0), qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS),
((12, 1), qwen4_exp_gemm.QWEN4_EXP_SM121_GEMM_PLANS),
((8, 0), {}),
],
)
Expand Down Expand Up @@ -939,6 +946,29 @@ def test_qwen4_exp_sm100_selected_shapes(
assert cosine > 0.999


@pytest.mark.parametrize("n,k,num_tokens,config", QWEN4_EXP_SM121_CASES)
def test_qwen4_exp_sm121_selected_shapes(
n: int,
k: int,
num_tokens: int,
config: SkinnyGemmConfig,
) -> None:
_require_capability_and_cute((12, 1))
torch.manual_seed(42 + num_tokens)
x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")

selected = qwen4_exp_gemm.QWEN4_EXP_SM121_GEMM_PLANS[(n, k)][num_tokens]
assert selected == config
output = qwen4_exp_gemm._qwen4_exp_low_latency_gemm(x, weight)

reference = torch.nn.functional.linear(x, weight)
cosine = torch.nn.functional.cosine_similarity(
output.float().flatten(), reference.float().flatten(), dim=0
).item()
assert cosine > 0.999


def test_glm52_q_b_nonpacked_single_row_falls_back() -> None:
_require_sm103_and_cute()
spec = glm52_gemm.GLM52_Q_B_PROJECTION
Expand Down
106 changes: 106 additions & 0 deletions vllm/models/qwen4_exp/nvidia/low_latency_gemm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -207,6 +207,110 @@
}


# DGX Spark (GB10) plans, TP=1.
QWEN4_EXP_SM121_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
# Shared-expert gate.
(1, 2560): {
1: SkinnyGemmConfig(1, 64, 1, k_unroll=2, static_k=2560),
2: SkinnyGemmConfig(2, 64, 1, k_unroll=2, vector_width=4, static_k=2560),
4: SkinnyGemmConfig(4, 128, 1, vector_width=2, static_k=2560),
8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
16: SkinnyGemmConfig(16, 128, 1, vector_width=4, static_k=2560),
},
# GDN fused B/A projection.
(96, 2560): {
1: SkinnyGemmConfig(1, 128, 4, k_unroll=2, vector_width=2, static_k=2560),
2: SkinnyGemmConfig(2, 64, 4, vector_width=2, static_k=2560),
4: SkinnyGemmConfig(4, 128, 4, vector_width=2, static_k=2560),
8: SkinnyGemmConfig(8, 128, 4, k_unroll=2, vector_width=4, static_k=2560),
16: SkinnyGemmConfig(16, 64, 1),
},
# HC merged down/injection projection.
(336, 10240): {
1: SkinnyGemmConfig(1, 256, 1, vector_width=4, static_k=10240),
2: SkinnyGemmConfig(2, 256, 1, k_unroll=4, vector_width=4, static_k=10240),
4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, static_k=10240),
8: SkinnyGemmConfig(8, 128, 4, k_unroll=4, static_k=10240),
16: SkinnyGemmConfig(16, 128, 4, k_unroll=4, static_k=10240),
},
# Router.
(512, 2560): {
1: SkinnyGemmConfig(1, 256, 1, k_unroll=4, vector_width=2, static_k=2560),
2: SkinnyGemmConfig(2, 256, 1, k_unroll=4, vector_width=2, static_k=2560),
4: SkinnyGemmConfig(4, 64, 1, k_unroll=2, static_k=2560),
8: SkinnyGemmConfig(8, 64, 2, static_k=2560),
16: SkinnyGemmConfig(16, 64, 2, k_unroll=2, static_k=2560),
},
# QSA indexer Q/K projection.
(640, 2560): {
1: SkinnyGemmConfig(1, 256, 1, vector_width=2, static_k=2560),
2: SkinnyGemmConfig(2, 256, 1, vector_width=2, static_k=2560),
4: SkinnyGemmConfig(4, 128, 1, k_unroll=2, vector_width=4, static_k=2560),
8: SkinnyGemmConfig(8, 128, 1, vector_width=4, static_k=2560),
16: SkinnyGemmConfig(16, 64, 1, static_k=2560),
},
# Shared-expert fused gate/up projection.
(1280, 2560): {
1: SkinnyGemmConfig(1, 32, 1, k_unroll=2, static_k=2560),
2: SkinnyGemmConfig(2, 64, 1, vector_width=4, static_k=2560),
4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, vector_width=2, static_k=2560),
8: SkinnyGemmConfig(8, 128, 1, k_unroll=2, vector_width=2, static_k=2560),
16: SkinnyGemmConfig(16, 32, 1, k_unroll=2, static_k=2560),
},
# Shared-expert down projection.
(2560, 640): {
1: SkinnyGemmConfig(1, 64, 1, k_unroll=4, vector_width=2, static_k=640),
2: SkinnyGemmConfig(2, 64, 1, vector_width=2, static_k=640),
4: SkinnyGemmConfig(4, 64, 1, vector_width=2, static_k=640),
8: SkinnyGemmConfig(8, 32, 1, k_unroll=4, vector_width=4, static_k=640),
16: SkinnyGemmConfig(16, 32, 1, k_unroll=2, vector_width=4, static_k=640),
},
# GDN and QSA output projections.
(2560, 6144): {
1: SkinnyGemmConfig(1, 64, 2, vector_width=4, static_k=6144),
2: SkinnyGemmConfig(2, 64, 1, k_unroll=4, vector_width=4, static_k=6144),
4: SkinnyGemmConfig(4, 64, 1, k_unroll=4, vector_width=4, static_k=6144),
8: SkinnyGemmConfig(8, 128, 1, vector_width=2, static_k=6144),
16: SkinnyGemmConfig(16, 32, 1, static_k=6144),
},
# HC up projection.
(10240, 320): {
1: SkinnyGemmConfig(1, 32, 1, k_unroll=2, vector_width=2, static_k=320),
2: SkinnyGemmConfig(2, 32, 1, vector_width=2, static_k=320),
4: SkinnyGemmConfig(4, 32, 1, k_unroll=2, vector_width=2, static_k=320),
8: SkinnyGemmConfig(8, 32, 1, vector_width=2, static_k=320),
},
# QSA fused QKV/gate projection.
(13312, 2560): {
1: SkinnyGemmConfig(1, 32, 1, vector_width=4, static_k=2560),
2: SkinnyGemmConfig(2, 32, 1, vector_width=4, static_k=2560),
4: SkinnyGemmConfig(4, 32, 1, vector_width=4, static_k=2560),
8: SkinnyGemmConfig(8, 64, 1, vector_width=2, static_k=2560),
16: SkinnyGemmConfig(16, 32, 1, vector_width=4, static_k=2560),
},
# GDN fused QKVZ projection.
(16384, 2560): {
1: SkinnyGemmConfig(1, 32, 1, k_unroll=4, vector_width=2, static_k=2560),
2: SkinnyGemmConfig(2, 32, 1, vector_width=2, static_k=2560),
4: SkinnyGemmConfig(4, 32, 1, k_unroll=2, vector_width=2, static_k=2560),
8: SkinnyGemmConfig(8, 64, 1, vector_width=2, static_k=2560),
16: SkinnyGemmConfig(16, 32, 1, k_unroll=2, vector_width=2, static_k=2560),
},
# LM head.
(248320, 2560): {
1: SkinnyGemmConfig(1, 128, 1, k_unroll=2, vector_width=4, static_k=2560),
2: SkinnyGemmConfig(2, 128, 1, k_unroll=2, vector_width=4, static_k=2560),
4: SkinnyGemmConfig(4, 32, 1, vector_width=4, static_k=2560),
8: SkinnyGemmConfig(8, 64, 1, k_unroll=2, vector_width=4, static_k=2560),
16: SkinnyGemmConfig(16, 128, 1, k_unroll=2, vector_width=4, static_k=2560),
},
}


def _is_sm121() -> bool:
return current_platform.is_device_capability((12, 1))


def _is_sm100() -> bool:
return current_platform.is_device_capability((10, 0))

Expand All @@ -226,6 +330,8 @@ def _gemm_plans() -> dict[tuple[int, int], dict[int, SkinnyGemmConfig]]:
return QWEN4_EXP_SM100_GEMM_PLANS
if _is_sm90():
return QWEN4_EXP_SM90_GEMM_PLANS
if _is_sm121():
return QWEN4_EXP_SM121_GEMM_PLANS
return {}


Expand Down
Loading