Skip to content

[Bugfix][Kernel] Only set PDL launch attr for MoE top-k gating on SM90+ - #65059

Merged
gau-nernst merged 4 commits into
vllm-project:mainfrom
srkim-business:fix/topk-gating-pdl-pre-sm90
Nov 2, 2026
Merged

gau-nernst merged 4 commits into
vllm-project:mainfrom
srkim-business:fix/topk-gating-pdl-pre-sm90

Conversation

@srkim-business

Copy link
Copy Markdown
Contributor

Overview

Fixes #65055. Since #59737, topkGating is launched with the PDL attribute on every CUDA arch. Pre-SM90 GPUs cannot use PDL, but stream capture still records the attribute as a programmatic edge in the CUDA graph, and replaying those graphs is slower. This PR sets the attribute only when the device is SM90+. SM90+ keeps the #59737 behaviour; older GPUs launch without the attribute, which measured the same as the v0.32 <<<>>> launch. The standalone effect was measured on sm80, sm86 and sm89; end-to-end numbers are A100 only.

Claims

  • 4× A100, TP4 + EP + MTP (our deployment config, private NVFP4 checkpoint of Qwen3.8-Flash-Next; setup below): per-request decode is 6.5–13.6% faster in 17 of 18 prompt × concurrency (1–16) cells, and 2.1% in the 18th. At concurrency 16: prose 74.3 → 80.1 and code 88.7 → 94.8 tok/s, close to the last commit before [Perf] Enable PDL for fused MoE top-k gating #59737 (3b05805: 81.5 and 95.3, one run).
  • Standalone CUDA (no vLLM): the attribute adds one programmatic edge per launch and 0.6–1.1 µs of graph replay per edge on A100 (sm80), RTX A4000 (sm86) and RTX 4090 (sm89), in every round; without the attribute cudaLaunchKernelEx matches <<<>>>.
  • 1× A100, CUDA graph replay of a router GEMM → topk_softmax chain: 6.9–9.0% faster for Qwen3.8-Flash-Next, Qwen3-30B-A3B and OLMoE shapes at 1/16/48 tokens, in every round against a no-change rebuild.
  • 1× A100 end to end, vllm bench latency: batch 1 is 0.7% faster for OLMoE-1B-7B (every round vs a no-change rebuild, 2 of 3 rounds vs the shipped build) and 1.2% faster for Qwen3-30B-A3B-GPTQ-Int4 (every round). Batch 16 is within noise for both.
  • No accuracy change: GSM8K 5-shot on Qwen3-30B-A3B-GPTQ-Int4 is 85.90% (base) vs 86.13% (fix). The kernel body is identical.
  • SM90+ is unchanged by construction (the same attribute is set as before); not measured.

Validation

Hardware: A100-PCIE-40GB, driver 590.48.01, CUDA 13.0, torch 2.14.0+cu130, vLLM v0.33.0 (4ba787e). base = the v0.33.0 _moe_C as shipped. rebuilt = the same source built with the same toolchain, no change. fix = rebuilt + this diff. Builds are _moe_C only, for sm80, overlaid on the v0.33.0 install. The diff applies unchanged to main (the file is identical).

Tests

  • pytest tests/kernels/moe/test_fused_topk.py on A100: base 864 passed, fix 864 passed.
  • GSM8K 5-shot, 1319 questions (tests/evals/gsm8k/gsm8k_eval.py): Qwen3-30B-A3B-GPTQ-Int4 85.90% → 86.13%; OLMoE-1B-7B (base model, near the floor with this script) 9.93% → 10.01%. Invalid rate 0–0.8%.
  • CI: the L4 (sm89) "MoE Kernels" step runs tests/kernels/moe on csrc/moe/ changes and covers the launch without the attribute; the SM90+ launch is unchanged. No new test: there is no functional change, and a regression check would need graph-edge inspection or timing, which does not belong in tests/.

Op chain, CUDA graph, µs per replay (median of 6 rounds, A/B/B/A order; Δ = median paired fix − rebuilt)

