veloxquant_mlx.quantizers.nsnquant module

NSNQuant quantizer — calibration-free universal-codebook vector quantization.

Inspired by “NSNQuant: A Double Normalization Approach for Calibration-Free Low-Bit Vector Quantization of KV Cache” (Donghyun Son, Euntae Choi, Sungjoo Yoo — NeurIPS 2025; arXiv:2505.18231, OpenReview id boNYskaXnO). Documented as “NSNQuant-adapted (VeloxQuant-MLX implementation)” — not a faithful port.

How it differs from the repo’s other VQ methods: every existing VQ method either fits its codebook to the data at hand (RVQ’s per-sequence k-means, CommVQ) or uses a data-independent geometric code (RaBitQ sign codes, VecInfer binary, PolarQuant polar grids, QJL sketches). NSNQuant inverts the relationship — it reshapes the data to match a fixed code: a Normalize-Shift-Normalize (NSN) transform plus a Hadamard rotation maps K/V token vectors onto (approximately) the standard normal distribution, so one codebook built offline from synthetic Gaussian samples — never from model activations — quantizes any model, layer, or dataset. Calibration-free by construction.

Core mechanism (per chunk of tokens):

  1. Normalize (token-wise): scale each token to norm sqrt(d); keep the scale s1.

  2. Shift (channel-wise): subtract the chunk’s own channel mean o (online statistics — never global, never calibrated).

  3. Normalize again (token-wise): rescale to norm sqrt(d); keep s2. (The second normalization slightly perturbs the zero mean; the paper shows the deviation is negligible.)

  4. Hadamard transform: rotate so channels decorrelate and the empirical distribution closely matches an isotropic standard normal.

  5. Vector quantization: 8-dim subvectors matched by cosine against a universal codebook. 2-bit: uint8 sign mask + uint8 index into a positive-orthant “magnitude” codebook (2 bits/element). 1-bit: uint8 index into a “signed” codebook only (1 bit/element).

  6. Restoration: x_hat = s1 * (s2 * x_nsn + o) after inverse Hadamard.

What we do NOT implement (see paper/NEW_METHOD_SURVEY_V11.md for the full rationale):

  • Pre-RoPE key handling + the paper’s RoPE-aware attention kernel — our cache wrappers receive post-RoPE keys from update_and_fetch, so NSN + Hadamard run post-RoPE on keys. This is the central simplification of the adaptation.

  • Value-side Hadamard fused into the projection layers (model surgery) — we apply the Hadamard explicitly to cached values instead.

  • Gradient fine-tuning of the codebook — deterministic seeded spherical k-means on synthetic standard-normal samples only.

  • 4-bit double quantization of the NSN metadata — s1/s2/o are stored fp16 and counted honestly in the byte accounting.

veloxquant_mlx.quantizers.nsnquant.nsn_transform(x)[source]

Apply the Normalize-Shift-Normalize transform to a chunk of tokens.

Parameters:

x (array) – (..., T, d) array — a self-contained chunk. The channel mean o is computed from this chunk only (online statistics).

Return type:

tuple[array, array, array, array]

Returns:

(x_nsn, s1, o, s2) where x_nsn is fp32 (..., T, d) with per-token norm sqrt(d), s1/s2 are fp16 (..., T, 1) per-token scales, and o is the fp16 (..., 1, d) per-chunk channel mean.

veloxquant_mlx.quantizers.nsnquant.nsn_inverse(x_nsn, s1, o, s2)[source]

Restore a chunk from its NSN form: x_hat = s1 * (s2 * x_nsn + o).

Exact (to fp16 metadata precision) when x_nsn is unquantized.

Return type:

array

Parameters:

  • x_nsn (mlx.core.array)

  • s1 (mlx.core.array)

  • o (mlx.core.array)

  • s2 (mlx.core.array)

veloxquant_mlx.quantizers.nsnquant.build_universal_codebook(codebook_size=256, subvector_dim=8, seed=1234, n_samples=262144, iters=25, kind='signed')[source]

Build the model-independent codebook via seeded spherical k-means.

Samples are synthetic standard-normal vectors (np.random.default_rng) — never model activations, which is what keeps NSNQuant calibration-free. Deterministic: identical arguments yield bitwise identical output. The paper additionally fine-tunes its codebook with gradient descent on a cosine objective; we skip that step (adaptation decision — expect a slightly worse codebook than the paper’s).

Parameters:

  • codebook_size (int) – Number of centroids (256 -> uint8 indices).

  • subvector_dim (int) – Dimension of each subvector (8 in the paper).

  • seed (int) – RNG seed.

  • n_samples (int) – Synthetic training samples.

  • iters (int) – Spherical k-means iterations.

  • kind (str) – "signed" (1-bit: k-means on raw samples) or "magnitude" (2-bit: k-means on |samples| — the codebook lives in the positive orthant and a per-subvector sign mask restores orientation).

Return type:

ndarray

Returns:

Float32 array of shape (codebook_size, subvector_dim) with unit-norm rows.

veloxquant_mlx.quantizers.nsnquant.vq_encode(x_nsn, codebook, bits)[source]

Encode NSN-space tokens against the universal codebook.

Parameters:

  • x_nsn (array) – (..., T, d) fp32 array (NSN + Hadamard domain).

  • codebook (ndarray) – (K, sub_d) unit-norm float32 codebook — kind="magnitude" for bits=2, kind="signed" for bits=1 (see build_universal_codebook()).

  • bits (int) – 2 (sign mask + index) or 1 (index only).

Return type:

dict

Returns:

Dict with idx (uint8 (..., T, d/sub_d)), signs (uint8 bitmask, 2-bit only), and the shape metadata needed by vq_decode().

veloxquant_mlx.quantizers.nsnquant.vq_decode(encoded, codebook)[source]

Decode vq_encode() output back to NSN-space tokens.

Codebook entries are unit-subvector approximations, so after lookup each token is renormalized to norm sqrt(d) — the scale-adjustment analog that keeps the stored s2 valid on restoration.

Return type:

array

Parameters:
veloxquant_mlx.quantizers.nsnquant.hadamard_forward(x)[source]

Plain (sign-free) normalized Hadamard rotation on the last axis.

mx.hadamard_transform is normalized (self-inverse) and norm-preserving, which is what lets the Hadamard run after NSN without disturbing the s1/o/s2 semantics: token norms are unchanged, so restoration in the rotated domain stays exact. Equivalent to HadamardPreconditioner with an all-ones diagonal — NSNQuant uses the paper’s practical plain-Hadamard choice, not the randomized-sign variant.

Return type:

array

Parameters:

x (mlx.core.array)

veloxquant_mlx.quantizers.nsnquant.hadamard_inverse(y)[source]

Inverse of hadamard_forward() (the transform is self-inverse).

Return type:

array

Parameters:

y (mlx.core.array)