[ATOM] accelerate lm cache: load the CPU prefix when the Qwen3.5 HBM hit is off the chunk grid

atom/kv_transfer/offload/hybrid/kimi_k3/connector.py

@@ -799,15 +799,17 @@ def _decide_load_after_alloc(self, seq, ls):
         joint = int(seq.offload_joint.boundary_tokens or 0)
         if joint <= hbm:
             return False, "per_req_cache_state_boundary", hbm, lmc, lmc - hbm, chunk
-        # Where the transfer starts -- NOT where the request may call itself
-        # cached. `allocate` claimed every matched block, not just resumable
-        # ones, so the KV below this is already resident; asking LMCache to
-        # resend it would land a second copy in HBM (`publish_loaded_prefix`
-        # keeps the canonical mapping, fresh blocks stay private). Floored to the
-        # chunk grid by `_joint_kv_boundary`, so aligned whenever `hbm` was.
+        # Where the transfer starts -- not where the request may call itself
+        # cached. After allocate, `claim_tokens` is at or above
+        # `num_cached_tokens` and on the chunk grid. Blocks strictly below
+        # `start` are the shared prefix; resending them would land a second
+        # copy (`publish_loaded_prefix` keeps the canonical mapping). An
+        # unaligned hit has already lowered `num_cached_tokens` to the chunk
+        # floor, and the partial chunk `[floor, hbm)` is private and sits at
+        # `start`, which is what this leg writes.
         start = max(hbm, int(seq.offload_joint.claim_tokens or 0))
-        # The KV leg moves whole chunks and the blocks below `start` are shared,
-        # so an unaligned start cannot be rounded down.
+        # Whole chunks only. Blocks below `start` are shared, so an unaligned
+        # start cannot be rounded down onto them.
         if start % chunk != 0:
             return False, "joint_unaligned_hbm_prefill", start, lmc, lmc - start, chunk
         # Transfer the chunk covering the boundary, claim only the boundary.

atom/model_engine/scheduler.py

@@ -35,7 +35,11 @@
 from atom.config import Config
 from atom.kv_transfer.disaggregation import KVConnectorOutput, kv_config_has_producer
 from atom.metrics.scheduler import SchedulerMetrics
-from atom.model_engine.block_manager import BlockManager, draft_kv_lookahead
+from atom.model_engine.block_manager import (
+    AdmissionProbe,
+    BlockManager,
+    draft_kv_lookahead,
+)
 from atom.model_engine.dynamic_chunking import (
     CALIBRATION_SWEEP_RATIO,
     USEFUL_SHRINK_REFERENCE_CHUNKS,
@@ -2291,7 +2295,7 @@ def _schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]] | None:
                 and self._tier_hit(seq) != probed_tier_kv
             ):
                 hbm_probe = self._probe_hbm_prefix(seq)
-            num_cached_blocks, block_hashes = hbm_probe
+            num_cached_blocks = hbm_probe.num_cached_blocks
             if num_cached_blocks < 0:
                 self.waiting.appendleft(seq)
                 break
@@ -2351,7 +2355,7 @@ def _schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]] | None:
             ):
                 self.waiting.appendleft(seq)
                 break
-            self.block_manager.record_allocation(seq, num_cached_blocks, block_hashes)
+            self.block_manager.record_allocation(seq, hbm_probe)
             if not self.block_manager.allocate(seq, num_cached_blocks):
                 # A state-less joint boundary could not be privatised without
                 # risking a shared decoding sequence's blocks (finding #2).
@@ -2386,7 +2390,17 @@ def _schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]] | None:
             # Guard: PD decode consumer inherits hit from prefill node;
             # don't clobber with the local hit (always 0 on consumer).
             if not seq.prefix_cache_hit_tokens:
