@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 226 lines 7.36 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from types import SimpleNamespace
5
6 import pytest
7 import torch
8
9 from vllm.platforms import current_platform
10 from vllm.triton_utils import tl, triton
11 from vllm.utils.torch_utils import set_random_seed
12 from vllm.v1.attention.backends.mla import rocm_aiter_mla_sparse as sparse_mod
13 from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
14 _use_rocm_sparse_triton,
15 fit_kpool_indices_to_aiter,
16 )
17 from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
18 _sparse_kv_row_offset,
19 _validate_dsv4_sparse_dims,
20 _validate_sparse_dims,
21 )
22
23
24 @triton.jit
25 def _store_sparse_kv_row_offset_kernel(slot_ptr, output_ptr, stride: tl.constexpr):
26 slot = tl.load(slot_ptr)
27 tl.store(output_ptr, _sparse_kv_row_offset(slot, stride))
28
29
30 def _fit_kpool_indices_reference(
31 token_indices: torch.Tensor, topk_tokens: int
32 ) -> torch.Tensor:
33 history = token_indices[:, :topk_tokens]
34 tail = token_indices[:, topk_tokens:]
35 valid_history = (history >= 0).sum(dim=1)
36 valid_tail = (tail >= 0).sum(dim=1)
37 keep_history = torch.minimum(valid_history, topk_tokens - valid_tail)
38
39 columns = torch.arange(topk_tokens, device=token_indices.device).unsqueeze(0)
40 tail_offsets = columns - keep_history.unsqueeze(1)
41 tail_values = torch.gather(
42 tail, 1, tail_offsets.clamp(min=0, max=tail.shape[1] - 1)
43 )
44 output = torch.where(columns < keep_history.unsqueeze(1), history, tail_values)
45 valid_output = columns < (keep_history + valid_tail).unsqueeze(1)
46 return torch.where(valid_output, output, -1)
47
48
49 @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
50 def test_fit_kpool_indices_preserves_tail_and_best_history():
51 token_indices = torch.tensor(
52 [
53 [10, 9, 8, 7, 6, 5, 100, 101],
54 [10, 9, 8, -1, -1, -1, 100, -1],
55 [-1, -1, -1, -1, -1, -1, -1, -1],
56 ],
57 dtype=torch.int32,
58 device="cuda",
59 )
60
61 fitted = fit_kpool_indices_to_aiter(token_indices, topk_tokens=6)
62
63 assert fitted.tolist() == [
64 [10, 9, 8, 7, 100, 101],
65 [10, 9, 8, 100, -1, -1],
66 [-1, -1, -1, -1, -1, -1],
67 ]
68
69
70 @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
71 @pytest.mark.parametrize(
72 ("num_tokens", "topk_tokens", "tail_width"),
73 [(1, 4, 1), (7, 6, 2), (13, 128, 128), (3, 2048, 128)],
74 )
75 def test_fit_kpool_indices_matches_eager_reference(num_tokens, topk_tokens, tail_width):
76 """The Triton fit must reproduce the eager packing rule exactly."""
77 set_random_seed(topk_tokens + num_tokens)
78 width = topk_tokens + tail_width
79 token_indices = torch.randint(0, 4096, (num_tokens, width), dtype=torch.int32)
80 valid = torch.rand((num_tokens, width)) < 0.5
81 valid[0] = False
82 valid[-1] = True
83 token_indices = torch.where(valid, token_indices, -1).cuda()
84
85 fitted = fit_kpool_indices_to_aiter(token_indices, topk_tokens=topk_tokens)
86
87 torch.testing.assert_close(
88 fitted, _fit_kpool_indices_reference(token_indices, topk_tokens)
89 )
90
91
92 def test_fit_kpool_indices_exact_width_is_noop():
93 token_indices = torch.tensor([[3, 2, 1, -1]], dtype=torch.int32)
94
95 fitted = fit_kpool_indices_to_aiter(token_indices, topk_tokens=4)
96
97 assert fitted.data_ptr() == token_indices.data_ptr()
98
99
100 def test_fit_kpool_indices_rejects_narrow_input():
101 with pytest.raises(ValueError, match="at least topk_tokens"):
102 fit_kpool_indices_to_aiter(
103 torch.zeros((1, 3), dtype=torch.int32), topk_tokens=4
104 )
105
106
107 @pytest.mark.parametrize(
108 (
109 "kv_cache_dtype",
110 "head_size",
111 "num_prefills",
112 "num_decodes",
113 "num_decode_tokens",
114 "max_query_len",
115 "expected",
116 ),
117 [
118 ("auto", 512, 1, 0, 0, 32, True),
119 ("auto", 512, 1, 2, 2, 32, True),
120 ("auto", 512, 0, 2, 2, 1, True),
121 ("fp8", 512, 1, 0, 0, 32, False),
122 ("auto", 576, 1, 0, 0, 32, False),
123 ("auto", 512, 0, 2, 4, 2, True),
124 ("auto", 512, 0, 2, 12, 6, True),
125 ("auto", 512, 0, 0, 0, 0, False),
126 ],
127 )
128 def test_rocm_sparse_triton_route(
129 kv_cache_dtype,
130 head_size,
131 num_prefills,
132 num_decodes,
133 num_decode_tokens,
134 max_query_len,
135 expected,
136 ):
137 """Validate Triton routing for prefill, decode, and MTP verification."""
138 assert (
139 _use_rocm_sparse_triton(
140 kv_cache_dtype=kv_cache_dtype,
141 head_size=head_size,
142 kv_lora_rank=512,
143 num_prefills=num_prefills,
144 num_decodes=num_decodes,
145 num_decode_tokens=num_decode_tokens,
146 max_query_len=max_query_len,
147 )
148 is expected
149 )
150
151
152 @pytest.mark.parametrize("num_heads", [8, 12])
153 def test_rocm_sparse_triton_route_preserves_padded_sinks(monkeypatch, num_heads):
154 captured = {}
155
156 def fake_rocm_sparse_attn_prefill(**kwargs):
157 output = kwargs["output"]
158 captured["attn_sink"] = kwargs["attn_sink"]
159 output.copy_(
160 captured["attn_sink"].to(output.dtype).view(1, -1, 1).expand_as(output)
161 )
162
163 monkeypatch.setattr(
164 sparse_mod, "rocm_sparse_attn_prefill", fake_rocm_sparse_attn_prefill
165 )
166
167 impl = object.__new__(sparse_mod.ROCMAiterMLASparseImpl)
168 impl.num_heads = num_heads
169 impl.kv_lora_rank = 512
170 impl.kv_cache_dtype = "auto"
171 impl.scale = 512**-0.5
172 impl.sinks = torch.arange(num_heads, dtype=torch.float32)
173
174 q = torch.zeros(2, 16, 512, dtype=torch.bfloat16)
175 kv = torch.zeros(4, 1, 512, dtype=torch.bfloat16)
176 metadata = SimpleNamespace(
177 attn_out_dtype=torch.bfloat16,
178 num_prefills=1,
179 num_decodes=0,
180 num_decode_tokens=0,
181 max_query_len=2,
182 paged_kv_indices=torch.empty(0, dtype=torch.int32),
183 paged_kv_indptr=torch.zeros(3, dtype=torch.int32),
184 )
185
186 output, lse = impl._forward_mla(SimpleNamespace(), q, kv, metadata)
187
188 if num_heads == 8:
189 expected_sinks = impl.sinks.repeat_interleave(2)
190 else:
191 expected_sinks = torch.cat((impl.sinks, impl.sinks[:4]))
192 torch.testing.assert_close(captured["attn_sink"], expected_sinks)
193 assert output.shape == (2, num_heads, 512)
194 torch.testing.assert_close(
195 output[:, :, 0].float(),
196 impl.sinks.expand(2, -1),
197 )
198 assert lse is None
199
200
201 def test_rocm_sparse_attention_accepts_glm_nope_dimensions():
202 _validate_sparse_dims(512, 512, 0, "test")
203
204
205 def test_rocm_sparse_attention_rejects_inconsistent_dimensions():
206 with pytest.raises(AssertionError, match="expected head_dim"):
207 _validate_sparse_dims(511, 512, 0, "test")
208
209
210 def test_dsv4_sparse_attention_keeps_layout_constraint():
211 _validate_dsv4_sparse_dims(512, 448, 64, "test")
212 with pytest.raises(AssertionError, match="expects 448 NoPE dims"):
213 _validate_dsv4_sparse_dims(512, 512, 0, "test")
214
215
216 @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
217 def test_sparse_prefill_kv_row_offset_does_not_overflow_int32():
218 # GLM's 640-token pages cross the signed-int32 address boundary at block
219 # 6554 for a 512-element KV row. The production kernel must promote the
220 # slot before multiplying by the row stride.
221 slot = torch.tensor([6554 * 640], dtype=torch.int32, device="cuda")
222 output = torch.empty(1, dtype=torch.int64, device="cuda")
223
224 _store_sparse_kv_row_offset_kernel[(1,)](slot, output, stride=512)
225
226 assert output.item() == 6554 * 640 * 512