Function turboquant_compress_3bit

turboquant_cutile/compress.py:66–126

3-bit MSE (8 centroids) + 1-bit QJL. total_bits=4.

(
    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,
)

@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))