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

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 1,335 lines 52.3 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3	"""Unified (prefill + decode) Triton attention kernel for TurboQuant.
4
5	Structure is ported from ``vllm/v1/attention/ops/triton_unified_attention.py``
6	(the AITER unified attention kernel already upstreamed into vLLM). Only the K
7	and V load sites are replaced: instead of reading raw fp16 keys/values from
8	two contiguous caches, this kernel reads TurboQuant-packed bytes from a single
9	combined cache and dequantizes on the fly inside the tile loop.
10
11	Benefits over the v1/v2 decode-only kernels:
12
13	1. GQA heads are stacked into ``BLOCK_M`` and the Q·K and P·V ops are proper
14	``tl.dot`` tensor-core operations (MFMA on MI300X).
15	2. The same kernel handles decode (``BLOCK_Q=1``) and prefill
16	(``BLOCK_Q>1``) -- no more Python per-request for-loop for continuation
17	chunks.
18	3. Only the subset of features exercised by the current TQ paths is kept
19	(causal, GQA). Sinks / softcap / ALiBi / sliding-window / qq-bias /
20	mm-prefix are deferred to follow-ups.
21
22	This is an opt-in v3 path behind ``VLLM_TQ_DECODE_V3``.
23	"""
24
25	from __future__ import annotations
26
27	import math
28	from typing import Any
29
30	import torch
31
32	from vllm.platforms import current_platform
33	from vllm.triton_utils import tl, triton
34
35	# reduce_segments is KV-format-agnostic: by the time it runs, K/V have been
36	# consumed and only (max, expsum, partial_output) triples remain. Reuse the
37	# baseline's implementation verbatim to avoid code duplication.
38	from vllm.v1.attention.ops.triton_unified_attention import reduce_segments
39
40	from .triton_turboquant_decode import _use_fp8_e4b15
41	from .triton_turboquant_decode_v2 import build_pair_lut
42
43	_is_hip = current_platform.is_rocm()
44
45
46	# ---------------------------------------------------------------------------
47	# Helper: find which sequence in the batch a global q-block index belongs to.
48	# Identical to the unified kernel's binary search.
49	# ---------------------------------------------------------------------------
50
51
52	@triton.jit
53	def _find_seq_idx(
54	query_start_len_ptr,
55	target_idx,
56	num_seqs,
57	BLOCK_Q: tl.constexpr,
58	use_q_block_mode: tl.constexpr,
59	):
60	left: tl.int32 = 0
61	right = num_seqs
62	while left < right:
63	mid = (left + right) // 2
64	val = tl.load(query_start_len_ptr + mid)
65	mid_val = val // BLOCK_Q + mid if use_q_block_mode else val
66	if mid_val <= target_idx:
67	left = mid + 1
68	else:
69	right = mid
70	return left - 1
71
72
73	# ---------------------------------------------------------------------------
74	# Fused Q-rotation prologue (Opt: eliminates the launcher-side
75	# q_rot = (query.float() @ PiT).to(dtype).contiguous()
76	# chain, which cost ~50-60 us/step at D=64 — all launch-bound, not compute-
77	# bound. The fused version does one small MFMA inside the attention kernel,
78	# once per program, before the KV-tile loop. PiT is [D,D] fp32 and fits in
79	# L1 easily (16KB at D=64, 64KB at D=128).
80	#
81	# Algebra is identical to the launcher path: both compute Q_rot = Q @ PiT
82	# with fp32 accumulate. Differences are rounding-order only (MFMA tiling
83	# vs rocBLAS tiling), well within the 4-bit KV quant noise floor.
84	# ---------------------------------------------------------------------------
85
86
87	@triton.jit
88	def _tq_fuse_q_rotation(
89	Q, # [BLOCK_M, HEAD_SIZE_PADDED] — raw query in Q.dtype
90	PiT_ptr,
91	PiT_stride_0: tl.int64,
92	PiT_stride_1: tl.int64,
93	dim_mask, # [HEAD_SIZE_PADDED] int1 — valid head dims
94	HEAD_SIZE_PADDED: tl.constexpr,
95	):
96	"""Fused Q @ PiT prologue. Called once per program for the MSE-key path
97	when the launcher has passed the raw (un-rotated) query. Returns the
98	rotated Q in the original dtype.
99	"""
100	d_offs = tl.arange(0, HEAD_SIZE_PADDED)
101	pit_offsets = d_offs[:, None] * PiT_stride_0 + d_offs[None, :] * PiT_stride_1
102	pit_mask = dim_mask[:, None] & dim_mask[None, :]
103	PiT_tile = tl.load(PiT_ptr + pit_offsets, mask=pit_mask, other=0.0).to(tl.float32)
104	# input_precision="ieee" pins both inputs to full fp32 MFMA (no TF32
105	# truncation). allow_tf32 is intentionally omitted — Triton rejects
106	# passing both. This matches the launcher's rocBLAS fp32 GEMM in
107	# algebra; rounding-order differs, diff is <= a few fp16 ulp.
108	Q_rot = tl.dot(
109	Q.to(tl.float32),
110	PiT_tile,
111	input_precision="ieee",
112	)
113	return Q_rot.to(Q.dtype)
114
115
116	# ---------------------------------------------------------------------------
117	# TQ K-tile dequant: returns K_T : [HEAD_SIZE_PADDED, TILE_SIZE] in Q.dtype
118	# ready for tl.dot(Q, K_T).
119	#
120	# Packed slot layout (per (block, pos, kv_head)) is the same as v1/v2:
121	# MSE path: [ key_idx_bytes (MSE_BYTES) \| key_norm_fp16 (2B) \| ...V... ]
122	# FP8 path: [ fp8_key_bytes (D) \| ...V... ]
123	# ---------------------------------------------------------------------------
124
125
126	@triton.jit
127	def _tq_load_k_tile(
128	KV_cache_ptr,
129	KV_cache_u16_ptr, # uint16 view of KV_cache (same storage)
130	data_bases, # [TILE_SIZE] int64 — byte offset to each token's DATA region
131	knorm_u16_addrs, # [TILE_SIZE] int64 — u16 element index for each token's K-norm
132	d_offs, # [HEAD_SIZE_PADDED]
133	d_mask, # [HEAD_SIZE_PADDED] int1
134	tile_mask, # [TILE_SIZE] int1
135	Centroids_ptr,
136	Pair_lut_ptr,
137	OUT_DTYPE: tl.constexpr, # tl.float16 or tl.bfloat16
138	HEAD_DIM: tl.constexpr,
139	BLOCK_D: tl.constexpr, # HEAD_SIZE_PADDED — needed for pair-LUT reshape
140	MSE_BITS: tl.constexpr,
141	N_CENTROIDS: tl.constexpr,
142	KEY_FP8: tl.constexpr,
143	USE_PAIR_LUT: tl.constexpr,
144	NORM_CORRECTION: tl.constexpr,
145	FP8_E4B15: tl.constexpr,
146	TILE_SIZE: tl.constexpr,
147	):
148	"""Load + dequantize a TILE_SIZE × HEAD_SIZE block of keys and return
149	the transposed tile K_T : [HEAD_SIZE_PADDED, TILE_SIZE].
150
151	Opt#3 SoA layout: packed K data is at `data_bases[t] + [0, MSE_BYTES)`
152	for MSE keys (or `[0, D)` for FP8). The per-token K-norm lives in the
153	per-block SoA metadata region; `knorm_u16_addrs` already encodes its
154	u16 element index. For decode tiles aligned with blocks, these addresses
155	are contiguous → one coalesced wide load replaces TILE_SIZE scattered
156	loads (the whole point of Opt#3).
157	"""
158	if KEY_FP8:
159	k_addrs = data_bases[:, None] + d_offs[None, :]
160	k_raw = tl.load(
161	KV_cache_ptr + k_addrs,
162	mask=tile_mask[:, None] & d_mask[None, :],
163	other=0,
164	)
165	if FP8_E4B15:
166	k_f32 = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32)
167	else:
168	k_f32 = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32)
169	K = k_f32 # [TILE_SIZE, HEAD_SIZE_PADDED]
170	else:
171	# MSE path: gather packed key indices + centroid LUT.
172	if MSE_BITS == 4 and USE_PAIR_LUT:
173	# FLUTE pair-LUT fast path — load each packed byte once, decode
174	# both nibbles, single 3-D gather returns (T[lo], T[hi]) per byte.
175	HALF_D: tl.constexpr = BLOCK_D // 2
176	half_offs = tl.arange(0, HALF_D)
177	byte_mask = (half_offs * 2) < HEAD_DIM
178	byte_addrs = data_bases[:, None] + half_offs[None, :]
179	byte_raw = tl.load(
180	KV_cache_ptr + byte_addrs,
181	mask=tile_mask[:, None] & byte_mask[None, :],
182	other=0,
183	).to(tl.int32)
184	lo_idx = byte_raw & 0xF
185	hi_idx = (byte_raw >> 4) & 0xF
186	pair_key = lo_idx * N_CENTROIDS + hi_idx
187	pair_slot = tl.arange(0, 2)
188	c_pair = tl.load(
189	Pair_lut_ptr + pair_key[:, :, None] * 2 + pair_slot[None, None, :],
190	mask=(tile_mask[:, None, None] & byte_mask[None, :, None]),
191	other=0.0,
192	)
193	c_vals = tl.reshape(c_pair, [TILE_SIZE, BLOCK_D])
194	elif MSE_BITS == 4:
195	half_idx = d_offs // 2
196	nibble_shift = (d_offs % 2) * 4
197	mse_addrs = data_bases[:, None] + half_idx[None, :]
198	mse_raw = tl.load(
199	KV_cache_ptr + mse_addrs,
200	mask=tile_mask[:, None] & d_mask[None, :],
201	other=0,
202	).to(tl.int32)
203	mse_idx = (mse_raw >> nibble_shift[None, :]) & 0xF
204	c_vals = tl.load(
205	Centroids_ptr + mse_idx,
206	mask=tile_mask[:, None] & d_mask[None, :],
207	other=0.0,
208	)
209	else:
210	# Generic bit extraction (3-bit, etc.)
211	mse_bit_off = d_offs * MSE_BITS
212	mse_byte_idx = mse_bit_off // 8
213	mse_bit_shift = mse_bit_off % 8
214	mse_mask_val = (1 << MSE_BITS) - 1
215	mse_addrs0 = data_bases[:, None] + mse_byte_idx[None, :]
216	mse_raw0 = tl.load(
217	KV_cache_ptr + mse_addrs0,
218	mask=tile_mask[:, None] & d_mask[None, :],
219	other=0,
220	).to(tl.int32)
221	mse_raw1 = tl.load(
222	KV_cache_ptr + mse_addrs0 + 1,
223	mask=tile_mask[:, None] & d_mask[None, :],
224	other=0,
225	).to(tl.int32)
226	raw16 = mse_raw0 \| (mse_raw1 << 8)
227	mse_idx = (raw16 >> mse_bit_shift[None, :]) & mse_mask_val
228	c_vals = tl.load(
229	Centroids_ptr + mse_idx,
230	mask=tile_mask[:, None] & d_mask[None, :],
231	other=0.0,
232	)
233
234	# Opt#1: 1/\|\|c_vec\|\| is pre-folded into the stored K-norm at store
235	# time, so the kernel doesn't recompute norm-correction here.
236	# Opt#3: K-norms for a tile are contiguous in the per-block SoA
237	# region when the tile lies within one block (always true for
238	# aligned decode tiles with TILE_SIZE == BLOCK_SIZE) — one coalesced
239	# u16 load replaces TILE_SIZE scattered 2-byte loads.
240	norm_u16 = tl.load(KV_cache_u16_ptr + knorm_u16_addrs, mask=tile_mask, other=0)
241	vec_norms = norm_u16.to(tl.float16, bitcast=True).to(tl.float32)
242	K = c_vals * vec_norms[:, None] # [TILE_SIZE, HEAD_SIZE_PADDED]
243
244	K_T = tl.trans(K.to(OUT_DTYPE)) # [HEAD_SIZE_PADDED, TILE_SIZE]
245	_ = HEAD_DIM
246	return K_T
247
248
249	# ---------------------------------------------------------------------------
250	# TQ V-tile dequant: returns V : [TILE_SIZE, HEAD_SIZE_PADDED] in Q.dtype
251	# ready for tl.dot(P, V).
252	# ---------------------------------------------------------------------------
253
254
255	@triton.jit
256	def _tq_load_v_tile(
257	KV_cache_ptr,
258	KV_cache_u16_ptr, # uint16 view of KV_cache (same storage)
259	val_bases, # [TILE_SIZE] int64 — byte offset to each token's V-data
260	vscale_u16_addrs, # [TILE_SIZE] int64 — u16 element index for V-scale
261	vzero_u16_addrs, # [TILE_SIZE] int64 — u16 element index for V-zero
262	d_offs,
263	d_mask,
264	tile_mask,
265	OUT_DTYPE: tl.constexpr,
266	HEAD_DIM: tl.constexpr,
267	VQB: tl.constexpr,
268	):
269	"""Load + dequantize a TILE_SIZE × HEAD_SIZE block of values.
270
271	Opt#3 SoA layout: packed V data is at `val_bases[t] + [0, VAL_DATA_BYTES)`
272	(V-data immediately follows K-data within the slot's data region, and
273	`val_bases` is precomputed by the caller as `data_base + KEY_DATA_BYTES`).
274	V-scale / V-zero live in the per-block SoA metadata region at indices
275	`vscale_u16_addrs` / `vzero_u16_addrs`. For tiles aligned to block
276	boundaries, those addresses are contiguous → one coalesced wide load
277	per field instead of TILE_SIZE scattered 2-byte loads.
278	"""
279	if VQB == 4:
280	vb_idx = d_offs // 2
281	vb_shift = (d_offs % 2) * 4
282	val_addrs = val_bases[:, None] + vb_idx[None, :]
283	val_raw = tl.load(
284	KV_cache_ptr + val_addrs,
285	mask=tile_mask[:, None] & d_mask[None, :],
286	other=0,
287	).to(tl.int32)
288	v_idx = ((val_raw >> vb_shift[None, :]) & 0xF).to(tl.float32)
289	else: # VQB == 3
290	val_bit_off = d_offs * 3
291	val_byte_idx = val_bit_off // 8
292	val_bit_shift = val_bit_off % 8
293	val_addrs0 = val_bases[:, None] + val_byte_idx[None, :]
294	val_raw0 = tl.load(
295	KV_cache_ptr + val_addrs0,
296	mask=tile_mask[:, None] & d_mask[None, :],
297	other=0,
298	).to(tl.int32)
299	val_raw1 = tl.load(
300	KV_cache_ptr + val_addrs0 + 1,
301	mask=tile_mask[:, None] & d_mask[None, :],
302	other=0,
303	).to(tl.int32)
304	raw16 = val_raw0 \| (val_raw1 << 8)
305	v_idx = ((raw16 >> val_bit_shift[None, :]) & 0x7).to(tl.float32)
306
307	# SoA scale / zero loads — coalesced on aligned decode tiles.
308	scale_u16 = tl.load(KV_cache_u16_ptr + vscale_u16_addrs, mask=tile_mask, other=0)
309	zero_u16 = tl.load(KV_cache_u16_ptr + vzero_u16_addrs, mask=tile_mask, other=0)
310	v_scales = scale_u16.to(tl.float16, bitcast=True).to(tl.float32)
311	v_zeros = zero_u16.to(tl.float16, bitcast=True).to(tl.float32)
312
313	V = v_idx * v_scales[:, None] + v_zeros[:, None] # [TILE_SIZE, HEAD_SIZE_PADDED]
314	_ = HEAD_DIM
315	return V.to(OUT_DTYPE)
316
317
318	# ---------------------------------------------------------------------------
319	# Unified 2D attention kernel over a TurboQuant KV cache (causal only, GQA).
320	# Structure follows ``kernel_unified_attention_2d`` from
321	# triton_unified_attention.py with:
322	# * Q-head GQA stacking into BLOCK_M
323	# * tensor-core tl.dot(Q, K) and tl.dot(P, V)
324	# * causal + query-padding masking
325	# * online softmax
326	# and K/V loads replaced by the TQ helpers above.
327	# ---------------------------------------------------------------------------
328
329
330	@triton.jit
331	def kernel_tq_unified_attention_2d(
332	output_ptr, # [num_tokens, Hq, D]
333	query_ptr, # [num_tokens, Hq, D] — raw if FUSE_Q_ROT else rotated
334	KV_cache_ptr, # [num_blocks, block_size, Hk, padded_slot] uint8 - TQ packed
335	KV_cache_u16_ptr, # uint16 view of same storage — Opt#2 wide metadata loads
336	Centroids_ptr, # [n_centroids] fp32
337	Pair_lut_ptr, # [N*N, 2] fp32 (pair-LUT for 4-bit MSE) or dummy
338	PiT_ptr, # [D, D] fp32 — only dereferenced when FUSE_Q_ROT == 1
339	block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]
340	seq_lens_ptr, # [num_seqs]
341	query_start_len_ptr, # [num_seqs+1]
342	sinks_ptr, # [Hq] fp32 — per-head sink logits; dereferenced only when USE_SINKS
343	scale, # float32
344	num_query_heads: tl.constexpr,
345	num_queries_per_kv: tl.constexpr,
346	block_table_stride: tl.int64,
347	query_stride_0: tl.int64,
348	query_stride_1: tl.int64,
349	output_stride_0: tl.int64,
350	output_stride_1: tl.int64,
351	stride_cache_block: tl.int64,
352	pit_stride_0: tl.int64,
353	pit_stride_1: tl.int64,
354	BLOCK_SIZE: tl.constexpr,
355	TILE_SIZE: tl.constexpr,
356	HEAD_SIZE: tl.constexpr,
357	HEAD_SIZE_PADDED: tl.constexpr,
358	BLOCK_Q: tl.constexpr,
359	BLOCK_M: tl.constexpr,
360	num_seqs: tl.int32,
361	# TQ layout constants
362	MSE_BITS: tl.constexpr,
363	MSE_BYTES: tl.constexpr,
364	VQB: tl.constexpr,
365	VAL_DATA_BYTES: tl.constexpr,
366	N_CENTROIDS: tl.constexpr,
367	KEY_FP8: tl.constexpr,
368	USE_PAIR_LUT: tl.constexpr,
369	# Opt#3 SoA layout constants
370	NUM_KV_HEADS: tl.constexpr,
371	KEY_DATA_BYTES: tl.constexpr, # MSE_BYTES for MSE, HEAD_SIZE for FP8
372	META_REGION_OFFSET: tl.constexpr, # bytes: bs * H * (KD+VD)
373	NUM_SOA_FIELDS: tl.constexpr, # 3 for MSE, 2 for FP8
374	SOA_K_NORM: tl.constexpr, # 0 for MSE, unused for FP8
375	SOA_V_SCALE: tl.constexpr, # 1 (MSE) / 0 (FP8)
376	SOA_V_ZERO: tl.constexpr, # 2 (MSE) / 1 (FP8)
377	NORM_CORRECTION: tl.constexpr = 0,
378	FP8_E4B15: tl.constexpr = 0,
379	FUSE_Q_ROT: tl.constexpr = 0, # Fused Q @ PiT prologue when 1 (MSE path)
380	USE_SINKS: tl.constexpr = 0, # Per-head sink logit folded into softmax denom
381	):
382	q_block_global_idx = tl.program_id(0)
383	kv_head_idx = tl.program_id(1)
384
385	seq_idx = _find_seq_idx(
386	query_start_len_ptr, q_block_global_idx, num_seqs, BLOCK_Q, True
387	)
388	q_block_start_idx = tl.load(query_start_len_ptr + seq_idx) // BLOCK_Q + seq_idx
389	q_block_local_idx = q_block_global_idx - q_block_start_idx
390
391	cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx)
392	cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + 1)
393	cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index
394
395	if q_block_local_idx * BLOCK_Q >= cur_batch_query_len:
396	return
397
398	offs_m = tl.arange(0, BLOCK_M)
399	offs_d = tl.arange(0, HEAD_SIZE_PADDED)
400	offs_t = tl.arange(0, TILE_SIZE)
401	query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv
402
403	query_offset_0 = cur_batch_in_all_start_index + query_pos
404	query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv
405	query_offset = (
406	query_offset_0[:, None] * query_stride_0
407	+ query_offset_1[:, None] * query_stride_1
408	+ offs_d[None, :]
409	)
410
411	dim_mask = tl.where(offs_d < HEAD_SIZE, 1, 0).to(tl.int1)
412	query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1)
413	query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1)
414
415	# Q : [BLOCK_M, HEAD_SIZE_PADDED]
416	Q = tl.load(
417	query_ptr + query_offset,
418	mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
419	other=0.0,
420	)
421
422	if FUSE_Q_ROT:
423	Q = _tq_fuse_q_rotation(
424	Q,
425	PiT_ptr,
426	pit_stride_0,
427	pit_stride_1,
428	dim_mask,
429	HEAD_SIZE_PADDED,
430	)
431
432	block_table_offset = seq_idx * block_table_stride
433
434	# Online softmax state.
435	# When USE_SINKS=1 we fold a per-head sink logit s_h into the softmax
436	# denominator by initializing M = s_h, L = 1.0 (== exp(s_h - s_h)).
437	# The standard online-softmax update then correctly rescales the sink
438	# contribution as real tokens update M, and the final acc / L gives
439	# out_h = sum_i exp(q.k_i)*V_i / (exp(s_h) + sum_i exp(q.k_i)).
440	# Masked rows (query_mask_1 == 0) get -inf so the sink cannot pollute
441	# invalid (q_token, q_head) positions outside the real query range.
442	if USE_SINKS:
443	M = tl.load(
444	sinks_ptr + query_offset_1, mask=query_mask_1, other=float("-inf")
445	).to(tl.float32)
446	else:
447	M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
448	L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
449	acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32)
450
451	seq_len = tl.load(seq_lens_ptr + seq_idx)
452	context_len = seq_len - cur_batch_query_len
453
454	# Longest key prefix any query row in this q-block can attend to.
455	max_seq_prefix_len = (
456	context_len
457	+ q_block_local_idx * BLOCK_Q
458	+ (BLOCK_M - 1) // num_queries_per_kv
459	+ 1
460	)
461	max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len)
462	num_tiles = tl.cdiv(max_seq_prefix_len, TILE_SIZE)
463
464	for j in range(0, num_tiles):
465	seq_offset = j * TILE_SIZE + offs_t
466	tile_mask = seq_offset < max_seq_prefix_len
467
468	physical_block_idx = tl.load(
469	block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE
470	).to(tl.int64)
471
472	# Opt#3 SoA addressing: compute data + SoA metadata addresses once
473	# per tile; the load helpers below issue coalesced wide loads from
474	# these addresses for decode tiles that align with block boundaries.
475	slot_within_block = (seq_offset % BLOCK_SIZE).to(tl.int64)
476	block_base = physical_block_idx * stride_cache_block
477	DATA_BYTES_PER_SLOT: tl.constexpr = KEY_DATA_BYTES + VAL_DATA_BYTES
478	data_bases = (
479	block_base
480	+ slot_within_block * (NUM_KV_HEADS * DATA_BYTES_PER_SLOT)
481	+ tl.cast(kv_head_idx, tl.int64) * DATA_BYTES_PER_SLOT
482	)
483	val_bases = data_bases + KEY_DATA_BYTES
484
485	head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + tl.cast(
486	kv_head_idx, tl.int64
487	) * (NUM_SOA_FIELDS * BLOCK_SIZE)
488	knorm_u16_addrs = (
489	head_meta_u16_base + SOA_K_NORM * BLOCK_SIZE + slot_within_block
490	)
491	vscale_u16_addrs = (
492	head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + slot_within_block
493	)
494	vzero_u16_addrs = (
495	head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + slot_within_block
496	)
497
498	# ---- K : [HEAD_SIZE_PADDED, TILE_SIZE] in Q.dtype ----
499	K_T = _tq_load_k_tile(
500	KV_cache_ptr,
501	KV_cache_u16_ptr,
502	data_bases,
503	knorm_u16_addrs,
504	offs_d,
505	dim_mask,
506	tile_mask,
507	Centroids_ptr,
508	Pair_lut_ptr,
509	OUT_DTYPE=Q.dtype,
510	HEAD_DIM=HEAD_SIZE,
511	BLOCK_D=HEAD_SIZE_PADDED,
512	MSE_BITS=MSE_BITS,
513	N_CENTROIDS=N_CENTROIDS,
514	KEY_FP8=KEY_FP8,
515	USE_PAIR_LUT=USE_PAIR_LUT,
516	NORM_CORRECTION=NORM_CORRECTION,
517	FP8_E4B15=FP8_E4B15,
518	TILE_SIZE=TILE_SIZE,
519	)
520
521	# ---- V : [TILE_SIZE, HEAD_SIZE_PADDED] in Q.dtype ----
522	V = _tq_load_v_tile(
523	KV_cache_ptr,
524	KV_cache_u16_ptr,
525	val_bases,
526	vscale_u16_addrs,
527	vzero_u16_addrs,
528	offs_d,
529	dim_mask,
530	tile_mask,
531	OUT_DTYPE=Q.dtype,
532	HEAD_DIM=HEAD_SIZE,
533	VQB=VQB,
534	)
535
536	# S : [BLOCK_M, TILE_SIZE]
537	S = scale * tl.dot(Q, K_T)
538
539	# Causal + query-padding mask
540	query_abs_pos = context_len + query_pos[:, None]
541	seq_mask = seq_offset[None, :] <= query_abs_pos
542	S = tl.where(
543	query_mask_1[:, None] & query_mask_0[:, None] & seq_mask,
544	S,
545	float("-inf"),
546	)
547
548	# Online softmax
549	m_j = tl.maximum(M, tl.max(S, axis=1))
550	m_j = tl.where(m_j > float("-inf"), m_j, 0.0)
551	P = tl.exp(S - m_j[:, None])
552	l_j = tl.sum(P, axis=1)
553	alpha = tl.exp(M - m_j)
554	acc = acc * alpha[:, None]
555	L = L * alpha + l_j
556	M = m_j
557
558	# acc += P · V
559	acc += tl.dot(P.to(V.dtype), V)
560
561	# Epilogue: normalize and store
562	acc = acc / L[:, None]
563
564	output_offset = (
565	query_offset_0[:, None] * output_stride_0
566	+ query_offset_1[:, None] * output_stride_1
567	+ offs_d[None, :]
568	)
569	tl.store(
570	output_ptr + output_offset,
571	acc,
572	mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
573	)
574
575
576	# ---------------------------------------------------------------------------
577	# Unified 3D (split-KV) attention kernel over a TurboQuant KV cache.
578	#
579	# Same loop body as the 2D kernel above, but with an extra ``segm_idx`` grid
580	# axis: each program walks only ``cdiv(seq_len, NUM_SEGMENTS_PER_SEQ * TILE)``
581	# tiles of the KV history and writes a partial (max, expsum, acc) triple to
582	# a scratch buffer. A subsequent ``reduce_segments`` kernel merges the
583	# ``NUM_SEGMENTS_PER_SEQ`` partials per (q_token, q_head) via the standard
584	# online-softmax merge.
585	#
586	# This gives the decode path ``num_seqs × Hk × NUM_SEGMENTS_PER_SEQ`` CTAs
587	# instead of just ``num_seqs × Hk``, which is what v2 already uses (and what
588	# lets v2 run flat at ~0.15 ms across KV lengths on MI300X).
589	# ---------------------------------------------------------------------------
590
591
592	@triton.jit
593	def kernel_tq_unified_attention_3d(
594	segm_output_ptr, # [num_tokens, Hq, NUM_SEGMENTS_PER_SEQ, HEAD_SIZE_PADDED] fp32
595	segm_max_ptr, # [num_tokens, Hq, NUM_SEGMENTS_PER_SEQ] fp32
596	segm_expsum_ptr, # [num_tokens, Hq, NUM_SEGMENTS_PER_SEQ] fp32
597	query_ptr, # [num_tokens, Hq, D] — raw if FUSE_Q_ROT else rotated
598	KV_cache_ptr, # [num_blocks, block_size, Hk, padded_slot] uint8
599	KV_cache_u16_ptr, # uint16 view of same storage — Opt#2 wide metadata loads
600	Centroids_ptr,
601	Pair_lut_ptr,
602	PiT_ptr, # [D, D] fp32 — only dereferenced when FUSE_Q_ROT == 1
603	block_tables_ptr,
604	seq_lens_ptr,
605	query_start_len_ptr,
606	sinks_ptr, # [Hq] fp32 — per-head sink logits; dereferenced only when USE_SINKS
607	scale,
608	num_query_heads: tl.constexpr,
609	num_queries_per_kv: tl.constexpr,
610	block_table_stride: tl.int64,
611	query_stride_0: tl.int64,
612	query_stride_1: tl.int64,
613	stride_cache_block: tl.int64,
614	pit_stride_0: tl.int64,
615	pit_stride_1: tl.int64,
616	BLOCK_SIZE: tl.constexpr,
617	TILE_SIZE: tl.constexpr,
618	HEAD_SIZE: tl.constexpr,
619	HEAD_SIZE_PADDED: tl.constexpr,
620	BLOCK_Q: tl.constexpr,
621	BLOCK_M: tl.constexpr,
622	num_seqs: tl.int32,
623	NUM_SEGMENTS_PER_SEQ: tl.constexpr,
624	# TQ layout constants (same as 2D)
625	MSE_BITS: tl.constexpr,
626	MSE_BYTES: tl.constexpr,
627	VQB: tl.constexpr,
628	VAL_DATA_BYTES: tl.constexpr,
629	N_CENTROIDS: tl.constexpr,
630	KEY_FP8: tl.constexpr,
631	USE_PAIR_LUT: tl.constexpr,
632	# Opt#3 SoA layout constants
633	NUM_KV_HEADS: tl.constexpr,
634	KEY_DATA_BYTES: tl.constexpr,
635	META_REGION_OFFSET: tl.constexpr,
636	NUM_SOA_FIELDS: tl.constexpr,
637	SOA_K_NORM: tl.constexpr,
638	SOA_V_SCALE: tl.constexpr,
639	SOA_V_ZERO: tl.constexpr,
640	NORM_CORRECTION: tl.constexpr = 0,
641	FP8_E4B15: tl.constexpr = 0,
642	FUSE_Q_ROT: tl.constexpr = 0, # Fused Q @ PiT prologue when 1 (MSE path)
643	USE_SINKS: tl.constexpr = 0, # Per-head sink logit folded into softmax denom
644	):
645	q_block_global_idx = tl.program_id(0)
646	kv_head_idx = tl.program_id(1)
647	segm_idx = tl.program_id(2)
648
649	seq_idx = _find_seq_idx(
650	query_start_len_ptr, q_block_global_idx, num_seqs, BLOCK_Q, True
651	)
652	q_block_start_idx = tl.load(query_start_len_ptr + seq_idx) // BLOCK_Q + seq_idx
653	q_block_local_idx = q_block_global_idx - q_block_start_idx
654
655	cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx)
656	cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + 1)
657	cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index
658
659	if q_block_local_idx * BLOCK_Q >= cur_batch_query_len:
660	return
661
662	seq_len = tl.load(seq_lens_ptr + seq_idx)
663	tiles_per_segment = tl.cdiv(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE)
664
665	# Segment empty for this sequence — skip.
666	if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len:
667	return
668
669	offs_m = tl.arange(0, BLOCK_M)
670	offs_d = tl.arange(0, HEAD_SIZE_PADDED)
671	offs_t = tl.arange(0, TILE_SIZE)
672	query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv
673
674	query_offset_0 = cur_batch_in_all_start_index + query_pos
675	query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv
676	query_offset = (
677	query_offset_0[:, None] * query_stride_0
678	+ query_offset_1[:, None] * query_stride_1
679	+ offs_d[None, :]
680	)
681
682	dim_mask = tl.where(offs_d < HEAD_SIZE, 1, 0).to(tl.int1)
683	query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1)
684	query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1)
685
686	Q = tl.load(
687	query_ptr + query_offset,
688	mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
689	other=0.0,
690	)
691
692	if FUSE_Q_ROT:
693	Q = _tq_fuse_q_rotation(
694	Q,
695	PiT_ptr,
696	pit_stride_0,
697	pit_stride_1,
698	dim_mask,
699	HEAD_SIZE_PADDED,
700	)
701
702	block_table_offset = seq_idx * block_table_stride
703
704	# Online softmax state. For the split-KV (3D) kernel we must fold the
705	# sink logit into only segment 0's (M, L) so that the stage-2 reducer
706	# — which is a standard online-softmax merge — counts the sink exactly
707	# once. Segments with segm_idx > 0 keep the plain -inf / 1.0 init and
708	# therefore carry no sink contribution; stage-2 then combines them.
709	# This mirrors the v1 sink pattern from PR #40663 (sid==0 init trick).
710	if USE_SINKS and segm_idx == 0:
711	M = tl.load(
712	sinks_ptr + query_offset_1, mask=query_mask_1, other=float("-inf")
713	).to(tl.float32)
714	else:
715	M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
716	L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
717	acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32)
718
719	context_len = seq_len - cur_batch_query_len
720	max_seq_prefix_len = (
721	context_len
722	+ q_block_local_idx * BLOCK_Q
723	+ (BLOCK_M - 1) // num_queries_per_kv
724	+ 1
725	)
726	max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len)
727	num_tiles = tl.cdiv(max_seq_prefix_len, TILE_SIZE)
728
729	# Segment bounds: clip to the causal prefix range.
730	tile_lo = segm_idx * tiles_per_segment
731	tile_hi = tl.minimum((segm_idx + 1) * tiles_per_segment, num_tiles)
732
733	for j in range(tile_lo, tile_hi):
734	seq_offset = j * TILE_SIZE + offs_t
735	tile_mask = seq_offset < max_seq_prefix_len
736
737	physical_block_idx = tl.load(
738	block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE
739	).to(tl.int64)
740
741	# Opt#3 SoA addressing (see 2D kernel for the layout invariants).
742	slot_within_block = (seq_offset % BLOCK_SIZE).to(tl.int64)
743	block_base = physical_block_idx * stride_cache_block
744	DATA_BYTES_PER_SLOT: tl.constexpr = KEY_DATA_BYTES + VAL_DATA_BYTES
745	data_bases = (
746	block_base
747	+ slot_within_block * (NUM_KV_HEADS * DATA_BYTES_PER_SLOT)
748	+ tl.cast(kv_head_idx, tl.int64) * DATA_BYTES_PER_SLOT
749	)
750	val_bases = data_bases + KEY_DATA_BYTES
751
752	head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + tl.cast(
753	kv_head_idx, tl.int64
754	) * (NUM_SOA_FIELDS * BLOCK_SIZE)
755	knorm_u16_addrs = (
756	head_meta_u16_base + SOA_K_NORM * BLOCK_SIZE + slot_within_block
757	)
758	vscale_u16_addrs = (
759	head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + slot_within_block
760	)
761	vzero_u16_addrs = (
762	head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + slot_within_block
763	)
764
765	K_T = _tq_load_k_tile(
766	KV_cache_ptr,
767	KV_cache_u16_ptr,
768	data_bases,
769	knorm_u16_addrs,
770	offs_d,
771	dim_mask,
772	tile_mask,
773	Centroids_ptr,
774	Pair_lut_ptr,
775	OUT_DTYPE=Q.dtype,
776	HEAD_DIM=HEAD_SIZE,
777	BLOCK_D=HEAD_SIZE_PADDED,
778	MSE_BITS=MSE_BITS,
779	N_CENTROIDS=N_CENTROIDS,
780	KEY_FP8=KEY_FP8,
781	USE_PAIR_LUT=USE_PAIR_LUT,
782	NORM_CORRECTION=NORM_CORRECTION,
783	FP8_E4B15=FP8_E4B15,
784	TILE_SIZE=TILE_SIZE,
785	)
786
787	V = _tq_load_v_tile(
788	KV_cache_ptr,
789	KV_cache_u16_ptr,
790	val_bases,
791	vscale_u16_addrs,
792	vzero_u16_addrs,
793	offs_d,
794	dim_mask,
795	tile_mask,
796	OUT_DTYPE=Q.dtype,
797	HEAD_DIM=HEAD_SIZE,
798	VQB=VQB,
799	)
800
801	S = scale * tl.dot(Q, K_T)
802
803	query_abs_pos = context_len + query_pos[:, None]
804	seq_mask = seq_offset[None, :] <= query_abs_pos
805	S = tl.where(
806	query_mask_1[:, None] & query_mask_0[:, None] & seq_mask,
807	S,
808	float("-inf"),
809	)
810
811	m_j = tl.maximum(M, tl.max(S, axis=1))
812	m_j = tl.where(m_j > float("-inf"), m_j, 0.0)
813	P = tl.exp(S - m_j[:, None])
814	l_j = tl.sum(P, axis=1)
815	alpha = tl.exp(M - m_j)
816	acc = acc * alpha[:, None]
817	L = L * alpha + l_j
818	M = m_j
819
820	acc += tl.dot(P.to(V.dtype), V)
821
822	# -------- Write partial (acc, M, L) to scratch; no normalization --------
823	# Layout matches baseline reduce_segments so we can reuse it.
824	segm_output_offset = (
825	query_offset_0[:, None].to(tl.int64)
826	* (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
827	+ query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
828	+ segm_idx * HEAD_SIZE_PADDED
829	+ tl.arange(0, HEAD_SIZE_PADDED)[None, :]
830	)
831	tl.store(
832	segm_output_ptr + segm_output_offset,
833	acc,
834	mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
835	)
836	segm_offset = (
837	query_offset_0.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ)
838	+ query_offset_1 * NUM_SEGMENTS_PER_SEQ
839	+ segm_idx
840	)
841	tl.store(segm_max_ptr + segm_offset, M, mask=query_mask_0 & query_mask_1)
842	tl.store(segm_expsum_ptr + segm_offset, L, mask=query_mask_0 & query_mask_1)
843
844
845	# ---------------------------------------------------------------------------
846	# Launcher
847	# ---------------------------------------------------------------------------
848
849
850	_layout_cache: dict[tuple, dict[str, int]] = {}
851
852
853	def _get_layout(D: int, mse_bits: int, value_quant_bits: int) -> dict[str, int]:
854	key = (D, mse_bits, value_quant_bits)
855	cfg = _layout_cache.get(key)
856	if cfg is None:
857	cfg = {
858	"mse_bytes": math.ceil(D * mse_bits / 8),
859	"val_data_bytes": math.ceil(D * value_quant_bits / 8),
860	"BLOCK_D": triton.next_power_of_2(D),
861	}
862	_layout_cache[key] = cfg
863	return cfg
864
865
866	def _get_pair_lut(centroids: torch.Tensor) -> torch.Tensor:
867	"""Return a fresh pair-LUT for ``centroids`` on each call.
868
869	Previously cached by ``(data_ptr, device)`` — that key is unsafe
870	because the CUDA allocator can reuse freed addresses for different
871	centroid tensors, yielding a stale (wrong-values) LUT. The LUT is tiny
872	(NN2 fp32, ~2KB at N=16) so unconditional rebuild is essentially
873	free compared to attention work. If this shows up on a profile,
874	replace with a hash-of-values fingerprint, not data_ptr.
875	"""
876	return build_pair_lut(centroids)
877
878
879	def triton_turboquant_unified_attention(
880	query: torch.Tensor, # [num_tokens, Hq, D] - fp16/bf16
881	kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8
882	block_table: torch.Tensor, # [num_seqs, max_num_blocks] int32
883	seq_lens: torch.Tensor, # [num_seqs] int32
884	query_start_loc: torch.Tensor, # [num_seqs+1] int32
885	Pi: torch.Tensor, # [D, D] fp32
886	centroids: torch.Tensor, # [n_centroids] fp32
887	scale: float,
888	mse_bits: int,
889	key_packed_size: int,
890	value_quant_bits: int,
891	value_packed_size: int, # unused; kept for signature parity with v2
892	key_fp8: bool = False,
893	norm_correction: bool = False,
894	PiT: torch.Tensor \| None = None,
895	output: torch.Tensor \| None = None,
896	tile_size: int \| None = None,
897	max_query_len: int \| None = None,
898	max_seq_len: int \| None = None,
899	num_kv_splits: int \| None = None,
900	force_2d: bool = False,
901	fuse_q_rot: bool = True,
902	sinks: torch.Tensor \| None = None, # [Hq] float — per-head sink logits
903	) -> torch.Tensor:
904	"""Launch unified TQ attention (v3).
905
906	``query`` carries raw query vectors. For the MSE-key path the query
907	has to be rotated by ``PiT`` before it can be multiplied against the
908	(already-rotated) stored K. By default (``fuse_q_rot=True``) that
909	rotation is done inside the attention kernel prologue as a single
910	small MFMA — no extra dispatch, no HBM round-trip. Setting
911	``fuse_q_rot=False`` restores the original launcher path (fp32 rocBLAS
912	GEMM + casts + .contiguous()), which is kept as an A/B toggle for
913	bench harnesses; numerical results match the fused path to within a
914	few ulp of fp32 round-off. The FP8-key path never rotates Q regardless
915	of this flag.
916
917	``tile_size`` defaults to ``32`` for prefill (``max_query_len > 1``) and
918	``16`` for pure decode (``max_query_len == 1``). Callers may override.
919
920	Dispatch rule:
921	* Prefill / chunked-prefill (any query block with ``BLOCK_Q > 1``) and
922	``num_tokens > num_seqs``: 2D kernel (plenty of CTAs already).
923	* Pure decode with long KV: 3D split-KV kernel + ``reduce_segments``
924	for extra parallelism; this is the same pattern v2 uses.
925	* ``force_2d=True``: force the 2D path regardless (used by the
926	apples-to-apples bench to isolate pure TQ dequant overhead).
927
928	``num_kv_splits`` controls the 3D-split segment count (default ``16``).
929
930	Returns ``output`` of shape ``[num_tokens, Hq, D]`` in ``query.dtype``.
931	"""
932	assert query.dim() == 3, f"query must be [N, Hq, D], got {query.shape}"
933	num_tokens, Hq, D = query.shape
934	Hk = kv_cache.shape[2]
935	block_size = kv_cache.shape[1]
936	kv_group_size = Hq // Hk
937	num_seqs = int(query_start_loc.shape[0] - 1)
938	device = query.device
939
940	cfg = _get_layout(D, mse_bits, value_quant_bits)
941	_ = value_packed_size # unused
942
943	# Q-rotation strategy:
944	# * FP8-key path: no rotation at all (keys stored as fp8, no codebook).
945	# * MSE-key path + fuse_q_rot=True (default): pass the raw query to the
946	# kernel together with PiT and let the kernel's prologue do one
947	# fp32 tl.dot(Q, PiT). Eliminates the launcher-side
948	# (cast, rocBLAS GEMM, cast, contiguous) chain (~50-60 us/step at
949	# D=64, all launch-bound — see bottleneck report §10).
950	# * MSE-key path + fuse_q_rot=False: legacy launcher rotation. Kept
951	# as an A/B toggle for bench harnesses.
952	#
953	# PiT is always materialized as fp32 D x D contiguous; the kernel loads
954	# it as fp32 (cheap, 16KB at D=64) and computes Q @ PiT with IEEE fp32
955	# accumulate via tl.dot.
956	if key_fp8:
957	q_rot = query.contiguous()
958	apply_fuse_q_rot = False
959	else:
960	if PiT is None:
961	PiT = Pi.T.contiguous()
962	apply_fuse_q_rot = bool(fuse_q_rot)
963	if apply_fuse_q_rot:
964	q_rot = query.contiguous()
965	else:
966	q_rot = (query.float() @ PiT).to(query.dtype).contiguous()
967
968	# PiT in fp32, contiguous. For the fused path this is what the kernel
969	# loads; for the legacy path it's passed through as a harmless tensor
970	# (never dereferenced under the FUSE_Q_ROT constexpr guard). On FP8
971	# path Pi may be unused at the call site, so we fall back to an
972	# arbitrary non-null tensor (reuse centroids) to satisfy Triton's
973	# non-null pointer requirement.
974	if (not key_fp8) and PiT is not None:
975	PiT_f32 = PiT if PiT.dtype == torch.float32 else PiT.to(torch.float32)
976	if not PiT_f32.is_contiguous():
977	PiT_f32 = PiT_f32.contiguous()
978	pit_stride_0 = PiT_f32.stride(0)
979	pit_stride_1 = PiT_f32.stride(1)
980	else:
981	PiT_f32 = centroids # harmless dummy; not dereferenced when FUSE_Q_ROT=0
982	pit_stride_0 = 0
983	pit_stride_1 = 0
984
985	# Sinks: per-head fp32 logits, contiguous. The kernel only dereferences
986	# this pointer when USE_SINKS=1, so when the caller passes None we bind
987	# a harmless non-null tensor (centroids) to satisfy Triton's non-null
988	# pointer requirement. See the sink design notes in the 2D/3D kernel
989	# bodies and the v1 precedent from PR #40663 (sid==0 init trick).
990	if sinks is not None:
991	sinks_f32 = sinks if sinks.dtype == torch.float32 else sinks.to(torch.float32)
992	if not sinks_f32.is_contiguous():
993	sinks_f32 = sinks_f32.contiguous()
994	assert sinks_f32.numel() == Hq, (
995	f"sinks must have shape [Hq={Hq}], got numel={sinks_f32.numel()}"
996	)
997	use_sinks = True
998	else:
999	sinks_f32 = centroids # harmless dummy; not dereferenced when USE_SINKS=0
1000	use_sinks = False
1001
1002	if output is None:
1003	output = torch.empty_like(query)
1004
1005	# BLOCK_M heuristic (TQ-specific; diverges from stock unified_attention).
1006	#
1007	# Stock fp16 picks BLOCK_M=16 for small GQA because fp16 is ALU-saturated
1008	# at any tile size and smaller BLOCK_M = better occupancy. TQ is a
1009	# DIFFERENT story: the dequant chain (packed-byte loads + centroid gather
1010	# + pair-LUT + norm/scale bitcasts) has big per-block fixed overhead, and
1011	# the pair-LUT gather pattern is latency-bound (NOT HBM-bandwidth-bound,
1012	# confirmed via rocprof: MemUnitStalled=0.01%, VALUBusy=56% at BM=16).
1013	#
1014	# For prefill/chunked (max_query_len > 1) we want BLOCK_M as large as
1015	# possible to (a) amortize the fixed dequant overhead over more query
1016	# tokens per CTA and (b) let MFMA saturate with a 128x32 tile shape.
1017	# Empirical sweep on MI300X across gpt-oss (D=64) and llama (D=128):
1018	# prefill B=1 Q=4k : BM=128 is 3.4x faster than BM=16
1019	# chunked B=64 Q=1k C=8: BM=128 is 5.6x faster than BM=16
1020	# With BM=128 v3 matches or BEATS the fp16 baseline on prefill/chunked.
1021	#
1022	# For pure decode (max_query_len == 1) BLOCK_M stays small because
1023	# BLOCK_Q > 1 just pads rows — each sequence contributes at most 1 query
1024	# token and larger BLOCK_M wastes lanes on masked-out rows.
1025	if max_query_len is not None:
1026	is_prefill_like = max_query_len > 1
1027	else:
1028	is_prefill_like = num_tokens > num_seqs
1029
1030	if is_prefill_like:
1031	# Prefill / chunked: use BLOCK_M=128 (rounded up to a multiple of
1032	# kv_group_size). Each CTA does very large work (BLOCK_M x TILE_SIZE
1033	# x KV_len MMA + dequant), so even small grids (~17 CTAs) saturate
1034	# MI300X just fine. We verified empirically: BM=128 wins across all
1035	# prefill/chunked sweep points, including small (B=1 Q=256 C=8k) and
1036	# large (B=64 Q=1k C=8k; 5.6x speedup vs BM=16). BLOCK_M must be a
1037	# multiple of kv_group_size so BLOCK_Q is an integer.
1038	BLOCK_M = max(128, triton.next_power_of_2(kv_group_size))
1039	else:
1040	# Decode: BLOCK_Q > 1 in decode just pads rows (at most 1 query
1041	# token per sequence), so there's little to amortize. Inherit
1042	# stock's small-BLOCK_M heuristic.
1043	BLOCK_M = 16 if kv_group_size <= 16 else triton.next_power_of_2(kv_group_size)
1044	BLOCK_Q = BLOCK_M // kv_group_size
1045
1046	# Grid: at most ceil(N / BLOCK_Q) + num_seqs q-blocks total, like unified.
1047	total_num_q_blocks = num_tokens // BLOCK_Q + num_seqs
1048
1049	# TILE_SIZE heuristic (matches stock unified_attention's _get_tile_size):
1050	# prefill (max_query_len > 1): 32
1051	# decode (max_query_len == 1): 16
1052	if tile_size is None:
1053	tile_size = 32 if is_prefill_like else 16
1054
1055	# Pair-LUT fast path for 4-bit MSE keys. Skipped for FP8 and non-4-bit.
1056	# When USE_PAIR_LUT==0 the kernel never dereferences pair_lut, but Triton
1057	# still requires a tensor pointer, so we fall back to reusing `centroids`.
1058	use_pair_lut = (not key_fp8) and (mse_bits == 4)
1059	pair_lut = _get_pair_lut(centroids) if use_pair_lut else centroids
1060
1061	fp8_e4b15 = _use_fp8_e4b15(device.index or 0)
1062	num_stages = 1 if _is_hip else 2
1063
1064	# uint16-aliased view of the cache. Under the Opt#3 SoA layout, K-norm /
1065	# V-scale / V-zero live in a contiguous per-block metadata region and are
1066	# always 2-byte aligned, so single-instruction u16 loads replace the
1067	# original 2× uint8 + OR sequence for every per-token metadata fetch.
1068	kv_cache_u16 = kv_cache.view(torch.uint16)
1069
1070	# Opt#3 SoA layout constants (derived locally; matches the store-side
1071	# computation so the launcher signature stays unchanged). Invariant:
1072	# data_bytes_per_slot + meta_bytes_per_slot == slot_size_aligned.
1073	mse_bytes = cfg["mse_bytes"]
1074	val_data_bytes = cfg["val_data_bytes"]
1075	key_data_bytes = D if key_fp8 else mse_bytes
1076	data_bytes_per_slot = key_data_bytes + val_data_bytes
1077	meta_region_offset = block_size * Hk * data_bytes_per_slot
1078	num_soa_fields = 2 if key_fp8 else 3
1079	soa_k_norm = 0 # unused for FP8; harmless constant
1080	soa_v_scale = 0 if key_fp8 else 1
1081	soa_v_zero = 1 if key_fp8 else 2
1082
1083	# ------------------------------------------------------------------
1084	# Dispatch: 2D for prefill / chunked; 3D split-KV for pure decode with
1085	# long KV. Falls back to 2D if the sequences are short or if the caller
1086	# forces it. ``is_prefill_like`` was already computed above for the
1087	# BLOCK_M / tile_size heuristic; reuse it here.
1088	# ------------------------------------------------------------------
1089
1090	# Use 3D only for pure decode where the KV history is long enough that
1091	# splitting meaningfully increases SM occupancy. Threshold is empirically
1092	# tuned on MI300X:
1093	# seq < 1024 -> 2D path (short KV, reduce overhead > parallelism gain)
1094	# seq >= 1024 -> 3D path (v3/v2 ratio 1.04-1.27x across shapes)
1095	# Below 1024 KV tokens the reduce_segments + scratch-alloc overhead
1096	# exceeds the parallelism gain from splitting. Above it, the single-CTA
1097	# 2D path is serial over the KV history and starves MI300X's 304 CUs.
1098	#
1099	# IMPORTANT: max_seq_len is expected to come from the caller (backends
1100	# track it from the block table). We avoid calling seq_lens.max().item()
1101	# here because that forces a GPU->CPU sync on every launcher call, which
1102	# adds ~150us of overhead to pure-decode steps.
1103	if max_seq_len is None:
1104	# Fallback: infer from block_table (host-side, no sync). This is an
1105	# upper bound, sufficient for the dispatch decision.
1106	max_seq_len_hint = int(block_table.shape[1]) * int(block_size)
1107	else:
1108	max_seq_len_hint = int(max_seq_len)
1109	use_3d = (not force_2d) and (not is_prefill_like) and max_seq_len_hint >= 1024
1110
1111	if not use_3d:
1112	kernel_tq_unified_attention_2d[(total_num_q_blocks, Hk)](
1113	output_ptr=output,
1114	query_ptr=q_rot,
1115	KV_cache_ptr=kv_cache,
1116	KV_cache_u16_ptr=kv_cache_u16,
1117	Centroids_ptr=centroids,
1118	Pair_lut_ptr=pair_lut,
1119	PiT_ptr=PiT_f32,
1120	block_tables_ptr=block_table,
1121	seq_lens_ptr=seq_lens,
1122	query_start_len_ptr=query_start_loc,
1123	sinks_ptr=sinks_f32,
1124	scale=scale,
1125	num_query_heads=Hq,
1126	num_queries_per_kv=kv_group_size,
1127	block_table_stride=block_table.stride(0),
1128	query_stride_0=q_rot.stride(0),
1129	query_stride_1=q_rot.stride(1),
1130	output_stride_0=output.stride(0),
1131	output_stride_1=output.stride(1),
1132	stride_cache_block=kv_cache.stride(0),
1133	pit_stride_0=pit_stride_0,
1134	pit_stride_1=pit_stride_1,
1135	BLOCK_SIZE=block_size,
1136	TILE_SIZE=tile_size,
1137	HEAD_SIZE=D,
1138	HEAD_SIZE_PADDED=cfg["BLOCK_D"],
1139	BLOCK_Q=BLOCK_Q,
1140	BLOCK_M=BLOCK_M,
1141	num_seqs=num_seqs,
1142	MSE_BITS=mse_bits,
1143	MSE_BYTES=mse_bytes,
1144	VQB=value_quant_bits,
1145	VAL_DATA_BYTES=val_data_bytes,
1146	N_CENTROIDS=int(centroids.numel()),
1147	KEY_FP8=1 if key_fp8 else 0,
1148	USE_PAIR_LUT=1 if use_pair_lut else 0,
1149	NUM_KV_HEADS=Hk,
1150	KEY_DATA_BYTES=key_data_bytes,
1151	META_REGION_OFFSET=meta_region_offset,
1152	NUM_SOA_FIELDS=num_soa_fields,
1153	SOA_K_NORM=soa_k_norm,
1154	SOA_V_SCALE=soa_v_scale,
1155	SOA_V_ZERO=soa_v_zero,
1156	NORM_CORRECTION=1 if norm_correction else 0,
1157	FP8_E4B15=fp8_e4b15,
1158	FUSE_Q_ROT=1 if apply_fuse_q_rot else 0,
1159	USE_SINKS=1 if use_sinks else 0,
1160	num_warps=4,
1161	num_stages=num_stages,
1162	)
1163	return output
1164
1165	# ---------------- 3D split-KV path ----------------
1166	# Pick segment count. 16 is a good default for MI300X (304 CUs): at
1167	# Hk=8, num_seqs=1, total CTAs = 1 * 8 * 16 = 128, which saturates when
1168	# combined with warp-level parallelism. Capped at ceil(max_seq_len /
1169	# TILE_SIZE) so we don't launch empty segments.
1170	if num_kv_splits is None:
1171	num_kv_splits = 16
1172	max_possible_splits = max(1, (max_seq_len_hint + tile_size - 1) // tile_size)
1173	num_segments = max(1, min(num_kv_splits, max_possible_splits))
1174
1175	HEAD_SIZE_PADDED = cfg["BLOCK_D"]
1176	# Scratch buffers for segment partials. fp32 to match the kernel's
1177	# accumulator dtype. Allocated per-call for now; a production backend
1178	# would hoist these to persistent buffers reused across layers.
1179	segm_output = torch.empty(
1180	(num_tokens, Hq, num_segments, HEAD_SIZE_PADDED),
1181	dtype=torch.float32,
1182	device=device,
1183	)
1184	segm_max = torch.empty(
1185	(num_tokens, Hq, num_segments), dtype=torch.float32, device=device
1186	)
1187	segm_expsum = torch.empty(
1188	(num_tokens, Hq, num_segments), dtype=torch.float32, device=device
1189	)
1190
1191	kernel_tq_unified_attention_3d[(total_num_q_blocks, Hk, num_segments)](
1192	segm_output_ptr=segm_output,
1193	segm_max_ptr=segm_max,
1194	segm_expsum_ptr=segm_expsum,
1195	query_ptr=q_rot,
1196	KV_cache_ptr=kv_cache,
1197	KV_cache_u16_ptr=kv_cache_u16,
1198	Centroids_ptr=centroids,
1199	Pair_lut_ptr=pair_lut,
1200	PiT_ptr=PiT_f32,
1201	block_tables_ptr=block_table,
1202	seq_lens_ptr=seq_lens,
1203	query_start_len_ptr=query_start_loc,
1204	sinks_ptr=sinks_f32,
1205	scale=scale,
1206	num_query_heads=Hq,
1207	num_queries_per_kv=kv_group_size,
1208	block_table_stride=block_table.stride(0),
1209	query_stride_0=q_rot.stride(0),
1210	query_stride_1=q_rot.stride(1),
1211	stride_cache_block=kv_cache.stride(0),
1212	pit_stride_0=pit_stride_0,
1213	pit_stride_1=pit_stride_1,
1214	BLOCK_SIZE=block_size,
1215	TILE_SIZE=tile_size,
1216	HEAD_SIZE=D,
1217	HEAD_SIZE_PADDED=HEAD_SIZE_PADDED,
1218	BLOCK_Q=BLOCK_Q,
1219	BLOCK_M=BLOCK_M,
1220	num_seqs=num_seqs,
1221	NUM_SEGMENTS_PER_SEQ=num_segments,
1222	MSE_BITS=mse_bits,
1223	MSE_BYTES=mse_bytes,
1224	VQB=value_quant_bits,
1225	VAL_DATA_BYTES=val_data_bytes,
1226	N_CENTROIDS=int(centroids.numel()),
1227	KEY_FP8=1 if key_fp8 else 0,
1228	USE_PAIR_LUT=1 if use_pair_lut else 0,
1229	NUM_KV_HEADS=Hk,
1230	KEY_DATA_BYTES=key_data_bytes,
1231	META_REGION_OFFSET=meta_region_offset,
1232	NUM_SOA_FIELDS=num_soa_fields,
1233	SOA_K_NORM=soa_k_norm,
1234	SOA_V_SCALE=soa_v_scale,
1235	SOA_V_ZERO=soa_v_zero,
1236	NORM_CORRECTION=1 if norm_correction else 0,
1237	FP8_E4B15=fp8_e4b15,
1238	FUSE_Q_ROT=1 if apply_fuse_q_rot else 0,
1239	USE_SINKS=1 if use_sinks else 0,
1240	num_warps=4,
1241	num_stages=num_stages,
1242	)
1243
1244	# Reduce: merge num_segments partials per (q_token, q_head) into the
1245	# final output via online-softmax merge. Reusing the baseline kernel
1246	# (KV-format-agnostic).
1247	reduce_segments[(num_tokens, Hq)](
1248	output_ptr=output,
1249	segm_output_ptr=segm_output,
1250	segm_max_ptr=segm_max,
1251	segm_expsum_ptr=segm_expsum,
1252	seq_lens_ptr=seq_lens,
1253	num_seqs=num_seqs,
1254	num_query_heads=Hq,
1255	out_scale_inv=1.0,
1256	output_stride_0=output.stride(0),
1257	output_stride_1=output.stride(1),
1258	block_table_stride=block_table.stride(0),
1259	TILE_SIZE=tile_size,
1260	HEAD_SIZE=D,
1261	HEAD_SIZE_PADDED=HEAD_SIZE_PADDED,
1262	query_start_len_ptr=query_start_loc,
1263	BLOCK_Q=BLOCK_Q,
1264	NUM_SEGMENTS_PER_SEQ=num_segments,
1265	USE_FP8=False,
1266	)
1267	return output
1268
1269
1270	# ---------------------------------------------------------------------------
1271	# Thin decode adapter: lets callers who use the AoS decode launcher signature
1272	# switch to the SoA decode without rebuilding query_start_loc themselves.
1273	# ---------------------------------------------------------------------------
1274
1275
1276	def triton_turboquant_decode_attention_soa(
1277	query: torch.Tensor, # [B, Hq, D]
1278	kv_cache: torch.Tensor,
1279	block_table: torch.Tensor,
1280	seq_lens: torch.Tensor,
1281	Pi: torch.Tensor,
1282	centroids: torch.Tensor,
1283	scale: float,
1284	mse_bits: int,
1285	key_packed_size: int,
1286	value_quant_bits: int,
1287	value_packed_size: int,
1288	key_fp8: bool = False,
1289	norm_correction: bool = False,
1290	PiT: torch.Tensor \| None = None,
1291	# kept for signature parity with v1/v2 decode launchers; unused here.
1292	max_seq_len: int = 0,
1293	mid_o_buf: Any = None,
1294	output_buf: torch.Tensor \| None = None,
1295	lse_buf: Any = None,
1296	buf_holder: Any = None,
1297	max_num_kv_splits: int = 32,
1298	sinks: torch.Tensor \| None = None,
1299	) -> torch.Tensor:
1300	"""Decode-only convenience wrapper around ``triton_turboquant_unified_attention``.
1301
1302	Treats a rank-3 ``[B, Hq, D]`` query as one token per request (query_len = 1)
1303	and synthesizes a ``query_start_loc`` of ``[0, 1, 2, ..., B]``.
1304	``max_num_kv_splits`` is forwarded to the unified launcher as the 3D
1305	split-KV segment count (capped per-call against ``max_seq_len``).
1306	``sinks`` (optional ``[Hq]`` fp32) are forwarded to the kernel which
1307	folds them into the softmax denominator via the init-time trick.
1308	"""
1309	del mid_o_buf, lse_buf, buf_holder
1310	B = query.shape[0]
1311	cu_seqlens_q = torch.arange(B + 1, device=query.device, dtype=seq_lens.dtype)
1312
1313	out = triton_turboquant_unified_attention(
1314	query=query.contiguous(),
1315	kv_cache=kv_cache,
1316	block_table=block_table,
1317	seq_lens=seq_lens,
1318	query_start_loc=cu_seqlens_q,
1319	Pi=Pi,
1320	centroids=centroids,
1321	scale=scale,
1322	mse_bits=mse_bits,
1323	key_packed_size=key_packed_size,
1324	value_quant_bits=value_quant_bits,
1325	value_packed_size=value_packed_size,
1326	key_fp8=key_fp8,
1327	norm_correction=norm_correction,
1328	PiT=PiT,
1329	output=output_buf[:B] if output_buf is not None else None,
1330	max_query_len=1,
1331	max_seq_len=max_seq_len if max_seq_len > 0 else None,
1332	num_kv_splits=max_num_kv_splits,
1333	sinks=sinks,
1334	)
1335	return out