| 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() |