From 169c2b187cee7b8f0f9c1f262c368d83cfc1d8cb Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Tue, 15 Sep 2026 20:06:58 +0000 Subject: [PATCH 1/6] Refactor KV cache entries as dataclasses with quantized metadata --- sdm/cache.py | 61 +++++++++++++++++++++++++++++----- sdm/models/kumo/tabular/icl.py | 25 ++++++-------- test/test_cache.py | 52 ++++++++++++++++++++++++++++- 3 files changed, 114 insertions(+), 24 deletions(-) diff --git a/sdm/cache.py b/sdm/cache.py index 59c6c0db0..38bc33333 100644 --- a/sdm/cache.py +++ b/sdm/cache.py @@ -9,8 +9,9 @@ Mapping, MutableMapping, ) +from dataclasses import dataclass, replace from enum import StrEnum -from typing import NamedTuple, Self, cast +from typing import Self, cast import torch from torch import Tensor @@ -19,12 +20,8 @@ from sdm.tensor.mixin import DeviceMixin -class _KVCacheEntry(NamedTuple): - key: Tensor - value: Tensor - - -class KVCacheEntry(_KVCacheEntry, DeviceMixin): +@dataclass +class KVCacheEntry(DeviceMixin): r"""Cached key/value projections for a single transformer block. Args: @@ -32,12 +29,60 @@ class KVCacheEntry(_KVCacheEntry, DeviceMixin): value: Cached value projection tensor. """ + key: Tensor + value: Tensor + + def select_heads(self, num_heads: int) -> Self: + """Retain the leading key/value heads for grouped-query attention.""" + return replace( + self, + key=self.key[..., :num_heads, :].clone(), + value=self.value[..., :num_heads, :].clone(), + ) + def _tensors(self) -> Iterator[Tensor]: yield self.key yield self.value def _apply_tensor(self, fn: Callable[[Tensor], Tensor]) -> Self: - return self.__class__(key=fn(self.key), value=fn(self.value)) + return replace(self, key=fn(self.key), value=fn(self.value)) + + +@dataclass +class QuantizedKVCacheEntry(KVCacheEntry): + """FP8 attention projections and their per-head dequantization scales. + + Keys and values have shape ``[..., rows, heads, channels]``. Scales have + shape ``[..., 1, heads, 1]``. The query scale is learned from context + queries and reused for prediction. ``dtype`` is the attention output dtype. + """ + + key_scale: Tensor + value_scale: Tensor + query_scale: Tensor + dtype: torch.dtype + + def select_heads(self, num_heads: int) -> Self: + """Retain leading K/V heads and scales, keeping all query scales.""" + return replace( + super().select_heads(num_heads), + key_scale=self.key_scale[..., :num_heads, :].clone(), + value_scale=self.value_scale[..., :num_heads, :].clone(), + ) + + def _tensors(self) -> Iterator[Tensor]: + yield from super()._tensors() + yield self.key_scale + yield self.value_scale + yield self.query_scale + + def _apply_tensor(self, fn: Callable[[Tensor], Tensor]) -> Self: + return replace( + super()._apply_tensor(fn), + key_scale=fn(self.key_scale), + value_scale=fn(self.value_scale), + query_scale=fn(self.query_scale), + ) class Cache(MutableMapping[Hashable, object], DeviceMixin): diff --git a/sdm/models/kumo/tabular/icl.py b/sdm/models/kumo/tabular/icl.py index efb75ff3f..37f9c79c5 100644 --- a/sdm/models/kumo/tabular/icl.py +++ b/sdm/models/kumo/tabular/icl.py @@ -97,21 +97,19 @@ def forward( ) if cache is not None and cache.is_recording: - x, (key, value) = result - kv_heads = self.kv_heads - if kv_heads is not None and kv_heads < key.size(-2): - key, value = _copy_key_value( - key=key[..., :kv_heads, :], - value=value[..., :kv_heads, :], - ) - cache[cache_key] = KVCacheEntry(key, value) - del key, value + x, entry = result + cache[cache_key] = ( + entry.select_heads(self.kv_heads) + if self.kv_heads is not None + else entry + ) + del entry else: x = result del result continue - x_context, (key, value) = layer( + x_context, entry = layer( query=x[..., :0, :] if last_layer else x[..., :R_train, :], key_value=x[..., :R_train, :], return_key_value=True, @@ -121,11 +119,8 @@ def forward( if last_layer else x[..., :R_train, :], ) - key_value = KVCacheEntry( - key=key[..., : self.kv_heads, :].contiguous(), - value=value[..., : self.kv_heads, :].contiguous(), - ) - del key, value + key_value = entry.select_heads(self.kv_heads) + del entry x_query = layer( query=x[..., R_train:, :], key_value=key_value, diff --git a/test/test_cache.py b/test/test_cache.py index d05102f09..d33a6bafb 100644 --- a/test/test_cache.py +++ b/test/test_cache.py @@ -7,7 +7,7 @@ import pytest import torch -from sdm.cache import Cache, KVCacheEntry +from sdm.cache import Cache, KVCacheEntry, QuantizedKVCacheEntry from sdm.testing import onlyCUDA, withCUDA @@ -118,3 +118,53 @@ def test_cache_serialization_without_cuda( torch.testing.assert_close( restored["tensor"], cache["tensor"], rtol=0, atol=0 ) + + +@withCUDA +@pytest.mark.parametrize("quantized", [False, True]) +def test_cache_head_selection_and_transfer( + device: torch.device, quantized: bool +) -> None: + key = torch.arange(192).reshape(2, 3, 4, 8).float() + entry = ( + QuantizedKVCacheEntry( + key=key.to(torch.float8_e4m3fn), + value=(key / 2).to(torch.float8_e4m3fn), + key_scale=torch.arange(8).reshape(2, 1, 4, 1).float() + 1, + value_scale=torch.ones(2, 1, 4, 1), + query_scale=torch.ones(2, 1, 8, 1), + dtype=torch.float32, + ) + if quantized + else KVCacheEntry(key, key / 2) + ) + selected = entry.select_heads(2) + restored = selected.to(device).cpu() + assert isinstance(restored, KVCacheEntry) + assert entry.key.size(-2) == 4 + for name in ("key", "value"): + actual = getattr(restored, name) + expected = getattr(entry, name)[..., :2, :] + torch.testing.assert_close(actual.float(), expected.float()) + assert actual.dtype == expected.dtype + assert ( + actual.untyped_storage().nbytes() + == actual.numel() * actual.element_size() + ) + expected_bytes = restored.key.numel() * restored.key.element_size() * 2 + if isinstance(entry, QuantizedKVCacheEntry): + assert isinstance(restored, QuantizedKVCacheEntry) + assert restored.dtype == entry.dtype + torch.testing.assert_close( + restored.key_scale, entry.key_scale[..., :2, :] + ) + torch.testing.assert_close( + restored.value_scale, entry.value_scale[..., :2, :] + ) + torch.testing.assert_close(restored.query_scale, entry.query_scale) + expected_bytes += ( + restored.key_scale.numel() + + restored.value_scale.numel() + + restored.query_scale.numel() + ) * 4 + assert Cache(entry=restored).size() == expected_bytes From cf98373f3f34162835777b0c8be185b8c9d6d72e Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Tue, 15 Sep 2026 20:23:51 +0000 Subject: [PATCH 2/6] Describe quantized cache metadata independently of storage format --- sdm/cache.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/sdm/cache.py b/sdm/cache.py index 38bc33333..67a021ec0 100644 --- a/sdm/cache.py +++ b/sdm/cache.py @@ -50,11 +50,12 @@ def _apply_tensor(self, fn: Callable[[Tensor], Tensor]) -> Self: @dataclass class QuantizedKVCacheEntry(KVCacheEntry): - """FP8 attention projections and their per-head dequantization scales. + """Quantized attention projections with per-head scaling metadata. Keys and values have shape ``[..., rows, heads, channels]``. Scales have shape ``[..., 1, heads, 1]``. The query scale is learned from context - queries and reused for prediction. ``dtype`` is the attention output dtype. + queries and reused for prediction. The key/value tensor dtypes identify + their storage formats; ``dtype`` is the attention output dtype. """ key_scale: Tensor From 070c605d51ef9bca505bffdf4395f84916b3ed6e Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Tue, 15 Sep 2026 20:38:49 +0000 Subject: [PATCH 3/6] Preserve contiguous head-selection behavior in cache entries --- sdm/cache.py | 8 ++++---- test/test_cache.py | 13 +++++++++---- 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/sdm/cache.py b/sdm/cache.py index 67a021ec0..cca32cb0e 100644 --- a/sdm/cache.py +++ b/sdm/cache.py @@ -36,8 +36,8 @@ def select_heads(self, num_heads: int) -> Self: """Retain the leading key/value heads for grouped-query attention.""" return replace( self, - key=self.key[..., :num_heads, :].clone(), - value=self.value[..., :num_heads, :].clone(), + key=self.key[..., :num_heads, :].contiguous(), + value=self.value[..., :num_heads, :].contiguous(), ) def _tensors(self) -> Iterator[Tensor]: @@ -67,8 +67,8 @@ def select_heads(self, num_heads: int) -> Self: """Retain leading K/V heads and scales, keeping all query scales.""" return replace( super().select_heads(num_heads), - key_scale=self.key_scale[..., :num_heads, :].clone(), - value_scale=self.value_scale[..., :num_heads, :].clone(), + key_scale=self.key_scale[..., :num_heads, :].contiguous(), + value_scale=self.value_scale[..., :num_heads, :].contiguous(), ) def _tensors(self) -> Iterator[Tensor]: diff --git a/test/test_cache.py b/test/test_cache.py index d33a6bafb..34f809ea0 100644 --- a/test/test_cache.py +++ b/test/test_cache.py @@ -147,10 +147,7 @@ def test_cache_head_selection_and_transfer( expected = getattr(entry, name)[..., :2, :] torch.testing.assert_close(actual.float(), expected.float()) assert actual.dtype == expected.dtype - assert ( - actual.untyped_storage().nbytes() - == actual.numel() * actual.element_size() - ) + assert actual.is_contiguous() expected_bytes = restored.key.numel() * restored.key.element_size() * 2 if isinstance(entry, QuantizedKVCacheEntry): assert isinstance(restored, QuantizedKVCacheEntry) @@ -168,3 +165,11 @@ def test_cache_head_selection_and_transfer( + restored.query_scale.numel() ) * 4 assert Cache(entry=restored).size() == expected_bytes + + +def test_cache_head_selection_preserves_contiguous_storage() -> None: + entry = KVCacheEntry(torch.ones(1, 1, 4, 8), torch.zeros(1, 1, 4, 8)) + selected = entry.select_heads(2) + assert selected.key.shape == selected.value.shape == (1, 1, 2, 8) + assert selected.key.data_ptr() == entry.key.data_ptr() + assert selected.value.data_ptr() == entry.value.data_ptr() From 308bfec01056ae3d8417901888619f7f18a645fa Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Tue, 15 Sep 2026 20:51:24 +0000 Subject: [PATCH 4/6] Simplify quantized cache tests --- test/test_cache.py | 62 +++++++++++++++++----------------------------- 1 file changed, 23 insertions(+), 39 deletions(-) diff --git a/test/test_cache.py b/test/test_cache.py index 34f809ea0..8733ce37c 100644 --- a/test/test_cache.py +++ b/test/test_cache.py @@ -121,50 +121,34 @@ def test_cache_serialization_without_cuda( @withCUDA -@pytest.mark.parametrize("quantized", [False, True]) -def test_cache_head_selection_and_transfer( - device: torch.device, quantized: bool +def test_quantized_cache_head_selection_and_transfer( + device: torch.device, ) -> None: - key = torch.arange(192).reshape(2, 3, 4, 8).float() - entry = ( - QuantizedKVCacheEntry( - key=key.to(torch.float8_e4m3fn), - value=(key / 2).to(torch.float8_e4m3fn), - key_scale=torch.arange(8).reshape(2, 1, 4, 1).float() + 1, - value_scale=torch.ones(2, 1, 4, 1), - query_scale=torch.ones(2, 1, 8, 1), - dtype=torch.float32, - ) - if quantized - else KVCacheEntry(key, key / 2) + key = torch.arange(192).reshape(2, 3, 4, 8).to(torch.float8_e4m3fn) + scale = torch.arange(8).reshape(2, 1, 4, 1).float() + 1 + entry = QuantizedKVCacheEntry( + key=key, + value=key, + key_scale=scale, + value_scale=scale, + query_scale=torch.ones(2, 1, 8, 1), + dtype=torch.float32, ) - selected = entry.select_heads(2) - restored = selected.to(device).cpu() - assert isinstance(restored, KVCacheEntry) + selected = entry.select_heads(2).to(device) + assert isinstance(selected, QuantizedKVCacheEntry) + assert selected.dtype == entry.dtype assert entry.key.size(-2) == 4 - for name in ("key", "value"): - actual = getattr(restored, name) - expected = getattr(entry, name)[..., :2, :] - torch.testing.assert_close(actual.float(), expected.float()) + for name in ("key", "value", "key_scale", "value_scale", "query_scale"): + expected = getattr(entry, name) + if name != "query_scale": + expected = expected[..., :2, :] + actual = getattr(selected, name) + assert actual.device == device assert actual.dtype == expected.dtype assert actual.is_contiguous() - expected_bytes = restored.key.numel() * restored.key.element_size() * 2 - if isinstance(entry, QuantizedKVCacheEntry): - assert isinstance(restored, QuantizedKVCacheEntry) - assert restored.dtype == entry.dtype - torch.testing.assert_close( - restored.key_scale, entry.key_scale[..., :2, :] - ) - torch.testing.assert_close( - restored.value_scale, entry.value_scale[..., :2, :] - ) - torch.testing.assert_close(restored.query_scale, entry.query_scale) - expected_bytes += ( - restored.key_scale.numel() - + restored.value_scale.numel() - + restored.query_scale.numel() - ) * 4 - assert Cache(entry=restored).size() == expected_bytes + torch.testing.assert_close(actual.cpu().float(), expected.float()) + # 192 K/V bytes + 32 K/V scale bytes + 64 query scale bytes. + assert Cache(entry=selected).size() == 288 def test_cache_head_selection_preserves_contiguous_storage() -> None: From 3ca491113abe51bb46073cff63f0764ae67979ba Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Wed, 16 Sep 2026 03:41:44 +0000 Subject: [PATCH 5/6] Document cache entry metadata and head selection --- sdm/cache.py | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/sdm/cache.py b/sdm/cache.py index cca32cb0e..90befd5a4 100644 --- a/sdm/cache.py +++ b/sdm/cache.py @@ -33,7 +33,14 @@ class KVCacheEntry(DeviceMixin): value: Tensor def select_heads(self, num_heads: int) -> Self: - """Retain the leading key/value heads for grouped-query attention.""" + """Retain the leading key/value heads for grouped-query attention. + + Args: + num_heads: Number of leading K/V heads to retain. + + Returns: + An entry with contiguous K/V slices, reusing storage when possible. + """ return replace( self, key=self.key[..., :num_heads, :].contiguous(), @@ -56,6 +63,14 @@ class QuantizedKVCacheEntry(KVCacheEntry): shape ``[..., 1, heads, 1]``. The query scale is learned from context queries and reused for prediction. The key/value tensor dtypes identify their storage formats; ``dtype`` is the attention output dtype. + + Args: + key: Quantized key projections. + value: Quantized value projections. + key_scale: Per-head key dequantization scales. + value_scale: Per-head value dequantization scales. + query_scale: Per-query-head scales derived from the context. + dtype: Attention output dtype. """ key_scale: Tensor @@ -64,7 +79,14 @@ class QuantizedKVCacheEntry(KVCacheEntry): dtype: torch.dtype def select_heads(self, num_heads: int) -> Self: - """Retain leading K/V heads and scales, keeping all query scales.""" + """Retain leading K/V heads and scales, keeping all query scales. + + Args: + num_heads: Number of leading K/V heads to retain. + + Returns: + A quantized entry with selected K/V heads and matching scales. + """ return replace( super().select_heads(num_heads), key_scale=self.key_scale[..., :num_heads, :].contiguous(), From aaf473534ea173d87ba3220afa823459a94e15c2 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Sat, 24 Oct 2026 15:43:20 +0000 Subject: [PATCH 6/6] Preserve compiled cache-copy behavior after rebasing on main --- sdm/cache.py | 25 ++++++++++++++++++++----- sdm/models/kumo/tabular/icl.py | 13 ------------- test/test_cache.py | 9 ++++++--- 3 files changed, 26 insertions(+), 21 deletions(-) diff --git a/sdm/cache.py b/sdm/cache.py index 90befd5a4..a11e37571 100644 --- a/sdm/cache.py +++ b/sdm/cache.py @@ -39,13 +39,15 @@ def select_heads(self, num_heads: int) -> Self: num_heads: Number of leading K/V heads to retain. Returns: - An entry with contiguous K/V slices, reusing storage when possible. + An entry with contiguous K/V copies when reducing the head count. """ - return replace( - self, - key=self.key[..., :num_heads, :].contiguous(), - value=self.value[..., :num_heads, :].contiguous(), + if num_heads >= self.key.size(-2): + return self + key, value = _copy_key_value( + key=self.key[..., :num_heads, :], + value=self.value[..., :num_heads, :], ) + return replace(self, key=key, value=value) def _tensors(self) -> Iterator[Tensor]: yield self.key @@ -266,3 +268,16 @@ def _apply(value: object) -> object: ) out._mode = self._mode return out + + +def _copy_key_value(key: Tensor, value: Tensor) -> tuple[Tensor, Tensor]: + return ( + key.clone(memory_format=torch.contiguous_format), + value.clone(memory_format=torch.contiguous_format), + ) + + +# Keep the guard for PyTorch 2.7; 2.8+ can use the decorator directly. +if hasattr(torch.compiler, "nested_compile_region"): + # Compile cache copies separately to free each layer's full K/V tensors. + _copy_key_value = torch.compiler.nested_compile_region(_copy_key_value) diff --git a/sdm/models/kumo/tabular/icl.py b/sdm/models/kumo/tabular/icl.py index 37f9c79c5..94531d5fb 100644 --- a/sdm/models/kumo/tabular/icl.py +++ b/sdm/models/kumo/tabular/icl.py @@ -132,16 +132,3 @@ def forward( x = torch.cat([x_context, x_query], -2) return self.head(self.norm(x)) - - -def _copy_key_value(key: Tensor, value: Tensor) -> tuple[Tensor, Tensor]: - return ( - key.clone(memory_format=torch.contiguous_format), - value.clone(memory_format=torch.contiguous_format), - ) - - -# Keep the guard for PyTorch 2.7; 2.8+ can use the decorator directly. -if hasattr(torch.compiler, "nested_compile_region"): - # Compile cache copies separately to free each layer's full K/V tensors. - _copy_key_value = torch.compiler.nested_compile_region(_copy_key_value) diff --git a/test/test_cache.py b/test/test_cache.py index 8733ce37c..0edbcaff0 100644 --- a/test/test_cache.py +++ b/test/test_cache.py @@ -151,9 +151,12 @@ def test_quantized_cache_head_selection_and_transfer( assert Cache(entry=selected).size() == 288 -def test_cache_head_selection_preserves_contiguous_storage() -> None: +def test_cache_head_selection_releases_unused_heads() -> None: entry = KVCacheEntry(torch.ones(1, 1, 4, 8), torch.zeros(1, 1, 4, 8)) selected = entry.select_heads(2) assert selected.key.shape == selected.value.shape == (1, 1, 2, 8) - assert selected.key.data_ptr() == entry.key.data_ptr() - assert selected.value.data_ptr() == entry.value.data_ptr() + assert selected.key.data_ptr() != entry.key.data_ptr() + assert selected.value.data_ptr() != entry.value.data_ptr() + assert selected.key.is_contiguous() + assert selected.value.is_contiguous() + assert entry.select_heads(4) is entry