Preserve fitted cache pinning across serialization

sdm/cache.py

@@ -10,10 +10,12 @@
     MutableMapping,
 )
 from enum import StrEnum
-from typing import NamedTuple, Self
+from typing import NamedTuple, Self, cast
 
+import torch
 from torch import Tensor
 
+from sdm._memory import set_pinned_max_round_threshold
 from sdm.tensor.mixin import DeviceMixin
 
 
@@ -126,6 +128,41 @@ def __len__(self) -> int:
     def __repr__(self) -> str:
         return repr(self._items)
 
+    def __getstate__(self) -> dict[str, object]:
+        state = self.__dict__.copy()
+        # Stable serialization metadata: positions in the tensor traversal.
+        state["_pinned_tensor_indices"] = tuple(
+            index
+            for index, tensor in enumerate(self._tensors())
+            if tensor.is_pinned()
+        )
+        return state
+
+    def __setstate__(self, state: dict[str, object]) -> None:
+        pinned = set(
+            cast(tuple[int, ...], state.pop("_pinned_tensor_indices", ()))
+        )
+        self.__dict__.update(state)
+        if not pinned or not torch.cuda.is_available():
+            return
+
+        replacements: dict[int, Tensor] = {}
+        for index, tensor in enumerate(self._tensors()):
+            if (
+                index in pinned
+                and tensor.is_cpu
+                and not tensor.is_pinned()
+                and id(tensor) not in replacements
+            ):
+                if not replacements:
+                    set_pinned_max_round_threshold(mb=1)
+                replacements[id(tensor)] = tensor.pin_memory()
+
+        if replacements:
+            self._items = self._apply_tensor(
+                lambda tensor: replacements.get(id(tensor), tensor)
+            )._items
+
     def _tensors(self) -> Iterator[Tensor]:
         def _iter_tensors(value: object) -> Iterator[Tensor]:
             if isinstance(value, Tensor):

test/test_cache.py

@@ -1,12 +1,14 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import io
 from typing import cast
 
+import pytest
 import torch
 
 from sdm.cache import Cache, KVCacheEntry
-from sdm.testing import withCUDA
+from sdm.testing import onlyCUDA, withCUDA
 
 
 @withCUDA
@@ -45,3 +47,74 @@ def test_cache_size() -> None:
     )
 
     assert cache.size() == 3 * 4 + 2 * 8 + 5 * 1 + 4 * 2
+
+
+@withCUDA
+@pytest.mark.parametrize("map_to_device", [False, True])
+def test_cache_serialization(
+    device: torch.device,
+    map_to_device: bool,
+) -> None:
+    nested = Cache(
+        entry=KVCacheEntry(
+            key=torch.arange(6).view(2, 3), value=torch.ones(2, 3)
+        )
+    )
+    if device.type == "cuda":
+        nested = nested.pin_memory()
+    cache = Cache(foo="foo", tensor=torch.arange(3), nested=nested).freeze()
+
+    stream = io.BytesIO()
+    torch.save(cache, stream)
+    stream.seek(0)
+    restored = torch.load(
+        stream,
+        map_location=device if map_to_device else None,
+        weights_only=False,
+    )
+
+    assert restored is not cache
+    assert restored["foo"] == "foo"
+    assert restored.is_replaying
+    restored_nested = cast(Cache, restored["nested"])
+    assert restored_nested.is_replaying
+    assert restored_nested.is_pinned() == (
+        device.type == "cuda" and not map_to_device
+    )
+    assert not cast(torch.Tensor, restored["tensor"]).is_pinned()
+
+    load_device = device if map_to_device else torch.device("cpu")
+    expected_entry = cast(KVCacheEntry, nested["entry"])
+    restored_entry = cast(KVCacheEntry, restored_nested["entry"])
+    for actual, expected in (
+        (
+            cast(torch.Tensor, restored["tensor"]),
+            cast(torch.Tensor, cache["tensor"]),
+        ),
+        (restored_entry.key, expected_entry.key),
+        (restored_entry.value, expected_entry.value),
+    ):
+        assert actual.device == load_device
+        torch.testing.assert_close(
+            actual, expected.to(load_device), rtol=0, atol=0
+        )
+
+
+@onlyCUDA
+def test_cache_serialization_without_cuda(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    cache = Cache(tensor=torch.arange(3).pin_memory())
+    stream = io.BytesIO()
+    torch.save(cache, stream)
+    stream.seek(0)
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    restored = torch.load(stream, map_location="cpu", weights_only=False)
+
+    assert restored.is_recording
+    assert restored.is_cpu
+    assert not restored.is_pinned()
+    torch.testing.assert_close(
+        restored["tensor"], cache["tensor"], rtol=0, atol=0
+    )