From ebcef33766f8bcce6344f2bafa245d47097d3c31 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Grunwald?= Date: Tue, 28 Jul 2026 00:51:17 +0200 Subject: [PATCH] Fix MQA with tensor parallelism on transformers modeling backend (#49987) Signed-off-by: microslaw Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .../models/transformers/fusers/test_linear.py | 163 +++++++++++++++++- .../models/transformers/fuser.py | 7 +- .../models/transformers/fusers/__init__.py | 9 +- .../models/transformers/fusers/base.py | 79 +++++---- .../models/transformers/fusers/packed_qkv.py | 163 ++++++++++++++++++ .../models/transformers/fusers/qkv.py | 7 +- .../models/transformers/fx_utils.py | 12 +- 7 files changed, 396 insertions(+), 44 deletions(-) create mode 100644 vllm/model_executor/models/transformers/fusers/packed_qkv.py diff --git a/tests/models/transformers/fusers/test_linear.py b/tests/models/transformers/fusers/test_linear.py index eff78fa290d..842a1f54794 100644 --- a/tests/models/transformers/fusers/test_linear.py +++ b/tests/models/transformers/fusers/test_linear.py @@ -11,7 +11,11 @@ import torch.nn as nn import torch.nn.functional as F from vllm.model_executor.models.transformers.fuser import get_fuser -from vllm.model_executor.models.transformers.fusers import GLUFuser, QKVFuser +from vllm.model_executor.models.transformers.fusers import ( + GLUFuser, + PackedQKVFuser, + QKVFuser, +) class SiluAndMulStub(nn.Module): @@ -203,6 +207,103 @@ class PerHeadQKNormAttention(FakeAttention): 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 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` looked up in `attention_instances`.""" @@ -215,6 +316,14 @@ class FakeSelfAttn(nn.Module): 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_fuser.cache_clear() @@ -267,6 +376,15 @@ def _apply_qkv_fuser_with_stubs(module: nn.Module, fuser: QKVFuser): 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): @@ -366,6 +484,49 @@ def test_qkv_identifies_output_projection(): # Norm children (q_norm/k_norm) must not disturb o_proj identification. assert get_fuser(QKNormAttention()).o_name == "o_proj" assert get_fuser(PerHeadQKNormAttention()).o_name == "o_proj" + # A module between o_proj and the return is transparent. + assert get_fuser(ResidDropoutAttention()).o_name == "o_proj" + + +@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) + 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) + attention_instances = {3: FakeMQASelfAttn()} + expected, _ = real(x, attention_instances=attention_instances) + 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, attention_instances=attention_instances) + 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()) is None def test_fuser_is_cached_per_class_and_structure(): diff --git a/vllm/model_executor/models/transformers/fuser.py b/vllm/model_executor/models/transformers/fuser.py index 13d84a245d9..cbcb24e2a95 100644 --- a/vllm/model_executor/models/transformers/fuser.py +++ b/vllm/model_executor/models/transformers/fuser.py @@ -18,9 +18,10 @@ from vllm.logger import init_logger from vllm.model_executor.models.transformers.fusers import ( BaseFuser, GLUFuser, + PackedQKVFuser, QKVFuser, + RewriteFuser, RMSNormFuser, - StackedFuser, ) from vllm.model_executor.models.transformers.fx_utils import trace @@ -47,9 +48,9 @@ def get_fuser(module: nn.Module) -> BaseFuser | None: return None if (graph := trace(module)) is None: return None - for fuser_cls in (GLUFuser, QKVFuser, RMSNormFuser): + for fuser_cls in (GLUFuser, QKVFuser, PackedQKVFuser, RMSNormFuser): if (fuser := fuser_cls.match(graph, module)) is not None: - if isinstance(fuser, StackedFuser): + if isinstance(fuser, RewriteFuser): try: fuser.update_forward(module) except Exception as exc: diff --git a/vllm/model_executor/models/transformers/fusers/__init__.py b/vllm/model_executor/models/transformers/fusers/__init__.py index 58910b0ecc3..ed4da5155cf 100644 --- a/vllm/model_executor/models/transformers/fusers/__init__.py +++ b/vllm/model_executor/models/transformers/fusers/__init__.py @@ -2,17 +2,24 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Concrete fusers for the Transformers modeling backend.""" -from vllm.model_executor.models.transformers.fusers.base import BaseFuser, StackedFuser +from vllm.model_executor.models.transformers.fusers.base import ( + BaseFuser, + RewriteFuser, + StackedFuser, +) from vllm.model_executor.models.transformers.fusers.glu import GLUFuser from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser +from vllm.model_executor.models.transformers.fusers.packed_qkv import PackedQKVFuser from vllm.model_executor.models.transformers.fusers.qkv import QKVFuser from vllm.model_executor.models.transformers.fusers.rms_norm import RMSNormFuser __all__ = [ "BaseFuser", + "RewriteFuser", "StackedFuser", "GLUFuser", "MoEBlockFuser", + "PackedQKVFuser", "QKVFuser", "RMSNormFuser", ] diff --git a/vllm/model_executor/models/transformers/fusers/base.py b/vllm/model_executor/models/transformers/fusers/base.py index abeff73e3ed..920e63101ac 100644 --- a/vllm/model_executor/models/transformers/fusers/base.py +++ b/vllm/model_executor/models/transformers/fusers/base.py @@ -57,27 +57,61 @@ class BaseFuser(ABC): return {} +def local_output_sizes(merged_name: str) -> str: + """Source for the per-rank widths of the merged linear `self.`.""" + merged = f"self.{merged_name}" + return f"[s // {merged}.tp_size for s in {merged}.output_sizes]" + + @dataclass -class StackedFuser(BaseFuser): - """A fuser that merges sibling projections into one stacked linear and - rewrites the forward to call it. +class RewriteFuser(BaseFuser): + """A fuser that rewrites the module's forward and rebinds it. - `match` and `update_forward` analyse the class once; `fuse` builds the merged - submodule and binds the compiled forward on an instance in place, so it keeps - its class and any attribute the fusion does not consume. + `match` and `update_forward` analyse the class once; `fuse` swaps the + submodules and binds the compiled forward on an instance in place, so it + keeps its class and any attribute the fusion does not consume. """ - merged_name: ClassVar[str] - """Attribute name of the merged module created by `update_attrs`.""" - merged_cls: ClassVar[str] - """Name of the vLLM class the merged projection becomes (for logging).""" - source_cls: str """Class of the HF module the fused projections belonged to (for logging).""" fused_forward: Callable = field(init=False, repr=False) """The compiled rewritten forward, set by `update_forward`.""" + @abstractmethod + def update_forward(self, module: nn.Module) -> None: + """Rewrite and compile `type(module)`'s forward source. + + Raises if the source does not admit the rewrite (fusion is then skipped). + """ + + @abstractmethod + def update_attrs( + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" + ) -> None: + """Replace `module`'s submodules with their vLLM equivalents.""" + + def fuse( + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" + ) -> nn.Module: + """Fuse an already-validated `module` in place (see `Fusers.__getitem__`). + + Builds the merged submodule and binds the compiled forward.""" + self.update_attrs(module, prefix, vllm_config) + module.forward = types.MethodType(self.fused_forward, module) + return module + + +@dataclass +class StackedFuser(RewriteFuser): + """A fuser that merges sibling projections into one stacked linear and + rewrites the forward to call it.""" + + merged_name: ClassVar[str] + """Attribute name of the merged module created by `update_attrs`.""" + merged_cls: ClassVar[str] + """Name of the vLLM class the merged projection becomes (for logging).""" + def info(self, name: str) -> str: sources = " + ".join(shard for shard, _ in self.shards) return ( @@ -108,26 +142,3 @@ class StackedFuser(BaseFuser): """`{merged_name: [projection names]}` so quantization can unpack the fused layer into its per-shard configs.""" return {self.merged_name: [name for name, _ in self.shards]} - - @abstractmethod - def update_forward(self, module: nn.Module) -> None: - """Rewrite and compile `type(module)`'s forward source. - - Raises if the source does not admit the rewrite (fusion is then skipped). - """ - - @abstractmethod - def update_attrs( - self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" - ) -> None: - """Replace `module`'s submodules with the merged module.""" - - def fuse( - self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" - ) -> nn.Module: - """Fuse an already-validated `module` in place (see `Fusers.__getitem__`). - - Builds the merged submodule and binds the compiled forward.""" - self.update_attrs(module, prefix, vllm_config) - module.forward = types.MethodType(self.fused_forward, module) - return module diff --git a/vllm/model_executor/models/transformers/fusers/packed_qkv.py b/vllm/model_executor/models/transformers/fusers/packed_qkv.py new file mode 100644 index 00000000000..ae5612fa816 --- /dev/null +++ b/vllm/model_executor/models/transformers/fusers/packed_qkv.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Packed-QKV fuser: `c_attn(x).split((q, kv, kv))` -> a `QKVParallelLinear`.""" + +import ast +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from torch import fx, nn + +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import QKVParallelLinear +from vllm.model_executor.models.transformers.fusers.base import ( + RewriteFuser, + local_output_sizes, +) +from vllm.model_executor.models.transformers.fx_utils import ( + compile_forward, + is_method, + recover_forward, + returned_linear, + upstream_linear, +) +from vllm.model_executor.models.transformers.utils import ( + log_replacement, + replace_linear_class, +) +from vllm.model_executor.models.utils import maybe_prefix + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +logger = init_logger(__name__) + + +@dataclass +class PackedQKVFuser(RewriteFuser): + """Fuser for attention with q, k and v packed into one projection.""" + + qkv_name: str + o_name: str | None + q_size: int + kv_size: int + + def info(self, name: str) -> str: + return ( + f"Fused: {self.qkv_name} ({name}: {self.source_cls}) -> QKVParallelLinear" + ) + + @staticmethod + def _packed_sizes(node: fx.Node) -> tuple[int, int] | None: + """`(q, kv)` from a `split((q, kv, kv), ...)` call, if it is one.""" + if not is_method(node, "split") or len(node.args) < 2: + return None + sizes = node.args[1] + if not isinstance(sizes, (tuple, list)) or len(sizes) != 3: + return None + if not all(isinstance(size, int) for size in sizes): + return None + q_size, k_size, v_size = sizes + if k_size != v_size or q_size < k_size: + return None + return q_size, k_size + + @classmethod + def match(cls, graph: fx.Graph, module: nn.Module) -> "PackedQKVFuser | None": + for node in graph.nodes: + if (sizes := cls._packed_sizes(node)) is None: + continue + q_size, kv_size = sizes + qkv_node = upstream_linear(node.args[0], module) + if qkv_node is None: + continue + qkv_name = str(qkv_node.target) + # The split must consume the whole projection. + if module.get_submodule(qkv_name).out_features != q_size + 2 * kv_size: + continue + # o_proj produces the module's output and consumes the query width. + o_name = returned_linear(graph, module) + if o_name == qkv_name or ( + o_name is not None + and module.get_submodule(o_name).in_features != q_size + ): + o_name = None + return cls( + source_cls=type(module).__name__, + qkv_name=qkv_name, + o_name=o_name, + q_size=q_size, + kv_size=kv_size, + ) + return None + + def _split_call(self, funcdef: ast.FunctionDef) -> ast.Call: + """The unique `self.(...)....split((a, b, c), ...)` call.""" + calls = [ + node + for node in ast.walk(funcdef) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "split" + and node.args + and isinstance(node.args[0], (ast.Tuple, ast.List)) + and len(node.args[0].elts) == 3 + and any( + isinstance(inner, ast.Attribute) and inner.attr == self.qkv_name + for inner in ast.walk(node.func.value) + ) + ] + if len(calls) != 1: + raise ValueError(f"{self.qkv_name} has {len(calls)} three-way splits") + return calls[0] + + def update_forward(self, module: nn.Module) -> None: + """Rewrite the split sizes to the sharded projection's per-rank widths.""" + funcdef, fn = recover_forward(type(module)) + split = self._split_call(funcdef) + # (q, kv, kv) -> [s // qkv.tp_size for s in qkv.output_sizes] + sections = local_output_sizes(self.qkv_name) + split.args[0] = ast.parse(sections, mode="eval").body + self.fused_forward = compile_forward(funcdef, fn) + + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: + """Shapes must be compatible with a head-sharded packed GEMM.""" + head_size = vllm_config.model_config.get_head_size() + qkv = module.get_submodule(self.qkv_name) + compatible = ( + self.q_size % head_size == 0 + and self.kv_size % head_size == 0 + and qkv.out_features == self.q_size + 2 * self.kv_size + ) + if not compatible: + logger.debug("%s is not compatible with packed QKV fusion", type(module)) + return compatible + + def update_attrs( + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" + ) -> None: + quant_config = vllm_config.quant_config + head_size = vllm_config.model_config.get_head_size() + qkv_prefix = maybe_prefix(prefix, self.qkv_name) + qkv = module.get_submodule(self.qkv_name) + merged = QKVParallelLinear( + hidden_size=qkv.in_features, + head_size=head_size, + total_num_heads=self.q_size // head_size, + total_num_kv_heads=self.kv_size // head_size, + bias=qkv.bias is not None, + quant_config=quant_config, + prefix=qkv_prefix, + return_bias=False, + ) + setattr(module, self.qkv_name, merged) + log_replacement(qkv_prefix, qkv, merged) + # If there is an output projection, we know it must be rowwise. + if self.o_name is not None: + o_prefix = maybe_prefix(prefix, self.o_name) + o_proj = module.get_submodule(self.o_name) + new_o = replace_linear_class( + o_proj, "rowwise", quant_config, prefix=o_prefix + ) + setattr(module, self.o_name, new_o) + log_replacement(o_prefix, o_proj, new_o) diff --git a/vllm/model_executor/models/transformers/fusers/qkv.py b/vllm/model_executor/models/transformers/fusers/qkv.py index b9fe9771fe1..370c0cdf4fa 100644 --- a/vllm/model_executor/models/transformers/fusers/qkv.py +++ b/vllm/model_executor/models/transformers/fusers/qkv.py @@ -10,7 +10,10 @@ from torch import fx, nn from vllm.logger import init_logger from vllm.model_executor.layers.linear import QKVParallelLinear -from vllm.model_executor.models.transformers.fusers.base import StackedFuser +from vllm.model_executor.models.transformers.fusers.base import ( + StackedFuser, + local_output_sizes, +) from vllm.model_executor.models.transformers.fx_utils import ( compile_forward, innermost_block, @@ -134,7 +137,7 @@ class QKVFuser(StackedFuser): if names & set(temps): raise ValueError("fused temporaries would shadow existing names") merged = f"self.{self.merged_name}" - sections = f"[s // {merged}.tp_size for s in {merged}.output_sizes]" + sections = local_output_sizes(self.merged_name) template = f"{', '.join(temps)} = {merged}(__arg__).split({sections}, -1)" assign = ast.parse(template).body[0] arg = next( diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py index 2dbc3a499de..d665bcfe8cf 100644 --- a/vllm/model_executor/models/transformers/fx_utils.py +++ b/vllm/model_executor/models/transformers/fx_utils.py @@ -394,8 +394,10 @@ def output_value(graph: fx.Graph) -> object | None: def upstream_linear(node: object, module: nn.Module) -> fx.Node | None: """Nearest linear producing `node`, walking back through splits/reshapes. - Never walks through a leaf call (e.g. an attention interface): its inputs - are what attention consumes, not what produced the value.""" + Non-linear submodules are transparent too (e.g. the dropout GPT-style + attentions apply after their output projection). Never walks through a leaf + call (e.g. an attention interface): its inputs are what attention consumes, + not what produced the value.""" stack = [node] seen: set[fx.Node] = set() while stack: @@ -405,7 +407,11 @@ def upstream_linear(node: object, module: nn.Module) -> fx.Node | None: seen.add(current) if is_linear(current, module): return current - if current.op in ("call_function", "call_method") and not is_leaf_call(current): + if current.op in ( + "call_function", + "call_method", + "call_module", + ) and not is_leaf_call(current): stack.extend(current.args) return None