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, 09:00 UTC
fused_moe_kernel - cast accumulator after applying router weights (#32002)
Signed-off-by: gnovack <gnovack@amazon.com>
gnovack committed
Jan 10, 2026 at 12:36 UTC
d1fd802fa3df9641e9b6cc838f055f81c7619dc3
1 file changed
+5
-8
vllm/model_executor/layers/fused_moe/fused_moe.py
+5
-8
@@ -539,20 +539,17 @@ def fused_moe_kernel(
539
accumulator = accumulator * moe_weight[:, None]
540
541
if use_int8_w8a16:
542
- accumulator = (accumulator * b_scale).to(compute_type)
543
- elif use_fp8_w8a8 or use_int8_w8a8:
544
- if group_k > 0 and group_n > 0:
545
- accumulator = accumulator.to(compute_type)
546
- else:
547
- accumulator = (accumulator * a_scale * b_scale).to(compute_type)
548
- else:
549
- accumulator = accumulator.to(compute_type)
542
+ accumulator = accumulator * b_scale
543
+ elif (use_fp8_w8a8 or use_int8_w8a8) and not (group_k > 0 and group_n > 0):
544
+ accumulator = accumulator * a_scale * b_scale
545
546
# Bias is added AFTER dequantization since bias is typically stored in
547
# the output dtype and should not be scaled by quantization factors.
548
if HAS_BIAS:
549
accumulator = accumulator + bias[None, :]
550
551
+ accumulator = accumulator.to(compute_type)
552
+
553
# -----------------------------------------------------------
554
# Write back the block of the output
555
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)