vllm.model_executor.models.transformers.fx_utils ¶

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 = [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

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

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

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

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

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

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

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)

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
    }

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

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

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

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("__")

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
        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 in assigns:
            if not _aliasing_reads(node.value) & names:
                continue
            targets = getattr(node, "targets", None) or [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

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 []

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

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]

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

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)

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)

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 {}

def is_fn(node: object, target: Callable) -> bool:
    """Is node `<target>()`."""
    return (
        isinstance(node, fx.Node)
        and node.op == "call_function"
        and node.target is target
    )

def is_leaf_call(node: object) -> bool:
    """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)

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
    )

def is_method(node: object, name: str) -> bool:
    """Is node `.<name>()`."""
    return (
        isinstance(node, fx.Node) and node.op == "call_method" and node.target == name
    )

def is_op(node: object, name: str) -> bool:
    """
    Is node `torch.<name>()`, `F.<name>()`, `operator.<name>()`, or `Tensor.<name>()`.
    """
    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))

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]

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

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

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)

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)

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