@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e614ab7806
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

Separate MLAAttention class from Attention (#25103)

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>

Naveenraj Kamalakannan committed Oct 8, 2025 at 20:11 UTC e614ab780645b06540b9f0427c26cb77ef34b558
10 files changed +502 -163
vllm/attention/backends/abstract.py
+26
@@ -6,6 +6,7 @@ from typing import Generic, Optional, Protocol, TypeVar
6
7 import torch
8
9 +from vllm.model_executor.layers.linear import ColumnParallelLinear
10 from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
11
12
@@ -184,6 +185,31 @@ class AttentionImpl(ABC, Generic[T]):
185
186
187 class MLAAttentionImpl(AttentionImpl[T], Generic[T]):
188 + @abstractmethod
189 + def __init__(
190 + self,
191 + num_heads: int,
192 + head_size: int,
193 + scale: float,
194 + num_kv_heads: int,
195 + alibi_slopes: Optional[list[float]],
196 + sliding_window: Optional[int],
197 + kv_cache_dtype: str,
198 + logits_soft_cap: Optional[float],
199 + attn_type: str,
200 + kv_sharing_target_layer_name: Optional[str],
201 + # MLA Specific Arguments
202 + q_lora_rank: Optional[int],
203 + kv_lora_rank: int,
204 + qk_nope_head_dim: int,
205 + qk_rope_head_dim: int,
206 + qk_head_dim: int,
207 + v_head_dim: int,
208 + kv_b_proj: ColumnParallelLinear,
209 + indexer: Optional[object] = None,
210 + ) -> None:
211 + raise NotImplementedError
212 +
213 @abstractmethod
214 def forward(
215 self,
vllm/attention/layer.py
+318 -22
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Attention layer."""
4
5 -from typing import Callable, Optional
5 +from typing import Callable, Optional, cast
6
7 import torch
8 import torch.nn as nn
@@ -10,7 +10,7 @@ import torch.nn.functional as F
10
11 import vllm.envs as envs
12 from vllm.attention import AttentionType
13 -from vllm.attention.backends.abstract import AttentionBackend
13 +from vllm.attention.backends.abstract import AttentionBackend, MLAAttentionImpl
14 from vllm.attention.backends.registry import _Backend, backend_name_to_enum
15 from vllm.attention.selector import get_attn_backend
16 from vllm.attention.utils.kv_sharing_utils import validate_kv_sharing_target
@@ -23,7 +23,10 @@ from vllm.distributed.kv_transfer import (
23 from vllm.forward_context import ForwardContext, get_forward_context
24 from vllm.logger import init_logger
25 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
26 -from vllm.model_executor.layers.linear import UnquantizedLinearMethod
26 +from vllm.model_executor.layers.linear import (
27 + ColumnParallelLinear,
28 + UnquantizedLinearMethod,
29 +)
30 from vllm.model_executor.layers.quantization import QuantizationConfig
31 from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
32 from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
@@ -131,8 +134,6 @@ class Attention(nn.Module, AttentionLayerBase):
134 quant_config: Optional[QuantizationConfig] = None,
135 logits_soft_cap: Optional[float] = None,
136 per_layer_sliding_window: Optional[int] = None,
134 - use_mla: bool = False,
135 - use_sparse: bool = False,
137 prefix: str = "",
138 attn_type: str = AttentionType.DECODER,
139 kv_sharing_target_layer_name: Optional[str] = None,
@@ -192,8 +193,6 @@ class Attention(nn.Module, AttentionLayerBase):
193 # the quant op after this attention layer.
194 self._o_scale_float: Optional[float] = None
195
195 - self.use_mla = use_mla
196 - self.use_sparse = use_sparse
196 self.num_heads = num_heads
197 self.head_size = head_size
198 self.num_kv_heads = num_kv_heads
@@ -229,9 +228,8 @@ class Attention(nn.Module, AttentionLayerBase):
228 dtype,
229 kv_cache_dtype,
230 block_size,
232 - use_mla=use_mla,
231 + use_mla=False,
232 has_sink=self.has_sink,
234 - use_sparse=use_sparse,
233 )
234 else:
235 self.attn_backend = attn_backend
@@ -349,19 +347,15 @@ class Attention(nn.Module, AttentionLayerBase):
347 output_shape = output_shape if output_shape is not None else query.shape
348 output = torch.zeros(output_shape, dtype=output_dtype, device=query.device)
349 hidden_size = output_shape[-1]
352 - # We skip reshaping query, key and value tensors for the MLA
353 - # backend since these tensors have different semantics and are
354 - # processed differently.
355 - if not self.use_mla:
356 - # Reshape the query, key, and value tensors.
357 - # NOTE(woosuk): We do this outside the custom op to minimize the
358 - # CPU overheads from the non-CUDA-graph regions.
359 - query = query.view(-1, self.num_heads, self.head_size)
360 - output = output.view(-1, self.num_heads, self.head_size)
361 - if key is not None:
362 - key = key.view(-1, self.num_kv_heads, self.head_size)
363 - if value is not None:
364 - value = value.view(-1, self.num_kv_heads, self.head_size)
350 + # Reshape the query, key, and value tensors.
351 + # NOTE(woosuk): We do this outside the custom op to minimize the
352 + # CPU overheads from the non-CUDA-graph regions.
353 + query = query.view(-1, self.num_heads, self.head_size)
354 + output = output.view(-1, self.num_heads, self.head_size)
355 + if key is not None:
356 + key = key.view(-1, self.num_kv_heads, self.head_size)
357 + if value is not None:
358 + value = value.view(-1, self.num_kv_heads, self.head_size)
359 if self.use_direct_call:
360 forward_context: ForwardContext = get_forward_context()
361 attn_metadata = forward_context.attn_metadata
@@ -570,6 +564,218 @@ class MultiHeadAttention(nn.Module):
564 return out.reshape(bsz, q_len, -1)
565
566
567 +class MLAAttention(nn.Module, AttentionLayerBase):
568 + """Multi-Head Latent Attention layer.
569 +
570 + This class takes query, and compressed key/value tensors as input.
571 + The class does the following:
572 +
573 + 1. Store the input key and value tensors in the KV cache.
574 + 2. Perform (multi-head/multi-query/grouped-query) attention.
575 + 3. Return the output tensor.
576 + """
577 +
578 + def __init__(
579 + self,
580 + num_heads: int,
581 + scale: float,
582 + qk_nope_head_dim: int,
583 + qk_rope_head_dim: int,
584 + v_head_dim: int,
585 + q_lora_rank: Optional[int],
586 + kv_lora_rank: int,
587 + kv_b_proj: ColumnParallelLinear,
588 + cache_config: Optional[CacheConfig] = None,
589 + quant_config: Optional[QuantizationConfig] = None,
590 + prefix: str = "",
591 + use_sparse: bool = False,
592 + indexer: Optional[object] = None,
593 + ):
594 + super().__init__()
595 + self.num_heads = num_heads
596 + self.scale = scale
597 + self.qk_nope_head_dim = qk_nope_head_dim
598 + self.qk_rope_head_dim = qk_rope_head_dim
599 + self.v_head_dim = v_head_dim
600 + self.q_lora_rank = q_lora_rank
601 + self.kv_lora_rank = kv_lora_rank
602 + self.head_size = kv_lora_rank + qk_rope_head_dim
603 + self.layer_name = prefix
604 +
605 + if cache_config is not None:
606 + kv_cache_dtype = cache_config.cache_dtype
607 + block_size = cache_config.block_size
608 + calculate_kv_scales = cache_config.calculate_kv_scales
609 + else:
610 + kv_cache_dtype = "auto"
611 + block_size = 16
612 + calculate_kv_scales = False
613 + self.kv_cache_dtype = kv_cache_dtype
614 +
615 + dtype = torch.get_default_dtype()
616 + self.attn_backend = get_attn_backend(
617 + self.head_size,
618 + dtype,
619 + kv_cache_dtype,
620 + block_size,
621 + use_mla=True,
622 + use_sparse=use_sparse,
623 + )
624 + impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls())
625 + self.impl = impl_cls(
626 + num_heads=self.num_heads,
627 + head_size=self.head_size,
628 + scale=self.scale,
629 + num_kv_heads=1,
630 + alibi_slopes=None,
631 + sliding_window=None,
632 + kv_cache_dtype=self.kv_cache_dtype,
633 + logits_soft_cap=None,
634 + attn_type=AttentionType.DECODER,
635 + kv_sharing_target_layer_name=None,
636 + # MLA Args
637 + q_lora_rank=self.q_lora_rank,
638 + kv_lora_rank=self.kv_lora_rank,
639 + qk_nope_head_dim=self.qk_nope_head_dim,
640 + qk_rope_head_dim=self.qk_rope_head_dim,
641 + qk_head_dim=self.qk_nope_head_dim + self.qk_rope_head_dim,
642 + v_head_dim=self.v_head_dim,
643 + kv_b_proj=kv_b_proj,
644 + indexer=indexer,
645 + )
646 +
647 + self.use_direct_call = not current_platform.opaque_attention_op()
648 +
649 + compilation_config = get_current_vllm_config().compilation_config
650 + if prefix in compilation_config.static_forward_context:
651 + raise ValueError(f"Duplicate layer name: {prefix}")
652 + compilation_config.static_forward_context[prefix] = self
653 +
654 + self.kv_cache = [
655 + torch.tensor([])
656 + for _ in range(
657 + get_current_vllm_config().parallel_config.pipeline_parallel_size
658 + )
659 + ]
660 +
661 + # Align with Attention's scale attributes for MLA backends.
662 +
663 + self.calculate_kv_scales = calculate_kv_scales
664 + self._k_scale = torch.tensor(1.0, dtype=torch.float32)
665 + self._v_scale = torch.tensor(1.0, dtype=torch.float32)
666 + self._q_scale = torch.tensor(1.0, dtype=torch.float32)
667 + self._prob_scale = torch.tensor(1.0, dtype=torch.float32)
668 +
669 + # Host-side mirrors used by some attention backends
670 + self._q_scale_float = 1.0
671 + self._k_scale_float = 1.0
672 + self._v_scale_float = 1.0
673 + self._o_scale_float: Optional[float] = None
674 +
675 + self.use_sparse = use_sparse
676 +
677 + # Initialize q/k/v range constants.
678 + try:
679 + self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)
680 + self.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)
681 + self.v_range = torch.tensor(envs.V_SCALE_CONSTANT, dtype=torch.float32)
682 + except torch.cuda.OutOfMemoryError:
683 + # Keep defaults if allocation fails; not critical for init.
684 + pass
685 +
686 + def forward(
687 + self,
688 + q: torch.Tensor,
689 + kv_c_normed: torch.Tensor,
690 + k_pe: torch.Tensor,
691 + output_shape: Optional[torch.Size] = None,
692 + ) -> torch.Tensor:
693 + if self.use_direct_call:
694 + forward_context: ForwardContext = get_forward_context()
695 + attn_metadata = forward_context.attn_metadata
696 + if isinstance(attn_metadata, dict):
697 + attn_metadata = attn_metadata[self.layer_name]
698 + self_kv_cache = self.kv_cache[forward_context.virtual_engine]
699 +
700 + # Mirror Attention.forward scale calculation path
701 + if self.calculate_kv_scales and getattr(
702 + attn_metadata, "enable_kv_scales_calculation", False
703 + ):
704 + self.calc_kv_scales(q, kv_c_normed, k_pe)
705 +
706 + if self.attn_backend.accept_output_buffer:
707 + output = torch.zeros(output_shape, dtype=q.dtype, device=q.device)
708 + self.impl.forward(
709 + self,
710 + q,
711 + kv_c_normed,
712 + k_pe,
713 + self_kv_cache,
714 + attn_metadata,
715 + output=output,
716 + )
717 + return output
718 + else:
719 + return self.impl.forward(
720 + self, q, kv_c_normed, k_pe, self_kv_cache, attn_metadata
721 + )
722 + else:
723 + if self.attn_backend.accept_output_buffer:
724 + output = torch.zeros(output_shape, dtype=q.dtype, device=q.device)
725 + torch.ops.vllm.unified_mla_attention_with_output(
726 + q,
727 + kv_c_normed,
728 + k_pe,
729 + output,
730 + self.layer_name,
731 + )
732 + return output
733 + else:
734 + # We can still access forward context to check calculation flag
735 + if self.calculate_kv_scales:
736 + forward_context = get_forward_context()
737 + attn_metadata = forward_context.attn_metadata
738 + if isinstance(attn_metadata, dict):
739 + attn_metadata = attn_metadata[self.layer_name]
740 + if getattr(attn_metadata, "enable_kv_scales_calculation", False):
741 + self.calc_kv_scales(q, kv_c_normed, k_pe)
742 + return torch.ops.vllm.unified_mla_attention(
743 + q,
744 + kv_c_normed,
745 + k_pe,
746 + self.layer_name,
747 + )
748 +
749 + def process_weights_after_loading(self, act_dtype: torch.dtype):
750 + if hasattr(self.impl, "process_weights_after_loading"):
751 + self.impl.process_weights_after_loading(act_dtype)
752 +
753 + def calc_kv_scales(
754 + self, q: torch.Tensor, kv_c_normed: torch.Tensor, k_pe: torch.Tensor
755 + ) -> None:
756 + """Optional scale calculation for MLA inputs.
757 +
758 + Mirrors Attention.calc_kv_scales. Not all MLA backends require this
759 + """
760 + # Use safe defaults if ranges are not present
761 + q_range = getattr(self, "q_range", torch.tensor(1.0))
762 + k_range = getattr(self, "k_range", torch.tensor(1.0))
763 + v_range = getattr(self, "v_range", torch.tensor(1.0))
764 +
765 + self._q_scale.copy_(torch.abs(q).max() / q_range)
766 + # kv_c_normed is the compressed KV representation; use it for k/v
767 + kv_abs_max = torch.abs(kv_c_normed).max()
768 + self._k_scale.copy_(kv_abs_max / k_range)
769 + self._v_scale.copy_(kv_abs_max / v_range)
770 + self._q_scale_float = self._q_scale.item()
771 + self._k_scale_float = self._k_scale.item()
772 + self._v_scale_float = self._v_scale.item()
773 + self.calculate_kv_scales = False
774 +
775 + def get_attn_backend(self) -> type[AttentionBackend]:
776 + return self.attn_backend
777 +
778 +
779 def wait_for_kv_layer_from_connector(layer_name: str):
780 if not has_kv_transfer_group() or not is_v1_kv_transfer_group():
781 return
@@ -726,3 +932,93 @@ direct_register_custom_op(
932 fake_impl=unified_attention_with_output_fake,
933 tags=tag_cudagraph_unsafe,
934 )
935 +
936 +
937 +def unified_mla_attention(
938 + q: torch.Tensor,
939 + kv_c_normed: torch.Tensor,
940 + k_pe: torch.Tensor,
941 + layer_name: str,
942 +) -> torch.Tensor:
943 + wait_for_kv_layer_from_connector(layer_name)
944 +
945 + forward_context: ForwardContext = get_forward_context()
946 + attn_metadata = forward_context.attn_metadata
947 + if isinstance(attn_metadata, dict):
948 + attn_metadata = attn_metadata[layer_name]
949 + self: MLAAttention = forward_context.no_compile_layers[layer_name]
950 + kv_cache = self.kv_cache[forward_context.virtual_engine]
951 + output = self.impl.forward(self, q, kv_c_normed, k_pe, kv_cache, attn_metadata)
952 +
953 + maybe_save_kv_layer_to_connector(layer_name, kv_cache)
954 + return output
955 +
956 +
957 +def unified_mla_attention_fake(
958 + q: torch.Tensor,
959 + kv_c_normed: torch.Tensor,
960 + k_pe: torch.Tensor,
961 + layer_name: str,
962 +) -> torch.Tensor:
963 + return torch.empty_like(q).contiguous()
964 +
965 +
966 +direct_register_custom_op(
967 + op_name="unified_mla_attention",
968 + op_func=unified_mla_attention,
969 + mutates_args=[],
970 + fake_impl=unified_mla_attention_fake,
971 + dispatch_key=current_platform.dispatch_key,
972 +)
973 +
974 +
975 +def unified_mla_attention_with_output(
976 + q: torch.Tensor,
977 + kv_c_normed: torch.Tensor,
978 + k_pe: torch.Tensor,
979 + output: torch.Tensor,
980 + layer_name: str,
981 + output_scale: Optional[torch.Tensor] = None,
982 + output_block_scale: Optional[torch.Tensor] = None,
983 +) -> None:
984 + wait_for_kv_layer_from_connector(layer_name)
985 + forward_context: ForwardContext = get_forward_context()
986 + attn_metadata = forward_context.attn_metadata
987 + if isinstance(attn_metadata, dict):
988 + attn_metadata = attn_metadata[layer_name]
989 + self: MLAAttention = forward_context.no_compile_layers[layer_name]
990 + kv_cache = self.kv_cache[forward_context.virtual_engine]
991 + self.impl.forward(
992 + self,
993 + q,
994 + kv_c_normed,
995 + k_pe,
996 + kv_cache,
997 + attn_metadata,
998 + output=output,
999 + output_scale=output_scale,
1000 + output_block_scale=output_block_scale,
1001 + )
1002 +
1003 + maybe_save_kv_layer_to_connector(layer_name, kv_cache)
1004 +
1005 +
1006 +def unified_mla_attention_with_output_fake(
1007 + q: torch.Tensor,
1008 + kv_c_normed: torch.Tensor,
1009 + k_pe: torch.Tensor,
1010 + output: torch.Tensor,
1011 + layer_name: str,
1012 + output_scale: Optional[torch.Tensor] = None,
1013 + output_block_scale: Optional[torch.Tensor] = None,
1014 +) -> None:
1015 + return
1016 +
1017 +
1018 +direct_register_custom_op(
1019 + op_name="unified_mla_attention_with_output",
1020 + op_func=unified_mla_attention_with_output,
1021 + mutates_args=["output", "output_block_scale"],
1022 + fake_impl=unified_mla_attention_with_output_fake,
1023 + dispatch_key=current_platform.dispatch_key,
1024 +)
vllm/config/compilation.py
+2
@@ -380,6 +380,8 @@ class CompilationConfig:
380 _attention_ops: ClassVar[list[str]] = [
381 "vllm.unified_attention",
382 "vllm.unified_attention_with_output",
383 + "vllm.unified_mla_attention",
384 + "vllm.unified_mla_attention_with_output",
385 "vllm.mamba_mixer2",
386 "vllm.mamba_mixer",
387 "vllm.short_conv",
vllm/model_executor/layers/mla.py
+11 -21
@@ -5,7 +5,7 @@ from typing import Optional
5
6 import torch
7
8 -from vllm.attention import Attention
8 +from vllm.attention.layer import MLAAttention
9 from vllm.config import CacheConfig
10 from vllm.model_executor.custom_op import CustomOp
11 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -30,8 +30,9 @@ class MLAModules:
30
31
32 @CustomOp.register("multi_head_latent_attention")
33 -class MultiHeadLatentAttention(CustomOp):
34 - """MLA layer registered as CustomOp.
33 +class MultiHeadLatentAttentionWrapper(CustomOp):
34 + """MLA layer registered as CustomOp to allow OOT backends to add
35 + custom implementations of the outer MLA layer (including rope & o_proj).
36 Note that currently MLA ignores the enable/disable mechanism of CustomOp
37 because there is only one in-tree implementation in forward_native.
38 TODO: implement this with a new PluggableLayer mechanism.
@@ -87,30 +88,19 @@ class MultiHeadLatentAttention(CustomOp):
88 self.topk_tokens = self.indexer.topk_tokens
89 self.topk_indices_buffer = mla_modules.topk_indices_buffer
90
90 - # In the MLA backend, kv_cache includes both k_c and
91 - # pe (i.e. decoupled position embeddings). In particular,
92 - # the concat_and_cache_mla op requires
93 - # k_c.size(1) + k_pe.size(1) == kv_cache.size(2)
94 - # i.e.
95 - # kv_lora_rank + qk_rope_head_dim == head_size
96 - self.mla_attn = Attention(
91 + self.mla_attn = MLAAttention(
92 num_heads=self.num_heads,
98 - head_size=self.kv_lora_rank + self.qk_rope_head_dim,
93 scale=scale,
100 - num_kv_heads=1,
101 - cache_config=cache_config,
102 - quant_config=quant_config,
103 - prefix=f"{prefix}.attn",
104 - use_mla=True,
105 - use_sparse=mla_modules.is_sparse,
106 - # MLA Args
107 - q_lora_rank=self.q_lora_rank,
108 - kv_lora_rank=self.kv_lora_rank,
94 qk_nope_head_dim=self.qk_nope_head_dim,
95 qk_rope_head_dim=self.qk_rope_head_dim,
111 - qk_head_dim=self.qk_head_dim,
96 v_head_dim=self.v_head_dim,
97 + q_lora_rank=self.q_lora_rank,
98 + kv_lora_rank=self.kv_lora_rank,
99 + cache_config=cache_config,
100 + quant_config=quant_config,
101 + prefix=f"{prefix}.attn",
102 kv_b_proj=self.kv_b_proj,
103 + use_sparse=self.is_sparse,
104 indexer=self.indexer,
105 )
106
vllm/model_executor/model_loader/utils.py
+4 -4
@@ -14,6 +14,7 @@ from torch import nn
14 from typing_extensions import assert_never
15
16 from vllm.attention import Attention
17 +from vllm.attention.layer import MLAAttention
18 from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
19 from vllm.logger import init_logger
20 from vllm.model_executor.layers.linear import QKVCrossParallelLinear
@@ -122,11 +123,10 @@ def process_weights_after_loading(
123 with device_loading_context(module, target_device):
124 quant_method.process_weights_after_loading(module)
125
125 - # Currently only used by MLA.
126 - # NOTE: This intentionally happens after other modules so we can easily
127 - # decompress the weights for MLA.
126 + # Initialize post-load attention weights for both Attention and MLA.
127 + # NOTE: Happens after other modules so we can easily decompress weights.
128 for _, module in model.named_modules():
129 - if isinstance(module, Attention) and hasattr(
129 + if isinstance(module, (Attention, MLAAttention)) and hasattr(
130 module, "process_weights_after_loading"
131 ):
132 # TODO(lucas): see if there is a way to unify the signatures
vllm/model_executor/models/deepseek_v2.py
+2 -2
@@ -58,7 +58,7 @@ from vllm.model_executor.layers.linear import (
58 RowParallelLinear,
59 )
60 from vllm.model_executor.layers.logits_processor import LogitsProcessor
61 -from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttention
61 +from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper
62 from vllm.model_executor.layers.quantization import QuantizationConfig
63 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
64 per_token_group_quant_fp8,
@@ -1038,7 +1038,7 @@ class DeepseekV2MLAAttention(nn.Module):
1038 topk_indices_buffer=topk_indices_buffer,
1039 )
1040
1041 - self.mla_attn = MultiHeadLatentAttention(
1041 + self.mla_attn = MultiHeadLatentAttentionWrapper(
1042 self.hidden_size,
1043 self.num_local_heads,
1044 self.scaling,
vllm/v1/attention/backends/utils.py
+2 -2
@@ -32,11 +32,11 @@ if TYPE_CHECKING:
32
33 import vllm.envs as envs
34 from vllm.attention.backends.abstract import AttentionBackend, AttentionMetadata
35 -from vllm.attention.layer import Attention
35 from vllm.distributed.kv_transfer.kv_connector.utils import (
36 get_kv_connector_cache_layout,
37 )
38 from vllm.logger import init_logger
39 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
40 from vllm.v1.kv_cache_interface import AttentionSpec
41 from vllm.v1.worker.ubatch_utils import UBatchSlice
42
@@ -408,7 +408,7 @@ def get_per_layer_parameters(
408 to use during `plan`.
409 """
410
411 - layers = get_layers_from_vllm_config(vllm_config, Attention, layer_names)
411 + layers = get_layers_from_vllm_config(vllm_config, AttentionLayerBase, layer_names)
412 per_layer_params: dict[str, PerLayerParameters] = {}
413
414 for key, layer in layers.items():
vllm/v1/spec_decode/eagle.py
+3 -3
@@ -9,11 +9,11 @@ import numpy as np
9 import torch
10 import torch.nn as nn
11
12 -from vllm.attention.layer import Attention
12 from vllm.config import CompilationLevel, VllmConfig, get_layers_from_vllm_config
13 from vllm.distributed.parallel_state import get_pp_group
14 from vllm.forward_context import set_forward_context
15 from vllm.logger import init_logger
16 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
17 from vllm.model_executor.model_loader import get_model
18 from vllm.model_executor.models import supports_multimodal
19 from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache
@@ -880,7 +880,7 @@ class EagleProposer:
880 def load_model(self, target_model: nn.Module) -> None:
881 draft_model_config = self.vllm_config.speculative_config.draft_model_config
882 target_attn_layer_names = set(
883 - get_layers_from_vllm_config(self.vllm_config, Attention).keys()
883 + get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase).keys()
884 )
885 # FIXME: support hybrid kv for draft model
886 target_indexer_layer_names = set(
@@ -897,7 +897,7 @@ class EagleProposer:
897 )
898
899 draft_attn_layer_names = (
900 - get_layers_from_vllm_config(self.vllm_config, Attention).keys()
900 + get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase).keys()
901 - target_attn_layer_names
902 )
903 indexer_layers = get_layers_from_vllm_config(
vllm/v1/worker/gpu_model_runner.py
+72 -69
@@ -20,6 +20,7 @@ from typing_extensions import TypeAlias
20 import vllm.envs as envs
21 from vllm.attention import Attention, AttentionType
22 from vllm.attention.backends.abstract import AttentionBackend
23 +from vllm.attention.layer import MLAAttention
24 from vllm.attention.layers.chunked_local_attention import ChunkedLocalAttention
25 from vllm.compilation.counter import compilation_counter
26 from vllm.compilation.cuda_graph import CUDAGraphWrapper
@@ -4388,98 +4389,100 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4389 use_mla = self.vllm_config.model_config.use_mla
4390 cache_dtype_str = self.vllm_config.cache_config.cache_dtype
4391 kv_cache_spec: dict[str, KVCacheSpec] = {}
4391 - attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention)
4392 + attn_layers = get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase)
4393 for layer_name, attn_module in attn_layers.items():
4393 - if (kv_tgt_layer := attn_module.kv_sharing_target_layer_name) is not None:
4394 - # The layer doesn't need its own KV cache and will use that of
4395 - # the target layer. We skip creating a KVCacheSpec for it, so
4396 - # that KV cache management logic will act as this layer does
4397 - # not exist, and doesn't allocate KV cache for the layer. This
4398 - # enables the memory saving of cross-layer kv sharing, allowing
4399 - # a given amount of memory to accommodate longer context lengths
4400 - # or enable more requests to be processed simultaneously.
4401 - self.shared_kv_cache_layers[layer_name] = kv_tgt_layer
4402 - continue
4394 + if isinstance(attn_module, Attention):
4395 + if (
4396 + kv_tgt_layer := attn_module.kv_sharing_target_layer_name
4397 + ) is not None:
4398 + # The layer doesn't need its own KV cache and will use that of
4399 + # the target layer. We skip creating a KVCacheSpec for it, so
4400 + # that KV cache management logic will act as this layer does
4401 + # not exist, and doesn't allocate KV cache for the layer. This
4402 + # enables the memory saving of cross-layer kv sharing, allowing
4403 + # a given amount of memory to accommodate longer context lengths
4404 + # or enable more requests to be processed simultaneously.
4405 + self.shared_kv_cache_layers[layer_name] = kv_tgt_layer
4406 + continue
4407
4404 - # TODO(lucas): move the attention specs into the model layers like
4405 - # the attention backends
4406 - if attn_module.attn_type == AttentionType.DECODER:
4407 - if attn_module.sliding_window is not None:
4408 - assert not use_mla, "MLA is not supported for slidingwindow"
4409 - kv_cache_spec[layer_name] = SlidingWindowSpec(
4410 - block_size=block_size,
4411 - num_kv_heads=attn_module.num_kv_heads,
4412 - head_size=attn_module.head_size,
4413 - dtype=self.kv_cache_dtype,
4414 - sliding_window=attn_module.sliding_window,
4415 - )
4416 - elif use_mla:
4417 - kv_cache_spec[layer_name] = MLAAttentionSpec(
4408 + # TODO(lucas): move the attention specs into the model layers like
4409 + # the attention backends
4410 + if attn_module.attn_type == AttentionType.DECODER:
4411 + if attn_module.sliding_window is not None:
4412 + assert not use_mla, "MLA is not supported for slidingwindow"
4413 + kv_cache_spec[layer_name] = SlidingWindowSpec(
4414 + block_size=block_size,
4415 + num_kv_heads=attn_module.num_kv_heads,
4416 + head_size=attn_module.head_size,
4417 + dtype=self.kv_cache_dtype,
4418 + sliding_window=attn_module.sliding_window,
4419 + )
4420 + elif self.attention_chunk_size is not None and isinstance(
4421 + attn_module, ChunkedLocalAttention
4422 + ):
4423 + kv_cache_spec[layer_name] = ChunkedLocalAttentionSpec(
4424 + block_size=block_size,
4425 + num_kv_heads=attn_module.num_kv_heads,
4426 + head_size=attn_module.head_size,
4427 + dtype=self.kv_cache_dtype,
4428 + attention_chunk_size=self.attention_chunk_size,
4429 + )
4430 + else:
4431 + kv_cache_spec[layer_name] = FullAttentionSpec(
4432 + block_size=block_size,
4433 + num_kv_heads=attn_module.num_kv_heads,
4434 + head_size=attn_module.head_size,
4435 + dtype=self.kv_cache_dtype,
4436 + )
4437 + elif attn_module.attn_type == AttentionType.ENCODER_DECODER:
4438 + kv_cache_spec[layer_name] = CrossAttentionSpec(
4439 block_size=block_size,
4440 num_kv_heads=attn_module.num_kv_heads,
4441 head_size=attn_module.head_size,
4442 dtype=self.kv_cache_dtype,
4422 - cache_dtype_str=cache_dtype_str,
4443 )
4424 - elif self.attention_chunk_size is not None and isinstance(
4425 - attn_module, ChunkedLocalAttention
4444 + elif attn_module.attn_type in (
4445 + AttentionType.ENCODER,
4446 + AttentionType.ENCODER_ONLY,
4447 ):
4427 - kv_cache_spec[layer_name] = ChunkedLocalAttentionSpec(
4428 - block_size=block_size,
4429 - num_kv_heads=attn_module.num_kv_heads,
4430 - head_size=attn_module.head_size,
4431 - dtype=self.kv_cache_dtype,
4432 - attention_chunk_size=self.attention_chunk_size,
4433 - )
4448 + # encoder-only attention does not need KV cache.
4449 + continue
4450 else:
4435 - kv_cache_spec[layer_name] = FullAttentionSpec(
4436 - block_size=block_size,
4437 - num_kv_heads=attn_module.num_kv_heads,
4438 - head_size=attn_module.head_size,
4439 - dtype=self.kv_cache_dtype,
4440 - )
4441 - elif attn_module.attn_type == AttentionType.ENCODER_DECODER:
4442 - kv_cache_spec[layer_name] = CrossAttentionSpec(
4451 + raise ValueError(f"Unknown attention type: {attn_module.attn_type}")
4452 +
4453 + elif isinstance(attn_module, MLAAttention):
4454 + kv_cache_spec[layer_name] = MLAAttentionSpec(
4455 block_size=block_size,
4444 - num_kv_heads=attn_module.num_kv_heads,
4456 + num_kv_heads=1,
4457 head_size=attn_module.head_size,
4458 dtype=self.kv_cache_dtype,
4459 + cache_dtype_str=cache_dtype_str,
4460 )
4448 - elif attn_module.attn_type in (
4449 - AttentionType.ENCODER,
4450 - AttentionType.ENCODER_ONLY,
4451 - ):
4452 - # encoder-only attention does not need KV cache.
4453 - continue
4454 - else:
4455 - raise ValueError(f"Unknown attention type: {attn_module.attn_type}")
4456 -
4457 - mamba_layers = get_layers_from_vllm_config(self.vllm_config, MambaBase)
4458 - if len(mamba_layers) > 0:
4459 - if (
4460 - self.vllm_config.speculative_config is not None
4461 - and self.vllm_config.model_config.hf_config.model_type
4462 - not in ["qwen3_next"]
4463 - ):
4464 - raise NotImplementedError(
4465 - "Mamba with speculative decoding is not supported yet."
4466 - )
4467 - mamba_block_size = self.vllm_config.cache_config.mamba_block_size
4468 - page_size_padded = self.vllm_config.cache_config.mamba_page_size_padded
4461
4470 - for layer_name, mamba_module in mamba_layers.items():
4462 + elif isinstance(attn_module, MambaBase):
4463 + if (
4464 + self.vllm_config.speculative_config is not None
4465 + and self.vllm_config.model_config.hf_config.model_type
4466 + not in ["qwen3_next"]
4467 + ):
4468 + raise NotImplementedError(
4469 + "Mamba with speculative decoding is not supported yet."
4470 + )
4471 + mamba_block_size = self.vllm_config.cache_config.mamba_block_size
4472 + page_size_padded = self.vllm_config.cache_config.mamba_page_size_padded
4473 kv_cache_spec[layer_name] = MambaSpec(
4472 - shapes=mamba_module.get_state_shape(),
4473 - dtypes=mamba_module.get_state_dtype(),
4474 + shapes=attn_module.get_state_shape(),
4475 + dtypes=attn_module.get_state_dtype(),
4476 block_size=mamba_block_size,
4477 page_size_padded=page_size_padded,
4476 - mamba_type=mamba_module.mamba_type,
4478 + mamba_type=attn_module.mamba_type,
4479 num_speculative_blocks=(
4480 self.speculative_config.num_speculative_tokens
4481 if self.speculative_config
4482 else 0
4483 ),
4484 )
4485 +
4486 ds_indexer_layers = get_layers_from_vllm_config(
4487 self.vllm_config, DeepseekV32IndexerCache
4488 )
vllm/v1/worker/tpu_model_runner.py
+62 -40
@@ -19,6 +19,7 @@ import torch_xla.runtime as xr
19 import vllm.envs as envs
20 from vllm.attention import Attention
21 from vllm.attention.backends.abstract import AttentionType
22 +from vllm.attention.layer import MLAAttention
23 from vllm.attention.layers.chunked_local_attention import ChunkedLocalAttention
24 from vllm.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
25 from vllm.config import (
@@ -32,6 +33,7 @@ from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks
33 from vllm.forward_context import set_forward_context
34 from vllm.logger import init_logger
35 from vllm.lora.layers import BaseLayerWithLoRA
36 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
37 from vllm.model_executor.model_loader import get_model_loader
38 from vllm.model_executor.model_loader.tpu import TPUModelLoader
39 from vllm.model_executor.models.interfaces import (
@@ -63,6 +65,7 @@ from vllm.v1.kv_cache_interface import (
65 FullAttentionSpec,
66 KVCacheConfig,
67 KVCacheSpec,
68 + MLAAttentionSpec,
69 SlidingWindowSpec,
70 )
71 from vllm.v1.outputs import (
@@ -561,52 +564,71 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
564 format. Layers that do not need KV cache are not included.
565 """
566
564 - layers = get_layers_from_vllm_config(self.vllm_config, Attention)
567 + layers = get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase)
568 block_size = self.vllm_config.cache_config.block_size
569 + cache_dtype_str = self.vllm_config.cache_config.cache_dtype
570 +
571 kv_cache_spec: dict[str, KVCacheSpec] = {}
572 for layer_name, attn_module in layers.items():
568 - if (kv_tgt_layer := attn_module.kv_sharing_target_layer_name) is not None:
569 - # The layer doesn't need its own KV cache and will use that of
570 - # the target layer. We skip creating a KVCacheSpec for it, so
571 - # that KV cache management logic will act as this layer does
572 - # not exist, and doesn't allocate KV cache for the layer. This
573 - # enables the memory saving of cross-layer kv sharing, allowing
574 - # a given amount of memory to accommodate longer context lengths
575 - # or enable more requests to be processed simultaneously.
576 - self.shared_kv_cache_layers[layer_name] = kv_tgt_layer
577 - continue
573 + # Classic Attention path
574 + if isinstance(attn_module, Attention):
575 + if (
576 + kv_tgt_layer := attn_module.kv_sharing_target_layer_name
577 + ) is not None:
578 + # The layer doesn't need its own KV cache and will use that of
579 + # the target layer. We skip creating a KVCacheSpec for it, so
580 + # that KV cache management logic will act as this layer does
581 + # not exist, and doesn't allocate KV cache for the layer. This
582 + # enables the memory saving of cross-layer kv sharing, allowing
583 + # a given amount of memory to accommodate longer context lengths
584 + # or enable more requests to be processed simultaneously.
585 + self.shared_kv_cache_layers[layer_name] = kv_tgt_layer
586 + continue
587
579 - if attn_module.attn_type == AttentionType.DECODER:
580 - if isinstance(attn_module, ChunkedLocalAttention):
581 - logger.warning_once(
582 - "Using irope in Pallas is not supported yet, it "
583 - "will fall back to global attention for long context."
584 - )
585 - if attn_module.sliding_window is not None:
586 - kv_cache_spec[layer_name] = SlidingWindowSpec(
587 - block_size=block_size,
588 - num_kv_heads=attn_module.num_kv_heads,
589 - head_size=attn_module.head_size,
590 - dtype=self.kv_cache_dtype,
591 - sliding_window=attn_module.sliding_window,
592 - )
588 + if attn_module.attn_type == AttentionType.DECODER:
589 + if isinstance(attn_module, ChunkedLocalAttention):
590 + logger.warning_once(
591 + "Using irope in Pallas is not supported yet, it "
592 + "will fall back to global attention for long context."
593 + )
594 + if attn_module.sliding_window is not None:
595 + kv_cache_spec[layer_name] = SlidingWindowSpec(
596 + block_size=block_size,
597 + num_kv_heads=attn_module.num_kv_heads,
598 + head_size=attn_module.head_size,
599 + dtype=self.kv_cache_dtype,
600 + sliding_window=attn_module.sliding_window,
601 + )
602 + else:
603 + kv_cache_spec[layer_name] = FullAttentionSpec(
604 + block_size=block_size,
605 + num_kv_heads=attn_module.num_kv_heads,
606 + head_size=attn_module.head_size,
607 + dtype=self.kv_cache_dtype,
608 + )
609 + elif attn_module.attn_type in (
610 + AttentionType.ENCODER,
611 + AttentionType.ENCODER_ONLY,
612 + ):
613 + # encoder-only attention does not need KV cache.
614 + continue
615 + elif attn_module.attn_type == AttentionType.ENCODER_DECODER:
616 + raise NotImplementedError
617 else:
594 - kv_cache_spec[layer_name] = FullAttentionSpec(
595 - block_size=block_size,
596 - num_kv_heads=attn_module.num_kv_heads,
597 - head_size=attn_module.head_size,
598 - dtype=self.kv_cache_dtype,
599 - )
600 - elif attn_module.attn_type in (
601 - AttentionType.ENCODER,
602 - AttentionType.ENCODER_ONLY,
603 - ):
604 - # encoder-only attention does not need KV cache.
605 - continue
606 - elif attn_module.attn_type == AttentionType.ENCODER_DECODER:
607 - raise NotImplementedError
618 + raise ValueError(f"Unknown attention type: {attn_module.attn_type}")
619 + # MLAAttention path
620 + elif isinstance(attn_module, MLAAttention):
621 + if layer_name in kv_cache_spec:
622 + continue
623 + kv_cache_spec[layer_name] = MLAAttentionSpec(
624 + block_size=block_size,
625 + num_kv_heads=1,
626 + head_size=attn_module.head_size,
627 + dtype=self.kv_cache_dtype,
628 + cache_dtype_str=cache_dtype_str,
629 + )
630 else:
609 - raise ValueError(f"Unknown attention type: {attn_module.attn_type}")
631 + continue
632
633 return kv_cache_spec
634