| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import TYPE_CHECKING, ClassVar, cast |
| 4 | |
| 5 | import torch |
| 6 | import torch.nn as nn |
| 7 | from transformers import DeepseekV2Config, DeepseekV3Config |
| 8 | |
| 9 | from vllm.compilation.breakable_cudagraph import eager_break_during_capture |
| 10 | from vllm.config import CacheConfig, CUDAGraphMode, VllmConfig |
| 11 | from vllm.distributed import get_tensor_model_parallel_world_size |
| 12 | from vllm.distributed.parallel_state import get_tp_group |
| 13 | from vllm.forward_context import get_forward_context |
| 14 | from vllm.model_executor.layers.attention import MLAAttention |
| 15 | from vllm.model_executor.layers.attention.attention import get_attention_context |
| 16 | from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm |
| 17 | from vllm.model_executor.layers.linear import ( |
| 18 | ColumnParallelLinear, |
| 19 | MergedColumnParallelLinear, |
| 20 | ReplicatedLinear, |
| 21 | RowParallelLinear, |
| 22 | ) |
| 23 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 24 | from vllm.model_executor.layers.rotary_embedding import get_rope |
| 25 | from vllm.model_executor.layers.sparse_attn_indexer import ( |
| 26 | SparseAttnIndexer, |
| 27 | sparse_attn_indexer, |
| 28 | ) |
| 29 | from vllm.model_executor.models.deepseek_v2 import ( |
| 30 | DeepSeekV2FusedQkvAProjLinear, |
| 31 | DeepseekV32IndexerCache, |
| 32 | yarn_get_mscale, |
| 33 | ) |
| 34 | from vllm.model_executor.models.utils import extract_layer_index |
| 35 | from vllm.models.deepseek_v32.common.kernels import fused_norm_rope, fused_q |
| 36 | from vllm.platforms import current_platform |
| 37 | from vllm.utils.torch_utils import is_quantized_kv_cache |
| 38 | from vllm.v1.attention.ops.pcp import ( |
| 39 | finalize_mla_pcp_decode, |
| 40 | maybe_gather_mla_latent_cache_inputs, |
| 41 | ) |
| 42 | |
| 43 | if TYPE_CHECKING: |
| 44 | from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata |
| 45 | from vllm.v1.attention.backends.mla.index_group import ( |
| 46 | SparseMLAIndexGroupBuilder, |
| 47 | ) |
| 48 | |
| 49 | |
| 50 | class DeepseekV32Indexer(nn.Module): |
| 51 | indexer_cache_cls = DeepseekV32IndexerCache |
| 52 | |
| 53 | def __init__( |
| 54 | self, |
| 55 | vllm_config: VllmConfig, |
| 56 | config: DeepseekV2Config | DeepseekV3Config, |
| 57 | hidden_size: int, |
| 58 | q_lora_rank: int, |
| 59 | quant_config: QuantizationConfig | None, |
| 60 | cache_config: CacheConfig | None, |
| 61 | topk_indices_buffer: torch.Tensor | None, |
| 62 | prefix: str = "", |
| 63 | ): |
| 64 | super().__init__() |
| 65 | self.topk_tokens = config.index_topk |
| 66 | self.n_head = config.index_n_heads |
| 67 | self.head_dim = config.index_head_dim |
| 68 | self.rope_dim = config.qk_rope_head_dim |
| 69 | self.q_lora_rank = q_lora_rank |
| 70 | |
| 71 | self.wq_b = ReplicatedLinear( |
| 72 | self.q_lora_rank, |
| 73 | self.head_dim * self.n_head, |
| 74 | bias=False, |
| 75 | quant_config=quant_config, |
| 76 | prefix=f"{prefix}.wq_b", |
| 77 | ) |
| 78 | |
| 79 | self.wk_weights_proj = MergedColumnParallelLinear( |
| 80 | hidden_size, |
| 81 | [self.head_dim, self.n_head], |
| 82 | bias=False, |
| 83 | quant_config=None, |
| 84 | disable_tp=True, |
| 85 | prefix=f"{prefix}.wk_weights_proj", |
| 86 | ) |
| 87 | self.k_norm = LayerNorm(self.head_dim, eps=1e-6, dtype=torch.float32) |
| 88 | self.softmax_scale = self.head_dim**-0.5 |
| 89 | |
| 90 | self.scale_fmt = "ue8m0" |
| 91 | self.quant_block_size = 128 |
| 92 | self.topk_indices_buffer = topk_indices_buffer |
| 93 | |
| 94 | assert cache_config is not None, "DeepSeek V3.2 indexer requires cache_config" |
| 95 | self.k_cache = type(self).indexer_cache_cls( |
| 96 | head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4, |
| 97 | dtype=torch.uint8, |
| 98 | prefix=f"{prefix}.k_cache", |
| 99 | cache_config=cache_config, |
| 100 | ) |
| 101 | self.max_model_len = vllm_config.model_config.max_model_len |
| 102 | self.prefix = prefix |
| 103 | |
| 104 | from vllm.v1.attention.backends.mla.indexer import ( |
| 105 | get_max_prefill_buffer_size, |
| 106 | ) |
| 107 | |
| 108 | self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config) |
| 109 | self.indexer_op = SparseAttnIndexer( |
| 110 | self.k_cache, |
| 111 | self.quant_block_size, |
| 112 | self.scale_fmt, |
| 113 | self.topk_tokens, |
| 114 | self.head_dim, |
| 115 | self.max_model_len, |
| 116 | self.max_total_seq_len, |
| 117 | self.topk_indices_buffer, |
| 118 | ) |
| 119 | |
| 120 | |
| 121 | class DeepseekV32Attention(MLAAttention): |
| 122 | indexer: "DeepseekV32Indexer | None" |
| 123 | indexer_cls: "type[DeepseekV32Indexer]" = DeepseekV32Indexer |
| 124 | supports_pcp_dcp: ClassVar[bool] = True |
| 125 | |
| 126 | def __init__( |
| 127 | self, |
| 128 | vllm_config: VllmConfig, |
| 129 | config: DeepseekV2Config | DeepseekV3Config, |
| 130 | prefix: str, |
| 131 | topk_indices_buffer: torch.Tensor | None = None, |
| 132 | attn_backend: "type | None" = None, |
| 133 | index_group_builder: "SparseMLAIndexGroupBuilder | None" = None, |
| 134 | ) -> None: |
| 135 | quant_config = vllm_config.quant_config |
| 136 | cache_config = vllm_config.cache_config |
| 137 | |
| 138 | hidden_size = config.hidden_size |
| 139 | qk_nope_head_dim = config.qk_nope_head_dim |
| 140 | qk_rope_head_dim = config.qk_rope_head_dim |
| 141 | v_head_dim = config.v_head_dim |
| 142 | q_lora_rank = config.q_lora_rank |
| 143 | kv_lora_rank = config.kv_lora_rank |
| 144 | num_heads = config.num_attention_heads |
| 145 | |
| 146 | tp_size = get_tensor_model_parallel_world_size() |
| 147 | assert num_heads % tp_size == 0 |
| 148 | num_local_heads = num_heads // tp_size |
| 149 | qk_head_dim = qk_nope_head_dim + qk_rope_head_dim |
| 150 | scaling = qk_head_dim**-0.5 |
| 151 | max_position_embeddings = getattr(config, "max_position_embeddings", 8192) |
| 152 | |
| 153 | # DSA checkpoints may use plain ("default") or yarn-scaled RoPE. |
| 154 | if config.rope_parameters["rope_type"] != "default": |
| 155 | config.rope_parameters["rope_type"] = ( |
| 156 | "deepseek_llama_scaling" |
| 157 | if config.rope_parameters.get("attention_factor") == 1.0 |
| 158 | else "deepseek_yarn" |
| 159 | ) |
| 160 | if config.rope_parameters["rope_type"] == "deepseek_yarn": |
| 161 | mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False) |
| 162 | scaling_factor = config.rope_parameters["factor"] |
| 163 | mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) |
| 164 | scaling = scaling * mscale * mscale |
| 165 | |
| 166 | layer_id = extract_layer_index(prefix) |
| 167 | index_topk_freq = getattr(config, "index_topk_freq", 1) |
| 168 | index_topk_pattern = getattr(config, "index_topk_pattern", None) |
| 169 | index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2) |
| 170 | if index_topk_pattern is None: |
| 171 | skip_topk = ( |
| 172 | max(layer_id - index_skip_topk_offset + 1, 0) % index_topk_freq != 0 |
| 173 | ) |
| 174 | elif 0 <= layer_id < len(index_topk_pattern): |
| 175 | skip_topk = index_topk_pattern[layer_id] == "S" |
| 176 | else: |
| 177 | skip_topk = False |
| 178 | |
| 179 | num_hidden_layers = getattr(config, "num_hidden_layers", None) |
| 180 | is_mtp_layer = num_hidden_layers is not None and layer_id >= num_hidden_layers |
| 181 | |
| 182 | kv_b_proj = ColumnParallelLinear( |
| 183 | kv_lora_rank, |
| 184 | num_heads * (qk_nope_head_dim + v_head_dim), |
| 185 | bias=False, |
| 186 | quant_config=quant_config, |
| 187 | prefix=f"{prefix}.kv_b_proj", |
| 188 | ) |
| 189 | indexer = None |
| 190 | if not skip_topk or is_mtp_layer: |
| 191 | indexer = type(self).indexer_cls( |
| 192 | vllm_config, |
| 193 | config, |
| 194 | hidden_size, |
| 195 | q_lora_rank, |
| 196 | quant_config, |
| 197 | cache_config, |
| 198 | topk_indices_buffer, |
| 199 | prefix=f"{prefix}.indexer", |
| 200 | ) |
| 201 | |
| 202 | super().__init__( |
| 203 | num_heads=num_local_heads, |
| 204 | scale=scaling, |
| 205 | qk_nope_head_dim=qk_nope_head_dim, |
| 206 | qk_rope_head_dim=qk_rope_head_dim, |
| 207 | v_head_dim=v_head_dim, |
| 208 | q_lora_rank=q_lora_rank, |
| 209 | kv_lora_rank=kv_lora_rank, |
| 210 | kv_b_proj=kv_b_proj, |
| 211 | cache_config=cache_config, |
| 212 | quant_config=quant_config, |
| 213 | prefix=f"{prefix}.attn", |
| 214 | use_sparse=True, |
| 215 | indexer=indexer, |
| 216 | topk_indices_buffer=topk_indices_buffer, |
| 217 | index_group_builder=index_group_builder, |
| 218 | attn_backend=attn_backend, |
| 219 | ) |
| 220 | |
| 221 | self.num_local_heads = num_local_heads |
| 222 | self.qk_head_dim = qk_head_dim |
| 223 | self.indexer = indexer |
| 224 | self.topk_indices_buffer = topk_indices_buffer |
| 225 | |
| 226 | self.skip_topk = False |
| 227 | enable_short_prefill_scoring_skip = ( |
| 228 | not is_mtp_layer |
| 229 | and not skip_topk |
| 230 | and not self.use_pcp |
| 231 | and current_platform.is_cuda() |
| 232 | and self.supports_dense_mha_prefill |
| 233 | ) |
| 234 | self._dense_mha_metadata_layer_name = ( |
| 235 | self.layer_name if enable_short_prefill_scoring_skip else "" |
| 236 | ) |
| 237 | |
| 238 | fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype) |
| 239 | self._fp8_query = fp8_attention and self.impl.supports_quant_query_input |
| 240 | self._fp8_kv_needs_view = fp8_attention and not self.kv_cache_dtype.endswith( |
| 241 | "_ds_mla" |
| 242 | ) |
| 243 | |
| 244 | self._index_rope_interleave = getattr(config, "indexer_rope_interleave", False) |
| 245 | |
| 246 | # Remaining MLA projections (registered on this module). |
| 247 | self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear( |
| 248 | hidden_size, |
| 249 | [q_lora_rank, kv_lora_rank + qk_rope_head_dim], |
| 250 | quant_config=quant_config, |
| 251 | prefix=f"{prefix}.fused_qkv_a_proj", |
| 252 | ) |
| 253 | self.q_a_layernorm = RMSNorm(q_lora_rank, eps=config.rms_norm_eps) |
| 254 | self.q_b_proj = ColumnParallelLinear( |
| 255 | q_lora_rank, |
| 256 | num_heads * qk_head_dim, |
| 257 | bias=False, |
| 258 | quant_config=quant_config, |
| 259 | prefix=f"{prefix}.q_b_proj", |
| 260 | ) |
| 261 | self.kv_a_layernorm = RMSNorm(kv_lora_rank, eps=config.rms_norm_eps) |
| 262 | |
| 263 | self.o_proj = RowParallelLinear( |
| 264 | num_heads * v_head_dim, |
| 265 | hidden_size, |
| 266 | bias=False, |
| 267 | reduce_results=False, |
| 268 | quant_config=quant_config, |
| 269 | prefix=f"{prefix}.o_proj", |
| 270 | ) |
| 271 | |
| 272 | self.rotary_emb = get_rope( |
| 273 | qk_rope_head_dim, |
| 274 | max_position=max_position_embeddings, |
| 275 | rope_parameters=config.rope_parameters, |
| 276 | is_neox_style=False, |
| 277 | ) |
| 278 | # Lightning indexer uses its own RoPE; interleave maps to non-NeoX. |
| 279 | self.indexer_rope_emb = get_rope( |
| 280 | qk_rope_head_dim, |
| 281 | max_position=max_position_embeddings, |
| 282 | rope_parameters=config.rope_parameters, |
| 283 | is_neox_style=not getattr(config, "indexer_rope_interleave", False), |
| 284 | ) |
| 285 | |
| 286 | def _should_prepare_mqa_query( |
| 287 | self, |
| 288 | attn_metadata: "MLACommonMetadata | None", |
| 289 | cudagraph_runtime_mode: CUDAGraphMode, |
| 290 | ) -> bool: |
| 291 | return ( |
| 292 | attn_metadata is None |
| 293 | or not self._use_sparse_mha(attn_metadata) |
| 294 | or cudagraph_runtime_mode == CUDAGraphMode.FULL |
| 295 | or torch.cuda.is_current_stream_capturing() |
| 296 | ) |
| 297 | |
| 298 | def forward( # type: ignore[override] |
| 299 | self, |
| 300 | positions: torch.Tensor, |
| 301 | hidden_states: torch.Tensor, |
| 302 | ) -> torch.Tensor: |
| 303 | # Captured: A-projections (+ indexer A-GEMM on indexer layers). |
| 304 | qkv_lora = self.fused_qkv_a_proj(hidden_states)[0] |
| 305 | q_c, kv_c, k_pe = qkv_lora.split( |
| 306 | [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 |
| 307 | ) |
| 308 | |
| 309 | if self.indexer is not None and not self.skip_topk: |
| 310 | kw = self.indexer.wk_weights_proj(hidden_states)[0] |
| 311 | index_k = kw[:, : self.indexer.head_dim] |
| 312 | index_weights = kw[:, self.indexer.head_dim :] |
| 313 | else: |
| 314 | index_k = None |
| 315 | index_weights = None |
| 316 | |
| 317 | num_tokens = hidden_states.shape[0] |
| 318 | output = torch.empty( |
| 319 | (num_tokens, self.num_local_heads * self.v_head_dim), |
| 320 | dtype=hidden_states.dtype, |
| 321 | device=hidden_states.device, |
| 322 | ) |
| 323 | forward_context = get_forward_context() |
| 324 | slot_mapping = forward_context.slot_mapping |
| 325 | assert isinstance(slot_mapping, dict) |
| 326 | mla_slot = slot_mapping.get(self.layer_name) |
| 327 | indexer_slot = ( |
| 328 | slot_mapping.get(self.indexer.k_cache.prefix) |
| 329 | if self.indexer is not None |
| 330 | else None |
| 331 | ) |
| 332 | hisparse_cache = self.hisparse_cache |
| 333 | layer_attn_metadata, _, _, _ = get_attention_context(self.layer_name) |
| 334 | self.impl.prepare_for_batch(layer_attn_metadata) |
| 335 | |
| 336 | prepare_mqa_query = self._should_prepare_mqa_query( |
| 337 | layer_attn_metadata, forward_context.cudagraph_runtime_mode |
| 338 | ) |
| 339 | if self.indexer is not None and not self.skip_topk: |
| 340 | has_indexer = True |
| 341 | indexer_k_norm_w = self.indexer.k_norm.weight |
| 342 | indexer_k_norm_bias = self.indexer.k_norm.bias |
| 343 | indexer_k_norm_eps = self.indexer.k_norm.eps |
| 344 | indexer_k_rope_cos_sin_cache = self.indexer_rope_emb.cos_sin_cache |
| 345 | indexer_cache_shuffled = self.indexer.k_cache.uses_shuffled_layout |
| 346 | indexer_k_cache = None if self.use_pcp else self.indexer.k_cache.kv_cache |
| 347 | index_k_out = torch.empty_like(index_k) if self.use_pcp else None |
| 348 | indexer_softmax_scale = self.indexer.softmax_scale |
| 349 | indexer_n_head_scale = self.indexer.n_head**-0.5 |
| 350 | else: |
| 351 | has_indexer = False |
| 352 | indexer_k_norm_w = None |
| 353 | indexer_k_norm_bias = None |
| 354 | indexer_k_norm_eps = 1e-6 |
| 355 | indexer_k_rope_cos_sin_cache = None |
| 356 | indexer_k_cache = None |
| 357 | indexer_cache_shuffled = False |
| 358 | index_k_out = None |
| 359 | indexer_softmax_scale = 0.0 |
| 360 | indexer_n_head_scale = 0.0 |
| 361 | |
| 362 | if forward_context.attn_metadata is None or self.use_pcp: |
| 363 | mla_kv_cache = None |
| 364 | mla_k_scale = None |
| 365 | indexer_k_cache = None |
| 366 | mla_slot = None |
| 367 | indexer_slot = None |
| 368 | else: |
| 369 | mla_kv_cache = None if hisparse_cache is not None else self.kv_cache |
| 370 | mla_k_scale = self._k_scale |
| 371 | |
| 372 | if self.use_pcp or hisparse_cache is None: |
| 373 | kv_c_out = torch.empty_like(kv_c) |
| 374 | k_pe_out = torch.empty_like(k_pe) |
| 375 | else: |
| 376 | kv_c_out = torch.empty_like(kv_c) |
| 377 | k_pe_out = torch.empty_like(k_pe) |
| 378 | q_c = fused_norm_rope( |
| 379 | positions, |
| 380 | q_c, |
| 381 | self.q_a_layernorm.weight, |
| 382 | self.q_a_layernorm.variance_epsilon, |
| 383 | kv_c, |
| 384 | self.kv_a_layernorm.weight, |
| 385 | self.kv_a_layernorm.variance_epsilon, |
| 386 | k_pe, |
| 387 | self.rotary_emb.cos_sin_cache, |
| 388 | index_k, |
| 389 | indexer_k_norm_w, |
| 390 | indexer_k_norm_bias, |
| 391 | indexer_k_norm_eps, |
| 392 | indexer_k_rope_cos_sin_cache, |
| 393 | self.topk_indices_buffer, |
| 394 | slot_mapping=mla_slot, |
| 395 | indexer_slot_mapping=indexer_slot, |
| 396 | indexer_k_cache=indexer_k_cache, |
| 397 | indexer_cache_shuffled=indexer_cache_shuffled, |
| 398 | mla_kv_cache=mla_kv_cache, |
| 399 | mla_kv_cache_dtype=self.kv_cache_dtype, |
| 400 | mla_k_scale=mla_k_scale, |
| 401 | has_indexer=has_indexer, |
| 402 | index_rope_interleave=self._index_rope_interleave, |
| 403 | kv_c_out=kv_c_out, |
| 404 | k_pe_out=k_pe_out, |
| 405 | index_k_out=index_k_out, |
| 406 | ) |
| 407 | |
| 408 | if hisparse_cache is not None and mla_slot is not None: |
| 409 | assert kv_c_out is not None and k_pe_out is not None |
| 410 | self.update_kv_cache( |
| 411 | kv_c_out, |
| 412 | k_pe_out, |
| 413 | self.kv_cache, |
| 414 | mla_slot, |
| 415 | layer_attn_metadata, |
| 416 | self.kv_cache_dtype, |
| 417 | self._k_scale, |
| 418 | ) |
| 419 | |
| 420 | q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim) |
| 421 | q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) |
| 422 | if prepare_mqa_query: |
| 423 | ql_nope = torch.bmm(q_nope.transpose(0, 1), self.W_UK_T).transpose(0, 1) |
| 424 | else: |
| 425 | ql_nope = q_nope |
| 426 | |
| 427 | if self.indexer is not None and not self.skip_topk: |
| 428 | index_q = self.indexer.wq_b(q_c)[0] |
| 429 | index_q = index_q.view(-1, self.indexer.n_head, self.indexer.head_dim) |
| 430 | else: |
| 431 | index_q = None |
| 432 | |
| 433 | index_q_fp8, index_weights_out, mqa_q = fused_q( |
| 434 | positions, |
| 435 | q_pe, |
| 436 | self.rotary_emb.cos_sin_cache, |
| 437 | index_q, |
| 438 | self.indexer_rope_emb.cos_sin_cache if has_indexer else None, |
| 439 | ql_nope, |
| 440 | self._q_scale, |
| 441 | index_weights, |
| 442 | indexer_softmax_scale, |
| 443 | indexer_n_head_scale, |
| 444 | has_indexer=has_indexer, |
| 445 | index_rope_interleave=self._index_rope_interleave, |
| 446 | quantize_mqa=self._fp8_query and prepare_mqa_query, |
| 447 | ) |
| 448 | |
| 449 | self._sparse_indexer_and_attn( |
| 450 | positions, |
| 451 | q_c, |
| 452 | q_nope, |
| 453 | q_pe, |
| 454 | index_q_fp8, |
| 455 | index_k_out, |
| 456 | index_weights_out, |
| 457 | kv_c_out, |
| 458 | k_pe_out, |
| 459 | ql_nope, |
| 460 | mqa_q, |
| 461 | output, |
| 462 | ) |
| 463 | return self.o_proj(output)[0] |
| 464 | |
| 465 | @eager_break_during_capture |
| 466 | def _sparse_indexer_and_attn( |
| 467 | self, |
| 468 | positions: torch.Tensor, |
| 469 | q_c: torch.Tensor, |
| 470 | q_nope: torch.Tensor, |
| 471 | q_pe: torch.Tensor, |
| 472 | index_q_fp8: torch.Tensor | None, |
| 473 | index_k: torch.Tensor | None, |
| 474 | index_weights_out: torch.Tensor | None, |
| 475 | kv_c: torch.Tensor | None, |
| 476 | k_pe: torch.Tensor | None, |
| 477 | ql_nope: torch.Tensor, |
| 478 | mqa_q: torch.Tensor, |
| 479 | output: torch.Tensor, |
| 480 | ) -> None: |
| 481 | if self.indexer is not None and not self.skip_topk: |
| 482 | assert index_q_fp8 is not None |
| 483 | assert index_weights_out is not None |
| 484 | if self.use_pcp: |
| 485 | assert index_k is not None |
| 486 | sparse_attn_indexer( |
| 487 | q_c, |
| 488 | self.indexer.k_cache.prefix, |
| 489 | self.indexer.k_cache.kv_cache, |
| 490 | index_q_fp8, |
| 491 | None, |
| 492 | index_k, |
| 493 | index_weights_out, |
| 494 | self.indexer.quant_block_size, |
| 495 | self.indexer.scale_fmt, |
| 496 | self.indexer.topk_tokens, |
| 497 | self.indexer.head_dim, |
| 498 | self.indexer.max_model_len, |
| 499 | self.indexer.max_total_seq_len, |
| 500 | self.topk_indices_buffer, |
| 501 | skip_k_cache_insert=not self.use_pcp, |
| 502 | use_pcp=self.use_pcp, |
| 503 | pcp_shard_decode_requests=self.pcp_shard_decode_requests, |
| 504 | dense_mha_metadata_layer_name=self._dense_mha_metadata_layer_name, |
| 505 | dcp_rank=( |
| 506 | self.dcp_manager.group.rank_in_group |
| 507 | if self.dcp_manager is not None |
| 508 | else 0 |
| 509 | ), |
| 510 | dcp_world_size=( |
| 511 | self._vllm_config.parallel_config.decode_context_parallel_size |
| 512 | ), |
| 513 | cp_kv_cache_interleave_size=( |
| 514 | self._vllm_config.parallel_config.cp_kv_cache_interleave_size |
| 515 | ), |
| 516 | skip_topk_buffer_clear=True, |
| 517 | topk_backend=self.indexer.indexer_op.topk_backend, |
| 518 | ) |
| 519 | self.impl.record_logical_topk_ready() # type: ignore[attr-defined] |
| 520 | |
| 521 | attn_metadata, _, kv_cache, layer_slot_mapping = get_attention_context( |
| 522 | self.layer_name |
| 523 | ) |
| 524 | if attn_metadata is None: |
| 525 | output.zero_() |
| 526 | return |
| 527 | attn_metadata = cast("MLACommonMetadata", attn_metadata) |
| 528 | self.impl.prepare_for_batch(attn_metadata) |
| 529 | |
| 530 | if self.use_pcp: |
| 531 | assert kv_c is not None and k_pe is not None |
| 532 | kv_for_cache, kpe_for_cache, cache_slot_mapping = ( |
| 533 | maybe_gather_mla_latent_cache_inputs( |
| 534 | kv_c, |
| 535 | k_pe.unsqueeze(1), |
| 536 | layer_slot_mapping, |
| 537 | attn_metadata.num_decode_tokens, |
| 538 | True, |
| 539 | pcp_shard_decode_requests=self.pcp_shard_decode_requests, |
| 540 | ) |
| 541 | ) |
| 542 | self.impl.do_kv_cache_update( # type: ignore[attr-defined] |
| 543 | kv_for_cache, |
| 544 | kpe_for_cache, |
| 545 | kv_cache, |
| 546 | cache_slot_mapping, |
| 547 | self.kv_cache_dtype, |
| 548 | self._k_scale, |
| 549 | ) |
| 550 | |
| 551 | num_actual = attn_metadata.num_actual_tokens # type: ignore[attr-defined] |
| 552 | if num_actual == 0: |
| 553 | output.zero_() |
| 554 | return |
| 555 | |
| 556 | if self._use_sparse_mha(attn_metadata): |
| 557 | assert kv_c is not None and k_pe is not None |
| 558 | mha_q_pe = ( |
| 559 | self.rotary_emb(positions, q_pe)[0] |
| 560 | if mqa_q.dtype != q_pe.dtype |
| 561 | else mqa_q |
| 562 | ) |
| 563 | mha_q = torch.cat((q_nope, mha_q_pe), dim=-1) |
| 564 | self.forward_impl( |
| 565 | mha_q, |
| 566 | kv_c, |
| 567 | k_pe.unsqueeze(1), |
| 568 | kv_cache, |
| 569 | attn_metadata, |
| 570 | output, |
| 571 | ) |
| 572 | return |
| 573 | |
| 574 | if self._fp8_kv_needs_view: |
| 575 | kv_cache = kv_cache.view(current_platform.fp8_dtype()) |
| 576 | if self._fp8_query: |
| 577 | # FlashInfer sparse: single packed fp8 query. |
| 578 | mqa_q_arg: torch.Tensor | tuple[torch.Tensor, torch.Tensor] = mqa_q[ |
| 579 | :num_actual |
| 580 | ] |
| 581 | else: |
| 582 | mqa_q_arg = (ql_nope[:num_actual], mqa_q[:num_actual]) |
| 583 | |
| 584 | if self.use_pcp and self.impl.dcp_world_size > self.impl.pcp_world_size: |
| 585 | if isinstance(mqa_q_arg, tuple): |
| 586 | mqa_q_arg = torch.cat(mqa_q_arg, dim=-1) |
| 587 | mqa_q_arg = get_tp_group().all_gather(mqa_q_arg, dim=1) |
| 588 | elif not self.use_pcp and self.impl.dcp_world_size > 1: |
| 589 | assert self.dcp_manager is not None |
| 590 | if isinstance(mqa_q_arg, tuple): |
| 591 | mqa_q_arg = torch.cat(mqa_q_arg, dim=-1) |
| 592 | assert self.dcp_manager.query_gather is not None |
| 593 | mqa_q_arg = self.dcp_manager.query_gather(mqa_q_arg) |
| 594 | attn_out, lse = self.impl.forward_mqa( # type: ignore[attr-defined] |
| 595 | mqa_q_arg, kv_cache, attn_metadata, self |
| 596 | ) |
| 597 | |
| 598 | if self.impl.dcp_world_size > 1: |
| 599 | assert lse is not None and self.dcp_manager is not None |
| 600 | seq_lens: torch.Tensor | None |
| 601 | query_start_loc: torch.Tensor | None |
| 602 | if self.use_pcp: |
| 603 | decode_metadata = getattr(attn_metadata, "decode", None) |
| 604 | if decode_metadata is not None: |
| 605 | seq_lens = decode_metadata.seq_lens |
| 606 | else: |
| 607 | all_seq_lens = cast( |
| 608 | torch.Tensor, |
| 609 | attn_metadata.seq_lens, # type: ignore[attr-defined] |
| 610 | ) |
| 611 | seq_lens = all_seq_lens[: attn_metadata.num_decodes] |
| 612 | query_start_loc = attn_metadata.query_start_loc[ |
| 613 | : attn_metadata.num_decodes + 1 |
| 614 | ] |
| 615 | else: |
| 616 | # The backend emits (0, -inf) for empty local shards, so no |
| 617 | # PCP-only empty-shard metadata is needed. |
| 618 | seq_lens = None |
| 619 | query_start_loc = None |
| 620 | # Under PCP+DCP the prefill rows attended over the gathered KV, so |
| 621 | # only the decode rows carry an LSE and take part in the merge. |
| 622 | num_merge_rows = lse.shape[0] |
| 623 | if num_merge_rows == attn_out.shape[0]: |
| 624 | attn_out = self.dcp_manager.combine( |
| 625 | attn_out, |
| 626 | lse, |
| 627 | seq_lens=seq_lens, # type: ignore[arg-type] |
| 628 | query_start_loc=query_start_loc, # type: ignore[arg-type] |
| 629 | ) |
| 630 | elif num_merge_rows > 0: |
| 631 | attn_out[:num_merge_rows] = self.dcp_manager.combine( |
| 632 | attn_out[:num_merge_rows], |
| 633 | lse, |
| 634 | seq_lens=seq_lens, # type: ignore[arg-type] |
| 635 | query_start_loc=query_start_loc, # type: ignore[arg-type] |
| 636 | ) |
| 637 | if self.use_pcp: |
| 638 | attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads) |
| 639 | |
| 640 | # NOTE(woosuk): While the below does not need to be in the eager region, |
| 641 | # we put it here to avoid copying the attention output. Move this back to the |
| 642 | # captured region once forward_mqa supports `out` argument. |
| 643 | x = attn_out.view( |
| 644 | num_actual, self.num_local_heads, self.kv_lora_rank |
| 645 | ).transpose(0, 1) |
| 646 | out = ( |
| 647 | output[:num_actual] |
| 648 | .view(num_actual, self.num_local_heads, self.v_head_dim) |
| 649 | .transpose(0, 1) |
| 650 | ) |
| 651 | torch.bmm(x, self.W_UV, out=out) |
| 652 | if self.use_pcp and num_actual < output.shape[0]: |
| 653 | output[num_actual:].zero_() |