Skip to content

vllm.model_executor.models.transformers.fx_utils

fx tracing and forward-source rewriting for the Transformers backend fusers.

A small engine, independent of any particular pattern: trace a module's forward with torch.fx (tolerating a partial graph), inspect the resulting nodes, and rewrite the forward's source (AST) so only matched calls change while the rest stays live Python. fusion.py builds the concrete fusion patterns on top.

Functions:

  • aliasing_names –

    seed plus every name that may alias one of them.

  • block_chain –

    Path of (statement list, index) pairs from block down to node.

  • bypass_existence_guard –

    Fold a guard on self.<name>'s existence to its constant value.

  • compile_forward –

    Compile funcdef in fn's module so tracebacks point at the source.

  • downstream_linear –

    Nearest linear consuming node's output, walking through casts/scalings.

  • find_node –

    The first node in graph matching predicate, or None.

  • forward_input_count –

    The number of tensor inputs cls.forward declares, excluding self and

  • forward_parameters –

    cls.forward's signature parameters, or empty if uninspectable.

  • is_fn –

    Is node <target>().

  • is_leaf_call –

    Is node a call recorded by _as_leaf_call (e.g. an attention interface).

  • is_linear –

    Is node nn.Linear.__call__().

  • is_method –

    Is node .<name>().

  • is_op –

    Is node <mod>.<name>() for torch, F, operator, or Tensor.

  • output_value –

    The value the graph's output node returns, if the trace reached one.

  • peel –

    Strip dtype-cast wrappers (.to(...), .float(), .type_as(...)).

  • recover_forward –

    Parse the source of cls.forward, ready for rewriting.

  • replace_expr –

    Replace the expression old (by identity) with new within module.

  • returned_linear –

    Name of the Linear producing the graph's (first) output value.

  • self_call_and_refs –

    The unique self.<name>(arg) call, plus every other self.<name> reference.

  • single_self_call –

    The unique self.<name>(arg) call in funcdef.

  • trace –

    Trace module.forward, returning the partial graph on failure.

  • upstream_linear –

    Nearest linear producing node, walking back through splits/reshapes.

  • written_names –

    Names region may rebind or write through.

_MODULE_CALL = nn.Module.__call__ module-attribute

The unpatched nn.Module.__call__. During tracing fx patches it to record call_module nodes; meta execution must call modules for real.

_UNKNOWN = object() module-attribute

Sentinel meta value for proxies whose concrete value could not be inferred. Distinct from None, which is a valid concrete value (e.g. attn_weights).

_AllLeafTracer

Bases: Tracer

Tracer that treats every submodule as a leaf.

Each child stays one call_module node, so matching sees the module's own forward structure (activations aren't decomposed into e.g. sigmoid * x). Every traced op is also executed on meta tensors (see _MetaProxy) so shape unpacks and *-splats trace through; anything else untraceable ends the trace early and the partial graph is matched.

Attributes:

  • varkw (str | None) –

    Name of the traced forward's **kwargs parameter, if any.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
class _AllLeafTracer(fx.Tracer):
    """Tracer that treats every submodule as a leaf.

    Each child stays one `call_module` node, so matching sees the module's own
    forward structure (activations aren't decomposed into e.g. `sigmoid * x`).
    Every traced op is also executed on meta tensors (see `_MetaProxy`) so
    shape unpacks and `*`-splats trace through; anything else untraceable ends
    the trace early and the partial graph is matched.
    """

    varkw: str | None = None
    """Name of the traced forward's `**kwargs` parameter, if any."""

    def is_leaf_module(self, m: nn.Module, module_qualified_name: str) -> bool:
        return True

    def proxy(self, node: fx.Node) -> fx.Proxy:
        return _MetaProxy(node, self)

    def create_proxy(self, kind, target, args, kwargs, *extra, **extra_kwargs):
        proxy = super().create_proxy(kind, target, args, kwargs, *extra, **extra_kwargs)
        if isinstance(proxy, _MetaProxy) and proxy.meta is _UNKNOWN:
            # A failure stays _UNKNOWN; only fatal if a shape question is asked.
            with contextlib.suppress(Exception):
                proxy.meta = self._infer_meta(kind, target, args, kwargs)
        return proxy

    def _infer_meta(self, kind: str, target: object, args: tuple, kwargs: dict):
        """Execute the op on meta tensors; PyTorch infers the output value."""
        if kind == "placeholder":
            # vLLM always feeds the model [1, seq_len, hidden_size] hidden states.
            weight = _reference_weight(self.root)
            if str(target) == "hidden_states" and weight is not None:
                return torch.empty(
                    1, 8, weight.shape[-1], dtype=weight.dtype, device="meta"
                )
            return _UNKNOWN
        if kind == "get_attr":
            value = operator.attrgetter(str(target))(self.root)
            if isinstance(value, torch.Tensor):
                value = torch.empty_like(value, device="meta")
            return value
        unknown = False

        def meta_of(arg: object) -> object:
            nonlocal unknown
            if isinstance(arg, fx.Proxy):
                meta = getattr(arg, "meta", _UNKNOWN)
                unknown = unknown or meta is _UNKNOWN
                return meta
            return arg

        meta_args = fx.node.map_aggregate(args, meta_of)
        meta_kwargs = fx.node.map_aggregate(kwargs, meta_of)
        if unknown:
            return _UNKNOWN
        if kind == "call_function":
            assert callable(target)
            return target(*meta_args, **meta_kwargs)
        if kind == "call_method":
            receiver, *rest = meta_args
            return getattr(receiver, str(target))(*rest, **meta_kwargs)
        if kind == "call_module":
            # Run the child's forward with all its state on "meta", without
            # mutating it (at match time params may be meta but buffers real).
            # fx patches `nn.Module.__call__` while tracing; restore the real
            # one so this execution is not itself recorded.
            child = self.root.get_submodule(str(target))
            state = {
                name: torch.empty_like(tensor, device="meta")
                for name, tensor in chain(
                    child.named_parameters(), child.named_buffers()
                )
            }
            with mock.patch.object(nn.Module, "__call__", _MODULE_CALL):
                return torch.func.functional_call(child, state, meta_args, meta_kwargs)
        return _UNKNOWN

    def _is_varkw(self, node: object) -> bool:
        return (
            isinstance(node, fx.Node)
            and node.op == "placeholder"
            and str(node.target).lstrip("*") == self.varkw
        )

    def iter(self, obj: fx.Proxy):
        # Assume kwargs is always empty to simplify tracing.
        node = obj.node
        if self._is_varkw(node) or (
            node.op == "call_method"
            and node.target == "keys"
            and self._is_varkw(node.args[0])
        ):
            return iter(())
        meta = getattr(obj, "meta", _UNKNOWN)
        if meta is _UNKNOWN:
            return super().iter(obj)
        return iter([obj[i] for i in range(len(meta))])

