veloxquant_mlx.quantizers.gear module¶
GEAR quantizer — error-feedback: residual low-rank + sparse outlier correction.
Inspired by “GEAR: An Efficient KV Cache Compression Recipe for Near-Lossless Generative Inference of LLM” (Kang et al., arXiv:2403.05527). Documented as “GEAR-adapted (VeloxQuant-MLX implementation)” — not a faithful port.
What GEAR adds that the repo did not have: error feedback. Every other method in the suite picks a bit-width (or a cache layout) and lives with the quantization error. GEAR makes any ultra-low-bit base quantizer near-lossless by reconstructing what it threw away, via the three-part decomposition
X ~= Quant_b(X) + L . R + S
Quant_b(X)— base: most entries at ultra-low precision. The paper’s backbone is KCVT: keys quantized per-channel (§2, “Group-wise Quantization”), values quantized per-token, both without KIVI’s fine-grained residual buffer. We reuse the repo’s asymmetric min/max group quant (cachegen.quantize_to_codes) but apply it along the axis KCVT specifies —base_axis="channel"for keys,"token"for values — so GEAR composes over an existing, already-tested quantizer with the paper’s actual backbone rather than a generic per-token quantizer for both.
L . R— a low-rank approximation of the quantization residualE = X - dequant(Quant_b(X)). The residual of a coherent KV matrix is itself low-rank, so a small rank recovers most of the lost signal cheaply.
S— a sparse matrix correcting the top-rhooutlier entries by magnitude in the post-low-rank residual, the few the low-rank term could not absorb.
- Adaptation:
The residual SVD is computed on the prefill batch (the cache wrapper owns the prefill-vs-decode orchestration), reusing the SVDq/PALU prefill-SVD pattern. The paper additionally uses a smaller rank (r=2) for streaming decode-only updates via a token buffer (§3, “Streaming Buffer”) purely as a speed optimization — the reconstruction target is unaffected, so we compute the same low-rank recovery on whatever batch the cache holds. Not ported: no perf claim depends on it.
We do not ship GEAR’s fused streaming-dequant CUDA kernel. The wrapper reconstructs fp16 then calls MLX SDPA; stored size shrinks, attend-time peak memory does not. Documented as a known simplification.
Reconstruction is a real lossy reconstruction (unlike CacheGen’s lossless byte model): the reported
compressed_*_bytesAND the values the model sees both reflect the error-feedback layers.
This module holds the pure numerics: base quant (borrowed, KCVT-axis-aware), residual, residual-low-rank, sparse-outlier extraction, full compress/reconstruct, and an honest byte estimator. The cache wrapper owns the per-layer prefill/decode state.
- class veloxquant_mlx.quantizers.gear.GEARState(codes, scale, zero, L, R, sp_idx, sp_val, n_rows, bits, rank, axis='token', d_cols=0)[source]¶
Bases:
NamedTupleA GEAR-compressed tensor: base codes + low-rank residual + sparse outliers.
- Variables:
codes – [n_groups, group_size, D] fp32 integer base codes.
scale – [n_groups, 1, D] fp32 per-group step.
zero – [n_groups, 1, D] fp32 per-group min.
L – [N, r] fp32 low-rank residual left factor (or None if rank==0).
R – [r, D] fp32 low-rank residual right factor (or None if rank==0).
sp_idx – [nnz] int32 flattened indices into the [N, D] residual (or None).
sp_val – [nnz] fp16 outlier residual values (or None).
n_rows – int original (pre-pad) token count (the
X.shape[0]GEAR was given — NOT necessarily the row count ofcodes, seeaxis).bits – int base bit-width.
rank – int residual low-rank (0 = no low-rank term).
axis – str base quant axis — “channel” (KCVT key scheme) or “token” (KCVT value scheme).
codes/scale/zeroare transposed (row count ==d_cols) whenaxis == "token".d_cols – int original channel count
D— the row count ofcodeswhenaxis == "token"(unused, equalscodes.shape[-1], whenaxis == "channel").
- Parameters:
- codes: mx.array¶
Alias for field number 0
- scale: mx.array¶
Alias for field number 1
- zero: mx.array¶
Alias for field number 2
- veloxquant_mlx.quantizers.gear.quantize_base(x, bits, group_size=32, axis='token')[source]¶
Base ultra-low-bit group quant. Returns
(CodeStream, base_recon[N, D]).Borrows the repo’s shared asymmetric min/max group quant (the same math KIVI/CacheGen use) so GEAR’s base layer is not new quant logic — only the axis it groups along is chosen here, per the paper’s KCVT backbone.
- Parameters:
x (
array) –[N, D]one head’s keys or values.bits (
int) – base bit-width.group_size (
int) – group size along the quantization axis.axis (
str) –"channel"(paper’s Key scheme — group along the token axis, one(scale, zero)per channel) or"token"(paper’s Value scheme — group along the channel axis, one(scale, zero)per token; also the default, matching this module’s pre-KCVT behavior for callers that do not passaxis).
- veloxquant_mlx.quantizers.gear.residual(x, base_recon)[source]¶
Quantization residual
E = x - base_reconas fp32[N, D].- Return type:
array- Parameters:
x (mlx.core.array)
base_recon (mlx.core.array)
- veloxquant_mlx.quantizers.gear.lowrank_error(E, rank, energy_threshold=0.9)[source]¶
Low-rank factors
(L, R)of the residual via truncated SVD.E ~= L @ RwithL = U_r * s_r([N, r]) andR = Vt_r([r, D]). Returns(None, None)whenrank == 0(pure-sparse / base-only mode).- Parameters:
E (mx.array) – Residual
[N, D]fp32.rank (int | None) – Explicit low-rank
r; 0 disables the low-rank term; None → chosen byenergy_threshold.energy_threshold (float) – Singular-value energy to retain when
rankis None.
- Return type:
tuple[mx.array | None, mx.array | None]
- Returns:
(L, R)fp32, or(None, None)if no low-rank term.
- veloxquant_mlx.quantizers.gear.sparse_outliers(resid, frac)[source]¶
Top-
fracentries ofresidby magnitude →(flat_idx, values).The outlier correction GEAR applies to the residual the low-rank term could not absorb. Selection is by absolute value over the flattened
[N, D].- Parameters:
resid (mx.array) – Post-low-rank residual
[N, D]fp32.frac (float) – Fraction of entries to keep exact (0 → no sparse term).
- Return type:
tuple[mx.array | None, mx.array | None]
- Returns:
(flat_idx[int32], values[fp16])or(None, None)if no sparse term.
- veloxquant_mlx.quantizers.gear.sparse_outliers_batched(resid, frac)[source]¶
Batched-leading-axis equivalent of
sparse_outliers()(#570).Every row of a
_compress_and_accountbatch shares the same[N, D]shape (sameS/Dfor the whole call), sonnz— unlike SVD rank — is already uniform across the batch; no ragged-width padding/truncation is needed here (contrast_truncated_svd_batched(), whose per-row rank genuinely varies and must be handled that way).- Parameters:
resid (mx.array) – Post-low-rank residual
[BH, N, D]fp32.frac (float) – Fraction of entries to keep exact (0 → no sparse term).
- Return type:
tuple[mx.array | None, mx.array | None]
- Returns:
(flat_idx[int32], values[fp16])each[BH, nnz], or(None, None)if no sparse term — matchingsparse_outliers()’ per-row contract, just with a leadingBHaxis.
- veloxquant_mlx.quantizers.gear.gear_compress(x, bits=2, rank=None, sparse_frac=0.01, group_size=32, energy_threshold=0.9, base_axis='token')[source]¶
Full GEAR compression of one head’s K/V matrix
[N, D]→GEARState.Pipeline: base group quant (KCVT-axis-aware) → residual → low-rank residual → sparse outliers of the post-low-rank residual.
- veloxquant_mlx.quantizers.gear.gear_reconstruct(state, base=None)[source]¶
Reconstruct fp16
[n_rows, D]from a GEARState (base + low-rank + sparse).- Parameters:
state (GEARState) – The compressed tensor.
base (mx.array | None) – Optional precomputed base-layer dequant (fp32,
[n_rows, D]), i.e.dequant(Quant_b(x))for this exactstate. Callers that already computed this (e.g.GEARKVCache._compress_and_account, whose Pass 1 batches the base dequant across allB*Hheads before Pass 3 calls this function once per head) should pass it through here instead of leaving it to be recomputed — this function otherwise redoes the identical per-head_dequant_axismath a second time for no reason. Recomputed fromstatewhen omitted, preserving the original behavior for standalone callers (e.g.gear_quant_dequant()).
- Return type:
mx.array
- veloxquant_mlx.quantizers.gear.gear_reconstruct_batched(base, L, R, sp_idx, sp_val, out_dtype=mlx.core.float16)[source]¶
Batched-leading-axis equivalent of
gear_reconstruct()(#570).Replaces the
for idx in range(B * H): gear_reconstruct(...)loop — the largest remaining per-head cost inGEARKVCache._compress_and_accountafter #504 batched Passes 1-2 (~50% of wall time in isolated profiling, per that class’s module docstring). The low-rank add and sparse scatter for head(b, h)only ever touch that head’s own factors/mask — no cross-head coupling — so both are pure vectorizations over the batch axis.L/Rhere are the padded, batch-max-rank factors straight out of_truncated_svd_batched(already zero-masked beyond each row’s own rank), NOT per-row-truncated toranks[idx]— unlike storage/byte accounting, the padded columns are numerically inert for this batched matmul (zero singular values contribute exactly 0 toL @ R), so skipping the per-row slice here is correctness-preserving and avoids rebuilding a ragged Python list just to re-stack it.- Parameters:
base (mx.array) – Batched base-layer dequant, fp32
[BH, N, D](Pass 1’s output — already computed once, batched, upstream).L (mx.array | None) –
[BH, N, max_rank]fp32 or None (no low-rank term for the batch).R (mx.array | None) –
[BH, max_rank, D]fp32 or None.sp_idx (mx.array | None) –
[BH, nnz]int32 flattened indices into each row’s[N, D], or None (no sparse term).sp_val (mx.array | None) –
[BH, nnz]fp16 outlier values, or None.out_dtype (mx.Dtype)
- Return type:
mx.array
- Returns:
fp16
[BH, N, D]reconstruction, matchingmx.stack([gear_reconstruct(...) for idx in range(BH)], axis=0).
- veloxquant_mlx.quantizers.gear.gear_bytes(state)[source]¶
Honest stored size (bytes) of a GEARState.
base codes (fixed-width packed) + fp16 group params + fp16
L,Rfactors + sparse triples (int32 index + fp16 value per nnz).
- veloxquant_mlx.quantizers.gear.base_only_bytes(state)[source]¶
Stored size of the base codes alone (no error-feedback) — the baseline.