Skip to content

feat: asynchronous hierarchical KV swap - #185

Open
drunkcoding wants to merge 21 commits into
mainfrom
plan/async-hierarchical-kv
Open

feat: asynchronous hierarchical KV swap#185
drunkcoding wants to merge 21 commits into
mainfrom
plan/async-hierarchical-kv

Conversation

@drunkcoding

@drunkcoding drunkcoding commented Aug 21, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Description

Implements event-driven GPU↔pinned-host KV swap for serving while retaining the default synchronous pageable-buffer path.

Scope

  • Bounded pinned-buffer and in-flight byte caps with backpressure
  • Generation-safe transfer records, block ownership, cancellation tombstones, retries, and shutdown drain
  • Group-atomic scheduler polling and recovery
  • Shared native coordinator transfer backend with exact byte/error accounting
  • Operator telemetry, Prometheus exposition, CUDA race coverage, benchmark tooling, and rollout documentation
  • ExternalKVStore remains a protocol-only extension boundary

Non-goals

No external/distributed store, SSD/RDMA/object-storage backend, multi-node protocol, or KV quantization.

Verification

  • pre-commit run --all-files: PASS
  • Focused Task 10 CPU suite: 154 passed
  • Full serving suite: 181 passed, 2 skipped (CUDA/FlashInfer prerequisites)
  • CUDA async KV suite: 9 passed
  • 1,000-cycle fake-backend ownership/leak stress: PASS (1,000 tickets retired)

Environment blockers

  • Full integration collection requires an OpenAI-compatible server on localhost:8000; test_oai_chat_completions.py and test_oai_completions.py fail collection with ConnectError: [Errno 111] Connection refused when that external prerequisite is absent.
  • Full sync/async DeepSeek-V2-Lite A/B commands produced status=BLOCKED because the shared Hugging Face cache path is not writable (PermissionError); parser, config conversion, percentile, and schema tests pass.

Type of Change

  • Production implementation
  • Tests
  • Documentation
  • Benchmark tooling

Checklist

  • Default sync path remains blocking/pageable and does not consume pinned-pool accounting
  • Async block/lease ownership survives cancellation until ticket retirement
  • Scheduler queue/status transitions are group-atomic
  • Native coordinator owns CUDA tickets; unified scheduler owns queue/results only
  • PR remains draft pending review and environment-dependent A/B evidence

drunkcoding and others added 2 commits August 21, 2026 14:54
drunkcoding @sisyphus-dev-ai
docs(plan): add async hierarchical KV plan
aba21d9 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
docs(plan): address Momus review
0a49bfe 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
@drunkcoding

drunkcoding commented Aug 28, 2026

Copy link
Copy Markdown
Contributor Author

Momus plan review: [OKAY] after 2 review rounds. Executable documentation QA was added in commit 0a49bfe.

drunkcoding and others added 17 commits August 29, 2026 00:42
drunkcoding
feat: add bounded KV transfer primitives
39bdac0
drunkcoding
feat: track asynchronous serving KV residency
ff8ced9 
Replace the blocking swap-out/swap-in buffers with per-generation KV
records, block ownership tracking, and event-driven transfer submission
in PagedKVCache. Async mode reserves pinned leases or destination blocks,
holds GPU blocks until DMA completion, and uses generation-keyed
tombstones so cancellation never targets recycled storage. Sync mode
keeps the existing blocking pageable round-trip and never constructs a
pinned pool.

Adapt three pre-existing tests that probed the removed internal fields
to the new record structure.
drunkcoding
feat: gate serving scheduling on KV transfer events
50bb629 
Replace the blocking preempt/recover path with an event-driven,
group-atomic scheduler. schedule() now polls KV transfer completions
once per cycle and advances per-request swap group records through
OUT_IN_FLIGHT -> HOST_RESIDENT -> IN_IN_FLIGHT phases before any queue
transition. Preemption reserves and submits a whole group at once,
returns backpressure without mutating state on reservation failure, and
never frees GPU blocks while an async D2H is in flight (sync completions
free immediately). Swap-in recovery runs after waiting admission so
resident work is not starved, restores each member's saved prior status,
and evicts a group to _waiting on retry exhaustion. has_work() now
accounts for swapped groups and pending transfers; has_runnable_work()
gates the engine loop on GPU-ready sequences.

Add fake-async-backend scheduler tests covering evicting-block reuse,
H2D-gated decode readiness, and the six group atomicity/retry rules, and
adapt the OOM recovery edge-case test to the new recovery entrypoint.
drunkcoding
feat: wire async KV lifecycle and fallback
d5883ad
drunkcoding @sisyphus-dev-ai
feat: report unified transfer results
f72f7b6 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
refactor: reuse KV transfer backend in native coordinator
f282598 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
feat: wire native KV offload to paged tensors
589f4f5 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
test: cover CUDA KV transfer ordering and races
a44d048 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
feat: collect asynchronous KV swap telemetry
40fceeb 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
feat: expose asynchronous KV swap metrics
2b8c550 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
test: make fake KV swap pool CPU-safe
318011e 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
bench: measure async KV swap tail latency
a2761aa 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
docs: add async KV swap operations guide
cd7bcff 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
docs: add KV swap benchmark rollout notes
a5feb49 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
test: keep async KV fixtures CPU-only
8a67a58 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
test: make CUDA transfer delay deterministic
8837d1d 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
drunkcoding @sisyphus-dev-ai
test: warm CUDA KV transfer path before gating
1168a5f 
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
@drunkcoding drunkcoding changed the title WIP: plan asynchronous hierarchical KV swap WIP: asynchronous hierarchical KV swap Aug 29, 2026
@drunkcoding

drunkcoding commented Aug 29, 2026

Copy link
Copy Markdown
Contributor Author

Task 5–10 implementation and verification update (1168a5f):

  • pre-commit: PASS
  • focused Task 10 CPU suite: 154 passed
  • serving suite: 181 passed, 2 skipped (CUDA/FlashInfer prerequisites)
  • CUDA-gated async KV suite: 9 passed
  • 1,000-cycle fake-backend stress: PASS, 1,000 tickets retired, zero blocks/leases/inflight/host-resident/retiring records

External prerequisites/blockers:

  • Full integration collection reaches test_oai_chat_completions.py and test_oai_completions.py, which require an OpenAI-compatible server on localhost:8000; without it both stop at collection with ConnectError: [Errno 111] Connection refused. This is not a product-test regression.
  • The full DeepSeek-V2-Lite sync/async A/B commands were attempted on the CUDA host, but both JSON reports are status=BLOCKED because the shared Hugging Face cache path is not writable (PermissionError). Benchmark unit coverage (parser/config conversion/percentiles/schema) is green.

The PR remains draft; no merge/ready/force operation was performed.

drunkcoding added 2 commits August 29, 2026 14:17
drunkcoding
refactor: drop dead KV selector and make native offload install test …
df150e6 
…hermetic
drunkcoding
fix(serving): read kv_swap args defensively in _initialize_model
4e425b8 
_initialize_model read args.kv_swap_* directly, so any caller whose args
Namespace predates the kv-swap flags (e.g. existing watchdog init paths)
crashed with AttributeError, aborting model init. Use getattr with the
documented defaults, matching the existing getattr at the reload path.
@drunkcoding
drunkcoding marked this pull request as ready for review September 2, 2026 21:53
@drunkcoding drunkcoding changed the title WIP: asynchronous hierarchical KV swap feat: asynchronous hierarchical KV swap Sep 2, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@drunkcoding