-                seq.prefix_cache_hit_tokens = seq.num_cached_tokens
+                # The gated HBM hit, not the chunk floor `allocate` may have
+                # lowered `num_cached_tokens` to. `[floor, hbm)` is refilled
+                # from CPU but was already resident; recording the floor makes
+                # that slice look like a CPU-tier load.
+                recorded = seq.num_cached_tokens
+                if (
+                    seq.offload_joint.boundary_tokens
+                    and seq.offload_joint.hbm_hit_tokens > recorded
+                ):
+                    recorded = seq.offload_joint.hbm_hit_tokens
+                seq.prefix_cache_hit_tokens = recorded
 
             self._lookup_detach_declined.discard(seq.id)
             self._notify_connector_after_prefill_alloc(seq)
@@ -2425,31 +2439,42 @@ def _schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]] | None:
 
             if seq.offload_joint.boundary_tokens:
                 # The KV leg was refused after the state leg was aimed at the
-                # joint boundary, so the state claims a prefix whose KV nobody
-                # will fetch. Disown it; `_decide_load_after_alloc` logs why.
-                logger.debug(
-                    "[JOINT-DISOWN] seq %s: KV leg refused at boundary %d; "
-                    "recomputing from %d",
-                    seq.id,
-                    seq.offload_joint.boundary_tokens,
-                    seq.num_cached_tokens,
-                )
-                # Privatise the claimed prefix (finding #2) so recompute-from-0
-                # cannot tear a shared decode. cancel_state_load does it when a
-                # CPU state load was pending; otherwise disown directly.
-                if seq.offload_joint.load_hash != -1:
-                    disowned = self.block_manager.cancel_state_load(seq)
+                # joint boundary. An unaligned joint resumes at the HBM hit
+                # when that hit is still recoverable; otherwise disown to 0.
+                # `_decide_load_after_alloc` logs why the leg was refused.
+                if self.block_manager.retreat_joint_to_hbm(seq):
+                    logger.debug(
+                        "[JOINT-RETREAT] seq %s: KV leg refused; resuming at %d",
+                        seq.id,
+                        seq.num_cached_tokens,
+                    )
                 else:
-                    disowned = self.block_manager.disown_claimed_prefix(seq)
-                if not disowned:
-                    self.block_manager.deallocate(seq)
-                    self.waiting.appendleft(seq)
-                    break
-                seq.num_cached_tokens = 0
-                # Partial reset by design: the KV/claim spans are left as-is,
-                # matching the pre-record code -- not a full `reset_joint`.
-                seq.offload_joint.boundary_tokens = 0
-                seq.offload_joint.boundary_hash = -1
+                    logger.debug(
+                        "[JOINT-DISOWN] seq %s: KV leg refused at boundary %d; "
+                        "recomputing from %d",
+                        seq.id,
+                        seq.offload_joint.boundary_tokens,
+                        seq.num_cached_tokens,
+                    )
+                    # Privatise the claimed prefix (finding #2) so
+                    # recompute-from-0 cannot tear a shared decode.
+                    # cancel_state_load does it when a CPU state load was
+                    # pending; otherwise disown directly.
+                    if seq.offload_joint.load_hash != -1:
+                        disowned = self.block_manager.cancel_state_load(seq)
+                    else:
+                        disowned = self.block_manager.disown_claimed_prefix(seq)
+                    if not disowned:
+                        self.block_manager.deallocate(seq)
+                        self.waiting.appendleft(seq)
+                        break
+                    seq.num_cached_tokens = 0
+                    # Partial reset by design: the KV/claim spans are left
+                    # as-is, matching the pre-record code -- not a full
+                    # `reset_joint`.
+                    seq.offload_joint.boundary_tokens = 0
+                    seq.offload_joint.boundary_hash = -1
+                    seq.offload_joint.hbm_hit_tokens = 0
 
             if seq.offload_joint.load_hash != -1:
                 # The kept state boundary is in LMCache, not HBM. It parks
@@ -2864,24 +2889,35 @@ def _consume_failed_remote_kv(self, seq: Sequence) -> bool:
         if not self._connector_flag("is_offload"):
             self._uncount_inflight_load(seq)
         if seq.offload_joint.load_hash != -1 or seq.offload_joint.boundary_tokens:
