执行摘要

• 一句话：修复 FusedMoE 加载 shape-(1,) per-tensor scale 崩溃
• 推荐动作：值得精读：该 PR 虽然改动小，但展示了处理上下游数据格式不一致的典型模式：提取统一转换函数、扩展至所有潜在调用点、添加针对性的回归测试。设计决策如使用 reshape(()) 而非 squeeze() 或 view([]) 也体现了对维度安全的考虑（numel > 1 时失败）。推荐作为技术参考。

功能与动机

Issue #43297 报告：llm-compressor NVFP4 预设产生的 per-tensor weight/input scale 为 shape-(1,) 张量，而 FusedMoE 权重加载器假定其为 0-D 标量，导致 RuntimeError。官方 DeepSeek-V4-Flash 等模型有约 66,000 个 affected tensors，server 完全不可用。本 PR 修复确保这些模型能正常加载。

实现拆解

1. 在 RoutedExperts 类中新增 _to_scalar 静态方法，内部调用 loaded_weight.reshape(()) 将 shape-(1,) 或 0-D 张量坍缩为 0-D 标量，并在 numel > 1 时抛出 RuntimeError。
2. 修改 _load_per_tensor_weight_scale 方法：在 w1/w3 和 w2 两个分支的赋值前调用 _to_scalar，替换直接赋值。
3. 修改 _load_single_value 方法：同样在赋值前调用 _to_scalar，覆盖 input_scale 等 scalar 加载路径。
4. 将 weight_loader 中 ModelOpt 路径的 loaded_weight.reshape(()) 替换为 _to_scalar，统一 coercion 入口。
5. 在 tests/kernels/moe/test_moe_weight_loading_padded.py 中添加 TestPerTensorScaleCoercion 测试类，包含两个用例：验证 shape-(1,) 和 0-D 张量都能正确降维，验证 numel > 1 的张量会引发 RuntimeError。

关键文件：

• vllm/model_executor/layers/fused_moe/routed_experts.py（模块 权重加载；类别 source；类型 core-logic；符号 _to_scalar, _load_per_tensor_weight_scale, _load_single_value）: 核心修复文件，新增 _to_scalar 方法和修改三个调用点以确保 shape-(1,) 的 per-tensor scale 能被正确加载。
• tests/kernels/moe/test_moe_weight_loading_padded.py（模块 MoE测试；类别 test；类型 test-coverage；符号 TestPerTensorScaleCoercion, test_collapses_to_scalar, test_rejects_non_scalar）: 新增 TestPerTensorScaleCoercion 测试类，回归验证 _to_scalar 能正确处理 shape-(1,) 和 0-D 标量，且对 numel > 1 时会报错。

关键符号：_to_scalar, _load_per_tensor_weight_scale, _load_single_value

关键源码片段

vllm/model_executor/layers/fused_moe/routed_experts.py

核心修复文件，新增 _to_scalar 方法和修改三个调用点以确保 shape-(1,) 的 per-tensor scale 能被正确加载。

@staticmethod
def _to_scalar(loaded_weight: torch.Tensor) -> torch.Tensor:
    # 处理 per-tensor scale 形状不一致的问题
    # llm-compressor NVFP4 预设输出 shape-(1,) 而代码预期 0-D 标量
    # `reshape(())` 将任何单元素张量坍缩为 0-D 标量
    # 如果 numel > 1 则会抛出 RuntimeError 避免静默广播
    return loaded_weight.reshape(())
​
def _load_per_tensor_weight_scale(
    self,
    shard_id: str,
    param: torch.nn.Parameter,
    loaded_weight: torch.Tensor,
    expert_id: int,
):
    param_data = param.data
    if shard_id in ('w1', 'w3'):
        # w1 和 w3 共享同一个参数 需要保留两个 scale 以便后续重量化
        idx = 0 if shard_id == 'w1' else 1
        param_data[expert_id][idx] = self._to_scalar(loaded_weight)
    elif shard_id == 'w2':
        # down_proj 对应的 w2 直接赋值
        param_data[expert_id] = self._to_scalar(loaded_weight)

tests/kernels/moe/test_moe_weight_loading_padded.py

新增 TestPerTensorScaleCoercion 测试类，回归验证 _to_scalar 能正确处理 shape-(1,) 和 0-D 标量，且对 numel > 1 时会报错。

class TestPerTensorScaleCoercion:
    """Regression test for shape-(1,) per-tensor scales (issue #43297)."""
​
    def test_collapses_to_scalar(self):
        # 验证 shape-(1,) 和 0-D 张量都能正确变为 0-D 标量 且数值正确
        for loaded_weight in (torch.tensor([0.5]), torch.tensor(0.5)):
            scalar = RoutedExperts._to_scalar(loaded_weight)
            assert scalar.shape == ()
            assert scalar.item() == pytest.approx(0.5)
​
    def test_rejects_non_scalar(self):
        # 验证 numel > 1 的张量会抛出 RuntimeError 不会静默广播
        with pytest.raises(RuntimeError):
            RoutedExperts._to_scalar(torch.tensor([0.1, 0.2]))

评论区精华

reviewer gemini-code-assist[bot] 指出 _load_single_value 也存在同样问题，作者随后扩展了修复。reviewer mgoin 建议将 coercion 提取为统一 helper 并与 ModelOpt 的现有 reshape(()) 合并，作者实现了 _to_scalar 并在三处调用点复用。mgoin 认为测试有些多余但无伤大雅，作者将测试精简聚焦于新 helper。最终 mgoin 给出 APPROVED。

• 扩展 _load_single_value 修复 (correctness): 作者接受建议，在 _load_single_value 的赋值前也调用了 _to_scalar。
• 统一 scalar 降维 helper (design): 作者实现 _to_scalar 静态方法，并在 _load_per_tensor_weight_scale、_load_single_value 和 ModelOpt 路径中统一使用。
• 测试精简 (testing): 作者将测试简化，聚焦于 _to_scalar helper 的两种场景，去掉多余重复。

风险与影响

• 风险：变更范围集中在 MoE 权重加载的 scalar 赋值路径，使用 reshape(()) 降维：
  • 若 loaded_weight 元素数为 1，降维操作无拷贝、无副作用。
  • 若元素数 > 1，reshape(()) 会抛出 RuntimeError，比原行为（广播导致静默错误）更安全。
  • 风险极低：所有调用均已覆盖，不影响非 scalar 权重加载流程。
  • 潜在回归：如果其他量化工具提供形状不符的 scalar（如 shape-(2,)）但以前被广播容忍，现在会报错。但这属于正确的行为变化，不应视为负回归。
  • 影响：
    • 用户：使用 NVFP4 量化和 MoE 模型的用户 server 可以正常启动，不再崩溃。
    • 系统：权重加载路径增加极小的 reshape 开销，对运行时性能无影响。
    • 团队：统一了 scalar 降维入口，便于后续维护和扩展。
    • 影响范围：直接修复约 66,000 个 tensor 的加载失败，间接为其他可能产生 shape-(1,) 张量的量化工具提供了兼容性。
    • 风险标记：核心路径变更, 窄修复低风险

关联脉络

• 暂无明显关联 PR