Repository navigation

[7/N][warmup][DSv4] Migrate router GEMM kernels - #50177

Open

LopezCastroRoberto wants to merge 14 commits into

vllm-project:mainfrom

LopezCastroRoberto:warmup_migration/dsv4_4

Open

[7/N][warmup][DSv4] Migrate router GEMM kernels#50177
LopezCastroRoberto wants to merge 14 commits into
vllm-project:mainfrom
LopezCastroRoberto:warmup_migration/dsv4_4

LopezCastroRoberto commented Jul 28, 2026 •

edited

Loading

Contributor

Depends on: #53566

For more details, see parent PR: #49627 and tracking issue #49349

Description

This PR migrates the DSv4 router GEMM kernels to the shared JIT warmup contract.

What Changed

Migrated low-latency BF16 and BF16x3 router GEMMs.
Moved specialization selection into kernel-owned dispatch methods.
Added compile-only CuTeDSL and Triton paths.
Preserved existing tuning tables, runtime heuristics, and performance specializations.
Removed the corresponding legacy router warmup orchestration.
Added dispatch-equivalence and kernel-specific warmup tests.

LopezCastroRoberto requested review from AndreasKaratzas, DarkLight1337, ProExpertProg, WoosukKwon, houseroad, mgoin, pavanimajety, robertgshaw2-redhat, tlrmchlsmth, yewentao256, youkaichao, ywang96 and zyongye as code owners

July 28, 2026 19:26

claude Bot reviewed

View reviewed changes

claude Bot left a comment

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

LopezCastroRoberto changed the title ~~[2/N][warmup][DSv4] Migrate MoE kernels~~ [3/N][warmup][DSv4] Migrate MoE kernels

mergify Bot added deepseek nvidia v1 labels

github-project-automation Bot added this to NVIDIA

This was referenced Jul 28, 2026

[DO NOT MERGE][Perf][Feat] DSv4 de-JITification #49627

Open

Zero JIT compilation during runtime #49349

Open

mergify Bot commented Jul 29, 2026

Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @LopezCastroRoberto.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

mergify Bot added the needs-rebase label

LopezCastroRoberto marked this pull request as draft

July 29, 2026 17:11

LopezCastroRoberto mentioned this pull request

[3/N][Feat][Perf] Add new warmup infrastructure for JITs. Add provider registry and orchestration for JIT warmup #50174

Merged

mergify Bot added inkling DSv4 labels

LopezCastroRoberto force-pushed the warmup_migration/dsv4_4 branch from aec29f7 to 5efe9bd Compare

August 24, 2026 10:17

LopezCastroRoberto changed the title ~~[3/N][warmup][DSv4] Migrate MoE kernels~~ [6/N][warmup][DSv4] Migrate router GEMM kernels

mergify Bot removed the needs-rebase label

mergify Bot removed the needs-rebase label


          update

75341a9

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>

mergify Bot commented Sep 11, 2026

Contributor

Hi @LopezCastroRoberto, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.


          update

25b2148

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>

mergify Bot commented Sep 11, 2026

Contributor

Hi @LopezCastroRoberto, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.


          update

141540e

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>

mergify Bot commented Sep 11, 2026

Contributor

Hi @LopezCastroRoberto, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.


          update

17ef361

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>

mergify Bot commented Sep 11, 2026

Contributor

Hi @LopezCastroRoberto, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

