Prhub

#51249 [Bugfix][Model] Add missing fused_qkv_a_proj to Kimi-Linear packed_modules_mapping

原始 PR 作者 JianDan0212 合并时间 2026-08-06 19:39 文件变更 1 提交数 1 评论 3 代码增减 +1 / -0

执行摘要

补全 Kimi-Linear 融合映射,修复注意力层误量化

作者在测试 Kimi-K3 W4A8(RTN)量化 checkpoint 时,发现若干 Linear 层被有意保留 bf16 不量化,但其中一个 bf16 的注意力投影仍被当作量化层处理。根因是 #50500 引入的 KimiLinearModel.packed_modules_mapping 缺少 fused_qkv_a_proj 条目,导致 compressed-tensors 的 should_ignore_layer() 无法把融合层名映射回未融合的 checkpoint 名称(q_a_proj / kv_a_proj_with_mqa),最终这些本应豁免的层被默认量化目标选中。

该 PR 本体只有 1 行配置修复,不值得单独精读;但若你参与 compressed-tensors 集成或 Kimi 系列模型维护,值得了解 packed_modules_mapping 与 should_ignore_layer() 的联动机制,以及后续是否为此类映射补充单元测试以覆盖回归。

讨论亮点

该 PR 没有产生人工 review 评论。copilot-pull-request-reviewer 自动评审确认修复方向是“教会 packed_modules_mapping 如何将 fused_qkv_a_proj 映射回原始 checkpoint(未融合)名称”;作者在 issue 评论中说明这是 #50500 的小型后续,并请 @kylesayrs @mgoin 复核;维护者 jeejeelee 触发 CI 后直接批准合并。

实现拆解

  1. 定位触发链路:Kimi-K3 的 MLA 注意力在 config 设置 q_lora_rank 时,会将 q_a_proj 与 kv_a_proj_with_mqa 融合为 fused_qkv_a_proj(见 mla.py 的 MultiHeadLatentAttention)。compressed-tensors 依赖 packed_modules_mapping 反查 checkpoint 中的层名,缺失该融合条目导致 ignore 规则失效。
  2. 修改映射表:在 vllm/models/kimi_k3/nvidia/model.py 的 KimiLinearModel.packed_modules_mapping 字典中新增 "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],与 mla.py 中的融合约定对齐。
  3. 验证方式:作者使用真实 Kimi-K3 W4A8(RTN)checkpoint 验证 should_ignore_layer() 现在能正确将 fused_qkv_a_proj 解析为“ignored”(bf16),而修复前会落入默认量化 Linear 目标并导致内核选择失败。
  4. 配套情况:本 PR 未附带单元测试,依赖真实 checkpoint 的集成验证和 CI(Buildkite #82670)覆盖;建议后续为 packed_modules_mapping 与 ignore 规则联动补充测试。
文件 模块 状态 重要度
vllm/models/kimi_k3/nvidia/model.py Kimi 模型 modified 4.16

关键符号

packed_modules_mapping

关键源码片段

vllm/models/kimi_k3/nvidia/model.py data-contract

PR 唯一改动文件,KimiLinearModel.packed_modules_mapping 新增 fused_qkv_a_proj 条目,修复量化 checkpoint 中 MLA 注意力层 ignore 规则解析失效的问题。

# vllm/models/kimi_k3/nvidia/model.py
​
class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
    # 融合参数名 -> 未融合 checkpoint 参数名的映射表,
    # 供 compressed-tensors 反向解析 ignore 规则中的层名。
    # 缺少 fused_qkv_a_proj 时,量化 checkpoint 中显式豁免
    # 的 q_a_proj / kv_a_proj_with_mqa 无法被识别为“忽略”,
    # 导致 MLA 注意力层被误量化或内核选择失败。
    packed_modules_mapping = {
        # MoE / dense MLP 的 gate 与 up 投影融合
        "gate_up_proj": ["gate_proj", "up_proj"],
        # KDA 线性注意力层的前置投影融合(q / k / v / b / f_a)
        "in_proj_qkvgfab": ["q_proj", "k_proj", "v_proj", "b_proj", "f_a_proj"],
        # 一维卷积融合(q / k / v 三个 conv1d)
        "conv1d": ["q_conv1d", "k_conv1d", "v_conv1d"],
        # MLA 注意力在 q_lora_rank 设置时的 q / kv 融合(本次新增)
        "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
    }

评论区精华

复核与 CI 触发 question

作者在 issue 评论中 @kylesayrs @mgoin 请求复核,并说明这是 #50500 的小型后续;copilot 自动评审概述了变更影响;维护者 jeejeelee 触发 CI 后批准。

结论:jeejeelee 批准合并,CI #82670 已触发;无未解决技术争议。 · 已解决

风险与影响

风险很低,但存在两个关注点:一是无配套单元测试,packed_modules_mapping 是压缩张量集成的数据契约,若后续 mla.py 调整融合命名,需要同步维护该映射;二是本 PR 只修改了 vllm/models/kimi_k3/nvidia/model.py,vllm/models/kimi_k3/amd/ 路径下若存在类似的 KimiLinearModel 定义,需确认是否有相同问题。此外,新增条目会改变 fused_qkv_a_proj 在量化目标选择中的默认行为,但修复后更符合 checkpoint 的显式豁免意图,回归风险可控。

影响范围限定在启用 q_lora_rank 的 Kimi-K3 量化 checkpoint(如 W4A8/RTN)加载路径:修复后 bf16 豁免按 checkpoint 意图生效,避免 MLA 注意力层精度退化与内核选择失败;对未使用量化或未显式豁免这些层的用户无行为变化。系统层面仅 +1 行配置映射,无性能与安全影响。该 PR 被标记于 v0.27.0 cherry picks,预计随版本发布。

缺少测试覆盖 量化路径回归风险 数据契约变更

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论