[DSV4] Add PAGE-backed state checkpoints

[yhl-amd]

Summary

• Keep each resident request on the existing contiguous, fixed-address Active Slot.
• Store immutable DSV4 state checkpoints in an ordered set of arbitrary PAGE-sized units, so checkpoint allocation no longer requires a contiguous 13-20-PAGE extent.
• Make PAGE-backed checkpointing the only path for StateTransfer.copy(); remove ATOM_V4_PAGED_STATE_CHECKPOINTS, the old slot-backed copy checkpoint fallback, and STATE_CKPT_EXTRA_ENTRIES Active Slot headroom.
• Carry runtime geometry as a frozen PagedStateCheckpointSpec (page-unit bytes, slot bytes, layout id) instead of mutating Config; units-per-checkpoint is derived and never transmitted.
• Make StateTransfer the single backend checkpoint capability: copy(layout_id), fork(tokens), and none() are invariant-checked. StateRuntime combines that capability with the optional PagedStateCheckpointSpec, validates the complete combination/layout once, and crosses runner RPC as one nested typed payload.
• Name Active Slot relocation and PAGE checkpoint scatter/gather separately as relocate_state_slots and execute_paged_state_copies; PD/RDMA get_kv_transfer_tensors remains independent.
• Add a CheckpointRecord control plane with COPYING/READY/EVICTING states, generation checks, reader pins, LRU, and atomic whole-record eviction.
• Carry one typed StateMaintenanceOps bundle through ScheduledBatch; its relocations, checkpoint stores, and checkpoint restores are drained once for a real batch and executed before the forward. DSV4 Main/Indexer scatter/gather uses one descriptor-driven Triton byte-copy launch.

Allocation and correctness invariants

• One PAGE remains one CompositePageUnit; one state image owns K arbitrary unit ids in canonical order.
• A prefix hash points to one READY checkpoint record, never to individual fragments.
• COPYING records are not visible to prefix lookup.
• PAGE allocation can reclaim only an entire unpinned checkpoint; no fragment is independently reused.
• A checkpoint hit first allocates a complete Active Slot and pins its source units before allocating fresh PAGE blocks. PAGE fragments are never adopted as a runtime slot.
• Main KV and FP8/FP4 Indexer regions share the same block-id ownership transition.
• Copy-transfer backends publish stable PAGE/state layout and geometry through the explicit runner -> engine -> BlockManager runtime spec. Missing or drifted geometry fails at startup; there is no slot-backed fallback.
• GDN StateTransfer.fork() checkpointing remains group-backed and is not part of the removed DSV4 copy path.

Implementation notes

• BlockPool.reserve_units/release_units records checkpoint identity, generation, and piece index for each reserved block id.
• ModelRunner derives the checkpoint spec from PoolPlan plus transfer.paged_layout_id and constructs one validated StateRuntime; EngineCore rebuilds that runtime from its nested wire payload and passes the same object explicitly through Scheduler to BlockManager.
• PageUnitCheckpointStore owns hash lookup, ordered units, LRU, pins, two-phase publication, and deferred cleanup.
• The copy planner intersects the Active Slot and Composite PAGE logical byte streams into physical spans, expands them into 4 KiB tiles, and launches one raw-byte Triton copy for all store/restore ops in the batch.
• Pipeline parallel and RapidServe configurations fail fast until they have a cross-stage/cross-scheduler ACK and rollback protocol.

Validation

Passed in the current workspace:

• python3 -m compileall -q atom ...
• git diff --check
• Black 26.5.1 and Ruff on the changed paged-state control-plane files
• 19 pure control-plane tests for StateTransfer, StateRuntime, StateMaintenanceOps, PagedStateCheckpointSpec, and PageUnitCheckpointStore
• control-plane smoke covering mandatory PAGE construction, store -> READY, restore into a distinct contiguous Active Slot, and removal of the slot-backed fallback
• explicit allocation from non-contiguous free PAGE ids

Added unit coverage for:

• arbitrary-unit reservation and atomic release
• COPYING visibility and READY publication
• multi-reader pins and whole-record LRU eviction
• unindex/clear during in-flight copies
• protected-hit admission accounting
• StateRuntime COPY/FORK/NONE combinations, illegal transfer/spec/layout pairs, nested wire revalidation, derived unit count, and explicit Scheduler/BlockManager propagation
• empty batches retaining queued state maintenance, and real batches carrying relocation/store/restore exactly once
• segmented-stream boundary and partial-tail planning
• GPU random-byte scatter/gather round trip when a CUDA/ROCm device is available

