#26884 [AMD] Fix GPT-OSS MXFP4 accuracy on ROCm AITER path

原始 PR 作者 kkHuang-amd 合并时间 2026-06-02 13:30 文件变更 5 提交数 4 评论 4 代码增减 +87 / -20

执行摘要

修复 AMD ROCm AITER 上 GPT-OSS MXFP4 精度错误

根据 PR 描述，在 ROCm AITER 上，GPT-OSS MXFP4 的 fused-MoE 路径一直在静默输出错误结果。有两个根本原因：

1) 当前的 shuffle_weight_a16w4 产生了 gate/up 交织的 tile 布局，而 AITER 的 MXFP4 fused-MoE kernel（包括 FlyDSL 和 CK）期望的是分隔布局；
2) 权重打乱后 is_shuffled 标志未正确传递（.data 赋值丢弃了 Python 自定义属性），导致 AITER 分配器选择了错误的 preshuffle_off codegen 分支。这两个 bug 相互作用，最终导致 GSM8K 评分从 ~0.88 跌至 ~0.21。

本 PR 展示了系统性的精度调试过程，涉及跨栈的权重布局、kernel dispatch 和属性传播问题。建议团队内对涉及 AITER MXFP4 的后续开发仔细审查 is_shuffled 标志的传播和 GateMode 配置。值得精读，尤其是在理解量化内核集成和平台适配方面。

讨论亮点

在 review 过程中，审核者 HaiShaw 提出关键要求：“Let try to set SGLANG_USE_AITER_MOE_GU_ITLV to false in gptoss + hip + aiter path?” 此要求被采纳，通过后续提交在 server_args.py 中自动设置该环境变量为 False，确保 GPT-OSS 路径默认使用分隔布局。审核者在后续迭代后批准了该 PR。

实现拆解

实现过程包含以下步骤：

修复权重布局与标志（python/sglang/srt/layers/quantization/mxfp4.py）：在 Mxfp4MoEMethod.process_weights_after_loading 和 apply 方法中，先对 w13 权重进行 gate/up 解交织（de-interleave）以匹配 AITER kernel 期望的分隔布局，然后用 shuffle_weight(is_guinterleave=False, gate_up=...) 和 shuffle_scale(is_guinterleave=False, gate_up=...) 替代旧的 shuffle_weight_a16w4 和 shuffle_scale_a16w4。同时显式设置 layer.w13_weight.is_shuffled = True 和 layer.w2_weight.is_shuffled = True，并确保 view 后的新张量重新标记 is_shuffled。
新增环境变量控制布局（python/sglang/srt/environ.py）：在 Envs 类中添加 SGLANG_USE_AITER_MOE_GU_ITLV = EnvBool(True)，默认启用交织布局以保持向后兼容。
调整 AITER run 方法（python/sglang/srt/layers/moe/moe_runner/aiter.py）：在 AiterRunnerCore.run 中，根据该环境变量选择 GateMode.SEPARATED 或 INTERLEAVE，并传入 swiglu_limit。
自动为 GPT-OSS 选择分隔布局（python/sglang/srt/server_args.py）：在 _handle_model_specific_adjustments 中，当检测到 ROCm、AITER 和 MXFP4 量化时，自动将 SGLANG_USE_AITER_MOE_GU_ITLV 设为 False，确保 GPT-OSS 路径始终使用分隔布局。
更新测试（test/registered/amd/accuracy/mi35x/test_gpt_oss_eval_mi35x.py）：在模型配置中显式设置 SGLANG_USE_AITER_MOE_GU_ITLV=0，使测试与实际部署行为一致。

这些修改涉及核心量化逻辑、运行时 kernel 分发、环境变量配置和服务端自动调优，共同解决了两个根本原因。

文件	模块	状态	重要度
`python/sglang/srt/layers/quantization/mxfp4.py`	量化层	modified	6.82
`python/sglang/srt/layers/moe/moe_runner/aiter.py`	MOE 执行器	modified	5.99
`python/sglang/srt/environ.py`	环境变量	modified	5.22
`python/sglang/srt/server_args.py`	服务配置	modified	5.11
`test/registered/amd/accuracy/mi35x/test_gpt_oss_eval_mi35x.py`	端到端测试	modified	4.63

关键符号

Mxfp4MoEMethod.process_weights_after_loading Mxfp4MoEMethod.apply AiterRunnerCore.run _handle_model_specific_adjustments __post_init__

