执行摘要

• 一句话：支持按 KV 缓存组分别选择注意力后端
• 推荐动作：值得精读。此 PR 展示了如何在已有配置框架上以最小侵入性扩展新功能：通过 Pydantic 验证器解析配置、复用枚举类型、在核心选择路径插入查找逻辑。设计思路对类似组件级配置（如不同 MoE 专家使用不同 kernel）有参考价值。

功能与动机

PR body 说明：模型将层拆分到多个 KV 缓存组时，被迫只能指定一个后端，尽管运行时 AttentionGroup 已支持异构后端。用户无法表达『full attention 层用 FlashAttention，sliding-window 层用 Triton』。此 PR 通过新增 backend_per_kind 映射填补了 UX 缺口。关联的 Issue #48011 先添加了滑动窗口能力标记，为此 PR 提供基础。

实现拆解

1. 配置模型扩展（vllm/config/attention.py）：在 AttentionConfig 中新增 backend_per_kind 字段（类型 dict[str, AttentionBackendEnum]），默认空字典。添加 Pydantic 验证器 validate_backend_per_kind_before，将字符串键/值解析为对应的枚举，并校验键是否属于 KVCacheSpecKind 的合法值。

2. 种类推导函数（vllm/v1/attention/selector.py）：新增 get_attn_spec_kind(use_mla, has_sliding_window, attn_type) 函数，根据层的属性（是否 MLA、是否滑动窗口、编解码类型）映射到对应的 KVCacheSpecKind 枚举值。该函数是 get_kv_cache_spec_kind 的镜像，但基于构建时的输入信号而非已生成的 KVCacheSpec。

3. 后端选择集成（vllm/v1/attention/selector.py）：在 get_attn_backend 中，构建 AttentionSelectorConfig 之后、传递给缓存函数之前，检查 vllm_config.attention_config.backend_per_kind。若不为空，则调用 get_attn_spec_kind 获取当前层的 kind，若该 kind 在映射中则使用对应的后端，否则回退到全局 backend。

4. 测试覆盖：

   • tests/v1/attention/test_backend_per_kind.py：单元测试验证 get_attn_spec_kind 对解码器层（四种 MLA/SW 组合）和编码器/交叉注意力层的种类推导；验证 AttentionConfig 对字符串的解析、非法 kind 的拒绝和默认空值。
   • tests/v1/e2e/general/test_attention_backend_per_kind.py：端到端测试，使用 google/gemma-3-1b-it（同时有 full_attention 和 sliding_window 组），通过 collective_rpc 收集运行时各组实际使用的后端名称，断言与配置一致。参数化测试交换映射，证明因果关系。

关键文件：

• vllm/v1/attention/selector.py（模块 注意力选择器；类别 source；类型 core-logic；符号 get_attn_spec_kind, get_attn_backend）: 核心实现：新增 get_attn_spec_kind 函数将层属性映射到 KVCacheSpecKind，并在 get_attn_backend 中集成 per-kind 后端查找逻辑。
• vllm/config/attention.py（模块 配置；类别 source；类型 configuration；符号 validate_backend_per_kind_before, backend_per_kind）: 配置模型扩展：新增 backend_per_kind 字段及其 Pydantic 验证器，实现配置解析与校验。
• tests/v1/attention/test_backend_per_kind.py（模块 单元测试；类别 test；类型 test-coverage；符号 test_get_attn_spec_kind_decoder, test_get_attn_spec_kind_attn_type, test_backend_per_kind_parses_strings, test_backend_per_kind_rejects_unknown_kind）: 单元测试：覆盖 get_attn_spec_kind 的种类推导、AttentionConfig 的解析和验证。
• tests/v1/e2e/general/test_attention_backend_per_kind.py（模块 端到端测试；类别 test；类型 test-coverage；符号 _collect_group_backends, test_backend_per_kind_splits_groups, test_backend_per_kind_overrides_global_backend）: 端到端测试：在真实模型上验证 per-kind 后端选择生效，通过 swapped 参数化证明因果关系。

关键符号：get_attn_spec_kind, validate_backend_per_kind_before, _collect_group_backends

关键源码片段

vllm/v1/attention/selector.py

核心实现：新增 get_attn_spec_kind 函数将层属性映射到 KVCacheSpecKind，并在 get_attn_backend 中集成 per-kind 后端查找逻辑。

# 文件：vllm/v1/attention/selector.py
​
from typing import TYPE_CHECKING
​
if TYPE_CHECKING:
    from vllm.v1.kv_cache_interface import KVCacheSpecKind
​
def get_attn_spec_kind(
    use_mla: bool,
    has_sliding_window: bool,
    attn_type: str,
) -> "KVCacheSpecKind":
    """根据层的属性派生 KVCacheSpecKind（KV 缓存组种类）。
​
    该函数是 get_kv_cache_spec_kind 的镜像，基于构建时的输入信号
    而非已生成的 KVCacheSpec，便于用户按种类配置 backend_per_kind。
    """
    from vllm.v1.kv_cache_interface import KVCacheSpecKind
​
    # 处理编码器 / 解码器注意力类型
    if attn_type == AttentionType.ENCODER_ONLY:
        return KVCacheSpecKind.ENCODER_ONLY_ATTENTION
    if attn_type == AttentionType.ENCODER_DECODER:
        return KVCacheSpecKind.CROSS_ATTENTION
