Prhub

#50533 [Bugfix][TurboQuant] Add KV quant mode for turboquant

原始 PR 作者 skavulya 合并时间 2026-08-01 01:42 文件变更 4 提交数 2 评论 0 代码增减 +32 / -2

执行摘要

修复 TurboQuant KV 缓存 dtype 误为 auto 崩溃

PR body 指出:TurboQuant KV cache 的 dtype 字符串会被错误地设置为 auto,导致 vllm serve Qwen/Qwen3-30B-A3B-Thinking-2507 --kv-cache-dtype turboquant_4bit_nc 崩溃,堆栈显示 TurboQuantConfig.from_cache_dtype(cache_dtype_str='auto') 抛出 Unknown TurboQuant cache dtype。根因是 TQFullAttentionSpec 构造时缺少 kv_quant_mode,默认落入 KVQuantMode.NONE,使 v1 worker 的 KV 缓存整形逻辑误判为未量化层。

值得精读,特别是 vllm/v1/kv_cache_interface.py 的 KVQuantMode 映射与 attention.py 的 get_kv_cache_spec 分支:它展示了如何通过枚举统一量化模式、避免字符串特判的接入方式。建议结合 PR body 引用的前序 TurboQuant 工作,理解 v1 KV 缓存规格化的演进脉络。

讨论亮点

本 PR 没有实质性的技术讨论:claude[bot] 自动审查提示该 PR 来自 fork 且默认禁用自动 review,维护者 mgoin 直接批准合并。缺少 review 交锋,因此依赖测试与代码自身的清晰度;后续若扩展 TurboQuant KV cache,建议补一次人工 review。

实现拆解

  1. 在 vllm/v1/kv_cache_interface.py 中为 KVQuantMode 枚举新增 TURBOQUANT = 6,并添加 is_turboquant 属性;get_kv_quant_mode 将任何以 turboquant_ 开头的 kv_cache_dtype 映射为 KVQuantMode.TURBOQUANT,这是整个修复的基础。
  2. 在 vllm/model_executor/layers/attention/attention.py 的 Attention.get_kv_cache_spec 中,构造 TQFullAttentionSpec 时补传 kv_quant_mode=quant_mode,使 KV cache spec 携带正确的量化模式标识。
  3. 在 vllm/v1/worker/gpu/attn_utils.py 的 _reshape_kv_cache 中删除对 TQFullAttentionSpec 的 import 与 isinstance 特判:layer_cache_dtype 现在完全由 kv_quant_mode != KVQuantMode.NONE 决定,turboquant 层会获得真实 cache_dtype 而不是 auto。
  4. 在 tests/quantization/test_turboquant.py 中新增 TestTurboQuantKVCacheSpec 类,用 ALL_PRESETS 参数化验证 Attention.get_kv_cache_spec 产物为 TQFullAttentionSpec 且 kv_quant_mode 为 TURBOQUANT。
    无配置、部署或文档配套改动。
文件 模块 状态 重要度
vllm/v1/kv_cache_interface.py 缓存接口 modified 6.23
vllm/model_executor/layers/attention/attention.py 注意力层 modified 4.96
vllm/v1/worker/gpu/attn_utils.py 缓存整形 modified 4.32
tests/quantization/test_turboquant.py 量化测试 modified 5.56

关键符号

is_turboquant get_kv_quant_mode Attention.get_kv_cache_spec _reshape_kv_cache test_kv_cache_spec_sets_kv_quant_mode

关键源码片段

vllm/v1/kv_cache_interface.py core-logic

核心修改文件:新增 KVQuantMode.TURBOQUANT 枚举、is_turboquant 属性,并在 get_kv_quant_mode 中按 turboquant_ 前缀完成映射,是修复的基础。

class KVQuantMode(IntEnum):
    """KV 缓存量化模式,attention 后端与内核据此分发量化逻辑,
    避免对 ``kv_cache_dtype`` 字符串做硬编码匹配。"""
​
    NONE = 0
    FP8_PER_TENSOR = 1 # per-tensor scales(当前 fp8 路径)
    INT8_PER_TOKEN_HEAD = 2 # per-token-head 动态 scale(int8)
    FP8_PER_TOKEN_HEAD = 3 # per-token-head 动态 scale(fp8)
    INT4_PER_TOKEN_HEAD = 4 # packed 2×int4/byte,RHT + 非对称 zero point
    NVFP4 = 5 # packed fp4 数据 + fp8 block scales
    # 新增:Hadamard 旋转 Lloyd-Max 量化,每个 slot 打包 K+V
    TURBOQUANT = 6
​
    @property
    def is_per_token_head(self) -> bool:
        """是否为 per-token-head 量化模式。"""
        return self in (
            KVQuantMode.INT8_PER_TOKEN_HEAD,
            KVQuantMode.FP8_PER_TOKEN_HEAD,
            KVQuantMode.INT4_PER_TOKEN_HEAD,
        )
​
    @property
    def is_nvfp4(self) -> bool:
        """是否为 NVFP4 打包量化模式。"""
        return self == KVQuantMode.NVFP4
​
    @property
    def is_turboquant(self) -> bool:
        """是否为 TurboQuant 量化模式。"""
        return self == KVQuantMode.TURBOQUANT
