vllm.model_executor.warmup.jit_warmup ¶

class JitWarmupRegistry:
    """Collect and compile JIT kernels selected during runner setup."""

    _active: ContextVar[JitWarmupRegistry | None] = ContextVar(
        "active_jit_warmup_registry",
        default=None,
    )

    def __init__(self, vllm_config: Any) -> None:
        self.vllm_config = vllm_config
        self._registrations: dict[
            VllmJitKernel[Any],
            list[tuple[tuple[Any, ...], dict[str, Any]]],
        ] = {}

    @contextmanager
    def activate(self) -> Iterator[None]:
        """Collect registrations made in this context."""
        token = self._active.set(self)
        try:
            yield
        finally:
            self._active.reset(token)

    @classmethod
    def register(
        cls,
        kernel: VllmJitKernel[Any],
        *args: Any,
        **kwargs: Any,
    ) -> None:
        """Register a kernel with the active registry, if one exists."""
        registry = cls._active.get()
        if registry is not None:
            registry._add(kernel, args, kwargs)

    def _add(
        self,
        kernel: VllmJitKernel[Any],
        args: tuple[Any, ...],
        kwargs: dict[str, Any],
    ) -> None:
        registrations = self._registrations.setdefault(kernel, [])
        if (
            not args
            and not kwargs
            and any(
                not registered_args and not registered_kwargs
                for registered_args, registered_kwargs in registrations
            )
        ):
            return
        registrations.append((args, kwargs))

    def __len__(self) -> int:
        return sum(len(registrations) for registrations in self._registrations.values())

    def warmup(self) -> None:
        """Expand registrations and compile each wrapper/key pair once."""
        from tqdm import tqdm

        from vllm.distributed import is_global_first_rank

        kernel_items: list[tuple[VllmJitKernel[Any], dict[Any, None]]] = []
        for kernel, registrations in self._registrations.items():
            compile_keys: dict[Any, None] = {}
            for args, kwargs in registrations:
                if not args and not kwargs:
                    args = (self.vllm_config,)
                for compile_key in kernel.get_warmup_keys(*args, **kwargs):
                    compile_keys[compile_key] = None
            if compile_keys:
                kernel_items.append((kernel, compile_keys))

        if not kernel_items:
            return

        total_keys = sum(len(compile_keys) for _, compile_keys in kernel_items)
        with tqdm(
            kernel_items,
            desc=f"JIT kernel warmup ({total_keys} compile keys)",
            disable=not is_global_first_rank(),
            dynamic_ncols=True,
            unit="kernel",
        ) as progress:
            for kernel, compile_keys in progress:
                progress.set_postfix_str(
                    f"{kernel.__class__.__name__} ({len(compile_keys)} keys)",
                    refresh=False,
                )
                for compile_key in compile_keys:
                    kernel.compile(compile_key)

class VllmJitKernel(Generic[CompileKeyT], ABC):
    """Kernel wrapper that owns dispatch, warmup keys, and compilation."""

    CompileKey: type[CompileKeyT]

    def __init__(self) -> None:
        self._dispatch_trace = _trace_compile_key_dispatch(self.dispatch)
        self._compiled_cache: dict[Any, Any] = {}

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

    def _get_or_compile(
        self,
        compile_key: CompileKeyT,
        *,
        runtime_context: Mapping[str, Any] | None = None,
    ) -> Any:
        """Return a cached executor, compiling it on a monitored cache miss."""
        if compile_key not in self._compiled_cache:
            self.compile(compile_key)

        try:
            return self._compiled_cache[compile_key]
        except KeyError as exc:
            details = [f"compile_key={compile_key!r}"]
            if runtime_context:
                details.append(f"runtime_context={dict(runtime_context)!r}")
            raise RuntimeError(
                f"{type(self).__name__}.compile(...) did not cache its JIT "
                f"executor ({', '.join(details)})"
            ) from exc

    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
            )
            # Unmatched **kwargs fields also belong to the expansion space.
            available_names = set(kwargs).union(
                *(group.rows[0] for group in input_groups)
            )
            input_names = compile_key_dispatch_trace.input_names_for(available_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 register_warmup(self, *args: Any, **kwargs: Any) -> None:
        """Register this kernel with the active runner's warmup registry."""
        JitWarmupRegistry.register(self, *args, **kwargs)

    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], ...]