| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | |
| 7 | pytest.importorskip("cutlass") |
| 8 | |
| 9 | from vllm.model_executor.layers.attention.sparse_mla_attention import ( |
| 10 | _build_topk_mask, |
| 11 | _topk_mask_shape, |
| 12 | ) |
| 13 | |
| 14 | |
| 15 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 16 | def test_build_topk_mask_single_request_matches_generic_path() -> None: |
| 17 | topk = torch.tensor( |
| 18 | [[0, 31, 32, 63, -1], [1, 64, 127, -1, -1]], |
| 19 | dtype=torch.int32, |
| 20 | device="cuda", |
| 21 | ) |
| 22 | |
| 23 | num_words = (128 + 31) // 32 |
| 24 | single_out = torch.zeros(1, 2, num_words, dtype=torch.int32, device="cuda") |
| 25 | generic_out = torch.zeros(2, 1, num_words, dtype=torch.int32, device="cuda") |
| 26 | single_req = _build_topk_mask([topk], [2], 2, 128, single_out) |
| 27 | generic = _build_topk_mask([topk[:1], topk[1:]], [1, 1], 1, 128, generic_out) |
| 28 | |
| 29 | torch.testing.assert_close(single_req[0, 0], generic[0, 0]) |
| 30 | torch.testing.assert_close(single_req[0, 1], generic[1, 0]) |
| 31 | |
| 32 | |
| 33 | def test_topk_mask_rows_are_aligned_for_vectorized_loads() -> None: |
| 34 | assert _topk_mask_shape(2, 129, 129) == (2, 256, 8) |
| 35 | assert _topk_mask_shape(2, 129, 129, reserve_key_starts_word=True) == ( |
| 36 | 2, |
| 37 | 256, |
| 38 | 8, |
| 39 | ) |
| 40 | |
| 41 | |
| 42 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 43 | def test_build_topk_mask_preserves_aligned_row_storage() -> None: |
| 44 | shape = _topk_mask_shape(1, 1, 129) |
| 45 | out = torch.zeros(shape, dtype=torch.int32, device="cuda") |
| 46 | topk = torch.tensor([[0, 128]], dtype=torch.int32, device="cuda") |
| 47 | |
| 48 | mask = _build_topk_mask([topk], [1], 1, 129, out) |
| 49 | |
| 50 | assert mask.shape == (1, 1, 8) |
| 51 | assert mask[0, 0].tolist() == [1, 0, 0, 0, 1, 0, 0, 0] |