tpt-kv-quant

KV-cache quantization for attention inference: the storage layout ([KvBlock]), streaming (append-time) quantization, a recency-tiered precision policy ([TieredKvPolicy]), and a GPU (wgpu) backend — built on the shared, frozen-at-1.0 math in tpt-quant-core.

Everything outside backend/wgpu.rs is meant to reach "stable and done" and stay frozen. The GPU kernels are ongoing performance work and live entirely behind the [KvBackend] trait, so they can be rewritten without changing — or bumping the version of — anything else in the crate.

• License: MIT OR Apache-2.0
• MSRV: 1.85.0 (edition 2024)
• no_std: yes (core + alloc). The only feature that pulls std and a GPU dependency is wgpu, which is default-off.

Non-goals

Written into this README on day one, and audited before every release:

• No model loading / format code. No .gguf / .safetensors parsing. The crate takes raw &[f32] and returns &[f32]. Full stop.
• No attention math beyond dequant-and-feed. dequantize_for_attention produces the dequantized K/V an attention kernel consumes; the QKᵀ / softmax / V matmul stays in the consumer.
• No CUDA / Metal / ROCm backends in v1. wgpu-only. (The trait seam leaves room for more backends later.)
• No other tiering strategies. v1 ships exactly one fixed 3-tier (recent / sink / mid / far) policy.

Quick start

use tpt_kv_quant::{
    backend::Tier, CpuRefBackend, KvBackend, KvScheme, TieredKvCache, TieredKvPolicy,
};
use tpt_quant_core::scheme::BitWidth;

// K/V shapes for this model.
const HEAD_DIM: usize = 64;
const NUM_KV_HEADS: usize = 8;
const DIM: usize = HEAD_DIM * NUM_KV_HEADS;

// Per-channel K (group_size = head_dim), per-token V (group_size = full token).
let scheme = KvScheme::new(BitWidth::Int8, BitWidth::Int8, DIM, true);

// Recent-window of 256 tokens unquantized; first 4 are pinned sinks;
// older tokens quantized at Int4 (mid) then Int2 (far).
let policy = TieredKvPolicy {
    recent_window: 256,
    sink_tokens: 4,
    mid_bits: BitWidth::Int4,
    far_bits: BitWidth::Int2,
};

let mut cache = TieredKvCache::new(CpuRefBackend, HEAD_DIM, NUM_KV_HEADS, policy, scheme, 1024);

// As tokens are produced, append them (token-major `f32`):
let k: Vec<f32> = vec![0.1; DIM];
let v: Vec<f32> = vec![-0.1; DIM];
cache.append(&k, &v).unwrap();

// Feed attention: reconstruct the whole sequence as dequantized f32.
let attn = cache.dequantize().unwrap();
assert_eq!(attn.seq_len, 1);
assert_eq!(attn.k.len(), DIM);

For a GPU backend, enable the wgpu feature and use WgpuBackend instead of CpuRefBackend:

# #[cfg(feature = "wgpu")]
# {
use tpt_kv_quant::backend::WgpuBackend;
let mut cache = TieredKvCache::new(
    WgpuBackend::new_headless(),
    HEAD_DIM,
    NUM_KV_HEADS,
    policy,
    scheme,
    1024,
);
# }

API overview

Item	Purpose
KvScheme	wraps a tpt_quant_core::QuantScheme for K and V
KvBlock	the storage format (STABLE, freeze early)
TieredKvPolicy	recency-based Recent/Sink/Mid/Far precision tiers
TieredKvCache<B>	append-time, backend-agnostic tiered quantization
KvBackend	the stability seam trait (quantize_append, dequantize_for_attention)
CpuRefBackend	the scalar no_std correctness oracle
WgpuBackend (feature wgpu)	the GPU backend that churns
AttentionInput	dequantized K/V an attention kernel consumes
differential	the differential-testing harness (cpu_ref vs anything)

The stability seam

backend/wgpu.rs is the only part of the crate expected to keep changing — different group sizes, fusion strategies, tuned workgroup sizes. Everything else (the layout, the tiering policy, the streaming append logic) stays frozen and is tested on CPU. The differential harness makes the GPU churn safe: every WGSL change is validated by running the same input through CpuRefBackend and WgpuBackend and asserting they agree within an epsilon.

Tiering model

[ sink_tokens )            -> Sink   (unquantized, pinned forever)
[ sink_tokens .. mid_start ) -> Far    (quantized at `far_bits`)
[ mid_start .. recent_start ) -> Mid   (quantized at `mid_bits`)
[ recent_start .. len )     -> Recent (unquantized)

where recent_start = len - recent_window and the Mid band width equals recent_window (clamped to the quantized region). A token only ever moves Recent → Mid → Far, so migration always re-quantizes from the true f32 source (no compounding error).

Development

cargo test --tests                       # unit + property + integration tests
cargo test --no-default-features         # confirms the no_std core builds
cargo clippy --all-targets --all-features -- -D warnings
cargo fmt --all --check
cargo miri test --no-default-features --lib   # UB check for core logic

The wgpu feature is validated in CI against CpuRefBackend via the differential harness, headless on the lavapipe software Vulkan adapter:

WGPU_ADAPTER_NAME=lavapipe cargo test --features wgpu

Status

Pre-1.0. The core (layout, tiering, streaming, CPU oracle, differential harness) is complete and tested on CPU; the wgpu backend is functional and expected to churn. See the roadmap (todo.md) for the path to a stable 1.0 and the integration work into tpt-spark, tpt-gpu-runtime, and tpt-abyss-engine.