diff --git a/tests/models/transformers/fusers/test_linear.py b/tests/models/transformers/fusers/test_linear.py index cc1af349fa..2e5b0fde38 100644 --- a/tests/models/transformers/fusers/test_linear.py +++ b/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 ` 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.:` 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 = [ diff --git a/tests/models/transformers/fusers/test_mla.py b/tests/models/transformers/fusers/test_mla.py index d61e54e923..9f51ab5662 100644 --- a/tests/models/transformers/fusers/test_mla.py +++ b/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`.""" diff --git a/vllm/model_executor/models/transformers/fusers/base.py b/vllm/model_executor/models/transformers/fusers/base.py index 7c7bfe273e..ad24f2542a 100644 --- a/vllm/model_executor/models/transformers/fusers/base.py +++ b/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.(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, diff --git a/vllm/model_executor/models/transformers/fusers/glu.py b/vllm/model_executor/models/transformers/fusers/glu.py index 6b81777dc7..15d886c473 100644 --- a/vllm/model_executor/models/transformers/fusers/glu.py +++ b/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]): diff --git a/vllm/model_executor/models/transformers/fusers/merged_column.py b/vllm/model_executor/models/transformers/fusers/merged_column.py index cd335438f0..4a2b9c2b09 100644 --- a/vllm/model_executor/models/transformers/fusers/merged_column.py +++ b/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) diff --git a/vllm/model_executor/models/transformers/fusers/mla.py b/vllm/model_executor/models/transformers/fusers/mla.py index c0d4127830..fa4b34750c 100644 --- a/vllm/model_executor/models/transformers/fusers/mla.py +++ b/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`), diff --git a/vllm/model_executor/models/transformers/fusers/qkv.py b/vllm/model_executor/models/transformers/fusers/qkv.py index e222d655e2..7b5b727d10 100644 --- a/vllm/model_executor/models/transformers/fusers/qkv.py +++ b/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. 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: diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py index 089f5870cc..332bd5414f 100644 --- a/vllm/model_executor/models/transformers/fx_utils.py +++ b/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.(arg)` call, plus every other `self.` reference. + + Raises unless exactly one single-argument `self.(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.(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.`).""" + 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.`'s existence to its constant value. + + The fuser deletes `self.` 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. is None` -> `False`, + `self. 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.:`, `... if self. else ...`, + `not self.`): 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. 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.(...)` 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`."""