Skip to content

KV Cache API

Named KV cache storage formats.

Byte layouts for compressed KV blocks and pages.

estimate_persistent_bytes 

estimate_persistent_bytes(geometry: CacheGeometry, layout: PackedKVLayout) -> dict[str, int]

Estimate persistent cache bytes for a geometry/layout pair.

Source code in turboquant/kv/layout.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
def estimate_persistent_bytes(geometry: CacheGeometry, layout: PackedKVLayout) -> dict[str, int]:
    """Estimate persistent cache bytes for a geometry/layout pair."""
    vectors = geometry.vectors
    return {
        "dense_kv_bytes": geometry.dense_kv_bytes,
        "compressed_k_bytes": vectors * layout.k_payload_bytes_per_vector(),
        "compressed_v_bytes": vectors * layout.v_payload_bytes_per_vector(),
        "norms_bytes": vectors * layout.k_norm_bytes_per_vector(),
        "scales_bytes": vectors * (
            layout.value_groups * layout.scale_bytes
            + (layout.scale_bytes if get_k_format(layout.k_format).uses_scales else 0)
        ),
        "zeros_bytes": vectors * (
            layout.value_groups * layout.zero_bytes
            + (layout.zero_bytes if get_k_format(layout.k_format).uses_zeros else 0)
        ),
    }

Key storage codecs for cache write paths.

Grouped value quantization and packed storage.

quantize_values 

quantize_values(values: Tensor, *, bits: int, group_size: int = 32, scale_layout: ScaleLayout = 'group', allow_v3: bool = False, scale_dtype: dtype = torch.float16) -> ValueData

Quantize values with per-group or per-token affine metadata.

Source code in turboquant/kv/values.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def quantize_values(
    values: torch.Tensor,
    *,
    bits: int,
    group_size: int = 32,
    scale_layout: ScaleLayout = "group",
    allow_v3: bool = False,
    scale_dtype: torch.dtype = torch.float16,
) -> ValueData:
    """Quantize values with per-group or per-token affine metadata."""
    if bits not in (2, 3, 4, 8):
        raise ValueError(f"value bits must be one of 2, 3, 4, 8, got {bits}")
    validate_v_format(f"V{bits}", allow_v3=allow_v3)

    if values.ndim < 1:
        raise ValueError("values must have at least one dimension")
    dim = int(values.shape[-1])
    effective_group = dim if scale_layout == "token" else group_size
    if scale_layout not in ("group", "token"):
        raise ValueError(f"unknown scale layout: {scale_layout}")

    dim_pad = padded_dim(dim, effective_group)
    if dim_pad != dim:
        values_pad = F.pad(values, (0, dim_pad - dim))
    else:
        values_pad = values

    groups = dim_pad // effective_group
    grouped = values_pad.float().reshape(*values.shape[:-1], groups, effective_group)
    valid = torch.arange(dim_pad, device=values.device) < dim
    valid = valid.reshape(groups, effective_group)
    view_shape = (1,) * (grouped.ndim - 2) + valid.shape
    valid = valid.reshape(view_shape)
    grouped_for_min = grouped.masked_fill(~valid, float("inf"))
    grouped_for_max = grouped.masked_fill(~valid, float("-inf"))
    v_min = grouped_for_min.min(dim=-1, keepdim=True).values
    v_max = grouped_for_max.max(dim=-1, keepdim=True).values
    levels = (1 << bits) - 1
    scales = ((v_max - v_min) / levels).clamp(min=1e-12)
    zeros = v_min
    q = ((grouped - zeros) / scales).round().clamp(0, levels).to(torch.int64)
    q = q.masked_fill(~valid, 0)
    q_flat = q.reshape(*values.shape[:-1], dim_pad)
    data = pack_unsigned(q_flat, bits)

    return ValueData(
        data=data,
        scales=scales.squeeze(-1).to(scale_dtype),
        zeros=zeros.squeeze(-1).to(scale_dtype),
        bits=bits,
        dim=dim,
        group_size=effective_group,
    )

dequantize_values 

dequantize_values(q: ValueData, *, dtype: dtype = torch.float32) -> torch.Tensor

Reconstruct values from packed affine metadata.

Source code in turboquant/kv/values.py 
83
84
85
86
87
88
89
def dequantize_values(q: ValueData, *, dtype: torch.dtype = torch.float32) -> torch.Tensor:
    """Reconstruct values from packed affine metadata."""
    dim_pad = int(q.scales.shape[-1]) * q.group_size
    indices = unpack_unsigned(q.data, q.bits, dim_pad).float()
    grouped = indices.reshape(*indices.shape[:-1], q.scales.shape[-1], q.group_size)
    values = grouped * q.scales.float().unsqueeze(-1) + q.zeros.float().unsqueeze(-1)
    return values.reshape(*indices.shape[:-1], dim_pad)[..., : q.dim].to(dtype)

value_data_nbytes 

value_data_nbytes(q: ValueData) -> dict[str, int]

Return byte counts for a ValueData payload.

Source code in turboquant/kv/values.py 
92
93
94
95
96
97
98
def value_data_nbytes(q: ValueData) -> dict[str, int]:
    """Return byte counts for a ValueData payload."""
    return {
        "compressed_v_bytes": q.data.numel() * q.data.element_size(),
        "scales_bytes": q.scales.numel() * q.scales.element_size(),
        "zeros_bytes": q.zeros.numel() * q.zeros.element_size(),
    }

value_formula_nbytes 

value_formula_nbytes(shape: tuple[int, ...], *, bits: int, group_size: int, scale_dtype_bytes: int = 2, zero_dtype_bytes: int = 2, scale_layout: ScaleLayout = 'group') -> dict[str, int]

Compute expected value bytes for a shape without allocating tensors.

