custom_generate/tq_value_mse_decode.py

#!/usr/bin/env python3
"""Decode side of the Lloyd-Max value codec.

Two shapes are needed, mirroring the two paths in vLLM's TurboQuant backend:

``unpack_values_rotated``
    Reads the packed indices and the fp16 norm and returns value vectors in
    *rotated* space.  This is what the fused decode kernel accumulates against:
    the attention weight carries the norm and the whole output is rotated back
    once, so no cached vector is ever rotated individually.

``unpack_values_model_space``
    The materialising form used by the continuation-prefill path, where the
    inverse rotation is one batched GEMM over the whole dequantised tensor.

The unpack itself is the key path's MSE branch applied at the value offset:
gather centroids by index, optionally renormalise (norm correction), scale by
the stored vector norm.
"""

from __future__ import annotations

import torch
import triton
import triton.language as tl


@triton.jit
def _zenit_unpack_values(
    KV_cache_ptr,
    Slot_mapping_ptr,
    Centroids_ptr,
    Out_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_PACKED: tl.constexpr,
    VAL_BYTES: tl.constexpr,
    VAL_BITS: tl.constexpr,
    NORM_CORRECTION: tl.constexpr,
):
    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
    )
    val_base = slot_base + KEY_PACKED

    d_offs = tl.arange(0, BLOCK_D)
    d_mask = d_offs < D

    if VAL_BITS == 4:
        byte_idx = d_offs // 2
        shift = (d_offs % 2) * 4
        raw = tl.load(KV_cache_ptr + val_base + byte_idx, mask=d_mask, other=0).to(
            tl.int32
        )
        idx = (raw >> shift) & 0xF
    elif VAL_BITS == 3:
        bit_off = d_offs * 3
        byte_idx = bit_off // 8
        shift = bit_off % 8
        raw0 = tl.load(KV_cache_ptr + val_base + byte_idx, mask=d_mask, other=0).to(
            tl.int32
        )
        raw1 = tl.load(KV_cache_ptr + val_base + byte_idx + 1, mask=d_mask, other=0).to(
            tl.int32
        )
        idx = ((raw0 | (raw1 << 8)) >> shift) & 0x7
    else:  # VAL_BITS == 2
        byte_idx = d_offs // 4
        shift = (d_offs % 4) * 2
        raw = tl.load(KV_cache_ptr + val_base + byte_idx, mask=d_mask, other=0).to(
            tl.int32
        )
        idx = (raw >> shift) & 0x3

    unit = tl.load(Centroids_ptr + idx, mask=d_mask, other=0.0)

    if NORM_CORRECTION == 1:
        energy = tl.sum(tl.where(d_mask, unit * unit, 0.0), axis=0)
        unit = unit / tl.sqrt(tl.maximum(energy, 1e-24))

    n_lo = tl.load(KV_cache_ptr + val_base + VAL_BYTES).to(tl.uint16)
    n_hi = tl.load(KV_cache_ptr + val_base + VAL_BYTES + 1).to(tl.uint16)
    vec_norm = (n_lo | (n_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)

    tl.store(Out_ptr + pid * D + d_offs, (unit * vec_norm).to(tl.float16), mask=d_mask)


def unpack_values_rotated(
    kv_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
    centroids: torch.Tensor,
    *,
    num_tokens: int,
    num_heads: int,
    head_dim: int,
    key_packed_size: int,
    value_bits: int,
    norm_correction: bool = True,
) -> torch.Tensor:
    """Value vectors in rotated space, [num_tokens * num_heads, head_dim]."""
    import math

    val_bytes = math.ceil(head_dim * value_bits / 8)
    out = torch.empty(
        num_tokens * num_heads, head_dim, dtype=torch.float16, device=kv_cache.device
    )
    _zenit_unpack_values[(num_tokens * num_heads,)](
        kv_cache.view(-1),
        slot_mapping,
        centroids,
        out,
        stride_cache_block=kv_cache.stride(0),
        stride_cache_pos=kv_cache.stride(1),
        stride_cache_head=kv_cache.stride(2),
        D=head_dim,
        H=num_heads,
        BLOCK_SIZE=kv_cache.shape[1],
        BLOCK_D=triton.next_power_of_2(head_dim),
        KEY_PACKED=key_packed_size,
        VAL_BYTES=val_bytes,
        VAL_BITS=value_bits,
        NORM_CORRECTION=1 if norm_correction else 0,
        num_warps=4,
        num_stages=1,
    )
    return out


def unpack_values_model_space(
    kv_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
    centroids: torch.Tensor,
    rotation: torch.Tensor,
    **kwargs,
) -> torch.Tensor:
    """Materialising form: one batched GEMM undoes the rotation for everyone."""
    rotated = unpack_values_rotated(kv_cache, slot_mapping, centroids, **kwargs)
    return (rotated.float() @ rotation.T).to(torch.float16)


def rotate_attention_output(output: torch.Tensor, rotation: torch.Tensor) -> torch.Tensor:
    """Fused form: the only inverse rotation the decode path ever performs.

    ``output`` is whatever the split-KV reduction produced while accumulating
    against rotated values, shaped [..., head_dim].
    """
    shape = output.shape
    flat = output.reshape(-1, shape[-1]).float()
    return (flat @ rotation.T).reshape(shape).to(output.dtype)


__all__ = [
    "rotate_attention_output",
    "unpack_values_model_space",
    "unpack_values_rotated",
]