DeepSeek-V4.1-Flash: LMCache tier offload through the MP server (vLLM plugin)

atom/config.py

@@ -2131,6 +2131,11 @@ class Config:
     torch_dtype: torch.dtype = field(init=False)
     speculative_config: SpeculativeConfig | None = None
     kv_transfer_config: dict = field(default_factory=dict)
+    # A frontend's own stable form of its speculative config, which the offload
+    # PAGE namespace hashes in place of `speculative_config`: only the frontend
+    # knows which of its fields differ between processes. None hashes
+    # `speculative_config` itself, as native ATOM does.
+    offload_speculative_config: Any = None
     kv_events_config: KVEventsConfig = field(default_factory=KVEventsConfig.from_env)
     # DSpark runtime knobs. Built once in the parent from --dspark-config (see
     # EngineArgs) and pickled into every worker. Read sites use `config.dspark.*`

atom/kv_transfer/disaggregation/pp_kv_aggregator.py

@@ -19,6 +19,7 @@
     ReqId,
     SaveOperationId,
     SaveSourceGroupId,
+    req_id_of,
 )
 
 
@@ -134,7 +135,7 @@ def ingest(self, pp_rank: int, output: KVConnectorOutput) -> KVConnectorOutput:
         )
 
     def _save_is_abandoned(self, operation) -> bool:
-        return str(getattr(operation, "req_id", operation)) in self._abandoned_saves
+        return req_id_of(operation) in self._abandoned_saves
 
     def forget(self, rid: ReqId) -> None:
         """Retire a request's saves after the scheduler abandons them.
@@ -155,7 +156,7 @@ def forget(self, rid: ReqId) -> None:
             self._abandoned_saves.popitem(last=False)
 
         def _owned_by(completion) -> bool:
-            return str(getattr(completion, "req_id", completion)) == sid
+            return req_id_of(completion) == sid
 
         for key in [k for k in self._saving if _owned_by(k)]:
             self._saving.pop(key, None)

atom/kv_transfer/disaggregation/types.py

@@ -102,6 +102,17 @@ def __post_init__(self) -> None:
 LoadCompletionId = ReqId | LoadOperationId
 
 
+def raw_req_id(completion: Any) -> ReqId:
+    """The request id of a completion id, bare or tagged with a generation, in
+    the type the engine compares against (`seq.id`)."""
+    return getattr(completion, "req_id", completion)
+
+
+def req_id_of(completion: Any) -> str:
+    """`raw_req_id` as a string key."""
+    return str(raw_req_id(completion))
+
+
 @dataclass(frozen=True)
 class StateStoreOperationId:
     """Exact identity of one hand-out of a state checkpoint to the CPU tier.

atom/kv_transfer/offload/README.md

@@ -31,8 +31,12 @@ present (DSV4 PAGE+SLOT), and `dense` otherwise.
 `kv_transfer_config.offload_layout` can override compatible choices without
 giving scheduler and worker different connector names. MiniMax-M3 cannot be
 overridden away from `m3`, because the other codecs do not preserve its NSA
-index cache. A GDN or K3 model cannot be overridden onto a layout that owns
-no state tier (`dense`), because that would restore KV over a stale slot.
+index cache. A model with per-request state (GDN `qwen`, K3 `kimi_k3`,
+DeepSeek-V4.1 `dsv41`) accepts only its own layout as an override: any other
+one, `qwen` and `kimi_k3` for each other included, has no tier for its state
+and would restore KV over a stale slot. DeepSeek-V4.1 resolves to `dsv41` (its
+PAGE+STATE pool, offloaded over `lmcache_mp` only), and no other model may
+name `dsv41`.
 
 It is the **ATOM-native, in-engine** offload path: the connector plugs straight
 into ATOM's scheduler/worker via the shared
@@ -104,7 +108,8 @@ Four rules carry the module:
 | `atom_lmcache_staging.py` | Per-thread CUDA streams, staging buffer, ready/free events, env helpers. |
 | `mp/connector.py` | Capability-selected public `lmcache_mp` worker/scheduler shells. |
 | `mp/deployment.py` | LMCache MP configuration, topology, model namespace and server adapters. |
-| `mp/worker.py`, `mp/scheduler.py` | Generic PAGE-only LMCache MP connector halves. |
+| `mp/worker.py`, `mp/scheduler.py` | LMCache MP connector halves. The worker registers PAGE views, plus checkpoint image groups when the backend publishes them; the vLLM plugin builds it directly. |
+| `mp/transfer_thread.py` | The transfer thread: submits off the forward thread under `OFFLOAD_MP_TRANSFER_THREAD`, and also polls for a worker built with `transfers_off_step`; `_ProducerFences`, the reused producer events. |
 | `mp/lookup.py` | Scheduler-side MP lookups and read-lock bookkeeping. |
 | `mp/transfer.py` | Transfer identity, terminal detection and the fail-stop transfer deadline. |
 | `mp/page_views.py` | Shared validation of backend-published PAGE views. |
