[Transformers backend] Enable QKV-Fuser for Gemma4 (#55690)

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
Signed-off-by: Thomas Ortner <bohnstinglthomas@gmail.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

tests/models/transformers/fusers/test_linear.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Unit tests for the Transformers modeling backend's linear fusers."""
 
+import ast
 import inspect
 from types import MethodType, SimpleNamespace
 
@@ -19,7 +20,10 @@ from vllm.model_executor.models.transformers.fusers import (
     packed_qkv,
     qkv,
 )
-from vllm.model_executor.models.transformers.fx_utils import trace
+from vllm.model_executor.models.transformers.fx_utils import (
+    bypass_existence_guard,
+    trace,
+)
 
 
 class SiluAndMulStub(nn.Module):
@@ -202,6 +206,14 @@ class FourParallelLinears(nn.Module):
         return self.proj_a(x), self.proj_b(x), self.proj_c(x), self.proj_d(x)
 
 
+class GuardedParallelLinears(FourParallelLinears):
+    """A projection referenced by an existence guard as well as its call."""
+
+    def forward(self, x):
+        d = self.proj_d(x) if self.proj_d is not None else None
+        return self.proj_a(x), self.proj_b(x), self.proj_c(x), d
+
+
 class MutatingParallelLinears(FourParallelLinears):
     def forward(self, x):
         a = self.proj_a(x)
@@ -209,6 +221,34 @@ class MutatingParallelLinears(FourParallelLinears):
         return a, self.proj_b(x), self.proj_c(x), self.proj_d(x)
 
 
+class AliasMutatingParallelLinears(FourParallelLinears):
+    """A view of the input is mutated between the projections."""
+
+    def forward(self, x):
+        a = self.proj_a(x)
+        alias = x.view(-1)
+        alias.mul_(2)
+        return a, self.proj_b(x), self.proj_c(x), self.proj_d(x)
+
+
+class FunctionalInplaceParallelLinears(FourParallelLinears):
+    """The input is mutated by a functional in-place call between projections."""
+
+    def forward(self, x):
+        a = self.proj_a(x)
+        F.relu(x, inplace=True)
+        return a, self.proj_b(x), self.proj_c(x), self.proj_d(x)
+
+
+class UnderscoreInplaceParallelLinears(FourParallelLinears):
+    """The input is mutated by a free in-place function between projections."""
+
+    def forward(self, x):
+        a = self.proj_a(x)
+        torch.relu_(x)
+        return a, self.proj_b(x), self.proj_c(x), self.proj_d(x)
+
+
 class PlainQKV(nn.Module):
     """q/k/v returned as a bare tuple: no reshape stands between the calls and
     the return, so both QKVFuser and MergedColumnParallelFuser can rewrite it."""
@@ -294,6 +334,306 @@ class ResidDropoutAttention(FakeAttention):
         return self.resid_dropout(self.o_proj(attn_output)), None
 
 
+class GuardedVAttention(FakeAttention):
+    """Gemma 4-style: `v_proj` is called once but also read in a `None` guard.
+
+    The guard is a fusion invariant (the fuser only binds instances where the
+    projection exists), so it folds to `True` and the call still fuses."""
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        v = (
+            self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+            if self.v_proj is not None
+            else k
+        )
+        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, None
+        )
+        attn_output, _ = attention_interface(
+            self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
+        )
+        return self.o_proj(attn_output.reshape(*input_shape, -1).contiguous()), None
+
+
+class TruthyGuardAttention(FakeAttention):
+    """`v_proj` guarded by a bare truthiness test instead of `is not None`.
+
+    An `nn.Linear` is always truthy, so the guard folds to `True` and the call
+    still fuses (the projection is a proven invariant, exactly as for a `None`
+    guard)."""
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        v = (
+            self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+            if self.v_proj
+            else k
+        )
+        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, None
+        )
+        attn_output, _ = attention_interface(
+            self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
+        )
+        return self.o_proj(attn_output.reshape(*input_shape, -1).contiguous()), None
+
+
+class StatementGuardVAttention(FakeAttention):
+    """`v_proj` guarded by a statement-form `if self.v_proj:` block.
+
+    The `if`'s test is not an `ast.Compare`, so the identity-check scan skips it;
+    `_in_boolean_context` folds the test to `True` (an `nn.Linear` is always
+    truthy) and the call inside the now-`if True:` body still fuses, with the
+    live `else` surviving the rewrite."""
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        if self.v_proj:
+            v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        else:
+            v = k
+        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, None
+        )
+        attn_output, _ = attention_interface(
+            self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
+        )
+        return self.o_proj(attn_output.reshape(*input_shape, -1).contiguous()), None
+
+
+class BranchedKVAttention(FakeAttention):
+    """Gemma 4-style: q at body level, k/v inside an untaken `else` branch.
+
+    The three calls do not share a block, so the fused GEMM must be hoisted to
+    the innermost block that dominates all of them (the function body). The
+    live `if`/`else` must survive the rewrite."""
+
+    is_kv_shared_layer = False
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        if self.is_kv_shared_layer:
+            k = v = None
+        else:
+            k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+            v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, None
+        )
+        attn_output, _ = attention_interface(
+            self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
+        )
+        return self.o_proj(attn_output.reshape(*input_shape, -1).contiguous()), None
+
+
+class RebindArgAttention(FakeAttention):
+    """The projection input is rebound (in an untaken branch) before k/v.
+
+    The branch is dead at trace time so the match succeeds, but the source
+    rewrite must refuse: hoisting the GEMM above the rebind would feed k/v a
+    different input. Result: no fusion."""
+
+    recompute = False
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        if self.recompute:
+            hidden_states = hidden_states * 2
+        k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, None
+        )
+        attn_output, _ = attention_interface(
+            self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
+        )
+        return self.o_proj(attn_output.reshape(*input_shape, -1)), None
+
+
+class InPlaceMutatedArgAttention(FakeAttention):
+    """The shared input is mutated in place (in an untaken branch) before k/v.
+
+    Like `RebindArgAttention`, but the mutation keeps the name in a `Load`
+    context (`hidden_states.mul_(2)`), so a bare `Name`-store check would miss it.
+    The rewrite must still refuse: hoisting the GEMM above the mutation changes
+    the value k/v see. Result: no fusion."""
+
+    recompute = False
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        if self.recompute:
+            hidden_states.mul_(2)
+        k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, None
+        )
+        attn_output, _ = attention_interface(
+            self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
+        )
+        return self.o_proj(attn_output.reshape(*input_shape, -1)), None
+
+
+class AliasMutatedArgAttention(FakeAttention):
+    """A view of the shared input is mutated before k/v.
+
+    `hidden_states.view(-1)` shares storage with its base, so `alias.mul_(2)`
+    changes what k/v read without ever naming `hidden_states` as a target. Only
+    tracking names that may alias the input catches it. Result: no fusion."""
+
+    recompute = False
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        if self.recompute:
+            alias = hidden_states.view(-1)
+            alias.mul_(2)
+        k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, None
+        )
+        attn_output, _ = attention_interface(
+            self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
+        )
+        return self.o_proj(attn_output.reshape(*input_shape, -1)), None
+
+
+class FunctionalInplaceArgAttention(FakeAttention):
+    """The shared input is mutated by a functional in-place call before k/v.
+
+    `F.relu(hidden_states, inplace=True)` writes through an argument, so the
+    name appears only in a `Load` context inside a call that is not a method on
+    it -- invisible to a trailing-underscore method check. Result: no fusion."""
+
+    recompute = False
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        if self.recompute:
+            F.relu(hidden_states, inplace=True)
+        k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, None
+        )
+        attn_output, _ = attention_interface(
+            self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
+        )
+        return self.o_proj(attn_output.reshape(*input_shape, -1)), None
+
+
+class SubscriptMutatedArgAttention(FakeAttention):
+    """The shared input is written through a subscript (in an untaken branch).
+
+    `hidden_states[..., 0] = 0` leaves the name in a `Load` context on the
+    subscript's value, so it too evades a bare `Name`-store check. The rewrite
+    must refuse for the same reason. Result: no fusion."""
+
+    recompute = False
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        if self.recompute:
+            hidden_states[..., 0] = 0
+        k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, None
+        )
+        attn_output, _ = attention_interface(
+            self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
+        )
+        return self.o_proj(attn_output.reshape(*input_shape, -1)), None
+
+
+class NonFoldableRefAttention(FakeAttention):
+    """A surviving `v_proj` reference that is not an existence guard -> no fusion.
+
+    Deleting `v_proj` would break `isinstance(self.v_proj, nn.Linear)`: it reads
+    the projection's type, not just whether it exists, so the fuser cannot fold
+    it and must refuse rather than rewrite it."""
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        if isinstance(self.v_proj, nn.Linear):
+            v = v.contiguous()
+        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, None
+        )
+        attn_output, _ = attention_interface(
+            self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
+        )
+        return self.o_proj(attn_output.reshape(*input_shape, -1)), None
+
+
 class PackedQKVAttention(nn.Module):
     """GPTBigCode-style: one packed projection split into q/k/v in the forward."""
 