shape (experts, top-k, hidden, layers) tokens base rebuilt fix Δ
Qwen3.8-Flash-Next (512, 10, 2560, 48) 1 / 16 / 48 843 / 893 / 917 846 / 894 / 916 782 / 831 / 850 −7.6% / −6.9% / −7.2%
Qwen3-30B-A3B (128, 8, 2048, 48) 1 / 16 / 48 592 / 672 / 688 596 / 674 / 691 542 / 620 / 636 −9.0% / −7.9% / −8.0%
OLMoE-1B-7B (64, 8, 2048, 16) 1 / 16 / 48 194 / 209 / 221 194 / 209 / 221 177 / 192 / 204 −8.7% / −8.4% / −7.8%

End to end, vllm bench latency (input 128, output 256, 3 rounds, ms per step)

model batch base rebuilt fix Δ
OLMoE-1B-7B-0924 1 2.952 2.950 2.932 −0.7% (3/3 rounds vs rebuilt, 2/3 vs base)
OLMoE-1B-7B-0924 16 7.789 7.731 7.802 within noise
Qwen3-30B-A3B-GPTQ-Int4 1 5.805 – 5.722 −1.2% (3/3 rounds)
Qwen3-30B-A3B-GPTQ-Int4 16 12.188 – 12.136 within noise (2/3 rounds faster)
TP4 setup

v0.33.0, 4× A100-PCIE-40GB, --tensor-parallel-size 4 --enable-expert-parallel, NVFP4 weights, MTP with 2 draft tokens, --compilation-config '{"mode": 0, "cudagraph_mode": "FULL_DECODE_ONLY"}', bf16 KV, a custom logits processor. Client: c concurrent streaming chat requests, 384 output tokens, model-default sampling; per-request decode = (tokens − 1) / (last − first token time), averaged. Prose / code / numbers prompts × c = 1, 2, 4, 8, 12, 16; two server starts per arm, A/B/B/A. With MTP, c = 1 varies by about ±10% from run to run through the draft acceptance rate.

Standalone CUDA reproduction (no vLLM) and commands

48 × (plain kernel → consumer kernel); only the consumer launch changes. A100, 21 rounds:

consumer launch graph replay, GPU µs eager, GPU µs
<<<>>> 128.7 239.8
cudaLaunchKernelEx, no attribute 128.7 239.3
cudaLaunchKernelEx + PDL attribute 181.8 239.8
PDL attribute on the first consumer only 130.3 239.4

Same harness, graph replay, other pre-Hopper GPUs (21 rounds, PDL attribute slower in every round):

GPU (arch) <<<>>> µs + PDL attribute per edge cudaLaunchKernelEx, no attribute
A100-PCIE-40GB (sm80) 128.7 +53.1 (+41%) 1.1 µs +0.03
RTX A4000 (sm86) 83.6 +29.9 (+36%) 0.62 µs +0.11
RTX 4090 (sm89) 79.0 +30.0 (+38%) 0.63 µs −0.10

cudaGraphGetEdges reports one programmatic edge per PDL launch on all three. The standalone harness and the op-chain script are in https://gist.github.com/srkim-business/ec5db57f7b78a9c9f06af65c891216c2.

# end to end (OLMoE-1B-7B-0924 and Qwen3-30B-A3B-GPTQ-Int4, batch 1 and 16)
vllm bench latency --model <model> --batch-size <1|16> --input-len 128 --output-len 256 \
  --num-iters-warmup 5 --num-iters 15 --seed 0 --max-model-len 1024 --gpu-memory-utilization 0.85
# GSM8K
vllm serve <model> --max-model-len 4096 --gpu-memory-utilization 0.85 --seed 0
python tests/evals/gsm8k/gsm8k_eval.py --port 8000 --seed 42 --num-questions 1319 --max-concurrency 64
# kernel tests
pytest tests/kernels/moe/test_fused_topk.py

