| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass |
| 4 | from typing import TYPE_CHECKING |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.config import CacheConfig |
| 9 | from vllm.model_executor.custom_op import PluggableLayer |
| 10 | from vllm.model_executor.layers.attention import MLAAttention |
| 11 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 12 | from vllm.models.common.ops import fused_q_kv_rmsnorm |
| 13 | from vllm.platforms import current_platform |
| 14 | |
| 15 | if TYPE_CHECKING: |
| 16 | from vllm.v1.attention.backends.mla.index_group import ( |
| 17 | SparseMLAIndexGroupBuilder, |
| 18 | ) |
| 19 | |
| 20 | |
| 21 | @dataclass |
| 22 | class MLAModules: |
| 23 | """Modules used in MLA.""" |
| 24 | |
| 25 | kv_a_layernorm: torch.nn.Module |
| 26 | kv_b_proj: torch.nn.Module |
| 27 | rotary_emb: torch.nn.Module | None |
| 28 | o_proj: torch.nn.Module |
| 29 | fused_qkv_a_proj: torch.nn.Module | None |
| 30 | kv_a_proj_with_mqa: torch.nn.Module | None |
| 31 | q_a_layernorm: torch.nn.Module | None |
| 32 | q_b_proj: torch.nn.Module | None |
| 33 | q_proj: torch.nn.Module | None |
| 34 | indexer: torch.nn.Module | None |
| 35 | is_sparse: bool |
| 36 | topk_indices_buffer: torch.Tensor | None |
| 37 | indexer_rotary_emb: torch.nn.Module | None = None |
| 38 | g_proj: torch.nn.Module | None = None |
| 39 | index_group_builder: "SparseMLAIndexGroupBuilder | None" = None |
| 40 | |
| 41 | |
| 42 | # --8<-- [start:multi_head_latent_attention] |
| 43 | @PluggableLayer.register("multi_head_latent_attention") |
| 44 | class MultiHeadLatentAttentionWrapper(PluggableLayer): |
| 45 | """Pluggable MLA layer which allows OOT backends to add |
| 46 | custom implementations of the outer MLA layer (including rope & o_proj). |
| 47 | Note that currently oot platforms can still use CustomOp.register_oot to |
| 48 | replace MLA layer entirely, although we use PluggableLayer to register |
| 49 | this layer now. |
| 50 | |
| 51 | This class takes positions and hidden_states as input. |
| 52 | The input tensors can either contain prefill tokens or decode tokens. |
| 53 | The class does the following: |
| 54 | |
| 55 | 1. MLA Preprocess. |
| 56 | 2. Perform multi-head attention to prefill tokens and |
| 57 | multi-query attention to decode tokens separately. |
| 58 | 3. Return the output tensor. |
| 59 | """ |
| 60 | |
| 61 | # --8<-- [end:multi_head_latent_attention] |
| 62 | |
| 63 | def __init__( |
| 64 | self, |
| 65 | hidden_size: int, |
| 66 | num_heads: int, |
| 67 | scale: float, |
| 68 | qk_nope_head_dim: int, |
| 69 | qk_rope_head_dim: int, |
| 70 | v_head_dim: int, |
| 71 | q_lora_rank: int | None, |
| 72 | kv_lora_rank: int, |
| 73 | mla_modules: MLAModules, |
| 74 | cache_config: CacheConfig | None = None, |
| 75 | quant_config: QuantizationConfig | None = None, |
| 76 | prefix: str = "", |
| 77 | skip_topk: bool = False, |
| 78 | non_causal_multi_token_decode: bool = False, |
| 79 | allow_short_prefill_indexer_scoring_skip: bool = False, |
| 80 | fuse_qkv_rmsnorm: bool = False, |
| 81 | ) -> None: |
| 82 | super().__init__() |
| 83 | self.hidden_size = hidden_size |
| 84 | self.qk_nope_head_dim = qk_nope_head_dim |
| 85 | self.qk_rope_head_dim = qk_rope_head_dim |
| 86 | self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim |
| 87 | self.v_head_dim = v_head_dim |
| 88 | self.q_lora_rank = q_lora_rank |
| 89 | self.kv_lora_rank = kv_lora_rank |
| 90 | self.num_heads = num_heads |
| 91 | self.fused_qkv_a_proj = mla_modules.fused_qkv_a_proj |
| 92 | self.kv_a_proj_with_mqa = mla_modules.kv_a_proj_with_mqa |
| 93 | self.q_a_layernorm = mla_modules.q_a_layernorm |
| 94 | self.q_b_proj = mla_modules.q_b_proj |
| 95 | self.q_proj = mla_modules.q_proj |
| 96 | self.kv_a_layernorm = mla_modules.kv_a_layernorm |
| 97 | self.kv_b_proj = mla_modules.kv_b_proj |
| 98 | self.rotary_emb = mla_modules.rotary_emb |
| 99 | self.o_proj = mla_modules.o_proj |
| 100 | self.indexer = mla_modules.indexer |
| 101 | self.indexer_rope_emb = mla_modules.indexer_rotary_emb |
| 102 | self.is_sparse = mla_modules.is_sparse |
| 103 | self.g_proj = mla_modules.g_proj |
| 104 | |
| 105 | # Whether to skip top-k token selection computation in this layer. |
| 106 | # When True, the indexer will not be called, and the layer will reuse |
| 107 | # the topk_tokens buffer written by a previous layer in the same pass. |
| 108 | # Refer: https://arxiv.org/abs/2603.12201 for more details. |
| 109 | self.skip_topk = skip_topk |
| 110 | # When True, fuse the q_a and kv_a RMSNorms into a single kernel launch |
| 111 | # (MLA layers with q-LoRA). Opt-in; default False preserves other models. |
| 112 | self.fuse_qkv_rmsnorm = fuse_qkv_rmsnorm |
| 113 | # qrep is active when the query projection is a DCP-group-sharded layer |
| 114 | # that materializes the full group head set locally. |
| 115 | q_proj_layer = self.q_b_proj if self.q_lora_rank is not None else self.q_proj |
| 116 | self.dcp_q_replicate = getattr(q_proj_layer, "qrep_active", False) |
| 117 | if self.indexer is not None: |
| 118 | assert hasattr(self.indexer, "topk_tokens") |
| 119 | self.topk_tokens = self.indexer.topk_tokens |
| 120 | self.topk_indices_buffer = mla_modules.topk_indices_buffer |
| 121 | |
| 122 | self.mla_attn = MLAAttention( |
| 123 | num_heads=self.num_heads, |
| 124 | scale=scale, |
| 125 | qk_nope_head_dim=self.qk_nope_head_dim, |
| 126 | qk_rope_head_dim=self.qk_rope_head_dim, |
| 127 | v_head_dim=self.v_head_dim, |
| 128 | q_lora_rank=self.q_lora_rank, |
| 129 | kv_lora_rank=self.kv_lora_rank, |
| 130 | cache_config=cache_config, |
| 131 | quant_config=quant_config, |
| 132 | prefix=f"{prefix}.attn", |
| 133 | kv_b_proj=self.kv_b_proj, |
| 134 | dcp_q_replicate=self.dcp_q_replicate, |
| 135 | use_sparse=self.is_sparse, |
| 136 | indexer=self.indexer, |
| 137 | topk_indices_buffer=mla_modules.topk_indices_buffer, |
| 138 | index_group_builder=mla_modules.index_group_builder, |
| 139 | non_causal_multi_token_decode=non_causal_multi_token_decode, |
| 140 | ) |
| 141 | indexer_op = getattr(self.indexer, "indexer_op", None) |
| 142 | if indexer_op is not None and hasattr( |
| 143 | indexer_op, "dense_mha_metadata_layer_name" |
| 144 | ): |
| 145 | enable_short_prefill_scoring_skip = ( |
| 146 | allow_short_prefill_indexer_scoring_skip |
| 147 | and not self.skip_topk |
| 148 | and not getattr(indexer_op, "use_pcp", False) |
| 149 | and current_platform.is_cuda() |
| 150 | ) |
| 151 | # The indexer and main MLA use independent decode thresholds and |
| 152 | # may classify the same short extend differently. Bind the main |
| 153 | # MLA layer name so the eager indexer op can check whether the |
| 154 | # batch's top-k indices will be consumed. |
| 155 | # PCP is excluded because indexer cache/scoring ownership differs |
| 156 | # across ranks and the no-consumer invariant has not been |
| 157 | # established there. |
| 158 | indexer_op.dense_mha_metadata_layer_name = ( |
| 159 | self.mla_attn.layer_name if enable_short_prefill_scoring_skip else "" |
| 160 | ) |
| 161 | self.prefix = prefix |
| 162 | |
| 163 | def forward( |
| 164 | self, |
| 165 | positions: torch.Tensor, |
| 166 | hidden_states: torch.Tensor, |
| 167 | llama_4_scaling: torch.Tensor | None = None, |
| 168 | ) -> torch.Tensor: |
| 169 | q_c = None |
| 170 | kv_lora = None |
| 171 | |
| 172 | if self.q_lora_rank is not None: |
| 173 | assert self.fused_qkv_a_proj is not None, ( |
| 174 | "fused_qkv_a_proj is required when q_lora_rank is not None" |
| 175 | ) |
| 176 | assert self.q_a_layernorm is not None, ( |
| 177 | "q_a_layernorm is required when q_lora_rank is not None" |
| 178 | ) |
| 179 | assert self.q_b_proj is not None, ( |
| 180 | "q_b_proj is required when q_lora_rank is not None" |
| 181 | ) |
| 182 | |
| 183 | qkv_lora = self.fused_qkv_a_proj(hidden_states)[0] |
| 184 | q_c, kv_lora = qkv_lora.split( |
| 185 | [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], |
| 186 | dim=-1, |
| 187 | ) |
| 188 | q_proj_layer = self.q_b_proj |
| 189 | if not self.fuse_qkv_rmsnorm: |
| 190 | q_c = self.q_a_layernorm(q_c) |
| 191 | q_proj_input = q_c |
| 192 | else: |
| 193 | assert self.kv_a_proj_with_mqa is not None, ( |
| 194 | "kv_a_proj_with_mqa is required when q_lora_rank is None" |
| 195 | ) |
| 196 | assert self.q_proj is not None, ( |
| 197 | "q_proj is required when q_lora_rank is None" |
| 198 | ) |
| 199 | kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0] |
| 200 | q_proj_layer = self.q_proj |
| 201 | q_proj_input = hidden_states |
| 202 | |
| 203 | kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) |
| 204 | if self.fuse_qkv_rmsnorm and q_c is not None: |
| 205 | # Fuse q_a + kv_a RMSNorm into one kernel launch (q_c is still the |
| 206 | # pre-norm projection output here). |
| 207 | assert self.q_a_layernorm is not None |
| 208 | q_proj_input, kv_c_normed = fused_q_kv_rmsnorm( |
| 209 | q_c, |
| 210 | kv_c, |
| 211 | self.q_a_layernorm.weight.data, |
| 212 | self.kv_a_layernorm.weight.data, |
| 213 | self.q_a_layernorm.variance_epsilon, |
| 214 | ) |
| 215 | q_c = q_proj_input |
| 216 | else: |
| 217 | kv_c_normed = self.kv_a_layernorm(kv_c) |
| 218 | # Add head dim of 1 to k_pe |
| 219 | k_pe = k_pe.unsqueeze(1) |
| 220 | |
| 221 | q = q_proj_layer(q_proj_input)[0] |
| 222 | heads = self.num_heads |
| 223 | if self.dcp_q_replicate: |
| 224 | heads *= q_proj_layer.group_size |
| 225 | q = q.view(-1, heads, self.qk_head_dim) |
| 226 | |
| 227 | if self.rotary_emb is not None: |
| 228 | q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb( |
| 229 | positions, q[..., self.qk_nope_head_dim :], k_pe |
| 230 | ) |
| 231 | |
| 232 | if self.indexer and self.is_sparse and not self.skip_topk: |
| 233 | self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb) |
| 234 | if self.is_sparse: |
| 235 | self.mla_attn.impl.record_logical_topk_ready() # type: ignore[attr-defined] |
| 236 | |
| 237 | if llama_4_scaling is not None: |
| 238 | q *= llama_4_scaling |
| 239 | |
| 240 | q_dcp_replicated = None |
| 241 | if self.dcp_q_replicate: |
| 242 | q_dcp_replicated, q = q, q_proj_layer._local_view(q) |
| 243 | |
| 244 | attn_out = self.mla_attn( |
| 245 | q, |
| 246 | kv_c_normed, |
| 247 | k_pe, |
| 248 | output_shape=(hidden_states.shape[0], self.num_heads * self.v_head_dim), |
| 249 | q_dcp_replicated=q_dcp_replicated, |
| 250 | ) |
| 251 | |
| 252 | if self.g_proj is not None: |
| 253 | attn_out = attn_out * self.g_proj(hidden_states)[0].sigmoid() |
| 254 | |
| 255 | return self.o_proj(attn_out)[0] |