执行摘要

• 一句话：新增 GraniteSWA/MoeSWA 支持，复用现有 Granite 实现
• 推荐动作：值得精读。该 PR 对模型扩展和权重加载的工程化处理有参考价值，尤其是 granite_layer_attn_params 的逐层解析逻辑和 hf_to_vllm_mapper 的复用方式。review 中关于配置格式的讨论也提供了每层配置设计的不同视角。

功能与动机

PR body 明确说明要支持 IBM 发布的 GraniteSWA 与 GraniteMoeSWA 模型，其中 GraniteSWA 是带每层滑动窗口注意力和可学习 attention sink 的 Granite，GraniteMoeSWA 是 GraniteMoeShared 风格 MoE 结合 SWA 注意力。与原始提案 #48270 不同，本实现选择折叠进现有模型类而非新建独立模型。同时指出 transformers 后端无法正确处理这些模型（sink token 被静默丢弃导致输出错误，见 #52156），因此需要 vLLM 原生实现。

实现拆解

实现分五步完成：

1. 在 granite.py 中新增 granite_layer_attn_params 函数：接收 config 和层索引，返回每层的滑动窗口大小、RoPE 基频（0 表示 NoPE）和是否存在 attention sink。同时改造 GraniteAttention.__init__，不再直接使用全局 rope_parameters，而是按层解析后构造 RoPE 与 sinks 参数，并传给 Attention。

2. 在 granitemoe.py 中同步重构 GraniteMoeAttention：构造函数改为接收 config 对象，复用 granite_layer_attn_params 解析每层参数；当 has_sink 时创建 nn.Parameter 并绑定 sharded_weight_loader(0)，否则置为 None；Attention 传入 per_layer_sliding_window 与 sinks。同时删除了旧的 fused_moe_make_expert_params_mapping 手动权重拆分，改用自动加载。

3. 在 granitemoeshared.py 中删除手写 load_weights 的 MoE 权重展开逻辑：原逻辑将 block_sparse_moe.input_linear.weight 拆分为各 expert 的 w1/w3，新实现直接复用 GraniteMoeModel.hf_to_vllm_mapper 和 GraniteMoeModel.load_weights，兼容新旧 HF 检查点命名，并保持 tie_word_embeddings 的跳过前缀逻辑。

4. 注册与文档：vllm/model_executor/models/registry.py 将 GraniteSWAForCausalLM 映射到 granite.GraniteForCausalLM，GraniteMoeSWAForCausalLM 映射到 granitemoeshared.GraniteMoeSharedForCausalLM；tests/models/registry.py 添加 _HfExamplesInfo 条目；docs/models/supported_models.md 补充支持列表。

5. 测试配套：tests/models/language/generation/test_granite.py 将模型列表改为字典并加入 ibm-granite/granite-swash-2b 和 ibm-granite/granite-swash-3b-a600m，通过 check_transformers_version 约束最低 transformers 版本（5.15.1）；新增两个单元测试验证无 SWA 配置时特性关闭、以及 layer_types 和 layer_rope_theta 逐层解析正确。

关键文件：

• vllm/model_executor/models/granite.py（模块 模型实现；类别 source；类型 data-contract；符号 granite_layer_attn_params）: 核心实现，新增 granite_layer_attn_params 解析每层滑动窗口、RoPE 与 sink 配置，并改造 GraniteAttention 支持这些特性。
• vllm/model_executor/models/granitemoe.py（模块 模型实现；类别 source；类型 data-contract；符号 _load_weights）: 重构 GraniteMoeAttention 以支持 SWA 特性，并清理权重加载逻辑，改动量最大（+57/-150）。
• vllm/model_executor/models/granitemoeshared.py（模块 模型实现；类别 source；类型 data-contract；符号 load_weights）: 删除手写 load_weights 的 MoE 权重拆分逻辑，复用 GraniteMoeModel 的映射与加载器，兼容新旧 HF 命名。
• tests/models/language/generation/test_granite.py（模块 模型测试；类别 test；类型 test-coverage；符号 test_granite_swa_features_are_off_without_swa_config, test_granite_swa_features_resolve_per_layer）: 添加 GraniteSWA 模型 parity 测试与逐层配置单元测试，验证新功能正确性及默认关闭行为。
• vllm/model_executor/models/registry.py（模块 注册表；类别 source；类型 data-contract）: 注册两个新 ForCausalLM 类到现有实现，是让模型可被 vLLM 发现的关键入口。
• tests/models/registry.py（模块 测试注册；类别 test；类型 test-coverage）: 为测试工具提供新模型的 HF examples 信息与最低 transformers 版本。
• docs/models/supported_models.md（模块 文档；类别 docs；类型 documentation）: 更新支持模型文档，让用户知晓新模型可用。

关键符号：granite_layer_attn_params, GraniteMoeAttention.init, GraniteMoeSharedForCausalLM.load_weights, GraniteForCausalLM.load_weights

关键源码片段

vllm/model_executor/models/granite.py

核心实现，新增 granite_layer_attn_params 解析每层滑动窗口、RoPE 与 sink 配置，并改造 GraniteAttention 支持这些特性。

def granite_layer_attn_params(
    config: PretrainedConfig, layer_idx: int
) -> tuple[int | None, float, bool]:
    """解析某一层的滑动窗口、RoPE 基频与 attention sink 使用情况。
​
    普通 Granite 配置没有 SWA 字段，回退到全注意力、全局 RoPE 基频且无 sink。
    HF SWA checkpoint 没有专门的 sink 标志，因此当使用 `layer_types` 时默认启用
    sink，并允许 `attention_sinks` 覆盖该决定。
​
    返回：
        滑动窗口大小（`None` 表示全注意力）、RoPE 基频 theta（`0` 表示 NoPE）、
        以及是否存在 attention sink。
    """
    layer_types = getattr(config, "layer_types", None)
    # 仅当该层为 sliding_attention 时才启用滑动窗口
    sliding_window = (
        config.sliding_window
        if layer_types is not None and layer_types[layer_idx] == "sliding_attention"
        else None
    )
