[MM] Fix compiled ViT attention output layouts (#58182)
Signed-off-by: Jakub Zakrzewski <jzakrzewski@nvidia.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@mistral.ai>
Jakub Zakrzewski committed
Sep 29, 2026 at 15:27 UTC
cd947f5c629f16fb1ac8ac243481ce82bfa6f5a2
2 files changed
+116
-15
tests/v1/attention/test_vit_attn_wrappers.py
new
+87
@@ -0,0 +1,87 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import sys
5
+from types import ModuleType, SimpleNamespace
6
+
7
+import pytest
8
+import torch
9
+from torch._subclasses.fake_tensor import FakeTensorMode
10
+
11
+from vllm.platforms import current_platform
12
+from vllm.v1.attention.ops import vit_attn_wrappers
13
+
14
+
15
+def _noncontiguous_bshd(device: torch.device | None = None) -> torch.Tensor:
16
+ return (
17
+ torch.arange(120, dtype=torch.float32, device=device)
18
+ .reshape(1, 3, 5, 8)
19
+ .transpose(1, 2)
20
+ )
21
+
22
+
23
+def test_registered_attention_wrapper_fake_layouts():
24
+ with FakeTensorMode():
25
+ q = _noncontiguous_bshd()
26
+ cu_seqlens = torch.tensor([0, 5], dtype=torch.int32)
27
+ max_seqlen = torch.tensor(5, dtype=torch.int32)
28
+ flashinfer_cu_seqlens = torch.tensor([0, 5, 0, 5], dtype=torch.int32)
29
+ sequence_lengths = torch.tensor([5], dtype=torch.int32)
30
+
31
+ outputs = (
32
+ torch.ops.vllm.flash_attn_maxseqlen_wrapper(
33
+ q, q, q, 1, False, None, None, cu_seqlens, max_seqlen
34
+ ),
35
+ torch.ops.vllm.triton_attn_wrapper(
36
+ q, q, q, 1, None, cu_seqlens, max_seqlen
37
+ ),
38
+ torch.ops.vllm.torch_sdpa_wrapper(q, q, q, None, None, False),
39
+ torch.ops.vllm.flashinfer_wrapper(
40
+ q,
41
+ q,
42
+ q,
43
+ 1.0,
44
+ torch.empty(1, dtype=torch.uint8),
45
+ flashinfer_cu_seqlens,
46
+ max_seqlen,
47
+ sequence_lengths,
48
+ None,
49
+ None,
50
+ None,
51
+ None,
52
+ ),
53
+ )
54
+
55
+ assert all(output.is_contiguous() for output in outputs)
56
+
57
+
58
+def test_aiter_flash_attn_omits_out_and_returns_contiguous_output(
59
+ monkeypatch: pytest.MonkeyPatch,
60
+):
61
+ aiter = ModuleType("aiter")
62
+
63
+ def flash_attn_varlen_func(q, k, v, **kwargs):
64
+ assert "out" not in kwargs
65
+ return q
66
+
67
+ aiter.__dict__["flash_attn_varlen_func"] = flash_attn_varlen_func
68
+ monkeypatch.setitem(sys.modules, "aiter", aiter)
69
+ monkeypatch.setattr(
70
+ vit_attn_wrappers, "current_platform", SimpleNamespace(is_rocm=lambda: True)
71
+ )
72
+ q = _noncontiguous_bshd()
73
+ assert not q.is_contiguous()
74
+
75
+ output = vit_attn_wrappers.flash_attn_maxseqlen_wrapper(q, q, q, 1, True, None)
76
+
77
+ assert torch.equal(output, q)
78
+ assert output.is_contiguous()
79
+
80
+
81
+@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires GPU")
82
+def test_sdpa_output_layout_without_cu_seqlens():
83
+ q = _noncontiguous_bshd(torch.device("cuda"))
84
+
85
+ output = vit_attn_wrappers.vit_torch_sdpa_wrapper(q, q, q)
86
+
87
+ assert output.is_contiguous()
vllm/v1/attention/ops/vit_attn_wrappers.py
+29
-15
@@ -58,13 +58,15 @@ def flash_attn_maxseqlen_wrapper(
58
max_seqlen = max_seqlen.item()
59
60
q, k, v = (einops.rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
61
-
62
- if current_platform.is_cuda() and torch.cuda.is_current_stream_capturing():
63
- # workaround for encoder CUDA-graph replay paddings with NaNs edge case
64
- kwargs["out"] = torch.zeros(
65
- q.shape[0], q.shape[1], v.shape[-1], dtype=q.dtype, device=q.device
66
- )
67
-
61
+ supports_out = not current_platform.is_rocm()
62
+ if supports_out:
63
+ if current_platform.is_cuda() and torch.cuda.is_current_stream_capturing():
64
+ # workaround for encoder CUDA-graph replay paddings with NaNs edge case
65
+ kwargs["out"] = torch.zeros(
66
+ q.shape[0], q.shape[1], v.shape[-1], dtype=q.dtype, device=q.device
67
+ )
68
+ else:
69
+ kwargs["out"] = torch.empty_like(q, memory_format=torch.contiguous_format)
70
output = flash_attn_varlen_func(
71
q,
72
k,
@@ -78,6 +80,8 @@ def flash_attn_maxseqlen_wrapper(
80
softmax_scale=scale,
81
**kwargs,
82
)
83
+ if not supports_out:
84
+ output = output.contiguous()
85
context_layer = einops.rearrange(output, "(b s) h d -> b s h d", b=batch_size)
86
return context_layer
87
@@ -93,7 +97,7 @@ def flash_attn_maxseqlen_wrapper_fake(
97
cu_seqlens: torch.Tensor | None = None,
98
max_seqlen: torch.Tensor | None = None,
99
) -> torch.Tensor:
96
- return torch.empty_like(q)
100
+ return torch.empty_like(q, memory_format=torch.contiguous_format)
101
102
103
direct_register_custom_op(
@@ -179,7 +183,7 @@ def triton_attn_wrapper(
183
max_seqlen = max_seqlen.item()
184
185
q, k, v = (einops.rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
182
- output = torch.empty_like(q)
186
+ output = torch.empty_like(q, memory_format=torch.contiguous_format)
187
context_attention_fwd(
188
q,
189
k,
@@ -207,7 +211,7 @@ def triton_attn_wrapper_fake(
211
cu_seqlens: torch.Tensor | None = None,
212
max_seqlen: torch.Tensor | None = None,
213
) -> torch.Tensor:
210
- return torch.empty_like(q)
214
+ return torch.empty_like(q, memory_format=torch.contiguous_format)
215
216
217
direct_register_custom_op(
@@ -281,7 +285,7 @@ def torch_sdpa_wrapper(
285
v = v.contiguous()
286
287
if cu_seqlens is None:
284
- return apply_sdpa(q, k, v, scale=scale, enable_gqa=enable_gqa)
288
+ return apply_sdpa(q, k, v, scale=scale, enable_gqa=enable_gqa).contiguous()
289
290
outputs = []
291
@@ -302,11 +306,11 @@ def torch_sdpa_wrapper_fake(
306
q: torch.Tensor,
307
k: torch.Tensor,
308
v: torch.Tensor,
305
- scale: float | None,
306
- cu_seqlens: torch.Tensor | None,
309
+ scale: float | None = None,
310
+ cu_seqlens: torch.Tensor | None = None,
311
enable_gqa: bool = False,
312
) -> torch.Tensor:
309
- return torch.empty_like(q)
313
+ return torch.empty_like(q, memory_format=torch.contiguous_format)
314
315
316
direct_register_custom_op(
@@ -368,6 +372,11 @@ def flashinfer_wrapper(
372
with gpu_sync_allowed():
373
max_seqlen = max_seqlen.item()
374
375
+ output_buffer = torch.empty_like(
376
+ q,
377
+ dtype=o_data_type,
378
+ memory_format=torch.contiguous_format,
379
+ )
380
output, _ = cudnn_batch_prefill_with_kv_cache(
381
q,
382
k,
@@ -388,6 +397,7 @@ def flashinfer_wrapper(
397
k_scale=k_scale,
398
v_scale=v_scale,
399
o_data_type=o_data_type,
400
+ out=output_buffer,
401
)
402
403
if is_reshaped:
@@ -410,7 +420,11 @@ def vit_flashinfer_wrapper_fake(
420
v_scale: torch.Tensor | None = None,
421
o_data_type: torch.dtype | None = None,
422
) -> torch.Tensor:
413
- return torch.empty_like(q, dtype=o_data_type or q.dtype)
423
+ return torch.empty_like(
424
+ q,
425
+ dtype=o_data_type,
426
+ memory_format=torch.contiguous_format,
427
+ )
428
429
430
direct_register_custom_op(