执行摘要

• 一句话：重构 ROCm MXFP4 W4A8 MoE，引入 oracle 后端选择
• 推荐动作：该 PR 很有价值，建议精读。重点关注 oracle 后端选择架构的演化（select_mxfp4_moe_backend 如何根据 activation_key 参数区分 W4A8 和 W4A16），以及 AiterW4A8ExpertsMonolithic 类如何通过 _supports_quant_scheme 声明自身能力。团队在后续添加新量化方案时可遵循此模式。

功能与动机

根据 PR body，此变更旨在移除已被 oracle 替代的 QuarkOCP_MX_MoEMethod_OSS，并为 ROCm 平台添加 AITER w4a8 后端。PR 作者在 body 中说明：'Remove QuarkOCP_MX_MoEMethod_OSS and add aiter w4a8 backend.' 以及 'Add unittest cases for rocm w4a16, w4a8 fused moe.' 这表明统一量化方案的后端选择是关键目标。

实现拆解

1. 提取 AITER W4A8 MoE 内核：创建新文件 aiter_mxfp4_w4a8_moe.py，包含 aiter_triton_kernel_w4a8_moe_forward 和 triton_kernel_fused_mxfp4_w4a8_experts 函数，以及 AiterW4A8ExpertsMonolithic 类。该类继承自 FusedMoEExpertsMonolithic，实现 _supports_current_device、_supports_no_act_and_mul、_supports_quant_scheme 等接口，明确声明对 W4A8（MXFP4 权重 + 静态 FP8 激活）方案的支持。此前 w4a8 内核内联在 gpt_oss_triton_kernels_moe.py 中，现完全移除。

2. 重构 oracle MXFP4 后端选择：在 oracle/mxfp4.py 中，将 AITER 拆分为 AITER_MXFP4_BF16（W4A16，CK 内核）和 AITER_MXFP4_FP8（W4A8，Triton 内核）两个后端。修改 select_gpt_oss_mxfp4_moe_backend 为 select_mxfp4_moe_backend，新增 activation_key 参数，使其能根据量化配置选择适合的后端。更新 backend_to_kernel_cls 映射，将 AITER_MXFP4_FP8 映射到 AiterW4A8ExpertsMonolithic。

3. 统一 Quark OCP MX 方案：在 quark_moe.py 中，删除 QuarkOCP_MX_MoEMethod_OSS 类及其引用，将所有 OCP MX 方案（w4a16、w4a8 等）统一由 QuarkOCP_MX_MoEMethod 处理。该类的 __init__ 方法根据 self.ocp_mx_scheme（如 w_mxfp4 或 w_mxfp4_a_fp8）调用 select_mxfp4_moe_backend，通过 activation_key 参数区分 w4a16 和 w4a8。同时移除对 _swizzle_mxfp4 的导入，添加 kFp8StaticTensorSym 等新导入。

4. 更新外围模块 mxfp4.py：调整 GptOssMxfp4MoEMethod 和 Mxfp4MoEMethod 的初始化为使用新的选择函数：前者调用 select_mxfp4_moe_backend，后者调用 select_deepseek_v4_mxfp4_moe_backend（原 select_mxfp4_moe_backend 更名）。

5. 添加 ROCm 集成测试：在 test_ocp_mx_moe.py 中新增 test_rocm_mxfp4_moe_oracle 函数，通过 pytest.mark.parametrize 覆盖 w4a16 和 w4a8 两种配置，创建模拟层并调用 oracle 选择的专家类进行正向传播，与参考实现对比精度。参考实现新增 swigluoai 激活函数，支持交错布局，用于 w4a8 内核。

关键文件：

• vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py（模块 MoE 核；类别 source；类型 core-logic；符号 aiter_triton_kernel_w4a8_moe_forward, triton_kernel_fused_mxfp4_w4a8_experts, AiterW4A8ExpertsMonolithic, init）: 核心新文件，实现了 AITER 的 W4A8 MoE 专家类 AiterW4A8ExpertsMonolithic 及其配套内核函数，是 oracle 后端选择机制的消费端。
• vllm/model_executor/layers/quantization/quark/quark_moe.py（模块 量化层；类别 source；类型 core-logic；符号 _setup_kernel_via_oracle, _setup_kernel, QuarkOCP_MX_MoEMethod_OSS, init）: 核心量化层文件，移除了 QuarkOCP_MX_MoEMethod_OSS，统一所有 OCP MX 方案到 QuarkOCP_MX_MoEMethod 并通过 select_mxfp4_moe_backend 选择后端。
• vllm/model_executor/layers/fused_moe/oracle/mxfp4.py（模块 后端选择；类别 source；类型 core-logic；符号 select_gpt_oss_mxfp4_moe_backend, select_mxfp4_moe_backend, select_deepseek_v4_mxfp4_moe_backend）: oracle 后端选择核心文件，新增 AITER_MXFP4_FP8 后端，修改 select_mxfp4_moe_backend 支持 activation_key 参数。
• vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py（模块 MoE 核；类别 source；类型 refactor；符号 triton_kernel_fused_mxfp4_w4a8_experts）: 原 w4a8 内核内联于此，现被移除并提取到新文件。
• tests/kernels/moe/test_ocp_mx_moe.py（模块 测试；类别 test；类型 test-coverage；符号 swigluoai, used, test_rocm_mxfp4_moe_oracle, MockLayer）: 新增 ROCm oracle 单元测试，覆盖 w4a16 和 w4a8 方案精度验证。

