vLLM/vllm/model_executor/warmup/jit_warmup.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Shared interfaces and tracing helpers for explicit JIT warmup keys."""

from __future__ import annotations

import ast
import inspect
import itertools
import operator
import textwrap
from abc import ABC, abstractmethod
from collections.abc import Callable, Iterable, Mapping
from dataclasses import dataclass
from typing import Any, Generic, TypeVar

__all__ = [
    "VllmJitKernel",
    "WarmupIntRange",
    "get_ast_full_name",
    "get_function_source_node",
    "zip_inputs",
]


CompileKeyT = TypeVar("CompileKeyT")


@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


WarmupValues = Any
CompileKeyDispatchFn = Callable[..., CompileKeyT]
WarmupPredicateFn = Callable[..., bool]
_LocalExprs = tuple[tuple[str, ast.AST], ...]


def _eval_local_exprs(
    local_exprs: _LocalExprs,
    values: Mapping[str, Any],
    globals_: Mapping[str, Any],
) -> dict[str, Any]:
    evaluated_values = dict(values)
    for name, expr in local_exprs:
        evaluated_values[name] = _eval_dispatch_expr(expr, evaluated_values, globals_)
    return evaluated_values


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

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


def _expand_warmup_values(values: WarmupValues) -> tuple[Any, ...]:
    if isinstance(values, WarmupIntRange):
        if values.advance is None:
            return tuple(range(values.start, values.stop, values.step))
        if values.step != 1:
            raise ValueError("WarmupIntRange cannot set both step and advance")

        expanded: list[int] = []
        value = values.start
        while value < values.stop:
            expanded.append(value)
            next_value = values.advance(value)
            if next_value <= value:
                raise ValueError("WarmupIntRange.advance must return a greater value")
            value = next_value
        return tuple(expanded)
    if isinstance(values, (list, tuple)):
        return tuple(values)
    return (values,)


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))


def _expand_warmup_input_rows(
    rows: tuple[Mapping[str, WarmupValues], ...],
    input_names: frozenset[str],
) -> tuple[dict[str, Any], ...]:
    active_names = frozenset(name for name in rows[0] if name in input_names)
    if not active_names:
        return ({},)

    return tuple(
        {name: value for name, value in row.items() if name in active_names}
        for row in rows
    )


def _merge_warmup_kwargs(parts: Iterable[Mapping[str, Any]]) -> dict[str, Any]:
    merged: dict[str, Any] = {}
    for part in parts:
        for name, value in part.items():
            if name in merged:
                raise ValueError(
                    f"Warmup dispatch input '{name}' is specified more than once"
                )
            merged[name] = value
    return merged


@dataclass(frozen=True)
class _CompileKeyDispatchTrace:
    local_exprs: _LocalExprs
    field_exprs: tuple[tuple[str, ast.AST], ...]
    globals: Mapping[str, Any]
    input_names: frozenset[str]
    defaults: Mapping[str, Any]

    def compile_key(
        self,
        compile_key_type: type[CompileKeyT],
        kwargs: Mapping[str, Any],
    ) -> CompileKeyT:
        dispatch_values = _eval_local_exprs(
            self.local_exprs, {**self.defaults, **kwargs}, self.globals
        )
        return compile_key_type(
            **{
                field: _eval_dispatch_expr(expr, dispatch_values, self.globals)
                for field, expr in self.field_exprs
            }
        )


@dataclass(frozen=True)
class _WarmupPredicateTrace:
    local_exprs: _LocalExprs
    return_expr: ast.AST
    globals: Mapping[str, Any]
    input_names: frozenset[str]
    defaults: Mapping[str, Any]

    def matches(self, kwargs: Mapping[str, Any]) -> bool:
        dispatch_values = _eval_local_exprs(
            self.local_exprs, {**self.defaults, **kwargs}, self.globals
        )
        return bool(
            _eval_dispatch_expr(self.return_expr, dispatch_values, self.globals)
        )


_BIN_OPS: dict[type[ast.operator], Callable[[Any, Any], Any]] = {
    ast.Add: operator.add,
    ast.Sub: operator.sub,
    ast.Mult: operator.mul,
    ast.FloorDiv: operator.floordiv,
    ast.Mod: operator.mod,
    ast.Pow: operator.pow,
}
_CMP_OPS: dict[type[ast.cmpop], Callable[[Any, Any], bool]] = {
    ast.Eq: operator.eq,
    ast.NotEq: operator.ne,
    ast.Lt: operator.lt,
    ast.LtE: operator.le,
    ast.Gt: operator.gt,
    ast.GtE: operator.ge,
}


def _dispatch_expr_source(node: ast.AST) -> str:
    try:
        return ast.unparse(node)
    except Exception:
        return ast.dump(node)


def _dispatch_expr_error(node: ast.AST, reason: str) -> ValueError:
    return ValueError(
        f"{reason}: {_dispatch_expr_source(node)}. "
        "Supported dispatch expressions are names, constants, attributes, "
        "tuple/list literals, conditional expressions, comparisons, boolean "
        "operators, unary not/minus, arithmetic, and calls without **kwargs."
    )


class _DispatchExprEvaluator(ast.NodeVisitor):
    def __init__(
        self,
        values: Mapping[str, Any],
        globals_: Mapping[str, Any],
    ) -> None:
        self.values = values
        self.globals = globals_

    def eval(self, node: ast.AST) -> Any:
        return self.visit(node)

    def generic_visit(self, node: ast.AST) -> Any:
        raise _dispatch_expr_error(node, "Unsupported dispatch expression")

    def visit_Name(self, node: ast.Name) -> Any:
        if node.id in self.values:
            return self.values[node.id]
        if node.id in self.globals:
            return self.globals[node.id]
        raise _dispatch_expr_error(node, f"Unknown dispatch name '{node.id}'")

    def visit_Constant(self, node: ast.Constant) -> Any:
        return node.value

    def visit_IfExp(self, node: ast.IfExp) -> Any:
        return self.visit(node.body if self.visit(node.test) else node.orelse)

    def visit_Tuple(self, node: ast.Tuple) -> tuple[Any, ...]:
        return tuple(self.visit(elt) for elt in node.elts)

    def visit_List(self, node: ast.List) -> list[Any]:
        return [self.visit(elt) for elt in node.elts]

    def visit_BoolOp(self, node: ast.BoolOp) -> Any:
        if isinstance(node.op, ast.And):
            result = None
            for value in node.values:
                result = self.visit(value)
                if not result:
                    return result
            return result
        if isinstance(node.op, ast.Or):
            result = None
            for value in node.values:
                result = self.visit(value)
                if result:
                    return result
            return result
        raise _dispatch_expr_error(node, "Unsupported dispatch boolean operator")

    def visit_Compare(self, node: ast.Compare) -> bool:
        left = self.visit(node.left)
        for op_node, comparator in zip(node.ops, node.comparators):
            right = self.visit(comparator)
            op = _CMP_OPS.get(type(op_node))
            if op is None:
                raise _dispatch_expr_error(
                    node, "Unsupported dispatch comparison operator"
                )
            if not op(left, right):
                return False
            left = right
        return True

    def visit_UnaryOp(self, node: ast.UnaryOp) -> Any:
        operand = self.visit(node.operand)
        if isinstance(node.op, ast.Not):
            return not operand
        if isinstance(node.op, ast.USub):
            return -operand
        raise _dispatch_expr_error(node, "Unsupported dispatch unary operator")

    def visit_BinOp(self, node: ast.BinOp) -> Any:
        op = _BIN_OPS.get(type(node.op))
        if op is None:
            raise _dispatch_expr_error(node, "Unsupported dispatch binary operator")
        return op(self.visit(node.left), self.visit(node.right))

    def visit_Call(self, node: ast.Call) -> Any:
        args = [self.visit(arg) for arg in node.args]
        fn = self.visit(node.func)
        call_kwargs: dict[str, Any] = {}
        for keyword in node.keywords:
            if keyword.arg is None:
                raise _dispatch_expr_error(
                    node, "Dispatch helper calls cannot use **kwargs"
                )
            call_kwargs[keyword.arg] = self.visit(keyword.value)
        return fn(*args, **call_kwargs)

    def visit_Attribute(self, node: ast.Attribute) -> Any:
        return getattr(self.visit(node.value), node.attr)


def get_ast_full_name(node: ast.AST) -> str | None:
    if isinstance(node, ast.Name):
        return node.id
    if isinstance(node, ast.Attribute):
        parent = get_ast_full_name(node.value)
        if parent is not None:
            return f"{parent}.{node.attr}"
    return None


def get_function_source_node(fn: Callable[..., Any]) -> ast.FunctionDef:
    source_fn = getattr(fn, "fn", fn)
    source = textwrap.dedent(inspect.getsource(source_fn))
    tree = ast.parse(source)
    function_defs = [node for node in tree.body if isinstance(node, ast.FunctionDef)]
    if len(function_defs) != 1:
        name = getattr(source_fn, "__name__", type(source_fn).__name__)
        raise ValueError(f"Expected one function in {name}, found {len(function_defs)}")
    return function_defs[0]


def _eval_dispatch_expr(
    node: ast.AST,
    kwargs: Mapping[str, Any],
    globals_: Mapping[str, Any],
) -> Any:
    return _DispatchExprEvaluator(kwargs, globals_).eval(node)


def _collect_input_names(
    node: ast.AST,
    candidate_names: set[str],
    local_names: set[str] | None = None,
) -> set[str]:
    if local_names is None:
        local_names = set()
    return {
        child.id
        for child in ast.walk(node)
        if (
            isinstance(child, ast.Name)
            and child.id in candidate_names
            and child.id not in local_names
        )
    }


def _collect_expression_body(
    fn: Callable[..., Any],
    function_def: ast.FunctionDef,
) -> tuple[list[tuple[str, ast.AST]], ast.AST]:
    local_exprs: list[tuple[str, ast.AST]] = []
    for statement in function_def.body:
        if (
            isinstance(statement, ast.Expr)
            and isinstance(statement.value, ast.Constant)
            and isinstance(statement.value.value, str)
        ):
            continue

        if isinstance(statement, ast.Assign):
            if len(statement.targets) != 1 or not isinstance(
                statement.targets[0], ast.Name
            ):
                raise _dispatch_expr_error(
                    statement,
                    "Traced warmup helper assignments must target one local name",
                )
            local_exprs.append((statement.targets[0].id, statement.value))
            continue

        if isinstance(statement, ast.AnnAssign):
            if statement.value is None:
                raise _dispatch_expr_error(
                    statement, "Traced warmup helper annotations must assign a value"
                )
            if not isinstance(statement.target, ast.Name):
                raise _dispatch_expr_error(
                    statement,
                    "Traced warmup helper assignments must target one local name",
                )
            local_exprs.append((statement.target.id, statement.value))
            continue

        if isinstance(statement, ast.Return) and statement.value is not None:
            return local_exprs, statement.value

        if isinstance(statement, ast.Return):
            raise _dispatch_expr_error(
                statement, "Traced warmup helper must return an expression"
            )

        raise _dispatch_expr_error(
            statement,
            "Traced warmup helper may only contain local assignments before return",
        )

    raise ValueError(f"Expected {fn.__name__} to return an expression")


def _function_trace_inputs(
    fn: Callable[..., Any],
) -> tuple[dict[str, Any], set[str]]:
    signature = inspect.signature(fn)
    defaults = {
        name: parameter.default
        for name, parameter in signature.parameters.items()
        if parameter.default is not inspect.Parameter.empty
    }
    bound_self = getattr(fn, "__self__", None)
    if bound_self is not None:
        defaults["self"] = bound_self
    return defaults, set(signature.parameters)


def _trace_compile_key_dispatch(
    fn: CompileKeyDispatchFn[Any],
) -> _CompileKeyDispatchTrace:
    source_fn = getattr(fn, "__func__", fn)
    globals_ = source_fn.__globals__
    function_def = get_function_source_node(fn)

    local_exprs, return_expr = _collect_expression_body(fn, function_def)
    if not isinstance(return_expr, ast.Call):
        raise _dispatch_expr_error(
            return_expr,
            "Dispatch must return one CompileKey(...) call",
        )

    field_exprs: list[tuple[str, ast.AST]] = []
    defaults, candidate_names = _function_trace_inputs(fn)
    input_names: set[str] = set()
    local_names = {name for name, _ in local_exprs}
    for _, expr in local_exprs:
        input_names.update(_collect_input_names(expr, candidate_names))
    for keyword in return_expr.keywords:
        if keyword.arg is None:
            raise ValueError(f"{fn.__name__} cannot use **kwargs in CompileKey")
        field_exprs.append((keyword.arg, keyword.value))
        input_names.update(
            _collect_input_names(keyword.value, candidate_names, local_names)
        )

    return _CompileKeyDispatchTrace(
        tuple(local_exprs),
        tuple(field_exprs),
        globals_,
        frozenset(input_names),
        defaults,
    )


def _trace_warmup_predicate(
    fn: WarmupPredicateFn,
) -> _WarmupPredicateTrace:
    source_fn = getattr(fn, "__func__", fn)
    globals_ = source_fn.__globals__
    function_def = get_function_source_node(fn)

    local_exprs, return_expr = _collect_expression_body(fn, function_def)
    defaults, candidate_names = _function_trace_inputs(fn)
    input_names: set[str] = set()
    local_names = {name for name, _ in local_exprs}
    for _, expr in local_exprs:
        input_names.update(_collect_input_names(expr, candidate_names))
    input_names.update(_collect_input_names(return_expr, candidate_names, local_names))

    return _WarmupPredicateTrace(
        tuple(local_exprs),
        return_expr,
        globals_,
        frozenset(input_names),
        defaults,
    )


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)