varkw = None class-attribute instance-attribute

Name of the traced forward's **kwargs parameter, if any.

_infer_meta(kind, target, args, kwargs)

Execute the op on meta tensors; PyTorch infers the output value.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def _infer_meta(self, kind: str, target: object, args: tuple, kwargs: dict):
    """Execute the op on meta tensors; PyTorch infers the output value."""
    if kind == "placeholder":
        # vLLM always feeds the model [1, seq_len, hidden_size] hidden states.
        weight = _reference_weight(self.root)
        if str(target) == "hidden_states" and weight is not None:
            return torch.empty(
                1, 8, weight.shape[-1], dtype=weight.dtype, device="meta"
            )
        return _UNKNOWN
    if kind == "get_attr":
        value = operator.attrgetter(str(target))(self.root)
        if isinstance(value, torch.Tensor):
            value = torch.empty_like(value, device="meta")
        return value
    unknown = False

    def meta_of(arg: object) -> object:
        nonlocal unknown
        if isinstance(arg, fx.Proxy):
            meta = getattr(arg, "meta", _UNKNOWN)
            unknown = unknown or meta is _UNKNOWN
            return meta
        return arg

    meta_args = fx.node.map_aggregate(args, meta_of)
    meta_kwargs = fx.node.map_aggregate(kwargs, meta_of)
    if unknown:
        return _UNKNOWN
    if kind == "call_function":
        assert callable(target)
        return target(*meta_args, **meta_kwargs)
    if kind == "call_method":
        receiver, *rest = meta_args
        return getattr(receiver, str(target))(*rest, **meta_kwargs)
    if kind == "call_module":
        # Run the child's forward with all its state on "meta", without
        # mutating it (at match time params may be meta but buffers real).
        # fx patches `nn.Module.__call__` while tracing; restore the real
        # one so this execution is not itself recorded.
        child = self.root.get_submodule(str(target))
        state = {
            name: torch.empty_like(tensor, device="meta")
            for name, tensor in chain(
                child.named_parameters(), child.named_buffers()
            )
        }
        with mock.patch.object(nn.Module, "__call__", _MODULE_CALL):
            return torch.func.functional_call(child, state, meta_args, meta_kwargs)
    return _UNKNOWN

_MetaAttribute

Bases: _MetaProxy, Attribute

Attribute proxy (e.g. x.shape) carrying its meta value.

Proxy.__getattr__ constructs Attribute directly, bypassing Tracer.proxy, so the meta value must be grafted on here too.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
83
84
85
86
87
88
89
90
91
92
93
94
class _MetaAttribute(_MetaProxy, fx.proxy.Attribute):
    """Attribute proxy (e.g. `x.shape`) carrying its meta value.

    `Proxy.__getattr__` constructs `Attribute` directly, bypassing
    `Tracer.proxy`, so the meta value must be grafted on here too."""

    def __init__(self, root: fx.Proxy, attr: str):
        super().__init__(root, attr)
        root_meta = getattr(root, "meta", _UNKNOWN)
        if root_meta is not _UNKNOWN:
            with contextlib.suppress(Exception):
                self.meta = getattr(root_meta, attr)

_MetaProxy

Bases: Proxy

Proxy carrying the meta-tensor value of the traced expression.

