# SPDX-License-Identifier: Apache-2.0 # 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 import pytest import torch import torch.nn as nn import torch.nn.functional as F from vllm.model_executor.models.transformers.fuser import Fusers, get_fuser, get_fusers from vllm.model_executor.models.transformers.fusers import ( GLUFuser, MergedColumnParallelFuser, PackedQKVFuser, QKVFuser, packed_qkv, qkv, ) from vllm.model_executor.models.transformers.fx_utils import ( bypass_existence_guard, trace, ) class SiluAndMulStub(nn.Module): """Stand-in for vLLM's `SiluAndMul` (no vLLM config required).""" def forward(self, x: torch.Tensor) -> torch.Tensor: d = x.shape[-1] // 2 return F.silu(x[..., :d]) * x[..., d:] class NoDownGLU(nn.Module): """`act(gate(x)) * up(x)` with no output projection -> `down_name` is None.""" def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False): super().__init__() self.gate_proj = nn.Linear(hidden, inter, bias=bias) self.up_proj = nn.Linear(hidden, inter, bias=bias) self.act_fn = nn.SiLU() def forward(self, x): return self.act_fn(self.gate_proj(x)) * self.up_proj(x) class GLUMLP(NoDownGLU): """`down(act(gate(x)) * up(x))` — the canonical HF GLU MLP.""" def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False): super().__init__(hidden, inter, bias) self.down_proj = nn.Linear(inter, hidden, bias=bias) def forward(self, x): return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) class ReversedGLUMLP(GLUMLP): """`up(x) * act(gate(x))` — operands swapped (multiply is commutative).""" def forward(self, x): return self.down_proj(self.up_proj(x) * self.act_fn(self.gate_proj(x))) class NotAnMLP(nn.Module): """Two linears but no activation*linear multiply -> must not match.""" def __init__(self): super().__init__() self.fc1 = nn.Linear(8, 8) self.fc2 = nn.Linear(8, 8) def forward(self, x): return self.fc2(self.fc1(x)) class NotAnActGLUMLP(GLUMLP): """GLU-shaped, but the "activation" is not a known activation module.""" def __init__(self): super().__init__() self.act_fn = nn.Dropout() class BroadcastGLU(NoDownGLU): """A broadcast multiply cannot use an equal-halves AndMul kernel.""" def __init__(self): super().__init__() self.gate_proj = nn.Linear(16, 1, bias=False) class ExtraProjGLU(GLUMLP): """A third linear on the same input (e.g. a router) -> gate/up still fuse.""" def __init__(self): super().__init__() self.router = nn.Linear(16, 4, bias=False) def forward(self, x): gated = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) return gated, self.router(x) class UntraceableMLP(GLUMLP): """Data-dependent control flow *before* the GLU -> no match.""" def forward(self, x): if x.sum() > 0: # noqa: SIM108 - intentionally untraceable return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) return x class UntraceableTailGLUMLP(GLUMLP): """Data-dependent control flow *after* the GLU -> still fusable.""" def forward(self, x): y = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) if y.sum() > torch.inf: # intentionally untraceable y = y * 0 return y class FakeAttention(nn.Module): """HF v5-style attention: shape unpacking, dead KV branch, kwargs interface.""" is_causal = True def __init__( self, hidden: int = 32, head_dim: int = 8, heads: int = 4, kv_heads: int = 4, bias: bool = False, layer_idx: int = 0, ): super().__init__() self.config = SimpleNamespace(_attn_implementation="vllm") self.layer_idx = layer_idx self.head_dim = head_dim self.scaling = head_dim**-0.5 self.q_proj = nn.Linear(hidden, heads * head_dim, bias=bias) self.k_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias) self.v_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias) self.o_proj = nn.Linear(heads * head_dim, hidden, bias=bias) 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 past_key_values is not None: k, v = past_key_values.update(k, v, self.layer_idx) attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, None ) attn_output, attn_weights = attention_interface( self, q, k, v, attention_mask, scaling=self.scaling, **kwargs ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() return self.o_proj(attn_output), attn_weights class ReversedFakeAttention(FakeAttention): """Projections computed in (v, k, q) order — q must still be identified.""" 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) v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) q = self.q_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 FourParallelLinears(nn.Module): """Four same-input projections of different widths -> one merged linear.""" def __init__(self): super().__init__() self.proj_a = nn.Linear(8, 4) self.proj_b = nn.Linear(8, 6) self.proj_c = nn.Linear(8, 8) self.proj_d = nn.Linear(8, 10) def forward(self, x): 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) x.add_(1) 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.""" def __init__(self, hidden: int = 32, heads: int = 4, head_dim: int = 8): super().__init__() self.q_proj = nn.Linear(hidden, heads * head_dim, bias=False) self.k_proj = nn.Linear(hidden, heads * head_dim, bias=False) self.v_proj = nn.Linear(hidden, heads * head_dim, bias=False) self.head_dim = head_dim def forward(self, x): return self.q_proj(x), self.k_proj(x), self.v_proj(x) class ExtraProjAttention(FakeAttention): """A second non-qkv linear of a different width -> `o_proj` still found.""" def __init__(self, **kwargs): super().__init__(**kwargs) self.sink_proj = nn.Linear(self.head_dim, self.head_dim, bias=False) class QKNormAttention(FakeAttention): """OLMoE-style: a full-dim norm applied to the whole q/k projection output.""" def __init__(self, **kwargs): super().__init__(**kwargs) self.q_norm = nn.RMSNorm(self.q_proj.out_features) self.k_norm = nn.RMSNorm(self.k_proj.out_features) def forward( self, hidden_states, attention_mask=None, past_key_values=None, **kwargs ): q = self.q_norm(self.q_proj(hidden_states)) k = self.k_norm(self.k_proj(hidden_states)) v = self.v_proj(hidden_states) return self.o_proj(q + k + v), None class PerHeadQKNormAttention(FakeAttention): """Qwen3-style: a per-head norm (`head_dim`) applied after the head reshape.""" def __init__(self, **kwargs): super().__init__(**kwargs) self.q_norm = nn.RMSNorm(self.head_dim) self.k_norm = nn.RMSNorm(self.head_dim) def forward( self, hidden_states, attention_mask=None, past_key_values=None, **kwargs ): shape = (*hidden_states.shape[:-1], -1, self.head_dim) q = self.q_norm(self.q_proj(hidden_states).view(shape)) k = self.k_norm(self.k_proj(hidden_states).view(shape)) v = self.v_proj(hidden_states).view(shape) return self.o_proj((q + k + v).flatten(-2)), None class ResidDropoutAttention(FakeAttention): """GPT-style dropout after `o_proj` -> the output projection is still found.""" def __init__(self, **kwargs): super().__init__(**kwargs) self.resid_dropout = nn.Dropout(0.0) 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) 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 ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() 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.""" is_causal = True def __init__( self, hidden: int = 32, head_dim: int = 8, heads: int = 4, kv_heads: int = 1, bias: bool = False, layer_idx: int = 0, ): super().__init__() self.config = SimpleNamespace(_attn_implementation="vllm") self.layer_idx = layer_idx self.head_dim = head_dim self.scaling = head_dim**-0.5 self.embed_dim = heads * head_dim self.kv_dim = kv_heads * head_dim self.c_attn = nn.Linear(hidden, self.embed_dim + 2 * self.kv_dim, bias=bias) self.c_proj = nn.Linear(self.embed_dim, hidden, bias=bias) self.resid_dropout = nn.Dropout(0.0) 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] q, k, v = ( self.c_attn(hidden_states) .unsqueeze(1) .split((self.embed_dim, self.kv_dim, self.kv_dim), dim=3) ) q = q.view(*input_shape, -1, self.head_dim).transpose(1, 2) if past_key_values is not None: k, v = past_key_values.update(k, v, self.layer_idx) attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, None ) attn_output, attn_weights = attention_interface( self, q, k, v, attention_mask, scaling=self.scaling, **kwargs ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() return self.resid_dropout(self.c_proj(attn_output)), attn_weights class PerHeadSplitAttention(nn.Module): """A packed projection reshaped and split *per head* -> not a q/k/v split.""" def __init__(self, hidden: int = 32, head_dim: int = 8, heads: int = 4): super().__init__() self.head_dim = head_dim self.heads = heads self.c_attn = nn.Linear(hidden, 3 * heads * head_dim) self.c_proj = nn.Linear(heads * head_dim, hidden) def forward(self, hidden_states): shape = (*hidden_states.shape[:2], self.heads, 3 * self.head_dim) q, k, v = ( self.c_attn(hidden_states) .view(shape) .transpose(1, 2) .split((self.head_dim, self.head_dim, self.head_dim), dim=3) ) return self.c_proj((q + k + v).transpose(1, 2).flatten(-2)) class FakeSelfAttn(nn.Module): """Stand-in for the vLLM `Attention` attached to the dispatching module.""" def __init__(self): super().__init__() self.impl = SimpleNamespace(scale=None) def forward(self, q, k, v): # MHA-shaped stub: any deterministic combination of q/k/v will do return q + 2 * k + 3 * v class FakeMQASelfAttn(FakeSelfAttn): """Stand-in for grouped/multi-query layouts, where `k`/`v` are narrower.""" def forward(self, q, k, v): groups = q.shape[-1] // k.shape[-1] return q + (2 * k + 3 * v).repeat(1, groups) @pytest.fixture(autouse=True) def _clear_fuser_cache(): get_fusers.cache_clear() yield get_fusers.cache_clear() def _apply_glu_fuser_with_stubs(module: nn.Module, fuser: GLUFuser): """Apply a fuser using plain stand-ins (merged `nn.Linear` + silu AndMul).""" gate = module.get_submodule(fuser.gate_name) up = module.get_submodule(fuser.up_name) merged = nn.Linear( gate.in_features, gate.out_features + up.out_features, bias=gate.bias is not None, ) with torch.no_grad(): merged.weight.copy_(torch.cat([gate.weight, up.weight], dim=0)) if gate.bias is not None: merged.bias.copy_(torch.cat([gate.bias, up.bias], dim=0)) setattr(module, fuser.merged_name, merged) setattr(module, fuser.act_name, SiluAndMulStub()) delattr(module, fuser.gate_name) delattr(module, fuser.up_name) module.forward = MethodType(fuser.fused_forward, module) return module def _apply_qkv_fuser_with_stubs(module: nn.Module, fuser: QKVFuser): """Apply a fuser using a plain merged `nn.Linear` (no TP sharding).""" q, k, v = ( module.get_submodule(name) for name in (fuser.q_name, fuser.k_name, fuser.v_name) ) merged = nn.Linear( q.in_features, q.out_features + k.out_features + v.out_features, bias=q.bias is not None, ) with torch.no_grad(): merged.weight.copy_(torch.cat([q.weight, k.weight, v.weight], dim=0)) if q.bias is not None: merged.bias.copy_(torch.cat([q.bias, k.bias, v.bias], dim=0)) merged.output_sizes = [q.out_features, k.out_features, v.out_features] merged.tp_size = 1 setattr(module, fuser.merged_name, merged) for name in (fuser.q_name, fuser.k_name, fuser.v_name): delattr(module, name) module.forward = MethodType(fuser.fused_forward, module) return module def _apply_packed_qkv_fuser_with_stubs(module: nn.Module, fuser: PackedQKVFuser): """Apply a fuser at `tp_size == 1`, where the rewritten split is unchanged.""" qkv = module.get_submodule(fuser.qkv_name) qkv.output_sizes = [fuser.q_size, fuser.kv_size, fuser.kv_size] qkv.tp_size = 1 module.forward = MethodType(fuser.fused_forward, module) return module @pytest.mark.parametrize("mlp_cls", [GLUMLP, ReversedGLUMLP]) @pytest.mark.parametrize("bias", [False, True]) def test_detects_and_rewrites_glu(mlp_cls, bias): with torch.device("meta"): meta = mlp_cls(bias=bias) fuser = get_fuser(meta, GLUFuser) assert isinstance(fuser, GLUFuser) assert ( fuser.gate_name, fuser.up_name, fuser.act_name, fuser.down_name, ) == ("gate_proj", "up_proj", "act_fn", "down_proj") # The rewritten forward references the merged projection instead of the # sources; the rest of the forward is untouched. names = fuser.fused_forward.__code__.co_names assert "gate_up_proj" in names and "act_fn" in names and "down_proj" in names assert not {"gate_proj", "up_proj"} & set(names) # Numerics: the fused forward must match the original on a real instance. real = mlp_cls(bias=bias) for p in real.parameters(): nn.init.normal_(p, std=0.05) x = torch.randn(4, 16) expected = real(x) fused = _apply_glu_fuser_with_stubs(real, fuser) # Fusion is in place: the module keeps its class and other attributes assert fused is real and type(fused) is mlp_cls torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5) def test_glu_fuses_alongside_unrelated_projections(): """Sibling linears the GLU does not consume must not block the fusion.""" with torch.device("meta"): fuser = get_fuser(ExtraProjGLU(), GLUFuser) assert fuser is not None assert (fuser.gate_name, fuser.up_name) == ("gate_proj", "up_proj") real = ExtraProjGLU() x = torch.randn(4, 16) expected = real(x) fused = _apply_glu_fuser_with_stubs(real, fuser) torch.testing.assert_close(fused(x), expected) def test_glu_identifies_down_projection(): """The row projection consuming `act(gate(x)) * up(x)` is identified. It is forced to `RowParallelLinear` in `update_attrs` so its sharded input matches the column-parallel merged gate/up; `None` when there is no such projection to force (fusion of gate/up still applies).""" with torch.device("meta"): assert get_fuser(GLUMLP(), GLUFuser).down_name == "down_proj" assert get_fuser(ReversedGLUMLP(), GLUFuser).down_name == "down_proj" assert get_fuser(NoDownGLU(), GLUFuser).down_name is None @pytest.mark.parametrize("attn_cls", [FakeAttention, ReversedFakeAttention]) @pytest.mark.parametrize("kv_heads", [4, 2]) def test_detects_and_rewrites_qkv(attn_cls, kv_heads): if attn_cls is ReversedFakeAttention and kv_heads == 4: pytest.skip("MHA q/k/v assignment is order-based by design") with torch.device("meta"): meta = attn_cls(kv_heads=kv_heads) fuser = get_fuser(meta, QKVFuser) assert isinstance(fuser, QKVFuser) # q (sharded differently under TP) must be identified exactly; k/v may be # swapped for non-canonical compute order, which is numerically consistent # because the weight mapping and the split indices follow the same # assignment. assert fuser.q_name == "q_proj" assert {fuser.k_name, fuser.v_name} == {"k_proj", "v_proj"} assert fuser.o_name == "o_proj" # The projections are merged; everything else stays live Python with its # original semantics (branches, kwargs, attribute reads) code = fuser.fused_forward.__code__ names = code.co_names assert "qkv_proj" in names and "output_sizes" in names and "o_proj" in names assert "tp_size" in names assert not {"q_proj", "k_proj", "v_proj"} & set(names) if attn_cls is FakeAttention: assert "update" in names # the cache branch survives assert code.co_flags & inspect.CO_VARKEYWORDS # **kwargs survives # Numerics: the fused forward must match the original on a real instance, # with a different layer_idx than the traced instance (kv_heads == heads so # the q/k/v stub combination is shape-compatible). 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) # Fusion is in place: the module keeps its class and other attributes assert fused is real and type(fused) is attn_cls assert fused.layer_idx == 3 and fused.is_causal and fused.config is not None out, _ = fused(x) torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5) def test_qkv_identifies_output_projection(): with torch.device("meta"): assert get_fuser(FakeAttention(), QKVFuser).o_name == "o_proj" assert get_fuser(ReversedFakeAttention(), QKVFuser).o_name == "o_proj" assert get_fuser(ExtraProjAttention(), QKVFuser).o_name == "o_proj" # Norm children (q_norm/k_norm) must not disturb o_proj identification. assert get_fuser(QKNormAttention(), QKVFuser).o_name == "o_proj" assert get_fuser(PerHeadQKNormAttention(), QKVFuser).o_name == "o_proj" # A module between o_proj and the return is transparent. assert get_fuser(ResidDropoutAttention(), QKVFuser).o_name == "o_proj" @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"): 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 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( cls, default_vllm_config, monkeypatch ): from vllm.model_executor import parameter from vllm.model_executor.layers import linear from vllm.model_executor.layers.linear import MergedColumnParallelLinear from vllm.model_executor.layers.utils import default_unquantized_gemm from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper monkeypatch.setattr(linear, "get_tensor_model_parallel_rank", lambda: 0) monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: 1) monkeypatch.setattr( linear, "dispatch_unquantized_gemm", lambda linear_backend="auto": default_unquantized_gemm, ) 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 = cls() fuser = MergedColumnParallelFuser.match(trace(meta), meta) assert fuser is not None 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 = cls() x = torch.randn(2, 8) expected = real(x) weights = [ (name, weight.detach().clone()) for name, weight in real.named_parameters() ] fused = fuser.fuse(real, "parallel", default_vllm_config) root = nn.Module() root.parallel = fused mapper = WeightsMapper(orig_to_new_stacked=fuser.orig_to_new_stacked("parallel")) AutoWeightsLoader(root).load_weights( ((f"parallel.{name}", weight) for name, weight in weights), mapper=mapper ) merged = getattr(fused, fuser.merged_name) assert isinstance(merged, MergedColumnParallelLinear) for actual, reference in zip(fused(x), expected): torch.testing.assert_close(actual, reference) def test_merged_name_is_unique_per_linear_names(): """The computed name must not collide across different projections, so stacking two unrelated fusions can never clobber packed_modules_mapping.""" ab_c = MergedColumnParallelFuser(source_cls="M", linear_names=("a_b", "c")) a_bc = MergedColumnParallelFuser(source_cls="M", linear_names=("a", "b_c")) assert ab_c.merged_name != a_bc.merged_name # Same names -> same (and therefore safely re-mergeable) name. dup = MergedColumnParallelFuser(source_cls="M", linear_names=("a_b", "c")) assert ab_c.merged_name == dup.merged_name def test_heterogeneous_instances_fall_back_to_merged_column_fuser( default_vllm_config, ): """Same class and structure -> both fusers are cached as candidates. A layer whose out_features aren't a multiple of its own head_dim (as a misconfigured or heterogeneous checkpoint might have) fails QKVFuser's validation and must still fuse via the generic fallback, while a well-formed sibling of the same class keeps using QKVFuser.""" with torch.device("meta"): container = nn.Module() container.compatible = PlainQKV(head_dim=8) container.incompatible = PlainQKV(head_dim=8) container.incompatible.head_dim = 5 # 32 % 5 != 0, unlike q_proj's 32 % 8 candidates = get_fusers(container.compatible) assert any(isinstance(f, QKVFuser) for f in candidates) assert any(isinstance(f, MergedColumnParallelFuser) for f in candidates) fusers = Fusers(container, default_vllm_config) assert isinstance(fusers[container.compatible][0], QKVFuser) assert isinstance(fusers[container.incompatible][0], MergedColumnParallelFuser) @pytest.mark.parametrize("kv_heads", [1, 2]) def test_detects_and_rewrites_packed_qkv(kv_heads): """A single projection split into q/k/v must be re-sharded, not merged. Only the split sizes change: `QKVParallelLinear` loads the packed checkpoint weight as-is, and shards q by heads while replicating k/v.""" with torch.device("meta"): meta = PackedQKVAttention(kv_heads=kv_heads) fuser = get_fuser(meta, PackedQKVFuser) assert isinstance(fuser, PackedQKVFuser) assert (fuser.qkv_name, fuser.o_name) == ("c_attn", "c_proj") assert (fuser.q_size, fuser.kv_size) == (32, 8 * kv_heads) # The hard-coded widths become the per-rank widths of the sharded linear names = fuser.fused_forward.__code__.co_names assert "output_sizes" in names and "tp_size" in names assert "kv_dim" not in names and "embed_dim" not in names # Numerics: the rewritten forward must match the original on a real instance real = PackedQKVAttention(kv_heads=kv_heads, layer_idx=3) for p in real.parameters(): nn.init.normal_(p, std=0.05) x = torch.randn(1, 5, 32) real.attn = FakeMQASelfAttn() expected, _ = real(x) fused = _apply_packed_qkv_fuser_with_stubs(real, fuser) # Fusion is in place: the module keeps its class and other attributes assert fused is real and type(fused) is PackedQKVAttention assert fused.layer_idx == 3 and fused.is_causal out, _ = fused(x) torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5) def test_per_head_split_is_not_packed_qkv(): """The split must consume the whole projection, else its sizes are head widths and re-sharding by them would be wrong.""" with torch.device("meta"): assert get_fuser(PerHeadSplitAttention(), PackedQKVFuser) is None def test_fuser_is_cached_per_class_and_structure(): with torch.device("meta"): fuser_a = get_fuser(GLUMLP(), GLUFuser) fuser_b = get_fuser(GLUMLP(), GLUFuser) assert fuser_a is fuser_b assert any(key[0] is GLUMLP for key in get_fusers.cache) @pytest.mark.parametrize("cls", [NotAnMLP, UntraceableMLP]) def test_non_matching_modules_return_none(cls): with torch.device("meta"): module = cls() assert get_fuser(module, GLUFuser) is None def test_untraceable_tail_still_fuses(): with torch.device("meta"): meta = UntraceableTailGLUMLP() fuser = get_fuser(meta, GLUFuser) assert isinstance(fuser, GLUFuser) # Numerics: the live tail must survive the rewrite real = UntraceableTailGLUMLP() for p in real.parameters(): nn.init.normal_(p, std=0.05) x = torch.randn(4, 16) expected = real(x) fused = _apply_glu_fuser_with_stubs(real, fuser) torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5) def test_weight_mappings_are_scoped_to_fused_prefixes(): from vllm.model_executor.models.utils import WeightsMapper with torch.device("meta"): glu_fuser = get_fuser(GLUMLP(), GLUFuser) qkv_fuser = get_fuser(FakeAttention(), QKVFuser) mapper = WeightsMapper() for prefix in ("model.layers.0.mlp", "model.layers.1.mlp"): mapper.orig_to_new_stacked.update(glu_fuser.orig_to_new_stacked(prefix)) mapper.orig_to_new_stacked.update( qkv_fuser.orig_to_new_stacked("model.layers.0.self_attn") ) names = [ "model.layers.0.mlp.gate_proj.weight", "model.layers.0.mlp.up_proj.weight", "model.layers.1.mlp.gate_proj.weight", "model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.k_proj.weight", "model.layers.0.self_attn.v_proj.weight", # Unfused modules at other prefixes must be left untouched. "model.layers.2.mlp.experts.0.gate_proj.weight", "model.layers.1.self_attn.q_proj.weight", ] # `apply` rewrites the name and stamps the shard id onto each tensor. weights = [(name, torch.empty(0)) for name in names] mapped = list(mapper.apply(weights)) mapped_names = [name for name, _ in mapped] shard_ids = [getattr(data, "shard_id", None) for _, data in mapped] assert mapped_names == [ "model.layers.0.mlp.gate_up_proj.weight", "model.layers.0.mlp.gate_up_proj.weight", "model.layers.1.mlp.gate_up_proj.weight", "model.layers.0.self_attn.qkv_proj.weight", "model.layers.0.self_attn.qkv_proj.weight", "model.layers.0.self_attn.qkv_proj.weight", # Only the exact fused layers are remapped; everything else is untouched. "model.layers.2.mlp.experts.0.gate_proj.weight", "model.layers.1.self_attn.q_proj.weight", ] assert shard_ids == [0, 1, 0, "q", "k", "v", None, None] # The fused layers are exposed to the quantization machinery via their # original constituent projection names (what the checkpoint stores). assert glu_fuser.packed_modules_mapping == { "gate_up_proj": ["gate_proj", "up_proj"], } assert qkv_fuser.packed_modules_mapping == { "qkv_proj": ["q_proj", "k_proj", "v_proj"], } @pytest.mark.parametrize("cls", [NotAnMLP, NotAnActGLUMLP, BroadcastGLU]) def test_unfusable_modules_are_not_fused(cls, default_vllm_config): with torch.device("meta"): module = cls() fuser = get_fuser(module, GLUFuser) # Either no pattern matches the class, or this instance fails validation # (`recursive_replace` gates fusion and its weight mappings on `validate`) assert fuser is None or not fuser.validate(module, default_vllm_config) def test_act_and_mul_derived_from_module(default_vllm_config): from transformers.activations import GELUTanh, SiLUActivation from vllm.model_executor.layers.activation import GeluAndMul, SiluAndMul assert isinstance(GLUFuser._get_act_and_mul(nn.SiLU()), SiluAndMul) assert isinstance(GLUFuser._get_act_and_mul(SiLUActivation()), SiluAndMul) gelu_tanh = GLUFuser._get_act_and_mul(GELUTanh()) assert isinstance(gelu_tanh, GeluAndMul) and gelu_tanh.approximate == "tanh" gelu = GLUFuser._get_act_and_mul(nn.GELU()) assert isinstance(gelu, GeluAndMul) and gelu.approximate == "none" # Not activations at all -> no fusion assert GLUFuser._get_act_and_mul_name(nn.Dropout()) is None assert GLUFuser._get_act_and_mul_name(nn.LayerNorm(8)) is None with pytest.raises(ValueError, match="No AndMul equivalent"): GLUFuser._get_act_and_mul(nn.Dropout()) def _wider_model_config(head_dim: int) -> SimpleNamespace: """A model whose global head size is twice `head_dim`, as a wider layer elsewhere in a heterogeneous checkpoint would make it.""" return SimpleNamespace( model_config=SimpleNamespace(get_head_size=lambda: 2 * head_dim), quant_config=None, ) @pytest.mark.parametrize( "cls, fuser_module, fuser_cls", [ (FakeAttention, qkv, QKVFuser), (PackedQKVAttention, packed_qkv, PackedQKVFuser), ], ) def test_head_counts_come_from_the_module_not_the_model( cls, fuser_module, fuser_cls, monkeypatch ): """A layer narrower than the model-wide head size must not be miscounted. On a heterogeneous checkpoint (Gemma 4) the model-wide head size is the largest across layers, so deriving `total_num_heads = out_features // head_size` from it undercounts heads on a narrower layer. The widths still add up, so nothing raises below TP=4: the layer is just sharded wrong. """ head_dim, heads, kv_heads = 8, 8, 4 vllm_config = _wider_model_config(head_dim) with torch.device("meta"): module = cls(hidden=32, head_dim=head_dim, heads=heads, kv_heads=kv_heads) # Both replacements shard, so they need a TP group; only the head counts # the fuser derives are under test here. captured = {} monkeypatch.setattr( fuser_module, "QKVParallelLinear", lambda **kwargs: captured.update(kwargs) or nn.Identity(), ) monkeypatch.setattr( fuser_module, "replace_linear_class", lambda *a, **kw: nn.Identity() ) fuser = get_fuser(module, fuser_cls) assert fuser is not None and fuser.validate(module, vllm_config) fuser.update_attrs(module, "model.layers.0.self_attn", vllm_config) assert captured["head_size"] == head_dim assert captured["total_num_heads"] == heads assert captured["total_num_kv_heads"] == kv_heads def test_validate_accepts_a_layer_the_model_wide_head_size_would_reject(): """`validate` gates fusion, so a wrong head size silently disables it.""" head_dim, heads, kv_heads = 8, 8, 3 vllm_config = _wider_model_config(head_dim) with torch.device("meta"): module = FakeAttention( hidden=32, head_dim=head_dim, heads=heads, kv_heads=kv_heads ) # kv width is 24, not a multiple of the model-wide 16, but is of this # layer's 8. fuser = get_fuser(module, QKVFuser) assert fuser is not None and fuser.validate(module, vllm_config)