feat: add DeepSeek V4 L2 KV cache offload and perf optimize.#620
feat: add DeepSeek V4 L2 KV cache offload and perf optimize.#620SimonCqk wants to merge 7 commits into
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 44005f44d5
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 771d911ca1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
SimonCqk
force-pushed
the
feat/ds-v4-l2-offload
branch
2 times, most recently
from
eb78fa6 to
19fcd96
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 19fcd96726
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7b52a25c91
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| raise ValueError( | ||
| f"draft_device_pool only supports DSA, MHA and MLA, " | ||
| f"got {type(actual_draft_pool)}" | ||
| if isinstance(actual_pool, DeepseekV4TokenToKVPool): |
There was a problem hiding this comment.
When the target is V4, this branch never handles draft_device_pool. V4 MTP keeps target and draft KV in independent tensors but feeds both the same scheduler page IDs, so after a host-prefix loadback only target state is restored and the draft reads stale/reused KV. Target verification protects committed tokens, but MTP acceptance can collapse. Please mirror/copy draft tensors in each paged transfer (and budget them), or reject V4 KVStore + MTP explicitly.
There was a problem hiding this comment.
Startup now rejects KVStore when grouped paged caches and a speculative draft KV pool coexist, before host-pool allocation. --disable-kvstore, configurations without a draft pool, and non-paged KVStore draft paths remain allowed. The behavior is covered by the parametrized grouped-KVStore draft-pool test.
And MTP with enabled L2 kvstore will be in next coming follow up seperate PR.
Implements the first complete host-side L2 offload path for DeepSeek V4's group-paged KV cache while preserving the existing scheduler overlap model. Key changes: - Add DeepseekV4TokenToKVPoolHost and DeepseekV4CachePool to mirror SWA KV, compressed KV, compressor state, and CSA indexer/compressor state in host memory with per-group page accounting. - Extend HybridPrefixCache, scheduler operations, FSM events, and Python bindings with paged-cache host snapshots, writeback/loadback plans, ack/failure handling, host eviction, and prefix-reuse metadata. - Wire MemoryExecutor, HostExecutor, event_loop, and scheduler_utils so enabled kvstore submits DeepSeek V4 L2 work through cache streams and LayerDoneCounter, while --disable-kvstore keeps the old path lean. - Add low-frequency paged L2 transfer summaries plus design docs and focused tests for scheduler lifecycle, runtime pool mapping, host executor dispatch, and disable-kvstore behavior. Signed-off-by: SimonCqk <cqk0100@gmail.com>
Reduces the host-side overhead of DeepSeek V4 L2 cache traffic in capacity-bound prefix-reuse scenarios without changing scheduler forward or overlap semantics. Key changes: - Coalesce scheduler paged-cache fragments into prepared group-level transfers so HostExecutor submits merged page ranges instead of repeatedly rebuilding the same plan per layer. - Chunk H2D loadback by layer range and keep LayerDoneCounter completion at layer granularity, allowing forward consumers to preserve the existing loadback synchronization contract. - Add a TVM FFI-safe pointer-table H2D scatter helper under tokenspeed-kernel, with Python fallback and tests for symbol-missing, threshold, bucketed pointer tables, and CUDA smoke coverage. - Keep D2H writeback on the direct path while recording low-frequency transfer summaries and backend statistics for pages, spans, tensor refs, effective copy calls, bytes, buckets, and kernel launches. Signed-off-by: SimonCqk <cqk0100@gmail.com>
Reduces cached-prefill latency spikes in the DeepSeek V4 HCA path by shrinking cache-insert work to active rows and avoiding per-request Triton specialization churn. Key changes: - Compute active HCA compressed rows from CPU-side prefill metadata so cached-prefill tails only materialize rows that can produce new compressed cache entries. - Add a direct HCA compress-and-cache-insert path that uses active token indices and avoids the older dense paged materialization path for the optimized case. - Keep the paged insert path as a correctness reference in tests while validating active-row and direct cache insert outputs against dense/paged behavior. - Collapse cold Triton JIT variants by making per-request block-table widths and cache block sizes runtime scalars or non-alignment-specialized arguments where they do not affect code shape. Signed-off-by: SimonCqk <cqk0100@gmail.com>
Allows PrefixReuse to keep a valid host-side history-family hit when continuation-only state groups are absent, while preserving exact terminal-state requirements for StateRecovery. Key changes: - Remove the early rejection that treated missing transport-only continuation state as a complete host-hit miss in the PrefixReuse path. - Keep StateRecovery stricter: exact terminal continuation state is still required before using a host state-recovery hit. - Add a C++ scheduler regression test showing that a host snapshot missing the continuation-only state group still returns a history-family host prefix hit with only replayable group pages. Signed-off-by: SimonCqk <cqk0100@gmail.com>
Keeps the L2 offload tests focused on stable interfaces and removes test-only dead work that was keeping shallow helper paths alive. Key changes: - Replace broad paged-pool tests with narrower prepared-transfer, prepared writeback, prepared loadback, and failure-path cases so tests no longer depend on writeback_paged/loadback_paged wrappers. - Move the Python binding shape smoke test out of the prefix metadata fixture class to avoid unrelated scheduler setup for that case. - Consolidate repeated MemoryExecutor fake cache op declarations and keep the HostExecutor tests centered on observable dispatch, ack, and LayerDoneCounter behavior. - Preserve focused coverage for HCA active rows, scheduler cache-event aggregation, prefix-cache metadata, and kvcacheio scatter fallback/bucket behavior. Signed-off-by: SimonCqk <cqk0100@gmail.com>
Fix the runtime ruff failures by using collections.abc.Iterable and dropping an unused logging import. Normalize KV prefix-cache lock targets to the nearest node that actually owns the requested KV resource. Paged-cache adjunct matches may extend MatchResult nodes deeper than KV device/host resources, and locking those deeper nodes can dereference missing resources during prefill or retracted decode recovery. Preserve request-local paged-cache tables when a retracted request already has them, so late forward recovery does not release the tables before rebuilding decode state. Update the stale history-only replay expectation to match the intended fallback behavior. Signed-off-by: SimonCqk <cqk0100@gmail.com>
SimonCqk
force-pushed
the
feat/ds-v4-l2-offload
branch
from
7b52a25 to
d8ed85f
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d8ed85ff92
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
SimonCqk
force-pushed
the
feat/ds-v4-l2-offload
branch
from
d8ed85f to
f554576
Compare
Reuse stream-affine H2D scatter metadata, lazily pin transfer indices, and keep pruning and allocation deterministic without adding CUDA synchronization. Preserve borrowed snapshots across demotion and radix splits, require exact continuation state for reusable prefixes, and abort one deterministic victim only after in-flight forward and cache work is retired. Pin paged-loadback host snapshots until the final layer H2D completion and retire them through the existing LoadBackDone lifecycle without changing ordinary radix KV or Mamba loadbacks. Deliver scheduler-owned aborts to runtime clients, retain generic retract capacity under --disable-kvstore, and reject unsupported paged draft-pool KVStore configurations. Remove dead diagnostic and test surfaces while retaining focused transfer, binding, lifecycle, and failure-path coverage. Signed-off-by: SimonCqk <cqk0100@gmail.com>
SimonCqk
force-pushed
the
feat/ds-v4-l2-offload
branch
from
f554576 to
d5e7b6f
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d5e7b6fb88
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
2 participants
Summary
This PR adds DeepSeek V4 L2 KV cache offload support for host-side CPU KV cache reuse.
Main changes:
Performance Setup
For the benchmark comparison we use
dp=1, tp=4.Reason: each DP rank owns its own device/host KV cache pool, so
dp=4shards requests across independent cache pools and can dilute prefix reuse.dp=1, tp=4keeps reuse in one KV pool while still using tensor parallelism.To construct a capacity-bound scenario, the server is started with:
Baseline without kvstore:
L2 kvstore enabled:
Overall performance comparison:
Test Regression