Shape questions (len, iteration, .shape unpacks) are answered by executing each op on the meta values, so PyTorch's meta kernels are the single source of shape inference — no per-op rules.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
class _MetaProxy(fx.Proxy):
    """Proxy carrying the meta-tensor value of the traced expression.

    Shape questions (`len`, iteration, `.shape` unpacks) are answered by
    executing each op on the meta values, so PyTorch's meta kernels are the
    single source of shape inference — no per-op rules."""

    meta: Any = _UNKNOWN

    def __len__(self) -> int:
        if self.meta is not _UNKNOWN:
            return len(self.meta)
        return super().__len__()  # type: ignore[misc]

    def __getattr__(self, k: str) -> "_MetaAttribute":
        return _MetaAttribute(self, k)

    def __setitem__(self, key: object, value: object) -> None:
        # Record `obj[key] = value` so tracing continues past in-place writes
        # into containers (e.g. Gemma 4's `shared_kv_states[...] = ...`).
        self.tracer.create_proxy(
            "call_function", operator.setitem, (self, key, value), {}
        )

_aliasing_reads(node)

Names node reads in a way that could yield a view of them.

Metadata (x.shape, x.size(0), x.dtype) is ints and tuples sharing no storage with x, so reaching a name only through metadata cannot alias it. Shape arithmetic is pervasive between projections, so propagating through it would taint nearly every later name.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
def _aliasing_reads(node: ast.expr) -> set[str]:
    """Names `node` reads in a way that could yield a view of them.

    Metadata (`x.shape`, `x.size(0)`, `x.dtype`) is ints and tuples sharing no
    storage with `x`, so reaching a name only through metadata cannot alias it.
    Shape arithmetic is pervasive between projections, so propagating through it
    would taint nearly every later name.
    """
    names: set[str] = set()
    stack: list[ast.AST] = [node]
    while stack:
        current = stack.pop()
        if isinstance(current, ast.Attribute) and current.attr in _METADATA_ATTRS:
            continue
        if (
            isinstance(current, ast.Call)
            and isinstance(current.func, ast.Attribute)
            and current.func.attr in _METADATA_ATTRS
        ):
            stack.extend(current.args)
            stack.extend(keyword.value for keyword in current.keywords)
            continue
        if isinstance(current, ast.Name) and isinstance(current.ctx, ast.Load):
            names.add(current.id)
        stack.extend(ast.iter_child_nodes(current))
    return names

_arg_names(call)

Every name call reads through its arguments.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
519
520
521
522
523
524
525
526
527
def _arg_names(call: ast.Call) -> set[str]:
    """Every name `call` reads through its arguments."""
    reads = [*call.args, *(kw.value for kw in call.keywords)]
    return {
        node.id
        for arg in reads
        for node in ast.walk(arg)
        if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load)
    }

_as_leaf_call(fn, length=None)

Wrap any callable so tracing records it as one opaque call_function node.

Lets the trace continue past untraceable bodies. Only the proxy arguments carry into the node's dataflow; the rest are dropped rather than lifted into the graph. length declares how many values the callable returns, so unpacking its result also traces. Called without proxies (i.e. outside tracing), the wrapper is a passthrough.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
def _as_leaf_call(fn: Callable, length: int | None = None) -> Callable:
    """Wrap any callable so tracing records it as one opaque `call_function` node.

    Lets the trace continue past untraceable bodies. Only the proxy arguments carry into
    the node's dataflow; the rest are dropped rather than lifted into the graph.
    `length` declares how many values the callable returns, so unpacking its result also
    traces. Called without proxies (i.e. outside tracing), the wrapper is a passthrough.
    """

    def leaf(*args, **kwargs):
        proxies = tuple(arg for arg in args if isinstance(arg, fx.Proxy))
        if not proxies:
            return fn(*args, **kwargs)
        proxy = proxies[0].tracer.create_proxy("call_function", fn, proxies, {})
        proxy.node.meta["leaf_call"] = True
        if length is not None:
            # The body never executes, so fabricate a value of the declared length.
            proxy.meta = (_UNKNOWN,) * length
        return proxy

    return leaf

_base_name(node)

The root Name of an attribute/subscript chain (a.b[c] -> a).

Source code in vllm/model_executor/models/transformers/fx_utils.py 
463
464
465
466
467
def _base_name(node: ast.expr) -> str | None:
    """The root `Name` of an attribute/subscript chain (`a.b[c]` -> `a`)."""
    while isinstance(node, (ast.Attribute, ast.Subscript)):
        node = node.value
    return node.id if isinstance(node, ast.Name) else None

_call_name(call)

The called function's own name, without its qualifier.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
512
513
514
515
516
def _call_name(call: ast.Call) -> str:
    """The called function's own name, without its qualifier."""
    if isinstance(call.func, ast.Attribute):
        return call.func.attr
    return call.func.id if isinstance(call.func, ast.Name) else ""

_in_boolean_context(funcdef, ref)

Is ref used only for its truth value (a test, or a not operand)?

In these positions the object's identity never escapes, so a reference that is always truthy can be replaced by True. and/or are excluded: they yield an operand, so the module could escape (x and self.<name>).

Source code in vllm/model_executor/models/transformers/fx_utils.py 
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
def _in_boolean_context(funcdef: ast.FunctionDef, ref: ast.expr) -> bool:
    """Is `ref` used only for its truth value (a test, or a `not` operand)?

    In these positions the object's identity never escapes, so a reference that
    is always truthy can be replaced by `True`. `and`/`or` are excluded: they
    yield an operand, so the module could escape (`x and self.<name>`)."""
    for node in ast.walk(funcdef):
        if (
            isinstance(node, (ast.If, ast.IfExp, ast.While, ast.Assert))
            and node.test is ref
        ):
            return True
        if (
            isinstance(node, ast.UnaryOp)
            and isinstance(node.op, ast.Not)
            and node.operand is ref
        ):
            return True
    return False