​
​
def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
    """把 ``kv_cache_dtype`` 字符串映射为 KVQuantMode。"""
    if kv_cache_dtype == "int4_per_token_head":
        return KVQuantMode.INT4_PER_TOKEN_HEAD
    if kv_cache_dtype == "int8_per_token_head":
        return KVQuantMode.INT8_PER_TOKEN_HEAD
    if kv_cache_dtype == "fp8_per_token_head":
        return KVQuantMode.FP8_PER_TOKEN_HEAD
    if kv_cache_dtype == "nvfp4":
        return KVQuantMode.NVFP4
    # 新增:所有 turboquant 预设(如 turboquant_4bit_nc)统一归为 TURBOQUANT
    if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("turboquant_"):
        return KVQuantMode.TURBOQUANT
    if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("fp8"):
        return KVQuantMode.FP8_PER_TENSOR
    return KVQuantMode.NONE
vllm/model_executor/layers/attention/attention.py data-contract

在 TQFullAttentionSpec 构造时补传 kv_quant_mode=quant_mode,这是让 KV 缓存规格携带正确量化模式的关键数据契约改动。

def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
    # quant_mode 由 get_kv_quant_mode(self.kv_cache_dtype) 得出
    # (前面的 sliding window 分支同样使用它)
    ...
    elif self.kv_cache_dtype.startswith("turboquant_"):
        from vllm.model_executor.layers.quantization.turboquant.config import (
            TurboQuantConfig,
        )
        from vllm.v1.kv_cache_interface import TQFullAttentionSpec
​
        tq_config = TurboQuantConfig.from_cache_dtype(
            self.kv_cache_dtype, self.head_size
        )
        # 关键变更:显式传入 kv_quant_mode=quant_mode,让 vllm/v1 的
        # KV 缓存整形逻辑能通过 kv_quant_mode != NONE 推断真实量化 dtype,
        # 避免退化为 auto。
        return TQFullAttentionSpec(
            block_size=block_size,
            num_kv_heads=self.num_kv_heads,
            head_size=self.head_size,
            head_size_v=self.head_size,
            dtype=self.kv_cache_torch_dtype,
            kv_quant_mode=quant_mode,
            tq_slot_size=tq_config.slot_size_aligned,
        )
    else:
        return FullAttentionSpec(
            block_size=block_size,
            num_kv_heads=self.num_kv_heads,
            head_size=self.head_size,
            head_size_v=self.head_size_v,
            dtype=self.kv_cache_torch_dtype,
            kv_quant_mode=quant_mode,
        )
tests/quantization/test_turboquant.py test-coverage

新增 TestTurboQuantKVCacheSpec 参数化测试,覆盖所有 turboquant 预设,验证 spec 生成时 kv_quant_mode 被正确设置为 TURBOQUANT。

class TestTurboQuantKVCacheSpec:
    @pytest.mark.parametrize("preset", ALL_PRESETS)
    def test_kv_cache_spec_sets_kv_quant_mode(self, preset):
        from vllm.model_executor.layers.attention.attention import Attention
        from vllm.v1.kv_cache_interface import KVQuantMode, TQFullAttentionSpec
​
        layer = SimpleNamespace(
            attn_type="decoder",
            kv_cache_dtype=preset,
            kv_cache_torch_dtype=torch.uint8,
            head_size=128,
            head_size_v=128,
            num_kv_heads=4,
            sliding_window=None,
        )
        vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=32))
​
        spec = Attention.get_kv_cache_spec(layer, vllm_config)
​
        # 每个 turboquant 预设都必须生成 TQFullAttentionSpec,
        # 且 kv_quant_mode 是 TURBOQUANT 而不是默认的 NONE。
        assert isinstance(spec, TQFullAttentionSpec)
        assert spec.kv_quant_mode == KVQuantMode.TURBOQUANT

评论区精华

Fork PR 自动审查被禁用 other

claude[bot] 指出该 PR 来自 fork,自动 review 被禁用,需要维护者手动触发。

结论:维护者 mgoin 直接批准合并,未触发额外人工 review。 · 已解决

风险与影响

attn_utils.py 中 layer_cache_dtype 的判断被简化为只看 kv_quant_mode == KVQuantMode.NONE,若未来出现其他量化模式未登记 KVQuantMode 而依赖特判,将会回归;get_kv_quant_mode 依赖 turboquant_ 前缀匹配,若新预设命名不符合前缀(如 tq4bit)会静默落到 NONE。测试仅覆盖 KVCacheSpec 生成阶段,未覆盖实际 _reshape_kv_cache 的 dtype 传递与 get_kv_cache_shape 调用路径,但改动较小,回归面可控。

直接影响使用 --kv-cache-dtype turboquant_* 的 v1 引擎用户,修复启动崩溃;KVQuantMode 枚举扩展为 6 个量化模式,turboquant 被纳入统一的量化模式分发架构。对非 turboquant 的 KV 缓存路径无行为变化,无新增依赖或配置,团队回归成本低。

v1 核心路径变更 前缀字符串匹配 测试未覆盖实机缓存初始化 改动小、回归面窄

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论