| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Parity test for the fused sconv seq-metadata kernel. |
| 4 | |
| 5 | ``sconv_seq_metadata`` fills the per-token ``seq_idx`` (owning request) and |
| 6 | ``query_start`` (first x-row of that request) buffers in a single launch. Actual |
| 7 | tokens must match the searchsorted-based reference, while CUDA graph padding |
| 8 | must be initialized to safe zero values. |
| 9 | """ |
| 10 | |
| 11 | import pytest |
| 12 | import torch |
| 13 | |
| 14 | from vllm.models.inkling.nvidia.ops.sconv import sconv_seq_metadata |
| 15 | from vllm.models.inkling.nvidia.sconv_swa_attn import InklingSconvMetadataBuilder |
| 16 | |
| 17 | CASES = [ |
| 18 | # (query_lens, extra_pad_tokens) |
| 19 | ([1], 0), # bsz1 decode |
| 20 | ([1] * 8, 0), # uniform decode |
| 21 | ([1] * 8, 3), # uniform decode, padded tokens past the last request |
| 22 | ([2] * 4, 0), # uniform spec-decode |
| 23 | ([2048], 0), # single prefill |
| 24 | ([517, 1, 1, 33, 1, 128], 0), # mixed prefill/decode |
| 25 | ([517, 1, 1, 33, 1, 128], 5), # mixed, padded |
| 26 | ([1] * 500, 0), # many requests (deep binary search) |
| 27 | ] |
| 28 | |
| 29 | |
| 30 | def _ref(query_start_loc: torch.Tensor, num_reqs: int, num_tokens: int): |
| 31 | cu_seqlens = query_start_loc[: num_reqs + 1].to(torch.int64) |
| 32 | token_idx = torch.arange(num_tokens, device=cu_seqlens.device, dtype=torch.int64) |
| 33 | seq_idx = (torch.searchsorted(cu_seqlens, token_idx, right=True) - 1).clamp( |
| 34 | max=num_reqs - 1 |
| 35 | ) |
| 36 | return seq_idx.to(torch.int32), cu_seqlens[seq_idx].to(torch.int32) |
| 37 | |
| 38 | |
| 39 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 40 | @pytest.mark.parametrize("query_lens,extra_pad", CASES) |
| 41 | def test_sconv_seq_metadata_matches_searchsorted(query_lens, extra_pad): |
| 42 | device = "cuda" |
| 43 | num_reqs = len(query_lens) |
| 44 | query_start_loc = torch.tensor( |
| 45 | [0] + list(torch.tensor(query_lens).cumsum(0)), dtype=torch.int32 |
| 46 | ).to(device) |
| 47 | num_actual_tokens = int(query_start_loc[-1]) |
| 48 | num_padded_tokens = num_actual_tokens + extra_pad |
| 49 | |
| 50 | ref_seq, ref_qs = _ref(query_start_loc, num_reqs, num_actual_tokens) |
| 51 | |
| 52 | seq_idx = torch.full((num_padded_tokens,), -1, dtype=torch.int32, device=device) |
| 53 | query_start = torch.full_like(seq_idx, -1) |
| 54 | sconv_seq_metadata( |
| 55 | query_start_loc, |
| 56 | num_reqs, |
| 57 | num_actual_tokens, |
| 58 | seq_idx, |
| 59 | query_start, |
| 60 | num_padded_tokens, |
| 61 | ) |
| 62 | |
| 63 | torch.testing.assert_close(seq_idx[:num_actual_tokens], ref_seq, rtol=0, atol=0) |
| 64 | torch.testing.assert_close(query_start[:num_actual_tokens], ref_qs, rtol=0, atol=0) |
| 65 | assert torch.count_nonzero(seq_idx[num_actual_tokens:]) == 0 |
| 66 | assert torch.count_nonzero(query_start[num_actual_tokens:]) == 0 |
| 67 | |
| 68 | |
| 69 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 70 | def test_sconv_metadata_reuses_padded_static_buffers(): |
| 71 | device = torch.device("cuda") |
| 72 | builder = object.__new__(InklingSconvMetadataBuilder) |
| 73 | builder.seq_idx_buffer = torch.empty(8, dtype=torch.int32, device=device) |
| 74 | builder.query_start_buffer = torch.empty(8, dtype=torch.int32, device=device) |
| 75 | |
| 76 | class CommonMetadata: |
| 77 | num_reqs = 1 |
| 78 | num_actual_tokens = 8 |
| 79 | query_start_loc = torch.tensor([0, 3], dtype=torch.int32, device=device) |
| 80 | query_start_loc_cpu = torch.tensor([0, 3], dtype=torch.int32) |
| 81 | block_table_tensor = torch.zeros((1, 1), dtype=torch.int32, device=device) |
| 82 | slot_mapping = torch.tensor( |
| 83 | [0, 1, 2, -1, -1, -1, -1, -1], |
| 84 | dtype=torch.int64, |
| 85 | device=device, |
| 86 | ) |
| 87 | |
| 88 | common = CommonMetadata() |
| 89 | first = builder.build(0, common) |
| 90 | pointers = tuple( |
| 91 | tensor.data_ptr() |
| 92 | for tensor in ( |
| 93 | first.block_table, |
| 94 | first.slot_mapping, |
| 95 | first.seq_idx, |
| 96 | first.query_start, |
| 97 | ) |
| 98 | ) |
| 99 | |
| 100 | common.query_start_loc[1] = 5 |
| 101 | common.query_start_loc_cpu[1] = 5 |
| 102 | common.slot_mapping[:5] = torch.arange(5, dtype=torch.int64, device=device) |
| 103 | second = builder.build(0, common) |
| 104 | |
| 105 | assert second.slot_mapping.shape == (8,) |
| 106 | assert second.seq_idx.shape == (8,) |
| 107 | assert second.query_start.shape == (8,) |
| 108 | assert pointers == tuple( |
| 109 | tensor.data_ptr() |
| 110 | for tensor in ( |
| 111 | second.block_table, |
| 112 | second.slot_mapping, |
| 113 | second.seq_idx, |
| 114 | second.query_start, |
| 115 | ) |
| 116 | ) |
| 117 | assert torch.count_nonzero(second.seq_idx[5:]) == 0 |
| 118 | assert torch.count_nonzero(second.query_start[5:]) == 0 |
| 119 | assert torch.all(second.slot_mapping[5:] == -1) |