@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 346 lines 11.9 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import random
4
5 import pytest
6 import torch
7
8 from vllm import _custom_ops as ops
9 from vllm.platforms import current_platform
10
11 DEVICE_TYPE = current_platform.device_type
12
13
14 def round_up(x, base):
15 return ((x + base - 1) // base) * base
16
17
18 def CEILDIV(x, y):
19 return (x + y - 1) // y
20
21
22 def sample_data(num_experts, max_loras, num_tokens, topk_num):
23 topk_ids = torch.zeros((num_tokens, topk_num), dtype=torch.int32)
24 token_lora_mapping = torch.zeros((num_tokens,), dtype=torch.int32)
25
26 for i in range(num_tokens):
27 pool = list(range(num_experts))
28 random.shuffle(pool)
29 for j in range(topk_num):
30 topk_ids[i, j] = pool[j]
31 token_lora_mapping[i] = random.randint(0, max_loras - 1)
32
33 return topk_ids.to(DEVICE_TYPE), token_lora_mapping.to(DEVICE_TYPE)
34
35
36 @pytest.mark.parametrize("num_tokens", [100, 200, 1024, 4096]) # 81920
37 @pytest.mark.parametrize("topk_num", [6])
38 @pytest.mark.parametrize("num_experts", [64, 128, 256, 512])
39 @pytest.mark.parametrize("max_loras", [2, 32])
40 @pytest.mark.parametrize("block_size", [16])
41 def test_moe_lora_align_block_size(
42 num_tokens, topk_num, num_experts, max_loras, block_size
43 ):
44 _run_moe_lora_align_block_size_case(
45 num_tokens, topk_num, num_experts, max_loras, block_size
46 )
47
48
49 def test_moe_lora_align_block_size_large_path_counts():
50 _run_moe_lora_align_block_size_case(100, 12, 64, 4, 16)
51
52
53 def _run_moe_lora_align_block_size_case(
54 num_tokens, topk_num, num_experts, max_loras, block_size
55 ):
56 # sample data
57 random.seed(1)
58 topk_ids, token_lora_mapping = sample_data(
59 num_experts, max_loras, num_tokens, topk_num
60 )
61
62 # compute paddings
63 max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
64 max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
65 if topk_ids.numel() < num_experts:
66 max_num_tokens_padded = topk_ids.numel() * block_size
67 max_num_m_blocks = CEILDIV(max_num_tokens_padded, block_size)
68
69 # init output tensors
70 sorted_token_ids = torch.full(
71 (max_loras * max_num_tokens_padded,),
72 topk_ids.numel(),
73 dtype=torch.int32,
74 device=DEVICE_TYPE,
75 )
76 expert_ids = torch.full(
77 (max_loras * max_num_m_blocks,),
78 num_experts,
79 dtype=torch.int32,
80 device=DEVICE_TYPE,
81 )
82 num_tokens_post_pad = torch.zeros(
83 (max_loras,), dtype=torch.int32, device=DEVICE_TYPE
84 )
85 adapter_enabled = torch.ones(
86 (max_loras + 1,), dtype=torch.int32, device=DEVICE_TYPE
87 )
88 lora_ids = torch.arange(max_loras + 2, dtype=torch.int32, device=DEVICE_TYPE)
89
90 # call kernel
91 ops.moe_lora_align_block_size(
92 topk_ids,
93 token_lora_mapping,
94 num_experts,
95 block_size,
96 max_loras,
97 max_num_tokens_padded,
98 max_num_m_blocks,
99 sorted_token_ids,
100 expert_ids,
101 num_tokens_post_pad,
102 adapter_enabled,
103 lora_ids,
104 )
105
106 # verify values
107 expert_ids = expert_ids.view(max_loras, -1)
108 sorted_token_ids = sorted_token_ids.view(max_loras, -1, block_size)
109 topk_ids_cpu = topk_ids.cpu()
110 token_lora_mapping_cpu = token_lora_mapping.cpu()
111
112 for lora_idx in range(max_loras):
113 expert_counts = torch.bincount(
114 topk_ids_cpu[token_lora_mapping_cpu == lora_idx].flatten(),
115 minlength=num_experts,
116 )
117 expected_post_pad = round_up(expert_counts, block_size).sum().item()
118 assert num_tokens_post_pad[lora_idx].item() == expected_post_pad
119
120 for token_idx in range(sorted_token_ids.size(1)):
121 block = sorted_token_ids[lora_idx][token_idx]
122 indices = block[block != topk_ids.numel()]
123 if indices.numel() > 0:
124 expert_id = expert_ids[lora_idx][token_idx]
125 assert torch.all(topk_ids.view(-1)[indices] == expert_id)
126
127
128 # Sentinel values for the regression tests below. Distinctive out-of-domain
129 # ints so that "kernel never wrote this slot" is directly observable: the
130 # kernel only ever writes a real expert id in [0, num_experts) or -1
131 # (expert_ids), a token index or the `numel` padding value (sorted_token_ids),
132 # and a block-aligned cumsum count (num_tokens_post_pad).
133 SENTINEL_EXPERT = -2
134 SENTINEL_TOKEN = -7
135 SENTINEL_NPAD = -13
136
137
138 def _build_and_run_align(
139 *,
140 num_lora_tokens,
141 num_base_tokens,
142 max_loras,
143 num_experts=64,
144 topk_num=6,
145 block_size=16,
146 lora_ids_override=None,
147 disabled_slots=(),
148 seed=1,
149 ):
150 """Build inputs the way ``LoRAKernelMeta.prepare_tensors`` does, run
151 ``moe_lora_align_block_size``, and return a dict of result tensors plus
152 derived sizes. Output buffers are pre-filled with ``SENTINEL_*`` so
153 callers can assert which slots the kernel did / did not touch.
154
155 Tokens are assigned to LoRA slot 0 (first ``num_lora_tokens``) then -1
156 (remaining ``num_base_tokens``), matching the "mixed base + 1 LoRA"
157 shape used to repro vllm-project/vllm#32235.
158
159 ``lora_ids_override``: optional 1-D int tensor of length ``max_loras+1``
160 used verbatim. Default mirrors ``prepare_tensors`` (sorted-unique into
161 the head, -1 tail).
162 ``disabled_slots``: iterable of slot indices to clear in ``adapter_enabled``.
163 """
164 random.seed(seed)
165 num_tokens = num_lora_tokens + num_base_tokens
166 assert num_tokens > 0, "test requires at least one token"
167
168 topk_ids = torch.zeros((num_tokens, topk_num), dtype=torch.int32)
169 token_lora_mapping = torch.empty((num_tokens,), dtype=torch.int32)
170 for i in range(num_tokens):
171 pool = list(range(num_experts))
172 random.shuffle(pool)
173 for j in range(topk_num):
174 topk_ids[i, j] = pool[j]
175 token_lora_mapping[i] = 0 if i < num_lora_tokens else -1
176 topk_ids = topk_ids.to(DEVICE_TYPE)
177 token_lora_mapping = token_lora_mapping.to(DEVICE_TYPE)
178
179 max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
180 max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
181 if topk_ids.numel() < num_experts:
182 max_num_tokens_padded = topk_ids.numel() * block_size
183 max_num_m_blocks = CEILDIV(max_num_tokens_padded, block_size)
184
185 if lora_ids_override is None:
186 lora_ids = torch.full(
187 (max_loras + 1,), -1, dtype=torch.int32, device=DEVICE_TYPE
188 )
189 unique_ids = torch.unique(token_lora_mapping, sorted=True)
190 lora_ids[: unique_ids.numel()] = unique_ids.to(torch.int32)
191 else:
192 assert lora_ids_override.numel() == max_loras + 1
193 lora_ids = lora_ids_override.to(dtype=torch.int32, device=DEVICE_TYPE)
194
195 adapter_enabled = torch.ones(
196 (max_loras + 1,), dtype=torch.int32, device=DEVICE_TYPE
197 )
198 for slot in disabled_slots:
199 adapter_enabled[slot] = 0
200
201 sorted_token_ids = torch.full(
202 (max_loras * max_num_tokens_padded,),
203 SENTINEL_TOKEN,
204 dtype=torch.int32,
205 device=DEVICE_TYPE,
206 )
207 expert_ids = torch.full(
208 (max_loras * max_num_m_blocks,),
209 SENTINEL_EXPERT,
210 dtype=torch.int32,
211 device=DEVICE_TYPE,
212 )
213 num_tokens_post_pad = torch.full(
214 (max_loras,), SENTINEL_NPAD, dtype=torch.int32, device=DEVICE_TYPE
215 )
216
217 ops.moe_lora_align_block_size(
218 topk_ids,
219 token_lora_mapping,
220 num_experts,
221 block_size,
222 max_loras,
223 max_num_tokens_padded,
224 max_num_m_blocks,
225 sorted_token_ids,
226 expert_ids,
227 num_tokens_post_pad,
228 adapter_enabled,
229 lora_ids,
230 )
231
232 return {
233 "lora_ids": lora_ids,
234 "sorted_token_ids": sorted_token_ids,
235 "expert_ids": expert_ids,
236 "num_tokens_post_pad": num_tokens_post_pad,
237 "max_num_tokens_padded": max_num_tokens_padded,
238 "block_size": block_size,
239 "max_loras": max_loras,
240 }
241
242
243 @pytest.mark.parametrize(
244 "max_loras",
245 [
246 1,
247 2,
248 ],
249 )
250 def test_moe_lora_align_block_size_mixed_base_and_lora(max_loras):
251 """Regression test for issue #32235: real LoRA slot must not be skipped
252 when ``active_lora_ids`` has -1 at position 0."""
253 out = _build_and_run_align(
254 num_lora_tokens=8, num_base_tokens=8, max_loras=max_loras
255 )
256
257 # Sanity check on the layout being tested.
258 assert out["lora_ids"][0].item() == -1, (
259 "prepare_tensors layout mismatch: -1 expected at position 0 for mixed batch"
260 )
261
262 real_slot = 0
263 post_pad = out["num_tokens_post_pad"][real_slot].item()
264 assert post_pad != SENTINEL_NPAD, (
265 f"num_tokens_post_pad[{real_slot}] was never written by the kernel; "
266 "the align kernel skipped the real LoRA slot."
267 )
268 assert (
269 0 < post_pad <= out["max_num_tokens_padded"]
270 and post_pad % out["block_size"] == 0
271 ), f"num_tokens_post_pad[{real_slot}]={post_pad} is not a valid block-aligned count"
272
273 expert_row = out["expert_ids"].view(max_loras, -1)[real_slot]
274 assert (expert_row != SENTINEL_EXPERT).all(), (
275 f"expert_ids row for slot {real_slot} has unwritten sentinel entries; "
276 "the align kernel skipped the real LoRA slot."
277 )
278
279 sorted_row = out["sorted_token_ids"].view(max_loras, -1)[real_slot]
280 assert (sorted_row != SENTINEL_TOKEN).all(), (
281 f"sorted_token_ids row for slot {real_slot} has unwritten sentinel "
282 "entries; the align kernel skipped the real LoRA slot."
283 )
284
285
286 def test_moe_lora_align_block_size_disabled_adapter_untouched():
287 """Disabled-adapter slot rows must remain untouched by all three align
288 kernels. Pins the invariant protected by the ``adapter_enabled`` guard
289 in ``lora_count_and_sort_expert_tokens_kernel``: without it the sort
290 kernel reads uninitialized ``token_mask`` values for disabled slots and
291 pollutes ``sorted_token_ids`` / ``cumsum_buffer``."""
292 max_loras = 1
293 out = _build_and_run_align(
294 num_lora_tokens=16,
295 num_base_tokens=0,
296 max_loras=max_loras,
297 disabled_slots=(0,),
298 )
299 # Sanity: slot 0 IS present in active_lora_ids (otherwise we would only
300 # exercise the lora_id == -1 / >= max_loras guards).
301 assert (out["lora_ids"] == 0).any().item()
302
303 assert out["num_tokens_post_pad"][0].item() == SENTINEL_NPAD, (
304 "num_tokens_post_pad[0] was modified for a disabled adapter slot."
305 )
306 expert_row = out["expert_ids"].view(max_loras, -1)[0]
307 assert (expert_row == SENTINEL_EXPERT).all(), (
308 "expert_ids row for disabled slot 0 was partially written."
309 )
310 # Row specifically protected by the sort-kernel adapter_enabled guard.
311 sorted_row = out["sorted_token_ids"].view(max_loras, -1)[0]
312 assert (sorted_row == SENTINEL_TOKEN).all(), (
313 "sorted_token_ids row for disabled slot 0 was polluted by the sort "
314 "kernel; lora_count_and_sort_expert_tokens_kernel must skip "
315 "adapter_enabled == 0 slots."
316 )
317
318
319 def test_moe_lora_align_block_size_lora_id_oob_guard():
320 """Regression test for the ``lora_id >= max_loras`` guard.
321
322 Production ``LoRAKernelMeta.prepare_tensors`` pre-fills the tail of
323 ``active_lora_ids`` with -1, so the existing ``lora_id == -1`` check
324 covers the extra slot. This test bypasses that invariant and injects
325 an out-of-range value (5 with max_loras=1) at the tail to verify the
326 explicit guard prevents OOB reads against ``adapter_enabled`` and
327 OOB writes against the max_loras-sized output buffers. Without the
328 guard, an illegal-memory-access would surface on the next CUDA sync.
329 """
330 max_loras = 1
331 lora_ids_override = torch.tensor([0, 5], dtype=torch.int32)
332 out = _build_and_run_align(
333 num_lora_tokens=16,
334 num_base_tokens=0,
335 max_loras=max_loras,
336 lora_ids_override=lora_ids_override,
337 )
338 # The .item() call below syncs and would surface any async
339 # illegal-memory-access from the OOB iteration.
340 assert out["num_tokens_post_pad"][0].item() != SENTINEL_NPAD, (
341 "real LoRA slot 0 was skipped by the align kernel"
342 )
343
344
345 if __name__ == "__main__":
346 pytest.main([__file__])