@@ -133,7 +138,14 @@ by the attention backend:
   transfer.
 
 Selection does not inspect model names or layout-ID prefixes. A new backend can
-reuse the native path by implementing these shared contracts. The native path
+reuse the native path by implementing these shared contracts.
+
+The vLLM plugin does not go through this shell. For DeepSeek-V4.1 it builds
+`LMCacheMPConnector` (`mp/worker.py`) directly with `transfers_off_step=True`:
+the backend publishes PAGE views plus checkpoint image groups
+(`recurrent_page_groups`), a load restores the image from the PAGE units it
+fills, the transfer thread submits and polls, and vLLM's preemption fence
+(`fence`) waits on that thread's rounds. The native path
 leases existing READY checkpoint PAGE units; it does not snapshot the request's
 live Active SLOT again. Every running request therefore keeps its normal fixed
 SLOT.
@@ -199,6 +211,12 @@ python -m atom.entrypoints.openai_server \
   }'
 ```
 
+`tp_rank_collapse: auto` decides from the config alone, before any backend is
+built, and leaves every state-owning layout (`kimi_k3`, `qwen`, `dsv41`) one
+copy per rank. DeepSeek-V4.1's backend declares its PAGE and STATE identical
+across TP, so `true` stores one copy for the TP group; the worker still refuses
+to register a collapse its backend does not declare.
+
 ATOM's configured chunk size must equal the MP server chunk size, and both must
 align to ATOM's PAGE/hash block size. A native prefix is loadable only where
 PAGE KV and a complete STATE checkpoint exist at the same boundary on every TP
@@ -246,7 +264,10 @@ The model namespace includes PAGE/model geometry, TP and speculation settings,
 native layout/image sizes, the Hugging Face commit when available, and the
 optional `lmcache.mp.model_revision`. Set that revision when weights are
 replaced in an existing local model directory; a directory name alone cannot
-identify changed contents.
+identify changed contents. DeepSeek-V4.1 resolved to `hybrid` before it had a
+layout of its own and keeps the namespace that gave it, so native ATOM's V4.1
+tiers stay reachable across the change. The vLLM plugin, which had no V4.1
+tiers before, also hashes V4.1's `sliding_window` and `candidate_block_size`.
 
 ### MP lifetime and representation
 

atom/kv_transfer/offload/_offload_common.py

@@ -29,6 +29,7 @@
     KVConnectorOutput,
     LoadCompletionId,
     SaveCompletionId,
+    raw_req_id,
 )
 from atom.kv_transfer.offload import config as offcfg
 from atom.utils import envs
@@ -364,9 +365,9 @@ def _guard(self, kind: str, fn, req, **fn_kwargs) -> None:
             logger.exception(
                 "offload %s failed for %s",
                 getattr(fn, "__name__", kind),
-                getattr(req, "req_id", req),
+                raw_req_id(req),
             )
-            rid = getattr(req, "req_id", req)
+            rid = raw_req_id(req)
             if kind == "load":
                 self._on_load_fail(rid)
                 with self._lock:
@@ -516,12 +517,12 @@ def process_completions(self, output: KVConnectorOutput) -> KVConnectorOutput:
         """Apply offload-specific completions and expose plain request IDs."""
 
         loaded = {
-            value.req_id if hasattr(value, "req_id") else value
+            raw_req_id(value)
             for value in output.finished_loading
             if self.load_finished(value) is not False
         }
         failed = {
-            value.req_id if hasattr(value, "req_id") else value
+            raw_req_id(value)
             for value in output.failed_loading
             if self.load_failed(value) is not False
         }
@@ -540,10 +541,10 @@ def process_completions(self, output: KVConnectorOutput) -> KVConnectorOutput:
             # non-terminal milestone such as PAGE source-safety.
             if handled is True:
                 value = completion.operation_id
-                terminal_saves.add(value.req_id if hasattr(value, "req_id") else value)
+                terminal_saves.add(raw_req_id(value))
         for value in output.finished_saving:
             self.save_finished(value)
-            terminal_saves.add(value.req_id if hasattr(value, "req_id") else value)
+            terminal_saves.add(raw_req_id(value))
 
         output.finished_loading = loaded
         output.failed_loading = failed
@@ -810,9 +811,9 @@ def _ensure_lookup_pin(self, seq, sid: str, spec) -> bool:
         A transfer does: it must not read an entry that nothing is holding. So
         the step that turns a load spec into a dispatched retrieve asks the tier
         for real, and the load stands only if the tier still holds at least what
-        the spec promised. If it holds less -- evicted in the steps since -- the
-        load is dropped and the request prefills, which is what would have
-        happened with no tier at all.
+        the spec promised. If it holds less -- evicted in the steps since -- this
+        answers False and the caller drops the load (reopening its saves): the
+        request prefills, which is what would have happened with no tier at all.
 
         Confirming rather than re-deriving is deliberate: by this point the spec
         may have been shaped by something this mixin does not own (the unaligned
@@ -847,7 +848,6 @@ def _ensure_lookup_pin(self, seq, sid: str, spec) -> bool:
                 hit,
                 int(spec.lmcache_cached_tokens),
             )
-            self._clear_pending_load(sid)
             return False
         return True
 
@@ -875,6 +875,19 @@ def _record_failed_load_attempt(self, sid: str) -> None:
         if failed_seq is not None:
             self._load_failed_seqs[sid] = failed_seq
 
+    def _reopen_saves(self, sid: str) -> None:
+        """A tier hit that is not loaded (the tier lost it, or the load failed)
+        is recomputed from the HBM frontier the dispatch recorded: let
+        `[hbm, hit)` be stored again rather than treated as persisted.
+
+        A promise withdrawn before dispatch records none and reopens nothing:
+        the tier still holds what the request recomputes.
+        """
+        floor = self._load_save_floors.get(sid)
+        entry = self._save_tracker.get(sid)
+        if floor is not None and entry is not None:
+            entry[1] = min(int(entry[1]), int(floor))
+
     def _release_failed_load_attempt(self, sid: str, seq) -> None:
         """Drop the mark when its sequence is done with the request ID."""
 

atom/kv_transfer/offload/chunked_scheduler.py

@@ -592,8 +592,10 @@ def _clear_pending_load(self, sid: str) -> None:
 
     def _release_lookup(self, sid: str) -> None:
         """Clear this request's lookup once no load spec needs it."""
-        # clear_lookup_status only clears the client's memo; it does not
-        # release worker pins. Keep the ID for the next metadata dispatch.
+        # The in-process client only clears its memo: the worker pins are
+        # released by the next metadata dispatch, which keeps the ID. The MP
+        # client also frees the lookup's read locks, unless a submitted
+        # retrieve has taken them over (`complete_retrieve`).
         if self._lookup_client is not None:
             try:
                 self._lookup_client.clear_lookup_status(sid)
@@ -663,6 +665,48 @@ def _build_save_request(
             save_operation=operation,
         )
 
+    def _retire_finished_save(self, sid: str, seq) -> None:
+        """Retire a finished request's save that will never be emitted.
+
+        Gives back the suffix leased at teardown, except blocks another
+        operation of this request may still read; `_release_operation_lease`
+        frees those when that operation becomes source-safe.
+        """
+        entry = self._save_tracker.get(sid)
+        if entry is None or entry[0] is not seq:
+            return
+        self._save_tracker.pop(sid, None)
+        lease_key = id(seq)
+        leased = self._save_lease_blocks.get(lease_key)
+        if not leased:
+            return
+        still_read: set[int] = set()
+        for operation, owner in self._save_operation_owner.items():
+            if owner is seq:
+                still_read.update(
+                    self._save_operation_blocks.get(operation, {}).values()
+                )
+        self._return_leased_blocks(lease_key, leased - still_read)
+
+    def _return_leased_blocks(self, lease_key, blocks) -> None:
+        """Give `blocks` back from the request lease `lease_key` names.
+
+        The ones the lease still holds go to the next source-safe release, and
+        the lease is dropped once nothing is left in it. `abandon_save` and
+        `reclaim_stale_leases` drop a whole lease without it.
+        """
+        leased = self._save_lease_blocks.get(lease_key)
+        releasable = set(blocks) & leased if leased is not None else set()
+        if not releasable:
+            return
+        leased.difference_update(releasable)
+        self._pending_source_safe_releases.append(frozenset(releasable))
+        self.total_source_safe_released_blocks += len(releasable)
+        if not leased:
+            self._save_lease_blocks.pop(lease_key, None)
+            self._save_lease_at.pop(lease_key, None)
+            self._save_lease_owner.pop(lease_key, None)
+
     def _build_load_request(
         self,
         seq,
@@ -773,11 +817,16 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata:
                 self._mark_load_skip(seq, reason, hbm, lmc, need, chunk)
                 self._clear_pending_load(sid)
                 continue
+            # The HBM frontier a load that fails from here on reopens saves to.
+            self._load_save_floors[sid] = self._chunk_floor(hbm)
             # This is the one place a retrieve is handed to the worker, so it
             # is where the load's lookup pin has to be live -- a spec derived
             # from a remembered hit has none yet.
             if not self._ensure_lookup_pin(seq, sid, ls):
                 self._mark_load_skip(seq, "tier_lost_prefix", hbm, lmc, need, chunk)
+                # The tier no longer holds the prefix this request recomputes.
+                # `_ensure_lookup_pin` only answers; the drop is done here.
+                self._reopen_saves(sid)
                 self._clear_pending_load(sid)
                 continue
             # num_cached after load = max(HBM, offload); never drop below HBM.
@@ -800,7 +849,6 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata:
                 len(list(seq.block_table)),
             )
             loading_sids.add(sid)
-            self._load_save_floors[sid] = self._chunk_floor(hbm)
             load_operation = self._new_load_operation(seq)
             seq._load_operation = load_operation
             self._active_load_operations[sid] = (seq, load_operation)
@@ -1202,17 +1250,7 @@ def _source_group_finished(self, identity: SaveSourceGroupId) -> None:
             return
         sid = str(operation.req_id)
         owner = self._save_operation_owner.get(operation)
-        lease_key = id(owner) if owner is not None else None
-        leased = self._save_lease_blocks.get(lease_key)
-        releasable = newly_safe & leased if leased is not None else set()
-        if releasable:
-            leased.difference_update(releasable)
-            self._pending_source_safe_releases.append(frozenset(releasable))
-            self.total_source_safe_released_blocks += len(releasable)
-            if not leased:
-                self._save_lease_blocks.pop(lease_key, None)
-                self._save_lease_at.pop(lease_key, None)
-                self._save_lease_owner.pop(lease_key, None)
+        self._return_leased_blocks(id(owner) if owner is not None else None, newly_safe)
         all_safe = set(block_map.values()).issubset(safe)
         if all_safe:
             # Every rank has finished reading every source block of this save,
@@ -1297,16 +1335,7 @@ def _retry_or_retire_failed_save(self, operation: SaveOperationId) -> bool:
         if getattr(entry[0], "_offload_finished", False):
             # Teardown may already have returned the source blocks to the pool.
             # A missing cache entry is safer than a retry against recycled KV.
-            self._save_tracker.pop(sid, None)
-            lease_key = id(entry[0])
-            remaining = self._save_lease_blocks.pop(lease_key, None)
-            self._save_lease_at.pop(lease_key, None)
-            self._save_lease_owner.pop(lease_key, None)
-            if remaining:
-                # Unbound (plugin) teardown also leased the not-yet-emitted
-                # suffix; no further save reads it once the request retires.
-                self._pending_source_safe_releases.append(frozenset(remaining))
-                self.total_source_safe_released_blocks += len(remaining)
+            self._retire_finished_save(sid, entry[0])
             return False
         entry[1] = min(int(entry[1]), previous)
         logger.warning(
@@ -1322,17 +1351,9 @@ def _release_operation_lease(self, operation) -> None:
         block_map = self._save_operation_blocks.pop(operation, {})
         self._save_operation_safe.pop(operation, None)
         owner = self._save_operation_owner.pop(operation, None)
-        lease_key = id(owner) if owner is not None else None
-        leased = self._save_lease_blocks.get(lease_key)
-        releasable = set(block_map.values()) & leased if leased is not None else set()
-        if releasable:
-            leased.difference_update(releasable)
-            self._pending_source_safe_releases.append(frozenset(releasable))
-            self.total_source_safe_released_blocks += len(releasable)
-            if not leased:
-                self._save_lease_blocks.pop(lease_key, None)
-                self._save_lease_at.pop(lease_key, None)
-                self._save_lease_owner.pop(lease_key, None)
+        self._return_leased_blocks(
+            id(owner) if owner is not None else None, block_map.values()
+        )
 
     def _finish_retired_request(self, sid: str) -> None:
         if sid in self._save_retry_blocked:
@@ -1386,13 +1407,8 @@ def load_failed(self, req_id) -> bool:
             # ID cannot complete it (including after request-ID reuse).
             return False
         self._finish_load_statistics(req_id, succeeded=False)
-        floor = self._load_save_floors.get(sid)
-        entry = self._save_tracker.get(sid)
-        if floor is not None and entry is not None:
-            # The LMCache hit was not actually loaded. Let the recomputed
-            # [HBM, LMC) chunks be saved again instead of permanently treating
-            # them as already persisted.
-            entry[1] = self._chunk_floor(floor)
+        # The LMCache hit was not actually loaded.
+        self._reopen_saves(sid)
         self._record_failed_load_attempt(sid)
         # Drop the hit too. `_repeat_load_suppressed` already answers for this
         # request, but the memo must not outlive the load spec it would imply