关键源码片段

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

核心量化量化方法，修复权重布局和标志设置，是修复的主要位置。

if _use_aiter:
    # Bias must be fp32 for the AITER kernels.
    if layer.w13_weight_bias is not None:
        layer.w13_weight_bias.data = layer.w13_weight_bias.data.to(torch.float32)
    if layer.w2_weight_bias is not None:
        layer.w2_weight_bias.data = layer.w2_weight_bias.data.to(torch.float32)

    # HF GPT-OSS 将 w13 存储为 gate/up 交织的行对 [(g0, u0), (g1, u1), ...]。
    # AITER MXFP4 fused MoE 内核（FlyDSL `gate_mode="separated"` 和 CK `preshuffle_on`）
    # 期望分隔布局 [gate_0..gate_{N-1}, up_0..up_{N-1}]。
    # 在 tile shuffle 之前解交织 weights、scales 和 bias，
    # 使 shuffle 后的字节落在内核读取的正确布局中。
    e, n, k = layer.w13_weight.shape
    layer.w13_weight.view(torch.uint8).copy_(
        layer.w13_weight.data.view(torch.uint8)
        .view(e, n // 2, 2, k)
        .permute(0, 2, 1, 3)
        .contiguous()
        .view(e, n, k)
    )
    layer.w13_weight_scale.data = (
        layer.w13_weight_scale.data.view(e, n // 2, 2, -1)
        .permute(0, 2, 1, 3)
        .contiguous()
        .view(e, n, -1)
    )
    layer.w13_weight_bias.data = (
        layer.w13_weight_bias.data.view(-1, n // 2, 2)
        .permute(0, 2, 1)
        .contiguous()
        .view(-1, n)
    )

    # 使用 `is_guinterleave=False` 的 shuffle 以匹配 ATOM 和 tuned FlyDSL 配置。
    # 旧的 `shuffle_weight_a16w4` 使用了 gate/up 交织的 tile 布局，导致精度丢失。
    layer.w13_weight.data = shuffle_weight(
        layer.w13_weight, is_guinterleave=False, gate_up=True
    )
    shuffled_w13_scale = shuffle_scale(
        layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]),
        experts_cnt=self.num_experts,
        is_guinterleave=False,
        gate_up=True,
    )
    layer.w2_weight.data = shuffle_weight(
        layer.w2_weight, is_guinterleave=False, gate_up=False
    )
    shuffled_w2_scale = shuffle_scale(
        layer.w2_weight_scale.view(-1, layer.w2_weight_scale.shape[-1]),
        experts_cnt=self.num_experts,
        is_guinterleave=False,
        gate_up=False,
    )

    # 必须显式设置 is_shuffled 属性，否则 AITER fused_moe
    # 会通过 getattr(w1, "is_shuffled", False) 判断为 False，
    # 从而选择错误的 preshuffle_off codegen 分支。
    layer.w13_weight.is_shuffled = True
    layer.w2_weight.is_shuffled = True

    layer.w13_weight_scale = torch.nn.Parameter(
        shuffled_w13_scale, requires_grad=False
    )
    layer.w2_weight_scale = torch.nn.Parameter(
        shuffled_w2_scale, requires_grad=False
    )

python/sglang/srt/layers/moe/moe_runner/aiter.py dependency-wiring

AITER MoE 运行时核心，根据环境变量路由 gate_mode。

if quant_info.swiglu_limit > 0:
    # 默认使用 INTERLEAVE 以保持向后兼容，但可通过环境变量切换为 SEPARATED。
    # Mxfp4MoEMethod (gpt-oss MXFP4) 和 tuned FlyDSL 内核使用 SEPARATED 布局。
    extra["gate_mode"] = (
        GateMode.INTERLEAVE.value
        if envs.SGLANG_USE_AITER_MOE_GU_ITLV.get()
        else GateMode.SEPARATED.value
    )
    extra["swiglu_limit"] = quant_info.swiglu_limit

python/sglang/srt/environ.py core-logic

声明新的环境变量，控制 gate/up tile 布局默认值。

# AMD & ROCm
SGLANG_USE_AITER = EnvBool(False)
SGLANG_USE_AITER_UNIFIED_ATTN = EnvBool(False)
# 选择 AITER MoE gate/up tile 布局：True -> 交织（匹配 FlyDSL `gate_mode="interleave"`），
# False -> 分隔（匹配 `gate_mode="separated"`，用于 gptoss_fp4 tuned 配置和 Mxfp4MoEMethod 的权重 shuffle）。
SGLANG_USE_AITER_MOE_GU_ITLV = EnvBool(True)
SGLANG_ROCM_FUSED_DECODE_MLA = EnvBool(False)

评论区精华

设置 SGLANG_USE_AITER_MOE_GU_ITLV 默认值 设计

HaiShaw 建议在 gptoss+hip+aiter 路径中将 SGLANG_USE_AITER_MOE_GU_ITLV 设置为 false。

结论：已采纳，在 server_args.py 中自动设为 False。 · 已解决

风险与影响

主要风险包括：

1) 布局不兼容：环境变量 SGLANG_USE_AITER_MOE_GU_ITLV 默认 True（交织布局），但 GPT-OSS 路径通过 server_args.py 自动设为 False。如果其他模型也使用 AITER MXFP4 且需要交织布局，现有关键路径仍保持默认，但任何新增的 MXFP4 模型必须确保 shuffle 阶段与 gate_mode 对齐。
2) 属性传播遗漏：虽然本 PR 修复了 .data 赋值丢失 is_shuffled 属性的问题，但其他类似模式（如 view 后未重新标记）可能存在于代码其他部分。
3) 测试覆盖：当前测试仅覆盖 GPT-OSS 模型和特定的 FlyDSL/CK kernel 组合，未覆盖其他 AITER MXFP4 路径（如其他模型的交织布局）。
4) 环境变量配置项增加：新增的环境变量需要团队理解其含义，否则可能误用。

