[yizhang-nv]

Description

Enable sleep/wakeup with KV-cache manager V2 while retaining native cache objects and their logical state. Wakeup uses ordinary page and scratch allocation. GPU pool virtual addresses and metadata buffers stay stable; refreshed KV/SSM indices let captured CUDA graphs address the restored pages even when slot IDs change.

• NONE/MEMSET discard reusable prefixes and release GPU and HOST pool backing without backing discarded contents up to disk. Retained cold pages are initialized through the existing codec; MEMSET zeros restored GPU contents.
• CPU/PINNED preserve contents through native host tiers and migration. Host-capacity checks run only when a preserving KV release is requested; weights-only and discard modes do not impose a global backup quota.
• Restore HOST pools before pages. Materialize discarded pages in one transaction across caches, cold tiers before hot, preserving original tier occupancy and rollback. Scratch is acquired by normal resume; scratch left by a failed resume is released before sleep.
• Refresh KV/SSM indices, sparse metadata, guard pages, and KDA replay history after remapping. Python retains active request IDs without copying request implementation fields into a sleep snapshot.
• Preserve small replay control tensors (accepted-token counts, buffer selectors, fixed row indices and the optional GDN pointer/stride table) with the existing allocator policy. NONE/MEMSET can discard model state without invalidating kernel indices or captured metadata addresses.
• Keep admission closed during transitions, coordinate MPI rank and Ray actor validation before mutation, pause pool maintenance while asleep, and reject KV sleep for registered transfer buffers that cannot be safely remapped.

Sleep drains user requests; internal dummy/padding caches survive. HOST base addresses and individual slot IDs may change. The CUDA Graph contract is stable GPU pool/metadata addresses with current indices. Transition failures retain the existing fail-stop service contract.

Test Coverage

Validation uses a shared B200 development container (umbriel-b200-043) and the standard SM100 build. Full stdout/stderr are saved outside the repository under /home/scratch.yizhan_sw_1/logs/2026-10-16/umbriel-b200-043/trim_tests_*. Tests remain selected by the existing CI lists.

The test changes are limited to migrating the existing sleep tests to real V2 execution (asserting the live manager is KVCacheManagerV2, on every rank for TP2) plus one check per new contract. Results at a1fac4449c:

• C++: cold-page/codec/tier 76 passed, host memory 9 passed.
• Native V2 Python: discard/MEMSET/preserve sleep state and cold-codec round trips 5 passed.
• VMM/CUDA graph: 10 passed, including the V1 control, V2 NONE/MEMSET/CPU, a cold GPU tier, and graph replay reading a KV page whose slot changed across sleep.
• Mamba hybrid sleep cases (logical slots, KDA replay relocation, replay control metadata): 8 passed.
• Executor/policy/routing/guard units (test_kvcm2_integration, test_py_executor, test_kv_pool_rebalance, test_seq_slot_sizing, test_ray_sleep_state, test_sleep_collective_rpc_guards): 684 passed (the guard file was re-run after a stub fix: 86 passed).
• Ray single-GPU actual-model sleep: 4 passed (KV-only CPU restore with prefix reuse preserved, all tags with NONE restore and reuse invalidated, existing V1 prefix-reuse control, weight discard).
• MPI TP2 actual-model sleep: 2 passed, including repeated KV-only sleep with per-device memory checks and output equivalence.
• Ray TP2 async release/resume: 2 passed.

Independent native-state and Python-integration reviews, fresh skeptical review, and actual claude-opus-5-5 review plus focused follow-up cover the final implementation. Review fixes retain failed-resume scratch cleanup, globally order cold-page initialization, correct the graph test's K/V-page indexing and preserve replay controls. The later test-only trimming was not re-run through the external reviews. Wakeup latency is unmeasured; full hybrid-model graph replay and the optional non-affine GDN descriptor fallback have not been exercised at runtime.

PR Checklist

• Description explains the final behavior and scope.
• Changes follow the coding guidelines and use signed commits.
• Relevant tests and CI selection are retained; no waive/skip/fallback masks V2 coverage.
• Documentation updated; no new dependency or public configuration field.
• No ownership or architecture-diagram change required.
• Latest revision V2 runtime validation complete; unrelated baseline failures are disclosed above.

[chienchunhung]

@yizhang-nv I opened draft successor #21221 from my fork to drive this work forward. It preserves all 30 original commits, authorship and sign-offs, with the same head af7661fed0f5; no implementation changes are part of the handoff.

The scope remains native V2 sleep/wake. Snapshot-specific integration and transfer reconstruction will stay separate, and the current transceiver restriction remains in place. Rebase, review and fresh verification are pending in the successor.

Can we use #21221 for follow-up development and keep this PR as the original reference? I have left this PR open. Please flag any pending work or review context to carry forward; targeted KV-manager review would be welcome.

[yizhang-nv]

@yizhang-nv I opened draft successor #21221 from my fork to drive this work forward. It preserves all 30 original commits, authorship and sign-offs, with the same head af7661fed0f5; no implementation changes are part of the handoff.

The scope remains native V2 sleep/wake. Snapshot-specific integration and transfer reconstruction will stay separate, and the current transceiver restriction remains in place. Rebase, review and fresh verification are pending in the successor.

Can we use #21221 for follow-up development and keep this PR as the original reference? I have left this PR open. Please flag any pending work or review context to carry forward; targeted KV-manager review would be welcome.

@chienchunhung LGTM