#!/usr/bin/env python3 """TurboQuant store with the Lloyd-Max codec on the value path. Keys are handled exactly as upstream: rotate, unit-normalise, binary-search the midpoints, pack indices, store the fp16 norm. Values now take the same path instead of uniform asymmetric quantization, which is both more accurate and two bytes smaller per vector (one norm replaces the scale/zero pair). Slot layout written here: [ key indices | key norm (2B) | value indices | value norm (2B) ] mse_bytes 2 val_bytes 2 The decode side never rotates a cached value back: with values stored rotated and unit-normalised, P @ V_hat == ((P * norms) @ V_unit_rot) @ R^T, so the norm folds into the attention weight and the inverse rotation is one matmul on the output. """ from __future__ import annotations import math import torch import triton import triton.language as tl @triton.jit def _bucketize_pack( Y_ptr, Midpoints_ptr, KV_cache_ptr, base, dest, d_offs, d_mask, D: tl.constexpr, BITS: tl.constexpr, N_CENTROIDS: tl.constexpr, BLOCK_D: tl.constexpr, BLOCK_GRP: tl.constexpr, NBYTES: tl.constexpr, ): """Binary-search the midpoints and pack the indices at ``dest``.""" y_vec = tl.load(Y_ptr + base + d_offs, mask=d_mask, other=0.0) lo = tl.zeros([BLOCK_D], dtype=tl.int32) hi = tl.full([BLOCK_D], N_CENTROIDS - 1, dtype=tl.int32) for _ in range(BITS): mid = (lo + hi) >> 1 safe_mid = tl.minimum(mid, N_CENTROIDS - 2) mid_val = tl.load(Midpoints_ptr + safe_mid, mask=d_mask, other=0.0) lo = tl.where(y_vec >= mid_val, mid + 1, lo) hi = tl.where(y_vec >= mid_val, hi, mid) idx = tl.minimum(lo, N_CENTROIDS - 1) if BITS == 4: idx_pairs = tl.reshape(idx, [BLOCK_D // 2, 2]) shifts_4 = tl.arange(0, 2) * 4 packed = tl.sum((idx_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8) offs = tl.arange(0, BLOCK_D // 2) tl.store(KV_cache_ptr + dest + offs, packed, mask=offs < NBYTES) elif BITS == 3: grp_offs = tl.arange(0, BLOCK_GRP) grp_mask = grp_offs < (D // 8) idx_grp = tl.reshape(idx, [BLOCK_GRP, 8]) shifts_3 = tl.arange(0, 8) * 3 packed_24 = tl.sum((idx_grp & 0x7) << shifts_3[None, :], axis=1) tl.store( KV_cache_ptr + dest + grp_offs * 3, (packed_24 & 0xFF).to(tl.uint8), mask=grp_mask ) tl.store( KV_cache_ptr + dest + grp_offs * 3 + 1, ((packed_24 >> 8) & 0xFF).to(tl.uint8), mask=grp_mask, ) tl.store( KV_cache_ptr + dest + grp_offs * 3 + 2, ((packed_24 >> 16) & 0xFF).to(tl.uint8), mask=grp_mask, ) elif BITS == 2: idx_quads = tl.reshape(idx, [BLOCK_D // 4, 4]) shifts_2 = tl.arange(0, 4) * 2 packed = tl.sum((idx_quads & 0x3) << shifts_2[None, :], axis=1).to(tl.uint8) offs = tl.arange(0, BLOCK_D // 4) tl.store(KV_cache_ptr + dest + offs, packed, mask=offs < NBYTES) @triton.jit def _store_fp16_norm(Norms_ptr, KV_cache_ptr, pid, dest): value = tl.load(Norms_ptr + pid).to(tl.float16) bits = value.to(tl.uint16, bitcast=True) tl.store(KV_cache_ptr + dest, (bits & 0xFF).to(tl.uint8)) tl.store(KV_cache_ptr + dest + 1, ((bits >> 8) & 0xFF).to(tl.uint8)) @triton.jit def _zenit_store_kv_mse( YK_ptr, KNorms_ptr, YV_ptr, VNorms_ptr, KMid_ptr, VMid_ptr, KV_cache_ptr, Slot_mapping_ptr, stride_cache_block: tl.constexpr, stride_cache_pos: tl.constexpr, stride_cache_head: tl.constexpr, D: tl.constexpr, H: tl.constexpr, BLOCK_SIZE: tl.constexpr, BLOCK_D: tl.constexpr, KEY_BYTES: tl.constexpr, VAL_BYTES: tl.constexpr, KEY_PACKED: tl.constexpr, KEY_BITS: tl.constexpr, VAL_BITS: tl.constexpr, KEY_CENTROIDS: tl.constexpr, VAL_CENTROIDS: tl.constexpr, BLOCK_GRP: tl.constexpr = 32, ): pid = tl.program_id(0) token_idx = pid // H head_idx = pid % H slot = tl.load(Slot_mapping_ptr + token_idx) if slot < 0: return blk = (slot // BLOCK_SIZE).to(tl.int64) off = (slot % BLOCK_SIZE).to(tl.int64) slot_base = ( blk * stride_cache_block + off * stride_cache_pos + tl.cast(head_idx, tl.int64) * stride_cache_head ) base = pid * D d_offs = tl.arange(0, BLOCK_D) d_mask = d_offs < D _bucketize_pack( YK_ptr, KMid_ptr, KV_cache_ptr, base, slot_base, d_offs, d_mask, D=D, BITS=KEY_BITS, N_CENTROIDS=KEY_CENTROIDS, BLOCK_D=BLOCK_D, BLOCK_GRP=BLOCK_GRP, NBYTES=KEY_BYTES, ) _store_fp16_norm(KNorms_ptr, KV_cache_ptr, pid, slot_base + KEY_BYTES) _bucketize_pack( YV_ptr, VMid_ptr, KV_cache_ptr, base, slot_base + KEY_PACKED, d_offs, d_mask, D=D, BITS=VAL_BITS, N_CENTROIDS=VAL_CENTROIDS, BLOCK_D=BLOCK_D, BLOCK_GRP=BLOCK_GRP, NBYTES=VAL_BYTES, ) _store_fp16_norm(VNorms_ptr, KV_cache_ptr, pid, slot_base + KEY_PACKED + VAL_BYTES) def store_kv_value_mse( key: torch.Tensor, value: torch.Tensor, kv_cache: torch.Tensor, slot_mapping: torch.Tensor, rotation: torch.Tensor, key_midpoints: torch.Tensor, value_midpoints: torch.Tensor, key_bits: int, value_bits: int, ): """Host side: two rotations by cuBLAS, one fused kernel.""" N, H, D = key.shape NH = N * H block_size = kv_cache.shape[1] key_bytes = math.ceil(D * key_bits / 8) val_bytes = math.ceil(D * value_bits / 8) key_packed = key_bytes + 2 k_flat = key.float().reshape(NH, D) k_norms = k_flat.norm(dim=1, keepdim=True) yk = (k_flat / (k_norms + 1e-8)) @ rotation v_flat = value.float().reshape(NH, D) v_norms = v_flat.norm(dim=1, keepdim=True) yv = (v_flat / (v_norms + 1e-8)) @ rotation _zenit_store_kv_mse[(NH,)]( yk, k_norms.reshape(-1), yv, v_norms.reshape(-1), key_midpoints, value_midpoints, kv_cache.view(-1), slot_mapping, stride_cache_block=kv_cache.stride(0), stride_cache_pos=kv_cache.stride(1), stride_cache_head=kv_cache.stride(2), D=D, H=H, BLOCK_SIZE=block_size, BLOCK_D=triton.next_power_of_2(D), KEY_BYTES=key_bytes, VAL_BYTES=val_bytes, KEY_PACKED=key_packed, KEY_BITS=key_bits, VAL_BITS=value_bits, KEY_CENTROIDS=2**key_bits, VAL_CENTROIDS=2**value_bits, num_warps=4, num_stages=1, ) return { "key_bytes": key_bytes, "value_bytes": val_bytes, "key_packed_size": key_packed, "value_packed_size": val_bytes + 2, "slot_size": key_packed + val_bytes + 2, } __all__ = ["store_kv_value_mse"]