Skip to content

Distributed block-table serialization (kv_cache_serde) #125

Closed

Description

@inureyes

Context

Disaggregated / decode-node handoff must transfer paged block tables and pool blocks, not dense snapshots, once the unified store is the source of truth.

Tasks

  • Extend src/distributed/kv_cache_serde/ to serialize PagedSequenceState block tables plus the referenced pool blocks and restore them onto the decode-node pool. (Fp16 pool-backed path; the block table is restored via the metadata path while pool block CONTENTS are transferred and re-materialized on fresh blocks. Turbo sidecars are out of scope for the Fp16-only path.)
  • Validate a prefill-node to decode-node handoff reconstructs identical KV and identical first-token logits. (Covered by the #[ignore] real-model parity test tests/paged_kv_serialize_parity.rs, run on Metal.)
  • Confirm refcount and block ownership are correct after restore (no double-free, no leak). (CPU unit tests assert per-block refcount == 1 and live-block-count parity after restore.)

Acceptance criteria

  • A disaggregated run reproduces single-node output within tolerance using the paged block-transfer path. (Covered by the #[ignore] parity test, run on Metal; the live disaggregated scheduler/transport wiring is the epic capstone End-to-end validation, parity, benchmarks, and docs #126.)
  • Block accounting after restore matches the originating node. (CPU unit tests assert paged_stats() + live_block_count() parity between origin and decode pools.)

Implementation status

Block-content transfer ("Approach A") is implemented in the kv_cache_serde module's public flow plus the supporting mlxcel-core pool/cache primitives. Wiring serialize/restore into the live disaggregated scheduler and a real transport is deferred to #126. See the PR for details.

Dependencies

Blocked by Phase 4 (unification).

Part of #116

Activity

  1. added this to the 0.2.0 milestone on May 27, 2026
  2. added
    status:backlogIn the backlog, not yet ready
    type:enhancementNew features, capabilities, or significant additions
    area:coremlxcel-core: MLX FFI, primitives, KV cache, layers
    on May 27, 2026
  3. added 2 commits that reference this issue on Jun 8, 2026
    92273c6
    c26c484
  4. self-assigned this
    on Aug 31, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

area:coremlxcel-core: MLX FFI, primitives, KV cache, layerspriority:highHigh prioritystatus:doneCompletedtype:enhancementNew features, capabilities, or significant additions

Type

No type

Projects

No projects

    Milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions