turboquant-gpu 0.1.0__py3-none-any.whl

turboquant_gpu/__init__.py

@@ -0,0 +1,9 @@
+__version__ = "0.1.0"
+
+def __getattr__(name):
+    if name == "TurboQuantEngine":
+        from .host import TurboQuantEngine
+        return TurboQuantEngine
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+__all__ = ["__version__", "TurboQuantEngine"]

turboquant_gpu/attention.py

@@ -0,0 +1,386 @@
+"""
+asymmetric attention kernels for compressed keys.
+
+score(q, k) ~ <q, k_mse> + ||r|| * sqrt(pi/2)/m * <S*q, signs>
+
+score-only variant outputs raw scores.
+fused variant does online softmax + V accumulation in one pass.
+USE_SWIZZLE enables exp2 ftz, approx div, prefetched loads, and block
+interleaving on architectures that support them.
+"""
+
+import math
+
+try:
+    import cuda.tile as ct
+    from cuda.tile import RoundingMode as RMd
+    _HAS_APPROX = hasattr(RMd, "APPROX")
+except ImportError:
+    import cutile as ct  # type: ignore
+    RMd = None
+    _HAS_APPROX = False
+
+from .constants import BLOCK_Q, BLOCK_KV, HEAD_DIM
+
+INV_LOG_2 = 1.0 / math.log(2)
+ConstBool = ct.Constant[bool]
+
+
+# ── score-only kernel ────────────────────────────────────────────────
+
+@ct.kernel(occupancy=2)
+def turboquant_attention_scores(
+    Q, K_mse, Signs, R_norms, Q_proj, Output,
+    scale: float,
+    correction_scale: float,
+    seq_k: int,
+    USE_SWIZZLE: ConstBool,
+):
+    """one program per query block, streams over all KV blocks."""
+    q_block = ct.bid(0)
+    zero_pad = ct.PaddingMode.ZERO
+
+    q_tile  = ct.load(Q,      index=(q_block, 0), shape=(BLOCK_Q, HEAD_DIM),
+                      padding_mode=zero_pad)
+    qp_tile = ct.load(Q_proj, index=(q_block, 0), shape=(BLOCK_Q, HEAD_DIM),
+                      padding_mode=zero_pad)
+
+    num_kv_blocks = ct.num_tiles(K_mse, axis=0, shape=(BLOCK_KV, HEAD_DIM))
+
+    for kv_block in range(num_kv_blocks):
+        if USE_SWIZZLE:
+            k_tile = ct.load(K_mse, index=(kv_block, 0),
+                             shape=(BLOCK_KV, HEAD_DIM), padding_mode=zero_pad,
+                             latency=2)
+        else:
+            k_tile = ct.load(K_mse, index=(kv_block, 0),
+                             shape=(BLOCK_KV, HEAD_DIM), padding_mode=zero_pad)
+
+        # mse term: Q @ K_mse^T
+        term1 = ct.mma(q_tile, ct.transpose(k_tile),
+                       ct.zeros((BLOCK_Q, BLOCK_KV), dtype=ct.float32))
+
+        # qjl bias-correction term
+        s_tile = ct.load(Signs, index=(kv_block, 0),
+                         shape=(BLOCK_KV, HEAD_DIM), padding_mode=zero_pad)
+        s_float = ct.astype(s_tile, ct.float16)
+
+        qjl_ip = ct.mma(qp_tile, ct.transpose(s_float),
+                        ct.zeros((BLOCK_Q, BLOCK_KV), dtype=ct.float32))
+
+        rn     = ct.load(R_norms, index=(kv_block,), shape=(BLOCK_KV,),
+                         padding_mode=zero_pad)
+        rn_f32 = ct.astype(rn, ct.float32)
+
+        term2 = correction_scale * qjl_ip * ct.expand_dims(rn_f32, axis=0)
+
+        scores = (term1 + term2) * scale
+        ct.store(Output, index=(q_block, kv_block), tile=scores)
+
+
+# ── fused attention (pre-decompressed V) ─────────────────────────────
+
+@ct.kernel(occupancy=2)
+def turboquant_fused_attention(
+    Q, K_mse, Signs, R_norms, Q_proj, V, Output,
+    scale: float,
+    correction_scale: float,
+    seq_k: int,
+    USE_SWIZZLE: ConstBool,
+):
+    """fused attention + online softmax.  V is already fp16."""
+    zero_pad = ct.PaddingMode.ZERO
+
+    if USE_SWIZZLE:
+        half = ct.cdiv(ct.num_blocks(0), 2)
+        q_block = ct.bid(0) // 2 + (ct.bid(0) % 2) * half
+    else:
+        q_block = ct.bid(0)
+
+    q_tile  = ct.load(Q,      index=(q_block, 0), shape=(BLOCK_Q, HEAD_DIM),
+                      padding_mode=zero_pad)
+    qp_tile = ct.load(Q_proj, index=(q_block, 0), shape=(BLOCK_Q, HEAD_DIM),
+                      padding_mode=zero_pad)
+
+    # online softmax accumulators
+    m_i = ct.full((BLOCK_Q,), -1e30, dtype=ct.float32)
+    l_i = ct.zeros((BLOCK_Q,), dtype=ct.float32)
+    acc = ct.zeros((BLOCK_Q, HEAD_DIM), dtype=ct.float32)
+
+    scale_log2 = scale * INV_LOG_2
+    num_kv_blocks = ct.num_tiles(K_mse, axis=0, shape=(BLOCK_KV, HEAD_DIM))
+
+    for kv_block in range(num_kv_blocks):
+        if USE_SWIZZLE:
+            k_tile = ct.load(K_mse, index=(kv_block, 0),
+                             shape=(BLOCK_KV, HEAD_DIM), padding_mode=zero_pad,
+                             latency=2)
+            v_tile = ct.load(V, index=(kv_block, 0),
+                             shape=(BLOCK_KV, HEAD_DIM), padding_mode=zero_pad,
+                             latency=4)
+        else:
+            k_tile = ct.load(K_mse, index=(kv_block, 0),
+                             shape=(BLOCK_KV, HEAD_DIM), padding_mode=zero_pad)
+            v_tile = ct.load(V, index=(kv_block, 0),
+                             shape=(BLOCK_KV, HEAD_DIM), padding_mode=zero_pad)
+
+        s_tile = ct.load(Signs, index=(kv_block, 0),
+                         shape=(BLOCK_KV, HEAD_DIM), padding_mode=zero_pad)
+        rn     = ct.load(R_norms, index=(kv_block,), shape=(BLOCK_KV,),
+                         padding_mode=zero_pad)
+
+        term1 = ct.mma(q_tile, ct.transpose(k_tile),
+                       ct.zeros((BLOCK_Q, BLOCK_KV), dtype=ct.float32))
+
+        s_float = ct.astype(s_tile, ct.float16)
+        qjl_ip  = ct.mma(qp_tile, ct.transpose(s_float),
+                         ct.zeros((BLOCK_Q, BLOCK_KV), dtype=ct.float32))
+        rn_f32  = ct.astype(rn, ct.float32)
+        term2   = correction_scale * qjl_ip * ct.expand_dims(rn_f32, axis=0)
+
+        raw_scores = term1 + term2
+
+        # online softmax update
+        if USE_SWIZZLE:
+            m_new = ct.maximum(m_i, ct.max(raw_scores, axis=1) * scale_log2)
+            alpha = ct.exp2(m_i - m_new, flush_to_zero=True)
+            p     = ct.exp2(raw_scores * scale_log2 - ct.expand_dims(m_new, axis=1),
+                            flush_to_zero=True)
+        else:
+            scores = raw_scores * scale
+            m_new  = ct.maximum(m_i, ct.max(scores, axis=1))
+            alpha  = ct.exp(m_i - m_new)
+            p      = ct.exp(scores - ct.expand_dims(m_new, axis=1))
+
+        l_i = alpha * l_i + ct.sum(p, axis=1)
+
+        p_fp16 = ct.astype(p, ct.float16)
+        acc = ct.expand_dims(alpha, axis=1) * acc + ct.mma(
+            p_fp16, v_tile, ct.zeros((BLOCK_Q, HEAD_DIM), dtype=ct.float32))
+
+        m_i = m_new
+
+    if USE_SWIZZLE and _HAS_APPROX:
+        out = ct.truediv(acc, ct.expand_dims(l_i, axis=1),
+                         flush_to_zero=True, rounding_mode=RMd.APPROX)
+    else:
+        out = acc / ct.expand_dims(l_i, axis=1)
+
+    ct.store(Output, index=(q_block, 0), tile=out)
+
+
+# ── fused attention (on-chip 3-bit V decompression) ──────────────────
+
+@ct.kernel(occupancy=2)
+def turboquant_fused_attention_vfused_3bit(
+    Q, K_mse, Signs, R_norms, Q_proj,
+    V_Indices, V_Norms, Pi, Output,
+    scale: float,
+    correction_scale: float,
+    seq_k: int,
+    vc0: float, vc1: float, vc2: float, vc3: float,
+    vc4: float, vc5: float, vc6: float, vc7: float,
+    USE_SWIZZLE: ConstBool,
+):
+    """fused attention that decompresses 3-bit V on-chip per block."""
+    zero_pad = ct.PaddingMode.ZERO
+
+    if USE_SWIZZLE:
+        half = ct.cdiv(ct.num_blocks(0), 2)
+        q_block = ct.bid(0) // 2 + (ct.bid(0) % 2) * half
+    else:
+        q_block = ct.bid(0)
+
+    q_tile  = ct.load(Q,      index=(q_block, 0), shape=(BLOCK_Q, HEAD_DIM),
+                      padding_mode=zero_pad)
+    qp_tile = ct.load(Q_proj, index=(q_block, 0), shape=(BLOCK_Q, HEAD_DIM),
+                      padding_mode=zero_pad)
+
+    pi_tile = ct.load(Pi, index=(0, 0), shape=(HEAD_DIM, HEAD_DIM))
+
+    m_i = ct.full((BLOCK_Q,), -1e30, dtype=ct.float32)
+    l_i = ct.zeros((BLOCK_Q,), dtype=ct.float32)
+    acc = ct.zeros((BLOCK_Q, HEAD_DIM), dtype=ct.float32)
+
+    scale_log2 = scale * INV_LOG_2
+    num_kv_blocks = ct.num_tiles(K_mse, axis=0, shape=(BLOCK_KV, HEAD_DIM))
+
+    for kv_block in range(num_kv_blocks):
+        if USE_SWIZZLE:
+            k_tile = ct.load(K_mse, index=(kv_block, 0),
+                             shape=(BLOCK_KV, HEAD_DIM), padding_mode=zero_pad,
+                             latency=2)
+        else:
+            k_tile = ct.load(K_mse, index=(kv_block, 0),
+                             shape=(BLOCK_KV, HEAD_DIM), padding_mode=zero_pad)
+
+        # decompress V tile on-chip
+        v_idx = ct.load(V_Indices, index=(kv_block, 0),
+                        shape=(BLOCK_KV, HEAD_DIM), padding_mode=zero_pad)
+        v_nrm = ct.load(V_Norms, index=(kv_block,),
+                        shape=(BLOCK_KV,), padding_mode=zero_pad)
+
+        vi_f32 = ct.astype(v_idx, ct.float32)
+        y_hat = ct.full((BLOCK_KV, HEAD_DIM), vc0, dtype=ct.float32)
+        y_hat = ct.where(vi_f32 > 0.5, vc1, y_hat)
+        y_hat = ct.where(vi_f32 > 1.5, vc2, y_hat)
+        y_hat = ct.where(vi_f32 > 2.5, vc3, y_hat)
+        y_hat = ct.where(vi_f32 > 3.5, vc4, y_hat)
+        y_hat = ct.where(vi_f32 > 4.5, vc5, y_hat)
+        y_hat = ct.where(vi_f32 > 5.5, vc6, y_hat)
+        y_hat = ct.where(vi_f32 > 6.5, vc7, y_hat)
+
+        v_recon = ct.mma(ct.astype(y_hat, ct.float16), pi_tile,
+                         ct.zeros((BLOCK_KV, HEAD_DIM), dtype=ct.float32))
+        v_nrm_f32 = ct.astype(v_nrm, ct.float32)
+        v_tile = ct.astype(
+            v_recon * ct.expand_dims(v_nrm_f32, axis=1), ct.float16)
+
+        # attention scores
+        s_tile = ct.load(Signs, index=(kv_block, 0),
+                         shape=(BLOCK_KV, HEAD_DIM), padding_mode=zero_pad)
+        rn     = ct.load(R_norms, index=(kv_block,), shape=(BLOCK_KV,),
+                         padding_mode=zero_pad)
+
+        term1 = ct.mma(q_tile, ct.transpose(k_tile),
+                       ct.zeros((BLOCK_Q, BLOCK_KV), dtype=ct.float32))
+
+        s_float = ct.astype(s_tile, ct.float16)
+        qjl_ip  = ct.mma(qp_tile, ct.transpose(s_float),
+                         ct.zeros((BLOCK_Q, BLOCK_KV), dtype=ct.float32))
+        rn_f32  = ct.astype(rn, ct.float32)
+        term2   = correction_scale * qjl_ip * ct.expand_dims(rn_f32, axis=0)
+
+        raw_scores = term1 + term2
+
+        if USE_SWIZZLE:
+            m_new = ct.maximum(m_i, ct.max(raw_scores, axis=1) * scale_log2)
+            alpha = ct.exp2(m_i - m_new, flush_to_zero=True)
+            p     = ct.exp2(raw_scores * scale_log2 - ct.expand_dims(m_new, axis=1),
+                            flush_to_zero=True)
+        else:
+            scores = raw_scores * scale
+            m_new  = ct.maximum(m_i, ct.max(scores, axis=1))
+            alpha  = ct.exp(m_i - m_new)
+            p      = ct.exp(scores - ct.expand_dims(m_new, axis=1))
+
+        l_i = alpha * l_i + ct.sum(p, axis=1)
+
+        p_fp16 = ct.astype(p, ct.float16)
+        acc = ct.expand_dims(alpha, axis=1) * acc + ct.mma(
+            p_fp16, v_tile, ct.zeros((BLOCK_Q, HEAD_DIM), dtype=ct.float32))
+
+        m_i = m_new
+
+    if USE_SWIZZLE and _HAS_APPROX:
+        out = ct.truediv(acc, ct.expand_dims(l_i, axis=1),
+                         flush_to_zero=True, rounding_mode=RMd.APPROX)
+    else:
+        out = acc / ct.expand_dims(l_i, axis=1)
+
+    ct.store(Output, index=(q_block, 0), tile=out)
+
+
+# ── fused attention (on-chip 2-bit V decompression) ──────────────────
+
+@ct.kernel(occupancy=2)
+def turboquant_fused_attention_vfused_2bit(
+    Q, K_mse, Signs, R_norms, Q_proj,
+    V_Indices, V_Norms, Pi, Output,
+    scale: float,
+    correction_scale: float,
+    seq_k: int,
+    vc0: float, vc1: float, vc2: float, vc3: float,
+    USE_SWIZZLE: ConstBool,
+):
+    """fused attention that decompresses 2-bit V on-chip per block."""
+    zero_pad = ct.PaddingMode.ZERO
+
+    if USE_SWIZZLE:
+        half = ct.cdiv(ct.num_blocks(0), 2)
+        q_block = ct.bid(0) // 2 + (ct.bid(0) % 2) * half
+    else:
+        q_block = ct.bid(0)
+
+    q_tile  = ct.load(Q,      index=(q_block, 0), shape=(BLOCK_Q, HEAD_DIM),
+                      padding_mode=zero_pad)
+    qp_tile = ct.load(Q_proj, index=(q_block, 0), shape=(BLOCK_Q, HEAD_DIM),
+                      padding_mode=zero_pad)
+
+    pi_tile = ct.load(Pi, index=(0, 0), shape=(HEAD_DIM, HEAD_DIM))
+
+    m_i = ct.full((BLOCK_Q,), -1e30, dtype=ct.float32)
+    l_i = ct.zeros((BLOCK_Q,), dtype=ct.float32)
+    acc = ct.zeros((BLOCK_Q, HEAD_DIM), dtype=ct.float32)
+
+    scale_log2 = scale * INV_LOG_2
+    num_kv_blocks = ct.num_tiles(K_mse, axis=0, shape=(BLOCK_KV, HEAD_DIM))
+
+    for kv_block in range(num_kv_blocks):
+        if USE_SWIZZLE:
+            k_tile = ct.load(K_mse, index=(kv_block, 0),
+                             shape=(BLOCK_KV, HEAD_DIM), padding_mode=zero_pad,
+                             latency=2)
+        else:
+            k_tile = ct.load(K_mse, index=(kv_block, 0),
+                             shape=(BLOCK_KV, HEAD_DIM), padding_mode=zero_pad)
+
+        v_idx = ct.load(V_Indices, index=(kv_block, 0),
+                        shape=(BLOCK_KV, HEAD_DIM), padding_mode=zero_pad)
+        v_nrm = ct.load(V_Norms, index=(kv_block,),
+                        shape=(BLOCK_KV,), padding_mode=zero_pad)
+
+        vi_f32 = ct.astype(v_idx, ct.float32)
+        y_hat = ct.full((BLOCK_KV, HEAD_DIM), vc0, dtype=ct.float32)
+        y_hat = ct.where(vi_f32 > 0.5, vc1, y_hat)
+        y_hat = ct.where(vi_f32 > 1.5, vc2, y_hat)
+        y_hat = ct.where(vi_f32 > 2.5, vc3, y_hat)
+
+        v_recon = ct.mma(ct.astype(y_hat, ct.float16), pi_tile,
+                         ct.zeros((BLOCK_KV, HEAD_DIM), dtype=ct.float32))
+        v_nrm_f32 = ct.astype(v_nrm, ct.float32)
+        v_tile = ct.astype(
+            v_recon * ct.expand_dims(v_nrm_f32, axis=1), ct.float16)
+
+        s_tile = ct.load(Signs, index=(kv_block, 0),
+                         shape=(BLOCK_KV, HEAD_DIM), padding_mode=zero_pad)
+        rn     = ct.load(R_norms, index=(kv_block,), shape=(BLOCK_KV,),
+                         padding_mode=zero_pad)
+
+        term1 = ct.mma(q_tile, ct.transpose(k_tile),
+                       ct.zeros((BLOCK_Q, BLOCK_KV), dtype=ct.float32))
+
+        s_float = ct.astype(s_tile, ct.float16)
+        qjl_ip  = ct.mma(qp_tile, ct.transpose(s_float),
+                         ct.zeros((BLOCK_Q, BLOCK_KV), dtype=ct.float32))
+        rn_f32  = ct.astype(rn, ct.float32)
+        term2   = correction_scale * qjl_ip * ct.expand_dims(rn_f32, axis=0)
+
+        raw_scores = term1 + term2
+
+        if USE_SWIZZLE:
+            m_new = ct.maximum(m_i, ct.max(raw_scores, axis=1) * scale_log2)
+            alpha = ct.exp2(m_i - m_new, flush_to_zero=True)
+            p     = ct.exp2(raw_scores * scale_log2 - ct.expand_dims(m_new, axis=1),
+                            flush_to_zero=True)
+        else:
+            scores = raw_scores * scale
+            m_new  = ct.maximum(m_i, ct.max(scores, axis=1))
+            alpha  = ct.exp(m_i - m_new)
+            p      = ct.exp(scores - ct.expand_dims(m_new, axis=1))
+
+        l_i = alpha * l_i + ct.sum(p, axis=1)
+
+        p_fp16 = ct.astype(p, ct.float16)
+        acc = ct.expand_dims(alpha, axis=1) * acc + ct.mma(
+            p_fp16, v_tile, ct.zeros((BLOCK_Q, HEAD_DIM), dtype=ct.float32))
+
+        m_i = m_new
+
+    if USE_SWIZZLE and _HAS_APPROX:
+        out = ct.truediv(acc, ct.expand_dims(l_i, axis=1),
+                         flush_to_zero=True, rounding_mode=RMd.APPROX)
+    else:
+        out = acc / ct.expand_dims(l_i, axis=1)
+
+    ct.store(Output, index=(q_block, 0), tile=out)

