fix(offload): make budgeted priority save admission mandatory

atom/entrypoints/openai/metrics.py

@@ -375,11 +375,82 @@ def collect(self) -> Iterable[GaugeMetricFamily | CounterMetricFamily]:
                 "Number of tokens saved to LMCache.",
                 offload.get("saved_tokens", 0),
             ),
+            (
+                "atom:lmcache_save_admitted",
+                "Number of LMCache save operations admitted by the scheduler.",
+                offload.get("save_admitted", 0),
+            ),
+            (
+                "atom:lmcache_early_released_blocks",
+                "Number of GPU KV blocks released before save terminal.",
+                offload.get("early_released_blocks", 0),
+            ),
+            (
+                "atom:lmcache_leased_source_blocks",
+                "Number of GPU KV source blocks ever leased to saves.",
+                offload.get("leased_source_blocks", 0),
+            ),
+            (
+                "atom:lmcache_source_safe_released_blocks",
+                "Number of leased GPU KV blocks released at source-safe.",
+                offload.get("source_safe_released_blocks", 0),
+            ),
+            (
+                "atom:lmcache_abnormal_lease_reclaims",
+                "Number of source leases reclaimed after an abnormal timeout.",
+                offload.get("abnormal_lease_reclaims", 0),
+            ),
+            (
+                "atom:lmcache_save_budget_rejected_total",
+                "Number of LMCache save admissions rejected by a hard budget.",
+                offload.get("save_budget_rejected", 0),
+            ),
+            (
+                "atom:lmcache_save_budget_rejected_blocks_total",
+                "Requested PAGE blocks in budget-rejected LMCache saves.",
+                offload.get("save_budget_rejected_blocks", 0),
+            ),
+            (
+                "atom:lmcache_save_budget_evicted_total",
+                "Committed LMCache saves evicted for higher-priority admissions.",
+                offload.get("save_budget_evicted", 0),
+            ),
+            (
+                "atom:lmcache_save_budget_evicted_blocks_total",
+                "Reserved PAGE blocks evicted for higher-priority admissions.",
+                offload.get("save_budget_evicted_blocks", 0),
+            ),
+            (
+                "atom:lmcache_save_oversized_total",
+                "LMCache saves whose source alone exceeded the PAGE budget.",
+                offload.get("save_oversized", 0),
+            ),
         ):
             metric = CounterMetricFamily(name, documentation)
             metric.add_metric([], float(value))
             yield metric
 
+        metric = CounterMetricFamily(
+            "atom:lmcache_save_dropped",
+            "Number of unadmitted LMCache saves dropped by reason.",
+            labels=["reason"],
+        )
+        for reason in ("capacity", "low_value", "terminal_failure", "stale"):
+            metric.add_metric([reason], float(offload.get(f"save_dropped_{reason}", 0)))
+        yield metric
+
+        metric = CounterMetricFamily(
+            "atom:lmcache_save_dropped_tokens",
+            "Number of dirty LMCache save tokens dropped by reason.",
+            labels=["reason"],
+        )
+        for reason in ("capacity", "low_value", "terminal_failure", "stale"):
+            metric.add_metric(
+                [reason],
+                float(offload.get(f"save_dropped_tokens_{reason}", 0)),
+            )
+        yield metric
+
         for name, documentation, value in (
             (
                 "atom:lmcache_loads_pending",
@@ -391,6 +462,76 @@ def collect(self) -> Iterable[GaugeMetricFamily | CounterMetricFamily]:
                 "Number of LMCache saves currently in flight.",
                 offload.get("saves_pending", 0),
             ),
+            (
+                "atom:lmcache_save_candidates",
+                "Number of uncommitted LMCache save candidates.",
+                offload.get("save_candidates", 0),
+            ),
+            (
+                "atom:lmcache_save_candidates_finished",
+                "Number of finished requests still present as save candidates.",
+                offload.get("save_candidates_finished", 0),
+            ),
+            (
+                "atom:lmcache_save_committed",
+                "Number of LMCache saves reserved but not yet dispatched.",
+                offload.get("save_committed", 0),
+            ),
+            (
+                "atom:lmcache_save_candidate_wait_seconds",
+                "Oldest current LMCache save candidate age in seconds.",
+                offload.get("save_candidate_wait_seconds", 0),
+            ),
+            (
+                "atom:lmcache_save_priority_score",
+                "Highest current LMCache save admission score.",
+                offload.get("save_priority_score", 0),
+            ),
+            (
+                "atom:lmcache_save_inflight_wait_seconds",
+                "Oldest current LMCache save operation age in seconds.",
+                offload.get("save_inflight_wait_seconds", 0),
+            ),
+            (
+                "atom:lmcache_save_pin_budget_blocks",
+                "Scheduler-local PAGE block budget for admitted LMCache saves.",
+                offload.get("save_pin_budget_blocks", 0),
+            ),
+            (
+                "atom:lmcache_save_reserved_blocks",
+                "PAGE blocks reserved by committed or active LMCache saves.",
+                offload.get("save_reserved_blocks", 0),
+            ),
+            (
+                "atom:lmcache_save_pinned_blocks",
+                "Unique physical GPU KV blocks held by finished admitted saves.",
+                offload.get("save_pinned_blocks", 0),
+            ),
+            (
+                "atom:lmcache_save_pinned_ratio",
+                "Fraction of the local KV pool reserved or pinned by saves.",
+                offload.get("save_pinned_ratio", 0),
+            ),
+            (
+                "atom:lmcache_save_budget_available_blocks",
+                "Unused scheduler-local LMCache save PAGE budget.",
+                offload.get("save_budget_available_blocks", 0),
+            ),
+            (
+                "atom:lmcache_save_pinned_tokens",
+                "GPU KV token capacity held by finished admitted saves.",
+                offload.get("save_pinned_tokens", 0),
+            ),
+            (
+                "atom:lmcache_deferred_free_requests",
+                "Finished requests awaiting an admitted LMCache save.",
+                offload.get("deferred_free_requests", 0),
+            ),
+            (
+                "atom:lmcache_blocks_waiting_for_store",
+                "Source-safe GPU KV blocks awaiting store terminal.",
+                offload.get("blocks_waiting_for_store", 0),
+            ),
         ):
             metric = GaugeMetricFamily(name, documentation)
             metric.add_metric([], float(value))

