@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 49905ad94d

[3/N][Feat][Perf] Add new warmup infrastructure for JITs. Add provider registry and orchestration for JIT warmup (#50174)

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Co-authored-by: Codex <codex@openai.com>

Roberto L. Castro committed Aug 17, 2026 at 19:11 UTC 49905ad94dfc58ba94cca8b173dba31b820d6bc6
12 files changed +1067 -168
docs/contributing/README.md
+1
@@ -71,6 +71,7 @@ uv pip install -e . --no-build-isolation
71 For more details about installing from source and installing for other hardware, check out the [installation instructions](../getting_started/installation/README.md) for your hardware and head to the "Build wheel from source" section.
72
73 For an optimized workflow when iterating on C++/CUDA kernels, see the [Incremental Compilation Workflow](./incremental_build.md) for recommendations.
74 +For JIT kernel warmup conventions, see [JIT Kernel Warmup](./jit_kernel_warmup.md).
75
76 !!! tip
77 vLLM is compatible with Python versions 3.10 to 3.13. However, vLLM's default [Dockerfile](../../docker/Dockerfile) ships with Python 3.12 and tests in CI (except `mypy`) are run with Python 3.12.
docs/contributing/jit_kernel_warmup.md new
+308
@@ -0,0 +1,308 @@
1 +# JIT Kernel Warmup
2 +
3 +vLLM uses JIT-generated kernels from Triton, CuTeDSL, TileLang, and other backends. This contract makes their required specializations available during startup, before the first request, by warming the kernel's **compile-key space** without dummy runtime launches or real tensor allocation.
4 +
5 +Use it when adding a warmable JIT kernel or migrating an existing warmup path.
6 +
7 +## In This Guide
8 +
9 +- [1. Quickstart](#1-quickstart): for contributors adding or migrating a warmable kernel.
10 +- [2. Search-Space Reference](#2-search-space-reference): additional details regarding warmup input expansion and traced dispatch rules.
11 +
12 +## 1. Quickstart
13 +
14 +Each warmable kernel defines its compile-key mapping and compile-only entry point beside its normal runtime implementation. The startup registry then warms only the wrappers selected by the current engine configuration.
15 +
16 +### Define the Kernel Wrapper
17 +
18 +Here, a **kernel wrapper** (or just **wrapper**) is an instance of a concrete `VllmJitKernel` subclass.
19 +
20 +Expose one wrapper near the kernel's normal runtime entry point. Prefer this shape:
21 +
22 +```python
23 +class MyKernel(VllmJitKernel["MyKernel.CompileKey"]):
24 +
25 + @dataclass(frozen=True)
26 + class CompileKey:
27 + ...
28 +
29 + @staticmethod
30 + def kernel(...):
31 + ...
32 +
33 + def dispatch(self, ...) -> CompileKey:
34 + return self.CompileKey(...)
35 +
36 + def get_warmup_keys(self, ...) -> list[CompileKey]:
37 + return self._trace_dispatch(self.dispatch)(...)
38 +
39 + def compile(self, compile_key: CompileKey) -> None:
40 + ...
41 +
42 + def __call__(self, ...):
43 + return self.kernel(...)
44 +
45 +
46 +MY_KERNEL = MyKernel()
47 +```
48 +
49 +`CompileKey`, `dispatch(...)`, and `get_warmup_keys(...)` are backend-agnostic. Backend-specific behavior belongs in `kernel(...)`, `compile(...)`, and `__call__(...)`.
50 +
51 +The module-level singleton should be used by warmup and by the runtime call path. This keeps dispatch behavior shared instead of duplicated.
52 +
53 +`VllmJitKernel.warmup(...)` compiles every key returned by `get_warmup_keys(...)`; wrappers should not reimplement it.
54 +
55 +### Choose Compile-Key Fields
56 +
57 +`CompileKey` must be frozen and hashable. Include only fields on which the backend specializes, such as tile sizes, head dimensions, dtypes, pointer alignment classes, or backend selectors; exclude runtime-only values. When unsure, inspect the backend cache key, specialization arguments, or verbose JIT-monitor output.
58 +
59 +### Generate Warmup Keys
60 +
61 +Use `_trace_dispatch(self.dispatch)` to describe representative inputs. The tracer maps them through the same specialization logic and deduplicates equal keys:
62 +
63 +```python
64 +def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
65 + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
66 + return self._trace_dispatch(self.dispatch)(
67 + num_tokens=WarmupIntRange(1, max_tokens + 1),
68 + )
69 +```
70 +
71 +Use independent ranges or alternatives for cartesian products, `zip_inputs(...)` for coupled rows, and `_when` for validity constraints. The complete syntax is documented in [Search-Space Reference](#2-search-space-reference).
72 +
73 +### Compile Without Launching
74 +
75 +`compile(compile_key)` means "make this specialization available". Depending on the backend, that may compile from source, call a compile-only API, load an already-built artifact, or compile on cache miss.
76 +
77 +`compile(...)` should not launch a real inference workload or allocate real tensors. Each DSL should expose fake tensor/spec descriptors suitable for compilation only.
78 +
79 +### Register the Selected Wrapper
80 +
81 +Register the wrapper where the runtime implementation is selected:
82 +
83 +```python
84 +MY_KERNEL.register_warmup()
85 +```
86 +
87 +Registration records metadata only. It does not compile or launch the kernel. Repeated registrations from equivalent layers are allowed and deduplicated later.
88 +
89 +### Review Checklist
90 +
91 +- Warm actual compile keys rather than representative non-key inputs.
92 +- Keep specialization mapping in `dispatch(...)` instead of duplicating it in warmup code.
93 +- Use fake tensors or backend compile-only descriptors; never perform a dummy runtime launch.
94 +- Keep registration metadata-only so model construction remains cheap and side-effect free.
95 +- Compile registered kernels only through `kernel_warmup()`.
96 +- Keep runtime execution and startup compilation separate and easy to review.
97 +- Use one module-level wrapper instance for registration and runtime calls.
98 +
99 +## 2. Search-Space Reference
100 +
101 +### How Tracing Works
102 +
103 +`_trace_dispatch(...)` expands the inputs declared by `get_warmup_keys(...)`. Each concrete combination becomes a `dispatch_values` mapping from input names to selected values. `_when` may reject that mapping; otherwise the tracer evaluates `dispatch(...)` to construct one `CompileKey`. Equal keys are deduplicated after all combinations are evaluated.
104 +
105 +One call to `dispatch(...)` returns one key, but many input points may map to the same key. Prefer this traced mapping over manually reconstructing keys in warmup code; `dispatch(...)` should express the same specialization logic used by the runtime path.
106 +
107 +### Define Input Spaces
108 +
109 +Use ranges and alternatives for independent axes, `zip_inputs(...)` for coupled rows, and `_when` for validity constraints.
110 +
111 +#### Integer Ranges
112 +
113 +Use `WarmupIntRange` for integer ranges:
114 +
115 +```python
116 +return self._trace_dispatch(self.dispatch)(
117 + num_prefills=WarmupIntRange(1, max_prefills + 1),
118 +)
119 +```
120 +
121 +`WarmupIntRange(start, stop, step)` follows Python `range(...)` semantics: `start` is inclusive, `stop` is exclusive, and `step` defaults to 1.
122 +
123 +For non-linear integer sequences, use `advance` to provide the action that computes each next value:
124 +
125 +```python
126 +return self._trace_dispatch(self.dispatch)(
127 + num_tokens=WarmupIntRange(
128 + 1,
129 + max_tokens + 1,
130 + advance=lambda value: next_power_of_2(value) + 1,
131 + ),
132 +)
133 +```
134 +
135 +This is useful for traversing specialization boundaries without enumerating every integer. `advance` cannot be combined with a non-default `step`, and it must return a value greater than its input so expansion always makes forward progress.
136 +
137 +#### Independent Alternatives
138 +
139 +Use tuples or lists for independent alternatives. Multiple expanded inputs form a cartesian product:
140 +
141 +```python
142 +return self._trace_dispatch(self.dispatch)(
143 + query_slice_start=WarmupIntRange(0, 2),
144 + query_slice_stop=(1, 2 * max_tokens - 1, 2 * max_tokens),
145 + COMPRESS_RATIO=list(compress_ratios),
146 +)
147 +```
148 +
149 +#### Coupled Inputs
150 +
151 +Use `zip_inputs(...)` when values must vary together row-by-row:
152 +
153 +```python
154 +WARMUP_INPUTS = zip_inputs(
155 + dict(compress_ratio=1, topk=0, topk_width=512),
156 + dict(compress_ratio=4, topk=512, topk_width=512),
157 +)
158 +
159 +
160 +return self._trace_dispatch(self.dispatch)(
161 + WARMUP_INPUTS,
162 + WINDOW_SIZE=window_size,
163 +)
164 +```
165 +
166 +Multiple `zip_inputs(...)` groups may be passed as positional arguments. The tracer forms the cartesian product across groups while preserving row-wise coupling inside each group.
167 +
168 +Every row in a `zip_inputs(...)` group must use the same string keys. A `zip_inputs(...)` group cannot specify a field that is also specified as a keyword input to `_trace_dispatch(...)`.
169 +
170 +#### Conditional Filtering
171 +
172 +Use `_when=...` to filter generated input points before they are passed to `dispatch(...)`. This is useful when independent ranges contain invalid combinations, but the validity rule belongs with the kernel warmup definition.
173 +
174 +```python
175 +def _is_valid_warmup_input(
176 + self,
177 + *,
178 + query_len: int,
179 + num_reqs: int,
180 + max_num_batched_tokens: int,
181 +) -> bool:
182 + return query_len + num_reqs - 1 <= max_num_batched_tokens
183 +
184 +
185 +return self._trace_dispatch(self.dispatch)(
186 + query_len=WarmupIntRange(1, max_tokens + 1),
187 + num_reqs=WarmupIntRange(1, max_reqs + 1),
188 + max_num_batched_tokens=max_tokens,
189 + _when=self._is_valid_warmup_input,
190 +)
191 +```
192 +
193 +`_when` accepts a function, bound method, or lambda and supports the same AST subset as `dispatch(...)`, including local assignments in function predicates.
194 +
195 +The predicate is evaluated on the expanded warmup inputs. If it returns `False`, that input point is skipped and no `CompileKey` is produced for it.
196 +
197 +### Write Dispatch Rules
198 +
199 +#### Local Assignments
200 +
201 +The traced body may contain local assignments, optionally annotated, followed by one `return self.CompileKey(...)` call. Local assignments let a kernel name intermediate specialization choices once and reuse them across fields:
202 +
203 +```python
204 +def dispatch(
205 + self,
206 + *,
207 + num_tokens: int,
208 + vectorized: bool,
209 +) -> CompileKey:
210 + block_size = next_power_of_2(num_tokens)
211 + return self.CompileKey(
212 + BLOCK_SIZE=block_size,
213 + VECTOR_WIDTH=4 if vectorized and block_size >= 4 else 1,
214 + )
215 +```
216 +
217 +#### Supported Expressions
218 +
219 +The evaluator supports these expressions inside local assignments and `CompileKey(...)` fields:
220 +
221 +| Feature | What It Allows |
222 +| --- | --- |
223 +| Names | Read dispatch inputs, local assignments, defaults, and module globals. |
224 +| Constants | Use literals such as integers, strings, booleans, and `None`. |
225 +| Attributes | Read structured values such as `cfg.block_size` or `mla_dims.v_head_dim`. |
226 +| Subscriptions | Read sequence positions or mapping values such as `config[0]` and `config["block_size"]`. |
227 +| Tuple/list literals | Build shapes, strides, and other small structured fields. |
228 +| Conditional expressions | Select a field with `x if condition else y`. |
229 +| Boolean expressions | Combine predicates with `and`, `or`, and `not`. |
230 +| Comparisons | Use `==`, `!=`, `<`, `<=`, `>`, `>=`, `in`, `not in`, `is`, and `is not`. |
231 +| Arithmetic | Use `+`, `-`, `*`, `//`, `%`, and `**`. |
232 +| Unary minus | Build negative sentinel values or signed descriptors. |
233 +| Helper calls | Call helpers with positional and explicit keyword arguments. |
234 +
235 +Python builtins such as `min(...)`, `max(...)`, and `len(...)` are resolved unless the name is overridden locally or globally.
236 +
237 +#### Helper Calls
238 +
239 +Helpers are useful for small specialization rules:
240 +
241 +```python
242 +def dispatch(self, *, num_tokens: int, block_size: int) -> CompileKey:
243 + return self.CompileKey(
244 + PADDED_TOKENS=round_up(num_tokens, multiple=block_size),
245 + )
246 +```
247 +
248 +`_trace_dispatch(...)` does not inspect helper bodies. It evaluates the call arguments and invokes the helper as ordinary Python, so control flow inside that helper is outside the AST interpreter's scope. Keep helpers deterministic and side-effect free.
249 +
250 +#### Direct Keyword Forwarding
251 +
252 +For many direct pass-through fields, the dispatch `**kwargs` parameter may be unpacked into `CompileKey(...)`:
253 +
254 +```python
255 +def dispatch(
256 + self,
257 + *,
258 + num_tokens: int,
259 + **compile_key_fields: int,
260 +) -> CompileKey:
261 + return self.CompileKey(
262 + **compile_key_fields,
263 + block_size=next_power_of_2(num_tokens),
264 + )
265 +```
266 +
267 +Unmatched dispatch arguments become compile-key fields and warmup inputs. Keep transformed inputs named and explicit. The unpacking must use the dispatch method's own `**kwargs` parameter directly and exactly once; arbitrary mappings, repeated unpacking, and helper-call `**kwargs` are rejected. The fully explicit form remains supported and is often clearer for non-trivial mappings.
268 +
269 +#### Unsupported Syntax
270 +
271 +Conditional expressions (`x if condition else y`) are supported, but statement-level `if` blocks are not supported directly inside traced `dispatch(...)` or `_when` bodies. The tracer expects a straight-line sequence of local assignments followed by one return expression. Small, pure helpers called by traced expressions execute as normal Python with concrete values and may use ordinary control flow, including `if` blocks. Do not put loops, mutation, side effects, or backend imports directly inside traced functions. Put environment and model gating in `get_warmup_keys(...)` or the outer warmup entry point.
272 +
273 +### Compile-Key Deduplication
274 +
275 +`_trace_dispatch(...)` deduplicates the resulting keys while preserving order. This is important when many runtime-like inputs map to the same static bucket.
276 +
277 +For example, this warmup range expands every token count, but the compile key only depends on the power-of-two bucket:
278 +
279 +```python
280 +def dispatch(
281 + self,
282 + *,
283 + num_tokens: int,
284 +) -> CompileKey:
285 + return self.CompileKey(
286 + BLOCK_SIZE=next_power_of_2(num_tokens),
287 + )
288 +
289 +
290 +def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
291 + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
292 + return self._trace_dispatch(self.dispatch)(
293 + num_tokens=WarmupIntRange(1, max_tokens + 1),
294 + )
295 +```
296 +
297 +For `max_tokens == 8`, the expanded inputs are `1, 2, 3, 4, 5, 6, 7, 8`, but the returned keys are:
298 +
299 +```python
300 +[
301 + CompileKey(BLOCK_SIZE=1),
302 + CompileKey(BLOCK_SIZE=2),
303 + CompileKey(BLOCK_SIZE=4),
304 + CompileKey(BLOCK_SIZE=8),
305 +]
306 +```
307 +
308 +Deduplication happens after `dispatch(...)` is evaluated, so the warmup system removes duplicate compile keys, not duplicate input values. `CompileKey` must be hashable for this to work; using `@dataclass(frozen=True)` is the standard pattern.
tests/model_executor/test_jit_warmup.py
+230 -1
@@ -9,11 +9,15 @@ from typing import Any, cast
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:
@@ -94,6 +98,35 @@ class RecordingToyKernel(ToyKernel):
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
@@ -287,6 +320,31 @@ def test_helper_calls_support_keywords_and_reject_star_kwargs() -> None:
320 StarKwargsKernel().compile_key({"tokens": 5, "block_size": 4})
321
322
323 +def test_dispatch_helper_calls_resolve_python_builtins() -> None:
324 + class BuiltinKernel(VllmJitKernel["BuiltinKernel.CompileKey"]):
325 + @dataclass(frozen=True)
326 + class CompileKey:
327 + value: int
328 +
329 + def dispatch( # type: ignore[override]
330 + self,
331 + *,
332 + tokens: int,
333 + limit: int,
334 + ) -> CompileKey:
335 + return self.CompileKey(value=max(1, min(tokens, limit)))
336 +
337 + def get_warmup_keys(self) -> list[CompileKey]:
338 + return []
339 +
340 + def compile(self, compile_key: CompileKey) -> None:
341 + pass
342 +
343 + assert BuiltinKernel().compile_key({"tokens": 8, "limit": 4}) == (
344 + BuiltinKernel.CompileKey(value=4)
345 + )
346 +
347 +
348 def test_dispatch_body_must_be_local_assignments_then_compile_key_return() -> None:
349 class BranchKernel(VllmJitKernel["BranchKernel.CompileKey"]):
350 @dataclass(frozen=True)
@@ -320,10 +378,85 @@ def test_dispatch_body_must_be_local_assignments_then_compile_key_return() -> No
378
379 with pytest.raises(ValueError, match="local assignments"):
380 BranchKernel()
323 - with pytest.raises(ValueError, match=r"cannot use \*\*kwargs in CompileKey"):
381 + with pytest.raises(
382 + ValueError,
383 + match=r"may unpack only its own \*\*kwargs parameter once",
384 + ):
385 KwargsReturnKernel()
386
387
388 +def test_dispatch_can_forward_compile_key_fields() -> None:
389 + class ForwardingKernel(VllmJitKernel["ForwardingKernel.CompileKey"]):
390 + @dataclass(frozen=True)
391 + class CompileKey:
392 + mode: int
393 + block_size: int
394 +
395 + def dispatch( # type: ignore[override]
396 + self,
397 + *,
398 + tokens: int,
399 + **compile_key_fields: int,
400 + ) -> CompileKey:
401 + return self.CompileKey(
402 + **compile_key_fields,
403 + block_size=_round_up(tokens, multiple=4),
404 + )
405 +
406 + def get_warmup_keys(self) -> list[CompileKey]:
407 + return self._trace_dispatch(self.dispatch)(
408 + tokens=(1, 5),
409 + mode=(2, 3),
410 + )
411 +
412 + def compile(self, compile_key: CompileKey) -> None:
413 + pass
414 +
415 + kernel = ForwardingKernel()
416 + expected = kernel.CompileKey(
417 + mode=2,
418 + block_size=8,
419 + )
420 + assert kernel.dispatch(tokens=5, mode=2) == expected
421 + assert kernel.compile_key({"tokens": 5, "mode": 2}) == expected
422 + assert kernel.get_warmup_keys() == [
423 + kernel.CompileKey(mode=2, block_size=4),
424 + kernel.CompileKey(mode=3, block_size=4),
425 + kernel.CompileKey(mode=2, block_size=8),
426 + kernel.CompileKey(mode=3, block_size=8),
427 + ]
428 + with pytest.raises(TypeError, match="field 'block_size' is specified twice"):
429 + kernel.compile_key({"tokens": 5, "mode": 2, "block_size": 4})
430 + with pytest.raises(TypeError, match="unexpected keyword argument 'extra'"):
431 + kernel.compile_key({"tokens": 5, "mode": 2, "extra": 1})
432 +
433 +
434 +def test_dispatch_supports_tuple_and_mapping_subscriptions() -> None:
435 + class SubscriptKernel(VllmJitKernel["SubscriptKernel.CompileKey"]):
436 + @dataclass(frozen=True)
437 + class CompileKey:
438 + first: int
439 + named: int
440 +
441 + def dispatch( # type: ignore[override]
442 + self,
443 + *,
444 + values: tuple[int, ...],
445 + config: dict[str, int],
446 + ) -> CompileKey:
447 + return self.CompileKey(first=values[0], named=config["named"])
448 +
449 + def get_warmup_keys(self) -> list[CompileKey]:
450 + return []
451 +
452 + def compile(self, compile_key: CompileKey) -> None:
453 + pass
454 +
455 + assert SubscriptKernel().compile_key(
456 + {"values": (3, 5), "config": {"named": 7}}
457 + ) == SubscriptKernel.CompileKey(first=3, named=7)
458 +
459 +
460 def test_dispatch_reports_unsupported_expression_with_context() -> None:
461 class UnsupportedKernel(VllmJitKernel["UnsupportedKernel.CompileKey"]):
462 @dataclass(frozen=True)
@@ -361,6 +494,102 @@ def test_warmup_compiles_all_returned_keys_in_order() -> None:
494 ]
495
496
497 +def test_runtime_cache_miss_compiles_and_caches_executor() -> None:
498 + class CachedKernel(VllmJitKernel["CachedKernel.CompileKey"]):
499 + @dataclass(frozen=True)
500 + class CompileKey:
501 + value: int
502 +
503 + def __init__(self) -> None:
504 + self.compiled: list[CachedKernel.CompileKey] = []
505 + super().__init__()
506 +
507 + def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
508 + return self.CompileKey(value=value)
509 +
510 + def get_warmup_keys(self) -> list[CompileKey]:
511 + return self._trace_dispatch(self.dispatch)(value=1)
512 +
513 + def compile(self, compile_key: CompileKey) -> None:
514 + self.compiled.append(compile_key)
515 + self._compiled_cache[compile_key] = object()
516 +
517 + def __call__(self, value: int) -> None:
518 + compile_key = self.dispatch(value=value)
519 + self._get_or_compile(compile_key)
520 +
521 + kernel = CachedKernel()
522 + kernel.warmup()
523 + kernel(1)
524 + kernel(2)
525 +
526 + assert kernel.compiled == [
527 + CachedKernel.CompileKey(value=1),
528 + CachedKernel.CompileKey(value=2),
529 + ]
530 +
531 +
532 +def test_registry_records_only_inside_model_setup_context() -> None:
533 + registry = JitWarmupRegistry(_config())
534 + kernel = RecordingToyKernel()
535 +
536 + kernel.register_warmup(3, _config())
537 + with registry.activate():
538 + kernel.register_warmup(3, _config())
539 +
540 + assert len(registry) == 1
541 + assert kernel.compiled == []
542 +
543 +
544 +def test_registry_expands_requests_and_deduplicates_owner_keys() -> None:
545 + registry = JitWarmupRegistry(_config())
546 + kernel = RecordingToyKernel()
547 +
548 + with registry.activate():
549 + kernel.register_warmup(3, _config())
550 + kernel.register_warmup(5, _config())
551 +
552 + registry.warmup()
553 +
554 + assert kernel.compiled == [
555 + ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
556 + ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
557 + ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
558 + ToyKernel.CompileKey(8, 8, 1, ("base", "default", -8, 2, 64), True),
559 + ]
560 +
561 +
562 +def test_registry_passes_vllm_config_to_default_requests() -> None:
563 + class ConfigKernel(VllmJitKernel["ConfigKernel.CompileKey"]):
564 + @dataclass(frozen=True)
565 + class CompileKey:
566 + value: int
567 +
568 + def __init__(self) -> None:
569 + self.compiled: list[ConfigKernel.CompileKey] = []
570 + super().__init__()
571 +
572 + def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
573 + return self.CompileKey(value=value)
574 +
575 + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
576 + return [self.dispatch(value=vllm_config.bias)]
577 +
578 + def compile(self, compile_key: CompileKey) -> None:
579 + self.compiled.append(compile_key)
580 +
581 + registry = JitWarmupRegistry(_config(bias=7))
582 + kernel = ConfigKernel()
583 +
584 + with registry.activate():
585 + kernel.register_warmup()
586 + kernel.register_warmup()
587 + assert len(registry) == 1
588 + registry.warmup()
589 +
590 + assert kernel.compiled == [ConfigKernel.CompileKey(value=7)]
591 +
592 +
593 def test_get_ast_full_name_handles_names_attributes_and_other_nodes() -> None:
594 dotted_expr = ast.parse("foo.bar.baz").body[0]
595 call_expr = ast.parse("foo()").body[0]
tests/v1/worker/test_jit_warmup_migration.py new
+52
@@ -0,0 +1,52 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Validate the registry reference kernel against runtime dispatch."""
4 +
5 +import pytest
6 +
7 +from vllm.platforms import current_platform
8 +
9 +if not current_platform.is_cuda_alike():
10 + pytest.skip("NVIDIA dispatch tests require CUDA", allow_module_level=True)
11 +
12 +from vllm.v1.worker.block_table import ComputeSlotMappingKernel
13 +
14 +
15 +@pytest.mark.parametrize(
16 + ("kv_cache_block_size", "blocks_per_kv_block", "block_size", "block_size_rep"),
17 + [
18 + (256, 1, 256, 16),
19 + (256, 4, 64, 16),
20 + (64, 1, 64, 16),
21 + (8, 1, 8, 2),
22 + (4, 1, 4, 2),
23 + ],
24 +)
25 +def test_compute_slot_mapping_warmup_matches_runtime_specializations(
26 + kv_cache_block_size: int,
27 + blocks_per_kv_block: int,
28 + block_size: int,
29 + block_size_rep: int,
30 +) -> None:
31 + kernel = ComputeSlotMappingKernel()
32 + kwargs = dict(
33 + kv_cache_block_size=kv_cache_block_size,
34 + blocks_per_kv_block=blocks_per_kv_block,
35 + total_cp_world_size=1,
36 + total_cp_rank=0,
37 + cp_kv_cache_interleave_size=1,
38 + block_table_stride=32768,
39 + block_size=block_size,
40 + )
41 + expected = kernel.CompileKey(
42 + kv_cache_block_size=kv_cache_block_size,
43 + blocks_per_kv_block=blocks_per_kv_block,
44 + total_cp_world_size=1,
45 + total_cp_rank=0,
46 + cp_kv_cache_interleave_size=1,
47 + block_table_stride=16,
48 + block_size=block_size_rep,
49 + )
50 +
51 + assert kernel.dispatch(**kwargs) == expected
52 + assert kernel.get_warmup_keys(**kwargs) == [expected]
vllm/model_executor/warmup/jit_warmup.py
+211 -18
@@ -5,16 +5,20 @@
5 from __future__ import annotations
6
7 import ast
8 +import builtins
9 import inspect
10 import itertools
11 import operator
12 import textwrap
13 from abc import ABC, abstractmethod
13 -from collections.abc import Callable, Iterable, Mapping
14 +from collections.abc import Callable, Iterable, Iterator, Mapping
15 +from contextlib import contextmanager
16 +from contextvars import ContextVar
17 from dataclasses import dataclass
18 from typing import Any, Generic, TypeVar
19
20 __all__ = [
21 + "JitWarmupRegistry",
22 "VllmJitKernel",
23 "WarmupIntRange",
24 "get_ast_full_name",
@@ -136,8 +140,15 @@ class _CompileKeyDispatchTrace:
140 field_exprs: tuple[tuple[str, ast.AST], ...]
141 globals: Mapping[str, Any]
142 input_names: frozenset[str]
143 + # Named parameters are excluded from direct **kwargs forwarding.
144 + named_parameters: frozenset[str] | None
145 defaults: Mapping[str, Any]
146
147 + def input_names_for(self, available_names: set[str]) -> frozenset[str]:
148 + if self.named_parameters is None:
149 + return self.input_names
150 + return self.input_names | (available_names - self.named_parameters)
151 +
152 def compile_key(
153 self,
154 compile_key_type: type[CompileKeyT],
@@ -146,12 +157,20 @@ class _CompileKeyDispatchTrace:
157 dispatch_values = _eval_local_exprs(
158 self.local_exprs, {**self.defaults, **kwargs}, self.globals
159 )
149 - return compile_key_type(
150 - **{
151 - field: _eval_dispatch_expr(expr, dispatch_values, self.globals)
152 - for field, expr in self.field_exprs
160 + named_parameters = self.named_parameters
161 + # Materialize direct fields before evaluating named AST expressions.
162 + fields: dict[str, Any] = {}
163 + if named_parameters is not None:
164 + fields = {
165 + name: value
166 + for name, value in kwargs.items()
167 + if name not in named_parameters
168 }
154 - )
169 + for field, expr in self.field_exprs:
170 + if field in fields:
171 + raise TypeError(f"CompileKey field '{field}' is specified twice")
172 + fields[field] = _eval_dispatch_expr(expr, dispatch_values, self.globals)
173 + return compile_key_type(**fields)
174
175
176 @dataclass(frozen=True)
@@ -186,6 +205,10 @@ _CMP_OPS: dict[type[ast.cmpop], Callable[[Any, Any], bool]] = {
205 ast.LtE: operator.le,
206 ast.Gt: operator.gt,
207 ast.GtE: operator.ge,
208 + ast.In: lambda left, right: left in right,
209 + ast.NotIn: lambda left, right: left not in right,
210 + ast.Is: operator.is_,
211 + ast.IsNot: operator.is_not,
212 }
213
214
@@ -200,8 +223,9 @@ def _dispatch_expr_error(node: ast.AST, reason: str) -> ValueError:
223 return ValueError(
224 f"{reason}: {_dispatch_expr_source(node)}. "
225 "Supported dispatch expressions are names, constants, attributes, "
203 - "tuple/list literals, conditional expressions, comparisons, boolean "
204 - "operators, unary not/minus, arithmetic, and calls without **kwargs."
226 + "subscriptions, tuple/list literals, conditional expressions, "
227 + "comparisons, boolean operators, unary not/minus, arithmetic, and "
228 + "calls without **kwargs."
229 )
230
231
@@ -225,6 +249,8 @@ class _DispatchExprEvaluator(ast.NodeVisitor):
249 return self.values[node.id]
250 if node.id in self.globals:
251 return self.globals[node.id]
252 + if hasattr(builtins, node.id):
253 + return getattr(builtins, node.id)
254 raise _dispatch_expr_error(node, f"Unknown dispatch name '{node.id}'")
255
256 def visit_Constant(self, node: ast.Constant) -> Any:
@@ -299,6 +325,9 @@ class _DispatchExprEvaluator(ast.NodeVisitor):
325 def visit_Attribute(self, node: ast.Attribute) -> Any:
326 return getattr(self.visit(node.value), node.attr)
327
328 + def visit_Subscript(self, node: ast.Subscript) -> Any:
329 + return self.visit(node.value)[self.visit(node.slice)]
330 +
331
332 def get_ast_full_name(node: ast.AST) -> str | None:
333 if isinstance(node, ast.Name):
@@ -310,15 +339,23 @@ def get_ast_full_name(node: ast.AST) -> str | None:
339 return None
340
341
313 -def get_function_source_node(fn: Callable[..., Any]) -> ast.FunctionDef:
342 +def get_function_source_node(fn: Callable[..., Any]) -> ast.FunctionDef | ast.Lambda:
343 source_fn = getattr(fn, "fn", fn)
344 source = textwrap.dedent(inspect.getsource(source_fn))
345 tree = ast.parse(source)
346 function_defs = [node for node in tree.body if isinstance(node, ast.FunctionDef)]
318 - if len(function_defs) != 1:
319 - name = getattr(source_fn, "__name__", type(source_fn).__name__)
320 - raise ValueError(f"Expected one function in {name}, found {len(function_defs)}")
321 - return function_defs[0]
347 + if len(function_defs) == 1:
348 + return function_defs[0]
349 +
350 + lambdas = [node for node in ast.walk(tree) if isinstance(node, ast.Lambda)]
351 + if len(lambdas) == 1:
352 + return lambdas[0]
353 +
354 + name = getattr(source_fn, "__name__", type(source_fn).__name__)
355 + raise ValueError(
356 + f"Expected one function or lambda in {name}, found "
357 + f"{len(function_defs)} functions and {len(lambdas)} lambdas"
358 + )
359
360
361 def _eval_dispatch_expr(
@@ -349,8 +386,11 @@ def _collect_input_names(
386
387 def _collect_expression_body(
388 fn: Callable[..., Any],
352 - function_def: ast.FunctionDef,
389 + function_def: ast.FunctionDef | ast.Lambda,
390 ) -> tuple[list[tuple[str, ast.AST]], ast.AST]:
391 + if isinstance(function_def, ast.Lambda):
392 + return [], function_def.body
393 +
394 local_exprs: list[tuple[str, ast.AST]] = []
395 for statement in function_def.body:
396 if (
@@ -421,6 +461,10 @@ def _trace_compile_key_dispatch(
461 source_fn = getattr(fn, "__func__", fn)
462 globals_ = source_fn.__globals__
463 function_def = get_function_source_node(fn)
464 + if isinstance(function_def, ast.Lambda):
465 + raise _dispatch_expr_error(
466 + function_def, "Dispatch must be a function definition"
467 + )
468
469 local_exprs, return_expr = _collect_expression_body(fn, function_def)
470 if not isinstance(return_expr, ast.Call):
@@ -431,13 +475,37 @@ def _trace_compile_key_dispatch(
475
476 field_exprs: list[tuple[str, ast.AST]] = []
477 defaults, candidate_names = _function_trace_inputs(fn)
478 + # Fields captured by dispatch **kwargs are forwarded, not AST-evaluated.
479 + signature = inspect.signature(fn)
480 + variadic_keyword = next(
481 + (
482 + name
483 + for name, parameter in signature.parameters.items()
484 + if parameter.kind is inspect.Parameter.VAR_KEYWORD
485 + ),
486 + None,
487 + )
488 + candidate_names.discard(variadic_keyword)
489 input_names: set[str] = set()
490 local_names = {name for name, _ in local_exprs}
491 for _, expr in local_exprs:
492 input_names.update(_collect_input_names(expr, candidate_names))
493 + named_parameters: frozenset[str] | None = None
494 for keyword in return_expr.keywords:
495 + # CompileKey may unpack only that **kwargs parameter, once.
496 if keyword.arg is None:
440 - raise ValueError(f"{fn.__name__} cannot use **kwargs in CompileKey")
497 + if (
498 + named_parameters is not None
499 + or variadic_keyword is None
500 + or not isinstance(keyword.value, ast.Name)
501 + or keyword.value.id != variadic_keyword
502 + ):
503 + raise ValueError(
504 + f"{fn.__name__} may unpack only its own **kwargs parameter "
505 + "once in CompileKey"
506 + )
507 + named_parameters = frozenset(candidate_names)
508 + continue
509 field_exprs.append((keyword.arg, keyword.value))
510 input_names.update(
511 _collect_input_names(keyword.value, candidate_names, local_names)
@@ -448,6 +516,7 @@ def _trace_compile_key_dispatch(
516 tuple(field_exprs),
517 globals_,
518 frozenset(input_names),
519 + named_parameters,
520 defaults,
521 )
522
@@ -482,10 +551,32 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):
551 CompileKey: type[CompileKeyT]
552
553 def __init__(self) -> None:
485 - self.compile_key_dispatch_trace = _trace_compile_key_dispatch(self.dispatch)
554 + self._dispatch_trace = _trace_compile_key_dispatch(self.dispatch)
555 + self._compiled_cache: dict[Any, Any] = {}
556
557 def compile_key(self, kwargs: Mapping[str, Any]) -> CompileKeyT:
488 - return self.compile_key_dispatch_trace.compile_key(self.CompileKey, kwargs)
558 + return self._dispatch_trace.compile_key(self.CompileKey, kwargs)
559 +
560 + def _get_or_compile(
561 + self,
562 + compile_key: CompileKeyT,
563 + *,
564 + runtime_context: Mapping[str, Any] | None = None,
565 + ) -> Any:
566 + """Return a cached executor, compiling it on a monitored cache miss."""
567 + if compile_key not in self._compiled_cache:
568 + self.compile(compile_key)
569 +
570 + try:
571 + return self._compiled_cache[compile_key]
572 + except KeyError as exc:
573 + details = [f"compile_key={compile_key!r}"]
574 + if runtime_context:
575 + details.append(f"runtime_context={dict(runtime_context)!r}")
576 + raise RuntimeError(
577 + f"{type(self).__name__}.compile(...) did not cache its JIT "
578 + f"executor ({', '.join(details)})"
579 + ) from exc
580
581 def _trace_dispatch(
582 self, dispatch: CompileKeyDispatchFn[CompileKeyT]
@@ -506,7 +597,11 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):
597 predicate_trace = (
598 _trace_warmup_predicate(_when) if _when is not None else None
599 )
509 - input_names = compile_key_dispatch_trace.input_names
600 + # Unmatched **kwargs fields also belong to the expansion space.
601 + available_names = set(kwargs).union(
602 + *(group.rows[0] for group in input_groups)
603 + )
604 + input_names = compile_key_dispatch_trace.input_names_for(available_names)
605 if predicate_trace is not None:
606 input_names = input_names | predicate_trace.input_names
607 expanded_input_groups = tuple(
@@ -565,7 +660,105 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):
660 """Compile one warmup key."""
661 raise NotImplementedError
662
663 + def register_warmup(self, *args: Any, **kwargs: Any) -> None:
664 + """Register this kernel with the active runner's warmup registry."""
665 + JitWarmupRegistry.register(self, *args, **kwargs)
666 +
667 def warmup(self, *args: Any, **kwargs: Any) -> None:
668 """Compile this kernel's warmup keys."""
669 for compile_key in self.get_warmup_keys(*args, **kwargs):
670 self.compile(compile_key)
671 +
672 +
673 +class JitWarmupRegistry:
674 + """Collect and compile JIT kernels selected during runner setup."""
675 +
676 + _active: ContextVar[JitWarmupRegistry | None] = ContextVar(
677 + "active_jit_warmup_registry",
678 + default=None,
679 + )
680 +
681 + def __init__(self, vllm_config: Any) -> None:
682 + self.vllm_config = vllm_config
683 + self._registrations: dict[
684 + VllmJitKernel[Any],
685 + list[tuple[tuple[Any, ...], dict[str, Any]]],
686 + ] = {}
687 +
688 + @contextmanager
689 + def activate(self) -> Iterator[None]:
690 + """Collect registrations made in this context."""
691 + token = self._active.set(self)
692 + try:
693 + yield
694 + finally:
695 + self._active.reset(token)
696 +
697 + @classmethod
698 + def register(
699 + cls,
700 + kernel: VllmJitKernel[Any],
701 + *args: Any,
702 + **kwargs: Any,
703 + ) -> None:
704 + """Register a kernel with the active registry, if one exists."""
705 + registry = cls._active.get()
706 + if registry is not None:
707 + registry._add(kernel, args, kwargs)
708 +
709 + def _add(
710 + self,
711 + kernel: VllmJitKernel[Any],
712 + args: tuple[Any, ...],
713 + kwargs: dict[str, Any],
714 + ) -> None:
715 + registrations = self._registrations.setdefault(kernel, [])
716 + if (
717 + not args
718 + and not kwargs
719 + and any(
720 + not registered_args and not registered_kwargs
721 + for registered_args, registered_kwargs in registrations
722 + )
723 + ):
724 + return
725 + registrations.append((args, kwargs))
726 +
727 + def __len__(self) -> int:
728 + return sum(len(registrations) for registrations in self._registrations.values())
729 +
730 + def warmup(self) -> None:
731 + """Expand registrations and compile each wrapper/key pair once."""
732 + from tqdm import tqdm
733 +
734 + from vllm.distributed import is_global_first_rank
735 +
736 + kernel_items: list[tuple[VllmJitKernel[Any], dict[Any, None]]] = []
737 + for kernel, registrations in self._registrations.items():
738 + compile_keys: dict[Any, None] = {}
739 + for args, kwargs in registrations:
740 + if not args and not kwargs:
741 + args = (self.vllm_config,)
742 + for compile_key in kernel.get_warmup_keys(*args, **kwargs):
743 + compile_keys[compile_key] = None
744 + if compile_keys:
745 + kernel_items.append((kernel, compile_keys))
746 +
747 + if not kernel_items:
748 + return
749 +
750 + total_keys = sum(len(compile_keys) for _, compile_keys in kernel_items)
751 + with tqdm(
752 + kernel_items,
753 + desc=f"JIT kernel warmup ({total_keys} compile keys)",
754 + disable=not is_global_first_rank(),
755 + dynamic_ncols=True,
756 + unit="kernel",
757 + ) as progress:
758 + for kernel, compile_keys in progress:
759 + progress.set_postfix_str(
760 + f"{kernel.__class__.__name__} ({len(compile_keys)} keys)",
761 + refresh=False,
762 + )
763 + for compile_key in compile_keys:
764 + kernel.compile(compile_key)
vllm/model_executor/warmup/jit_warmup_triton_helper.py
+38
@@ -12,6 +12,42 @@ from vllm.model_executor.warmup.jit_warmup import (
12 )
13
14
15 +def triton_scalar_specialization_rep(value: int) -> int:
16 + """Return an integer with the same default Triton JIT specialization.
17 +
18 + For an ordinary integer argument, Triton's cache key contains its inferred
19 + type (``i32``, ``i64``, or ``u64``) and one of three value classes:
20 +
21 + * ``1`` is specialized as the exact constant ``1``.
22 + * Multiples of 16 receive a ``tt.divisibility = 16`` attribute.
23 + * All other values have no value specialization.
24 +
25 + Warmup only needs one concrete value for each cache-key class. This helper
26 + returns ``1`` for the exact-one class and otherwise returns a divisible or
27 + generic representative while preserving the inferred integer type.
28 +
29 + This applies only to non-``constexpr`` integer arguments using Triton's
30 + default specialization. Do not use it for arguments listed in
31 + ``do_not_specialize`` or ``do_not_specialize_on_alignment``.
32 + """
33 + if value == 1:
34 + return 1
35 +
36 + if -(1 << 31) <= value < (1 << 31):
37 + divisible_rep = 16
38 + generic_rep = 2
39 + elif -(1 << 63) <= value < (1 << 63):
40 + divisible_rep = 1 << 31
41 + generic_rep = (1 << 31) + 1
42 + elif 0 <= value < (1 << 64):
43 + divisible_rep = 1 << 63
44 + generic_rep = (1 << 63) + 1
45 + else:
46 + raise OverflowError(f"Integer {value} is outside Triton's scalar range")
47 +
48 + return divisible_rep if value % 16 == 0 else generic_rep
49 +
50 +
51 @dataclass(frozen=True)
52 class TritonWarmupTensor:
53 # Compile-only tensor descriptor for Triton pointer specialization.
@@ -167,6 +203,8 @@ def trace_triton_kernel_specialization_args(
203 kernel: Callable[..., Any],
204 ) -> tuple[str, ...]:
205 function_def = get_function_source_node(kernel)
206 + if not isinstance(function_def, ast.FunctionDef):
207 + raise ValueError("Expected Triton kernel to be defined as a function")
208 source_fn = getattr(kernel, "fn", kernel)
209 arg_names = tuple(inspect.signature(source_fn).parameters)
210 constexpr_args = _triton_constexpr_arg_names(kernel, function_def, arg_names)
vllm/model_executor/warmup/kernel_warmup.py
+17 -6
@@ -6,6 +6,7 @@ This is useful specifically for JIT'ed kernels as we don't want JIT'ing to
6 happen during model execution.
7 """
8
9 +import time
10 from typing import TYPE_CHECKING
11
12 import torch
@@ -36,9 +37,6 @@ from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup
37 from vllm.model_executor.warmup.sparse_mla_triton_warmup import (
38 sparse_mla_triton_warmup,
39 )
39 -from vllm.model_executor.warmup.v1_block_table_warmup import (
40 - warm_v1_block_table_kernels,
41 -)
40 from vllm.platforms import current_platform
41 from vllm.utils.deep_gemm import is_deep_gemm_supported
42 from vllm.utils.flashinfer import has_flashinfer
@@ -103,15 +101,28 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
101 )
102
103 if not worker.use_v2_model_runner:
106 - # Pooling models do not use the generation slot-mapping path.
107 - if not worker.model_runner.is_pooling_model:
108 - warm_v1_block_table_kernels(worker.model_runner)
104 # The KV-block zeroing kernel is driven by the scheduler's
105 # `new_block_ids_to_zero`, so no dummy run ever reaches it.
106 zeroer = getattr(worker.model_runner, "_kv_block_zeroer", None)
107 if zeroer is not None:
108 zeroer.warmup(worker.model_runner.kv_cache_config.num_blocks)
109
110 + if worker.vllm_config.kernel_config.enable_jit_warmup:
111 + logger.info("JIT kernel warmup starting.")
112 + jit_warmup_start = time.perf_counter()
113 + try:
114 + worker.model_runner.jit_warmup_registry.warmup()
115 + except Exception:
116 + logger.exception(
117 + "JIT kernel warmup failed after %.2fs.",
118 + time.perf_counter() - jit_warmup_start,
119 + )
120 + raise
121 + logger.info(
122 + "JIT kernel warmup finished in %.2fs.",
123 + time.perf_counter() - jit_warmup_start,
124 + )
125 +
126 qwen_triton_warmup(worker.model_runner, worker.vllm_config.model_config)
127
128 compilation_config = worker.vllm_config.compilation_config
vllm/model_executor/warmup/v1_block_table_warmup.py deleted
-29
@@ -1,29 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Warm up v1 block-table Triton kernels."""
4 -
5 -from typing import TYPE_CHECKING
6 -
7 -import torch
8 -
9 -if TYPE_CHECKING:
10 - from vllm.v1.worker.gpu_model_runner import GPUModelRunner
11 -
12 -_SLOT_MAPPING_WARMUP_TOKENS = 8
13 -
14 -
15 -def warm_v1_block_table_kernels(runner: "GPUModelRunner") -> None:
16 - """JIT-compile ``_compute_slot_mapping_kernel`` for the real block tables."""
17 -
18 - device = runner.device
19 - block_table = runner.input_batch.block_table
20 - num_tokens = min(
21 - _SLOT_MAPPING_WARMUP_TOKENS,
22 - runner.scheduler_config.max_num_batched_tokens,
23 - )
24 - if num_tokens <= 0:
25 - return
26 -
27 - query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
28 - positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
29 - block_table.compute_slot_mapping(1, query_start_loc, positions)
vllm/v1/worker/block_table.py
+154 -67
@@ -2,13 +2,22 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import math
5 +from dataclasses import dataclass
6 from enum import Enum
7 +from typing import Any
8
9 import numpy as np
10 import torch
11
12 from vllm.distributed import get_dcp_group, get_pcp_group
13 from vllm.logger import init_logger
14 +from vllm.model_executor.warmup.jit_warmup import (
15 + VllmJitKernel,
16 +)
17 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (
18 + TritonWarmupTensor,
19 + triton_scalar_specialization_rep,
20 +)
21 from vllm.triton_utils import tl, triton
22 from vllm.utils.math_utils import cdiv
23 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
@@ -134,6 +143,16 @@ class BlockTable:
143 self.dcp_rank = 0
144 self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size
145 self.slot_mapping_mode = slot_mapping_mode
146 + if self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT:
147 + _COMPUTE_SLOT_MAPPING_KERNEL.register_warmup(
148 + kv_cache_block_size=self.kv_cache_block_size,
149 + blocks_per_kv_block=self.blocks_per_kv_block,
150 + total_cp_world_size=self.dcp_world_size,
151 + total_cp_rank=self.dcp_rank,
152 + cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
153 + block_table_stride=self.block_table.gpu.stride(0),
154 + block_size=self.block_size,
155 + )
156
157 def append_row(
158 self,
@@ -192,7 +211,8 @@ class BlockTable:
211 return
212 assert self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT
213
195 - _compute_slot_mapping_kernel[(num_reqs + 1,)](
214 + _COMPUTE_SLOT_MAPPING_KERNEL(
215 + num_reqs,
216 num_tokens,
217 self.max_num_batched_tokens,
218 query_start_loc,
@@ -201,13 +221,11 @@ class BlockTable:
221 self.block_table.gpu.stride(0),
222 self.block_size,
223 self.slot_mapping.gpu,
204 - KV_CACHE_BLOCK_SIZE=self.kv_cache_block_size,
205 - BLOCKS_PER_KV_BLOCK=self.blocks_per_kv_block,
206 - TOTAL_CP_WORLD_SIZE=self.dcp_world_size,
207 - TOTAL_CP_RANK=self.dcp_rank,
208 - CP_KV_CACHE_INTERLEAVE_SIZE=self.cp_kv_cache_interleave_size,
209 - PAD_ID=PAD_SLOT_ID,
210 - BLOCK_SIZE=1024,
224 + self.kv_cache_block_size,
225 + self.blocks_per_kv_block,
226 + self.dcp_world_size,
227 + self.dcp_rank,
228 + self.cp_kv_cache_interleave_size,
229 )
230
231 def commit_block_table(self, num_reqs: int) -> None:
@@ -376,67 +394,136 @@ class MultiGroupBlockTable:
394 return self.block_tables[idx]
395
396
379 -@triton.jit(do_not_specialize=["num_tokens", "max_num_tokens"])
380 -def _compute_slot_mapping_kernel(
381 - num_tokens,
382 - max_num_tokens,
383 - query_start_loc_ptr, # [num_reqs + 1], int32
384 - positions_ptr, # [num_tokens], int64
385 - block_table_ptr, # [max_num_reqs, max_num_blocks_per_req], int32 (flat)
386 - block_table_stride, # max_num_blocks_per_req
387 - block_size,
388 - slot_mapping_ptr, # [max_num_tokens], int64
389 - KV_CACHE_BLOCK_SIZE: tl.constexpr,
390 - BLOCKS_PER_KV_BLOCK: tl.constexpr,
391 - TOTAL_CP_WORLD_SIZE: tl.constexpr,
392 - TOTAL_CP_RANK: tl.constexpr,
393 - CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr,
394 - PAD_ID: tl.constexpr,
395 - BLOCK_SIZE: tl.constexpr,
396 -):
397 - req_idx = tl.program_id(0)
398 -
399 - if req_idx == tl.num_programs(0) - 1:
400 - # Pad remaining slots for CUDA graph compatibility.
401 - for i in range(num_tokens, max_num_tokens, BLOCK_SIZE):
397 +class ComputeSlotMappingKernel(VllmJitKernel["ComputeSlotMappingKernel.CompileKey"]):
398 + triton_block_size = 1024
399 +
400 + @dataclass(frozen=True)
401 + class CompileKey:
402 + kv_cache_block_size: int
403 + blocks_per_kv_block: int
404 + total_cp_world_size: int
405 + total_cp_rank: int
406 + cp_kv_cache_interleave_size: int
407 + block_table_stride: int
408 + block_size: int
409 +
410 + @staticmethod
411 + @triton.jit(do_not_specialize=["num_tokens", "max_num_tokens"])
412 + def kernel(
413 + num_tokens,
414 + max_num_tokens,
415 + query_start_loc_ptr, # [num_reqs + 1], int32
416 + positions_ptr, # [num_tokens], int64
417 + block_table_ptr, # [max_num_reqs, max_num_blocks_per_req], int32 (flat)
418 + block_table_stride, # max_num_blocks_per_req
419 + block_size,
420 + slot_mapping_ptr, # [max_num_tokens], int64
421 + KV_CACHE_BLOCK_SIZE: tl.constexpr,
422 + BLOCKS_PER_KV_BLOCK: tl.constexpr,
423 + TOTAL_CP_WORLD_SIZE: tl.constexpr,
424 + TOTAL_CP_RANK: tl.constexpr,
425 + CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr,
426 + PAD_ID: tl.constexpr,
427 + BLOCK_SIZE: tl.constexpr,
428 + ):
429 + req_idx = tl.program_id(0)
430 +
431 + if req_idx == tl.num_programs(0) - 1:
432 + # Pad remaining slots for CUDA graph compatibility.
433 + for i in range(num_tokens, max_num_tokens, BLOCK_SIZE):
434 + offsets = i + tl.arange(0, BLOCK_SIZE)
435 + tl.store(
436 + slot_mapping_ptr + offsets,
437 + PAD_ID,
438 + mask=offsets < max_num_tokens,
439 + )
440 + return
441 +
442 + start_idx = tl.load(query_start_loc_ptr + req_idx).to(tl.int64)
443 + end_idx = tl.load(query_start_loc_ptr + req_idx + 1).to(tl.int64)
444 +
445 + virtual_block_size = KV_CACHE_BLOCK_SIZE * TOTAL_CP_WORLD_SIZE
446 + row_offset = req_idx * block_table_stride
447 + for i in range(start_idx, end_idx, BLOCK_SIZE):
448 offsets = i + tl.arange(0, BLOCK_SIZE)
403 - tl.store(
404 - slot_mapping_ptr + offsets,
405 - PAD_ID,
406 - mask=offsets < max_num_tokens,
449 + mask = offsets < end_idx
450 + pos = tl.load(positions_ptr + offsets, mask=mask, other=0)
451 + virtual_block_indices = pos // virtual_block_size
452 + virtual_block_offsets = pos - virtual_block_indices * virtual_block_size
453 + is_local = (
454 + virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE
455 + ) % TOTAL_CP_WORLD_SIZE == TOTAL_CP_RANK
456 + local_block_offsets = (
457 + virtual_block_offsets
458 + // (TOTAL_CP_WORLD_SIZE * CP_KV_CACHE_INTERLEAVE_SIZE)
459 + ) * CP_KV_CACHE_INTERLEAVE_SIZE + (
460 + virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE
461 + )
462 +
463 + block_indices = (
464 + virtual_block_indices * BLOCKS_PER_KV_BLOCK
465 + + local_block_offsets // block_size
466 )
408 - return
409 -
410 - start_idx = tl.load(query_start_loc_ptr + req_idx).to(tl.int64)
411 - end_idx = tl.load(query_start_loc_ptr + req_idx + 1).to(tl.int64)
412 -
413 - virtual_block_size = KV_CACHE_BLOCK_SIZE * TOTAL_CP_WORLD_SIZE
414 - row_offset = req_idx * block_table_stride
415 - for i in range(start_idx, end_idx, BLOCK_SIZE):
416 - offsets = i + tl.arange(0, BLOCK_SIZE)
417 - mask = offsets < end_idx
418 - pos = tl.load(positions_ptr + offsets, mask=mask, other=0)
419 - virtual_block_indices = pos // virtual_block_size
420 - virtual_block_offsets = pos - virtual_block_indices * virtual_block_size
421 - is_local = (
422 - virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE
423 - ) % TOTAL_CP_WORLD_SIZE == TOTAL_CP_RANK
424 - local_block_offsets = (
425 - virtual_block_offsets // (TOTAL_CP_WORLD_SIZE * CP_KV_CACHE_INTERLEAVE_SIZE)
426 - ) * CP_KV_CACHE_INTERLEAVE_SIZE + (
427 - virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE
467 + block_numbers = tl.load(
468 + block_table_ptr + row_offset + block_indices,
469 + mask=mask & is_local,
470 + other=0,
471 + ).to(tl.int64)
472 + slot_offsets = local_block_offsets % block_size
473 + slot_ids = block_numbers * block_size + slot_offsets
474 + slot_ids = tl.where(is_local, slot_ids, PAD_ID)
475 + tl.store(slot_mapping_ptr + offsets, slot_ids, mask=mask)
476 +
477 + def dispatch( # type: ignore[override]
478 + self,
479 + *,
480 + block_table_stride: int,
481 + block_size: int,
482 + **compile_key_fields: int,
483 + ) -> CompileKey:
484 + return self.CompileKey(
485 + **compile_key_fields,
486 + block_table_stride=triton_scalar_specialization_rep(block_table_stride),
487 + block_size=triton_scalar_specialization_rep(block_size),
488 )
489
430 - block_indices = (
431 - virtual_block_indices * BLOCKS_PER_KV_BLOCK
432 - + local_block_offsets // block_size
490 + def get_warmup_keys(self, **dispatch_kwargs: int) -> list[CompileKey]:
491 + return self._trace_dispatch(self.dispatch)(**dispatch_kwargs)
492 +
493 + def compile(self, compile_key: CompileKey) -> None:
494 + warmup = getattr(self.kernel, "warmup", None)
495 + assert warmup is not None
496 + int32_ptr = TritonWarmupTensor(torch.int32)
497 + int64_ptr = TritonWarmupTensor(torch.int64)
498 + warmup(
499 + 2, # arbitrary, num_tokens in do_not_specialize
500 + 2, # arbitrary, max_num_tokens in do_not_specialize
501 + int32_ptr,
502 + int64_ptr,
503 + int32_ptr,
504 + compile_key.block_table_stride,
505 + compile_key.block_size,
506 + int64_ptr,
507 + KV_CACHE_BLOCK_SIZE=compile_key.kv_cache_block_size,
508 + BLOCKS_PER_KV_BLOCK=compile_key.blocks_per_kv_block,
509 + TOTAL_CP_WORLD_SIZE=compile_key.total_cp_world_size,
510 + TOTAL_CP_RANK=compile_key.total_cp_rank,
511 + CP_KV_CACHE_INTERLEAVE_SIZE=compile_key.cp_kv_cache_interleave_size,
512 + PAD_ID=PAD_SLOT_ID,
513 + BLOCK_SIZE=self.triton_block_size,
514 + grid=(2,),
515 )
434 - block_numbers = tl.load(
435 - block_table_ptr + row_offset + block_indices,
436 - mask=mask & is_local,
437 - other=0,
438 - ).to(tl.int64)
439 - slot_offsets = local_block_offsets % block_size
440 - slot_ids = block_numbers * block_size + slot_offsets
441 - slot_ids = tl.where(is_local, slot_ids, PAD_ID)
442 - tl.store(slot_mapping_ptr + offsets, slot_ids, mask=mask)
516 +
517 + def __call__(
518 + self,
519 + num_reqs: int,
520 + *args: Any,
521 + ) -> None:
522 + self.kernel[(num_reqs + 1,)](
523 + *args,
524 + PAD_ID=PAD_SLOT_ID,
525 + BLOCK_SIZE=self.triton_block_size,
526 + )
527 +
528 +
529 +_COMPUTE_SLOT_MAPPING_KERNEL = ComputeSlotMappingKernel()
vllm/v1/worker/cpu_model_runner.py
+1 -1
@@ -74,7 +74,7 @@ class CPUModelRunner(GPUModelRunner):
74
75 import vllm.v1.worker.block_table
76
77 - vllm.v1.worker.block_table._compute_slot_mapping_kernel = (
77 + vllm.v1.worker.block_table._COMPUTE_SLOT_MAPPING_KERNEL.kernel = (
78 cpu_tl.compute_slot_mapping_kernel
79 )
80
vllm/v1/worker/gpu/model_runner.py
+2
@@ -51,6 +51,7 @@ from vllm.model_executor.offloader import (
51 get_offloader,
52 set_offloader,
53 )
54 +from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry
55 from vllm.multimodal import MULTIMODAL_REGISTRY
56 from vllm.multimodal.encoder_budget import (
57 MultiModalBudget,
@@ -170,6 +171,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
171 self.speculative_config is not None and self.speculative_config.use_dspark()
172 )
173 self.observability_config = vllm_config.observability_config
174 + self.jit_warmup_registry = JitWarmupRegistry(vllm_config)
175
176 self.device = device
177 self.dtype = self.model_config.dtype
vllm/v1/worker/gpu_model_runner.py
+53 -46
@@ -100,6 +100,7 @@ from vllm.model_executor.offloader import (
100 get_offloader,
101 set_offloader,
102 )
103 +from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry
104 from vllm.multimodal import MULTIMODAL_REGISTRY
105 from vllm.multimodal.encoder_budget import MultiModalBudget
106 from vllm.multimodal.inputs import (
@@ -516,6 +517,7 @@ class GPUModelRunner(
517 self.scheduler_config = vllm_config.scheduler_config
518 self.speculative_config = vllm_config.speculative_config
519 self.observability_config = vllm_config.observability_config
520 + self.jit_warmup_registry = JitWarmupRegistry(vllm_config)
521
522 model_config = self.model_config
523 cache_config = self.cache_config
@@ -750,34 +752,37 @@ class GPUModelRunner(
752 self._init_kernel_block_sizes = [placeholder_block_size]
753 self._init_max_num_blocks = [placeholder_max_num_blocks]
754 self._init_slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT]
753 - self.input_batch = InputBatch(
754 - max_num_reqs=self.max_num_reqs,
755 - # We need to use the encoder length for encoder-decoder
756 - # because of KV cache for cross-attention.
757 - max_model_len=max(self.max_model_len, self.max_encoder_len),
758 - max_num_batched_tokens=self.max_num_tokens,
759 - device=self.device,
760 - vocab_size=self.model_config.get_vocab_size(),
761 - block_sizes=[placeholder_block_size],
762 - kernel_block_sizes=[placeholder_block_size],
763 - max_num_blocks_per_req=[placeholder_max_num_blocks],
764 - num_spec_tokens=self.num_spec_tokens,
765 - logitsprocs=build_logitsprocs(
766 - self.vllm_config,
767 - self.device,
768 - PIN_MEMORY,
769 - self.is_pooling_model,
770 - custom_logitsprocs,
771 - ),
772 - # We currently don't know whether a particular custom logits processor
773 - # uses output token ids so we set this conservatively. Thinking-budget
774 - # tracking is requested dynamically when a budgeted request is in the batch.
775 - logitsprocs_need_output_token_ids=bool(custom_logitsprocs),
776 - is_pooling_model=self.is_pooling_model,
777 - cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,
778 - reasoning_config=self.vllm_config.reasoning_config,
779 - use_replayssm=self.cache_config.use_replayssm,
780 - )
755 + # Capture warmup providers registered by the initial placeholder InputBatch
756 + with self.jit_warmup_registry.activate():
757 + self.input_batch = InputBatch(
758 + max_num_reqs=self.max_num_reqs,
759 + # We need to use the encoder length for encoder-decoder
760 + # because of KV cache for cross-attention.
761 + max_model_len=max(self.max_model_len, self.max_encoder_len),
762 + max_num_batched_tokens=self.max_num_tokens,
763 + device=self.device,
764 + vocab_size=self.model_config.get_vocab_size(),
765 + block_sizes=[placeholder_block_size],
766 + kernel_block_sizes=[placeholder_block_size],
767 + max_num_blocks_per_req=[placeholder_max_num_blocks],
768 + num_spec_tokens=self.num_spec_tokens,
769 + logitsprocs=build_logitsprocs(
770 + self.vllm_config,
771 + self.device,
772 + PIN_MEMORY,
773 + self.is_pooling_model,
774 + custom_logitsprocs,
775 + ),
776 + # We currently don't know whether a particular custom logits processor
777 + # uses output token ids so we set this conservatively. Thinking-budget
778 + # tracking is requested dynamically when a budgeted request is in the
779 + # batch.
780 + logitsprocs_need_output_token_ids=bool(custom_logitsprocs),
781 + is_pooling_model=self.is_pooling_model,
782 + cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,
783 + reasoning_config=self.vllm_config.reasoning_config,
784 + use_replayssm=self.cache_config.use_replayssm,
785 + )
786
787 # Separate cuda stream for overlapping transfer of sampled token ids from
788 # GPU to CPU when async scheduling is enabled.
@@ -7416,24 +7421,26 @@ class GPUModelRunner(
7421 self._init_kernel_block_sizes = kernel_block_sizes
7422 self._init_max_num_blocks = max_num_blocks
7423 self._init_slot_mapping_modes = slot_mapping_modes
7419 - self.input_batch = InputBatch(
7420 - max_num_reqs=self.max_num_reqs,
7421 - max_model_len=max_model_len,
7422 - max_num_batched_tokens=self.max_num_tokens,
7423 - device=self.device,
7424 - vocab_size=self.model_config.get_vocab_size(),
7425 - block_sizes=block_sizes,
7426 - kernel_block_sizes=kernel_block_sizes,
7427 - max_num_blocks_per_req=max_num_blocks,
7428 - num_spec_tokens=self.num_spec_tokens,
7429 - logitsprocs=self.input_batch.logitsprocs,
7430 - logitsprocs_need_output_token_ids=self.input_batch.logitsprocs_need_output_token_ids,
7431 - is_pooling_model=self.is_pooling_model,
7432 - cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,
7433 - reasoning_config=self.vllm_config.reasoning_config,
7434 - use_replayssm=self.cache_config.use_replayssm,
7435 - slot_mapping_modes=slot_mapping_modes,
7436 - )
7424 + # Capture warmup providers registered after final KV-cache geometry is known
7425 + with self.jit_warmup_registry.activate():
7426 + self.input_batch = InputBatch(
7427 + max_num_reqs=self.max_num_reqs,
7428 + max_model_len=max_model_len,
7429 + max_num_batched_tokens=self.max_num_tokens,
7430 + device=self.device,
7431 + vocab_size=self.model_config.get_vocab_size(),
7432 + block_sizes=block_sizes,
7433 + kernel_block_sizes=kernel_block_sizes,
7434 + max_num_blocks_per_req=max_num_blocks,
7435 + num_spec_tokens=self.num_spec_tokens,
7436 + logitsprocs=self.input_batch.logitsprocs,
7437 + logitsprocs_need_output_token_ids=self.input_batch.logitsprocs_need_output_token_ids,
7438 + is_pooling_model=self.is_pooling_model,
7439 + cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,
7440 + reasoning_config=self.vllm_config.reasoning_config,
7441 + use_replayssm=self.cache_config.use_replayssm,
7442 + slot_mapping_modes=slot_mapping_modes,
7443 + )
7444
7445 assert self._init_block_sizes == block_sizes, (
7446 f"InputBatch block_sizes {self._init_block_sizes} != "