diff --git a/atom/kv_transfer/disaggregation/multi/multi_connector.py b/atom/kv_transfer/disaggregation/multi/multi_connector.py index dc892cc502..f228cbcd63 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( @@ -489,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 @@ -646,6 +670,25 @@ def protected_block_ids(self, seq: Any): protected.update(blocks) return frozenset(protected) + def can_partially_deallocate_state(self, seq: Any) -> bool: + """True only if every sub still deferring this request guarantees it. + + 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 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: callback = getattr(connector, "activate_block_leases", None) 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 d1a1a5ec68..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]; 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. - 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 @@ -257,17 +270,62 @@ 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 + # 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 # the same number. Set through `set_block_count`. num_blocks: int = field(init=False, default=0) + @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. + """ + if not other.pages: + return + self.pages.extend(other.pages) + self.tp_replication_factor = math.gcd( + self.tp_replication_factor, other.tp_replication_factor + ) + 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 65231813b5..73affd5e08 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,197 @@ 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/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: `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 \ + --eviction-policy LRU +``` + +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. +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. 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. 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 +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. 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. + +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: @@ -141,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 @@ -848,8 +1042,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..dba285ab7c 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) @@ -517,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.""" @@ -643,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 @@ -667,16 +659,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 +728,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 +1039,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 +1056,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..4a2a93982d 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,17 @@ 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 + + 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 +180,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 +210,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 +231,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 +404,88 @@ 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 _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 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 + ) + 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 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 + # 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): + # retire the request rather than emit an empty save forever. + 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 +542,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 +566,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 +595,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 +608,33 @@ 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_released", False) and ( + self._reacquires_late_source() + ): + 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( + 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 + ) + 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 +654,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) @@ -597,19 +694,21 @@ 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 - full block table and computed frontier so the next metadata build can - still dispatch that save after unrelated blocks have been released. + 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 - 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,20 +719,29 @@ 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]) + 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)) + 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: """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()) @@ -713,7 +821,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 +1031,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 +1040,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. + # 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 @@ -969,9 +1077,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: @@ -1108,8 +1214,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/connector.py b/atom/kv_transfer/offload/connector.py index 77060978f6..9b496e4aaa 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: @@ -327,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/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 deleted file mode 100644 index 12e309f342..0000000000 --- a/atom/kv_transfer/offload/mp/backend.py +++ /dev/null @@ -1,1068 +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 logging -import os -import threading -import time -from collections import deque -from dataclasses import dataclass, replace -from typing import Any - -import torch - -from atom.kv_transfer.disaggregation.base import KVConnectorBase -from atom.kv_transfer.disaggregation.types import ( - KVConnectorOutput, - LoadCompletionId, - LoadOperationId, - SaveCompletionId, -) -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, LMCacheReqMeta - -logger = logging.getLogger("atom") - -_MP_LAYOUT_VERSION = 2 -_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 _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") - 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, - ) - 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") - - _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. - """ - - hf_config = getattr(config, "hf_config", None) - hf_config = getattr(hf_config, "text_config", hf_config) - return getattr(hf_config, "kv_lora_rank", None) is not None - - -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, -) -> int: - """Validate the attention backend's whole-PAGE replication declaration.""" - - factor = offcfg._strict_integer( - "KV PAGE TP replication factor", - getattr(transfer_tensors, "tp_replication_factor", 1), - minimum=1, - ) - if tp_size % factor: - raise ValueError( - f"KV PAGE 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}" - ) - 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) -> 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, - ) - return f"{page_namespace}::lmcache-mp-v{_MP_LAYOUT_VERSION}" - - -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) -> 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), - 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: - 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), - 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") - if num_blocks <= 0: - raise ValueError("lmcache_mp num_blocks must be positive") - - 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)}" - ) - - 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") - - 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:])) - indices_by_layout.setdefault(layout, []).append(index) - devices.add(view.device) - bytes_per_block += 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()), - 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 - - -@dataclass -class _PendingSave: - completion: SaveCompletionId - future: Any | None - - -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 rather than risk block reuse. - 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()) - - -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: - if timeout <= 0 or poll_interval <= 0: - raise ValueError("LMCache MP lookup timeout and poll interval must be > 0") - self._adapter = adapter - 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 - self._adapter.maybe_submit_lookup_request(lookup_id, token_ids) - deadline = time.monotonic() + self._timeout - while True: - result = self._adapter.check_lookup_result(lookup_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) - 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=lookup_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=lookup_id, - ) - if state is not None: - state.retrieve_start = start - state.retrieve_end = end - self._adapter.cleanup_lookup_result(lookup_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(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) - if state is not None and state.hit is None: - result = self._adapter.check_lookup_result(lookup_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(lookup_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=lookup_id, - ) - self._adapter.cleanup_lookup_result(lookup_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() - self._immediate_saves: set[SaveCompletionId] = set() - 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 = str(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, - ) - transfer = self._adapter.submit_retrieve_request( - request_id, - op, - event, - ) - except Exception: - logger.exception( - "LMCache MP load submission failed 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 - 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 = str(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) - 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.add(completion) - 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) - _remember_operation_tombstone( - operation_id, - self._completed_save_operations, - self._completed_save_operation_order, - ) - self._immediate_saves.add(completion) - 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, - ) - transfer = self._adapter.submit_store_request( - request_id, - op, - event, - ) - except Exception: - logger.exception( - "LMCache MP save submission failed 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._immediate_saves.add(completion) - return - with self._lock: - self._submitting_saves.discard(operation_id) - self._pending_saves[operation_id] = _PendingSave( - completion=completion, - future=transfer, - ) - - 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() - 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()): - terminal, _result = _terminal_future_result(pending.future) - 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) - - for operation_id, pending in list(self._pending_loads.items()): - terminal, result = _terminal_future_result(pending.future) - if not terminal: - 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) - 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, - ) - - -class LMCacheMPConnectorScheduler(ChunkedOffloadSchedulerBase): - """Scheduler-side LMCache MP connector for generic PAGE offload.""" - - def __init__(self, config: Any) -> 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) - 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, - timeout=timeout, - poll_interval=poll_interval, - ) - self._mp_adapter = adapter - 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), - ) - 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 request_finished(self, seq: Any) -> None: - super().request_finished(seq) - try: - self._mp_adapter.end_session(str(seq.id)) - except Exception: - logger.warning( - "LMCache MP end_session failed for request %s", - seq.id, - exc_info=True, - ) - - -__all__ = [ - "LMCacheMPConnector", - "LMCacheMPConnectorScheduler", -] diff --git a/atom/kv_transfer/offload/mp/connector.py b/atom/kv_transfer/offload/mp/connector.py index d417b70e45..d918856448 100644 --- a/atom/kv_transfer/offload/mp/connector.py +++ b/atom/kv_transfer/offload/mp/connector.py @@ -1,16 +1,201 @@ # 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 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: + """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: + deployment._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 = mp_worker.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: + deployment._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 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 = mp_scheduler.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/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_layout.py b/atom/kv_transfer/offload/mp/native_state_layout.py new file mode 100644 index 0000000000..9f78ae95ed --- /dev/null +++ b/atom/kv_transfer/offload/mp/native_state_layout.py @@ -0,0 +1,208 @@ +# 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 + +from atom.kv_transfer.offload.mp.page_views import validate_page_views + + +@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") + + # 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) + 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") + 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") + + # 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[:state_count], regions[:state_count], 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..a72efc4bb9 --- /dev/null +++ b/atom/kv_transfer/offload/mp/native_state_scheduler.py @@ -0,0 +1,567 @@ +# 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 + +import logging +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.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 + +logger = logging.getLogger("atom") + +_MAX_SAVE_ATTEMPTS = 3 + +_T = TypeVar("_T") + + +@dataclass +class _NativeSave: + seq: Any + source: StateStoreOperationId + saved_before: int + boundary: int + + +@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) + # 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 " + "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. 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 + 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 + ) + floor = self._save_floor(exhausted) + # Every tracked request asks every step. The answer changes only when + # 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[: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): + 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 = (frontier, floor, store.generation, found) + return found + + 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, self._save_floor(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 + 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( + 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),)) + ) + # 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) + 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_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 + prefix_hash = self._boundary_hash(seq, lmc) + 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 = 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 + 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: + 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 + 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 _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 new file mode 100644 index 0000000000..b97c17bdc5 --- /dev/null +++ b/atom/kv_transfer/offload/mp/native_state_worker.py @@ -0,0 +1,473 @@ +# 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 +import time +from dataclasses import dataclass, field +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.deployment import ( + _make_worker_adapter, + _mp_session_id, + _published_tp_replication_factor, + _storage_kv_transfer_config, + _tp_replication_factor, + _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 + +logger = logging.getLogger("atom") +NATIVE_STATE_MP_STORE_CHANNEL = "native_state_mp_store" + + +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}" + ) + + +@dataclass +class _NativePending: + request: LMCacheReqMeta + future: Any + restore_event: Any = None + restore_succeeded: bool = False + descriptor_slot: int | None = None + immediate_success: bool = False + started_at: float = field(default_factory=lambda: time.monotonic()) + + +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) + ) + # 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", + 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: + # 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) + 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 + # 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 + 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 unprovable; 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 + # 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( + (), + ( + 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_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 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 + # `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: + for operation_id, pending in list(self._native_saves.items()): + self._poll_native_save(output, operation_id, pending) + for operation_id, pending in list(self._native_loads.items()): + 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 + ) + 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 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__ = [ + "NATIVE_STATE_MP_STORE_CHANNEL", + "NativeStateLMCacheMPConnector", + "require_native_state_server", +] 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..4e84a4db30 --- /dev/null +++ b/atom/kv_transfer/offload/mp/page_views.py @@ -0,0 +1,158 @@ +# 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 + + +@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..7da1dc651f --- /dev/null +++ b/atom/kv_transfer/offload/mp/scheduler.py @@ -0,0 +1,245 @@ +# 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 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: + # 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 (deadline %.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/atom/model_engine/block_manager.py b/atom/model_engine/block_manager.py index caa9504d47..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, @@ -2499,7 +2506,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 +2528,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 +2548,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 = seq.cache_seed + 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..34cd360faf 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -2103,9 +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_regions = self.draft_kv_builder.get_kv_transfer_tensors() - if draft_regions: - transfer_tensors.block_regions.extend(draft_regions) + 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_engine/page_unit_checkpoint.py b/atom/model_engine/page_unit_checkpoint.py index fb11d6ad70..ceb9356d95 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 @@ -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 @@ -116,6 +118,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 +144,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: @@ -153,6 +166,12 @@ 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. 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() @@ -382,6 +401,89 @@ 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 _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: + """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 + self._note_change(record.prefix_hash) + return True + def complete_inflight(self) -> None: stores, self._inflight_stores = self._inflight_stores, [] for checkpoint_id in stores: @@ -401,6 +503,7 @@ def complete_inflight(self) -> None: continue record.state = READY self.hash_to_checkpoint[record.prefix_hash] = checkpoint_id + self._note_change(record.prefix_hash) self._lru[checkpoint_id] = None self._queue_offload_store(checkpoint_id, record) @@ -491,13 +594,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 +607,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 +742,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) @@ -678,6 +820,7 @@ def _release_restore_pin(self, checkpoint_id: int) -> None: self._release_record(checkpoint_id) def unindex(self, prefix_hash: int) -> bool: + 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) @@ -694,6 +837,7 @@ def unindex(self, prefix_hash: int) -> bool: return True def clear(self) -> None: + self._note_change(None) self.hash_to_checkpoint.clear() self._pending_by_hash.clear() self._lru.clear() @@ -708,6 +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._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 @@ -717,6 +862,7 @@ def _evict(self, checkpoint_id: int) -> None: def _release_record(self, checkpoint_id: int) -> None: 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] @@ -773,6 +919,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 +1070,62 @@ 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.get(owner) + if units is None: + return False + 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 + ) -> 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 +1158,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..0968cb0669 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. @@ -1152,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. @@ -1300,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 @@ -2655,19 +2678,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 +3453,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..d494a90adb 100644 --- a/atom/model_ops/attentions/aiter_attention.py +++ b/atom/model_ops/attentions/aiter_attention.py @@ -773,29 +773,25 @@ 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.page_region import page_region + from atom.kv_transfer.disaggregation.types import KVTransferTensors if not self.kv_pools: return None - # 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}", - ) + # 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() ], - slot_regions=[], + # 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/aiter_mla.py b/atom/model_ops/attentions/aiter_mla.py index 94656f0ac2..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,32 +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. - 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: list[PageRegion] = [] + 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.append(page_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, + pages.append( + page_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_region_consumer_indices = None index_cache_layer_ids = getattr(runner, "index_cache_layer_ids", ()) @@ -1635,9 +1625,7 @@ def gather_sharded_index( return slot.data_ptr(), pages return KVTransferTensors( - block_regions=block_regions, - slot_regions=[], - block_tensor_views=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/backends.py b/atom/model_ops/attentions/backends.py index 220977426d..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 ( @@ -342,13 +343,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 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. + """ + self._checkpoint_staging().reserve(descriptor_slots) + + 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_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 + `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, + separately fenced buffer (see `DescriptorStaging`). + """ + 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) + staging = self._checkpoint_staging_cache = DescriptorStaging( + max_ops * plan.num_spans, self._checkpoint_descriptor_device() + ) + 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 32ac5e42de..18ca129c4e 100644 --- a/atom/model_ops/attentions/deepseek_v41/backend.py +++ b/atom/model_ops/attentions/deepseek_v41/backend.py @@ -260,8 +260,11 @@ 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 reserve_checkpoint_descriptors(self, descriptor_slots): + 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 cdc18c63df..8723f7ca6b 100644 --- a/atom/model_ops/attentions/deepseek_v41/checkpoints.py +++ b/atom/model_ops/attentions/deepseek_v41/checkpoints.py @@ -5,10 +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 StateCopies: @@ -33,15 +33,10 @@ 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 = DescriptorStaging( + 2 * max_batch * self.plan.num_spans, cache.pool.device ) - self.upload_done = torch.cuda.Event() if cache.pool.is_cuda else None - self.pending = False + self.staging.reserve([0]) def entry(self, slot): self.cache.require_committed() @@ -73,7 +68,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 +90,10 @@ def execute(self, stores, restores): ) at += take return + rows = self.staging.rows(descriptor_slot) total = (len(stores) + len(restores)) * self.plan.num_spans - if total > self.staging.np.shape[0]: + if total > rows.shape[0]: raise ValueError("Checkpoint copy batch exceeds descriptor capacity") - if self.pending: - self.upload_done.synchronize() at = 0 for ops, storing in ((stores, True), (restores, False)): if not ops: @@ -124,12 +118,10 @@ def execute(self, stores, restores): dtype=np.int64, ) self.plan.write_descriptor( - self.staging.np[at:end], slot_bases, page_bases, forward=storing + rows[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 = 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) @@ -148,11 +140,8 @@ def warmup(self): at += size dst = np.asarray([[slot.data_ptr() + off for off in offsets]], dtype=np.int64) self.plan.write_descriptor( - self.staging.np[: self.plan.num_spans], + self.staging.rows(0)[: 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 - 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 fb5214d4f2..2e757f236f 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_descriptor: CpuGpuBuffer | None = None + self._checkpoint_staging_cache: DescriptorStaging | None = None @property def prep_stream(self): @@ -969,6 +970,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,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() - 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: @@ -1008,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 @@ -1159,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_descriptor = None + 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. @@ -1182,42 +1185,16 @@ 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_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. @@ -1913,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 @@ -1945,7 +1924,7 @@ def get_kv_transfer_tensors(self): geo = self.pool_geometry elem_fp32 = 4 - block_regions: list[KVTransferRegion] = [] + pages: list[PageRegion] = [] swa_block_regions: list[KVTransferRegion] = [] slot_regions: list[KVTransferRegion] = [] @@ -1971,12 +1950,14 @@ def get_kv_transfer_tensors(self): ) ) for plane, row_bytes, role in planes: - block_regions.append( - KVTransferRegion( - plane.data_ptr(), - self.num_blocks * geo.block_bytes(row_bytes), - geo.block_bytes(row_bytes), + # A plane also holds SLOT rows after its PAGE blocks; publish only + # the blocks. + 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), ) ) @@ -1986,24 +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): - view = pool[pos] - if not view.is_contiguous(): - raise RuntimeError( - "a CSA indexer layer must be contiguous to be transferred" - ) - 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.append( + page_region( + pool[pos], semantic_role=f"{role_prefix}.layer_{layer_id}" ) ) 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: " @@ -2060,8 +2033,12 @@ 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, + pages=pages, swa_block_regions=swa_block_regions, slot_regions=slot_regions, num_slots=num_slots, @@ -2070,6 +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, + 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(pages), ) # ------------------------------------------------------------------ # diff --git a/atom/model_ops/attentions/gdn_attn.py b/atom/model_ops/attentions/gdn_attn.py index e4c53b52c3..e350d4acc9 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,14 +832,15 @@ 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() - 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: @@ -878,7 +854,7 @@ def execute_paged_state_copies(self, store_ops, restore_ops) -> None: 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. @@ -901,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/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/plugin/vllm/kv_transfer/seq_view.py b/atom/plugin/vllm/kv_transfer/seq_view.py index 187ab8e3d9..a9639d5746 100644 --- a/atom/plugin/vllm/kv_transfer/seq_view.py +++ b/atom/plugin/vllm/kv_transfer/seq_view.py @@ -20,15 +20,22 @@ 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. + # 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", + "_offload_released", "_request", "block_table", "offload_handoff_boundary_tokens", @@ -113,8 +120,10 @@ 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", "_offload_finished_block_ids", "_offload_finished_cached_tokens", + "_offload_released", ): if hasattr(self, frozen): delattr(self, frozen) diff --git a/atom/spec_decode/draft_kv.py b/atom/spec_decode/draft_kv.py index 96166a91dd..e9e5eb0578 100644 --- a/atom/spec_decode/draft_kv.py +++ b/atom/spec_decode/draft_kv.py @@ -187,17 +187,18 @@ 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.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, + ) diff --git a/atom/utils/envs.py b/atom/utils/envs.py index fda4d28faa..92ac89408f 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -45,6 +45,81 @@ 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(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 or raw == "": + 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: + """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: + 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. @@ -826,6 +901,74 @@ 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) --- + # 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_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. + "OFFLOAD_MAX_PENDING_SAVES": lambda: _optional_int_env("OFFLOAD_MAX_PENDING_SAVES"), + # 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) + ), + "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 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"), } @@ -866,9 +1009,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/docs/environment_variables.md b/docs/environment_variables.md index 2dc46774a5..c99c8caf29 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -288,21 +288,34 @@ flag below. Details in the state-checkpoint section of the ### LMCache offload tier +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 | |----------|------|---------|-------------| | **ATOM_KV_OFFLOAD** | str | "" (off) | Enables LMCache KV offload without `--kv-transfer-config`, so a launcher that owns that flag for P/D can still add offload. `lmcache` selects the in-process `lmcache_offload` connector, `lmcache_mp` the standalone-server `lmcache_mp` connector (start `lmcache server` first). With a P/D connector in `--kv-transfer-config`, both run behind a `multi` connector. Setting it while `--kv-transfer-config` already names an offload connector is an error. | | **ATOM_KV_OFFLOAD_EXTRA_CONFIG** | JSON object | "" | The offload connector's `kv_connector_extra_config`: `lmcache.` overrides (e.g. `{"lmcache.chunk_size": 256}`) and, for `lmcache_mp`, `lmcache.mp.*` options (e.g. `{"lmcache.mp.port": 5556}`). Requires `ATOM_KV_OFFLOAD`. | - -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. - -| 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/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_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_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_envs.py b/tests/test_envs.py index bd2ce7a524..76068c3963 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -277,3 +277,51 @@ 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] == [] + + +@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 diff --git a/tests/test_lmcache_mp.py b/tests/test_lmcache_mp.py index 28814e426b..7c2ea89be7 100644 --- a/tests/test_lmcache_mp.py +++ b/tests/test_lmcache_mp.py @@ -6,27 +6,39 @@ import sys import types from collections import deque -from dataclasses import dataclass +from dataclasses import dataclass, replace from types import SimpleNamespace 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.page_region import page_region from atom.kv_transfer.disaggregation.types import ( + ConnectorCompletion, KVTransferRegion, KVTransferTensors, LoadOperationId, + PageRegion, 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( @@ -37,7 +49,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 +78,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,32 +99,48 @@ 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"), ], ) 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, ) +@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 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, ) @@ -112,11 +148,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 "}) ) @@ -151,9 +187,11 @@ 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( + 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 @@ -172,7 +210,7 @@ def __init__(self, **kwargs): ) strategies = [ - mp_connector._parallel_strategy(_config(tp=8), 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} @@ -180,6 +218,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 = deployment._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): @@ -194,7 +257,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) ] @@ -203,9 +266,51 @@ 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 deployment._tp_replication_factor(config) == 8 + + config.speculative_config = SimpleNamespace( + method="dspark", draft_model_hf_config=SimpleNamespace() + ) + assert deployment._tp_replication_factor(config) == 1 + + owns_pool["value"] = False + assert deployment._tp_replication_factor(config) == 8 + config.speculative_config.method = "mtp" + owns_pool["value"] = True + assert deployment._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) + 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 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, @@ -215,7 +320,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}) ) @@ -248,9 +353,11 @@ 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( + 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) @@ -259,30 +366,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" @@ -290,24 +397,40 @@ 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" + assert deployment._model_namespace(config) == ( + "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(deployment.offcfg, "build_lmcache_config", lambda _kvc: cfg) + monkeypatch.setattr( + 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 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): config = _config(role="not-a-role") connected = False - def connect(_config): + def connect(_config, **_kwargs): nonlocal connected 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 @@ -323,13 +446,13 @@ def shutdown(self): adapter = Adapter() monkeypatch.setattr( - mp_connector, + mp_scheduler, "_make_scheduler_adapter", - lambda _config: 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 @@ -351,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) @@ -361,7 +485,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 = page_views._build_cache_views(transfer_tensors, num_blocks=2) assert list(views.tensors) == [ "page.0.primary.0", @@ -371,31 +496,66 @@ 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(pages=[PageRegion(region, page)]) + transfer.set_block_count(2) + + views = page_views._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(): 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"): - 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) + 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"): - 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 + 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"): - mp_connector._build_cache_views(unsupported_rank, num_blocks=2) + page_views._build_cache_views(unsupported_rank, num_blocks=2) @pytest.mark.parametrize( @@ -427,7 +587,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: @@ -457,10 +617,11 @@ 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, poll_interval=0.01, ) @@ -471,15 +632,19 @@ 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 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, poll_interval=0.01, ) @@ -499,8 +664,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( @@ -520,10 +685,11 @@ 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, poll_interval=0.01, ) @@ -538,11 +704,12 @@ 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, poll_interval=0.01, ) @@ -553,16 +720,17 @@ 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): 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, poll_interval=0.01, ) @@ -570,20 +738,21 @@ 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 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__( @@ -663,7 +832,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() @@ -677,7 +846,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): @@ -703,13 +875,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__( @@ -722,7 +895,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, @@ -764,6 +937,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 +949,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,19 +1007,22 @@ 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 -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, @@ -851,7 +1033,7 @@ def _finish_load( def _finish_save( - worker: mp_connector.LMCacheMPConnector, + worker: mp_worker.LMCacheMPConnector, operation_id: str, *, result: bool = True, @@ -879,6 +1061,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 @@ -1054,6 +1237,121 @@ def test_worker_pre_submit_drops_are_immediately_terminal( assert worker._pending_saves == {} +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, 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 unprovable(_request_id, _op, _event): + raise ConnectionError("server may have received request") + + 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(transfer.time, "monotonic", lambda: now[0]) + worker = _worker(adapter) + 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 + + now[0] += 2 + with pytest.raises(transfer.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(transfer.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(transfer.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()) + + output = worker.get_finished() + 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 + if completion.channel == DENSE_PAGE_STORE_CHANNEL + ] + assert not store.succeeded + + +def test_transfer_deadline_defaults_to_twenty_minutes(): + assert transfer._transfer_deadline_s(_config()) == 1200.0 + assert ( + transfer._transfer_deadline_s( + _config(extra={"lmcache.mp.transfer_deadline_s": 30}) + ) + == 30.0 + ) + with pytest.raises(ValueError, match="transfer_deadline_s"): + transfer._transfer_deadline_s( + _config(extra={"lmcache.mp.transfer_deadline_s": 0}) + ) + + def test_worker_save_slices_chunk_blocks_and_preserves_operation( fake_lmcache_modules, ): @@ -1069,6 +1367,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 +1377,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 == { + ConnectorCompletion( + DENSE_PAGE_SOURCE_SAFE_CHANNEL, + SaveSourceGroupId(operation, ((8, 16),)), + True, + ) + } + + def test_non_writer_completes_save_without_submitting(fake_lmcache_modules): adapter = _WorkerAdapter() worker = _worker(adapter) @@ -1094,11 +1420,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 == 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 == 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 == DENSE_PAGE_SOURCE_SAFE_CHANNEL + } == {((8, 16),)} + + def test_worker_tracks_two_load_generations_for_one_raw_request( fake_lmcache_modules, ): @@ -1286,7 +1654,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"]) @@ -1295,7 +1663,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) @@ -1335,7 +1703,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"] @@ -1348,21 +1716,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, ) @@ -1407,11 +1770,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( {}, @@ -1449,11 +1812,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), @@ -1478,7 +1841,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}, @@ -1500,3 +1863,75 @@ def test_factory_registers_lmcache_mp_alias_without_pd_staging(): ) is False ) + + +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 + fp16 = torch.zeros(3, 4, 2, dtype=torch.float16) + 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 + 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_page_region_rejects_what_it_cannot_alias(): + with pytest.raises(ValueError, match="contiguous"): + page_region(torch.zeros(4, 3).t(), semantic_role="strided") + with pytest.raises(ValueError, match="cannot publish"): + page_region( + torch.zeros(8, dtype=torch.uint8), + semantic_role="short", + unit_bytes=4, + total_bytes=12, + ) + + +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_builder_views.py b/tests/test_lmcache_mp_builder_views.py index db847d3265..d1c175d1d2 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 @@ -15,8 +16,10 @@ 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 _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): @@ -143,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", (), {}), @@ -171,6 +202,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 @@ -200,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..``). @@ -219,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) diff --git a/tests/test_lmcache_mp_native_layout.py b/tests/test_lmcache_mp_native_layout.py new file mode 100644 index 0000000000..38b8fa8add --- /dev/null +++ b/tests/test_lmcache_mp_native_layout.py @@ -0,0 +1,408 @@ +# 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, + PageRegion, +) +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( + 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)) + ], + paged_state_checkpoint_spec=spec, + execute_paged_state_copies=lambda stores, restores, descriptor_slot=0: 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 + # 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), + (2,), + (3, 4), + (5,), + (6, 7), + (8,), + (9,), + ) + assert [group.engine_group_id for group in layout.kernel_groups] == [ + 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.pages[0] = replace(transfer.pages[0], view=transfer.pages[0].view.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 + + +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), + ) + + +@pytest.mark.parametrize( + ("breakage", "message"), + [ + (lambda t: setattr(t.block_regions[1], "unit_bytes", 9), "byte geometry"), + ( + 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"), + ], +) +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.page_views 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) diff --git a/tests/test_lmcache_mp_native_scheduler.py b/tests/test_lmcache_mp_native_scheduler.py new file mode 100644 index 0000000000..5a8fc30408 --- /dev/null +++ b/tests/test_lmcache_mp_native_scheduler.py @@ -0,0 +1,886 @@ +# 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, + SaveSourceGroupId, +) +from atom.kv_transfer.offload.chunked_scheduler import DENSE_PAGE_SOURCE_SAFE_CHANNEL +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, +) +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( + deployment.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", + min_save_tokens=0, +): + monkeypatch.setenv("OFFLOAD_MAX_PENDING_SAVES", str(capacity)) + 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 = [] + + def connect(config, *, checkpoint_spec=None): + connections.append(checkpoint_spec) + return adapter + + monkeypatch.setattr(mp_scheduler, "_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, + enable_prefix_caching=True, + 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] + 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_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_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_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 + + 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): + 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, +): + 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): + """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 + 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 + assert scheduler._pinned_state_bytes == 30 + assert scheduler.should_defer_free(seq) + terminal(scheduler, request.save_operation) + 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 + 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(transfer.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(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(transfer.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) + 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_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) + 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_MIN_SAVE_TOKENS", "0") + monkeypatch.setenv("OFFLOAD_MAX_PENDING_SAVES", "2") + monkeypatch.setattr( + mp_scheduler, "_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..fcb0609eb4 --- /dev/null +++ b/tests/test_lmcache_mp_native_worker.py @@ -0,0 +1,597 @@ +# 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, + PageRegion, + 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.deployment 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_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( + pages=[ + PageRegion( + KVTransferRegion( + base_addr=page.data_ptr(), unit_bytes=32, total_bytes=page.numel() + ), + page, + ) + ], + paged_state_checkpoint_spec=spec, + execute_paged_state_copies=lambda stores, restores, descriptor_slot=0: 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_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 + + +@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 transfer + + now = [1000.0] + monkeypatch.setattr(transfer.time, "monotonic", lambda: now[0]) + + def unprovable(*_): + raise ConnectionError("server may have received request") + + 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 + with pytest.raises(transfer.LMCacheTransferUnprovable): + worker.get_finished() + assert len(pending) == 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_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()) + + 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 second.connector_completions == { + ConnectorCompletion( + DENSE_PAGE_SOURCE_SAFE_CHANNEL, + SaveSourceGroupId(req.save_operation, ((8, 16),)), + True, + ) + } + + +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_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),)} + + +class _Event: + def __init__(self): + self.recorded_on = None + self.done = False + + 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) + + def wait_event(self, event): + self.waited_events.append(event) + + +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=[] + ) + + 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() + ) + monkeypatch.setattr( + 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 cuda.restore + assert cuda.restore.waited_streams == [cuda.compute] + + worker.start_load_kv(object()) + 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): + """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( + pages=[ + PageRegion( + KVTransferRegion( + base_addr=page.data_ptr(), unit_bytes=32, total_bytes=page.numel() + ), + page, + ) + ], + paged_state_checkpoint_spec=PagedStateCheckpointSpec( + 32, 128, "native-test-v1", 80 + ), + execute_paged_state_copies=lambda stores, restores, descriptor_slot=0: 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_raising_restore_keeps_its_descriptor_slot_until_the_deadline( + worker, monkeypatch +): + """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 native_state_worker, transfer + + now = [1000.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( + 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 + 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._transfer_deadline_s + with pytest.raises(transfer.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 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 deployment + + monkeypatch.setattr(deployment.offcfg, "build_lmcache_config", lambda _: object()) + monkeypatch.setattr(deployment.offcfg, "lmcache_replica_world_size", lambda _: 2) + monkeypatch.setattr( + deployment.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..29a156505d 100644 --- a/tests/test_lmcache_mp_shell.py +++ b/tests/test_lmcache_mp_shell.py @@ -1,12 +1,17 @@ # 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 ( + 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, @@ -22,9 +27,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 mp_worker.LMCacheMPConnector + assert LMCacheMPConnectorScheduler is not mp_scheduler.LMCacheMPConnectorScheduler def test_factory_registration_resolves_public_connectors(monkeypatch): @@ -56,3 +61,182 @@ 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(mp_worker, "LMCacheMPConnector", PageWorker) + config = _config() + worker = LMCacheMPConnector(config) + tensors = SimpleNamespace( + paged_state_checkpoint_spec=object(), + execute_paged_state_copies=lambda stores, restores, descriptor_slot=0: 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(mp_scheduler, "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(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) + ) + 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(mp_worker, "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(mp_worker, "LMCacheMPConnector", PageWorker) + builder = SimpleNamespace( + state_transfer=lambda: SimpleNamespace(copies=False), + ) + LMCacheMPConnector(_config()).register_kv_caches( + {}, 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( + mp_scheduler, + "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_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..1f36f87ca2 100644 --- a/tests/test_multi_connector.py +++ b/tests/test_multi_connector.py @@ -267,6 +267,40 @@ 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, + ) + + 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 + + # --------------------------------------------------------------------------- # Scheduler-side # --------------------------------------------------------------------------- @@ -959,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 9428ce780d..0e1cb2b7cf 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): @@ -88,6 +89,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 +318,214 @@ 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( + enable_prefix_caching=True, + 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_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, caplog + ): + scheduler = _early_release_scheduler(monkeypatch, chunk_size=8) + 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[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 + 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.""" + 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 + + 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 + + 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: @@ -550,6 +742,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..9fa46b896f --- /dev/null +++ b/tests/test_page_unit_checkpoint_leases.py @@ -0,0 +1,280 @@ +# SPDX-License-Identifier: MIT + +"""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, + 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_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) + 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) + + +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 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"] 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 == [] diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 0ae1bb0642..bc0bd810f3 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( @@ -2734,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.