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, 15:15 UTC
[Bugfix][Kernel] fix bias adding in triton kernel implemented fused moe (#31676)
Signed-off-by: xuebwang-amd <xuebwang@amd.com>
xuebwang-amd committed
Jan 7, 2026 at 15:36 UTC
0dd5dee9b9bc88453f5f3eacfde751e6b9ba4871
1 file changed
+8
-5
vllm/model_executor/layers/fused_moe/fused_moe.py
+8
-5
@@ -518,11 +518,7 @@ def fused_moe_kernel(
518
# Advance the ptrs to the next K block.
519
a_ptrs += BLOCK_SIZE_K * stride_ak
520
b_ptrs += BLOCK_SIZE_K * stride_bk
521
- if HAS_BIAS:
522
- accumulator = accumulator + bias[None, :]
523
- if MUL_ROUTED_WEIGHT:
524
- moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
525
- accumulator = accumulator * moe_weight[:, None]
521
+
522
if use_int8_w8a16:
523
accumulator = (accumulator * b_scale).to(compute_type)
524
elif use_fp8_w8a8 or use_int8_w8a8:
@@ -533,6 +529,13 @@ def fused_moe_kernel(
529
else:
530
accumulator = accumulator.to(compute_type)
531
532
+ # Since bias is typically not quantized, it's added after dequantization.
533
+ if HAS_BIAS:
534
+ 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]
538
+
539
# -----------------------------------------------------------
540
# Write back the block of the output
541
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)