From 62893d7530352395e8987f8e8123295ee6ef3f61 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Fri, 25 Sep 2026 15:46:31 +0000 Subject: [PATCH 01/28] feat(offload): support native DSv4 checkpoints with LMCache MP Save and restore DSv4 PAGE KV together with the matching recurrent STATE checkpoint through a standalone LMCache multiprocess server. The lmcache_mp connector selects its implementation from backend capabilities: backends that publish a PagedStateCheckpointSpec and execute_paged_state_copies use the native PAGE/STATE path, others keep the PAGE-only transport. Saves lease immutable READY checkpoint PAGE units instead of snapshotting the live SLOT; restores reserve fresh units and adopt them as a local READY checkpoint. Also: - allow single-host DP and DP-attention, scoping MP sessions per replica; - release source PAGE blocks per chunk and reacquire a finished request's still-resident prefix at save admission; - allow partial release of recurrent-state requests only when the connector guarantees an independent state lease; - read offload env knobs through atom.utils.envs; - refuse PAGE-copied state backends that lack the native contract. Requires LMCache dev@05fc77a (LMCache/LMCache#5132), pinned on main by #2395. Start the server with --null-block-id -1 --separate-object-groups. Co-authored-by: Claude Opus 5.5 (1M context) --- .../disaggregation/multi/multi_connector.py | 23 + atom/kv_transfer/disaggregation/types.py | 12 + atom/kv_transfer/offload/README.md | 186 ++++- .../offload/_block_gpu_connector.py | 13 +- atom/kv_transfer/offload/_offload_common.py | 72 +- .../offload/atom_lmcache_staging.py | 35 - atom/kv_transfer/offload/chunked_scheduler.py | 159 +++-- atom/kv_transfer/offload/connector.py | 9 + .../offload/hybrid/dsv4/connector.py | 56 +- .../kv_transfer/offload/hybrid/dsv4/policy.py | 13 +- .../offload/hybrid/kimi_k3/connector.py | 4 +- atom/kv_transfer/offload/metadata.py | 11 + atom/kv_transfer/offload/mp/backend.py | 253 +++++-- atom/kv_transfer/offload/mp/connector.py | 187 ++++- .../offload/mp/native_state_layout.py | 222 ++++++ .../offload/mp/native_state_scheduler.py | 516 ++++++++++++++ .../offload/mp/native_state_worker.py | 416 +++++++++++ atom/model_engine/block_manager.py | 76 +- atom/model_engine/block_pool.py | 29 + atom/model_engine/model_runner.py | 18 +- atom/model_engine/page_unit_checkpoint.py | 190 ++++- atom/model_engine/scheduler.py | 62 +- atom/model_ops/attentions/aiter_attention.py | 17 +- atom/model_ops/attentions/backends.py | 55 +- atom/model_ops/attentions/deepseek_v4_attn.py | 75 +- atom/model_ops/attentions/gdn_attn.py | 37 +- atom/plugin/vllm/kv_transfer/seq_view.py | 16 +- atom/spec_decode/draft_kv.py | 39 +- atom/utils/envs.py | 130 +++- tests/plugin/test_vllm_offload_seq_view.py | 4 +- tests/test_block_manager.py | 30 + tests/test_draft_kv_builder.py | 34 + tests/test_lmcache_mp.py | 244 ++++++- tests/test_lmcache_mp_builder_views.py | 63 ++ tests/test_lmcache_mp_native_layout.py | 339 +++++++++ tests/test_lmcache_mp_native_scheduler.py | 649 ++++++++++++++++++ tests/test_lmcache_mp_native_worker.py | 427 ++++++++++++ tests/test_lmcache_mp_shell.py | 160 ++++- tests/test_lmcache_offload_connector.py | 6 +- tests/test_lmcache_offload_v4_page_slot.py | 12 +- tests/test_multi_connector.py | 24 + tests/test_offload_early_block_release.py | 109 ++- tests/test_page_unit_checkpoint_leases.py | 246 +++++++ tests/test_scheduler.py | 26 + 44 files changed, 4857 insertions(+), 447 deletions(-) create mode 100644 atom/kv_transfer/offload/mp/native_state_layout.py create mode 100644 atom/kv_transfer/offload/mp/native_state_scheduler.py create mode 100644 atom/kv_transfer/offload/mp/native_state_worker.py create mode 100644 tests/test_lmcache_mp_native_layout.py create mode 100644 tests/test_lmcache_mp_native_scheduler.py create mode 100644 tests/test_lmcache_mp_native_worker.py create mode 100644 tests/test_page_unit_checkpoint_leases.py diff --git a/atom/kv_transfer/disaggregation/multi/multi_connector.py b/atom/kv_transfer/disaggregation/multi/multi_connector.py index dc892cc502..5e565ca8bb 100644 --- a/atom/kv_transfer/disaggregation/multi/multi_connector.py +++ b/atom/kv_transfer/disaggregation/multi/multi_connector.py @@ -348,6 +348,13 @@ def close(self) -> None: class MultiConnectorScheduler(KVConnectorSchedulerBase): """Scheduler-side composite connector.""" + def bind_block_manager(self, block_manager: Any) -> None: + """Bind native checkpoint owners before any request can be admitted.""" + for connector in self._connectors: + bind = getattr(connector, "bind_block_manager", None) + if callable(bind): + bind(block_manager) + def __init__(self, config: Any) -> None: self._connectors = _build_subconnectors(config, role="scheduler") self.is_producer = any( @@ -646,6 +653,22 @@ def protected_block_ids(self, seq: Any): protected.update(blocks) return frozenset(protected) + def can_partially_deallocate_state(self, seq: Any) -> bool: + """Forward an explicit state-lifetime guarantee from its owner. + + Missing methods are not treated as approval. The separate + ``protected_block_ids`` check still requires every deferring sub to + narrow its PAGE reads before the scheduler takes this path. + """ + for connector in self._connectors: + should_defer = getattr(connector, "should_defer_free", None) + if not callable(should_defer) or not should_defer(seq): + continue + callback = getattr(connector, "can_partially_deallocate_state", None) + if callable(callback) and callback(seq) is True: + return True + return False + def activate_block_leases(self, seq: Any, block_ids) -> None: for connector in self._connectors: callback = getattr(connector, "activate_block_leases", None) diff --git a/atom/kv_transfer/disaggregation/types.py b/atom/kv_transfer/disaggregation/types.py index d1a1a5ec68..8b4bb70b4b 100644 --- a/atom/kv_transfer/disaggregation/types.py +++ b/atom/kv_transfer/disaggregation/types.py @@ -257,6 +257,18 @@ class KVTransferTensors: gather_sharded_index: Callable[..., tuple[int, int]] | None = None # Appended after the original staging fields for positional compatibility. prepare_sharded_index: Callable[..., Any] | None = None + # Native PAGE-backed state image contract. The spec names the exact codec + # and image geometry; the callback restores that image into an Active SLOT + # after a remote retrieve has finished writing its PAGE units. These are + # intentionally untyped model-engine objects to keep this module free of + # model-engine imports. The PAGE tensor views retain the source allocations; + # neither field asks the connector to snapshot a complete live SLOT. + paged_state_checkpoint_spec: object | None = None + execute_paged_state_copies: Callable[..., None] | None = None + # Same declaration for PAGE-backed native state checkpoint images. Kept + # separate so future layouts can describe PAGE and recurrent-state + # replication independently. Appended for positional compatibility. + native_state_tp_replication_factor: int = 1 # Scheduler blocks the PAGE regions are addressed in. `init=False` because # a backend cannot answer it: `req.block_ids` is the scheduler's id space, # and a backend counts in its own page -- a different unit even where it is diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index 7f608457bc..f439614697 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -1,8 +1,14 @@ # LMCache CPU/NVMe KV Cache Offload (ATOM standalone) This module adds a **CPU DRAM (L2) and optional NVMe (L3) cache tier** on top of -ATOM's native HBM prefix and state caches. Dense MHA/MLA continues to offload -opaque KV blocks. Stateful DeepSeek V4 (DSV4) uses a stricter pair: +ATOM's native HBM prefix and state caches. It exposes two standalone connector +paths: + +- `lmcache_offload` embeds an LMCache engine in each ATOM worker; +- `lmcache_mp` connects ATOM workers to an external LMCache multiprocess server. + +The in-process path keeps dense MHA/MLA data as opaque KV blocks. Stateful +DeepSeek V4 (DSV4) uses a stricter pair: - **PAGE** — page-major compressed KV from `KVTransferTensors.block_regions`, persisted incrementally through standard `LMCacheEngine.store/retrieve`. @@ -12,8 +18,9 @@ opaque KV blocks. Stateful DeepSeek V4 (DSV4) uses a stricter pair: A DSV4 boundary is reusable only when both PAGE and SLOT restore successfully. Missing, incompatible, or corrupt sidecar data fails closed to recomputation. -The public configuration remains `kv_connector: "lmcache_offload"`. The thin -top-level shell resolves one of four layouts: `m3` for MiniMax-M3 PAGE regions +For the in-process path, the public configuration remains +`kv_connector: "lmcache_offload"`. The thin top-level shell resolves one of four +layouts: `m3` for MiniMax-M3 PAGE regions (including its NSA index cache), `kimi_k3` when the text config has `model_type == "kimi_linear"` (dense paged MLA KV plus a KDA per-request state tier), `hybrid` when `hf_config.compress_ratios` is present (DSV4 PAGE+SLOT), @@ -41,7 +48,7 @@ the byte-level deep dives ([Key Modules](#key-modules-in-depth), [Relationship to LMCache](#relationship-to-lmcache-reuse-vs-override)) come later. Unfamiliar terms are in the [Glossary](#glossary). -## Design at a Glance +## In-process `lmcache_offload` Design at a Glance Four rules carry the module: @@ -96,13 +103,176 @@ Four rules carry the module: | `hybrid/kimi_k3/state_object.py` | One state checkpoint as a single opaque object keyed by ATOM's own hash, bypassing LMCache's `ChunkedTokenDatabase` (state bytes are not token-sliceable). | | `hybrid/kimi_k3/state_tier.py` | Worker-side store/load driver for the state tier on its own executor; reports store/finished/failed hash sets for the engine-side `StateOffloadIndex` to apply. | | `atom_lmcache_staging.py` | Per-thread CUDA streams, staging buffer, ready/free events, env helpers. | +| `mp/connector.py` | Capability-selected public `lmcache_mp` worker/scheduler shells. | +| `mp/backend.py` | Generic PAGE-only LMCache multiprocess transport and adapter integration. | +| `mp/native_state_{layout,scheduler,worker}.py` | PAGE-backed native-state registration, scheduler leases, and worker transfer/restore lifecycle. | The engine-side counterpart of the state tier lives outside this directory: `atom/model_engine/state_offload.py` holds `StateOffloadIndex`, which applies the store/finished/failed hash sets that `hybrid/kimi_k3/state_tier.py` reports back, so a later prefix hit knows which checkpoints the CPU tier can actually serve. -## Architecture +## LMCache Multiprocess (`lmcache_mp`) + +The `lmcache_mp` connector selects one of two paths from capabilities published +by the attention backend: + +- ordinary backends publish PAGE views and use PAGE-only transfer; +- stateful backends additionally publish a `PagedStateCheckpointSpec` and an + `execute_paged_state_copies` callback, enabling a combined PAGE/native-STATE + transfer. + +Selection does not inspect model names or layout-ID prefixes. A new backend can +reuse the native path by implementing these shared contracts. The native path +leases existing READY checkpoint PAGE units; it does not snapshot the request's +live Active SLOT again. Every running request therefore keeps its normal fixed +SLOT. + +The PAGE-only path covers both sparse MLA and MHA/GQA layouts. GLM-5.2's MLA KV +and index cache are byte-identical across TP, so automatic rank collapse stores +one copy while every rank retrieves it. MiniMax-M3 publishes its GQA KV, scale, +and NSA index-cache planes as zero-copy PAGE views and deliberately stores one +shard per TP rank. Both use chunk-completion events to release source PAGE +blocks before the remote store becomes terminal. + +### Running the MP server + +The matching LMCache build must support the global `--null-block-id` setting +and sparse null handling. Run the MP server on the same host, with GPU IPC +access to the ATOM worker allocations: + +```bash +lmcache server --host 127.0.0.1 --port 5555 \ + --chunk-size 256 \ + --null-block-id -1 --separate-object-groups \ + --supported-transfer-mode lmcache_driven --l1-size-gb 64 +``` + +For example, add the following settings to a DSv4 launch that publishes the +native-state contract: + +```bash +export LMCACHE_CHUNK_SIZE=256 +export OFFLOAD_MAX_PENDING_SAVES=2 +export OFFLOAD_MIN_LOAD_TOKENS=8192 +export OFFLOAD_MIN_SAVE_TOKENS=8192 + +python -m atom.entrypoints.openai_server \ + --model deepseek-ai/DeepSeek-V4-Pro --kv_cache_dtype fp8 -tp 8 \ + --enable_prefix_caching --state-checkpoint-interval-tokens 8192 \ + --kv-transfer-config '{ + "kv_connector": "lmcache_mp", + "kv_role": "offload", + "kv_connector_extra_config": { + "lmcache.mp.host": "tcp://127.0.0.1", + "lmcache.mp.port": 5555, + "lmcache.mp.tp_rank_collapse": true + } + }' +``` + +ATOM's configured chunk size must equal the MP server chunk size, and both must +align to ATOM's PAGE/hash block size. A native prefix is loadable only where +PAGE KV and a complete STATE checkpoint exist at the same boundary on every TP +rank. + +`lmcache.mp.max_pinned_state_bytes` optionally limits native checkpoint sources +and temporary restore images together. By default it is +`OFFLOAD_MAX_PENDING_SAVES * units_per_checkpoint * page_unit_bytes` for each TP +worker. Candidates consume no PAGE/image pin until admission. If a request has +already finished, admission resolves its token/hash chain through the live +prefix index and stores only the still-resident contiguous prefix. +`OFFLOAD_MIN_SAVE_TOKENS` suppresses a late save whose remaining prefix is too +small. Unless configured otherwise, the shared save limit is +`max(2, 2 * OFFLOAD_COPY_WORKERS)`. + +The model namespace includes PAGE/model geometry, TP and speculation settings, +native layout/image sizes, the Hugging Face commit when available, and the +optional `lmcache.mp.model_revision`. Set that revision when weights are +replaced in an existing local model directory; a directory name alone cannot +identify changed contents. + +### MP lifetime and representation + +Engine group `0` contains ordinary PAGE views, where block ID `0` is valid +data. Native image ordinal `j` uses engine group `1+j`, aliases the same PAGE +allocation, and declares a one-chunk recurrent window. Earlier chunks use `-1` +for every STATE group and all STATE ordinals become present together at the +checkpoint endpoint. The server-wide `--null-block-id -1` distinguishes those +absent chunks, while `--separate-object-groups` separates PAGE and STATE +objects. The final STATE alias is trimmed at `image_bytes` without changing the +underlying PAGE stride. + +The scheduler dispatches one combined PAGE/STATE generation at a time per +request, using round-robin admission plus count and byte bounds. It leases the +exact READY state image only after admission. Source-safe events release PAGE +leases chunk by chunk and release the state image when its endpoint is safe; +terminal completion settles the logical operation. Failed saves roll back the +watermark for at most three attempts at one boundary. An uncertain remote DMA +is never reclaimed by elapsed time alone. + +Restore transfers PAGE KV only for `[hbm, lmcache)` and restores the endpoint's +full native image into the request's already allocated SLOT. An unaligned HBM +hit is first prefetched to the next chunk boundary; the aligned remainder is +loaded only when it meets `OFFLOAD_MIN_LOAD_TOKENS`. Native restoration uses a +dedicated stream and descriptor slot, and completion is polled through a CUDA +event without synchronizing the compute stream. On success, the temporary +STATE PAGE units are atomically adopted as an unpinned READY checkpoint for +normal local reuse and LRU eviction. + +```text +SAVE + live request PAGEs ---- chunk-safe events ----> early PAGE release + READY state PAGEs ---- endpoint-safe event ---> state lease release + remote store terminal ------------------------> operation settlement + +RESTORE + remote PAGE/STATE -> temporary PAGE units -> Active SLOT + -> adopt units as READY local checkpoint +``` + +### MP scope and constraints + +- One MP server supports TP, single-host DP, and single-host DP-attention with + EP. DP replicas use private request-session IDs but share the same + content-addressed model namespace. +- Multi-node DP, PP, PCP, DCP, and engine-driven transfers are rejected. + Multi-node DP needs one GPU-local LMCache server per host plus routing. +- PAGE tensors are registered as zero-copy `uint8` views, preserving FP8/BF16 + bit patterns and avoiding numerical dtype conversion in ROCm pointer paths. +- DSv4 declares PAGE and native STATE fully TP-replicated. In `auto` mode, one + rank stores and every rank retrieves (`num_kv_readers=TP`). GLM-5.2 follows + the same replicated sparse-MLA rule. MiniMax-M3 remains sharded and stores on + every rank. +- External native restore supports zero-HBM and incremental local-prefix cases; + native lookup truncates the query to + `floor((prompt_tokens - 1) / chunk_size) * chunk_size`, and PAGE and STATE + must resolve to the same real endpoint. +- Running requests never give their Active SLOT to the MP connector. Pending + saves are unpinned; admitted saves retain only source-unsafe PAGE chunks and + the exact READY state image. +- LMCache may still allocate internal GPU transfer buffers. Direct aliases + remove an additional ATOM SLOT image, not every transport copy. + +### MP validation + +CPU tests cover READY leases, generation replay, eviction/reset, byte budgets, +fair admission, cancellation/failure, full-prompt boundaries, alias byte order, +and strided-tail registration. LMCache tests cover null markers, serialization, +sparse STATE lookup, and capability negotiation. + +The independent-process GPU transport contract can be tested from the matching +LMCache checkout: + +```bash +python -m pytest -xvs tests/v1/multiprocess/test_native_state_alias_gpu.py +``` + +The DSv4-Pro TP8 integration test additionally exercises a cold save followed +by full remote restore, READY reuse, incremental restore, exact next-token +comparison, per-rank retrieve counts, one-writer storage, and promoted ranges. + +## In-process `lmcache_offload` Architecture The connector is split across two processes, mirroring ATOM's P/D split: @@ -848,8 +1018,8 @@ staging. Dense and DSV4 instantiate this shared adapter directly. Small but load-bearing. `_ThreadTransferState` lazily creates, per thread, the two CUDA streams (`pack_stream`, `copy_stream`) and a `_StagingBuffer` holding the device tensor plus `ready`/`free` CUDA events that gate the pipeline hand-off. -Also the `_env_flag/_env_int/_env_optional_int` helpers that parse the `OFFLOAD_*` -knobs. This per-thread isolation is exactly why load and save never contend. +The `OFFLOAD_*` knobs are parsed in `atom/utils/envs.py` like every other ATOM +env var. This per-thread isolation is exactly why load and save never contend. ### `config.py` & `metadata.py` — wiring and descriptors diff --git a/atom/kv_transfer/offload/_block_gpu_connector.py b/atom/kv_transfer/offload/_block_gpu_connector.py index efeb809874..ef741e0748 100644 --- a/atom/kv_transfer/offload/_block_gpu_connector.py +++ b/atom/kv_transfer/offload/_block_gpu_connector.py @@ -25,14 +25,13 @@ from atom.kv_transfer.disaggregation.types import SaveOperationId, SaveSourceGroupId from atom.kv_transfer.offload.atom_lmcache_staging import ( - _env_flag, - _env_optional_int, _PipelineStage, _StagingBuffer, _ThreadTransferState, memory_object_as_uint8, run_staged_pipeline, ) +from atom.utils import envs logger = logging.getLogger("atom") @@ -46,7 +45,7 @@ # handoff a no-op of in-order execution and gives up an overlap that does not # exist. Correctness is unaffected -- one stream is strictly more ordered than # two, and the terminal synchronize is unchanged. -_SINGLE_STREAM = _env_flag("OFFLOAD_SINGLE_STREAM") +_SINGLE_STREAM = envs.OFFLOAD_SINGLE_STREAM # The default staging buffer is denominated in bytes, not in LMCache chunks -- # see `_default_staging_buffer_chunks` for why. 48 MiB: on GLM-5.2, whose chunk @@ -260,10 +259,10 @@ def __init__( self._quarantined_staging_tensors: list[torch.Tensor] = [] self._quarantined_block_id_owners: list[Any] = [] self._quarantined_staging_lock = threading.Lock() - requested_buffer_chunks = _env_optional_int("OFFLOAD_GPU_STAGING_CHUNKS") + requested_buffer_chunks = envs.OFFLOAD_GPU_STAGING_CHUNKS if requested_buffer_chunks is None: requested_buffer_chunks = self._default_staging_buffer_chunks() - max_staging_bytes = _env_optional_int("OFFLOAD_GPU_STAGING_MAX_BYTES") + max_staging_bytes = envs.OFFLOAD_GPU_STAGING_MAX_BYTES if max_staging_bytes is not None: if max_staging_bytes < self._gpu_staging_chunk_bytes: raise ValueError( @@ -280,8 +279,8 @@ def __init__( self._gpu_staging_buffer_bytes = ( self._staging_buffer_chunks * self._gpu_staging_chunk_bytes ) - self._release_gpu_staging_after_transfer = _env_flag( - "OFFLOAD_RELEASE_GPU_STAGING_AFTER_TRANSFER" + self._release_gpu_staging_after_transfer = ( + envs.OFFLOAD_RELEASE_GPU_STAGING_AFTER_TRANSFER ) def _default_staging_buffer_chunks(self) -> int: diff --git a/atom/kv_transfer/offload/_offload_common.py b/atom/kv_transfer/offload/_offload_common.py index 8402f73bff..a2d0e7ec5f 100644 --- a/atom/kv_transfer/offload/_offload_common.py +++ b/atom/kv_transfer/offload/_offload_common.py @@ -29,6 +29,7 @@ SaveCompletionId, ) from atom.kv_transfer.offload import config as offcfg +from atom.utils import envs logger = logging.getLogger("atom") _VALID_KV_ROLES = {"offload", "kv_both", "kv_producer", "kv_consumer"} @@ -154,9 +155,7 @@ def _init_worker_common( # behind fire-and-forget saves. OFFLOAD_COPY_WORKERS tunes the save pool, # OFFLOAD_LOAD_WORKERS the load pool. n_save = ( - int(os.environ.get("OFFLOAD_COPY_WORKERS", "1")) - if save_workers is None - else int(save_workers) + envs.OFFLOAD_COPY_WORKERS if save_workers is None else int(save_workers) ) if n_save <= 0: raise ValueError("offload save worker count must be positive") @@ -169,9 +168,7 @@ def _init_worker_common( # threads are independent; each costs one more # `gpu_staging_buffer_bytes` allocation per rank. n_load = ( - int(os.environ.get("OFFLOAD_LOAD_WORKERS", "1")) - if load_workers is None - else int(load_workers) + envs.OFFLOAD_LOAD_WORKERS if load_workers is None else int(load_workers) ) if n_load <= 0: raise ValueError("offload load worker count must be positive") @@ -326,14 +323,7 @@ def _lookup_unpin(self, req_id) -> None: @staticmethod def _profile_enabled() -> bool: - # Same env-flag semantics as `atom_lmcache_staging._env_flag` (kept - # inline rather than imported: that module pulls in torch and this one - # is torch-free). Strip, and read the empty string as OFF -- `VAR=` is - # how a shell clears a flag inline, and a bare membership test would - # read "" as ON (not in the false set), the opposite of what the - # operator wrote; `VAR="off "` had the same trap. - raw = os.environ.get("OFFLOAD_PROFILE", "0").strip().lower() - return bool(raw) and raw not in {"0", "false", "no", "off"} + return envs.OFFLOAD_PROFILE def _last_gpu_connector_transfer_stats(self) -> dict[str, int | float]: gpu_connector = getattr(getattr(self, "_engine", None), "gpu_connector", None) @@ -667,16 +657,16 @@ 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 + a `kv_connector_extra_config` `"max_pending_saves"` override, else the + `OFFLOAD_MAX_PENDING_SAVES` setting, else 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) @@ -736,27 +726,8 @@ def _init_tier_hit_memo(self) -> None: # sid -> (weak ref to the sequence asked about, hit or None, budget). self._tier_hit_memo: dict[str, tuple[object, int | None, int]] = {} - self._tier_memo_steps = self._positive_env("OFFLOAD_LOOKUP_MEMO_STEPS", 32) - self._tier_retry_steps = self._positive_env("OFFLOAD_LOOKUP_RETRY_STEPS", 32) - - @staticmethod - def _positive_env(name: str, default: int) -> int: - raw = os.environ.get(name) - if raw is None: - return default - try: - value = int(raw) - except ValueError: - value = -1 - if value < 0: - logger.warning( - "LMCache offload scheduler: invalid %s=%r; using %d", - name, - raw, - default, - ) - return default - return value + self._tier_memo_steps = envs.OFFLOAD_LOOKUP_MEMO_STEPS + self._tier_retry_steps = envs.OFFLOAD_LOOKUP_RETRY_STEPS @staticmethod def _weak_seq(seq): @@ -1066,7 +1037,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]) @@ -1083,14 +1054,9 @@ def max_pending_saves(kvc, save_workers: int) -> int: 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 + capacity = envs.OFFLOAD_MAX_PENDING_SAVES + if capacity is None: + capacity = max(2, 2 * save_workers) else: if isinstance(configured, bool) or not isinstance(configured, int): raise ValueError("max pending saves must be a positive integer") diff --git a/atom/kv_transfer/offload/atom_lmcache_staging.py b/atom/kv_transfer/offload/atom_lmcache_staging.py index 21e255ff90..fe27d593c5 100644 --- a/atom/kv_transfer/offload/atom_lmcache_staging.py +++ b/atom/kv_transfer/offload/atom_lmcache_staging.py @@ -5,7 +5,6 @@ from __future__ import annotations -import os from collections.abc import Callable, Sequence from dataclasses import dataclass from typing import Any @@ -61,40 +60,6 @@ def __init__(self, use_cuda: bool) -> None: self.free_event = torch.cuda.Event(blocking=False) -def _env_flag(name: str, default: str = "0") -> bool: - # Stripped, and empty reads as off: `VAR=` is how a shell script clears a - # flag inline, and a bare membership test reads the empty string as ON -- - # the opposite of what the operator wrote. `VAR="off "` did the same. - raw = os.environ.get(name, default).strip().lower() - return bool(raw) and raw not in ("0", "false", "no", "off") - - -def _env_int(name: str, default: int, *, min_value: int = 1) -> int: - raw = os.environ.get(name) - if raw is None: - return default - try: - value = int(raw) - except ValueError as exc: - raise ValueError(f"{name} must be an integer, got {raw!r}") from exc - if value < min_value: - raise ValueError(f"{name} must be >= {min_value}, got {value}") - return value - - -def _env_optional_int(name: str, *, min_value: int = 1) -> int | None: - raw = os.environ.get(name) - if raw is None or raw == "": - return None - try: - value = int(raw) - except ValueError as exc: - raise ValueError(f"{name} must be an integer, got {raw!r}") from exc - if value < min_value: - raise ValueError(f"{name} must be >= {min_value}, got {value}") - return value - - class _ThreadTransferState: def __init__( self, diff --git a/atom/kv_transfer/offload/chunked_scheduler.py b/atom/kv_transfer/offload/chunked_scheduler.py index e65865d236..4c3b10cf9d 100644 --- a/atom/kv_transfer/offload/chunked_scheduler.py +++ b/atom/kv_transfer/offload/chunked_scheduler.py @@ -6,7 +6,6 @@ from __future__ import annotations import logging -import os import time from atom.kv_transfer.disaggregation.base import KVConnectorSchedulerBase @@ -28,6 +27,7 @@ LoadSpec, SaveSpec, ) +from atom.utils import envs logger = logging.getLogger("atom") @@ -157,21 +157,18 @@ def __init__( self._lookup_in_step: list[str] = [] self._lookup_results: dict[str, tuple[object, int]] = {} self._handoff_loads: set[str] = set() + self._block_manager = None # Unaligned handoff is always on: when the HBM prefix-cache hit is not # chunk-aligned, recompute the misaligned head up to the next chunk # boundary, then load the aligned remainder from CPU. (Previously gated # by the OFFLOAD_UNALIGNED_HANDOFF env var; now unconditional.) - try: - self._min_load_tokens = max( - 0, int(os.environ.get("OFFLOAD_MIN_LOAD_TOKENS", "8192")) - ) - except ValueError: - logger.warning( - "LMCache offload scheduler: invalid OFFLOAD_MIN_LOAD_TOKENS=%r; " - "using 8192", - os.environ.get("OFFLOAD_MIN_LOAD_TOKENS"), - ) - self._min_load_tokens = 8192 + self._min_load_tokens = envs.OFFLOAD_MIN_LOAD_TOKENS + self._min_save_tokens = envs.OFFLOAD_MIN_SAVE_TOKENS + + 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 # -- match: how many extra tokens can come from CPU/NVMe ------------- def _begin_load_lifecycle(self, seq) -> None: @@ -184,6 +181,10 @@ def _begin_load_lifecycle(self, seq) -> None: self._forget_tier_hit(sid) self._load_lifecycles[sid] = seq + def _lookup_token_ids(self, seq) -> list[int]: + """The prompt extent this layout can safely ask the tier to resume.""" + return list(seq.token_ids[: seq.num_prompt_tokens]) + def install_hit_cap_hook(self, hook) -> None: """Let a hybrid connector shorten every hit this scheduler reports. @@ -210,6 +211,8 @@ def get_num_new_matched_tokens(self, seq) -> tuple[int, bool]: sid = str(seq.id) if self._repeat_load_suppressed(seq, sid): return 0, False + if not self._lookup_token_ids(seq): + return 0, False pending = self._lookup_results.get(sid) if pending is not None and pending[0] is not seq: # An older lifecycle still owns this worker-side pin. Its cleanup @@ -229,7 +232,7 @@ def _fresh_tier_lookup(self, seq, sid: str) -> int | None: """Ask the tier, take its pin, and remember the hit. None if it did not answer.""" num_prompt = seq.num_prompt_tokens - token_ids = list(seq.token_ids[:num_prompt]) + token_ids = self._lookup_token_ids(seq) if sid not in self._lookup_in_step: self._lookup_in_step.append(sid) self._lookup_results[sid] = (seq, 0) @@ -402,6 +405,57 @@ def _may_emit_save(self) -> bool: """Return whether another save may be emitted this scheduler step.""" return True + 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(claimed[:keep_count]) + drop = claimed[keep_count:] + 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 build_connector_meta(self) -> LMCacheOffloadMetadata: meta = LMCacheOffloadMetadata() @@ -458,8 +512,7 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: ) loading_sids.add(sid) self._load_save_floors[sid] = self._chunk_floor(hbm) - load_operation = LoadOperationId(seq.id, self._load_nonce) - self._load_nonce += 1 + load_operation = self._new_load_operation(seq) seq._load_operation = load_operation self._active_load_operations[sid] = (seq, load_operation) self._track_load_statistics(load_operation, lmc - hbm) @@ -483,7 +536,6 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: # 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) @@ -513,7 +565,7 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: int(seq.num_prompt_tokens), ) is_last_prefill = computed >= int(seq.num_prompt_tokens) - aligned = (computed // chunk) * chunk + aligned = self._save_frontier(seq) if aligned <= saved: continue logger.debug( @@ -526,20 +578,29 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: ) 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, + late_acquired = frozenset() + if getattr(seq, "_offload_finished", False) and not seq.block_table: + late_source = self._late_save_source(seq, saved, aligned) + if late_source is None: + self._save_tracker.pop(sid, None) + continue + aligned, block_ids, late_acquired = late_source + else: + block_ids = list(seq.block_table) + try: + request = self._build_save_request( + seq, 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) + continue + self._track_save_statistics(save_operation, aligned - saved) + meta.add_request(request) entry[1] = aligned self._save_previous_frontier[save_operation] = saved self._save_previous_owner[save_operation] = seq @@ -559,6 +620,8 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: } self._save_operation_safe[save_operation] = set() self._save_operation_owner[save_operation] = seq + if late_acquired: + self.activate_block_leases(seq, late_acquired) dispatched = set(meta.lookup_requests_in_step) for sid in dispatched: self._lookup_results.pop(sid, None) @@ -598,18 +661,16 @@ def protected_block_ids(self, seq) -> frozenset | None: touches HBM blocks this connector does not track per-range) -- the scheduler falls back to deferring the whole request. This includes a final save that has not been emitted yet: request teardown freezes its - full block table and computed frontier so the next metadata build can - still dispatch that save after unrelated blocks have been released. + token identity and computed frontier, but protects no physical PAGE + until admission reacquires the still-canonical prefix. """ if not self._early_release or self._has_active_load(seq): return None - sid = str(seq.id) - table = list(getattr(seq, "_offload_finished_block_ids", seq.block_table)) - if not hasattr(seq, "_offload_finished_block_ids"): - seq._offload_finished_block_ids = table - seq._offload_finished_cached_tokens = min( - int(getattr(seq, "num_cached_tokens", 0)), int(seq.num_prompt_tokens) - ) + seq._offload_finished = True + if not hasattr(seq, "_offload_finished_cached_tokens"): + seq._offload_finished_cached_tokens = min( + int(getattr(seq, "num_cached_tokens", 0)), int(seq.num_prompt_tokens) + ) protected: set[int] = set() for operation, blocks in self._save_operation_blocks.items(): @@ -620,14 +681,6 @@ def protected_block_ids(self, seq) -> frozenset | None: block_id for block_id in blocks.values() if block_id not in safe ) - entry = self._save_tracker.get(sid) - if entry is not None and entry[0] is seq: - saved = int(entry[1]) - aligned = self._save_frontier(seq) - if aligned > saved: - start_block = saved // self.virtual_block_size - end_block = -(-aligned // self.virtual_block_size) - protected.update(table[start_block:end_block]) return frozenset(protected) def activate_block_leases(self, seq, block_ids: frozenset[int]) -> None: @@ -713,7 +766,7 @@ def has_pending_work(self) -> bool: a load or cleanup dispatchable in ``build_connector_meta``. """ pending_finished_save = getattr(self, "_early_release", False) and any( - hasattr(entry[0], "_offload_finished_block_ids") + getattr(entry[0], "_offload_finished", False) and self._has_pending_save(entry[0]) for entry in self._save_tracker.values() ) @@ -923,7 +976,7 @@ def _retry_or_retire_failed_save(self, operation: SaveOperationId) -> None: return if owner is not None and entry[0] is not owner: return # request ID was reused - if hasattr(entry[0], "_offload_finished_block_ids"): + if getattr(entry[0], "_offload_finished", False): # Teardown may already have returned the source blocks to the pool. # A missing cache entry is safer than a retry against recycled KV. self._save_tracker.pop(sid, None) @@ -932,8 +985,8 @@ def _retry_or_retire_failed_save(self, operation: SaveOperationId) -> None: self._save_lease_at.pop(lease_key, None) self._save_lease_owner.pop(lease_key, None) if remaining: - # Teardown also leased any not-yet-emitted suffix. No further - # save will read it once this request is retired. + # Teardown leased only emitted saves' sources; the unemitted + # suffix was never leased and is dropped with the request. self._pending_source_safe_releases.append(frozenset(remaining)) self.total_source_safe_released_blocks += len(remaining) return @@ -969,9 +1022,7 @@ def _finish_retired_request(self, sid: str) -> None: if entry is None: return seq = entry[0] - if hasattr(seq, "_offload_finished_block_ids") and not self._has_pending_save( - seq - ): + if getattr(seq, "_offload_finished", False) and not self._has_pending_save(seq): self._save_tracker.pop(sid, None) def blocks_waiting_for_store(self) -> int: diff --git a/atom/kv_transfer/offload/connector.py b/atom/kv_transfer/offload/connector.py index 77060978f6..a111936740 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. @@ -223,6 +228,10 @@ def protected_block_ids(self, seq): callback = getattr(self._impl, "protected_block_ids", None) return callback(seq) if callback is not None else None + def can_partially_deallocate_state(self, seq) -> bool: + callback = getattr(self._impl, "can_partially_deallocate_state", None) + return callable(callback) and callback(seq) is True + def activate_block_leases(self, seq, block_ids) -> None: callback = getattr(self._impl, "activate_block_leases", None) if callback is not None: diff --git a/atom/kv_transfer/offload/hybrid/dsv4/connector.py b/atom/kv_transfer/offload/hybrid/dsv4/connector.py index 7324ae9e9b..6db327ed65 100644 --- a/atom/kv_transfer/offload/hybrid/dsv4/connector.py +++ b/atom/kv_transfer/offload/hybrid/dsv4/connector.py @@ -27,7 +27,6 @@ from __future__ import annotations import logging -import os import threading import time from contextlib import nullcontext @@ -87,6 +86,7 @@ SlotLoadSpec, SlotSaveSpec, ) +from atom.utils import envs logger = logging.getLogger("atom") @@ -131,24 +131,6 @@ def _wait_for_publication( sleep(min(poll_interval_s, remaining)) -def _env_nonnegative_float(name: str, default: float) -> float: - value = float(os.environ.get(name, str(default))) - if not isfinite(value): - raise ValueError(f"{name} must be finite") - if value < 0: - raise ValueError(f"{name} must be nonnegative") - return value - - -def _env_positive_float(name: str, default: float) -> float: - value = float(os.environ.get(name, str(default))) - if not isfinite(value): - raise ValueError(f"{name} must be finite") - if value <= 0: - raise ValueError(f"{name} must be positive") - return value - - @dataclass(frozen=True) class _SlotSaveSnapshot: """RPC-thread result carried to one background sidecar save.""" @@ -210,14 +192,8 @@ def __init__(self, config) -> None: self.virtual_block_size: int | None = None self.profile = None self.chunk_size: int | None = None - self._publication_timeout_s = _env_nonnegative_float( - "OFFLOAD_PUBLICATION_TIMEOUT_S", - 5.0, - ) - self._publication_poll_interval_s = _env_positive_float( - "OFFLOAD_PUBLICATION_POLL_INTERVAL_S", - 0.01, - ) + self._publication_timeout_s = envs.OFFLOAD_PUBLICATION_TIMEOUT_S + self._publication_poll_interval_s = envs.OFFLOAD_PUBLICATION_POLL_INTERVAL_S self._publication_clock = time.monotonic self._publication_sleep = time.sleep @@ -228,14 +204,14 @@ def __init__(self, config) -> None: # reload sat behind ~N filler saves -> request hung well past timeout). # 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")) + n_save_workers = envs.OFFLOAD_COPY_WORKERS # The SLOT load path is *not* thread-safe against itself: a worker batch # shares one staging row across its loads (`_SlotLoadBatchReservation`) # precisely because the load executor runs them in submission order. A # second load thread would run two loads through the same staging row # concurrently and corrupt both. OFFLOAD_LOAD_WORKERS therefore does not # apply to DSV4 -- honour it loudly rather than silently. - n_load_env = int(os.environ.get("OFFLOAD_LOAD_WORKERS", "1")) + n_load_env = envs.OFFLOAD_LOAD_WORKERS if n_load_env != 1: logger.warning( "ATOM DSV4 offload: ignoring OFFLOAD_LOAD_WORKERS=%d; the SLOT " @@ -450,13 +426,9 @@ def _slot_staging_slots(self) -> int: extra = kvc.get("kv_connector_extra_config", kvc) or {} configured = extra.get("slot_sidecar_staging_slots") if configured is None: - configured = os.environ.get("OFFLOAD_SLOT_STAGING_SLOTS", "1") - try: - count = int(configured) - except (TypeError, ValueError) as exc: - raise ValueError( - "SLOT sidecar staging count must be an integer" - ) from exc + count = envs.OFFLOAD_SLOT_STAGING_SLOTS + if count is None: + count = 1 else: if isinstance(configured, bool) or not isinstance(configured, Integral): raise ValueError("SLOT sidecar staging count must be an integer") @@ -1934,17 +1906,7 @@ def __init__(self, config) -> None: # chunk-aligned, recompute the misaligned head up to the next chunk # boundary, then load the aligned remainder from CPU. (Previously gated # by the OFFLOAD_UNALIGNED_HANDOFF env var; now unconditional.) - try: - self._min_load_tokens = max( - 0, int(os.environ.get("OFFLOAD_MIN_LOAD_TOKENS", "8192")) - ) - except ValueError: - logger.warning( - "LMCache offload scheduler: invalid OFFLOAD_MIN_LOAD_TOKENS=%r; " - "using 8192", - os.environ.get("OFFLOAD_MIN_LOAD_TOKENS"), - ) - self._min_load_tokens = 8192 + self._min_load_tokens = envs.OFFLOAD_MIN_LOAD_TOKENS world = offcfg.lmcache_replica_world_size(config) meta = offcfg.build_lmcache_metadata(config, cfg, world, 0) diff --git a/atom/kv_transfer/offload/hybrid/dsv4/policy.py b/atom/kv_transfer/offload/hybrid/dsv4/policy.py index 875861b9b9..599554613a 100644 --- a/atom/kv_transfer/offload/hybrid/dsv4/policy.py +++ b/atom/kv_transfer/offload/hybrid/dsv4/policy.py @@ -14,7 +14,6 @@ import hashlib import heapq import json -import os import threading from collections import OrderedDict from collections.abc import Iterator, MutableSet @@ -22,6 +21,8 @@ from math import lcm from numbers import Integral +from atom.utils import envs + @dataclass(frozen=True) class DSV4OffloadProfile: @@ -287,13 +288,9 @@ def _committed_sidecar_capacity(kvc) -> int: extra = (kvc or {}).get("kv_connector_extra_config", kvc or {}) or {} configured = extra.get("committed_sidecar_index_capacity") if configured is None: - raw = os.environ.get("OFFLOAD_COMMITTED_SIDECAR_CAPACITY", "65536") - try: - capacity = int(raw) - except (TypeError, ValueError) as exc: - raise ValueError( - "committed sidecar index capacity must be a positive integer" - ) from exc + capacity = envs.OFFLOAD_COMMITTED_SIDECAR_CAPACITY + if capacity is None: + capacity = 65536 else: if isinstance(configured, bool) or not isinstance(configured, int): raise ValueError( diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py index d32b4f3aba..538c92e56a 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py @@ -11,7 +11,6 @@ from __future__ import annotations import logging -import os import time from atom.kv_transfer.disaggregation.types import ( @@ -37,6 +36,7 @@ _JointPark, ) from atom.kv_transfer.offload.metadata import LMCacheOffloadMetadata +from atom.utils import envs logger = logging.getLogger("atom") @@ -611,7 +611,7 @@ def __init__(self, config) -> None: # 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), + envs.OFFLOAD_COPY_WORKERS, ) self._save_inflight_since: dict[str, float] = {} self._save_stalled = False diff --git a/atom/kv_transfer/offload/metadata.py b/atom/kv_transfer/offload/metadata.py index 1958df796c..f10d3b4815 100644 --- a/atom/kv_transfer/offload/metadata.py +++ b/atom/kv_transfer/offload/metadata.py @@ -127,6 +127,16 @@ class SlotLoadSpec: destination_group: int +@dataclass(frozen=True) +class NativeStateTransfer: + """One immutable native checkpoint image sharing a PAGE transfer operation.""" + + unit_ids: tuple[int, ...] + boundary_tokens: int + prefix_hash: int + destination_slot: int | None = None + + @dataclass class LMCacheReqMeta: """Everything the worker needs to load/save one request's KV this step.""" @@ -150,6 +160,7 @@ class LMCacheReqMeta: save_operation: SaveOperationId | None = None # Appended for positional compatibility with existing metadata producers. load_operation: LoadOperationId | None = None + native_state: NativeStateTransfer | None = None class LMCacheOffloadMetadata(ConnectorMetadata): diff --git a/atom/kv_transfer/offload/mp/backend.py b/atom/kv_transfer/offload/mp/backend.py index 12e309f342..ecf0e7a89d 100644 --- a/atom/kv_transfer/offload/mp/backend.py +++ b/atom/kv_transfer/offload/mp/backend.py @@ -11,11 +11,13 @@ from __future__ import annotations +import hashlib +import json import logging -import os import threading import time from collections import deque +from collections.abc import Iterable, Sequence from dataclasses import dataclass, replace from typing import Any @@ -23,19 +25,27 @@ from atom.kv_transfer.disaggregation.base import KVConnectorBase from atom.kv_transfer.disaggregation.types import ( + ConnectorCompletion, KVConnectorOutput, LoadCompletionId, LoadOperationId, SaveCompletionId, + SaveOperationId, + SaveSourceGroupId, ) from atom.kv_transfer.offload import config as offcfg from atom.kv_transfer.offload._offload_common import validated_kv_role -from atom.kv_transfer.offload.chunked_scheduler import ChunkedOffloadSchedulerBase +from atom.kv_transfer.offload.chunked_scheduler import ( + DENSE_PAGE_SOURCE_SAFE_CHANNEL, + DENSE_PAGE_STORE_CHANNEL, + ChunkedOffloadSchedulerBase, +) from atom.kv_transfer.offload.metadata import LMCacheOffloadMetadata, LMCacheReqMeta +from atom.utils import envs logger = logging.getLogger("atom") -_MP_LAYOUT_VERSION = 2 +_MP_LAYOUT_VERSION = 3 _OPERATION_TOMBSTONE_LIMIT = 4096 @@ -61,11 +71,17 @@ def _storage_kv_transfer_config(config: Any) -> dict[str, Any]: return kvc +def _mp_session_id(config: Any, request_id: Any) -> str: + """Scope one LMCache MP request session to its global DP replica.""" + + return f"{offcfg.lmcache_engine_id(config)}:{request_id}" + + def _transfer_mode(config: Any) -> str: extra = _extra_config(config) configured_mode = extra.get("lmcache.mp.mp_transfer_mode") if configured_mode is None: - configured_mode = os.environ.get("LMCACHE_MP_TRANSFER_MODE", "auto") + configured_mode = envs.LMCACHE_MP_TRANSFER_MODE transfer_mode = str(configured_mode).strip().lower() if transfer_mode not in ("auto", "lmcache_driven", "engine_driven"): raise ValueError( @@ -114,14 +130,24 @@ def _validate_mp_config(config: Any) -> tuple[int, int]: or 1, minimum=1, ) + dp_size_local = offcfg._strict_integer( + "data_parallel_size_local", + getattr(parallel_config, "data_parallel_size_local", dp_size) or dp_size, + minimum=1, + ) if pp_size != 1: raise NotImplementedError("lmcache_mp does not support PP yet") if dcp_size != 1: raise NotImplementedError("lmcache_mp does not support DCP yet") if pcp_size != 1: raise NotImplementedError("lmcache_mp does not support PCP yet") - if dp_size != 1 or bool(getattr(config, "enable_dp_attention", False)): - raise NotImplementedError("lmcache_mp currently supports TP-only deployments") + if dp_size_local != dp_size: + raise NotImplementedError( + "lmcache_mp supports DP and DP-attention only within one host; " + "multi-node DP requires one LMCache server per host and local " + "server routing " + f"(data_parallel_size={dp_size}, local={dp_size_local})" + ) _transfer_mode(config) return tp_size, pp_size @@ -140,6 +166,11 @@ def _config_has_fully_replicated_tp_pages(config: Any) -> bool: """ hf_config = getattr(config, "hf_config", None) + # MiniMax-M3 is GQA. Some TP ranks can happen to own the same KV head when + # TP exceeds the global KV-head count, but the complete PAGE object is not + # replicated across the whole TP group and must remain one shard per rank. + if offcfg._is_minimax_m3(hf_config): + return False hf_config = getattr(hf_config, "text_config", hf_config) return getattr(hf_config, "kv_lora_rank", None) is not None @@ -168,22 +199,29 @@ def _published_tp_replication_factor( transfer_tensors: Any, *, tp_size: int, + native_state: bool = False, ) -> int: - """Validate the attention backend's whole-PAGE replication declaration.""" + """Validate a backend's whole-object TP replication declaration.""" + attribute = ( + "native_state_tp_replication_factor" + if native_state + else "tp_replication_factor" + ) + label = "native STATE" if native_state else "KV PAGE" factor = offcfg._strict_integer( - "KV PAGE TP replication factor", - getattr(transfer_tensors, "tp_replication_factor", 1), + f"{label} TP replication factor", + getattr(transfer_tensors, attribute, 1), minimum=1, ) if tp_size % factor: raise ValueError( - f"KV PAGE TP replication factor {factor} must divide TP size {tp_size}" + f"{label} TP replication factor {factor} must divide TP size {tp_size}" ) if factor not in (1, tp_size): raise NotImplementedError( "lmcache_mp currently supports only sharded or fully TP-replicated " - f"PAGE layouts, got replication factor {factor} for TP size {tp_size}" + f"{label} layouts, got replication factor {factor} for TP size {tp_size}" ) return factor @@ -218,7 +256,7 @@ def _server_urls(config: Any) -> list[str]: return urls -def _model_namespace(config: Any) -> str: +def _model_namespace(config: Any, *, checkpoint_spec: Any = None) -> str: """Build a model/layout namespace shared by scheduler and workers.""" cfg = offcfg.build_lmcache_config(_storage_kv_transfer_config(config)) @@ -228,7 +266,21 @@ def _model_namespace(config: Any) -> str: cfg, world_size, ) - return f"{page_namespace}::lmcache-mp-v{_MP_LAYOUT_VERSION}" + namespace = f"{page_namespace}::lmcache-mp-v{_MP_LAYOUT_VERSION}" + if checkpoint_spec is not None: + hf = getattr(config, "hf_config", None) + hf = getattr(hf, "text_config", hf) + document = { + "checkpoint": checkpoint_spec.to_wire(), + "hf_commit": getattr(hf, "_commit_hash", None), + "revision": getattr(config, "revision", None), + "model_revision": _extra_config(config).get("lmcache.mp.model_revision"), + } + fingerprint = hashlib.sha256( + json.dumps(document, sort_keys=True).encode() + ).hexdigest()[:32] + namespace += f"::native-state-v1-{fingerprint}" + return namespace def _parallel_strategy(config: Any, worker_id: int) -> Any: @@ -247,7 +299,7 @@ def _parallel_strategy(config: Any, worker_id: int) -> Any: ) -def _make_scheduler_adapter(config: Any) -> Any: +def _make_scheduler_adapter(config: Any, *, checkpoint_spec: Any = None) -> Any: import zmq from lmcache.integration.atom import AtomMPSchedulerAdapter @@ -264,14 +316,14 @@ def _create_key(self, *args: Any, **kwargs: Any) -> Any: return _ReaderAwareSchedulerAdapter( server_url=_server_urls(config)[0], context=zmq.Context.instance(), - model_name=_model_namespace(config), + model_name=_model_namespace(config, checkpoint_spec=checkpoint_spec), block_size=int(config.kv_cache_block_size), parallel_config=_parallel_strategy(config, 0), mq_timeout=float(extra.get("lmcache.mp.mq_timeout", 300.0)), ) -def _make_worker_adapter(config: Any, rank: int) -> Any: +def _make_worker_adapter(config: Any, rank: int, *, checkpoint_spec: Any = None) -> Any: import zmq from lmcache.integration.atom import AtomMPWorkerAdapter @@ -279,7 +331,7 @@ def _make_worker_adapter(config: Any, rank: int) -> Any: return AtomMPWorkerAdapter( server_url=_server_urls(config)[0], context=zmq.Context.instance(), - model_name=_model_namespace(config), + model_name=_model_namespace(config, checkpoint_spec=checkpoint_spec), block_size=int(config.kv_cache_block_size), parallel_config=_parallel_strategy(config, rank), mq_timeout=float(extra.get("lmcache.mp.mq_timeout", 300.0)), @@ -372,9 +424,16 @@ def _build_cache_views( if bool(getattr(region, "reverse_indexed", False)): raise ValueError("lmcache_mp PAGE regions cannot be reverse-indexed") + # LMCache receives these tensors as opaque PAGE storage, not numerical + # values. Publish a zero-copy byte view so every transfer path copies + # the exact bit pattern. This is especially important on ROCm, where + # LMCache's Python raw-pointer fallback cannot express FP8 through the + # CUDA array interface and would otherwise reconstruct the destination + # as uint8 while keeping the staging object as FP8. + byte_view = view.view(torch.uint8) role = str(getattr(region, "semantic_role", None) or f"plane_{index}") - tensors[f"page.{index}.{role}"] = view - layout = (view.dtype, tuple(int(dim) for dim in view.shape[1:])) + tensors[f"page.{index}.{role}"] = byte_view + layout = (byte_view.dtype, tuple(int(dim) for dim in byte_view.shape[1:])) indices_by_layout.setdefault(layout, []).append(index) devices.add(view.device) bytes_per_block += unit_bytes @@ -406,6 +465,8 @@ class _PendingLoad: class _PendingSave: completion: SaveCompletionId future: Any | None + start: int + end: int def _terminal_future_result(future: Any | None) -> tuple[bool, Any]: @@ -450,15 +511,47 @@ def _remember_operation_tombstone( tombstones.discard(tombstone_order.popleft()) +def _chunk_ranges(start: int, end: int, chunk_size: int) -> tuple[tuple[int, int], ...]: + """Split ``[start, end)`` into LMCache chunk token ranges.""" + + return tuple( + (chunk_start, min(chunk_start + chunk_size, end)) + for chunk_start in range(start, end, chunk_size) + ) + + +def _source_safe_completions( + operation: SaveOperationId, ranges: Iterable[Sequence[int]] +) -> set[ConnectorCompletion]: + """One PAGE source-safe completion per finished token range of a save.""" + + return { + ConnectorCompletion( + DENSE_PAGE_SOURCE_SAFE_CHANNEL, + SaveSourceGroupId(operation, (tuple(token_range),)), + True, + ) + for token_range in ranges + } + + class _MPLookupClient: """Synchronous lookup facade used by ATOM's existing scheduler policy.""" token_database = None - def __init__(self, adapter: Any, *, timeout: float, poll_interval: float) -> None: + def __init__( + self, + adapter: Any, + *, + config: Any, + timeout: float, + poll_interval: float, + ) -> None: if timeout <= 0 or poll_interval <= 0: raise ValueError("LMCache MP lookup timeout and poll interval must be > 0") self._adapter = adapter + self._config = config self._timeout = timeout self._poll_interval = poll_interval self._lookups: dict[str, _LookupState] = {} @@ -473,10 +566,11 @@ def lookup(self, token_ids: list[int], lookup_id: str) -> int | None: state = _LookupState(token_ids=list(token_ids)) self._lookups[lookup_id] = state - self._adapter.maybe_submit_lookup_request(lookup_id, token_ids) + request_id = _mp_session_id(self._config, lookup_id) + self._adapter.maybe_submit_lookup_request(request_id, token_ids) deadline = time.monotonic() + self._timeout while True: - result = self._adapter.check_lookup_result(lookup_id) + result = self._adapter.check_lookup_result(request_id) if result is not None: hit = int(result) state.hit = hit @@ -513,6 +607,7 @@ def prepare_retrieve( f"invalid retrieve range for {lookup_id}: start={start}, end={end}" ) state = self._lookups.get(lookup_id) + request_id = _mp_session_id(self._config, lookup_id) if ( state is not None and state.hit is not None @@ -527,7 +622,7 @@ def prepare_retrieve( token_ids=state.token_ids, start=0, end=min(start, state.hit), - request_id=lookup_id, + request_id=request_id, ) if ( state is not None @@ -539,12 +634,12 @@ def prepare_retrieve( token_ids=state.token_ids, start=end, end=state.hit, - request_id=lookup_id, + request_id=request_id, ) if state is not None: state.retrieve_start = start state.retrieve_end = end - self._adapter.cleanup_lookup_result(lookup_id) + self._adapter.cleanup_lookup_result(request_id) def complete_retrieve(self, lookup_id: str, *, succeeded: bool) -> None: # Once a retrieve is submitted, LMCache owns the remaining lookup @@ -554,7 +649,7 @@ def complete_retrieve(self, lookup_id: str, *, succeeded: bool) -> None: # scheduler cannot distinguish a pre-submit failure; that rarer case is # left to request end_session/server TTL cleanup instead. self._lookups.pop(lookup_id, None) - self._adapter.cleanup_lookup_result(lookup_id) + self._adapter.cleanup_lookup_result(_mp_session_id(self._config, lookup_id)) def hit_tokens(self, lookup_id: str) -> int | None: state = self._lookups.get(lookup_id) @@ -562,8 +657,9 @@ def hit_tokens(self, lookup_id: str) -> int | None: def clear_lookup_status(self, lookup_id: str) -> None: state = self._lookups.pop(lookup_id, None) + request_id = _mp_session_id(self._config, lookup_id) if state is not None and state.hit is None: - result = self._adapter.check_lookup_result(lookup_id) + result = self._adapter.check_lookup_result(request_id) if result is None: logger.warning( "LMCache MP lookup for request %s is still pending during " @@ -571,7 +667,7 @@ def clear_lookup_status(self, lookup_id: str) -> None: "cleanup releases any eventual locks", lookup_id, ) - self._adapter.cleanup_lookup_result(lookup_id) + self._adapter.cleanup_lookup_result(request_id) return state.hit = int(result) # Once retrieve has started, the transfer owns the remaining read @@ -586,9 +682,9 @@ def clear_lookup_status(self, lookup_id: str) -> None: token_ids=state.token_ids, start=0, end=state.hit, - request_id=lookup_id, + request_id=request_id, ) - self._adapter.cleanup_lookup_result(lookup_id) + self._adapter.cleanup_lookup_result(request_id) class LMCacheMPConnector(KVConnectorBase): @@ -621,7 +717,12 @@ def __init__(self, config: Any) -> None: self._completed_load_operations: set[str] = set() self._completed_save_operation_order: deque[str] = deque() self._completed_load_operation_order: deque[str] = deque() - self._immediate_saves: set[SaveCompletionId] = set() + # Immediate successes include collapsed-TP non-writers. Keep their + # logical token range so they can report PAGE source-safety too: the TP + # aggregator must receive the same chunk completion from every rank + # before releasing the writer's source blocks early. + self._immediate_saves: dict[SaveCompletionId, tuple[int, int] | None] = {} + self._immediate_save_failures: set[SaveCompletionId] = set() self._immediate_load_failures: set[LoadCompletionId] = set() self._lock = threading.Lock() @@ -761,7 +862,7 @@ def _submit_load(self, req: LMCacheReqMeta, event: Any) -> None: assert req.load_spec is not None completion = req.load_operation or req.req_id - request_id = str(req.req_id) + request_id = _mp_session_id(self._config, req.req_id) operation_id = _transfer_operation_id("load", completion) with self._lock: if ( @@ -823,7 +924,7 @@ def _submit_save(self, req: LMCacheReqMeta, event: Any) -> None: assert req.save_spec is not None completion = req.save_operation or req.req_id - request_id = str(req.req_id) + request_id = _mp_session_id(self._config, req.req_id) operation_id = _transfer_operation_id("save", completion) with self._lock: if ( @@ -835,6 +936,10 @@ def _submit_save(self, req: LMCacheReqMeta, event: Any) -> None: f"duplicate LMCache MP save operation {operation_id!r}" ) self._submitting_saves.add(operation_id) + end = (len(req.token_ids) // self.chunk_size) * self.chunk_size + start = ( + int(req.save_spec.skip_leading_tokens) // self.chunk_size + ) * self.chunk_size if not self._is_kv_writer: with self._lock: self._submitting_saves.discard(operation_id) @@ -843,12 +948,10 @@ def _submit_save(self, req: LMCacheReqMeta, event: Any) -> None: self._completed_save_operations, self._completed_save_operation_order, ) - self._immediate_saves.add(completion) + self._immediate_saves[completion] = ( + (start, end) if start < end else None + ) return - end = (len(req.token_ids) // self.chunk_size) * self.chunk_size - start = ( - int(req.save_spec.skip_leading_tokens) // self.chunk_size - ) * self.chunk_size if start >= end: with self._lock: self._submitting_saves.discard(operation_id) @@ -857,7 +960,7 @@ def _submit_save(self, req: LMCacheReqMeta, event: Any) -> None: self._completed_save_operations, self._completed_save_operation_order, ) - self._immediate_saves.add(completion) + self._immediate_saves[completion] = None return try: block_ids = self._block_slice(req, start, end) @@ -867,7 +970,12 @@ def _submit_save(self, req: LMCacheReqMeta, event: Any) -> None: start=start, end=end, ) - transfer = self._adapter.submit_store_request( + submit = getattr( + self._adapter, + "submit_store_request_with_chunk_events", + self._adapter.submit_store_request, + ) + transfer = submit( request_id, op, event, @@ -884,13 +992,15 @@ def _submit_save(self, req: LMCacheReqMeta, event: Any) -> None: self._completed_save_operations, self._completed_save_operation_order, ) - self._immediate_saves.add(completion) + self._immediate_save_failures.add(completion) return with self._lock: self._submitting_saves.discard(operation_id) self._pending_saves[operation_id] = _PendingSave( completion=completion, future=transfer, + start=start, + end=end, ) def get_finished(self) -> KVConnectorOutput: @@ -899,6 +1009,7 @@ def get_finished(self) -> KVConnectorOutput: done_load: set[LoadCompletionId] = set() failed_load: set[LoadCompletionId] = set() done_save: set[SaveCompletionId] = set() + connector_completions: set[ConnectorCompletion] = set() with self._lock: # Heartbeat health is a control-plane signal, not proof that GPU # work submitted before the failure has quiesced. Keep every real @@ -906,7 +1017,20 @@ def get_finished(self) -> KVConnectorOutput: # represented by None, while LMCache's missing-registration path # returns an event-free terminal False future. for operation_id, pending in list(self._pending_saves.items()): - terminal, _result = _terminal_future_result(pending.future) + take_ranges = getattr(pending.future, "take_completed_ranges", None) + if callable(take_ranges) and isinstance( + pending.completion, SaveOperationId + ): + try: + connector_completions |= _source_safe_completions( + pending.completion, take_ranges() + ) + except Exception: + logger.warning( + "LMCache MP source-safe event polling failed", + exc_info=True, + ) + terminal, result = _terminal_future_result(pending.future) if terminal: self._pending_saves.pop(operation_id, None) _remember_operation_tombstone( @@ -915,6 +1039,20 @@ def get_finished(self) -> KVConnectorOutput: self._completed_save_operation_order, ) done_save.add(pending.completion) + if isinstance(pending.completion, SaveOperationId): + connector_completions |= _source_safe_completions( + pending.completion, + _chunk_ranges( + pending.start, pending.end, int(self.chunk_size) + ), + ) + connector_completions.add( + ConnectorCompletion( + DENSE_PAGE_STORE_CHANNEL, + pending.completion, + result is True, + ) + ) for operation_id, pending in list(self._pending_loads.items()): terminal, result = _terminal_future_result(pending.future) @@ -931,20 +1069,40 @@ def get_finished(self) -> KVConnectorOutput: else: done_load.add(pending.completion) done_save.update(self._immediate_saves) + for completion, token_range in self._immediate_saves.items(): + if isinstance(completion, SaveOperationId): + if token_range is not None: + connector_completions |= _source_safe_completions( + completion, + _chunk_ranges(*token_range, int(self.chunk_size)), + ) + connector_completions.add( + ConnectorCompletion(DENSE_PAGE_STORE_CHANNEL, completion, True) + ) + for completion in self._immediate_save_failures: + done_save.add(completion) + if isinstance(completion, SaveOperationId): + connector_completions.add( + ConnectorCompletion(DENSE_PAGE_STORE_CHANNEL, completion, False) + ) failed_load.update(self._immediate_load_failures) self._immediate_saves.clear() + self._immediate_save_failures.clear() self._immediate_load_failures.clear() return KVConnectorOutput( finished_loading=done_load, failed_loading=failed_load, finished_saving=done_save, + connector_completions=connector_completions, ) class LMCacheMPConnectorScheduler(ChunkedOffloadSchedulerBase): """Scheduler-side LMCache MP connector for generic PAGE offload.""" - def __init__(self, config: Any) -> None: + _supports_early_block_release = True + + def __init__(self, config: Any, *, checkpoint_spec: Any = None) -> None: _validate_mp_config(config) kvc = getattr(config, "kv_transfer_config", {}) or {} validated_kv_role(kvc) @@ -953,13 +1111,14 @@ def __init__(self, config: Any) -> None: config.kv_cache_block_size, minimum=1, ) - adapter = _make_scheduler_adapter(config) + adapter = _make_scheduler_adapter(config, checkpoint_spec=checkpoint_spec) try: extra = _extra_config(config) timeout = float(extra.get("lmcache.mp.lookup_timeout", 30.0)) poll_interval = float(extra.get("lmcache.mp.lookup_poll_interval", 0.01)) lookup_client = _MPLookupClient( adapter, + config=config, timeout=timeout, poll_interval=poll_interval, ) @@ -975,6 +1134,10 @@ def __init__(self, config: Any) -> None: shutdown() raise + def save_abandon_timeout_s(self) -> float: + """A timeout cannot prove that a remote MP DMA stopped reading HBM.""" + return 0.0 + def get_num_new_matched_tokens(self, seq: Any) -> tuple[int, bool]: matched = super().get_num_new_matched_tokens(seq) sid = str(seq.id) @@ -1053,7 +1216,7 @@ def load_failed(self, req_id: Any) -> bool: def request_finished(self, seq: Any) -> None: super().request_finished(seq) try: - self._mp_adapter.end_session(str(seq.id)) + self._mp_adapter.end_session(_mp_session_id(self._config, seq.id)) except Exception: logger.warning( "LMCache MP end_session failed for request %s", diff --git a/atom/kv_transfer/offload/mp/connector.py b/atom/kv_transfer/offload/mp/connector.py index d417b70e45..3cc1a4e3a6 100644 --- a/atom/kv_transfer/offload/mp/connector.py +++ b/atom/kv_transfer/offload/mp/connector.py @@ -1,16 +1,189 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Public LMCache multiprocess connector. +"""Capability-selected public LMCache multiprocess connector. -The implementation is layout-neutral: attention backends describe their PAGE -storage through ``KVTransferTensors`` and this connector transports those views -without selecting a model-specific implementation. +The configured connector name stays ``lmcache_mp``. The scheduler selects its +implementation when ``BlockManager`` exposes a PAGE-backed checkpoint +coordinator; each worker makes the same choice when the attention backend +publishes checkpoint geometry plus its copy callback. No model name or +architecture string participates in the decision. """ -from atom.kv_transfer.offload.mp.backend import ( - LMCacheMPConnector, - LMCacheMPConnectorScheduler, +from __future__ import annotations + +from typing import Any + +from atom.kv_transfer.disaggregation.base import ( + KVConnectorBase, + KVConnectorSchedulerBase, ) +from atom.kv_transfer.disaggregation.types import KVConnectorOutput +from atom.kv_transfer.offload._offload_common import validated_kv_role +from atom.kv_transfer.offload.mp import backend + + +def _publishes_native_state(transfer_tensors: Any) -> bool: + """Validate and detect the reusable PAGE-backed native-state contract.""" + + spec = getattr(transfer_tensors, "paged_state_checkpoint_spec", None) + copy = getattr(transfer_tensors, "execute_paged_state_copies", None) + if spec is None and copy is None: + return False + if spec is None: + raise ValueError( + "execute_paged_state_copies requires paged_state_checkpoint_spec" + ) + if not callable(copy): + raise TypeError( + "paged_state_checkpoint_spec requires a callable execute_paged_state_copies" + ) + return True + + +def _copies_paged_state(transfer_tensors: Any) -> bool: + """Whether the backend keeps PAGE-backed state checkpoints. + + This is the predicate `BlockManager` builds `paged_state_checkpoints` from, + and that coordinator is what selects the native-state scheduler. + """ + + builder = getattr(transfer_tensors, "state_backend", None) + state_transfer = getattr(builder, "state_transfer", None) + return callable(state_transfer) and bool(state_transfer().copies) + + +class LMCacheMPConnector(KVConnectorBase): + """Worker shell selecting PAGE-only or native-state transport by capability.""" + + is_producer = False + + def __init__(self, config: Any) -> None: + backend._validate_mp_config(config) + validated_kv_role(getattr(config, "kv_transfer_config", {}) or {}) + self._config = config + self._impl: KVConnectorBase | None = None + + def _require_impl(self) -> KVConnectorBase: + if self._impl is None: + raise RuntimeError("lmcache_mp KV caches are not registered") + return self._impl + + def register_kv_caches( + self, + kv_caches: dict[str, Any], + transfer_tensors: Any = None, + num_blocks: int | None = None, + ) -> None: + if self._impl is not None: + raise RuntimeError("lmcache_mp KV caches are already registered") + if _publishes_native_state(transfer_tensors): + from atom.kv_transfer.offload.mp.native_state_worker import ( + NativeStateLMCacheMPConnector, + ) + + impl: KVConnectorBase = NativeStateLMCacheMPConnector(self._config) + elif _copies_paged_state(transfer_tensors): + # The scheduler is already native-state for this backend. A PAGE-only + # worker would store under a namespace it never looks up. + raise NotImplementedError( + "lmcache_mp: this backend keeps PAGE-backed state checkpoints " + "but publishes no paged_state_checkpoint_spec and " + "execute_paged_state_copies for the native-state transfer" + ) + else: + impl = backend.LMCacheMPConnector(self._config) + impl.register_kv_caches(kv_caches, transfer_tensors, num_blocks) + self._impl = impl + + def start_load_kv(self, metadata: Any) -> None: + self._require_impl().start_load_kv(metadata) + + def get_finished(self) -> KVConnectorOutput: + if self._impl is None: + return KVConnectorOutput() + return self._impl.get_finished() + + def get_finished_recv_blocks(self) -> list[int]: + if self._impl is None: + return [] + return self._impl.get_finished_recv_blocks() + + def __getattr__(self, name: str) -> Any: + impl = self.__dict__.get("_impl") + if impl is None: + raise AttributeError(name) + return getattr(impl, name) + + +class LMCacheMPConnectorScheduler(KVConnectorSchedulerBase): + """Scheduler shell selecting transport after checkpoint capability exists.""" + + is_producer = False + is_offload = True + + def __init__(self, config: Any) -> None: + backend._validate_mp_config(config) + validated_kv_role(getattr(config, "kv_transfer_config", {}) or {}) + self._config = config + self._impl: KVConnectorSchedulerBase | None = None + self._block_manager: Any = None + + def _require_impl(self) -> KVConnectorSchedulerBase: + if self._impl is None: + raise RuntimeError("lmcache_mp scheduler is not bound to BlockManager") + return self._impl + + def bind_block_manager(self, block_manager: Any) -> None: + if self._block_manager is block_manager: + return + if self._block_manager is not None: + raise RuntimeError( + "lmcache_mp scheduler is already bound to a block manager" + ) + coordinator = getattr(block_manager, "paged_state_checkpoints", None) + if coordinator is None: + impl: KVConnectorSchedulerBase = backend.LMCacheMPConnectorScheduler( + self._config + ) + else: + from atom.kv_transfer.offload.mp.native_state_scheduler import ( + NativeStateLMCacheMPConnectorScheduler, + ) + + impl = NativeStateLMCacheMPConnectorScheduler(self._config) + bind = getattr(impl, "bind_block_manager", None) + if callable(bind): + bind(block_manager) + self._impl = impl + self._block_manager = block_manager + + def get_num_new_matched_tokens(self, seq: Any) -> tuple[int, bool]: + return self._require_impl().get_num_new_matched_tokens(seq) + + def build_connector_meta(self) -> Any: + return self._require_impl().build_connector_meta() + + def update_state_after_alloc(self, seq: Any) -> None: + self._require_impl().update_state_after_alloc(seq) + + def request_finished(self, seq: Any) -> None: + self._require_impl().request_finished(seq) + + def should_defer_free(self, seq: Any) -> bool: + return self._require_impl().should_defer_free(seq) + + def send_finished(self, req_id: Any) -> None: + self._require_impl().send_finished(req_id) + + def source_blocks_released(self, seq: Any) -> None: + self._require_impl().source_blocks_released(seq) + + def __getattr__(self, name: str) -> Any: + impl = self.__dict__.get("_impl") + if impl is None: + raise AttributeError(name) + return getattr(impl, name) + __all__ = ["LMCacheMPConnector", "LMCacheMPConnectorScheduler"] diff --git a/atom/kv_transfer/offload/mp/native_state_layout.py b/atom/kv_transfer/offload/mp/native_state_layout.py new file mode 100644 index 0000000000..cf932237a5 --- /dev/null +++ b/atom/kv_transfer/offload/mp/native_state_layout.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Owned tensor aliases for generic PAGE-backed native checkpoint images. + +The native codec streams PAGE units in ordinal order, with regions inside +each unit. Every ordinal gets its own engine block-id list while aliasing the +original PAGE allocation. No Active SLOT or extra state allocation is exposed. +LMCache imports stay behind ``engine_group_infos`` so layout validation and +copy-plan tests run without a GPU or LMCache's native extension. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import torch + + +@dataclass(frozen=True) +class NativeStateMPKernelGroup: + """One actual dtype/shape/engine-address-space registration identity.""" + + tensor_indices: tuple[int, ...] + engine_group_id: int + tokens_per_block: int + sw_size_tokens: int = -1 + recurrent_state: bool = False + extra_object_group_tag: int = 0 + + +@dataclass(frozen=True) +class NativeStateMPLayout: + tensors: tuple[torch.Tensor, ...] + kernel_groups: tuple[NativeStateMPKernelGroup, ...] + checkpoint_spec: Any + + def engine_group_infos(self) -> list[Any]: + """Convert the neutral plan at the actual LMCache registration boundary.""" + from lmcache.v1.multiprocess.group_view import EngineGroupInfo + + return [ + EngineGroupInfo( + engine_group_id=group.engine_group_id, + layer_indices=group.tensor_indices, + tokens_per_block=group.tokens_per_block, + sw_size_tokens=group.sw_size_tokens, + recurrent_state=group.recurrent_state, + extra_object_group_tag=group.extra_object_group_tag, + ) + for group in self.kernel_groups + ] + + def validate_unit_ids(self, unit_ids: Sequence[int]) -> tuple[int, ...]: + """Validate the PAGE IDs addressed by one native checkpoint image.""" + ids = tuple(unit_ids) + units_per_checkpoint = int(self.checkpoint_spec.units_per_checkpoint) + if len(ids) != units_per_checkpoint: + raise ValueError( + f"native image needs {units_per_checkpoint} unit IDs, got {len(ids)}" + ) + num_blocks = self.tensors[0].shape[0] + for unit_id in ids: + if type(unit_id) is not int or not 0 <= unit_id < num_blocks: + raise ValueError(f"invalid native checkpoint unit ID: {unit_id!r}") + if len(set(ids)) != len(ids): + raise ValueError("native checkpoint unit IDs must be distinct") + return ids + + +def _positive_int(name: str, value: Any) -> int: + if type(value) is not int or value <= 0: + raise ValueError(f"{name} must be a positive integer") + return value + + +def build_native_state_mp_layout( + transfer_tensors: Any, + *, + block_size: int, + chunk_size: int, + num_blocks: int | None = None, +) -> NativeStateMPLayout: + """Validate and alias PAGE plus compact native STATE for LMCache MP. + + PAGE uses engine group zero. STATE ordinal ``j`` uses group ``1+j`` and one + logical chunk per block. The MP server's global null block ID must be + configured as ``-1``. The final ordinal's final region ends exactly at + ``image_bytes``; its dim-0 stride continues to address the original full + PAGE region. + """ + block_size = _positive_int("native PAGE block size", block_size) + chunk_size = _positive_int("native checkpoint chunk size", chunk_size) + if chunk_size % block_size: + raise ValueError("native checkpoint chunk size must be divisible by block size") + if transfer_tensors is None: + raise ValueError("native-state LMCache MP requires KVTransferTensors") + spec = getattr(transfer_tensors, "paged_state_checkpoint_spec", None) + if spec is None: + raise ValueError("native-state LMCache MP requires paged checkpoint geometry") + layout_id = str(getattr(spec, "layout_id", "")).strip() + if not layout_id: + raise ValueError("native checkpoint layout_id must be non-empty") + page_bytes = _positive_int("native PAGE unit bytes", spec.page_unit_bytes) + image_bytes = _positive_int("native checkpoint image bytes", spec.image_bytes) + slot_bytes = _positive_int("native Active SLOT bytes", spec.slot_bytes) + if image_bytes > slot_bytes: + raise ValueError("native image bytes exceed Active SLOT bytes") + units = _positive_int("native checkpoint units", spec.units_per_checkpoint) + if units != (image_bytes + page_bytes - 1) // page_bytes: + raise ValueError("native checkpoint units disagree with image geometry") + if not callable(getattr(transfer_tensors, "execute_paged_state_copies", None)): + raise TypeError("native-state LMCache MP requires execute_paged_state_copies") + if num_blocks is None: + num_blocks = getattr(transfer_tensors, "num_blocks", None) + num_blocks = _positive_int("native scheduler block count", num_blocks) + published_blocks = getattr(transfer_tensors, "num_blocks", 0) + if published_blocks not in (0, num_blocks): + raise ValueError("native published block count disagrees with scheduler") + + regions = list(getattr(transfer_tensors, "block_regions", None) or []) + page_views = list(getattr(transfer_tensors, "block_tensor_views", None) or []) + if not regions or len(regions) != len(page_views): + raise ValueError( + "native-state LMCache MP needs one owned tensor view per PAGE region" + ) + devices = set() + actual_page_bytes = 0 + for index, (view, region) in enumerate(zip(page_views, regions, strict=True)): + if not isinstance(view, torch.Tensor): + raise TypeError(f"native PAGE view {index} must be a Tensor") + if view.ndim != 3 or view.shape[0] != num_blocks or view.numel() == 0: + raise ValueError(f"native PAGE view {index} has invalid block geometry") + if not view[0].is_contiguous(): + raise ValueError(f"native PAGE view {index} has non-contiguous inner rows") + if block_size % view.shape[1]: + raise ValueError( + f"native PAGE view {index} physical slots must divide block size" + ) + unit_bytes = _positive_int( + f"native PAGE region {index} bytes", region.unit_bytes + ) + if ( + view[0].numel() * view.element_size() != unit_bytes + or view.stride(0) * view.element_size() != unit_bytes + or region.total_bytes != num_blocks * unit_bytes + ): + raise ValueError(f"native PAGE view {index} byte geometry mismatch") + if view.data_ptr() != region.base_addr: + raise ValueError(f"native PAGE view {index} does not alias its region") + if region.reverse_indexed: + raise ValueError("native PAGE regions cannot be reverse-indexed") + devices.add(view.device) + actual_page_bytes += unit_bytes + if len(devices) != 1: + raise ValueError("native PAGE views must share one device") + if actual_page_bytes != page_bytes: + raise ValueError("PAGE regions do not cover the native PAGE unit") + + tensors = list(page_views) + engine_ids = [0] * len(tensors) + image_offset = 0 + for ordinal in range(units): + for page_view, region in zip(page_views, regions, strict=True): + nbytes = min(region.unit_bytes, image_bytes - image_offset) + if not nbytes: + break + # Retyping and reshaping contiguous inner dimensions retain the + # allocation owner. The explicit shape retains full PAGE stride, + # including for a partial final region. + byte_view = page_view.view(torch.uint8).view(num_blocks, 1, -1) + # Give the singleton physical-slot dimension its canonical tight + # stride too: LMCache's padded-layout validation checks stride(1) + # even for a size-one dimension. A plain final-axis slice retains + # the full region width there and is rejected at registration. + alias = byte_view.as_strided( + (num_blocks, 1, nbytes), (byte_view.stride(0), nbytes, 1) + ) + tensors.append(alias) + engine_ids.append(1 + ordinal) + image_offset += nbytes + if image_offset != image_bytes: + raise ValueError("PAGE aliases do not cover the native image") + + by_identity: dict[tuple, list[int]] = {} + for index, (tensor, engine_id) in enumerate(zip(tensors, engine_ids, strict=True)): + identity = tensor.dtype, tuple(tensor.shape[1:]), engine_id + members = by_identity.setdefault(identity, []) + if members and tensors[members[0]].stride(0) != tensor.stride(0): + # LMCache currently stamps one representative stride per identity. + # Splitting metadata cannot fix this: it re-coalesces the identity + # during registration. Reject instead of addressing another PAGE. + raise ValueError( + "native-state LMCache MP equal-shape regions have different " + "block strides; " + "LMCache requires one stride per physical kernel identity" + ) + members.append(index) + groups = tuple( + NativeStateMPKernelGroup( + tensor_indices=tuple(indices), + engine_group_id=identity[2], + tokens_per_block=chunk_size if identity[2] else block_size, + sw_size_tokens=chunk_size if identity[2] else -1, + recurrent_state=bool(identity[2]), + ) + for identity, indices in by_identity.items() + ) + return NativeStateMPLayout( + tensors=tuple(tensors), + kernel_groups=groups, + checkpoint_spec=spec, + ) + + +__all__ = [ + "NativeStateMPKernelGroup", + "NativeStateMPLayout", + "build_native_state_mp_layout", +] diff --git a/atom/kv_transfer/offload/mp/native_state_scheduler.py b/atom/kv_transfer/offload/mp/native_state_scheduler.py new file mode 100644 index 0000000000..b69848ef8b --- /dev/null +++ b/atom/kv_transfer/offload/mp/native_state_scheduler.py @@ -0,0 +1,516 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Bounded transfers of native READY checkpoint images over LMCache MP.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, TypeVar + +from atom.kv_transfer.disaggregation.types import ( + ConnectorCompletion, + LoadOperationId, + SaveOperationId, + SaveSourceGroupId, + StateStoreOperationId, +) +from atom.kv_transfer.offload import config as offcfg +from atom.kv_transfer.offload._offload_common import ( + max_pending_saves, + validated_kv_role, +) +from atom.kv_transfer.offload.metadata import NativeStateTransfer +from atom.kv_transfer.offload.mp.backend import ( + LMCacheMPConnectorScheduler, + _extra_config, + _validate_mp_config, +) +from atom.kv_transfer.offload.mp.native_state_worker import ( + NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL, + NATIVE_STATE_MP_STORE_CHANNEL, + require_native_state_server, +) +from atom.model_engine.page_unit_checkpoint import SuspendedCheckpointRestore +from atom.utils import envs + +_MAX_SAVE_ATTEMPTS = 3 + +_T = TypeVar("_T") + + +@dataclass +class _NativeSave: + seq: Any + source: StateStoreOperationId + saved_before: int + boundary: int + source_safe: bool = False + + +@dataclass +class _NativeLoad: + seq: Any + transfer: NativeStateTransfer + hbm_cached_tokens: int + local_restore: SuspendedCheckpointRestore | None = None + dispatched: bool = False + + +class NativeStateLMCacheMPConnectorScheduler(LMCacheMPConnectorScheduler): + """Pair PAGE KV with an exact native state image under one operation ID. + + Native sources are acquired only after save admission. Loads reserve raw + PAGE units after request allocation and before parking; their release waits + for the worker's terminal transfer-and-restore report. Neither side recycles + these units in response to elapsed time. + """ + + _supports_early_block_release = True + + def __init__(self, config: Any) -> None: + # The layout is known only after BlockManager builds the native pool. + # Defer connecting so scheduler and workers use the same namespace. + _validate_mp_config(config) + self._config = config + self.kv_role = validated_kv_role( + getattr(config, "kv_transfer_config", {}) or {} + ) + self._do_save = self.kv_role in ("offload", "kv_both", "kv_producer") + self._do_load = self.kv_role in ("offload", "kv_both", "kv_consumer") + self._block_manager = None + + def bind_block_manager(self, block_manager: Any) -> None: + if self._block_manager is block_manager: + return + if self._block_manager is not None: + raise RuntimeError( + "native-state LMCache MP scheduler is already bound to a block manager" + ) + coordinator = getattr(block_manager, "paged_state_checkpoints", None) + if coordinator is None: + raise ValueError( + "native-state LMCache MP requires PAGE checkpoint geometry" + ) + super().__init__(self._config, checkpoint_spec=coordinator.store.spec) + try: + require_native_state_server(self._mp_adapter, self._config) + self._hash_block_size = int(block_manager.hash_block_size) + if self.chunk_size % self._hash_block_size: + raise ValueError( + "native-state LMCache MP chunk size must align to native " + "hash blocks" + ) + self._max_pending_saves = max_pending_saves( + getattr(self._config, "kv_transfer_config", {}) or {}, + envs.OFFLOAD_COPY_WORKERS, + ) + self._image_reservation_bytes = ( + coordinator.store.spec.units_per_checkpoint + * coordinator.store.spec.page_unit_bytes + ) + self._max_pinned_state_bytes = offcfg._strict_integer( + "lmcache.mp.max_pinned_state_bytes", + _extra_config(self._config).get( + "lmcache.mp.max_pinned_state_bytes", + self._max_pending_saves * self._image_reservation_bytes, + ), + minimum=self._image_reservation_bytes, + ) + except Exception: + shutdown = getattr(self._mp_adapter, "shutdown", None) + if callable(shutdown): + shutdown() + raise + self._block_manager = block_manager + self._checkpoints = coordinator + self._pinned_state_bytes = 0 + self._native_saves: dict[SaveOperationId, _NativeSave] = {} + self._native_loads: dict[LoadOperationId, _NativeLoad] = {} + self._native_load_operations: dict[str, LoadOperationId] = {} + # One retry record per tracked request: a newer boundary can retry the + # complete unsaved prefix, but a permanently failing boundary cannot + # hold a finished request forever. + self._save_failures: dict[str, tuple[Any, int, int]] = {} + self._retired_requests: dict[str, Any] = {} + + def can_partially_deallocate_state(self, seq: Any) -> bool: + """The checkpoint coordinator owns state sources past request free. + + Native-state saves never read the request's active state slot after + admission: ``acquire_checkpoint_source`` gives each emitted save an + independent checkpoint PAGE lease. Waiting candidates likewise refer + to READY checkpoint content and acquire that lease only when admitted. + """ + return ( + bool(getattr(seq, "has_per_req_cache", False)) + and self._block_manager is not None + and getattr(self, "_checkpoints", None) + is getattr(self._block_manager, "paged_state_checkpoints", None) + ) + + def _boundary_hash(self, seq: Any, boundary: int) -> int: + count = boundary // self._hash_block_size + if boundary <= 0 or boundary % self._hash_block_size: + raise ValueError("native checkpoint boundary must align to hash blocks") + chain = getattr(seq, "block_hashes", ()) + if len(chain) >= count: + return int(chain[count - 1]) + # Some checkpoint producers do not reserve midstep checkpoints, so + # BlockManager may leave Sequence.block_hashes empty. Use its exact + # public hashing algorithm and token slices, caching only this request's + # immutable prompt chain. + chain = getattr(seq, "_mp_checkpoint_hashes", None) + if chain is None: + chain = [] + seq._mp_checkpoint_hashes = chain + prefix = chain[-1] if chain else -1 + for index in range(len(chain), count): + start = index * self._hash_block_size + prefix = self._block_manager.compute_hash( + seq.token_ids[start : start + self._hash_block_size], prefix + ) + chain.append(prefix) + return int(chain[count - 1]) + + def _lookup_token_ids(self, seq: Any) -> list[int]: + # Truncate the QUERY, not its answer: MP recurrent/window keys and read + # locks belong to the requested endpoint. A full-prompt query followed + # by rounding the result would pair the earlier KV with later state. + boundary = self._chunk_floor(max(0, int(seq.num_prompt_tokens) - 1)) + return list(seq.token_ids[:boundary]) + + def get_num_new_matched_tokens(self, seq: Any) -> tuple[int, bool]: + if not getattr(seq, "has_per_req_cache", False): + return 0, False + previous = self._native_load_operations.get(str(seq.id)) + if previous is not None: + return 0, False + return super().get_num_new_matched_tokens(seq) + + def _has_state_budget(self) -> bool: + return ( + self._pinned_state_bytes + self._image_reservation_bytes + <= self._max_pinned_state_bytes + ) + + def _charge_state_image(self, acquire: Callable[[], _T | None]) -> _T | None: + """Charge one image to the pinned-state budget around ``acquire()``. + + The charge lands before ``acquire`` takes any pin and is refunded when + it returns None or raises, so a refused source or reservation leaves the + budget as it found it. All callers run on the scheduler thread. + """ + self._pinned_state_bytes += self._image_reservation_bytes + try: + result = acquire() + except BaseException: + self._refund_state_image() + raise + if result is None: + self._refund_state_image() + return result + + def _refund_state_image(self) -> None: + self._pinned_state_bytes -= self._image_reservation_bytes + + def _save_frontier(self, seq: Any) -> int: + if not getattr(seq, "has_per_req_cache", False): + return 0 + frontier = super()._save_frontier(seq) + failed = self._save_failures.get(str(seq.id)) + exhausted = ( + failed[1] + if failed is not None + and failed[0] is seq + and failed[2] >= _MAX_SAVE_ATTEMPTS + else 0 + ) + for boundary in range(frontier, exhausted, -self.chunk_size): + if self._checkpoints.contains(self._boundary_hash(seq, boundary)): + return boundary + return 0 + + def _late_save_frontier(self, seq: Any, saved: int, available: int) -> int: + available = self._chunk_floor(available) + for boundary in range(available, saved, -self.chunk_size): + if self._checkpoints.contains(self._boundary_hash(seq, boundary)): + return boundary + return saved + + def _may_emit_save(self) -> bool: + return ( + len(self._save_inflight) < self._max_pending_saves + and self._has_state_budget() + ) + + def _build_save_request( + self, seq, saved, aligned, operation, block_ids, is_last_prefill + ): + if not self._may_emit_save(): + return None + prefix_hash = self._boundary_hash(seq, aligned) + source = self._charge_state_image( + lambda: self._checkpoints.acquire_checkpoint_source( + prefix_hash, max_inflight=self._max_pending_saves + ) + ) + if source is None: + return None + state_operation, unit_ids = source + request = super()._build_save_request( + seq, saved, aligned, operation, block_ids, is_last_prefill + ) + request.native_state = NativeStateTransfer(unit_ids, aligned, prefix_hash) + self._native_saves[operation] = _NativeSave( + seq, state_operation, saved, aligned + ) + return request + + def _complete_native_save( + self, operation: SaveOperationId, *, succeeded: bool + ) -> None: + lease = self._native_saves.pop(operation, None) + if lease is None: + return + # A terminal MP event is a backstop source fence even if an older + # server did not provide per-chunk milestones, or those notifications + # were coalesced/lost before reaching the scheduler. + self._source_group_finished( + SaveSourceGroupId(operation, ((lease.saved_before, lease.boundary),)) + ) + if not lease.source_safe: + self._checkpoints.release_offload_store_source(lease.source) + self._checkpoints.settle_offload_store(lease.source) + self._refund_state_image() + sid = str(operation.req_id) + if succeeded: + self._save_failures.pop(sid, None) + else: + failed = self._save_failures.get(sid) + attempts = ( + failed[2] + 1 + if failed is not None and failed[:2] == (lease.seq, lease.boundary) + else 1 + ) + self._save_failures[sid] = (lease.seq, lease.boundary, attempts) + entry = self._save_tracker.get(sid) + if entry is not None and entry[0] is lease.seq: + entry[1] = min(int(entry[1]), lease.saved_before) + self._store_finished(operation, succeeded=succeeded) + + def save_finished(self, req_id: Any) -> None: + if isinstance(req_id, SaveOperationId): + self._complete_native_save(req_id, succeeded=True) + + def connector_completion(self, completion: ConnectorCompletion) -> bool | None: + if completion.channel == NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL: + operation = completion.operation_id + if not isinstance(operation, SaveOperationId): + return False + lease = self._native_saves.get(operation) + if lease is not None and not lease.source_safe: + self._checkpoints.release_offload_store_source(lease.source) + lease.source_safe = True + return None + if completion.channel != NATIVE_STATE_MP_STORE_CHANNEL: + return super().connector_completion(completion) + if not isinstance(completion.operation_id, SaveOperationId): + return False + self._complete_native_save( + completion.operation_id, succeeded=completion.succeeded + ) + return True + + def _decide_load_after_alloc(self, seq: Any, load_spec): + decision = super()._decide_load_after_alloc(seq, load_spec) + should_load, _reason, hbm, lmc, need, chunk = decision + if not should_load: + return decision + if not getattr(seq, "has_per_req_cache", False) or seq.state_slot < 0: + return False, "native_destination_missing", hbm, lmc, need, chunk + if lmc % chunk or lmc >= int(seq.num_prompt_tokens): + return False, "native_boundary_unaligned", hbm, lmc, need, chunk + sid = str(seq.id) + operation = self._native_load_operations.get(sid) + if operation is not None: + lease = self._native_loads[operation] + if lease.seq is seq: + return decision + return False, "native_load_id_busy", hbm, lmc, need, chunk + if not self._has_state_budget(): + return False, "native_state_budget", hbm, lmc, need, chunk + operation = LoadOperationId(seq.id, self._load_nonce) + self._load_nonce += 1 + units = self._charge_state_image( + lambda: self._checkpoints.reserve_transfer_units(operation) + ) + if units is None: + return False, "native_state_units", hbm, lmc, need, chunk + local_restore = self._checkpoints.suspend_queued_restore(int(seq.state_slot)) + self._native_loads[operation] = _NativeLoad( + seq, + NativeStateTransfer( + units, + lmc, + self._boundary_hash(seq, lmc), + destination_slot=int(seq.state_slot), + ), + hbm, + local_restore, + ) + self._native_load_operations[sid] = operation + self._active_load_operations[sid] = (seq, operation) + seq._load_operation = operation + return decision + + def _new_load_operation(self, seq: Any) -> LoadOperationId: + operation = self._native_load_operations[str(seq.id)] + self._native_loads[operation].dispatched = True + # The scheduler publishes the loaded PAGE prefix only after terminal + # success. Its hash chain must exist before suffix prefill checkpoints. + seq.offload_load_start_tokens = self._native_loads[operation].hbm_cached_tokens + return operation + + def build_connector_meta(self): + metadata = super().build_connector_meta() + for request in metadata.requests: + if request.load_operation is not None: + request.native_state = self._native_loads[ + request.load_operation + ].transfer + return metadata + + def _release_native_load( + self, operation: LoadOperationId, *, release_units: bool = True + ) -> _NativeLoad | None: + lease = self._native_loads.pop(operation, None) + if lease is not None: + if release_units: + self._checkpoints.release_transfer_units(operation) + self._refund_state_image() + sid = str(operation.req_id) + if self._native_load_operations.get(sid) == operation: + del self._native_load_operations[sid] + return lease + + def _clear_pending_load(self, sid: str) -> None: + operation = self._native_load_operations.get(sid) + lease = self._native_loads.get(operation) + if lease is not None and not lease.dispatched: + if lease.local_restore is not None: + self._checkpoints.resume_suspended_restore(lease.local_restore) + lease.local_restore = None + self._release_native_load(operation) + if self._active_load_operations.get(sid) == (lease.seq, operation): + self._active_load_operations.pop(sid, None) + if getattr(lease.seq, "_load_operation", None) == operation: + delattr(lease.seq, "_load_operation") + super()._clear_pending_load(sid) + + def _finish_native_load(self, operation: Any, *, succeeded: bool) -> bool: + if ( + not isinstance(operation, LoadOperationId) + or operation not in self._native_loads + ): + return False + lease = self._native_loads[operation] + if succeeded: + self._checkpoints.adopt_transfer_units( + operation, lease.transfer.prefix_hash + ) + if lease.local_restore is not None: + self._checkpoints.release_suspended_restore(lease.local_restore) + lease.local_restore = None + self._release_native_load(operation, release_units=False) + else: + if lease.local_restore is not None: + self._checkpoints.resume_suspended_restore(lease.local_restore) + lease.local_restore = None + self._release_native_load(operation) + finished = ( + super().load_finished(operation) + if succeeded + else super().load_failed(operation) + ) + if self._retired_requests.get(str(operation.req_id)) is lease.seq: + self.request_finished(lease.seq) + return finished + + def load_finished(self, req_id: Any) -> bool: + return self._finish_native_load(req_id, succeeded=True) + + def load_failed(self, req_id: Any) -> bool: + return self._finish_native_load(req_id, succeeded=False) + + def _has_active_load(self, seq: Any) -> bool: + # The native lease also retains a cancelled/reused request lifecycle + # whose generic request-ID entry may have been replaced. + return any( + lease.seq is seq for lease in self._native_loads.values() + ) or super()._has_active_load(seq) + + def cancel_pending_load(self, seq: Any) -> None: + operation = self._native_load_operations.get(str(seq.id)) + lease = self._native_loads.get(operation) + if lease is not None and lease.seq is seq and lease.dispatched: + return # A request cancellation does not cancel an MP DMA. + super().cancel_pending_load(seq) + + def request_finished(self, seq: Any) -> None: + sid = str(seq.id) + self._retired_requests[sid] = seq + operation = self._native_load_operations.get(sid) + lease = self._native_loads.get(operation) + if lease is not None and lease.seq is seq and lease.dispatched: + return + super().request_finished(seq) + self._finish_retired_request(sid) + + def _finish_retired_request(self, sid: str) -> None: + super()._finish_retired_request(sid) + seq = self._retired_requests.get(sid) + if seq is not None and not self.should_defer_free(seq): + entry = self._save_tracker.get(sid) + if entry is not None and entry[0] is seq: + self._save_tracker.pop(sid, None) + self._save_failures.pop(sid, None) + self._retired_requests.pop(sid, None) + + def has_pending_work(self) -> bool: + return ( + bool(self._native_loads) + or bool(self._retired_requests) + or super().has_pending_work() + ) + + def process_completions(self, output): + output = super().process_completions(output) + # A retired request may have waited behind admission with an unpinned + # READY candidate. If eviction spends it before dispatch, no worker + # completion exists to wake deferred-free cleanup. Report that local + # terminal condition only for this same request lifecycle, after all + # actual transfers and every remaining candidate are gone. + for sid, seq in list(self._retired_requests.items()): + entry = self._save_tracker.get(sid) + lifecycle = self._load_lifecycles.get(sid) + if (entry is not None and entry[0] is not seq) or ( + lifecycle is not None and lifecycle is not seq + ): + continue + if not self.should_defer_free(seq): + output.finished_saving.add(seq.id) + self._finish_retired_request(sid) + return output + + def abandon_save(self, req_id: Any) -> None: + # This legacy callback carries no proof that the MP reader stopped. + # Explicit terminal failures already settle the exact operation above. + return None + + def reclaim_stale_leases(self, timeout_s: float) -> list[frozenset]: + return [] + + +__all__ = ["NativeStateLMCacheMPConnectorScheduler"] diff --git a/atom/kv_transfer/offload/mp/native_state_worker.py b/atom/kv_transfer/offload/mp/native_state_worker.py new file mode 100644 index 0000000000..70c33dfa2a --- /dev/null +++ b/atom/kv_transfer/offload/mp/native_state_worker.py @@ -0,0 +1,416 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Native PAGE-backed checkpoint transport over standalone LMCache.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Any + +import torch + +from atom.kv_transfer.disaggregation.types import ( + ConnectorCompletion, + KVConnectorOutput, +) +from atom.kv_transfer.offload import config as offcfg +from atom.kv_transfer.offload._offload_common import max_pending_saves +from atom.kv_transfer.offload.metadata import LMCacheReqMeta, NativeStateTransfer +from atom.kv_transfer.offload.mp.backend import ( + LMCacheMPConnector, + _chunk_ranges, + _make_worker_adapter, + _mp_session_id, + _published_tp_replication_factor, + _remember_operation_tombstone, + _source_safe_completions, + _storage_kv_transfer_config, + _terminal_future_result, + _tp_replication_factor, + _transfer_operation_id, + _validate_mp_config, +) +from atom.kv_transfer.offload.mp.native_state_layout import ( + build_native_state_mp_layout, +) +from atom.model_engine.page_unit_checkpoint import CheckpointRestoreOp +from atom.utils import envs + +logger = logging.getLogger("atom") +NATIVE_STATE_MP_STORE_CHANNEL = "native_state_mp_store" +NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL = "native_state_mp_state_source_safe" + + +def require_native_state_server(adapter: Any, config: Any) -> None: + """Validate native transfer geometry shared with the LMCache server.""" + configured_chunk = int( + offcfg.build_lmcache_config(_storage_kv_transfer_config(config)).chunk_size + ) + if configured_chunk != int(adapter.lmcache_tokens_per_chunk): + raise ValueError( + "Native-state LMCache configured chunk size must match the MP server: " + f"configured={configured_chunk}, server={adapter.lmcache_tokens_per_chunk}" + ) + + +class _UncertainSubmission: + """A transport exception cannot prove that a remote DMA stopped.""" + + def query(self) -> bool: + return False + + +@dataclass +class _NativePending: + request: LMCacheReqMeta + future: Any + restore_event: Any = None + restore_succeeded: bool = False + descriptor_slot: int | None = None + immediate_success: bool = False + state_source_safe: bool = False + + +class NativeStateLMCacheMPConnector(LMCacheMPConnector): + """Transfer pinned native images; never read a request's active SLOT to save.""" + + def __init__(self, config: Any) -> None: + super().__init__(config) + self._native_saves: dict[str, _NativePending] = {} + self._native_loads: dict[str, _NativePending] = {} + self._native_layout = None + self._native_copy = None + self._restore_stream = None + self._restore_descriptor_slots: list[int] = [] + self._max_pending_saves = max_pending_saves( + getattr(config, "kv_transfer_config", {}) or {}, + envs.OFFLOAD_COPY_WORKERS, + ) + + def register_kv_caches( + self, + _kv_caches: dict[str, Any], + transfer_tensors: Any = None, + num_blocks: int | None = None, + ) -> None: + from aiter.dist.parallel_state import get_tp_group + + spec = getattr(transfer_tensors, "paged_state_checkpoint_spec", None) + native_copy = getattr(transfer_tensors, "execute_paged_state_copies", None) + if spec is None or not callable(native_copy): + raise ValueError( + "native-state LMCache MP needs checkpoint geometry and a copy callback" + ) + tp_size, _ = _validate_mp_config(self._config) + rank = int(get_tp_group().rank_in_group) + requested_replication = _tp_replication_factor(self._config) + published_page_replication = _published_tp_replication_factor( + transfer_tensors, tp_size=tp_size + ) + published_state_replication = _published_tp_replication_factor( + transfer_tensors, tp_size=tp_size, native_state=True + ) + if requested_replication > min( + published_page_replication, published_state_replication + ): + raise ValueError( + "LMCache MP native TP rank collapse was requested, but the " + "attention backend did not declare both PAGE and STATE fully " + "replicated" + ) + self._is_kv_writer = rank % requested_replication == 0 + adapter = _make_worker_adapter(self._config, rank, checkpoint_spec=spec) + try: + require_native_state_server(adapter, self._config) + chunk_size = int(adapter.lmcache_tokens_per_chunk) + layout = build_native_state_mp_layout( + transfer_tensors, + block_size=self.block_size, + chunk_size=chunk_size, + num_blocks=num_blocks, + ) + adapter.register_kv_caches( + { + f"native_state.{i}": tensor + for i, tensor in enumerate(layout.tensors) + }, + engine_group_infos=layout.engine_group_infos(), + ) + except Exception: + adapter.shutdown() + raise + self._adapter = adapter + self.chunk_size = chunk_size + self._native_layout = layout + self._native_copy = native_copy + self._restore_stream = torch.cuda.Stream(device=torch.cuda.current_device()) + self._restore_descriptor_slots = list( + range(1, max(2, self._max_pending_saves) + 1) + ) + logger.info( + "LMCache MP native state registered rank=%d native_image=%d " + "units=%d groups=%d chunk=%d", + rank, + spec.image_bytes, + spec.units_per_checkpoint, + len(layout.kernel_groups), + chunk_size, + ) + + def _native_block_ids( + self, + req: LMCacheReqMeta, + start: int, + end: int, + *, + loading: bool, + ) -> list[list[int]]: + state = req.native_state + if state is None: + raise ValueError( + "native-state LMCache MP transfers require an exact checkpoint" + ) + if start % self.chunk_size or end % self.chunk_size or start >= end: + raise ValueError( + "native-state LMCache MP requires a nonempty chunk-aligned range" + ) + if state.boundary_tokens != end or len(req.token_ids) != end: + raise ValueError("native STATE and PAGE endpoints must match") + if loading and state.destination_slot is None: + raise ValueError("native-state restore requires a destination SLOT") + if not loading and state.destination_slot is not None: + raise ValueError("native-state save must refer to immutable PAGE units") + unit_ids = self._native_layout.validate_unit_ids(state.unit_ids) + pages = self._block_slice(req, start, end) + if set(pages) & set(unit_ids): + raise ValueError("checkpoint PAGE units must not overlap KV PAGE blocks") + count = (end - start) // self.chunk_size + # Earlier chunks have PAGE only. A boundary image is indivisible: all + # its ordinal groups are present at exactly the same final chunk. + return [pages] + [[-1] * (count - 1) + [unit] for unit in unit_ids] + + def _submit_native(self, req: LMCacheReqMeta, event: Any, *, loading: bool) -> None: + from lmcache.integration.atom import AtomMPTransferSpec + + completion = req.load_operation if loading else req.save_operation + if completion is None: + raise ValueError( + "native-state LMCache MP transfers require exact operation generations" + ) + operation_id = _transfer_operation_id("load" if loading else "save", completion) + pending = self._native_loads if loading else self._native_saves + completed = ( + self._completed_load_operations + if loading + else self._completed_save_operations + ) + with self._lock: + if operation_id in pending or operation_id in completed: + raise RuntimeError( + f"duplicate native-state LMCache MP operation {operation_id!r}" + ) + if not loading and not self._is_kv_writer: + pending[operation_id] = _NativePending( + req, None, immediate_success=True + ) + return + if not loading and len(pending) >= self._max_pending_saves: + # The scheduler has the same bound. Refuse before transport; + # a terminal False safely returns the logical admission credit. + pending[operation_id] = _NativePending(req, None) + return + end = len(req.token_ids) + start = ( + req.load_spec.hbm_cached_tokens + if loading + else req.save_spec.skip_leading_tokens + ) + try: + groups = self._native_block_ids(req, start, end, loading=loading) + spec = AtomMPTransferSpec( + token_ids=list(req.token_ids), + block_ids=groups, + start=start, + end=end, + ) + except Exception: + logger.exception( + "Invalid native-state LMCache MP descriptor %s", operation_id + ) + pending[operation_id] = _NativePending(req, None) + return + entry = _NativePending(req, _UncertainSubmission()) + pending[operation_id] = entry + submit = ( + self._adapter.submit_retrieve_request + if loading + else getattr( + self._adapter, + "submit_store_request_with_chunk_events", + self._adapter.submit_store_request, + ) + ) + try: + future = submit(_mp_session_id(self._config, req.req_id), spec, event) + except Exception: + # Retain the exact source/destination lease. The server might have + # received the request before the connection raised an exception. + logger.exception( + "Native-state LMCache MP submission uncertain; retaining lease %s", + operation_id, + ) + return + with self._lock: + entry.future = future + + def _submit_load(self, req: LMCacheReqMeta, event: Any) -> None: + self._submit_native(req, event, loading=True) + + def _submit_save(self, req: LMCacheReqMeta, event: Any) -> None: + self._submit_native(req, event, loading=False) + + def _begin_restore(self, pending: _NativePending) -> bool: + if not self._restore_descriptor_slots: + return False + state: NativeStateTransfer = pending.request.native_state + spec = self._native_layout.checkpoint_spec + event = torch.cuda.Event() + descriptor_slot = self._restore_descriptor_slots.pop() + pending.restore_event = event + pending.descriptor_slot = descriptor_slot + with torch.cuda.stream(self._restore_stream): + try: + self._native_copy( + (), + ( + CheckpointRestoreOp( + dst_slot=state.destination_slot, + unit_ids=state.unit_ids, + total_bytes=spec.image_bytes, + layout_id=spec.layout_id, + ), + ), + descriptor_slot=descriptor_slot, + ) + pending.restore_succeeded = True + except Exception: + logger.exception("LMCache MP native restore failed") + finally: + event.record(self._restore_stream) + return True + + def _emit_native_source_safe( + self, + output: KVConnectorOutput, + pending: _NativePending, + ranges: tuple[tuple[int, int], ...], + ) -> None: + operation = pending.request.save_operation + if operation is None or not ranges: + return + output.connector_completions |= _source_safe_completions(operation, ranges) + boundary = int(pending.request.native_state.boundary_tokens) + if not pending.state_source_safe and any(end >= boundary for _, end in ranges): + output.connector_completions.add( + ConnectorCompletion( + NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL, + operation, + True, + ) + ) + pending.state_source_safe = True + + def get_finished(self) -> KVConnectorOutput: + output = KVConnectorOutput() + with self._lock: + for operation_id, pending in list(self._native_saves.items()): + take_ranges = getattr(pending.future, "take_completed_ranges", None) + if callable(take_ranges): + try: + self._emit_native_source_safe( + output, pending, tuple(take_ranges()) + ) + except Exception: + logger.warning( + "LMCache MP source-safe event polling failed", + exc_info=True, + ) + if pending.immediate_success: + terminal, result = True, True + else: + terminal, result = _terminal_future_result(pending.future) + if not terminal: + continue + save_spec = pending.request.save_spec + start = 0 if save_spec is None else int(save_spec.skip_leading_tokens) + end = len(pending.request.token_ids) + self._emit_native_source_safe( + output, pending, _chunk_ranges(start, end, self.chunk_size) + ) + output.connector_completions.add( + ConnectorCompletion( + NATIVE_STATE_MP_STORE_CHANNEL, + pending.request.save_operation, + succeeded=result is True, + ) + ) + del self._native_saves[operation_id] + _remember_operation_tombstone( + operation_id, + self._completed_save_operations, + self._completed_save_operation_order, + ) + for operation_id, pending in list(self._native_loads.items()): + if pending.restore_event is None: + terminal, result = _terminal_future_result(pending.future) + if not terminal: + continue + if result is True: + try: + if not self._begin_restore(pending): + continue + except Exception: + logger.exception( + "LMCache MP native restore safety unknown; " + "retaining lease" + ) + pending.future = _UncertainSubmission() + pending.restore_event = None + continue + if pending.restore_event is not None: + try: + if not pending.restore_event.query(): + continue + except Exception: + logger.exception( + "LMCache MP native restore event safety unknown; " + "retaining lease" + ) + continue + completion = pending.request.load_operation + if pending.descriptor_slot is not None: + self._restore_descriptor_slots.append(pending.descriptor_slot) + pending.descriptor_slot = None + target = ( + output.finished_loading + if pending.restore_succeeded + else output.failed_loading + ) + target.add(completion) + del self._native_loads[operation_id] + _remember_operation_tombstone( + operation_id, + self._completed_load_operations, + self._completed_load_operation_order, + ) + return output + + +__all__ = [ + "NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL", + "NATIVE_STATE_MP_STORE_CHANNEL", + "NativeStateLMCacheMPConnector", + "require_native_state_server", +] diff --git a/atom/model_engine/block_manager.py b/atom/model_engine/block_manager.py index caa9504d47..32712c448c 100644 --- a/atom/model_engine/block_manager.py +++ b/atom/model_engine/block_manager.py @@ -2499,7 +2499,11 @@ def register_received_prefix(self, seq: Sequence) -> int: return num_full - start def deallocate_partial( - self, seq: Sequence, protected_block_ids: frozenset[int] + self, + seq: Sequence, + protected_block_ids: frozenset[int], + *, + per_request_state_safe: bool = False, ) -> None: """Deallocate `seq` now, except block IDs a pending offload save reads. @@ -2517,20 +2521,15 @@ def deallocate_partial( request is deallocated normally. This makes ownership explicit: shared prefix blocks retain their other owners, while the save owns exactly one refcount share until ``free_leased_blocks`` releases it. + + Per-request recurrent state is fail-closed: its active slots may be + released here only when the connector explicitly confirms that every + state source needed after request teardown has an independent lease. """ - if seq.has_per_req_cache: - # Per-request recurrent state (GDN/hybrid checkpoints) has release - # ordering this simplified path never replicates (orphan load - # slots, `state_offload.abandon_load`, fork-source pins -- see - # `deallocate` below). Not reachable today: every offload connector - # that can drive early release sets `_permit_per_request_state = - # False` and rejects such a model at `register_kv_caches`. Kept as - # an explicit guard rather than a silent skip, so a future - # hybrid connector that both permits per-request state and enables - # early release fails loudly here instead of freeing a leased PAGE - # or leaking a state slot. + if seq.has_per_req_cache and not per_request_state_safe: raise RuntimeError( - "partial PAGE deallocation is unsupported for per-request state" + "partial PAGE deallocation is unsupported for per-request " + "state without an explicit connector safety capability" ) table = set(seq.block_table) unknown = set(protected_block_ids) - table @@ -2542,6 +2541,57 @@ def deallocate_partial( self.kv.claim(block_id) self.deallocate(seq) + def acquire_offload_prefix( + self, + seq: Sequence, + start_tokens: int, + end_tokens: int, + ) -> tuple[list[int], int, tuple[int, ...]]: + """Claim the still-resident contiguous prefix used by a late save. + + A finished request no longer owns its old physical block table. Each + block is therefore resolved again through the content hash index and + token-checked before it is claimed. The scan stops at the first gap so + the returned range can never describe a cache object with a hole. + """ + hbs = self._hash_block_size() + if ( + start_tokens < 0 + or end_tokens < start_tokens + or start_tokens % hbs + or end_tokens % hbs + ): + raise ValueError( + f"offload source range [{start_tokens}, {end_tokens}) must " + f"align to hash block size {hbs}" + ) + end_block = end_tokens // hbs + start_block = start_tokens // hbs + block_ids = [-1] * end_block + claimed: list[int] = [] + parent_hash = -1 + available_end = start_tokens + for index in range(end_block): + token_ids = self._hash_block_tokens(seq, index) + block_hash = self.compute_hash(token_ids, parent_hash) + parent_hash = block_hash + if index < start_block: + continue + block_id = self.kv.lookup(block_hash) + if block_id < 0: + break + block = self.kv.block(block_id) + if block.hash != block_hash or block.token_ids != token_ids: + break + self.kv.claim(block_id) + claimed.append(block_id) + block_ids[index] = block_id + available_end = (index + 1) * hbs + # The caller may have to trim a partial LMCache chunk from the tail. + # Preserve token order until that trim is complete; a set is suitable + # only after the exact ordered prefix has been selected. + return block_ids, available_end, tuple(claimed) + def free_leased_blocks(self, block_ids) -> None: """Return block IDs a `deallocate_partial` lease held back, to the pool. diff --git a/atom/model_engine/block_pool.py b/atom/model_engine/block_pool.py index b2a2c4a427..e12ff0f7bc 100644 --- a/atom/model_engine/block_pool.py +++ b/atom/model_engine/block_pool.py @@ -365,6 +365,35 @@ def release_units(self, unit_ids: Iterable[int], owner: Hashable) -> None: del self._raw_unit_owner[block_id] self.free(block_id) + def rekey_units( + self, + unit_ids: Iterable[int], + old_owner: Hashable, + new_owner: Hashable, + ) -> None: + """Atomically transfer a raw reservation between logical owners. + + The units remain allocated throughout the hand-off. This is required + when an external-transfer scratch image becomes a reusable checkpoint: + releasing and re-reserving would briefly expose the units through the + free list, allowing another allocation to overwrite bytes that the new + checkpoint index already advertises. + """ + ids = list(unit_ids) + if len(ids) != len(set(ids)): + raise ValueError("a raw-unit ownership transfer contains duplicate ids") + if old_owner is None or new_owner is None: + raise ValueError("raw-unit ownership transfer needs both owners") + for piece_index, block_id in enumerate(ids): + actual = self._raw_unit_owner.get(block_id) + expected = (old_owner, piece_index) + if actual != expected: + raise AssertionError( + f"raw unit {block_id} belongs to {actual!r}, not {expected!r}" + ) + for piece_index, block_id in enumerate(ids): + self._raw_unit_owner[block_id] = (new_owner, piece_index) + # ------------------------------ resizing ------------------------------- # def extend(self, count: int) -> int: """Grow the pool by up to `count` blocks; returns how many it took. diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index 4129209400..7562e163e9 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -2103,9 +2103,21 @@ def allocate_kv_cache(self, num_kvcache_blocks): # place the builder and the connector are both in scope. transfer_tensors.state_backend = self.attn_metadata_builder if hasattr(self, "draft_kv_builder") and transfer_tensors is not None: - draft_regions = self.draft_kv_builder.get_kv_transfer_tensors() - if draft_regions: - transfer_tensors.block_regions.extend(draft_regions) + draft_transfer_tensors = self.draft_kv_builder.get_kv_transfer_tensors() + if draft_transfer_tensors.block_regions: + transfer_tensors.block_regions.extend( + draft_transfer_tensors.block_regions + ) + transfer_tensors.block_tensor_views.extend( + draft_transfer_tensors.block_tensor_views + ) + # Replication is a whole-PAGE declaration. One sharded draft + # region makes the combined target+draft object sharded even + # when the target alone is replicated. + transfer_tensors.tp_replication_factor = math.gcd( + transfer_tensors.tp_replication_factor, + draft_transfer_tensors.tp_replication_factor, + ) if transfer_tensors is not None: # After the draft's regions are in, and here because this is the # only place holding both the complete region list and the diff --git a/atom/model_engine/page_unit_checkpoint.py b/atom/model_engine/page_unit_checkpoint.py index fb11d6ad70..3416d3e647 100644 --- a/atom/model_engine/page_unit_checkpoint.py +++ b/atom/model_engine/page_unit_checkpoint.py @@ -7,7 +7,7 @@ import logging from collections import OrderedDict, deque -from collections.abc import Iterator, Mapping +from collections.abc import Hashable, Iterator, Mapping from dataclasses import dataclass from time import monotonic @@ -116,6 +116,14 @@ class CheckpointRecord: pin_count: int = 0 +@dataclass(frozen=True) +class SuspendedCheckpointRestore: + """A queued restore removed from execution while retaining its source pin.""" + + checkpoint_id: int + operation: CheckpointRestoreOp + + @dataclass class _OffloadPin: """A store that has been dispatched to the worker but not yet indexed. @@ -134,6 +142,9 @@ class _OffloadPin: checkpoint_id: int pinned_at: float source_released: bool = False + # Explicit connector leases may still have DMA reading their PAGE units + # after a timeout. Only a release or terminal completion can end them. + timeout_reclaimable: bool = True class PageUnitCheckpointStore: @@ -382,6 +393,65 @@ def cancel_queued_restore(self, dst_slot: int) -> None: kept.append((checkpoint_id, op)) self._queued_restores = kept + def suspend_queued_restore( + self, dst_slot: int + ) -> SuspendedCheckpointRestore | None: + """Remove one queued restore without releasing its checkpoint pin. + + An external full-snapshot restore uses this to prevent an older local + checkpoint copy from overwriting the destination SLOT later. The + caller must eventually either resume or release the returned lease. + """ + found: SuspendedCheckpointRestore | None = None + kept: list[tuple[int, CheckpointRestoreOp]] = [] + for checkpoint_id, op in self._queued_restores: + if op.dst_slot == dst_slot: + if found is not None: + raise AssertionError( + f"multiple queued checkpoint restores target slot {dst_slot}" + ) + found = SuspendedCheckpointRestore(checkpoint_id, op) + else: + kept.append((checkpoint_id, op)) + self._queued_restores = kept + return found + + def resume_suspended_restore(self, restore: SuspendedCheckpointRestore) -> None: + """Put a suspended restore back on the next maintenance batch.""" + record = self.records.get(restore.checkpoint_id) + if record is None or record.pin_count <= 0: + raise AssertionError("suspended checkpoint restore lost its source pin") + self._queued_restores.append((restore.checkpoint_id, restore.operation)) + + def release_suspended_restore(self, restore: SuspendedCheckpointRestore) -> None: + """Release a suspended restore whose destination was filled elsewhere.""" + self._release_restore_pin(restore.checkpoint_id) + + def adopt_units( + self, unit_ids: tuple[int, ...], owner: Hashable, prefix_hash: int + ) -> bool: + """Publish reserved units as a READY checkpoint without a copy. + + The units change owner in place, so no other allocation can observe + them free in between. When `prefix_hash` is already READY or pending, + the units are released instead and False is returned. The adopted + checkpoint is not nominated for offload: it just arrived from that tier. + """ + if self.contains_or_pending(prefix_hash): + self.pool.release_units(unit_ids, owner) + return False + checkpoint_id = self._new_identity() + self.pool.rekey_units(unit_ids, owner, ("state-checkpoint", checkpoint_id)) + record = CheckpointRecord( + prefix_hash=int(prefix_hash), + unit_ids=tuple(unit_ids), + state=READY, + ) + self.records[checkpoint_id] = record + self.hash_to_checkpoint[record.prefix_hash] = checkpoint_id + self._lru[checkpoint_id] = None + return True + def complete_inflight(self) -> None: stores, self._inflight_stores = self._inflight_stores, [] for checkpoint_id in stores: @@ -491,13 +561,9 @@ def take_offload_stores( # cannot spin this call. deferred.append(checkpoint_id) continue - self._offload_generation += 1 - op = StateStoreOperationId( - int(record.prefix_hash), self._offload_generation + out.append( + self._pin_offload_source(checkpoint_id, timeout_reclaimable=True) ) - record.pin_count += 1 - self._offload_pins[op] = _OffloadPin(checkpoint_id, monotonic()) - out.append((op, record.unit_ids)) # Re-queue the deferred nominations for the next drain, once the # colliding in-flight generation has had a chance to settle. The drain # is `popleft()` (oldest at the left), so `extend` would append these @@ -508,6 +574,43 @@ def take_offload_stores( self._offload_ready.extendleft(reversed(deferred)) return out + def acquire_checkpoint_source( + self, prefix_hash: int, *, max_inflight: int | None = None + ) -> tuple[StateStoreOperationId, tuple[int, ...]] | None: + """Lease the existing READY image for an admitted external transfer. + + The caller must finish scheduler admission and reserve its transfer + byte budget before acquiring this source. A lease pins the image's + existing PAGE units; it neither copies an Active Slot nor queues a + checkpoint. Missing, COPYING, or already dispatched hashes return + None, as does a full optional operation limit. The limit includes + operations whose source was released but whose result is outstanding. + + Release with `release_offload_store_source` once every reader has + stopped, then `settle_offload_store` on terminal completion. Settlement + can also release a source for a transfer known to have stopped or never + started. These leases are never reclaimed by elapsed time: a missing + report cannot prove that another process has stopped its DMA. + """ + if max_inflight is not None and len(self._offload_pins) >= max_inflight: + return None + checkpoint_id = self.lookup(prefix_hash) + if checkpoint_id < 0 or self._hash_in_flight(prefix_hash): + return None + return self._pin_offload_source(checkpoint_id, timeout_reclaimable=False) + + def _pin_offload_source( + self, checkpoint_id: int, *, timeout_reclaimable: bool + ) -> tuple[StateStoreOperationId, tuple[int, ...]]: + record = self.records[checkpoint_id] + self._offload_generation += 1 + op = StateStoreOperationId(int(record.prefix_hash), self._offload_generation) + record.pin_count += 1 + self._offload_pins[op] = _OffloadPin( + checkpoint_id, monotonic(), timeout_reclaimable=timeout_reclaimable + ) + return op, record.unit_ids + def _hash_in_flight(self, prefix_hash: int) -> bool: """Whether some generation of `prefix_hash` is already pinned. @@ -606,12 +709,18 @@ def reclaim_stale_offload_pins(self, timeout_s: float) -> int: `BlockManager.settle_state_store` refuses to index it. The reclaim recovers the memory and forfeits the entry, which is the only pair of outcomes this can honestly offer. + + Explicit `acquire_checkpoint_source` leases are excluded. Their + connector must confirm source release or terminal completion; the + legacy nomination path's timeout policy does not apply to them. """ if timeout_s <= 0 or not self._offload_pins: return 0 cutoff = monotonic() - timeout_s stale = [ - op for op, pin in self._offload_pins.items() if pin.pinned_at <= cutoff + op + for op, pin in self._offload_pins.items() + if pin.timeout_reclaimable and pin.pinned_at <= cutoff ] for op in stale: pin = self._offload_pins.pop(op) @@ -773,6 +882,10 @@ def __init__( # for -- that is one boundary reached twice, not two boundaries. self._pending: dict[tuple[int, int], tuple[Sequence, int]] = {} self._store_ops: list[CheckpointStoreOp] = [] + # External loads write temporary PAGE images before a runner gathers + # them into an Active Slot. Owners include a transfer generation, so a + # late release can never release a newer reservation. + self._transfer_units: dict[Hashable, tuple[int, ...]] = {} self.checkpoints_kept = 0 self.checkpoints_dropped = 0 self.checkpoints_orphaned = 0 @@ -920,6 +1033,59 @@ def has_available_units( def ensure_free_units(self, count: int) -> bool: return self.store.ensure_free_units(count) + def reserve_transfer_units(self, owner: Hashable) -> tuple[int, ...] | None: + """Reserve one raw checkpoint image for an admitted external load. + + `owner` must uniquely identify this transfer attempt, including its + generation. Duplicate active owners and insufficient available units + are refused. The caller reserves its byte budget before this call and + releases these units only after both the transfer and restore stop + using them. No checkpoint is published and no Active Slot is touched. + """ + if owner is None: + raise ValueError("a transfer reservation needs an owner") + if owner in self._transfer_units: + return None + if not self.store.ensure_free_units(self.store.units_per_checkpoint): + return None + units = self.store.pool.reserve_units( + self.store.units_per_checkpoint, ("state-transfer", owner) + ) + if units is None: + return None + self._transfer_units[owner] = tuple(units) + return self._transfer_units[owner] + + def release_transfer_units(self, owner: Hashable) -> None: + """Release a completed or safely cancelled load reservation, once.""" + units = self._transfer_units.pop(owner, None) + if units is not None: + self.store.pool.release_units(units, ("state-transfer", owner)) + + def adopt_transfer_units(self, owner: Hashable, prefix_hash: int) -> bool: + """Publish a completed external image without exposing it to reuse. + + Returns ``True`` when the transfer's units became the canonical READY + checkpoint. If the same hash was published concurrently, the incoming + units are released and ``False`` is returned. The adopted checkpoint is + deliberately not nominated for offload: it just arrived from that tier. + """ + units = self._transfer_units.pop(owner, None) + if units is None: + return False + return self.store.adopt_units(units, ("state-transfer", owner), prefix_hash) + + def suspend_queued_restore( + self, dst_slot: int + ) -> SuspendedCheckpointRestore | None: + return self.store.suspend_queued_restore(dst_slot) + + def resume_suspended_restore(self, restore: SuspendedCheckpointRestore) -> None: + self.store.resume_suspended_restore(restore) + + def release_suspended_restore(self, restore: SuspendedCheckpointRestore) -> None: + self.store.release_suspended_restore(restore) + def contains(self, h: int) -> bool: """Whether HBM holds a READY image for `h` right now. @@ -952,6 +1118,14 @@ def take_offload_stores( """`(operation, unit_ids)` to hand the tier now. See the store.""" return self.store.take_offload_stores(max_inflight) + def acquire_checkpoint_source( + self, prefix_hash: int, *, max_inflight: int | None = None + ) -> tuple[StateStoreOperationId, tuple[int, ...]] | None: + """Lease an admitted transfer's exact READY image. See the store.""" + return self.store.acquire_checkpoint_source( + prefix_hash, max_inflight=max_inflight + ) + def release_offload_store_source(self, op: StateStoreOperationId) -> None: """The gather drained; hand the units back but keep the pin. See store.""" self.store.release_offload_store_source(op) diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 9e1233acb0..4eac92d157 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -20,7 +20,6 @@ from __future__ import annotations import logging -import os import struct import threading import time @@ -85,14 +84,12 @@ def _offload_max_pending_saves() -> int: global _MAX_PENDING_OFFLOAD if _MAX_PENDING_OFFLOAD is not None: return _MAX_PENDING_OFFLOAD - raw = os.environ.get("OFFLOAD_MAX_PENDING_SAVES", "2") try: - value = int(raw) - except ValueError: - logger.warning( - "invalid OFFLOAD_MAX_PENDING_SAVES=%r; using 2 for the state tier", - raw, - ) + value = envs.OFFLOAD_MAX_PENDING_SAVES + except ValueError as exc: + logger.warning("%s; using 2 for the state tier", exc) + value = None + if value is None: value = 2 _MAX_PENDING_OFFLOAD = max(1, value) return _MAX_PENDING_OFFLOAD @@ -624,6 +621,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, @@ -1123,6 +1123,16 @@ def _connector_protected_block_ids(self, seq: Sequence) -> frozenset | None: return None return callback(seq) + def _connector_can_partially_deallocate_state(self, seq: Sequence) -> bool: + """Require an explicit connector guarantee before releasing state. + + A PAGE block lease alone says nothing about a recurrent-state source. + Connectors without this capability therefore remain on the original + whole-request deferred-free path. + """ + callback = getattr(self.kv_connector, "can_partially_deallocate_state", None) + return callable(callback) and callback(seq) is True + def _drain_source_safe_releases(self) -> None: """Free block IDs the offload connector reports as newly source-safe. @@ -2655,19 +2665,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) @@ -3432,14 +3440,22 @@ def postprocess( # it takes is a side effect. if self._connector_should_defer_free(seq): protected = self._connector_protected_block_ids(seq) - if protected is not None: + state_safe = ( + not seq.has_per_req_cache + or self._connector_can_partially_deallocate_state(seq) + ) + if protected is not None and state_safe: # Early block release: only the save's exact source # blocks stay pinned; everything else in this # finished request's block_table -- decode blocks, an # unaligned prompt tail, already-saved ranges -- frees # now instead of waiting behind `deferred_free_blocks`. before = len(seq.block_table) - self.block_manager.deallocate_partial(seq, protected) + # `state_safe` holds on this branch; restate it so the + # block manager's own fail-closed check stays armed. + self.block_manager.deallocate_partial( + seq, protected, per_request_state_safe=True + ) activate = getattr( self.kv_connector, "activate_block_leases", None ) diff --git a/atom/model_ops/attentions/aiter_attention.py b/atom/model_ops/attentions/aiter_attention.py index 04f29691e9..55a1b1be8d 100644 --- a/atom/model_ops/attentions/aiter_attention.py +++ b/atom/model_ops/attentions/aiter_attention.py @@ -780,6 +780,11 @@ def get_kv_transfer_tensors(self): if not self.kv_pools: return None + region_tensors = [ + (geometry, role, tensor) + for geometry, pool in self.kv_pools.items() + for role, tensor in pool.region_tensors() + ] # Every field of every pool, indexer cache included: a region per # (pool, field, layer), in declared order. return KVTransferTensors( @@ -792,10 +797,18 @@ def get_kv_transfer_tensors(self): # per-layer regions are otherwise named alike. semantic_role=f"mha.{geometry}.{role}", ) - for geometry, pool in self.kv_pools.items() - for role, tensor in pool.region_tensors() + for geometry, role, tensor in region_tensors ], slot_regions=[], + # MHA arenas expose one contiguous byte row per scheduler block. + # LMCache MP deliberately treats the last two dimensions as opaque + # copy geometry, so adding a singleton physical-slot dimension is + # a zero-copy view of the exact region declared above. + block_tensor_views=[tensor.unsqueeze(1) for _, _, tensor in region_tensors], + # GQA/MQA KV heads are sharded or only partially replicated across + # TP. In particular MiniMax-M3 must keep one stored shard per rank; + # whole-object TP collapse is unsafe for this PAGE layout. + tp_replication_factor=1, ) def refresh_flydsl_plan(self, context_lens, *, create=False): diff --git a/atom/model_ops/attentions/backends.py b/atom/model_ops/attentions/backends.py index 220977426d..fac3a559a8 100644 --- a/atom/model_ops/attentions/backends.py +++ b/atom/model_ops/attentions/backends.py @@ -342,13 +342,66 @@ def execute_paged_state_copies( self, store_ops: Sequence[CheckpointStoreOp], restore_ops: Sequence[CheckpointRestoreOp], + descriptor_slot: int = 0, ) -> None: - """Copy checkpoints between Active Slots and arbitrary PAGEs.""" + """Copy checkpoints between Active Slots and arbitrary PAGEs. + + Stateful backends use ``descriptor_slot`` to isolate concurrent copy + descriptors. The default implementation only handles the empty no-op + case, so it intentionally leaves that interface argument unused. + """ if store_ops or restore_ops: raise NotImplementedError( f"{type(self).__name__} does not implement PAGE-backed state copy" ) + def _checkpoint_descriptor_device(self) -> torch.device: + """Device of the planes `execute_paged_state_copies` copies between.""" + raise NotImplementedError( + f"{type(self).__name__} does not implement PAGE-backed state copy" + ) + + def _checkpoint_descriptor_buffer(self, descriptor_slot: int = 0) -> CpuGpuBuffer: + """Pinned staging for a step's whole descriptor, sized for the worst step. + + Pinned because the alternative synchronizes: a pageable H2D from + `build()` makes the host wait out the forward already enqueued, which + measured 2.9 ms behind 4 ms of work against 0.1 ms staged. Reused + because allocating pinned memory is itself a synchronizing call. + + A step can carry at most one store and one restore per sequence, so two + per sequence bounds it. The caller checks that bound rather than growing + on demand: a descriptor that did not fit would otherwise be silently + truncated into a copy of the wrong shape. + + The store half of that bound is not a property of the batch -- it is + held by `PagedStateCheckpointCoordinator._supersede`, which keeps one + pending boundary per sequence. A change that let two of a sequence's + boundaries drain together would raise from `build()` here, and would be + storing one of them from the wrong slot besides; the two constraints + have the same owner and move together. + + Slot 0 serves `build()`; each other `descriptor_slot` is a separate + buffer, so an out-of-band restore never shares staging with a step. + """ + if descriptor_slot < 0: + raise ValueError("checkpoint descriptor slot must be non-negative") + descriptors = getattr(self, "_checkpoint_descriptors", None) + if descriptors is None: + descriptors = self._checkpoint_descriptors = {} + descriptor = descriptors.get(descriptor_slot) + if descriptor is None: + plan = self._checkpoint_copy_plan() + max_ops = 2 * int(self.model_runner.config.max_num_seqs) + descriptor = CpuGpuBuffer( + max_ops * plan.num_spans, + 3, + dtype=torch.int64, + device=self._checkpoint_descriptor_device(), + ) + descriptors[descriptor_slot] = descriptor + return descriptor + def warmup_per_req_cache(self) -> None: """Pay whatever the first checkpoint copy would pay, before serving. diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index fb5214d4f2..610c68d3bc 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -730,7 +730,7 @@ def __init__(self, model_runner): self._page_unit_region_owners: tuple[int, ...] = () self._checkpoint_plan_cache: SegmentedCopyPlan | None = None self._checkpoint_slot_base_cache: np.ndarray | None = None - self._checkpoint_descriptor: CpuGpuBuffer | None = None + self._checkpoint_descriptors: dict[int, CpuGpuBuffer] = {} @property def prep_stream(self): @@ -969,6 +969,7 @@ def execute_paged_state_copies( self, store_ops: Sequence[CheckpointStoreOp], restore_ops: Sequence[CheckpointRestoreOp], + descriptor_slot: int = 0, ) -> None: """Copy raw checkpoint bytes between slots and non-contiguous PAGEs. @@ -987,7 +988,7 @@ def execute_paged_state_copies( slot_bases = self._checkpoint_slot_bases() per_op = plan.num_spans total = (len(store_ops) + len(restore_ops)) * per_op - staging = self._checkpoint_descriptor_buffer() + staging = self._checkpoint_descriptor_buffer(descriptor_slot) if total > staging.np.shape[0]: raise RuntimeError( f"a step asked to copy {total // per_op} checkpoints, more " @@ -1159,7 +1160,7 @@ def _invalidate_pool_caches(self) -> None: self._checkpoint_range_cache = None self._checkpoint_plan_cache = None self._checkpoint_slot_base_cache = None - self._checkpoint_descriptor = None + self._checkpoint_descriptors = {} def warmup_per_req_cache(self) -> None: """Run one checkpoint copy now, so the first real one is only a copy. @@ -1190,34 +1191,8 @@ def warmup_per_req_cache(self) -> None: ) launch_copy_descriptor(staging.copy_to_gpu(plan.num_spans), plan) - def _checkpoint_descriptor_buffer(self) -> CpuGpuBuffer: - """Pinned staging for a step's whole descriptor, sized for the worst step. - - Pinned because the alternative synchronizes: a pageable H2D from - `build()` makes the host wait out the forward already enqueued, which - measured 2.9 ms behind 4 ms of work against 0.1 ms staged. Reused - because allocating pinned memory is itself a synchronizing call. - - A step can carry at most one store and one restore per sequence, so - two per Active Slot bounds it -- 1.6 MB at the shipped geometry. The - caller checks that bound rather than growing on demand: a descriptor - that did not fit would otherwise be silently truncated into a copy of - the wrong shape. - - The store half is held by `PagedStateCheckpointCoordinator._supersede`, - which keeps one pending boundary per sequence -- see the longer note on - `GDNStateMixin._checkpoint_descriptor_buffer`. - """ - if self._checkpoint_descriptor is None: - plan = self._checkpoint_copy_plan() - max_ops = 2 * int(self.model_runner.config.max_num_seqs) - self._checkpoint_descriptor = CpuGpuBuffer( - max_ops * plan.num_spans, - 3, - dtype=torch.int64, - device=self._kv_planes()[0].device, - ) - return self._checkpoint_descriptor + def _checkpoint_descriptor_device(self) -> torch.device: + return self._kv_planes()[0].device def _checkpoint_copy_plan(self) -> SegmentedCopyPlan: """Where a slot's checkpoint ranges meet a whole image's PAGE regions. @@ -1946,6 +1921,7 @@ def get_kv_transfer_tensors(self): elem_fp32 = 4 block_regions: list[KVTransferRegion] = [] + block_tensor_sources: list[torch.Tensor] = [] swa_block_regions: list[KVTransferRegion] = [] slot_regions: list[KVTransferRegion] = [] @@ -1971,6 +1947,9 @@ def get_kv_transfer_tensors(self): ) ) for plane, row_bytes, role in planes: + if not plane.is_contiguous(): + raise RuntimeError("a KV plane must be contiguous to be transferred") + block_tensor_sources.append(plane) block_regions.append( KVTransferRegion( plane.data_ptr(), @@ -1991,6 +1970,7 @@ def get_kv_transfer_tensors(self): raise RuntimeError( "a CSA indexer layer must be contiguous to be transferred" ) + block_tensor_sources.append(view) block_regions.append( KVTransferRegion( view.data_ptr(), @@ -2011,6 +1991,30 @@ def get_kv_transfer_tensors(self): f"checkpoint={checkpoint_spec.page_unit_bytes}" ) + # LMCache MP needs one block-major tensor for every PAGE region. Build + # byte views only after validating the declared checkpoint geometry so + # malformed layouts fail with the useful region-size error above. The + # source tensors are not guaranteed to expose rows as dimension zero; + # flattening their byte representation also handles the one-dimensional + # allocation owners used by the transfer-geometry tests. + block_tensor_views: list[torch.Tensor] = [] + for source, region in zip(block_tensor_sources, block_regions, strict=True): + byte_view = source.view(torch.uint8).reshape(-1) + required_bytes = self.num_blocks * region.unit_bytes + if byte_view.numel() < required_bytes: + raise RuntimeError( + f"DSV4 PAGE tensor for {region.semantic_role!r} has " + f"{byte_view.numel()} bytes, fewer than the " + f"{required_bytes} bytes declared by its region" + ) + block_tensor_views.append( + byte_view[:required_bytes].view( + self.num_blocks, + 1, + region.unit_bytes, + ) + ) + # Full per-request SLOT (legacy field name: `swa_block_regions`) is one # slot per plane — compressor state and SWA, every layer at once — and # the id the connector zips over is its pool group. @@ -2060,6 +2064,10 @@ def get_kv_transfer_tensors(self): pool, stride, runner.v4_state_arena, self.pool_geometry ) + # DSv4 PAGE KV and its compact C4/SWA checkpoint image are MLA state + # before any TP-sharded projection, so both are byte-identical on every + # TP rank. One writer is sufficient; all ranks still read. + tp_size = int(getattr(runner.config, "tensor_parallel_size", 1) or 1) return KVTransferTensors( block_regions=block_regions, swa_block_regions=swa_block_regions, @@ -2070,6 +2078,11 @@ def get_kv_transfer_tensors(self): staging_pool_size=pool_size if staging_region else 0, gather_slot=gather_slot, scatter_slot=scatter_slot, + block_tensor_views=block_tensor_views, + tp_replication_factor=tp_size, + native_state_tp_replication_factor=tp_size, + paged_state_checkpoint_spec=checkpoint_spec, + execute_paged_state_copies=self.execute_paged_state_copies, ) # ------------------------------------------------------------------ # diff --git a/atom/model_ops/attentions/gdn_attn.py b/atom/model_ops/attentions/gdn_attn.py index e4c53b52c3..cde3a03876 100644 --- a/atom/model_ops/attentions/gdn_attn.py +++ b/atom/model_ops/attentions/gdn_attn.py @@ -775,35 +775,8 @@ def _checkpoint_copy_plan(self) -> SegmentedCopyPlan: ) return self._checkpoint_plan_cache - def _checkpoint_descriptor_buffer(self) -> CpuGpuBuffer: - """Pinned staging for a step's whole descriptor, sized for the worst step. - - Pinned because the alternative synchronizes: a pageable H2D from - `build()` makes the host wait out the forward already enqueued. Reused - because allocating pinned memory is itself a synchronizing call. - - A step can carry at most one store and one restore per sequence, so two - per sequence bounds it. The caller checks that bound rather than growing - on demand: a descriptor that did not fit would otherwise be silently - truncated into a copy of the wrong shape. - - The store half of that bound is not a property of the batch -- it is - held by `PagedStateCheckpointCoordinator._supersede`, which keeps one - pending boundary per sequence. A change that let two of a sequence's - boundaries drain together would raise from `build()` here, and would be - storing one of them from the wrong slot besides; the two constraints - have the same owner and move together. - """ - if getattr(self, "_checkpoint_descriptor", None) is None: - plan = self._checkpoint_copy_plan() - max_ops = 2 * int(self.model_runner.config.max_num_seqs) - self._checkpoint_descriptor = CpuGpuBuffer( - max_ops * plan.num_spans, - 3, - dtype=torch.int64, - device=self.model_runner.mamba_k_cache.device, - ) - return self._checkpoint_descriptor + def _checkpoint_descriptor_device(self) -> torch.device: + return self.model_runner.mamba_k_cache.device def _validate_paged_state_op(self, op) -> None: """Refuse an op this worker cannot honour, before it addresses memory. @@ -839,7 +812,9 @@ def _validate_paged_state_op(self, op) -> None: if any(unit < 0 or unit >= num_blocks for unit in op.unit_ids): raise RuntimeError("state checkpoint PAGE unit is out of range") - def execute_paged_state_copies(self, store_ops, restore_ops) -> None: + def execute_paged_state_copies( + self, store_ops, restore_ops, descriptor_slot: int = 0 + ) -> None: """Copy raw checkpoint bytes between slots and non-contiguous PAGEs. Every op of either direction goes into one descriptor and one launch. @@ -857,7 +832,7 @@ def execute_paged_state_copies(self, store_ops, restore_ops) -> None: slot_bases = self._checkpoint_slot_bases() per_op = plan.num_spans total = (len(store_ops) + len(restore_ops)) * per_op - staging = self._checkpoint_descriptor_buffer() + staging = self._checkpoint_descriptor_buffer(descriptor_slot) if total > staging.np.shape[0]: raise RuntimeError( f"a step asked to copy {total // per_op} checkpoints, more " diff --git a/atom/plugin/vllm/kv_transfer/seq_view.py b/atom/plugin/vllm/kv_transfer/seq_view.py index 187ab8e3d9..7760b99e97 100644 --- a/atom/plugin/vllm/kv_transfer/seq_view.py +++ b/atom/plugin/vllm/kv_transfer/seq_view.py @@ -20,14 +20,18 @@ class SeqView: """One vLLM request, shaped like an ATOM ``Sequence`` for offload.""" __slots__ = ( + # Stamped by the shared offload scheduler when it refreshes a parked + # deferred save; read back with ``getattr(seq, ..., None)``, so an + # unset slot reads as "not parked". + "_deferred_save_at", "_load_operation", "_num_cached_tokens", # Written by the chunked scheduler's early-block-release path, which - # freezes a finished request's placement so a final save can still be - # dispatched after vLLM has handed the blocks back. Deliberately left - # unset in ``__init__``: that path tests for them with ``hasattr``, and - # an unset slot is absent the same way a missing attribute is. - "_offload_finished_block_ids", + # marks a finished request and freezes its computed prefix so a final + # save can still be dispatched after vLLM has handed the blocks back. + # Deliberately left unset in ``__init__``: an unset slot reads as absent + # the same way a missing attribute does. + "_offload_finished", "_offload_finished_cached_tokens", "_request", "block_table", @@ -113,7 +117,7 @@ def reset_for_preemption(self) -> None: # Frozen placement from a previous finish is placement too, and a # preempted request's is as stale as the live block table. for frozen in ( - "_offload_finished_block_ids", + "_offload_finished", "_offload_finished_cached_tokens", ): if hasattr(self, frozen): diff --git a/atom/spec_decode/draft_kv.py b/atom/spec_decode/draft_kv.py index 96166a91dd..c91b1ba20e 100644 --- a/atom/spec_decode/draft_kv.py +++ b/atom/spec_decode/draft_kv.py @@ -187,17 +187,28 @@ def build_kv_cache_tensor(self, module): v_scale=getattr(module, "v_scale", None), ) - def get_kv_transfer_tensors(self) -> list: - from atom.kv_transfer.disaggregation.types import KVTransferRegion - - return [ - KVTransferRegion( - base_addr=t.data_ptr(), - total_bytes=t.numel() * t.element_size(), - unit_bytes=t.stride(0) * t.element_size(), - # The draft's rows sit in the same block ids as the target's, - # so its regions need a name that says which stack they are. - semantic_role=f"draft.{role}", - ) - for role, t in self.kv_pool.region_tensors() - ] + def get_kv_transfer_tensors(self): + from atom.kv_transfer.disaggregation.types import ( + KVTransferRegion, + KVTransferTensors, + ) + + region_tensors = self.kv_pool.region_tensors() + return KVTransferTensors( + block_regions=[ + KVTransferRegion( + base_addr=t.data_ptr(), + total_bytes=t.numel() * t.element_size(), + unit_bytes=t.stride(0) * t.element_size(), + # The draft's rows sit in the same block ids as the target's, + # so its regions need a name that says which stack they are. + semantic_role=f"draft.{role}", + ) + for role, t in region_tensors + ], + slot_regions=[], + block_tensor_views=[t.unsqueeze(1) for _, t in region_tensors], + # A paged MHA/GQA draft is a per-rank shard. Appending it to an + # otherwise replicated target must disable whole-object collapse. + tp_replication_factor=1, + ) diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 68d98aec82..2af0a54da8 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -45,6 +45,68 @@ def _positive_float_env(name: str, default: str) -> float: return float(default) +def _flag_env(name: str, default: str = "0") -> bool: + # Stripped, and empty reads as off: `VAR=` is how a shell script clears a + # flag inline, and a bare membership test reads the empty string as ON -- + # the opposite of what the operator wrote. `VAR="off "` did the same. + raw = os.getenv(name, default).strip().lower() + return bool(raw) and raw not in ("0", "false", "no", "off") + + +def _optional_int_env(name: str, *, min_value: int | None = None) -> int | None: + """Unset or empty reads as None; anything else must be an integer.""" + raw = os.getenv(name) + if raw is None or raw == "": + return None + try: + value = int(raw) + except ValueError as exc: + raise ValueError(f"{name} must be an integer, got {raw!r}") from exc + if min_value is not None and value < min_value: + raise ValueError(f"{name} must be >= {min_value}, got {value}") + return value + + +def _int_env_or_default(name: str, default: int) -> int: + raw = os.getenv(name) + if raw is None: + return default + try: + return int(raw) + except ValueError: + logger.warning( + "Invalid %s=%r: expected an integer; using default %d", + name, + raw, + default, + ) + return default + + +def _nonnegative_int_env(name: str, default: int) -> int: + value = _int_env_or_default(name, default) + if value < 0: + logger.warning( + "Invalid %s=%r: expected a nonnegative integer; using default %d", + name, + value, + default, + ) + return default + return value + + +def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float: + value = float(os.getenv(name, str(default))) + if not math.isfinite(value): + raise ValueError(f"{name} must be finite") + if allow_zero and value < 0: + raise ValueError(f"{name} must be nonnegative") + if not allow_zero and value <= 0: + raise ValueError(f"{name} must be positive") + return value + + environment_variables: dict[str, Callable[[], Any]] = { # Forward metadata transport: direct or packed. Both keep source checks. # Single-member groups and strided bindings retain direct copies. @@ -818,6 +880,68 @@ def _positive_float_env(name: str, default: str) -> float: "ATOM_ENGRAM_OVERLAP": lambda: os.getenv("ATOM_ENGRAM_OVERLAP", "1") == "1", # Fuse FP32 post-wkv gating and residual addition; 0 selects the torch reference. "ATOM_ENGRAM_FUSED_GATE": lambda: os.getenv("ATOM_ENGRAM_FUSED_GATE", "1") == "1", + # --- KV offload (LMCache) --- + # Save / load executor widths of the offload worker. DSV4 ignores + # OFFLOAD_LOAD_WORKERS: its SLOT load path needs a serial load executor. + "OFFLOAD_COPY_WORKERS": lambda: int(os.getenv("OFFLOAD_COPY_WORKERS", "1") or "1"), + "OFFLOAD_LOAD_WORKERS": lambda: int(os.getenv("OFFLOAD_LOAD_WORKERS", "1") or "1"), + # Running-plus-queued save bound. None means unset: the connector then + # derives max(2, 2 * OFFLOAD_COPY_WORKERS) and the scheduler's state tier + # uses 2. A kv_connector_extra_config "max_pending_saves" takes precedence. + "OFFLOAD_MAX_PENDING_SAVES": lambda: _optional_int_env("OFFLOAD_MAX_PENDING_SAVES"), + # Minimum external-tier hit worth loading, and minimum remaining prefix + # worth a late save, in tokens. + "OFFLOAD_MIN_LOAD_TOKENS": lambda: max( + 0, _int_env_or_default("OFFLOAD_MIN_LOAD_TOKENS", 8192) + ), + "OFFLOAD_MIN_SAVE_TOKENS": lambda: max( + 0, _int_env_or_default("OFFLOAD_MIN_SAVE_TOKENS", 8192) + ), + # Scheduler steps a memoised tier-lookup answer is replayed, and steps a + # failed lookup suppresses the next attempt. + "OFFLOAD_LOOKUP_MEMO_STEPS": lambda: _nonnegative_int_env( + "OFFLOAD_LOOKUP_MEMO_STEPS", 32 + ), + "OFFLOAD_LOOKUP_RETRY_STEPS": lambda: _nonnegative_int_env( + "OFFLOAD_LOOKUP_RETRY_STEPS", 32 + ), + # Per-transfer offload profiling logs. + "OFFLOAD_PROFILE": lambda: _flag_env("OFFLOAD_PROFILE"), + # Experimental: run the staging pack and copy legs on one stream. + "OFFLOAD_SINGLE_STREAM": lambda: _flag_env("OFFLOAD_SINGLE_STREAM"), + # GPU staging buffer size in LMCache chunks, and an upper bound in bytes. + # Unset derives the size from the KV geometry. + "OFFLOAD_GPU_STAGING_CHUNKS": lambda: _optional_int_env( + "OFFLOAD_GPU_STAGING_CHUNKS", min_value=1 + ), + "OFFLOAD_GPU_STAGING_MAX_BYTES": lambda: _optional_int_env( + "OFFLOAD_GPU_STAGING_MAX_BYTES", min_value=1 + ), + # Free the GPU staging buffer after each transfer instead of keeping it. + "OFFLOAD_RELEASE_GPU_STAGING_AFTER_TRANSFER": lambda: _flag_env( + "OFFLOAD_RELEASE_GPU_STAGING_AFTER_TRANSFER" + ), + # DSV4 SLOT sidecar staging rows; None means unset (1). A + # kv_connector_extra_config "slot_sidecar_staging_slots" takes precedence. + "OFFLOAD_SLOT_STAGING_SLOTS": lambda: _optional_int_env( + "OFFLOAD_SLOT_STAGING_SLOTS" + ), + # DSV4 committed SLOT sidecar index capacity; None means unset (65536). A + # kv_connector_extra_config "committed_sidecar_index_capacity" takes + # precedence. + "OFFLOAD_COMMITTED_SIDECAR_CAPACITY": lambda: _optional_int_env( + "OFFLOAD_COMMITTED_SIDECAR_CAPACITY" + ), + # DSV4 wait for a saved SLOT sidecar to become visible, and its poll period. + "OFFLOAD_PUBLICATION_TIMEOUT_S": lambda: _finite_float_env( + "OFFLOAD_PUBLICATION_TIMEOUT_S", 5.0, allow_zero=True + ), + "OFFLOAD_PUBLICATION_POLL_INTERVAL_S": lambda: _finite_float_env( + "OFFLOAD_PUBLICATION_POLL_INTERVAL_S", 0.01, allow_zero=False + ), + # LMCache MP transfer mode: auto, lmcache_driven or engine_driven. A + # kv_connector_extra_config "lmcache.mp.mp_transfer_mode" takes precedence. + "LMCACHE_MP_TRANSFER_MODE": lambda: os.getenv("LMCACHE_MP_TRANSFER_MODE", "auto"), } @@ -858,9 +982,3 @@ def __getattr__(name: str): # env itself -- it asks the connector, via # save_abandon_timeout_s. ATOM does not own the # knob, hence no default of its own here. -# OFFLOAD_MAX_PENDING_SAVES — offload connector queue-depth bound; -# defined/defaulted in -# kv_transfer/offload/_offload_common.py and -# documented in kv_transfer/offload/README.md. -# The state tier shares it (scheduler.py) -# rather than adding a second knob. diff --git a/tests/plugin/test_vllm_offload_seq_view.py b/tests/plugin/test_vllm_offload_seq_view.py index 6d55b09a4e..d917799719 100644 --- a/tests/plugin/test_vllm_offload_seq_view.py +++ b/tests/plugin/test_vllm_offload_seq_view.py @@ -154,12 +154,12 @@ def test_preemption_forgets_the_frozen_placement_too(): """Frozen placement is placement; a preempted request's is equally stale.""" reg = SeqViewRegistry() view = reg.get_or_create(_request()) - view._offload_finished_block_ids = [7, 8, 9] + view._offload_finished = True view._offload_finished_cached_tokens = 384 view.reset_for_preemption() - assert not hasattr(view, "_offload_finished_block_ids") + assert not hasattr(view, "_offload_finished") assert not hasattr(view, "_offload_finished_cached_tokens") diff --git a/tests/test_block_manager.py b/tests/test_block_manager.py index 9b5dbd7d55..984f2c3408 100644 --- a/tests/test_block_manager.py +++ b/tests/test_block_manager.py @@ -3,6 +3,7 @@ import logging +import pytest from conftest import MockConfig from atom.model_engine.block_manager import BlockManager @@ -264,6 +265,35 @@ def test_frees_everything_outside_the_protected_set(self, seq_factory): for block_id in protected: assert bm.kv.block(block_id).ref_count >= 1 + def test_per_request_state_requires_explicit_connector_capability( + self, seq_factory + ): + cfg = MockConfig(num_kvcache_blocks=32, kv_cache_block_size=4) + bm = BlockManager(cfg) + seq = seq_factory(list(range(16))) + bm.allocate(seq) + table = list(seq.block_table) + seq.has_per_req_cache = True + + with pytest.raises(RuntimeError, match="explicit connector safety capability"): + bm.deallocate_partial(seq, frozenset(table[:2])) + + assert list(seq.block_table) == table + assert all(bm.kv.block(block_id).ref_count == 1 for block_id in table) + + def test_explicit_connector_capability_allows_per_request_state(self, seq_factory): + cfg = MockConfig(num_kvcache_blocks=32, kv_cache_block_size=4) + bm = BlockManager(cfg) + seq = seq_factory(list(range(16))) + bm.allocate(seq) + protected = frozenset(seq.block_table[:2]) + seq.has_per_req_cache = True + + bm.deallocate_partial(seq, protected, per_request_state_safe=True) + + assert not seq.block_table + assert bm.kv.num_used == len(protected) + def test_free_leased_blocks_returns_the_rest_to_the_pool(self, seq_factory): cfg = MockConfig(num_kvcache_blocks=32, kv_cache_block_size=4) bm = BlockManager(cfg) diff --git a/tests/test_draft_kv_builder.py b/tests/test_draft_kv_builder.py index b299a862fe..ce445df6cd 100644 --- a/tests/test_draft_kv_builder.py +++ b/tests/test_draft_kv_builder.py @@ -19,8 +19,10 @@ from types import SimpleNamespace import pytest +import torch from torch import nn +from atom.model_ops.attentions.mha_kv_pool import MhaKvPool from atom.spec_decode.draft_kv import DRAFT_KV_ROWS, DraftKvBuilder @@ -162,6 +164,38 @@ def test_releasing_before_the_first_build_builds_nothing(): assert asked == [] +def test_draft_exports_zero_copy_mp_views_and_disables_tp_collapse(): + pool = MhaKvPool( + layers=1, + block_size=128, + num_kv_heads=4, + head_dim=128, + kv_dtype=torch.bfloat16, + ) + pool.allocate(2, "cpu") + builder = DraftKvBuilder.__new__(DraftKvBuilder) + builder._kv_pool = pool + + transfer = builder.get_kv_transfer_tensors() + transfer.set_block_count(2) + + assert transfer.tp_replication_factor == 1 + assert [region.semantic_role for region in transfer.block_regions] == [ + "draft.k.layer_0", + "draft.v.layer_0", + ] + assert [tuple(view.shape[:2]) for view in transfer.block_tensor_views] == [ + (2, 1), + (2, 1), + ] + assert all( + region.base_addr == view.data_ptr() + for region, view in zip( + transfer.block_regions, transfer.block_tensor_views, strict=True + ) + ) + + def test_rows_are_the_pool_layers_the_modules_index(): """The two answers this class is about, taken from one walk: every row index a module gets has to fall inside the pool built for them.""" diff --git a/tests/test_lmcache_mp.py b/tests/test_lmcache_mp.py index 28814e426b..d5f3eb9977 100644 --- a/tests/test_lmcache_mp.py +++ b/tests/test_lmcache_mp.py @@ -12,6 +12,7 @@ import pytest import torch +from atom.kv_transfer.disaggregation.aggregator import KVOutputAggregator from atom.kv_transfer.disaggregation.factory import KVConnectorFactory from atom.kv_transfer.disaggregation.types import ( KVTransferRegion, @@ -37,7 +38,10 @@ def _config( dcp: int = 1, pcp: int = 1, dp: int = 1, + dp_local: int | None = None, + dp_rank: int = 0, enable_dp_attention: bool = False, + enable_expert_parallel: bool = False, role: str = "offload", extra: dict | None = None, kv_lora_rank: int | None = None, @@ -63,8 +67,13 @@ def _config( decode_context_parallel_size=dcp, prefill_context_parallel_size=pcp, enable_dp_attention=enable_dp_attention, + enable_expert_parallel=enable_expert_parallel, speculative_config=None, - parallel_config=SimpleNamespace(data_parallel_size=dp), + parallel_config=SimpleNamespace( + data_parallel_size=dp, + data_parallel_size_local=dp if dp_local is None else dp_local, + data_parallel_rank=dp_rank, + ), kv_transfer_config={ "kv_connector": "lmcache_mp", "kv_role": role, @@ -79,8 +88,7 @@ def _config( ({"pp": 2}, "does not support PP"), ({"dcp": 2}, "does not support DCP"), ({"pcp": 2}, "does not support PCP"), - ({"dp": 2}, "TP-only"), - ({"enable_dp_attention": True}, "TP-only"), + ({"dp": 2, "dp_local": 1}, "only within one host"), ({"tp": 1.5}, "tensor_parallel_size must be an integer"), ], ) @@ -96,6 +104,23 @@ def test_mp_config_accepts_arbitrary_model_type(): ) +@pytest.mark.parametrize( + "kwargs", + [ + {"dp": 2, "dp_rank": 1}, + {"dp": 2, "enable_dp_attention": True}, + { + "tp": 1, + "dp": 8, + "enable_dp_attention": True, + "enable_expert_parallel": True, + }, + ], +) +def test_mp_config_accepts_single_host_dp_and_dpa(kwargs): + assert mp_connector._validate_mp_config(_config(**kwargs))[1] == 1 + + def test_mp_scheduler_is_not_a_dense_transport_scheduler(): from atom.kv_transfer.offload.dense.connector import DenseOffloadScheduler @@ -151,7 +176,9 @@ def __init__(self, **kwargs): adapter_module.AtomMPWorkerAdapter = AtomMPWorkerAdapter monkeypatch.setitem(sys.modules, "lmcache.integration.atom", adapter_module) monkeypatch.setenv("LMCACHE_MP_TRANSFER_MODE", environment_mode) - monkeypatch.setattr(mp_connector, "_model_namespace", lambda _config: "test") + monkeypatch.setattr( + mp_connector, "_model_namespace", lambda _config, **_kwargs: "test" + ) adapter = mp_connector._make_worker_adapter(_config(extra=extra), rank=1) @@ -172,7 +199,7 @@ def __init__(self, **kwargs): ) strategies = [ - mp_connector._parallel_strategy(_config(tp=8), rank) for rank in range(8) + mp_connector._parallel_strategy(_config(tp=8, dp=2), rank) for rank in range(8) ] assert {strategy.world_size for strategy in strategies} == {8} @@ -180,6 +207,31 @@ def __init__(self, **kwargs): assert {strategy.tp_size for strategy in strategies} == {8} +def test_parallel_strategy_uses_one_worker_per_dpa_engine(monkeypatch): + class AtomMPParallelConfig: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + adapter_module = types.ModuleType("lmcache.integration.atom") + adapter_module.AtomMPParallelConfig = AtomMPParallelConfig + monkeypatch.setitem(sys.modules, "lmcache.integration.atom", adapter_module) + + strategy = mp_connector._parallel_strategy( + _config( + tp=1, + dp=8, + dp_rank=5, + enable_dp_attention=True, + enable_expert_parallel=True, + ), + 0, + ) + + assert strategy.world_size == 1 + assert strategy.worker_id == 0 + assert strategy.tp_size == 1 + + def test_parallel_strategy_collapses_fully_replicated_mla(monkeypatch): class AtomMPParallelConfig: def __init__(self, **kwargs): @@ -203,6 +255,21 @@ def __init__(self, **kwargs): assert {strategy.tp_size for strategy in strategies} == {8} +def test_auto_rank_collapse_distinguishes_glm52_mla_from_minimax_m3_gqa(): + glm52 = _config(model_type="glm_moe_dsa", tp=8, kv_lora_rank=512) + minimax = _config(model_type="minimax_m3_vl", tp=8) + minimax.hf_config.architectures = ["MiniMaxM3SparseForConditionalGeneration"] + minimax.hf_config.text_config = SimpleNamespace( + num_key_value_heads=4, + # Stay fail-closed for this model family even if a wrapper grows an + # unrelated field with the same name in the future. + kv_lora_rank=512, + ) + + assert mp_connector._tp_replication_factor(glm52) == 8 + assert mp_connector._tp_replication_factor(minimax) == 1 + + def test_tp_rank_collapse_can_be_disabled_and_rejects_bad_values(): assert ( mp_connector._tp_replication_factor( @@ -248,7 +315,9 @@ def _create_key(self, *_args, **_kwargs): adapter_module.AtomMPParallelConfig = AtomMPParallelConfig adapter_module.AtomMPSchedulerAdapter = AtomMPSchedulerAdapter monkeypatch.setitem(sys.modules, "lmcache.integration.atom", adapter_module) - monkeypatch.setattr(mp_connector, "_model_namespace", lambda _config: "test") + monkeypatch.setattr( + mp_connector, "_model_namespace", lambda _config, **_kwargs: "test" + ) adapter = mp_connector._make_scheduler_adapter( _config(tp=8, kv_lora_rank=kv_lora_rank) @@ -291,16 +360,32 @@ def test_model_namespace_reuses_generic_page_namespace(monkeypatch): config = _config(model_type="ordinary_mha", tp=4) assert mp_connector._model_namespace(config) == ( - "model::atom-page-v2-layout::lmcache-mp-v2" + "model::atom-page-v2-layout::lmcache-mp-v3" ) assert calls == [(config, cfg, 4)] +def test_dp_replicas_share_model_namespace_but_not_request_sessions(monkeypatch): + cfg = object() + monkeypatch.setattr(mp_connector.offcfg, "build_lmcache_config", lambda _kvc: cfg) + monkeypatch.setattr( + mp_connector.offcfg, + "build_page_namespace", + lambda _config, _lmcache_cfg, world: f"page-tp{world}", + ) + rank0 = _config(tp=4, dp=2, dp_rank=0) + rank1 = _config(tp=4, dp=2, dp_rank=1) + + assert mp_connector._model_namespace(rank0) == mp_connector._model_namespace(rank1) + assert mp_connector._mp_session_id(rank0, 0) == "atom-offload-dp0:0" + assert mp_connector._mp_session_id(rank1, 0) == "atom-offload-dp1:0" + + def test_scheduler_validates_role_before_connecting(monkeypatch): config = _config(role="not-a-role") connected = False - def connect(_config): + def connect(_config, **_kwargs): nonlocal connected connected = True raise AssertionError("must not connect") @@ -325,7 +410,7 @@ def shutdown(self): monkeypatch.setattr( mp_connector, "_make_scheduler_adapter", - lambda _config: adapter, + lambda _config, **_kwargs: adapter, ) with pytest.raises(ValueError, match="LMCache chunk size"): @@ -361,7 +446,8 @@ def _transfer_tensors(*, tp_replication_factor: int = 1) -> KVTransferTensors: def test_build_cache_views_groups_opaque_layouts(): - views = mp_connector._build_cache_views(_transfer_tensors(), num_blocks=2) + transfer_tensors = _transfer_tensors() + views = mp_connector._build_cache_views(transfer_tensors, num_blocks=2) assert list(views.tensors) == [ "page.0.primary.0", @@ -371,6 +457,43 @@ def test_build_cache_views_groups_opaque_layouts(): ] assert views.layer_groups == ((0, 1), (2, 3)) assert views.bytes_per_block == 2 * (4 * 32 * 2 + 4 * 16) + assert all(tensor.dtype == torch.uint8 for tensor in views.tensors.values()) + assert tuple(views.tensors["page.0.primary.0"].shape) == (2, 4, 64) + assert tuple(views.tensors["page.2.sidecar.0"].shape) == (2, 4, 16) + assert views.tensors["page.0.primary.0"].data_ptr() == ( + transfer_tensors.block_tensor_views[0].data_ptr() + ) + + +def test_build_cache_views_publishes_float8_pages_as_raw_bytes(): + fp8 = getattr(torch, "float8_e4m3fn", None) + if fp8 is None: + pytest.skip("torch build has no float8_e4m3fn") + + page = torch.empty((2, 4, 32), dtype=fp8) + page.view(torch.uint8).copy_( + torch.arange(page.numel(), dtype=torch.uint8).reshape(page.shape) + ) + region = KVTransferRegion( + base_addr=page.data_ptr(), + total_bytes=page.numel(), + unit_bytes=page[0].numel(), + semantic_role="latent", + ) + transfer = KVTransferTensors( + block_regions=[region], + slot_regions=[], + block_tensor_views=[page], + ) + transfer.set_block_count(2) + + views = mp_connector._build_cache_views(transfer, num_blocks=2) + published = views.tensors["page.0.latent"] + + assert published.dtype == torch.uint8 + assert published.shape == page.shape + assert published.data_ptr() == page.data_ptr() + assert torch.equal(published, page.view(torch.uint8)) def test_build_cache_views_rejects_missing_or_bad_geometry(): @@ -461,6 +584,7 @@ def test_mp_lookup_releases_only_hbm_prefix_after_retrieve_handoff(monkeypatch): adapter = _LookupAdapter([None, 8]) client = mp_connector._MPLookupClient( adapter, + config=_config(), timeout=10.0, poll_interval=0.01, ) @@ -471,7 +595,10 @@ def test_mp_lookup_releases_only_hbm_prefix_after_retrieve_handoff(monkeypatch): # LMCache owns and releases [4, 8) once retrieve is submitted, including # on terminal failure. The scheduler releases only the HBM-resident prefix. + assert adapter.submissions == [("atom-offload-dp0:req", list(range(8)))] assert [(call["start"], call["end"]) for call in adapter.freed] == [(0, 4)] + assert {call["request_id"] for call in adapter.freed} == {"atom-offload-dp0:req"} + assert adapter.cleaned == ["atom-offload-dp0:req", "atom-offload-dp0:req"] assert client.hit_tokens("req") is None @@ -480,6 +607,7 @@ def test_mp_lookup_releases_hit_suffix_outside_retrieve_range(monkeypatch): adapter = _LookupAdapter([12]) client = mp_connector._MPLookupClient( adapter, + config=_config(), timeout=10.0, poll_interval=0.01, ) @@ -524,6 +652,7 @@ def test_mp_lookup_rejects_retrieve_beyond_lookup_hit(monkeypatch): adapter = _LookupAdapter([8]) client = mp_connector._MPLookupClient( adapter, + config=_config(), timeout=10.0, poll_interval=0.01, ) @@ -543,6 +672,7 @@ def test_mp_lookup_timeout_defers_cleanup_until_result(monkeypatch): adapter = _LookupAdapter([None, None]) client = mp_connector._MPLookupClient( adapter, + config=_config(), timeout=1.0, poll_interval=0.01, ) @@ -553,7 +683,7 @@ def test_mp_lookup_timeout_defers_cleanup_until_result(monkeypatch): adapter.results.append(8) client.clear_lookup_status("req") assert [(call["start"], call["end"]) for call in adapter.freed] == [(0, 8)] - assert adapter.cleaned == ["req"] + assert adapter.cleaned == ["atom-offload-dp0:req"] def test_mp_lookup_pending_cleanup_drops_adapter_bookkeeping(monkeypatch): @@ -563,6 +693,7 @@ def test_mp_lookup_pending_cleanup_drops_adapter_bookkeeping(monkeypatch): adapter = _LookupAdapter([None, None]) client = mp_connector._MPLookupClient( adapter, + config=_config(), timeout=1.0, poll_interval=0.01, ) @@ -570,7 +701,7 @@ def test_mp_lookup_pending_cleanup_drops_adapter_bookkeeping(monkeypatch): assert client.lookup(list(range(8)), "req") is None client.clear_lookup_status("req") - assert adapter.cleaned == ["req"] + assert adapter.cleaned == ["atom-offload-dp0:req"] assert client.hit_tokens("req") is None @@ -579,6 +710,7 @@ def _full_prompt_hit_scheduler(monkeypatch, chunk_size): adapter = _LookupAdapter([8]) lookup = mp_connector._MPLookupClient( adapter, + config=_config(), timeout=1.0, poll_interval=0.01, ) @@ -677,7 +809,10 @@ def test_mp_lookup_timeout_is_not_recorded_as_a_tier_miss(monkeypatch): assert scheduler._tier_hit_memo["7"][1] is None assert scheduler.get_num_new_matched_tokens(seq) == (4, True) - assert [request_id for request_id, _tokens in adapter.submissions] == ["7", "7"] + assert [request_id for request_id, _tokens in adapter.submissions] == [ + "atom-offload-dp0:7", + "atom-offload-dp0:7", + ] def test_block_size_one_full_prompt_hit_stays_declined_when_asked_again(monkeypatch): @@ -705,6 +840,7 @@ def test_stale_load_failure_does_not_release_current_generation_locks(): adapter = _LookupAdapter([]) lookup = mp_connector._MPLookupClient( adapter, + config=_config(), timeout=1.0, poll_interval=0.01, ) @@ -764,6 +900,7 @@ def __init__(self, result=True) -> None: self.ready = False self.value = result self.query_error = None + self.completed_ranges = [] def query(self): if self.query_error is not None: @@ -775,6 +912,11 @@ def result(self, timeout=None): raise TimeoutError("future is not ready") return self.value + def take_completed_ranges(self): + ranges = self.completed_ranges + self.completed_ranges = [] + return ranges + @pytest.fixture def fake_lmcache_modules(monkeypatch): @@ -828,6 +970,9 @@ def submit_store_request(self, request_id, op, event): self.saves.append((request_id, op, event, future)) return future + def submit_store_request_with_chunk_events(self, request_id, op, event): + return self.submit_store_request(request_id, op, event) + def shutdown(self): self.shutdown_called = True @@ -879,6 +1024,7 @@ def test_worker_uses_transfer_boundary_and_exact_completion(fake_lmcache_modules ) worker._submit_load(request, object()) + assert adapter.loads[0][0] == "atom-offload-dp0:5" submitted = adapter.loads[0][1] assert submitted.start == 0 assert submitted.end == 8 @@ -1069,6 +1215,7 @@ def test_worker_save_slices_chunk_blocks_and_preserves_operation( ) worker._submit_save(request, object()) + assert adapter.saves[0][0] == "atom-offload-dp0:8" submitted = adapter.saves[0][1] assert submitted.start == 8 assert submitted.end == 16 @@ -1078,6 +1225,33 @@ def test_worker_save_slices_chunk_blocks_and_preserves_operation( assert worker.get_finished().finished_saving == {operation} +def test_worker_reports_chunk_source_safe_before_store_terminal(fake_lmcache_modules): + adapter = _WorkerAdapter() + worker = _worker(adapter) + operation = SaveOperationId(req_id=81, generation=2) + request = LMCacheReqMeta( + req_id=81, + token_ids=list(range(16)), + block_ids=[30, 31, 32, 33], + save_spec=SaveSpec(skip_leading_tokens=0), + save_operation=operation, + ) + + worker._submit_save(request, object()) + future = worker._pending_saves["save:81:2"].future + future.completed_ranges = [(8, 16)] + output = worker.get_finished() + + assert output.finished_saving == set() + assert output.connector_completions == { + mp_connector.ConnectorCompletion( + mp_connector.DENSE_PAGE_SOURCE_SAFE_CHANNEL, + mp_connector.SaveSourceGroupId(operation, ((8, 16),)), + True, + ) + } + + def test_non_writer_completes_save_without_submitting(fake_lmcache_modules): adapter = _WorkerAdapter() worker = _worker(adapter) @@ -1094,11 +1268,53 @@ def test_non_writer_completes_save_without_submitting(fake_lmcache_modules): worker._submit_save(request, object()) assert adapter.saves == [] - assert worker.get_finished().finished_saving == {operation} + output = worker.get_finished() + assert output.finished_saving == {operation} + assert { + completion.operation_id.ranges + for completion in output.connector_completions + if completion.channel == mp_connector.DENSE_PAGE_SOURCE_SAFE_CHANNEL + } == {((0, 8),)} with pytest.raises(RuntimeError, match="duplicate LMCache MP save"): worker._submit_save(request, object()) +def test_collapsed_tp_early_release_waits_only_for_the_single_writer_dma( + fake_lmcache_modules, +): + operation = SaveOperationId(req_id=91, generation=1) + request = LMCacheReqMeta( + req_id=91, + token_ids=list(range(16)), + block_ids=[40, 41, 42, 43], + save_spec=SaveSpec(skip_leading_tokens=0), + save_operation=operation, + ) + writer = _worker(_WorkerAdapter()) + non_writer = _worker(_WorkerAdapter()) + non_writer._is_kv_writer = False + writer._submit_save(request, object()) + non_writer._submit_save(request, object()) + writer_future = writer._pending_saves["save:91:1"].future + writer_future.completed_ranges = [(0, 8)] + + aggregator = KVOutputAggregator(world_size=2) + first = aggregator.aggregate([writer.get_finished(), non_writer.get_finished()]) + assert { + completion.operation_id.ranges + for completion in first.connector_completions + if completion.channel == mp_connector.DENSE_PAGE_SOURCE_SAFE_CHANNEL + } == {((0, 8),)} + + writer_future.completed_ranges = [(8, 16)] + second = aggregator.aggregate([writer.get_finished(), non_writer.get_finished()]) + assert { + completion.operation_id.ranges + for completion in second.connector_completions + if completion.channel == mp_connector.DENSE_PAGE_SOURCE_SAFE_CHANNEL + } == {((8, 16),)} + + def test_worker_tracks_two_load_generations_for_one_raw_request( fake_lmcache_modules, ): diff --git a/tests/test_lmcache_mp_builder_views.py b/tests/test_lmcache_mp_builder_views.py index db847d3265..8c9297e7ea 100644 --- a/tests/test_lmcache_mp_builder_views.py +++ b/tests/test_lmcache_mp_builder_views.py @@ -5,6 +5,7 @@ from __future__ import annotations +import ast import importlib.util import sys import types @@ -16,7 +17,9 @@ import torch from atom.kv_transfer.offload.mp.backend import _build_cache_views +from atom.model_ops.attentions.mha_kv_pool import MhaKvPool from atom.model_ops.attentions.mla_kv_pool import MlaKvPool +from atom.model_ops.attentions.pool_layout.entry_arena import EntryField _MISSING = object() _REPO_ROOT = Path(__file__).parents[1] @@ -79,6 +82,31 @@ def _sub_pool_module(): ) +@pytest.fixture(scope="module") +def mha_export_method(): + # Execute only the shipped exporter: importing the full backend requires a + # GPU AITER build, while this PAGE-view contract is plain torch geometry. + source = _REPO_ROOT / "atom/model_ops/attentions/aiter_attention.py" + module = ast.parse(source.read_text()) + builder = next( + node + for node in module.body + if isinstance(node, ast.ClassDef) + and node.name == "AiterAttentionMetadataBuilder" + ) + method = next( + node + for node in builder.body + if isinstance(node, ast.FunctionDef) and node.name == "get_kv_transfer_tensors" + ) + namespace = {} + exec( # noqa: S102 -- execute the local exporter under the CPU fixture + compile(ast.Module(body=[method], type_ignores=[]), str(source), "exec"), + namespace, + ) + return namespace["get_kv_transfer_tensors"] + + @pytest.fixture(scope="module") def mla_builder_cls(): def noop(*args, **kwargs): @@ -171,6 +199,41 @@ def _assert_region_view_geometry(transfer): assert view.numel() * view.element_size() == region.total_bytes +def test_minimax_m3_builder_publishes_gqa_and_index_views_without_tp_collapse( + mha_export_method, +): + num_blocks = 3 + pool = MhaKvPool( + layers=2, + block_size=128, + num_kv_heads=1, + head_dim=128, + kv_dtype=torch.float8_e4m3fnuz, + extra_fields=(EntryField("index", 1, (128 * 128,), torch.uint8),), + ) + pool.allocate(num_blocks, "cpu") + builder = SimpleNamespace(kv_pools={"gqa4": pool}) + + transfer = mha_export_method(builder) + transfer.set_block_count(num_blocks) + + assert transfer.tp_replication_factor == 1 + assert len(transfer.block_regions) == 4 * 2 + 1 + assert all( + region.semantic_role.startswith("mha.gqa4.") + for region in transfer.block_regions + ) + assert all( + view.shape[:2] == (num_blocks, 1) for view in transfer.block_tensor_views + ) + _assert_region_view_geometry(transfer) + cache_views = _build_cache_views(transfer, num_blocks=num_blocks) + assert len(cache_views.tensors) == len(transfer.block_regions) + + transfer.block_tensor_views[-1][1].fill_(7) + assert torch.all(pool.field_view("index", 0, torch.uint8, (num_blocks, -1))[1] == 7) + + @pytest.mark.parametrize("index_layers,index_rows", [(0, 0), (2, 2), (1, 1)]) def test_mla_builder_publishes_latent_and_index_views( mla_builder_cls, index_layers, index_rows diff --git a/tests/test_lmcache_mp_native_layout.py b/tests/test_lmcache_mp_native_layout.py new file mode 100644 index 0000000000..91e0f34abf --- /dev/null +++ b/tests/test_lmcache_mp_native_layout.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""CPU byte contracts for direct native checkpoint registration.""" + +from __future__ import annotations + +import ast +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from atom.kv_transfer.disaggregation.types import KVTransferRegion, KVTransferTensors +from atom.kv_transfer.offload.mp.native_state_layout import ( + build_native_state_mp_layout, +) +from atom.model_engine.page_unit_checkpoint import PagedStateCheckpointSpec + + +def _transfer(*, widths=(8, 8, 2), image_bytes=39, num_blocks=7): + views = [] + for index, width in enumerate(widths): + # Nonzero allocation offsets exercise the actual shared-arena views. + raw = torch.arange(num_blocks * width + 16, dtype=torch.uint8)[8:-8] + raw.add_(index * 40) + dtype = torch.bfloat16 if index == 0 and width % 2 == 0 else torch.uint8 + views.append(raw.view(dtype).view(num_blocks, 1, -1)) + spec = PagedStateCheckpointSpec( + page_unit_bytes=sum(widths), + slot_bytes=max(image_bytes, 128), + image_bytes=image_bytes, + layout_id="native-test-v1", + ) + transfer = KVTransferTensors( + block_regions=[ + KVTransferRegion( + view.data_ptr(), + num_blocks * width, + width, + semantic_role=f"region.{i}", + ) + for i, (view, width) in enumerate(zip(views, widths, strict=True)) + ], + slot_regions=[], + block_tensor_views=views, + paged_state_checkpoint_spec=spec, + execute_paged_state_copies=lambda stores, restores: None, + ) + transfer.set_block_count(num_blocks) + return transfer + + +def _layout(transfer): + return build_native_state_mp_layout(transfer, block_size=4, chunk_size=16) + + +def _state_alias_targets(layout, transfer, ids): + ids = layout.validate_unit_ids(ids) + page_count = len(transfer.block_tensor_views) + engine_group_by_tensor = { + tensor_index: group.engine_group_id + for group in layout.kernel_groups + for tensor_index in group.tensor_indices + } + return [ + (layout.tensors[tensor_index], ids[engine_group_by_tensor[tensor_index] - 1]) + for tensor_index in range(page_count, len(layout.tensors)) + ] + + +def _gather(layout, transfer, ids): + return torch.cat( + [ + tensor[block_id].flatten() + for tensor, block_id in _state_alias_targets(layout, transfer, ids) + ] + ) + + +def test_native_image_round_trip_uses_arbitrary_unit_ids_and_valid_page_zero(): + source = _transfer() + source_layout = _layout(source) + source_ids = [4, 0, 3] + actual = _gather(source_layout, source, source_ids) + expected = torch.cat( + [ + view[unit_id].view(torch.uint8).flatten() + for unit_id in source_ids + for view in source.block_tensor_views + ] + )[:39] + assert torch.equal(actual, expected) + + destination = _transfer() + for view in destination.block_tensor_views: + view.view(torch.uint8).fill_(0xCD) + destination_layout = _layout(destination) + destination_ids = [1, 5, 2] + offset = 0 + for tensor, block_id in _state_alias_targets( + destination_layout, destination, destination_ids + ): + target = tensor[block_id].flatten() + target.copy_(actual[offset : offset + target.numel()]) + offset += target.numel() + assert offset == actual.numel() + assert torch.equal( + _gather(destination_layout, destination, destination_ids), expected + ) + # The partial final unit owns only the first three bytes of region zero. + assert torch.all( + destination.block_tensor_views[0][2].view(torch.uint8).flatten()[3:] == 0xCD + ) + assert torch.all(destination.block_tensor_views[1][2].view(torch.uint8) == 0xCD) + for view in destination.block_tensor_views: + assert torch.all(view[0].view(torch.uint8) == 0xCD) + assert torch.all(view[6].view(torch.uint8) == 0xCD) + + +def test_layout_coalesces_equal_shapes_inside_ordinal_and_preserves_trim_stride(): + transfer = _transfer() + layout = _layout(transfer) + assert layout.checkpoint_spec.units_per_checkpoint == 3 + assert layout.checkpoint_spec.page_unit_bytes == 18 + assert len(layout.tensors) == 10 + assert len(layout.kernel_groups) == 8 + assert tuple(group.tensor_indices for group in layout.kernel_groups) == ( + (0,), + (1,), + (2,), + (3, 4), + (5,), + (6, 7), + (8,), + (9,), + ) + assert [group.engine_group_id for group in layout.kernel_groups] == [ + 0, + 0, + 0, + 1, + 1, + 2, + 2, + 3, + ] + for group in layout.kernel_groups: + if group.engine_group_id: + assert group.tokens_per_block == group.sw_size_tokens == 16 + assert group.recurrent_state is True + assert group.extra_object_group_tag == 0 + else: + assert group.tokens_per_block == 4 + assert group.sw_size_tokens == -1 + assert group.recurrent_state is False + expected_aliases = [(0, 8), (1, 8), (2, 2), (0, 8), (1, 8), (2, 2), (0, 3)] + tail = layout.tensors[-1] + assert tuple(tail.shape) == (7, 1, 3) + assert tail.stride(0) == 8 + assert not tail.is_contiguous() + for tensor, (region_index, nbytes) in zip( + layout.tensors[len(transfer.block_tensor_views) :], + expected_aliases, + strict=True, + ): + owner = transfer.block_tensor_views[region_index] + assert tensor.shape[-1] == nbytes + assert tensor.untyped_storage().data_ptr() == owner.untyped_storage().data_ptr() + assert tensor.data_ptr() == owner.data_ptr() + assert tensor.storage_offset() == owner.storage_offset() * owner.element_size() + assert tensor.stride(0) == owner.stride(0) * owner.element_size() + + +@pytest.mark.parametrize("image_bytes", [1, 8, 16, 18, 19, 36, 39, 54]) +def test_image_trims_at_region_and_unit_boundaries(image_bytes): + transfer = _transfer(image_bytes=image_bytes) + layout = _layout(transfer) + state_tensors = layout.tensors[len(transfer.block_tensor_views) :] + assert sum(tensor.shape[-1] for tensor in state_tensors) == image_bytes + assert { + group.engine_group_id for group in layout.kernel_groups if group.recurrent_state + } == set(range(1, layout.checkpoint_spec.units_per_checkpoint + 1)) + + +@pytest.mark.parametrize("ids", [[1], [0, -1, 2], [0, 7, 2], [0, 0, 2], [True, 2, 3]]) +def test_unit_validation_rejects_missing_null_out_of_range_and_duplicate_units(ids): + with pytest.raises(ValueError, match="unit ID"): + _layout(_transfer()).validate_unit_ids(ids) + + +def test_equal_shape_unequal_stride_is_rejected_before_registration(): + # Last ordinal contains 4 bytes from each region: the second region has + # stride 8, so LMCache cannot coalesce it with the first region's stride 4. + with pytest.raises(ValueError, match="different block strides"): + _layout(_transfer(widths=(4, 8), image_bytes=20)) + + +def test_registration_rejects_mismatched_native_geometry_and_missing_restore(): + transfer = _transfer() + transfer.execute_paged_state_copies = None + with pytest.raises(TypeError, match="execute_paged_state_copies"): + _layout(transfer) + transfer = _transfer() + transfer.paged_state_checkpoint_spec = replace( + transfer.paged_state_checkpoint_spec, page_unit_bytes=17 + ) + with pytest.raises(ValueError, match="do not cover"): + _layout(transfer) + transfer = _transfer() + transfer.block_tensor_views[0] = transfer.block_tensor_views[0].clone() + with pytest.raises(ValueError, match="does not alias"): + _layout(transfer) + + +@pytest.fixture +def export_method(): + # Execute the shipped exporter without importing attention kernels. Its + # global dependencies are only torch and its local transfer-types import; + # model construction and AITER imports are outside this CPU contract. + source = Path(__file__).parents[1] / "atom/model_ops/attentions/deepseek_v4_attn.py" + module = ast.parse(source.read_text()) + builder = next( + node + for node in module.body + if isinstance(node, ast.ClassDef) + and node.name == "DeepseekV4AttentionMetadataBuilder" + ) + method = next( + node + for node in builder.body + if isinstance(node, ast.FunctionDef) and node.name == "get_kv_transfer_tensors" + ) + namespace = {"torch": torch} + exec( # noqa: S102 -- execute the local exporter under the CPU fixture + compile(ast.Module(body=[method], type_ignores=[]), str(source), "exec"), + namespace, + ) + return namespace["get_kv_transfer_tensors"] + + +@pytest.mark.parametrize("indexer_fp4", [False, True]) +def test_attention_export_retains_page_owners_native_spec_and_restore_callback( + export_method, indexer_fp4 +): + num_blocks, num_slots, envelope_rows = 5, 2, 3 + planes = [ + torch.arange( + (num_blocks * envelope_rows + num_slots * 2) * 4, dtype=torch.uint8 + ).view(-1, 4), + torch.zeros( + num_blocks * envelope_rows + num_slots * 2, 2, dtype=torch.bfloat16 + ), + ] + index_data = torch.zeros(2, num_blocks, 2, 4, dtype=torch.uint8) + index_scale = torch.zeros(2, num_blocks, 2, 1, dtype=torch.uint8) + pools = [(index_data, "dsv4.indexer.data")] + if indexer_fp4: + pools.append((index_scale, "dsv4.indexer.scale")) + page_bytes = 24 + 16 + (4 if indexer_fp4 else 0) + spec = PagedStateCheckpointSpec(page_bytes, 128, "dsv4-paged-state-v3:test", 45) + + def restore(stores, restores): + return None + + geo = SimpleNamespace( + envelope_rows=envelope_rows, + block_bytes=lambda row_bytes: envelope_rows * row_bytes, + slot_bytes=lambda row_bytes: 2 * row_bytes, + physical_slot=lambda slot: slot, + slot_span=lambda slot: (num_blocks * envelope_rows, 0), + ) + builder = SimpleNamespace( + model_runner=SimpleNamespace( + v4_unified_kv=[], + config=SimpleNamespace(kv_transfer_config={"kv_connector": "lmcache_mp"}), + state_runtime=SimpleNamespace(checkpoint_spec=spec), + ), + _indexer_fp4=indexer_fp4, + num_state_slots=num_slots, + num_blocks=num_blocks, + pool_geometry=geo, + _plane_fields=[SimpleNamespace(name="main"), SimpleNamespace(name="rope")], + _kv_planes=lambda: planes, + _plane_row_widths=lambda: [4, 4], + _indexer_page_pools=lambda: pools, + csa_layers=[2, 7], + execute_paged_state_copies=restore, + ) + export_method.__globals__["_uses_pd_staging"] = lambda config: False + export_method.__globals__["_validate_fp4_indexer_transfer"] = lambda config: None + transfer = export_method(builder) + transfer.set_block_count(num_blocks) + assert transfer.paged_state_checkpoint_spec is spec + assert transfer.execute_paged_state_copies is restore + assert len(transfer.block_tensor_views) == 4 + (2 if indexer_fp4 else 0) + for region, view in zip( + transfer.block_regions, transfer.block_tensor_views, strict=True + ): + assert view.data_ptr() == region.base_addr + assert view.shape[:2] == (num_blocks, 1) + assert view.numel() * view.element_size() == region.total_bytes + assert view.stride(0) * view.element_size() == region.unit_bytes + transfer.block_tensor_views[0][4].fill_(0xAB) + assert torch.all(planes[0][12:15] == 0xAB) + assert torch.any(planes[0][15:] != 0xAB) + _layout(transfer) + + +def test_actual_lmcache_registration_preserves_native_aliases_and_stride(): + pytest.importorskip("lmcache.lmcache_native", exc_type=ImportError) + from lmcache.utils import EngineType + from lmcache.v1.gpu_connector.utils import ( + normalize_and_discover_per_layer_formats, + ) + from lmcache.v1.kv_layer_groups import KVLayerGroupsManager + + layout = _layout(_transfer()) + groups = layout.engine_group_infos() + layer_groups = tuple(group.tensor_indices for group in layout.kernel_groups) + normalized, formats = normalize_and_discover_per_layer_formats( + list(layout.tensors), layer_groups, EngineType.ATOM + ) + manager = KVLayerGroupsManager( + normalized, + formats, + groups, + 16, + separate_object_groups=True, + ) + assert [tuple(group.layer_indices) for group in manager.kernel_groups] == list( + layer_groups + ) + assert manager.kernel_groups[-1].shape_desc.block_stride_elems == 8 + assert manager.kernel_groups[-1].shape_desc.hs == 3 + assert len(manager.object_groups) == 2 + assert manager.object_groups[-1].sw_size_chunks == 1 diff --git a/tests/test_lmcache_mp_native_scheduler.py b/tests/test_lmcache_mp_native_scheduler.py new file mode 100644 index 0000000000..09d0348551 --- /dev/null +++ b/tests/test_lmcache_mp_native_scheduler.py @@ -0,0 +1,649 @@ +# SPDX-License-Identifier: MIT + +"""Native-state MP admission and PAGE-image lifetime contracts.""" + +from types import SimpleNamespace + +import pytest +from conftest import MockConfig + +from atom.kv_transfer.disaggregation.types import ( + ConnectorCompletion, + KVConnectorOutput, + LoadOperationId, + SaveOperationId, +) +from atom.kv_transfer.offload.mp import backend +from atom.kv_transfer.offload.mp.connector import LMCacheMPConnectorScheduler +from atom.kv_transfer.offload.mp.native_state_scheduler import ( + NativeStateLMCacheMPConnectorScheduler, +) +from atom.kv_transfer.offload.mp.native_state_worker import ( + NATIVE_STATE_MP_STORE_CHANNEL, +) +from atom.model_engine.block_manager import BlockManager +from atom.model_engine.block_pool import BlockPool +from atom.model_engine.page_unit_checkpoint import ( + PagedStateCheckpointCoordinator, + PagedStateCheckpointSpec, +) +from atom.model_engine.scheduler import ScheduledBatch, ScheduledBatchOutput, Scheduler +from atom.model_engine.sequence import Sequence, SequenceStatus, SequenceType +from atom.model_engine.state_runtime import StateRuntime, StateTransfer + + +@pytest.fixture(autouse=True) +def storage_config(monkeypatch): + # The transport is a double; retain config/chunk validation without needing + # the optional LMCache package or a standalone server on a CPU test runner. + monkeypatch.setattr( + backend.offcfg, + "build_lmcache_config", + lambda kvc: SimpleNamespace( + chunk_size=kvc["kv_connector_extra_config"]["lmcache.chunk_size"] + ), + ) + + +class Adapter: + lmcache_tokens_per_chunk = 8 + + def __init__(self): + self.queries = [] + self.freed = [] + self.ended = [] + self.closed = False + + def maybe_submit_lookup_request(self, request_id, token_ids): + self.queries.append((request_id, list(token_ids))) + + def check_lookup_result(self, request_id): + return len(self.queries[-1][1]) + + def free_lookup_locks(self, **kwargs): + self.freed.append(kwargs) + + def cleanup_lookup_result(self, request_id): + pass + + def end_session(self, request_id): + self.ended.append(request_id) + + def shutdown(self): + self.closed = True + + +def make_scheduler(monkeypatch, *, capacity=2, budget=60, units=30, role="offload"): + monkeypatch.setenv("OFFLOAD_MAX_PENDING_SAVES", str(capacity)) + adapter = Adapter() + connections = [] + + def connect(config, *, checkpoint_spec=None): + connections.append(checkpoint_spec) + return adapter + + monkeypatch.setattr(backend, "_make_scheduler_adapter", connect) + config = SimpleNamespace( + kv_cache_block_size=4, + kv_transfer_config={ + "kv_role": role, + "kv_connector_extra_config": { + "lmcache.mp.max_pinned_state_bytes": budget, + "lmcache.chunk_size": 8, + }, + }, + ) + scheduler = NativeStateLMCacheMPConnectorScheduler(config) + assert connections == [] + checkpoints = PagedStateCheckpointCoordinator( + BlockPool(units), + PagedStateCheckpointSpec(10, 50, "native-test", image_bytes=25), + enabled=True, + ) + manager = SimpleNamespace( + paged_state_checkpoints=checkpoints, + hash_block_size=4, + compute_hash=BlockManager.compute_hash, + ) + scheduler.bind_block_manager(manager) + scheduler.bind_block_manager(manager) + assert connections == [checkpoints.store.spec] + scheduler._min_load_tokens = 0 + return scheduler, checkpoints, adapter + + +def test_native_scheduler_explicitly_guarantees_partial_state_deallocation( + monkeypatch, +): + scheduler, _, _ = make_scheduler(monkeypatch) + seq = Sequence(list(range(24)), 4, has_per_req_cache=True) + + assert scheduler.can_partially_deallocate_state(seq) is True + + +def sequence(request_id=1, *, count=24, computed=16, token_offset=0): + seq = Sequence( + list(range(token_offset, token_offset + count)), + 4, + id=request_id, + has_per_req_cache=True, + ) + seq.num_cached_tokens = computed + seq.state_slots = [5] + seq.block_table = list(range((count + 3) // 4)) + return seq + + +def checkpoint(scheduler, checkpoints, seq, boundary, *, ready=True): + prefix_hash = scheduler._boundary_hash(seq, boundary) + op = checkpoints.store.begin_store(prefix_hash, src_slot=seq.state_slot) + assert op is not None + if ready: + checkpoints.complete_previous_batch() + return prefix_hash, op + + +def terminal(scheduler, operation, *, succeeded=True): + output = KVConnectorOutput( + finished_saving={operation}, + connector_completions={ + ConnectorCompletion(NATIVE_STATE_MP_STORE_CHANNEL, operation, succeeded) + }, + ) + return scheduler.process_completions(output) + + +def test_save_uses_ready_native_hash_and_never_copies_active_slot(monkeypatch): + scheduler, checkpoints, _ = make_scheduler(monkeypatch) + seq = sequence() + assert seq.block_hashes == [] + prefix_hash, image = checkpoint(scheduler, checkpoints, seq, 16, ready=False) + scheduler.update_state_after_alloc(seq) + assert scheduler.build_connector_meta().requests == [] + assert scheduler.should_defer_free(seq) is False + + checkpoints.complete_previous_batch() + seq.state_slots = [99] # The native image is independent of this live slot. + [request] = scheduler.build_connector_meta().requests + assert isinstance(request.save_operation, SaveOperationId) + assert request.native_state.prefix_hash == prefix_hash + assert request.native_state.boundary_tokens == 16 + assert request.native_state.unit_ids == image.unit_ids + assert request.native_state.destination_slot is None + assert request.token_ids == list(seq.token_ids[:16]) + assert request.save_spec.skip_leading_tokens == 0 + assert checkpoints.take_checkpoint_ops() == ((), ()) + assert scheduler._pinned_state_bytes == 30 + checkpoints.clear_index() + assert checkpoints.store.pool.num_free == 27 + + terminal(scheduler, request.save_operation) + assert checkpoints.store.pool.num_free == 30 + assert scheduler._pinned_state_bytes == 0 + + +def test_budget_smaller_than_one_native_image_is_rejected(monkeypatch): + with pytest.raises(ValueError, match="max_pinned_state_bytes"): + make_scheduler(monkeypatch, budget=29) + + +@pytest.mark.parametrize(("capacity", "budget"), [(1, 60), (2, 30)]) +def test_save_credit_and_bytes_are_reserved_before_native_pin( + monkeypatch, capacity, budget +): + scheduler, checkpoints, _ = make_scheduler( + monkeypatch, capacity=capacity, budget=budget + ) + seqs = [sequence(1), sequence(2, token_offset=100)] + ids = [] + for seq in seqs: + prefix_hash, _ = checkpoint(scheduler, checkpoints, seq, 16) + ids.append(checkpoints.store.lookup(prefix_hash)) + scheduler.update_state_after_alloc(seq) + [first] = scheduler.build_connector_meta().requests + assert first.req_id == 1 + assert checkpoints.store.records[ids[0]].pin_count == 1 + assert checkpoints.store.records[ids[1]].pin_count == 0 + assert scheduler._save_tracker["2"][1] == 0 + assert scheduler.build_connector_meta().requests == [] + terminal(scheduler, first.save_operation) + [second] = scheduler.build_connector_meta().requests + assert second.req_id == 2 + assert second.native_state.unit_ids != first.native_state.unit_ids + + +def test_save_frontier_selects_existing_checkpoint_below_computed_tokens(monkeypatch): + scheduler, checkpoints, _ = make_scheduler(monkeypatch) + seq = sequence(computed=24) + checkpoint(scheduler, checkpoints, seq, 8) + scheduler.update_state_after_alloc(seq) + [request] = scheduler.build_connector_meta().requests + assert len(request.token_ids) == 8 + assert request.native_state.boundary_tokens == 8 + assert scheduler._save_tracker["1"][1] == 8 + + +def test_store_failure_rolls_back_for_bounded_retry_and_ignores_stale_reports( + monkeypatch, +): + scheduler, checkpoints, _ = make_scheduler(monkeypatch) + seq = sequence() + checkpoint(scheduler, checkpoints, seq, 16) + scheduler.update_state_after_alloc(seq) + [first] = scheduler.build_connector_meta().requests + terminal(scheduler, first.save_operation, succeeded=False) + assert scheduler._save_tracker["1"][1] == 0 + [second] = scheduler.build_connector_meta().requests + assert second.save_operation != first.save_operation + terminal(scheduler, first.save_operation) + assert scheduler._pinned_state_bytes == 30 + assert scheduler._save_inflight["1"] == second.save_operation + terminal(scheduler, second.save_operation, succeeded=False) + [third] = scheduler.build_connector_meta().requests + terminal(scheduler, third.save_operation, succeeded=False) + assert scheduler.build_connector_meta().requests == [] + assert scheduler.should_defer_free(seq) is False + assert scheduler.get_statistics()["saved_tokens"] == 0 + scheduler.request_finished(seq) + assert scheduler._save_tracker == {} + assert scheduler._save_failures == {} + + +def test_no_timeout_or_abandon_recycles_dispatched_source(monkeypatch): + scheduler, checkpoints, _ = make_scheduler(monkeypatch) + seq = sequence() + checkpoint(scheduler, checkpoints, seq, 16) + scheduler.update_state_after_alloc(seq) + [request] = scheduler.build_connector_meta().requests + scheduler.abandon_save(request.save_operation) + assert scheduler.save_abandon_timeout_s() == 0 + assert scheduler.reclaim_stale_leases(1e-9) == [] + assert checkpoints.reclaim_stale_offload_pins(1e-9) == 0 + assert scheduler._pinned_state_bytes == 30 + assert scheduler.should_defer_free(seq) + terminal(scheduler, request.save_operation) + assert scheduler._pinned_state_bytes == 0 + + +@pytest.mark.parametrize("prompt", [8, 16, 19, 24]) +def test_load_queries_safe_chunk_boundary_before_reserving_state(monkeypatch, prompt): + scheduler, checkpoints, adapter = make_scheduler(monkeypatch) + seq = sequence(count=prompt, computed=0) + expected = ((prompt - 1) // 8) * 8 + assert scheduler.get_num_new_matched_tokens(seq) == (expected, expected > 0) + assert scheduler._pinned_state_bytes == 0 + if expected == 0: + assert adapter.queries == [] + return + assert adapter.queries == [ + (f"atom-offload-dp0:{seq.id}", list(seq.token_ids[:expected])) + ] + scheduler.update_state_after_alloc(seq) + assert scheduler.should_park_for_load_after_alloc(seq) + operation = seq._load_operation + assert isinstance(operation, LoadOperationId) + assert scheduler._pinned_state_bytes == 30 + assert checkpoints.store.pool.num_free == 27 + [request] = scheduler.build_connector_meta().requests + assert request.load_operation == operation + assert request.native_state.boundary_tokens == expected + assert request.native_state.destination_slot == seq.state_slot + assert request.load_spec.lmcache_cached_tokens == expected + assert request.load_spec.transfer_end_tokens is None + assert request.token_ids == list(seq.token_ids[:expected]) + assert ( + scheduler.load_finished(LoadOperationId(seq.id, operation.generation + 1)) + is False + ) + assert scheduler._pinned_state_bytes == 30 + assert scheduler.load_finished(operation) is True + assert scheduler._pinned_state_bytes == 0 + # The completed transfer is now a reusable, unpinned READY checkpoint. + assert checkpoints.store.pool.num_free == 27 + assert checkpoints.contains(request.native_state.prefix_hash) + checkpoint_id = checkpoints.store.lookup(request.native_state.prefix_hash) + assert checkpoints.store.records[checkpoint_id].pin_count == 0 + assert scheduler.load_finished(operation) is False + + +def test_aligned_hbm_hit_uses_incremental_native_restore(monkeypatch): + scheduler, checkpoints, _ = make_scheduler(monkeypatch) + seq = sequence(computed=0) + assert scheduler.get_num_new_matched_tokens(seq) == (16, True) + seq.num_cached_tokens = 8 + scheduler.update_state_after_alloc(seq) + assert scheduler.should_park_for_load_after_alloc(seq) is True + [request] = scheduler.build_connector_meta().requests + assert request.load_spec.hbm_cached_tokens == 8 + assert request.load_spec.lmcache_cached_tokens == 16 + assert request.token_ids == list(seq.token_ids[:16]) + assert seq.offload_load_start_tokens == 8 + assert scheduler._pinned_state_bytes == 30 + assert checkpoints.store.pool.num_free == 27 + + +def test_unaligned_hbm_hit_prefills_to_chunk_then_uses_native_restore(monkeypatch): + scheduler, checkpoints, _ = make_scheduler(monkeypatch) + seq = sequence(computed=0) + assert scheduler.get_num_new_matched_tokens(seq) == (16, True) + seq.num_cached_tokens = 4 + scheduler.update_state_after_alloc(seq) + assert scheduler.should_park_for_load_after_alloc(seq) is False + assert scheduler.build_connector_meta().requests == [] + assert scheduler._handoff_loads == {str(seq.id)} + assert scheduler.adjust_prefill_chunk_after_alloc(seq, 16) == 4 + + # The local prefill reaches the next LMCache chunk boundary. The same + # lookup lease is then handed off to an incremental [8, 16) retrieve. + seq.num_cached_tokens = 8 + assert scheduler.should_park_partial_prefill_for_load(seq) is True + [request] = scheduler.build_connector_meta().requests + assert request.load_spec.hbm_cached_tokens == 8 + assert request.load_spec.lmcache_cached_tokens == 16 + assert seq.offload_load_start_tokens == 8 + assert scheduler._handoff_loads == set() + assert checkpoints.store.pool.num_free == 27 + + +def test_incremental_restore_supersedes_queued_local_state_restore(monkeypatch): + scheduler, checkpoints, _ = make_scheduler(monkeypatch) + seq = sequence(computed=0) + local_hash, _ = checkpoint(scheduler, checkpoints, seq, 8) + assert checkpoints.begin_restore(local_hash, seq.state_slot) + local_id = checkpoints.store.lookup(local_hash) + assert checkpoints.store.records[local_id].pin_count == 1 + + assert scheduler.get_num_new_matched_tokens(seq) == (16, True) + seq.num_cached_tokens = 8 + scheduler.update_state_after_alloc(seq) + assert scheduler.should_park_for_load_after_alloc(seq) + assert not checkpoints.restore_queued_for(seq.state_slot) + [request] = scheduler.build_connector_meta().requests + + assert scheduler.load_finished(request.load_operation) + assert checkpoints.store.records[local_id].pin_count == 0 + assert checkpoints.contains(request.native_state.prefix_hash) + + +def test_incremental_restore_failure_resumes_queued_local_state_restore(monkeypatch): + scheduler, checkpoints, _ = make_scheduler(monkeypatch) + seq = sequence(computed=0) + local_hash, _ = checkpoint(scheduler, checkpoints, seq, 8) + assert checkpoints.begin_restore(local_hash, seq.state_slot) + + scheduler.get_num_new_matched_tokens(seq) + seq.num_cached_tokens = 8 + scheduler.update_state_after_alloc(seq) + assert scheduler.should_park_for_load_after_alloc(seq) + [request] = scheduler.build_connector_meta().requests + assert not checkpoints.restore_queued_for(seq.state_slot) + + assert scheduler.load_failed(request.load_operation) + assert checkpoints.restore_queued_for(seq.state_slot) + [restore] = checkpoints.take_checkpoint_ops()[1] + assert restore.dst_slot == seq.state_slot + + +def test_load_capacity_failure_never_parks_or_claims_missing_state(monkeypatch): + scheduler, checkpoints, _ = make_scheduler(monkeypatch, units=2) + seq = sequence(computed=0) + assert scheduler.get_num_new_matched_tokens(seq) == (16, True) + scheduler.update_state_after_alloc(seq) + assert scheduler.should_park_for_load_after_alloc(seq) is False + assert seq.offload_loaded_tokens == 0 + assert scheduler._native_loads == {} + assert scheduler._pinned_state_bytes == 0 + assert checkpoints.store.pool.num_free == 2 + + +def test_load_failure_and_cancellation_wait_for_exact_terminal_report(monkeypatch): + scheduler, checkpoints, adapter = make_scheduler(monkeypatch) + seq = sequence(computed=0) + scheduler.get_num_new_matched_tokens(seq) + scheduler.update_state_after_alloc(seq) + assert scheduler.should_park_for_load_after_alloc(seq) + [request] = scheduler.build_connector_meta().requests + scheduler.cancel_pending_load(seq) + scheduler.request_finished(seq) + assert scheduler.should_defer_free(seq) + assert scheduler.has_pending_work() + assert scheduler._pinned_state_bytes == 30 + assert adapter.ended == [] + assert scheduler.load_failed(request.load_operation) + assert checkpoints.store.pool.num_free == 30 + assert scheduler._pinned_state_bytes == 0 + assert adapter.ended == [f"atom-offload-dp0:{seq.id}"] + assert not scheduler.should_defer_free(seq) + + +def test_cancel_before_dispatch_releases_reserved_load_units(monkeypatch): + scheduler, checkpoints, _ = make_scheduler(monkeypatch) + seq = sequence(computed=0) + scheduler.get_num_new_matched_tokens(seq) + scheduler.update_state_after_alloc(seq) + assert scheduler.should_park_for_load_after_alloc(seq) + scheduler.cancel_pending_load(seq) + assert scheduler._native_loads == {} + assert scheduler._pinned_state_bytes == 0 + assert checkpoints.store.pool.num_free == 30 + assert scheduler.build_connector_meta().requests == [] + + +def test_load_and_save_share_the_native_byte_budget(monkeypatch): + scheduler, checkpoints, _ = make_scheduler(monkeypatch, budget=30) + saving = sequence(1) + checkpoint(scheduler, checkpoints, saving, 16) + scheduler.update_state_after_alloc(saving) + [request] = scheduler.build_connector_meta().requests + loading = sequence(2, computed=0, token_offset=100) + scheduler.get_num_new_matched_tokens(loading) + scheduler.update_state_after_alloc(loading) + assert scheduler.should_park_for_load_after_alloc(loading) is False + assert scheduler._pinned_state_bytes == 30 + assert scheduler._native_loads == {} + terminal(scheduler, request.save_operation) + assert scheduler._pinned_state_bytes == 0 + + +def engine_scheduler(monkeypatch): + adapter = Adapter() + monkeypatch.setenv("OFFLOAD_MIN_LOAD_TOKENS", "0") + monkeypatch.setenv("OFFLOAD_MAX_PENDING_SAVES", "2") + monkeypatch.setattr( + backend, "_make_scheduler_adapter", lambda _config, **kwargs: adapter + ) + config = MockConfig( + num_kvcache_blocks=40, + enable_prefix_caching=True, + pool_entries={"state": 4}, + state_checkpoint_interval_tokens=0, + state_checkpoint_demand=False, + kv_transfer_config={ + "kv_connector": "lmcache_mp", + "kv_role": "offload", + "kv_connector_extra_config": { + "lmcache.mp.max_pinned_state_bytes": 60, + "lmcache.chunk_size": 8, + }, + }, + ) + connector = LMCacheMPConnectorScheduler(config) + monkeypatch.setattr( + "atom.utils.forward_context.get_kvconnector", lambda _role, _config: connector + ) + spec = PagedStateCheckpointSpec(10, 50, "native-engine-test", image_bytes=25) + engine = Scheduler( + config, + state_runtime=StateRuntime( + transfer=StateTransfer.copy(spec.layout_id), checkpoint_spec=spec + ), + ) + assert connector._block_manager is engine.block_manager + return engine, connector, adapter + + +def start_engine_load(engine): + seq = Sequence(list(range(24)), 4, has_per_req_cache=True) + engine.add(seq) + batch, scheduled = engine.schedule() + assert scheduled == {} + assert seq.status == SequenceStatus.WAITING_FOR_REMOTE_KVS + assert seq.num_cached_tokens == 0 + assert seq.state_slot >= 0 + assert len(seq.block_table) == 6 + [request] = batch.connector_meta_output.requests + assert request.load_operation == seq._load_operation + assert request.native_state.destination_slot == seq.state_slot + assert set(request.native_state.unit_ids).isdisjoint(seq.block_table) + return seq, request + + +def test_engine_allocates_then_parks_and_wakes_at_exact_native_boundary(monkeypatch): + engine, connector, adapter = engine_scheduler(monkeypatch) + seq, request = start_engine_load(engine) + block_table = list(seq.block_table) + state_slot = seq.state_slot + assert engine.block_manager.kv.num_free == 31 + assert connector._pinned_state_bytes == 30 + + batch, scheduled = engine.schedule() + assert scheduled == {} and batch.connector_meta_output.requests == [] + engine._update_from_kv_xfer_finished( + KVConnectorOutput( + finished_loading={ + LoadOperationId(seq.id, request.load_operation.generation + 1) + } + ) + ) + assert seq.status == SequenceStatus.WAITING_FOR_REMOTE_KVS + assert connector._pinned_state_bytes == 30 + + engine._update_from_kv_xfer_finished( + KVConnectorOutput(finished_loading={request.load_operation}) + ) + assert connector._pinned_state_bytes == 0 + # The three transfer units remain as an unpinned READY checkpoint. + assert engine.block_manager.kv.num_free == 31 + assert engine.block_manager.paged_state_checkpoints.contains( + request.native_state.prefix_hash + ) + batch, scheduled = engine.schedule() + assert scheduled[seq.id] is seq + assert seq.num_cached_tokens == 16 + assert list(batch.num_cached_tokens) == [16] + assert seq.offload_promoted_tokens == 16 + assert list(seq.block_table) == block_table + assert seq.state_slot == state_slot + assert engine._num_parked_remote_kv == 0 + assert len(adapter.queries) == 1 + + +@pytest.mark.parametrize("succeeded", [True, False]) +def test_engine_aborted_load_keeps_page_units_and_slot_until_terminal( + monkeypatch, succeeded +): + engine, connector, _ = engine_scheduler(monkeypatch) + seq, request = start_engine_load(engine) + block_table = list(seq.block_table) + state_slot = seq.state_slot + seq.status = SequenceStatus.ABORTED + engine.schedule() + assert seq.status == SequenceStatus.FINISHED + assert engine.deferred_free_blocks[seq.id] is seq + assert list(seq.block_table) == block_table + assert seq.state_slot == state_slot + assert engine.block_manager.kv.num_free == 31 + assert connector._pinned_state_bytes == 30 + + kwargs = { + "finished_loading" if succeeded else "failed_loading": {request.load_operation} + } + engine._update_from_kv_xfer_finished(KVConnectorOutput(**kwargs)) + assert connector._pinned_state_bytes == 0 + assert seq.id not in engine.deferred_free_blocks + assert list(seq.block_table) == [] + assert seq.state_slot == -1 + expected_free = 37 if succeeded else 40 + assert engine.block_manager.kv.num_free == expected_free + assert ( + engine.block_manager.paged_state_checkpoints.contains( + request.native_state.prefix_hash + ) + is succeeded + ) + assert engine._num_parked_remote_kv == 0 + + +def test_engine_failed_native_load_recomputes_from_zero_without_reallocation( + monkeypatch, +): + engine, connector, adapter = engine_scheduler(monkeypatch) + seq, request = start_engine_load(engine) + state_slot = seq.state_slot + block_table = list(seq.block_table) + engine._update_from_kv_xfer_finished( + KVConnectorOutput(failed_loading={request.load_operation}) + ) + assert connector._pinned_state_bytes == 0 + assert connector._save_tracker[str(seq.id)][1] == 0 + batch, scheduled = engine.schedule() + assert scheduled[seq.id] is seq + assert seq.num_cached_tokens == 0 + assert list(batch.num_cached_tokens) == [0] + assert seq.offload_load_failed is True + assert list(seq.block_table) == block_table + assert seq.state_slot == state_slot + assert engine._num_parked_remote_kv == 0 + assert len(adapter.queries) == 1 + + +def test_engine_releases_retired_request_when_unpinned_candidate_is_evicted( + monkeypatch, +): + engine, connector, _ = engine_scheduler(monkeypatch) + bm = engine.block_manager + seq = Sequence(list(range(24)), 4, has_per_req_cache=True) + assert bm.allocate(seq, bm.can_allocate(seq)) + bm.hash_blocks(seq, 16) + seq.num_cached_tokens = 16 + checkpoints = bm.paged_state_checkpoints + prefix_hash, _ = checkpoint(connector, checkpoints, seq, 16) + connector.update_state_after_alloc(seq) + seq.type = SequenceType.PREFILL + seq.status = SequenceStatus.RUNNING + engine.running.append(seq) + batch = ScheduledBatch( + {seq.id: seq}, + [8], + 8, + total_tokens_num_prefill=8, + total_seqs_num=1, + total_seqs_num_prefill=1, + num_cached_tokens=[16], + is_final_chunk=[True], + ) + engine.postprocess( + [seq], + ScheduledBatchOutput([seq.id], [(engine.eos_token_id,)], None, None, None), + batch=batch, + ) + assert seq.status == SequenceStatus.FINISHED + # A not-yet-admitted candidate owns no PAGE lease. The request's KV blocks + # and active SLOT are therefore released immediately. + assert seq.id not in engine.deferred_free_blocks + assert seq.state_slot == -1 + assert list(seq.block_table) == [] + assert connector._native_saves == {} + checkpoint_id = checkpoints.store.lookup(prefix_hash) + assert checkpoints.store.records[checkpoint_id].pin_count == 0 + + # The next admission may evict this waiting candidate. It has never been + # sent, so there will be no GPU completion naming this finished request. + checkpoints.unindex(prefix_hash) + assert connector.has_pending_work() + engine._update_from_kv_xfer_finished(KVConnectorOutput()) + assert seq.id not in engine.deferred_free_blocks + assert bm.kv.num_free == 40 + assert connector._retired_requests == {} + assert connector._save_tracker == {} + assert not connector.has_pending_work() diff --git a/tests/test_lmcache_mp_native_worker.py b/tests/test_lmcache_mp_native_worker.py new file mode 100644 index 0000000000..aec40e436a --- /dev/null +++ b/tests/test_lmcache_mp_native_worker.py @@ -0,0 +1,427 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Worker contracts for reusable native-state LMCache MP transfers.""" + +import sys +import types +from dataclasses import dataclass, replace +from types import SimpleNamespace + +import pytest +import torch + +from atom.kv_transfer.disaggregation.types import ( + ConnectorCompletion, + KVTransferRegion, + KVTransferTensors, + LoadOperationId, + SaveOperationId, + SaveSourceGroupId, +) +from atom.kv_transfer.offload.chunked_scheduler import DENSE_PAGE_SOURCE_SAFE_CHANNEL +from atom.kv_transfer.offload.metadata import ( + LMCacheReqMeta, + LoadSpec, + NativeStateTransfer, + SaveSpec, +) +from atom.kv_transfer.offload.mp.backend import _model_namespace, _tp_replication_factor +from atom.kv_transfer.offload.mp.native_state_layout import ( + build_native_state_mp_layout, +) +from atom.kv_transfer.offload.mp.native_state_worker import ( + NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL, + NATIVE_STATE_MP_STORE_CHANNEL, + NativeStateLMCacheMPConnector, + require_native_state_server, +) +from atom.model_engine.page_unit_checkpoint import PagedStateCheckpointSpec + + +@dataclass +class _TransferSpec: + token_ids: list[int] + block_ids: list[list[int]] + start: int = 0 + end: int = 0 + + +@pytest.fixture(autouse=True) +def fake_lmcache_transfer_spec(monkeypatch): + """Keep transport-contract tests independent of optional LMCache installs.""" + lmcache = types.ModuleType("lmcache") + integration = types.ModuleType("lmcache.integration") + atom = types.ModuleType("lmcache.integration.atom") + atom.AtomMPTransferSpec = _TransferSpec + lmcache.integration = integration + integration.atom = atom + monkeypatch.setitem(sys.modules, "lmcache", lmcache) + monkeypatch.setitem(sys.modules, "lmcache.integration", integration) + monkeypatch.setitem(sys.modules, "lmcache.integration.atom", atom) + + +class Future: + def __init__(self, value=True, ready=False, source_ranges=()): + self.value, self.ready = value, ready + self.source_ranges = list(source_ranges) + + def query(self): + return self.ready + + def result(self, timeout=0): + return self.value + + def take_completed_ranges(self): + ranges, self.source_ranges = self.source_ranges, [] + return tuple(ranges) + + +def config(**extra): + return SimpleNamespace( + kv_cache_block_size=4, + tensor_parallel_size=2, + hf_config=SimpleNamespace(model_type="reusable_test_model", kv_lora_rank=512), + kv_transfer_config={ + "kv_connector": "lmcache_mp", + "kv_role": "offload", + "kv_connector_extra_config": extra, + }, + ) + + +@pytest.fixture +def worker(): + instance = NativeStateLMCacheMPConnector(config()) + page = torch.zeros((32, 1, 32), dtype=torch.uint8) + spec = PagedStateCheckpointSpec(32, 128, "native-test-v1", 80) + tensors = KVTransferTensors( + block_regions=[ + KVTransferRegion( + base_addr=page.data_ptr(), unit_bytes=32, total_bytes=page.numel() + ) + ], + slot_regions=[], + block_tensor_views=[page], + paged_state_checkpoint_spec=spec, + execute_paged_state_copies=lambda *_: None, + ) + tensors.set_block_count(32) + instance._native_layout = build_native_state_mp_layout( + tensors, block_size=4, chunk_size=8 + ) + instance.chunk_size = 8 + instance.submitted = [] + instance.submitted_request_ids = [] + instance.future = Future() + + def submit(request_id, op, event): + instance.submitted_request_ids.append(request_id) + instance.submitted.append(op) + return instance.future + + instance._adapter = SimpleNamespace( + submit_store_request=submit, submit_retrieve_request=submit + ) + return instance + + +def request(*, loading=False, units=(0, 25, 31), generation=1, hbm=0): + return LMCacheReqMeta( + req_id=7, + token_ids=list(range(16)), + block_ids=[1, 2, 3, 4], + load_spec=LoadSpec(hbm, 16) if loading else None, + save_spec=None if loading else SaveSpec(0), + save_operation=None if loading else SaveOperationId(7, generation), + load_operation=LoadOperationId(7, generation) if loading else None, + native_state=NativeStateTransfer(units, 16, 998, 2 if loading else None), + ) + + +def test_store_transmits_page_zero_as_real_native_unit(worker): + req = request() + worker._submit_save(req, object()) + assert worker.submitted_request_ids == ["atom-offload-dp0:7"] + assert worker.submitted[0].block_ids == [ + [1, 2, 3, 4], + [-1, 0], + [-1, 25], + [-1, 31], + ] + assert not worker.get_finished().connector_completions + worker.future.ready = True + finished = worker.get_finished() + assert {completion.channel for completion in finished.connector_completions} == { + DENSE_PAGE_SOURCE_SAFE_CHANNEL, + NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL, + NATIVE_STATE_MP_STORE_CHANNEL, + } + terminals = [ + completion + for completion in finished.connector_completions + if completion.channel == NATIVE_STATE_MP_STORE_CHANNEL + ] + assert terminals == [ + ConnectorCompletion(NATIVE_STATE_MP_STORE_CHANNEL, req.save_operation, True) + ] + page_ranges = { + completion.operation_id.ranges + for completion in finished.connector_completions + if completion.channel == DENSE_PAGE_SOURCE_SAFE_CHANNEL + } + assert page_ranges == {((0, 8),), ((8, 16),)} + assert not finished.finished_saving # one quorum channel for the entire pair + + +def test_native_state_groups_share_one_presence_mask(worker): + req = replace( + request(), + token_ids=list(range(32)), + block_ids=list(range(1, 9)), + native_state=NativeStateTransfer((0, 25, 31), 32, 998, None), + ) + + block_groups = worker._native_block_ids(req, 0, 32, loading=False) + state_presence = [ + [block_id != -1 for block_id in group] for group in block_groups[1:] + ] + + assert state_presence == [[False, False, False, True]] * 3 + + +@pytest.mark.parametrize( + "units", [(0, -1, 31), (0, 31), (0, 0, 31), (0, 25, 32), (0, 2, 31)] +) +def test_invalid_native_source_fails_before_transport(worker, units): + worker._submit_save(request(units=units), object()) + assert not worker.submitted + completions = worker.get_finished().connector_completions + [terminal] = [ + completion + for completion in completions + if completion.channel == NATIVE_STATE_MP_STORE_CHANNEL + ] + assert not terminal.succeeded + + +def test_uncertain_remote_submission_retains_lease(worker): + def uncertain(*_): + raise ConnectionError("server may have received request") + + worker._adapter.submit_store_request = uncertain + worker._submit_save(request(), object()) + assert not worker.get_finished().connector_completions + assert len(worker._native_saves) == 1 + + +def test_failed_retrieve_does_not_restore_or_report_success(worker, monkeypatch): + monkeypatch.setattr( + worker, "_begin_restore", lambda _: pytest.fail("unexpected restore") + ) + worker.future.value, worker.future.ready = False, True + req = request(loading=True) + worker._submit_load(req, object()) + finished = worker.get_finished() + assert finished.failed_loading == {req.load_operation} + assert not finished.finished_loading + + +def test_load_completion_waits_for_native_restore(worker, monkeypatch): + restore_event = Future(ready=False) + restored = [] + + def restore(pending): + restored.append(pending.request.native_state.destination_slot) + pending.restore_event = restore_event + pending.restore_succeeded = True + + monkeypatch.setattr(worker, "_begin_restore", restore) + req = request(loading=True) + worker._submit_load(req, object()) + assert not worker.get_finished().finished_loading + assert restored == [] + worker.future.ready = True + assert not worker.get_finished().finished_loading + assert restored == [2] + restore_event.ready = True + assert worker.get_finished().finished_loading == {req.load_operation} + assert restored == [2] + + +def test_query_exception_never_releases_dma_source(worker): + def broken(): + raise RuntimeError("IPC event unavailable") + + worker.future.query = broken + worker._submit_save(request(), object()) + assert not worker.get_finished().connector_completions + assert worker._native_saves + + +def test_worker_admission_bound_rejects_before_dma(worker): + worker._max_pending_saves = 1 + worker._submit_save(request(), object()) + worker._submit_save(request(generation=2), object()) + assert len(worker.submitted) == 1 + completions = worker.get_finished().connector_completions + [terminal] = [ + completion + for completion in completions + if completion.channel == NATIVE_STATE_MP_STORE_CHANNEL + ] + assert terminal.operation_id == SaveOperationId(7, 2) + assert not terminal.succeeded + + +def test_exact_completed_generation_cannot_replay(worker): + worker.future.ready = True + worker._submit_save(request(), object()) + worker.get_finished() + with pytest.raises(RuntimeError, match="duplicate"): + worker._submit_save(request(), object()) + + +def test_native_state_uses_same_tp_rank_collapse_as_page(): + assert _tp_replication_factor(config()) == 2 + assert _tp_replication_factor(config(**{"lmcache.mp.tp_rank_collapse": True})) == 2 + + +def test_incremental_load_transfers_only_page_suffix_but_full_native_image(worker): + req = request(loading=True, hbm=8) + worker._submit_load(req, object()) + [submitted] = worker.submitted + assert submitted.start == 8 + assert submitted.end == 16 + assert submitted.block_ids == [[3, 4], [0], [25], [31]] + + +def test_source_safe_ranges_are_reported_before_terminal(worker): + req = request() + worker.future.source_ranges = [(0, 8)] + worker._submit_save(req, object()) + + first = worker.get_finished() + assert first.connector_completions == { + ConnectorCompletion( + DENSE_PAGE_SOURCE_SAFE_CHANNEL, + SaveSourceGroupId(req.save_operation, ((0, 8),)), + True, + ) + } + assert worker._native_saves + + worker.future.source_ranges = [(8, 16)] + second = worker.get_finished() + assert ( + ConnectorCompletion( + DENSE_PAGE_SOURCE_SAFE_CHANNEL, + SaveSourceGroupId(req.save_operation, ((8, 16),)), + True, + ) + in second.connector_completions + ) + assert ( + ConnectorCompletion( + NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL, req.save_operation, True + ) + in second.connector_completions + ) + assert all( + completion.channel != NATIVE_STATE_MP_STORE_CHANNEL + for completion in second.connector_completions + ) + + +def test_collapsed_tp_non_writer_skips_transport_and_reports_safe_success(worker): + worker._is_kv_writer = False + req = request() + worker._submit_save(req, object()) + assert worker.submitted == [] + + completions = worker.get_finished().connector_completions + assert ( + ConnectorCompletion( + NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL, req.save_operation, True + ) + in completions + ) + assert ( + ConnectorCompletion(NATIVE_STATE_MP_STORE_CHANNEL, req.save_operation, True) + in completions + ) + assert { + completion.operation_id.ranges + for completion in completions + if completion.channel == DENSE_PAGE_SOURCE_SAFE_CHANNEL + } == {((0, 8),), ((8, 16),)} + + +def test_restore_uses_independent_descriptor_slot_without_stream_sync( + worker, monkeypatch +): + from atom.kv_transfer.offload.mp import native_state_worker + + class Event: + def __init__(self): + self.recorded_on = None + + def record(self, stream): + self.recorded_on = stream + + class StreamContext: + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + restore_stream = object() + event = Event() + copied = [] + worker._restore_stream = restore_stream + worker._restore_descriptor_slots = [3] + worker._native_copy = lambda stores, restores, descriptor_slot=0: copied.append( + (stores, restores, descriptor_slot) + ) + monkeypatch.setattr(native_state_worker.torch.cuda, "Event", lambda: event) + monkeypatch.setattr( + native_state_worker.torch.cuda, "stream", lambda stream: StreamContext() + ) + + req = request(loading=True) + worker._submit_load(req, object()) + pending = next(iter(worker._native_loads.values())) + assert worker._begin_restore(pending) + assert pending.descriptor_slot == 3 + assert worker._restore_descriptor_slots == [] + assert copied[0][2] == 3 + assert event.recorded_on is restore_stream + + +def test_native_server_chunk_mismatch_fails_before_registration(monkeypatch): + from atom.kv_transfer.offload.mp import native_state_worker + + monkeypatch.setattr( + native_state_worker.offcfg, + "build_lmcache_config", + lambda _: SimpleNamespace(chunk_size=256), + ) + adapter = SimpleNamespace(lmcache_tokens_per_chunk=512) + with pytest.raises(ValueError, match="must match"): + require_native_state_server(adapter, config()) + + +def test_native_namespace_changes_with_image_codec(monkeypatch): + from atom.kv_transfer.offload.mp import backend + + monkeypatch.setattr(backend.offcfg, "build_lmcache_config", lambda _: object()) + monkeypatch.setattr(backend.offcfg, "lmcache_replica_world_size", lambda _: 2) + monkeypatch.setattr( + backend.offcfg, "build_page_namespace", lambda *_: "page-config" + ) + first = PagedStateCheckpointSpec(32, 128, "native-test-v1", 80) + second = replace(first, image_bytes=81) + assert _model_namespace(config(), checkpoint_spec=first) != _model_namespace( + config(), checkpoint_spec=second + ) diff --git a/tests/test_lmcache_mp_shell.py b/tests/test_lmcache_mp_shell.py index a3fb8e23e3..35905fe27f 100644 --- a/tests/test_lmcache_mp_shell.py +++ b/tests/test_lmcache_mp_shell.py @@ -1,12 +1,16 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Public contracts for the layout-neutral LMCache MP connector.""" +"""Public contracts for the capability-selected LMCache MP connector.""" from types import SimpleNamespace from atom.kv_transfer.disaggregation.factory import KVConnectorFactory -from atom.kv_transfer.offload.mp import backend +from atom.kv_transfer.offload.mp import ( + backend, + native_state_scheduler, + native_state_worker, +) from atom.kv_transfer.offload.mp.connector import ( LMCacheMPConnector, LMCacheMPConnectorScheduler, @@ -22,9 +26,9 @@ def _config(): ) -def test_public_connector_exports_generic_implementation(): - assert LMCacheMPConnector is backend.LMCacheMPConnector - assert LMCacheMPConnectorScheduler is backend.LMCacheMPConnectorScheduler +def test_public_connector_is_a_layout_neutral_shell(): + assert LMCacheMPConnector is not backend.LMCacheMPConnector + assert LMCacheMPConnectorScheduler is not backend.LMCacheMPConnectorScheduler def test_factory_registration_resolves_public_connectors(monkeypatch): @@ -56,3 +60,149 @@ def test_factory_registration_resolves_public_connectors(monkeypatch): assert worker.config is config assert scheduler.config is config assert KVConnectorFactory.canonical_name("LMCacheMPConnector") == "lmcache_mp" + + +def test_worker_selects_native_state_from_published_contract(monkeypatch): + selected = [] + + class NativeWorker: + def __init__(self, config): + selected.append(("native", config)) + + def register_kv_caches(self, caches, tensors, num_blocks): + self.registered = (caches, tensors, num_blocks) + + class PageWorker(NativeWorker): + def __init__(self, config): + selected.append(("page", config)) + + monkeypatch.setattr( + native_state_worker, "NativeStateLMCacheMPConnector", NativeWorker + ) + monkeypatch.setattr(backend, "LMCacheMPConnector", PageWorker) + config = _config() + worker = LMCacheMPConnector(config) + tensors = SimpleNamespace( + paged_state_checkpoint_spec=object(), + execute_paged_state_copies=lambda *_: None, + ) + worker.register_kv_caches({}, tensors, 7) + assert selected == [("native", config)] + assert worker._impl.registered == ({}, tensors, 7) + + +def test_scheduler_selects_native_state_from_block_manager(monkeypatch): + selected = [] + + class NativeScheduler: + def __init__(self, config): + selected.append(("native", config)) + + def bind_block_manager(self, manager): + self.manager = manager + + class PageScheduler: + def __init__(self, config): + selected.append(("page", config)) + + monkeypatch.setattr( + native_state_scheduler, + "NativeStateLMCacheMPConnectorScheduler", + NativeScheduler, + ) + monkeypatch.setattr(backend, "LMCacheMPConnectorScheduler", PageScheduler) + config = _config() + scheduler = LMCacheMPConnectorScheduler(config) + manager = SimpleNamespace(paged_state_checkpoints=object()) + scheduler.bind_block_manager(manager) + assert selected == [("native", config)] + assert scheduler._impl.manager is manager + + +def test_scheduler_shell_forwards_block_lifecycle_hooks(): + calls = [] + scheduler = LMCacheMPConnectorScheduler(_config()) + scheduler._impl = SimpleNamespace( + should_defer_free=lambda seq: seq == "held", + send_finished=lambda req_id: calls.append(("sent", req_id)), + source_blocks_released=lambda seq: calls.append(("released", seq)), + ) + + assert scheduler.should_defer_free("held") is True + assert scheduler.should_defer_free("free") is False + scheduler.send_finished("request-1") + scheduler.source_blocks_released("request-2") + assert calls == [("sent", "request-1"), ("released", "request-2")] + + +def test_plain_page_layout_uses_generic_implementations(monkeypatch): + selected = [] + + class PageWorker: + def __init__(self, config): + selected.append("worker") + + def register_kv_caches(self, *_): + pass + + class PageScheduler: + def __init__(self, config): + selected.append("scheduler") + + monkeypatch.setattr(backend, "LMCacheMPConnector", PageWorker) + monkeypatch.setattr(backend, "LMCacheMPConnectorScheduler", PageScheduler) + LMCacheMPConnector(_config()).register_kv_caches({}, SimpleNamespace(), 7) + LMCacheMPConnectorScheduler(_config()).bind_block_manager( + SimpleNamespace(paged_state_checkpoints=None) + ) + assert selected == ["worker", "scheduler"] + + +def test_worker_rejects_paged_state_without_native_contract(monkeypatch): + """A backend whose state is PAGE-copied but that publishes no native contract. + + `BlockManager` builds `paged_state_checkpoints` from the same + `state_transfer().copies`, so the scheduler shell selects the native-state + implementation. A PAGE-only worker would then store under a namespace the + native scheduler never looks up: offload silently never hits. + """ + import pytest + + selected = [] + + class PageWorker: + def __init__(self, config): + selected.append("page") + + def register_kv_caches(self, *_): + pass + + monkeypatch.setattr(backend, "LMCacheMPConnector", PageWorker) + builder = SimpleNamespace( + state_transfer=lambda: SimpleNamespace(copies=True), + ) + tensors = SimpleNamespace(state_backend=builder) + + with pytest.raises(NotImplementedError, match="paged_state_checkpoint_spec"): + LMCacheMPConnector(_config()).register_kv_caches({}, tensors, 7) + assert selected == [] + + +def test_worker_keeps_page_only_path_for_forked_state(monkeypatch): + selected = [] + + class PageWorker: + def __init__(self, config): + selected.append("page") + + def register_kv_caches(self, *_): + pass + + monkeypatch.setattr(backend, "LMCacheMPConnector", PageWorker) + builder = SimpleNamespace( + state_transfer=lambda: SimpleNamespace(copies=False), + ) + LMCacheMPConnector(_config()).register_kv_caches( + {}, SimpleNamespace(state_backend=builder), 7 + ) + assert selected == ["page"] diff --git a/tests/test_lmcache_offload_connector.py b/tests/test_lmcache_offload_connector.py index 4f04118d90..6d50955780 100644 --- a/tests/test_lmcache_offload_connector.py +++ b/tests/test_lmcache_offload_connector.py @@ -306,7 +306,10 @@ def test_committed_sidecar_capacity_precedence(monkeypatch): def test_committed_sidecar_capacity_rejects_invalid_env(monkeypatch, value): monkeypatch.setenv("OFFLOAD_COMMITTED_SIDECAR_CAPACITY", value) - with pytest.raises(ValueError, match="committed sidecar index capacity"): + with pytest.raises( + ValueError, + match="committed sidecar index capacity|OFFLOAD_COMMITTED_SIDECAR_CAPACITY", + ): connector_module._committed_sidecar_capacity({}) @@ -6409,6 +6412,7 @@ def test_every_member_the_scheduler_reads_is_reachable_through_the_shell(): # the same change, deliberately. must_be_seen = { "should_defer_free", + "can_partially_deallocate_state", "abandon_save", "save_abandon_timeout_s", "release_stalled_save", diff --git a/tests/test_lmcache_offload_v4_page_slot.py b/tests/test_lmcache_offload_v4_page_slot.py index 22d8b9af3d..e1d37c1a31 100644 --- a/tests/test_lmcache_offload_v4_page_slot.py +++ b/tests/test_lmcache_offload_v4_page_slot.py @@ -32,8 +32,6 @@ from atom.kv_transfer.offload.hybrid.dsv4.connector import ( DSV4_CHECKPOINT_SAVE_CHANNEL, DSV4_PAGE_SAVE_CHANNEL, - _env_nonnegative_float, - _env_positive_float, _wait_for_publication, ) from atom.kv_transfer.offload.hybrid.dsv4.connector import ( @@ -50,6 +48,7 @@ SlotLoadSpec, SlotSaveSpec, ) +from atom.utils import envs _FINGERPRINT = bytes.fromhex("00112233445566778899aabbccddeeff") _PAYLOAD = b"\x07\x08\x09\xff" @@ -1638,11 +1637,10 @@ def _probe() -> bool: @pytest.mark.parametrize( - ("name", "parser", "value"), + ("name", "value"), [ pytest.param( "OFFLOAD_PUBLICATION_TIMEOUT_S", - _env_nonnegative_float, value, id=f"timeout-{label}", ) @@ -1655,7 +1653,6 @@ def _probe() -> bool: + [ pytest.param( "OFFLOAD_PUBLICATION_POLL_INTERVAL_S", - _env_positive_float, value, id=f"interval-{label}", ) @@ -1669,19 +1666,18 @@ def _probe() -> bool: def test_publication_env_parsers_reject_nonfinite_values( monkeypatch, name, - parser, value, ): monkeypatch.setenv(name, value) with pytest.raises(ValueError, match=rf"{name} must be finite"): - parser(name, 1.0) + getattr(envs, name) def test_publication_timeout_env_parser_accepts_zero(monkeypatch): monkeypatch.setenv("OFFLOAD_PUBLICATION_TIMEOUT_S", "0") - assert _env_nonnegative_float("OFFLOAD_PUBLICATION_TIMEOUT_S", 5.0) == 0.0 + assert envs.OFFLOAD_PUBLICATION_TIMEOUT_S == 0.0 @pytest.mark.parametrize( diff --git a/tests/test_multi_connector.py b/tests/test_multi_connector.py index 6e0ddb1189..defad56490 100644 --- a/tests/test_multi_connector.py +++ b/tests/test_multi_connector.py @@ -267,6 +267,30 @@ def _save_operation_meta(*operations): return meta +def test_partial_state_deallocation_requires_an_explicit_capable_subconnector(): + seq = SimpleNamespace(has_per_req_cache=True) + unsupported = SimpleNamespace(should_defer_free=lambda _seq: True) + declined = SimpleNamespace( + should_defer_free=lambda _seq: True, + can_partially_deallocate_state=lambda _seq: False, + ) + inactive_capable = SimpleNamespace( + should_defer_free=lambda _seq: False, + can_partially_deallocate_state=lambda _seq: True, + ) + capable = SimpleNamespace( + should_defer_free=lambda _seq: True, + can_partially_deallocate_state=lambda _seq: True, + ) + + assert _sched([unsupported, declined]).can_partially_deallocate_state(seq) is False + assert ( + _sched([unsupported, inactive_capable]).can_partially_deallocate_state(seq) + is False + ) + assert _sched([unsupported, capable]).can_partially_deallocate_state(seq) is True + + # --------------------------------------------------------------------------- # Scheduler-side # --------------------------------------------------------------------------- diff --git a/tests/test_offload_early_block_release.py b/tests/test_offload_early_block_release.py index 9428ce780d..3742946fef 100644 --- a/tests/test_offload_early_block_release.py +++ b/tests/test_offload_early_block_release.py @@ -31,6 +31,7 @@ DenseOffloadScheduler, ) from atom.model_engine.block_manager import BlockManager +from atom.model_engine.sequence import Sequence def _config(role="kv_producer", *, block_size=4): @@ -43,6 +44,7 @@ def _config(role="kv_producer", *, block_size=4): def _early_release_scheduler(monkeypatch, role="kv_producer", *, chunk_size=8): + monkeypatch.setenv("OFFLOAD_MIN_SAVE_TOKENS", "0") monkeypatch.setattr( offcfg, "build_lmcache_config", @@ -88,6 +90,23 @@ def _finish_and_lease(scheduler, seq): return protected +def _resident_sequence(scheduler, req_id, num_prompt_tokens, num_blocks): + """Build a real hashed prefix so late-save admission can reacquire it.""" + bm = BlockManager( + MockConfig( + num_kvcache_blocks=max(32, num_blocks + 8), + kv_cache_block_size=4, + enable_prefix_caching=True, + ) + ) + seq = Sequence(list(range(num_prompt_tokens)), 4, id=req_id) + assert bm.allocate(seq, bm.can_allocate(seq)) + table = list(seq.block_table) + bm.hash_blocks(seq, num_prompt_tokens, start_tokens=0) + scheduler.bind_block_manager(bm) + return bm, seq, table + + class TestBlockPoolLeaseOwnership: def test_12_block_request_releases_b9_through_b12_immediately(self, seq_factory): bm = BlockManager(MockConfig(num_kvcache_blocks=16, kv_cache_block_size=1)) @@ -300,40 +319,101 @@ def fail(*_args, **_kwargs): class TestIncrementalLeaseRelease: def test_b1_b2_release_while_b3_b8_remain_protected(self, monkeypatch): scheduler = _early_release_scheduler(monkeypatch, chunk_size=8) - seq = _seq(100, num_prompt_tokens=48, num_blocks=12) + bm, seq, table = _resident_sequence(scheduler, 100, 48, 12) scheduler.update_state_after_alloc(seq) seq.num_cached_tokens = 8 op1 = scheduler.build_connector_meta().requests[0].save_operation seq.num_cached_tokens = 32 protected = _finish_and_lease(scheduler, seq) - assert protected == frozenset(range(8)) + assert protected == frozenset(table[:2]) + bm.deallocate_partial(seq, protected) assert scheduler.connector_completion(_source_safe(op1, (0, 8))) is None - assert scheduler.take_source_safe_releases() == [frozenset({0, 1})] - assert scheduler.protected_block_ids(seq) == frozenset(range(2, 8)) + released = scheduler.take_source_safe_releases() + assert released == [frozenset(table[:2])] + bm.free_leased_blocks(released[0]) + assert scheduler.protected_block_ids(seq) == frozenset() assert scheduler.connector_completion(_store_terminal(op1)) is True op2 = scheduler.build_connector_meta().requests[0].save_operation + assert scheduler.protected_block_ids(seq) == frozenset(table[2:8]) assert scheduler.connector_completion(_source_safe(op2, (8, 32))) is None - assert scheduler.take_source_safe_releases() == [frozenset(range(2, 8))] + assert scheduler.take_source_safe_releases() == [frozenset(table[2:8])] scheduler.connector_completion(_store_terminal(op2)) assert scheduler.protected_block_ids(seq) == frozenset() def test_final_pending_save_survives_request_block_table_clear(self, monkeypatch): scheduler = _early_release_scheduler(monkeypatch, chunk_size=8) - seq = _seq(101, num_prompt_tokens=32, num_blocks=8) + bm, seq, table = _resident_sequence(scheduler, 101, 32, 8) scheduler.update_state_after_alloc(seq) seq.num_cached_tokens = 32 protected = _finish_and_lease(scheduler, seq) - assert protected == frozenset(range(8)) + assert protected == frozenset() assert scheduler.has_pending_work() is True - seq.block_table.clear() - seq.num_cached_tokens = 0 + bm.deallocate_partial(seq, protected) + assert list(seq.block_table) == [] + assert bm.kv.num_used == 0 request = scheduler.build_connector_meta().requests[0] - assert request.block_ids == list(range(8)) + assert request.block_ids == table assert request.token_ids == list(range(32)) + assert scheduler.protected_block_ids(seq) == frozenset(table) + assert bm.kv.num_used == 8 + + def test_late_acquire_stops_at_first_evicted_hash_gap(self, monkeypatch): + scheduler = _early_release_scheduler(monkeypatch, chunk_size=8) + bm, seq, table = _resident_sequence(scheduler, 102, 32, 8) + scheduler.update_state_after_alloc(seq) + seq.num_cached_tokens = 32 + protected = _finish_and_lease(scheduler, seq) + bm.deallocate_partial(seq, protected) + + # Reuse the fifth block for unrelated content. The first four blocks + # remain canonical, but the save must not splice blocks after this gap. + bm.kv.allocate(table[4]) + [request] = scheduler.build_connector_meta().requests + assert request.token_ids == list(range(16)) + assert request.block_ids[:4] == table[:4] + assert request.save_spec.skip_leading_tokens == 0 + assert scheduler.protected_block_ids(seq) == frozenset(table[:4]) + + def test_late_acquire_trims_claims_in_token_order(self, monkeypatch): + scheduler = _early_release_scheduler(monkeypatch, chunk_size=8) + released = [] + scheduler._block_manager = SimpleNamespace( + acquire_offload_prefix=lambda *_args: ( + [1, 33, 65, -1], + 12, + (1, 33, 65), + ), + free_leased_blocks=lambda blocks: released.append(tuple(blocks)), + ) + seq = _seq(104, num_prompt_tokens=16, num_blocks=4) + + target, block_ids, protected = scheduler._late_save_source(seq, 0, 16) + + assert target == 8 + assert block_ids[:2] == [1, 33] + assert protected == frozenset({1, 33}) + assert released == [(65,)] + + def test_late_acquire_below_save_threshold_is_released_and_skipped( + self, monkeypatch + ): + scheduler = _early_release_scheduler(monkeypatch, chunk_size=8) + scheduler._min_save_tokens = 24 + bm, seq, table = _resident_sequence(scheduler, 103, 32, 8) + scheduler.update_state_after_alloc(seq) + seq.num_cached_tokens = 32 + protected = _finish_and_lease(scheduler, seq) + bm.deallocate_partial(seq, protected) + bm.kv.allocate(table[4]) + + assert scheduler.build_connector_meta().requests == [] + assert str(seq.id) not in scheduler._save_tracker + assert scheduler.protected_block_ids(seq) == frozenset() + assert bm.kv.num_used == 1 class TestTPQuorum: @@ -432,20 +512,22 @@ def test_mixed_tp_outcome_waits_for_all_source_quiescent_reports(self, monkeypat assert scheduler.take_source_safe_releases() == [frozenset({0, 1})] assert scheduler.build_connector_meta().requests == [] - def test_retired_failure_also_releases_unemitted_suffix(self, monkeypatch): + def test_retired_failure_releases_its_lease_and_drops_the_suffix(self, monkeypatch): scheduler = _early_release_scheduler(monkeypatch, chunk_size=8) seq = _seq(295, num_prompt_tokens=16, num_blocks=4) scheduler.update_state_after_alloc(seq) seq.num_cached_tokens = 8 operation = scheduler.build_connector_meta().requests[0].save_operation seq.num_cached_tokens = 16 - assert _finish_and_lease(scheduler, seq) == frozenset({0, 1, 2, 3}) + # Teardown leases only the emitted save's source; the unemitted suffix + # is left for late-save admission to reacquire. + assert _finish_and_lease(scheduler, seq) == frozenset({0, 1}) scheduler.connector_completion(_store_terminal(operation, False)) scheduler.connector_completion(_source_quiescent(operation)) released = scheduler.take_source_safe_releases() - assert sorted(block for group in released for block in group) == [0, 1, 2, 3] + assert sorted(block for group in released for block in group) == [0, 1] assert scheduler.build_connector_meta().requests == [] assert scheduler.has_pending_work() is False @@ -550,6 +632,7 @@ def test_timeout_drops_an_unemitted_finished_save_before_freeing(self, monkeypat seq = _seq(399, num_prompt_tokens=16, num_blocks=4) scheduler.update_state_after_alloc(seq) seq.num_cached_tokens = 8 + scheduler.build_connector_meta() _finish_and_lease(scheduler, seq) scheduler._save_lease_at[id(seq)] = time.monotonic() - 10 diff --git a/tests/test_page_unit_checkpoint_leases.py b/tests/test_page_unit_checkpoint_leases.py new file mode 100644 index 0000000000..776ea7574c --- /dev/null +++ b/tests/test_page_unit_checkpoint_leases.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: MIT + +"""Native PAGE image leases used by external checkpoint transfers.""" + +from atom.model_engine.block_pool import BlockPool +from atom.model_engine.page_unit_checkpoint import ( + PagedStateCheckpointCoordinator, + PagedStateCheckpointSpec, +) + + +def coordinator(num_units=9): + return PagedStateCheckpointCoordinator( + BlockPool(num_units), + PagedStateCheckpointSpec(10, 50, "lease-layout", image_bytes=25), + enabled=True, + ) + + +def ready(c, prefix_hash): + op = c.store.begin_store(prefix_hash, src_slot=7) + assert op is not None + c.complete_previous_batch() + return op.unit_ids + + +def test_source_lease_exposes_only_an_existing_exact_ready_image(): + c = coordinator() + assert c.acquire_checkpoint_source(101) is None + original = c.store.begin_store(101, src_slot=7) + assert original is not None + assert c.acquire_checkpoint_source(101) is None + + c.complete_previous_batch() + leased = c.acquire_checkpoint_source(101) + assert leased is not None + operation, units = leased + assert operation.prefix_hash == 101 + assert units == original.unit_ids + assert c.acquire_checkpoint_source(102) is None + assert c.take_checkpoint_ops() == ((), ()) + assert c.store.pool.num_used == 3 + assert len(c.store.records) == 1 + assert c.checkpoints_kept == 0 + + +def test_source_lease_blocks_eviction_until_reader_release(): + c = coordinator(num_units=3) + units = ready(c, 101) + operation, leased_units = c.acquire_checkpoint_source(101) + assert leased_units == units + assert not c.has_available_units(1) + assert c.store.begin_store(202, src_slot=8) is None + + c.unindex(101) + assert not c.contains(101) + assert c.store.pool.num_free == 0 + c.release_offload_store_source(operation) + assert c.store.pool.num_free == 3 + assert c.has_offload_pins() + c.release_offload_store_source(operation) + c.settle_offload_store(operation) + c.settle_offload_store(operation) + assert c.store.pool.num_free == 3 + assert not c.has_offload_pins() + + +def test_source_release_keeps_duplicate_guard_until_settlement(): + c = coordinator(num_units=3) + ready(c, 101) + first, _ = c.acquire_checkpoint_source(101) + assert c.acquire_checkpoint_source(101) is None + c.release_offload_store_source(first) + assert c.acquire_checkpoint_source(101) is None + + c.settle_offload_store(first) + second, _ = c.acquire_checkpoint_source(101) + assert second.prefix_hash == first.prefix_hash + assert second.generation > first.generation + c.release_offload_store_source(first) + c.settle_offload_store(first) + assert not c.has_available_units(1) + assert c.has_offload_pins() + c.settle_offload_store(second) + assert c.has_available_units(3) + + +def test_operation_limit_includes_source_released_results_still_pending(): + c = coordinator() + ready(c, 101) + ready(c, 202) + assert c.acquire_checkpoint_source(101, max_inflight=0) is None + first, _ = c.acquire_checkpoint_source(101, max_inflight=1) + assert c.acquire_checkpoint_source(202, max_inflight=1) is None + c.release_offload_store_source(first) + assert c.acquire_checkpoint_source(202, max_inflight=1) is None + c.settle_offload_store(first) + assert c.acquire_checkpoint_source(202, max_inflight=1) is not None + + +def test_external_source_never_times_out_while_legacy_policy_is_preserved( + monkeypatch, +): + monkeypatch.setattr("atom.model_engine.page_unit_checkpoint.monotonic", lambda: 1.0) + c = coordinator(num_units=6) + c.store._offload_sink = True + ready(c, 101) + ready(c, 202) + external, _ = c.acquire_checkpoint_source(101) + [(legacy, _)] = c.take_offload_stores(max_inflight=2) + assert legacy.prefix_hash == 202 + assert legacy.generation > external.generation + + monkeypatch.setattr( + "atom.model_engine.page_unit_checkpoint.monotonic", lambda: 1000.0 + ) + assert c.reclaim_stale_offload_pins(timeout_s=1) == 1 + assert c.was_reclaimed(legacy) + assert not c.was_reclaimed(external) + c.unindex(101) + assert c.store.pool.num_free == 0 + assert c.reclaim_stale_offload_pins(timeout_s=1) == 0 + assert c.has_offload_pins() + c.release_offload_store_source(external) + assert c.store.pool.num_free == 3 + assert c.reclaim_stale_offload_pins(timeout_s=1) == 0 + assert c.has_offload_pins() + c.settle_offload_store(external) + assert not c.has_offload_pins() + + +def test_source_lease_survives_cache_reset_until_terminal_completion(): + c = coordinator(num_units=3) + ready(c, 101) + operation, _ = c.acquire_checkpoint_source(101) + c.clear_index() + assert c.store.pool.num_free == 0 + assert c.acquire_checkpoint_source(101) is None + # A terminal cancellation before DMA starts may settle in one phase. + c.settle_offload_store(operation) + assert c.store.pool.num_free == 3 + + +def test_external_load_reserves_raw_image_units_without_publishing_checkpoint(): + c = coordinator(num_units=3) + owner = ("request-a", 1) + units = c.reserve_transfer_units(owner) + assert units is not None and len(units) == 3 + assert c.store.pool.num_free == 0 + assert c.reserve_transfer_units(owner) is None + assert c.reserve_transfer_units(("request-b", 1)) is None + assert c.store.records == {} + assert c.take_checkpoint_ops() == ((), ()) + + c.clear_index() + c.complete_previous_batch() + assert c.store.pool.num_free == 0 + c.release_transfer_units(("request-a", 0)) + assert c.store.pool.num_free == 0 + c.release_transfer_units(owner) + c.release_transfer_units(owner) + assert c.store.pool.num_free == 3 + + new_owner = ("request-a", 2) + assert c.reserve_transfer_units(new_owner) is not None + c.release_transfer_units(owner) + assert c.store.pool.num_free == 0 + c.release_transfer_units(new_owner) + assert c.store.pool.num_free == 3 + + +def test_completed_external_load_is_adopted_as_ready_without_free_list_window(): + c = coordinator(num_units=3) + owner = ("request-a", 1) + units = c.reserve_transfer_units(owner) + assert units is not None + assert c.store.pool.num_free == 0 + + assert c.adopt_transfer_units(owner, 101) + assert c.contains(101) + checkpoint_id = c.store.lookup(101) + record = c.store.records[checkpoint_id] + assert record.unit_ids == units + assert record.pin_count == 0 + assert c.store.pool.num_free == 0 + assert not c.store._offload_ready + + c.unindex(101) + assert c.store.pool.num_free == 3 + + +def test_duplicate_external_load_adoption_releases_incoming_units(): + c = coordinator(num_units=6) + canonical = ready(c, 101) + owner = ("request-a", 1) + incoming = c.reserve_transfer_units(owner) + assert incoming is not None and incoming != canonical + assert c.store.pool.num_free == 0 + + assert not c.adopt_transfer_units(owner, 101) + assert c.store.records[c.store.lookup(101)].unit_ids == canonical + assert c.store.pool.num_free == 3 + + +def test_suspended_restore_can_resume_or_release_its_source_pin(): + c = coordinator(num_units=3) + ready(c, 101) + checkpoint_id = c.store.lookup(101) + + assert c.begin_restore(101, dst_slot=9) + suspended = c.suspend_queued_restore(9) + assert suspended is not None + assert c.take_checkpoint_ops()[1] == () + assert c.store.records[checkpoint_id].pin_count == 1 + + c.resume_suspended_restore(suspended) + [restore] = c.take_checkpoint_ops()[1] + assert restore.dst_slot == 9 + c.complete_previous_batch() + assert c.store.records[checkpoint_id].pin_count == 0 + + assert c.begin_restore(101, dst_slot=10) + suspended = c.suspend_queued_restore(10) + assert suspended is not None + c.release_suspended_restore(suspended) + assert c.store.records[checkpoint_id].pin_count == 0 + assert c.take_checkpoint_ops()[1] == () + + +def test_external_load_evicts_only_available_checkpoint_sources(): + c = coordinator(num_units=6) + ready(c, 101) + ready(c, 202) + operation, source_units = c.acquire_checkpoint_source(101) + incoming = c.reserve_transfer_units(("request-a", 1)) + assert incoming is not None + assert set(incoming).isdisjoint(source_units) + assert c.contains(101) + assert not c.contains(202) + assert c.store.evictions == 1 + assert c.reserve_transfer_units(("request-b", 1)) is None + assert c.store.evictions == 1 + + c.release_transfer_units(("request-a", 1)) + c.settle_offload_store(operation) + assert c.has_available_units(6) diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 0ae1bb0642..92346a0675 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -2040,6 +2040,32 @@ def test_producer_retains_blocks_until_send_finishes(self, seq_factory, pp_size) assert not seq.block_table assert sched.is_finished() + def test_per_request_state_without_capability_keeps_whole_request_deferred( + self, seq_factory + ): + sched = Scheduler( + MockConfig( + pool_entries={"state": 2}, + pool_entries_per_req={"state": 1}, + ) + ) + seq = self._prefill( + sched, + seq_factory([1, 2, 3, 4], has_per_req_cache=True), + ) + pending: set[str] = set() + sched.kv_connector = SimpleNamespace( + request_finished=lambda s: pending.add(str(s.id)), + should_defer_free=lambda s: str(s.id) in pending, + protected_block_ids=lambda _s: frozenset(), + ) + + sched.postprocess([seq], self._output(seq.id, [sched.eos_token_id])) + + assert sched.deferred_free_blocks[seq.id] is seq + assert seq.block_table + assert seq.state_slot >= 0 + @pytest.mark.parametrize("pp_size", [1, 4]) @pytest.mark.parametrize("streaming", [False, True]) @pytest.mark.parametrize( From 19d92e03280abdda28f9110692074a82addbb0f7 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Sat, 26 Sep 2026 03:24:51 +0000 Subject: [PATCH 02/28] docs(offload): pass --eviction-policy to the LMCache MP server example The pinned LMCache (05fc77a) requires --eviction-policy on `lmcache server`; without it the documented command exits 2 during argument parsing. Reported-by: kvnloo Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/kv_transfer/offload/README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index f439614697..ed827728fa 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -145,7 +145,8 @@ access to the ATOM worker allocations: lmcache server --host 127.0.0.1 --port 5555 \ --chunk-size 256 \ --null-block-id -1 --separate-object-groups \ - --supported-transfer-mode lmcache_driven --l1-size-gb 64 + --supported-transfer-mode lmcache_driven --l1-size-gb 64 \ + --eviction-policy LRU ``` For example, add the following settings to a DSv4 launch that publishes the From ad850b5eee8dd79379fd18b18fe9acb12142ee39 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Sat, 26 Sep 2026 14:06:59 +0000 Subject: [PATCH 03/28] fix(lmcache_mp): skip native saves shorter than OFFLOAD_MIN_SAVE_TOKENS Native MP stored every READY checkpoint boundary, including whole short prompts, although a prefix below OFFLOAD_MIN_LOAD_TOKENS (8192 by default, equal to OFFLOAD_MIN_SAVE_TOKENS) can never be loaded back. On DSv4-Pro TP8 with a 1K/1K workload at concurrency 64, offload stored 257 such prefixes and cost 12.4% output throughput; skipping them stores none and leaves 1.6-2.2%, within run-to-run noise. 16K prompts still save as before. Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/kv_transfer/offload/README.md | 7 ++++-- .../offload/mp/native_state_scheduler.py | 6 ++++- atom/utils/envs.py | 5 ++-- tests/test_lmcache_mp_native_scheduler.py | 24 ++++++++++++++++++- 4 files changed, 36 insertions(+), 6 deletions(-) diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index ed827728fa..fe78bee88e 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -183,8 +183,11 @@ and temporary restore images together. By default it is worker. Candidates consume no PAGE/image pin until admission. If a request has already finished, admission resolves its token/hash chain through the live prefix index and stores only the still-resident contiguous prefix. -`OFFLOAD_MIN_SAVE_TOKENS` suppresses a late save whose remaining prefix is too -small. Unless configured otherwise, the shared save limit is +Native MP never stores a prefix shorter than `OFFLOAD_MIN_SAVE_TOKENS`: with +the default equal to `OFFLOAD_MIN_LOAD_TOKENS` it could never be loaded back, +and on a 1K-token workload skipping those saves took the offload throughput +cost from 12% to within run-to-run noise. The same threshold also suppresses a +late save whose remaining prefix is too small. Unless configured otherwise, the shared save limit is `max(2, 2 * OFFLOAD_COPY_WORKERS)`. The model namespace includes PAGE/model geometry, TP and speculation settings, diff --git a/atom/kv_transfer/offload/mp/native_state_scheduler.py b/atom/kv_transfer/offload/mp/native_state_scheduler.py index b69848ef8b..70aeefdd91 100644 --- a/atom/kv_transfer/offload/mp/native_state_scheduler.py +++ b/atom/kv_transfer/offload/mp/native_state_scheduler.py @@ -227,7 +227,11 @@ def _save_frontier(self, seq: Any) -> int: and failed[2] >= _MAX_SAVE_ATTEMPTS else 0 ) - for boundary in range(frontier, exhausted, -self.chunk_size): + # A prefix shorter than OFFLOAD_MIN_SAVE_TOKENS is not worth its copy: + # with the default equal to OFFLOAD_MIN_LOAD_TOKENS it could never be + # loaded back anyway. + floor = max(exhausted, self._min_save_tokens - 1, 0) + for boundary in range(frontier, floor, -self.chunk_size): if self._checkpoints.contains(self._boundary_hash(seq, boundary)): return boundary return 0 diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 2af0a54da8..1f9f79aaa0 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -889,8 +889,9 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float: # derives max(2, 2 * OFFLOAD_COPY_WORKERS) and the scheduler's state tier # uses 2. A kv_connector_extra_config "max_pending_saves" takes precedence. "OFFLOAD_MAX_PENDING_SAVES": lambda: _optional_int_env("OFFLOAD_MAX_PENDING_SAVES"), - # Minimum external-tier hit worth loading, and minimum remaining prefix - # worth a late save, in tokens. + # Minimum external-tier hit worth loading, and minimum prefix worth saving + # (native MP's shortest stored boundary and a late save's remainder), in + # tokens. "OFFLOAD_MIN_LOAD_TOKENS": lambda: max( 0, _int_env_or_default("OFFLOAD_MIN_LOAD_TOKENS", 8192) ), diff --git a/tests/test_lmcache_mp_native_scheduler.py b/tests/test_lmcache_mp_native_scheduler.py index 09d0348551..0eb0ef4ff1 100644 --- a/tests/test_lmcache_mp_native_scheduler.py +++ b/tests/test_lmcache_mp_native_scheduler.py @@ -73,8 +73,17 @@ def shutdown(self): self.closed = True -def make_scheduler(monkeypatch, *, capacity=2, budget=60, units=30, role="offload"): +def make_scheduler( + monkeypatch, + *, + capacity=2, + budget=60, + units=30, + role="offload", + min_save_tokens=0, +): monkeypatch.setenv("OFFLOAD_MAX_PENDING_SAVES", str(capacity)) + monkeypatch.setenv("OFFLOAD_MIN_SAVE_TOKENS", str(min_save_tokens)) adapter = Adapter() connections = [] @@ -223,6 +232,18 @@ def test_save_frontier_selects_existing_checkpoint_below_computed_tokens(monkeyp assert scheduler._save_tracker["1"][1] == 8 +def test_save_frontier_skips_checkpoints_shorter_than_min_save_tokens(monkeypatch): + scheduler, checkpoints, _ = make_scheduler(monkeypatch, min_save_tokens=16) + seq = sequence(computed=24) + checkpoint(scheduler, checkpoints, seq, 8) + scheduler.update_state_after_alloc(seq) + assert scheduler.build_connector_meta().requests == [] + + checkpoint(scheduler, checkpoints, seq, 16) + [request] = scheduler.build_connector_meta().requests + assert request.native_state.boundary_tokens == 16 + + def test_store_failure_rolls_back_for_bounded_retry_and_ignores_stale_reports( monkeypatch, ): @@ -448,6 +469,7 @@ def test_load_and_save_share_the_native_byte_budget(monkeypatch): def engine_scheduler(monkeypatch): adapter = Adapter() monkeypatch.setenv("OFFLOAD_MIN_LOAD_TOKENS", "0") + monkeypatch.setenv("OFFLOAD_MIN_SAVE_TOKENS", "0") monkeypatch.setenv("OFFLOAD_MAX_PENDING_SAVES", "2") monkeypatch.setattr( backend, "_make_scheduler_adapter", lambda _config, **kwargs: adapter From ec3fe93bb395bd2ffe1a6e4728efbae860c583bb Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Sun, 27 Sep 2026 03:44:52 +0000 Subject: [PATCH 04/28] fix(offload): keep early release safe without a BlockManager; one save threshold MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups (#2250 §1, §4). - Unbound schedulers (the vLLM plugin, where vLLM owns the blocks) cannot reacquire a finished request's prefix, so teardown again freezes the block table and leases the unemitted suffix, and the final save reads exactly those leased blocks. The late-save reacquire path now runs only when a BlockManager is bound, instead of keying off an empty block table that the plugin never clears. - The shared base no longer reads OFFLOAD_MIN_SAVE_TOKENS, so dense and hybrid behave as before this PR. A late save with nothing savable resident now retires the request instead of emitting an empty save forever. - Native MP applies OFFLOAD_MIN_SAVE_TOKENS to the absolute boundary for normal and late saves alike, so a long request's short tail is still stored, and the floor can never reach boundary 0. Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/kv_transfer/offload/chunked_scheduler.py | 43 +++++++++++---- .../offload/mp/native_state_scheduler.py | 16 +++--- atom/plugin/vllm/kv_transfer/seq_view.py | 5 +- tests/test_lmcache_mp_native_scheduler.py | 26 ++++++++- tests/test_offload_early_block_release.py | 53 +++++++++++++++---- 5 files changed, 113 insertions(+), 30 deletions(-) diff --git a/atom/kv_transfer/offload/chunked_scheduler.py b/atom/kv_transfer/offload/chunked_scheduler.py index 4c3b10cf9d..21a7dab25a 100644 --- a/atom/kv_transfer/offload/chunked_scheduler.py +++ b/atom/kv_transfer/offload/chunked_scheduler.py @@ -163,7 +163,6 @@ def __init__( # boundary, then load the aligned remainder from CPU. (Previously gated # by the OFFLOAD_UNALIGNED_HANDOFF env var; now unconditional.) self._min_load_tokens = envs.OFFLOAD_MIN_LOAD_TOKENS - self._min_save_tokens = envs.OFFLOAD_MIN_SAVE_TOKENS def bind_block_manager(self, block_manager) -> None: if self._block_manager is not None and self._block_manager is not block_manager: @@ -450,7 +449,9 @@ def _late_save_source( drop = claimed[keep_count:] if drop: self._block_manager.free_leased_blocks(drop) - if target - saved < self._min_save_tokens: + if target <= saved: + # Nothing savable is still resident (or no boundary qualifies): + # retire the request rather than emit an empty save forever. if keep: self._block_manager.free_leased_blocks(keep) return None @@ -579,14 +580,19 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: save_operation = SaveOperationId(seq.id, self._save_nonce) self._save_nonce += 1 late_acquired = frozenset() - if getattr(seq, "_offload_finished", False) and not seq.block_table: + if ( + getattr(seq, "_offload_finished", False) + and self._block_manager is not None + ): late_source = self._late_save_source(seq, saved, aligned) if late_source is None: self._save_tracker.pop(sid, None) continue aligned, block_ids, late_acquired = late_source else: - block_ids = list(seq.block_table) + block_ids = list( + getattr(seq, "_offload_finished_block_ids", seq.block_table) + ) try: request = self._build_save_request( seq, saved, aligned, save_operation, block_ids, is_last_prefill @@ -660,9 +666,13 @@ def protected_block_ids(self, seq) -> frozenset | None: does not support exact leases, or a load is in flight, since a load also touches HBM blocks this connector does not track per-range) -- the scheduler falls back to deferring the whole request. This includes a - final save that has not been emitted yet: request teardown freezes its - token identity and computed frontier, but protects no physical PAGE - until admission reacquires the still-canonical prefix. + final save that has not been emitted yet. With a bound `BlockManager` + (the ATOM engine) teardown freezes its token identity and computed + frontier but protects no physical PAGE: admission reacquires the + still-canonical prefix through the hash index. Without one (the vLLM + plugin, where vLLM owns the blocks) there is nothing to reacquire + through, so teardown freezes the block table and leases the unemitted + suffix, and the final save reads exactly those leased blocks. """ if not self._early_release or self._has_active_load(seq): return None @@ -681,6 +691,17 @@ def protected_block_ids(self, seq) -> frozenset | None: block_id for block_id in blocks.values() if block_id not in safe ) + if self._block_manager is None: + table = list(getattr(seq, "_offload_finished_block_ids", seq.block_table)) + seq._offload_finished_block_ids = table + entry = self._save_tracker.get(str(seq.id)) + if entry is not None and entry[0] is seq: + saved = int(entry[1]) + aligned = self._save_frontier(seq) + if aligned > saved: + start_block = saved // self.virtual_block_size + end_block = -(-aligned // self.virtual_block_size) + protected.update(table[start_block:end_block]) return frozenset(protected) def activate_block_leases(self, seq, block_ids: frozenset[int]) -> None: @@ -985,8 +1006,8 @@ def _retry_or_retire_failed_save(self, operation: SaveOperationId) -> None: self._save_lease_at.pop(lease_key, None) self._save_lease_owner.pop(lease_key, None) if remaining: - # Teardown leased only emitted saves' sources; the unemitted - # suffix was never leased and is dropped with the request. + # Unbound (plugin) teardown also leased the not-yet-emitted + # suffix; no further save reads it once the request retires. self._pending_source_safe_releases.append(frozenset(remaining)) self.total_source_safe_released_blocks += len(remaining) return @@ -1159,8 +1180,8 @@ def request_finished(self, seq) -> None: # 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`. + # reacquires the prefix through the bound BlockManager, or + # (unbound) reads the block table `protected_block_ids` froze. seq._offload_finished_cached_tokens = min( int(getattr(seq, "num_cached_tokens", 0)), int(seq.num_prompt_tokens), diff --git a/atom/kv_transfer/offload/mp/native_state_scheduler.py b/atom/kv_transfer/offload/mp/native_state_scheduler.py index 70aeefdd91..306839c3ad 100644 --- a/atom/kv_transfer/offload/mp/native_state_scheduler.py +++ b/atom/kv_transfer/offload/mp/native_state_scheduler.py @@ -97,6 +97,10 @@ def bind_block_manager(self, block_manager: Any) -> None: try: require_native_state_server(self._mp_adapter, self._config) self._hash_block_size = int(block_manager.hash_block_size) + # The shortest prefix worth storing, in absolute tokens. With the + # default equal to OFFLOAD_MIN_LOAD_TOKENS a shorter prefix could + # never be loaded back. Normal and late saves apply the same rule. + self._min_save_tokens = envs.OFFLOAD_MIN_SAVE_TOKENS if self.chunk_size % self._hash_block_size: raise ValueError( "native-state LMCache MP chunk size must align to native " @@ -227,18 +231,18 @@ def _save_frontier(self, seq: Any) -> int: and failed[2] >= _MAX_SAVE_ATTEMPTS else 0 ) - # A prefix shorter than OFFLOAD_MIN_SAVE_TOKENS is not worth its copy: - # with the default equal to OFFLOAD_MIN_LOAD_TOKENS it could never be - # loaded back anyway. - floor = max(exhausted, self._min_save_tokens - 1, 0) - for boundary in range(frontier, floor, -self.chunk_size): + for boundary in range(frontier, self._save_floor(exhausted), -self.chunk_size): if self._checkpoints.contains(self._boundary_hash(seq, boundary)): return boundary return 0 + def _save_floor(self, lower: int) -> int: + """Exclusive lower bound for boundaries worth storing (never 0).""" + return max(lower, self._min_save_tokens - 1, 0) + def _late_save_frontier(self, seq: Any, saved: int, available: int) -> int: available = self._chunk_floor(available) - for boundary in range(available, saved, -self.chunk_size): + for boundary in range(available, self._save_floor(saved), -self.chunk_size): if self._checkpoints.contains(self._boundary_hash(seq, boundary)): return boundary return saved diff --git a/atom/plugin/vllm/kv_transfer/seq_view.py b/atom/plugin/vllm/kv_transfer/seq_view.py index 7760b99e97..2727c8a4fc 100644 --- a/atom/plugin/vllm/kv_transfer/seq_view.py +++ b/atom/plugin/vllm/kv_transfer/seq_view.py @@ -27,11 +27,13 @@ class SeqView: "_load_operation", "_num_cached_tokens", # Written by the chunked scheduler's early-block-release path, which - # marks a finished request and freezes its computed prefix so a final + # marks a finished request and freezes its computed prefix (and, with no + # BlockManager to reacquire through, its block table) so a final # save can still be dispatched after vLLM has handed the blocks back. # Deliberately left unset in ``__init__``: an unset slot reads as absent # the same way a missing attribute does. "_offload_finished", + "_offload_finished_block_ids", "_offload_finished_cached_tokens", "_request", "block_table", @@ -118,6 +120,7 @@ def reset_for_preemption(self) -> None: # preempted request's is as stale as the live block table. for frozen in ( "_offload_finished", + "_offload_finished_block_ids", "_offload_finished_cached_tokens", ): if hasattr(self, frozen): diff --git a/tests/test_lmcache_mp_native_scheduler.py b/tests/test_lmcache_mp_native_scheduler.py index 0eb0ef4ff1..95ff8dfef6 100644 --- a/tests/test_lmcache_mp_native_scheduler.py +++ b/tests/test_lmcache_mp_native_scheduler.py @@ -83,7 +83,10 @@ def make_scheduler( min_save_tokens=0, ): monkeypatch.setenv("OFFLOAD_MAX_PENDING_SAVES", str(capacity)) - monkeypatch.setenv("OFFLOAD_MIN_SAVE_TOKENS", str(min_save_tokens)) + if min_save_tokens is None: + monkeypatch.delenv("OFFLOAD_MIN_SAVE_TOKENS", raising=False) + else: + monkeypatch.setenv("OFFLOAD_MIN_SAVE_TOKENS", str(min_save_tokens)) adapter = Adapter() connections = [] @@ -244,6 +247,27 @@ def test_save_frontier_skips_checkpoints_shorter_than_min_save_tokens(monkeypatc assert request.native_state.boundary_tokens == 16 +def test_default_min_save_tokens_skips_a_short_prompt(monkeypatch): + scheduler, checkpoints, _ = make_scheduler(monkeypatch, min_save_tokens=None) + seq = sequence(computed=24) + checkpoint(scheduler, checkpoints, seq, 16) + scheduler.update_state_after_alloc(seq) + assert scheduler._min_save_tokens == 8192 + assert scheduler.build_connector_meta().requests == [] + + +def test_late_save_threshold_is_the_absolute_boundary_not_the_remainder( + monkeypatch, +): + scheduler, checkpoints, _ = make_scheduler(monkeypatch, min_save_tokens=16) + seq = sequence(computed=24) + checkpoint(scheduler, checkpoints, seq, 24) + # Already saved 16; the 8-token remainder still reaches boundary 24 >= 16. + assert scheduler._late_save_frontier(seq, 16, 24) == 24 + scheduler._min_save_tokens = 32 + assert scheduler._late_save_frontier(seq, 16, 24) == 16 + + def test_store_failure_rolls_back_for_bounded_retry_and_ignores_stale_reports( monkeypatch, ): diff --git a/tests/test_offload_early_block_release.py b/tests/test_offload_early_block_release.py index 3742946fef..60919a3d82 100644 --- a/tests/test_offload_early_block_release.py +++ b/tests/test_offload_early_block_release.py @@ -44,7 +44,6 @@ def _config(role="kv_producer", *, block_size=4): def _early_release_scheduler(monkeypatch, role="kv_producer", *, chunk_size=8): - monkeypatch.setenv("OFFLOAD_MIN_SAVE_TOKENS", "0") monkeypatch.setattr( offcfg, "build_lmcache_config", @@ -398,23 +397,57 @@ def test_late_acquire_trims_claims_in_token_order(self, monkeypatch): assert protected == frozenset({1, 33}) assert released == [(65,)] - def test_late_acquire_below_save_threshold_is_released_and_skipped( - self, monkeypatch - ): + def test_late_acquire_with_nothing_resident_retires_the_request(self, monkeypatch): scheduler = _early_release_scheduler(monkeypatch, chunk_size=8) - scheduler._min_save_tokens = 24 bm, seq, table = _resident_sequence(scheduler, 103, 32, 8) scheduler.update_state_after_alloc(seq) seq.num_cached_tokens = 32 protected = _finish_and_lease(scheduler, seq) bm.deallocate_partial(seq, protected) - bm.kv.allocate(table[4]) + bm.kv.allocate(table[0]) # evict the first block: nothing is savable assert scheduler.build_connector_meta().requests == [] assert str(seq.id) not in scheduler._save_tracker + assert scheduler.build_connector_meta().requests == [] assert scheduler.protected_block_ids(seq) == frozenset() assert bm.kv.num_used == 1 + def test_late_save_persists_the_short_tail_of_a_long_request(self, monkeypatch): + """The tail after earlier saves is stored however short it is.""" + scheduler = _early_release_scheduler(monkeypatch, chunk_size=8) + bm, seq, table = _resident_sequence(scheduler, 105, 32, 8) + scheduler.update_state_after_alloc(seq) + seq.num_cached_tokens = 24 + [first] = scheduler.build_connector_meta().requests + scheduler.save_finished(first.save_operation) + seq.num_cached_tokens = 32 + protected = _finish_and_lease(scheduler, seq) + bm.deallocate_partial(seq, protected) + + [tail] = scheduler.build_connector_meta().requests + assert tail.save_spec.skip_leading_tokens == 24 + assert tail.token_ids == list(range(32)) + assert tail.block_ids[6:8] == table[6:8] + + def test_unbound_teardown_leases_the_blocks_its_final_save_reads(self, monkeypatch): + """No BlockManager (the vLLM plugin): vLLM frees every unleased block at + teardown and the table is not cleared, so the final save must read only + blocks the teardown lease kept.""" + scheduler = _early_release_scheduler(monkeypatch, chunk_size=8) + seq = _seq(106, num_prompt_tokens=16, num_blocks=4) + scheduler.update_state_after_alloc(seq) + seq.num_cached_tokens = 8 + first = scheduler.build_connector_meta().requests[0].save_operation + seq.num_cached_tokens = 16 + leased = _finish_and_lease(scheduler, seq) + scheduler.save_finished(first) + + [final] = scheduler.build_connector_meta().requests + start = final.save_spec.skip_leading_tokens // 4 + end = len(final.token_ids) // 4 + assert final.block_ids[start:end] == [2, 3] + assert set(final.block_ids[start:end]) <= leased + class TestTPQuorum: def test_one_incomplete_rank_prevents_logical_group_release(self, monkeypatch): @@ -512,22 +545,20 @@ def test_mixed_tp_outcome_waits_for_all_source_quiescent_reports(self, monkeypat assert scheduler.take_source_safe_releases() == [frozenset({0, 1})] assert scheduler.build_connector_meta().requests == [] - def test_retired_failure_releases_its_lease_and_drops_the_suffix(self, monkeypatch): + def test_retired_failure_also_releases_unemitted_suffix(self, monkeypatch): scheduler = _early_release_scheduler(monkeypatch, chunk_size=8) seq = _seq(295, num_prompt_tokens=16, num_blocks=4) scheduler.update_state_after_alloc(seq) seq.num_cached_tokens = 8 operation = scheduler.build_connector_meta().requests[0].save_operation seq.num_cached_tokens = 16 - # Teardown leases only the emitted save's source; the unemitted suffix - # is left for late-save admission to reacquire. - assert _finish_and_lease(scheduler, seq) == frozenset({0, 1}) + assert _finish_and_lease(scheduler, seq) == frozenset({0, 1, 2, 3}) scheduler.connector_completion(_store_terminal(operation, False)) scheduler.connector_completion(_source_quiescent(operation)) released = scheduler.take_source_safe_releases() - assert sorted(block for group in released for block in group) == [0, 1] + assert sorted(block for group in released for block in group) == [0, 1, 2, 3] assert scheduler.build_connector_meta().requests == [] assert scheduler.has_pending_work() is False From c944f617f327684c748f265375016a8ed70d8747 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Sun, 27 Sep 2026 03:49:39 +0000 Subject: [PATCH 05/28] fix(lmcache_mp): bound unprovable transfers; fence native restores MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups (#2250 §3, §7). - A submission that raises, a restore that raises, or a restore event whose query raises may still be touching engine memory, so its lease is kept -- but no longer forever. _UncertainSubmission now fails the transfer after lmcache.mp.uncertain_transfer_timeout_s (default: twice lmcache.mp.mq_timeout) with a warning, so the save or load settles and its lease, budget, pending slot and descriptor slot are released. Before, one ZMQ hiccup could stop all native saves and loads and keep EngineCore busy-looping. - PAGE-only MP treats a raising save submission the same way instead of as a definite failure. The failure reported no quiescence, so under #2339 the save was never retried and its lease was never released. The now-dead _immediate_save_failures set is removed. - Native restores are fenced both ways: the restore stream waits for work already on the compute stream (write-after-read on the SLOT's previous occupant), and every step's compute stream waits on in-flight restore events (read-after-write, including SLOT relocations in build()). The fence is on the GPU, so the host does not block. Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/kv_transfer/offload/mp/backend.py | 75 +++++++++++++++---- .../offload/mp/native_state_worker.py | 42 ++++++++--- tests/test_lmcache_mp.py | 50 +++++++++++++ tests/test_lmcache_mp_native_worker.py | 71 +++++++++++++++++- 4 files changed, 211 insertions(+), 27 deletions(-) diff --git a/atom/kv_transfer/offload/mp/backend.py b/atom/kv_transfer/offload/mp/backend.py index ecf0e7a89d..a54babc498 100644 --- a/atom/kv_transfer/offload/mp/backend.py +++ b/atom/kv_transfer/offload/mp/backend.py @@ -469,6 +469,55 @@ class _PendingSave: end: int +def _uncertain_transfer_timeout_s(config: Any) -> float: + """How long an unprovable transfer keeps its lease before it is failed. + + Nothing in the MP protocol can prove that a submission whose outcome is + unknown has stopped touching engine memory, so the bound is a policy: by + default twice the adapter's own request timeout. + """ + + extra = _extra_config(config) + default = 2.0 * float(extra.get("lmcache.mp.mq_timeout", 300.0)) + timeout = float(extra.get("lmcache.mp.uncertain_transfer_timeout_s", default)) + if not timeout > 0: + raise ValueError("lmcache.mp.uncertain_transfer_timeout_s must be > 0") + return timeout + + +class _UncertainSubmission: + """A transfer whose outcome is unknown: retained, then failed at a deadline. + + A transport exception cannot prove that a remote DMA stopped, so the lease + is kept while this reports not-done. Past ``timeout_s`` it reports a failed + terminal result, so the operation settles and releases its lease and budget + instead of wedging them for the process lifetime. + """ + + def __init__(self, timeout_s: float, operation: str) -> None: + self._deadline = time.monotonic() + timeout_s + self._timeout_s = timeout_s + self._operation = operation + self._expired = False + + def query(self) -> bool: + if time.monotonic() < self._deadline: + return False + if not self._expired: + self._expired = True + logger.warning( + "LMCache MP transfer %s stayed unprovable for %.0fs; " + "treating it as failed and releasing its lease", + self._operation, + self._timeout_s, + ) + return True + + def result(self, timeout: float | None = None) -> bool: + del timeout + return False + + def _terminal_future_result(future: Any | None) -> tuple[bool, Any]: """Return ``(terminal, result)`` without blocking on a device future.""" @@ -722,7 +771,7 @@ def __init__(self, config: Any) -> None: # aggregator must receive the same chunk completion from every rank # before releasing the writer's source blocks early. self._immediate_saves: dict[SaveCompletionId, tuple[int, int] | None] = {} - self._immediate_save_failures: set[SaveCompletionId] = set() + self._uncertain_timeout_s = _uncertain_transfer_timeout_s(config) self._immediate_load_failures: set[LoadCompletionId] = set() self._lock = threading.Lock() @@ -981,18 +1030,23 @@ def _submit_save(self, req: LMCacheReqMeta, event: Any) -> None: event, ) except Exception: + # The server may have received the request before the connection + # raised: keep the source leased until the transfer is provably + # over or the uncertainty bound expires. logger.exception( - "LMCache MP save submission failed for %s", + "LMCache MP save submission uncertain for %s", req.req_id, ) with self._lock: self._submitting_saves.discard(operation_id) - _remember_operation_tombstone( - operation_id, - self._completed_save_operations, - self._completed_save_operation_order, + self._pending_saves[operation_id] = _PendingSave( + completion=completion, + future=_UncertainSubmission( + self._uncertain_timeout_s, operation_id + ), + start=start, + end=end, ) - self._immediate_save_failures.add(completion) return with self._lock: self._submitting_saves.discard(operation_id) @@ -1079,15 +1133,8 @@ def get_finished(self) -> KVConnectorOutput: connector_completions.add( ConnectorCompletion(DENSE_PAGE_STORE_CHANNEL, completion, True) ) - for completion in self._immediate_save_failures: - done_save.add(completion) - if isinstance(completion, SaveOperationId): - connector_completions.add( - ConnectorCompletion(DENSE_PAGE_STORE_CHANNEL, completion, False) - ) failed_load.update(self._immediate_load_failures) self._immediate_saves.clear() - self._immediate_save_failures.clear() self._immediate_load_failures.clear() return KVConnectorOutput( finished_loading=done_load, diff --git a/atom/kv_transfer/offload/mp/native_state_worker.py b/atom/kv_transfer/offload/mp/native_state_worker.py index 70c33dfa2a..cdcaccd82f 100644 --- a/atom/kv_transfer/offload/mp/native_state_worker.py +++ b/atom/kv_transfer/offload/mp/native_state_worker.py @@ -30,6 +30,7 @@ _terminal_future_result, _tp_replication_factor, _transfer_operation_id, + _UncertainSubmission, _validate_mp_config, ) from atom.kv_transfer.offload.mp.native_state_layout import ( @@ -55,13 +56,6 @@ def require_native_state_server(adapter: Any, config: Any) -> None: ) -class _UncertainSubmission: - """A transport exception cannot prove that a remote DMA stopped.""" - - def query(self) -> bool: - return False - - @dataclass class _NativePending: request: LMCacheReqMeta @@ -241,7 +235,9 @@ def _submit_native(self, req: LMCacheReqMeta, event: Any, *, loading: bool) -> N ) pending[operation_id] = _NativePending(req, None) return - entry = _NativePending(req, _UncertainSubmission()) + entry = _NativePending( + req, _UncertainSubmission(self._uncertain_timeout_s, operation_id) + ) pending[operation_id] = entry submit = ( self._adapter.submit_retrieve_request @@ -280,6 +276,10 @@ def _begin_restore(self, pending: _NativePending) -> bool: descriptor_slot = self._restore_descriptor_slots.pop() pending.restore_event = event pending.descriptor_slot = descriptor_slot + # Write-after-read: the destination SLOT's previous occupant may still + # have compute work enqueued. Order the restore after everything the + # compute stream has been given so far. + self._restore_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(self._restore_stream): try: self._native_copy( @@ -322,6 +322,23 @@ def _emit_native_source_safe( ) pending.state_source_safe = True + def start_load_kv(self, metadata: Any) -> None: + # Read-after-write: this step's compute (including SLOT relocations in + # `build()`) must not touch a SLOT while its native restore is still + # writing. The host only polls the event, so fence the compute stream + # on the GPU instead of blocking here. + with self._lock: + restores = [ + pending.restore_event + for pending in self._native_loads.values() + if pending.restore_event is not None + ] + if restores: + compute = torch.cuda.current_stream() + for event in restores: + compute.wait_event(event) + super().start_load_kv(metadata) + def get_finished(self) -> KVConnectorOutput: output = KVConnectorOutput() with self._lock: @@ -376,7 +393,9 @@ def get_finished(self) -> KVConnectorOutput: "LMCache MP native restore safety unknown; " "retaining lease" ) - pending.future = _UncertainSubmission() + pending.future = _UncertainSubmission( + self._uncertain_timeout_s, operation_id + ) pending.restore_event = None continue if pending.restore_event is not None: @@ -388,6 +407,11 @@ def get_finished(self) -> KVConnectorOutput: "LMCache MP native restore event safety unknown; " "retaining lease" ) + pending.future = _UncertainSubmission( + self._uncertain_timeout_s, operation_id + ) + pending.restore_event = None + pending.restore_succeeded = False continue completion = pending.request.load_operation if pending.descriptor_slot is not None: diff --git a/tests/test_lmcache_mp.py b/tests/test_lmcache_mp.py index d5f3eb9977..79e09c7d55 100644 --- a/tests/test_lmcache_mp.py +++ b/tests/test_lmcache_mp.py @@ -1200,6 +1200,56 @@ def test_worker_pre_submit_drops_are_immediately_terminal( assert worker._pending_saves == {} +def test_worker_uncertain_save_submission_is_bounded_then_quiescent( + fake_lmcache_modules, + monkeypatch, +): + """A raising submit may have reached the server, so the source stays leased; + past the bound the save reports every chunk source-safe and the store + failed, which lets the scheduler retry or retire it instead of wedging.""" + adapter = _WorkerAdapter() + + def uncertain(_request_id, _op, _event): + raise ConnectionError("server may have received request") + + monkeypatch.setattr(adapter, "submit_store_request", uncertain) + monkeypatch.setattr(adapter, "submit_store_request_with_chunk_events", uncertain) + now = [1000.0] + monkeypatch.setattr(mp_connector.time, "monotonic", lambda: now[0]) + worker = _worker(adapter) + operation = SaveOperationId(req_id=16, generation=1) + worker._submit_save( + LMCacheReqMeta( + req_id=16, + token_ids=list(range(16)), + block_ids=[100, 101, 102, 103], + save_spec=SaveSpec(skip_leading_tokens=0), + save_operation=operation, + ), + object(), + ) + + assert not worker.get_finished().connector_completions + assert len(worker._pending_saves) == 1 + + now[0] += worker._uncertain_timeout_s + 1 + output = worker.get_finished() + assert output.finished_saving == {operation} + assert worker._pending_saves == {} + safe = { + completion.operation_id.ranges + for completion in output.connector_completions + if completion.channel == mp_connector.DENSE_PAGE_SOURCE_SAFE_CHANNEL + } + assert safe == {((0, 8),), ((8, 16),)} + [store] = [ + completion + for completion in output.connector_completions + if completion.channel == mp_connector.DENSE_PAGE_STORE_CHANNEL + ] + assert not store.succeeded + + def test_worker_save_slices_chunk_blocks_and_preserves_operation( fake_lmcache_modules, ): diff --git a/tests/test_lmcache_mp_native_worker.py b/tests/test_lmcache_mp_native_worker.py index aec40e436a..df5014615e 100644 --- a/tests/test_lmcache_mp_native_worker.py +++ b/tests/test_lmcache_mp_native_worker.py @@ -215,6 +215,48 @@ def uncertain(*_): assert len(worker._native_saves) == 1 +def test_uncertain_submission_fails_after_the_uncertainty_bound(worker, monkeypatch): + """Retaining the lease is right, but not forever: past the bound the save + settles as failed so its lease, budget and pending slot are released.""" + from atom.kv_transfer.offload.mp import backend + + now = [1000.0] + monkeypatch.setattr(backend.time, "monotonic", lambda: now[0]) + + def uncertain(*_): + raise ConnectionError("server may have received request") + + worker._adapter.submit_store_request = uncertain + worker._submit_save(request(), object()) + now[0] += worker._uncertain_timeout_s - 1 + assert not worker.get_finished().connector_completions + assert len(worker._native_saves) == 1 + + now[0] += 2 + completions = worker.get_finished().connector_completions + [terminal] = [ + completion + for completion in completions + if completion.channel == NATIVE_STATE_MP_STORE_CHANNEL + ] + assert not terminal.succeeded + assert worker._native_saves == {} + + +def test_uncertainty_bound_defaults_to_twice_the_adapter_timeout(): + from atom.kv_transfer.offload.mp import backend + + assert backend._uncertain_transfer_timeout_s(config()) == 600.0 + assert ( + backend._uncertain_transfer_timeout_s(config(**{"lmcache.mp.mq_timeout": 5})) + == 10.0 + ) + with pytest.raises(ValueError, match="uncertain_transfer_timeout_s"): + backend._uncertain_transfer_timeout_s( + config(**{"lmcache.mp.uncertain_transfer_timeout_s": 0}) + ) + + def test_failed_retrieve_does_not_restore_or_report_success(worker, monkeypatch): monkeypatch.setattr( worker, "_begin_restore", lambda _: pytest.fail("unexpected restore") @@ -357,9 +399,10 @@ def test_collapsed_tp_non_writer_skips_transport_and_reports_safe_success(worker } == {((0, 8),), ((8, 16),)} -def test_restore_uses_independent_descriptor_slot_without_stream_sync( - worker, monkeypatch -): +def test_restore_is_fenced_against_the_compute_stream_both_ways(worker, monkeypatch): + """The restore uses its own descriptor slot and stream, and orders against + compute both ways: it waits for work already on the compute stream (the + SLOT's previous occupant), and the next step's compute waits for it.""" from atom.kv_transfer.offload.mp import native_state_worker class Event: @@ -369,6 +412,18 @@ def __init__(self): def record(self, stream): self.recorded_on = stream + class Stream: + def __init__(self, name): + self.name = name + self.waited_streams = [] + self.waited_events = [] + + def wait_stream(self, other): + self.waited_streams.append(other) + + def wait_event(self, event): + self.waited_events.append(event) + class StreamContext: def __enter__(self): return self @@ -376,7 +431,8 @@ def __enter__(self): def __exit__(self, *_args): return False - restore_stream = object() + compute = Stream("compute") + restore_stream = Stream("restore") event = Event() copied = [] worker._restore_stream = restore_stream @@ -388,6 +444,9 @@ def __exit__(self, *_args): monkeypatch.setattr( native_state_worker.torch.cuda, "stream", lambda stream: StreamContext() ) + monkeypatch.setattr( + native_state_worker.torch.cuda, "current_stream", lambda: compute + ) req = request(loading=True) worker._submit_load(req, object()) @@ -397,6 +456,10 @@ def __exit__(self, *_args): assert worker._restore_descriptor_slots == [] assert copied[0][2] == 3 assert event.recorded_on is restore_stream + assert restore_stream.waited_streams == [compute] + + worker.start_load_kv(object()) + assert compute.waited_events == [event] def test_native_server_chunk_mismatch_fails_before_registration(monkeypatch): From e93d0ef6356a7b22cb1c4f8825905534db654c2c Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Sun, 27 Sep 2026 03:52:57 +0000 Subject: [PATCH 06/28] fix(lmcache_mp): end sessions after the last save; unwind failed admissions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups (#2250 §7). - The MP scheduler ended a request's session in request_finished, but with early release its final save is emitted and submitted under that session afterwards. Session ends are now deferred until no save of the request is tracked or in flight. They are swept on retirement and at every metadata build. - Native save admission returns the checkpoint lease and budget charge if anything after the acquire raises; the pin is never timeout-reclaimable, so it used to leak for the process lifetime. Load admission computes the boundary hash before reserving units, and releases the units, the budget and any suspended local restore if the rest raises. Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/kv_transfer/offload/mp/backend.py | 42 +++++++++++--- .../offload/mp/native_state_scheduler.py | 58 +++++++++++++------ tests/test_lmcache_mp_native_scheduler.py | 57 ++++++++++++++++++ 3 files changed, 130 insertions(+), 27 deletions(-) diff --git a/atom/kv_transfer/offload/mp/backend.py b/atom/kv_transfer/offload/mp/backend.py index a54babc498..b1eb8ef926 100644 --- a/atom/kv_transfer/offload/mp/backend.py +++ b/atom/kv_transfer/offload/mp/backend.py @@ -1220,6 +1220,9 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: int(req.load_spec.hbm_cached_tokens), int(end), ) + # Also catches retirements that bypass _finish_retired_request, such as + # a late save with nothing left to store. + self._end_finished_sessions() return metadata def load_finished(self, req_id: Any) -> bool: @@ -1262,14 +1265,37 @@ def load_failed(self, req_id: Any) -> bool: def request_finished(self, seq: Any) -> None: super().request_finished(seq) - try: - self._mp_adapter.end_session(_mp_session_id(self._config, seq.id)) - except Exception: - logger.warning( - "LMCache MP end_session failed for request %s", - seq.id, - exc_info=True, - ) + self._pending_session_ends()[str(seq.id)] = seq + self._end_finished_sessions() + + def _finish_retired_request(self, sid: str) -> None: + super()._finish_retired_request(sid) + self._end_finished_sessions() + + def _pending_session_ends(self) -> dict[str, Any]: + """Finished requests whose MP session must outlive their last save. + + A late save is emitted, and submitted under the request's session, + after request_finished, so the session ends only once none can follow. + """ + return self.__dict__.setdefault("_sessions_to_end", {}) + + def _end_finished_sessions(self) -> None: + """End each finished request's session once no save of it can follow.""" + pending = self._pending_session_ends() + for sid, seq in list(pending.items()): + entry = self._save_tracker.get(sid) + if (entry is not None and entry[0] is seq) or sid in self._save_inflight: + continue + del pending[sid] + try: + self._mp_adapter.end_session(_mp_session_id(self._config, seq.id)) + except Exception: + logger.warning( + "LMCache MP end_session failed for request %s", + seq.id, + exc_info=True, + ) __all__ = [ diff --git a/atom/kv_transfer/offload/mp/native_state_scheduler.py b/atom/kv_transfer/offload/mp/native_state_scheduler.py index 306839c3ad..f481c6b6d8 100644 --- a/atom/kv_transfer/offload/mp/native_state_scheduler.py +++ b/atom/kv_transfer/offload/mp/native_state_scheduler.py @@ -267,13 +267,22 @@ def _build_save_request( if source is None: return None state_operation, unit_ids = source - request = super()._build_save_request( - seq, saved, aligned, operation, block_ids, is_last_prefill - ) - request.native_state = NativeStateTransfer(unit_ids, aligned, prefix_hash) - self._native_saves[operation] = _NativeSave( - seq, state_operation, saved, aligned - ) + try: + request = super()._build_save_request( + seq, saved, aligned, operation, block_ids, is_last_prefill + ) + request.native_state = NativeStateTransfer(unit_ids, aligned, prefix_hash) + self._native_saves[operation] = _NativeSave( + seq, state_operation, saved, aligned + ) + except BaseException: + # Nothing was dispatched: return the lease and the budget charge + # the same way a settled save does, or they leak for the process + # lifetime (the pin is never timeout-reclaimable). + self._checkpoints.release_offload_store_source(state_operation) + self._checkpoints.settle_offload_store(state_operation) + self._refund_state_image() + raise return request def _complete_native_save( @@ -349,6 +358,7 @@ def _decide_load_after_alloc(self, seq: Any, load_spec): return False, "native_load_id_busy", hbm, lmc, need, chunk if not self._has_state_budget(): return False, "native_state_budget", hbm, lmc, need, chunk + prefix_hash = self._boundary_hash(seq, lmc) operation = LoadOperationId(seq.id, self._load_nonce) self._load_nonce += 1 units = self._charge_state_image( @@ -356,18 +366,28 @@ def _decide_load_after_alloc(self, seq: Any, load_spec): ) if units is None: return False, "native_state_units", hbm, lmc, need, chunk - local_restore = self._checkpoints.suspend_queued_restore(int(seq.state_slot)) - self._native_loads[operation] = _NativeLoad( - seq, - NativeStateTransfer( - units, - lmc, - self._boundary_hash(seq, lmc), - destination_slot=int(seq.state_slot), - ), - hbm, - local_restore, - ) + local_restore = None + try: + local_restore = self._checkpoints.suspend_queued_restore( + int(seq.state_slot) + ) + self._native_loads[operation] = _NativeLoad( + seq, + NativeStateTransfer( + units, + lmc, + prefix_hash, + destination_slot=int(seq.state_slot), + ), + hbm, + local_restore, + ) + except BaseException: + if local_restore is not None: + self._checkpoints.resume_suspended_restore(local_restore) + self._checkpoints.release_transfer_units(operation) + self._refund_state_image() + raise self._native_load_operations[sid] = operation self._active_load_operations[sid] = (seq, operation) seq._load_operation = operation diff --git a/tests/test_lmcache_mp_native_scheduler.py b/tests/test_lmcache_mp_native_scheduler.py index 95ff8dfef6..8ef81c5c5c 100644 --- a/tests/test_lmcache_mp_native_scheduler.py +++ b/tests/test_lmcache_mp_native_scheduler.py @@ -441,6 +441,63 @@ def test_load_capacity_failure_never_parks_or_claims_missing_state(monkeypatch): assert checkpoints.store.pool.num_free == 2 +def test_session_outlives_a_save_emitted_after_request_finished(monkeypatch): + """The final save is emitted and submitted under the request's MP session + after request_finished, so the session must not end before it settles.""" + scheduler, checkpoints, adapter = make_scheduler(monkeypatch) + seq = sequence() + checkpoint(scheduler, checkpoints, seq, 16) + scheduler.update_state_after_alloc(seq) + scheduler.request_finished(seq) + assert adapter.ended == [] + + [request] = scheduler.build_connector_meta().requests + assert adapter.ended == [] + + terminal(scheduler, request.save_operation) + scheduler.build_connector_meta() + assert adapter.ended == [f"atom-offload-dp0:{seq.id}"] + + +def test_save_admission_returns_lease_and_budget_when_it_raises(monkeypatch): + from atom.kv_transfer.offload.chunked_scheduler import ( + ChunkedOffloadSchedulerBase, + ) + + scheduler, checkpoints, _ = make_scheduler(monkeypatch) + seq = sequence() + checkpoint(scheduler, checkpoints, seq, 16) + scheduler.update_state_after_alloc(seq) + + def boom(*_args, **_kwargs): + raise RuntimeError("request construction failed") + + monkeypatch.setattr(ChunkedOffloadSchedulerBase, "_build_save_request", boom) + with pytest.raises(RuntimeError, match="request construction failed"): + scheduler.build_connector_meta() + assert scheduler._pinned_state_bytes == 0 + assert scheduler._native_saves == {} + assert checkpoints.store._offload_pins == {} + + +def test_load_admission_returns_units_and_budget_when_it_raises(monkeypatch): + scheduler, checkpoints, _ = make_scheduler(monkeypatch) + free_before = checkpoints.store.pool.num_free + seq = sequence(computed=0) + scheduler.get_num_new_matched_tokens(seq) + + def boom(_slot): + raise AssertionError("two queued restores target one slot") + + monkeypatch.setattr(scheduler._checkpoints, "suspend_queued_restore", boom) + scheduler.update_state_after_alloc(seq) + with pytest.raises(AssertionError, match="two queued restores"): + scheduler.should_park_for_load_after_alloc(seq) + assert scheduler._pinned_state_bytes == 0 + assert scheduler._native_loads == {} + assert checkpoints.store.pool.num_free == free_before + + def test_load_failure_and_cancellation_wait_for_exact_terminal_report(monkeypatch): scheduler, checkpoints, adapter = make_scheduler(monkeypatch) seq = sequence(computed=0) From a608c55c19114c57a84610bb0d34b77acdb60d92 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Sun, 27 Sep 2026 03:56:33 +0000 Subject: [PATCH 07/28] fix(lmcache_mp): keep draft PAGE regions out of the native image MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups (#2250 §5, §7). - DSv4 now publishes paged_state_region_count, the number of its own PAGE regions. A draft with a pool of its own appends regions after them, and the native layout validates checkpoint coverage and builds STATE aliases from the leading regions only. Draft rows are registered as ordinary PAGE KV and never folded into the checkpoint image. Before, DSv4 with such a draft failed registration with "PAGE regions do not cover the native PAGE unit". - Auto rank collapse is off when a DSpark draft's backend owns a KV pool: that draft publishes per-rank PAGE regions, so the complete PAGE object is not replicated and registration would otherwise reject the requested collapse. - The native layout registers its PAGE group as uint8 views, like the PAGE-only path and like its own STATE aliases. LMCache's ROCm raw-pointer fallback cannot express FP8 through the CUDA array interface. Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/kv_transfer/disaggregation/types.py | 5 +++ atom/kv_transfer/offload/mp/backend.py | 21 +++++++++ .../offload/mp/native_state_layout.py | 19 ++++++-- atom/model_ops/attentions/deepseek_v4_attn.py | 1 + tests/test_lmcache_mp.py | 27 +++++++++++ tests/test_lmcache_mp_native_layout.py | 45 +++++++++++++++++-- 6 files changed, 111 insertions(+), 7 deletions(-) diff --git a/atom/kv_transfer/disaggregation/types.py b/atom/kv_transfer/disaggregation/types.py index 8b4bb70b4b..2dafedf4e1 100644 --- a/atom/kv_transfer/disaggregation/types.py +++ b/atom/kv_transfer/disaggregation/types.py @@ -269,6 +269,11 @@ class KVTransferTensors: # separate so future layouts can describe PAGE and recurrent-state # replication independently. Appended for positional compatibility. native_state_tp_replication_factor: int = 1 + # How many leading `block_regions` make up one native checkpoint PAGE unit. + # A draft with a pool of its own appends its regions after the target's; + # those are ordinary PAGE KV, never part of the state image. None means + # every region. Appended for positional compatibility. + paged_state_region_count: int | None = None # Scheduler blocks the PAGE regions are addressed in. `init=False` because # a backend cannot answer it: `req.block_ids` is the scheduler's id space, # and a backend counts in its own page -- a different unit even where it is diff --git a/atom/kv_transfer/offload/mp/backend.py b/atom/kv_transfer/offload/mp/backend.py index b1eb8ef926..559d2dd594 100644 --- a/atom/kv_transfer/offload/mp/backend.py +++ b/atom/kv_transfer/offload/mp/backend.py @@ -165,6 +165,11 @@ def _config_has_fully_replicated_tp_pages(config: Any) -> bool: registering any cache memory. """ + if _config_has_own_pool_draft(config): + # A DSpark draft whose backend owns a KV pool appends per-rank-sharded + # PAGE regions (`draft_kv.py` declares factor 1), so the complete PAGE + # object is not replicated even when the target's is. + return False hf_config = getattr(config, "hf_config", None) # MiniMax-M3 is GQA. Some TP ranks can happen to own the same KV head when # TP exceeds the global KV-head count, but the complete PAGE object is not @@ -175,6 +180,22 @@ def _config_has_fully_replicated_tp_pages(config: Any) -> bool: return getattr(hf_config, "kv_lora_rank", None) is not None +def _config_has_own_pool_draft(config: Any) -> bool: + """Whether a speculative draft caches into a KV pool of its own. + + Mirrors `spec_decode.draft_kv.draft_kv_builder`, which only DSpark calls: + the draft's own backend answers through `DRAFT_OWNS_KV_POOL`. + """ + + speculative = getattr(config, "speculative_config", None) + draft_hf = getattr(speculative, "draft_model_hf_config", None) + if getattr(speculative, "method", None) != "dspark" or draft_hf is None: + return False + from atom.utils.selector import attn_family, get_attn_backend + + return bool(get_attn_backend(attn_family(draft_hf)).DRAFT_OWNS_KV_POOL) + + def _tp_replication_factor(config: Any) -> int: """Return the PAGE rank-collapse factor selected before workers start. diff --git a/atom/kv_transfer/offload/mp/native_state_layout.py b/atom/kv_transfer/offload/mp/native_state_layout.py index cf932237a5..f7eed93981 100644 --- a/atom/kv_transfer/offload/mp/native_state_layout.py +++ b/atom/kv_transfer/offload/mp/native_state_layout.py @@ -126,6 +126,14 @@ def build_native_state_mp_layout( raise ValueError( "native-state LMCache MP needs one owned tensor view per PAGE region" ) + # Only the leading regions form a checkpoint PAGE unit; a draft's own pool + # is appended after them and is stored as ordinary PAGE KV. + state_count = getattr(transfer_tensors, "paged_state_region_count", None) + if state_count is None: + state_count = len(regions) + state_count = _positive_int("native state region count", state_count) + if state_count > len(regions): + raise ValueError("native state region count exceeds the PAGE regions") devices = set() actual_page_bytes = 0 for index, (view, region) in enumerate(zip(page_views, regions, strict=True)): @@ -153,17 +161,22 @@ def build_native_state_mp_layout( if region.reverse_indexed: raise ValueError("native PAGE regions cannot be reverse-indexed") devices.add(view.device) - actual_page_bytes += unit_bytes + if index < state_count: + actual_page_bytes += unit_bytes if len(devices) != 1: raise ValueError("native PAGE views must share one device") if actual_page_bytes != page_bytes: raise ValueError("PAGE regions do not cover the native PAGE unit") - tensors = list(page_views) + # Registered as bytes, like the PAGE-only path: LMCache's ROCm raw-pointer + # fallback cannot express FP8 through the CUDA array interface. + tensors = [view.view(torch.uint8) for view in page_views] engine_ids = [0] * len(tensors) image_offset = 0 for ordinal in range(units): - for page_view, region in zip(page_views, regions, strict=True): + for page_view, region in zip( + page_views[:state_count], regions[:state_count], strict=True + ): nbytes = min(region.unit_bytes, image_bytes - image_offset) if not nbytes: break diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index 610c68d3bc..6d86721ee9 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -2083,6 +2083,7 @@ def get_kv_transfer_tensors(self): native_state_tp_replication_factor=tp_size, paged_state_checkpoint_spec=checkpoint_spec, execute_paged_state_copies=self.execute_paged_state_copies, + paged_state_region_count=len(block_regions), ) # ------------------------------------------------------------------ # diff --git a/tests/test_lmcache_mp.py b/tests/test_lmcache_mp.py index 79e09c7d55..7c4fc87329 100644 --- a/tests/test_lmcache_mp.py +++ b/tests/test_lmcache_mp.py @@ -255,6 +255,33 @@ def __init__(self, **kwargs): assert {strategy.tp_size for strategy in strategies} == {8} +def test_auto_rank_collapse_is_off_for_a_draft_with_its_own_pool(monkeypatch): + """A DSpark draft whose backend owns a KV pool appends per-rank PAGE + regions, so a replicated MLA target no longer collapses.""" + from atom.utils import selector + + owns_pool = {"value": True} + monkeypatch.setattr(selector, "attn_family", lambda _hf: "draft-family") + monkeypatch.setattr( + selector, + "get_attn_backend", + lambda _family: SimpleNamespace(DRAFT_OWNS_KV_POOL=owns_pool["value"]), + ) + config = _config(tp=8, kv_lora_rank=512) + assert mp_connector._tp_replication_factor(config) == 8 + + config.speculative_config = SimpleNamespace( + method="dspark", draft_model_hf_config=SimpleNamespace() + ) + assert mp_connector._tp_replication_factor(config) == 1 + + owns_pool["value"] = False + assert mp_connector._tp_replication_factor(config) == 8 + config.speculative_config.method = "mtp" + owns_pool["value"] = True + assert mp_connector._tp_replication_factor(config) == 8 + + def test_auto_rank_collapse_distinguishes_glm52_mla_from_minimax_m3_gqa(): glm52 = _config(model_type="glm_moe_dsa", tp=8, kv_lora_rank=512) minimax = _config(model_type="minimax_m3_vl", tp=8) diff --git a/tests/test_lmcache_mp_native_layout.py b/tests/test_lmcache_mp_native_layout.py index 91e0f34abf..be107bfdec 100644 --- a/tests/test_lmcache_mp_native_layout.py +++ b/tests/test_lmcache_mp_native_layout.py @@ -126,10 +126,11 @@ def test_layout_coalesces_equal_shapes_inside_ordinal_and_preserves_trim_stride( assert layout.checkpoint_spec.units_per_checkpoint == 3 assert layout.checkpoint_spec.page_unit_bytes == 18 assert len(layout.tensors) == 10 - assert len(layout.kernel_groups) == 8 + # PAGE registers as bytes, so the BF16 and uint8 8-byte regions share one + # physical kernel identity just as the STATE aliases do. + assert len(layout.kernel_groups) == 7 assert tuple(group.tensor_indices for group in layout.kernel_groups) == ( - (0,), - (1,), + (0, 1), (2,), (3, 4), (5,), @@ -138,7 +139,6 @@ def test_layout_coalesces_equal_shapes_inside_ordinal_and_preserves_trim_stride( (9,), ) assert [group.engine_group_id for group in layout.kernel_groups] == [ - 0, 0, 0, 1, @@ -337,3 +337,40 @@ def test_actual_lmcache_registration_preserves_native_aliases_and_stride(): assert manager.kernel_groups[-1].shape_desc.hs == 3 assert len(manager.object_groups) == 2 assert manager.object_groups[-1].sw_size_chunks == 1 + + +def test_page_group_registers_every_region_as_bytes(): + transfer = _transfer() + layout = _layout(transfer) + page = layout.tensors[: len(transfer.block_tensor_views)] + assert transfer.block_tensor_views[0].dtype == torch.bfloat16 + assert [tensor.dtype for tensor in page] == [torch.uint8] * len(page) + assert [tensor.data_ptr() for tensor in page] == [ + view.data_ptr() for view in transfer.block_tensor_views + ] + + +def test_draft_regions_after_the_state_regions_stay_ordinary_page(): + """A draft with its own pool appends regions after the target's. They are + registered as PAGE KV but never folded into the checkpoint image.""" + target_only = _transfer(widths=(8, 8, 2)) + with_draft = _transfer(widths=(8, 8, 2, 6)) + with_draft.paged_state_checkpoint_spec = target_only.paged_state_checkpoint_spec + + with pytest.raises(ValueError, match="do not cover the native PAGE unit"): + _layout(with_draft) + + with_draft.paged_state_region_count = 3 + layout = _layout(with_draft) + page_groups = [group for group in layout.kernel_groups if not group.engine_group_id] + assert sorted(i for group in page_groups for i in group.tensor_indices) == [ + 0, + 1, + 2, + 3, + ] + ids = (0, 5, 6) + assert torch.equal( + _gather(layout, with_draft, ids), + _gather(_layout(target_only), target_only, ids), + ) From 13721f0dfee24c388fbb89574d8f38f7967c4423 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Sun, 27 Sep 2026 04:00:19 +0000 Subject: [PATCH 08/28] fix(attention): take the descriptor slot in DSv4.1 paged-state copies MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-up (#2250 §5). descriptor_slot reached the base class, DSv4 and GDN but not DSv4.1, whose execute_paged_state_copies(stores, restores) would raise TypeError inside the native restore's exception handler and be reported as a failed restore. StateCopies now keeps an independent pinned staging buffer and upload fence per descriptor slot, like the other backends. A new AST sweep test fails if any attention backend's implementation lacks the parameter. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../attentions/deepseek_v41/backend.py | 4 +- .../attentions/deepseek_v41/checkpoints.py | 61 +++++++++++++------ tests/attentions/deepseek_v41/test_cache.py | 16 +++++ tests/test_paged_state_copy_signature.py | 37 +++++++++++ 4 files changed, 96 insertions(+), 22 deletions(-) create mode 100644 tests/test_paged_state_copy_signature.py diff --git a/atom/model_ops/attentions/deepseek_v41/backend.py b/atom/model_ops/attentions/deepseek_v41/backend.py index 32ac5e42de..79199d510b 100644 --- a/atom/model_ops/attentions/deepseek_v41/backend.py +++ b/atom/model_ops/attentions/deepseek_v41/backend.py @@ -260,8 +260,8 @@ def state_entry_views(self, slot): def relocate_state_slots(self, pairs): self.copies.relocate(pairs) - def execute_paged_state_copies(self, stores, restores): - self.copies.execute(stores, restores) + def execute_paged_state_copies(self, stores, restores, descriptor_slot=0): + self.copies.execute(stores, restores, descriptor_slot=descriptor_slot) def warmup_per_req_cache(self): self.copies.warmup() diff --git a/atom/model_ops/attentions/deepseek_v41/checkpoints.py b/atom/model_ops/attentions/deepseek_v41/checkpoints.py index cdc18c63df..d7fc009705 100644 --- a/atom/model_ops/attentions/deepseek_v41/checkpoints.py +++ b/atom/model_ops/attentions/deepseek_v41/checkpoints.py @@ -11,6 +11,17 @@ from atom.utils import CpuGpuBuffer +class _DescriptorStaging: + """One pinned descriptor buffer and the event fencing its last upload.""" + + def __init__(self, rows, device, pin_memory): + self.buffer = CpuGpuBuffer( + rows, 3, dtype=torch.int64, device=device, pin_memory=pin_memory + ) + self.upload_done = torch.cuda.Event() if pin_memory else None + self.pending = False + + class StateCopies: def __init__(self, cache, spec, max_batch): self.cache, self.spec = cache, spec @@ -33,15 +44,23 @@ def __init__(self, cache, spec, max_batch): ], spec.image_bytes, ) - self.staging = CpuGpuBuffer( - 2 * max_batch * self.plan.num_spans, - 3, - dtype=torch.int64, - device=cache.pool.device, - pin_memory=cache.pool.is_cuda, + self._staging_rows = 2 * max_batch * self.plan.num_spans + # Slot 0 serves `build()`; each other descriptor slot is an independent + # buffer, so an out-of-band restore never overwrites a step's staging. + self._stagings = {0: self._new_staging()} + + def _new_staging(self): + return _DescriptorStaging( + self._staging_rows, self.cache.pool.device, self.cache.pool.is_cuda ) - self.upload_done = torch.cuda.Event() if cache.pool.is_cuda else None - self.pending = False + + def staging(self, descriptor_slot=0): + if descriptor_slot < 0: + raise ValueError("checkpoint descriptor slot must be non-negative") + staging = self._stagings.get(descriptor_slot) + if staging is None: + staging = self._stagings[descriptor_slot] = self._new_staging() + return staging def entry(self, slot): self.cache.require_committed() @@ -73,7 +92,7 @@ def _validate(self, op, storing): if any(unit < 0 or unit >= self.cache.num_pages for unit in op.unit_ids): raise IndexError("Checkpoint PAGE unit is out of range") - def execute(self, stores, restores): + def execute(self, stores, restores, descriptor_slot=0): if not stores and not restores: return for ops, storing in ((stores, True), (restores, False)): @@ -95,11 +114,12 @@ def execute(self, stores, restores): ) at += take return + staging = self.staging(descriptor_slot) total = (len(stores) + len(restores)) * self.plan.num_spans - if total > self.staging.np.shape[0]: + if total > staging.buffer.np.shape[0]: raise ValueError("Checkpoint copy batch exceeds descriptor capacity") - if self.pending: - self.upload_done.synchronize() + if staging.pending: + staging.upload_done.synchronize() at = 0 for ops, storing in ((stores, True), (restores, False)): if not ops: @@ -124,12 +144,12 @@ def execute(self, stores, restores): dtype=np.int64, ) self.plan.write_descriptor( - self.staging.np[at:end], slot_bases, page_bases, forward=storing + staging.buffer.np[at:end], slot_bases, page_bases, forward=storing ) at = end - descriptor = self.staging.copy_to_gpu(total) - self.upload_done.record() - self.pending = True + descriptor = staging.buffer.copy_to_gpu(total) + staging.upload_done.record() + staging.pending = True cut = len(stores) * self.plan.num_spans # Restore may consume an image stored in this same maintenance batch. launch_copy_descriptor(descriptor[:cut], self.plan) @@ -147,12 +167,13 @@ def warmup(self): offsets.append(at) at += size dst = np.asarray([[slot.data_ptr() + off for off in offsets]], dtype=np.int64) + staging = self.staging(0) self.plan.write_descriptor( - self.staging.np[: self.plan.num_spans], + staging.buffer.np[: self.plan.num_spans], np.asarray([[slot.data_ptr()]], dtype=np.int64), dst, ) - descriptor = self.staging.copy_to_gpu(self.plan.num_spans) - self.upload_done.record() - self.pending = True + descriptor = staging.buffer.copy_to_gpu(self.plan.num_spans) + staging.upload_done.record() + staging.pending = True launch_copy_descriptor(descriptor, self.plan) diff --git a/tests/attentions/deepseek_v41/test_cache.py b/tests/attentions/deepseek_v41/test_cache.py index 40f0776087..fecd5705fd 100644 --- a/tests/attentions/deepseek_v41/test_cache.py +++ b/tests/attentions/deepseek_v41/test_cache.py @@ -451,3 +451,19 @@ def test_shared_indptr_launch_matches_row_counts_and_replay(decode, width, ratio dtype=torch.int32, ) torch.testing.assert_close(extend, expected, rtol=0, atol=0) + + +def test_state_copy_descriptor_slots_have_independent_staging(): + from atom.model_ops.attentions.deepseek_v41.checkpoints import StateCopies + + copies = StateCopies.__new__(StateCopies) + copies._staging_rows = 3 + copies.cache = type( + "Cache", + (), + {"pool": type("Pool", (), {"device": "cpu", "is_cuda": False})()}, + )() + copies._stagings = {0: copies._new_staging()} + assert copies.staging(0) is copies.staging(0) + assert copies.staging(1) is not copies.staging(0) + assert copies.staging(1) is copies.staging(1) diff --git a/tests/test_paged_state_copy_signature.py b/tests/test_paged_state_copy_signature.py new file mode 100644 index 0000000000..c8cb7ea2ac --- /dev/null +++ b/tests/test_paged_state_copy_signature.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Every PAGE-backed state copy implementation takes a descriptor slot. + +The native LMCache MP restore always passes ``descriptor_slot=`` by keyword, +so an implementation left on the old signature fails only once a restore +reaches it, inside an exception handler that reports "restore failed". +""" + +import ast +import pathlib + +ATTENTIONS = pathlib.Path(__file__).parents[1] / "atom/model_ops/attentions" + + +def _implementations(): + for path in sorted(ATTENTIONS.rglob("*.py")): + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if ( + isinstance(node, ast.FunctionDef) + and node.name == "execute_paged_state_copies" + ): + args = node.args.args + node.args.kwonlyargs + yield path.relative_to(ATTENTIONS.parents[2]), node, args + + +def test_every_implementation_accepts_descriptor_slot(): + found = list(_implementations()) + assert len(found) >= 4 + missing = [ + f"{path}:{node.lineno}" + for path, node, args in found + if "descriptor_slot" not in {arg.arg for arg in args} + ] + assert missing == [] From 971823cb0523e31351de910a42ca0a85d429c29b Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Sun, 27 Sep 2026 04:02:05 +0000 Subject: [PATCH 09/28] fix(offload): seed native and late-save hash chains from cache_seed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-up (#2250 §2). The native boundary-hash fallback and BlockManager.acquire_offload_prefix both started their chains at -1, while every BlockManager chain starts at seq.cache_seed, which multimodal requests set. For those requests every native save was skipped, every restored checkpoint was published under a key no loader looks up, and late saves found nothing resident. The native fallback now extends its chain through the new BlockManager.prefix_hash_chain (the manager's own _chain_to), so both _boundary_hash branches share one seed, algorithm and token slicing. acquire_offload_prefix seeds from seq.cache_seed. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../offload/mp/native_state_scheduler.py | 20 +++++--------- atom/model_engine/block_manager.py | 9 ++++++- tests/test_lmcache_mp_native_scheduler.py | 26 +++++++++++++++++++ tests/test_offload_early_block_release.py | 22 ++++++++++++++++ 4 files changed, 63 insertions(+), 14 deletions(-) diff --git a/atom/kv_transfer/offload/mp/native_state_scheduler.py b/atom/kv_transfer/offload/mp/native_state_scheduler.py index f481c6b6d8..f1ebcba458 100644 --- a/atom/kv_transfer/offload/mp/native_state_scheduler.py +++ b/atom/kv_transfer/offload/mp/native_state_scheduler.py @@ -162,20 +162,14 @@ def _boundary_hash(self, seq: Any, boundary: int) -> int: if len(chain) >= count: return int(chain[count - 1]) # Some checkpoint producers do not reserve midstep checkpoints, so - # BlockManager may leave Sequence.block_hashes empty. Use its exact - # public hashing algorithm and token slices, caching only this request's - # immutable prompt chain. - chain = getattr(seq, "_mp_checkpoint_hashes", None) - if chain is None: - chain = [] + # BlockManager may leave Sequence.block_hashes empty. Extend the chain + # through BlockManager itself -- its seed (`seq.cache_seed`, which + # multimodal requests set), algorithm and token slices -- caching only + # this request's immutable prompt chain. + chain = getattr(seq, "_mp_checkpoint_hashes", None) or [] + if len(chain) < count: + chain = self._block_manager.prefix_hash_chain(seq, chain, count) seq._mp_checkpoint_hashes = chain - prefix = chain[-1] if chain else -1 - for index in range(len(chain), count): - start = index * self._hash_block_size - prefix = self._block_manager.compute_hash( - seq.token_ids[start : start + self._hash_block_size], prefix - ) - chain.append(prefix) return int(chain[count - 1]) def _lookup_token_ids(self, seq: Any) -> list[int]: diff --git a/atom/model_engine/block_manager.py b/atom/model_engine/block_manager.py index 32712c448c..e598082344 100644 --- a/atom/model_engine/block_manager.py +++ b/atom/model_engine/block_manager.py @@ -867,6 +867,13 @@ def _chain_to( chain.append(h) return chain + def prefix_hash_chain( + self, seq: Sequence, block_hashes: list[int], blocks: int + ) -> list[int]: + """`block_hashes` continued to `blocks` entries, for callers outside + the manager. Seeded from `seq.cache_seed` like every chain it mints.""" + return self._chain_to(seq, block_hashes, blocks) + def can_allocate( self, seq: Sequence, @@ -2569,7 +2576,7 @@ def acquire_offload_prefix( start_block = start_tokens // hbs block_ids = [-1] * end_block claimed: list[int] = [] - parent_hash = -1 + parent_hash = seq.cache_seed available_end = start_tokens for index in range(end_block): token_ids = self._hash_block_tokens(seq, index) diff --git a/tests/test_lmcache_mp_native_scheduler.py b/tests/test_lmcache_mp_native_scheduler.py index 8ef81c5c5c..792c2ac314 100644 --- a/tests/test_lmcache_mp_native_scheduler.py +++ b/tests/test_lmcache_mp_native_scheduler.py @@ -117,6 +117,13 @@ def connect(config, *, checkpoint_spec=None): hash_block_size=4, compute_hash=BlockManager.compute_hash, ) + # BlockManager's own chain walk (seed, algorithm, token slices). + manager._hash_block_tokens = lambda seq, index: BlockManager._hash_block_tokens( + manager, seq, index + ) + manager.prefix_hash_chain = lambda seq, hashes, blocks: BlockManager._chain_to( + manager, seq, hashes, blocks + ) scheduler.bind_block_manager(manager) scheduler.bind_block_manager(manager) assert connections == [checkpoints.store.spec] @@ -247,6 +254,25 @@ def test_save_frontier_skips_checkpoints_shorter_than_min_save_tokens(monkeypatc assert request.native_state.boundary_tokens == 16 +def test_boundary_hash_seeds_from_cache_seed_like_block_manager(monkeypatch): + """Multimodal requests seed their chain with cache_seed != -1; the fallback + walk must agree with BlockManager's own chain from block 0 onward.""" + scheduler, _, _ = make_scheduler(monkeypatch) + manager = scheduler._block_manager + seq = sequence() + seq.cache_seed = 987654321 + seq.block_hashes = [] + fallback = scheduler._boundary_hash(seq, 8) + expected = manager.prefix_hash_chain(seq, [], 2) + assert fallback == expected[1] + assert fallback != manager.prefix_hash_chain(sequence(), [], 2)[1] + + chained = sequence(request_id=2) + chained.cache_seed = seq.cache_seed + chained.block_hashes = list(expected) + assert scheduler._boundary_hash(chained, 8) == fallback + + def test_default_min_save_tokens_skips_a_short_prompt(monkeypatch): scheduler, checkpoints, _ = make_scheduler(monkeypatch, min_save_tokens=None) seq = sequence(computed=24) diff --git a/tests/test_offload_early_block_release.py b/tests/test_offload_early_block_release.py index 60919a3d82..9107d15ca8 100644 --- a/tests/test_offload_early_block_release.py +++ b/tests/test_offload_early_block_release.py @@ -397,6 +397,28 @@ def test_late_acquire_trims_claims_in_token_order(self, monkeypatch): assert protected == frozenset({1, 33}) assert released == [(65,)] + def test_late_acquire_finds_a_multimodal_prefix(self, monkeypatch): + """BlockManager seeds a multimodal chain with cache_seed; reacquiring + must use the same seed or block 0 never matches.""" + scheduler = _early_release_scheduler(monkeypatch, chunk_size=8) + bm = BlockManager( + MockConfig( + num_kvcache_blocks=24, kv_cache_block_size=4, enable_prefix_caching=True + ) + ) + seq = Sequence(list(range(16)), 4, id=107) + seq.cache_seed = 424242 + assert bm.allocate(seq, bm.can_allocate(seq)) + table = list(seq.block_table) + bm.hash_blocks(seq, 16, start_tokens=0) + scheduler.bind_block_manager(bm) + bm.deallocate(seq) + + block_ids, available, claimed = bm.acquire_offload_prefix(seq, 0, 16) + assert available == 16 + assert block_ids == table + assert sorted(claimed) == sorted(table) + def test_late_acquire_with_nothing_resident_retires_the_request(self, monkeypatch): scheduler = _early_release_scheduler(monkeypatch, chunk_size=8) bm, seq, table = _resident_sequence(scheduler, 103, 32, 8) From 7cdef48d09d4fcb7ccccd6d3841021a3ace4507d Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Sun, 27 Sep 2026 04:03:27 +0000 Subject: [PATCH 10/28] fix(offload): fail closed on disabled paged-state checkpoints and composite state release MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups (#2250 §6). - The MP scheduler shell picked the native-state scheduler whenever the checkpoint coordinator existed, but the coordinator exists without prefix caching and then can never hold a READY image, so offload silently did nothing. A PAGE-only fallback would restore KV under stale recurrent state. Binding now fails with a message naming --enable-prefix-caching. - MultiConnectorScheduler.can_partially_deallocate_state now requires every still-deferring sub to guarantee state safety. Before, one declaring sub was enough, and the state slot could be recycled while a non-declaring sub (e.g. a dense PAGE sub) still needed the request alive. The test now separates all-of from any-of. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../disaggregation/multi/multi_connector.py | 17 ++++++---- atom/kv_transfer/offload/mp/connector.py | 10 ++++++ tests/test_lmcache_mp_shell.py | 33 +++++++++++++++++++ tests/test_multi_connector.py | 22 +++++++++---- 4 files changed, 69 insertions(+), 13 deletions(-) diff --git a/atom/kv_transfer/disaggregation/multi/multi_connector.py b/atom/kv_transfer/disaggregation/multi/multi_connector.py index 5e565ca8bb..d02e4236f5 100644 --- a/atom/kv_transfer/disaggregation/multi/multi_connector.py +++ b/atom/kv_transfer/disaggregation/multi/multi_connector.py @@ -654,20 +654,23 @@ def protected_block_ids(self, seq: Any): return frozenset(protected) def can_partially_deallocate_state(self, seq: Any) -> bool: - """Forward an explicit state-lifetime guarantee from its owner. + """True only if every sub still deferring this request guarantees it. - Missing methods are not treated as approval. The separate - ``protected_block_ids`` check still requires every deferring sub to - narrow its PAGE reads before the scheduler takes this path. + Fail-closed like ``protected_block_ids``: releasing the state slot is + safe only when no deferring sub still needs the request alive, and a + missing method is not approval. ``protected_block_ids`` narrows PAGE + blocks only, so it cannot stand in for a sub's state guarantee. """ + deferring = 0 for connector in self._connectors: should_defer = getattr(connector, "should_defer_free", None) if not callable(should_defer) or not should_defer(seq): continue + deferring += 1 callback = getattr(connector, "can_partially_deallocate_state", None) - if callable(callback) and callback(seq) is True: - return True - return False + if not callable(callback) or callback(seq) is not True: + return False + return deferring > 0 def activate_block_leases(self, seq: Any, block_ids) -> None: for connector in self._connectors: diff --git a/atom/kv_transfer/offload/mp/connector.py b/atom/kv_transfer/offload/mp/connector.py index 3cc1a4e3a6..0588886528 100644 --- a/atom/kv_transfer/offload/mp/connector.py +++ b/atom/kv_transfer/offload/mp/connector.py @@ -142,6 +142,16 @@ def bind_block_manager(self, block_manager: Any) -> None: "lmcache_mp scheduler is already bound to a block manager" ) coordinator = getattr(block_manager, "paged_state_checkpoints", None) + if coordinator is not None and not getattr(coordinator, "enabled", True): + # The backend keeps PAGE-backed state checkpoints, but none can + # ever become READY. A native scheduler would never save, and a + # PAGE-only one would restore KV under stale recurrent state. + raise ValueError( + "lmcache_mp: this backend keeps PAGE-backed state checkpoints, " + "which are disabled (they need --enable-prefix-caching and at " + "least one state slot); native-state offload cannot run and " + "PAGE-only offload would restore KV without its state" + ) if coordinator is None: impl: KVConnectorSchedulerBase = backend.LMCacheMPConnectorScheduler( self._config diff --git a/tests/test_lmcache_mp_shell.py b/tests/test_lmcache_mp_shell.py index 35905fe27f..ad8a90f5cc 100644 --- a/tests/test_lmcache_mp_shell.py +++ b/tests/test_lmcache_mp_shell.py @@ -206,3 +206,36 @@ def register_kv_caches(self, *_): {}, SimpleNamespace(state_backend=builder), 7 ) assert selected == ["page"] + + +def test_scheduler_rejects_disabled_paged_state_checkpoints(monkeypatch): + """A coordinator exists whenever the backend copies PAGE-backed state, but + it can hold a READY image only with prefix caching on. Neither transport + is correct without it, so binding fails instead of silently offloading + nothing (native) or restoring KV under stale state (PAGE-only).""" + import pytest + + selected = [] + + class NativeScheduler: + def __init__(self, config): + selected.append("native") + + monkeypatch.setattr( + native_state_scheduler, + "NativeStateLMCacheMPConnectorScheduler", + NativeScheduler, + ) + monkeypatch.setattr( + backend, + "LMCacheMPConnectorScheduler", + lambda config: selected.append("page"), + ) + disabled = SimpleNamespace(paged_state_checkpoints=SimpleNamespace(enabled=False)) + with pytest.raises(ValueError, match="enable-prefix-caching"): + LMCacheMPConnectorScheduler(_config()).bind_block_manager(disabled) + assert selected == [] + + enabled = SimpleNamespace(paged_state_checkpoints=SimpleNamespace(enabled=True)) + LMCacheMPConnectorScheduler(_config()).bind_block_manager(enabled) + assert selected == ["native"] diff --git a/tests/test_multi_connector.py b/tests/test_multi_connector.py index defad56490..9a051fb48a 100644 --- a/tests/test_multi_connector.py +++ b/tests/test_multi_connector.py @@ -283,12 +283,22 @@ def test_partial_state_deallocation_requires_an_explicit_capable_subconnector(): can_partially_deallocate_state=lambda _seq: True, ) - assert _sched([unsupported, declined]).can_partially_deallocate_state(seq) is False - assert ( - _sched([unsupported, inactive_capable]).can_partially_deallocate_state(seq) - is False - ) - assert _sched([unsupported, capable]).can_partially_deallocate_state(seq) is True + inactive_unsupported = SimpleNamespace(should_defer_free=lambda _seq: False) + + def allowed(*subs): + return _sched(list(subs)).can_partially_deallocate_state(seq) + + assert allowed(unsupported, declined) is False + assert allowed(unsupported, inactive_capable) is False + # A deferring sub without the guarantee still needs the request alive, even + # though another deferring sub declares its own state safe (all, not any). + assert allowed(unsupported, capable) is False + assert allowed(capable, declined) is False + # Subs that no longer defer do not veto. + assert allowed(inactive_unsupported, capable) is True + assert allowed(capable, capable) is True + # Nothing deferring is not a guarantee either. + assert allowed(inactive_unsupported, inactive_capable) is False # --------------------------------------------------------------------------- From 17918abdd39cb09eb1583e7a9b81dfa1b09be5f6 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Sun, 27 Sep 2026 04:07:40 +0000 Subject: [PATCH 11/28] perf(lmcache_mp): cache the native save frontier; reserve restore slots early MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups (#2250 §5, §8). - Native _save_frontier walked the prompt checkpoint by checkpoint for every tracked request on every scheduler step. PageUnitCheckpointStore now has a generation that is bumped whenever the READY set can change, and the answer is cached per request on (frontier, floor, generation). - Restore descriptor slots >= 1 were first allocated as pinned memory on the connector thread mid-serving. The native worker now reserves them at registration through the new builder hook reserve_checkpoint_descriptors (base builders allocate their descriptor buffers, DSv4.1 its per-slot staging). - Document why single-host DP replicas share one (model_name, worker_id, world_size) identity. It is the content-addressed storage namespace, while the server registers GPU memory per instance_id, refcounts layout descriptors per (model_name, world_size), and _mp_session_id scopes sessions per replica. Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/kv_transfer/offload/mp/backend.py | 7 +++ .../offload/mp/native_state_scheduler.py | 17 +++++- .../offload/mp/native_state_worker.py | 7 +++ atom/model_engine/page_unit_checkpoint.py | 9 +++ atom/model_ops/attentions/backends.py | 10 ++++ .../attentions/deepseek_v41/backend.py | 4 ++ tests/test_lmcache_mp_native_scheduler.py | 23 ++++++++ tests/test_lmcache_mp_native_worker.py | 56 +++++++++++++++++++ 8 files changed, 130 insertions(+), 3 deletions(-) diff --git a/atom/kv_transfer/offload/mp/backend.py b/atom/kv_transfer/offload/mp/backend.py index 559d2dd594..2d5d21ba1c 100644 --- a/atom/kv_transfer/offload/mp/backend.py +++ b/atom/kv_transfer/offload/mp/backend.py @@ -141,6 +141,13 @@ def _validate_mp_config(config: Any) -> tuple[int, int]: raise NotImplementedError("lmcache_mp does not support DCP yet") if pcp_size != 1: raise NotImplementedError("lmcache_mp does not support PCP yet") + # Single-host DP replicas deliberately share one (model_name, worker_id, + # world_size) identity: it is the content-addressed storage namespace, so + # replicas deduplicate identical prefixes. It is not a registration key. + # The server registers GPU memory per unique instance_id, refcounts layout + # descriptors per (model_name, world_size) (all replicas publish the same + # one), and _mp_session_id scopes request sessions and their locks per + # replica. if dp_size_local != dp_size: raise NotImplementedError( "lmcache_mp supports DP and DP-attention only within one host; " diff --git a/atom/kv_transfer/offload/mp/native_state_scheduler.py b/atom/kv_transfer/offload/mp/native_state_scheduler.py index f1ebcba458..6119a00c72 100644 --- a/atom/kv_transfer/offload/mp/native_state_scheduler.py +++ b/atom/kv_transfer/offload/mp/native_state_scheduler.py @@ -225,10 +225,21 @@ def _save_frontier(self, seq: Any) -> int: and failed[2] >= _MAX_SAVE_ATTEMPTS else 0 ) - for boundary in range(frontier, self._save_floor(exhausted), -self.chunk_size): + floor = self._save_floor(exhausted) + # Every tracked request asks every step. The answer changes only when + # its frontier moves or a checkpoint is published or dropped, so + # rescan only then instead of walking the prompt per step. + key = (frontier, floor, self._checkpoints.store.generation) + memo = getattr(seq, "_mp_save_frontier_memo", None) + if memo is not None and memo[0] == key: + return memo[1] + found = 0 + for boundary in range(frontier, floor, -self.chunk_size): if self._checkpoints.contains(self._boundary_hash(seq, boundary)): - return boundary - return 0 + found = boundary + break + seq._mp_save_frontier_memo = (key, found) + return found def _save_floor(self, lower: int) -> int: """Exclusive lower bound for boundaries worth storing (never 0).""" diff --git a/atom/kv_transfer/offload/mp/native_state_worker.py b/atom/kv_transfer/offload/mp/native_state_worker.py index cdcaccd82f..f60ee67276 100644 --- a/atom/kv_transfer/offload/mp/native_state_worker.py +++ b/atom/kv_transfer/offload/mp/native_state_worker.py @@ -143,6 +143,13 @@ def register_kv_caches( self._restore_descriptor_slots = list( range(1, max(2, self._max_pending_saves) + 1) ) + # Pinned staging allocation synchronizes the device. Reserve every + # restore slot now, at registration, rather than on the connector + # thread the first time a restore reaches it mid-serving. + builder = getattr(transfer_tensors, "state_backend", None) + reserve = getattr(builder, "reserve_checkpoint_descriptors", None) + if callable(reserve): + reserve(self._restore_descriptor_slots) logger.info( "LMCache MP native state registered rank=%d native_image=%d " "units=%d groups=%d chunk=%d", diff --git a/atom/model_engine/page_unit_checkpoint.py b/atom/model_engine/page_unit_checkpoint.py index 3416d3e647..acaa147788 100644 --- a/atom/model_engine/page_unit_checkpoint.py +++ b/atom/model_engine/page_unit_checkpoint.py @@ -164,6 +164,9 @@ def __init__( # hold 127 blocks per checkpoint out of the pool forever. self._offload_sink = offload_sink self.hash_to_checkpoint: dict[int, int] = {} + # Bumped whenever the READY set can change (publish, unindex, evict, + # release, clear), so callers can cache `contains` answers. + self.generation = 0 self.records: dict[int, CheckpointRecord] = {} self._pending_by_hash: dict[int, int] = {} self._lru: OrderedDict[int, None] = OrderedDict() @@ -450,6 +453,7 @@ def adopt_units( self.records[checkpoint_id] = record self.hash_to_checkpoint[record.prefix_hash] = checkpoint_id self._lru[checkpoint_id] = None + self.generation += 1 return True def complete_inflight(self) -> None: @@ -471,6 +475,7 @@ def complete_inflight(self) -> None: continue record.state = READY self.hash_to_checkpoint[record.prefix_hash] = checkpoint_id + self.generation += 1 self._lru[checkpoint_id] = None self._queue_offload_store(checkpoint_id, record) @@ -787,6 +792,7 @@ def _release_restore_pin(self, checkpoint_id: int) -> None: self._release_record(checkpoint_id) def unindex(self, prefix_hash: int) -> bool: + self.generation += 1 checkpoint_id = self.hash_to_checkpoint.pop(prefix_hash, -1) if checkpoint_id < 0: checkpoint_id = self._pending_by_hash.pop(prefix_hash, -1) @@ -803,6 +809,7 @@ def unindex(self, prefix_hash: int) -> bool: return True def clear(self) -> None: + self.generation += 1 self.hash_to_checkpoint.clear() self._pending_by_hash.clear() self._lru.clear() @@ -817,6 +824,7 @@ def _evict(self, checkpoint_id: int) -> None: record = self.records[checkpoint_id] if record.state != READY or record.pin_count: raise AssertionError("only an unpinned READY checkpoint is evictable") + self.generation += 1 if self.hash_to_checkpoint.get(record.prefix_hash) == checkpoint_id: del self.hash_to_checkpoint[record.prefix_hash] record.state = EVICTING @@ -825,6 +833,7 @@ def _evict(self, checkpoint_id: int) -> None: self.evictions += 1 def _release_record(self, checkpoint_id: int) -> None: + self.generation += 1 record = self.records.pop(checkpoint_id) self._lru.pop(checkpoint_id, None) if self.hash_to_checkpoint.get(record.prefix_hash) == checkpoint_id: diff --git a/atom/model_ops/attentions/backends.py b/atom/model_ops/attentions/backends.py index fac3a559a8..fb0fd16555 100644 --- a/atom/model_ops/attentions/backends.py +++ b/atom/model_ops/attentions/backends.py @@ -355,6 +355,16 @@ def execute_paged_state_copies( f"{type(self).__name__} does not implement PAGE-backed state copy" ) + def reserve_checkpoint_descriptors(self, descriptor_slots) -> None: + """Allocate staging for out-of-band copy slots before serving. + + Allocating pinned memory synchronizes, so a slot first touched by a + mid-serving restore would stall its thread. A connector that will pass + these slots to `execute_paged_state_copies` reserves them up front. + """ + for descriptor_slot in descriptor_slots: + self._checkpoint_descriptor_buffer(descriptor_slot) + def _checkpoint_descriptor_device(self) -> torch.device: """Device of the planes `execute_paged_state_copies` copies between.""" raise NotImplementedError( diff --git a/atom/model_ops/attentions/deepseek_v41/backend.py b/atom/model_ops/attentions/deepseek_v41/backend.py index 79199d510b..f47181f81f 100644 --- a/atom/model_ops/attentions/deepseek_v41/backend.py +++ b/atom/model_ops/attentions/deepseek_v41/backend.py @@ -263,6 +263,10 @@ def relocate_state_slots(self, pairs): def execute_paged_state_copies(self, stores, restores, descriptor_slot=0): self.copies.execute(stores, restores, descriptor_slot=descriptor_slot) + def reserve_checkpoint_descriptors(self, descriptor_slots): + for descriptor_slot in descriptor_slots: + self.copies.staging(descriptor_slot) + def warmup_per_req_cache(self): self.copies.warmup() diff --git a/tests/test_lmcache_mp_native_scheduler.py b/tests/test_lmcache_mp_native_scheduler.py index 792c2ac314..b79828201c 100644 --- a/tests/test_lmcache_mp_native_scheduler.py +++ b/tests/test_lmcache_mp_native_scheduler.py @@ -273,6 +273,29 @@ def test_boundary_hash_seeds_from_cache_seed_like_block_manager(monkeypatch): assert scheduler._boundary_hash(chained, 8) == fallback +def test_save_frontier_rescans_only_when_frontier_or_checkpoints_change( + monkeypatch, +): + scheduler, checkpoints, _ = make_scheduler(monkeypatch) + seq = sequence(computed=24) + checkpoint(scheduler, checkpoints, seq, 8) + calls = [] + contains = checkpoints.contains + monkeypatch.setattr( + checkpoints, "contains", lambda h: calls.append(h) or contains(h) + ) + + assert scheduler._save_frontier(seq) == 8 + scanned = len(calls) + assert scanned > 0 + assert scheduler._save_frontier(seq) == 8 + assert len(calls) == scanned # same frontier, same checkpoints: cached + + checkpoint(scheduler, checkpoints, seq, 16) # a newer checkpoint publishes + assert scheduler._save_frontier(seq) == 16 + assert len(calls) > scanned + + def test_default_min_save_tokens_skips_a_short_prompt(monkeypatch): scheduler, checkpoints, _ = make_scheduler(monkeypatch, min_save_tokens=None) seq = sequence(computed=24) diff --git a/tests/test_lmcache_mp_native_worker.py b/tests/test_lmcache_mp_native_worker.py index df5014615e..948e8059ea 100644 --- a/tests/test_lmcache_mp_native_worker.py +++ b/tests/test_lmcache_mp_native_worker.py @@ -462,6 +462,62 @@ def __exit__(self, *_args): assert compute.waited_events == [event] +def test_registration_reserves_every_restore_descriptor_slot(monkeypatch): + """Pinned staging allocation synchronizes; the restore slots are reserved + while registering, not on the connector thread mid-serving.""" + import sys + import types + + from atom.kv_transfer.offload.mp import native_state_layout, native_state_worker + + parallel_state = types.ModuleType("aiter.dist.parallel_state") + parallel_state.get_tp_group = lambda: SimpleNamespace(rank_in_group=0) + monkeypatch.setitem(sys.modules, "aiter.dist.parallel_state", parallel_state) + adapter = SimpleNamespace( + lmcache_tokens_per_chunk=8, + register_kv_caches=lambda tensors, engine_group_infos: None, + shutdown=lambda: None, + ) + monkeypatch.setattr( + native_state_worker, "_make_worker_adapter", lambda *_a, **_k: adapter + ) + monkeypatch.setattr( + native_state_worker, "require_native_state_server", lambda *_a: None + ) + monkeypatch.setattr( + native_state_layout.NativeStateMPLayout, "engine_group_infos", lambda _s: [] + ) + monkeypatch.setattr(native_state_worker.torch.cuda, "Stream", lambda **_k: None) + monkeypatch.setattr(native_state_worker.torch.cuda, "current_device", lambda: 0) + + reserved = [] + page = torch.zeros((32, 1, 32), dtype=torch.uint8) + tensors = KVTransferTensors( + block_regions=[ + KVTransferRegion( + base_addr=page.data_ptr(), unit_bytes=32, total_bytes=page.numel() + ) + ], + slot_regions=[], + block_tensor_views=[page], + paged_state_checkpoint_spec=PagedStateCheckpointSpec( + 32, 128, "native-test-v1", 80 + ), + execute_paged_state_copies=lambda *_: None, + ) + tensors.set_block_count(32) + tensors.state_backend = SimpleNamespace( + reserve_checkpoint_descriptors=lambda slots: reserved.extend(slots) + ) + connector = NativeStateLMCacheMPConnector( + config(**{"lmcache.mp.tp_rank_collapse": False}) + ) + connector.register_kv_caches({}, tensors, 32) + + assert reserved == connector._restore_descriptor_slots + assert reserved and 0 not in reserved + + def test_native_server_chunk_mismatch_fails_before_registration(monkeypatch): from atom.kv_transfer.offload.mp import native_state_worker From 2f0d60605c0bc3f39a96fd6c206f116b7ac665a2 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Sun, 27 Sep 2026 04:09:22 +0000 Subject: [PATCH 12/28] docs(offload): document every offload env var; guard it with a test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups (#2250 §8, §9). - docs/environment_variables.md listed only OFFLOAD_MAX_PENDING_SAVES and the LMCache pin timeout, and still said they bypass atom.utils.envs. All 16 ATOM-owned offload knobs are now documented with type, default, precedence and invalid-value behavior. A new test fails if an OFFLOAD_*/LMCACHE_* var registered in envs.py is missing from the reference. - README and envs.py describe OFFLOAD_MIN_SAVE_TOKENS as the native absolute boundary for normal and late saves, and the README documents the uncertain-transfer bound. - Test that a restore raising after it took a descriptor slot fails the load and returns the slot once the uncertainty bound expires. Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/kv_transfer/offload/README.md | 8 +++++-- atom/utils/envs.py | 6 +++--- docs/environment_variables.md | 30 ++++++++++++++++++++------ tests/test_envs.py | 19 ++++++++++++++++ tests/test_lmcache_mp_native_worker.py | 26 ++++++++++++++++++++++ 5 files changed, 77 insertions(+), 12 deletions(-) diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index fe78bee88e..86e939b4cd 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -186,8 +186,12 @@ prefix index and stores only the still-resident contiguous prefix. Native MP never stores a prefix shorter than `OFFLOAD_MIN_SAVE_TOKENS`: with the default equal to `OFFLOAD_MIN_LOAD_TOKENS` it could never be loaded back, and on a 1K-token workload skipping those saves took the offload throughput -cost from 12% to within run-to-run noise. The same threshold also suppresses a -late save whose remaining prefix is too small. Unless configured otherwise, the shared save limit is +cost from 12% to within run-to-run noise. It is an absolute boundary for normal +and late saves alike, so the short tail of a long request is still stored. +Other connectors ignore it. A transfer whose outcome is unknown (its submission +or restore raised) keeps its lease for `lmcache.mp.uncertain_transfer_timeout_s` +(default twice `lmcache.mp.mq_timeout`, i.e. 600 s), then fails so the lease +and budget are released. Unless configured otherwise, the shared save limit is `max(2, 2 * OFFLOAD_COPY_WORKERS)`. The model namespace includes PAGE/model geometry, TP and speculation settings, diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 1f9f79aaa0..c2a7f13006 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -889,9 +889,9 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float: # derives max(2, 2 * OFFLOAD_COPY_WORKERS) and the scheduler's state tier # uses 2. A kv_connector_extra_config "max_pending_saves" takes precedence. "OFFLOAD_MAX_PENDING_SAVES": lambda: _optional_int_env("OFFLOAD_MAX_PENDING_SAVES"), - # Minimum external-tier hit worth loading, and minimum prefix worth saving - # (native MP's shortest stored boundary and a late save's remainder), in - # tokens. + # Minimum external-tier hit worth loading, in tokens. And the shortest + # prefix native lmcache_mp stores: an absolute boundary, for normal and + # late saves alike (other connectors ignore it). "OFFLOAD_MIN_LOAD_TOKENS": lambda: max( 0, _int_env_or_default("OFFLOAD_MIN_LOAD_TOKENS", 8192) ), diff --git a/docs/environment_variables.md b/docs/environment_variables.md index 62963b04da..78ca5c6cd5 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -288,16 +288,32 @@ flag below. Details in the state-checkpoint section of the ### LMCache offload tier -The LMCache source-pin timeout and ATOM's pending-save bound are read directly -via `os.environ` rather than through `atom.utils.envs`. Their behavior is -defined in `atom/kv_transfer/offload/_offload_common.py` and documented in full -in `atom/kv_transfer/offload/README.md`; they are also listed here so they are -discoverable from the central env reference despite bypassing the registry. +ATOM's own offload knobs are defined in `atom/utils/envs.py`; their behavior is +documented in `atom/kv_transfer/offload/README.md`. Where a +`kv_connector_extra_config` key also exists, it takes precedence over the env +var. `LMCACHE_EC_PIN_TIMEOUT_SEC` belongs to LMCache and is read only to +derive a bound. | Variable | Type | Default | Description | |----------|------|---------|-------------| -| **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_COPY_WORKERS** | int | 1 | Save executor threads per offload worker. Also scales the default `OFFLOAD_MAX_PENDING_SAVES`. | +| **OFFLOAD_LOAD_WORKERS** | int | 1 | Load executor threads per offload worker. DSV4's in-process path ignores it (its SLOT load path needs a serial load executor). | +| **OFFLOAD_MAX_PENDING_SAVES** | int | unset: `max(2, 2 × OFFLOAD_COPY_WORKERS)` for connectors, **2** for the scheduler's state tier | Bound on running-plus-queued saves. KV and state saves share it because both pin the same pool. A non-integer raises on the connector path and warns (using 2) on the state-tier path. Overridden by `max_pending_saves` in `kv_connector_extra_config`. | +| **OFFLOAD_MIN_LOAD_TOKENS** | int | 8192 | Smallest external-tier hit worth loading; shorter hits are recomputed. Negative values clamp to 0. | +| **OFFLOAD_MIN_SAVE_TOKENS** | int | 8192 | Shortest prefix native `lmcache_mp` stores, as an absolute boundary for normal and late saves. With the default equal to `OFFLOAD_MIN_LOAD_TOKENS`, a shorter prefix could never be loaded back. Other connectors ignore it. | +| **OFFLOAD_LOOKUP_MEMO_STEPS** | int | 32 | Scheduler steps a memoised tier-lookup answer is replayed before the tier is asked again. | +| **OFFLOAD_LOOKUP_RETRY_STEPS** | int | 32 | Scheduler steps a failed tier lookup suppresses the next attempt. | +| **OFFLOAD_PROFILE** | bool | 0 | Emit `[OFFLOAD-SAVE-PROF]` / `[OFFLOAD-LOAD-PROF]` per-transfer records. An empty value reads as off. | +| **OFFLOAD_SINGLE_STREAM** | bool | 0 | Experimental: run the staging pack and copy legs on one stream. | +| **OFFLOAD_GPU_STAGING_CHUNKS** | int | derived from KV geometry | GPU staging buffer size in LMCache chunks (≥ 1). | +| **OFFLOAD_GPU_STAGING_MAX_BYTES** | int | unset | Upper bound on the GPU staging buffer in bytes; must hold at least one chunk. | +| **OFFLOAD_RELEASE_GPU_STAGING_AFTER_TRANSFER** | bool | 0 | Free the GPU staging buffer after each transfer instead of keeping it. | +| **OFFLOAD_SLOT_STAGING_SLOTS** | int | 1 | DSV4 in-process SLOT sidecar staging rows. Overridden by `slot_sidecar_staging_slots`. | +| **OFFLOAD_COMMITTED_SIDECAR_CAPACITY** | int | 65536 | DSV4 in-process committed SLOT sidecar index capacity. Overridden by `committed_sidecar_index_capacity`. | +| **OFFLOAD_PUBLICATION_TIMEOUT_S** | float | 5.0 | DSV4 in-process wait for a saved SLOT sidecar to become visible (finite, ≥ 0). | +| **OFFLOAD_PUBLICATION_POLL_INTERVAL_S** | float | 0.01 | Poll period of that wait (finite, > 0). | +| **LMCACHE_MP_TRANSFER_MODE** | str | auto | `lmcache_mp` transfer mode: `auto` or `lmcache_driven` (`engine_driven` is rejected). Overridden by `lmcache.mp.mp_transfer_mode`. | +| **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 and, when unset, assumes LMCache's. | ## KV cache events diff --git a/tests/test_envs.py b/tests/test_envs.py index bd2ce7a524..3b8d880fee 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -277,3 +277,22 @@ def test_mla_fp8_prefill_flag(monkeypatch): for value, expected in [("0", False), ("1", True), ("true", False)]: monkeypatch.setenv(name, value) assert getattr(_get_envs(), name) is expected + + +def test_offload_env_vars_are_documented(): + """Every offload knob registered in envs.py appears in the central env + reference, so a new one cannot land undocumented.""" + import pathlib + + from atom.utils import envs + + doc = ( + pathlib.Path(__file__).parents[1] / "docs" / "environment_variables.md" + ).read_text() + offload = [ + name + for name in envs.environment_variables + if name.startswith(("OFFLOAD_", "LMCACHE_")) + ] + assert offload + assert [name for name in offload if f"**{name}**" not in doc] == [] diff --git a/tests/test_lmcache_mp_native_worker.py b/tests/test_lmcache_mp_native_worker.py index 948e8059ea..f9833bacb1 100644 --- a/tests/test_lmcache_mp_native_worker.py +++ b/tests/test_lmcache_mp_native_worker.py @@ -518,6 +518,32 @@ def test_registration_reserves_every_restore_descriptor_slot(monkeypatch): assert reserved and 0 not in reserved +def test_raising_restore_returns_its_descriptor_slot_after_the_bound( + worker, monkeypatch +): + """A restore that raises after taking a descriptor slot keeps it only until + the uncertainty bound, then the load fails and the slot is reusable.""" + from atom.kv_transfer.offload.mp import backend + + now = [1000.0] + monkeypatch.setattr(backend.time, "monotonic", lambda: now[0]) + worker._restore_stream = object() # no wait_stream: raises inside + worker._restore_descriptor_slots = [3] + worker.future.value, worker.future.ready = True, True + req = request(loading=True) + worker._submit_load(req, object()) + + first = worker.get_finished() + assert not first.finished_loading and not first.failed_loading + assert worker._restore_descriptor_slots == [] + + now[0] += worker._uncertain_timeout_s + 1 + finished = worker.get_finished() + assert finished.failed_loading == {req.load_operation} + assert worker._restore_descriptor_slots == [3] + assert worker._native_loads == {} + + def test_native_server_chunk_mismatch_fails_before_registration(monkeypatch): from atom.kv_transfer.offload.mp import native_state_worker From 94c07412945db7100826ea605439bb0e009d6aa3 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Sun, 27 Sep 2026 04:48:18 +0000 Subject: [PATCH 13/28] test(lmcache_mp): make the raising-restore test independent of CUDA On a CPU-only torch, torch.cuda.Event is a dummy class that raises when it is built, which happens before the restore takes its descriptor slot, so the test's "slot is held until the bound" assertion failed in CI. Stub the Event so the restore fails deterministically at the stream fence, after the slot is taken, on every runner. Co-Authored-By: Claude Opus 5.5 (1M context) --- tests/test_lmcache_mp_native_worker.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/test_lmcache_mp_native_worker.py b/tests/test_lmcache_mp_native_worker.py index f9833bacb1..13e2bc25fd 100644 --- a/tests/test_lmcache_mp_native_worker.py +++ b/tests/test_lmcache_mp_native_worker.py @@ -523,10 +523,15 @@ def test_raising_restore_returns_its_descriptor_slot_after_the_bound( ): """A restore that raises after taking a descriptor slot keeps it only until the uncertainty bound, then the load fails and the slot is reusable.""" - from atom.kv_transfer.offload.mp import backend + from atom.kv_transfer.offload.mp import backend, native_state_worker now = [1000.0] monkeypatch.setattr(backend.time, "monotonic", lambda: now[0]) + # A real Event cannot be built on a CPU runner; fail deterministically + # after the slot is taken instead, at the stream fence. + monkeypatch.setattr( + native_state_worker.torch.cuda, "Event", lambda: SimpleNamespace() + ) worker._restore_stream = object() # no wait_stream: raises inside worker._restore_descriptor_slots = [3] worker.future.value, worker.future.ready = True, True From 83f515a89be74b2051239410a989738b38d20fda Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Mon, 28 Sep 2026 03:28:33 +0000 Subject: [PATCH 14/28] fix(lmcache_mp): release transfer memory only on a report, fail stop at a deadline The MP path answered "what settles a transfer whose outcome is unprovable" three ways: never (native checkpoint pins, a zeroed abandon timeout), after 600 s (the uncertainty bound, which then freed load destinations and PAGE sources a server might still be using), and never-unbounded (a future whose poll keeps raising). One rule now: memory under an MP transfer is released only on a terminal report, and a transfer still not terminal after `lmcache.mp.transfer_deadline_s` (default 1200 s) raises `LMCacheTransferUnprovable`, stopping the engine. - Worker: every pending PAGE-only and native load/save carries its start time and is checked against the deadline, which bounds a raising submission, a future whose poll keeps raising, and a restore whose event cannot be queried. A descriptor that cannot be built is provably unsent and still fails immediately; only the transport call is unprovable. - Scheduler: a watchdog over dispatched saves, loads and native checkpoint sources catches a completion that never arrives (lost report, silent TP rank), with a margin so the worker fails first. - `save_abandon_timeout_s` is inherited again, so the engine-wide stalled save, state-pin and orphan-load-slot reclaimers stay on; MP opts out through its own `abandon_save` / `reclaim_stale_leases` instead. Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/kv_transfer/offload/README.md | 18 +- atom/kv_transfer/offload/mp/backend.py | 205 +++++++++++++----- .../offload/mp/native_state_scheduler.py | 18 +- .../offload/mp/native_state_worker.py | 38 ++-- tests/test_lmcache_mp.py | 134 ++++++++---- tests/test_lmcache_mp_native_scheduler.py | 47 +++- tests/test_lmcache_mp_native_worker.py | 80 +++---- 7 files changed, 367 insertions(+), 173 deletions(-) diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index 86e939b4cd..2f5e1c95bd 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -188,11 +188,19 @@ the default equal to `OFFLOAD_MIN_LOAD_TOKENS` it could never be loaded back, and on a 1K-token workload skipping those saves took the offload throughput cost from 12% to within run-to-run noise. It is an absolute boundary for normal and late saves alike, so the short tail of a long request is still stored. -Other connectors ignore it. A transfer whose outcome is unknown (its submission -or restore raised) keeps its lease for `lmcache.mp.uncertain_transfer_timeout_s` -(default twice `lmcache.mp.mq_timeout`, i.e. 600 s), then fails so the lease -and budget are released. Unless configured otherwise, the shared save limit is -`max(2, 2 * OFFLOAD_COPY_WORKERS)`. +Other connectors ignore it. Unless configured otherwise, the shared save limit +is `max(2, 2 * OFFLOAD_COPY_WORKERS)`. + +Engine memory under an MP transfer (PAGE sources, restore destinations, +checkpoint pins, descriptor slots) is released only on a terminal report from +the server. A timeout cannot prove that a remote DMA stopped, so no clock ever +frees it. Instead, a transfer that is still not terminal after +`lmcache.mp.transfer_deadline_s` (default 1200 s) is fatal: the worker, or the +scheduler for a report that never arrives, raises `LMCacheTransferUnprovable` +and the engine stops. This covers a submission that raised, a future whose +poll keeps raising, a restore whose event cannot be queried, and a lost or +never-aggregated completion. The engine-wide stalled-save reclaim leaves MP +leases alone for the same reason. The model namespace includes PAGE/model geometry, TP and speculation settings, native layout/image sizes, the Hugging Face commit when available, and the diff --git a/atom/kv_transfer/offload/mp/backend.py b/atom/kv_transfer/offload/mp/backend.py index 2d5d21ba1c..3c796f697d 100644 --- a/atom/kv_transfer/offload/mp/backend.py +++ b/atom/kv_transfer/offload/mp/backend.py @@ -18,7 +18,7 @@ import time from collections import deque from collections.abc import Iterable, Sequence -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace from typing import Any import torch @@ -487,6 +487,7 @@ class _LookupState: class _PendingLoad: completion: LoadCompletionId future: Any | None + started_at: float = field(default_factory=lambda: time.monotonic()) @dataclass @@ -495,51 +496,67 @@ class _PendingSave: future: Any | None start: int end: int + started_at: float = field(default_factory=lambda: time.monotonic()) -def _uncertain_transfer_timeout_s(config: Any) -> float: - """How long an unprovable transfer keeps its lease before it is failed. +_DEFAULT_TRANSFER_DEADLINE_S = 1200.0 +# The scheduler watches the same operations from dispatch, a step or more +# before a worker submits them. The margin lets the worker, which can name the +# exact transfer, fail first. +_SCHEDULER_DEADLINE_MARGIN_S = 60.0 - Nothing in the MP protocol can prove that a submission whose outcome is - unknown has stopped touching engine memory, so the bound is a policy: by - default twice the adapter's own request timeout. + +class LMCacheTransferUnprovable(RuntimeError): + """An MP transfer never reached a terminal state within the deadline.""" + + +def _transfer_deadline_s(config: Any) -> float: + """Seconds any MP transfer may stay non-terminal before the engine stops. + + Engine memory under a transfer is released only on a terminal report: a + timeout cannot prove that the server stopped its DMA, so freeing on a + clock could hand live blocks to another request. A transfer that never + reports is therefore fatal once this deadline passes -- fail-stop, rather + than corrupt memory or wedge the pool with no fault to point at. """ extra = _extra_config(config) - default = 2.0 * float(extra.get("lmcache.mp.mq_timeout", 300.0)) - timeout = float(extra.get("lmcache.mp.uncertain_transfer_timeout_s", default)) - if not timeout > 0: - raise ValueError("lmcache.mp.uncertain_transfer_timeout_s must be > 0") - return timeout + deadline = float( + extra.get("lmcache.mp.transfer_deadline_s", _DEFAULT_TRANSFER_DEADLINE_S) + ) + if not deadline > 0: + raise ValueError("lmcache.mp.transfer_deadline_s must be > 0") + return deadline -class _UncertainSubmission: - """A transfer whose outcome is unknown: retained, then failed at a deadline. +def _enforce_transfer_deadline( + operation: Any, started_at: float, deadline_s: float +) -> None: + """Raise ``LMCacheTransferUnprovable`` once a transfer outlives the bound.""" - A transport exception cannot prove that a remote DMA stopped, so the lease - is kept while this reports not-done. Past ``timeout_s`` it reports a failed - terminal result, so the operation settles and releases its lease and budget - instead of wedging them for the process lifetime. - """ + elapsed = time.monotonic() - started_at + if elapsed < deadline_s: + return + message = ( + f"LMCache MP transfer {operation} is still not terminal after " + f"{elapsed:.0f}s. The memory it reads or writes cannot be released " + "without proof that the server stopped, so the engine is stopping. " + "Raise lmcache.mp.transfer_deadline_s if the tier is only slow." + ) + logger.error(message) + raise LMCacheTransferUnprovable(message) - def __init__(self, timeout_s: float, operation: str) -> None: - self._deadline = time.monotonic() + timeout_s - self._timeout_s = timeout_s - self._operation = operation - self._expired = False + +class _UnprovableSubmission: + """Stand-in future for a submission whose transport call raised. + + The server may have taken the request before the connection failed, so + this never becomes terminal on its own: the lease is kept and only the + transfer deadline ends it. + """ def query(self) -> bool: - if time.monotonic() < self._deadline: - return False - if not self._expired: - self._expired = True - logger.warning( - "LMCache MP transfer %s stayed unprovable for %.0fs; " - "treating it as failed and releasing its lease", - self._operation, - self._timeout_s, - ) - return True + return False def result(self, timeout: float | None = None) -> bool: del timeout @@ -557,7 +574,8 @@ def _terminal_future_result(future: Any | None) -> tuple[bool, Any]: return True, future.result(timeout=0) except Exception: # A query/result exception is not proof that the remote GPU stream has - # quiesced. Keep the operation pending rather than risk block reuse. + # quiesced. Keep the operation pending; the transfer deadline bounds a + # future that keeps raising. logger.warning( "LMCache MP transfer future could not prove terminal state", exc_info=True, @@ -799,7 +817,7 @@ def __init__(self, config: Any) -> None: # aggregator must receive the same chunk completion from every rank # before releasing the writer's source blocks early. self._immediate_saves: dict[SaveCompletionId, tuple[int, int] | None] = {} - self._uncertain_timeout_s = _uncertain_transfer_timeout_s(config) + self._transfer_deadline_s = _transfer_deadline_s(config) self._immediate_load_failures: set[LoadCompletionId] = set() self._lock = threading.Lock() @@ -970,14 +988,10 @@ def _submit_load(self, req: LMCacheReqMeta, event: Any) -> None: start=start, end=end, ) - transfer = self._adapter.submit_retrieve_request( - request_id, - op, - event, - ) except Exception: + # Nothing was sent: the failure is provable and terminal. logger.exception( - "LMCache MP load submission failed for %s", + "Invalid LMCache MP load descriptor for %s", req.req_id, ) with self._lock: @@ -989,6 +1003,20 @@ def _submit_load(self, req: LMCacheReqMeta, event: Any) -> None: ) self._immediate_load_failures.add(completion) return + try: + transfer = self._adapter.submit_retrieve_request( + request_id, + op, + event, + ) + except Exception: + # The server may have taken the request before the connection + # raised, and may still write the destination blocks. + logger.exception( + "LMCache MP load submission unprovable for %s", + req.req_id, + ) + transfer = _UnprovableSubmission() with self._lock: self._submitting_loads.discard(operation_id) self._pending_loads[operation_id] = _PendingLoad( @@ -1047,6 +1075,22 @@ def _submit_save(self, req: LMCacheReqMeta, event: Any) -> None: start=start, end=end, ) + except Exception: + # Nothing was sent: report a terminal failure so the save settles. + logger.exception( + "Invalid LMCache MP save descriptor for %s", + req.req_id, + ) + with self._lock: + self._submitting_saves.discard(operation_id) + self._pending_saves[operation_id] = _PendingSave( + completion=completion, + future=None, + start=start, + end=end, + ) + return + try: submit = getattr( self._adapter, "submit_store_request_with_chunk_events", @@ -1059,23 +1103,13 @@ def _submit_save(self, req: LMCacheReqMeta, event: Any) -> None: ) except Exception: # The server may have received the request before the connection - # raised: keep the source leased until the transfer is provably - # over or the uncertainty bound expires. + # raised: keep the source leased until a terminal report, which + # the transfer deadline turns into a fail-stop if it never comes. logger.exception( - "LMCache MP save submission uncertain for %s", + "LMCache MP save submission unprovable for %s", req.req_id, ) - with self._lock: - self._submitting_saves.discard(operation_id) - self._pending_saves[operation_id] = _PendingSave( - completion=completion, - future=_UncertainSubmission( - self._uncertain_timeout_s, operation_id - ), - start=start, - end=end, - ) - return + transfer = _UnprovableSubmission() with self._lock: self._submitting_saves.discard(operation_id) self._pending_saves[operation_id] = _PendingSave( @@ -1113,6 +1147,10 @@ def get_finished(self) -> KVConnectorOutput: exc_info=True, ) terminal, result = _terminal_future_result(pending.future) + if not terminal: + _enforce_transfer_deadline( + operation_id, pending.started_at, self._transfer_deadline_s + ) if terminal: self._pending_saves.pop(operation_id, None) _remember_operation_tombstone( @@ -1139,6 +1177,9 @@ def get_finished(self) -> KVConnectorOutput: for operation_id, pending in list(self._pending_loads.items()): terminal, result = _terminal_future_result(pending.future) if not terminal: + _enforce_transfer_deadline( + operation_id, pending.started_at, self._transfer_deadline_s + ) continue self._pending_loads.pop(operation_id, None) _remember_operation_tombstone( @@ -1198,6 +1239,10 @@ def __init__(self, config: Any, *, checkpoint_spec: Any = None) -> None: poll_interval=poll_interval, ) self._mp_adapter = adapter + self._scheduler_deadline_s = ( + _transfer_deadline_s(config) + _SCHEDULER_DEADLINE_MARGIN_S + ) + self._transfer_seen_at: dict[Any, float] = {} super().__init__( config, chunk_size=int(adapter.lmcache_tokens_per_chunk), @@ -1209,10 +1254,6 @@ def __init__(self, config: Any, *, checkpoint_spec: Any = None) -> None: shutdown() raise - def save_abandon_timeout_s(self) -> float: - """A timeout cannot prove that a remote MP DMA stopped reading HBM.""" - return 0.0 - def get_num_new_matched_tokens(self, seq: Any) -> tuple[int, bool]: matched = super().get_num_new_matched_tokens(seq) sid = str(seq.id) @@ -1291,6 +1332,49 @@ def load_failed(self, req_id: Any) -> bool: ) return super().load_failed(req_id) + def abandon_save(self, req_id: Any) -> None: + # The engine's clock-based reclaim cannot prove the MP server stopped + # reading the source, so an MP save keeps its lease until a terminal + # report. `_enforce_transfer_deadlines` bounds a report that never comes. + logger.warning( + "LMCache MP keeps the source of request %s leased until its save " + "reports; a transfer that never reports stops the engine after %.0fs", + req_id, + self._scheduler_deadline_s, + ) + + def reclaim_stale_leases(self, timeout_s: float) -> list[frozenset]: + del timeout_s # See abandon_save: no release without a report. + return [] + + def _live_transfers(self) -> set[Any]: + """Dispatched operations still waiting for a terminal worker report.""" + live: set[Any] = set(self._save_inflight.values()) + live.update(self._save_operation_owner) + live.update(operation for _, operation in self._active_load_operations.values()) + return live + + def _enforce_transfer_deadlines(self) -> None: + """Fail-stop when a dispatched transfer never reports. + + The worker bounds what it submitted; this catches what it cannot see, + such as a lost completion or a TP rank that never reports. + """ + live = self._live_transfers() + seen = self._transfer_seen_at + for operation in [op for op in seen if op not in live]: + del seen[operation] + now = time.monotonic() + for operation in live: + _enforce_transfer_deadline( + operation, seen.setdefault(operation, now), self._scheduler_deadline_s + ) + + def process_completions(self, output: KVConnectorOutput) -> KVConnectorOutput: + output = super().process_completions(output) + self._enforce_transfer_deadlines() + return output + def request_finished(self, seq: Any) -> None: super().request_finished(seq) self._pending_session_ends()[str(seq.id)] = seq @@ -1329,4 +1413,5 @@ def _end_finished_sessions(self) -> None: __all__ = [ "LMCacheMPConnector", "LMCacheMPConnectorScheduler", + "LMCacheTransferUnprovable", ] diff --git a/atom/kv_transfer/offload/mp/native_state_scheduler.py b/atom/kv_transfer/offload/mp/native_state_scheduler.py index 6119a00c72..7fd5a294fd 100644 --- a/atom/kv_transfer/offload/mp/native_state_scheduler.py +++ b/atom/kv_transfer/offload/mp/native_state_scheduler.py @@ -537,13 +537,17 @@ def process_completions(self, output): self._finish_retired_request(sid) return output - def abandon_save(self, req_id: Any) -> None: - # This legacy callback carries no proof that the MP reader stopped. - # Explicit terminal failures already settle the exact operation above. - return None - - def reclaim_stale_leases(self, timeout_s: float) -> list[frozenset]: - return [] + def _live_transfers(self) -> set[Any]: + live = super()._live_transfers() + # A native save holds its checkpoint pin and state budget until its + # STORE terminal; nothing but that report may release them. + live.update(self._native_saves) + live.update( + operation + for operation, lease in self._native_loads.items() + if lease.dispatched + ) + return live __all__ = ["NativeStateLMCacheMPConnectorScheduler"] diff --git a/atom/kv_transfer/offload/mp/native_state_worker.py b/atom/kv_transfer/offload/mp/native_state_worker.py index f60ee67276..6dcec709b7 100644 --- a/atom/kv_transfer/offload/mp/native_state_worker.py +++ b/atom/kv_transfer/offload/mp/native_state_worker.py @@ -6,7 +6,8 @@ from __future__ import annotations import logging -from dataclasses import dataclass +import time +from dataclasses import dataclass, field from typing import Any import torch @@ -21,6 +22,7 @@ from atom.kv_transfer.offload.mp.backend import ( LMCacheMPConnector, _chunk_ranges, + _enforce_transfer_deadline, _make_worker_adapter, _mp_session_id, _published_tp_replication_factor, @@ -30,7 +32,7 @@ _terminal_future_result, _tp_replication_factor, _transfer_operation_id, - _UncertainSubmission, + _UnprovableSubmission, _validate_mp_config, ) from atom.kv_transfer.offload.mp.native_state_layout import ( @@ -65,6 +67,7 @@ class _NativePending: descriptor_slot: int | None = None immediate_success: bool = False state_source_safe: bool = False + started_at: float = field(default_factory=lambda: time.monotonic()) class NativeStateLMCacheMPConnector(LMCacheMPConnector): @@ -242,9 +245,10 @@ def _submit_native(self, req: LMCacheReqMeta, event: Any, *, loading: bool) -> N ) pending[operation_id] = _NativePending(req, None) return - entry = _NativePending( - req, _UncertainSubmission(self._uncertain_timeout_s, operation_id) - ) + # Installed before the transport call: if it raises, the server may + # still have taken the request, so the lease stays until the + # transfer deadline. + entry = _NativePending(req, _UnprovableSubmission()) pending[operation_id] = entry submit = ( self._adapter.submit_retrieve_request @@ -261,7 +265,7 @@ def _submit_native(self, req: LMCacheReqMeta, event: Any, *, loading: bool) -> N # Retain the exact source/destination lease. The server might have # received the request before the connection raised an exception. logger.exception( - "Native-state LMCache MP submission uncertain; retaining lease %s", + "Native-state LMCache MP submission unprovable; retaining lease %s", operation_id, ) return @@ -366,6 +370,9 @@ def get_finished(self) -> KVConnectorOutput: else: terminal, result = _terminal_future_result(pending.future) if not terminal: + _enforce_transfer_deadline( + operation_id, pending.started_at, self._transfer_deadline_s + ) continue save_spec = pending.request.save_spec start = 0 if save_spec is None else int(save_spec.skip_leading_tokens) @@ -390,6 +397,9 @@ def get_finished(self) -> KVConnectorOutput: if pending.restore_event is None: terminal, result = _terminal_future_result(pending.future) if not terminal: + _enforce_transfer_deadline( + operation_id, pending.started_at, self._transfer_deadline_s + ) continue if result is True: try: @@ -400,25 +410,25 @@ def get_finished(self) -> KVConnectorOutput: "LMCache MP native restore safety unknown; " "retaining lease" ) - pending.future = _UncertainSubmission( - self._uncertain_timeout_s, operation_id - ) + pending.future = _UnprovableSubmission() pending.restore_event = None continue if pending.restore_event is not None: try: - if not pending.restore_event.query(): - continue + restored = pending.restore_event.query() except Exception: logger.exception( "LMCache MP native restore event safety unknown; " "retaining lease" ) - pending.future = _UncertainSubmission( - self._uncertain_timeout_s, operation_id - ) + pending.future = _UnprovableSubmission() pending.restore_event = None pending.restore_succeeded = False + restored = False + if not restored: + _enforce_transfer_deadline( + operation_id, pending.started_at, self._transfer_deadline_s + ) continue completion = pending.request.load_operation if pending.descriptor_slot is not None: diff --git a/tests/test_lmcache_mp.py b/tests/test_lmcache_mp.py index 7c4fc87329..03d009ec22 100644 --- a/tests/test_lmcache_mp.py +++ b/tests/test_lmcache_mp.py @@ -6,7 +6,7 @@ import sys import types from collections import deque -from dataclasses import dataclass +from dataclasses import dataclass, replace from types import SimpleNamespace import pytest @@ -1227,48 +1227,99 @@ def test_worker_pre_submit_drops_are_immediately_terminal( assert worker._pending_saves == {} -def test_worker_uncertain_save_submission_is_bounded_then_quiescent( +def _unprovable_request(kind: str) -> LMCacheReqMeta: + if kind == "load": + return LMCacheReqMeta( + req_id=16, + token_ids=list(range(16)), + block_ids=[100, 101, 102, 103], + load_spec=LoadSpec(0, 16, can_load=True), + load_operation=LoadOperationId(req_id=16, generation=1), + ) + return LMCacheReqMeta( + req_id=16, + token_ids=list(range(16)), + block_ids=[100, 101, 102, 103], + save_spec=SaveSpec(skip_leading_tokens=0), + save_operation=SaveOperationId(req_id=16, generation=1), + ) + + +@pytest.mark.parametrize("kind", ["load", "save"]) +def test_worker_unprovable_submission_is_never_released_on_a_clock( fake_lmcache_modules, monkeypatch, + kind, ): - """A raising submit may have reached the server, so the source stays leased; - past the bound the save reports every chunk source-safe and the store - failed, which lets the scheduler retry or retire it instead of wedging.""" + """A raising submit may have reached the server, which may still read the + source or write the destination. Nothing settles, so nothing is released; + past the transfer deadline the worker stops the engine instead.""" adapter = _WorkerAdapter() - def uncertain(_request_id, _op, _event): + def unprovable(_request_id, _op, _event): raise ConnectionError("server may have received request") - monkeypatch.setattr(adapter, "submit_store_request", uncertain) - monkeypatch.setattr(adapter, "submit_store_request_with_chunk_events", uncertain) + for name in ( + "submit_retrieve_request", + "submit_store_request", + "submit_store_request_with_chunk_events", + ): + monkeypatch.setattr(adapter, name, unprovable) now = [1000.0] monkeypatch.setattr(mp_connector.time, "monotonic", lambda: now[0]) worker = _worker(adapter) - operation = SaveOperationId(req_id=16, generation=1) - worker._submit_save( - LMCacheReqMeta( - req_id=16, - token_ids=list(range(16)), - block_ids=[100, 101, 102, 103], - save_spec=SaveSpec(skip_leading_tokens=0), - save_operation=operation, - ), - object(), - ) + request = _unprovable_request(kind) + if kind == "load": + worker._submit_load(request, object()) + pending = worker._pending_loads + else: + worker._submit_save(request, object()) + pending = worker._pending_saves + + now[0] += worker._transfer_deadline_s - 1 + output = worker.get_finished() + assert not output.connector_completions + assert not output.failed_loading and not output.finished_saving + assert len(pending) == 1 - assert not worker.get_finished().connector_completions - assert len(worker._pending_saves) == 1 + now[0] += 2 + with pytest.raises(mp_connector.LMCacheTransferUnprovable): + worker.get_finished() + assert len(pending) == 1 + + +def test_worker_future_that_keeps_raising_is_bounded_by_the_deadline( + fake_lmcache_modules, + monkeypatch, +): + now = [1000.0] + monkeypatch.setattr(mp_connector.time, "monotonic", lambda: now[0]) + worker = _worker(_WorkerAdapter()) + worker._submit_save(_unprovable_request("save"), object()) + [pending] = worker._pending_saves.values() + pending.future.query_error = ConnectionError("IPC context torn down") + + assert not worker.get_finished().finished_saving + now[0] += worker._transfer_deadline_s + with pytest.raises(mp_connector.LMCacheTransferUnprovable): + worker.get_finished() + + +def test_worker_invalid_descriptor_fails_before_transport( + fake_lmcache_modules, +): + """A descriptor that cannot be built is provably unsent: it settles now.""" + adapter = _WorkerAdapter() + worker = _worker(adapter) + load = replace(_unprovable_request("load"), block_ids=[100]) + save = replace(_unprovable_request("save"), block_ids=[100]) + worker._submit_load(load, object()) + worker._submit_save(save, object()) - now[0] += worker._uncertain_timeout_s + 1 output = worker.get_finished() - assert output.finished_saving == {operation} - assert worker._pending_saves == {} - safe = { - completion.operation_id.ranges - for completion in output.connector_completions - if completion.channel == mp_connector.DENSE_PAGE_SOURCE_SAFE_CHANNEL - } - assert safe == {((0, 8),), ((8, 16),)} + assert not adapter.loads and not adapter.saves + assert output.failed_loading == {load.load_operation} + assert output.finished_saving == {save.save_operation} [store] = [ completion for completion in output.connector_completions @@ -1277,6 +1328,20 @@ def uncertain(_request_id, _op, _event): assert not store.succeeded +def test_transfer_deadline_defaults_to_twenty_minutes(): + assert mp_connector._transfer_deadline_s(_config()) == 1200.0 + assert ( + mp_connector._transfer_deadline_s( + _config(extra={"lmcache.mp.transfer_deadline_s": 30}) + ) + == 30.0 + ) + with pytest.raises(ValueError, match="transfer_deadline_s"): + mp_connector._transfer_deadline_s( + _config(extra={"lmcache.mp.transfer_deadline_s": 0}) + ) + + def test_worker_save_slices_chunk_blocks_and_preserves_operation( fake_lmcache_modules, ): @@ -1641,21 +1706,16 @@ def request(generation): def test_worker_immediate_operation_tombstones_reject_replay( fake_lmcache_modules, - monkeypatch, ): adapter = _WorkerAdapter() - - def fail_load(_request_id, _op, _event): - raise RuntimeError("submission failed") - - monkeypatch.setattr(adapter, "submit_retrieve_request", fail_load) worker = _worker(adapter) load_operation = LoadOperationId(req_id=25, generation=1) save_operation = SaveOperationId(req_id=26, generation=1) + # Too few blocks for the range: provably never sent, so immediately terminal. load_request = LMCacheReqMeta( req_id=25, token_ids=list(range(8)), - block_ids=[120, 121], + block_ids=[120], load_spec=LoadSpec(0, 8, can_load=True), load_operation=load_operation, ) diff --git a/tests/test_lmcache_mp_native_scheduler.py b/tests/test_lmcache_mp_native_scheduler.py index b79828201c..c93cf0f4d1 100644 --- a/tests/test_lmcache_mp_native_scheduler.py +++ b/tests/test_lmcache_mp_native_scheduler.py @@ -344,13 +344,15 @@ def test_store_failure_rolls_back_for_bounded_retry_and_ignores_stale_reports( def test_no_timeout_or_abandon_recycles_dispatched_source(monkeypatch): + """The engine's clock-based reclaimers stay on for everything else, but + none of them may release what a dispatched MP transfer still uses.""" scheduler, checkpoints, _ = make_scheduler(monkeypatch) seq = sequence() checkpoint(scheduler, checkpoints, seq, 16) scheduler.update_state_after_alloc(seq) [request] = scheduler.build_connector_meta().requests + assert scheduler.save_abandon_timeout_s() > 0 scheduler.abandon_save(request.save_operation) - assert scheduler.save_abandon_timeout_s() == 0 assert scheduler.reclaim_stale_leases(1e-9) == [] assert checkpoints.reclaim_stale_offload_pins(1e-9) == 0 assert scheduler._pinned_state_bytes == 30 @@ -359,6 +361,49 @@ def test_no_timeout_or_abandon_recycles_dispatched_source(monkeypatch): assert scheduler._pinned_state_bytes == 0 +@pytest.mark.parametrize("kind", ["save", "load"]) +def test_a_transfer_that_never_reports_stops_the_engine(monkeypatch, kind): + """A lost completion or a TP rank that never reports cannot be told from + a slow DMA, so past the deadline the scheduler fails stop -- holding the + lease and budget rather than recycling them.""" + now = [1000.0] + monkeypatch.setattr(backend.time, "monotonic", lambda: now[0]) + scheduler, checkpoints, _ = make_scheduler(monkeypatch) + seq = sequence(computed=0 if kind == "load" else 16) + if kind == "save": + checkpoint(scheduler, checkpoints, seq, 16) + scheduler.update_state_after_alloc(seq) + else: + scheduler.get_num_new_matched_tokens(seq) + scheduler.update_state_after_alloc(seq) + assert scheduler.should_park_for_load_after_alloc(seq) + [request] = scheduler.build_connector_meta().requests + assert (request.save_operation or request.load_operation) is not None + + scheduler.process_completions(KVConnectorOutput()) + now[0] += scheduler._scheduler_deadline_s - 1 + scheduler.process_completions(KVConnectorOutput()) + now[0] += 2 + with pytest.raises(backend.LMCacheTransferUnprovable): + scheduler.process_completions(KVConnectorOutput()) + assert scheduler._pinned_state_bytes == 30 + + +def test_a_transfer_that_reports_leaves_the_watchdog(monkeypatch): + now = [1000.0] + monkeypatch.setattr(backend.time, "monotonic", lambda: now[0]) + scheduler, checkpoints, _ = make_scheduler(monkeypatch) + seq = sequence() + checkpoint(scheduler, checkpoints, seq, 16) + scheduler.update_state_after_alloc(seq) + [request] = scheduler.build_connector_meta().requests + scheduler.process_completions(KVConnectorOutput()) + terminal(scheduler, request.save_operation) + assert scheduler._transfer_seen_at == {} + now[0] += scheduler._scheduler_deadline_s + 1 + scheduler.process_completions(KVConnectorOutput()) + + @pytest.mark.parametrize("prompt", [8, 16, 19, 24]) def test_load_queries_safe_chunk_boundary_before_reserving_state(monkeypatch, prompt): scheduler, checkpoints, adapter = make_scheduler(monkeypatch) diff --git a/tests/test_lmcache_mp_native_worker.py b/tests/test_lmcache_mp_native_worker.py index 13e2bc25fd..c87e525b92 100644 --- a/tests/test_lmcache_mp_native_worker.py +++ b/tests/test_lmcache_mp_native_worker.py @@ -205,56 +205,37 @@ def test_invalid_native_source_fails_before_transport(worker, units): assert not terminal.succeeded -def test_uncertain_remote_submission_retains_lease(worker): - def uncertain(*_): - raise ConnectionError("server may have received request") - - worker._adapter.submit_store_request = uncertain - worker._submit_save(request(), object()) - assert not worker.get_finished().connector_completions - assert len(worker._native_saves) == 1 - - -def test_uncertain_submission_fails_after_the_uncertainty_bound(worker, monkeypatch): - """Retaining the lease is right, but not forever: past the bound the save - settles as failed so its lease, budget and pending slot are released.""" +@pytest.mark.parametrize("loading", [False, True]) +def test_unprovable_submission_retains_lease_until_the_deadline_stops_the_engine( + worker, monkeypatch, loading +): + """No clock releases a lease the server may still be using: the operation + stays pending, and past the transfer deadline the worker fails stop.""" from atom.kv_transfer.offload.mp import backend now = [1000.0] monkeypatch.setattr(backend.time, "monotonic", lambda: now[0]) - def uncertain(*_): + def unprovable(*_): raise ConnectionError("server may have received request") - worker._adapter.submit_store_request = uncertain - worker._submit_save(request(), object()) - now[0] += worker._uncertain_timeout_s - 1 - assert not worker.get_finished().connector_completions - assert len(worker._native_saves) == 1 + worker._adapter.submit_store_request = unprovable + worker._adapter.submit_retrieve_request = unprovable + pending = worker._native_loads if loading else worker._native_saves + if loading: + worker._submit_load(request(loading=True), object()) + else: + worker._submit_save(request(), object()) + now[0] += worker._transfer_deadline_s - 1 + output = worker.get_finished() + assert not output.connector_completions + assert not output.finished_loading and not output.failed_loading + assert len(pending) == 1 now[0] += 2 - completions = worker.get_finished().connector_completions - [terminal] = [ - completion - for completion in completions - if completion.channel == NATIVE_STATE_MP_STORE_CHANNEL - ] - assert not terminal.succeeded - assert worker._native_saves == {} - - -def test_uncertainty_bound_defaults_to_twice_the_adapter_timeout(): - from atom.kv_transfer.offload.mp import backend - - assert backend._uncertain_transfer_timeout_s(config()) == 600.0 - assert ( - backend._uncertain_transfer_timeout_s(config(**{"lmcache.mp.mq_timeout": 5})) - == 10.0 - ) - with pytest.raises(ValueError, match="uncertain_transfer_timeout_s"): - backend._uncertain_transfer_timeout_s( - config(**{"lmcache.mp.uncertain_transfer_timeout_s": 0}) - ) + with pytest.raises(backend.LMCacheTransferUnprovable): + worker.get_finished() + assert len(pending) == 1 def test_failed_retrieve_does_not_restore_or_report_success(worker, monkeypatch): @@ -518,11 +499,12 @@ def test_registration_reserves_every_restore_descriptor_slot(monkeypatch): assert reserved and 0 not in reserved -def test_raising_restore_returns_its_descriptor_slot_after_the_bound( +def test_raising_restore_keeps_its_descriptor_slot_until_the_deadline( worker, monkeypatch ): - """A restore that raises after taking a descriptor slot keeps it only until - the uncertainty bound, then the load fails and the slot is reusable.""" + """A restore that raises after taking a descriptor slot may have queued a + copy that still reads the slot's staging buffer: the slot and lease stay + held, and the transfer deadline stops the engine instead of recycling.""" from atom.kv_transfer.offload.mp import backend, native_state_worker now = [1000.0] @@ -542,11 +524,11 @@ def test_raising_restore_returns_its_descriptor_slot_after_the_bound( assert not first.finished_loading and not first.failed_loading assert worker._restore_descriptor_slots == [] - now[0] += worker._uncertain_timeout_s + 1 - finished = worker.get_finished() - assert finished.failed_loading == {req.load_operation} - assert worker._restore_descriptor_slots == [3] - assert worker._native_loads == {} + now[0] += worker._transfer_deadline_s + with pytest.raises(backend.LMCacheTransferUnprovable): + worker.get_finished() + assert worker._restore_descriptor_slots == [] + assert len(worker._native_loads) == 1 def test_native_server_chunk_mismatch_fails_before_registration(monkeypatch): From 5cc76a87f0560f26a7c0ac06739b073801102636 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Mon, 28 Sep 2026 03:32:38 +0000 Subject: [PATCH 15/28] fix(offload): lease the final save source when prefix caching is off The late final save reacquired a finished request's prefix through the hash index whenever a BlockManager was bound. With prefix caching off nothing is indexed, so the lookup found nothing and the save was dropped without a warning -- the combination the offload README ships for `lmcache_offload`. Hash reacquire now requires a bound manager with prefix caching; otherwise teardown leases the unemitted suffix of the frozen table, as on main. A late save that finds part of its prefix evicted is counted in `truncated_late_saves`. Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/kv_transfer/offload/_offload_common.py | 2 + atom/kv_transfer/offload/chunked_scheduler.py | 49 ++++++++++++++----- tests/test_lmcache_mp_native_scheduler.py | 1 + tests/test_offload_early_block_release.py | 31 ++++++++++++ 4 files changed, 70 insertions(+), 13 deletions(-) diff --git a/atom/kv_transfer/offload/_offload_common.py b/atom/kv_transfer/offload/_offload_common.py index a2d0e7ec5f..dba285ab7c 100644 --- a/atom/kv_transfer/offload/_offload_common.py +++ b/atom/kv_transfer/offload/_offload_common.py @@ -507,6 +507,7 @@ def _init_offload_statistics(self) -> None: self.total_leased_source_blocks = 0 # ever protected by a save lease self.total_source_safe_released_blocks = 0 # freed once their save reported self.total_abnormal_lease_reclaims = 0 # freed by stall timeout, no report + self.total_truncated_late_saves = 0 # final save lost evicted prefix blocks def process_completions(self, output: KVConnectorOutput) -> KVConnectorOutput: """Apply offload-specific completions and expose plain request IDs.""" @@ -633,6 +634,7 @@ def get_statistics(self) -> dict[str, int]: source_safe_released_blocks=self.total_source_safe_released_blocks, blocks_waiting_for_store=self.blocks_waiting_for_store(), abnormal_lease_reclaims=self.total_abnormal_lease_reclaims, + truncated_late_saves=self.total_truncated_late_saves, ) return statistics diff --git a/atom/kv_transfer/offload/chunked_scheduler.py b/atom/kv_transfer/offload/chunked_scheduler.py index 21a7dab25a..86be439d6f 100644 --- a/atom/kv_transfer/offload/chunked_scheduler.py +++ b/atom/kv_transfer/offload/chunked_scheduler.py @@ -433,12 +433,26 @@ def _late_save_frontier(self, seq, saved: int, available: int) -> int: del seq, saved return self._chunk_floor(available) + def _reacquires_late_source(self) -> bool: + """Whether a finished request's final save reacquires its prefix by hash. + + Only a bound `BlockManager` with prefix caching has a hash index to + reacquire through. Without one (the vLLM plugin, or + `--no-enable_prefix_caching`) teardown instead leases the unemitted + suffix of the frozen block table, and the final save reads exactly + those leased blocks. + """ + manager = self._block_manager + return manager is not None and bool( + getattr(manager, "enable_prefix_caching", False) + ) + 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") + if not self._reacquires_late_source(): + raise RuntimeError("late offload save requires a prefix-cache index") block_ids, available, claimed = self._block_manager.acquire_offload_prefix( seq, saved, aligned ) @@ -449,6 +463,16 @@ def _late_save_source( drop = claimed[keep_count:] if drop: self._block_manager.free_leased_blocks(drop) + if available < aligned: + # A block of the prefix was evicted (or reused) between teardown + # and admission, so the final save stores less than was computed. + self.total_truncated_late_saves += 1 + logger.debug( + "Late offload save of seq %s found %d of %d tokens resident", + seq.id, + available, + aligned, + ) if target <= saved: # Nothing savable is still resident (or no boundary qualifies): # retire the request rather than emit an empty save forever. @@ -580,9 +604,8 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: save_operation = SaveOperationId(seq.id, self._save_nonce) self._save_nonce += 1 late_acquired = frozenset() - if ( - getattr(seq, "_offload_finished", False) - and self._block_manager is not None + if getattr(seq, "_offload_finished", False) and ( + self._reacquires_late_source() ): late_source = self._late_save_source(seq, saved, aligned) if late_source is None: @@ -666,13 +689,13 @@ def protected_block_ids(self, seq) -> frozenset | None: does not support exact leases, or a load is in flight, since a load also touches HBM blocks this connector does not track per-range) -- the scheduler falls back to deferring the whole request. This includes a - final save that has not been emitted yet. With a bound `BlockManager` - (the ATOM engine) teardown freezes its token identity and computed - frontier but protects no physical PAGE: admission reacquires the - still-canonical prefix through the hash index. Without one (the vLLM - plugin, where vLLM owns the blocks) there is nothing to reacquire - through, so teardown freezes the block table and leases the unemitted - suffix, and the final save reads exactly those leased blocks. + final save that has not been emitted yet. When the prefix can be + reacquired by hash (`_reacquires_late_source`) teardown freezes its + token identity and computed frontier but protects no physical PAGE: + admission reacquires the still-canonical prefix through the index. + Otherwise there is nothing to reacquire through, so teardown freezes + the block table and leases the unemitted suffix, and the final save + reads exactly those leased blocks. """ if not self._early_release or self._has_active_load(seq): return None @@ -691,7 +714,7 @@ def protected_block_ids(self, seq) -> frozenset | None: block_id for block_id in blocks.values() if block_id not in safe ) - if self._block_manager is None: + if not self._reacquires_late_source(): table = list(getattr(seq, "_offload_finished_block_ids", seq.block_table)) seq._offload_finished_block_ids = table entry = self._save_tracker.get(str(seq.id)) diff --git a/tests/test_lmcache_mp_native_scheduler.py b/tests/test_lmcache_mp_native_scheduler.py index c93cf0f4d1..e4a667aead 100644 --- a/tests/test_lmcache_mp_native_scheduler.py +++ b/tests/test_lmcache_mp_native_scheduler.py @@ -114,6 +114,7 @@ def connect(config, *, checkpoint_spec=None): ) manager = SimpleNamespace( paged_state_checkpoints=checkpoints, + enable_prefix_caching=True, hash_block_size=4, compute_hash=BlockManager.compute_hash, ) diff --git a/tests/test_offload_early_block_release.py b/tests/test_offload_early_block_release.py index 9107d15ca8..459553d34d 100644 --- a/tests/test_offload_early_block_release.py +++ b/tests/test_offload_early_block_release.py @@ -381,6 +381,7 @@ def test_late_acquire_trims_claims_in_token_order(self, monkeypatch): scheduler = _early_release_scheduler(monkeypatch, chunk_size=8) released = [] scheduler._block_manager = SimpleNamespace( + enable_prefix_caching=True, acquire_offload_prefix=lambda *_args: ( [1, 33, 65, -1], 12, @@ -433,6 +434,7 @@ def test_late_acquire_with_nothing_resident_retires_the_request(self, monkeypatc assert scheduler.build_connector_meta().requests == [] assert scheduler.protected_block_ids(seq) == frozenset() assert bm.kv.num_used == 1 + assert scheduler.get_statistics()["truncated_late_saves"] == 1 def test_late_save_persists_the_short_tail_of_a_long_request(self, monkeypatch): """The tail after earlier saves is stored however short it is.""" @@ -470,6 +472,35 @@ def test_unbound_teardown_leases_the_blocks_its_final_save_reads(self, monkeypat assert final.block_ids[start:end] == [2, 3] assert set(final.block_ids[start:end]) <= leased + def test_without_prefix_caching_teardown_leases_the_final_save_source( + self, monkeypatch + ): + """With prefix caching off nothing is hash-indexed, so reacquiring would + find nothing and silently drop the final save; teardown leases the + unemitted suffix instead and the save reads those blocks.""" + scheduler = _early_release_scheduler(monkeypatch, chunk_size=8) + bm = BlockManager( + MockConfig( + num_kvcache_blocks=24, + kv_cache_block_size=4, + enable_prefix_caching=False, + ) + ) + scheduler.bind_block_manager(bm) + seq = Sequence(list(range(16)), 4, id=108) + assert bm.allocate(seq, bm.can_allocate(seq)) + table = list(seq.block_table) + scheduler.update_state_after_alloc(seq) + seq.num_cached_tokens = 16 + leased = _finish_and_lease(scheduler, seq) + assert leased == frozenset(table) + bm.deallocate_partial(seq, leased) + + [final] = scheduler.build_connector_meta().requests + assert final.save_spec.skip_leading_tokens == 0 + assert final.block_ids[:4] == table + assert scheduler.get_statistics()["truncated_late_saves"] == 0 + class TestTPQuorum: def test_one_incomplete_rank_prevents_logical_group_release(self, monkeypatch): From eaac7eb491a8500c9bc7dc27be9f692e9524e288 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Mon, 28 Sep 2026 03:33:57 +0000 Subject: [PATCH 16/28] fix(offload): reacquire a final save source only after a partial release `protected_block_ids` marked the request finished before the scheduler decided whether it could release it partially. When per-request state made that unsafe, the request was deferred whole and kept its table, but the next metadata build still took the late-save path and claimed every block a second time through the hash index (dropping the save if a hash was missing). `activate_block_leases`, which only the partial-release branch calls, now marks the release, and only that mark selects the reacquire path. Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/kv_transfer/offload/chunked_scheduler.py | 10 ++++++-- atom/plugin/vllm/kv_transfer/seq_view.py | 2 ++ tests/test_offload_early_block_release.py | 23 +++++++++++++++++++ 3 files changed, 33 insertions(+), 2 deletions(-) diff --git a/atom/kv_transfer/offload/chunked_scheduler.py b/atom/kv_transfer/offload/chunked_scheduler.py index 86be439d6f..6a7eeda5b9 100644 --- a/atom/kv_transfer/offload/chunked_scheduler.py +++ b/atom/kv_transfer/offload/chunked_scheduler.py @@ -604,7 +604,7 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata: save_operation = SaveOperationId(seq.id, self._save_nonce) self._save_nonce += 1 late_acquired = frozenset() - if getattr(seq, "_offload_finished", False) and ( + if getattr(seq, "_offload_released", False) and ( self._reacquires_late_source() ): late_source = self._late_save_source(seq, saved, aligned) @@ -730,7 +730,13 @@ 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 + # Only a partial deallocation reaches here. `protected_block_ids` runs + # before the scheduler decides, and a request deferred whole still owns + # its table, so only from now on must a final save reacquire its source. + seq._offload_released = True + if not block_ids: return lease_key = id(seq) leased = self._save_lease_blocks.setdefault(lease_key, set()) diff --git a/atom/plugin/vllm/kv_transfer/seq_view.py b/atom/plugin/vllm/kv_transfer/seq_view.py index 2727c8a4fc..a9639d5746 100644 --- a/atom/plugin/vllm/kv_transfer/seq_view.py +++ b/atom/plugin/vllm/kv_transfer/seq_view.py @@ -35,6 +35,7 @@ class SeqView: "_offload_finished", "_offload_finished_block_ids", "_offload_finished_cached_tokens", + "_offload_released", "_request", "block_table", "offload_handoff_boundary_tokens", @@ -122,6 +123,7 @@ def reset_for_preemption(self) -> None: "_offload_finished", "_offload_finished_block_ids", "_offload_finished_cached_tokens", + "_offload_released", ): if hasattr(self, frozen): delattr(self, frozen) diff --git a/tests/test_offload_early_block_release.py b/tests/test_offload_early_block_release.py index 459553d34d..22e4a4f5e2 100644 --- a/tests/test_offload_early_block_release.py +++ b/tests/test_offload_early_block_release.py @@ -501,6 +501,29 @@ def test_without_prefix_caching_teardown_leases_the_final_save_source( assert final.block_ids[:4] == table assert scheduler.get_statistics()["truncated_late_saves"] == 0 + def test_a_request_deferred_whole_saves_from_the_table_it_still_owns( + self, monkeypatch + ): + """`protected_block_ids` runs before the scheduler decides. A request + whose state cannot be released partially is deferred whole and keeps + its table, so its final save reads that table instead of taking a + second claim on every block through the hash index.""" + scheduler = _early_release_scheduler(monkeypatch, chunk_size=8) + bm, seq, table = _resident_sequence(scheduler, 109, 16, 4) + scheduler.update_state_after_alloc(seq) + seq.num_cached_tokens = 16 + scheduler.request_finished(seq) + assert scheduler.protected_block_ids(seq) is not None + # Deferred whole: no deallocate_partial, no activate_block_leases. + monkeypatch.setattr( + bm, + "acquire_offload_prefix", + lambda *_args: pytest.fail("reacquired a table the request owns"), + ) + + [final] = scheduler.build_connector_meta().requests + assert final.block_ids[:4] == table + class TestTPQuorum: def test_one_incomplete_rank_prevents_logical_group_release(self, monkeypatch): From 818a555a33fcb56d8bd2c7876f0e9e4ed150231d Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Mon, 28 Sep 2026 03:35:11 +0000 Subject: [PATCH 17/28] fix(lmcache_mp): keep the native STATE image pinned until the STORE terminal The worker released a save's checkpoint image as soon as any PAGE chunk milestone reached the boundary. Chunk milestones report PAGE token ranges; nothing in LMCache promises that a chunk completes only after every engine group registered at it, STATE groups included, has been read. The pinned LMCache emits no chunk events, so this path never ran, but it would have been unsound on the first server that did. The STATE source-safe channel is removed: the image pin is released at the STORE terminal, and chunk milestones release PAGE leases only. Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/kv_transfer/offload/README.md | 5 ++-- .../offload/mp/native_state_scheduler.py | 16 ++--------- .../offload/mp/native_state_worker.py | 26 +++++------------ tests/test_lmcache_mp_native_scheduler.py | 27 ++++++++++++++++++ tests/test_lmcache_mp_native_worker.py | 28 ++++--------------- 5 files changed, 46 insertions(+), 56 deletions(-) diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index 2f5e1c95bd..416e68fa7b 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -222,8 +222,9 @@ underlying PAGE stride. The scheduler dispatches one combined PAGE/STATE generation at a time per request, using round-robin admission plus count and byte bounds. It leases the exact READY state image only after admission. Source-safe events release PAGE -leases chunk by chunk and release the state image when its endpoint is safe; -terminal completion settles the logical operation. Failed saves roll back the +leases chunk by chunk. The state image stays pinned until the STORE terminal: +a PAGE chunk milestone does not prove the server has read the STATE groups. +Terminal completion also settles the logical operation. Failed saves roll back the watermark for at most three attempts at one boundary. An uncertain remote DMA is never reclaimed by elapsed time alone. diff --git a/atom/kv_transfer/offload/mp/native_state_scheduler.py b/atom/kv_transfer/offload/mp/native_state_scheduler.py index 7fd5a294fd..7c6ccd91f5 100644 --- a/atom/kv_transfer/offload/mp/native_state_scheduler.py +++ b/atom/kv_transfer/offload/mp/native_state_scheduler.py @@ -28,7 +28,6 @@ _validate_mp_config, ) from atom.kv_transfer.offload.mp.native_state_worker import ( - NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL, NATIVE_STATE_MP_STORE_CHANNEL, require_native_state_server, ) @@ -46,7 +45,6 @@ class _NativeSave: source: StateStoreOperationId saved_before: int boundary: int - source_safe: bool = False @dataclass @@ -302,8 +300,9 @@ def _complete_native_save( self._source_group_finished( SaveSourceGroupId(operation, ((lease.saved_before, lease.boundary),)) ) - if not lease.source_safe: - self._checkpoints.release_offload_store_source(lease.source) + # The image pin is released on the terminal only: no PAGE chunk + # milestone proves the server has finished reading the STATE groups. + self._checkpoints.release_offload_store_source(lease.source) self._checkpoints.settle_offload_store(lease.source) self._refund_state_image() sid = str(operation.req_id) @@ -327,15 +326,6 @@ def save_finished(self, req_id: Any) -> None: self._complete_native_save(req_id, succeeded=True) def connector_completion(self, completion: ConnectorCompletion) -> bool | None: - if completion.channel == NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL: - operation = completion.operation_id - if not isinstance(operation, SaveOperationId): - return False - lease = self._native_saves.get(operation) - if lease is not None and not lease.source_safe: - self._checkpoints.release_offload_store_source(lease.source) - lease.source_safe = True - return None if completion.channel != NATIVE_STATE_MP_STORE_CHANNEL: return super().connector_completion(completion) if not isinstance(completion.operation_id, SaveOperationId): diff --git a/atom/kv_transfer/offload/mp/native_state_worker.py b/atom/kv_transfer/offload/mp/native_state_worker.py index 6dcec709b7..e2cce74206 100644 --- a/atom/kv_transfer/offload/mp/native_state_worker.py +++ b/atom/kv_transfer/offload/mp/native_state_worker.py @@ -43,7 +43,6 @@ logger = logging.getLogger("atom") NATIVE_STATE_MP_STORE_CHANNEL = "native_state_mp_store" -NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL = "native_state_mp_state_source_safe" def require_native_state_server(adapter: Any, config: Any) -> None: @@ -66,7 +65,6 @@ class _NativePending: restore_succeeded: bool = False descriptor_slot: int | None = None immediate_success: bool = False - state_source_safe: bool = False started_at: float = field(default_factory=lambda: time.monotonic()) @@ -312,26 +310,17 @@ def _begin_restore(self, pending: _NativePending) -> bool: event.record(self._restore_stream) return True - def _emit_native_source_safe( + def _emit_page_source_safe( self, output: KVConnectorOutput, pending: _NativePending, ranges: tuple[tuple[int, int], ...], ) -> None: + # PAGE only. A chunk milestone says nothing about when the server has + # read the STATE groups, so the image pin waits for the STORE terminal. operation = pending.request.save_operation - if operation is None or not ranges: - return - output.connector_completions |= _source_safe_completions(operation, ranges) - boundary = int(pending.request.native_state.boundary_tokens) - if not pending.state_source_safe and any(end >= boundary for _, end in ranges): - output.connector_completions.add( - ConnectorCompletion( - NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL, - operation, - True, - ) - ) - pending.state_source_safe = True + if operation is not None and ranges: + output.connector_completions |= _source_safe_completions(operation, ranges) def start_load_kv(self, metadata: Any) -> None: # Read-after-write: this step's compute (including SLOT relocations in @@ -357,7 +346,7 @@ def get_finished(self) -> KVConnectorOutput: take_ranges = getattr(pending.future, "take_completed_ranges", None) if callable(take_ranges): try: - self._emit_native_source_safe( + self._emit_page_source_safe( output, pending, tuple(take_ranges()) ) except Exception: @@ -377,7 +366,7 @@ def get_finished(self) -> KVConnectorOutput: save_spec = pending.request.save_spec start = 0 if save_spec is None else int(save_spec.skip_leading_tokens) end = len(pending.request.token_ids) - self._emit_native_source_safe( + self._emit_page_source_safe( output, pending, _chunk_ranges(start, end, self.chunk_size) ) output.connector_completions.add( @@ -450,7 +439,6 @@ def get_finished(self) -> KVConnectorOutput: __all__ = [ - "NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL", "NATIVE_STATE_MP_STORE_CHANNEL", "NativeStateLMCacheMPConnector", "require_native_state_server", diff --git a/tests/test_lmcache_mp_native_scheduler.py b/tests/test_lmcache_mp_native_scheduler.py index e4a667aead..7e59260b83 100644 --- a/tests/test_lmcache_mp_native_scheduler.py +++ b/tests/test_lmcache_mp_native_scheduler.py @@ -12,7 +12,9 @@ KVConnectorOutput, LoadOperationId, SaveOperationId, + SaveSourceGroupId, ) +from atom.kv_transfer.offload.chunked_scheduler import DENSE_PAGE_SOURCE_SAFE_CHANNEL from atom.kv_transfer.offload.mp import backend from atom.kv_transfer.offload.mp.connector import LMCacheMPConnectorScheduler from atom.kv_transfer.offload.mp.native_state_scheduler import ( @@ -362,6 +364,31 @@ def test_no_timeout_or_abandon_recycles_dispatched_source(monkeypatch): assert scheduler._pinned_state_bytes == 0 +def test_page_source_safe_never_releases_the_state_image(monkeypatch): + """PAGE chunk milestones do not prove the server read the STATE groups: + the image stays pinned, units included, until the STORE terminal.""" + scheduler, checkpoints, _ = make_scheduler(monkeypatch) + seq = sequence() + checkpoint(scheduler, checkpoints, seq, 16) + scheduler.update_state_after_alloc(seq) + [request] = scheduler.build_connector_meta().requests + [pin] = checkpoints.store._offload_pins.values() + scheduler.process_completions( + KVConnectorOutput( + connector_completions={ + ConnectorCompletion( + DENSE_PAGE_SOURCE_SAFE_CHANNEL, + SaveSourceGroupId(request.save_operation, ((0, 8), (8, 16))), + True, + ) + } + ) + ) + assert not pin.source_released + terminal(scheduler, request.save_operation) + assert not checkpoints.has_offload_pins() + + @pytest.mark.parametrize("kind", ["save", "load"]) def test_a_transfer_that_never_reports_stops_the_engine(monkeypatch, kind): """A lost completion or a TP rank that never reports cannot be told from diff --git a/tests/test_lmcache_mp_native_worker.py b/tests/test_lmcache_mp_native_worker.py index c87e525b92..65d4aee4e0 100644 --- a/tests/test_lmcache_mp_native_worker.py +++ b/tests/test_lmcache_mp_native_worker.py @@ -31,7 +31,6 @@ build_native_state_mp_layout, ) from atom.kv_transfer.offload.mp.native_state_worker import ( - NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL, NATIVE_STATE_MP_STORE_CHANNEL, NativeStateLMCacheMPConnector, require_native_state_server, @@ -154,7 +153,6 @@ def test_store_transmits_page_zero_as_real_native_unit(worker): finished = worker.get_finished() assert {completion.channel for completion in finished.connector_completions} == { DENSE_PAGE_SOURCE_SAFE_CHANNEL, - NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL, NATIVE_STATE_MP_STORE_CHANNEL, } terminals = [ @@ -319,7 +317,10 @@ def test_incremental_load_transfers_only_page_suffix_but_full_native_image(worke assert submitted.block_ids == [[3, 4], [0], [25], [31]] -def test_source_safe_ranges_are_reported_before_terminal(worker): +def test_page_ranges_are_source_safe_before_terminal_but_state_is_not(worker): + """Chunk milestones cover PAGE only: the final chunk being done says + nothing about the STATE groups, so nothing about the image is reported + before the STORE terminal.""" req = request() worker.future.source_ranges = [(0, 8)] worker._submit_save(req, object()) @@ -336,24 +337,13 @@ def test_source_safe_ranges_are_reported_before_terminal(worker): worker.future.source_ranges = [(8, 16)] second = worker.get_finished() - assert ( + assert second.connector_completions == { ConnectorCompletion( DENSE_PAGE_SOURCE_SAFE_CHANNEL, SaveSourceGroupId(req.save_operation, ((8, 16),)), True, ) - in second.connector_completions - ) - assert ( - ConnectorCompletion( - NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL, req.save_operation, True - ) - in second.connector_completions - ) - assert all( - completion.channel != NATIVE_STATE_MP_STORE_CHANNEL - for completion in second.connector_completions - ) + } def test_collapsed_tp_non_writer_skips_transport_and_reports_safe_success(worker): @@ -363,12 +353,6 @@ def test_collapsed_tp_non_writer_skips_transport_and_reports_safe_success(worker assert worker.submitted == [] completions = worker.get_finished().connector_completions - assert ( - ConnectorCompletion( - NATIVE_STATE_MP_SOURCE_SAFE_CHANNEL, req.save_operation, True - ) - in completions - ) assert ( ConnectorCompletion(NATIVE_STATE_MP_STORE_CHANNEL, req.save_operation, True) in completions From c8156e82728faad97682d685ac95c877a2674156 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Mon, 28 Sep 2026 03:37:30 +0000 Subject: [PATCH 18/28] fix(envs): read an empty offload knob as its default everywhere `OFFLOAD_PUBLICATION_TIMEOUT_S=` made `float('')` raise out of connector init, and `OFFLOAD_COPY_WORKERS` / `OFFLOAD_LOAD_WORKERS` died with a bare `invalid literal for int()`. The offload section now states one policy: unset or empty is the default; a set but unusable value warns and falls back for knobs that only tune reuse, and is rejected at startup, naming the variable, for widths, sizes and timeouts. The transfer-mode comment no longer lists `engine_driven` as valid. Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/utils/envs.py | 30 ++++++++++++++++++++++++------ tests/test_envs.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 6 deletions(-) diff --git a/atom/utils/envs.py b/atom/utils/envs.py index c2a7f13006..507cf5af14 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -67,9 +67,15 @@ def _optional_int_env(name: str, *, min_value: int | None = None) -> int | None: return value +def _int_env(name: str, default: int) -> int: + """Unset or empty reads as ``default``; anything else must be an integer.""" + value = _optional_int_env(name) + return default if value is None else value + + def _int_env_or_default(name: str, default: int) -> int: raw = os.getenv(name) - if raw is None: + if raw is None or raw == "": return default try: return int(raw) @@ -97,7 +103,14 @@ def _nonnegative_int_env(name: str, default: int) -> int: def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float: - value = float(os.getenv(name, str(default))) + """Unset or empty reads as ``default``; a set value must be a valid float.""" + raw = os.getenv(name) + if raw is None or raw == "": + return default + try: + value = float(raw) + except ValueError as exc: + raise ValueError(f"{name} must be a number, got {raw!r}") from exc if not math.isfinite(value): raise ValueError(f"{name} must be finite") if allow_zero and value < 0: @@ -881,10 +894,14 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float: # Fuse FP32 post-wkv gating and residual addition; 0 selects the torch reference. "ATOM_ENGRAM_FUSED_GATE": lambda: os.getenv("ATOM_ENGRAM_FUSED_GATE", "1") == "1", # --- KV offload (LMCache) --- + # Unset or empty always means the default, so a knob can be cleared inline. + # A set but unusable value warns and falls back for the token thresholds + # and lookup step counts, which only tune reuse; it is rejected at startup + # for widths, sizes and timeouts, which shape memory and concurrency. # Save / load executor widths of the offload worker. DSV4 ignores # OFFLOAD_LOAD_WORKERS: its SLOT load path needs a serial load executor. - "OFFLOAD_COPY_WORKERS": lambda: int(os.getenv("OFFLOAD_COPY_WORKERS", "1") or "1"), - "OFFLOAD_LOAD_WORKERS": lambda: int(os.getenv("OFFLOAD_LOAD_WORKERS", "1") or "1"), + "OFFLOAD_COPY_WORKERS": lambda: _int_env("OFFLOAD_COPY_WORKERS", 1), + "OFFLOAD_LOAD_WORKERS": lambda: _int_env("OFFLOAD_LOAD_WORKERS", 1), # Running-plus-queued save bound. None means unset: the connector then # derives max(2, 2 * OFFLOAD_COPY_WORKERS) and the scheduler's state tier # uses 2. A kv_connector_extra_config "max_pending_saves" takes precedence. @@ -940,8 +957,9 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float: "OFFLOAD_PUBLICATION_POLL_INTERVAL_S": lambda: _finite_float_env( "OFFLOAD_PUBLICATION_POLL_INTERVAL_S", 0.01, allow_zero=False ), - # LMCache MP transfer mode: auto, lmcache_driven or engine_driven. A - # kv_connector_extra_config "lmcache.mp.mp_transfer_mode" takes precedence. + # LMCache MP transfer mode: auto or lmcache_driven (engine_driven is + # rejected as not implemented). A kv_connector_extra_config + # "lmcache.mp.mp_transfer_mode" takes precedence. "LMCACHE_MP_TRANSFER_MODE": lambda: os.getenv("LMCACHE_MP_TRANSFER_MODE", "auto"), } diff --git a/tests/test_envs.py b/tests/test_envs.py index 3b8d880fee..76068c3963 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -296,3 +296,32 @@ def test_offload_env_vars_are_documented(): ] assert offload assert [name for name in offload if f"**{name}**" not in doc] == [] + + +@pytest.mark.parametrize( + ("name", "default"), + [ + ("OFFLOAD_PUBLICATION_TIMEOUT_S", 5.0), + ("OFFLOAD_PUBLICATION_POLL_INTERVAL_S", 0.01), + ("OFFLOAD_COPY_WORKERS", 1), + ("OFFLOAD_LOAD_WORKERS", 1), + ("OFFLOAD_MIN_SAVE_TOKENS", 8192), + ], +) +def test_empty_offload_knob_reads_as_its_default(monkeypatch, name, default): + """`VAR=` is how a knob is cleared inline; it must never crash startup.""" + monkeypatch.setenv(name, "") + assert getattr(_get_envs(), name) == default + + +@pytest.mark.parametrize("name", ["OFFLOAD_COPY_WORKERS", "OFFLOAD_LOAD_WORKERS"]) +def test_malformed_offload_worker_width_names_the_variable(monkeypatch, name): + monkeypatch.setenv(name, "two") + with pytest.raises(ValueError, match=f"{name} must be an integer"): + getattr(_get_envs(), name) + + +def test_malformed_offload_timeout_names_the_variable(monkeypatch): + monkeypatch.setenv("OFFLOAD_PUBLICATION_TIMEOUT_S", "soon") + with pytest.raises(ValueError, match="OFFLOAD_PUBLICATION_TIMEOUT_S must be"): + _ = _get_envs().OFFLOAD_PUBLICATION_TIMEOUT_S From 9b8f7798cb5691065a78bb3635611beae51f0908 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Mon, 28 Sep 2026 03:37:30 +0000 Subject: [PATCH 19/28] fix(lmcache_mp): keep failed adoptions releasable and surface broken invariants - `adopt_transfer_units` popped the transfer record before adopting; if adoption raised, the units were reserved but unreachable by any release path. The record is now removed only after adoption returns. - A native restore that loses the publish race to an identical image was dropped silently; it is now logged as a deduplication. - The worker's own pending-save refusal is unreachable while the scheduler enforces the same bound; if it ever fires it now logs an error instead of passing as an ordinary save failure. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../offload/mp/native_state_scheduler.py | 13 +++++++++++-- .../offload/mp/native_state_worker.py | 11 +++++++++-- atom/model_engine/page_unit_checkpoint.py | 7 +++++-- tests/test_page_unit_checkpoint_leases.py | 19 +++++++++++++++++++ 4 files changed, 44 insertions(+), 6 deletions(-) diff --git a/atom/kv_transfer/offload/mp/native_state_scheduler.py b/atom/kv_transfer/offload/mp/native_state_scheduler.py index 7c6ccd91f5..9ea10fdc2e 100644 --- a/atom/kv_transfer/offload/mp/native_state_scheduler.py +++ b/atom/kv_transfer/offload/mp/native_state_scheduler.py @@ -5,6 +5,7 @@ from __future__ import annotations +import logging from collections.abc import Callable from dataclasses import dataclass from typing import Any, TypeVar @@ -34,6 +35,8 @@ from atom.model_engine.page_unit_checkpoint import SuspendedCheckpointRestore from atom.utils import envs +logger = logging.getLogger("atom") + _MAX_SAVE_ATTEMPTS = 3 _T = TypeVar("_T") @@ -440,9 +443,15 @@ def _finish_native_load(self, operation: Any, *, succeeded: bool) -> bool: return False lease = self._native_loads[operation] if succeeded: - self._checkpoints.adopt_transfer_units( + if not self._checkpoints.adopt_transfer_units( operation, lease.transfer.prefix_hash - ) + ): + # Another request published the same image first. Ours was + # released; the SLOT already holds identical restored state. + logger.debug( + "Native restore for %s deduplicated against a READY image", + operation, + ) if lease.local_restore is not None: self._checkpoints.release_suspended_restore(lease.local_restore) lease.local_restore = None diff --git a/atom/kv_transfer/offload/mp/native_state_worker.py b/atom/kv_transfer/offload/mp/native_state_worker.py index e2cce74206..9ce13eaff6 100644 --- a/atom/kv_transfer/offload/mp/native_state_worker.py +++ b/atom/kv_transfer/offload/mp/native_state_worker.py @@ -219,8 +219,15 @@ def _submit_native(self, req: LMCacheReqMeta, event: Any, *, loading: bool) -> N ) return if not loading and len(pending) >= self._max_pending_saves: - # The scheduler has the same bound. Refuse before transport; - # a terminal False safely returns the logical admission credit. + # Unreachable while the scheduler enforces the same bound: every + # save held here is still in its `_save_inflight`. Refuse before + # transport; a terminal False returns the admission credit. + logger.error( + "Native-state LMCache MP save %s exceeds the scheduler's " + "pending bound %d; refusing it", + operation_id, + self._max_pending_saves, + ) pending[operation_id] = _NativePending(req, None) return end = len(req.token_ids) diff --git a/atom/model_engine/page_unit_checkpoint.py b/atom/model_engine/page_unit_checkpoint.py index acaa147788..97e97a20ba 100644 --- a/atom/model_engine/page_unit_checkpoint.py +++ b/atom/model_engine/page_unit_checkpoint.py @@ -1079,10 +1079,13 @@ def adopt_transfer_units(self, owner: Hashable, prefix_hash: int) -> bool: units are released and ``False`` is returned. The adopted checkpoint is deliberately not nominated for offload: it just arrived from that tier. """ - units = self._transfer_units.pop(owner, None) + units = self._transfer_units.get(owner) if units is None: return False - return self.store.adopt_units(units, ("state-transfer", owner), prefix_hash) + adopted = self.store.adopt_units(units, ("state-transfer", owner), prefix_hash) + # Only now: if adoption raised, the units stay reachable for release. + del self._transfer_units[owner] + return adopted def suspend_queued_restore( self, dst_slot: int diff --git a/tests/test_page_unit_checkpoint_leases.py b/tests/test_page_unit_checkpoint_leases.py index 776ea7574c..1c9b700ff8 100644 --- a/tests/test_page_unit_checkpoint_leases.py +++ b/tests/test_page_unit_checkpoint_leases.py @@ -2,6 +2,8 @@ """Native PAGE image leases used by external checkpoint transfers.""" +import pytest + from atom.model_engine.block_pool import BlockPool from atom.model_engine.page_unit_checkpoint import ( PagedStateCheckpointCoordinator, @@ -202,6 +204,23 @@ def test_duplicate_external_load_adoption_releases_incoming_units(): assert c.store.pool.num_free == 3 +def test_failed_adoption_keeps_transfer_units_releasable(monkeypatch): + """If adoption raises, the units must still be owned by the transfer + record, or no release path could ever return them to the pool.""" + c = coordinator(num_units=3) + owner = ("request-a", 1) + assert c.reserve_transfer_units(owner) is not None + + def broken(*_args): + raise AssertionError("owner mismatch") + + monkeypatch.setattr(c.store, "adopt_units", broken) + with pytest.raises(AssertionError): + c.adopt_transfer_units(owner, 101) + c.release_transfer_units(owner) + assert c.store.pool.num_free == 3 + + def test_suspended_restore_can_resume_or_release_its_source_pin(): c = coordinator(num_units=3) ready(c, 101) From 45943d16fe07ce185c74bc68ee85f1b6de55bec7 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Mon, 28 Sep 2026 03:38:45 +0000 Subject: [PATCH 20/28] test(lmcache_mp): run a native restore end to end and fix drifted copy doubles The restore-copy doubles were positional-only, so the production call with `descriptor_slot=` would have raised inside `_begin_restore` and been reported as a failed restore; the success tests replaced `_begin_restore` outright, so no test ran the happy path. The doubles now take the production signature, and a new test drives a terminal retrieve through `get_finished` into the real `_begin_restore`, then asserts the copy's descriptor slot, that the load finishes only once the restore event is done, and that the slot is returned. The CUDA stand-ins are shared with the fence test. Co-Authored-By: Claude Opus 5.5 (1M context) --- tests/test_lmcache_mp_native_layout.py | 2 +- tests/test_lmcache_mp_native_worker.py | 128 +++++++++++++++++-------- tests/test_lmcache_mp_shell.py | 2 +- 3 files changed, 90 insertions(+), 42 deletions(-) diff --git a/tests/test_lmcache_mp_native_layout.py b/tests/test_lmcache_mp_native_layout.py index be107bfdec..dece1311e4 100644 --- a/tests/test_lmcache_mp_native_layout.py +++ b/tests/test_lmcache_mp_native_layout.py @@ -47,7 +47,7 @@ def _transfer(*, widths=(8, 8, 2), image_bytes=39, num_blocks=7): slot_regions=[], block_tensor_views=views, paged_state_checkpoint_spec=spec, - execute_paged_state_copies=lambda stores, restores: None, + execute_paged_state_copies=lambda stores, restores, descriptor_slot=0: None, ) transfer.set_block_count(num_blocks) return transfer diff --git a/tests/test_lmcache_mp_native_worker.py b/tests/test_lmcache_mp_native_worker.py index 65d4aee4e0..d62a040c69 100644 --- a/tests/test_lmcache_mp_native_worker.py +++ b/tests/test_lmcache_mp_native_worker.py @@ -103,7 +103,7 @@ def worker(): slot_regions=[], block_tensor_views=[page], paged_state_checkpoint_spec=spec, - execute_paged_state_copies=lambda *_: None, + execute_paged_state_copies=lambda stores, restores, descriptor_slot=0: None, ) tensors.set_block_count(32) instance._native_layout = build_native_state_mp_layout( @@ -364,67 +364,115 @@ def test_collapsed_tp_non_writer_skips_transport_and_reports_safe_success(worker } == {((0, 8),), ((8, 16),)} -def test_restore_is_fenced_against_the_compute_stream_both_ways(worker, monkeypatch): - """The restore uses its own descriptor slot and stream, and orders against - compute both ways: it waits for work already on the compute stream (the - SLOT's previous occupant), and the next step's compute waits for it.""" - from atom.kv_transfer.offload.mp import native_state_worker +class _Event: + def __init__(self): + self.recorded_on = None + self.done = False - class Event: - def __init__(self): - self.recorded_on = None + def record(self, stream): + self.recorded_on = stream - def record(self, stream): - self.recorded_on = stream + def query(self): + return self.done - class Stream: - def __init__(self, name): - self.name = name - self.waited_streams = [] - self.waited_events = [] - def wait_stream(self, other): - self.waited_streams.append(other) +class _Stream: + def __init__(self, name): + self.name = name + self.waited_streams = [] + self.waited_events = [] - def wait_event(self, event): - self.waited_events.append(event) + def wait_stream(self, other): + self.waited_streams.append(other) - class StreamContext: - def __enter__(self): - return self + def wait_event(self, event): + self.waited_events.append(event) - def __exit__(self, *_args): - return False - compute = Stream("compute") - restore_stream = Stream("restore") - event = Event() - copied = [] - worker._restore_stream = restore_stream - worker._restore_descriptor_slots = [3] - worker._native_copy = lambda stores, restores, descriptor_slot=0: copied.append( - (stores, restores, descriptor_slot) +class _StreamContext: + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + +@pytest.fixture +def cuda(worker, monkeypatch): + """CPU stand-ins for the restore stream, the compute stream and events.""" + from atom.kv_transfer.offload.mp import native_state_worker + + stubs = SimpleNamespace( + compute=_Stream("compute"), restore=_Stream("restore"), events=[] ) - monkeypatch.setattr(native_state_worker.torch.cuda, "Event", lambda: event) + + def event(): + stubs.events.append(_Event()) + return stubs.events[-1] + + worker._restore_stream = stubs.restore + monkeypatch.setattr(native_state_worker.torch.cuda, "Event", event) monkeypatch.setattr( - native_state_worker.torch.cuda, "stream", lambda stream: StreamContext() + native_state_worker.torch.cuda, "stream", lambda stream: _StreamContext() ) monkeypatch.setattr( - native_state_worker.torch.cuda, "current_stream", lambda: compute + native_state_worker.torch.cuda, "current_stream", lambda: stubs.compute + ) + return stubs + + +def test_restore_is_fenced_against_the_compute_stream_both_ways(worker, cuda): + """The restore uses its own descriptor slot and stream, and orders against + compute both ways: it waits for work already on the compute stream (the + SLOT's previous occupant), and the next step's compute waits for it.""" + copied = [] + worker._restore_descriptor_slots = [3] + worker._native_copy = lambda stores, restores, descriptor_slot=0: copied.append( + (stores, restores, descriptor_slot) ) req = request(loading=True) worker._submit_load(req, object()) pending = next(iter(worker._native_loads.values())) assert worker._begin_restore(pending) + [event] = cuda.events assert pending.descriptor_slot == 3 assert worker._restore_descriptor_slots == [] assert copied[0][2] == 3 - assert event.recorded_on is restore_stream - assert restore_stream.waited_streams == [compute] + assert event.recorded_on is cuda.restore + assert cuda.restore.waited_streams == [cuda.compute] worker.start_load_kv(object()) - assert compute.waited_events == [event] + assert cuda.compute.waited_events == [event] + + +def test_native_restore_reports_success_and_returns_its_slot(worker, cuda): + """End to end through `get_finished`: a terminal retrieve starts the real + restore with the production copy signature, and only once its event is + done does the load finish and the descriptor slot come back.""" + copied = [] + worker._restore_descriptor_slots = [3] + worker._native_copy = lambda stores, restores, descriptor_slot=0: copied.append( + (stores, restores, descriptor_slot) + ) + worker.future.value, worker.future.ready = True, True + req = request(loading=True) + worker._submit_load(req, object()) + + first = worker.get_finished() + assert not first.finished_loading and not first.failed_loading + [(stores, (restore,), slot)] = copied + assert stores == () and slot == 3 + assert restore.dst_slot == 2 + assert tuple(restore.unit_ids) == (0, 25, 31) + assert worker._restore_descriptor_slots == [] + + cuda.events[-1].done = True + second = worker.get_finished() + assert second.finished_loading == {req.load_operation} + assert not second.failed_loading + assert worker._restore_descriptor_slots == [3] + assert worker._native_loads == {} def test_registration_reserves_every_restore_descriptor_slot(monkeypatch): @@ -468,7 +516,7 @@ def test_registration_reserves_every_restore_descriptor_slot(monkeypatch): paged_state_checkpoint_spec=PagedStateCheckpointSpec( 32, 128, "native-test-v1", 80 ), - execute_paged_state_copies=lambda *_: None, + execute_paged_state_copies=lambda stores, restores, descriptor_slot=0: None, ) tensors.set_block_count(32) tensors.state_backend = SimpleNamespace( diff --git a/tests/test_lmcache_mp_shell.py b/tests/test_lmcache_mp_shell.py index ad8a90f5cc..3a8124e95f 100644 --- a/tests/test_lmcache_mp_shell.py +++ b/tests/test_lmcache_mp_shell.py @@ -84,7 +84,7 @@ def __init__(self, config): worker = LMCacheMPConnector(config) tensors = SimpleNamespace( paged_state_checkpoint_spec=object(), - execute_paged_state_copies=lambda *_: None, + execute_paged_state_copies=lambda stores, restores, descriptor_slot=0: None, ) worker.register_kv_caches({}, tensors, 7) assert selected == [("native", config)] From 4f0eeda69bb9843007fb5781c2a63f393c76e48a Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Mon, 28 Sep 2026 03:40:27 +0000 Subject: [PATCH 21/28] perf(lmcache_mp): invalidate a request's save frontier only on its own checkpoints The per-request `_save_frontier` memo was keyed on the store-wide generation, so any request's publish or eviction forced every tracked request to rescan its prompt on the next step. The store now keeps a bounded log of which prefix hash each generation bump touched; a request rescans only when a change hits one of the boundaries it scanned (its answer or anything above it), or when the log cannot tell. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../offload/mp/native_state_scheduler.py | 32 ++++++++++---- atom/model_engine/page_unit_checkpoint.py | 42 +++++++++++++++---- tests/test_lmcache_mp_native_scheduler.py | 10 +++++ tests/test_page_unit_checkpoint_leases.py | 15 +++++++ 4 files changed, 85 insertions(+), 14 deletions(-) diff --git a/atom/kv_transfer/offload/mp/native_state_scheduler.py b/atom/kv_transfer/offload/mp/native_state_scheduler.py index 9ea10fdc2e..c908810805 100644 --- a/atom/kv_transfer/offload/mp/native_state_scheduler.py +++ b/atom/kv_transfer/offload/mp/native_state_scheduler.py @@ -228,18 +228,36 @@ def _save_frontier(self, seq: Any) -> int: ) floor = self._save_floor(exhausted) # Every tracked request asks every step. The answer changes only when - # its frontier moves or a checkpoint is published or dropped, so - # rescan only then instead of walking the prompt per step. - key = (frontier, floor, self._checkpoints.store.generation) + # its frontier moves or one of ITS boundary checkpoints is published or + # dropped, so rescan only then -- not whenever any request's + # checkpoint traffic moves the store. + store = self._checkpoints.store memo = getattr(seq, "_mp_save_frontier_memo", None) - if memo is not None and memo[0] == key: - return memo[1] + if memo is not None and memo[:2] == (frontier, floor): + changed = store.changed_since(memo[2]) + boundaries = getattr(seq, "_mp_boundary_by_hash", {}) + if changed is not None and not any( + floor < boundaries.get(prefix_hash, 0) <= frontier + for prefix_hash in changed + ): + seq._mp_save_frontier_memo = ( + frontier, + floor, + store.generation, + memo[3], + ) + return memo[3] found = 0 + boundaries = getattr(seq, "_mp_boundary_by_hash", None) + if boundaries is None: + boundaries = seq._mp_boundary_by_hash = {} for boundary in range(frontier, floor, -self.chunk_size): - if self._checkpoints.contains(self._boundary_hash(seq, boundary)): + prefix_hash = self._boundary_hash(seq, boundary) + boundaries[prefix_hash] = boundary + if self._checkpoints.contains(prefix_hash): found = boundary break - seq._mp_save_frontier_memo = (key, found) + seq._mp_save_frontier_memo = (frontier, floor, store.generation, found) return found def _save_floor(self, lower: int) -> int: diff --git a/atom/model_engine/page_unit_checkpoint.py b/atom/model_engine/page_unit_checkpoint.py index 97e97a20ba..ceb9356d95 100644 --- a/atom/model_engine/page_unit_checkpoint.py +++ b/atom/model_engine/page_unit_checkpoint.py @@ -15,6 +15,8 @@ from atom.model_engine.block_pool import BlockPool from atom.model_engine.sequence import Sequence +_CHANGE_LOG_LIMIT = 4096 + logger = logging.getLogger("atom") #: How many reclaimed store operations to remember. Only has to outlive a @@ -165,8 +167,11 @@ def __init__( self._offload_sink = offload_sink self.hash_to_checkpoint: dict[int, int] = {} # Bumped whenever the READY set can change (publish, unindex, evict, - # release, clear), so callers can cache `contains` answers. + # release, clear), so callers can cache `contains` answers. The log + # names the hash behind each recent bump (None: every hash), so a + # caller can invalidate only the answers that bump can affect. self.generation = 0 + self._changes: deque[tuple[int, int | None]] = deque(maxlen=_CHANGE_LOG_LIMIT) self.records: dict[int, CheckpointRecord] = {} self._pending_by_hash: dict[int, int] = {} self._lru: OrderedDict[int, None] = OrderedDict() @@ -430,6 +435,29 @@ def release_suspended_restore(self, restore: SuspendedCheckpointRestore) -> None """Release a suspended restore whose destination was filled elsewhere.""" self._release_restore_pin(restore.checkpoint_id) + def _note_change(self, prefix_hash: int | None) -> None: + self.generation += 1 + self._changes.append((self.generation, prefix_hash)) + + def changed_since(self, generation: int) -> set[int] | None: + """Hashes whose READY state may have changed after ``generation``. + + None when that cannot be told: the log no longer reaches back that + far, or a `clear` touched every hash. + """ + if generation >= self.generation: + return set() + if not self._changes or self._changes[0][0] > generation + 1: + return None + changed: set[int] = set() + for bumped_at, prefix_hash in reversed(self._changes): + if bumped_at <= generation: + break + if prefix_hash is None: + return None + changed.add(prefix_hash) + return changed + def adopt_units( self, unit_ids: tuple[int, ...], owner: Hashable, prefix_hash: int ) -> bool: @@ -453,7 +481,7 @@ def adopt_units( self.records[checkpoint_id] = record self.hash_to_checkpoint[record.prefix_hash] = checkpoint_id self._lru[checkpoint_id] = None - self.generation += 1 + self._note_change(record.prefix_hash) return True def complete_inflight(self) -> None: @@ -475,7 +503,7 @@ def complete_inflight(self) -> None: continue record.state = READY self.hash_to_checkpoint[record.prefix_hash] = checkpoint_id - self.generation += 1 + self._note_change(record.prefix_hash) self._lru[checkpoint_id] = None self._queue_offload_store(checkpoint_id, record) @@ -792,7 +820,7 @@ def _release_restore_pin(self, checkpoint_id: int) -> None: self._release_record(checkpoint_id) def unindex(self, prefix_hash: int) -> bool: - self.generation += 1 + self._note_change(prefix_hash) checkpoint_id = self.hash_to_checkpoint.pop(prefix_hash, -1) if checkpoint_id < 0: checkpoint_id = self._pending_by_hash.pop(prefix_hash, -1) @@ -809,7 +837,7 @@ def unindex(self, prefix_hash: int) -> bool: return True def clear(self) -> None: - self.generation += 1 + self._note_change(None) self.hash_to_checkpoint.clear() self._pending_by_hash.clear() self._lru.clear() @@ -824,7 +852,7 @@ def _evict(self, checkpoint_id: int) -> None: record = self.records[checkpoint_id] if record.state != READY or record.pin_count: raise AssertionError("only an unpinned READY checkpoint is evictable") - self.generation += 1 + self._note_change(record.prefix_hash) if self.hash_to_checkpoint.get(record.prefix_hash) == checkpoint_id: del self.hash_to_checkpoint[record.prefix_hash] record.state = EVICTING @@ -833,8 +861,8 @@ def _evict(self, checkpoint_id: int) -> None: self.evictions += 1 def _release_record(self, checkpoint_id: int) -> None: - self.generation += 1 record = self.records.pop(checkpoint_id) + self._note_change(record.prefix_hash) self._lru.pop(checkpoint_id, None) if self.hash_to_checkpoint.get(record.prefix_hash) == checkpoint_id: del self.hash_to_checkpoint[record.prefix_hash] diff --git a/tests/test_lmcache_mp_native_scheduler.py b/tests/test_lmcache_mp_native_scheduler.py index 7e59260b83..c9a2ac7092 100644 --- a/tests/test_lmcache_mp_native_scheduler.py +++ b/tests/test_lmcache_mp_native_scheduler.py @@ -294,9 +294,19 @@ def test_save_frontier_rescans_only_when_frontier_or_checkpoints_change( assert scheduler._save_frontier(seq) == 8 assert len(calls) == scanned # same frontier, same checkpoints: cached + other = sequence(2, token_offset=100) + checkpoint(scheduler, checkpoints, other, 8) # another request's traffic + assert scheduler._save_frontier(seq) == 8 + assert len(calls) == scanned # moves the store, but none of our boundaries + checkpoint(scheduler, checkpoints, seq, 16) # a newer checkpoint publishes assert scheduler._save_frontier(seq) == 16 assert len(calls) > scanned + rescanned = len(calls) + + checkpoints.store.unindex(scheduler._boundary_hash(seq, 16)) # ours drops + assert scheduler._save_frontier(seq) == 8 + assert len(calls) > rescanned def test_default_min_save_tokens_skips_a_short_prompt(monkeypatch): diff --git a/tests/test_page_unit_checkpoint_leases.py b/tests/test_page_unit_checkpoint_leases.py index 1c9b700ff8..9fa46b896f 100644 --- a/tests/test_page_unit_checkpoint_leases.py +++ b/tests/test_page_unit_checkpoint_leases.py @@ -263,3 +263,18 @@ def test_external_load_evicts_only_available_checkpoint_sources(): c.release_transfer_units(("request-a", 1)) c.settle_offload_store(operation) assert c.has_available_units(6) + + +def test_change_log_names_the_hashes_behind_each_generation(): + c = coordinator(num_units=6) + start = c.store.generation + ready(c, 101) + ready(c, 202) + assert c.store.changed_since(start) >= {101, 202} + assert c.store.changed_since(c.store.generation) == set() + middle = c.store.generation + c.unindex(101) + assert 101 in c.store.changed_since(middle) + assert 202 not in c.store.changed_since(middle) + c.store.clear() + assert c.store.changed_since(middle) is None From 28dda85e5e63cc9dc190edff1a884cf997e987f2 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Mon, 28 Sep 2026 04:22:41 +0000 Subject: [PATCH 22/28] refactor(attentions): one fenced descriptor staging for every checkpoint copier DSv4.1 kept its own `_DescriptorStaging`, a third copy of the per-slot pinned descriptor pool the base builder shares with DSv4 and GDN, and the only one that fenced reuse: a non-blocking H2D reads its pinned rows when the stream reaches it, so refilling a buffer whose last upload is still queued rewrites the descriptor under an earlier copy. `DescriptorStaging` in `pool_layout/paged_state_copy.py` now owns the per-slot buffers and the fence, and the base builder (DSv4, GDN) and DSv4.1's `StateCopies` both use it, so the base path gains the fence too. Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/model_ops/attentions/backends.py | 29 +++------ .../attentions/deepseek_v41/backend.py | 3 +- .../attentions/deepseek_v41/checkpoints.py | 52 +++------------- atom/model_ops/attentions/deepseek_v4_attn.py | 22 ++++--- atom/model_ops/attentions/gdn_attn.py | 17 +++--- .../pool_layout/paged_state_copy.py | 61 +++++++++++++++++++ tests/attentions/deepseek_v41/test_cache.py | 16 ----- tests/test_paged_state_copy_planner.py | 39 ++++++++++++ 8 files changed, 142 insertions(+), 97 deletions(-) diff --git a/atom/model_ops/attentions/backends.py b/atom/model_ops/attentions/backends.py index fb0fd16555..d7e58c98c4 100644 --- a/atom/model_ops/attentions/backends.py +++ b/atom/model_ops/attentions/backends.py @@ -23,6 +23,7 @@ from atom.model_engine.scheduler import ScheduledBatch from atom.model_engine.state_runtime import StateTransfer from atom.model_ops.attention_mla import MLAModules +from atom.model_ops.attentions.pool_layout.paged_state_copy import DescriptorStaging from atom.model_ops.attentions.pool_layout.pool_rows import PoolRowsMixin from atom.model_ops.attentions.pool_layout.sub_pool_spec import SubPoolSpec from atom.model_ops.attentions.token_layout.batch_ids import ( @@ -362,8 +363,7 @@ def reserve_checkpoint_descriptors(self, descriptor_slots) -> None: mid-serving restore would stall its thread. A connector that will pass these slots to `execute_paged_state_copies` reserves them up front. """ - for descriptor_slot in descriptor_slots: - self._checkpoint_descriptor_buffer(descriptor_slot) + self._checkpoint_staging().reserve(descriptor_slots) def _checkpoint_descriptor_device(self) -> torch.device: """Device of the planes `execute_paged_state_copies` copies between.""" @@ -371,7 +371,7 @@ def _checkpoint_descriptor_device(self) -> torch.device: f"{type(self).__name__} does not implement PAGE-backed state copy" ) - def _checkpoint_descriptor_buffer(self, descriptor_slot: int = 0) -> CpuGpuBuffer: + def _checkpoint_staging(self) -> DescriptorStaging: """Pinned staging for a step's whole descriptor, sized for the worst step. Pinned because the alternative synchronizes: a pageable H2D from @@ -391,26 +391,17 @@ def _checkpoint_descriptor_buffer(self, descriptor_slot: int = 0) -> CpuGpuBuffe storing one of them from the wrong slot besides; the two constraints have the same owner and move together. - Slot 0 serves `build()`; each other `descriptor_slot` is a separate - buffer, so an out-of-band restore never shares staging with a step. + Slot 0 serves `build()`; each other descriptor slot is a separate, + separately fenced buffer (see `DescriptorStaging`). """ - if descriptor_slot < 0: - raise ValueError("checkpoint descriptor slot must be non-negative") - descriptors = getattr(self, "_checkpoint_descriptors", None) - if descriptors is None: - descriptors = self._checkpoint_descriptors = {} - descriptor = descriptors.get(descriptor_slot) - if descriptor is None: + staging = getattr(self, "_checkpoint_staging_cache", None) + if staging is None: plan = self._checkpoint_copy_plan() max_ops = 2 * int(self.model_runner.config.max_num_seqs) - descriptor = CpuGpuBuffer( - max_ops * plan.num_spans, - 3, - dtype=torch.int64, - device=self._checkpoint_descriptor_device(), + staging = self._checkpoint_staging_cache = DescriptorStaging( + max_ops * plan.num_spans, self._checkpoint_descriptor_device() ) - descriptors[descriptor_slot] = descriptor - return descriptor + return staging def warmup_per_req_cache(self) -> None: """Pay whatever the first checkpoint copy would pay, before serving. diff --git a/atom/model_ops/attentions/deepseek_v41/backend.py b/atom/model_ops/attentions/deepseek_v41/backend.py index f47181f81f..18ca129c4e 100644 --- a/atom/model_ops/attentions/deepseek_v41/backend.py +++ b/atom/model_ops/attentions/deepseek_v41/backend.py @@ -264,8 +264,7 @@ def execute_paged_state_copies(self, stores, restores, descriptor_slot=0): self.copies.execute(stores, restores, descriptor_slot=descriptor_slot) def reserve_checkpoint_descriptors(self, descriptor_slots): - for descriptor_slot in descriptor_slots: - self.copies.staging(descriptor_slot) + self.copies.staging.reserve(descriptor_slots) def warmup_per_req_cache(self): self.copies.warmup() diff --git a/atom/model_ops/attentions/deepseek_v41/checkpoints.py b/atom/model_ops/attentions/deepseek_v41/checkpoints.py index d7fc009705..8723f7ca6b 100644 --- a/atom/model_ops/attentions/deepseek_v41/checkpoints.py +++ b/atom/model_ops/attentions/deepseek_v41/checkpoints.py @@ -5,21 +5,10 @@ import torch from atom.model_ops.attentions.pool_layout.paged_state_copy import ( + DescriptorStaging, launch_copy_descriptor, plan_segmented_copy, ) -from atom.utils import CpuGpuBuffer - - -class _DescriptorStaging: - """One pinned descriptor buffer and the event fencing its last upload.""" - - def __init__(self, rows, device, pin_memory): - self.buffer = CpuGpuBuffer( - rows, 3, dtype=torch.int64, device=device, pin_memory=pin_memory - ) - self.upload_done = torch.cuda.Event() if pin_memory else None - self.pending = False class StateCopies: @@ -44,23 +33,10 @@ def __init__(self, cache, spec, max_batch): ], spec.image_bytes, ) - self._staging_rows = 2 * max_batch * self.plan.num_spans - # Slot 0 serves `build()`; each other descriptor slot is an independent - # buffer, so an out-of-band restore never overwrites a step's staging. - self._stagings = {0: self._new_staging()} - - def _new_staging(self): - return _DescriptorStaging( - self._staging_rows, self.cache.pool.device, self.cache.pool.is_cuda + self.staging = DescriptorStaging( + 2 * max_batch * self.plan.num_spans, cache.pool.device ) - - def staging(self, descriptor_slot=0): - if descriptor_slot < 0: - raise ValueError("checkpoint descriptor slot must be non-negative") - staging = self._stagings.get(descriptor_slot) - if staging is None: - staging = self._stagings[descriptor_slot] = self._new_staging() - return staging + self.staging.reserve([0]) def entry(self, slot): self.cache.require_committed() @@ -114,12 +90,10 @@ def execute(self, stores, restores, descriptor_slot=0): ) at += take return - staging = self.staging(descriptor_slot) + rows = self.staging.rows(descriptor_slot) total = (len(stores) + len(restores)) * self.plan.num_spans - if total > staging.buffer.np.shape[0]: + if total > rows.shape[0]: raise ValueError("Checkpoint copy batch exceeds descriptor capacity") - if staging.pending: - staging.upload_done.synchronize() at = 0 for ops, storing in ((stores, True), (restores, False)): if not ops: @@ -144,12 +118,10 @@ def execute(self, stores, restores, descriptor_slot=0): dtype=np.int64, ) self.plan.write_descriptor( - staging.buffer.np[at:end], slot_bases, page_bases, forward=storing + rows[at:end], slot_bases, page_bases, forward=storing ) at = end - descriptor = staging.buffer.copy_to_gpu(total) - staging.upload_done.record() - staging.pending = True + descriptor = self.staging.upload(descriptor_slot, total) cut = len(stores) * self.plan.num_spans # Restore may consume an image stored in this same maintenance batch. launch_copy_descriptor(descriptor[:cut], self.plan) @@ -167,13 +139,9 @@ def warmup(self): offsets.append(at) at += size dst = np.asarray([[slot.data_ptr() + off for off in offsets]], dtype=np.int64) - staging = self.staging(0) self.plan.write_descriptor( - staging.buffer.np[: self.plan.num_spans], + self.staging.rows(0)[: self.plan.num_spans], np.asarray([[slot.data_ptr()]], dtype=np.int64), dst, ) - descriptor = staging.buffer.copy_to_gpu(self.plan.num_spans) - staging.upload_done.record() - staging.pending = True - launch_copy_descriptor(descriptor, self.plan) + launch_copy_descriptor(self.staging.upload(0, self.plan.num_spans), self.plan) diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index 6d86721ee9..a701135fd9 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -80,6 +80,7 @@ plan_regions, ) from atom.model_ops.attentions.pool_layout.paged_state_copy import ( + DescriptorStaging, SegmentedCopyPlan, launch_copy_descriptor, plan_segmented_copy, @@ -730,7 +731,7 @@ def __init__(self, model_runner): self._page_unit_region_owners: tuple[int, ...] = () self._checkpoint_plan_cache: SegmentedCopyPlan | None = None self._checkpoint_slot_base_cache: np.ndarray | None = None - self._checkpoint_descriptors: dict[int, CpuGpuBuffer] = {} + self._checkpoint_staging_cache: DescriptorStaging | None = None @property def prep_stream(self): @@ -988,14 +989,15 @@ def execute_paged_state_copies( slot_bases = self._checkpoint_slot_bases() per_op = plan.num_spans total = (len(store_ops) + len(restore_ops)) * per_op - staging = self._checkpoint_descriptor_buffer(descriptor_slot) - if total > staging.np.shape[0]: + staging = self._checkpoint_staging() + rows = staging.rows(descriptor_slot) + if total > rows.shape[0]: raise RuntimeError( f"a step asked to copy {total // per_op} checkpoints, more " - f"than the {staging.np.shape[0] // per_op} its descriptor was " + f"than the {rows.shape[0] // per_op} its descriptor was " "sized for" ) - descriptor = staging.np[:total] + descriptor = rows[:total] at = 0 for ops, storing in ((store_ops, True), (restore_ops, False)): if not ops: @@ -1009,7 +1011,7 @@ def execute_paged_state_copies( forward=storing, ) at = end - launch_copy_descriptor(staging.copy_to_gpu(total), plan) + launch_copy_descriptor(staging.upload(descriptor_slot, total), plan) def _validate_paged_state_op( self, op: CheckpointStoreOp | CheckpointRestoreOp @@ -1160,7 +1162,7 @@ def _invalidate_pool_caches(self) -> None: self._checkpoint_range_cache = None self._checkpoint_plan_cache = None self._checkpoint_slot_base_cache = None - self._checkpoint_descriptors = {} + self._checkpoint_staging_cache = None def warmup_per_req_cache(self) -> None: """Run one checkpoint copy now, so the first real one is only a copy. @@ -1183,13 +1185,13 @@ def warmup_per_req_cache(self) -> None: if not plan.num_spans: return units = self.model_runner.state_runtime.checkpoint_spec.units_per_checkpoint - staging = self._checkpoint_descriptor_buffer() + staging = self._checkpoint_staging() plan.write_descriptor( - staging.np[: plan.num_spans], + staging.rows(0)[: plan.num_spans], self._checkpoint_slot_bases()[:1], self._page_unit_bases([list(range(units))]), ) - launch_copy_descriptor(staging.copy_to_gpu(plan.num_spans), plan) + launch_copy_descriptor(staging.upload(0, plan.num_spans), plan) def _checkpoint_descriptor_device(self) -> torch.device: return self._kv_planes()[0].device diff --git a/atom/model_ops/attentions/gdn_attn.py b/atom/model_ops/attentions/gdn_attn.py index cde3a03876..e350d4acc9 100644 --- a/atom/model_ops/attentions/gdn_attn.py +++ b/atom/model_ops/attentions/gdn_attn.py @@ -832,14 +832,15 @@ def execute_paged_state_copies( slot_bases = self._checkpoint_slot_bases() per_op = plan.num_spans total = (len(store_ops) + len(restore_ops)) * per_op - staging = self._checkpoint_descriptor_buffer(descriptor_slot) - if total > staging.np.shape[0]: + staging = self._checkpoint_staging() + rows = staging.rows(descriptor_slot) + if total > rows.shape[0]: raise RuntimeError( f"a step asked to copy {total // per_op} checkpoints, more " - f"than the {staging.np.shape[0] // per_op} its descriptor was " + f"than the {rows.shape[0] // per_op} its descriptor was " "sized for" ) - descriptor = staging.np[:total] + descriptor = rows[:total] at = 0 for ops, storing in ((store_ops, True), (restore_ops, False)): if not ops: @@ -853,7 +854,7 @@ def execute_paged_state_copies( forward=storing, ) at = end - launch_copy_descriptor(staging.copy_to_gpu(total), plan) + launch_copy_descriptor(staging.upload(descriptor_slot, total), plan) def warmup_per_req_cache(self) -> None: """Run one checkpoint copy now, so the first real one is only a copy. @@ -876,13 +877,13 @@ def warmup_per_req_cache(self) -> None: plan = self._checkpoint_copy_plan() if not plan.num_spans: return - staging = self._checkpoint_descriptor_buffer() + staging = self._checkpoint_staging() plan.write_descriptor( - staging.np[: plan.num_spans], + staging.rows(0)[: plan.num_spans], self._checkpoint_slot_bases()[:1], self._page_unit_bases([list(range(spec.units_per_checkpoint))]), ) - launch_copy_descriptor(staging.copy_to_gpu(plan.num_spans), plan) + launch_copy_descriptor(staging.upload(0, plan.num_spans), plan) def state_entry_views(self, slot: int) -> list[torch.Tensor]: """One contiguous slice per (cache, layer) — the slot's whole state. diff --git a/atom/model_ops/attentions/pool_layout/paged_state_copy.py b/atom/model_ops/attentions/pool_layout/paged_state_copy.py index 3907511c9b..b066436a11 100644 --- a/atom/model_ops/attentions/pool_layout/paged_state_copy.py +++ b/atom/model_ops/attentions/pool_layout/paged_state_copy.py @@ -260,6 +260,67 @@ def _copy_tiles_kernel( _copy_tiles_kernel = None +class DescriptorStaging: + """Pinned copy-descriptor rows, one buffer per slot, each fenced on its upload. + + A non-blocking H2D reads its pinned source when the stream reaches the + copy, not when it is enqueued. Refilling a buffer whose last upload is + still queued would rewrite the descriptor under an earlier copy, so each + slot waits for its own previous upload before handing its rows out again. + + Slot 0 serves `build()`; every other slot is an independent buffer, so an + out-of-band restore never shares staging with a step. Buffers are reused + because allocating pinned memory synchronizes; `reserve` allocates them + before serving. + """ + + def __init__(self, rows: int, device: torch.device | str) -> None: + self._rows = rows + self._device = torch.device(device) + # slot -> (pinned host rows, device rows, upload fence) + self._slots: dict[ + int, tuple[torch.Tensor, torch.Tensor, torch.cuda.Event | None] + ] = {} + self._pending: set[int] = set() + + def reserve(self, slots: Sequence[int]) -> None: + for slot in slots: + self._slot(slot) + + def rows(self, slot: int = 0) -> np.ndarray: + """This slot's host rows, safe to overwrite.""" + host, _, uploaded = self._slot(slot) + if slot in self._pending: + uploaded.synchronize() + self._pending.discard(slot) + return host.numpy() + + def upload(self, slot: int, count: int) -> torch.Tensor: + """Copy the first ``count`` rows to the device on the current stream.""" + host, device, uploaded = self._slot(slot) + descriptor = device[:count].copy_(host[:count], non_blocking=True) + if uploaded is not None: + uploaded.record() + self._pending.add(slot) + return descriptor + + def _slot( + self, slot: int + ) -> tuple[torch.Tensor, torch.Tensor, torch.cuda.Event | None]: + if slot < 0: + raise ValueError("checkpoint descriptor slot must be non-negative") + entry = self._slots.get(slot) + if entry is None: + cuda = self._device.type == "cuda" + host = torch.zeros(self._rows, 3, dtype=torch.int64, pin_memory=cuda) + entry = self._slots[slot] = ( + host, + torch.zeros_like(host, device=self._device), + torch.cuda.Event() if cuda else None, + ) + return entry + + def launch_copy_descriptor(descriptor: torch.Tensor, plan: SegmentedCopyPlan) -> None: """Copy every span a resident descriptor names, in one launch. diff --git a/tests/attentions/deepseek_v41/test_cache.py b/tests/attentions/deepseek_v41/test_cache.py index fecd5705fd..40f0776087 100644 --- a/tests/attentions/deepseek_v41/test_cache.py +++ b/tests/attentions/deepseek_v41/test_cache.py @@ -451,19 +451,3 @@ def test_shared_indptr_launch_matches_row_counts_and_replay(decode, width, ratio dtype=torch.int32, ) torch.testing.assert_close(extend, expected, rtol=0, atol=0) - - -def test_state_copy_descriptor_slots_have_independent_staging(): - from atom.model_ops.attentions.deepseek_v41.checkpoints import StateCopies - - copies = StateCopies.__new__(StateCopies) - copies._staging_rows = 3 - copies.cache = type( - "Cache", - (), - {"pool": type("Pool", (), {"device": "cpu", "is_cuda": False})()}, - )() - copies._stagings = {0: copies._new_staging()} - assert copies.staging(0) is copies.staging(0) - assert copies.staging(1) is not copies.staging(0) - assert copies.staging(1) is copies.staging(1) diff --git a/tests/test_paged_state_copy_planner.py b/tests/test_paged_state_copy_planner.py index 1388c9ef56..f6393b0fde 100644 --- a/tests/test_paged_state_copy_planner.py +++ b/tests/test_paged_state_copy_planner.py @@ -13,6 +13,7 @@ import torch from atom.model_ops.attentions.pool_layout.paged_state_copy import ( + DescriptorStaging, launch_copy_descriptor, plan_segmented_copy, ) @@ -277,3 +278,41 @@ def test_a_flat_dst_bases_is_refused_by_this_guard_not_by_numpy(): with pytest.raises(ValueError, match="one row of bases per copy"): plan.write_descriptor(out, three, np.zeros(3, dtype=np.int64)) + + +def test_descriptor_slots_have_independent_staging(): + """Slot 0 is the step's; an out-of-band restore's slot never shares it.""" + staging = DescriptorStaging(3, "cpu") + staging.rows(0)[0, 0] = 7 + assert staging.rows(0)[0, 0] == 7 + assert staging.rows(1)[0, 0] == 0 + assert staging.rows(1).shape == (3, 3) + staging.rows(1)[:2] = 5 + assert staging.upload(1, 2).tolist() == [[5, 5, 5], [5, 5, 5]] + with pytest.raises(ValueError, match="non-negative"): + staging.rows(-1) + + +def test_descriptor_staging_waits_for_its_last_upload_before_reuse(monkeypatch): + """A pinned host buffer is rewritten only after the H2D that read it ran.""" + events = [] + + class Event: + def __init__(self): + self.log = [] + events.append(self) + + def record(self): + self.log.append("record") + + def synchronize(self): + self.log.append("synchronize") + + staging = DescriptorStaging(2, "cpu") + host = torch.zeros(2, 3, dtype=torch.int64) + staging._slots[4] = (host, torch.zeros_like(host), Event()) + staging.rows(4) + staging.upload(4, 1) + staging.rows(4) + staging.rows(4) + assert events[0].log == ["record", "synchronize"] From ec7e2eca7e6f003e8026d5204f701511b962b309 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Mon, 28 Sep 2026 04:24:19 +0000 Subject: [PATCH 23/28] refactor(lmcache_mp): validate PAGE views once for both registrations `build_native_state_mp_layout` and `_build_cache_views` each validated the backend-published PAGE views, and had already drifted: PAGE-only required full contiguity and compared total bytes against the view, native checked inner contiguity plus the block stride against the region. Both now call `validate_page_views` (new `mp/page_views.py`), which applies the stricter union -- shape, tight block-major stride, unit and total bytes, aliasing, forward indexing, one device -- and, for native, that a block's physical slots divide the block size. Each builder keeps only what is its own: PAGE-only rejects stateful fields; native checks the image coverage. Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/kv_transfer/offload/mp/backend.py | 50 +--------- .../offload/mp/native_state_layout.py | 47 ++-------- atom/kv_transfer/offload/mp/page_views.py | 94 +++++++++++++++++++ tests/test_lmcache_mp_native_layout.py | 24 +++++ 4 files changed, 132 insertions(+), 83 deletions(-) create mode 100644 atom/kv_transfer/offload/mp/page_views.py diff --git a/atom/kv_transfer/offload/mp/backend.py b/atom/kv_transfer/offload/mp/backend.py index 3c796f697d..d41ffa74ad 100644 --- a/atom/kv_transfer/offload/mp/backend.py +++ b/atom/kv_transfer/offload/mp/backend.py @@ -41,6 +41,7 @@ ChunkedOffloadSchedulerBase, ) from atom.kv_transfer.offload.metadata import LMCacheOffloadMetadata, LMCacheReqMeta +from atom.kv_transfer.offload.mp.page_views import validate_page_views from atom.utils import envs logger = logging.getLogger("atom") @@ -386,8 +387,6 @@ def _build_cache_views( if transfer_tensors is None: raise ValueError("lmcache_mp requires KVTransferTensors") - if num_blocks <= 0: - raise ValueError("lmcache_mp num_blocks must be positive") stateful_fields = { "num_slots": getattr(transfer_tensors, "num_slots", 0), @@ -409,49 +408,11 @@ def _build_cache_views( f"published through {', '.join(populated_stateful_fields)}" ) - regions = list(getattr(transfer_tensors, "block_regions", None) or []) - views = list(getattr(transfer_tensors, "block_tensor_views", None) or []) - if not regions or len(views) != len(regions): - raise ValueError( - "lmcache_mp requires one block_tensor_view per block region: " - f"views={len(views)} regions={len(regions)}" - ) - tensors: dict[str, torch.Tensor] = {} indices_by_layout: dict[tuple[torch.dtype, tuple[int, ...]], list[int]] = {} - devices: set[torch.device] = set() bytes_per_block = 0 - for index, (view, region) in enumerate(zip(views, regions, strict=True)): - if not isinstance(view, torch.Tensor): - raise TypeError(f"lmcache_mp block tensor view {index} is not a Tensor") - if view.numel() == 0 or not view.is_contiguous(): - raise ValueError( - f"lmcache_mp block tensor view {index} must be non-empty and contiguous" - ) - if view.ndim != 3 or int(view.shape[0]) != num_blocks: - raise ValueError( - f"lmcache_mp block tensor view {index} has shape={tuple(view.shape)}, " - f"expected [{num_blocks}, physical_slots, opaque_width]" - ) - - unit_bytes = int(region.unit_bytes) - total_bytes = int(region.total_bytes) - actual_unit_bytes = view[0].numel() * view.element_size() - actual_total_bytes = view.numel() * view.element_size() - if actual_unit_bytes != unit_bytes or actual_total_bytes != total_bytes: - raise ValueError( - f"lmcache_mp block tensor view {index} byte geometry mismatch: " - f"unit={actual_unit_bytes}/{unit_bytes} " - f"total={actual_total_bytes}/{total_bytes}" - ) - if view.data_ptr() != int(region.base_addr): - raise ValueError( - f"lmcache_mp block tensor view {index} does not start at its " - "declared region address" - ) - if bool(getattr(region, "reverse_indexed", False)): - raise ValueError("lmcache_mp PAGE regions cannot be reverse-indexed") - + for page in validate_page_views(transfer_tensors, num_blocks=num_blocks): + index, view, region = page.index, page.view, page.region # LMCache receives these tensors as opaque PAGE storage, not numerical # values. Publish a zero-copy byte view so every transfer path copies # the exact bit pattern. This is especially important on ROCm, where @@ -463,11 +424,8 @@ def _build_cache_views( tensors[f"page.{index}.{role}"] = byte_view layout = (byte_view.dtype, tuple(int(dim) for dim in byte_view.shape[1:])) indices_by_layout.setdefault(layout, []).append(index) - devices.add(view.device) - bytes_per_block += unit_bytes + bytes_per_block += page.unit_bytes - if len(devices) != 1: - raise ValueError("lmcache_mp block tensor views must share one device") return _CacheViews( tensors=tensors, layer_groups=tuple(tuple(indices) for indices in indices_by_layout.values()), diff --git a/atom/kv_transfer/offload/mp/native_state_layout.py b/atom/kv_transfer/offload/mp/native_state_layout.py index f7eed93981..9f78ae95ed 100644 --- a/atom/kv_transfer/offload/mp/native_state_layout.py +++ b/atom/kv_transfer/offload/mp/native_state_layout.py @@ -18,6 +18,8 @@ import torch +from atom.kv_transfer.offload.mp.page_views import validate_page_views + @dataclass(frozen=True) class NativeStateMPKernelGroup: @@ -120,12 +122,13 @@ def build_native_state_mp_layout( if published_blocks not in (0, num_blocks): raise ValueError("native published block count disagrees with scheduler") - regions = list(getattr(transfer_tensors, "block_regions", None) or []) - page_views = list(getattr(transfer_tensors, "block_tensor_views", None) or []) - if not regions or len(regions) != len(page_views): - raise ValueError( - "native-state LMCache MP needs one owned tensor view per PAGE region" - ) + # STATE aliases take one chunk per block, so a block's physical slots + # must divide the PAGE block size as well. + pages = validate_page_views( + transfer_tensors, num_blocks=num_blocks, block_size=block_size + ) + regions = [page.region for page in pages] + page_views = [page.view for page in pages] # Only the leading regions form a checkpoint PAGE unit; a draft's own pool # is appended after them and is stored as ordinary PAGE KV. state_count = getattr(transfer_tensors, "paged_state_region_count", None) @@ -134,37 +137,7 @@ def build_native_state_mp_layout( state_count = _positive_int("native state region count", state_count) if state_count > len(regions): raise ValueError("native state region count exceeds the PAGE regions") - devices = set() - actual_page_bytes = 0 - for index, (view, region) in enumerate(zip(page_views, regions, strict=True)): - if not isinstance(view, torch.Tensor): - raise TypeError(f"native PAGE view {index} must be a Tensor") - if view.ndim != 3 or view.shape[0] != num_blocks or view.numel() == 0: - raise ValueError(f"native PAGE view {index} has invalid block geometry") - if not view[0].is_contiguous(): - raise ValueError(f"native PAGE view {index} has non-contiguous inner rows") - if block_size % view.shape[1]: - raise ValueError( - f"native PAGE view {index} physical slots must divide block size" - ) - unit_bytes = _positive_int( - f"native PAGE region {index} bytes", region.unit_bytes - ) - if ( - view[0].numel() * view.element_size() != unit_bytes - or view.stride(0) * view.element_size() != unit_bytes - or region.total_bytes != num_blocks * unit_bytes - ): - raise ValueError(f"native PAGE view {index} byte geometry mismatch") - if view.data_ptr() != region.base_addr: - raise ValueError(f"native PAGE view {index} does not alias its region") - if region.reverse_indexed: - raise ValueError("native PAGE regions cannot be reverse-indexed") - devices.add(view.device) - if index < state_count: - actual_page_bytes += unit_bytes - if len(devices) != 1: - raise ValueError("native PAGE views must share one device") + actual_page_bytes = sum(page.unit_bytes for page in pages[:state_count]) if actual_page_bytes != page_bytes: raise ValueError("PAGE regions do not cover the native PAGE unit") diff --git a/atom/kv_transfer/offload/mp/page_views.py b/atom/kv_transfer/offload/mp/page_views.py new file mode 100644 index 0000000000..d921ca4a5f --- /dev/null +++ b/atom/kv_transfer/offload/mp/page_views.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""The one check that backend-published PAGE views are what they claim to be. + +Both LMCache MP registrations -- PAGE-only and native state -- hand the same +block-major views to the server, so they share this validation rather than +keeping two copies that drift apart. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch + + +@dataclass(frozen=True) +class PageView: + """One validated block-major PAGE view and the region it aliases.""" + + index: int + view: torch.Tensor + region: Any + unit_bytes: int + + +def validate_page_views( + transfer_tensors: Any, + *, + num_blocks: int, + block_size: int | None = None, +) -> list[PageView]: + """Validate ``block_tensor_views`` against ``block_regions``. + + Each view must be ``[num_blocks, physical_slots, opaque_width]``, tightly + block-major (one unit per block, contiguous inside it), exactly cover and + start at its region, be forward-indexed, and share one device with the + rest. ``block_size``, when given, also requires the physical slots of a + block to divide it. + """ + if transfer_tensors is None: + raise ValueError("lmcache_mp requires KVTransferTensors") + if type(num_blocks) is not int or num_blocks <= 0: + raise ValueError("lmcache_mp num_blocks must be a positive integer") + regions = list(getattr(transfer_tensors, "block_regions", None) or []) + views = list(getattr(transfer_tensors, "block_tensor_views", None) or []) + if not regions or len(views) != len(regions): + raise ValueError( + "lmcache_mp requires one block_tensor_view per block region: " + f"views={len(views)} regions={len(regions)}" + ) + + validated: list[PageView] = [] + devices: set[torch.device] = set() + for index, (view, region) in enumerate(zip(views, regions, strict=True)): + name = f"LMCache MP PAGE view {index}" + if not isinstance(view, torch.Tensor): + raise TypeError(f"{name} is not a Tensor") + if view.ndim != 3 or int(view.shape[0]) != num_blocks: + raise ValueError( + f"{name} has shape={tuple(view.shape)}, " + f"expected [{num_blocks}, physical_slots, opaque_width]" + ) + if view.numel() == 0 or not view[0].is_contiguous(): + raise ValueError(f"{name} must be non-empty and contiguous") + if block_size is not None and block_size % int(view.shape[1]): + raise ValueError(f"{name} physical slots must divide block size") + unit_bytes = int(region.unit_bytes) + actual_unit_bytes = view[0].numel() * view.element_size() + if ( + unit_bytes <= 0 + or actual_unit_bytes != unit_bytes + or view.stride(0) * view.element_size() != unit_bytes + or int(region.total_bytes) != num_blocks * unit_bytes + ): + raise ValueError( + f"{name} byte geometry mismatch: unit={actual_unit_bytes}/" + f"{unit_bytes} stride={view.stride(0) * view.element_size()} " + f"total={region.total_bytes}/{num_blocks * unit_bytes}" + ) + if view.data_ptr() != int(region.base_addr): + raise ValueError(f"{name} does not alias its declared region") + if bool(getattr(region, "reverse_indexed", False)): + raise ValueError("lmcache_mp PAGE regions cannot be reverse-indexed") + devices.add(view.device) + validated.append(PageView(index, view, region, unit_bytes)) + if len(devices) != 1: + raise ValueError("lmcache_mp PAGE views must share one device") + return validated + + +__all__ = ["PageView", "validate_page_views"] diff --git a/tests/test_lmcache_mp_native_layout.py b/tests/test_lmcache_mp_native_layout.py index dece1311e4..e7a29bac0f 100644 --- a/tests/test_lmcache_mp_native_layout.py +++ b/tests/test_lmcache_mp_native_layout.py @@ -374,3 +374,27 @@ def test_draft_regions_after_the_state_regions_stay_ordinary_page(): _gather(layout, with_draft, ids), _gather(_layout(target_only), target_only, ids), ) + + +@pytest.mark.parametrize( + ("breakage", "message"), + [ + (lambda t: setattr(t.block_regions[1], "unit_bytes", 9), "byte geometry"), + (lambda t: t.block_tensor_views.pop(), "one block_tensor_view per"), + (lambda t: setattr(t.block_regions[2], "reverse_indexed", True), "reverse"), + ], +) +def test_native_and_page_only_registration_share_page_validation(breakage, message): + """One validator for both registrations, so they cannot drift apart.""" + from atom.kv_transfer.offload.mp.backend import _build_cache_views + + native = _transfer() + breakage(native) + with pytest.raises(ValueError, match=message): + _layout(native) + page_only = _transfer() + page_only.paged_state_checkpoint_spec = None + page_only.execute_paged_state_copies = None + breakage(page_only) + with pytest.raises(ValueError, match=message): + _build_cache_views(page_only, num_blocks=7) From cf1fea47bd63e6d9f3bd74d7a38f665795abf2c0 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Mon, 28 Sep 2026 04:25:00 +0000 Subject: [PATCH 24/28] refactor(lmcache_mp): split the native worker's completion poll `get_finished` was 91 lines at six levels of nesting, with three near-identical recovery blocks. It is now a loop over `_poll_native_save` and `_poll_native_load`; the load's retrieve-then-restore progression lives in `_advance_native_load`, and both unprovable-restore paths share `_hold_unprovable_restore`. No behaviour change. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../offload/mp/native_state_worker.py | 199 ++++++++++-------- 1 file changed, 109 insertions(+), 90 deletions(-) diff --git a/atom/kv_transfer/offload/mp/native_state_worker.py b/atom/kv_transfer/offload/mp/native_state_worker.py index 9ce13eaff6..06667a3dcb 100644 --- a/atom/kv_transfer/offload/mp/native_state_worker.py +++ b/atom/kv_transfer/offload/mp/native_state_worker.py @@ -350,99 +350,118 @@ def get_finished(self) -> KVConnectorOutput: output = KVConnectorOutput() with self._lock: for operation_id, pending in list(self._native_saves.items()): - take_ranges = getattr(pending.future, "take_completed_ranges", None) - if callable(take_ranges): - try: - self._emit_page_source_safe( - output, pending, tuple(take_ranges()) - ) - except Exception: - logger.warning( - "LMCache MP source-safe event polling failed", - exc_info=True, - ) - if pending.immediate_success: - terminal, result = True, True - else: - terminal, result = _terminal_future_result(pending.future) - if not terminal: - _enforce_transfer_deadline( - operation_id, pending.started_at, self._transfer_deadline_s - ) - continue - save_spec = pending.request.save_spec - start = 0 if save_spec is None else int(save_spec.skip_leading_tokens) - end = len(pending.request.token_ids) - self._emit_page_source_safe( - output, pending, _chunk_ranges(start, end, self.chunk_size) - ) - output.connector_completions.add( - ConnectorCompletion( - NATIVE_STATE_MP_STORE_CHANNEL, - pending.request.save_operation, - succeeded=result is True, - ) - ) - del self._native_saves[operation_id] - _remember_operation_tombstone( - operation_id, - self._completed_save_operations, - self._completed_save_operation_order, - ) + self._poll_native_save(output, operation_id, pending) for operation_id, pending in list(self._native_loads.items()): - if pending.restore_event is None: - terminal, result = _terminal_future_result(pending.future) - if not terminal: - _enforce_transfer_deadline( - operation_id, pending.started_at, self._transfer_deadline_s - ) - continue - if result is True: - try: - if not self._begin_restore(pending): - continue - except Exception: - logger.exception( - "LMCache MP native restore safety unknown; " - "retaining lease" - ) - pending.future = _UnprovableSubmission() - pending.restore_event = None - continue - if pending.restore_event is not None: - try: - restored = pending.restore_event.query() - except Exception: - logger.exception( - "LMCache MP native restore event safety unknown; " - "retaining lease" - ) - pending.future = _UnprovableSubmission() - pending.restore_event = None - pending.restore_succeeded = False - restored = False - if not restored: - _enforce_transfer_deadline( - operation_id, pending.started_at, self._transfer_deadline_s - ) - continue - completion = pending.request.load_operation - if pending.descriptor_slot is not None: - self._restore_descriptor_slots.append(pending.descriptor_slot) - pending.descriptor_slot = None - target = ( - output.finished_loading - if pending.restore_succeeded - else output.failed_loading + self._poll_native_load(output, operation_id, pending) + return output + + def _poll_native_save( + self, output: KVConnectorOutput, operation_id: str, pending: _NativePending + ) -> None: + take_ranges = getattr(pending.future, "take_completed_ranges", None) + if callable(take_ranges): + try: + self._emit_page_source_safe(output, pending, tuple(take_ranges())) + except Exception: + logger.warning( + "LMCache MP source-safe event polling failed", exc_info=True ) - target.add(completion) - del self._native_loads[operation_id] - _remember_operation_tombstone( - operation_id, - self._completed_load_operations, - self._completed_load_operation_order, + if pending.immediate_success: + terminal, result = True, True + else: + terminal, result = _terminal_future_result(pending.future) + if not terminal: + _enforce_transfer_deadline( + operation_id, pending.started_at, self._transfer_deadline_s + ) + return + save_spec = pending.request.save_spec + start = 0 if save_spec is None else int(save_spec.skip_leading_tokens) + end = len(pending.request.token_ids) + self._emit_page_source_safe( + output, pending, _chunk_ranges(start, end, self.chunk_size) + ) + output.connector_completions.add( + ConnectorCompletion( + NATIVE_STATE_MP_STORE_CHANNEL, + pending.request.save_operation, + succeeded=result is True, + ) + ) + del self._native_saves[operation_id] + _remember_operation_tombstone( + operation_id, + self._completed_save_operations, + self._completed_save_operation_order, + ) + + def _poll_native_load( + self, output: KVConnectorOutput, operation_id: str, pending: _NativePending + ) -> None: + if not self._advance_native_load(operation_id, pending): + return + if pending.descriptor_slot is not None: + self._restore_descriptor_slots.append(pending.descriptor_slot) + pending.descriptor_slot = None + target = ( + output.finished_loading + if pending.restore_succeeded + else output.failed_loading + ) + target.add(pending.request.load_operation) + del self._native_loads[operation_id] + _remember_operation_tombstone( + operation_id, + self._completed_load_operations, + self._completed_load_operation_order, + ) + + def _advance_native_load(self, operation_id: str, pending: _NativePending) -> bool: + """Move a load toward its outcome; True once that outcome is final. + + A retrieve that fails is final at once. One that succeeds starts the + native restore, and is final only once the restore's event is done. + """ + if pending.restore_event is None: + terminal, result = _terminal_future_result(pending.future) + if not terminal: + _enforce_transfer_deadline( + operation_id, pending.started_at, self._transfer_deadline_s ) - return output + return False + if result is not True: + return True + try: + if not self._begin_restore(pending): + return False # every descriptor slot is busy; retry + except Exception: + logger.exception( + "LMCache MP native restore safety unknown; retaining lease" + ) + self._hold_unprovable_restore(pending) + return False + try: + restored = pending.restore_event.query() + except Exception: + logger.exception( + "LMCache MP native restore event safety unknown; retaining lease" + ) + self._hold_unprovable_restore(pending) + return False + if not restored: + _enforce_transfer_deadline( + operation_id, pending.started_at, self._transfer_deadline_s + ) + return restored + + @staticmethod + def _hold_unprovable_restore(pending: _NativePending) -> None: + """A restore whose GPU progress cannot be proven keeps its SLOT, lease + and descriptor slot (a queued copy may still read its staging); only + the transfer deadline ends it.""" + pending.future = _UnprovableSubmission() + pending.restore_event = None + pending.restore_succeeded = False __all__ = [ From ac5f2ddfa2f3562c01c3ea3c393359650a74ca3f Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Mon, 28 Sep 2026 04:29:06 +0000 Subject: [PATCH 25/28] refactor(lmcache_mp): split mp/backend.py by responsibility `mp/backend.py` had grown to 1332 lines holding configuration, PAGE view validation, transfer bookkeeping, lookups and both connector halves. It is split without behaviour change into: - `deployment.py`: config validation, TP/DP topology and rank collapse, model namespace, server adapters - `transfer.py`: operation identity, terminal detection, the transfer deadline and `LMCacheTransferUnprovable` - `lookup.py`: the scheduler's lookup client and read-lock bookkeeping - `page_views.py`: PAGE-only `_build_cache_views`, beside the shared view validation it uses - `worker.py` / `scheduler.py`: the PAGE-only connector halves The native modules, the public shell and the tests import from the new homes; tests patch each name where it is looked up. Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/kv_transfer/offload/README.md | 6 +- atom/kv_transfer/offload/mp/backend.py | 1375 ----------------- atom/kv_transfer/offload/mp/connector.py | 12 +- atom/kv_transfer/offload/mp/deployment.py | 342 ++++ atom/kv_transfer/offload/mp/lookup.py | 175 +++ .../offload/mp/native_state_scheduler.py | 7 +- .../offload/mp/native_state_worker.py | 20 +- atom/kv_transfer/offload/mp/page_views.py | 64 + atom/kv_transfer/offload/mp/scheduler.py | 236 +++ atom/kv_transfer/offload/mp/transfer.py | 179 +++ atom/kv_transfer/offload/mp/worker.py | 490 ++++++ tests/test_lmcache_mp.py | 205 +-- tests/test_lmcache_mp_builder_views.py | 2 +- tests/test_lmcache_mp_native_layout.py | 2 +- tests/test_lmcache_mp_native_scheduler.py | 15 +- tests/test_lmcache_mp_native_worker.py | 25 +- tests/test_lmcache_mp_shell.py | 21 +- 17 files changed, 1653 insertions(+), 1523 deletions(-) delete mode 100644 atom/kv_transfer/offload/mp/backend.py create mode 100644 atom/kv_transfer/offload/mp/deployment.py create mode 100644 atom/kv_transfer/offload/mp/lookup.py create mode 100644 atom/kv_transfer/offload/mp/scheduler.py create mode 100644 atom/kv_transfer/offload/mp/transfer.py create mode 100644 atom/kv_transfer/offload/mp/worker.py diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index 416e68fa7b..31709d9c0f 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -104,7 +104,11 @@ Four rules carry the module: | `hybrid/kimi_k3/state_tier.py` | Worker-side store/load driver for the state tier on its own executor; reports store/finished/failed hash sets for the engine-side `StateOffloadIndex` to apply. | | `atom_lmcache_staging.py` | Per-thread CUDA streams, staging buffer, ready/free events, env helpers. | | `mp/connector.py` | Capability-selected public `lmcache_mp` worker/scheduler shells. | -| `mp/backend.py` | Generic PAGE-only LMCache multiprocess transport and adapter integration. | +| `mp/deployment.py` | LMCache MP configuration, topology, model namespace and server adapters. | +| `mp/worker.py`, `mp/scheduler.py` | Generic PAGE-only LMCache MP connector halves. | +| `mp/lookup.py` | Scheduler-side MP lookups and read-lock bookkeeping. | +| `mp/transfer.py` | Transfer identity, terminal detection and the fail-stop transfer deadline. | +| `mp/page_views.py` | Shared validation of backend-published PAGE views. | | `mp/native_state_{layout,scheduler,worker}.py` | PAGE-backed native-state registration, scheduler leases, and worker transfer/restore lifecycle. | The engine-side counterpart of the state tier lives outside this directory: diff --git a/atom/kv_transfer/offload/mp/backend.py b/atom/kv_transfer/offload/mp/backend.py deleted file mode 100644 index d41ffa74ad..0000000000 --- a/atom/kv_transfer/offload/mp/backend.py +++ /dev/null @@ -1,1375 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""Generic PAGE adapter for LMCache's standalone multiprocess server. - -Attention backends publish block-major tensor views through -``KVTransferTensors``. LMCache MP can attach several physical kernel groups to -the same engine block-id space, so this adapter registers those opaque views -without knowing which model or attention implementation produced them. -""" - -from __future__ import annotations - -import hashlib -import json -import logging -import threading -import time -from collections import deque -from collections.abc import Iterable, Sequence -from dataclasses import dataclass, field, replace -from typing import Any - -import torch - -from atom.kv_transfer.disaggregation.base import KVConnectorBase -from atom.kv_transfer.disaggregation.types import ( - ConnectorCompletion, - KVConnectorOutput, - LoadCompletionId, - LoadOperationId, - SaveCompletionId, - SaveOperationId, - SaveSourceGroupId, -) -from atom.kv_transfer.offload import config as offcfg -from atom.kv_transfer.offload._offload_common import validated_kv_role -from atom.kv_transfer.offload.chunked_scheduler import ( - DENSE_PAGE_SOURCE_SAFE_CHANNEL, - DENSE_PAGE_STORE_CHANNEL, - ChunkedOffloadSchedulerBase, -) -from atom.kv_transfer.offload.metadata import LMCacheOffloadMetadata, LMCacheReqMeta -from atom.kv_transfer.offload.mp.page_views import validate_page_views -from atom.utils import envs - -logger = logging.getLogger("atom") - -_MP_LAYOUT_VERSION = 3 -_OPERATION_TOMBSTONE_LIMIT = 4096 - - -def _extra_config(config: Any) -> dict[str, Any]: - kvc = getattr(config, "kv_transfer_config", {}) or {} - extra = kvc.get("kv_connector_extra_config", {}) or {} - if not isinstance(extra, dict): - raise TypeError("kv_connector_extra_config must be a dictionary") - return extra - - -def _storage_kv_transfer_config(config: Any) -> dict[str, Any]: - """Remove MP transport-only options before LMCache storage parsing.""" - - kvc = dict(getattr(config, "kv_transfer_config", {}) or {}) - extra = kvc.get("kv_connector_extra_config") - if isinstance(extra, dict): - kvc["kv_connector_extra_config"] = { - key: value - for key, value in extra.items() - if not (isinstance(key, str) and key.startswith("lmcache.mp.")) - } - return kvc - - -def _mp_session_id(config: Any, request_id: Any) -> str: - """Scope one LMCache MP request session to its global DP replica.""" - - return f"{offcfg.lmcache_engine_id(config)}:{request_id}" - - -def _transfer_mode(config: Any) -> str: - extra = _extra_config(config) - configured_mode = extra.get("lmcache.mp.mp_transfer_mode") - if configured_mode is None: - configured_mode = envs.LMCACHE_MP_TRANSFER_MODE - transfer_mode = str(configured_mode).strip().lower() - if transfer_mode not in ("auto", "lmcache_driven", "engine_driven"): - raise ValueError( - "LMCache MP transfer mode must be 'auto', 'lmcache_driven', or " - f"'engine_driven', got {configured_mode!r}" - ) - if transfer_mode == "engine_driven": - raise NotImplementedError( - "ATOM lmcache_mp requires LMCache's lmcache_driven transfer path " - "because engine_driven does not support multiple physical " - "cache groups" - ) - return transfer_mode - - -def _validate_mp_config(config: Any) -> tuple[int, int]: - """Validate topology constraints shared by generic PAGE layouts.""" - tp_size = offcfg._strict_integer( - "tensor_parallel_size", - getattr(config, "tensor_parallel_size", 1) or 1, - minimum=1, - ) - pp_size = offcfg._strict_integer( - "pipeline_parallel_size", - getattr(config, "pipeline_parallel_size", 1) or 1, - minimum=1, - ) - dcp_size = offcfg._strict_integer( - "decode_context_parallel_size", - getattr(config, "decode_context_parallel_size", 1) or 1, - minimum=1, - ) - pcp_size = offcfg._strict_integer( - "prefill_context_parallel_size", - getattr(config, "prefill_context_parallel_size", 1) or 1, - minimum=1, - ) - parallel_config = getattr(config, "parallel_config", None) - dp_size = offcfg._strict_integer( - "data_parallel_size", - getattr( - parallel_config, - "data_parallel_size", - getattr(config, "data_parallel_size", 1), - ) - or 1, - minimum=1, - ) - dp_size_local = offcfg._strict_integer( - "data_parallel_size_local", - getattr(parallel_config, "data_parallel_size_local", dp_size) or dp_size, - minimum=1, - ) - if pp_size != 1: - raise NotImplementedError("lmcache_mp does not support PP yet") - if dcp_size != 1: - raise NotImplementedError("lmcache_mp does not support DCP yet") - if pcp_size != 1: - raise NotImplementedError("lmcache_mp does not support PCP yet") - # Single-host DP replicas deliberately share one (model_name, worker_id, - # world_size) identity: it is the content-addressed storage namespace, so - # replicas deduplicate identical prefixes. It is not a registration key. - # The server registers GPU memory per unique instance_id, refcounts layout - # descriptors per (model_name, world_size) (all replicas publish the same - # one), and _mp_session_id scopes request sessions and their locks per - # replica. - if dp_size_local != dp_size: - raise NotImplementedError( - "lmcache_mp supports DP and DP-attention only within one host; " - "multi-node DP requires one LMCache server per host and local " - "server routing " - f"(data_parallel_size={dp_size}, local={dp_size_local})" - ) - - _transfer_mode(config) - return tp_size, pp_size - - -def _config_has_fully_replicated_tp_pages(config: Any) -> bool: - """Conservatively identify configs whose complete PAGE cache is TP-replicated. - - MLA caches the shared latent before the TP-sharded KV-B projection. ATOM's - sparse MLA index-key projection is replicated too, so its auxiliary PAGE - plane has the same property. - - The worker validates this config-time prediction against the attention - backend's ``KVTransferTensors.tp_replication_factor`` declaration before - registering any cache memory. - """ - - if _config_has_own_pool_draft(config): - # A DSpark draft whose backend owns a KV pool appends per-rank-sharded - # PAGE regions (`draft_kv.py` declares factor 1), so the complete PAGE - # object is not replicated even when the target's is. - return False - hf_config = getattr(config, "hf_config", None) - # MiniMax-M3 is GQA. Some TP ranks can happen to own the same KV head when - # TP exceeds the global KV-head count, but the complete PAGE object is not - # replicated across the whole TP group and must remain one shard per rank. - if offcfg._is_minimax_m3(hf_config): - return False - hf_config = getattr(hf_config, "text_config", hf_config) - return getattr(hf_config, "kv_lora_rank", None) is not None - - -def _config_has_own_pool_draft(config: Any) -> bool: - """Whether a speculative draft caches into a KV pool of its own. - - Mirrors `spec_decode.draft_kv.draft_kv_builder`, which only DSpark calls: - the draft's own backend answers through `DRAFT_OWNS_KV_POOL`. - """ - - speculative = getattr(config, "speculative_config", None) - draft_hf = getattr(speculative, "draft_model_hf_config", None) - if getattr(speculative, "method", None) != "dspark" or draft_hf is None: - return False - from atom.utils.selector import attn_family, get_attn_backend - - return bool(get_attn_backend(attn_family(draft_hf)).DRAFT_OWNS_KV_POOL) - - -def _tp_replication_factor(config: Any) -> int: - """Return the PAGE rank-collapse factor selected before workers start. - - ``auto`` uses only structural cache information available in the shared - engine config. An explicit boolean is useful for a new attention backend: - ``True`` requests full TP collapse, but worker registration still fails - closed unless that backend declares every PAGE region byte-identical. - """ - - tp_size, _ = _validate_mp_config(config) - configured = _extra_config(config).get("lmcache.mp.tp_rank_collapse", "auto") - if isinstance(configured, str) and configured.strip().lower() == "auto": - collapse = _config_has_fully_replicated_tp_pages(config) - elif type(configured) is bool: - collapse = configured - else: - raise TypeError("lmcache.mp.tp_rank_collapse must be true, false, or 'auto'") - return tp_size if collapse else 1 - - -def _published_tp_replication_factor( - transfer_tensors: Any, - *, - tp_size: int, - native_state: bool = False, -) -> int: - """Validate a backend's whole-object TP replication declaration.""" - - attribute = ( - "native_state_tp_replication_factor" - if native_state - else "tp_replication_factor" - ) - label = "native STATE" if native_state else "KV PAGE" - factor = offcfg._strict_integer( - f"{label} TP replication factor", - getattr(transfer_tensors, attribute, 1), - minimum=1, - ) - if tp_size % factor: - raise ValueError( - f"{label} TP replication factor {factor} must divide TP size {tp_size}" - ) - if factor not in (1, tp_size): - raise NotImplementedError( - "lmcache_mp currently supports only sharded or fully TP-replicated " - f"{label} layouts, got replication factor {factor} for TP size {tp_size}" - ) - return factor - - -def _server_urls(config: Any) -> list[str]: - extra = _extra_config(config) - configured = extra.get("lmcache.mp.server_urls") - if configured is not None: - if isinstance(configured, (list, tuple)): - urls = [str(value).strip() for value in configured if str(value).strip()] - else: - urls = [ - value.strip() for value in str(configured).split(",") if value.strip() - ] - else: - host = str(extra.get("lmcache.mp.host", "tcp://localhost")).strip() - if not host: - raise ValueError("lmcache.mp.host must be non-empty") - port = offcfg._strict_integer( - "lmcache.mp.port", - extra.get("lmcache.mp.port", 5555), - minimum=1, - ) - if not 1 <= port <= 65535: - raise ValueError("lmcache.mp.port must be in [1, 65535]") - urls = [f"{host}:{port}"] - urls = [url if "://" in url else f"tcp://{url}" for url in urls] - if len(urls) != 1: - raise NotImplementedError( - "lmcache_mp currently supports exactly one LMCache server" - ) - return urls - - -def _model_namespace(config: Any, *, checkpoint_spec: Any = None) -> str: - """Build a model/layout namespace shared by scheduler and workers.""" - - cfg = offcfg.build_lmcache_config(_storage_kv_transfer_config(config)) - world_size = offcfg.lmcache_replica_world_size(config) - page_namespace = offcfg.build_page_namespace( - config, - cfg, - world_size, - ) - namespace = f"{page_namespace}::lmcache-mp-v{_MP_LAYOUT_VERSION}" - if checkpoint_spec is not None: - hf = getattr(config, "hf_config", None) - hf = getattr(hf, "text_config", hf) - document = { - "checkpoint": checkpoint_spec.to_wire(), - "hf_commit": getattr(hf, "_commit_hash", None), - "revision": getattr(config, "revision", None), - "model_revision": _extra_config(config).get("lmcache.mp.model_revision"), - } - fingerprint = hashlib.sha256( - json.dumps(document, sort_keys=True).encode() - ).hexdigest()[:32] - namespace += f"::native-state-v1-{fingerprint}" - return namespace - - -def _parallel_strategy(config: Any, worker_id: int) -> Any: - from lmcache.integration.atom import AtomMPParallelConfig - - tp_size, pp_size = _validate_mp_config(config) - if worker_id < 0 or worker_id >= tp_size: - raise ValueError( - f"LMCache MP worker rank {worker_id} is outside [0, {tp_size})" - ) - replication_factor = _tp_replication_factor(config) - return AtomMPParallelConfig( - world_size=tp_size * pp_size // replication_factor, - worker_id=worker_id // replication_factor, - tp_size=tp_size, - ) - - -def _make_scheduler_adapter(config: Any, *, checkpoint_spec: Any = None) -> Any: - import zmq - from lmcache.integration.atom import AtomMPSchedulerAdapter - - num_kv_readers = _tp_replication_factor(config) - - class _ReaderAwareSchedulerAdapter(AtomMPSchedulerAdapter): - """Reserve one LMCache read lock for every collapsed TP consumer.""" - - def _create_key(self, *args: Any, **kwargs: Any) -> Any: - key = super()._create_key(*args, **kwargs) - return replace(key, num_kv_readers=num_kv_readers) - - extra = _extra_config(config) - return _ReaderAwareSchedulerAdapter( - server_url=_server_urls(config)[0], - context=zmq.Context.instance(), - model_name=_model_namespace(config, checkpoint_spec=checkpoint_spec), - block_size=int(config.kv_cache_block_size), - parallel_config=_parallel_strategy(config, 0), - mq_timeout=float(extra.get("lmcache.mp.mq_timeout", 300.0)), - ) - - -def _make_worker_adapter(config: Any, rank: int, *, checkpoint_spec: Any = None) -> Any: - import zmq - from lmcache.integration.atom import AtomMPWorkerAdapter - - extra = _extra_config(config) - return AtomMPWorkerAdapter( - server_url=_server_urls(config)[0], - context=zmq.Context.instance(), - model_name=_model_namespace(config, checkpoint_spec=checkpoint_spec), - block_size=int(config.kv_cache_block_size), - parallel_config=_parallel_strategy(config, rank), - mq_timeout=float(extra.get("lmcache.mp.mq_timeout", 300.0)), - heartbeat_interval=float(extra.get("lmcache.mp.heartbeat_interval", 10.0)), - transfer_mode=_transfer_mode(config), - ) - - -@dataclass(frozen=True) -class _CacheViews: - """Validated block-major tensors and their copy-kernel group indices.""" - - tensors: dict[str, torch.Tensor] - layer_groups: tuple[tuple[int, ...], ...] - bytes_per_block: int - - -def _build_cache_views( - transfer_tensors: Any, - *, - num_blocks: int, -) -> _CacheViews: - """Validate backend-published PAGE views without inspecting model internals.""" - - if transfer_tensors is None: - raise ValueError("lmcache_mp requires KVTransferTensors") - - stateful_fields = { - "num_slots": getattr(transfer_tensors, "num_slots", 0), - "slot_regions": getattr(transfer_tensors, "slot_regions", None), - "swa_block_regions": getattr(transfer_tensors, "swa_block_regions", None), - "staging_region": getattr(transfer_tensors, "staging_region", None), - "gather_slot": getattr(transfer_tensors, "gather_slot", None), - "scatter_slot": getattr(transfer_tensors, "scatter_slot", None), - "expected_full_slot_region_count": getattr( - transfer_tensors, "expected_full_slot_region_count", None - ), - } - populated_stateful_fields = [ - name for name, value in stateful_fields.items() if value - ] - if populated_stateful_fields: - raise NotImplementedError( - "lmcache_mp supports PAGE-only layouts; stateful SLOT data was " - f"published through {', '.join(populated_stateful_fields)}" - ) - - tensors: dict[str, torch.Tensor] = {} - indices_by_layout: dict[tuple[torch.dtype, tuple[int, ...]], list[int]] = {} - bytes_per_block = 0 - for page in validate_page_views(transfer_tensors, num_blocks=num_blocks): - index, view, region = page.index, page.view, page.region - # LMCache receives these tensors as opaque PAGE storage, not numerical - # values. Publish a zero-copy byte view so every transfer path copies - # the exact bit pattern. This is especially important on ROCm, where - # LMCache's Python raw-pointer fallback cannot express FP8 through the - # CUDA array interface and would otherwise reconstruct the destination - # as uint8 while keeping the staging object as FP8. - byte_view = view.view(torch.uint8) - role = str(getattr(region, "semantic_role", None) or f"plane_{index}") - tensors[f"page.{index}.{role}"] = byte_view - layout = (byte_view.dtype, tuple(int(dim) for dim in byte_view.shape[1:])) - indices_by_layout.setdefault(layout, []).append(index) - bytes_per_block += page.unit_bytes - - return _CacheViews( - tensors=tensors, - layer_groups=tuple(tuple(indices) for indices in indices_by_layout.values()), - bytes_per_block=bytes_per_block, - ) - - -@dataclass -class _LookupState: - token_ids: list[int] - hit: int | None = None - retrieve_start: int | None = None - retrieve_end: int | None = None - - -@dataclass -class _PendingLoad: - completion: LoadCompletionId - future: Any | None - started_at: float = field(default_factory=lambda: time.monotonic()) - - -@dataclass -class _PendingSave: - completion: SaveCompletionId - future: Any | None - start: int - end: int - started_at: float = field(default_factory=lambda: time.monotonic()) - - -_DEFAULT_TRANSFER_DEADLINE_S = 1200.0 -# The scheduler watches the same operations from dispatch, a step or more -# before a worker submits them. The margin lets the worker, which can name the -# exact transfer, fail first. -_SCHEDULER_DEADLINE_MARGIN_S = 60.0 - - -class LMCacheTransferUnprovable(RuntimeError): - """An MP transfer never reached a terminal state within the deadline.""" - - -def _transfer_deadline_s(config: Any) -> float: - """Seconds any MP transfer may stay non-terminal before the engine stops. - - Engine memory under a transfer is released only on a terminal report: a - timeout cannot prove that the server stopped its DMA, so freeing on a - clock could hand live blocks to another request. A transfer that never - reports is therefore fatal once this deadline passes -- fail-stop, rather - than corrupt memory or wedge the pool with no fault to point at. - """ - - extra = _extra_config(config) - deadline = float( - extra.get("lmcache.mp.transfer_deadline_s", _DEFAULT_TRANSFER_DEADLINE_S) - ) - if not deadline > 0: - raise ValueError("lmcache.mp.transfer_deadline_s must be > 0") - return deadline - - -def _enforce_transfer_deadline( - operation: Any, started_at: float, deadline_s: float -) -> None: - """Raise ``LMCacheTransferUnprovable`` once a transfer outlives the bound.""" - - elapsed = time.monotonic() - started_at - if elapsed < deadline_s: - return - message = ( - f"LMCache MP transfer {operation} is still not terminal after " - f"{elapsed:.0f}s. The memory it reads or writes cannot be released " - "without proof that the server stopped, so the engine is stopping. " - "Raise lmcache.mp.transfer_deadline_s if the tier is only slow." - ) - logger.error(message) - raise LMCacheTransferUnprovable(message) - - -class _UnprovableSubmission: - """Stand-in future for a submission whose transport call raised. - - The server may have taken the request before the connection failed, so - this never becomes terminal on its own: the lease is kept and only the - transfer deadline ends it. - """ - - def query(self) -> bool: - return False - - def result(self, timeout: float | None = None) -> bool: - del timeout - return False - - -def _terminal_future_result(future: Any | None) -> tuple[bool, Any]: - """Return ``(terminal, result)`` without blocking on a device future.""" - - if future is None: - return True, None - try: - if not future.query(): - return False, None - return True, future.result(timeout=0) - except Exception: - # A query/result exception is not proof that the remote GPU stream has - # quiesced. Keep the operation pending; the transfer deadline bounds a - # future that keeps raising. - logger.warning( - "LMCache MP transfer future could not prove terminal state", - exc_info=True, - ) - return False, None - - -def _transfer_operation_id(kind: str, completion: Any) -> str: - """Return a stable LMCache handle ID for one ATOM completion generation.""" - - request_id = getattr(completion, "req_id", completion) - generation = getattr(completion, "generation", None) - if generation is None: - return f"{kind}:{request_id}" - return f"{kind}:{request_id}:{generation}" - - -def _remember_operation_tombstone( - operation_id: str, - tombstones: set[str], - tombstone_order: deque[str], -) -> None: - if operation_id in tombstones: - return - tombstones.add(operation_id) - tombstone_order.append(operation_id) - while len(tombstone_order) > _OPERATION_TOMBSTONE_LIMIT: - tombstones.discard(tombstone_order.popleft()) - - -def _chunk_ranges(start: int, end: int, chunk_size: int) -> tuple[tuple[int, int], ...]: - """Split ``[start, end)`` into LMCache chunk token ranges.""" - - return tuple( - (chunk_start, min(chunk_start + chunk_size, end)) - for chunk_start in range(start, end, chunk_size) - ) - - -def _source_safe_completions( - operation: SaveOperationId, ranges: Iterable[Sequence[int]] -) -> set[ConnectorCompletion]: - """One PAGE source-safe completion per finished token range of a save.""" - - return { - ConnectorCompletion( - DENSE_PAGE_SOURCE_SAFE_CHANNEL, - SaveSourceGroupId(operation, (tuple(token_range),)), - True, - ) - for token_range in ranges - } - - -class _MPLookupClient: - """Synchronous lookup facade used by ATOM's existing scheduler policy.""" - - token_database = None - - def __init__( - self, - adapter: Any, - *, - config: Any, - timeout: float, - poll_interval: float, - ) -> None: - if timeout <= 0 or poll_interval <= 0: - raise ValueError("LMCache MP lookup timeout and poll interval must be > 0") - self._adapter = adapter - self._config = config - self._timeout = timeout - self._poll_interval = poll_interval - self._lookups: dict[str, _LookupState] = {} - - def lookup(self, token_ids: list[int], lookup_id: str) -> int | None: - """Hit length for this prompt, or None if the tier never answered. - - None is a non-answer, not an empty answer: the caller must ask again - rather than record "this tier has nothing" for a prompt the tier may - well hold. - """ - - state = _LookupState(token_ids=list(token_ids)) - self._lookups[lookup_id] = state - request_id = _mp_session_id(self._config, lookup_id) - self._adapter.maybe_submit_lookup_request(request_id, token_ids) - deadline = time.monotonic() + self._timeout - while True: - result = self._adapter.check_lookup_result(request_id) - if result is not None: - hit = int(result) - state.hit = hit - return hit - if time.monotonic() >= deadline: - logger.warning( - "LMCache MP lookup timed out after %.1fs for request %s", - self._timeout, - lookup_id, - ) - # The MP API has no cancel-prefetch call. Keep the adapter job - # intact so request_finished() can release locks if the result - # becomes available; eagerly cleaning it here would orphan the - # server-side lookup and its locks. `state.hit` stays None, - # which is what tells clear_lookup_status() the job is still - # outstanding. - return None - time.sleep(self._poll_interval) - - def prepare_retrieve( - self, - lookup_id: str, - start: int, - end: int | None = None, - ) -> None: - """Hand one subrange of a lookup hit to the worker retrieve. - - The worker owns and releases locks in ``[start, end)``. Any hit prefix - already resident in HBM and any hit suffix beyond ``end`` will not be - consumed by that retrieve, so release those locks here. - """ - if start < 0 or (end is not None and end < start): - raise ValueError( - f"invalid retrieve range for {lookup_id}: start={start}, end={end}" - ) - state = self._lookups.get(lookup_id) - request_id = _mp_session_id(self._config, lookup_id) - if ( - state is not None - and state.hit is not None - and end is not None - and end > state.hit - ): - raise ValueError( - f"retrieve end {end} exceeds lookup hit {state.hit} for {lookup_id}" - ) - if state is not None and state.hit is not None and start > 0: - self._adapter.free_lookup_locks( - token_ids=state.token_ids, - start=0, - end=min(start, state.hit), - request_id=request_id, - ) - if ( - state is not None - and state.hit is not None - and end is not None - and end < state.hit - ): - self._adapter.free_lookup_locks( - token_ids=state.token_ids, - start=end, - end=state.hit, - request_id=request_id, - ) - if state is not None: - state.retrieve_start = start - state.retrieve_end = end - self._adapter.cleanup_lookup_result(request_id) - - def complete_retrieve(self, lookup_id: str, *, succeeded: bool) -> None: - # Once a retrieve is submitted, LMCache owns the remaining lookup - # locks. Its lmcache-driven transfer releases them on both success and - # failure (including partial failures). Releasing the range again from - # the scheduler would decrement every TP rank's read locks twice. The - # scheduler cannot distinguish a pre-submit failure; that rarer case is - # left to request end_session/server TTL cleanup instead. - self._lookups.pop(lookup_id, None) - self._adapter.cleanup_lookup_result(_mp_session_id(self._config, lookup_id)) - - def hit_tokens(self, lookup_id: str) -> int | None: - state = self._lookups.get(lookup_id) - return None if state is None else state.hit - - def clear_lookup_status(self, lookup_id: str) -> None: - state = self._lookups.pop(lookup_id, None) - request_id = _mp_session_id(self._config, lookup_id) - if state is not None and state.hit is None: - result = self._adapter.check_lookup_result(request_id) - if result is None: - logger.warning( - "LMCache MP lookup for request %s is still pending during " - "cleanup; dropping local state while server TTL/session " - "cleanup releases any eventual locks", - lookup_id, - ) - self._adapter.cleanup_lookup_result(request_id) - return - state.hit = int(result) - # Once retrieve has started, the transfer owns the remaining read - # locks. Failed terminal loads call complete_retrieve(False) first. - if ( - state is not None - and state.retrieve_start is None - and state.hit - and state.hit > 0 - ): - self._adapter.free_lookup_locks( - token_ids=state.token_ids, - start=0, - end=state.hit, - request_id=request_id, - ) - self._adapter.cleanup_lookup_result(request_id) - - -class LMCacheMPConnector(KVConnectorBase): - """Worker-side LMCache MP connector for backend-published PAGE views.""" - - is_producer = False - - def __init__(self, config: Any) -> None: - _validate_mp_config(config) - self._config = config - kvc = getattr(config, "kv_transfer_config", {}) or {} - self.kv_role = validated_kv_role(kvc) - self._do_save = self.kv_role in ("offload", "kv_both", "kv_producer") - self._do_load = self.kv_role in ("offload", "kv_both", "kv_consumer") - self.block_size = offcfg._strict_integer( - "LMCache MP block size", - config.kv_cache_block_size, - minimum=1, - ) - self.chunk_size: int | None = None - self._adapter: Any = None - self._is_kv_writer = True - self._pending_saves: dict[str, _PendingSave] = {} - self._pending_loads: dict[str, _PendingLoad] = {} - self._submitting_saves: set[str] = set() - self._submitting_loads: set[str] = set() - # Keep in-flight IDs in the unbounded pending/submitting collections. - # Only completed IDs enter these bounded replay tombstones. - self._completed_save_operations: set[str] = set() - self._completed_load_operations: set[str] = set() - self._completed_save_operation_order: deque[str] = deque() - self._completed_load_operation_order: deque[str] = deque() - # Immediate successes include collapsed-TP non-writers. Keep their - # logical token range so they can report PAGE source-safety too: the TP - # aggregator must receive the same chunk completion from every rank - # before releasing the writer's source blocks early. - self._immediate_saves: dict[SaveCompletionId, tuple[int, int] | None] = {} - self._transfer_deadline_s = _transfer_deadline_s(config) - self._immediate_load_failures: set[LoadCompletionId] = set() - self._lock = threading.Lock() - - def register_kv_caches( - self, - _kv_caches: dict[str, Any], - transfer_tensors: Any = None, - num_blocks: int | None = None, - ) -> None: - if num_blocks is None: - num_blocks = getattr(transfer_tensors, "num_blocks", None) - if num_blocks is None: - raise ValueError("lmcache_mp requires the scheduler block count") - normalized_num_blocks = offcfg._strict_integer( - "LMCache MP block count", - num_blocks, - minimum=1, - ) - - from aiter.dist.parallel_state import get_tp_group - from lmcache.v1.multiprocess.group_view import EngineGroupInfo - - tp = get_tp_group() - rank = int(tp.rank_in_group) - tp_size, _ = _validate_mp_config(self._config) - requested_replication = _tp_replication_factor(self._config) - published_replication = _published_tp_replication_factor( - transfer_tensors, - tp_size=tp_size, - ) - if requested_replication > published_replication: - raise ValueError( - "LMCache MP TP rank collapse was requested, but the attention " - "backend did not declare the complete PAGE layout fully " - f"replicated (published factor={published_replication}, " - f"TP size={tp_size})" - ) - self._is_kv_writer = rank % requested_replication == 0 - views = _build_cache_views( - transfer_tensors, - num_blocks=normalized_num_blocks, - ) - block_regions = getattr(transfer_tensors, "block_regions", None) or [] - expected = sum(int(region.unit_bytes) for region in block_regions) - if expected != views.bytes_per_block: - raise ValueError( - "lmcache_mp block geometry mismatch: " - f"views={views.bytes_per_block} transfer_regions={expected}" - ) - - adapter = _make_worker_adapter(self._config, rank) - groups = [ - EngineGroupInfo( - engine_group_id=0, - layer_indices=layer_indices, - tokens_per_block=self.block_size, - ) - for layer_indices in views.layer_groups - ] - try: - adapter.register_kv_caches(views.tensors, engine_group_infos=groups) - chunk_size = offcfg._strict_integer( - "LMCache MP chunk size", - adapter.lmcache_tokens_per_chunk, - minimum=1, - ) - if chunk_size % self.block_size: - raise ValueError( - f"LMCache MP chunk size {chunk_size} must be divisible by " - f"ATOM block size {self.block_size}" - ) - except Exception: - shutdown = getattr(adapter, "shutdown", None) - if callable(shutdown): - shutdown() - raise - self._adapter = adapter - self.chunk_size = chunk_size - logger.info( - "LMCache MP registered rank=%d tensors=%d groups=%d " - "bytes_per_block=%d chunk=%d tp_replication=%d writer=%s " - "save=%s load=%s", - rank, - len(views.tensors), - len(views.layer_groups), - views.bytes_per_block, - self.chunk_size, - requested_replication, - self._is_kv_writer, - self._do_save, - self._do_load, - ) - - def start_load_kv(self, metadata: Any) -> None: - if not isinstance(metadata, LMCacheOffloadMetadata): - return - if self._adapter is None or self.chunk_size is None: - raise RuntimeError("lmcache_mp KV caches are not registered") - - requests = [ - req - for req in metadata.requests - if (req.load_spec is not None and self._do_load) - or (req.save_spec is not None and self._do_save) - ] - if not requests: - return - event = torch.cuda.Event(interprocess=True) - event.record(torch.cuda.current_stream()) - for req in requests: - if req.load_spec is not None and self._do_load: - self._submit_load(req, event) - if req.save_spec is not None and self._do_save: - self._submit_save(req, event) - - def _block_slice(self, req: LMCacheReqMeta, start: int, end: int) -> list[int]: - if start < 0 or end < start: - raise ValueError(f"invalid LMCache MP token range [{start}, {end})") - if start % self.block_size or end % self.block_size: - raise ValueError( - f"LMCache MP token range [{start}, {end}) must align to " - f"block size {self.block_size}" - ) - block_ids = list( - req.block_ids[start // self.block_size : end // self.block_size] - ) - expected = (end - start) // self.block_size - if len(block_ids) != expected: - raise ValueError( - f"LMCache MP request {req.req_id} needs {expected} blocks for " - f"[{start}, {end}), got {len(block_ids)}" - ) - return block_ids - - def _submit_load(self, req: LMCacheReqMeta, event: Any) -> None: - from lmcache.integration.atom import AtomMPTransferSpec - - assert req.load_spec is not None - completion = req.load_operation or req.req_id - request_id = _mp_session_id(self._config, req.req_id) - operation_id = _transfer_operation_id("load", completion) - with self._lock: - if ( - operation_id in self._completed_load_operations - or operation_id in self._pending_loads - or operation_id in self._submitting_loads - ): - raise RuntimeError( - f"duplicate LMCache MP load operation {operation_id!r}" - ) - self._submitting_loads.add(operation_id) - start = int(req.load_spec.hbm_cached_tokens) - end = ( - int(req.load_spec.lmcache_cached_tokens) - if req.load_spec.transfer_end_tokens is None - else int(req.load_spec.transfer_end_tokens) - ) - try: - if start % self.chunk_size or end % self.chunk_size: - raise ValueError( - f"load range [{start}, {end}) is not LMCache chunk aligned " - f"({self.chunk_size})" - ) - block_ids = self._block_slice(req, start, end) - op = AtomMPTransferSpec( - token_ids=list(req.token_ids), - block_ids=[block_ids], - start=start, - end=end, - ) - except Exception: - # Nothing was sent: the failure is provable and terminal. - logger.exception( - "Invalid LMCache MP load descriptor for %s", - req.req_id, - ) - with self._lock: - self._submitting_loads.discard(operation_id) - _remember_operation_tombstone( - operation_id, - self._completed_load_operations, - self._completed_load_operation_order, - ) - self._immediate_load_failures.add(completion) - return - try: - transfer = self._adapter.submit_retrieve_request( - request_id, - op, - event, - ) - except Exception: - # The server may have taken the request before the connection - # raised, and may still write the destination blocks. - logger.exception( - "LMCache MP load submission unprovable for %s", - req.req_id, - ) - transfer = _UnprovableSubmission() - with self._lock: - self._submitting_loads.discard(operation_id) - self._pending_loads[operation_id] = _PendingLoad( - completion=completion, - future=transfer, - ) - - def _submit_save(self, req: LMCacheReqMeta, event: Any) -> None: - from lmcache.integration.atom import AtomMPTransferSpec - - assert req.save_spec is not None - completion = req.save_operation or req.req_id - request_id = _mp_session_id(self._config, req.req_id) - operation_id = _transfer_operation_id("save", completion) - with self._lock: - if ( - operation_id in self._completed_save_operations - or operation_id in self._pending_saves - or operation_id in self._submitting_saves - ): - raise RuntimeError( - f"duplicate LMCache MP save operation {operation_id!r}" - ) - self._submitting_saves.add(operation_id) - end = (len(req.token_ids) // self.chunk_size) * self.chunk_size - start = ( - int(req.save_spec.skip_leading_tokens) // self.chunk_size - ) * self.chunk_size - if not self._is_kv_writer: - with self._lock: - self._submitting_saves.discard(operation_id) - _remember_operation_tombstone( - operation_id, - self._completed_save_operations, - self._completed_save_operation_order, - ) - self._immediate_saves[completion] = ( - (start, end) if start < end else None - ) - return - if start >= end: - with self._lock: - self._submitting_saves.discard(operation_id) - _remember_operation_tombstone( - operation_id, - self._completed_save_operations, - self._completed_save_operation_order, - ) - self._immediate_saves[completion] = None - return - try: - block_ids = self._block_slice(req, start, end) - op = AtomMPTransferSpec( - token_ids=list(req.token_ids), - block_ids=[block_ids], - start=start, - end=end, - ) - except Exception: - # Nothing was sent: report a terminal failure so the save settles. - logger.exception( - "Invalid LMCache MP save descriptor for %s", - req.req_id, - ) - with self._lock: - self._submitting_saves.discard(operation_id) - self._pending_saves[operation_id] = _PendingSave( - completion=completion, - future=None, - start=start, - end=end, - ) - return - try: - submit = getattr( - self._adapter, - "submit_store_request_with_chunk_events", - self._adapter.submit_store_request, - ) - transfer = submit( - request_id, - op, - event, - ) - except Exception: - # The server may have received the request before the connection - # raised: keep the source leased until a terminal report, which - # the transfer deadline turns into a fail-stop if it never comes. - logger.exception( - "LMCache MP save submission unprovable for %s", - req.req_id, - ) - transfer = _UnprovableSubmission() - with self._lock: - self._submitting_saves.discard(operation_id) - self._pending_saves[operation_id] = _PendingSave( - completion=completion, - future=transfer, - start=start, - end=end, - ) - - def get_finished(self) -> KVConnectorOutput: - if self._adapter is None: - return KVConnectorOutput() - done_load: set[LoadCompletionId] = set() - failed_load: set[LoadCompletionId] = set() - done_save: set[SaveCompletionId] = set() - connector_completions: set[ConnectorCompletion] = set() - with self._lock: - # Heartbeat health is a control-plane signal, not proof that GPU - # work submitted before the failure has quiesced. Keep every real - # future until its device event is terminal. A pre-submit drop is - # represented by None, while LMCache's missing-registration path - # returns an event-free terminal False future. - for operation_id, pending in list(self._pending_saves.items()): - take_ranges = getattr(pending.future, "take_completed_ranges", None) - if callable(take_ranges) and isinstance( - pending.completion, SaveOperationId - ): - try: - connector_completions |= _source_safe_completions( - pending.completion, take_ranges() - ) - except Exception: - logger.warning( - "LMCache MP source-safe event polling failed", - exc_info=True, - ) - terminal, result = _terminal_future_result(pending.future) - if not terminal: - _enforce_transfer_deadline( - operation_id, pending.started_at, self._transfer_deadline_s - ) - if terminal: - self._pending_saves.pop(operation_id, None) - _remember_operation_tombstone( - operation_id, - self._completed_save_operations, - self._completed_save_operation_order, - ) - done_save.add(pending.completion) - if isinstance(pending.completion, SaveOperationId): - connector_completions |= _source_safe_completions( - pending.completion, - _chunk_ranges( - pending.start, pending.end, int(self.chunk_size) - ), - ) - connector_completions.add( - ConnectorCompletion( - DENSE_PAGE_STORE_CHANNEL, - pending.completion, - result is True, - ) - ) - - for operation_id, pending in list(self._pending_loads.items()): - terminal, result = _terminal_future_result(pending.future) - if not terminal: - _enforce_transfer_deadline( - operation_id, pending.started_at, self._transfer_deadline_s - ) - continue - self._pending_loads.pop(operation_id, None) - _remember_operation_tombstone( - operation_id, - self._completed_load_operations, - self._completed_load_operation_order, - ) - if result is not True: - failed_load.add(pending.completion) - else: - done_load.add(pending.completion) - done_save.update(self._immediate_saves) - for completion, token_range in self._immediate_saves.items(): - if isinstance(completion, SaveOperationId): - if token_range is not None: - connector_completions |= _source_safe_completions( - completion, - _chunk_ranges(*token_range, int(self.chunk_size)), - ) - connector_completions.add( - ConnectorCompletion(DENSE_PAGE_STORE_CHANNEL, completion, True) - ) - failed_load.update(self._immediate_load_failures) - self._immediate_saves.clear() - self._immediate_load_failures.clear() - return KVConnectorOutput( - finished_loading=done_load, - failed_loading=failed_load, - finished_saving=done_save, - connector_completions=connector_completions, - ) - - -class LMCacheMPConnectorScheduler(ChunkedOffloadSchedulerBase): - """Scheduler-side LMCache MP connector for generic PAGE offload.""" - - _supports_early_block_release = True - - def __init__(self, config: Any, *, checkpoint_spec: Any = None) -> None: - _validate_mp_config(config) - kvc = getattr(config, "kv_transfer_config", {}) or {} - validated_kv_role(kvc) - offcfg._strict_integer( - "LMCache MP block size", - config.kv_cache_block_size, - minimum=1, - ) - adapter = _make_scheduler_adapter(config, checkpoint_spec=checkpoint_spec) - try: - extra = _extra_config(config) - timeout = float(extra.get("lmcache.mp.lookup_timeout", 30.0)) - poll_interval = float(extra.get("lmcache.mp.lookup_poll_interval", 0.01)) - lookup_client = _MPLookupClient( - adapter, - config=config, - timeout=timeout, - poll_interval=poll_interval, - ) - self._mp_adapter = adapter - self._scheduler_deadline_s = ( - _transfer_deadline_s(config) + _SCHEDULER_DEADLINE_MARGIN_S - ) - self._transfer_seen_at: dict[Any, float] = {} - super().__init__( - config, - chunk_size=int(adapter.lmcache_tokens_per_chunk), - lookup_client=lookup_client, - ) - except Exception: - shutdown = getattr(adapter, "shutdown", None) - if callable(shutdown): - shutdown() - raise - - def get_num_new_matched_tokens(self, seq: Any) -> tuple[int, bool]: - matched = super().get_num_new_matched_tokens(seq) - sid = str(seq.id) - num_prompt = int(seq.num_prompt_tokens) - # The base's remembered hit, not `hit_tokens`: a step answered from the - # memo runs no lookup, so there is no client-side state to read back. - hit = self._last_tier_hit(seq, sid) - if hit != num_prompt or num_prompt % self.chunk_size: - return matched - - load_spec = self._load_specs.get(sid) - if load_spec is None: - return matched - if self.block_size == 1: - # Allocating prompt-1 one-token blocks cannot provide a destination - # for the complete final LMCache chunk. Recompute instead. - self._clear_pending_load(sid) - return 0, False - - load_spec.transfer_end_tokens = num_prompt - self._hit_save_floors[sid] = num_prompt - return matched - - def build_connector_meta(self) -> LMCacheOffloadMetadata: - metadata = super().build_connector_meta() - for req in metadata.requests: - if req.load_spec is not None: - end = req.load_spec.transfer_end_tokens - if end is None: - end = req.load_spec.lmcache_cached_tokens - self._lookup_client.prepare_retrieve( - str(req.req_id), - int(req.load_spec.hbm_cached_tokens), - int(end), - ) - # Also catches retirements that bypass _finish_retired_request, such as - # a late save with nothing left to store. - self._end_finished_sessions() - return metadata - - def load_finished(self, req_id: Any) -> bool: - finished = super().load_finished(req_id) - if finished: - raw_id = req_id.req_id if hasattr(req_id, "req_id") else req_id - try: - self._lookup_client.complete_retrieve( - str(raw_id), - succeeded=True, - ) - except Exception: - logger.warning( - "LMCache MP successful-load cleanup failed for request %s", - raw_id, - exc_info=True, - ) - return finished - - def load_failed(self, req_id: Any) -> bool: - raw_id = req_id.req_id if hasattr(req_id, "req_id") else req_id - active = self._active_load_operations.get(str(raw_id)) - if isinstance(req_id, LoadOperationId): - is_current = active is not None and active[1] == req_id - else: - is_current = active is None - if is_current: - try: - self._lookup_client.complete_retrieve( - str(raw_id), - succeeded=False, - ) - except Exception: - logger.warning( - "LMCache MP failed-load cleanup failed for request %s", - raw_id, - exc_info=True, - ) - return super().load_failed(req_id) - - def abandon_save(self, req_id: Any) -> None: - # The engine's clock-based reclaim cannot prove the MP server stopped - # reading the source, so an MP save keeps its lease until a terminal - # report. `_enforce_transfer_deadlines` bounds a report that never comes. - logger.warning( - "LMCache MP keeps the source of request %s leased until its save " - "reports; a transfer that never reports stops the engine after %.0fs", - req_id, - self._scheduler_deadline_s, - ) - - def reclaim_stale_leases(self, timeout_s: float) -> list[frozenset]: - del timeout_s # See abandon_save: no release without a report. - return [] - - def _live_transfers(self) -> set[Any]: - """Dispatched operations still waiting for a terminal worker report.""" - live: set[Any] = set(self._save_inflight.values()) - live.update(self._save_operation_owner) - live.update(operation for _, operation in self._active_load_operations.values()) - return live - - def _enforce_transfer_deadlines(self) -> None: - """Fail-stop when a dispatched transfer never reports. - - The worker bounds what it submitted; this catches what it cannot see, - such as a lost completion or a TP rank that never reports. - """ - live = self._live_transfers() - seen = self._transfer_seen_at - for operation in [op for op in seen if op not in live]: - del seen[operation] - now = time.monotonic() - for operation in live: - _enforce_transfer_deadline( - operation, seen.setdefault(operation, now), self._scheduler_deadline_s - ) - - def process_completions(self, output: KVConnectorOutput) -> KVConnectorOutput: - output = super().process_completions(output) - self._enforce_transfer_deadlines() - return output - - def request_finished(self, seq: Any) -> None: - super().request_finished(seq) - self._pending_session_ends()[str(seq.id)] = seq - self._end_finished_sessions() - - def _finish_retired_request(self, sid: str) -> None: - super()._finish_retired_request(sid) - self._end_finished_sessions() - - def _pending_session_ends(self) -> dict[str, Any]: - """Finished requests whose MP session must outlive their last save. - - A late save is emitted, and submitted under the request's session, - after request_finished, so the session ends only once none can follow. - """ - return self.__dict__.setdefault("_sessions_to_end", {}) - - def _end_finished_sessions(self) -> None: - """End each finished request's session once no save of it can follow.""" - pending = self._pending_session_ends() - for sid, seq in list(pending.items()): - entry = self._save_tracker.get(sid) - if (entry is not None and entry[0] is seq) or sid in self._save_inflight: - continue - del pending[sid] - try: - self._mp_adapter.end_session(_mp_session_id(self._config, seq.id)) - except Exception: - logger.warning( - "LMCache MP end_session failed for request %s", - seq.id, - exc_info=True, - ) - - -__all__ = [ - "LMCacheMPConnector", - "LMCacheMPConnectorScheduler", - "LMCacheTransferUnprovable", -] diff --git a/atom/kv_transfer/offload/mp/connector.py b/atom/kv_transfer/offload/mp/connector.py index 0588886528..d918856448 100644 --- a/atom/kv_transfer/offload/mp/connector.py +++ b/atom/kv_transfer/offload/mp/connector.py @@ -20,7 +20,9 @@ ) from atom.kv_transfer.disaggregation.types import KVConnectorOutput from atom.kv_transfer.offload._offload_common import validated_kv_role -from atom.kv_transfer.offload.mp import backend +from atom.kv_transfer.offload.mp import deployment +from atom.kv_transfer.offload.mp import scheduler as mp_scheduler +from atom.kv_transfer.offload.mp import worker as mp_worker def _publishes_native_state(transfer_tensors: Any) -> bool: @@ -59,7 +61,7 @@ class LMCacheMPConnector(KVConnectorBase): is_producer = False def __init__(self, config: Any) -> None: - backend._validate_mp_config(config) + deployment._validate_mp_config(config) validated_kv_role(getattr(config, "kv_transfer_config", {}) or {}) self._config = config self._impl: KVConnectorBase | None = None @@ -92,7 +94,7 @@ def register_kv_caches( "execute_paged_state_copies for the native-state transfer" ) else: - impl = backend.LMCacheMPConnector(self._config) + impl = mp_worker.LMCacheMPConnector(self._config) impl.register_kv_caches(kv_caches, transfer_tensors, num_blocks) self._impl = impl @@ -123,7 +125,7 @@ class LMCacheMPConnectorScheduler(KVConnectorSchedulerBase): is_offload = True def __init__(self, config: Any) -> None: - backend._validate_mp_config(config) + deployment._validate_mp_config(config) validated_kv_role(getattr(config, "kv_transfer_config", {}) or {}) self._config = config self._impl: KVConnectorSchedulerBase | None = None @@ -153,7 +155,7 @@ def bind_block_manager(self, block_manager: Any) -> None: "PAGE-only offload would restore KV without its state" ) if coordinator is None: - impl: KVConnectorSchedulerBase = backend.LMCacheMPConnectorScheduler( + impl: KVConnectorSchedulerBase = mp_scheduler.LMCacheMPConnectorScheduler( self._config ) else: diff --git a/atom/kv_transfer/offload/mp/deployment.py b/atom/kv_transfer/offload/mp/deployment.py new file mode 100644 index 0000000000..c9d8661872 --- /dev/null +++ b/atom/kv_transfer/offload/mp/deployment.py @@ -0,0 +1,342 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""How this engine meets LMCache's standalone multiprocess server. + +Configuration validation, TP/DP topology and rank collapse, the model +namespace, and the adapters both connector halves open to the server. +""" + +from __future__ import annotations + +import hashlib +import json +import logging +from dataclasses import replace +from typing import Any + +from atom.kv_transfer.offload import config as offcfg +from atom.utils import envs + +logger = logging.getLogger("atom") + +_MP_LAYOUT_VERSION = 3 + + +def _extra_config(config: Any) -> dict[str, Any]: + kvc = getattr(config, "kv_transfer_config", {}) or {} + extra = kvc.get("kv_connector_extra_config", {}) or {} + if not isinstance(extra, dict): + raise TypeError("kv_connector_extra_config must be a dictionary") + return extra + + +def _storage_kv_transfer_config(config: Any) -> dict[str, Any]: + """Remove MP transport-only options before LMCache storage parsing.""" + + kvc = dict(getattr(config, "kv_transfer_config", {}) or {}) + extra = kvc.get("kv_connector_extra_config") + if isinstance(extra, dict): + kvc["kv_connector_extra_config"] = { + key: value + for key, value in extra.items() + if not (isinstance(key, str) and key.startswith("lmcache.mp.")) + } + return kvc + + +def _mp_session_id(config: Any, request_id: Any) -> str: + """Scope one LMCache MP request session to its global DP replica.""" + + return f"{offcfg.lmcache_engine_id(config)}:{request_id}" + + +def _transfer_mode(config: Any) -> str: + extra = _extra_config(config) + configured_mode = extra.get("lmcache.mp.mp_transfer_mode") + if configured_mode is None: + configured_mode = envs.LMCACHE_MP_TRANSFER_MODE + transfer_mode = str(configured_mode).strip().lower() + if transfer_mode not in ("auto", "lmcache_driven", "engine_driven"): + raise ValueError( + "LMCache MP transfer mode must be 'auto', 'lmcache_driven', or " + f"'engine_driven', got {configured_mode!r}" + ) + if transfer_mode == "engine_driven": + raise NotImplementedError( + "ATOM lmcache_mp requires LMCache's lmcache_driven transfer path " + "because engine_driven does not support multiple physical " + "cache groups" + ) + return transfer_mode + + +def _validate_mp_config(config: Any) -> tuple[int, int]: + """Validate topology constraints shared by generic PAGE layouts.""" + tp_size = offcfg._strict_integer( + "tensor_parallel_size", + getattr(config, "tensor_parallel_size", 1) or 1, + minimum=1, + ) + pp_size = offcfg._strict_integer( + "pipeline_parallel_size", + getattr(config, "pipeline_parallel_size", 1) or 1, + minimum=1, + ) + dcp_size = offcfg._strict_integer( + "decode_context_parallel_size", + getattr(config, "decode_context_parallel_size", 1) or 1, + minimum=1, + ) + pcp_size = offcfg._strict_integer( + "prefill_context_parallel_size", + getattr(config, "prefill_context_parallel_size", 1) or 1, + minimum=1, + ) + parallel_config = getattr(config, "parallel_config", None) + dp_size = offcfg._strict_integer( + "data_parallel_size", + getattr( + parallel_config, + "data_parallel_size", + getattr(config, "data_parallel_size", 1), + ) + or 1, + minimum=1, + ) + dp_size_local = offcfg._strict_integer( + "data_parallel_size_local", + getattr(parallel_config, "data_parallel_size_local", dp_size) or dp_size, + minimum=1, + ) + if pp_size != 1: + raise NotImplementedError("lmcache_mp does not support PP yet") + if dcp_size != 1: + raise NotImplementedError("lmcache_mp does not support DCP yet") + if pcp_size != 1: + raise NotImplementedError("lmcache_mp does not support PCP yet") + # Single-host DP replicas deliberately share one (model_name, worker_id, + # world_size) identity: it is the content-addressed storage namespace, so + # replicas deduplicate identical prefixes. It is not a registration key. + # The server registers GPU memory per unique instance_id, refcounts layout + # descriptors per (model_name, world_size) (all replicas publish the same + # one), and _mp_session_id scopes request sessions and their locks per + # replica. + if dp_size_local != dp_size: + raise NotImplementedError( + "lmcache_mp supports DP and DP-attention only within one host; " + "multi-node DP requires one LMCache server per host and local " + "server routing " + f"(data_parallel_size={dp_size}, local={dp_size_local})" + ) + + _transfer_mode(config) + return tp_size, pp_size + + +def _config_has_fully_replicated_tp_pages(config: Any) -> bool: + """Conservatively identify configs whose complete PAGE cache is TP-replicated. + + MLA caches the shared latent before the TP-sharded KV-B projection. ATOM's + sparse MLA index-key projection is replicated too, so its auxiliary PAGE + plane has the same property. + + The worker validates this config-time prediction against the attention + backend's ``KVTransferTensors.tp_replication_factor`` declaration before + registering any cache memory. + """ + + if _config_has_own_pool_draft(config): + # A DSpark draft whose backend owns a KV pool appends per-rank-sharded + # PAGE regions (`draft_kv.py` declares factor 1), so the complete PAGE + # object is not replicated even when the target's is. + return False + hf_config = getattr(config, "hf_config", None) + # MiniMax-M3 is GQA. Some TP ranks can happen to own the same KV head when + # TP exceeds the global KV-head count, but the complete PAGE object is not + # replicated across the whole TP group and must remain one shard per rank. + if offcfg._is_minimax_m3(hf_config): + return False + hf_config = getattr(hf_config, "text_config", hf_config) + return getattr(hf_config, "kv_lora_rank", None) is not None + + +def _config_has_own_pool_draft(config: Any) -> bool: + """Whether a speculative draft caches into a KV pool of its own. + + Mirrors `spec_decode.draft_kv.draft_kv_builder`, which only DSpark calls: + the draft's own backend answers through `DRAFT_OWNS_KV_POOL`. + """ + + speculative = getattr(config, "speculative_config", None) + draft_hf = getattr(speculative, "draft_model_hf_config", None) + if getattr(speculative, "method", None) != "dspark" or draft_hf is None: + return False + from atom.utils.selector import attn_family, get_attn_backend + + return bool(get_attn_backend(attn_family(draft_hf)).DRAFT_OWNS_KV_POOL) + + +def _tp_replication_factor(config: Any) -> int: + """Return the PAGE rank-collapse factor selected before workers start. + + ``auto`` uses only structural cache information available in the shared + engine config. An explicit boolean is useful for a new attention backend: + ``True`` requests full TP collapse, but worker registration still fails + closed unless that backend declares every PAGE region byte-identical. + """ + + tp_size, _ = _validate_mp_config(config) + configured = _extra_config(config).get("lmcache.mp.tp_rank_collapse", "auto") + if isinstance(configured, str) and configured.strip().lower() == "auto": + collapse = _config_has_fully_replicated_tp_pages(config) + elif type(configured) is bool: + collapse = configured + else: + raise TypeError("lmcache.mp.tp_rank_collapse must be true, false, or 'auto'") + return tp_size if collapse else 1 + + +def _published_tp_replication_factor( + transfer_tensors: Any, + *, + tp_size: int, + native_state: bool = False, +) -> int: + """Validate a backend's whole-object TP replication declaration.""" + + attribute = ( + "native_state_tp_replication_factor" + if native_state + else "tp_replication_factor" + ) + label = "native STATE" if native_state else "KV PAGE" + factor = offcfg._strict_integer( + f"{label} TP replication factor", + getattr(transfer_tensors, attribute, 1), + minimum=1, + ) + if tp_size % factor: + raise ValueError( + f"{label} TP replication factor {factor} must divide TP size {tp_size}" + ) + if factor not in (1, tp_size): + raise NotImplementedError( + "lmcache_mp currently supports only sharded or fully TP-replicated " + f"{label} layouts, got replication factor {factor} for TP size {tp_size}" + ) + return factor + + +def _server_urls(config: Any) -> list[str]: + extra = _extra_config(config) + configured = extra.get("lmcache.mp.server_urls") + if configured is not None: + if isinstance(configured, (list, tuple)): + urls = [str(value).strip() for value in configured if str(value).strip()] + else: + urls = [ + value.strip() for value in str(configured).split(",") if value.strip() + ] + else: + host = str(extra.get("lmcache.mp.host", "tcp://localhost")).strip() + if not host: + raise ValueError("lmcache.mp.host must be non-empty") + port = offcfg._strict_integer( + "lmcache.mp.port", + extra.get("lmcache.mp.port", 5555), + minimum=1, + ) + if not 1 <= port <= 65535: + raise ValueError("lmcache.mp.port must be in [1, 65535]") + urls = [f"{host}:{port}"] + urls = [url if "://" in url else f"tcp://{url}" for url in urls] + if len(urls) != 1: + raise NotImplementedError( + "lmcache_mp currently supports exactly one LMCache server" + ) + return urls + + +def _model_namespace(config: Any, *, checkpoint_spec: Any = None) -> str: + """Build a model/layout namespace shared by scheduler and workers.""" + + cfg = offcfg.build_lmcache_config(_storage_kv_transfer_config(config)) + world_size = offcfg.lmcache_replica_world_size(config) + page_namespace = offcfg.build_page_namespace( + config, + cfg, + world_size, + ) + namespace = f"{page_namespace}::lmcache-mp-v{_MP_LAYOUT_VERSION}" + if checkpoint_spec is not None: + hf = getattr(config, "hf_config", None) + hf = getattr(hf, "text_config", hf) + document = { + "checkpoint": checkpoint_spec.to_wire(), + "hf_commit": getattr(hf, "_commit_hash", None), + "revision": getattr(config, "revision", None), + "model_revision": _extra_config(config).get("lmcache.mp.model_revision"), + } + fingerprint = hashlib.sha256( + json.dumps(document, sort_keys=True).encode() + ).hexdigest()[:32] + namespace += f"::native-state-v1-{fingerprint}" + return namespace + + +def _parallel_strategy(config: Any, worker_id: int) -> Any: + from lmcache.integration.atom import AtomMPParallelConfig + + tp_size, pp_size = _validate_mp_config(config) + if worker_id < 0 or worker_id >= tp_size: + raise ValueError( + f"LMCache MP worker rank {worker_id} is outside [0, {tp_size})" + ) + replication_factor = _tp_replication_factor(config) + return AtomMPParallelConfig( + world_size=tp_size * pp_size // replication_factor, + worker_id=worker_id // replication_factor, + tp_size=tp_size, + ) + + +def _make_scheduler_adapter(config: Any, *, checkpoint_spec: Any = None) -> Any: + import zmq + from lmcache.integration.atom import AtomMPSchedulerAdapter + + num_kv_readers = _tp_replication_factor(config) + + class _ReaderAwareSchedulerAdapter(AtomMPSchedulerAdapter): + """Reserve one LMCache read lock for every collapsed TP consumer.""" + + def _create_key(self, *args: Any, **kwargs: Any) -> Any: + key = super()._create_key(*args, **kwargs) + return replace(key, num_kv_readers=num_kv_readers) + + extra = _extra_config(config) + return _ReaderAwareSchedulerAdapter( + server_url=_server_urls(config)[0], + context=zmq.Context.instance(), + model_name=_model_namespace(config, checkpoint_spec=checkpoint_spec), + block_size=int(config.kv_cache_block_size), + parallel_config=_parallel_strategy(config, 0), + mq_timeout=float(extra.get("lmcache.mp.mq_timeout", 300.0)), + ) + + +def _make_worker_adapter(config: Any, rank: int, *, checkpoint_spec: Any = None) -> Any: + import zmq + from lmcache.integration.atom import AtomMPWorkerAdapter + + extra = _extra_config(config) + return AtomMPWorkerAdapter( + server_url=_server_urls(config)[0], + context=zmq.Context.instance(), + model_name=_model_namespace(config, checkpoint_spec=checkpoint_spec), + block_size=int(config.kv_cache_block_size), + parallel_config=_parallel_strategy(config, rank), + mq_timeout=float(extra.get("lmcache.mp.mq_timeout", 300.0)), + heartbeat_interval=float(extra.get("lmcache.mp.heartbeat_interval", 10.0)), + transfer_mode=_transfer_mode(config), + ) diff --git a/atom/kv_transfer/offload/mp/lookup.py b/atom/kv_transfer/offload/mp/lookup.py new file mode 100644 index 0000000000..026e342253 --- /dev/null +++ b/atom/kv_transfer/offload/mp/lookup.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Scheduler-side LMCache MP lookups and their read-lock bookkeeping.""" + +from __future__ import annotations + +import logging +import time +from dataclasses import dataclass +from typing import Any + +from atom.kv_transfer.offload.mp.deployment import _mp_session_id + +logger = logging.getLogger("atom") + + +@dataclass +class _LookupState: + token_ids: list[int] + hit: int | None = None + retrieve_start: int | None = None + retrieve_end: int | None = None + + +class _MPLookupClient: + """Synchronous lookup facade used by ATOM's existing scheduler policy.""" + + token_database = None + + def __init__( + self, + adapter: Any, + *, + config: Any, + timeout: float, + poll_interval: float, + ) -> None: + if timeout <= 0 or poll_interval <= 0: + raise ValueError("LMCache MP lookup timeout and poll interval must be > 0") + self._adapter = adapter + self._config = config + self._timeout = timeout + self._poll_interval = poll_interval + self._lookups: dict[str, _LookupState] = {} + + def lookup(self, token_ids: list[int], lookup_id: str) -> int | None: + """Hit length for this prompt, or None if the tier never answered. + + None is a non-answer, not an empty answer: the caller must ask again + rather than record "this tier has nothing" for a prompt the tier may + well hold. + """ + + state = _LookupState(token_ids=list(token_ids)) + self._lookups[lookup_id] = state + request_id = _mp_session_id(self._config, lookup_id) + self._adapter.maybe_submit_lookup_request(request_id, token_ids) + deadline = time.monotonic() + self._timeout + while True: + result = self._adapter.check_lookup_result(request_id) + if result is not None: + hit = int(result) + state.hit = hit + return hit + if time.monotonic() >= deadline: + logger.warning( + "LMCache MP lookup timed out after %.1fs for request %s", + self._timeout, + lookup_id, + ) + # The MP API has no cancel-prefetch call. Keep the adapter job + # intact so request_finished() can release locks if the result + # becomes available; eagerly cleaning it here would orphan the + # server-side lookup and its locks. `state.hit` stays None, + # which is what tells clear_lookup_status() the job is still + # outstanding. + return None + time.sleep(self._poll_interval) + + def prepare_retrieve( + self, + lookup_id: str, + start: int, + end: int | None = None, + ) -> None: + """Hand one subrange of a lookup hit to the worker retrieve. + + The worker owns and releases locks in ``[start, end)``. Any hit prefix + already resident in HBM and any hit suffix beyond ``end`` will not be + consumed by that retrieve, so release those locks here. + """ + if start < 0 or (end is not None and end < start): + raise ValueError( + f"invalid retrieve range for {lookup_id}: start={start}, end={end}" + ) + state = self._lookups.get(lookup_id) + request_id = _mp_session_id(self._config, lookup_id) + if ( + state is not None + and state.hit is not None + and end is not None + and end > state.hit + ): + raise ValueError( + f"retrieve end {end} exceeds lookup hit {state.hit} for {lookup_id}" + ) + if state is not None and state.hit is not None and start > 0: + self._adapter.free_lookup_locks( + token_ids=state.token_ids, + start=0, + end=min(start, state.hit), + request_id=request_id, + ) + if ( + state is not None + and state.hit is not None + and end is not None + and end < state.hit + ): + self._adapter.free_lookup_locks( + token_ids=state.token_ids, + start=end, + end=state.hit, + request_id=request_id, + ) + if state is not None: + state.retrieve_start = start + state.retrieve_end = end + self._adapter.cleanup_lookup_result(request_id) + + def complete_retrieve(self, lookup_id: str, *, succeeded: bool) -> None: + # Once a retrieve is submitted, LMCache owns the remaining lookup + # locks. Its lmcache-driven transfer releases them on both success and + # failure (including partial failures). Releasing the range again from + # the scheduler would decrement every TP rank's read locks twice. The + # scheduler cannot distinguish a pre-submit failure; that rarer case is + # left to request end_session/server TTL cleanup instead. + self._lookups.pop(lookup_id, None) + self._adapter.cleanup_lookup_result(_mp_session_id(self._config, lookup_id)) + + def hit_tokens(self, lookup_id: str) -> int | None: + state = self._lookups.get(lookup_id) + return None if state is None else state.hit + + def clear_lookup_status(self, lookup_id: str) -> None: + state = self._lookups.pop(lookup_id, None) + request_id = _mp_session_id(self._config, lookup_id) + if state is not None and state.hit is None: + result = self._adapter.check_lookup_result(request_id) + if result is None: + logger.warning( + "LMCache MP lookup for request %s is still pending during " + "cleanup; dropping local state while server TTL/session " + "cleanup releases any eventual locks", + lookup_id, + ) + self._adapter.cleanup_lookup_result(request_id) + return + state.hit = int(result) + # Once retrieve has started, the transfer owns the remaining read + # locks. Failed terminal loads call complete_retrieve(False) first. + if ( + state is not None + and state.retrieve_start is None + and state.hit + and state.hit > 0 + ): + self._adapter.free_lookup_locks( + token_ids=state.token_ids, + start=0, + end=state.hit, + request_id=request_id, + ) + self._adapter.cleanup_lookup_result(request_id) diff --git a/atom/kv_transfer/offload/mp/native_state_scheduler.py b/atom/kv_transfer/offload/mp/native_state_scheduler.py index c908810805..a72efc4bb9 100644 --- a/atom/kv_transfer/offload/mp/native_state_scheduler.py +++ b/atom/kv_transfer/offload/mp/native_state_scheduler.py @@ -23,15 +23,12 @@ validated_kv_role, ) from atom.kv_transfer.offload.metadata import NativeStateTransfer -from atom.kv_transfer.offload.mp.backend import ( - LMCacheMPConnectorScheduler, - _extra_config, - _validate_mp_config, -) +from atom.kv_transfer.offload.mp.deployment import _extra_config, _validate_mp_config from atom.kv_transfer.offload.mp.native_state_worker import ( NATIVE_STATE_MP_STORE_CHANNEL, require_native_state_server, ) +from atom.kv_transfer.offload.mp.scheduler import LMCacheMPConnectorScheduler from atom.model_engine.page_unit_checkpoint import SuspendedCheckpointRestore from atom.utils import envs diff --git a/atom/kv_transfer/offload/mp/native_state_worker.py b/atom/kv_transfer/offload/mp/native_state_worker.py index 06667a3dcb..b97c17bdc5 100644 --- a/atom/kv_transfer/offload/mp/native_state_worker.py +++ b/atom/kv_transfer/offload/mp/native_state_worker.py @@ -19,25 +19,27 @@ from atom.kv_transfer.offload import config as offcfg from atom.kv_transfer.offload._offload_common import max_pending_saves from atom.kv_transfer.offload.metadata import LMCacheReqMeta, NativeStateTransfer -from atom.kv_transfer.offload.mp.backend import ( - LMCacheMPConnector, - _chunk_ranges, - _enforce_transfer_deadline, +from atom.kv_transfer.offload.mp.deployment import ( _make_worker_adapter, _mp_session_id, _published_tp_replication_factor, - _remember_operation_tombstone, - _source_safe_completions, _storage_kv_transfer_config, - _terminal_future_result, _tp_replication_factor, - _transfer_operation_id, - _UnprovableSubmission, _validate_mp_config, ) from atom.kv_transfer.offload.mp.native_state_layout import ( build_native_state_mp_layout, ) +from atom.kv_transfer.offload.mp.transfer import ( + _chunk_ranges, + _enforce_transfer_deadline, + _remember_operation_tombstone, + _source_safe_completions, + _terminal_future_result, + _transfer_operation_id, + _UnprovableSubmission, +) +from atom.kv_transfer.offload.mp.worker import LMCacheMPConnector from atom.model_engine.page_unit_checkpoint import CheckpointRestoreOp from atom.utils import envs diff --git a/atom/kv_transfer/offload/mp/page_views.py b/atom/kv_transfer/offload/mp/page_views.py index d921ca4a5f..4e84a4db30 100644 --- a/atom/kv_transfer/offload/mp/page_views.py +++ b/atom/kv_transfer/offload/mp/page_views.py @@ -91,4 +91,68 @@ def validate_page_views( return validated +@dataclass(frozen=True) +class _CacheViews: + """Validated block-major tensors and their copy-kernel group indices.""" + + tensors: dict[str, torch.Tensor] + layer_groups: tuple[tuple[int, ...], ...] + bytes_per_block: int + + +def _build_cache_views( + transfer_tensors: Any, + *, + num_blocks: int, +) -> _CacheViews: + """Validate backend-published PAGE views without inspecting model internals.""" + + if transfer_tensors is None: + raise ValueError("lmcache_mp requires KVTransferTensors") + + stateful_fields = { + "num_slots": getattr(transfer_tensors, "num_slots", 0), + "slot_regions": getattr(transfer_tensors, "slot_regions", None), + "swa_block_regions": getattr(transfer_tensors, "swa_block_regions", None), + "staging_region": getattr(transfer_tensors, "staging_region", None), + "gather_slot": getattr(transfer_tensors, "gather_slot", None), + "scatter_slot": getattr(transfer_tensors, "scatter_slot", None), + "expected_full_slot_region_count": getattr( + transfer_tensors, "expected_full_slot_region_count", None + ), + } + populated_stateful_fields = [ + name for name, value in stateful_fields.items() if value + ] + if populated_stateful_fields: + raise NotImplementedError( + "lmcache_mp supports PAGE-only layouts; stateful SLOT data was " + f"published through {', '.join(populated_stateful_fields)}" + ) + + tensors: dict[str, torch.Tensor] = {} + indices_by_layout: dict[tuple[torch.dtype, tuple[int, ...]], list[int]] = {} + bytes_per_block = 0 + for page in validate_page_views(transfer_tensors, num_blocks=num_blocks): + index, view, region = page.index, page.view, page.region + # LMCache receives these tensors as opaque PAGE storage, not numerical + # values. Publish a zero-copy byte view so every transfer path copies + # the exact bit pattern. This is especially important on ROCm, where + # LMCache's Python raw-pointer fallback cannot express FP8 through the + # CUDA array interface and would otherwise reconstruct the destination + # as uint8 while keeping the staging object as FP8. + byte_view = view.view(torch.uint8) + role = str(getattr(region, "semantic_role", None) or f"plane_{index}") + tensors[f"page.{index}.{role}"] = byte_view + layout = (byte_view.dtype, tuple(int(dim) for dim in byte_view.shape[1:])) + indices_by_layout.setdefault(layout, []).append(index) + bytes_per_block += page.unit_bytes + + return _CacheViews( + tensors=tensors, + layer_groups=tuple(tuple(indices) for indices in indices_by_layout.values()), + bytes_per_block=bytes_per_block, + ) + + __all__ = ["PageView", "validate_page_views"] diff --git a/atom/kv_transfer/offload/mp/scheduler.py b/atom/kv_transfer/offload/mp/scheduler.py new file mode 100644 index 0000000000..27c07c9fe6 --- /dev/null +++ b/atom/kv_transfer/offload/mp/scheduler.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Scheduler-side LMCache MP connector for generic PAGE offload.""" + +from __future__ import annotations + +import logging +import time +from typing import Any + +from atom.kv_transfer.disaggregation.types import ( + KVConnectorOutput, + LoadOperationId, +) +from atom.kv_transfer.offload import config as offcfg +from atom.kv_transfer.offload._offload_common import validated_kv_role +from atom.kv_transfer.offload.chunked_scheduler import ( + ChunkedOffloadSchedulerBase, +) +from atom.kv_transfer.offload.metadata import LMCacheOffloadMetadata +from atom.kv_transfer.offload.mp.deployment import ( + _extra_config, + _make_scheduler_adapter, + _mp_session_id, + _validate_mp_config, +) +from atom.kv_transfer.offload.mp.lookup import _MPLookupClient +from atom.kv_transfer.offload.mp.transfer import ( + _SCHEDULER_DEADLINE_MARGIN_S, + LMCacheTransferUnprovable, + _enforce_transfer_deadline, + _transfer_deadline_s, +) + +logger = logging.getLogger("atom") + + +class LMCacheMPConnectorScheduler(ChunkedOffloadSchedulerBase): + """Scheduler-side LMCache MP connector for generic PAGE offload.""" + + _supports_early_block_release = True + + def __init__(self, config: Any, *, checkpoint_spec: Any = None) -> None: + _validate_mp_config(config) + kvc = getattr(config, "kv_transfer_config", {}) or {} + validated_kv_role(kvc) + offcfg._strict_integer( + "LMCache MP block size", + config.kv_cache_block_size, + minimum=1, + ) + adapter = _make_scheduler_adapter(config, checkpoint_spec=checkpoint_spec) + try: + extra = _extra_config(config) + timeout = float(extra.get("lmcache.mp.lookup_timeout", 30.0)) + poll_interval = float(extra.get("lmcache.mp.lookup_poll_interval", 0.01)) + lookup_client = _MPLookupClient( + adapter, + config=config, + timeout=timeout, + poll_interval=poll_interval, + ) + self._mp_adapter = adapter + self._scheduler_deadline_s = ( + _transfer_deadline_s(config) + _SCHEDULER_DEADLINE_MARGIN_S + ) + self._transfer_seen_at: dict[Any, float] = {} + super().__init__( + config, + chunk_size=int(adapter.lmcache_tokens_per_chunk), + lookup_client=lookup_client, + ) + except Exception: + shutdown = getattr(adapter, "shutdown", None) + if callable(shutdown): + shutdown() + raise + + def get_num_new_matched_tokens(self, seq: Any) -> tuple[int, bool]: + matched = super().get_num_new_matched_tokens(seq) + sid = str(seq.id) + num_prompt = int(seq.num_prompt_tokens) + # The base's remembered hit, not `hit_tokens`: a step answered from the + # memo runs no lookup, so there is no client-side state to read back. + hit = self._last_tier_hit(seq, sid) + if hit != num_prompt or num_prompt % self.chunk_size: + return matched + + load_spec = self._load_specs.get(sid) + if load_spec is None: + return matched + if self.block_size == 1: + # Allocating prompt-1 one-token blocks cannot provide a destination + # for the complete final LMCache chunk. Recompute instead. + self._clear_pending_load(sid) + return 0, False + + load_spec.transfer_end_tokens = num_prompt + self._hit_save_floors[sid] = num_prompt + return matched + + def build_connector_meta(self) -> LMCacheOffloadMetadata: + metadata = super().build_connector_meta() + for req in metadata.requests: + if req.load_spec is not None: + end = req.load_spec.transfer_end_tokens + if end is None: + end = req.load_spec.lmcache_cached_tokens + self._lookup_client.prepare_retrieve( + str(req.req_id), + int(req.load_spec.hbm_cached_tokens), + int(end), + ) + # Also catches retirements that bypass _finish_retired_request, such as + # a late save with nothing left to store. + self._end_finished_sessions() + return metadata + + def load_finished(self, req_id: Any) -> bool: + finished = super().load_finished(req_id) + if finished: + raw_id = req_id.req_id if hasattr(req_id, "req_id") else req_id + try: + self._lookup_client.complete_retrieve( + str(raw_id), + succeeded=True, + ) + except Exception: + logger.warning( + "LMCache MP successful-load cleanup failed for request %s", + raw_id, + exc_info=True, + ) + return finished + + def load_failed(self, req_id: Any) -> bool: + raw_id = req_id.req_id if hasattr(req_id, "req_id") else req_id + active = self._active_load_operations.get(str(raw_id)) + if isinstance(req_id, LoadOperationId): + is_current = active is not None and active[1] == req_id + else: + is_current = active is None + if is_current: + try: + self._lookup_client.complete_retrieve( + str(raw_id), + succeeded=False, + ) + except Exception: + logger.warning( + "LMCache MP failed-load cleanup failed for request %s", + raw_id, + exc_info=True, + ) + return super().load_failed(req_id) + + def abandon_save(self, req_id: Any) -> None: + # The engine's clock-based reclaim cannot prove the MP server stopped + # reading the source, so an MP save keeps its lease until a terminal + # report. `_enforce_transfer_deadlines` bounds a report that never comes. + logger.warning( + "LMCache MP keeps the source of request %s leased until its save " + "reports; a transfer that never reports stops the engine after %.0fs", + req_id, + self._scheduler_deadline_s, + ) + + def reclaim_stale_leases(self, timeout_s: float) -> list[frozenset]: + del timeout_s # See abandon_save: no release without a report. + return [] + + def _live_transfers(self) -> set[Any]: + """Dispatched operations still waiting for a terminal worker report.""" + live: set[Any] = set(self._save_inflight.values()) + live.update(self._save_operation_owner) + live.update(operation for _, operation in self._active_load_operations.values()) + return live + + def _enforce_transfer_deadlines(self) -> None: + """Fail-stop when a dispatched transfer never reports. + + The worker bounds what it submitted; this catches what it cannot see, + such as a lost completion or a TP rank that never reports. + """ + live = self._live_transfers() + seen = self._transfer_seen_at + for operation in [op for op in seen if op not in live]: + del seen[operation] + now = time.monotonic() + for operation in live: + _enforce_transfer_deadline( + operation, seen.setdefault(operation, now), self._scheduler_deadline_s + ) + + def process_completions(self, output: KVConnectorOutput) -> KVConnectorOutput: + output = super().process_completions(output) + self._enforce_transfer_deadlines() + return output + + def request_finished(self, seq: Any) -> None: + super().request_finished(seq) + self._pending_session_ends()[str(seq.id)] = seq + self._end_finished_sessions() + + def _finish_retired_request(self, sid: str) -> None: + super()._finish_retired_request(sid) + self._end_finished_sessions() + + def _pending_session_ends(self) -> dict[str, Any]: + """Finished requests whose MP session must outlive their last save. + + A late save is emitted, and submitted under the request's session, + after request_finished, so the session ends only once none can follow. + """ + return self.__dict__.setdefault("_sessions_to_end", {}) + + def _end_finished_sessions(self) -> None: + """End each finished request's session once no save of it can follow.""" + pending = self._pending_session_ends() + for sid, seq in list(pending.items()): + entry = self._save_tracker.get(sid) + if (entry is not None and entry[0] is seq) or sid in self._save_inflight: + continue + del pending[sid] + try: + self._mp_adapter.end_session(_mp_session_id(self._config, seq.id)) + except Exception: + logger.warning( + "LMCache MP end_session failed for request %s", + seq.id, + exc_info=True, + ) + + +__all__ = ["LMCacheMPConnectorScheduler", "LMCacheTransferUnprovable"] diff --git a/atom/kv_transfer/offload/mp/transfer.py b/atom/kv_transfer/offload/mp/transfer.py new file mode 100644 index 0000000000..1d8d881a88 --- /dev/null +++ b/atom/kv_transfer/offload/mp/transfer.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Bookkeeping shared by every LMCache MP transfer: operation identity, +terminal detection, and the deadline past which an unprovable transfer +stops the engine instead of releasing memory the server may still use. +""" + +from __future__ import annotations + +import logging +import time +from collections import deque +from collections.abc import Iterable, Sequence +from dataclasses import dataclass, field +from typing import Any + +from atom.kv_transfer.disaggregation.types import ( + ConnectorCompletion, + LoadCompletionId, + SaveCompletionId, + SaveOperationId, + SaveSourceGroupId, +) +from atom.kv_transfer.offload.chunked_scheduler import ( + DENSE_PAGE_SOURCE_SAFE_CHANNEL, +) +from atom.kv_transfer.offload.mp.deployment import _extra_config + +logger = logging.getLogger("atom") + +_OPERATION_TOMBSTONE_LIMIT = 4096 + + +@dataclass +class _PendingLoad: + completion: LoadCompletionId + future: Any | None + started_at: float = field(default_factory=lambda: time.monotonic()) + + +@dataclass +class _PendingSave: + completion: SaveCompletionId + future: Any | None + start: int + end: int + started_at: float = field(default_factory=lambda: time.monotonic()) + + +_DEFAULT_TRANSFER_DEADLINE_S = 1200.0 +# The scheduler watches the same operations from dispatch, a step or more +# before a worker submits them. The margin lets the worker, which can name the +# exact transfer, fail first. +_SCHEDULER_DEADLINE_MARGIN_S = 60.0 + + +class LMCacheTransferUnprovable(RuntimeError): + """An MP transfer never reached a terminal state within the deadline.""" + + +def _transfer_deadline_s(config: Any) -> float: + """Seconds any MP transfer may stay non-terminal before the engine stops. + + Engine memory under a transfer is released only on a terminal report: a + timeout cannot prove that the server stopped its DMA, so freeing on a + clock could hand live blocks to another request. A transfer that never + reports is therefore fatal once this deadline passes -- fail-stop, rather + than corrupt memory or wedge the pool with no fault to point at. + """ + + extra = _extra_config(config) + deadline = float( + extra.get("lmcache.mp.transfer_deadline_s", _DEFAULT_TRANSFER_DEADLINE_S) + ) + if not deadline > 0: + raise ValueError("lmcache.mp.transfer_deadline_s must be > 0") + return deadline + + +def _enforce_transfer_deadline( + operation: Any, started_at: float, deadline_s: float +) -> None: + """Raise ``LMCacheTransferUnprovable`` once a transfer outlives the bound.""" + + elapsed = time.monotonic() - started_at + if elapsed < deadline_s: + return + message = ( + f"LMCache MP transfer {operation} is still not terminal after " + f"{elapsed:.0f}s. The memory it reads or writes cannot be released " + "without proof that the server stopped, so the engine is stopping. " + "Raise lmcache.mp.transfer_deadline_s if the tier is only slow." + ) + logger.error(message) + raise LMCacheTransferUnprovable(message) + + +class _UnprovableSubmission: + """Stand-in future for a submission whose transport call raised. + + The server may have taken the request before the connection failed, so + this never becomes terminal on its own: the lease is kept and only the + transfer deadline ends it. + """ + + def query(self) -> bool: + return False + + def result(self, timeout: float | None = None) -> bool: + del timeout + return False + + +def _terminal_future_result(future: Any | None) -> tuple[bool, Any]: + """Return ``(terminal, result)`` without blocking on a device future.""" + + if future is None: + return True, None + try: + if not future.query(): + return False, None + return True, future.result(timeout=0) + except Exception: + # A query/result exception is not proof that the remote GPU stream has + # quiesced. Keep the operation pending; the transfer deadline bounds a + # future that keeps raising. + logger.warning( + "LMCache MP transfer future could not prove terminal state", + exc_info=True, + ) + return False, None + + +def _transfer_operation_id(kind: str, completion: Any) -> str: + """Return a stable LMCache handle ID for one ATOM completion generation.""" + + request_id = getattr(completion, "req_id", completion) + generation = getattr(completion, "generation", None) + if generation is None: + return f"{kind}:{request_id}" + return f"{kind}:{request_id}:{generation}" + + +def _remember_operation_tombstone( + operation_id: str, + tombstones: set[str], + tombstone_order: deque[str], +) -> None: + if operation_id in tombstones: + return + tombstones.add(operation_id) + tombstone_order.append(operation_id) + while len(tombstone_order) > _OPERATION_TOMBSTONE_LIMIT: + tombstones.discard(tombstone_order.popleft()) + + +def _chunk_ranges(start: int, end: int, chunk_size: int) -> tuple[tuple[int, int], ...]: + """Split ``[start, end)`` into LMCache chunk token ranges.""" + + return tuple( + (chunk_start, min(chunk_start + chunk_size, end)) + for chunk_start in range(start, end, chunk_size) + ) + + +def _source_safe_completions( + operation: SaveOperationId, ranges: Iterable[Sequence[int]] +) -> set[ConnectorCompletion]: + """One PAGE source-safe completion per finished token range of a save.""" + + return { + ConnectorCompletion( + DENSE_PAGE_SOURCE_SAFE_CHANNEL, + SaveSourceGroupId(operation, (tuple(token_range),)), + True, + ) + for token_range in ranges + } diff --git a/atom/kv_transfer/offload/mp/worker.py b/atom/kv_transfer/offload/mp/worker.py new file mode 100644 index 0000000000..fd16615b09 --- /dev/null +++ b/atom/kv_transfer/offload/mp/worker.py @@ -0,0 +1,490 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Worker-side LMCache MP connector for backend-published PAGE views. + +Attention backends publish block-major tensor views through +``KVTransferTensors``. LMCache MP can attach several physical kernel groups to +the same engine block-id space, so this connector registers those opaque views +without knowing which model or attention implementation produced them. +""" + +from __future__ import annotations + +import logging +import threading +from collections import deque +from typing import Any + +import torch + +from atom.kv_transfer.disaggregation.base import KVConnectorBase +from atom.kv_transfer.disaggregation.types import ( + ConnectorCompletion, + KVConnectorOutput, + LoadCompletionId, + SaveCompletionId, + SaveOperationId, +) +from atom.kv_transfer.offload import config as offcfg +from atom.kv_transfer.offload._offload_common import validated_kv_role +from atom.kv_transfer.offload.chunked_scheduler import ( + DENSE_PAGE_STORE_CHANNEL, +) +from atom.kv_transfer.offload.metadata import LMCacheOffloadMetadata, LMCacheReqMeta +from atom.kv_transfer.offload.mp.deployment import ( + _make_worker_adapter, + _mp_session_id, + _published_tp_replication_factor, + _tp_replication_factor, + _validate_mp_config, +) +from atom.kv_transfer.offload.mp.page_views import _build_cache_views +from atom.kv_transfer.offload.mp.transfer import ( + _chunk_ranges, + _enforce_transfer_deadline, + _PendingLoad, + _PendingSave, + _remember_operation_tombstone, + _source_safe_completions, + _terminal_future_result, + _transfer_deadline_s, + _transfer_operation_id, + _UnprovableSubmission, +) + +logger = logging.getLogger("atom") + + +class LMCacheMPConnector(KVConnectorBase): + """Worker-side LMCache MP connector for backend-published PAGE views.""" + + is_producer = False + + def __init__(self, config: Any) -> None: + _validate_mp_config(config) + self._config = config + kvc = getattr(config, "kv_transfer_config", {}) or {} + self.kv_role = validated_kv_role(kvc) + self._do_save = self.kv_role in ("offload", "kv_both", "kv_producer") + self._do_load = self.kv_role in ("offload", "kv_both", "kv_consumer") + self.block_size = offcfg._strict_integer( + "LMCache MP block size", + config.kv_cache_block_size, + minimum=1, + ) + self.chunk_size: int | None = None + self._adapter: Any = None + self._is_kv_writer = True + self._pending_saves: dict[str, _PendingSave] = {} + self._pending_loads: dict[str, _PendingLoad] = {} + self._submitting_saves: set[str] = set() + self._submitting_loads: set[str] = set() + # Keep in-flight IDs in the unbounded pending/submitting collections. + # Only completed IDs enter these bounded replay tombstones. + self._completed_save_operations: set[str] = set() + self._completed_load_operations: set[str] = set() + self._completed_save_operation_order: deque[str] = deque() + self._completed_load_operation_order: deque[str] = deque() + # Immediate successes include collapsed-TP non-writers. Keep their + # logical token range so they can report PAGE source-safety too: the TP + # aggregator must receive the same chunk completion from every rank + # before releasing the writer's source blocks early. + self._immediate_saves: dict[SaveCompletionId, tuple[int, int] | None] = {} + self._transfer_deadline_s = _transfer_deadline_s(config) + self._immediate_load_failures: set[LoadCompletionId] = set() + self._lock = threading.Lock() + + def register_kv_caches( + self, + _kv_caches: dict[str, Any], + transfer_tensors: Any = None, + num_blocks: int | None = None, + ) -> None: + if num_blocks is None: + num_blocks = getattr(transfer_tensors, "num_blocks", None) + if num_blocks is None: + raise ValueError("lmcache_mp requires the scheduler block count") + normalized_num_blocks = offcfg._strict_integer( + "LMCache MP block count", + num_blocks, + minimum=1, + ) + + from aiter.dist.parallel_state import get_tp_group + from lmcache.v1.multiprocess.group_view import EngineGroupInfo + + tp = get_tp_group() + rank = int(tp.rank_in_group) + tp_size, _ = _validate_mp_config(self._config) + requested_replication = _tp_replication_factor(self._config) + published_replication = _published_tp_replication_factor( + transfer_tensors, + tp_size=tp_size, + ) + if requested_replication > published_replication: + raise ValueError( + "LMCache MP TP rank collapse was requested, but the attention " + "backend did not declare the complete PAGE layout fully " + f"replicated (published factor={published_replication}, " + f"TP size={tp_size})" + ) + self._is_kv_writer = rank % requested_replication == 0 + views = _build_cache_views( + transfer_tensors, + num_blocks=normalized_num_blocks, + ) + block_regions = getattr(transfer_tensors, "block_regions", None) or [] + expected = sum(int(region.unit_bytes) for region in block_regions) + if expected != views.bytes_per_block: + raise ValueError( + "lmcache_mp block geometry mismatch: " + f"views={views.bytes_per_block} transfer_regions={expected}" + ) + + adapter = _make_worker_adapter(self._config, rank) + groups = [ + EngineGroupInfo( + engine_group_id=0, + layer_indices=layer_indices, + tokens_per_block=self.block_size, + ) + for layer_indices in views.layer_groups + ] + try: + adapter.register_kv_caches(views.tensors, engine_group_infos=groups) + chunk_size = offcfg._strict_integer( + "LMCache MP chunk size", + adapter.lmcache_tokens_per_chunk, + minimum=1, + ) + if chunk_size % self.block_size: + raise ValueError( + f"LMCache MP chunk size {chunk_size} must be divisible by " + f"ATOM block size {self.block_size}" + ) + except Exception: + shutdown = getattr(adapter, "shutdown", None) + if callable(shutdown): + shutdown() + raise + self._adapter = adapter + self.chunk_size = chunk_size + logger.info( + "LMCache MP registered rank=%d tensors=%d groups=%d " + "bytes_per_block=%d chunk=%d tp_replication=%d writer=%s " + "save=%s load=%s", + rank, + len(views.tensors), + len(views.layer_groups), + views.bytes_per_block, + self.chunk_size, + requested_replication, + self._is_kv_writer, + self._do_save, + self._do_load, + ) + + def start_load_kv(self, metadata: Any) -> None: + if not isinstance(metadata, LMCacheOffloadMetadata): + return + if self._adapter is None or self.chunk_size is None: + raise RuntimeError("lmcache_mp KV caches are not registered") + + requests = [ + req + for req in metadata.requests + if (req.load_spec is not None and self._do_load) + or (req.save_spec is not None and self._do_save) + ] + if not requests: + return + event = torch.cuda.Event(interprocess=True) + event.record(torch.cuda.current_stream()) + for req in requests: + if req.load_spec is not None and self._do_load: + self._submit_load(req, event) + if req.save_spec is not None and self._do_save: + self._submit_save(req, event) + + def _block_slice(self, req: LMCacheReqMeta, start: int, end: int) -> list[int]: + if start < 0 or end < start: + raise ValueError(f"invalid LMCache MP token range [{start}, {end})") + if start % self.block_size or end % self.block_size: + raise ValueError( + f"LMCache MP token range [{start}, {end}) must align to " + f"block size {self.block_size}" + ) + block_ids = list( + req.block_ids[start // self.block_size : end // self.block_size] + ) + expected = (end - start) // self.block_size + if len(block_ids) != expected: + raise ValueError( + f"LMCache MP request {req.req_id} needs {expected} blocks for " + f"[{start}, {end}), got {len(block_ids)}" + ) + return block_ids + + def _submit_load(self, req: LMCacheReqMeta, event: Any) -> None: + from lmcache.integration.atom import AtomMPTransferSpec + + assert req.load_spec is not None + completion = req.load_operation or req.req_id + request_id = _mp_session_id(self._config, req.req_id) + operation_id = _transfer_operation_id("load", completion) + with self._lock: + if ( + operation_id in self._completed_load_operations + or operation_id in self._pending_loads + or operation_id in self._submitting_loads + ): + raise RuntimeError( + f"duplicate LMCache MP load operation {operation_id!r}" + ) + self._submitting_loads.add(operation_id) + start = int(req.load_spec.hbm_cached_tokens) + end = ( + int(req.load_spec.lmcache_cached_tokens) + if req.load_spec.transfer_end_tokens is None + else int(req.load_spec.transfer_end_tokens) + ) + try: + if start % self.chunk_size or end % self.chunk_size: + raise ValueError( + f"load range [{start}, {end}) is not LMCache chunk aligned " + f"({self.chunk_size})" + ) + block_ids = self._block_slice(req, start, end) + op = AtomMPTransferSpec( + token_ids=list(req.token_ids), + block_ids=[block_ids], + start=start, + end=end, + ) + except Exception: + # Nothing was sent: the failure is provable and terminal. + logger.exception( + "Invalid LMCache MP load descriptor for %s", + req.req_id, + ) + with self._lock: + self._submitting_loads.discard(operation_id) + _remember_operation_tombstone( + operation_id, + self._completed_load_operations, + self._completed_load_operation_order, + ) + self._immediate_load_failures.add(completion) + return + try: + transfer = self._adapter.submit_retrieve_request( + request_id, + op, + event, + ) + except Exception: + # The server may have taken the request before the connection + # raised, and may still write the destination blocks. + logger.exception( + "LMCache MP load submission unprovable for %s", + req.req_id, + ) + transfer = _UnprovableSubmission() + with self._lock: + self._submitting_loads.discard(operation_id) + self._pending_loads[operation_id] = _PendingLoad( + completion=completion, + future=transfer, + ) + + def _submit_save(self, req: LMCacheReqMeta, event: Any) -> None: + from lmcache.integration.atom import AtomMPTransferSpec + + assert req.save_spec is not None + completion = req.save_operation or req.req_id + request_id = _mp_session_id(self._config, req.req_id) + operation_id = _transfer_operation_id("save", completion) + with self._lock: + if ( + operation_id in self._completed_save_operations + or operation_id in self._pending_saves + or operation_id in self._submitting_saves + ): + raise RuntimeError( + f"duplicate LMCache MP save operation {operation_id!r}" + ) + self._submitting_saves.add(operation_id) + end = (len(req.token_ids) // self.chunk_size) * self.chunk_size + start = ( + int(req.save_spec.skip_leading_tokens) // self.chunk_size + ) * self.chunk_size + if not self._is_kv_writer: + with self._lock: + self._submitting_saves.discard(operation_id) + _remember_operation_tombstone( + operation_id, + self._completed_save_operations, + self._completed_save_operation_order, + ) + self._immediate_saves[completion] = ( + (start, end) if start < end else None + ) + return + if start >= end: + with self._lock: + self._submitting_saves.discard(operation_id) + _remember_operation_tombstone( + operation_id, + self._completed_save_operations, + self._completed_save_operation_order, + ) + self._immediate_saves[completion] = None + return + try: + block_ids = self._block_slice(req, start, end) + op = AtomMPTransferSpec( + token_ids=list(req.token_ids), + block_ids=[block_ids], + start=start, + end=end, + ) + except Exception: + # Nothing was sent: report a terminal failure so the save settles. + logger.exception( + "Invalid LMCache MP save descriptor for %s", + req.req_id, + ) + with self._lock: + self._submitting_saves.discard(operation_id) + self._pending_saves[operation_id] = _PendingSave( + completion=completion, + future=None, + start=start, + end=end, + ) + return + try: + submit = getattr( + self._adapter, + "submit_store_request_with_chunk_events", + self._adapter.submit_store_request, + ) + transfer = submit( + request_id, + op, + event, + ) + except Exception: + # The server may have received the request before the connection + # raised: keep the source leased until a terminal report, which + # the transfer deadline turns into a fail-stop if it never comes. + logger.exception( + "LMCache MP save submission unprovable for %s", + req.req_id, + ) + transfer = _UnprovableSubmission() + with self._lock: + self._submitting_saves.discard(operation_id) + self._pending_saves[operation_id] = _PendingSave( + completion=completion, + future=transfer, + start=start, + end=end, + ) + + def get_finished(self) -> KVConnectorOutput: + if self._adapter is None: + return KVConnectorOutput() + done_load: set[LoadCompletionId] = set() + failed_load: set[LoadCompletionId] = set() + done_save: set[SaveCompletionId] = set() + connector_completions: set[ConnectorCompletion] = set() + with self._lock: + # Heartbeat health is a control-plane signal, not proof that GPU + # work submitted before the failure has quiesced. Keep every real + # future until its device event is terminal. A pre-submit drop is + # represented by None, while LMCache's missing-registration path + # returns an event-free terminal False future. + for operation_id, pending in list(self._pending_saves.items()): + take_ranges = getattr(pending.future, "take_completed_ranges", None) + if callable(take_ranges) and isinstance( + pending.completion, SaveOperationId + ): + try: + connector_completions |= _source_safe_completions( + pending.completion, take_ranges() + ) + except Exception: + logger.warning( + "LMCache MP source-safe event polling failed", + exc_info=True, + ) + terminal, result = _terminal_future_result(pending.future) + if not terminal: + _enforce_transfer_deadline( + operation_id, pending.started_at, self._transfer_deadline_s + ) + if terminal: + self._pending_saves.pop(operation_id, None) + _remember_operation_tombstone( + operation_id, + self._completed_save_operations, + self._completed_save_operation_order, + ) + done_save.add(pending.completion) + if isinstance(pending.completion, SaveOperationId): + connector_completions |= _source_safe_completions( + pending.completion, + _chunk_ranges( + pending.start, pending.end, int(self.chunk_size) + ), + ) + connector_completions.add( + ConnectorCompletion( + DENSE_PAGE_STORE_CHANNEL, + pending.completion, + result is True, + ) + ) + + for operation_id, pending in list(self._pending_loads.items()): + terminal, result = _terminal_future_result(pending.future) + if not terminal: + _enforce_transfer_deadline( + operation_id, pending.started_at, self._transfer_deadline_s + ) + continue + self._pending_loads.pop(operation_id, None) + _remember_operation_tombstone( + operation_id, + self._completed_load_operations, + self._completed_load_operation_order, + ) + if result is not True: + failed_load.add(pending.completion) + else: + done_load.add(pending.completion) + done_save.update(self._immediate_saves) + for completion, token_range in self._immediate_saves.items(): + if isinstance(completion, SaveOperationId): + if token_range is not None: + connector_completions |= _source_safe_completions( + completion, + _chunk_ranges(*token_range, int(self.chunk_size)), + ) + connector_completions.add( + ConnectorCompletion(DENSE_PAGE_STORE_CHANNEL, completion, True) + ) + failed_load.update(self._immediate_load_failures) + self._immediate_saves.clear() + self._immediate_load_failures.clear() + return KVConnectorOutput( + finished_loading=done_load, + failed_loading=failed_load, + finished_saving=done_save, + connector_completions=connector_completions, + ) + + +__all__ = ["LMCacheMPConnector"] diff --git a/tests/test_lmcache_mp.py b/tests/test_lmcache_mp.py index 03d009ec22..c73a12669f 100644 --- a/tests/test_lmcache_mp.py +++ b/tests/test_lmcache_mp.py @@ -15,19 +15,28 @@ from atom.kv_transfer.disaggregation.aggregator import KVOutputAggregator from atom.kv_transfer.disaggregation.factory import KVConnectorFactory from atom.kv_transfer.disaggregation.types import ( + ConnectorCompletion, KVTransferRegion, KVTransferTensors, LoadOperationId, SaveOperationId, + SaveSourceGroupId, +) +from atom.kv_transfer.offload.chunked_scheduler import ( + DENSE_PAGE_SOURCE_SAFE_CHANNEL, + DENSE_PAGE_STORE_CHANNEL, + ChunkedOffloadSchedulerBase, ) -from atom.kv_transfer.offload.chunked_scheduler import ChunkedOffloadSchedulerBase from atom.kv_transfer.offload.metadata import ( LMCacheOffloadMetadata, LMCacheReqMeta, LoadSpec, SaveSpec, ) -from atom.kv_transfer.offload.mp import backend as mp_connector +from atom.kv_transfer.offload.mp import deployment, page_views, transfer +from atom.kv_transfer.offload.mp import lookup as mp_lookup +from atom.kv_transfer.offload.mp import scheduler as mp_scheduler +from atom.kv_transfer.offload.mp import worker as mp_worker def _config( @@ -94,11 +103,11 @@ def _config( ) def test_mp_config_rejects_unsupported_topologies(kwargs, message): with pytest.raises((NotImplementedError, ValueError), match=message): - mp_connector._validate_mp_config(_config(**kwargs)) + deployment._validate_mp_config(_config(**kwargs)) def test_mp_config_accepts_arbitrary_model_type(): - assert mp_connector._validate_mp_config(_config(model_type="ordinary_mha")) == ( + assert deployment._validate_mp_config(_config(model_type="ordinary_mha")) == ( 2, 1, ) @@ -118,18 +127,18 @@ def test_mp_config_accepts_arbitrary_model_type(): ], ) def test_mp_config_accepts_single_host_dp_and_dpa(kwargs): - assert mp_connector._validate_mp_config(_config(**kwargs))[1] == 1 + assert deployment._validate_mp_config(_config(**kwargs))[1] == 1 def test_mp_scheduler_is_not_a_dense_transport_scheduler(): from atom.kv_transfer.offload.dense.connector import DenseOffloadScheduler assert issubclass( - mp_connector.LMCacheMPConnectorScheduler, + mp_scheduler.LMCacheMPConnectorScheduler, ChunkedOffloadSchedulerBase, ) assert not issubclass( - mp_connector.LMCacheMPConnectorScheduler, + mp_scheduler.LMCacheMPConnectorScheduler, DenseOffloadScheduler, ) @@ -137,11 +146,11 @@ def test_mp_scheduler_is_not_a_dense_transport_scheduler(): def test_mp_config_rejects_engine_driven_transfer(monkeypatch): monkeypatch.setenv("LMCACHE_MP_TRANSFER_MODE", " EnGiNe_DrIvEn ") with pytest.raises(NotImplementedError, match="multiple physical"): - mp_connector._validate_mp_config(_config()) + deployment._validate_mp_config(_config()) monkeypatch.setenv("LMCACHE_MP_TRANSFER_MODE", "auto") with pytest.raises(NotImplementedError, match="multiple physical"): - mp_connector._validate_mp_config( + deployment._validate_mp_config( _config(extra={"lmcache.mp.mp_transfer_mode": " EnGiNe_DrIvEn "}) ) @@ -177,10 +186,10 @@ def __init__(self, **kwargs): monkeypatch.setitem(sys.modules, "lmcache.integration.atom", adapter_module) monkeypatch.setenv("LMCACHE_MP_TRANSFER_MODE", environment_mode) monkeypatch.setattr( - mp_connector, "_model_namespace", lambda _config, **_kwargs: "test" + deployment, "_model_namespace", lambda _config, **_kwargs: "test" ) - adapter = mp_connector._make_worker_adapter(_config(extra=extra), rank=1) + adapter = deployment._make_worker_adapter(_config(extra=extra), rank=1) assert adapter.transfer_mode == expected @@ -199,7 +208,7 @@ def __init__(self, **kwargs): ) strategies = [ - mp_connector._parallel_strategy(_config(tp=8, dp=2), rank) for rank in range(8) + deployment._parallel_strategy(_config(tp=8, dp=2), rank) for rank in range(8) ] assert {strategy.world_size for strategy in strategies} == {8} @@ -216,7 +225,7 @@ def __init__(self, **kwargs): adapter_module.AtomMPParallelConfig = AtomMPParallelConfig monkeypatch.setitem(sys.modules, "lmcache.integration.atom", adapter_module) - strategy = mp_connector._parallel_strategy( + strategy = deployment._parallel_strategy( _config( tp=1, dp=8, @@ -246,7 +255,7 @@ def __init__(self, **kwargs): ) strategies = [ - mp_connector._parallel_strategy(_config(tp=8, kv_lora_rank=512), rank) + deployment._parallel_strategy(_config(tp=8, kv_lora_rank=512), rank) for rank in range(8) ] @@ -268,18 +277,18 @@ def test_auto_rank_collapse_is_off_for_a_draft_with_its_own_pool(monkeypatch): lambda _family: SimpleNamespace(DRAFT_OWNS_KV_POOL=owns_pool["value"]), ) config = _config(tp=8, kv_lora_rank=512) - assert mp_connector._tp_replication_factor(config) == 8 + assert deployment._tp_replication_factor(config) == 8 config.speculative_config = SimpleNamespace( method="dspark", draft_model_hf_config=SimpleNamespace() ) - assert mp_connector._tp_replication_factor(config) == 1 + assert deployment._tp_replication_factor(config) == 1 owns_pool["value"] = False - assert mp_connector._tp_replication_factor(config) == 8 + assert deployment._tp_replication_factor(config) == 8 config.speculative_config.method = "mtp" owns_pool["value"] = True - assert mp_connector._tp_replication_factor(config) == 8 + assert deployment._tp_replication_factor(config) == 8 def test_auto_rank_collapse_distinguishes_glm52_mla_from_minimax_m3_gqa(): @@ -293,13 +302,13 @@ def test_auto_rank_collapse_distinguishes_glm52_mla_from_minimax_m3_gqa(): kv_lora_rank=512, ) - assert mp_connector._tp_replication_factor(glm52) == 8 - assert mp_connector._tp_replication_factor(minimax) == 1 + assert deployment._tp_replication_factor(glm52) == 8 + assert deployment._tp_replication_factor(minimax) == 1 def test_tp_rank_collapse_can_be_disabled_and_rejects_bad_values(): assert ( - mp_connector._tp_replication_factor( + deployment._tp_replication_factor( _config( tp=8, kv_lora_rank=512, @@ -309,7 +318,7 @@ def test_tp_rank_collapse_can_be_disabled_and_rejects_bad_values(): == 1 ) with pytest.raises(TypeError, match="true, false, or 'auto'"): - mp_connector._tp_replication_factor( + deployment._tp_replication_factor( _config(extra={"lmcache.mp.tp_rank_collapse": 1}) ) @@ -343,10 +352,10 @@ def _create_key(self, *_args, **_kwargs): adapter_module.AtomMPSchedulerAdapter = AtomMPSchedulerAdapter monkeypatch.setitem(sys.modules, "lmcache.integration.atom", adapter_module) monkeypatch.setattr( - mp_connector, "_model_namespace", lambda _config, **_kwargs: "test" + deployment, "_model_namespace", lambda _config, **_kwargs: "test" ) - adapter = mp_connector._make_scheduler_adapter( + adapter = mp_scheduler._make_scheduler_adapter( _config(tp=8, kv_lora_rank=kv_lora_rank) ) key = adapter._create_key([], 0, 0, "req", None) @@ -355,30 +364,30 @@ def _create_key(self, *_args, **_kwargs): def test_server_url_normalization_and_single_server_limit(): - assert mp_connector._server_urls(_config()) == ["tcp://localhost:5555"] - assert mp_connector._server_urls( + assert deployment._server_urls(_config()) == ["tcp://localhost:5555"] + assert deployment._server_urls( _config(extra={"lmcache.mp.host": "cache-host", "lmcache.mp.port": 6555}) ) == ["tcp://cache-host:6555"] - assert mp_connector._server_urls( + assert deployment._server_urls( _config(extra={"lmcache.mp.server_urls": "tcp://cache-host:6555"}) ) == ["tcp://cache-host:6555"] with pytest.raises(NotImplementedError, match="exactly one"): - mp_connector._server_urls( + deployment._server_urls( _config(extra={"lmcache.mp.server_urls": "host-a:1,host-b:2"}) ) with pytest.raises(NotImplementedError, match="exactly one"): - mp_connector._server_urls(_config(extra={"lmcache.mp.server_urls": []})) + deployment._server_urls(_config(extra={"lmcache.mp.server_urls": []})) with pytest.raises(ValueError, match=r"\[1, 65535\]"): - mp_connector._server_urls(_config(extra={"lmcache.mp.port": 70000})) + deployment._server_urls(_config(extra={"lmcache.mp.port": 70000})) def test_model_namespace_reuses_generic_page_namespace(monkeypatch): calls = [] cfg = object() - monkeypatch.setattr(mp_connector.offcfg, "build_lmcache_config", lambda _kvc: cfg) + monkeypatch.setattr(deployment.offcfg, "build_lmcache_config", lambda _kvc: cfg) monkeypatch.setattr( - mp_connector.offcfg, + deployment.offcfg, "build_page_namespace", lambda config, lmcache_cfg, world: ( calls.append((config, lmcache_cfg, world)) or "model::atom-page-v2-layout" @@ -386,7 +395,7 @@ def test_model_namespace_reuses_generic_page_namespace(monkeypatch): ) config = _config(model_type="ordinary_mha", tp=4) - assert mp_connector._model_namespace(config) == ( + assert deployment._model_namespace(config) == ( "model::atom-page-v2-layout::lmcache-mp-v3" ) assert calls == [(config, cfg, 4)] @@ -394,18 +403,18 @@ def test_model_namespace_reuses_generic_page_namespace(monkeypatch): def test_dp_replicas_share_model_namespace_but_not_request_sessions(monkeypatch): cfg = object() - monkeypatch.setattr(mp_connector.offcfg, "build_lmcache_config", lambda _kvc: cfg) + monkeypatch.setattr(deployment.offcfg, "build_lmcache_config", lambda _kvc: cfg) monkeypatch.setattr( - mp_connector.offcfg, + deployment.offcfg, "build_page_namespace", lambda _config, _lmcache_cfg, world: f"page-tp{world}", ) rank0 = _config(tp=4, dp=2, dp_rank=0) rank1 = _config(tp=4, dp=2, dp_rank=1) - assert mp_connector._model_namespace(rank0) == mp_connector._model_namespace(rank1) - assert mp_connector._mp_session_id(rank0, 0) == "atom-offload-dp0:0" - assert mp_connector._mp_session_id(rank1, 0) == "atom-offload-dp1:0" + assert deployment._model_namespace(rank0) == deployment._model_namespace(rank1) + assert deployment._mp_session_id(rank0, 0) == "atom-offload-dp0:0" + assert deployment._mp_session_id(rank1, 0) == "atom-offload-dp1:0" def test_scheduler_validates_role_before_connecting(monkeypatch): @@ -417,9 +426,9 @@ def connect(_config, **_kwargs): connected = True raise AssertionError("must not connect") - monkeypatch.setattr(mp_connector, "_make_scheduler_adapter", connect) + monkeypatch.setattr(mp_scheduler, "_make_scheduler_adapter", connect) with pytest.raises(ValueError, match="invalid kv_role"): - mp_connector.LMCacheMPConnectorScheduler(config) + mp_scheduler.LMCacheMPConnectorScheduler(config) assert connected is False @@ -435,13 +444,13 @@ def shutdown(self): adapter = Adapter() monkeypatch.setattr( - mp_connector, + mp_scheduler, "_make_scheduler_adapter", lambda _config, **_kwargs: adapter, ) with pytest.raises(ValueError, match="LMCache chunk size"): - mp_connector.LMCacheMPConnectorScheduler(_config()) + mp_scheduler.LMCacheMPConnectorScheduler(_config()) assert adapter.closed is True @@ -474,7 +483,7 @@ def _transfer_tensors(*, tp_replication_factor: int = 1) -> KVTransferTensors: def test_build_cache_views_groups_opaque_layouts(): transfer_tensors = _transfer_tensors() - views = mp_connector._build_cache_views(transfer_tensors, num_blocks=2) + views = page_views._build_cache_views(transfer_tensors, num_blocks=2) assert list(views.tensors) == [ "page.0.primary.0", @@ -514,7 +523,7 @@ def test_build_cache_views_publishes_float8_pages_as_raw_bytes(): ) transfer.set_block_count(2) - views = mp_connector._build_cache_views(transfer, num_blocks=2) + views = page_views._build_cache_views(transfer, num_blocks=2) published = views.tensors["page.0.latent"] assert published.dtype == torch.uint8 @@ -527,25 +536,25 @@ def test_build_cache_views_rejects_missing_or_bad_geometry(): missing_view = _transfer_tensors() missing_view.block_tensor_views.pop() with pytest.raises(ValueError, match="one block_tensor_view per block region"): - mp_connector._build_cache_views(missing_view, num_blocks=2) + page_views._build_cache_views(missing_view, num_blocks=2) bad_geometry = _transfer_tensors() bad_geometry.block_regions[0].unit_bytes += 1 with pytest.raises(ValueError, match="byte geometry mismatch"): - mp_connector._build_cache_views(bad_geometry, num_blocks=2) + page_views._build_cache_views(bad_geometry, num_blocks=2) noncontiguous = _transfer_tensors() noncontiguous.block_tensor_views[0] = torch.zeros(2, 32, 4).transpose(1, 2) assert not noncontiguous.block_tensor_views[0].is_contiguous() with pytest.raises(ValueError, match="non-empty and contiguous"): - mp_connector._build_cache_views(noncontiguous, num_blocks=2) + page_views._build_cache_views(noncontiguous, num_blocks=2) unsupported_rank = _transfer_tensors() unsupported_rank.block_tensor_views[0] = torch.zeros( 2, 4, 4, 8, dtype=torch.float16 ) with pytest.raises(ValueError, match="physical_slots, opaque_width"): - mp_connector._build_cache_views(unsupported_rank, num_blocks=2) + page_views._build_cache_views(unsupported_rank, num_blocks=2) @pytest.mark.parametrize( @@ -577,7 +586,7 @@ def test_build_cache_views_rejects_stateful_slot_layouts(field, value): NotImplementedError, match=rf"PAGE-only layouts.*{field}", ): - mp_connector._build_cache_views(transfer_tensors, num_blocks=2) + page_views._build_cache_views(transfer_tensors, num_blocks=2) class _LookupAdapter: @@ -607,9 +616,9 @@ def end_session(self, request_id): def test_mp_lookup_releases_only_hbm_prefix_after_retrieve_handoff(monkeypatch): - monkeypatch.setattr(mp_connector.time, "sleep", lambda _seconds: None) + monkeypatch.setattr(transfer.time, "sleep", lambda _seconds: None) adapter = _LookupAdapter([None, 8]) - client = mp_connector._MPLookupClient( + client = mp_lookup._MPLookupClient( adapter, config=_config(), timeout=10.0, @@ -630,9 +639,9 @@ def test_mp_lookup_releases_only_hbm_prefix_after_retrieve_handoff(monkeypatch): def test_mp_lookup_releases_hit_suffix_outside_retrieve_range(monkeypatch): - monkeypatch.setattr(mp_connector.time, "sleep", lambda _seconds: None) + monkeypatch.setattr(transfer.time, "sleep", lambda _seconds: None) adapter = _LookupAdapter([12]) - client = mp_connector._MPLookupClient( + client = mp_lookup._MPLookupClient( adapter, config=_config(), timeout=10.0, @@ -654,8 +663,8 @@ def test_mp_lookup_releases_hit_suffix_outside_retrieve_range(monkeypatch): ), ) ) - scheduler = mp_connector.LMCacheMPConnectorScheduler.__new__( - mp_connector.LMCacheMPConnectorScheduler + scheduler = mp_scheduler.LMCacheMPConnectorScheduler.__new__( + mp_scheduler.LMCacheMPConnectorScheduler ) scheduler._lookup_client = client monkeypatch.setattr( @@ -675,9 +684,9 @@ def test_mp_lookup_releases_hit_suffix_outside_retrieve_range(monkeypatch): def test_mp_lookup_rejects_retrieve_beyond_lookup_hit(monkeypatch): - monkeypatch.setattr(mp_connector.time, "sleep", lambda _seconds: None) + monkeypatch.setattr(transfer.time, "sleep", lambda _seconds: None) adapter = _LookupAdapter([8]) - client = mp_connector._MPLookupClient( + client = mp_lookup._MPLookupClient( adapter, config=_config(), timeout=10.0, @@ -694,10 +703,10 @@ def test_mp_lookup_rejects_retrieve_beyond_lookup_hit(monkeypatch): def test_mp_lookup_timeout_defers_cleanup_until_result(monkeypatch): ticks = iter([0.0, 0.0, 2.0]) - monkeypatch.setattr(mp_connector.time, "monotonic", lambda: next(ticks)) - monkeypatch.setattr(mp_connector.time, "sleep", lambda _seconds: None) + monkeypatch.setattr(transfer.time, "monotonic", lambda: next(ticks)) + monkeypatch.setattr(transfer.time, "sleep", lambda _seconds: None) adapter = _LookupAdapter([None, None]) - client = mp_connector._MPLookupClient( + client = mp_lookup._MPLookupClient( adapter, config=_config(), timeout=1.0, @@ -715,10 +724,10 @@ def test_mp_lookup_timeout_defers_cleanup_until_result(monkeypatch): def test_mp_lookup_pending_cleanup_drops_adapter_bookkeeping(monkeypatch): ticks = iter([0.0, 0.0, 2.0]) - monkeypatch.setattr(mp_connector.time, "monotonic", lambda: next(ticks)) - monkeypatch.setattr(mp_connector.time, "sleep", lambda _seconds: None) + monkeypatch.setattr(transfer.time, "monotonic", lambda: next(ticks)) + monkeypatch.setattr(transfer.time, "sleep", lambda _seconds: None) adapter = _LookupAdapter([None, None]) - client = mp_connector._MPLookupClient( + client = mp_lookup._MPLookupClient( adapter, config=_config(), timeout=1.0, @@ -733,16 +742,16 @@ def test_mp_lookup_pending_cleanup_drops_adapter_bookkeeping(monkeypatch): def _full_prompt_hit_scheduler(monkeypatch, chunk_size): - monkeypatch.setattr(mp_connector.time, "sleep", lambda _seconds: None) + monkeypatch.setattr(transfer.time, "sleep", lambda _seconds: None) adapter = _LookupAdapter([8]) - lookup = mp_connector._MPLookupClient( + lookup = mp_lookup._MPLookupClient( adapter, config=_config(), timeout=1.0, poll_interval=0.01, ) - scheduler = mp_connector.LMCacheMPConnectorScheduler.__new__( - mp_connector.LMCacheMPConnectorScheduler + scheduler = mp_scheduler.LMCacheMPConnectorScheduler.__new__( + mp_scheduler.LMCacheMPConnectorScheduler ) scheduler._mp_adapter = adapter ChunkedOffloadSchedulerBase.__init__( @@ -822,7 +831,7 @@ def test_mp_lookup_timeout_is_not_recorded_as_a_tier_miss(monkeypatch): """ ticks = iter([0.0, 0.0, 2.0, 0.0]) - monkeypatch.setattr(mp_connector.time, "monotonic", lambda: next(ticks)) + monkeypatch.setattr(transfer.time, "monotonic", lambda: next(ticks)) scheduler, lookup = _full_prompt_hit_scheduler(monkeypatch, chunk_size=4) adapter = lookup._adapter adapter.results.clear() @@ -865,14 +874,14 @@ def test_block_size_one_full_prompt_hit_stays_declined_when_asked_again(monkeypa def test_stale_load_failure_does_not_release_current_generation_locks(): adapter = _LookupAdapter([]) - lookup = mp_connector._MPLookupClient( + lookup = mp_lookup._MPLookupClient( adapter, config=_config(), timeout=1.0, poll_interval=0.01, ) - scheduler = mp_connector.LMCacheMPConnectorScheduler.__new__( - mp_connector.LMCacheMPConnectorScheduler + scheduler = mp_scheduler.LMCacheMPConnectorScheduler.__new__( + mp_scheduler.LMCacheMPConnectorScheduler ) scheduler._mp_adapter = adapter ChunkedOffloadSchedulerBase.__init__( @@ -885,7 +894,7 @@ def test_stale_load_failure_does_not_release_current_generation_locks(): current = LoadOperationId(req_id=7, generation=2) stale = LoadOperationId(req_id=7, generation=1) scheduler._active_load_operations["7"] = (seq, current) - lookup._lookups["7"] = mp_connector._LookupState( + lookup._lookups["7"] = mp_lookup._LookupState( token_ids=list(range(8)), hit=8, retrieve_start=0, @@ -1004,15 +1013,15 @@ def shutdown(self): self.shutdown_called = True -def _worker(adapter: _WorkerAdapter) -> mp_connector.LMCacheMPConnector: - worker = mp_connector.LMCacheMPConnector(_config()) +def _worker(adapter: _WorkerAdapter) -> mp_worker.LMCacheMPConnector: + worker = mp_worker.LMCacheMPConnector(_config()) worker._adapter = adapter worker.chunk_size = 8 return worker def _finish_load( - worker: mp_connector.LMCacheMPConnector, + worker: mp_worker.LMCacheMPConnector, operation_id: str, *, result: bool = True, @@ -1023,7 +1032,7 @@ def _finish_load( def _finish_save( - worker: mp_connector.LMCacheMPConnector, + worker: mp_worker.LMCacheMPConnector, operation_id: str, *, result: bool = True, @@ -1266,7 +1275,7 @@ def unprovable(_request_id, _op, _event): ): monkeypatch.setattr(adapter, name, unprovable) now = [1000.0] - monkeypatch.setattr(mp_connector.time, "monotonic", lambda: now[0]) + monkeypatch.setattr(transfer.time, "monotonic", lambda: now[0]) worker = _worker(adapter) request = _unprovable_request(kind) if kind == "load": @@ -1283,7 +1292,7 @@ def unprovable(_request_id, _op, _event): assert len(pending) == 1 now[0] += 2 - with pytest.raises(mp_connector.LMCacheTransferUnprovable): + with pytest.raises(transfer.LMCacheTransferUnprovable): worker.get_finished() assert len(pending) == 1 @@ -1293,7 +1302,7 @@ def test_worker_future_that_keeps_raising_is_bounded_by_the_deadline( monkeypatch, ): now = [1000.0] - monkeypatch.setattr(mp_connector.time, "monotonic", lambda: now[0]) + monkeypatch.setattr(transfer.time, "monotonic", lambda: now[0]) worker = _worker(_WorkerAdapter()) worker._submit_save(_unprovable_request("save"), object()) [pending] = worker._pending_saves.values() @@ -1301,7 +1310,7 @@ def test_worker_future_that_keeps_raising_is_bounded_by_the_deadline( assert not worker.get_finished().finished_saving now[0] += worker._transfer_deadline_s - with pytest.raises(mp_connector.LMCacheTransferUnprovable): + with pytest.raises(transfer.LMCacheTransferUnprovable): worker.get_finished() @@ -1323,21 +1332,21 @@ def test_worker_invalid_descriptor_fails_before_transport( [store] = [ completion for completion in output.connector_completions - if completion.channel == mp_connector.DENSE_PAGE_STORE_CHANNEL + if completion.channel == DENSE_PAGE_STORE_CHANNEL ] assert not store.succeeded def test_transfer_deadline_defaults_to_twenty_minutes(): - assert mp_connector._transfer_deadline_s(_config()) == 1200.0 + assert transfer._transfer_deadline_s(_config()) == 1200.0 assert ( - mp_connector._transfer_deadline_s( + transfer._transfer_deadline_s( _config(extra={"lmcache.mp.transfer_deadline_s": 30}) ) == 30.0 ) with pytest.raises(ValueError, match="transfer_deadline_s"): - mp_connector._transfer_deadline_s( + transfer._transfer_deadline_s( _config(extra={"lmcache.mp.transfer_deadline_s": 0}) ) @@ -1386,9 +1395,9 @@ def test_worker_reports_chunk_source_safe_before_store_terminal(fake_lmcache_mod assert output.finished_saving == set() assert output.connector_completions == { - mp_connector.ConnectorCompletion( - mp_connector.DENSE_PAGE_SOURCE_SAFE_CHANNEL, - mp_connector.SaveSourceGroupId(operation, ((8, 16),)), + ConnectorCompletion( + DENSE_PAGE_SOURCE_SAFE_CHANNEL, + SaveSourceGroupId(operation, ((8, 16),)), True, ) } @@ -1415,7 +1424,7 @@ def test_non_writer_completes_save_without_submitting(fake_lmcache_modules): assert { completion.operation_id.ranges for completion in output.connector_completions - if completion.channel == mp_connector.DENSE_PAGE_SOURCE_SAFE_CHANNEL + if completion.channel == DENSE_PAGE_SOURCE_SAFE_CHANNEL } == {((0, 8),)} with pytest.raises(RuntimeError, match="duplicate LMCache MP save"): worker._submit_save(request, object()) @@ -1445,7 +1454,7 @@ def test_collapsed_tp_early_release_waits_only_for_the_single_writer_dma( assert { completion.operation_id.ranges for completion in first.connector_completions - if completion.channel == mp_connector.DENSE_PAGE_SOURCE_SAFE_CHANNEL + if completion.channel == DENSE_PAGE_SOURCE_SAFE_CHANNEL } == {((0, 8),)} writer_future.completed_ranges = [(8, 16)] @@ -1453,7 +1462,7 @@ def test_collapsed_tp_early_release_waits_only_for_the_single_writer_dma( assert { completion.operation_id.ranges for completion in second.connector_completions - if completion.channel == mp_connector.DENSE_PAGE_SOURCE_SAFE_CHANNEL + if completion.channel == DENSE_PAGE_SOURCE_SAFE_CHANNEL } == {((8, 16),)} @@ -1644,7 +1653,7 @@ def submit_with_replay(*args): def test_worker_operation_tombstone_limit_is_4096(): - assert mp_connector._OPERATION_TOMBSTONE_LIMIT == 4096 + assert transfer._OPERATION_TOMBSTONE_LIMIT == 4096 @pytest.mark.parametrize("kind", ["load", "save"]) @@ -1653,7 +1662,7 @@ def test_worker_terminal_operation_tombstones_are_bounded( monkeypatch, kind, ): - monkeypatch.setattr(mp_connector, "_OPERATION_TOMBSTONE_LIMIT", 2) + monkeypatch.setattr(transfer, "_OPERATION_TOMBSTONE_LIMIT", 2) adapter = _WorkerAdapter() worker = _worker(adapter) @@ -1693,7 +1702,7 @@ def request(generation): if kind == "load" else worker._completed_save_operation_order ) - assert mp_connector._OPERATION_TOMBSTONE_LIMIT == 2 + assert transfer._OPERATION_TOMBSTONE_LIMIT == 2 assert seen == {f"{kind}:24:2", f"{kind}:24:3"} assert list(order) == [f"{kind}:24:2", f"{kind}:24:3"] @@ -1760,11 +1769,11 @@ def test_registers_multiple_layouts_as_views_of_one_engine_group( adapter = _WorkerAdapter() monkeypatch.setattr( - mp_connector, + mp_worker, "_make_worker_adapter", lambda _config, _rank: adapter, ) - worker = mp_connector.LMCacheMPConnector(_config(model_type="ordinary_mha")) + worker = mp_worker.LMCacheMPConnector(_config(model_type="ordinary_mha")) transfer_tensors = _transfer_tensors() worker.register_kv_caches( {}, @@ -1802,11 +1811,11 @@ def test_registration_collapses_only_backend_declared_tp_replicas( adapter = _WorkerAdapter() monkeypatch.setattr( - mp_connector, + mp_worker, "_make_worker_adapter", lambda _config, _rank: adapter, ) - worker = mp_connector.LMCacheMPConnector(_config(tp=8, kv_lora_rank=512)) + worker = mp_worker.LMCacheMPConnector(_config(tp=8, kv_lora_rank=512)) worker.register_kv_caches( {}, transfer_tensors=_transfer_tensors(tp_replication_factor=8), @@ -1831,7 +1840,7 @@ def test_registration_rejects_rank_collapse_without_backend_declaration( monkeypatch.setitem(sys.modules, "aiter.dist", dist) monkeypatch.setitem(sys.modules, "aiter.dist.parallel_state", parallel_state) - worker = mp_connector.LMCacheMPConnector( + worker = mp_worker.LMCacheMPConnector( _config( tp=8, extra={"lmcache.mp.tp_rank_collapse": True}, diff --git a/tests/test_lmcache_mp_builder_views.py b/tests/test_lmcache_mp_builder_views.py index 8c9297e7ea..1fa46f3aa1 100644 --- a/tests/test_lmcache_mp_builder_views.py +++ b/tests/test_lmcache_mp_builder_views.py @@ -16,7 +16,7 @@ import pytest import torch -from atom.kv_transfer.offload.mp.backend import _build_cache_views +from atom.kv_transfer.offload.mp.page_views import _build_cache_views from atom.model_ops.attentions.mha_kv_pool import MhaKvPool from atom.model_ops.attentions.mla_kv_pool import MlaKvPool from atom.model_ops.attentions.pool_layout.entry_arena import EntryField diff --git a/tests/test_lmcache_mp_native_layout.py b/tests/test_lmcache_mp_native_layout.py index e7a29bac0f..9bc5c74351 100644 --- a/tests/test_lmcache_mp_native_layout.py +++ b/tests/test_lmcache_mp_native_layout.py @@ -386,7 +386,7 @@ def test_draft_regions_after_the_state_regions_stay_ordinary_page(): ) def test_native_and_page_only_registration_share_page_validation(breakage, message): """One validator for both registrations, so they cannot drift apart.""" - from atom.kv_transfer.offload.mp.backend import _build_cache_views + from atom.kv_transfer.offload.mp.page_views import _build_cache_views native = _transfer() breakage(native) diff --git a/tests/test_lmcache_mp_native_scheduler.py b/tests/test_lmcache_mp_native_scheduler.py index c9a2ac7092..2fc6e3a7b6 100644 --- a/tests/test_lmcache_mp_native_scheduler.py +++ b/tests/test_lmcache_mp_native_scheduler.py @@ -15,7 +15,8 @@ SaveSourceGroupId, ) from atom.kv_transfer.offload.chunked_scheduler import DENSE_PAGE_SOURCE_SAFE_CHANNEL -from atom.kv_transfer.offload.mp import backend +from atom.kv_transfer.offload.mp import deployment, transfer +from atom.kv_transfer.offload.mp import scheduler as mp_scheduler from atom.kv_transfer.offload.mp.connector import LMCacheMPConnectorScheduler from atom.kv_transfer.offload.mp.native_state_scheduler import ( NativeStateLMCacheMPConnectorScheduler, @@ -39,7 +40,7 @@ def storage_config(monkeypatch): # The transport is a double; retain config/chunk validation without needing # the optional LMCache package or a standalone server on a CPU test runner. monkeypatch.setattr( - backend.offcfg, + deployment.offcfg, "build_lmcache_config", lambda kvc: SimpleNamespace( chunk_size=kvc["kv_connector_extra_config"]["lmcache.chunk_size"] @@ -96,7 +97,7 @@ def connect(config, *, checkpoint_spec=None): connections.append(checkpoint_spec) return adapter - monkeypatch.setattr(backend, "_make_scheduler_adapter", connect) + monkeypatch.setattr(mp_scheduler, "_make_scheduler_adapter", connect) config = SimpleNamespace( kv_cache_block_size=4, kv_transfer_config={ @@ -405,7 +406,7 @@ def test_a_transfer_that_never_reports_stops_the_engine(monkeypatch, kind): a slow DMA, so past the deadline the scheduler fails stop -- holding the lease and budget rather than recycling them.""" now = [1000.0] - monkeypatch.setattr(backend.time, "monotonic", lambda: now[0]) + monkeypatch.setattr(transfer.time, "monotonic", lambda: now[0]) scheduler, checkpoints, _ = make_scheduler(monkeypatch) seq = sequence(computed=0 if kind == "load" else 16) if kind == "save": @@ -422,14 +423,14 @@ def test_a_transfer_that_never_reports_stops_the_engine(monkeypatch, kind): now[0] += scheduler._scheduler_deadline_s - 1 scheduler.process_completions(KVConnectorOutput()) now[0] += 2 - with pytest.raises(backend.LMCacheTransferUnprovable): + with pytest.raises(transfer.LMCacheTransferUnprovable): scheduler.process_completions(KVConnectorOutput()) assert scheduler._pinned_state_bytes == 30 def test_a_transfer_that_reports_leaves_the_watchdog(monkeypatch): now = [1000.0] - monkeypatch.setattr(backend.time, "monotonic", lambda: now[0]) + monkeypatch.setattr(transfer.time, "monotonic", lambda: now[0]) scheduler, checkpoints, _ = make_scheduler(monkeypatch) seq = sequence() checkpoint(scheduler, checkpoints, seq, 16) @@ -685,7 +686,7 @@ def engine_scheduler(monkeypatch): monkeypatch.setenv("OFFLOAD_MIN_SAVE_TOKENS", "0") monkeypatch.setenv("OFFLOAD_MAX_PENDING_SAVES", "2") monkeypatch.setattr( - backend, "_make_scheduler_adapter", lambda _config, **kwargs: adapter + mp_scheduler, "_make_scheduler_adapter", lambda _config, **kwargs: adapter ) config = MockConfig( num_kvcache_blocks=40, diff --git a/tests/test_lmcache_mp_native_worker.py b/tests/test_lmcache_mp_native_worker.py index d62a040c69..a701b98aac 100644 --- a/tests/test_lmcache_mp_native_worker.py +++ b/tests/test_lmcache_mp_native_worker.py @@ -26,7 +26,10 @@ NativeStateTransfer, SaveSpec, ) -from atom.kv_transfer.offload.mp.backend import _model_namespace, _tp_replication_factor +from atom.kv_transfer.offload.mp.deployment import ( + _model_namespace, + _tp_replication_factor, +) from atom.kv_transfer.offload.mp.native_state_layout import ( build_native_state_mp_layout, ) @@ -209,10 +212,10 @@ def test_unprovable_submission_retains_lease_until_the_deadline_stops_the_engine ): """No clock releases a lease the server may still be using: the operation stays pending, and past the transfer deadline the worker fails stop.""" - from atom.kv_transfer.offload.mp import backend + from atom.kv_transfer.offload.mp import transfer now = [1000.0] - monkeypatch.setattr(backend.time, "monotonic", lambda: now[0]) + monkeypatch.setattr(transfer.time, "monotonic", lambda: now[0]) def unprovable(*_): raise ConnectionError("server may have received request") @@ -231,7 +234,7 @@ def unprovable(*_): assert len(pending) == 1 now[0] += 2 - with pytest.raises(backend.LMCacheTransferUnprovable): + with pytest.raises(transfer.LMCacheTransferUnprovable): worker.get_finished() assert len(pending) == 1 @@ -537,10 +540,10 @@ def test_raising_restore_keeps_its_descriptor_slot_until_the_deadline( """A restore that raises after taking a descriptor slot may have queued a copy that still reads the slot's staging buffer: the slot and lease stay held, and the transfer deadline stops the engine instead of recycling.""" - from atom.kv_transfer.offload.mp import backend, native_state_worker + from atom.kv_transfer.offload.mp import native_state_worker, transfer now = [1000.0] - monkeypatch.setattr(backend.time, "monotonic", lambda: now[0]) + monkeypatch.setattr(transfer.time, "monotonic", lambda: now[0]) # A real Event cannot be built on a CPU runner; fail deterministically # after the slot is taken instead, at the stream fence. monkeypatch.setattr( @@ -557,7 +560,7 @@ def test_raising_restore_keeps_its_descriptor_slot_until_the_deadline( assert worker._restore_descriptor_slots == [] now[0] += worker._transfer_deadline_s - with pytest.raises(backend.LMCacheTransferUnprovable): + with pytest.raises(transfer.LMCacheTransferUnprovable): worker.get_finished() assert worker._restore_descriptor_slots == [] assert len(worker._native_loads) == 1 @@ -577,12 +580,12 @@ def test_native_server_chunk_mismatch_fails_before_registration(monkeypatch): def test_native_namespace_changes_with_image_codec(monkeypatch): - from atom.kv_transfer.offload.mp import backend + from atom.kv_transfer.offload.mp import deployment - monkeypatch.setattr(backend.offcfg, "build_lmcache_config", lambda _: object()) - monkeypatch.setattr(backend.offcfg, "lmcache_replica_world_size", lambda _: 2) + monkeypatch.setattr(deployment.offcfg, "build_lmcache_config", lambda _: object()) + monkeypatch.setattr(deployment.offcfg, "lmcache_replica_world_size", lambda _: 2) monkeypatch.setattr( - backend.offcfg, "build_page_namespace", lambda *_: "page-config" + deployment.offcfg, "build_page_namespace", lambda *_: "page-config" ) first = PagedStateCheckpointSpec(32, 128, "native-test-v1", 80) second = replace(first, image_bytes=81) diff --git a/tests/test_lmcache_mp_shell.py b/tests/test_lmcache_mp_shell.py index 3a8124e95f..29a156505d 100644 --- a/tests/test_lmcache_mp_shell.py +++ b/tests/test_lmcache_mp_shell.py @@ -7,10 +7,11 @@ from atom.kv_transfer.disaggregation.factory import KVConnectorFactory from atom.kv_transfer.offload.mp import ( - backend, native_state_scheduler, native_state_worker, ) +from atom.kv_transfer.offload.mp import scheduler as mp_scheduler +from atom.kv_transfer.offload.mp import worker as mp_worker from atom.kv_transfer.offload.mp.connector import ( LMCacheMPConnector, LMCacheMPConnectorScheduler, @@ -27,8 +28,8 @@ def _config(): def test_public_connector_is_a_layout_neutral_shell(): - assert LMCacheMPConnector is not backend.LMCacheMPConnector - assert LMCacheMPConnectorScheduler is not backend.LMCacheMPConnectorScheduler + assert LMCacheMPConnector is not mp_worker.LMCacheMPConnector + assert LMCacheMPConnectorScheduler is not mp_scheduler.LMCacheMPConnectorScheduler def test_factory_registration_resolves_public_connectors(monkeypatch): @@ -79,7 +80,7 @@ def __init__(self, config): monkeypatch.setattr( native_state_worker, "NativeStateLMCacheMPConnector", NativeWorker ) - monkeypatch.setattr(backend, "LMCacheMPConnector", PageWorker) + monkeypatch.setattr(mp_worker, "LMCacheMPConnector", PageWorker) config = _config() worker = LMCacheMPConnector(config) tensors = SimpleNamespace( @@ -110,7 +111,7 @@ def __init__(self, config): "NativeStateLMCacheMPConnectorScheduler", NativeScheduler, ) - monkeypatch.setattr(backend, "LMCacheMPConnectorScheduler", PageScheduler) + monkeypatch.setattr(mp_scheduler, "LMCacheMPConnectorScheduler", PageScheduler) config = _config() scheduler = LMCacheMPConnectorScheduler(config) manager = SimpleNamespace(paged_state_checkpoints=object()) @@ -149,8 +150,8 @@ class PageScheduler: def __init__(self, config): selected.append("scheduler") - monkeypatch.setattr(backend, "LMCacheMPConnector", PageWorker) - monkeypatch.setattr(backend, "LMCacheMPConnectorScheduler", PageScheduler) + monkeypatch.setattr(mp_worker, "LMCacheMPConnector", PageWorker) + monkeypatch.setattr(mp_scheduler, "LMCacheMPConnectorScheduler", PageScheduler) LMCacheMPConnector(_config()).register_kv_caches({}, SimpleNamespace(), 7) LMCacheMPConnectorScheduler(_config()).bind_block_manager( SimpleNamespace(paged_state_checkpoints=None) @@ -177,7 +178,7 @@ def __init__(self, config): def register_kv_caches(self, *_): pass - monkeypatch.setattr(backend, "LMCacheMPConnector", PageWorker) + monkeypatch.setattr(mp_worker, "LMCacheMPConnector", PageWorker) builder = SimpleNamespace( state_transfer=lambda: SimpleNamespace(copies=True), ) @@ -198,7 +199,7 @@ def __init__(self, config): def register_kv_caches(self, *_): pass - monkeypatch.setattr(backend, "LMCacheMPConnector", PageWorker) + monkeypatch.setattr(mp_worker, "LMCacheMPConnector", PageWorker) builder = SimpleNamespace( state_transfer=lambda: SimpleNamespace(copies=False), ) @@ -227,7 +228,7 @@ def __init__(self, config): NativeScheduler, ) monkeypatch.setattr( - backend, + mp_scheduler, "LMCacheMPConnectorScheduler", lambda config: selected.append("page"), ) From ecffc999deda3a2d15a275102225a6fb7a728cf0 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Mon, 28 Sep 2026 04:59:45 +0000 Subject: [PATCH 26/28] fix(offload): keep clock reclaim off MP requests; warn on truncated late saves - With the abandon window inherited again, a slow but legitimate MP save deferred past it was "abandoned" (a no-op for MP) and, a minute later, logged as a wedged P/D send -- both long before MP's own transfer deadline. A connector can now answer `waits_for_transfer_report(seq)`; the stalled-save reclaim skips such requests. LMCache MP answers True, the in-process shell forwards (default False), and the composite answers True only if every sub still deferring the request does, so P/D sends and in-process saves keep their abandon path. - A late save that finds part of its prefix evicted now logs a warning with the running count, not a debug line: early release returns the tail's blocks at teardown, and this is the cost of that trade. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../disaggregation/multi/multi_connector.py | 17 +++++++++++ atom/kv_transfer/offload/chunked_scheduler.py | 9 ++++-- atom/kv_transfer/offload/connector.py | 6 ++++ atom/kv_transfer/offload/mp/scheduler.py | 19 ++++++++---- atom/model_engine/scheduler.py | 13 +++++++++ tests/test_lmcache_mp_native_scheduler.py | 1 + tests/test_multi_connector.py | 26 +++++++++++++++++ tests/test_offload_early_block_release.py | 5 +++- tests/test_scheduler.py | 29 +++++++++++++++++++ 9 files changed, 117 insertions(+), 8 deletions(-) diff --git a/atom/kv_transfer/disaggregation/multi/multi_connector.py b/atom/kv_transfer/disaggregation/multi/multi_connector.py index d02e4236f5..f228cbcd63 100644 --- a/atom/kv_transfer/disaggregation/multi/multi_connector.py +++ b/atom/kv_transfer/disaggregation/multi/multi_connector.py @@ -496,6 +496,23 @@ def source_blocks_released(self, seq: Any) -> None: if callable(fn): fn(seq) + def waits_for_transfer_report(self, seq: Any) -> bool: + """True only if every sub still deferring this request waits for reports. + + Clock-based reclaim is skipped only when no deferring sub would act on + it; a P/D send or an in-process save still gets its abandon path. + """ + deferring = 0 + for connector in self._connectors: + should_defer = getattr(connector, "should_defer_free", None) + if not callable(should_defer) or not should_defer(seq): + continue + deferring += 1 + callback = getattr(connector, "waits_for_transfer_report", None) + if not callable(callback) or callback(seq) is not True: + return False + return deferring > 0 + def abandon_save(self, req_id: Any) -> None: # Reclamation of a stalled offload save (see # `DenseOffloadConnector.abandon_save`). Only the offload sub tracks diff --git a/atom/kv_transfer/offload/chunked_scheduler.py b/atom/kv_transfer/offload/chunked_scheduler.py index 6a7eeda5b9..4a2a93982d 100644 --- a/atom/kv_transfer/offload/chunked_scheduler.py +++ b/atom/kv_transfer/offload/chunked_scheduler.py @@ -467,11 +467,16 @@ def _late_save_source( # A block of the prefix was evicted (or reused) between teardown # and admission, so the final save stores less than was computed. self.total_truncated_late_saves += 1 - logger.debug( - "Late offload save of seq %s found %d of %d tokens resident", + # Early release trades this: the tail's blocks went back to the pool + # at teardown, so eviction before admission shortens the final save. + logger.warning( + "Late offload save of seq %s found %d of %d computed tokens still " + "resident; the evicted tail is not stored " + "(total truncated_late_saves=%d)", seq.id, available, aligned, + self.total_truncated_late_saves, ) if target <= saved: # Nothing savable is still resident (or no boundary qualifies): diff --git a/atom/kv_transfer/offload/connector.py b/atom/kv_transfer/offload/connector.py index a111936740..9b496e4aaa 100644 --- a/atom/kv_transfer/offload/connector.py +++ b/atom/kv_transfer/offload/connector.py @@ -336,6 +336,12 @@ def take_state_source_releases(self) -> set: callback = getattr(self._impl, "take_state_source_releases", None) return callback() if callback is not None else set() + def waits_for_transfer_report(self, seq) -> bool: + # In-process offload proves its copies finished, so the engine's + # clock-based reclaim applies unless the impl says otherwise. + callback = getattr(self._impl, "waits_for_transfer_report", None) + return bool(callback(seq)) if callback is not None else False + def save_abandon_timeout_s(self) -> float: # Plain forward: the abstract lifecycle contract guarantees every _impl # defines this (concrete on OffloadSchedulerMixin). The scheduler sources diff --git a/atom/kv_transfer/offload/mp/scheduler.py b/atom/kv_transfer/offload/mp/scheduler.py index 27c07c9fe6..7da1dc651f 100644 --- a/atom/kv_transfer/offload/mp/scheduler.py +++ b/atom/kv_transfer/offload/mp/scheduler.py @@ -155,13 +155,22 @@ def load_failed(self, req_id: Any) -> bool: ) return super().load_failed(req_id) + def waits_for_transfer_report(self, seq: Any) -> bool: + """MP releases memory only on a report; see `_enforce_transfer_deadlines`. + + The engine's clock-based reclaim cannot prove the MP server stopped + reading or writing, so it leaves MP requests alone instead of logging + them as abandoned or wedged; a report that never comes stops the engine. + """ + del seq + return True + def abandon_save(self, req_id: Any) -> None: - # The engine's clock-based reclaim cannot prove the MP server stopped - # reading the source, so an MP save keeps its lease until a terminal - # report. `_enforce_transfer_deadlines` bounds a report that never comes. - logger.warning( + # Reached only through a composite connector, which still abandons its + # other legs: an MP save keeps its lease until a terminal report. + logger.debug( "LMCache MP keeps the source of request %s leased until its save " - "reports; a transfer that never reports stops the engine after %.0fs", + "reports (deadline %.0fs)", req_id, self._scheduler_deadline_s, ) diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 4eac92d157..0968cb0669 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -1162,6 +1162,17 @@ def _connector_release_stalled_save(self, seq: Sequence) -> None: if callable(callback): callback(seq) + def _connector_waits_for_report(self, seq: Sequence) -> bool: + """Whether the connector refuses clock-based reclaim for this request. + + A connector that cannot prove a stalled transfer stopped touching the + blocks keeps them until a terminal report and bounds a missing report + itself (LMCache MP fails stop at its transfer deadline). Abandoning it + here would release nothing and only log a misleading wedge. + """ + callback = getattr(self.kv_connector, "waits_for_transfer_report", None) + return bool(callback(seq)) if callable(callback) else False + def _connector_abandon_save(self, seq: Sequence) -> None: """Tell the connector to drop a save this reclaim just abandoned. @@ -1310,6 +1321,8 @@ def _reconcile_stalled_deferred_saves(self) -> int: abandoned = 0 wedged: list = [] for seq in stalled: + if self._connector_waits_for_report(seq): + continue # The save is abandoned once; the release attempt is not. Whether # the blocks come back is a separate question -- a send that has # not reported still claims them -- and giving up on the first diff --git a/tests/test_lmcache_mp_native_scheduler.py b/tests/test_lmcache_mp_native_scheduler.py index 2fc6e3a7b6..5a8fc30408 100644 --- a/tests/test_lmcache_mp_native_scheduler.py +++ b/tests/test_lmcache_mp_native_scheduler.py @@ -366,6 +366,7 @@ def test_no_timeout_or_abandon_recycles_dispatched_source(monkeypatch): scheduler.update_state_after_alloc(seq) [request] = scheduler.build_connector_meta().requests assert scheduler.save_abandon_timeout_s() > 0 + assert scheduler.waits_for_transfer_report(seq) scheduler.abandon_save(request.save_operation) assert scheduler.reclaim_stale_leases(1e-9) == [] assert checkpoints.reclaim_stale_offload_pins(1e-9) == 0 diff --git a/tests/test_multi_connector.py b/tests/test_multi_connector.py index 9a051fb48a..1f36f87ca2 100644 --- a/tests/test_multi_connector.py +++ b/tests/test_multi_connector.py @@ -993,3 +993,29 @@ def test_should_defer_free_holds_while_any_sub_claims_the_source(): assert sched.should_defer_free(seq) is True, "the send still claims it" producer.defer = False assert sched.should_defer_free(seq) is False + + +def test_waits_for_transfer_report_needs_every_deferring_sub(): + """Clock reclaim is skipped only if no deferring sub would act on it.""" + from types import SimpleNamespace + + from atom.kv_transfer.disaggregation.multi.multi_connector import ( + MultiConnectorScheduler, + ) + + def sub(defers, waits=None): + fields = {"should_defer_free": lambda _seq: defers} + if waits is not None: + fields["waits_for_transfer_report"] = lambda _seq: waits + return SimpleNamespace(**fields) + + def multi(*subs): + connector = object.__new__(MultiConnectorScheduler) + connector._connectors = list(subs) + return connector + + seq = object() + assert multi(sub(True, True), sub(False)).waits_for_transfer_report(seq) + assert not multi(sub(True, True), sub(True)).waits_for_transfer_report(seq) + assert not multi(sub(True, True), sub(True, False)).waits_for_transfer_report(seq) + assert not multi(sub(False, True)).waits_for_transfer_report(seq) diff --git a/tests/test_offload_early_block_release.py b/tests/test_offload_early_block_release.py index 22e4a4f5e2..0e1cb2b7cf 100644 --- a/tests/test_offload_early_block_release.py +++ b/tests/test_offload_early_block_release.py @@ -420,7 +420,9 @@ def test_late_acquire_finds_a_multimodal_prefix(self, monkeypatch): assert block_ids == table assert sorted(claimed) == sorted(table) - def test_late_acquire_with_nothing_resident_retires_the_request(self, monkeypatch): + def test_late_acquire_with_nothing_resident_retires_the_request( + self, monkeypatch, caplog + ): scheduler = _early_release_scheduler(monkeypatch, chunk_size=8) bm, seq, table = _resident_sequence(scheduler, 103, 32, 8) scheduler.update_state_after_alloc(seq) @@ -435,6 +437,7 @@ def test_late_acquire_with_nothing_resident_retires_the_request(self, monkeypatc assert scheduler.protected_block_ids(seq) == frozenset() assert bm.kv.num_used == 1 assert scheduler.get_statistics()["truncated_late_saves"] == 1 + assert any("truncated_late_saves=1" in r.getMessage() for r in caplog.records) def test_late_save_persists_the_short_tail_of_a_long_request(self, monkeypatch): """The tail after earlier saves is stored however short it is.""" diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 92346a0675..bc0bd810f3 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -2760,6 +2760,35 @@ def test_save_timeout_does_not_release_a_pending_producer_send(self, monkeypatch assert not s.deferred_free_blocks assert seq._deferred_save_at is None + def test_a_connector_that_waits_for_reports_is_left_alone( + self, monkeypatch, caplog + ): + """LMCache MP cannot prove a stalled DMA stopped, keeps the blocks until + a report and fails stop itself at its deadline. Reclaim must neither + abandon it nor, after the retries, log it as a wedged P/D send.""" + import logging + import time as _time + + import atom.model_engine.scheduler as sched_mod + + seq = SimpleNamespace(id=1, _deferred_save_at=_time.monotonic() - 500.0) + abandoned: list = [] + connector = SimpleNamespace( + save_abandon_timeout_s=lambda: 100.0, + abandon_save=abandoned.append, + should_defer_free=lambda _seq: True, + waits_for_transfer_report=lambda _seq: True, + ) + s, freed = self._sched(monkeypatch, [seq], connector=connector) + with caplog.at_level(logging.WARNING, logger="atom"): + for _ in range(sched_mod._RELEASE_WEDGE_ATTEMPTS + 1): + s._next_save_reconcile_at = 0.0 + assert s._reconcile_stalled_deferred_saves() == 0 + assert not abandoned and not freed + assert s._abandoned_saves == 0 + assert seq.id in s.deferred_free_blocks + assert not caplog.records + def test_a_claim_nobody_will_retire_is_escalated(self, monkeypatch, caplog): """The retry is not an escape hatch on its own, so it must be visible. From 069718a519ba3c28f99af6b2c0dd162376015391 Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Mon, 28 Sep 2026 06:57:32 +0000 Subject: [PATCH 27/28] refactor(kv_transfer): publish each PAGE region and its byte view together Every backend kept `block_regions` and `block_tensor_views` as two lists paired by hand (DSV4 even collected a third list of sources to zip later). `KVTransferTensors.add_block_region(tensor, semantic_role=...)` now appends both from one tensor: the region's addresses and a zero-copy `uint8 [num_units, 1, unit_bytes]` alias of exactly those bytes, cut from a larger allocation when `total_bytes` says so. DSV4, MHA, the MHA draft and MLA all publish through it. MHA and the draft already published this shape. MLA's views change from `[n, block_size, width]` to the same byte form; LMCache stores the same bytes in the same order either way, and neither its object key nor ATOM's model namespace depends on view shape, so existing cache entries stay valid. The MLA builder test's module stubs had drifted since #2399 (`atom.utils.block_tables`); they are fixed so the test runs again. Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/kv_transfer/disaggregation/types.py | 58 +++++++++++++++-- atom/model_ops/attentions/aiter_attention.py | 41 ++++-------- atom/model_ops/attentions/aiter_mla.py | 33 +++------- atom/model_ops/attentions/deepseek_v4_attn.py | 63 ++++--------------- atom/spec_decode/draft_kv.py | 22 +++---- tests/test_lmcache_mp.py | 44 +++++++++++++ tests/test_lmcache_mp_builder_views.py | 14 +++-- 7 files changed, 150 insertions(+), 125 deletions(-) diff --git a/atom/kv_transfer/disaggregation/types.py b/atom/kv_transfer/disaggregation/types.py index 2dafedf4e1..dbdce4bb2a 100644 --- a/atom/kv_transfer/disaggregation/types.py +++ b/atom/kv_transfer/disaggregation/types.py @@ -231,11 +231,11 @@ class KVTransferTensors: scatter_slot: Callable[[int, int], None] | None = None # Appended for positional compatibility with existing generic descriptors. expected_full_slot_region_count: int | None = None - # Zero-copy, physical-order views paired with block regions. - # Each view is [num_units, physical_slots_per_unit, opaque_width]; it may - # retain its native dtype, and one dim-0 unit must cover exactly the paired - # region's unit_bytes. The last two dimensions describe copy geometry, not - # a semantic token/head layout. + # Zero-copy, physical-order views paired with block regions. Each view is + # [num_units, physical_slots_per_unit, opaque_width] and one dim-0 unit + # covers exactly the paired region's unit_bytes; the last two dimensions + # describe copy geometry, not a token/head layout. Backends publish both + # through `add_block_region`, which keeps the pair in step. block_tensor_views: list[Any] = field(default_factory=list) # Number of TP workers whose complete PAGE layout is byte-identical. # ``1`` means no cross-rank deduplication is safe. @@ -280,6 +280,54 @@ class KVTransferTensors: # the same number. Set through `set_block_count`. num_blocks: int = field(init=False, default=0) + def add_block_region( + self, + tensor: Any, + *, + semantic_role: str, + unit_bytes: int | None = None, + total_bytes: int | None = None, + ) -> None: + """Publish one PAGE region and its block-major byte view, built together. + + The region (addresses for P/D and dense offload) and the view (a tensor + for LMCache MP) describe the same bytes, so they come from one place + rather than two lists kept in step by hand. The view is a zero-copy + ``uint8 [num_units, 1, unit_bytes]`` alias of ``tensor``: the last two + dimensions are opaque copy geometry, and LMCache stores the same bytes + in the same order whatever shape a block is given. + + ``unit_bytes`` defaults to ``tensor``'s row stride and ``total_bytes`` + to all of it. A larger allocation that holds more than the PAGE units + (DSV4's planes also hold SLOT rows after them) passes both. + """ + import torch # this module stays importable without torch + + if not tensor.is_contiguous(): + raise ValueError(f"PAGE tensor {semantic_role!r} must be contiguous") + element_size = tensor.element_size() + if unit_bytes is None: + unit_bytes = tensor.stride(0) * element_size + if total_bytes is None: + total_bytes = tensor.numel() * element_size + available = tensor.numel() * element_size + if unit_bytes <= 0 or total_bytes % unit_bytes or total_bytes > available: + raise ValueError( + f"PAGE tensor {semantic_role!r} has {available} bytes; cannot " + f"publish {total_bytes} bytes in units of {unit_bytes}" + ) + # `view`, never `reshape`: a copy would publish bytes nobody writes. + byte_view = tensor.view(torch.uint8).view(-1)[:total_bytes] + self.block_regions.append( + KVTransferRegion( + base_addr=tensor.data_ptr(), + total_bytes=total_bytes, + unit_bytes=unit_bytes, + semantic_role=semantic_role, + ) + ) + self.block_tensor_views.append(byte_view.view(-1, 1, unit_bytes)) + def set_block_count(self, num_blocks: int) -> None: """Fix the block id space, and check every region is in it. diff --git a/atom/model_ops/attentions/aiter_attention.py b/atom/model_ops/attentions/aiter_attention.py index 55a1b1be8d..d5e0551d89 100644 --- a/atom/model_ops/attentions/aiter_attention.py +++ b/atom/model_ops/attentions/aiter_attention.py @@ -773,43 +773,28 @@ def _get_sparse_attention_block_tables( ) def get_kv_transfer_tensors(self): - from atom.kv_transfer.disaggregation.types import ( - KVTransferRegion, - KVTransferTensors, - ) + from atom.kv_transfer.disaggregation.types import KVTransferTensors if not self.kv_pools: return None - region_tensors = [ - (geometry, role, tensor) - for geometry, pool in self.kv_pools.items() - for role, tensor in pool.region_tensors() - ] - # Every field of every pool, indexer cache included: a region per - # (pool, field, layer), in declared order. - return KVTransferTensors( - block_regions=[ - KVTransferRegion( - base_addr=tensor.data_ptr(), - total_bytes=tensor.numel() * tensor.element_size(), - unit_bytes=tensor.stride(0) * tensor.element_size(), - # The geometry, because a hybrid declares two pools whose - # per-layer regions are otherwise named alike. - semantic_role=f"mha.{geometry}.{role}", - ) - for geometry, role, tensor in region_tensors - ], + transfer = KVTransferTensors( + block_regions=[], slot_regions=[], - # MHA arenas expose one contiguous byte row per scheduler block. - # LMCache MP deliberately treats the last two dimensions as opaque - # copy geometry, so adding a singleton physical-slot dimension is - # a zero-copy view of the exact region declared above. - block_tensor_views=[tensor.unsqueeze(1) for _, _, tensor in region_tensors], # GQA/MQA KV heads are sharded or only partially replicated across # TP. In particular MiniMax-M3 must keep one stored shard per rank; # whole-object TP collapse is unsafe for this PAGE layout. tp_replication_factor=1, ) + # Every field of every pool, indexer cache included: a region per + # (pool, field, layer), in declared order. The geometry is in the role + # because a hybrid declares two pools whose per-layer regions are + # otherwise named alike. + for geometry, pool in self.kv_pools.items(): + for role, tensor in pool.region_tensors(): + transfer.add_block_region( + tensor, semantic_role=f"mha.{geometry}.{role}" + ) + return transfer def refresh_flydsl_plan(self, context_lens, *, create=False): """Build or refresh aiter #5546's work plan for this forward. diff --git a/atom/model_ops/attentions/aiter_mla.py b/atom/model_ops/attentions/aiter_mla.py index 94656f0ac2..338fdb7124 100644 --- a/atom/model_ops/attentions/aiter_mla.py +++ b/atom/model_ops/attentions/aiter_mla.py @@ -1447,32 +1447,19 @@ def get_kv_transfer_tensors(self): # and no per-field override to keep in step with the pooling ones. # DCP PD still dispatches on the two collapsed roles (`mla.kv` / # `dsa.index_cache`) rather than the pool's per-layer names. - region_tensors = self.kv_pool.region_tensors() - block_regions: list[KVTransferRegion] = [] - for role, t in region_tensors: - bpb = t.stride(0) * t.element_size() + pages = KVTransferTensors(block_regions=[], slot_regions=[]) + for role, t in self.kv_pool.region_tensors(): if role.startswith("index."): - block_regions.append( - KVTransferRegion( - base_addr=t.data_ptr(), - total_bytes=t.numel() * t.element_size(), - unit_bytes=bpb, - semantic_role=INDEX_CACHE_ROLE, - ) - ) + pages.add_block_region(t, semantic_role=INDEX_CACHE_ROLE) index_tensors.append(t) continue - block_regions.append( - KVTransferRegion( - base_addr=t.data_ptr(), - total_bytes=t.numel() * t.element_size(), - unit_bytes=bpb, - semantic_role=( - MLA_KV_ROLE if role.startswith("kv.") else f"mla.{role}" - ), - ) + pages.add_block_region( + t, + semantic_role=( + MLA_KV_ROLE if role.startswith("kv.") else f"mla.{role}" + ), ) - block_tensor_views = [t for _, t in region_tensors] + block_regions = pages.block_regions block_region_consumer_indices = None index_cache_layer_ids = getattr(runner, "index_cache_layer_ids", ()) @@ -1637,7 +1624,7 @@ def gather_sharded_index( return KVTransferTensors( block_regions=block_regions, slot_regions=[], - block_tensor_views=block_tensor_views, + block_tensor_views=pages.block_tensor_views, block_region_consumer_indices=block_region_consumer_indices, index_staging_region=index_staging_region, index_staging_pool_size=index_staging_pool_size, diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index a701135fd9..bfe7ed6fbc 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -1922,8 +1922,8 @@ def get_kv_transfer_tensors(self): geo = self.pool_geometry elem_fp32 = 4 - block_regions: list[KVTransferRegion] = [] - block_tensor_sources: list[torch.Tensor] = [] + # PAGE regions and their LMCache MP byte views, published as pairs. + pages = KVTransferTensors(block_regions=[], slot_regions=[]) swa_block_regions: list[KVTransferRegion] = [] slot_regions: list[KVTransferRegion] = [] @@ -1949,16 +1949,13 @@ def get_kv_transfer_tensors(self): ) ) for plane, row_bytes, role in planes: - if not plane.is_contiguous(): - raise RuntimeError("a KV plane must be contiguous to be transferred") - block_tensor_sources.append(plane) - block_regions.append( - KVTransferRegion( - plane.data_ptr(), - self.num_blocks * geo.block_bytes(row_bytes), - geo.block_bytes(row_bytes), - semantic_role=role, - ) + # A plane also holds SLOT rows after its PAGE blocks; publish only + # the blocks. + pages.add_block_region( + plane, + semantic_role=role, + unit_bytes=geo.block_bytes(row_bytes), + total_bytes=self.num_blocks * geo.block_bytes(row_bytes), ) # Compressed PAGE regions: one per CSA layer in each indexer pool. @@ -1967,20 +1964,10 @@ def get_kv_transfer_tensors(self): # this remains correct for both the FP8 row layout and FP4 tiles. for pool, role_prefix in self._indexer_page_pools(): for pos, layer_id in enumerate(self.csa_layers): - view = pool[pos] - if not view.is_contiguous(): - raise RuntimeError( - "a CSA indexer layer must be contiguous to be transferred" - ) - block_tensor_sources.append(view) - block_regions.append( - KVTransferRegion( - view.data_ptr(), - view.numel() * view.element_size(), - view.stride(0) * view.element_size(), - semantic_role=f"{role_prefix}.layer_{layer_id}", - ) + pages.add_block_region( + pool[pos], semantic_role=f"{role_prefix}.layer_{layer_id}" ) + block_regions = pages.block_regions checkpoint_spec = runner.state_runtime.checkpoint_spec if checkpoint_spec is None: @@ -1993,30 +1980,6 @@ def get_kv_transfer_tensors(self): f"checkpoint={checkpoint_spec.page_unit_bytes}" ) - # LMCache MP needs one block-major tensor for every PAGE region. Build - # byte views only after validating the declared checkpoint geometry so - # malformed layouts fail with the useful region-size error above. The - # source tensors are not guaranteed to expose rows as dimension zero; - # flattening their byte representation also handles the one-dimensional - # allocation owners used by the transfer-geometry tests. - block_tensor_views: list[torch.Tensor] = [] - for source, region in zip(block_tensor_sources, block_regions, strict=True): - byte_view = source.view(torch.uint8).reshape(-1) - required_bytes = self.num_blocks * region.unit_bytes - if byte_view.numel() < required_bytes: - raise RuntimeError( - f"DSV4 PAGE tensor for {region.semantic_role!r} has " - f"{byte_view.numel()} bytes, fewer than the " - f"{required_bytes} bytes declared by its region" - ) - block_tensor_views.append( - byte_view[:required_bytes].view( - self.num_blocks, - 1, - region.unit_bytes, - ) - ) - # Full per-request SLOT (legacy field name: `swa_block_regions`) is one # slot per plane — compressor state and SWA, every layer at once — and # the id the connector zips over is its pool group. @@ -2080,7 +2043,7 @@ def get_kv_transfer_tensors(self): staging_pool_size=pool_size if staging_region else 0, gather_slot=gather_slot, scatter_slot=scatter_slot, - block_tensor_views=block_tensor_views, + block_tensor_views=pages.block_tensor_views, tp_replication_factor=tp_size, native_state_tp_replication_factor=tp_size, paged_state_checkpoint_spec=checkpoint_spec, diff --git a/atom/spec_decode/draft_kv.py b/atom/spec_decode/draft_kv.py index c91b1ba20e..2c92ec8935 100644 --- a/atom/spec_decode/draft_kv.py +++ b/atom/spec_decode/draft_kv.py @@ -189,26 +189,18 @@ def build_kv_cache_tensor(self, module): def get_kv_transfer_tensors(self): from atom.kv_transfer.disaggregation.types import ( - KVTransferRegion, KVTransferTensors, ) - region_tensors = self.kv_pool.region_tensors() - return KVTransferTensors( - block_regions=[ - KVTransferRegion( - base_addr=t.data_ptr(), - total_bytes=t.numel() * t.element_size(), - unit_bytes=t.stride(0) * t.element_size(), - # The draft's rows sit in the same block ids as the target's, - # so its regions need a name that says which stack they are. - semantic_role=f"draft.{role}", - ) - for role, t in region_tensors - ], + transfer = KVTransferTensors( + block_regions=[], slot_regions=[], - block_tensor_views=[t.unsqueeze(1) for _, t in region_tensors], # A paged MHA/GQA draft is a per-rank shard. Appending it to an # otherwise replicated target must disable whole-object collapse. tp_replication_factor=1, ) + # The draft's rows sit in the same block ids as the target's, so its + # regions need a name that says which stack they are. + for role, t in self.kv_pool.region_tensors(): + transfer.add_block_region(t, semantic_role=f"draft.{role}") + return transfer diff --git a/tests/test_lmcache_mp.py b/tests/test_lmcache_mp.py index c73a12669f..ab1d4d2379 100644 --- a/tests/test_lmcache_mp.py +++ b/tests/test_lmcache_mp.py @@ -1862,3 +1862,47 @@ def test_factory_registers_lmcache_mp_alias_without_pd_staging(): ) is False ) + + +def test_add_block_region_publishes_a_region_and_its_byte_view_together(): + """One call yields both halves, zero-copy, over exactly the PAGE bytes: + a larger allocation (DSV4's planes also hold SLOT rows) is cut to them.""" + arena = torch.arange(3 * 8 + 4, dtype=torch.int16) # 3 blocks + a tail + transfer = KVTransferTensors(block_regions=[], slot_regions=[]) + transfer.add_block_region( + arena, semantic_role="plane", unit_bytes=16, total_bytes=3 * 16 + ) + fp16 = torch.zeros(3, 4, 2, dtype=torch.float16) + transfer.add_block_region(fp16, semantic_role="rows") + + [plane, rows] = transfer.block_regions + assert (plane.base_addr, plane.total_bytes, plane.unit_bytes) == ( + arena.data_ptr(), + 48, + 16, + ) + assert (rows.total_bytes, rows.unit_bytes) == (48, 16) + for region, view in zip( + transfer.block_regions, transfer.block_tensor_views, strict=True + ): + assert view.dtype == torch.uint8 + assert tuple(view.shape) == (3, 1, 16) + assert view.data_ptr() == region.base_addr + transfer.block_tensor_views[1][2].fill_(1) + assert torch.all(fp16[2].view(torch.uint8) == 1) + transfer.set_block_count(3) + assert page_views._build_cache_views(transfer, num_blocks=3).bytes_per_block == 32 + + +def test_add_block_region_rejects_what_it_cannot_alias(): + transfer = KVTransferTensors(block_regions=[], slot_regions=[]) + with pytest.raises(ValueError, match="contiguous"): + transfer.add_block_region(torch.zeros(4, 3).t(), semantic_role="strided") + with pytest.raises(ValueError, match="cannot publish"): + transfer.add_block_region( + torch.zeros(8, dtype=torch.uint8), + semantic_role="short", + unit_bytes=4, + total_bytes=12, + ) + assert transfer.block_regions == [] and transfer.block_tensor_views == [] diff --git a/tests/test_lmcache_mp_builder_views.py b/tests/test_lmcache_mp_builder_views.py index 1fa46f3aa1..d1c175d1d2 100644 --- a/tests/test_lmcache_mp_builder_views.py +++ b/tests/test_lmcache_mp_builder_views.py @@ -171,6 +171,9 @@ def noop(*args, **kwargs): kv_indices_generate_triton=noop, mtp_prepare_decode_mla_kernel=noop, ), + "atom.utils.block_tables": _module( + "atom.utils.block_tables", block_table_state=noop + ), "atom.utils.forward_context": _module( "atom.utils.forward_context", AttentionMetaData=type("AttentionMetaData", (), {}), @@ -263,10 +266,13 @@ def test_mla_builder_publishes_latent_and_index_views( assert transfer.num_blocks == 0 transfer.set_block_count(2) + # Published as block-major bytes, like every other backend: one fp16 + # [2 rows, 7] block is 28 bytes, one index block index_rows * 3. + assert all(view.dtype == torch.uint8 for view in transfer.block_tensor_views) assert [tuple(view.shape) for view in transfer.block_tensor_views] == [ - (2, 2, 7), - (2, 2, 7), - ] + [(2, index_rows, 3)] * index_layers + (2, 1, 28), + (2, 1, 28), + ] + [(2, 1, index_rows * 3)] * index_layers # Main's DCP transfer contract intentionally collapses physical rows into # transport roles. MP still keeps them distinct through the plane index in # each tensor key (``page..``). @@ -282,4 +288,4 @@ def test_mla_builder_publishes_latent_and_index_views( assert cache_views.bytes_per_block == 2 * 2 * 7 * 2 + index_layers * index_rows * 3 # Transfer writes must update the allocation used by attention kernels. transfer.block_tensor_views[0][1].fill_(7) - assert torch.all(pool.layer("kv", 0)[1] == 7) + assert torch.all(pool.layer("kv", 0)[1].view(torch.uint8) == 7) From 603db8d070f6b9f677635fae9f58f3b3498584fa Mon Sep 17 00:00:00 2001 From: Honglie Yi Date: Mon, 28 Sep 2026 07:12:33 +0000 Subject: [PATCH 28/28] refactor(kv_transfer): make a PAGE unit one value, its region and view together `add_block_region` paired a region and its view at one call site, but the pair was only a convention: both lists stayed public and mutable, the constructor still took them separately, DSV4 and MLA used a throwaway `KVTransferTensors` as a scratch builder, the draft merge extended the two lists by hand, and tensor code lived in the torch-free `types` contract. - `PageRegion(region, view)` is a frozen value; `KVTransferTensors.pages` holds them. `block_regions` and `block_tensor_views` are read-only tuples derived from it, so no code can add to one without the other. - `page_region(tensor, ...)` in the new `disaggregation/page_region.py` builds one from the owning tensor (zero-copy byte view, contiguity and size checks); `types.py` no longer touches torch. - Builders pass `pages=[...]` straight to the final object: no scratch instance. The address-only producer (Qwen4 exp) publishes `PageRegion(region)` with no view, which LMCache MP refuses as before. - `merge_pages(other)` replaces the draft merge in `ModelRunner`, carrying the gcd replication rule with it. `validate_page_views` stays on the LMCache MP side: it is the consumer checking what it is handed, not a second copy of construction. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../kv_transfer/disaggregation/page_region.py | 59 ++++++++++ atom/kv_transfer/disaggregation/types.py | 107 ++++++++---------- atom/kv_transfer/offload/README.md | 5 +- atom/model_engine/model_runner.py | 18 +-- atom/model_ops/attentions/aiter_attention.py | 24 ++-- atom/model_ops/attentions/aiter_mla.py | 23 ++-- atom/model_ops/attentions/deepseek_v4_attn.py | 31 ++--- atom/model_ops/attentions/qwen4_exp_attn.py | 6 +- atom/spec_decode/draft_kv.py | 22 ++-- tests/test_deepseek_v4_transfer_regions.py | 32 +++--- tests/test_lmcache_mp.py | 77 +++++++++---- tests/test_lmcache_mp_native_layout.py | 30 +++-- tests/test_lmcache_mp_native_worker.py | 23 ++-- 13 files changed, 272 insertions(+), 185 deletions(-) create mode 100644 atom/kv_transfer/disaggregation/page_region.py diff --git a/atom/kv_transfer/disaggregation/page_region.py b/atom/kv_transfer/disaggregation/page_region.py new file mode 100644 index 0000000000..22204c8b7f --- /dev/null +++ b/atom/kv_transfer/disaggregation/page_region.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Build a `PageRegion` from the tensor that owns its bytes. + +Kept out of `types`, which stays a torch-free contract: this is where a +backend's allocation becomes both halves of a PAGE unit at once. +""" + +from __future__ import annotations + +import torch + +from atom.kv_transfer.disaggregation.types import KVTransferRegion, PageRegion + + +def page_region( + tensor: torch.Tensor, + *, + semantic_role: str, + unit_bytes: int | None = None, + total_bytes: int | None = None, +) -> PageRegion: + """The region of ``tensor``'s PAGE blocks and a byte view aliasing them. + + The view is a zero-copy ``uint8 [num_units, 1, unit_bytes]`` alias: the + last two dimensions are opaque copy geometry, and LMCache stores the same + bytes in the same order whatever shape a block is given. + + ``unit_bytes`` defaults to ``tensor``'s row stride and ``total_bytes`` to + all of it. An allocation that holds more than its PAGE blocks (DSV4's + planes also hold SLOT rows after them) passes both. + """ + if not tensor.is_contiguous(): + raise ValueError(f"PAGE tensor {semantic_role!r} must be contiguous") + available = tensor.numel() * tensor.element_size() + if unit_bytes is None: + unit_bytes = tensor.stride(0) * tensor.element_size() + if total_bytes is None: + total_bytes = available + if unit_bytes <= 0 or total_bytes % unit_bytes or total_bytes > available: + raise ValueError( + f"PAGE tensor {semantic_role!r} has {available} bytes; cannot " + f"publish {total_bytes} bytes in units of {unit_bytes}" + ) + # `view`, never `reshape`: a copy would publish bytes nobody writes. + byte_view = tensor.view(torch.uint8).view(-1)[:total_bytes] + return PageRegion( + region=KVTransferRegion( + base_addr=tensor.data_ptr(), + total_bytes=total_bytes, + unit_bytes=unit_bytes, + semantic_role=semantic_role, + ), + view=byte_view.view(-1, 1, unit_bytes), + ) + + +__all__ = ["page_region"] diff --git a/atom/kv_transfer/disaggregation/types.py b/atom/kv_transfer/disaggregation/types.py index dbdce4bb2a..540e842a9d 100644 --- a/atom/kv_transfer/disaggregation/types.py +++ b/atom/kv_transfer/disaggregation/types.py @@ -11,6 +11,7 @@ from __future__ import annotations +import math from collections.abc import Callable from dataclasses import dataclass, field from typing import Any @@ -189,11 +190,29 @@ def unit_addr(self, index: int) -> int: return self.base_addr + index * self.unit_bytes +@dataclass(frozen=True) +class PageRegion: + """One PAGE region and the tensor that aliases exactly its bytes. + + The region is the address form every transport reads (P/D, dense offload); + the view is the tensor form LMCache MP registers. They describe the same + bytes, so they travel as one value and cannot fall out of step. A producer + that only has addresses leaves ``view`` None; LMCache MP then refuses the + layout rather than register a partial one. Build one from a tensor with + `atom.kv_transfer.disaggregation.page_region.page_region`. + """ + + region: KVTransferRegion + view: Any = None + + @dataclass class KVTransferTensors: """Physical PAGE, SLOT, and compressor-only staging region contract. - ``block_regions`` contain forward-indexed, block-indexed PAGE units. + ``pages`` are forward-indexed, block-indexed PAGE units, each with its + region and (when published) its tensor view; ``block_regions`` and + ``block_tensor_views`` read them back. ``swa_block_regions`` is a legacy field name for complete reverse-indexed per-request SLOT units, including both compressor state and SWA. ``expected_full_slot_region_count`` makes a stateful layout's complete @@ -213,9 +232,9 @@ class KVTransferTensors: protocol can describe replication per region group. """ - # Block-indexed PAGE regions, indexed forward by block id. - block_regions: list[KVTransferRegion] - slot_regions: list[KVTransferRegion] + # Block-indexed PAGE regions with their views, indexed forward by block id. + pages: list[PageRegion] = field(default_factory=list) + slot_regions: list[KVTransferRegion] = field(default_factory=list) num_slots: int = 0 # Optional producer-local -> consumer-global mapping for non-uniform block # region layouts. Uniform per-layer groups leave this unset and use the @@ -231,12 +250,6 @@ class KVTransferTensors: scatter_slot: Callable[[int, int], None] | None = None # Appended for positional compatibility with existing generic descriptors. expected_full_slot_region_count: int | None = None - # Zero-copy, physical-order views paired with block regions. Each view is - # [num_units, physical_slots_per_unit, opaque_width] and one dim-0 unit - # covers exactly the paired region's unit_bytes; the last two dimensions - # describe copy geometry, not a token/head layout. Backends publish both - # through `add_block_region`, which keeps the pair in step. - block_tensor_views: list[Any] = field(default_factory=list) # Number of TP workers whose complete PAGE layout is byte-identical. # ``1`` means no cross-rank deduplication is safe. tp_replication_factor: int = 1 @@ -280,59 +293,39 @@ class KVTransferTensors: # the same number. Set through `set_block_count`. num_blocks: int = field(init=False, default=0) - def add_block_region( - self, - tensor: Any, - *, - semantic_role: str, - unit_bytes: int | None = None, - total_bytes: int | None = None, - ) -> None: - """Publish one PAGE region and its block-major byte view, built together. - - The region (addresses for P/D and dense offload) and the view (a tensor - for LMCache MP) describe the same bytes, so they come from one place - rather than two lists kept in step by hand. The view is a zero-copy - ``uint8 [num_units, 1, unit_bytes]`` alias of ``tensor``: the last two - dimensions are opaque copy geometry, and LMCache stores the same bytes - in the same order whatever shape a block is given. - - ``unit_bytes`` defaults to ``tensor``'s row stride and ``total_bytes`` - to all of it. A larger allocation that holds more than the PAGE units - (DSV4's planes also hold SLOT rows after them) passes both. + @property + def block_regions(self) -> tuple[KVTransferRegion, ...]: + """Every PAGE region, in publication order (read-only; add `pages`).""" + return tuple(page.region for page in self.pages) + + @property + def block_tensor_views(self) -> tuple[Any, ...]: + """The published views, each a zero-copy byte alias of its region. + + ``uint8 [num_units, 1, unit_bytes]``: the last two dimensions are + opaque copy geometry for LMCache MP, not a token/head layout. Empty when + the producer published addresses only. + """ + return tuple(page.view for page in self.pages if page.view is not None) + + def merge_pages(self, other: KVTransferTensors) -> None: + """Append another contributor's PAGE units (a draft's own pool). + + Replication is a whole-PAGE declaration: one sharded contributor makes + the combined object sharded even when this one alone is replicated. """ - import torch # this module stays importable without torch - - if not tensor.is_contiguous(): - raise ValueError(f"PAGE tensor {semantic_role!r} must be contiguous") - element_size = tensor.element_size() - if unit_bytes is None: - unit_bytes = tensor.stride(0) * element_size - if total_bytes is None: - total_bytes = tensor.numel() * element_size - available = tensor.numel() * element_size - if unit_bytes <= 0 or total_bytes % unit_bytes or total_bytes > available: - raise ValueError( - f"PAGE tensor {semantic_role!r} has {available} bytes; cannot " - f"publish {total_bytes} bytes in units of {unit_bytes}" - ) - # `view`, never `reshape`: a copy would publish bytes nobody writes. - byte_view = tensor.view(torch.uint8).view(-1)[:total_bytes] - self.block_regions.append( - KVTransferRegion( - base_addr=tensor.data_ptr(), - total_bytes=total_bytes, - unit_bytes=unit_bytes, - semantic_role=semantic_role, - ) + if not other.pages: + return + self.pages.extend(other.pages) + self.tp_replication_factor = math.gcd( + self.tp_replication_factor, other.tp_replication_factor ) - self.block_tensor_views.append(byte_view.view(-1, 1, unit_bytes)) def set_block_count(self, num_blocks: int) -> None: """Fix the block id space, and check every region is in it. - Called once the last contributor's regions are in the list; a draft - appends its own after construction, so this cannot be a `__post_init__`. + Called once the last contributor's regions are in; a draft's are merged + after construction (`merge_pages`), so this cannot be a `__post_init__`. A region that does not divide into exactly `num_blocks` units was registered in some other unit. Both ends would still agree on diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index 31709d9c0f..cdd11a43a6 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -332,7 +332,10 @@ flowchart LR different concepts. They must not be substituted for one another: ```text -block_regions PAGE: physical-block units, forward indexed +pages PAGE: physical-block units, forward indexed; each a + PageRegion (its region, plus the byte view LMCache MP + registers). `block_regions` / `block_tensor_views` read + the two halves back swa_block_regions SLOT: complete request-slot units, reverse indexed (the field name is historical) staging_region compressor-only P/D staging; never a full sidecar source diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index 7562e163e9..34cd360faf 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -2103,21 +2103,9 @@ def allocate_kv_cache(self, num_kvcache_blocks): # place the builder and the connector are both in scope. transfer_tensors.state_backend = self.attn_metadata_builder if hasattr(self, "draft_kv_builder") and transfer_tensors is not None: - draft_transfer_tensors = self.draft_kv_builder.get_kv_transfer_tensors() - if draft_transfer_tensors.block_regions: - transfer_tensors.block_regions.extend( - draft_transfer_tensors.block_regions - ) - transfer_tensors.block_tensor_views.extend( - draft_transfer_tensors.block_tensor_views - ) - # Replication is a whole-PAGE declaration. One sharded draft - # region makes the combined target+draft object sharded even - # when the target alone is replicated. - transfer_tensors.tp_replication_factor = math.gcd( - transfer_tensors.tp_replication_factor, - draft_transfer_tensors.tp_replication_factor, - ) + transfer_tensors.merge_pages( + self.draft_kv_builder.get_kv_transfer_tensors() + ) if transfer_tensors is not None: # After the draft's regions are in, and here because this is the # only place holding both the complete region list and the diff --git a/atom/model_ops/attentions/aiter_attention.py b/atom/model_ops/attentions/aiter_attention.py index d5e0551d89..d494a90adb 100644 --- a/atom/model_ops/attentions/aiter_attention.py +++ b/atom/model_ops/attentions/aiter_attention.py @@ -773,28 +773,26 @@ def _get_sparse_attention_block_tables( ) def get_kv_transfer_tensors(self): + from atom.kv_transfer.disaggregation.page_region import page_region from atom.kv_transfer.disaggregation.types import KVTransferTensors if not self.kv_pools: return None - transfer = KVTransferTensors( - block_regions=[], - slot_regions=[], + return KVTransferTensors( + # Every field of every pool, indexer cache included: a unit per + # (pool, field, layer), in declared order. The geometry is in the + # role because a hybrid declares two pools whose per-layer regions + # are otherwise named alike. + pages=[ + page_region(tensor, semantic_role=f"mha.{geometry}.{role}") + for geometry, pool in self.kv_pools.items() + for role, tensor in pool.region_tensors() + ], # GQA/MQA KV heads are sharded or only partially replicated across # TP. In particular MiniMax-M3 must keep one stored shard per rank; # whole-object TP collapse is unsafe for this PAGE layout. tp_replication_factor=1, ) - # Every field of every pool, indexer cache included: a region per - # (pool, field, layer), in declared order. The geometry is in the role - # because a hybrid declares two pools whose per-layer regions are - # otherwise named alike. - for geometry, pool in self.kv_pools.items(): - for role, tensor in pool.region_tensors(): - transfer.add_block_region( - tensor, semantic_role=f"mha.{geometry}.{role}" - ) - return transfer def refresh_flydsl_plan(self, context_lens, *, create=False): """Build or refresh aiter #5546's work plan for this forward. diff --git a/atom/model_ops/attentions/aiter_mla.py b/atom/model_ops/attentions/aiter_mla.py index 338fdb7124..1c7d9ce960 100644 --- a/atom/model_ops/attentions/aiter_mla.py +++ b/atom/model_ops/attentions/aiter_mla.py @@ -1395,11 +1395,13 @@ def build_kv_cache_tensor(self, module): ) def get_kv_transfer_tensors(self): + from atom.kv_transfer.disaggregation.page_region import page_region from atom.kv_transfer.disaggregation.types import ( INDEX_CACHE_ROLE, MLA_KV_ROLE, KVTransferRegion, KVTransferTensors, + PageRegion, ) runner = self.model_runner @@ -1447,19 +1449,20 @@ def get_kv_transfer_tensors(self): # and no per-field override to keep in step with the pooling ones. # DCP PD still dispatches on the two collapsed roles (`mla.kv` / # `dsa.index_cache`) rather than the pool's per-layer names. - pages = KVTransferTensors(block_regions=[], slot_regions=[]) + pages: list[PageRegion] = [] for role, t in self.kv_pool.region_tensors(): if role.startswith("index."): - pages.add_block_region(t, semantic_role=INDEX_CACHE_ROLE) + pages.append(page_region(t, semantic_role=INDEX_CACHE_ROLE)) index_tensors.append(t) continue - pages.add_block_region( - t, - semantic_role=( - MLA_KV_ROLE if role.startswith("kv.") else f"mla.{role}" - ), + pages.append( + page_region( + t, + semantic_role=( + MLA_KV_ROLE if role.startswith("kv.") else f"mla.{role}" + ), + ) ) - block_regions = pages.block_regions block_region_consumer_indices = None index_cache_layer_ids = getattr(runner, "index_cache_layer_ids", ()) @@ -1622,9 +1625,7 @@ def gather_sharded_index( return slot.data_ptr(), pages return KVTransferTensors( - block_regions=block_regions, - slot_regions=[], - block_tensor_views=pages.block_tensor_views, + pages=pages, block_region_consumer_indices=block_region_consumer_indices, index_staging_region=index_staging_region, index_staging_pool_size=index_staging_pool_size, diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index bfe7ed6fbc..2e757f236f 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -1890,9 +1890,11 @@ def get_kv_transfer_tensors(self): ``staging_region`` and ``gather_slot`` describe only compressor-state PD staging and must never be used as the source of a SLOT sidecar. """ + from atom.kv_transfer.disaggregation.page_region import page_region from atom.kv_transfer.disaggregation.types import ( KVTransferRegion, KVTransferTensors, + PageRegion, ) runner = self.model_runner @@ -1922,8 +1924,7 @@ def get_kv_transfer_tensors(self): geo = self.pool_geometry elem_fp32 = 4 - # PAGE regions and their LMCache MP byte views, published as pairs. - pages = KVTransferTensors(block_regions=[], slot_regions=[]) + pages: list[PageRegion] = [] swa_block_regions: list[KVTransferRegion] = [] slot_regions: list[KVTransferRegion] = [] @@ -1951,11 +1952,13 @@ def get_kv_transfer_tensors(self): for plane, row_bytes, role in planes: # A plane also holds SLOT rows after its PAGE blocks; publish only # the blocks. - pages.add_block_region( - plane, - semantic_role=role, - unit_bytes=geo.block_bytes(row_bytes), - total_bytes=self.num_blocks * geo.block_bytes(row_bytes), + pages.append( + page_region( + plane, + semantic_role=role, + unit_bytes=geo.block_bytes(row_bytes), + total_bytes=self.num_blocks * geo.block_bytes(row_bytes), + ) ) # Compressed PAGE regions: one per CSA layer in each indexer pool. @@ -1964,15 +1967,16 @@ def get_kv_transfer_tensors(self): # this remains correct for both the FP8 row layout and FP4 tiles. for pool, role_prefix in self._indexer_page_pools(): for pos, layer_id in enumerate(self.csa_layers): - pages.add_block_region( - pool[pos], semantic_role=f"{role_prefix}.layer_{layer_id}" + pages.append( + page_region( + pool[pos], semantic_role=f"{role_prefix}.layer_{layer_id}" + ) ) - block_regions = pages.block_regions checkpoint_spec = runner.state_runtime.checkpoint_spec if checkpoint_spec is None: raise RuntimeError("DSV4 PAGE/state checkpoint sizing spec is missing") - transfer_page_bytes = sum(region.unit_bytes for region in block_regions) + transfer_page_bytes = sum(page.region.unit_bytes for page in pages) if transfer_page_bytes != checkpoint_spec.page_unit_bytes: raise RuntimeError( "DSV4 PAGE transfer regions do not cover the sized PAGE unit: " @@ -2034,7 +2038,7 @@ def get_kv_transfer_tensors(self): # TP rank. One writer is sufficient; all ranks still read. tp_size = int(getattr(runner.config, "tensor_parallel_size", 1) or 1) return KVTransferTensors( - block_regions=block_regions, + pages=pages, swa_block_regions=swa_block_regions, slot_regions=slot_regions, num_slots=num_slots, @@ -2043,12 +2047,11 @@ def get_kv_transfer_tensors(self): staging_pool_size=pool_size if staging_region else 0, gather_slot=gather_slot, scatter_slot=scatter_slot, - block_tensor_views=pages.block_tensor_views, tp_replication_factor=tp_size, native_state_tp_replication_factor=tp_size, paged_state_checkpoint_spec=checkpoint_spec, execute_paged_state_copies=self.execute_paged_state_copies, - paged_state_region_count=len(block_regions), + paged_state_region_count=len(pages), ) # ------------------------------------------------------------------ # diff --git a/atom/model_ops/attentions/qwen4_exp_attn.py b/atom/model_ops/attentions/qwen4_exp_attn.py index a583f0f2c1..7196ec5509 100644 --- a/atom/model_ops/attentions/qwen4_exp_attn.py +++ b/atom/model_ops/attentions/qwen4_exp_attn.py @@ -323,6 +323,7 @@ def get_kv_transfer_tensors(self): from atom.kv_transfer.disaggregation.types import ( KVTransferRegion, KVTransferTensors, + PageRegion, ) arena = getattr(self, "qsa_arena", None) @@ -355,7 +356,10 @@ def region(tensor, role): ple = getattr(runner, name, None) if ple is not None: slots.append(region(ple, role)) - return KVTransferTensors(block_regions=blocks, slot_regions=slots) + # Addresses only: no byte views, so LMCache MP refuses this layout. + return KVTransferTensors( + pages=[PageRegion(block) for block in blocks], slot_regions=slots + ) def allocate_per_req_cache(self, entries: dict[str, int]) -> dict[str, object]: caches = super().allocate_per_req_cache(entries) diff --git a/atom/spec_decode/draft_kv.py b/atom/spec_decode/draft_kv.py index 2c92ec8935..e9e5eb0578 100644 --- a/atom/spec_decode/draft_kv.py +++ b/atom/spec_decode/draft_kv.py @@ -188,19 +188,17 @@ def build_kv_cache_tensor(self, module): ) def get_kv_transfer_tensors(self): - from atom.kv_transfer.disaggregation.types import ( - KVTransferTensors, - ) - - transfer = KVTransferTensors( - block_regions=[], - slot_regions=[], + from atom.kv_transfer.disaggregation.page_region import page_region + from atom.kv_transfer.disaggregation.types import KVTransferTensors + + return KVTransferTensors( + # The draft's rows sit in the same block ids as the target's, so + # its regions need a name that says which stack they are. + pages=[ + page_region(t, semantic_role=f"draft.{role}") + for role, t in self.kv_pool.region_tensors() + ], # A paged MHA/GQA draft is a per-rank shard. Appending it to an # otherwise replicated target must disable whole-object collapse. tp_replication_factor=1, ) - # The draft's rows sit in the same block ids as the target's, so its - # regions need a name that says which stack they are. - for role, t in self.kv_pool.region_tensors(): - transfer.add_block_region(t, semantic_role=f"draft.{role}") - return transfer diff --git a/tests/test_deepseek_v4_transfer_regions.py b/tests/test_deepseek_v4_transfer_regions.py index 638fb807f6..76b97fa4d6 100644 --- a/tests/test_deepseek_v4_transfer_regions.py +++ b/tests/test_deepseek_v4_transfer_regions.py @@ -18,6 +18,7 @@ INDEX_CACHE_FP4_PREFIX, KVTransferRegion, KVTransferTensors, + PageRegion, ) from atom.kv_transfer.offload.hybrid.dsv4.codec import DSV4PageSlotCodec from atom.model_engine.kv_block import STATE_SLOT_CLASS @@ -42,7 +43,7 @@ def test_generic_transfer_tensors_default_to_no_full_slot_expectation(): - transfer = KVTransferTensors(block_regions=[], slot_regions=[]) + transfer = KVTransferTensors() assert transfer.expected_full_slot_region_count is None @@ -70,11 +71,10 @@ class TestTheBlockIdSpacePageRegionsAreAddressedIn: def test_regions_in_the_scheduler_block_are_accepted(self): transfer = KVTransferTensors( - block_regions=[ - _page_region(8, 2080, "kv.layer_0"), - _page_region(8, 132, "index.layer_0"), - ], - slot_regions=[], + pages=[ + PageRegion(_page_region(8, 2080, "kv.layer_0")), + PageRegion(_page_region(8, 132, "index.layer_0")), + ] ) transfer.set_block_count(8) @@ -87,11 +87,10 @@ def test_a_region_registered_in_another_unit_is_refused(self): `1/block_ratio` the bytes -- the same total, which is why no byte count catches it.""" transfer = KVTransferTensors( - block_regions=[ - _page_region(8, 2080, "kv.layer_0"), - _page_region(8 * 16, 2080 // 16, "kv.layer_1"), - ], - slot_regions=[], + pages=[ + PageRegion(_page_region(8, 2080, "kv.layer_0")), + PageRegion(_page_region(8 * 16, 2080 // 16, "kv.layer_1")), + ] ) with pytest.raises(ValueError, match="kv.layer_1 holds 128 blocks"): @@ -99,8 +98,9 @@ def test_a_region_registered_in_another_unit_is_refused(self): def test_a_region_that_does_not_divide_into_whole_blocks_is_refused(self): transfer = KVTransferTensors( - block_regions=[KVTransferRegion(0x1000, 100, 32, semantic_role="ragged")], - slot_regions=[], + pages=[ + PageRegion(KVTransferRegion(0x1000, 100, 32, semantic_role="ragged")) + ] ) with pytest.raises(ValueError, match="does not divide into whole blocks"): @@ -110,7 +110,7 @@ def test_the_count_is_not_something_a_backend_can_pass(self): """`init=False` is the enforcement: a backend cannot state a count that its own regions contradict, because it cannot state one at all.""" with pytest.raises(TypeError): - KVTransferTensors(block_regions=[], slot_regions=[], num_blocks=8) + KVTransferTensors(num_blocks=8) @contextmanager @@ -789,8 +789,8 @@ def test_fp4_pd_copies_data_scales_and_swa_without_touching_other_blocks( owners, src, dst, conn, req = _fp4_pd_pair( v4_builder_cls, mooncake_page_writer, kv_dtype, coalesce ) - src_regions = src.block_regions + src.swa_block_regions - dst_regions = dst.block_regions + dst.swa_block_regions + src_regions = [*src.block_regions, *src.swa_block_regions] + dst_regions = [*dst.block_regions, *dst.swa_block_regions] expected = [] for n, (sr, dr) in enumerate(zip(src_regions, dst_regions, strict=True)): # Different layers and byte offsets expose missing scales, swapped diff --git a/tests/test_lmcache_mp.py b/tests/test_lmcache_mp.py index ab1d4d2379..7c2ea89be7 100644 --- a/tests/test_lmcache_mp.py +++ b/tests/test_lmcache_mp.py @@ -14,11 +14,13 @@ from atom.kv_transfer.disaggregation.aggregator import KVOutputAggregator from atom.kv_transfer.disaggregation.factory import KVConnectorFactory +from atom.kv_transfer.disaggregation.page_region import page_region from atom.kv_transfer.disaggregation.types import ( ConnectorCompletion, KVTransferRegion, KVTransferTensors, LoadOperationId, + PageRegion, SaveOperationId, SaveSourceGroupId, ) @@ -472,9 +474,10 @@ def _transfer_tensors(*, tp_replication_factor: int = 1) -> KVTransferTensors: for tensor, role in zip(tensors, roles, strict=True) ] transfer = KVTransferTensors( - block_regions=regions, - slot_regions=[], - block_tensor_views=tensors, + pages=[ + PageRegion(region, tensor) + for region, tensor in zip(regions, tensors, strict=True) + ], tp_replication_factor=tp_replication_factor, ) transfer.set_block_count(2) @@ -516,11 +519,7 @@ def test_build_cache_views_publishes_float8_pages_as_raw_bytes(): unit_bytes=page[0].numel(), semantic_role="latent", ) - transfer = KVTransferTensors( - block_regions=[region], - slot_regions=[], - block_tensor_views=[page], - ) + transfer = KVTransferTensors(pages=[PageRegion(region, page)]) transfer.set_block_count(2) views = page_views._build_cache_views(transfer, num_blocks=2) @@ -534,7 +533,7 @@ def test_build_cache_views_publishes_float8_pages_as_raw_bytes(): def test_build_cache_views_rejects_missing_or_bad_geometry(): missing_view = _transfer_tensors() - missing_view.block_tensor_views.pop() + missing_view.pages[-1] = PageRegion(missing_view.pages[-1].region) with pytest.raises(ValueError, match="one block_tensor_view per block region"): page_views._build_cache_views(missing_view, num_blocks=2) @@ -544,14 +543,16 @@ def test_build_cache_views_rejects_missing_or_bad_geometry(): page_views._build_cache_views(bad_geometry, num_blocks=2) noncontiguous = _transfer_tensors() - noncontiguous.block_tensor_views[0] = torch.zeros(2, 32, 4).transpose(1, 2) + noncontiguous.pages[0] = replace( + noncontiguous.pages[0], view=torch.zeros(2, 32, 4).transpose(1, 2) + ) assert not noncontiguous.block_tensor_views[0].is_contiguous() with pytest.raises(ValueError, match="non-empty and contiguous"): page_views._build_cache_views(noncontiguous, num_blocks=2) unsupported_rank = _transfer_tensors() - unsupported_rank.block_tensor_views[0] = torch.zeros( - 2, 4, 4, 8, dtype=torch.float16 + unsupported_rank.pages[0] = replace( + unsupported_rank.pages[0], view=torch.zeros(2, 4, 4, 8, dtype=torch.float16) ) with pytest.raises(ValueError, match="physical_slots, opaque_width"): page_views._build_cache_views(unsupported_rank, num_blocks=2) @@ -1864,18 +1865,19 @@ def test_factory_registers_lmcache_mp_alias_without_pd_staging(): ) -def test_add_block_region_publishes_a_region_and_its_byte_view_together(): +def test_page_region_is_a_region_and_its_byte_view_built_together(): """One call yields both halves, zero-copy, over exactly the PAGE bytes: a larger allocation (DSV4's planes also hold SLOT rows) is cut to them.""" arena = torch.arange(3 * 8 + 4, dtype=torch.int16) # 3 blocks + a tail - transfer = KVTransferTensors(block_regions=[], slot_regions=[]) - transfer.add_block_region( - arena, semantic_role="plane", unit_bytes=16, total_bytes=3 * 16 - ) fp16 = torch.zeros(3, 4, 2, dtype=torch.float16) - transfer.add_block_region(fp16, semantic_role="rows") + transfer = KVTransferTensors( + pages=[ + page_region(arena, semantic_role="plane", unit_bytes=16, total_bytes=48), + page_region(fp16, semantic_role="rows"), + ] + ) - [plane, rows] = transfer.block_regions + plane, rows = transfer.block_regions assert (plane.base_addr, plane.total_bytes, plane.unit_bytes) == ( arena.data_ptr(), 48, @@ -1894,15 +1896,42 @@ def test_add_block_region_publishes_a_region_and_its_byte_view_together(): assert page_views._build_cache_views(transfer, num_blocks=3).bytes_per_block == 32 -def test_add_block_region_rejects_what_it_cannot_alias(): - transfer = KVTransferTensors(block_regions=[], slot_regions=[]) +def test_page_region_rejects_what_it_cannot_alias(): with pytest.raises(ValueError, match="contiguous"): - transfer.add_block_region(torch.zeros(4, 3).t(), semantic_role="strided") + page_region(torch.zeros(4, 3).t(), semantic_role="strided") with pytest.raises(ValueError, match="cannot publish"): - transfer.add_block_region( + page_region( torch.zeros(8, dtype=torch.uint8), semantic_role="short", unit_bytes=4, total_bytes=12, ) - assert transfer.block_regions == [] and transfer.block_tensor_views == [] + + +def test_pages_are_read_only_through_the_derived_lists(): + """`block_regions` / `block_tensor_views` are views of `pages`, so nothing + can append to one without the other.""" + transfer = KVTransferTensors( + pages=[page_region(torch.zeros(2, 4, dtype=torch.uint8), semantic_role="p")] + ) + with pytest.raises(AttributeError): + transfer.block_regions.append(transfer.block_regions[0]) + with pytest.raises(AttributeError): + transfer.block_tensor_views.append(transfer.block_tensor_views[0]) + + +def test_merge_pages_appends_a_draft_and_takes_the_gcd_replication(): + target = KVTransferTensors( + pages=[page_region(torch.zeros(2, 4, dtype=torch.uint8), semantic_role="t")], + tp_replication_factor=8, + ) + target.merge_pages(KVTransferTensors(tp_replication_factor=1)) + assert target.tp_replication_factor == 8 # nothing merged, nothing changes + draft = KVTransferTensors( + pages=[page_region(torch.zeros(2, 2, dtype=torch.uint8), semantic_role="d")], + tp_replication_factor=1, + ) + target.merge_pages(draft) + assert [r.semantic_role for r in target.block_regions] == ["t", "d"] + assert len(target.block_tensor_views) == 2 + assert target.tp_replication_factor == 1 diff --git a/tests/test_lmcache_mp_native_layout.py b/tests/test_lmcache_mp_native_layout.py index 9bc5c74351..38b8fa8add 100644 --- a/tests/test_lmcache_mp_native_layout.py +++ b/tests/test_lmcache_mp_native_layout.py @@ -13,7 +13,11 @@ import pytest import torch -from atom.kv_transfer.disaggregation.types import KVTransferRegion, KVTransferTensors +from atom.kv_transfer.disaggregation.types import ( + KVTransferRegion, + KVTransferTensors, + PageRegion, +) from atom.kv_transfer.offload.mp.native_state_layout import ( build_native_state_mp_layout, ) @@ -35,17 +39,18 @@ def _transfer(*, widths=(8, 8, 2), image_bytes=39, num_blocks=7): layout_id="native-test-v1", ) transfer = KVTransferTensors( - block_regions=[ - KVTransferRegion( - view.data_ptr(), - num_blocks * width, - width, - semantic_role=f"region.{i}", + pages=[ + PageRegion( + KVTransferRegion( + view.data_ptr(), + num_blocks * width, + width, + semantic_role=f"region.{i}", + ), + view, ) for i, (view, width) in enumerate(zip(views, widths, strict=True)) ], - slot_regions=[], - block_tensor_views=views, paged_state_checkpoint_spec=spec, execute_paged_state_copies=lambda stores, restores, descriptor_slot=0: None, ) @@ -210,7 +215,7 @@ def test_registration_rejects_mismatched_native_geometry_and_missing_restore(): with pytest.raises(ValueError, match="do not cover"): _layout(transfer) transfer = _transfer() - transfer.block_tensor_views[0] = transfer.block_tensor_views[0].clone() + transfer.pages[0] = replace(transfer.pages[0], view=transfer.pages[0].view.clone()) with pytest.raises(ValueError, match="does not alias"): _layout(transfer) @@ -380,7 +385,10 @@ def test_draft_regions_after_the_state_regions_stay_ordinary_page(): ("breakage", "message"), [ (lambda t: setattr(t.block_regions[1], "unit_bytes", 9), "byte geometry"), - (lambda t: t.block_tensor_views.pop(), "one block_tensor_view per"), + ( + lambda t: t.pages.__setitem__(-1, PageRegion(t.pages[-1].region)), + "one block_tensor_view per", + ), (lambda t: setattr(t.block_regions[2], "reverse_indexed", True), "reverse"), ], ) diff --git a/tests/test_lmcache_mp_native_worker.py b/tests/test_lmcache_mp_native_worker.py index a701b98aac..fcb0609eb4 100644 --- a/tests/test_lmcache_mp_native_worker.py +++ b/tests/test_lmcache_mp_native_worker.py @@ -16,6 +16,7 @@ KVTransferRegion, KVTransferTensors, LoadOperationId, + PageRegion, SaveOperationId, SaveSourceGroupId, ) @@ -98,13 +99,14 @@ def worker(): page = torch.zeros((32, 1, 32), dtype=torch.uint8) spec = PagedStateCheckpointSpec(32, 128, "native-test-v1", 80) tensors = KVTransferTensors( - block_regions=[ - KVTransferRegion( - base_addr=page.data_ptr(), unit_bytes=32, total_bytes=page.numel() + pages=[ + PageRegion( + KVTransferRegion( + base_addr=page.data_ptr(), unit_bytes=32, total_bytes=page.numel() + ), + page, ) ], - slot_regions=[], - block_tensor_views=[page], paged_state_checkpoint_spec=spec, execute_paged_state_copies=lambda stores, restores, descriptor_slot=0: None, ) @@ -509,13 +511,14 @@ def test_registration_reserves_every_restore_descriptor_slot(monkeypatch): reserved = [] page = torch.zeros((32, 1, 32), dtype=torch.uint8) tensors = KVTransferTensors( - block_regions=[ - KVTransferRegion( - base_addr=page.data_ptr(), unit_bytes=32, total_bytes=page.numel() + pages=[ + PageRegion( + KVTransferRegion( + base_addr=page.data_ptr(), unit_bytes=32, total_bytes=page.numel() + ), + page, ) ], - slot_regions=[], - block_tensor_views=[page], paged_state_checkpoint_spec=PagedStateCheckpointSpec( 32, 128, "native-test-v1", 80 ),