Details

  • Root cause: on pre-SM90 the attribute does nothing for eager launches, but stream capture records it as a programmatic edge (cudaGraphGetEdges counts one per launch on sm80). Each edge adds about 1.1 µs of GPU time per graph replay on A100.
  • [Perf] Enable PDL for fused MoE top-k gating #59737 first had this host check (uncached); review asked to drop it for consistency with topk_softplus_sqrt / per_token_group_quant and to avoid cudaGetDevice + call_once per eager launch. The measurements above are the reason to bring it back. It is cached in a function-local static as in fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu (numAttrs = (sm_version >= 90) ? 1 : 0), so eager launches pay nothing per call. This also matches the scope in [Perf] Enable PDL for fused MoE top-k gating #59737's own description ("Enable PDL for the consumer on SM90+"). Python launch sites already gate PDL on current_platform.is_arch_support_pdl().
  • The static is initialized once per template instantiation, on the device the caller's DeviceGuard selects. It assumes all GPUs in a process share the major arch, which holds for vLLM workers (one GPU each); a mismatch would only affect speed, not results.
  • The 1-GPU gain (about 1%) matches the GPU-side edge cost. The TP4 + EP + MTP gain does not: at concurrency 16 a decode step goes from about 29.5 to 27.3 ms, while 48 edges × 1.1 µs is about 0.05 ms per forward; we have not located the rest. The TP4 A/B swaps only _moe_C (shipped → this diff built from the same source); on one GPU a no-change rebuild is within 0.8% (median) of the shipped build, so we attribute the gain to this diff.
  • Other C++ launchers still set the attribute on all archs, for example topk_softplus_sqrt_kernels.cu and per_token_group_quant.cu (the latter can run on sm89 for FP8). They are left out to keep this fix small enough to cherry-pick; a shared helper can follow if maintainers prefer.
  • Not a duplicate: no open issue or PR changes this launch attribute (searched topkGating PDL, topk_softmax PDL, 59737). [Perf][MoE] Select full-warp top-k rows with redux.sync and keep selections in registers (SM80+) #62912 (open) changes top-k selection in the same file and keeps config.numAttrs = 1 as context next to this change, so whichever lands second needs a small rebase; it does not touch the attribute.
  • Not covered: Hopper/Blackwell (unchanged code path, not measured), sm75 and older, end to end on GPUs other than A100, and a full main build (only _moe_C was rebuilt and overlaid on v0.33.0).
  • This is a performance regression against v0.32.0; we'd like it cherry-picked into a v0.33.x patch release if one is planned.
  • I used AI assistance (Claude) for the investigation, the code and the tests. I reviewed every changed line and ran the tests and GPU checks above.

Pull Request Checklist
  • I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

  • AI assistance was used during the creation of this PR.

  • Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

  • Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

  • Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

  • Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

On pre-SM90 GPUs the PDL launch attribute added in vllm-project#59737 is still
recorded as a programmatic edge in captured CUDA graphs and slows graph
replay. Set it only on SM90+ and cache the check in a function-local
static.

Fixes vllm-project#65055

Co-authored-by: Claude <noreply@anthropic.com>
Signed-off-by: Seungryong Kim <137975233+srkim-business@users.noreply.github.com>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@mergify mergify Bot added the bug Something isn't working label Nov 2, 2026
@srkim-business

Copy link
Copy Markdown
Contributor Author

/ci run

@github-actions

github-actions Bot commented Nov 2, 2026

Copy link
Copy Markdown

✅ Triggered Buildkite CI #99104 for commit 303fe0338295.

ℹ️ This PR is 30 commits behind upstream main. Merged upstream main into this branch as 303fe0338295 before starting CI. Pull the branch before pushing more changes.

@gau-nernst gau-nernst added the ready ONLY add when PR is ready to merge/full CI is needed label Nov 2, 2026
@github-actions

github-actions Bot commented Nov 2, 2026

Copy link
Copy Markdown

✅ @srkim-business, CI is now available for this PR.

  • /ci run starts upstream CI; /amd-ci run starts AMD CI only.
  • Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
  • /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
  • /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
  • /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

@BESzbing

BESzbing commented Nov 2, 2026

Copy link
Copy Markdown

Independent Ada (sm_89) A/B for this PR — the reporter's standalone effect was measured on a 4090, but there was no Ada end-to-end number yet, and upstream CI does not run consumer sm_89.

Setup. 1× RTX 4090 24 GB, driver 580.95.05, CUDA 13.0, torch 2.14.1+cu132. Two _moe_C builds from the same tree at 0e5b2ba118 + this diff: stock = master, fix = master + the three-line change (only difference is numAttrs). Same shape as the reporter's "vLLM op chain": per layer router GEMM (bf16 torch.mm) → torch.ops._moe_C.topk_softmax → elementwise consumer that feeds the next layer, all layers in one captured CUDA graph, replay() timed with CUDA events (21 replays per process, median). 6 rounds, order alternated A/B / B/A; Δ is the median of 6 paired differences.

shape (experts, top-k, hidden, layers) tokens stock fix Δ Δ% per edge
OLMoE-1B-7B (64, 8, 2048, 16) 1 189.4 179.2 −10.2 −5.4% 0.64 µs
16 225.3 215.0 −10.2 −4.5% 0.64 µs
48 234.5 223.7 −10.8 −4.6% 0.67 µs
Qwen3-30B-A3B (128, 8, 2048, 48) 1 575.0 542.7 −32.3 −5.6% 0.67 µs
16 686.1 652.3 −33.8 −4.9% 0.70 µs
48 710.7 677.9 −32.8 −4.6% 0.68 µs
Qwen3.8-Flash-Next (512, 10, 2560, 48) 1 815.1 777.7 −37.4 −4.6% 0.78 µs
16 919.6 882.7 −37.9 −4.1% 0.79 µs
48 985.1 945.2 −40.4 −4.1% 0.84 µs
  • All 54 paired deltas are negative; per-round spread is −9.2…−12.3 µs at 16 layers and −30.7…−43.0 µs at 48 layers.
  • Per-edge cost on sm_89 is 0.64–0.84 µs, i.e. the same 0.63 µs/edge the report has standalone on this GPU, so the mechanism is confirmed on Ada.
  • Anchor: the pre-[Perf] Enable PDL for fused MoE top-k gating #59737 shipped wheel's _moe_C (v0.32.0) on the same harness gives 184.3 / 218.1 / 225.3 µs for OLMoE at 1/16/48 tokens. fix lands at or slightly below the pre-regression level, so the change adds no measurable overhead on sm_89 either.
  • Correctness unchanged: topk_softmax ids identical to torch.topk(gating.softmax(-1)) and weights max-abs-err ≤ 6e-8 on both builds for all three shapes.
  • Eager could not be resolved with this harness — that path is host-bound here (3.4–5.2 ms vs 0.8–1.0 ms for the graph), so a ~30 µs GPU effect is buried in dispatch noise. Consistent with "eager is not affected", but not evidence for it.

Two questions while I had the tree open.

  1. Coverage. cudaLaunchAttributeProgrammaticStreamSerialization is set at 28 sites in csrc/; ten are still unconditional on pre-Hopper, and several of those run per layer on common models: fused_qknorm_rope_kernel.cu:594, quantization/w8a8/fp8/per_token_group_quant.cu:411,456,819, dsv3_fused_a_gemm.cu:694, moe/topk_softplus_sqrt_kernels.cu:212,692, fp32_router_gemm.cu:188, kimi_k3/fused_kda_decode_kernel.cu:759, kimi_k3/attn_res_kernel.cu:915, custom_all_gather_reduce_scatter.cu:63,155, beam/allreduce_norm.cu:450. At ~0.6 µs per edge per replay, each of those still costs what this PR removes for one kernel. Would it be better to land a shared helper (mirroring enable_pdl in moe/grouped_topk_kernels.cu:1303) and convert the hot per-layer ones in the same PR, or is a follow-up the plan?

  2. static const bool use_pdl. The cache is per template instantiation and captures the device active on first call. get_device_prop() is already cached per device, so dropping the static only costs a cudaGetDevice(), and it avoids the mixed-device footgun.

Happy to rerun this harness on any follow-up patch — the two builds and the script are reproducible here.