@@ -561,18 +901,246 @@ def test_qkv_identifies_output_projection():
         assert get_fuser(ResidDropoutAttention(), QKVFuser).o_name == "o_proj"
 
 
-def test_merged_column_fuser_rejects_input_mutation():
-    """The later projections must not be moved ahead of an input mutation."""
+@pytest.mark.parametrize(
+    "attn_cls",
+    [
+        GuardedVAttention,
+        TruthyGuardAttention,
+        StatementGuardVAttention,
+        BranchedKVAttention,
+    ],
+)
+def test_fuses_gemma4_qkv_obstacles(attn_cls):
+    """A guarded `v_proj` (obstacle 1) or branch-split k/v (obstacle 2) fuses.
+
+    Both patterns leave `v_proj` referenced beyond its single call, or split the
+    calls across blocks; the rewrite folds the guard / hoists the GEMM and the
+    numerics still match. The guard may be a `None` comparison or a bare
+    truthiness test (ternary or statement form); all are proven invariants, so
+    folding them is sound. The naive "relax the reference count" fix fuses the
+    guarded cases but silently drops the guard's semantics.
+    """
     with torch.device("meta"):
-        module = MutatingParallelLinears()
+        meta = attn_cls()
+    fuser = get_fuser(meta, QKVFuser)
+    assert isinstance(fuser, QKVFuser)
+    assert (fuser.q_name, fuser.o_name) == ("q_proj", "o_proj")
+
+    names = set(fuser.fused_forward.__code__.co_names)
+    assert "qkv_proj" in names
+    assert not {"q_proj", "k_proj", "v_proj"} & names
+    if attn_cls is BranchedKVAttention:
+        assert "is_kv_shared_layer" in names  # the live branch survives
+
+    real = attn_cls(kv_heads=4, layer_idx=3)
+    for p in real.parameters():
+        nn.init.normal_(p, std=0.05)
+    x = torch.randn(1, 5, 32)
+    real.attn = FakeSelfAttn()
+    expected, _ = real(x)
+    fused = _apply_qkv_fuser_with_stubs(real, fuser)
+    out, _ = fused(x)
+    torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)
+
+
+@pytest.mark.parametrize(
+    "attn_cls",
+    [
+        RebindArgAttention,
+        InPlaceMutatedArgAttention,
+        SubscriptMutatedArgAttention,
+        AliasMutatedArgAttention,
+        FunctionalInplaceArgAttention,
+        NonFoldableRefAttention,
+    ],
+)
+def test_qkv_refuses_unsound_rewrites(attn_cls):
+    """The match succeeds but the source rewrite must refuse, leaving no fuser.
+
+    `RebindArgAttention` rebinds the shared input before k/v; the `*Mutated*` and
+    `FunctionalInplace*` cases mutate it in place, directly (`x.mul_(...)`,
+    `x[...] = 0`), through a view that shares its storage, or through a call that
+    writes an argument (`F.relu(x, inplace=True)`) -- each would change what k/v
+    see if the GEMM were hoisted above them.
+    `NonFoldableRefAttention` reads `v_proj` outside an existence guard (the fuser
+    cannot fold it away). All fail closed in `update_forward`, so `get_fuser`
+    returns `None`.
+    """
+    with torch.device("meta"):
+        meta = attn_cls()
+    assert get_fuser(meta, QKVFuser) is None
+
+
+def _guard_funcdef(expr: str) -> tuple[ast.FunctionDef, ast.Attribute]:
+    """A `def f(self, a, b, T): return <expr>` and its `self.v_proj` reference."""
+    funcdef = ast.parse(f"def f(self, a, b, T):\n    return {expr}").body[0]
+    assert isinstance(funcdef, ast.FunctionDef)
+    ref = next(
+        node
+        for node in ast.walk(funcdef)
+        if isinstance(node, ast.Attribute) and node.attr == "v_proj"
+    )
+    return funcdef, ref
+
+
+@pytest.mark.parametrize(
+    "expr, expected",
+    [
+        ("a if self.v_proj is not None else b", True),
+        ("a if self.v_proj is None else b", False),
+        ("a if None is not self.v_proj else b", True),  # None on the left
+    ],
+)
+def test_bypass_existence_guard_folds_identity_none_check(expr, expected):
+    """`is`/`is not None`, either operand order, fold to a bool.
+
+    The projection is a proven invariant (it exists as an `nn.Linear`), so an
+    `is (not) None` guard has a constant truth value; folding it lets the call be
+    rewritten away without changing the guard's outcome. `==`/`!=` are excluded
+    (see the refusal test) because a subclass may override `__eq__`/`__ne__`."""
+    funcdef, ref = _guard_funcdef(expr)
+    bypass_existence_guard(funcdef, ref, "v_proj")
+    test = next(node for node in ast.walk(funcdef) if isinstance(node, ast.IfExp)).test
+    assert isinstance(test, ast.Constant) and test.value is expected
+
+
+@pytest.mark.parametrize(
+    "expr",
+    [
+        "a if self.v_proj else b",  # a ternary test
+        "not self.v_proj",  # a `not` operand
+    ],
+)
+def test_bypass_existence_guard_folds_bare_truthiness(expr):
+    """A bare truthiness test folds the reference itself to `True`.
+
+    An `nn.Linear` has no `__bool__`/`__len__`, so it is always truthy; in a pure
+    boolean position (a test, or a `not` operand) the reference can be replaced by
+    `True`, letting the projection be deleted."""
+    funcdef, ref = _guard_funcdef(expr)
+    bypass_existence_guard(funcdef, ref, "v_proj")
+    # No `self.v_proj` reference survives; a literal `True` took its place.
+    attrs = [n for n in ast.walk(funcdef) if isinstance(n, ast.Attribute)]
+    assert not any(n.attr == "v_proj" for n in attrs)
+    constants = [n for n in ast.walk(funcdef) if isinstance(n, ast.Constant)]
+    assert any(n.value is True for n in constants)
+
+
+def test_bypass_existence_guard_folds_bare_if_statement():
+    """A statement-form `if self.<name>:` folds its test to `True`.
+
+    The `len(node.ops) == 1` `Compare` scan skips it (it is not a comparison at
+    all); `_in_boolean_context` then matches the `if`'s test and folds it, so the
+    bare-truthiness statement form is handled, not silently dropped."""
+    funcdef = ast.parse(
+        "def f(self, a, b):\n    if self.v_proj:\n        return a\n    return b"
+    ).body[0]
+    assert isinstance(funcdef, ast.FunctionDef)
+    ref = next(
+        node
+        for node in ast.walk(funcdef)
+        if isinstance(node, ast.Attribute) and node.attr == "v_proj"
+    )
+    bypass_existence_guard(funcdef, ref, "v_proj")
+    test = next(node for node in ast.walk(funcdef) if isinstance(node, ast.If)).test
+    assert isinstance(test, ast.Constant) and test.value is True
+
+
+@pytest.mark.parametrize(
+    "expr",
+    [
+        "a if isinstance(self.v_proj, T) else b",  # a different predicate
+        "self.v_proj.weight",  # an attribute read, not a guard
+        "a and self.v_proj",  # `and` yields the operand, which may escape
+        "self.v_proj or a",  # `or` yields the module itself when truthy
+        "a if self.v_proj == None else b",  # `__eq__` may not track identity
+        "a if self.v_proj != None else b",  # `__ne__` may not track identity
+    ],
+)
+def test_bypass_existence_guard_refuses_non_guard_reference(expr):
+    """A surviving reference that is not an identity guard cannot be folded away.
+
+    Rewriting it is unsafe (its value depends on more than the projection's
+    existence), so the fuser refuses rather than change semantics. `==`/`!=` are
+    refused too: `is_linear` accepts `nn.Linear` subclasses that could override
+    `__eq__`/`__ne__`, so equality need not agree with `is (not) None`."""
+    funcdef, ref = _guard_funcdef(expr)
+    with pytest.raises(ValueError, match="outside an existence guard"):
+        bypass_existence_guard(funcdef, ref, "v_proj")
+
+
+@pytest.mark.parametrize("layer_idx", [0, 1])
+def test_fuses_real_gemma4_attention(layer_idx):
+    """The real Gemma 4 attention (both obstacles at once) fuses end to end."""
+    pytest.importorskip("transformers.models.gemma4")
+    from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
+    from transformers.models.gemma4.modeling_gemma4 import Gemma4TextAttention
+
+    config = Gemma4TextConfig(
+        hidden_size=64,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        num_hidden_layers=2,
+        head_dim=16,
+        intermediate_size=128,
+    )
+    with torch.device("meta"):
+        meta = Gemma4TextAttention(config, layer_idx=layer_idx)
+    fuser = get_fuser(meta, QKVFuser)
+    assert isinstance(fuser, QKVFuser)
+    assert (fuser.q_name, fuser.k_name, fuser.v_name) == (
+        "q_proj",
+        "k_proj",
+        "v_proj",
+    )
+    assert fuser.o_name == "o_proj"
+    names = set(fuser.fused_forward.__code__.co_names)
+    assert not {"q_proj", "k_proj", "v_proj"} & names
+
+
+def test_merged_column_fuser_folds_existence_guard():
+    """A guarded projection fuses: the merged linear replaces the deleted names.
+
+    Guard folding lives on `StackedFuser`, not just `QKVFuser`, because every
+    stacked fuser deletes the projections it merges -- a surviving reference to
+    one would fail at runtime."""
+    with torch.device("meta"):
+        module = GuardedParallelLinears()
     fuser = MergedColumnParallelFuser.match(trace(module), module)
     assert fuser is not None