turboquant_gpu/codebook.py

@@ -0,0 +1,70 @@
+"""
+lloyd-max optimal scalar quantizer.
+
+after rotation each coordinate ~ N(0, 1/d).  we solve lloyd-max
+(1d k-means against that pdf) to get optimal centroids + boundaries.
+runs once on cpu at init — the codebook is tiny.
+"""
+
+import math
+import torch
+from scipy import integrate
+
+
+def _gaussian_pdf(x, sigma):
+    return (1.0 / (math.sqrt(2 * math.pi) * sigma)) * math.exp(
+        -x * x / (2 * sigma * sigma)
+    )
+
+
+def solve_lloyd_max(d, bits, max_iter=200, tol=1e-10):
+    """return (centroids, boundaries) as sorted float32 tensors."""
+    n_levels = 1 << bits
+    sigma = 1.0 / math.sqrt(d)
+    pdf = lambda x: _gaussian_pdf(x, sigma)
+
+    lo, hi = -3.5 * sigma, 3.5 * sigma
+    centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)]
+
+    for _ in range(max_iter):
+        boundaries = [
+            (centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)
+        ]
+        edges = [lo * 3] + boundaries + [hi * 3]
+        new_centroids = []
+        for i in range(n_levels):
+            a, b = edges[i], edges[i + 1]
+            num, _ = integrate.quad(lambda x: x * pdf(x), a, b)
+            den, _ = integrate.quad(pdf, a, b)
+            new_centroids.append(num / den if den > 1e-15 else centroids[i])
+        if max(abs(new_centroids[i] - centroids[i]) for i in range(n_levels)) < tol:
+            break
+        centroids = new_centroids
+
+    boundaries = [
+        (centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)
+    ]
+    return (
+        torch.tensor(centroids, dtype=torch.float32),
+        torch.tensor(boundaries, dtype=torch.float32),
+    )
+
+
+class LloydMaxCodebook:
+    """pre-solved lloyd-max codebook for a given (d, bits) pair."""
+
+    def __init__(self, d, bits):
+        self.d = d
+        self.bits = bits
+        self.n_levels = 1 << bits
+        self.centroids, self.boundaries = solve_lloyd_max(d, bits)
+
+    def quantize(self, x):
+        diffs = x.unsqueeze(-1) - self.centroids.to(x.device)
+        return diffs.abs().argmin(dim=-1)
+
+    def dequantize(self, indices):
+        return self.centroids.to(indices.device)[indices.long()]
+
+    def __repr__(self):
+        return f"LloydMaxCodebook(d={self.d}, bits={self.bits}, levels={self.n_levels})"

