@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 59b2fdfc4e

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

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Harry Mellor committed Aug 4, 2026 at 13:24 UTC 59b2fdfc4e237794c2b26f0781054ced73f5b8df
9 files changed +699 -52
tests/models/transformers/fusers/test_mla.py new
+161
@@ -0,0 +1,161 @@
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 MLA fuser.
4 +
5 +The fuser must discover every MLA submodule structurally (never by assuming the
6 +Transformers attribute names) and, when the query is low-rank, merge the checkpoint's
7 +separate `q_a_proj`/`kv_a_proj_with_mqa` weights into the single fused down-projection.
8 +"""
9 +
10 +import pytest
11 +import torch
12 +import torch.nn as nn
13 +import torch.nn.functional as F
14 +
15 +from vllm.model_executor.models.transformers.fuser import get_fuser
16 +from vllm.model_executor.models.transformers.fusers import MLAFuser
17 +from vllm.model_executor.models.transformers.fx_utils import trace
18 +
19 +_FUSED_QKV_A_PROJ = MLAFuser.merged_name
20 +
21 +
22 +def _match(q_lora_rank: int | None) -> MLAFuser | None:
23 + """Match a meta `DeepseekV2Attention` directly (bypassing the per-class
24 + `get_fuser` cache, so both q_lora variants of the same class are seen)."""
25 + pytest.importorskip("transformers.models.deepseek_v2.modeling_deepseek_v2")
26 + from transformers.models.deepseek_v2.configuration_deepseek_v2 import (
27 + DeepseekV2Config,
28 + )
29 + from transformers.models.deepseek_v2.modeling_deepseek_v2 import (
30 + DeepseekV2Attention,
31 + )
32 +
33 + cfg = DeepseekV2Config(
34 + hidden_size=256,
35 + num_attention_heads=16,
36 + kv_lora_rank=128,
37 + qk_rope_head_dim=32,
38 + qk_nope_head_dim=32,
39 + v_head_dim=32,
40 + q_lora_rank=q_lora_rank,
41 + num_hidden_layers=1,
42 + )
43 + with torch.device("meta"):
44 + attn = DeepseekV2Attention(cfg, layer_idx=0)
45 + return MLAFuser.match(trace(attn), attn)
46 +
47 +
48 +def test_discovers_modules_without_q_lora():
49 + fuser = _match(q_lora_rank=None)
50 + assert isinstance(fuser, MLAFuser)
51 + assert not fuser.has_q_lora
52 + assert fuser.q_proj_name == "q_proj"
53 + assert fuser.kv_a_proj_name == "kv_a_proj_with_mqa"
54 + assert fuser.kv_a_layernorm_name == "kv_a_layernorm"
55 + assert fuser.kv_b_proj_name == "kv_b_proj"
56 + assert fuser.o_proj_name == "o_proj"
57 + assert fuser.q_a_proj_name is None
58 + # Nothing is stacked without a query LoRA.
59 + assert fuser.packed_modules_mapping == {}
60 + assert fuser.orig_to_new_stacked("model.layers.0.self_attn") == {}
61 +
62 +
63 +def test_discovers_modules_with_q_lora():
64 + fuser = _match(q_lora_rank=64)
65 + assert isinstance(fuser, MLAFuser)
66 + assert fuser.has_q_lora
67 + assert fuser.q_a_proj_name == "q_a_proj"
68 + assert fuser.q_a_layernorm_name == "q_a_layernorm"
69 + assert fuser.q_b_proj_name == "q_b_proj"
70 + assert fuser.kv_a_proj_name == "kv_a_proj_with_mqa"
71 + assert fuser.kv_a_layernorm_name == "kv_a_layernorm"
72 + assert fuser.kv_b_proj_name == "kv_b_proj"
73 + assert fuser.o_proj_name == "o_proj"
74 + assert fuser.q_proj_name is None
75 +
76 +
77 +def test_q_lora_stacks_qkv_a_proj():
78 + """The MLA layer reads `q_a_proj` and `kv_a_proj_with_mqa` fused into one
79 + down-projection, so both checkpoint weights must remap into it."""
80 + fuser = _match(q_lora_rank=64)
81 + assert isinstance(fuser, MLAFuser)
82 + prefix = "model.layers.0.self_attn"
83 + merged = f"{prefix}.{_FUSED_QKV_A_PROJ}"
84 + assert fuser.packed_modules_mapping == {
85 + _FUSED_QKV_A_PROJ: ["q_a_proj", "kv_a_proj_with_mqa"]
86 + }
87 + assert fuser.orig_to_new_stacked(prefix) == {
88 + f"{prefix}.q_a_proj": (merged, 0),
89 + f"{prefix}.kv_a_proj_with_mqa": (merged, 1),
90 + }
91 +
92 +
93 +class _Norm(nn.Module):
94 + """A real RMSNorm computation: `match` verifies chain norms via `RMSNormFuser`."""
95 +
96 + def __init__(self, n: int):
97 + super().__init__()
98 + self.weight = nn.Parameter(torch.ones(n))
99 +
100 + def forward(self, x):
101 + variance = x.pow(2).mean(-1, keepdim=True)
102 + return self.weight * (x * torch.rsqrt(variance + 1e-6))
103 +
104 +
105 +class RenamedMLA(nn.Module):
106 + """An MLA-shaped attention whose children have non-standard names, proving
107 + discovery is by structure and not attribute name."""
108 +
109 + def __init__(self):
110 + super().__init__()
111 + heads, kv_lora, rope, nope, v, hidden = 4, 32, 8, 8, 16, 64
112 + self.alpha = nn.Linear(hidden, heads * (nope + rope)) # q_proj
113 + self.beta = nn.Linear(hidden, kv_lora + rope) # kv_a_proj_with_mqa
114 + self.gamma = _Norm(kv_lora) # kv_a_layernorm
115 + self.delta = nn.Linear(kv_lora, heads * (nope + v)) # kv_b_proj
116 + self.omega = nn.Linear(heads * v, hidden) # o_proj (uncalled)
117 + self.kv_lora, self.rope = kv_lora, rope
118 +
119 + def forward(self, hidden_states):
120 + q = self.alpha(hidden_states)
121 + kv_lora, k_pe = torch.split(
122 + self.beta(hidden_states), [self.kv_lora, self.rope], dim=-1
123 + )
124 + expanded = self.delta(self.gamma(kv_lora))
125 + # Stand-in for the attention interface; `match` finds `o_proj` (omega) from
126 + # the source as the Linear producing the returned value.
127 + attn_output = expanded.sum() + q.sum() + k_pe.sum()
128 + attn_output = self.omega(attn_output)
129 + return attn_output
130 +
131 +
132 +def test_discovers_modules_under_arbitrary_names():
133 + """Discovery is purely structural: `RenamedMLA` gives its children non-standard
134 + names, and `match` still locates each projection by dataflow."""
135 + with torch.device("meta"):
136 + module = RenamedMLA()
137 + fuser = MLAFuser.match(trace(module), module)
138 + assert isinstance(fuser, MLAFuser)
139 + assert not fuser.has_q_lora
140 + assert fuser.q_proj_name == "alpha"
141 + assert fuser.kv_a_proj_name == "beta"
142 + assert fuser.kv_a_layernorm_name == "gamma"
143 + assert fuser.kv_b_proj_name == "delta"
144 + assert fuser.o_proj_name == "omega"
145 +
146 +
147 +class GLU(nn.Module):
148 + """A gated MLP: matches the GLU fuser, never MLA."""
149 +
150 + def __init__(self):
151 + super().__init__()
152 + self.gate = nn.Linear(16, 16)
153 + self.up = nn.Linear(16, 16)
154 +
155 + def forward(self, x):
156 + return self.up(F.silu(self.gate(x)))
157 +
158 +
159 +def test_non_mla_is_not_matched():
160 + with torch.device("meta"):
161 + assert not isinstance(get_fuser(GLU()), MLAFuser)
tests/models/transformers/test_backend.py
+42 -11
@@ -35,6 +35,12 @@ def get_num_fused(model) -> tuple[int, int]:
35 return glu, qkv
36
37
38 +def count_mla_layers(model) -> int:
39 + from vllm.model_executor.layers.attention import MLAAttention
40 +
41 + return sum(isinstance(m, MLAAttention) for m in model.attention_instances.values())
42 +
43 +
44 def check_implementation(
45 runner_ref: type[HfRunner | VllmRunner],
46 runner_test: type[VllmRunner],
@@ -97,17 +103,6 @@ def test_models(
103 model_impl: str,
104 num_fused: tuple[int, int],
105 ) -> None:
100 - import transformers
101 - from packaging.version import Version
102 -
103 - installed = Version(transformers.__version__)
104 - required = Version("5.0.0")
105 - if model == "allenai/OLMoE-1B-7B-0924" and installed < required:
106 - pytest.skip(
107 - "MoE models with the Transformers modeling backend require "
108 - f"transformers>={required}, but got {installed}"
109 - )
110 -
106 check_implementation(
107 hf_runner,
108 vllm_runner,
@@ -132,6 +127,42 @@ def test_hybrid_attention(vllm_runner: type[VllmRunner]) -> None:
127 )
128
129
130 +def test_mla(vllm_runner: type[VllmRunner], example_prompts: list[str]) -> None:
131 + import transformers
132 + from packaging.version import Version
133 +
134 + installed = Version(transformers.__version__)
135 + required = Version("5.15.0.dev0")
136 + if installed < required:
137 + pytest.skip(
138 + "MLA models with the Transformers modeling backend require "
139 + f"transformers>={required}, but got {installed}"
140 + )
141 +
142 + model = get_model("DeepseekV2ForCausalLM") # DeepSeek-V2-Lite, MLA + MoE
143 + args = (example_prompts, 32, 5)
144 + kwargs: dict[str, Any] = {"max_model_len": 2048, "enforce_eager": True}
145 +
146 + with vllm_runner(
147 + model, model_impl="transformers", trust_remote_code=False, **kwargs
148 + ) as model_test:
149 + model_config = model_test.llm.llm_engine.model_config
150 + assert model_config.using_transformers_backend()
151 + num_layers = model_config.hf_config.get_text_config().num_hidden_layers
152 + assert model_test.apply_model(count_mla_layers) == [num_layers]
153 + outputs_test = model_test.generate_greedy_logprobs(*args)
154 +
155 + with vllm_runner(model, model_impl="auto") as model_ref:
156 + outputs_ref = model_ref.generate_greedy_logprobs(*args)
157 +
158 + check_logprobs_close(
159 + outputs_0_lst=outputs_ref,
160 + outputs_1_lst=outputs_test,
161 + name_0="native",
162 + name_1="transformers",
163 + )
164 +
165 +
166 @multi_gpu_test(num_gpus=2)
167 def test_distributed(
168 hf_runner: type[HfRunner],
vllm/config/model.py
+7 -1
@@ -1789,7 +1789,13 @@ class ModelConfig:
1789
1790 @property
1791 def use_mla(self) -> bool:
1792 - return self.is_deepseek_mla and not envs.VLLM_MLA_DISABLE
1792 + if envs.VLLM_MLA_DISABLE:
1793 + return False
1794 + if self.using_transformers_backend():
1795 + # kv_lora_rank indicates that a Transformers model implementation uses MLA
1796 + return getattr(self.hf_text_config, "kv_lora_rank", None) is not None
1797 + # Manually maintained list of model types for vLLM model implementations
1798 + return self.is_deepseek_mla
1799
1800 @property
1801 def is_matryoshka(self) -> bool:
vllm/model_executor/models/transformers/__init__.py
+34 -3
@@ -40,7 +40,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
40 if TYPE_CHECKING:
41 import torch
42
43 - from vllm.model_executor.layers.attention import Attention
43 + from vllm.model_executor.layers.attention import Attention, MLAAttention
44
45
46 def vllm_attention_forward(
@@ -53,7 +53,7 @@ def vllm_attention_forward(
53 # Transformers kwargs
54 scaling: float | None = None,
55 # vLLM kwargs
56 - attention_instances: dict[int, "Attention"] | None = None,
56 + attention_instances: "dict[int, Attention] | None" = None,
57 **kwargs,
58 ):
59 self_attn = attention_instances[module.layer_idx]
@@ -78,7 +78,38 @@ def vllm_attention_forward(
78 return attn_output, None
79
80
81 -ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_attention_forward
81 +def vllm_mla_attention_forward(
82 + # Transformers args
83 + module: "torch.nn.Module",
84 + query: "torch.Tensor",
85 + kv_c_normed: "torch.Tensor",
86 + k_pe: "torch.Tensor",
87 + attention_mask: "torch.Tensor",
88 + # Transformers kwargs
89 + scaling: float | None = None,
90 + # vLLM kwargs
91 + attention_instances: "dict[int, MLAAttention] | None" = None,
92 + **kwargs,
93 +):
94 + self_attn = attention_instances[module.layer_idx]
95 + # [batch=1, heads, num_tokens, qk_head_dim] -> [num_tokens, heads, qk_head_dim]
96 + query = query.transpose(1, 2).flatten(0, 1)
97 + num_tokens, num_heads = query.shape[:2]
98 + # [batch=1, num_tokens, kv_lora_rank] -> [num_tokens, kv_lora_rank]
99 + kv_c_normed = kv_c_normed.reshape(-1, kv_c_normed.shape[-1])
100 + # [batch=1, heads=1, num_tokens, qk_rope] -> [num_tokens, 1, qk_rope]
101 + k_pe = k_pe.reshape(-1, 1, k_pe.shape[-1])
102 + attn_output = self_attn.forward(
103 + query,
104 + kv_c_normed,
105 + k_pe,
106 + output_shape=(num_tokens, num_heads * self_attn.v_head_dim),
107 + )
108 + return attn_output, None
109 +
110 +
111 +ALL_ATTENTION_FUNCTIONS.register("vllm", vllm_attention_forward)
112 +ALL_ATTENTION_FUNCTIONS.register("vllm_mla", vllm_mla_attention_forward)
113
114
115 # Text only models
vllm/model_executor/models/transformers/base.py
+104 -33
@@ -16,6 +16,7 @@
16 # limitations under the License.
17 """Transformers modeling backend base class."""
18
19 +import os
20 from collections.abc import Callable, Iterable
21 from itertools import chain
22 from operator import attrgetter
@@ -40,7 +41,10 @@ from vllm.logger import init_logger
41 from vllm.model_executor.layers.attention import (
42 Attention,
43 EncoderOnlyAttention,
44 + MLAAttention,
45 )
46 +from vllm.model_executor.layers.attention.mla_attention import get_mla_dims
47 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
48 from vllm.model_executor.layers.fused_moe import MoERunner
49 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
50 from vllm.model_executor.models.interfaces import (
@@ -52,11 +56,13 @@ from vllm.model_executor.models.interfaces import (
56 )
57 from vllm.model_executor.models.interfaces_base import VllmModel
58 from vllm.model_executor.models.transformers.fuser import BaseFuser, Fusers
59 +from vllm.model_executor.models.transformers.fusers import MLAFuser
60 from vllm.model_executor.models.transformers.utils import (
61 can_enable_torch_compile,
62 get_feature_request_tip,
63 init_on_device_without_buffers,
64 log_replacement,
65 + named_state,
66 replace_conv_class,
67 replace_linear_class,
68 )
@@ -64,11 +70,11 @@ from vllm.model_executor.models.utils import (
70 AutoWeightsLoader,
71 PPMissingLayer,
72 WeightsMapper,
73 + extract_layer_index,
74 make_empty_intermediate_tensors_factory,
75 maybe_prefix,
76 )
77 from vllm.sequence import IntermediateTensors
71 -from vllm.v1.attention.backend import AttentionType
78
79 if TYPE_CHECKING:
80 from transformers import PreTrainedModel
@@ -129,6 +135,9 @@ class Base(
135 self.packed_modules_mapping: dict[str, list[str]] = {}
136 """Fused module -> constituent projections, populated by `recursive_replace`
137 for the quantization machinery and loaders (e.g. bitsandbytes)."""
138 + self.fusers: dict[str, BaseFuser] = {}
139 + """Module qualname -> the fuser applied to it, populated
140 + by `recursive_replace` for `create_attention_instances`."""
141
142 # Attrs for Eagle3 (see self.set_aux_hidden_state_layers)
143 self._target_class: type[nn.Module] = nn.Module
@@ -192,6 +201,9 @@ class Base(
201 # Initialize any parameters that have not had their modules replaced
202 self.init_parameters(self.model)
203
204 + # Upcast weights Transformers always keeps in fp32
205 + self.keep_in_fp32(self.model)
206 +
207 # Pipeline parallel intermediate tensors
208 self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
209 ["hidden_states"], self.text_config.hidden_size
@@ -457,6 +469,8 @@ class Base(
469
470 def register_fusion(fuser: BaseFuser, prefix: str):
471 """Register a fused layer's mappings just before it is built."""
472 + self.fusers[prefix] = fuser
473 +
474 orig_to_new_stacked = fuser.orig_to_new_stacked(prefix)
475 self.hf_to_vllm_mapper.orig_to_new_stacked.update(orig_to_new_stacked)
476
@@ -522,60 +536,105 @@ class Base(
536 """
537 Create `Attention` instances to inform KV cache allocation.
538 """
539 + mla_fusers = {}
540 + attention_instances = {}
541 text_config = self.text_config
542 + attn_cls = self._get_attn_cls()
543 +
544 + # kv_lora_rank indicates that this is an MLA model
545 + if getattr(text_config, "kv_lora_rank", None) is not None:
546 + mla_fusers = {
547 + extract_layer_index(prefix): (prefix, fuser)
548 + for prefix, fuser in self.fusers.items()
549 + if isinstance(fuser, MLAFuser)
550 + }
551 + if attn_cls is MLAAttention:
552 + text_config._attn_implementation = "vllm_mla"
553 + else:
554 + # MLA model not using MLAAttention: recompute head_size for full attn
555 + qk_nope_head_dim = getattr(text_config, "qk_nope_head_dim", 0)
556 + qk_rope_head_dim = getattr(text_config, "qk_rope_head_dim", 0)
557 + if qk_head_dim := qk_nope_head_dim + qk_rope_head_dim:
558 + self.model_config.model_arch_config.head_size = qk_head_dim
559
560 num_heads = self.model_config.get_num_attention_heads(self.parallel_config)
561 head_size = self.model_config.get_head_size()
562 num_kv_heads = self.model_config.get_num_kv_heads(self.parallel_config)
563 logits_soft_cap = getattr(text_config, "attn_logit_softcapping", None)
564
532 - # In encoder models, the attention layers will have `is_causal=False`
533 - is_encoder = lambda module: not getattr(module, "is_causal", True)
534 - has_encoder = lambda model: any(is_encoder(m) for m in model.modules())
535 - is_multimodal = lambda config: config != config.get_text_config()
536 - # vLLM does not support encoder-decoder models, so if any encoder layer is
537 - # found in a text only model, we assume the whole model is an encoder model
538 - if has_encoder(self.model) and not is_multimodal(self.config):
539 - self.check_version("5.0.0", "encoder models support")
540 - attn_type = AttentionType.ENCODER_ONLY
541 - else:
542 - attn_type = AttentionType.DECODER
543 -
565 pp_rank = self.pp_group.rank_in_group
566 pp_size = self.pp_group.world_size
567 start, end = get_pp_indices(text_config.num_hidden_layers, pp_rank, pp_size)
568
548 - attention_instances = {}
569 for i in range(start, end):
550 - # Handle interleaved sliding window attention
551 - per_layer_sliding_window = None
552 - if (
553 - hasattr(self.config, "layer_types")
554 - and self.config.layer_types[i] == "sliding_attention"
555 - ):
556 - per_layer_sliding_window = self.config.sliding_window
557 -
558 - attn_cls = (
559 - EncoderOnlyAttention
560 - if attn_type == AttentionType.ENCODER_ONLY
561 - else Attention
562 - )
563 - attention_instances[i] = attn_cls(
570 + kwargs = dict(
571 num_heads=num_heads,
565 - head_size=head_size,
572 # NOTE: We use Llama scale as default, if it's set by
573 # Transformers, it's updated in vllm_attention_forward
574 scale=head_size**-0.5,
569 - num_kv_heads=num_kv_heads,
575 cache_config=self.cache_config,
576 quant_config=self.quant_config,
572 - logits_soft_cap=logits_soft_cap,
573 - per_layer_sliding_window=per_layer_sliding_window,
577 prefix=f"{i}.attn",
575 - attn_type=attn_type,
578 )
579 +
580 + if attn_cls is MLAAttention:
581 + prefix, fuser = mla_fusers[i]
582 + mla_module = self.get_submodule(prefix)
583 + dims = get_mla_dims(self.model_config)
584 + kwargs.update(
585 + scale=mla_module.scaling,
586 + qk_nope_head_dim=dims.qk_nope_head_dim,
587 + qk_rope_head_dim=dims.qk_rope_head_dim,
588 + v_head_dim=dims.v_head_dim,
589 + q_lora_rank=dims.q_lora_rank,
590 + kv_lora_rank=dims.kv_lora_rank,
591 + kv_b_proj=mla_module.get_submodule(fuser.kv_b_proj_name),
592 + )
593 + else:
594 + kwargs.update(
595 + head_size=head_size,
596 + num_kv_heads=num_kv_heads,
597 + logits_soft_cap=logits_soft_cap,
598 + )
599 +
600 + # Handle interleaved sliding window attention
601 + if (
602 + hasattr(text_config, "layer_types")
603 + and text_config.layer_types[i] == "sliding_attention"
604 + ):
605 + kwargs["per_layer_sliding_window"] = text_config.sliding_window
606 +
607 + attn_instance = attn_cls(**kwargs)
608 + if attn_cls is MLAAttention:
609 + # Attach MLA attn_instance to mla_module so it appears in
610 + # model.named_modules() and runs its process_weights_after_loading
611 + mla_module._vllm_mla_attn = attn_instance
612 + attention_instances[i] = attn_instance
613 return attention_instances
614
615 + def _get_attn_cls(self) -> type[AttentionLayerBase]:
616 + """Return the `Attention` class to use for this model's layers."""
617 + # In encoder models, the attention layers will have `is_causal=False`
618 + is_encoder = lambda module: not getattr(module, "is_causal", True)
619 + has_encoder = lambda model: any(is_encoder(m) for m in model.modules())
620 + is_multimodal = lambda config: config != config.get_text_config()
621 + # vLLM does not support encoder-decoder models, so if any encoder layer is
622 + # found in a text only model, we assume the whole model is an encoder model
623 + if has_encoder(self.model) and not is_multimodal(self.config):
624 + self.check_version("5.0.0", "encoder models support")
625 + return EncoderOnlyAttention
626 + if self.model_config.use_mla:
627 + self.check_version("5.15.0.dev0", "optimized MLA support")
628 + if any(isinstance(fuser, MLAFuser) for fuser in self.fusers.values()):
629 + return MLAAttention
630 + logger.warning_once(
631 + "This model uses MLA but `MLAFuser` failed to match and/or fuse any "
632 + "MLA attention layers. Falling back to full attention with a padded "
633 + "`value` head dimension."
634 + )
635 + os.environ["VLLM_MLA_DISABLE"] = "1"
636 + return Attention
637 +
638 def init_parameters(self, module: nn.Module, dtype: torch.dtype | None = None):
639 """
640 If a `parameter` is on the `meta` device, then its parent
@@ -604,6 +663,18 @@ class Base(
663
664 _init_parameters(module)
665
666 + def keep_in_fp32(self, module: nn.Module):
667 + """Honor `_keep_in_fp32_modules_strict` as `from_pretrained` would."""
668 + if self.model_config.dtype not in (torch.float16, torch.bfloat16):
669 + return
670 + fragments = getattr(module, "_keep_in_fp32_modules_strict", None)
671 + if not fragments:
672 + return
673 + pattern = re.compile("|".join(re.escape(f) for f in fragments))
674 + for name, tensor in named_state(module):
675 + if not hasattr(tensor, "weight_loader") and pattern.search(name):
676 + tensor.data = tensor.data.to(torch.float32)
677 +
678 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
679 return self.model.get_input_embeddings()(input_ids)
680
vllm/model_executor/models/transformers/fuser.py
+5 -4
@@ -18,6 +18,7 @@ from vllm.logger import init_logger
18 from vllm.model_executor.models.transformers.fusers import (
19 BaseFuser,
20 GLUFuser,
21 + MLAFuser,
22 PackedQKVFuser,
23 QKVFuser,
24 RewriteFuser,
@@ -48,7 +49,7 @@ def get_fuser(module: nn.Module) -> BaseFuser | None:
49 return None
50 if (graph := trace(module)) is None:
51 return None
51 - for fuser_cls in (GLUFuser, QKVFuser, PackedQKVFuser, RMSNormFuser):
52 + for fuser_cls in (MLAFuser, GLUFuser, QKVFuser, PackedQKVFuser, RMSNormFuser):
53 if (fuser := fuser_cls.match(graph, module)) is not None:
54 if isinstance(fuser, RewriteFuser):
55 try:
@@ -74,14 +75,14 @@ def get_fuser(module: nn.Module) -> BaseFuser | None:
75
76
77 class Fusers(UserDict):
77 - """Mapping from module class to fuser, for all fusable classes in a model."""
78 + """Mapping from module class and shape to fuser, for all fusable modules."""
79
80 def __init__(self, model: nn.Module, vllm_config: "VllmConfig"):
81 self.vllm_config = vllm_config
81 - super().__init__({type(m): get_fuser(m) for m in model.modules()})
82 + super().__init__({key(m): get_fuser(m) for m in model.modules()})
83
84 def __getitem__(self, m: nn.Module) -> BaseFuser | None:
84 - fuser = self.data.get(type(m))
85 + fuser = self.data.get(key(m))
86 if fuser is not None and fuser.validate(m, self.vllm_config):
87 return fuser
88 return None
vllm/model_executor/models/transformers/fusers/__init__.py
+2
@@ -8,6 +8,7 @@ from vllm.model_executor.models.transformers.fusers.base import (
8 StackedFuser,
9 )
10 from vllm.model_executor.models.transformers.fusers.glu import GLUFuser
11 +from vllm.model_executor.models.transformers.fusers.mla import MLAFuser
12 from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser
13 from vllm.model_executor.models.transformers.fusers.packed_qkv import PackedQKVFuser
14 from vllm.model_executor.models.transformers.fusers.qkv import QKVFuser
@@ -18,6 +19,7 @@ __all__ = [
19 "RewriteFuser",
20 "StackedFuser",
21 "GLUFuser",
22 + "MLAFuser",
23 "MoEBlockFuser",
24 "PackedQKVFuser",
25 "QKVFuser",
vllm/model_executor/models/transformers/fusers/mla.py new
+325
@@ -0,0 +1,325 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""MLA fuser: adapt a Transformers MLA attention module for vLLM's MLA layer."""
4 +
5 +import ast
6 +import inspect
7 +import textwrap
8 +from dataclasses import dataclass
9 +from typing import TYPE_CHECKING, ClassVar
10 +
11 +from torch import fx, nn
12 +
13 +from vllm.logger import init_logger
14 +from vllm.model_executor.layers.linear import MergedColumnParallelLinear
15 +from vllm.model_executor.models.transformers.fusers.base import StackedFuser
16 +from vllm.model_executor.models.transformers.fusers.rms_norm import RMSNormFuser
17 +from vllm.model_executor.models.transformers.fx_utils import (
18 + compile_forward,
19 + downstream_linear,
20 + is_linear,
21 + recover_forward,
22 + replace_expr,
23 + returned_linear,
24 + single_self_call,
25 + trace,
26 + upstream_linear,
27 +)
28 +from vllm.model_executor.models.transformers.utils import (
29 + log_replacement,
30 + replace_linear_class,
31 +)
32 +from vllm.model_executor.models.utils import ShardId, maybe_prefix
33 +
34 +if TYPE_CHECKING:
35 + from vllm.config import VllmConfig
36 +
37 +logger = init_logger(__name__)
38 +
39 +# Temporaries the fused down-projection binds in the rewritten forward.
40 +_Q_A_TEMP = "__q_a_fused"
41 +_KV_A_TEMP = "__kv_a_fused"
42 +
43 +
44 +def _consumes_placeholder(node: fx.Node) -> bool:
45 + """Whether `node` is a linear applied directly to a `forward` input."""
46 + return (
47 + len(node.args) == 1
48 + and isinstance(node.args[0], fx.Node)
49 + and node.args[0].op == "placeholder"
50 + )
51 +
52 +
53 +def _norm_size(norm: nn.Module) -> int:
54 + weight = getattr(norm, "weight", None)
55 + return weight.numel() if weight is not None else -1
56 +
57 +
58 +def _is_rms_norm(module: nn.Module) -> bool:
59 + """Whether `module` computes an RMSNorm, verified by `RMSNormFuser`'s matcher."""
60 + graph = trace(module)
61 + return graph is not None and RMSNormFuser.match(graph, module) is not None
62 +
63 +
64 +def _top_level_index(funcdef: ast.FunctionDef, node: ast.AST) -> int:
65 + """Index in `funcdef.body` of the top-level statement containing `node`."""
66 + for index, stmt in enumerate(funcdef.body):
67 + if any(child is node for child in ast.walk(stmt)):
68 + return index
69 + raise ValueError("node is not in the function body")
70 +
71 +
72 +def _single_expand_call(
73 + funcdef: ast.FunctionDef, module: nn.Module, kv_b_proj_name: str
74 +) -> ast.Call:
75 + """The KV expansion call `self.<method>(kv_c_normed, k_pe)`: the unique
76 + two-argument call to a method of `module` with the expansion's signature.
77 +
78 + The expansion is `kv_c_normed, k_pe -> key, value`: two inputs, every
79 + `return` a pair, and it applies `kv_b_proj` -- the projection `MLAAttention`
80 + owns after the bypass, absorbed into the attention computation.
81 + """
82 +
83 + def is_expansion_method(name: str) -> bool:
84 + method = getattr(type(module), name, None)
85 + if method is None:
86 + return False
87 + code = getattr(inspect.unwrap(method), "__code__", None)
88 + if (
89 + code is None
90 + or code.co_argcount != 3 # self, kv_c_normed, k_pe
91 + or kv_b_proj_name not in code.co_names
92 + ):
93 + return False
94 + try:
95 + tree = ast.parse(textwrap.dedent(inspect.getsource(method)))
96 + except (OSError, TypeError):
97 + return False
98 + returns = [node for node in ast.walk(tree) if isinstance(node, ast.Return)]
99 + return bool(returns) and all(
100 + isinstance(ret.value, ast.Tuple) and len(ret.value.elts) == 2
101 + for ret in returns
102 + )
103 +
104 + calls = [
105 + node
106 + for node in ast.walk(funcdef)
107 + if isinstance(node, ast.Call)
108 + and isinstance(node.func, ast.Attribute)
109 + and isinstance(node.func.value, ast.Name)
110 + and node.func.value.id == "self"
111 + and len(node.args) == 2
112 + and not node.keywords
113 + and is_expansion_method(node.func.attr)
114 + ]
115 + if len(calls) != 1:
116 + raise ValueError("expected exactly one KV expansion method call")
117 + return calls[0]
118 +
119 +
120 +@dataclass
121 +class MLAFuser(StackedFuser):
122 + """Fuser for the MLA attention pattern."""
123 +
124 + q_proj_name: str | None
125 + q_a_proj_name: str | None
126 + q_a_layernorm_name: str | None
127 + q_b_proj_name: str | None
128 + kv_a_proj_name: str
129 + kv_a_layernorm_name: str
130 + kv_b_proj_name: str
131 + o_proj_name: str | None
132 + merged_name: ClassVar[str] = "fused_qkv_a_proj"
133 + merged_cls: ClassVar[str] = "MergedColumnParallelLinear"
134 +
135 + @property
136 + def has_q_lora(self) -> bool:
137 + return self.q_a_proj_name is not None
138 +
139 + def info(self, name: str) -> str:
140 + info_str = (
141 + f"Fused: {name} ({self.source_cls}) -> MLAAttention (attention interface)"
142 + )
143 + if self.has_q_lora:
144 + info_str += "; " + super().info(name).removeprefix("Fused: ")
145 + return info_str
146 +
147 + @property
148 + def shards(self) -> list[tuple[str, ShardId]]:
149 + """`q_a_proj` and `kv_a_proj_with_mqa` stack into one down-projection."""
150 + if self.has_q_lora:
151 + return [(self.q_a_proj_name, 0), (self.kv_a_proj_name, 1)]
152 + return []
153 +
154 + @property
155 + def packed_modules_mapping(self) -> dict[str, list[str]]:
156 + if self.has_q_lora:
157 + return super().packed_modules_mapping
158 + return {}
159 +
160 + @classmethod
161 + def match(cls, graph: fx.Graph, module: nn.Module) -> "MLAFuser | None":
162 + # Find all `rms_norm(linear(placeholder))` chains.
163 + chains = []
164 + for node in graph.nodes:
165 + if node.op != "call_module" or is_linear(node, module) or not node.args:
166 + continue
167 + source = upstream_linear(node.args[0], module)
168 + if source is None or not _consumes_placeholder(source):
169 + continue
170 + if not _is_rms_norm(module.get_submodule(node.target)):
171 + continue
172 + chains.append((source, node))
173 +
174 + # Tell the chains apart by width.
175 + def is_kv_a(chain) -> bool:
176 + source, rms_norm = chain
177 + source_mod = module.get_submodule(source.target)
178 + rms_norm_mod = module.get_submodule(rms_norm.target)
179 + return source_mod.out_features != _norm_size(rms_norm_mod)
180 +
181 + kv_a_chains = [chain for chain in chains if is_kv_a(chain)]
182 + q_a_chains = [chain for chain in chains if not is_kv_a(chain)]
183 + # Exactly one KV chain is MLA's signature. The query chain is optional.
184 + if len(kv_a_chains) != 1 or len(q_a_chains) > 1:
185 + return None
186 + kv_a_proj, kv_a_layernorm = kv_a_chains[0]
187 +
188 + # Linear children claimed for a role so far; the rest resolve by elimination.
189 + claimed_linears = {kv_a_proj.target}
190 + q_proj_name = q_a_proj = q_a_layernorm = q_b_proj = None
191 + if q_a_chains:
192 + # Find `q_b_proj(q_a_layernorm(q_a_proj(placeholder)))`.
193 + q_a_proj, q_a_layernorm = q_a_chains[0]
194 + q_b_proj = downstream_linear(q_a_layernorm, module)
195 + if q_b_proj is None:
196 + return None
197 + claimed_linears |= {q_a_proj.target, q_b_proj.target}
198 + else:
199 + # Find `q_proj(placeholder)`.
200 + placeholder_linears = {
201 + node.target
202 + for node in graph.nodes
203 + if is_linear(node, module) and _consumes_placeholder(node)
204 + }
205 + if len(q_proj_candidates := placeholder_linears - claimed_linears) != 1:
206 + return None
207 + q_proj_name = next(iter(q_proj_candidates))
208 + claimed_linears.add(q_proj_name)
209 +
210 + # Find `kv_b_proj(kv_a_layernorm(...))`.
211 + kv_b_proj = downstream_linear(kv_a_layernorm, module)
212 + if kv_b_proj is None or kv_b_proj.target in claimed_linears:
213 + return None
214 + claimed_linears.add(kv_b_proj.target)
215 +
216 + # Find `o_proj` if it is returned by the forward graph.
217 + o_proj_name = returned_linear(graph, module)
218 + if o_proj_name in claimed_linears:
219 + o_proj_name = None
220 +
221 + return cls(
222 + source_cls=type(module).__name__,
223 + q_proj_name=q_proj_name,
224 + q_a_proj_name=q_a_proj.target if q_a_proj else None,
225 + q_a_layernorm_name=q_a_layernorm.target if q_a_layernorm else None,
226 + q_b_proj_name=q_b_proj.target if q_b_proj else None,
227 + kv_a_proj_name=kv_a_proj.target,
228 + kv_a_layernorm_name=kv_a_layernorm.target,
229 + kv_b_proj_name=kv_b_proj.target,
230 + o_proj_name=o_proj_name,
231 + )
232 +
233 + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
234 + return vllm_config.model_config.use_mla
235 +
236 + def update_forward(self, module: nn.Module) -> None:
237 + """Merge `q_a_proj` and `kv_a_proj` into one fused down proj then split.
238 + Bypass the KV expansion method so the compressed latent reaches the `vllm_mla`
239 + attention interface unexpanded."""
240 + funcdef, fn = recover_forward(type(module))
241 + if self.has_q_lora:
242 + # q_a_proj is usually inside the `else` of `if self.q_lora_rank is None`.
243 + # The fused call is inserted at the top-level statement preceding both.
244 + q_call = single_self_call(funcdef, self.q_a_proj_name)
245 + kv_call = single_self_call(funcdef, self.kv_a_proj_name)
246 + if ast.dump(q_call.args[0]) != ast.dump(kv_call.args[0]):
247 + raise ValueError("down-projections read different inputs")
248 + names = {n.id for n in ast.walk(funcdef) if isinstance(n, ast.Name)}
249 + if names & {_Q_A_TEMP, _KV_A_TEMP}:
250 + raise ValueError("fused temporaries would shadow existing names")
251 +
252 + merged = f"self.{self.merged_name}"
253 + targets = f"{_Q_A_TEMP}, {_KV_A_TEMP}"
254 + sections = f"[s // {merged}.tp_size for s in {merged}.output_sizes]"
255 + source = f"{targets} = {merged}(__arg__).split({sections}, -1)"
256 + assign = ast.parse(source=source).body[0]
257 + placeholder = next(
258 + node
259 + for node in ast.walk(assign)
260 + if isinstance(node, ast.Name) and node.id == "__arg__"
261 + )
262 + replace_expr(assign, placeholder, q_call.args[0])
263 +
264 + index = min(_top_level_index(funcdef, call) for call in (q_call, kv_call))
265 + ast.copy_location(assign, funcdef.body[index])
266 + funcdef.body.insert(index, assign)
267 + replace_expr(funcdef, q_call, ast.Name(id=_Q_A_TEMP, ctx=ast.Load()))
268 + replace_expr(funcdef, kv_call, ast.Name(id=_KV_A_TEMP, ctx=ast.Load()))
269 +
270 + # Transformers expands the latent into full key/value in a dedicated method.
271 + # `MLAAttention` consumes the latent directly (absorbing `kv_b_proj`),
272 + # so replace the expansion call with its own arguments so `kv_c_normed, k_pe`
273 + # flow to the interface in place of `key, value`.
274 + expand_call = _single_expand_call(funcdef, module, self.kv_b_proj_name)
275 + replace_expr(
276 + funcdef, expand_call, ast.Tuple(elts=list(expand_call.args), ctx=ast.Load())
277 + )
278 + self.fused_forward = compile_forward(funcdef, fn)
279 +
280 + def update_attrs(self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"):
281 + quant_config = vllm_config.quant_config
282 +
283 + def replace_linear_by_name(name: str, style: str):
284 + linear = module.get_submodule(name)
285 + replacement = replace_linear_class(
286 + linear, style, quant_config, prefix=maybe_prefix(prefix, name)
287 + )
288 + setattr(module, name, replacement)
289 + log_replacement(maybe_prefix(prefix, name), linear, replacement)
290 +
291 + if self.has_q_lora:
292 + q_a = module.get_submodule(self.q_a_proj_name)
293 + kv_a = module.get_submodule(self.kv_a_proj_name)
294 + merged = MergedColumnParallelLinear(
295 + input_size=q_a.in_features,
296 + output_sizes=[q_a.out_features, kv_a.out_features],
297 + bias=q_a.bias is not None,
298 + quant_config=quant_config,
299 + prefix=maybe_prefix(prefix, self.merged_name),
300 + return_bias=False,
301 + disable_tp=True,
302 + )
303 + logger.debug(
304 + "%s: %s, %s: %s -> %s: %s",
305 + self.q_a_proj_name,
306 + q_a,
307 + self.kv_a_proj_name,
308 + kv_a,
309 + self.merged_name,
310 + merged,
311 + )
312 + setattr(module, self.merged_name, merged)
313 + # The rewritten forward calls the merged projection instead.
314 + delattr(module, self.q_a_proj_name)
315 + delattr(module, self.kv_a_proj_name)
316 + replace_linear_by_name(self.q_b_proj_name, "colwise")
317 + else:
318 + replace_linear_by_name(self.kv_a_proj_name, "replicate")
319 + replace_linear_by_name(self.q_proj_name, "colwise")
320 +
321 + replace_linear_by_name(self.kv_b_proj_name, "colwise")
322 + # MLAAttention calls kv_b_proj and expects vLLM's default return_bias=True
323 + module.get_submodule(self.kv_b_proj_name).return_bias = True
324 + if self.o_proj_name is not None:
325 + replace_linear_by_name(self.o_proj_name, "rowwise")
vllm/model_executor/models/transformers/fx_utils.py
+19
@@ -416,6 +416,25 @@ def upstream_linear(node: object, module: nn.Module) -> fx.Node | None:
416 return None
417
418
419 +def downstream_linear(node: fx.Node, module: nn.Module) -> fx.Node | None:
420 + """Nearest linear consuming `node`'s output, walking through casts/scalings.
421 +
422 + Never walks through a leaf call (e.g. an attention interface): what crosses
423 + it is consumed by the attention computation, not projected."""
424 + queue = list(node.users)
425 + seen: set[fx.Node] = set()
426 + while queue:
427 + current = queue.pop(0)
428 + if current in seen:
429 + continue
430 + seen.add(current)
431 + if is_linear(current, module):
432 + return current
433 + if current.op in ("call_function", "call_method") and not is_leaf_call(current):
434 + queue.extend(current.users)
435 + return None
436 +
437 +
438 def returned_linear(graph: fx.Graph, module: nn.Module) -> str | None:
439 """Name of the Linear producing the graph's (first) output value."""
440 value = output_value(graph)