-    with pytest.raises(ValueError, match="cross other operations"):
+    fuser.update_forward(module)
+    names = set(fuser.fused_forward.__code__.co_names)
+    assert fuser.merged_name in names
+    assert not {"proj_a", "proj_b", "proj_c", "proj_d"} & names
+
+
+@pytest.mark.parametrize(
+    "cls",
+    [
+        MutatingParallelLinears,
+        AliasMutatingParallelLinears,
+        FunctionalInplaceParallelLinears,
+        UnderscoreInplaceParallelLinears,
+    ],
+)
+def test_merged_column_fuser_rejects_input_mutation(cls):
+    """The later projections must not be moved ahead of an input mutation.
+
+    The mutation may be direct (`x.add_(1)`), through a view that shares the
+    input's storage, or through a call that writes an argument in place
+    (`F.relu(x, inplace=True)`, `torch.relu_(x)`). Each one silently changes the
+    fused result, so the rewrite must refuse rather than fuse."""
+    with torch.device("meta"):
+        module = cls()
+    fuser = MergedColumnParallelFuser.match(trace(module), module)
+    assert fuser is not None
+    with pytest.raises(ValueError, match="rebound or mutated"):
         fuser.update_forward(module)
 
 
+@pytest.mark.parametrize("cls", [FourParallelLinears, GuardedParallelLinears])
 def test_merged_column_fuser_supports_any_number_of_linears(
-    default_vllm_config, monkeypatch
+    cls, default_vllm_config, monkeypatch
 ):
     from vllm.model_executor import parameter
     from vllm.model_executor.layers import linear
