Repository navigation
[Bugfix][Kernel] Only set PDL launch attr for MoE top-k gating on SM90+ - #65059
Conversation
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>
|
/ci run |
|
✅ Triggered Buildkite CI #99104 for commit ℹ️ This PR is 30 commits behind upstream |
|
✅ @srkim-business, CI is now available for this PR.
|
|
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
Two questions while I had the tree open.
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. |
|
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-filesThen, commit the changes and push to your branch. For future commits, |
|
Pre-commit failure will be fixed by #65080. Please only re-run CI after that PR is merged. |
|
/ci run |
|
✅ Triggered Buildkite CI #99117 for commit |
|
@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).
Written with AI assistance (Claude Code). |
|
/ci retry |
|
✅ Queued 2 failed job(s) for retry in Buildkite CI #99117. |
|
Would you be open to a separate benchmark follow-up based on your I checked the existing coverage: 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 I haven't run GPU measurements yet; I'd validate the agreed scope on a pre-Hopper GPU before opening a PR. |
|
/ci retry |
|
✅ Triggered Buildkite CI #99186 for commit |
|
@Ondre119 Sure, feel free to adapt |
Overview
Fixes #65055. Since #59737,
topkGatingis 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
cudaLaunchKernelExmatches<<<>>>.topk_softmaxchain: 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.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.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_Cas shipped.rebuilt= the same source built with the same toolchain, no change.fix=rebuilt+ this diff. Builds are_moe_Conly, 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.pyon A100: base 864 passed, fix 864 passed.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%.tests/kernels/moeoncsrc/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 intests/.Op chain, CUDA graph, µs per replay (median of 6 rounds, A/B/B/A order; Δ = median paired
fix − rebuilt)End to end,
vllm bench latency(input 128, output 256, 3 rounds, ms per step)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:
<<<>>>cudaLaunchKernelEx, no attributecudaLaunchKernelEx+ PDL attributeSame harness, graph replay, other pre-Hopper GPUs (21 rounds, PDL attribute slower in every round):
<<<>>>µscudaLaunchKernelEx, no attributecudaGraphGetEdgesreports 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.Details
cudaGraphGetEdgescounts one per launch on sm80). Each edge adds about 1.1 µs of GPU time per graph replay on A100.topk_softplus_sqrt/per_token_group_quantand to avoidcudaGetDevice+call_onceper eager launch. The measurements above are the reason to bring it back. It is cached in a function-local static as infused_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 oncurrent_platform.is_arch_support_pdl().DeviceGuardselects. 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._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.topk_softplus_sqrt_kernels.cuandper_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.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 keepsconfig.numAttrs = 1as context next to this change, so whichever lands second needs a small rebase; it does not touch the attribute._moe_Cwas rebuilt and overlaid on v0.33.0).Pull Request Checklist
I used vLLM's
/pr-checklistskill. (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.