Repository navigation
Conversation
| ) | ||
|
|
||
| elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_FP8: |
There was a problem hiding this comment.
this is moved without changes
| w13_bias=w13_bias, | ||
| w2_bias=w2_bias, | ||
| _cache_permute_indices=_cache_permute_indices, | ||
| elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4: |
There was a problem hiding this comment.
this is moved without changes
|
This pull request has merge conflicts that must be resolved before it can be |
e28e449 to
2d7822b
Compare
BowenBao
changed the title
fxmarty-amd
left a comment
fxmarty-amd
left a comment
There was a problem hiding this comment.
Not sure about the interleave/contiguous terminology given the extra swizzling done in convert_weight_to_mxfp4_moe_kernel_format ?
Is the semantic meaning of W13Layout => "layout expected by convert_weight_to_mxfp4_moe_kernel_format prior to any extra swizzling/interleaving" ?
| is_gfx1250 = False | ||
| if current_platform.is_rocm(): | ||
| from vllm.platforms.rocm import on_gfx1250 | ||
|
|
||
| is_gfx1250 = on_gfx1250() |
| class W13Layout(Enum): | ||
| """Layout of the fused gate/up (w13) weight tensor.""" | ||
|
|
||
| CONTIGUOUS_W1W3 = "contiguous_w1w3" | ||
| CONTIGUOUS_W3W1 = "contiguous_w3w1" | ||
| INTERLEAVED_W1W3 = "interleaved_w1w3" | ||
| INTERLEAVED_W3W1 = "interleaved_w3w1" |
There was a problem hiding this comment.
I am not sure about the meaning of these given that there are still block_scale_interleave/nvfp4_block_scale_interleave/interleave_moe_weights_for_sm90_mixed_gemm/rocm_aiter_ops.shuffle_weight_a16w4/_swizzle_mxfp4/get_w2_permute_indices_with_cache calls as part of the huge if/elif/elif block in convert_weight_to_mxfp4_moe_kernel_format:
- https://github.com/BowenBao/vllm/blob/2d7822bf3d1c53144c411c3cacdc639d5b3eacd5/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py#L1249-L1251
- https://github.com/BowenBao/vllm/blob/2d7822bf3d1c53144c411c3cacdc639d5b3eacd5/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py#L1235-L1237
- https://github.com/BowenBao/vllm/blob/2d7822bf3d1c53144c411c3cacdc639d5b3eacd5/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py#L1058
- https://github.com/BowenBao/vllm/blob/2d7822bf3d1c53144c411c3cacdc639d5b3eacd5/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py#L1095-L1098
- https://github.com/BowenBao/vllm/blob/2d7822bf3d1c53144c411c3cacdc639d5b3eacd5/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py#L957-L963
There was a problem hiding this comment.
So eventually e.g. CONTIGUOUS_W1W3 could mean a different thing depending on the backend that is picked?
There was a problem hiding this comment.
Some of these are overlooked duplicated pre-processes. The rest (the aiter side links) are swizzling which is another weight pre-process but orthogonal to gate-up interleave.
There was a problem hiding this comment.
On a higher level, a more developer friendly and new kernel integration friendly approach would be to have a canonical format of weights, then let the selected backend to apply process_weight_after_loading-like function. It seems there shouldn't be a standard for this interleaving definition or how the weights are expected. New backend and implementation are practically unknown and can have different layout.
There was a problem hiding this comment.
makes sense. Relaxing the PR to mainly focus on clarifying the incoming weight layout.
|
This pull request has merge conflicts that must be resolved before it can be |
2d7822b to
3a24f68
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
|
This pull request has merge conflicts that must be resolved before it can be |
b4267d7 to
979b814
Compare
…ormat Add an is_w13_interleaved flag so GPT-OSS checkpoints (w13 gate/up rows interleaved) can share the generic MXFP4 conversion. Each backend branch converts w13 to the layout its kernel needs: - TRTLLM, FlashInfer CUTLASS, AITER CK W4A16 and AITER W4A4 de-interleave first and reuse the contiguous path. - TRITON, AITER Triton W4A16 and AITER W4A8 skip their interleave step when w13 is already interleaved. - Other backends consume either layout as-is. The AITER W4A4/W4A8 branches move here from the GPT-OSS converter (W4A8 now interleaves contiguous w13 itself), and OCP MX emulation reduces per-expert static input scales like the GPT-OSS converter did. Nothing calls the new flag yet. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Bowen Bao <bowenbao@amd.com>
Remove convert_gpt_oss_weight_to_mxfp4_moe_kernel_format. GptOssMxfp4MoEMethod and Quark OCP MX (for gpt_oss models) call convert_weight_to_mxfp4_moe_kernel_format(..., is_w13_interleaved=True); other Quark models now go through the contiguous path instead of the GPT-OSS one. The ROCm oracle test feeds interleaved w13 for SWIGLUOAI and checks it against a contiguous reference, and gains GELU, SWIGLUOAI_UNINTERLEAVE and emulation cases. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Bowen Bao <bowenbao@amd.com>
Remove select_deepseek_v4_mxfp4_moe_backend. Mxfp4MoEMethod now calls select_mxfp4_moe_backend(..., use_deepseek_v4_priority=True), which keeps the DeepSeek-V4 behavior: its priority list, explicit aliases trying every variant, and no user activation override. Other callers are unchanged. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Bowen Bao <bowenbao@amd.com>
979b814 to
11b62e2
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
The changes are easier to review by going over the 3 commits separately. Most changes are only moving code around, which is not as significant as the default github diff.
fxmarty-amd
left a comment
fxmarty-amd
left a comment
There was a problem hiding this comment.
Overall LGTM with the comments below, and assuming that:
- Only gpt-oss uses
GptOssMxfp4MoEMethod - All
convert_gpt_oss_weight_to_mxfp4_moe_kernel_formatprocessing code was duplicate and/or update to useconvert_weight_to_mxfp4_moe_kernel_format+is_w13_interleaved
(did not use agent review)
| if not use_deepseek_v4_priority: | ||
| # Select kernels in order of backend. | ||
| AVAILABLE_BACKENDS = _filter_by_activation( | ||
| _get_priority_backends_for_gpt_oss(), requested_activation_key |
There was a problem hiding this comment.
This name is misleading now?
gpt-oss / other models should have the same default priority?
| def select_mxfp4_moe_backend( | ||
| config: FusedMoEConfig, | ||
| activation_key: QuantKey | None = None, | ||
| *, |
| config: FusedMoEConfig, | ||
| activation_key: QuantKey | None = None, | ||
| *, | ||
| use_deepseek_v4_priority: bool = False, |
There was a problem hiding this comment.
Removing select_deepseek_v4_mxfp4_moe_backend seems unrelated to PR title (?), but having a single select_mxfp4_moe_backend that does not duplicate code seems like a good change to me, sounds good to have it here.
As per 11b62e2, this is only a refactor.
| "rtol": 0.3, | ||
| "percent": 0.95, | ||
| "requires_aiter": False, | ||
| "requires_gfx950": False, | ||
| "backend_override": "TRITON_UNFUSED", | ||
| }, | ||
| "TRITON_UNFUSED__SWIGLUOAI_UNINTERLEAVE": { | ||
| "activation": "SWIGLUOAI_UNINTERLEAVE", | ||
| "rtol": 0.5, | ||
| "percent": 0.8, |
There was a problem hiding this comment.
These tolerances are very loose.
We should make sure the tolerances in these tests are not just set loose just because AITER/vLLM is bugged, see:
- [MOE][Quantization] Extend correctness test of
AITER_MXFP4_BF16MOE backend and temporarily fix correctness issues with Swiglu + no bias code path #59766 - [CK] [bug fix] Fix wrong cktile 2stages
moe_gemm1_heuristic_dispatcherdispatch in case a template is not built, add template specialization for(Swiglu, HasBias=false, split_k=1)ROCm/aiter#6022 - [cktile BF16-MXFP4 MOE][bug] Fix wrong output with
model_dimnot multiple of 512 with swiglu + separated gate/up ROCm/aiter#6086
This can be done in followup PRs.
I am fine if these atols are enough to prove no-regression in this PR.
github-project-automation
Bot
moved this from To Triage
to Ready
in gpt-oss Issues & Enhancements
|
This pull request has merge conflicts that must be resolved before it can be |
Background
GPT-OSS checkpoints store fused gate/up weights in interleaved row order [g0, u0, g1, u1, ...], while standard MoE checkpoints (DeepSeek, Kimi-K3, and most of the ckpts from different quantizers etc.) store them contiguously [gate; up].
In the beginning when mxfp4 oracle was initially created, this interleaved weight format is overlooked. Some backends accepts interleaved format while others accepts contiguous format. They are all mixed together in the same
convert_weight_*andselect_mxfp4_moe_backendcall.Later when DSV4 was added, two parallel functions (convert_gpt_oss_weight_* and convert_weight_*) were created. The original
convert_weight_*was renamedconvert_gpt_oss_weight_*, despite it handles more models than just gpt-oss. Now several backends shared identical code between them, meaning they silently accepted either layout without converting, which is incorrect for backends that assume a specific layout.Audit DeepDive
This PR starts a complete, deepdive audit to the existing state of all the backends. NOTE: The audit is produced by Claude Code, and verified carefully by the author. There could be mistakes.
Terminology
[gate_0..gate_N; up_0..up_N](standard vLLM checkpoint layout)[up_0..up_N; gate_0..gate_N](FlashInfer CUTLASS convention, halves swapped)[gate_0, up_0, gate_1, up_1, ...](GPT-OSS checkpoint layout)[up_0, gate_0, up_1, gate_1, ...](TRTLLM convention, w1/w3 swapped within pairs)w13 Layout Expected by Backend x Activation
fp8_utils.py:143:hidden_size: tl.constexpr = N // 2;fp8_utils.py:165-178: loads gate frominput_ptr + offset, up frominput_ptr + hidden_size + offsetlayers/activation.py:37-38:gate = tl.load(x_row_ptr + offsets),up = tl.load(x_row_ptr + offsets + d)withd = n // 2at line 55activation_kernels.cu:483-484:gate_ptr = input + row * 2 * d,up_ptr = gate_ptr + dmxfp4.py:1310-1317: comment "TRTLLM kernel expects interleaved [w3_0, w1_0, ...]";torch.stack([w3_weight, w1_weight], dim=2).reshape()produces[up_0, gate_0, up_1, gate_1, ...]trtllm_mxfp4_moe.py:139-146: SITU support gated behindhas_flashinfer_situ_activation()flashinfer_utils.py:53-56:swap_w13_to_w31reshapes to(-1, 2, half, cols)and.flip(dims=[1]), swapping gate and up halvesmxfp4.py:1535-1539: comment "Standard checkpoints store [w1; w3], FlashInfer CUTLASS consumes [w3; w1]"flashinfer_cutlass_moe.py:191-198:_supports_activationlists SILU, GELU_TANH, RELU2_NO_MUL, SWIGLUOAI, SWIGLUOAI_UNINTERLEAVE_swiglu.py:52:gelu, linear = tl.split(tl.reshape(input, (input.shape[0], input.shape[1] // 2, 2)))— reshapes to(M, N//2, 2)then splits, pairing consecutive elementsgpt_oss_triton_kernels_moe.py:967-968:_supports_activationreturnsactivation == MoEActivation.SWIGLUOAImxfp4.py:1477-1485:shuffle_weightconverts contiguous[w1;w3]to interleaved viatorch.stack((first, second), dim=-1).reshape(shape)activation_kernels.cu:410: comment// Interleaved gate/up: input has [gate0, up0, gate1, up1, ...]; vectorized path at line 448:ACT_FN(vp[2*j], vp[2*j+1], ...)fused_moe/activation.py:262-268: dispatches tosilu_and_mul_with_clamp→activation_kernels.cu:102-108:x_ptr = input + token_idx * 2 * d,y_ptr = x_ptr + dfused_moe/activation.py:230-236: dispatches tosilu_and_muletc. →act_and_mul_kernelatactivation_kernels.cu:104:input // [..., 2, d]midpoint splitgpt_oss_triton_kernels_moe.py:1057-1064:_supports_activationlists SILU, GELU, SWIGLUOAI, SWIGLUSTEP, SWIGLUOAI_UNINTERLEAVErocm_aiter_moe.py:384-385:gate_mode = GateMode.INTERLEAVE.valuewhenuse_mxfp4_w4a16=Truemxfp4.py:1429-1434: weight shuffle withis_guinterleave=True, gate_up=Truerocm_aiter_moe.py:374-383: SITU defaults toGateMode.SEPARATED.value(unless env-varVLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1enables interleaved a8w4 path)mxfp4.py:1429-1434always usesis_guinterleave=Trueregardless of activation —is_guinterleavecontrols CK tile layout, not logical gate/up ordering. Runtimegate_modecontrols which kernel variant is dispatched.mxfp4.py:1689-1691:_interleave_w13(w13_weight, w13_weight_scale, w13_bias)converts contiguous → interleaved before_swizzle_mxfp4moe_op_gemm_a8w4.py:719-722:a_gelu = a[..., ::2],a_linear = a[..., 1::2]inswiglu_torchaiter_mxfp4_w4a8_moe.py:305: only supportsMoEActivation.SWIGLUOAIrocm_aiter_moe.py:374: defaultgate_mode = ""(not set for W4A4 path);_aiter_ops.py:228:gate_mode: str = ""meansgate_modenot passed to kernelmxfp4.py:1624-1662: no_interleave_w13call; onlye8m0_shuffle+shuffle_weightsfused_moe/activation.py:260-261:torch.ops._C.swigluoai_and_mul(output, input)→activation_kernels.cu:410-466: interleavedgate = in_ptr[2*idx],up = in_ptr[2*idx+1]fused_moe/activation.py:230-268: dispatches tosilu_and_muletc. →activation_kernels.cu:102-108:x_ptr = input + token_idx * 2 * d,y_ptr = x_ptr + dmarlin_moe.py:170-177: callsapply_moe_activationafter GEMMxpu_moe.py:170:activation=activation.valuepassed to externalXpuFusedMoefromvllm_xpu_kernels; cannot trace kernel sourcemoe.cpp:454-455:nb_upper = nb, nb_lower = nb + NB—NBis block count in first half,nb + NBindexes second half (midpoint split)silu_and_mulandswigluactivation paths use identicalnb_upper/nb_loweraddressing. SWIGLUOAI differs only in activation math (clamped sigmoid), not in layout.apply_moe_activation)ocp_mx_emulation_moe.py:38: inheritsTritonExperts;ocp_mx_emulation_moe.py:188: dequantizes then callssuper().apply()→apply_moe_activationapply_moe_activation)fused_humming_moe.py:302-305:_supports_activationdelegates toapply_moe_activation_supported(activation);fused_humming_moe.py:553-574:apply_activation()callsself.activation()→apply_moe_activationThe AITER backends are in a better shape after this change. Other backends are updated with best effort and would need careful community review.
Next steps
convert_gpt_oss_weight_*andselect_gpt_oss_mxfp4_moe_backend. Extendselect_mxfp4_moe_backendto differentiate source gate-up layout.Local tests