turboquant_gpu/compress.py

@@ -0,0 +1,210 @@
+"""
+key and value compression kernels.
+
+pipeline per key token:
+  normalize → rotate (Pi^T) → lloyd-max quantize → store indices + norms
+  residual → QJL project (S^T) → store signs + residual norms
+
+value compression is the same minus the QJL step.
+"""
+
+try:
+    import cuda.tile as ct
+except ImportError:
+    import cutile as ct  # type: ignore
+
+from .constants import BLOCK_S, HEAD_DIM
+
+
+# ── key compression ──────────────────────────────────────────────────
+
+@ct.kernel
+def turboquant_compress_2bit(
+    K, Pi_T, Pi, S_T,
+    Indices, Signs, Norms, RNorms,
+    c0: float, c1: float, c2: float, c3: float,
+    b1: float, b2: float, b3: float,
+    seq_k: int,
+):
+    """2-bit mse (4 centroids) + 1-bit qjl.  total_bits=3."""
+    block_id = ct.bid(0)
+    zero_pad = ct.PaddingMode.ZERO
+
+    k_tile = ct.load(K, index=(block_id, 0), shape=(BLOCK_S, HEAD_DIM),
+                     padding_mode=zero_pad)
+    pi_t = ct.load(Pi_T, index=(0, 0), shape=(HEAD_DIM, HEAD_DIM))
+    pi   = ct.load(Pi,   index=(0, 0), shape=(HEAD_DIM, HEAD_DIM))
+    s_t  = ct.load(S_T,  index=(0, 0), shape=(HEAD_DIM, HEAD_DIM))
+
+    k_f32 = ct.astype(k_tile, ct.float32)
+    norms = ct.sqrt(ct.sum(k_f32 * k_f32, axis=1))
+    safe_norms = ct.where(norms > 1e-8, norms, 1e-8)
+    k_normed = k_f32 / ct.expand_dims(safe_norms, axis=1)
+
+    # rotate into quantization basis
+    y = ct.mma(ct.astype(k_normed, ct.float16), pi_t,
+               ct.zeros((BLOCK_S, HEAD_DIM), dtype=ct.float32))
+
+    # lloyd-max assignment via boundary comparisons
+    idx = ct.zeros((BLOCK_S, HEAD_DIM), dtype=ct.float32)
+    idx = ct.where(y > b1, 1.0, idx)
+    idx = ct.where(y > b2, 2.0, idx)
+    idx = ct.where(y > b3, 3.0, idx)
+
+    # dequantize and un-rotate for mse reconstruction
+    y_hat = ct.full((BLOCK_S, HEAD_DIM), c0, dtype=ct.float32)
+    y_hat = ct.where(idx > 0.5, c1, y_hat)
+    y_hat = ct.where(idx > 1.5, c2, y_hat)
+    y_hat = ct.where(idx > 2.5, c3, y_hat)
+
+    k_bar_hat = ct.mma(ct.astype(y_hat, ct.float16), pi,
+                       ct.zeros((BLOCK_S, HEAD_DIM), dtype=ct.float32))
+    k_mse = k_bar_hat * ct.expand_dims(norms, axis=1)
+
+    # qjl: project residual for bias-correction signs
+    residual = k_f32 - k_mse
+    r_norms = ct.sqrt(ct.sum(residual * residual, axis=1))
+    projected = ct.mma(ct.astype(residual, ct.float16), s_t,
+                       ct.zeros((BLOCK_S, HEAD_DIM), dtype=ct.float32))
+    signs = ct.where(projected >= 0.0, 1.0, -1.0)
+
+    ct.store(Indices, index=(block_id, 0), tile=ct.astype(idx, ct.uint8))
+    ct.store(Signs,   index=(block_id, 0), tile=ct.astype(signs, ct.int8))
+    ct.store(Norms,   index=(block_id,),   tile=ct.astype(norms, ct.float16))
+    ct.store(RNorms,  index=(block_id,),   tile=ct.astype(r_norms, ct.float16))
+
+
+@ct.kernel
+def turboquant_compress_3bit(
+    K, Pi_T, Pi, S_T,
+    Indices, Signs, Norms, RNorms,
+    c0: float, c1: float, c2: float, c3: float,
+    c4: float, c5: float, c6: float, c7: float,
+    b1: float, b2: float, b3: float, b4: float,
+    b5: float, b6: float, b7: float,
+    seq_k: int,
+):
+    """3-bit mse (8 centroids) + 1-bit qjl.  total_bits=4."""
+    block_id = ct.bid(0)
+    zero_pad = ct.PaddingMode.ZERO
+
+    k_tile = ct.load(K, index=(block_id, 0), shape=(BLOCK_S, HEAD_DIM),
+                     padding_mode=zero_pad)
+    pi_t = ct.load(Pi_T, index=(0, 0), shape=(HEAD_DIM, HEAD_DIM))
+    pi   = ct.load(Pi,   index=(0, 0), shape=(HEAD_DIM, HEAD_DIM))
+    s_t  = ct.load(S_T,  index=(0, 0), shape=(HEAD_DIM, HEAD_DIM))
+
+    k_f32 = ct.astype(k_tile, ct.float32)
+    norms = ct.sqrt(ct.sum(k_f32 * k_f32, axis=1))
+    safe_norms = ct.where(norms > 1e-8, norms, 1e-8)
+    k_normed = k_f32 / ct.expand_dims(safe_norms, axis=1)
+
+    y = ct.mma(ct.astype(k_normed, ct.float16), pi_t,
+               ct.zeros((BLOCK_S, HEAD_DIM), dtype=ct.float32))
+
+    idx = ct.zeros((BLOCK_S, HEAD_DIM), dtype=ct.float32)
+    idx = ct.where(y > b1, 1.0, idx)
+    idx = ct.where(y > b2, 2.0, idx)
+    idx = ct.where(y > b3, 3.0, idx)
+    idx = ct.where(y > b4, 4.0, idx)
+    idx = ct.where(y > b5, 5.0, idx)
+    idx = ct.where(y > b6, 6.0, idx)
+    idx = ct.where(y > b7, 7.0, idx)
+
+    y_hat = ct.full((BLOCK_S, HEAD_DIM), c0, dtype=ct.float32)
+    y_hat = ct.where(idx > 0.5, c1, y_hat)
+    y_hat = ct.where(idx > 1.5, c2, y_hat)
+    y_hat = ct.where(idx > 2.5, c3, y_hat)
+    y_hat = ct.where(idx > 3.5, c4, y_hat)
+    y_hat = ct.where(idx > 4.5, c5, y_hat)
+    y_hat = ct.where(idx > 5.5, c6, y_hat)
+    y_hat = ct.where(idx > 6.5, c7, y_hat)
+
+    k_bar_hat = ct.mma(ct.astype(y_hat, ct.float16), pi,
+                       ct.zeros((BLOCK_S, HEAD_DIM), dtype=ct.float32))
+    k_mse = k_bar_hat * ct.expand_dims(norms, axis=1)
+
+    residual = k_f32 - k_mse
+    r_norms = ct.sqrt(ct.sum(residual * residual, axis=1))
+    projected = ct.mma(ct.astype(residual, ct.float16), s_t,
+                       ct.zeros((BLOCK_S, HEAD_DIM), dtype=ct.float32))
+    signs = ct.where(projected >= 0.0, 1.0, -1.0)
+
+    ct.store(Indices, index=(block_id, 0), tile=ct.astype(idx, ct.uint8))
+    ct.store(Signs,   index=(block_id, 0), tile=ct.astype(signs, ct.int8))
+    ct.store(Norms,   index=(block_id,),   tile=ct.astype(norms, ct.float16))
+    ct.store(RNorms,  index=(block_id,),   tile=ct.astype(r_norms, ct.float16))
+
+
+# ── value compression (mse only, no qjl) ────────────────────────────
+
+@ct.kernel
+def turboquant_compress_values_3bit(
+    V, Pi_T,
+    Indices, Norms,
+    c0: float, c1: float, c2: float, c3: float,
+    c4: float, c5: float, c6: float, c7: float,
+    b1: float, b2: float, b3: float, b4: float,
+    b5: float, b6: float, b7: float,
+    seq_v: int,
+):
+    """3-bit value compression (8 levels)."""
+    block_id = ct.bid(0)
+    zero_pad = ct.PaddingMode.ZERO
+
+    v_tile = ct.load(V, index=(block_id, 0), shape=(BLOCK_S, HEAD_DIM),
+                     padding_mode=zero_pad)
+    pi_t = ct.load(Pi_T, index=(0, 0), shape=(HEAD_DIM, HEAD_DIM))
+
+    v_f32 = ct.astype(v_tile, ct.float32)
+    norms = ct.sqrt(ct.sum(v_f32 * v_f32, axis=1))
+    safe_norms = ct.where(norms > 1e-8, norms, 1e-8)
+    v_normed = v_f32 / ct.expand_dims(safe_norms, axis=1)
+
+    y = ct.mma(ct.astype(v_normed, ct.float16), pi_t,
+               ct.zeros((BLOCK_S, HEAD_DIM), dtype=ct.float32))
+
+    idx = ct.zeros((BLOCK_S, HEAD_DIM), dtype=ct.float32)
+    idx = ct.where(y > b1, 1.0, idx)
+    idx = ct.where(y > b2, 2.0, idx)
+    idx = ct.where(y > b3, 3.0, idx)
+    idx = ct.where(y > b4, 4.0, idx)
+    idx = ct.where(y > b5, 5.0, idx)
+    idx = ct.where(y > b6, 6.0, idx)
+    idx = ct.where(y > b7, 7.0, idx)
+
+    ct.store(Indices, index=(block_id, 0), tile=ct.astype(idx, ct.uint8))
+    ct.store(Norms,   index=(block_id,),   tile=ct.astype(norms, ct.float16))
+
+
+@ct.kernel
+def turboquant_compress_values_2bit(
+    V, Pi_T,
+    Indices, Norms,
+    c0: float, c1: float, c2: float, c3: float,
+    b1: float, b2: float, b3: float,
+    seq_v: int,
+):
+    """2-bit value compression (4 levels)."""
+    block_id = ct.bid(0)
+    zero_pad = ct.PaddingMode.ZERO
+
+    v_tile = ct.load(V, index=(block_id, 0), shape=(BLOCK_S, HEAD_DIM),
+                     padding_mode=zero_pad)
+    pi_t = ct.load(Pi_T, index=(0, 0), shape=(HEAD_DIM, HEAD_DIM))
+
+    v_f32 = ct.astype(v_tile, ct.float32)
+    norms = ct.sqrt(ct.sum(v_f32 * v_f32, axis=1))
+    safe_norms = ct.where(norms > 1e-8, norms, 1e-8)
+    v_normed = v_f32 / ct.expand_dims(safe_norms, axis=1)
+
+    y = ct.mma(ct.astype(v_normed, ct.float16), pi_t,
+               ct.zeros((BLOCK_S, HEAD_DIM), dtype=ct.float32))
+
+    idx = ct.zeros((BLOCK_S, HEAD_DIM), dtype=ct.float32)
+    idx = ct.where(y > b1, 1.0, idx)
+    idx = ct.where(y > b2, 2.0, idx)
+    idx = ct.where(y > b3, 3.0, idx)
+
+    ct.store(Indices, index=(block_id, 0), tile=ct.astype(idx, ct.uint8))
+    ct.store(Norms,   index=(block_id,),   tile=ct.astype(norms, ct.float16))