-            # The state never arrived, so the boundary is not this request's
-            # history. Disown it exactly as `BlockManager.allocate` does at
-            # admission, otherwise the resume runs with `num_cached_tokens > 0`
-            # over a group nobody filled -- `has_initial_state` reads that as
-            # "the recurrence continues". Before `offload_loaded_tokens` is set
-            # below, so that field records the disowned figure. Privatise the
+            # The state or the joint KV never arrived. An unaligned joint
+            # resumes at the HBM hit when that hit is still recoverable.
+            # Otherwise disown exactly as `BlockManager.allocate` does:
+            # `num_cached_tokens > 0` would tell `has_initial_state` the
+            # recurrence continues over a group nobody filled. Privatise the
             # claimed prefix (finding #2): the offload resume path reuses this
-            # exact block_table without re-allocating, so recompute-from-0 would
-            # otherwise tear a shared decode.
+            # block_table, so recompute-from-0 would tear a shared decode.
+            if self.block_manager.retreat_joint_to_hbm(seq):
+                seq.offload_loaded = False
+                seq.offload_loaded_tokens = seq.num_cached_tokens
+                seq.offload_load_start_tokens = None
+                if seq.offload_joint.load_hash != -1:
+                    # The hit's own state is a tier load. Stay parked for it;
+                    # this failure was the joint, not that load.
+                    seq.offload_load_failed = False
+                    seq.status = SequenceStatus.WAITING_FOR_REMOTE_KVS
+                    return False
+                # The HBM hit's blocks are already in the table. Mark the
+                # load failed so admission resumes that prefill instead of
+                # calling allocate on a non-empty block table.
+                seq.offload_load_failed = True
+                return True
             disowned = self.block_manager.disown_claimed_prefix(seq)
             seq.num_cached_tokens = 0
             seq.offload_joint.load_hash = -1
-            # A joint boundary that failed on either leg is not this request's
-            # history any more than a state-only one is, and leaving it set
-            # would let the connector clamp a later lookup to it. Partial reset
-            # by design: the KV/claim spans are left as-is, as before.
+            # Partial reset by design: the KV/claim spans are left as-is.
             seq.offload_joint.boundary_tokens = 0
             seq.offload_joint.boundary_hash = -1
+            seq.offload_joint.hbm_hit_tokens = 0
             if not disowned:
                 # Cannot back private copies; drop the whole table so the
                 # re-admission rebuilds a clean one via can_allocate/allocate
@@ -3065,21 +3101,20 @@ def _lookup_detach_enabled(self, seq: Sequence) -> bool:
             or (seq.kv_transfer_params or {}).get("do_remote_prefill")
         )
 
-    def _probe_hbm_prefix(self, seq: Sequence) -> tuple[int, list[int]]:
+    def _probe_hbm_prefix(self, seq: Sequence) -> AdmissionProbe:
         """`can_allocate(record=False)`: the HBM hit in blocks (-1 if the
-        request does not fit) and the hash chain it walked.
+        request does not fit), the hash chain it walked and the demand its fit
+        passed, which `record_allocation` commits.
 
         Done once per admission attempt; the result is shared by the detach
         decision and the admission itself, which run in the same pass with no
         pool change in between.
         """
-        block_hashes: list[int] = []
-        num_cached_blocks = self.block_manager.can_allocate(
-            seq,
-            record=False,
-            block_hashes=block_hashes,
+        probe = AdmissionProbe()
+        probe.num_cached_blocks = self.block_manager.can_allocate(
+            seq, record=False, probe=probe
         )
-        return num_cached_blocks, block_hashes
+        return probe
 
     @staticmethod
     def _tier_hit(seq: Sequence) -> tuple[int, int]:
@@ -3091,17 +3126,15 @@ def _tier_hit(seq: Sequence) -> tuple[int, int]:
             int(getattr(joint, "kv_only_tokens", 0) or 0),
         )
 
