[Bugfix][MiMo] Apply vision attention sinks in the window attention path (#49815)
Signed-off-by: almogtavor <almogtavor@gmail.com>
Almog Tavor committed
Aug 11, 2026 at 01:36 UTC
c3cac8c63d917759a38e94dfa933e71196de40bb
3 files changed
+132
-12
tests/models/multimodal/test_mimo_v2_omni.py
new
+91
@@ -0,0 +1,91 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""MiMo-V2 vision window attention has to apply the per-head sink logits."""
4
+
5
+import pytest
6
+import torch
7
+
8
+from tests.utils import ensure_current_vllm_config
9
+from vllm.distributed.parallel_state import (
10
+ init_distributed_environment,
11
+ initialize_model_parallel,
12
+)
13
+from vllm.platforms import current_platform
14
+from vllm.utils.network_utils import get_open_port
15
+
16
+EMBED_DIM = 256
17
+NUM_HEADS = 4
18
+HEAD_DIM = 64
19
+WINDOW = 8
20
+# One sequence shorter than the window, one longer.
21
+SEQ_LENS = [5, 37]
22
+
23
+
24
+@pytest.fixture(scope="module")
25
+def vision_attn_env():
26
+ init_distributed_environment(
27
+ world_size=1,
28
+ rank=0,
29
+ local_rank=0,
30
+ distributed_init_method=f"tcp://127.0.0.1:{get_open_port()}",
31
+ backend="nccl",
32
+ )
33
+ default_dtype = torch.get_default_dtype()
34
+ torch.set_default_dtype(torch.bfloat16)
35
+ with ensure_current_vllm_config():
36
+ initialize_model_parallel(tensor_model_parallel_size=1)
37
+ yield
38
+ torch.set_default_dtype(default_dtype)
39
+
40
+
41
+def _reference(q, k, v, cu_seqlens, sinks, scale):
42
+ """Dense windowed softmax with the sink added to each sequence's key 0."""
43
+ groups = q.shape[1] // k.shape[1]
44
+ out = torch.empty_like(q, dtype=torch.float32)
45
+ for start, end in zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist()):
46
+ qs = q[start:end].float()
47
+ ks = k[start:end].float().repeat_interleave(groups, dim=1)
48
+ vs = v[start:end].float().repeat_interleave(groups, dim=1)
49
+ scores = torch.einsum("qhd,khd->hqk", qs, ks) * scale
50
+ scores[..., 0] += sinks.float().view(-1, 1)
51
+ pos = torch.arange(end - start, device=q.device)
52
+ outside = (pos.view(-1, 1) - pos.view(1, -1)).abs() > WINDOW
53
+ scores.masked_fill_(outside, -torch.inf)
54
+ out[start:end] = torch.einsum("hqk,khd->qhd", scores.softmax(-1), vs)
55
+ return out
56
+
57
+
58
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires flash-attn")
59
+@pytest.mark.parametrize("num_kv_heads", [NUM_HEADS, NUM_HEADS // 2])
60
+def test_window_attention_applies_sinks(vision_attn_env, num_kv_heads):
61
+ from vllm.model_executor.models.mimo_v2_omni import MiMoVisionAttention
62
+
63
+ torch.manual_seed(0)
64
+ attn = MiMoVisionAttention(
65
+ embed_dim=EMBED_DIM,
66
+ num_heads=NUM_HEADS,
67
+ num_kv_heads=num_kv_heads,
68
+ qk_channels=HEAD_DIM,
69
+ kv_channels=HEAD_DIM,
70
+ use_sink=True,
71
+ visual_token_window_size=WINDOW,
72
+ ).cuda()
73
+ attn.sinks.data.normal_()
74
+
75
+ total = sum(SEQ_LENS)
76
+ opts = dict(device="cuda", dtype=torch.bfloat16)
77
+ q = torch.randn(total, NUM_HEADS, HEAD_DIM, **opts)
78
+ k = torch.randn(total, num_kv_heads, HEAD_DIM, **opts)
79
+ v = torch.randn(total, num_kv_heads, HEAD_DIM, **opts)
80
+ cu_seqlens = torch.tensor(
81
+ [0, *torch.tensor(SEQ_LENS).cumsum(0).tolist()],
82
+ device="cuda",
83
+ dtype=torch.int32,
84
+ )
85
+
86
+ out = attn._forward_window_attn(q, k, v, cu_seqlens, max(SEQ_LENS))
87
+ ref = _reference(q, k, v, cu_seqlens, attn.sinks, attn.scale)
88
+
89
+ # bf16 attention lands at ~2e-3 here; dropping the sinks lands at ~1e-1.
90
+ error = ((out.float() - ref).norm() / ref.norm()).item()
91
+ assert error < 1e-2, f"sink-corrected output is off by {error:.2e}"
vllm/model_executor/models/mimo_v2_omni.py
+27
-10
@@ -193,7 +193,7 @@ class MiMoVisionAttention(nn.Module):
193
# Rotary embeddings applied separately to Q and K
194
self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True)
195
196
- # Sink attention weights (loaded but not used in vLLM flash_attn)
196
+ # Per-head sink logits, applied in the window attention path.
197
# The checkpoint stores these only for non-full-attention blocks
198
self.use_sink = use_sink
199
if use_sink:
@@ -214,21 +214,38 @@ class MiMoVisionAttention(nn.Module):
214
cu_seqlens: torch.Tensor,
215
max_seqlen: torch.Tensor,
216
) -> torch.Tensor:
217
- """Window attention via flash_attn_varlen_func with window_size."""
218
- from vllm.vllm_flash_attn import flash_attn_varlen_func
217
+ """Window attention with the per-head sink applied to key 0.
218
+
219
+ The reference adds ``sinks[h]`` to the logit of each sequence's first
220
+ key, which the Triton prefill kernel supports directly, so the softmax
221
+ normalizes over the biased scores in one pass.
222
+ """
223
+ from vllm.v1.attention.ops.triton_prefill_attention import (
224
+ context_attention_fwd,
225
+ )
226
227
w = self.visual_token_window_size
221
- output = flash_attn_varlen_func(
228
+ output = torch.empty_like(q)
229
+ head_start = self.tp_rank * self.num_heads_per_partition
230
+ sinks = (
231
+ self.sinks[head_start : head_start + self.num_heads_per_partition]
232
+ if self.sinks is not None
233
+ else None
234
+ )
235
+ context_attention_fwd(
236
q,
237
k,
238
v,
225
- cu_seqlens_q=cu_seqlens,
226
- cu_seqlens_k=cu_seqlens,
227
- max_seqlen_q=max_seqlen,
228
- max_seqlen_k=max_seqlen,
239
+ output,
240
+ b_start_loc=cu_seqlens[:-1],
241
+ b_seq_len=cu_seqlens[1:] - cu_seqlens[:-1],
242
+ max_input_len=max_seqlen,
243
+ is_causal=False,
244
softmax_scale=self.scale,
230
- causal=False,
231
- window_size=[w, w],
245
+ sliding_window_q=w,
246
+ sliding_window_k=w,
247
+ sinks=sinks,
248
+ sinks_bias_key0=True,
249
)
250
return output
251
vllm/v1/attention/ops/triton_prefill_attention.py
+14
-2
@@ -58,6 +58,7 @@ def _fwd_kernel(
58
IS_CAUSAL: tl.constexpr,
59
SLIDING_WINDOW_Q: tl.constexpr,
60
SLIDING_WINDOW_K: tl.constexpr,
61
+ SINKS_BIAS_KEY0: tl.constexpr,
62
USE_SINKS: tl.constexpr,
63
Lk: tl.constexpr,
64
):
@@ -98,8 +99,15 @@ def _fwd_kernel(
99
# initialize pointer to m and l
100
if USE_SINKS:
101
sink = tl.load(Sinks + cur_head) * 1.4426950408889634
101
- m_i = tl.full([BLOCK_M], sink, dtype=tl.float32)
102
- l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
102
+ if SINKS_BIAS_KEY0:
103
+ # Sinks bias the logit of key 0, so the softmax starts empty and
104
+ # normalizes over the biased scores as usual.
105
+ m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
106
+ l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
107
+ else:
108
+ # Sinks are a null logit that only inflates the denominator.
109
+ m_i = tl.full([BLOCK_M], sink, dtype=tl.float32)
110
+ l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
111
else:
112
m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
113
l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
@@ -160,6 +168,8 @@ def _fwd_kernel(
168
169
qk = tl.dot(q, k)
170
qk = tl.where(mask, qk * sm_scale, -1.0e8)
171
+ if USE_SINKS and SINKS_BIAS_KEY0:
172
+ qk = tl.where(mask & (pos_k == 0), qk + sink, qk)
173
m_ij = tl.maximum(m_i, tl.max(qk, 1))
174
qk -= m_ij[:, None]
175
p = tl.math.exp2(qk)
@@ -217,6 +227,7 @@ def context_attention_fwd(
227
sliding_window_q: int | None = None,
228
sliding_window_k: int | None = None,
229
sinks: torch.Tensor | None = None,
230
+ sinks_bias_key0: bool = False,
231
):
232
"""
233
q, k, v: [b * s, head, head_dim]
@@ -268,6 +279,7 @@ def context_attention_fwd(
279
SLIDING_WINDOW_Q=sliding_window_q,
280
SLIDING_WINDOW_K=sliding_window_k,
281
USE_SINKS=sinks is not None,
282
+ SINKS_BIAS_KEY0=sinks_bias_key0,
283
num_warps=num_warps,
284
num_stages=1,
285
Lk=Lk,