执行摘要
- 一句话:FP4 MoE 对齐降到 128,权重减约 40 GiB
- 推荐动作:值得精读:对于在 AMD 平台上部署 DeepSeek-V4 的团队,值得理解这个 padding 对齐与 kernel tile 大小的关系;对于一般读者,这是一个"一行常量修改 + 外部依赖联动"的典型案例,可学习如何在量化路径中通过 kernel 能力约束来压缩权重内存。关注点:aiter 版本升级的落地状态;若后续要推广到其他模型,注意不同 MoE kernel 的 tile_k 支持差异;建议补充参数化单测。
功能与动机
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 显存。
实现拆解
变更入口:唯一改动文件是 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)会进入。
- 对齐常量调整:将 fp4_k_align 从 256 改为 128,并更新注释:DeepSeek-V4 MoE 由 FlyDSL kernel 实现,支持 tile_k=128,因此中间维度只需填充到 128 的倍数;shuffle_scale 对非 256 形状的支持自 aiter PR#4130 起可用。此改动同步影响 padded_inter 的计算结果。
- 填充逻辑复用:后续逻辑不变——仍按 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 分支不再触发。
- 配套依赖:本 PR 单独合并不完整,必须同步升级 aiter 到包含 PR#4130 的版本,否则非 256 形状下 shuffle_scale 会出错;评论中 bingxche 还发现 AITER scout 存在疑似回归 commit,需先排查。
- 验证配套:未新增单测;作者提供 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(模块 量化层;类别 source;类型 core-logic;符号 process_weights_after_loading_block_quant): 唯一改动文件:在 process_weights_after_loading_block_quant() 中将 fp4_k_align 从 256 改为 128,直接决定 DeepSeek-V4 FP4 expert 权重是否被填充到 512,是内存收益和性能权衡的核心。
关键符号:process_weights_after_loading_block_quant
关键源码片段
python/sglang/srt/layers/quantization/fp8.py
唯一改动文件:在 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 分支。
评论区精华
评论区围绕两个问题展开:
风险与影响
- 风险:外部依赖风险:本 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 回归待排查
关联脉络
- PR #33090 [AMD][Fix] Restore aiter-padded MoE weight dims for serialized checkpoints: 同为 AMD aiter MoE 权重维度对齐问题,一个修序列化 checkpoint 的维度恢复,一个修 padding 粒度。
- PR #32315 [AMD] Speed up DSV4 MoE weight loading from mmap views: 同为 AMD 平台 DeepSeek-V4 MoE 权重加载优化,改善加载内存与速度。
- PR #31221 [AMD] Derive AITER verify tokens-per-req from input shape: 同为 AMD AITER 后端修复,共享 aiter 依赖演进背景。
参与讨论