feat(cache): persist GDN checkpoints as SSD sidecars

omlx/cache/type_registry.py

@@ -145,6 +145,13 @@ def is_rotating_family(cls, class_name: str) -> bool:
             CacheType.BATCH_ROTATING_KVCACHE,
         )
 
+    @classmethod
+    def is_arrays_family(cls, class_name: str) -> bool:
+        """Check whether a class name belongs to the ArraysCache family."""
+        if class_name == "SizedArraysCache":
+            return True
+        return cls._class_name_map.get(class_name) == CacheType.ARRAYS_CACHE
+
     @classmethod
     def detect_cache_type(cls, cache_obj: Any) -> CacheType:
         """Detect cache type from object.

omlx/config.py

@@ -110,6 +110,8 @@ class PagedSSDCacheConfig:
     cache_dir: Optional[Path] = None
     max_size: str = "100GB"
     hot_cache_max_size: str = "0"  # "0" = disabled, e.g. "8GB"
+    gdn_ssd_split_enabled: bool = False
+    gdn_ssd_pending_max_size: str = "512MB"
 
     @property
     def max_size_bytes(self) -> int:
@@ -180,6 +182,13 @@ def from_env(cls) -> "OMLXConfig":
 
         # Paged SSD cache settings
         config.paged_ssd_cache.hot_cache_only = os.getenv("OMLX_HOT_CACHE_ONLY", "false").lower() == "true"
+        config.paged_ssd_cache.gdn_ssd_split_enabled = os.getenv(
+            "OMLX_GDN_SSD_SPLIT_ENABLED", "false"
+        ).lower() in ("true", "1", "yes")
+        config.paged_ssd_cache.gdn_ssd_pending_max_size = os.getenv(
+            "OMLX_GDN_SSD_PENDING_MAX_SIZE",
+            config.paged_ssd_cache.gdn_ssd_pending_max_size,
+        )
         paged_ssd_dir = os.getenv("OMLX_PAGED_SSD_CACHE_DIR")
         if paged_ssd_dir:
             config.paged_ssd_cache.enabled = True
@@ -297,5 +306,18 @@ def validate(self) -> List[str]:
         if self.paged_ssd_cache.enabled:
             if not self.paged_ssd_cache.cache_dir:
                 errors.append("Paged SSD cache enabled but no cache_dir specified")
+        if (
+            self.paged_ssd_cache.gdn_ssd_split_enabled
+            and self.paged_ssd_cache.hot_cache_only
+        ):
+            errors.append(
+                "gdn_ssd_split_enabled cannot be used with hot_cache_only"
+            )
+        try:
+            pending_size = parse_size(self.paged_ssd_cache.gdn_ssd_pending_max_size)
+            if pending_size <= 0:
+                errors.append("gdn_ssd_pending_max_size must be positive")
+        except (AttributeError, TypeError, ValueError) as exc:
+            errors.append(f"Invalid gdn_ssd_pending_max_size: {exc}")
 
         return errors

omlx/scheduler.py

@@ -22,7 +22,7 @@
 from array import array
 from collections import OrderedDict, defaultdict, deque
 from collections.abc import Callable
-from contextlib import contextmanager
+from contextlib import contextmanager, suppress
 from dataclasses import dataclass, field
 from enum import Enum
 from pathlib import Path
@@ -1433,6 +1433,10 @@ class SchedulerConfig:
     paged_ssd_cache_max_size: int = 100 * 1024 * 1024 * 1024  # 100GB default
     hot_cache_max_size: int = 0  # In-memory hot cache size in bytes (0 = disabled)
     hot_cache_budget: Any | None = None  # Shared process-wide hot cache budget
+    # Store top-level ArraysCache recurrent state as SSD sidecars while the
+    # ordinary block retains only KV/sliceable payloads.
+    gdn_ssd_split_enabled: bool = False
+    gdn_ssd_pending_max_bytes: int = 512 * 1024 * 1024
 
     # Model identification (for cache isolation between different models)
     model_name: str = ""  # OpenAI API model name (e.g., "mlx-community/Llama-3.2-3B")
@@ -1481,11 +1485,13 @@ def __init__(
         request_id: str,
         valid_tcs: list[int],
         in_memory_snapshots: dict[int, Any],
+        paged_ssd_manager: Any | None = None,
     ) -> None:
         self._store = store
         self._request_id = request_id
         self._valid_tcs = set(valid_tcs)
         self._in_memory = in_memory_snapshots
+        self._paged_ssd_manager = paged_ssd_manager
 
     def __contains__(self, tc: int) -> bool:
         return tc in self._valid_tcs
@@ -1511,6 +1517,46 @@ def iter_in_memory_extracted(self):
             if snap is not None:
                 yield snap
 
+    def commit_gdn_checkpoint(
+        self,
+        token_count: int,
+        source_block_hash: bytes,
+        *,
+        layer_cache_types: list[str] | None,
+        layer_meta_states: list[Any] | None,
+        model_name: str,
+        block_size: int,
+    ) -> bool:
+        """Promote one request-local boundary snapshot to durable sidecar storage."""
+        if self._store is None or self._paged_ssd_manager is None:
+            return False
+        staged_path = self._store.take_staged_file(self._request_id, token_count)
+        if staged_path is None:
+            return False
+        try:
+            signature = self._paged_ssd_manager.cache_signature_for(
+                model_name=model_name,
+                num_layers=len(layer_cache_types or []),
+                block_size=block_size,
+                layer_cache_types=layer_cache_types or [],
+            )
+            committed = self._paged_ssd_manager.commit_gdn_checkpoint_file(
+                source_block_hash,
+                staged_path,
+                token_count=token_count,
+                model_name=model_name,
+                cache_signature=signature,
+                block_size=block_size,
+            )
+            if committed is None:
+                with suppress(OSError):
+                    staged_path.unlink()
+            return committed is not None
+        except Exception:
+            with suppress(OSError):
+                staged_path.unlink()
+            return False
+
 
 class Scheduler:
     """
