# PR #54773 完整报告

- 仓库：`vllm-project/vllm`
- 标题：[ROCm][MoE] Fix gfx950 block scale swizzle for AITER Triton MXFP4 W4A16
- 合并时间：2026-09-02 03:48
- 原文链接：http://prhub.com.cn/vllm-project/vllm/pull/54773

---

# 执行摘要

- 一句话：修复 gfx950 上 AITER Triton MXFP4 W4A16 的 block scale swizzle 问题。
- 推荐动作：该 PR 值得精读，特别是对于维护 ROCm 后端和 MoE 内核的工程师。它展示了如何通过共享标志确保设备特定行为的一致性，设计决策清晰且向后兼容。

# 功能与动机

PR body 中描述：'The AITER Triton MXFP4 W4A16 backend crashes the process on gfx950 (MI355) with a GPU memory access fault, because it tells the kernel its block scales are not swizzled while weight loading has already swizzled them.' 修复了因 swizzle 标志不一致导致的内存访问故障。

# 实现拆解

1. **导入共享标志**：在文件 `vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py` 中，导入 `should_use_cdna4_mx_scale_swizzle` 函数，该函数封装了判断是否应使用 CDNA4 布局的逻辑。
2. **动态计算 swizzle 标志**：在 `aiter_triton_kernel_w4a16_moe_forward` 函数中，根据设备类型和 TP 大小计算 swizzle 标志：`swz = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None`。
3. **修改内核调用**：将两个 `moe_gemm_a16w4` GEMM 调用的 `swizzle_mx_scale` 参数从硬编码的 `None` 替换为动态计算的 `swz`，确保内核知道 block scales 的布局。
4. **添加注释说明**：在变更处添加注释，强调 load-time swizzle 和 kernel 参数必须一致，避免内核索引错误。
5. **测试验证**：PR body 提供了在 MI355X 硬件上的测试结果（12 passed, 0 skipped），但代码中没有直接测试文件变更。

关键文件：
- `vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py`（模块 MoE 专家层；类别 source；类型 data-contract；符号 aiter_triton_kernel_w4a16_moe_forward）: 核心变更文件，修复了内核调用中 block scale swizzle 标志的传递，解决了 gfx950 上的 GPU 内存访问故障。

关键符号：aiter_triton_kernel_w4a16_moe_forward

## 关键源码片段

### `vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py`

核心变更文件，修复了内核调用中 block scale swizzle 标志的传递，解决了 gfx950 上的 GPU 内存访问故障。

```python
def aiter_triton_kernel_w4a16_moe_forward(
    hidden_states: torch.Tensor,
    w1,
    w2,
    gating_output: torch.Tensor,
    topk: int,
    renormalize: bool,
    activation: MoEActivation = MoEActivation.SWIGLUOAI,
    quant_config: FusedMoEQuantConfig | None = None,
    # ... 其他参数 ...
):
    assert quant_config is not None and rocm_aiter_ops.is_enabled()
    from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
        should_use_cdna4_mx_scale_swizzle,
    )
    from vllm.platforms.rocm import on_gfx1250

    # ... 路由和激活设置 ...

    # Weight loading already swizzled the block scales on gfx950, so the kernel
    # has to be told; the load-time swizzle and this gate must agree or the
    # kernel indexes a swizzled buffer as if it were linear.
    swz = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None

    # 第一个 GEMM 调用：修复 swizzle_mx_scale 从硬编码 None 改为动态计算的 swz
    intermediate = moe_gemm_a16w4(
        hidden_states,
        w1_data,
        None,
        w1_wscale,
        None,
        None,
        quant_config.w1_bias,
        routing_data,
        gather_indx=gather_idx,
        gammas=gammas if apply_router_weight_on_input else None,
        swizzle_mx_scale=swz,  # 关键修复点
        apply_swiglu=True,
        alpha=swiglu_alpha,
        limit=swiglu_limit,
        swiglu_add_residual=swiglu_add_residual,
        unpadded_N=unpadded_N_w1,
        unpadded_K=unpadded_K_w1,
    )

    # 第二个 GEMM 调用：同样修复 swizzle_mx_scale
    out = moe_gemm_a16w4(
        intermediate,
        w2_data,
        None,
        w2_wscale,
        None,
        None,
        quant_config.w2_bias,
        routing_data,
        scatter_indx=scatter_idx,
        gammas=None if apply_router_weight_on_input else gammas,
        swizzle_mx_scale=swz,  # 关键修复点
        unpadded_N=unpadded_N_w2,
        unpadded_K=unpadded_K_w2,
    )

    return out

```

# 评论区精华

Review 中无实质性技术讨论；AndreasKaratzas 批准了 PR，claude[bot] 指出自动化 review 被禁用。

- 暂无高价值评论线程

# 风险与影响

- 风险：风险较低：修复仅影响 gfx950 设备上的 AITER Triton MXFP4 W4A16 路径，通过条件判断确保非 gfx950 设备行为不变（`swz` 为 `None`）。未引入新功能或破坏性变更，但需注意后续对 `should_use_cdna4_mx_scale_swizzle` 逻辑的修改可能影响所有设备。
- 影响：影响范围：修复了 gfx950 上的特定 crash bug，提升了 ROCm 硬件支持的稳定性。对用户：消除了 MI355 上的运行时故障；对系统：无广泛影响；对团队：需要关注设备特定逻辑的一致性，避免类似不匹配问题。
- 风险标记：硬件特定行为 , 数据契约不一致

# 关联脉络

- PR #50622 [Feature][ROCm] Add AITER Triton MXFP4 W4A16 MoE backend: 该 PR 引入了 AITER Triton MXFP4 W4A16 MoE 后端，本 PR 修复了其在 gfx950 上的 swizzle 问题。