[FP8 1/5] Refactor K/V cache types and add quantized entries

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,25 +20,94 @@
 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:
         key: Cached key projection tensor.
         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.
+
+        Args:
+            num_heads: Number of leading K/V heads to retain.
+
+        Returns:
+            An entry with contiguous K/V copies when reducing the head count.
+        """
+        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
         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):
+    """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. 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
+    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.
+
+        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(),
+            value_scale=self.value_scale[..., :num_heads, :].contiguous(),
+        )
+
+    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):
@@ -198,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)

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,
@@ -137,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)

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,45 @@ def test_cache_serialization_without_cuda(
     torch.testing.assert_close(
         restored["tensor"], cache["tensor"], rtol=0, atol=0
     )
+
+
+@withCUDA
+def test_quantized_cache_head_selection_and_transfer(
+    device: torch.device,
+) -> None:
+    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).to(device)
+    assert isinstance(selected, QuantizedKVCacheEntry)
+    assert selected.dtype == entry.dtype
+    assert entry.key.size(-2) == 4
+    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()
+        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_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.is_contiguous()
+    assert selected.value.is_contiguous()
+    assert entry.select_heads(4) is entry