GitHub Pre Checkin, including Black, Ruff, schema validation, and the full non-GPU unit suite, passes for the validated StateRuntime revision. Offline inference also passes.

DSV4 GPU BF16/FP8/FP4 bitwise and performance validation still requires a ROCm/Triton machine.

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 1894 --add-label <label>

[valarLip]

Reviewed at 891952b2b. The design direction is right — PAGE-backing is what finally decouples checkpoint capacity from max_num_seqs, and it sidesteps the layer-major/entry-major ordering problem that made an elastic KV↔arena boundary intractable. Two correctness holes should block the merge though, plus one hardening item. Both blockers have a red repro at the bottom of this comment.

Both blockers come from the same asymmetry:

	where it begins	when it is completed
store	begin_store, inside take_checkpoint_ops — only ever for a batch that really runs	complete_inflight
restore	begin_restore, from BlockManager._attach_state_group during allocate() — at admission	complete_inflight, unconditionally

The reader pin and the queued gather are taken together but completed apart.

---

Blocker 1 — an empty batch unpins a gather that has not run yet

Scheduler.schedule calls complete_previous_state_batch() unconditionally at the top, but skips take_state_maintenance_ops() when the batch is empty (deliberately — test_empty_batch_does_not_drain_state_maintenance pins that). So a restore queued on a tick that produces no batch has its source unpinned on the next tick while its op is still in _restore_ops.

The record then becomes READY with pin_count == 0, which is exactly what ensure_free_units looks for. Its units are recycled into ordinary KV blocks by _fresh_block, and whenever the queued gather finally rides a batch it copies those recycled bytes into the resuming request's Active Slot — stale compressor ring and stale SWA window, the #1417 shape the two-phase publication exists to prevent.

Reachable path (scheduler.py):

1261  self.block_manager.allocate(seq, num_cached_blocks)
        └─ _attach_state_group → begin_restore → pin += 1, op appended to _restore_ops
1272  needs_remote_load = self._confirm_remote_load_after_alloc(seq, needs_remote_load)
1277  self._park_for_remote_load(seq, skipped_waiting_requests)
1278  continue

_park_for_remote_load only appends to skipped_waiting_requests and sets WAITING_FOR_REMOTE_KVS — it does not deallocate. The seq never reaches _schedule_prefill_seq, so it never enters scheduled_seqs. With nothing else running, the batch is empty.

This is not gated away: model_runner.py:1674-1685 rejects pipeline_parallel_size > 1 and enable_rapidserve, but not a KV transfer config, and self.kv_connector = get_kvconnector("scheduler", config) (scheduler.py:708) is independent of RapidServe. So the PD/connector path and PAGE checkpoints coexist.

Worth noting the scope is narrow: a non-empty batch is fine — the op is drained into it and the gather executes before the parked seq ever runs. Only the "allocated but the batch came out empty" tick bites.

Verified end to end on the fixture below by asserting the buggy behaviour instead of the invariant: after complete_previous_batch() the pin is 0, ensure_free_units evicts the record, pool.num_free rises by its whole unit count, and the still-queued op goes on naming those units.

Blocker 2 — deallocate does not cancel a gather aimed at the released slot

BlockManager.deallocate calls paged_state_checkpoints.forget_pending(seq), whose whole body is self._pending.pop(id(seq), None) — store intents only. A CheckpointRestoreOp already queued in _restore_ops with dst_slot == seq.per_req_cache_group survives, and the very next line returns that slot to the free list via self.state.release(...).

The next request to pop the slot gets the dead request's checkpoint image scattered over its Active Slot before its own forward — wrong compressor state and wrong sliding window for a request that never asked to resume. If two restores name the same recycled slot in one batch, the two gathers race inside a single Triton launch.

Triggers are all ordinary: a request is admitted on a checkpoint hit (op queued), then aborted, or its remote KV load fails, or it is preempted before the op drains.

Repro output:

AssertionError: a gather into a slot that has gone back on the free list must not
survive deallocate; got (CheckpointRestoreOp(dst_slot=4, unit_ids=(0, 1, 2),
total_bytes=25, layout_id='layout-v1'),)