-    def _can_detach_tier_lookup(
-        self, seq: Sequence, hbm_probe: tuple[int, list[int]]
-    ) -> bool:
+    def _can_detach_tier_lookup(self, seq: Sequence, hbm_probe: AdmissionProbe) -> bool:
         """Whether admission may proceed without waiting for `seq`'s lookup.
 
         Side-effect free. Requires an HBM hit and the connector's consent
         (`can_detach_lookup`, which applies `OFFLOAD_MIN_LOAD_TOKENS`). The
         probe is the one admission uses, so it sees the HBM prefix
         `_decide_load_after_alloc` will see.
         """
-        hit_blocks = hbm_probe[0]
+        hit_blocks = hbm_probe.num_cached_blocks
         if hit_blocks <= 0:
             return False
         num_hbm_tokens = hit_blocks * self.block_manager.hash_block_size

atom/model_engine/sequence.py

@@ -219,29 +219,39 @@ class OffloadJointRecord:
     kv_only_tokens: int = 0
     # The boundary both legs of a joint load are aimed at, or 0. Chosen by
     # `can_allocate`: the rightmost checkpoint rung the LMCache KV prefix
-    # covers. `num_cached_tokens` stays at the HBM prefix until both legs
-    # report, which keeps the forward honest if either fails.
+    # covers. An aligned hit leaves `num_cached_tokens` at the HBM prefix
+    # until both legs report. An unaligned hit lowers it to the chunk floor
+    # at allocate time and refills `[floor, hbm)` privately. A failed leg
+    # returns to that HBM hit (`hbm_hit_tokens`) when the hit's blocks and
+    # state are still recoverable, and disowns to token 0 only when they
+    # are not.
     boundary_tokens: int = 0
     # Content hash of that boundary's last block, for the state leg.
     boundary_hash: int = -1
     # How far the KV leg transfers, in tokens: the LMCache chunk that *covers*
     # the boundary above, which is at or past it. The request still only claims
     # the boundary -- see `_claim_after_load`.
     kv_tokens: int = 0
-    # How far `allocate` claimed straight out of the HBM prefix cache, in
-    # tokens. Above `num_cached_tokens`, not below it: the blocks in
-    # `(hit, compressed_hit]` hold this prompt's KV and are still indexed --
-    # the state gate cut the hit, the KV never went anywhere. Claiming them is
-    # what keeps the KV leg from paying LMCache to resend what the pool already
-    # has, and it is where that leg starts.
+    # How far `allocate` claimed out of the HBM prefix cache, in tokens, and
+    # where the KV leg starts. After allocate this is at or above
+    # `num_cached_tokens`. A compressed hit past the gated hit claims resident
+    # blocks above the HBM prefix. An unaligned hit lowers `num_cached_tokens`
+    # to this floor, so the partial chunk is private rather than shared. The
+    # blocks in `(hit, compressed_hit]` still hold this prompt's KV; claiming
+    # them is what keeps the KV leg from paying LMCache to resend what the
+    # pool already has.
     claim_tokens: int = 0