_inplace_target(node)

Base name node mutates in place, if any.

A write through the name (x[i] = ..., x.attr = ...) or an in-place method call (x.mul_(...)) leaves the base name in a Load context, so it is not a plain Name store (see _rebound_names).

Source code in vllm/model_executor/models/transformers/fx_utils.py 
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
def _inplace_target(node: ast.AST) -> str | None:
    """Base name `node` mutates in place, if any.

    A write through the name (`x[i] = ...`, `x.attr = ...`) or an in-place method
    call (`x.mul_(...)`) leaves the base name in a `Load` context, so it is not a
    plain `Name` store (see `_rebound_names`)."""
    if isinstance(node, (ast.Attribute, ast.Subscript)) and isinstance(
        node.ctx, (ast.Store, ast.Del)
    ):
        return _base_name(node)
    if (
        isinstance(node, ast.Call)
        and isinstance(node.func, ast.Attribute)
        and node.func.attr.endswith("_")
        and not node.func.attr.endswith("__")
    ):
        return _base_name(node.func.value)
    return None

_leaf_attention_interfaces()

Patch AttentionInterface.get_interface so traced forwards see a leaf node.

vllm_attention_function needs runtime context so it is untraceable. Every interface returns (attn_output, attn_weights).

Source code in vllm/model_executor/models/transformers/fx_utils.py 
220
221
222
223
224
225
226
227
228
229
230
231
232
def _leaf_attention_interfaces():
    """Patch `AttentionInterface.get_interface` so traced forwards see a leaf node.

    `vllm_attention_function` needs runtime context so it is untraceable.
    Every interface returns `(attn_output, attn_weights)`."""
    from transformers.modeling_utils import AttentionInterface

    original = AttentionInterface.get_interface

    def get_interface(self, *args, **kwargs):
        return _as_leaf_call(original(self, *args, **kwargs), length=2)

    return mock.patch.object(AttentionInterface, "get_interface", get_interface)

_mutated_names(region)

Names mutated in place within region (see _inplace_target).

Source code in vllm/model_executor/models/transformers/fx_utils.py 
502
503
504
505
506
507
508
509
def _mutated_names(region: list[ast.stmt]) -> set[str]:
    """Names mutated in place within `region` (see `_inplace_target`)."""
    return {
        base
        for stmt in region
        for node in ast.walk(stmt)
        if (base := _inplace_target(node)) is not None
    }

_rebound_names(region)

Names rebound outright within region (x = ..., del x).

These are Name nodes in a Store/Del context.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
470
471
472
473
474
475
476
477
478
479
def _rebound_names(region: list[ast.stmt]) -> set[str]:
    """Names rebound outright within `region` (`x = ...`, `del x`).

    These are `Name` nodes in a `Store`/`Del` context."""
    return {
        node.id
        for stmt in region
        for node in ast.walk(stmt)
        if isinstance(node, ast.Name) and isinstance(node.ctx, (ast.Store, ast.Del))
    }

_reference_weight(module)

A weight whose trailing dim is the module's hidden size.

Linears and 2-D gate weights are [out, hidden]; norm weights are [hidden]. Used to fabricate a placeholder input of matching size/dtype.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
44
45
46
47
48
49
50
51
52
53
54
55
def _reference_weight(module: nn.Module) -> torch.Tensor | None:
    """A weight whose trailing dim is the module's hidden size.

    Linears and 2-D gate weights are `[out, hidden]`; norm weights are
    `[hidden]`. Used to fabricate a placeholder input of matching size/dtype."""
    for child in module.modules():
        if isinstance(child, nn.Linear):
            return child.weight
    for param in module.parameters():
        if param.ndim in (1, 2):
            return param
    return None

_returns_fresh(node, module)

Does evaluating node allocate, rather than return a view of its input?

A nn.Linear writes its output into new storage, so its result aliases nothing that went in, and a chain of view ops on top of it aliases only that fresh tensor. Walking the receiver chain down to a self.<linear>(...) call therefore proves the value cannot alias the projection input. A literal aliases nothing either, and a branch or sequence is fresh when every part is -- which is how a guarded projection (p(x) if p is not None else None) stays untracked. A call on an attribute that is missing or None cannot run at all, so a branch selecting a different configuration is fresh too.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
def _returns_fresh(node: ast.expr, module: nn.Module) -> bool:
    """Does evaluating `node` allocate, rather than return a view of its input?

    A `nn.Linear` writes its output into new storage, so its result aliases
    nothing that went in, and a chain of view ops on top of it aliases only that
    fresh tensor. Walking the receiver chain down to a `self.<linear>(...)` call
    therefore proves the value cannot alias the projection input. A literal
    aliases nothing either, and a branch or sequence is fresh when every part is
    -- which is how a guarded projection (`p(x) if p is not None else None`)
    stays untracked. A call on an attribute that is missing or `None` cannot run
    at all, so a branch selecting a different configuration is fresh too.
    """
    if isinstance(node, ast.Constant):
        return True
    if isinstance(node, ast.IfExp):
        if isinstance(node.test, ast.Constant):
            # A folded existence guard leaves the untaken branch as dead code.
            taken = node.body if node.test.value else node.orelse
            return _returns_fresh(taken, module)
        return _returns_fresh(node.body, module) and _returns_fresh(node.orelse, module)
    if isinstance(node, (ast.Tuple, ast.List)):
        return all(_returns_fresh(elt, module) for elt in node.elts)
    while isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
        func = node.func
        if isinstance(func.value, ast.Name) and func.value.id == "self":
            attribute = getattr(module, func.attr, None)
            if attribute is None or isinstance(attribute, nn.Linear):
                # A missing or `None` attribute cannot be called, so the branch
                # selecting it never runs (`q_proj` on a q-LoRA layer).
                return True
            # Anything else can only return storage it was handed, so it is
            # fresh exactly when its inputs are (`q_norm(q_proj(x))`).
            return bool(node.args) and all(
                _returns_fresh(arg, module) for arg in node.args
            )
        node = func.value
    return False

_writes_through_args(call)

Does call write through an argument rather than return a new value?

Torch spells this three ways: an out= destination, inplace=True, and the trailing-underscore convention (torch.relu_(x)). The underscore is checked on the function's own name, so it catches the free-function form that _inplace_target cannot (there the base name is the module, torch).

Source code in vllm/model_executor/models/transformers/fx_utils.py 
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
def _writes_through_args(call: ast.Call) -> bool:
    """Does `call` write through an argument rather than return a new value?

    Torch spells this three ways: an `out=` destination, `inplace=True`, and the
    trailing-underscore convention (`torch.relu_(x)`). The underscore is checked
    on the function's own name, so it catches the free-function form that
    `_inplace_target` cannot (there the base name is the module, `torch`).
    """
    for keyword in call.keywords:
        if keyword.arg == "out":
            return True
        if keyword.arg == "inplace" and not (
            isinstance(keyword.value, ast.Constant) and keyword.value.value is False
        ):
            return True
    name = _call_name(call)
    return name.endswith("_") and not name.endswith("__")

aliasing_names(funcdef, seed, module)

seed plus every name that may alias one of them.

A view shares storage with its base (h = x.view(-1)), so mutating h mutates x; tracking only seed would miss that. Any assignment reading a tracked name therefore taints its targets, transitively, unless its value is freshly allocated (_returns_fresh).

Source code in vllm/model_executor/models/transformers/fx_utils.py 
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
def aliasing_names(
    funcdef: ast.FunctionDef, seed: set[str], module: nn.Module
) -> set[str]:
    """`seed` plus every name that may alias one of them.

    A view shares storage with its base (`h = x.view(-1)`), so mutating `h`
    mutates `x`; tracking only `seed` would miss that. Any assignment reading a
    tracked name therefore taints its targets, transitively, unless its value is
    freshly allocated (`_returns_fresh`).
    """
    names = set(seed)
    assigns = [
        (node, node.value)
        for node in ast.walk(funcdef)
        if isinstance(node, (ast.Assign, ast.AnnAssign))
        and node.value is not None
        and not _returns_fresh(node.value, module)
    ]
    changed = True
    while changed:
        changed = False
        for node, value in assigns:
            if not _aliasing_reads(value) & names:
                continue
            targets = node.targets if isinstance(node, ast.Assign) else [node.target]
            stores = {
                child.id
                for target in targets
                for child in ast.walk(target)
                if isinstance(child, ast.Name) and isinstance(child.ctx, ast.Store)
            }
            if stores - names:
                names |= stores
                changed = True
    return names

block_chain(block, node)

Path of (statement list, index) pairs from block down to node.

Each pair names a nested block and the index in it of the statement containing node; the last pair is node's innermost block. Empty if node is not in block. The longest common prefix of several nodes' chains is the innermost block that dominates them all.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
def block_chain(
    block: list[ast.stmt], node: ast.AST
) -> list[tuple[list[ast.stmt], int]]:
    """Path of (statement list, index) pairs from `block` down to `node`.

    Each pair names a nested block and the index in it of the statement
    containing `node`; the last pair is `node`'s innermost block. Empty if
    `node` is not in `block`. The longest common prefix of several nodes' chains
    is the innermost block that dominates them all.
    """
    for index, stmt in enumerate(block):
        if not any(child is node for child in ast.walk(stmt)):
            continue
        child_blocks = [
            getattr(stmt, fld, None) for fld in ("body", "orelse", "finalbody")
        ]
        child_blocks += [h.body for h in getattr(stmt, "handlers", [])]
        child_blocks += [c.body for c in getattr(stmt, "cases", [])]
        for child_block in child_blocks:
            if (
                isinstance(child_block, list)
                and child_block
                and (tail := block_chain(child_block, node))
            ):
                return [(block, index), *tail]
        return [(block, index)]
    return []

bypass_existence_guard(funcdef, ref, name)

Fold a guard on self.<name>'s existence to its constant value.

The fuser deletes self.<name> and binds only instances where it exists as a truthy nn.Linear (guaranteed by the match, the cache key, and validate), so a guard testing its presence is a fusion invariant. Two forms are folded:

  • an identity None check self.<name> is None -> False, self.<name> is not None -> True (either operand order). ==/!= are not folded: is_linear accepts nn.Linear subclasses, which may override __eq__/__ne__, so equality is not guaranteed to track identity.
  • a bare truthiness test (if self.<name>:, ... if self.<name> else ..., not self.<name>): the reference itself to True.

