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

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 1,311 lines 45.7 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3	"""Fused compressor + FP8/MXFP4 UE8M0 quantization + KV cache insert kernels.
4
5	Three specialized kernels:
6	- _fused_kv_compress_norm_rope_insert_sparse_attn:
7	head=512, nope=448 FP8 + rope=64 bf16
8	- _fused_kv_compress_norm_rope_insert_indexer_attn:
9	head=128, all FP8, 1 block/token
10	- _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn:
11	head=128, MXFP4 (block=32), 4 ue8m0 bytes
12
13	RoPE is register-based via tl.reshape -> tl.split -> tl.interleave (or the
14	even/odd halves are consumed directly for MXFP4, no interleave needed).
15	FP8 UE8M0 quant uses tl.reshape to tile [N_QUANT_BLOCKS, QUANT_BLOCK] for
16	per-block absmax entirely in registers. MXFP4 does the same tiling on the
17	even/odd halves, producing (N_QUANT_BLOCKS, MXFP4_BLOCK/2) packed nibbles
18	and N_QUANT_BLOCKS ue8m0 bytes.
19	"""
20
21	from dataclasses import dataclass
22	from functools import lru_cache
23	from types import SimpleNamespace
24	from typing import Any
25
26	import torch
27
28	from vllm.model_executor.warmup.jit_warmup import kernel_launcher
29	from vllm.model_executor.warmup.jit_warmup_triton_helper import (
30	LaunchSpec,
31	TritonWarmupTensor,
32	VllmTritonJitKernel,
33	)
34	from vllm.platforms import current_platform
35	from vllm.triton_utils import tl, triton
36	from vllm.utils.math_utils import round_up
37
38	if current_platform.is_rocm():
39	from vllm.platforms.rocm import _ON_GFX950
40	else:
41	_ON_GFX950 = False
42
43	from .fused_indexer_q import _fp32x2_to_fp4x2
44
45
46	def compress_norm_rope_store_triton(
47	state_cache: torch.Tensor,
48	num_actual: int,
49	token_to_req_indices: torch.Tensor,
50	positions: torch.Tensor,
51	slot_mapping: torch.Tensor,
52	block_table: torch.Tensor,
53	block_size: int,
54	state_width: int,
55	cos_sin_cache: torch.Tensor,
56	kv_cache: torch.Tensor,
57	k_cache_metadata: Any,
58	pdl_kwargs: dict,
59	head_dim: int,
60	rope_head_dim: int,
61	compress_ratio: int,
62	overlap: bool,
63	use_fp4_cache: bool,
64	rms_norm_weight: torch.Tensor,
65	rms_norm_eps: float,
66	quant_block: int,
67	token_stride: int,
68	scale_dim: int,
69	) -> None:
70	"""Shared triton launcher for the fused compress+norm+RoPE+insert path.
71
72	Picks one of the three kernels in this module based on ``head_dim`` and
73	``use_fp4_cache``. Identical launch signature for all three.
74	"""
75	if head_dim == 512:
76	kernel = _fused_kv_compress_norm_rope_insert_sparse_attn
77	num_warps = 4
78	kernel_kwargs = {"SANITIZE_CACHE_NANS": _ON_GFX950}
79	else:
80	_FUSED_KV_COMPRESS_NORM_ROPE_INSERT_INDEXER_TRITON_KERNEL(
81	state_cache=state_cache,
82	num_actual=num_actual,
83	token_to_req_indices=token_to_req_indices,
84	positions=positions,
85	slot_mapping=slot_mapping,
86	block_table=block_table,
87	block_size=block_size,
88	state_width=state_width,
89	cos_sin_cache=cos_sin_cache,
90	kv_cache=kv_cache,
91	k_cache_metadata=k_cache_metadata,
92	pdl_kwargs=pdl_kwargs,
93	head_dim=head_dim,
94	rope_head_dim=rope_head_dim,
95	compress_ratio=compress_ratio,
96	overlap=overlap,
97	use_fp4_cache=use_fp4_cache,
98	rms_norm_weight=rms_norm_weight,
99	rms_norm_eps=rms_norm_eps,
100	quant_block=quant_block,
101	token_stride=token_stride,
102	scale_dim=scale_dim,
103	)
104	return
105
106	kernel[(num_actual,)](
107	# state cache
108	state_cache,
109	state_cache.stride(0),
110	state_cache.stride(1),
111	# metadata
112	token_to_req_indices,
113	positions,
114	slot_mapping,
115	block_table,
116	block_table.stride(0),
117	block_size,
118	# RMSNorm
119	rms_norm_weight,
120	rms_norm_eps,
121	# RoPE
122	cos_sin_cache,
123	cos_sin_cache.stride(0),
124	# KV cache
125	kv_cache,
126	k_cache_metadata.slot_mapping,
127	kv_cache.shape[1], # paged KV cache block size (tokens per block)
128	# constexprs
129	HEAD_SIZE=head_dim,
130	TRITON_BLOCK_SIZE=triton.next_power_of_2(head_dim),
131	STATE_WIDTH=state_width,
132	COMPRESS_RATIO=compress_ratio,
133	OVERLAP=overlap,
134	ROPE_HEAD_DIM=rope_head_dim,
135	FP8_MAX=448.0,
136	QUANT_BLOCK=quant_block,
137	TOKEN_STRIDE=token_stride,
138	SCALE_DIM=scale_dim,
139	KV_BLOCK_STRIDE=kv_cache.stride(0),
140	num_warps=num_warps,
141	**kernel_kwargs,
142	**pdl_kwargs,
143	)
144
145
146	# =============================================================================
147	# DeepseekV4 Attention path (head=512, nope=448 FP8 + rope=64 bf16)
148	# =============================================================================
149	@triton.jit
150	def _fused_kv_compress_norm_rope_insert_sparse_attn(
151	# ── state cache (compressor internal state) ──
152	state_cache_ptr,
153	state_cache_stride0,
154	state_cache_stride1,
155	# ── metadata ──
156	token_to_req_indices_ptr,
157	positions_ptr,
158	slot_mapping_ptr,
159	block_table_ptr,
160	block_table_stride,
161	block_size,
162	# ── RMSNorm ──
163	rms_norm_weight_ptr,
164	rms_norm_eps,
165	# ── RoPE ──
166	cos_sin_cache_ptr,
167	cos_sin_stride,
168	# ── KV cache output ──
169	k_cache_ptr,
170	kv_slot_mapping_ptr,
171	kv_cache_block_size,
172	# ── constexprs ──
173	HEAD_SIZE: tl.constexpr,
174	TRITON_BLOCK_SIZE: tl.constexpr,
175	STATE_WIDTH: tl.constexpr,
176	COMPRESS_RATIO: tl.constexpr,
177	OVERLAP: tl.constexpr,
178	ROPE_HEAD_DIM: tl.constexpr,
179	FP8_MAX: tl.constexpr, # 448.0
180	QUANT_BLOCK: tl.constexpr, # 64 for DeepseekV4
181	TOKEN_STRIDE: tl.constexpr, # 576 for DeepseekV4
182	SCALE_DIM: tl.constexpr, # 8 for DeepseekV4 (7 real + 1 pad)
183	KV_BLOCK_STRIDE: tl.constexpr,
184	SANITIZE_CACHE_NANS: tl.constexpr,
185	):
186	"""Fused compress → RMSNorm → FP8 quant (nope) → RoPE → bf16 store (rope).
187
188	One program per token; early-exits for non-boundary positions.
189
190	Cache block layout (``block_size`` tokens):
191	[0, bs*576): token data (448 fp8 + 128 bf16 each)
192	[bs576, +bs8): uint8 UE8M0 scales (7 real + 1 pad each)
193	"""
194	token_idx = tl.program_id(0)
195
196	slot_id = tl.load(slot_mapping_ptr + token_idx)
197	if slot_id < 0:
198	return
199
200	position = tl.load(positions_ptr + token_idx)
201	if (position + 1) % COMPRESS_RATIO != 0:
202	return
203
204	req_idx = tl.load(token_to_req_indices_ptr + token_idx)
205
206	# ── Gather state cache entries ────────────────────────────────────
207	start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1
208	tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO)
209	pos = start + tokens
210	mask_pos = pos >= 0
211
212	block_indices = pos // block_size
213	block_numbers = tl.load(
214	block_table_ptr + req_idx * block_table_stride + block_indices,
215	mask=mask_pos,
216	other=0,
217	)
218	block_offsets = pos % block_size
219	head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE
220
221	block = tl.arange(0, TRITON_BLOCK_SIZE)
222	mask = block < HEAD_SIZE
223	block_numbers_i64 = block_numbers.to(tl.int64)
224
225	# Precomputed row base shared by score and kv loads
226	row_base = (
227	state_cache_ptr
228	+ block_numbers_i64 * state_cache_stride0
229	+ block_offsets * state_cache_stride1
230	+ head_offset
231	)
232
233	combined_mask = mask_pos[:, None] & mask[None, :]
234
235	# ── Softmax + weighted sum ───────────────────────────────────────
236	score = tl.load(
237	row_base[:, None] + STATE_WIDTH + block[None, :],
238	mask=combined_mask,
239	other=float("-inf"),
240	)
241	score = tl.softmax(score, dim=0)
242
243	kv = tl.load(
244	row_base[:, None] + block[None, :],
245	mask=combined_mask,
246	other=0.0,
247	)
248
249	compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32
250
251	# ── RMSNorm (fp32 throughout) ──────────────────────────────────────
252	rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0)
253	variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_SIZE
254	rrms = tl.rsqrt(variance + rms_norm_eps)
255	normed = compressed_kv * rrms * rms_w
256
257	# ── KV cache pointers ────────────────────────────────────────────
258	kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx)
259	if kv_slot_idx < 0:
260	return
261	kv_block_idx = kv_slot_idx // kv_cache_block_size
262	kv_pos_in_block = kv_slot_idx % kv_cache_block_size
263
264	cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE
265	fp8_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE
266	scale_ptr = (
267	cache_block_ptr
268	+ kv_cache_block_size * TOKEN_STRIDE
269	+ kv_pos_in_block * SCALE_DIM
270	)
271
272	NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM # 448
273	HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2 # 32
274
275	# FP8 UE8M0 quant: cast fp32 → bf16 → fp32 before quant to match reference.
276	N_QUANT_BLOCKS: tl.constexpr = TRITON_BLOCK_SIZE // QUANT_BLOCK
277	N_NOPE_BLOCKS: tl.constexpr = NOPE_HEAD_DIM // QUANT_BLOCK # 7
278	INV_FP8_MAX: tl.constexpr = 1.0 / FP8_MAX
279
280	quant_input = normed.to(tl.bfloat16).to(tl.float32)
281	quant_2d = tl.reshape(quant_input, (N_QUANT_BLOCKS, QUANT_BLOCK))
282	abs_2d = tl.abs(quant_2d)
283	block_absmax = tl.max(abs_2d, axis=1) # [N_QUANT_BLOCKS] fp32
284	block_absmax = tl.maximum(block_absmax, 1e-4)
285
286	raw_scales = block_absmax * INV_FP8_MAX
287	exponents = tl.ceil(tl.log2(raw_scales))
288	inv_scales = tl.exp2(-exponents)
289	inv_scales_col = tl.reshape(inv_scales, (N_QUANT_BLOCKS, 1))
290	x_scaled = quant_2d * inv_scales_col
291	x_clamped = tl.clamp(x_scaled, -FP8_MAX, FP8_MAX)
292	x_fp8 = x_clamped.to(tl.float8e4nv)
293	x_uint8 = x_fp8.to(tl.uint8, bitcast=True)
294	x_uint8_flat = tl.reshape(x_uint8, (TRITON_BLOCK_SIZE,))
295
296	nope_mask = block < NOPE_HEAD_DIM
297	tl.store(fp8_ptr + block, x_uint8_flat, mask=nope_mask)
298
299	scale_idx = tl.arange(0, N_QUANT_BLOCKS)
300	encoded = exponents + 127.0
301	max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0
302	encoded = tl.maximum(tl.minimum(encoded, max_encoded), 0.0)
303	tl.store(
304	scale_ptr + scale_idx,
305	encoded.to(tl.uint8),
306	mask=scale_idx < N_NOPE_BLOCKS,
307	)
308	tl.store(scale_ptr + N_NOPE_BLOCKS, tl.zeros((), dtype=tl.uint8))
309
310	# Register-based GPT-J RoPE in fp32.
311	NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2
312	NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2
313
314	pair_2d = tl.reshape(normed, (NUM_PAIRS, 2))
315	even, odd = tl.split(pair_2d) # each [NUM_PAIRS] fp32
316
317	pair_idx = tl.arange(0, NUM_PAIRS)
318	rope_pair_local = pair_idx - NOPE_PAIRS
319	is_rope_pair = rope_pair_local >= 0
320	cs_idx = tl.maximum(rope_pair_local, 0)
321
322	compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO
323	cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride
324	cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0)
325	sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0)
326
327	new_even = even * cos_v - odd * sin_v
328	new_odd = odd * cos_v + even * sin_v
329	result = tl.interleave(new_even, new_odd) # [TRITON_BLOCK_SIZE] fp32
330	if SANITIZE_CACHE_NANS:
331	result = tl.where(result == result, result, 0.0)
332
333	# Store rotated rope portion as bf16 into the cache's bf16 area.
334	bf16_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16))
335	rope_local = block - NOPE_HEAD_DIM
336	is_rope = (block >= NOPE_HEAD_DIM) & mask
337	tl.store(bf16_ptr + rope_local, result.to(tl.bfloat16), mask=is_rope)
338
339
340	# =============================================================================
341	# Split kernels variant of the head=512 compressor (deep cr=128 gather).
342	# - compress gather: instead of launching one program per token, split along
343	# the head dimension to maximize CU occupancy. The head dimension split
344	# does not require cross-group reduction
345	# - finalize norm rope quant store: same as the single pass kernel due to its
346	# per-token nature
347	# Mirrors the CUDA cutedsl split kernel where num_splits is occupancy-targeted.
348	# Currently only tested and validated on ROCm gfx950
349	# =============================================================================
350	@lru_cache(maxsize=1)
351	def _n_cu() -> int:
352	return torch.cuda.get_device_properties(0).multi_processor_count
353
354
355	def _pick_compress_num_splits(
356	num_actual: int, compress_ratio: int, head_dim: int
357	) -> int:
358	"""Occupancy-targeted column splits for the cr>=128 head=512 compressor.
359
360	Sizes the per-token fan-out so (estimated computing tokens) * num_splits ~
361	#CU, capped by head tiling at a 32-wide min tile, as a power-of-2 divisor of
362	head_dim.
363	"""
364	max_splits = head_dim // 32
365	est_compute = max(1, num_actual // compress_ratio)
366	target = -(-_n_cu() // est_compute) # ceil(#CU / est_compute)
367	ns = 1
368	while ns * 2 <= min(target, max_splits) and head_dim % (ns * 2) == 0:
369	ns *= 2
370	return ns
371
372
373	@triton.jit
374	def _compress_gather_split_sparse_attn(
375	state_cache_ptr,
376	state_cache_stride0,
377	state_cache_stride1,
378	positions_ptr,
379	slot_mapping_ptr,
380	token_to_req_indices_ptr,
381	block_table_ptr,
382	block_table_stride,
383	block_size,
384	scratch_ptr,
385	scratch_stride,
386	HEAD_SIZE: tl.constexpr,
387	STATE_WIDTH: tl.constexpr,
388	COMPRESS_RATIO: tl.constexpr,
389	NUM_SPLITS: tl.constexpr,
390	HEAD_TILE: tl.constexpr, # HEAD_SIZE // NUM_SPLITS
391	):
392	"""Stage 1: per-(token, head-split) compress gather, write to fp32 scratch.
393
394	No-overlap gather (cr>=128) on rows [0, COMPRESS_RATIO)
395	"""
396	pid = tl.program_id(0)
397	token_idx = pid // NUM_SPLITS
398	split_idx = pid % NUM_SPLITS
399
400	slot_id = tl.load(slot_mapping_ptr + token_idx)
401	if slot_id < 0:
402	return
403	position = tl.load(positions_ptr + token_idx)
404	if (position + 1) % COMPRESS_RATIO != 0:
405	return
406	req_idx = tl.load(token_to_req_indices_ptr + token_idx)
407
408	start = position - COMPRESS_RATIO + 1
409	rows = tl.arange(0, COMPRESS_RATIO)
410	pos = start + rows
411	mask_pos = pos >= 0
412	block_numbers = tl.load(
413	block_table_ptr + req_idx * block_table_stride + pos // block_size,
414	mask=mask_pos,
415	other=0,
416	).to(tl.int64)
417	block_offsets = pos % block_size
418
419	col = split_idx * HEAD_TILE + tl.arange(0, HEAD_TILE)
420	row_base = (
421	state_cache_ptr
422	+ block_numbers * state_cache_stride0
423	+ block_offsets * state_cache_stride1
424	)
425	cmask = mask_pos[:, None]
426
427	score = tl.load(
428	row_base[:, None] + STATE_WIDTH + col[None, :],
429	mask=cmask,
430	other=float("-inf"),
431	)
432	score = tl.softmax(score, dim=0)
433	kv = tl.load(row_base[:, None] + col[None, :], mask=cmask, other=0.0)
434	compressed = tl.sum(kv * score, axis=0) # [HEAD_TILE] fp32
435	tl.store(scratch_ptr + token_idx * scratch_stride + col, compressed)
436
437
438	@triton.jit
439	def _finalize_norm_rope_quant_store_sparse_attn(
440	scratch_ptr,
441	scratch_stride,
442	positions_ptr,
443	slot_mapping_ptr,
444	rms_norm_weight_ptr,
445	rms_norm_eps,
446	cos_sin_cache_ptr,
447	cos_sin_stride,
448	k_cache_ptr,
449	kv_slot_mapping_ptr,
450	kv_cache_block_size,
451	HEAD_SIZE: tl.constexpr,
452	TRITON_BLOCK_SIZE: tl.constexpr,
453	COMPRESS_RATIO: tl.constexpr,
454	ROPE_HEAD_DIM: tl.constexpr,
455	FP8_MAX: tl.constexpr,
456	QUANT_BLOCK: tl.constexpr,
457	TOKEN_STRIDE: tl.constexpr,
458	SCALE_DIM: tl.constexpr,
459	KV_BLOCK_STRIDE: tl.constexpr,
460	SANITIZE_CACHE_NANS: tl.constexpr,
461	):
462	"""Stage 2: read compressed_kv[512] from scratch buffer, then
463	RMSNorm + FP8 quant (nope) + RoPE + bf16 store
464	"""
465	token_idx = tl.program_id(0)
466	slot_id = tl.load(slot_mapping_ptr + token_idx)
467	if slot_id < 0:
468	return
469	position = tl.load(positions_ptr + token_idx)
470	if (position + 1) % COMPRESS_RATIO != 0:
471	return
472
473	block = tl.arange(0, TRITON_BLOCK_SIZE)
474	mask = block < HEAD_SIZE
475	compressed_kv = tl.load(
476	scratch_ptr + token_idx * scratch_stride + block, mask=mask, other=0.0
477	)
478
479	rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0)
480	variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_SIZE
481	rrms = tl.rsqrt(variance + rms_norm_eps)
482	normed = compressed_kv * rrms * rms_w
483
484	kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx)
485	if kv_slot_idx < 0:
486	return
487	kv_block_idx = kv_slot_idx // kv_cache_block_size
488	kv_pos_in_block = kv_slot_idx % kv_cache_block_size
489	cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE
490	fp8_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE
491	scale_ptr = (
492	cache_block_ptr
493	+ kv_cache_block_size * TOKEN_STRIDE
494	+ kv_pos_in_block * SCALE_DIM
495	)
496
497	NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM
498	HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2
499	N_QUANT_BLOCKS: tl.constexpr = TRITON_BLOCK_SIZE // QUANT_BLOCK
500	N_NOPE_BLOCKS: tl.constexpr = NOPE_HEAD_DIM // QUANT_BLOCK
501	INV_FP8_MAX: tl.constexpr = 1.0 / FP8_MAX
502
503	quant_input = normed.to(tl.bfloat16).to(tl.float32)
504	quant_2d = tl.reshape(quant_input, (N_QUANT_BLOCKS, QUANT_BLOCK))
505	block_absmax = tl.maximum(tl.max(tl.abs(quant_2d), axis=1), 1e-4)
506	raw_scales = block_absmax * INV_FP8_MAX
507	exponents = tl.ceil(tl.log2(raw_scales))
508	inv_scales = tl.exp2(-exponents)
509	x_scaled = quant_2d * tl.reshape(inv_scales, (N_QUANT_BLOCKS, 1))
510	x_clamped = tl.clamp(x_scaled, -FP8_MAX, FP8_MAX)
511	x_uint8 = tl.reshape(
512	x_clamped.to(tl.float8e4nv).to(tl.uint8, bitcast=True),
513	(TRITON_BLOCK_SIZE,),
514	)
515	tl.store(fp8_ptr + block, x_uint8, mask=block < NOPE_HEAD_DIM)
516
517	scale_idx = tl.arange(0, N_QUANT_BLOCKS)
518	max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0
519	encoded = tl.maximum(tl.minimum(exponents + 127.0, max_encoded), 0.0)
520	tl.store(
521	scale_ptr + scale_idx, encoded.to(tl.uint8), mask=scale_idx < N_NOPE_BLOCKS
522	)
523	tl.store(scale_ptr + N_NOPE_BLOCKS, tl.zeros((), dtype=tl.uint8))
524
525	NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2
526	NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2
527	even, odd = tl.split(tl.reshape(normed, (NUM_PAIRS, 2)))
528	pair_idx = tl.arange(0, NUM_PAIRS)
529	rope_pair_local = pair_idx - NOPE_PAIRS
530	is_rope_pair = rope_pair_local >= 0
531	cs_idx = tl.maximum(rope_pair_local, 0)
532	compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO
533	cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride
534	cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0)
535	sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0)
536	new_even = even * cos_v - odd * sin_v
537	new_odd = odd * cos_v + even * sin_v
538	result = tl.interleave(new_even, new_odd)
539	if SANITIZE_CACHE_NANS:
540	result = tl.where(result == result, result, 0.0)
541	bf16_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16))
542	rope_local = block - NOPE_HEAD_DIM
543	is_rope = (block >= NOPE_HEAD_DIM) & mask
544	tl.store(bf16_ptr + rope_local, result.to(tl.bfloat16), mask=is_rope)
545
546
547	def _launch_two_stage_sparse_attn_compressor(
548	state_cache: torch.Tensor,
549	token_to_req_indices: torch.Tensor,
550	positions: torch.Tensor,
551	slot_mapping: torch.Tensor,
552	block_table: torch.Tensor,
553	block_size: int,
554	state_width: int,
555	compress_ratio: int,
556	cos_sin_cache: torch.Tensor,
557	kv_cache: torch.Tensor,
558	kv_slot_mapping: torch.Tensor,
559	rms_norm_weight: torch.Tensor,
560	rms_norm_eps: float,
561	quant_block: int,
562	token_stride: int,
563	scale_dim: int,
564	head_dim: int,
565	rope_head_dim: int,
566	num_actual: int,
567	compress_scratch: torch.Tensor,
568	) -> None:
569	num_splits = _pick_compress_num_splits(num_actual, compress_ratio, head_dim)
570	head_tile = head_dim // num_splits
571	scratch = compress_scratch[:num_actual]
572	_compress_gather_split_sparse_attn[(num_actual * num_splits,)](
573	state_cache,
574	state_cache.stride(0),
575	state_cache.stride(1),
576	positions,
577	slot_mapping,
578	token_to_req_indices,
579	block_table,
580	block_table.stride(0),
581	block_size,
582	scratch,
583	scratch.stride(0),
584	HEAD_SIZE=head_dim,
585	STATE_WIDTH=state_width,
586	COMPRESS_RATIO=compress_ratio,
587	NUM_SPLITS=num_splits,
588	HEAD_TILE=head_tile,
589	)
590	_finalize_norm_rope_quant_store_sparse_attn[(num_actual,)](
591	scratch,
592	scratch.stride(0),
593	positions,
594	slot_mapping,
595	rms_norm_weight,
596	rms_norm_eps,
597	cos_sin_cache,
598	cos_sin_cache.stride(0),
599	kv_cache,
600	kv_slot_mapping,
601	kv_cache.shape[1],
602	HEAD_SIZE=head_dim,
603	TRITON_BLOCK_SIZE=triton.next_power_of_2(head_dim),
604	COMPRESS_RATIO=compress_ratio,
605	ROPE_HEAD_DIM=rope_head_dim,
606	FP8_MAX=448.0,
607	QUANT_BLOCK=quant_block,
608	TOKEN_STRIDE=token_stride,
609	SCALE_DIM=scale_dim,
610	KV_BLOCK_STRIDE=kv_cache.stride(0),
611	SANITIZE_CACHE_NANS=_ON_GFX950,
612	)
613
614
615	def compress_norm_rope_store_two_stage_triton(
616	state_cache: torch.Tensor,
617	num_actual: int,
618	token_to_req_indices: torch.Tensor,
619	positions: torch.Tensor,
620	slot_mapping: torch.Tensor,
621	block_table: torch.Tensor,
622	block_size: int,
623	state_width: int,
624	cos_sin_cache: torch.Tensor,
625	kv_cache: torch.Tensor,
626	k_cache_metadata: Any,
627	pdl_kwargs: dict,
628	head_dim: int,
629	rope_head_dim: int,
630	compress_ratio: int,
631	overlap: bool,
632	use_fp4_cache: bool,
633	rms_norm_weight: torch.Tensor,
634	rms_norm_eps: float,
635	quant_block: int,
636	token_stride: int,
637	scale_dim: int,
638	num_decode_tokens: int,
639	compress_scratch: torch.Tensor,
640	) -> None:
641	"""Two-stage split compressor dispatch for head=512 cr>=128 (no-overlap).
642
643	Run the occupancy-fanned two-stage split for prefill [num_decodee_tokens:]
644	to fill the CUs, and use the original single-pass launcher
645	for decode [0, num_decode_tokens)
646	"""
647	num_decodes = min(max(num_decode_tokens, 0), num_actual)
648	num_prefills = num_actual - num_decodes
649	if num_prefills > 0:
650	_launch_two_stage_sparse_attn_compressor(
651	state_cache=state_cache,
652	token_to_req_indices=token_to_req_indices[num_decodes:],
653	positions=positions[num_decodes:],
654	slot_mapping=slot_mapping[num_decodes:],
655	block_table=block_table,
656	block_size=block_size,
657	state_width=state_width,
658	compress_ratio=compress_ratio,
659	cos_sin_cache=cos_sin_cache,
660	kv_cache=kv_cache,
661	kv_slot_mapping=k_cache_metadata.slot_mapping[num_decodes:],
662	rms_norm_weight=rms_norm_weight,
663	rms_norm_eps=rms_norm_eps,
664	quant_block=quant_block,
665	token_stride=token_stride,
666	scale_dim=scale_dim,
667	head_dim=head_dim,
668	rope_head_dim=rope_head_dim,
669	num_actual=num_prefills,
670	compress_scratch=compress_scratch,
671	)
672	if num_decodes > 0:
673	compress_norm_rope_store_triton(
674	state_cache=state_cache,
675	num_actual=num_decodes,
676	token_to_req_indices=token_to_req_indices,
677	positions=positions,
678	slot_mapping=slot_mapping,
679	block_table=block_table,
680	block_size=block_size,
681	state_width=state_width,
682	cos_sin_cache=cos_sin_cache,
683	kv_cache=kv_cache,
684	k_cache_metadata=k_cache_metadata,
685	pdl_kwargs=pdl_kwargs,
686	head_dim=head_dim,
687	rope_head_dim=rope_head_dim,
688	compress_ratio=compress_ratio,
689	overlap=overlap,
690	use_fp4_cache=use_fp4_cache,
691	rms_norm_weight=rms_norm_weight,
692	rms_norm_eps=rms_norm_eps,
693	quant_block=quant_block,
694	token_stride=token_stride,
695	scale_dim=scale_dim,
696	)
697
698
699	# =============================================================================
700	# Indexer path (head=128, all FP8, single quant block)
701	# =============================================================================
702	@triton.jit(
703	do_not_specialize=[
704	"state_cache_stride0",
705	"state_cache_stride1",
706	"block_table_stride",
707	"rms_norm_eps",
708	"cos_sin_stride",
709	]
710	)
711	def _fused_kv_compress_norm_rope_insert_indexer_attn(
712	# ── state cache (compressor internal state) ──
713	state_cache_ptr,
714	state_cache_stride0,
715	state_cache_stride1,
716	# ── metadata ──
717	token_to_req_indices_ptr,
718	positions_ptr,
719	slot_mapping_ptr,
720	block_table_ptr,
721	block_table_stride,
722	block_size,
723	# ── RMSNorm ──
724	rms_norm_weight_ptr,
725	rms_norm_eps,
726	# ── RoPE ──
727	cos_sin_cache_ptr,
728	cos_sin_stride,
729	# ── KV cache output ──
730	k_cache_ptr,
731	kv_slot_mapping_ptr,
732	kv_cache_block_size,
733	# ── constexprs ──
734	HEAD_SIZE: tl.constexpr,
735	TRITON_BLOCK_SIZE: tl.constexpr,
736	STATE_WIDTH: tl.constexpr,
737	COMPRESS_RATIO: tl.constexpr,
738	OVERLAP: tl.constexpr,
739	ROPE_HEAD_DIM: tl.constexpr,
740	FP8_MAX: tl.constexpr, # 448.0
741	QUANT_BLOCK: tl.constexpr, # 128 for indexer
742	TOKEN_STRIDE: tl.constexpr, # 128 for indexer
743	SCALE_DIM: tl.constexpr, # 4 for indexer (1 float32)
744	KV_BLOCK_STRIDE: tl.constexpr,
745	):
746	"""Fused compress → RMSNorm → RoPE → FP8 quant → store.
747
748	One program per token; early-exits for non-boundary positions.
749
750	Cache block layout:
751	[0, bs*128): FP8 data (128 bytes/token)
752	[bs128, +bs4): float32 scales (4 bytes/token)
753
754	For head_dim=128 we have exactly one quant block, so we skip the
755	[N_QUANT_BLOCKS, QUANT_BLOCK] reshape entirely and use a flat
756	``tl.max`` reduction.
757	"""
758	token_idx = tl.program_id(0)
759
760	slot_id = tl.load(slot_mapping_ptr + token_idx)
761	if slot_id < 0:
762	return
763
764	position = tl.load(positions_ptr + token_idx)
765	if (position + 1) % COMPRESS_RATIO != 0:
766	return
767
768	req_idx = tl.load(token_to_req_indices_ptr + token_idx)
769
770	# ── Gather state cache entries ────────────────────────────────────
771	start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1
772	tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO)
773	pos = start + tokens
774	mask_pos = pos >= 0
775
776	block_indices = pos // block_size
777	block_numbers = tl.load(
778	block_table_ptr + req_idx * block_table_stride + block_indices,
779	mask=mask_pos,
780	other=0,
781	)
782	block_offsets = pos % block_size
783	head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE
784
785	block = tl.arange(0, TRITON_BLOCK_SIZE)
786	mask = block < HEAD_SIZE
787	block_numbers_i64 = block_numbers.to(tl.int64)
788
789	row_base = (
790	state_cache_ptr
791	+ block_numbers_i64 * state_cache_stride0
792	+ block_offsets * state_cache_stride1
793	+ head_offset
794	)
795
796	combined_mask = mask_pos[:, None] & mask[None, :]
797
798	score = tl.load(
799	row_base[:, None] + STATE_WIDTH + block[None, :],
800	mask=combined_mask,
801	other=float("-inf"),
802	)
803	score = tl.softmax(score, dim=0)
804
805	kv = tl.load(
806	row_base[:, None] + block[None, :],
807	mask=combined_mask,
808	other=0.0,
809	)
810
811	compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32
812
813	# ── RMSNorm (fp32 throughout) ──────────────────────────────────────
814	rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0)
815	variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_SIZE
816	rrms = tl.rsqrt(variance + rms_norm_eps)
817	normed = compressed_kv * rrms * rms_w
818
819	# ── KV cache pointers ────────────────────────────────────────────
820	kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx)
821	if kv_slot_idx < 0:
822	return
823	kv_block_idx = kv_slot_idx // kv_cache_block_size
824	kv_pos_in_block = kv_slot_idx % kv_cache_block_size
825
826	cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE
827	fp8_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE
828	scale_ptr = (
829	cache_block_ptr
830	+ kv_cache_block_size * TOKEN_STRIDE
831	+ kv_pos_in_block * SCALE_DIM
832	)
833
834	NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM
835	HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2
836
837	# ── Register-based GPT-J forward RoPE in fp32 ─────────────────────
838	NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2
839	NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2
840
841	normed_2d = tl.reshape(normed, (NUM_PAIRS, 2))
842	even, odd = tl.split(normed_2d) # each [NUM_PAIRS] fp32
843
844	pair_idx = tl.arange(0, NUM_PAIRS)
845	rope_pair_local = pair_idx - NOPE_PAIRS
846	is_rope_pair = rope_pair_local >= 0
847	cs_idx = tl.maximum(rope_pair_local, 0)
848
849	compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO
850	cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride
851	cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0)
852	sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0)
853
854	new_even = even * cos_v - odd * sin_v
855	new_odd = odd * cos_v + even * sin_v
856	result = tl.interleave(new_even, new_odd) # fp32
857
858	# ── FP8 UE8M0 quant: single block, flat reduction ────────────────
859	tl.static_assert(
860	TRITON_BLOCK_SIZE == QUANT_BLOCK,
861	"Indexer expects one quant block (QUANT_BLOCK == TRITON_BLOCK_SIZE)",
862	)
863	INV_FP8_MAX: tl.constexpr = 1.0 / FP8_MAX
864
865	result_bf16 = result.to(tl.bfloat16).to(tl.float32)
866	absmax = tl.max(tl.abs(result_bf16), axis=0) # scalar
867	absmax = tl.maximum(absmax, 1e-4)
868	raw_scale = absmax * INV_FP8_MAX
869	exponent = tl.ceil(tl.log2(raw_scale))
870	inv_scale = tl.exp2(-exponent)
871
872	x_scaled = result_bf16 * inv_scale
873	x_clamped = tl.clamp(x_scaled, -FP8_MAX, FP8_MAX)
874	x_fp8 = x_clamped.to(tl.float8e4nv)
875	x_uint8 = x_fp8.to(tl.uint8, bitcast=True)
876
877	tl.store(fp8_ptr + block, x_uint8, mask=mask)
878
879	# Single float32 scale
880	scale_val = tl.exp2(exponent)
881	tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), scale_val)
882
883
884	# =============================================================================
885	# Indexer path (head=128, MXFP4: 2 nibbles/byte + ue8m0 per 32-elem block)
886	# =============================================================================
887	@triton.jit(
888	do_not_specialize=[
889	"state_cache_stride0",
890	"state_cache_stride1",
891	"block_table_stride",
892	"rms_norm_eps",
893	"cos_sin_stride",
894	]
895	)
896	def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn(
897	# ── state cache (compressor internal state) ──
898	state_cache_ptr,
899	state_cache_stride0,
900	state_cache_stride1,
901	# ── metadata ──
902	token_to_req_indices_ptr,
903	positions_ptr,
904	slot_mapping_ptr,
905	block_table_ptr,
906	block_table_stride,
907	block_size,
908	# ── RMSNorm ──
909	rms_norm_weight_ptr,
910	rms_norm_eps,
911	# ── RoPE ──
912	cos_sin_cache_ptr,
913	cos_sin_stride,
914	# ── KV cache output ──
915	k_cache_ptr,
916	kv_slot_mapping_ptr,
917	kv_cache_block_size,
918	# ── constexprs ──
919	HEAD_SIZE: tl.constexpr,
920	TRITON_BLOCK_SIZE: tl.constexpr,
921	STATE_WIDTH: tl.constexpr,
922	COMPRESS_RATIO: tl.constexpr,
923	OVERLAP: tl.constexpr,
924	ROPE_HEAD_DIM: tl.constexpr,
925	FP8_MAX: tl.constexpr, # unused for MXFP4 (kept for signature parity)
926	QUANT_BLOCK: tl.constexpr, # 32 for MXFP4
927	TOKEN_STRIDE: tl.constexpr, # HEAD_SIZE // 2 = 64 packed bytes/token
928	SCALE_DIM: tl.constexpr, # HEAD_SIZE // QUANT_BLOCK = 4 ue8m0 bytes/token
929	KV_BLOCK_STRIDE: tl.constexpr,
930	):
931	"""Fused compress → RMSNorm → RoPE → MXFP4 quant → store.
932
933	One program per token; early-exits for non-boundary positions.
934
935	Cache block layout (``block_size`` tokens per cache block):
936	[0, bs*TOKEN_STRIDE): packed MXFP4 nibbles (2 values/byte)
937	[bsTOKEN_STRIDE, +bsSCALE_DIM): ue8m0 scale bytes (one per 32-elem block)
938
939	MXFP4 format:
940	- E2M1 4-bit values packed two per byte (low nibble first, then high).
941	- Per-32-element block scale = 2^ceil(log2(amax / 6.0)), stored ue8m0
942	(byte = exponent + 127).
943	- Max representable magnitude = 6.0.
944	"""
945	token_idx = tl.program_id(0)
946
947	slot_id = tl.load(slot_mapping_ptr + token_idx)
948	if slot_id < 0:
949	return
950
951	position = tl.load(positions_ptr + token_idx)
952	if (position + 1) % COMPRESS_RATIO != 0:
953	return
954
955	req_idx = tl.load(token_to_req_indices_ptr + token_idx)
956
957	# ── Gather state cache entries ────────────────────────────────────
958	start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1
959	tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO)
960	pos = start + tokens
961	mask_pos = pos >= 0
962
963	block_indices = pos // block_size
964	block_numbers = tl.load(
965	block_table_ptr + req_idx * block_table_stride + block_indices,
966	mask=mask_pos,
967	other=0,
968	)
969	block_offsets = pos % block_size
970	head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE
971
972	block = tl.arange(0, TRITON_BLOCK_SIZE)
973	mask = block < HEAD_SIZE
974	block_numbers_i64 = block_numbers.to(tl.int64)
975
976	row_base = (
977	state_cache_ptr
978	+ block_numbers_i64 * state_cache_stride0
979	+ block_offsets * state_cache_stride1
980	+ head_offset
981	)
982
983	combined_mask = mask_pos[:, None] & mask[None, :]
984
985	score = tl.load(
986	row_base[:, None] + STATE_WIDTH + block[None, :],
987	mask=combined_mask,
988	other=float("-inf"),
989	)
990	score = tl.softmax(score, dim=0)
991
992	kv = tl.load(
993	row_base[:, None] + block[None, :],
994	mask=combined_mask,
995	other=0.0,
996	)
997
998	compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32
999
1000	# ── RMSNorm (fp32 throughout) ──────────────────────────────────────
1001	rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0)
1002	variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_SIZE
1003	rrms = tl.rsqrt(variance + rms_norm_eps)
1004	normed = compressed_kv * rrms * rms_w
1005
1006	# ── KV cache pointers (segregated: values first, then scales) ────
1007	kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx)
1008	if kv_slot_idx < 0:
1009	return
1010	kv_block_idx = kv_slot_idx // kv_cache_block_size
1011	kv_pos_in_block = kv_slot_idx % kv_cache_block_size
1012
1013	cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE
1014	val_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE
1015	scale_ptr = (
1016	cache_block_ptr
1017	+ kv_cache_block_size * TOKEN_STRIDE
1018	+ kv_pos_in_block * SCALE_DIM
1019	)
1020
1021	NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM
1022	HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2
1023
1024	# ── Register-based GPT-J forward RoPE in fp32 ─────────────────────
1025	# We keep the even/odd halves (no tl.interleave afterwards) because the
1026	# MXFP4 per-block absmax / pack naturally operates on (even, odd) pairs.
1027	NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2
1028	NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2
1029
1030	normed_2d = tl.reshape(normed, (NUM_PAIRS, 2))
1031	even, odd = tl.split(normed_2d) # each [NUM_PAIRS] fp32
1032
1033	pair_idx = tl.arange(0, NUM_PAIRS)
1034	rope_pair_local = pair_idx - NOPE_PAIRS
1035	is_rope_pair = rope_pair_local >= 0
1036	cs_idx = tl.maximum(rope_pair_local, 0)
1037
1038	compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO
1039	cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride
1040	cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0)
1041	sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0)
1042
1043	new_even = even * cos_v - odd * sin_v
1044	new_odd = odd * cos_v + even * sin_v
1045
1046	# bf16 roundtrip for parity with reference / Q-side kernel numerics.
1047	new_even = new_even.to(tl.bfloat16).to(tl.float32)
1048	new_odd = new_odd.to(tl.bfloat16).to(tl.float32)
1049
1050	# ── MXFP4 quant: tile even/odd halves into (N_BLOCKS, HALF_BLOCK) ──
1051	# Each MXFP4 block of QUANT_BLOCK elements = HALF_BLOCK consecutive pairs,
1052	# so (N_BLOCKS, HALF_BLOCK) rows of even/odd each land exactly one block.
1053	N_QUANT_BLOCKS: tl.constexpr = HEAD_SIZE // QUANT_BLOCK
1054	HALF_BLOCK: tl.constexpr = QUANT_BLOCK // 2
1055	tl.static_assert(TRITON_BLOCK_SIZE == HEAD_SIZE)
1056	tl.static_assert(HEAD_SIZE % QUANT_BLOCK == 0)
1057	tl.static_assert(TOKEN_STRIDE == HEAD_SIZE // 2)
1058	tl.static_assert(SCALE_DIM == N_QUANT_BLOCKS)
1059
1060	even_2d = tl.reshape(new_even, (N_QUANT_BLOCKS, HALF_BLOCK))
1061	odd_2d = tl.reshape(new_odd, (N_QUANT_BLOCKS, HALF_BLOCK))
1062
1063	amax = tl.maximum(
1064	tl.max(tl.abs(even_2d), axis=1),
1065	tl.max(tl.abs(odd_2d), axis=1),
1066	)
1067	amax = tl.maximum(amax, 6.0 * (2**-126))
1068
1069	# ue8m0 block scale: 2^ceil(log2(amax / 6.0)), stored as (exp + 127) byte.
1070	log2_ratio = tl.ceil(tl.log2(amax * (1.0 / 6.0)))
1071	log2_ratio = tl.minimum(tl.maximum(log2_ratio, -127.0), 127.0)
1072	inv_scale = tl.exp2(-log2_ratio)
1073	ue8m0 = (log2_ratio + 127.0).to(tl.uint8) # [N_QUANT_BLOCKS]
1074
1075	inv_scale_col = tl.reshape(inv_scale, (N_QUANT_BLOCKS, 1))
1076	packed = _fp32x2_to_fp4x2(
1077	even_2d * inv_scale_col, odd_2d * inv_scale_col
1078	) # (N_BLOCKS, HALF_BLOCK) uint8
1079	packed_flat = tl.reshape(packed, (TOKEN_STRIDE,))
1080
1081	tl.store(val_ptr + tl.arange(0, TOKEN_STRIDE), packed_flat)
1082	tl.store(scale_ptr + tl.arange(0, SCALE_DIM), ue8m0)
1083
1084
1085	class FusedKVCompressNormRopeInsertIndexerTritonKernel(
1086	VllmTritonJitKernel["FusedKVCompressNormRopeInsertIndexerTritonKernel.CompileKey"]
1087	):
1088	@dataclass(frozen=True)
1089	class CompileKey:
1090	dtype: torch.dtype
1091	use_fp4_cache: bool
1092	head_size: int
1093	triton_block_size: int
1094	state_width: int
1095	compress_ratio: int
1096	overlap: bool
1097	rope_head_dim: int
1098	fp8_max: float
1099	quant_block: int
1100	token_stride: int
1101	scale_dim: int
1102	block_size: int
1103	kv_cache_block_size: int
1104	kv_block_stride: int
1105
1106	# Canonical kernel for the launcher's arg_names mapping; the launched
1107	# variant (fp8/mxfp4, same arg_names) is picked per call in __call__.
1108	kernel = staticmethod(_fused_kv_compress_norm_rope_insert_indexer_attn)
1109
1110	def dispatch( # type: ignore[override]
1111	self,
1112	*,
1113	dtype: torch.dtype,
1114	use_fp4_cache: bool,
1115	head_dim: int,
1116	rope_head_dim: int,
1117	compress_ratio: int,
1118	cache_block_size: int,
1119	cache_alignment: int,
1120	runtime_state_width: int \| None = None,
1121	runtime_quant_block: int \| None = None,
1122	runtime_token_stride: int \| None = None,
1123	runtime_scale_dim: int \| None = None,
1124	runtime_kv_block_stride: int \| None = None,
1125	) -> CompileKey:
1126	overlap = compress_ratio == 4
1127	default_quant_block = 32 if use_fp4_cache else 128
1128	quant_block = (
1129	runtime_quant_block
1130	if runtime_quant_block is not None
1131	else default_quant_block
1132	)
1133	default_token_stride = head_dim // 2 if use_fp4_cache else head_dim
1134	token_stride = (
1135	runtime_token_stride
1136	if runtime_token_stride is not None
1137	else default_token_stride
1138	)
1139	default_scale_dim = head_dim // quant_block if use_fp4_cache else 4
1140	scale_dim = (
1141	runtime_scale_dim if runtime_scale_dim is not None else default_scale_dim
1142	)
1143	raw_kv_cache_block_size = cache_block_size // compress_ratio
1144	kv_cache_block_size = (
1145	raw_kv_cache_block_size if raw_kv_cache_block_size >= 1 else 1
1146	)
1147	state_block_size = (
1148	4
1149	if compress_ratio == 4
1150	else 8
1151	if compress_ratio == 128
1152	else cache_block_size
1153	)
1154	default_kv_block_stride = round_up(
1155	kv_cache_block_size * (token_stride + scale_dim),
1156	cache_alignment,
1157	)
1158	kv_block_stride = (
1159	runtime_kv_block_stride
1160	if runtime_kv_block_stride is not None
1161	else default_kv_block_stride
1162	)
1163	state_width = (
1164	runtime_state_width
1165	if runtime_state_width is not None
1166	else head_dim * (1 + overlap)
1167	)
1168	return self.CompileKey(
1169	dtype=dtype,
1170	use_fp4_cache=use_fp4_cache,
1171	head_size=head_dim,
1172	triton_block_size=triton.next_power_of_2(head_dim),
1173	state_width=state_width,
1174	compress_ratio=compress_ratio,
1175	overlap=overlap,
1176	rope_head_dim=rope_head_dim,
1177	fp8_max=448.0,
1178	quant_block=quant_block,
1179	token_stride=token_stride,
1180	scale_dim=scale_dim,
1181	block_size=state_block_size,
1182	kv_cache_block_size=kv_cache_block_size,
1183	kv_block_stride=kv_block_stride,
1184	)
1185
1186	def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
1187	hf_config = vllm_config.model_config.hf_config
1188	if hf_config is None:
1189	return []
1190
1191	cache_block_size = vllm_config.cache_config.block_size
1192	from vllm.v1.attention.backends.mla.indexer import (
1193	dsa_indexer_uses_fp4,
1194	)
1195
1196	use_fp4_cache = dsa_indexer_uses_fp4(vllm_config)
1197	head_dim = int(getattr(hf_config, "index_head_dim", 0) or 0)
1198	rope_head_dim = int(getattr(hf_config, "qk_rope_head_dim", 0) or 0)
1199	if head_dim <= 0 or rope_head_dim <= 0 or cache_block_size <= 0:
1200	return []
1201
1202	cache_dtype = vllm_config.cache_config.cache_dtype
1203	cache_alignment = 576 if cache_dtype == "fp8_ds_mla" else 512
1204	return self._trace_dispatch(self.dispatch)(
1205	dtype=vllm_config.model_config.dtype,
1206	use_fp4_cache=use_fp4_cache,
1207	head_dim=head_dim,
1208	rope_head_dim=rope_head_dim,
1209	compress_ratio=4,
1210	cache_block_size=cache_block_size,
1211	cache_alignment=cache_alignment,
1212	)
1213
1214	def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
1215	fp32_ptr = TritonWarmupTensor(torch.float32, shape=(1, 1))
1216	int32_ptr = TritonWarmupTensor(torch.int32)
1217	int64_ptr = TritonWarmupTensor(torch.int64)
1218	# kv_cache supplies kv_cache_block_size (shape[1]) and KV_BLOCK_STRIDE
1219	# (stride(0)); reproduce both so warmup and launch specialize identically.
1220	kv_cache = TritonWarmupTensor(
1221	torch.uint8,
1222	shape=(1, compile_key.kv_cache_block_size, 1),
1223	strides=(compile_key.kv_block_stride, 1, 1),
1224	)
1225	return dict(
1226	state_cache=fp32_ptr,
1227	num_actual=1,
1228	token_to_req_indices=int32_ptr,
1229	positions=int64_ptr,
1230	slot_mapping=int64_ptr,
1231	block_table=int32_ptr,
1232	block_size=compile_key.block_size,
1233	state_width=compile_key.state_width,
1234	cos_sin_cache=fp32_ptr,
1235	kv_cache=kv_cache,
1236	k_cache_metadata=SimpleNamespace(slot_mapping=int64_ptr),
1237	pdl_kwargs={},
1238	head_dim=compile_key.head_size,
1239	rope_head_dim=compile_key.rope_head_dim,
1240	compress_ratio=compile_key.compress_ratio,
1241	overlap=compile_key.overlap,
1242	use_fp4_cache=compile_key.use_fp4_cache,
1243	rms_norm_weight=TritonWarmupTensor(
1244	compile_key.dtype,
1245	shape=(compile_key.head_size,),
1246	),
1247	rms_norm_eps=1e-6,
1248	quant_block=compile_key.quant_block,
1249	token_stride=compile_key.token_stride,
1250	scale_dim=compile_key.scale_dim,
1251	)
1252
1253	@kernel_launcher
1254	def __call__(
1255	self,
1256	*,
1257	state_cache: torch.Tensor,
1258	num_actual: int,
1259	token_to_req_indices: torch.Tensor,
1260	positions: torch.Tensor,
1261	slot_mapping: torch.Tensor,
1262	block_table: torch.Tensor,
1263	block_size: int,
1264	state_width: int,
1265	cos_sin_cache: torch.Tensor,
1266	kv_cache: torch.Tensor,
1267	k_cache_metadata: Any,
1268	pdl_kwargs: dict,
1269	head_dim: int,
1270	rope_head_dim: int,
1271	compress_ratio: int,
1272	overlap: bool,
1273	use_fp4_cache: bool,
1274	rms_norm_weight: torch.Tensor,
1275	rms_norm_eps: float,
1276	quant_block: int,
1277	token_stride: int,
1278	scale_dim: int,
1279	) -> LaunchSpec:
1280	return (num_actual,), dict(
1281	kernel=(
1282	_fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
1283	if use_fp4_cache
1284	else _fused_kv_compress_norm_rope_insert_indexer_attn
1285	),
1286	state_cache_stride0=state_cache.stride(0),
1287	state_cache_stride1=state_cache.stride(1),
1288	block_table_stride=block_table.stride(0),
1289	cos_sin_stride=cos_sin_cache.stride(0),
1290	k_cache_ptr=kv_cache,
1291	kv_slot_mapping_ptr=k_cache_metadata.slot_mapping,
1292	kv_cache_block_size=kv_cache.shape[1],
1293	HEAD_SIZE=head_dim,
1294	TRITON_BLOCK_SIZE=triton.next_power_of_2(head_dim),
1295	STATE_WIDTH=state_width,
1296	COMPRESS_RATIO=compress_ratio,
1297	OVERLAP=overlap,
1298	ROPE_HEAD_DIM=rope_head_dim,
1299	FP8_MAX=448.0,
1300	QUANT_BLOCK=quant_block,
1301	TOKEN_STRIDE=token_stride,
1302	SCALE_DIM=scale_dim,
1303	KV_BLOCK_STRIDE=kv_cache.stride(0),
1304	num_warps=1,
1305	**pdl_kwargs,
1306	)
1307
1308
1309	_FUSED_KV_COMPRESS_NORM_ROPE_INSERT_INDEXER_TRITON_KERNEL = (
1310	FusedKVCompressNormRopeInsertIndexerTritonKernel()
1311	)