turboquant_gpu/constants.py

@@ -0,0 +1,8 @@
+HEAD_DIM = 128
+BLOCK_Q  = 16
+BLOCK_KV = 64
+BLOCK_S  = 64
+
+SUPPORTED_MSE_BITS = {1, 2, 3, 4}
+DEFAULT_TOTAL_BITS = 3
+DEFAULT_SEED = 42

turboquant_gpu/decompress.py

@@ -0,0 +1,77 @@
+"""
+value decompression kernels.
+
+indices → dequantize via centroids → un-rotate (Pi) → scale by norms.
+"""
+
+try:
+    import cuda.tile as ct
+except ImportError:
+    import cutile as ct  # type: ignore
+
+from .constants import BLOCK_S, HEAD_DIM
+
+
+@ct.kernel
+def turboquant_decompress_3bit(
+    Indices, Norms, Pi,
+    Output,
+    c0: float, c1: float, c2: float, c3: float,
+    c4: float, c5: float, c6: float, c7: float,
+    seq_k: int,
+):
+    block_id = ct.bid(0)
+    zero_pad = ct.PaddingMode.ZERO
+
+    idx_tile  = ct.load(Indices, index=(block_id, 0), shape=(BLOCK_S, HEAD_DIM),
+                        padding_mode=zero_pad)
+    norm_tile = ct.load(Norms, index=(block_id,), shape=(BLOCK_S,),
+                        padding_mode=zero_pad)
+    pi = ct.load(Pi, index=(0, 0), shape=(HEAD_DIM, HEAD_DIM))
+
+    idx_f32 = ct.astype(idx_tile, ct.float32)
+
+    y_hat = ct.full((BLOCK_S, HEAD_DIM), c0, dtype=ct.float32)
+    y_hat = ct.where(idx_f32 > 0.5, c1, y_hat)
+    y_hat = ct.where(idx_f32 > 1.5, c2, y_hat)
+    y_hat = ct.where(idx_f32 > 2.5, c3, y_hat)
+    y_hat = ct.where(idx_f32 > 3.5, c4, y_hat)
+    y_hat = ct.where(idx_f32 > 4.5, c5, y_hat)
+    y_hat = ct.where(idx_f32 > 5.5, c6, y_hat)
+    y_hat = ct.where(idx_f32 > 6.5, c7, y_hat)
+
+    x_hat  = ct.mma(ct.astype(y_hat, ct.float16), pi,
+                    ct.zeros((BLOCK_S, HEAD_DIM), dtype=ct.float32))
+    result = x_hat * ct.expand_dims(ct.astype(norm_tile, ct.float32), axis=1)
+
+    ct.store(Output, index=(block_id, 0), tile=ct.astype(result, ct.float16))
+
+
+@ct.kernel
+def turboquant_decompress_2bit(
+    Indices, Norms, Pi,
+    Output,
+    c0: float, c1: float, c2: float, c3: float,
+    seq_k: int,
+):
+    block_id = ct.bid(0)
+    zero_pad = ct.PaddingMode.ZERO
+
+    idx_tile  = ct.load(Indices, index=(block_id, 0), shape=(BLOCK_S, HEAD_DIM),
+                        padding_mode=zero_pad)
+    norm_tile = ct.load(Norms, index=(block_id,), shape=(BLOCK_S,),
+                        padding_mode=zero_pad)
+    pi = ct.load(Pi, index=(0, 0), shape=(HEAD_DIM, HEAD_DIM))
+
+    idx_f32 = ct.astype(idx_tile, ct.float32)
+
+    y_hat = ct.full((BLOCK_S, HEAD_DIM), c0, dtype=ct.float32)
+    y_hat = ct.where(idx_f32 > 0.5, c1, y_hat)
+    y_hat = ct.where(idx_f32 > 1.5, c2, y_hat)
+    y_hat = ct.where(idx_f32 > 2.5, c3, y_hat)
+
+    x_hat  = ct.mma(ct.astype(y_hat, ct.float16), pi,
+                    ct.zeros((BLOCK_S, HEAD_DIM), dtype=ct.float32))
+    result = x_hat * ct.expand_dims(ct.astype(norm_tile, ct.float32), axis=1)
+
+    ct.store(Output, index=(block_id, 0), tile=ct.astype(result, ct.float16))

