mixle.experimental.kv_cache_quant module

I2: KV-cache quantization + E2 tails – int8/fp8 KV for inference, with E2’s cluster structure supplying “quantized exact outliers + G4 parametric tails” for the far-field bank’s own outlier bookkeeping.

What this is. Two related but separate quantization seams, both built on top of already-existing mechanisms rather than inventing a new quantizer:

  1. quantize_kv_cache() / dequantize_kv_cache() – ordinary affine int8 (or native fp8_e4m3) quantization of E1’s exact near-field KV window (SlidingWindowState.cache_k / cache_v, or any (..., d_head) K/V tensor at inference time). This is the literal “int8/fp8 KV for inference” half of the roadmap card – a standard per-tensor affine round-trip, not the sorted-profile machinery. It is scoped to the near-field window because that is what “the KV cache” means operationally: the thing every attention step reads on every token.

  2. quantize_cluster_outliers() – G4’s sorted-profile quantizer (mixle.models.sorted_profile_quantizer) applied to the E2 ClusterBank’s own outlier/tail bookkeeping. E2 already separates, per cluster, per chunk (birth_and_merge’s receipt["per_cluster_outlier_tokens"]): tokens whose residual against the cluster’s Gaussian-affine fit was largest (the outlier_top_k highest-residual tokens per cluster, currently a plain dense fp32 tensor – E2’s own module docstring calls this out as the “I2/G4 storage seam”, see E2_UNAVAILABLE_PIECES["I2/G4"] in moment_closure_attention.py). This module closes that seam: those flagged outlier tokens are int8-quantized (“quantized exact outliers” – exact in the sense of being carved out and identified individually, not exact in the sense of full float32 precision), while the surrounding non-outlier K/V population of the same chunk goes through G4’s fit_sorted_profile (head-exact top-k + parametric Gaussian tail fit, its own KS-receipt-gated dense fallback) – the “G4 parametric tails” half of the card.

Both halves reuse existing machinery on purpose: (1) is deliberately NOT routed through G4 (a KS-fit-gated parametric quantizer is the wrong tool for “quantize this window on every single token” – it is a per-tensor batch operation with real fitting cost, appropriate for the once-per-chunk ClusterBank outlier snapshot in (2), not for a per-step cache write), and (2) is deliberately NOT a new int8 scheme – it calls quantize_kv_cache() for the outlier half and mixle.models.sorted_profile_quantizer.fit_sorted_profile verbatim for the tail half, so there is exactly one int8 implementation and exactly one parametric-tail implementation in this codebase, both reused rather than duplicated.

Honest scope. fp8 support here is gated on torch.float8_e4m3fn (available on this environment’s torch 2.12 build, CPU-only – no fp8 hardware acceleration is claimed or exercised, this is a numerical round-trip test of the dtype’s representable grid, not a throughput benchmark). No custom Triton/CUDA kernels are written; this module is receipts-and-correctness scoped, not a speed optimization.

class AffineQuantized(codes, scale, mode)[source]

Bases: object

Round-tripped quantized tensor: quantized codes plus the (per-tensor) scale needed to dequantize.

Parameters:
codes

torch.int8 (int8 mode) or torch.float8_e4m3fn (fp8 mode) tensor, same shape as the input.

Type:

Any

scale

For int8, max(|x|) / 127 – dequant = codes.float() * scale. For fp8, always 1.0 (fp8’s own exponent field already spans the input’s dynamic range for the K/V magnitudes this module targets; see quantize_kv_cache()’s docstring for the honest caveat about very large-magnitude tensors).

Type:

float

mode

Which quantization scheme produced this.

Type:

Literal[‘int8’, ‘fp8’]

class QuantizedClusterOutliers(cluster_id, outlier_k, outlier_v, outlier_indices, tail_k, tail_v)[source]

Bases: object

Storage format for one birth_and_merge chunk’s per-cluster outlier tokens (E2’s “I2/G4 storage seam”, see moment_closure_attention.E2_UNAVAILABLE_PIECES["I2/G4"]): the flagged outlier tokens’ K/V get quantize_kv_cache()’d (“quantized exact outliers” – exact positions, quantized values); the surrounding non-outlier chunk population gets G4’s fit_sorted_profile() (“G4 parametric tails”).

Parameters:

  • cluster_id (int)

  • outlier_k (AffineQuantized | None)

  • outlier_v (AffineQuantized | None)

  • outlier_indices (ndarray | None)

  • tail_k (SortedProfileEncoding | None)

  • tail_v (SortedProfileEncoding | None)

