diff --git a/tests/models/transformers/fusers/test_mla.py b/tests/models/transformers/fusers/test_mla.py new file mode 100644 index 0000000000..2358b5c59a --- /dev/null +++ b/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) diff --git a/tests/models/transformers/test_backend.py b/tests/models/transformers/test_backend.py index a3eea1783f..1945fd706b 100644 --- a/tests/models/transformers/test_backend.py +++ b/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], diff --git a/vllm/config/model.py b/vllm/config/model.py index 91eb3c9bce..e6c3dd944d 100644 --- a/vllm/config/model.py +++ b/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: diff --git a/vllm/model_executor/models/transformers/__init__.py b/vllm/model_executor/models/transformers/__init__.py index ff4bb6cd43..7a7aab7266 100644 --- a/vllm/model_executor/models/transformers/__init__.py +++ b/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 diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index 35e5c9d1a0..b54cb5aa25 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/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) diff --git a/vllm/model_executor/models/transformers/fuser.py b/vllm/model_executor/models/transformers/fuser.py index cbcb24e2a9..52a08ac90b 100644 --- a/vllm/model_executor/models/transformers/fuser.py +++ b/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 diff --git a/vllm/model_executor/models/transformers/fusers/__init__.py b/vllm/model_executor/models/transformers/fusers/__init__.py index ed4da5155c..a0772ed4a2 100644 --- a/vllm/model_executor/models/transformers/fusers/__init__.py +++ b/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", diff --git a/vllm/model_executor/models/transformers/fusers/mla.py b/vllm/model_executor/models/transformers/fusers/mla.py new file mode 100644 index 0000000000..8e17bff490 --- /dev/null +++ b/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.(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") diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py index d665bcfe8c..089f5870cc 100644 --- a/vllm/model_executor/models/transformers/fx_utils.py +++ b/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)