用户影响：修复了 ROCm AITER 上 GPT-OSS MXFP4 路径的精度问题，GSM8K 指标恢复到与 ATOM 基线一致（0.88）。对不涉及 ROCm/AITER 的用户无影响。系统影响：无性能回退，仅改变权重布局和 kernel 选择；新增环境变量默认保留原有行为。团队影响：维护者需注意新增的环境变量及其在不同模型组中的默认值设定；社区贡献者引入新的 AITER MXFP4 模型时必须确认布局一致性。总体影响限于特定硬件和模型组合，但修复了关键精度问题。

核心路径变更环境变量配置依赖平台特定修复布局对齐风险

关联 Issue

未识别关联 Issue

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

完整报告

执行摘要

一句话：修复 AMD ROCm AITER 上 GPT-OSS MXFP4 精度错误
推荐动作：本 PR 展示了系统性的精度调试过程，涉及跨栈的权重布局、kernel dispatch 和属性传播问题。建议团队内对涉及 AITER MXFP4 的后续开发仔细审查 is_shuffled 标志的传播和 GateMode 配置。值得精读，尤其是在理解量化内核集成和平台适配方面。

功能与动机

根据 PR 描述，在 ROCm AITER 上，GPT-OSS MXFP4 的 fused-MoE 路径一直在静默输出错误结果。有两个根本原因：

实现拆解

实现过程包含以下步骤：

修复权重布局与标志（python/sglang/srt/layers/quantization/mxfp4.py）：在 Mxfp4MoEMethod.process_weights_after_loading 和 apply 方法中，先对 w13 权重进行 gate/up 解交织（de-interleave）以匹配 AITER kernel 期望的分隔布局，然后用 shuffle_weight(is_guinterleave=False, gate_up=...) 和 shuffle_scale(is_guinterleave=False, gate_up=...) 替代旧的 shuffle_weight_a16w4 和 shuffle_scale_a16w4。同时显式设置 layer.w13_weight.is_shuffled = True 和 layer.w2_weight.is_shuffled = True，并确保 view 后的新张量重新标记 is_shuffled。
新增环境变量控制布局（python/sglang/srt/environ.py）：在 Envs 类中添加 SGLANG_USE_AITER_MOE_GU_ITLV = EnvBool(True)，默认启用交织布局以保持向后兼容。
调整 AITER run 方法（python/sglang/srt/layers/moe/moe_runner/aiter.py）：在 AiterRunnerCore.run 中，根据该环境变量选择 GateMode.SEPARATED 或 INTERLEAVE，并传入 swiglu_limit。
自动为 GPT-OSS 选择分隔布局（python/sglang/srt/server_args.py）：在 _handle_model_specific_adjustments 中，当检测到 ROCm、AITER 和 MXFP4 量化时，自动将 SGLANG_USE_AITER_MOE_GU_ITLV 设为 False，确保 GPT-OSS 路径始终使用分隔布局。
更新测试（test/registered/amd/accuracy/mi35x/test_gpt_oss_eval_mi35x.py）：在模型配置中显式设置 SGLANG_USE_AITER_MOE_GU_ITLV=0，使测试与实际部署行为一致。

