A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 06:30 UTC
[Kernel][MoE] fix computation order of MoE weight multiplication and improve flow (#31962)
Signed-off-by: xuebwang-amd <xuebwang@amd.com>
xuebwang-amd committed
Jan 13, 2026 at 06:17 UTC
629584bfc9fc5c88c24f6672ca536fae7f75b835
1 file changed
+24
-9
vllm/model_executor/layers/fused_moe/fused_moe.py
+24
-9
@@ -531,22 +531,37 @@ def fused_moe_kernel(
531
a_ptrs += BLOCK_SIZE_K * stride_ak
532
b_ptrs += BLOCK_SIZE_K * stride_bk
533
534
- # Router weight multiplication MUST happen in float32 before precision
535
- # conversion for numerical stability (especially critical on ROCm).
536
- if MUL_ROUTED_WEIGHT:
537
- moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
538
- accumulator = accumulator * moe_weight[:, None]
539
-
534
+ # Dequantization for supported quantization schemes:
535
+ # - int8_w8a16
536
+ # - fp8_w8a8
537
+ # - int8_w8a8
538
+ # Accumulator and scalings are in float32 to preserve numerical accuracy.
539
if use_int8_w8a16:
540
accumulator = accumulator * b_scale
541
elif (use_fp8_w8a8 or use_int8_w8a8) and not (group_k > 0 and group_n > 0):
542
accumulator = accumulator * a_scale * b_scale
543
545
- # Bias is added AFTER dequantization since bias is typically stored in
546
- # the output dtype and should not be scaled by quantization factors.
544
+ # Bias addition:
545
+ # Bias must be applied after dequantization:
546
+ # - Since bias is typically not quantized
547
+ # - Bias should not be scaled by quantization factors
548
if HAS_BIAS:
548
- accumulator = accumulator + bias[None, :]
549
+ accumulator += bias[None, :]
550
+
551
+ # Router (MoE) weight multiplication:
552
+ # This multiplication MUST be performed in float32 before any precision
553
+ # conversion to ensure numerical stability, which is especially critical
554
+ # on ROCm platforms.
555
+ if MUL_ROUTED_WEIGHT:
556
+ moe_weight = tl.load(
557
+ topk_weights_ptr + offs_token,
558
+ mask=token_mask,
559
+ other=0,
560
+ )
561
+ accumulator *= moe_weight[:, None]
562
563
+ # Final precision conversion:
564
+ # Cast once at the end to the desired compute/output dtype.
565
accumulator = accumulator.to(compute_type)
566
567
# -----------------------------------------------------------