atom/kv_transfer/disaggregation/multi/multi_connector.py

@@ -439,6 +439,12 @@ def __init__(self, config: Any) -> None:
         # hosts the state tier -- see `_load_owner` (review round 5, finding 1).
         self._load_winner: dict = {}
 
+    def bind_block_manager(self, block_manager) -> None:
+        for connector in self._connectors:
+            callback = getattr(connector, "bind_block_manager", None)
+            if callable(callback):
+                callback(block_manager)
+
     def _load_owner(self, seq: Any):
         """The sub that armed this request's KV load, or None.
 

atom/kv_transfer/offload/README.md

@@ -878,6 +878,14 @@ Connector-specific tuning (env):
 | `OFFLOAD_MIN_LOAD_TOKENS` | 8192 | Don't reload a hit smaller than this; recompute is cheaper. |
 | `OFFLOAD_COPY_WORKERS` | 1 | SAVE daemon threads. LOAD is always a single thread (TTFT-critical). |
 | `OFFLOAD_MAX_PENDING_SAVES` | `max(2, 2 × OFFLOAD_COPY_WORKERS)` | Positive integer bound on total admitted worker saves (running + queued), acquired before SLOT snapshot or executor submission. |
+| `OFFLOAD_SAVE_MIN_OBSERVED_COUNT` | 2 | Hard minimum rank-local prefix observations for save admission. |
+| `OFFLOAD_SAVE_MAX_PINNED_RATIO` | 0.20 | Maximum fraction of this scheduler's physical KV block pool that committed reservations plus unsafe save leases may occupy. Valid range: 0.0–0.30. It is local per DP scheduler and is never multiplied by TP size. |
+| `OFFLOAD_SAVE_MAX_PINNED_BLOCKS` | — | Optional absolute PAGE-block clamp. The effective local budget is `min(floor(total_blocks × ratio), absolute)`. |
+| `OFFLOAD_SAVE_AGING_WEIGHT` | 0.01 | Priority score added per candidate wait-second. It cannot bypass the minimum observation count. |
+| `OFFLOAD_SAVE_RELEASE_WEIGHT` | 1.0 | Priority weight for blocks released by a finished save. |
+| `OFFLOAD_SAVE_DEMAND_BLOCK_TOKENS` | 8192 | Coarse cumulative-prefix size for the independent rank-local demand tracker. |
+| `OFFLOAD_SAVE_DEMAND_MAX_ENTRIES` | 65536 | Per-rank bound for save-demand history. |
+| `OFFLOAD_SAVE_DEMAND_TTL_SECONDS` | 600 | Quiet-time expiry for save-demand history. |
 | `OFFLOAD_GPU_STAGING_CHUNKS` | 2 | Chunks per bounded GPU staging buffer. Sizes **each** buffer — load and save own separate ones, so resident HBM ≈ `(1 + OFFLOAD_COPY_WORKERS) × chunks × chunk_bytes`. |
 | `OFFLOAD_GPU_STAGING_MAX_BYTES` | — | Hard cap on staging bytes (clamps the chunk count). |
 | `OFFLOAD_RELEASE_GPU_STAGING_AFTER_TRANSFER` | 0 | Free the staging buffer after each transfer (lower idle HBM, higher churn). |
@@ -897,8 +905,26 @@ still falls back safely to per-group ID preparation and blocking host copies.
 `"lmcache.<field>": value` extra. The actual connector extra
 `"slot_sidecar_staging_slots": N` overrides `OFFLOAD_SLOT_STAGING_SLOTS`, and
 `"committed_sidecar_index_capacity": N` overrides
-`OFFLOAD_COMMITTED_SIDECAR_CAPACITY`. `"max_pending_saves": N` overrides
-`OFFLOAD_MAX_PENDING_SAVES`. All apply only to that connector.
+`OFFLOAD_COMMITTED_SIDECAR_CAPACITY`. These apply only to that connector.
+`max_pending_saves` is intentionally process-wide and is configured only with
+`OFFLOAD_MAX_PENDING_SAVES`, so the scheduler and every worker use one bound.
+
+The chunked and DSV4 save lifecycle is explicitly
+`candidate -> committed -> inflight`. A finished candidate must reserve one of
+the pending-save slots immediately or be dropped atomically (PAGE and SLOT
+together); only committed/inflight saves may keep its source blocks alive. An
+inflight save that finishes with a residual tail gets one more admission
+decision, and a failed finished save is dropped rather than retried forever.
+The demand counter is intentionally separate from DP prefix-routing ownership:
+it observes requests at allocation time and does not publish an HBM owner.
+
+Admission also reserves physical block IDs from the scheduler-local
+`BlockManager` against the effective ratio/absolute budget. Shared IDs are
+counted once. Admission may replace one or more lower-scored,
+undispatched committed saves, but it simulates the complete victim set before
+mutating state and never evicts an inflight PAGE or SLOT operation. DSV4 counts
+the full retained request block table because this branch does not yet release
+individual PAGE source ranges before terminal completion.
 
 AOS1 follows the engine's location policy exactly. Submission passes
 `store_location` to `StorageManager.batched_put`; discovery searches only

atom/kv_transfer/offload/_offload_common.py

@@ -566,16 +566,15 @@ def save_abandon_timeout_s(self) -> float:
     def max_pending_saves(self) -> int | None:
         """Running-plus-queued save bound this connector enforces, else None.
 