​
    layer_rope_theta = getattr(config, "layer_rope_theta", None)
    # 每层可覆盖全局 RoPE 基频，0 表示该层不旋转（NoPE）
    rope_theta = (
        layer_rope_theta[layer_idx]
        if layer_rope_theta is not None
        else config.rope_parameters["rope_theta"]
    )
​
    # HF 约定：存在 layer_types 时默认有 sink，除非显式关闭
    has_sink = getattr(config, "attention_sinks", layer_types is not None)
    return sliding_window, rope_theta, has_sink

vllm/model_executor/models/granitemoe.py

重构 GraniteMoeAttention 以支持 SWA 特性，并清理权重加载逻辑，改动量最大（+57/-150）。

class GraniteMoeAttention(nn.Module):
    def __init__(
        self,
        config: PretrainedConfig,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        max_position: int = 4096 * 32,
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
        attention_multiplier: float | None = None,
        prefix: str = "",
    ) -> None:
        # 省略：qkv_proj / o_proj 等常规初始化
        # ...
​
        # 通过 granite_layer_attn_params 解析每层的滑动窗口、RoPE 基频和 sink
        sliding_window, rope_theta, has_sink = granite_layer_attn_params(
            config, extract_layer_index(prefix)
        )
​
        self.use_rope = rope_theta != 0 # theta 为 0 表示 NoPE
        if self.use_rope:
            self.rotary_emb = get_rope(
                self.head_dim,
                max_position=max_position,
                rope_parameters={**config.rope_parameters, "rope_theta": rope_theta},
                is_neox_style=True,
            )
​
        if has_sink:
            # 可学习的 per-head sink，权重加载时按 rank 0 分片
            self.sinks = nn.Parameter(torch.empty(self.num_heads), requires_grad=False)
            set_weight_attrs(self.sinks, {"weight_loader": sharded_weight_loader(0)})
        else:
            self.sinks = None
​
        self.attn = Attention(
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_kv_heads,
            cache_config=cache_config,
            quant_config=quant_config,
            per_layer_sliding_window=sliding_window,
            prefix=f"{prefix}.attn",
            sinks=self.sinks,
        )

评论区精华

核心讨论集中在两点：

1. 权重加载简化：hmellor 在 granitemoe.py 评论中指出，旧的权重拆分逻辑没有必要，RoutedExperts 可以直接加载融合权重，并表示“我会推送正确的权重加载”。随后提交了清理提交，用 hf_to_vllm_mapper 替代手动展开，作者 daviswer 接受。

2. 每层配置的格式取舍：hmellor 建议在 granite.py 使用 nested rope_parameters（类似 Gemma 3），daviswer 解释因为 layer type 和 rope 参数各自独立按层索引决定，扁平列表更清晰；hmellor 进一步建议按 layer type 嵌套，但 daviswer 说明这两个维度并不绑定，最终维持扁平列表。hmellor 还提到了 Nvidia 的 per_layer_config API 作为未来改进方向，daviswer 表示感谢。

• 权重加载拆分逻辑是否必要 (design): hmellor 提交了清理，用 hf_to_vllm_mapper 和 AutoWeightsLoader 替换手动拆分，daviswer 接受，问题解决。
• 每层 RoPE 配置的 schema 设计 (design): 维持当前扁平列表实现，但记录 per_layer_config API 作为未来改进参考，讨论友好结束。

风险与影响

• 风险：主要风险点：
  • transformers 后端不兼容：PR body 明确声明 transformers 后端会静默丢弃 sink token，产生错误输出（#52156）。若用户错误使用该后端（例如通过 --model-impl 或自动回退），可能得到错误结果而不知情。
  • 权重加载逻辑重写：granitemoeshared.py 删除手写拆分后依赖 hf_to_vllm_mapper，需保证新旧 HF 检查点命名均正确映射；hmellor 用 ibm/PowerMoE-3b 验证生成一致，但覆盖面有限。
  • sinks 参数加载：sharded_weight_loader(0) 要求权重名字和维度正确，若 checkpoint 格式变化可能导致加载失败或静默错位。
  • 依赖版本：测试依赖 transformers 5.15.1+，旧环境可能跳过或失败。
  • 影响：用户可直接加载 GraniteSWA 和 GraniteMoeSWA 系列检查点，无需等待独立模型实现；现有 Granite/GraniteMoeShared 检查点行为不受影响（新特性默认关闭）。对团队而言，该 PR 展示了在已有模型类上扩展变体架构的模式，并清理了冗余的权重加载代码，降低了后续维护成本。同时将 transformers 后端不支持的模型场景固化到文档和测试中，有助于避免误用。
  • 风险标记：transformers 后端不兼容, 权重加载逻辑重写, 依赖 transformers>=5.15.1, 新特性默认关闭

关联脉络

• PR #48270 [Model] Add GraniteSWA and GraniteMoeSWA: 本 PR 的原始提案，被本 PR 取代并关闭；review 中引用了该 PR 的评论。
• PR #52156 [Bug?] GraniteSWA wrong outputs with transformers backend: PR body 中提及 transformers 后端无法正确处理 sink token 的问题，是本 PR 需要原生实现的重要动机。