执行摘要

• 一句话：修复Quark MXFP4 emulation MoE维度block对齐
• 推荐动作：该 PR 修复了实际 bug 且重构方向正确，值得精读。review 中的共享 helper 设计思路体现了代码复用和统一管理的重要性。

功能与动机

Quark-quantized MXFP4/FP8 GPT-OSS loading的emulation路径在TP下缺失OCP block对齐，导致RuntimeError: The size of tensor a (1440) must match the size of tensor b (1472)。需要修复以使emulation后端正常工作。

实现拆解

1. 在 vllm/model_executor/layers/fused_moe/oracle/mxfp4.py 中导入 OCP_MX_BLOCK_SIZE，并在 mxfp4_round_up_hidden_size_and_intermediate_size 函数中增加 Mxfp4MoeBackend.EMULATION 分支，将 hidden_size 和 intermediate_size 向上 round up 到 OCP_MX_BLOCK_SIZE（32），确保 scale 缓冲区维度不会被截断。
2. 在 vllm/model_executor/layers/quantization/quark/quark_moe.py 的 maybe_roundup_sizes 方法中，移除原先针对 emulation 后端的排除条件 self.mxfp4_backend != Mxfp4MoeBackend.EMULATION，使所有非 None 后端都调用共享对齐函数，保持逻辑统一。
3. 在 tests/kernels/moe/test_ocp_mx_moe.py 中新增参数化测试 test_mxfp4_emulation_rounds_up_to_block_size，覆盖多种 TP 分片后维度组合（如 2880/720, 2880/360, 2880/2880, 90/90），验证 roundup 结果与预期一致且能被 OCP_MX_BLOCK_SIZE 整除。测试通过 @pytest.mark.skipif(not ROCM_AVAILABLE) 限定仅在 ROCm 下执行。

关键文件：

• vllm/model_executor/layers/fused_moe/oracle/mxfp4.py（模块 MoE量化；类别 source；类型 data-contract；符号 mxfp4_round_up_hidden_size_and_intermediate_size）: 核心文件：在共享对齐函数中添加 EMULATION 分支，并导入 OCP_MX_BLOCK_SIZE
• vllm/model_executor/layers/quantization/quark/quark_moe.py（模块 量化层；类别 source；类型 data-contract；符号 maybe_roundup_sizes）: 移除排除emulation后端的条件，使emulation也调用共享对齐函数，确保一致性。
• tests/kernels/moe/test_ocp_mx_moe.py（模块 测试；类别 test；类型 test-coverage；符号 test_mxfp4_emulation_rounds_up_to_block_size）: 新增单元测试，验证emulation后端维度对齐的roundup逻辑，覆盖关键shard场景。

关键符号：mxfp4_round_up_hidden_size_and_intermediate_size, maybe_roundup_sizes, test_mxfp4_emulation_rounds_up_to_block_size

关键源码片段

vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

核心文件：在共享对齐函数中添加 EMULATION 分支，并导入 OCP_MX_BLOCK_SIZE

# 在文件开头导入 OCP_MX_BLOCK_SIZE（32）
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
    OCP_MX_BLOCK_SIZE,
)
​
def mxfp4_round_up_hidden_size_and_intermediate_size(
    backend: Mxfp4MoeBackend, hidden_size: int, intermediate_size: int
) -> tuple[int, int]:
    """
    Round up hidden_size and intermediate_size based on backend requirements.
    模拟后端只需要 OCP MX block 对齐，以防止 per-block scale 缓冲区被截断。
    """
    if backend == Mxfp4MoeBackend.EMULATION:
        # Emulation has no kernel tile; it only needs OCP MX block alignment.
        intermediate_size = round_up(intermediate_size, OCP_MX_BLOCK_SIZE)
        hidden_size = round_up(hidden_size, OCP_MX_BLOCK_SIZE)
    elif backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
        # DeepGEMM requires M/N/K alignment
        intermediate_size = round_up(intermediate_size, 128)
        hidden_size = round_up(hidden_size, 128)
    # ... 其他后端分支（Marlin, TRTLLM, FlashInfer, ROCm, CPU 等）略
    else:
        # 默认对齐到 64
        intermediate_size = round_up(intermediate_size, 64)
    return hidden_size, intermediate_size

