| 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 vllm.platforms import current_platform |
| 9 | |
| 10 | EMBED_DIM = 256 |
| 11 | NUM_HEADS = 4 |
| 12 | HEAD_DIM = 64 |
| 13 | WINDOW = 8 |
| 14 | # One sequence shorter than the window, one longer. |
| 15 | SEQ_LENS = [5, 37] |
| 16 | |
| 17 | |
| 18 | @pytest.fixture |
| 19 | def vision_attn_env(dist_init): |
| 20 | default_dtype = torch.get_default_dtype() |
| 21 | torch.set_default_dtype(torch.bfloat16) |
| 22 | yield |
| 23 | torch.set_default_dtype(default_dtype) |
| 24 | |
| 25 | |
| 26 | def _reference(q, k, v, cu_seqlens, sinks, scale): |
| 27 | """Dense windowed softmax with the sink added to each sequence's key 0.""" |
| 28 | groups = q.shape[1] // k.shape[1] |
| 29 | out = torch.empty_like(q, dtype=torch.float32) |
| 30 | for start, end in zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist()): |
| 31 | qs = q[start:end].float() |
| 32 | ks = k[start:end].float().repeat_interleave(groups, dim=1) |
| 33 | vs = v[start:end].float().repeat_interleave(groups, dim=1) |
| 34 | scores = torch.einsum("qhd,khd->hqk", qs, ks) * scale |
| 35 | scores[..., 0] += sinks.float().view(-1, 1) |
| 36 | pos = torch.arange(end - start, device=q.device) |
| 37 | outside = (pos.view(-1, 1) - pos.view(1, -1)).abs() > WINDOW |
| 38 | scores.masked_fill_(outside, -torch.inf) |
| 39 | out[start:end] = torch.einsum("hqk,khd->qhd", scores.softmax(-1), vs) |
| 40 | return out |
| 41 | |
| 42 | |
| 43 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="requires flash-attn") |
| 44 | @pytest.mark.parametrize("num_kv_heads", [NUM_HEADS, NUM_HEADS // 2]) |
| 45 | def test_window_attention_applies_sinks(vision_attn_env, num_kv_heads): |
| 46 | from vllm.model_executor.models.mimo_v2_omni import MiMoVisionAttention |
| 47 | |
| 48 | torch.manual_seed(0) |
| 49 | attn = MiMoVisionAttention( |
| 50 | embed_dim=EMBED_DIM, |
| 51 | num_heads=NUM_HEADS, |
| 52 | num_kv_heads=num_kv_heads, |
| 53 | qk_channels=HEAD_DIM, |
| 54 | kv_channels=HEAD_DIM, |
| 55 | use_sink=True, |
| 56 | visual_token_window_size=WINDOW, |
| 57 | ).cuda() |
| 58 | attn.sinks.data.normal_() |
| 59 | |
| 60 | total = sum(SEQ_LENS) |
| 61 | opts = dict(device="cuda", dtype=torch.bfloat16) |
| 62 | q = torch.randn(total, NUM_HEADS, HEAD_DIM, **opts) |
| 63 | k = torch.randn(total, num_kv_heads, HEAD_DIM, **opts) |
| 64 | v = torch.randn(total, num_kv_heads, HEAD_DIM, **opts) |
| 65 | cu_seqlens = torch.tensor( |
| 66 | [0, *torch.tensor(SEQ_LENS).cumsum(0).tolist()], |
| 67 | device="cuda", |
| 68 | dtype=torch.int32, |
| 69 | ) |
| 70 | |
| 71 | out = attn._forward_window_attn(q, k, v, cu_seqlens, max(SEQ_LENS)) |
| 72 | ref = _reference(q, k, v, cu_seqlens, attn.sinks, attn.scale) |
| 73 | |
| 74 | # bf16 attention lands at ~2e-3 here; dropping the sinks lands at ~1e-1. |
| 75 | error = ((out.float() - ref).norm() / ref.norm()).item() |
| 76 | assert error < 1e-2, f"sink-corrected output is off by {error:.2e}" |