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

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 1,155 lines 37.2 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3	"""Sub-byte packed (INT4) per-token-head KV cache mode.
4
5	INT4 packs two 4-bit values per cache byte, pre-rotates with a single RHT,
6	and hides a 4-bit zero-point in the scale's low mantissa bits — too
7	different from the core kernel to share it. Owns the whole mode: nibble
8	pack/unpack, the reshape (write) kernel, the split-dot attention (read)
9	kernel, the RHT transform, and the public ``reshape_and_cache_int4`` /
10	``unified_attention_int4`` entry points.
11	"""
12
13	from __future__ import annotations
14
15	from typing import Any
16
17	import torch
18
19	from vllm.platforms import current_platform
20	from vllm.triton_utils import tl, triton
21	from vllm.v1.attention.ops.triton_attention_helpers import (
22	apply_alibi_to_score,
23	apply_softcap,
24	cdiv_fn,
25	compute_kv_seq_mask,
26	compute_tile_loop_bounds,
27	init_softmax_M,
28	load_qq_bias_tile,
29	resolve_seq_and_query_len,
30	softmax_step,
31	store_segm_reduce_scalars,
32	)
33	from vllm.v1.attention.ops.triton_unified_attention import reduce_segments
34
35	float8_info = torch.finfo(current_platform.fp8_dtype())
36
37	# 2 x int4 packed per storage byte.
38	_INT4_PACKING_FACTOR = 2
39
40
41	# ----------------------------------------------------------------------
42	# Nibble pack / unpack (shared write+read format)
43	# ----------------------------------------------------------------------
44
45
46	@triton.jit
47	def pack_int4_nibbles(lo, hi):
48	"""Pack two uint8 values (each in [0, 15]) into one byte."""
49	return (lo & 0xF) \| ((hi & 0xF) << 4)
50
51
52	@triton.jit
53	def unpack_int4_nibbles(packed):
54	"""Split one packed byte into the (low, high) nibble pair as uint8."""
55	return packed & 0xF, (packed >> 4) & 0xF
56
57
58	# ----------------------------------------------------------------------
59	# Write path: RHT + pack + per-(token, head) scale
60	# ----------------------------------------------------------------------
61
62
63	@triton.jit
64	def _reshape_cache_int4_kernel(
65	key_ptr,
66	value_ptr,
67	key_cache_ptr,
68	value_cache_ptr,
69	k_scale_cache_ptr,
70	v_scale_cache_ptr,
71	slot_mapping_ptr,
72	stride_key_tok: tl.int64,
73	stride_key_head: tl.int64,
74	stride_val_tok: tl.int64,
75	stride_val_head: tl.int64,
76	stride_kc_blk: tl.int64,
77	stride_kc_slot: tl.int64,
78	stride_kc_head: tl.int64,
79	stride_vc_blk: tl.int64,
80	stride_vc_slot: tl.int64,
81	stride_vc_head: tl.int64,
82	stride_ks_blk: tl.int64,
83	stride_ks_slot: tl.int64,
84	stride_ks_head: tl.int64,
85	stride_vs_blk: tl.int64,
86	stride_vs_slot: tl.int64,
87	stride_vs_head: tl.int64,
88	block_size: tl.constexpr,
89	head_size: tl.constexpr,
90	head_size_v: tl.constexpr,
91	PACKED_HEAD_PADDED: tl.constexpr,
92	):
93	"""INT4 asymmetric quantization with zero-point steganography."""
94	tok = tl.program_id(0)
95	head = tl.program_id(1)
96
97	slot = tl.load(slot_mapping_ptr + tok).to(tl.int64)
98	if slot < 0:
99	return
100
101	blk = slot // block_size
102	slot_in_blk = slot % block_size
103
104	half_offs = tl.arange(0, PACKED_HEAD_PADDED)
105	even_offs = half_offs * 2
106	odd_offs = half_offs * 2 + 1
107
108	half_k = head_size // 2
109	even_k_mask = even_offs < head_size
110	odd_k_mask = odd_offs < head_size
111	key_base = key_ptr + tok * stride_key_tok + head * stride_key_head
112
113	k_even = tl.load(key_base + even_offs, mask=even_k_mask, other=0.0).to(tl.float32)
114	k_odd = tl.load(key_base + odd_offs, mask=odd_k_mask, other=0.0).to(tl.float32)
115
116	k_min = tl.minimum(
117	tl.min(tl.where(even_k_mask, k_even, float("inf"))),
118	tl.min(tl.where(odd_k_mask, k_odd, float("inf"))),
119	)
120	k_max = tl.maximum(
121	tl.max(tl.where(even_k_mask, k_even, float("-inf"))),
122	tl.max(tl.where(odd_k_mask, k_odd, float("-inf"))),
123	)
124	k_scale = tl.maximum((k_max - k_min) / 15.0, 1e-6)
125	k_zp_f = tl.clamp(
126	tl.where(
127	-k_min / k_scale >= 0,
128	(-k_min / k_scale + 0.5).to(tl.int32),
129	(-k_min / k_scale - 0.5).to(tl.int32),
130	).to(tl.float32),
131	0.0,
132	15.0,
133	)
134
135	inv_k = 1.0 / k_scale
136	k_even_s = k_even * inv_k + k_zp_f
137	k_odd_s = k_odd * inv_k + k_zp_f
138	k_even_q = tl.clamp(
139	tl.where(
140	k_even_s >= 0,
141	(k_even_s + 0.5).to(tl.int32),
142	(k_even_s - 0.5).to(tl.int32),
143	).to(tl.float32),
144	0.0,
145	15.0,
146	)
147	k_odd_q = tl.clamp(
148	tl.where(
149	k_odd_s >= 0,
150	(k_odd_s + 0.5).to(tl.int32),
151	(k_odd_s - 0.5).to(tl.int32),
152	).to(tl.float32),
153	0.0,
154	15.0,
155	)
156
157	k_zp_int = k_zp_f.to(tl.int32)
158	k_scale_bits = k_scale.to(tl.int32, bitcast=True)
159	k_scale_packed = ((k_scale_bits & -16) \| (k_zp_int & 0xF)).to(
160	tl.float32, bitcast=True
161	)
162
163	tl.store(
164	k_scale_cache_ptr
165	+ blk * stride_ks_blk
166	+ slot_in_blk * stride_ks_slot
167	+ head * stride_ks_head,
168	k_scale_packed,
169	)
170
171	k_packed = pack_int4_nibbles(k_even_q.to(tl.uint8), k_odd_q.to(tl.uint8))
172	tl.store(
173	key_cache_ptr
174	+ blk * stride_kc_blk
175	+ slot_in_blk * stride_kc_slot
176	+ head * stride_kc_head
177	+ half_offs,
178	k_packed,
179	mask=half_offs < half_k,
180	)
181
182	half_v = head_size_v // 2
183	even_v_mask = even_offs < head_size_v
184	odd_v_mask = odd_offs < head_size_v
185	val_base = value_ptr + tok * stride_val_tok + head * stride_val_head
186
187	v_even = tl.load(val_base + even_offs, mask=even_v_mask, other=0.0).to(tl.float32)
188	v_odd = tl.load(val_base + odd_offs, mask=odd_v_mask, other=0.0).to(tl.float32)
189
190	v_min = tl.minimum(
191	tl.min(tl.where(even_v_mask, v_even, float("inf"))),
192	tl.min(tl.where(odd_v_mask, v_odd, float("inf"))),
193	)
194	v_max = tl.maximum(
195	tl.max(tl.where(even_v_mask, v_even, float("-inf"))),
196	tl.max(tl.where(odd_v_mask, v_odd, float("-inf"))),
197	)
198	v_scale = tl.maximum((v_max - v_min) / 15.0, 1e-6)
199	v_zp_f = tl.clamp(
200	tl.where(
201	-v_min / v_scale >= 0,
202	(-v_min / v_scale + 0.5).to(tl.int32),
203	(-v_min / v_scale - 0.5).to(tl.int32),
204	).to(tl.float32),
205	0.0,
206	15.0,
207	)
208
209	inv_v = 1.0 / v_scale
210	v_even_s = v_even * inv_v + v_zp_f
211	v_odd_s = v_odd * inv_v + v_zp_f
212	v_even_q = tl.clamp(
213	tl.where(
214	v_even_s >= 0,
215	(v_even_s + 0.5).to(tl.int32),
216	(v_even_s - 0.5).to(tl.int32),
217	).to(tl.float32),
218	0.0,
219	15.0,
220	)
221	v_odd_q = tl.clamp(
222	tl.where(
223	v_odd_s >= 0,
224	(v_odd_s + 0.5).to(tl.int32),
225	(v_odd_s - 0.5).to(tl.int32),
226	).to(tl.float32),
227	0.0,
228	15.0,
229	)
230
231	v_zp_int = v_zp_f.to(tl.int32)
232	v_scale_bits = v_scale.to(tl.int32, bitcast=True)
233	v_scale_packed = ((v_scale_bits & -16) \| (v_zp_int & 0xF)).to(
234	tl.float32, bitcast=True
235	)
236
237	tl.store(
238	v_scale_cache_ptr
239	+ blk * stride_vs_blk
240	+ slot_in_blk * stride_vs_slot
241	+ head * stride_vs_head,
242	v_scale_packed,
243	)
244
245	v_packed = pack_int4_nibbles(v_even_q.to(tl.uint8), v_odd_q.to(tl.uint8))
246	tl.store(
247	value_cache_ptr
248	+ blk * stride_vc_blk
249	+ slot_in_blk * stride_vc_slot
250	+ head * stride_vc_head
251	+ half_offs,
252	v_packed,
253	mask=half_offs < half_v,
254	)
255
256
257	def _run_reshape_kernel(
258	kernel,
259	*,
260	key: torch.Tensor,
261	value: torch.Tensor,
262	key_cache: torch.Tensor,
263	value_cache: torch.Tensor,
264	k_scale_cache: torch.Tensor,
265	v_scale_cache: torch.Tensor,
266	slot_mapping: torch.Tensor,
267	packing_factor: int,
268	) -> None:
269	"""Launch the packed INT4 reshape kernel."""
270	num_tokens, num_kv_heads, head_size = key.shape
271	head_size_v = value.shape[2]
272	assert head_size % packing_factor == 0 and head_size_v % packing_factor == 0
273	packed_padded = triton.next_power_of_2(
274	max(head_size, head_size_v) // packing_factor
275	)
276	if current_platform.is_rocm() or current_platform.is_xpu():
277	num_warps = 4
278	else:
279	num_warps = min(16, max(1, packed_padded // 32))
280
281	kernel[(num_tokens, num_kv_heads)](
282	key_ptr=key,
283	value_ptr=value,
284	key_cache_ptr=key_cache,
285	value_cache_ptr=value_cache,
286	k_scale_cache_ptr=k_scale_cache,
287	v_scale_cache_ptr=v_scale_cache,
288	slot_mapping_ptr=slot_mapping,
289	stride_key_tok=key.stride(0),
290	stride_key_head=key.stride(1),
291	stride_val_tok=value.stride(0),
292	stride_val_head=value.stride(1),
293	stride_kc_blk=key_cache.stride(0),
294	stride_kc_slot=key_cache.stride(1),
295	stride_kc_head=key_cache.stride(2),
296	stride_vc_blk=value_cache.stride(0),
297	stride_vc_slot=value_cache.stride(1),
298	stride_vc_head=value_cache.stride(2),
299	stride_ks_blk=k_scale_cache.stride(0),
300	stride_ks_slot=k_scale_cache.stride(1),
301	stride_ks_head=k_scale_cache.stride(2),
302	stride_vs_blk=v_scale_cache.stride(0),
303	stride_vs_slot=v_scale_cache.stride(1),
304	stride_vs_head=v_scale_cache.stride(2),
305	block_size=key_cache.shape[1],
306	head_size=head_size,
307	head_size_v=head_size_v,
308	PACKED_HEAD_PADDED=packed_padded,
309	num_warps=num_warps,
310	)
311
312
313	# ----------------------------------------------------------------------
314	# Read path: split-dot attention over the packed cache
315	# ----------------------------------------------------------------------
316
317
318	@triton.jit
319	def _attn_packed(
320	# Output destinations. In 2D mode the final result is written into
321	# ``output_ptr``; in 3D mode per-segment partials go into the three
322	# ``segm_*`` tensors and ``output_ptr`` is unused.
323	output_ptr,
324	segm_output_ptr,
325	segm_max_ptr,
326	segm_expsum_ptr,
327	query_ptr,
328	key_cache_ptr,
329	value_cache_ptr,
330	sink_ptr,
331	block_tables_ptr,
332	seq_lens_ptr,
333	alibi_slopes_ptr,
334	qq_bias_ptr,
335	scale,
336	out_scale,
337	softcap,
338	k_scale_cache_ptr,
339	v_scale_cache_ptr,
340	num_query_heads: tl.constexpr,
341	num_queries_per_kv: tl.constexpr,
342	block_table_stride: tl.int64,
343	query_stride_0: tl.int64,
344	query_stride_1: tl.int64,
345	output_stride_0: tl.int64,
346	output_stride_1: tl.int64,
347	qq_bias_stride_0: tl.int64,
348	BLOCK_SIZE: tl.constexpr,
349	TILE_SIZE: tl.constexpr,
350	HEAD_SIZE: tl.constexpr,
351	HEAD_SIZE_PADDED: tl.constexpr,
352	PACKED_HEAD_PADDED: tl.constexpr, # HEAD_SIZE / PACKING_FACTOR, rounded up
353	USE_ALIBI_SLOPES: tl.constexpr,
354	USE_ALIBI_SQRT: tl.constexpr,
355	USE_QQ_BIAS: tl.constexpr,
356	USE_SOFTCAP: tl.constexpr,
357	USE_SINKS: tl.constexpr,
358	SLIDING_WINDOW: tl.constexpr,
359	USE_MM_PREFIX: tl.constexpr,
360	MAX_MM_RANGES: tl.constexpr,
361	mm_prefix_range_ptr,
362	stride_k_cache_0: tl.int64,
363	stride_k_cache_1: tl.int64,
364	stride_k_cache_2: tl.int64,
365	stride_k_cache_3: tl.constexpr,
366	stride_v_cache_0: tl.int64,
367	stride_v_cache_1: tl.int64,
368	stride_v_cache_2: tl.int64,
369	stride_v_cache_3: tl.constexpr,
370	stride_ks_blk: tl.int64,
371	stride_ks_slot: tl.int64,
372	stride_ks_head: tl.int64,
373	stride_vs_blk: tl.int64,
374	stride_vs_slot: tl.int64,
375	stride_vs_head: tl.int64,
376	query_start_len_ptr,
377	BLOCK_Q: tl.constexpr,
378	num_seqs: tl.int32,
379	BLOCK_M: tl.constexpr,
380	NUM_SEGMENTS_PER_SEQ: tl.constexpr,
381	USE_FP8: tl.constexpr,
382	IS_3D: tl.constexpr,
383	# 2 → INT4 nibble pair (asymmetric + zp). The packed KV cache stores
384	# one byte per ``packed_offs``, holding PACKING_FACTOR values.
385	PACKING_FACTOR: tl.constexpr,
386	FP8_MIN: tl.constexpr = float8_info.min,
387	FP8_MAX: tl.constexpr = float8_info.max,
388	):
389	# Shared prologue: sequence lookup, q-block bounds, early returns.
390	q_block_global_idx = tl.program_id(0)
391	kv_head_idx = tl.program_id(1)
392	segm_idx = tl.program_id(2) if IS_3D else 0
393
394	(
395	seq_idx,
396	q_block_local_idx,
397	cur_batch_in_all_start_index,
398	cur_batch_query_len,
399	seq_len,
400	) = resolve_seq_and_query_len(
401	query_start_len_ptr, seq_lens_ptr, q_block_global_idx, num_seqs, BLOCK_Q
402	)
403
404	if q_block_local_idx * BLOCK_Q >= cur_batch_query_len:
405	return
406
407	if IS_3D:
408	tiles_per_segment = cdiv_fn(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE)
409	if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len:
410	return
411	else:
412	tiles_per_segment = 0
413
414	offs_m = tl.arange(0, BLOCK_M)
415	offs_t = tl.arange(0, TILE_SIZE)
416	query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv
417
418	query_offset_0 = cur_batch_in_all_start_index + query_pos
419	query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv
420
421	query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1)
422	query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1)
423
424	# Split-Q prologue: PACKING_FACTOR interleaved streams of Q.
425	# INT4 uses 2 streams (even / odd). The packed KV cache stores one
426	# byte per ``packed_offs``, which holds PACKING_FACTOR values.
427	packed_offs = tl.arange(0, PACKED_HEAD_PADDED)
428	offs_s0 = packed_offs * PACKING_FACTOR
429	offs_s1 = packed_offs * PACKING_FACTOR + 1
430	mask_s0 = tl.where(offs_s0 < HEAD_SIZE, 1, 0).to(tl.int1)
431	mask_s1 = tl.where(offs_s1 < HEAD_SIZE, 1, 0).to(tl.int1)
432	packed_dim_mask = tl.where(packed_offs < HEAD_SIZE // PACKING_FACTOR, 1, 0).to(
433	tl.int1
434	)
435	q_base = (
436	query_offset_0[:, None] * query_stride_0
437	+ query_offset_1[:, None] * query_stride_1
438	)
439	q_mask = query_mask_0[:, None] & query_mask_1[:, None]
440	Q_s0 = tl.load(
441	query_ptr + q_base + offs_s0[None, :],
442	mask=mask_s0[None, :] & q_mask,
443	other=0.0,
444	).to(tl.float32)
445	Q_s1 = tl.load(
446	query_ptr + q_base + offs_s1[None, :],
447	mask=mask_s1[None, :] & q_mask,
448	other=0.0,
449	).to(tl.float32)
450
451	# INT4 asymmetric correction needs sum(Q) per row.
452	Q_sum = tl.sum(Q_s0, axis=1) + tl.sum(Q_s1, axis=1)
453
454	block_table_offset = seq_idx * block_table_stride
455
456	# Online-softmax state + optional feature loads.
457	M = init_softmax_M(
458	sink_ptr, query_offset_1, query_mask_1, segm_idx, BLOCK_M, USE_SINKS, IS_3D
459	)
460	L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
461	acc_s0 = tl.zeros([BLOCK_M, PACKED_HEAD_PADDED], dtype=tl.float32)
462	acc_s1 = tl.zeros([BLOCK_M, PACKED_HEAD_PADDED], dtype=tl.float32)
463
464	context_len = seq_len - cur_batch_query_len
465
466	if USE_ALIBI_SLOPES:
467	alibi_slope = tl.load(
468	alibi_slopes_ptr + query_offset_1, mask=query_mask_1, other=0.0
469	)
470
471	if USE_QQ_BIAS:
472	qq_bias_row_ptrs = qq_bias_ptr + query_pos[:, None] * qq_bias_stride_0
473
474	loop_lo, loop_hi, max_seq_prefix_len = compute_tile_loop_bounds(
475	context_len,
476	seq_len,
477	cur_batch_query_len,
478	q_block_local_idx,
479	segm_idx,
480	tiles_per_segment,
481	TILE_SIZE,
482	BLOCK_M,
483	BLOCK_Q,
484	num_queries_per_kv,
485	SLIDING_WINDOW,
486	USE_MM_PREFIX,
487	IS_3D,
488	)
489
490	# Tile loop. Per-tile: load packed KV + scales, dequantize into
491	# PACKING_FACTOR streams, compute the split dot, run the shared
492	# softmax step, accumulate per stream.
493	for j in range(loop_lo, loop_hi):
494	seq_offset = j * TILE_SIZE + offs_t
495	tile_mask = seq_offset < max_seq_prefix_len
496
497	physical_block_idx = tl.load(
498	block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE
499	).to(tl.int64)
500
501	slot_in_blk = seq_offset % BLOCK_SIZE
502	k_off = (
503	physical_block_idx[None, :] * stride_k_cache_0
504	+ kv_head_idx * stride_k_cache_2
505	+ packed_offs[:, None] * stride_k_cache_3
506	+ slot_in_blk[None, :] * stride_k_cache_1
507	)
508	K_packed = tl.load(
509	key_cache_ptr + k_off,
510	mask=packed_dim_mask[:, None] & tile_mask[None, :],
511	other=0,
512	)
513	v_off = (
514	physical_block_idx[:, None] * stride_v_cache_0
515	+ kv_head_idx * stride_v_cache_2
516	+ packed_offs[None, :] * stride_v_cache_3
517	+ slot_in_blk[:, None] * stride_v_cache_1
518	)
519	V_packed = tl.load(
520	value_cache_ptr + v_off,
521	mask=packed_dim_mask[None, :] & tile_mask[:, None],
522	other=0,
523	)
524	# Dequantize KV. INT4 unpacks nibbles as plain uint [0..15];
525	# the zero-point is applied on the score side.
526	K_s0_u, K_s1_u = unpack_int4_nibbles(K_packed)
527	K_s0 = K_s0_u.to(tl.float32)
528	K_s1 = K_s1_u.to(tl.float32)
529	V_s0_u, V_s1_u = unpack_int4_nibbles(V_packed)
530	V_s0 = V_s0_u.to(tl.float32)
531	V_s1 = V_s1_u.to(tl.float32)
532
533	ks_idx = (
534	physical_block_idx * stride_ks_blk
535	+ slot_in_blk * stride_ks_slot
536	+ kv_head_idx * stride_ks_head
537	)
538	ks_raw = tl.load(k_scale_cache_ptr + ks_idx, mask=tile_mask, other=0)
539	vs_idx = (
540	physical_block_idx * stride_vs_blk
541	+ slot_in_blk * stride_vs_slot
542	+ kv_head_idx * stride_vs_head
543	)
544	vs_raw = tl.load(v_scale_cache_ptr + vs_idx, mask=tile_mask, other=0)
545
546	# INT4 steganographs the 4-bit zero-point in the low 4 bits of
547	# the float32 scale's mantissa.
548	ks_bits = ks_raw.to(tl.int32, bitcast=True)
549	k_zp = (ks_bits & 0xF).to(tl.float32)
550	k_token_head_scales = (ks_bits & -16).to(tl.float32, bitcast=True)
551	vs_bits = vs_raw.to(tl.int32, bitcast=True)
552	v_zp = (vs_bits & 0xF).to(tl.float32)
553	v_token_head_scales = (vs_bits & -16).to(tl.float32, bitcast=True)
554
555	query_abs_pos = context_len + query_pos[:, None]
556	seq_mask = compute_kv_seq_mask(
557	query_abs_pos,
558	seq_offset,
559	seq_idx,
560	seq_len,
561	mm_prefix_range_ptr,
562	SLIDING_WINDOW,
563	USE_MM_PREFIX,
564	MAX_MM_RANGES,
565	)
566
567	# Score: split-dot across the 2 INT4 streams; fused
568	# softmax_scale * per-(token, head) k_scale in one mul. INT4
569	# subtracts the ``zp * sum(Q)`` correction term.
570	S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32)
571	raw_dot = tl.dot(Q_s0, K_s0) + tl.dot(Q_s1, K_s1)
572	S += (raw_dot - Q_sum[:, None] * k_zp[None, :]) * (
573	scale * k_token_head_scales[None, :]
574	)
575
576	if USE_SOFTCAP:
577	S = apply_softcap(S, softcap)
578
579	S = tl.where(
580	query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, S, float("-inf")
581	)
582
583	if USE_ALIBI_SLOPES:
584	S = apply_alibi_to_score(
585	S, alibi_slope, seq_offset, context_len, query_pos, USE_ALIBI_SQRT
586	)
587
588	if USE_QQ_BIAS:
589	S += load_qq_bias_tile(
590	qq_bias_row_ptrs, seq_offset, context_len, qq_bias_stride_0
591	)
592
593	M, L, P, alpha = softmax_step(S, M, L)
594	acc_s0 = acc_s0 * alpha[:, None]
595	acc_s1 = acc_s1 * alpha[:, None]
596
597	if SLIDING_WINDOW:
598	qpos_lo = q_block_local_idx * BLOCK_Q
599	sw_mask = (context_len + qpos_lo - seq_offset) < SLIDING_WINDOW
600	V_s0 = tl.where(sw_mask[:, None], V_s0, 0.0)
601	V_s1 = tl.where(sw_mask[:, None], V_s1, 0.0)
602
603	# Fuse v per-(token, head) scale into P. INT4 also subtracts
604	# the v-zero-point contribution from each stream once.
605	P_v = (P * v_token_head_scales[None, :]).to(tl.float32)
606	Pv_zp_sum = tl.sum(P_v * v_zp[None, :], axis=1)
607	acc_s0 += tl.dot(P_v, V_s0) - Pv_zp_sum[:, None]
608	acc_s1 += tl.dot(P_v, V_s1) - Pv_zp_sum[:, None]
609
610	# Epilogue. 2D writes the final output with optional FP8 clamp;
611	# 3D writes the per-segment partials (output / max / expsum) for
612	# ``reduce_segments`` to finalize. Each stream writes its own
613	# stripe in the output layout.
614	out_mask = query_mask_0[:, None] & query_mask_1[:, None]
615	if IS_3D:
616	segm_base = (
617	query_offset_0[:, None].to(tl.int64)
618	* (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
619	+ query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
620	+ segm_idx * HEAD_SIZE_PADDED
621	)
622	tl.store(
623	segm_output_ptr + segm_base + offs_s0[None, :],
624	acc_s0,
625	mask=mask_s0[None, :] & out_mask,
626	)
627	tl.store(
628	segm_output_ptr + segm_base + offs_s1[None, :],
629	acc_s1,
630	mask=mask_s1[None, :] & out_mask,
631	)
632	store_segm_reduce_scalars(
633	segm_max_ptr,
634	segm_expsum_ptr,
635	query_offset_0,
636	query_offset_1,
637	segm_idx,
638	M,
639	L,
640	query_mask_0,
641	query_mask_1,
642	num_query_heads,
643	NUM_SEGMENTS_PER_SEQ,
644	)
645	else:
646	acc_s0 = acc_s0 / L[:, None]
647	acc_s1 = acc_s1 / L[:, None]
648	if USE_FP8:
649	out_s = tl.load(out_scale)
650	acc_s0 = tl.clamp(acc_s0 * out_s, FP8_MIN, FP8_MAX)
651	acc_s1 = tl.clamp(acc_s1 * out_s, FP8_MIN, FP8_MAX)
652	out_base = (
653	query_offset_0[:, None] * output_stride_0
654	+ query_offset_1[:, None] * output_stride_1
655	)
656	tl.store(
657	output_ptr + out_base + offs_s0[None, :],
658	acc_s0,
659	mask=mask_s0[None, :] & out_mask,
660	)
661	tl.store(
662	output_ptr + out_base + offs_s1[None, :],
663	acc_s1,
664	mask=mask_s1[None, :] & out_mask,
665	)
666
667
668	def _launch_packed_attn(
669	*,
670	q,
671	k_cache,
672	v_cache,
673	out,
674	cu_seqlens_q,
675	max_seqlen_q,
676	seqused_k,
677	softmax_scale,
678	window_size,
679	block_table,
680	softcap,
681	sinks,
682	alibi_slopes,
683	use_alibi_sqrt,
684	qq_bias,
685	output_scale,
686	mm_prefix_range,
687	k_scale_cache,
688	v_scale_cache,
689	seq_threshold_3D,
690	num_par_softmax_segments,
691	softmax_segm_output,
692	softmax_segm_max,
693	softmax_segm_expsum,
694	packing_factor: int,
695	):
696	"""Launch ``_attn_packed`` for one of the sub-byte modes.
697
698	Handles 2D-vs-3D dispatch, placeholder pointers for the unused side
699	of that split, and the trailing ``reduce_segments`` pass. Writes
700	into ``out`` (directly for 2D; via the segm buffers for 3D).
701	"""
702	import vllm.envs as envs
703	from vllm.v1.attention.ops.triton_unified_attention import _get_tile_size
704
705	is_batch_invariant = envs.VLLM_BATCH_INVARIANT
706
707	use_mm_prefix = False
708	max_mm_ranges = 0
709	if mm_prefix_range is not None:
710	assert mm_prefix_range.ndim == 3, (
711	f"Unsupported mm_prefix_range shape: {mm_prefix_range.shape}"
712	)
713	use_mm_prefix = True
714	max_mm_ranges = mm_prefix_range.shape[1]
715
716	block_size = v_cache.shape[1]
717	num_seqs = len(seqused_k)
718	num_query_heads = q.shape[1]
719	num_kv_heads = k_cache.shape[2]
720	num_queries_per_kv = num_query_heads // num_kv_heads
721	head_size = q.shape[2]
722
723	BLOCK_M = (
724	16 if num_queries_per_kv <= 16 else triton.next_power_of_2(num_queries_per_kv)
725	)
726	BLOCK_Q = BLOCK_M // num_queries_per_kv
727	total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs
728	sliding_window_val = 1 + window_size[0] if window_size[0] >= 0 else 0
729	TILE_SIZE_PREFILL = _get_tile_size(
730	head_size, sliding_window_val, q.element_size(), is_prefill=True
731	)
732	TILE_SIZE_DECODE = _get_tile_size(
733	head_size, sliding_window_val, q.element_size(), is_prefill=False
734	)
735
736	use_3d = not (
737	seq_threshold_3D is None
738	or num_par_softmax_segments is None
739	or softmax_segm_output is None
740	or softmax_segm_max is None
741	or softmax_segm_expsum is None
742	or max_seqlen_q > 1
743	or num_seqs > seq_threshold_3D
744	or is_batch_invariant
745	)
746
747	# 3D never reads ``output_ptr`` and 2D never reads the segm tensors,
748	# but Triton needs a non-null pointer everywhere; reuse ``out`` as
749	# the placeholder for the unused side.
750	segm_output_ptr = softmax_segm_output if use_3d else out
751	segm_max_ptr = softmax_segm_max if use_3d else out
752	segm_expsum_ptr = softmax_segm_expsum if use_3d else out
753	num_segments = num_par_softmax_segments if use_3d else 1
754
755	grid: tuple[Any, ...]
756	if use_3d:
757	grid = (total_num_q_blocks, num_kv_heads, num_par_softmax_segments)
758	tile_size = TILE_SIZE_DECODE
759	else:
760	grid = (total_num_q_blocks, num_kv_heads)
761	tile_size = TILE_SIZE_PREFILL
762
763	_attn_packed[grid](
764	output_ptr=out,
765	segm_output_ptr=segm_output_ptr,
766	segm_max_ptr=segm_max_ptr,
767	segm_expsum_ptr=segm_expsum_ptr,
768	query_ptr=q,
769	key_cache_ptr=k_cache,
770	value_cache_ptr=v_cache,
771	sink_ptr=sinks,
772	block_tables_ptr=block_table,
773	seq_lens_ptr=seqused_k,
774	alibi_slopes_ptr=alibi_slopes,
775	qq_bias_ptr=qq_bias,
776	scale=softmax_scale,
777	out_scale=1 / output_scale if output_scale is not None else 1.0,
778	softcap=softcap,
779	k_scale_cache_ptr=k_scale_cache,
780	v_scale_cache_ptr=v_scale_cache,
781	num_query_heads=num_query_heads,
782	num_queries_per_kv=num_queries_per_kv,
783	block_table_stride=block_table.stride(0),
784	query_stride_0=q.stride(0),
785	query_stride_1=q.stride(1),
786	output_stride_0=out.stride(0),
787	output_stride_1=out.stride(1),
788	qq_bias_stride_0=qq_bias.stride(0) if qq_bias is not None else 0,
789	BLOCK_SIZE=block_size,
790	TILE_SIZE=tile_size,
791	HEAD_SIZE=head_size,
792	HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
793	PACKED_HEAD_PADDED=triton.next_power_of_2(head_size) // packing_factor,
794	USE_ALIBI_SLOPES=alibi_slopes is not None,
795	USE_ALIBI_SQRT=use_alibi_sqrt,
796	USE_QQ_BIAS=qq_bias is not None,
797	USE_SOFTCAP=(softcap > 0),
798	USE_SINKS=(sinks is not None),
799	SLIDING_WINDOW=(1 + window_size[0]),
800	USE_MM_PREFIX=use_mm_prefix,
801	MAX_MM_RANGES=max_mm_ranges,
802	mm_prefix_range_ptr=mm_prefix_range,
803	stride_k_cache_0=k_cache.stride(0),
804	stride_k_cache_1=k_cache.stride(1),
805	stride_k_cache_2=k_cache.stride(2),
806	stride_k_cache_3=k_cache.stride(3),
807	stride_v_cache_0=v_cache.stride(0),
808	stride_v_cache_1=v_cache.stride(1),
809	stride_v_cache_2=v_cache.stride(2),
810	stride_v_cache_3=v_cache.stride(3),
811	stride_ks_blk=k_scale_cache.stride(0),
812	stride_ks_slot=k_scale_cache.stride(1),
813	stride_ks_head=k_scale_cache.stride(2),
814	stride_vs_blk=v_scale_cache.stride(0),
815	stride_vs_slot=v_scale_cache.stride(1),
816	stride_vs_head=v_scale_cache.stride(2),
817	query_start_len_ptr=cu_seqlens_q,
818	BLOCK_Q=BLOCK_Q,
819	num_seqs=num_seqs,
820	BLOCK_M=BLOCK_M,
821	NUM_SEGMENTS_PER_SEQ=num_segments,
822	USE_FP8=output_scale is not None,
823	IS_3D=use_3d,
824	PACKING_FACTOR=packing_factor,
825	)
826
827	if use_3d:
828	reduce_segments[(q.shape[0], num_query_heads)](
829	output_ptr=out,
830	segm_output_ptr=softmax_segm_output,
831	segm_max_ptr=softmax_segm_max,
832	segm_expsum_ptr=softmax_segm_expsum,
833	seq_lens_ptr=seqused_k,
834	num_seqs=num_seqs,
835	num_query_heads=num_query_heads,
836	out_scale_inv=1 / output_scale if output_scale is not None else 1.0,
837	output_stride_0=out.stride(0),
838	output_stride_1=out.stride(1),
839	block_table_stride=block_table.stride(0),
840	TILE_SIZE=TILE_SIZE_DECODE,
841	HEAD_SIZE=head_size,
842	HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
843	query_start_len_ptr=cu_seqlens_q,
844	BLOCK_Q=BLOCK_Q,
845	NUM_SEGMENTS_PER_SEQ=num_par_softmax_segments,
846	USE_FP8=output_scale is not None,
847	)
848
849
850	# ----------------------------------------------------------------------
851	# Public entry points
852	# ----------------------------------------------------------------------
853
854
855	def reshape_and_cache_int4(
856	key: torch.Tensor,
857	value: torch.Tensor,
858	key_cache: torch.Tensor,
859	value_cache: torch.Tensor,
860	slot_mapping: torch.Tensor,
861	*,
862	k_scale_cache: torch.Tensor,
863	v_scale_cache: torch.Tensor,
864	) -> None:
865	"""Pre-rotate (RHT), pack to INT4 and write into the paged cache."""
866	key = single_rht(key.float()).to(key.dtype)
867	value = single_rht(value.float()).to(value.dtype)
868	_run_reshape_kernel(
869	_reshape_cache_int4_kernel,
870	key=key,
871	value=value,
872	key_cache=key_cache,
873	value_cache=value_cache,
874	k_scale_cache=k_scale_cache,
875	v_scale_cache=v_scale_cache,
876	slot_mapping=slot_mapping,
877	packing_factor=_INT4_PACKING_FACTOR,
878	)
879
880
881	def unified_attention_int4(
882	q: torch.Tensor,
883	k_cache: torch.Tensor,
884	v_cache: torch.Tensor,
885	out: torch.Tensor,
886	*,
887	cu_seqlens_q: torch.Tensor,
888	max_seqlen_q: int,
889	seqused_k: torch.Tensor,
890	max_seqlen_k: int,
891	softmax_scale: float,
892	window_size: tuple[int, int],
893	block_table: torch.Tensor,
894	softcap: float,
895	sinks: torch.Tensor \| None,
896	alibi_slopes: torch.Tensor \| None,
897	use_alibi_sqrt: bool,
898	qq_bias: torch.Tensor \| None,
899	output_scale: torch.Tensor \| None,
900	mm_prefix_range: torch.Tensor \| None,
901	k_scale_cache: torch.Tensor,
902	v_scale_cache: torch.Tensor,
903	seq_threshold_3D: int \| None = None,
904	num_par_softmax_segments: int \| None = None,
905	softmax_segm_output: torch.Tensor \| None = None,
906	softmax_segm_max: torch.Tensor \| None = None,
907	softmax_segm_expsum: torch.Tensor \| None = None,
908	) -> None:
909	"""Paged attention over the INT4 packed cache, writing into out.
910
911	The forward RHT has norm ``sqrt(head_size)``, so ``softmax_scale`` is
912	divided by ``head_size`` and the inverse RHT divides the output by
913	``head_size`` as well.
914	"""
915	q_orig_dtype = q.dtype
916	q = single_rht(q.float()).to(q_orig_dtype)
917	head_size = q.shape[2]
918	softmax_scale = softmax_scale / head_size
919
920	_launch_packed_attn(
921	q=q,
922	k_cache=k_cache,
923	v_cache=v_cache,
924	out=out,
925	cu_seqlens_q=cu_seqlens_q,
926	max_seqlen_q=max_seqlen_q,
927	seqused_k=seqused_k,
928	softmax_scale=softmax_scale,
929	window_size=window_size,
930	block_table=block_table,
931	softcap=softcap,
932	sinks=sinks,
933	alibi_slopes=alibi_slopes,
934	use_alibi_sqrt=use_alibi_sqrt,
935	qq_bias=qq_bias,
936	output_scale=output_scale,
937	mm_prefix_range=mm_prefix_range,
938	k_scale_cache=k_scale_cache,
939	v_scale_cache=v_scale_cache,
940	seq_threshold_3D=seq_threshold_3D,
941	num_par_softmax_segments=num_par_softmax_segments,
942	softmax_segm_output=softmax_segm_output,
943	softmax_segm_max=softmax_segm_max,
944	softmax_segm_expsum=softmax_segm_expsum,
945	packing_factor=_INT4_PACKING_FACTOR,
946	)
947
948	out_f = single_rht(out.float(), inverse=True) / head_size
949	out.copy_(out_f.to(q_orig_dtype))
950
951
952	# ----------------------------------------------------------------------
953	# Randomized Hadamard Transform (RHT) — gaussianizes K/V before INT4
954	# quantization; applied on write and to Q before the read kernels.
955	# ----------------------------------------------------------------------
956
957	# Hadacore (CUDA tensor core kernel) availability check
958	# Hadacore's CUDA impl is only registered when built for sm_80+, but the
959	# schema def is unconditional — on ROCm ``hasattr`` is True yet dispatch
960	# would crash, so we also gate on ``is_cuda()`` and the sm_80 capability.
961	_HADACORE_AVAILABLE: bool \| None = None
962
963
964	def _hadacore_available() -> bool:
965	global _HADACORE_AVAILABLE
966	if _HADACORE_AVAILABLE is None:
967	_HADACORE_AVAILABLE = (
968	current_platform.is_cuda()
969	and current_platform.has_device_capability(80)
970	and hasattr(torch.ops._C, "hadacore_transform")
971	)
972	return _HADACORE_AVAILABLE
973
974
975	# Cached Hadamard matrices (one per (size, dtype, device) tuple)
976	_HADAMARD_MATRIX_CACHE: dict[tuple[int, torch.dtype, str], torch.Tensor] = {}
977
978
979	def _get_hadamard_matrix(
980	d: int, dtype: torch.dtype, device: torch.device
981	) -> torch.Tensor:
982	key = (d, dtype, str(device))
983	cached = _HADAMARD_MATRIX_CACHE.get(key)
984	if cached is None:
985	H = torch.ones(1, 1, dtype=torch.float32, device=device)
986	while H.shape[0] < d:
987	H = torch.cat(
988	[
989	torch.cat([H, H], dim=1),
990	torch.cat([H, -H], dim=1),
991	],
992	dim=0,
993	)
994	cached = H.to(dtype).contiguous()
995	_HADAMARD_MATRIX_CACHE[key] = cached
996	return cached
997
998
999	# Triton MMA Hadamard kernel (Tier 2)
1000	@triton.jit
1001	def _hadamard_mma_kernel(
1002	x_ptr,
1003	h_ptr,
1004	out_ptr,
1005	n_rows,
1006	stride_x_row: tl.int64,
1007	stride_x_col: tl.int64,
1008	stride_o_row: tl.int64,
1009	stride_o_col: tl.int64,
1010	BLOCK_M: tl.constexpr,
1011	D: tl.constexpr,
1012	):
1013	pid = tl.program_id(0)
1014	rows = pid * BLOCK_M + tl.arange(0, BLOCK_M)
1015	cols = tl.arange(0, D)
1016	row_mask = rows < n_rows
1017
1018	x = tl.load(
1019	x_ptr + rows[:, None] * stride_x_row + cols[None, :] * stride_x_col,
1020	mask=row_mask[:, None],
1021	other=0.0,
1022	)
1023	H = tl.load(h_ptr + cols[:, None] * D + cols[None, :])
1024
1025	out = tl.dot(x, H, out_dtype=tl.float32).to(x.dtype)
1026
1027	tl.store(
1028	out_ptr + rows[:, None] * stride_o_row + cols[None, :] * stride_o_col,
1029	out,
1030	mask=row_mask[:, None],
1031	)
1032
1033
1034	# H is D×D bf16 = 2·D² bytes of LDS. AMD CDNA has 64 KiB LDS, so D ≤ 128
1035	# (32 KiB) leaves room for input + accumulator. Larger D falls back.
1036	_TRITON_HADAMARD_MIN_D = 16
1037	_TRITON_HADAMARD_MAX_D = 128
1038
1039
1040	def _triton_hadamard_transform(x: torch.Tensor) -> torch.Tensor:
1041	d = x.shape[-1]
1042	orig_shape = x.shape
1043	orig_dtype = x.dtype
1044
1045	work_dtype = torch.bfloat16 if orig_dtype == torch.float32 else orig_dtype
1046	x2d = x.contiguous().to(work_dtype).reshape(-1, d)
1047	out2d = torch.empty_like(x2d)
1048	n_rows = x2d.shape[0]
1049	H_mat = _get_hadamard_matrix(d, work_dtype, x.device)
1050
1051	BLOCK_M = 16
1052	grid = (triton.cdiv(n_rows, BLOCK_M),)
1053	# num_stages=1: the kernel has no loop, so default 3-stage pipelining
1054	# would triple-buffer H and blow the AMD LDS budget.
1055	_hadamard_mma_kernel[grid](
1056	x2d,
1057	H_mat,
1058	out2d,
1059	n_rows,
1060	x2d.stride(0),
1061	x2d.stride(1),
1062	out2d.stride(0),
1063	out2d.stride(1),
1064	BLOCK_M=BLOCK_M,
1065	D=d,
1066	num_stages=1,
1067	num_warps=4,
1068	)
1069	return out2d.reshape(orig_shape).to(orig_dtype)
1070
1071
1072	# Public API
1073	def fast_hadamard_transform(x: torch.Tensor) -> torch.Tensor:
1074	"""Unnormalized Walsh-Hadamard Transform along the last dimension.
1075
1076	H_d × x where H_d × H_d = d × I. Last dim must be a power of 2.
1077
1078	Three-tier dispatch:
1079	1. Hadacore CUDA Tensor Core kernel (sm_80+).
1080	2. Triton MMA matmul kernel (CUDA fallback + ROCm MFMA/WMMA path).
1081	3. PyTorch butterfly (CPU and any GPU/dtype combo Triton can't take).
1082	"""
1083	d = x.shape[-1]
1084	assert d & (d - 1) == 0, f"Requires power-of-2 dim, got {d}"
1085
1086	# Tier 1 — hadacore on CUDA.
1087	if _hadacore_available() and 0 < d <= (1 << 15):
1088	from vllm import _custom_ops as ops
1089
1090	# hadacore returns x @ (H/√d); rescale to the unnormalized H × x
1091	# convention the INT4 scale math is calibrated to.
1092	rescale = d**0.5
1093	if x.dtype in (torch.float16, torch.bfloat16):
1094	y = ops.hadacore_transform(x.contiguous().clone(), inplace=True)
1095	return y * rescale
1096	# fp32 → bf16 round-trip; precision loss is irrelevant before
1097	# INT4 quantization.
1098	orig_dtype = x.dtype
1099	x_bf16 = x.contiguous().to(torch.bfloat16)
1100	y_bf16 = ops.hadacore_transform(x_bf16, inplace=True)
1101	return y_bf16.to(orig_dtype) * rescale
1102
1103	# Tier 2 — Triton MMA kernel (covers ROCm via MFMA/WMMA codegen, and
1104	# also CUDA when hadacore is unavailable).
1105	if (
1106	x.is_cuda
1107	and _TRITON_HADAMARD_MIN_D <= d <= _TRITON_HADAMARD_MAX_D
1108	and x.dtype in (torch.float16, torch.bfloat16, torch.float32)
1109	):
1110	return _triton_hadamard_transform(x)
1111
1112	# Tier 3 — PyTorch butterfly (CPU / unsupported dtype / D < 16).
1113	h = 1
1114	while h < d:
1115	xv = x.view(x.shape[:-1], d // (2 h), 2, h)
1116	a = xv[..., 0, :]
1117	b = xv[..., 1, :]
1118	x = torch.stack([a + b, a - b], dim=-2).reshape(x.shape)
1119	h <<= 1
1120	return x
1121
1122
1123	# Randomized Hadamard Transform (used by INT4)
1124	# Deterministic ±1 signs for Randomized Hadamard Transform.
1125	# RHT = H × D × x (sign flip + Hadamard). Breaks residual structure
1126	# in KV vectors, improving quantization quality.
1127	_RHT_SIGNS_CACHE: dict[tuple[int, int, str], torch.Tensor] = {}
1128
1129
1130	def _get_rht_signs(d: int, round_idx: int, device: torch.device) -> torch.Tensor:
1131	"""Return a cached deterministic ±1 sign vector of length d."""
1132	key = (d, round_idx, str(device))
1133	if key not in _RHT_SIGNS_CACHE:
1134	gen = torch.Generator(device="cpu")
1135	gen.manual_seed(0x9E3779B9 + round_idx * 0x517CC1B7)
1136	signs = (
1137	2.0 * torch.bernoulli(torch.full((d,), 0.5, device="cpu"), generator=gen)
1138	- 1.0
1139	)
1140	_RHT_SIGNS_CACHE[key] = signs.to(device)
1141	return _RHT_SIGNS_CACHE[key]
1142
1143
1144	def single_rht(x: torch.Tensor, inverse: bool = False) -> torch.Tensor:
1145	"""Single Randomized Hadamard Transform: H × D₁ × x.
1146
1147	Used by INT4 per-token-head quantization to gaussianize data
1148	before asymmetric quantization.
1149	"""
1150	d = x.shape[-1]
1151	d1 = _get_rht_signs(d, 0, x.device)
1152	if inverse:
1153	return fast_hadamard_transform(x) * d1
1154	else:
1155	return fast_hadamard_transform(x * d1)