chore: import upstream snapshot with attribution

python/tokenspeed/runtime/cache/executor/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Cache-operation executors for host and storage transfers."""

python/tokenspeed/runtime/cache/executor/flat_memory_executor.py

@@ -0,0 +1,445 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Flat host-tier executor (M15 Phase D2): drives FlatWriteBack/FlatLoadBack
+page-id pairs against the byte-blind :class:`FlatHostMirror`, replacing the
+radix ``MemoryExecutor`` when serving with a flat-built scheduler ext and the
+kvstore enabled. Unlike the radix host executor it ACKS loadbacks: the flat
+C++ scheduler pins source host pages and destination device blocks until a
+``Cache.LoadBackDoneEvent`` retires the op.
+"""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from collections.abc import Iterable, Sequence
+
+import psutil
+import torch
+from tokenspeed_scheduler import Cache
+
+from tokenspeed.runtime.cache.executor.host_executor import (
+    _Ack,
+    _cache_stream_priorities,
+    _new_cache_stream,
+    _ordered_unique,
+)
+from tokenspeed.runtime.cache.flat_host_mirror import (
+    FlatHostMirror,
+    flat_bytes_per_host_page,
+)
+from tokenspeed.runtime.cache.kvstore_controller import LayerDoneCounter
+from tokenspeed.runtime.cache.transfer.types import CacheKind
+from tokenspeed.runtime.execution.cuda_graph_wrapper import get_is_capture_mode
+from tokenspeed.runtime.utils import get_colorful_logger
+
+logger = get_colorful_logger(__name__)
+
+_HOST_MEM_HEADROOM_BYTES = 10 * (1024**3)
+
+
+def flat_num_host_pages(
+    *,
+    bytes_per_host_page: int,
+    device_pool_size: int,
+    page_size: int,
+    host_ratio: float,
+    host_size_gb: float,
+) -> int:
+    """Host page budget from the kvstore sizing knobs (same knobs the radix
+    ``HostKVCache`` resolves, kv_cache_host.py:91-102, budget arithmetic only):
+
+    - ``host_size_gb > 0``: explicit byte budget, floor to whole mirror pages
+      (never exceeds the requested bytes):
+      ``host_size_gb * 1e9 // bytes_per_host_page``.
+    - otherwise ratio sizing, mirroring the radix token->page align-up:
+      ``int(device_pool_size * host_ratio) // page_size + 1``.
+    """
+    if bytes_per_host_page <= 0:
+        raise ValueError(f"bytes_per_host_page must be > 0, got {bytes_per_host_page}")
+    if page_size <= 0:
+        raise ValueError(f"page_size must be > 0, got {page_size}")
+    if host_size_gb > 0:
+        num_pages = int(host_size_gb * 1e9 // bytes_per_host_page)
+    else:
+        num_pages = int(device_pool_size * host_ratio) // page_size + 1
+    if num_pages <= 0:
+        raise ValueError(
+            "flat host tier resolved to zero host pages "
+            f"(host_size_gb={host_size_gb}, host_ratio={host_ratio}, "
+            f"bytes_per_host_page={bytes_per_host_page}); increase the "
+            "kvstore size."
+        )
+    return num_pages
+
+
+class FlatMemoryExecutor:
+    """Slim replacement for ``MemoryExecutor`` under the flat host tier.
+
+    Exposes the exact surface ``EventLoop`` drives: ``submit_plan`` /
+    ``poll_results`` / ``get_producer_index`` / ``set_consumer`` (plus the
+    ``host_exec.pools`` attribute walk in ``_setup_layerwise_loadback``).
+    No host pool, no storage executor, no mamba: the flat scheduler config
+    validation already rejects those setups.
+    """
+
+    # EventLoop keys per-op inflight accounting off this: flat loadbacks are
+    # acked (LoadBackDoneEvent), radix loadbacks are not.
+    emits_loadback_acks = True
+
+    def __init__(self, device_pool, *, host_ratio: float, host_size_gb: float):
+        self.page_size = int(device_pool.page_size)
+        self.layer_num = len(device_pool.k_buffer)
+
+        bytes_per_host_page = flat_bytes_per_host_page(device_pool)
+        num_host_pages = flat_num_host_pages(
+            bytes_per_host_page=bytes_per_host_page,
+            device_pool_size=int(device_pool.size),
+            page_size=self.page_size,
+            host_ratio=host_ratio,
+            host_size_gb=host_size_gb,
+        )
+        requested_bytes = num_host_pages * bytes_per_host_page
+        available_bytes = psutil.virtual_memory().available - _HOST_MEM_HEADROOM_BYTES
+        if requested_bytes > available_bytes:
+            raise ValueError(
+                f"Not enough host memory for the flat host tier. Requesting "
+                f"{requested_bytes / 1e9:.2f} GB but only have "
+                f"{available_bytes / 1e9:.2f} GB free. Please reduce the "
+                f"size of the KVStore."
+            )
+        logger.info(
+            "Allocating %.2f GB pinned host memory for the flat host tier "
+            "(num_host_pages=%s bytes_per_host_page=%s host_size_gb=%r "
+            "host_ratio=%r device_pool.size=%r)",
+            requested_bytes / 1e9,
+            num_host_pages,
+            bytes_per_host_page,
+            host_size_gb,
+            host_ratio,
+            device_pool.size,
+        )
+        self.mirror = FlatHostMirror(device_pool, num_host_pages)
+        self.num_host_pages = num_host_pages
+
+        # Layerwise loadback fencing: register the counter where the radix
+        # KVCachePool would, so pool.get_key_buffer/get_value_buffer gate on
+        # the same wait_until(layer_id) machinery.
+        self._counter = LayerDoneCounter(self.layer_num)
+        device_pool.register_layer_transfer_counter(self._counter)
+        # _start_loading maps layer -> mirror V-tensor event and relies on
+        # load_events[-1] (LayerLoadingEvent.finish_event, reuse fence in
+        # update_producer) covering EVERY copy: it pins the last layer to
+        # the op's last per-tensor event, which without state slabs is the
+        # last layer's V event only if that V mirror is the last KV tensor
+        # pair. Holds for both layouts (legacy: identity; slab: last layer
+        # is the last occurrence of its group). A state last layer has no
+        # KV mirror at all -- its copies (state slabs trail every KV
+        # tensor) are covered by the events[-1] pin in _start_loading.
+        assert (
+            self.mirror.state_tensor_indices_of_layer(self.layer_num - 1) is not None
+            or self.mirror.tensor_index_of_layer(self.layer_num - 1)
+            == self.mirror.num_k_tensors - 1
+        ), "flat host tier: last layer's V mirror is not the last KV tensor pair"
+
+        write_priority, load_priority = _cache_stream_priorities()
+        self.write_stream = _new_cache_stream(write_priority)
+        self.load_stream = _new_cache_stream(load_priority)
+
+        # (device_page, host_page) pairs staged between submit() and flush().
+        self._pending_write_pairs: list[tuple[int, int]] = []
+        self._pending_write_op_ids: list[int] = []
+        self._pending_load_pairs: list[tuple[int, int]] = []
+        self._pending_load_op_ids: list[int] = []
+        self.ack_write_queue: list[_Ack] = []
+        self.ack_load_queue: list[_Ack] = []
+        # Ops whose page lists were empty on the wire (C++ dedups transfers
+        # across ops of one batched operation) and no batch event covers them.
+        self._immediate_write_op_ids: list[int] = []
+        self._immediate_load_op_ids: list[int] = []
+
+        self._producer_map: OrderedDict[int, int] = OrderedDict()
+        self._producer_map_limit = 1024
+
+        # Surface for EventLoop._setup_layerwise_loadback, which walks
+        # memory_executor.host_exec.pools to enumerate fencing kinds.
+        self.host_exec = self
+        self.pools = {CacheKind.KV: self.mirror}
+
+    # ------------------------------------------------------------------
+    # Submission (wire shape: batched Flat{WriteBack,LoadBack}Operation)
+    # ------------------------------------------------------------------
+
+    def submit_plan(self, plan) -> None:
+        if plan.cache:
+            logger.debug("[cache_op] flat submit_plan: %s cache ops", len(plan.cache))
+        for op in plan.cache:
+            self.submit(op)
+        self.flush()
+
+    def submit(self, op) -> None:
+        if isinstance(op, Cache.WriteBackOp):
+            self.submit_writeback(op.op_ids, op.src_pages, op.dst_pages)
+        elif isinstance(op, Cache.LoadBackOp):
+            self.submit_loadback(op.op_ids, op.src_pages, op.dst_pages)
+        else:
+            raise ValueError(
+                f"flat host tier: unsupported cache op kind {type(op).__name__}"
+            )
+
+    def _submit(
+        self,
+        op_ids: Sequence[int],
+        src_pages: Sequence[Sequence[int]],
+        dst_pages: Sequence[Sequence[int]],
+        *,
+        pending_op_ids: list[int],
+        pending_pairs: list[tuple[int, int]],
+        src_is_device: bool,
+    ) -> None:
+        """Stage copies as (device_page, host_page) pairs; fail loud on a
+        ragged wire payload instead of silently dropping trailing ops."""
+        assert len(op_ids) == len(src_pages) == len(dst_pages), (
+            f"flat host tier: ragged cache-op payload (op_ids={len(op_ids)}, "
+            f"src_pages={len(src_pages)}, dst_pages={len(dst_pages)})"
+        )
+        for op_id, src, dst in zip(op_ids, src_pages, dst_pages):
+            assert len(src) == len(dst), (
+                f"flat host tier: op {op_id} src/dst page lists differ "
+                f"({len(src)} vs {len(dst)})"
+            )
+            pending_op_ids.append(int(op_id))
+            device_pages, host_pages = (src, dst) if src_is_device else (dst, src)
+            pending_pairs.extend(
+                (int(d), int(h)) for d, h in zip(device_pages, host_pages)
+            )
+
+    def submit_writeback(
+        self,
+        op_ids: Sequence[int],
+        src_pages: Sequence[Sequence[int]],
+        dst_pages: Sequence[Sequence[int]],
+    ) -> None:
+        """Stage device->host copies: src=device pages, dst=host pages."""
+        self._submit(
+            op_ids,
+            src_pages,
+            dst_pages,
+            pending_op_ids=self._pending_write_op_ids,
+            pending_pairs=self._pending_write_pairs,
+            src_is_device=True,
+        )
+
+    def submit_loadback(
+        self,
+        op_ids: Sequence[int],
+        src_pages: Sequence[Sequence[int]],
+        dst_pages: Sequence[Sequence[int]],
+    ) -> None:
+        """Stage host->device copies: src=host pages, dst=device pages."""
+        self._submit(
+            op_ids,
+            src_pages,
+            dst_pages,
+            pending_op_ids=self._pending_load_op_ids,
+            pending_pairs=self._pending_load_pairs,
+            src_is_device=False,
+        )
+
+    def flush(self) -> None:
+        self._start_loading()
+        self._start_writing()
+
+    def _start_writing(self) -> None:
+        if not self._pending_write_op_ids:
+            return
+        op_ids = _ordered_unique(self._pending_write_op_ids)
+        pairs = self._pending_write_pairs
+        self._pending_write_op_ids = []
+        self._pending_write_pairs = []
+        if not pairs:
+            self._immediate_write_op_ids.extend(op_ids)
+            return
+        # Order the D2H copies after already-enqueued default-stream work
+        # (same fence the radix _start_writing places).
+        start_event = torch.cuda.Event()
+        start_event.record()
+        start_event.wait(self.write_stream)
+        self.mirror.store_pages(pairs, self.write_stream)
+        finish_event = torch.cuda.Event()
+        finish_event.record(self.write_stream)
+        self.ack_write_queue.append(_Ack(finish_event, op_ids))
+
+    def _start_loading(self) -> None:
+        if not self._pending_load_op_ids:
+            return
+        assert (
+            not get_is_capture_mode()
+        ), "cache loadback must run in eager admission iter"
+        op_ids = _ordered_unique(self._pending_load_op_ids)
+        pairs = self._pending_load_pairs
+        self._pending_load_op_ids = []
+        self._pending_load_pairs = []
+        if not pairs:
+            self._immediate_load_op_ids.extend(op_ids)
+            return
+
+        producer_id = self._counter.update_producer()
+        producer_event = self._counter.events[producer_id]
+        producer_event.start_event.record()
+        producer_event.start_event.wait(self.load_stream)
+
+        events = self.mirror.load_pages_with_events(pairs, self.load_stream)
+        # Layer fence: layer L is readable once its V mirror copy lands; the
+        # load stream is serial (all K copies precede all V copies), so the
+        # V-tensor event also covers L's K copy. Paired slab layers share the
+        # slab event -- correct by design. State layers instead fence on
+        # their ssm event: conv precedes ssm in tensor_pairs order (and both
+        # follow every KV tensor), so on the serial stream the ssm event
+        # covers the conv copy and the layer's KV copies -- the same
+        # K-before-V reasoning as above.
+        num_k = self.mirror.num_k_tensors
+        for layer_id in range(self.layer_num):
+            state_indices = self.mirror.state_tensor_indices_of_layer(layer_id)
+            if state_indices is not None:
+                producer_event.load_events[layer_id] = events[state_indices[1]]
+            else:
+                producer_event.load_events[layer_id] = events[
+                    num_k + self.mirror.tensor_index_of_layer(layer_id)
+                ]
+        # finish_event (== load_events[-1]) is the producer-slot reuse fence
+        # in update_producer and must cover EVERY copy of the op; state
+        # tensors copy after all KV tensors, so pin the last layer to the
+        # op's last per-tensor event. A no-op without state slabs: events[-1]
+        # is then the last layer's V event (ctor assert).
+        producer_event.load_events[self.layer_num - 1] = events[-1]
+        # events[-1] is also the reassigned finish_event, so the ack covers
+        # every copy.
+        self.ack_load_queue.append(_Ack(events[-1], op_ids))
+        for op_id in op_ids:
+            self._producer_map[op_id] = producer_id
+        while len(self._producer_map) > self._producer_map_limit:
+            self._producer_map.popitem(last=False)
+
+    # ------------------------------------------------------------------
+    # Ack draining
+    # ------------------------------------------------------------------
+
+    def poll_results(self) -> list:
+        results: list = []
+        for op_id in self._immediate_write_op_ids:
+            results.append(self._write_done(op_id))
+        self._immediate_write_op_ids.clear()
+        for op_id in self._immediate_load_op_ids:
+            results.append(self._load_done(op_id))
+        self._immediate_load_op_ids.clear()
+
+        remaining_writes = []
+        for ack in self.ack_write_queue:
+            if ack.finish_event.query():
+                results.extend(self._write_done(op_id) for op_id in ack.op_ids)
+            else:
+                remaining_writes.append(ack)
+        self.ack_write_queue[:] = remaining_writes
+
+        remaining_loads = []
+        for ack in self.ack_load_queue:
+            if ack.finish_event.query():
+                results.extend(self._load_done(op_id) for op_id in ack.op_ids)
+            else:
+                remaining_loads.append(ack)
+        self.ack_load_queue[:] = remaining_loads
+
+        if results:
+            for r in results:
+                logger.debug(
+                    "[cache_op] flat done op_id=%s success=%s type=%s",
+                    r.op_id,
+                    r.success,
+                    type(r).__name__,
+                )
+        return results
+
+    @staticmethod
+    def _write_done(op_id: int):
+        evt = Cache.WriteBackDoneEvent()
+        evt.op_id = op_id
+        evt.success = True
+        return evt
+
+    @staticmethod
+    def _load_done(op_id: int):
+        evt = Cache.LoadBackDoneEvent()
+        evt.op_id = op_id
+        evt.success = True
+        return evt
+
+    # ------------------------------------------------------------------
+    # Layerwise loadback fencing (EventLoop._setup_layerwise_loadback)
+    # ------------------------------------------------------------------
+
+    def get_producer_index(
+        self, kind_or_op_id: CacheKind | str | int, op_id: int | None = None
+    ) -> int | None:
+        if op_id is None:
+            op_id = int(kind_or_op_id)
+        return self._producer_map.pop(int(op_id), None)
+
+    def set_consumer(
+        self,
+        kind_or_producer_index: CacheKind | str | int | Iterable[int],
+        producer_index: int | Iterable[int] | None = None,
+    ) -> None:
+        if producer_index is None:
+            producer_index = kind_or_producer_index
+        self._counter.set_consumer(producer_index)
+
+    # ------------------------------------------------------------------
+    # MemoryExecutor surface stubs
+    # ------------------------------------------------------------------
+
+    def set_mamba_layerwise_cow(self, cow_dst_pages_by_src) -> None:
+        assert not cow_dst_pages_by_src, (
+            "flat host tier has no mamba L2: the flat scheduler config "
+            "validation rejects state-only pools"
+        )
+
+    def query_l3_pages(self, hashes: list[str]) -> int:
+        # No L3 storage tier under the flat build (EventLoop refuses a
+        # storage backend up front); report zero hits.
+        return 0
+
+    def shutdown(self) -> None:
+        self.write_stream.synchronize()
+        self.load_stream.synchronize()
+
+    def reset(self) -> None:
+        self.write_stream.synchronize()
+        self.load_stream.synchronize()
+        self._pending_write_pairs.clear()
+        self._pending_write_op_ids.clear()
+        self._pending_load_pairs.clear()
+        self._pending_load_op_ids.clear()
+        self.ack_write_queue.clear()
+        self.ack_load_queue.clear()
+        self._immediate_write_op_ids.clear()
+        self._immediate_load_op_ids.clear()
+        self._producer_map.clear()
+        self._counter.reset()

python/tokenspeed/runtime/cache/executor/host_executor.py

@@ -0,0 +1,516 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Host-side executor for cache writeback and loadback operations."""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from collections.abc import Iterable
+from typing import NamedTuple
+
+import torch
+from tokenspeed_scheduler import Cache
+
+from tokenspeed.runtime.cache.transfer.kv_pool import KVCachePool
+from tokenspeed.runtime.cache.transfer.pool import CachePool
+from tokenspeed.runtime.cache.transfer.types import CacheKind, Location, TransferUnit
+from tokenspeed.runtime.execution.cuda_graph_wrapper import get_is_capture_mode
+from tokenspeed.runtime.utils import get_colorful_logger, get_device_module
+
+logger = get_colorful_logger(__name__)
+device_module = get_device_module()
+CONCURRENT_WRITEBACK_BLOCK_QUOTA = 2
+
+
+def _cache_stream_priorities() -> tuple[int | None, int | None]:
+    priority_range = getattr(device_module.Stream, "priority_range", None)
+    if priority_range is None:
+        return None, None
+    try:
+        least_priority, greatest_priority = priority_range()
+    except (RuntimeError, TypeError):
+        return None, None
+    return least_priority, greatest_priority
+
+
+def _new_cache_stream(priority: int | None = None):
+    if priority is None:
+        return device_module.Stream()
+    try:
+        return device_module.Stream(priority=priority)
+    except (RuntimeError, TypeError):
+        return device_module.Stream()
+
+
+def page_ids_to_token_indices(
+    page_ids: list[int],
+    page_size: int,
+    device: str = "cpu",
+) -> torch.Tensor:
+    if len(page_ids) == 0:
+        return torch.empty((0,), dtype=torch.int64, device=device)
+    pages = torch.tensor(page_ids, dtype=torch.int64, device=device)
+    offsets = torch.arange(page_size, dtype=torch.int64, device=device)
+    return (pages[:, None] * page_size + offsets[None, :]).reshape(-1)
+
+
+def _dedupe_page_pairs(
+    src_pages: Iterable[int],
+    dst_pages: Iterable[int],
+) -> tuple[list[int], list[int]]:
+    seen = set()
+    deduped_src = []
+    deduped_dst = []
+    for src_page, dst_page in zip(src_pages, dst_pages):
+        pair = (int(src_page), int(dst_page))
+        if pair in seen:
+            continue
+        seen.add(pair)
+        deduped_src.append(pair[0])
+        deduped_dst.append(pair[1])
+    return deduped_src, deduped_dst
+
+
+def _ordered_unique(values: Iterable[int]) -> list[int]:
+    seen = set()
+    result = []
+    for value in values:
+        value = int(value)
+        if value in seen:
+            continue
+        seen.add(value)
+        result.append(value)
+    return result
+
+
+class _Ack(NamedTuple):
+    finish_event: object  # device_module.Event
+    op_ids: list[int]
+
+
+class HostExecutor:
+    def __init__(
+        self,
+        page_size: int | None = None,
+        device_pool=None,
+        host_pool=None,
+        io_backend: str = "kernel",
+        layer_num: int | None = None,
+        draft_device_pool=None,
+        draft_host_pool=None,
+        draft_layer_num: int = 0,
+        pools: list[CachePool] | None = None,
+    ):
+        self.io_backend = io_backend
+        if pools is None:
+            if (
+                page_size is None
+                or device_pool is None
+                or host_pool is None
+                or layer_num is None
+            ):
+                raise ValueError("HostExecutor requires either pools or KV pool inputs")
+            pools = [
+                KVCachePool(
+                    device_pool=device_pool,
+                    host_pool=host_pool,
+                    io_backend=io_backend,
+                    layer_num=layer_num,
+                    draft_device_pool=draft_device_pool,
+                    draft_host_pool=draft_host_pool,
+                    draft_layer_num=draft_layer_num,
+                )
+            ]
+        if not pools:
+            raise ValueError("HostExecutor requires at least one cache pool")
+
+        self.pools = {CacheKind(pool.kind): pool for pool in pools}
+        self.device = next(iter(self.pools.values())).device
+
+        write_priority, load_priority = _cache_stream_priorities()
+        self.write_stream = _new_cache_stream(write_priority)
+        self.load_stream = _new_cache_stream(load_priority)
+        self._writeback_block_quota: int | None = None
+
+        self.write_queues: dict[CacheKind, list[TransferUnit]] = {
+            kind: [] for kind in self.pools
+        }
+        self.load_queues: dict[CacheKind, list[TransferUnit]] = {
+            kind: [] for kind in self.pools
+        }
+
+        self.ack_write_queue: list[_Ack] = []
+        self.ack_load_queue: list[_Ack] = []
+        self.completed_writebacks: list[int] = []
+
+        self._counters = {
+            kind: pool.get_layer_done_counter() for kind, pool in self.pools.items()
+        }
+        self._producer_map: dict[CacheKind, OrderedDict[int, int]] = {
+            kind: OrderedDict() for kind in self.pools
+        }
+        self._producer_map_limit = 1024
+
+    def enqueue_writeback(
+        self,
+        op_id,
+        src_pages,
+        dst_pages,
+        is_retract: bool = False,
+        kind: CacheKind | str = CacheKind.KV,
+    ) -> None:
+        kind = CacheKind(kind)
+        pool = self.pools[kind]
+        src_pages, dst_pages = _dedupe_page_pairs(src_pages, dst_pages)
+        if not src_pages:
+            self.completed_writebacks.append(op_id)
+            return
+        device_indices = page_ids_to_token_indices(src_pages, pool.page_size(), "cpu")
+        host_indices = page_ids_to_token_indices(dst_pages, pool.page_size(), "cpu")
+        self.write_queues[kind].append(
+            TransferUnit(
+                kind=kind,
+                src_loc=Location.DEVICE,
+                dst_loc=Location.HOST,
+                src_indices=device_indices,
+                dst_indices=host_indices,
+                op_id=op_id,
+                is_retract=is_retract,
+            )
+        )
+
+    def enqueue_loadback(
+        self,
+        op_id,
+        src_pages,
+        dst_pages,
+        kind: CacheKind | str = CacheKind.KV,
+        layerwise_cow_dst_pages_by_src: dict[int, list[int]] | None = None,
+    ) -> None:
+        kind = CacheKind(kind)
+        pool = self.pools[kind]
+        src_pages, dst_pages = _dedupe_page_pairs(src_pages, dst_pages)
+        if not src_pages:
+            return
+        host_indices = page_ids_to_token_indices(src_pages, pool.page_size(), "cpu")
+        device_indices = page_ids_to_token_indices(dst_pages, pool.page_size(), "cpu")
+        cow_src_indices = None
+        cow_dst_indices = None
+        if layerwise_cow_dst_pages_by_src:
+            cow_src_pages: list[int] = []
+            cow_dst_pages: list[int] = []
+            for dst_page in dst_pages:
+                for cow_dst in layerwise_cow_dst_pages_by_src.get(int(dst_page), []):
+                    cow_src_pages.append(int(dst_page))
+                    cow_dst_pages.append(int(cow_dst))
+            if cow_src_pages:
+                cow_src_indices = page_ids_to_token_indices(
+                    cow_src_pages, pool.page_size(), "cpu"
+                )
+                cow_dst_indices = page_ids_to_token_indices(
+                    cow_dst_pages, pool.page_size(), "cpu"
+                )
+        self.load_queues[kind].append(
+            TransferUnit(
+                kind=kind,
+                src_loc=Location.HOST,
+                dst_loc=Location.DEVICE,
+                src_indices=host_indices,
+                dst_indices=device_indices,
+                op_id=op_id,
+                layerwise_cow_src_indices=cow_src_indices,
+                layerwise_cow_dst_indices=cow_dst_indices,
+            )
+        )
+
+    def flush(self) -> None:
+        throttle_writeback = self._has_work(self.load_queues) and not any(
+            unit.is_retract for units in self.write_queues.values() for unit in units
+        )
+        writeback_block_quota = (
+            CONCURRENT_WRITEBACK_BLOCK_QUOTA if throttle_writeback else None
+        )
+        previous_writeback_block_quota = getattr(self, "_writeback_block_quota", None)
+        self._writeback_block_quota = writeback_block_quota
+        try:
+            self._start_loading()
+            self._start_writing()
+        finally:
+            self._writeback_block_quota = previous_writeback_block_quota
+
+    def _start_writing(self) -> None:
+        if not self._has_work(self.write_queues):
+            return
+
+        start_event = device_module.Event()
+        finish_event = device_module.Event()
+        op_ids: list[int] = []
+
+        start_event.record()
+        with device_module.stream(self.write_stream):
+            start_event.wait(self.write_stream)
+            for kind, units in self.write_queues.items():
+                if not units:
+                    continue
+                pool = self.pools[kind]
+                unit = self._merge_units(units)
+                src_indices, dst_indices = self._prepare_indices(unit, pool)
+                self._pool_writeback(
+                    pool, src_indices.to(torch.int64), dst_indices.to(torch.int64)
+                )
+                self._record_if_cuda(src_indices, self.write_stream)
+                self._record_if_cuda(dst_indices, self.write_stream)
+                op_ids.extend(unit.op_id for unit in units)
+            finish_event.record()
+
+        self._clear_queues(self.write_queues)
+        self.ack_write_queue.append(_Ack(finish_event, _ordered_unique(op_ids)))
+
+    def _start_loading(self) -> None:
+        if not self._has_work(self.load_queues):
+            return
+        assert (
+            not get_is_capture_mode()
+        ), "cache loadback must run in eager admission iter"
+
+        with device_module.stream(self.load_stream):
+            for kind, units in self.load_queues.items():
+                if not units:
+                    continue
+                pool = self.pools[kind]
+                counter = self._counters[kind]
+                producer_id = counter.update_producer()
+                producer_event = counter.events[producer_id]
+                producer_event.start_event.record()
+                producer_event.start_event.wait(self.load_stream)
+
+                unit = self._merge_units(units)
+                src_indices, dst_indices = self._prepare_indices(unit, pool)
+                layerwise_copy = getattr(pool, "copy_layer", None)
+                cow_src_indices = unit.layerwise_cow_src_indices
+                cow_dst_indices = unit.layerwise_cow_dst_indices
+                for layer_index in range(pool.num_layers()):
+                    pool.loadback(
+                        src_indices.to(torch.int64),
+                        dst_indices.to(torch.int64),
+                        layer_index,
+                    )
+                    if (
+                        layerwise_copy is not None
+                        and cow_src_indices is not None
+                        and cow_dst_indices is not None
+                    ):
+                        layerwise_copy(
+                            cow_src_indices.to(torch.int64),
+                            cow_dst_indices.to(torch.int64),
+                            layer_index,
+                        )
+                    producer_event.complete(layer_index)
+                self._record_if_cuda(src_indices, self.load_stream)
+                self._record_if_cuda(dst_indices, self.load_stream)
+                if cow_src_indices is not None:
+                    self._record_if_cuda(cow_src_indices, self.load_stream)
+                if cow_dst_indices is not None:
+                    self._record_if_cuda(cow_dst_indices, self.load_stream)
+
+                op_ids = _ordered_unique(unit.op_id for unit in units)
+                self.ack_load_queue.append(_Ack(producer_event.finish_event, op_ids))
+                producer_map = self._producer_map[kind]
+                for op_id in op_ids:
+                    producer_map[op_id] = producer_id
+                while len(producer_map) > self._producer_map_limit:
+                    producer_map.popitem(last=False)
+
+        self._clear_queues(self.load_queues)
+
+    @staticmethod
+    def _has_work(queues: dict[CacheKind, list[TransferUnit]]) -> bool:
+        return any(bool(units) for units in queues.values())
+
+    @staticmethod
+    def _clear_queues(queues: dict[CacheKind, list[TransferUnit]]) -> None:
+        for units in queues.values():
+            units.clear()
+
+    @staticmethod
+    def _merge_units(units: list[TransferUnit]) -> TransferUnit:
+        assert units
+        if len(units) == 1:
+            return units[0]
+        first = units[0]
+        cow_src_indices = [
+            unit.layerwise_cow_src_indices
+            for unit in units
+            if unit.layerwise_cow_src_indices is not None
+        ]
+        cow_dst_indices = [
+            unit.layerwise_cow_dst_indices
+            for unit in units
+            if unit.layerwise_cow_dst_indices is not None
+        ]
+        return TransferUnit(
+            kind=first.kind,
+            src_loc=first.src_loc,
+            dst_loc=first.dst_loc,
+            src_indices=torch.cat([unit.src_indices for unit in units]),
+            dst_indices=torch.cat([unit.dst_indices for unit in units]),
+            op_id=-1,
+            is_retract=any(unit.is_retract for unit in units),
+            layerwise_cow_src_indices=(
+                torch.cat(cow_src_indices) if cow_src_indices else None
+            ),
+            layerwise_cow_dst_indices=(
+                torch.cat(cow_dst_indices) if cow_dst_indices else None
+            ),
+        )
+
+    def _prepare_indices(
+        self, unit: TransferUnit, pool: CachePool
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if unit.src_loc == Location.HOST:
+            host_indices = unit.src_indices
+            device_indices = unit.dst_indices
+        elif unit.dst_loc == Location.HOST:
+            host_indices = unit.dst_indices
+            device_indices = unit.src_indices
+        else:
+            raise ValueError(f"unsupported transfer direction: {unit.direction}")
+
+        io_backend = getattr(pool, "io_backend", self.io_backend)
+        if io_backend == "kernel":
+            target_device = pool.device
+            if device_indices.device != target_device:
+                device_indices = device_indices.to(target_device, non_blocking=True)
+            if host_indices.device != target_device:
+                host_indices = host_indices.to(target_device, non_blocking=True)
+        elif io_backend == "direct":
+            if pool.host_layout == "layer_first":
+                device_indices = device_indices.cpu()
+                host_indices, idx = host_indices.sort()
+                device_indices = device_indices.index_select(0, idx)
+            else:
+                raise ValueError(f"Unsupported host layout: {pool.host_layout}")
+        else:
+            raise ValueError(f"Unsupported io_backend={io_backend}")
+
+        if unit.src_loc == Location.HOST:
+            return host_indices, device_indices
+        return device_indices, host_indices
+
+    def _pool_writeback(
+        self, pool: CachePool, src_indices: torch.Tensor, dst_indices: torch.Tensor
+    ) -> None:
+        try:
+            pool.writeback(
+                src_indices, dst_indices, block_quota=self._writeback_block_quota
+            )
+        except TypeError as exc:
+            if "block_quota" not in str(exc):
+                raise
+            pool.writeback(src_indices, dst_indices)
+
+    @staticmethod
+    def _record_if_cuda(tensor: torch.Tensor, stream) -> None:
+        if tensor.is_cuda:
+            tensor.record_stream(stream)
+
+    def drain(self) -> list:
+        results: list = []
+        results.extend(self._poll_write_acks())
+        results.extend(self._poll_load_acks())
+        return results
+
+    def _poll_write_acks(self) -> list:
+        results = []
+        completed_writebacks = getattr(self, "completed_writebacks", [])
+        for op_id in completed_writebacks:
+            logger.debug("[cache_op] writeback done op_id=%s immediate=True", op_id)
+            evt = Cache.WriteBackDoneEvent()
+            evt.op_id = op_id
+            evt.success = True
+            results.append(evt)
+        completed_writebacks.clear()
+        remaining = []
+        for ack in self.ack_write_queue:
+            if ack.finish_event.query():
+                logger.debug(
+                    "[cache_op] writeback done op_ids=%s immediate=False", ack.op_ids
+                )
+                for op_id in ack.op_ids:
+                    evt = Cache.WriteBackDoneEvent()
+                    evt.op_id = op_id
+                    evt.success = True
+                    results.append(evt)
+            else:
+                remaining.append(ack)
+        self.ack_write_queue[:] = remaining
+        return results
+
+    def _poll_load_acks(self) -> list:
+        results = []
+        remaining = []
+        for ack in self.ack_load_queue:
+            if not ack.finish_event.query():
+                remaining.append(ack)
+        self.ack_load_queue[:] = remaining
+        return results
+
+    def get_producer_index(
+        self, kind_or_op_id: CacheKind | str | int, op_id: int | None = None
+    ) -> int | None:
+        if op_id is None:
+            kind = CacheKind.KV
+            op_id = int(kind_or_op_id)
+        else:
+            kind = CacheKind(kind_or_op_id)
+        return self._producer_map[kind].pop(int(op_id), None)
+
+    def set_consumer(
+        self,
+        kind_or_producer_index: CacheKind | str | int | Iterable[int],
+        producer_index: int | Iterable[int] | None = None,
+    ) -> None:
+        if producer_index is None:
+            kind = CacheKind.KV
+            producer_index = kind_or_producer_index
+        else:
+            kind = CacheKind(kind_or_producer_index)
+        self._counters[kind].set_consumer(producer_index)
+
+    def shutdown(self) -> None:
+        self.write_stream.synchronize()
+        self.load_stream.synchronize()
+        for pool in self.pools.values():
+            shutdown = getattr(pool, "shutdown", None)
+            if shutdown is not None:
+                shutdown()
+
+    def reset(self) -> None:
+        self.write_stream.synchronize()
+        self.load_stream.synchronize()
+        self._clear_queues(self.write_queues)
+        self._clear_queues(self.load_queues)
+        self.ack_write_queue.clear()
+        self.ack_load_queue.clear()
+        for producer_map in self._producer_map.values():
+            producer_map.clear()
+        for counter in self._counters.values():
+            counter.reset()

python/tokenspeed/runtime/cache/executor/memory_executor.py

@@ -0,0 +1,500 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Top-level memory executor that coordinates host and storage executors."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass
+
+from tokenspeed_scheduler import Cache
+
+from tokenspeed.runtime.cache.executor.host_executor import HostExecutor
+from tokenspeed.runtime.cache.executor.storage_executor import StorageExecutor
+from tokenspeed.runtime.cache.kv_cache_host import (
+    DSATokenToKVPoolHost,
+    MHATokenToKVPoolHost,
+    MLATokenToKVPoolHost,
+    get_available_host_memory_bytes,
+)
+from tokenspeed.runtime.cache.mamba_cache_host import MambaPoolHost
+from tokenspeed.runtime.cache.transfer.kv_pool import KVCachePool
+from tokenspeed.runtime.cache.transfer.mamba_pool import MambaCachePool
+from tokenspeed.runtime.cache.transfer.types import CacheKind
+from tokenspeed.runtime.layers.attention.kv_cache.dsa import DSATokenToKVPool
+from tokenspeed.runtime.layers.attention.kv_cache.mha import MHATokenToKVPool
+from tokenspeed.runtime.layers.attention.kv_cache.mla import MLATokenToKVPool
+from tokenspeed.runtime.utils import get_colorful_logger
+
+logger = get_colorful_logger(__name__)
+
+
+@dataclass(slots=True)
+class MemoryExecutorConfig:
+    layer_num: int
+    page_size: int = 64
+    host_ratio: float = 2.0
+    host_size_gb: int = 0
+    host_parallel_count: int = 1
+    host_reserve_gb: float = 10.0
+    io_backend: str = "kernel"
+    host_layout: str = "layer_first"
+    storage_backend: str | None = "mooncake"
+    storage_backend_extra_config: str | None = None
+    model_name: str | None = None
+    enable_mamba_l2: bool = False
+    mamba_l2_host_slots: int = 0
+    mamba_l2_layout: str = "layer_first"
+    mamba_l2_io_backend: str = "kernel"
+
+
+def _aligned_token_count(size: int, page_size: int) -> int:
+    return (size // page_size + 1) * page_size
+
+
+def _pool_size_per_token(pool) -> int:
+    dtype_size = pool.store_dtype.itemsize
+    if isinstance(pool, DSATokenToKVPool):
+        latent_size = (
+            (pool.kv_lora_rank + pool.qk_rope_head_dim) * dtype_size * pool.layer_num
+        )
+        return latent_size + pool.index_k_row_bytes * pool.layer_num
+    if isinstance(pool, MHATokenToKVPool):
+        return pool.head_dim * pool.head_num * pool.layer_num * dtype_size * 2
+    if isinstance(pool, MLATokenToKVPool):
+        return (pool.kv_lora_rank + pool.qk_rope_head_dim) * dtype_size * pool.layer_num
+    raise ValueError(f"Unsupported KV pool type for host budget: {type(pool)}")
+
+
+def _auto_capped_host_size_tokens(
+    *,
+    requested_tokens: int,
+    page_size: int,
+    size_per_token: int,
+    available_host_memory_bytes: int,
+    host_parallel_count: int,
+) -> int:
+    """Return a HostKVCache host_size_tokens override, or 0 when no cap is needed."""
+
+    aligned_requested_tokens = _aligned_token_count(requested_tokens, page_size)
+    requested_bytes = aligned_requested_tokens * size_per_token
+    per_rank_budget = available_host_memory_bytes // max(host_parallel_count, 1)
+    if requested_bytes <= per_rank_budget:
+        return 0
+
+    budget_tokens = per_rank_budget // max(size_per_token, 1)
+    max_aligned_tokens = (budget_tokens // page_size) * page_size
+    if max_aligned_tokens <= page_size:
+        raise ValueError(
+            "Not enough host memory available for KVStore after cgroup-aware "
+            f"budgeting: per-rank budget={per_rank_budget / 1e9:.2f} GB, "
+            f"size_per_token={size_per_token}."
+        )
+    return max_aligned_tokens - page_size
+
+
+class MemoryExecutor:
+    """Coordinate host-memory and storage-backed cache operations."""
+
+    def __init__(
+        self,
+        device_pool,
+        config: MemoryExecutorConfig,
+        is_dp_attention_enabled: bool,
+        tp_group=None,
+        draft_device_pool=None,
+        mamba_pool=None,
+    ):
+        self.page_size = config.page_size
+        kv_pool_types = (DSATokenToKVPool, MHATokenToKVPool, MLATokenToKVPool)
+
+        # Unwrap LayerMappedKVPool (hybrid GDN models) to get the inner MHA pool.
+        actual_pool = device_pool
+        if hasattr(device_pool, "inner") and not isinstance(device_pool, kv_pool_types):
+            actual_pool = device_pool.inner
+
+        actual_draft_pool = None
+        if draft_device_pool is not None:
+            actual_draft_pool = draft_device_pool
+            if hasattr(draft_device_pool, "inner") and not isinstance(
+                draft_device_pool, kv_pool_types
+            ):
+                actual_draft_pool = draft_device_pool.inner
+            if not isinstance(actual_draft_pool, kv_pool_types):
+                raise ValueError(
+                    f"draft_device_pool only supports DSA, MHA and MLA, "
+                    f"got {type(actual_draft_pool)}"
+                )
+
+        host_size_tokens = 0
+        if config.host_size_gb == 0:
+            target_size_per_token = _pool_size_per_token(actual_pool)
+            draft_size_per_token = (
+                _pool_size_per_token(actual_draft_pool)
+                if actual_draft_pool is not None
+                else 0
+            )
+            combined_size_per_token = target_size_per_token + draft_size_per_token
+            reserve_bytes = int(config.host_reserve_gb * (1024**3))
+            available_bytes, _, cgroup_available = get_available_host_memory_bytes(
+                reserve_bytes
+            )
+            requested_tokens = int(actual_pool.size * config.host_ratio)
+            host_size_tokens = _auto_capped_host_size_tokens(
+                requested_tokens=requested_tokens,
+                page_size=config.page_size,
+                size_per_token=combined_size_per_token,
+                available_host_memory_bytes=available_bytes,
+                host_parallel_count=config.host_parallel_count,
+            )
+            if host_size_tokens > 0:
+                capped_tokens = _aligned_token_count(host_size_tokens, config.page_size)
+                requested_tokens_aligned = _aligned_token_count(
+                    requested_tokens, config.page_size
+                )
+                logger.warning(
+                    "Capping KVStore host pool for cgroup budget: "
+                    "tokens %s -> %s, total bytes %.2f GB -> %.2f GB "
+                    "(parallel_count=%s, available=%.2f GB, cgroup_available=%s)",
+                    requested_tokens_aligned,
+                    capped_tokens,
+                    requested_tokens_aligned * combined_size_per_token / 1e9,
+                    capped_tokens * combined_size_per_token / 1e9,
+                    config.host_parallel_count,
+                    available_bytes / 1e9,
+                    (
+                        f"{cgroup_available / 1e9:.2f} GB"
+                        if cgroup_available is not None
+                        else "unlimited"
+                    ),
+                )
+
+        # DSA subclasses MLA, so it must be matched before the MLA branch.
+        if isinstance(actual_pool, DSATokenToKVPool):
+            self.host_pool = DSATokenToKVPoolHost(
+                actual_pool,
+                config.host_ratio,
+                config.host_size_gb,
+                config.page_size,
+                config.host_layout,
+                host_size_tokens=host_size_tokens,
+            )
+        elif isinstance(actual_pool, MHATokenToKVPool):
+            self.host_pool = MHATokenToKVPoolHost(
+                actual_pool,
+                config.host_ratio,
+                config.host_size_gb,
+                config.page_size,
+                config.host_layout,
+                host_size_tokens=host_size_tokens,
+            )
+        elif isinstance(actual_pool, MLATokenToKVPool):
+            self.host_pool = MLATokenToKVPoolHost(
+                actual_pool,
+                config.host_ratio,
+                config.host_size_gb,
+                config.page_size,
+                config.host_layout,
+                host_size_tokens=host_size_tokens,
+            )
+        else:
+            raise ValueError(
+                f"host_pool only supports DSA, MHA and MLA, got {type(actual_pool)} "
+                f"from module {type(actual_pool).__module__}"
+            )
+
+        # Draft model L2 cache: draft shares the same page mapping as the base
+        # model, so its host pool must hold exactly the same number of tokens.
+        # Pass host_size_tokens directly to bypass ratio/GB recalculation.
+        if actual_draft_pool is not None:
+            if isinstance(actual_draft_pool, DSATokenToKVPool):
+                self.draft_host_pool = DSATokenToKVPoolHost(
+                    actual_draft_pool,
+                    config.host_ratio,
+                    config.host_size_gb,
+                    config.page_size,
+                    config.host_layout,
+                    host_size_tokens=self.host_pool.size,
+                )
+            elif isinstance(actual_draft_pool, MHATokenToKVPool):
+                self.draft_host_pool = MHATokenToKVPoolHost(
+                    actual_draft_pool,
+                    config.host_ratio,
+                    config.host_size_gb,
+                    config.page_size,
+                    config.host_layout,
+                    host_size_tokens=self.host_pool.size,
+                )
+            elif isinstance(actual_draft_pool, MLATokenToKVPool):
+                self.draft_host_pool = MLATokenToKVPoolHost(
+                    actual_draft_pool,
+                    config.host_ratio,
+                    config.host_size_gb,
+                    config.page_size,
+                    config.host_layout,
+                    host_size_tokens=self.host_pool.size,
+                )
+            else:
+                raise ValueError(
+                    f"draft_device_pool only supports DSA, MHA and MLA, "
+                    f"got {type(actual_draft_pool)}"
+                )
+            draft_host_bytes = (
+                self.draft_host_pool.size * self.draft_host_pool.size_per_token
+            )
+            logger.info(
+                "Allocating %.2f GB host memory for draft model L2 cache (pool_type=%s size_tokens=%s size_per_token=%s layer_num=%s)",
+                draft_host_bytes / 1e9,
+                type(self.draft_host_pool).__name__,
+                self.draft_host_pool.size,
+                self.draft_host_pool.size_per_token,
+                actual_draft_pool.layer_num,
+            )
+            draft_layer_num = actual_draft_pool.layer_num
+        else:
+            self.draft_host_pool = None
+            draft_layer_num = 0
+
+        pools = None
+        self.mamba_host_pool = None
+        if (
+            config.enable_mamba_l2
+            and mamba_pool is not None
+            and config.mamba_l2_host_slots > 0
+        ):
+            self.mamba_host_pool = MambaPoolHost(
+                mamba_pool,
+                host_size_slots=config.mamba_l2_host_slots,
+                layout=config.mamba_l2_layout,
+            )
+            pools = [
+                KVCachePool(
+                    device_pool=device_pool,
+                    host_pool=self.host_pool,
+                    io_backend=config.io_backend,
+                    layer_num=actual_pool.layer_num,
+                    draft_device_pool=(
+                        actual_draft_pool if draft_device_pool is not None else None
+                    ),
+                    draft_host_pool=self.draft_host_pool,
+                    draft_layer_num=draft_layer_num,
+                ),
+                MambaCachePool(
+                    device_pool=mamba_pool,
+                    host_pool=self.mamba_host_pool,
+                    io_backend=config.mamba_l2_io_backend,
+                ),
+            ]
+            logger.debug(
+                "[cache_op] MemoryExecutor init pools=%s host_pools=%s draft=%s mamba=%s io_backend=%s host_layout=%s",
+                [pool.kind.value for pool in pools],
+                [type(self.host_pool).__name__, type(self.mamba_host_pool).__name__],
+                self.draft_host_pool is not None,
+                True,
+                config.io_backend,
+                config.host_layout,
+            )
+
+        if pools is not None:
+            self.host_exec = HostExecutor(pools=pools, io_backend=config.io_backend)
+        else:
+            self.host_exec = HostExecutor(
+                page_size=config.page_size,
+                device_pool=device_pool,
+                host_pool=self.host_pool,
+                io_backend=config.io_backend,
+                layer_num=actual_pool.layer_num,
+                draft_device_pool=(
+                    actual_draft_pool if draft_device_pool is not None else None
+                ),
+                draft_host_pool=self.draft_host_pool,
+                draft_layer_num=draft_layer_num,
+            )
+        self.storage_exec = StorageExecutor(
+            page_size=config.page_size,
+            device_pool=device_pool,
+            host_pool=self.host_pool,
+            storage_backend_type=config.storage_backend,
+            storage_backend_extra_config=config.storage_backend_extra_config,
+            model_name=config.model_name,
+            is_dp_attention_enabled=is_dp_attention_enabled,
+            tp_group=tp_group,
+        )
+        self._pending_mamba_layerwise_cow: dict[int, list[int]] | None = None
+
+    @staticmethod
+    def _page_groups_by_kind(op) -> dict[CacheKind, tuple[list, list]]:
+        src_by_kind = getattr(op, "src_pages_by_kind", None)
+        dst_by_kind = getattr(op, "dst_pages_by_kind", None)
+        if src_by_kind is None or dst_by_kind is None:
+            return {CacheKind.KV: (op.src_pages, op.dst_pages)}
+        groups: dict[CacheKind, tuple[list, list]] = {}
+        for kind in CacheKind:
+            src_pages = src_by_kind.get(kind.value, [])
+            dst_pages = dst_by_kind.get(kind.value, [])
+            groups[kind] = (src_pages, dst_pages)
+        return groups
+
+    def set_mamba_layerwise_cow(
+        self, cow_dst_pages_by_src: dict[int, list[int]] | None
+    ) -> None:
+        self._pending_mamba_layerwise_cow = cow_dst_pages_by_src or None
+
+    def submit_plan(self, plan) -> None:
+        if plan.cache:
+            logger.debug("[cache_op] submit_plan: %s cache ops", len(plan.cache))
+        try:
+            for op in plan.cache:
+                self.submit(op)
+            self.host_exec.flush()
+        finally:
+            self._pending_mamba_layerwise_cow = None
+
+    def submit(self, op) -> None:
+        if isinstance(op, Cache.WriteBackOp):
+            logger.debug(
+                "[cache_op] writeback op_id=%s src_pages=%s dst_pages=%s",
+                op.op_ids,
+                len(op.src_pages),
+                len(op.dst_pages),
+            )
+            groups = self._page_groups_by_kind(op)
+            for i in range(len(op.op_ids)):
+                op_id = op.op_ids[i]
+                is_retract = bool(getattr(op, "is_retract", [False])[i])
+                for kind, (src_groups, dst_groups) in groups.items():
+                    if kind not in self.host_exec.pools:
+                        continue
+                    src_pages = src_groups[i] if i < len(src_groups) else []
+                    dst_pages = dst_groups[i] if i < len(dst_groups) else []
+                    if not src_pages:
+                        continue
+                    if kind == CacheKind.MAMBA:
+                        logger.debug(
+                            "[cache_op][mamba_l2] writeback schedule "
+                            "op_id=%s slots=%s device_slots=%s host_slots=%s "
+                            "is_retract=%s",
+                            op_id,
+                            len(src_pages),
+                            src_pages[:8],
+                            dst_pages[:8],
+                            is_retract,
+                        )
+                    self.host_exec.enqueue_writeback(
+                        op_id,
+                        src_pages,
+                        dst_pages,
+                        is_retract=is_retract,
+                        kind=kind,
+                    )
+                if all(
+                    i >= len(src_groups) or not src_groups[i]
+                    for kind, (src_groups, _) in groups.items()
+                    if kind in self.host_exec.pools
+                ):
+                    self.host_exec.completed_writebacks.append(op_id)
+        elif isinstance(op, Cache.LoadBackOp):
+            logger.debug(
+                "[cache_op] loadback op_id=%s src_pages=%s dst_pages=%s",
+                op.op_ids,
+                len(op.src_pages),
+                len(op.dst_pages),
+            )
+            groups = self._page_groups_by_kind(op)
+            for i in range(len(op.op_ids)):
+                op_id = op.op_ids[i]
+                for kind, (src_groups, dst_groups) in groups.items():
+                    if kind not in self.host_exec.pools:
+                        continue
+                    src_pages = src_groups[i] if i < len(src_groups) else []
+                    dst_pages = dst_groups[i] if i < len(dst_groups) else []
+                    if not src_pages:
+                        continue
+                    if kind == CacheKind.MAMBA:
+                        logger.debug(
+                            "[cache_op][mamba_l2] loadback schedule "
+                            "op_id=%s slots=%s host_slots=%s device_slots=%s",
+                            op_id,
+                            len(src_pages),
+                            src_pages[:8],
+                            dst_pages[:8],
+                        )
+                    loadback_kwargs = {}
+                    mamba_layerwise_cow = getattr(
+                        self, "_pending_mamba_layerwise_cow", None
+                    )
+                    if kind == CacheKind.MAMBA and mamba_layerwise_cow:
+                        loadback_kwargs["layerwise_cow_dst_pages_by_src"] = (
+                            mamba_layerwise_cow
+                        )
+                    self.host_exec.enqueue_loadback(
+                        op_id, src_pages, dst_pages, kind=kind, **loadback_kwargs
+                    )
+
+        elif isinstance(op, Cache.PrefetchOp):
+            logger.debug(
+                "[cache_op] prefetch op_id=%s dst_pages=%s", op.op_id, len(op.dst_pages)
+            )
+            self.storage_exec.submit_prefetch(op)
+        elif isinstance(op, Cache.BackUpOp):
+            logger.debug(
+                "[cache_op] backup op_id=%s src_pages=%s", op.op_id, len(op.src_pages)
+            )
+            self.storage_exec.submit_backup(op)
+        else:
+            raise ValueError("unsupported cache op kind")
+
+    def poll_results(self) -> list:
+        results: list = []
+        results.extend(self.host_exec.drain())
+        results.extend(self.storage_exec.drain())
+        if results:
+            for r in results:
+                logger.debug(
+                    "[cache_op] done op_id=%s success=%s type=%s",
+                    r.op_id,
+                    r.success,
+                    type(r).__name__,
+                )
+        return results
+
+    def get_producer_index(
+        self, kind_or_op_id: CacheKind | str | int, op_id: int | None = None
+    ) -> int | None:
+        return self.host_exec.get_producer_index(kind_or_op_id, op_id)
+
+    def set_consumer(
+        self,
+        kind_or_producer_index: CacheKind | str | int | Iterable[int],
+        producer_index: int | Iterable[int] | None = None,
+    ) -> None:
+        self.host_exec.set_consumer(kind_or_producer_index, producer_index)
+
+    def query_l3_pages(self, hashes: list[str]) -> int:
+        return self.storage_exec.query_exists(hashes)
+
+    def shutdown(self) -> None:
+        self.host_exec.shutdown()
+        self.storage_exec.shutdown()
+
+    def reset(self) -> None:
+        self.host_exec.reset()
+        self.storage_exec.drain()

python/tokenspeed/runtime/cache/executor/storage_executor.py

@@ -0,0 +1,487 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Storage-backed executor for cache prefetch and backup operations."""
+
+from __future__ import annotations
+
+import json
+import threading
+import time
+from concurrent.futures import Future, ThreadPoolExecutor
+from functools import partial
+from queue import Empty, Queue
+
+import torch
+import torch.distributed as dist
+from tokenspeed_scheduler import Cache
+
+from tokenspeed.runtime.cache.executor.host_executor import (
+    page_ids_to_token_indices,
+)
+from tokenspeed.runtime.cache.kvstore_storage import KVStoreStorageConfig
+from tokenspeed.runtime.cache.storage import StorageBackendFactory
+from tokenspeed.runtime.distributed.process_group_manager import _make_all_groups
+from tokenspeed.runtime.layers.attention.kv_cache.mla import MLATokenToKVPool
+from tokenspeed.runtime.utils import get_colorful_logger
+
+logger = get_colorful_logger(__name__)
+
+
+def _parse_storage_backend_extra_config(raw: str | None):
+    extra_config = {}
+    if raw:
+        try:
+            extra_config = json.loads(raw)
+        except json.JSONDecodeError as exc:
+            logger.error("Invalid backend extra config JSON: %s", exc)
+            raise
+
+    prefetch_threshold = extra_config.pop("prefetch_threshold", 256)
+    prefetch_timeout_base = extra_config.pop("prefetch_timeout_base", 1)
+    prefetch_timeout_per_ki_token = extra_config.pop(
+        "prefetch_timeout_per_ki_token", 0.25
+    )
+
+    if not isinstance(prefetch_threshold, int):
+        raise ValueError(
+            f"prefetch_threshold must be int, got {type(prefetch_threshold).__name__}"
+        )
+    if not isinstance(prefetch_timeout_base, (int, float)):
+        raise ValueError(
+            f"prefetch_timeout_base must be number, got {type(prefetch_timeout_base).__name__}"
+        )
+    if not isinstance(prefetch_timeout_per_ki_token, (int, float)):
+        raise ValueError(
+            f"prefetch_timeout_per_ki_token must be number, got {type(prefetch_timeout_per_ki_token).__name__}"
+        )
+
+    return (
+        extra_config,
+        prefetch_threshold,
+        float(prefetch_timeout_base),
+        float(prefetch_timeout_per_ki_token),
+    )
+
+
+def _generate_storage_config(
+    device_pool,
+    host_pool,
+    model_name: str | None,
+    extra_config_dict: dict,
+    is_dp_attention_enabled: bool,
+) -> KVStoreStorageConfig:
+    tp_rank = dist.get_rank() if dist.is_initialized() else 0
+    tp_size = dist.get_world_size() if dist.is_initialized() else 1
+
+    is_mla_backend = isinstance(device_pool, MLATokenToKVPool)
+
+    return KVStoreStorageConfig(
+        tp_rank=tp_rank,
+        tp_size=tp_size,
+        is_mla_model=is_mla_backend,
+        is_page_first_layout=host_pool.layout == "page_first",
+        model_name=model_name,
+        extra_config=extra_config_dict,
+    )
+
+
+class StorageExecutor:
+    """Execute L3 storage prefetch and backup operations asynchronously."""
+
+    def __init__(
+        self,
+        page_size: int,
+        device_pool,
+        host_pool,
+        storage_backend_type: str | None,
+        storage_backend_extra_config: str | None = None,
+        model_name: str | None = None,
+        is_dp_attention_enabled: bool = False,
+        storage_batch_size: int = 128,
+        tp_group=None,
+    ):
+        self.page_size = page_size
+        self.host_pool = host_pool
+        self.storage_batch_size = storage_batch_size
+        (
+            extra_config_dict,
+            _prefetch_threshold,
+            prefetch_timeout_base,
+            prefetch_timeout_per_ki_token,
+        ) = _parse_storage_backend_extra_config(storage_backend_extra_config)
+
+        self.prefetch_timeout_base = prefetch_timeout_base
+        self.prefetch_timeout_per_page = (
+            page_size / 1024 * prefetch_timeout_per_ki_token
+        )
+
+        self.storage_backend = None
+        if storage_backend_type is not None:
+            storage_config = _generate_storage_config(
+                device_pool,
+                host_pool,
+                model_name,
+                extra_config_dict,
+                is_dp_attention_enabled,
+            )
+            try:
+                self.storage_backend = StorageBackendFactory.create_backend(
+                    storage_backend_type, storage_config, host_pool
+                )
+            except ValueError as exc:
+                raise ValueError(f"Failed to create storage backend: {exc}") from exc
+
+            self.storage_backend.register_mem_pool_host(host_pool)
+        self.tp_size = (
+            torch.distributed.get_world_size(group=tp_group)
+            if tp_group is not None
+            else 1
+        )
+        # Dedicated subgroup for kvstore collectives. The shared ``tp_group``
+        # is used by the engine main thread elsewhere (event_loop,
+        # request_handler). Issuing collectives on it from
+        # the aggregator thread would race those callers — different threads
+        # on the same rank issuing on the same group can pair up out-of-order
+        # across ranks and hang.
+        #
+        # ``new_group`` is itself a world-wide collective, so when there is
+        # more than one TP-shaped group in the world (e.g. DP attention with
+        # attn_tp groups [0..3] and [4..7]) we have to call it for *every*
+        # such group in the same deterministic order on every rank — calling
+        # it only with the local rank's TP group would deadlock other ranks.
+        # ``_make_all_groups`` enumerates the full set with the same size and
+        # stride pattern, mirroring ``pg_manager.init_process_group``.
+        self._tp_group = None
+        if tp_group is not None and self.tp_size > 1:
+            my_ranks = tuple(torch.distributed.get_process_group_ranks(tp_group))
+            for g in _make_all_groups(my_ranks):
+                pg = torch.distributed.new_group(ranks=list(g), backend="gloo")
+                if g == my_ranks:
+                    self._tp_group = pg
+        self._results: Queue = Queue()
+        self._prefetch_op_to_rid: dict = {}  # op_id → request_id
+        self._executor: ThreadPoolExecutor | None = None
+        # All collectives on the dedicated kvstore subgroup are funneled
+        # through a single aggregator thread so they issue in deterministic
+        # submit order across ranks (Gloo and NCCL groups both require
+        # callers to agree on issuance order; concurrent issuance from
+        # worker threads can deadlock when ops finish in different order on
+        # each rank).
+        self._aggregator_pending: Queue = Queue()
+        self._aggregator_stop = threading.Event()
+        self._aggregator_thread: threading.Thread | None = None
+        if self.storage_backend is not None:
+            self._executor = ThreadPoolExecutor(
+                max_workers=2,
+                thread_name_prefix="tokenspeed-mem-l3-io",
+            )
+            self._aggregator_thread = threading.Thread(
+                target=self._aggregator_loop,
+                name="tokenspeed-mem-l3-aggr",
+                daemon=True,
+            )
+            self._aggregator_thread.start()
+
+    @property
+    def enabled(self) -> bool:
+        return self.storage_backend is not None
+
+    def submit_prefetch(self, op) -> None:
+        # Extract request_id from the op and remember mapping
+        rid = op.request_id
+        self._prefetch_op_to_rid[op.op_id] = rid
+
+        if not self.enabled:
+            evt = Cache.PrefetchDoneEvent()
+            evt.op_id = op.op_id
+            evt.request_id = rid
+            evt.success = False
+            evt.completed_pages = 0
+            self._results.put(evt)
+            return
+        future = self._executor.submit(self._run_prefetch, op)
+        # Enqueue at submit time (not completion time) so both ranks see the
+        # same aggregator order, which guarantees the per-op TP all_reduce
+        # pairs up correctly.
+        self._aggregator_pending.put(("prefetch", op.op_id, rid, future))
+
+    def submit_backup(self, op) -> None:
+        if not self.enabled:
+            evt = Cache.BackUpDoneEvent()
+            evt.op_id = op.op_id
+            evt.success = False
+            self._results.put(evt)
+            return
+        future = self._executor.submit(self._run_backup, op)
+        future.add_done_callback(partial(self._on_backup_done, op.op_id))
+
+    def _prefetch_deadline(self, n_pages: int) -> float:
+        return (
+            time.monotonic()
+            + self.prefetch_timeout_base
+            + n_pages * self.prefetch_timeout_per_page
+        )
+
+    def _run_prefetch(self, op) -> int:
+        hashes = op.rolling_page_hashes
+        if not hashes:
+            logger.debug("[cache_op] prefetch_exec op_id=%s no hashes, skip", op.op_id)
+            return 0
+        dst_pages = op.dst_pages
+        if len(hashes) != len(dst_pages):
+            raise ValueError(
+                f"prefetch key/page mismatch: {len(hashes)} hashes "
+                f"vs {len(dst_pages)} dst_pages"
+            )
+
+        deadline = self._prefetch_deadline(len(hashes))
+        completed_pages = 0
+
+        for i in range(0, len(hashes), self.storage_batch_size):
+            if time.monotonic() > deadline:
+                logger.warning(
+                    "prefetch op %s timed out after %s/%s pages",
+                    op.op_id,
+                    completed_pages,
+                    len(hashes),
+                )
+                break
+
+            batch_hashes = hashes[i : i + self.storage_batch_size]
+            batch_dst = dst_pages[i : i + len(batch_hashes)]
+            host_indices = page_ids_to_token_indices(batch_dst, self.page_size, "cpu")
+            try:
+                results = self.storage_backend.batch_get_v1(batch_hashes, host_indices)
+            except Exception as exc:
+                logger.warning(
+                    "prefetch op %s batch IO error at offset %s: %s", op.op_id, i, exc
+                )
+                break
+            result_ok = 0
+            for ok in results:
+                if not ok:
+                    break
+                result_ok += 1
+            completed_pages += result_ok
+            if result_ok < len(results):
+                logger.warning(
+                    "prefetch op %s: %s/%s pages missed in batch at offset %s",
+                    op.op_id,
+                    len(results) - result_ok,
+                    len(results),
+                    i,
+                )
+                break
+
+        # NOTE: TP all_reduce moved to the aggregator thread; the worker only
+        # performs storage I/O and returns the local count. See _aggregator_loop.
+        logger.debug(
+            "[cache_op] prefetch_exec op_id=%s completed %s/%s pages (local)",
+            op.op_id,
+            completed_pages,
+            len(hashes),
+        )
+        return completed_pages
+
+    def _run_backup(self, op) -> None:
+        hashes = op.rolling_page_hashes
+        src_pages = op.src_pages
+        total_num_pages = len(hashes)
+        logger.debug(
+            "[cache_op] backup_exec op_id=%s pages=%s", op.op_id, total_num_pages
+        )
+        completed_pages = 0
+        for i in range(0, total_num_pages, self.storage_batch_size):
+            batch_hashes = hashes[i : i + self.storage_batch_size]
+            batch_src = src_pages[i : i + len(batch_hashes)]
+            host_indices = page_ids_to_token_indices(batch_src, self.page_size, "cpu")
+            results = self.storage_backend.batch_set_v1(batch_hashes, host_indices)
+            result_ok = sum(1 for ok in results if ok)
+            completed_pages += result_ok
+            if result_ok < len(results):
+                failed_count = len(results) - result_ok
+                logger.warning(
+                    "backup op %s: %s/%s pages failed in batch at offset %s",
+                    op.op_id,
+                    failed_count,
+                    len(results),
+                    i,
+                )
+                raise RuntimeError(
+                    f"backup op {op.op_id}: {completed_pages}/{total_num_pages} pages ok, "
+                    f"{failed_count} failed in batch at offset {i}"
+                )
+        logger.debug(
+            "[cache_op] backup_exec op_id=%s done, all %s/%s pages ok",
+            op.op_id,
+            completed_pages,
+            total_num_pages,
+        )
+
+    def _aggregator_loop(self) -> None:
+        """Single thread that owns all collectives on ``self._tp_group``.
+
+        ``self._tp_group`` is a dedicated kvstore-only Gloo subgroup, separate
+        from the ``tp_group`` used by the main thread elsewhere in the engine
+        — so this thread's collectives don't race main-thread collectives.
+
+        Both prefetch completions (which need a TP MIN all_reduce on
+        ``completed_pages``) and ``query_exists`` calls are funneled through
+        ``_aggregator_pending`` in submit order. Because both ranks observe
+        the same submit order (deterministic from the C++ scheduler +
+        main-thread engine loop), the aggregator on each rank issues
+        collectives in the same order, so the MIN all_reduces pair correctly
+        across ranks.
+        """
+        while not self._aggregator_stop.is_set():
+            try:
+                item = self._aggregator_pending.get(block=True, timeout=1)
+            except Empty:
+                continue
+            kind = item[0]
+            if kind == "prefetch":
+                _, op_id, rid, future = item
+                evt = Cache.PrefetchDoneEvent()
+                evt.op_id = op_id
+                evt.request_id = self._prefetch_op_to_rid.pop(op_id, rid)
+                # Encode local outcome with a -1 sentinel for "this rank
+                # failed locally". A bare ``continue`` on local failure would
+                # skip the all_reduce below and deadlock peers that *did*
+                # succeed locally — the collective must run in lockstep on
+                # every rank. ReduceOp.MIN propagates the sentinel across all
+                # ranks, so every rank agrees the op failed.
+                try:
+                    local = future.result()
+                except Exception as exc:
+                    logger.error("prefetch op %s local I/O failed: %s", op_id, exc)
+                    local = -1
+                completed_pages = local
+                if self.tp_size > 1:
+                    try:
+                        t = torch.tensor(local, dtype=torch.int)
+                        torch.distributed.all_reduce(
+                            t,
+                            op=torch.distributed.ReduceOp.MIN,
+                            group=self._tp_group,
+                        )
+                        completed_pages = t.item()
+                    except Exception as exc:
+                        logger.warning("prefetch op %s TP sync failed: %s", op_id, exc)
+                        completed_pages = -1
+                if completed_pages < 0:
+                    evt.success = False
+                    evt.completed_pages = 0
+                else:
+                    evt.success = True
+                    evt.completed_pages = completed_pages
+                logger.debug(
+                    "[prefetch_done] op_id=%s request_id=%s success=%s completed_pages=%s",
+                    op_id,
+                    evt.request_id,
+                    evt.success,
+                    evt.completed_pages,
+                )
+                self._results.put(evt)
+            elif kind == "query_exists":
+                _, hashes, result_future = item
+                # Same reasoning as the prefetch branch: never short-circuit
+                # before the collective.
+                try:
+                    local = self._query_exists_local(hashes)
+                except Exception as exc:
+                    logger.error("query_exists local I/O failed: %s", exc)
+                    local = -1
+                total_hit = local
+                if self.tp_size > 1:
+                    try:
+                        t = torch.tensor(local, dtype=torch.int)
+                        torch.distributed.all_reduce(
+                            t,
+                            op=torch.distributed.ReduceOp.MIN,
+                            group=self._tp_group,
+                        )
+                        total_hit = t.item()
+                    except Exception as exc:
+                        logger.warning("query_exists TP sync failed: %s", exc)
+                        total_hit = -1
+                if total_hit < 0:
+                    result_future.set_exception(
+                        RuntimeError("query_exists failed on at least one rank")
+                    )
+                else:
+                    result_future.set_result(total_hit)
+            else:
+                logger.error("unknown aggregator item kind: %s", kind)
+
+    def _on_backup_done(self, op_id: int, future) -> None:
+        evt = Cache.BackUpDoneEvent()
+        evt.op_id = op_id
+        try:
+            future.result()
+            evt.success = True
+        except Exception as exc:
+            evt.success = False
+            logger.error("backup op %s failed: %s", op_id, exc)
+        self._results.put(evt)
+
+    def drain(self) -> list:
+        results = []
+        while True:
+            try:
+                results.append(self._results.get_nowait())
+            except Empty:
+                return results
+
+    def query_exists(self, hashes: list[str]) -> int:
+        if not self.enabled or not hashes:
+            return 0
+        if self.tp_size <= 1 or self._aggregator_thread is None:
+            return self._query_exists_local(hashes)
+        # Route through the aggregator so the all_reduce on ``tp_group`` is
+        # serialized with prefetch-completion all_reduces on the same group.
+        result_future: Future = Future()
+        self._aggregator_pending.put(("query_exists", list(hashes), result_future))
+        return result_future.result()
+
+    def _query_exists_local(self, hashes: list[str]) -> int:
+        total_hit = 0
+        for i in range(0, len(hashes), self.storage_batch_size):
+            batch = hashes[i : i + self.storage_batch_size]
+            hit = self.storage_backend.batch_exists(batch)
+            total_hit += hit
+            if hit < len(batch):
+                break
+        return total_hit
+
+    def shutdown(self) -> None:
+        if self._aggregator_thread is not None:
+            self._aggregator_stop.set()
+            self._aggregator_thread.join(timeout=10)
+            self._aggregator_thread = None
+        if self._executor is not None:
+            self._executor.shutdown(wait=True)
+            self._executor = None
+        if self.storage_backend is not None and hasattr(self.storage_backend, "close"):
+            try:
+                self.storage_backend.close()
+            except Exception:
+                logger.exception("Failed to close storage backend")
+            self.storage_backend = None