vllm.model_executor.warmup.jit_warmup ¶

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,
            _when: WarmupPredicateFn | None = None,
            **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"
                    )
            predicate_trace = (
                _trace_warmup_predicate(_when) if _when is not None else None
            )
            input_names = compile_key_dispatch_trace.input_names
            if predicate_trace is not None:
                input_names = input_names | predicate_trace.input_names
            expanded_input_groups = tuple(
                _expand_warmup_input_rows(group.rows, input_names)
                for group in input_groups
            )
            # Expand independent keyword inputs into cartesian-product axes.
            expanded_kwarg_axes = tuple(
                (name, _expand_warmup_values(value))
                for name, value in kwargs.items()
                if name in input_names
            )
            dispatch_value_axes = (
                *expanded_input_groups,
                *(values for _, values in expanded_kwarg_axes),
            )
            input_group_count = len(expanded_input_groups)
            kwarg_names = tuple(name for name, _ in expanded_kwarg_axes)
            compile_keys: dict[CompileKeyT, None] = {}
            for dispatch_value_set in itertools.product(*dispatch_value_axes):
                dispatch_values = _merge_warmup_kwargs(
                    (
                        *dispatch_value_set[:input_group_count],
                        dict(
                            zip(
                                kwarg_names,
                                dispatch_value_set[input_group_count:],
                            )
                        ),
                    )
                )
                if predicate_trace is not None and not predicate_trace.matches(
                    dispatch_values
                ):
                    continue
                compile_key = compile_key_dispatch_trace.compile_key(
                    self.CompileKey, dispatch_values
                )
                compile_keys[compile_key] = None
            return list(compile_keys)

        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)

@dataclass(frozen=True)
class WarmupIntRange:
    """Expand integers with range semantics or a custom monotonic progression."""

    start: int
    stop: int
    step: int = 1
    advance: Callable[[int], int] | None = None

@dataclass(frozen=True)
class _WarmupInputRows:
    """Warmup dispatch inputs expanded in lockstep."""

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