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). This module closes that storage 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

Positive finite multiplier used by dequant = codes.float() * scale. Int8 uses max(|x|) / 127 (or 1 for an all-zero tensor). FP8 uses 1 when the source fits natively and otherwise scales its largest magnitude to the dtype’s finite maximum before casting.

Type:

float

mode

Which quantization scheme produced this.

Type:

Literal[‘int8’, ‘fp8’]

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

Bases: object

Storage format for one birth_and_merge chunk’s per-cluster outlier tokens: the flagged 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)

  • member_indices (ndarray)

  • outlier_indices (ndarray)

  • tail_indices (ndarray)

  • 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

member_indices

Every flat token index assigned to this cluster.

Type:

numpy.ndarray

outlier_indices

The cluster members stored individually.

Type:

numpy.ndarray

tail_indices

The remaining cluster members stored in the tail encoding.

Type:

numpy.ndarray

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

class ClusterTokenConservation(input_count, member_count, outlier_count, tail_count)[source]

Bases: object

Verified token-count receipt for one clustered chunk.

Parameters:

  • input_count (int)

  • member_count (int)

  • outlier_count (int)

  • tail_count (int)

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: casts to torch.float8_e4m3fn (4 exponent bits, 3 mantissa bits). Values outside that dtype’s finite range are first scaled by max(|x|) / finfo.max so the cast cannot silently produce NaNs. 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": ..., "v": ..., "indices": ..., "member_indices": ...}}.

  • 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 (members assigned to that cluster but not in its outlier 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 (members assigned to that cluster but not in its outlier 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.

The returned dictionary includes the member, outlier, and tail token indices needed to place each reconstructed population without duplicating tokens from other clusters.

Parameters:

q (QuantizedClusterOutliers)

Return type:

dict[str, Any]

verify_cluster_token_conservation(quantized, n_tokens)[source]

Verify that cluster members partition the chunk and each member occurs in exactly one storage tier.

Parameters:

  • quantized (dict[int, QuantizedClusterOutliers])

  • n_tokens (int)

Return type:

ClusterTokenConservation

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