​
    # MLA 层分支
    if use_mla:
        return (
            KVCacheSpecKind.SLIDING_WINDOW_MLA 
            if has_sliding_window 
            else KVCacheSpecKind.MLA_ATTENTION
        )
​
    # 非 MLA 解码器层分支
    return (
        KVCacheSpecKind.SLIDING_WINDOW
        if has_sliding_window
        else KVCacheSpecKind.FULL_ATTENTION
    )

# 在 get_attn_backend 中插入后端覆盖逻辑
​
    # ... 前面的配置构建 ...
    attn_type = attn_type or AttentionType.DECODER
    attn_selector_config = AttentionSelectorConfig(
        head_size=head_size,
        dtype=dtype,
        kv_cache_dtype=cast(CacheDType | None, kv_cache_dtype),
        block_size=block_size,
        use_mla=use_mla,
        has_sink=has_sink,
        use_sparse=use_sparse,
        use_mm_prefix=use_mm_prefix,
        use_per_head_quant_scales=use_per_head_quant_scales,
        attn_type=attn_type,
        has_sliding_window=has_sliding_window,
        use_non_causal=vllm_config.attention_config.use_non_causal,
        use_batch_invariant=envs.VLLM_BATCH_INVARIANT,
        use_kv_connector=use_kv_connector,
    )
​
    # 若配置了 per-kind 覆盖，则根据当前层的种类查找对应的后端
    attention_config = vllm_config.attention_config
    backend = attention_config.backend
    if attention_config.backend_per_kind:
        kind = get_attn_spec_kind(
            use_mla=use_mla,
            has_sliding_window=has_sliding_window,
            attn_type=attn_type,
        )
        # 优先使用 per-kind 映射，不存在则 fallback 到全局 backend
        backend = attention_config.backend_per_kind.get(kind.value, backend)
​
    return _cached_get_attn_backend(
        backend=backend,
        attn_selector_config=attn_selector_config,
        num_heads=num_heads,
    )

vllm/config/attention.py

配置模型扩展：新增 backend_per_kind 字段及其 Pydantic 验证器，实现配置解析与校验。

# 文件：vllm/config/attention.py
​
from dataclasses import field
from typing import Any
from pydantic import field_validator
​
@config
class AttentionConfig:
    # ... 其他字段 ...
​
    backend_per_kind: dict[str, AttentionBackendEnum] = field(default_factory=dict)
    """Per-KV-cache-group attention backend overrides, keyed by
    KVCacheSpecKind (e.g., {"mla_attention": "FLASHINFER_MLA",
    "sliding_window_mla": "TRITON_MLA"}). This lets a model that splits its
    layers across multiple KV-cache groups use a different backend per group.
​
    An entry overrides `backend` for layers of the matching kind; kinds not
    listed fall back to `backend` (or automatic selection)."""
​
    @field_validator("backend_per_kind", mode="before")
    @classmethod
    def validate_backend_per_kind_before(cls, value: Any) -> Any:
        """Parse the backend_per_kind map from strings."""
        from vllm.v1.kv_cache_interface import KVCacheSpecKind
​
        if not isinstance(value, dict):
            return value
        valid_kinds = {kind.value for kind in KVCacheSpecKind}
        parsed: dict[str, AttentionBackendEnum] = {}
        for kind, backend in value.items():
            if kind not in valid_kinds:
                raise ValueError(
                    f"Unknown KV cache group kind '{kind}' in "
                    f"backend_per_kind. Valid kinds are: "
                    f"{', '.join(sorted(valid_kinds))}."
                )
            if isinstance(backend, str):
                backend = AttentionBackendEnum[backend.upper()]
            parsed[kind] = backend
        return parsed

评论区精华

Review 过程简单，MatthewBonanni 直接批准。作者在 PR body 中指出了已知设计取舍：sink_full_attention 和 chunked_local_attention 种类不可通过 backend_per_kind 单独指定，因为它们对应的层不暴露区分信号给选择器，会被解析为 full_attention。这是一个待未来扩展的局限性，当前行为不会产生错误配置。

• sink_full_attention 和 chunked_local_attention 不可单独指定 (design): 接受为当前限制，未来可能扩展。

风险与影响

• 风险：风险较低：

  • 向后兼容：默认 backend_per_kind 为空字典，不改变任何现有行为。
  • 配置错误：验证器会拒绝非法 kind 名称，错误会在启动时立即暴露，避免运行时意外。
  • 一致性风险：get_attn_spec_kind 与 get_kv_cache_spec_kind 需要保持同步；若未来新增 KVCacheSpecKind 枚举值，必须在此函数中添加对应分支。测试覆盖了当前所有组合，但新增时需要更新测试。
  • 性能：每次 get_attn_backend 调用都会查字典，但次数有限（每层一次启动时），影响极小。
  • 影响：用户影响：为高级用户提供了细粒度控制注意力后端的能力，尤其在混合注意力模型（如 Gemma-3、DeepSeek-V3 等）中可分别调优不同层的性能/精度。
    系统影响：无部署或 API 变化，仅新增一个异步配置项。
    团队影响：扩展了配置模型，需要维护 get_attn_spec_kind 与 KVCacheSpecKind 的一致性。

• 风险标记：配置一致性风险, 核心路径扩展

关联脉络

• PR #48011 [Attention] Make sliding-window support an explicit backend capability: 此 PR 是 #48012 的前置依赖，为 sliding-window 添加了显式能力标记，使得 backend_per_kind 可以基于 has_sliding_window 进行种类区分。