@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 447 lines 16.5 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Fused Triton kernels for TurboQuant KV store.
4
5 Two kernels:
6 1. _tq_fused_store_fp8: FP8 key scatter + value uniform quantization.
7 2. _tq_fused_store_mse: Fused binary-search bucketize + MSE index
8 packing + value quantization.
9
10 The launcher `triton_turboquant_store` selects the appropriate kernel.
11 """
12
13 import math
14
15 import torch
16
17 from vllm.triton_utils import tl, triton
18 from vllm.v1.attention.ops.triton_turboquant_decode import _use_fp8_e4b15
19
20 # ═══════════════════════════════════════════════════════════════════════
21 # Shared: value uniform quantization + pack + scale/zero store
22 # ═══════════════════════════════════════════════════════════════════════
23
24
25 @triton.jit
26 def _store_quantized_value(
27 Value_ptr,
28 KV_cache_ptr,
29 base, # pid * D offset into Value_ptr
30 slot_base, # byte offset into KV_cache_ptr for this slot+head
31 d_offs, # tl.arange(0, BLOCK_D)
32 d_mask, # d_offs < D
33 D: tl.constexpr,
34 KPS: tl.constexpr,
35 VQB: tl.constexpr,
36 VAL_DATA_BYTES: tl.constexpr,
37 BLOCK_D: tl.constexpr,
38 BLOCK_VAL: tl.constexpr,
39 BLOCK_GRP: tl.constexpr,
40 ):
41 """Uniform quantization of values to VQB bits, pack, and store with scale/zero."""
42 val_cache_offset = KPS
43
44 if VQB == 3:
45 val_vec = tl.load(Value_ptr + base + d_offs, mask=d_mask, other=0.0).to(
46 tl.float32
47 )
48 val_min = tl.min(tl.where(d_mask, val_vec, float("inf")), axis=0)
49 val_max = tl.max(tl.where(d_mask, val_vec, -float("inf")), axis=0)
50 v_scale = (val_max - val_min) / 7.0
51 v_scale = tl.where(v_scale > 1e-8, v_scale, 1e-8)
52
53 q_vals = tl.minimum(
54 tl.maximum(((val_vec - val_min) / v_scale + 0.5).to(tl.int32), 0), 7
55 )
56
57 grp_offs = tl.arange(0, BLOCK_GRP)
58 grp_mask = grp_offs < (D // 8)
59 q_grp = tl.reshape(q_vals, [BLOCK_GRP, 8])
60 shifts_3bit = tl.arange(0, 8) * 3
61 packed_24 = tl.sum(q_grp << shifts_3bit[None, :], axis=1)
62 b0 = (packed_24 & 0xFF).to(tl.uint8)
63 b1 = ((packed_24 >> 8) & 0xFF).to(tl.uint8)
64 b2 = ((packed_24 >> 16) & 0xFF).to(tl.uint8)
65 tl.store(
66 KV_cache_ptr + slot_base + val_cache_offset + grp_offs * 3,
67 b0,
68 mask=grp_mask,
69 )
70 tl.store(
71 KV_cache_ptr + slot_base + val_cache_offset + grp_offs * 3 + 1,
72 b1,
73 mask=grp_mask,
74 )
75 tl.store(
76 KV_cache_ptr + slot_base + val_cache_offset + grp_offs * 3 + 2,
77 b2,
78 mask=grp_mask,
79 )
80
81 sc_offset = val_cache_offset + VAL_DATA_BYTES
82 sc_f16 = v_scale.to(tl.float16)
83 sc_u16 = sc_f16.to(tl.uint16, bitcast=True)
84 tl.store(KV_cache_ptr + slot_base + sc_offset, (sc_u16 & 0xFF).to(tl.uint8))
85 tl.store(
86 KV_cache_ptr + slot_base + sc_offset + 1,
87 ((sc_u16 >> 8) & 0xFF).to(tl.uint8),
88 )
89 zr_f16 = val_min.to(tl.float16)
90 zr_u16 = zr_f16.to(tl.uint16, bitcast=True)
91 tl.store(KV_cache_ptr + slot_base + sc_offset + 2, (zr_u16 & 0xFF).to(tl.uint8))
92 tl.store(
93 KV_cache_ptr + slot_base + sc_offset + 3,
94 ((zr_u16 >> 8) & 0xFF).to(tl.uint8),
95 )
96
97 else: # VQB == 4
98 val_vec = tl.load(Value_ptr + base + d_offs, mask=d_mask, other=0.0).to(
99 tl.float32
100 )
101 val_min = tl.min(tl.where(d_mask, val_vec, float("inf")), axis=0)
102 val_max = tl.max(tl.where(d_mask, val_vec, -float("inf")), axis=0)
103 v_scale = (val_max - val_min) / 15.0
104 v_scale = tl.where(v_scale > 1e-8, v_scale, 1e-8)
105
106 # Quantize all D elements from register (no re-load)
107 q_all = tl.minimum(
108 tl.maximum(((val_vec - val_min) / v_scale + 0.5).to(tl.int32), 0), 15
109 )
110 # Reshape to pairs and pack two 4-bit values per byte
111 q_pairs = tl.reshape(q_all, [BLOCK_D // 2, 2])
112 shifts_4 = tl.arange(0, 2) * 4
113 packed_val = tl.sum((q_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8)
114 val_offs = tl.arange(0, BLOCK_D // 2)
115 val_mask = val_offs < VAL_DATA_BYTES
116 tl.store(
117 KV_cache_ptr + slot_base + val_cache_offset + val_offs,
118 packed_val,
119 mask=val_mask,
120 )
121
122 sc_offset = val_cache_offset + VAL_DATA_BYTES
123 sc_f16 = v_scale.to(tl.float16)
124 sc_u16 = sc_f16.to(tl.uint16, bitcast=True)
125 tl.store(KV_cache_ptr + slot_base + sc_offset, (sc_u16 & 0xFF).to(tl.uint8))
126 tl.store(
127 KV_cache_ptr + slot_base + sc_offset + 1,
128 ((sc_u16 >> 8) & 0xFF).to(tl.uint8),
129 )
130 zr_f16 = val_min.to(tl.float16)
131 zr_u16 = zr_f16.to(tl.uint16, bitcast=True)
132 tl.store(KV_cache_ptr + slot_base + sc_offset + 2, (zr_u16 & 0xFF).to(tl.uint8))
133 tl.store(
134 KV_cache_ptr + slot_base + sc_offset + 3,
135 ((zr_u16 >> 8) & 0xFF).to(tl.uint8),
136 )
137
138
139 # ═══════════════════════════════════════════════════════════════════════
140 # FP8 key store + value uniform quantization
141 # ═══════════════════════════════════════════════════════════════════════
142
143
144 @triton.jit
145 def _tq_fused_store_fp8(
146 Key_ptr, # [NH, D] float16/bfloat16 — raw keys
147 Value_ptr, # [NH, D] float16/bfloat16 — raw values
148 KV_cache_ptr, # [total_bytes] uint8 (flattened view)
149 Slot_mapping_ptr, # [N] int32 — per-token slot indices
150 # Cache strides (for computing byte offsets)
151 stride_cache_block: tl.constexpr,
152 stride_cache_pos: tl.constexpr,
153 stride_cache_head: tl.constexpr,
154 # Dimensions
155 D: tl.constexpr,
156 H: tl.constexpr,
157 BLOCK_SIZE: tl.constexpr,
158 BLOCK_D: tl.constexpr,
159 # TQ layout
160 KPS: tl.constexpr,
161 # Value quantization
162 VQB: tl.constexpr,
163 VAL_DATA_BYTES: tl.constexpr,
164 # Packing block sizes
165 BLOCK_VAL: tl.constexpr,
166 BLOCK_GRP: tl.constexpr = 16,
167 FP8_E4B15: tl.constexpr = 0, # 1 = e4b15 (Ampere/Ada), 0 = e4nv (Hopper+)
168 ):
169 """FP8 key cast+scatter + value uniform quantization."""
170 pid = tl.program_id(0)
171 token_idx = pid // H
172 head_idx = pid % H
173
174 slot = tl.load(Slot_mapping_ptr + token_idx)
175 if slot < 0:
176 return
177 blk = (slot // BLOCK_SIZE).to(tl.int64)
178 off = (slot % BLOCK_SIZE).to(tl.int64)
179 head_idx_i64 = tl.cast(head_idx, tl.int64)
180 slot_base = (
181 blk * stride_cache_block
182 + off * stride_cache_pos
183 + head_idx_i64 * stride_cache_head
184 )
185
186 base = pid * D
187
188 # ── FP8 KEY: cast to FP8 in-kernel and store ─────────────────
189 d_offs = tl.arange(0, BLOCK_D)
190 d_mask = d_offs < D
191 k_vals = tl.load(Key_ptr + base + d_offs, mask=d_mask, other=0.0).to(tl.float32)
192 k_fp8 = k_vals.to(tl.float8e4b15) if FP8_E4B15 else k_vals.to(tl.float8e4nv)
193 k_bytes = k_fp8.to(tl.uint8, bitcast=True)
194 tl.store(KV_cache_ptr + slot_base + d_offs, k_bytes, mask=d_mask)
195
196 # ── VALUE QUANTIZE + PACK ───────────────────────────────────────
197 _store_quantized_value(
198 Value_ptr,
199 KV_cache_ptr,
200 base,
201 slot_base,
202 d_offs,
203 d_mask,
204 D=D,
205 KPS=KPS,
206 VQB=VQB,
207 VAL_DATA_BYTES=VAL_DATA_BYTES,
208 BLOCK_D=BLOCK_D,
209 BLOCK_VAL=BLOCK_VAL,
210 BLOCK_GRP=BLOCK_GRP,
211 )
212
213
214 # ═══════════════════════════════════════════════════════════════════════
215 # Fused MSE store: bucketize + MSE index pack + norm store + value pack
216 # (eliminates 4 PyTorch kernel launches per layer vs pack-only kernel)
217 # ═══════════════════════════════════════════════════════════════════════
218
219
220 @triton.jit
221 def _tq_fused_store_mse(
222 # Post-rotation inputs
223 Y_ptr, # [NH, D] float32 — rotated normalized keys (x_hat @ PiT)
224 Norms_ptr, # [NH] float32 — key vector norms (||k||)
225 Value_ptr, # [NH, D] float32 — raw values
226 # Quantization tables
227 Midpoints_ptr, # [n_centroids-1] float32
228 # Cache and indexing
229 KV_cache_ptr, # [total_bytes] uint8 (flattened view)
230 Slot_mapping_ptr, # [N] int32 — per-token slot indices
231 # Cache strides
232 stride_cache_block: tl.constexpr,
233 stride_cache_pos: tl.constexpr,
234 stride_cache_head: tl.constexpr,
235 # Dimensions
236 D: tl.constexpr,
237 H: tl.constexpr,
238 BLOCK_SIZE: tl.constexpr,
239 BLOCK_D: tl.constexpr,
240 # TQ layout
241 MSE_BYTES: tl.constexpr,
242 KPS: tl.constexpr,
243 # Value quantization
244 VQB: tl.constexpr,
245 VAL_DATA_BYTES: tl.constexpr,
246 # Packing block sizes
247 BLOCK_VAL: tl.constexpr,
248 # MSE params
249 MSE_BITS: tl.constexpr,
250 N_CENTROIDS: tl.constexpr,
251 BLOCK_GRP: tl.constexpr = 16,
252 ):
253 """Fused MSE quantize + pack + store.
254
255 Performs binary-search bucketize, MSE index packing, norm storage,
256 and value quantization in one kernel.
257 """
258 pid = tl.program_id(0)
259 token_idx = pid // H
260 head_idx = pid % H
261
262 slot = tl.load(Slot_mapping_ptr + token_idx)
263 if slot < 0:
264 return
265 blk = (slot // BLOCK_SIZE).to(tl.int64)
266 off = (slot % BLOCK_SIZE).to(tl.int64)
267 head_idx_i64 = tl.cast(head_idx, tl.int64)
268 slot_base = (
269 blk * stride_cache_block
270 + off * stride_cache_pos
271 + head_idx_i64 * stride_cache_head
272 )
273
274 base = pid * D
275 d_offs = tl.arange(0, BLOCK_D)
276 d_mask = d_offs < D
277
278 # ── 1. BINARY SEARCH BUCKETIZE ───────────────────────────────────
279 # Midpoints are sorted (N_CENTROIDS-1 values); binary search finds
280 # insertion point in MSE_BITS iterations vs N_CENTROIDS-1 for linear.
281 y_vec = tl.load(Y_ptr + base + d_offs, mask=d_mask, other=0.0)
282 lo = tl.zeros([BLOCK_D], dtype=tl.int32)
283 hi = tl.full([BLOCK_D], N_CENTROIDS - 1, dtype=tl.int32)
284 for _ in range(MSE_BITS):
285 mid = (lo + hi) >> 1
286 # Clamp to valid midpoint index [0, N_CENTROIDS-2] for load safety;
287 # the search result (lo) is still correct since converged lanes
288 # don't change.
289 safe_mid = tl.minimum(mid, N_CENTROIDS - 2)
290 mid_val = tl.load(Midpoints_ptr + safe_mid, mask=d_mask, other=0.0)
291 lo = tl.where(y_vec >= mid_val, mid + 1, lo)
292 hi = tl.where(y_vec >= mid_val, hi, mid)
293 idx = tl.minimum(lo, N_CENTROIDS - 1)
294
295 # ── 2. PACK MSE INDICES from register idx ─────────────────────────
296 if MSE_BITS == 4:
297 idx_pairs = tl.reshape(idx, [BLOCK_D // 2, 2])
298 shifts_4 = tl.arange(0, 2) * 4
299 packed = tl.sum((idx_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8)
300 mse_offs = tl.arange(0, BLOCK_D // 2)
301 mse_mask = mse_offs < MSE_BYTES
302 tl.store(KV_cache_ptr + slot_base + mse_offs, packed, mask=mse_mask)
303
304 elif MSE_BITS == 3:
305 grp_offs = tl.arange(0, BLOCK_GRP)
306 grp_mask = grp_offs < (D // 8)
307 idx_grp = tl.reshape(idx, [BLOCK_GRP, 8])
308 shifts_3 = tl.arange(0, 8) * 3
309 packed_24 = tl.sum((idx_grp & 0x7) << shifts_3[None, :], axis=1)
310 b0 = (packed_24 & 0xFF).to(tl.uint8)
311 b1 = ((packed_24 >> 8) & 0xFF).to(tl.uint8)
312 b2 = ((packed_24 >> 16) & 0xFF).to(tl.uint8)
313 tl.store(KV_cache_ptr + slot_base + grp_offs * 3, b0, mask=grp_mask)
314 tl.store(KV_cache_ptr + slot_base + grp_offs * 3 + 1, b1, mask=grp_mask)
315 tl.store(KV_cache_ptr + slot_base + grp_offs * 3 + 2, b2, mask=grp_mask)
316
317 # ── 3. STORE vec_norm (fp16, 2 bytes) ─────────────────────────────
318 norm_offset = MSE_BYTES
319
320 vn_f16 = tl.load(Norms_ptr + pid).to(tl.float16)
321 vn_u16 = vn_f16.to(tl.uint16, bitcast=True)
322 tl.store(KV_cache_ptr + slot_base + norm_offset, (vn_u16 & 0xFF).to(tl.uint8))
323 tl.store(
324 KV_cache_ptr + slot_base + norm_offset + 1, ((vn_u16 >> 8) & 0xFF).to(tl.uint8)
325 )
326
327 # ── 4. VALUE QUANTIZE + PACK ──────────────────────────────────────
328 _store_quantized_value(
329 Value_ptr,
330 KV_cache_ptr,
331 base,
332 slot_base,
333 d_offs,
334 d_mask,
335 D=D,
336 KPS=KPS,
337 VQB=VQB,
338 VAL_DATA_BYTES=VAL_DATA_BYTES,
339 BLOCK_D=BLOCK_D,
340 BLOCK_VAL=BLOCK_VAL,
341 BLOCK_GRP=BLOCK_GRP,
342 )
343
344
345 # ═══════════════════════════════════════════════════════════════════════
346 # Launcher
347 # ═══════════════════════════════════════════════════════════════════════
348
349
350 def triton_turboquant_store(
351 key: torch.Tensor, # [N, H, D] — raw keys (post-RoPE)
352 value: torch.Tensor, # [N, H, D] — raw values
353 kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8
354 slot_mapping: torch.Tensor, # [N] int32
355 PiT: torch.Tensor, # [D, D] float32
356 midpoints: torch.Tensor, # [n_centroids-1] float32
357 mse_bits: int,
358 key_packed_size: int,
359 value_quant_bits: int,
360 key_fp8: bool = False,
361 ):
362 """Launch TQ store kernel (FP8 or MSE path)."""
363 N, H, D = key.shape
364 NH = N * H
365 block_size = kv_cache.shape[1]
366 BLOCK_D = triton.next_power_of_2(D)
367 mse_bytes = math.ceil(D * mse_bits / 8)
368 n_centroids = 2**mse_bits
369
370 val_data_bytes = math.ceil(D * value_quant_bits / 8)
371
372 BLOCK_VAL = triton.next_power_of_2(val_data_bytes)
373
374 # Cache strides (element_size=1 for uint8, so stride in bytes = stride())
375 stride_block = kv_cache.stride(0)
376 stride_pos = kv_cache.stride(1)
377 stride_head = kv_cache.stride(2)
378
379 block_grp = triton.next_power_of_2(D // 8) if D >= 8 else 1
380
381 # ── FP8 PATH: in-kernel FP8 cast + scatter via fp8 kernel ──
382 if key_fp8:
383 k_flat = key.reshape(NH, D).contiguous()
384 v_flat = value.reshape(NH, D).contiguous()
385
386 fp8_e4b15 = _use_fp8_e4b15(key.device.index or 0)
387
388 grid = (NH,)
389 _tq_fused_store_fp8[grid](
390 k_flat,
391 v_flat,
392 kv_cache,
393 slot_mapping,
394 stride_cache_block=stride_block,
395 stride_cache_pos=stride_pos,
396 stride_cache_head=stride_head,
397 D=D,
398 H=H,
399 BLOCK_SIZE=block_size,
400 BLOCK_D=BLOCK_D,
401 KPS=key_packed_size,
402 VQB=value_quant_bits,
403 VAL_DATA_BYTES=val_data_bytes,
404 BLOCK_VAL=BLOCK_VAL,
405 BLOCK_GRP=block_grp,
406 FP8_E4B15=fp8_e4b15,
407 num_warps=4,
408 num_stages=1,
409 )
410 return
411
412 # ── MSE PATH: external GEMM + fused bucketize/pack kernel ──
413 # Normalize + rotation GEMM externally (cuBLAS is faster than in-kernel)
414 k_flat = key.float().reshape(NH, D)
415 norms = k_flat.norm(dim=1, keepdim=True)
416 x_hat = k_flat / (norms + 1e-8)
417 y = x_hat @ PiT
418
419 v_flat = value.float().reshape(NH, D)
420
421 # Fused kernel: bucketize + MSE index pack + norm store + value pack
422 grid = (NH,)
423 _tq_fused_store_mse[grid](
424 y,
425 norms.squeeze(1),
426 v_flat,
427 midpoints,
428 kv_cache,
429 slot_mapping,
430 stride_cache_block=stride_block,
431 stride_cache_pos=stride_pos,
432 stride_cache_head=stride_head,
433 D=D,
434 H=H,
435 BLOCK_SIZE=block_size,
436 BLOCK_D=BLOCK_D,
437 MSE_BYTES=mse_bytes,
438 KPS=key_packed_size,
439 VQB=value_quant_bits,
440 VAL_DATA_BYTES=val_data_bytes,
441 BLOCK_VAL=BLOCK_VAL,
442 MSE_BITS=mse_bits,
443 N_CENTROIDS=n_centroids,
444 BLOCK_GRP=block_grp,
445 num_warps=4,
446 num_stages=1,
447 )