Any other surviving reference escapes the projection's value, which no longer exists after fusion, so refuse rather than change semantics.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
def bypass_existence_guard(
    funcdef: ast.FunctionDef, ref: ast.Attribute, name: str
) -> None:
    """Fold a guard on `self.<name>`'s existence to its constant value.

    The fuser deletes `self.<name>` and binds only instances where it exists as a
    truthy `nn.Linear` (guaranteed by the match, the cache key, and `validate`),
    so a guard testing its presence is a fusion invariant. Two forms are folded:

    - an identity `None` check `self.<name> is None` -> `False`,
      `self.<name> is not None` -> `True` (either operand order). `==`/`!=` are
      *not* folded: `is_linear` accepts `nn.Linear` subclasses, which may override
      `__eq__`/`__ne__`, so equality is not guaranteed to track identity.
    - a bare truthiness test (`if self.<name>:`, `... if self.<name> else ...`,
      `not self.<name>`): the reference itself to `True`.

    Any other surviving reference escapes the projection's value, which no longer
    exists after fusion, so refuse rather than change semantics."""
    for node in ast.walk(funcdef):
        # `self.<name> is (not) None`, either operand order.
        if not (isinstance(node, ast.Compare) and len(node.ops) == 1):
            continue
        (op,), (right,) = node.ops, node.comparators
        if not isinstance(op, (ast.Is, ast.IsNot)):
            continue
        if (ref is node.left and _is_none(right)) or (
            ref is right and _is_none(node.left)
        ):
            value = isinstance(op, ast.IsNot)
            replace_expr(funcdef, node, ast.copy_location(ast.Constant(value), node))
            return
    # A bare truthiness test (statement `if`, ternary, or `not`); an `nn.Linear`
    # is always truthy, so the reference folds to `True`.
    if _in_boolean_context(funcdef, ref):
        replace_expr(funcdef, ref, ast.copy_location(ast.Constant(True), ref))
        return
    raise ValueError(f"{name} is referenced outside an existence guard")

compile_forward(funcdef, fn)

Compile funcdef in fn's module so tracebacks point at the source.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
314
315
316
317
318
319
320
321
322
def compile_forward(funcdef: ast.FunctionDef, fn: Callable) -> Callable:
    """Compile `funcdef` in `fn`'s module so tracebacks point at the source."""
    module = ast.Module(body=[funcdef], type_ignores=[])
    ast.fix_missing_locations(module)
    ast.increment_lineno(module, fn.__code__.co_firstlineno - 1)
    code = compile(module, fn.__code__.co_filename, "exec")
    namespace: dict = {}
    exec(code, fn.__globals__, namespace)
    return namespace[funcdef.name]

downstream_linear(node, module)

Nearest linear consuming node's output, walking through casts/scalings.

Never walks through a leaf call (e.g. an attention interface): what crosses it is consumed by the attention computation, not projected.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
def downstream_linear(node: fx.Node, module: nn.Module) -> fx.Node | None:
    """Nearest linear consuming `node`'s output, walking through casts/scalings.

    Never walks through a leaf call (e.g. an attention interface): what crosses
    it is consumed by the attention computation, not projected."""
    queue = list(node.users)
    seen: set[fx.Node] = set()
    while queue:
        current = queue.pop(0)
        if current in seen:
            continue
        seen.add(current)
        if is_linear(current, module):
            return current
        if current.op in ("call_function", "call_method") and not is_leaf_call(current):
            queue.extend(current.users)
    return None

find_node(graph, predicate)

The first node in graph matching predicate, or None.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
689
690
691
def find_node(graph: fx.Graph, predicate: Callable[[fx.Node], bool]) -> fx.Node | None:
    """The first node in `graph` matching `predicate`, or `None`."""
    return next((n for n in graph.nodes if predicate(n)), None)

forward_input_count(cls)