关键符号：aiter_triton_kernel_w4a8_moe_forward, triton_kernel_fused_mxfp4_w4a8_experts, AiterW4A8ExpertsMonolithic._supports_quant_scheme, AiterW4A8ExpertsMonolithic.apply, select_mxfp4_moe_backend, select_deepseek_v4_mxfp4_moe_backend, QuarkOCP_MX_MoEMethod.init, QuarkOCP_MX_MoEMethod._setup_kernel_via_oracle, backend_to_kernel_cls, test_rocm_mxfp4_moe_oracle

评论区精华

Review 中的核心讨论包括：

• 后端名称映射清理：BowenBao 在 map_mxfp4_backend 处添加 TODO，期望 runner_backend 只包含后端名而不含 dtype 后缀，dtype 应由量化配置推断。已在 sig-quantization 中讨论。
• activation_key 参数的必要性：robertgshaw2-redhat 认为该参数可从 config 派生，不应显式传入。BowenBao 回应目前量化信息不在 FusedMoEConfig 中，因此需要传递激活键。后离线同步决定保留该参数。
• emulation 逻辑简化：gemini-code-assist 提议定义清晰的 is_native_backend_available 辅助函数简化 emulation 条件判断，BowenBao 指出已存在 TODO，将在后续 PR 中重构。
• 新专家类的文件归属：reviewer 质疑 AiterW4A8ExpertsMonolithic 为何放在 gpt_oss_triton_kernels_moe.py，后 BowenBao 将其移至独立文件 aiter_mxfp4_w4a8_moe.py。

• 测试精度阈值：AndreasKaratzas 建议添加统计打印（max error、mean error）并收紧精度阈值，BowenBao 添加了统计打印但表示阈值宽松是因参考实现不完全对齐，将在后续跟进。

• 后端名称映射清理 (design): 作为 TODO 记录，不在此 PR 解决。

• activation_key 参数设计 (design): 保留参数，待后续将量化信息并入 config 后再重构。
• emulation 逻辑简化 (design): 已添加 TODO，后续 w4a4 迁移完成后再简化。

风险与影响

• 风险：主要技术风险包括：
  1. 核心路径变更风险：quark_moe.py 中的 QuarkOCP_MX_MoEMethod_OSS 被移除，所有 OCP MX 方案统一走 QuarkOCP_MX_MoEMethod。如果新 oracle 选择逻辑有缺陷，可能导致 w4a8 或 w4a16 方案回退到错误后端（如 emulation），影响推理精度。虽然测试覆盖了 w4a16 和 w4a8，但 w4a4 等方案尚未通过 oracle 路径，可能被遗漏。
  2. ROCm 平台专属：新 AITER 内核仅限 ROCm GFX950 架构，在其他平台上会通过 _supports_current_device 返回 False，由 oracle 回退到其他后端（如 emulation），不影响现有用户。
  3. 性能风险：新内核依赖于 aiter 和 triton_kernels 库的版本兼容性，如果依赖项升级可能引入性能回退。测试中已包含端到端精度验证，但缺乏性能回归测试。
     - 影响：
     • 用户影响：ROCm 用户在使用 w4a8 量化方案时将自动启用新的 AITER Triton 内核，提升推理速度。所有 MXFP4 方案的后端选择统一由 oracle 接管，用户无需手动指定。
     • 系统影响：代码结构更清晰，QuarkOCP_MX_MoEMethod_OSS 冗余类移除，降低了维护成本。oracle 后端选择框架得到扩展，支持 activation_key 过滤，为后续添加更多量化方案奠定基础。
     • 团队影响：AMD ROCm 团队可直接通过添加专家类和 oracle 分支来支持新的量化方案，开发效率提升。
     • 风险标记：核心路径变更, ROCm 平台专用, 依赖库版本兼容

关联脉络

• PR #37481 [XPU] enable is_act_and_mul for xpu: 均涉及 fused_moe 模块的激活函数扩展，本 PR 新增 swigluoai 参考实现用于 w4a8 测试。
• PR #39931 [Feature] TurboQuant: support hybrid models and uniform quantization: 同为量化方案扩展，本 PR 将 w4a8 后端选择通过 oracle 统一，与 TurboQuant 引入的新量化方案有架构关联。
• PR #41569 [ROCm][CI] Fix MLA prefill scale for DeepSeek GSM8K: 同为 ROCm 平台改进，本 PR 新增的 AITER 内核仅对 ROCm GFX950 生效，与 ROCm CI 稳定性相关。