cluster_id

Which live cluster slot this chunk’s outliers/tail came from.

Type:

int

outlier_k

Quantized exact K values of the flagged outlier tokens.

Type:

AffineQuantized | None

outlier_v

Quantized exact V values of the flagged outlier tokens.

Type:

AffineQuantized | None

outlier_indices

Flat (batch*time) token indices the outliers came from.

Type:

np.ndarray | None

tail_k

G4 parametric-tail encoding of the non-outlier K population.

Type:

SortedProfileEncoding | None

tail_v

G4 parametric-tail encoding of the non-outlier V population.

Type:

SortedProfileEncoding | None

quantize_kv_cache(x, *, mode='int8')[source]

Quantize a K or V tensor (any shape, real-valued) to int8 or fp8 for inference-time KV-cache storage.

int8: symmetric per-tensor affine quantization, scale = max(|x|) / 127, codes = round(x / scale) clamped to [-127, 127]. Per-tensor (not per-channel/per-head) scale is the deliberately simple baseline this module ships; a per-head scale would shrink error further at the cost of n_head extra floats stored per cache write; the perplexity receipt in mixle/tests/kv_cache_quant_test.py reports the per-tensor baseline honestly rather than tuning against a stronger scheme this module does not implement.

fp8: a direct cast to torch.float8_e4m3fn (4 exponent bits, 3 mantissa bits) and back – no scale computation needed since fp8’s floating exponent already tracks the input’s dynamic range (unlike int8’s fixed-point grid). Requires torch.float8_e4m3fn (torch >= 2.1); raises if unavailable rather than silently falling back to int8.

Returns:

AffineQuantized

Parameters:
Return type:

AffineQuantized

dequantize_kv_cache(q)[source]

Inverse of quantize_kv_cache(): returns a float32 tensor, same shape as the original input.

Parameters:

q (AffineQuantized)

Return type:

Any

quantize_cluster_outliers(per_cluster_outlier_tokens, flat_k, flat_v, *, mode='int8', tail_family=None, tail_top_k=0)[source]

Close E2’s I2/G4 storage seam for one birth_and_merge chunk.

Parameters:

  • per_cluster_outlier_tokens (dict) – birth_and_merge’s receipt["per_cluster_outlier_tokens"] – {cluster_id: {"k": (n_out, n_head, d_head), "v": ..., "indices": (n_out,)}}.

  • flat_k (Any) – The full chunk’s (b*t, n_head, d_head) K/V tensors (the same tensors birth_and_merge computed flat_k/flat_v from) – used to build the non-outlier tail population per cluster (every token NOT in that cluster’s indices).

  • flat_v (Any) – The full chunk’s (b*t, n_head, d_head) K/V tensors (the same tensors birth_and_merge computed flat_k/flat_v from) – used to build the non-outlier tail population per cluster (every token NOT in that cluster’s indices).

  • mode (Literal['int8', 'fp8']) – Quantization mode for the outlier half (see quantize_kv_cache()).

  • tail_family (Any) – Passed through to fit_sorted_profile() for the tail half (default GaussianEstimator()).

  • tail_top_k (int) – Head-exact top-k within the tail fit itself (default 0 – the “head-exact” carve-out is already handled by this function’s own outlier/tail split, so G4’s internal top-k defaults off to avoid double-carving the same outliers twice).

Returns:

dict[int, QuantizedClusterOutliers], keyed by cluster id.

Return type:

dict[int, QuantizedClusterOutliers]

dequantize_cluster_outliers(q)[source]

Inverse of one cluster’s QuantizedClusterOutliers: returns {"outlier_k": tensor|None, "outlier_v": tensor|None, "tail_k": np.ndarray|None, "tail_v": np.ndarray|None}.

Parameters:

q (QuantizedClusterOutliers)

Return type:

dict[str, Any]

quantization_error_per_token(x, *, mode='int8')[source]

Per-token (leading-axis) mean absolute quantize/dequantize round-trip error of x under quantize_kv_cache() – x: (n_tokens, ...), returns (n_tokens,).

Used by the receipt-correlation acceptance test (roadmap I2, “receipt correlation inside E2”) to ask whether E2’s own per-cluster misfit signal lines up with where naive KV quantization error is largest.

Parameters:
Return type:

Any