Skip to content

vllm.model_executor.warmup.jit_warmup

Shared interfaces and tracing helpers for explicit JIT warmup keys.

Classes:

  • VllmJitKernel –

    Kernel wrapper that owns dispatch, warmup keys, and compilation.

Functions:

  • zip_inputs –

    Group row-wise dispatch inputs that should be expanded in lockstep.

VllmJitKernel

Bases: Generic[CompileKeyT], ABC

Kernel wrapper that owns dispatch, warmup keys, and compilation.

Methods:

  • compile –

    Compile one warmup key.

  • dispatch –

    Build one compile key from one concrete dispatch point.

  • get_warmup_keys –

    Return compile keys that should be warmed for this kernel.

  • warmup –

    Compile this kernel's warmup keys.

Source code in vllm/model_executor/warmup/jit_warmup.py 
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
class VllmJitKernel(Generic[CompileKeyT], ABC):
    """Kernel wrapper that owns dispatch, warmup keys, and compilation."""

    CompileKey: type[CompileKeyT]

    def __init__(self) -> None:
        self.compile_key_dispatch_trace = _trace_compile_key_dispatch(self.dispatch)

    def compile_key(self, kwargs: Mapping[str, Any]) -> CompileKeyT:
        return self.compile_key_dispatch_trace.compile_key(self.CompileKey, kwargs)

    def _trace_dispatch(
        self, dispatch: CompileKeyDispatchFn[CompileKeyT]
    ) -> Callable[..., list[CompileKeyT]]:
        compile_key_dispatch_trace = _trace_compile_key_dispatch(dispatch)

        def traced(
            *input_groups: _WarmupInputRows,
            **kwargs: WarmupValues,
        ) -> list[CompileKeyT]:
            for group in input_groups:
                if not isinstance(group, _WarmupInputRows):
                    raise TypeError(
                        "_trace_dispatch positional arguments must be "
                        "zip_inputs(...) groups"
                    )
            expanded_input_groups = tuple(
                _expand_warmup_input_rows(
                    group.rows, compile_key_dispatch_trace.input_names
                )
                for group in input_groups
            )
            expanded_kwargs = _expand_warmup_value_grid(
                kwargs, compile_key_dispatch_trace.input_names
            )
            dispatch_value_groups = (*expanded_input_groups, expanded_kwargs)
            return list(
                dict.fromkeys(
                    compile_key_dispatch_trace.compile_key(
                        self.CompileKey, _merge_warmup_kwargs(dispatch_values)
                    )
                    for dispatch_values in itertools.product(*dispatch_value_groups)
                )
            )

        return traced

    @abstractmethod
    def dispatch(self, **kwargs: Any) -> CompileKeyT:
        """Build one compile key from one concrete dispatch point."""
        raise NotImplementedError

    @abstractmethod
    def get_warmup_keys(self, *args: Any, **kwargs: Any) -> list[CompileKeyT]:
        """Return compile keys that should be warmed for this kernel."""
        raise NotImplementedError

    @abstractmethod
    def compile(self, compile_key: CompileKeyT) -> None:
        """Compile one warmup key."""
        raise NotImplementedError

    def warmup(self, *args: Any, **kwargs: Any) -> None:
        """Compile this kernel's warmup keys."""
        for compile_key in self.get_warmup_keys(*args, **kwargs):
            self.compile(compile_key)

compile(compile_key) abstractmethod

Compile one warmup key.

Source code in vllm/model_executor/warmup/jit_warmup.py 
466
467
468
469
@abstractmethod
def compile(self, compile_key: CompileKeyT) -> None:
    """Compile one warmup key."""
    raise NotImplementedError

dispatch(**kwargs) abstractmethod

Build one compile key from one concrete dispatch point.

Source code in vllm/model_executor/warmup/jit_warmup.py 
456
457
458
459
@abstractmethod
def dispatch(self, **kwargs: Any) -> CompileKeyT:
    """Build one compile key from one concrete dispatch point."""
    raise NotImplementedError

get_warmup_keys(*args, **kwargs) abstractmethod

Return compile keys that should be warmed for this kernel.

Source code in vllm/model_executor/warmup/jit_warmup.py 
461
462
463
464
@abstractmethod
def get_warmup_keys(self, *args: Any, **kwargs: Any) -> list[CompileKeyT]:
    """Return compile keys that should be warmed for this kernel."""
    raise NotImplementedError

warmup(*args, **kwargs)

Compile this kernel's warmup keys.

Source code in vllm/model_executor/warmup/jit_warmup.py 
471
472
473
474
def warmup(self, *args: Any, **kwargs: Any) -> None:
    """Compile this kernel's warmup keys."""
    for compile_key in self.get_warmup_keys(*args, **kwargs):
        self.compile(compile_key)

_WarmupInputRows dataclass

Warmup dispatch inputs expanded in lockstep.

Source code in vllm/model_executor/warmup/jit_warmup.py 
40
41
42
43
44
@dataclass(frozen=True)
class _WarmupInputRows:
    """Warmup dispatch inputs expanded in lockstep."""

    rows: tuple[Mapping[str, WarmupValues], ...]

zip_inputs(*rows)

Group row-wise dispatch inputs that should be expanded in lockstep.

Source code in vllm/model_executor/warmup/jit_warmup.py 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
def zip_inputs(*rows: Mapping[str, WarmupValues]) -> _WarmupInputRows:
    """Group row-wise dispatch inputs that should be expanded in lockstep."""
    if not rows:
        raise ValueError("zip_inputs requires at least one dispatch input row")
    if not all(isinstance(row, Mapping) for row in rows):
        raise ValueError("zip_inputs rows must be mappings")

    first_names = frozenset(rows[0])
    if not first_names:
        raise ValueError("zip_inputs rows require at least one dispatch input name")
    if not all(isinstance(name, str) for name in first_names):
        raise ValueError("zip_inputs dispatch input names must be strings")

    input_rows: list[Mapping[str, WarmupValues]] = []
    for row in rows:
        names = frozenset(row)
        if names != first_names:
            raise ValueError("zip_inputs rows must use the same dispatch input names")
        input_rows.append(dict(row))

    return _WarmupInputRows(rows=tuple(input_rows))