这些修改涉及核心量化逻辑、运行时 kernel 分发、环境变量配置和服务端自动调优，共同解决了两个根本原因。

关键文件：

python/sglang/srt/layers/quantization/mxfp4.py（模块量化层；类别 source；类型 core-logic；符号 Mxfp4MoEMethod.process_weights_after_loading, Mxfp4MoEMethod.apply）: 核心量化量化方法，修复权重布局和标志设置，是修复的主要位置。
python/sglang/srt/layers/moe/moe_runner/aiter.py（模块 MOE执行器；类别 source；类型 dependency-wiring；符号 AiterRunnerCore.run）: AITER MoE 运行时核心，根据环境变量路由 gate_mode。
python/sglang/srt/environ.py（模块环境变量；类别 source；类型 core-logic；符号 SGLANG_USE_AITER_MOE_GU_ITLV）: 声明新的环境变量，控制 gate/up tile 布局默认值。
python/sglang/srt/server_args.py（模块服务配置；类别 source；类型 core-logic；符号 _handle_model_specific_adjustments）: 自动为 GPT-OSS 模型设置环境变量为 False，确保使用分隔布局。
test/registered/amd/accuracy/mi35x/test_gpt_oss_eval_mi35x.py（模块端到端测试；类别 test；类型 test-coverage；符号 post_init）: 更新测试配置，显式设置环境变量以覆盖自动化默认值，验证修复。

关键符号：Mxfp4MoEMethod.process_weights_after_loading, Mxfp4MoEMethod.apply, AiterRunnerCore.run, _handle_model_specific_adjustments, post_init

关键源码片段

`python/sglang/srt/layers/quantization/mxfp4.py`

核心量化量化方法，修复权重布局和标志设置，是修复的主要位置。

