[MiniMax-M3][LMCache] Drive ATOM's byte-level KV offload from the vLLM plugin

atom/kv_transfer/offload/dense/connector.py

@@ -742,6 +742,33 @@ def save_finished(self, req_id) -> None:
         self._save_inflight.pop(sid, None)
         self._finish_save_statistics(req_id)
 
+    def save_finished_by_request(self, req_id) -> None:
+        """Complete a save when only the plain request id is available.
+
+        `save_finished` refuses a raw id once the lifecycle has an exact
+        `SaveOperationId`, so a delayed report cannot complete a newer
+        lifecycle. A vLLM-plugin scheduler cannot satisfy that: vLLM's
+        `KVConnectorOutput` carries request ids as plain strings, so the exact
+        identity never survives the trip back from the worker.
+
+        Resolving the parked identity here keeps the guard meaningful instead of
+        weakening `save_finished` -- and without it the entry never clears, so
+        `_save_inflight` grows for the life of the process and
+        `has_pending_work()` never goes quiet.
+        """
+        sid = str(req_id)
+        active = self._save_inflight.get(sid)
+        self.save_finished(active if active is not None else sid)
+
+    def load_finished_by_request(self, req_id) -> bool:
+        """`load_finished` for a caller that has only the plain request id.
+
+        Same reason as `save_finished_by_request`.
+        """
+        sid = str(req_id)
+        entry = self._active_load_operations.get(sid)
+        return self.load_finished(entry[1] if entry is not None else sid)
+
     def abandon_save(self, req_id) -> None:
         """Force-drop a save the scheduler reclaimed after it stalled.
 

atom/plugin/vllm/attention/minimax_m3_attnetion.py

@@ -340,6 +340,42 @@ def _ensure_fp8_scales(self, kv_cache: torch.Tensor):
             self.v_scale = self.kv_scale[1]
         return self.k_scale, self.v_scale
 
+    def get_kv_transfer_scales(
+        self, kv_cache: "torch.Tensor | None" = None
+    ) -> tuple["torch.Tensor | None", "torch.Tensor | None"]:
+        """The fp8 scales a KV transfer must carry alongside this layer's bytes.
+
+        The sparse cache stores fp8 mantissas whose scale is per token AND per
+        head -- `(num_blocks, num_kv_heads, block_size)`, one fp32 per element
+        of the paged cache. Move the mantissas without them and a restored
+        block is dequantised against whatever the block's previous occupant
+        left behind: fluent-looking garbage, no error anywhere. So any tier
+        that moves this layer's KV has to move these too, and it can only know
+        that by asking -- vLLM's `kv_caches` registration carries the KV
+        tensors alone, and these live on the layer.
+
+        Named after the native path's `get_kv_transfer_tensors`, which reports
+        the same regions off `runner.kv_scale`.
+
+        `kv_cache` overrides the layer's own tensor, for callers that hold it
+        before the layer does -- a connector registering at engine start runs
+        before the first forward, and the scales are allocated lazily. Passing
+        the tensor vLLM registered is also what keeps the allocation stable:
+        `_ensure_fp8_scales` reallocates on a shape or device change, which
+        would strand a pointer a tier had already registered.
+
+        Returns `(None, None)` when the cache is not fp8 -- nothing to carry.
+        """
+        cache = self.kv_cache if kv_cache is None else kv_cache
+        if self.kv_cache_dtype != "fp8":
+            return None, None
+        if cache is None or cache.numel() == 0:
+            raise RuntimeError(
+                f"{self.layer_name}: cannot size the fp8 KV scales before the "
+                "KV cache is allocated"
+            )
+        return self._ensure_fp8_scales(cache)
+
     def _page16_shuffle_cache_for_sparse_kernel(
         self,
     ) -> tuple[torch.Tensor, torch.Tensor, object, object]:

atom/plugin/vllm/kv_transfer/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: MIT
+"""vLLM-plugin side of ATOM's KV transfer paths."""