Support MLA properly in the Transformers modeling backend (#48250)

Signed-off-by: Harry Mellor <[email protected]>

tests/models/transformers/fusers/test_mla.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the Transformers modeling backend's MLA fuser.
+
+The fuser must discover every MLA submodule structurally (never by assuming the
+Transformers attribute names) and, when the query is low-rank, merge the checkpoint's
+separate `q_a_proj`/`kv_a_proj_with_mqa` weights into the single fused down-projection.
+"""
+
+import pytest
+import torch
+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 MLAFuser
+from vllm.model_executor.models.transformers.fx_utils import trace
+
+_FUSED_QKV_A_PROJ = MLAFuser.merged_name
+
+
+def _match(q_lora_rank: int | None) -> MLAFuser | None:
+    """Match a meta `DeepseekV2Attention` directly (bypassing the per-class
+    `get_fuser` cache, so both q_lora variants of the same class are seen)."""
+    pytest.importorskip("transformers.models.deepseek_v2.modeling_deepseek_v2")
+    from transformers.models.deepseek_v2.configuration_deepseek_v2 import (
+        DeepseekV2Config,
+    )
+    from transformers.models.deepseek_v2.modeling_deepseek_v2 import (
+        DeepseekV2Attention,
+    )
+
+    cfg = DeepseekV2Config(
+        hidden_size=256,
+        num_attention_heads=16,
+        kv_lora_rank=128,
+        qk_rope_head_dim=32,
+        qk_nope_head_dim=32,
+        v_head_dim=32,
+        q_lora_rank=q_lora_rank,
+        num_hidden_layers=1,
+    )
+    with torch.device("meta"):
+        attn = DeepseekV2Attention(cfg, layer_idx=0)
+    return MLAFuser.match(trace(attn), attn)
+
+
+def test_discovers_modules_without_q_lora():
+    fuser = _match(q_lora_rank=None)
+    assert isinstance(fuser, MLAFuser)
+    assert not fuser.has_q_lora
+    assert fuser.q_proj_name == "q_proj"
+    assert fuser.kv_a_proj_name == "kv_a_proj_with_mqa"
+    assert fuser.kv_a_layernorm_name == "kv_a_layernorm"
+    assert fuser.kv_b_proj_name == "kv_b_proj"
+    assert fuser.o_proj_name == "o_proj"
+    assert fuser.q_a_proj_name is None
+    # Nothing is stacked without a query LoRA.
+    assert fuser.packed_modules_mapping == {}
+    assert fuser.orig_to_new_stacked("model.layers.0.self_attn") == {}
+
+
+def test_discovers_modules_with_q_lora():
+    fuser = _match(q_lora_rank=64)
+    assert isinstance(fuser, MLAFuser)
+    assert fuser.has_q_lora
+    assert fuser.q_a_proj_name == "q_a_proj"
+    assert fuser.q_a_layernorm_name == "q_a_layernorm"
+    assert fuser.q_b_proj_name == "q_b_proj"
+    assert fuser.kv_a_proj_name == "kv_a_proj_with_mqa"
+    assert fuser.kv_a_layernorm_name == "kv_a_layernorm"
+    assert fuser.kv_b_proj_name == "kv_b_proj"
+    assert fuser.o_proj_name == "o_proj"
+    assert fuser.q_proj_name is None
+
+
+def test_q_lora_stacks_qkv_a_proj():
+    """The MLA layer reads `q_a_proj` and `kv_a_proj_with_mqa` fused into one
+    down-projection, so both checkpoint weights must remap into it."""
+    fuser = _match(q_lora_rank=64)
+    assert isinstance(fuser, MLAFuser)
+    prefix = "model.layers.0.self_attn"
+    merged = f"{prefix}.{_FUSED_QKV_A_PROJ}"
+    assert fuser.packed_modules_mapping == {
+        _FUSED_QKV_A_PROJ: ["q_a_proj", "kv_a_proj_with_mqa"]
+    }
+    assert fuser.orig_to_new_stacked(prefix) == {
+        f"{prefix}.q_a_proj": (merged, 0),
+        f"{prefix}.kv_a_proj_with_mqa": (merged, 1),
+    }
+
+
+class _Norm(nn.Module):
+    """A real RMSNorm computation: `match` verifies chain norms via `RMSNormFuser`."""
+
+    def __init__(self, n: int):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(n))
+
+    def forward(self, x):
+        variance = x.pow(2).mean(-1, keepdim=True)
+        return self.weight * (x * torch.rsqrt(variance + 1e-6))
+
+
+class RenamedMLA(nn.Module):
+    """An MLA-shaped attention whose children have non-standard names, proving
+    discovery is by structure and not attribute name."""
+
+    def __init__(self):
+        super().__init__()
+        heads, kv_lora, rope, nope, v, hidden = 4, 32, 8, 8, 16, 64
+        self.alpha = nn.Linear(hidden, heads * (nope + rope))  # q_proj
+        self.beta = nn.Linear(hidden, kv_lora + rope)  # kv_a_proj_with_mqa
+        self.gamma = _Norm(kv_lora)  # kv_a_layernorm
+        self.delta = nn.Linear(kv_lora, heads * (nope + v))  # kv_b_proj
+        self.omega = nn.Linear(heads * v, hidden)  # o_proj (uncalled)
+        self.kv_lora, self.rope = kv_lora, rope
+
+    def forward(self, hidden_states):
+        q = self.alpha(hidden_states)
+        kv_lora, k_pe = torch.split(
+            self.beta(hidden_states), [self.kv_lora, self.rope], dim=-1
+        )
+        expanded = self.delta(self.gamma(kv_lora))
+        # Stand-in for the attention interface; `match` finds `o_proj` (omega) from
+        # the source as the Linear producing the returned value.
+        attn_output = expanded.sum() + q.sum() + k_pe.sum()
+        attn_output = self.omega(attn_output)
+        return attn_output
+
+
+def test_discovers_modules_under_arbitrary_names():
+    """Discovery is purely structural: `RenamedMLA` gives its children non-standard
+    names, and `match` still locates each projection by dataflow."""
+    with torch.device("meta"):
+        module = RenamedMLA()
+        fuser = MLAFuser.match(trace(module), module)
+    assert isinstance(fuser, MLAFuser)
+    assert not fuser.has_q_lora
+    assert fuser.q_proj_name == "alpha"
+    assert fuser.kv_a_proj_name == "beta"
+    assert fuser.kv_a_layernorm_name == "gamma"
+    assert fuser.kv_b_proj_name == "delta"
+    assert fuser.o_proj_name == "omega"
+
+
+class GLU(nn.Module):
+    """A gated MLP: matches the GLU fuser, never MLA."""
+
+    def __init__(self):
+        super().__init__()
+        self.gate = nn.Linear(16, 16)
+        self.up = nn.Linear(16, 16)
+
+    def forward(self, x):
+        return self.up(F.silu(self.gate(x)))
+
+
+def test_non_mla_is_not_matched():
+    with torch.device("meta"):
+        assert not isinstance(get_fuser(GLU()), MLAFuser)

tests/models/transformers/test_backend.py

@@ -35,6 +35,12 @@ def get_num_fused(model) -> tuple[int, int]:
     return glu, qkv
 
 
+def count_mla_layers(model) -> int:
+    from vllm.model_executor.layers.attention import MLAAttention
+
+    return sum(isinstance(m, MLAAttention) for m in model.attention_instances.values())
+
+
 def check_implementation(
     runner_ref: type[HfRunner | VllmRunner],
     runner_test: type[VllmRunner],
@@ -97,17 +103,6 @@ def test_models(
     model_impl: str,
     num_fused: tuple[int, int],
 ) -> None:
-    import transformers
-    from packaging.version import Version
-
-    installed = Version(transformers.__version__)
-    required = Version("5.0.0")
-    if model == "allenai/OLMoE-1B-7B-0924" and installed < required:
-        pytest.skip(
-            "MoE models with the Transformers modeling backend require "
-            f"transformers>={required}, but got {installed}"
-        )
-
     check_implementation(
         hf_runner,
         vllm_runner,
@@ -132,6 +127,42 @@ def test_hybrid_attention(vllm_runner: type[VllmRunner]) -> None:
     )
 
 
+def test_mla(vllm_runner: type[VllmRunner], example_prompts: list[str]) -> None:
+    import transformers
+    from packaging.version import Version
+
+    installed = Version(transformers.__version__)
+    required = Version("5.15.0.dev0")
+    if installed < required:
+        pytest.skip(
+            "MLA models with the Transformers modeling backend require "
+            f"transformers>={required}, but got {installed}"
+        )
+
+    model = get_model("DeepseekV2ForCausalLM")  # DeepSeek-V2-Lite, MLA + MoE
+    args = (example_prompts, 32, 5)
+    kwargs: dict[str, Any] = {"max_model_len": 2048, "enforce_eager": True}
+
+    with vllm_runner(
+        model, model_impl="transformers", trust_remote_code=False, **kwargs
+    ) as model_test:
+        model_config = model_test.llm.llm_engine.model_config
+        assert model_config.using_transformers_backend()
+        num_layers = model_config.hf_config.get_text_config().num_hidden_layers
+        assert model_test.apply_model(count_mla_layers) == [num_layers]
+        outputs_test = model_test.generate_greedy_logprobs(*args)
+
+    with vllm_runner(model, model_impl="auto") as model_ref:
+        outputs_ref = model_ref.generate_greedy_logprobs(*args)
+
+    check_logprobs_close(
+        outputs_0_lst=outputs_ref,
+        outputs_1_lst=outputs_test,
+        name_0="native",
+        name_1="transformers",
+    )
+
+
 @multi_gpu_test(num_gpus=2)
 def test_distributed(
     hf_runner: type[HfRunner],

vllm/config/model.py

@@ -1789,7 +1789,13 @@ class ModelConfig:
 
     @property
     def use_mla(self) -> bool:
-        return self.is_deepseek_mla and not envs.VLLM_MLA_DISABLE
+        if envs.VLLM_MLA_DISABLE:
+            return False
+        if self.using_transformers_backend():
+            # kv_lora_rank indicates that a Transformers model implementation uses MLA
+            return getattr(self.hf_text_config, "kv_lora_rank", None) is not None
+        # Manually maintained list of model types for vLLM model implementations
+        return self.is_deepseek_mla
 
     @property
     def is_matryoshka(self) -> bool:

vllm/model_executor/models/transformers/__init__.py

@@ -40,7 +40,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
 if TYPE_CHECKING:
     import torch
 
-    from vllm.model_executor.layers.attention import Attention
+    from vllm.model_executor.layers.attention import Attention, MLAAttention
 
 
 def vllm_attention_forward(
@@ -53,7 +53,7 @@ def vllm_attention_forward(
     # Transformers kwargs
     scaling: float | None = None,
     # vLLM kwargs
-    attention_instances: dict[int, "Attention"] | None = None,
+    attention_instances: "dict[int, Attention] | None" = None,
     **kwargs,
 ):
     self_attn = attention_instances[module.layer_idx]
@@ -78,7 +78,38 @@ def vllm_attention_forward(
     return attn_output, None
 
 
-ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_attention_forward
+def vllm_mla_attention_forward(
+    # Transformers args
+    module: "torch.nn.Module",
+    query: "torch.Tensor",
+    kv_c_normed: "torch.Tensor",
+    k_pe: "torch.Tensor",
+    attention_mask: "torch.Tensor",
+    # Transformers kwargs
+    scaling: float | None = None,
+    # vLLM kwargs
+    attention_instances: "dict[int, MLAAttention] | None" = None,
+    **kwargs,
+):
+    self_attn = attention_instances[module.layer_idx]
+    # [batch=1, heads, num_tokens, qk_head_dim] -> [num_tokens, heads, qk_head_dim]
+    query = query.transpose(1, 2).flatten(0, 1)
+    num_tokens, num_heads = query.shape[:2]
+    # [batch=1, num_tokens, kv_lora_rank] -> [num_tokens, kv_lora_rank]
+    kv_c_normed = kv_c_normed.reshape(-1, kv_c_normed.shape[-1])
+    # [batch=1, heads=1, num_tokens, qk_rope] -> [num_tokens, 1, qk_rope]
+    k_pe = k_pe.reshape(-1, 1, k_pe.shape[-1])
+    attn_output = self_attn.forward(
+        query,
+        kv_c_normed,
+        k_pe,
+        output_shape=(num_tokens, num_heads * self_attn.v_head_dim),
+    )
+    return attn_output, None
+
+
+ALL_ATTENTION_FUNCTIONS.register("vllm", vllm_attention_forward)
+ALL_ATTENTION_FUNCTIONS.register("vllm_mla", vllm_mla_attention_forward)
 
 
 # Text only models

vllm/model_executor/models/transformers/base.py

@@ -16,6 +16,7 @@
 # limitations under the License.
 """Transformers modeling backend base class."""
 
+import os
 from collections.abc import Callable, Iterable
 from itertools import chain
 from operator import attrgetter
@@ -40,7 +41,10 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.attention import (
     Attention,
     EncoderOnlyAttention,
+    MLAAttention,
 )
+from vllm.model_executor.layers.attention.mla_attention import get_mla_dims
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.fused_moe import MoERunner
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from vllm.model_executor.models.interfaces import (
@@ -52,11 +56,13 @@ from vllm.model_executor.models.interfaces import (
 )
 from vllm.model_executor.models.interfaces_base import VllmModel
 from vllm.model_executor.models.transformers.fuser import BaseFuser, Fusers
+from vllm.model_executor.models.transformers.fusers import MLAFuser
 from vllm.model_executor.models.transformers.utils import (
     can_enable_torch_compile,
     get_feature_request_tip,
     init_on_device_without_buffers,
     log_replacement,
+    named_state,
     replace_conv_class,
     replace_linear_class,
 )
@@ -64,11 +70,11 @@ from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     PPMissingLayer,
     WeightsMapper,
+    extract_layer_index,
     make_empty_intermediate_tensors_factory,
     maybe_prefix,
 )
 from vllm.sequence import IntermediateTensors
-from vllm.v1.attention.backend import AttentionType
 
 if TYPE_CHECKING:
     from transformers import PreTrainedModel
@@ -129,6 +135,9 @@ class Base(
         self.packed_modules_mapping: dict[str, list[str]] = {}
         """Fused module -> constituent projections, populated by `recursive_replace`
         for the quantization machinery and loaders (e.g. bitsandbytes)."""
+        self.fusers: dict[str, BaseFuser] = {}
+        """Module qualname -> the fuser applied to it, populated
+        by `recursive_replace` for `create_attention_instances`."""
 
         # Attrs for Eagle3 (see self.set_aux_hidden_state_layers)
         self._target_class: type[nn.Module] = nn.Module
@@ -192,6 +201,9 @@ class Base(
         # Initialize any parameters that have not had their modules replaced
         self.init_parameters(self.model)
 
+        # Upcast weights Transformers always keeps in fp32
+        self.keep_in_fp32(self.model)
+
         # Pipeline parallel intermediate tensors
         self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
             ["hidden_states"], self.text_config.hidden_size
@@ -457,6 +469,8 @@ class Base(
 
         def register_fusion(fuser: BaseFuser, prefix: str):
             """Register a fused layer's mappings just before it is built."""
+            self.fusers[prefix] = fuser
+
             orig_to_new_stacked = fuser.orig_to_new_stacked(prefix)
             self.hf_to_vllm_mapper.orig_to_new_stacked.update(orig_to_new_stacked)
 
@@ -522,13 +536,84 @@ class Base(
         """
         Create `Attention` instances to inform KV cache allocation.
         """
+        mla_fusers = {}
+        attention_instances = {}
         text_config = self.text_config
+        attn_cls = self._get_attn_cls()
+
+        # kv_lora_rank indicates that this is an MLA model
+        if getattr(text_config, "kv_lora_rank", None) is not None:
+            mla_fusers = {
+                extract_layer_index(prefix): (prefix, fuser)
+                for prefix, fuser in self.fusers.items()
+                if isinstance(fuser, MLAFuser)
+            }
+            if attn_cls is MLAAttention:
+                text_config._attn_implementation = "vllm_mla"
+            else:
+                # MLA model not using MLAAttention: recompute head_size for full attn
+                qk_nope_head_dim = getattr(text_config, "qk_nope_head_dim", 0)
+                qk_rope_head_dim = getattr(text_config, "qk_rope_head_dim", 0)
+                if qk_head_dim := qk_nope_head_dim + qk_rope_head_dim:
+                    self.model_config.model_arch_config.head_size = qk_head_dim
 
         num_heads = self.model_config.get_num_attention_heads(self.parallel_config)
         head_size = self.model_config.get_head_size()
         num_kv_heads = self.model_config.get_num_kv_heads(self.parallel_config)
         logits_soft_cap = getattr(text_config, "attn_logit_softcapping", None)
 
+        pp_rank = self.pp_group.rank_in_group
+        pp_size = self.pp_group.world_size
+        start, end = get_pp_indices(text_config.num_hidden_layers, pp_rank, pp_size)
+
+        for i in range(start, end):
+            kwargs = dict(
+                num_heads=num_heads,
+                # NOTE: We use Llama scale as default, if it's set by
+                # Transformers, it's updated in vllm_attention_forward
+                scale=head_size**-0.5,
+                cache_config=self.cache_config,
+                quant_config=self.quant_config,
+                prefix=f"{i}.attn",
+            )
+
+            if attn_cls is MLAAttention:
+                prefix, fuser = mla_fusers[i]
+                mla_module = self.get_submodule(prefix)
+                dims = get_mla_dims(self.model_config)
+                kwargs.update(
+                    scale=mla_module.scaling,
+                    qk_nope_head_dim=dims.qk_nope_head_dim,
+                    qk_rope_head_dim=dims.qk_rope_head_dim,
+                    v_head_dim=dims.v_head_dim,
+                    q_lora_rank=dims.q_lora_rank,
+                    kv_lora_rank=dims.kv_lora_rank,
+                    kv_b_proj=mla_module.get_submodule(fuser.kv_b_proj_name),
+                )
+            else:
+                kwargs.update(
+                    head_size=head_size,
+                    num_kv_heads=num_kv_heads,
+                    logits_soft_cap=logits_soft_cap,
+                )
+
+                # Handle interleaved sliding window attention
+                if (
+                    hasattr(text_config, "layer_types")
+                    and text_config.layer_types[i] == "sliding_attention"
+                ):
+                    kwargs["per_layer_sliding_window"] = text_config.sliding_window
+
+            attn_instance = attn_cls(**kwargs)
+            if attn_cls is MLAAttention:
+                # Attach MLA attn_instance to mla_module so it appears in
+                # model.named_modules() and runs its process_weights_after_loading
+                mla_module._vllm_mla_attn = attn_instance
+            attention_instances[i] = attn_instance
+        return attention_instances
+
+    def _get_attn_cls(self) -> type[AttentionLayerBase]:
+        """Return the `Attention` class to use for this model's layers."""
         # In encoder models, the attention layers will have `is_causal=False`
         is_encoder = lambda module: not getattr(module, "is_causal", True)
         has_encoder = lambda model: any(is_encoder(m) for m in model.modules())
@@ -537,44 +622,18 @@ class Base(
         # found in a text only model, we assume the whole model is an encoder model
         if has_encoder(self.model) and not is_multimodal(self.config):
             self.check_version("5.0.0", "encoder models support")
-            attn_type = AttentionType.ENCODER_ONLY
-        else:
-            attn_type = AttentionType.DECODER
-
-        pp_rank = self.pp_group.rank_in_group
-        pp_size = self.pp_group.world_size
-        start, end = get_pp_indices(text_config.num_hidden_layers, pp_rank, pp_size)
-
-        attention_instances = {}
-        for i in range(start, end):
-            # Handle interleaved sliding window attention
-            per_layer_sliding_window = None
-            if (
-                hasattr(self.config, "layer_types")
-                and self.config.layer_types[i] == "sliding_attention"
-            ):
-                per_layer_sliding_window = self.config.sliding_window
-
-            attn_cls = (
-                EncoderOnlyAttention
-                if attn_type == AttentionType.ENCODER_ONLY
-                else Attention
+            return EncoderOnlyAttention
+        if self.model_config.use_mla:
+            self.check_version("5.15.0.dev0", "optimized MLA support")
+            if any(isinstance(fuser, MLAFuser) for fuser in self.fusers.values()):
+                return MLAAttention
+            logger.warning_once(
+                "This model uses MLA but `MLAFuser` failed to match and/or fuse any "
+                "MLA attention layers. Falling back to full attention with a padded "
+                "`value` head dimension."
             )
-            attention_instances[i] = attn_cls(
-                num_heads=num_heads,
-                head_size=head_size,
-                # NOTE: We use Llama scale as default, if it's set by
-                # Transformers, it's updated in vllm_attention_forward
-                scale=head_size**-0.5,
-                num_kv_heads=num_kv_heads,
-                cache_config=self.cache_config,
-                quant_config=self.quant_config,
-                logits_soft_cap=logits_soft_cap,
-                per_layer_sliding_window=per_layer_sliding_window,
-                prefix=f"{i}.attn",
-                attn_type=attn_type,
-            )
-        return attention_instances
+            os.environ["VLLM_MLA_DISABLE"] = "1"
+        return Attention
 
     def init_parameters(self, module: nn.Module, dtype: torch.dtype | None = None):
         """
@@ -604,6 +663,18 @@ class Base(
 
         _init_parameters(module)
 
+    def keep_in_fp32(self, module: nn.Module):
+        """Honor `_keep_in_fp32_modules_strict` as `from_pretrained` would."""
+        if self.model_config.dtype not in (torch.float16, torch.bfloat16):
+            return
+        fragments = getattr(module, "_keep_in_fp32_modules_strict", None)
+        if not fragments:
+            return
+        pattern = re.compile("|".join(re.escape(f) for f in fragments))
+        for name, tensor in named_state(module):
+            if not hasattr(tensor, "weight_loader") and pattern.search(name):
+                tensor.data = tensor.data.to(torch.float32)
+
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.get_input_embeddings()(input_ids)
 

vllm/model_executor/models/transformers/fuser.py

@@ -18,6 +18,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.models.transformers.fusers import (
     BaseFuser,
     GLUFuser,
+    MLAFuser,
     PackedQKVFuser,
     QKVFuser,
     RewriteFuser,
@@ -48,7 +49,7 @@ def get_fuser(module: nn.Module) -> BaseFuser | None:
         return None
     if (graph := trace(module)) is None:
         return None
-    for fuser_cls in (GLUFuser, QKVFuser, PackedQKVFuser, RMSNormFuser):
+    for fuser_cls in (MLAFuser, GLUFuser, QKVFuser, PackedQKVFuser, RMSNormFuser):
         if (fuser := fuser_cls.match(graph, module)) is not None:
             if isinstance(fuser, RewriteFuser):
                 try:
@@ -74,14 +75,14 @@ def get_fuser(module: nn.Module) -> BaseFuser | None:
 
 
 class Fusers(UserDict):
-    """Mapping from module class to fuser, for all fusable classes in a model."""
+    """Mapping from module class and shape to fuser, for all fusable modules."""
 
     def __init__(self, model: nn.Module, vllm_config: "VllmConfig"):
         self.vllm_config = vllm_config
-        super().__init__({type(m): get_fuser(m) for m in model.modules()})
+        super().__init__({key(m): get_fuser(m) for m in model.modules()})
 
     def __getitem__(self, m: nn.Module) -> BaseFuser | None:
-        fuser = self.data.get(type(m))
+        fuser = self.data.get(key(m))
         if fuser is not None and fuser.validate(m, self.vllm_config):
             return fuser
         return None

vllm/model_executor/models/transformers/fusers/__init__.py

@@ -8,6 +8,7 @@ from vllm.model_executor.models.transformers.fusers.base import (
     StackedFuser,
 )
 from vllm.model_executor.models.transformers.fusers.glu import GLUFuser
+from vllm.model_executor.models.transformers.fusers.mla import MLAFuser
 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
@@ -18,6 +19,7 @@ __all__ = [
     "RewriteFuser",
     "StackedFuser",
     "GLUFuser",
+    "MLAFuser",
     "MoEBlockFuser",
     "PackedQKVFuser",
     "QKVFuser",

vllm/model_executor/models/transformers/fusers/mla.py

@@ -0,0 +1,325 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MLA fuser: adapt a Transformers MLA attention module for vLLM's MLA layer."""
+
+import ast
+import inspect
+import textwrap
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, ClassVar
+
+from torch import fx, nn
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import MergedColumnParallelLinear
+from vllm.model_executor.models.transformers.fusers.base import StackedFuser
+from vllm.model_executor.models.transformers.fusers.rms_norm import RMSNormFuser
+from vllm.model_executor.models.transformers.fx_utils import (
+    compile_forward,
+    downstream_linear,
+    is_linear,
+    recover_forward,
+    replace_expr,
+    returned_linear,
+    single_self_call,
+    trace,
+    upstream_linear,
+)
+from vllm.model_executor.models.transformers.utils import (
+    log_replacement,
+    replace_linear_class,
+)
+from vllm.model_executor.models.utils import ShardId, maybe_prefix
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+logger = init_logger(__name__)
+
+# Temporaries the fused down-projection binds in the rewritten forward.
+_Q_A_TEMP = "__q_a_fused"
+_KV_A_TEMP = "__kv_a_fused"
+
+
+def _consumes_placeholder(node: fx.Node) -> bool:
+    """Whether `node` is a linear applied directly to a `forward` input."""
+    return (
+        len(node.args) == 1
+        and isinstance(node.args[0], fx.Node)
+        and node.args[0].op == "placeholder"
+    )
+
+
+def _norm_size(norm: nn.Module) -> int:
+    weight = getattr(norm, "weight", None)
+    return weight.numel() if weight is not None else -1
+
+
+def _is_rms_norm(module: nn.Module) -> bool:
+    """Whether `module` computes an RMSNorm, verified by `RMSNormFuser`'s matcher."""
+    graph = trace(module)
+    return graph is not None and RMSNormFuser.match(graph, module) is not None
+
+
+def _top_level_index(funcdef: ast.FunctionDef, node: ast.AST) -> int:
+    """Index in `funcdef.body` of the top-level statement containing `node`."""
+    for index, stmt in enumerate(funcdef.body):
+        if any(child is node for child in ast.walk(stmt)):
+            return index
+    raise ValueError("node is not in the function body")
+
+
+def _single_expand_call(
+    funcdef: ast.FunctionDef, module: nn.Module, kv_b_proj_name: str
+) -> ast.Call:
+    """The KV expansion call `self.<method>(kv_c_normed, k_pe)`: the unique
+    two-argument call to a method of `module` with the expansion's signature.
+
+    The expansion is `kv_c_normed, k_pe -> key, value`: two inputs, every
+    `return` a pair, and it applies `kv_b_proj` -- the projection `MLAAttention`
+    owns after the bypass, absorbed into the attention computation.
+    """
+
+    def is_expansion_method(name: str) -> bool:
+        method = getattr(type(module), name, None)
+        if method is None:
+            return False
+        code = getattr(inspect.unwrap(method), "__code__", None)
+        if (
+            code is None
+            or code.co_argcount != 3  # self, kv_c_normed, k_pe
+            or kv_b_proj_name not in code.co_names
+        ):
+            return False
+        try:
+            tree = ast.parse(textwrap.dedent(inspect.getsource(method)))
+        except (OSError, TypeError):
+            return False
+        returns = [node for node in ast.walk(tree) if isinstance(node, ast.Return)]
+        return bool(returns) and all(
+            isinstance(ret.value, ast.Tuple) and len(ret.value.elts) == 2
+            for ret in returns
+        )
+
+    calls = [
+        node
+        for node in ast.walk(funcdef)
+        if isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Attribute)
+        and isinstance(node.func.value, ast.Name)
+        and node.func.value.id == "self"
+        and len(node.args) == 2
+        and not node.keywords
+        and is_expansion_method(node.func.attr)
+    ]
+    if len(calls) != 1:
+        raise ValueError("expected exactly one KV expansion method call")
+    return calls[0]
+
+
+@dataclass
+class MLAFuser(StackedFuser):
+    """Fuser for the MLA attention pattern."""
+
+    q_proj_name: str | None
+    q_a_proj_name: str | None
+    q_a_layernorm_name: str | None
+    q_b_proj_name: str | None
+    kv_a_proj_name: str
+    kv_a_layernorm_name: str
+    kv_b_proj_name: str
+    o_proj_name: str | None
+    merged_name: ClassVar[str] = "fused_qkv_a_proj"
+    merged_cls: ClassVar[str] = "MergedColumnParallelLinear"
+
+    @property
+    def has_q_lora(self) -> bool:
+        return self.q_a_proj_name is not None
+
+    def info(self, name: str) -> str:
+        info_str = (
+            f"Fused: {name} ({self.source_cls}) -> MLAAttention (attention interface)"
+        )
+        if self.has_q_lora:
+            info_str += "; " + super().info(name).removeprefix("Fused: ")
+        return info_str
+
+    @property
+    def shards(self) -> list[tuple[str, ShardId]]:
+        """`q_a_proj` and `kv_a_proj_with_mqa` stack into one down-projection."""
+        if self.has_q_lora:
+            return [(self.q_a_proj_name, 0), (self.kv_a_proj_name, 1)]
+        return []
+
+    @property
+    def packed_modules_mapping(self) -> dict[str, list[str]]:
+        if self.has_q_lora:
+            return super().packed_modules_mapping
+        return {}
+
+    @classmethod
+    def match(cls, graph: fx.Graph, module: nn.Module) -> "MLAFuser | None":
+        # Find all `rms_norm(linear(placeholder))` chains.
+        chains = []
+        for node in graph.nodes:
+            if node.op != "call_module" or is_linear(node, module) or not node.args:
+                continue
+            source = upstream_linear(node.args[0], module)
+            if source is None or not _consumes_placeholder(source):
+                continue
+            if not _is_rms_norm(module.get_submodule(node.target)):
+                continue
+            chains.append((source, node))
+
+        # Tell the chains apart by width.
+        def is_kv_a(chain) -> bool:
+            source, rms_norm = chain
+            source_mod = module.get_submodule(source.target)
+            rms_norm_mod = module.get_submodule(rms_norm.target)
+            return source_mod.out_features != _norm_size(rms_norm_mod)
+
+        kv_a_chains = [chain for chain in chains if is_kv_a(chain)]
+        q_a_chains = [chain for chain in chains if not is_kv_a(chain)]
+        # Exactly one KV chain is MLA's signature. The query chain is optional.
+        if len(kv_a_chains) != 1 or len(q_a_chains) > 1:
+            return None
+        kv_a_proj, kv_a_layernorm = kv_a_chains[0]
+
+        # Linear children claimed for a role so far; the rest resolve by elimination.
+        claimed_linears = {kv_a_proj.target}
+        q_proj_name = q_a_proj = q_a_layernorm = q_b_proj = None
+        if q_a_chains:
+            # Find `q_b_proj(q_a_layernorm(q_a_proj(placeholder)))`.
+            q_a_proj, q_a_layernorm = q_a_chains[0]
+            q_b_proj = downstream_linear(q_a_layernorm, module)
+            if q_b_proj is None:
+                return None
+            claimed_linears |= {q_a_proj.target, q_b_proj.target}
+        else:
+            # Find `q_proj(placeholder)`.
+            placeholder_linears = {
+                node.target
+                for node in graph.nodes
+                if is_linear(node, module) and _consumes_placeholder(node)
+            }
+            if len(q_proj_candidates := placeholder_linears - claimed_linears) != 1:
+                return None
+            q_proj_name = next(iter(q_proj_candidates))
+            claimed_linears.add(q_proj_name)
+
+        # Find `kv_b_proj(kv_a_layernorm(...))`.
+        kv_b_proj = downstream_linear(kv_a_layernorm, module)
+        if kv_b_proj is None or kv_b_proj.target in claimed_linears:
+            return None
+        claimed_linears.add(kv_b_proj.target)
+
+        # Find `o_proj` if it is returned by the forward graph.
+        o_proj_name = returned_linear(graph, module)
+        if o_proj_name in claimed_linears:
+            o_proj_name = None
+
+        return cls(
+            source_cls=type(module).__name__,
+            q_proj_name=q_proj_name,
+            q_a_proj_name=q_a_proj.target if q_a_proj else None,
+            q_a_layernorm_name=q_a_layernorm.target if q_a_layernorm else None,
+            q_b_proj_name=q_b_proj.target if q_b_proj else None,
+            kv_a_proj_name=kv_a_proj.target,
+            kv_a_layernorm_name=kv_a_layernorm.target,
+            kv_b_proj_name=kv_b_proj.target,
+            o_proj_name=o_proj_name,
+        )
+
+    def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
+        return vllm_config.model_config.use_mla
+
+    def update_forward(self, module: nn.Module) -> None:
+        """Merge `q_a_proj` and `kv_a_proj` into one fused down proj then split.
+        Bypass the KV expansion method so the compressed latent reaches the `vllm_mla`
+        attention interface unexpanded."""
+        funcdef, fn = recover_forward(type(module))
+        if self.has_q_lora:
+            # q_a_proj is usually inside the `else` of `if self.q_lora_rank is None`.
+            # The fused call is inserted at the top-level statement preceding both.
+            q_call = single_self_call(funcdef, self.q_a_proj_name)
+            kv_call = single_self_call(funcdef, self.kv_a_proj_name)
+            if ast.dump(q_call.args[0]) != ast.dump(kv_call.args[0]):
+                raise ValueError("down-projections read different inputs")
+            names = {n.id for n in ast.walk(funcdef) if isinstance(n, ast.Name)}
+            if names & {_Q_A_TEMP, _KV_A_TEMP}:
+                raise ValueError("fused temporaries would shadow existing names")
+
+            merged = f"self.{self.merged_name}"
+            targets = f"{_Q_A_TEMP}, {_KV_A_TEMP}"
+            sections = f"[s // {merged}.tp_size for s in {merged}.output_sizes]"
+            source = f"{targets} = {merged}(__arg__).split({sections}, -1)"
+            assign = ast.parse(source=source).body[0]
+            placeholder = next(
+                node
+                for node in ast.walk(assign)
+                if isinstance(node, ast.Name) and node.id == "__arg__"
+            )
+            replace_expr(assign, placeholder, q_call.args[0])
+
+            index = min(_top_level_index(funcdef, call) for call in (q_call, kv_call))
+            ast.copy_location(assign, funcdef.body[index])
+            funcdef.body.insert(index, assign)
+            replace_expr(funcdef, q_call, ast.Name(id=_Q_A_TEMP, ctx=ast.Load()))
+            replace_expr(funcdef, kv_call, ast.Name(id=_KV_A_TEMP, ctx=ast.Load()))
+
+        # Transformers expands the latent into full key/value in a dedicated method.
+        # `MLAAttention` consumes the latent directly (absorbing `kv_b_proj`),
+        # so replace the expansion call with its own arguments so `kv_c_normed, k_pe`
+        # flow to the interface in place of `key, value`.
+        expand_call = _single_expand_call(funcdef, module, self.kv_b_proj_name)
+        replace_expr(
+            funcdef, expand_call, ast.Tuple(elts=list(expand_call.args), ctx=ast.Load())
+        )
+        self.fused_forward = compile_forward(funcdef, fn)
+
+    def update_attrs(self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"):
+        quant_config = vllm_config.quant_config
+
+        def replace_linear_by_name(name: str, style: str):
+            linear = module.get_submodule(name)
+            replacement = replace_linear_class(
+                linear, style, quant_config, prefix=maybe_prefix(prefix, name)
+            )
+            setattr(module, name, replacement)
+            log_replacement(maybe_prefix(prefix, name), linear, replacement)
+
+        if self.has_q_lora:
+            q_a = module.get_submodule(self.q_a_proj_name)
+            kv_a = module.get_submodule(self.kv_a_proj_name)
+            merged = MergedColumnParallelLinear(
+                input_size=q_a.in_features,
+                output_sizes=[q_a.out_features, kv_a.out_features],
+                bias=q_a.bias is not None,
+                quant_config=quant_config,
+                prefix=maybe_prefix(prefix, self.merged_name),
+                return_bias=False,
+                disable_tp=True,
+            )
+            logger.debug(
+                "%s: %s, %s: %s -> %s: %s",
+                self.q_a_proj_name,
+                q_a,
+                self.kv_a_proj_name,
+                kv_a,
+                self.merged_name,
+                merged,
+            )
+            setattr(module, self.merged_name, merged)
+            # The rewritten forward calls the merged projection instead.
+            delattr(module, self.q_a_proj_name)
+            delattr(module, self.kv_a_proj_name)
+            replace_linear_by_name(self.q_b_proj_name, "colwise")
+        else:
+            replace_linear_by_name(self.kv_a_proj_name, "replicate")
+            replace_linear_by_name(self.q_proj_name, "colwise")
+
+        replace_linear_by_name(self.kv_b_proj_name, "colwise")
+        # MLAAttention calls kv_b_proj and expects vLLM's default return_bias=True
+        module.get_submodule(self.kv_b_proj_name).return_bias = True
+        if self.o_proj_name is not None:
+            replace_linear_by_name(self.o_proj_name, "rowwise")

vllm/model_executor/models/transformers/fx_utils.py

@@ -416,6 +416,25 @@ def upstream_linear(node: object, module: nn.Module) -> fx.Node | None:
     return None
 
 
+def downstream_linear(node: fx.Node, module: nn.Module) -> fx.Node | None:
+    """Nearest linear consuming `node`'s output, walking through casts/scalings.
+
+    Never walks through a leaf call (e.g. an attention interface): what crosses
+    it is consumed by the attention computation, not projected."""
+    queue = list(node.users)
+    seen: set[fx.Node] = set()
+    while queue:
+        current = queue.pop(0)
+        if current in seen:
+            continue
+        seen.add(current)
+        if is_linear(current, module):
+            return current
+        if current.op in ("call_function", "call_method") and not is_leaf_call(current):
+            queue.extend(current.users)
+    return None
+
+
 def returned_linear(graph: fx.Graph, module: nn.Module) -> str | None:
     """Name of the Linear producing the graph's (first) output value."""
     value = output_value(graph)