-        The public read of the per-connector `_max_pending_saves` that
-        `max_pending_saves(kvc, save_workers)` computes from
-        `kv_connector_extra_config` and `OFFLOAD_COPY_WORKERS`. The state leg
+        The public read of the connector's `_max_pending_saves`, computed from
+        the process-wide `OFFLOAD_MAX_PENDING_SAVES` setting (or the default
+        derived from `OFFLOAD_COPY_WORKERS`). The state leg
         (`Scheduler._state_store_pending_cap`) shares this exact number with the
-        KV leg's `_may_emit_save` so both legs pin the same slice of the pool,
-        and honours a per-connector `"max_pending_saves"` override the env reader
-        never sees. None when the connector does not bound its save queue
-        (`_may_emit_save` always True, as on dense) -- the scheduler then falls
-        back to the env reader. Exposed so the scheduler never reaches through
-        the delegating shell's `_impl` for it.
+        KV leg's `_may_emit_save` so both legs pin the same slice of the pool.
+        None when the connector does not bound its save queue (`_may_emit_save`
+        always True, as on dense) -- the scheduler then falls back to the env
+        reader. Exposed so the scheduler never reaches through the delegating
+        shell's `_impl` for it.
         """
         return getattr(self, "_max_pending_saves", None)
 
@@ -744,7 +743,7 @@ def _save_frontier(self, seq) -> int:
     def _has_pending_save(self, seq) -> bool:
         sid = str(seq.id)
         entry = self._save_tracker.get(sid)
-        if entry is None:
+        if entry is None or entry[0] is not seq:
             return False
         return self._save_frontier(seq) > int(entry[1])
 
@@ -755,24 +754,17 @@ def _has_active_load(self, seq) -> bool:
         return active is not None and active[0] is seq
 
 
-def max_pending_saves(kvc, save_workers: int) -> int:
-    """Return the maximum running-plus-queued worker save operations."""
+def max_pending_saves(save_workers: int) -> int:
+    """Return the process-wide running-plus-queued worker save bound."""
 
-    extra = (kvc or {}).get("kv_connector_extra_config", kvc or {}) or {}
-    configured = extra.get("max_pending_saves")
-    if configured is None:
-        configured = os.environ.get(
-            "OFFLOAD_MAX_PENDING_SAVES",
-            str(max(2, 2 * save_workers)),
-        )
-        try:
-            capacity = int(configured)
-        except (TypeError, ValueError) as exc:
-            raise ValueError("max pending saves must be a positive integer") from exc
-    else:
-        if isinstance(configured, bool) or not isinstance(configured, int):
-            raise ValueError("max pending saves must be a positive integer")
-        capacity = configured
+    configured = os.environ.get(
+        "OFFLOAD_MAX_PENDING_SAVES",
+        str(max(2, 2 * save_workers)),
+    )
+    try:
+        capacity = int(configured)
+    except (TypeError, ValueError) as exc:
+        raise ValueError("max pending saves must be a positive integer") from exc
     if capacity <= 0:
         raise ValueError("max pending saves must be a positive integer")
     return capacity