if _use_aiter:
    # Bias must be fp32 for the AITER kernels.
    if layer.w13_weight_bias is not None:
        layer.w13_weight_bias.data = layer.w13_weight_bias.data.to(torch.float32)
    if layer.w2_weight_bias is not None:
        layer.w2_weight_bias.data = layer.w2_weight_bias.data.to(torch.float32)

    # HF GPT-OSS 将 w13 存储为 gate/up 交织的行对 [(g0, u0), (g1, u1), ...]。
    # AITER MXFP4 fused MoE 内核（FlyDSL `gate_mode="separated"` 和 CK `preshuffle_on`）
    # 期望分隔布局 [gate_0..gate_{N-1}, up_0..up_{N-1}]。
    # 在 tile shuffle 之前解交织 weights、scales 和 bias，
    # 使 shuffle 后的字节落在内核读取的正确布局中。
    e, n, k = layer.w13_weight.shape
    layer.w13_weight.view(torch.uint8).copy_(
        layer.w13_weight.data.view(torch.uint8)
        .view(e, n // 2, 2, k)
        .permute(0, 2, 1, 3)
        .contiguous()
        .view(e, n, k)
    )
    layer.w13_weight_scale.data = (
        layer.w13_weight_scale.data.view(e, n // 2, 2, -1)
        .permute(0, 2, 1, 3)
        .contiguous()
        .view(e, n, -1)
    )
    layer.w13_weight_bias.data = (
        layer.w13_weight_bias.data.view(-1, n // 2, 2)
        .permute(0, 2, 1)
        .contiguous()
        .view(-1, n)
    )

    # 使用 `is_guinterleave=False` 的 shuffle 以匹配 ATOM 和 tuned FlyDSL 配置。
    # 旧的 `shuffle_weight_a16w4` 使用了 gate/up 交织的 tile 布局，导致精度丢失。
    layer.w13_weight.data = shuffle_weight(
        layer.w13_weight, is_guinterleave=False, gate_up=True
    )
    shuffled_w13_scale = shuffle_scale(
        layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]),
        experts_cnt=self.num_experts,
        is_guinterleave=False,
        gate_up=True,
    )
    layer.w2_weight.data = shuffle_weight(
        layer.w2_weight, is_guinterleave=False, gate_up=False
    )
    shuffled_w2_scale = shuffle_scale(
        layer.w2_weight_scale.view(-1, layer.w2_weight_scale.shape[-1]),
        experts_cnt=self.num_experts,
        is_guinterleave=False,
        gate_up=False,
    )

    # 必须显式设置 is_shuffled 属性，否则 AITER fused_moe
    # 会通过 getattr(w1, "is_shuffled", False) 判断为 False，
    # 从而选择错误的 preshuffle_off codegen 分支。
    layer.w13_weight.is_shuffled = True
    layer.w2_weight.is_shuffled = True

    layer.w13_weight_scale = torch.nn.Parameter(
        shuffled_w13_scale, requires_grad=False
    )
    layer.w2_weight_scale = torch.nn.Parameter(
        shuffled_w2_scale, requires_grad=False
    )

`python/sglang/srt/layers/moe/moe_runner/aiter.py`

AITER MoE 运行时核心，根据环境变量路由 gate_mode。

if quant_info.swiglu_limit > 0:
    # 默认使用 INTERLEAVE 以保持向后兼容，但可通过环境变量切换为 SEPARATED。
    # Mxfp4MoEMethod (gpt-oss MXFP4) 和 tuned FlyDSL 内核使用 SEPARATED 布局。
    extra["gate_mode"] = (
        GateMode.INTERLEAVE.value
        if envs.SGLANG_USE_AITER_MOE_GU_ITLV.get()
        else GateMode.SEPARATED.value
    )
    extra["swiglu_limit"] = quant_info.swiglu_limit

`python/sglang/srt/environ.py`

声明新的环境变量，控制 gate/up tile 布局默认值。

# AMD & ROCm
SGLANG_USE_AITER = EnvBool(False)
SGLANG_USE_AITER_UNIFIED_ATTN = EnvBool(False)
# 选择 AITER MoE gate/up tile 布局：True -> 交织（匹配 FlyDSL `gate_mode="interleave"`），
# False -> 分隔（匹配 `gate_mode="separated"`，用于 gptoss_fp4 tuned 配置和 Mxfp4MoEMethod 的权重 shuffle）。
SGLANG_USE_AITER_MOE_GU_ITLV = EnvBool(True)
SGLANG_ROCM_FUSED_DECODE_MLA = EnvBool(False)

评论区精华

设置 SGLANG_USE_AITER_MOE_GU_ITLV 默认值 (design): 已采纳，在 server_args.py 中自动设为 False。

风险与影响

风险：主要风险包括：
1) 布局不兼容：环境变量 SGLANG_USE_AITER_MOE_GU_ITLV 默认 True（交织布局），但 GPT-OSS 路径通过 server_args.py 自动设为 False。如果其他模型也使用 AITER MXFP4 且需要交织布局，现有关键路径仍保持默认，但任何新增的 MXFP4 模型必须确保 shuffle 阶段与 gate_mode 对齐。
2) 属性传播遗漏：虽然本 PR 修复了 .data 赋值丢失 is_shuffled 属性的问题，但其他类似模式（如 view 后未重新标记）可能存在于代码其他部分。
3) 测试覆盖：当前测试仅覆盖 GPT-OSS 模型和特定的 FlyDSL/CK kernel 组合，未覆盖其他 AITER MXFP4 路径（如其他模型的交织布局）。
4) 环境变量配置项增加：新增的环境变量需要团队理解其含义，否则可能误用。
- 影响：用户影响：修复了 ROCm AITER 上 GPT-OSS MXFP4 路径的精度问题，GSM8K 指标恢复到与 ATOM 基线一致（0.88）。对不涉及 ROCm/AITER 的用户无影响。系统影响：无性能回退，仅改变权重布局和 kernel 选择；新增环境变量默认保留原有行为。团队影响：维护者需注意新增的环境变量及其在不同模型组中的默认值设定；社区贡献者引入新的 AITER MXFP4 模型时必须确认布局一致性。总体影响限于特定硬件和模型组合，但修复了关键精度问题。
- 风险标记：核心路径变更, 环境变量配置依赖, 平台特定修复, 布局对齐风险

关联脉络

暂无明显关联 PR

#26884 [AMD] Fix GPT-OSS MXFP4 accuracy on ROCm AITER path

执行摘要

修复 AMD ROCm AITER 上 GPT-OSS MXFP4 精度错误

实现拆解

评论区精华

风险与影响

关联 Issue

未识别关联 Issue

完整报告

参与讨论