[Bugfix][MiMo] Apply vision attention sinks in the window attention path (#49815)

Signed-off-by: almogtavor <[email protected]>

tests/models/multimodal/test_mimo_v2_omni.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiMo-V2 vision window attention has to apply the per-head sink logits."""
+
+import pytest
+import torch
+
+from tests.utils import ensure_current_vllm_config
+from vllm.distributed.parallel_state import (
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm.platforms import current_platform
+from vllm.utils.network_utils import get_open_port
+
+EMBED_DIM = 256
+NUM_HEADS = 4
+HEAD_DIM = 64
+WINDOW = 8
+# One sequence shorter than the window, one longer.
+SEQ_LENS = [5, 37]
+
+
+@pytest.fixture(scope="module")
+def vision_attn_env():
+    init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method=f"tcp://127.0.0.1:{get_open_port()}",
+        backend="nccl",
+    )
+    default_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(torch.bfloat16)
+    with ensure_current_vllm_config():
+        initialize_model_parallel(tensor_model_parallel_size=1)
+        yield
+    torch.set_default_dtype(default_dtype)
+
+
+def _reference(q, k, v, cu_seqlens, sinks, scale):
+    """Dense windowed softmax with the sink added to each sequence's key 0."""
+    groups = q.shape[1] // k.shape[1]
+    out = torch.empty_like(q, dtype=torch.float32)
+    for start, end in zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist()):
+        qs = q[start:end].float()
+        ks = k[start:end].float().repeat_interleave(groups, dim=1)
+        vs = v[start:end].float().repeat_interleave(groups, dim=1)
+        scores = torch.einsum("qhd,khd->hqk", qs, ks) * scale
+        scores[..., 0] += sinks.float().view(-1, 1)
+        pos = torch.arange(end - start, device=q.device)
+        outside = (pos.view(-1, 1) - pos.view(1, -1)).abs() > WINDOW
+        scores.masked_fill_(outside, -torch.inf)
+        out[start:end] = torch.einsum("hqk,khd->qhd", scores.softmax(-1), vs)
+    return out
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires flash-attn")
+@pytest.mark.parametrize("num_kv_heads", [NUM_HEADS, NUM_HEADS // 2])
+def test_window_attention_applies_sinks(vision_attn_env, num_kv_heads):
+    from vllm.model_executor.models.mimo_v2_omni import MiMoVisionAttention
+
+    torch.manual_seed(0)
+    attn = MiMoVisionAttention(
+        embed_dim=EMBED_DIM,
+        num_heads=NUM_HEADS,
+        num_kv_heads=num_kv_heads,
+        qk_channels=HEAD_DIM,
+        kv_channels=HEAD_DIM,
+        use_sink=True,
+        visual_token_window_size=WINDOW,
+    ).cuda()
+    attn.sinks.data.normal_()
+
+    total = sum(SEQ_LENS)
+    opts = dict(device="cuda", dtype=torch.bfloat16)
+    q = torch.randn(total, NUM_HEADS, HEAD_DIM, **opts)
+    k = torch.randn(total, num_kv_heads, HEAD_DIM, **opts)
+    v = torch.randn(total, num_kv_heads, HEAD_DIM, **opts)
+    cu_seqlens = torch.tensor(
+        [0, *torch.tensor(SEQ_LENS).cumsum(0).tolist()],
+        device="cuda",
+        dtype=torch.int32,
+    )
+
+    out = attn._forward_window_attn(q, k, v, cu_seqlens, max(SEQ_LENS))
+    ref = _reference(q, k, v, cu_seqlens, attn.sinks, attn.scale)
+
+    # bf16 attention lands at ~2e-3 here; dropping the sinks lands at ~1e-1.
+    error = ((out.float() - ref).norm() / ref.norm()).item()
+    assert error < 1e-2, f"sink-corrected output is off by {error:.2e}"

vllm/model_executor/models/mimo_v2_omni.py

@@ -193,7 +193,7 @@ class MiMoVisionAttention(nn.Module):
         # Rotary embeddings applied separately to Q and K
         self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True)
 
-        # Sink attention weights (loaded but not used in vLLM flash_attn)
+        # Per-head sink logits, applied in the window attention path.
         # The checkpoint stores these only for non-full-attention blocks
         self.use_sink = use_sink
         if use_sink:
@@ -214,21 +214,38 @@ class MiMoVisionAttention(nn.Module):
         cu_seqlens: torch.Tensor,
         max_seqlen: torch.Tensor,
     ) -> torch.Tensor:
