@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e509d32b5a

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

Thomas Ortner committed Sep 9, 2026 at 15:46 UTC e509d32b5a2f3d81dc949ff89f85c4c8f860ae29
8 files changed +1027 -82
tests/models/transformers/fusers/test_linear.py
+578 -15
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Unit tests for the Transformers modeling backend's linear fusers."""
4
5 +import ast
6 import inspect
7 from types import MethodType, SimpleNamespace
8
@@ -19,7 +20,10 @@ from vllm.model_executor.models.transformers.fusers import (
20 packed_qkv,
21 qkv,
22 )
22 -from vllm.model_executor.models.transformers.fx_utils import trace
23 +from vllm.model_executor.models.transformers.fx_utils import (
24 + bypass_existence_guard,
25 + trace,
26 +)
27
28
29 class SiluAndMulStub(nn.Module):
@@ -202,6 +206,14 @@ class FourParallelLinears(nn.Module):
206 return self.proj_a(x), self.proj_b(x), self.proj_c(x), self.proj_d(x)
207
208
209 +class GuardedParallelLinears(FourParallelLinears):
210 + """A projection referenced by an existence guard as well as its call."""
211 +
212 + def forward(self, x):
213 + d = self.proj_d(x) if self.proj_d is not None else None
214 + return self.proj_a(x), self.proj_b(x), self.proj_c(x), d
215 +
216 +
217 class MutatingParallelLinears(FourParallelLinears):
218 def forward(self, x):
219 a = self.proj_a(x)
@@ -209,6 +221,34 @@ class MutatingParallelLinears(FourParallelLinears):
221 return a, self.proj_b(x), self.proj_c(x), self.proj_d(x)
222
223
224 +class AliasMutatingParallelLinears(FourParallelLinears):
225 + """A view of the input is mutated between the projections."""
226 +
227 + def forward(self, x):
228 + a = self.proj_a(x)
229 + alias = x.view(-1)
230 + alias.mul_(2)
231 + return a, self.proj_b(x), self.proj_c(x), self.proj_d(x)
232 +
233 +
234 +class FunctionalInplaceParallelLinears(FourParallelLinears):
235 + """The input is mutated by a functional in-place call between projections."""
236 +
237 + def forward(self, x):
238 + a = self.proj_a(x)
239 + F.relu(x, inplace=True)
240 + return a, self.proj_b(x), self.proj_c(x), self.proj_d(x)
241 +
242 +
243 +class UnderscoreInplaceParallelLinears(FourParallelLinears):
244 + """The input is mutated by a free in-place function between projections."""
245 +
246 + def forward(self, x):
247 + a = self.proj_a(x)
248 + torch.relu_(x)
249 + return a, self.proj_b(x), self.proj_c(x), self.proj_d(x)
250 +
251 +
252 class PlainQKV(nn.Module):
253 """q/k/v returned as a bare tuple: no reshape stands between the calls and
254 the return, so both QKVFuser and MergedColumnParallelFuser can rewrite it."""
@@ -294,6 +334,306 @@ class ResidDropoutAttention(FakeAttention):
334 return self.resid_dropout(self.o_proj(attn_output)), None
335
336
337 +class GuardedVAttention(FakeAttention):
338 + """Gemma 4-style: `v_proj` is called once but also read in a `None` guard.
339 +
340 + The guard is a fusion invariant (the fuser only binds instances where the
341 + projection exists), so it folds to `True` and the call still fuses."""
342 +
343 + def forward(
344 + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
345 + ):
346 + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
347 +
348 + input_shape = hidden_states.shape[:-1]
349 + hidden_shape = (*input_shape, -1, self.head_dim)
350 + q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
351 + k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
352 + v = (
353 + self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
354 + if self.v_proj is not None
355 + else k
356 + )
357 + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
358 + self.config._attn_implementation, None
359 + )
360 + attn_output, _ = attention_interface(
361 + self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
362 + )
363 + return self.o_proj(attn_output.reshape(*input_shape, -1).contiguous()), None
364 +
365 +
366 +class TruthyGuardAttention(FakeAttention):
367 + """`v_proj` guarded by a bare truthiness test instead of `is not None`.
368 +
369 + An `nn.Linear` is always truthy, so the guard folds to `True` and the call
370 + still fuses (the projection is a proven invariant, exactly as for a `None`
371 + guard)."""
372 +
373 + def forward(
374 + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
375 + ):
376 + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
377 +
378 + input_shape = hidden_states.shape[:-1]
379 + hidden_shape = (*input_shape, -1, self.head_dim)
380 + q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
381 + k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
382 + v = (
383 + self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
384 + if self.v_proj
385 + else k
386 + )
387 + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
388 + self.config._attn_implementation, None
389 + )
390 + attn_output, _ = attention_interface(
391 + self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
392 + )
393 + return self.o_proj(attn_output.reshape(*input_shape, -1).contiguous()), None
394 +
395 +
396 +class StatementGuardVAttention(FakeAttention):
397 + """`v_proj` guarded by a statement-form `if self.v_proj:` block.
398 +
399 + The `if`'s test is not an `ast.Compare`, so the identity-check scan skips it;
400 + `_in_boolean_context` folds the test to `True` (an `nn.Linear` is always
401 + truthy) and the call inside the now-`if True:` body still fuses, with the
402 + live `else` surviving the rewrite."""
403 +
404 + def forward(
405 + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
406 + ):
407 + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
408 +
409 + input_shape = hidden_states.shape[:-1]
410 + hidden_shape = (*input_shape, -1, self.head_dim)
411 + q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
412 + k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
413 + if self.v_proj:
414 + v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
415 + else:
416 + v = k
417 + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
418 + self.config._attn_implementation, None
419 + )
420 + attn_output, _ = attention_interface(
421 + self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
422 + )
423 + return self.o_proj(attn_output.reshape(*input_shape, -1).contiguous()), None
424 +
425 +
426 +class BranchedKVAttention(FakeAttention):
427 + """Gemma 4-style: q at body level, k/v inside an untaken `else` branch.
428 +
429 + The three calls do not share a block, so the fused GEMM must be hoisted to
430 + the innermost block that dominates all of them (the function body). The
431 + live `if`/`else` must survive the rewrite."""
432 +
433 + is_kv_shared_layer = False
434 +
435 + def forward(
436 + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
437 + ):
438 + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
439 +
440 + input_shape = hidden_states.shape[:-1]
441 + hidden_shape = (*input_shape, -1, self.head_dim)
442 + q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
443 + if self.is_kv_shared_layer:
444 + k = v = None
445 + else:
446 + k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
447 + v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
448 + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
449 + self.config._attn_implementation, None
450 + )
451 + attn_output, _ = attention_interface(
452 + self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
453 + )
454 + return self.o_proj(attn_output.reshape(*input_shape, -1).contiguous()), None
455 +
456 +
457 +class RebindArgAttention(FakeAttention):
458 + """The projection input is rebound (in an untaken branch) before k/v.
459 +
460 + The branch is dead at trace time so the match succeeds, but the source
461 + rewrite must refuse: hoisting the GEMM above the rebind would feed k/v a
462 + different input. Result: no fusion."""
463 +
464 + recompute = False
465 +
466 + def forward(
467 + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
468 + ):
469 + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
470 +
471 + input_shape = hidden_states.shape[:-1]
472 + hidden_shape = (*input_shape, -1, self.head_dim)
473 + q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
474 + if self.recompute:
475 + hidden_states = hidden_states * 2
476 + k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
477 + v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
478 + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
479 + self.config._attn_implementation, None
480 + )
481 + attn_output, _ = attention_interface(
482 + self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
483 + )
484 + return self.o_proj(attn_output.reshape(*input_shape, -1)), None
485 +
486 +
487 +class InPlaceMutatedArgAttention(FakeAttention):
488 + """The shared input is mutated in place (in an untaken branch) before k/v.
489 +
490 + Like `RebindArgAttention`, but the mutation keeps the name in a `Load`
491 + context (`hidden_states.mul_(2)`), so a bare `Name`-store check would miss it.
492 + The rewrite must still refuse: hoisting the GEMM above the mutation changes
493 + the value k/v see. Result: no fusion."""
494 +
495 + recompute = False
496 +
497 + def forward(
498 + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
499 + ):
500 + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
501 +
502 + input_shape = hidden_states.shape[:-1]
503 + hidden_shape = (*input_shape, -1, self.head_dim)
504 + q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
505 + if self.recompute:
506 + hidden_states.mul_(2)
507 + k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
508 + v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
509 + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
510 + self.config._attn_implementation, None
511 + )
512 + attn_output, _ = attention_interface(
513 + self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
514 + )
515 + return self.o_proj(attn_output.reshape(*input_shape, -1)), None
516 +
517 +
518 +class AliasMutatedArgAttention(FakeAttention):
519 + """A view of the shared input is mutated before k/v.
520 +
521 + `hidden_states.view(-1)` shares storage with its base, so `alias.mul_(2)`
522 + changes what k/v read without ever naming `hidden_states` as a target. Only
523 + tracking names that may alias the input catches it. Result: no fusion."""
524 +
525 + recompute = False
526 +
527 + def forward(
528 + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
529 + ):
530 + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
531 +
532 + input_shape = hidden_states.shape[:-1]
533 + hidden_shape = (*input_shape, -1, self.head_dim)
534 + q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
535 + if self.recompute:
536 + alias = hidden_states.view(-1)
537 + alias.mul_(2)
538 + k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
539 + v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
540 + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
541 + self.config._attn_implementation, None
542 + )
543 + attn_output, _ = attention_interface(
544 + self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
545 + )
546 + return self.o_proj(attn_output.reshape(*input_shape, -1)), None
547 +
548 +
549 +class FunctionalInplaceArgAttention(FakeAttention):
550 + """The shared input is mutated by a functional in-place call before k/v.
551 +
552 + `F.relu(hidden_states, inplace=True)` writes through an argument, so the
553 + name appears only in a `Load` context inside a call that is not a method on
554 + it -- invisible to a trailing-underscore method check. Result: no fusion."""
555 +
556 + recompute = False
557 +
558 + def forward(
559 + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
560 + ):
561 + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
562 +
563 + input_shape = hidden_states.shape[:-1]
564 + hidden_shape = (*input_shape, -1, self.head_dim)
565 + q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
566 + if self.recompute:
567 + F.relu(hidden_states, inplace=True)
568 + k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
569 + v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
570 + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
571 + self.config._attn_implementation, None
572 + )
573 + attn_output, _ = attention_interface(
574 + self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
575 + )
576 + return self.o_proj(attn_output.reshape(*input_shape, -1)), None
577 +
578 +
579 +class SubscriptMutatedArgAttention(FakeAttention):
580 + """The shared input is written through a subscript (in an untaken branch).
581 +
582 + `hidden_states[..., 0] = 0` leaves the name in a `Load` context on the
583 + subscript's value, so it too evades a bare `Name`-store check. The rewrite
584 + must refuse for the same reason. Result: no fusion."""
585 +
586 + recompute = False
587 +
588 + def forward(
589 + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
590 + ):
591 + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
592 +
593 + input_shape = hidden_states.shape[:-1]
594 + hidden_shape = (*input_shape, -1, self.head_dim)
595 + q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
596 + if self.recompute:
597 + hidden_states[..., 0] = 0
598 + k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
599 + v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
600 + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
601 + self.config._attn_implementation, None
602 + )
603 + attn_output, _ = attention_interface(
604 + self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
605 + )
606 + return self.o_proj(attn_output.reshape(*input_shape, -1)), None
607 +
608 +
609 +class NonFoldableRefAttention(FakeAttention):
610 + """A surviving `v_proj` reference that is not an existence guard -> no fusion.
611 +
612 + Deleting `v_proj` would break `isinstance(self.v_proj, nn.Linear)`: it reads
613 + the projection's type, not just whether it exists, so the fuser cannot fold
614 + it and must refuse rather than rewrite it."""
615 +
616 + def forward(
617 + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
618 + ):
619 + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
620 +
621 + input_shape = hidden_states.shape[:-1]
622 + hidden_shape = (*input_shape, -1, self.head_dim)
623 + q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
624 + k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
625 + v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
626 + if isinstance(self.v_proj, nn.Linear):
627 + v = v.contiguous()
628 + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
629 + self.config._attn_implementation, None
630 + )
631 + attn_output, _ = attention_interface(
632 + self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
633 + )
634 + return self.o_proj(attn_output.reshape(*input_shape, -1)), None
635 +
636 +
637 class PackedQKVAttention(nn.Module):
638 """GPTBigCode-style: one packed projection split into q/k/v in the forward."""
639
@@ -561,18 +901,246 @@ def test_qkv_identifies_output_projection():
901 assert get_fuser(ResidDropoutAttention(), QKVFuser).o_name == "o_proj"
902
903
564 -def test_merged_column_fuser_rejects_input_mutation():
565 - """The later projections must not be moved ahead of an input mutation."""
904 +@pytest.mark.parametrize(
905 + "attn_cls",
906 + [
907 + GuardedVAttention,
908 + TruthyGuardAttention,
909 + StatementGuardVAttention,
910 + BranchedKVAttention,
911 + ],
912 +)
913 +def test_fuses_gemma4_qkv_obstacles(attn_cls):
914 + """A guarded `v_proj` (obstacle 1) or branch-split k/v (obstacle 2) fuses.
915 +
916 + Both patterns leave `v_proj` referenced beyond its single call, or split the
917 + calls across blocks; the rewrite folds the guard / hoists the GEMM and the
918 + numerics still match. The guard may be a `None` comparison or a bare
919 + truthiness test (ternary or statement form); all are proven invariants, so
920 + folding them is sound. The naive "relax the reference count" fix fuses the
921 + guarded cases but silently drops the guard's semantics.
922 + """
923 + with torch.device("meta"):
924 + meta = attn_cls()
925 + fuser = get_fuser(meta, QKVFuser)
926 + assert isinstance(fuser, QKVFuser)
927 + assert (fuser.q_name, fuser.o_name) == ("q_proj", "o_proj")
928 +
929 + names = set(fuser.fused_forward.__code__.co_names)
930 + assert "qkv_proj" in names
931 + assert not {"q_proj", "k_proj", "v_proj"} & names
932 + if attn_cls is BranchedKVAttention:
933 + assert "is_kv_shared_layer" in names # the live branch survives
934 +
935 + real = attn_cls(kv_heads=4, layer_idx=3)
936 + for p in real.parameters():
937 + nn.init.normal_(p, std=0.05)
938 + x = torch.randn(1, 5, 32)
939 + real.attn = FakeSelfAttn()
940 + expected, _ = real(x)
941 + fused = _apply_qkv_fuser_with_stubs(real, fuser)
942 + out, _ = fused(x)
943 + torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)
944 +
945 +
946 +@pytest.mark.parametrize(
947 + "attn_cls",
948 + [
949 + RebindArgAttention,
950 + InPlaceMutatedArgAttention,
951 + SubscriptMutatedArgAttention,
952 + AliasMutatedArgAttention,
953 + FunctionalInplaceArgAttention,
954 + NonFoldableRefAttention,
955 + ],
956 +)
957 +def test_qkv_refuses_unsound_rewrites(attn_cls):
958 + """The match succeeds but the source rewrite must refuse, leaving no fuser.
959 +
960 + `RebindArgAttention` rebinds the shared input before k/v; the `*Mutated*` and
961 + `FunctionalInplace*` cases mutate it in place, directly (`x.mul_(...)`,
962 + `x[...] = 0`), through a view that shares its storage, or through a call that
963 + writes an argument (`F.relu(x, inplace=True)`) -- each would change what k/v
964 + see if the GEMM were hoisted above them.
965 + `NonFoldableRefAttention` reads `v_proj` outside an existence guard (the fuser
966 + cannot fold it away). All fail closed in `update_forward`, so `get_fuser`
967 + returns `None`.
968 + """
969 with torch.device("meta"):
567 - module = MutatingParallelLinears()
970 + meta = attn_cls()
971 + assert get_fuser(meta, QKVFuser) is None
972 +
973 +
974 +def _guard_funcdef(expr: str) -> tuple[ast.FunctionDef, ast.Attribute]:
975 + """A `def f(self, a, b, T): return <expr>` and its `self.v_proj` reference."""
976 + funcdef = ast.parse(f"def f(self, a, b, T):\n return {expr}").body[0]
977 + assert isinstance(funcdef, ast.FunctionDef)
978 + ref = next(
979 + node
980 + for node in ast.walk(funcdef)
981 + if isinstance(node, ast.Attribute) and node.attr == "v_proj"
982 + )
983 + return funcdef, ref
984 +
985 +
986 +@pytest.mark.parametrize(
987 + "expr, expected",
988 + [
989 + ("a if self.v_proj is not None else b", True),
990 + ("a if self.v_proj is None else b", False),
991 + ("a if None is not self.v_proj else b", True), # None on the left
992 + ],
993 +)
994 +def test_bypass_existence_guard_folds_identity_none_check(expr, expected):
995 + """`is`/`is not None`, either operand order, fold to a bool.
996 +
997 + The projection is a proven invariant (it exists as an `nn.Linear`), so an
998 + `is (not) None` guard has a constant truth value; folding it lets the call be
999 + rewritten away without changing the guard's outcome. `==`/`!=` are excluded
1000 + (see the refusal test) because a subclass may override `__eq__`/`__ne__`."""
1001 + funcdef, ref = _guard_funcdef(expr)
1002 + bypass_existence_guard(funcdef, ref, "v_proj")
1003 + test = next(node for node in ast.walk(funcdef) if isinstance(node, ast.IfExp)).test
1004 + assert isinstance(test, ast.Constant) and test.value is expected
1005 +
1006 +
1007 +@pytest.mark.parametrize(
1008 + "expr",
1009 + [
1010 + "a if self.v_proj else b", # a ternary test
1011 + "not self.v_proj", # a `not` operand
1012 + ],
1013 +)
1014 +def test_bypass_existence_guard_folds_bare_truthiness(expr):
1015 + """A bare truthiness test folds the reference itself to `True`.
1016 +
1017 + An `nn.Linear` has no `__bool__`/`__len__`, so it is always truthy; in a pure
1018 + boolean position (a test, or a `not` operand) the reference can be replaced by
1019 + `True`, letting the projection be deleted."""
1020 + funcdef, ref = _guard_funcdef(expr)
1021 + bypass_existence_guard(funcdef, ref, "v_proj")
1022 + # No `self.v_proj` reference survives; a literal `True` took its place.
1023 + attrs = [n for n in ast.walk(funcdef) if isinstance(n, ast.Attribute)]
1024 + assert not any(n.attr == "v_proj" for n in attrs)
1025 + constants = [n for n in ast.walk(funcdef) if isinstance(n, ast.Constant)]
1026 + assert any(n.value is True for n in constants)
1027 +
1028 +
1029 +def test_bypass_existence_guard_folds_bare_if_statement():
1030 + """A statement-form `if self.<name>:` folds its test to `True`.
1031 +
1032 + The `len(node.ops) == 1` `Compare` scan skips it (it is not a comparison at
1033 + all); `_in_boolean_context` then matches the `if`'s test and folds it, so the
1034 + bare-truthiness statement form is handled, not silently dropped."""
1035 + funcdef = ast.parse(
1036 + "def f(self, a, b):\n if self.v_proj:\n return a\n return b"
1037 + ).body[0]
1038 + assert isinstance(funcdef, ast.FunctionDef)
1039 + ref = next(
1040 + node
1041 + for node in ast.walk(funcdef)
1042 + if isinstance(node, ast.Attribute) and node.attr == "v_proj"
1043 + )
1044 + bypass_existence_guard(funcdef, ref, "v_proj")
1045 + test = next(node for node in ast.walk(funcdef) if isinstance(node, ast.If)).test
1046 + assert isinstance(test, ast.Constant) and test.value is True
1047 +
1048 +
1049 +@pytest.mark.parametrize(
1050 + "expr",
1051 + [
1052 + "a if isinstance(self.v_proj, T) else b", # a different predicate
1053 + "self.v_proj.weight", # an attribute read, not a guard
1054 + "a and self.v_proj", # `and` yields the operand, which may escape
1055 + "self.v_proj or a", # `or` yields the module itself when truthy
1056 + "a if self.v_proj == None else b", # `__eq__` may not track identity
1057 + "a if self.v_proj != None else b", # `__ne__` may not track identity
1058 + ],
1059 +)
1060 +def test_bypass_existence_guard_refuses_non_guard_reference(expr):
1061 + """A surviving reference that is not an identity guard cannot be folded away.
1062 +
1063 + Rewriting it is unsafe (its value depends on more than the projection's
1064 + existence), so the fuser refuses rather than change semantics. `==`/`!=` are
1065 + refused too: `is_linear` accepts `nn.Linear` subclasses that could override
1066 + `__eq__`/`__ne__`, so equality need not agree with `is (not) None`."""
1067 + funcdef, ref = _guard_funcdef(expr)
1068 + with pytest.raises(ValueError, match="outside an existence guard"):
1069 + bypass_existence_guard(funcdef, ref, "v_proj")
1070 +
1071 +
1072 +@pytest.mark.parametrize("layer_idx", [0, 1])
1073 +def test_fuses_real_gemma4_attention(layer_idx):
1074 + """The real Gemma 4 attention (both obstacles at once) fuses end to end."""
1075 + pytest.importorskip("transformers.models.gemma4")
1076 + from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
1077 + from transformers.models.gemma4.modeling_gemma4 import Gemma4TextAttention
1078 +
1079 + config = Gemma4TextConfig(
1080 + hidden_size=64,
1081 + num_attention_heads=4,
1082 + num_key_value_heads=2,
1083 + num_hidden_layers=2,
1084 + head_dim=16,
1085 + intermediate_size=128,
1086 + )
1087 + with torch.device("meta"):
1088 + meta = Gemma4TextAttention(config, layer_idx=layer_idx)
1089 + fuser = get_fuser(meta, QKVFuser)
1090 + assert isinstance(fuser, QKVFuser)
1091 + assert (fuser.q_name, fuser.k_name, fuser.v_name) == (
1092 + "q_proj",
1093 + "k_proj",
1094 + "v_proj",
1095 + )
1096 + assert fuser.o_name == "o_proj"
1097 + names = set(fuser.fused_forward.__code__.co_names)
1098 + assert not {"q_proj", "k_proj", "v_proj"} & names
1099 +
1100 +
1101 +def test_merged_column_fuser_folds_existence_guard():
1102 + """A guarded projection fuses: the merged linear replaces the deleted names.
1103 +
1104 + Guard folding lives on `StackedFuser`, not just `QKVFuser`, because every
1105 + stacked fuser deletes the projections it merges -- a surviving reference to
1106 + one would fail at runtime."""
1107 + with torch.device("meta"):
1108 + module = GuardedParallelLinears()
1109 + fuser = MergedColumnParallelFuser.match(trace(module), module)
1110 + assert fuser is not None
1111 + fuser.update_forward(module)
1112 + names = set(fuser.fused_forward.__code__.co_names)
1113 + assert fuser.merged_name in names
1114 + assert not {"proj_a", "proj_b", "proj_c", "proj_d"} & names
1115 +
1116 +
1117 +@pytest.mark.parametrize(
1118 + "cls",
1119 + [
1120 + MutatingParallelLinears,
1121 + AliasMutatingParallelLinears,
1122 + FunctionalInplaceParallelLinears,
1123 + UnderscoreInplaceParallelLinears,
1124 + ],
1125 +)
1126 +def test_merged_column_fuser_rejects_input_mutation(cls):
1127 + """The later projections must not be moved ahead of an input mutation.
1128 +
1129 + The mutation may be direct (`x.add_(1)`), through a view that shares the
1130 + input's storage, or through a call that writes an argument in place
1131 + (`F.relu(x, inplace=True)`, `torch.relu_(x)`). Each one silently changes the
1132 + fused result, so the rewrite must refuse rather than fuse."""
1133 + with torch.device("meta"):
1134 + module = cls()
1135 fuser = MergedColumnParallelFuser.match(trace(module), module)
1136 assert fuser is not None
570 - with pytest.raises(ValueError, match="cross other operations"):
1137 + with pytest.raises(ValueError, match="rebound or mutated"):
1138 fuser.update_forward(module)
1139
1140
1141 +@pytest.mark.parametrize("cls", [FourParallelLinears, GuardedParallelLinears])
1142 def test_merged_column_fuser_supports_any_number_of_linears(
575 - default_vllm_config, monkeypatch
1143 + cls, default_vllm_config, monkeypatch
1144 ):
1145 from vllm.model_executor import parameter
1146 from vllm.model_executor.layers import linear
@@ -585,19 +1153,14 @@ def test_merged_column_fuser_supports_any_number_of_linears(
1153 monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0)
1154 monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1)
1155 with torch.device("meta"):
588 - meta = FourParallelLinears()
1156 + meta = cls()
1157 fuser = MergedColumnParallelFuser.match(trace(meta), meta)
1158 assert fuser is not None
591 - assert fuser.linear_names == ("proj_a", "proj_b", "proj_c", "proj_d")
592 - assert fuser.shards == [
593 - ("proj_a", 0),
594 - ("proj_b", 1),
595 - ("proj_c", 2),
596 - ("proj_d", 3),
597 - ]
1159 + assert set(fuser.linear_names) == {"proj_a", "proj_b", "proj_c", "proj_d"}
1160 + assert fuser.shards == [(name, i) for i, name in enumerate(fuser.linear_names)]
1161 fuser.update_forward(meta)
1162
600 - real = FourParallelLinears()
1163 + real = cls()
1164 x = torch.randn(2, 8)
1165 expected = real(x)
1166 weights = [
tests/models/transformers/fusers/test_mla.py
+26 -4
@@ -19,9 +19,8 @@ from vllm.model_executor.models.transformers.fx_utils import trace
19 _FUSED_QKV_A_PROJ = MLAFuser.merged_name
20
21
22 -def _match(q_lora_rank: int | None) -> MLAFuser | None:
23 - """Match a meta `DeepseekV2Attention` directly (bypassing the per-class
24 - `get_fuser` cache, so both q_lora variants of the same class are seen)."""
22 +def _attention(q_lora_rank: int | None):
23 + """A meta `DeepseekV2Attention` in the requested q_lora configuration."""
24 pytest.importorskip("transformers.models.deepseek_v2.modeling_deepseek_v2")
25 from transformers.models.deepseek_v2.configuration_deepseek_v2 import (
26 DeepseekV2Config,
@@ -41,7 +40,13 @@ def _match(q_lora_rank: int | None) -> MLAFuser | None:
40 num_hidden_layers=1,
41 )
42 with torch.device("meta"):
44 - attn = DeepseekV2Attention(cfg, layer_idx=0)
43 + return DeepseekV2Attention(cfg, layer_idx=0)
44 +
45 +
46 +def _match(q_lora_rank: int | None) -> MLAFuser | None:
47 + """Match a meta `DeepseekV2Attention` directly (bypassing the per-class
48 + `get_fuser` cache, so both q_lora variants of the same class are seen)."""
49 + attn = _attention(q_lora_rank)
50 return MLAFuser.match(trace(attn), attn)
51
52
@@ -90,6 +95,23 @@ def test_q_lora_stacks_qkv_a_proj():
95 }
96
97
98 +@pytest.mark.parametrize("q_lora_rank", [64, None])
99 +def test_update_forward_rewrites_real_attention(q_lora_rank):
100 + """The source rewrite must survive the input-stability check.
101 +
102 + The q-LoRA path hoists the merged down-projection to a top-level statement
103 + while its calls sit inside the `q_lora_rank` branches, so the region it
104 + spans is wide; a check that over-approximates aliasing would silently refuse
105 + here and cost the fusion rather than fail loudly."""
106 + fuser = _match(q_lora_rank)
107 + assert isinstance(fuser, MLAFuser)
108 + fuser.update_forward(_attention(q_lora_rank))
109 + names = set(fuser.fused_forward.__code__.co_names)
110 + if fuser.has_q_lora:
111 + assert _FUSED_QKV_A_PROJ in names
112 + assert not {"q_a_proj", "kv_a_proj_with_mqa"} & names
113 +
114 +
115 class _Norm(nn.Module):
116 """A real RMSNorm computation: `match` verifies chain norms via `RMSNormFuser`."""
117
vllm/model_executor/models/transformers/fusers/base.py
+55 -2
@@ -5,13 +5,19 @@
5 import ast
6 import types
7 from abc import ABC, abstractmethod
8 -from collections.abc import Callable
8 +from collections.abc import Callable, Iterable
9 from dataclasses import dataclass, field
10 from typing import TYPE_CHECKING, ClassVar
11
12 from torch import fx, nn
13
14 -from vllm.model_executor.models.transformers.fx_utils import replace_expr
14 +from vllm.model_executor.models.transformers.fx_utils import (
15 + aliasing_names,
16 + bypass_existence_guard,
17 + replace_expr,
18 + self_call_and_refs,
19 + written_names,
20 +)
21 from vllm.model_executor.models.utils import ShardId, maybe_prefix
22
23 if TYPE_CHECKING:
@@ -160,6 +166,53 @@ class StackedFuser(RewriteFuser):
166 fused layer into its per-shard configs."""
167 return {self.merged_name: [name for name, _ in self.shards]}
168
169 + def _unguarded_calls(
170 + self, funcdef: ast.FunctionDef, names: Iterable[str]
171 + ) -> list[ast.Call]:
172 + """One `self.<name>(arg)` call per projection, existence guards folded.
173 +
174 + `update_attrs` deletes the projections it merges, so any reference to one beyond
175 + its call site must be a guard on its existence; folding those to their constant
176 + value keeps the rewritten forward off a name that no longer exists. A reference
177 + that is not such a guard raises, so fusion is skipped."""
178 + calls = []
179 + for name in names:
180 + call, refs = self_call_and_refs(funcdef, name)
181 + for ref in refs:
182 + bypass_existence_guard(funcdef, ref, name)
183 + calls.append(call)
184 + return calls
185 +
186 + def _check_input_stable(
187 + self,
188 + funcdef: ast.FunctionDef,
189 + module: nn.Module,
190 + calls: list[ast.Call],
191 + block: list[ast.stmt],
192 + indices: list[int],
193 + ) -> None:
194 + """Raise unless hoisting the merged GEMM preserves what it reads.
195 +
196 + Fusing moves every projection to one call at `min(indices)`, so the
197 + merged GEMM reads the input once, up front, where the last of `calls`
198 + would have read it later. That holds only if nothing in between changes
199 + the input, and a change need not name it: writing any view that shares
200 + its storage changes it too. Both halves of the check over-approximate,
201 + since a false hit costs a fusion while a miss returns wrong numbers.
202 + """
203 + arg_names = {
204 + node.id
205 + for node in ast.walk(calls[0].args[0])
206 + if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load)
207 + }
208 + # Protect the names the input reads and anything that may alias them.
209 + tracked = aliasing_names(funcdef, arg_names, module)
210 + # Any of them rebound, written through, or passed to a call that writes
211 + # its argument would leave the hoisted GEMM reading a different value.
212 + region = block[min(indices) : max(indices) + 1]
213 + if tracked & written_names(region):
214 + raise ValueError("projection input is rebound or mutated before all calls")
215 +
216 def _splice_merged_split(
217 self,
218 funcdef: ast.FunctionDef,
vllm/model_executor/models/transformers/fusers/glu.py
+2 -2
@@ -137,8 +137,8 @@ class GLUFuser(MergedColumnParallelFuser):
137 """Replace `act(gate(x)) * up(x)` with `act(gate_up(x))` in source."""
138 funcdef, fn = recover_forward(type(module))
139 act_call = single_self_call(funcdef, self.act_name)
140 - gate_call = single_self_call(funcdef, self.gate_name)
141 - up_call = single_self_call(funcdef, self.up_name)
140 + calls = self._unguarded_calls(funcdef, (self.gate_name, self.up_name))
141 + gate_call, up_call = calls
142 if act_call.args[0] is not gate_call:
143 raise ValueError("activation does not directly wrap the gate")
144 if ast.dump(gate_call.args[0]) != ast.dump(up_call.args[0]):
vllm/model_executor/models/transformers/fusers/merged_column.py
+8 -19
@@ -12,11 +12,10 @@ from vllm.logger import init_logger
12 from vllm.model_executor.layers.linear import MergedColumnParallelLinear
13 from vllm.model_executor.models.transformers.fusers.base import StackedFuser
14 from vllm.model_executor.models.transformers.fx_utils import (
15 + block_chain,
16 compile_forward,
16 - innermost_block,
17 is_linear,
18 recover_forward,
19 - single_self_call,
19 )
20 from vllm.model_executor.models.utils import ShardId, maybe_prefix
21
@@ -92,30 +91,20 @@ class MergedColumnParallelFuser(StackedFuser):
91 def update_forward(self, module: nn.Module) -> None:
92 """Replace the parallel calls with one merged call and split."""
93 funcdef, fn = recover_forward(type(module))
95 - calls = [single_self_call(funcdef, name) for name in self.linear_names]
94 + calls = self._unguarded_calls(funcdef, self.linear_names)
95 if len(set(ast.dump(call.args[0]) for call in calls)) != 1:
96 raise ValueError("parallel linears read different inputs")
98 - blocks = [innermost_block(funcdef.body, call) for call in calls]
99 - if any(found is None for found in blocks):
97 + chains = [block_chain(funcdef.body, call) for call in calls]
98 + if any(not chain for chain in chains):
99 raise ValueError("parallel linear calls not found in the function body")
100 + blocks = [chain[-1] for chain in chains]
101 if len(set(id(block) for block, _ in blocks)) != 1:
102 raise ValueError("parallel linear calls are in different blocks")
103
104 block = blocks[0][0]
105 - index = min(index for _, index in blocks)
106 - calls_have_name_arg = all(isinstance(call.args[0], ast.Name) for call in calls)
107 - # Moving projections must not cross operations that can change their input.
108 - for statement in block[index : max(index for _, index in blocks) + 1]:
109 - if not isinstance(statement, (ast.Assign, ast.Return)):
110 - raise ValueError("parallel projections cross other operations")
111 - if isinstance(statement, ast.Assign) and any(
112 - not isinstance(target, ast.Name) for target in statement.targets
113 - ):
114 - raise ValueError("parallel projection assignment has side effects")
115 - value = statement.value
116 - values = value.elts if isinstance(value, ast.Tuple) else [value]
117 - if any(value not in calls for value in values) or not calls_have_name_arg:
118 - raise ValueError("parallel projections cross other operations")
105 + indices = [index for _, index in blocks]
106 + index = min(indices)
107 + self._check_input_stable(funcdef, module, calls, block, indices)
108
109 # l1(x), l2(x), ... -> merged(x).split(merged.output_sizes / merged.tp_size, -1)
110 self._splice_merged_split(funcdef, calls, block, index)
vllm/model_executor/models/transformers/fusers/mla.py
+4 -4
@@ -21,7 +21,6 @@ from vllm.model_executor.models.transformers.fx_utils import (
21 recover_forward,
22 replace_expr,
23 returned_linear,
24 - single_self_call,
24 upstream_linear,
25 )
26 from vllm.model_executor.models.transformers.utils import (
@@ -239,11 +238,12 @@ class MLAFuser(StackedFuser):
238 # q_a_proj is usually inside the `else` of `if self.q_lora_rank is None`.
239 # The fused call is inserted at the top-level statement preceding both.
240 names = [self.q_a_proj_name, self.kv_a_proj_name]
242 - calls = [single_self_call(funcdef, name) for name in names]
241 + calls = self._unguarded_calls(funcdef, names)
242 if ast.dump(calls[0].args[0]) != ast.dump(calls[1].args[0]):
243 raise ValueError("down-projections read different inputs")
245 - index = min(_top_level_index(funcdef, call) for call in calls)
246 - self._splice_merged_split(funcdef, calls, funcdef.body, index)
244 + indices = [_top_level_index(funcdef, call) for call in calls]
245 + self._check_input_stable(funcdef, module, calls, funcdef.body, indices)
246 + self._splice_merged_split(funcdef, calls, funcdef.body, min(indices))
247
248 # Transformers expands the latent into full key/value in a dedicated method.
249 # `MLAAttention` consumes the latent directly (absorbing `kv_b_proj`),
vllm/model_executor/models/transformers/fusers/qkv.py
+28 -15
@@ -15,12 +15,11 @@ from vllm.model_executor.models.transformers.fusers.base import (
15 fused_head_size,
16 )
17 from vllm.model_executor.models.transformers.fx_utils import (
18 + block_chain,
19 compile_forward,
19 - innermost_block,
20 is_linear,
21 recover_forward,
22 returned_linear,
23 - single_self_call,
23 )
24 from vllm.model_executor.models.transformers.utils import (
25 log_replacement,
@@ -99,12 +98,18 @@ class QKVFuser(StackedFuser):
98 return cls(source_cls=type(module).__name__, **names)
99
100 def update_forward(self, module: nn.Module) -> None:
102 - """Replace `q(x), k(x), v(x)` with `qkv(x).split(sizes, -1)` in source."""
101 + """Replace `q(x), k(x), v(x)` with `qkv(x).split(sizes, -1)` in source.
102 +
103 + A projection may be guarded by an existence check -- a
104 + `self.<proj> is (not) None` comparison (e.g. Gemma 4's `v_proj`) or a bare
105 + truthiness test -- which is folded to its constant value (see
106 + `bypass_existence_guard`). The calls may sit in different branches, so the
107 + fused GEMM is inserted before the earliest of them, in the innermost block
108 + that dominates all three.
109 + """
110 funcdef, fn = recover_forward(type(module))
104 - calls = [
105 - single_self_call(funcdef, name)
106 - for name in (self.q_name, self.k_name, self.v_name)
107 - ]
111 + proj_names = (self.q_name, self.k_name, self.v_name)
112 + calls = self._unguarded_calls(funcdef, proj_names)
113 arg_dumps = {ast.dump(call.args[0]) for call in calls}
114 if len(arg_dumps) != 1:
115 raise ValueError("projection inputs are written differently")
@@ -115,7 +120,7 @@ class QKVFuser(StackedFuser):
120 name
121 for name, child in module.named_children()
122 if isinstance(child, nn.Linear)
118 - } - {self.q_name, self.k_name, self.v_name}
123 + } - set(proj_names)
124 for node in ast.walk(funcdef):
125 if (
126 isinstance(node, ast.Call)
@@ -124,16 +129,24 @@ class QKVFuser(StackedFuser):
129 and any(ast.dump(arg) in arg_dumps for arg in node.args)
130 ):
131 raise ValueError("another linear consumes the same input")
127 - blocks = [innermost_block(funcdef.body, call) for call in calls]
128 - if any(found is None for found in blocks):
132 +
133 + # Insert the fused GEMM before the earliest call, in the innermost block
134 + # common to all three (the calls may be split across branches).
135 + chains = [block_chain(funcdef.body, call) for call in calls]
136 + if any(not chain for chain in chains):
137 raise ValueError("projection calls not found in the function body")
130 - if len({id(block) for block, _ in blocks}) != 1:
131 - raise ValueError("projection calls are in different blocks")
138 + depth = 0
139 + for level in zip(*chains):
140 + if len({id(block) for block, _ in level}) != 1:
141 + break
142 + depth += 1
143 + block = chains[0][depth - 1][0]
144 + indices = [chain[depth - 1][1] for chain in chains]
145 + insert_index = min(indices)
146 + self._check_input_stable(funcdef, module, calls, block, indices)
147
148 # q(x), k(x), v(x) -> q, k, v = qkv(x).split(qkv.output_sizes / qkv.tp_size, -1)
134 - block = blocks[0][0]
135 - index = min(index for _, index in blocks)
136 - self._splice_merged_split(funcdef, calls, block, index)
149 + self._splice_merged_split(funcdef, calls, block, insert_index)
150 self.fused_forward = compile_forward(funcdef, fn)
151
152 def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
vllm/model_executor/models/transformers/fx_utils.py
+326 -21
@@ -70,6 +70,13 @@ class _MetaProxy(fx.Proxy):
70 def __getattr__(self, k: str) -> "_MetaAttribute":
71 return _MetaAttribute(self, k)
72
73 + def __setitem__(self, key: object, value: object) -> None:
74 + # Record `obj[key] = value` so tracing continues past in-place writes
75 + # into containers (e.g. Gemma 4's `shared_kv_states[...] = ...`).
76 + self.tracer.create_proxy(
77 + "call_function", operator.setitem, (self, key, value), {}
78 + )
79 +
80
81 class _MetaAttribute(_MetaProxy, fx.proxy.Attribute):
82 """Attribute proxy (e.g. `x.shape`) carrying its meta value.
@@ -312,41 +319,125 @@ def compile_forward(funcdef: ast.FunctionDef, fn: Callable) -> Callable:
319 return namespace[funcdef.name]
320
321
315 -def single_self_call(funcdef: ast.FunctionDef, name: str) -> ast.Call:
316 - """The unique `self.<name>(arg)` call in `funcdef`.
322 +def self_call_and_refs(
323 + funcdef: ast.FunctionDef, name: str
324 +) -> tuple[ast.Call, list[ast.Attribute]]:
325 + """The unique `self.<name>(arg)` call, plus every other `self.<name>` reference.
326
318 - Raises unless `name` appears exactly once, as such a call, so the source
319 - rewrite agrees with the fx match.
327 + Raises unless exactly one single-argument `self.<name>(arg)` call exists, so
328 + one fx `call_module` node maps to one syntactic call site. The remaining
329 + references (e.g. `is not None` guards) are returned for the caller to resolve.
330 """
331 uses = [
332 node
333 for node in ast.walk(funcdef)
324 - if isinstance(node, ast.Attribute) and node.attr == name
334 + if isinstance(node, ast.Attribute)
335 + and node.attr == name
336 + and isinstance(node.value, ast.Name)
337 + and node.value.id == "self"
338 ]
326 - if len(uses) != 1:
327 - raise ValueError(f"{name} is referenced {len(uses)} times")
339 calls = [
340 node
341 for node in ast.walk(funcdef)
342 if isinstance(node, ast.Call)
332 - and node.func is uses[0]
343 + and any(node.func is use for use in uses)
344 and len(node.args) == 1
345 and not isinstance(node.args[0], ast.Starred)
346 and not node.keywords
347 ]
337 - if (
338 - len(calls) != 1
339 - or not isinstance(uses[0].value, ast.Name)
340 - or uses[0].value.id != "self"
341 - ):
342 - raise ValueError(f"{name} is not a single-argument call on self")
343 - return calls[0]
348 + if len(calls) != 1:
349 + raise ValueError(f"{name} is not called exactly once as self.{name}(arg)")
350 + call = calls[0]
351 + other_refs = [use for use in uses if use is not call.func]
352 + return call, other_refs
353 +
354 +
355 +def single_self_call(funcdef: ast.FunctionDef, name: str) -> ast.Call:
356 + """The unique `self.<name>(arg)` call in `funcdef`.
357 +
358 + Raises unless `name` appears exactly once, as such a call, so the source
359 + rewrite agrees with the fx match.
360 + """
361 + call, other_refs = self_call_and_refs(funcdef, name)
362 + if other_refs:
363 + raise ValueError(f"{name} is referenced {len(other_refs) + 1} times")
364 + return call
365 +
366 +
367 +def _is_none(node: ast.expr) -> bool:
368 + return isinstance(node, ast.Constant) and node.value is None
369 +
370 +
371 +def _in_boolean_context(funcdef: ast.FunctionDef, ref: ast.expr) -> bool:
372 + """Is `ref` used only for its truth value (a test, or a `not` operand)?
373
374 + In these positions the object's identity never escapes, so a reference that
375 + is always truthy can be replaced by `True`. `and`/`or` are excluded: they
376 + yield an operand, so the module could escape (`x and self.<name>`)."""
377 + for node in ast.walk(funcdef):
378 + if (
379 + isinstance(node, (ast.If, ast.IfExp, ast.While, ast.Assert))
380 + and node.test is ref
381 + ):
382 + return True
383 + if (
384 + isinstance(node, ast.UnaryOp)
385 + and isinstance(node.op, ast.Not)
386 + and node.operand is ref
387 + ):
388 + return True
389 + return False
390 +
391 +
392 +def bypass_existence_guard(
393 + funcdef: ast.FunctionDef, ref: ast.Attribute, name: str
394 +) -> None:
395 + """Fold a guard on `self.<name>`'s existence to its constant value.
396
346 -def innermost_block(
397 + The fuser deletes `self.<name>` and binds only instances where it exists as a
398 + truthy `nn.Linear` (guaranteed by the match, the cache key, and `validate`),
399 + so a guard testing its presence is a fusion invariant. Two forms are folded:
400 +
401 + - an identity `None` check `self.<name> is None` -> `False`,
402 + `self.<name> is not None` -> `True` (either operand order). `==`/`!=` are
403 + *not* folded: `is_linear` accepts `nn.Linear` subclasses, which may override
404 + `__eq__`/`__ne__`, so equality is not guaranteed to track identity.
405 + - a bare truthiness test (`if self.<name>:`, `... if self.<name> else ...`,
406 + `not self.<name>`): the reference itself to `True`.
407 +
408 + Any other surviving reference escapes the projection's value, which no longer
409 + exists after fusion, so refuse rather than change semantics."""
410 + for node in ast.walk(funcdef):
411 + # `self.<name> is (not) None`, either operand order.
412 + if not (isinstance(node, ast.Compare) and len(node.ops) == 1):
413 + continue
414 + (op,), (right,) = node.ops, node.comparators
415 + if not isinstance(op, (ast.Is, ast.IsNot)):
416 + continue
417 + if (ref is node.left and _is_none(right)) or (
418 + ref is right and _is_none(node.left)
419 + ):
420 + value = isinstance(op, ast.IsNot)
421 + replace_expr(funcdef, node, ast.copy_location(ast.Constant(value), node))
422 + return
423 + # A bare truthiness test (statement `if`, ternary, or `not`); an `nn.Linear`
424 + # is always truthy, so the reference folds to `True`.
425 + if _in_boolean_context(funcdef, ref):
426 + replace_expr(funcdef, ref, ast.copy_location(ast.Constant(True), ref))
427 + return
428 + raise ValueError(f"{name} is referenced outside an existence guard")
429 +
430 +
431 +def block_chain(
432 block: list[ast.stmt], node: ast.AST
348 -) -> tuple[list[ast.stmt], int] | None:
349 - """The innermost statement list containing `node`, and the index within."""
433 +) -> list[tuple[list[ast.stmt], int]]:
434 + """Path of (statement list, index) pairs from `block` down to `node`.
435 +
436 + Each pair names a nested block and the index in it of the statement
437 + containing `node`; the last pair is `node`'s innermost block. Empty if
438 + `node` is not in `block`. The longest common prefix of several nodes' chains
439 + is the innermost block that dominates them all.
440 + """
441 for index, stmt in enumerate(block):
442 if not any(child is node for child in ast.walk(stmt)):
443 continue
@@ -359,13 +450,227 @@ def innermost_block(
450 if (
451 isinstance(child_block, list)
452 and child_block
362 - and (found := innermost_block(child_block, node)) is not None
453 + and (tail := block_chain(child_block, node))
454 ):
364 - return found
365 - return block, index
455 + return [(block, index), *tail]
456 + return [(block, index)]
457 + return []
458 +
459 +
460 +def _base_name(node: ast.expr) -> str | None:
461 + """The root `Name` of an attribute/subscript chain (`a.b[c]` -> `a`)."""
462 + while isinstance(node, (ast.Attribute, ast.Subscript)):
463 + node = node.value
464 + return node.id if isinstance(node, ast.Name) else None
465 +
466 +
467 +def _rebound_names(region: list[ast.stmt]) -> set[str]:
468 + """Names rebound outright within `region` (`x = ...`, `del x`).
469 +
470 + These are `Name` nodes in a `Store`/`Del` context."""
471 + return {
472 + node.id
473 + for stmt in region
474 + for node in ast.walk(stmt)
475 + if isinstance(node, ast.Name) and isinstance(node.ctx, (ast.Store, ast.Del))
476 + }
477 +
478 +
479 +def _inplace_target(node: ast.AST) -> str | None:
480 + """Base name `node` mutates in place, if any.
481 +
482 + A write through the name (`x[i] = ...`, `x.attr = ...`) or an in-place method
483 + call (`x.mul_(...)`) leaves the base name in a `Load` context, so it is not a
484 + plain `Name` store (see `_rebound_names`)."""
485 + if isinstance(node, (ast.Attribute, ast.Subscript)) and isinstance(
486 + node.ctx, (ast.Store, ast.Del)
487 + ):
488 + return _base_name(node)
489 + if (
490 + isinstance(node, ast.Call)
491 + and isinstance(node.func, ast.Attribute)
492 + and node.func.attr.endswith("_")
493 + and not node.func.attr.endswith("__")
494 + ):
495 + return _base_name(node.func.value)
496 return None
497
498
499 +def _mutated_names(region: list[ast.stmt]) -> set[str]:
500 + """Names mutated in place within `region` (see `_inplace_target`)."""
501 + return {
502 + base
503 + for stmt in region
504 + for node in ast.walk(stmt)
505 + if (base := _inplace_target(node)) is not None
506 + }
507 +
508 +
509 +def _call_name(call: ast.Call) -> str:
510 + """The called function's own name, without its qualifier."""
511 + if isinstance(call.func, ast.Attribute):
512 + return call.func.attr
513 + return call.func.id if isinstance(call.func, ast.Name) else ""
514 +
515 +
516 +def _arg_names(call: ast.Call) -> set[str]:
517 + """Every name `call` reads through its arguments."""
518 + reads = [*call.args, *(kw.value for kw in call.keywords)]
519 + return {
520 + node.id
521 + for arg in reads
522 + for node in ast.walk(arg)
523 + if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load)
524 + }
525 +
526 +
527 +def _writes_through_args(call: ast.Call) -> bool:
528 + """Does `call` write through an argument rather than return a new value?
529 +
530 + Torch spells this three ways: an `out=` destination, `inplace=True`, and the
531 + trailing-underscore convention (`torch.relu_(x)`). The underscore is checked
532 + on the function's own name, so it catches the free-function form that
533 + `_inplace_target` cannot (there the base name is the module, `torch`).
534 + """
535 + for keyword in call.keywords:
536 + if keyword.arg == "out":
537 + return True
538 + if keyword.arg == "inplace" and not (
539 + isinstance(keyword.value, ast.Constant) and keyword.value.value is False
540 + ):
541 + return True
542 + name = _call_name(call)
543 + return name.endswith("_") and not name.endswith("__")
544 +
545 +
546 +def written_names(region: list[ast.stmt]) -> set[str]:
547 + """Names `region` may rebind or write through.
548 +
549 + Beyond plain rebinding and writes through the name itself (`_rebound_names`,
550 + `_mutated_names`), a call can write through an argument without naming it as
551 + a target: `F.relu(x, inplace=True)`, `torch.add(x, y, out=x)`,
552 + `torch.relu_(x)`. A statement-level call whose result is discarded is treated
553 + the same way, since it can only be there for a side effect.
554 + """
555 + written = _rebound_names(region) | _mutated_names(region)
556 + for stmt in region:
557 + if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Call):
558 + written |= _arg_names(stmt.value)
559 + for node in ast.walk(stmt):
560 + if isinstance(node, ast.Call) and _writes_through_args(node):
561 + written |= _arg_names(node)
562 + return written
563 +
564 +
565 +_METADATA_ATTRS = frozenset(
566 + {"shape", "dtype", "device", "ndim", "size", "dim", "numel", "item"}
567 +)
568 +
569 +
570 +def _aliasing_reads(node: ast.expr) -> set[str]:
571 + """Names `node` reads in a way that could yield a view of them.
572 +
573 + Metadata (`x.shape`, `x.size(0)`, `x.dtype`) is ints and tuples sharing no
574 + storage with `x`, so reaching a name only through metadata cannot alias it.
575 + Shape arithmetic is pervasive between projections, so propagating through it
576 + would taint nearly every later name.
577 + """
578 + names: set[str] = set()
579 + stack = [node]
580 + while stack:
581 + current = stack.pop()
582 + if isinstance(current, ast.Attribute) and current.attr in _METADATA_ATTRS:
583 + continue
584 + if (
585 + isinstance(current, ast.Call)
586 + and isinstance(current.func, ast.Attribute)
587 + and current.func.attr in _METADATA_ATTRS
588 + ):
589 + stack.extend(current.args)
590 + stack.extend(keyword.value for keyword in current.keywords)
591 + continue
592 + if isinstance(current, ast.Name) and isinstance(current.ctx, ast.Load):
593 + names.add(current.id)
594 + stack.extend(ast.iter_child_nodes(current))
595 + return names
596 +
597 +
598 +def _returns_fresh(node: ast.expr, module: nn.Module) -> bool:
599 + """Does evaluating `node` allocate, rather than return a view of its input?
600 +
601 + A `nn.Linear` writes its output into new storage, so its result aliases
602 + nothing that went in, and a chain of view ops on top of it aliases only that
603 + fresh tensor. Walking the receiver chain down to a `self.<linear>(...)` call
604 + therefore proves the value cannot alias the projection input. A literal
605 + aliases nothing either, and a branch or sequence is fresh when every part is
606 + -- which is how a guarded projection (`p(x) if p is not None else None`)
607 + stays untracked. A call on an attribute that is missing or `None` cannot run
608 + at all, so a branch selecting a different configuration is fresh too.
609 + """
610 + if isinstance(node, ast.Constant):
611 + return True
612 + if isinstance(node, ast.IfExp):
613 + if isinstance(node.test, ast.Constant):
614 + # A folded existence guard leaves the untaken branch as dead code.
615 + taken = node.body if node.test.value else node.orelse
616 + return _returns_fresh(taken, module)
617 + return _returns_fresh(node.body, module) and _returns_fresh(node.orelse, module)
618 + if isinstance(node, (ast.Tuple, ast.List)):
619 + return all(_returns_fresh(elt, module) for elt in node.elts)
620 + while isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
621 + func = node.func
622 + if isinstance(func.value, ast.Name) and func.value.id == "self":
623 + attribute = getattr(module, func.attr, None)
624 + if attribute is None or isinstance(attribute, nn.Linear):
625 + # A missing or `None` attribute cannot be called, so the branch
626 + # selecting it never runs (`q_proj` on a q-LoRA layer).
627 + return True
628 + # Anything else can only return storage it was handed, so it is
629 + # fresh exactly when its inputs are (`q_norm(q_proj(x))`).
630 + return bool(node.args) and all(
631 + _returns_fresh(arg, module) for arg in node.args
632 + )
633 + node = func.value
634 + return False
635 +
636 +
637 +def aliasing_names(
638 + funcdef: ast.FunctionDef, seed: set[str], module: nn.Module
639 +) -> set[str]:
640 + """`seed` plus every name that may alias one of them.
641 +
642 + A view shares storage with its base (`h = x.view(-1)`), so mutating `h`
643 + mutates `x`; tracking only `seed` would miss that. Any assignment reading a
644 + tracked name therefore taints its targets, transitively, unless its value is
645 + freshly allocated (`_returns_fresh`).
646 + """
647 + names = set(seed)
648 + assigns = [
649 + node
650 + for node in ast.walk(funcdef)
651 + if isinstance(node, (ast.Assign, ast.AnnAssign))
652 + and node.value is not None
653 + and not _returns_fresh(node.value, module)
654 + ]
655 + changed = True
656 + while changed:
657 + changed = False
658 + for node in assigns:
659 + if not _aliasing_reads(node.value) & names:
660 + continue
661 + targets = getattr(node, "targets", None) or [node.target]
662 + stores = {
663 + child.id
664 + for target in targets
665 + for child in ast.walk(target)
666 + if isinstance(child, ast.Name) and isinstance(child.ctx, ast.Store)
667 + }
668 + if stores - names:
669 + names |= stores
670 + changed = True
671 + return names
672 +
673 +
674 def replace_expr(module: ast.AST, old: ast.expr, new: ast.expr) -> None:
675 """Replace the expression `old` (by identity) with `new` within `module`."""
676