Suggested direction for both

Same fix shape: tie the pin and the cancellation to the moment the op actually leaves the queue, not to the tick. Concretely, have complete_inflight's restore half only settle restores that were handed out by a previous take_checkpoint_ops, and have forget_pending also drop _restore_ops entries whose dst_slot matches the sequence's group.

---

Hardening — block_manager.py:443

assert restored, "gated PAGE checkpoint disappeared before attach"

python -O erases this, and the request then proceeds on a freshly popped slot still holding the previous occupant's compressor ring and SWA window — silently. The same file already makes this call the other way at __init__ lines 141-143 ("rather than asserting, which python -O would drop"), where the interval snap became a warning. Either raise unconditionally, or fall back to "no resume" and drop the hit.

---

Non-blocking notes

• Checkpoint stores now draw from the live KV pool with no watermark. begin_store → reserve_units → pool.pop() + allocate() takes from the same pool that serves may_append, and the resulting COPYING record is refused by both has_available_units and ensure_free_units (both require READY) for a full tick. Removing STATE_CKPT_EXTRA_ENTRIES at the same time leaves no headroom knob at all. The failure mode changes character here: today a checkpoint that cannot find a free slot is simply dropped and live requests are untouched; after this PR a burst of simultaneous rung crossings (uniform-length requests under a benchmark all cross the ladder together) can lock N × units_per_checkpoint blocks that neither can_append nor _fresh_block can reclaim. On a 256-seq V4 config that measures 14 units per checkpoint, 64 concurrent stores is ~1.4 GB unreclaimable for a tick. A reserve fraction on checkpoint units, or making COPYING cancellable, would keep the upside without the new starvation mode.

