RotorQuant KV-Cache Quantization

AI & Machine Learning Foundations advanced How AI Works ~35 min read

Overview

Long-context inference is often limited by the key-value cache rather than by model weights. Every generated token leaves key and value vectors behind for every attention layer, and the cache grows linearly with sequence length. A model that fits comfortably at short context can become memory-bound when it serves many concurrent sessions or a single very long prompt. KV-cache quantization attacks that bottleneck by storing those vectors in fewer bits while preserving the attention scores computed from them.

RotorQuant is a rotation-based KV-cache quantization technique that replaces a dense global random rotation with many small block-diagonal rotations. Its technical move is to split each attention-head vector into tiny groups, rotate each group with a compact geometric-algebra rotor, then apply scalar quantization and invert the rotation during reconstruction. The result aims to keep the useful decorrelation property of TurboQuant while removing the dense matrix multiply and most of the rotation parameters.

The important engineering lesson is not that every runtime should adopt a specific Clifford algebra implementation. The transferable pattern is: when a theoretically strong compression method depends on an expensive dense transform, look for a structured transform that preserves the geometric property you need, maps cleanly to registers, and can be fused with quantization into a single device kernel.

Core Concepts

KV-Cache Pressure

During autoregressive decoding, attention does not recompute keys and values for all prior tokens. Instead, it stores them in a cache and appends the vectors for each new token. That cache is read repeatedly by future attention steps, so its representation affects memory footprint, bandwidth, and attention-score fidelity. Quantizing weights alone does not solve this because the cache is generated dynamically at inference time.

Rotation Before Scalar Quantization

Scalar quantizers work best when each coordinate has a predictable distribution and does not carry extreme correlation with neighboring coordinates. TurboQuant uses a dense random orthogonal transform to mix vector coordinates before Lloyd-Max scalar quantization, then uses a QJL residual correction to reduce inner-product bias. RotorQuant keeps the rotate-then-quantize framing but changes the rotation operator.

Clifford Rotors

A rotor is a compact way to represent a rotation in geometric algebra. For 3D blocks, the rotor has a small number of active components and rotates a vector through a sandwich product: apply the rotor, multiply by the embedded vector, then apply the reverse rotor. This produces an orthogonal rotation without storing or multiplying by a full dense matrix.

Key Concept: Structured Decorrelation

RotorQuant trades global coordinate mixing for local block mixing. Instead of one large head-size-by-head-size rotation matrix, it uses independent rotations over small vector chunks. This reduces parameter count and arithmetic while still giving the scalar quantizer a rotated coordinate system that is easier to encode.

How It Works

Chunk the Head Vector

A typical attention head might have 128 dimensions. RotorQuant partitions that vector into groups of three dimensions, pads the final group if needed, and assigns a rotor to each group. Each group can be processed independently, which creates natural parallelism for a GPU kernel: one thread or lane handles a batch/group pair, keeps the group in registers, and writes the packed quantized result.

def rotor_quantize_head(x, rotors, codebooks):
    # x is one attention-head vector, already split into 3D groups.
    codes = []
    for group, rotor in zip(chunks(x, 3), rotors):
        mv = embed_as_multivector(group)
        rotated = sandwich(rotor, mv)       # R * v * reverse(R)
        code = lloyd_max_quantize(rotated, codebooks)
        codes.append(code)
    return pack(codes)

Exploit Sparse Products

A naive geometric product over all multivector components would waste work on zeros. The RotorQuant report shows a sparse product specialized to rotor structure: only scalar and bivector terms participate, so the implementation can use a fixed sequence of fused multiply-add operations. This is the difference between a mathematically interesting transform and one that can compete with tuned inference kernels.

def sparse_rotor_step(x, s, p12, p13, p23):
    # Representative structure: scalar plus three bivector coefficients.
    r0 = s*x[0] - p12*x[4] - p13*x[5] - p23*x[6]
    r1 = s*x[1] + p12*x[2] + p13*x[3] + p23*x[7]
    r2 = s*x[2] - p12*x[1] + p23*x[3] - p13*x[7]
    return r0, r1, r2

Fuse the Device Pipeline

The full pipeline is only efficient if embedding, rotation, quantization, inverse rotation, and extraction avoid intermediate global-memory writes. A fused CUDA, Triton, or Metal kernel can load a small rotor and codebook data, operate on registers, and write compact codes. That kernel design is why structured rotations matter: the transform is not just smaller on paper, it fits the execution model of the hardware.

Practical Applications

Long-Context Chat Serving

Chat servers with long context windows keep large per-session KV caches. Compressing the cache can increase the number of concurrent sessions that fit on a GPU, reduce cache paging, or let a smaller device serve longer prompts. Rotor-style quantizers are most attractive when cache memory and bandwidth dominate latency.

On-Device Inference

Apple Silicon, laptops, and edge GPUs have tighter memory and bandwidth budgets than data-center accelerators. A blockwise rotation with a compact parameter set is easier to keep in fast memory than a dense rotation matrix. Fusing the transform into Metal or CUDA avoids round-trips through framework-level tensor operators.

Vector Compression Research

The same design question appears in ANN search and embedding storage: can a structured transform make low-bit quantization easier without expensive training or codebook fitting? RotorQuant is a useful case study because it demonstrates a structured alternative to dense random rotations, then measures both distortion and downstream attention behavior.

Implementation Patterns

Separate Algorithm Choice from Runtime Format

Treat the cache type as a runtime capability, not as a model architecture change. The model emits normal key and value vectors; the cache layer decides whether to store them as FP16, int8, planar rotations, quaternion rotations, or rotor-quantized codes. This makes it possible to switch cache policies per deployment and to fall back safely when a backend lacks kernel support.

type KVCachePolicy struct {
    KeyFormat   string // "f16", "int8", "rotor3", "planar3"
    ValueFormat string
    GroupSize   int
    Bits        int
    Backend     string // "cuda", "metal", "cpu"
}

func (p KVCachePolicy) Validate() error {
    if p.Bits < 2 || p.Bits > 8 {
        return fmt.Errorf("unsupported cache bit width")
    }
    if p.GroupSize != 3 && p.GroupSize != 4 {
        return fmt.Errorf("unsupported rotation group size")
    }
    return nil
}

Benchmark Attention Fidelity, Not Just MSE

Reconstruction error is useful, but the application cares about attention scores and generated output quality. A KV quantizer can have worse raw mean-squared error yet preserve nearest keys or attention rankings well enough for inference. Measure cosine similarity against the original attention output, top-k key retrieval, perplexity, and long-context task quality.

Keep a Backend-Specific Kernel Boundary

Prototype the math in Python or a tensor framework, but move the hot path behind a backend boundary early. The useful production kernel is the fused kernel, not a chain of tiny operations launched from a high-level framework. A clean boundary also lets CUDA, Metal, and CPU fallbacks share the same cache-policy API while using different implementations.

Performance Considerations

Common Pitfalls

Further Reading