-        """Window attention via flash_attn_varlen_func with window_size."""
-        from vllm.vllm_flash_attn import flash_attn_varlen_func
+        """Window attention with the per-head sink applied to key 0.
+
+        The reference adds ``sinks[h]`` to the logit of each sequence's first
+        key, which the Triton prefill kernel supports directly, so the softmax
+        normalizes over the biased scores in one pass.
+        """
+        from vllm.v1.attention.ops.triton_prefill_attention import (
+            context_attention_fwd,
+        )
 
         w = self.visual_token_window_size
-        output = flash_attn_varlen_func(
+        output = torch.empty_like(q)
+        head_start = self.tp_rank * self.num_heads_per_partition
+        sinks = (
+            self.sinks[head_start : head_start + self.num_heads_per_partition]
+            if self.sinks is not None
+            else None
+        )
+        context_attention_fwd(
             q,
             k,
             v,
-            cu_seqlens_q=cu_seqlens,
-            cu_seqlens_k=cu_seqlens,
-            max_seqlen_q=max_seqlen,
-            max_seqlen_k=max_seqlen,
+            output,
+            b_start_loc=cu_seqlens[:-1],
+            b_seq_len=cu_seqlens[1:] - cu_seqlens[:-1],
+            max_input_len=max_seqlen,
+            is_causal=False,
             softmax_scale=self.scale,
-            causal=False,
-            window_size=[w, w],
+            sliding_window_q=w,
+            sliding_window_k=w,
+            sinks=sinks,
+            sinks_bias_key0=True,
         )
         return output
 

vllm/v1/attention/ops/triton_prefill_attention.py

@@ -58,6 +58,7 @@ def _fwd_kernel(
     IS_CAUSAL: tl.constexpr,
     SLIDING_WINDOW_Q: tl.constexpr,
     SLIDING_WINDOW_K: tl.constexpr,
+    SINKS_BIAS_KEY0: tl.constexpr,
     USE_SINKS: tl.constexpr,
     Lk: tl.constexpr,
 ):
@@ -98,8 +99,15 @@ def _fwd_kernel(
     # initialize pointer to m and l
     if USE_SINKS:
         sink = tl.load(Sinks + cur_head) * 1.4426950408889634
-        m_i = tl.full([BLOCK_M], sink, dtype=tl.float32)
-        l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
+        if SINKS_BIAS_KEY0:
+            # Sinks bias the logit of key 0, so the softmax starts empty and
+            # normalizes over the biased scores as usual.
+            m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+            l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+        else:
+            # Sinks are a null logit that only inflates the denominator.
+            m_i = tl.full([BLOCK_M], sink, dtype=tl.float32)
+            l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
     else:
         m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
         l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
@@ -160,6 +168,8 @@ def _fwd_kernel(
 
         qk = tl.dot(q, k)
         qk = tl.where(mask, qk * sm_scale, -1.0e8)
+        if USE_SINKS and SINKS_BIAS_KEY0:
+            qk = tl.where(mask & (pos_k == 0), qk + sink, qk)
         m_ij = tl.maximum(m_i, tl.max(qk, 1))
         qk -= m_ij[:, None]
         p = tl.math.exp2(qk)
@@ -217,6 +227,7 @@ def context_attention_fwd(
     sliding_window_q: int | None = None,
     sliding_window_k: int | None = None,
     sinks: torch.Tensor | None = None,
+    sinks_bias_key0: bool = False,
 ):
     """
     q, k, v: [b * s, head, head_dim]
@@ -268,6 +279,7 @@ def context_attention_fwd(
         SLIDING_WINDOW_Q=sliding_window_q,
         SLIDING_WINDOW_K=sliding_window_k,
         USE_SINKS=sinks is not None,
+        SINKS_BIAS_KEY0=sinks_bias_key0,
         num_warps=num_warps,
         num_stages=1,
         Lk=Lk,