• launch_copy_spans builds the descriptor in Python on the pre-forward critical path. Spans are expanded to one entry per 4 KiB tile, so a ~21 MB Active Slot is ~5,200 tiles × 3 arrays (src_ptrs int64, dst_ptrs int64, valid_bytes) ≈ 15k Python ints marshalled through a synchronizing pageable torch.tensor(..., device=cuda) per op, inside CommonAttentionBuilder.build(). Since the Active Slot is contiguous and PAGE units are equal-sized, the mapping is affine — the only per-op input is the K unit ids. Precomputing the logical span table once and passing just the unit-id table (letting the kernel do unit = offset // unit_bytes) turns this into a block table, the same shape the KV path already uses.

• BlockPool.reserve_units is documented as atomic but re-enters through the eviction callback. reserve_units → pop → allocate → _unindex → on_evict → _record_evicted → unindex → _release_record → release_units → free mutates _free, _vacant, _cached and _raw_unit_owner in the middle of the reservation loop. Benign today (the re-entrant path only adds free blocks, and a pending hash always names a block held by a live sequence so it cannot be the popped victim), but the "Atomically reserve" docstring is not accurate. Minor: the # Validate ownership before releasing any unit. comment sits below the loop it describes.

Checked, no issue found

Wire format (StateRuntime.to_wire/from_wire → EngineCore → Scheduler → BlockManager) is consistent; h if num_cached_blocks > 0 else -1 at block_manager.py:405 short-circuits fine on a cold start; the byte-stream ordering is symmetric between _active_slot_segments and _page_unit_segments for store and restore; block_bytes/slot_bytes are linear in row width so the sizing cross-check in allocate_per_req_cache holds; slot rows and block envelopes provably do not overlap; no stale references remain to state_copy_pairs, copy_state_entries, release_state_pins, state_transfer_kind, state_fork_tokens, pending_checkpoint, or STATE_CKPT_EXTRA_ENTRIES; TBO's build_ubatch_* paths derive from the single build() result so maintenance ops still execute exactly once per batch; _checkpoint_room's block_manager.state.transfer.forks correctly evaluates False under PAGE checkpointing, so DSV4 spec-decode boundaries stay checkpointable.

Failing PP and RapidServe fast rather than half-supporting them is the right call — PP plus stateful attention has never actually started on our box, and the deferred-hash path there would have stored the wrong bytes silently.

On validation

The PR notes that DSV4 GPU bitwise and performance validation still needs a ROCm/Triton machine. Worth flagging that this PR replaces the path #1417 was filed against, and the slot-backed scheme it replaces was signed off with 1319-question GSM8K at n≥3 per arm plus the #1417 coherence probe — that model needs n=3 because same-code run-to-range reaches 1.6pp. Happy to run gate A (V4-Flash-DSpark tp2, 1319 × 3) and the #1417 prefix-hit coherence probe here once the two blockers are fixed; running them before that would just produce dirty numbers.

---

Repro — tests/test_page_unit_checkpoint_restore_lifetime.py (both tests fail at 891952b2b)

# SPDX-License-Identifier: MIT

"""Repro for the two restore-lifetime holes in PAGE-backed state checkpoints.

Both come from the same asymmetry. A *store* is begun inside
`take_checkpoint_ops`, so it only ever exists for a batch that really runs. A
*restore* is begun from `BlockManager._attach_state_group` during `allocate()`,
so its reader pin and its queued gather are taken at admission — one for the
tick, the other for whenever a batch next carries it. Nothing keeps the two
together, and both of these tests show them coming apart.
"""

import pytest

from atom.model_engine.block_pool import BlockPool
from atom.model_engine.page_unit_checkpoint import (
    READY,
    PagedStateCheckpointCoordinator,
    PagedStateCheckpointSpec,
)


class _Seq:
    """The only thing `forget_pending` reads off a sequence is its identity."""

    def __init__(self, group):
        self.per_req_cache_group = group
        self.has_per_req_cache = True


def make_coordinator(num_units=20, unit_bytes=10, slot_bytes=25):
    pool = BlockPool(num_units)
    spec = PagedStateCheckpointSpec(
        page_unit_bytes=unit_bytes,
        slot_bytes=slot_bytes,
        layout_id="layout-v1",
    )
    return pool, PagedStateCheckpointCoordinator(pool, spec, enabled=True)


def publish(coordinator, prefix_hash, src_slot=0):
    """A checkpoint whose store has already ridden a batch."""
    assert coordinator.store.begin_store(prefix_hash, src_slot=src_slot) is not None
    coordinator.complete_previous_batch()
    checkpoint_id = next(
        cid
        for cid, record in coordinator.store.records.items()
        if record.prefix_hash == prefix_hash
    )
    assert coordinator.store.records[checkpoint_id].state == READY
    return checkpoint_id


def test_an_empty_batch_unpins_a_gather_that_has_not_run_yet():
    """The pin has to outlive the gather, not the tick.

    `Scheduler.schedule` calls `complete_previous_state_batch` unconditionally
    but skips `take_state_maintenance_ops` when the batch is empty. A request
    admitted on a tick that produces no batch — the PD consumer path parks it
    in `_park_for_remote_load` after `allocate` — therefore has its source
    unpinned while its gather is still queued.

    What that costs, measured on this fixture by asserting the buggy behaviour
    instead of the invariant: once the pin is gone the record is READY and
    unpinned, `ensure_free_units` evicts it, `pool.num_free` rises by its whole
    unit count, and the still-queued gather goes on naming units that
    `_fresh_block` is now free to hand to anyone. That is the #1417 shape.
    """
    _, coordinator = make_coordinator()
    checkpoint_id = publish(coordinator, prefix_hash=7)

    # Admission takes the pin and queues the gather.
    assert coordinator.begin_restore(7, dst_slot=1)
    assert coordinator.store.records[checkpoint_id].pin_count == 1

    # The tick produced an empty batch, so the bundle is NOT drained. The next
    # tick completes the previous batch all the same.
    coordinator.complete_previous_batch()

    assert coordinator.store.records[checkpoint_id].pin_count == 1, (
        "a gather that has not ridden a batch must keep its source pinned; "
        "the pin belongs to the op, not to the tick"
    )


def test_deallocate_does_not_cancel_a_gather_aimed_at_the_released_slot():
    """`forget_pending` drops store intents only; restores keep their target.

    `BlockManager.deallocate` calls `forget_pending(seq)` and then returns
    `seq.per_req_cache_group` to the state pool's free list. A gather already
    queued against that slot survives both, so the next request to pop the slot
    gets the dead request's checkpoint image written over its Active Slot
    before its own forward.
    """
    _, coordinator = make_coordinator()
    publish(coordinator, prefix_hash=7)

    seq = _Seq(group=4)
    assert coordinator.begin_restore(7, dst_slot=seq.per_req_cache_group)

    # The request is aborted / its remote load fails / it is preempted.
    coordinator.forget_pending(seq)

    _, restores = coordinator.take_checkpoint_ops()
    assert restores == (), (
        "a gather into a slot that has gone back on the free list must not "
        f"survive deallocate; got {restores}"
    )

[yhl-amd]

Addressed the maintainer review blockers in 895827b:

• Restore reservations now remain queued until take_checkpoint_ops hands them to a real batch; complete_previous_batch settles inflight restores only.
• deallocate cancels any queued gather targeting the released Active Slot and drops its reader pin before the slot returns to the freelist.
• A disappeared gated checkpoint now returns the newly allocated slot and raises RuntimeError unconditionally.
• Corrected the reserve_units atomic wording and ownership-validation comment placement.

Added regressions for the empty-batch lifetime, queued cancellation including EVICTING release, recycled-slot safety, and the fail-fast path. Local validation: 291 related BlockManager/Scheduler/state tests passed; Black, changed-file Ruff, compileall, and diff check passed.

The watermark and copy-descriptor observations remain non-blocking follow-up work pending stress/performance data.

[yhl-amd]

E2E validation on 895827b8 completed with 8x MI355X using image rocm/atom-dev@sha256:bd915ac97e249e1cf43282bc9232b79e18be76bece0b8788caa463e3501a3ec6.

• Full GSM8K, 64-shot, 1,319 requests per round; every prompt exceeded the 8,192-token checkpoint interval (min 9,098, max 12,021, mean 10,463.28).
• Round 1: strict 0.9613343442, flexible 0.9605761941, 548 s. The 1,319 token-ID records were fsynced; SHA256 6192978780a1fe3602ea904cc05218656733808c168c856818e2e117b35cb155.
• Round 2 replay on the same server: 1,319/1,319 HTTP 200 on attempt 1, 381.42 s. Re-evaluated with the same lm-eval filters: strict/flexible 0.9620924943 (1,269/1,319), so no accuracy regression.
• Round-2 cache delta: 8,249,600 cached / 13,801,066 full tokens = 59.7751% realized hit; 74.7982% compressed hit; 15.0231% lost to the 8K checkpoint boundary. Checkpoints kept +1,320, dropped 0, evicted 0.
• Startup reported PAGE-backed checkpoints with 1,521,920-byte units, 25,395,200-byte slots, and 17 units/checkpoint. Server logs show 8,192-token restores (cached: [8192]); no traceback/assert/runtime/OOM/GPU-fault patterns, and the container was not OOM-killed or restarted.

[yhl-amd]

Follow-up MTP/DSpark coverage on the same 895827b8 image and 64-shot GSM8K workload:

• Re-ran both rounds with real MTP-1 enabled: --method mtp --num-speculative-tokens 1; startup confirmed SpeculativeConfig(method=mtp, num_spec_tokens=1) and loaded the mtp.0.* draft layer.
• MTP changes the state geometry from ring 128 / 25,395,200-byte slot / 17 PAGE units to ring 129 / 26,782,720-byte slot / 18 PAGE units.
• Round 1: strict 0.9545109932 (1,259/1,319), flexible 0.9537528431; 534 s; acceptance 95.9064% (60,398/62,976), 1.9591 tok/fwd.
• Round 2: 1,319/1,319 first-attempt success; strict 0.9583017437 (1,264/1,319), flexible 0.9575435936; 366.57 s; acceptance 95.8537% (60,292/62,900). Accuracy improved by 5 questions vs the MTP cold round, so the 18-unit checkpoint restore path did not regress accuracy.
• Round-2 cache delta: 7,365,120 cached / 13,801,066 full = 53.3663% realized hit; 66.5734% compressed hit; 13.2071% lost to the 8K boundary. Checkpoints kept +1,321, dropped 0, evicted 0. No fatal/OOM/GPU-fault patterns.
• Compared with non-MTP, MTP was ~2.6% faster on round 1 and ~3.9% faster on round 2. Scores varied by ~0.38-0.68 pp across the distinct batching/spec paths; the within-MTP checkpoint-hit round did not decline.

DSpark is not validly testable with this local checkpoint: it has num_nextn_predict_layers=1 and mtp.0.*, but no dspark_block_size, DSpark target-layer metadata, or separate DSpark draft checkpoint. DSpark is a different block drafter, not a runtime toggle that can be safely forced onto serial-MTP weights. A V4-Pro-DSpark checkpoint is required for that follow-up.