AI assistance was used to write the harness and this comment; I ran both builds and checked the numbers against the raw per-round files.

@mergify

mergify Bot commented Nov 2, 2026

Copy link
Copy Markdown
Contributor

Hi @srkim-business, 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.

@gau-nernst

Copy link
Copy Markdown
Contributor

Pre-commit failure will be fixed by #65080. Please only re-run CI after that PR is merged.

@srkim-business

Copy link
Copy Markdown
Contributor Author

/ci run

@github-actions

github-actions Bot commented Nov 2, 2026

Copy link
Copy Markdown

✅ Triggered Buildkite CI #99117 for commit 0390140c6b86.

@srkim-business

Copy link
Copy Markdown
Contributor Author

@BESzbing Thanks for the independent sm_89 A/B. The per-edge cost matches what I measured standalone on a 4090 (about 0.62 µs).

  1. Coverage: I'd like to keep this PR to the one launcher I measured end to end (A100, TP4 + EP MoE decode). A shared helper for the remaining sites sounds like the right shape, but it needs per-site measurements, so I'd treat it as separate work rather than extend this PR.

  2. static: this is deliberate. In [Perf] Enable PDL for fused MoE top-k gating #59737 the review asked to drop a per-launch host check because of the extra device query on every eager launch, so I kept the check out of the launch path. In normal vLLM deployments each worker process drives one GPU, so one process launching this kernel on GPUs of different architectures should not happen. If the reviewers prefer a per-device value without static, I'm happy to change it.

Written with AI assistance (Claude Code).

@srkim-business

Copy link
Copy Markdown
Contributor Author

/ci retry

@github-actions

github-actions Bot commented Nov 2, 2026

Copy link
Copy Markdown

✅ Queued 2 failed job(s) for retry in Buildkite CI #99117.

@gau-nernst
gau-nernst enabled auto-merge (squash) November 2, 2026 10:32
@Ondre119

Ondre119 commented Nov 2, 2026

Copy link
Copy Markdown

Would you be open to a separate benchmark follow-up based on your topk_chain_bench.py gist?

I checked the existing coverage: benchmark_fused_topk.py times standalone top-k in eager mode, while benchmark_moe.py already captures top-k plus expert computation in a CUDA graph. A smaller router GEMM -> top-k -> consumer chain would help isolate launch/graph overhead without the expert computation.

I'd like to adapt your existing harness, with attribution, into the repository benchmark rather than build another standalone harness. The scope would be eager and graph measurements of that chain, configurable shapes/dtypes/layer counts, and raw samples with build identifiers for paired comparisons. Correctness would reuse the existing fused-top-k tests.

Would extending benchmark_fused_topk.py be the right home for this, or would you prefer another existing benchmark? I'd keep this as separate benchmark work, leaving this PR focused on the measured launcher.

I haven't run GPU measurements yet; I'd validate the agreed scope on a pre-Hopper GPU before opening a PR.

@gau-nernst

Copy link
Copy Markdown
Contributor

/ci retry

@github-actions

github-actions Bot commented Nov 2, 2026

Copy link
Copy Markdown

✅ Triggered Buildkite CI #99186 for commit 220fa30cac3e, running 1 failed step(s) from Buildkite CI #99117.

@gau-nernst
gau-nernst merged commit 187f1a2 into vllm-project:main Nov 2, 2026
15 checks passed
@srkim-business

Copy link
Copy Markdown
Contributor Author

@Ondre119 Sure, feel free to adapt topk_chain_bench.py, and thanks for asking first. If it ends up close to the original, a Co-authored-by: Seungryong Kim <137975233+srkim-business@users.noreply.github.com> trailer would be appreciated. Where it should live is the maintainers' call. Extending benchmark_fused_topk.py sounds reasonable to me, since it already covers this op. Note that #62912 also changes that file (it adds --dtype), so it may be easier to start after that PR lands.

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

Labels

bug Something isn't working ready ONLY add when PR is ready to merge/full CI is needed

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Performance]: MoE top-k gating PDL launch (#59737) slows CUDA graph replay on pre-Hopper GPUs

5 participants