Prhub

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

原始 PR 作者 1am9trash 合并时间 2026-08-02 14:42 文件变更 1 提交数 1 评论 7 代码增减 +5 / -4

执行摘要

FP4 MoE 对齐降到 128,权重减约 40 GiB

PR body 明确指出:AMD M355 上每层 MoE 权重占 2.105 GB,而同层在 NVIDIA B200 上仅 1.582 GB,差距约 25%。根因是 AMD 路径将 w13 专家权重按 256 的倍数填充,DeepSeek-V4-Pro 在 TP8 下每分区中间维度 3072/8=384 被抬到 512,使 routed-expert 权重膨胀约 1/3。作者目标是消除冗余填充,使权重占用接近 NVIDIA 水平,并为 AMD 部署节省约 40 GiB 显存。

值得精读:对于在 AMD 平台上部署 DeepSeek-V4 的团队,值得理解这个 padding 对齐与 kernel tile 大小的关系;对于一般读者,这是一个"一行常量修改 + 外部依赖联动"的典型案例,可学习如何在量化路径中通过 kernel 能力约束来压缩权重内存。关注点:aiter 版本升级的落地状态;若后续要推广到其他模型,注意不同 MoE kernel 的 tile_k 支持差异;建议补充参数化单测。

讨论亮点

评论区围绕两个问题展开:

  • aiter 依赖:HaiShaw 要求 @bingxche 确认 aiter 升级,并请求 aiter scout 落地 PR#4130 后的 commit;bingxche 反馈"发现一个疑似回归 commit 影响 DSv32 和 DSv4-pro,需先解决再升级"。HaiShaw 最后追问 "aiter scout green?",该闭环未在评论中明确回复。
  • 性能权衡:1am9trash 实测:send_one 无差异;小并发(< cc16)无差异;cc16 约 1% 下降。kernel 剖析显示主要差异在 moe1:w/ padding(inter=512)的 mfma_moe1_...t32x128x256 平均 35.15 us,w/o padding(inter=384)的 mfma_moe1..._t32x64x256 平均 39.31 us;moe2 基本持平。作者随后在 aiter 侧为 dim=384 补充调优配置。
  • 审核结论:kkHuang-amd 与 HaiShaw 均批准(LGTM),无 inline review 评论;gemini-code-assist bot 无反馈。

实现拆解

变更入口:唯一改动文件是 python/sglang/srt/layers/quantization/fp8.py,函数为 process_weights_after_loading_block_quant(),属于 FP4 专家权重加载后的后处理路径,由 _use_aiter and self.is_fp4_expert 守卫,仅 DeepSeek-V4(经 model_config.py 的 is_deepseek_v4 + SGLANG_DSV4_FP4_EXPERTS)会进入。

  1. 对齐常量调整:将 fp4_k_align 从 256 改为 128,并更新注释:DeepSeek-V4 MoE 由 FlyDSL kernel 实现,支持 tile_k=128,因此中间维度只需填充到 128 的倍数;shuffle_scale 对非 256 形状的支持自 aiter PR#4130 起可用。此改动同步影响 padded_inter 的计算结果。
  2. 填充逻辑复用:后续逻辑不变——仍按 padded_inter 与 inter_per_part 的差值记录 layer.intermediate_pad、layer.hidden_pad,并在需要时把 w13_weight 从 (E, 2inter, K_packed) 零填充到 (E, 2padded, K_packed),w2_weight 同理。TP8 下 384 已为 128 的倍数,整条 padding 分支不再触发。
  3. 配套依赖:本 PR 单独合并不完整,必须同步升级 aiter 到包含 PR#4130 的版本,否则非 256 形状下 shuffle_scale 会出错;评论中 bingxche 还发现 AITER scout 存在疑似回归 commit,需先排查。
  4. 验证配套:未新增单测;作者提供 GSM8k 精度(tp8 0.940 / tp8+dp8 0.942)与 8k1k 压测(权重加载后显存 159.07 → 112.36 GB,ITL/TTFT/TTT 基本持平)作为验证;aiter 侧额外为 dim=384 补充了 dsv4_fp8fp4_tuned_fmoe.csv 调优。
文件 模块 状态 重要度
python/sglang/srt/layers/quantization/fp8.py 量化层 modified 5.88

关键符号

process_weights_after_loading_block_quant

关键源码片段

python/sglang/srt/layers/quantization/fp8.py core-logic

唯一改动文件:在 process_weights_after_loading_block_quant() 中将 fp4_k_align 从 256 改为 128,直接决定 DeepSeek-V4 FP4 expert 权重是否被填充到 512,是内存收益和性能权衡的核心。