turboquant_gpu/host.py

@@ -0,0 +1,288 @@
+"""
+host-side engine: compress / decompress KV caches, run generation.
+
+tries cutile kernels first; falls back to pytorch if the driver or
+import is unavailable.
+"""
+
+import math
+import torch
+
+from .codebook import LloydMaxCodebook
+from .constants import BLOCK_S, DEFAULT_SEED, DEFAULT_TOTAL_BITS, HEAD_DIM
+
+
+def _rotation_matrix(d, seed, device="cpu"):
+    gen = torch.Generator(device="cpu")
+    gen.manual_seed(seed)
+    G = torch.randn(d, d, generator=gen)
+    Q, R = torch.linalg.qr(G)
+    diag_sign = torch.sign(torch.diag(R))
+    diag_sign[diag_sign == 0] = 1.0
+    return (Q * diag_sign.unsqueeze(0)).to(device)
+
+
+def _qjl_matrix(d, seed, device="cpu"):
+    gen = torch.Generator(device="cpu")
+    gen.manual_seed(seed + 10000)
+    return torch.randn(d, d, generator=gen).to(device)
+
+
+class TurboQuantEngine:
+    def __init__(self, head_dim=HEAD_DIM, total_bits=DEFAULT_TOTAL_BITS,
+                 seed=DEFAULT_SEED, device="cpu"):
+        self.head_dim   = head_dim
+        self.total_bits = total_bits
+        self.mse_bits   = max(total_bits - 1, 1)
+        self.device     = device
+
+        self.Pi  = _rotation_matrix(head_dim, seed, device)
+        self.PiT = self.Pi.T.contiguous()
+        self.S   = _qjl_matrix(head_dim, seed, device)
+        self.ST  = self.S.T.contiguous()
+
+        self.key_codebook = LloydMaxCodebook(head_dim, self.mse_bits)
+        self.val_codebook = LloydMaxCodebook(head_dim, total_bits)
+
+    # ── public api ───────────────────────────────────────────────────
+
+    @torch.no_grad()
+    def compress_kv_cache(self, past_key_values):
+        """compress a full KV cache from a huggingface model forward pass."""
+        kv_keys, kv_vals = self._extract_kv(past_key_values)
+
+        layers = []
+        for li in range(len(kv_keys)):
+            n_heads = kv_keys[li].shape[1]
+            ck = [self._compress_keys(kv_keys[li][0, h].half().contiguous())
+                  for h in range(n_heads)]
+            cv = [self._compress_values(kv_vals[li][0, h].half().contiguous())
+                  for h in range(n_heads)]
+            layers.append((ck, cv))
+
+        return {"layers": layers}
+
+    @torch.no_grad()
+    def build_cache(self, compressed):
+        """reconstruct a DynamicCache from compressed K (via k_mse) and V."""
+        from transformers import DynamicCache
+        cache = DynamicCache()
+        for li, (ck_list, cv_list) in enumerate(compressed["layers"]):
+            k_heads = [ck["k_mse"] for ck in ck_list]
+            k_layer = torch.stack(k_heads).unsqueeze(0)
+            v_heads = [self._decompress_values(cv) for cv in cv_list]
+            v_layer = torch.stack(v_heads).unsqueeze(0)
+            cache.update(k_layer, v_layer, li)
+        return cache
+
+    @torch.no_grad()
+    def generate(self, model, tokenizer, prompt, max_new_tokens=100,
+                 repetition_penalty=1.3):
+        """prefill → compress KV → decode with repetition penalty."""
+        inputs = tokenizer(prompt, return_tensors="pt").to(self.device)
+        out = model(**inputs, use_cache=True)
+
+        compressed = self.compress_kv_cache(out.past_key_values)
+        stats = self.compression_stats(out.past_key_values)
+        cache = self.build_cache(compressed)
+
+        seq_len = inputs["input_ids"].shape[1]
+        all_ids = inputs["input_ids"][0].tolist()
+
+        next_tok = out.logits[:, -1:].argmax(dim=-1)
+        all_ids.append(next_tok.item())
+
+        for step in range(max_new_tokens - 1):
+            o = model(input_ids=next_tok, past_key_values=cache,
+                      position_ids=torch.tensor([[seq_len + step]],
+                                                device=self.device),
+                      use_cache=True)
+            cache = o.past_key_values
+            logits = o.logits[:, -1, :]
+
+            if repetition_penalty != 1.0:
+                for tid in set(all_ids):
+                    if logits[0, tid] > 0:
+                        logits[0, tid] /= repetition_penalty
+                    else:
+                        logits[0, tid] *= repetition_penalty
+
+            next_tok = logits.argmax(dim=-1, keepdim=True)
+            all_ids.append(next_tok.item())
+            if next_tok.item() == tokenizer.eos_token_id:
+                break
+
+        n_new = len(all_ids) - seq_len
+        text = tokenizer.decode(all_ids, skip_special_tokens=True)
+        return {"text": text, "tokens": n_new, "stats": stats}
+
+    def compression_stats(self, past_key_values):
+        """return compression ratio and byte counts."""
+        kv_keys, kv_vals = self._extract_kv(past_key_values)
+        n_layers = len(kv_keys)
+        n_heads  = kv_keys[0].shape[1]
+        seq_len  = kv_keys[0].shape[2]
+
+        fp16_bytes = sum(k.nbytes + v.nbytes for k, v in zip(kv_keys, kv_vals))
+        tq_bytes   = self._compressed_bytes(seq_len) * n_heads * n_layers
+
+        return {
+            "seq_len": seq_len, "n_layers": n_layers, "n_heads": n_heads,
+            "fp16_bytes": fp16_bytes, "tq_bytes": tq_bytes,
+            "ratio": fp16_bytes / tq_bytes,
+        }
+
+    # ── internals ────────────────────────────────────────────────────
+
+    @staticmethod
+    def _extract_kv(past_key_values):
+        """handle both DynamicCache and list-of-tuples formats."""
+        try:
+            return past_key_values.key_cache, past_key_values.value_cache
+        except AttributeError:
+            keys = [kv[0] for kv in past_key_values]
+            vals = [kv[1] for kv in past_key_values]
+            return keys, vals
+
+    def _compress_keys(self, K):
+        try:
+            from .compress import turboquant_compress_2bit, turboquant_compress_3bit
+            import cuda.tile as ct
+
+            seq_k, d = K.shape
+            indices = torch.empty(seq_k, d, dtype=torch.uint8, device=K.device)
+            signs   = torch.empty(seq_k, d, dtype=torch.int8,  device=K.device)
+            norms   = torch.empty(seq_k, dtype=torch.float16,  device=K.device)
+            r_norms = torch.empty(seq_k, dtype=torch.float16,  device=K.device)
+
+            grid = (self._cdiv(seq_k, BLOCK_S), 1, 1)
+            c = self.key_codebook.centroids.tolist()
+            b = self.key_codebook.boundaries.tolist()
+            stream = torch.cuda.current_stream()
+
+            if self.mse_bits == 2:
+                ct.launch(stream, grid, turboquant_compress_2bit, (
+                    K, self.PiT.half(), self.Pi.half(), self.ST.half(),
+                    indices, signs, norms, r_norms, *c, *b, seq_k))
+            elif self.mse_bits == 3:
+                ct.launch(stream, grid, turboquant_compress_3bit, (
+                    K, self.PiT.half(), self.Pi.half(), self.ST.half(),
+                    indices, signs, norms, r_norms, *c, *b, seq_k))
+            else:
+                return self._compress_keys_pt(K)
+
+            k_mse = self._dequant_keys(indices, norms)
+            return {"indices": indices, "k_mse": k_mse, "qjl_signs": signs,
+                    "vec_norms": norms, "residual_norms": r_norms}
+        except (ImportError, RuntimeError):
+            return self._compress_keys_pt(K)
+
+    def _compress_values(self, V):
+        try:
+            from .compress import (turboquant_compress_values_3bit,
+                                   turboquant_compress_values_2bit)
+            import cuda.tile as ct
+
+            seq_v, d = V.shape
+            indices = torch.empty(seq_v, d, dtype=torch.uint8, device=V.device)
+            norms   = torch.empty(seq_v, dtype=torch.float16,  device=V.device)
+
+            grid = (self._cdiv(seq_v, BLOCK_S), 1, 1)
+            c = self.val_codebook.centroids.tolist()
+            b = self.val_codebook.boundaries.tolist()
+            stream = torch.cuda.current_stream()
+
+            if self.total_bits == 3:
+                ct.launch(stream, grid, turboquant_compress_values_3bit, (
+                    V, self.PiT.half(), indices, norms, *c, *b, seq_v))
+            elif self.total_bits == 2:
+                ct.launch(stream, grid, turboquant_compress_values_2bit, (
+                    V, self.PiT.half(), indices, norms, *c, *b, seq_v))
+            else:
+                return self._compress_values_pt(V)
+
+            return {"indices": indices, "vec_norms": norms}
+        except (ImportError, RuntimeError):
+            return self._compress_values_pt(V)
+
+    def _decompress_values(self, cv):
+        try:
+            from .decompress import (turboquant_decompress_3bit,
+                                     turboquant_decompress_2bit)
+            import cuda.tile as ct
+
+            indices = cv["indices"]
+            norms   = cv["vec_norms"]
+            seq_v   = indices.shape[0]
+            output  = torch.empty(seq_v, self.head_dim, dtype=torch.float16,
+                                  device=indices.device)
+
+            grid = (self._cdiv(seq_v, BLOCK_S), 1, 1)
+            c = self.val_codebook.centroids.tolist()
+            stream = torch.cuda.current_stream()
+
+            if self.total_bits == 3:
+                ct.launch(stream, grid, turboquant_decompress_3bit, (
+                    indices, norms, self.Pi.half(), output, *c, seq_v))
+            elif self.total_bits == 2:
+                ct.launch(stream, grid, turboquant_decompress_2bit, (
+                    indices, norms, self.Pi.half(), output, *c, seq_v))
+            else:
+                return self._decompress_values_pt(cv)
+            return output
+        except (ImportError, RuntimeError):
+            return self._decompress_values_pt(cv)
+
+    # ── pytorch fallbacks ────────────────────────────────────────────
+
+    def _compress_keys_pt(self, K):
+        K_f = K.float()
+        norms = torch.norm(K_f, dim=-1, keepdim=True)
+        K_normed = K_f / (norms + 1e-8)
+        rotated = K_normed @ self.PiT.float()
+
+        c = self.key_codebook.centroids.to(K.device)
+        indices = (rotated.unsqueeze(-1) - c).abs().argmin(dim=-1).to(torch.uint8)
+
+        y_hat = c[indices.long()]
+        k_mse = (y_hat @ self.Pi.float()) * norms
+        residual = K_f - k_mse
+        r_norms = torch.norm(residual, dim=-1)
+
+        signs = torch.sign(residual @ self.ST.float()).to(torch.int8)
+        signs[signs == 0] = 1
+
+        return {"indices": indices, "k_mse": k_mse.half(), "qjl_signs": signs,
+                "vec_norms": norms.squeeze(-1).half(),
+                "residual_norms": r_norms.half()}
+
+    def _compress_values_pt(self, V):
+        V_f = V.float()
+        norms = torch.norm(V_f, dim=-1, keepdim=True)
+        V_normed = V_f / (norms + 1e-8)
+        rotated = V_normed @ self.PiT.float()
+
+        c = self.val_codebook.centroids.to(V.device)
+        indices = (rotated.unsqueeze(-1) - c).abs().argmin(dim=-1).to(torch.uint8)
+
+        return {"indices": indices, "vec_norms": norms.squeeze(-1).half()}
+
+    def _decompress_values_pt(self, cv):
+        c = self.val_codebook.centroids.to(cv["indices"].device)
+        y_hat = c[cv["indices"].long()]
+        norms = cv["vec_norms"].float().unsqueeze(-1)
+        return ((y_hat @ self.Pi.float()) * norms).half()
+
+    def _dequant_keys(self, indices, norms):
+        c = self.key_codebook.centroids.to(indices.device)
+        y_hat = c[indices.long()]
+        return ((y_hat.float() @ self.Pi.float()) * norms.float().unsqueeze(-1)).half()
+
+    def _cdiv(self, a, b):
+        return (a + b - 1) // b
+
+    def _compressed_bytes(self, seq_len):
+        d = self.head_dim
+        key_bytes = (seq_len * d * self.mse_bits + seq_len * d + seq_len * 32) / 8
+        val_bytes = (seq_len * d * self.total_bits + seq_len * 16) / 8
+        return key_bytes + val_bytes