tests/kernels/moe/test_ocp_mx_moe.py

新增单元测试，验证emulation后端维度对齐的roundup逻辑，覆盖关键shard场景。

# 参数化测试，验证 emulation 后端的 round up 行为
@pytest.mark.skipif(not ROCM_AVAILABLE, reason="emulation backend targets ROCm")
@pytest.mark.parametrize(
    "hidden_size,intermediate_size,expected_hidden,expected_intermediate",
    [
        (2880, 720, 2880, 736), # GPT-OSS TP=4 shard: 720 -> round_up(720, 32)
        (2880, 360, 2880, 384), # GPT-OSS TP=8 shard: 360 -> 384
        (2880, 2880, 2880, 2880), # 已经 block-aligned，保持不变
        (90, 90, 96, 96), # 两个维度都未对齐
    ],
)
def test_mxfp4_emulation_rounds_up_to_block_size(
    hidden_size: int,
    intermediate_size: int,
    expected_hidden: int,
    expected_intermediate: int,
):
    """确保 emulation 后端将 each per-partition dim 向上对齐到 OCP_MX_BLOCK_SIZE"""
    from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
        Mxfp4MoeBackend,
        mxfp4_round_up_hidden_size_and_intermediate_size,
    )
    from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
        OCP_MX_BLOCK_SIZE,
    )
​
    rounded_hidden, rounded_intermediate = (
        mxfp4_round_up_hidden_size_and_intermediate_size(
            Mxfp4MoeBackend.EMULATION, hidden_size, intermediate_size
        )
    )
​
    assert rounded_hidden == expected_hidden
    assert rounded_intermediate == expected_intermediate
    # scale 缓冲区必须不产生 floor-truncation
    assert rounded_hidden % OCP_MX_BLOCK_SIZE == 0
    assert rounded_intermediate % OCP_MX_BLOCK_SIZE == 0

评论区精华

1. 对齐逻辑位置： BowenBao 和 fxmarty-amd 在 review 中指出 emulation 的对齐逻辑应统一放入共享函数，而非留在 quark_moe.py。作者采纳，为 EMULATION 添加显式分支并移除排除条件。讨论后合并。（设计）
2. 测试平台限制： tjtanaa 询问测试是否应限于 ROCm，作者确认 emulation 后端当前仅在 ROCm 使用，追加 @pytest.mark.skipif(not ROCM_AVAILABLE)。（测试）

• Emulation对齐逻辑应统一放入共享helper (design): 作者为EMULATION添加显式分支，并移除quark_moe.py中的排除条件，所有后端roundup统一在共享函数中处理。
• 测试是否应限于ROCm平台 (testing): 作者确认emulation后端目前仅用于ROCm，因此加入@pytest.mark.skipif(not ROCM_AVAILABLE)限定。

风险与影响

• 风险：
  • 回归风险：修改了 MoE 对齐函数的控制流，但添加的是新分支，不影响现有后端，风险较低。
  • 测试覆盖局限：新测试仅限 ROCm 运行，若未来 emulation 后端在其他平台启用，可能遗漏。
  • 依赖风险：修复依赖 OCP_MX_BLOCK_SIZE 常量的正确性。
• 影响：
  • 用户影响：修复了 ROCm 用户使用 --moe-backend emulation 加载 Quark MXFP4 量化 GPT-OSS 模型时的阻断性错误。
  • 系统影响：改动局限在维度对齐路径，不涉及 kernel 或运行时。
  • 团队影响：代码结构改进有助于后续维护。
  • 风险标记：权重加载路径变更, 测试限ROCm, 回归风险

关联脉络

• 暂无明显关联 PR