| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 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 | |
| 9 | import pytest |
| 10 | import torch |
| 11 | import torch.nn as nn |
| 12 | import torch.nn.functional as F |
| 13 | |
| 14 | from vllm.model_executor.models.transformers.fuser import Fusers, get_fuser, get_fusers |
| 15 | from vllm.model_executor.models.transformers.fusers import ( |
| 16 | GLUFuser, |
| 17 | MergedColumnParallelFuser, |
| 18 | PackedQKVFuser, |
| 19 | QKVFuser, |
| 20 | packed_qkv, |
| 21 | qkv, |
| 22 | ) |
| 23 | from vllm.model_executor.models.transformers.fx_utils import ( |
| 24 | bypass_existence_guard, |
| 25 | trace, |
| 26 | ) |
| 27 | |
| 28 | |
| 29 | class SiluAndMulStub(nn.Module): |
| 30 | """Stand-in for vLLM's `SiluAndMul` (no vLLM config required).""" |
| 31 | |
| 32 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 33 | d = x.shape[-1] // 2 |
| 34 | return F.silu(x[..., :d]) * x[..., d:] |
| 35 | |
| 36 | |
| 37 | class NoDownGLU(nn.Module): |
| 38 | """`act(gate(x)) * up(x)` with no output projection -> `down_name` is None.""" |
| 39 | |
| 40 | def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False): |
| 41 | super().__init__() |
| 42 | self.gate_proj = nn.Linear(hidden, inter, bias=bias) |
| 43 | self.up_proj = nn.Linear(hidden, inter, bias=bias) |
| 44 | self.act_fn = nn.SiLU() |
| 45 | |
| 46 | def forward(self, x): |
| 47 | return self.act_fn(self.gate_proj(x)) * self.up_proj(x) |
| 48 | |
| 49 | |
| 50 | class GLUMLP(NoDownGLU): |
| 51 | """`down(act(gate(x)) * up(x))` — the canonical HF GLU MLP.""" |
| 52 | |
| 53 | def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False): |
| 54 | super().__init__(hidden, inter, bias) |
| 55 | self.down_proj = nn.Linear(inter, hidden, bias=bias) |
| 56 | |
| 57 | def forward(self, x): |
| 58 | return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) |
| 59 | |
| 60 | |
| 61 | class ReversedGLUMLP(GLUMLP): |
| 62 | """`up(x) * act(gate(x))` — operands swapped (multiply is commutative).""" |
| 63 | |
| 64 | def forward(self, x): |
| 65 | return self.down_proj(self.up_proj(x) * self.act_fn(self.gate_proj(x))) |
| 66 | |
| 67 | |
| 68 | class NotAnMLP(nn.Module): |
| 69 | """Two linears but no activation*linear multiply -> must not match.""" |
| 70 | |
| 71 | def __init__(self): |
| 72 | super().__init__() |
| 73 | self.fc1 = nn.Linear(8, 8) |
| 74 | self.fc2 = nn.Linear(8, 8) |
| 75 | |
| 76 | def forward(self, x): |
| 77 | return self.fc2(self.fc1(x)) |
| 78 | |
| 79 | |
| 80 | class NotAnActGLUMLP(GLUMLP): |
| 81 | """GLU-shaped, but the "activation" is not a known activation module.""" |
| 82 | |
| 83 | def __init__(self): |
| 84 | super().__init__() |
| 85 | self.act_fn = nn.Dropout() |
| 86 | |
| 87 | |
| 88 | class BroadcastGLU(NoDownGLU): |
| 89 | """A broadcast multiply cannot use an equal-halves AndMul kernel.""" |
| 90 | |
| 91 | def __init__(self): |
| 92 | super().__init__() |
| 93 | self.gate_proj = nn.Linear(16, 1, bias=False) |
| 94 | |
| 95 | |
| 96 | class ExtraProjGLU(GLUMLP): |
| 97 | """A third linear on the same input (e.g. a router) -> gate/up still fuse.""" |
| 98 | |
| 99 | def __init__(self): |
| 100 | super().__init__() |
| 101 | self.router = nn.Linear(16, 4, bias=False) |
| 102 | |
| 103 | def forward(self, x): |
| 104 | gated = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) |
| 105 | return gated, self.router(x) |
| 106 | |
| 107 | |
| 108 | class UntraceableMLP(GLUMLP): |
| 109 | """Data-dependent control flow *before* the GLU -> no match.""" |
| 110 | |
| 111 | def forward(self, x): |
| 112 | if x.sum() > 0: # noqa: SIM108 - intentionally untraceable |
| 113 | return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) |
| 114 | return x |
| 115 | |
| 116 | |
| 117 | class UntraceableTailGLUMLP(GLUMLP): |
| 118 | """Data-dependent control flow *after* the GLU -> still fusable.""" |
| 119 | |
| 120 | def forward(self, x): |
| 121 | y = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) |
| 122 | if y.sum() > torch.inf: # intentionally untraceable |
| 123 | y = y * 0 |
| 124 | return y |
| 125 | |
| 126 | |
| 127 | class FakeAttention(nn.Module): |
| 128 | """HF v5-style attention: shape unpacking, dead KV branch, kwargs interface.""" |
| 129 | |
| 130 | is_causal = True |
| 131 | |
| 132 | def __init__( |
| 133 | self, |
| 134 | hidden: int = 32, |
| 135 | head_dim: int = 8, |
| 136 | heads: int = 4, |
| 137 | kv_heads: int = 4, |
| 138 | bias: bool = False, |
| 139 | layer_idx: int = 0, |
| 140 | ): |
| 141 | super().__init__() |
| 142 | self.config = SimpleNamespace(_attn_implementation="vllm") |
| 143 | self.layer_idx = layer_idx |
| 144 | self.head_dim = head_dim |
| 145 | self.scaling = head_dim**-0.5 |
| 146 | self.q_proj = nn.Linear(hidden, heads * head_dim, bias=bias) |
| 147 | self.k_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias) |
| 148 | self.v_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias) |
| 149 | self.o_proj = nn.Linear(heads * head_dim, hidden, bias=bias) |
| 150 | |
| 151 | def forward( |
| 152 | self, hidden_states, attention_mask=None, past_key_values=None, **kwargs |
| 153 | ): |
| 154 | from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS |
| 155 | |
| 156 | input_shape = hidden_states.shape[:-1] |
| 157 | hidden_shape = (*input_shape, -1, self.head_dim) |
| 158 | q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) |
| 159 | k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) |
| 160 | v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) |
| 161 | if past_key_values is not None: |
| 162 | k, v = past_key_values.update(k, v, self.layer_idx) |
| 163 | attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface( |
| 164 | self.config._attn_implementation, None |
| 165 | ) |
| 166 | attn_output, attn_weights = attention_interface( |
| 167 | self, q, k, v, attention_mask, scaling=self.scaling, **kwargs |
| 168 | ) |
| 169 | attn_output = attn_output.reshape(*input_shape, -1).contiguous() |
| 170 | return self.o_proj(attn_output), attn_weights |
| 171 | |
| 172 | |
| 173 | class ReversedFakeAttention(FakeAttention): |
| 174 | """Projections computed in (v, k, q) order — q must still be identified.""" |
| 175 | |
| 176 | def forward( |
| 177 | self, hidden_states, attention_mask=None, past_key_values=None, **kwargs |
| 178 | ): |
| 179 | from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS |
| 180 | |
| 181 | input_shape = hidden_states.shape[:-1] |
| 182 | hidden_shape = (*input_shape, -1, self.head_dim) |
| 183 | v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) |
| 184 | k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) |
| 185 | q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) |
| 186 | attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface( |
| 187 | self.config._attn_implementation, None |
| 188 | ) |
| 189 | attn_output, _ = attention_interface( |
| 190 | self, q, k, v, attention_mask, scaling=self.scaling, **kwargs |
| 191 | ) |
| 192 | return self.o_proj(attn_output.reshape(*input_shape, -1)), None |
| 193 | |
| 194 | |
| 195 | class FourParallelLinears(nn.Module): |
| 196 | """Four same-input projections of different widths -> one merged linear.""" |
| 197 | |
| 198 | def __init__(self): |
| 199 | super().__init__() |
| 200 | self.proj_a = nn.Linear(8, 4) |
| 201 | self.proj_b = nn.Linear(8, 6) |
| 202 | self.proj_c = nn.Linear(8, 8) |
| 203 | self.proj_d = nn.Linear(8, 10) |
| 204 | |
| 205 | def forward(self, x): |
| 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) |
| 220 | x.add_(1) |
| 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.""" |
| 255 | |
| 256 | def __init__(self, hidden: int = 32, heads: int = 4, head_dim: int = 8): |
| 257 | super().__init__() |
| 258 | self.q_proj = nn.Linear(hidden, heads * head_dim, bias=False) |
| 259 | self.k_proj = nn.Linear(hidden, heads * head_dim, bias=False) |
| 260 | self.v_proj = nn.Linear(hidden, heads * head_dim, bias=False) |
| 261 | self.head_dim = head_dim |
| 262 | |
| 263 | def forward(self, x): |
| 264 | return self.q_proj(x), self.k_proj(x), self.v_proj(x) |
| 265 | |
| 266 | |
| 267 | class ExtraProjAttention(FakeAttention): |
| 268 | """A second non-qkv linear of a different width -> `o_proj` still found.""" |
| 269 | |
| 270 | def __init__(self, **kwargs): |
| 271 | super().__init__(**kwargs) |
| 272 | self.sink_proj = nn.Linear(self.head_dim, self.head_dim, bias=False) |
| 273 | |
| 274 | |
| 275 | class QKNormAttention(FakeAttention): |
| 276 | """OLMoE-style: a full-dim norm applied to the whole q/k projection output.""" |
| 277 | |
| 278 | def __init__(self, **kwargs): |
| 279 | super().__init__(**kwargs) |
| 280 | self.q_norm = nn.RMSNorm(self.q_proj.out_features) |
| 281 | self.k_norm = nn.RMSNorm(self.k_proj.out_features) |
| 282 | |
| 283 | def forward( |
| 284 | self, hidden_states, attention_mask=None, past_key_values=None, **kwargs |
| 285 | ): |
| 286 | q = self.q_norm(self.q_proj(hidden_states)) |
| 287 | k = self.k_norm(self.k_proj(hidden_states)) |
| 288 | v = self.v_proj(hidden_states) |
| 289 | return self.o_proj(q + k + v), None |
| 290 | |
| 291 | |
| 292 | class PerHeadQKNormAttention(FakeAttention): |
| 293 | """Qwen3-style: a per-head norm (`head_dim`) applied after the head reshape.""" |
| 294 | |
| 295 | def __init__(self, **kwargs): |
| 296 | super().__init__(**kwargs) |
| 297 | self.q_norm = nn.RMSNorm(self.head_dim) |
| 298 | self.k_norm = nn.RMSNorm(self.head_dim) |
| 299 | |
| 300 | def forward( |
| 301 | self, hidden_states, attention_mask=None, past_key_values=None, **kwargs |
| 302 | ): |
| 303 | shape = (*hidden_states.shape[:-1], -1, self.head_dim) |
| 304 | q = self.q_norm(self.q_proj(hidden_states).view(shape)) |
| 305 | k = self.k_norm(self.k_proj(hidden_states).view(shape)) |
| 306 | v = self.v_proj(hidden_states).view(shape) |
| 307 | return self.o_proj((q + k + v).flatten(-2)), None |
| 308 | |
| 309 | |
| 310 | class ResidDropoutAttention(FakeAttention): |
| 311 | """GPT-style dropout after `o_proj` -> the output projection is still found.""" |
| 312 | |
| 313 | def __init__(self, **kwargs): |
| 314 | super().__init__(**kwargs) |
| 315 | self.resid_dropout = nn.Dropout(0.0) |
| 316 | |
| 317 | def forward( |
| 318 | self, hidden_states, attention_mask=None, past_key_values=None, **kwargs |
| 319 | ): |
| 320 | from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS |
| 321 | |
| 322 | input_shape = hidden_states.shape[:-1] |
| 323 | hidden_shape = (*input_shape, -1, self.head_dim) |
| 324 | q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) |
| 325 | k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) |
| 326 | v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) |
| 327 | attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface( |
| 328 | self.config._attn_implementation, None |
| 329 | ) |
| 330 | attn_output, _ = attention_interface( |
| 331 | self, q, k, v, attention_mask, scaling=self.scaling, **kwargs |
| 332 | ) |
| 333 | attn_output = attn_output.reshape(*input_shape, -1).contiguous() |
| 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 | |
| 640 | is_causal = True |
| 641 | |
| 642 | def __init__( |
| 643 | self, |
| 644 | hidden: int = 32, |
| 645 | head_dim: int = 8, |
| 646 | heads: int = 4, |
| 647 | kv_heads: int = 1, |
| 648 | bias: bool = False, |
| 649 | layer_idx: int = 0, |
| 650 | ): |
| 651 | super().__init__() |
| 652 | self.config = SimpleNamespace(_attn_implementation="vllm") |
| 653 | self.layer_idx = layer_idx |
| 654 | self.head_dim = head_dim |
| 655 | self.scaling = head_dim**-0.5 |
| 656 | self.embed_dim = heads * head_dim |
| 657 | self.kv_dim = kv_heads * head_dim |
| 658 | self.c_attn = nn.Linear(hidden, self.embed_dim + 2 * self.kv_dim, bias=bias) |
| 659 | self.c_proj = nn.Linear(self.embed_dim, hidden, bias=bias) |
| 660 | self.resid_dropout = nn.Dropout(0.0) |
| 661 | |
| 662 | def forward( |
| 663 | self, hidden_states, attention_mask=None, past_key_values=None, **kwargs |
| 664 | ): |
| 665 | from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS |
| 666 | |
| 667 | input_shape = hidden_states.shape[:-1] |
| 668 | q, k, v = ( |
| 669 | self.c_attn(hidden_states) |
| 670 | .unsqueeze(1) |
| 671 | .split((self.embed_dim, self.kv_dim, self.kv_dim), dim=3) |
| 672 | ) |
| 673 | q = q.view(*input_shape, -1, self.head_dim).transpose(1, 2) |
| 674 | if past_key_values is not None: |
| 675 | k, v = past_key_values.update(k, v, self.layer_idx) |
| 676 | attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface( |
| 677 | self.config._attn_implementation, None |
| 678 | ) |
| 679 | attn_output, attn_weights = attention_interface( |
| 680 | self, q, k, v, attention_mask, scaling=self.scaling, **kwargs |
| 681 | ) |
| 682 | attn_output = attn_output.reshape(*input_shape, -1).contiguous() |
| 683 | return self.resid_dropout(self.c_proj(attn_output)), attn_weights |
| 684 | |
| 685 | |
| 686 | class PerHeadSplitAttention(nn.Module): |
| 687 | """A packed projection reshaped and split *per head* -> not a q/k/v split.""" |
| 688 | |
| 689 | def __init__(self, hidden: int = 32, head_dim: int = 8, heads: int = 4): |
| 690 | super().__init__() |
| 691 | self.head_dim = head_dim |
| 692 | self.heads = heads |
| 693 | self.c_attn = nn.Linear(hidden, 3 * heads * head_dim) |
| 694 | self.c_proj = nn.Linear(heads * head_dim, hidden) |
| 695 | |
| 696 | def forward(self, hidden_states): |
| 697 | shape = (*hidden_states.shape[:2], self.heads, 3 * self.head_dim) |
| 698 | q, k, v = ( |
| 699 | self.c_attn(hidden_states) |
| 700 | .view(shape) |
| 701 | .transpose(1, 2) |
| 702 | .split((self.head_dim, self.head_dim, self.head_dim), dim=3) |
| 703 | ) |
| 704 | return self.c_proj((q + k + v).transpose(1, 2).flatten(-2)) |
| 705 | |
| 706 | |
| 707 | class FakeSelfAttn(nn.Module): |
| 708 | """Stand-in for the vLLM `Attention` attached to the dispatching module.""" |
| 709 | |
| 710 | def __init__(self): |
| 711 | super().__init__() |
| 712 | self.impl = SimpleNamespace(scale=None) |
| 713 | |
| 714 | def forward(self, q, k, v): |
| 715 | # MHA-shaped stub: any deterministic combination of q/k/v will do |
| 716 | return q + 2 * k + 3 * v |
| 717 | |
| 718 | |
| 719 | class FakeMQASelfAttn(FakeSelfAttn): |
| 720 | """Stand-in for grouped/multi-query layouts, where `k`/`v` are narrower.""" |
| 721 | |
| 722 | def forward(self, q, k, v): |
| 723 | groups = q.shape[-1] // k.shape[-1] |
| 724 | return q + (2 * k + 3 * v).repeat(1, groups) |
| 725 | |
| 726 | |
| 727 | @pytest.fixture(autouse=True) |
| 728 | def _clear_fuser_cache(): |
| 729 | get_fusers.cache_clear() |
| 730 | yield |
| 731 | get_fusers.cache_clear() |
| 732 | |
| 733 | |
| 734 | def _apply_glu_fuser_with_stubs(module: nn.Module, fuser: GLUFuser): |
| 735 | """Apply a fuser using plain stand-ins (merged `nn.Linear` + silu AndMul).""" |
| 736 | gate = module.get_submodule(fuser.gate_name) |
| 737 | up = module.get_submodule(fuser.up_name) |
| 738 | merged = nn.Linear( |
| 739 | gate.in_features, |
| 740 | gate.out_features + up.out_features, |
| 741 | bias=gate.bias is not None, |
| 742 | ) |
| 743 | with torch.no_grad(): |
| 744 | merged.weight.copy_(torch.cat([gate.weight, up.weight], dim=0)) |
| 745 | if gate.bias is not None: |
| 746 | merged.bias.copy_(torch.cat([gate.bias, up.bias], dim=0)) |
| 747 | setattr(module, fuser.merged_name, merged) |
| 748 | setattr(module, fuser.act_name, SiluAndMulStub()) |
| 749 | delattr(module, fuser.gate_name) |
| 750 | delattr(module, fuser.up_name) |
| 751 | module.forward = MethodType(fuser.fused_forward, module) |
| 752 | return module |
| 753 | |
| 754 | |
| 755 | def _apply_qkv_fuser_with_stubs(module: nn.Module, fuser: QKVFuser): |
| 756 | """Apply a fuser using a plain merged `nn.Linear` (no TP sharding).""" |
| 757 | q, k, v = ( |
| 758 | module.get_submodule(name) |
| 759 | for name in (fuser.q_name, fuser.k_name, fuser.v_name) |
| 760 | ) |
| 761 | merged = nn.Linear( |
| 762 | q.in_features, |
| 763 | q.out_features + k.out_features + v.out_features, |
| 764 | bias=q.bias is not None, |
| 765 | ) |
| 766 | with torch.no_grad(): |
| 767 | merged.weight.copy_(torch.cat([q.weight, k.weight, v.weight], dim=0)) |
| 768 | if q.bias is not None: |
| 769 | merged.bias.copy_(torch.cat([q.bias, k.bias, v.bias], dim=0)) |
| 770 | merged.output_sizes = [q.out_features, k.out_features, v.out_features] |
| 771 | merged.tp_size = 1 |
| 772 | setattr(module, fuser.merged_name, merged) |
| 773 | for name in (fuser.q_name, fuser.k_name, fuser.v_name): |
| 774 | delattr(module, name) |
| 775 | module.forward = MethodType(fuser.fused_forward, module) |
| 776 | return module |
| 777 | |
| 778 | |
| 779 | def _apply_packed_qkv_fuser_with_stubs(module: nn.Module, fuser: PackedQKVFuser): |
| 780 | """Apply a fuser at `tp_size == 1`, where the rewritten split is unchanged.""" |
| 781 | qkv = module.get_submodule(fuser.qkv_name) |
| 782 | qkv.output_sizes = [fuser.q_size, fuser.kv_size, fuser.kv_size] |
| 783 | qkv.tp_size = 1 |
| 784 | module.forward = MethodType(fuser.fused_forward, module) |
| 785 | return module |
| 786 | |
| 787 | |
| 788 | @pytest.mark.parametrize("mlp_cls", [GLUMLP, ReversedGLUMLP]) |
| 789 | @pytest.mark.parametrize("bias", [False, True]) |
| 790 | def test_detects_and_rewrites_glu(mlp_cls, bias): |
| 791 | with torch.device("meta"): |
| 792 | meta = mlp_cls(bias=bias) |
| 793 | fuser = get_fuser(meta, GLUFuser) |
| 794 | assert isinstance(fuser, GLUFuser) |
| 795 | assert ( |
| 796 | fuser.gate_name, |
| 797 | fuser.up_name, |
| 798 | fuser.act_name, |
| 799 | fuser.down_name, |
| 800 | ) == ("gate_proj", "up_proj", "act_fn", "down_proj") |
| 801 | |
| 802 | # The rewritten forward references the merged projection instead of the |
| 803 | # sources; the rest of the forward is untouched. |
| 804 | names = fuser.fused_forward.__code__.co_names |
| 805 | assert "gate_up_proj" in names and "act_fn" in names and "down_proj" in names |
| 806 | assert not {"gate_proj", "up_proj"} & set(names) |
| 807 | |
| 808 | # Numerics: the fused forward must match the original on a real instance. |
| 809 | real = mlp_cls(bias=bias) |
| 810 | for p in real.parameters(): |
| 811 | nn.init.normal_(p, std=0.05) |
| 812 | x = torch.randn(4, 16) |
| 813 | expected = real(x) |
| 814 | fused = _apply_glu_fuser_with_stubs(real, fuser) |
| 815 | |
| 816 | # Fusion is in place: the module keeps its class and other attributes |
| 817 | assert fused is real and type(fused) is mlp_cls |
| 818 | torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5) |
| 819 | |
| 820 | |
| 821 | def test_glu_fuses_alongside_unrelated_projections(): |
| 822 | """Sibling linears the GLU does not consume must not block the fusion.""" |
| 823 | with torch.device("meta"): |
| 824 | fuser = get_fuser(ExtraProjGLU(), GLUFuser) |
| 825 | assert fuser is not None |
| 826 | assert (fuser.gate_name, fuser.up_name) == ("gate_proj", "up_proj") |
| 827 | real = ExtraProjGLU() |
| 828 | x = torch.randn(4, 16) |
| 829 | expected = real(x) |
| 830 | fused = _apply_glu_fuser_with_stubs(real, fuser) |
| 831 | torch.testing.assert_close(fused(x), expected) |
| 832 | |
| 833 | |
| 834 | def test_glu_identifies_down_projection(): |
| 835 | """The row projection consuming `act(gate(x)) * up(x)` is identified. |
| 836 | |
| 837 | It is forced to `RowParallelLinear` in `update_attrs` so its sharded input |
| 838 | matches the column-parallel merged gate/up; `None` when there is no such |
| 839 | projection to force (fusion of gate/up still applies).""" |
| 840 | with torch.device("meta"): |
| 841 | assert get_fuser(GLUMLP(), GLUFuser).down_name == "down_proj" |
| 842 | assert get_fuser(ReversedGLUMLP(), GLUFuser).down_name == "down_proj" |
| 843 | assert get_fuser(NoDownGLU(), GLUFuser).down_name is None |
| 844 | |
| 845 | |
| 846 | @pytest.mark.parametrize("attn_cls", [FakeAttention, ReversedFakeAttention]) |
| 847 | @pytest.mark.parametrize("kv_heads", [4, 2]) |
| 848 | def test_detects_and_rewrites_qkv(attn_cls, kv_heads): |
| 849 | if attn_cls is ReversedFakeAttention and kv_heads == 4: |
| 850 | pytest.skip("MHA q/k/v assignment is order-based by design") |
| 851 | with torch.device("meta"): |
| 852 | meta = attn_cls(kv_heads=kv_heads) |
| 853 | fuser = get_fuser(meta, QKVFuser) |
| 854 | assert isinstance(fuser, QKVFuser) |
| 855 | # q (sharded differently under TP) must be identified exactly; k/v may be |
| 856 | # swapped for non-canonical compute order, which is numerically consistent |
| 857 | # because the weight mapping and the split indices follow the same |
| 858 | # assignment. |
| 859 | assert fuser.q_name == "q_proj" |
| 860 | assert {fuser.k_name, fuser.v_name} == {"k_proj", "v_proj"} |
| 861 | assert fuser.o_name == "o_proj" |
| 862 | |
| 863 | # The projections are merged; everything else stays live Python with its |
| 864 | # original semantics (branches, kwargs, attribute reads) |
| 865 | code = fuser.fused_forward.__code__ |
| 866 | names = code.co_names |
| 867 | assert "qkv_proj" in names and "output_sizes" in names and "o_proj" in names |
| 868 | assert "tp_size" in names |
| 869 | assert not {"q_proj", "k_proj", "v_proj"} & set(names) |
| 870 | if attn_cls is FakeAttention: |
| 871 | assert "update" in names # the cache branch survives |
| 872 | assert code.co_flags & inspect.CO_VARKEYWORDS # **kwargs survives |
| 873 | |
| 874 | # Numerics: the fused forward must match the original on a real instance, |
| 875 | # with a different layer_idx than the traced instance (kv_heads == heads so |
| 876 | # the q/k/v stub combination is shape-compatible). |
| 877 | real = attn_cls(kv_heads=4, layer_idx=3) |
| 878 | for p in real.parameters(): |
| 879 | nn.init.normal_(p, std=0.05) |
| 880 | x = torch.randn(1, 5, 32) |
| 881 | real.attn = FakeSelfAttn() |
| 882 | expected, _ = real(x) |
| 883 | fused = _apply_qkv_fuser_with_stubs(real, fuser) |
| 884 | |
| 885 | # Fusion is in place: the module keeps its class and other attributes |
| 886 | assert fused is real and type(fused) is attn_cls |
| 887 | assert fused.layer_idx == 3 and fused.is_causal and fused.config is not None |
| 888 | out, _ = fused(x) |
| 889 | torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5) |
| 890 | |
| 891 | |
| 892 | def test_qkv_identifies_output_projection(): |
| 893 | with torch.device("meta"): |
| 894 | assert get_fuser(FakeAttention(), QKVFuser).o_name == "o_proj" |
| 895 | assert get_fuser(ReversedFakeAttention(), QKVFuser).o_name == "o_proj" |
| 896 | assert get_fuser(ExtraProjAttention(), QKVFuser).o_name == "o_proj" |
| 897 | # Norm children (q_norm/k_norm) must not disturb o_proj identification. |
| 898 | assert get_fuser(QKNormAttention(), QKVFuser).o_name == "o_proj" |
| 899 | assert get_fuser(PerHeadQKNormAttention(), QKVFuser).o_name == "o_proj" |
| 900 | # A module between o_proj and the return is transparent. |
| 901 | assert get_fuser(ResidDropoutAttention(), QKVFuser).o_name == "o_proj" |
| 902 | |
| 903 | |
| 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"): |
| 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 |
| 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( |
| 1143 | cls, default_vllm_config, monkeypatch |
| 1144 | ): |
| 1145 | from vllm.model_executor import parameter |
| 1146 | from vllm.model_executor.layers import linear |
| 1147 | from vllm.model_executor.layers.linear import MergedColumnParallelLinear |
| 1148 | from vllm.model_executor.layers.utils import default_unquantized_gemm |
| 1149 | from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper |
| 1150 | |
| 1151 | monkeypatch.setattr(linear, "get_tensor_model_parallel_rank", lambda: 0) |
| 1152 | monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: 1) |
| 1153 | monkeypatch.setattr( |
| 1154 | linear, |
| 1155 | "dispatch_unquantized_gemm", |
| 1156 | lambda linear_backend="auto": default_unquantized_gemm, |
| 1157 | ) |
| 1158 | monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0) |
| 1159 | monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1) |
| 1160 | with torch.device("meta"): |
| 1161 | meta = cls() |
| 1162 | fuser = MergedColumnParallelFuser.match(trace(meta), meta) |
| 1163 | assert fuser is not None |
| 1164 | assert set(fuser.linear_names) == {"proj_a", "proj_b", "proj_c", "proj_d"} |
| 1165 | assert fuser.shards == [(name, i) for i, name in enumerate(fuser.linear_names)] |
| 1166 | fuser.update_forward(meta) |
| 1167 | |
| 1168 | real = cls() |
| 1169 | x = torch.randn(2, 8) |
| 1170 | expected = real(x) |
| 1171 | weights = [ |
| 1172 | (name, weight.detach().clone()) for name, weight in real.named_parameters() |
| 1173 | ] |
| 1174 | fused = fuser.fuse(real, "parallel", default_vllm_config) |
| 1175 | root = nn.Module() |
| 1176 | root.parallel = fused |
| 1177 | mapper = WeightsMapper(orig_to_new_stacked=fuser.orig_to_new_stacked("parallel")) |
| 1178 | AutoWeightsLoader(root).load_weights( |
| 1179 | ((f"parallel.{name}", weight) for name, weight in weights), mapper=mapper |
| 1180 | ) |
| 1181 | |
| 1182 | merged = getattr(fused, fuser.merged_name) |
| 1183 | assert isinstance(merged, MergedColumnParallelLinear) |
| 1184 | for actual, reference in zip(fused(x), expected): |
| 1185 | torch.testing.assert_close(actual, reference) |
| 1186 | |
| 1187 | |
| 1188 | def test_merged_name_is_unique_per_linear_names(): |
| 1189 | """The computed name must not collide across different projections, so |
| 1190 | stacking two unrelated fusions can never clobber packed_modules_mapping.""" |
| 1191 | ab_c = MergedColumnParallelFuser(source_cls="M", linear_names=("a_b", "c")) |
| 1192 | a_bc = MergedColumnParallelFuser(source_cls="M", linear_names=("a", "b_c")) |
| 1193 | assert ab_c.merged_name != a_bc.merged_name |
| 1194 | # Same names -> same (and therefore safely re-mergeable) name. |
| 1195 | dup = MergedColumnParallelFuser(source_cls="M", linear_names=("a_b", "c")) |
| 1196 | assert ab_c.merged_name == dup.merged_name |
| 1197 | |
| 1198 | |
| 1199 | def test_heterogeneous_instances_fall_back_to_merged_column_fuser( |
| 1200 | default_vllm_config, |
| 1201 | ): |
| 1202 | """Same class and structure -> both fusers are cached as candidates. A |
| 1203 | layer whose out_features aren't a multiple of its own head_dim (as a |
| 1204 | misconfigured or heterogeneous checkpoint might have) fails QKVFuser's |
| 1205 | validation and must still fuse via the generic fallback, while a |
| 1206 | well-formed sibling of the same class keeps using QKVFuser.""" |
| 1207 | with torch.device("meta"): |
| 1208 | container = nn.Module() |
| 1209 | container.compatible = PlainQKV(head_dim=8) |
| 1210 | container.incompatible = PlainQKV(head_dim=8) |
| 1211 | container.incompatible.head_dim = 5 # 32 % 5 != 0, unlike q_proj's 32 % 8 |
| 1212 | |
| 1213 | candidates = get_fusers(container.compatible) |
| 1214 | assert any(isinstance(f, QKVFuser) for f in candidates) |
| 1215 | assert any(isinstance(f, MergedColumnParallelFuser) for f in candidates) |
| 1216 | |
| 1217 | fusers = Fusers(container, default_vllm_config) |
| 1218 | assert isinstance(fusers[container.compatible][0], QKVFuser) |
| 1219 | assert isinstance(fusers[container.incompatible][0], MergedColumnParallelFuser) |
| 1220 | |
| 1221 | |
| 1222 | @pytest.mark.parametrize("kv_heads", [1, 2]) |
| 1223 | def test_detects_and_rewrites_packed_qkv(kv_heads): |
| 1224 | """A single projection split into q/k/v must be re-sharded, not merged. |
| 1225 | |
| 1226 | Only the split sizes change: `QKVParallelLinear` loads the packed |
| 1227 | checkpoint weight as-is, and shards q by heads while replicating k/v.""" |
| 1228 | with torch.device("meta"): |
| 1229 | meta = PackedQKVAttention(kv_heads=kv_heads) |
| 1230 | fuser = get_fuser(meta, PackedQKVFuser) |
| 1231 | assert isinstance(fuser, PackedQKVFuser) |
| 1232 | assert (fuser.qkv_name, fuser.o_name) == ("c_attn", "c_proj") |
| 1233 | assert (fuser.q_size, fuser.kv_size) == (32, 8 * kv_heads) |
| 1234 | |
| 1235 | # The hard-coded widths become the per-rank widths of the sharded linear |
| 1236 | names = fuser.fused_forward.__code__.co_names |
| 1237 | assert "output_sizes" in names and "tp_size" in names |
| 1238 | assert "kv_dim" not in names and "embed_dim" not in names |
| 1239 | |
| 1240 | # Numerics: the rewritten forward must match the original on a real instance |
| 1241 | real = PackedQKVAttention(kv_heads=kv_heads, layer_idx=3) |
| 1242 | for p in real.parameters(): |
| 1243 | nn.init.normal_(p, std=0.05) |
| 1244 | x = torch.randn(1, 5, 32) |
| 1245 | real.attn = FakeMQASelfAttn() |
| 1246 | expected, _ = real(x) |
| 1247 | fused = _apply_packed_qkv_fuser_with_stubs(real, fuser) |
| 1248 | |
| 1249 | # Fusion is in place: the module keeps its class and other attributes |
| 1250 | assert fused is real and type(fused) is PackedQKVAttention |
| 1251 | assert fused.layer_idx == 3 and fused.is_causal |
| 1252 | out, _ = fused(x) |
| 1253 | torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5) |
| 1254 | |
| 1255 | |
| 1256 | def test_per_head_split_is_not_packed_qkv(): |
| 1257 | """The split must consume the whole projection, else its sizes are head |
| 1258 | widths and re-sharding by them would be wrong.""" |
| 1259 | with torch.device("meta"): |
| 1260 | assert get_fuser(PerHeadSplitAttention(), PackedQKVFuser) is None |
| 1261 | |
| 1262 | |
| 1263 | def test_fuser_is_cached_per_class_and_structure(): |
| 1264 | with torch.device("meta"): |
| 1265 | fuser_a = get_fuser(GLUMLP(), GLUFuser) |
| 1266 | fuser_b = get_fuser(GLUMLP(), GLUFuser) |
| 1267 | assert fuser_a is fuser_b |
| 1268 | assert any(key[0] is GLUMLP for key in get_fusers.cache) |
| 1269 | |
| 1270 | |
| 1271 | @pytest.mark.parametrize("cls", [NotAnMLP, UntraceableMLP]) |
| 1272 | def test_non_matching_modules_return_none(cls): |
| 1273 | with torch.device("meta"): |
| 1274 | module = cls() |
| 1275 | assert get_fuser(module, GLUFuser) is None |
| 1276 | |
| 1277 | |
| 1278 | def test_untraceable_tail_still_fuses(): |
| 1279 | with torch.device("meta"): |
| 1280 | meta = UntraceableTailGLUMLP() |
| 1281 | fuser = get_fuser(meta, GLUFuser) |
| 1282 | assert isinstance(fuser, GLUFuser) |
| 1283 | |
| 1284 | # Numerics: the live tail must survive the rewrite |
| 1285 | real = UntraceableTailGLUMLP() |
| 1286 | for p in real.parameters(): |
| 1287 | nn.init.normal_(p, std=0.05) |
| 1288 | x = torch.randn(4, 16) |
| 1289 | expected = real(x) |
| 1290 | fused = _apply_glu_fuser_with_stubs(real, fuser) |
| 1291 | torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5) |
| 1292 | |
| 1293 | |
| 1294 | def test_weight_mappings_are_scoped_to_fused_prefixes(): |
| 1295 | from vllm.model_executor.models.utils import WeightsMapper |
| 1296 | |
| 1297 | with torch.device("meta"): |
| 1298 | glu_fuser = get_fuser(GLUMLP(), GLUFuser) |
| 1299 | qkv_fuser = get_fuser(FakeAttention(), QKVFuser) |
| 1300 | |
| 1301 | mapper = WeightsMapper() |
| 1302 | for prefix in ("model.layers.0.mlp", "model.layers.1.mlp"): |
| 1303 | mapper.orig_to_new_stacked.update(glu_fuser.orig_to_new_stacked(prefix)) |
| 1304 | mapper.orig_to_new_stacked.update( |
| 1305 | qkv_fuser.orig_to_new_stacked("model.layers.0.self_attn") |
| 1306 | ) |
| 1307 | |
| 1308 | names = [ |
| 1309 | "model.layers.0.mlp.gate_proj.weight", |
| 1310 | "model.layers.0.mlp.up_proj.weight", |
| 1311 | "model.layers.1.mlp.gate_proj.weight", |
| 1312 | "model.layers.0.self_attn.q_proj.weight", |
| 1313 | "model.layers.0.self_attn.k_proj.weight", |
| 1314 | "model.layers.0.self_attn.v_proj.weight", |
| 1315 | # Unfused modules at other prefixes must be left untouched. |
| 1316 | "model.layers.2.mlp.experts.0.gate_proj.weight", |
| 1317 | "model.layers.1.self_attn.q_proj.weight", |
| 1318 | ] |
| 1319 | # `apply` rewrites the name and stamps the shard id onto each tensor. |
| 1320 | weights = [(name, torch.empty(0)) for name in names] |
| 1321 | mapped = list(mapper.apply(weights)) |
| 1322 | mapped_names = [name for name, _ in mapped] |
| 1323 | shard_ids = [getattr(data, "shard_id", None) for _, data in mapped] |
| 1324 | |
| 1325 | assert mapped_names == [ |
| 1326 | "model.layers.0.mlp.gate_up_proj.weight", |
| 1327 | "model.layers.0.mlp.gate_up_proj.weight", |
| 1328 | "model.layers.1.mlp.gate_up_proj.weight", |
| 1329 | "model.layers.0.self_attn.qkv_proj.weight", |
| 1330 | "model.layers.0.self_attn.qkv_proj.weight", |
| 1331 | "model.layers.0.self_attn.qkv_proj.weight", |
| 1332 | # Only the exact fused layers are remapped; everything else is untouched. |
| 1333 | "model.layers.2.mlp.experts.0.gate_proj.weight", |
| 1334 | "model.layers.1.self_attn.q_proj.weight", |
| 1335 | ] |
| 1336 | assert shard_ids == [0, 1, 0, "q", "k", "v", None, None] |
| 1337 | |
| 1338 | # The fused layers are exposed to the quantization machinery via their |
| 1339 | # original constituent projection names (what the checkpoint stores). |
| 1340 | assert glu_fuser.packed_modules_mapping == { |
| 1341 | "gate_up_proj": ["gate_proj", "up_proj"], |
| 1342 | } |
| 1343 | assert qkv_fuser.packed_modules_mapping == { |
| 1344 | "qkv_proj": ["q_proj", "k_proj", "v_proj"], |
| 1345 | } |
| 1346 | |
| 1347 | |
| 1348 | @pytest.mark.parametrize("cls", [NotAnMLP, NotAnActGLUMLP, BroadcastGLU]) |
| 1349 | def test_unfusable_modules_are_not_fused(cls, default_vllm_config): |
| 1350 | with torch.device("meta"): |
| 1351 | module = cls() |
| 1352 | fuser = get_fuser(module, GLUFuser) |
| 1353 | # Either no pattern matches the class, or this instance fails validation |
| 1354 | # (`recursive_replace` gates fusion and its weight mappings on `validate`) |
| 1355 | assert fuser is None or not fuser.validate(module, default_vllm_config) |
| 1356 | |
| 1357 | |
| 1358 | def test_act_and_mul_derived_from_module(default_vllm_config): |
| 1359 | from transformers.activations import GELUTanh, SiLUActivation |
| 1360 | |
| 1361 | from vllm.model_executor.layers.activation import GeluAndMul, SiluAndMul |
| 1362 | |
| 1363 | assert isinstance(GLUFuser._get_act_and_mul(nn.SiLU()), SiluAndMul) |
| 1364 | assert isinstance(GLUFuser._get_act_and_mul(SiLUActivation()), SiluAndMul) |
| 1365 | gelu_tanh = GLUFuser._get_act_and_mul(GELUTanh()) |
| 1366 | assert isinstance(gelu_tanh, GeluAndMul) and gelu_tanh.approximate == "tanh" |
| 1367 | gelu = GLUFuser._get_act_and_mul(nn.GELU()) |
| 1368 | assert isinstance(gelu, GeluAndMul) and gelu.approximate == "none" |
| 1369 | # Not activations at all -> no fusion |
| 1370 | assert GLUFuser._get_act_and_mul_name(nn.Dropout()) is None |
| 1371 | assert GLUFuser._get_act_and_mul_name(nn.LayerNorm(8)) is None |
| 1372 | with pytest.raises(ValueError, match="No AndMul equivalent"): |
| 1373 | GLUFuser._get_act_and_mul(nn.Dropout()) |
| 1374 | |
| 1375 | |
| 1376 | def _wider_model_config(head_dim: int) -> SimpleNamespace: |
| 1377 | """A model whose global head size is twice `head_dim`, as a wider layer |
| 1378 | elsewhere in a heterogeneous checkpoint would make it.""" |
| 1379 | return SimpleNamespace( |
| 1380 | model_config=SimpleNamespace(get_head_size=lambda: 2 * head_dim), |
| 1381 | quant_config=None, |
| 1382 | ) |
| 1383 | |
| 1384 | |
| 1385 | @pytest.mark.parametrize( |
| 1386 | "cls, fuser_module, fuser_cls", |
| 1387 | [ |
| 1388 | (FakeAttention, qkv, QKVFuser), |
| 1389 | (PackedQKVAttention, packed_qkv, PackedQKVFuser), |
| 1390 | ], |
| 1391 | ) |
| 1392 | def test_head_counts_come_from_the_module_not_the_model( |
| 1393 | cls, fuser_module, fuser_cls, monkeypatch |
| 1394 | ): |
| 1395 | """A layer narrower than the model-wide head size must not be miscounted. |
| 1396 | |
| 1397 | On a heterogeneous checkpoint (Gemma 4) the model-wide head size is the |
| 1398 | largest across layers, so deriving `total_num_heads = out_features // |
| 1399 | head_size` from it undercounts heads on a narrower layer. The widths still |
| 1400 | add up, so nothing raises below TP=4: the layer is just sharded wrong. |
| 1401 | """ |
| 1402 | head_dim, heads, kv_heads = 8, 8, 4 |
| 1403 | vllm_config = _wider_model_config(head_dim) |
| 1404 | with torch.device("meta"): |
| 1405 | module = cls(hidden=32, head_dim=head_dim, heads=heads, kv_heads=kv_heads) |
| 1406 | |
| 1407 | # Both replacements shard, so they need a TP group; only the head counts |
| 1408 | # the fuser derives are under test here. |
| 1409 | captured = {} |
| 1410 | monkeypatch.setattr( |
| 1411 | fuser_module, |
| 1412 | "QKVParallelLinear", |
| 1413 | lambda **kwargs: captured.update(kwargs) or nn.Identity(), |
| 1414 | ) |
| 1415 | monkeypatch.setattr( |
| 1416 | fuser_module, "replace_linear_class", lambda *a, **kw: nn.Identity() |
| 1417 | ) |
| 1418 | fuser = get_fuser(module, fuser_cls) |
| 1419 | assert fuser is not None and fuser.validate(module, vllm_config) |
| 1420 | fuser.update_attrs(module, "model.layers.0.self_attn", vllm_config) |
| 1421 | |
| 1422 | assert captured["head_size"] == head_dim |
| 1423 | assert captured["total_num_heads"] == heads |
| 1424 | assert captured["total_num_kv_heads"] == kv_heads |
| 1425 | |
| 1426 | |
| 1427 | def test_validate_accepts_a_layer_the_model_wide_head_size_would_reject(): |
| 1428 | """`validate` gates fusion, so a wrong head size silently disables it.""" |
| 1429 | head_dim, heads, kv_heads = 8, 8, 3 |
| 1430 | vllm_config = _wider_model_config(head_dim) |
| 1431 | with torch.device("meta"): |
| 1432 | module = FakeAttention( |
| 1433 | hidden=32, head_dim=head_dim, heads=heads, kv_heads=kv_heads |
| 1434 | ) |
| 1435 | |
| 1436 | # kv width is 24, not a multiple of the model-wide 16, but is of this |
| 1437 | # layer's 8. |
| 1438 | fuser = get_fuser(module, QKVFuser) |
| 1439 | assert fuser is not None and fuser.validate(module, vllm_config) |