Prhub

#43261 [Bug] Fix ci issue `assert output_size is not None` AssertionError

原始 PR 作者 yewentao256 合并时间 2026-05-21 16:58 文件变更 5 提交数 3 评论 11 代码增减 +34 / -3

bugfix fp8 nvidia kernel

执行摘要

修复 FP8 线性层 padding 后 weight_loader 缺失导致的 CI 断言错误

参考 PR #43261 的 body：该 PR 修复了 pytest tests/v1/spec_decode/test_speculators_correctness.py::test_speculators_model[dflash] 因 AssertionError: assert output_size is not None 导致的失败。该断言位于 vllm/model_executor/kernels/linear/scaled_mm/cutlass.py:227 apply_scaled_mm 方法中，原因是 self.logical_output_size 未在调用前被正确赋值，根因在于权重填充后未通过 process_weights_after_loading 同步更新 weight_loader 及 logical_output_size。

该 PR 是典型的 bug 修复与长期可维护性改进的结合。值得阅读 cutlass.py 中 padded_weight_loader 的设计——它展示了如何处理参数张量由于 padding 导致的形状不兼容，并保持加载器可重入。同时关注量化方法与内核后处理之间的调用拓扑，确保嵌套调用正确。建议在后续类似变更中延续此模式。

讨论亮点

haosdent 在 review 中提出：「Should we also patch online/fp8.py?」——该建议被采纳，最终 online/fp8.py 也加入了相同调用。
haosdent 指出该变更可能无法解决 AttributeError: 'CutlassFP8ScaledMMLinearKernel' object has no attribute 'logical_output_size'；Isotr0py 回应：「It's fixed by https://github.com/vllm-project/vllm/pull/43268 instead」，即该错误由另一 PR 单独修复。
haosdent 进一步质疑：「Seems even #43261 + #43268 could not handle cases like test_fp8_reloading」；Isotr0py 解释 reloading 测试通过新增的 padded_weight_loader 得到修复。

实现拆解

  1. 新增 padded_weight_loader 静态方法（cutlass.py）：当加载的权重形状与参数形状不一致时（因padding导致尺寸变大），通过切片拷贝实现部分更新，确保参数张量的padding区域不被覆盖。
  2. 在 process_weights_after_loading 中设置自定义加载器（cutlass.py）：在对权重和权重尺度完成 padding 后，使用 set_weight_attrs 将 padded_weight_loader 注册为参数的 weight_loader 属性。这保证了后续权重重载（如 test_fp8_reloading）时能正确处理已 padding 的参数。
  3. 确保内核后处理被上游量化方法调用（online/fp8.py 和 fp8.py）：在 process_weights_after_loading 方法末尾添加 self.fp8_linear.process_weights_after_loading(layer)，使得 CutlassFP8ScaledMMLinearKernel 的 padding 等后处理逻辑能被执行，从而初始化 logical_output_size。
  4. 同步更新测试：
    • test_cutlass_scaled_mm.py：在 test_cutlass_fp8_gemm_padded 中模拟了 process_weights_after_loading 的 padding 行为并设置 kernel.logical_output_size，使单测通过。
    • test_fp8.py：在 test_fp8_reloading 中移除了对 weight_loader 与原加载器是同一对象的断言，因为现在权重可能挂载自定义 padded_weight_loader。
文件 模块 状态 重要度
vllm/model_executor/kernels/linear/scaled_mm/cutlass.py 量化内核 modified 6.76
vllm/model_executor/layers/quantization/online/fp8.py 量化层 modified 4.53
vllm/model_executor/layers/quantization/fp8.py 量化层 modified 4.39
tests/kernels/quantization/test_cutlass_scaled_mm.py CUTLASS 测试 modified 4.41
tests/quantization/test_fp8.py FP8 测试 modified 3.43

关键符号

padded_weight_loader process_weights_after_loading (cutlass.py) process_weights_after_loading (online/fp8.py) process_weights_after_loading (fp8.py)

关键源码片段

vllm/model_executor/kernels/linear/scaled_mm/cutlass.py data-contract

核心修复文件：新增 `padded_weight_loader` 并在填充后设置 `weight_loader`，解决 `output_size` 未初始化问题。

