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, 12:45 UTC
[ROCm][LoRA] Fix MoE accuracy regression by preserving float32 router weight scaling (#31931)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jan 7, 2026 at 22:17 UTC
c4041f37a4a4b70c8821c35df9c79298ee5fcb7b
1 file changed
+8
-4
vllm/model_executor/layers/fused_moe/fused_moe.py
+8
-4
@@ -519,6 +519,12 @@ def fused_moe_kernel(
519
a_ptrs += BLOCK_SIZE_K * stride_ak
520
b_ptrs += BLOCK_SIZE_K * stride_bk
521
522
+ # Router weight multiplication MUST happen in float32 before precision
523
+ # conversion for numerical stability (especially critical on ROCm).
524
+ if MUL_ROUTED_WEIGHT:
525
+ moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
526
+ accumulator = accumulator * moe_weight[:, None]
527
+
528
if use_int8_w8a16:
529
accumulator = (accumulator * b_scale).to(compute_type)
530
elif use_fp8_w8a8 or use_int8_w8a8:
@@ -529,12 +535,10 @@ def fused_moe_kernel(
535
else:
536
accumulator = accumulator.to(compute_type)
537
532
- # Since bias is typically not quantized, it's added after dequantization.
538
+ # Bias is added AFTER dequantization since bias is typically stored in
539
+ # the output dtype and should not be scaled by quantization factors.
540
if HAS_BIAS:
541
accumulator = accumulator + bias[None, :]
535
- if MUL_ROUTED_WEIGHT:
536
- moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
537
- accumulator = accumulator * moe_weight[:, None]
542
543
# -----------------------------------------------------------
544
# Write back the block of the output