执行摘要
- 一句话:新增 MXFP8 MoRI dispatch,AMD DSV4 吞吐提升约 6-10%
- 推荐动作:值得精读。核心看点有三:一是“发送侧量化、按 live token 而非填充缓冲计数”的性能洞察,直击 AITER 量化网格按 padded 行数分配的开销来源;二是用 scale dtype(fp32 vs e8m0)而非 dispatch 枚举来区分输入格式,解耦了调度器与 runner 的枚举依赖;三是用启动期 inspect 签名探测 + bf16 回退管理跨仓库依赖,把“静默出错”的字节布局问题转化为显式警告。配套的契约测试思路(不依赖 GPU 却能钉住最容易静默回归的布局常量)也值得借鉴。
功能与动机
DeepSeek-V4 的 MoE 使用 per_1x32(MXFP4 权重)量化,AITER 需要携带 group-32 e8m0 微缩放的 fp8 激活,而 MoRI 现有三种 dispatch dtype 都无法产出该格式:bf16 无缩放(接收侧必须量化)、fp8 是 group-128 fp32 缩放(group 大小错误触发 fp8->bf16 upscale 往返)、fp4 缩放正确但 payload 需 upscale_mxfp4。bf16 只是因为无需 upscale kernel 才成为默认,但把 bf16->fp8 转换推给了接收侧——AITER 的量化网格按输入行数(MoRI 填充后的接收缓冲,3072*8=24576 行)而非实际 ~112 个 live token 分配,开销集中在接收侧。PR body 明确提出“quantizes on the send side, over live tokens, emitting the fp8 + group-32 e8m0 layout the MoE kernels consume directly”。
实现拆解
- 扩展 DispatchDtype 枚举与布局计算(python/sglang/srt/layers/moe/token_dispatcher/moriep.py):新增 mxfp8 = "mxfp8_blockwise" 成员;在 init_mori_op 中为 mxfp8 计算 scale_dim = hidden_size // MXFP4_BLOCK_SIZE(每 32 通道一个 scale,DSV4 隐藏维度 7168 → 224 个 scale),scale_type_size 取 torch.float8_e8m0fnu.itemsize(1 字节),确保 dispatch 缓冲按 e8m0 布局精确分配,不欠分配也不浪费。
- 环境变量接线与 aiter 能力探测(moriep.py):_apply_dispatch_dtype_override 解析 SGLANG_MORI_DISPATCH_DTYPE=mxfp8 时,先调用新增的 _aiter_supports_mxfp8_dispatch()(lru_cache 缓存,用 inspect 检查 aiter.get_hip_quant 及其 per_1x32 实例是否接受 scale_type 参数)验证 aiter 能消费该布局;不支持则 warning_once 并回退 bf16。默认仍是 bf16,不改变既有行为。
- 接收侧跳过 upscale(python/sglang/srt/layers/moe/moe_runner/aiter.py 的 _pre_permute_deepep_to_aiter):新增 is_mx_fp8_dispatch 判断——a1_scale 存在、dtype 为 float8_e8m0fnu 且非 fp4 payload;命中时跳过 is_w4a4 分支的 fp8->bf16 upscale,fp8 + e8m0 数据直接交给 fused_moe。与普通 fp8 dispatch 的区分依据是 scale 的 dtype(fp32 vs e8m0),而非 dispatch 枚举本身。
- 测试配套:新增 test/registered/unit/layers/test_moriep_mxfp8_dispatch.py(81 行,9 个 CPU 用例),锁定枚举唯一性、MXFP4_BLOCK_SIZE=32、7168//32=224、e8m0 1 字节、env 映射到枚举成员、空 live token 批的 scale 形状 (0, 224);经 register_cpu_ci 注册到 base-a-test-cpu 套件。
- 性能验证:MI355X(gfx950)、TP8 + DP8 attention + EP8/MoRI + EAGLE MTP、8192 in / 1024 out,同构建仅切换环境变量,每 conc 跑一轮;所有数据点均超出 0.5-2.7% 的 run-to-run 噪声带。
| conc |
吞吐 bf16 (tok/s) |
吞吐 mxfp8 (tok/s) |
吞吐提升 |
TPOT bf16 (ms) |
TPOT mxfp8 (ms) |
TPOT 改善 |
| 4 |
1022.03 |
1121.97 |
+9.8% |
33.30 |
30.19 |
-9.3% |
| 8 |
1818.01 |
1969.84 |
+8.4% |
35.52 |
32.73 |
-7.9% |
| 16 |
3080.51 |
3277.43 |
+6.4% |
42.12 |
39.52 |
-6.2% |
| 32 |
4527.58 |
4905.24 |
+8.3% |
57.79 |
53.17 |
-8.0% |
| 64 |
6264.62 |
6827.92 |
+9.0% |
85.21 |
78.19 |
-8.2% |
关键文件:
python/sglang/srt/layers/moe/token_dispatcher/moriep.py(模块 专家分发;类别 source;类型 dependency-wiring;符号 DispatchDtype, _aiter_supports_mxfp8_dispatch, init_mori_op, _apply_dispatch_dtype_override): 变更入口与核心接线:新增 DispatchDtype.mxfp8 枚举、_aiter_supports_mxfp8_dispatch 启动探测、init_mori_op 中 e8m0 布局计算(scale_dim = hidden_size // 32、1 字节 scale)、_apply_dispatch_dtype_override 的环境变量解析与 bf16 回退。
python/sglang/srt/layers/moe/moe_runner/aiter.py(模块 专家执行;类别 source;类型 core-logic;符号 _pre_permute_deepep_to_aiter): 接收侧核心逻辑:_pre_permute_deepep_to_aiter 中通过 a1_scale.dtype == float8_e8m0fnu 识别 mxfp8 dispatch,跳过 fp8->bf16 upscale,让 fp8 数据直达 fused_moe。
test/registered/unit/layers/test_moriep_mxfp8_dispatch.py(模块 单元测试;类别 test;类型 test-coverage;符号 test_mxfp8_member_exists_and_is_distinct, test_scale_group_size_is_32, test_scale_dim_matches_group_32_layout, test_e8m0_scale_is_one_byte): 9 个 CPU 契约用例锁定字节布局与 env 接线,覆盖最容易静默回归的点(scale group、scale 数量、e8m0 字节数、空 token 批形状);注册进 base-a-test-cpu CI 套件。
关键符号:_aiter_supports_mxfp8_dispatch, _apply_dispatch_dtype_override, init_mori_op, _pre_permute_deepep_to_aiter
关键源码片段
python/sglang/srt/layers/moe/token_dispatcher/moriep.py
变更入口与核心接线:新增 DispatchDtype.mxfp8 枚举、_aiter_supports_mxfp8_dispatch 启动探测、init_mori_op 中 e8m0 布局计算(scale_dim = hidden_size // 32、1 字节 scale)、_apply_dispatch_dtype_override 的环境变量解析与 bf16 回退。
# moriep.py:mxfp8 dispatch 模式的核心接线
# 1) 探测当前 aiter 能否消费 group-32 e8m0 布局;用 lru_cache 缓存,
# 避免每次初始化都重复 import + inspect。
@functools.lru_cache(maxsize=1)
def _aiter_supports_mxfp8_dispatch() -> bool:
"""确认 aiter 是否支持 fp8 激活 + group-32 e8m0 微缩放。
这里必须探测而不是假设:旧版 per_1x32 量化照样返回 fp8,
但附带连续 fp32 缩放,MoE 侧会把 fp32 字节误读为 e8m0,
产生垃圾输出而非异常。通过 inspect 检查签名,把问题拦在
启动期并回退到 bf16。
"""
try:
import inspect
from aiter import get_hip_quant
return "scale_type" in inspect.signature(get_hip_quant).parameters or any(
"scale_type" in inspect.signature(f).parameters
for f in (get_hip_quant(QuantType.per_1x32),)
)
except Exception:
return False
# 2) 枚举新增 mxfp8 成员:fp8 载荷 + group-32 e8m0 微缩放,
# 与 per_1x32(MXFP4 权重)MoE kernel 的输入格式完全一致。
class DispatchDtype(Enum):
bf16 = "bfloat16"
fp8 = "float8_blockwise"
fp4 = "mxfp4_blockwise"
mxfp8 = "mxfp8_blockwise"
# 3) 环境变量接线(_apply_dispatch_dtype_override 内):默认仍是 bf16;
# mxfp8 仅在 aiter 支持时启用,否则回退 bf16 并给出明确警告,
# 避免静默的字节布局错配。
elif dispatch_dtype == "mxfp8":
if _aiter_supports_mxfp8_dispatch():
self.dispatch_dtype = DispatchDtype.mxfp8
else:
logger.warning_once(
"SGLANG_MORI_DISPATCH_DTYPE=mxfp8 requires an aiter "
"build whose per_1x32 quant accepts scale_type "
"(for the group-32 e8m0 byte layout the MoE kernels "
"consume). This aiter does not, so the send-side "
"quant would emit continuous fp32 scales and the "
"MoE would read them as e8m0 bytes. Falling back to "
"bf16 dispatch."
)
python/sglang/srt/layers/moe/moe_runner/aiter.py
接收侧核心逻辑:_pre_permute_deepep_to_aiter 中通过 a1_scale.dtype == float8_e8m0fnu 识别 mxfp8 dispatch,跳过 fp8->bf16 upscale,让 fp8 数据直达 fused_moe。
# aiter.py:_pre_permute_deepep_to_aiter 中判断是否跳过 upscale 往返。
# mxfp8 dispatch 已携带 fp8 数据 + group-32 e8m0 缩放,正是 per_1x32
# 想要的输入格式,可直接交给 fused_moe;只有 fp8 dispatch 的
# group-128/fp32 缩放才需要 dequant 往返,两者的区分依据是
# a1_scale 的 dtype(fp32 vs e8m0)。
is_mx_fp8_dispatch = (
a1_scale is not None
and a1_scale.dtype == torch.float8_e8m0fnu
and not is_fp4_dispatch
)
if (
is_w4a4
and a1_scale is not None
and not is_fp4_dispatch
and not is_mx_fp8_dispatch
):
# W4A4 权重 + 普通 FP8 dispatch:先 dequant FP8 -> BF16;
# FP4 per_1x32 路径需要 BF16 输入。
hidden_states = upscale(
hidden_states, a1_scale, num_local_tokens, output_dtype
)
a1_scale = None
test/registered/unit/layers/test_moriep_mxfp8_dispatch.py
9 个 CPU 契约用例锁定字节布局与 env 接线,覆盖最容易静默回归的点(scale group、scale 数量、e8m0 字节数、空 token 批形状);注册进 base-a-test-cpu CI 套件。
# 契约测试:锁定 mxfp8 dispatch 的字节布局,防止“能跑但悄悄退化”。
# 错误 group 大小或错误 scale dtype 的 fp8 payload 仍能运行,
# 但会重新引入本模式要消除的 upscale 往返,症状只有吞吐损失。
def test_scale_dim_matches_group_32_layout():
"""每 32 通道一个 scale;错配会欠分配 scale buffer,
导致 kernel 越界读取。"""
assert HIDDEN % MXFP4_BLOCK_SIZE == 0
assert HIDDEN // MXFP4_BLOCK_SIZE == 224
def test_empty_token_batch_scale_shape():
"""decode 可能给某 rank 0 个 live token;空分支也必须产出
形状正确的 scale tensor,否则 all-to-all 会失步。"""
scale = torch.empty((0, HIDDEN // MXFP4_BLOCK_SIZE), dtype=torch.float8_e8m0fnu)
assert scale.shape == (0, 224)
assert scale.dtype == torch.float8_e8m0fnu
评论区精华
该 PR 没有 review 行内评论,核心讨论集中在 issue 评论与合入前的跨仓库协调:
- karverma-amd 在 issue 评论中说明配套 aiter PR(ROCm/aiter#4954)为 fused_moe 增加 12 行 a8w4 mxfp8 passthrough 分支,两边可任意顺序合并;若设置了 SGLANG_MORI_DISPATCH_DTYPE=mxfp8 但 aiter 缺少配套改动,8b0cd59 新增的启动探测会检测到 per_1x32 量化不接受 scale_type,记录警告并回退 bf16。
- 作者明确指出“Probed rather than assumed”的原因:旧版 per_1x32 量化仍返回 fp8 但带连续 fp32 缩放,MoE 侧会把 fp32 字节误读为 e8m0 产生垃圾输出而非异常,因此必须把问题拦截在启动期。
-
HaiShaw 直接 APPROVED 并通过 /tag-and-rerun-ci 触发 CI 复跑;4 个 commit 中后两个为格式化与 CI 注册的收尾(pre-commit、CPU CI 注册),说明过程经过了 lint 与测试基建的打磨。
-
aiter 配套改动、合并顺序与回退机制 (design): 通过启动期签名探测 + bf16 回退消除跨仓库依赖风险;HaiShaw 批准并触发 /tag-and-rerun-ci。
风险与影响
- 风险:
- 跨仓库依赖:mxfp8 模式的收益依赖 ROCm/aiter#4954 合入;缺省时虽有启动探测兜底回退,但若 aiter 侧把 scale_type 参数改名或改为 keyword-only 意外路径,inspect 探测可能误判,需要保持两侧联动维护。
- 字节布局静默回归:fp8 payload 配上错误 scale group(128)或错误 scale dtype(fp32)仍能运行,只是悄悄退回 upscale 往返,症状仅是吞吐损失;本 PR 用 CPU 契约测试钉住布局,但 GPU 端真实 kernel 路径没有被自动化单测覆盖(测试不依赖 GPU)。
- scale dtype 启发式判定:aiter.py 用 a1_scale.dtype == torch.float8_e8m0fnu 区分 mxfp8 与普通 fp8 dispatch;若未来 fp8 dispatch 更改 scale dtype 或出现新的 e8m0 变体,该判定可能误判分支。
- 空 live token 批:decode 可能给某 rank 0 个 token,空分支必须产出形状正确的 scale tensor,否则 all-to-all 失步;测试覆盖了 (0, 224) 形状,但真机多 rank 场景仍依赖 benchmark 验证。
- CI 状态:PR Test (Base) 显示失败而 Extra 通过,材料未给出失败原因,合入前应确认该失败与本 PR 无关。
- 影响:
- 用户影响:仅 AMD(MI355X/gfx950 等)且启用 MoRI 后端 + DSV4 类 per_1x32 权重的用户,显式设置 SGLANG_MORI_DISPATCH_DTYPE=mxfp8 可获得 6.4%-9.8% 吞吐提升与 6.2%-9.3% TPOT 改善;默认为 bf16,其余用户零影响。
- 系统影响:改动集中在 MoE dispatch 与 aiter runner 两条窄路径,不触碰调度器、KV cache 等核心模块;新增 1 个 CPU CI 测试文件,CI 负担约 5 秒。
- 团队影响:需要与 ROCm/aiter 仓库协调发布节奏(aiter 侧先合入则立即受益,后合入则自动回退),是跨仓库协作的典型样例。
- 风险标记:依赖配套 aiter 改动, 字节布局静默回归, GPU 路径无单测覆盖, 默认 bf16 的 opt-in 特性
关联脉络
- PR #35374 [Kernel] Add H200 MoE configs for Qwen3.5 and Qwen3.6: 同属 MoE 性能调优脉络(H200 FP8 吞吐提升 vs 本 PR MI355X MXFP8 吞吐提升),说明 MoE 量化/调度是持续优化重点,可对照不同硬件上的收益路径。
- PR #36529 [Fix][XPU/ROCm/NPU] Defer sgl_kernel.quantization import in expert_pack: 同为 ROCm/XPU 侧 MoE 量化路径的兼容性修复,与本次 aiter 依赖接线守护同一类“导入/量化时点”问题。
- PR #36736 [AMD][CI] Merge the four MI35x DeepSeek-V3.2 nightly jobs into two to save runtime: AMD MI35x 平台 CI 建设与本 PR 的 AMD DSV4 验证呼应,且本 PR 的单元测试也注册进了 CPU CI 套件。
参与讨论