@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 697 lines 21.4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import ast
5 from dataclasses import dataclass
6 from types import SimpleNamespace
7 from typing import Any, cast
8
9 import pytest
10
11 from vllm.model_executor.warmup.jit_warmup import (
12 JitWarmupRegistry,
13 VllmJitKernel,
14 WarmupIntRange,
15 get_ast_full_name,
16 zip_inputs,
17 )
18 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
19 triton_scalar_specialization_rep,
20 )
21
22
23 def _next_power_of_2(value: int) -> int:
24 return 1 << max(0, value - 1).bit_length()
25
26
27 def _round_up(value: int, *, multiple: int) -> int:
28 return ((value + multiple - 1) // multiple) * multiple
29
30
31 def _config(
32 *,
33 bias: int = 0,
34 disabled: bool = False,
35 name: str = "base",
36 vectorized: bool = False,
37 ) -> SimpleNamespace:
38 return SimpleNamespace(
39 bias=bias,
40 disabled=disabled,
41 name=name,
42 vectorized=vectorized,
43 )
44
45
46 class ToyKernel(VllmJitKernel["ToyKernel.CompileKey"]):
47 @dataclass(frozen=True)
48 class CompileKey:
49 block_size: int
50 work: int
51 vector_width: int
52 descriptor: tuple[object, ...]
53 enabled: bool
54
55 def dispatch( # type: ignore[override]
56 self,
57 *,
58 tokens: int,
59 cfg: Any,
60 lanes: int = 1,
61 mode: str = "default",
62 debug: int = 0,
63 ) -> CompileKey:
64 block_size = _next_power_of_2(tokens)
65 work: int = block_size * lanes + cfg.bias
66 return self.CompileKey(
67 block_size=block_size,
68 work=work,
69 vector_width=4 if cfg.vectorized and block_size >= 4 else 1,
70 descriptor=(
71 cfg.name,
72 mode,
73 -block_size,
74 block_size % 3,
75 block_size**2,
76 ),
77 enabled=not cfg.disabled,
78 )
79
80 def get_warmup_keys(self, max_tokens: int, cfg: Any) -> list[CompileKey]:
81 return self._trace_dispatch(self.dispatch)(
82 tokens=WarmupIntRange(1, max_tokens + 1),
83 cfg=cfg,
84 # This argument is intentionally unused by dispatch expressions.
85 debug=WarmupIntRange(0, 100),
86 )
87
88 def compile(self, compile_key: CompileKey) -> None:
89 pass
90
91
92 class RecordingToyKernel(ToyKernel):
93 def __init__(self) -> None:
94 self.compiled: list[ToyKernel.CompileKey] = []
95 super().__init__()
96
97 def compile(self, compile_key: ToyKernel.CompileKey) -> None:
98 self.compiled.append(compile_key)
99
100
101 @pytest.mark.parametrize(
102 ("value", "expected"),
103 [
104 (-(1 << 63), 1 << 31),
105 (-(1 << 31) - 1, (1 << 31) + 1),
106 (-(1 << 31), 16),
107 (0, 16),
108 (1, 1),
109 (2, 2),
110 (16, 16),
111 ((1 << 31) - 1, 2),
112 (1 << 31, 1 << 31),
113 ((1 << 31) + 1, (1 << 31) + 1),
114 ((1 << 63) - 1, (1 << 31) + 1),
115 (1 << 63, 1 << 63),
116 ((1 << 63) + 1, (1 << 63) + 1),
117 ((1 << 64) - 1, (1 << 63) + 1),
118 ],
119 )
120 def test_triton_scalar_specialization_rep(value: int, expected: int) -> None:
121 assert triton_scalar_specialization_rep(value) == expected
122
123
124 @pytest.mark.parametrize("value", [-(1 << 63) - 1, 1 << 64])
125 def test_triton_scalar_specialization_rep_rejects_out_of_range(value: int) -> None:
126 with pytest.raises(OverflowError, match="outside Triton's scalar range"):
127 triton_scalar_specialization_rep(value)
128
129
130 def test_trace_dispatch_expands_ranges_dedupes_and_ignores_unused_inputs() -> None:
131 cfg = _config()
132
133 assert ToyKernel().get_warmup_keys(5, cfg) == [
134 ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
135 ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
136 ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
137 ToyKernel.CompileKey(8, 8, 1, ("base", "default", -8, 2, 64), True),
138 ]
139
140
141 def test_warmup_range_uses_custom_advancement() -> None:
142 keys = ToyKernel()._trace_dispatch(ToyKernel().dispatch)(
143 tokens=WarmupIntRange(
144 1,
145 10,
146 advance=lambda value: _next_power_of_2(value) + 1,
147 ),
148 cfg=_config(),
149 )
150
151 assert [key.block_size for key in keys] == [1, 2, 4, 8, 16]
152
153
154 def test_warmup_range_validates_custom_advancement() -> None:
155 trace_dispatch = ToyKernel()._trace_dispatch(ToyKernel().dispatch)
156
157 with pytest.raises(ValueError, match="both step and advance"):
158 trace_dispatch(
159 tokens=WarmupIntRange(
160 1,
161 10,
162 step=2,
163 advance=lambda value: value + 1,
164 ),
165 cfg=_config(),
166 )
167 with pytest.raises(ValueError, match="must return a greater value"):
168 trace_dispatch(
169 tokens=WarmupIntRange(1, 10, advance=lambda value: value),
170 cfg=_config(),
171 )
172
173
174 def test_warmup_cases_support_lambda_advancement() -> None:
175 def warmup_inputs() -> dict[str, Any]:
176 value: Any = WarmupIntRange(1, 9, advance=lambda value: value * 2)
177 return dict(value=value)
178
179 assert list(ToyKernel()._expand_warmup_cases(warmup_inputs)) == [
180 {"value": 1},
181 {"value": 2},
182 {"value": 4},
183 {"value": 8},
184 ]
185
186
187 def test_compile_key_uses_defaults_locals_attributes_and_expressions() -> None:
188 cfg = _config(bias=3, disabled=True, name="cfg", vectorized=True)
189
190 assert ToyKernel().compile_key(
191 {
192 "tokens": 4,
193 "cfg": cfg,
194 "lanes": 2,
195 }
196 ) == ToyKernel.CompileKey(
197 block_size=4,
198 work=11,
199 vector_width=4,
200 descriptor=("cfg", "default", -4, 1, 16),
201 enabled=False,
202 )
203
204
205 def test_trace_dispatch_combines_zipped_rows_with_independent_values() -> None:
206 cfg = _config(vectorized=True)
207
208 keys = ToyKernel()._trace_dispatch(ToyKernel().dispatch)(
209 zip_inputs(
210 dict(tokens=1, mode="small"),
211 dict(tokens=4, mode="wide"),
212 ),
213 cfg=cfg,
214 lanes=(1, 2),
215 )
216
217 assert keys == [
218 ToyKernel.CompileKey(1, 1, 1, ("base", "small", -1, 1, 1), True),
219 ToyKernel.CompileKey(1, 2, 1, ("base", "small", -1, 1, 1), True),
220 ToyKernel.CompileKey(4, 4, 4, ("base", "wide", -4, 1, 16), True),
221 ToyKernel.CompileKey(4, 8, 4, ("base", "wide", -4, 1, 16), True),
222 ]
223
224
225 def test_trace_dispatch_filters_with_traced_predicate() -> None:
226 class PredicateKernel(ToyKernel):
227 def _is_valid_warmup_dispatch(
228 self,
229 *,
230 tokens: int,
231 lanes: int,
232 max_work: int,
233 ) -> bool:
234 block_size = _next_power_of_2(tokens)
235 return block_size * lanes <= max_work
236
237 def get_warmup_keys(
238 self, max_tokens: int, cfg: Any
239 ) -> list[ToyKernel.CompileKey]:
240 return self._trace_dispatch(self.dispatch)(
241 tokens=WarmupIntRange(1, max_tokens + 1),
242 lanes=(1, 2),
243 cfg=cfg,
244 max_work=4,
245 _when=self._is_valid_warmup_dispatch,
246 )
247
248 assert PredicateKernel().get_warmup_keys(5, _config()) == [
249 ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
250 ToyKernel.CompileKey(1, 2, 1, ("base", "default", -1, 1, 1), True),
251 ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
252 ToyKernel.CompileKey(2, 4, 1, ("base", "default", -2, 2, 4), True),
253 ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
254 ]
255
256
257 def test_zip_inputs_validates_input_rows() -> None:
258 with pytest.raises(ValueError, match="requires at least one"):
259 zip_inputs()
260 with pytest.raises(ValueError, match="rows must be mappings"):
261 zip_inputs(cast(Any, ("tokens", 1)))
262 with pytest.raises(ValueError, match="at least one dispatch input name"):
263 zip_inputs({})
264 with pytest.raises(ValueError, match="dispatch input names must be strings"):
265 zip_inputs(cast(Any, {1: 2}))
266 with pytest.raises(ValueError, match="same dispatch input names"):
267 zip_inputs({"tokens": 1}, {"mode": "small"})
268
269
270 def test_trace_dispatch_rejects_bad_positional_groups_and_duplicates() -> None:
271 kernel = ToyKernel()
272
273 with pytest.raises(TypeError, match="zip_inputs"):
274 kernel._trace_dispatch(kernel.dispatch)(
275 cast(Any, {"tokens": 1}),
276 cfg=_config(),
277 )
278
279 with pytest.raises(ValueError, match="specified more than once"):
280 kernel._trace_dispatch(kernel.dispatch)(
281 zip_inputs(dict(tokens=1, mode="small")),
282 tokens=2,
283 cfg=_config(),
284 )
285
286
287 def test_helper_calls_support_keywords_and_reject_star_kwargs() -> None:
288 class HelperKernel(VllmJitKernel["HelperKernel.CompileKey"]):
289 @dataclass(frozen=True)
290 class CompileKey:
291 value: int
292
293 def dispatch( # type: ignore[override]
294 self,
295 *,
296 tokens: int,
297 block_size: int,
298 ) -> CompileKey:
299 return self.CompileKey(value=_round_up(tokens, multiple=block_size))
300
301 def get_warmup_keys(self) -> list[CompileKey]:
302 return []
303
304 def compile(self, compile_key: CompileKey) -> None:
305 pass
306
307 class StarKwargsKernel(VllmJitKernel["StarKwargsKernel.CompileKey"]):
308 @dataclass(frozen=True)
309 class CompileKey:
310 value: int
311
312 def dispatch( # type: ignore[override]
313 self,
314 *,
315 tokens: int,
316 block_size: int,
317 ) -> CompileKey:
318 return self.CompileKey(value=_round_up(tokens, **{"multiple": block_size}))
319
320 def get_warmup_keys(self) -> list[CompileKey]:
321 return []
322
323 def compile(self, compile_key: CompileKey) -> None:
324 pass
325
326 assert HelperKernel().compile_key(
327 {
328 "tokens": 5,
329 "block_size": 4,
330 }
331 ) == HelperKernel.CompileKey(value=8)
332 with pytest.raises(ValueError, match=r"cannot use \*\*kwargs"):
333 StarKwargsKernel().compile_key({"tokens": 5, "block_size": 4})
334
335
336 def test_dispatch_helper_calls_resolve_python_builtins() -> None:
337 class BuiltinKernel(VllmJitKernel["BuiltinKernel.CompileKey"]):
338 @dataclass(frozen=True)
339 class CompileKey:
340 value: int
341
342 def dispatch( # type: ignore[override]
343 self,
344 *,
345 tokens: int,
346 limit: int,
347 ) -> CompileKey:
348 return self.CompileKey(value=max(1, min(tokens, limit)))
349
350 def get_warmup_keys(self) -> list[CompileKey]:
351 return []
352
353 def compile(self, compile_key: CompileKey) -> None:
354 pass
355
356 assert BuiltinKernel().compile_key({"tokens": 8, "limit": 4}) == (
357 BuiltinKernel.CompileKey(value=4)
358 )
359
360
361 def test_dispatch_body_must_be_local_assignments_then_compile_key_return() -> None:
362 class BranchKernel(VllmJitKernel["BranchKernel.CompileKey"]):
363 @dataclass(frozen=True)
364 class CompileKey:
365 value: int
366
367 def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
368 if value > 0:
369 value = 1
370 return self.CompileKey(value=value)
371
372 def get_warmup_keys(self) -> list[CompileKey]:
373 return []
374
375 def compile(self, compile_key: CompileKey) -> None:
376 pass
377
378 class KwargsReturnKernel(VllmJitKernel["KwargsReturnKernel.CompileKey"]):
379 @dataclass(frozen=True)
380 class CompileKey:
381 value: int
382
383 def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
384 return self.CompileKey(**{"value": value})
385
386 def get_warmup_keys(self) -> list[CompileKey]:
387 return []
388
389 def compile(self, compile_key: CompileKey) -> None:
390 pass
391
392 with pytest.raises(ValueError, match="local assignments"):
393 BranchKernel()
394 with pytest.raises(
395 ValueError,
396 match=r"may unpack only its own \*\*kwargs parameter once",
397 ):
398 KwargsReturnKernel()
399
400
401 def test_dispatch_can_forward_compile_key_fields() -> None:
402 class ForwardingKernel(VllmJitKernel["ForwardingKernel.CompileKey"]):
403 @dataclass(frozen=True)
404 class CompileKey:
405 mode: int
406 block_size: int
407
408 def dispatch( # type: ignore[override]
409 self,
410 *,
411 tokens: int,
412 **compile_key_fields: int,
413 ) -> CompileKey:
414 return self.CompileKey(
415 **compile_key_fields,
416 block_size=_round_up(tokens, multiple=4),
417 )
418
419 def get_warmup_keys(self) -> list[CompileKey]:
420 return self._trace_dispatch(self.dispatch)(
421 tokens=(1, 5),
422 mode=(2, 3),
423 )
424
425 def compile(self, compile_key: CompileKey) -> None:
426 pass
427
428 kernel = ForwardingKernel()
429 expected = kernel.CompileKey(
430 mode=2,
431 block_size=8,
432 )
433 assert kernel.dispatch(tokens=5, mode=2) == expected
434 assert kernel.compile_key({"tokens": 5, "mode": 2}) == expected
435 assert kernel.get_warmup_keys() == [
436 kernel.CompileKey(mode=2, block_size=4),
437 kernel.CompileKey(mode=3, block_size=4),
438 kernel.CompileKey(mode=2, block_size=8),
439 kernel.CompileKey(mode=3, block_size=8),
440 ]
441 with pytest.raises(TypeError, match="field 'block_size' is specified twice"):
442 kernel.compile_key({"tokens": 5, "mode": 2, "block_size": 4})
443 with pytest.raises(TypeError, match="unexpected keyword argument 'extra'"):
444 kernel.compile_key({"tokens": 5, "mode": 2, "extra": 1})
445
446
447 def test_forwarded_compile_key_fields_exclude_predicate_only_inputs() -> None:
448 class ForwardingKernel(VllmJitKernel["ForwardingKernel.CompileKey"]):
449 @dataclass(frozen=True)
450 class CompileKey:
451 mode: int
452 block_size: int
453
454 def dispatch( # type: ignore[override]
455 self,
456 *,
457 tokens: int,
458 **compile_key_fields: int,
459 ) -> CompileKey:
460 return self.CompileKey(
461 **compile_key_fields,
462 block_size=_round_up(tokens, multiple=4),
463 )
464
465 def get_warmup_keys(self) -> list[CompileKey]:
466 return self._trace_dispatch(self.dispatch)(
467 tokens=(1, 5),
468 mode=2,
469 max_tokens=4,
470 _when=lambda *, tokens, max_tokens: tokens <= max_tokens,
471 )
472
473 def compile(self, compile_key: CompileKey) -> None:
474 pass
475
476 assert ForwardingKernel().get_warmup_keys() == [
477 ForwardingKernel.CompileKey(mode=2, block_size=4)
478 ]
479
480
481 def test_dispatch_supports_tuple_and_mapping_subscriptions() -> None:
482 class SubscriptKernel(VllmJitKernel["SubscriptKernel.CompileKey"]):
483 @dataclass(frozen=True)
484 class CompileKey:
485 first: int
486 named: int
487
488 def dispatch( # type: ignore[override]
489 self,
490 *,
491 values: tuple[int, ...],
492 config: dict[str, int],
493 ) -> CompileKey:
494 return self.CompileKey(first=values[0], named=config["named"])
495
496 def get_warmup_keys(self) -> list[CompileKey]:
497 return []
498
499 def compile(self, compile_key: CompileKey) -> None:
500 pass
501
502 assert SubscriptKernel().compile_key(
503 {"values": (3, 5), "config": {"named": 7}}
504 ) == SubscriptKernel.CompileKey(first=3, named=7)
505
506
507 def test_dispatch_reports_unsupported_expression_with_context() -> None:
508 class UnsupportedKernel(VllmJitKernel["UnsupportedKernel.CompileKey"]):
509 @dataclass(frozen=True)
510 class CompileKey:
511 value: object
512
513 def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
514 return self.CompileKey(value={value})
515
516 def get_warmup_keys(self) -> list[CompileKey]:
517 return []
518
519 def compile(self, compile_key: CompileKey) -> None:
520 pass
521
522 with pytest.raises(ValueError) as exc_info:
523 UnsupportedKernel().compile_key({"value": 1})
524
525 message = str(exc_info.value)
526 assert "Unsupported dispatch expression" in message
527 assert "{value}" in message
528 assert "Supported dispatch expressions" in message
529
530
531 def test_warmup_compiles_all_returned_keys_in_order() -> None:
532 kernel = RecordingToyKernel()
533 cfg = _config()
534
535 kernel.warmup(3, cfg)
536
537 assert kernel.compiled == [
538 ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
539 ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
540 ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
541 ]
542
543
544 def test_runtime_cache_miss_compiles_and_caches_executor() -> None:
545 class CachedKernel(VllmJitKernel["CachedKernel.CompileKey"]):
546 @dataclass(frozen=True)
547 class CompileKey:
548 value: int
549
550 def __init__(self) -> None:
551 self.compiled: list[CachedKernel.CompileKey] = []
552 super().__init__()
553
554 def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
555 return self.CompileKey(value=value)
556
557 def get_warmup_keys(self) -> list[CompileKey]:
558 return self._trace_dispatch(self.dispatch)(value=1)
559
560 def compile(self, compile_key: CompileKey) -> None:
561 self.compiled.append(compile_key)
562 self._compiled_cache[compile_key] = object()
563
564 def __call__(self, value: int) -> None:
565 compile_key = self.dispatch(value=value)
566 self._get_or_compile(compile_key)
567
568 kernel = CachedKernel()
569 kernel.warmup()
570 kernel(1)
571 kernel(2)
572
573 assert kernel.compiled == [
574 CachedKernel.CompileKey(value=1),
575 CachedKernel.CompileKey(value=2),
576 ]
577
578
579 def test_registry_records_only_inside_model_setup_context() -> None:
580 registry = JitWarmupRegistry(_config())
581 kernel = RecordingToyKernel()
582 config = _config()
583
584 kernel.register_warmup(3, _config())
585 with registry.activate():
586 kernel.register_warmup(3, config)
587 kernel.register_warmup(3, config)
588
589 assert len(registry) == 1
590 assert kernel.compiled == []
591
592
593 def test_registry_capture() -> None:
594 class Owner:
595 @JitWarmupRegistry.capture
596 def __init__(
597 self,
598 vllm_config: Any,
599 kernel: RecordingToyKernel,
600 ) -> None:
601 kernel.register_warmup(3, vllm_config)
602
603 kernel = RecordingToyKernel()
604 config = _config()
605 owner = Owner(config, kernel)
606
607 assert len(owner.jit_warmup_registry) == 1 # type: ignore[attr-defined]
608 kernel.register_warmup(5, config)
609 assert len(owner.jit_warmup_registry) == 1 # type: ignore[attr-defined]
610
611
612 def test_registry_expands_requests_and_deduplicates_owner_keys() -> None:
613 registry = JitWarmupRegistry(_config())
614 kernel = RecordingToyKernel()
615
616 with registry.activate():
617 kernel.register_warmup(3, _config())
618 kernel.register_warmup(5, _config())
619
620 registry.warmup()
621
622 assert kernel.compiled == [
623 ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
624 ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
625 ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
626 ToyKernel.CompileKey(8, 8, 1, ("base", "default", -8, 2, 64), True),
627 ]
628
629
630 def test_registry_passes_vllm_config_to_default_requests() -> None:
631 class ConfigKernel(VllmJitKernel["ConfigKernel.CompileKey"]):
632 @dataclass(frozen=True)
633 class CompileKey:
634 value: int
635
636 def __init__(self) -> None:
637 self.compiled: list[ConfigKernel.CompileKey] = []
638 super().__init__()
639
640 def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
641 return self.CompileKey(value=value)
642
643 def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
644 return [self.dispatch(value=vllm_config.bias)]
645
646 def compile(self, compile_key: CompileKey) -> None:
647 self.compiled.append(compile_key)
648
649 registry = JitWarmupRegistry(_config(bias=7))
650 kernel = ConfigKernel()
651
652 with registry.activate():
653 kernel.register_warmup()
654 kernel.register_warmup()
655 assert len(registry) == 1
656 registry.warmup()
657
658 assert kernel.compiled == [ConfigKernel.CompileKey(value=7)]
659
660
661 def test_registry_default_request_supports_config_free_kernel() -> None:
662 class ConfigFreeKernel(VllmJitKernel["ConfigFreeKernel.CompileKey"]):
663 @dataclass(frozen=True)
664 class CompileKey:
665 value: int
666
667 def __init__(self) -> None:
668 self.compiled: list[ConfigFreeKernel.CompileKey] = []
669 super().__init__()
670
671 def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
672 return self.CompileKey(value=value)
673
674 def get_warmup_keys(self) -> list[CompileKey]:
675 return [self.dispatch(value=7)]
676
677 def compile(self, compile_key: CompileKey) -> None:
678 self.compiled.append(compile_key)
679
680 registry = JitWarmupRegistry(_config())
681 kernel = ConfigFreeKernel()
682
683 with registry.activate():
684 kernel.register_warmup()
685 registry.warmup()
686
687 assert kernel.compiled == [ConfigFreeKernel.CompileKey(value=7)]
688
689
690 def test_get_ast_full_name_handles_names_attributes_and_other_nodes() -> None:
691 dotted_expr = ast.parse("foo.bar.baz").body[0]
692 call_expr = ast.parse("foo()").body[0]
693 assert isinstance(dotted_expr, ast.Expr)
694 assert isinstance(call_expr, ast.Expr)
695
696 assert get_ast_full_name(dotted_expr.value) == "foo.bar.baz"
697 assert get_ast_full_name(call_expr.value) is None