# vllm/model_executor/kernels/linear/scaled_mm/cutlass.py
​
@staticmethod
def padded_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
    """
    自定义权重加载器，处理 padding 后张量形状不匹配的情况。
    当加载的权重形状与参数形状不一致时（常见于权重被 padding 至 16 对齐），
    使用切片拷贝仅覆盖逻辑区域，保留 padding 区域。
    """
    if loaded_weight.shape != param.shape:
        # 计算切片，仅拷贝 loaded_weight 形状对应的部分
        slices = tuple(slice(0, s) for s in loaded_weight.shape)
        param.data[slices].copy_(loaded_weight)
    else:
        # 形状一致时直接拷贝
        param.data.copy_(loaded_weight.view(param.shape))
​
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
    weight_name, weight_scale_name, _, _ = self.layer_param_names
    weight = getattr(layer, weight_name)
    self.logical_output_size = weight.shape[1]
    pad_k = (16 - weight.shape[0] % 16) % 16
    pad_n = (16 - weight.shape[1] % 16) % 16
    if pad_k == 0 and pad_n == 0:
        return
    padded_weight = torch.nn.functional.pad(
        weight.t().contiguous(),
        (0, pad_k, 0, pad_n),
    ).t()
    replace_parameter(layer, weight_name, padded_weight.data)
    set_weight_attrs(
        getattr(layer, weight_name),
        {"weight_loader": self.padded_weight_loader},
    )
    weight_scale = getattr(layer, weight_scale_name, None)
    if weight_scale is not None and pad_n > 0 and weight_scale.numel() > 1:
        flat_scale = weight_scale.reshape(-1)
        padded_scale = self._pad_to_alignment(
            flat_scale, dim=0, alignment=16, value=1.0
        ).view(-1, *weight_scale.shape[1:])
        replace_parameter(layer, weight_scale_name, padded_scale.data)
        set_weight_attrs(
            getattr(layer, weight_name),
            {"weight_loader": self.padded_weight_loader},
        )
vllm/model_executor/layers/quantization/online/fp8.py core-logic

确保 FP8 在线量化方法在权重后处理中调用内核的 process_weights_after_loading，使 padding 和 weight_loader 设置生效

# vllm/model_executor/layers/quantization/online/fp8.py
​
def process_weights_after_loading(self, layer: Module) -> None:
    if getattr(layer, "_already_called_process_weights_after_loading", False):
        return
    layer.input_scale = None
    qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None)
    replace_parameter(layer, "weight", qweight.t().data)
    replace_parameter(layer, "weight_scale", weight_scale.data)
    # 调用内核层的后处理（padding + weight_loader 设置）
    self.fp8_linear.process_weights_after_loading(layer)
    layer._already_called_process_weights_after_loading = True

评论区精华

需要同步修改 online/fp8.py 设计

haosdent 在 review 中建议：'Should we also patch online/fp8.py?'

结论：Isotr0py 采纳建议，在后续 commit 中添加了对应修改。 · 已解决

AttributeError: logical_output_size 未初始化 正确性

haosdent 指出当前修改无法解决 `AttributeError: 'CutlassFP8ScaledMMLinearKernel' object has no attribute 'logical_output_size'`；Isotr0py 回复该问题由 #43268 独立修复。

结论：该问题由另一 PR 修复，本 PR 不直接处理。 · 已解决

test_fp8_reloading 用例的兼容性 测试

haosdent 质疑即使合并 #43261 和 #43268 也无法处理 `test_fp8_reloading`；Isotr0py 解释通过引入 `padded_weight_loader` 确保重载时权重正确拷贝。

结论：padded_weight_loader 解决了重载测试。 · 已解决

风险与影响

padded_weight_loader 仅在参数形状不匹配时通过切片拷贝加载，若其他加载场景（如分片加载）也触发该逻辑，可能出现尺寸判断错误。process_weights_after_loading 被调用两次的防护已存在（_already_called_process_weights_after_loading），但若 guard 失效，可能导致重复 padding 或设置。自定义 weight_loader 替换原有加载器，可能影响 checkpoint 加载的兼容性，特别是当模型使用其他量化方法或非 FP8 时，本变更不会触发。测试覆盖限于单测和特定 CI 场景，生产环境中非 16 对齐的模型（如 Qwen2.5-VL）的真实行为缺乏全面验证。

用户：无直接交互变更，但修复了 CI 失败，保障了代码质量。使用 FP8 量化和 CUTLASS 内核的模型（特别是需 padding 的非对齐维度）会正确初始化 logical_output_size，避免运行时断言错误。系统：权重加载后处理流程更健壮，自定义加载器机制可复用。团队：减少了维护负担，统一了量化层与内核层后处理的调用模式。

核心路径变更 权重加载兼容性

关联 Issue

未识别关联 Issue

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

完整报告

Markdown PDF

参与讨论