+    # The gated HBM hit, in tokens, when `allocate` lowered
+    # `num_cached_tokens` below it to refill a partial chunk. 0 when the hit
+    # was not lowered. A failed joint resumes here instead of at token 0.
+    hbm_hit_tokens: int = 0
 
     def reset_joint(self) -> None:
         """Drop the joint boundary and both of its transfer spans.
 
-        The four joint fields move together -- a boundary is meaningless
-        without the KV span aimed at it and the prefix claimed below it -- so
-        the sites that abandon a joint load in full clear them through here
+        The joint fields move together -- a boundary is meaningless without
+        the KV span aimed at it and the prefix claimed below it -- so the
+        sites that abandon a joint load in full clear them through here
         rather than by hand. Leaves `load_hash` and `kv_prefix_tokens`
         untouched: they are set on their own paths and are not part of the
         boundary this drops.
@@ -250,6 +260,7 @@ def reset_joint(self) -> None:
         self.boundary_hash = -1
         self.kv_tokens = 0
         self.claim_tokens = 0
+        self.hbm_hit_tokens = 0
 
 
 class Sequence:

tests/test_admission_host_overhead.py

@@ -25,7 +25,7 @@
     stateful_seq,
 )
 
-from atom.model_engine.block_manager import BlockManager
+from atom.model_engine.block_manager import AdmissionProbe, BlockManager
 from atom.model_engine.prefix_hash import (
     attach_shipped_chains,
     chained_block_hashes,
@@ -74,11 +74,12 @@ def pool_state(bm):
 
 def admit(bm, seq, chunk=None):
     """Phase 2: probe, record, allocate, forward."""
-    hashes = []
-    hit = bm.can_allocate(seq, record=False, block_hashes=hashes)
+    probe = AdmissionProbe()
+    hit = bm.can_allocate(seq, record=False, probe=probe)
+    hashes = probe.block_hashes
     if hit < 0:
         return (hit, hashes)
-    bm.record_allocation(seq, hit, hashes)
+    bm.record_allocation(seq, probe)
     assert bm.allocate(seq, hit)
     replay = seq.replays_hit, seq.replay_kv_load
     todo = seq.num_tokens - seq.num_cached_tokens

tests/test_offload_skip_lookup.py

@@ -515,17 +515,18 @@ def test_admission_walks_the_prefix_once_when_the_detach_is_probed(monkeypatch):
     monkeypatch.setattr(sched_mod.envs, "OFFLOAD_SKIP_LOOKUP_ON_HBM_HIT", True)
     walks = []
 
-    def can_allocate(seq, *, record, block_hashes):
+    def can_allocate(seq, *, record, probe):
         walks.append(record)
-        block_hashes.extend([11, 12])
+        probe.block_hashes.extend([11, 12])
         return 2
 
     s = Scheduler.__new__(Scheduler)
     s.kv_connector = _Conn()
     s._local_prefill_coalescing = False
     s.block_manager = SimpleNamespace(hash_block_size=4, can_allocate=can_allocate)
     probe = s._probe_hbm_prefix(SimpleNamespace())
-    assert probe == (2, [11, 12]) and walks == [False]
+    assert (probe.num_cached_blocks, probe.block_hashes) == (2, [11, 12])
+    assert walks == [False]
     assert s._can_detach_tier_lookup(SimpleNamespace(), probe)
     assert walks == [False]  # the decision itself does not walk again
 

tests/test_scheduler.py

@@ -1044,6 +1044,43 @@ def test_offload_parked_count_released_only_after_resume(self, seq_factory):
         assert failed.block_table == [0]
         assert sched._num_parked_remote_kv == 0
 
+    def test_a_retreated_joint_resumes_on_its_existing_blocks(self, seq_factory):
+        sched = Scheduler(
+            MockConfig(
+                max_num_seqs=2,
+                max_num_batched_tokens=64,
+                num_kvcache_blocks=100,
+            )
+        )
+        sched.kv_connector = SimpleNamespace(
+            is_offload=True,
+            build_connector_meta=lambda: None,
+        )
+        seq = seq_factory(list(range(32)))
+        seq.num_cached_tokens = 8
+        seq.block_table = [7, 8]
+        seq.offload_joint.boundary_tokens = 256
+        seq.offload_joint.load_hash = 99
+
+        def retreat(retreated):
+            retreated.num_cached_tokens = 16
+            retreated.offload_joint.load_hash = -1
+            retreated.offload_joint.reset_joint()
+            return True
+
+        sched.block_manager.retreat_joint_to_hbm = retreat
+        sched._park_for_remote_load(seq, deque())
+        sched.failed_recving_kv_req_ids.append(seq.id)
+        sched.waiting.append(seq)
+
+        sched.schedule()
+
+        assert seq.offload_load_failed is True
+        assert seq.block_table == [7, 8]
+        assert seq.num_cached_tokens == 16
+        assert seq in sched.running
+        assert sched._num_parked_remote_kv == 0
+
     def test_a_resumed_offload_prefill_reports_the_hit_the_load_gave_it(
         self, seq_factory
     ):