Source code in turboquant/kv/values.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
def value_formula_nbytes(
    shape: tuple[int, ...],
    *,
    bits: int,
    group_size: int,
    scale_dtype_bytes: int = 2,
    zero_dtype_bytes: int = 2,
    scale_layout: ScaleLayout = "group",
) -> dict[str, int]:
    """Compute expected value bytes for a shape without allocating tensors."""
    if not shape:
        raise ValueError("shape must have at least one dimension")
    dim = int(shape[-1])
    effective_group = dim if scale_layout == "token" else group_size
    dim_pad = padded_dim(dim, effective_group)
    outer = math.prod(shape[:-1]) if len(shape) > 1 else 1
    groups = dim_pad // effective_group
    payload = outer * ((dim_pad * bits + 7) // 8)
    meta = outer * groups
    return {
        "compressed_v_bytes": payload,
        "scales_bytes": meta * scale_dtype_bytes,
        "zeros_bytes": meta * zero_dtype_bytes,
        "padded_dim": dim_pad,
        "groups": groups,
    }

Compressed KV cache ownership and write paths.

TurboQuant compressed KV store - owns the quantized historical segment.

Design rules: - Chunks are stored in lists; concatenation is deferred (lazy flatten). - Flat cache is materialized on first read and invalidated on write. - No per-token overhead; all writes are chunk-based.

ValueQuantized dataclass 

ValueQuantized(data: Tensor, scales: Tensor, zeros: Tensor, bits: int = 2)

Quantized value cache (bit-packed).

FlatCache

Bases: NamedTuple

Flattened view of compressed KV for fast read access.

CompressedKVStore 

CompressedKVStore(head_dim: int, num_kv_heads: int, key_bits: int = 3, value_bits: int = 2, value_group_size: int = 32, device: device = None, layer_idx: int = 0)

Chunked compressed KV store with lazy flattening.

Keys are quantized via TurboQuantCompressorV2 (unbiased inner-product estimator). Values use TurboQuantCompressorMSE. Chunks are kept in lists until a flat view is requested.

Source code in turboquant/kv/store.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
def __init__(
    self,
    head_dim: int,
    num_kv_heads: int,
    key_bits: int = 3,
    value_bits: int = 2,
    value_group_size: int = 32,
    device: torch.device = None,
    layer_idx: int = 0,
):
    self.head_dim = head_dim
    self.num_kv_heads = num_kv_heads
    self.key_bits = key_bits
    self.value_bits = value_bits
    self.value_group_size = min(value_group_size, head_dim)
    self.device = device or torch.device("cuda")
    self.layer_idx = layer_idx

    self.key_compressor = TurboQuantCompressorV2(
        head_dim=head_dim,
        bits=key_bits,
        seed=42 + layer_idx * 7,
        device=str(self.device),
    )
    self.value_compressor = TurboQuantCompressorMSE(
        head_dim=head_dim,
        bits=value_bits,
        seed=42 + layer_idx * 7 + 500,
        device=str(self.device),
    )

    self._key_chunks: list[CompressedKeys] = []
    self._value_chunks: list[CompressedValues] = []
    self._chunk_lengths: list[int] = []

    self._flat: Optional[FlatCache] = None

num_tokens property 

num_tokens: int

Total number of compressed tokens.

num_chunks property 

num_chunks: int

Number of stored chunks.

append_chunk 

append_chunk(key: Tensor, value: Tensor)

Quantize and store a chunk of KV pairs.

key/value: (chunk_len, num_kv_heads, head_dim)

Source code in turboquant/kv/store.py 
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
def append_chunk(self, key: torch.Tensor, value: torch.Tensor):
    """Quantize and store a chunk of KV pairs.

    key/value: (chunk_len, num_kv_heads, head_dim)
    """
    chunk_len = key.shape[0]

    # Reshape to (1, num_kv_heads, chunk_len, head_dim) for compressor
    k = key.transpose(0, 1).unsqueeze(0)  # (1, H, T, D)
    v = value.transpose(0, 1).unsqueeze(0)

    key_compressed = self.key_compressor.compress(k)
    value_compressed = self.value_compressor.compress(v)

    self._key_chunks.append(key_compressed)
    self._value_chunks.append(value_compressed)
    self._chunk_lengths.append(chunk_len)
    self._flat = None  # Invalidate cached flat view

get_flat_cache 

get_flat_cache() -> Optional[FlatCache]

Return a flattened view of all compressed tokens. Cached until next write.

Source code in turboquant/kv/store.py 
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
def get_flat_cache(self) -> Optional[FlatCache]:
    """Return a flattened view of all compressed tokens. Cached until next write."""
    if not self._key_chunks:
        return None

    if self._flat is not None:
        return self._flat

    if len(self._key_chunks) == 1:
        flat_k = self._flatten_keys(self._key_chunks[0])
        flat_v = self._flatten_values(self._value_chunks[0])
    else:
        flat_k = self._concat_keys([self._flatten_keys(c) for c in self._key_chunks])
        flat_v = self._concat_values([self._flatten_values(c) for c in self._value_chunks])

    self._flat = FlatCache(
        key_compressed=flat_k,
        value_compressed=flat_v,
        num_tokens=self.num_tokens,
    )
    return self._flat

memory_bytes 

memory_bytes() -> int

Estimate GPU memory used by compressed data.

Source code in turboquant/kv/store.py 
264
265
266
267
268
269
270
271
272
273
274
275
def memory_bytes(self) -> int:
    """Estimate GPU memory used by compressed data."""
    total = 0
    for ck in self._key_chunks:
        total += ck.indices.nelement()  # uint8
        total += ck.norms.nelement() * 2  # fp16
        total += ck.qjl_signs.nelement()  # int8
        total += ck.residual_norms.nelement() * 2  # fp16
    for cv in self._value_chunks:
        total += cv.indices.nelement()  # uint8
        total += cv.norms.nelement() * 2  # fp16
    return total

reset 

reset()

Clear all stored chunks.

Source code in turboquant/kv/store.py 
277
278
279
280
281
282
def reset(self):
    """Clear all stored chunks."""
    self._key_chunks.clear()
    self._value_chunks.clear()
    self._chunk_lengths.clear()
    self._flat = None

TurboQuantKVCache 

TurboQuantKVCache(head_dim: int, key_bits: int = 3, value_bits: int = 2, value_group_size: int = 32, buffer_size: int = 128, device: device = None, dtype: dtype = torch.float16, layer_idx: int = 0)

KV cache using TurboQuant for keys and group quantization for values.

Drop-in replacement concept for a standard KV cache with compression.

Usage: cache = TurboQuantKVCache(head_dim=128, key_bits=3, value_bits=2)

# During prefill:
cache.prefill(key_states, value_states)

# During decode (one token at a time):
cache.append(new_key, new_value)

# Compute attention:
scores = cache.attention_scores(query_states)
output = cache.attend(query_states, scores_after_softmax)
Source code in turboquant/kv/store.py 
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
def __init__(
    self,
    head_dim: int,
    key_bits: int = 3,
    value_bits: int = 2,
    value_group_size: int = 32,
    buffer_size: int = 128,
    device: torch.device = None,
    dtype: torch.dtype = torch.float16,
    layer_idx: int = 0,
):
    self.head_dim = head_dim
    self.key_bits = key_bits
    self.value_bits = value_bits
    self.value_group_size = value_group_size
    self.buffer_size = buffer_size
    self.device = device or torch.device("cuda")
    self.dtype = dtype
    self.layer_idx = layer_idx

    self.key_compressor = TurboQuantCompressorV2(
        head_dim=head_dim,
        bits=key_bits,
        seed=42 + layer_idx * 7,
        device=str(self.device),
    )
    self.value_compressor = TurboQuantCompressorMSE(
        head_dim=head_dim,
        bits=value_bits,
        seed=42 + layer_idx * 7 + 500,
        device=str(self.device),
    )

    # State
    self.seq_len: int = 0
    self.key_quantized: Optional[CompressedKeys] = None
    self.value_quantized: Optional[CompressedValues] = None

    # Buffer for recent unquantized tokens
    self.key_buffer: Optional[torch.Tensor] = None
    self.value_buffer: Optional[torch.Tensor] = None

prefill 

prefill(keys: Tensor, values: Tensor)

Process prefill tokens.

Args: keys: (batch, n_heads, seq_len, head_dim) values: (batch, n_heads, seq_len, head_dim)

Source code in turboquant/kv/store.py 
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
def prefill(self, keys: torch.Tensor, values: torch.Tensor):
    """Process prefill tokens.

    Args:
        keys: (batch, n_heads, seq_len, head_dim)
        values: (batch, n_heads, seq_len, head_dim)
    """
    seq_len = keys.shape[-2]
    self.seq_len = seq_len

    if seq_len <= self.buffer_size:
        # Everything fits in buffer
        self.key_buffer = keys
        self.value_buffer = values
        return

    # Split into quantized portion and buffer
    n_quant = seq_len - self.buffer_size

    keys_to_quant = keys[..., :n_quant, :]
    values_to_quant = values[..., :n_quant, :]

    self.key_buffer = keys[..., n_quant:, :]
    self.value_buffer = values[..., n_quant:, :]

    # Quantize
    self.key_quantized = self.key_compressor.compress(keys_to_quant)
    self.value_quantized = self.value_compressor.compress(values_to_quant)

append 

append(key: Tensor, value: Tensor)

Append a single decode token.

Args: key: (batch, n_heads, 1, head_dim) value: (batch, n_heads, 1, head_dim)

Source code in turboquant/kv/store.py 
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
def append(self, key: torch.Tensor, value: torch.Tensor):
    """Append a single decode token.

    Args:
        key: (batch, n_heads, 1, head_dim)
        value: (batch, n_heads, 1, head_dim)
    """
    self.seq_len += 1

    if self.key_buffer is not None:
        self.key_buffer = torch.cat([self.key_buffer, key], dim=-2)
        self.value_buffer = torch.cat([self.value_buffer, value], dim=-2)
    else:
        self.key_buffer = key
        self.value_buffer = value

    # If buffer exceeds size, flush oldest chunk
    if self.key_buffer.shape[-2] > self.buffer_size:
        self._flush_buffer()

attention_scores 

attention_scores(query: Tensor, scale: float = None) -> torch.Tensor

Compute attention logits using asymmetric estimator.

Args: query: (batch, n_heads, n_q, head_dim) scale: attention scale factor (default: 1/sqrt(head_dim))

Returns: scores: (batch, n_heads, n_q, seq_len)

Source code in turboquant/kv/store.py 
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
def attention_scores(self, query: torch.Tensor, scale: float = None) -> torch.Tensor:
    """Compute attention logits using asymmetric estimator.

    Args:
        query: (batch, n_heads, n_q, head_dim)
        scale: attention scale factor (default: 1/sqrt(head_dim))

    Returns:
        scores: (batch, n_heads, n_q, seq_len)
    """
    import math
    if scale is None:
        scale = 1.0 / math.sqrt(self.head_dim)

    scores_parts = []

    # Quantized portion - use asymmetric estimator
    if self.key_quantized is not None:
        scores_quant = self.key_compressor.asymmetric_attention_scores(query, self.key_quantized)
        scores_parts.append(scores_quant * scale)

    # Buffer portion (full precision)
    if self.key_buffer is not None:
        scores_buf = torch.matmul(query, self.key_buffer.transpose(-2, -1))
        scores_parts.append(scores_buf * scale)

    return torch.cat(scores_parts, dim=-1)

attend 

attend(attn_weights: Tensor) -> torch.Tensor

Compute attention output: out = softmax(scores) @ values.

Args: attn_weights: (batch, n_heads, n_q, seq_len) - already softmaxed

Returns: output: (batch, n_heads, n_q, head_dim)

Source code in turboquant/kv/store.py 
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
def attend(self, attn_weights: torch.Tensor) -> torch.Tensor:
    """Compute attention output: out = softmax(scores) @ values.

    Args:
        attn_weights: (batch, n_heads, n_q, seq_len) - already softmaxed

    Returns:
        output: (batch, n_heads, n_q, head_dim)
    """
    output_parts = []
    col_offset = 0

    # Quantized values
    if self.value_quantized is not None:
        n_quant = self.value_quantized.indices.shape[-2]
        w_quant = attn_weights[..., col_offset:col_offset + n_quant]
        v_dequant = self.value_compressor.decompress(self.value_quantized)
        output_parts.append(torch.matmul(w_quant, v_dequant))
        col_offset += n_quant

    # Buffer values (full precision)
    if self.value_buffer is not None:
        n_buf = self.value_buffer.shape[-2]
        w_buf = attn_weights[..., col_offset:col_offset + n_buf]
        output_parts.append(torch.matmul(w_buf, self.value_buffer))

    return sum(output_parts)

memory_bytes 

memory_bytes() -> dict

Estimate memory usage of the cache.

Source code in turboquant/kv/store.py 
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
def memory_bytes(self) -> dict:
    """Estimate memory usage of the cache."""
    info = {"quantized_keys": 0, "quantized_values": 0, "buffer": 0, "total": 0}

    if self.key_quantized is not None:
        info["quantized_keys"] += self.key_quantized.indices.nelement()
        info["quantized_keys"] += self.key_quantized.qjl_signs.nelement()
        info["quantized_keys"] += self.key_quantized.residual_norms.nelement() * 2
        info["quantized_keys"] += self.key_quantized.norms.nelement() * 2

    if self.value_quantized is not None:
        info["quantized_values"] += self.value_quantized.indices.nelement()
        info["quantized_values"] += self.value_quantized.norms.nelement() * 2

    if self.key_buffer is not None:
        info["buffer"] += self.key_buffer.nelement() * 2  # fp16
    if self.value_buffer is not None:
        info["buffer"] += self.value_buffer.nelement() * 2

    info["total"] = info["quantized_keys"] + info["quantized_values"] + info["buffer"]
    return info

get_seq_length 

get_seq_length() -> int

Return current sequence length.

Source code in turboquant/kv/store.py 
505
506
507
def get_seq_length(self) -> int:
    """Return current sequence length."""
    return self.seq_len

reset 

reset()

Clear all state.

Source code in turboquant/kv/store.py 
509
510
511
512
513
514
515
def reset(self):
    """Clear all state."""
    self.seq_len = 0
    self.key_quantized = None
    self.value_quantized = None
    self.key_buffer = None
    self.value_buffer = None

unpack_value_data 

unpack_value_data(vq: ValueQuantized) -> torch.Tensor

Unpack bit-packed value data to uint8 per-element.

Source code in turboquant/kv/store.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def unpack_value_data(vq: ValueQuantized) -> torch.Tensor:
    """Unpack bit-packed value data to uint8 per-element."""
    bits = vq.bits
    packed = vq.data
    if bits == 2:
        v0 = packed & 0x03
        v1 = (packed >> 2) & 0x03
        v2 = (packed >> 4) & 0x03
        v3 = (packed >> 6) & 0x03
        return torch.stack([v0, v1, v2, v3], dim=-1).reshape(*packed.shape[:-1], packed.shape[-1] * 4)
    elif bits == 4:
        v0 = packed & 0x0F
        v1 = (packed >> 4) & 0x0F
        return torch.stack([v0, v1], dim=-1).reshape(*packed.shape[:-1], packed.shape[-1] * 2)
    return packed

quantize_value_chunk 

quantize_value_chunk(v: Tensor, bits: int = 2, group_size: int = 32) -> ValueQuantized

Symmetric group quantization for value vectors.

Args: v: (..., seq_len, d) value vectors bits: quantization bits (2 or 4) group_size: number of elements per quantization group

Source code in turboquant/kv/store.py 
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def quantize_value_chunk(
    v: torch.Tensor,
    bits: int = 2,
    group_size: int = 32,
) -> ValueQuantized:
    """Symmetric group quantization for value vectors.

    Args:
        v: (..., seq_len, d) value vectors
        bits: quantization bits (2 or 4)
        group_size: number of elements per quantization group
    """
    orig_shape = v.shape
    d = orig_shape[-1]
    n_groups = d // group_size
    if d % group_size != 0:
        # Pad to group_size
        pad_size = group_size - (d % group_size)
        v = torch.nn.functional.pad(v, (0, pad_size))
        d = v.shape[-1]
        n_groups = d // group_size

    # Reshape to groups
    v_grouped = v.reshape(*orig_shape[:-1], n_groups, group_size)

    # Compute scale and zero per group (asymmetric)
    v_min = v_grouped.min(dim=-1, keepdim=True).values
    v_max = v_grouped.max(dim=-1, keepdim=True).values

    n_levels = 2**bits - 1
    scale = (v_max - v_min) / n_levels
    scale = scale.clamp(min=1e-10)
    zero = v_min

    # Quantize
    v_q = ((v_grouped - zero) / scale).round().clamp(0, n_levels).to(torch.uint8)
    v_q_flat = v_q.reshape(*orig_shape[:-1], d)

    # Bit-pack
    if bits == 2:
        assert d % 4 == 0
        v_4 = v_q_flat.reshape(*orig_shape[:-1], d // 4, 4)
        packed = v_4[..., 0] | (v_4[..., 1] << 2) | (v_4[..., 2] << 4) | (v_4[..., 3] << 6)
        v_q_flat = packed
    elif bits == 4:
        assert d % 2 == 0
        v_2 = v_q_flat.reshape(*orig_shape[:-1], d // 2, 2)
        packed = v_2[..., 0] | (v_2[..., 1] << 4)
        v_q_flat = packed

    return ValueQuantized(
        data=v_q_flat,
        scales=scale.squeeze(-1),
        zeros=zero.squeeze(-1),
        bits=bits,
    )

dequantize_value_chunk 

dequantize_value_chunk(vq: ValueQuantized, group_size: int = 32) -> torch.Tensor

Dequantize value vectors from bit-packed format.

Source code in turboquant/kv/store.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
def dequantize_value_chunk(
    vq: ValueQuantized,
    group_size: int = 32,
) -> torch.Tensor:
    """Dequantize value vectors from bit-packed format."""
    data = unpack_value_data(vq).float()
    d = data.shape[-1]
    batch_shape = data.shape[:-1]

    n_groups = d // group_size
    data = data.reshape(*batch_shape, n_groups, group_size)
    scales = vq.scales.unsqueeze(-1)
    zeros = vq.zeros.unsqueeze(-1)

    v = data * scales + zeros
    return v.reshape(*batch_shape, d)

Production compressors for TurboQuant KV cache.

This module provides the main compressor implementations: - TurboQuantCompressorV2: Full compressor with asymmetric attention score estimation - TurboQuantCompressorMSE: Simpler MSE-only compressor for values - MSECompressor: Single-stage MSE-optimal compressor (no QJL) - TurboQuantV3: Community-informed compressor with residual windowing

Key insight from the TurboQuant paper: ~= + ||r_k|| * sqrt(pi/2)/m * <S@q, sign(S@r_k)>

This is unbiased with variance O(1/d), even though k_mse itself has high per-vector error. The estimator works because QJL corrects the bias in the inner product space, not in the vector space.

CompressedKeys dataclass 

CompressedKeys(indices: Tensor, norms: Tensor, qjl_signs: Tensor, residual_norms: Tensor, original_dtype: dtype = torch.float16)

Compressed key representation for asymmetric attention.

CompressedValues dataclass 

CompressedValues(indices: Tensor, norms: Tensor, original_dtype: dtype = torch.float16)

Compressed value representation.

TurboQuantCompressorV2 

TurboQuantCompressorV2(head_dim: int, bits: int, seed: int, device: str = 'cpu')

Compressor supporting direct inner product computation without full decompression.

Stores compressed representations AND supports asymmetric attention scores. This is the full TurboQuant algorithm from the paper.

Source code in turboquant/kv/compressors.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def __init__(self, head_dim: int, bits: int, seed: int, device: str = "cpu"):
    self.head_dim = head_dim
    self.bits = bits
    self.mse_bits = max(bits - 1, 1)
    self.device = device

    # Rotation matrix (Haar-distributed random orthogonal)
    gen = torch.Generator(device="cpu")
    gen.manual_seed(seed)
    G = torch.randn(head_dim, head_dim, generator=gen)
    Q, R = torch.linalg.qr(G)
    diag_sign = torch.sign(torch.diag(R))
    diag_sign[diag_sign == 0] = 1.0
    self.Pi = (Q * diag_sign.unsqueeze(0)).to(device)

    # Lloyd-Max codebook
    self.centroids = self._solve_codebook(head_dim, self.mse_bits).to(device)

    # QJL projection matrix
    gen2 = torch.Generator(device="cpu")
    gen2.manual_seed(seed + 10000)
    self.S = torch.randn(head_dim, head_dim, generator=gen2).to(device)

    # Precompute transpose for fast dequant
    self.PiT = self.Pi.T.contiguous()

compress 

compress(states: Tensor) -> CompressedKeys

Compress states: (batch, heads, seq, head_dim) -> CompressedKeys.

Source code in turboquant/kv/compressors.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
@torch.no_grad()
def compress(self, states: torch.Tensor) -> CompressedKeys:
    """Compress states: (batch, heads, seq, head_dim) -> CompressedKeys."""
    B, H, S, D = states.shape
    flat = states.reshape(-1, D).float()

    # Store original norms
    vec_norms = torch.norm(flat, dim=-1, keepdim=True)  # (N, 1)
    flat_norm = flat / (vec_norms + 1e-8)

    # Rotate and quantize
    rotated = flat_norm @ self.Pi.T
    diffs = rotated.unsqueeze(-1) - self.centroids
    indices = diffs.abs().argmin(dim=-1).to(torch.uint8)

    # MSE reconstruction in original space
    reconstructed_rotated = self.centroids[indices.long()]
    k_mse = (reconstructed_rotated @ self.Pi) * vec_norms  # (N, D)

    # Residual in original space
    residual = flat - k_mse
    residual_norm = torch.norm(residual, dim=-1)  # (N,)

    # QJL signs of residual
    projected = residual @ self.S.T
    signs = (projected >= 0).to(torch.int8) * 2 - 1  # {-1, +1}

    return CompressedKeys(
        indices=indices.reshape(B, H, S, D),
        norms=vec_norms.squeeze(-1).to(torch.float16).reshape(B, H, S),
        qjl_signs=signs.reshape(B, H, S, D),
        residual_norms=residual_norm.to(torch.float16).reshape(B, H, S),
        original_dtype=states.dtype,
    )

decompress 

decompress(compressed: CompressedKeys) -> torch.Tensor

Decompress to full tensors (MSE component only).

Source code in turboquant/kv/compressors.py 
149
150
151
152
153
154
155
156
157
@torch.no_grad()
def decompress(self, compressed: CompressedKeys) -> torch.Tensor:
    """Decompress to full tensors (MSE component only)."""
    B, H, S, D = compressed.indices.shape
    indices = compressed.indices.reshape(-1, D).long()
    norms = compressed.norms.reshape(-1, 1).float()

    reconstructed = self.centroids[indices] @ self.Pi
    return (reconstructed * norms).reshape(B, H, S, D).to(compressed.original_dtype)

asymmetric_attention_scores 

asymmetric_attention_scores(queries: Tensor, compressed: CompressedKeys) -> torch.Tensor

Compute attention scores directly from compressed K.

Uses the asymmetric estimator: ~= + ||r_k|| * sqrt(pi/2)/m * <S@q, signs_k>

Args: queries: (batch, heads, seq_q, head_dim) compressed: CompressedKeys from compress()

Returns: scores: (batch, heads, seq_q, seq_k)

Source code in turboquant/kv/compressors.py 
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
@torch.no_grad()
def asymmetric_attention_scores(self, queries: torch.Tensor, compressed: CompressedKeys) -> torch.Tensor:
    """Compute attention scores <Q, K> directly from compressed K.

    Uses the asymmetric estimator:
        <q, k> ~= <q, k_mse> + ||r_k|| * sqrt(pi/2)/m * <S@q, signs_k>

    Args:
        queries: (batch, heads, seq_q, head_dim)
        compressed: CompressedKeys from compress()

    Returns:
        scores: (batch, heads, seq_q, seq_k)
    """
    # Decompress MSE component
    k_mse = self.decompress(compressed).float()  # (B, H, S_k, D)
    signs = compressed.qjl_signs.float()          # (B, H, S_k, D)
    r_norm = compressed.residual_norms.float()    # (B, H, S_k)

    # Term 1: Q @ K_mse^T
    term1 = torch.matmul(queries.float(), k_mse.transpose(-2, -1))  # (B, H, S_q, S_k)

    # Term 2: QJL correction
    q_projected = torch.matmul(queries.float(), self.S.T)  # (B, H, S_q, D)
    qjl_ip = torch.matmul(q_projected, signs.transpose(-2, -1))  # (B, H, S_q, S_k)

    m = self.S.shape[0]
    correction_scale = math.sqrt(math.pi / 2) / m
    term2 = correction_scale * qjl_ip * r_norm.unsqueeze(-2)

    return term1 + term2

TurboQuantCompressorMSE 

TurboQuantCompressorMSE(head_dim: int, bits: int, seed: int, device: str = 'cpu')

Simpler MSE-only compressor for values (no QJL needed).

Source code in turboquant/kv/compressors.py 
195
196
197
198
199
200
201
202
203
204
205
206
207
def __init__(self, head_dim: int, bits: int, seed: int, device: str = "cpu"):
    self.head_dim = head_dim
    self.bits = bits
    self.device = device

    gen = torch.Generator(device="cpu")
    gen.manual_seed(seed)
    G = torch.randn(head_dim, head_dim, generator=gen)
    Q, R = torch.linalg.qr(G)
    diag_sign = torch.sign(torch.diag(R))
    diag_sign[diag_sign == 0] = 1.0
    self.Pi = (Q * diag_sign.unsqueeze(0)).to(device)
    self.centroids = self._solve_codebook(head_dim, bits).to(device)

compress 

compress(states: Tensor) -> CompressedValues

Compress states: (batch, heads, seq, head_dim) -> CompressedValues.

Source code in turboquant/kv/compressors.py 
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
@torch.no_grad()
def compress(self, states: torch.Tensor) -> CompressedValues:
    """Compress states: (batch, heads, seq, head_dim) -> CompressedValues."""
    B, H, S, D = states.shape
    flat = states.reshape(-1, D).float()
    vec_norms = torch.norm(flat, dim=-1, keepdim=True)
    flat_norm = flat / (vec_norms + 1e-8)
    rotated = flat_norm @ self.Pi.T
    diffs = rotated.unsqueeze(-1) - self.centroids
    indices = diffs.abs().argmin(dim=-1).to(torch.uint8)
    return CompressedValues(
        indices=indices.reshape(B, H, S, D),
        norms=vec_norms.squeeze(-1).to(torch.float16).reshape(B, H, S),
        original_dtype=states.dtype,
    )

decompress 

decompress(compressed: CompressedValues) -> torch.Tensor

Decompress back to full tensors.

Source code in turboquant/kv/compressors.py 
258
259
260
261
262
263
264
265
@torch.no_grad()
def decompress(self, compressed: CompressedValues) -> torch.Tensor:
    """Decompress back to full tensors."""
    B, H, S, D = compressed.indices.shape
    indices = compressed.indices.reshape(-1, D).long()
    norms = compressed.norms.reshape(-1, 1).float()
    reconstructed = self.centroids[indices] @ self.Pi
    return (reconstructed * norms).reshape(B, H, S, D).to(compressed.original_dtype)

MSECompressor 

MSECompressor(head_dim: int, bits: int, seed: int, device: str = 'cpu')

Single-stage MSE-optimal compressor with bit packing.

Used for both keys and values. No QJL — all bits go to reconstruction quality. This is the core building block for TurboQuantV3.

Source code in turboquant/kv/compressors.py 
275
276
277
278
279
280
281
282
283
284
285
286
287
def __init__(self, head_dim: int, bits: int, seed: int, device: str = "cpu"):
    self.head_dim = head_dim
    self.bits = bits
    self.device = device

    gen = torch.Generator(device="cpu")
    gen.manual_seed(seed)
    G = torch.randn(head_dim, head_dim, generator=gen)
    Q, R = torch.linalg.qr(G)
    diag_sign = torch.sign(torch.diag(R))
    diag_sign[diag_sign == 0] = 1.0
    self.Pi = (Q * diag_sign.unsqueeze(0)).to(device)
    self.centroids = self._solve_codebook(head_dim, bits).to(device)

compress 

compress(states: Tensor) -> dict

Compress (B, H, S, D) -> dict with bit-packed indices + norms.

Source code in turboquant/kv/compressors.py 
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
@torch.no_grad()
def compress(self, states: torch.Tensor) -> dict:
    """Compress (B, H, S, D) -> dict with bit-packed indices + norms."""
    B, H, S, D = states.shape
    N = B * H * S
    flat = states.reshape(N, D).float()

    # Normalize to unit sphere
    vec_norms = torch.norm(flat, dim=-1)  # (N,)
    flat_norm = flat / (vec_norms.unsqueeze(-1) + 1e-8)

    # Rotate + quantize
    rotated = flat_norm @ self.Pi.T
    diffs = rotated.unsqueeze(-1) - self.centroids
    indices = diffs.abs().argmin(dim=-1).to(torch.uint8)  # (N, D)

    # Bit-pack indices
    indices_per_byte = 8 // self.bits
    idx_pad = (indices_per_byte - D % indices_per_byte) % indices_per_byte
    idx_flat = indices.long()
    if idx_pad:
        idx_flat = F.pad(idx_flat, (0, idx_pad))
    n_groups = idx_flat.shape[-1] // indices_per_byte
    idx_powers = torch.tensor(
        [2 ** (self.bits * i) for i in range(indices_per_byte - 1, -1, -1)],
        dtype=torch.long, device=idx_flat.device
    )
    idx_bytes = (idx_flat.reshape(N, n_groups, indices_per_byte) * idx_powers).sum(-1).to(torch.uint8)

    return {
        "idx_bytes": idx_bytes.reshape(B, H, S, n_groups),
        "vec_norms": vec_norms.to(torch.float16).reshape(B, H, S),
        "shape": (B, H, S, D),
        "idx_pad": idx_pad,
    }

decompress 

decompress(compressed: dict) -> torch.Tensor

Decompress back to (B, H, S, D) tensor.

Source code in turboquant/kv/compressors.py 
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
@torch.no_grad()
def decompress(self, compressed: dict) -> torch.Tensor:
    """Decompress back to (B, H, S, D) tensor."""
    B, H, S, D = compressed["shape"]
    N = B * H * S
    idx_bytes = compressed["idx_bytes"].reshape(N, -1)
    vec_norms = compressed["vec_norms"].reshape(N, 1).float()
    idx_pad = compressed["idx_pad"]

    # Unpack indices
    indices_per_byte = 8 // self.bits
    mask = (1 << self.bits) - 1
    idx_shifts = torch.tensor(
        [self.bits * i for i in range(indices_per_byte - 1, -1, -1)],
        dtype=torch.long, device=idx_bytes.device
    )
    indices = ((idx_bytes.long().unsqueeze(-1) >> idx_shifts) & mask).reshape(N, -1)
    if idx_pad:
        indices = indices[:, :D]

    # Reconstruct
    reconstructed = (self.centroids[indices] @ self.Pi) * vec_norms
    return reconstructed.reshape(B, H, S, D)

memory_bytes 

memory_bytes(B: int, H: int, S: int) -> dict

Actual memory usage in bytes.

Source code in turboquant/kv/compressors.py 
355
356
357
358
359
360
361
362
363
364
365
366
367
368
def memory_bytes(self, B: int, H: int, S: int) -> dict:
    """Actual memory usage in bytes."""
    D = self.head_dim
    N = B * H * S
    indices_per_byte = 8 // self.bits
    idx_bytes = N * math.ceil(D / indices_per_byte)
    norm_bytes = N * 2  # fp16
    compressed = idx_bytes + norm_bytes
    fp16 = N * D * 2
    return {
        "compressed_bytes": compressed,
        "fp16_bytes": fp16,
        "compression_ratio": fp16 / compressed if compressed > 0 else 0,
    }

TurboQuantV3 

TurboQuantV3(head_dim: int, key_bits: int = 4, value_bits: int = 2, residual_window: int = 128, layer_idx: int = 0, n_layers: int = 36, protected_layers: int = 4, protected_bits: int = 8, seed: int = 42, device: str = 'cpu')

Community-informed KV cache compressor.

Key improvements over V2: - MSE-only: no QJL, all bits go to reconstruction quality - Asymmetric: separate bit-widths for keys vs values - Residual window: recent tokens kept in fp16 - Layer-adaptive: configurable per-layer bit overrides

Usage: compressor = TurboQuantV3(head_dim=128, key_bits=4, value_bits=2, device="cuda") compressed_k, compressed_v = compressor.compress_kv(keys, values) keys_out, values_out = compressor.decompress_kv(compressed_k, compressed_v)

Source code in turboquant/kv/compressors.py 
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
def __init__(
    self,
    head_dim: int,
    key_bits: int = 4,
    value_bits: int = 2,
    residual_window: int = 128,
    layer_idx: int = 0,
    n_layers: int = 36,
    protected_layers: int = 4,
    protected_bits: int = 8,
    seed: int = 42,
    device: str = "cpu",
):
    self.head_dim = head_dim
    self.residual_window = residual_window
    self.device = device

    # Layer-adaptive: first/last N layers get more bits
    is_protected = layer_idx < protected_layers or layer_idx >= (n_layers - protected_layers)
    effective_key_bits = protected_bits if is_protected else key_bits
    effective_value_bits = protected_bits if is_protected else value_bits

    # Cap at 8 bits
    self.key_bits = min(effective_key_bits, 8)
    self.value_bits = min(effective_value_bits, 8)

    seed_base = seed + layer_idx * 1000
    self.key_compressor = MSECompressor(head_dim, self.key_bits, seed=seed_base, device=device)
    self.val_compressor = MSECompressor(head_dim, self.value_bits, seed=seed_base + 500, device=device)

compress_kv 

compress_kv(keys: Tensor, values: Tensor) -> Tuple[dict, dict]

Compress key and value tensors.

Input: keys, values - both (B, H, S, D)

If S > residual_window, the last residual_window tokens are kept in fp16 (uncompressed) for generation quality.

Source code in turboquant/kv/compressors.py 
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
@torch.no_grad()
def compress_kv(
    self, keys: torch.Tensor, values: torch.Tensor
) -> Tuple[dict, dict]:
    """Compress key and value tensors.

    Input: keys, values - both (B, H, S, D)

    If S > residual_window, the last `residual_window` tokens are kept
    in fp16 (uncompressed) for generation quality.
    """
    B, H, S, D = keys.shape
    rw = self.residual_window

    if S <= rw:
        # Short sequence - keep everything in fp16
        return (
            {"fp16": keys, "compressed": None, "shape": (B, H, S, D), "split_at": S},
            {"fp16": values, "compressed": None, "shape": (B, H, S, D), "split_at": S},
        )

    # Split: compress old tokens, keep recent in fp16
    split_at = S - rw

    old_keys = keys[:, :, :split_at, :]
    recent_keys = keys[:, :, split_at:, :]
    old_values = values[:, :, :split_at, :]
    recent_values = values[:, :, split_at:, :]

    compressed_k = {
        "compressed": self.key_compressor.compress(old_keys),
        "fp16": recent_keys,
        "shape": (B, H, S, D),
        "split_at": split_at,
    }
    compressed_v = {
        "compressed": self.val_compressor.compress(old_values),
        "fp16": recent_values,
        "shape": (B, H, S, D),
        "split_at": split_at,
    }
    return compressed_k, compressed_v

decompress_kv 

decompress_kv(compressed_k: dict, compressed_v: dict) -> Tuple[torch.Tensor, torch.Tensor]

Decompress back to full tensors.

Source code in turboquant/kv/compressors.py 
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
@torch.no_grad()
def decompress_kv(
    self, compressed_k: dict, compressed_v: dict
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Decompress back to full tensors."""
    if compressed_k["compressed"] is None:
        return compressed_k["fp16"], compressed_v["fp16"]

    # Decompress old tokens
    old_keys = self.key_compressor.decompress(compressed_k["compressed"])
    old_values = self.val_compressor.decompress(compressed_v["compressed"])

    # Concatenate with fp16 recent tokens
    dtype = compressed_k["fp16"].dtype
    keys = torch.cat([old_keys.to(dtype), compressed_k["fp16"]], dim=2)
    values = torch.cat([old_values.to(dtype), compressed_v["fp16"]], dim=2)
    return keys, values

memory_bytes 

memory_bytes(B: int, H: int, S: int) -> dict

Report actual memory usage including residual window.

Source code in turboquant/kv/compressors.py 
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
def memory_bytes(self, B: int, H: int, S: int) -> dict:
    """Report actual memory usage including residual window."""
    rw = min(self.residual_window, S)
    compressed_S = max(S - rw, 0)
    fp16_S = rw

    if compressed_S > 0:
        k_mem = self.key_compressor.memory_bytes(B, H, compressed_S)
        v_mem = self.val_compressor.memory_bytes(B, H, compressed_S)
        compressed_bytes = k_mem["compressed_bytes"] + v_mem["compressed_bytes"]
    else:
        compressed_bytes = 0

    fp16_window_bytes = B * H * fp16_S * self.head_dim * 2 * 2  # keys + values, fp16
    total_compressed = compressed_bytes + fp16_window_bytes
    total_fp16 = B * H * S * self.head_dim * 2 * 2

    return {
        "compressed_bytes": total_compressed,
        "fp16_bytes": total_fp16,
        "compression_ratio": total_fp16 / total_compressed if total_compressed > 0 else 0,
        "compressed_tokens": compressed_S,
        "fp16_tokens": fp16_S,
        "key_bits": self.key_bits,
        "value_bits": self.value_bits,
    }

KV cache byte accounting.

ByteLedger 

ByteLedger(dense_kv_bytes: int = 0)

Mutable builder for MemoryReport.

Source code in turboquant/kv/memory.py 
100
101
def __init__(self, dense_kv_bytes: int = 0):
    self._report = MemoryReport(dense_kv_bytes=dense_kv_bytes)

bytes_from_tensors 

bytes_from_tensors(tensors: Mapping[str, Tensor]) -> dict[str, int]

Return a byte ledger for named tensors.

Source code in turboquant/kv/memory.py 
92
93
94
def bytes_from_tensors(tensors: Mapping[str, torch.Tensor]) -> dict[str, int]:
    """Return a byte ledger for named tensors."""
    return {name: tensor_nbytes(tensor) for name, tensor in tensors.items()}

report_from_components 

report_from_components(*, dense_kv_bytes: int, compressed_k: Tensor | None = None, compressed_v: Tensor | None = None, scales: Tensor | None = None, zeros: Tensor | None = None, norms: Mapping[str, Tensor] | Tensor | None = None, centroids: Mapping[str, Tensor] | Tensor | None = None, rotation: Mapping[str, Tensor] | Tensor | None = None, metadata_bytes: int = 0, recent_window: Tensor | None = None, boundary: Tensor | None = None, sparse: Tensor | None = None, outlier: Tensor | None = None, workspace: Mapping[str, Tensor] | Tensor | None = None, peak_allocated_bytes: int = 0, formula_total_bytes: int = 0) -> MemoryReport

Build a MemoryReport from actual tensors.

Source code in turboquant/kv/memory.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def report_from_components(
    *,
    dense_kv_bytes: int,
    compressed_k: torch.Tensor | None = None,
    compressed_v: torch.Tensor | None = None,
    scales: torch.Tensor | None = None,
    zeros: torch.Tensor | None = None,
    norms: Mapping[str, torch.Tensor] | torch.Tensor | None = None,
    centroids: Mapping[str, torch.Tensor] | torch.Tensor | None = None,
    rotation: Mapping[str, torch.Tensor] | torch.Tensor | None = None,
    metadata_bytes: int = 0,
    recent_window: torch.Tensor | None = None,
    boundary: torch.Tensor | None = None,
    sparse: torch.Tensor | None = None,
    outlier: torch.Tensor | None = None,
    workspace: Mapping[str, torch.Tensor] | torch.Tensor | None = None,
    peak_allocated_bytes: int = 0,
    formula_total_bytes: int = 0,
) -> MemoryReport:
    """Build a MemoryReport from actual tensors."""
    ledger = ByteLedger(dense_kv_bytes=dense_kv_bytes)
    optional = [
        ("compressed_k_bytes", compressed_k),
        ("compressed_v_bytes", compressed_v),
        ("scales_bytes", scales),
        ("zeros_bytes", zeros),
        ("recent_window_bytes", recent_window),
        ("boundary_bytes", boundary),
        ("sparse_bytes", sparse),
        ("outlier_bytes", outlier),
    ]
    for field, tensor in optional:
        if tensor is not None:
            ledger.add_tensor(field, tensor)

    for field, item in [
        ("norms_bytes", norms),
        ("centroids_bytes", centroids),
        ("rotation_bytes", rotation),
        ("workspace_bytes", workspace),
    ]:
        if isinstance(item, torch.Tensor):
            ledger.add_tensor(field, item)
        elif item is not None:
            ledger.add_tensors(field, item)

    ledger.add_bytes("metadata_bytes", metadata_bytes)
    ledger.add_bytes("peak_allocated_bytes", peak_allocated_bytes)
    ledger.add_bytes("formula_total_bytes", formula_total_bytes)
    return ledger.to_report()

Model-aware cache policy selection.

Dense recent-window storage.