@@ -585,19 +1153,14 @@ def test_merged_column_fuser_supports_any_number_of_linears(
     monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0)
     monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1)
     with torch.device("meta"):
-        meta = FourParallelLinears()
+        meta = cls()
     fuser = MergedColumnParallelFuser.match(trace(meta), meta)
     assert fuser is not None
-    assert fuser.linear_names == ("proj_a", "proj_b", "proj_c", "proj_d")
-    assert fuser.shards == [
-        ("proj_a", 0),
-        ("proj_b", 1),
-        ("proj_c", 2),
-        ("proj_d", 3),
-    ]
+    assert set(fuser.linear_names) == {"proj_a", "proj_b", "proj_c", "proj_d"}
+    assert fuser.shards == [(name, i) for i, name in enumerate(fuser.linear_names)]
     fuser.update_forward(meta)
 
-    real = FourParallelLinears()
+    real = cls()
     x = torch.randn(2, 8)
     expected = real(x)
     weights = [

tests/models/transformers/fusers/test_mla.py

@@ -19,9 +19,8 @@ from vllm.model_executor.models.transformers.fx_utils import trace
 _FUSED_QKV_A_PROJ = MLAFuser.merged_name
 
 
-def _match(q_lora_rank: int | None) -> MLAFuser | None:
-    """Match a meta `DeepseekV2Attention` directly (bypassing the per-class
-    `get_fuser` cache, so both q_lora variants of the same class are seen)."""
+def _attention(q_lora_rank: int | None):
+    """A meta `DeepseekV2Attention` in the requested q_lora configuration."""
     pytest.importorskip("transformers.models.deepseek_v2.modeling_deepseek_v2")
     from transformers.models.deepseek_v2.configuration_deepseek_v2 import (
         DeepseekV2Config,
@@ -41,7 +40,13 @@ def _match(q_lora_rank: int | None) -> MLAFuser | None:
         num_hidden_layers=1,
     )
     with torch.device("meta"):
-        attn = DeepseekV2Attention(cfg, layer_idx=0)
+        return DeepseekV2Attention(cfg, layer_idx=0)
+
+
+def _match(q_lora_rank: int | None) -> MLAFuser | None:
+    """Match a meta `DeepseekV2Attention` directly (bypassing the per-class
+    `get_fuser` cache, so both q_lora variants of the same class are seen)."""
+    attn = _attention(q_lora_rank)
     return MLAFuser.match(trace(attn), attn)
 
 
@@ -90,6 +95,23 @@ def test_q_lora_stacks_qkv_a_proj():
     }
 
 
