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

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 4,026 lines 143 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3	"""Unit tests for the sparse MLA backends and utilities."""
4
5	import math
6	from collections import deque
7	from types import MethodType, SimpleNamespace
8	from unittest.mock import MagicMock
9
10	import pytest
11	import torch
12
13	from tests.v1.attention.test_mla_backends import (
14	BATCH_SPECS,
15	BatchSpec,
16	MockSparseMLAAttentionLayer,
17	create_and_prepopulate_kv_cache,
18	)
19	from tests.v1.attention.utils import (
20	create_common_attn_metadata,
21	create_standard_kv_cache_spec,
22	create_vllm_config,
23	)
24	from vllm import _custom_ops as ops
25	from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
26	from vllm.config import (
27	CUDAGraphMode,
28	HiSparseConfig,
29	SpeculativeConfig,
30	set_current_vllm_config,
31	)
32	from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
33	HiSparseConnectorWorker,
34	)
35	from vllm.forward_context import set_forward_context
36	from vllm.model_executor.layers.attention import mla_attention
37	from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha
38	from vllm.model_executor.layers.attention.sparse_mla_attention import (
39	GLOBAL_TOPK_MASK_MAX_BYTES,
40	SparseMLACommonImpl,
41	SparseMLAPrefillMetadata,
42	_is_masked_mha_available,
43	_masked_mha_workspace_fits,
44	_topk_mask_shape,
45	_use_dense_mha_prefill,
46	)
47	from vllm.model_executor.layers.linear import ColumnParallelLinear
48	from vllm.platforms import current_platform
49
50	# TODO: Integrate ROCMAiterMLASparseBackend for ROCm.
51	# The ROCm sparse MLA backend (rocm_aiter_mla_sparse.py) has a compatible
52	# forward_mqa interface but needs validation on ROCm hardware.
53	if not current_platform.is_cuda():
54	pytest.skip(
55	"Sparse MLA backend tests currently only support CUDA. "
56	"ROCm support requires integrating ROCMAiterMLASparseBackend.",
57	allow_module_level=True,
58	)
59
60	from vllm.model_executor.layers.attention.mla_attention import (
61	_canonicalize_sparse_mla_kv_cache_dtype,
62	)
63	from vllm.model_executor.layers.attention.sparse_mla_attention import (
64	SharedTopkIndicesBuffer,
65	)
66	from vllm.utils.math_utils import cdiv
67	from vllm.utils.torch_utils import current_stream
68	from vllm.v1.attention.backends.mla import index_group as index_group_module
69	from vllm.v1.attention.backends.mla.flashattn_mla_sparse import (
70	FlashAttnMLASparseImpl,
71	FlashAttnMLASparseMetadataBuilder,
72	)
73	from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
74	FlashInferMLASparseImpl,
75	FlashInferMLASparseTRTLLMBackend,
76	)
77	from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import (
78	FlashInferMLASparseSM120Impl,
79	)
80	from vllm.v1.attention.backends.mla.flashmla_sparse import (
81	FlashMLASparseBackend,
82	FlashMLASparseImpl,
83	FlashMLASparseMetadata,
84	FlashMLASparseMetadataBuilder,
85	triton_convert_req_index_to_global_index,
86	)
87	from vllm.v1.attention.backends.mla.index_group import (
88	HiSparseMLAIndexGroup,
89	SparseMLAIndexGroupBuilder,
90	)
91	from vllm.v1.attention.backends.mla.indexer import (
92	DeepseekV32IndexerMetadataBuilder,
93	)
94	from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
95	from vllm.v1.attention.backends.utils import (
96	split_decodes_and_prefills,
97	split_prefill_chunks,
98	)
99	from vllm.v1.attention.ops import flashmla
100	from vllm.v1.hisparse import runtime as hisparse_runtime
101	from vllm.v1.hisparse.runtime import (
102	HiSparseCacheHandle,
103	HiSparseRuntime,
104	ResolvedHiSparseConfig,
105	_has_hisparse_ops,
106	build_hisparse_prefill_staging_plan,
107	hisparse_prefill_staging_remap,
108	)
109	from vllm.v1.hisparse.types import SparseKVRowMirror
110
111	SPARSE_BACKEND_BATCH_SPECS = {
112	name: BATCH_SPECS[name]
113	for name in [
114	"mixed_small",
115	"mixed_medium",
116	"small_prefill",
117	"medium_prefill",
118	"single_prefill",
119	]
120	}
121
122	SPARSE_BACKEND_BATCH_SPECS["large_q_prefill"] = BatchSpec(
123	seq_lens=[1024] * 2, query_lens=[256] * 2
124	)
125	SPARSE_BACKEND_BATCH_SPECS["large_q_pure_prefill"] = BatchSpec(
126	seq_lens=[256] * 2, query_lens=[256] * 2
127	)
128
129	DEVICE_TYPE = current_platform.device_type
130
131
132	def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch):
133	"""Weight absorption must not change the model's attention temperature."""
134	model_scale = 256**-0.5
135	kv_lora_rank = 512
136	topk = torch.zeros((1, 1), dtype=torch.int32)
137	metadata = SimpleNamespace(
138	req_id_per_token=torch.zeros(1, dtype=torch.int32),
139	block_table=torch.zeros((1, 1), dtype=torch.int32),
140	block_size=1,
141	)
142	recorded_scale = None
143
144	impl = object.__new__(FlashInferMLASparseImpl)
145	impl.scale = model_scale
146	impl.qk_nope_head_dim = 256
147	impl.kv_lora_rank = kv_lora_rank
148	impl.qk_rope_head_dim = 0
149	impl.kv_cache_dtype = "auto"
150	impl.topk_indices_buffer = topk
151	impl.dcp_world_size = 1
152	impl._workspace_buffer = torch.empty(1)
153	impl.bmm1_scale = None
154	impl.bmm2_scale = None
155	impl.is_nope_mla = True
156	impl.need_to_return_lse_for_decode = False
157	impl.index_group = MagicMock()
158	impl.index_group_index = 0
159	impl.index_group.convert_logical_to_physical_topk.return_value = (
160	topk,
161	torch.ones(1, dtype=torch.int32),
162	)
163
164	import flashinfer.decode
165
166	def fake_flashinfer(**kwargs):
167	nonlocal recorded_scale
168	recorded_scale = kwargs["bmm1_scale"]
169	return torch.zeros((1, 1, 1, kv_lora_rank))
170
171	monkeypatch.setattr(
172	flashinfer.decode,
173	"trtllm_batch_decode_with_kv_cache_mla",
174	fake_flashinfer,
175	)
176	impl.forward_mqa(
177	torch.zeros(1, 1, kv_lora_rank),
178	torch.zeros(1, 1, kv_lora_rank),
179	metadata,
180	SimpleNamespace(),
181	)
182
183	assert recorded_scale == model_scale
184	assert recorded_scale != kv_lora_rank**-0.5
185
186
187	def test_hisparse_routes_prefill_to_sparse_mqa():
188	layer = SimpleNamespace(hisparse_cache=object())
189
190	assert not mla_attention.MLAAttention._use_sparse_mha(layer, SimpleNamespace())
191
192
193	def test_hisparse_metadata_keeps_short_prefill_indexer_enabled():
194	config = SimpleNamespace(
195	attention_config=SimpleNamespace(
196	hisparse_config=object(), sparse_mla_force_mqa=False
197	)
198	)
199
200	assert not _use_dense_mha_prefill(config, prefill_max_seq_len=512, topk_tokens=2048)
201
202
203	def test_unified_mla_prepares_batch_without_slot_mapping(monkeypatch):
204	layer = SimpleNamespace(
205	prepare_kv_cache_update=MagicMock(),
206	update_kv_cache=MagicMock(),
207	)
208	metadata = object()
209	monkeypatch.setattr(
210	mla_attention,
211	"get_attention_context",
212	lambda _: (metadata, layer, torch.empty(0), None),
213	)
214	tensor = torch.empty(0, device=DEVICE_TYPE)
215
216	mla_attention.unified_mla_kv_cache_update(tensor, tensor, "layer", "auto", tensor)
217
218	layer.prepare_kv_cache_update.assert_called_once_with(metadata)
219	layer.update_kv_cache.assert_not_called()
220
221
222	def _float_to_e8m0_truncate(f: float) -> float:
223	"""Simulate SM100's float -> e8m0 -> bf16 scale conversion.
224	e8m0 format only stores the exponent (power of 2).
225	cudaRoundZero truncates toward zero, meaning we round down to the
226	nearest power of 2.
227	"""
228	if f <= 0:
229	return 0.0
230	# e8m0 = floor(log2(f)), then 2^(e8m0)
231	# This is equivalent to truncating to the nearest power of 2 below f
232	exp = math.floor(math.log2(f))
233	return 2.0**exp
234
235
236	def _dequantize_fp8_ds_mla_entry(
237	cache_slice: torch.Tensor,
238	kv_lora_rank: int,
239	rope_dim: int,
240	dtype: torch.dtype,
241	simulate_sm100_e8m0_scales: bool = False,
242	) -> tuple[torch.Tensor, torch.Tensor]:
243	"""Dequantize a single fp8_ds_mla cache entry back to latent + rope.
244
245	Args:
246	simulate_sm100_e8m0_scales: If True, simulate the SM100 kernel's
247	float -> e8m0 -> bf16 scale conversion path.
248
249	"""
250	# The first kv_lora_rank bytes store FP8 latent values with one scale per
251	# 128 element tile written as float32 right after the latent payload.
252	scales = cache_slice.view(torch.float32)[kv_lora_rank // 4 : kv_lora_rank // 4 + 4]
253	latent = torch.empty(kv_lora_rank, dtype=torch.float16, device=cache_slice.device)
254	for tile_idx in range(4):
255	tile_start = tile_idx * 128
256	tile_end = tile_start + 128
257	scale_val = float(scales[tile_idx].item())
258	if simulate_sm100_e8m0_scales:
259	# Simulate the lossy float -> e8m0 -> bf16 conversion
260	scale_val = _float_to_e8m0_truncate(scale_val)
261	ops.convert_fp8(
262	latent[tile_start:tile_end],
263	cache_slice[tile_start:tile_end],
264	scale_val,
265	kv_dtype="fp8",
266	)
267	latent = latent.to(dtype)
268
269	rope_offset = kv_lora_rank // 2 + 8
270	rope_vals = cache_slice.view(dtype)[rope_offset : rope_offset + rope_dim]
271	return latent, rope_vals.clone()
272
273
274	def _quantize_dequantize_fp8_ds_mla(
275	kv_c: torch.Tensor,
276	k_pe: torch.Tensor,
277	block_size: int,
278	scale: torch.Tensor,
279	simulate_sm100_e8m0_scales: bool = False,
280	) -> tuple[torch.Tensor, torch.Tensor]:
281	"""Round-trip kv_c/k_pe though the fp8_ds_mla cache layout.
282
283	Args:
284	simulate_sm100_e8m0_scales: If True, simulate the SM100 kernel's
285	float -> e8m0 -> bf16 scale conversion in dequantization.
286
287	"""
288	if kv_c.numel() == 0:
289	return kv_c.clone(), k_pe.clone()
290
291	kv_lora_rank = kv_c.shape[-1]
292	rope_dim = k_pe.shape[-1]
293	num_tokens = kv_c.shape[0]
294	num_blocks = max(1, math.ceil(num_tokens / block_size))
295	entry_size = kv_lora_rank + 4 * 4 + 2 * rope_dim
296
297	tmp_cache = torch.zeros(
298	num_blocks, block_size, entry_size, dtype=torch.uint8, device=kv_c.device
299	)
300	slot_mapping = torch.arange(num_tokens, dtype=torch.long, device=kv_c.device)
301
302	ops.concat_and_cache_mla(
303	kv_c, k_pe, tmp_cache, slot_mapping, kv_cache_dtype="fp8_ds_mla", scale=scale
304	)
305
306	dequant_kv_c = torch.empty_like(kv_c)
307	dequant_k_pe = torch.empty_like(k_pe)
308
309	for token_idx in range(num_tokens):
310	slot = slot_mapping[token_idx].item()
311	block_idx = slot // block_size
312	block_offset = slot % block_size
313	cache_slice = tmp_cache[block_idx, block_offset]
314	latent, rope_vals = _dequantize_fp8_ds_mla_entry(
315	cache_slice,
316	kv_lora_rank,
317	rope_dim,
318	kv_c.dtype,
319	simulate_sm100_e8m0_scales=simulate_sm100_e8m0_scales,
320	)
321	dequant_kv_c[token_idx] = latent
322	dequant_k_pe[token_idx] = rope_vals
323
324	return dequant_kv_c, dequant_k_pe
325
326
327	_E2M1_TABLE = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
328
329
330	def _quantize_dequantize_nvfp4_ds_mla(
331	kv_c: torch.Tensor,
332	k_pe: torch.Tensor,
333	block_size: int,
334	kv_cache_dtype: str,
335	scale: torch.Tensor,
336	) -> tuple[torch.Tensor, torch.Tensor]:
337	"""Round-trip kv_c/k_pe through the nvfp4_ds_mla cache layout.
338
339	Layout (per token, uint8): [e2m1-packed NoPE \| unscaled e4m3 RoPE \|
340	per-16 e4m3 NoPE SFs]. The SF bytes are stored permuted (an 8x4 -> 4x8
341	transpose: the scale for element block s lives at byte 8 * (s & 3) +
342	(s >> 2)) so that the 8 scales one FlashMLA dequant thread needs are
343	contiguous. Dequant: x = value * float(sf), which is exact in bf16, so
344	this python dequant matches the kernels bit-for-bit.
345	"""
346	if kv_c.numel() == 0:
347	return kv_c.clone(), k_pe.clone()
348
349	kv_lora_rank = kv_c.shape[-1]
350	rope_dim = k_pe.shape[-1]
351	num_tokens = kv_c.shape[0]
352	num_blocks = max(1, math.ceil(num_tokens / block_size))
353	entry_size = (
354	math.ceil((kv_lora_rank // 2 + rope_dim + kv_lora_rank // 16) / 16) * 16
355	)
356	sf_nope_off = kv_lora_rank // 2 + rope_dim
357
358	tmp_cache = torch.zeros(
359	num_blocks, block_size, entry_size, dtype=torch.uint8, device=kv_c.device
360	)
361	slot_mapping = torch.arange(num_tokens, dtype=torch.long, device=kv_c.device)
362	ops.concat_and_cache_mla(
363	kv_c, k_pe, tmp_cache, slot_mapping, kv_cache_dtype=kv_cache_dtype, scale=scale
364	)
365
366	tokens = tmp_cache.view(-1, entry_size)[:num_tokens]
367	table = torch.tensor(
368	_E2M1_TABLE + [-v for v in _E2M1_TABLE],
369	dtype=torch.float32,
370	device=kv_c.device,
371	)
372
373	def unpack_e2m1(packed: torch.Tensor) -> torch.Tensor:
374	lo = (packed & 0xF).long()
375	hi = (packed >> 4).long()
376	return table[torch.stack([lo, hi], dim=-1).flatten(-2)]
377
378	nope_vals = unpack_e2m1(tokens[:, : kv_lora_rank // 2])
379	num_nope_sf = kv_lora_rank // 16
380	nope_sf = (
381	# undo the on-wire SF permutation -> element-block order
382	tokens[:, sf_nope_off : sf_nope_off + num_nope_sf]
383	.unflatten(-1, (4, num_nope_sf // 4))
384	.transpose(-1, -2)
385	.flatten(-2)
386	.view(torch.float8_e4m3fn)
387	.float()
388	)
389	dequant_kv_c = (
390	(nope_vals.unflatten(-1, (-1, 16)) * nope_sf.unsqueeze(-1)).flatten(-2)
391	).to(kv_c.dtype)
392
393	# RoPE: plain e4m3, no scale factor
394	rope_raw = tokens[:, kv_lora_rank // 2 : sf_nope_off]
395	dequant_k_pe = rope_raw.view(torch.float8_e4m3fn).to(k_pe.dtype)
396
397	return dequant_kv_c, dequant_k_pe
398
399
400	@pytest.mark.parametrize(
401	"backend_cls",
402	[FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend],
403	ids=["FlashMLA", "FlashInferTRTLLM"],
404	)
405	@pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys()))
406	@pytest.mark.parametrize(
407	"kv_cache_dtype",
408	["auto", "fp8", "fp8_ds_mla", "nvfp4_ds_mla"],
409	)
410	@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
411	@pytest.mark.parametrize("block_size", [32, 64])
412	@pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)])
413	def test_sparse_backend_decode_correctness(
414	default_vllm_config,
415	dist_init,
416	backend_cls,
417	batch_name,
418	kv_cache_dtype,
419	tensor_parallel_size,
420	block_size,
421	workspace_init,
422	q_scale: float,
423	k_scale: float,
424	monkeypatch,
425	):
426	if (
427	batch_name == "large_q_pure_prefill"
428	and backend_cls == FlashMLASparseBackend
429	and kv_cache_dtype == "fp8_ds_mla"
430	and tensor_parallel_size == 4
431	):
432	monkeypatch.setattr(
433	"vllm.v1.attention.backends.mla.flashmla_sparse.split_prefill_chunks",
434	lambda rows, capacity: [(i, i + 1) for i in range(len(rows))],
435	)
436	if kv_cache_dtype not in backend_cls.supported_kv_cache_dtypes:
437	pytest.skip(f"{backend_cls.get_name()} does not support {kv_cache_dtype}")
438
439	if (
440	backend_cls == FlashMLASparseBackend
441	and kv_cache_dtype.startswith("fp8")
442	and kv_cache_dtype != "fp8_ds_mla"
443	):
444	pytest.skip(
445	"FlashMLA Sparse Attention backend fp8 only supports "
446	"fp8_ds_mla kv-cache dtype"
447	)
448
449	if kv_cache_dtype == "nvfp4_ds_mla":
450	# Unlike "fp8" (an alias canonicalized to "fp8_ds_mla" above), the
451	# NVFP4 DS-MLA dtype must reach the backend unchanged.
452	assert (
453	_canonicalize_sparse_mla_kv_cache_dtype(backend_cls, kv_cache_dtype)
454	== kv_cache_dtype
455	)
456	device_capability = current_platform.get_device_capability()
457	if device_capability is None or device_capability.major != 10:
458	pytest.skip("The NVFP4 DS-MLA kv-cache dtype requires SM 10.x")
459
460	supported_block_sizes = backend_cls.get_supported_kernel_block_sizes()
461	if block_size not in supported_block_sizes:
462	pytest.skip(
463	f"{backend_cls.get_name()} does not support block_size={block_size}"
464	)
465
466	if backend_cls == FlashMLASparseBackend:
467	ok, reason = flashmla.is_flashmla_sparse_supported()
468	if not ok:
469	pytest.skip(reason)
470	elif backend_cls == FlashInferMLASparseTRTLLMBackend:
471	device_capability = current_platform.get_device_capability()
472	if device_capability is None or not backend_cls.supports_compute_capability(
473	device_capability
474	):
475	pytest.skip("FlashInferMLASparseTRTLLMBackend requires SM 10.x capability")
476
477	batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name]
478	use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla"
479	use_nvfp4_ds_mla_quantization = kv_cache_dtype == "nvfp4_ds_mla"
480
481	device = torch.device(DEVICE_TYPE)
482	dtype = torch.bfloat16
483
484	# Model hyper-parameters (kept intentionally small for the unit test)
485	total_num_heads = 128
486	# Compute per-rank heads for simulated TP
487	num_heads = max(1, total_num_heads // tensor_parallel_size)
488
489	kv_lora_rank = 512
490	qk_nope_head_dim = 128
491	qk_rope_head_dim = 64
492	v_head_dim = 128
493	head_size = kv_lora_rank + qk_rope_head_dim
494	topk_tokens = 128
495
496	max_seqlen = max(batch_spec.seq_lens)
497	total_cache_tokens = sum(batch_spec.seq_lens)
498
499	# Note: We use TP=1 to avoid multi-GPU requirements in CI.
500	# The test simulates head partitioning via mocked methods below.
501	vllm_config = create_vllm_config(
502	model_name="deepseek-ai/DeepSeek-V2-Lite-Chat",
503	tensor_parallel_size=1,
504	max_model_len=max_seqlen,
505	num_gpu_blocks=max(2048, cdiv(total_cache_tokens, block_size) + 1),
506	block_size=block_size,
507	hf_config_override={
508	"index_topk": topk_tokens,
509	"attn_module_list_cfg": [{"topk_tokens": topk_tokens}],
510	},
511	)
512	model_config = vllm_config.model_config
513	model_config.hf_text_config = SimpleNamespace(
514	index_topk=topk_tokens,
515	q_lora_rank=None,
516	kv_lora_rank=kv_lora_rank,
517	qk_nope_head_dim=qk_nope_head_dim,
518	qk_rope_head_dim=qk_rope_head_dim,
519	v_head_dim=v_head_dim,
520	model_type="deepseek_v2",
521	)
522	del model_config.hf_config.index_topk # Composite configs only nest this field.
523	model_config.dtype = dtype
524	model_config.get_num_attention_heads = MethodType(
525	lambda self, parallel_config: num_heads,
526	model_config,
527	)
528	model_config.get_num_kv_heads = MethodType(
529	lambda self, parallel_config: 1, model_config
530	)
531	model_config.get_head_size = MethodType(lambda self: head_size, model_config)
532	model_config.get_sliding_window = MethodType(lambda self: None, model_config)
533
534	kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
535
536	torch.manual_seed(0)
537
538	scale = 1.0 / math.sqrt(head_size)
539
540	# Shared MLA projection weights to keep reference and backend in sync
541	W_UK = torch.rand(
542	kv_lora_rank, num_heads, qk_nope_head_dim, dtype=dtype, device=device
543	)
544	W_UV = torch.rand(kv_lora_rank, num_heads, v_head_dim, dtype=dtype, device=device)
545
546	# Build synthetic decode-only workload
547	seq_lens = batch_spec.seq_lens
548	query_lens = batch_spec.query_lens
549
550	# Pre-compute positions and sparse indices for all tokens.
551	# We need these BEFORE computing the reference to use sparse attention masks.
552	total_query_tokens = sum(query_lens)
553	positions = []
554	for i in range(batch_spec.batch_size):
555	s_len = seq_lens[i]
556	q_len = query_lens[i]
557	ctx_len = s_len - q_len
558	for q_idx in range(q_len):
559	positions.append(ctx_len + q_idx)
560
561	# Create sparse indices with UNIQUE per-token offsets to catch bugs where
562	# the kernel uses wrong indices for some tokens (e.g., due to incorrect
563	# tensor shapes like [1, num_tokens, ...] instead of [num_tokens, 1, ...]).
564	# Also include -1 masked indices to verify the kernel handles them correctly.
565	sparse_indices = torch.empty(
566	total_query_tokens, topk_tokens, dtype=torch.int32, device=device
567	)
568	for tok_idx in range(total_query_tokens):
569	max_valid_idx = positions[tok_idx]
570	offset = tok_idx * 7 # Prime number for varied offsets
571	# Use only half the topk indices as valid, mask the rest with -1
572	# This tests that the kernel correctly ignores -1 indices
573	num_valid = min(topk_tokens // 2, max_valid_idx + 1)
574	if num_valid > 0:
575	valid_range = torch.arange(num_valid, device=device, dtype=torch.int32)
576	tok_indices = (valid_range + offset) % (max_valid_idx + 1)
577	# Pad with -1 for the remaining positions
578	tok_indices = torch.cat(
579	[
580	tok_indices,
581	torch.full(
582	(topk_tokens - num_valid,), -1, device=device, dtype=torch.int32
583	),
584	]
585	)
586	else:
587	tok_indices = torch.full(
588	(topk_tokens,), -1, device=device, dtype=torch.int32
589	)
590	tok_indices[0] = 0 # At least one valid index
591	sparse_indices[tok_idx] = tok_indices
592
593	all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], []
594	kv_c_contexts, k_pe_contexts = [], []
595	reference_outputs = []
596
597	kv_cache_scale = torch.tensor(k_scale, dtype=torch.float32, device=device)
598	global_token_idx = 0
599
600	for i in range(batch_spec.batch_size):
601	s_len = seq_lens[i]
602	q_len = query_lens[i]
603	ctx_len = s_len - q_len
604
605	q_c = torch.rand(
606	q_len,
607	num_heads,
608	qk_nope_head_dim + qk_rope_head_dim,
609	dtype=dtype,
610	device=device,
611	)
612	kv_c_full = torch.rand(s_len, kv_lora_rank, dtype=dtype, device=device)
613	k_pe_full = torch.rand(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device)
614
615	if use_fp8_ds_mla_quantization:
616	is_sm100 = torch.cuda.get_device_capability()[0] >= 10
617	kv_c_full, k_pe_squeezed = _quantize_dequantize_fp8_ds_mla(
618	kv_c_full,
619	k_pe_full.squeeze(1),
620	block_size=block_size,
621	scale=kv_cache_scale,
622	simulate_sm100_e8m0_scales=is_sm100,
623	)
624	k_pe_full = k_pe_squeezed.unsqueeze(1)
625	elif use_nvfp4_ds_mla_quantization:
626	kv_c_full, k_pe_squeezed = _quantize_dequantize_nvfp4_ds_mla(
627	kv_c_full,
628	k_pe_full.squeeze(1),
629	block_size=block_size,
630	kv_cache_dtype=kv_cache_dtype,
631	scale=kv_cache_scale,
632	)
633	k_pe_full = k_pe_squeezed.unsqueeze(1)
634
635	q_nope, q_pe = q_c.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)
636	ql_nope = torch.einsum("qnh,lnh->qnl", q_nope, W_UK)
637	q_mqa = torch.cat([ql_nope, q_pe], dim=-1)
638
639	k_mqa = torch.cat([kv_c_full, k_pe_full.squeeze(1)], dim=-1)
640	v_mqa = kv_c_full
641
642	# Compute sparse SDPA reference per query token using its sparse indices
643	for q_idx in range(q_len):
644	tok_sparse_idx = sparse_indices[global_token_idx]
645	valid_mask = tok_sparse_idx >= 0
646	valid_indices = tok_sparse_idx[valid_mask].long()
647
648	q_tok = q_mqa[q_idx : q_idx + 1] # [1, num_heads, head_dim]
649	k_sparse = k_mqa[valid_indices] # [num_valid, head_dim]
650	v_sparse = v_mqa[valid_indices] # [num_valid, kv_lora_rank]
651
652	k_sparse = k_sparse.unsqueeze(1).expand(-1, num_heads, -1)
653	v_sparse = v_sparse.unsqueeze(1).expand(-1, num_heads, -1)
654
655	# SDPA: [1, num_heads, 1, head_dim] x [1, num_heads, num_valid, head_dim]
656	q_sdpa_in = q_tok.unsqueeze(0).transpose(1, 2)
657	k_sdpa_in = k_sparse.unsqueeze(0).transpose(1, 2)
658	v_sdpa_in = v_sparse.unsqueeze(0).transpose(1, 2)
659
660	sdpa_out = torch.nn.functional.scaled_dot_product_attention(
661	q_sdpa_in, k_sdpa_in, v_sdpa_in, scale=scale
662	)
663	sdpa_out = sdpa_out.transpose(1, 2).squeeze(
664	0
665	) # [1, num_heads, kv_lora_rank]
666
667	sdpa_out = torch.einsum("qnl,lnv->qnv", sdpa_out, W_UV)
668	reference_outputs.append(sdpa_out.flatten(start_dim=-2))
669
670	global_token_idx += 1
671
672	all_q_vllm.append(q_c)
673	all_kv_c_vllm.append(kv_c_full[ctx_len:])
674	all_k_pe_vllm.append(k_pe_full[ctx_len:])
675	kv_c_contexts.append(kv_c_full[: ctx_len + 1])
676	k_pe_contexts.append(k_pe_full[: ctx_len + 1])
677
678	query_vllm = torch.cat(all_q_vllm, dim=0)
679	kv_c_vllm = torch.cat(all_kv_c_vllm, dim=0)
680	k_pe_vllm = torch.cat(all_k_pe_vllm, dim=0)
681	sdpa_reference = torch.cat(reference_outputs, dim=0)
682
683	vllm_config.cache_config.cache_dtype = kv_cache_dtype
684	vllm_config.model_config.hf_text_config.index_topk = topk_tokens
685
686	common_attn_metadata = create_common_attn_metadata(
687	batch_spec,
688	vllm_config.cache_config.block_size,
689	device,
690	arange_block_indices=True,
691	)
692
693	kv_cache = create_and_prepopulate_kv_cache(
694	kv_c_contexts=kv_c_contexts,
695	k_pe_contexts=k_pe_contexts,
696	block_size=vllm_config.cache_config.block_size,
697	head_size=head_size,
698	dtype=dtype,
699	device=device,
700	num_blocks=vllm_config.cache_config.num_gpu_blocks,
701	common_attn_metadata=common_attn_metadata,
702	randomize_blocks=False,
703	kv_cache_dtype=kv_cache_dtype,
704	scale=kv_cache_scale,
705	)
706
707	# The sparse builder clones the layer's dense-MHA prefill backend from
708	# static_forward_context; register a mock layer carrying one.
709	prefill_backend = get_mla_prefill_backend(vllm_config)(
710	num_heads=num_heads,
711	scale=scale,
712	kv_lora_rank=kv_lora_rank,
713	qk_nope_head_dim=qk_nope_head_dim,
714	qk_rope_head_dim=qk_rope_head_dim,
715	v_head_dim=v_head_dim,
716	vllm_config=vllm_config,
717	)
718	vllm_config.compilation_config.static_forward_context["placeholder"] = (
719	SimpleNamespace(prefill_backend=prefill_backend)
720	)
721
722	builder_cls = backend_cls.get_builder_cls()
723	builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device)
724	metadata = builder.build(
725	common_prefix_len=0, common_attn_metadata=common_attn_metadata
726	)
727
728	# Use the pre-computed sparse_indices for the mock indexer
729	mock_indexer = SimpleNamespace(topk_indices_buffer=sparse_indices)
730
731	kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1)
732	kv_b_proj_weight = kv_b_proj_weight.view(
733	kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim)
734	)
735
736	mock_kv_b_proj = ColumnParallelLinear(
737	input_size=kv_lora_rank,
738	output_size=num_heads * (qk_nope_head_dim + v_head_dim),
739	bias=False,
740	).to(device=device, dtype=dtype)
741	mock_kv_b_proj.weight = torch.nn.Parameter(kv_b_proj_weight.T.contiguous())
742
743	impl_cls = backend_cls.get_impl_cls()
744	with set_current_vllm_config(vllm_config):
745	impl = impl_cls(
746	num_heads=num_heads,
747	head_size=head_size,
748	scale=scale,
749	num_kv_heads=1,
750	alibi_slopes=None,
751	sliding_window=None,
752	kv_cache_dtype=vllm_config.cache_config.cache_dtype,
753	logits_soft_cap=None,
754	attn_type="decoder",
755	kv_sharing_target_layer_name=None,
756	q_lora_rank=None,
757	kv_lora_rank=kv_lora_rank,
758	qk_nope_head_dim=qk_nope_head_dim,
759	qk_rope_head_dim=qk_rope_head_dim,
760	qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,
761	v_head_dim=v_head_dim,
762	kv_b_proj=mock_kv_b_proj,
763	indexer=mock_indexer,
764	)
765
766	impl.process_weights_after_loading(dtype)
767
768	# Create mock sparse MLA layer with weight matrices
769	mock_layer = MockSparseMLAAttentionLayer(
770	impl=impl,
771	num_heads=num_heads,
772	qk_nope_head_dim=qk_nope_head_dim,
773	qk_rope_head_dim=qk_rope_head_dim,
774	v_head_dim=v_head_dim,
775	kv_lora_rank=kv_lora_rank,
776	device=device,
777	W_UK=W_UK,
778	W_UV=W_UV,
779	q_scale=q_scale,
780	k_scale=k_scale,
781	)
782
783	out_buffer = torch.empty(
784	metadata.num_actual_tokens, num_heads * v_head_dim, dtype=dtype, device=device
785	)
786
787	with torch.inference_mode():
788	if backend_cls == FlashMLASparseBackend and kv_cache_dtype == "fp8_ds_mla":
789	from vllm.v1.worker.workspace import current_workspace_manager
790
791	current_workspace_manager().lock()
792	backend_output = mock_layer.forward_impl(
793	query_vllm,
794	kv_c_vllm,
795	k_pe_vllm,
796	kv_cache,
797	metadata,
798	out_buffer,
799	)
800
801	assert backend_output.shape == sdpa_reference.shape
802	assert backend_output.dtype == sdpa_reference.dtype
803	assert torch.isfinite(backend_output).all()
804
805	# FP8 quantization introduces some error, but should be within reasonable bounds
806	# BF16 (auto) should be very accurate, FP8 allows slightly more tolerance
807	if kv_cache_dtype.startswith("fp8"):
808	torch.testing.assert_close(
809	backend_output, sdpa_reference, rtol=0.065, atol=0.05
810	)
811	else:
812	torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.01, atol=0.01)
813
814
815	def _triton_convert_reference_impl(
816	req_ids: torch.Tensor,
817	block_table: torch.Tensor,
818	token_indices: torch.Tensor,
819	block_size: int,
820	num_topk_tokens: int,
821	HAS_PREFILL_WORKSPACE: bool = False,
822	prefill_workspace_request_ids: torch.Tensor \| None = None,
823	prefill_workspace_starts: torch.Tensor \| None = None,
824	) -> torch.Tensor:
825	"""Reference implementation for triton_convert_req_index_to_global_index."""
826	num_tokens = req_ids.shape[0]
827	max_blocks_per_req = block_table.shape[1]
828	result = torch.empty(
829	num_tokens, num_topk_tokens, dtype=torch.int32, device=req_ids.device
830	)
831
832	for token_id in range(num_tokens):
833	req_id = req_ids[token_id].item()
834
835	# Determine if this token uses workspace or paged cache
836	use_prefill_workspace = False
837	workspace_start = 0
838	if HAS_PREFILL_WORKSPACE and prefill_workspace_request_ids is not None:
839	assert prefill_workspace_starts is not None
840	prefill_req_id = prefill_workspace_request_ids[token_id].item()
841	if prefill_req_id >= 0:
842	use_prefill_workspace = True
843	workspace_start = prefill_workspace_starts[prefill_req_id].item()
844
845	for idx_id in range(num_topk_tokens):
846	token_idx = token_indices[token_id, idx_id].item()
847
848	if token_idx == -1:
849	result[token_id, idx_id] = -1
850	elif use_prefill_workspace:
851	# Prefill + using prefill workspace: map to workspace offset
852	result[token_id, idx_id] = workspace_start + token_idx
853	else:
854	# Decode: map to paged cache
855	block_id = token_idx // block_size
856	if block_id >= max_blocks_per_req:
857	result[token_id, idx_id] = -1
858	else:
859	block_num = block_table[req_id, block_id].item()
860	offset = token_idx % block_size
861	result[token_id, idx_id] = block_num * block_size + offset
862
863	return result
864
865
866	@pytest.mark.parametrize("block_size", [16, 64, 128])
867	@pytest.mark.parametrize("num_topk_tokens", [128, 256, 512])
868	@pytest.mark.skipif(
869	torch.cuda.get_device_capability() < (9, 0),
870	reason="FlashMLASparseBackend requires CUDA 9.0 or higher",
871	)
872	def test_triton_convert_req_index_to_global_index_decode_only(
873	block_size, num_topk_tokens
874	):
875	device = torch.device(DEVICE_TYPE)
876	num_tokens = 8
877	num_requests = 4
878	max_blocks_per_req = 10
879
880	req_id = torch.randint(
881	0, num_requests, (num_tokens,), dtype=torch.int32, device=device
882	)
883	block_table = torch.randint(
884	0, 100, (num_requests, max_blocks_per_req), dtype=torch.int32, device=device
885	)
886
887	token_indices = torch.randint(
888	0,
889	block_size * max_blocks_per_req,
890	(num_tokens, num_topk_tokens),
891	dtype=torch.int32,
892	device=device,
893	)
894
895	# Set some to -1 to test masking
896	token_indices[0, :10] = -1
897	token_indices[3, 50:60] = -1
898
899	# Set some to out of bounds
900	token_indices[2, 100:110] = max_blocks_per_req * block_size
901	token_indices[6, 150:160] = max_blocks_per_req * block_size
902
903	result = triton_convert_req_index_to_global_index(
904	req_id,
905	block_table,
906	token_indices,
907	BLOCK_SIZE=block_size,
908	NUM_TOPK_TOKENS=num_topk_tokens,
909	)
910
911	reference_result = _triton_convert_reference_impl(
912	req_id,
913	block_table,
914	token_indices,
915	block_size,
916	num_topk_tokens,
917	)
918
919	torch.testing.assert_close(result, reference_result, rtol=0, atol=0)
920
921
922	def test_index_group_convert_during_piecewise_capture():
923	"""In piecewise cudagraph mode the indexer is captured while the convert
924	runs in the following eager break. The convert's side stream must not
925	wait on ``logical_topk_ready`` there: an event recorded inside a captured
926	segment is graph-local, and an eager wait on it raises
927	cudaErrorInvalidValue."""
928	device = torch.device(DEVICE_TYPE)
929	num_tokens, num_topk, num_requests, blocks_per_req, block_size = 8, 128, 4, 4, 16
930
931	logical_topk_indices = torch.randint(
932	0,
933	block_size * blocks_per_req,
934	(num_tokens, num_topk),
935	dtype=torch.int32,
936	device=device,
937	)
938	builder = SparseMLAIndexGroupBuilder(logical_topk_indices)
939	group, layer_index = builder.register_layer(is_index_producing_layer=True)
940	assert layer_index == 0 and group.has_indexer
941
942	req_id = torch.arange(num_tokens, dtype=torch.int32, device=device) % num_requests
943	block_table = torch.arange(
944	num_requests * blocks_per_req, dtype=torch.int32, device=device
945	).view(num_requests, blocks_per_req)
946	attn_metadata = SimpleNamespace(
947	req_id_per_token=req_id, block_table=block_table, block_size=block_size
948	)
949	expected = triton_convert_req_index_to_global_index(
950	req_id,
951	block_table,
952	logical_topk_indices,
953	BLOCK_SIZE=block_size,
954	NUM_TOPK_TOKENS=num_topk,
955	)
956
957	vllm_config = create_vllm_config(model_name="Qwen/Qwen3.5-0.8B")
958	marker = torch.zeros(1, device=device)
959	capture_stream = torch.cuda.Stream(device=device)
960	# torch.cuda.stream routes through vLLM's patched torch.cuda.set_stream,
961	# so current_stream() tracks the capture stream as in production.
962	with (
963	torch.cuda.stream(capture_stream),
964	set_forward_context(
965	None, vllm_config, cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE
966	),
967	):
968	capture = BreakableCUDAGraphCapture()
969	with capture:
970	marker.add_(1)
971	# Record inside a captured segment, convert in the eager
972	# break — mirrors mla.py -> unified_mla_attention_with_output.
973	group.set_logical_topk_ready(layer_index)
974	capture.add_eager(
975	lambda: group.convert_logical_to_physical_topk(
976	layer_index,
977	logical_topk_indices,
978	attn_metadata,
979	block_stride_rows=None,
980	return_valid_counts=False,
981	)
982	)
983	marker.add_(1)
984	capture.replay()
985	capture_stream.synchronize()
986
987	result = group.physical_topk_indices[:num_tokens]
988	torch.testing.assert_close(result, expected, rtol=0, atol=0)
989
990
991	@pytest.mark.parametrize("block_size", [16])
992	@pytest.mark.skipif(
993	torch.cuda.get_device_capability() < (9, 0),
994	reason="FlashMLASparseBackend requires CUDA 9.0 or higher",
995	)
996	def test_triton_convert_req_index_to_global_index_with_prefill_workspace(block_size):
997	device = torch.device(DEVICE_TYPE)
998	num_requests = 4
999	max_blocks_per_req = 8
1000	num_topk_tokens = 128
1001
1002	# First 6 tokens are decode (reqs 0, 1), last 6 are prefill (reqs 2, 3)
1003	req_id = torch.tensor(
1004	[0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3], dtype=torch.int32, device=device
1005	)
1006	prefill_workspace_request_ids = torch.tensor(
1007	[-1, -1, -1, -1, -1, -1, 0, 0, 0, 1, 1, 1], dtype=torch.int32, device=device
1008	)
1009
1010	# Workspace starts for the 2 prefill reqs: req 2 starts at 0, req 3 starts at 100
1011	prefill_workspace_starts = torch.tensor([0, 100], dtype=torch.int32, device=device)
1012
1013	block_table = torch.randint(
1014	0, 50, (num_requests, max_blocks_per_req), dtype=torch.int32, device=device
1015	)
1016	token_indices = torch.randint(
1017	0,
1018	block_size * max_blocks_per_req,
1019	(req_id.shape[0], num_topk_tokens),
1020	dtype=torch.int32,
1021	device=device,
1022	)
1023
1024	# Set some to -1 to test masking
1025	token_indices[0, :10] = -1
1026	token_indices[3, 50:60] = -1
1027
1028	# Set some to out of bounds
1029	token_indices[2, 100:110] = max_blocks_per_req * block_size
1030	token_indices[6, 150:160] = max_blocks_per_req * block_size
1031
1032	result = triton_convert_req_index_to_global_index(
1033	req_id,
1034	block_table,
1035	token_indices,
1036	BLOCK_SIZE=block_size,
1037	NUM_TOPK_TOKENS=num_topk_tokens,
1038	HAS_PREFILL_WORKSPACE=True,
1039	prefill_workspace_request_ids=prefill_workspace_request_ids,
1040	prefill_workspace_starts=prefill_workspace_starts,
1041	)
1042
1043	reference_result = _triton_convert_reference_impl(
1044	req_id,
1045	block_table,
1046	token_indices,
1047	block_size,
1048	num_topk_tokens,
1049	HAS_PREFILL_WORKSPACE=True,
1050	prefill_workspace_request_ids=prefill_workspace_request_ids,
1051	prefill_workspace_starts=prefill_workspace_starts,
1052	)
1053
1054	torch.testing.assert_close(result, reference_result, rtol=0, atol=0)
1055
1056
1057	@pytest.mark.skipif(
1058	torch.cuda.get_device_capability() < (9, 0),
1059	reason="FlashMLASparseBackend requires CUDA 9.0 or higher",
1060	)
1061	def test_triton_convert_rejects_req_id_longer_than_token_indices():
1062	"""Guard against the #47327 regression: the kernel grid is sized by
1063	req_id but the output is allocated like token_indices, so a full-batch
1064	req_id combined with an MQA-subset token_indices wrote past the end of
1065	the output buffer. The wrapper must reject the length mismatch instead
1066	of corrupting memory."""
1067	device = torch.device(DEVICE_TYPE)
1068	num_topk_tokens = 128
1069	block_size = 64
1070	block_table = torch.arange(40, dtype=torch.int32, device=device).view(4, 10)
1071
1072	# Full batch: 2 decode tokens + 10 prefill tokens
1073	req_id_full = torch.tensor(
1074	[0, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3], dtype=torch.int32, device=device
1075	)
1076	num_mqa_tokens = 2
1077	token_indices = torch.randint(
1078	0,
1079	block_size * 10,
1080	(num_mqa_tokens, num_topk_tokens),
1081	dtype=torch.int32,
1082	device=device,
1083	)
1084
1085	with pytest.raises(AssertionError, match="must cover the same tokens"):
1086	triton_convert_req_index_to_global_index(
1087	req_id_full,
1088	block_table,
1089	token_indices,
1090	BLOCK_SIZE=block_size,
1091	NUM_TOPK_TOKENS=num_topk_tokens,
1092	)
1093
1094	# The sliced call is the intended usage and must match the reference.
1095	result = triton_convert_req_index_to_global_index(
1096	req_id_full[:num_mqa_tokens],
1097	block_table,
1098	token_indices,
1099	BLOCK_SIZE=block_size,
1100	NUM_TOPK_TOKENS=num_topk_tokens,
1101	)
1102	reference = _triton_convert_reference_impl(
1103	req_id_full[:num_mqa_tokens],
1104	block_table,
1105	token_indices,
1106	block_size,
1107	num_topk_tokens,
1108	)
1109	torch.testing.assert_close(result, reference, rtol=0, atol=0)
1110
1111
1112	@pytest.mark.skipif(
1113	torch.cuda.get_device_capability() < (9, 0),
1114	reason="FlashMLASparseBackend requires CUDA 9.0 or higher",
1115	)
1116	def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens():
1117	"""Guard against the #47327 regression: when the dense-MHA prefill split
1118	is active, forward_mqa only receives the leading decode tokens, but
1119	_forward_bf16_kv passed the full-batch req_id_per_token to the index
1120	conversion, making it write past the end of its output buffer. The call
1121	site must slice req_id_per_token to the MQA tokens."""
1122	device = torch.device(DEVICE_TYPE)
1123	num_topk_tokens = 128
1124	block_size = 64
1125	num_batch_tokens = 12
1126	num_mqa_tokens = 2
1127
1128	attn_metadata = SimpleNamespace(
1129	req_id_per_token=torch.tensor(
1130	[0, 1] + [2] * 5 + [3] * 5, dtype=torch.int32, device=device
1131	),
1132	block_table=torch.arange(40, dtype=torch.int32, device=device).view(4, 10),
1133	block_size=block_size,
1134	num_decode_tokens=num_mqa_tokens,
1135	num_actual_tokens=num_mqa_tokens,
1136	)
1137	assert attn_metadata.req_id_per_token.shape[0] == num_batch_tokens
1138
1139	q = torch.zeros(num_mqa_tokens, 4, 576, dtype=torch.bfloat16, device=device)
1140	kv_cache = torch.zeros(40, block_size, 576, dtype=torch.bfloat16, device=device)
1141	topk_indices = torch.randint(
1142	0,
1143	block_size * 10,
1144	(num_mqa_tokens, num_topk_tokens),
1145	dtype=torch.int32,
1146	device=device,
1147	)
1148
1149	captured = {}
1150
1151	def _stub_kernel(q, kv, indices, lengths, actual_num_heads):
1152	captured["indices"] = indices
1153	captured["actual_num_heads"] = actual_num_heads
1154	return (
1155	torch.zeros(q.shape[0], q.shape[1], 512, dtype=q.dtype, device=q.device),
1156	None,
1157	)
1158
1159	def _convert_topk(indices, metadata, **kwargs):
1160	return triton_convert_req_index_to_global_index(
1161	metadata.req_id_per_token[: indices.shape[0]],
1162	metadata.block_table,
1163	indices,
1164	BLOCK_SIZE=metadata.block_size,
1165	BLOCK_STRIDE_ROWS=kwargs["block_stride_rows"],
1166	NUM_TOPK_TOKENS=indices.shape[1],
1167	return_valid_counts=kwargs["return_valid_counts"],
1168	)
1169
1170	stub_impl = SimpleNamespace(
1171	_bf16_flash_mla_kernel=_stub_kernel,
1172	_convert_logical_to_physical_topk=_convert_topk,
1173	index_group=None,
1174	index_group_index=0,
1175	pcp_dcp_kv_gather=False,
1176	)
1177
1178	out, _ = FlashMLASparseImpl._forward_bf16_kv(
1179	stub_impl, q, kv_cache, topk_indices, attn_metadata, q.shape[1]
1180	)
1181
1182	assert out.shape[0] == num_mqa_tokens
1183	assert captured["indices"].shape[0] == num_mqa_tokens
1184	assert captured["actual_num_heads"] == q.shape[1]
1185	reference = _triton_convert_reference_impl(
1186	attn_metadata.req_id_per_token[:num_mqa_tokens],
1187	attn_metadata.block_table,
1188	topk_indices,
1189	block_size,
1190	num_topk_tokens,
1191	)
1192	torch.testing.assert_close(captured["indices"], reference, rtol=0, atol=0)
1193
1194
1195	@pytest.mark.parametrize(
1196	"seq_lens,max_buf,expected",
1197	[
1198	# Basic split: totals per chunk ≤ max_buf
1199	(torch.tensor([2, 3, 4, 2]), 5, [(0, 2), (2, 3), (3, 4)]),
1200	# Exact fits should split between items when adding the next would overflow
1201	(torch.tensor([5, 5, 5]), 5, [(0, 1), (1, 2), (2, 3)]),
1202	# All requests fit in a single chunk
1203	(torch.tensor([1, 1, 1]), 10, [(0, 3)]),
1204	# Large buffer
1205	(torch.tensor([4, 4, 4]), 100, [(0, 3)]),
1206	],
1207	)
1208	def test_split_prefill_chunks(seq_lens, max_buf, expected):
1209	out = split_prefill_chunks(seq_lens, max_buf)
1210	assert out == expected
1211
1212
1213	@pytest.mark.parametrize(
1214	("max_query_len", "expected"),
1215	[(32768, True), (33024, False)],
1216	)
1217	def test_masked_mha_workspace_fits_single_request_boundary(max_query_len, expected):
1218	"""A 32K prefill needs the default workspace exactly; shrinking it would
1219	push a supported request onto MQA."""
1220	assert (
1221	_masked_mha_workspace_fits(
1222	batch_size=1,
1223	max_query_len=max_query_len,
1224	max_context_chunk_seq_len=0,
1225	workspace_numel=GLOBAL_TOPK_MASK_MAX_BYTES // torch.int32.itemsize,
1226	)
1227	is expected
1228	)
1229
1230
1231	@pytest.mark.parametrize(
1232	("backend_name", "tensor_parallel_size", "query_len"),
1233	[
1234	("FLASHMLA_SPARSE", 4, 48 * 1024),
1235	("FLASHMLA_SPARSE", 8, 112 * 1024),
1236	("FLASHINFER_MLA_SPARSE", 4, 36 * 1024),
1237	("FLASHINFER_MLA_SPARSE", 8, 64 * 1024),
1238	],
1239	)
1240	def test_masked_mha_workspace_guards_long_routing_policy(
1241	backend_name, tensor_parallel_size, query_len
1242	):
1243	assert _use_masked_mha(
1244	backend_name=backend_name,
1245	tensor_parallel_size=tensor_parallel_size,
1246	qk_head_dim=256,
1247	v_head_dim=256,
1248	query_len=query_len,
1249	seq_len=query_len,
1250	has_context=False,
1251	)
1252	assert not _masked_mha_workspace_fits(
1253	batch_size=1,
1254	max_query_len=query_len,
1255	max_context_chunk_seq_len=0,
1256	workspace_numel=GLOBAL_TOPK_MASK_MAX_BYTES // torch.int32.itemsize,
1257	)
1258
1259
1260	@pytest.mark.parametrize(
1261	("model_dims", "kv_cache_dtype", "fa_version", "expected"),
1262	[
1263	pytest.param((128, 512, 128, 64, 128), "auto", 4, True, id="deepseek_v32"),
1264	pytest.param((64, 512, 192, 64, 256), "auto", 4, True, id="glm5"),
1265	pytest.param((64, 512, 256, 0, 256), "auto", 4, True, id="glm53_flash_nope"),
1266	pytest.param((64, 512, 256, 0, 256), "fp8", 4, False, id="quantized_kv"),
1267	pytest.param((64, 512, 256, 0, 256), "auto", 3, False, id="no_fa4"),
1268	pytest.param((64, 512, 256, 64, 256), "auto", 4, False, id="rope_320"),
1269	pytest.param((128, 512, 256, 0, 256), "auto", 4, False, id="wrong_heads"),
1270	],
1271	)
1272	def test_is_masked_mha_available_model_dims(
1273	monkeypatch, model_dims, kv_cache_dtype, fa_version, expected
1274	):
1275	"""The allow-list gates masked MHA per exact model geometry: the DeepSeek-V3.2,
1276	GLM-5 and NoPE GLM-5.3-Flash layouts on an SM100-family GPU with FA4 and an
1277	unquantized KV cache, nothing else."""
1278	import vllm.model_executor.layers.attention.sparse_mla_attention as mod
1279
1280	monkeypatch.setattr(
1281	mod.current_platform, "is_device_capability_family", lambda family: True
1282	)
1283	monkeypatch.setattr(mod, "get_flash_attn_version", lambda **kwargs: fa_version)
1284	assert _is_masked_mha_available(*model_dims, kv_cache_dtype) is expected
1285
1286
1287	def test_masked_mha_workspace_fits_accounts_for_batch_and_context():
1288	"""Request count and context chunk length are independent multipliers."""
1289	base = dict(batch_size=2, max_query_len=2048, max_context_chunk_seq_len=2048)
1290	exact = math.prod(_topk_mask_shape(2, 2048, 2048))
1291
1292	assert _masked_mha_workspace_fits(**base, workspace_numel=exact)
1293	assert not _masked_mha_workspace_fits(
1294	{base, "batch_size": 3}, workspace_numel=exact
1295	)
1296	assert not _masked_mha_workspace_fits(
1297	{base, "max_context_chunk_seq_len": 4096}, workspace_numel=exact
1298	)
1299
1300
1301	PREFILL_BATCH_SPECS = {
1302	"short_dense_mha": BatchSpec(seq_lens=[64, 128], query_lens=[64, 128]),
1303	"short_context_dense_mha": BatchSpec(seq_lens=[128, 160], query_lens=[64, 32]),
1304	"masked_mha": BatchSpec(seq_lens=[256], query_lens=[256]),
1305	"masked_mha_chunked_context": BatchSpec(seq_lens=[448, 384], query_lens=[256, 256]),
1306	}
1307
1308
1309	@pytest.mark.skipif(
1310	torch.cuda.get_device_capability()[0] < 10,
1311	reason="Sparse MLA forward_mha requires FA4 (SM100+)",
1312	)
1313	@pytest.mark.parametrize("batch_name", list(PREFILL_BATCH_SPECS.keys()))
1314	@pytest.mark.parametrize("kv_cache_dtype", ["auto"])
1315	@pytest.mark.parametrize(
1316	("num_heads", "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim"),
1317	[
1318	pytest.param(128, 128, 64, 128, id="deepseek_hd192_v128"),
1319	pytest.param(64, 192, 64, 256, id="glm5_hd256_v256"),
1320	pytest.param(64, 256, 0, 256, id="glm53_flash_nope_hd256_v256"),
1321	],
1322	)
1323	def test_sparse_backend_prefill_correctness(
1324	default_vllm_config,
1325	dist_init,
1326	batch_name,
1327	kv_cache_dtype,
1328	num_heads,
1329	qk_nope_head_dim,
1330	qk_rope_head_dim,
1331	v_head_dim,
1332	workspace_init,
1333	):
1334	"""Test dense and masked MHA across supported sparse MLA dimensions."""
1335	backend_cls = FlashMLASparseBackend
1336	batch_spec = PREFILL_BATCH_SPECS[batch_name]
1337
1338	device = torch.device("cuda")
1339	dtype = torch.bfloat16
1340	block_size = 64
1341
1342	kv_lora_rank = 512
1343	head_size = kv_lora_rank + qk_rope_head_dim
1344	masked_mha = batch_name.startswith("masked_mha")
1345	topk_tokens = 200 if masked_mha else 512
1346
1347	max_seqlen = max(batch_spec.seq_lens)
1348	total_cache_tokens = sum(batch_spec.seq_lens)
1349
1350	vllm_config = create_vllm_config(
1351	model_name="deepseek-ai/DeepSeek-V2-Lite-Chat",
1352	tensor_parallel_size=1,
1353	max_model_len=max_seqlen,
1354	num_gpu_blocks=max(2048, cdiv(total_cache_tokens, block_size) + 1),
1355	block_size=block_size,
1356	hf_config_override={
1357	"index_topk": topk_tokens,
1358	"attn_module_list_cfg": [{"topk_tokens": topk_tokens}],
1359	},
1360	)
1361	model_config = vllm_config.model_config
1362	model_config.hf_text_config = SimpleNamespace(
1363	index_topk=topk_tokens,
1364	q_lora_rank=None,
1365	kv_lora_rank=kv_lora_rank,
1366	qk_nope_head_dim=qk_nope_head_dim,
1367	qk_rope_head_dim=qk_rope_head_dim,
1368	v_head_dim=v_head_dim,
1369	model_type="deepseek_v2",
1370	)
1371	del model_config.hf_config.index_topk # Composite configs only nest this field.
1372	model_config.dtype = dtype
1373	model_config.model_arch_config.total_num_attention_heads = num_heads
1374	model_config.get_num_attention_heads = MethodType(
1375	lambda self, parallel_config: num_heads, model_config
1376	)
1377	model_config.get_num_kv_heads = MethodType(
1378	lambda self, parallel_config: 1, model_config
1379	)
1380	model_config.get_head_size = MethodType(lambda self: head_size, model_config)
1381	model_config.get_sliding_window = MethodType(lambda self: None, model_config)
1382
1383	kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
1384	scale = 1.0 / math.sqrt(qk_nope_head_dim + qk_rope_head_dim)
1385
1386	torch.manual_seed(42)
1387
1388	W_UK = torch.rand(
1389	kv_lora_rank, num_heads, qk_nope_head_dim, dtype=dtype, device=device
1390	)
1391	W_UV = torch.rand(kv_lora_rank, num_heads, v_head_dim, dtype=dtype, device=device)
1392
1393	seq_lens = batch_spec.seq_lens
1394	query_lens = batch_spec.query_lens
1395
1396	# Compute dense reference outputs.
1397	total_query_tokens = sum(query_lens)
1398	sparse_indices = torch.full(
1399	(total_query_tokens, topk_tokens), -1, dtype=torch.int32, device=device
1400	)
1401
1402	all_q, all_kv_c_new, all_k_pe_new = [], [], []
1403	kv_c_contexts, k_pe_contexts = [], []
1404	reference_outputs = []
1405	global_token_idx = 0
1406
1407	for i in range(batch_spec.batch_size):
1408	s_len = seq_lens[i]
1409	q_len = query_lens[i]
1410	ctx_len = s_len - q_len
1411
1412	q_mha = torch.rand(
1413	q_len,
1414	num_heads,
1415	qk_nope_head_dim + qk_rope_head_dim,
1416	dtype=dtype,
1417	device=device,
1418	)
1419	kv_c_full = torch.rand(s_len, kv_lora_rank, dtype=dtype, device=device)
1420	k_pe_full = torch.rand(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device)
1421
1422	# Decompress all KV for reference
1423	kv_b_weight = torch.cat([W_UK, W_UV], dim=-1).view(
1424	kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim)
1425	)
1426	kv_decompressed = (kv_c_full @ kv_b_weight).view(
1427	s_len, num_heads, qk_nope_head_dim + v_head_dim
1428	)
1429	k_nope_all, v_all = kv_decompressed.split(
1430	[qk_nope_head_dim, v_head_dim], dim=-1
1431	)
1432	k_pe_expanded = k_pe_full.expand(-1, num_heads, -1)
1433	k_all = torch.cat([k_nope_all, k_pe_expanded], dim=-1)
1434
1435	for j in range(q_len):
1436	attend_end = ctx_len + j + 1
1437	q_tok = q_mha[j : j + 1] # (1, H, D_qk)
1438	if masked_mha:
1439	actual_topk = min(topk_tokens, attend_end)
1440	attend_indices = torch.randperm(attend_end, device=device)[:actual_topk]
1441	sparse_indices[global_token_idx, :actual_topk] = attend_indices
1442	k_attend = k_all[attend_indices]
1443	v_attend = v_all[attend_indices]
1444	else:
1445	k_attend = k_all[:attend_end] # (N, H, D_qk)
1446	v_attend = v_all[:attend_end] # (N, H, D_v)
1447
1448	q_sdpa = q_tok.unsqueeze(0).transpose(1, 2).float()
1449	k_sdpa = k_attend.unsqueeze(0).transpose(1, 2).float()
1450	v_sdpa = v_attend.unsqueeze(0).transpose(1, 2).float()
1451
1452	out = torch.nn.functional.scaled_dot_product_attention(
1453	q_sdpa, k_sdpa, v_sdpa, scale=scale
1454	)
1455	out = out.transpose(1, 2).squeeze(0) # (1, H, D_v)
1456	reference_outputs.append(out.to(dtype).flatten(start_dim=-2))
1457	global_token_idx += 1
1458
1459	all_q.append(q_mha)
1460	all_kv_c_new.append(kv_c_full[ctx_len:])
1461	all_k_pe_new.append(k_pe_full[ctx_len:])
1462	kv_c_contexts.append(kv_c_full)
1463	k_pe_contexts.append(k_pe_full)
1464
1465	query_cat = torch.cat(all_q, dim=0)
1466	kv_c_cat = torch.cat(all_kv_c_new, dim=0)
1467	k_pe_cat = torch.cat(all_k_pe_new, dim=0)
1468	ref_output = torch.cat(reference_outputs, dim=0)
1469
1470	vllm_config.cache_config.cache_dtype = kv_cache_dtype
1471	vllm_config.model_config.hf_text_config.index_topk = topk_tokens
1472
1473	common_attn_metadata = create_common_attn_metadata(
1474	batch_spec,
1475	vllm_config.cache_config.block_size,
1476	device,
1477	arange_block_indices=True,
1478	)
1479
1480	kv_cache = create_and_prepopulate_kv_cache(
1481	kv_c_contexts=kv_c_contexts,
1482	k_pe_contexts=k_pe_contexts,
1483	block_size=block_size,
1484	head_size=head_size,
1485	dtype=dtype,
1486	device=device,
1487	num_blocks=vllm_config.cache_config.num_gpu_blocks,
1488	common_attn_metadata=common_attn_metadata,
1489	randomize_blocks=False,
1490	kv_cache_dtype=kv_cache_dtype,
1491	)
1492
1493	# The sparse builder clones the layer's dense-MHA prefill backend from
1494	# static_forward_context; register a mock layer carrying one.
1495	prefill_backend = get_mla_prefill_backend(vllm_config)(
1496	num_heads=num_heads,
1497	scale=scale,
1498	kv_lora_rank=kv_lora_rank,
1499	qk_nope_head_dim=qk_nope_head_dim,
1500	qk_rope_head_dim=qk_rope_head_dim,
1501	v_head_dim=v_head_dim,
1502	vllm_config=vllm_config,
1503	)
1504	vllm_config.compilation_config.static_forward_context["placeholder"] = (
1505	SimpleNamespace(prefill_backend=prefill_backend)
1506	)
1507
1508	builder_cls = backend_cls.get_builder_cls()
1509	builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device)
1510	if batch_name == "masked_mha_chunked_context":
1511	builder.chunked_prefill_workspace_size = block_size * batch_spec.batch_size
1512	builder.chunked_prefill_workspace = torch.empty(
1513	(builder.chunked_prefill_workspace_size, head_size),
1514	dtype=dtype,
1515	device=device,
1516	)
1517	# Drive the queries through the dense-MHA prefill path directly (the routing
1518	# threshold would otherwise classify these short queries as MQA decodes).
1519	builder.reorder_batch_threshold = 1
1520	metadata = builder.build(
1521	common_prefix_len=0, common_attn_metadata=common_attn_metadata
1522	)
1523
1524	mock_indexer = SimpleNamespace(topk_indices_buffer=sparse_indices)
1525
1526	kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1).view(
1527	kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim)
1528	)
1529
1530	mock_kv_b_proj = ColumnParallelLinear(
1531	input_size=kv_lora_rank,
1532	output_size=num_heads * (qk_nope_head_dim + v_head_dim),
1533	bias=False,
1534	).to(device=device, dtype=dtype)
1535	mock_kv_b_proj.weight = torch.nn.Parameter(kv_b_proj_weight.T.contiguous())
1536
1537	impl_cls = backend_cls.get_impl_cls()
1538	with set_current_vllm_config(vllm_config):
1539	impl = impl_cls(
1540	num_heads=num_heads,
1541	head_size=head_size,
1542	scale=scale,
1543	num_kv_heads=1,
1544	alibi_slopes=None,
1545	sliding_window=None,
1546	kv_cache_dtype=kv_cache_dtype,
1547	logits_soft_cap=None,
1548	attn_type="decoder",
1549	kv_sharing_target_layer_name=None,
1550	q_lora_rank=None,
1551	kv_lora_rank=kv_lora_rank,
1552	qk_nope_head_dim=qk_nope_head_dim,
1553	qk_rope_head_dim=qk_rope_head_dim,
1554	qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,
1555	v_head_dim=v_head_dim,
1556	kv_b_proj=mock_kv_b_proj,
1557	indexer=mock_indexer,
1558	)
1559	impl.process_weights_after_loading(dtype)
1560
1561	out_buffer = torch.empty(
1562	total_query_tokens, num_heads * v_head_dim, dtype=dtype, device=device
1563	)
1564
1565	with torch.inference_mode():
1566	impl.forward_mha(
1567	q=query_cat,
1568	kv_c_normed=kv_c_cat,
1569	k_pe=k_pe_cat,
1570	# Impls see the bind-time-squeezed [B, N, C] cache; mirror bind_kv_cache.
1571	kv_c_and_k_pe_cache=kv_cache.squeeze(1),
1572	attn_metadata=metadata,
1573	k_scale=torch.tensor(1.0, device=device),
1574	output=out_buffer,
1575	)
1576
1577	assert out_buffer.shape == ref_output.shape
1578	assert torch.isfinite(out_buffer).all(), "Non-finite values in output"
1579	torch.testing.assert_close(out_buffer, ref_output, rtol=0.01, atol=0.01)
1580
1581
1582	@pytest.mark.parametrize(
1583	"seq_lens,query_lens,workspace_size,max_logits_bytes,expected",
1584	[
1585	(
1586	torch.tensor([0]),
1587	torch.tensor([0]),
1588	100,
1589	1000,
1590	[],
1591	),
1592	# Logits constraint triggers split (M*N exceeds budget)
1593	# req0: M=10, N=100 -> 1000 elems (4000 bytes) - fits in 5000
1594	# req1: adding M=10, N=100 -> new_M=20, new_N=200 -> 4000 elems > 1250
1595	(
1596	torch.tensor([100, 100, 100]),
1597	torch.tensor([10, 10, 10]),
1598	1000, # workspace allows all
1599	5000, # 1250 float32 elems -> forces split
1600	[
1601	(slice(0, 1), slice(0, 10)),
1602	(slice(1, 2), slice(0, 10)),
1603	(slice(2, 3), slice(0, 10)),
1604	],
1605	),
1606	# Both constraints satisfied - all fit in one chunk
1607	(
1608	torch.tensor([10, 10, 10]),
1609	torch.tensor([5, 5, 5]),
1610	100,
1611	10000, # 2500 elems, MN = 1530 = 450 < 2500
1612	[(slice(0, 3), slice(0, 15))],
1613	),
1614	# Workspace constraint triggers first
1615	(
1616	torch.tensor([50, 50, 50]),
1617	torch.tensor([1, 1, 1]),
1618	50, # workspace only fits one at a time
1619	1000000, # logits budget is huge
1620	[
1621	(slice(0, 1), slice(0, 1)),
1622	(slice(1, 2), slice(0, 1)),
1623	(slice(2, 3), slice(0, 1)),
1624	],
1625	),
1626	# Greedy filling: first two fit, third doesn't
1627	# req0: M=5, N=10 -> 50 elems
1628	# req0+1: M=10, N=20 -> 200 elems <= 250
1629	# req0+1+2: M=15, N=30 -> 450 elems > 250
1630	(
1631	torch.tensor([10, 10, 10]),
1632	torch.tensor([5, 5, 5]),
1633	100,
1634	1000, # 250 elems
1635	[(slice(0, 2), slice(0, 10)), (slice(2, 3), slice(0, 5))],
1636	),
1637	],
1638	)
1639	def test_split_indexer_prefill_chunks(
1640	seq_lens, query_lens, workspace_size, max_logits_bytes, expected
1641	):
1642	out = DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks(
1643	seq_lens,
1644	query_lens,
1645	workspace_size,
1646	max_logits_bytes,
1647	)
1648	assert out == expected
1649
1650
1651	def test_split_indexer_prefill_chunks_single_request_overflow():
1652	"""Test that single request exceeding budget is sub-chunked on query dim."""
1653	seq_lens = torch.tensor([1000, 50])
1654	query_lens = torch.tensor([100, 5])
1655
1656	out = DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks(
1657	seq_lens, query_lens, 2000, 1000
1658	)
1659	# max_logits_elems = 250, N=1000 -> max_q = 1 -> 100 query sub-chunks
1660	expected = [(slice(0, 1), slice(i, i + 1)) for i in range(100)]
1661	# req1: M=5, N=50 -> 250 elems fits budget
1662	expected.append((slice(1, 2), slice(0, 5)))
1663	assert out == expected
1664
1665
1666	# Power-of-two, GLM's padded tile, and atomic fallback, with reused buffers.
1667	@pytest.mark.parametrize(
1668	"num_topk_tokens,reuse_buffers",
1669	[(128, False), (2176, False), (2176, True), (4224, False), (4224, True)],
1670	)
1671	def test_triton_convert_returns_valid_counts(num_topk_tokens: int, reuse_buffers: bool):
1672	"""Test that return_valid_counts correctly counts non-negative indices."""
1673	device = torch.device(DEVICE_TYPE)
1674	num_tokens = 8
1675	num_requests = 2
1676	max_blocks_per_req = 10
1677	block_size = 64
1678
1679	req_id = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1], dtype=torch.int32, device=device)
1680	block_table = torch.arange(
1681	num_requests * max_blocks_per_req, dtype=torch.int32, device=device
1682	).view(num_requests, max_blocks_per_req)
1683
1684	# Include partly valid, fully valid, and empty selections for each request.
1685	token_indices = torch.full(
1686	(num_tokens, num_topk_tokens), -1, dtype=torch.int32, device=device
1687	)
1688	valid_counts_per_token = [
1689	num_topk_tokens // 2,
1690	num_topk_tokens // 4,
1691	num_topk_tokens,
1692	0,
1693	] * 2
1694	expected_valid = []
1695	for i in range(num_tokens):
1696	num_valid = valid_counts_per_token[i]
1697	token_indices[i, :num_valid] = torch.arange(
1698	num_valid, dtype=torch.int32, device=device
1699	) % (block_size * max_blocks_per_req)
1700	expected_valid.append(num_valid)
1701
1702	expected_valid_tensor = torch.tensor(
1703	expected_valid, dtype=torch.int32, device=device
1704	)
1705
1706	buffers = {}
1707	if reuse_buffers:
1708	buffers = dict(
1709	out=torch.full_like(token_indices, -99),
1710	valid_counts_out=torch.full(
1711	(num_tokens,), 99, dtype=torch.int32, device=device
1712	),
1713	)
1714
1715	# Test with return_valid_counts=True
1716	result, valid_counts = triton_convert_req_index_to_global_index(
1717	req_id,
1718	block_table,
1719	token_indices,
1720	BLOCK_SIZE=block_size,
1721	NUM_TOPK_TOKENS=num_topk_tokens,
1722	return_valid_counts=True,
1723	**buffers,
1724	)
1725
1726	torch.testing.assert_close(valid_counts, expected_valid_tensor, rtol=0, atol=0)
1727
1728	# Test that return_valid_counts=False returns only the indices
1729	result_only = triton_convert_req_index_to_global_index(
1730	req_id,
1731	block_table,
1732	token_indices,
1733	BLOCK_SIZE=block_size,
1734	NUM_TOPK_TOKENS=num_topk_tokens,
1735	return_valid_counts=False,
1736	)
1737	assert isinstance(result_only, torch.Tensor)
1738	for row, num_valid in enumerate(expected_valid):
1739	compact_valid = result[row, :num_valid].sort().values
1740	original_valid = result_only[row][result_only[row] >= 0].sort().values
1741	torch.testing.assert_close(compact_valid, original_valid, rtol=0, atol=0)
1742	assert torch.all(result[row, num_valid:] == -1)
1743
1744
1745	def test_sparse_mla_index_group_converts_decode_indices_once(monkeypatch):
1746	device = torch.device(DEVICE_TYPE)
1747	logical = torch.full((2, 128), -1, dtype=torch.int32, device=device)
1748	logical[:, :4] = torch.tensor(
1749	[[0, 2, 5, -1], [1, 3, 7, -1]], dtype=torch.int32, device=device
1750	)
1751	builder = SparseMLAIndexGroupBuilder(logical)
1752	group, leader_index = builder.register_layer(True)
1753	follower_group, follower_index = builder.register_layer(False)
1754	assert follower_group is group
1755
1756	block_table = torch.tensor([[2, 3], [4, 5]], dtype=torch.int32, device=device)
1757	metadata = SimpleNamespace(
1758	block_table=block_table,
1759	block_size=4,
1760	req_id_per_token=torch.tensor([0, 1], dtype=torch.int32, device=device),
1761	num_decode_tokens=2,
1762	num_actual_tokens=2,
1763	)
1764	calls = []
1765	convert = index_group_module.triton_convert_req_index_to_global_index
1766
1767	def convert_spy(args, *kwargs):
1768	calls.append((args, kwargs))
1769	return convert(args, *kwargs)
1770
1771	monkeypatch.setattr(
1772	index_group_module,
1773	"triton_convert_req_index_to_global_index",
1774	convert_spy,
1775	)
1776	group.set_logical_topk_ready(leader_index)
1777	leader_result = group.convert_logical_to_physical_topk(
1778	leader_index,
1779	logical,
1780	metadata,
1781	block_stride_rows=None,
1782	return_valid_counts=True,
1783	)
1784	follower_result = group.convert_logical_to_physical_topk(
1785	follower_index,
1786	logical,
1787	metadata,
1788	block_stride_rows=None,
1789	return_valid_counts=True,
1790	)
1791	torch.accelerator.synchronize()
1792
1793	assert len(calls) == 1
1794	leader_indices, leader_counts = leader_result
1795	follower_indices, follower_counts = follower_result
1796	torch.testing.assert_close(follower_indices, leader_indices)
1797	torch.testing.assert_close(follower_counts, leader_counts)
1798
1799
1800	def test_sparse_mla_index_group_falls_back_for_prefill_sized_batch(monkeypatch):
1801	device = torch.device(DEVICE_TYPE)
1802	logical = torch.full((4, 128), -1, dtype=torch.int32, device=device)
1803	logical[:, 0] = torch.tensor([0, 1, 2, 3], dtype=torch.int32, device=device)
1804	builder = SparseMLAIndexGroupBuilder(logical, max_decode_rows=2)
1805	group, leader_index = builder.register_layer(True)
1806	follower_group, follower_index = builder.register_layer(False)
1807	assert follower_group is group
1808
1809	metadata = SimpleNamespace(
1810	block_table=torch.tensor(
1811	[[2], [3], [4], [5]], dtype=torch.int32, device=device
1812	),
1813	block_size=4,
1814	req_id_per_token=torch.arange(4, dtype=torch.int32, device=device),
1815	)
1816	calls = 0
1817	convert = index_group_module.triton_convert_req_index_to_global_index
1818
1819	def convert_spy(args, *kwargs):
1820	nonlocal calls
1821	calls += 1
1822	return convert(args, *kwargs)
1823
1824	monkeypatch.setattr(
1825	index_group_module,
1826	"triton_convert_req_index_to_global_index",
1827	convert_spy,
1828	)
1829	leader_result = group.convert_logical_to_physical_topk(
1830	leader_index,
1831	logical,
1832	metadata,
1833	block_stride_rows=None,
1834	return_valid_counts=False,
1835	)
1836	follower_result = group.convert_logical_to_physical_topk(
1837	follower_index,
1838	logical,
1839	metadata,
1840	block_stride_rows=None,
1841	return_valid_counts=False,
1842	)
1843	torch.accelerator.synchronize()
1844
1845	assert calls == 2
1846	torch.testing.assert_close(follower_result, leader_result)
1847
1848
1849	def test_sparse_mla_index_groups_own_distinct_physical_buffers():
1850	logical = torch.empty((8, 128), dtype=torch.int32, device=DEVICE_TYPE)
1851	builder = SparseMLAIndexGroupBuilder(logical, max_decode_rows=2)
1852
1853	first, _ = builder.register_layer(True)
1854	follower, _ = builder.register_layer(False)
1855	second, _ = builder.register_layer(True)
1856
1857	assert follower is first
1858	assert second is not first
1859	assert first.physical_topk_indices[:2].shape == (2, 128)
1860	assert second.physical_topk_indices[:2].shape == (2, 128)
1861	assert first.physical_topk_indices.data_ptr() != (
1862	second.physical_topk_indices.data_ptr()
1863	)
1864
1865
1866	# HiSparse is host-resident-only and kernel-only: runtime construction
1867	# raises without the compiled CUDA ops.
1868	requires_hisparse_ops = pytest.mark.skipif(
1869	not _has_hisparse_ops(),
1870	reason="HiSparse CUDA ops not compiled",
1871	)
1872
1873
1874	def _make_hisparse_runtime(
1875	*,
1876	top_k: int = 4,
1877	device_buffer_size: int = 5,
1878	max_num_reqs: int = 2,
1879	row_width: int = 8,
1880	block_size: int = 64,
1881	max_swap_rows: int \| None = None,
1882	index_group: hisparse_runtime.HiSparseIndexGroup \| None = None,
1883	max_union_rows: int \| None = None,
1884	) -> HiSparseRuntime:
1885	runtime = HiSparseRuntime(
1886	config=ResolvedHiSparseConfig(
1887	top_k=top_k,
1888	device_buffer_size=device_buffer_size,
1889	max_union_rows=max_union_rows or device_buffer_size,
1890	),
1891	max_num_reqs=max_num_reqs,
1892	row_width=row_width,
1893	kv_dtype=torch.float32,
1894	device=DEVICE_TYPE,
1895	max_swap_rows=max_swap_rows,
1896	index_group=index_group,
1897	)
1898	blocks_per_request = cdiv(runtime.region_stride, block_size)
1899	# Leave one extra physical block for resident-tier tests. HMA group layouts
1900	# alias this backing, but distinct group allocations use distinct block IDs.
1901	num_blocks = max_num_reqs * blocks_per_request + 1
1902	raw = torch.zeros(
1903	num_blocks * block_size * row_width,
1904	dtype=torch.float32,
1905	device=DEVICE_TYPE,
1906	).view(torch.int8)
1907	block_table = torch.arange(
1908	max_num_reqs * blocks_per_request,
1909	dtype=torch.int32,
1910	device=DEVICE_TYPE,
1911	).view(max_num_reqs, blocks_per_request)
1912	runtime.bind_hot_cache(
1913	raw,
1914	byte_offset=0,
1915	block_stride=block_size * row_width * torch.float32.itemsize,
1916	num_blocks=num_blocks,
1917	block_size=block_size,
1918	block_table=block_table,
1919	)
1920	runtime.request_state_indices = torch.arange(
1921	max_num_reqs, dtype=torch.int32, device=runtime.device
1922	)
1923	return runtime
1924
1925
1926	def _make_hisparse_cache_handle(
1927	*,
1928	top_k: int = 4,
1929	device_buffer_size: int = 5,
1930	max_num_reqs: int = 2,
1931	row_width: int = 8,
1932	block_size: int = 64,
1933	max_swap_rows: int \| None = None,
1934	) -> HiSparseCacheHandle:
1935	runtime = _make_hisparse_runtime(
1936	top_k=top_k,
1937	device_buffer_size=device_buffer_size,
1938	max_num_reqs=max_num_reqs,
1939	row_width=row_width,
1940	block_size=block_size,
1941	max_swap_rows=max_swap_rows,
1942	)
1943	return HiSparseCacheHandle(runtime)
1944
1945
1946	def _valid_rows(num_rows: int) -> torch.Tensor:
1947	return torch.tensor([num_rows], dtype=torch.int32, device=DEVICE_TYPE)
1948
1949
1950	@requires_hisparse_ops
1951	def test_hisparse_uses_graph_stable_request_state_mapping():
1952	device = torch.device(DEVICE_TYPE)
1953	block_size, row_width = 64, 8
1954	runtime = _make_hisparse_runtime(
1955	top_k=1,
1956	device_buffer_size=2,
1957	max_num_reqs=2,
1958	row_width=row_width,
1959	block_size=block_size,
1960	)
1961	runtime.bind_source_cache(
1962	torch.arange(2 * block_size * row_width, dtype=torch.float32)
1963	.view(2, block_size, row_width)
1964	.pin_memory()
1965	)
1966	runtime.request_state_indices = torch.tensor([1], dtype=torch.int32, device=device)
1967
1968	cache_handle = HiSparseCacheHandle(runtime)
1969	cache_handle.all_context_pages_resident = False
1970	runtime.begin_forward()
1971	cache_handle.swap_in(
1972	req_id_per_token=torch.tensor([0], dtype=torch.int32, device=device),
1973	block_table=torch.tensor([[1]], dtype=torch.int32, device=device),
1974	logical_topk_indices=torch.tensor([[0]], dtype=torch.int32, device=device),
1975	block_size=block_size,
1976	num_valid_rows=_valid_rows(1),
1977	)
1978	torch.accelerator.synchronize()
1979
1980	assert (runtime.index_group.device_global_indices[0] == -1).all()
1981	assert (runtime.index_group.device_global_indices[1] == block_size).any()
1982
1983
1984	@requires_hisparse_ops
1985	def test_hisparse_maps_speculative_rows_through_request_state():
1986	"""Multiple verification rows for one request share its persistent state."""
1987	device = torch.device(DEVICE_TYPE)
1988	block_size, row_width = 64, 8
1989	runtime = _make_hisparse_runtime(
1990	top_k=1,
1991	device_buffer_size=2,
1992	max_num_reqs=1,
1993	row_width=row_width,
1994	block_size=block_size,
1995	max_swap_rows=4,
1996	)
1997	runtime.bind_source_cache(
1998	torch.zeros(1, block_size, row_width, dtype=torch.float32).pin_memory()
1999	)
2000
2001	cache_handle = HiSparseCacheHandle(runtime)
2002	resident_table = torch.tensor([[1]], dtype=torch.int32, device=device)
2003	cache_handle.bind_cache(
2004	runtime.hot.cache.view(torch.int8),
2005	byte_offset=0,
2006	block_stride=block_size * row_width * torch.float32.itemsize,
2007	num_blocks=2,
2008	block_size=block_size,
2009	block_table=resident_table,
2010	slot_mapping=torch.arange(4, dtype=torch.int64, device=device),
2011	)
2012	cache_handle.all_context_pages_resident = True
2013	runtime.begin_forward()
2014	physical = cache_handle.swap_in(
2015	req_id_per_token=torch.zeros(4, dtype=torch.int32, device=device),
2016	block_table=torch.zeros((1, 1), dtype=torch.int32, device=device),
2017	logical_topk_indices=torch.zeros((4, 1), dtype=torch.int32, device=device),
2018	block_size=block_size,
2019	num_valid_rows=_valid_rows(4),
2020	)
2021	torch.accelerator.synchronize()
2022
2023	assert physical.tolist() == [[block_size]] * 4
2024
2025
2026	@requires_hisparse_ops
2027	@pytest.mark.parametrize(
2028	"num_rows,row_stride,top_k",
2029	[(8, 128, 128), (4, 32, 128), (8, 2048, 2048)],
2030	ids=["disjoint", "overlap", "long_mtp"],
2031	)
2032	def test_hisparse_speculative_rows_resolve_host_misses_consistently(
2033	num_rows, row_stride, top_k
2034	):
2035	"""Verification rows of one request must agree on its hot-buffer slots.
2036
2037	Each row resolved residency in its own thread block against the request's
2038	shared slot ownership and LRU, so rows missing on different host rows could
2039	claim the same free slot. One row's hot index then pointed at a slot holding
2040	another row's KV. Overlapping rows must share one slot per host row, and
2041	long_mtp's rows span several resolution passes.
2042	"""
2043	device = torch.device(DEVICE_TYPE)
2044	block_size, row_width = 64, 64
2045	num_blocks = 2 * num_rows * top_k // block_size
2046
2047	kv_pool = torch.randn(
2048	(num_blocks + 1, block_size, row_width), dtype=torch.float32
2049	).pin_memory()
2050	flat_pool = kv_pool.reshape(-1, row_width)
2051	runtime = _make_hisparse_runtime(
2052	top_k=top_k,
2053	device_buffer_size=num_rows * top_k,
2054	max_num_reqs=1,
2055	row_width=row_width,
2056	block_size=block_size,
2057	max_swap_rows=num_rows,
2058	)
2059	runtime.bind_source_cache(kv_pool)
2060	cache = HiSparseCacheHandle(runtime)
2061	cache.all_context_pages_resident = False
2062
2063	block_table = torch.arange(
2064	1, num_blocks + 1, dtype=torch.int32, device=device
2065	).view(1, num_blocks)
2066	req_ids = torch.zeros(num_rows, dtype=torch.int32, device=device)
2067	rows = torch.arange(num_rows, dtype=torch.int32, device=device)[:, None]
2068	base = torch.arange(top_k, dtype=torch.int32, device=device)[None, :]
2069	for step in range(4):
2070	runtime.begin_forward()
2071	# Rows start row_stride apart (disjoint at row_stride == top_k) and
2072	# shift by half a top-k per step: all misses first, then a mix.
2073	topk = (rows * row_stride + base + step * (top_k // 2)) % (
2074	num_blocks * block_size
2075	)
2076	hot_indices = cache.swap_in(
2077	req_id_per_token=req_ids,
2078	block_table=block_table,
2079	logical_topk_indices=topk.clone(),
2080	block_size=block_size,
2081	num_valid_rows=_valid_rows(num_rows),
2082	)
2083	torch.accelerator.synchronize()
2084
2085	global_ref = _triton_convert_reference_impl(
2086	req_ids, block_table, topk, block_size, top_k
2087	)
2088	gathered = runtime.hot.attention_cache.reshape(-1, row_width)[
2089	hot_indices.to(torch.long)
2090	].cpu()
2091	torch.testing.assert_close(gathered, flat_pool[global_ref.cpu().to(torch.long)])
2092
2093
2094	@requires_hisparse_ops
2095	def test_hisparse_speculative_rows_count_a_shared_load_once():
2096	"""A host row two verification rows select loads once; the reuse is a hit."""
2097	device = torch.device(DEVICE_TYPE)
2098	block_size, row_width, top_k = 64, 64, 128
2099	num_blocks = top_k // block_size
2100	kv_pool = torch.randn(
2101	(num_blocks + 1, block_size, row_width), dtype=torch.float32
2102	).pin_memory()
2103	runtime = _make_hisparse_runtime(
2104	top_k=top_k,
2105	device_buffer_size=2 * top_k,
2106	max_num_reqs=1,
2107	row_width=row_width,
2108	block_size=block_size,
2109	max_swap_rows=2,
2110	)
2111	runtime.bind_source_cache(kv_pool)
2112	cache = HiSparseCacheHandle(runtime)
2113	cache.all_context_pages_resident = False
2114	block_table = torch.arange(
2115	1, num_blocks + 1, dtype=torch.int32, device=device
2116	).view(1, num_blocks)
2117	topk = torch.arange(top_k, dtype=torch.int32, device=device).repeat(2, 1)
2118
2119	runtime.begin_forward()
2120	cache.swap_in(
2121	req_id_per_token=torch.zeros(2, dtype=torch.int32, device=device),
2122	block_table=block_table,
2123	logical_topk_indices=topk,
2124	block_size=block_size,
2125	num_valid_rows=_valid_rows(2),
2126	)
2127	torch.accelerator.synchronize()
2128
2129	assert runtime.index_group.swap_stats.cpu().tolist() == [top_k, top_k]
2130
2131
2132	@requires_hisparse_ops
2133	def test_hisparse_resolver_reads_current_request_ids():
2134	"""The resolver must see this call's request ids, not the previous call's.
2135
2136	It runs on the copy stream after waiting only on ``logical_topk_ready``,
2137	which the indexer records before attention writes the ids. An MTP drafter
2138	alternates four rows per request with one, so stale ids resolve rows
2139	against other requests' KV. Stalling the compute stream exposes the race.
2140	"""
2141	device = torch.device(DEVICE_TYPE)
2142	block_size, row_width, top_k, blocks_per_req, num_reqs = 64, 16, 64, 4, 2
2143	kv_pool = torch.randn(
2144	(num_reqs * blocks_per_req + 1, block_size, row_width), dtype=torch.float32
2145	).pin_memory()
2146	runtime = _make_hisparse_runtime(
2147	top_k=top_k,
2148	device_buffer_size=4 * top_k,
2149	max_num_reqs=num_reqs,
2150	row_width=row_width,
2151	block_size=block_size,
2152	max_swap_rows=4 * num_reqs,
2153	)
2154	runtime.bind_source_cache(kv_pool)
2155	ready = torch.Event()
2156	runtime.index_group.logical_topk_ready = ready
2157	cache = HiSparseCacheHandle(runtime)
2158	cache.all_context_pages_resident = False
2159	cache.source_block_table = torch.arange(
2160	1, num_reqs * blocks_per_req + 1, dtype=torch.int32, device=device
2161	).view(num_reqs, blocks_per_req)
2162
2163	index_group = object.__new__(HiSparseMLAIndexGroup)
2164	index_group.caches = [cache]
2165	index_group.physical_topk_indices = torch.empty(
2166	(4 * num_reqs + 1, top_k), dtype=torch.int32, device=device
2167	)
2168
2169	# Four rows per request, then one per request plus a CUDA-graph padding
2170	# row, which maps to request 0 and must be skipped.
2171	for rows_per_req, num_padding in ((4, 0), (1, 1)):
2172	num_rows = num_reqs * rows_per_req
2173	req_ids = torch.arange(num_reqs, dtype=torch.int32, device=device)
2174	req_ids = torch.cat(
2175	[
2176	req_ids.repeat_interleave(rows_per_req),
2177	torch.zeros(num_padding, dtype=torch.int32, device=device),
2178	]
2179	)
2180	topk = torch.stack(
2181	[
2182	torch.randperm(blocks_per_req * block_size, device=device)[:top_k]
2183	for _ in range(req_ids.numel())
2184	]
2185	).to(torch.int32)
2186	query_start_loc = torch.arange(
2187	0, num_rows + 1, rows_per_req, dtype=torch.int32, device=device
2188	)
2189	attn_metadata = SimpleNamespace(
2190	req_id_per_token=req_ids,
2191	query_start_loc=torch.cat([query_start_loc, query_start_loc[-1:]]),
2192	block_size=block_size,
2193	)
2194	runtime.begin_forward()
2195	ready.record(current_stream())
2196	torch.cuda._sleep(200_000_000)
2197	hot_indices = index_group.convert_logical_to_physical_topk(
2198	0, topk, attn_metadata, block_stride_rows=None, return_valid_counts=False
2199	)
2200	torch.accelerator.synchronize()
2201
2202	assert (hot_indices[num_rows:] == -1).all()
2203	global_ref = _triton_convert_reference_impl(
2204	req_ids[:num_rows],
2205	cache.source_block_table,
2206	topk[:num_rows],
2207	block_size,
2208	top_k,
2209	)
2210	gathered = runtime.hot.attention_cache.reshape(-1, row_width)[
2211	hot_indices[:num_rows].to(torch.long)
2212	].cpu()
2213	torch.testing.assert_close(
2214	gathered, kv_pool.reshape(-1, row_width)[global_ref.cpu().to(torch.long)]
2215	)
2216
2217
2218	@requires_hisparse_ops
2219	def test_hisparse_union_table_scales_with_union_bound():
2220	"""Shared memory follows the per-request union bound, not the hot buffer.
2221
2222	Sizing the table from device_buffer_size needs ~340 KB here and fails.
2223	"""
2224	device = torch.device(DEVICE_TYPE)
2225	block_size, row_width, top_k = 64, 64, 128
2226	num_blocks = top_k // block_size
2227	kv_pool = torch.randn(
2228	(num_blocks + 1, block_size, row_width), dtype=torch.float32
2229	).pin_memory()
2230	runtime = _make_hisparse_runtime(
2231	top_k=top_k,
2232	device_buffer_size=32768,
2233	max_union_rows=top_k,
2234	max_num_reqs=1,
2235	row_width=row_width,
2236	block_size=block_size,
2237	max_swap_rows=1,
2238	)
2239	runtime.bind_source_cache(kv_pool)
2240	cache = HiSparseCacheHandle(runtime)
2241	cache.all_context_pages_resident = False
2242	req_ids = torch.zeros(1, dtype=torch.int32, device=device)
2243	block_table = torch.arange(
2244	1, num_blocks + 1, dtype=torch.int32, device=device
2245	).view(1, num_blocks)
2246	topk = torch.arange(top_k, dtype=torch.int32, device=device).view(1, top_k)
2247
2248	runtime.begin_forward()
2249	hot_indices = cache.swap_in(
2250	req_id_per_token=req_ids,
2251	block_table=block_table,
2252	logical_topk_indices=topk,
2253	block_size=block_size,
2254	num_valid_rows=_valid_rows(req_ids.numel()),
2255	)
2256	torch.accelerator.synchronize()
2257
2258	global_ref = _triton_convert_reference_impl(
2259	req_ids, block_table, topk, block_size, top_k
2260	)
2261	gathered = runtime.hot.attention_cache.reshape(-1, row_width)[
2262	hot_indices.to(torch.long)
2263	].cpu()
2264	torch.testing.assert_close(
2265	gathered, kv_pool.reshape(-1, row_width)[global_ref.cpu().to(torch.long)]
2266	)
2267
2268
2269	@requires_hisparse_ops
2270	def test_hisparse_rejects_union_bound_beyond_shared_memory():
2271	"""An oversized union bound fails at construction, not on the first forward."""
2272	with pytest.raises(ValueError, match="shared memory"):
2273	_make_hisparse_runtime(top_k=2048, device_buffer_size=32768)
2274
2275
2276	@requires_hisparse_ops
2277	def test_hisparse_resident_rows_bypass_hot_lru():
2278	device = torch.device(DEVICE_TYPE)
2279	block_size, row_width = 64, 8
2280	cache_handle = _make_hisparse_cache_handle(
2281	top_k=4,
2282	device_buffer_size=5,
2283	max_num_reqs=1,
2284	row_width=row_width,
2285	block_size=block_size,
2286	)
2287	raw = cache_handle.runtime.hot.cache.view(torch.int8)
2288	resident_table = torch.tensor([[1]], dtype=torch.int32, device=device)
2289	resident_slots = torch.tensor([block_size], dtype=torch.int64, device=device)
2290	cache_handle.bind_cache(
2291	raw,
2292	byte_offset=0,
2293	block_stride=block_size * row_width * torch.float32.itemsize,
2294	num_blocks=2,
2295	block_size=block_size,
2296	block_table=resident_table,
2297	slot_mapping=resident_slots,
2298	)
2299
2300	host = torch.randn(1, block_size, row_width).pin_memory()
2301	cache_handle.runtime.bind_source_cache(host)
2302	resident_rows = host[0].to(device).add(7)
2303	cache_handle.runtime.hot.cache[1].copy_(resident_rows)
2304	source_table = torch.tensor([[5]], dtype=torch.int32, device=device)
2305	request_ids = torch.zeros(1, dtype=torch.int32, device=device)
2306	topk = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32, device=device)
2307	lru_before = cache_handle.runtime.index_group.lru_slots.clone()
2308
2309	cache_handle.runtime.begin_forward()
2310	indices, valid_counts = cache_handle.swap_in(
2311	req_id_per_token=request_ids,
2312	block_table=source_table,
2313	logical_topk_indices=topk,
2314	block_size=block_size,
2315	num_valid_rows=_valid_rows(request_ids.numel()),
2316	return_valid_counts=True,
2317	)
2318	torch.accelerator.synchronize()
2319
2320	expected_indices = topk + block_size
2321	torch.testing.assert_close(indices, expected_indices)
2322	torch.testing.assert_close(
2323	valid_counts, torch.tensor([topk.shape[1]], dtype=torch.int32, device=device)
2324	)
2325	torch.testing.assert_close(cache_handle.runtime.index_group.lru_slots, lru_before)
2326	torch.testing.assert_close(
2327	cache_handle.runtime.hot.attention_cache.view(-1, row_width)[
2328	indices.to(torch.long)
2329	],
2330	resident_rows[topk.to(torch.long)],
2331	)
2332
2333
2334	@requires_hisparse_ops
2335	def test_hisparse_swap_in_preserves_rows_across_eviction():
2336	device = torch.device(DEVICE_TYPE)
2337	block_size = 64
2338	row_width = 64
2339	num_blocks = 64
2340	top_k = 128
2341	buf = 256
2342	num_reqs = 4
2343
2344	kv_pool = torch.randn(
2345	(num_blocks + 1, block_size, row_width), dtype=torch.float32
2346	).pin_memory()
2347	flat_pool = kv_pool.reshape(-1, row_width)
2348
2349	runtime = _make_hisparse_runtime(
2350	top_k=top_k,
2351	device_buffer_size=buf,
2352	max_num_reqs=num_reqs,
2353	row_width=row_width,
2354	)
2355	runtime.bind_source_cache(kv_pool)
2356	cache = HiSparseCacheHandle(runtime)
2357	cache.all_context_pages_resident = False
2358	cache.runtime.begin_forward()
2359
2360	blocks_per_req = num_blocks // num_reqs
2361	block_table = torch.arange(
2362	1, num_blocks + 1, dtype=torch.int32, device=device
2363	).view(num_reqs, blocks_per_req)
2364	req_ids = torch.arange(num_reqs, dtype=torch.int32, device=device)
2365	seq_len = blocks_per_req * block_size
2366	base = torch.arange(top_k, dtype=torch.int32, device=device)
2367	for step in range(3):
2368	cache.runtime.begin_forward()
2369	topk = torch.stack(
2370	[(base + step * top_k + row * 17) % seq_len for row in range(num_reqs)]
2371	)
2372	topk[:, -1] = -1
2373
2374	hot_indices, valid_counts = cache.swap_in(
2375	req_id_per_token=req_ids,
2376	block_table=block_table,
2377	logical_topk_indices=topk.clone(),
2378	block_size=block_size,
2379	num_valid_rows=_valid_rows(req_ids.numel()),
2380	return_valid_counts=True,
2381	)
2382	torch.accelerator.synchronize()
2383
2384	global_ref = _triton_convert_reference_impl(
2385	req_ids, block_table, topk, block_size, top_k
2386	)
2387	valid = hot_indices >= 0
2388	torch.testing.assert_close(
2389	valid_counts,
2390	(global_ref >= 0).sum(dim=1, dtype=torch.int32),
2391	rtol=0,
2392	atol=0,
2393	)
2394	gathered = cache.runtime.hot.attention_cache.reshape(-1, row_width)[
2395	hot_indices[valid].to(torch.long)
2396	].cpu()
2397	expected = flat_pool[global_ref[valid].cpu().to(torch.long)]
2398	torch.testing.assert_close(gathered, expected)
2399
2400
2401	@requires_hisparse_ops
2402	def test_hisparse_multi_step_swaps_match_independent():
2403	"""The shared LRU must retain every speculative step until followers run."""
2404	device = torch.device(DEVICE_TYPE)
2405	torch.manual_seed(0)
2406	block_size = 64
2407	row_width = 64
2408	num_blocks = 64
2409	top_k = 128
2410	buf = 3 * top_k
2411	num_reqs = 4
2412
2413	kv_pool = torch.randn(
2414	(num_blocks, block_size, row_width), dtype=torch.float32
2415	).pin_memory()
2416
2417	def make(
2418	index_group: hisparse_runtime.HiSparseIndexGroup \| None = None,
2419	) -> HiSparseCacheHandle:
2420	runtime = _make_hisparse_runtime(
2421	top_k=top_k,
2422	max_num_reqs=num_reqs,
2423	device_buffer_size=buf,
2424	row_width=row_width,
2425	block_size=block_size,
2426	max_swap_rows=2 * num_reqs,
2427	index_group=index_group,
2428	)
2429	runtime.bind_source_cache(kv_pool)
2430	cache = HiSparseCacheHandle(runtime)
2431	cache.all_context_pages_resident = False
2432	return cache
2433
2434	blocks_per_req = num_blocks // num_reqs
2435	block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).view(
2436	num_reqs, blocks_per_req
2437	)
2438	req_ids = torch.arange(num_reqs, dtype=torch.int32, device=device)
2439	seq_len = blocks_per_req * block_size
2440	producer = make()
2441	shared = make(producer.runtime.index_group)
2442	indep = make()
2443	base = torch.arange(top_k, dtype=torch.int32, device=device)
2444	for iteration in range(2):
2445	producer.runtime.begin_forward()
2446	shared.runtime.begin_forward()
2447	indep.runtime.begin_forward()
2448	topks = [
2449	torch.stack(
2450	[
2451	(base + (iteration * 2 + step) * top_k + row * 17) % seq_len
2452	for row in range(num_reqs)
2453	]
2454	)
2455	for step in range(2)
2456	]
2457	kw = dict(
2458	req_id_per_token=req_ids,
2459	block_table=block_table,
2460	block_size=block_size,
2461	num_valid_rows=_valid_rows(req_ids.numel()),
2462	)
2463	producer_indices = [
2464	producer.swap_in(logical_topk_indices=topk.clone(), **kw) for topk in topks
2465	]
2466	independent_indices = [
2467	indep.swap_in(logical_topk_indices=topk.clone(), **kw) for topk in topks
2468	]
2469	follower_indices = [
2470	shared.swap_in(logical_topk_indices=topk.clone(), **kw) for topk in topks
2471	]
2472	torch.accelerator.synchronize()
2473
2474	for idx_full, idx_indep, idx_shared in zip(
2475	producer_indices, independent_indices, follower_indices
2476	):
2477	torch.testing.assert_close(idx_shared, idx_full, rtol=0, atol=0)
2478	torch.testing.assert_close(idx_indep, idx_full, rtol=0, atol=0)
2479	torch.testing.assert_close(shared.runtime.hot.cache, indep.runtime.hot.cache)
2480
2481	shared_rows = shared.runtime.hot.attention_cache.reshape(-1, row_width)
2482	host_rows = kv_pool.reshape(-1, row_width)
2483	for physical_rows, logical_rows in zip(follower_indices, topks):
2484	global_rows = _triton_convert_reference_impl(
2485	req_ids, block_table, logical_rows, block_size, top_k
2486	)
2487	valid = physical_rows >= 0
2488	torch.testing.assert_close(
2489	shared_rows[physical_rows[valid].to(torch.long)].cpu(),
2490	host_rows[global_rows[valid].cpu().to(torch.long)],
2491	)
2492
2493
2494	@requires_hisparse_ops
2495	def test_hisparse_multi_step_writes_request_major_output():
2496	device = torch.device(DEVICE_TYPE)
2497	block_size = 64
2498	row_width = 8
2499	num_reqs = 2
2500	query_len = 3
2501	top_k = 4
2502	cache = _make_hisparse_cache_handle(
2503	top_k=top_k,
2504	device_buffer_size=query_len * top_k,
2505	max_num_reqs=num_reqs,
2506	row_width=row_width,
2507	block_size=block_size,
2508	max_swap_rows=num_reqs * query_len,
2509	)
2510	host = torch.arange(
2511	(num_reqs + 1) * block_size * row_width, dtype=torch.float32
2512	).view(num_reqs + 1, block_size, row_width)
2513	cache.runtime.bind_source_cache(host.pin_memory())
2514	block_table = torch.arange(1, num_reqs + 1, dtype=torch.int32, device=device).view(
2515	-1, 1
2516	)
2517	request_ids = torch.arange(num_reqs, dtype=torch.int32, device=device)
2518	logical = torch.tensor(
2519	[
2520	[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]],
2521	[[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]],
2522	],
2523	dtype=torch.int32,
2524	device=device,
2525	)
2526	physical = torch.full_like(logical, -1)
2527	valid_counts = torch.zeros((num_reqs, query_len), dtype=torch.int32, device=device)
2528
2529	cache.runtime.begin_forward()
2530	for step in range(query_len):
2531	cache.swap_in(
2532	request_ids,
2533	block_table,
2534	logical[:, step],
2535	block_size=block_size,
2536	num_valid_rows=_valid_rows(request_ids.numel()),
2537	return_valid_counts=True,
2538	attention_indices_out=physical[:, step],
2539	valid_counts_out=valid_counts[:, step],
2540	)
2541	torch.accelerator.synchronize()
2542
2543	torch.testing.assert_close(
2544	valid_counts,
2545	torch.full_like(valid_counts, top_k),
2546	)
2547	global_rows = _triton_convert_reference_impl(
2548	torch.arange(num_reqs, device=device, dtype=torch.int32).repeat_interleave(
2549	query_len
2550	),
2551	block_table,
2552	logical.view(num_reqs * query_len, top_k),
2553	block_size,
2554	top_k,
2555	)
2556	hot_rows = cache.runtime.hot.attention_cache.reshape(-1, row_width)
2557	host_rows = host.reshape(-1, row_width)
2558	torch.testing.assert_close(
2559	hot_rows[physical.view(num_reqs * query_len, top_k).to(torch.long)].cpu(),
2560	host_rows[global_rows.cpu().to(torch.long)],
2561	)
2562
2563
2564	@requires_hisparse_ops
2565	def test_hisparse_kv_update_writes_resident_and_staging_caches():
2566	device = torch.device(DEVICE_TYPE)
2567	block_size = 4
2568	row_width = 8
2569	cache_handle = _make_hisparse_cache_handle(
2570	max_num_reqs=1,
2571	row_width=row_width,
2572	block_size=block_size,
2573	)
2574	resident_slots = torch.tensor([4, 5, -1], dtype=torch.int64, device=device)
2575	cache_handle.bind_cache(
2576	cache_handle.runtime.hot.cache.view(torch.int8),
2577	byte_offset=0,
2578	block_stride=block_size * row_width * torch.float32.itemsize,
2579	num_blocks=cache_handle.runtime.hot.cache.shape[0],
2580	block_size=block_size,
2581	block_table=torch.tensor([[1]], dtype=torch.int32, device=device),
2582	slot_mapping=resident_slots,
2583	)
2584	cache_handle.mirror_staging_cache = torch.empty(
2585	(1, block_size, row_width), dtype=torch.float32, device=device
2586	)
2587	cache_handle.mirror_staging_slots = torch.arange(
2588	block_size, dtype=torch.int64, device=device
2589	)
2590	slots = torch.tensor([3, 7, -1], dtype=torch.int64, device=device)
2591	kv_c = torch.randn(8, row_width - 2, device=device)
2592	k_pe = torch.randn(8, 1, 2, device=device)
2593	cache_handle.num_actual_tokens = slots.numel()
2594	cache_handle.decode_batch = False
2595	cache_handle.host_mirror_required = True
2596	source_cache = torch.zeros_like(cache_handle.view.cache)
2597	impl = object.__new__(FlashMLASparseImpl)
2598	layer = SimpleNamespace(
2599	hisparse_cache=cache_handle,
2600	use_pcp=False,
2601	pcp_shard_decode_requests=False,
2602	impl=impl,
2603	)
2604	mla_attention.MLAAttention.update_kv_cache(
2605	layer,
2606	kv_c,
2607	k_pe,
2608	source_cache,
2609	slots,
2610	SimpleNamespace(num_decode_tokens=0),
2611	"auto",
2612	torch.tensor(1.0, device=device),
2613	)
2614	torch.accelerator.synchronize()
2615
2616	expected = torch.cat([kv_c[:2], k_pe[:2, 0]], dim=-1).cpu()
2617	torch.testing.assert_close(
2618	cache_handle.view.cache.view(-1, row_width)[resident_slots[:2]],
2619	expected.to(device),
2620	)
2621	torch.testing.assert_close(source_cache, torch.zeros_like(source_cache))
2622	staged = cache_handle.mirror_staging_cache.view(-1, row_width)
2623	staged_expected = torch.cat([kv_c[:3], k_pe[:3, 0]], dim=-1)
2624	torch.testing.assert_close(staged[:3], staged_expected)
2625
2626
2627	@requires_hisparse_ops
2628	def test_hisparse_row_dma_copies_discontiguous_spans_across_layers():
2629	device = torch.device(DEVICE_TYPE)
2630	block_size = 4
2631	row_width = 8
2632	num_rows = 16
2633	resident_caches = tuple(
2634	(
2635	torch.arange(num_rows * row_width, dtype=torch.uint8, device=device)
2636	.add_(layer_index * 17)
2637	.view(num_rows // block_size, block_size, row_width)
2638	)
2639	for layer_index in range(2)
2640	)
2641	host_caches = tuple(
2642	torch.zeros((num_rows, row_width), dtype=torch.uint8, pin_memory=True)
2643	for _ in resident_caches
2644	)
2645	worker = object.__new__(HiSparseConnectorWorker)
2646	worker.is_host_writer = True
2647	worker.kernel_block_size = block_size
2648	worker.resident_caches = resident_caches
2649	worker.host_caches = host_caches
2650	worker.cache_handles = [
2651	SimpleNamespace(
2652	decode_batch=True,
2653	runtime=SimpleNamespace(resident_source_index=0),
2654	)
2655	for _ in resident_caches
2656	]
2657	worker.dma_stream = torch.cuda.Stream(device=device)
2658	worker.host_write_event = torch.Event()
2659	worker._dma_free_descriptors = []
2660	worker._pending_dma_descriptors = deque()
2661	worker._pending_transfer_events = deque()
2662	worker._enqueued_transfer_ids = []
2663	worker._dma_submitted = False
2664	worker._set_row_mirrors(
2665	(
2666	SparseKVRowMirror((1,), 5, 2),
2667	SparseKVRowMirror((8,), 10, 4),
2668	)
2669	)
2670
2671	worker._enqueue_row_dma(range(2))
2672	current_stream().wait_event(worker.host_write_event)
2673	torch.accelerator.synchronize()
2674
2675	for resident, host in zip(resident_caches, host_caches):
2676	flat_resident = resident.view(num_rows, row_width).cpu()
2677	torch.testing.assert_close(host[5:7], flat_resident[1:3])
2678	torch.testing.assert_close(host[10:14], flat_resident[8:12])
2679
2680
2681	@requires_hisparse_ops
2682	def test_hisparse_remaps_strided_hma_rows_for_attention():
2683	device = torch.device(DEVICE_TYPE)
2684	block_size, row_width = 4, 8
2685	runtime = _make_hisparse_runtime(
2686	max_num_reqs=1,
2687	block_size=block_size,
2688	row_width=row_width,
2689	)
2690	page_elements = block_size * row_width
2691	stride_elements = 3 * page_elements
2692	byte_offset = page_elements * torch.float32.itemsize
2693	raw = torch.zeros(2 * stride_elements, dtype=torch.float32, device=device).view(
2694	torch.int8
2695	)
2696	runtime.bind_hot_cache(
2697	raw,
2698	byte_offset=byte_offset,
2699	block_stride=stride_elements * torch.float32.itemsize,
2700	num_blocks=2,
2701	block_size=block_size,
2702	block_table=torch.tensor([[0, 1]], dtype=torch.int32, device=device),
2703	)
2704
2705	source = (
2706	torch.arange(36 * row_width, dtype=torch.float32)
2707	.view(9, block_size, row_width)
2708	.pin_memory()
2709	)
2710	runtime.bind_source_cache(source)
2711	cache = HiSparseCacheHandle(runtime)
2712	cache.all_context_pages_resident = False
2713	runtime.begin_forward()
2714	attention_indices = cache.swap_in(
2715	req_id_per_token=torch.tensor([0], dtype=torch.int32, device=device),
2716	block_table=torch.arange(1, 9, dtype=torch.int32, device=device).view(1, 8),
2717	logical_topk_indices=torch.tensor(
2718	[[0, 1, 2, 3]], dtype=torch.int32, device=device
2719	),
2720	block_size=block_size,
2721	num_valid_rows=_valid_rows(1),
2722	)
2723
2724	assert runtime.hot.attention_cache is not None
2725	assert runtime.hot.attention_cache.is_contiguous()
2726	physical_indices = runtime.index_group.shared_topk.device_topk_rows[0]
2727	expected = (
2728	physical_indices // block_size * (stride_elements // row_width)
2729	+ physical_indices % block_size
2730	)
2731	torch.testing.assert_close(attention_indices[0], expected)
2732
2733
2734	@requires_hisparse_ops
2735	def test_hisparse_gather_zeroes_unaligned_destination():
2736	"""Invalid host rows are safely zeroed even at a byte-unaligned address."""
2737	device = torch.device(DEVICE_TYPE)
2738	row_bytes = 16
2739	host_cache = torch.ones(1, row_bytes, dtype=torch.uint8).pin_memory()
2740	storage = torch.full((2 * row_bytes + 1,), 7, dtype=torch.uint8, device=device)
2741	hot_cache = torch.as_strided(
2742	storage,
2743	size=(2, 1, row_bytes),
2744	stride=(row_bytes + 1, row_bytes, 1),
2745	)
2746	global_indices = torch.tensor([[1]], dtype=torch.int32, device=device)
2747	hot_indices = torch.tensor([[1]], dtype=torch.int32, device=device)
2748	miss_mask = torch.ones_like(global_indices)
2749
2750	torch.ops._C_cache_ops.hisparse_gather_plan(
2751	host_cache,
2752	hot_cache,
2753	global_indices,
2754	hot_indices,
2755	miss_mask,
2756	None,
2757	)
2758	torch.accelerator.synchronize()
2759
2760	torch.testing.assert_close(hot_cache[1, 0], torch.zeros_like(hot_cache[1, 0]))
2761	assert storage[row_bytes].item() == 7
2762
2763
2764	@requires_hisparse_ops
2765	def test_hisparse_gather_rejects_short_attention_block_stride():
2766	"""Reject a stride that would let the gather kernel cross block bounds."""
2767	device = torch.device(DEVICE_TYPE)
2768	host_cache = torch.ones(1, 16, dtype=torch.uint8).pin_memory()
2769	hot_cache = torch.empty((1, 2, 16), dtype=torch.uint8, device=device)
2770	global_indices = torch.zeros((1, 1), dtype=torch.int32, device=device)
2771	hot_indices = torch.zeros_like(global_indices)
2772	miss_mask = torch.ones_like(global_indices)
2773	attention_indices = torch.empty_like(global_indices)
2774
2775	with pytest.raises(RuntimeError, match="stride must cover one hot block"):
2776	torch.ops._C_cache_ops.hisparse_gather_plan(
2777	host_cache,
2778	hot_cache,
2779	global_indices,
2780	hot_indices,
2781	miss_mask,
2782	None,
2783	attention_indices,
2784	1,
2785	)
2786
2787
2788	@requires_hisparse_ops
2789	def test_hisparse_newest_write_and_recycled_slot_invalidation():
2790	"""Newest writes clamp padding and invalidated rows are loaded again."""
2791	device = torch.device(DEVICE_TYPE)
2792	block_size = 4
2793	row_width = 8
2794	num_blocks = 8
2795
2796	kv_pool = (
2797	torch.arange(num_blocks * block_size * row_width, dtype=torch.float32)
2798	.view(num_blocks, block_size, row_width)
2799	.pin_memory()
2800	)
2801	flat_pool = kv_pool.reshape(-1, row_width)
2802
2803	cache_handle = _make_hisparse_cache_handle(block_size=block_size)
2804	cache_handle.runtime.eager_host_mirror = True
2805
2806	block_table = torch.tensor([[2, 0, 4]], dtype=torch.int32, device=device)
2807	resident_slots = torch.tensor([4, -1], dtype=torch.int64, device=device)
2808	cache_handle.bind_cache(
2809	cache_handle.runtime.hot.cache.view(torch.int8),
2810	byte_offset=0,
2811	block_stride=block_size * row_width * torch.float32.itemsize,
2812	num_blocks=cache_handle.runtime.hot.cache.shape[0],
2813	block_size=block_size,
2814	block_table=torch.tensor([[0, 0, 1]], dtype=torch.int32, device=device),
2815	slot_mapping=resident_slots,
2816	)
2817	req_ids = torch.tensor([0], dtype=torch.int32, device=device)
2818	newest_global = 4 * block_size
2819	padded_global = 1
2820	padded_row = flat_pool[padded_global].clone()
2821	slot_mapping = torch.tensor(
2822	[newest_global, padded_global], dtype=torch.int64, device=device
2823	)
2824	request_state_indices = torch.tensor([0, -1], dtype=torch.int32, device=device)
2825	cache_handle.runtime.request_state_indices = request_state_indices
2826
2827	kv_c = torch.randn(3, row_width - 2, device=device)
2828	k_pe = torch.randn(3, 1, 2, device=device)
2829	cache_handle.runtime.bind_source_cache(kv_pool)
2830	cache_handle.num_actual_tokens = slot_mapping.numel()
2831	ops.concat_and_cache_mla(
2832	kv_c[:1],
2833	k_pe[:1, 0],
2834	cache_handle.view.cache,
2835	resident_slots[:1],
2836	kv_cache_dtype="auto",
2837	scale=torch.tensor(1.0, device=device),
2838	)
2839	torch.accelerator.synchronize()
2840	expected_row = torch.cat([kv_c[0], k_pe[0, 0]]).cpu()
2841	flat_pool[newest_global].copy_(expected_row)
2842	torch.testing.assert_close(flat_pool[newest_global], expected_row)
2843	torch.testing.assert_close(flat_pool[padded_global], padded_row)
2844	flat_hot = cache_handle.runtime.hot.cache.view(-1, row_width)
2845	torch.testing.assert_close(flat_hot[resident_slots[0]].cpu(), expected_row)
2846
2847	topk = torch.tensor([[0, -1, -1, -1]], dtype=torch.int32, device=device)
2848	cache_handle.all_context_pages_resident = False
2849	cache_handle.runtime.begin_forward()
2850	hot_indices = cache_handle.swap_in(
2851	req_id_per_token=req_ids,
2852	block_table=block_table,
2853	logical_topk_indices=topk,
2854	block_size=block_size,
2855	num_valid_rows=_valid_rows(req_ids.numel()),
2856	)
2857	torch.accelerator.synchronize()
2858	stale_hot_slot = hot_indices[0, 0].item()
2859	stale_row = flat_hot[stale_hot_slot].clone()
2860
2861	flat_pool[8] += 1000
2862	cache_handle.runtime.invalidate_written_slots(
2863	torch.tensor([8], dtype=torch.int64, device=device), req_ids
2864	)
2865	cache_handle.runtime.begin_forward()
2866	hot_indices = cache_handle.swap_in(
2867	req_id_per_token=req_ids,
2868	block_table=block_table,
2869	logical_topk_indices=topk,
2870	block_size=block_size,
2871	num_valid_rows=_valid_rows(req_ids.numel()),
2872	)
2873	torch.accelerator.synchronize()
2874	idx = hot_indices.cpu().tolist()[0][0]
2875	assert not torch.equal(flat_hot[idx], stale_row)
2876	torch.testing.assert_close(flat_hot[idx].cpu(), flat_pool[8])
2877
2878
2879	@requires_hisparse_ops
2880	def test_hisparse_mixed_batch_bf16_row_split(
2881	default_vllm_config, dist_init, workspace_init
2882	):
2883	"""Host-resident mixed batch on the bf16 path is row-split.
2884
2885	Two long-context speculative-decode requests + one short local-prefill
2886	chunk: every decode step must be served from the bounded hot buffer before
2887	it is reused, while only the prefill rows' blocks are staged host->GPU.
2888	"""
2889	ok, reason = flashmla.is_flashmla_sparse_supported()
2890	if not ok:
2891	pytest.skip(reason)
2892
2893	device = torch.device(DEVICE_TYPE)
2894	dtype = torch.bfloat16
2895	torch.manual_seed(0)
2896
2897	num_heads = 64
2898	kv_lora_rank = 512
2899	qk_nope_head_dim = 128
2900	qk_rope_head_dim = 64
2901	v_head_dim = 128
2902	head_size = kv_lora_rank + qk_rope_head_dim
2903	topk_tokens = 128
2904	block_size = 64
2905
2906	# Long decode contexts + a short prefill chunk (router shortcut shape).
2907	batch_spec = BatchSpec(seq_lens=[2048, 2048, 192], query_lens=[2, 2, 64])
2908	max_seqlen = max(batch_spec.seq_lens)
2909	total_cache_tokens = sum(batch_spec.seq_lens)
2910	total_tokens = batch_spec.compute_num_tokens()
2911
2912	vllm_config = create_vllm_config(
2913	model_name="deepseek-ai/DeepSeek-V2-Lite-Chat",
2914	tensor_parallel_size=1,
2915	max_model_len=max_seqlen,
2916	num_gpu_blocks=max(2048, cdiv(total_cache_tokens, block_size) + 1),
2917	block_size=block_size,
2918	hf_config_override={
2919	"index_topk": topk_tokens,
2920	"attn_module_list_cfg": [{"topk_tokens": topk_tokens}],
2921	},
2922	)
2923	vllm_config.attention_config.hisparse_config = HiSparseConfig(
2924	device_buffer_size=2 * topk_tokens,
2925	)
2926	vllm_config.speculative_config = SpeculativeConfig(
2927	method="ngram", num_speculative_tokens=1
2928	)
2929	model_config = vllm_config.model_config
2930	model_config.hf_text_config = SimpleNamespace(
2931	index_topk=topk_tokens,
2932	q_lora_rank=None,
2933	kv_lora_rank=kv_lora_rank,
2934	qk_nope_head_dim=qk_nope_head_dim,
2935	qk_rope_head_dim=qk_rope_head_dim,
2936	v_head_dim=v_head_dim,
2937	model_type="deepseek_v2",
2938	)
2939	model_config.dtype = dtype
2940	model_config.get_num_attention_heads = MethodType(
2941	lambda self, parallel_config: num_heads, model_config
2942	)
2943	model_config.get_num_kv_heads = MethodType(
2944	lambda self, parallel_config: 1, model_config
2945	)
2946	model_config.get_head_size = MethodType(lambda self: head_size, model_config)
2947	model_config.get_sliding_window = MethodType(lambda self: None, model_config)
2948	# create_hisparse_cache_handle sizes hot buffers per layer.
2949	model_config.get_num_layers = MethodType(
2950	lambda self, parallel_config: 1, model_config
2951	)
2952
2953	kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
2954	common_attn_metadata = create_common_attn_metadata(
2955	batch_spec,
2956	vllm_config.cache_config.block_size,
2957	device,
2958	arange_block_indices=True,
2959	)
2960	# Prepopulate every position of every sequence so the forward needs no
2961	# KV-cache update (the row split itself is what is under test).
2962	kv_c_contexts = [
2963	torch.rand(s_len, kv_lora_rank, dtype=dtype, device=device)
2964	for s_len in batch_spec.seq_lens
2965	]
2966	k_pe_contexts = [
2967	torch.rand(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device)
2968	for s_len in batch_spec.seq_lens
2969	]
2970	kv_cache = create_and_prepopulate_kv_cache(
2971	kv_c_contexts=kv_c_contexts,
2972	k_pe_contexts=k_pe_contexts,
2973	block_size=vllm_config.cache_config.block_size,
2974	head_size=head_size,
2975	dtype=dtype,
2976	device=device,
2977	num_blocks=vllm_config.cache_config.num_gpu_blocks,
2978	common_attn_metadata=common_attn_metadata,
2979	randomize_blocks=False,
2980	kv_cache_dtype="auto",
2981	)
2982
2983	prefill_backend = get_mla_prefill_backend(vllm_config)(
2984	num_heads=num_heads,
2985	scale=1.0 / math.sqrt(head_size),
2986	kv_lora_rank=kv_lora_rank,
2987	qk_nope_head_dim=qk_nope_head_dim,
2988	qk_rope_head_dim=qk_rope_head_dim,
2989	v_head_dim=v_head_dim,
2990	vllm_config=vllm_config,
2991	)
2992	vllm_config.compilation_config.static_forward_context["placeholder"] = (
2993	SimpleNamespace(prefill_backend=prefill_backend)
2994	)
2995
2996	builder_cls = FlashMLASparseBackend.get_builder_cls()
2997	builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device)
2998	metadata = builder.build(
2999	common_prefix_len=0, common_attn_metadata=common_attn_metadata
3000	)
3001	assert isinstance(metadata.prefill, SparseMLAPrefillMetadata)
3002
3003	# Per-token sparse indices bounded by each token's position, with unique
3004	# offsets and -1 padding (same construction as the decode parity test).
3005	positions: list[int] = []
3006	for i in range(batch_spec.batch_size):
3007	ctx_len = batch_spec.seq_lens[i] - batch_spec.query_lens[i]
3008	positions.extend(ctx_len + q_idx for q_idx in range(batch_spec.query_lens[i]))
3009	sparse_indices = torch.empty(
3010	total_tokens, topk_tokens, dtype=torch.int32, device=device
3011	)
3012	for tok_idx in range(total_tokens):
3013	max_valid_idx = positions[tok_idx]
3014	offset = tok_idx * 7 # Prime number for varied offsets
3015	num_valid = min(topk_tokens // 2, max_valid_idx + 1)
3016	valid_range = torch.arange(num_valid, device=device, dtype=torch.int32)
3017	sparse_indices[tok_idx] = torch.cat(
3018	[
3019	(valid_range + offset) % (max_valid_idx + 1),
3020	torch.full(
3021	(topk_tokens - num_valid,), -1, device=device, dtype=torch.int32
3022	),
3023	]
3024	)
3025
3026	q = torch.rand(total_tokens, num_heads, head_size, dtype=dtype, device=device)
3027	mock_indexer = SimpleNamespace(
3028	topk_indices_buffer=sparse_indices, topk_tokens=topk_tokens
3029	)
3030
3031	impl_cls = FlashMLASparseBackend.get_impl_cls()
3032	with set_current_vllm_config(vllm_config):
3033	impl = impl_cls(
3034	num_heads=num_heads,
3035	head_size=head_size,
3036	scale=1.0 / math.sqrt(head_size),
3037	num_kv_heads=1,
3038	alibi_slopes=None,
3039	sliding_window=None,
3040	kv_cache_dtype="auto",
3041	logits_soft_cap=None,
3042	attn_type="decoder",
3043	kv_sharing_target_layer_name=None,
3044	q_lora_rank=None,
3045	kv_lora_rank=kv_lora_rank,
3046	qk_nope_head_dim=qk_nope_head_dim,
3047	qk_rope_head_dim=qk_rope_head_dim,
3048	qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,
3049	v_head_dim=v_head_dim,
3050	kv_b_proj=None,
3051	indexer=mock_indexer,
3052	)
3053	assert isinstance(impl.index_group, HiSparseMLAIndexGroup)
3054	cache_handle = impl.index_group.cache(impl.index_group_index)
3055	blocks_per_request = cdiv(cache_handle.runtime.region_stride, block_size)
3056	num_hot_blocks = cache_handle.runtime.max_num_reqs * blocks_per_request
3057	hot_cache = torch.zeros(
3058	num_hot_blocks * block_size * head_size,
3059	dtype=dtype,
3060	device=device,
3061	).view(torch.int8)
3062	cache_handle.runtime.bind_hot_cache(
3063	hot_cache,
3064	byte_offset=0,
3065	block_stride=block_size * head_size * dtype.itemsize,
3066	num_blocks=num_hot_blocks,
3067	block_size=block_size,
3068	block_table=torch.arange(num_hot_blocks, dtype=torch.int32, device=device).view(
3069	cache_handle.runtime.max_num_reqs, blocks_per_request
3070	),
3071	)
3072	cache_handle.runtime.request_state_indices = torch.arange(
3073	cache_handle.runtime.max_num_reqs, dtype=torch.int32, device=device
3074	)
3075	impl.prepare_for_batch(metadata)
3076
3077	# Device-resident reference: the whole batch converted against the full
3078	# block table and run as one kernel call over the GPU cache.
3079	ref_topk, ref_topk_length = triton_convert_req_index_to_global_index(
3080	metadata.req_id_per_token,
3081	metadata.block_table,
3082	sparse_indices,
3083	BLOCK_SIZE=metadata.block_size,
3084	NUM_TOPK_TOKENS=topk_tokens,
3085	return_valid_counts=True,
3086	)
3087	reference, _ = impl._bf16_flash_mla_kernel(q, kv_cache, ref_topk, ref_topk_length)
3088
3089	# Host-resident pool with identical contents.
3090	kv_pool = kv_cache.squeeze(1).cpu().pin_memory()
3091	cache_handle.runtime.bind_source_cache(kv_pool)
3092	cache_handle.source_block_table = metadata.block_table
3093
3094	staging_calls = []
3095	original_gather = cache_handle.runtime.gather_prefill_cache
3096
3097	def spy_gather(self, kv, plan, resident_cache=None):
3098	staged = original_gather(kv, plan, resident_cache)
3099	staging_calls.append((plan, staged.shape))
3100	return staged
3101
3102	cache_handle.runtime.gather_prefill_cache = MethodType(
3103	spy_gather, cache_handle.runtime
3104	)
3105
3106	backend_output, _ = impl._forward_bf16_kv(
3107	q, kv_pool, sparse_indices, metadata, q.shape[1]
3108	)
3109	torch.accelerator.synchronize()
3110
3111	# Only the prefill rows' blocks were staged: the decode rows' 2048-token
3112	# contexts (32 blocks each) must stay off the staging gather.
3113	assert len(staging_calls) == 1
3114	plan, staged_shape = staging_calls[0]
3115	assert plan is metadata.prefill.host_staging_plan
3116	prefill_blocks = cdiv(batch_spec.seq_lens[-1], block_size)
3117	assert staged_shape[0] <= prefill_blocks + 1 # +1: block-0 tail padding
3118
3119	torch.testing.assert_close(backend_output, reference, rtol=0.01, atol=0.01)
3120
3121
3122	def test_hisparse_prefill_staging_remap():
3123	"""Compacted staging references the same host rows as direct indexing."""
3124	block_size = 4
3125	block_table = torch.tensor(
3126	[[5, 2, -1, -1], [2, 7, 3, -1], [9, -1, -1, -1]], dtype=torch.int32
3127	)
3128
3129	new_bt, row_ids = hisparse_prefill_staging_remap(block_table, block_size)
3130
3131	assert new_bt.shape == block_table.shape
3132	assert row_ids.dtype == torch.int32 and new_bt.dtype == torch.int32
3133	n_unique = len({0, 2, 3, 5, 7, 9}) # -1 clamps to block 0
3134	flat_rows = row_ids.flatten()
3135	assert (flat_rows >= 0).sum() == n_unique * block_size
3136	for i in range(block_table.shape[0]):
3137	for j in range(block_table.shape[1]):
3138	orig = max(int(block_table[i, j]), 0)
3139	staged = int(new_bt[i, j])
3140	for k in range(block_size):
3141	assert int(flat_rows[staged * block_size + k]) == orig * block_size + k
3142
3143
3144	def test_hisparse_prefill_staging_plan_masks_unused_blocks():
3145	block_table = torch.tensor([[5, 2, 7], [9, 3, 4]], dtype=torch.int32)
3146	plan = build_hisparse_prefill_staging_plan(
3147	block_table,
3148	seq_lens=torch.tensor([5, 8], dtype=torch.int32),
3149	block_size=4,
3150	staging_block_capacity=4,
3151	)
3152
3153	assert (plan.block_table[:, 2] == 0).all()
3154	valid_rows = plan.row_ids[plan.row_ids >= 0]
3155	assert set((valid_rows[::4] // 4).tolist()) == {0, 2, 3, 5, 9}
3156	for row, used in enumerate((2, 2)):
3157	for column in range(used):
3158	staged_block = int(plan.block_table[row, column])
3159	host_block = int(block_table[row, column])
3160	assert int(valid_rows[staged_block * 4]) == host_block * 4
3161
3162
3163	def test_hisparse_prefill_staging_plan_resolves_resident_sources():
3164	"""Per-page resident hits become device sources; misses stay host DMAs."""
3165	block_size = 4
3166	resident_block_size = 2
3167	block_table = torch.tensor([[5, 2, 0], [9, 3, 0]], dtype=torch.int32)
3168	seq_lens = torch.tensor([5, 8], dtype=torch.int32)
3169	plan = build_hisparse_prefill_staging_plan(
3170	block_table, seq_lens, block_size, staging_block_capacity=4
3171	)
3172	# Two resident pages per host block; 0 entries are null (not resident).
3173	resident_table = torch.tensor(
3174	[[11, 12, 0, 13, 0, 0], [21, 0, 22, 23, 0, 0]], dtype=torch.int32
3175	)
3176
3177	plan.ensure_gpu_sources(resident_table, resident_block_size)
3178
3179	assert plan.gpu_row_ids is not None
3180	unique_hosts = (plan.row_ids[0].view(-1, block_size)[:, 0] // block_size).tolist()
3181	gpu_rows = plan.gpu_row_ids[0].view(-1, block_size)
3182	reps = {5: (0, 0), 2: (0, 1), 9: (1, 0), 3: (1, 1)}
3183	for u, host_id in enumerate(unique_hosts):
3184	for t in range(block_size):
3185	if host_id < 0:
3186	assert int(gpu_rows[u, t]) == -1
3187	continue
3188	if host_id == 0:
3189	assert int(gpu_rows[u, t]) == -1
3190	continue
3191	row, col = reps[host_id]
3192	page = t // resident_block_size
3193	res_block = int(resident_table[row, col * 2 + page])
3194	expected = (
3195	res_block * resident_block_size + t % resident_block_size
3196	if res_block > 0
3197	else -1
3198	)
3199	assert int(gpu_rows[u, t]) == expected
3200	valid_rows = plan.row_ids >= 0
3201	torch.testing.assert_close(
3202	plan.miss_mask,
3203	((plan.gpu_row_ids < 0) & valid_rows).int(),
3204	check_dtype=False,
3205	)
3206
3207	plan.ensure_gpu_sources(torch.zeros_like(resident_table), resident_block_size)
3208
3209	assert plan.gpu_row_ids is not None
3210	assert (plan.gpu_row_ids == -1).all()
3211
3212
3213	@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
3214	def test_hisparse_gather_prefill_cache_prefers_resident_rows():
3215	"""Staged rows come from the resident cache when a shadow page exists."""
3216	if not _has_hisparse_ops():
3217	pytest.skip("hisparse CUDA ops unavailable")
3218	device = torch.device("cuda")
3219	block_size, resident_block_size, row_width = 4, 2, 16
3220	block_table = torch.tensor([[5, 2, 0], [9, 3, 0]], dtype=torch.int32, device=device)
3221	seq_lens = torch.tensor([5, 8], dtype=torch.int32, device=device)
3222	plan = build_hisparse_prefill_staging_plan(
3223	block_table, seq_lens, block_size, staging_block_capacity=4
3224	)
3225	resident_table = torch.tensor(
3226	[[11, 12, 0, 13, 0, 0], [21, 0, 22, 23, 0, 0]],
3227	dtype=torch.int32,
3228	device=device,
3229	)
3230	plan.ensure_gpu_sources(resident_table, resident_block_size)
3231
3232	num_host_blocks, num_res_blocks = 10, 24
3233	host_cache = (
3234	torch.arange(num_host_blocks * block_size * row_width, dtype=torch.float32)
3235	.view(num_host_blocks, block_size, row_width)
3236	.pin_memory()
3237	)
3238	resident_cache = (
3239	(
3240	-torch.arange(
3241	num_res_blocks * resident_block_size * row_width, dtype=torch.float32
3242	)
3243	- 1.0
3244	)
3245	.view(num_res_blocks, resident_block_size, row_width)
3246	.to(device)
3247	)
3248
3249	staged = HiSparseRuntime.gather_prefill_cache(
3250	None, host_cache, plan, resident_cache=resident_cache
3251	)
3252
3253	staged_flat = staged.view(-1, row_width).cpu()
3254	host_flat = host_cache.view(-1, row_width)
3255	resident_flat = resident_cache.reshape(-1, row_width).cpu()
3256	gpu_rows = plan.gpu_row_ids[0].cpu()
3257	host_rows = plan.row_ids[0].cpu()
3258	for i in range(staged_flat.shape[0]):
3259	if int(host_rows[i]) < 0:
3260	assert int(plan.miss_mask[0, i]) == 0
3261	continue
3262	expected = (
3263	resident_flat[int(gpu_rows[i])]
3264	if int(gpu_rows[i]) >= 0
3265	else host_flat[int(host_rows[i])]
3266	)
3267	torch.testing.assert_close(staged_flat[i], expected)
3268
3269
3270	def test_hisparse_fp8_decode_resolves_rows_once_then_runs_batched_attention():
3271	"""Verification rows resolve in one swap_in; FP8 attention sees them batched."""
3272	device = torch.device("cpu")
3273	num_decodes = 2
3274	query_len = 3
3275	num_tokens = num_decodes * query_len
3276	q = torch.randn(num_tokens, 2, 4, device=device)
3277	topk = (
3278	torch.arange(num_tokens, dtype=torch.int32, device=device)
3279	.view(-1, 1)
3280	.expand(-1, 4)
3281	)
3282	kernel_shapes: list[torch.Size] = []
3283
3284	def swap_in(req_ids, , logical_topk_indices, *kwargs): # noqa: ARG001
3285	return logical_topk_indices + 10
3286
3287	def run_kernel(self, , q, *kwargs): # noqa: ARG001
3288	kernel_shapes.append(q.shape)
3289	return q[..., :1], None
3290
3291	cache = SimpleNamespace(
3292	runtime=SimpleNamespace(
3293	hot=SimpleNamespace(attention_cache=torch.empty(1, device=device))
3294	),
3295	source_block_table=torch.empty(
3296	num_decodes, 1, dtype=torch.int32, device=device
3297	),
3298	swap_in=MagicMock(side_effect=swap_in),
3299	)
3300	index_group = object.__new__(HiSparseMLAIndexGroup)
3301	index_group.caches = [cache]
3302	index_group.physical_topk_indices = torch.empty(
3303	(num_tokens + 1, 4), dtype=torch.int32, device=device
3304	)
3305	impl = SimpleNamespace(
3306	kv_lora_rank=1,
3307	index_group=index_group,
3308	index_group_index=0,
3309	)
3310	impl._fp8_flash_mla_kernel = MethodType(run_kernel, impl)
3311	metadata = SimpleNamespace(
3312	query_start_loc=torch.arange(
3313	0, num_tokens + 1, query_len, dtype=torch.int32, device=device
3314	),
3315	req_id_per_token=torch.arange(
3316	num_decodes, dtype=torch.int32, device=device
3317	).repeat_interleave(query_len),
3318	block_size=64,
3319	)
3320
3321	output = FlashMLASparseImpl._host_backed_fp8_decode(
3322	impl, q, topk, metadata, SimpleNamespace(), num_decodes
3323	)
3324
3325	cache.swap_in.assert_called_once()
3326	torch.testing.assert_close(
3327	cache.swap_in.call_args.args[0], metadata.req_id_per_token
3328	)
3329	torch.testing.assert_close(
3330	cache.swap_in.call_args.kwargs["logical_topk_indices"], topk
3331	)
3332	assert kernel_shapes == [(num_decodes, query_len, 2, 4)]
3333	assert output.shape == (num_tokens, 2, 1)
3334
3335
3336	def test_hisparse_keeps_speculative_rows_as_decodes():
3337	"""HiSparse limits decodes to verification requests, not short extends.
3338
3339	FlashAttn sparse MLA used to classify MTP verification as prefill under
3340	HiSparse, whose residency-dependent route a FULL CUDA graph freezes at
3341	capture.
3342	"""
3343	builder = object.__new__(FlashAttnMLASparseMetadataBuilder)
3344	builder.vllm_config = SimpleNamespace(
3345	attention_config=SimpleNamespace(hisparse_config=object()),
3346	speculative_config=SimpleNamespace(
3347	num_speculative_tokens=3, parallel_drafting=False
3348	),
3349	parallel_config=SimpleNamespace(decode_context_parallel_size=1),
3350	use_v2_model_runner=True,
3351	)
3352
3353	builder._init_reorder_batch_threshold(256, supports_spec_as_decode=True)
3354
3355	assert builder.reorder_batch_threshold == 4
3356
3357
3358	def test_flashinfer_hisparse_decode_runs_batched_attention():
3359	device = torch.device("cpu")
3360	num_tokens = 6
3361	q = torch.randn(num_tokens, 2, 4, device=device)
3362	topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=device)
3363	physical_topk = topk + 10
3364	valid_counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=device)
3365	kernel_shapes: list[torch.Size] = []
3366
3367	def convert_decode(self, args, *kwargs): # noqa: ARG001
3368	return physical_topk, valid_counts
3369
3370	def prepare_kernel(self, args, *kwargs): # noqa: ARG001
3371	pass
3372
3373	def run_kernel(self, q, cache, indices, counts): # noqa: ARG001
3374	kernel_shapes.append(q.shape)
3375	return q[..., :1], None
3376
3377	cache_handle = SimpleNamespace(
3378	runtime=SimpleNamespace(
3379	hot=SimpleNamespace(attention_cache=torch.empty(1, device=device))
3380	),
3381	)
3382	index_group = object.__new__(HiSparseMLAIndexGroup)
3383	index_group.caches = [cache_handle]
3384	index_group.convert_logical_to_physical_topk = MethodType(
3385	convert_decode, index_group
3386	)
3387	impl = object.__new__(FlashInferMLASparseImpl)
3388	impl.topk_indices_buffer = topk
3389	impl.index_group = index_group
3390	impl.index_group_index = 0
3391	impl._prepare_mqa_kernel = MethodType(prepare_kernel, impl)
3392	impl._run_mqa_kernel = MethodType(run_kernel, impl)
3393	metadata = SimpleNamespace(num_decode_tokens=num_tokens)
3394
3395	output, lse = FlashInferMLASparseImpl.forward_mqa(
3396	impl,
3397	q,
3398	torch.empty(1, device=device),
3399	metadata,
3400	SimpleNamespace(),
3401	)
3402
3403	assert kernel_shapes == [q.shape]
3404	assert output.shape == (num_tokens, 2, 1)
3405	assert lse is None
3406
3407
3408	def test_flashattn_hisparse_decode_uses_index_group():
3409	num_tokens = 4
3410	q_nope = torch.empty(num_tokens, 2, 3, device=DEVICE_TYPE)
3411	q_rope = torch.empty(num_tokens, 2, 1, device=DEVICE_TYPE)
3412	topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=DEVICE_TYPE)
3413	physical = topk.clone()
3414	counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE)
3415	index_group = object.__new__(HiSparseMLAIndexGroup)
3416	index_group.convert_logical_to_physical_topk = MagicMock(
3417	return_value=(physical, counts)
3418	)
3419	index_group.physical_kv_cache = MagicMock(
3420	return_value=torch.empty(1, device=DEVICE_TYPE)
3421	)
3422	impl = object.__new__(FlashAttnMLASparseImpl)
3423	impl.topk_indices_buffer = topk
3424	impl.index_group = index_group
3425	impl.index_group_index = 0
3426	impl._run_mqa_kernel = MagicMock(return_value=q_nope[..., :1])
3427	metadata = SimpleNamespace(num_decode_tokens=num_tokens, block_size=64)
3428
3429	output, lse = FlashAttnMLASparseImpl.forward_mqa(
3430	impl,
3431	(q_nope, q_rope),
3432	torch.empty(1, device=DEVICE_TYPE),
3433	metadata,
3434	SimpleNamespace(),
3435	)
3436
3437	assert output.shape == (num_tokens, 2, 1)
3438	assert lse is None
3439	index_group.convert_logical_to_physical_topk.assert_called_once()
3440	impl._run_mqa_kernel.assert_called_once()
3441
3442
3443	def test_flashinfer_sm120_hisparse_decode_uses_index_group():
3444	num_tokens = 4
3445	q = torch.empty(num_tokens, 2, 4, device=DEVICE_TYPE)
3446	topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=DEVICE_TYPE)
3447	physical = topk.clone()
3448	index_group = object.__new__(HiSparseMLAIndexGroup)
3449	index_group.convert_logical_to_physical_topk = MagicMock(return_value=physical)
3450	index_group.physical_kv_cache = MagicMock(
3451	return_value=torch.empty(1, device=DEVICE_TYPE)
3452	)
3453	impl = object.__new__(FlashInferMLASparseSM120Impl)
3454	impl.topk_indices_buffer = topk
3455	impl.index_group = index_group
3456	impl.index_group_index = 0
3457	impl._run_mqa_kernel = MagicMock(return_value=q[..., :1])
3458	metadata = SimpleNamespace(
3459	num_decode_tokens=num_tokens,
3460	topk_tokens=topk.shape[1],
3461	)
3462
3463	output, lse = FlashInferMLASparseSM120Impl.forward_mqa(
3464	impl,
3465	q,
3466	torch.empty(1, device=DEVICE_TYPE),
3467	metadata,
3468	SimpleNamespace(),
3469	)
3470
3471	assert output.shape == (num_tokens, 2, 1)
3472	assert lse is None
3473	index_group.convert_logical_to_physical_topk.assert_called_once()
3474	impl._run_mqa_kernel.assert_called_once()
3475
3476
3477	def test_hisparse_resident_prefill_uses_attention_block_stride():
3478	expected = torch.tensor([[19]], dtype=torch.int32)
3479	cache_handle = SimpleNamespace(
3480	all_context_pages_resident=True,
3481	view=SimpleNamespace(block_size=64, attention_block_stride=832),
3482	block_table=torch.tensor([[3]], dtype=torch.int32),
3483	)
3484	index_group = object.__new__(HiSparseMLAIndexGroup)
3485	index_group.caches = [cache_handle]
3486	index_group.physical_topk_indices = torch.empty((1, 1), dtype=torch.int32)
3487	index_group._convert_once = MagicMock(return_value=expected)
3488	topk = torch.zeros((2, 1), dtype=torch.int32)
3489	metadata = SimpleNamespace(
3490	req_id_per_token=torch.zeros(2, dtype=torch.int32),
3491	)
3492
3493	result = index_group.convert_logical_to_physical_topk(
3494	0,
3495	topk,
3496	metadata,
3497	block_stride_rows=None,
3498	return_valid_counts=False,
3499	)
3500
3501	assert result is expected
3502	assert index_group._convert_once.call_args.kwargs["block_stride_rows"] == 832
3503
3504
3505	def test_hisparse_flashmla_reorders_full_speculative_window_as_decode():
3506	builder = object.__new__(FlashMLASparseMetadataBuilder)
3507	builder.vllm_config = SimpleNamespace(
3508	attention_config=SimpleNamespace(hisparse_config=object()),
3509	speculative_config=SimpleNamespace(
3510	num_speculative_tokens=3,
3511	parallel_drafting=False,
3512	),
3513	parallel_config=SimpleNamespace(decode_context_parallel_size=1),
3514	)
3515
3516	builder._init_reorder_batch_threshold(
3517	256,
3518	supports_spec_as_decode=True,
3519	supports_dcp_with_varlen=True,
3520	)
3521
3522	assert builder.reorder_batch_threshold == 4
3523
3524
3525	def test_hisparse_mixed_single_token_batch_is_not_decode_only():
3526	cache = SimpleNamespace(
3527	runtime=SimpleNamespace(
3528	begin_forward=lambda: None,
3529	eager_host_mirror=True,
3530	),
3531	)
3532	metadata = SimpleNamespace(
3533	num_decode_tokens=1,
3534	num_actual_tokens=2,
3535	max_query_len=1,
3536	num_reqs=2,
3537	req_id_per_token=torch.arange(2),
3538	)
3539
3540	HiSparseCacheHandle._prepare_for_batch(cache, metadata)
3541
3542	assert not cache.decode_batch
3543
3544
3545	def test_flashmla_cache_dtype_aliases_use_ds_layout():
3546	# kv-cache dtype aliases are canonicalized to fp8_ds_mla before the layer
3547	# stores kv_cache_dtype, so they cannot bypass the gate.
3548	for alias in ("fp8", "fp8_e4m3"):
3549	assert (
3550	_canonicalize_sparse_mla_kv_cache_dtype(FlashMLASparseBackend, alias)
3551	== "fp8_ds_mla"
3552	)
3553
3554
3555	def test_flashmla_common_metadata_requires_uniform_decodes():
3556	common_metadata = SimpleNamespace(
3557	max_query_len=3,
3558	num_reqs=3,
3559	num_actual_tokens=6,
3560	query_start_loc_cpu=torch.tensor([0, 1, 3, 6]),
3561	is_prefilling=None,
3562	)
3563
3564	split = split_decodes_and_prefills(
3565	common_metadata,
3566	decode_threshold=128,
3567	require_uniform=FlashMLASparseMetadataBuilder.require_uniform_decodes,
3568	)
3569
3570	assert split == (1, 2, 1, 5)
3571
3572
3573	def test_flashmla_fp8_metadata_excludes_zero_token_decode_padding(monkeypatch):
3574	monkeypatch.setattr(
3575	"vllm.v1.attention.backends.mla.flashmla_sparse.get_mla_metadata",
3576	lambda: (object(), None),
3577	)
3578	builder = SimpleNamespace(
3579	device=torch.device(DEVICE_TYPE),
3580	dummy_block_table=torch.zeros(7, 1, device=DEVICE_TYPE),
3581	max_model_len_tensor=torch.zeros(7, device=DEVICE_TYPE),
3582	pcp_dcp_kv_gather=False,
3583	)
3584	query_start_loc_cpu = torch.tensor([0, 110, 220, 330, 440, 550, 660, 660])
3585	common_metadata = SimpleNamespace(
3586	num_actual_tokens=660,
3587	query_start_loc_cpu=query_start_loc_cpu,
3588	seq_lens=torch.arange(7, device=DEVICE_TYPE),
3589	)
3590	metadata = FlashMLASparseMetadata(
3591	num_reqs=7,
3592	max_query_len=110,
3593	max_seq_len=110,
3594	num_actual_tokens=660,
3595	query_start_loc=query_start_loc_cpu.to(DEVICE_TYPE),
3596	slot_mapping=torch.arange(660, device=DEVICE_TYPE),
3597	block_table=torch.zeros(7, 1, dtype=torch.int32, device=DEVICE_TYPE),
3598	req_id_per_token=torch.zeros(660, dtype=torch.int32, device=DEVICE_TYPE),
3599	num_decodes=7,
3600	num_prefills=0,
3601	num_decode_tokens=660,
3602	)
3603
3604	fp8_metadata = FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode(
3605	builder, common_metadata, metadata
3606	)
3607
3608	assert fp8_metadata.num_decodes == 6
3609	assert fp8_metadata.num_decode_tokens == 660
3610	assert fp8_metadata.decode is not None
3611	assert fp8_metadata.decode.decode_query_len == 110
3612	torch.testing.assert_close(
3613	fp8_metadata.decode.seq_lens, torch.arange(6, device=DEVICE_TYPE)
3614	)
3615
3616
3617	@pytest.mark.parametrize("use_mixed_batch", [False, True])
3618	def test_flashmla_fp8_paths_accept_decode_subset(monkeypatch, use_mixed_batch: bool):
3619	num_decode_tokens = 2
3620	num_batch_tokens = 5
3621	q = torch.empty(num_decode_tokens, 2, 3, device=DEVICE_TYPE)
3622	topk_indices = torch.empty(num_decode_tokens, 4, device=DEVICE_TYPE)
3623	kernel_q_shapes = []
3624
3625	def convert_indices(args, *kwargs): # noqa: ARG001
3626	assert not kwargs.get("HAS_PREFILL_WORKSPACE", False)
3627	if not kwargs.get("return_valid_counts", False):
3628	return topk_indices
3629	valid_counts = torch.full(
3630	(num_decode_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE
3631	)
3632	return topk_indices, valid_counts
3633
3634	monkeypatch.setattr(
3635	"vllm.v1.attention.backends.mla.flashmla_sparse."
3636	"triton_convert_req_index_to_global_index",
3637	convert_indices,
3638	)
3639
3640	def run_kernel(**kwargs):
3641	kernel_q_shapes.append(kwargs["q"].shape)
3642	return kwargs["q"][..., :1], None
3643
3644	if use_mixed_batch:
3645	fp8_metadata = FlashMLASparseMetadata.FP8KernelMetadata(
3646	scheduler_metadata=object(),
3647	dummy_block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE),
3648	cache_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE),
3649	)
3650	else:
3651	FP8Meta = FlashMLASparseMetadata.FP8SeparatePrefillDecode
3652	fp8_metadata = FP8Meta(
3653	num_decodes=1,
3654	num_prefills=1,
3655	num_decode_tokens=num_decode_tokens,
3656	num_prefill_tokens=num_batch_tokens - num_decode_tokens,
3657	decode=FP8Meta.Decode(
3658	seq_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE),
3659	kernel_metadata=object(),
3660	decode_query_len=num_decode_tokens,
3661	),
3662	prefill=FP8Meta.Prefill(
3663	request_ids=torch.empty(
3664	num_batch_tokens, dtype=torch.int32, device=DEVICE_TYPE
3665	),
3666	workspace_starts=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE),
3667	chunks=[],
3668	),
3669	)
3670	metadata = SimpleNamespace(
3671	fp8_extra_metadata=fp8_metadata,
3672	fp8_use_mixed_batch=use_mixed_batch,
3673	num_actual_tokens=num_batch_tokens,
3674	num_decode_tokens=num_decode_tokens,
3675	req_id_per_token=torch.empty(
3676	num_batch_tokens, dtype=torch.int32, device=DEVICE_TYPE
3677	),
3678	block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE),
3679	block_size=64,
3680	)
3681	impl = SimpleNamespace(
3682	kv_cache_dtype="fp8_ds_mla",
3683	topk_indices_buffer=topk_indices,
3684	num_heads=2,
3685	kv_lora_rank=1,
3686	index_group=None,
3687	index_group_index=0,
3688	dcp_world_size=1,
3689	pcp_dcp_kv_gather=False,
3690	need_to_return_lse_for_decode=False,
3691	_fp8_flash_mla_kernel=run_kernel,
3692	_convert_logical_to_physical_topk=(
3693	lambda indices, metadata, **kwargs: convert_indices(
3694	metadata.req_id_per_token[: indices.shape[0]],
3695	metadata.block_table,
3696	indices,
3697	return_valid_counts=kwargs["return_valid_counts"],
3698	)
3699	),
3700	)
3701	impl._forward_fp8_kv_mixed_batch = MethodType(
3702	FlashMLASparseImpl._forward_fp8_kv_mixed_batch, impl
3703	)
3704	impl._forward_fp8_kv_separate_prefill_decode = MethodType(
3705	FlashMLASparseImpl._forward_fp8_kv_separate_prefill_decode, impl
3706	)
3707
3708	output, lse = FlashMLASparseImpl.forward_mqa(
3709	impl,
3710	q,
3711	torch.empty(0, device=DEVICE_TYPE),
3712	metadata,
3713	None,
3714	)
3715
3716	assert kernel_q_shapes == [(1, num_decode_tokens, 2, 3)]
3717	assert output.shape == (num_decode_tokens, 2, 1)
3718	assert lse is None
3719
3720
3721	def _build_sparse_dcp_vllm_config(
3722	local_heads: int,
3723	dcp_world_size: int,
3724	comm_backend: str = "ag_rs",
3725	):
3726	"""Minimal sparse-MLA VllmConfig for the FlashMLASparse DCP head-envelope
3727	guard. TP is simulated by mocking ``get_num_attention_heads`` to return the
3728	per-rank head count, as the decode-correctness test above does.
3729	"""
3730	kv_lora_rank = 512
3731	qk_nope_head_dim = 128
3732	qk_rope_head_dim = 64
3733	v_head_dim = 128
3734	head_size = kv_lora_rank + qk_rope_head_dim
3735	topk_tokens = 128
3736
3737	vllm_config = create_vllm_config(
3738	model_name="deepseek-ai/DeepSeek-V2-Lite-Chat",
3739	tensor_parallel_size=1,
3740	max_model_len=4096,
3741	block_size=64,
3742	hf_config_override={
3743	"index_topk": topk_tokens,
3744	"attn_module_list_cfg": [{"topk_tokens": topk_tokens}],
3745	},
3746	)
3747	model_config = vllm_config.model_config
3748	model_config.dtype = torch.bfloat16
3749	model_config.hf_text_config = SimpleNamespace(
3750	index_topk=topk_tokens,
3751	q_lora_rank=None,
3752	kv_lora_rank=kv_lora_rank,
3753	qk_nope_head_dim=qk_nope_head_dim,
3754	qk_rope_head_dim=qk_rope_head_dim,
3755	v_head_dim=v_head_dim,
3756	model_type="deepseek_v2",
3757	)
3758	del model_config.hf_config.index_topk # Composite configs only nest this field.
3759	model_config.get_num_attention_heads = MethodType(
3760	lambda self, parallel_config: local_heads, model_config
3761	)
3762	model_config.get_num_kv_heads = MethodType(
3763	lambda self, parallel_config: 1, model_config
3764	)
3765	model_config.get_head_size = MethodType(lambda self: head_size, model_config)
3766	model_config.get_sliding_window = MethodType(lambda self: None, model_config)
3767
3768	vllm_config.cache_config.cache_dtype = "fp8_ds_mla"
3769	vllm_config.parallel_config.decode_context_parallel_size = dcp_world_size
3770	vllm_config.parallel_config.dcp_comm_backend = comm_backend
3771	# The base builder clones the layer's dense-MHA prefill backend from
3772	# static_forward_context; the guard tests never run prefill.
3773	vllm_config.compilation_config.static_forward_context["placeholder"] = (
3774	SimpleNamespace(prefill_backend=None)
3775	)
3776	return vllm_config
3777
3778
3779	@pytest.mark.skipif(
3780	torch.cuda.get_device_capability() < (9, 0),
3781	reason="FlashMLASparseBackend requires CUDA 9.0 or higher",
3782	)
3783	@pytest.mark.parametrize(
3784	"local_heads,dcp_world_size,should_raise",
3785	[
3786	(16, 8, True),
3787	(24, 4, True),
3788	(16, 4, False),
3789	(16, 1, False),
3790	],
3791	)
3792	def test_fp8_dcp_head_envelope_guard(local_heads, dcp_world_size, should_raise):
3793	"""The fp8 decode envelope (head padding + tile-scheduler metadata) is
3794	sized from the local head count while the kernel runs on the DCP-gathered
3795	heads, so the builder must reject configs where the two pad differently.
3796	"""
3797	device = torch.device(DEVICE_TYPE)
3798	vllm_config = _build_sparse_dcp_vllm_config(local_heads, dcp_world_size)
3799	kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
3800	builder_cls = FlashMLASparseBackend.get_builder_cls()
3801
3802	if should_raise:
3803	with pytest.raises(NotImplementedError, match="envelope"):
3804	builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device)
3805	else:
3806	builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device)
3807	gathered_heads = local_heads * dcp_world_size
3808	local_pad = 64 if local_heads <= 64 else 128
3809	gathered_pad = 64 if gathered_heads <= 64 else 128
3810	assert builder.fp8_decode_padded_heads == local_pad
3811	assert local_pad == gathered_pad
3812
3813
3814	def test_fp8_mixed_batch_dcp_neutralizes_empty_rows(monkeypatch):
3815	"""A decode row whose top-k shard holds no local candidates (all -1) has
3816	undefined kernel out/lse; it must come back as (0, -inf), the identity of
3817	the cross-rank LSE merge, or a NaN would survive the merge even at zero
3818	weight (0 * NaN = NaN)."""
3819	num_tokens, num_heads, head_dim = 3, 2, 3
3820	q = torch.empty(num_tokens, num_heads, head_dim, device=DEVICE_TYPE)
3821	local_indices = torch.tensor(
3822	[[0, 1, -1, -1], [-1, -1, -1, -1], [2, -1, 3, -1]],
3823	dtype=torch.int32,
3824	device=DEVICE_TYPE,
3825	)
3826
3827	monkeypatch.setattr(
3828	"vllm.v1.attention.backends.mla.flashmla_sparse."
3829	"triton_filter_and_convert_dcp_index",
3830	lambda args, *kwargs: local_indices,
3831	)
3832
3833	def run_kernel(**kwargs):
3834	out = torch.full(
3835	(1, num_tokens, num_heads, 1), float("nan"), device=DEVICE_TYPE
3836	)
3837	lse = torch.full((1, num_heads, num_tokens), float("nan"), device=DEVICE_TYPE)
3838	for token_id in (0, 2): # rows with local candidates get real values
3839	out[0, token_id] = float(token_id + 1)
3840	lse[0, :, token_id] = float(token_id + 1)
3841	return out, lse
3842
3843	metadata = SimpleNamespace(
3844	fp8_extra_metadata=FlashMLASparseMetadata.FP8KernelMetadata(
3845	scheduler_metadata=object(),
3846	dummy_block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE),
3847	cache_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE),
3848	),
3849	req_id_per_token=torch.empty(num_tokens, dtype=torch.int32, device=DEVICE_TYPE),
3850	block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE),
3851	block_size=64,
3852	cp_kv_cache_interleave_size=1,
3853	)
3854	impl = SimpleNamespace(
3855	index_group=None,
3856	index_group_index=0,
3857	dcp_world_size=2,
3858	dcp_rank=0,
3859	need_to_return_lse_for_decode=True,
3860	_fp8_flash_mla_kernel=run_kernel,
3861	)
3862
3863	out, lse = FlashMLASparseImpl._forward_fp8_kv_mixed_batch(
3864	impl, q, torch.empty(0, device=DEVICE_TYPE), local_indices, metadata
3865	)
3866
3867	assert torch.equal(out[1], torch.zeros_like(out[1]))
3868	assert torch.isneginf(lse[1]).all()
3869	for token_id in (0, 2):
3870	assert torch.equal(out[token_id], torch.full_like(out[token_id], token_id + 1))
3871	assert torch.equal(lse[token_id], torch.full_like(lse[token_id], token_id + 1))
3872	assert out.is_contiguous()
3873	assert not out.isnan().any()
3874	assert not lse.isnan().any()
3875
3876
3877	def test_hisparse_prefill_reuses_builder_staging_plan():
3878	"""Every layer must reuse the batch plan instead of synchronizing to dedupe."""
3879	plan = SimpleNamespace(
3880	block_table=torch.tensor([[0]], dtype=torch.int32),
3881	ensure_gpu_sources=MagicMock(),
3882	)
3883	staged = torch.empty((1, 1, 8))
3884	calls = []
3885
3886	def gather(kv_cache, staging_plan, resident_cache=None):
3887	calls.append((kv_cache, staging_plan, resident_cache))
3888	return staged
3889
3890	resident_cache = torch.empty((1, 1, 8))
3891	resident_block_table = torch.tensor([[1]], dtype=torch.int32)
3892	cache = SimpleNamespace(
3893	runtime=SimpleNamespace(
3894	gather_prefill_cache=gather,
3895	),
3896	view=SimpleNamespace(cache=resident_cache, block_size=1),
3897	block_table=resident_block_table,
3898	)
3899	index_group = object.__new__(HiSparseMLAIndexGroup)
3900	index_group.caches = [cache]
3901	source = torch.empty((1, 1, 8))
3902	metadata = SimpleNamespace(
3903	num_decodes=0,
3904	num_decode_tokens=0,
3905	seq_lens=torch.tensor([1], dtype=torch.int32),
3906	prefill=SimpleNamespace(host_staging_plan=plan),
3907	req_id_per_token=torch.tensor([0], dtype=torch.int32),
3908	)
3909
3910	result, block_table, request_ids = index_group.stage_prefill_rows(
3911	0, source, metadata
3912	)
3913
3914	assert result is staged
3915	assert block_table is plan.block_table
3916	torch.testing.assert_close(request_ids, metadata.req_id_per_token)
3917	plan.ensure_gpu_sources.assert_called_once()
3918	args = plan.ensure_gpu_sources.call_args.args
3919	torch.testing.assert_close(args[0], resident_block_table)
3920	assert args[1] == 1
3921	assert calls == [(source, plan, resident_cache)]
3922
3923
3924	def test_hisparse_fp8_prefill_gather_uses_dedicated_stream(monkeypatch):
3925	compute_stream = object()
3926	prefill_stream = MagicMock()
3927	gather = MagicMock()
3928	monkeypatch.setattr(index_group_module, "current_stream", lambda: compute_stream)
3929	monkeypatch.setattr(
3930	index_group_module.ops, "cp_gather_and_upconvert_fp8_kv_cache", gather
3931	)
3932
3933	resident_cache = torch.empty((2, 4, 656), dtype=torch.uint8)
3934	ready = MagicMock()
3935	cache = SimpleNamespace(
3936	view=SimpleNamespace(cache=resident_cache, block_size=4),
3937	block_table=torch.tensor([[2, 3]], dtype=torch.int32),
3938	)
3939	index_group = object.__new__(HiSparseMLAIndexGroup)
3940	index_group.caches = [cache]
3941	index_group.prefill_stream = prefill_stream
3942	index_group.prefill_ready_events = [ready]
3943	source = torch.empty((3, 4, 656), dtype=torch.uint8)
3944	dst = torch.empty((8, 576), dtype=torch.bfloat16)
3945	block_table = torch.tensor([[0, 1]], dtype=torch.int32)
3946	workspace_starts = torch.tensor([0], dtype=torch.int32)
3947	plan = SimpleNamespace(
3948	block_table=torch.tensor([[0, 1]], dtype=torch.int32),
3949	row_ids=torch.arange(8, dtype=torch.int32).view(1, -1),
3950	gpu_row_ids=torch.full((1, 8), -1, dtype=torch.int32),
3951	ensure_gpu_sources=MagicMock(),
3952	)
3953	metadata = SimpleNamespace(
3954	prefill=SimpleNamespace(host_staging_plan=plan),
3955	num_decodes=0,
3956	)
3957
3958	result = index_group.gather_fp8_prefill(
3959	0,
3960	source,
3961	dst,
3962	block_table,
3963	workspace_starts,
3964	1,
3965	metadata,
3966	0,
3967	)
3968
3969	prefill_stream.wait_stream.assert_called_once_with(compute_stream)
3970	ready.record.assert_called_once_with(prefill_stream)
3971	assert result is ready
3972	gather.assert_called_once()
3973	args = gather.call_args
3974	assert args.args[0] is resident_cache
3975	assert args.args[1] is dst
3976	torch.testing.assert_close(args.args[2], plan.block_table)
3977	assert args.args[3] is workspace_starts
3978	assert args.args[4] == 1
3979	assert args.kwargs["host_cache"].data_ptr() == source.data_ptr()
3980	assert args.kwargs["host_row_ids"] is plan.row_ids
3981	assert args.kwargs["device_row_ids"] is plan.gpu_row_ids
3982	plan.ensure_gpu_sources.assert_called_once()
3983	ensure_args = plan.ensure_gpu_sources.call_args.args
3984	torch.testing.assert_close(ensure_args[0], cache.block_table)
3985	assert ensure_args[1] == 4
3986
3987
3988	def test_sparse_impl_observes_repointed_indexer_buffer():
3989	"""The MTP proposer repoints the draft's indexer at the target model's buffer
3990	after the backend impl is built, so the impl must resolve the buffer per read.
3991	Snapshotting it in __init__ leaves the layer reading indices nothing writes."""
3992	impl = object.__new__(FlashInferMLASparseImpl)
3993	own = torch.zeros(4, 8, dtype=torch.int32)
3994	target = torch.ones(4, 8, dtype=torch.int32)
3995	indexer = SimpleNamespace(topk_indices_buffer=own)
3996	impl.init_topk_indices_buffer(indexer, None)
3997
3998	assert impl.topk_indices_buffer is own
3999
4000	indexer.topk_indices_buffer = target
4001
4002	assert impl.topk_indices_buffer is target
4003
4004
4005	def test_explicit_topk_buffer_supersedes_indexer():
4006	"""Backbone skip-topk layers have no indexer, and the proposer also assigns the
4007	shared buffer directly onto draft submodules."""
4008	impl = object.__new__(FlashInferMLASparseImpl)
4009	indexer = SimpleNamespace(
4010	topk_indices_buffer=torch.zeros(2, 2, dtype=torch.int32),
4011	)
4012	impl.init_topk_indices_buffer(indexer, None)
4013	explicit = torch.ones(2, 2, dtype=torch.int32)
4014
4015	impl.topk_indices_buffer = explicit
4016
4017	assert impl.topk_indices_buffer is explicit
4018
4019
4020	def test_sparse_mla_common_impl_resolves_buffer_lazily():
4021	"""Guards the whole SparseMLACommonImpl family at once: a backend that
4022	snapshots the buffer instead silently loses MTP buffer sharing."""
4023	assert issubclass(SparseMLACommonImpl, SharedTopkIndicesBuffer)
4024	assert isinstance(SparseMLACommonImpl.topk_indices_buffer, property), (
4025	"topk_indices_buffer must stay a lazily-resolved property"
4026	)