diff --git a/sdm/cache.py b/sdm/cache.py index 08ab744c7..59c6c0db0 100644 --- a/sdm/cache.py +++ b/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): diff --git a/test/test_cache.py b/test/test_cache.py index 36865b137..d05102f09 100644 --- a/test/test_cache.py +++ b/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 + )