From 846a103bd9f4a93018ab6191a080c130b4aadde0 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Tue, 15 Sep 2026 09:13:07 +0000 Subject: [PATCH 1/4] fix(offload): bound DSV4 save admission --- .../offload/hybrid/dsv4/connector.py | 52 +++++++- tests/test_lmcache_offload_connector.py | 120 ++++++++++++++---- 2 files changed, 142 insertions(+), 30 deletions(-) diff --git a/atom/kv_transfer/offload/hybrid/dsv4/connector.py b/atom/kv_transfer/offload/hybrid/dsv4/connector.py index bfac4c1d91..0b7ca370f3 100644 --- a/atom/kv_transfer/offload/hybrid/dsv4/connector.py +++ b/atom/kv_transfer/offload/hybrid/dsv4/connector.py @@ -1882,6 +1882,18 @@ def __init__(self, config) -> None: # prefix is stored to LMCache once prefill computes it # (seq.prefix_hashes_published flips True), chunk by chunk. self._save_tracker: dict[str, list] = {} + # The worker has the same bound as a last-resort guard, but admission + # belongs here: rejecting after dispatch lets TP ranks make different + # decisions and forces the scheduler to roll the PAGE watermark back. + # Keep one shared limit for PAGE-only, SLOT-only, and PAGE+SLOT saves. + self._max_pending_saves = max_pending_saves( + kvc, + int(os.environ.get("OFFLOAD_COPY_WORKERS", "1") or 1), + ) + # Resume after the last admitted request when capacity becomes free so + # a long request at the head of the insertion-ordered tracker cannot + # monopolize the bounded save queue. + self._save_rr_last: str | None = None # Scheduler-lifetime completion generation: every emitted save gets a # distinct SaveOperationId, so late TP notifications cannot complete a # later PAGE/SLOT save after request cleanup or request-ID reuse. @@ -2171,6 +2183,25 @@ def _next_save_operation(self, seq) -> SaveOperationId: self._save_nonce += 1 return operation + def _may_emit_save(self) -> bool: + """Return whether the scheduler may dispatch one more save operation. + + A PAGE+SLOT request is one worker operation even though both completion + channels retain its identity. Counting distinct operation IDs keeps the + scheduler's view equal to the worker semaphore while old multi-inflight + state, if any, drains safely. + """ + + operations = { + operation + for inflight in self._save_inflight.values() + for operation in inflight + } + operations.update( + inflight[0] for inflight in self._sidecar_save_inflight.values() + ) + return len(operations) < self._max_pending_saves + def _clear_lookup_status(self, sid: str) -> None: if self._lookup_client is None: return @@ -2365,12 +2396,23 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: # chunked prefill, seq.num_cached_tokens advances after each prefill # chunk's forward has completed; use it as the D2H-safe frontier. chunk = self.chunk_size or 256 - for sid, entry in self._save_tracker.items(): + tracker_sids = list(self._save_tracker.keys()) + if tracker_sids and self._save_rr_last in self._save_tracker: + start = (tracker_sids.index(self._save_rr_last) + 1) % len(tracker_sids) + tracker_sids = tracker_sids[start:] + tracker_sids[:start] + for sid in tracker_sids: if not self._do_save: continue + if not self._may_emit_save(): + break + entry = self._save_tracker[sid] seq, saved = entry if sid in self._reqs_need_recv or sid in loading_sids: continue # loading this step; defer its save + if sid in self._save_inflight or sid in self._sidecar_save_inflight: + # A request's PAGE and SLOT describe one ordered checkpoint + # stream. Do not let a later boundary overtake an earlier save. + continue computed = min( int(getattr(seq, "num_cached_tokens", 0)), int(seq.num_prompt_tokens), @@ -2378,12 +2420,7 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: is_last_prefill = computed >= int(seq.num_prompt_tokens) aligned = (computed // chunk) * chunk sidecar_candidate = self._sidecar_save_candidate(seq, computed) - boundary_needs_page = ( - sidecar_candidate is not None and sidecar_candidate[0] > saved - ) - page_save_due = aligned > saved and ( - sid not in self._save_inflight or boundary_needs_page - ) + page_save_due = aligned > saved if not page_save_due and sidecar_candidate is None: continue slot_save_spec = None @@ -2440,6 +2477,7 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: sidecar_candidate[0], sidecar_candidate[1], ) + self._save_rr_last = sid dispatched = set(meta.lookup_requests_in_step) self._lookup_in_step = [ sid for sid in self._lookup_in_step if sid not in dispatched diff --git a/tests/test_lmcache_offload_connector.py b/tests/test_lmcache_offload_connector.py index 5a4739dd6b..e26df5a071 100644 --- a/tests/test_lmcache_offload_connector.py +++ b/tests/test_lmcache_offload_connector.py @@ -155,6 +155,8 @@ def _scheduler() -> LMCacheOffloadConnectorScheduler: sched._load_save_floors = {} sched._hit_save_floors = {} sched._save_tracker = {} + sched._max_pending_saves = 4 + sched._save_rr_last = None sched._save_nonce = 0 sched._load_nonce = 0 sched._load_lifecycles = {} @@ -2484,6 +2486,14 @@ def test_slot_save_waits_until_post_allocation_state_copy_forward_completes(): destination_slot["value"] = "initialized-from-checkpoint" seq._state_initialized_after_alloc = True + # The state is now eligible, but its SLOT-only publication must wait for + # the earlier PAGE operation for this request to retire. + assert sched.build_connector_meta().requests == [] + assert sched.connector_completion( + _page_completion(before_request.save_operation, succeeded=True) + ) + sched.save_finished(before_request.save_operation) + after_forward = sched.build_connector_meta() after_request = after_forward.requests[0] if after_request.slot_save_spec is not None: @@ -2552,7 +2562,7 @@ def test_sidecar_only_save_emits_when_page_boundary_is_already_stored(): assert request.slot_save_spec == SlotSaveSpec(8192, boundary_hash, 2) -def test_page_save_inflight_still_cuts_and_emits_exact_sidecar_boundary(): +def test_page_save_inflight_cuts_but_defers_exact_sidecar_boundary(): sched = _stateful_scheduler(hit=0) seq = _stateful_seq( req_id=721, @@ -2570,16 +2580,20 @@ def test_page_save_inflight_still_cuts_and_emits_exact_sidecar_boundary(): seq.num_cached_tokens = 8192 meta = sched.build_connector_meta() - assert len(meta.requests) == 1 - request = meta.requests[0] + # Stop exactly at the state boundary, but do not let the PAGE+SLOT save + # overtake the earlier PAGE save for the same request. + assert meta.requests == [] + assert sched._save_tracker["721"][1] == 4096 + assert sched._save_inflight["721"] == {prior_operation} + assert "721" not in sched._sidecar_save_inflight + + sched.save_finished(prior_operation) + request = sched.build_connector_meta().requests[0] assert request.save_spec.skip_leading_tokens == 4096 assert request.token_ids == seq.token_ids[:8192] assert request.slot_save_spec.boundary_tokens == 8192 assert sched._save_tracker["721"][1] == 8192 - assert sched._save_inflight["721"] == { - prior_operation, - request.save_operation, - } + assert sched._save_inflight["721"] == {request.save_operation} assert sched._sidecar_save_inflight["721"][:2] == ( request.save_operation, 8192, @@ -2603,18 +2617,19 @@ def test_save_callbacks_clear_only_matching_operation_generation(): assert page.save_operation == SaveOperationId(seq.id, 0) seq.num_cached_tokens = 8192 + assert sched.build_connector_meta().requests == [] + sched.save_finished(page.save_operation) + boundary = sched.build_connector_meta().requests[0] assert boundary.save_operation == SaveOperationId(seq.id, 1) - assert sched._save_inflight["725"] == { - SaveOperationId(seq.id, 0), - SaveOperationId(seq.id, 1), - } + assert sched._save_inflight["725"] == {SaveOperationId(seq.id, 1)} assert sched._sidecar_save_inflight["725"][0] == SaveOperationId(seq.id, 1) sched.sidecar_save_finished(page.save_operation) assert "725" in sched._sidecar_save_inflight assert sched._committed_sidecar_hashes == set() + # A duplicate late callback for generation 0 cannot retire generation 1. sched.save_finished(page.save_operation) assert sched._save_inflight["725"] == {SaveOperationId(seq.id, 1)} @@ -2893,6 +2908,8 @@ def test_tp_cross_generation_reports_do_not_clear_or_commit_until_matched(): sched._save_tracker["726"] = [seq, 0] page = sched.build_connector_meta().requests[0] seq.num_cached_tokens = 8192 + assert sched.build_connector_meta().requests == [] + sched.save_finished(page.save_operation) boundary = sched.build_connector_meta().requests[0] boundary_hash = boundary.slot_save_spec.boundary_block_hash aggregator = KVOutputAggregator(world_size=2) @@ -2919,10 +2936,7 @@ def test_tp_cross_generation_reports_do_not_clear_or_commit_until_matched(): ) host._update_from_kv_xfer_finished(mixed) - assert sched._save_inflight["726"] == { - page.save_operation, - boundary.save_operation, - } + assert sched._save_inflight["726"] == {boundary.save_operation} assert boundary_hash not in sched._committed_sidecar_hashes matched = aggregator.aggregate( @@ -2961,11 +2975,19 @@ def test_earlier_save_completion_clears_only_its_page_generation(): sched._save_nonce = 1 sched._save_inflight["722"] = {prior_operation} meta = sched.build_connector_meta() - assert meta.requests[0].save_spec.skip_leading_tokens == 4096 - boundary_operation = meta.requests[0].save_operation + assert meta.requests == [] sched.save_finished(prior_operation) + boundary = sched.build_connector_meta().requests[0] + assert boundary.save_spec.skip_leading_tokens == 4096 + boundary_operation = boundary.save_operation + + assert sched._save_inflight["722"] == {boundary_operation} + assert "722" in sched._sidecar_save_inflight + # Replaying the earlier generation after the next one was armed must not + # clear either half of the new PAGE+SLOT operation. + sched.save_finished(prior_operation) assert sched._save_inflight["722"] == {boundary_operation} assert "722" in sched._sidecar_save_inflight @@ -2988,8 +3010,16 @@ def test_collapsed_sidecar_and_save_completion_clears_page_inflight(): sched._save_tracker["723"] = [seq, 4096] prior_operation = SaveOperationId(seq.id, 0) sched._save_nonce = 1 - sched._save_inflight["723"] = {prior_operation} - boundary_operation = sched.build_connector_meta().requests[0].save_operation + boundary_operation = SaveOperationId(seq.id, 1) + # Exercise cleanup of state left by an older scheduler that allowed two + # generations for one request to overlap. New emissions never create this. + sched._save_inflight["723"] = {prior_operation, boundary_operation} + boundary_hash = _chained_prefix_hashes(seq.token_ids, 256)[8192] + sched._sidecar_save_inflight["723"] = ( + boundary_operation, + 8192, + boundary_hash, + ) host = Scheduler.__new__(Scheduler) host.kv_connector = sched host.deferred_free_blocks = {} @@ -3335,12 +3365,10 @@ def test_terminal_sidecar_is_emitted_after_earlier_inflight_boundary_completes() seq.num_cached_tokens = 16_384 sched.request_finished(seq) assert sched.should_defer_free(seq) is True - page_only = sched.build_connector_meta() - assert len(page_only.requests) == 1 - assert page_only.requests[0].slot_save_spec is None - sched.save_finished(page_only.requests[0].save_operation) + assert sched.build_connector_meta().requests == [] sched.sidecar_save_finished(first_operation) + sched.save_finished(first_operation) assert sched.should_defer_free(seq) is True terminal = sched.build_connector_meta() @@ -4374,6 +4402,52 @@ def test_chunked_prefill_save_uses_computed_frontier_and_serializes_inflight(): assert meta3.requests[0].is_last_prefill is True +def test_dsv4_save_admission_is_bounded_and_round_robin(): + sched = _scheduler() + sched._max_pending_saves = 2 + for req_id in (100, 101, 102): + seq = SimpleNamespace( + id=req_id, + token_ids=list(range(8)), + block_table=[3, 4], + num_prompt_tokens=8, + num_cached_tokens=8, + has_per_req_cache=False, + ) + sched._save_tracker[str(req_id)] = [seq, 0] + + first = sched.build_connector_meta() + assert [request.req_id for request in first.requests] == [100, 101] + assert sched._may_emit_save() is False + + # The worker has one slot again. Resume after the last admitted request, + # rather than letting request 100 win the queue repeatedly. + sched.save_finished(first.requests[0].save_operation) + second = sched.build_connector_meta() + assert [request.req_id for request in second.requests] == [102] + assert sched._may_emit_save() is False + + +def test_dsv4_page_and_sidecar_share_one_admission_slot(): + sched = _stateful_scheduler(hit=0) + sched._max_pending_saves = 1 + seq = _stateful_seq( + req_id=103, + num_prompt_tokens=8192, + num_cached_tokens=8192, + group=2, + ) + sched._save_tracker["103"] = [seq, 0] + + request = sched.build_connector_meta().requests[0] + assert request.save_spec is not None + assert request.slot_save_spec is not None + assert sched._save_inflight["103"] == {request.save_operation} + assert sched._sidecar_save_inflight["103"][0] == request.save_operation + # PAGE+SLOT is one worker task, not two units of scheduler capacity. + assert sched._may_emit_save() is False + + def test_finished_saving_releases_deferred_free_with_string_req_id(): class _BlockManager: def __init__(self) -> None: From a8c377a554b4fbaa8f0ccd3a95a972ef929f90e5 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Tue, 15 Sep 2026 18:32:05 +0000 Subject: [PATCH 2/4] fix(offload): make save admission process-wide --- atom/kv_transfer/offload/README.md | 5 ++- atom/kv_transfer/offload/_offload_common.py | 44 ++++++++----------- .../offload/hybrid/dsv4/connector.py | 6 +-- .../offload/hybrid/kimi_k3/connector.py | 3 +- atom/model_engine/scheduler.py | 24 +++++----- tests/test_lmcache_offload_v4_page_slot.py | 29 ++++-------- 6 files changed, 43 insertions(+), 68 deletions(-) diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index d38594e70f..3822e4e694 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -897,8 +897,9 @@ still falls back safely to per-group ID preparation and blocking host copies. `"lmcache.": 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. AOS1 follows the engine's location policy exactly. Submission passes `store_location` to `StorageManager.batched_put`; discovery searches only diff --git a/atom/kv_transfer/offload/_offload_common.py b/atom/kv_transfer/offload/_offload_common.py index 8e6e52cdcd..192d9eb86b 100644 --- a/atom/kv_transfer/offload/_offload_common.py +++ b/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) @@ -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 diff --git a/atom/kv_transfer/offload/hybrid/dsv4/connector.py b/atom/kv_transfer/offload/hybrid/dsv4/connector.py index 0b7ca370f3..b77acf7dfd 100644 --- a/atom/kv_transfer/offload/hybrid/dsv4/connector.py +++ b/atom/kv_transfer/offload/hybrid/dsv4/connector.py @@ -200,7 +200,6 @@ class DSV4OffloadConnector(OffloadWorkerMixin, KVConnectorBase): def __init__(self, config) -> None: self._config = config - kvc = getattr(config, "kv_transfer_config", {}) or {} raw_block_size = config.kv_cache_block_size if isinstance(raw_block_size, bool) or not isinstance(raw_block_size, Integral): # Preserve the public configuration error contract. @@ -228,7 +227,7 @@ def __init__(self, config) -> None: # The ATOM LMCache GPU connector owns per-thread staging streams. # OFFLOAD_COPY_WORKERS tunes the SAVE pool only. n_save_workers = int(os.environ.get("OFFLOAD_COPY_WORKERS", "1")) - self._max_pending_saves = max_pending_saves(kvc, n_save_workers) + self._max_pending_saves = max_pending_saves(n_save_workers) self._save_admission = threading.BoundedSemaphore(self._max_pending_saves) # Terminal PAGE-save outcomes, drained by ``get_finished`` onto # ``DSV4_PAGE_SAVE_CHANNEL``. Every request carrying a ``save_spec`` @@ -1887,8 +1886,7 @@ def __init__(self, config) -> None: # decisions and forces the scheduler to roll the PAGE watermark back. # Keep one shared limit for PAGE-only, SLOT-only, and PAGE+SLOT saves. self._max_pending_saves = max_pending_saves( - kvc, - int(os.environ.get("OFFLOAD_COPY_WORKERS", "1") or 1), + int(os.environ.get("OFFLOAD_COPY_WORKERS", "1") or 1) ) # Resume after the last admitted request when capacity becomes free so # a long request at the head of the insertion-ordered tracker cannot diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py index d32b4f3aba..8669994cd1 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py @@ -610,8 +610,7 @@ def __init__(self, config) -> None: # a slow backend can hold. Same knob and default the DSV4 layout bounds # its worker queue with, read from the other end. self._max_pending_saves = max_pending_saves( - getattr(config, "kv_transfer_config", None) or {}, - int(os.environ.get("OFFLOAD_COPY_WORKERS", "1") or 1), + int(os.environ.get("OFFLOAD_COPY_WORKERS", "1") or 1) ) self._save_inflight_since: dict[str, float] = {} self._save_stalled = False diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index a4d96228e7..f0b703ff17 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -2213,19 +2213,17 @@ def _state_store_pending_cap(self) -> int: """The running-plus-queued cap for state-tier stores. Read off the connector's public `max_pending_saves` -- the canonical - `_offload_common.max_pending_saves` value it computed from - `kv_connector_extra_config` and `OFFLOAD_COPY_WORKERS`, the exact bound - the KV leg's `_may_emit_save` enforces. Sharing that number rather than - re-reading the bare `OFFLOAD_MAX_PENDING_SAVES` default of 2 keeps the - two legs from pinning different amounts of the same pool, and honours a - per-connector `"max_pending_saves"` override the env reader never sees. - Falls back to the env reader for a connector (or test double) that never - bounded its save queue (`max_pending_saves` is None, as on dense) and, - under `kv_connector: multi`, reaches the bound through the state-tier sub - -- `MultiConnectorScheduler` bounds nothing of its own, so without the - sub probe the composite silently caps the state leg at the env default of - 2 while the KV leg keeps the sub's real bound. The public accessor means - neither read reaches through the delegating shell's `_impl` any more. + `_offload_common.max_pending_saves` value it computed from the + process-wide `OFFLOAD_MAX_PENDING_SAVES` setting and + `OFFLOAD_COPY_WORKERS`, and the exact bound the KV leg's + `_may_emit_save` enforces. Falls back to the env reader for a connector + (or test double) that never bounded its save queue (`max_pending_saves` + is None, as on dense) and, under `kv_connector: multi`, reaches the + bound through the state-tier sub -- `MultiConnectorScheduler` bounds + nothing of its own, so without the sub probe the composite can silently + use a different default from the state-tier implementation. The public + accessor means neither read reaches through the delegating shell's + `_impl` any more. """ conn = self.kv_connector cap = getattr(conn, "max_pending_saves", None) diff --git a/tests/test_lmcache_offload_v4_page_slot.py b/tests/test_lmcache_offload_v4_page_slot.py index 6d0fa5dbc4..203416b17a 100644 --- a/tests/test_lmcache_offload_v4_page_slot.py +++ b/tests/test_lmcache_offload_v4_page_slot.py @@ -1713,32 +1713,19 @@ def _unexpected_executor(*args, **kwargs): LMCacheOffloadConnector(config) -@pytest.mark.parametrize("value", [True, False, 0, -1, 1.5, "1.5", "bad"]) -def test_max_pending_saves_rejects_invalid_connector_extra(value): +@pytest.mark.parametrize("value", ["", "0", "-1", "1.5", "true", "false", "bad"]) +def test_max_pending_saves_rejects_invalid_env(monkeypatch, value): + monkeypatch.setenv("OFFLOAD_MAX_PENDING_SAVES", value) with pytest.raises(ValueError, match="max pending saves"): - connector_module.max_pending_saves( - { - "kv_connector_extra_config": { - "max_pending_saves": value, - } - }, - 2, - ) + connector_module.max_pending_saves(2) -def test_max_pending_saves_precedence_and_worker_default(monkeypatch): +def test_max_pending_saves_env_and_worker_default(monkeypatch): monkeypatch.setenv("OFFLOAD_MAX_PENDING_SAVES", "7") - assert connector_module.max_pending_saves({}, 3) == 7 - assert ( - connector_module.max_pending_saves( - {"kv_connector_extra_config": {"max_pending_saves": 5}}, - 3, - ) - == 5 - ) + assert connector_module.max_pending_saves(3) == 7 monkeypatch.delenv("OFFLOAD_MAX_PENDING_SAVES") - assert connector_module.max_pending_saves({}, 3) == 6 - assert connector_module.max_pending_saves({}, 1) == 2 + assert connector_module.max_pending_saves(3) == 6 + assert connector_module.max_pending_saves(1) == 2 def test_wait_for_publication_times_out_with_bounded_positive_sleeps(): From 8cad8a4302868b8a6c7197aa28c119d186ecd1ab Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Mon, 21 Sep 2026 07:36:42 +0000 Subject: [PATCH 3/4] feat(offload): add budgeted priority save queue --- atom/entrypoints/openai/metrics.py | 141 +++ .../disaggregation/multi/multi_connector.py | 6 + atom/kv_transfer/offload/README.md | 26 + atom/kv_transfer/offload/_offload_common.py | 2 +- atom/kv_transfer/offload/chunked_scheduler.py | 1035 +++++++++++++++-- atom/kv_transfer/offload/connector.py | 5 + .../offload/hybrid/dsv4/connector.py | 1013 ++++++++++++++-- .../offload/hybrid/kimi_k3/connector.py | 3 +- atom/kv_transfer/offload/save_admission.py | 243 ++++ atom/model_engine/block_manager.py | 11 + atom/model_engine/llm_engine.py | 38 + atom/model_engine/scheduler.py | 3 + docs/environment_variables.md | 9 + tests/entrypoints/test_metrics.py | 108 +- tests/test_lmcache_offload_connector.py | 567 ++++++++- 15 files changed, 3022 insertions(+), 188 deletions(-) create mode 100644 atom/kv_transfer/offload/save_admission.py diff --git a/atom/entrypoints/openai/metrics.py b/atom/entrypoints/openai/metrics.py index 0f0136c455..7031159a52 100644 --- a/atom/entrypoints/openai/metrics.py +++ b/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)) diff --git a/atom/kv_transfer/disaggregation/multi/multi_connector.py b/atom/kv_transfer/disaggregation/multi/multi_connector.py index 901211d58b..12deb50424 100644 --- a/atom/kv_transfer/disaggregation/multi/multi_connector.py +++ b/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. diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index 3822e4e694..79a8bc593c 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -878,6 +878,15 @@ 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_POLICY` | `round_robin` | Chunked LMCache save selection, including `lmcache_mp` and DSV4. `priority` uses reuse demand and dirty cost, and hard-drops a finished request's unadmitted save instead of pinning its KV indefinitely. | +| `OFFLOAD_SAVE_MIN_OBSERVED_COUNT` | 2 | Priority mode's hard minimum rank-local prefix observations. | +| `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). | @@ -901,6 +910,23 @@ still falls back safely to per-group ID preparation and blocking host copies. `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. +In priority mode 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. Priority 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 `retrieve_locations`, and reads target the backend location returned by that diff --git a/atom/kv_transfer/offload/_offload_common.py b/atom/kv_transfer/offload/_offload_common.py index 192d9eb86b..a07365d2c8 100644 --- a/atom/kv_transfer/offload/_offload_common.py +++ b/atom/kv_transfer/offload/_offload_common.py @@ -743,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]) diff --git a/atom/kv_transfer/offload/chunked_scheduler.py b/atom/kv_transfer/offload/chunked_scheduler.py index 81214afe32..9a2c8560c3 100644 --- a/atom/kv_transfer/offload/chunked_scheduler.py +++ b/atom/kv_transfer/offload/chunked_scheduler.py @@ -8,6 +8,8 @@ import logging import os import time +from dataclasses import dataclass, replace +from math import ceil from atom.kv_transfer.disaggregation.base import KVConnectorSchedulerBase from atom.kv_transfer.disaggregation.types import ( @@ -28,6 +30,12 @@ LoadSpec, SaveSpec, ) +from atom.kv_transfer.offload.save_admission import ( + PrefixDemandKey, + PrefixDemandTracker, + SaveBlockReservation, + load_save_admission_config, +) logger = logging.getLogger("atom") @@ -35,6 +43,24 @@ DENSE_PAGE_STORE_CHANNEL = "dense.page.store" +@dataclass(frozen=True) +class _ChunkedSaveCandidate: + sid: str + seq: object + generation: int + saved: int + aligned: int + observed_count: int + reusable_tokens: int + enqueued_at: float + finished: bool + held_blocks: int + + @property + def dirty_tokens(self) -> int: + return max(0, self.aligned - self.saved) + + class ChunkedOffloadSchedulerBase(OffloadSchedulerMixin, KVConnectorSchedulerBase): """Transport- and layout-neutral policy for chunk-aligned KV offload.""" @@ -102,6 +128,46 @@ def __init__( # save. Subclasses may bound the number of outstanding saves through # `_may_emit_save`; resuming after this sid prevents starvation. self._save_rr_last: str | None = None + save_admission = load_save_admission_config() + self._save_policy = save_admission.policy + self._save_min_observed_count = save_admission.min_observed_count + self._save_aging_weight = save_admission.aging_weight + self._save_release_weight = save_admission.release_weight + self._save_max_pinned_ratio = save_admission.max_pinned_ratio + self._save_max_pinned_blocks = save_admission.max_pinned_blocks + self._save_pin_total_blocks = 0 + self._save_pin_budget_blocks: int | None = None + self._save_block_reservations: dict[ + str | SaveOperationId, SaveBlockReservation + ] = {} + self._prefix_demand = PrefixDemandTracker( + block_tokens=save_admission.demand_block_tokens, + max_entries=save_admission.demand_max_entries, + ttl_seconds=save_admission.demand_ttl_seconds, + ) + self._save_demand_keys: dict[ + str, tuple[object, tuple[PrefixDemandKey, ...]] + ] = {} + self._save_candidate_since: dict[str, tuple[object, float]] = {} + self._save_candidate_generation: dict[str, tuple[object, int]] = {} + self._save_candidate_nonce = 0 + self._save_committed: dict[str, object] = {} + self._finished_save_requests: dict[str, object] = {} + self._finished_save_failed: set[str] = set() + self._save_inflight_since: dict[SaveOperationId, float] = {} + self.total_save_admitted = 0 + self.total_save_budget_rejected = 0 + self.total_save_budget_rejected_blocks = 0 + self.total_save_budget_evicted = 0 + self.total_save_budget_evicted_blocks = 0 + self.total_save_oversized = 0 + self._save_drop_totals = { + "capacity": 0, + "low_value": 0, + "terminal_failure": 0, + "stale": 0, + } + self._save_drop_token_totals = dict.fromkeys(self._save_drop_totals, 0) # sid -> exact save generation. Exact matching prevents a delayed TP # notification for an older request lifecycle from releasing the # current request's deferred blocks. @@ -142,6 +208,31 @@ def __init__( os.environ.get("OFFLOAD_MIN_LOAD_TOKENS"), ) self._min_load_tokens = 8192 + try: + self._min_save_tokens = max( + 0, int(os.environ.get("OFFLOAD_MIN_SAVE_TOKENS", "8192")) + ) + except ValueError: + logger.warning( + "LMCache offload scheduler: invalid OFFLOAD_MIN_SAVE_TOKENS=%r; " + "using 8192", + os.environ.get("OFFLOAD_MIN_SAVE_TOKENS"), + ) + self._min_save_tokens = 8192 + + def bind_block_manager(self, block_manager) -> None: + if self._block_manager is not None and self._block_manager is not block_manager: + raise RuntimeError("offload scheduler is already bound to a block manager") + self._block_manager = block_manager + total_blocks = int(block_manager.total_allocatable_kv_blocks) + if total_blocks <= 0: + raise ValueError("BlockManager must expose a positive KV block capacity") + ratio_budget = int(total_blocks * self._save_max_pinned_ratio) + absolute = self._save_max_pinned_blocks + self._save_pin_total_blocks = total_blocks + self._save_pin_budget_blocks = ( + ratio_budget if absolute is None else min(ratio_budget, absolute) + ) # -- match: how many extra tokens can come from CPU/NVMe ------------- def _begin_load_lifecycle(self, seq) -> None: @@ -247,6 +338,27 @@ def update_state_after_alloc(self, seq) -> None: entry = self._save_tracker.get(sid) if entry is None or entry[0] is not seq: self._save_tracker[sid] = [seq, initial_saved] + now = time.monotonic() + demand_keys = ( + self._prefix_demand.observe( + seq.token_ids, + int(seq.num_prompt_tokens), + now, + ) + if self._save_policy == "priority" + else () + ) + self._save_demand_keys[sid] = (seq, demand_keys) + self._save_candidate_generation[sid] = ( + seq, + self._save_candidate_nonce, + ) + self._save_candidate_nonce += 1 + previous_committed = self._save_committed.pop(sid, None) + if previous_committed is not None: + self._release_save_reservation(sid, owner=previous_committed) + self._finished_save_requests.pop(sid, None) + self._finished_save_failed.discard(sid) else: entry[1] = max(int(entry[1]), initial_saved) @@ -303,7 +415,584 @@ def _may_emit_save(self) -> bool: """Return whether another save may be emitted this scheduler step.""" return True + def _ensure_save_admission_state(self) -> None: + """Supply compatibility defaults for lightweight/manual schedulers.""" + + if not hasattr(self, "_save_policy"): + self._save_policy = "round_robin" + if not hasattr(self, "_save_min_observed_count"): + self._save_min_observed_count = 2 + if not hasattr(self, "_save_aging_weight"): + self._save_aging_weight = 0.01 + if not hasattr(self, "_save_release_weight"): + self._save_release_weight = 1.0 + if not hasattr(self, "_save_max_pinned_ratio"): + self._save_max_pinned_ratio = 0.20 + if not hasattr(self, "_save_max_pinned_blocks"): + self._save_max_pinned_blocks = None + if not hasattr(self, "_save_pin_total_blocks"): + self._save_pin_total_blocks = 0 + if not hasattr(self, "_save_pin_budget_blocks"): + self._save_pin_budget_blocks = None + if not hasattr(self, "_save_block_reservations"): + self._save_block_reservations = {} + if not hasattr(self, "_save_lease_blocks"): + self._save_lease_blocks = {} + if not hasattr(self, "_prefix_demand"): + self._prefix_demand = PrefixDemandTracker() + for name in ( + "_save_demand_keys", + "_save_candidate_since", + "_save_candidate_generation", + "_save_committed", + "_finished_save_requests", + "_save_inflight_since", + ): + if not hasattr(self, name): + setattr(self, name, {}) + if not hasattr(self, "_save_candidate_nonce"): + self._save_candidate_nonce = 0 + if not hasattr(self, "_finished_save_failed"): + self._finished_save_failed = set() + if not hasattr(self, "total_save_admitted"): + self.total_save_admitted = 0 + for name in ( + "total_save_budget_rejected", + "total_save_budget_rejected_blocks", + "total_save_budget_evicted", + "total_save_budget_evicted_blocks", + "total_save_oversized", + ): + if not hasattr(self, name): + setattr(self, name, 0) + if not hasattr(self, "_save_drop_totals"): + self._save_drop_totals = { + "capacity": 0, + "low_value": 0, + "terminal_failure": 0, + "stale": 0, + } + if not hasattr(self, "_save_drop_token_totals"): + self._save_drop_token_totals = dict.fromkeys(self._save_drop_totals, 0) + + def _reserved_save_resources_allow(self, committed_after: int) -> bool: + """Layout-specific admission hook evaluated before any mutation.""" + + del committed_after + return True + + def _save_count_allows(self, committed_after: int) -> bool: + limit = getattr(self, "_max_pending_saves", None) + return limit is None or len(self._save_inflight) + committed_after <= limit + + def _pinned_save_block_ids(self) -> set[int]: + pinned: set[int] = set() + for blocks in self._save_lease_blocks.values(): + pinned.update(blocks) + return pinned + + def _save_block_usage( + self, + reservations: dict[str | SaveOperationId, SaveBlockReservation] | None = None, + ) -> tuple[int, int, int]: + """Return ``(reserved, pinned, total_budget_usage)`` without double count.""" + + current = ( + self._save_block_reservations if reservations is None else reservations + ) + pinned_ids = self._pinned_save_block_ids() + reserved_ids: set[int] = set() + estimated = 0 + for reservation in current.values(): + reserved_ids.update(reservation.block_ids) + estimated += int(reservation.estimated_blocks) + reserved = len(reserved_ids - pinned_ids) + estimated + return reserved, len(pinned_ids), len(pinned_ids | reserved_ids) + estimated + + def _candidate_block_reservation( + self, + candidate: _ChunkedSaveCandidate, + *, + priority_score: float, + ) -> SaveBlockReservation: + source_block_size = int(getattr(self, "virtual_block_size", self.block_size)) + start_block = candidate.saved // source_block_size + end_block = ceil(candidate.aligned / source_block_size) + required = max(0, end_block - start_block) + table = list(getattr(candidate.seq, "block_table", ())) + # A frozen table records old identities only. Once early deallocation + # clears the live table those IDs may already belong to another request, + # so retain a conservative count rather than a false exact reservation. + if candidate.finished and not table: + block_ids = frozenset() + estimated = required + else: + available = table[start_block : min(end_block, len(table))] + block_ids = frozenset( + physical_id + for block_id in available + if (physical_id := int(block_id)) >= 0 + ) + estimated = max(0, required - len(available)) + return SaveBlockReservation( + sid=candidate.sid, + seq=candidate.seq, + generation=candidate.generation, + block_ids=block_ids, + estimated_blocks=estimated, + priority_score=float(priority_score), + committed=True, + ) + + @staticmethod + def _reservation_footprint(reservation: SaveBlockReservation) -> int: + return len(reservation.block_ids) + int(reservation.estimated_blocks) + + def _reservation_fits( + self, + reservation: SaveBlockReservation, + *, + victim_keys: tuple[str, ...] = (), + ) -> bool: + projected = dict(self._save_block_reservations) + for key in victim_keys: + projected.pop(key, None) + projected[reservation.sid] = reservation + committed_after = sum(item.committed for item in projected.values()) + if not self._save_count_allows(committed_after): + return False + if not self._reserved_save_resources_allow(committed_after): + return False + budget = self._save_pin_budget_blocks + return budget is None or self._save_block_usage(projected)[2] <= budget + + def _try_reserve_save_candidate( + self, + candidate: _ChunkedSaveCandidate, + *, + allow_eviction: bool, + now: float | None = None, + ) -> bool: + """Atomically reserve PAGE/count/state budget for one candidate. + + Victims are simulated first and are restricted to lower-priority, + undispatched committed saves owned by this same scheduler. Inflight + operation reservations are never considered. + """ + + existing = self._save_block_reservations.get(candidate.sid) + if ( + existing is not None + and existing.seq is candidate.seq + and existing.generation == candidate.generation + and existing.committed + ): + return True + + score = self._candidate_priority( + candidate, + time.monotonic() if now is None else now, + ) + reservation = self._candidate_block_reservation(candidate, priority_score=score) + victims: list[SaveBlockReservation] = [] + if not self._reservation_fits(reservation): + if allow_eviction: + eligible = sorted( + ( + item + for key, item in self._save_block_reservations.items() + if isinstance(key, str) + and item.committed + and item.priority_score < reservation.priority_score + and self._save_committed.get(key) is item.seq + and key not in self._save_inflight + and (entry := self._save_tracker.get(key)) is not None + and entry[0] is item.seq + ), + key=lambda item: ( + item.priority_score, + item.sid, + item.generation, + ), + ) + for victim in eligible: + victims.append(victim) + if self._reservation_fits( + reservation, + victim_keys=tuple(item.sid for item in victims), + ): + break + else: + victims = [] + if not victims and not self._reservation_fits(reservation): + footprint = self._reservation_footprint(reservation) + self.total_save_budget_rejected += 1 + self.total_save_budget_rejected_blocks += footprint + budget = self._save_pin_budget_blocks + if budget is not None and footprint > budget: + self.total_save_oversized += 1 + return False + + # Every condition was checked against the same simulated state. Apply + # the replacement only now so a candidate that still cannot fit never + # destroys useful committed work. + for victim in victims: + self.total_save_budget_evicted += 1 + self.total_save_budget_evicted_blocks += self._reservation_footprint(victim) + self.drop_unadmitted_save(victim.seq, reason="capacity") + self._save_committed[candidate.sid] = candidate.seq + self._save_block_reservations[candidate.sid] = reservation + self.total_save_admitted += 1 + return True + + def _release_save_reservation( + self, key: str | SaveOperationId, *, owner=None + ) -> SaveBlockReservation | None: + reservation = self._save_block_reservations.get(key) + if reservation is None or (owner is not None and reservation.seq is not owner): + return None + return self._save_block_reservations.pop(key) + + def _candidate_generation_for(self, sid: str, seq) -> int: + record = self._save_candidate_generation.get(sid) + if record is not None and record[0] is seq: + return record[1] + generation = self._save_candidate_nonce + self._save_candidate_nonce += 1 + self._save_candidate_generation[sid] = (seq, generation) + return generation + + def _candidate_since_for(self, sid: str, seq, now: float) -> float: + record = self._save_candidate_since.get(sid) + if record is not None and record[0] is seq: + return record[1] + self._save_candidate_since[sid] = (seq, now) + return now + + def _build_save_candidate( + self, sid: str, *, now: float | None = None + ) -> _ChunkedSaveCandidate | None: + self._ensure_save_admission_state() + entry = self._save_tracker.get(sid) + if entry is None: + return None + seq, saved = entry + aligned = self._save_frontier(seq) + if aligned <= int(saved): + return None + candidate_now = time.monotonic() if now is None else float(now) + demand = self._save_demand_keys.get(sid) + observed_count, reusable_tokens = (0, 0) + if demand is not None and demand[0] is seq: + observed_count, reusable_tokens = self._prefix_demand.heat( + demand[1], + max_tokens=aligned, + now=candidate_now, + ) + finished = self._finished_save_requests.get(sid) is seq + return _ChunkedSaveCandidate( + sid=sid, + seq=seq, + generation=self._candidate_generation_for(sid, seq), + saved=int(saved), + aligned=aligned, + observed_count=observed_count, + reusable_tokens=reusable_tokens, + enqueued_at=self._candidate_since_for(sid, seq, candidate_now), + finished=finished, + held_blocks=( + len( + getattr( + seq, + "_offload_finished_block_ids", + getattr(seq, "block_table", ()), + ) + ) + if finished + else 0 + ), + ) + + def _candidate_priority( + self, candidate: _ChunkedSaveCandidate, now: float + ) -> float: + chunk = max(1, int(self.chunk_size or 256)) + cost_units = max(1, ceil(candidate.dirty_tokens / chunk)) + reusable = max(candidate.reusable_tokens, chunk) + benefit = candidate.observed_count * reusable + release_bonus = ( + self._save_release_weight + * candidate.held_blocks + * max(1, int(getattr(self, "virtual_block_size", self.block_size))) + ) + age = max(0.0, now - candidate.enqueued_at) + return (benefit + release_bonus) / cost_units + self._save_aging_weight * age + + def _candidate_sort_key( + self, candidate: _ChunkedSaveCandidate, now: float + ) -> tuple[float, int, float, str, int]: + return ( + -self._candidate_priority(candidate, now), + candidate.dirty_tokens, + candidate.enqueued_at, + candidate.sid, + candidate.generation, + ) + + def _meets_save_value_threshold(self, candidate: _ChunkedSaveCandidate) -> bool: + return candidate.observed_count >= self._save_min_observed_count + + def _commit_finished_save(self, candidate: _ChunkedSaveCandidate) -> bool: + if not candidate.finished or not self._meets_save_value_threshold(candidate): + return False + entry = self._save_tracker.get(candidate.sid) + if entry is None or entry[0] is not candidate.seq: + return False + return self._try_reserve_save_candidate( + candidate, + allow_eviction=self._save_policy == "priority", + ) + + def _forget_save_candidate(self, sid: str, seq) -> None: + for mapping in ( + self._save_demand_keys, + self._save_candidate_since, + self._save_candidate_generation, + ): + record = mapping.get(sid) + if record is not None and record[0] is seq: + mapping.pop(sid, None) + + def drop_unadmitted_save(self, seq, *, reason: str) -> bool: + """Atomically remove a candidate that no worker can be reading.""" + + sid = str(seq.id) + if sid in self._save_inflight: + return False + entry = self._save_tracker.get(sid) + if entry is None or entry[0] is not seq: + return False + candidate = self._build_save_candidate(sid) + dropped_tokens = 0 if candidate is None else candidate.dirty_tokens + self._save_tracker.pop(sid, None) + if self._save_committed.get(sid) is seq: + self._save_committed.pop(sid, None) + self._release_save_reservation(sid, owner=seq) + if self._finished_save_requests.get(sid) is seq: + self._finished_save_requests.pop(sid, None) + self._finished_save_failed.discard(sid) + self._forget_save_candidate(sid, seq) + if reason not in self._save_drop_totals: + reason = "stale" + self._save_drop_totals[reason] += 1 + self._save_drop_token_totals[reason] += dropped_tokens + logger.info( + "LMCache offload: dropped unadmitted save req=%s reason=%s tokens=%d", + sid, + reason, + dropped_tokens, + ) + return True + + def _settle_finished_save(self, sid: str, *, failed: bool = False) -> None: + seq = self._finished_save_requests.get(sid) + if seq is None or sid in self._save_inflight or sid in self._save_committed: + return + if failed: + self._finished_save_failed.add(sid) + if sid in self._finished_save_failed: + self.drop_unadmitted_save(seq, reason="terminal_failure") + return + candidate = self._build_save_candidate(sid) + if candidate is None: + entry = self._save_tracker.get(sid) + if entry is not None and entry[0] is seq: + self._save_tracker.pop(sid, None) + self._finished_save_requests.pop(sid, None) + self._finished_save_failed.discard(sid) + self._forget_save_candidate(sid, seq) + return + if self._commit_finished_save(candidate): + return + reason = ( + "low_value" + if not self._meets_save_value_threshold(candidate) + else "capacity" + ) + self.drop_unadmitted_save(seq, reason=reason) + + def _new_load_operation(self, seq) -> LoadOperationId: + operation = LoadOperationId(seq.id, self._load_nonce) + self._load_nonce += 1 + return operation + + def _build_save_request( + self, + seq, + saved: int, + aligned: int, + operation: SaveOperationId, + block_ids: list[int], + is_last_prefill: bool, + ) -> LMCacheReqMeta | None: + """Admit any layout-specific sources before advancing the watermark.""" + return LMCacheReqMeta( + req_id=seq.id, + token_ids=list(seq.token_ids[:aligned]), + block_ids=block_ids, + save_spec=SaveSpec(skip_leading_tokens=saved, can_save=True), + is_last_prefill=is_last_prefill, + save_operation=operation, + ) + + def _late_save_frontier(self, seq, saved: int, available: int) -> int: + """Return the largest layout-valid boundary for a late-acquired source.""" + del seq, saved + return self._chunk_floor(available) + + def _late_save_source( + self, seq, saved: int, aligned: int + ) -> tuple[int, list[int], frozenset[int]] | None: + """Reacquire a finished request's still-resident prefix at admission.""" + if self._block_manager is None: + raise RuntimeError("late offload save requires a bound block manager") + block_ids, available, claimed = self._block_manager.acquire_offload_prefix( + seq, saved, aligned + ) + target = self._late_save_frontier(seq, saved, available) + source_block_size = int(getattr(self, "virtual_block_size", self.block_size)) + keep_count = max(0, (target - saved) // source_block_size) + keep = frozenset(list(claimed)[:keep_count]) + drop = set(claimed) - set(keep) + if drop: + self._block_manager.free_leased_blocks(drop) + if target - saved < self._min_save_tokens: + if keep: + self._block_manager.free_leased_blocks(keep) + return None + return target, block_ids, keep + + def _emit_save_candidate( + self, + meta: LMCacheOffloadMetadata, + candidate: _ChunkedSaveCandidate, + ) -> bool: + sid = candidate.sid + seq = candidate.seq + entry = self._save_tracker.get(sid) + if entry is None or entry[0] is not seq: + return False + committed_reservation = self._save_block_reservations.get(sid) + if ( + committed_reservation is None + or committed_reservation.seq is not seq + or committed_reservation.generation != candidate.generation + or not committed_reservation.committed + ): + return False + computed = min( + int( + getattr( + seq, + "_offload_finished_cached_tokens", + getattr(seq, "num_cached_tokens", 0), + ) + ), + int(seq.num_prompt_tokens), + ) + is_last_prefill = computed >= int(seq.num_prompt_tokens) + aligned = candidate.aligned + save_operation = SaveOperationId(seq.id, self._save_nonce) + self._save_nonce += 1 + late_acquired = frozenset() + if hasattr(seq, "_offload_finished_block_ids") and not seq.block_table: + late_source = self._late_save_source(seq, candidate.saved, aligned) + if late_source is None: + if self._save_policy == "priority" and candidate.finished: + self.drop_unadmitted_save(seq, reason="stale") + else: + self._save_committed.pop(sid, None) + self._release_save_reservation(sid, owner=seq) + self._save_tracker.pop(sid, None) + self._forget_save_candidate(sid, seq) + return False + aligned, block_ids, late_acquired = late_source + else: + block_ids = list(seq.block_table) + try: + request = self._build_save_request( + seq, + candidate.saved, + aligned, + save_operation, + block_ids, + is_last_prefill, + ) + except Exception: + if late_acquired: + self._block_manager.free_leased_blocks(late_acquired) + raise + if request is None: + if late_acquired: + self._block_manager.free_leased_blocks(late_acquired) + self._save_committed.pop(sid, None) + self._release_save_reservation(sid, owner=seq) + if candidate.finished: + self.drop_unadmitted_save(seq, reason="capacity") + return False + + logger.debug( + "[OFFLOAD-SAVE-EMIT] seq=%s computed=%d num_prompt=%d " + "aligned=%d saved=%d policy=%s observed=%d", + seq.id, + computed, + int(seq.num_prompt_tokens), + aligned, + candidate.saved, + self._save_policy, + candidate.observed_count, + ) + self._track_save_statistics(save_operation, aligned - candidate.saved) + self._save_inflight_since[save_operation] = time.monotonic() + meta.add_request(request) + entry[1] = aligned + self._save_committed.pop(sid, None) + self._release_save_reservation(sid, owner=seq) + self._save_inflight[sid] = save_operation + if getattr(self, "_early_release", False): + # Freeze the exact token-index -> block-id mapping before a + # finished request clears its block table. The lease itself is + # activated only by `activate_block_leases` at teardown. + source_block_size = getattr(self, "virtual_block_size", self.block_size) + start_block = candidate.saved // source_block_size + end_block = -(-aligned // source_block_size) # ceil div + block_map = { + index: block_ids[index] + for index in range(start_block, min(end_block, len(block_ids))) + } + self._save_operation_blocks[save_operation] = block_map + self._save_operation_safe[save_operation] = set() + self._save_operation_owner[save_operation] = seq + self._save_block_reservations[save_operation] = replace( + committed_reservation, + block_ids=frozenset( + int(block_id) + for block_id in block_map.values() + if int(block_id) >= 0 + ), + estimated_blocks=max(0, end_block - start_block - len(block_map)), + committed=False, + ) + if late_acquired: + self.activate_block_leases(seq, late_acquired) + else: + self._save_block_reservations[save_operation] = replace( + committed_reservation, committed=False + ) + return True + def build_connector_meta(self) -> LMCacheOffloadMetadata: + self._ensure_save_admission_state() meta = LMCacheOffloadMetadata() # Loads @@ -374,79 +1063,98 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: for sid in self._lookup_in_step if sid in loading_sids or sid not in self._load_specs ] - # Saves: store fully computed prompt chunks. Under scheduler-side - # chunked prefill, seq.num_cached_tokens advances after each prefill - # chunk's forward has completed; use it as the D2H-safe frontier. - chunk = self.chunk_size or 256 - tracker_sids = list(self._save_tracker.keys()) - if tracker_sids and self._save_rr_last in self._save_tracker: - start = (tracker_sids.index(self._save_rr_last) + 1) % len(tracker_sids) - tracker_sids = tracker_sids[start:] + tracker_sids[:start] - for sid in tracker_sids: - entry = self._save_tracker[sid] - if not self._do_save: - continue - if not self._may_emit_save(): - break - seq, saved = entry - if sid in self._reqs_need_recv or sid in loading_sids: - continue # loading this step; defer its save - if sid in self._save_inflight: - continue # keep at most one save per request in flight - computed = min( - int( - getattr( - seq, - "_offload_finished_cached_tokens", - getattr(seq, "num_cached_tokens", 0), + # Saves. Priority mode admits/replaces the complete candidate set before + # dispatch so a newly hotter candidate can still evict lower-value + # committed work. Dispatching first would make that work inflight and + # therefore correctly, but prematurely, non-evictable. + if self._do_save: + now = time.monotonic() + for sid, seq in list(self._save_committed.items()): + entry = self._save_tracker.get(sid) + if entry is None or entry[0] is not seq: + self._save_committed.pop(sid, None) + self._release_save_reservation(sid, owner=seq) + continue + candidate = self._build_save_candidate(sid, now=now) + if candidate is None: + self._save_committed.pop(sid, None) + self._release_save_reservation(sid, owner=seq) + self._settle_finished_save(sid) + + tracker_sids = list(self._save_tracker.keys()) + if self._save_policy == "round_robin": + if tracker_sids and self._save_rr_last in self._save_tracker: + start = (tracker_sids.index(self._save_rr_last) + 1) % len( + tracker_sids ) - ), - int(seq.num_prompt_tokens), - ) - is_last_prefill = computed >= int(seq.num_prompt_tokens) - aligned = (computed // chunk) * chunk - if aligned <= saved: - continue - logger.debug( - "[OFFLOAD-SAVE-EMIT] seq=%s computed=%d num_prompt=%d aligned=%d saved=%d", - seq.id, - computed, - int(seq.num_prompt_tokens), - aligned, - saved, - ) - save_operation = SaveOperationId(seq.id, self._save_nonce) - self._save_nonce += 1 - self._track_save_statistics(save_operation, aligned - saved) - block_ids = list( - getattr(seq, "_offload_finished_block_ids", seq.block_table) - ) - meta.add_request( - LMCacheReqMeta( - req_id=seq.id, - token_ids=list(seq.token_ids[:aligned]), - block_ids=block_ids, - save_spec=SaveSpec(skip_leading_tokens=saved, can_save=True), - is_last_prefill=is_last_prefill, - save_operation=save_operation, + tracker_sids = tracker_sids[start:] + tracker_sids[:start] + candidates = [ + candidate + for sid in tracker_sids + if (candidate := self._build_save_candidate(sid, now=now)) + is not None + ] + else: + candidates = [ + candidate + for sid in tracker_sids + if (candidate := self._build_save_candidate(sid, now=now)) + is not None + and self._meets_save_value_threshold(candidate) + ] + candidates.sort( + key=lambda candidate: self._candidate_sort_key(candidate, now) ) - ) - entry[1] = aligned - self._save_inflight[sid] = save_operation - self._save_rr_last = sid - if getattr(self, "_early_release", False): - # Freeze the exact token-index -> block-id mapping before a - # finished request clears its block table. The lease itself is - # activated only by `activate_block_leases` at teardown. - source_block_size = getattr(self, "virtual_block_size", self.block_size) - start_block = saved // source_block_size - end_block = -(-aligned // source_block_size) # ceil div - self._save_operation_blocks[save_operation] = { - index: block_ids[index] - for index in range(start_block, min(end_block, len(block_ids))) - } - self._save_operation_safe[save_operation] = set() - self._save_operation_owner[save_operation] = seq + + if self._save_policy == "priority": + for candidate in candidates: + sid = candidate.sid + if sid in self._save_committed or sid in self._save_inflight: + continue + if sid in self._reqs_need_recv or sid in loading_sids: + continue + if not self._try_reserve_save_candidate( + candidate, + allow_eviction=True, + now=now, + ): + self.drop_unadmitted_save(candidate.seq, reason="capacity") + + committed = [] + for sid, seq in list(self._save_committed.items()): + candidate = self._build_save_candidate(sid, now=now) + if candidate is None or candidate.seq is not seq: + self._save_committed.pop(sid, None) + self._release_save_reservation(sid, owner=seq) + continue + committed.append(candidate) + committed.sort( + key=lambda candidate: self._candidate_sort_key(candidate, now) + ) + for candidate in committed: + sid = candidate.sid + if sid in self._reqs_need_recv or sid in loading_sids: + continue + if sid in self._save_inflight: + continue + self._emit_save_candidate(meta, candidate) + else: + for candidate in candidates: + sid = candidate.sid + if not self._may_emit_save(): + break + if sid in self._save_committed or sid in self._save_inflight: + continue + if sid in self._reqs_need_recv or sid in loading_sids: + continue + if not self._try_reserve_save_candidate( + candidate, + allow_eviction=False, + now=now, + ): + break + if self._emit_save_candidate(meta, candidate): + self._save_rr_last = sid dispatched = set(meta.lookup_requests_in_step) for sid in dispatched: self._lookup_results.pop(sid, None) @@ -457,6 +1165,7 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: return meta def should_defer_free(self, seq) -> bool: + self._ensure_save_admission_state() if self._has_active_load(seq): return True if not self._do_save: @@ -465,6 +1174,16 @@ def should_defer_free(self, seq) -> bool: operation_blocks = getattr(self, "_save_operation_blocks", {}) operation_safe = getattr(self, "_save_operation_safe", {}) operation_owner = getattr(self, "_save_operation_owner", {}) + active_operation = self._save_inflight.get(sid) + active_owner = operation_owner.get(active_operation) + active_entry = self._save_tracker.get(sid) + active_save = active_operation is not None and ( + active_owner is seq + or ( + active_owner is None + and (active_entry is None or active_entry[0] is seq) + ) + ) unsafe_retired_operation = any( owner is seq and not set(operation_blocks.get(operation, {}).values()).issubset( @@ -473,7 +1192,8 @@ def should_defer_free(self, seq) -> bool: for operation, owner in operation_owner.items() ) return ( - sid in self._save_inflight + active_save + or self._save_committed.get(sid) is seq or self._has_pending_save(seq) or unsafe_retired_operation ) @@ -521,7 +1241,30 @@ def protected_block_ids(self, seq) -> frozenset | None: def activate_block_leases(self, seq, block_ids: frozenset[int]) -> None: """Record the refcount shares transferred at request deallocation.""" - if not self._early_release or not block_ids: + if not self._early_release: + return + sid = str(seq.id) + committed = self._save_block_reservations.get(sid) + if committed is not None and committed.seq is seq and committed.committed: + # The scheduler is about to clear the request's live block table. + # Its old physical IDs can be reused, so a not-yet-dispatched save + # keeps only a conservative count reservation from this point on. + self._save_block_reservations[sid] = replace( + committed, + block_ids=frozenset(), + estimated_blocks=self._reservation_footprint(committed), + ) + for operation, reservation in list(self._save_block_reservations.items()): + if ( + isinstance(operation, SaveOperationId) + and reservation.seq is seq + and not reservation.committed + ): + # The future reservation becomes an actual save-owned physical + # lease. Removing it before adding the lease avoids double + # charging the same IDs during the ownership transition. + self._save_block_reservations.pop(operation, None) + if not block_ids: return lease_key = id(seq) leased = self._save_lease_blocks.setdefault(lease_key, set()) @@ -562,6 +1305,7 @@ def reclaim_stale_leases(self, timeout_s: float) -> list[frozenset]: operation = self._save_inflight.get(sid) if sid is not None else None if operation in owned_operations: self._save_inflight.pop(sid, None) + self._save_inflight_since.pop(operation, None) self._cancel_save_statistics(operation) for candidate in [ op for op in self._save_operation_blocks if op in owned_operations @@ -570,10 +1314,19 @@ def reclaim_stale_leases(self, timeout_s: float) -> list[frozenset]: self._save_operation_safe.pop(candidate, None) self._save_operation_owner.pop(candidate, None) self._source_safe_waiting_for_store.pop(candidate, None) + self._release_save_reservation(candidate, owner=owner) if sid is not None: entry = self._save_tracker.get(sid) if entry is not None and entry[0] is owner: self._save_tracker.pop(sid, None) + if self._save_committed.get(sid) is owner: + self._save_committed.pop(sid, None) + self._release_save_reservation(sid, owner=owner) + if self._finished_save_requests.get(sid) is owner: + self._finished_save_requests.pop(sid, None) + self._finished_save_failed.discard(sid) + if owner is not None: + self._forget_save_candidate(sid, owner) if blocks: released.append(frozenset(blocks)) self.total_abnormal_lease_reclaims += len(blocks) @@ -604,6 +1357,7 @@ def has_pending_work(self) -> bool: ) return ( bool(self._reqs_need_recv) + or bool(self._save_committed) or bool(self._save_inflight) or bool(getattr(self, "_save_lease_blocks", {})) or pending_finished_save @@ -709,6 +1463,12 @@ def _source_group_finished(self, identity: SaveSourceGroupId) -> None: safe.update(newly_safe) if not newly_safe: return + reservation = self._save_block_reservations.get(operation) + if reservation is not None: + remaining = reservation.block_ids - newly_safe + self._save_block_reservations[operation] = replace( + reservation, block_ids=frozenset(remaining) + ) sid = str(operation.req_id) owner = self._save_operation_owner.get(operation) lease_key = id(owner) if owner is not None else None @@ -730,10 +1490,12 @@ def _source_group_finished(self, identity: SaveSourceGroupId) -> None: self._save_operation_blocks.pop(operation, None) self._save_operation_safe.pop(operation, None) self._save_operation_owner.pop(operation, None) + self._release_save_reservation(operation) def _store_finished(self, operation: SaveOperationId, *, succeeded: bool) -> None: sid = str(operation.req_id) active = self._save_inflight.get(sid) + owner = self._save_operation_owner.get(operation) if ( active != operation and operation not in self._save_operation_blocks @@ -750,11 +1512,27 @@ def _store_finished(self, operation: SaveOperationId, *, succeeded: bool) -> Non else: self._cancel_save_statistics(operation) # Failed stores keep unsafe ranges leased until abandon timeout. + block_map = self._save_operation_blocks.get(operation, {}) + safe = self._save_operation_safe.get(operation, set()) + if set(block_map.values()).issubset(safe): + # A terminal failure can arrive after every source range became + # safe. Retire that exact operation now; otherwise native MP, + # whose terminal is also its final source fence, leaks owner and + # block-map bookkeeping forever. + self._release_operation_lease(operation) + self._save_inflight_since.pop(operation, None) + if ( + self._save_policy == "priority" + and owner is not None + and self._finished_save_requests.get(sid) is owner + ): + self._settle_finished_save(sid, failed=not succeeded) self._finish_retired_request(sid) def _release_operation_lease(self, operation) -> None: if not isinstance(operation, SaveOperationId): return + self._release_save_reservation(operation) block_map = self._save_operation_blocks.pop(operation, {}) self._save_operation_safe.pop(operation, None) owner = self._save_operation_owner.pop(operation, None) @@ -779,6 +1557,13 @@ def _finish_retired_request(self, sid: str) -> None: seq ): self._save_tracker.pop(sid, None) + if self._save_committed.get(sid) is seq: + self._save_committed.pop(sid, None) + self._release_save_reservation(sid, owner=seq) + if self._finished_save_requests.get(sid) is seq: + self._finished_save_requests.pop(sid, None) + self._finished_save_failed.discard(sid) + self._forget_save_candidate(sid, seq) def blocks_waiting_for_store(self) -> int: return sum( @@ -790,8 +1575,9 @@ def abandon_save(self, req_id) -> None: sid = str(req_id.req_id if isinstance(req_id, SaveOperationId) else req_id) operation = self._save_inflight.pop(sid, None) if operation is not None: + self._save_inflight_since.pop(operation, None) self._cancel_save_statistics(operation) - self._save_tracker.pop(sid, None) + tracker = self._save_tracker.pop(sid, None) owner = self._save_operation_owner.pop(operation, None) lease_key = id(owner) if owner is not None else None self._save_lease_at.pop(lease_key, None) @@ -801,10 +1587,88 @@ def abandon_save(self, req_id) -> None: self._save_operation_blocks.pop(operation, None) self._save_operation_safe.pop(operation, None) self._source_safe_waiting_for_store.pop(operation, None) + self._release_save_reservation(operation, owner=owner) + committed = self._save_committed.pop(sid, None) + if committed is not None: + self._release_save_reservation(sid, owner=committed) + finished = self._finished_save_requests.pop(sid, None) + seq = finished if finished is not None else committed + if seq is None and tracker is not None: + seq = tracker[0] + self._finished_save_failed.discard(sid) + if seq is not None: + self._forget_save_candidate(sid, seq) if blocks: self._pending_source_safe_releases.append(frozenset(blocks)) self.total_abnormal_lease_reclaims += len(blocks) + def get_statistics(self) -> dict[str, int | float]: + """Return save-admission gauges in addition to transfer counters.""" + + self._ensure_save_admission_state() + statistics = super().get_statistics() + now = time.monotonic() + candidates = [] + for sid, entry in self._save_tracker.items(): + seq = entry[0] + operation = self._save_inflight.get(sid) + if ( + self._save_operation_owner.get(operation) is seq + or self._save_committed.get(sid) is seq + ): + continue + candidate = self._build_save_candidate(sid, now=now) + if candidate is not None: + candidates.append(candidate) + + finished = [candidate for candidate in candidates if candidate.finished] + candidate_scores = [ + self._candidate_priority(candidate, now) for candidate in candidates + ] + reserved_blocks, pinned_blocks, budget_used = self._save_block_usage() + budget = self._save_pin_budget_blocks + total_blocks = self._save_pin_total_blocks + statistics.update( + save_candidates=len(candidates), + save_candidates_finished=len(finished), + save_committed=len(self._save_committed), + save_admitted=self.total_save_admitted, + save_candidate_wait_seconds=max( + (max(0.0, now - candidate.enqueued_at) for candidate in candidates), + default=0.0, + ), + save_priority_score=max(candidate_scores, default=0.0), + save_inflight_wait_seconds=max( + ( + max(0.0, now - started) + for started in self._save_inflight_since.values() + ), + default=0.0, + ), + save_pinned_blocks=pinned_blocks, + save_pinned_tokens=pinned_blocks * int(self.virtual_block_size), + save_pin_budget_blocks=0 if budget is None else budget, + save_reserved_blocks=reserved_blocks, + save_pinned_ratio=(budget_used / total_blocks if total_blocks > 0 else 0.0), + save_budget_available_blocks=( + 0 if budget is None else max(0, budget - budget_used) + ), + save_budget_rejected=self.total_save_budget_rejected, + save_budget_rejected_blocks=self.total_save_budget_rejected_blocks, + save_budget_evicted=self.total_save_budget_evicted, + save_budget_evicted_blocks=self.total_save_budget_evicted_blocks, + save_oversized=self.total_save_oversized, + deferred_free_requests=sum( + bool(blocks) for blocks in self._save_lease_blocks.values() + ), + ) + for reason, total in self._save_drop_totals.items(): + statistics[f"save_dropped_{reason}"] = total + statistics[f"save_dropped_tokens_{reason}"] = self._save_drop_token_totals[ + reason + ] + return statistics + def load_failed(self, req_id) -> bool: sid = str(req_id.req_id if isinstance(req_id, LoadOperationId) else req_id) active = self._active_load_operations.get(sid) @@ -864,7 +1728,27 @@ def request_finished(self, seq) -> None: self._load_lifecycles.pop(sid, None) entry = self._save_tracker.get(sid) if entry is not None and entry[0] is seq: - if self._early_release: + if self._save_policy == "priority": + self._finished_save_requests[sid] = seq + # Freeze the final source identity before request teardown. + # Priority reservations dispatch on a later metadata build and + # reacquire only the still-canonical prefix blocks. + if self._early_release: + seq._offload_finished_cached_tokens = min( + int(getattr(seq, "num_cached_tokens", 0)), + int(seq.num_prompt_tokens), + ) + if sid in self._save_inflight or self._save_committed.get(sid) is seq: + pass + else: + candidate = self._build_save_candidate(sid) + if candidate is None: + self._settle_finished_save(sid) + elif not self._meets_save_value_threshold(candidate): + self.drop_unadmitted_save(seq, reason="low_value") + elif not self._commit_finished_save(candidate): + self.drop_unadmitted_save(seq, reason="capacity") + elif self._early_release: # Freeze the final computed frontier before BlockManager # clears it during partial deallocation. Keep the tracker when # a final chunk still needs emission; a later metadata build @@ -878,6 +1762,7 @@ def request_finished(self, seq) -> None: self._save_tracker.pop(sid, None) elif not self.should_defer_free(seq): self._save_tracker.pop(sid, None) + self._forget_save_candidate(sid, seq) if hasattr(seq, "_load_operation"): delattr(seq, "_load_operation") diff --git a/atom/kv_transfer/offload/connector.py b/atom/kv_transfer/offload/connector.py index 8f5b169368..514b233536 100644 --- a/atom/kv_transfer/offload/connector.py +++ b/atom/kv_transfer/offload/connector.py @@ -177,6 +177,11 @@ class LMCacheOffloadConnectorScheduler(KVConnectorSchedulerBase): def __init__(self, config) -> None: self._impl = _build_scheduler(config) + def bind_block_manager(self, block_manager) -> None: + callback = getattr(self._impl, "bind_block_manager", None) + if callable(callback): + callback(block_manager) + @property def has_state_tier(self) -> bool: """True when the selected impl actually hosts the KDA state tier. diff --git a/atom/kv_transfer/offload/hybrid/dsv4/connector.py b/atom/kv_transfer/offload/hybrid/dsv4/connector.py index b77acf7dfd..c3c560c6dc 100644 --- a/atom/kv_transfer/offload/hybrid/dsv4/connector.py +++ b/atom/kv_transfer/offload/hybrid/dsv4/connector.py @@ -31,8 +31,8 @@ import threading import time from contextlib import nullcontext -from dataclasses import dataclass -from math import isfinite +from dataclasses import dataclass, replace +from math import ceil, isfinite from numbers import Integral import torch @@ -86,6 +86,12 @@ SlotLoadSpec, SlotSaveSpec, ) +from atom.kv_transfer.offload.save_admission import ( + PrefixDemandKey, + PrefixDemandTracker, + SaveBlockReservation, + load_save_admission_config, +) logger = logging.getLogger("atom") @@ -148,6 +154,30 @@ def _env_positive_float(name: str, default: float) -> float: return value +@dataclass(frozen=True) +class _SaveCandidate: + sid: str + seq: object + generation: int + saved: int + computed: int + aligned: int + sidecar_candidate: tuple[int, int] | None + observed_count: int + reusable_tokens: int + enqueued_at: float + finished: bool + held_blocks: int + + @property + def dirty_tokens(self) -> int: + return max(0, self.aligned - self.saved) + + @property + def page_save_due(self) -> bool: + return self.aligned > self.saved + + @dataclass(frozen=True) class _SlotSaveSnapshot: """RPC-thread result carried to one background sidecar save.""" @@ -1892,6 +1922,59 @@ def __init__(self, config) -> None: # a long request at the head of the insertion-ordered tracker cannot # monopolize the bounded save queue. self._save_rr_last: str | None = None + # Priority mode separates a lightweight candidate from a capacity- + # holding reservation. This matters at request finish: connector + # metadata for the current step has already been built, so the request + # cannot be dispatched immediately. A committed entry reserves one + # worker slot until the next metadata build emits it. + save_admission = load_save_admission_config() + self._save_policy = save_admission.policy + self._save_max_pinned_ratio = save_admission.max_pinned_ratio + self._save_max_pinned_blocks = save_admission.max_pinned_blocks + self._block_manager = None + self._save_pin_total_blocks = 0 + self._save_pin_budget_blocks: int | None = None + self._save_block_reservations: dict[ + str | SaveOperationId, SaveBlockReservation + ] = {} + if self._save_policy == "priority": + self._save_min_observed_count = save_admission.min_observed_count + self._save_aging_weight = save_admission.aging_weight + self._save_release_weight = save_admission.release_weight + self._prefix_demand = PrefixDemandTracker( + block_tokens=save_admission.demand_block_tokens, + max_entries=save_admission.demand_max_entries, + ttl_seconds=save_admission.demand_ttl_seconds, + ) + else: + self._save_min_observed_count = 2 + self._save_aging_weight = 0.01 + self._save_release_weight = 1.0 + self._prefix_demand = PrefixDemandTracker() + self._save_demand_keys: dict[ + str, tuple[object, tuple[PrefixDemandKey, ...]] + ] = {} + self._save_candidate_since: dict[str, tuple[object, float]] = {} + self._save_candidate_generation: dict[str, tuple[object, int]] = {} + self._save_candidate_nonce = 0 + self._save_committed: dict[str, object] = {} + self._finished_save_requests: dict[str, object] = {} + self._finished_save_failed: set[str] = set() + self._save_inflight_since: dict[SaveOperationId, float] = {} + self._save_operation_owners: dict[SaveOperationId, object] = {} + self.total_save_admitted = 0 + self.total_save_budget_rejected = 0 + self.total_save_budget_rejected_blocks = 0 + self.total_save_budget_evicted = 0 + self.total_save_budget_evicted_blocks = 0 + self.total_save_oversized = 0 + self._save_drop_totals = { + "capacity": 0, + "low_value": 0, + "terminal_failure": 0, + "stale": 0, + } + self._save_drop_token_totals = dict.fromkeys(self._save_drop_totals, 0) # Scheduler-lifetime completion generation: every emitted save gets a # distinct SaveOperationId, so late TP notifications cannot complete a # later PAGE/SLOT save after request cleanup or request-ID reuse. @@ -1904,7 +1987,9 @@ def __init__(self, config) -> None: # ``build_connector_meta`` advances the saved watermark where it *emits* # a save, so the advance is a prediction, not a fact. Keep what each one # was contingent on so a save that never lands can take its advance back. - self._save_watermark_rollback: dict[str, dict[SaveOperationId, int]] = {} + self._save_watermark_rollback: dict[ + str, dict[SaveOperationId, tuple[object, int]] + ] = {} # Stateful PAGE/SLOT protocol. Sidecar commits are session-local because # worker-side sidecar storage is not queried by the scheduler. self._committed_sidecar_hashes = _BoundedLRUSet( @@ -1953,6 +2038,20 @@ def __init__(self, config) -> None: "LMCache offload scheduler: lookup client unavailable: %s", e ) + def bind_block_manager(self, block_manager) -> None: + if self._block_manager is not None and self._block_manager is not block_manager: + raise RuntimeError("offload scheduler is already bound to a block manager") + self._block_manager = block_manager + total_blocks = int(block_manager.total_allocatable_kv_blocks) + if total_blocks <= 0: + raise ValueError("BlockManager must expose a positive KV block capacity") + ratio_budget = int(total_blocks * self._save_max_pinned_ratio) + absolute = self._save_max_pinned_blocks + self._save_pin_total_blocks = total_blocks + self._save_pin_budget_blocks = ( + ratio_budget if absolute is None else min(ratio_budget, absolute) + ) + # -- match: how many extra tokens can come from CPU/NVMe ------------- def _begin_load_lifecycle(self, seq) -> None: sid = str(seq.id) @@ -2082,6 +2181,27 @@ def update_state_after_alloc(self, seq) -> None: entry = self._save_tracker.get(sid) if entry is None or entry[0] is not seq: self._save_tracker[sid] = [seq, initial_saved] + now = time.monotonic() + demand_keys = ( + self._prefix_demand.observe( + seq.token_ids, + int(seq.num_prompt_tokens), + now, + ) + if self._save_policy == "priority" + else () + ) + self._save_demand_keys[sid] = (seq, demand_keys) + self._save_candidate_generation[sid] = ( + seq, + self._save_candidate_nonce, + ) + self._save_candidate_nonce += 1 + previous_committed = self._save_committed.pop(sid, None) + if previous_committed is not None: + self._release_save_reservation(sid, owner=previous_committed) + self._finished_save_requests.pop(sid, None) + self._finished_save_failed.discard(sid) self._sidecar_hash_cache.pop(sid, None) self._failed_sidecar_saves.pop(sid, None) else: @@ -2181,24 +2301,516 @@ def _next_save_operation(self, seq) -> SaveOperationId: self._save_nonce += 1 return operation + def _inflight_save_operations(self) -> set[SaveOperationId]: + operations = { + operation + for inflight in self._save_inflight.values() + for operation in inflight + } + operations.update( + inflight[0] for inflight in self._sidecar_save_inflight.values() + ) + return operations + + def _retire_save_operation_owner(self, operation: SaveOperationId) -> None: + if any(operation in inflight for inflight in self._save_inflight.values()): + return + if any( + inflight[0] == operation + for inflight in self._sidecar_save_inflight.values() + ): + return + self._save_inflight_since.pop(operation, None) + owner = self._save_operation_owners.pop(operation, None) + self._release_save_reservation(operation, owner=owner) + def _may_emit_save(self) -> bool: """Return whether the scheduler may dispatch one more save operation. A PAGE+SLOT request is one worker operation even though both completion channels retain its identity. Counting distinct operation IDs keeps the scheduler's view equal to the worker semaphore while old multi-inflight - state, if any, drains safely. + state, if any, drains safely. A committed finished save also consumes a + slot: it is guaranteed dispatch on the next metadata build and must not + be overbooked by an active request in the meantime. """ - operations = { - operation - for inflight in self._save_inflight.values() - for operation in inflight - } - operations.update( - inflight[0] for inflight in self._sidecar_save_inflight.values() + return ( + len(self._inflight_save_operations()) + len(self._save_committed) + < self._max_pending_saves + ) + + def _save_count_allows(self, committed_after: int) -> bool: + return ( + len(self._inflight_save_operations()) + committed_after + <= self._max_pending_saves + ) + + def _pinned_save_block_ids( + self, + reservations: dict[str | SaveOperationId, SaveBlockReservation], + ) -> set[int]: + pinned: set[int] = set() + for reservation in reservations.values(): + if self._finished_save_requests.get(reservation.sid) is reservation.seq: + pinned.update(reservation.block_ids) + return pinned + + def _save_block_usage( + self, + reservations: dict[str | SaveOperationId, SaveBlockReservation] | None = None, + ) -> tuple[int, int, int]: + """Return ``(reserved, pinned, total_budget_usage)`` without double count.""" + + current = ( + self._save_block_reservations if reservations is None else reservations + ) + pinned_ids = self._pinned_save_block_ids(current) + reserved_ids: set[int] = set() + pinned_estimated = 0 + reserved_estimated = 0 + for reservation in current.values(): + if self._finished_save_requests.get(reservation.sid) is reservation.seq: + pinned_estimated += int(reservation.estimated_blocks) + else: + reserved_ids.update(reservation.block_ids) + reserved_estimated += int(reservation.estimated_blocks) + reserved = len(reserved_ids - pinned_ids) + reserved_estimated + pinned = len(pinned_ids) + pinned_estimated + total = len(pinned_ids | reserved_ids) + pinned_estimated + reserved_estimated + return reserved, pinned, total + + def _candidate_block_reservation( + self, + candidate: _SaveCandidate, + *, + priority_score: float, + ) -> SaveBlockReservation: + # DSV4 in this branch cannot release PAGE source ranges independently. + # If the request finishes while its PAGE or SLOT save is outstanding, + # teardown retains the complete block table. Reserve that exact local + # physical footprint up front so admission cannot pretend only the + # dirty PAGE interval is pinned. + table = list(getattr(candidate.seq, "block_table", ())) + physical_ids = [int(block_id) for block_id in table if int(block_id) >= 0] + block_ids = frozenset(physical_ids) + required_tokens = max( + candidate.aligned if candidate.page_save_due else 0, + ( + candidate.sidecar_candidate[0] + if candidate.sidecar_candidate is not None + else 0 + ), + ) + required_blocks = ceil(required_tokens / max(1, int(self.block_size))) + estimated = max(0, required_blocks - len(physical_ids)) + return SaveBlockReservation( + sid=candidate.sid, + seq=candidate.seq, + generation=candidate.generation, + block_ids=block_ids, + estimated_blocks=estimated, + priority_score=float(priority_score), + committed=True, + ) + + @staticmethod + def _reservation_footprint(reservation: SaveBlockReservation) -> int: + return len(reservation.block_ids) + int(reservation.estimated_blocks) + + def _reservation_fits( + self, + reservation: SaveBlockReservation, + *, + victim_keys: tuple[str, ...] = (), + ) -> bool: + projected = dict(self._save_block_reservations) + for key in victim_keys: + projected.pop(key, None) + projected[reservation.sid] = reservation + committed_after = sum(item.committed for item in projected.values()) + if not self._save_count_allows(committed_after): + return False + budget = self._save_pin_budget_blocks + return budget is None or self._save_block_usage(projected)[2] <= budget + + def _try_reserve_save_candidate( + self, + candidate: _SaveCandidate, + *, + allow_eviction: bool, + now: float | None = None, + ) -> bool: + """Atomically reserve one DSV4 PAGE+SLOT save operation.""" + + existing = self._save_block_reservations.get(candidate.sid) + if ( + existing is not None + and existing.seq is candidate.seq + and existing.generation == candidate.generation + and existing.committed + ): + return True + + score = self._candidate_priority( + candidate, + time.monotonic() if now is None else now, + ) + reservation = self._candidate_block_reservation(candidate, priority_score=score) + victims: list[SaveBlockReservation] = [] + if not self._reservation_fits(reservation): + if allow_eviction: + eligible = sorted( + ( + item + for key, item in self._save_block_reservations.items() + if isinstance(key, str) + and item.committed + and item.priority_score < reservation.priority_score + and self._save_committed.get(key) is item.seq + and key not in self._save_inflight + and key not in self._sidecar_save_inflight + and (entry := self._save_tracker.get(key)) is not None + and entry[0] is item.seq + ), + key=lambda item: ( + item.priority_score, + item.sid, + item.generation, + ), + ) + for victim in eligible: + victims.append(victim) + if self._reservation_fits( + reservation, + victim_keys=tuple(item.sid for item in victims), + ): + break + else: + victims = [] + if not victims and not self._reservation_fits(reservation): + footprint = self._reservation_footprint(reservation) + self.total_save_budget_rejected += 1 + self.total_save_budget_rejected_blocks += footprint + budget = self._save_pin_budget_blocks + if budget is not None and footprint > budget: + self.total_save_oversized += 1 + return False + + # Simulate the whole victim set before mutating anything. PAGE and SLOT + # share one reservation and therefore can only be replaced together. + for victim in victims: + self.total_save_budget_evicted += 1 + self.total_save_budget_evicted_blocks += self._reservation_footprint(victim) + self.drop_unadmitted_save(victim.seq, reason="capacity") + self._save_committed[candidate.sid] = candidate.seq + self._save_block_reservations[candidate.sid] = reservation + self.total_save_admitted += 1 + return True + + def _release_save_reservation( + self, + key: str | SaveOperationId, + *, + owner=None, + ) -> SaveBlockReservation | None: + reservation = self._save_block_reservations.get(key) + if reservation is None or (owner is not None and reservation.seq is not owner): + return None + return self._save_block_reservations.pop(key) + + def _refresh_finished_save_reservations(self, seq) -> None: + """Charge every physical block retained by a now-finished request.""" + + block_ids = frozenset( + physical_id + for block_id in getattr(seq, "block_table", ()) + if (physical_id := int(block_id)) >= 0 + ) + for key, reservation in list(self._save_block_reservations.items()): + if reservation.seq is seq: + self._save_block_reservations[key] = replace( + reservation, + block_ids=block_ids, + estimated_blocks=0, + ) + + def _candidate_generation_for(self, sid: str, seq) -> int: + record = self._save_candidate_generation.get(sid) + if record is not None and record[0] is seq: + return record[1] + generation = self._save_candidate_nonce + self._save_candidate_nonce += 1 + self._save_candidate_generation[sid] = (seq, generation) + return generation + + def _candidate_since_for(self, sid: str, seq, now: float) -> float: + record = self._save_candidate_since.get(sid) + if record is not None and record[0] is seq: + return record[1] + self._save_candidate_since[sid] = (seq, now) + return now + + def _candidate_demand( + self, + sid: str, + seq, + *, + aligned: int, + now: float, + ) -> tuple[int, int]: + record = self._save_demand_keys.get(sid) + if record is None or record[0] is not seq: + return 0, 0 + return self._prefix_demand.heat( + record[1], + max_tokens=aligned, + now=now, + ) + + def _build_save_candidate( + self, sid: str, *, now: float | None = None + ) -> _SaveCandidate | None: + entry = self._save_tracker.get(sid) + if entry is None: + return None + seq, saved = entry + computed = min( + int(getattr(seq, "num_cached_tokens", 0)), + int(getattr(seq, "num_prompt_tokens", 0)), + ) + chunk = int(self.chunk_size or 256) + aligned = (computed // chunk) * chunk + sidecar_candidate = self._sidecar_save_candidate(seq, computed) + if aligned <= int(saved) and sidecar_candidate is None: + return None + candidate_now = time.monotonic() if now is None else float(now) + observed_count, reusable_tokens = self._candidate_demand( + sid, + seq, + aligned=aligned, + now=candidate_now, + ) + finished = self._finished_save_requests.get(sid) is seq + return _SaveCandidate( + sid=sid, + seq=seq, + generation=self._candidate_generation_for(sid, seq), + saved=int(saved), + computed=computed, + aligned=aligned, + sidecar_candidate=sidecar_candidate, + observed_count=observed_count, + reusable_tokens=reusable_tokens, + enqueued_at=self._candidate_since_for(sid, seq, candidate_now), + finished=finished, + held_blocks=len(getattr(seq, "block_table", ())) if finished else 0, + ) + + def _candidate_priority(self, candidate: _SaveCandidate, now: float) -> float: + chunk = max(1, int(self.chunk_size or 256)) + cost_units = max(1, ceil(candidate.dirty_tokens / chunk)) + reusable = max(candidate.reusable_tokens, chunk) + benefit = candidate.observed_count * reusable + release_bonus = ( + self._save_release_weight + * candidate.held_blocks + * max(1, int(self.block_size)) + ) + age = max(0.0, now - candidate.enqueued_at) + return (benefit + release_bonus) / cost_units + self._save_aging_weight * age + + def _candidate_sort_key( + self, candidate: _SaveCandidate, now: float + ) -> tuple[float, int, float, str, int]: + return ( + -self._candidate_priority(candidate, now), + candidate.dirty_tokens, + candidate.enqueued_at, + candidate.sid, + candidate.generation, + ) + + def _meets_save_value_threshold(self, candidate: _SaveCandidate) -> bool: + return candidate.observed_count >= self._save_min_observed_count + + def _commit_finished_save(self, candidate: _SaveCandidate) -> bool: + if not candidate.finished or not self._meets_save_value_threshold(candidate): + return False + entry = self._save_tracker.get(candidate.sid) + if entry is None or entry[0] is not candidate.seq: + return False + return self._try_reserve_save_candidate( + candidate, + allow_eviction=self._save_policy == "priority", + ) + + def _forget_save_candidate(self, sid: str, seq) -> None: + for mapping in ( + self._save_demand_keys, + self._save_candidate_since, + self._save_candidate_generation, + ): + record = mapping.get(sid) + if record is not None and record[0] is seq: + mapping.pop(sid, None) + + def drop_unadmitted_save(self, seq, *, reason: str) -> bool: + """Drop save work that has never been handed to a worker. + + PAGE and SLOT form one checkpoint operation, so cleanup is all-or-none. + An inflight operation is never cancelled here; its source may already + be in use by a worker. + """ + + sid = str(seq.id) + if sid in self._save_inflight or sid in self._sidecar_save_inflight: + return False + entry = self._save_tracker.get(sid) + if entry is None or entry[0] is not seq: + return False + candidate = self._build_save_candidate(sid) + dropped_tokens = 0 if candidate is None else candidate.dirty_tokens + self._save_tracker.pop(sid, None) + if self._save_committed.get(sid) is seq: + self._save_committed.pop(sid, None) + self._release_save_reservation(sid, owner=seq) + if self._finished_save_requests.get(sid) is seq: + self._finished_save_requests.pop(sid, None) + self._finished_save_failed.discard(sid) + self._failed_sidecar_saves.pop(sid, None) + self._save_watermark_rollback.pop(sid, None) + cached = self._sidecar_hash_cache.get(sid) + if cached is not None and cached[0] is seq: + self._sidecar_hash_cache.pop(sid, None) + self._forget_save_candidate(sid, seq) + if reason not in self._save_drop_totals: + reason = "stale" + self._save_drop_totals[reason] += 1 + self._save_drop_token_totals[reason] += dropped_tokens + logger.info( + "LMCache offload: dropped unadmitted save req=%s reason=%s tokens=%d", + sid, + reason, + dropped_tokens, + ) + return True + + def _settle_finished_save(self, sid: str) -> None: + seq = self._finished_save_requests.get(sid) + if seq is None: + return + if ( + sid in self._save_inflight + or sid in self._sidecar_save_inflight + or sid in self._save_committed + ): + return + if sid in self._finished_save_failed: + self.drop_unadmitted_save(seq, reason="terminal_failure") + return + candidate = self._build_save_candidate(sid) + if candidate is None: + entry = self._save_tracker.get(sid) + if entry is not None and entry[0] is seq: + self._save_tracker.pop(sid, None) + self._finished_save_requests.pop(sid, None) + self._finished_save_failed.discard(sid) + self._failed_sidecar_saves.pop(sid, None) + self._save_watermark_rollback.pop(sid, None) + self._forget_save_candidate(sid, seq) + return + if self._commit_finished_save(candidate): + return + reason = ( + "low_value" + if not self._meets_save_value_threshold(candidate) + else "capacity" + ) + self.drop_unadmitted_save(seq, reason=reason) + + def _emit_save_candidate( + self, + meta: LMCacheOffloadMetadata, + candidate: _SaveCandidate, + ) -> None: + sid = candidate.sid + seq = candidate.seq + slot_save_spec = None + if candidate.sidecar_candidate is not None: + boundary, boundary_hash = candidate.sidecar_candidate + slot_save_spec = SlotSaveSpec( + boundary_tokens=boundary, + boundary_block_hash=boundary_hash, + source_group=int(seq.state_slot), + ) + token_end = max( + candidate.aligned if candidate.page_save_due else 0, + ( + candidate.sidecar_candidate[0] + if candidate.sidecar_candidate is not None + else 0 + ), + ) + logger.debug( + "[OFFLOAD-SAVE-EMIT] seq=%s computed=%d num_prompt=%d " + "aligned=%d saved=%d sidecar=%s policy=%s observed=%d", + seq.id, + candidate.computed, + int(seq.num_prompt_tokens), + candidate.aligned, + candidate.saved, + candidate.sidecar_candidate, + self._save_policy, + candidate.observed_count, + ) + save_operation = self._next_save_operation(seq) + committed_reservation = self._save_block_reservations.get(sid) + if ( + committed_reservation is None + or committed_reservation.seq is not seq + or committed_reservation.generation != candidate.generation + or not committed_reservation.committed + ): + return + self._track_save_statistics(save_operation, candidate.dirty_tokens) + self._save_inflight_since[save_operation] = time.monotonic() + self._save_operation_owners[save_operation] = seq + meta.add_request( + LMCacheReqMeta( + req_id=seq.id, + token_ids=list(seq.token_ids[:token_end]), + block_ids=list(seq.block_table), + save_spec=( + SaveSpec(skip_leading_tokens=candidate.saved, can_save=True) + if candidate.page_save_due + else None + ), + is_last_prefill=(candidate.computed >= int(seq.num_prompt_tokens)), + slot_save_spec=slot_save_spec, + save_operation=save_operation, + ) + ) + entry = self._save_tracker.get(sid) + if candidate.page_save_due and entry is not None and entry[0] is seq: + self._save_watermark_rollback.setdefault(sid, {})[save_operation] = ( + seq, + int(entry[1]), + ) + entry[1] = candidate.aligned + self._save_inflight.setdefault(sid, set()).add(save_operation) + if candidate.sidecar_candidate is not None: + self._sidecar_save_inflight[sid] = ( + save_operation, + candidate.sidecar_candidate[0], + candidate.sidecar_candidate[1], + ) + self._save_committed.pop(sid, None) + self._release_save_reservation(sid, owner=seq) + self._save_block_reservations[save_operation] = replace( + committed_reservation, + committed=False, ) - return len(operations) < self._max_pending_saves def _clear_lookup_status(self, sid: str) -> None: if self._lookup_client is None: @@ -2390,92 +3002,108 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: meta.lookup_requests_in_step = [ sid for sid in self._lookup_in_step if sid not in self._handoff_loads ] - # Saves: store fully computed prompt chunks. Under scheduler-side - # chunked prefill, seq.num_cached_tokens advances after each prefill - # chunk's forward has completed; use it as the D2H-safe frontier. - chunk = self.chunk_size or 256 - tracker_sids = list(self._save_tracker.keys()) - if tracker_sids and self._save_rr_last in self._save_tracker: - start = (tracker_sids.index(self._save_rr_last) + 1) % len(tracker_sids) - tracker_sids = tracker_sids[start:] + tracker_sids[:start] - for sid in tracker_sids: - if not self._do_save: - continue - if not self._may_emit_save(): - break - entry = self._save_tracker[sid] - seq, saved = entry - if sid in self._reqs_need_recv or sid in loading_sids: - continue # loading this step; defer its save - if sid in self._save_inflight or sid in self._sidecar_save_inflight: - # A request's PAGE and SLOT describe one ordered checkpoint - # stream. Do not let a later boundary overtake an earlier save. - continue - computed = min( - int(getattr(seq, "num_cached_tokens", 0)), - int(seq.num_prompt_tokens), - ) - is_last_prefill = computed >= int(seq.num_prompt_tokens) - aligned = (computed // chunk) * chunk - sidecar_candidate = self._sidecar_save_candidate(seq, computed) - page_save_due = aligned > saved - if not page_save_due and sidecar_candidate is None: - continue - slot_save_spec = None - if sidecar_candidate is not None: - boundary, boundary_hash = sidecar_candidate - slot_save_spec = SlotSaveSpec( - boundary_tokens=boundary, - boundary_block_hash=boundary_hash, - source_group=int(seq.state_slot), - ) - token_end = max( - aligned if page_save_due else 0, - sidecar_candidate[0] if sidecar_candidate is not None else 0, - ) - logger.debug( - "[OFFLOAD-SAVE-EMIT] seq=%s computed=%d num_prompt=%d " - "aligned=%d saved=%d sidecar=%s", - seq.id, - computed, - int(seq.num_prompt_tokens), - aligned, - saved, - sidecar_candidate, - ) - save_operation = self._next_save_operation(seq) - self._track_save_statistics( - save_operation, - aligned - saved if page_save_due else 0, - ) - meta.add_request( - LMCacheReqMeta( - req_id=seq.id, - token_ids=list(seq.token_ids[:token_end]), - block_ids=list(seq.block_table), - save_spec=( - SaveSpec(skip_leading_tokens=saved, can_save=True) - if page_save_due - else None - ), - is_last_prefill=is_last_prefill, - slot_save_spec=slot_save_spec, - save_operation=save_operation, - ) - ) - if page_save_due: - self._save_watermark_rollback.setdefault(sid, {})[save_operation] = int( - entry[1] + # Saves. Priority mode admits/replaces the complete candidate set before + # dispatch so a newly hotter candidate can still evict lower-value + # committed work. Dispatching first would make that work inflight and + # therefore correctly, but prematurely, non-evictable. + if self._do_save: + now = time.monotonic() + for sid, seq in list(self._save_committed.items()): + entry = self._save_tracker.get(sid) + if entry is None or entry[0] is not seq: + self._save_committed.pop(sid, None) + self._release_save_reservation(sid, owner=seq) + continue + candidate = self._build_save_candidate(sid, now=now) + if candidate is None: + self._save_committed.pop(sid, None) + self._release_save_reservation(sid, owner=seq) + self._settle_finished_save(sid) + + tracker_sids = list(self._save_tracker.keys()) + if self._save_policy == "round_robin": + if tracker_sids and self._save_rr_last in self._save_tracker: + start = (tracker_sids.index(self._save_rr_last) + 1) % len( + tracker_sids + ) + tracker_sids = tracker_sids[start:] + tracker_sids[:start] + candidates = [ + candidate + for sid in tracker_sids + if (candidate := self._build_save_candidate(sid, now=now)) + is not None + ] + else: + candidates = [ + candidate + for sid in tracker_sids + if (candidate := self._build_save_candidate(sid, now=now)) + is not None + and self._meets_save_value_threshold(candidate) + ] + candidates.sort( + key=lambda candidate: self._candidate_sort_key(candidate, now) ) - entry[1] = aligned - self._save_inflight.setdefault(sid, set()).add(save_operation) - if sidecar_candidate is not None: - self._sidecar_save_inflight[sid] = ( - save_operation, - sidecar_candidate[0], - sidecar_candidate[1], + + if self._save_policy == "priority": + for candidate in candidates: + sid = candidate.sid + if sid in self._save_committed: + continue + if sid in self._reqs_need_recv or sid in loading_sids: + continue + if sid in self._save_inflight or sid in self._sidecar_save_inflight: + continue + if not self._try_reserve_save_candidate( + candidate, + allow_eviction=True, + now=now, + ): + self.drop_unadmitted_save(candidate.seq, reason="capacity") + + committed = [] + for sid, seq in list(self._save_committed.items()): + candidate = self._build_save_candidate(sid, now=now) + if candidate is None or candidate.seq is not seq: + self._save_committed.pop(sid, None) + self._release_save_reservation(sid, owner=seq) + continue + committed.append(candidate) + committed.sort( + key=lambda candidate: self._candidate_sort_key(candidate, now) ) - self._save_rr_last = sid + for candidate in committed: + sid = candidate.sid + if sid in self._reqs_need_recv or sid in loading_sids: + continue + if sid in self._save_inflight or sid in self._sidecar_save_inflight: + continue + self._emit_save_candidate(meta, candidate) + else: + for candidate in candidates: + sid = candidate.sid + if not self._may_emit_save(): + break + if sid in self._save_committed: + continue + if sid in self._reqs_need_recv or sid in loading_sids: + continue # loading this step; defer its save + if sid in self._save_inflight or sid in self._sidecar_save_inflight: + # PAGE and SLOT describe one ordered checkpoint stream. + continue + if not self._try_reserve_save_candidate( + candidate, + allow_eviction=False, + now=now, + ): + if candidate.finished: + self.drop_unadmitted_save( + candidate.seq, + reason="capacity", + ) + break + self._emit_save_candidate(meta, candidate) + self._save_rr_last = sid dispatched = set(meta.lookup_requests_in_step) self._lookup_in_step = [ sid for sid in self._lookup_in_step if sid not in dispatched @@ -2502,6 +3130,7 @@ def should_defer_free(self, seq) -> bool: return ( sid in self._save_inflight or sid in self._sidecar_save_inflight + or self._save_committed.get(sid) is seq or self._has_pending_save(seq) or self._has_pending_sidecar_save(seq) ) @@ -2512,19 +3141,22 @@ def has_pending_work(self) -> bool: Feeds ``EngineCore.has_pending_kv_work()``, so it reads only state that clears itself: ``_reqs_need_recv`` is emptied by every ``build_connector_meta`` and ``_save_inflight`` by ``save_finished``. - Saves that are queued but not yet dispatched are covered there by the - scheduler's ``deferred_free_blocks``, which ``should_defer_free`` - keeps populated for exactly those requests. + A priority save committed after its request finished is included even + before its next-step dispatch, so an idle engine keeps producing + connector metadata until that reservation becomes an operation. """ return bool( - self._reqs_need_recv or self._save_inflight or self._sidecar_save_inflight + self._reqs_need_recv + or self._save_committed + or self._save_inflight + or self._sidecar_save_inflight ) @staticmethod def _watermark_sid(operation) -> str: return str(getattr(operation, "req_id", operation)) - def _forget_save_watermark(self, operation) -> int | None: + def _forget_save_watermark(self, operation) -> tuple[object, int] | None: """Drop one operation's contingency record, returning what it held.""" sid = self._watermark_sid(operation) pending = self._save_watermark_rollback.get(sid) @@ -2553,12 +3185,13 @@ def _rollback_save_watermark(self, operation) -> None: prefix has to be contiguous anyway, so everything beyond the hole was unusable regardless. """ - previous = self._forget_save_watermark(operation) - if previous is None: + rollback = self._forget_save_watermark(operation) + if rollback is None: return + seq, previous = rollback sid = self._watermark_sid(operation) entry = self._save_tracker.get(sid) - if entry is None: + if entry is None or entry[0] is not seq: return current = int(entry[1]) if current <= previous: @@ -2575,12 +3208,21 @@ def _rollback_save_watermark(self, operation) -> None: def save_finished(self, req_id) -> None: if isinstance(req_id, SaveOperationId): sid = str(req_id.req_id) + owner = self._save_operation_owners.get(req_id) inflight = self._save_inflight.get(sid) if inflight is not None: inflight.discard(req_id) if not inflight: self._save_inflight.pop(sid, None) + self._save_inflight_since.pop(req_id, None) self._finish_save_statistics(req_id) + self._retire_save_operation_owner(req_id) + if ( + self._save_policy == "priority" + and owner is not None + and self._finished_save_requests.get(sid) is owner + ): + self._settle_finished_save(sid) return sid = str(req_id) @@ -2591,8 +3233,14 @@ def save_finished(self, req_id) -> None: return if sid in self._sidecar_save_inflight: return - self._save_inflight.pop(sid, None) + retired = self._save_inflight.pop(sid, None) + if retired is not None: + for operation in retired: + self._save_inflight_since.pop(operation, None) + self._retire_save_operation_owner(operation) self._finish_save_statistics(req_id) + if self._save_policy == "priority": + self._settle_finished_save(sid) def abandon_save(self, req_id) -> None: """Force-drop a save the scheduler reclaimed after it stalled. @@ -2610,14 +3258,34 @@ def abandon_save(self, req_id) -> None: entry is dropped so the save loop cannot re-emit against freed blocks. """ sid = str(req_id.req_id if isinstance(req_id, SaveOperationId) else req_id) + retired_operations: set[SaveOperationId] = set() inflight = self._save_inflight.pop(sid, None) if inflight is not None: for operation in inflight: + retired_operations.add(operation) + self._save_inflight_since.pop(operation, None) + self._save_operation_owners.pop(operation, None) self._cancel_save_statistics(operation) sidecar = self._sidecar_save_inflight.pop(sid, None) if sidecar is not None: + retired_operations.add(sidecar[0]) + self._save_inflight_since.pop(sidecar[0], None) + self._save_operation_owners.pop(sidecar[0], None) self._cancel_save_statistics(sidecar[0]) + for operation in retired_operations: + self._release_save_reservation(operation) + committed = self._save_committed.pop(sid, None) + if committed is not None: + self._release_save_reservation(sid, owner=committed) + finished = self._finished_save_requests.pop(sid, None) + seq = finished if finished is not None else committed + if seq is None: + entry = self._save_tracker.get(sid) + seq = None if entry is None else entry[0] self._save_tracker.pop(sid, None) + self._finished_save_failed.discard(sid) + if seq is not None: + self._forget_save_candidate(sid, seq) # The tracker entry is gone, so there is no longer a watermark to take # back; keeping the records would only leak. self._save_watermark_rollback.pop(sid, None) @@ -2642,6 +3310,19 @@ def connector_completion(self, completion: ConnectorCompletion) -> bool: if completion.succeeded: self._release_save_watermark(completion.operation_id) else: + sid = self._watermark_sid(completion.operation_id) + inflight = self._save_inflight.get(sid, set()) + sidecar = self._sidecar_save_inflight.get(sid) + owns_operation = completion.operation_id in inflight or ( + sidecar is not None and sidecar[0] == completion.operation_id + ) + owner = self._save_operation_owners.get(completion.operation_id) + if ( + owns_operation + and owner is not None + and self._finished_save_requests.get(sid) is owner + ): + self._finished_save_failed.add(sid) self._rollback_save_watermark(completion.operation_id) # Nothing was persisted, so these bytes must not land in # `total_saved_tokens`; cancelling first makes the @@ -2667,6 +3348,8 @@ def sidecar_save_finished(self, req_id) -> None: inflight[0], SaveOperationId ): return + operation = inflight[0] + owner = self._save_operation_owners.get(operation) self._sidecar_save_inflight.pop(sid, None) identity = (inflight[1], inflight[2]) self._committed_sidecar_hashes.add(identity[1]) @@ -2675,6 +3358,13 @@ def sidecar_save_finished(self, req_id) -> None: failed.discard(identity) if not failed: self._failed_sidecar_saves.pop(sid, None) + self._retire_save_operation_owner(operation) + if ( + self._save_policy == "priority" + and owner is not None + and self._finished_save_requests.get(sid) is owner + ): + self._settle_finished_save(sid) def sidecar_save_failed(self, req_id) -> None: sid = str(req_id.req_id if isinstance(req_id, SaveOperationId) else req_id) @@ -2687,9 +3377,20 @@ def sidecar_save_failed(self, req_id) -> None: inflight[0], SaveOperationId ): return + operation = inflight[0] + owner = self._save_operation_owners.get(operation) + if owner is not None and self._finished_save_requests.get(sid) is owner: + self._finished_save_failed.add(sid) self._sidecar_save_inflight.pop(sid, None) identity = (inflight[1], inflight[2]) self._failed_sidecar_saves.setdefault(sid, set()).add(identity) + self._retire_save_operation_owner(operation) + if ( + self._save_policy == "priority" + and owner is not None + and self._finished_save_requests.get(sid) is owner + ): + self._settle_finished_save(sid) def load_failed(self, req_id) -> bool: sid = str(req_id.req_id if isinstance(req_id, LoadOperationId) else req_id) @@ -2728,6 +3429,78 @@ def load_finished(self, req_id) -> bool: self._load_save_floors.pop(sid, None) return True + def get_statistics(self) -> dict[str, int | float]: + statistics = super().get_statistics() + now = time.monotonic() + candidates = [] + for sid in self._save_tracker: + if ( + sid in self._save_inflight + or sid in self._sidecar_save_inflight + or sid in self._save_committed + ): + continue + candidate = self._build_save_candidate(sid, now=now) + if candidate is not None: + candidates.append(candidate) + + finished = [ + candidate + for candidate in candidates + if self._finished_save_requests.get(candidate.sid) is candidate.seq + ] + candidate_scores = [ + self._candidate_priority(candidate, now) for candidate in candidates + ] + reserved_blocks, pinned_blocks, budget_used = self._save_block_usage() + budget = self._save_pin_budget_blocks + total_blocks = self._save_pin_total_blocks + statistics.update( + save_candidates=len(candidates), + save_candidates_finished=len(finished), + save_committed=len(self._save_committed), + save_admitted=self.total_save_admitted, + save_candidate_wait_seconds=max( + (max(0.0, now - candidate.enqueued_at) for candidate in candidates), + default=0.0, + ), + save_priority_score=max(candidate_scores, default=0.0), + save_inflight_wait_seconds=max( + ( + max(0.0, now - started) + for started in self._save_inflight_since.values() + ), + default=0.0, + ), + save_pinned_blocks=pinned_blocks, + save_pinned_tokens=pinned_blocks * int(self.block_size), + save_pin_budget_blocks=0 if budget is None else budget, + save_reserved_blocks=reserved_blocks, + save_pinned_ratio=(budget_used / total_blocks if total_blocks > 0 else 0.0), + save_budget_available_blocks=( + 0 if budget is None else max(0, budget - budget_used) + ), + save_budget_rejected=self.total_save_budget_rejected, + save_budget_rejected_blocks=self.total_save_budget_rejected_blocks, + save_budget_evicted=self.total_save_budget_evicted, + save_budget_evicted_blocks=self.total_save_budget_evicted_blocks, + save_oversized=self.total_save_oversized, + deferred_free_requests=len( + { + reservation.sid + for reservation in self._save_block_reservations.values() + if self._finished_save_requests.get(reservation.sid) + is reservation.seq + } + ), + ) + for reason, total in self._save_drop_totals.items(): + statistics[f"save_dropped_{reason}"] = total + statistics[f"save_dropped_tokens_{reason}"] = self._save_drop_token_totals[ + reason + ] + return statistics + def request_finished(self, seq) -> None: sid = str(seq.id) if self._load_lifecycles.get(sid) is seq: @@ -2742,10 +3515,32 @@ def request_finished(self, seq) -> None: if cached is not None and cached[0] is seq: self._sidecar_hash_cache.pop(sid, None) entry = self._save_tracker.get(sid) - if entry is not None and entry[0] is seq and not self.should_defer_free(seq): - self._save_tracker.pop(sid, None) - self._failed_sidecar_saves.pop(sid, None) - self._save_watermark_rollback.pop(sid, None) + if entry is not None and entry[0] is seq: + self._finished_save_requests[sid] = seq + self._refresh_finished_save_reservations(seq) + if self._save_policy == "priority": + if ( + sid in self._save_inflight + or sid in self._sidecar_save_inflight + or self._save_committed.get(sid) is seq + ): + # Already committed/dispatched. Its terminal callbacks will + # either reserve one residual tail or drop it. + pass + else: + candidate = self._build_save_candidate(sid) + if candidate is None: + self._settle_finished_save(sid) + elif not self._meets_save_value_threshold(candidate): + self.drop_unadmitted_save(seq, reason="low_value") + elif not self._commit_finished_save(candidate): + self.drop_unadmitted_save(seq, reason="capacity") + elif not self.should_defer_free(seq): + self._save_tracker.pop(sid, None) + self._finished_save_requests.pop(sid, None) + self._failed_sidecar_saves.pop(sid, None) + self._save_watermark_rollback.pop(sid, None) + self._forget_save_candidate(sid, seq) if hasattr(seq, "_load_operation"): delattr(seq, "_load_operation") diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py index 8669994cd1..700cab4cb5 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py @@ -660,7 +660,8 @@ def _may_emit_save(self) -> bool: than `OFFLOAD_MAX_PENDING_SAVES` requests pinned at once.""" return ( not self._save_stalled - and len(self._save_inflight) < self._max_pending_saves + and len(self._save_inflight) + len(self._save_committed) + < self._max_pending_saves ) def has_pending_work(self) -> bool: diff --git a/atom/kv_transfer/offload/save_admission.py b/atom/kv_transfer/offload/save_admission.py new file mode 100644 index 0000000000..fdffcafeee --- /dev/null +++ b/atom/kv_transfer/offload/save_admission.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Shared, CPU-only policy helpers for value-aware offload save admission.""" + +from __future__ import annotations + +import array +import hashlib +import os +import time +from collections import OrderedDict +from dataclasses import dataclass +from typing import Any + +PrefixDemandKey = tuple[int, bytes] + + +def _nonnegative_int(name: str, default: int) -> int: + raw = os.environ.get(name, str(default)) + try: + value = int(raw) + except (TypeError, ValueError) as exc: + raise ValueError(f"{name} must be an integer") from exc + if value < 0: + raise ValueError(f"{name} must be nonnegative") + return value + + +def _nonnegative_float(name: str, default: float) -> float: + raw = os.environ.get(name, str(default)) + try: + value = float(raw) + except (TypeError, ValueError) as exc: + raise ValueError(f"{name} must be a number") from exc + if not 0 <= value < float("inf"): + raise ValueError(f"{name} must be finite and nonnegative") + return value + + +def _positive_float(name: str, default: float) -> float: + value = _nonnegative_float(name, default) + if value <= 0: + raise ValueError(f"{name} must be positive") + return value + + +def _optional_nonnegative_int(name: str) -> int | None: + raw = os.environ.get(name) + if raw is None or not raw.strip(): + return None + try: + value = int(raw) + except (TypeError, ValueError) as exc: + raise ValueError(f"{name} must be an integer") from exc + if value < 0: + raise ValueError(f"{name} must be nonnegative") + return value + + +@dataclass(frozen=True) +class SaveAdmissionConfig: + policy: str + min_observed_count: int + aging_weight: float + release_weight: float + demand_block_tokens: int + demand_max_entries: int + demand_ttl_seconds: float + max_pinned_ratio: float + max_pinned_blocks: int | None + + +@dataclass(frozen=True) +class SaveBlockReservation: + """One scheduler-local PAGE budget reservation. + + ``block_ids`` are physical IDs from this scheduler's own BlockManager. + ``estimated_blocks`` accounts for source blocks whose current physical IDs + are unavailable (notably a finished request after early deallocation). + The two fields are additive; exact IDs are unioned across reservations and + leases before the conservative estimate is added. + """ + + sid: str + seq: Any + generation: int + block_ids: frozenset[int] + estimated_blocks: int + priority_score: float + committed: bool + + +def load_save_admission_config() -> SaveAdmissionConfig: + """Read the shared save policy, preserving old defaults when disabled.""" + + policy = os.environ.get("OFFLOAD_SAVE_POLICY", "round_robin").strip().lower() + if policy not in {"round_robin", "priority"}: + raise ValueError( + f"OFFLOAD_SAVE_POLICY must be 'round_robin' or 'priority', got {policy!r}" + ) + max_pinned_ratio = _nonnegative_float("OFFLOAD_SAVE_MAX_PINNED_RATIO", 0.20) + if max_pinned_ratio > 0.30: + raise ValueError("OFFLOAD_SAVE_MAX_PINNED_RATIO must be at most 0.30") + max_pinned_blocks = _optional_nonnegative_int("OFFLOAD_SAVE_MAX_PINNED_BLOCKS") + if policy == "round_robin": + return SaveAdmissionConfig( + policy, + 2, + 0.01, + 1.0, + 8192, + 65536, + 600.0, + max_pinned_ratio, + max_pinned_blocks, + ) + return SaveAdmissionConfig( + policy=policy, + min_observed_count=_nonnegative_int("OFFLOAD_SAVE_MIN_OBSERVED_COUNT", 2), + aging_weight=_nonnegative_float("OFFLOAD_SAVE_AGING_WEIGHT", 0.01), + release_weight=_nonnegative_float("OFFLOAD_SAVE_RELEASE_WEIGHT", 1.0), + demand_block_tokens=_nonnegative_int("OFFLOAD_SAVE_DEMAND_BLOCK_TOKENS", 8192), + demand_max_entries=_nonnegative_int("OFFLOAD_SAVE_DEMAND_MAX_ENTRIES", 65536), + demand_ttl_seconds=_positive_float("OFFLOAD_SAVE_DEMAND_TTL_SECONDS", 600), + max_pinned_ratio=max_pinned_ratio, + max_pinned_blocks=max_pinned_blocks, + ) + + +@dataclass +class _PrefixDemandEntry: + count: int + last_seen: float + + +class PrefixDemandTracker: + """Bounded rank-local demand counts for coarse cumulative prefixes. + + These counts deliberately do not share DP prefix-route hint state. Route + hints are published only after prefill proves an HBM owner, whereas save + admission needs request demand at entry, before a prefix is stored. + """ + + def __init__( + self, + *, + block_tokens: int = 8192, + max_entries: int = 65536, + ttl_seconds: float = 600.0, + clock=time.monotonic, + ) -> None: + if block_tokens <= 0 or max_entries <= 0: + raise ValueError("prefix demand block size and capacity must be positive") + if not 0 < ttl_seconds < float("inf"): + raise ValueError("prefix demand TTL must be finite and positive") + self.block_tokens = int(block_tokens) + self.max_entries = int(max_entries) + self.ttl_seconds = float(ttl_seconds) + self._clock = clock + self._entries: OrderedDict[PrefixDemandKey, _PrefixDemandEntry] = OrderedDict() + + def fingerprints( + self, token_ids: array.array, prompt_tokens: int + ) -> tuple[PrefixDemandKey, ...]: + """Hash cumulative coarse prefixes without retaining token buffers.""" + + prompt_tokens = int(prompt_tokens) + if not 0 <= prompt_tokens <= len(token_ids): + raise ValueError("prompt length exceeds the supplied token buffer") + if not isinstance(token_ids, array.array): + token_ids = array.array("i", token_ids) + if token_ids.typecode != "i" or token_ids.itemsize != 4: + raise ValueError("prefix demand expects an int32 token buffer") + + raw = memoryview(token_ids).cast("B")[: prompt_tokens * 4] + block_bytes = self.block_tokens * 4 + digest = hashlib.blake2b(digest_size=16) + keys = [] + for end in range(block_bytes, len(raw) + 1, block_bytes): + digest.update(raw[end - block_bytes : end]) + keys.append((end // 4, digest.digest())) + if not keys and raw: + digest.update(raw) + keys.append((prompt_tokens, digest.digest())) + return tuple(keys) + + def observe( + self, token_ids: array.array, prompt_tokens: int, now: float | None = None + ) -> tuple[PrefixDemandKey, ...]: + """Record one request and return reusable keys for later scoring.""" + + observed_at = self._clock() if now is None else float(now) + keys = self.fingerprints(token_ids, prompt_tokens) + for key in keys: + entry = self._entries.get(key) + if entry is None or observed_at - entry.last_seen >= self.ttl_seconds: + entry = _PrefixDemandEntry(count=1, last_seen=observed_at) + self._entries[key] = entry + else: + entry.count += 1 + entry.last_seen = observed_at + self._entries.move_to_end(key) + if len(self._entries) > self.max_entries: + self._entries.popitem(last=False) + return keys + + def heat( + self, + keys: tuple[PrefixDemandKey, ...], + *, + max_tokens: int, + now: float | None = None, + ) -> tuple[int, int]: + """Return ``(observed_count, reusable_tokens)`` for the hottest key.""" + + checked_at = self._clock() if now is None else float(now) + reusable_limit = max(0, int(max_tokens)) + if reusable_limit == 0: + return 0, 0 + best = (0, 0) + for key in keys: + entry = self._entries.get(key) + if entry is None: + continue + if checked_at - entry.last_seen >= self.ttl_seconds: + del self._entries[key] + continue + self._entries.move_to_end(key) + best = max(best, (entry.count, min(int(key[0]), reusable_limit))) + return best + + def clear(self) -> None: + self._entries.clear() + + +__all__ = [ + "PrefixDemandKey", + "PrefixDemandTracker", + "SaveAdmissionConfig", + "SaveBlockReservation", + "load_save_admission_config", +] diff --git a/atom/model_engine/block_manager.py b/atom/model_engine/block_manager.py index 5e78520798..374336f4b2 100644 --- a/atom/model_engine/block_manager.py +++ b/atom/model_engine/block_manager.py @@ -575,6 +575,17 @@ def num_pool_blocks(self, seq_len: int) -> int: )[0] return int((local_len + self.block_size - 1) // self.block_size) + @property + def total_allocatable_kv_blocks(self) -> int: + """Physical KV blocks owned by this scheduler-local allocation pool. + + This value is already local to one EngineCore/DP scheduler. Tensor + parallel ranks shard the contents of these logical blocks, so callers + must neither aggregate another DP pool nor multiply by TP world size. + """ + + return int(self.kv.num_blocks) + @property def max_pool_tokens(self) -> int: """Longest prompt, in global tokens, whose KV fits an entirely empty pool. diff --git a/atom/model_engine/llm_engine.py b/atom/model_engine/llm_engine.py index c073accc82..80fe190a64 100644 --- a/atom/model_engine/llm_engine.py +++ b/atom/model_engine/llm_engine.py @@ -570,11 +570,49 @@ def summed_queues(key: str) -> int: "saved_tokens", "loads_pending", "saves_pending", + "early_released_blocks", + "leased_source_blocks", + "source_safe_released_blocks", + "blocks_waiting_for_store", + "abnormal_lease_reclaims", + "save_candidates", + "save_candidates_finished", + "save_committed", + "save_admitted", + "save_pin_budget_blocks", + "save_reserved_blocks", + "save_pinned_blocks", + "save_pinned_tokens", + "save_budget_available_blocks", + "save_budget_rejected", + "save_budget_rejected_blocks", + "save_budget_evicted", + "save_budget_evicted_blocks", + "save_oversized", + "deferred_free_requests", + "save_dropped_capacity", + "save_dropped_low_value", + "save_dropped_terminal_failure", + "save_dropped_stale", + "save_dropped_tokens_capacity", + "save_dropped_tokens_low_value", + "save_dropped_tokens_terminal_failure", + "save_dropped_tokens_stale", ) offload_totals = { key: sum(int(stats.get(key, 0)) for stats in offload_rank_stats) for key in offload_keys } + for key in ( + "save_candidate_wait_seconds", + "save_priority_score", + "save_inflight_wait_seconds", + "save_pinned_ratio", + ): + offload_totals[key] = max( + (float(stats.get(key, 0.0)) for stats in offload_rank_stats), + default=0.0, + ) return { "enabled": bool(rank_stats), diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index f0b703ff17..6ad191bfaa 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -561,6 +561,9 @@ def __init__( from atom.utils.forward_context import get_kvconnector self.kv_connector = get_kvconnector("scheduler", config) + bind_block_manager = getattr(self.kv_connector, "bind_block_manager", None) + if callable(bind_block_manager): + bind_block_manager(self.block_manager) from atom.distributed.kv_events import ( EventPublisher as _EventPublisher, diff --git a/docs/environment_variables.md b/docs/environment_variables.md index d0f0af9502..605d305d2c 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -204,6 +204,15 @@ discoverable from the central env reference despite bypassing the registry. |----------|------|---------|-------------| | **LMCACHE_EC_PIN_TIMEOUT_SEC** | float | LMCache's own (300) | LMCache's source-pin timeout. ATOM reads it only to derive the engine's save-abandon window (`pin + 30s`), so the two stay ordered — a lost store report is reclaimed only after LMCache would already have force-unpinned its source. Non-positive disables ATOM's reclamation. ATOM sets no default of its own; when unset it assumes LMCache's. | | **OFFLOAD_MAX_PENDING_SAVES** | int | **2**, flat, for the engine-side/state-tier reader (`scheduler.py`); `max(2, 2 × OFFLOAD_COPY_WORKERS)` for the KV-leg reader (`_offload_common.py`) | Bound on total in-flight offload transfers (running + queued) held before a SLOT snapshot or executor submission. A KV save and a state store both pin bytes out of the same pool while they run, so the KV leg and the K3 state tier share this one number rather than each carrying its own. Two readers compute it, though: the KV leg's canonical `_offload_common.max_pending_saves` derives the shown default from `OFFLOAD_COPY_WORKERS` and **raises** on an unparseable value, while the scheduler's state-tier reader (`_offload_max_pending_saves`) has a simpler fallback — a flat default of **2** (no `OFFLOAD_COPY_WORKERS` scaling) that **warns and uses 2** on an unparseable value rather than raising. Set the env to an explicit integer to pin both. | +| **OFFLOAD_SAVE_POLICY** | enum | `round_robin` | Chunked LMCache save admission policy, including `lmcache_mp` PAGE-only/native-state and DSV4. `round_robin` preserves the compatibility behavior. `priority` admits reusable prefixes by rank-local demand and dirty-token cost; a finished request that cannot reserve capacity immediately is dropped instead of pinning KV indefinitely. | +| **OFFLOAD_SAVE_MIN_OBSERVED_COUNT** | int | `2` | Priority mode's hard reuse-demand floor. Aging and finished-request release benefit do not bypass it. | +| **OFFLOAD_SAVE_MAX_PINNED_RATIO** | float | `0.20` | Maximum fraction of this scheduler-local physical KV pool reserved or pinned by admitted saves. Valid range is `0.0`–`0.30`; the pool is local to one DP scheduler and is not multiplied by TP size. | +| **OFFLOAD_SAVE_MAX_PINNED_BLOCKS** | int | unset | Optional nonnegative absolute PAGE-block clamp. The effective local budget is `min(floor(total_blocks × ratio), absolute)`. | +| **OFFLOAD_SAVE_AGING_WEIGHT** | float | `0.01` | Nonnegative per-second aging term in the priority score. | +| **OFFLOAD_SAVE_RELEASE_WEIGHT** | float | `1.0` | Nonnegative weight for GPU blocks released by completing a finished request's save. | +| **OFFLOAD_SAVE_DEMAND_BLOCK_TOKENS** | int | `8192` | Coarse cumulative-prefix boundary used by the independent rank-local demand tracker. | +| **OFFLOAD_SAVE_DEMAND_MAX_ENTRIES** | int | `65536` | Maximum number of coarse prefix-demand entries per rank. | +| **OFFLOAD_SAVE_DEMAND_TTL_SECONDS** | float | `600` | Positive quiet-time TTL for rank-local prefix-demand entries. | ## Profiling & debugging diff --git a/tests/entrypoints/test_metrics.py b/tests/entrypoints/test_metrics.py index 1e75ea49ff..1d473abb72 100644 --- a/tests/entrypoints/test_metrics.py +++ b/tests/entrypoints/test_metrics.py @@ -12,10 +12,12 @@ from __future__ import annotations import gc +from types import SimpleNamespace from prometheus_client import CollectorRegistry, generate_latest +from prometheus_client.parser import text_string_to_metric_families -from atom.entrypoints.openai.metrics import _gc_metrics +from atom.entrypoints.openai.metrics import AtomMetricsExporter, _gc_metrics def _render() -> str: @@ -82,3 +84,107 @@ def test_every_generation_is_labelled_rather_than_summed(): for generation in ("0", "1", "2"): assert f'atom:gc_collections_total{{generation="{generation}"}}' in exposition + + +def test_lmcache_save_admission_and_budget_metrics_are_exported(): + exporter = AtomMetricsExporter() + exporter.update( + { + "enabled": True, + "offload": { + "save_admitted": 7, + "save_candidates": 3, + "save_candidates_finished": 1, + "save_committed": 2, + "save_dropped_capacity": 4, + "save_dropped_tokens_capacity": 4096, + "save_pin_budget_blocks": 100, + "save_reserved_blocks": 20, + "save_pinned_blocks": 12, + "save_pinned_ratio": 0.32, + "save_budget_available_blocks": 68, + "save_budget_rejected": 5, + "save_budget_rejected_blocks": 40, + "save_budget_evicted": 2, + "save_budget_evicted_blocks": 16, + "save_oversized": 1, + "save_pinned_tokens": 3072, + "deferred_free_requests": 2, + }, + } + ) + + samples = [ + sample + for family in text_string_to_metric_families(exporter.render().decode()) + for sample in family.samples + if sample.name.startswith("atom:lmcache_save_") + or sample.name == "atom:lmcache_deferred_free_requests" + ] + by_name_and_labels = { + (sample.name, tuple(sorted(sample.labels.items()))): sample.value + for sample in samples + } + assert by_name_and_labels[("atom:lmcache_save_admitted_total", ())] == 7 + assert by_name_and_labels[("atom:lmcache_save_candidates", ())] == 3 + assert by_name_and_labels[("atom:lmcache_save_committed", ())] == 2 + assert ( + by_name_and_labels[ + ("atom:lmcache_save_dropped_total", (("reason", "capacity"),)) + ] + == 4 + ) + assert ( + by_name_and_labels[ + ( + "atom:lmcache_save_dropped_tokens_total", + (("reason", "capacity"),), + ) + ] + == 4096 + ) + assert by_name_and_labels[("atom:lmcache_save_pinned_blocks", ())] == 12 + assert by_name_and_labels[("atom:lmcache_save_reserved_blocks", ())] == 20 + assert by_name_and_labels[("atom:lmcache_save_pin_budget_blocks", ())] == 100 + assert by_name_and_labels[("atom:lmcache_save_pinned_ratio", ())] == 0.32 + assert by_name_and_labels[("atom:lmcache_save_budget_rejected_total", ())] == 5 + assert by_name_and_labels[("atom:lmcache_save_budget_evicted_total", ())] == 2 + assert by_name_and_labels[("atom:lmcache_deferred_free_requests", ())] == 2 + + +def test_lmcache_save_budget_aggregation_sums_dp_pools_but_uses_hottest_ratio(): + from atom.model_engine.llm_engine import LLMEngine + + engine = LLMEngine.__new__(LLMEngine) + engine.core_mgr = SimpleNamespace( + latest_metrics={ + 0: { + "enabled": True, + "offload": { + "save_pin_budget_blocks": 20, + "save_reserved_blocks": 4, + "save_pinned_blocks": 2, + "save_budget_available_blocks": 14, + "save_pinned_ratio": 0.06, + }, + }, + 1: { + "enabled": True, + "offload": { + "save_pin_budget_blocks": 20, + "save_reserved_blocks": 6, + "save_pinned_blocks": 3, + "save_budget_available_blocks": 11, + "save_pinned_ratio": 0.09, + }, + }, + }, + get_dp_router_statistics=dict, + ) + + offload = engine.get_metrics_statistics()["offload"] + assert offload["save_pin_budget_blocks"] == 40 + assert offload["save_reserved_blocks"] == 10 + assert offload["save_pinned_blocks"] == 5 + assert offload["save_budget_available_blocks"] == 25 + assert offload["save_pinned_ratio"] == 0.09 diff --git a/tests/test_lmcache_offload_connector.py b/tests/test_lmcache_offload_connector.py index e26df5a071..b876f80ee2 100644 --- a/tests/test_lmcache_offload_connector.py +++ b/tests/test_lmcache_offload_connector.py @@ -6,6 +6,7 @@ import logging import sys import threading +import time import types from collections import deque from contextlib import nullcontext @@ -85,6 +86,10 @@ SlotLoadSpec, SlotSaveSpec, ) +from atom.kv_transfer.offload.save_admission import ( + PrefixDemandTracker, + load_save_admission_config, +) from atom.model_engine.block_manager import BlockManager from atom.model_engine.scheduler import Scheduler from atom.model_engine.sequence import OffloadJointRecord, SequenceStatus @@ -157,6 +162,39 @@ def _scheduler() -> LMCacheOffloadConnectorScheduler: sched._save_tracker = {} sched._max_pending_saves = 4 sched._save_rr_last = None + sched._save_policy = "round_robin" + sched._save_min_observed_count = 2 + sched._save_aging_weight = 0.01 + sched._save_release_weight = 1.0 + sched._save_max_pinned_ratio = 0.20 + sched._save_max_pinned_blocks = None + sched._block_manager = None + sched._save_pin_total_blocks = 0 + sched._save_pin_budget_blocks = None + sched._save_block_reservations = {} + sched._prefix_demand = PrefixDemandTracker(block_tokens=4) + sched._save_demand_keys = {} + sched._save_candidate_since = {} + sched._save_candidate_generation = {} + sched._save_candidate_nonce = 0 + sched._save_committed = {} + sched._finished_save_requests = {} + sched._finished_save_failed = set() + sched._save_inflight_since = {} + sched._save_operation_owners = {} + sched.total_save_admitted = 0 + sched.total_save_budget_rejected = 0 + sched.total_save_budget_rejected_blocks = 0 + sched.total_save_budget_evicted = 0 + sched.total_save_budget_evicted_blocks = 0 + sched.total_save_oversized = 0 + sched._save_drop_totals = { + "capacity": 0, + "low_value": 0, + "terminal_failure": 0, + "stale": 0, + } + sched._save_drop_token_totals = dict.fromkeys(sched._save_drop_totals, 0) sched._save_nonce = 0 sched._load_nonce = 0 sched._load_lifecycles = {} @@ -198,6 +236,57 @@ def _stateful_scheduler(hit: int) -> LMCacheOffloadConnectorScheduler: return sched +def _priority_scheduler(*, capacity: int = 2) -> LMCacheOffloadConnectorScheduler: + sched = _scheduler() + sched._save_policy = "priority" + sched._max_pending_saves = capacity + sched._save_min_observed_count = 2 + return sched + + +def _priority_seq( + sched: LMCacheOffloadConnectorScheduler, + *, + req_id: int, + token_ids: list[int], + computed: int | None = None, +) -> SimpleNamespace: + seq = SimpleNamespace( + id=req_id, + token_ids=token_ids, + block_table=list( + range((len(token_ids) + sched.block_size - 1) // sched.block_size) + ), + num_prompt_tokens=len(token_ids), + num_cached_tokens=len(token_ids) if computed is None else computed, + has_per_req_cache=False, + ) + sched.update_state_after_alloc(seq) + return seq + + +def _set_save_budget( + sched: LMCacheOffloadConnectorScheduler, + *, + total_blocks: int, + ratio: float = 0.30, + absolute: int | None = None, +) -> None: + sched._save_max_pinned_ratio = ratio + sched._save_max_pinned_blocks = absolute + sched.bind_block_manager(SimpleNamespace(total_allocatable_kv_blocks=total_blocks)) + + +def _observe_priority( + sched: LMCacheOffloadConnectorScheduler, + seq: SimpleNamespace, + count: int, +) -> None: + # update_state_after_alloc records the first observation. + for _ in range(max(0, count - 1)): + sched._prefix_demand.observe(seq.token_ids, seq.num_prompt_tokens) + + def _stateful_seq( *, req_id: int, @@ -259,6 +348,23 @@ def test_bounded_commit_index_touching_duplicate_refreshes_recency(): assert len(index) == 2 +def test_prefix_demand_tracker_counts_without_publishing_route_ownership(): + now = [10.0] + tracker = PrefixDemandTracker( + block_tokens=4, + ttl_seconds=5, + clock=lambda: now[0], + ) + tokens = list(range(8)) + + keys = tracker.observe(tokens, len(tokens)) + tracker.observe(tokens, len(tokens)) + + assert tracker.heat(keys, max_tokens=8) == (2, 8) + now[0] = 16.0 + assert tracker.heat(keys, max_tokens=8) == (0, 0) + + @pytest.mark.parametrize( "value", [True, False, 0, -1, 1.5, "1.5", "not-an-int"], @@ -4428,6 +4534,449 @@ def test_dsv4_save_admission_is_bounded_and_round_robin(): assert sched._may_emit_save() is False +def test_dsv4_priority_save_admits_hotter_prefix_first(): + sched = _priority_scheduler(capacity=1) + hot = _priority_seq(sched, req_id=200, token_ids=list(range(8))) + _priority_seq(sched, req_id=201, token_ids=list(range(20, 28))) + sched._prefix_demand.observe(hot.token_ids, hot.num_prompt_tokens) + + meta = sched.build_connector_meta() + + assert [request.req_id for request in meta.requests] == [200] + assert sched.total_save_admitted == 1 + + +def test_dsv4_priority_save_prefers_less_dirty_work_at_equal_value(): + sched = _priority_scheduler(capacity=1) + short = _priority_seq(sched, req_id=202, token_ids=list(range(4))) + long = _priority_seq(sched, req_id=203, token_ids=list(range(20, 28))) + sched._prefix_demand.observe(short.token_ids, short.num_prompt_tokens) + sched._prefix_demand.observe(long.token_ids, long.num_prompt_tokens) + + meta = sched.build_connector_meta() + + assert [request.req_id for request in meta.requests] == [202] + + +def test_dsv4_priority_save_order_is_stable_for_equal_candidates(): + sched = _priority_scheduler(capacity=1) + later_id = _priority_seq(sched, req_id=205, token_ids=list(range(8))) + earlier_id = _priority_seq(sched, req_id=204, token_ids=list(range(20, 28))) + sched._prefix_demand.observe(later_id.token_ids, later_id.num_prompt_tokens) + sched._prefix_demand.observe(earlier_id.token_ids, earlier_id.num_prompt_tokens) + # Remove enqueue time as a tie-break so the stable request/generation key + # is the deciding order regardless of tracker insertion order. + same_time = time.monotonic() + sched._save_candidate_since["205"] = (later_id, same_time) + sched._save_candidate_since["204"] = (earlier_id, same_time) + + meta = sched.build_connector_meta() + + assert [request.req_id for request in meta.requests] == [204] + + +def test_dsv4_priority_aging_can_break_a_value_tie_without_bypassing_floor(): + sched = _priority_scheduler(capacity=1) + sched._save_aging_weight = 100 + old = _priority_seq(sched, req_id=212, token_ids=list(range(8))) + newer = _priority_seq(sched, req_id=213, token_ids=list(range(20, 28))) + sched._prefix_demand.observe(old.token_ids, old.num_prompt_tokens) + sched._prefix_demand.observe(newer.token_ids, newer.num_prompt_tokens) + sched._prefix_demand.observe(newer.token_ids, newer.num_prompt_tokens) + now = time.monotonic() + sched._save_candidate_since["212"] = (old, now - 1) + sched._save_candidate_since["213"] = (newer, now) + + meta = sched.build_connector_meta() + + assert [request.req_id for request in meta.requests] == [212] + + +def test_save_pin_ratio_is_validated_and_capped_at_thirty_percent(monkeypatch): + monkeypatch.setenv("OFFLOAD_SAVE_MAX_PINNED_RATIO", "0.30") + assert load_save_admission_config().max_pinned_ratio == pytest.approx(0.30) + + monkeypatch.setenv("OFFLOAD_SAVE_MAX_PINNED_RATIO", "0.300001") + with pytest.raises(ValueError, match="at most 0.30"): + load_save_admission_config() + + +def test_dsv4_save_budget_uses_dp_local_pool_and_absolute_clamp(): + sched = _priority_scheduler(capacity=8) + sched._save_max_pinned_ratio = 0.30 + sched._save_max_pinned_blocks = 17 + manager = SimpleNamespace( + total_allocatable_kv_blocks=100, + tensor_parallel_world_size=8, + ) + + sched.bind_block_manager(manager) + + assert sched._save_pin_total_blocks == 100 + assert sched._save_pin_budget_blocks == 17 + + +def test_dsv4_zero_save_budget_rejects_without_deferring_free(): + sched = _priority_scheduler(capacity=2) + _set_save_budget(sched, total_blocks=10, absolute=0) + seq = _priority_seq(sched, req_id=220, token_ids=list(range(8))) + _observe_priority(sched, seq, 2) + + sched.request_finished(seq) + + assert "220" not in sched._save_tracker + assert "220" not in sched._save_committed + assert sched.should_defer_free(seq) is False + stats = sched.get_statistics() + assert stats["save_budget_rejected"] == 1 + assert stats["save_budget_rejected_blocks"] == 2 + assert stats["save_oversized"] == 1 + + +def test_dsv4_round_robin_finished_budget_rejection_stops_deferring_free(): + sched = _scheduler() + _set_save_budget(sched, total_blocks=10, absolute=0) + seq = _priority_seq(sched, req_id=238, token_ids=list(range(8))) + sched.request_finished(seq) + assert sched.should_defer_free(seq) is True + + assert sched.build_connector_meta().requests == [] + + assert sched.should_defer_free(seq) is False + assert "238" not in sched._save_tracker + assert sched.total_save_budget_rejected == 1 + + +def test_dsv4_save_budget_counts_shared_physical_ids_once(): + sched = _priority_scheduler(capacity=2) + _set_save_budget(sched, total_blocks=10, absolute=3) + first = _priority_seq(sched, req_id=221, token_ids=list(range(8))) + first.block_table = [1, 2] + second = _priority_seq(sched, req_id=222, token_ids=list(range(20, 28))) + second.block_table = [2, 3] + _observe_priority(sched, first, 2) + _observe_priority(sched, second, 2) + + sched.request_finished(first) + sched.request_finished(second) + + assert set(sched._save_committed) == {"221", "222"} + stats = sched.get_statistics() + assert stats["save_pinned_blocks"] == 3 + assert stats["save_budget_available_blocks"] == 0 + + +def test_dsv4_higher_priority_candidate_replaces_lower_committed_save(): + sched = _priority_scheduler(capacity=2) + sched._save_aging_weight = 0 + _set_save_budget(sched, total_blocks=10, absolute=2) + low = _priority_seq(sched, req_id=223, token_ids=list(range(8))) + low.block_table = [1, 2] + _observe_priority(sched, low, 2) + sched.request_finished(low) + + high = _priority_seq(sched, req_id=224, token_ids=list(range(20, 28))) + high.block_table = [3, 4] + _observe_priority(sched, high, 3) + sched.request_finished(high) + + assert sched._save_committed == {"224": high} + assert "223" not in sched._save_tracker + assert sched.should_defer_free(low) is False + assert sched.total_save_budget_evicted == 1 + assert sched.total_save_budget_evicted_blocks == 2 + + +def test_dsv4_priority_replacement_can_evict_multiple_committed_saves(): + sched = _priority_scheduler(capacity=3) + sched._save_aging_weight = 0 + _set_save_budget(sched, total_blocks=20, absolute=4) + low_a = _priority_seq(sched, req_id=225, token_ids=list(range(8))) + low_a.block_table = [1, 2] + low_b = _priority_seq(sched, req_id=226, token_ids=list(range(20, 28))) + low_b.block_table = [3, 4] + for seq in (low_a, low_b): + _observe_priority(sched, seq, 2) + sched.request_finished(seq) + + high = _priority_seq(sched, req_id=227, token_ids=list(range(40, 52))) + high.block_table = [5, 6, 7] + _observe_priority(sched, high, 4) + sched.request_finished(high) + + assert sched._save_committed == {"227": high} + assert "225" not in sched._save_tracker + assert "226" not in sched._save_tracker + assert sched.total_save_budget_evicted == 2 + assert sched.total_save_budget_evicted_blocks == 4 + + +def test_dsv4_failed_replacement_simulation_evicts_nothing(): + sched = _priority_scheduler(capacity=3) + sched._save_aging_weight = 0 + _set_save_budget(sched, total_blocks=20, absolute=4) + low_a = _priority_seq(sched, req_id=228, token_ids=list(range(8))) + low_a.block_table = [1, 2] + low_b = _priority_seq(sched, req_id=229, token_ids=list(range(20, 28))) + low_b.block_table = [3, 4] + for seq in (low_a, low_b): + _observe_priority(sched, seq, 2) + sched.request_finished(seq) + + oversized = _priority_seq(sched, req_id=230, token_ids=list(range(40, 60))) + oversized.block_table = [5, 6, 7, 8, 9] + _observe_priority(sched, oversized, 5) + sched.request_finished(oversized) + + assert set(sched._save_committed) == {"228", "229"} + assert "230" not in sched._save_tracker + assert sched.total_save_budget_evicted == 0 + assert sched.total_save_budget_rejected == 1 + assert sched.total_save_oversized == 1 + + +def test_dsv4_equal_priority_candidate_cannot_replace_committed_save(): + sched = _priority_scheduler(capacity=2) + sched._save_aging_weight = 0 + _set_save_budget(sched, total_blocks=10, absolute=2) + first = _priority_seq(sched, req_id=231, token_ids=list(range(8))) + first.block_table = [1, 2] + _observe_priority(sched, first, 2) + sched.request_finished(first) + + equal = _priority_seq(sched, req_id=232, token_ids=list(range(20, 28))) + equal.block_table = [3, 4] + _observe_priority(sched, equal, 2) + sched.request_finished(equal) + + assert sched._save_committed == {"231": first} + assert "232" not in sched._save_tracker + assert sched.total_save_budget_evicted == 0 + + +def test_dsv4_inflight_page_slot_save_is_never_an_eviction_victim(): + sched = _stateful_scheduler(hit=0) + sched._save_policy = "priority" + sched._save_min_observed_count = 2 + sched._save_aging_weight = 0 + sched._max_pending_saves = 2 + _set_save_budget(sched, total_blocks=200, absolute=32) + low = _stateful_seq( + req_id=233, + num_prompt_tokens=8192, + num_cached_tokens=8192, + group=2, + ) + low.block_table = list(range(32)) + sched.update_state_after_alloc(low) + _observe_priority(sched, low, 2) + request = sched.build_connector_meta().requests[0] + assert request.save_spec is not None + assert request.slot_save_spec is not None + + high = _stateful_seq( + req_id=234, + num_prompt_tokens=8192, + num_cached_tokens=8192, + group=3, + ) + high.token_ids = list(range(10000, 18192)) + high.block_table = list(range(32, 64)) + sched.update_state_after_alloc(high) + _observe_priority(sched, high, 4) + sched.request_finished(high) + + assert request.save_operation in sched._save_inflight["233"] + assert sched._sidecar_save_inflight["233"][0] == request.save_operation + assert request.save_operation in sched._save_block_reservations + assert "234" not in sched._save_tracker + assert sched.total_save_budget_evicted == 0 + + +def test_dsv4_request_id_reuse_releases_old_committed_reservation(): + sched = _priority_scheduler(capacity=2) + _set_save_budget(sched, total_blocks=10, absolute=2) + old = _priority_seq(sched, req_id=235, token_ids=list(range(8))) + _observe_priority(sched, old, 2) + sched.request_finished(old) + assert "235" in sched._save_block_reservations + + new = _priority_seq(sched, req_id=235, token_ids=list(range(20, 28))) + + assert sched._save_tracker["235"][0] is new + assert "235" not in sched._save_committed + assert "235" not in sched._save_block_reservations + assert sched.should_defer_free(old) is False + + +def test_dsv4_page_slot_replacement_is_atomic(): + sched = _stateful_scheduler(hit=0) + sched._save_policy = "priority" + sched._save_min_observed_count = 2 + sched._save_aging_weight = 0 + sched._max_pending_saves = 2 + _set_save_budget(sched, total_blocks=200, absolute=32) + low = _stateful_seq( + req_id=236, + num_prompt_tokens=8192, + num_cached_tokens=8192, + group=2, + ) + low.block_table = list(range(32)) + sched.update_state_after_alloc(low) + _observe_priority(sched, low, 2) + assert sched._build_save_candidate("236").sidecar_candidate is not None + sched.request_finished(low) + + high = _stateful_seq( + req_id=237, + num_prompt_tokens=8192, + num_cached_tokens=8192, + group=3, + ) + high.token_ids = list(range(10000, 18192)) + high.block_table = list(range(32, 64)) + sched.update_state_after_alloc(high) + _observe_priority(sched, high, 4) + sched.request_finished(high) + + assert sched._save_committed == {"237": high} + assert "236" not in sched._save_tracker + assert "236" not in sched._sidecar_save_inflight + assert set(sched._save_block_reservations) == {"237"} + + +def test_dsv4_priority_finished_candidate_commits_then_dispatches(): + sched = _priority_scheduler(capacity=1) + seq = _priority_seq(sched, req_id=206, token_ids=list(range(8))) + sched._prefix_demand.observe(seq.token_ids, seq.num_prompt_tokens) + + sched.request_finished(seq) + + assert sched._save_committed == {"206": seq} + assert sched.should_defer_free(seq) is True + stats = sched.get_statistics() + assert stats["save_admitted"] == 1 + assert stats["save_committed"] == 1 + assert stats["save_pinned_blocks"] == 2 + assert stats["deferred_free_requests"] == 1 + request = sched.build_connector_meta().requests[0] + assert request.req_id == 206 + assert "206" not in sched._save_committed + assert request.save_operation in sched._save_inflight["206"] + + +def test_dsv4_priority_finished_candidate_drops_when_capacity_is_full(): + sched = _priority_scheduler(capacity=1) + busy = _priority_seq(sched, req_id=207, token_ids=list(range(8))) + sched._prefix_demand.observe(busy.token_ids, busy.num_prompt_tokens) + assert sched.build_connector_meta().requests[0].req_id == 207 + + seq = _priority_seq(sched, req_id=208, token_ids=list(range(20, 28))) + sched._prefix_demand.observe(seq.token_ids, seq.num_prompt_tokens) + sched.request_finished(seq) + + assert "208" not in sched._save_tracker + assert sched.should_defer_free(seq) is False + stats = sched.get_statistics() + assert stats["save_dropped_capacity"] == 1 + assert stats["save_dropped_tokens_capacity"] == 8 + + +def test_dsv4_priority_finished_inflight_residual_is_committed_once(): + sched = _priority_scheduler(capacity=1) + seq = _priority_seq( + sched, + req_id=209, + token_ids=list(range(8)), + computed=4, + ) + sched._prefix_demand.observe(seq.token_ids, seq.num_prompt_tokens) + first = sched.build_connector_meta().requests[0] + + seq.num_cached_tokens = 8 + sched.request_finished(seq) + assert sched.should_defer_free(seq) is True + + sched.connector_completion( + ConnectorCompletion(DSV4_PAGE_SAVE_CHANNEL, first.save_operation, True) + ) + sched.save_finished(first.save_operation) + + assert sched._save_committed == {"209": seq} + second = sched.build_connector_meta().requests[0] + assert second.save_spec.skip_leading_tokens == 4 + assert sched.total_save_admitted == 2 + + +def test_dsv4_priority_failed_finished_save_drops_residual_instead_of_retrying(): + sched = _priority_scheduler(capacity=1) + seq = _priority_seq( + sched, + req_id=210, + token_ids=list(range(8)), + computed=4, + ) + sched._prefix_demand.observe(seq.token_ids, seq.num_prompt_tokens) + first = sched.build_connector_meta().requests[0] + seq.num_cached_tokens = 8 + sched.request_finished(seq) + + sched.connector_completion( + ConnectorCompletion(DSV4_PAGE_SAVE_CHANNEL, first.save_operation, False) + ) + sched.save_finished(first.save_operation) + + assert "210" not in sched._save_tracker + assert "210" not in sched._save_committed + assert sched.should_defer_free(seq) is False + assert sched.get_statistics()["save_dropped_terminal_failure"] == 1 + + +def test_dsv4_priority_late_failure_does_not_touch_reused_request_id(): + sched = _priority_scheduler(capacity=1) + old = _priority_seq(sched, req_id=211, token_ids=list(range(8)), computed=4) + sched._prefix_demand.observe(old.token_ids, old.num_prompt_tokens) + operation = sched.build_connector_meta().requests[0].save_operation + old.num_cached_tokens = 8 + sched.request_finished(old) + + new = _priority_seq( + sched, + req_id=211, + token_ids=list(range(20, 28)), + computed=8, + ) + sched.connector_completion( + ConnectorCompletion(DSV4_PAGE_SAVE_CHANNEL, operation, False) + ) + sched.save_finished(operation) + + assert sched._save_tracker["211"] == [new, 0] + assert "211" not in sched._finished_save_failed + + +def test_dsv4_priority_drop_removes_page_and_sidecar_candidate_together(): + sched = _stateful_scheduler(hit=0) + sched._save_policy = "priority" + sched._save_min_observed_count = 2 + seq = _stateful_seq( + req_id=214, + num_prompt_tokens=8192, + num_cached_tokens=8192, + group=2, + ) + sched.update_state_after_alloc(seq) + + sched.request_finished(seq) + + assert "214" not in sched._save_tracker + assert "214" not in sched._save_committed + assert "214" not in sched._sidecar_save_inflight + assert sched._committed_sidecar_hashes == set() + assert sched.should_defer_free(seq) is False + + def test_dsv4_page_and_sidecar_share_one_admission_slot(): sched = _stateful_scheduler(hit=0) sched._max_pending_saves = 1 @@ -4812,7 +5361,19 @@ def test_scheduler_offload_statistics_are_cumulative(): sched.load_finished("1") sched.save_finished("2") - assert sched.get_statistics() == { + stats = sched.get_statistics() + assert { + key: stats[key] + for key in ( + "load_requests", + "loaded_tokens", + "load_failures", + "save_requests", + "saved_tokens", + "loads_pending", + "saves_pending", + ) + } == { "load_requests": 1, "loaded_tokens": 8192, "load_failures": 0, @@ -5077,6 +5638,7 @@ def _k3_scheduler() -> KimiK3OffloadScheduler: s = KimiK3OffloadScheduler.__new__(KimiK3OffloadScheduler) s.chunk_size = 256 s._save_inflight = {} + s._save_committed = {} s._save_tracker = {} s._save_rr_last = None s._pending_state_loads = [] @@ -5405,6 +5967,9 @@ def test_no_more_saves_go_out_than_the_pool_can_afford_to_pin(): s._save_inflight = {"1": object()} assert s._may_emit_save() is True + s._save_committed = {"2": object()} + assert s._may_emit_save() is False + s._save_committed = {} s._save_stalled = True assert s._may_emit_save() is False From 52062aa09585bd9b5c31a6f05e4cdf648e687a20 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Mon, 21 Sep 2026 08:11:07 +0000 Subject: [PATCH 4/4] refactor(offload): make priority save admission mandatory --- atom/kv_transfer/offload/README.md | 7 +- atom/kv_transfer/offload/chunked_scheduler.py | 193 +++++---------- .../offload/hybrid/dsv4/connector.py | 233 ++++++------------ atom/kv_transfer/offload/save_admission.py | 21 +- docs/environment_variables.md | 3 +- tests/test_lmcache_offload_connector.py | 80 +----- 6 files changed, 157 insertions(+), 380 deletions(-) diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index 79a8bc593c..4a5d806652 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -878,8 +878,7 @@ 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_POLICY` | `round_robin` | Chunked LMCache save selection, including `lmcache_mp` and DSV4. `priority` uses reuse demand and dirty cost, and hard-drops a finished request's unadmitted save instead of pinning its KV indefinitely. | -| `OFFLOAD_SAVE_MIN_OBSERVED_COUNT` | 2 | Priority mode's hard minimum rank-local prefix observations. | +| `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. | @@ -910,7 +909,7 @@ still falls back safely to per-group ID preparation and blocking host copies. `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. -In priority mode the chunked and DSV4 save lifecycle is explicitly +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 @@ -921,7 +920,7 @@ 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. Priority admission may replace one or more lower-scored, +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 diff --git a/atom/kv_transfer/offload/chunked_scheduler.py b/atom/kv_transfer/offload/chunked_scheduler.py index 9a2c8560c3..07a312fca8 100644 --- a/atom/kv_transfer/offload/chunked_scheduler.py +++ b/atom/kv_transfer/offload/chunked_scheduler.py @@ -124,12 +124,7 @@ def __init__( # prefix is stored to LMCache once prefill computes it # (seq.prefix_hashes_published flips True), chunk by chunk. self._save_tracker: dict[str, list] = {} - # Round-robin cursor over `_save_tracker`: the last sid that emitted a - # save. Subclasses may bound the number of outstanding saves through - # `_may_emit_save`; resuming after this sid prevents starvation. - self._save_rr_last: str | None = None save_admission = load_save_admission_config() - self._save_policy = save_admission.policy self._save_min_observed_count = save_admission.min_observed_count self._save_aging_weight = save_admission.aging_weight self._save_release_weight = save_admission.release_weight @@ -339,14 +334,10 @@ def update_state_after_alloc(self, seq) -> None: if entry is None or entry[0] is not seq: self._save_tracker[sid] = [seq, initial_saved] now = time.monotonic() - demand_keys = ( - self._prefix_demand.observe( - seq.token_ids, - int(seq.num_prompt_tokens), - now, - ) - if self._save_policy == "priority" - else () + demand_keys = self._prefix_demand.observe( + seq.token_ids, + int(seq.num_prompt_tokens), + now, ) self._save_demand_keys[sid] = (seq, demand_keys) self._save_candidate_generation[sid] = ( @@ -418,8 +409,6 @@ def _may_emit_save(self) -> bool: def _ensure_save_admission_state(self) -> None: """Supply compatibility defaults for lightweight/manual schedulers.""" - if not hasattr(self, "_save_policy"): - self._save_policy = "round_robin" if not hasattr(self, "_save_min_observed_count"): self._save_min_observed_count = 2 if not hasattr(self, "_save_aging_weight"): @@ -750,7 +739,7 @@ def _commit_finished_save(self, candidate: _ChunkedSaveCandidate) -> bool: return False return self._try_reserve_save_candidate( candidate, - allow_eviction=self._save_policy == "priority", + allow_eviction=True, ) def _forget_save_candidate(self, sid: str, seq) -> None: @@ -908,7 +897,7 @@ def _emit_save_candidate( if hasattr(seq, "_offload_finished_block_ids") and not seq.block_table: late_source = self._late_save_source(seq, candidate.saved, aligned) if late_source is None: - if self._save_policy == "priority" and candidate.finished: + if candidate.finished: self.drop_unadmitted_save(seq, reason="stale") else: self._save_committed.pop(sid, None) @@ -943,13 +932,12 @@ def _emit_save_candidate( logger.debug( "[OFFLOAD-SAVE-EMIT] seq=%s computed=%d num_prompt=%d " - "aligned=%d saved=%d policy=%s observed=%d", + "aligned=%d saved=%d observed=%d", seq.id, computed, int(seq.num_prompt_tokens), aligned, candidate.saved, - self._save_policy, candidate.observed_count, ) self._track_save_statistics(save_operation, aligned - candidate.saved) @@ -1063,7 +1051,7 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: for sid in self._lookup_in_step if sid in loading_sids or sid not in self._load_specs ] - # Saves. Priority mode admits/replaces the complete candidate set before + # Saves. Admit/replace the complete candidate set before # dispatch so a newly hotter candidate can still evict lower-value # committed work. Dispatching first would make that work inflight and # therefore correctly, but prematurely, non-evictable. @@ -1081,80 +1069,47 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: self._release_save_reservation(sid, owner=seq) self._settle_finished_save(sid) - tracker_sids = list(self._save_tracker.keys()) - if self._save_policy == "round_robin": - if tracker_sids and self._save_rr_last in self._save_tracker: - start = (tracker_sids.index(self._save_rr_last) + 1) % len( - tracker_sids - ) - tracker_sids = tracker_sids[start:] + tracker_sids[:start] - candidates = [ - candidate - for sid in tracker_sids - if (candidate := self._build_save_candidate(sid, now=now)) - is not None - ] - else: - candidates = [ - candidate - for sid in tracker_sids - if (candidate := self._build_save_candidate(sid, now=now)) - is not None - and self._meets_save_value_threshold(candidate) - ] - candidates.sort( - key=lambda candidate: self._candidate_sort_key(candidate, now) - ) + candidates = [ + candidate + for sid in self._save_tracker + if (candidate := self._build_save_candidate(sid, now=now)) is not None + and self._meets_save_value_threshold(candidate) + ] + candidates.sort( + key=lambda candidate: self._candidate_sort_key(candidate, now) + ) - if self._save_policy == "priority": - for candidate in candidates: - sid = candidate.sid - if sid in self._save_committed or sid in self._save_inflight: - continue - if sid in self._reqs_need_recv or sid in loading_sids: - continue - if not self._try_reserve_save_candidate( - candidate, - allow_eviction=True, - now=now, - ): - self.drop_unadmitted_save(candidate.seq, reason="capacity") - - committed = [] - for sid, seq in list(self._save_committed.items()): - candidate = self._build_save_candidate(sid, now=now) - if candidate is None or candidate.seq is not seq: - self._save_committed.pop(sid, None) - self._release_save_reservation(sid, owner=seq) - continue - committed.append(candidate) - committed.sort( - key=lambda candidate: self._candidate_sort_key(candidate, now) - ) - for candidate in committed: - sid = candidate.sid - if sid in self._reqs_need_recv or sid in loading_sids: - continue - if sid in self._save_inflight: - continue - self._emit_save_candidate(meta, candidate) - else: - for candidate in candidates: - sid = candidate.sid - if not self._may_emit_save(): - break - if sid in self._save_committed or sid in self._save_inflight: - continue - if sid in self._reqs_need_recv or sid in loading_sids: - continue - if not self._try_reserve_save_candidate( - candidate, - allow_eviction=False, - now=now, - ): - break - if self._emit_save_candidate(meta, candidate): - self._save_rr_last = sid + for candidate in candidates: + sid = candidate.sid + if sid in self._save_committed or sid in self._save_inflight: + continue + if sid in self._reqs_need_recv or sid in loading_sids: + continue + if not self._try_reserve_save_candidate( + candidate, + allow_eviction=True, + now=now, + ): + self.drop_unadmitted_save(candidate.seq, reason="capacity") + + committed = [] + for sid, seq in list(self._save_committed.items()): + candidate = self._build_save_candidate(sid, now=now) + if candidate is None or candidate.seq is not seq: + self._save_committed.pop(sid, None) + self._release_save_reservation(sid, owner=seq) + continue + committed.append(candidate) + committed.sort( + key=lambda candidate: self._candidate_sort_key(candidate, now) + ) + for candidate in committed: + sid = candidate.sid + if sid in self._reqs_need_recv or sid in loading_sids: + continue + if sid in self._save_inflight: + continue + self._emit_save_candidate(meta, candidate) dispatched = set(meta.lookup_requests_in_step) for sid in dispatched: self._lookup_results.pop(sid, None) @@ -1521,11 +1476,7 @@ def _store_finished(self, operation: SaveOperationId, *, succeeded: bool) -> Non # block-map bookkeeping forever. self._release_operation_lease(operation) self._save_inflight_since.pop(operation, None) - if ( - self._save_policy == "priority" - and owner is not None - and self._finished_save_requests.get(sid) is owner - ): + if owner is not None and self._finished_save_requests.get(sid) is owner: self._settle_finished_save(sid, failed=not succeeded) self._finish_retired_request(sid) @@ -1728,41 +1679,25 @@ def request_finished(self, seq) -> None: self._load_lifecycles.pop(sid, None) entry = self._save_tracker.get(sid) if entry is not None and entry[0] is seq: - if self._save_policy == "priority": - self._finished_save_requests[sid] = seq - # Freeze the final source identity before request teardown. - # Priority reservations dispatch on a later metadata build and - # reacquire only the still-canonical prefix blocks. - if self._early_release: - seq._offload_finished_cached_tokens = min( - int(getattr(seq, "num_cached_tokens", 0)), - int(seq.num_prompt_tokens), - ) - if sid in self._save_inflight or self._save_committed.get(sid) is seq: - pass - else: - candidate = self._build_save_candidate(sid) - if candidate is None: - self._settle_finished_save(sid) - elif not self._meets_save_value_threshold(candidate): - self.drop_unadmitted_save(seq, reason="low_value") - elif not self._commit_finished_save(candidate): - self.drop_unadmitted_save(seq, reason="capacity") - elif self._early_release: - # Freeze the final computed frontier before BlockManager - # clears it during partial deallocation. Keep the tracker when - # a final chunk still needs emission; a later metadata build - # uses the frozen block table recorded by - # `protected_block_ids`. + self._finished_save_requests[sid] = seq + # Freeze the final source identity before request teardown. A + # reservation dispatches on a later metadata build and reacquires + # only the still-canonical prefix blocks. + if self._early_release: seq._offload_finished_cached_tokens = min( int(getattr(seq, "num_cached_tokens", 0)), int(seq.num_prompt_tokens), ) - if not self.should_defer_free(seq): - self._save_tracker.pop(sid, None) - elif not self.should_defer_free(seq): - self._save_tracker.pop(sid, None) - self._forget_save_candidate(sid, seq) + if sid in self._save_inflight or self._save_committed.get(sid) is seq: + pass + else: + candidate = self._build_save_candidate(sid) + if candidate is None: + self._settle_finished_save(sid) + elif not self._meets_save_value_threshold(candidate): + self.drop_unadmitted_save(seq, reason="low_value") + elif not self._commit_finished_save(candidate): + self.drop_unadmitted_save(seq, reason="capacity") if hasattr(seq, "_load_operation"): delattr(seq, "_load_operation") diff --git a/atom/kv_transfer/offload/hybrid/dsv4/connector.py b/atom/kv_transfer/offload/hybrid/dsv4/connector.py index c3c560c6dc..a676a3a984 100644 --- a/atom/kv_transfer/offload/hybrid/dsv4/connector.py +++ b/atom/kv_transfer/offload/hybrid/dsv4/connector.py @@ -1918,17 +1918,12 @@ def __init__(self, config) -> None: self._max_pending_saves = max_pending_saves( int(os.environ.get("OFFLOAD_COPY_WORKERS", "1") or 1) ) - # Resume after the last admitted request when capacity becomes free so - # a long request at the head of the insertion-ordered tracker cannot - # monopolize the bounded save queue. - self._save_rr_last: str | None = None - # Priority mode separates a lightweight candidate from a capacity- + # Save admission separates a lightweight candidate from a capacity- # holding reservation. This matters at request finish: connector # metadata for the current step has already been built, so the request # cannot be dispatched immediately. A committed entry reserves one # worker slot until the next metadata build emits it. save_admission = load_save_admission_config() - self._save_policy = save_admission.policy self._save_max_pinned_ratio = save_admission.max_pinned_ratio self._save_max_pinned_blocks = save_admission.max_pinned_blocks self._block_manager = None @@ -1937,20 +1932,14 @@ def __init__(self, config) -> None: self._save_block_reservations: dict[ str | SaveOperationId, SaveBlockReservation ] = {} - if self._save_policy == "priority": - self._save_min_observed_count = save_admission.min_observed_count - self._save_aging_weight = save_admission.aging_weight - self._save_release_weight = save_admission.release_weight - self._prefix_demand = PrefixDemandTracker( - block_tokens=save_admission.demand_block_tokens, - max_entries=save_admission.demand_max_entries, - ttl_seconds=save_admission.demand_ttl_seconds, - ) - else: - self._save_min_observed_count = 2 - self._save_aging_weight = 0.01 - self._save_release_weight = 1.0 - self._prefix_demand = PrefixDemandTracker() + self._save_min_observed_count = save_admission.min_observed_count + self._save_aging_weight = save_admission.aging_weight + self._save_release_weight = save_admission.release_weight + self._prefix_demand = PrefixDemandTracker( + block_tokens=save_admission.demand_block_tokens, + max_entries=save_admission.demand_max_entries, + ttl_seconds=save_admission.demand_ttl_seconds, + ) self._save_demand_keys: dict[ str, tuple[object, tuple[PrefixDemandKey, ...]] ] = {} @@ -2182,14 +2171,10 @@ def update_state_after_alloc(self, seq) -> None: if entry is None or entry[0] is not seq: self._save_tracker[sid] = [seq, initial_saved] now = time.monotonic() - demand_keys = ( - self._prefix_demand.observe( - seq.token_ids, - int(seq.num_prompt_tokens), - now, - ) - if self._save_policy == "priority" - else () + demand_keys = self._prefix_demand.observe( + seq.token_ids, + int(seq.num_prompt_tokens), + now, ) self._save_demand_keys[sid] = (seq, demand_keys) self._save_candidate_generation[sid] = ( @@ -2642,7 +2627,7 @@ def _commit_finished_save(self, candidate: _SaveCandidate) -> bool: return False return self._try_reserve_save_candidate( candidate, - allow_eviction=self._save_policy == "priority", + allow_eviction=True, ) def _forget_save_candidate(self, sid: str, seq) -> None: @@ -2754,14 +2739,13 @@ def _emit_save_candidate( ) logger.debug( "[OFFLOAD-SAVE-EMIT] seq=%s computed=%d num_prompt=%d " - "aligned=%d saved=%d sidecar=%s policy=%s observed=%d", + "aligned=%d saved=%d sidecar=%s observed=%d", seq.id, candidate.computed, int(seq.num_prompt_tokens), candidate.aligned, candidate.saved, candidate.sidecar_candidate, - self._save_policy, candidate.observed_count, ) save_operation = self._next_save_operation(seq) @@ -3002,7 +2986,7 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: meta.lookup_requests_in_step = [ sid for sid in self._lookup_in_step if sid not in self._handoff_loads ] - # Saves. Priority mode admits/replaces the complete candidate set before + # Saves. Admit/replace the complete candidate set before # dispatch so a newly hotter candidate can still evict lower-value # committed work. Dispatching first would make that work inflight and # therefore correctly, but prematurely, non-evictable. @@ -3020,90 +3004,49 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: self._release_save_reservation(sid, owner=seq) self._settle_finished_save(sid) - tracker_sids = list(self._save_tracker.keys()) - if self._save_policy == "round_robin": - if tracker_sids and self._save_rr_last in self._save_tracker: - start = (tracker_sids.index(self._save_rr_last) + 1) % len( - tracker_sids - ) - tracker_sids = tracker_sids[start:] + tracker_sids[:start] - candidates = [ - candidate - for sid in tracker_sids - if (candidate := self._build_save_candidate(sid, now=now)) - is not None - ] - else: - candidates = [ - candidate - for sid in tracker_sids - if (candidate := self._build_save_candidate(sid, now=now)) - is not None - and self._meets_save_value_threshold(candidate) - ] - candidates.sort( - key=lambda candidate: self._candidate_sort_key(candidate, now) - ) + candidates = [ + candidate + for sid in self._save_tracker + if (candidate := self._build_save_candidate(sid, now=now)) is not None + and self._meets_save_value_threshold(candidate) + ] + candidates.sort( + key=lambda candidate: self._candidate_sort_key(candidate, now) + ) - if self._save_policy == "priority": - for candidate in candidates: - sid = candidate.sid - if sid in self._save_committed: - continue - if sid in self._reqs_need_recv or sid in loading_sids: - continue - if sid in self._save_inflight or sid in self._sidecar_save_inflight: - continue - if not self._try_reserve_save_candidate( - candidate, - allow_eviction=True, - now=now, - ): - self.drop_unadmitted_save(candidate.seq, reason="capacity") - - committed = [] - for sid, seq in list(self._save_committed.items()): - candidate = self._build_save_candidate(sid, now=now) - if candidate is None or candidate.seq is not seq: - self._save_committed.pop(sid, None) - self._release_save_reservation(sid, owner=seq) - continue - committed.append(candidate) - committed.sort( - key=lambda candidate: self._candidate_sort_key(candidate, now) - ) - for candidate in committed: - sid = candidate.sid - if sid in self._reqs_need_recv or sid in loading_sids: - continue - if sid in self._save_inflight or sid in self._sidecar_save_inflight: - continue - self._emit_save_candidate(meta, candidate) - else: - for candidate in candidates: - sid = candidate.sid - if not self._may_emit_save(): - break - if sid in self._save_committed: - continue - if sid in self._reqs_need_recv or sid in loading_sids: - continue # loading this step; defer its save - if sid in self._save_inflight or sid in self._sidecar_save_inflight: - # PAGE and SLOT describe one ordered checkpoint stream. - continue - if not self._try_reserve_save_candidate( - candidate, - allow_eviction=False, - now=now, - ): - if candidate.finished: - self.drop_unadmitted_save( - candidate.seq, - reason="capacity", - ) - break - self._emit_save_candidate(meta, candidate) - self._save_rr_last = sid + for candidate in candidates: + sid = candidate.sid + if sid in self._save_committed: + continue + if sid in self._reqs_need_recv or sid in loading_sids: + continue + if sid in self._save_inflight or sid in self._sidecar_save_inflight: + continue + if not self._try_reserve_save_candidate( + candidate, + allow_eviction=True, + now=now, + ): + self.drop_unadmitted_save(candidate.seq, reason="capacity") + + committed = [] + for sid, seq in list(self._save_committed.items()): + candidate = self._build_save_candidate(sid, now=now) + if candidate is None or candidate.seq is not seq: + self._save_committed.pop(sid, None) + self._release_save_reservation(sid, owner=seq) + continue + committed.append(candidate) + committed.sort( + key=lambda candidate: self._candidate_sort_key(candidate, now) + ) + for candidate in committed: + sid = candidate.sid + if sid in self._reqs_need_recv or sid in loading_sids: + continue + if sid in self._save_inflight or sid in self._sidecar_save_inflight: + continue + self._emit_save_candidate(meta, candidate) dispatched = set(meta.lookup_requests_in_step) self._lookup_in_step = [ sid for sid in self._lookup_in_step if sid not in dispatched @@ -3217,11 +3160,7 @@ def save_finished(self, req_id) -> None: self._save_inflight_since.pop(req_id, None) self._finish_save_statistics(req_id) self._retire_save_operation_owner(req_id) - if ( - self._save_policy == "priority" - and owner is not None - and self._finished_save_requests.get(sid) is owner - ): + if owner is not None and self._finished_save_requests.get(sid) is owner: self._settle_finished_save(sid) return @@ -3239,8 +3178,7 @@ def save_finished(self, req_id) -> None: self._save_inflight_since.pop(operation, None) self._retire_save_operation_owner(operation) self._finish_save_statistics(req_id) - if self._save_policy == "priority": - self._settle_finished_save(sid) + self._settle_finished_save(sid) def abandon_save(self, req_id) -> None: """Force-drop a save the scheduler reclaimed after it stalled. @@ -3359,11 +3297,7 @@ def sidecar_save_finished(self, req_id) -> None: if not failed: self._failed_sidecar_saves.pop(sid, None) self._retire_save_operation_owner(operation) - if ( - self._save_policy == "priority" - and owner is not None - and self._finished_save_requests.get(sid) is owner - ): + if owner is not None and self._finished_save_requests.get(sid) is owner: self._settle_finished_save(sid) def sidecar_save_failed(self, req_id) -> None: @@ -3385,11 +3319,7 @@ def sidecar_save_failed(self, req_id) -> None: identity = (inflight[1], inflight[2]) self._failed_sidecar_saves.setdefault(sid, set()).add(identity) self._retire_save_operation_owner(operation) - if ( - self._save_policy == "priority" - and owner is not None - and self._finished_save_requests.get(sid) is owner - ): + if owner is not None and self._finished_save_requests.get(sid) is owner: self._settle_finished_save(sid) def load_failed(self, req_id) -> bool: @@ -3518,29 +3448,22 @@ def request_finished(self, seq) -> None: if entry is not None and entry[0] is seq: self._finished_save_requests[sid] = seq self._refresh_finished_save_reservations(seq) - if self._save_policy == "priority": - if ( - sid in self._save_inflight - or sid in self._sidecar_save_inflight - or self._save_committed.get(sid) is seq - ): - # Already committed/dispatched. Its terminal callbacks will - # either reserve one residual tail or drop it. - pass - else: - candidate = self._build_save_candidate(sid) - if candidate is None: - self._settle_finished_save(sid) - elif not self._meets_save_value_threshold(candidate): - self.drop_unadmitted_save(seq, reason="low_value") - elif not self._commit_finished_save(candidate): - self.drop_unadmitted_save(seq, reason="capacity") - elif not self.should_defer_free(seq): - self._save_tracker.pop(sid, None) - self._finished_save_requests.pop(sid, None) - self._failed_sidecar_saves.pop(sid, None) - self._save_watermark_rollback.pop(sid, None) - self._forget_save_candidate(sid, seq) + if ( + sid in self._save_inflight + or sid in self._sidecar_save_inflight + or self._save_committed.get(sid) is seq + ): + # Already committed/dispatched. Its terminal callbacks will + # either reserve one residual tail or drop it. + pass + else: + candidate = self._build_save_candidate(sid) + if candidate is None: + self._settle_finished_save(sid) + elif not self._meets_save_value_threshold(candidate): + self.drop_unadmitted_save(seq, reason="low_value") + elif not self._commit_finished_save(candidate): + self.drop_unadmitted_save(seq, reason="capacity") if hasattr(seq, "_load_operation"): delattr(seq, "_load_operation") diff --git a/atom/kv_transfer/offload/save_admission.py b/atom/kv_transfer/offload/save_admission.py index fdffcafeee..d0e0c0d90f 100644 --- a/atom/kv_transfer/offload/save_admission.py +++ b/atom/kv_transfer/offload/save_admission.py @@ -60,7 +60,6 @@ def _optional_nonnegative_int(name: str) -> int | None: @dataclass(frozen=True) class SaveAdmissionConfig: - policy: str min_observed_count: int aging_weight: float release_weight: float @@ -92,31 +91,13 @@ class SaveBlockReservation: def load_save_admission_config() -> SaveAdmissionConfig: - """Read the shared save policy, preserving old defaults when disabled.""" + """Read the shared priority save-admission configuration.""" - policy = os.environ.get("OFFLOAD_SAVE_POLICY", "round_robin").strip().lower() - if policy not in {"round_robin", "priority"}: - raise ValueError( - f"OFFLOAD_SAVE_POLICY must be 'round_robin' or 'priority', got {policy!r}" - ) max_pinned_ratio = _nonnegative_float("OFFLOAD_SAVE_MAX_PINNED_RATIO", 0.20) if max_pinned_ratio > 0.30: raise ValueError("OFFLOAD_SAVE_MAX_PINNED_RATIO must be at most 0.30") max_pinned_blocks = _optional_nonnegative_int("OFFLOAD_SAVE_MAX_PINNED_BLOCKS") - if policy == "round_robin": - return SaveAdmissionConfig( - policy, - 2, - 0.01, - 1.0, - 8192, - 65536, - 600.0, - max_pinned_ratio, - max_pinned_blocks, - ) return SaveAdmissionConfig( - policy=policy, min_observed_count=_nonnegative_int("OFFLOAD_SAVE_MIN_OBSERVED_COUNT", 2), aging_weight=_nonnegative_float("OFFLOAD_SAVE_AGING_WEIGHT", 0.01), release_weight=_nonnegative_float("OFFLOAD_SAVE_RELEASE_WEIGHT", 1.0), diff --git a/docs/environment_variables.md b/docs/environment_variables.md index 605d305d2c..54285416ee 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -204,8 +204,7 @@ discoverable from the central env reference despite bypassing the registry. |----------|------|---------|-------------| | **LMCACHE_EC_PIN_TIMEOUT_SEC** | float | LMCache's own (300) | LMCache's source-pin timeout. ATOM reads it only to derive the engine's save-abandon window (`pin + 30s`), so the two stay ordered — a lost store report is reclaimed only after LMCache would already have force-unpinned its source. Non-positive disables ATOM's reclamation. ATOM sets no default of its own; when unset it assumes LMCache's. | | **OFFLOAD_MAX_PENDING_SAVES** | int | **2**, flat, for the engine-side/state-tier reader (`scheduler.py`); `max(2, 2 × OFFLOAD_COPY_WORKERS)` for the KV-leg reader (`_offload_common.py`) | Bound on total in-flight offload transfers (running + queued) held before a SLOT snapshot or executor submission. A KV save and a state store both pin bytes out of the same pool while they run, so the KV leg and the K3 state tier share this one number rather than each carrying its own. Two readers compute it, though: the KV leg's canonical `_offload_common.max_pending_saves` derives the shown default from `OFFLOAD_COPY_WORKERS` and **raises** on an unparseable value, while the scheduler's state-tier reader (`_offload_max_pending_saves`) has a simpler fallback — a flat default of **2** (no `OFFLOAD_COPY_WORKERS` scaling) that **warns and uses 2** on an unparseable value rather than raising. Set the env to an explicit integer to pin both. | -| **OFFLOAD_SAVE_POLICY** | enum | `round_robin` | Chunked LMCache save admission policy, including `lmcache_mp` PAGE-only/native-state and DSV4. `round_robin` preserves the compatibility behavior. `priority` admits reusable prefixes by rank-local demand and dirty-token cost; a finished request that cannot reserve capacity immediately is dropped instead of pinning KV indefinitely. | -| **OFFLOAD_SAVE_MIN_OBSERVED_COUNT** | int | `2` | Priority mode's hard reuse-demand floor. Aging and finished-request release benefit do not bypass it. | +| **OFFLOAD_SAVE_MIN_OBSERVED_COUNT** | int | `2` | Hard reuse-demand floor for priority-ranked save admission. Aging and finished-request release benefit do not bypass it. | | **OFFLOAD_SAVE_MAX_PINNED_RATIO** | float | `0.20` | Maximum fraction of this scheduler-local physical KV pool reserved or pinned by admitted saves. Valid range is `0.0`–`0.30`; the pool is local to one DP scheduler and is not multiplied by TP size. | | **OFFLOAD_SAVE_MAX_PINNED_BLOCKS** | int | unset | Optional nonnegative absolute PAGE-block clamp. The effective local budget is `min(floor(total_blocks × ratio), absolute)`. | | **OFFLOAD_SAVE_AGING_WEIGHT** | float | `0.01` | Nonnegative per-second aging term in the priority score. | diff --git a/tests/test_lmcache_offload_connector.py b/tests/test_lmcache_offload_connector.py index b876f80ee2..264d689a95 100644 --- a/tests/test_lmcache_offload_connector.py +++ b/tests/test_lmcache_offload_connector.py @@ -161,9 +161,10 @@ def _scheduler() -> LMCacheOffloadConnectorScheduler: sched._hit_save_floors = {} sched._save_tracker = {} sched._max_pending_saves = 4 - sched._save_rr_last = None - sched._save_policy = "round_robin" - sched._save_min_observed_count = 2 + # Most lifecycle tests below insert tracker entries directly instead of + # exercising demand observation. Keep their threshold open; focused + # priority-admission tests restore the production default of two. + sched._save_min_observed_count = 0 sched._save_aging_weight = 0.01 sched._save_release_weight = 1.0 sched._save_max_pinned_ratio = 0.20 @@ -238,7 +239,6 @@ def _stateful_scheduler(hit: int) -> LMCacheOffloadConnectorScheduler: def _priority_scheduler(*, capacity: int = 2) -> LMCacheOffloadConnectorScheduler: sched = _scheduler() - sched._save_policy = "priority" sched._max_pending_saves = capacity sched._save_min_observed_count = 2 return sched @@ -4508,7 +4508,7 @@ def test_chunked_prefill_save_uses_computed_frontier_and_serializes_inflight(): assert meta3.requests[0].is_last_prefill is True -def test_dsv4_save_admission_is_bounded_and_round_robin(): +def test_dsv4_save_admission_is_bounded_and_drops_unadmitted_work(): sched = _scheduler() sched._max_pending_saves = 2 for req_id in (100, 101, 102): @@ -4526,12 +4526,12 @@ def test_dsv4_save_admission_is_bounded_and_round_robin(): assert [request.req_id for request in first.requests] == [100, 101] assert sched._may_emit_save() is False - # The worker has one slot again. Resume after the last admitted request, - # rather than letting request 100 win the queue repeatedly. + # A candidate that could not reserve capacity was rejected rather than + # retained behind the admitted work. + assert "102" not in sched._save_tracker sched.save_finished(first.requests[0].save_operation) second = sched.build_connector_meta() - assert [request.req_id for request in second.requests] == [102] - assert sched._may_emit_save() is False + assert second.requests == [] def test_dsv4_priority_save_admits_hotter_prefix_first(): @@ -4633,14 +4633,11 @@ def test_dsv4_zero_save_budget_rejects_without_deferring_free(): assert stats["save_oversized"] == 1 -def test_dsv4_round_robin_finished_budget_rejection_stops_deferring_free(): +def test_dsv4_finished_budget_rejection_stops_deferring_free(): sched = _scheduler() _set_save_budget(sched, total_blocks=10, absolute=0) seq = _priority_seq(sched, req_id=238, token_ids=list(range(8))) sched.request_finished(seq) - assert sched.should_defer_free(seq) is True - - assert sched.build_connector_meta().requests == [] assert sched.should_defer_free(seq) is False assert "238" not in sched._save_tracker @@ -4756,7 +4753,6 @@ def test_dsv4_equal_priority_candidate_cannot_replace_committed_save(): def test_dsv4_inflight_page_slot_save_is_never_an_eviction_victim(): sched = _stateful_scheduler(hit=0) - sched._save_policy = "priority" sched._save_min_observed_count = 2 sched._save_aging_weight = 0 sched._max_pending_saves = 2 @@ -4811,7 +4807,6 @@ def test_dsv4_request_id_reuse_releases_old_committed_reservation(): def test_dsv4_page_slot_replacement_is_atomic(): sched = _stateful_scheduler(hit=0) - sched._save_policy = "priority" sched._save_min_observed_count = 2 sched._save_aging_weight = 0 sched._max_pending_saves = 2 @@ -4958,7 +4953,6 @@ def test_dsv4_priority_late_failure_does_not_touch_reused_request_id(): def test_dsv4_priority_drop_removes_page_and_sidecar_candidate_together(): sched = _stateful_scheduler(hit=0) - sched._save_policy = "priority" sched._save_min_observed_count = 2 seq = _stateful_seq( req_id=214, @@ -5640,7 +5634,6 @@ def _k3_scheduler() -> KimiK3OffloadScheduler: s._save_inflight = {} s._save_committed = {} s._save_tracker = {} - s._save_rr_last = None s._pending_state_loads = [] s._pending_state_stores = [] s._save_inflight_since = {} @@ -5667,59 +5660,6 @@ def _k3_seq(*, hbm: int, joint: int = 0, kv: int = 0, claimed: int = 0): ) -def test_bounded_saves_are_shared_round_robin_not_head_first(): - """With a bounded ``_may_emit_save`` (kimi_k3 caps outstanding saves), the - save scan must rotate over ``_save_tracker`` so a long multi-chunk request - at the insertion-ordered head cannot re-win the freed slot every step and - starve later requests -- whose blocks stay pinned by ``should_defer_free`` - until their save drains. Without the round-robin cursor the emission order - would be [100, 100, 100]; with it, each request gets a turn.""" - - class _Cap1(DenseOffloadScheduler): - def _may_emit_save(self): # one save outstanding at a time - return len(self._save_inflight) < 1 - - def _track_save_statistics(self, *a, **k): - pass - - s = _Cap1.__new__(_Cap1) - s.chunk_size = 256 - s.block_size = 256 - s._do_save = True - s._do_load = False - s._reqs_need_recv = {} - s._lookup_in_step = [] - s._save_tracker = {} - s._save_inflight = {} - s._save_nonce = 0 - s._save_rr_last = None - - seqs = {} - for i in range(3): - sid = str(100 + i) - seq = SimpleNamespace( - id=100 + i, - num_prompt_tokens=4096, - token_ids=list(range(4096)), - num_cached_tokens=0, - block_table=list(range(16)), - ) - s._save_tracker[sid] = [seq, 0] - seqs[sid] = seq - - emitted = [] - for _ in range(3): - for seq in seqs.values(): # every request computes one more chunk - seq.num_cached_tokens += 256 - meta = s.build_connector_meta() - saves = [r for r in meta.requests if r.save_spec is not None] - assert len(saves) == 1 # the cap admits exactly one per step - emitted.append(str(saves[0].req_id)) - s._save_inflight.clear() # that save completes, freeing the one slot - - assert emitted == ["100", "101", "102"] - - def test_layout_selection_picks_kimi_k3_for_a_kda_config(): """K3's paged KV is ordinary dense MLA, so the layout cannot be read off the KV shape -- the KDA state is what makes it its own family."""