LopezCastroRoberto added a commit to LopezCastroRoberto/vllm that referenced this pull request


          feat(warmup): apply router GEMM migration (vllm-project#50177)

7bc5658

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>

LopezCastroRoberto added a commit to LopezCastroRoberto/vllm that referenced this pull request


          fix(warmup): sync latest router migration (vllm-project#50177)

ba4165b

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>


          Merge branch 'main' into warmup_migration/dsv4_4

139b1c3

LucasWilkinson commented Sep 13, 2026

Contributor

/ci run

LucasWilkinson enabled auto-merge (squash)

September 13, 2026 20:49

github-actions Bot commented Sep 13, 2026

✅ Triggered Buildkite CI #88682 for commit 139b1c3f1772.

LucasWilkinson disabled auto-merge

September 13, 2026 20:49

mergify Bot commented Sep 23, 2026

Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @LopezCastroRoberto.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

mergify Bot added the needs-rebase label

lucifer1004 mentioned this pull request

[Warmup] Cover DSv4 sparse top-k and DFlash input-prep Triton specializations that JIT at first request #58559

Open

gau-nernst mentioned this pull request

[Perf] Faster BF16x3 router GEMM at large batch sizes on SM100 #60376

Draft

3 of 6 tasks

douglas-reid added a commit to douglas-reid/vllm that referenced this pull request


          [Perf][MoE] Wrap GateLinear ll_bf16_gemm in custom op to prevent torc…

268e68a

…h.compile branch freezing

In GateLinear.forward, Tier 1 (ll_bf16_gemm, the CuTe DSL low-latency
BF16->FP32 router GEMM for M <= 16 on SM90+) was gated by an inline
Python conditional on x.shape[0]. Unlike Tier 2
(torch.ops.vllm.fp32_router_gemm_dispatch) and Tier 4
(torch.ops.vllm.rocm_bf16x3_router_gemm_dispatch), Tier 1 did not wrap
its runtime x.shape[0] check in a custom op.

When GateLinear is called from @support_torch_compile model code,
vLLM's TorchCompileWithNoGuardsWrapper drops all Dynamo guards, so a
plain Python branch on x.shape[0] is frozen at first trace: tracing
first at M > 16 (e.g. M = 2048 during warmup) freezes the branch to
False and permanently falls back to Tier 5
torch.mm(x, weight.T, out_dtype=torch.float32) during decode, while an
M <= 16 first trace would freeze ll_bf16_gemm in at M = 8192.

Tier 1 requires bf16 weights + out_dtype=torch.float32 + bias=False +
unquantized + SM90+ and a gate called from @support_torch_compile model
code. On main that is Gemma 4 (Gemma4Router.proj), Nemotron-H
(force_fp32_compute=True keeps bf16 weights when specialized kernels
exist), and MiMo-V2 when moe_router_dtype is bfloat16. After this change,
Gemma 4 decode (M <= 16) and prefill (M > 16) use different GEMM kernels,
so router logits may differ in the last ULP across batch sizes; this is
Tier 1's pre-existing design (DeepSeek already behaves this way) and
VLLM_BATCH_INVARIANT gating is handled at init by vllm-project#59985. The c1_pdl
warmup gap is addressed by vllm-project#50177.

Changes:
1. Register torch.ops.vllm.ll_bf16_router_gemm_dispatch(x, weight,
   max_tokens) so 0 < x.shape[0] <= max_tokens is evaluated dynamically
   at runtime outside the compiled FX graph, calling
   ll_bf16_gemm(x.contiguous(), weight) for 0 < m <= max_tokens and
   falling back to Tier 5 torch.mm(x, weight.T, out_dtype=torch.float32)
   otherwise.
2. Update test_gate_linear_torch_compile_dynamic_m_dispatch to drop
   guards via guard_filter_fn and run decode calls under
   torch.compiler.set_stance("fail_on_recompile") so the test fails on
   main and passes with the custom op.

Signed-off-by: Douglas Reid <21148125+douglas-reid@users.noreply.github.com>

douglas-reid mentioned this pull request

[Perf][MoE] Wrap GateLinear ll_bf16_gemm in custom op to prevent torch.compile branch freezing #60681

Draft

LopezCastroRoberto added the warmup label

ischencheng mentioned this pull request

[Perf][Kernel] Add a CuteDSL BF16x3 router GEMM for fp32 router weights on Hopper #61903

Merged

6 tasks done

WoosukKwon mentioned this pull request

[Kernel] Add router GEMM kernels for (5120, 256) and BF16 output to ll_bf16_gemm #62042

Merged

6 tasks done

new-TonyWang mentioned this pull request

[Perf][Kernel] Enable split-K 16 for the SM90 BF16x3 router GEMM #64054

Closed

6 tasks done

This branch has not been deployed

No deployments

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

Labels

deepseek DSv4 inkling needs-rebase nvidia ready v1 warmup