From 7e509f690e90a985195d0741533bf404cc8b0d7d Mon Sep 17 00:00:00 2001 From: Thomas Wang Date: Sun, 2 Aug 2026 14:42:21 +0800 Subject: [PATCH] [AMD] Fix DeepSeek-V4 FP4 MoE expert memory bloat (#31450) AITER specific --- python/sglang/srt/layers/quantization/fp8.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 26fc7b94a..7fa220b32 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -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