Repository navigation
[Bugfix][MoE] Size the permute scratch from its inputs - #60447
Merged
Merged
Conversation
The Humming grouped and CUTLASS experts sized MoEPermuteScratch from FusedMoEConfig, predicting how many rows the all2all backend would hand them. That prediction missed PCP, DeepEP v2's power-of-two receive buckets, and (for CUTLASS) DP entirely, and it keyed the hidden buffer on the unquantized dtype. The scratch now holds only index metadata and grows with the inputs it sees until the workspace manager is locked, like the shared workspace. Permuted hidden states come from the caller's workspace; Humming reuses its idle w13 input buffer on the prequantized path. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: mgoin <mgoin64@gmail.com>
mergify
Bot
added
performance
mgoin
marked this pull request as ready for review
October 7, 2026 17:52
mgoin
requested review from
AndreasKaratzas,
WoosukKwon,
njhill,
pavanimajety,
tlrmchlsmth,
wzhao18,
yewentao256 and
zyongye
as code owners
October 7, 2026 17:52
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Member
Author
|
/ci run |
|
✅ Triggered Buildkite CI #93385 for commit |
Contributor
CI selector (shadow): 11 test steps (22 jobs) instead of 84 (109 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (11)
Would skip (today's rules run them) (77)
Would add (today's rules do not run them) (4)
AMD mirrors: would skip (85)
AMD mirrors: would add (0)none 7 changed files · base |
This was referenced Oct 8, 2026
Contributor
|
Thanks @mgoin for the general solution. I validated this PR with DeepSeek-V4.1-Flash PCP on one 8x H20 (SM90) node and it looks good. The test result can be found in #59857. Since this PR is a more general solution, I have closed #60423. Btw, could you take a look at #59857 when you have time? Any feedback or review comments would be greatly appreciated! |
6 tasks done
A weight reload rebuilds the experts after the workspace is locked, so a scratch owned by the experts instance came back empty and could not grow. Use the workspace manager's shared scratch, as Humming does. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: mgoin <mgoin64@gmail.com>
Member
Author
|
/ci run |
|
✅ Triggered Buildkite CI #93606 for commit |
bnellnm
approved these changes
Oct 8, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
Stop predicting the MoE permute scratch capacity from
FusedMoEConfig; size it from the inputs it actually receives, growing until the workspace manager is locked. This is intended as the root fix behind #55105 (DeepEP v2 + sequence parallelism) and #60423 (PCP), which each patch a different term of the same capacity formula for theassert n_token <= self.max_num_tokensfailure in the Humming grouped path.Claims
max_num_tokens * dp_size * experts_per_token, so DP, sequence parallelism, PCP, and DeepEP v2's power-of-two receive buckets cannot make the scratch too small.hidden_states.dtype == self.hidden_dtypemismatch on the prequantized (block-FP8 activation) path.CutlassExpertsFp8and the CUTLASS W4A8 experts get the same behavior. They sized the scratch withmax_num_tokensand nodp_sizefactor.(num_experts, num_local_experts, device)instead of one per(capacity, topk, hidden size, dtype).Validation
pytest tests/kernels/moe/test_moe_permute_unpermute.py tests/quantization/test_humming_mxfp4_block_fp8.py tests/kernels/moe/test_moe.py -k "permute or humming or cutlass or block_fp8"(1x B300)tests/kernels/moe/test_cutlass_moe.pyand-k "permute or humming or cutlass")nvidia/Qwen3-30B-A3B-NVFP4,--moe-backend humming,VLLM_HUMMING_MOE_GEMM_TYPE=grouped_contiguous, 1x B300, eager (existingconfigs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml)allgather_reducescatter, CUDA graphs,--max-num-batched-tokens 3000pytest tests/model_executor/model_loader/test_reload.py -k "test_reload_weights and FP8_DYNAMIC"(1x B300,VLLM_CUTLASSFP8 backend, in-place weight reload)mypy-3.12Not validated:
HummingGroupedExperts.applyis mock-based. It is guarded by dtype and shape assertions.Details
Root cause.
workspace13/workspace2are sized from the actualMand grow untillock_workspace().MoEPermuteScratchwas instead sized up front from config, which requires knowing how many rows every prepare/finalize backend can return in every topology. That bound is a property of the dispatcher, not of the experts, and it has been wrong for PCP, for DeepEP v2 with non-power-of-two capacities, and for CUTLASS with DP.Change.
MoEPermuteScratchholds only index metadata and gainsreserve(expanded_rows). It reallocates while the workspace manager is unlocked and raisesScratch growth is not allowed after lockingafterwards, matching the shared workspace's contract.workspace2slice. Humming's prequantized path now writes intoquanted_gate_up_input, which is allocated with the right shape and dtype and is otherwise unused when w13 quantization is skipped, so this costs no extra memory.get_simultaneous: a secondget_simultaneouscall insideapply()would alias the liveworkspace13/workspace2, andtoken_expert_indicesis anarangethat would otherwise be refilled every call. They are small (seven int32 arrays over the expanded rows plus the sort workspace).Tradeoffs and limitations.
workspace13.moe_permute(scratch=...)withoutpermuted_hidden_statesnow allocates that tensor per call. Both in-tree callers pass a buffer;benchmark_moe_permute_unpermute.py's scratch numbers are not comparable to earlier ones.get_persistent_resource's docstring is loosened from "must not resize" to "must not resize once the manager is locked".test_reload_weightsin CI build 93385).Relation to open PRs. Not a duplicate: #55105 and #60423 extend the capacity formula for one topology each and conflict with each other; this removes the formula. The DeepEP v2 handle-capacity change in #55105 (
next_power_of_2(ceil(max_num_tokens / sp_size))) is independent and still useful. Searches forMoEPermuteScratchandpermute scratch workspacefound no other open PR.AI assistance was used to analyze the capacity mismatch, implement the change, and run the validation above.
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.
🤖 Generated with Claude Code