| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """fx tracing and forward-source rewriting for the Transformers backend fusers. |
| 4 | |
| 5 | A small engine, independent of any particular pattern: trace a module's forward |
| 6 | with `torch.fx` (tolerating a partial graph), inspect the resulting nodes, and |
| 7 | rewrite the forward's *source* (AST) so only matched calls change while the rest |
| 8 | stays live Python. `fusion.py` builds the concrete fusion patterns on top. |
| 9 | """ |
| 10 | |
| 11 | import ast |
| 12 | import contextlib |
| 13 | import inspect |
| 14 | import operator |
| 15 | import textwrap |
| 16 | from collections.abc import Callable |
| 17 | from itertools import chain |
| 18 | from typing import Any |
| 19 | from unittest import mock |
| 20 | |
| 21 | import torch |
| 22 | from torch import fx, nn |
| 23 | from torch.nn import functional as F |
| 24 | from typing_extensions import TypeIs |
| 25 | |
| 26 | from vllm.logger import init_logger |
| 27 | |
| 28 | logger = init_logger(__name__) |
| 29 | |
| 30 | _UNKNOWN: Any = object() |
| 31 | """Sentinel meta value for proxies whose concrete value could not be inferred. |
| 32 | Distinct from `None`, which is a valid concrete value (e.g. `attn_weights`).""" |
| 33 | |
| 34 | _MODULE_CALL = nn.Module.__call__ |
| 35 | """The unpatched `nn.Module.__call__`. During tracing fx patches it to record |
| 36 | `call_module` nodes; meta execution must call modules for real.""" |
| 37 | |
| 38 | |
| 39 | def is_leaf_call(node: object) -> TypeIs[fx.Node]: |
| 40 | """Is node a call recorded by `_as_leaf_call` (e.g. an attention interface).""" |
| 41 | return isinstance(node, fx.Node) and node.meta.get("leaf_call", False) |
| 42 | |
| 43 | |
| 44 | def _reference_weight(module: nn.Module) -> torch.Tensor | None: |
| 45 | """A weight whose trailing dim is the module's hidden size. |
| 46 | |
| 47 | Linears and 2-D gate weights are `[out, hidden]`; norm weights are |
| 48 | `[hidden]`. Used to fabricate a placeholder input of matching size/dtype.""" |
| 49 | for child in module.modules(): |
| 50 | if isinstance(child, nn.Linear): |
| 51 | return child.weight |
| 52 | for param in module.parameters(): |
| 53 | if param.ndim in (1, 2): |
| 54 | return param |
| 55 | return None |
| 56 | |
| 57 | |
| 58 | class _MetaProxy(fx.Proxy): |
| 59 | """Proxy carrying the meta-tensor value of the traced expression. |
| 60 | |
| 61 | Shape questions (`len`, iteration, `.shape` unpacks) are answered by |
| 62 | executing each op on the meta values, so PyTorch's meta kernels are the |
| 63 | single source of shape inference — no per-op rules.""" |
| 64 | |
| 65 | meta: Any = _UNKNOWN |
| 66 | |
| 67 | def __len__(self) -> int: |
| 68 | if self.meta is not _UNKNOWN: |
| 69 | return len(self.meta) |
| 70 | return super().__len__() # type: ignore[misc] |
| 71 | |
| 72 | def __getattr__(self, k: str) -> "_MetaAttribute": |
| 73 | return _MetaAttribute(self, k) |
| 74 | |
| 75 | def __setitem__(self, key: object, value: object) -> None: |
| 76 | # Record `obj[key] = value` so tracing continues past in-place writes |
| 77 | # into containers (e.g. Gemma 4's `shared_kv_states[...] = ...`). |
| 78 | self.tracer.create_proxy( |
| 79 | "call_function", operator.setitem, (self, key, value), {} |
| 80 | ) |
| 81 | |
| 82 | |
| 83 | class _MetaAttribute(_MetaProxy, fx.proxy.Attribute): |
| 84 | """Attribute proxy (e.g. `x.shape`) carrying its meta value. |
| 85 | |
| 86 | `Proxy.__getattr__` constructs `Attribute` directly, bypassing |
| 87 | `Tracer.proxy`, so the meta value must be grafted on here too.""" |
| 88 | |
| 89 | def __init__(self, root: fx.Proxy, attr: str): |
| 90 | super().__init__(root, attr) |
| 91 | root_meta = getattr(root, "meta", _UNKNOWN) |
| 92 | if root_meta is not _UNKNOWN: |
| 93 | with contextlib.suppress(Exception): |
| 94 | self.meta = getattr(root_meta, attr) |
| 95 | |
| 96 | |
| 97 | class _AllLeafTracer(fx.Tracer): |
| 98 | """Tracer that treats every submodule as a leaf. |
| 99 | |
| 100 | Each child stays one `call_module` node, so matching sees the module's own |
| 101 | forward structure (activations aren't decomposed into e.g. `sigmoid * x`). |
| 102 | Every traced op is also executed on meta tensors (see `_MetaProxy`) so |
| 103 | shape unpacks and `*`-splats trace through; anything else untraceable ends |
| 104 | the trace early and the partial graph is matched. |
| 105 | """ |
| 106 | |
| 107 | varkw: str | None = None |
| 108 | """Name of the traced forward's `**kwargs` parameter, if any.""" |
| 109 | |
| 110 | def is_leaf_module(self, m: nn.Module, module_qualified_name: str) -> bool: |
| 111 | return True |
| 112 | |
| 113 | def proxy(self, node: fx.Node) -> fx.Proxy: |
| 114 | return _MetaProxy(node, self) |
| 115 | |
| 116 | def create_proxy(self, kind, target, args, kwargs, *extra, **extra_kwargs): |
| 117 | proxy = super().create_proxy(kind, target, args, kwargs, *extra, **extra_kwargs) |
| 118 | if isinstance(proxy, _MetaProxy) and proxy.meta is _UNKNOWN: |
| 119 | # A failure stays _UNKNOWN; only fatal if a shape question is asked. |
| 120 | with contextlib.suppress(Exception): |
| 121 | proxy.meta = self._infer_meta(kind, target, args, kwargs) |
| 122 | return proxy |
| 123 | |
| 124 | def _infer_meta(self, kind: str, target: object, args: tuple, kwargs: dict): |
| 125 | """Execute the op on meta tensors; PyTorch infers the output value.""" |
| 126 | if kind == "placeholder": |
| 127 | # vLLM always feeds the model [1, seq_len, hidden_size] hidden states. |
| 128 | weight = _reference_weight(self.root) |
| 129 | if str(target) == "hidden_states" and weight is not None: |
| 130 | return torch.empty( |
| 131 | 1, 8, weight.shape[-1], dtype=weight.dtype, device="meta" |
| 132 | ) |
| 133 | return _UNKNOWN |
| 134 | if kind == "get_attr": |
| 135 | value = operator.attrgetter(str(target))(self.root) |
| 136 | if isinstance(value, torch.Tensor): |
| 137 | value = torch.empty_like(value, device="meta") |
| 138 | return value |
| 139 | unknown = False |
| 140 | |
| 141 | def meta_of(arg: object) -> object: |
| 142 | nonlocal unknown |
| 143 | if isinstance(arg, fx.Proxy): |
| 144 | meta = getattr(arg, "meta", _UNKNOWN) |
| 145 | unknown = unknown or meta is _UNKNOWN |
| 146 | return meta |
| 147 | return arg |
| 148 | |
| 149 | meta_args = fx.node.map_aggregate(args, meta_of) |
| 150 | meta_kwargs = fx.node.map_aggregate(kwargs, meta_of) |
| 151 | if unknown: |
| 152 | return _UNKNOWN |
| 153 | if kind == "call_function": |
| 154 | assert callable(target) |
| 155 | return target(*meta_args, **meta_kwargs) |
| 156 | if kind == "call_method": |
| 157 | receiver, *rest = meta_args |
| 158 | return getattr(receiver, str(target))(*rest, **meta_kwargs) |
| 159 | if kind == "call_module": |
| 160 | # Run the child's forward with all its state on "meta", without |
| 161 | # mutating it (at match time params may be meta but buffers real). |
| 162 | # fx patches `nn.Module.__call__` while tracing; restore the real |
| 163 | # one so this execution is not itself recorded. |
| 164 | child = self.root.get_submodule(str(target)) |
| 165 | state = { |
| 166 | name: torch.empty_like(tensor, device="meta") |
| 167 | for name, tensor in chain( |
| 168 | child.named_parameters(), child.named_buffers() |
| 169 | ) |
| 170 | } |
| 171 | with mock.patch.object(nn.Module, "__call__", _MODULE_CALL): |
| 172 | return torch.func.functional_call(child, state, meta_args, meta_kwargs) |
| 173 | return _UNKNOWN |
| 174 | |
| 175 | def _is_varkw(self, node: object) -> bool: |
| 176 | return ( |
| 177 | isinstance(node, fx.Node) |
| 178 | and node.op == "placeholder" |
| 179 | and str(node.target).lstrip("*") == self.varkw |
| 180 | ) |
| 181 | |
| 182 | def iter(self, obj: fx.Proxy): |
| 183 | # Assume kwargs is always empty to simplify tracing. |
| 184 | node = obj.node |
| 185 | if self._is_varkw(node) or ( |
| 186 | node.op == "call_method" |
| 187 | and node.target == "keys" |
| 188 | and self._is_varkw(node.args[0]) |
| 189 | ): |
| 190 | return iter(()) |
| 191 | meta = getattr(obj, "meta", _UNKNOWN) |
| 192 | if meta is _UNKNOWN: |
| 193 | return super().iter(obj) |
| 194 | return iter([obj[i] for i in range(len(meta))]) |
| 195 | |
| 196 | |
| 197 | def _as_leaf_call(fn: Callable, length: int | None = None) -> Callable: |
| 198 | """Wrap any callable so tracing records it as one opaque `call_function` node. |
| 199 | |
| 200 | Lets the trace continue past untraceable bodies. Only the proxy arguments carry into |
| 201 | the node's dataflow; the rest are dropped rather than lifted into the graph. |
| 202 | `length` declares how many values the callable returns, so unpacking its result also |
| 203 | traces. Called without proxies (i.e. outside tracing), the wrapper is a passthrough. |
| 204 | """ |
| 205 | |
| 206 | def leaf(*args, **kwargs): |
| 207 | proxies = tuple(arg for arg in args if isinstance(arg, fx.Proxy)) |
| 208 | if not proxies: |
| 209 | return fn(*args, **kwargs) |
| 210 | proxy = proxies[0].tracer.create_proxy("call_function", fn, proxies, {}) |
| 211 | proxy.node.meta["leaf_call"] = True |
| 212 | if length is not None: |
| 213 | # The body never executes, so fabricate a value of the declared length. |
| 214 | proxy.meta = (_UNKNOWN,) * length |
| 215 | return proxy |
| 216 | |
| 217 | return leaf |
| 218 | |
| 219 | |
| 220 | def _leaf_attention_interfaces(): |
| 221 | """Patch `AttentionInterface.get_interface` so traced forwards see a leaf node. |
| 222 | |
| 223 | `vllm_attention_function` needs runtime context so it is untraceable. |
| 224 | Every interface returns `(attn_output, attn_weights)`.""" |
| 225 | from transformers.modeling_utils import AttentionInterface |
| 226 | |
| 227 | original = AttentionInterface.get_interface |
| 228 | |
| 229 | def get_interface(self, *args, **kwargs): |
| 230 | return _as_leaf_call(original(self, *args, **kwargs), length=2) |
| 231 | |
| 232 | return mock.patch.object(AttentionInterface, "get_interface", get_interface) |
| 233 | |
| 234 | |
| 235 | def trace(module: nn.Module) -> fx.Graph | None: |
| 236 | """Trace `module.forward`, returning the partial graph on failure.""" |
| 237 | parameters = forward_parameters(type(module)) |
| 238 | # vLLM never passes `past_key_values` so it is always the default value of `None`. |
| 239 | # Make this concrete to simplify tracing. |
| 240 | concrete_args = None |
| 241 | if "past_key_values" in parameters: |
| 242 | concrete_args = {"past_key_values": None} |
| 243 | # Get the name of the kwargs parameter passed to module.forward (usually "kwargs") |
| 244 | tracer = _AllLeafTracer() |
| 245 | tracer.varkw = next( |
| 246 | ( |
| 247 | p.name |
| 248 | for p in parameters.values() |
| 249 | if p.kind is inspect.Parameter.VAR_KEYWORD |
| 250 | ), |
| 251 | None, |
| 252 | ) |
| 253 | try: |
| 254 | with _leaf_attention_interfaces(): |
| 255 | return tracer.trace(module, concrete_args=concrete_args) |
| 256 | except Exception as exc: |
| 257 | logger.debug("Could not fully trace %s: %s", type(module), exc) |
| 258 | return getattr(tracer, "graph", None) |
| 259 | |
| 260 | |
| 261 | def recover_forward(cls: type[nn.Module]) -> tuple[ast.FunctionDef, Callable]: |
| 262 | """Parse the source of `cls.forward`, ready for rewriting.""" |
| 263 | fn = inspect.unwrap(cls.forward) |
| 264 | if fn.__code__.co_freevars: |
| 265 | raise ValueError("forward is a closure") |
| 266 | tree = ast.parse(textwrap.dedent(inspect.getsource(fn))) |
| 267 | funcdef = tree.body[0] |
| 268 | if not isinstance(funcdef, ast.FunctionDef): |
| 269 | raise ValueError("source is not a plain function definition") |
| 270 | # `fn` is already unwrapped; don't re-apply its decorators |
| 271 | funcdef.decorator_list.clear() |
| 272 | # Annotations may not evaluate outside the defining module (e.g. with |
| 273 | # postponed evaluation); they're not needed at runtime |
| 274 | funcdef.returns = None |
| 275 | args = funcdef.args |
| 276 | for arg in ( |
| 277 | *args.posonlyargs, |
| 278 | *args.args, |
| 279 | *args.kwonlyargs, |
| 280 | *filter(None, (args.vararg, args.kwarg)), |
| 281 | ): |
| 282 | arg.annotation = None |
| 283 | # Recompiling outside the class body would break name mangling |
| 284 | for node in ast.walk(funcdef): |
| 285 | name = getattr(node, "attr", None) or getattr(node, "id", None) |
| 286 | if name and name.startswith("__") and not name.endswith("__"): |
| 287 | raise ValueError(f"{name} would be name mangled") |
| 288 | return funcdef, fn |
| 289 | |
| 290 | |
| 291 | def forward_parameters(cls: type[nn.Module]) -> dict[str, inspect.Parameter]: |
| 292 | """`cls.forward`'s signature parameters, or empty if uninspectable.""" |
| 293 | try: |
| 294 | return dict(inspect.signature(cls.forward).parameters) |
| 295 | except (TypeError, ValueError): |
| 296 | return {} |
| 297 | |
| 298 | |
| 299 | def forward_input_count(cls: type[nn.Module]) -> int: |
| 300 | """The number of tensor inputs `cls.forward` declares, excluding `self` and |
| 301 | any `*args`/`**kwargs`. Read from the signature, so it is independent of |
| 302 | whether the trace completes (unlike counting placeholders).""" |
| 303 | params = list(forward_parameters(cls).values()) |
| 304 | if not params: |
| 305 | return 1 # uninspectable: assume a single input and let matching decide |
| 306 | fixed = ( |
| 307 | inspect.Parameter.POSITIONAL_ONLY, |
| 308 | inspect.Parameter.POSITIONAL_OR_KEYWORD, |
| 309 | inspect.Parameter.KEYWORD_ONLY, |
| 310 | ) |
| 311 | return sum(1 for p in params[1:] if p.kind in fixed) |
| 312 | |
| 313 | |
| 314 | def compile_forward(funcdef: ast.FunctionDef, fn: Callable) -> Callable: |
| 315 | """Compile `funcdef` in `fn`'s module so tracebacks point at the source.""" |
| 316 | module = ast.Module(body=[funcdef], type_ignores=[]) |
| 317 | ast.fix_missing_locations(module) |
| 318 | ast.increment_lineno(module, fn.__code__.co_firstlineno - 1) |
| 319 | code = compile(module, fn.__code__.co_filename, "exec") |
| 320 | namespace: dict = {} |
| 321 | exec(code, fn.__globals__, namespace) |
| 322 | return namespace[funcdef.name] |
| 323 | |
| 324 | |
| 325 | def self_call_and_refs( |
| 326 | funcdef: ast.FunctionDef, name: str |
| 327 | ) -> tuple[ast.Call, list[ast.Attribute]]: |
| 328 | """The unique `self.<name>(arg)` call, plus every other `self.<name>` reference. |
| 329 | |
| 330 | Raises unless exactly one single-argument `self.<name>(arg)` call exists, so |
| 331 | one fx `call_module` node maps to one syntactic call site. The remaining |
| 332 | references (e.g. `is not None` guards) are returned for the caller to resolve. |
| 333 | """ |
| 334 | uses = [ |
| 335 | node |
| 336 | for node in ast.walk(funcdef) |
| 337 | if isinstance(node, ast.Attribute) |
| 338 | and node.attr == name |
| 339 | and isinstance(node.value, ast.Name) |
| 340 | and node.value.id == "self" |
| 341 | ] |
| 342 | calls = [ |
| 343 | node |
| 344 | for node in ast.walk(funcdef) |
| 345 | if isinstance(node, ast.Call) |
| 346 | and any(node.func is use for use in uses) |
| 347 | and len(node.args) == 1 |
| 348 | and not isinstance(node.args[0], ast.Starred) |
| 349 | and not node.keywords |
| 350 | ] |
| 351 | if len(calls) != 1: |
| 352 | raise ValueError(f"{name} is not called exactly once as self.{name}(arg)") |
| 353 | call = calls[0] |
| 354 | other_refs = [use for use in uses if use is not call.func] |
| 355 | return call, other_refs |
| 356 | |
| 357 | |
| 358 | def single_self_call(funcdef: ast.FunctionDef, name: str) -> ast.Call: |
| 359 | """The unique `self.<name>(arg)` call in `funcdef`. |
| 360 | |
| 361 | Raises unless `name` appears exactly once, as such a call, so the source |
| 362 | rewrite agrees with the fx match. |
| 363 | """ |
| 364 | call, other_refs = self_call_and_refs(funcdef, name) |
| 365 | if other_refs: |
| 366 | raise ValueError(f"{name} is referenced {len(other_refs) + 1} times") |
| 367 | return call |
| 368 | |
| 369 | |
| 370 | def _is_none(node: ast.expr) -> bool: |
| 371 | return isinstance(node, ast.Constant) and node.value is None |
| 372 | |
| 373 | |
| 374 | def _in_boolean_context(funcdef: ast.FunctionDef, ref: ast.expr) -> bool: |
| 375 | """Is `ref` used only for its truth value (a test, or a `not` operand)? |
| 376 | |
| 377 | In these positions the object's identity never escapes, so a reference that |
| 378 | is always truthy can be replaced by `True`. `and`/`or` are excluded: they |
| 379 | yield an operand, so the module could escape (`x and self.<name>`).""" |
| 380 | for node in ast.walk(funcdef): |
| 381 | if ( |
| 382 | isinstance(node, (ast.If, ast.IfExp, ast.While, ast.Assert)) |
| 383 | and node.test is ref |
| 384 | ): |
| 385 | return True |
| 386 | if ( |
| 387 | isinstance(node, ast.UnaryOp) |
| 388 | and isinstance(node.op, ast.Not) |
| 389 | and node.operand is ref |
| 390 | ): |
| 391 | return True |
| 392 | return False |
| 393 | |
| 394 | |
| 395 | def bypass_existence_guard( |
| 396 | funcdef: ast.FunctionDef, ref: ast.Attribute, name: str |
| 397 | ) -> None: |
| 398 | """Fold a guard on `self.<name>`'s existence to its constant value. |
| 399 | |
| 400 | The fuser deletes `self.<name>` and binds only instances where it exists as a |
| 401 | truthy `nn.Linear` (guaranteed by the match, the cache key, and `validate`), |
| 402 | so a guard testing its presence is a fusion invariant. Two forms are folded: |
| 403 | |
| 404 | - an identity `None` check `self.<name> is None` -> `False`, |
| 405 | `self.<name> is not None` -> `True` (either operand order). `==`/`!=` are |
| 406 | *not* folded: `is_linear` accepts `nn.Linear` subclasses, which may override |
| 407 | `__eq__`/`__ne__`, so equality is not guaranteed to track identity. |
| 408 | - a bare truthiness test (`if self.<name>:`, `... if self.<name> else ...`, |
| 409 | `not self.<name>`): the reference itself to `True`. |
| 410 | |
| 411 | Any other surviving reference escapes the projection's value, which no longer |
| 412 | exists after fusion, so refuse rather than change semantics.""" |
| 413 | for node in ast.walk(funcdef): |
| 414 | # `self.<name> is (not) None`, either operand order. |
| 415 | if not (isinstance(node, ast.Compare) and len(node.ops) == 1): |
| 416 | continue |
| 417 | (op,), (right,) = node.ops, node.comparators |
| 418 | if not isinstance(op, (ast.Is, ast.IsNot)): |
| 419 | continue |
| 420 | if (ref is node.left and _is_none(right)) or ( |
| 421 | ref is right and _is_none(node.left) |
| 422 | ): |
| 423 | value = isinstance(op, ast.IsNot) |
| 424 | replace_expr(funcdef, node, ast.copy_location(ast.Constant(value), node)) |
| 425 | return |
| 426 | # A bare truthiness test (statement `if`, ternary, or `not`); an `nn.Linear` |
| 427 | # is always truthy, so the reference folds to `True`. |
| 428 | if _in_boolean_context(funcdef, ref): |
| 429 | replace_expr(funcdef, ref, ast.copy_location(ast.Constant(True), ref)) |
| 430 | return |
| 431 | raise ValueError(f"{name} is referenced outside an existence guard") |
| 432 | |
| 433 | |
| 434 | def block_chain( |
| 435 | block: list[ast.stmt], node: ast.AST |
| 436 | ) -> list[tuple[list[ast.stmt], int]]: |
| 437 | """Path of (statement list, index) pairs from `block` down to `node`. |
| 438 | |
| 439 | Each pair names a nested block and the index in it of the statement |
| 440 | containing `node`; the last pair is `node`'s innermost block. Empty if |
| 441 | `node` is not in `block`. The longest common prefix of several nodes' chains |
| 442 | is the innermost block that dominates them all. |
| 443 | """ |
| 444 | for index, stmt in enumerate(block): |
| 445 | if not any(child is node for child in ast.walk(stmt)): |
| 446 | continue |
| 447 | child_blocks = [ |
| 448 | getattr(stmt, fld, None) for fld in ("body", "orelse", "finalbody") |
| 449 | ] |
| 450 | child_blocks += [h.body for h in getattr(stmt, "handlers", [])] |
| 451 | child_blocks += [c.body for c in getattr(stmt, "cases", [])] |
| 452 | for child_block in child_blocks: |
| 453 | if ( |
| 454 | isinstance(child_block, list) |
| 455 | and child_block |
| 456 | and (tail := block_chain(child_block, node)) |
| 457 | ): |
| 458 | return [(block, index), *tail] |
| 459 | return [(block, index)] |
| 460 | return [] |
| 461 | |
| 462 | |
| 463 | def _base_name(node: ast.expr) -> str | None: |
| 464 | """The root `Name` of an attribute/subscript chain (`a.b[c]` -> `a`).""" |
| 465 | while isinstance(node, (ast.Attribute, ast.Subscript)): |
| 466 | node = node.value |
| 467 | return node.id if isinstance(node, ast.Name) else None |
| 468 | |
| 469 | |
| 470 | def _rebound_names(region: list[ast.stmt]) -> set[str]: |
| 471 | """Names rebound outright within `region` (`x = ...`, `del x`). |
| 472 | |
| 473 | These are `Name` nodes in a `Store`/`Del` context.""" |
| 474 | return { |
| 475 | node.id |
| 476 | for stmt in region |
| 477 | for node in ast.walk(stmt) |
| 478 | if isinstance(node, ast.Name) and isinstance(node.ctx, (ast.Store, ast.Del)) |
| 479 | } |
| 480 | |
| 481 | |
| 482 | def _inplace_target(node: ast.AST) -> str | None: |
| 483 | """Base name `node` mutates in place, if any. |
| 484 | |
| 485 | A write through the name (`x[i] = ...`, `x.attr = ...`) or an in-place method |
| 486 | call (`x.mul_(...)`) leaves the base name in a `Load` context, so it is not a |
| 487 | plain `Name` store (see `_rebound_names`).""" |
| 488 | if isinstance(node, (ast.Attribute, ast.Subscript)) and isinstance( |
| 489 | node.ctx, (ast.Store, ast.Del) |
| 490 | ): |
| 491 | return _base_name(node) |
| 492 | if ( |
| 493 | isinstance(node, ast.Call) |
| 494 | and isinstance(node.func, ast.Attribute) |
| 495 | and node.func.attr.endswith("_") |
| 496 | and not node.func.attr.endswith("__") |
| 497 | ): |
| 498 | return _base_name(node.func.value) |
| 499 | return None |
| 500 | |
| 501 | |
| 502 | def _mutated_names(region: list[ast.stmt]) -> set[str]: |
| 503 | """Names mutated in place within `region` (see `_inplace_target`).""" |
| 504 | return { |
| 505 | base |
| 506 | for stmt in region |
| 507 | for node in ast.walk(stmt) |
| 508 | if (base := _inplace_target(node)) is not None |
| 509 | } |
| 510 | |
| 511 | |
| 512 | def _call_name(call: ast.Call) -> str: |
| 513 | """The called function's own name, without its qualifier.""" |
| 514 | if isinstance(call.func, ast.Attribute): |
| 515 | return call.func.attr |
| 516 | return call.func.id if isinstance(call.func, ast.Name) else "" |
| 517 | |
| 518 | |
| 519 | def _arg_names(call: ast.Call) -> set[str]: |
| 520 | """Every name `call` reads through its arguments.""" |
| 521 | reads = [*call.args, *(kw.value for kw in call.keywords)] |
| 522 | return { |
| 523 | node.id |
| 524 | for arg in reads |
| 525 | for node in ast.walk(arg) |
| 526 | if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load) |
| 527 | } |
| 528 | |
| 529 | |
| 530 | def _writes_through_args(call: ast.Call) -> bool: |
| 531 | """Does `call` write through an argument rather than return a new value? |
| 532 | |
| 533 | Torch spells this three ways: an `out=` destination, `inplace=True`, and the |
| 534 | trailing-underscore convention (`torch.relu_(x)`). The underscore is checked |
| 535 | on the function's own name, so it catches the free-function form that |
| 536 | `_inplace_target` cannot (there the base name is the module, `torch`). |
| 537 | """ |
| 538 | for keyword in call.keywords: |
| 539 | if keyword.arg == "out": |
| 540 | return True |
| 541 | if keyword.arg == "inplace" and not ( |
| 542 | isinstance(keyword.value, ast.Constant) and keyword.value.value is False |
| 543 | ): |
| 544 | return True |
| 545 | name = _call_name(call) |
| 546 | return name.endswith("_") and not name.endswith("__") |
| 547 | |
| 548 | |
| 549 | def written_names(region: list[ast.stmt]) -> set[str]: |
| 550 | """Names `region` may rebind or write through. |
| 551 | |
| 552 | Beyond plain rebinding and writes through the name itself (`_rebound_names`, |
| 553 | `_mutated_names`), a call can write through an argument without naming it as |
| 554 | a target: `F.relu(x, inplace=True)`, `torch.add(x, y, out=x)`, |
| 555 | `torch.relu_(x)`. A statement-level call whose result is discarded is treated |
| 556 | the same way, since it can only be there for a side effect. |
| 557 | """ |
| 558 | written = _rebound_names(region) | _mutated_names(region) |
| 559 | for stmt in region: |
| 560 | if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Call): |
| 561 | written |= _arg_names(stmt.value) |
| 562 | for node in ast.walk(stmt): |
| 563 | if isinstance(node, ast.Call) and _writes_through_args(node): |
| 564 | written |= _arg_names(node) |
| 565 | return written |
| 566 | |
| 567 | |
| 568 | _METADATA_ATTRS = frozenset( |
| 569 | {"shape", "dtype", "device", "ndim", "size", "dim", "numel", "item"} |
| 570 | ) |
| 571 | |
| 572 | |
| 573 | def _aliasing_reads(node: ast.expr) -> set[str]: |
| 574 | """Names `node` reads in a way that could yield a view of them. |
| 575 | |
| 576 | Metadata (`x.shape`, `x.size(0)`, `x.dtype`) is ints and tuples sharing no |
| 577 | storage with `x`, so reaching a name only through metadata cannot alias it. |
| 578 | Shape arithmetic is pervasive between projections, so propagating through it |
| 579 | would taint nearly every later name. |
| 580 | """ |
| 581 | names: set[str] = set() |
| 582 | stack: list[ast.AST] = [node] |
| 583 | while stack: |
| 584 | current = stack.pop() |
| 585 | if isinstance(current, ast.Attribute) and current.attr in _METADATA_ATTRS: |
| 586 | continue |
| 587 | if ( |
| 588 | isinstance(current, ast.Call) |
| 589 | and isinstance(current.func, ast.Attribute) |
| 590 | and current.func.attr in _METADATA_ATTRS |
| 591 | ): |
| 592 | stack.extend(current.args) |
| 593 | stack.extend(keyword.value for keyword in current.keywords) |
| 594 | continue |
| 595 | if isinstance(current, ast.Name) and isinstance(current.ctx, ast.Load): |
| 596 | names.add(current.id) |
| 597 | stack.extend(ast.iter_child_nodes(current)) |
| 598 | return names |
| 599 | |
| 600 | |
| 601 | def _returns_fresh(node: ast.expr, module: nn.Module) -> bool: |
| 602 | """Does evaluating `node` allocate, rather than return a view of its input? |
| 603 | |
| 604 | A `nn.Linear` writes its output into new storage, so its result aliases |
| 605 | nothing that went in, and a chain of view ops on top of it aliases only that |
| 606 | fresh tensor. Walking the receiver chain down to a `self.<linear>(...)` call |
| 607 | therefore proves the value cannot alias the projection input. A literal |
| 608 | aliases nothing either, and a branch or sequence is fresh when every part is |
| 609 | -- which is how a guarded projection (`p(x) if p is not None else None`) |
| 610 | stays untracked. A call on an attribute that is missing or `None` cannot run |
| 611 | at all, so a branch selecting a different configuration is fresh too. |
| 612 | """ |
| 613 | if isinstance(node, ast.Constant): |
| 614 | return True |
| 615 | if isinstance(node, ast.IfExp): |
| 616 | if isinstance(node.test, ast.Constant): |
| 617 | # A folded existence guard leaves the untaken branch as dead code. |
| 618 | taken = node.body if node.test.value else node.orelse |
| 619 | return _returns_fresh(taken, module) |
| 620 | return _returns_fresh(node.body, module) and _returns_fresh(node.orelse, module) |
| 621 | if isinstance(node, (ast.Tuple, ast.List)): |
| 622 | return all(_returns_fresh(elt, module) for elt in node.elts) |
| 623 | while isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): |
| 624 | func = node.func |
| 625 | if isinstance(func.value, ast.Name) and func.value.id == "self": |
| 626 | attribute = getattr(module, func.attr, None) |
| 627 | if attribute is None or isinstance(attribute, nn.Linear): |
| 628 | # A missing or `None` attribute cannot be called, so the branch |
| 629 | # selecting it never runs (`q_proj` on a q-LoRA layer). |
| 630 | return True |
| 631 | # Anything else can only return storage it was handed, so it is |
| 632 | # fresh exactly when its inputs are (`q_norm(q_proj(x))`). |
| 633 | return bool(node.args) and all( |
| 634 | _returns_fresh(arg, module) for arg in node.args |
| 635 | ) |
| 636 | node = func.value |
| 637 | return False |
| 638 | |
| 639 | |
| 640 | def aliasing_names( |
| 641 | funcdef: ast.FunctionDef, seed: set[str], module: nn.Module |
| 642 | ) -> set[str]: |
| 643 | """`seed` plus every name that may alias one of them. |
| 644 | |
| 645 | A view shares storage with its base (`h = x.view(-1)`), so mutating `h` |
| 646 | mutates `x`; tracking only `seed` would miss that. Any assignment reading a |
| 647 | tracked name therefore taints its targets, transitively, unless its value is |
| 648 | freshly allocated (`_returns_fresh`). |
| 649 | """ |
| 650 | names = set(seed) |
| 651 | assigns = [ |
| 652 | (node, node.value) |
| 653 | for node in ast.walk(funcdef) |
| 654 | if isinstance(node, (ast.Assign, ast.AnnAssign)) |
| 655 | and node.value is not None |
| 656 | and not _returns_fresh(node.value, module) |
| 657 | ] |
| 658 | changed = True |
| 659 | while changed: |
| 660 | changed = False |
| 661 | for node, value in assigns: |
| 662 | if not _aliasing_reads(value) & names: |
| 663 | continue |
| 664 | targets = node.targets if isinstance(node, ast.Assign) else [node.target] |
| 665 | stores = { |
| 666 | child.id |
| 667 | for target in targets |
| 668 | for child in ast.walk(target) |
| 669 | if isinstance(child, ast.Name) and isinstance(child.ctx, ast.Store) |
| 670 | } |
| 671 | if stores - names: |
| 672 | names |= stores |
| 673 | changed = True |
| 674 | return names |
| 675 | |
| 676 | |
| 677 | def replace_expr(module: ast.AST, old: ast.expr, new: ast.expr) -> None: |
| 678 | """Replace the expression `old` (by identity) with `new` within `module`.""" |
| 679 | |
| 680 | class _Replacer(ast.NodeTransformer): |
| 681 | def visit(self, node: ast.AST) -> ast.AST: |
| 682 | if node is old: |
| 683 | return new |
| 684 | return super().generic_visit(node) |
| 685 | |
| 686 | _Replacer().visit(module) |
| 687 | |
| 688 | |
| 689 | def find_node(graph: fx.Graph, predicate: Callable[[fx.Node], bool]) -> fx.Node | None: |
| 690 | """The first node in `graph` matching `predicate`, or `None`.""" |
| 691 | return next((n for n in graph.nodes if predicate(n)), None) |
| 692 | |
| 693 | |
| 694 | def output_value(graph: fx.Graph) -> object | None: |
| 695 | """The value the graph's `output` node returns, if the trace reached one.""" |
| 696 | output = find_node(graph, lambda n: n.op == "output") |
| 697 | if output is None or not output.args: |
| 698 | return None |
| 699 | return output.args[0] |
| 700 | |
| 701 | |
| 702 | def upstream_linear(node: object, module: nn.Module) -> fx.Node | None: |
| 703 | """Nearest linear producing `node`, walking back through splits/reshapes. |
| 704 | |
| 705 | Non-linear submodules are transparent too (e.g. the dropout GPT-style |
| 706 | attentions apply after their output projection). Never walks through a leaf |
| 707 | call (e.g. an attention interface): its inputs are what attention consumes, |
| 708 | not what produced the value.""" |
| 709 | stack = [node] |
| 710 | seen: set[fx.Node] = set() |
| 711 | while stack: |
| 712 | current = stack.pop() |
| 713 | if not isinstance(current, fx.Node) or current in seen: |
| 714 | continue |
| 715 | seen.add(current) |
| 716 | if is_linear(current, module): |
| 717 | return current |
| 718 | if current.op in ( |
| 719 | "call_function", |
| 720 | "call_method", |
| 721 | "call_module", |
| 722 | ) and not is_leaf_call(current): |
| 723 | stack.extend(current.args) |
| 724 | return None |
| 725 | |
| 726 | |
| 727 | def downstream_linear(node: fx.Node, module: nn.Module) -> fx.Node | None: |
| 728 | """Nearest linear consuming `node`'s output, walking through casts/scalings. |
| 729 | |
| 730 | Never walks through a leaf call (e.g. an attention interface): what crosses |
| 731 | it is consumed by the attention computation, not projected.""" |
| 732 | queue = list(node.users) |
| 733 | seen: set[fx.Node] = set() |
| 734 | while queue: |
| 735 | current = queue.pop(0) |
| 736 | if current in seen: |
| 737 | continue |
| 738 | seen.add(current) |
| 739 | if is_linear(current, module): |
| 740 | return current |
| 741 | if current.op in ("call_function", "call_method") and not is_leaf_call(current): |
| 742 | queue.extend(current.users) |
| 743 | return None |
| 744 | |
| 745 | |
| 746 | def returned_linear(graph: fx.Graph, module: nn.Module) -> str | None: |
| 747 | """Name of the Linear producing the graph's (first) output value.""" |
| 748 | value = output_value(graph) |
| 749 | if isinstance(value, (tuple, list)) and value: |
| 750 | value = value[0] |
| 751 | linear = upstream_linear(value, module) |
| 752 | return None if linear is None else str(linear.target) |
| 753 | |
| 754 | |
| 755 | def is_linear(node: fx.Node, module: nn.Module) -> bool: |
| 756 | """Is node `nn.Linear.__call__()`.""" |
| 757 | return node.op == "call_module" and isinstance( |
| 758 | module.get_submodule(node.target), nn.Linear |
| 759 | ) |
| 760 | |
| 761 | |
| 762 | _DTYPE_CASTS = frozenset({"to", "float", "double", "half", "bfloat16", "type_as"}) |
| 763 | |
| 764 | |
| 765 | def peel(node: object) -> object: |
| 766 | """Strip dtype-cast wrappers (`.to(...)`, `.float()`, `.type_as(...)`).""" |
| 767 | while ( |
| 768 | isinstance(node, fx.Node) |
| 769 | and node.op == "call_method" |
| 770 | and node.target in _DTYPE_CASTS |
| 771 | ): |
| 772 | node = node.args[0] |
| 773 | return node |
| 774 | |
| 775 | |
| 776 | def is_fn(node: object, target: object) -> TypeIs[fx.Node]: |
| 777 | """Is node `<target>()`.""" |
| 778 | return ( |
| 779 | isinstance(node, fx.Node) |
| 780 | and node.op == "call_function" |
| 781 | and node.target is target |
| 782 | ) |
| 783 | |
| 784 | |
| 785 | def is_method(node: object, name: str) -> TypeIs[fx.Node]: |
| 786 | """Is node `.<name>()`.""" |
| 787 | return ( |
| 788 | isinstance(node, fx.Node) and node.op == "call_method" and node.target == name |
| 789 | ) |
| 790 | |
| 791 | |
| 792 | def is_op(node: object, name: str) -> TypeIs[fx.Node]: |
| 793 | """Is node `<mod>.<name>()` for torch, F, operator, or Tensor.""" |
| 794 | return any( |
| 795 | is_fn(node, getattr(module, name, None)) for module in (torch, F, operator) |
| 796 | ) or (hasattr(torch.Tensor, name) and is_method(node, name)) |