+@pytest.mark.parametrize("q_lora_rank", [64, None])
+def test_update_forward_rewrites_real_attention(q_lora_rank):
+    """The source rewrite must survive the input-stability check.
+
+    The q-LoRA path hoists the merged down-projection to a top-level statement
+    while its calls sit inside the `q_lora_rank` branches, so the region it
+    spans is wide; a check that over-approximates aliasing would silently refuse
+    here and cost the fusion rather than fail loudly."""
+    fuser = _match(q_lora_rank)
+    assert isinstance(fuser, MLAFuser)
+    fuser.update_forward(_attention(q_lora_rank))
+    names = set(fuser.fused_forward.__code__.co_names)
+    if fuser.has_q_lora:
+        assert _FUSED_QKV_A_PROJ in names
+        assert not {"q_a_proj", "kv_a_proj_with_mqa"} & names
+
+
 class _Norm(nn.Module):
     """A real RMSNorm computation: `match` verifies chain norms via `RMSNormFuser`."""
 

vllm/model_executor/models/transformers/fusers/base.py

@@ -5,13 +5,19 @@
 import ast
 import types
 from abc import ABC, abstractmethod
-from collections.abc import Callable
+from collections.abc import Callable, Iterable
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, ClassVar
 
 from torch import fx, nn
 