# python/sglang/srt/layers/quantization/fp8.py
def process_weights_after_loading_block_quant(self, layer: Module) -> None:
    # AMD FP4 experts: 走 aiter 原生 MXFP4 MoE 路径
    if _use_aiter and self.is_fp4_expert:
        # DeepSeek-V4 MoE 使用 FlyDSL kernel,支持 tile_k=128,
        # 因此中间维度只需对齐到 128 的倍数。此前按 256 对齐,
        # 导致 TP8 下 384 -> 512 的 25% 权重膨胀(PR#31450 修复)。
        # 注意:shuffle_scale 支持非 256 形状需要 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 # 每个 partition 的中间维度
        padded_inter = (
            (inter_per_part + fp4_k_align - 1) // fp4_k_align * fp4_k_align
        )
​
        # 记录 padding 量,aiter fused_moe 依赖 intermediate_pad 得知真实中间维度
        layer.intermediate_pad = padded_inter - inter_per_part
        layer.hidden_pad = 0
​
        if padded_inter != inter_per_part:
            pad_amount = padded_inter - inter_per_part
            fp4_block_k = 32
​
            # 将 w13_weight 从 (E, 2*inter, K_packed) 扩为 (E, 2*padded, K_packed)
            old_w13 = layer.w13_weight.data
            new_w13 = torch.zeros(
                E,
                2 * padded_inter,
                w13_K_packed,
                dtype=old_w13.dtype,
                device=old_w13.device,
            )
            # 前半段(gate/up 输出)拷贝到开头
            new_w13[:, :inter_per_part, :] = old_w13[:, :inter_per_part, :]
            # 后半段(down 输出)偏移到 padded 位置,保持与 kernel 期望布局一致
            new_w13[:, padded_inter : padded_inter + inter_per_part, :] = old_w13[
                :, inter_per_part:, :
            ]
            layer.w13_weight = torch.nn.Parameter(new_w13, requires_grad=False)
​
            # w2_weight 也按 padded_inter 做 packed 维度填充,逻辑与 w13 对称

注意:改动前后唯一实质差异是 fp4_k_align 的取值与注释;填充、intermediate_pad 记录等逻辑完全复用,因此对 384 这类 128 倍数维度会直接跳过 padding 分支。

评论区精华

aiter 版本升级与回归排查 other

HaiShaw 要求在合并前确认 aiter 升级,并请 bingxche 帮助 aiter scout 落地 PR#4130 之后的 commit;bingxche 反馈在最新 AITER scout 中发现疑似回归 commit 影响 DSv32 和 DSv4-pro,需先解决再升级。

结论:PR 合并前需完成 aiter 升级;AITER scout 回归需先排查,最终状态未在评论中闭环。 · 待处理

去 padding 后的性能回归评估 性能

1am9trash 实测:send_one 与小并发无明显差异,cc16 约 1% 性能下降;kernel 剖析显示 w/o padding 的 moe1 kernel 平均 39.31 us,高于 w/ padding 的 35.15 us,moe2 几乎持平。

结论:微小性能回退可接受,且 aiter 侧针对 dim=384 补充了调优配置。 · 已解决

风险与影响

外部依赖风险:本 PR 的 shuffle_scale 非 256 形状依赖 aiter PR#4130;若部署环境 aiter 版本未跟上,FP4 MoE 计算可能出错。且 AITER scout 中已知有疑似回归 commit,升级本身存在不确定性。性能风险:cc16 大并发下约 1% 吞吐下降,来自 moe1 kernel 未对齐时的耗时增加;低并发无差异。回归风险:改动仅触及 is_fp4_expert 路径(DeepSeek-V4 专属),其他模型不受影响;但非 256 对齐形状在其他 aiter MoE kernel(非 FlyDSL)上的支持未在此 PR 中验证。测试覆盖缺失:未附带单元测试,依赖手工 GSM8k 与压测;建议以后补充参数化单测,断言 384 这类输入不再触发 padding。

对用户/系统:AMD(M355)上部署 DeepSeek-V4-Pro(FP4 expert)的用户显存占用显著下降,权重加载后显存从 159.07 GB 降至 112.36 GB,可释放更多空间给 KV cache 等;TP8 下模型权重减少约 40 GiB。对团队:需要与 aiter 团队的版本联动(cherry-pick PR#4130),对 AMD 平台 CI 和依赖管理有影响。影响程度:仅 AMD + DeepSeek-V4 + FP4 expert 路径,其他模型/平台无影响;改动面极小(1 文件 9 行),但内存收益大。

依赖 aiter PR#4130 升级 cc16 下约 1% 性能下降 无单测配套 AITER scout 回归待排查

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论