The number of tensor inputs cls.forward declares, excluding self and any *args/**kwargs. Read from the signature, so it is independent of whether the trace completes (unlike counting placeholders).

Source code in vllm/model_executor/models/transformers/fx_utils.py 
299
300
301
302
303
304
305
306
307
308
309
310
311
def forward_input_count(cls: type[nn.Module]) -> int:
    """The number of tensor inputs `cls.forward` declares, excluding `self` and
    any `*args`/`**kwargs`. Read from the signature, so it is independent of
    whether the trace completes (unlike counting placeholders)."""
    params = list(forward_parameters(cls).values())
    if not params:
        return 1  # uninspectable: assume a single input and let matching decide
    fixed = (
        inspect.Parameter.POSITIONAL_ONLY,
        inspect.Parameter.POSITIONAL_OR_KEYWORD,
        inspect.Parameter.KEYWORD_ONLY,
    )
    return sum(1 for p in params[1:] if p.kind in fixed)

forward_parameters(cls)

cls.forward's signature parameters, or empty if uninspectable.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
291
292
293
294
295
296
def forward_parameters(cls: type[nn.Module]) -> dict[str, inspect.Parameter]:
    """`cls.forward`'s signature parameters, or empty if uninspectable."""
    try:
        return dict(inspect.signature(cls.forward).parameters)
    except (TypeError, ValueError):
        return {}

is_fn(node, target)

Is node <target>().

Source code in vllm/model_executor/models/transformers/fx_utils.py 
776
777
778
779
780
781
782
def is_fn(node: object, target: object) -> TypeIs[fx.Node]:
    """Is node `<target>()`."""
    return (
        isinstance(node, fx.Node)
        and node.op == "call_function"
        and node.target is target
    )

is_leaf_call(node)

Is node a call recorded by _as_leaf_call (e.g. an attention interface).

Source code in vllm/model_executor/models/transformers/fx_utils.py 
39
40
41
def is_leaf_call(node: object) -> TypeIs[fx.Node]:
    """Is node a call recorded by `_as_leaf_call` (e.g. an attention interface)."""
    return isinstance(node, fx.Node) and node.meta.get("leaf_call", False)

is_linear(node, module)

Is node nn.Linear.__call__().

Source code in vllm/model_executor/models/transformers/fx_utils.py 
755
756
757
758
759
def is_linear(node: fx.Node, module: nn.Module) -> bool:
    """Is node `nn.Linear.__call__()`."""
    return node.op == "call_module" and isinstance(
        module.get_submodule(node.target), nn.Linear
    )

is_method(node, name)

Is node .<name>().

Source code in vllm/model_executor/models/transformers/fx_utils.py 
785
786
787
788
789
def is_method(node: object, name: str) -> TypeIs[fx.Node]:
    """Is node `.<name>()`."""
    return (
        isinstance(node, fx.Node) and node.op == "call_method" and node.target == name
    )

is_op(node, name)

Is node <mod>.<name>() for torch, F, operator, or Tensor.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
792
793
794
795
796
def is_op(node: object, name: str) -> TypeIs[fx.Node]:
    """Is node `<mod>.<name>()` for torch, F, operator, or Tensor."""
    return any(
        is_fn(node, getattr(module, name, None)) for module in (torch, F, operator)
    ) or (hasattr(torch.Tensor, name) and is_method(node, name))

output_value(graph)

The value the graph's output node returns, if the trace reached one.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
694
695
696
697
698
699
def output_value(graph: fx.Graph) -> object | None:
    """The value the graph's `output` node returns, if the trace reached one."""
    output = find_node(graph, lambda n: n.op == "output")
    if output is None or not output.args:
        return None
    return output.args[0]

peel(node)

Strip dtype-cast wrappers (.to(...), .float(), .type_as(...)).

Source code in vllm/model_executor/models/transformers/fx_utils.py 
765
766
767
768
769
770
771
772
773
def peel(node: object) -> object:
    """Strip dtype-cast wrappers (`.to(...)`, `.float()`, `.type_as(...)`)."""
    while (
        isinstance(node, fx.Node)
        and node.op == "call_method"
        and node.target in _DTYPE_CASTS
    ):
        node = node.args[0]
    return node

recover_forward(cls)

Parse the source of cls.forward, ready for rewriting.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
def recover_forward(cls: type[nn.Module]) -> tuple[ast.FunctionDef, Callable]:
    """Parse the source of `cls.forward`, ready for rewriting."""
    fn = inspect.unwrap(cls.forward)
    if fn.__code__.co_freevars:
        raise ValueError("forward is a closure")
    tree = ast.parse(textwrap.dedent(inspect.getsource(fn)))
    funcdef = tree.body[0]
    if not isinstance(funcdef, ast.FunctionDef):
        raise ValueError("source is not a plain function definition")
    # `fn` is already unwrapped; don't re-apply its decorators
    funcdef.decorator_list.clear()
    # Annotations may not evaluate outside the defining module (e.g. with
    # postponed evaluation); they're not needed at runtime
    funcdef.returns = None
    args = funcdef.args
    for arg in (
        *args.posonlyargs,
        *args.args,
        *args.kwonlyargs,
        *filter(None, (args.vararg, args.kwarg)),
    ):
        arg.annotation = None
    # Recompiling outside the class body would break name mangling
    for node in ast.walk(funcdef):
        name = getattr(node, "attr", None) or getattr(node, "id", None)
        if name and name.startswith("__") and not name.endswith("__"):
            raise ValueError(f"{name} would be name mangled")
    return funcdef, fn

replace_expr(module, old, new)

Replace the expression old (by identity) with new within module.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
677
678
679
680
681
682
683
684
685
686
def replace_expr(module: ast.AST, old: ast.expr, new: ast.expr) -> None:
    """Replace the expression `old` (by identity) with `new` within `module`."""

    class _Replacer(ast.NodeTransformer):
        def visit(self, node: ast.AST) -> ast.AST:
            if node is old:
                return new
            return super().generic_visit(node)

    _Replacer().visit(module)

returned_linear(graph, module)

Name of the Linear producing the graph's (first) output value.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
746
747
748
749
750
751
752
def returned_linear(graph: fx.Graph, module: nn.Module) -> str | None:
    """Name of the Linear producing the graph's (first) output value."""
    value = output_value(graph)
    if isinstance(value, (tuple, list)) and value:
        value = value[0]
    linear = upstream_linear(value, module)
    return None if linear is None else str(linear.target)

self_call_and_refs(funcdef, name)

The unique self.<name>(arg) call, plus every other self.<name> reference.

Raises unless exactly one single-argument self.<name>(arg) call exists, so one fx call_module node maps to one syntactic call site. The remaining references (e.g. is not None guards) are returned for the caller to resolve.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
def self_call_and_refs(
    funcdef: ast.FunctionDef, name: str
) -> tuple[ast.Call, list[ast.Attribute]]:
    """The unique `self.<name>(arg)` call, plus every other `self.<name>` reference.

    Raises unless exactly one single-argument `self.<name>(arg)` call exists, so
    one fx `call_module` node maps to one syntactic call site. The remaining
    references (e.g. `is not None` guards) are returned for the caller to resolve.
    """
    uses = [
        node
        for node in ast.walk(funcdef)
        if isinstance(node, ast.Attribute)
        and node.attr == name
        and isinstance(node.value, ast.Name)
        and node.value.id == "self"
    ]
    calls = [
        node
        for node in ast.walk(funcdef)
        if isinstance(node, ast.Call)
        and any(node.func is use for use in uses)
        and len(node.args) == 1
        and not isinstance(node.args[0], ast.Starred)
        and not node.keywords
    ]
    if len(calls) != 1:
        raise ValueError(f"{name} is not called exactly once as self.{name}(arg)")
    call = calls[0]
    other_refs = [use for use in uses if use is not call.func]
    return call, other_refs

single_self_call(funcdef, name)

The unique self.<name>(arg) call in funcdef.

Raises unless name appears exactly once, as such a call, so the source rewrite agrees with the fx match.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
358
359
360
361
362
363
364
365
366
367
def single_self_call(funcdef: ast.FunctionDef, name: str) -> ast.Call:
    """The unique `self.<name>(arg)` call in `funcdef`.

    Raises unless `name` appears exactly once, as such a call, so the source
    rewrite agrees with the fx match.
    """
    call, other_refs = self_call_and_refs(funcdef, name)
    if other_refs:
        raise ValueError(f"{name} is referenced {len(other_refs) + 1} times")
    return call

trace(module)

Trace module.forward, returning the partial graph on failure.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
def trace(module: nn.Module) -> fx.Graph | None:
    """Trace `module.forward`, returning the partial graph on failure."""
    parameters = forward_parameters(type(module))
    # vLLM never passes `past_key_values` so it is always the default value of `None`.
    # Make this concrete to simplify tracing.
    concrete_args = None
    if "past_key_values" in parameters:
        concrete_args = {"past_key_values": None}
    # Get the name of the kwargs parameter passed to module.forward (usually "kwargs")
    tracer = _AllLeafTracer()
    tracer.varkw = next(
        (
            p.name
            for p in parameters.values()
            if p.kind is inspect.Parameter.VAR_KEYWORD
        ),
        None,
    )
    try:
        with _leaf_attention_interfaces():
            return tracer.trace(module, concrete_args=concrete_args)
    except Exception as exc:
        logger.debug("Could not fully trace %s: %s", type(module), exc)
        return getattr(tracer, "graph", None)

upstream_linear(node, module)

Nearest linear producing node, walking back through splits/reshapes.

Non-linear submodules are transparent too (e.g. the dropout GPT-style attentions apply after their output projection). Never walks through a leaf call (e.g. an attention interface): its inputs are what attention consumes, not what produced the value.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
def upstream_linear(node: object, module: nn.Module) -> fx.Node | None:
    """Nearest linear producing `node`, walking back through splits/reshapes.

    Non-linear submodules are transparent too (e.g. the dropout GPT-style
    attentions apply after their output projection). Never walks through a leaf
    call (e.g. an attention interface): its inputs are what attention consumes,
    not what produced the value."""
    stack = [node]
    seen: set[fx.Node] = set()
    while stack:
        current = stack.pop()
        if not isinstance(current, fx.Node) or current in seen:
            continue
        seen.add(current)
        if is_linear(current, module):
            return current
        if current.op in (
            "call_function",
            "call_method",
            "call_module",
        ) and not is_leaf_call(current):
            stack.extend(current.args)
    return None

written_names(region)

Names region may rebind or write through.

Beyond plain rebinding and writes through the name itself (_rebound_names, _mutated_names), a call can write through an argument without naming it as a target: F.relu(x, inplace=True), torch.add(x, y, out=x), torch.relu_(x). A statement-level call whose result is discarded is treated the same way, since it can only be there for a side effect.

Source code in vllm/model_executor/models/transformers/fx_utils.py 
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
def written_names(region: list[ast.stmt]) -> set[str]:
    """Names `region` may rebind or write through.

    Beyond plain rebinding and writes through the name itself (`_rebound_names`,
    `_mutated_names`), a call can write through an argument without naming it as
    a target: `F.relu(x, inplace=True)`, `torch.add(x, y, out=x)`,
    `torch.relu_(x)`. A statement-level call whose result is discarded is treated
    the same way, since it can only be there for a side effect.
    """
    written = _rebound_names(region) | _mutated_names(region)
    for stmt in region:
        if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Call):
            written |= _arg_names(stmt.value)
        for node in ast.walk(stmt):
            if isinstance(node, ast.Call) and _writes_through_args(node):
                written |= _arg_names(node)
    return written