[AMD] Fix DeepSeek-V4 FP4 MoE expert memory bloat (#31450)

AITER specific
This commit is contained in:
Thomas Wang
2026-08-01 23:42:21 -07:00
committed by GitHub
parent 056474cdb0
commit 7e509f690e
+5 -4
View File
@@ -1405,10 +1405,11 @@ class Fp8MoEMethod(FusedMoEMethodBase):
gu_intv = envs.SGLANG_USE_AITER_MOE_GU_ITLV.get()
fp4_weight_dtype = _require_fp4_dtype()
# CK FP4 MoE kernel requires K_packed divisible by 128
# (i.e., K_logical divisible by 256).
# Pad intermediate_size_per_partition if needed.
fp4_k_align = 256
# DeepSeek V4 MoE is implemented by the FlyDSL kernel, which supports
# tile_k=128, so we only need to pad dim to 128. This lets DeepSeek-V4-Pro
# at TP8 skip padding 384 -> 512, reducing routed-expert memory by ~25%.
# shuffle_scale also supports non-256 shapes since aiter PR#4130.
fp4_k_align = 128
E, w13_N, w13_K_packed = layer.w13_weight.shape
_, w2_N, w2_K_packed = layer.w2_weight.shape
inter_per_part = w13_N // 2