Skip to content

vllm.model_executor.models.granite

Inference-only IBM Granite model compatible with HuggingFace weights.

Also serves the granite_swa checkpoints (GraniteSWAForCausalLM), supporting three additional features: per-layer sliding window attention (layer_types), a learnable per-head attention sink (self_attn.sinks), and a per-layer RoPE base (layer_rope_theta, with 0 for NoPE).

Functions:

granite_layer_attn_params(config, layer_idx)

Resolve one layer's sliding window, RoPE base and sink usage.

Plain Granite configs carry no SWA fields and fall back to full attention, global RoPE base and no sink. HF SWA checkpoints use sinks without a dedicated flag, so assume true when layer_types is used, and allow attention_sinks to override that decision.

Returns:

  • int | None –

    Sliding window size (None for full attention), RoPE base theta (0

  • float –

    for NoPE), and attention sink presence/absence.

Source code in vllm/model_executor/models/granite.py 
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def granite_layer_attn_params(
    config: PretrainedConfig, layer_idx: int
) -> tuple[int | None, float, bool]:
    """Resolve one layer's sliding window, RoPE base and sink usage.

    Plain Granite configs carry no SWA fields and fall back to full
    attention, global RoPE base and no sink. HF SWA checkpoints use
    sinks without a dedicated flag, so assume true when `layer_types`
    is used, and allow `attention_sinks` to override that decision.

    Returns:
        Sliding window size (`None` for full attention), RoPE base theta (`0`
        for NoPE), and attention sink presence/absence.
    """
    layer_types = getattr(config, "layer_types", None)
    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_theta = (
        layer_rope_theta[layer_idx]
        if layer_rope_theta is not None
        else config.rope_parameters["rope_theta"]
    )

    has_sink = getattr(config, "attention_sinks", layer_types is not None)
    return sliding_window, rope_theta, has_sink