veloxquant_mlx.quantizers.kitty module

Kitty quantizer — dynamic channel-wise mixed-precision key quantization.

Inspired by “Kitty: Plug-and-Play Continuous Batching with Dynamic Token Selection” (arXiv:2511.18643, Nov 2025, unreviewed preprint). Documented as “Kitty-adapted (VeloxQuant-MLX implementation)” — not a faithful port.

Algorithm:
For a given key tensor K ∈ R^{S×D}:

  1. Compute per-channel sensitivity as the variance across the sequence dimension: σ²_j = Var(K[:, j]) for j in 0..D-1.

  2. Rank channels by σ² descending. Top hi_fraction → hi_bit quantization. Remaining channels → lo_bit quantization.

  3. Each channel group is independently quantized with asymmetric min/max group quantization (reusing _group_quant_dequant).

  4. The full key tensor is reconstructed in fp16 for downstream SDPA.

Operating in the original key space (no projection) makes Kitty zero- calibration: no SVD, no codebook training, sensitivity computed on-the-fly from the incoming K tensor itself.

Effective bit-width:
avg_bits = hi_fraction × hi_bit + (1 − hi_fraction) × lo_bit

= 0.25 × 4 + 0.75 × 2 = 2.5 bits/element (default settings)

Adaptation notes:

  • Paper uses static offline sensitivity ranking; we compute it online from each incoming key batch (prefill uses full batch variance; decode uses a running accumulator updated per token).

  • Values are left at fp16 — Kitty is a key-only method.

  • Static calibration path is NOT implemented (deferred; no calibration dataset requirement keeps the cache-only contract intact).

veloxquant_mlx.quantizers.kitty.rank_channels_by_sensitivity(keys, hi_fraction=0.25)[source]

Rank key channels by variance; return hi/lo channel index splits.

Parameters:

  • keys (array) – Shape [S, D] fp16 or fp32. S is sequence length, D is head dim.

  • hi_fraction (float) – Fraction of channels routed to high-bit quantization.

Return type:

tuple[list[int], list[int]]

Returns:

(hi_indices, lo_indices) — sorted ascending lists of column indices. hi_indices: top hi_fraction channels by variance. lo_indices: remaining channels.

veloxquant_mlx.quantizers.kitty.quantize_mixed_channels(keys, hi_indices, lo_indices, hi_bit=4, lo_bit=2, group_size=32)[source]

Apply mixed-precision quantization per channel group.

Hi channels (high variance) get hi_bit, lo channels get lo_bit. Quantization is asymmetric min/max grouped along the sequence axis.

Parameters:

  • keys (array) – [S, D] fp16 or fp32.

  • hi_indices (list[int]) – Column indices to quantize at hi_bit.

  • lo_indices (list[int]) – Column indices to quantize at lo_bit.

  • hi_bit (int) – Bit width for high-sensitivity channels.

  • lo_bit (int) – Bit width for low-sensitivity channels.

  • group_size (int) – Group size for quantization along axis 0.

Return type:

array

Returns:

Reconstructed keys [S, D] fp16.

veloxquant_mlx.quantizers.kitty.compute_running_variance(key_sum, key_sq_sum, n)[source]

Compute per-channel variance from running Welford-style accumulators.

Parameters:

  • key_sum (array) – [D] sum of all keys seen so far (fp32).

  • key_sq_sum (array) – [D] sum of squared keys seen so far (fp32).

  • n (int) – Number of tokens accumulated.

Return type:

array

Returns:

Per-channel variance [D] fp32. Returns zeros if n < 2.

veloxquant_mlx.quantizers.kitty.hi_mask_from_variance_batched(variance, hi_fraction)[source]

Batched-leading-axis channel ranking: variance [N, D] -> hi-bit mask [N, D] bool.

Replaces the per-row .tolist() + Python sorted() ranking in rank_channels_by_sensitivity() / the decode branch of KittyKVCache._quantize_keys with one host-sync-free mx.argsort call across all N rows (e.g. flattened B*H) at once. mx.argsort breaks ties the same way Python’s stable sorted() does (first index wins), so the selected channel set matches exactly.

Parameters:

  • variance (array) – [N, D] fp32 per-row, per-channel variance.

  • hi_fraction (float) – Fraction of channels routed to high-bit quantization.

Return type:

array

Returns:

[N, D] bool mask, True where the channel is in that row’s top-hi_fraction (by variance, descending; ties broken by lower channel index, matching Python’s stable sort).

veloxquant_mlx.quantizers.kitty.quantize_mixed_channels_batched(keys, hi_mask, hi_bit=4, lo_bit=2, group_size=32)[source]

Batched-leading-axis equivalent of quantize_mixed_channels().

Quantizes the full [N, S, D] tensor at hi_bit and at lo_bit independently (each a single batched group-quant call, see _group_quant_dequant_batched()), then selects per-row, per-channel between the two reconstructions with hi_mask — avoiding a ragged gather/scatter over per-row index lists of different composition. Verified bit-for-bit equivalent to looping quantize_mixed_channels() over axis 0 with the corresponding hi_indices/lo_indices.

Parameters:

  • keys (array) – [N, S, D] fp16 or fp32.

  • hi_mask (array) – [N, D] bool, True where that row’s channel is hi-bit (as returned by hi_mask_from_variance_batched()).

  • hi_bit (int) – Bit width for high-sensitivity channels.

  • lo_bit (int) – Bit width for low-sensitivity channels.

  • group_size (int) – Group size for quantization along the sequence axis.

Return type:

array

Returns:

Reconstructed keys [N, S, D] fp16.