Fix MQA with tensor parallelism on transformers modeling backend (#49987)

Signed-off-by: microslaw <milosz.grunwald@intel.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

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():

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:

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",
 ]

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_name>`."""
+    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

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.<qkv_name>(...)....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)

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(

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