-from vllm.model_executor.models.transformers.fx_utils import replace_expr
+from vllm.model_executor.models.transformers.fx_utils import (
+    aliasing_names,
+    bypass_existence_guard,
+    replace_expr,
+    self_call_and_refs,
+    written_names,
+)
 from vllm.model_executor.models.utils import ShardId, maybe_prefix
 
 if TYPE_CHECKING:
@@ -160,6 +166,53 @@ class StackedFuser(RewriteFuser):
         fused layer into its per-shard configs."""
         return {self.merged_name: [name for name, _ in self.shards]}
 
+    def _unguarded_calls(
+        self, funcdef: ast.FunctionDef, names: Iterable[str]
+    ) -> list[ast.Call]:
+        """One `self.<name>(arg)` call per projection, existence guards folded.
+
+        `update_attrs` deletes the projections it merges, so any reference to one beyond
+        its call site must be a guard on its existence; folding those to their constant
+        value keeps the rewritten forward off a name that no longer exists. A reference
+        that is not such a guard raises, so fusion is skipped."""
+        calls = []
+        for name in names:
+            call, refs = self_call_and_refs(funcdef, name)
+            for ref in refs:
+                bypass_existence_guard(funcdef, ref, name)
+            calls.append(call)
+        return calls
+
+    def _check_input_stable(
+        self,
+        funcdef: ast.FunctionDef,
+        module: nn.Module,
+        calls: list[ast.Call],
+        block: list[ast.stmt],
+        indices: list[int],
+    ) -> None:
+        """Raise unless hoisting the merged GEMM preserves what it reads.
+
+        Fusing moves every projection to one call at `min(indices)`, so the
+        merged GEMM reads the input once, up front, where the last of `calls`
+        would have read it later. That holds only if nothing in between changes
+        the input, and a change need not name it: writing any view that shares
+        its storage changes it too. Both halves of the check over-approximate,
+        since a false hit costs a fusion while a miss returns wrong numbers.
+        """
+        arg_names = {
+            node.id
+            for node in ast.walk(calls[0].args[0])
+            if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load)
+        }
+        # Protect the names the input reads and anything that may alias them.
+        tracked = aliasing_names(funcdef, arg_names, module)
+        # Any of them rebound, written through, or passed to a call that writes
+        # its argument would leave the hoisted GEMM reading a different value.
+        region = block[min(indices) : max(indices) + 1]
+        if tracked & written_names(region):
+            raise ValueError("projection input is rebound or mutated before all calls")
+
     def _splice_merged_split(
         self,
         funcdef: ast.FunctionDef,

vllm/model_executor/models/transformers/fusers/glu.py

@@ -137,8 +137,8 @@ class GLUFuser(MergedColumnParallelFuser):
         """Replace `act(gate(x)) * up(x)` with `act(gate_up(x))` in source."""
         funcdef, fn = recover_forward(type(module))
         act_call = single_self_call(funcdef, self.act_name)
-        gate_call = single_self_call(funcdef, self.gate_name)
-        up_call = single_self_call(funcdef, self.up_name)
+        calls = self._unguarded_calls(funcdef, (self.gate_name, self.up_name))
+        gate_call, up_call = calls
         if act_call.args[0] is not gate_call:
             raise ValueError("activation does not directly wrap the gate")
         if ast.dump(gate_call.args[0]) != ast.dump(up_call.args[0]):

vllm/model_executor/models/transformers/fusers/merged_column.py

@@ -12,11 +12,10 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import MergedColumnParallelLinear
 from vllm.model_executor.models.transformers.fusers.base import StackedFuser
 from vllm.model_executor.models.transformers.fx_utils import (
+    block_chain,
     compile_forward,
-    innermost_block,
     is_linear,
     recover_forward,
-    single_self_call,
 )
 from vllm.model_executor.models.utils import ShardId, maybe_prefix
 
@@ -92,30 +91,20 @@ class MergedColumnParallelFuser(StackedFuser):
     def update_forward(self, module: nn.Module) -> None:
         """Replace the parallel calls with one merged call and split."""
         funcdef, fn = recover_forward(type(module))
-        calls = [single_self_call(funcdef, name) for name in self.linear_names]
+        calls = self._unguarded_calls(funcdef, self.linear_names)
         if len(set(ast.dump(call.args[0]) for call in calls)) != 1:
             raise ValueError("parallel linears read different inputs")
-        blocks = [innermost_block(funcdef.body, call) for call in calls]
-        if any(found is None for found in blocks):
+        chains = [block_chain(funcdef.body, call) for call in calls]
+        if any(not chain for chain in chains):
             raise ValueError("parallel linear calls not found in the function body")
+        blocks = [chain[-1] for chain in chains]
         if len(set(id(block) for block, _ in blocks)) != 1:
             raise ValueError("parallel linear calls are in different blocks")
 
         block = blocks[0][0]
-        index = min(index for _, index in blocks)
-        calls_have_name_arg = all(isinstance(call.args[0], ast.Name) for call in calls)
-        # Moving projections must not cross operations that can change their input.
-        for statement in block[index : max(index for _, index in blocks) + 1]:
-            if not isinstance(statement, (ast.Assign, ast.Return)):
-                raise ValueError("parallel projections cross other operations")
-            if isinstance(statement, ast.Assign) and any(
-                not isinstance(target, ast.Name) for target in statement.targets
-            ):
-                raise ValueError("parallel projection assignment has side effects")
-            value = statement.value
-            values = value.elts if isinstance(value, ast.Tuple) else [value]
-            if any(value not in calls for value in values) or not calls_have_name_arg:
-                raise ValueError("parallel projections cross other operations")
+        indices = [index for _, index in blocks]
+        index = min(indices)
+        self._check_input_stable(funcdef, module, calls, block, indices)
 
         # l1(x), l2(x), ... -> merged(x).split(merged.output_sizes / merged.tp_size, -1)
         self._splice_merged_split(funcdef, calls, block, index)

vllm/model_executor/models/transformers/fusers/mla.py

@@ -21,7 +21,6 @@ from vllm.model_executor.models.transformers.fx_utils import (
     recover_forward,
     replace_expr,
     returned_linear,
-    single_self_call,
     upstream_linear,
 )
 from vllm.model_executor.models.transformers.utils import (
@@ -239,11 +238,12 @@ class MLAFuser(StackedFuser):
             # q_a_proj is usually inside the `else` of `if self.q_lora_rank is None`.
             # The fused call is inserted at the top-level statement preceding both.
             names = [self.q_a_proj_name, self.kv_a_proj_name]
-            calls = [single_self_call(funcdef, name) for name in names]
+            calls = self._unguarded_calls(funcdef, names)
             if ast.dump(calls[0].args[0]) != ast.dump(calls[1].args[0]):
                 raise ValueError("down-projections read different inputs")
-            index = min(_top_level_index(funcdef, call) for call in calls)
-            self._splice_merged_split(funcdef, calls, funcdef.body, index)
+            indices = [_top_level_index(funcdef, call) for call in calls]
+            self._check_input_stable(funcdef, module, calls, funcdef.body, indices)
+            self._splice_merged_split(funcdef, calls, funcdef.body, min(indices))
 
         # Transformers expands the latent into full key/value in a dedicated method.
         # `MLAAttention` consumes the latent directly (absorbing `kv_b_proj`),

vllm/model_executor/models/transformers/fusers/qkv.py

@@ -15,12 +15,11 @@ from vllm.model_executor.models.transformers.fusers.base import (
     fused_head_size,
 )
 from vllm.model_executor.models.transformers.fx_utils import (
+    block_chain,
     compile_forward,
-    innermost_block,
     is_linear,
     recover_forward,
     returned_linear,
-    single_self_call,
 )
 from vllm.model_executor.models.transformers.utils import (
     log_replacement,
@@ -99,12 +98,18 @@ class QKVFuser(StackedFuser):
         return cls(source_cls=type(module).__name__, **names)
 
     def update_forward(self, module: nn.Module) -> None:
-        """Replace `q(x), k(x), v(x)` with `qkv(x).split(sizes, -1)` in source."""
+        """Replace `q(x), k(x), v(x)` with `qkv(x).split(sizes, -1)` in source.
+
+        A projection may be guarded by an existence check -- a
+        `self.<proj> is (not) None` comparison (e.g. Gemma 4's `v_proj`) or a bare
+        truthiness test -- which is folded to its constant value (see
+        `bypass_existence_guard`). The calls may sit in different branches, so the
+        fused GEMM is inserted before the earliest of them, in the innermost block
+        that dominates all three.
+        """
         funcdef, fn = recover_forward(type(module))
-        calls = [
-            single_self_call(funcdef, name)
-            for name in (self.q_name, self.k_name, self.v_name)
-        ]
+        proj_names = (self.q_name, self.k_name, self.v_name)
+        calls = self._unguarded_calls(funcdef, proj_names)
         arg_dumps = {ast.dump(call.args[0]) for call in calls}
         if len(arg_dumps) != 1:
             raise ValueError("projection inputs are written differently")
@@ -115,7 +120,7 @@ class QKVFuser(StackedFuser):
             name
             for name, child in module.named_children()
             if isinstance(child, nn.Linear)
-        } - {self.q_name, self.k_name, self.v_name}
+        } - set(proj_names)
         for node in ast.walk(funcdef):
             if (
                 isinstance(node, ast.Call)
@@ -124,16 +129,24 @@ class QKVFuser(StackedFuser):
                 and any(ast.dump(arg) in arg_dumps for arg in node.args)
             ):
                 raise ValueError("another linear consumes the same input")
-        blocks = [innermost_block(funcdef.body, call) for call in calls]
-        if any(found is None for found in blocks):
+
+        # Insert the fused GEMM before the earliest call, in the innermost block
+        # common to all three (the calls may be split across branches).
+        chains = [block_chain(funcdef.body, call) for call in calls]
+        if any(not chain for chain in chains):
             raise ValueError("projection calls not found in the function body")
-        if len({id(block) for block, _ in blocks}) != 1:
-            raise ValueError("projection calls are in different blocks")
+        depth = 0
+        for level in zip(*chains):
+            if len({id(block) for block, _ in level}) != 1:
+                break
+            depth += 1
+        block = chains[0][depth - 1][0]
+        indices = [chain[depth - 1][1] for chain in chains]
+        insert_index = min(indices)
+        self._check_input_stable(funcdef, module, calls, block, indices)
 
         # q(x), k(x), v(x) -> q, k, v = qkv(x).split(qkv.output_sizes / qkv.tp_size, -1)
-        block = blocks[0][0]
-        index = min(index for _, index in blocks)
-        self._splice_merged_split(funcdef, calls, block, index)
+        self._splice_merged_split(funcdef, calls, block, insert_index)
         self.fused_forward = compile_forward(funcdef, fn)
 
     def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:

vllm/model_executor/models/transformers/fx_utils.py

@@ -70,6 +70,13 @@ class _MetaProxy(fx.Proxy):
     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), {}
+        )
+
 
 class _MetaAttribute(_MetaProxy, fx.proxy.Attribute):
     """Attribute proxy (e.g. `x.shape`) carrying its meta value.
@@ -312,41 +319,125 @@ def compile_forward(funcdef: ast.FunctionDef, fn: Callable) -> Callable:
     return namespace[funcdef.name]
 
 
+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
+
+
 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.
     """
-    uses = [
-        node
-        for node in ast.walk(funcdef)
-        if isinstance(node, ast.Attribute) and node.attr == name
-    ]
-    if len(uses) != 1:
-        raise ValueError(f"{name} is referenced {len(uses)} times")
-    calls = [
-        node
-        for node in ast.walk(funcdef)
-        if isinstance(node, ast.Call)
-        and node.func is uses[0]
-        and len(node.args) == 1
-        and not isinstance(node.args[0], ast.Starred)
-        and not node.keywords
-    ]
-    if (
-        len(calls) != 1
-        or not isinstance(uses[0].value, ast.Name)
-        or uses[0].value.id != "self"
-    ):
-        raise ValueError(f"{name} is not a single-argument call on self")
-    return calls[0]
+    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
 
 
-def innermost_block(
+def _is_none(node: ast.expr) -> bool:
+    return isinstance(node, ast.Constant) and node.value is None
+
+
+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 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 block_chain(
     block: list[ast.stmt], node: ast.AST
-) -> tuple[list[ast.stmt], int] | None:
-    """The innermost statement list containing `node`, and the index within."""
+) -> 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
@@ -359,13 +450,227 @@ def innermost_block(
             if (
                 isinstance(child_block, list)
                 and child_block
-                and (found := innermost_block(child_block, node)) is not None
+                and (tail := block_chain(child_block, node))
             ):
-                return found
-        return block, index
+                return [(block, index), *tail]
+        return [(block, index)]
+    return []
+
+
+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 _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 _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 _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 _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 _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 _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 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
+
+
+_METADATA_ATTRS = frozenset(
+    {"shape", "dtype", "device", "ndim", "size", "dim", "numel", "item"}
+)
+
+
+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 _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 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 replace_expr(module: ast.AST, old: ast.expr, new: ast.expr) -> None:
     """Replace the expression `old` (by identity) with `new` within `module`."""