@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 524 lines 21.2 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
19 from .triton_turboquant_decode import _use_fp8_e4b15
20
21 # ═══════════════════════════════════════════════════════════════════════
22 # Shared: value uniform quantization + pack + SoA scale/zero store
23 # ═══════════════════════════════════════════════════════════════════════
24
25
26 @triton.jit
27 def _store_quantized_value(
28 Value_ptr,
29 KV_cache_ptr,
30 KV_cache_u16_ptr, # Opt#3: u16 view for SoA scale/zero stores
31 base, # pid * D offset into Value_ptr
32 data_base, # byte offset into KV_cache_ptr for this slot+head's DATA region
33 vscale_u16_addr, # u16 element index for V-scale in the SoA region
34 vzero_u16_addr, # u16 element index for V-zero in the SoA region
35 d_offs, # tl.arange(0, BLOCK_D)
36 d_mask, # d_offs < D
37 D: tl.constexpr,
38 # Offset (bytes) from data_base to V-data region within the slot's data.
39 # For MSE keys this equals MSE_BYTES; for FP8 keys it equals D.
40 V_DATA_OFFSET: tl.constexpr,
41 VQB: tl.constexpr,
42 VAL_DATA_BYTES: tl.constexpr,
43 BLOCK_D: tl.constexpr,
44 BLOCK_VAL: tl.constexpr,
45 BLOCK_GRP: tl.constexpr,
46 ):
47 """Uniform quantization of values to VQB bits, pack, and store with
48 scale/zero in the per-block SoA metadata region (Opt#3 layout)."""
49 if VQB == 3:
50 val_vec = tl.load(Value_ptr + base + d_offs, mask=d_mask, other=0.0).to(
51 tl.float32
52 )
53 val_min = tl.min(tl.where(d_mask, val_vec, float("inf")), axis=0)
54 val_max = tl.max(tl.where(d_mask, val_vec, -float("inf")), axis=0)
55 v_scale = (val_max - val_min) / 7.0
56 v_scale = tl.where(v_scale > 1e-8, v_scale, 1e-8)
57
58 q_vals = tl.minimum(
59 tl.maximum(((val_vec - val_min) / v_scale + 0.5).to(tl.int32), 0), 7
60 )
61
62 grp_offs = tl.arange(0, BLOCK_GRP)
63 grp_mask = grp_offs < (D // 8)
64 q_grp = tl.reshape(q_vals, [BLOCK_GRP, 8])
65 shifts_3bit = tl.arange(0, 8) * 3
66 packed_24 = tl.sum(q_grp << shifts_3bit[None, :], axis=1)
67 b0 = (packed_24 & 0xFF).to(tl.uint8)
68 b1 = ((packed_24 >> 8) & 0xFF).to(tl.uint8)
69 b2 = ((packed_24 >> 16) & 0xFF).to(tl.uint8)
70 tl.store(
71 KV_cache_ptr + data_base + V_DATA_OFFSET + grp_offs * 3,
72 b0,
73 mask=grp_mask,
74 )
75 tl.store(
76 KV_cache_ptr + data_base + V_DATA_OFFSET + grp_offs * 3 + 1,
77 b1,
78 mask=grp_mask,
79 )
80 tl.store(
81 KV_cache_ptr + data_base + V_DATA_OFFSET + grp_offs * 3 + 2,
82 b2,
83 mask=grp_mask,
84 )
85 else: # VQB == 4
86 val_vec = tl.load(Value_ptr + base + d_offs, mask=d_mask, other=0.0).to(
87 tl.float32
88 )
89 val_min = tl.min(tl.where(d_mask, val_vec, float("inf")), axis=0)
90 val_max = tl.max(tl.where(d_mask, val_vec, -float("inf")), axis=0)
91 v_scale = (val_max - val_min) / 15.0
92 v_scale = tl.where(v_scale > 1e-8, v_scale, 1e-8)
93
94 # Quantize all D elements from register (no re-load)
95 q_all = tl.minimum(
96 tl.maximum(((val_vec - val_min) / v_scale + 0.5).to(tl.int32), 0), 15
97 )
98 q_pairs = tl.reshape(q_all, [BLOCK_D // 2, 2])
99 shifts_4 = tl.arange(0, 2) * 4
100 packed_val = tl.sum((q_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8)
101 val_offs = tl.arange(0, BLOCK_D // 2)
102 val_mask = val_offs < VAL_DATA_BYTES
103 tl.store(
104 KV_cache_ptr + data_base + V_DATA_OFFSET + val_offs,
105 packed_val,
106 mask=val_mask,
107 )
108
109 # ── SoA scale/zero store (Opt#3) ──────────────────────────────────
110 # Scale and zero each get one fp16 stored at their token's slot index
111 # within the per-head SoA region. Single u16 store per field vs the
112 # previous 2×u8 writes per field.
113 sc_u16 = v_scale.to(tl.float16).to(tl.uint16, bitcast=True)
114 zr_u16 = val_min.to(tl.float16).to(tl.uint16, bitcast=True)
115 tl.store(KV_cache_u16_ptr + vscale_u16_addr, sc_u16)
116 tl.store(KV_cache_u16_ptr + vzero_u16_addr, zr_u16)
117
118
119 # ═══════════════════════════════════════════════════════════════════════
120 # FP8 key store + value uniform quantization
121 # ═══════════════════════════════════════════════════════════════════════
122
123
124 @triton.jit
125 def _tq_fused_store_fp8(
126 Key_ptr, # [NH, D] float16/bfloat16 — raw keys
127 Value_ptr, # [NH, D] float16/bfloat16 — raw values
128 KV_cache_ptr, # [total_bytes] uint8 (flattened view)
129 KV_cache_u16_ptr, # uint16-aliased view of the same storage (Opt#3)
130 Slot_mapping_ptr, # [N] int32 — per-token slot indices
131 # Cache strides (stride_cache_block = bytes per block = bs*H*slot_aligned)
132 stride_cache_block: tl.constexpr,
133 # Dimensions
134 D: tl.constexpr,
135 H: tl.constexpr,
136 BLOCK_SIZE: tl.constexpr,
137 BLOCK_D: tl.constexpr,
138 # SoA layout (Opt#3): data region is [bs, H, KEY_DATA_BYTES+VAL_DATA_BYTES],
139 # followed by metadata region [H, NUM_SOA_FIELDS, bs] of fp16.
140 KEY_DATA_BYTES: tl.constexpr, # FP8: = D
141 VAL_DATA_BYTES: tl.constexpr,
142 META_REGION_OFFSET: tl.constexpr, # bytes from block start to SoA region
143 NUM_SOA_FIELDS: tl.constexpr, # FP8: 2 (V-scale, V-zero)
144 SOA_V_SCALE: tl.constexpr, # FP8: 0
145 SOA_V_ZERO: tl.constexpr, # FP8: 1
146 # Value quantization
147 VQB: tl.constexpr,
148 # Packing block sizes
149 BLOCK_VAL: tl.constexpr,
150 BLOCK_GRP: tl.constexpr = 16,
151 FP8_E4B15: tl.constexpr = 0, # 1 = e4b15 (Ampere/Ada), 0 = e4nv (Hopper+)
152 ):
153 """FP8 key cast+scatter + value uniform quantization (SoA layout)."""
154 pid = tl.program_id(0)
155 token_idx = pid // H
156 head_idx = pid % H
157
158 slot = tl.load(Slot_mapping_ptr + token_idx)
159 if slot < 0:
160 return
161 blk = (slot // BLOCK_SIZE).to(tl.int64)
162 off = (slot % BLOCK_SIZE).to(tl.int64)
163 head_idx_i64 = tl.cast(head_idx, tl.int64)
164
165 # Block start byte address — stride_cache_block is unchanged (same total
166 # bytes per block as the AoS layout).
167 block_base = blk * stride_cache_block
168
169 # Data region: per-slot stride = H * DATA_BYTES_PER_SLOT; per-head stride
170 # within a slot = DATA_BYTES_PER_SLOT.
171 DATA_BYTES_PER_SLOT: tl.constexpr = KEY_DATA_BYTES + VAL_DATA_BYTES
172 data_base = (
173 block_base
174 + off * (H * DATA_BYTES_PER_SLOT)
175 + head_idx_i64 * DATA_BYTES_PER_SLOT
176 )
177
178 # SoA metadata region (u16 element indexing): head strip = NUM_SOA_FIELDS
179 # consecutive bs-long fp16 arrays.
180 head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + head_idx_i64 * (
181 NUM_SOA_FIELDS * BLOCK_SIZE
182 )
183 vscale_u16_addr = head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + off
184 vzero_u16_addr = head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + off
185
186 base = pid * D
187
188 # ── FP8 KEY: cast to FP8 in-kernel and store at data_base ────────
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)
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 + data_base + d_offs, k_bytes, mask=d_mask)
195
196 # ── VALUE QUANTIZE + PACK + SoA scale/zero store ─────────────────
197 _store_quantized_value(
198 Value_ptr,
199 KV_cache_ptr,
200 KV_cache_u16_ptr,
201 base,
202 data_base,
203 vscale_u16_addr,
204 vzero_u16_addr,
205 d_offs,
206 d_mask,
207 D=D,
208 V_DATA_OFFSET=KEY_DATA_BYTES, # V follows K in the data region
209 VQB=VQB,
210 VAL_DATA_BYTES=VAL_DATA_BYTES,
211 BLOCK_D=BLOCK_D,
212 BLOCK_VAL=BLOCK_VAL,
213 BLOCK_GRP=BLOCK_GRP,
214 )
215
216
217 # ═══════════════════════════════════════════════════════════════════════
218 # Fused MSE store: bucketize + MSE index pack + norm store + value pack
219 # (eliminates 4 PyTorch kernel launches per layer vs pack-only kernel)
220 # ═══════════════════════════════════════════════════════════════════════
221
222
223 @triton.jit
224 def _tq_fused_store_mse(
225 # Post-rotation inputs
226 Y_ptr, # [NH, D] float32 — rotated normalized keys (x_hat @ PiT)
227 Norms_ptr, # [NH] float32 — key vector norms (||k||)
228 Value_ptr, # [NH, D] float32 — raw values
229 # Quantization tables
230 Midpoints_ptr, # [n_centroids-1] float32
231 Centroids_ptr, # [n_centroids] float32 — used only when NORM_CORRECTION=1
232 # Cache and indexing
233 KV_cache_ptr, # [total_bytes] uint8 (flattened view)
234 KV_cache_u16_ptr, # uint16-aliased view of the same storage (Opt#3)
235 Slot_mapping_ptr, # [N] int32 — per-token slot indices
236 # Cache strides (stride_cache_block = bs*H*slot_aligned, unchanged)
237 stride_cache_block: tl.constexpr,
238 # Dimensions
239 D: tl.constexpr,
240 H: tl.constexpr,
241 BLOCK_SIZE: tl.constexpr,
242 BLOCK_D: tl.constexpr,
243 # SoA layout (Opt#3): data region [bs, H, MSE_BYTES+VAL_DATA_BYTES],
244 # then metadata region [H, NUM_SOA_FIELDS, bs] of fp16.
245 MSE_BYTES: tl.constexpr,
246 VAL_DATA_BYTES: tl.constexpr,
247 META_REGION_OFFSET: tl.constexpr, # bytes
248 NUM_SOA_FIELDS: tl.constexpr, # MSE: 3 (k_norm, v_scale, v_zero)
249 SOA_K_NORM: tl.constexpr, # MSE: 0
250 SOA_V_SCALE: tl.constexpr, # MSE: 1
251 SOA_V_ZERO: tl.constexpr, # MSE: 2
252 # Value quantization
253 VQB: tl.constexpr,
254 # Packing block sizes
255 BLOCK_VAL: tl.constexpr,
256 # MSE params
257 MSE_BITS: tl.constexpr,
258 N_CENTROIDS: tl.constexpr,
259 # When 1, fold the per-token centroid-vector norm ||c_t|| into the
260 # stored scalar: store ||k_t|| / ||c_t|| instead of ||k_t||. This
261 # removes the per-tile sum+sqrt+divide cost from the attention kernel's
262 # K-load path. Load-side math becomes a plain multiply K = c * stored.
263 NORM_CORRECTION: tl.constexpr = 0,
264 BLOCK_GRP: tl.constexpr = 16,
265 ):
266 """Fused MSE quantize + pack + SoA-metadata store (Opt#3 layout).
267
268 Binary-search bucketize → MSE-index pack → K-norm store (SoA) → value
269 uniform quantize + pack → V scale/zero store (SoA), in one kernel.
270 """
271 pid = tl.program_id(0)
272 token_idx = pid // H
273 head_idx = pid % H
274
275 slot = tl.load(Slot_mapping_ptr + token_idx)
276 if slot < 0:
277 return
278 blk = (slot // BLOCK_SIZE).to(tl.int64)
279 off = (slot % BLOCK_SIZE).to(tl.int64)
280 head_idx_i64 = tl.cast(head_idx, tl.int64)
281
282 # Block start byte address — stride unchanged from AoS (same total
283 # bytes per block); layout inside the block differs.
284 block_base = blk * stride_cache_block
285
286 # Data region offset (packed K indices + packed V data).
287 DATA_BYTES_PER_SLOT: tl.constexpr = MSE_BYTES + VAL_DATA_BYTES
288 data_base = (
289 block_base
290 + off * (H * DATA_BYTES_PER_SLOT)
291 + head_idx_i64 * DATA_BYTES_PER_SLOT
292 )
293
294 # SoA metadata addresses (u16 element indexing).
295 head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + head_idx_i64 * (
296 NUM_SOA_FIELDS * BLOCK_SIZE
297 )
298 knorm_u16_addr = head_meta_u16_base + SOA_K_NORM * BLOCK_SIZE + off
299 vscale_u16_addr = head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + off
300 vzero_u16_addr = head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + off
301
302 base = pid * D
303 d_offs = tl.arange(0, BLOCK_D)
304 d_mask = d_offs < D
305
306 # ── 1. BINARY SEARCH BUCKETIZE ────────────────────────────────────
307 y_vec = tl.load(Y_ptr + base + d_offs, mask=d_mask, other=0.0)
308 lo = tl.zeros([BLOCK_D], dtype=tl.int32)
309 hi = tl.full([BLOCK_D], N_CENTROIDS - 1, dtype=tl.int32)
310 for _ in range(MSE_BITS):
311 mid = (lo + hi) >> 1
312 safe_mid = tl.minimum(mid, N_CENTROIDS - 2)
313 mid_val = tl.load(Midpoints_ptr + safe_mid, mask=d_mask, other=0.0)
314 lo = tl.where(y_vec >= mid_val, mid + 1, lo)
315 hi = tl.where(y_vec >= mid_val, hi, mid)
316 idx = tl.minimum(lo, N_CENTROIDS - 1)
317
318 # ── 2. PACK MSE INDICES into data region at data_base[0:MSE_BYTES] ──
319 if MSE_BITS == 4:
320 idx_pairs = tl.reshape(idx, [BLOCK_D // 2, 2])
321 shifts_4 = tl.arange(0, 2) * 4
322 packed = tl.sum((idx_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8)
323 mse_offs = tl.arange(0, BLOCK_D // 2)
324 mse_mask = mse_offs < MSE_BYTES
325 tl.store(KV_cache_ptr + data_base + mse_offs, packed, mask=mse_mask)
326
327 elif MSE_BITS == 3:
328 grp_offs = tl.arange(0, BLOCK_GRP)
329 grp_mask = grp_offs < (D // 8)
330 idx_grp = tl.reshape(idx, [BLOCK_GRP, 8])
331 shifts_3 = tl.arange(0, 8) * 3
332 packed_24 = tl.sum((idx_grp & 0x7) << shifts_3[None, :], axis=1)
333 b0 = (packed_24 & 0xFF).to(tl.uint8)
334 b1 = ((packed_24 >> 8) & 0xFF).to(tl.uint8)
335 b2 = ((packed_24 >> 16) & 0xFF).to(tl.uint8)
336 tl.store(KV_cache_ptr + data_base + grp_offs * 3, b0, mask=grp_mask)
337 tl.store(KV_cache_ptr + data_base + grp_offs * 3 + 1, b1, mask=grp_mask)
338 tl.store(KV_cache_ptr + data_base + grp_offs * 3 + 2, b2, mask=grp_mask)
339
340 # ── 3. STORE K-norm in SoA region (single u16 store) ─────────────
341 # Norm-correction (Opt#1) still folds 1/||c_vec|| into the stored value;
342 # only the destination moved to the SoA K-norm array.
343 vn_f32 = tl.load(Norms_ptr + pid) # ||k_t||, fp32
344
345 if NORM_CORRECTION:
346 safe_idx = tl.minimum(tl.maximum(idx, 0), N_CENTROIDS - 1)
347 c_vals = tl.load(Centroids_ptr + safe_idx, mask=d_mask, other=0.0)
348 c_norm_sq = tl.sum(tl.where(d_mask, c_vals * c_vals, 0.0), axis=0)
349 c_inv_norm = 1.0 / tl.sqrt(c_norm_sq + 1e-16)
350 vn_f32 = vn_f32 * c_inv_norm
351
352 vn_u16 = vn_f32.to(tl.float16).to(tl.uint16, bitcast=True)
353 tl.store(KV_cache_u16_ptr + knorm_u16_addr, vn_u16)
354
355 # ── 4. VALUE QUANTIZE + PACK + SoA scale/zero store ──────────────
356 _store_quantized_value(
357 Value_ptr,
358 KV_cache_ptr,
359 KV_cache_u16_ptr,
360 base,
361 data_base,
362 vscale_u16_addr,
363 vzero_u16_addr,
364 d_offs,
365 d_mask,
366 D=D,
367 V_DATA_OFFSET=MSE_BYTES, # V follows K indices in the data region
368 VQB=VQB,
369 VAL_DATA_BYTES=VAL_DATA_BYTES,
370 BLOCK_D=BLOCK_D,
371 BLOCK_VAL=BLOCK_VAL,
372 BLOCK_GRP=BLOCK_GRP,
373 )
374
375
376 # ═══════════════════════════════════════════════════════════════════════
377 # Launcher
378 # ═══════════════════════════════════════════════════════════════════════
379
380
381 def triton_turboquant_store(
382 key: torch.Tensor, # [N, H, D] — raw keys (post-RoPE)
383 value: torch.Tensor, # [N, H, D] — raw values
384 kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8
385 slot_mapping: torch.Tensor, # [N] int32
386 PiT: torch.Tensor, # [D, D] float32
387 midpoints: torch.Tensor, # [n_centroids-1] float32
388 mse_bits: int,
389 key_packed_size: int,
390 value_quant_bits: int,
391 key_fp8: bool = False,
392 centroids: torch.Tensor
393 | None = None, # [n_centroids] float32 — required when norm_correction=True
394 norm_correction: bool = False,
395 ):
396 """Launch TQ store kernel (FP8 or MSE path).
397
398 When ``norm_correction=True`` (default off; set by MSE norm-correction
399 presets), the centroid-vector norm ||c_t|| is folded into the stored
400 per-token scalar at store time (stored value = ||k_t|| / ||c_t||). The
401 load-side attention kernel no longer computes per-tile sum+sqrt+divide,
402 saving ~26% of TQ overhead on gpt-oss decode (ctx=4096). ``centroids``
403 must be provided in this mode.
404 """
405 if norm_correction and not key_fp8:
406 assert centroids is not None, (
407 "norm_correction=True requires centroids tensor; got None. "
408 "Pass the same [n_centroids] fp32 tensor used to derive midpoints."
409 )
410 N, H, D = key.shape
411 NH = N * H
412 block_size = kv_cache.shape[1]
413 BLOCK_D = triton.next_power_of_2(D)
414 mse_bytes = math.ceil(D * mse_bits / 8)
415 n_centroids = 2**mse_bits
416
417 val_data_bytes = math.ceil(D * value_quant_bits / 8)
418
419 BLOCK_VAL = triton.next_power_of_2(val_data_bytes)
420
421 # Cache strides (element_size=1 for uint8, so stride in bytes = stride()).
422 # Only stride_cache_block (== bs*H*slot_aligned) is used under Opt#3;
423 # per-slot and per-head strides are computed from DATA_BYTES_PER_SLOT.
424 stride_block = kv_cache.stride(0)
425
426 block_grp = triton.next_power_of_2(D // 8) if D >= 8 else 1
427
428 # ── Opt#3 SoA layout constants (derived locally so the launcher API
429 # stays unchanged). Invariant: key_packed_size + value_packed_size
430 # == slot_size_aligned == (key_data + val_data) + (k_meta + v_meta).
431 key_data_bytes = D if key_fp8 else mse_bytes
432 data_bytes_per_slot = key_data_bytes + val_data_bytes
433 # Data region occupies [0, block_size * H * data_bytes_per_slot) bytes
434 # at the start of each block; metadata region starts at that offset.
435 meta_region_offset = block_size * H * data_bytes_per_slot
436 num_soa_fields = 2 if key_fp8 else 3
437 soa_v_scale = 0 if key_fp8 else 1
438 soa_v_zero = 1 if key_fp8 else 2
439
440 # uint16-aliased view of the same storage — enables single-instruction
441 # u16 writes for the SoA K-norm / V-scale / V-zero stores. The cache
442 # bytes are always contiguous and 2-byte aligned under the Opt#3 layout
443 # (MSE_BYTES, VAL_DATA_BYTES, meta_region_offset are all even).
444 kv_cache_u16 = kv_cache.view(torch.uint16)
445
446 # ── FP8 PATH: in-kernel FP8 cast + scatter via fp8 kernel ──
447 if key_fp8:
448 k_flat = key.reshape(NH, D).contiguous()
449 v_flat = value.reshape(NH, D).contiguous()
450
451 fp8_e4b15 = _use_fp8_e4b15(key.device.index or 0)
452
453 grid = (NH,)
454 _tq_fused_store_fp8[grid](
455 k_flat,
456 v_flat,
457 kv_cache.view(-1),
458 kv_cache_u16.view(-1),
459 slot_mapping,
460 stride_cache_block=stride_block,
461 D=D,
462 H=H,
463 BLOCK_SIZE=block_size,
464 BLOCK_D=BLOCK_D,
465 KEY_DATA_BYTES=key_data_bytes,
466 VAL_DATA_BYTES=val_data_bytes,
467 META_REGION_OFFSET=meta_region_offset,
468 NUM_SOA_FIELDS=num_soa_fields,
469 SOA_V_SCALE=soa_v_scale,
470 SOA_V_ZERO=soa_v_zero,
471 VQB=value_quant_bits,
472 BLOCK_VAL=BLOCK_VAL,
473 BLOCK_GRP=block_grp,
474 FP8_E4B15=fp8_e4b15,
475 num_warps=4,
476 num_stages=1,
477 )
478 return
479
480 # ── MSE PATH: external GEMM + fused bucketize/pack kernel ──
481 # Normalize + rotation GEMM externally (cuBLAS is faster than in-kernel)
482 k_flat = key.float().reshape(NH, D)
483 norms = k_flat.norm(dim=1, keepdim=True)
484 x_hat = k_flat / (norms + 1e-8)
485 y = x_hat @ PiT
486
487 v_flat = value.float().reshape(NH, D)
488
489 # When norm_correction=True, the kernel folds 1/||c_t|| into the
490 # stored per-token K-norm so the load kernel can skip per-tile norm
491 # math (Opt#1). centroids is only dereferenced inside the kernel when
492 # NORM_CORRECTION=1; pass midpoints as a harmless placeholder otherwise.
493 centroids_ptr = centroids if centroids is not None else midpoints
494 grid = (NH,)
495 _tq_fused_store_mse[grid](
496 y,
497 norms.squeeze(1),
498 v_flat,
499 midpoints,
500 centroids_ptr,
501 kv_cache.view(-1),
502 kv_cache_u16.view(-1),
503 slot_mapping,
504 stride_cache_block=stride_block,
505 D=D,
506 H=H,
507 BLOCK_SIZE=block_size,
508 BLOCK_D=BLOCK_D,
509 MSE_BYTES=mse_bytes,
510 VAL_DATA_BYTES=val_data_bytes,
511 META_REGION_OFFSET=meta_region_offset,
512 NUM_SOA_FIELDS=num_soa_fields,
513 SOA_K_NORM=0,
514 SOA_V_SCALE=soa_v_scale,
515 SOA_V_ZERO=soa_v_zero,
516 VQB=value_quant_bits,
517 BLOCK_VAL=BLOCK_VAL,
518 MSE_BITS=mse_bits,
519 N_CENTROIDS=n_centroids,
520 NORM_CORRECTION=1 if norm_correction else 0,
521 BLOCK_GRP=block_grp,
522 num_warps=4,
523 num_stages=1,
524 )