[Bugfix] Apply attention sinks in the Transformers backend (#52156)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Thomas Parnell committed
Sep 8, 2026 at 15:54 UTC
414057a3d3aa47915b10103e4a181604e0b3562f
4 files changed
+127
-31
tests/models/transformers/test_backend.py
+55
-11
@@ -6,6 +6,7 @@ import ast
6
import contextlib
7
import os
8
import tempfile
9
+from functools import partial
10
from types import SimpleNamespace
11
from typing import Any
12
@@ -40,23 +41,21 @@ def get_model(arch: str) -> str:
41
return model_info.default
42
43
44
+def count_modules(model, cls: type) -> int:
45
+ return sum(isinstance(m, cls) for m in model.modules())
46
+
47
+
48
def get_num_fused(model) -> tuple[int, int]:
49
from vllm.model_executor.layers.linear import (
50
MergedColumnParallelLinear,
51
QKVParallelLinear,
52
)
53
49
- glu = sum(isinstance(m, MergedColumnParallelLinear) for m in model.modules())
50
- qkv = sum(isinstance(m, QKVParallelLinear) for m in model.modules())
54
+ glu = count_modules(model, MergedColumnParallelLinear)
55
+ qkv = count_modules(model, QKVParallelLinear)
56
return glu, qkv
57
58
54
-def count_mla_layers(model) -> int:
55
- from vllm.model_executor.layers.attention import MLAAttention
56
-
57
- return sum(isinstance(m, MLAAttention) for m in model.attention_instances.values())
58
-
59
-
59
def check_implementation(
60
runner_ref: type[HfRunner | VllmRunner],
61
runner_test: type[VllmRunner],
@@ -87,6 +86,12 @@ def check_implementation(
86
assert num_glu == expected_glu * num_layers
87
assert num_qkv == expected_qkv * num_layers
88
89
+ tp_size = kwargs_test.get("tensor_parallel_size", 1)
90
+ from vllm.model_executor.layers.attention import Attention
91
+
92
+ counts = model_test.apply_model(partial(count_modules, cls=Attention))
93
+ assert counts == [num_layers] * tp_size
94
+
95
outputs_test = model_test.generate_greedy_logprobs(*args)
96
97
with runner_ref(model, **kwargs_ref) as model_ref:
@@ -143,6 +148,42 @@ def test_hybrid_attention(vllm_runner: type[VllmRunner]) -> None:
148
)
149
150
151
+def get_sinks(model) -> dict[int, torch.Tensor]:
152
+ """The sink tensor each attention layer was handed, keyed by layer index."""
153
+ sinks = {}
154
+ for i, (prefix, fuser) in model.attention_fusers.items():
155
+ if (sink := fuser.sinks(model.get_submodule(prefix))) is not None:
156
+ sinks[i] = sink.float().cpu()
157
+ return sinks
158
+
159
+
160
+def test_sinks(hf_runner: type[HfRunner], vllm_runner: type[VllmRunner]) -> None:
161
+ """Learnable attention sinks must reach the attention layers.
162
+
163
+ Only the attention impl can apply them, so if they are not passed to
164
+ `Attention` they are dropped and every softmax is subtly wrong.
165
+ """
166
+ from vllm.platforms import current_platform
167
+
168
+ if not current_platform.has_device_capability(90):
169
+ pytest.skip("Attention sinks need FlashAttention 3 or TRTLLM attention")
170
+
171
+ model = "tiny-random/gpt-oss-bf16"
172
+ with vllm_runner(
173
+ model, model_impl="transformers", max_model_len=1024, enforce_eager=True
174
+ ) as model_test:
175
+ assert model_test.llm.llm_engine.model_config.using_transformers_backend()
176
+ sinks = model_test.apply_model(get_sinks)[0]
177
+
178
+ with hf_runner(model, dtype="bfloat16") as model_ref:
179
+ layers = model_ref.model.model.get_decoder().layers
180
+ expected = {i: layer.self_attn.sinks.float() for i, layer in enumerate(layers)}
181
+
182
+ assert sinks.keys() == expected.keys()
183
+ for i, sink in sinks.items():
184
+ torch.testing.assert_close(sink, expected[i].cpu())
185
+
186
+
187
def test_mla(vllm_runner: type[VllmRunner], example_prompts: list[str]) -> None:
188
import transformers
189
from packaging.version import Version
@@ -165,7 +206,10 @@ def test_mla(vllm_runner: type[VllmRunner], example_prompts: list[str]) -> None:
206
model_config = model_test.llm.llm_engine.model_config
207
assert model_config.using_transformers_backend()
208
num_layers = model_config.hf_config.get_text_config().num_hidden_layers
168
- assert model_test.apply_model(count_mla_layers) == [num_layers]
209
+ from vllm.model_executor.layers.attention import MLAAttention
210
+
211
+ counts = model_test.apply_model(partial(count_modules, cls=MLAAttention))
212
+ assert counts == [num_layers]
213
outputs_test = model_test.generate_greedy_logprobs(*args)
214
215
with vllm_runner(model, model_impl="auto") as model_ref:
@@ -627,7 +671,7 @@ def test_attention_dispatch_is_matched(model_type: str):
671
"""Exactly the decoder layers' attention modules match an `AttentionFuser`.
672
673
That match is what `recursive_replace` records, and so what
630
- `create_attention_instances` attaches vLLM's attention layer to.
674
+ `_create_attention_instances` attaches vLLM's attention layer to.
675
"""
676
model = build_model(model_type)
677
matched = {
@@ -682,7 +726,7 @@ def test_attention_dispatch_is_required(monkeypatch: pytest.MonkeyPatch):
726
ValueError,
727
match="Layer 0 does not dispatch through the Transformers attention interface",
728
):
685
- Base.create_attention_instances(model)
729
+ Base._create_attention_instances(model)
730
731
732
@pytest.mark.parametrize("model_type", ATTENTION_MODEL_TYPES)
vllm/model_executor/models/transformers/__init__.py
+25
@@ -40,6 +40,29 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
40
if TYPE_CHECKING:
41
import torch
42
43
+ from vllm.model_executor.layers.attention import Attention, MLAAttention
44
+
45
+
46
+def check_sinks(
47
+ module: "torch.nn.Module",
48
+ self_attn: "Attention | MLAAttention",
49
+ s_aux: "torch.Tensor | None",
50
+):
51
+ """Fail loudly if the model applies a sink the attention layer will not.
52
+
53
+ Only the attention impl can fold a sink into the softmax denominator, so a sink
54
+ that never reached `Attention` is dropped and every softmax is subtly wrong.
55
+ """
56
+ if s_aux is None or getattr(self_attn, "has_sink", False):
57
+ return
58
+ raise ValueError(
59
+ f"{type(module).__name__} applies attention sinks, but they were not passed "
60
+ f"to {type(self_attn).__name__}, so the output would be wrong. Either the "
61
+ "Transformers modeling backend could not find the parameter holding them, or "
62
+ "vLLM does not support sinks for this kind of attention. Please open an issue "
63
+ "at https://github.com/vllm-project/vllm/issues/new"
64
+ )
65
+
66
67
def vllm_attention_forward(
68
# Transformers args
@@ -51,6 +74,7 @@ def vllm_attention_forward(
74
**kwargs,
75
):
76
self_attn = getattr(module, VLLM_ATTN_ATTR)
77
+ check_sinks(module, self_attn, kwargs.get("s_aux"))
78
hidden = query.shape[-2]
79
head_dim_qk = query.shape[-1]
80
head_dim_v = value.shape[-1]
@@ -80,6 +104,7 @@ def vllm_mla_attention_forward(
104
**kwargs,
105
):
106
self_attn = getattr(module, VLLM_ATTN_ATTR)
107
+ check_sinks(module, self_attn, kwargs.get("s_aux"))
108
# [batch=1, heads, num_tokens, qk_head_dim] -> [num_tokens, heads, qk_head_dim]
109
query = query.transpose(1, 2).flatten(0, 1)
110
num_tokens, num_heads = query.shape[:2]
vllm/model_executor/models/transformers/base.py
+14
-18
@@ -139,11 +139,11 @@ class Base(
139
for the quantization machinery and loaders (e.g. bitsandbytes)."""
140
self.fusers: dict[str, list[BaseFuser]] = {}
141
"""Module qualname -> the fusers applied to it, populated
142
- by `recursive_replace` for `create_attention_instances`."""
142
+ by `recursive_replace` for `_create_attention_instances`."""
143
self.attention_fusers: dict[int, tuple[str, AttentionFuser]] = {}
144
"""`layer_idx` -> the qualname and fuser of the module computing that
145
layer's attention, populated by `recursive_replace` for
146
- `create_attention_instances`."""
146
+ `_create_attention_instances`."""
147
148
# Attrs for Eagle3 (see self.set_aux_hidden_state_layers)
149
self._target_class: type[nn.Module] = nn.Module
@@ -179,7 +179,7 @@ class Base(
179
# Substitute remaining layers with vLLM's layers as needed
180
self.recursive_replace()
181
# Create attention instances for KV cache allocation
182
- self.attention_instances = self.create_attention_instances()
182
+ self._create_attention_instances()
183
184
# Initialize any parameters that have not had their modules replaced
185
self.init_parameters(self.model)
@@ -197,7 +197,7 @@ class Base(
197
Patch the config to ensure that the model is created correctly:
198
199
- Sets the attention implementation to "vllm" so the attention instances from
200
- `create_attention_instances` are used
200
+ `_create_attention_instances` are used
201
- Sets the dtype to the default torch dtype set by vLLM because Transformers
202
uses the config dtype when creating the model
203
"""
@@ -570,11 +570,8 @@ class Base(
570
571
_recursive_replace(self.model, prefix="model")
572
573
- def create_attention_instances(self) -> dict[int, Attention]:
574
- """
575
- Create `Attention` instances to inform KV cache allocation.
576
- """
577
- attention_instances = {}
573
+ def _create_attention_instances(self):
574
+ """Create `Attention` instances to inform KV cache allocation."""
575
text_config = self.text_config
576
attn_cls = self._get_attn_cls()
577
@@ -614,13 +611,9 @@ class Base(
611
self.parallel_config, arch_config
612
)
613
head_size = arch_config.head_size
617
- scale = attn_fuser.scale(attn_module)
618
- # Default to Llama scale if AttentionFuser couldn't identify it
619
- if scale is None:
614
+ if (scale := attn_fuser.scale(attn_module)) is None:
615
+ # Default to Llama scale if AttentionFuser couldn't identify it
616
scale = head_size**-0.5
621
- num_kv_heads = self.model_config.get_num_kv_heads(
622
- self.parallel_config, arch_config
623
- )
617
618
kwargs = dict(
619
num_heads=num_heads,
@@ -651,13 +644,18 @@ class Base(
644
else:
645
kwargs.update(
646
head_size=head_size,
654
- num_kv_heads=num_kv_heads,
647
+ num_kv_heads=self.model_config.get_num_kv_heads(
648
+ self.parallel_config, arch_config
649
+ ),
650
logits_soft_cap=logits_soft_cap,
651
)
652
653
# Handle interleaved sliding window attention
654
if layer_types and layer_types[i] == "sliding_attention":
655
kwargs["per_layer_sliding_window"] = text_config.sliding_window
656
+ # Handle attention sinks
657
+ if (sinks := attn_fuser.sinks(attn_module)) is not None:
658
+ kwargs["sinks"] = sinks
659
660
attn_instance = attn_cls(**kwargs)
661
@@ -670,8 +668,6 @@ class Base(
668
# layer identity into the traced graph and so costs one compiled
669
# artifact per layer.
670
setattr(attn_module, VLLM_ATTN_ATTR, attn_instance)
673
- attention_instances[i] = attn_instance
674
- return attention_instances
671
672
def _get_attn_cls(self) -> type[AttentionLayerBase]:
673
"""Return the `Attention` class to use for this model's layers."""
vllm/model_executor/models/transformers/fusers/attention.py
+33
-2
@@ -9,9 +9,12 @@ from dataclasses import dataclass
9
from functools import cache
10
from typing import TYPE_CHECKING, ClassVar
11
12
+import torch
13
from torch import fx, nn
14
15
+from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
16
from vllm.model_executor.models.transformers.fusers.base import BaseFuser
17
+from vllm.model_executor.utils import set_weight_attrs
18
19
if TYPE_CHECKING:
20
from vllm.config import VllmConfig
@@ -82,8 +85,10 @@ class AttentionFuser(BaseFuser):
85
86
source_cls: str
87
"""Class of the HF module that dispatches (for logging)."""
85
- scale_expr: ast.expr | None
88
+ scale_expr: ast.expr | None = None
89
"""Source of the `scaling=` the module hands the interface, if it hands one."""
90
+ s_aux_expr: ast.expr | None = None
91
+ """Source of the `s_aux=` the module hands the interface, if it hands one."""
92
93
def info(self, name: str) -> str:
94
return f"Found: {name} ({self.source_cls}) -> attention interface"
@@ -96,7 +101,13 @@ class AttentionFuser(BaseFuser):
101
return None
102
scaling = [kw.value for kw in call.keywords if kw.arg == "scaling"]
103
scale_expr = scaling[0] if len(scaling) == 1 else None
99
- return cls(source_cls=type(module).__name__, scale_expr=scale_expr)
104
+ s_aux = [kw.value for kw in call.keywords if kw.arg == "s_aux"]
105
+ s_aux_expr = s_aux[0] if len(s_aux) == 1 else None
106
+ return cls(
107
+ source_cls=type(module).__name__,
108
+ scale_expr=scale_expr,
109
+ s_aux_expr=s_aux_expr,
110
+ )
111
112
def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
113
"""Whether `module` will actually dispatch to vLLM."""
@@ -108,6 +119,13 @@ class AttentionFuser(BaseFuser):
119
def fuse(
120
self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
121
) -> nn.Module:
122
+ if (sinks := self.sinks(module)) is not None:
123
+ size = sinks.numel() // vllm_config.parallel_config.tensor_parallel_size
124
+ device = vllm_config.device_config.device
125
+ data = torch.empty(size, dtype=sinks.dtype, device=device)
126
+ sinks_param = nn.Parameter(data, requires_grad=False)
127
+ set_weight_attrs(sinks_param, {"weight_loader": sharded_weight_loader(0)})
128
+ setattr(module, self.s_aux_expr.attr, sinks_param)
129
return module
130
131
def layer_index(self, module: nn.Module) -> int | None:
@@ -127,3 +145,16 @@ class AttentionFuser(BaseFuser):
145
f"{type(module).__name__}."
146
)
147
return float(scale)
148
+
149
+ def sinks(self, module: nn.Module) -> nn.Parameter | None:
150
+ """The per-head sink tensor `module` passes to the interface, or `None`."""
151
+ if self.s_aux_expr is None:
152
+ return None
153
+ s_aux = _resolve(self.s_aux_expr, module)
154
+ if not isinstance(s_aux, (nn.Parameter, type(None))):
155
+ expression = ast.unparse(self.s_aux_expr)
156
+ raise ValueError(
157
+ f"Cannot resolve attention s_aux expression {expression!r} in "
158
+ f"{type(module).__name__}."
159
+ )
160
+ return s_aux