@@ -1898,6 +1944,7 @@ def __init__(
             self.block_aware_cache = BlockAwarePrefixCache(
                 model=model,
                 paged_cache_manager=self.paged_cache_manager,
+                gdn_ssd_split_enabled=self.config.gdn_ssd_split_enabled,
             )
 
             # Initialize paged SSD cache. If the backing directory is not
@@ -6063,6 +6110,7 @@ def _get_boundary_store_override(
             request_id=request_id,
             valid_tcs=provider_tcs,
             in_memory_snapshots=extracted_in_memory,
+            paged_ssd_manager=self.paged_ssd_cache_manager,
         )
 
         token_sequence = (
@@ -8221,6 +8269,25 @@ def _do_abort_request(self, request_id: str) -> bool:
         if request is None:
             return False
 
+        # A finished request remains in self.requests while its async
+        # store_cache worker owns boundary snapshots and cache buffers. Do not
+        # run abort cleanup concurrently with that worker: the normal deferred
+        # drain will release the batch row, snapshots, and Request once the
+        # future completes. If it completed between steps, drain it now before
+        # deciding whether there is anything left to abort.
+        store_future = self._inflight_store_futures.get(request_id)
+        if store_future is not None:
+            if not store_future.done():
+                logger.debug(
+                    "Deferring abort cleanup for %s until async store_cache completes",
+                    request_id,
+                )
+                return False
+            self._drain_pending_async_removes()
+            request = self.requests.get(request_id)
+            if request is None:
+                return False
+
         self._clear_request_admission_bookkeeping(request_id)
 
         # Remove from waiting queue
@@ -11081,6 +11148,29 @@ def get_cache_stats(self) -> dict[str, Any] | None:
 
     def reset(self) -> None:
         """Reset the scheduler state."""
+        # A store_cache worker may still be loading request-local boundary
+        # snapshots or publishing blocks. reset() clears both namespaces, so
+        # use the same bounded teardown barrier as shutdown() before aborting
+        # requests or clearing caches. The drain performs the request-local
+        # cleanup only after every future has completed.
+        inflight = list(self._inflight_store_futures.values())
+        if inflight:
+            logger.info(
+                "Waiting for %d inflight async store_cache future(s) before reset...",
+                len(inflight),
+            )
+            _done, not_done = concurrent.futures.wait(
+                inflight, timeout=FATAL_TEARDOWN_TIMEOUT_S
+            )
+            if not_done:
+                fatal_exit(
+                    "Scheduler reset timed out after "
+                    f"{FATAL_TEARDOWN_TIMEOUT_S:.0f}s waiting for "
+                    f"{len(not_done)} async store_cache future(s)"
+                )
+                return
+            self._drain_pending_async_removes()
+
         # Drain any pending deferred aborts
         self._pending_abort_ids.clear()
 
@@ -11685,6 +11775,7 @@ def _init_tiered_cache(self) -> bool:
                 hot_cache_max_bytes=self.config.hot_cache_max_size,
                 hot_cache_only=self.config.hot_cache_only,
                 hot_cache_budget=self.config.hot_cache_budget,
+                gdn_ssd_split_enabled=self.config.gdn_ssd_split_enabled,
                 expected_model_name=self.config.model_name or "",
                 expected_num_layers=expected_num_layers,
                 expected_block_size=self.config.paged_cache_block_size,
@@ -11710,8 +11801,13 @@ def _init_tiered_cache(self) -> bool:
             if BoundarySnapshotSSDStore is not None and not self.config.hot_cache_only:
                 try:
                     self._boundary_snapshot_store = BoundarySnapshotSSDStore(
-                        base_dir=Path(self.config.paged_ssd_cache_dir)
+                        base_dir=Path(self.config.paged_ssd_cache_dir),
+                        pending_max_bytes=self.config.gdn_ssd_pending_max_bytes,
                     )
+                    if self.block_aware_cache is not None:
+                        self.block_aware_cache.set_gdn_checkpoint_loader(
+                            self._boundary_snapshot_store.load_file
+                        )
                 except Exception as e:
                     logger.debug(
                         "Failed to initialize boundary snapshot SSD store: %s", e

omlx/settings.py

@@ -312,6 +312,8 @@ class CacheSettings:
     ssd_cache_max_size: str = "auto"  # "auto" means 10% of SSD capacity
     hot_cache_max_size: str = "0"  # "0" = disabled, e.g. "8GB"
     initial_cache_blocks: int = 256  # Starting blocks (grows dynamically)
+    gdn_ssd_split_enabled: bool = False
+    gdn_ssd_pending_max_size: str = "512MB"
 
     def get_ssd_cache_dir(self, base_path: Path) -> Path:
         """
@@ -351,6 +353,8 @@ def to_dict(self) -> dict[str, Any]:
         return {
             "enabled": self.enabled,
             "hot_cache_only": self.hot_cache_only,
+            "gdn_ssd_split_enabled": self.gdn_ssd_split_enabled,
+            "gdn_ssd_pending_max_size": self.gdn_ssd_pending_max_size,
             "ssd_cache_dir": self.ssd_cache_dir,
             "ssd_cache_max_size": self.ssd_cache_max_size,
             "hot_cache_max_size": self.hot_cache_max_size,
@@ -367,6 +371,10 @@ def from_dict(cls, data: dict[str, Any]) -> CacheSettings:
         return cls(
             enabled=data.get("enabled", True),
             hot_cache_only=data.get("hot_cache_only", False),
+            gdn_ssd_split_enabled=data.get("gdn_ssd_split_enabled", False),
+            gdn_ssd_pending_max_size=data.get(
+                "gdn_ssd_pending_max_size", "512MB"
+            ),
             ssd_cache_dir=data.get("ssd_cache_dir"),
             ssd_cache_max_size=data.get("ssd_cache_max_size", "auto"),
             hot_cache_max_size=hot_cache_max_size,
@@ -963,6 +971,14 @@ def _apply_env_overrides(self) -> None:
             self.cache.ssd_cache_max_size = ssd_cache_max
         if hot_cache_only := os.getenv("OMLX_HOT_CACHE_ONLY"):
             self.cache.hot_cache_only = hot_cache_only.lower() in ("true", "1", "yes")
+        if gdn_ssd_split := os.getenv("OMLX_GDN_SSD_SPLIT_ENABLED"):
+            self.cache.gdn_ssd_split_enabled = gdn_ssd_split.lower() in (
+                "true",
+                "1",
+                "yes",
+            )
+        if gdn_ssd_pending_max := os.getenv("OMLX_GDN_SSD_PENDING_MAX_SIZE"):
+            self.cache.gdn_ssd_pending_max_size = gdn_ssd_pending_max
         if initial_blocks := os.getenv("OMLX_INITIAL_CACHE_BLOCKS"):
             try:
                 self.cache.initial_cache_blocks = int(initial_blocks)
@@ -1336,6 +1352,18 @@ def validate(self) -> list[str]:
             )
 
         # Cache validation
+        if self.cache.gdn_ssd_split_enabled and self.cache.hot_cache_only:
+            errors.append(
+                "gdn_ssd_split_enabled cannot be used with hot_cache_only"
+            )
+
+        try:
+            gdn_pending_size = parse_size(self.cache.gdn_ssd_pending_max_size)
+            if gdn_pending_size <= 0:
+                errors.append("gdn_ssd_pending_max_size must be positive")
+        except (AttributeError, TypeError, ValueError) as e:
+            errors.append(f"Invalid gdn_ssd_pending_max_size: {e}")
+
         if self.cache.ssd_cache_max_size.lower() != "auto":
             try:
                 size = parse_size(self.cache.ssd_cache_max_size)
@@ -1471,6 +1499,10 @@ def to_scheduler_config(self) -> SchedulerConfig:
                 self.base_path
             ),
             hot_cache_max_size=self.cache.get_hot_cache_max_size_bytes(),
+            gdn_ssd_split_enabled=self.cache.gdn_ssd_split_enabled,
+            gdn_ssd_pending_max_bytes=parse_size(
+                self.cache.gdn_ssd_pending_max_size
+            ),
         )
 
     def to_dict(self) -> dict[str, Any]: