执行摘要

• 一句话：统一 KV cache scale 加载，移除模型代码重复
• 推荐动作：建议合并。该 PR 设计清晰，自解释性强，是典型的“移除重复、统一入口”重构。值得关注的决策：通过 WeightsMapper 的 | 操作符合并多个映射规则，以及使用 KVCacheScaleParameter 作为标量的唯一容器。测试覆盖方面，由于删除了一些 Mock 测试，建议补充端到端集成测试（如加载含 KV scale 的 FP8 checkpoint）。

功能与动机

每个模型的 load_weights 方法都包含约 10 行相似的样板代码，用于从 checkpoint 中加载 KV cache 的缩放因子（q_scale/k_scale/v_scale）。这种重复不仅增加了维护成本，而且新模型很容易遗漏该逻辑。同时，Transformers 模型后端之前完全无法加载这些缩放因子。该变更旨在消除样板、统一入口，并让所有模型自动获得该能力。

实现拆解

1. 引入 KVCacheScaleParameter（vllm/model_executor/layers/quantization/kv_cache.py）：继承 torch.nn.Parameter，初始化为 -1.0（无效哨兵值）。内置 weight_loader 静态方法，只接受形状为 () 或 (1,) 的标量，并将 reshape 为标量后赋值，避免了每个模型手动 flatten 和标量判断。
2. 将 get_cache_scale 改为 get_cache_scale_mapper（fp8.py、quark.py、base_config.py）：原方法按名称拼接字符串，现改为返回一个 WeightsMapper 对象（名称后缀映射表），由 AutoWeightsLoader 统一应用。所有量化配置类均通过 get_cache_scale_mapper 提供映射规则。
3. 在 AutoWeightsLoader.load_weights 中集成 mapper（vllm/model_executor/models/utils.py）：从 self.module.quant_config 获取 cache_scale_mapper，若存在则与已有的 mapper 通过 | 操作符合并，随后在加载过程中自动应用。这样任何使用 AutoWeightsLoader 的模型都无需额外代码。
4. 删除所有模型中的手动 KV cache scale 加载逻辑：涉及 50 多个模型文件（如 gpt_oss.py、llama4.py、glm_ocr_mtp.py、mimo_v2.py 等），每个文件移除了约 10-15 行的 if scale_name 判断、weight_loader 调用和 loaded_params 更新。这些逻辑现在由第 3 步统一处理。
5. 清理测试（tests/model_executor/test_eagle_quantization.py）：删除了三个测试函数，它们仅 Mock 了 get_cache_scale 但未实际依赖 vLLM 内部逻辑，属于无效测试。

关键文件：

• vllm/model_executor/layers/quantization/kv_cache.py（模块 量化层；类别 source；类型 data-contract；符号 KVCacheScaleParameter, new, weight_loader）: 新增 KVCacheScaleParameter，封装标量参数行为，为所有 Attention 层提供统一的 KV cache scale 参数类。同时修改 create_weights 使用该类实例，实现自动标量初始化。
• vllm/model_executor/models/utils.py（模块 加载工具；类别 source；类型 data-contract；符号 load_weights）: 在 AutoWeightsLoader.load_weights 中集成 KV cache scale mapper：从 quant_config 获取 mapper，并与已有的 mapper 合并。这是整个重构的中心调度点，确保所有使用该加载器的模型自动获得 scale 重映射。
• vllm/model_executor/layers/quantization/fp8.py（模块 量化层；类别 source；类型 data-contract；符号 get_cache_scale, get_cache_scale_mapper）: 将原 get_cache_scale（字符串映射）转换为 get_cache_scale_mapper（返回 WeightsMapper）。此更改是所有量化后端的模板，演示从临时字符串构造到声明式字典的演进。

关键符号：KVCacheScaleParameter.new, KVCacheScaleParameter.weight_loader, BaseKVCacheMethod.create_weights, Fp8KVCacheMethod.get_cache_scale_mapper, QuarkConfig.get_cache_scale_mapper, AutoWeightsLoader.load_weights

关键源码片段

vllm/model_executor/layers/quantization/kv_cache.py

新增 KVCacheScaleParameter，封装标量参数行为，为所有 Attention 层提供统一的 KV cache scale 参数类。同时修改 create_weights 使用该类实例，实现自动标量初始化。

# vllm/model_executor/layers/quantization/kv_cache.py
import torch
from vllm.v1.kv_cache_interface import kv_cache_uses_per_token_head_scales
​
class KVCacheScaleParameter(torch.nn.Parameter):
    """专用于 KV cache 缩放因子的标量参数。
    初始化为 -1.0（无效值），后续由 checkpoint 中加载的权重覆盖。
    `weight_loader` 只接受形状为 () 或 (1,) 的标量，其他形状会触发错误。
    此路径不处理 per-head 的动态 scale（那部分走 compressed-tensors 的 `_tp_aware_loader`）。
    """
​
    def __new__(cls) -> "KVCacheScaleParameter":
        # 调用父类 `Parameter.__new__` 并传入默认张量 -1.0
        return super().__new__(cls, torch.tensor(-1.0), requires_grad=False)
​
    @staticmethod
    def weight_loader(param: torch.nn.Parameter, loaded_weight: torch.Tensor) -> None:
        """加载时确保权重为标量，并将 data 用 reshape 为标量赋值。"""
        if loaded_weight.numel() != 1:
            raise ValueError(
                f"KV-cache scale expects a scalar weight, got shape "
                f"{tuple(loaded_weight.shape)}"
            )
        param.data.copy_(loaded_weight.reshape(()))
​
​
class BaseKVCacheMethod(QuantizeMethodBase):
    def create_weights(self, layer: torch.nn.Module):
        # 使用 KVCacheScaleParameter 替代原始的 torch.nn.Parameter 构造
        layer.q_scale = KVCacheScaleParameter()
        layer.k_scale = KVCacheScaleParameter()
        layer.v_scale = KVCacheScaleParameter()
        layer.prob_scale = KVCacheScaleParameter()
​
    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
        # 略去初始化校验，核心逻辑不变
        ...
        if kv_cache_uses_per_token_head_scales(layer.kv_cache_dtype):
            # 动态 scale 场景，清除哨兵参数并设为 1.0
            ...
            del layer.k_scale

vllm/model_executor/models/utils.py

在 AutoWeightsLoader.load_weights 中集成 KV cache scale mapper：从 quant_config 获取 mapper，并与已有的 mapper 合并。这是整个重构的中心调度点，确保所有使用该加载器的模型自动获得 scale 重映射。

# vllm/model_executor/models/utils.py
class AutoWeightsLoader:
    # ...
    def load_weights(
        self,
        weights: Iterable[Tuple[str, torch.Tensor]],
        *,
        mapper: Optional[WeightsMapper] = None,
    ) -> Set[str]:
        # 从模块的 quant_config 获取 cache scale mapper
        quant_config = getattr(self.module, "quant_config", None)
        cache_scale_mapper = (
            quant_config.get_cache_scale_mapper()
            if quant_config is not None else None
        )
        if cache_scale_mapper is not None:
            # 与已有的 mapper 合并（WeightsMapper 实现了 __or__）
            mapper = mapper | cache_scale_mapper if mapper is not None else cache_scale_mapper
        # 之后继续原有的加载流程，mapper 会应用于所有权重名称
        # ...

vllm/model_executor/layers/quantization/fp8.py

将原 get_cache_scale（字符串映射）转换为 get_cache_scale_mapper（返回 WeightsMapper）。此更改是所有量化后端的模板，演示从临时字符串构造到声明式字典的演进。

# vllm/model_executor/layers/quantization/fp8.py
from vllm.model_executor.models.utils import WeightsMapper
​
class Fp8KVCacheMethod(BaseKVCacheMethod):
    # ...
    def get_cache_scale_mapper(self) -> "WeightsMapper":
        """返回从 compressed-tensors 格式到 vLLM 格式的 KV cache scale 名称映射。"""
        # 使用 WeightsMapper 的 orig_to_new_suffix 参数声明后缀替换规则
        return WeightsMapper(
            orig_to_new_suffix={
                ".k_proj.output_scale": ".attn.k_scale",
                ".v_proj.output_scale": ".attn.v_scale",
                ".q_proj.output_scale": ".attn.q_scale",
                ".self_attn.prob_output_scale": ".self_attn.attn.prob_scale",
            }
        )

评论区精华

1. gemini-code-assist[bot] 指出在 utils.py 中通过 | 合并 WeightsMapper 实例可能引发 TypeError，除非自定义类实现了 __or__。作者 hmellor 回应 WeightsMapper 确实实现了 __or__ 方法（在另一处定义中），因此该担忧不成立。
2. gemini-code-assist[bot] 质疑移除 commandr.py 中的 scale 加载逻辑会导致该模型无法加载 KV cache scale，因为该模型可能使用自定义加载循环。hmellor 澄清 CohereForCausalLM 实际上依赖 AutoWeightsLoader，而 commandr.py 中的类是 AutoWeightsLoader 的组成部分，因此映射逻辑仍在顶层生效。

• WeightsMapper | 操作符是否可用 (design): hmellor 回复 WeightsMapper 确实实现了 __or__ 方法，因此合并安全。
• commandr.py 中移除 scale 加载是否破坏功能 (correctness): hmellor 澄清 CohereForCausalLM 使用 AutoWeightsLoader，而 commandr.py 类是 AutoWeightsLoader 的一部分，因此全局 mapper 仍然生效。

风险与影响

• 风险：

  1. WeightsMapper 合并兼容性：若未来有自定义 WeightsMapper 子类未实现 __or__，合并时会崩溃。当前默认实现支持该操作，但需注意扩展。
  2. 模型未使用 AutoWeightsLoader：极少数模型可能直接手写加载循环而绕过了顶层 mapper，这些模型将丢失 KV cache scale 映射（例如某些 legacy 模型）。需逐一确认所有模型都通过 AutoWeightsLoader 加载权重。
  3. 测试覆盖减少：删除了三个测试，虽然它们仅测试 Mock 逻辑，但可能掩盖回归。建议新增集成测试，验证实际 checkpoint 中 KV cache scale 的加载。
  4. 形状假设破坏：KVCacheScaleParameter.weight_loader 强制标量检查，若未来出现非标量 scale（如 per-head 动态 scale），该 loader 会拒绝。但此类 scale 走的是另一条路径（compressed-tensors 的 _tp_aware_loader），因此风险可控。
     - 影响：用户：无直接影响，所有行为应保持兼容。模型加载故障可能表现为 KV cache 量化异常，但回归风险较低。
     系统：减少了约 730 行代码，降低了模型 load_weights 的重复度和出错概率。新模型加入时自动获得 KV cache scale 加载能力。
     团队：重构提升了代码一致性，方便后续量化方案扩展。需注意后续模型开发时不应再手写 scale 加载逻辑。

• 风险标记：核心路径变更, 测试覆盖降低, 兼容性风险

关联脉络

• PR #43567 : Reviewer Isotr0py 提议用此 PR 替代 #43567，说明两者解决同一问题或有重叠。