| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Shared interfaces and tracing helpers for explicit JIT warmup keys.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import ast |
| 8 | import builtins |
| 9 | import inspect |
| 10 | import itertools |
| 11 | import operator |
| 12 | import textwrap |
| 13 | from abc import ABC, abstractmethod |
| 14 | from collections.abc import Callable, Iterable, Iterator, Mapping |
| 15 | from contextlib import contextmanager |
| 16 | from contextvars import ContextVar |
| 17 | from dataclasses import dataclass, fields |
| 18 | from functools import wraps |
| 19 | from typing import Any, Generic, TypeVar, cast |
| 20 | |
| 21 | __all__ = [ |
| 22 | "JitWarmupRegistry", |
| 23 | "VllmJitKernel", |
| 24 | "WarmupChoices", |
| 25 | "WarmupIntRange", |
| 26 | "get_ast_full_name", |
| 27 | "get_function_source_node", |
| 28 | "kernel_launcher", |
| 29 | "zip_inputs", |
| 30 | ] |
| 31 | |
| 32 | |
| 33 | CompileKeyT = TypeVar("CompileKeyT") |
| 34 | |
| 35 | |
| 36 | def kernel_launcher(call_fn: Callable[..., Any]) -> Callable[..., Any]: |
| 37 | """Delegate a declarative launch specification to its backend owner.""" |
| 38 | signature = inspect.signature(call_fn) |
| 39 | |
| 40 | @wraps(call_fn) |
| 41 | def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: |
| 42 | launch_spec = call_fn(self, *args, **kwargs) |
| 43 | if not self.bind_launch_inputs: |
| 44 | return self.launch(launch_spec, {}) |
| 45 | |
| 46 | bound = signature.bind(self, *args, **kwargs) |
| 47 | bound.apply_defaults() |
| 48 | inputs = { |
| 49 | name: value for name, value in bound.arguments.items() if name != "self" |
| 50 | } |
| 51 | return self.launch(launch_spec, inputs) |
| 52 | |
| 53 | return wrapper |
| 54 | |
| 55 | |
| 56 | @dataclass(frozen=True) |
| 57 | class WarmupIntRange: |
| 58 | """Expand integers with range semantics or a custom monotonic progression.""" |
| 59 | |
| 60 | start: int |
| 61 | stop: int |
| 62 | step: int = 1 |
| 63 | advance: Callable[[int], int] | None = None |
| 64 | |
| 65 | |
| 66 | @dataclass(frozen=True) |
| 67 | class WarmupChoices: |
| 68 | """Expand an explicit finite set of values in a traced warmup method.""" |
| 69 | |
| 70 | values: tuple[Any, ...] |
| 71 | |
| 72 | def __init__(self, *values: Any) -> None: |
| 73 | object.__setattr__(self, "values", values) |
| 74 | |
| 75 | |
| 76 | def _when(condition: bool) -> None: |
| 77 | """Filter symbolic cases in an AST-traced warmup-cases method.""" |
| 78 | raise RuntimeError("_when() is only valid in an AST-traced warmup method") |
| 79 | |
| 80 | |
| 81 | WarmupValues = Any |
| 82 | CompileKeyDispatchFn = Callable[..., CompileKeyT] |
| 83 | WarmupPredicateFn = Callable[..., bool] |
| 84 | _LocalExprs = tuple[tuple[str, ast.AST], ...] |
| 85 | |
| 86 | |
| 87 | def _eval_local_exprs( |
| 88 | local_exprs: _LocalExprs, |
| 89 | values: Mapping[str, Any], |
| 90 | globals_: Mapping[str, Any], |
| 91 | ) -> dict[str, Any]: |
| 92 | evaluated_values = dict(values) |
| 93 | for name, expr in local_exprs: |
| 94 | evaluated_values[name] = _eval_dispatch_expr(expr, evaluated_values, globals_) |
| 95 | return evaluated_values |
| 96 | |
| 97 | |
| 98 | @dataclass(frozen=True) |
| 99 | class _WarmupInputRows: |
| 100 | """Warmup dispatch inputs expanded in lockstep.""" |
| 101 | |
| 102 | rows: tuple[Mapping[str, WarmupValues], ...] |
| 103 | |
| 104 | |
| 105 | def _expand_warmup_values(values: WarmupValues) -> tuple[Any, ...]: |
| 106 | if isinstance(values, WarmupChoices): |
| 107 | return values.values |
| 108 | if isinstance(values, WarmupIntRange): |
| 109 | if values.advance is None: |
| 110 | return tuple(range(values.start, values.stop, values.step)) |
| 111 | if values.step != 1: |
| 112 | raise ValueError("WarmupIntRange cannot set both step and advance") |
| 113 | |
| 114 | expanded: list[int] = [] |
| 115 | value = values.start |
| 116 | while value < values.stop: |
| 117 | expanded.append(value) |
| 118 | next_value = values.advance(value) |
| 119 | if next_value <= value: |
| 120 | raise ValueError("WarmupIntRange.advance must return a greater value") |
| 121 | value = next_value |
| 122 | return tuple(expanded) |
| 123 | if isinstance(values, (list, tuple)): |
| 124 | return tuple(values) |
| 125 | return (values,) |
| 126 | |
| 127 | |
| 128 | def zip_inputs(*rows: Mapping[str, WarmupValues]) -> _WarmupInputRows: |
| 129 | """Group row-wise dispatch inputs that should be expanded in lockstep.""" |
| 130 | if not rows: |
| 131 | raise ValueError("zip_inputs requires at least one dispatch input row") |
| 132 | if not all(isinstance(row, Mapping) for row in rows): |
| 133 | raise ValueError("zip_inputs rows must be mappings") |
| 134 | |
| 135 | first_names = frozenset(rows[0]) |
| 136 | if not first_names: |
| 137 | raise ValueError("zip_inputs rows require at least one dispatch input name") |
| 138 | if not all(isinstance(name, str) for name in first_names): |
| 139 | raise ValueError("zip_inputs dispatch input names must be strings") |
| 140 | |
| 141 | input_rows: list[Mapping[str, WarmupValues]] = [] |
| 142 | for row in rows: |
| 143 | names = frozenset(row) |
| 144 | if names != first_names: |
| 145 | raise ValueError("zip_inputs rows must use the same dispatch input names") |
| 146 | input_rows.append(dict(row)) |
| 147 | |
| 148 | return _WarmupInputRows(rows=tuple(input_rows)) |
| 149 | |
| 150 | |
| 151 | def _expand_warmup_input_rows( |
| 152 | rows: tuple[Mapping[str, WarmupValues], ...], |
| 153 | input_names: frozenset[str], |
| 154 | ) -> tuple[dict[str, Any], ...]: |
| 155 | active_names = frozenset(name for name in rows[0] if name in input_names) |
| 156 | if not active_names: |
| 157 | return ({},) |
| 158 | |
| 159 | return tuple( |
| 160 | {name: value for name, value in row.items() if name in active_names} |
| 161 | for row in rows |
| 162 | ) |
| 163 | |
| 164 | |
| 165 | def _merge_warmup_kwargs(parts: Iterable[Mapping[str, Any]]) -> dict[str, Any]: |
| 166 | merged: dict[str, Any] = {} |
| 167 | for part in parts: |
| 168 | for name, value in part.items(): |
| 169 | if name in merged: |
| 170 | raise ValueError( |
| 171 | f"Warmup dispatch input '{name}' is specified more than once" |
| 172 | ) |
| 173 | merged[name] = value |
| 174 | return merged |
| 175 | |
| 176 | |
| 177 | @dataclass(frozen=True) |
| 178 | class _CompileKeyDispatchTrace: |
| 179 | local_exprs: _LocalExprs |
| 180 | field_exprs: tuple[tuple[str, ast.AST], ...] |
| 181 | globals: Mapping[str, Any] |
| 182 | input_names: frozenset[str] |
| 183 | # Named parameters are excluded from direct **kwargs forwarding. |
| 184 | named_parameters: frozenset[str] | None |
| 185 | defaults: Mapping[str, Any] |
| 186 | |
| 187 | def input_names_for(self, available_names: set[str]) -> frozenset[str]: |
| 188 | if self.named_parameters is None: |
| 189 | return self.input_names |
| 190 | return self.input_names | (available_names - self.named_parameters) |
| 191 | |
| 192 | def compile_key( |
| 193 | self, |
| 194 | compile_key_type: type[CompileKeyT], |
| 195 | kwargs: Mapping[str, Any], |
| 196 | ) -> CompileKeyT: |
| 197 | dispatch_values = _eval_local_exprs( |
| 198 | self.local_exprs, {**self.defaults, **kwargs}, self.globals |
| 199 | ) |
| 200 | named_parameters = self.named_parameters |
| 201 | # Materialize direct fields before evaluating named AST expressions. |
| 202 | fields: dict[str, Any] = {} |
| 203 | if named_parameters is not None: |
| 204 | fields = { |
| 205 | name: value |
| 206 | for name, value in kwargs.items() |
| 207 | if name not in named_parameters |
| 208 | } |
| 209 | for field, expr in self.field_exprs: |
| 210 | if field in fields: |
| 211 | raise TypeError(f"CompileKey field '{field}' is specified twice") |
| 212 | fields[field] = _eval_dispatch_expr(expr, dispatch_values, self.globals) |
| 213 | return compile_key_type(**fields) |
| 214 | |
| 215 | |
| 216 | @dataclass(frozen=True) |
| 217 | class _WarmupPredicateTrace: |
| 218 | local_exprs: _LocalExprs |
| 219 | return_expr: ast.AST |
| 220 | globals: Mapping[str, Any] |
| 221 | input_names: frozenset[str] |
| 222 | defaults: Mapping[str, Any] |
| 223 | |
| 224 | def matches(self, kwargs: Mapping[str, Any]) -> bool: |
| 225 | dispatch_values = _eval_local_exprs( |
| 226 | self.local_exprs, {**self.defaults, **kwargs}, self.globals |
| 227 | ) |
| 228 | return bool( |
| 229 | _eval_dispatch_expr(self.return_expr, dispatch_values, self.globals) |
| 230 | ) |
| 231 | |
| 232 | |
| 233 | _BIN_OPS: dict[type[ast.operator], Callable[[Any, Any], Any]] = { |
| 234 | ast.Add: operator.add, |
| 235 | ast.Sub: operator.sub, |
| 236 | ast.Mult: operator.mul, |
| 237 | ast.FloorDiv: operator.floordiv, |
| 238 | ast.Mod: operator.mod, |
| 239 | ast.Pow: operator.pow, |
| 240 | } |
| 241 | _CMP_OPS: dict[type[ast.cmpop], Callable[[Any, Any], bool]] = { |
| 242 | ast.Eq: operator.eq, |
| 243 | ast.NotEq: operator.ne, |
| 244 | ast.Lt: operator.lt, |
| 245 | ast.LtE: operator.le, |
| 246 | ast.Gt: operator.gt, |
| 247 | ast.GtE: operator.ge, |
| 248 | ast.In: lambda left, right: left in right, |
| 249 | ast.NotIn: lambda left, right: left not in right, |
| 250 | ast.Is: operator.is_, |
| 251 | ast.IsNot: operator.is_not, |
| 252 | } |
| 253 | |
| 254 | |
| 255 | def _dispatch_expr_source(node: ast.AST) -> str: |
| 256 | try: |
| 257 | return ast.unparse(node) |
| 258 | except Exception: |
| 259 | return ast.dump(node) |
| 260 | |
| 261 | |
| 262 | def _dispatch_expr_error(node: ast.AST, reason: str) -> ValueError: |
| 263 | return ValueError( |
| 264 | f"{reason}: {_dispatch_expr_source(node)}. " |
| 265 | "Supported dispatch expressions are names, constants, attributes, " |
| 266 | "subscriptions, tuple/list literals, conditional expressions, " |
| 267 | "comparisons, boolean operators, unary not/minus, arithmetic, and " |
| 268 | "calls without **kwargs." |
| 269 | ) |
| 270 | |
| 271 | |
| 272 | class _DispatchExprEvaluator(ast.NodeVisitor): |
| 273 | def __init__( |
| 274 | self, |
| 275 | values: Mapping[str, Any], |
| 276 | globals_: Mapping[str, Any], |
| 277 | ) -> None: |
| 278 | self.values = values |
| 279 | self.globals = globals_ |
| 280 | |
| 281 | def eval(self, node: ast.AST) -> Any: |
| 282 | return self.visit(node) |
| 283 | |
| 284 | def generic_visit(self, node: ast.AST) -> Any: |
| 285 | raise _dispatch_expr_error(node, "Unsupported dispatch expression") |
| 286 | |
| 287 | def visit_Name(self, node: ast.Name) -> Any: |
| 288 | if node.id in self.values: |
| 289 | return self.values[node.id] |
| 290 | if node.id in self.globals: |
| 291 | return self.globals[node.id] |
| 292 | if hasattr(builtins, node.id): |
| 293 | return getattr(builtins, node.id) |
| 294 | raise _dispatch_expr_error(node, f"Unknown dispatch name '{node.id}'") |
| 295 | |
| 296 | def visit_Constant(self, node: ast.Constant) -> Any: |
| 297 | return node.value |
| 298 | |
| 299 | def visit_Lambda(self, node: ast.Lambda) -> Callable[..., Any]: |
| 300 | arguments = node.args |
| 301 | if ( |
| 302 | arguments.posonlyargs |
| 303 | or arguments.vararg is not None |
| 304 | or arguments.kwonlyargs |
| 305 | or arguments.kwarg is not None |
| 306 | or arguments.defaults |
| 307 | ): |
| 308 | raise _dispatch_expr_error(node, "Traced lambdas require positional args") |
| 309 | names = tuple(argument.arg for argument in arguments.args) |
| 310 | |
| 311 | def evaluate(*values: Any) -> Any: |
| 312 | if len(values) != len(names): |
| 313 | raise TypeError(f"Expected {len(names)} lambda arguments") |
| 314 | return type(self)( |
| 315 | dict(self.values) | dict(zip(names, values)), self.globals |
| 316 | ).eval(node.body) |
| 317 | |
| 318 | return evaluate |
| 319 | |
| 320 | def visit_IfExp(self, node: ast.IfExp) -> Any: |
| 321 | return self.visit(node.body if self.visit(node.test) else node.orelse) |
| 322 | |
| 323 | def visit_Tuple(self, node: ast.Tuple) -> tuple[Any, ...]: |
| 324 | return tuple(self.visit(elt) for elt in node.elts) |
| 325 | |
| 326 | def visit_List(self, node: ast.List) -> list[Any]: |
| 327 | return [self.visit(elt) for elt in node.elts] |
| 328 | |
| 329 | def visit_BoolOp(self, node: ast.BoolOp) -> Any: |
| 330 | if isinstance(node.op, ast.And): |
| 331 | result = None |
| 332 | for value in node.values: |
| 333 | result = self.visit(value) |
| 334 | if not result: |
| 335 | return result |
| 336 | return result |
| 337 | if isinstance(node.op, ast.Or): |
| 338 | result = None |
| 339 | for value in node.values: |
| 340 | result = self.visit(value) |
| 341 | if result: |
| 342 | return result |
| 343 | return result |
| 344 | raise _dispatch_expr_error(node, "Unsupported dispatch boolean operator") |
| 345 | |
| 346 | def visit_Compare(self, node: ast.Compare) -> bool: |
| 347 | left = self.visit(node.left) |
| 348 | for op_node, comparator in zip(node.ops, node.comparators): |
| 349 | right = self.visit(comparator) |
| 350 | op = _CMP_OPS.get(type(op_node)) |
| 351 | if op is None: |
| 352 | raise _dispatch_expr_error( |
| 353 | node, "Unsupported dispatch comparison operator" |
| 354 | ) |
| 355 | if not op(left, right): |
| 356 | return False |
| 357 | left = right |
| 358 | return True |
| 359 | |
| 360 | def visit_UnaryOp(self, node: ast.UnaryOp) -> Any: |
| 361 | operand = self.visit(node.operand) |
| 362 | if isinstance(node.op, ast.Not): |
| 363 | return not operand |
| 364 | if isinstance(node.op, ast.USub): |
| 365 | return -operand |
| 366 | raise _dispatch_expr_error(node, "Unsupported dispatch unary operator") |
| 367 | |
| 368 | def visit_BinOp(self, node: ast.BinOp) -> Any: |
| 369 | op = _BIN_OPS.get(type(node.op)) |
| 370 | if op is None: |
| 371 | raise _dispatch_expr_error(node, "Unsupported dispatch binary operator") |
| 372 | return op(self.visit(node.left), self.visit(node.right)) |
| 373 | |
| 374 | def visit_Call(self, node: ast.Call) -> Any: |
| 375 | args: list[Any] = [] |
| 376 | for arg in node.args: |
| 377 | if isinstance(arg, ast.Starred): |
| 378 | args.extend(self.visit(arg.value)) |
| 379 | else: |
| 380 | args.append(self.visit(arg)) |
| 381 | fn = self.visit(node.func) |
| 382 | call_kwargs: dict[str, Any] = {} |
| 383 | for keyword in node.keywords: |
| 384 | if keyword.arg is None: |
| 385 | raise _dispatch_expr_error( |
| 386 | node, "Dispatch helper calls cannot use **kwargs" |
| 387 | ) |
| 388 | call_kwargs[keyword.arg] = self.visit(keyword.value) |
| 389 | return fn(*args, **call_kwargs) |
| 390 | |
| 391 | def visit_Attribute(self, node: ast.Attribute) -> Any: |
| 392 | return getattr(self.visit(node.value), node.attr) |
| 393 | |
| 394 | def visit_Subscript(self, node: ast.Subscript) -> Any: |
| 395 | return self.visit(node.value)[self.visit(node.slice)] |
| 396 | |
| 397 | |
| 398 | def get_ast_full_name(node: ast.AST) -> str | None: |
| 399 | if isinstance(node, ast.Name): |
| 400 | return node.id |
| 401 | if isinstance(node, ast.Attribute): |
| 402 | parent = get_ast_full_name(node.value) |
| 403 | if parent is not None: |
| 404 | return f"{parent}.{node.attr}" |
| 405 | return None |
| 406 | |
| 407 | |
| 408 | def get_function_source_node(fn: Callable[..., Any]) -> ast.FunctionDef | ast.Lambda: |
| 409 | source_fn = getattr(fn, "fn", fn) |
| 410 | source = textwrap.dedent(inspect.getsource(source_fn)) |
| 411 | tree = ast.parse(source) |
| 412 | function_defs = [node for node in tree.body if isinstance(node, ast.FunctionDef)] |
| 413 | if len(function_defs) == 1: |
| 414 | return function_defs[0] |
| 415 | |
| 416 | lambdas = [node for node in ast.walk(tree) if isinstance(node, ast.Lambda)] |
| 417 | if len(lambdas) == 1: |
| 418 | return lambdas[0] |
| 419 | |
| 420 | name = getattr(source_fn, "__name__", type(source_fn).__name__) |
| 421 | raise ValueError( |
| 422 | f"Expected one function or lambda in {name}, found " |
| 423 | f"{len(function_defs)} functions and {len(lambdas)} lambdas" |
| 424 | ) |
| 425 | |
| 426 | |
| 427 | def _eval_dispatch_expr( |
| 428 | node: ast.AST, |
| 429 | kwargs: Mapping[str, Any], |
| 430 | globals_: Mapping[str, Any], |
| 431 | ) -> Any: |
| 432 | return _DispatchExprEvaluator(kwargs, globals_).eval(node) |
| 433 | |
| 434 | |
| 435 | def _validate_dispatch_expr(node: ast.AST) -> None: |
| 436 | """Enforce the dispatch AST subset before compiling traced code.""" |
| 437 | allowed_nodes = ( |
| 438 | ast.Name, |
| 439 | ast.Constant, |
| 440 | ast.Lambda, |
| 441 | ast.arguments, |
| 442 | ast.arg, |
| 443 | ast.IfExp, |
| 444 | ast.Tuple, |
| 445 | ast.List, |
| 446 | ast.BoolOp, |
| 447 | ast.And, |
| 448 | ast.Or, |
| 449 | ast.Compare, |
| 450 | *tuple(_CMP_OPS), |
| 451 | ast.UnaryOp, |
| 452 | ast.Not, |
| 453 | ast.USub, |
| 454 | ast.BinOp, |
| 455 | *tuple(_BIN_OPS), |
| 456 | ast.Call, |
| 457 | ast.keyword, |
| 458 | ast.Starred, |
| 459 | ast.Attribute, |
| 460 | ast.Subscript, |
| 461 | ast.Load, |
| 462 | ) |
| 463 | for child in ast.walk(node): |
| 464 | if not isinstance(child, allowed_nodes): |
| 465 | raise _dispatch_expr_error(child, "Unsupported dispatch expression") |
| 466 | if isinstance(child, ast.Call) and any( |
| 467 | keyword.arg is None for keyword in child.keywords |
| 468 | ): |
| 469 | raise _dispatch_expr_error( |
| 470 | child, "Dispatch helper calls cannot use **kwargs" |
| 471 | ) |
| 472 | |
| 473 | |
| 474 | def _collect_input_names( |
| 475 | node: ast.AST, |
| 476 | candidate_names: set[str], |
| 477 | local_names: set[str] | None = None, |
| 478 | ) -> set[str]: |
| 479 | if local_names is None: |
| 480 | local_names = set() |
| 481 | return { |
| 482 | child.id |
| 483 | for child in ast.walk(node) |
| 484 | if ( |
| 485 | isinstance(child, ast.Name) |
| 486 | and child.id in candidate_names |
| 487 | and child.id not in local_names |
| 488 | ) |
| 489 | } |
| 490 | |
| 491 | |
| 492 | def _named_assignment( |
| 493 | statement: ast.stmt, |
| 494 | subject: str, |
| 495 | ) -> tuple[str, ast.AST] | None: |
| 496 | match statement: |
| 497 | case ast.Assign(targets=[ast.Name(id=name)], value=value): |
| 498 | return name, value |
| 499 | case ast.AnnAssign(target=ast.Name(id=name), value=value) if value is not None: |
| 500 | return name, value |
| 501 | case ast.Assign() | ast.AnnAssign(): |
| 502 | raise _dispatch_expr_error( |
| 503 | statement, f"{subject} assignments require one name" |
| 504 | ) |
| 505 | case _: |
| 506 | return None |
| 507 | |
| 508 | |
| 509 | def _collect_expression_body( |
| 510 | fn: Callable[..., Any], |
| 511 | function_def: ast.FunctionDef | ast.Lambda, |
| 512 | ) -> tuple[list[tuple[str, ast.AST]], ast.AST]: |
| 513 | if isinstance(function_def, ast.Lambda): |
| 514 | return [], function_def.body |
| 515 | |
| 516 | local_exprs: list[tuple[str, ast.AST]] = [] |
| 517 | for statement in function_def.body: |
| 518 | if ( |
| 519 | isinstance(statement, ast.Expr) |
| 520 | and isinstance(statement.value, ast.Constant) |
| 521 | and isinstance(statement.value.value, str) |
| 522 | ): |
| 523 | continue |
| 524 | |
| 525 | if ( |
| 526 | assignment := _named_assignment(statement, "Traced warmup helper") |
| 527 | ) is not None: |
| 528 | local_exprs.append(assignment) |
| 529 | continue |
| 530 | |
| 531 | if isinstance(statement, ast.Return): |
| 532 | if statement.value is None: |
| 533 | raise _dispatch_expr_error( |
| 534 | statement, "Traced warmup helper must return an expression" |
| 535 | ) |
| 536 | return local_exprs, statement.value |
| 537 | |
| 538 | raise _dispatch_expr_error( |
| 539 | statement, |
| 540 | "Traced warmup helper may only contain local assignments before return", |
| 541 | ) |
| 542 | |
| 543 | raise ValueError(f"Expected {fn.__name__} to return an expression") |
| 544 | |
| 545 | |
| 546 | def _function_trace_inputs( |
| 547 | fn: Callable[..., Any], |
| 548 | ) -> tuple[dict[str, Any], set[str]]: |
| 549 | signature = inspect.signature(fn) |
| 550 | defaults = { |
| 551 | name: parameter.default |
| 552 | for name, parameter in signature.parameters.items() |
| 553 | if parameter.default is not inspect.Parameter.empty |
| 554 | } |
| 555 | bound_self = getattr(fn, "__self__", None) |
| 556 | if bound_self is not None: |
| 557 | defaults["self"] = bound_self |
| 558 | return defaults, set(signature.parameters) |
| 559 | |
| 560 | |
| 561 | def _trace_compile_key_dispatch( |
| 562 | fn: CompileKeyDispatchFn[Any], |
| 563 | ) -> _CompileKeyDispatchTrace: |
| 564 | source_fn = getattr(fn, "__func__", fn) |
| 565 | globals_ = source_fn.__globals__ |
| 566 | function_def = get_function_source_node(fn) |
| 567 | if isinstance(function_def, ast.Lambda): |
| 568 | raise _dispatch_expr_error( |
| 569 | function_def, "Dispatch must be a function definition" |
| 570 | ) |
| 571 | |
| 572 | local_exprs, return_expr = _collect_expression_body(fn, function_def) |
| 573 | if not isinstance(return_expr, ast.Call): |
| 574 | raise _dispatch_expr_error( |
| 575 | return_expr, |
| 576 | "Dispatch must return one CompileKey(...) call", |
| 577 | ) |
| 578 | |
| 579 | field_exprs: list[tuple[str, ast.AST]] = [] |
| 580 | defaults, candidate_names = _function_trace_inputs(fn) |
| 581 | # Fields captured by dispatch **kwargs are forwarded, not AST-evaluated. |
| 582 | signature = inspect.signature(fn) |
| 583 | variadic_keyword = next( |
| 584 | ( |
| 585 | name |
| 586 | for name, parameter in signature.parameters.items() |
| 587 | if parameter.kind is inspect.Parameter.VAR_KEYWORD |
| 588 | ), |
| 589 | None, |
| 590 | ) |
| 591 | candidate_names.discard(variadic_keyword) |
| 592 | input_names: set[str] = set() |
| 593 | local_names = {name for name, _ in local_exprs} |
| 594 | for _, expr in local_exprs: |
| 595 | input_names.update(_collect_input_names(expr, candidate_names)) |
| 596 | named_parameters: frozenset[str] | None = None |
| 597 | for keyword in return_expr.keywords: |
| 598 | # CompileKey may unpack only that **kwargs parameter, once. |
| 599 | if keyword.arg is None: |
| 600 | if ( |
| 601 | named_parameters is not None |
| 602 | or variadic_keyword is None |
| 603 | or not isinstance(keyword.value, ast.Name) |
| 604 | or keyword.value.id != variadic_keyword |
| 605 | ): |
| 606 | raise ValueError( |
| 607 | f"{fn.__name__} may unpack only its own **kwargs parameter " |
| 608 | "once in CompileKey" |
| 609 | ) |
| 610 | named_parameters = frozenset(candidate_names) |
| 611 | continue |
| 612 | field_exprs.append((keyword.arg, keyword.value)) |
| 613 | input_names.update( |
| 614 | _collect_input_names(keyword.value, candidate_names, local_names) |
| 615 | ) |
| 616 | |
| 617 | return _CompileKeyDispatchTrace( |
| 618 | tuple(local_exprs), |
| 619 | tuple(field_exprs), |
| 620 | globals_, |
| 621 | frozenset(input_names), |
| 622 | named_parameters, |
| 623 | defaults, |
| 624 | ) |
| 625 | |
| 626 | |
| 627 | def _trace_warmup_predicate( |
| 628 | fn: WarmupPredicateFn, |
| 629 | ) -> _WarmupPredicateTrace: |
| 630 | source_fn = getattr(fn, "__func__", fn) |
| 631 | globals_ = source_fn.__globals__ |
| 632 | function_def = get_function_source_node(fn) |
| 633 | |
| 634 | local_exprs, return_expr = _collect_expression_body(fn, function_def) |
| 635 | defaults, candidate_names = _function_trace_inputs(fn) |
| 636 | input_names: set[str] = set() |
| 637 | local_names = {name for name, _ in local_exprs} |
| 638 | for _, expr in local_exprs: |
| 639 | input_names.update(_collect_input_names(expr, candidate_names)) |
| 640 | input_names.update(_collect_input_names(return_expr, candidate_names, local_names)) |
| 641 | |
| 642 | return _WarmupPredicateTrace( |
| 643 | tuple(local_exprs), |
| 644 | return_expr, |
| 645 | globals_, |
| 646 | frozenset(input_names), |
| 647 | defaults, |
| 648 | ) |
| 649 | |
| 650 | |
| 651 | class VllmJitKernel(Generic[CompileKeyT], ABC): |
| 652 | """Kernel wrapper that owns dispatch, warmup keys, and compilation.""" |
| 653 | |
| 654 | CompileKey: type[CompileKeyT] |
| 655 | bind_launch_inputs = True |
| 656 | |
| 657 | def __init__(self) -> None: |
| 658 | dispatch = type(self).dispatch |
| 659 | self._dispatch_trace = ( |
| 660 | None |
| 661 | if dispatch is VllmJitKernel.dispatch |
| 662 | else _trace_compile_key_dispatch(self.dispatch) |
| 663 | ) |
| 664 | self._compiled_cache: dict[Any, Any] = {} |
| 665 | |
| 666 | def compile_key(self, kwargs: Mapping[str, Any]) -> CompileKeyT: |
| 667 | if self._dispatch_trace is None: |
| 668 | raise TypeError(f"{type(self).__name__} does not define dispatch()") |
| 669 | return self._dispatch_trace.compile_key(self.CompileKey, kwargs) |
| 670 | |
| 671 | def _get_or_compile( |
| 672 | self, |
| 673 | compile_key: CompileKeyT, |
| 674 | *, |
| 675 | runtime_context: Mapping[str, Any] | None = None, |
| 676 | ) -> Any: |
| 677 | """Return a cached executor, compiling it on a monitored cache miss.""" |
| 678 | if compile_key not in self._compiled_cache: |
| 679 | self.compile(compile_key) |
| 680 | |
| 681 | try: |
| 682 | return self._compiled_cache[compile_key] |
| 683 | except KeyError as exc: |
| 684 | details = [f"compile_key={compile_key!r}"] |
| 685 | if runtime_context: |
| 686 | details.append(f"runtime_context={dict(runtime_context)!r}") |
| 687 | raise RuntimeError( |
| 688 | f"{type(self).__name__}.compile(...) did not cache its JIT " |
| 689 | f"executor ({', '.join(details)})" |
| 690 | ) from exc |
| 691 | |
| 692 | def _trace_dispatch( |
| 693 | self, dispatch: CompileKeyDispatchFn[CompileKeyT] |
| 694 | ) -> Callable[..., list[CompileKeyT]]: |
| 695 | compile_key_dispatch_trace = _trace_compile_key_dispatch(dispatch) |
| 696 | |
| 697 | def traced( |
| 698 | *input_groups: _WarmupInputRows, |
| 699 | _when: WarmupPredicateFn | None = None, |
| 700 | **kwargs: WarmupValues, |
| 701 | ) -> list[CompileKeyT]: |
| 702 | for group in input_groups: |
| 703 | if not isinstance(group, _WarmupInputRows): |
| 704 | raise TypeError( |
| 705 | "_trace_dispatch positional arguments must be " |
| 706 | "zip_inputs(...) groups" |
| 707 | ) |
| 708 | predicate_trace = ( |
| 709 | _trace_warmup_predicate(_when) if _when is not None else None |
| 710 | ) |
| 711 | predicate_only_names: frozenset[str] = frozenset() |
| 712 | if predicate_trace is not None: |
| 713 | compile_key_fields = frozenset( |
| 714 | field.name for field in fields(cast(Any, self.CompileKey)) |
| 715 | ) |
| 716 | predicate_only_names = ( |
| 717 | predicate_trace.input_names |
| 718 | - compile_key_dispatch_trace.input_names |
| 719 | - compile_key_fields |
| 720 | ) |
| 721 | # Unmatched **kwargs fields also belong to the expansion space. |
| 722 | available_names = set(kwargs).union( |
| 723 | *(group.rows[0] for group in input_groups) |
| 724 | ) |
| 725 | input_names = compile_key_dispatch_trace.input_names_for(available_names) |
| 726 | if predicate_trace is not None: |
| 727 | input_names = input_names | predicate_trace.input_names |
| 728 | expanded_input_groups = tuple( |
| 729 | _expand_warmup_input_rows(group.rows, input_names) |
| 730 | for group in input_groups |
| 731 | ) |
| 732 | # Expand independent keyword inputs into cartesian-product axes. |
| 733 | expanded_kwarg_axes = tuple( |
| 734 | (name, _expand_warmup_values(value)) |
| 735 | for name, value in kwargs.items() |
| 736 | if name in input_names |
| 737 | ) |
| 738 | dispatch_value_axes = ( |
| 739 | *expanded_input_groups, |
| 740 | *(values for _, values in expanded_kwarg_axes), |
| 741 | ) |
| 742 | input_group_count = len(expanded_input_groups) |
| 743 | kwarg_names = tuple(name for name, _ in expanded_kwarg_axes) |
| 744 | compile_keys: dict[CompileKeyT, None] = {} |
| 745 | for dispatch_value_set in itertools.product(*dispatch_value_axes): |
| 746 | dispatch_values = _merge_warmup_kwargs( |
| 747 | ( |
| 748 | *dispatch_value_set[:input_group_count], |
| 749 | dict( |
| 750 | zip( |
| 751 | kwarg_names, |
| 752 | dispatch_value_set[input_group_count:], |
| 753 | ) |
| 754 | ), |
| 755 | ) |
| 756 | ) |
| 757 | if predicate_trace is not None and not predicate_trace.matches( |
| 758 | dispatch_values |
| 759 | ): |
| 760 | continue |
| 761 | compile_key = compile_key_dispatch_trace.compile_key( |
| 762 | self.CompileKey, |
| 763 | { |
| 764 | name: value |
| 765 | for name, value in dispatch_values.items() |
| 766 | if name not in predicate_only_names |
| 767 | }, |
| 768 | ) |
| 769 | compile_keys[compile_key] = None |
| 770 | return list(compile_keys) |
| 771 | |
| 772 | return traced |
| 773 | |
| 774 | def _expand_warmup_cases( |
| 775 | self, |
| 776 | cases_fn: Callable[..., Any], |
| 777 | *args: Any, |
| 778 | _value_expander: Callable[[WarmupValues], tuple[Any, ...]] = ( |
| 779 | _expand_warmup_values |
| 780 | ), |
| 781 | **kwargs: Any, |
| 782 | ) -> Iterator[Mapping[str, Any]]: |
| 783 | """Expand symbolic domains declared inside a warmup-cases method.""" |
| 784 | function_def = get_function_source_node(cases_fn) |
| 785 | if isinstance(function_def, ast.Lambda): |
| 786 | raise _dispatch_expr_error( |
| 787 | function_def, "Warmup cases must be a function definition" |
| 788 | ) |
| 789 | signature = inspect.signature(cases_fn) |
| 790 | bound = signature.bind(*args, **kwargs) |
| 791 | bound.apply_defaults() |
| 792 | static_values = dict(bound.arguments) |
| 793 | bound_self = getattr(cases_fn, "__self__", None) |
| 794 | if bound_self is not None: |
| 795 | static_values["self"] = bound_self |
| 796 | globals_ = getattr(cases_fn, "__func__", cases_fn).__globals__ |
| 797 | |
| 798 | domains: list[tuple[str, tuple[Any, ...]]] = [] |
| 799 | local_exprs: list[tuple[str, ast.AST]] = [] |
| 800 | predicates: list[ast.AST] = [] |
| 801 | return_expr: ast.AST | None = None |
| 802 | for statement in function_def.body: |
| 803 | if ( |
| 804 | isinstance(statement, ast.Expr) |
| 805 | and isinstance(statement.value, ast.Constant) |
| 806 | and isinstance(statement.value.value, str) |
| 807 | ): |
| 808 | continue |
| 809 | assignment = _named_assignment(statement, "Warmup case") |
| 810 | if assignment is not None: |
| 811 | name, value_expr = assignment |
| 812 | call_name = ( |
| 813 | (get_ast_full_name(value_expr.func) or "").split(".")[-1] |
| 814 | if isinstance(value_expr, ast.Call) |
| 815 | else "" |
| 816 | ) |
| 817 | if call_name in {"WarmupIntRange", "WarmupChoices"}: |
| 818 | value = _eval_dispatch_expr(value_expr, static_values, globals_) |
| 819 | domains.append((name, _value_expander(value))) |
| 820 | else: |
| 821 | local_exprs.append((name, value_expr)) |
| 822 | continue |
| 823 | if ( |
| 824 | isinstance(statement, ast.Expr) |
| 825 | and isinstance(statement.value, ast.Call) |
| 826 | and (get_ast_full_name(statement.value.func) or "").endswith("_when") |
| 827 | ): |
| 828 | if len(statement.value.args) != 1 or statement.value.keywords: |
| 829 | raise _dispatch_expr_error( |
| 830 | statement, "_when requires exactly one positional expression" |
| 831 | ) |
| 832 | predicates.append(statement.value.args[0]) |
| 833 | continue |
| 834 | if isinstance(statement, ast.Return): |
| 835 | if return_expr is not None or statement.value is None: |
| 836 | raise _dispatch_expr_error( |
| 837 | statement, "Warmup cases require one final return expression" |
| 838 | ) |
| 839 | return_expr = statement.value |
| 840 | continue |
| 841 | raise _dispatch_expr_error( |
| 842 | statement, |
| 843 | "AST-traced warmup cases support assignments, _when, and return", |
| 844 | ) |
| 845 | |
| 846 | if return_expr is None or not isinstance(return_expr, ast.Call): |
| 847 | raise ValueError("AST-traced warmup cases must return one case(...) call") |
| 848 | if any(keyword.arg is None for keyword in return_expr.keywords): |
| 849 | raise _dispatch_expr_error( |
| 850 | return_expr, "Warmup case expressions do not support **kwargs" |
| 851 | ) |
| 852 | |
| 853 | class _InlineDomainRewriter(ast.NodeTransformer): |
| 854 | def visit_Call(self, node: ast.Call) -> ast.AST: |
| 855 | call_name = (get_ast_full_name(node.func) or "").split(".")[-1] |
| 856 | if call_name not in {"WarmupIntRange", "WarmupChoices"}: |
| 857 | return self.generic_visit(node) |
| 858 | name = f"__warmup_domain_{len(domains)}" |
| 859 | domain = _eval_dispatch_expr(node, static_values, globals_) |
| 860 | domains.append((name, _value_expander(domain))) |
| 861 | return ast.copy_location(ast.Name(id=name, ctx=ast.Load()), node) |
| 862 | |
| 863 | return_expr = cast(ast.Call, _InlineDomainRewriter().visit(return_expr)) |
| 864 | |
| 865 | domain_names = tuple(name for name, _ in domains) |
| 866 | dynamic_names = set(domain_names) |
| 867 | dynamic_local_exprs: list[tuple[str, ast.AST]] = [] |
| 868 | for name, expr in local_exprs: |
| 869 | referenced_names = { |
| 870 | node.id for node in ast.walk(expr) if isinstance(node, ast.Name) |
| 871 | } |
| 872 | if referenced_names & dynamic_names: |
| 873 | dynamic_local_exprs.append((name, expr)) |
| 874 | dynamic_names.add(name) |
| 875 | else: |
| 876 | static_values[name] = _eval_dispatch_expr(expr, static_values, globals_) |
| 877 | |
| 878 | for _, expr in dynamic_local_exprs: |
| 879 | _validate_dispatch_expr(expr) |
| 880 | for predicate in predicates: |
| 881 | _validate_dispatch_expr(predicate) |
| 882 | _validate_dispatch_expr(return_expr) |
| 883 | |
| 884 | case_body: list[ast.stmt] = [ |
| 885 | ast.Assign( |
| 886 | targets=[ast.Name(id=name, ctx=ast.Store())], |
| 887 | value=cast(ast.expr, expr), |
| 888 | ) |
| 889 | for name, expr in dynamic_local_exprs |
| 890 | ] |
| 891 | if predicates: |
| 892 | predicate = ( |
| 893 | predicates[0] |
| 894 | if len(predicates) == 1 |
| 895 | else ast.BoolOp( |
| 896 | op=ast.And(), |
| 897 | values=[cast(ast.expr, value) for value in predicates], |
| 898 | ) |
| 899 | ) |
| 900 | case_body.append( |
| 901 | ast.If( |
| 902 | test=ast.UnaryOp(op=ast.Not(), operand=cast(ast.expr, predicate)), |
| 903 | body=[ast.Return(value=ast.Constant(value=None))], |
| 904 | orelse=[], |
| 905 | ) |
| 906 | ) |
| 907 | case_body.append(ast.Return(value=cast(ast.expr, return_expr))) |
| 908 | # FunctionDef fields vary across the supported Python versions, so no |
| 909 | # single constructor overload matches every mypy target. |
| 910 | case_function = ast.FunctionDef( # type: ignore[call-overload] |
| 911 | name="__vllm_warmup_case", |
| 912 | args=ast.arguments( |
| 913 | posonlyargs=[], |
| 914 | args=[ast.arg(arg=name) for name in domain_names], |
| 915 | kwonlyargs=[], |
| 916 | kw_defaults=[], |
| 917 | defaults=[], |
| 918 | ), |
| 919 | body=case_body, |
| 920 | decorator_list=[], |
| 921 | returns=None, |
| 922 | type_comment=None, |
| 923 | ) |
| 924 | case_globals = dict(globals_) |
| 925 | case_globals.update(static_values) |
| 926 | exec( |
| 927 | compile( |
| 928 | ast.fix_missing_locations( |
| 929 | ast.Module(body=[case_function], type_ignores=[]) |
| 930 | ), |
| 931 | "<jit-warmup>", |
| 932 | "exec", |
| 933 | ), |
| 934 | case_globals, |
| 935 | ) |
| 936 | evaluate_case = cast(Callable[..., Any], case_globals[case_function.name]) |
| 937 | |
| 938 | domain_values = tuple(values for _, values in domains) |
| 939 | for values in itertools.product(*domain_values): |
| 940 | case = evaluate_case(*values) |
| 941 | if case is None: |
| 942 | continue |
| 943 | if not isinstance(case, Mapping): |
| 944 | raise TypeError("AST-traced warmup cases must return a mapping") |
| 945 | if any( |
| 946 | isinstance(value, WarmupChoices | WarmupIntRange) |
| 947 | for value in case.values() |
| 948 | ): |
| 949 | raise TypeError("Warmup domains must be direct expressions") |
| 950 | yield case |
| 951 | |
| 952 | def dispatch(self, **kwargs: Any) -> CompileKeyT: |
| 953 | """Build one compile key from one concrete dispatch point.""" |
| 954 | raise NotImplementedError |
| 955 | |
| 956 | @abstractmethod |
| 957 | def get_warmup_keys(self, *args: Any, **kwargs: Any) -> list[CompileKeyT]: |
| 958 | """Return compile keys that should be warmed for this kernel.""" |
| 959 | raise NotImplementedError |
| 960 | |
| 961 | @abstractmethod |
| 962 | def compile(self, compile_key: CompileKeyT) -> None: |
| 963 | """Compile one warmup key.""" |
| 964 | raise NotImplementedError |
| 965 | |
| 966 | def compile_many(self, compile_keys: Iterable[CompileKeyT]) -> None: |
| 967 | """Compile a batch of warmup keys, allowing backend-specific scheduling.""" |
| 968 | for compile_key in compile_keys: |
| 969 | self.compile(compile_key) |
| 970 | |
| 971 | def register_warmup(self, *args: Any, **kwargs: Any) -> None: |
| 972 | """Register this kernel with the active runner's warmup registry.""" |
| 973 | JitWarmupRegistry.register(self, *args, **kwargs) |
| 974 | |
| 975 | def warmup(self, *args: Any, **kwargs: Any) -> None: |
| 976 | """Compile this kernel's warmup keys.""" |
| 977 | self.compile_many(self.get_warmup_keys(*args, **kwargs)) |
| 978 | |
| 979 | |
| 980 | def _same_value(left: Any, right: Any) -> bool: |
| 981 | """True if two registration values are the same object or compare equal. |
| 982 | |
| 983 | Identity is checked first so shared singletons (e.g. ``vllm_config``, torch |
| 984 | dtypes) short-circuit before any potentially deep or non-boolean ``__eq__``. |
| 985 | """ |
| 986 | if left is right: |
| 987 | return True |
| 988 | try: |
| 989 | return bool(left == right) |
| 990 | except (TypeError, ValueError, RuntimeError): |
| 991 | return False |
| 992 | |
| 993 | |
| 994 | def _same_registration( |
| 995 | left: tuple[tuple[Any, ...], dict[str, Any]], |
| 996 | right: tuple[tuple[Any, ...], dict[str, Any]], |
| 997 | ) -> bool: |
| 998 | """True if two ``(args, kwargs)`` registrations are equivalent.""" |
| 999 | left_args, left_kwargs = left |
| 1000 | right_args, right_kwargs = right |
| 1001 | if len(left_args) != len(right_args) or left_kwargs.keys() != right_kwargs.keys(): |
| 1002 | return False |
| 1003 | return all( |
| 1004 | _same_value(x, y) for x, y in zip(left_args, right_args, strict=True) |
| 1005 | ) and all(_same_value(left_kwargs[k], right_kwargs[k]) for k in left_kwargs) |
| 1006 | |
| 1007 | |
| 1008 | class JitWarmupRegistry: |
| 1009 | """Collect and compile JIT kernels selected during runner setup.""" |
| 1010 | |
| 1011 | _active: ContextVar[JitWarmupRegistry | None] = ContextVar( |
| 1012 | "active_jit_warmup_registry", |
| 1013 | default=None, |
| 1014 | ) |
| 1015 | |
| 1016 | def __init__(self, vllm_config: Any) -> None: |
| 1017 | self.vllm_config = vllm_config |
| 1018 | self._registrations: dict[ |
| 1019 | VllmJitKernel[Any], |
| 1020 | list[tuple[tuple[Any, ...], dict[str, Any]]], |
| 1021 | ] = {} |
| 1022 | |
| 1023 | @classmethod |
| 1024 | def capture(cls, init_fn: Callable[..., None]) -> Callable[..., None]: |
| 1025 | """Collect warmup registrations made while the decorated callable runs.""" |
| 1026 | |
| 1027 | @wraps(init_fn) |
| 1028 | def wrapped(instance: Any, vllm_config: Any, *args: Any, **kwargs: Any) -> None: |
| 1029 | registry = cls(vllm_config) |
| 1030 | instance.jit_warmup_registry = registry |
| 1031 | with registry.activate(): |
| 1032 | init_fn(instance, vllm_config, *args, **kwargs) |
| 1033 | |
| 1034 | return wrapped |
| 1035 | |
| 1036 | @contextmanager |
| 1037 | def activate(self) -> Iterator[None]: |
| 1038 | """Collect registrations made in this context.""" |
| 1039 | token = self._active.set(self) |
| 1040 | try: |
| 1041 | yield |
| 1042 | finally: |
| 1043 | self._active.reset(token) |
| 1044 | |
| 1045 | @classmethod |
| 1046 | def register( |
| 1047 | cls, |
| 1048 | kernel: VllmJitKernel[Any], |
| 1049 | *args: Any, |
| 1050 | **kwargs: Any, |
| 1051 | ) -> None: |
| 1052 | """Register a kernel with the active registry, if one exists.""" |
| 1053 | registry = cls._active.get() |
| 1054 | if registry is not None: |
| 1055 | registry._add(kernel, args, kwargs) |
| 1056 | |
| 1057 | def _add( |
| 1058 | self, |
| 1059 | kernel: VllmJitKernel[Any], |
| 1060 | args: tuple[Any, ...], |
| 1061 | kwargs: dict[str, Any], |
| 1062 | ) -> None: |
| 1063 | registrations = self._registrations.setdefault(kernel, []) |
| 1064 | # Every layer of a deep model appends an identical (args, kwargs) here |
| 1065 | # (e.g. a 61-layer DSA model registers the same pack (dtype, pad_value) |
| 1066 | # 61 times); each expands to the same compile keys, so tracing |
| 1067 | # get_warmup_keys once per distinct registration is sufficient. Dedup on |
| 1068 | # identity-or-equality -- identity short-circuits shared singletons like |
| 1069 | # vllm_config before any deep __eq__. |
| 1070 | if any( |
| 1071 | _same_registration(registered, (args, kwargs)) |
| 1072 | for registered in registrations |
| 1073 | ): |
| 1074 | return |
| 1075 | registrations.append((args, kwargs)) |
| 1076 | |
| 1077 | def __len__(self) -> int: |
| 1078 | return sum(len(registrations) for registrations in self._registrations.values()) |
| 1079 | |
| 1080 | def warmup(self) -> None: |
| 1081 | """Expand registrations and compile each wrapper/key pair once.""" |
| 1082 | from tqdm import tqdm |
| 1083 | |
| 1084 | from vllm.distributed import is_global_first_rank |
| 1085 | |
| 1086 | kernel_items: list[tuple[VllmJitKernel[Any], dict[Any, None]]] = [] |
| 1087 | for kernel, registrations in self._registrations.items(): |
| 1088 | compile_keys: dict[Any, None] = {} |
| 1089 | for args, kwargs in registrations: |
| 1090 | if ( |
| 1091 | not args |
| 1092 | and not kwargs |
| 1093 | and "vllm_config" |
| 1094 | in inspect.signature(kernel.get_warmup_keys).parameters |
| 1095 | ): |
| 1096 | kwargs = {"vllm_config": self.vllm_config} |
| 1097 | for compile_key in kernel.get_warmup_keys(*args, **kwargs): |
| 1098 | compile_keys[compile_key] = None |
| 1099 | if compile_keys: |
| 1100 | kernel_items.append((kernel, compile_keys)) |
| 1101 | |
| 1102 | if not kernel_items: |
| 1103 | return |
| 1104 | |
| 1105 | total_keys = sum(len(compile_keys) for _, compile_keys in kernel_items) |
| 1106 | with tqdm( |
| 1107 | kernel_items, |
| 1108 | desc=f"JIT kernel warmup ({total_keys} compile keys)", |
| 1109 | disable=not is_global_first_rank(), |
| 1110 | dynamic_ncols=True, |
| 1111 | unit="kernel", |
| 1112 | ) as progress: |
| 1113 | for kernel, compile_keys in progress: |
| 1114 | progress.set_postfix_str( |
| 1115 | f"{kernel.__class__.__name__} ({len(compile_keys)} keys)", |
| 1116 | refresh=False, |
| 1117 | ) |
| 1118 | kernel.compile_many(compile_keys) |