1119 lines
39 KiB
Python
1119 lines
39 KiB
Python
# 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 builtins
|
|
import inspect
|
|
import itertools
|
|
import operator
|
|
import textwrap
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import Callable, Iterable, Iterator, Mapping
|
|
from contextlib import contextmanager
|
|
from contextvars import ContextVar
|
|
from dataclasses import dataclass, fields
|
|
from functools import wraps
|
|
from typing import Any, Generic, TypeVar, cast
|
|
|
|
__all__ = [
|
|
"JitWarmupRegistry",
|
|
"VllmJitKernel",
|
|
"WarmupChoices",
|
|
"WarmupIntRange",
|
|
"get_ast_full_name",
|
|
"get_function_source_node",
|
|
"kernel_launcher",
|
|
"zip_inputs",
|
|
]
|
|
|
|
|
|
CompileKeyT = TypeVar("CompileKeyT")
|
|
|
|
|
|
def kernel_launcher(call_fn: Callable[..., Any]) -> Callable[..., Any]:
|
|
"""Delegate a declarative launch specification to its backend owner."""
|
|
signature = inspect.signature(call_fn)
|
|
|
|
@wraps(call_fn)
|
|
def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
|
|
launch_spec = call_fn(self, *args, **kwargs)
|
|
if not self.bind_launch_inputs:
|
|
return self.launch(launch_spec, {})
|
|
|
|
bound = signature.bind(self, *args, **kwargs)
|
|
bound.apply_defaults()
|
|
inputs = {
|
|
name: value for name, value in bound.arguments.items() if name != "self"
|
|
}
|
|
return self.launch(launch_spec, inputs)
|
|
|
|
return wrapper
|
|
|
|
|
|
@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 WarmupChoices:
|
|
"""Expand an explicit finite set of values in a traced warmup method."""
|
|
|
|
values: tuple[Any, ...]
|
|
|
|
def __init__(self, *values: Any) -> None:
|
|
object.__setattr__(self, "values", values)
|
|
|
|
|
|
def _when(condition: bool) -> None:
|
|
"""Filter symbolic cases in an AST-traced warmup-cases method."""
|
|
raise RuntimeError("_when() is only valid in an AST-traced warmup method")
|
|
|
|
|
|
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, WarmupChoices):
|
|
return values.values
|
|
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]
|
|
# Named parameters are excluded from direct **kwargs forwarding.
|
|
named_parameters: frozenset[str] | None
|
|
defaults: Mapping[str, Any]
|
|
|
|
def input_names_for(self, available_names: set[str]) -> frozenset[str]:
|
|
if self.named_parameters is None:
|
|
return self.input_names
|
|
return self.input_names | (available_names - self.named_parameters)
|
|
|
|
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
|
|
)
|
|
named_parameters = self.named_parameters
|
|
# Materialize direct fields before evaluating named AST expressions.
|
|
fields: dict[str, Any] = {}
|
|
if named_parameters is not None:
|
|
fields = {
|
|
name: value
|
|
for name, value in kwargs.items()
|
|
if name not in named_parameters
|
|
}
|
|
for field, expr in self.field_exprs:
|
|
if field in fields:
|
|
raise TypeError(f"CompileKey field '{field}' is specified twice")
|
|
fields[field] = _eval_dispatch_expr(expr, dispatch_values, self.globals)
|
|
return compile_key_type(**fields)
|
|
|
|
|
|
@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,
|
|
ast.In: lambda left, right: left in right,
|
|
ast.NotIn: lambda left, right: left not in right,
|
|
ast.Is: operator.is_,
|
|
ast.IsNot: operator.is_not,
|
|
}
|
|
|
|
|
|
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, "
|
|
"subscriptions, 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]
|
|
if hasattr(builtins, node.id):
|
|
return getattr(builtins, 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_Lambda(self, node: ast.Lambda) -> Callable[..., Any]:
|
|
arguments = node.args
|
|
if (
|
|
arguments.posonlyargs
|
|
or arguments.vararg is not None
|
|
or arguments.kwonlyargs
|
|
or arguments.kwarg is not None
|
|
or arguments.defaults
|
|
):
|
|
raise _dispatch_expr_error(node, "Traced lambdas require positional args")
|
|
names = tuple(argument.arg for argument in arguments.args)
|
|
|
|
def evaluate(*values: Any) -> Any:
|
|
if len(values) != len(names):
|
|
raise TypeError(f"Expected {len(names)} lambda arguments")
|
|
return type(self)(
|
|
dict(self.values) | dict(zip(names, values)), self.globals
|
|
).eval(node.body)
|
|
|
|
return evaluate
|
|
|
|
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: list[Any] = []
|
|
for arg in node.args:
|
|
if isinstance(arg, ast.Starred):
|
|
args.extend(self.visit(arg.value))
|
|
else:
|
|
args.append(self.visit(arg))
|
|
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 visit_Subscript(self, node: ast.Subscript) -> Any:
|
|
return self.visit(node.value)[self.visit(node.slice)]
|
|
|
|
|
|
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 | ast.Lambda:
|
|
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:
|
|
return function_defs[0]
|
|
|
|
lambdas = [node for node in ast.walk(tree) if isinstance(node, ast.Lambda)]
|
|
if len(lambdas) == 1:
|
|
return lambdas[0]
|
|
|
|
name = getattr(source_fn, "__name__", type(source_fn).__name__)
|
|
raise ValueError(
|
|
f"Expected one function or lambda in {name}, found "
|
|
f"{len(function_defs)} functions and {len(lambdas)} lambdas"
|
|
)
|
|
|
|
|
|
def _eval_dispatch_expr(
|
|
node: ast.AST,
|
|
kwargs: Mapping[str, Any],
|
|
globals_: Mapping[str, Any],
|
|
) -> Any:
|
|
return _DispatchExprEvaluator(kwargs, globals_).eval(node)
|
|
|
|
|
|
def _validate_dispatch_expr(node: ast.AST) -> None:
|
|
"""Enforce the dispatch AST subset before compiling traced code."""
|
|
allowed_nodes = (
|
|
ast.Name,
|
|
ast.Constant,
|
|
ast.Lambda,
|
|
ast.arguments,
|
|
ast.arg,
|
|
ast.IfExp,
|
|
ast.Tuple,
|
|
ast.List,
|
|
ast.BoolOp,
|
|
ast.And,
|
|
ast.Or,
|
|
ast.Compare,
|
|
*tuple(_CMP_OPS),
|
|
ast.UnaryOp,
|
|
ast.Not,
|
|
ast.USub,
|
|
ast.BinOp,
|
|
*tuple(_BIN_OPS),
|
|
ast.Call,
|
|
ast.keyword,
|
|
ast.Starred,
|
|
ast.Attribute,
|
|
ast.Subscript,
|
|
ast.Load,
|
|
)
|
|
for child in ast.walk(node):
|
|
if not isinstance(child, allowed_nodes):
|
|
raise _dispatch_expr_error(child, "Unsupported dispatch expression")
|
|
if isinstance(child, ast.Call) and any(
|
|
keyword.arg is None for keyword in child.keywords
|
|
):
|
|
raise _dispatch_expr_error(
|
|
child, "Dispatch helper calls cannot use **kwargs"
|
|
)
|
|
|
|
|
|
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 _named_assignment(
|
|
statement: ast.stmt,
|
|
subject: str,
|
|
) -> tuple[str, ast.AST] | None:
|
|
match statement:
|
|
case ast.Assign(targets=[ast.Name(id=name)], value=value):
|
|
return name, value
|
|
case ast.AnnAssign(target=ast.Name(id=name), value=value) if value is not None:
|
|
return name, value
|
|
case ast.Assign() | ast.AnnAssign():
|
|
raise _dispatch_expr_error(
|
|
statement, f"{subject} assignments require one name"
|
|
)
|
|
case _:
|
|
return None
|
|
|
|
|
|
def _collect_expression_body(
|
|
fn: Callable[..., Any],
|
|
function_def: ast.FunctionDef | ast.Lambda,
|
|
) -> tuple[list[tuple[str, ast.AST]], ast.AST]:
|
|
if isinstance(function_def, ast.Lambda):
|
|
return [], function_def.body
|
|
|
|
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 (
|
|
assignment := _named_assignment(statement, "Traced warmup helper")
|
|
) is not None:
|
|
local_exprs.append(assignment)
|
|
continue
|
|
|
|
if isinstance(statement, ast.Return):
|
|
if statement.value is None:
|
|
raise _dispatch_expr_error(
|
|
statement, "Traced warmup helper must return an expression"
|
|
)
|
|
return local_exprs, statement.value
|
|
|
|
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)
|
|
if isinstance(function_def, ast.Lambda):
|
|
raise _dispatch_expr_error(
|
|
function_def, "Dispatch must be a function definition"
|
|
)
|
|
|
|
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)
|
|
# Fields captured by dispatch **kwargs are forwarded, not AST-evaluated.
|
|
signature = inspect.signature(fn)
|
|
variadic_keyword = next(
|
|
(
|
|
name
|
|
for name, parameter in signature.parameters.items()
|
|
if parameter.kind is inspect.Parameter.VAR_KEYWORD
|
|
),
|
|
None,
|
|
)
|
|
candidate_names.discard(variadic_keyword)
|
|
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))
|
|
named_parameters: frozenset[str] | None = None
|
|
for keyword in return_expr.keywords:
|
|
# CompileKey may unpack only that **kwargs parameter, once.
|
|
if keyword.arg is None:
|
|
if (
|
|
named_parameters is not None
|
|
or variadic_keyword is None
|
|
or not isinstance(keyword.value, ast.Name)
|
|
or keyword.value.id != variadic_keyword
|
|
):
|
|
raise ValueError(
|
|
f"{fn.__name__} may unpack only its own **kwargs parameter "
|
|
"once in CompileKey"
|
|
)
|
|
named_parameters = frozenset(candidate_names)
|
|
continue
|
|
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),
|
|
named_parameters,
|
|
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]
|
|
bind_launch_inputs = True
|
|
|
|
def __init__(self) -> None:
|
|
dispatch = type(self).dispatch
|
|
self._dispatch_trace = (
|
|
None
|
|
if dispatch is VllmJitKernel.dispatch
|
|
else _trace_compile_key_dispatch(self.dispatch)
|
|
)
|
|
self._compiled_cache: dict[Any, Any] = {}
|
|
|
|
def compile_key(self, kwargs: Mapping[str, Any]) -> CompileKeyT:
|
|
if self._dispatch_trace is None:
|
|
raise TypeError(f"{type(self).__name__} does not define dispatch()")
|
|
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
|
|
)
|
|
predicate_only_names: frozenset[str] = frozenset()
|
|
if predicate_trace is not None:
|
|
compile_key_fields = frozenset(
|
|
field.name for field in fields(cast(Any, self.CompileKey))
|
|
)
|
|
predicate_only_names = (
|
|
predicate_trace.input_names
|
|
- compile_key_dispatch_trace.input_names
|
|
- compile_key_fields
|
|
)
|
|
# 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,
|
|
{
|
|
name: value
|
|
for name, value in dispatch_values.items()
|
|
if name not in predicate_only_names
|
|
},
|
|
)
|
|
compile_keys[compile_key] = None
|
|
return list(compile_keys)
|
|
|
|
return traced
|
|
|
|
def _expand_warmup_cases(
|
|
self,
|
|
cases_fn: Callable[..., Any],
|
|
*args: Any,
|
|
_value_expander: Callable[[WarmupValues], tuple[Any, ...]] = (
|
|
_expand_warmup_values
|
|
),
|
|
**kwargs: Any,
|
|
) -> Iterator[Mapping[str, Any]]:
|
|
"""Expand symbolic domains declared inside a warmup-cases method."""
|
|
function_def = get_function_source_node(cases_fn)
|
|
if isinstance(function_def, ast.Lambda):
|
|
raise _dispatch_expr_error(
|
|
function_def, "Warmup cases must be a function definition"
|
|
)
|
|
signature = inspect.signature(cases_fn)
|
|
bound = signature.bind(*args, **kwargs)
|
|
bound.apply_defaults()
|
|
static_values = dict(bound.arguments)
|
|
bound_self = getattr(cases_fn, "__self__", None)
|
|
if bound_self is not None:
|
|
static_values["self"] = bound_self
|
|
globals_ = getattr(cases_fn, "__func__", cases_fn).__globals__
|
|
|
|
domains: list[tuple[str, tuple[Any, ...]]] = []
|
|
local_exprs: list[tuple[str, ast.AST]] = []
|
|
predicates: list[ast.AST] = []
|
|
return_expr: ast.AST | None = None
|
|
for statement in function_def.body:
|
|
if (
|
|
isinstance(statement, ast.Expr)
|
|
and isinstance(statement.value, ast.Constant)
|
|
and isinstance(statement.value.value, str)
|
|
):
|
|
continue
|
|
assignment = _named_assignment(statement, "Warmup case")
|
|
if assignment is not None:
|
|
name, value_expr = assignment
|
|
call_name = (
|
|
(get_ast_full_name(value_expr.func) or "").split(".")[-1]
|
|
if isinstance(value_expr, ast.Call)
|
|
else ""
|
|
)
|
|
if call_name in {"WarmupIntRange", "WarmupChoices"}:
|
|
value = _eval_dispatch_expr(value_expr, static_values, globals_)
|
|
domains.append((name, _value_expander(value)))
|
|
else:
|
|
local_exprs.append((name, value_expr))
|
|
continue
|
|
if (
|
|
isinstance(statement, ast.Expr)
|
|
and isinstance(statement.value, ast.Call)
|
|
and (get_ast_full_name(statement.value.func) or "").endswith("_when")
|
|
):
|
|
if len(statement.value.args) != 1 or statement.value.keywords:
|
|
raise _dispatch_expr_error(
|
|
statement, "_when requires exactly one positional expression"
|
|
)
|
|
predicates.append(statement.value.args[0])
|
|
continue
|
|
if isinstance(statement, ast.Return):
|
|
if return_expr is not None or statement.value is None:
|
|
raise _dispatch_expr_error(
|
|
statement, "Warmup cases require one final return expression"
|
|
)
|
|
return_expr = statement.value
|
|
continue
|
|
raise _dispatch_expr_error(
|
|
statement,
|
|
"AST-traced warmup cases support assignments, _when, and return",
|
|
)
|
|
|
|
if return_expr is None or not isinstance(return_expr, ast.Call):
|
|
raise ValueError("AST-traced warmup cases must return one case(...) call")
|
|
if any(keyword.arg is None for keyword in return_expr.keywords):
|
|
raise _dispatch_expr_error(
|
|
return_expr, "Warmup case expressions do not support **kwargs"
|
|
)
|
|
|
|
class _InlineDomainRewriter(ast.NodeTransformer):
|
|
def visit_Call(self, node: ast.Call) -> ast.AST:
|
|
call_name = (get_ast_full_name(node.func) or "").split(".")[-1]
|
|
if call_name not in {"WarmupIntRange", "WarmupChoices"}:
|
|
return self.generic_visit(node)
|
|
name = f"__warmup_domain_{len(domains)}"
|
|
domain = _eval_dispatch_expr(node, static_values, globals_)
|
|
domains.append((name, _value_expander(domain)))
|
|
return ast.copy_location(ast.Name(id=name, ctx=ast.Load()), node)
|
|
|
|
return_expr = cast(ast.Call, _InlineDomainRewriter().visit(return_expr))
|
|
|
|
domain_names = tuple(name for name, _ in domains)
|
|
dynamic_names = set(domain_names)
|
|
dynamic_local_exprs: list[tuple[str, ast.AST]] = []
|
|
for name, expr in local_exprs:
|
|
referenced_names = {
|
|
node.id for node in ast.walk(expr) if isinstance(node, ast.Name)
|
|
}
|
|
if referenced_names & dynamic_names:
|
|
dynamic_local_exprs.append((name, expr))
|
|
dynamic_names.add(name)
|
|
else:
|
|
static_values[name] = _eval_dispatch_expr(expr, static_values, globals_)
|
|
|
|
for _, expr in dynamic_local_exprs:
|
|
_validate_dispatch_expr(expr)
|
|
for predicate in predicates:
|
|
_validate_dispatch_expr(predicate)
|
|
_validate_dispatch_expr(return_expr)
|
|
|
|
case_body: list[ast.stmt] = [
|
|
ast.Assign(
|
|
targets=[ast.Name(id=name, ctx=ast.Store())],
|
|
value=cast(ast.expr, expr),
|
|
)
|
|
for name, expr in dynamic_local_exprs
|
|
]
|
|
if predicates:
|
|
predicate = (
|
|
predicates[0]
|
|
if len(predicates) == 1
|
|
else ast.BoolOp(
|
|
op=ast.And(),
|
|
values=[cast(ast.expr, value) for value in predicates],
|
|
)
|
|
)
|
|
case_body.append(
|
|
ast.If(
|
|
test=ast.UnaryOp(op=ast.Not(), operand=cast(ast.expr, predicate)),
|
|
body=[ast.Return(value=ast.Constant(value=None))],
|
|
orelse=[],
|
|
)
|
|
)
|
|
case_body.append(ast.Return(value=cast(ast.expr, return_expr)))
|
|
# FunctionDef fields vary across the supported Python versions, so no
|
|
# single constructor overload matches every mypy target.
|
|
case_function = ast.FunctionDef( # type: ignore[call-overload]
|
|
name="__vllm_warmup_case",
|
|
args=ast.arguments(
|
|
posonlyargs=[],
|
|
args=[ast.arg(arg=name) for name in domain_names],
|
|
kwonlyargs=[],
|
|
kw_defaults=[],
|
|
defaults=[],
|
|
),
|
|
body=case_body,
|
|
decorator_list=[],
|
|
returns=None,
|
|
type_comment=None,
|
|
)
|
|
case_globals = dict(globals_)
|
|
case_globals.update(static_values)
|
|
exec(
|
|
compile(
|
|
ast.fix_missing_locations(
|
|
ast.Module(body=[case_function], type_ignores=[])
|
|
),
|
|
"<jit-warmup>",
|
|
"exec",
|
|
),
|
|
case_globals,
|
|
)
|
|
evaluate_case = cast(Callable[..., Any], case_globals[case_function.name])
|
|
|
|
domain_values = tuple(values for _, values in domains)
|
|
for values in itertools.product(*domain_values):
|
|
case = evaluate_case(*values)
|
|
if case is None:
|
|
continue
|
|
if not isinstance(case, Mapping):
|
|
raise TypeError("AST-traced warmup cases must return a mapping")
|
|
if any(
|
|
isinstance(value, WarmupChoices | WarmupIntRange)
|
|
for value in case.values()
|
|
):
|
|
raise TypeError("Warmup domains must be direct expressions")
|
|
yield case
|
|
|
|
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 compile_many(self, compile_keys: Iterable[CompileKeyT]) -> None:
|
|
"""Compile a batch of warmup keys, allowing backend-specific scheduling."""
|
|
for compile_key in compile_keys:
|
|
self.compile(compile_key)
|
|
|
|
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."""
|
|
self.compile_many(self.get_warmup_keys(*args, **kwargs))
|
|
|
|
|
|
def _same_value(left: Any, right: Any) -> bool:
|
|
"""True if two registration values are the same object or compare equal.
|
|
|
|
Identity is checked first so shared singletons (e.g. ``vllm_config``, torch
|
|
dtypes) short-circuit before any potentially deep or non-boolean ``__eq__``.
|
|
"""
|
|
if left is right:
|
|
return True
|
|
try:
|
|
return bool(left == right)
|
|
except (TypeError, ValueError, RuntimeError):
|
|
return False
|
|
|
|
|
|
def _same_registration(
|
|
left: tuple[tuple[Any, ...], dict[str, Any]],
|
|
right: tuple[tuple[Any, ...], dict[str, Any]],
|
|
) -> bool:
|
|
"""True if two ``(args, kwargs)`` registrations are equivalent."""
|
|
left_args, left_kwargs = left
|
|
right_args, right_kwargs = right
|
|
if len(left_args) != len(right_args) or left_kwargs.keys() != right_kwargs.keys():
|
|
return False
|
|
return all(
|
|
_same_value(x, y) for x, y in zip(left_args, right_args, strict=True)
|
|
) and all(_same_value(left_kwargs[k], right_kwargs[k]) for k in left_kwargs)
|
|
|
|
|
|
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]]],
|
|
] = {}
|
|
|
|
@classmethod
|
|
def capture(cls, init_fn: Callable[..., None]) -> Callable[..., None]:
|
|
"""Collect warmup registrations made while the decorated callable runs."""
|
|
|
|
@wraps(init_fn)
|
|
def wrapped(instance: Any, vllm_config: Any, *args: Any, **kwargs: Any) -> None:
|
|
registry = cls(vllm_config)
|
|
instance.jit_warmup_registry = registry
|
|
with registry.activate():
|
|
init_fn(instance, vllm_config, *args, **kwargs)
|
|
|
|
return wrapped
|
|
|
|
@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, [])
|
|
# Every layer of a deep model appends an identical (args, kwargs) here
|
|
# (e.g. a 61-layer DSA model registers the same pack (dtype, pad_value)
|
|
# 61 times); each expands to the same compile keys, so tracing
|
|
# get_warmup_keys once per distinct registration is sufficient. Dedup on
|
|
# identity-or-equality -- identity short-circuits shared singletons like
|
|
# vllm_config before any deep __eq__.
|
|
if any(
|
|
_same_registration(registered, (args, kwargs))
|
|
for registered 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
|
|
and "vllm_config"
|
|
in inspect.signature(kernel.get_warmup_keys).parameters
|
|
):
|
|
kwargs = {"vllm_config": 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,
|
|
)
|
|
kernel.compile_many(compile_keys)
|