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

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 2,471 lines 102 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4	import functools
5	from dataclasses import dataclass
6	from enum import Enum
7	from pathlib import Path
8	from typing import TYPE_CHECKING, ClassVar, Final
9
10	import torch
11
12	from vllm import envs
13	from vllm._aiter_ops import rocm_aiter_ops
14	from vllm.config import VllmConfig
15	from vllm.config.cache import CacheDType
16	from vllm.distributed import get_dcp_group
17	from vllm.logger import init_logger
18	from vllm.model_executor.layers.attention.mla_attention import (
19	MLACommonBackend,
20	MLACommonDecodeMetadata,
21	MLACommonImpl,
22	MLACommonMetadata,
23	MLACommonMetadataBuilder,
24	QueryLenSupport,
25	)
26	from vllm.triton_utils import tl, triton
27	from vllm.utils.gpu_sync_debug import gpu_sync_allowed
28	from vllm.utils.math_utils import cdiv, largest_power_of_2_divisor
29	from vllm.v1.attention.backend import (
30	AttentionCGSupport,
31	AttentionLayer,
32	CommonAttentionMetadata,
33	MultipleOf,
34	)
35	from vllm.v1.attention.backends.utils import (
36	get_dcp_local_seq_lens,
37	)
38	from vllm.v1.attention.ops.rocm_aiter_mla_merge import (
39	merge_mla_segments_triton,
40	)
41	from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache
42
43	if TYPE_CHECKING:
44	from vllm.platforms.interface import DeviceCapability
45
46	logger = init_logger(__name__)
47
48
49	def _segmented_mla_page_size(block_size: int) -> int:
50	"""Largest supported power-of-two subpage dividing a physical KV block.
51
52	The subpage size becomes the segmented kernel's ``TILE_SIZE``: the cache is
53	reinterpreted as ``(-1, page_size, 1, head_dim)`` before the call. It must
54	divide the physical block and be a power of two, since the kernel walks a
55	tile with ``tl.arange(0, TILE_SIZE)``. 128 is the largest tile that tiling
56	supports -- a tile holds ``TILE_SIZE x kv_lora_rank`` keys and its scores
57	are ``BLOCK_M x TILE_SIZE`` -- so a larger manager block is split into
58	several subpages instead of widening the tile.
59	"""
60	assert block_size > 0
61	return min(128, largest_power_of_2_divisor(block_size))
62
63
64	@functools.lru_cache(maxsize=1)
65	def _get_mla_gluon():
66	"""Load the small-head Gluon MLA entry point."""
67	unified_module = "aiter.ops.triton.gluon.mla_gluon"
68	try:
69	from aiter.ops.triton.gluon.mla_gluon import mla_gluon
70
71	return mla_gluon
72	except ModuleNotFoundError as unified_import_error:
73	if not unified_module.startswith(unified_import_error.name or ""):
74	raise
75	legacy_module = "aiter.ops.triton.gluon.mla_decode_gluon"
76	try:
77	from aiter.ops.triton.gluon.mla_decode_gluon import mla_decode_gluon
78
79	return mla_decode_gluon
80	except ModuleNotFoundError as legacy_import_error:
81	if not legacy_module.startswith(legacy_import_error.name or ""):
82	raise
83	raise RuntimeError(
84	"ROCM_AITER_MLA requires an AITER build with the small-head "
85	"Gluon MLA kernel (mla_gluon or mla_decode_gluon) when decode "
86	"heads are fewer than 16."
87	) from unified_import_error
88
89
90	@functools.lru_cache(maxsize=1)
91	def _fp8_mla_prefill_supported() -> bool:
92	"""Auto-detect FP8 MLA prefill via mla_prefill_ps_asm_fwd + mla_reduce_v1.
93
94	Requires gfx950 plus an AITER build that exports both kernels. When
95	either is missing we silently fall back to ``flash_attn_varlen_func``.
96	"""
97	try:
98	from vllm.platforms.rocm import on_gfx950
99	except Exception: # noqa: BLE001
100	return False
101	if not on_gfx950():
102	return False
103	try:
104	from aiter import mla_prefill_ps_asm_fwd, mla_reduce_v1 # noqa: F401
105	except Exception: # noqa: BLE001
106	return False
107	return True
108
109
110	@functools.lru_cache(maxsize=1)
111	def _aiter_mla_native_h24_reducer_supported() -> bool:
112	"""Whether AITER's JIT reducer supports the native H24/512 shape."""
113	try:
114	from aiter.jit.core import AITER_CSRC_DIR
115
116	reduce_source = Path(AITER_CSRC_DIR) / "kernels" / "mla" / "reduce.cu"
117	source = "".join(reduce_source.read_text(encoding="utf-8").split())
118	except (ImportError, OSError):
119	return False
120	return "MLA_REDUCE_CASE_EF(NUM_HEAD,24,HEAD_DIM,512," in source
121
122
123	@functools.lru_cache(maxsize=1)
124	def _aiter_mla_native_h24_metadata_supported() -> bool:
125	"""Whether AITER's fast MLA metadata planner accepts native H24.
126
127	The reducer and metadata planner have independent shape dispatch. Checking
128	only the reducer can route H24 into a planner that rejects it before the
129	attention kernel launches. Until AITER exposes a capability API, inspect
130	the shipped JIT source for an explicit native-H24 planner branch.
131	"""
132	try:
133	from aiter.jit.core import AITER_CSRC_DIR
134
135	metadata_source = (
136	Path(AITER_CSRC_DIR) / "kernels" / "mla" / "metadata" / "v1_2_device.cuh"
137	)
138	source = "".join(metadata_source.read_text(encoding="utf-8").split())
139	except (ImportError, OSError):
140	return False
141	return "num_heads==24" in source
142
143
144	def _aiter_mla_native_h24_supported() -> bool:
145	"""Whether the complete AITER decode path supports native H24."""
146	return (
147	_aiter_mla_native_h24_reducer_supported()
148	and _aiter_mla_native_h24_metadata_supported()
149	)
150
151
152	def _aiter_mla_non_causal_asm_kernels() -> bool:
153	"""Whether this arch ships non-causal MLA decode ASM kernels.
154
155	The Python `causal=` probe cannot see the per-arch manifest. gfx950 has
156	the kernels; gfx942 and everything else must fall through to another
157	backend.
158	"""
159	try:
160	from vllm.platforms.rocm import on_gfx950
161	except Exception: # noqa: BLE001
162	return False
163	return bool(on_gfx950())
164
165
166	@functools.lru_cache(maxsize=1)
167	def _gluon_mla_decode_supported() -> bool:
168	"""The small-head Gluon MLA decode kernel only has a gfx950 (CDNA4) build.
169
170	Its tiling needs ~160 KiB of LDS, which exceeds CDNA3's 64 KiB, so on
171	gfx942 there is no kernel to fall through to and selecting it asserts
172	(``mla_gluon requires gfx950``). Restrict Gluon decode to gfx950; other
173	archs use the asm persistent decode, which ``get_mla_padded_q`` makes
174	correct for any 1..15 heads.
175	"""
176	try:
177	from vllm.platforms.rocm import on_gfx950
178	except Exception: # noqa: BLE001
179	return False
180	return on_gfx950()
181
182
183	# Past 2 GiB, mla_gluon swaps its masked buffer_load for an unmasked
184	# global_load that reads outside the cache and aborts the process. Mirrors the
185	# kernel's own within_2gb bound, so the fallback engages exactly where it does.
186	_GLUON_MAX_KV_CACHE_BYTES = 1 << 31
187
188
189	def _gluon_kv_cache_in_bounds(kv_cache_bytes: int \| None) -> bool:
190	"""Whether this layer's KV cache keeps Gluon on its bounds-checked path.
191
192	The cache is one flat tensor per layer, so the bound disqualifies the layer
193	outright; no batch shape brings an oversized cache back in range. ``None``
194	is the not-yet-sized case during profiling, before any kernel runs.
195	"""
196	if kv_cache_bytes is None or kv_cache_bytes <= _GLUON_MAX_KV_CACHE_BYTES:
197	return True
198	logger.warning_once(
199	"KV cache is %.1f GiB per layer, past the 2 GiB where the Gluon MLA "
200	"kernel drops its KV bounds mask; using the padded ASM decode instead. "
201	"Lower --gpu-memory-utilization or --max-model-len to get Gluon back.",
202	kv_cache_bytes / (1 << 30),
203	)
204	return False
205
206
207	@functools.lru_cache(maxsize=1)
208	def _get_segmented_mla_decode():
209	"""Load AITER's segmented MLA decode with unreduced partial output."""
210	from aiter.ops.triton.attention.mla import mla_decode_fwd
211
212	return mla_decode_fwd
213
214
215	@functools.lru_cache(maxsize=1)
216	def _get_aiter_mla_decode():
217	from aiter.mla import mla_decode_fwd
218
219	return mla_decode_fwd
220
221
222	@functools.lru_cache(maxsize=1)
223	def _segmented_mla_decode_supported() -> bool:
224	"""Whether AITER exposes the segmented MLA decode used by DCP verify."""
225	try:
226	_get_segmented_mla_decode()
227	except Exception: # noqa: BLE001
228	return False
229	return True
230
231
232	def _segmented_dcp_verify_supported(dcp_world_size: int, cp_interleave: int) -> bool:
233	"""Whether this configuration can serve DCP verify on segmented MLA.
234
235	Configuration only -- whether a given batch takes the route additionally
236	depends on its query length. Round-robin interleaving other than 1 is
237	excluded because the per-row causal window has not been validated there.
238	"""
239	return (
240	dcp_world_size > 1 and cp_interleave == 1 and _segmented_mla_decode_supported()
241	)
242
243
244	# DCP decode routing
245	#
246	# With decode context parallelism each rank holds a round-robin shard of the KV
247	# cache (token t on rank t % dcp_world_size; requires interleave size 1) and
248	# attends over it with all DCP-gathered query heads. Partial outputs are merged
249	# across ranks by LSE. Each decode batch takes one route:
250	#
251	# * plain: qlen == 1 or a non-causal block. Every query sees every local token.
252	# * cprr: causal qlen >= _MIN_CPRR_QLEN, gfx950, speculative decoding. AITER's
253	# round-robin ASM decode rebuilds global positions from g_kv_indptr, cp_rank
254	# and cp_world_size and applies the causal window in-kernel. Kernels exist only
255	# at _NATIVE_CPRR_HEADS; other gathered head counts are padded up.
256	# * segmented: causal qlen > 1 when cprr is not selected or qlen is below
257	# _MIN_CPRR_QLEN. Each verify token becomes a single-query row whose KV length
258	# encodes its causal window.
259	#
260	# The plain kernel must never serve a causal qlen > 1 block: it applies
261	# causality on local indices, which is wrong over a round-robin shard.
262	_NATIVE_CPRR_HEADS: Final = (16, 32, 64, 128)
263	_MIN_CPRR_QLEN: Final = 3
264
265
266	class _DCPDecodeRoute(Enum):
267	PLAIN = "plain"
268	SEGMENTED = "segmented"
269	CPRR = "cprr"
270
271
272	def _asm_dcp_verify_heads(decode_num_heads: int) -> int:
273	"""Head count the cprr asm decode runs at, or 0 if it cannot serve this shape."""
274	if decode_num_heads in _NATIVE_CPRR_HEADS:
275	return decode_num_heads
276	for nat in _NATIVE_CPRR_HEADS:
277	if nat >= decode_num_heads:
278	return nat
279	return 0
280
281
282	def _select_dcp_decode_route(
283	supports_segmented: bool,
284	causal: bool,
285	max_qo_len: int,
286	asm_selected: bool,
287	) -> _DCPDecodeRoute:
288	"""Select one DCP decode route for this batch; see DCP decode routing."""
289	if not causal or max_qo_len <= 1:
290	return _DCPDecodeRoute.PLAIN
291	if asm_selected and max_qo_len >= _MIN_CPRR_QLEN:
292	return _DCPDecodeRoute.CPRR
293	if supports_segmented:
294	return _DCPDecodeRoute.SEGMENTED
295	raise RuntimeError(
296	"ROCM_AITER_MLA DCP multi-token verify requires either segmented "
297	"MLA or the round-robin asm decode."
298	)
299
300
301	def _asm_dcp_verify_configured(
302	dcp_world_size: int, cp_interleave: int, multi_token_decode: bool
303	) -> bool:
304	"""Whether asm cprr DCP verify is reachable in this process.
305
306	Requires DCP with interleave size 1, speculative decoding, and gfx950 (the
307	only arch AITER builds cprr kernels for).
308	"""
309	if not (dcp_world_size > 1 and cp_interleave == 1 and multi_token_decode):
310	return False
311	from vllm.platforms.rocm import on_gfx950
312
313	return bool(on_gfx950())
314
315
316	def _aiter_mla_small_head_mode() -> str:
317	"""Small-head (<16) MLA decode kernel selection.
318
319	Controlled by ``VLLM_ROCM_AITER_MLA_ASM_PADDING``:
320
321	- ``"auto"`` (default): let the arch decide -- divisor head counts keep the
322	Gluon decode where a build exists (gfx950), everything else (non-divisor
323	counts and all counts on gfx942) uses the padded persistent-scheduling
324	ASM decode.
325	- ``"gluon"``: prefer the Gluon path wherever a build exists.
326	- ``"asm"``: force the padded persistent-scheduling ASM decode.
327
328	On gfx942 (no Gluon build) the ASM path is always used regardless of this
329	setting; ``"gluon"`` there falls back to ASM with a one-time warning.
330	"""
331	import vllm.envs as envs
332
333	mode = (envs.VLLM_ROCM_AITER_MLA_ASM_PADDING or "auto").lower()
334	if mode == "gluon" and not _gluon_mla_decode_supported():
335	logger.warning_once(
336	"VLLM_ROCM_AITER_MLA_ASM_PADDING=gluon requested, but this device "
337	"has no Gluon MLA decode build (Gluon requires gfx950); using the "
338	"padded persistent-scheduling ASM decode instead."
339	)
340	return mode
341
342
343	class AiterMLABackend(MLACommonBackend):
344	supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
345	supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
346	"auto",
347	"float16",
348	"bfloat16",
349	"fp8",
350	"fp8_e4m3",
351	"fp8_e5m2",
352	]
353
354	@classmethod
355	def get_supported_head_sizes(cls) -> list[int]:
356	return []
357
358	@staticmethod
359	def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int \| MultipleOf]:
360	# The aiter MLA decode kernel always operates with page_size=1
361	# internally (the wrapper flattens kv_buffer via .view(-1, 1, 1, H)).
362	# We support any kernel_block_size by expanding block-level indices
363	# into per-token flat indices in the metadata builder.
364	return [MultipleOf(1)]
365
366	@staticmethod
367	def get_name() -> str:
368	return "ROCM_AITER_MLA"
369
370	@classmethod
371	def supports_non_causal(cls) -> bool:
372	# causal= on mla_decode_fwd is necessary but not sufficient: pinned
373	# AITER v0.1.21.post2 still exports that argument on gfx942, whose ASM
374	# manifest has no non-causal decode kernels. Selecting this backend
375	# there reaches `cannot get heuristic kernel!` instead of TRITON_MLA.
376	return bool(
377	rocm_aiter_ops.mla_decode_supports_non_causal()
378	and _aiter_mla_non_causal_asm_kernels()
379	)
380
381	@classmethod
382	def validate_configuration(
383	cls,
384	head_size: int,
385	dtype: torch.dtype,
386	kv_cache_dtype: CacheDType \| None,
387	block_size: int \| None,
388	use_mla: bool,
389	has_sink: bool,
390	use_sparse: bool,
391	use_mm_prefix: bool,
392	use_per_head_quant_scales: bool,
393	device_capability: "DeviceCapability",
394	attn_type: str,
395	has_sliding_window: bool = False,
396	use_non_causal: bool = False,
397	use_batch_invariant: bool = False,
398	use_kv_connector: bool = False,
399	use_pcp: bool = False,
400	use_adaptive_verification: bool = False,
401	use_dcp: bool = False,
402	use_rswa: bool = False,
403	) -> list[str]:
404	invalid_reasons = super().validate_configuration(
405	head_size=head_size,
406	dtype=dtype,
407	kv_cache_dtype=kv_cache_dtype,
408	block_size=block_size,
409	use_mla=use_mla,
410	has_sink=has_sink,
411	use_sparse=use_sparse,
412	use_mm_prefix=use_mm_prefix,
413	use_per_head_quant_scales=use_per_head_quant_scales,
414	device_capability=device_capability,
415	attn_type=attn_type,
416	has_sliding_window=has_sliding_window,
417	use_non_causal=use_non_causal,
418	use_batch_invariant=use_batch_invariant,
419	use_kv_connector=use_kv_connector,
420	use_pcp=use_pcp,
421	use_adaptive_verification=use_adaptive_verification,
422	use_dcp=use_dcp,
423	use_rswa=use_rswa,
424	)
425	# Arch/signature cannot express this: the backend still advertises
426	# fp16, so a DSpark draft with kv_cache_dtype="auto" would select
427	# AITER and abort on fp16 query. Causal fp16 is fine.
428	if (
429	use_non_causal
430	and dtype == torch.float16
431	and kv_cache_dtype in (None, "auto", "float16")
432	):
433	invalid_reasons.append("non-causal fp16 MLA decode not supported")
434	return invalid_reasons
435
436	@staticmethod
437	def get_impl_cls() -> type["AiterMLAImpl"]:
438	return AiterMLAImpl
439
440	@staticmethod
441	def get_builder_cls() -> type["AiterMLAMetadataBuilder"]:
442	return AiterMLAMetadataBuilder
443
444
445	@dataclass
446	class AiterMLADCPVerifyMetadata:
447	"""One paged-KV row per verify token, for segmented DCP verification.
448
449	These are produced together or not at all, so they travel as one value: its
450	presence on the decode metadata is the routing decision the builder made,
451	and the impl does not re-derive it.
452	"""
453
454	# Local KV length of each verify row, already causally bounded.
455	row_lens: torch.Tensor
456	# Subpage IDs each row reads, shape [num_rows, max_local_pages].
457	block_table: torch.Tensor
458	# One query per row, so this is an arange.
459	qo_indptr: torch.Tensor
460	# Subpage size the block table is expressed in, and the segmented kernel's
461	# TILE_SIZE. Carried so the impl reinterprets the cache exactly the way the
462	# block table was built, instead of re-deriving it from the cache shape.
463	page_size: int
464	# Kernel-visible KV bound. Whenever full graphs are enabled this is the
465	# configuration's maximum for every batch, not just during capture, so the
466	# page table keeps one shape across replays.
467	max_kv_seq_len: int
468
469
470	@dataclass
471	class AiterMLADecodeMetadata(MLACommonDecodeMetadata):
472	# The indptr of the paged kv cache, shape: [batch_size + 1]
473	paged_kv_indptr: torch.Tensor \| None = None
474	# The page indices of the paged kv cache
475	paged_kv_indices: torch.Tensor \| None = None
476	# The number of entries in the last page of each request in
477	# the paged kv cache, shape: [batch_size]
478	paged_kv_last_page_len: torch.Tensor \| None = None
479	# The query indptr, shape : [num_decode + 1]
480	qo_indptr: torch.Tensor \| None = None
481	# The dtype of MLA out tensor
482	attn_out_dtype: torch.dtype = torch.bfloat16
483	# The max query output length: int
484	max_qo_len: int \| None = None
485	# Minimum KV length used by Gluon to choose a safe split count.
486	min_kv_seq_len: int = 1
487	# Exactly one DCP decode route is selected by the metadata builder.
488	dcp_route: _DCPDecodeRoute = _DCPDecodeRoute.PLAIN
489	# Set exactly when this batch routes to segmented DCP verification.
490	dcp_verify: AiterMLADCPVerifyMetadata \| None = None
491	# cprr only: global per-request page indptr and this rank's position.
492	g_kv_indptr: torch.Tensor \| None = None
493	cp_world_size: int = 1
494	cp_rank: int = 0
495	# Padded head count cprr runs at; 0 when this batch is not on cprr.
496	asm_decode_num_heads: int = 0
497	# Must equal the split cap the reduce scratch was sized with, or the kernel
498	# writes past it and faults.
499	mla_num_kv_splits: int = 0
500	# Small-head decode uses Gluon (avoids padding to 16).
501	use_gluon_decode: bool = False
502	# Small-head non-DCP multi-token verify uses Gluon's native MTP entry.
503	use_gluon_verify: bool = False
504	# Whether persistent MLA metadata was computed
505	has_persistent_metadata: bool = False
506
507
508	@dataclass
509	class AiterMLAMetadata(MLACommonMetadata[AiterMLADecodeMetadata]):
510	work_meta_data: torch.Tensor \| None = None
511	work_indptr: torch.Tensor \| None = None
512	work_info_set: torch.Tensor \| None = None
513	reduce_indptr: torch.Tensor \| None = None
514	reduce_final_map: torch.Tensor \| None = None
515	reduce_partial_map: torch.Tensor \| None = None
516
517	# FP8 ASM prefill persistent-scheduling (PS) metadata. Populated by
518	# AiterMLAMetadataBuilder._build_fp8_prefill_ps_metadata when prefill
519	# tokens are present and FP8 MLA prefill is supported on the device.
520	# Left as None on hosts/configs that fall back to flash_attn_varlen_func.
521	fp8_prefill_qo_indptr: torch.Tensor \| None = None
522	fp8_prefill_kv_indptr: torch.Tensor \| None = None
523	fp8_prefill_kv_indices: torch.Tensor \| None = None
524	fp8_prefill_work_indptr: torch.Tensor \| None = None
525	fp8_prefill_work_info_set: torch.Tensor \| None = None
526	fp8_prefill_reduce_indptr: torch.Tensor \| None = None
527	fp8_prefill_reduce_final_map: torch.Tensor \| None = None
528	fp8_prefill_reduce_partial_map: torch.Tensor \| None = None
529	fp8_prefill_max_q_len: int \| None = None
530	fp8_prefill_num_partial_tiles: int \| None = None
531
532
533	# Tile size used by the mla_prefill_ps_asm_fwd assembly kernel.
534	_FP8_PREFILL_TILE_Q = 256
535
536
537	class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
538	# TODO(luka, lucas): audit this as part of:
539	# https://github.com/vllm-project/vllm/issues/22945
540	_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
541	query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
542	# Served by passing the mask to the kernel; _build_decode turns away the
543	# shapes AITER has no non-causal kernel for.
544	supports_non_causal_multi_token_decode: ClassVar[bool] = True
545	# Set from the common metadata every build; a batch is causal unless the
546	# drafter says otherwise.
547	_decode_causal: bool = True
548	# A non-causal draft needs no cross-shard causal window: the ordinary
549	# mask0 decode returns per-row LSE for the DCP merge.
550	supports_non_causal_multi_token_dcp: ClassVar[bool] = True
551
552	@staticmethod
553	def _uniform_padded_mtp_qo_len(
554	qo_len: torch.Tensor,
555	max_qo_len: int,
556	num_decode_tokens: int,
557	) -> int:
558	num_reqs = qo_len.numel()
559	if num_reqs == 0 or num_decode_tokens <= 0:
560	return 0
561
562	# Full-CG pads q to a captured token count while leaving
563	# query_start_loc flat for dummy requests. Only synthesize dummy rows
564	# when every padded request maps to the same qlen and the q buffer has
565	# exactly that many rows.
566	if num_decode_tokens <= int(qo_len.sum().item()):
567	return 0
568	if num_decode_tokens % num_reqs != 0:
569	return 0
570
571	uniform_qo_len = num_decode_tokens // num_reqs
572	if uniform_qo_len <= 1:
573	return 0
574
575	positive_qo_len = qo_len[qo_len > 0]
576	if positive_qo_len.numel() == qo_len.numel():
577	return 0
578	if positive_qo_len.numel() > 0:
579	if max_qo_len != uniform_qo_len:
580	return 0
581	if not torch.all(positive_qo_len == uniform_qo_len):
582	return 0
583
584	zero_positions = torch.nonzero(qo_len == 0, as_tuple=False).flatten()
585	if zero_positions.numel() > 0:
586	first_zero = int(zero_positions[0].item())
587	if torch.any(qo_len[first_zero:] > 0):
588	return 0
589
590	return uniform_qo_len
591
592	def __init__(
593	self,
594	kv_cache_spec: AttentionSpec,
595	layer_names: list[str],
596	vllm_config: VllmConfig,
597	device: torch.device,
598	):
599	parallel_config = vllm_config.parallel_config
600	supports_segmented_dcp_verify = _segmented_dcp_verify_supported(
601	parallel_config.decode_context_parallel_size,
602	parallel_config.cp_kv_cache_interleave_size,
603	)
604	asm_dcp_verify_config = (
605	_asm_dcp_verify_configured(
606	parallel_config.decode_context_parallel_size,
607	parallel_config.cp_kv_cache_interleave_size,
608	multi_token_decode=vllm_config.speculative_config is not None,
609	)
610	and envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY == "asm"
611	)
612	super().__init__(
613	kv_cache_spec,
614	layer_names,
615	vllm_config,
616	device,
617	AiterMLAMetadata,
618	supports_dcp_with_varlen=(
619	supports_segmented_dcp_verify or asm_dcp_verify_config
620	),
621	)
622	self._asm_dcp_verify = False
623	self._asm_dcp_verify_heads = 0
624	# The head-count half of the route decision needs self.num_heads, so only
625	# the config half is known before super().__init__().
626	if asm_dcp_verify_config:
627	self._asm_dcp_verify_heads = _asm_dcp_verify_heads(
628	self.num_heads * self.dcp_world_size
629	)
630	if not self._asm_dcp_verify_heads:
631	raise ValueError(
632	"VLLM_ROCM_AITER_MLA_DCP_VERIFY=asm, but the round-robin asm "
633	f"decode has no kernel for {self.num_heads * self.dcp_world_size} "
634	f"DCP-gathered heads (native counts: {_NATIVE_CPRR_HEADS}). "
635	"Set VLLM_ROCM_AITER_MLA_DCP_VERIFY=segmented."
636	)
637	self._asm_dcp_verify = True
638	self._supports_segmented_dcp_verify = supports_segmented_dcp_verify
639	self._mla_max_split_per_batch = 0
640	if self._asm_dcp_verify:
641	# Cap on KV splits per batch; at low batch the KV axis is the only
642	# parallelism. Buffer sizing and every runtime call must see this value.
643	self._mla_max_split_per_batch = torch.cuda.get_device_properties(
644	device
645	).multi_processor_count
646
647	self.compilation_config = vllm_config.compilation_config
648	self.decode_attn_out_dtype = vllm_config.model_config.dtype
649
650	# Needed to place a verify row's causal window on this rank's KV shard.
651	self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
652
653	# reorder_batch_threshold is the largest query length decode can be
654	# handed, and already accounts for the drafting scheme. A method-name
655	# whitelist sizes unlisted drafters for qlen=1, which closes the
656	# persistent gate below and makes aiter raise a KeyError mid-run.
657	self._mtp_decode_qlen = self.reorder_batch_threshold or 1
658
659	# Store the kernel block size from the spec. When kernel_block_size=1
660	# (no spec-dec), behavior is identical to the original. When > 1
661	# (e.g. 16 with Eagle3), we expand block-level indices into per-token
662	# flat indices since the aiter kernel always uses page_size=1 internally.
663	self.kernel_block_size = kv_cache_spec.block_size
664	self._segmented_page_size = _segmented_mla_page_size(self.kernel_block_size)
665
666	# In the flat view (.view(-1,1,1,H)), each token is its own page,
667	# so max_num_pages_per_req = max_model_len regardless of
668	# kernel_block_size.
669	max_num_pages_per_req = vllm_config.model_config.max_model_len
670	max_num_reqs = vllm_config.scheduler_config.max_num_seqs
671	max_num_pages = max_num_reqs * max_num_pages_per_req
672
673	# Preparing persistent buffers
674	# TODO: we can disambiguate between decode and mixed-prefill decode here
675	# so we can only use the persistent buffer if a cudagraph is actually
676	# being used.
677
678	# paged_kv_last_page_len is always 1s (the aiter kernel always sees
679	# page_size=1 after .view(-1,1,1,H) flattening), so we create it
680	# once and reuse slices in both eager and cudagraph modes.
681	self.paged_kv_last_page_len = torch.ones(
682	max_num_reqs, dtype=torch.int32, device=device
683	)
684
685	# Persistent buffer for paged_kv_indices to avoid blocking boolean mask
686	# indexing (block_table_tensor[mask]) which has data-dependent output size.
687	self.paged_kv_indices = torch.zeros(
688	max_num_pages, dtype=torch.int32, device=device
689	)
690
691	from aiter import dtypes, get_mla_metadata_info_v1
692
693	# Decode kernels consume the DCP-gathered query heads.
694	self._decode_num_heads = self.num_heads * self.dcp_world_size
695	# Keep metadata sizing consistent with the padded tensor shape passed
696	# to mla_decode_fwd, including native 24-head AITER builds.
697	self._num_attention_heads = AiterMLAHelper.get_actual_mla_num_heads(
698	self._decode_num_heads
699	)
700	# the cprr kernel exists only at the native head
701	# counts, so run (and size the persistent schedule) at the padded count.
702	# 96 is 16-aligned, so get_actual_mla_num_heads leaves it alone and the
703	# pad has to be applied here explicitly.
704	if self._asm_dcp_verify:
705	self._num_attention_heads = self._asm_dcp_verify_heads
706	kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto")
707	if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"):
708	kv_cache_dtype_str = "fp8"
709	kv_dtype = dtypes.fp8
710	else:
711	kv_dtype = {
712	torch.float16: dtypes.fp16,
713	torch.bfloat16: dtypes.bf16,
714	}[kv_cache_spec.dtype]
715	# _build_decode needs the cache dtype to pick the decode kernel; keep
716	# the normalized string instead of dropping it at the end of __init__.
717	self._kv_cache_dtype_str = kv_cache_dtype_str
718	# Sized per layer, matching the flat KV tensor the kernel is handed.
719	num_gpu_blocks = vllm_config.cache_config.num_gpu_blocks
720	self._kv_cache_bytes = (
721	None
722	if num_gpu_blocks is None
723	else num_gpu_blocks * kv_cache_spec.page_size_bytes
724	)
725	# MLAAttention quantizes decode Q to FP8 before calling this backend
726	# whenever the KV cache is FP8 and supports_quant_query_input is true.
727	q_dtype = (
728	dtypes.fp8 if kv_cache_dtype_str == "fp8" else self.decode_attn_out_dtype
729	)
730	# Persist for get_mla_metadata_v1 (decode build): omitting these causes
731	# wrong split/reduce metadata for the gfx950 fp8 nhead=32 fold path.
732	self._mla_q_dtype = q_dtype
733	self._mla_kv_dtype = kv_dtype
734	(
735	(work_meta_data_size, work_meta_data_type),
736	(work_indptr_size, work_indptr_type),
737	(work_info_set_size, work_info_set_type),
738	(reduce_indptr_size, reduce_indptr_type),
739	(reduce_final_map_size, reduce_final_map_type),
740	(reduce_partial_map_size, reduce_partial_map_type),
741	) = get_mla_metadata_info_v1(
742	max_num_reqs,
743	self._mtp_decode_qlen,
744	self._num_attention_heads,
745	q_dtype,
746	kv_dtype,
747	is_sparse=False,
748	fast_mode=True,
749	# Must match max_split_per_batch in get_mla_metadata_v1; without it
750	# aiter sizes the reduce scratch from capacity (~9 GiB).
751	**(
752	dict(max_split_per_batch=self._mla_max_split_per_batch)
753	if self._asm_dcp_verify
754	else {}
755	),
756	)
757	self._mla_work_meta_data = torch.empty(
758	work_meta_data_size, dtype=work_meta_data_type, device=device
759	)
760	self._mla_work_indptr = torch.empty(
761	work_indptr_size, dtype=work_indptr_type, device=device
762	)
763	self._mla_work_info_set = torch.empty(
764	work_info_set_size, dtype=work_info_set_type, device=device
765	)
766	self._mla_reduce_indptr = torch.empty(
767	reduce_indptr_size, dtype=reduce_indptr_type, device=device
768	)
769	self._mla_reduce_final_map = torch.empty(
770	reduce_final_map_size, dtype=reduce_final_map_type, device=device
771	)
772	self._mla_reduce_partial_map = torch.empty(
773	reduce_partial_map_size,
774	dtype=reduce_partial_map_type,
775	device=device,
776	)
777
778	# The assembly prefill requires FP8 KV, bf16 output, and 16-aligned
779	# heads. It writes bf16 through a raw output pointer, so fp16 must use
780	# the standard prefill path. Head counts that are not a multiple of 16
781	# are replicate-padded up to one in _mla_fp8_prefill_attn, so the gate
782	# is the same head-count predicate the decode path uses.
783	self._fp8_prefill_enabled = _fp8_mla_prefill_supported() and (
784	kv_cache_dtype_str == "fp8"
785	and vllm_config.model_config.dtype == torch.bfloat16
786	and AiterMLAHelper.is_valid_num_heads(self.num_heads)
787	)
788	if self._fp8_prefill_enabled:
789	max_prefill_qlen = min(
790	vllm_config.model_config.max_model_len,
791	vllm_config.scheduler_config.max_num_batched_tokens,
792	)
793	self._init_fp8_prefill_ps_buffers(
794	max_num_reqs,
795	max_prefill_qlen,
796	vllm_config.scheduler_config.max_num_batched_tokens,
797	vllm_config.model_config.dtype,
798	device,
799	)
800
801	# Full cudagraphs need a stable address for the cprr global page indptr.
802	# Zero-initialized; element 0 is the indptr base and is never rewritten.
803	self._g_kv_indptr_buf: torch.Tensor \| None = None
804	if self._asm_dcp_verify:
805	self._g_kv_indptr_buf = torch.zeros(
806	max_num_reqs + 1, dtype=torch.int32, device=device
807	)
808
809	# Persistent buffers for segmented DCP verification. Captured graphs
810	# require stable addresses while row lengths vary between replays.
811	self._dcp_verify_buffers: AiterMLADCPVerifyMetadata \| None = None
812	self._graph_seq_lens: torch.Tensor \| None = None
813	if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
814	self.paged_kv_indptr = torch.zeros(
815	max_num_reqs + 1, dtype=torch.int32, device=device
816	)
817
818	self.qo_indptr = torch.zeros(
819	max_num_reqs + 1, dtype=torch.int32, device=device
820	)
821
822	# Full graphs require a stable address after uniform-MTP padding.
823	self._graph_seq_lens = torch.zeros(
824	max_num_reqs, dtype=torch.int32, device=device
825	)
826
827	if self._supports_segmented_dcp_verify and self._mtp_decode_qlen > 1:
828	# Allocate even when CPRR is the preferred route: a later
829	# replay can fall below _MIN_CPRR_QLEN and needs these
830	# addresses. A DCP rank's shard of the longest sequence
831	# bounds every verify row, and full graphs need that bound
832	# to be constant.
833	num_dcp_partitions = (
834	self.dcp_world_size * self.cp_kv_cache_interleave_size
835	)
836	graph_max_kv_seq_len = (
837	cdiv(vllm_config.model_config.max_model_len, num_dcp_partitions)
838	* self.cp_kv_cache_interleave_size
839	)
840	max_verify_rows = max_num_reqs * self._mtp_decode_qlen
841	max_local_pages = cdiv(graph_max_kv_seq_len, self._segmented_page_size)
842	self._dcp_verify_buffers = AiterMLADCPVerifyMetadata(
843	row_lens=torch.zeros(
844	max_verify_rows, dtype=torch.int32, device=device
845	),
846	block_table=torch.zeros(
847	(max_verify_rows, max_local_pages),
848	dtype=torch.int32,
849	device=device,
850	),
851	qo_indptr=torch.arange(
852	max_verify_rows + 1, dtype=torch.int32, device=device
853	),
854	page_size=self._segmented_page_size,
855	max_kv_seq_len=graph_max_kv_seq_len,
856	)
857
858	def _init_fp8_prefill_ps_buffers(
859	self,
860	max_num_reqs: int,
861	max_prefill_qlen: int,
862	max_num_batched_tokens: int,
863	attn_out_dtype: torch.dtype,
864	device: torch.device,
865	) -> None:
866	"""Pre-allocate persistent buffers for FP8 MLA prefill PS metadata.
867
868	Uses ``get_ps_metadata_info_v1`` with max values so the buffers are
869	large enough for any batch. ``get_ps_metadata_v1`` fills them
870	per-batch in ``build()``. The FP8 prefill forward path also uses the
871	global workspace manager for per-call scratch, so reserve its maximum
872	shape here before the workspace manager is locked after warmup.
873
874	Args:
875	max_num_reqs: Maximum number of concurrent requests.
876	max_prefill_qlen: Maximum Q-length for a single request in one
877	prefill batch. Should be ``min(max_model_len,
878	max_num_batched_tokens)`` — a single request never exceeds
879	``max_model_len`` tokens, nor the per-batch token budget.
880	max_num_batched_tokens: Maximum number of tokens scheduled in one
881	batch. The ``final_lse`` scratch is sized by ``total_q`` (the
882	summed Q-length over all prefill requests in the batch), which
883	is bounded by this budget rather than by a single request's
884	``max_prefill_qlen`` — concurrent requests can sum to more than
885	``max_model_len`` when ``max_model_len < max_num_batched_tokens``.
886	attn_out_dtype: Dtype of the attention output buffer, used to size
887	the padded-head output scratch (small head counts only).
888	device: Target device for the buffers.
889
890	"""
891	from aiter import get_ps_metadata_info_v1
892
893	# After kv_b_proj decompression, K has num_heads heads (same as Q).
894	# So gqa_ratio=1 and num_head_k=num_heads for the PS kernel.
895	# Head counts that are not a multiple of 16 (K3: 12/rank at TP8) are
896	# replicate-padded up to one in _mla_fp8_prefill_attn; build the PS
897	# metadata for that same padded count so the work/reduce maps and the
898	# scratch reservations describe the width the kernel is handed.
899	#
900	# This was previously max(16, num_heads), which agrees with the helper
901	# at every head count reachable today (12 and 16 both give 16) and
902	# differs only above 16: at 24 it leaves num_head_k=24 while the
903	# forward pads to 32. That is not a correctness bug -- the 24-wide work
904	# maps still cover head-tiles 0..23, which are the real heads -- but it
905	# is expensive, because a lower head alignment yields more partial
906	# tiles: gcd-driven, 24 heads -> 64 tiles vs 32 heads -> 16, i.e.
907	# 193.6 MiB of reservations instead of 68.6 MiB at batch=1/qlen=512
908	# (measured on gfx950). Sizing both from one helper keeps the widths
909	# equal and takes the cheaper tiling.
910	num_head_k = AiterMLAHelper.get_fp8_prefill_num_heads(self.num_heads)
911	v_head_dim = self.mla_dims.v_head_dim
912	# gqa_ratio = 1
913	# qlen_granularity = _FP8_PREFILL_TILE_Q // max(gqa_ratio, 1)
914	qlen_granularity = _FP8_PREFILL_TILE_Q
915
916	(
917	(work_metadata_size, work_metadata_dtype),
918	(work_indptr_size, work_indptr_dtype),
919	(work_info_size, work_info_dtype),
920	(reduce_indptr_size, reduce_indptr_dtype),
921	(reduce_final_map_size, reduce_final_map_dtype),
922	(reduce_partial_map_size, reduce_partial_map_dtype),
923	) = get_ps_metadata_info_v1(
924	batch_size=max_num_reqs,
925	num_head_k=num_head_k,
926	max_qlen=max_prefill_qlen,
927	qlen_granularity=qlen_granularity,
928	)
929
930	self.fp8_ps_work_metadata = torch.empty(
931	work_metadata_size, dtype=work_metadata_dtype, device=device
932	)
933	self.fp8_ps_work_indptr = torch.empty(
934	work_indptr_size, dtype=work_indptr_dtype, device=device
935	)
936	self.fp8_ps_work_info = torch.empty(
937	*work_info_size, dtype=work_info_dtype, device=device
938	)
939	self.fp8_ps_reduce_indptr = torch.empty(
940	reduce_indptr_size, dtype=reduce_indptr_dtype, device=device
941	)
942	self.fp8_ps_reduce_final_map = torch.empty(
943	*reduce_final_map_size, dtype=reduce_final_map_dtype, device=device
944	)
945	self.fp8_ps_reduce_partial_map = torch.empty(
946	reduce_partial_map_size,
947	dtype=reduce_partial_map_dtype,
948	device=device,
949	)
950
951	# get_ps_metadata_v1 builds the plan on the host and writes its outputs with
952	# blocking hipMemcpy that is not ordered on the current stream. Writing the
953	# device buffers directly lets step N+1's build (async scheduling) overwrite a
954	# plan that step N's prefill kernels are still reading. Build into pinned host
955	# staging instead and copy on the current stream. Two staging slots, each
956	# reused only after its previous H2D copy has completed.
957	self._fp8_ps_device_outputs = (
958	self.fp8_ps_work_indptr,
959	self.fp8_ps_work_info,
960	self.fp8_ps_reduce_indptr,
961	self.fp8_ps_reduce_final_map,
962	self.fp8_ps_reduce_partial_map,
963	)
964	self._fp8_ps_staging = [
965	tuple(
966	torch.empty(t.shape, dtype=t.dtype, device="cpu", pin_memory=True)
967	for t in self._fp8_ps_device_outputs
968	)
969	for _ in range(2)
970	]
971	# One event per slot, recorded after that slot's H2D copy. synchronize()
972	# on an event that has not been recorded yet returns immediately.
973	self._fp8_ps_staging_free = [torch.cuda.Event(), torch.cuda.Event()]
974	self._fp8_ps_slot = 0
975
976	from vllm.platforms import current_platform
977	from vllm.v1.worker.workspace import current_workspace_manager
978
979	# AITER metadata sizing assumes all requests can carry max_prefill_qlen tokens,
980	# which is a loose worst case in the number of QO tiles. Rather, the sum of
981	# qlens is bounded by max_num_batched_tokens. Manually compute the number of
982	# qo tiles to avoid OOM at startup.
983	# TODO: AITER should give us this budget constrained value
984	qo_tile_cnt = (
985	cdiv(max_num_batched_tokens, _FP8_PREFILL_TILE_Q) + max_num_reqs - 1
986	)
987	max_num_partial_tiles = qo_tile_cnt + current_platform.num_compute_units()
988	reservations: list[tuple[tuple[int, ...], torch.dtype]] = [
989	(
990	(max_num_partial_tiles * _FP8_PREFILL_TILE_Q, num_head_k, v_head_dim),
991	torch.float32,
992	),
993	(
994	(max_num_partial_tiles * _FP8_PREFILL_TILE_Q, num_head_k),
995	torch.float32,
996	),
997	((max_num_batched_tokens, num_head_k), torch.float32),
998	]
999	if self.num_heads < num_head_k:
1000	# Padded head counts also take their kernel output buffer from the
1001	# workspace: the caller's [total_q, num_heads * v_head_dim] output
1002	# cannot back a num_head_k-head view (see _mla_fp8_prefill_attn).
1003	reservations.append(
1004	((max_num_batched_tokens, num_head_k, v_head_dim), attn_out_dtype)
1005	)
1006	current_workspace_manager().get_simultaneous(*reservations)
1007
1008	logger.info(
1009	"FP8 MLA prefill PS buffers allocated "
1010	"(max_batch=%d, max_qlen=%d, num_head_k=%d)",
1011	max_num_reqs,
1012	max_prefill_qlen,
1013	num_head_k,
1014	)
1015
1016	def _build_fp8_prefill_ps_metadata(
1017	self,
1018	metadata: AiterMLAMetadata,
1019	common_attn_metadata: CommonAttentionMetadata,
1020	) -> None:
1021	"""Build per-batch FP8 MLA prefill PS metadata and attach to metadata.
1022
1023	Called from ``build()`` when prefill tokens are present and
1024	FP8 MLA prefill is enabled (auto-detected via
1025	``_fp8_mla_prefill_supported()``).
1026	"""
1027	from aiter import get_ps_metadata_v1
1028
1029	prefill = metadata.prefill
1030	# Caller (build()) only invokes this when prefill tokens exist, so
1031	# metadata.prefill is guaranteed non-None. Assert to narrow for mypy.
1032	assert prefill is not None
1033	qo_indptr = prefill.query_start_loc
1034	kv_indptr = qo_indptr # new tokens: KV length == Q length
1035
1036	# Reuse the existing CPU view of query_start_loc instead of forcing a
1037	# device->host copy. Prefill batches sit at the tail of the request
1038	# list, so we slice from num_decodes onwards and rebase to zero, the
1039	# same transform the parent build applies on device tensors.
1040	num_decodes = metadata.num_decodes
1041	qsl_cpu = common_attn_metadata.query_start_loc_cpu
1042	qo_indptr_cpu = (qsl_cpu[num_decodes:] - qsl_cpu[num_decodes]).to(torch.int32)
1043	kv_indptr_cpu = qo_indptr_cpu.clone()
1044	seq_lens_cpu = (qo_indptr_cpu[1:] - qo_indptr_cpu[:-1]).to(torch.int32)
1045
1046	# Head counts that are not a multiple of 16 (K3: 12/rank at TP8) are
1047	# replicate-padded up to one in _mla_fp8_prefill_attn; build the PS
1048	# metadata for that same padded count so the work/reduce maps and the
1049	# scratch reservations describe the width the kernel is handed.
1050	#
1051	# This was previously max(16, num_heads), which agrees with the helper
1052	# at every head count reachable today (12 and 16 both give 16) and
1053	# differs only above 16: at 24 it leaves num_head_k=24 while the
1054	# forward pads to 32. That is not a correctness bug -- the 24-wide work
1055	# maps still cover head-tiles 0..23, which are the real heads -- but it
1056	# is expensive, because a lower head alignment yields more partial
1057	# tiles: gcd-driven, 24 heads -> 64 tiles vs 32 heads -> 16, i.e.
1058	# 193.6 MiB of reservations instead of 68.6 MiB at batch=1/qlen=512
1059	# (measured on gfx950). Sizing both from one helper keeps the widths
1060	# equal and takes the cheaper tiling.
1061	num_head_k = AiterMLAHelper.get_fp8_prefill_num_heads(self.num_heads)
1062	# gqa_ratio = 1
1063	# qhead_granularity = max(gqa_ratio, 1)
1064	# qlen_granularity = _FP8_PREFILL_TILE_Q // qhead_granularity
1065	gqa_ratio = 1
1066	qhead_granularity = 1
1067	qlen_granularity = _FP8_PREFILL_TILE_Q
1068	kvlen_granularity = 128
1069	block_size = 1 # non-paged: each "page" is one token
1070
1071	slot = self._fp8_ps_slot
1072	self._fp8_ps_slot ^= 1
1073	staging_free = self._fp8_ps_staging_free[slot]
1074	with gpu_sync_allowed():
1075	staging_free.synchronize()
1076	(
1077	work_indptr_host,
1078	work_info_host,
1079	reduce_indptr_host,
1080	reduce_final_map_host,
1081	reduce_partial_map_host,
1082	) = self._fp8_ps_staging[slot]
1083
1084	# work_metadata is not written by the host planner; it stays on device.
1085	get_ps_metadata_v1(
1086	qo_indptr_cpu,
1087	kv_indptr_cpu,
1088	seq_lens_cpu,
1089	gqa_ratio,
1090	num_head_k,
1091	self.fp8_ps_work_metadata,
1092	work_indptr_host,
1093	work_info_host,
1094	reduce_indptr_host,
1095	reduce_final_map_host,
1096	reduce_partial_map_host,
1097	qhead_granularity=qhead_granularity,
1098	qlen_granularity=qlen_granularity,
1099	kvlen_granularity=kvlen_granularity,
1100	block_size=block_size,
1101	is_causal=True,
1102	)
1103	for dst, src in zip(self._fp8_ps_device_outputs, self._fp8_ps_staging[slot]):
1104	dst.copy_(src, non_blocking=True)
1105	staging_free.record()
1106
1107	total_prefill_tokens = int(qo_indptr_cpu[-1].item())
1108	kv_indices = torch.arange(
1109	total_prefill_tokens, device=qo_indptr.device, dtype=torch.int32
1110	)
1111
1112	# The actual number of active partial tiles for this batch is the
1113	# final value of reduce_indptr, read from the host plan (no GPU sync).
1114	num_partial_tiles = int(reduce_indptr_host[-1])
1115
1116	# Attach PS metadata to the metadata object so forward_mha can read it.
1117	metadata.fp8_prefill_qo_indptr = qo_indptr
1118	metadata.fp8_prefill_kv_indptr = kv_indptr
1119	metadata.fp8_prefill_kv_indices = kv_indices
1120	metadata.fp8_prefill_work_indptr = self.fp8_ps_work_indptr
1121	metadata.fp8_prefill_work_info_set = self.fp8_ps_work_info
1122	metadata.fp8_prefill_reduce_indptr = self.fp8_ps_reduce_indptr
1123	metadata.fp8_prefill_reduce_final_map = self.fp8_ps_reduce_final_map
1124	metadata.fp8_prefill_reduce_partial_map = self.fp8_ps_reduce_partial_map
1125	metadata.fp8_prefill_max_q_len = prefill.max_query_len
1126	metadata.fp8_prefill_num_partial_tiles = num_partial_tiles
1127
1128	def _build_dcp_verify_row_view(
1129	self,
1130	qlen: int,
1131	block_table: torch.Tensor,
1132	dcp_tot_seq_lens: torch.Tensor,
1133	) -> AiterMLADCPVerifyMetadata:
1134	"""Build one paged-KV row per verify token for segmented DCP verification.
1135
1136	Every row is a single query (``qo_indptr`` is an arange), so the whole
1137	causal structure is carried by ``dcp_local_verify_row_lens`` and the
1138	kernel applies no tail of its own.
1139	"""
1140	assert self.dcp_world_size > 1
1141	num_reqs = dcp_tot_seq_lens.numel()
1142	row_lens = AiterMLAHelper.dcp_local_verify_row_lens(
1143	dcp_tot_seq_lens,
1144	qlen,
1145	self.dcp_world_size,
1146	self.dcp_rank,
1147	self.cp_kv_cache_interleave_size,
1148	)
1149	num_rows = row_lens.numel()
1150	buffers = self._dcp_verify_buffers
1151	# Persistent buffers exist exactly when full graphs are on, and they
1152	# already carry the static bound they were sized for.
1153	max_kv_seq_len = (
1154	buffers.max_kv_seq_len
1155	if buffers is not None
1156	else max(1, int(row_lens.max().item()))
1157	)
1158	page_size = self._segmented_page_size
1159	max_local_pages = cdiv(max_kv_seq_len, page_size)
1160	if buffers is not None:
1161	row_block_table = buffers.block_table[:num_rows]
1162	buffers.row_lens[:num_rows].copy_(row_lens, non_blocking=True)
1163	row_lens = buffers.row_lens[:num_rows]
1164	qo_indptr = buffers.qo_indptr[: num_rows + 1]
1165	else:
1166	row_block_table = torch.empty(
1167	(num_rows, max_local_pages),
1168	dtype=torch.int32,
1169	device=block_table.device,
1170	)
1171	qo_indptr = torch.arange(
1172	num_rows + 1,
1173	dtype=torch.int32,
1174	device=block_table.device,
1175	)
1176	self._fill_dcp_verify_page_table(row_block_table, block_table, num_reqs, qlen)
1177	return AiterMLADCPVerifyMetadata(
1178	row_lens=row_lens,
1179	block_table=row_block_table,
1180	qo_indptr=qo_indptr,
1181	page_size=page_size,
1182	max_kv_seq_len=max_kv_seq_len,
1183	)
1184
1185	def _fill_dcp_verify_page_table(
1186	self,
1187	row_block_table: torch.Tensor,
1188	block_table: torch.Tensor,
1189	num_reqs: int,
1190	qlen: int,
1191	) -> None:
1192	"""Expand each request's blocks into the subpages one verify row reads.
1193
1194	Every row of a request shares the request's shard, so the page list is
1195	built once per request and repeated; only ``row_lens`` distinguishes the
1196	rows. A DCP rank holds ``1/dcp_world_size`` of the sequence, so the
1197	request's block table is always wider than the pages a row can reach.
1198	"""
1199	kernel_block_size = self.kernel_block_size
1200	page_size = self._segmented_page_size
1201	assert kernel_block_size is not None
1202	assert page_size is not None
1203	pages_per_block = kernel_block_size // page_size
1204	max_local_pages = row_block_table.shape[1]
1205	max_local_blocks = cdiv(max_local_pages, pages_per_block)
1206	assert max_local_blocks <= block_table.shape[1], (
1207	f"DCP verify needs {max_local_blocks} blocks per request but the "
1208	f"block table only has {block_table.shape[1]}"
1209	)
1210	subpage_offsets = torch.arange(
1211	pages_per_block,
1212	dtype=block_table.dtype,
1213	device=block_table.device,
1214	)
1215	per_req_page_table = (
1216	block_table[:num_reqs, :max_local_blocks, None] * pages_per_block
1217	+ subpage_offsets
1218	).flatten(1)[:, :max_local_pages]
1219	# Broadcast the request's page list across its rows instead of
1220	# materializing a repeat_interleave copy.
1221	row_block_table.unflatten(0, (num_reqs, qlen)).copy_(
1222	per_req_page_table.unsqueeze(1)
1223	)
1224
1225	def _build_decode(
1226	self,
1227	block_table_tensor: torch.Tensor,
1228	seq_lens_device: torch.Tensor,
1229	max_seq_len: int,
1230	query_start_loc_cpu: torch.Tensor,
1231	query_start_loc_device: torch.Tensor,
1232	num_decode_tokens: int,
1233	max_query_len: int,
1234	dcp_tot_seq_lens_device: torch.Tensor \| None,
1235	) -> AiterMLADecodeMetadata:
1236	causal = self._decode_causal
1237	device = self.device
1238	num_reqs = seq_lens_device.size(0)
1239	qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
1240	max_qo_len = qo_len.max().item()
1241	padded_mtp_qo_len = self._uniform_padded_mtp_qo_len(
1242	qo_len, max_qo_len, num_decode_tokens
1243	)
1244	if padded_mtp_qo_len > 0:
1245	max_qo_len = padded_mtp_qo_len
1246	pad_uniform_mtp = padded_mtp_qo_len > 0
1247
1248	seq_lens_for_kernel = seq_lens_device
1249	# the global lengths need the same dummy-row
1250	# treatment as the local ones below, but only for g_kv_indptr -- the
1251	# metadata field keeps the true values, which is what the segmented
1252	# route wants (it derives its own row lens and gives a dummy row
1253	# length 0, correctly).
1254	g_tot_seq_lens = dcp_tot_seq_lens_device
1255	num_kernel_reqs = num_reqs
1256	if pad_uniform_mtp:
1257	qo_lens_device = (
1258	query_start_loc_device[1 : num_reqs + 1]
1259	- query_start_loc_device[:num_reqs]
1260	).to(torch.int32)
1261	seq_lens_for_kernel = torch.where(
1262	qo_lens_device > 0,
1263	seq_lens_for_kernel,
1264	seq_lens_for_kernel.new_full((), max_qo_len),
1265	)
1266	# A dummy row just got a LOCAL length of max_qo_len, so its GLOBAL
1267	# length has to say the same thing or the cprr mask is degenerate:
1268	# global_len 0 puts the row query positions at {-3,-2,-1} while its
1269	# local shard sits at {rank, W+rank, 2W+rank}, so every entry is
1270	# masked and the row softmaxes over nothing. max_qo_len * W is
1271	# exactly the global length whose round-robin shard is max_qo_len
1272	# on every rank.
1273	if g_tot_seq_lens is not None:
1274	g_tot_seq_lens = torch.where(
1275	qo_lens_device > 0,
1276	g_tot_seq_lens,
1277	g_tot_seq_lens.new_full((), max_qo_len * self.dcp_world_size),
1278	)
1279
1280	if self._graph_seq_lens is not None:
1281	self._graph_seq_lens[:num_kernel_reqs].copy_(
1282	seq_lens_for_kernel, non_blocking=True
1283	)
1284	seq_lens_for_kernel = self._graph_seq_lens[:num_kernel_reqs]
1285
1286	# The aiter kernel always operates with page_size=1 (the wrapper
1287	# flattens kv_buffer). last_page_len is always 1.
1288	paged_kv_last_page_len = self.paged_kv_last_page_len[:num_kernel_reqs]
1289
1290	# indptr: cumsum of seq_lens (one page per token in the flat view)
1291	paged_kv_indptr = torch.cat(
1292	[
1293	torch.zeros(1, dtype=torch.int32, device=device),
1294	seq_lens_for_kernel.cumsum(dim=0, dtype=torch.int32),
1295	]
1296	)
1297	use_gluon_decode = AiterMLAHelper.use_gluon_decode(
1298	self._decode_num_heads,
1299	int(max_qo_len),
1300	self._kv_cache_dtype_str,
1301	self._kv_cache_bytes,
1302	)
1303	use_gluon_verify = AiterMLAHelper.use_gluon_verify(
1304	self._decode_num_heads,
1305	int(max_qo_len),
1306	self._kv_cache_dtype_str,
1307	self.dcp_world_size,
1308	causal,
1309	self._kv_cache_bytes,
1310	)
1311	dcp_route = _DCPDecodeRoute.PLAIN
1312	if self.dcp_world_size > 1:
1313	dcp_route = _select_dcp_decode_route(
1314	self._supports_segmented_dcp_verify,
1315	causal,
1316	int(max_qo_len),
1317	self._asm_dcp_verify,
1318	)
1319
1320	# Segmented DCP verify carries its own per-row subpage table, so the
1321	# flat per-token view is dead work for it. Leave the buffer alone and
1322	# hand the metadata None, so a future reader cannot pick up whatever
1323	# the previous batch left behind.
1324	paged_kv_indices = None
1325	if dcp_route is not _DCPDecodeRoute.SEGMENTED:
1326	if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
1327	self.paged_kv_indices.fill_(-1)
1328
1329	# Expand block_table entries into per-token flat indices.
1330	# When kernel_block_size=1, this degrades to a direct copy (identical
1331	# to the original _copy_page_indices_kernel).
1332	# When kernel_block_size=K>1, block_table entry b covering K tokens
1333	# gets expanded to flat indices bK, bK+1, ..., b*K+(K-1).
1334	# Chunk count comes from the block table width, an upper bound on
1335	# tokens per request that is available host-side, so this adds no
1336	# device synchronisation. Programs whose chunk lies past
1337	# num_tokens mask out entirely.
1338	max_tokens_per_req = block_table_tensor.shape[1] * self.kernel_block_size
1339	num_chunks = max(1, cdiv(max_tokens_per_req, 1024))
1340	_expand_page_indices_kernel[(num_reqs, num_chunks)](
1341	self.paged_kv_indices,
1342	block_table_tensor,
1343	block_table_tensor.stride(0),
1344	block_table_tensor.stride(1),
1345	paged_kv_indptr,
1346	KERNEL_BLOCK_SIZE=self.kernel_block_size,
1347	BLOCK_SIZE=1024,
1348	)
1349	paged_kv_indices = self.paged_kv_indices
1350
1351	if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
1352	self.paged_kv_indptr[: 1 + num_kernel_reqs].copy_(
1353	paged_kv_indptr, non_blocking=True
1354	)
1355	self.paged_kv_indptr[1 + num_kernel_reqs :].fill_(paged_kv_indptr[-1])
1356	paged_kv_indptr = self.paged_kv_indptr[: 1 + num_kernel_reqs]
1357
1358	# paged_kv_last_page_len already uses the pre-initialized buffer slice
1359	# (set above), so no copy needed - buffer is always 1s.
1360
1361	if pad_uniform_mtp:
1362	qo_indptr_src = AiterMLAHelper.qo_indptr_for_uniform_qlen(
1363	num_kernel_reqs, int(max_qo_len), device
1364	)
1365	else:
1366	qo_indptr_src = query_start_loc_device[: 1 + num_kernel_reqs]
1367	self.qo_indptr[: 1 + num_kernel_reqs].copy_(
1368	qo_indptr_src, non_blocking=True
1369	)
1370	self.qo_indptr[1 + num_kernel_reqs :] = qo_indptr_src[-1]
1371	qo_indptr = self.qo_indptr[: 1 + num_kernel_reqs]
1372
1373	else:
1374	if max_qo_len == 1:
1375	qo_indptr = AiterMLAHelper.qo_indptr_for_uniform_qlen(
1376	num_kernel_reqs, 1, device
1377	)
1378	else:
1379	if pad_uniform_mtp:
1380	qo_indptr = AiterMLAHelper.qo_indptr_for_uniform_qlen(
1381	num_kernel_reqs, int(max_qo_len), device
1382	)
1383	else:
1384	qo_indptr = query_start_loc_device[: 1 + num_kernel_reqs]
1385
1386	g_kv_indptr = None
1387	if dcp_route is _DCPDecodeRoute.CPRR and g_tot_seq_lens is not None:
1388	assert self._g_kv_indptr_buf is not None
1389	_ngk = g_tot_seq_lens.shape[0]
1390	g_kv_indptr = self._g_kv_indptr_buf[: _ngk + 1]
1391	g_kv_indptr[1:].copy_(g_tot_seq_lens.cumsum(dim=0, dtype=torch.int32))
1392	# A FULL cudagraph replays at its capture-time request count, so a
1393	# smaller real batch leaves the tail holding the PREVIOUS step's
1394	# cumsum. Flatten it the way paged_kv_indptr is flattened below, so
1395	# padding rows read global_len == 0 and agree with their zero-length
1396	# local shard instead of a stale global length.
1397	self._g_kv_indptr_buf[_ngk + 1 :].fill_(g_kv_indptr[-1])
1398
1399	has_persistent_metadata = False
1400	# Only the asm decode consumes the schedule, so gate on the routing
1401	# rather than on num_heads >= 16, which denies it to a padded rank
1402	# running the same asm kernels. The predicates are disjoint -- decode
1403	# is qlen==1, verify is qlen>1 -- and cover both Gluon entries plus the
1404	# segmented DCP verify.
1405	use_persistent_metadata = (
1406	not use_gluon_decode
1407	and not use_gluon_verify
1408	and dcp_route is not _DCPDecodeRoute.SEGMENTED
1409	# A padded rank has no bf16 persistent kernel past qlen 4 where the
1410	# gfx950 fold is absent; the non-persistent entry covers it. fp8
1411	# keeps the schedule -- its fold rejects non-persistent outright.
1412	# A non-causal block keeps it too: what the fold drops past qlen 4
1413	# is the block's causal staircase, which a non-causal block does
1414	# not have, and the schedule is the only thing carrying its mask.
1415	and (
1416	not causal
1417	or self._decode_num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
1418	or max_qo_len <= AiterMLAHelper._ASM_PADDED_MAX_PS_QLEN
1419	or is_quantized_kv_cache(self._kv_cache_dtype_str)
1420	)
1421	and max_qo_len >= 1
1422	and max_qo_len <= self._mtp_decode_qlen
1423	)
1424	if (
1425	not causal
1426	and max_qo_len == 2
1427	and is_quantized_kv_cache(self._kv_cache_dtype_str)
1428	and not AiterMLAHelper.has_fp8_non_causal_qlen2_kernel(
1429	self._decode_num_heads
1430	)
1431	):
1432	# AITER's fp8 dispatch folds (gqa 16, qlen 3 or 4) onto the
1433	# qseqlen-4 kernel but never lists 2. Only 32/64/96/128 heads
1434	# at qlen 2 fold onto that 16-head / 4-token non-causal kernel;
1435	# 16-head qlen 2, and padded 48/80/112, keep Q2 and abort the
1436	# process rather than raising. The bf16 fold has no such hole.
1437	raise ValueError(
1438	"AITER has no non-causal fp8 MLA kernel for this 2-token "
1439	"query block. Pin the draft to TRITON_MLA for this "
1440	"speculative config."
1441	)
1442	if use_persistent_metadata:
1443	from aiter import get_mla_metadata_v1
1444
1445	uni_qo_len = (
1446	max_qo_len if pad_uniform_mtp or torch.all(qo_len == max_qo_len) else -1
1447	)
1448	cprr_kwargs: dict = {}
1449	is_causal = causal
1450	if self._asm_dcp_verify:
1451	cprr_kwargs = dict(
1452	max_split_per_batch=self._mla_max_split_per_batch,
1453	intra_batch_mode=False,
1454	)
1455	if dcp_route is _DCPDecodeRoute.CPRR:
1456	is_causal = False
1457	cprr_kwargs["is_cp_round_robin"] = True
1458	metadata_num_heads = (
1459	self._asm_dcp_verify_heads
1460	if dcp_route is _DCPDecodeRoute.CPRR
1461	else AiterMLAHelper.get_actual_mla_num_heads(self._decode_num_heads)
1462	)
1463	get_mla_metadata_v1(
1464	qo_indptr,
1465	paged_kv_indptr,
1466	paged_kv_last_page_len,
1467	metadata_num_heads,
1468	1,
1469	is_causal,
1470	self._mla_work_meta_data,
1471	self._mla_work_info_set,
1472	self._mla_work_indptr,
1473	self._mla_reduce_indptr,
1474	self._mla_reduce_final_map,
1475	self._mla_reduce_partial_map,
1476	page_size=1,
1477	kv_granularity=16,
1478	max_seqlen_qo=max_qo_len,
1479	uni_seqlen_qo=uni_qo_len,
1480	fast_mode=True,
1481	dtype_q=self._mla_q_dtype,
1482	dtype_kv=self._mla_kv_dtype,
1483	**cprr_kwargs,
1484	)
1485	has_persistent_metadata = True
1486
1487	# Small-head multi-token verify uses mla_gluon's 4-D MTP entry over the
1488	# ordinary per-request paged-KV view, so there is no expanded per-token
1489	# buffer to build here. mla_gluon still wants a lower bound on the KV
1490	# length it is asked to split; on the verify path that bound is over
1491	# active requests, not cudagraph padding rows pinned to max_qo_len.
1492	# The .item() below runs in the builder, outside the captured region,
1493	# so it does not abort HIP graph capture the way the per-layer syncs
1494	# in forward_mqa did. DCP verify never reaches here -- it carries its
1495	# own per-row lengths instead of a single split bound.
1496	min_kv_seq_len = 1
1497	if use_gluon_verify:
1498	per_req_len = paged_kv_indptr[1:] - paged_kv_indptr[:-1]
1499	if pad_uniform_mtp:
1500	active = qo_lens_device > 0
1501	if active.any():
1502	min_kv_seq_len = int(per_req_len[active].min().item())
1503	else:
1504	min_kv_seq_len = int(per_req_len.min().item())
1505
1506	dcp_verify = None
1507	if dcp_route is _DCPDecodeRoute.SEGMENTED:
1508	assert dcp_tot_seq_lens_device is not None
1509	dcp_verify = self._build_dcp_verify_row_view(
1510	int(max_qo_len),
1511	block_table_tensor,
1512	dcp_tot_seq_lens_device,
1513	)
1514
1515	attn_metadata = AiterMLADecodeMetadata(
1516	block_table=block_table_tensor,
1517	seq_lens=seq_lens_for_kernel,
1518	paged_kv_indptr=paged_kv_indptr,
1519	paged_kv_indices=paged_kv_indices,
1520	paged_kv_last_page_len=paged_kv_last_page_len,
1521	qo_indptr=qo_indptr,
1522	dcp_tot_seq_lens=dcp_tot_seq_lens_device,
1523	max_qo_len=max_qo_len,
1524	min_kv_seq_len=min_kv_seq_len,
1525	dcp_route=dcp_route,
1526	dcp_verify=dcp_verify,
1527	g_kv_indptr=g_kv_indptr,
1528	cp_world_size=self.dcp_world_size,
1529	cp_rank=self.dcp_rank,
1530	asm_decode_num_heads=(
1531	self._asm_dcp_verify_heads if dcp_route is _DCPDecodeRoute.CPRR else 0
1532	),
1533	mla_num_kv_splits=(
1534	self._mla_max_split_per_batch
1535	if dcp_route is _DCPDecodeRoute.CPRR
1536	else 0
1537	),
1538	use_gluon_decode=use_gluon_decode,
1539	use_gluon_verify=use_gluon_verify,
1540	attn_out_dtype=self.decode_attn_out_dtype,
1541	has_persistent_metadata=has_persistent_metadata,
1542	)
1543
1544	return attn_metadata
1545
1546	def build(
1547	self,
1548	common_prefix_len: int,
1549	common_attn_metadata: CommonAttentionMetadata,
1550	fast_build: bool = False,
1551	) -> AiterMLAMetadata:
1552	# The common MLA builder owns the causal flag but its _build_decode
1553	# contract is shared with CUDA backends. Keep this ROCm-specific state
1554	# on the per-ubatch Aiter builder instead of changing that shared API;
1555	# super().build invokes _build_decode synchronously.
1556	self._decode_causal = common_attn_metadata.causal
1557	attn_metadata = super().build(
1558	common_prefix_len, common_attn_metadata, fast_build
1559	)
1560	if (
1561	attn_metadata.decode is not None
1562	and attn_metadata.decode.has_persistent_metadata
1563	):
1564	attn_metadata.work_meta_data = self._mla_work_meta_data
1565	attn_metadata.work_indptr = self._mla_work_indptr
1566	attn_metadata.work_info_set = self._mla_work_info_set
1567	attn_metadata.reduce_indptr = self._mla_reduce_indptr
1568	attn_metadata.reduce_final_map = self._mla_reduce_final_map
1569	attn_metadata.reduce_partial_map = self._mla_reduce_partial_map
1570	if (
1571	self._fp8_prefill_enabled
1572	and attn_metadata.prefill is not None
1573	and attn_metadata.prefill.chunked_context is None
1574	):
1575	self._build_fp8_prefill_ps_metadata(attn_metadata, common_attn_metadata)
1576	return attn_metadata
1577
1578
1579	@triton.jit
1580	def _expand_page_indices_kernel(
1581	page_indices,
1582	block_table,
1583	block_table_stride_0,
1584	block_table_stride_1,
1585	cu_num_tokens,
1586	KERNEL_BLOCK_SIZE: tl.constexpr,
1587	BLOCK_SIZE: tl.constexpr,
1588	):
1589	"""Expand block table entries into per-token flat page indices.
1590
1591	The aiter MLA kernel always operates with page_size=1 internally
1592	(kv_buffer is flattened via .view(-1, 1, 1, H)). This kernel converts
1593	block-level indices from the block table into individual token positions
1594	in the flattened KV buffer.
1595
1596	When KERNEL_BLOCK_SIZE=1: block_idx=t, offset=0, flat=block_id
1597	(equivalent to a direct copy -- no regression from the original kernel).
1598
1599	When KERNEL_BLOCK_SIZE=K: block table entry b (covering K tokens)
1600	is expanded to flat indices bK, bK+1, ..., b*K+(K-1).
1601	"""
1602	# One program per (request, token-chunk). Parallelising over requests
1603	# alone degenerates at low concurrency: with num_reqs == 1 a single
1604	# workgroup walked the whole sequence in a serial loop.
1605	req_idx = tl.program_id(0)
1606	chunk_idx = tl.program_id(1)
1607	row_ptr = block_table + req_idx * block_table_stride_0
1608	start_idx = tl.load(cu_num_tokens + req_idx)
1609	num_tokens = tl.load(cu_num_tokens + req_idx + 1) - start_idx
1610
1611	# The grid is sized from the block table width, an upper bound over all
1612	# requests, so a ragged batch launches chunks past a short request's end.
1613	# Returning here keeps those programs from issuing masked-out loads and
1614	# stores at all.
1615	chunk_start = chunk_idx * BLOCK_SIZE
1616	if chunk_start >= num_tokens:
1617	return
1618
1619	token_offsets = chunk_start + tl.arange(0, BLOCK_SIZE)
1620	mask = token_offsets < num_tokens
1621
1622	# Which block in the block table does this token belong to?
1623	block_idx = token_offsets // KERNEL_BLOCK_SIZE
1624	# Offset within that block
1625	offset_in_block = token_offsets % KERNEL_BLOCK_SIZE
1626
1627	# Load the block ID from the block table
1628	# Both strides are taken from the caller: the block table is a view owned
1629	# elsewhere, so a unit column stride must not be assumed.
1630	block_ids = tl.load(row_ptr + block_idx * block_table_stride_1, mask=mask)
1631
1632	# Compute flat index in the flattened kv_buffer
1633	flat_indices = block_ids * KERNEL_BLOCK_SIZE + offset_in_block
1634
1635	tl.store(
1636	page_indices + start_idx + token_offsets,
1637	flat_indices,
1638	mask=mask,
1639	)
1640
1641
1642	class AiterMLAHelper:
1643	"""AITER MLA persistent (asm) decode requires a multiple of 16 heads. Unaligned
1644	head counts through 128 are padded to the next multiple of 16 by tiling the
1645	query heads and slicing to the padded size. Native H24 AITER builds bypass
1646	that padding. Small divisors of 16 retain the existing repeat_interleave and
1647	strided-unpad behavior. Native and aligned counts pass through without
1648	copies.
1649	"""
1650
1651	_AITER_MIN_MLA_HEADS: Final = 16
1652	_AITER_MAX_PADDED_MLA_HEADS: Final = 128
1653	# Largest qlen the padded gqa=16 asm decode has a bf16 persistent kernel
1654	# for. Above it only the non-persistent qseqlen=8 entry exists, and the
1655	# fold that reaches a persistent one is gfx950-only.
1656	_ASM_PADDED_MAX_PS_QLEN: Final = 4
1657	_AITER_UNSUPPORTED_HEADS: ClassVar[tuple[int, ...]] = ()
1658	# Pinned AITER v0.1.21.post2 folds these fp8 qlen-2 counts onto the
1659	# non-causal 16-head / 4-token kernel. Other multiples of 16 (48/80/112)
1660	# fold to H16 while keeping Q2, which has no non-causal fp8 entry.
1661	_AITER_FP8_NON_CAUSAL_QLEN2_HEADS: ClassVar[tuple[int, ...]] = (32, 64, 96, 128)
1662
1663	@staticmethod
1664	def qo_indptr_for_uniform_qlen(
1665	num_reqs: int,
1666	qlen: int,
1667	device: torch.device,
1668	dtype: torch.dtype = torch.int32,
1669	) -> torch.Tensor:
1670	"""Build ``[0, qlen, 2qlen, ..., num_reqsqlen]``."""
1671	return torch.arange(
1672	0,
1673	(num_reqs + 1) * qlen,
1674	step=qlen,
1675	dtype=dtype,
1676	device=device,
1677	)
1678
1679	@staticmethod
1680	def check_num_heads_validity(num_heads: int):
1681	assert AiterMLAHelper.is_valid_num_heads(num_heads), (
1682	"ROCM AITER MLA requires a positive multiple of 16 heads, or an "
1683	"unaligned head count up to 128 (padded to the next multiple of "
1684	f"16), but got {num_heads}.\n"
1685	f"Try adjusting tensor_parallel_size value."
1686	)
1687
1688	@staticmethod
1689	def is_valid_num_heads(num_heads: int) -> bool:
1690	return (
1691	num_heads > 0
1692	and num_heads not in AiterMLAHelper._AITER_UNSUPPORTED_HEADS
1693	and (
1694	num_heads <= AiterMLAHelper._AITER_MAX_PADDED_MLA_HEADS
1695	or num_heads % AiterMLAHelper._AITER_MIN_MLA_HEADS == 0
1696	)
1697	)
1698
1699	@staticmethod
1700	def has_fp8_non_causal_qlen2_kernel(num_heads: int) -> bool:
1701	"""Whether fp8 (num_heads, qlen=2) folds onto a non-causal ASM kernel.
1702
1703	Pinned AITER v0.1.21.post2 folds 32/64/96/128 heads at qlen 2 onto the
1704	non-causal 16-head / 4-token kernel. A 16-head (or padded-to-16) qlen-2
1705	block has no matching entry, and padded 48/80/112 keep Q2 after the
1706	H16 fold.
1707	"""
1708	kernel_heads = AiterMLAHelper.get_actual_mla_num_heads(num_heads)
1709	return kernel_heads in AiterMLAHelper._AITER_FP8_NON_CAUSAL_QLEN2_HEADS
1710
1711	@staticmethod
1712	def get_actual_mla_num_heads(num_heads: int) -> int:
1713	if num_heads == 24 and _aiter_mla_native_h24_supported():
1714	return num_heads
1715	m = AiterMLAHelper._AITER_MIN_MLA_HEADS
1716	return -(-num_heads // m) * m
1717
1718	@staticmethod
1719	def get_fp8_prefill_num_heads(num_heads: int) -> int:
1720	"""Head count the FP8 PS asm prefill runs at: the next multiple of 16.
1721
1722	Deliberately not ``get_actual_mla_num_heads``. That one carves out
1723	native H24 when ``_aiter_mla_native_h24_supported()``, which probes the
1724	asm decode reducer and metadata. The prefill is a different kernel
1725	pair (``mla_prefill_ps_asm_fwd`` + ``mla_reduce_v1``) with no such
1726	probe, so 24 heads pad to 32 here even on a native-H24 build.
1727
1728	The PS metadata in ``_init_fp8_prefill_ps_buffers``/``build()`` must be
1729	sized with this same function, or the work/reduce maps describe a
1730	different head count than the kernel is handed.
1731
1732	This function itself has no upper bound; the ceiling comes from the
1733	``is_valid_num_heads`` gate, which rejects counts above
1734	``_AITER_MAX_PADDED_MLA_HEADS`` (128) unless they are already
1735	16-aligned. That bound was established for the asm decode padding, so
1736	an architecture with, say, 136 heads per rank would be refused the
1737	prefill here for a reason that was never measured against this kernel
1738	pair. Revisit the constant rather than special-casing prefill.
1739	"""
1740	m = AiterMLAHelper._AITER_MIN_MLA_HEADS
1741	return -(-num_heads // m) * m
1742
1743	@staticmethod
1744	def get_mla_padded_q(
1745	num_heads: int, q: torch.Tensor, target_heads: int \| None = None
1746	) -> torch.Tensor:
1747	m = (
1748	target_heads
1749	if target_heads is not None
1750	else AiterMLAHelper.get_actual_mla_num_heads(num_heads)
1751	)
1752	if num_heads == m:
1753	return q
1754	if m % num_heads == 0:
1755	return q.repeat_interleave(m // num_heads, dim=1)
1756	# Non-divisor head counts cannot be padded by repeat_interleave. Tile
1757	# the query heads and slice to exactly m. MLA attention is independent
1758	# per query head over the shared KV, so padding heads cannot affect
1759	# heads [0:num_heads]; they are sliced back off the output.
1760	reps = -(-m // num_heads) # ceil(m / num_heads)
1761	# Slicing a tiled tensor yields a non-contiguous view. The asm decode
1762	# reads q as packed [tokens, m, head_dim], so materialize it.
1763	return q.repeat(1, reps, 1)[:, :m, :].contiguous()
1764
1765	@staticmethod
1766	def get_mla_unpadded_o(num_heads: int, o: torch.Tensor) -> torch.Tensor:
1767	return AiterMLAHelper._get_mla_unpadded_heads(num_heads, o)
1768
1769	@staticmethod
1770	def _get_mla_unpadded_heads(num_heads: int, tensor: torch.Tensor) -> torch.Tensor:
1771	m = AiterMLAHelper.get_actual_mla_num_heads(num_heads)
1772	if num_heads == m:
1773	return tensor
1774	if m % num_heads == 0:
1775	return tensor[:, :: m // num_heads, ...]
1776	# Undo the tile-padding from get_mla_padded_q: the real heads are the
1777	# first num_heads.
1778	return tensor[:, :num_heads, ...]
1779
1780	@staticmethod
1781	def get_mla_unpadded_lse(num_heads: int, lse: torch.Tensor) -> torch.Tensor:
1782	return AiterMLAHelper._get_mla_unpadded_heads(num_heads, lse)
1783
1784	@staticmethod
1785	def use_gluon_decode(
1786	num_heads: int,
1787	max_qo_len: int,
1788	kv_cache_dtype: str,
1789	kv_cache_bytes: int \| None = None,
1790	) -> bool:
1791	# Small-head (<16) single-token decode takes either the Gluon kernel or
1792	# the padded asm persistent decode, selected by
1793	# VLLM_ROCM_AITER_MLA_ASM_PADDING and the arch (Gluon is gfx950 only).
1794	m = AiterMLAHelper._AITER_MIN_MLA_HEADS
1795	if num_heads >= m or max_qo_len != 1:
1796	return False
1797	# Gluon's only fp8-KV regime, bh16bn128, is a bf16-query kernel with a
1798	# hardcoded scale that asserts batch_size == 1, so it cannot serve a
1799	# decode batch. Checked before the mode knob: an explicit "gluon"
1800	# request under fp8 would assert immediately.
1801	if is_quantized_kv_cache(kv_cache_dtype):
1802	return False
1803	mode = _aiter_mla_small_head_mode()
1804	if mode == "asm":
1805	return False
1806	if not _gluon_mla_decode_supported():
1807	return False
1808	if mode != "gluon" and m % num_heads != 0:
1809	return False
1810	# Last, so the size warning only fires where Gluon would have run.
1811	return _gluon_kv_cache_in_bounds(kv_cache_bytes)
1812
1813	@staticmethod
1814	def use_gluon_verify(
1815	num_heads: int,
1816	max_qo_len: int,
1817	kv_cache_dtype: str,
1818	dcp_world_size: int = 1,
1819	causal: bool = True,
1820	kv_cache_bytes: int \| None = None,
1821	) -> bool:
1822	"""Whether a small-head multi-token verify uses native Gluon MTP.
1823
1824	bf16 has no gqa<16, qseqlen>1 asm kernel, so verify goes through
1825	``mla_gluon``'s 4-D MTP entry (``q`` shaped ``[batch, qlen, nhead, dim]``)
1826	with ``use_2d_view=False``. fp8 has one via the q-row fold and must not
1827	come here: the MTP path hands Gluon the batch size its fp8 regime
1828	asserts against. A predicate rather than inline in forward_mqa so the
1829	builder sees the same answer the impl acts on.
1830
1831	DCP verify is excluded: its per-row causal windows are served by the
1832	segmented path, which Gluon's MTP entry cannot express.
1833
1834	Gluon masks the block causally with no way to turn it off, so a
1835	non-causal block takes the padded asm decode whatever the head count
1836	-- padding is what gives it a kernel there.
1837	"""
1838	if not causal:
1839	return False
1840	if max_qo_len <= 1 or dcp_world_size > 1:
1841	return False
1842	if is_quantized_kv_cache(kv_cache_dtype):
1843	return False
1844	if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS:
1845	return False
1846	if not _gluon_mla_decode_supported():
1847	return False
1848	# Same arch, mode and cache-size gating as use_gluon_decode.
1849	if _aiter_mla_small_head_mode() == "asm":
1850	return False
1851	return _gluon_kv_cache_in_bounds(kv_cache_bytes)
1852
1853	@staticmethod
1854	def dcp_local_verify_row_lens(
1855	tot_seq_lens: torch.Tensor,
1856	qlen: int,
1857	dcp_world_size: int,
1858	dcp_rank: int,
1859	cp_interleave: int,
1860	) -> torch.Tensor:
1861	"""Local KV length of every verify row, in global-causal order.
1862
1863	Row ``t`` of a qlen-token verify attends global positions
1864	``[0, seq_len - qlen + t]``, its own token included, so its local length
1865	is the round-robin count evaluated at that bound. Truncating causally in
1866	global coordinates has to happen before the shard mapping: the tokens a
1867	row drops sit on ``qlen - 1 - t`` different ranks, while subtracting the
1868	offset from a request's local length takes one from every rank.
1869	"""
1870	offsets = torch.arange(
1871	1,
1872	qlen + 1,
1873	device=tot_seq_lens.device,
1874	dtype=tot_seq_lens.dtype,
1875	)
1876	visible = (tot_seq_lens.unsqueeze(1) - qlen + offsets).clamp_(min=0)
1877	return get_dcp_local_seq_lens(
1878	visible,
1879	dcp_world_size,
1880	dcp_rank,
1881	cp_interleave,
1882	).flatten()
1883
1884
1885	class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
1886	# DCP decode paths return natural-log softmax LSE for the cross-rank merge.
1887	can_return_lse_for_decode: bool = True
1888	supports_dcp: bool = True
1889	# Measured on gfx950: aiter.mla.mla_decode_fwd(return_lse=True) matches
1890	# logsumexp to fp32 exactly, and merge_mla_segments_triton converts AITER's
1891	# base-2 segment statistics with LOGE2. Stated rather than inherited because
1892	# the DCP combine silently corrupts the softmax denominator if it disagrees.
1893	lse_base_on_e: bool = True
1894
1895	@property
1896	def _decode_num_heads(self) -> int:
1897	"""Return the query-head count after DCP gathering."""
1898	return self.num_heads * self.dcp_world_size
1899
1900	def __init__(
1901	self,
1902	num_heads: int,
1903	head_size: int,
1904	scale: float,
1905	num_kv_heads: int,
1906	alibi_slopes: list[float] \| None,
1907	sliding_window: int \| None,
1908	kv_cache_dtype: str,
1909	logits_soft_cap: float \| None,
1910	attn_type: str,
1911	kv_sharing_target_layer_name: str \| None,
1912	# MLA Specific Arguments
1913	**mla_args,
1914	) -> None:
1915	super().__init__(
1916	num_heads,
1917	head_size,
1918	scale,
1919	num_kv_heads,
1920	alibi_slopes,
1921	sliding_window,
1922	kv_cache_dtype,
1923	logits_soft_cap,
1924	attn_type,
1925	kv_sharing_target_layer_name,
1926	**mla_args,
1927	)
1928	AiterMLAHelper.check_num_heads_validity(num_heads)
1929	AiterMLAHelper.check_num_heads_validity(self._decode_num_heads)
1930
1931	unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
1932	if any(unsupported_features):
1933	raise NotImplementedError(
1934	"Aiter MLA does not support one of the following: "
1935	"alibi_slopes, sliding_window, logits_soft_cap"
1936	)
1937
1938	from aiter import flash_attn_varlen_func
1939
1940	self.flash_attn_varlen_func = flash_attn_varlen_func
1941
1942	# FP8 MLA prefill kernel imports (lazy, only when enabled).
1943	# Auto-enabled on gfx950 when AITER ships the kernels. Only runs when the
1944	# KV cache is FP8. Head counts that are not a multiple of 16 are
1945	# replicate-padded up to one (see _mla_fp8_prefill_attn).
1946	from vllm.utils.torch_utils import is_quantized_kv_cache
1947
1948	self._fp8_prefill_enabled = _fp8_mla_prefill_supported() and (
1949	is_quantized_kv_cache(kv_cache_dtype)
1950	and AiterMLAHelper.is_valid_num_heads(self.num_heads)
1951	)
1952	if self._fp8_prefill_enabled:
1953	from aiter import mla_prefill_ps_asm_fwd, mla_reduce_v1
1954
1955	self._mla_prefill_ps_asm_fwd = mla_prefill_ps_asm_fwd
1956	self._mla_reduce_v1 = mla_reduce_v1
1957
1958	def _flash_attn_varlen_diff_headdims(
1959	self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs
1960	):
1961	output = self.flash_attn_varlen_func( # type: ignore[call-arg]
1962	q=q,
1963	k=k,
1964	v=v,
1965	softmax_scale=softmax_scale,
1966	return_lse=return_softmax_lse,
1967	**kwargs,
1968	)
1969
1970	return output
1971
1972	def _mla_fp8_prefill_attn(
1973	self,
1974	q: torch.Tensor,
1975	k: torch.Tensor,
1976	v: torch.Tensor,
1977	attn_metadata: AiterMLAMetadata,
1978	out: torch.Tensor,
1979	) -> None:
1980	"""Run FP8 MLA prefill via mla_prefill_ps_asm_fwd + mla_reduce_v1.
1981
1982	Q, K, V are already decompressed (post-kv_b_proj), so K and V have
1983	``num_heads`` heads (same as Q) and gqa_ratio=1. Writes the
1984	result in-place to ``out``, which is the [total_q, nhead * v_head_dim]
1985	output buffer supplied by ``forward_mha``; no extra allocation or
1986	copy is required.
1987	"""
1988	from vllm.platforms import current_platform
1989	from vllm.v1.worker.workspace import current_workspace_manager
1990
1991	fp8_dtype = current_platform.fp8_dtype()
1992	total_q = q.shape[0]
1993	# PS asm prefill + mla_reduce_v1 require 16-aligned heads, and the PS
1994	# metadata is built for get_fp8_prefill_num_heads(num_heads). For head
1995	# counts that are not a multiple of 16 (K3 = 12/rank at TP8)
1996	# replicate-pad q/k/v up to that count, then slice the output back to
1997	# the real head count.
1998	#
1999	# Counts above 16 that are not multiples of 16 (24, 40, ...) are
2000	# handled by the same code but are not reached by any current model
2001	# and TP that fits: 96 heads would need TP4, whose weights exceed a
2002	# 288 GiB GPU, and 128-head models land on 128/64/32/16/8. The path is
2003	# kept general for future architectures and its numerics are covered by
2004	# test_fp8_prefill_matches_reference[num_heads=24]. Note the cost is
2005	# (padded - real)/real extra FLOPs and q/k/v bytes, which is worst just
2006	# above a multiple of 16 (17 heads pad to 32); a future arch landing
2007	# there should measure against the flash_attn_varlen_func fallback
2008	# rather than assume the asm path wins.
2009	#
2010	# Exact, not approximate: after kv_b_proj gqa_ratio is 1, so q, k and v
2011	# all carry num_heads heads and attention is independent per head.
2012	# Padding all three identically makes padded head j a duplicate of real
2013	# head j % num_heads, so the real heads [0:num_heads] are bit-identical
2014	# to the unpadded result. Same argument as the decode path; only the
2015	# target width differs.
2016	_real_nhead = self.num_heads
2017	nhead = AiterMLAHelper.get_fp8_prefill_num_heads(_real_nhead)
2018	_pad = nhead != _real_nhead
2019	if _pad:
2020	q = AiterMLAHelper.get_mla_padded_q(_real_nhead, q, nhead)
2021	k = AiterMLAHelper.get_mla_padded_q(_real_nhead, k, nhead)
2022	v = AiterMLAHelper.get_mla_padded_q(_real_nhead, v, nhead)
2023	v_head_dim = self.v_head_dim
2024	tile_q = _FP8_PREFILL_TILE_Q
2025
2026	# The FP8 ASM kernel expects FP8 inputs; the q_scale/k_scale/v_scale
2027	# parameters select per-tensor dequant scales. Q/K/V arrive as
2028	# bf16 from kv_b_proj, so cast here (one_scale=1.0 disables scaling).
2029	if q.dtype != fp8_dtype:
2030	q = q.to(fp8_dtype)
2031	if k.dtype != fp8_dtype:
2032	k = k.to(fp8_dtype)
2033	if v.dtype != fp8_dtype:
2034	v = v.to(fp8_dtype)
2035
2036	one_scale = torch.ones((), dtype=torch.float32, device=q.device)
2037
2038	# num_partial_tiles is resolved during metadata build to avoid an
2039	# in-forward .item() sync that would prevent CUDA Graph capture.
2040	# forward_mha gates the FP8 path on fp8_prefill_qo_indptr being set,
2041	# and the builder always sets every fp8_prefill_* field together, so
2042	# num_partial_tiles is non-None here.
2043	num_partial_tiles = attn_metadata.fp8_prefill_num_partial_tiles
2044	assert num_partial_tiles is not None
2045
2046	# Per-call scratch is served from the workspace manager so allocator
2047	# churn in the prefill hot path is bounded after warmup, matching the
2048	# pattern in PR #41002. The builder reserves the maximum shape of every
2049	# tensor requested here before the workspace is locked.
2050	scratch: list[tuple[tuple[int, ...], torch.dtype]] = [
2051	((num_partial_tiles * tile_q, nhead, v_head_dim), torch.float32),
2052	((num_partial_tiles * tile_q, nhead), torch.float32),
2053	((total_q, nhead), torch.float32),
2054	]
2055	if _pad:
2056	# The ASM and reduce kernels write a [total_q, nhead, v_head_dim]
2057	# buffer. With unpadded heads that aliases the caller's
2058	# [total_q, nhead * v_head_dim] output, so write straight into it;
2059	# padded heads do not fit that storage and need their own buffer.
2060	scratch.append(((total_q, nhead, v_head_dim), out.dtype))
2061
2062	workspace = current_workspace_manager()
2063	logits, attn_lse, final_lse, pad_out = workspace.get_simultaneous(scratch)
2064	out_3d = pad_out[0] if _pad else out.view(total_q, nhead, v_head_dim)
2065
2066	# Phase 1: persistent-scheduling assembly prefill kernel.
2067	self._mla_prefill_ps_asm_fwd(
2068	q,
2069	k,
2070	v,
2071	attn_metadata.fp8_prefill_qo_indptr,
2072	attn_metadata.fp8_prefill_kv_indptr,
2073	attn_metadata.fp8_prefill_kv_indices,
2074	attn_metadata.fp8_prefill_work_indptr,
2075	attn_metadata.fp8_prefill_work_info_set,
2076	attn_metadata.fp8_prefill_max_q_len,
2077	self.scale,
2078	True, # is_causal
2079	logits,
2080	attn_lse,
2081	out_3d,
2082	one_scale,
2083	one_scale,
2084	one_scale,
2085	)
2086
2087	# Phase 2: reduction across KV splits.
2088	self._mla_reduce_v1(
2089	logits,
2090	attn_lse,
2091	attn_metadata.fp8_prefill_reduce_indptr,
2092	attn_metadata.fp8_prefill_reduce_final_map,
2093	attn_metadata.fp8_prefill_reduce_partial_map,
2094	tile_q,
2095	# num_kv_splits added by ROCm/aiter#3391; 0 selects the kernel
2096	# default max(cu_num, 0) == cu_num, matching pre-#3391 behavior.
2097	0,
2098	out_3d,
2099	final_lse,
2100	)
2101
2102	if _pad:
2103	out.view(total_q, _real_nhead, v_head_dim).copy_(out_3d[:, :_real_nhead, :])
2104
2105	def forward_mha(
2106	self,
2107	q: torch.Tensor,
2108	kv_c_normed: torch.Tensor,
2109	k_pe: torch.Tensor,
2110	kv_c_and_k_pe_cache: torch.Tensor,
2111	attn_metadata: MLACommonMetadata,
2112	k_scale: torch.Tensor,
2113	output: torch.Tensor,
2114	output_scale: torch.Tensor \| None = None,
2115	) -> None:
2116	"""Dispatch prefill to the FP8 ASM kernel when available.
2117
2118	Falls back to the parent (``flash_attn_varlen_func``) when FP8
2119	MLA prefill is disabled, PS metadata is missing, or chunked
2120	context requires two-pass merge.
2121
2122	The annotation uses the base ``MLACommonMetadata`` to honour LSP
2123	with ``MLACommonImpl.forward_mha``; the AITER builder always
2124	produces ``AiterMLAMetadata`` instances at runtime, so we narrow
2125	with ``isinstance`` before reading the AITER-specific FP8 fields.
2126	"""
2127	if (
2128	not self._fp8_prefill_enabled
2129	or not isinstance(attn_metadata, AiterMLAMetadata)
2130	or attn_metadata.fp8_prefill_qo_indptr is None
2131	):
2132	return super().forward_mha(
2133	q,
2134	kv_c_normed,
2135	k_pe,
2136	kv_c_and_k_pe_cache,
2137	attn_metadata,
2138	k_scale,
2139	output,
2140	output_scale,
2141	)
2142
2143	assert attn_metadata.prefill is not None
2144	prefill_metadata = attn_metadata.prefill
2145	has_context = prefill_metadata.chunked_context is not None
2146
2147	if has_context:
2148	return super().forward_mha(
2149	q,
2150	kv_c_normed,
2151	k_pe,
2152	kv_c_and_k_pe_cache,
2153	attn_metadata,
2154	k_scale,
2155	output,
2156	output_scale,
2157	)
2158
2159	assert output_scale is None, (
2160	"fused FP8 output not supported by the AITER FP8 MLA prefill path"
2161	)
2162
2163	kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
2164	-1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
2165	)
2166	k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
2167	k = self._concat_k_nope_k_pe(k_nope, k_pe)
2168
2169	self._mla_fp8_prefill_attn(q, k, v, attn_metadata, output)
2170
2171	def _forward_segmented_dcp_verify(
2172	self,
2173	q_nope: torch.Tensor,
2174	q_pe: torch.Tensor,
2175	verify: AiterMLADCPVerifyMetadata,
2176	kv_c_and_k_pe_cache: torch.Tensor,
2177	layer: AttentionLayer,
2178	out_dtype: torch.dtype,
2179	) -> tuple[torch.Tensor, torch.Tensor]:
2180	"""Run segmented attention over this rank's shard of every verify row.
2181
2182	Each row's length already covers the current tokens this rank holds, so
2183	the cross-rank LSE merge in the MLA layer completes the causal block.
2184	"""
2185	q_mla = torch.cat([q_nope, q_pe], dim=-1)
2186	# skip_reduce with NUM_SEGMENTS>1 returns the partials and never writes
2187	# `out`. Pass None rather than aliasing q_mla: if that ever stops
2188	# holding, this fails loudly instead of scribbling over the query.
2189	segment_partials = _get_segmented_mla_decode()(
2190	q_mla,
2191	kv_c_and_k_pe_cache.view(
2192	-1,
2193	verify.page_size,
2194	1,
2195	kv_c_and_k_pe_cache.shape[-1],
2196	),
2197	None,
2198	verify.qo_indptr,
2199	verify.row_lens,
2200	verify.max_kv_seq_len,
2201	verify.block_table,
2202	self.scale,
2203	self.kv_lora_rank,
2204	self.qk_rope_head_dim,
2205	causal=True,
2206	q_descale=None,
2207	kv_descale=layer._k_scale,
2208	skip_reduce=True,
2209	)
2210	assert isinstance(segment_partials, tuple) and len(segment_partials) == 3, (
2211	"AITER segmented MLA verify must return segment partials "
2212	"when skip_reduce=True."
2213	)
2214	segm_output, segm_max, segm_expsum = segment_partials
2215	return merge_mla_segments_triton(
2216	segm_output,
2217	segm_max,
2218	segm_expsum,
2219	verify.row_lens,
2220	verify.page_size,
2221	out_dtype,
2222	)
2223
2224	def forward_mqa(
2225	self,
2226	q: torch.Tensor \| tuple[torch.Tensor, torch.Tensor],
2227	kv_c_and_k_pe_cache: torch.Tensor,
2228	attn_metadata: AiterMLAMetadata,
2229	layer: AttentionLayer,
2230	) -> tuple[torch.Tensor, torch.Tensor \| None]:
2231	assert kv_c_and_k_pe_cache.numel() > 0
2232	assert attn_metadata.decode is not None
2233
2234	decode = attn_metadata.decode
2235	assert decode.max_qo_len is not None
2236	if decode.use_gluon_decode:
2237	assert decode.paged_kv_indptr is not None
2238	assert decode.paged_kv_indices is not None
2239	if type(q) is tuple:
2240	q_nope, q_pe = q
2241	else:
2242	q_nope, q_pe = torch.split(
2243	q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
2244	)
2245	B, num_q_heads, _ = q_nope.shape
2246	o = torch.empty(
2247	B,
2248	num_q_heads,
2249	self.kv_lora_rank,
2250	dtype=decode.attn_out_dtype,
2251	device=q_nope.device,
2252	)
2253	kv_buffer = kv_c_and_k_pe_cache.reshape(-1, kv_c_and_k_pe_cache.shape[-1])
2254	mla_gluon = _get_mla_gluon()
2255	mla_gluon(
2256	q_nope=q_nope,
2257	q_pe=q_pe,
2258	kv_c=kv_buffer,
2259	o=o,
2260	page_table=decode.paged_kv_indices,
2261	seq_info=decode.paged_kv_indptr,
2262	sm_scale=self.scale,
2263	k_pe=None,
2264	kv_pe_offset=self.kv_lora_rank,
2265	use_2d_view=False,
2266	kv_scale=1.0,
2267	min_kv_seq_len=decode.min_kv_seq_len,
2268	)
2269	return o, None
2270
2271	# 12-head (<16) multi-token verify (DSpark): the asm path has no
2272	# gqa<16, qseqlen>1 kernel, so the block goes to the gluon kernel's 4-D
2273	# MTP entry, which serves a whole (1 + num_spec) block in one launch.
2274	# The block is causal -- the target is checking draft tokens, so
2275	# position t must not see t+1 -- and the kernel bounds each query
2276	# position's scores itself.
2277	# Arch, mode and dtype gating all live in use_gluon_verify, which the
2278	# builder evaluates and records on the metadata, so this branch acts on
2279	# the same answer the builder used to decide whether the asm decode
2280	# would run. DCP verify is routed to the segmented path below instead.
2281	if decode.use_gluon_verify:
2282	if type(q) is tuple:
2283	q_nope, q_pe = q
2284	else:
2285	q_nope, q_pe = torch.split(
2286	q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
2287	)
2288	B, num_q_heads, _ = q_nope.shape
2289	o = torch.empty(
2290	B,
2291	num_q_heads,
2292	self.kv_lora_rank,
2293	dtype=decode.attn_out_dtype,
2294	device=q_nope.device,
2295	)
2296	kv_buffer = kv_c_and_k_pe_cache.reshape(-1, kv_c_and_k_pe_cache.shape[-1])
2297	assert attn_metadata.causal, (
2298	"AITER MLA small-head verify MTP is causal-only"
2299	)
2300	# Hand mla_gluon its 4-D MTP entry instead of an expanded
2301	# per-verify-token paged-KV view. The flat query layout is already
2302	# row-major (r * qlen + t), so unflatten is a free view. mla_gluon
2303	# applies the per-position causal bound
2304	# score_end = min(split_kv_end, seq_len - qlen + q_pos + 1)
2305	# in-kernel, and seq_lens already spans the verify block, so that
2306	# bound is context_r + q_pos + 1 -- the same window the expanded
2307	# view supplied by truncating each row's page list.
2308	qlen = int(decode.max_qo_len)
2309	num_reqs = B // qlen
2310	if num_reqs * qlen != B:
2311	raise ValueError(
2312	f"verify block {B} rows is not a multiple of qlen {qlen}"
2313	)
2314	assert decode.paged_kv_indptr is not None
2315	assert decode.paged_kv_indices is not None
2316	mla_gluon = _get_mla_gluon()
2317	mla_gluon(
2318	q_nope=q_nope.unflatten(0, (num_reqs, qlen)),
2319	q_pe=q_pe.unflatten(0, (num_reqs, qlen)),
2320	kv_c=kv_buffer,
2321	o=o.unflatten(0, (num_reqs, qlen)),
2322	page_table=decode.paged_kv_indices,
2323	seq_info=decode.paged_kv_indptr,
2324	sm_scale=self.scale,
2325	k_pe=None,
2326	kv_pe_offset=self.kv_lora_rank,
2327	use_2d_view=False,
2328	kv_scale=1.0,
2329	min_kv_seq_len=decode.min_kv_seq_len,
2330	)
2331	return o, None
2332
2333	dcp_route = decode.dcp_route
2334	if dcp_route is _DCPDecodeRoute.SEGMENTED:
2335	verify = decode.dcp_verify
2336	assert verify is not None
2337	if type(q) is tuple:
2338	q_nope, q_pe = q
2339	else:
2340	q_nope, q_pe = torch.split(
2341	q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
2342	)
2343	if (
2344	is_quantized_kv_cache(self.kv_cache_dtype)
2345	and q_nope.dtype != torch.bfloat16
2346	):
2347	q_nope = q_nope.to(torch.bfloat16) * layer._q_scale
2348	q_pe = q_pe.to(torch.bfloat16) * layer._q_scale
2349	return self._forward_segmented_dcp_verify(
2350	q_nope,
2351	q_pe,
2352	verify,
2353	kv_c_and_k_pe_cache,
2354	layer,
2355	decode.attn_out_dtype,
2356	)
2357
2358	if type(q) is tuple:
2359	q = torch.cat(q, dim=-1)
2360
2361	assert isinstance(q, torch.Tensor)
2362	assert decode.paged_kv_indptr is not None
2363	assert decode.paged_kv_indices is not None
2364	B = q.shape[0]
2365
2366	assert q.shape[1] == self._decode_num_heads, (
2367	"ROCM_AITER_MLA decode expected the DCP-gathered query head count "
2368	f"{self._decode_num_heads}, got {q.shape[1]}"
2369	)
2370	# cprr has kernels only at native head counts, so pad q up (e.g. 96 ->
2371	# 128). Heads are independent in MLA; the extras are sliced off below.
2372	asm_dcp_heads = decode.asm_decode_num_heads
2373	if asm_dcp_heads > self._decode_num_heads:
2374	mla_num_heads = asm_dcp_heads
2375	mla_padded_q = AiterMLAHelper.get_mla_padded_q(
2376	self._decode_num_heads, q, target_heads=mla_num_heads
2377	)
2378	else:
2379	mla_padded_q = AiterMLAHelper.get_mla_padded_q(self._decode_num_heads, q)
2380	mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(
2381	self._decode_num_heads
2382	)
2383	o = torch.empty(
2384	B,
2385	mla_num_heads,
2386	self.kv_lora_rank,
2387	dtype=attn_metadata.decode.attn_out_dtype,
2388	device=q.device,
2389	)
2390	if decode.max_qo_len > 1 and not decode.has_persistent_metadata:
2391	# MTP verification can call the AITER MLA decode kernel with
2392	# qlen > 1. If that path is running without persistent metadata,
2393	# zero-fill so unwritten lanes cannot leak into logits.
2394	o.zero_()
2395
2396	kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2)
2397
2398	# Build kwargs for mla_decode_fwd. Pass persistent metadata only
2399	# when it was successfully computed.
2400	mla_kwargs = dict(
2401	q_scale=layer._q_scale,
2402	kv_scale=layer._k_scale,
2403	)
2404	if attn_metadata.work_meta_data is not None:
2405	mla_kwargs.update(
2406	work_meta_data=attn_metadata.work_meta_data,
2407	work_indptr=attn_metadata.work_indptr,
2408	work_info_set=attn_metadata.work_info_set,
2409	reduce_indptr=attn_metadata.reduce_indptr,
2410	reduce_final_map=attn_metadata.reduce_final_map,
2411	reduce_partial_map=attn_metadata.reduce_partial_map,
2412	)
2413
2414	lse = None
2415	if self.dcp_world_size > 1:
2416	if dcp_route is _DCPDecodeRoute.CPRR:
2417	# cprr only: global positions for the in-kernel causal window.
2418	assert decode.g_kv_indptr is not None
2419	mla_kwargs["g_kv_indptr"] = decode.g_kv_indptr
2420	mla_kwargs["cp_world_size"] = decode.cp_world_size
2421	mla_kwargs["cp_rank"] = decode.cp_rank
2422	# Must match the schedule get_mla_metadata_v1 was built with.
2423	assert decode.mla_num_kv_splits > 0
2424	mla_kwargs["num_kv_splits"] = decode.mla_num_kv_splits
2425	mla_kwargs["intra_batch_mode"] = False
2426	# The vLLM custom-op wrapper exposes only the in-place output and
2427	# drops aiter's final LSE, which the cross-shard merge needs, so go
2428	# through aiter's native entry point on the DCP path.
2429	_, lse = _get_aiter_mla_decode()(
2430	mla_padded_q,
2431	kv_buffer.view(-1, 1, 1, mla_padded_q.shape[-1]),
2432	o,
2433	decode.qo_indptr,
2434	decode.paged_kv_indptr,
2435	decode.paged_kv_indices,
2436	decode.paged_kv_last_page_len,
2437	decode.max_qo_len,
2438	sm_scale=self.scale,
2439	return_lse=True,
2440	causal=attn_metadata.causal,
2441	**mla_kwargs,
2442	)
2443	assert lse is not None, (
2444	"aiter mla_decode_fwd(return_lse=True) returned no LSE; upgrade "
2445	"aiter to a build with decode LSE support."
2446	)
2447	else:
2448	rocm_aiter_ops.mla_decode_fwd(
2449	mla_padded_q,
2450	kv_buffer,
2451	o,
2452	self.scale,
2453	decode.qo_indptr,
2454	decode.max_qo_len,
2455	decode.paged_kv_indptr,
2456	decode.paged_kv_indices,
2457	decode.paged_kv_last_page_len,
2458	causal=attn_metadata.causal,
2459	**mla_kwargs,
2460	)
2461
2462	if asm_dcp_heads > self._decode_num_heads:
2463	# Keep a view: a post-kernel copy races the symm-mem a2a combine.
2464	output = o[:, : self._decode_num_heads]
2465	if lse is not None:
2466	lse = lse[:, : self._decode_num_heads, ...]
2467	else:
2468	output = AiterMLAHelper.get_mla_unpadded_o(self._decode_num_heads, o)
2469	if lse is not None:
2470	lse = AiterMLAHelper.get_mla_unpadded_lse(self._decode_num_heads, lse)
2471	return output, lse