mla_attention.py · tduncan/Nvidia-Optimization-Stack-VLLM

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 3,329 lines 128 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3	"""# MLA Common Components
4
5	This file implements common components for MLA implementations.
6
7	First we define:
8
9	Sq as Q sequence length
10	Skv as KV sequence length
11
12	MLA has two possible ways of computing, a data-movement friendly approach and a
13	compute friendly approach. We generally want to use the compute friendly
14	approach for "prefill" (i.e. the ratio Sq / Skv is relatively large, often near
15	1) and the data-movement friendly approach for "decode" (i.e. the ratio
16	Sq / Skv is small, often near 0).
17
18	NOTE what we deem small and large is currently determined by if it is labelled
19	prefill or decode by the scheduler, but this is something we should probably
20	tune.
21
22	Main reference: DeepseekV2 paper, and FlashInfer Implementation
23	(https://arxiv.org/abs/2405.04434 and https://github.com/flashinfer-ai/flashinfer/pull/551).
24
25	Deepseek's MLA attention works the following way:
26	* Use a single latent vector to represent the per-token entry of the KV cache.
27	* For decode (i.e. the memory friendly approach) the attention "simulates" a
28	multi-head attention, while the compute is similar to multi-query attention.
29
30	Below is an example of both paths assuming batch size = 1
31
32	## More Extent Definitions:
33
34	C Context length, `Skv - Sq`
35	H hidden size
36	N number of attention heads
37	Lq latent dimension for Q 1536 in DSV3
38	Lkv latent dimension for K/V 512 in DSV3
39	P nope dimension, no rope. 128 in DSV3
40	R rope dimension, goes through rope. 64 in DSV3
41	V V head dim. 128 in DSV3
42
43	## Vector/Matrix Definitions
44
45	h_t hidden states (input to attention) shape [Sq, H]
46	q_c latent/compressed Q shape [Sq, Lq]
47	q_nope uncompressed Q (no-rope) shape [Sq, N, P]
48	q_pe uncompressed Q (rope) shape [Sq, N, R]
49	kv_c latent/compressed KV shape [Skv, Lkv]
50	k_pe decoupled k position embeddings shape [Skv, R]
51	new_kv_c new kv_c from current iter shape [Sq, Lkv]
52	new_k_pe new k_pe from current iter shape [Sq, R]
53	cache_kv_c cached k_c from previous iters shape [C, Lkv]
54	cache_k_pe cached k_pe from previous iters shape [C, R]
55	W_DQ project h_t to q_c shape [H, Lq]
56	W_UQ project q_c to q_nope shape [Lq, N * P]
57	W_QR project q_c to q_pe shape [Lq, N * R]
58	W_DKV project h_t to kv_c shape [H, Lkv]
59	W_UK project kv_c to k_nope shape [Lkv, N, P]
60	W_KR project h_t to k_pe shape [H, R]
61	W_UV project kv_c to v shape [Lkv, N, V]
62	W_O project v to h_t shape [N * V, H]
63
64
65	## Compute Friendly Approach (i.e. "forward_mha"):
66
67	q_c = h_t @ W_DQ
68	q_nope = (q_c @ W_UQ).view(Sq, N, P)
69	q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)
70	new_kv_c = h_t @ W_DKV
71	new_k_pe = RoPE(h_t @ W_KR)
72	kv_c = torch.cat([new_kv_c, cache_kv_c], dim=0)
73	k_pe = torch.cat([new_k_pe, cache_k_pe], dim=0)
74	k_nope = (kv_c @ W_UK.view(Lkv, N * P)).view(Skv, N, P)
75	v = (kv_c @ W_UV.view(Lkv, N * V)).view(Skv, N, V)
76
77	// MHA with QK headdim = P + R
78	// V headdim = V
79	// sdpa_o shape [Sq, N, V]
80	sdpa_o = scaled_dot_product_attention(
81	torch.cat([q_nope, q_pe], dim=-1),
82	torch.cat([k_nope, k_pe.unsqueeze(1).expand(-1, N, -1)], dim=-1),
83	v
84	)
85	return sdpa_o @ W_O
86
87	NOTE: in the actual code,
88	`kv_b_proj` is [W_UK; W_UV] concatenated per head
89	`q_b_proj` is [W_UQ; W_QR] concatenated per head
90	`out_proj` is W_O
91
92
93	## Data-Movement Friendly Approach (i.e. "forward_mqa"):
94
95	Runtime
96	q_c = h_t @ W_DQ
97	q_nope = (q_c @ W_UQ).view(-1, N, P)
98	ql_nope = einsum("snh,lnh->snl", q_nope, W_UK)
99	q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)
100	new_kv_c = h_t @ W_DKV
101	new_k_pe = RoPE(h_t @ W_KR)
102	kv_c = torch.cat([new_kv_c, cache_kv_c], dim=0)
103	k_pe = torch.cat([new_k_pe, cache_k_pe], dim=0)
104
105	// MQA with QK headdim = Lkv + R
106	// V headdim = Lkv
107	// sdpa_o shape [Sq, N, Lkv]
108	// NOTE: this is less compute-friendly since Lkv > P
109	// but is more data-movement friendly since its MQA vs MHA
110	sdpa_o = scaled_dot_product_attention(
111	torch.cat([ql_nope, q_pe], dim=-1),
112	torch.cat([kv_c, k_pe], dim=-1),
113	kv_c
114	)
115
116	o = einsum("snl,lnv->snv", sdpa_o.reshape(-1, N, Lkv), W_UV)
117	return o.view(-1, N * V) @ W_O
118
119
120	## Chunked Prefill
121
122	For chunked prefill we want to use the compute friendly algorithm. We are
123	assuming sufficiently large Sq / Skv ratio, in the future may want to switch to
124	the data-movement friendly approach if the chunk (i.e. `Sq`) is small.
125
126	However, the compute-friendly approach can potentially run out of memory if Skv
127	is large due to: `k_nope = (kv_c @ W_UK).view(Skv, N, P)`
128
129	To mitigate this, we chunk the computation of attention with respect to the
130	current context (i.e. `cache_kv_c` and `cache_k_pe`) so that we can used a
131	fixed workspace size.
132
133	The chunked prefill approach is as follows:
134
135	W Workspace rows, i.e. the context rows we may gather at once,
136	used to bound the memory usage
137
138	The context is scheduled per request: requests are packed in order, splitting
139	the next request on an aligned boundary when needed to fill `W`.
140	So a chunk covers a contiguous run of prefills and only ever the tokens and
141	context rows of those prefills — attention, up-projection and merging are
142	charged only to the requests it covers, and no chunk contains an empty context
143	span. See `plan_mla_context_chunks`.
144
145	Only a chunk's first request can be a continuation, so accumulation performs at
146	most one request-slice merge and one bulk write per chunk.
147
148	q_c = h_t @ W_DQ
149	q_nope = (q_c @ W_UQ).view(Sq, N, P)
150	q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)
151	new_kv_c = h_t @ W_DKV
152	new_k_pe = RoPE(h_t @ W_KR)
153	new_k_nope = (new_kv_c @ W_UK.view(Lkv, N * P)).view(Sq, N, P)
154	new_v = (new_kv_c @ W_UV.view(Lkv, N * V)).view(Sq, N, V)
155
156	// MHA between queries and new KV
157	// with QK headdim = P + R
158	// V headdim = V
159	// curr_o shape [Sq, N, V]
160	// curr_lse shape [N, Sq], this is just order FA returns
161	curr_o, curr_lse = scaled_dot_product_attention(
162	torch.cat([q_nope, q_pe], dim=-1),
163	torch.cat([new_k_nope, new_k_pe.unsqueeze(1).expand(-1, N, -1)], dim=-1),
164	new_v,
165	causal=True,
166	return_softmax_lse=True
167	)
168
169	// Compute attention with the already existing context. Shown for a single
170	// request; with a batch, a chunk covers a run of requests and both `q` and the
171	// gathered context below are sliced to that run.
172	for chunk_idx in range(cdiv(C, W)):
173	chunk_start = chunk_idx * W
174	chunk_end = min(chunk_start + W, C)
175	Sc = chunk_end - chunk_start
176	cache_kv_c_chunk = cache_kv_c[chunk_start:chunk_end]
177	cache_k_pe_chunk = cache_k_pe[chunk_start:chunk_end]
178	cache_k_nope_chunk = (cache_kv_c_chunk @ W_UK).view(-1, N, P)
179	cache_v_chunk = (cache_kv_c_chunk @ W_UV).view(-1, N, V)
180
181	chunk_o, chunk_lse = scaled_dot_product_attention(
182	torch.cat([q_nope, q_pe], dim=-1),
183	torch.cat([cache_k_nope_chunk,
184	cache_k_pe_chunk.unsqueeze(1).expand(-1, N, -1)],
185	dim=-1),
186	cache_v_chunk,
187	causal=False,
188	return_softmax_lse=True
189	)
190
191	curr_o, curr_lse = merge_attn_states(
192	suffix_output=curr_o,
193	suffix_lse=curr_lse,
194	prefix_output=chunk_o,
195	prefix_lse=chunk_lse,
196	)
197
198	return curr_o @ W_O
199	"""
200
201	import functools
202	import itertools
203	import math
204	from abc import abstractmethod
205	from collections.abc import Callable
206	from dataclasses import dataclass, replace
207	from enum import Enum
208	from math import lcm
209	from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast
210
211	import numpy as np
212	import torch
213	import torch.nn as nn
214	from tqdm import tqdm
215
216	import vllm.envs as envs
217	from vllm import _custom_ops as ops
218	from vllm._aiter_ops import rocm_aiter_ops
219	from vllm.compilation.breakable_cudagraph import eager_break_during_capture
220	from vllm.config import (
221	CacheConfig,
222	ModelConfig,
223	VllmConfig,
224	get_current_vllm_config,
225	)
226	from vllm.config.cache import CacheDType
227	from vllm.distributed.parallel_state import (
228	get_dcp_group,
229	get_tp_group,
230	is_global_first_rank,
231	)
232	from vllm.forward_context import ForwardContext, get_forward_context
233	from vllm.logger import init_logger
234	from vllm.model_executor.custom_op import CustomOp
235	from vllm.model_executor.layers.attention.attention import (
236	_init_kv_cache_quant,
237	get_attention_context,
238	set_default_quant_scales,
239	should_load_quant_weights,
240	)
241	from vllm.model_executor.layers.attention.kv_transfer_utils import (
242	maybe_transfer_kv_layer,
243	)
244	from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
245	from vllm.model_executor.layers.linear import (
246	ColumnParallelLinear,
247	)
248	from vllm.model_executor.layers.quantization import (
249	QuantizationConfig,
250	resolve_quant_method,
251	)
252	from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
253	from vllm.model_executor.layers.quantization.utils.quant_utils import (
254	GroupShape,
255	QuantKey,
256	get_and_maybe_dequant_weights,
257	kFp8Dynamic64Sym,
258	kFp8Dynamic128Sym,
259	kFp8StaticTensorSym,
260	kNvfp4Dynamic,
261	)
262	from vllm.model_executor.utils import replace_parameter
263	from vllm.platforms import current_platform
264	from vllm.utils.flashinfer import has_flashinfer
265	from vllm.utils.math_utils import cdiv, round_down, round_up
266	from vllm.utils.torch_utils import (
267	PIN_MEMORY,
268	LayerNameType,
269	_encode_layer_name,
270	_resolve_layer_name,
271	async_tensor_h2d,
272	direct_register_custom_op,
273	get_dtype_size,
274	is_quantized_kv_cache,
275	kv_cache_dtype_str_to_dtype,
276	)
277	from vllm.v1.attention.backend import (
278	AttentionBackend,
279	AttentionLayer,
280	AttentionMetadata,
281	AttentionMetadataBuilder,
282	AttentionType,
283	CommonAttentionMetadata,
284	MLAAttentionImpl,
285	)
286	from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
287	from vllm.v1.attention.backends.mla.prefill import (
288	MLAPrefillBackend,
289	get_mla_prefill_backend,
290	)
291	from vllm.v1.attention.backends.utils import (
292	get_dcp_local_seq_lens,
293	get_num_attention_heads_from_layers,
294	split_decodes_and_prefills,
295	)
296	from vllm.v1.attention.ops.dcp import MLADCPManager
297	from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
298	from vllm.v1.attention.ops.pcp import (
299	finalize_mla_pcp_decode,
300	maybe_gather_mla_latent_cache_inputs,
301	)
302	from vllm.v1.attention.selector import get_attn_backend
303	from vllm.v1.kv_cache_interface import (
304	AttentionSpec,
305	KVCacheSpec,
306	KVQuantMode,
307	MLAAttentionSpec,
308	SlidingWindowMLASpec,
309	get_kv_quant_mode,
310	)
311
312	if TYPE_CHECKING:
313	from vllm.v1.attention.backends.mla.index_group import (
314	SparseMLAIndexGroupBuilder,
315	)
316
317	logger = init_logger(__name__)
318
319	_FP8_DTYPE = current_platform.fp8_dtype()
320
321
322	def _detect_output_quant_key(
323	output: torch.Tensor,
324	output_scale: torch.Tensor \| None,
325	output_block_scale: torch.Tensor \| None,
326	output_dim: int,
327	) -> QuantKey \| None:
328	"""Detect the output quantization key from fusion pass parameters.
329
330	Returns the appropriate QuantKey, or None if no quantization is needed.
331	Detection is based on output dtype and which scale tensors are present.
332	"""
333	if output_scale is None and output_block_scale is None:
334	return None
335	if output_block_scale is not None:
336	if output.dtype == _FP8_DTYPE:
337	# Per-group FP8 uses block scales only, not a separate output_scale
338	assert output_scale is None
339	# Infer group size from scale shape
340	num_groups = output_block_scale.shape[-1]
341	group_size = output_dim // num_groups
342	if group_size == 128:
343	return kFp8Dynamic128Sym
344	elif group_size == 64:
345	return kFp8Dynamic64Sym
346	else:
347	raise ValueError(
348	f"Unsupported group FP8 group_size={group_size} "
349	f"(output_dim={output_dim}, num_groups={num_groups}). "
350	f"Only group_size 128 and 64 are supported."
351	)
352	# output_scale None implies MXFP4, not supported
353	assert output_scale is not None
354	return kNvfp4Dynamic
355	return kFp8StaticTensorSym
356
357
358	def _canonicalize_sparse_mla_kv_cache_dtype(
359	attn_backend: type[AttentionBackend],
360	kv_cache_dtype: CacheDType,
361	) -> CacheDType:
362	backend_name = attn_backend.get_name()
363	if backend_name == "FLASHMLA_SPARSE" and is_quantized_kv_cache(kv_cache_dtype):
364	# The NVFP4 DS-MLA format is used as-is; any other quantized dtype
365	# (fp8, fp8_e4m3, ...) is served via the fp8_ds_mla format.
366	if kv_cache_dtype == "nvfp4_ds_mla":
367	return kv_cache_dtype
368	return "fp8_ds_mla"
369	if backend_name == "FLASHINFER_MLA_SPARSE_SM120" and kv_cache_dtype in (
370	"auto",
371	"fp8",
372	"fp8_e4m3",
373	):
374	return "fp8_ds_mla"
375	return kv_cache_dtype
376
377
378	def _get_kv_b_proj_input_dtype(
379	kv_b_proj: ColumnParallelLinear, use_fp8_prefill: bool
380	) -> torch.dtype \| None:
381	weight = getattr(kv_b_proj, "weight", None)
382	weight_dtype = weight.dtype if weight is not None else kv_b_proj.params_dtype
383	if weight_dtype == torch.int32:
384	return kv_b_proj.params_dtype
385	if weight_dtype == torch.uint8:
386	return None
387	if weight_dtype == current_platform.fp8_dtype():
388	from vllm.model_executor.layers.quantization.modelopt import (
389	ModelOptLinearMethod,
390	)
391	from vllm.model_executor.layers.quantization.utils.quant_utils import (
392	kFp8Static128BlockSym,
393	)
394
395	quant_method = getattr(kv_b_proj, "quant_method", None)
396	# FP8_PB_WO dynamically quantizes BF16/FP16 inputs in the linear method.
397	if (
398	isinstance(quant_method, ModelOptLinearMethod)
399	and quant_method.spec.weight is kFp8Static128BlockSym
400	):
401	return quant_method.input_dtype
402	if not use_fp8_prefill:
403	return None
404	return weight_dtype
405
406
407	def split_kv_b_proj(
408	kv_b_proj: nn.Module,
409	out_dtype: torch.dtype,
410	kv_lora_rank: int,
411	num_heads: int,
412	qk_nope_head_dim: int,
413	v_head_dim: int,
414	) -> tuple[torch.Tensor, torch.Tensor]:
415	"""Dequantize ``kv_b_proj`` and return ``W_UK [L,N,P]``, ``W_UV [L,N,V]``."""
416	weight = get_and_maybe_dequant_weights(kv_b_proj, out_dtype=out_dtype).T
417	assert weight.shape == (
418	kv_lora_rank,
419	num_heads * (qk_nope_head_dim + v_head_dim),
420	), f"kv_b_proj weight {tuple(weight.shape)} vs {kv_lora_rank=} {num_heads=}"
421	weight = weight.view(kv_lora_rank, num_heads, qk_nope_head_dim + v_head_dim)
422	return weight.split([qk_nope_head_dim, v_head_dim], dim=-1)
423
424
425	class MLAAttention(nn.Module, AttentionLayerBase):
426	"""Multi-Head Latent Attention layer.
427
428	NOTE: Please read the comment at the top of the file before trying to
429	understand this class
430
431	This class takes query, and compressed key/value tensors as input.
432	The class does the following:
433
434	1. Store the input key and value tensors in the KV cache.
435	2. Perform (multi-head/multi-query/grouped-query) attention.
436	3. Return the output tensor.
437	"""
438
439	supports_dense_mha_prefill: ClassVar[bool] = True
440	# Under PCP+DCP only the decode rows carry an LSE; the base forward
441	# merges a full-batch LSE, so subclasses opt in with their own forward.
442	supports_pcp_dcp: ClassVar[bool] = False
443
444	def __init__(
445	self,
446	num_heads: int,
447	scale: float,
448	qk_nope_head_dim: int,
449	qk_rope_head_dim: int,
450	v_head_dim: int,
451	q_lora_rank: int \| None,
452	kv_lora_rank: int,
453	kv_b_proj: ColumnParallelLinear,
454	dcp_q_replicate: bool = False,
455	cache_config: CacheConfig \| None = None,
456	quant_config: QuantizationConfig \| None = None,
457	prefix: str = "",
458	attn_backend: type[AttentionBackend] \| None = None,
459	use_sparse: bool = False,
460	indexer: object \| None = None,
461	topk_indices_buffer: torch.Tensor \| None = None,
462	index_group_builder: "SparseMLAIndexGroupBuilder \| None" = None,
463	non_causal_multi_token_decode: bool = False,
464	sliding_window: int \| None = None,
465	prefill_backend_cls: type[MLAPrefillBackend] \| None = None,
466	**extra_impl_args,
467	):
468	super().__init__()
469	self.num_heads = num_heads
470	self.scale = scale
471	self.qk_nope_head_dim = qk_nope_head_dim
472	self.qk_rope_head_dim = qk_rope_head_dim
473	self.v_head_dim = v_head_dim
474	self.q_lora_rank = q_lora_rank
475	self.kv_lora_rank = kv_lora_rank
476	self.kv_b_proj = kv_b_proj
477	self.dcp_q_replicate = dcp_q_replicate
478	self.W_UK_T_dcp_qrep: torch.Tensor \| None = None
479	self.head_size = kv_lora_rank + qk_rope_head_dim
480	self.layer_name = prefix
481	self.indexer = indexer
482	self.non_causal_multi_token_decode = non_causal_multi_token_decode
483	self.sliding_window = sliding_window
484	self.num_kv_heads = 1
485	self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
486
487	if cache_config is not None:
488	kv_cache_dtype: CacheDType = cache_config.cache_dtype
489	else:
490	kv_cache_dtype = "auto"
491	self.quant_config = quant_config
492
493	if cache_config is not None and cache_config.kv_cache_dtype_skip_layers:
494	from vllm.model_executor.models.utils import extract_layer_index
495
496	layer_idx = extract_layer_index(prefix)
497	if str(layer_idx) in cache_config.kv_cache_dtype_skip_layers:
498	kv_cache_dtype = "auto"
499	logger.debug(
500	"Layer %s: kv_cache_dtype=%s",
501	prefix,
502	kv_cache_dtype,
503	)
504
505	dtype = torch.get_default_dtype()
506	if attn_backend is not None:
507	assert attn_backend.is_mla(), (
508	f"MLAAttention: attn_backend must be an MLA backend, "
509	f"got {attn_backend.get_name()} instead"
510	)
511	self.attn_backend = attn_backend
512	else:
513	self.attn_backend = get_attn_backend(
514	self.head_size,
515	dtype,
516	kv_cache_dtype,
517	use_mla=True,
518	use_sparse=use_sparse,
519	num_heads=self.num_heads,
520	)
521
522	normalized_kv_cache_dtype = _canonicalize_sparse_mla_kv_cache_dtype(
523	self.attn_backend, kv_cache_dtype
524	)
525	if normalized_kv_cache_dtype != kv_cache_dtype:
526	if cache_config is not None:
527	cache_config.cache_dtype = normalized_kv_cache_dtype
528	kv_cache_dtype = normalized_kv_cache_dtype
529	logger.info_once(
530	"Using %s KV cache format for %s backend.",
531	kv_cache_dtype,
532	self.attn_backend.get_name(),
533	)
534
535	if (
536	self.attn_backend.get_name() == "FLASHINFER_MLA_SPARSE"
537	and kv_cache_dtype != "fp8_ds_mla"
538	and is_quantized_kv_cache(kv_cache_dtype)
539	):
540	logger.info_once(
541	"Using standard fp8 KV cache format. To use DeepSeek's fp8_ds_mla "
542	"KV cache format, please set `--attention-backend FLASHMLA_SPARSE`"
543	)
544
545	# Initialize KV cache quantization attributes
546	self.kv_cache_dtype = kv_cache_dtype
547	_init_kv_cache_quant(self, quant_config, prefix)
548
549	if (
550	cache_config is not None
551	and cache_config.enable_prefix_caching
552	and envs.VLLM_BATCH_INVARIANT
553	and (
554	self.attn_backend.get_name() == "TRITON_MLA"
555	or self.attn_backend.get_name() == "FLASHINFER"
556	)
557	):
558	logger.warning_once(
559	"Disabling prefix caching for TRITON_MLA / FLASHINFER "
560	"with batch invariance, as it is not yet supported.",
561	)
562	cache_config.enable_prefix_caching = False
563
564	# Sparse MLA reads top-k indices from a shared buffer. Pass it
565	# explicitly so backbone "skip" layers (indexer=None) still find it.
566	if use_sparse:
567	extra_impl_args["topk_indices_buffer"] = topk_indices_buffer
568	if index_group_builder is not None:
569	extra_impl_args["index_group_builder"] = index_group_builder
570
571	impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls())
572	impl = impl_cls(
573	num_heads=self.num_heads,
574	head_size=self.head_size,
575	scale=self.scale,
576	num_kv_heads=1,
577	alibi_slopes=None,
578	sliding_window=sliding_window,
579	kv_cache_dtype=self.kv_cache_dtype,
580	logits_soft_cap=None,
581	attn_type=AttentionType.DECODER,
582	kv_sharing_target_layer_name=None,
583	# MLA Args
584	q_lora_rank=self.q_lora_rank,
585	kv_lora_rank=self.kv_lora_rank,
586	qk_nope_head_dim=self.qk_nope_head_dim,
587	qk_rope_head_dim=self.qk_rope_head_dim,
588	qk_head_dim=self.qk_nope_head_dim + self.qk_rope_head_dim,
589	v_head_dim=self.v_head_dim,
590	kv_b_proj=kv_b_proj,
591	indexer=indexer,
592	**extra_impl_args,
593	)
594	self.impl = impl # type: ignore[assignment]
595	index_group = getattr(impl, "index_group", None)
596	self.hisparse_cache = (
597	index_group.cache(cast(Any, impl).index_group_index)
598	if isinstance(index_group, HiSparseMLAIndexGroup)
599	else None
600	)
601	self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None)
602	self.is_amx_bmm_enabled = getattr(self.impl, "uses_amx_bmm", False)
603	# AMX reads kv_b_proj's weight directly and never calls it live; the
604	# reference CPU MLA backend calls it but isn't perf-critical. Skip
605	# the packed-kernel dispatch either way.
606	kv_b_proj._cpu_skip_gemm_dispatch = True
607	self.use_direct_call = not current_platform.opaque_attention_op()
608
609	vllm_config = get_current_vllm_config()
610	parallel_config = vllm_config.parallel_config
611	self.use_pcp = parallel_config.prefill_context_parallel_size > 1
612	self.pcp_shard_decode_requests = parallel_config.pcp_shard_decode_requests
613	compilation_config = vllm_config.compilation_config
614	if prefix in compilation_config.static_forward_context:
615	raise ValueError(f"Duplicate layer name: {prefix}")
616	compilation_config.static_forward_context[prefix] = self
617
618	self.prefill_backend: MLAPrefillBackend \| None
619	if self.impl.is_sparse and not (
620	self.impl.supports_dense_mha_prefill and self.supports_dense_mha_prefill
621	):
622	logger.warning_once(
623	"Sparse MLA layer has no dense-MHA prefill path; using the top-k "
624	"MQA path only."
625	)
626	self.prefill_backend = None
627	else:
628	try:
629	prefill_backend_cls = prefill_backend_cls or get_mla_prefill_backend(
630	vllm_config
631	)
632	except ValueError:
633	if (
634	not self.impl.is_sparse
635	or vllm_config.attention_config.mla_prefill_backend is not None
636	):
637	raise
638	logger.warning_once(
639	"No MLA prefill backend supports this model; sparse MLA will "
640	"use the top-k MQA path only (no dense-MHA prefill)."
641	)
642	self.prefill_backend = None
643	else:
644	self.prefill_backend = prefill_backend_cls(
645	num_heads=self.num_heads,
646	scale=self.scale,
647	kv_lora_rank=self.kv_lora_rank,
648	qk_nope_head_dim=self.qk_nope_head_dim,
649	qk_rope_head_dim=self.qk_rope_head_dim,
650	v_head_dim=self.v_head_dim,
651	vllm_config=vllm_config,
652	)
653
654	self.kv_cache = torch.tensor([])
655
656	self.use_sparse = use_sparse
657
658	if vllm_config.kernel_config.enable_jit_warmup:
659	backend_name = self.attn_backend.get_name()
660	if backend_name in (
661	"FLASHMLA_SPARSE",
662	"FLASHINFER_MLA_SPARSE",
663	"FLASHINFER_MLA_SPARSE_SM120",
664	"DEEPSEEK_V32_INDEXER",
665	):
666	from vllm.v1.attention.backends.mla.compressor_utils import (
667	_COMPRESSED_SLOT_MAPPING_KERNEL,
668	)
669	from vllm.v1.attention.backends.mla.indexer import (
670	_BUILD_PREFILL_CHUNK_METADATA_KERNEL,
671	_PREPARE_UNIFORM_DECODE_KERNEL,
672	)
673
674	_COMPRESSED_SLOT_MAPPING_KERNEL.register_warmup()
675	_PREPARE_UNIFORM_DECODE_KERNEL.register_warmup()
676	_BUILD_PREFILL_CHUNK_METADATA_KERNEL.register_warmup()
677
678	if backend_name != "DEEPSEEK_V32_INDEXER":
679	from vllm.v1.attention.backends.mla.sparse_swa import (
680	_COMPUTE_PREFILL_METADATA_KERNEL,
681	)
682
683	_COMPUTE_PREFILL_METADATA_KERNEL.register_warmup()
684
685	if self.use_pcp and self.impl.dcp_world_size > 1 and not self.supports_pcp_dcp:
686	raise NotImplementedError(
687	f"{type(self).__name__} does not support PCP+DCP."
688	)
689	self.dcp_manager: MLADCPManager \| None = None
690	if self.impl.dcp_world_size > 1:
691	query_dtype = (
692	current_platform.fp8_dtype()
693	if is_quantized_kv_cache(self.kv_cache_dtype)
694	and self.kv_cache_dtype != "fp8_ds_mla"
695	and self.impl.supports_quant_query_input
696	else dtype
697	)
698	self.dcp_manager = MLADCPManager(
699	vllm_config=vllm_config,
700	device=next(kv_b_proj.parameters()).device,
701	num_heads=self.num_heads,
702	query_head_dim=self.kv_lora_rank + self.qk_rope_head_dim,
703	output_head_dim=self.kv_lora_rank,
704	query_dtype=query_dtype,
705	output_dtype=dtype,
706	padded_num_heads=self.q_pad_num_heads,
707	is_lse_base_on_e=self.impl.lse_base_on_e,
708	use_pcp=self.use_pcp,
709	)
710
711	self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()
712
713	# If kv_b_proj_weight is unquantized, quantize it to mxfp4 if supported
714	self.is_aiter_triton_fp4_bmm_enabled = (
715	rocm_aiter_ops.is_fp4bmm_enabled()
716	and hasattr(self.kv_b_proj, "weight")
717	and self.kv_b_proj.weight.dtype == torch.bfloat16
718	)
719
720	# Attributes for forward_impl method
721	self._vllm_config = get_current_vllm_config()
722	self._chunked_prefill_workspace_size: int \| None = None
723	self._decode_concat_quant_fp8_op = _DecodeConcatQuantFP8(
724	static=True,
725	group_shape=GroupShape.PER_TENSOR,
726	compile_native=True,
727	)
728	self._quant_fp8_op = QuantFP8(
729	static=True,
730	group_shape=GroupShape.PER_TENSOR,
731	compile_native=True,
732	)
733
734	def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
735	# [B, H=1, N, C] -> [B, N, C]
736	self.kv_cache = kv_cache.squeeze(1)
737	if (
738	self._vllm_config.kernel_config.enable_jit_warmup
739	and self._uses_flat_kv_cache()
740	):
741	from vllm.v1.attention.backends.mla.sparse_utils import (
742	_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL,
743	)
744
745	row_width = self.kv_cache.shape[-1]
746	assert self.kv_cache.stride(0) % row_width == 0
747	_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL.register_warmup(
748	self._vllm_config,
749	block_stride_rows=self.kv_cache.stride(0) // row_width,
750	)
751
752	def _uses_flat_kv_cache(self) -> bool:
753	backend = self.attn_backend.get_name()
754	return backend == "FLASHINFER_MLA_SPARSE" or (
755	backend == "FLASHMLA_SPARSE"
756	and self.kv_cache_dtype not in ("fp8_ds_mla", "nvfp4_ds_mla")
757	)
758
759	@property
760	def chunked_prefill_workspace_size(self) -> int:
761	if self._chunked_prefill_workspace_size is None:
762	self._chunked_prefill_workspace_size = (
763	MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(
764	self._vllm_config
765	)
766	)
767	return self._chunked_prefill_workspace_size
768
769	def update_kv_cache(
770	self,
771	kv_c_normed: torch.Tensor,
772	k_pe: torch.Tensor,
773	kv_cache: torch.Tensor,
774	slot_mapping: torch.Tensor \| None,
775	attn_metadata: "MLACommonMetadata \| None",
776	kv_cache_dtype: str,
777	k_scale: torch.Tensor,
778	) -> None:
779	cache = self.hisparse_cache
780	if slot_mapping is None or (cache is not None and cache.dummy_batch):
781	return
782	kv_c_normed, k_pe, slot_mapping = maybe_gather_mla_latent_cache_inputs(
783	kv_c_normed,
784	k_pe,
785	slot_mapping,
786	attn_metadata.num_decode_tokens if attn_metadata is not None else None,
787	self.use_pcp,
788	pcp_shard_decode_requests=self.pcp_shard_decode_requests,
789	)
790	assert slot_mapping is not None
791	if cache is not None:
792	kv_cache, slot_mapping, num_rows = cache.write_target(
793	kv_c_normed.shape[0], slot_mapping.numel()
794	)
795	kv_c_normed = kv_c_normed[:num_rows]
796	k_pe = k_pe[:num_rows]
797	self.impl.do_kv_cache_update( # type: ignore[attr-defined]
798	kv_c_normed,
799	k_pe,
800	kv_cache,
801	slot_mapping,
802	kv_cache_dtype,
803	k_scale,
804	)
805	if cache is not None:
806	mirror_target = cache.mirror_write_target(kv_c_normed.shape[0])
807	if mirror_target is not None:
808	mirror_cache, mirror_slots = mirror_target
809	self.impl.do_kv_cache_update( # type: ignore[attr-defined]
810	kv_c_normed,
811	k_pe,
812	mirror_cache,
813	mirror_slots,
814	kv_cache_dtype,
815	k_scale,
816	)
817
818	def prepare_kv_cache_update(
819	self, attn_metadata: "MLACommonMetadata \| None"
820	) -> None:
821	cache = self.hisparse_cache
822	if cache is not None and cache.runtime.is_group_leader:
823	cache.prepare_group_for_batch(attn_metadata)
824
825	def forward(
826	self,
827	q: torch.Tensor,
828	kv_c_normed: torch.Tensor,
829	k_pe: torch.Tensor,
830	output_shape: torch.Size \| None = None,
831	q_dcp_replicated: torch.Tensor \| None = None,
832	) -> torch.Tensor:
833	if self.use_direct_call:
834	forward_context: ForwardContext = get_forward_context()
835	attn_metadata_raw = forward_context.attn_metadata
836	attn_metadata: MLACommonMetadata
837	if isinstance(attn_metadata_raw, dict):
838	attn_metadata = attn_metadata_raw[self.layer_name] # type: ignore[assignment]
839	elif isinstance(attn_metadata_raw, list):
840	# list[dict[str, AttentionMetadata]]: used in speculative decoding
841	# where [0] is the base-model (non-speculative) metadata dict.
842	attn_metadata = attn_metadata_raw[0][self.layer_name] # type: ignore[assignment]
843	else:
844	attn_metadata = attn_metadata_raw
845	self_kv_cache = self.kv_cache
846	slot_mapping = forward_context.slot_mapping
847
848	assert isinstance(slot_mapping, dict), (
849	f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
850	)
851	layer_slot_mapping = slot_mapping.get(self.layer_name)
852	self.prepare_kv_cache_update(attn_metadata)
853	self.update_kv_cache(
854	kv_c_normed,
855	k_pe,
856	self_kv_cache,
857	layer_slot_mapping,
858	attn_metadata,
859	self.kv_cache_dtype,
860	self._k_scale,
861	)
862	if self.hisparse_cache is not None:
863	self.hisparse_cache.finish_kv_update()
864	output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
865	self.forward_impl(
866	q,
867	kv_c_normed,
868	k_pe,
869	self_kv_cache,
870	attn_metadata,
871	output=output,
872	q_dcp_replicated=q_dcp_replicated,
873	)
874	return output
875	else:
876	encoded = _encode_layer_name(self.layer_name)
877	kv_cache_dummy_dep = torch.ops.vllm.unified_mla_kv_cache_update(
878	kv_c_normed,
879	k_pe,
880	encoded,
881	self.kv_cache_dtype,
882	self._k_scale,
883	)
884	output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
885	torch.ops.vllm.unified_mla_attention_with_output(
886	q,
887	kv_c_normed,
888	k_pe,
889	output,
890	encoded,
891	kv_cache_dummy_dep=kv_cache_dummy_dep,
892	q_dcp_replicated=q_dcp_replicated,
893	)
894	return output
895
896	def forward_impl(
897	self,
898	q: torch.Tensor,
899	k_c_normed: torch.Tensor, # key in unified attn
900	k_pe: torch.Tensor, # value in unified attn
901	kv_cache: torch.Tensor,
902	attn_metadata: "MLACommonMetadata",
903	output: torch.Tensor,
904	output_scale: torch.Tensor \| None = None,
905	output_block_scale: torch.Tensor \| None = None,
906	quant_group_size: int \| None = None,
907	quant_scale_ue8m0: bool \| None = None,
908	quant_col_major: bool \| None = None,
909	quant_tma_aligned: bool \| None = None,
910	q_dcp_replicated: torch.Tensor \| None = None,
911	) -> torch.Tensor:
912	assert output is not None, "Output tensor must be provided."
913
914	quant_key = _detect_output_quant_key(
915	output, output_scale, output_block_scale, self.num_heads * self.v_head_dim
916	)
917	if quant_key is not None:
918	# The fusion pass has allocated output with quantized dtype
919	# (FP8 or uint8 for FP4). We can't write into it directly,
920	# so we swap in a temp buffer for computation, then quantize
921	# into the real output at the end.
922	# NOTE(carlyou): this is temporary until kernels support fp8 output
923	quant_output = output
924	output = torch.empty(
925	output.shape[0],
926	self.num_heads * self.v_head_dim,
927	dtype=q.dtype,
928	device=output.device,
929	)
930
931	if attn_metadata is None:
932	# During the profile run try to simulate to worse case output size
933	# for `self.kv_b_proj(kv_c_normed)` in `_compute_prefill_context`
934	# since this can be large
935	_ = torch.empty(
936	(
937	self.chunked_prefill_workspace_size,
938	self.num_heads,
939	self.qk_nope_head_dim + self.v_head_dim,
940	),
941	device=k_c_normed.device,
942	dtype=k_c_normed.dtype,
943	)
944
945	# The zero fill is required when used with DP + EP
946	# to ensure all ranks within a DP group compute the
947	# same expert outputs.
948	if quant_key is not None:
949	return quant_output.fill_(0)
950	return output.fill_(0)
951
952	fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype)
953
954	num_actual_toks = attn_metadata.num_actual_tokens
955	if self.use_pcp and self.impl.dcp_world_size > 1 and quant_key is not None:
956	raise NotImplementedError(
957	"MRV2 MLA PCP+DCP does not support fused output quantization yet."
958	)
959
960	# Inputs and outputs may be padded for CUDA graphs
961	output_padded = output
962	output = output[:num_actual_toks, ...]
963	q = q[:num_actual_toks, ...]
964	if q_dcp_replicated is not None:
965	q_dcp_replicated = q_dcp_replicated[:num_actual_toks, ...]
966	k_c_normed = k_c_normed[:num_actual_toks, ...]
967	k_pe = k_pe[:num_actual_toks, ...]
968
969	if fp8_attention and self.kv_cache_dtype not in (
970	# Opaque per-token byte formats stay as raw uint8
971	"fp8_ds_mla",
972	"nvfp4_ds_mla",
973	):
974	kv_cache = kv_cache.view(current_platform.fp8_dtype())
975
976	assert (
977	attn_metadata.num_decodes is not None
978	and attn_metadata.num_prefills is not None
979	and attn_metadata.num_decode_tokens is not None
980	)
981	num_mqa_tokens = attn_metadata.num_decode_tokens
982	num_mha_tokens = q.size(0) - num_mqa_tokens
983	use_mha = True
984
985	if self.impl.is_sparse and num_mha_tokens > 0:
986	use_mha = self._use_sparse_mha(attn_metadata)
987	if not use_mha:
988	num_mqa_tokens = q.size(0)
989	num_mha_tokens = 0
990
991	mha_use_quant_output = (
992	quant_key is not None
993	and self.prefill_backend is not None
994	and self.prefill_backend.supports_quant_output(quant_key)
995	and (
996	not self.impl.is_sparse
997	or attn_metadata.prefill_max_seq_len # type: ignore[attr-defined]
998	<= attn_metadata.topk_tokens # type: ignore[attr-defined]
999	)
1000	and attn_metadata is not None
1001	and attn_metadata.prefill is not None
1002	and attn_metadata.prefill.chunked_context is None
1003	and self.impl.dcp_world_size <= 1
1004	)
1005
1006	if num_mha_tokens > 0:
1007	if mha_use_quant_output:
1008	mha_output = quant_output
1009	mha_output_scale = output_scale
1010	else:
1011	mha_output = output
1012	mha_output_scale = None
1013
1014	self.impl.forward_mha( # type: ignore[attr-defined]
1015	q[num_mqa_tokens:],
1016	k_c_normed[num_mqa_tokens:],
1017	k_pe[num_mqa_tokens:],
1018	kv_cache,
1019	attn_metadata,
1020	self._k_scale,
1021	output=mha_output[num_mqa_tokens:num_actual_toks],
1022	output_scale=mha_output_scale,
1023	)
1024
1025	if num_mqa_tokens > 0:
1026	if q_dcp_replicated is not None:
1027	mqa_q = q_dcp_replicated[:num_mqa_tokens]
1028	qrep_decode = True
1029	else:
1030	mqa_q = q[:num_mqa_tokens]
1031	qrep_decode = False
1032	mqa_output_slice = output[:num_mqa_tokens]
1033
1034	mqa_q_nope, mqa_q_pe = mqa_q.split(
1035	[self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
1036	)
1037
1038	# Convert from (B, N, P) to (N, B, P)
1039	mqa_q_nope = mqa_q_nope.transpose(0, 1)
1040
1041	if self.q_pad_num_heads is not None:
1042	B, N, L = mqa_q_pe.shape
1043	mqa_pe_padded = mqa_q_pe.new_empty((B, self.q_pad_num_heads, L))
1044	mqa_pe_padded.resize_((B, N, L))
1045	mqa_pe_padded.copy_(mqa_q_pe)
1046	mqa_q_pe = mqa_pe_padded
1047
1048	if self.is_aiter_triton_fp4_bmm_enabled:
1049	mqa_ql_nope = rocm_aiter_ops.batched_gemm_a16wfp4(
1050	mqa_q_nope,
1051	self.W_K,
1052	self.W_K_scale,
1053	transpose_bm=True,
1054	prequant=True,
1055	y_scale=self._q_scale if fp8_attention else None,
1056	)
1057	elif self.is_aiter_triton_fp8_bmm_enabled:
1058	# Multiply+Transpose (N, B, P)x(N, P, L)->(N, B, L)->(B, N, L)
1059	mqa_ql_nope = rocm_aiter_ops.triton_fp8_bmm(
1060	mqa_q_nope,
1061	self.W_K,
1062	self.W_K_scale,
1063	group_size=128,
1064	transpose_bm=True,
1065	)
1066	elif self.is_amx_bmm_enabled:
1067	# bmm_cpu computes out[n] = mat1[n] @ mat2[n]^T against
1068	# AMXMLAImpl's own (N, L, P) packed W_UK -- same as prefill.
1069	N, B, P = mqa_q_nope.shape
1070	L = self.kv_lora_rank
1071	mqa_ql_nope = mqa_q_nope.new_empty((N, B, L))
1072	ops.bmm_cpu(
1073	mqa_ql_nope,
1074	mqa_q_nope,
1075	self.impl._w_uk_packed, # type: ignore[attr-defined]
1076	True,
1077	self.impl._w_scale, # type: ignore[attr-defined]
1078	)
1079	mqa_ql_nope = mqa_ql_nope.transpose(0, 1)
1080	else:
1081	# Pads the head_dim if necessary (for the underlying kernel)
1082	N, B, P = mqa_q_nope.shape
1083	W_UK_T = self.W_UK_T_dcp_qrep if qrep_decode else self.W_UK_T
1084	assert W_UK_T is not None
1085	_, _, L = W_UK_T.shape
1086
1087	if self.q_pad_num_heads is not None:
1088	mqa_ql_nope = mqa_q_nope.new_empty((self.q_pad_num_heads, B, L))
1089	mqa_ql_nope.resize_((N, B, L))
1090	# Multiply (N, B, P) x (N, P, L) -> (N, B, L)
1091	torch.bmm(mqa_q_nope, W_UK_T, out=mqa_ql_nope)
1092	# Convert from (N, B, L) to (B, N, L)
1093	mqa_ql_nope = mqa_ql_nope.transpose(0, 1)
1094	else:
1095	# Write the (N, B, L) bmm result straight into a
1096	# token-major (B, N, L) buffer so the MQA query is already
1097	# contiguous; a NoPE model (qk_rope_head_dim == 0) then
1098	# needs no concat at all.
1099	mqa_ql_nope = mqa_q_nope.new_empty((B, N, L))
1100	torch.bmm(mqa_q_nope, W_UK_T, out=mqa_ql_nope.transpose(0, 1))
1101
1102	if fp8_attention and self.impl.supports_quant_query_input:
1103	assert mqa_ql_nope.shape[0] == mqa_q_pe.shape[0]
1104	assert mqa_ql_nope.shape[1] == mqa_q_pe.shape[1]
1105	mqa_q = self._decode_concat_quant_fp8_op(
1106	mqa_ql_nope, mqa_q_pe, self._q_scale
1107	)
1108	else:
1109	mqa_q = (mqa_ql_nope, mqa_q_pe)
1110	# concatenate nope + pe -> (B, N, L + P) (fp8 op above may have fused)
1111	if self.impl.dcp_world_size > 1:
1112	assert self.dcp_manager is not None
1113	if self.use_pcp:
1114	if self.impl.dcp_world_size > self.impl.pcp_world_size:
1115	if isinstance(mqa_q, tuple):
1116	mqa_q = torch.cat(mqa_q, dim=-1)
1117	mqa_q = get_tp_group().all_gather(mqa_q, dim=1)
1118	else:
1119	if isinstance(mqa_q, tuple):
1120	# concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P)
1121	mqa_q = torch.cat(mqa_q, dim=-1)
1122	if not qrep_decode:
1123	assert self.dcp_manager.query_gather is not None
1124	mqa_q = self.dcp_manager.query_gather(mqa_q)
1125
1126	# call decode attn
1127	if not self.impl.is_sparse:
1128	assert attn_metadata.decode is not None
1129	attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self) # type: ignore[attr-defined]
1130
1131	# correct dcp attn_out with lse.
1132	if self.impl.dcp_world_size > 1:
1133	assert lse is not None
1134	assert self.dcp_manager is not None
1135	decode_metadata = getattr(attn_metadata, "decode", None)
1136	if not use_mha:
1137	seq_lens = cast(torch.Tensor, attn_metadata.seq_lens) # type: ignore[attr-defined]
1138	query_start_loc = attn_metadata.query_start_loc
1139	else:
1140	seq_lens = (
1141	decode_metadata.seq_lens
1142	if decode_metadata is not None
1143	else cast(torch.Tensor, attn_metadata.seq_lens)[ # type: ignore[attr-defined]
1144	: attn_metadata.num_decodes
1145	]
1146	)
1147	query_start_loc = attn_metadata.query_start_loc[
1148	: attn_metadata.num_decodes + 1
1149	]
1150	attn_out = self.dcp_manager.combine(
1151	attn_out,
1152	lse,
1153	seq_lens=seq_lens,
1154	query_start_loc=query_start_loc,
1155	)
1156	if self.use_pcp:
1157	attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads)
1158
1159	# v_up projection
1160	self._v_up_proj(attn_out, out=mqa_output_slice)
1161
1162	if quant_key is not None:
1163	quant_idx = num_mqa_tokens if mha_use_quant_output else num_actual_toks
1164	if quant_idx == 0:
1165	return quant_output
1166	actual = output[:quant_idx]
1167	if quant_key == kNvfp4Dynamic:
1168	# NVFP4: two FP4 values packed into one uint8
1169	assert output_block_scale is not None
1170	fp4_data, fp4_scales = ops.scaled_fp4_quant(actual, output_scale)
1171	quant_output[:quant_idx].copy_(fp4_data)
1172	output_block_scale[: fp4_scales.shape[0]].copy_(fp4_scales)
1173	elif quant_key in (kFp8Dynamic128Sym, kFp8Dynamic64Sym):
1174	# Per-group FP8
1175	assert output_block_scale is not None
1176	assert quant_group_size is not None, (
1177	"Group FP8 output quant requested but "
1178	"quant_group_size not passed through custom op"
1179	)
1180	finfo = torch.finfo(_FP8_DTYPE)
1181	torch.ops._C.per_token_group_fp8_quant(
1182	actual,
1183	quant_output[:quant_idx],
1184	output_block_scale[:quant_idx],
1185	quant_group_size,
1186	1e-10, # eps
1187	finfo.min,
1188	finfo.max,
1189	quant_scale_ue8m0,
1190	quant_col_major,
1191	quant_tma_aligned,
1192	)
1193	elif quant_key == kFp8StaticTensorSym:
1194	# Static FP8 quantization
1195	fp8_data, _ = self._quant_fp8_op(actual, output_scale)
1196	quant_output[:quant_idx].copy_(fp8_data)
1197	else:
1198	raise ValueError(f"Unsupported quant_key: {quant_key}")
1199	return quant_output
1200
1201	if self.use_pcp and output_padded.shape[0] > num_actual_toks:
1202	output_padded[num_actual_toks:].zero_()
1203	return output_padded
1204
1205	def _use_sparse_mha(self, attn_metadata: "MLACommonMetadata") -> bool:
1206	if self.hisparse_cache is not None:
1207	return False
1208	prefill = attn_metadata.prefill
1209	if prefill is None:
1210	return False
1211	use_masked_mha = (
1212	self.prefill_backend is not None
1213	and self.impl.masked_mha_available # type: ignore[attr-defined]
1214	and self.impl.dcp_world_size <= 1
1215	and _use_masked_mha(
1216	backend_name=self.attn_backend.get_name(),
1217	tensor_parallel_size=(
1218	self._vllm_config.parallel_config.tensor_parallel_size
1219	),
1220	qk_head_dim=self.qk_nope_head_dim + self.qk_rope_head_dim,
1221	v_head_dim=self.v_head_dim,
1222	query_len=prefill.max_query_len,
1223	seq_len=attn_metadata.prefill_max_seq_len, # type: ignore[attr-defined]
1224	has_context=prefill.chunked_context is not None,
1225	)
1226	and self.impl.masked_mha_workspace_fits(prefill) # type: ignore[attr-defined]
1227	)
1228	return (prefill.use_dense_mha or use_masked_mha) and not (
1229	self._vllm_config.attention_config.sparse_mla_force_mqa
1230	)
1231
1232	def process_weights_after_loading(self, act_dtype: torch.dtype):
1233	# Let per-backend impls do their own weight packing first (no-op
1234	# unless overridden), mirroring Attention.process_weights_after_loading.
1235	self.impl.process_weights_after_loading(act_dtype)
1236
1237	if self.is_amx_bmm_enabled:
1238	# AMXMLAImpl already packed its own W_UK/W_UV above, for both
1239	# prefill and decode. Release the now-unused raw weight.
1240	self.kv_b_proj.weight = torch.nn.Parameter(
1241	torch.empty(0), requires_grad=False
1242	)
1243	return
1244
1245	if self.dcp_q_replicate:
1246	# qrep wired here: validate unsupported decode backends once.
1247	assert self.q_pad_num_heads in (None, self.num_heads), (
1248	"DCP query replication is unsupported on head-padding MLA "
1249	"backends (q_pad_num_heads)."
1250	)
1251	if (
1252	self.is_aiter_triton_fp4_bmm_enabled
1253	or self.is_aiter_triton_fp8_bmm_enabled
1254	):
1255	raise NotImplementedError(
1256	"DCP query replication is not implemented for the aiter "
1257	"FP4/FP8 MLA BMM paths."
1258	)
1259
1260	W_UK, W_UV = split_kv_b_proj(
1261	self.kv_b_proj,
1262	act_dtype,
1263	self.kv_lora_rank,
1264	self.num_heads,
1265	self.qk_nope_head_dim,
1266	self.v_head_dim,
1267	)
1268
1269	# If kv_b_proj_weight is unquantized, quantize it to mxfp4 if supported
1270	if self.is_aiter_triton_fp4_bmm_enabled:
1271	from vllm.model_executor.layers.quantization.quark.utils import (
1272	quark_quantize_weight_to_mxfp4,
1273	)
1274
1275	self.W_K, self.W_K_scale = quark_quantize_weight_to_mxfp4(W_UK)
1276	# Convert from (L, N, P) to (N, L, P)
1277	self.W_K = self.W_K.transpose(0, 1)
1278	self.W_K_scale = self.W_K_scale.transpose(0, 1)
1279
1280	self.W_V, self.W_V_scale = quark_quantize_weight_to_mxfp4(
1281	W_UV.permute(1, 2, 0)
1282	)
1283	elif self.is_aiter_triton_fp8_bmm_enabled:
1284	W_K = W_UK.transpose(0, 1) # 16 512 128
1285	W_V = W_UV.permute(1, 2, 0) # 16 128 512
1286	self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant(
1287	W_K, dtype=current_platform.fp8_dtype()
1288	)
1289	self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant(
1290	W_V, dtype=current_platform.fp8_dtype()
1291	)
1292
1293	# The kernel operates on non-padded inputs. Hence, pre-compiling
1294	# triton kernel to avoid runtime compilation for unseen batch sizes
1295	# Pre-compile for batch sizes 1 to 1024 to cover most use-cases.
1296	# On DS-R1, this step adds roughly 50s to the model loading time.
1297	max_batch_size = 1024 # [ToDo] Find the optimal upper limit
1298	pre_compilation_list = list(range(1, max_batch_size + 1))
1299	if is_global_first_rank():
1300	pre_compilation_list = tqdm(
1301	pre_compilation_list,
1302	desc="[Aiter Triton] Pre-compiling fp8 BMM kernel",
1303	total=max_batch_size,
1304	)
1305
1306	for m in pre_compilation_list:
1307	x = torch.empty(
1308	(self.W_K.shape[0], m, self.W_K.shape[2]),
1309	dtype=torch.bfloat16,
1310	device=self.W_K.device,
1311	)
1312	rocm_aiter_ops.triton_fp8_bmm(
1313	x, self.W_K, self.W_K_scale, group_size=128, transpose_bm=True
1314	)
1315
1316	x = torch.empty(
1317	(self.W_V.shape[0], m, self.W_V.shape[2]),
1318	dtype=torch.bfloat16,
1319	device=self.W_V.device,
1320	)
1321	rocm_aiter_ops.triton_fp8_bmm(
1322	x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True
1323	)
1324	else:
1325	# Convert from (L, N, V) to (N, L, V)
1326	replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True)
1327	# Convert from (L, N, P) to (N, P, L)
1328	replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True)
1329	if self.dcp_q_replicate:
1330	self.W_UK_T_dcp_qrep = get_dcp_group().all_gather(
1331	self.W_UK_T.contiguous(), dim=0
1332	)
1333
1334	# If we should not load quant weights, we initialize the scales to 1.0
1335	# as the default value. See [Note: Register q/k/v/prob scales in state dict]
1336	# for more details.
1337	quant_method = (
1338	resolve_quant_method(self.quant_config, self, prefix=self.layer_name)
1339	if self.quant_config
1340	else None
1341	)
1342	if not should_load_quant_weights(quant_method):
1343	set_default_quant_scales(self, register_buffer=False)
1344
1345	def get_attn_backend(self) -> type[AttentionBackend]:
1346	return self.attn_backend
1347
1348	def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
1349	kv_cache_dtype = kv_cache_dtype_str_to_dtype(
1350	self.kv_cache_dtype, vllm_config.model_config
1351	)
1352	common_kwargs = dict(
1353	block_size=vllm_config.cache_config.block_size,
1354	num_kv_heads=1,
1355	max_tp_shards=1,
1356	head_size=self.head_size,
1357	dtype=kv_cache_dtype,
1358	cache_dtype_str=self.kv_cache_dtype,
1359	kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
1360	# ds_mla layouts pack NoPE + RoPE + scales into one opaque per-token
1361	# blob, so the size is not derivable from head_size.
1362	# See flashmla_sparse.py.
1363	state_content_bytes={"fp8_ds_mla": 656, "nvfp4_ds_mla": 352}.get(
1364	self.kv_cache_dtype
1365	),
1366	)
1367	if self.sliding_window is not None:
1368	return SlidingWindowMLASpec(
1369	**common_kwargs,
1370	sliding_window=self.sliding_window,
1371	)
1372	spec = MLAAttentionSpec(
1373	**common_kwargs,
1374	is_index_group_leader=self.indexer is not None,
1375	non_causal_multi_token_decode=self.non_causal_multi_token_decode,
1376	)
1377	# SM100 FlashMLA paged kernels also express TMA coordinates in token rows.
1378	uses_tma_rows = (
1379	self.attn_backend.get_name() == "FLASHMLA_SPARSE"
1380	and self.kv_cache_dtype in ("fp8_ds_mla", "nvfp4_ds_mla")
1381	and current_platform.is_device_capability_family(100)
1382	)
1383	if self._uses_flat_kv_cache() or uses_tma_rows:
1384	spec = replace(spec, block_stride_alignment=spec.state_content_size_bytes)
1385	return spec
1386
1387	def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
1388	# Convert from (B, N, L) to (N, B, L)
1389	x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
1390	out = out.view(-1, self.num_heads, self.v_head_dim)
1391	if self.is_aiter_triton_fp4_bmm_enabled:
1392	out = rocm_aiter_ops.batched_gemm_a16wfp4(
1393	x,
1394	self.W_V,
1395	self.W_V_scale,
1396	out,
1397	transpose_bm=True,
1398	prequant=True,
1399	y_scale=None,
1400	)
1401	x = out.view(-1, self.num_heads * self.v_head_dim)
1402	elif self.is_aiter_triton_fp8_bmm_enabled:
1403	# Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)
1404	x = rocm_aiter_ops.triton_fp8_bmm(
1405	x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True, YQ=out
1406	)
1407	elif self.is_amx_bmm_enabled:
1408	# bmm_cpu computes out[n] = mat1[n] @ mat2[n]^T against
1409	# AMXMLAImpl's own (N, V, L) packed W_UV -- same as prefill.
1410	ops.bmm_cpu(
1411	out.transpose(0, 1),
1412	x,
1413	self.impl._w_uv_packed, # type: ignore[attr-defined]
1414	True,
1415	self.impl._w_scale, # type: ignore[attr-defined]
1416	)
1417	else:
1418	# Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)
1419	torch.bmm(x, self.W_UV, out=out.transpose(0, 1))
1420
1421
1422	def unified_mla_kv_cache_update(
1423	kv_c_normed: torch.Tensor,
1424	k_pe: torch.Tensor,
1425	layer_name: LayerNameType,
1426	kv_cache_dtype: str,
1427	k_scale: torch.Tensor,
1428	) -> torch.Tensor:
1429	"""Returns a dummy that is passed to unified_attention to signal a side effect and
1430	the data dependency between them to ensure torch.compile preserves ordering.
1431	"""
1432	layer_name = _resolve_layer_name(layer_name)
1433	attn_metadata, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(
1434	layer_name
1435	)
1436	attn_layer.prepare_kv_cache_update(attn_metadata)
1437	if layer_slot_mapping is not None:
1438	attn_layer.update_kv_cache(
1439	kv_c_normed,
1440	k_pe,
1441	kv_cache,
1442	layer_slot_mapping,
1443	attn_metadata,
1444	kv_cache_dtype,
1445	k_scale,
1446	)
1447
1448	return torch.empty(0, device=kv_c_normed.device, dtype=kv_c_normed.dtype)
1449
1450
1451	def unified_mla_kv_cache_update_fake(
1452	kv_c_normed: torch.Tensor,
1453	k_pe: torch.Tensor,
1454	layer_name: LayerNameType,
1455	kv_cache_dtype: str,
1456	k_scale: torch.Tensor,
1457	) -> torch.Tensor:
1458	return torch.empty(0, device=kv_c_normed.device, dtype=kv_c_normed.dtype)
1459
1460
1461	direct_register_custom_op(
1462	op_name="unified_mla_kv_cache_update",
1463	op_func=unified_mla_kv_cache_update,
1464	fake_impl=unified_mla_kv_cache_update_fake,
1465	)
1466
1467
1468	@eager_break_during_capture
1469	@maybe_transfer_kv_layer
1470	def unified_mla_attention_with_output(
1471	q: torch.Tensor,
1472	kv_c_normed: torch.Tensor,
1473	k_pe: torch.Tensor,
1474	output: torch.Tensor,
1475	layer_name: LayerNameType,
1476	output_scale: torch.Tensor \| None = None,
1477	output_block_scale: torch.Tensor \| None = None,
1478	kv_cache_dummy_dep: torch.Tensor \| None = None,
1479	quant_group_size: int \| None = None,
1480	quant_scale_ue8m0: bool \| None = None,
1481	quant_col_major: bool \| None = None,
1482	quant_tma_aligned: bool \| None = None,
1483	q_dcp_replicated: torch.Tensor \| None = None,
1484	) -> None:
1485	# kv_cache_dummy_dep is not used but accepting it creates a data dependency
1486	# that ensures torch.compile preserves ordering between KV cache update and
1487	# attention forward.
1488	del kv_cache_dummy_dep
1489	layer_name = _resolve_layer_name(layer_name)
1490	attn_metadata, layer, kv_cache, _ = get_attention_context(layer_name)
1491	if layer.hisparse_cache is not None:
1492	layer.hisparse_cache.finish_kv_update()
1493	layer.forward_impl(
1494	q,
1495	kv_c_normed,
1496	k_pe,
1497	kv_cache,
1498	attn_metadata,
1499	output=output,
1500	output_scale=output_scale,
1501	output_block_scale=output_block_scale,
1502	quant_group_size=quant_group_size,
1503	quant_scale_ue8m0=quant_scale_ue8m0,
1504	quant_col_major=quant_col_major,
1505	quant_tma_aligned=quant_tma_aligned,
1506	q_dcp_replicated=q_dcp_replicated,
1507	)
1508
1509
1510	direct_register_custom_op(
1511	op_name="unified_mla_attention_with_output",
1512	op_func=unified_mla_attention_with_output,
1513	mutates_args=["output", "output_block_scale"],
1514	dispatch_key=current_platform.dispatch_key,
1515	tags=(torch.Tag.flexible_layout,),
1516	)
1517
1518
1519	class QueryLenSupport(Enum):
1520	"""Defines the level of query length support for an attention backend's
1521	decode pipeline.
1522
1523	- SINGLE_ONLY: Decode pipeline only supports single-token queries
1524	(query_len=1)
1525	- UNIFORM: Decode pipeline supports uniform multi-token queries
1526	(all requests must have same query_len > 1)
1527	- VARLEN: Decode pipeline supports variable-length queries
1528	(mixed query lengths in same batch)
1529	"""
1530
1531	SINGLE_ONLY = "single_only"
1532	UNIFORM = "uniform"
1533	VARLEN = "varlen"
1534
1535
1536	def dynamic_per_batched_tensor_quant(
1537	x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn
1538	):
1539	DTYPE_MAX = torch.finfo(dtype).max
1540	min_val, max_val = x.aminmax()
1541	amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-10)
1542	scale = DTYPE_MAX / amax
1543	x_scl_sat = (x * scale).clamp(min=-DTYPE_MAX, max=DTYPE_MAX)
1544	return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal()
1545
1546
1547	@CustomOp.register(
1548	"mla_decode_concat_quant_fp8",
1549	dynamic_arg_dims={"decode_ql_nope": 0, "decode_q_pe": 0},
1550	)
1551	class _DecodeConcatQuantFP8(QuantFP8):
1552	"""QuantFP8 variant that concatenates decode_ql_nope and decode_q_pe before
1553	quantization. When disabled, forward_native is compiled via torch.compile,
1554	fusing cat/reshape/quant/view together.
1555	"""
1556
1557	def _make_forward(quant_fn): # noqa: N805
1558	"""Factory to create forward methods that concat before quantization."""
1559
1560	def forward(
1561	self,
1562	decode_ql_nope: torch.Tensor,
1563	decode_q_pe: torch.Tensor,
1564	scale: torch.Tensor,
1565	scale_ub: torch.Tensor \| None = None,
1566	) -> torch.Tensor:
1567	if decode_q_pe.shape[-1] == 0 and decode_ql_nope.is_contiguous():
1568	decode_q0 = decode_ql_nope
1569	else:
1570	decode_q0 = torch.cat((decode_ql_nope, decode_q_pe), dim=-1)
1571	decode_q_flat = decode_q0.reshape(decode_q0.shape[0], -1)
1572	decode_q, _ = quant_fn(self, decode_q_flat, scale, scale_ub)
1573	return decode_q.view(decode_q0.shape)
1574
1575	return forward
1576
1577	forward_native = _make_forward(QuantFP8.forward_native) # type: ignore[arg-type]
1578	forward_cuda = _make_forward(QuantFP8.forward_cuda) # type: ignore[arg-type]
1579	forward_hip = _make_forward(QuantFP8.forward_hip) # type: ignore[arg-type]
1580
1581
1582	class MLACommonBackend(AttentionBackend):
1583	@classmethod
1584	def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec":
1585	"""Per-token-head modes pack an inline fp32 scale pair after the
1586	latent data (single-sided: ``head_size_v == 0`` for MLA)."""
1587	spec = super().customize_spec(spec)
1588	mode = spec.kv_quant_mode
1589	if spec.state_content_bytes is not None or not mode.is_per_token_head:
1590	return spec
1591	head_size = spec.head_size
1592	if mode == KVQuantMode.INT4_PER_TOKEN_HEAD:
1593	head_size //= 2
1594	scale_bytes = get_dtype_size(torch.float32)
1595	content = head_size * get_dtype_size(spec.dtype) + 2 * scale_bytes
1596	return replace(spec, state_content_bytes=content)
1597
1598	@staticmethod
1599	def get_name() -> str:
1600	return "TRITON_MLA"
1601
1602	@staticmethod
1603	def get_builder_cls() -> type["MLACommonMetadataBuilder"]:
1604	return MLACommonMetadataBuilder
1605
1606	@classmethod
1607	def get_supported_head_sizes(cls) -> list[int]:
1608	return [320, 512, 576]
1609
1610	@classmethod
1611	def is_mla(cls) -> bool:
1612	return True
1613
1614
1615	@dataclass
1616	class MLACommonPrefillMetadata:
1617	"""Prefill Specific Metadata"""
1618
1619	@dataclass
1620	class ContextChunk:
1621	"""One workspace-sized slice of paged context for a run of prefills."""
1622
1623	index: int
1624	request_slice: slice
1625	token_slice: slice
1626	continuation_token_end: int
1627	is_continuation: bool
1628	num_context_tokens: int
1629	query_start_loc: torch.Tensor
1630	max_query_len: int
1631	cu_seq_lens: torch.Tensor
1632	starts: torch.Tensor
1633	max_seq_len: int
1634	seq_lens: torch.Tensor
1635	token_to_seq: torch.Tensor
1636	all_rows_active: bool
1637
1638	# for mla DCP
1639	padded_local_seq_lens: list[int] \| None = None
1640	local_context_lens_allranks: list[list[int]] \| None = None
1641	padded_local_cu_seq_lens: torch.Tensor \| None = None
1642	padded_local_token_to_seq: torch.Tensor \| None = None
1643	num_local_context_tokens: int = 0
1644	local_starts: list[int] \| None = None
1645
1646	@property
1647	def num_requests(self) -> int:
1648	return self.request_slice.stop - self.request_slice.start
1649
1650	@dataclass
1651	class ChunkedContextMetadata:
1652	context_lens: torch.Tensor
1653	workspace: torch.Tensor
1654	chunks: "list[MLACommonPrefillMetadata.ContextChunk]"
1655	context_lens_list: list[int]
1656	empty_token_slices: list[slice]
1657	dcp_manager: MLADCPManager \| None = None
1658
1659	block_table: torch.Tensor
1660	query_start_loc: torch.Tensor
1661	max_query_len: int
1662	chunked_context: ChunkedContextMetadata \| None = None
1663	q_data_type: torch.dtype \| None = None
1664	output_dtype: torch.dtype \| None = None
1665	prefill_backend: MLAPrefillBackend \| None = None
1666	query_lens_cpu: torch.Tensor \| None = None
1667	use_dense_mha: bool = False
1668	topk_mask_workspace: torch.Tensor \| None = None
1669
1670
1671	@dataclass
1672	class MLACommonDecodeMetadata:
1673	block_table: torch.Tensor
1674	seq_lens: torch.Tensor
1675	dcp_tot_seq_lens: torch.Tensor \| None
1676
1677
1678	D = TypeVar("D", bound=MLACommonDecodeMetadata)
1679
1680
1681	@dataclass
1682	class MLACommonMetadata(AttentionMetadata, Generic[D]):
1683	"""Metadata for MLACommon.
1684
1685	NOTE: Please read the comment at the top of the file before trying to
1686	understand this class
1687	"""
1688
1689	# NOTE(sang): Definition of context_len, query_len, and seq_len.
1690	# \|---------- N-1 iteration --------\|
1691	# \|---------------- N iteration ---------------------\|
1692	# \|- tokenA -\|......................\|-- newTokens ---\|
1693	# \|---------- context_len ----------\|
1694	# \|-------------------- seq_len ---------------------\|
1695	# \|-- query_len ---\|
1696
1697	num_reqs: int
1698	max_query_len: int
1699	max_seq_len: int
1700
1701	num_actual_tokens: int # Number of tokens excluding padding.
1702	query_start_loc: torch.Tensor
1703	slot_mapping: torch.Tensor
1704
1705	# New for MLA (compared to FlashAttention)
1706	# For handling prefill decode split
1707	num_decodes: int
1708	num_decode_tokens: int
1709	num_prefills: int
1710
1711	causal: bool = True
1712
1713	# The dimension of the attention heads
1714	head_dim: int \| None = None
1715
1716	prefill: MLACommonPrefillMetadata \| None = None
1717	decode: D \| None = None
1718
1719	def __post_init__(self):
1720	if self.head_dim is not None and not MLACommonBackend.supports_head_size(
1721	self.head_dim
1722	):
1723	raise ValueError(f"Head dimension {self.head_dim} is not supported by MLA.")
1724
1725
1726	M = TypeVar("M", bound=MLACommonMetadata)
1727	A = TypeVar("A", bound=AttentionMetadata)
1728
1729
1730	@dataclass
1731	class MLADims:
1732	q_lora_rank: int \| None
1733	kv_lora_rank: int
1734	qk_nope_head_dim: int
1735	qk_rope_head_dim: int
1736	v_head_dim: int
1737
1738
1739	def get_mla_dims(model_config: ModelConfig) -> MLADims:
1740	hf_text_config = model_config.hf_text_config
1741
1742	# Check if this is a DeepseekV4 config (uses unified head_dim + rope_head_dim)
1743	if hasattr(hf_text_config, "compress_ratios"):
1744	# DeepseekV4 style config: unified head_dim with rope_head_dim
1745	head_dim = hf_text_config.head_dim
1746	rope_head_dim = hf_text_config.qk_rope_head_dim
1747	return MLADims(
1748	q_lora_rank=hf_text_config.q_lora_rank,
1749	kv_lora_rank=head_dim,
1750	qk_nope_head_dim=head_dim - rope_head_dim,
1751	qk_rope_head_dim=rope_head_dim,
1752	v_head_dim=head_dim,
1753	)
1754
1755	# DeepseekV2/V3 style config
1756	return MLADims(
1757	q_lora_rank=getattr(hf_text_config, "q_lora_rank", None),
1758	kv_lora_rank=hf_text_config.kv_lora_rank,
1759	qk_nope_head_dim=hf_text_config.qk_nope_head_dim,
1760	qk_rope_head_dim=hf_text_config.qk_rope_head_dim,
1761	v_head_dim=hf_text_config.v_head_dim,
1762	)
1763
1764
1765	_MaskedMHARanges = tuple[tuple[int, int], ...]
1766
1767	_MASKED_MHA_THRESHOLDS: dict[
1768	tuple[int, int],
1769	dict[str, dict[int, tuple[_MaskedMHARanges, _MaskedMHARanges \| None]]],
1770	] = {
1771	(192, 128): {
1772	"FLASHMLA_SPARSE": {
1773	1: (((2048, 1536), (4096, 4096)), None),
1774	2: (((2048, 512), (4096, 1024), (8192, 4096)), None),
1775	4: (
1776	((2048, 512), (4096, 1024), (8192, 1536), (16384, 8192)),
1777	None,
1778	),
1779	8: (
1780	(
1781	(2048, 512),
1782	(4096, 512),
1783	(8192, 1024),
1784	(16384, 2048),
1785	(32768, 16384),
1786	),
1787	None,
1788	),
1789	},
1790	"FLASHINFER_MLA_SPARSE": {
1791	8: (
1792	(
1793	(2048, 512),
1794	(4096, 1024),
1795	(8192, 1024),
1796	(16384, 2048),
1797	(32768, 32768),
1798	),
1799	None,
1800	),
1801	},
1802	},
1803	(256, 256): {
1804	"FLASHMLA_SPARSE": {
1805	1: (((8 * 1024, 0),), ()),
1806	2: (
1807	((20 * 1024, 0),),
1808	((6 * 1024, 4 * 1024), (10 * 1024, 8 * 1024)),
1809	),
1810	4: (
1811	((48 * 1024, 0),),
1812	(
1813	(8 * 1024, 4 * 1024),
1814	(16 * 1024, 8 * 1024),
1815	(24 * 1024, 16 * 1024),
1816	(34 * 1024, 32 * 1024),
1817	),
1818	),
1819	8: (
1820	((112 * 1024, 0),),
1821	(
1822	(8 * 1024, 4 * 1024),
1823	(16 * 1024, 8 * 1024),
1824	(32 * 1024, 16 * 1024),
1825	(48 * 1024, 32 * 1024),
1826	),
1827	),
1828	},
1829	"FLASHINFER_MLA_SPARSE": {
1830	4: (
1831	((36 * 1024, 0),),
1832	(
1833	(12 * 1024, 4 * 1024),
1834	(16 * 1024, 8 * 1024),
1835	(24 * 1024, 16 * 1024),
1836	(28 * 1024, 24 * 1024),
1837	),
1838	),
1839	8: (
1840	((64 * 1024, 0),),
1841	(
1842	(20 * 1024, 4 * 1024),
1843	(24 * 1024, 8 * 1024),
1844	(32 * 1024, 16 * 1024),
1845	(48 * 1024, 32 * 1024),
1846	(56 * 1024, 40 * 1024),
1847	(64 * 1024, 48 * 1024),
1848	(72 * 1024, 56 * 1024),
1849	),
1850	),
1851	},
1852	},
1853	}
1854
1855
1856	def _use_masked_mha(
1857	*,
1858	backend_name: str,
1859	tensor_parallel_size: int,
1860	qk_head_dim: int,
1861	v_head_dim: int,
1862	query_len: int,
1863	seq_len: int,
1864	has_context: bool,
1865	) -> bool:
1866	thresholds = (
1867	_MASKED_MHA_THRESHOLDS.get((qk_head_dim, v_head_dim), {})
1868	.get(backend_name, {})
1869	.get(tensor_parallel_size)
1870	)
1871	if thresholds is None:
1872	return False
1873	pure_prefill_ranges, context_ranges = thresholds
1874	ranges = pure_prefill_ranges
1875	if has_context and context_ranges is not None:
1876	ranges = context_ranges
1877	for max_seq_len, min_query_len in ranges:
1878	if seq_len <= max_seq_len:
1879	return query_len >= min_query_len
1880	return False
1881
1882
1883	@functools.cache
1884	def backend_supports_prefill_query_quantization() -> bool:
1885	"""Check if the selected MLA prefill backend supports query quantization.
1886
1887	Currently supported backends:
1888	- FlashInfer
1889	- TRT-LLM Ragged
1890
1891	Not supported:
1892	- FlashAttention (FA3/FA4)
1893	- Non-GB200 devices (FP8 prefill requires device capability 100)
1894	"""
1895	# FP8 prefill query quantization requires GB200 (device capability 100)
1896	# for the necessary FP8 kernels at the moment.
1897	if not current_platform.is_device_capability_family(100):
1898	return False
1899
1900	from vllm.config import get_current_vllm_config
1901	from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
1902
1903	vllm_config = get_current_vllm_config()
1904	backend_cls = get_mla_prefill_backend(vllm_config)
1905	return backend_cls.get_name() in (
1906	"FLASHINFER",
1907	"TRTLLM_RAGGED",
1908	"TOKENSPEED_MLA",
1909	)
1910
1911
1912	@dataclass
1913	class _ContextChunkPlan:
1914	"""Request-space layout of one context chunk, before tensors are built."""
1915
1916	request_start: int
1917	request_end: int
1918	# Per-request context offset and row count, in request order.
1919	starts: list[int]
1920	seq_lens: list[int]
1921	is_continuation: bool
1922
1923
1924	def plan_mla_context_chunks(
1925	context_lens: list[int],
1926	row_budget: int,
1927	max_context_chunk: int,
1928	split_alignment: int,
1929	padded_rows: Callable[[int], int],
1930	) -> list[_ContextChunkPlan]:
1931	"""Pack per-request contexts into workspace-sized chunks."""
1932	assert max_context_chunk > 0
1933	assert split_alignment > 0
1934
1935	def aligned_split_len(start: int, remaining: int, available: int) -> int:
1936	max_split = min(
1937	max_context_chunk,
1938	round_down(remaining - 1, split_alignment),
1939	)
1940	low, high = 0, max_split // split_alignment
1941	while low < high:
1942	mid = (low + high + 1) // 2
1943	length = mid * split_alignment
1944	rows = padded_rows(start + length) - padded_rows(start)
1945	if rows <= available:
1946	low = mid
1947	else:
1948	high = mid - 1
1949	return low * split_alignment
1950
1951	plans: list[_ContextChunkPlan] = []
1952	num_requests = len(context_lens)
1953	request = 0
1954	start = 0
1955	while request < num_requests:
1956	context_len = context_lens[request]
1957	if context_len == 0:
1958	assert start == 0
1959	request += 1
1960	continue
1961
1962	request_start = request
1963	starts: list[int] = []
1964	seq_lens: list[int] = []
1965	rows = 0
1966	while request < num_requests and context_lens[request] > 0:
1967	remaining = context_lens[request] - start
1968	request_rows = padded_rows(start + remaining) - padded_rows(start)
1969	if rows + request_rows > row_budget:
1970	split_len = aligned_split_len(start, remaining, row_budget - rows)
1971	if split_len > 0:
1972	starts.append(start)
1973	seq_lens.append(split_len)
1974	rows += padded_rows(start + split_len) - padded_rows(start)
1975	start += split_len
1976	break
1977	rows += request_rows
1978	starts.append(start)
1979	seq_lens.append(remaining)
1980	request += 1
1981	start = 0
1982	assert seq_lens
1983	plans.append(
1984	_ContextChunkPlan(
1985	request_start=request_start,
1986	request_end=request_start + len(seq_lens),
1987	starts=starts,
1988	seq_lens=seq_lens,
1989	is_continuation=starts[0] > 0,
1990	)
1991	)
1992	return plans
1993
1994
1995	def align_mla_chunked_context_workspace_size(
1996	vllm_config: VllmConfig,
1997	workspace_size: int,
1998	) -> int:
1999	parallel_config = vllm_config.parallel_config
2000	alignment = vllm_config.cache_config.block_size
2001	if parallel_config.decode_context_parallel_size > 1:
2002	alignment = lcm(
2003	alignment,
2004	parallel_config.decode_context_parallel_size
2005	* parallel_config.cp_kv_cache_interleave_size,
2006	)
2007	return round_up(max(workspace_size, alignment), alignment)
2008
2009
2010	def build_mla_chunked_context_metadata(
2011	*,
2012	context_lens_cpu: torch.Tensor,
2013	prefill_query_start_loc_cpu: torch.Tensor,
2014	chunked_prefill_workspace: torch.Tensor,
2015	chunked_prefill_workspace_size: int,
2016	block_size: int,
2017	align_chunk_to_block: bool,
2018	device: torch.device,
2019	dcp_world_size: int,
2020	dcp_local_block_size: int,
2021	dcp_virtual_block_size: int,
2022	dcp_manager: MLADCPManager \| None = None,
2023	) -> "MLACommonPrefillMetadata.ChunkedContextMetadata \| None":
2024	"""Build chunked-context metadata for an MLA prefill.
2025
2026	Shared by dense and sparse builders. Packs the prefill contexts into a flat
2027	list of workspace-sized per-request chunks and, under DCP, plans the
2028	per-rank interleaved local chunks the all-gather reduction consumes.
2029
2030	Args:
2031	context_lens_cpu: Per-prefill context length (seq_len - query_len).
2032	prefill_query_start_loc_cpu: Prefill query cumulative offsets (0-based).
2033	chunked_prefill_workspace: Scratch buffer the context gather writes to.
2034	chunked_prefill_workspace_size: Row capacity of the workspace.
2035	block_size: KV cache page size for chunk-start alignment.
2036	align_chunk_to_block: Round the chunk size down to ``block_size``.
2037	device: Target device for the returned tensors.
2038	dcp_world_size: Decode-context-parallel world size (1 if disabled).
2039	dcp_local_block_size: Per-rank interleave block size for DCP.
2040	dcp_virtual_block_size: ``dcp_local_block_size * dcp_world_size``.
2041	dcp_manager: Shared MLA DCP collective manager.
2042
2043	Returns:
2044	The chunked-context metadata, or None when no prefill has any context.
2045
2046	"""
2047	# NOTE: it is recommended you read the `Chunked Prefill` section in the
2048	# comment at the top of the file before trying to understand this code.
2049	context_lens = context_lens_cpu.tolist()
2050	if max(context_lens, default=0) <= 0:
2051	return None
2052
2053	# The `gather_and_maybe_dequant_cache` kernel cannot handle chunk starts
2054	# that are not aligned to block_size, so a split request advances in
2055	# block-aligned steps.
2056	chunk_alignment = block_size if align_chunk_to_block else 1
2057	if dcp_world_size > 1:
2058	# Each rank gathers only its own shard, so a chunk's workspace cost is
2059	# 1/dcp_world_size of its context rows, rounded up to the interleave
2060	# block size. Sub-chunks are additionally aligned to the virtual block
2061	# size so that every rank's local start stays interleave-aligned.
2062	row_budget = chunked_prefill_workspace_size // dcp_world_size
2063	chunk_alignment = math.lcm(chunk_alignment, dcp_virtual_block_size)
2064
2065	def padded_rows(rows: int) -> int:
2066	return cdiv(rows, dcp_virtual_block_size) * dcp_local_block_size
2067	else:
2068	row_budget = chunked_prefill_workspace_size
2069
2070	def padded_rows(rows: int) -> int:
2071	return rows
2072
2073	max_context_chunk = round_down(chunked_prefill_workspace_size, chunk_alignment)
2074	assert max_context_chunk > 0, (
2075	f"chunked prefill workspace ({chunked_prefill_workspace_size} rows) is "
2076	f"smaller than the context chunk alignment ({chunk_alignment} rows)"
2077	)
2078
2079	plans = plan_mla_context_chunks(
2080	context_lens,
2081	row_budget,
2082	max_context_chunk,
2083	chunk_alignment,
2084	padded_rows,
2085	)
2086
2087	query_start_loc = prefill_query_start_loc_cpu.tolist()
2088	empty_token_slices = [
2089	slice(query_start_loc[i], query_start_loc[i + 1])
2090	for i, context_len in enumerate(context_lens)
2091	if context_len == 0
2092	]
2093
2094	use_dcp = dcp_world_size > 1
2095	local_context_lens_allranks = (
2096	get_dcp_local_seq_lens(
2097	context_lens_cpu, dcp_world_size, None, dcp_local_block_size
2098	).tolist()
2099	if use_dcp
2100	else None
2101	)
2102
2103	# Concatenate the per-chunk fields so each one costs a single
2104	# host-to-device transfer, and remember where every chunk's slice lands.
2105	starts_flat: list[int] = []
2106	seq_lens_flat: list[int] = []
2107	cu_seq_lens_flat: list[int] = []
2108	cu_seqlens_q_flat: list[int] = []
2109	token_to_seq_parts: list[np.ndarray] = []
2110	padded_local_cu_seq_lens_flat: list[int] = []
2111	padded_local_token_to_seq_parts: list[np.ndarray] = []
2112	local_starts_per_chunk: list[list[int]] = []
2113	local_seq_lens_per_chunk: list[list[int]] = []
2114	layouts: list[tuple[slice, slice, slice, slice]] = []
2115
2116	request_offset = boundary_offset = token_offset = local_token_offset = 0
2117	for plan in plans:
2118	num_requests = len(plan.seq_lens)
2119	num_tokens = sum(plan.seq_lens)
2120	num_local_tokens = num_tokens
2121
2122	seq_lens_flat.extend(plan.seq_lens)
2123	cu_seq_lens_flat.extend(itertools.accumulate(plan.seq_lens, initial=0))
2124	query_base = query_start_loc[plan.request_start]
2125	cu_seqlens_q_flat.extend(
2126	query_start_loc[request] - query_base
2127	for request in range(plan.request_start, plan.request_end + 1)
2128	)
2129	token_to_seq_parts.append(
2130	np.repeat(np.arange(num_requests, dtype=np.int32), plan.seq_lens)
2131	)
2132
2133	if use_dcp:
2134	# A request's local rows are its context rows sharded across ranks
2135	# and rounded up to the interleave block size, so the local start of
2136	# a continuation is the padded row count of the context before it.
2137	local_starts = [padded_rows(start) for start in plan.starts]
2138	local_seq_lens = [
2139	padded_rows(start + length) - local_start
2140	for start, length, local_start in zip(
2141	plan.starts, plan.seq_lens, local_starts
2142	)
2143	]
2144	local_starts_per_chunk.append(local_starts)
2145	local_seq_lens_per_chunk.append(local_seq_lens)
2146	padded_local_cu_seq_lens_flat.extend(
2147	itertools.accumulate(local_seq_lens, initial=0)
2148	)
2149	padded_local_token_to_seq_parts.append(
2150	np.repeat(np.arange(num_requests, dtype=np.int32), local_seq_lens)
2151	)
2152	num_local_tokens = sum(local_seq_lens)
2153	# The gather takes per-rank local offsets under DCP.
2154	starts_flat.extend(local_starts)
2155	else:
2156	starts_flat.extend(plan.starts)
2157	assert num_local_tokens <= row_budget
2158
2159	layouts.append(
2160	(
2161	slice(request_offset, request_offset + num_requests),
2162	slice(boundary_offset, boundary_offset + num_requests + 1),
2163	slice(token_offset, token_offset + num_tokens),
2164	slice(local_token_offset, local_token_offset + num_local_tokens),
2165	)
2166	)
2167	request_offset += num_requests
2168	boundary_offset += num_requests + 1
2169	token_offset += num_tokens
2170	local_token_offset += num_local_tokens
2171
2172	seq_lens_cpu = torch.tensor(
2173	seq_lens_flat, dtype=torch.int32, device="cpu", pin_memory=PIN_MEMORY
2174	)
2175	starts_and_context_lens = async_tensor_h2d(
2176	starts_flat + context_lens, device=device, dtype=torch.int32
2177	)
2178	starts = starts_and_context_lens[: len(starts_flat)]
2179	context_lens_gpu = starts_and_context_lens[len(starts_flat) :]
2180	cu_seq_lens = async_tensor_h2d(cu_seq_lens_flat, device=device, dtype=torch.int32)
2181	cu_seqlens_q = async_tensor_h2d(cu_seqlens_q_flat, device=device, dtype=torch.int32)
2182	token_to_seq = async_tensor_h2d(
2183	np.concatenate(token_to_seq_parts), device=device, dtype=torch.int32
2184	)
2185	if use_dcp:
2186	padded_local_cu_seq_lens = async_tensor_h2d(
2187	padded_local_cu_seq_lens_flat, device=device, dtype=torch.int32
2188	)
2189	padded_local_token_to_seq = async_tensor_h2d(
2190	np.concatenate(padded_local_token_to_seq_parts),
2191	device=device,
2192	dtype=torch.int32,
2193	)
2194
2195	chunks: list[MLACommonPrefillMetadata.ContextChunk] = []
2196	for index, (plan, layout) in enumerate(zip(plans, layouts)):
2197	request_slice, boundary_slice, token_slice, local_token_slice = layout
2198	query_lens = [
2199	query_start_loc[request + 1] - query_start_loc[request]
2200	for request in range(plan.request_start, plan.request_end)
2201	]
2202	chunk = MLACommonPrefillMetadata.ContextChunk(
2203	index=index,
2204	request_slice=slice(plan.request_start, plan.request_end),
2205	token_slice=slice(
2206	query_start_loc[plan.request_start], query_start_loc[plan.request_end]
2207	),
2208	continuation_token_end=query_start_loc[plan.request_start + 1],
2209	is_continuation=plan.is_continuation,
2210	num_context_tokens=token_slice.stop - token_slice.start,
2211	query_start_loc=cu_seqlens_q[boundary_slice],
2212	max_query_len=max(query_lens),
2213	cu_seq_lens=cu_seq_lens[boundary_slice],
2214	starts=starts[request_slice],
2215	max_seq_len=max(plan.seq_lens),
2216	seq_lens=seq_lens_cpu[request_slice],
2217	token_to_seq=token_to_seq[token_slice],
2218	all_rows_active=all(length > 0 for length in query_lens)
2219	and all(length > 0 for length in plan.seq_lens),
2220	num_local_context_tokens=local_token_slice.stop - local_token_slice.start,
2221	)
2222	if use_dcp:
2223	assert local_context_lens_allranks is not None
2224	chunk.padded_local_seq_lens = local_seq_lens_per_chunk[index]
2225	chunk.local_context_lens_allranks = local_context_lens_allranks[
2226	plan.request_start : plan.request_end
2227	]
2228	chunk.padded_local_cu_seq_lens = padded_local_cu_seq_lens[boundary_slice]
2229	chunk.padded_local_token_to_seq = padded_local_token_to_seq[
2230	local_token_slice
2231	]
2232	chunk.local_starts = local_starts_per_chunk[index]
2233	chunks.append(chunk)
2234
2235	return MLACommonPrefillMetadata.ChunkedContextMetadata(
2236	context_lens=context_lens_gpu,
2237	workspace=chunked_prefill_workspace,
2238	chunks=chunks,
2239	context_lens_list=context_lens,
2240	empty_token_slices=empty_token_slices,
2241	dcp_manager=dcp_manager,
2242	)
2243
2244
2245	class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
2246	"""NOTE: Please read the comment at the top of the file before trying to
2247	understand this class
2248	"""
2249
2250	kv_cache_spec: AttentionSpec
2251
2252	# Defines the level of query length support for this backend.
2253	# - SINGLE_ONLY: Only single-token queries (no spec decode support)
2254	# - UNIFORM: Supports uniform multi-token queries (spec decode with uniform lengths)
2255	# - VARLEN: Supports variable-length queries (spec decode with mixed lengths)
2256	# If set to UNIFORM or VARLEN, this will increase `reorder_batch_threshold` when
2257	# speculative decoding is enabled.
2258	query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.SINGLE_ONLY
2259
2260	# Whether this builder can flatten a non-causal query block into decode rows.
2261	supports_non_causal_multi_token_decode: ClassVar[bool] = False
2262
2263	# Whether can support non-causal multi-token decode with DCP KV cache.
2264	supports_non_causal_multi_token_dcp: ClassVar[bool] = False
2265
2266	# The threshold for reordering the batch into decode and prefill requests.
2267	# If > 1, the batch will be reordered such that requests with
2268	# query length <= threshold are classified as decode requests.
2269	# Use `query_len_support` (above) to set this automatically
2270	# when speculative decoding is enabled.
2271	reorder_batch_threshold: int = 1
2272
2273	def _validate_dspark_dcp_support(self, supports_dcp_with_varlen: bool) -> None:
2274	speculative_config = getattr(self.vllm_config, "speculative_config", None)
2275	parallel_config = self.vllm_config.parallel_config
2276	if (
2277	speculative_config is None
2278	or getattr(speculative_config, "method", None) != "dspark"
2279	or parallel_config.decode_context_parallel_size <= 1
2280	):
2281	return
2282
2283	if self.non_causal_multi_token_decode:
2284	supported = self.supports_non_causal_multi_token_dcp
2285	query_mode = "non-causal draft"
2286	else:
2287	supported = supports_dcp_with_varlen
2288	query_mode = "causal multi-token"
2289
2290	if not supported:
2291	raise ValueError(
2292	f"{type(self).__name__} does not support {query_mode} MLA "
2293	"attention for DSpark with decode context parallelism. Select "
2294	"a backend with explicit DSpark DCP support or set "
2295	"decode_context_parallel_size=1."
2296	)
2297
2298	@staticmethod
2299	def determine_chunked_prefill_workspace_size(vllm_config: VllmConfig) -> int:
2300	scheduler_config = vllm_config.scheduler_config
2301	cache_config = vllm_config.cache_config
2302	model_config = vllm_config.model_config
2303
2304	chunked_prefill_workspace_size = min(
2305	# Try for 8 full length request or at least 4 pages per-request
2306	max(
2307	8 * model_config.max_model_len,
2308	4 * scheduler_config.max_num_seqs * cache_config.block_size,
2309	),
2310	# For long-context models try not to over-allocate limiting
2311	# kv-cache space, limiting it to 64k tokens,
2312	# which would result in the workspace being:
2313	# 2(576)(64*1024) = 144mb
2314	# (assuming 576 MLA head dim, and fp16)
2315	# which would result in up-projected context being
2316	# 2(192128)(641024) = 3gb
2317	# (assuming 192 QK head dim, 128 heads, and fp16)
2318	64 * 1024,
2319	)
2320
2321	return align_mla_chunked_context_workspace_size(
2322	vllm_config,
2323	chunked_prefill_workspace_size,
2324	)
2325
2326	@staticmethod
2327	def determine_prefill_query_data_type(
2328	vllm_config: VllmConfig,
2329	model_dtype: torch.dtype,
2330	) -> torch.dtype:
2331	"""Determine the query data type for prefill queries.
2332	Return FP8 dtype if cache is FP8 and prefill query quantization
2333	is enabled, else model dtype.
2334	"""
2335	use_fp8 = (
2336	is_quantized_kv_cache(vllm_config.cache_config.cache_dtype)
2337	and vllm_config.attention_config.use_prefill_query_quantization
2338	and backend_supports_prefill_query_quantization()
2339	)
2340
2341	if use_fp8:
2342	fp8_dtype = current_platform.fp8_dtype()
2343	logger.info_once("FP8 prefill attention enabled: query data type is FP8")
2344	return fp8_dtype
2345	elif vllm_config.attention_config.use_prefill_query_quantization:
2346	logger.info_once(
2347	"Unable to perform FP8 prefill attention when"
2348	" use_prefill_query_quantization is enabled. Please"
2349	" ensure that --kv-cache-dtype is set to fp8 and your prefill"
2350	" backend is compatible with FP8 attention.",
2351	)
2352	return model_dtype
2353	elif (
2354	is_quantized_kv_cache(vllm_config.cache_config.cache_dtype)
2355	and backend_supports_prefill_query_quantization()
2356	):
2357	logger.warning_once(
2358	"FP8 KV cache is enabled but prefill queries are not "
2359	"quantized to FP8. For long-context workloads (ISL >= 4K), "
2360	"enabling FP8 prefill attention can significantly optimize "
2361	"prefill latency. To enable, add: "
2362	'--attention-config \'{"use_prefill_query_quantization"'
2363	": true}'",
2364	)
2365
2366	return model_dtype
2367
2368	def __init__(
2369	self,
2370	kv_cache_spec: AttentionSpec,
2371	layer_names: list[str],
2372	vllm_config: VllmConfig,
2373	device: torch.device,
2374	metadata_cls: type[M] \| None = None,
2375	supports_dcp_with_varlen: bool = False,
2376	):
2377	self.metadata_cls = (
2378	metadata_cls if metadata_cls is not None else MLACommonMetadata
2379	)
2380	self.kv_cache_spec = kv_cache_spec
2381	self.model_config = vllm_config.model_config
2382	parallel_config = vllm_config.parallel_config
2383	self.compilation_config = vllm_config.compilation_config
2384	self.vllm_config = vllm_config
2385	self.device = device
2386	self.use_pcp = parallel_config.prefill_context_parallel_size > 1
2387	self.non_causal_multi_token_decode = getattr(
2388	kv_cache_spec, "non_causal_multi_token_decode", False
2389	)
2390	self._validate_dspark_dcp_support(supports_dcp_with_varlen)
2391
2392	# A draft cache group can have a different head count from the target.
2393	self.num_heads = get_num_attention_heads_from_layers(
2394	vllm_config, layer_names
2395	) or self.model_config.get_num_attention_heads(parallel_config)
2396	# Hybrid MLA models may use different latent dimensions per KV group.
2397	layer = self.compilation_config.static_forward_context[layer_names[0]]
2398	self.mla_dims = MLADims(
2399	q_lora_rank=layer.q_lora_rank,
2400	kv_lora_rank=layer.kv_lora_rank,
2401	qk_nope_head_dim=layer.qk_nope_head_dim,
2402	qk_rope_head_dim=layer.qk_rope_head_dim,
2403	v_head_dim=layer.v_head_dim,
2404	)
2405	self.aot_schedule = current_platform.is_cuda()
2406
2407	self.kv_cache_spec = kv_cache_spec
2408	self.q_data_type = self.determine_prefill_query_data_type(
2409	vllm_config, self.model_config.dtype
2410	)
2411	attention_layer = self.compilation_config.static_forward_context[layer_names[0]]
2412
2413	try:
2414	self.dcp_world_size = get_dcp_group().world_size
2415	except AssertionError:
2416	# DCP might not be initialized in testing
2417	self.dcp_world_size = 1
2418	self.dcp_local_block_size = parallel_config.cp_kv_cache_interleave_size
2419	self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size
2420	self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
2421
2422	self.page_size = self.kv_cache_spec.block_size
2423
2424	self.chunked_prefill_workspace_size = (
2425	self.determine_chunked_prefill_workspace_size(vllm_config)
2426	)
2427
2428	use_packed_fp8_cache = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"
2429	self.dcp_manager: MLADCPManager \| None = None
2430	if self.dcp_world_size > 1:
2431	# Note(hc): The local kvcache is incomplete when DCP is triggered,
2432	# an additional kvcache allgather across the DCP group is therefore
2433	# required, so the workspace has to be enlarged by 1/DCP relative
2434	# to the original TP allocation.
2435	assert self.chunked_prefill_workspace_size % self.dcp_world_size == 0
2436	self.chunked_prefill_workspace = torch.empty(
2437	(
2438	self.chunked_prefill_workspace_size
2439	+ self.chunked_prefill_workspace_size // self.dcp_world_size,
2440	self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim,
2441	),
2442	dtype=torch.bfloat16
2443	if use_packed_fp8_cache
2444	else self.model_config.dtype,
2445	device=device,
2446	)
2447	self.dcp_manager = getattr(attention_layer, "dcp_manager", None)
2448	assert isinstance(self.dcp_manager, MLADCPManager)
2449	self.dcp_manager.init_kv_gather(
2450	self.chunked_prefill_workspace,
2451	self.chunked_prefill_workspace_size,
2452	)
2453	else:
2454	self.chunked_prefill_workspace = torch.empty(
2455	(
2456	self.chunked_prefill_workspace_size,
2457	self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim,
2458	),
2459	dtype=torch.bfloat16 if use_packed_fp8_cache else self.q_data_type,
2460	device=device,
2461	)
2462
2463	# Metadata builders are created per ubatch when DBO is enabled. MLA
2464	# prefill backends keep the prepared metadata on the backend object, so
2465	# each builder needs its own backend instance to avoid cross-ubatch races.
2466	self._prefill_backend = attention_layer.prefill_backend.clone()
2467
2468	supports_spec_decode = self.query_len_support != QueryLenSupport.SINGLE_ONLY
2469	self._init_reorder_batch_threshold(
2470	self.reorder_batch_threshold, supports_spec_decode, supports_dcp_with_varlen
2471	)
2472
2473	if self.query_len_support == QueryLenSupport.SINGLE_ONLY:
2474	assert self.reorder_batch_threshold == 1, (
2475	f"reorder_batch_threshold must be 1 when query_len_support is "
2476	f"SINGLE_ONLY, got {self.reorder_batch_threshold}"
2477	)
2478
2479	def _build_decode(
2480	self,
2481	block_table_tensor: torch.Tensor,
2482	seq_lens_device: torch.Tensor,
2483	max_seq_len: int,
2484	query_start_loc_cpu: torch.Tensor,
2485	query_start_loc_device: torch.Tensor,
2486	num_decode_tokens: int,
2487	max_query_len: int,
2488	dcp_tot_seq_lens_device: torch.Tensor \| None,
2489	) -> MLACommonDecodeMetadata:
2490	return MLACommonDecodeMetadata(
2491	block_table=block_table_tensor,
2492	seq_lens=seq_lens_device,
2493	dcp_tot_seq_lens=dcp_tot_seq_lens_device,
2494	)
2495
2496	def build_for_cudagraph_capture(
2497	self, common_attn_metadata: CommonAttentionMetadata
2498	) -> M:
2499	"""This method builds the metadata for full cudagraph capture.
2500	Currently, only decode is supported for full cudagraphs with MLA.
2501	"""
2502	m = common_attn_metadata
2503	assert m.num_reqs <= (m.num_actual_tokens * self.reorder_batch_threshold), (
2504	"MLA only supports decode-only full CUDAGraph capture. "
2505	"Make sure all cudagraph capture sizes <= max_num_seq."
2506	)
2507
2508	assert m.max_query_len <= self.reorder_batch_threshold # decode only
2509
2510	return self.build(0, m)
2511
2512	def build(
2513	self,
2514	common_prefix_len: int,
2515	common_attn_metadata: CommonAttentionMetadata,
2516	fast_build: bool = False,
2517	) -> M:
2518	num_reqs = common_attn_metadata.num_reqs
2519	num_tokens = common_attn_metadata.num_actual_tokens
2520	max_query_len = common_attn_metadata.max_query_len
2521	max_seq_len = common_attn_metadata.max_seq_len
2522
2523	# Note(simon): be careful about the CPU <> GPU memory movement in this
2524	# function. We should avoid GPU -> CPU sync as much as possible because
2525	# it blocks on all previous kernels.
2526	device = self.device
2527	block_table_tensor = common_attn_metadata.block_table_tensor
2528	slot_mapping = common_attn_metadata.slot_mapping
2529
2530	query_start_loc = common_attn_metadata.query_start_loc
2531	query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
2532	seq_lens = common_attn_metadata.seq_lens
2533	dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens
2534
2535	non_causal_decode = common_attn_metadata.causal is False
2536	if non_causal_decode:
2537	if not (
2538	self.supports_non_causal_multi_token_decode
2539	and self.non_causal_multi_token_decode
2540	):
2541	raise ValueError(
2542	"Non-causal multi-token MLA requires an explicitly supported "
2543	"attention group."
2544	)
2545	query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
2546	num_active_reqs = int(torch.count_nonzero(query_lens > 0))
2547	uniform_active_queries = num_active_reqs > 0 and bool(
2548	torch.all(query_lens[:num_active_reqs] == query_lens[0])
2549	)
2550	trailing_graph_padding = bool(torch.all(query_lens[num_active_reqs:] == 0))
2551	if not (uniform_active_queries and trailing_graph_padding):
2552	raise ValueError(
2553	"Non-causal MLA requires a uniform query block; got query "
2554	f"lengths {query_lens.tolist()}."
2555	)
2556	# Use exact GPU sequence lengths instead of the prefill path's CPU
2557	# context-length upper bounds.
2558	num_decodes = num_reqs
2559	num_prefills = 0
2560	num_decode_tokens = num_tokens
2561	num_prefill_tokens = 0
2562	else:
2563	num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
2564	split_decodes_and_prefills(
2565	common_attn_metadata,
2566	decode_threshold=self.reorder_batch_threshold,
2567	require_uniform=(self.query_len_support != QueryLenSupport.VARLEN),
2568	treat_short_extends_as_decodes=not self.use_pcp,
2569	)
2570	)
2571
2572	assert num_decodes + num_prefills == num_reqs
2573	assert num_decode_tokens + num_prefill_tokens == num_tokens
2574
2575	prefill_metadata = None
2576	if num_prefills > 0:
2577	reqs_start = num_decodes # prefill_start
2578
2579	# Upper bound is exact for prefill rows (no D2H sync).
2580	seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
2581	assert seq_lens_cpu is not None
2582	prefill_query_lens_cpu = (
2583	query_start_loc_cpu[reqs_start + 1 : num_reqs + 1]
2584	- query_start_loc_cpu[reqs_start:num_reqs]
2585	)
2586	context_lens_cpu = torch.empty(
2587	num_prefills, dtype=seq_lens_cpu.dtype, pin_memory=PIN_MEMORY
2588	)
2589	torch.subtract(
2590	seq_lens_cpu[reqs_start:num_reqs],
2591	prefill_query_lens_cpu,
2592	out=context_lens_cpu,
2593	)
2594	prefill_query_start_loc = (
2595	query_start_loc[reqs_start:] - query_start_loc[reqs_start]
2596	)
2597	prefill_query_start_loc_cpu = (
2598	query_start_loc_cpu[reqs_start:] - query_start_loc_cpu[reqs_start]
2599	)
2600
2601	chunked_context_metadata = build_mla_chunked_context_metadata(
2602	context_lens_cpu=context_lens_cpu,
2603	prefill_query_start_loc_cpu=prefill_query_start_loc_cpu,
2604	chunked_prefill_workspace=self.chunked_prefill_workspace,
2605	chunked_prefill_workspace_size=self.chunked_prefill_workspace_size,
2606	block_size=self.page_size,
2607	align_chunk_to_block=True,
2608	device=device,
2609	dcp_world_size=self.dcp_world_size,
2610	dcp_local_block_size=self.dcp_local_block_size,
2611	dcp_virtual_block_size=self.dcp_virtual_block_size,
2612	dcp_manager=self.dcp_manager,
2613	)
2614
2615	prefill_metadata = MLACommonPrefillMetadata(
2616	block_table=block_table_tensor[reqs_start:, ...],
2617	query_start_loc=prefill_query_start_loc,
2618	max_query_len=max_query_len,
2619	chunked_context=chunked_context_metadata,
2620	output_dtype=self.model_config.dtype,
2621	q_data_type=self.q_data_type,
2622	prefill_backend=self._prefill_backend,
2623	query_lens_cpu=prefill_query_lens_cpu,
2624	)
2625
2626	self._prefill_backend.prepare_metadata(prefill_metadata)
2627
2628	decode_metadata = None
2629	if num_decodes > 0:
2630	dcp_tot_seq_lens_device = None
2631	if self.dcp_world_size > 1:
2632	dcp_tot_seq_lens_device = seq_lens[:num_decodes]
2633	seq_lens = dcp_local_seq_lens
2634
2635	# After DCP distribution, the maximum number of tokens for any rank is
2636	# ceil(L / (N * I)) * I, where L is max_seq_len, N is dcp_world_size,
2637	# and I is cp_kv_cache_interleave_size.
2638	# This eliminates GPU->CPU sync while minimizing workspace
2639	# over-allocation.
2640	num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size
2641	max_seq_len = (
2642	(max_seq_len + num_partitions - 1) // num_partitions
2643	) * self.cp_kv_cache_interleave_size
2644
2645	decode_metadata = self._build_decode(
2646	block_table_tensor=block_table_tensor[:num_decodes, ...],
2647	seq_lens_device=seq_lens[:num_decodes],
2648	max_seq_len=max_seq_len,
2649	query_start_loc_cpu=query_start_loc_cpu[: num_decodes + 1],
2650	query_start_loc_device=query_start_loc[: num_decodes + 1],
2651	num_decode_tokens=num_decode_tokens,
2652	max_query_len=min(
2653	common_attn_metadata.max_query_len, self.reorder_batch_threshold
2654	),
2655	dcp_tot_seq_lens_device=dcp_tot_seq_lens_device,
2656	)
2657
2658	attn_metadata = self.metadata_cls(
2659	num_reqs=common_attn_metadata.num_reqs,
2660	max_query_len=common_attn_metadata.max_query_len,
2661	max_seq_len=max_seq_len,
2662	num_actual_tokens=num_tokens,
2663	query_start_loc=query_start_loc,
2664	slot_mapping=slot_mapping,
2665	head_dim=self.model_config.get_head_size(),
2666	# MLACommonMetadata Chunk prefill specific
2667	num_decodes=num_decodes,
2668	num_decode_tokens=num_decode_tokens,
2669	num_prefills=num_prefills,
2670	causal=not non_causal_decode,
2671	prefill=prefill_metadata,
2672	decode=decode_metadata,
2673	)
2674
2675	return attn_metadata # type: ignore[return-value]
2676
2677
2678	def reorg_kvcache(
2679	allgatered_kv_c_normed: torch.Tensor,
2680	allgatered_k_pe: torch.Tensor,
2681	padded_local_chunk_seq_lens_lst: list[int],
2682	local_context_lens_allranks: list[list[int]],
2683	local_starts: list[int],
2684	sum_seq_len: int,
2685	max_seq_len: int,
2686	toks: int,
2687	) -> tuple[torch.Tensor, torch.Tensor]:
2688	"""Reorg and unpad kvcache after cp local gather to tp layout for attn kernel.
2689	e.g.
2690	allgatered_kv_c_normed = [T0_0, T0_1, T0_2, T0_3, T1_0, T1_1, ...,
2691	T0_4, T0_5, pad, pad, T1_2, pad, ...]
2692	-> reorganized_kv_c_normed = [T0_0, T0_1, T0_2, T0_3, T0_4, T0_5,
2693	T1_0, T1_1, T1_2, ...]
2694
2695	Args:
2696	allgatered_kv_c_normed: all-gathered, padded latent KV cache.
2697	allgatered_k_pe: all-gathered, padded RoPE key cache.
2698	padded_local_chunk_seq_lens_lst: local chunk context lengths
2699	under current CP rank.
2700	local_context_lens_allranks: local context lengths on each CP rank.
2701	local_starts: per-request local offset into the context this chunk
2702	starts at.
2703	sum_seq_len: the sum of cp_chunk_seq_lens_lst.
2704	max_seq_len: the max value of cp_chunk_seq_lens_lst.
2705	toks: the number of tokens for local gather cache.
2706
2707	"""
2708	kv_c_segments = []
2709	k_pe_segments = []
2710	src_token_idx = 0
2711	max_seq_len_check = 0
2712	for padded_local_chunk_seq_len, local_context_lens, local_start in zip(
2713	padded_local_chunk_seq_lens_lst, local_context_lens_allranks, local_starts
2714	):
2715	cur_seq_len = 0
2716	for rank, local_context_len in enumerate(local_context_lens):
2717	# Note(qcs): We split the context into multiple chunks,
2718	# depending on the size of the workspace.
2719	# local_context in dcp0: \|-----------------\|
2720	# local_context in dcp1: \|--------------\|
2721	# n*padded_local_chunk: \|-----\|-----\|-----\|
2722	# local_chunk_len in dcp1: \|-----\|-----\|--\|
2723	# so we need update the last chunk length in dcp1.
2724	local_chunk_len = min(
2725	max(0, local_context_len - local_start),
2726	padded_local_chunk_seq_len,
2727	)
2728	if local_chunk_len != 0:
2729	kv_c_segment = allgatered_kv_c_normed[
2730	rank * toks + src_token_idx : rank * toks
2731	+ src_token_idx
2732	+ local_chunk_len
2733	]
2734	k_pe_segment = allgatered_k_pe[
2735	rank * toks + src_token_idx : rank * toks
2736	+ src_token_idx
2737	+ local_chunk_len
2738	]
2739	kv_c_segments.append(kv_c_segment)
2740	k_pe_segments.append(k_pe_segment)
2741	cur_seq_len += local_chunk_len
2742	max_seq_len_check = max(max_seq_len_check, cur_seq_len)
2743	src_token_idx += padded_local_chunk_seq_len
2744	reorganized_kv_c_normed = torch.cat(kv_c_segments, dim=0)
2745	reorganized_k_pe = torch.cat(k_pe_segments, dim=0)
2746	assert reorganized_kv_c_normed.shape[0] == sum_seq_len
2747	assert reorganized_k_pe.shape[0] == sum_seq_len
2748	assert max_seq_len_check == max_seq_len
2749	return reorganized_kv_c_normed, reorganized_k_pe
2750
2751
2752	def neutralize_empty_context_partials(
2753	chunked_context: "MLACommonPrefillMetadata.ChunkedContextMetadata",
2754	output: torch.Tensor,
2755	output_lse: torch.Tensor,
2756	) -> None:
2757	"""Neutralize the partial of every prefill that no chunk covers.
2758
2759	A prefill without context is never gathered, so nothing would write its rows;
2760	a zero output with an ``-inf`` lse carries no weight into the final merge
2761	against the suffix partial.
2762	"""
2763	for token_slice in chunked_context.empty_token_slices:
2764	output[token_slice].zero_()
2765	output_lse[:, token_slice].fill_(float("-inf"))
2766
2767
2768	def init_mla_context_partial(
2769	chunked_context: "MLACommonPrefillMetadata.ChunkedContextMetadata",
2770	attn_output: torch.Tensor,
2771	attn_softmax_lse: torch.Tensor,
2772	num_tokens: int,
2773	) -> tuple[torch.Tensor, torch.Tensor]:
2774	"""Allocate the running context partial over all prefill tokens.
2775
2776	Laid out like the chunk partials so the final whole-batch merge against the
2777	suffix partial sees matching head strides. Callers whose backend honors an
2778	``out`` tensor already know that layout and can allocate directly, pairing it
2779	with ``neutralize_empty_context_partials``.
2780	"""
2781	output = torch.empty(
2782	(num_tokens, *attn_output.shape[1:]),
2783	dtype=attn_output.dtype,
2784	device=attn_output.device,
2785	)
2786	output_lse = torch.empty(
2787	(attn_softmax_lse.shape[0], num_tokens),
2788	dtype=attn_softmax_lse.dtype,
2789	device=attn_softmax_lse.device,
2790	)
2791	neutralize_empty_context_partials(chunked_context, output, output_lse)
2792	return output, output_lse
2793
2794
2795	def accumulate_mla_context_chunk(
2796	chunk: "MLACommonPrefillMetadata.ContextChunk",
2797	attn_output: torch.Tensor,
2798	attn_softmax_lse: torch.Tensor,
2799	output: torch.Tensor,
2800	output_lse: torch.Tensor,
2801	output_written: bool = False,
2802	) -> None:
2803	"""Fold one chunk's partial into the running context partial.
2804
2805	Only the first request may be a continuation; its tokens are merged and the
2806	remaining token range is initialized.
2807
2808	Args:
2809	chunk: The context chunk being folded in.
2810	attn_output: The chunk's attention output.
2811	attn_softmax_lse: The chunk's log-sum-exp values.
2812	output: Running context partial, updated in place.
2813	output_lse: Running log-sum-exp, updated in place.
2814	output_written: The chunk's attention output already landed in
2815	``output[chunk.token_slice]`` because the backend was handed it as
2816	``out``, leaving only the lse to fold. Invalid for a continuation
2817	chunk, whose leading tokens must be merged rather than overwritten.
2818
2819	"""
2820	token_start = chunk.token_slice.start
2821	token_end = chunk.token_slice.stop
2822	init_start = token_start
2823	if chunk.is_continuation:
2824	assert not output_written, (
2825	"a continuation chunk must not write over the partial it merges with"
2826	)
2827	init_start = chunk.continuation_token_end
2828	num_merged = init_start - token_start
2829	merge_attn_states(
2830	output=output[token_start:init_start],
2831	output_lse=output_lse[:, token_start:init_start],
2832	prefix_output=output[token_start:init_start],
2833	prefix_lse=output_lse[:, token_start:init_start],
2834	suffix_output=attn_output[:num_merged],
2835	suffix_lse=attn_softmax_lse[:, :num_merged],
2836	)
2837	if init_start < token_end:
2838	written = init_start - token_start
2839	if not output_written:
2840	output[init_start:token_end].copy_(attn_output[written:])
2841	output_lse[:, init_start:token_end].copy_(attn_softmax_lse[:, written:])
2842
2843
2844	class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2845	"""Shared MLA base providing dense-MHA prefill (via the selected
2846	MLAPrefillBackend) for both dense and sparse impls; subclasses add decode
2847	(``forward_mqa``).
2848	"""
2849
2850	_use_flashinfer_concat_mla_k: bool
2851
2852	def __init__(
2853	self,
2854	num_heads: int,
2855	head_size: int,
2856	scale: float,
2857	num_kv_heads: int,
2858	kv_cache_dtype: str,
2859	kv_lora_rank: int,
2860	qk_nope_head_dim: int,
2861	qk_rope_head_dim: int,
2862	qk_head_dim: int,
2863	v_head_dim: int,
2864	kv_b_proj: ColumnParallelLinear,
2865	) -> None:
2866	self.num_heads = num_heads
2867	self.head_size = head_size
2868	self.scale = float(scale)
2869	self.num_kv_heads = num_kv_heads
2870	self.kv_cache_dtype = kv_cache_dtype
2871	self.kv_lora_rank = kv_lora_rank
2872	self.qk_nope_head_dim = qk_nope_head_dim
2873	self.qk_rope_head_dim = qk_rope_head_dim
2874	self.qk_head_dim = qk_head_dim
2875	self.v_head_dim = v_head_dim
2876	self.kv_b_proj = kv_b_proj
2877
2878	def _concat_k_nope_k_pe(
2879	self, k_nope: torch.Tensor, k_pe: torch.Tensor
2880	) -> torch.Tensor:
2881	"""Efficiently concatenate k_nope and k_pe tensors along the last dimension.
2882
2883	This function avoids the performance penalty of torch.cat with expanded
2884	non-contiguous tensors by pre-allocating the output and using direct copies.
2885
2886	Args:
2887	k_nope: Tensor of shape [..., nope_dim]
2888	k_pe: Tensor to broadcast and concatenate, typically shape [..., 1, pe_dim]
2889	or [..., pe_dim]
2890
2891	Returns:
2892	Tensor of shape [..., nope_dim + pe_dim]
2893
2894	"""
2895	if k_pe.shape[-1] == 0:
2896	# NoPE MLA: nothing to append, so no copy either.
2897	return k_nope
2898
2899	k = torch.empty(
2900	(*k_nope.shape[:-1], k_nope.shape[-1] + k_pe.shape[-1]),
2901	dtype=k_nope.dtype,
2902	device=k_nope.device,
2903	)
2904
2905	if self._use_flashinfer_concat_mla_k:
2906	torch.ops.vllm.flashinfer_concat_mla_k(k, k_nope, k_pe)
2907	else:
2908	# Fallback: Direct copies with efficient broadcasting
2909	k[..., : k_nope.shape[-1]] = k_nope
2910	k[..., k_nope.shape[-1] :] = k_pe
2911	return k
2912
2913	def _compute_prefill_context(
2914	self,
2915	q: torch.Tensor,
2916	kv_c_and_k_pe_cache: torch.Tensor,
2917	attn_metadata: MLACommonMetadata,
2918	k_scale: torch.Tensor,
2919	):
2920	assert attn_metadata.prefill is not None
2921	prefill_metadata = attn_metadata.prefill
2922	assert prefill_metadata.prefill_backend is not None
2923	chunked_context = prefill_metadata.chunked_context
2924	assert chunked_context is not None
2925
2926	use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()
2927	kv_b_proj_input_dtype = _get_kv_b_proj_input_dtype(
2928	self.kv_b_proj, use_fp8_prefill
2929	)
2930	workspace = chunked_context.workspace
2931
2932	if use_fp8_prefill:
2933	q = q.to(prefill_metadata.q_data_type)
2934
2935	output = None
2936	output_lse = None
2937	for chunk in chunked_context.chunks:
2938	toks = chunk.num_context_tokens
2939	chunk_block_table = prefill_metadata.block_table[chunk.request_slice]
2940	if self.kv_cache_dtype == "fp8_ds_mla":
2941	ops.cp_gather_and_upconvert_fp8_kv_cache(
2942	src_cache=kv_c_and_k_pe_cache,
2943	dst=workspace[:toks],
2944	block_table=chunk_block_table,
2945	workspace_starts=chunk.cu_seq_lens,
2946	batch_size=chunk.num_requests,
2947	seq_starts=chunk.starts,
2948	)
2949	elif current_platform.is_cpu():
2950	assert not is_quantized_kv_cache(self.kv_cache_dtype), (
2951	"CPU MLA context gather fallback only supports "
2952	f"non-quantized KV cache, got {self.kv_cache_dtype}"
2953	)
2954	ops.gather_mla_context_cache_cpu(
2955	src_cache=kv_c_and_k_pe_cache,
2956	dst=workspace[:toks],
2957	block_table=chunk_block_table,
2958	starts=chunk.starts,
2959	cu_seq_lens=chunk.cu_seq_lens,
2960	)
2961	elif not use_fp8_prefill:
2962	ops.gather_and_maybe_dequant_cache(
2963	src_cache=kv_c_and_k_pe_cache,
2964	dst=workspace,
2965	block_table=chunk_block_table,
2966	cu_seq_lens=chunk.cu_seq_lens,
2967	token_to_seq=chunk.token_to_seq,
2968	num_tokens=toks,
2969	kv_cache_dtype=self.kv_cache_dtype,
2970	scale=k_scale,
2971	seq_starts=chunk.starts,
2972	)
2973	else:
2974	# FP8 path: gather cache without dequantization
2975	ops.cp_gather_cache(
2976	src_cache=kv_c_and_k_pe_cache,
2977	dst=workspace[:toks],
2978	block_table=chunk_block_table,
2979	cu_seq_lens=chunk.cu_seq_lens,
2980	batch_size=chunk.num_requests,
2981	seq_starts=chunk.starts,
2982	)
2983
2984	# Extract kv_c_normed from workspace
2985	kv_c_normed = workspace[:toks][..., : self.kv_lora_rank]
2986	if kv_b_proj_input_dtype is not None:
2987	kv_c_normed = kv_c_normed.to(kv_b_proj_input_dtype)
2988
2989	k_pe = workspace[:toks][..., self.kv_lora_rank :].unsqueeze(1)
2990	kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
2991	-1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
2992	)
2993
2994	# To Do: Use epilogue of kv_b_proj to generate fp8 kv_nope.
2995	if use_fp8_prefill:
2996	kv_nope = kv_nope.to(prefill_metadata.q_data_type)
2997	k_pe = k_pe.to(prefill_metadata.q_data_type)
2998	k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
2999
3000	k = self._concat_k_nope_k_pe(k_nope, k_pe)
3001
3002	attn_output, attn_softmax_lse = (
3003	prefill_metadata.prefill_backend.run_prefill_context_chunk(
3004	chunk=chunk,
3005	q=q[chunk.token_slice],
3006	k=k,
3007	v=v,
3008	)
3009	)
3010
3011	if output is None:
3012	if (
3013	len(chunked_context.chunks) == 1
3014	and not chunked_context.empty_token_slices
3015	):
3016	return attn_output, attn_softmax_lse
3017	output, output_lse = init_mla_context_partial(
3018	chunked_context,
3019	attn_output,
3020	attn_softmax_lse,
3021	num_tokens=q.shape[0],
3022	)
3023	accumulate_mla_context_chunk(
3024	chunk, attn_output, attn_softmax_lse, output, output_lse
3025	)
3026
3027	return output, output_lse
3028
3029	def _context_parallel_compute_prefill_context(
3030	self,
3031	q: torch.Tensor,
3032	kv_c_and_k_pe_cache: torch.Tensor,
3033	attn_metadata: MLACommonMetadata,
3034	k_scale: torch.Tensor,
3035	dcp_world_size: int,
3036	):
3037	assert attn_metadata.prefill is not None
3038	prefill_metadata = attn_metadata.prefill
3039	assert prefill_metadata.prefill_backend is not None
3040	chunked_context = prefill_metadata.chunked_context
3041	assert chunked_context is not None
3042
3043	use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()
3044	kv_b_proj_input_dtype = _get_kv_b_proj_input_dtype(
3045	self.kv_b_proj, use_fp8_prefill
3046	)
3047	output = None
3048	output_lse = None
3049	workspace = chunked_context.workspace
3050
3051	for chunk in chunked_context.chunks:
3052	assert chunk.padded_local_seq_lens is not None
3053	assert chunk.local_context_lens_allranks is not None
3054	assert chunk.padded_local_cu_seq_lens is not None
3055	assert chunk.padded_local_token_to_seq is not None
3056	assert chunk.local_starts is not None
3057
3058	toks = chunk.num_local_context_tokens
3059	padded_local_cu_seq_lens = chunk.padded_local_cu_seq_lens
3060	block_table = prefill_metadata.block_table[chunk.request_slice]
3061	if self.kv_cache_dtype == "fp8_ds_mla":
3062	ops.cp_gather_and_upconvert_fp8_kv_cache(
3063	src_cache=kv_c_and_k_pe_cache,
3064	dst=workspace[:toks],
3065	block_table=block_table,
3066	workspace_starts=padded_local_cu_seq_lens,
3067	batch_size=chunk.num_requests,
3068	seq_starts=chunk.starts,
3069	)
3070	elif is_quantized_kv_cache(self.kv_cache_dtype):
3071	assert k_scale is not None
3072	ops.gather_and_maybe_dequant_cache(
3073	src_cache=kv_c_and_k_pe_cache,
3074	dst=workspace,
3075	block_table=block_table,
3076	cu_seq_lens=padded_local_cu_seq_lens,
3077	token_to_seq=chunk.padded_local_token_to_seq,
3078	num_tokens=toks,
3079	kv_cache_dtype=self.kv_cache_dtype,
3080	scale=k_scale,
3081	seq_starts=chunk.starts,
3082	)
3083	else:
3084	ops.cp_gather_cache(
3085	src_cache=kv_c_and_k_pe_cache,
3086	dst=workspace[:toks],
3087	block_table=block_table,
3088	cu_seq_lens=padded_local_cu_seq_lens,
3089	batch_size=chunk.num_requests,
3090	seq_starts=chunk.starts,
3091	)
3092	# workspace
3093	# \|------- N tokens --------\|--------- N*dcp_size tokens ----------\|
3094	# \|<- use for local_gather ->\|<--------- use for allgather -------->\|
3095	allgather_offset = workspace.shape[0] // (dcp_world_size + 1)
3096	assert allgather_offset * (dcp_world_size + 1) == workspace.shape[0]
3097	assert toks <= allgather_offset
3098	local_gathered_kvcache = workspace[:toks]
3099	cur_allgather_workspace = workspace[
3100	allgather_offset : allgather_offset * (1 + dcp_world_size)
3101	]
3102	assert toks * dcp_world_size <= cur_allgather_workspace.shape[0]
3103	cur_allgather_kvcache = cur_allgather_workspace[: toks * dcp_world_size]
3104	dcp_manager = cast(MLADCPManager, chunked_context.dcp_manager)
3105	dcp_manager.kv_gather(cur_allgather_kvcache, local_gathered_kvcache)
3106	assert (
3107	cur_allgather_kvcache.shape[-1]
3108	== self.kv_lora_rank + self.qk_rope_head_dim
3109	)
3110	allgatered_kv_c_normed, allgatered_k_pe = cur_allgather_kvcache.unsqueeze(
3111	1
3112	).split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
3113
3114	kv_c_normed, k_pe = reorg_kvcache(
3115	allgatered_kv_c_normed,
3116	allgatered_k_pe,
3117	padded_local_chunk_seq_lens_lst=chunk.padded_local_seq_lens,
3118	local_context_lens_allranks=chunk.local_context_lens_allranks,
3119	local_starts=chunk.local_starts,
3120	sum_seq_len=chunk.num_context_tokens,
3121	max_seq_len=chunk.max_seq_len,
3122	toks=toks,
3123	)
3124	if kv_b_proj_input_dtype is not None:
3125	kv_c_normed = kv_c_normed.to(kv_b_proj_input_dtype)
3126
3127	kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
3128	-1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
3129	)
3130	if use_fp8_prefill:
3131	kv_nope = kv_nope.to(prefill_metadata.q_data_type)
3132	k_pe = k_pe.to(prefill_metadata.q_data_type)
3133	k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
3134	k = self._concat_k_nope_k_pe(k_nope, k_pe)
3135
3136	attn_output, attn_softmax_lse = (
3137	prefill_metadata.prefill_backend.run_prefill_context_chunk(
3138	chunk=chunk,
3139	q=q[chunk.token_slice],
3140	k=k,
3141	v=v,
3142	)
3143	)
3144
3145	if output is None:
3146	if (
3147	len(chunked_context.chunks) == 1
3148	and not chunked_context.empty_token_slices
3149	):
3150	return attn_output, attn_softmax_lse
3151	output, output_lse = init_mla_context_partial(
3152	chunked_context,
3153	attn_output,
3154	attn_softmax_lse,
3155	num_tokens=q.shape[0],
3156	)
3157	accumulate_mla_context_chunk(
3158	chunk, attn_output, attn_softmax_lse, output, output_lse
3159	)
3160
3161	return output, output_lse
3162
3163	def forward_mha( # type: ignore[override]
3164	self,
3165	q: torch.Tensor,
3166	kv_c_normed: torch.Tensor,
3167	k_pe: torch.Tensor,
3168	kv_c_and_k_pe_cache: torch.Tensor,
3169	attn_metadata: MLACommonMetadata,
3170	k_scale: torch.Tensor,
3171	output: torch.Tensor,
3172	output_scale: torch.Tensor \| None = None,
3173	) -> None:
3174	assert attn_metadata.prefill is not None
3175	prefill_metadata = attn_metadata.prefill
3176	assert prefill_metadata.prefill_backend is not None
3177	use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()
3178
3179	# Convert q to FP8 if FP8 prefill attention is enabled
3180	if use_fp8_prefill:
3181	q = q.to(prefill_metadata.q_data_type)
3182
3183	has_context = prefill_metadata.chunked_context is not None
3184	assert output_scale is None or not has_context, (
3185	"Fused FP8 output is only wired for the non-chunked-context path"
3186	)
3187
3188	kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
3189	-1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
3190	)
3191	k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
3192	k = self._concat_k_nope_k_pe(k_nope, k_pe)
3193
3194	if use_fp8_prefill:
3195	k = k.to(prefill_metadata.q_data_type)
3196	v = v.to(prefill_metadata.q_data_type)
3197
3198	output_prefill = prefill_metadata.prefill_backend.run_prefill_new_tokens(
3199	q=q,
3200	k=k,
3201	v=v,
3202	return_softmax_lse=has_context,
3203	out=(
3204	output.view(-1, self.num_heads, self.v_head_dim)
3205	if output_scale is not None
3206	else None
3207	),
3208	output_scale=output_scale,
3209	)
3210
3211	if has_context:
3212	assert prefill_metadata.chunked_context is not None
3213	suffix_output, suffix_lse = output_prefill
3214	if self.dcp_world_size > 1:
3215	context_output, context_lse = (
3216	self._context_parallel_compute_prefill_context(
3217	q,
3218	kv_c_and_k_pe_cache,
3219	attn_metadata,
3220	k_scale=k_scale,
3221	dcp_world_size=self.dcp_world_size,
3222	)
3223	)
3224	else:
3225	context_output, context_lse = self._compute_prefill_context(
3226	q, kv_c_and_k_pe_cache, attn_metadata, k_scale
3227	)
3228
3229	context_output = context_output[..., : self.v_head_dim]
3230	suffix_output = suffix_output[..., : self.v_head_dim]
3231
3232	output = output.view(-1, self.num_heads, self.v_head_dim)
3233	merge_attn_states(
3234	output=output,
3235	prefix_output=context_output,
3236	prefix_lse=context_lse,
3237	suffix_output=suffix_output,
3238	suffix_lse=suffix_lse,
3239	)
3240	elif output_scale is None:
3241	# With output_scale set, backend already wrote into `output` in place.
3242	assert isinstance(output_prefill, torch.Tensor)
3243	output_prefill = output_prefill[..., : self.v_head_dim]
3244	output_prefill = output_prefill.flatten(start_dim=-2)
3245	output.copy_(output_prefill)
3246
3247
3248	class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
3249	"""NOTE: Please read the comment at the top of the file before trying to
3250	understand this class
3251	"""
3252
3253	def fused_output_quant_supported(self, quant_key):
3254	return quant_key in (
3255	kFp8StaticTensorSym,
3256	kNvfp4Dynamic,
3257	kFp8Dynamic128Sym,
3258	kFp8Dynamic64Sym,
3259	)
3260
3261	def __init__(
3262	self,
3263	num_heads: int,
3264	head_size: int,
3265	scale: float,
3266	num_kv_heads: int,
3267	alibi_slopes: list[float] \| None,
3268	sliding_window: int \| None,
3269	kv_cache_dtype: str,
3270	logits_soft_cap: float \| None,
3271	attn_type: str,
3272	kv_sharing_target_layer_name: str \| None,
3273	# MLA Specific Arguments
3274	q_lora_rank: int \| None,
3275	kv_lora_rank: int,
3276	qk_nope_head_dim: int,
3277	qk_rope_head_dim: int,
3278	qk_head_dim: int,
3279	v_head_dim: int,
3280	kv_b_proj: ColumnParallelLinear,
3281	# DSV3.2 MLA Specific Arguments
3282	indexer: object \| None = None,
3283	q_pad_num_heads: int \| None = None,
3284	) -> None:
3285	if kv_sharing_target_layer_name is not None:
3286	raise NotImplementedError("KV sharing is not supported for MLA")
3287
3288	super().__init__(
3289	num_heads,
3290	head_size,
3291	scale,
3292	num_kv_heads,
3293	kv_cache_dtype,
3294	kv_lora_rank,
3295	qk_nope_head_dim,
3296	qk_rope_head_dim,
3297	qk_head_dim,
3298	v_head_dim,
3299	kv_b_proj,
3300	)
3301	self.q_lora_rank = q_lora_rank
3302	self.indexer = indexer
3303	self.q_pad_num_heads = q_pad_num_heads
3304	self.supports_quant_query_input = True
3305	self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()
3306
3307	# Use flashinfer's optimized concat_mla_k kernel when available.
3308	# The kernel is optimized for DeepSeek V3 dimensions:
3309	# num_heads=128, nope_dim=128, rope_dim=64
3310	self._use_flashinfer_concat_mla_k = (
3311	has_flashinfer()
3312	and (self.num_heads == 128)
3313	and (self.qk_nope_head_dim == 128)
3314	and (self.qk_rope_head_dim == 64)
3315	)
3316
3317	parallel_config = get_current_vllm_config().parallel_config
3318	# Avoid requiring an initialized DCP group in tests.
3319	self.dcp_world_size: int = parallel_config.decode_context_parallel_size
3320
3321	@abstractmethod
3322	def forward_mqa(
3323	self,
3324	q: torch.Tensor \| tuple[torch.Tensor, torch.Tensor],
3325	kv_c_and_k_pe_cache: torch.Tensor,
3326	attn_metadata: M,
3327	layer: AttentionLayer,
3328	) -> tuple[torch.Tensor, torch.Tensor \| None]:
3329	raise NotImplementedError