turboquant_gpu-0.1.0.dist-info/METADATA

@@ -0,0 +1,96 @@
+Metadata-Version: 2.4
+Name: turboquant-gpu
+Version: 0.1.0
+Summary: TurboQuant KV cache compression for LLM inference — cuTile GPU kernels
+Author: Anirudh Bharadwaj Vangara
+License-Expression: MIT
+Keywords: quantization,kv-cache,llm,inference,cutile,cuda,gpu,attention,blackwell,hopper,h100,b200
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: scipy
+Provides-Extra: gpu
+Requires-Dist: cuda-tile; extra == "gpu"
+Dynamic: license-file
+
+# turboquant-gpu
+
+**5.02x KV cache compression for LLM inference** — GPU-accelerated cuTile kernels with PyTorch fallback.
+
+```
+pip install turboquant-gpu
+```
+
+## quick start
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from turboquant_gpu import TurboQuantEngine
+import torch
+
+model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-1.5B", torch_dtype=torch.float16, device_map="cuda")
+tok   = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B")
+
+engine = TurboQuantEngine(head_dim=128, total_bits=3, device="cuda")
+result = engine.generate(model, tok, "The key to efficient LLM inference is")
+
+print(result["text"])
+print(f"{result['tokens']} tokens | {result['stats']['ratio']:.1f}x compression")
+```
+
+## how it works
+
+Implements the [TurboQuant](https://arxiv.org/abs/2501.09747) algorithm:
+
+1. **normalize + rotate** — random orthogonal rotation (Pi) makes coordinates near-Gaussian
+2. **Lloyd-Max quantize** — optimal scalar quantization against N(0, 1/d)
+3. **QJL bias correction** — 1-bit sign sketch of the residual for unbiased key scores
+
+At 3-bit (2-bit MSE + 1-bit QJL) this gives ~5x compression with negligible quality loss.
+
+## step-by-step api
+
+```python
+engine = TurboQuantEngine(head_dim=128, total_bits=3, device="cuda")
+
+# after model prefill:
+compressed = engine.compress_kv_cache(out.past_key_values)
+cache      = engine.build_cache(compressed)
+stats      = engine.compression_stats(out.past_key_values)
+
+# or just do it all in one call:
+result = engine.generate(model, tokenizer, "your prompt here")
+```
+
+## gpu support
+
+Written in [cuTile](https://docs.nvidia.com/cuda/cutile-python/) for cross-architecture portability.
+Falls back to PyTorch if cuTile or a compatible driver isn't available.
+
+| GPU family | Architecture | Status |
+|------------|-------------|--------|
+| A100       | Ampere      | supported (PyTorch fallback) |
+| H100       | Hopper      | supported |
+| B200/B300  | Blackwell   | supported + swizzle fast path |
+
+## kernels
+
+| kernel | what it does |
+|--------|-------------|
+| `compress_keys` | normalize → rotate(Pi^T) → Lloyd-Max quantize → QJL signs |
+| `compress_values` | normalize → rotate(Pi^T) → Lloyd-Max quantize |
+| `decompress_values` | dequantize → un-rotate(Pi) → scale by norms |
+| `attention_scores` | asymmetric dot product with QJL correction |
+| `fused_attention` | scores + online softmax + V accumulation |
+
+## license
+
+MIT

turboquant_gpu-0.1.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+turboquant_gpu/__init__.py,sha256=hVDM-NQiB_PzgSyqQXi2KoMEUspka_PKvw55yGPhODk,278
+turboquant_gpu/attention.py,sha256=O13thg9bO7h9OIySgmhYnv61aWidWmC5wkCKKVuh4L8,15147
+turboquant_gpu/codebook.py,sha256=cJAOa_qnUBAJqMwNrHIK2igbmX61cQ4gbVXGfr9hyRk,2240
+turboquant_gpu/compress.py,sha256=yP6BYrk29yAG8dMxkT2rUet1Z8qThOLRb8AWpwwbu6s,8003
+turboquant_gpu/constants.py,sha256=VSmWSTSlxO0e0QTfvtxso41EvJTR0gCKJK4QL7POXpM,133
+turboquant_gpu/decompress.py,sha256=b8VKHC3ZmmzFC_CKgP_BmlL0OGooFI30ECYWM2DXbPo,2581
+turboquant_gpu/host.py,sha256=mVzm4Vhl3NT0clN4XXR4Z3VpEZ9PyFRqwxtexmTTvng,11848
+turboquant_gpu-0.1.0.dist-info/licenses/LICENSE,sha256=jdrItB2_D55fIzIsSe7Esi-l2ykHOQS1CmAr0Y_fBAA,1082
+turboquant_gpu-0.1.0.dist-info/METADATA,sha256=9dqKsJja0ErrdAffd9rXw4r3e--y4SoSkGeHG2e2dKI,3314
+turboquant_gpu-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+turboquant_gpu-0.1.0.dist-info/top_level.txt,sha256=ZrloYBosuQLyIN7iBgLJS9NENKJMkYszJU3zi7aCIBI,15
+turboquant_gpu-0.1.0.dist-info/RECORD,,

turboquant_gpu-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

turboquant_gpu-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Anirudh Bharadwaj Vangara
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

turboquant_gpu-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+turboquant_gpu