refactor(qwen3): migrate prefix/offload path to openinfer-kv-store - #829
refactor(qwen3): migrate prefix/offload path to openinfer-kv-store#829xiaguan wants to merge 2 commits into
Conversation
qwen3's prefix-restore and KV offload now ride openinfer-kv-store
end to end; the crate drops its openinfer-kv-cache / openinfer-kv-offload
deps, and the vLLM P/D compat layer is deleted in the same cut.
- executor: Option<OffloadEngine> -> Option<Arc<KvStore>>; the prefetch
state machine (PrefetchState/PrefetchPhase, begin/drain/wait_kv_prefetch,
remote_fetch.rs, scheduler loading queue) is deleted outright — no
compat switch. seal/retire/flush ride KvStore; the flush-on-finish
Finished barrier becomes store.flush_saves_then.
- scheduler: ResolveHub resolves prefixes on the store runtime at intake
(CacheScope{lora}, default policy, TokenSink as CancelProbe; echo
requests skip resolve). KvPrefix RAII hold covers queue time and drops
at the first PrefillRequestResult; admission credits
resolved_prefix_blocks (hit_tokens/block_size) and excludes
store.pinned_blocks(0) from the budget. Both submit variants wired.
- kv-store absorbs the physical KV layer (KvBuffer/KvLayout/KvView/
KvCacheManager/BlockPool-with-events) and gains the embedding API
(runtime_handle, set_resolve_enabled / set_l1_retention_disabled for
the runtime prefix-cache toggles, flush_saves_then/blocking).
- vLLM P/D compat deleted: vllm_compat state machine, VllmBlockHasher
(xxh3_128 CBOR chain), miss_wait window, page_first registration,
Qwen3VllmCompatOptions, --kv-pd-* server flags. PagedKvPageFirst in
kernels is an unrelated LayoutTag and stays.
- tests: kv_offload_cpu_hit rewritten to the new two-act flow
(seal -> flush -> evict -> resolve_prefix -> prefill match; logits
tolerances unchanged). docs synced: design.md migration status,
pegaflow-offload-integration rewired, pd-disaggregation-m2 superseded.
GPU gates (5090, Qwen3-4B) follow on this branch.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: xiaguan <751080330@qq.com>
5090 GPU gate 结果(Qwen3-4B,/data/Qwen3-4B,GPU4,sm_120,2b55e031)
其余验收项:
|
![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: 2b55e03186
ℹ️ 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".
| let fresh_needed = footprint.saturating_sub(resolved_prefix_blocks(&req, block_size)); | ||
| if fresh_needed <= budget && decode_slots > 0 { |
There was a problem hiding this comment.
When several CPU-tier hits resolve before admission, every KvPrefix hold removes its pages from available_blocks(), but this loop credits only the request currently being considered while leaving all other deferred holds pinned. For example, in a 100-page pool, two requests with 80-page lifetime footprints and 40-page resolved prefixes leave 20 pages free; each still needs 40 pages, so neither is admitted and neither hold is ever released. Bound concurrent reservations or release/defer later holds so at least the head request can make progress.
Useful? React with 👍 / 👎.
| self.runtime.spawn(async move { | ||
| if let Err(err) = tier.flush().await { |
There was a problem hiding this comment.
With chunked prefill and concurrent finishes, an earlier tier.flush() can drain one request's early-chunk save handle from PegaflowTier::pending_saves; when that request later finishes, this independently spawned flush sees only its remaining handles and can invoke then before the earlier barrier completes. Under --kv-p2p-flush-on-finish, the HTTP response can therefore advertise KV readiness while part of the prefix is not yet registered, producing a short peer hit. Serialize or chain barriers so each callback transitively covers handles drained by predecessors.
Useful? React with 👍 / 👎.
| if let Err(err) = tier.flush().await { | ||
| log::warn!("kv-store flush before callback failed: {err:#}"); | ||
| } | ||
| then(); |
There was a problem hiding this comment.
When P2P registration or a save task stalls, tier.flush().await has no timeout, so the callback documented as always running is never reached. Qwen3's P/D prefill role with --kv-p2p-flush-on-finish withholds the request's Finished event behind this callback, meaning a MetaServer or transfer stall can leave the client request hanging indefinitely instead of degrading to a peer recompute. Wrap the barrier in a bounded timeout and run then on timeout as well.
Useful? React with 👍 / 👎.
| // 2. Once idle, apply pending control commands before admitting newer | ||
| // generation requests that arrived behind them. | ||
| if active.is_empty() && deferred.is_empty() && prefilling.is_empty() { | ||
| drain_idle_control(&mut executor, &mut pending_control); |
There was a problem hiding this comment.
With LoRA plus KV offload enabled, a Generate using adapter A is moved out of all three checked queues while its prefix resolves. If an Unload(A) command arrives immediately afterward, this condition treats the scheduler as idle and unloads the adapter before the earlier generation request returns, after which that request is rejected as using an unknown adapter. Include resolve_hub.in_flight() == 0 in the idle condition so control commands cannot overtake previously submitted generation work.
Useful? React with 👍 / 👎.
| executor | ||
| .available_blocks() | ||
| .saturating_sub(resolve_hub.pinned_blocks()), |
There was a problem hiding this comment.
During KV offload saves, each KvBlockGuard keeps its slot active, and BlockManager::available_blocks() already excludes active slots; subtracting pinned_blocks() again double-charges every page being copied. Under sustained offload or near-capacity workloads this can defer otherwise feasible requests until D2H completes and materially reduce throughput. Pass the pool's available count directly, or change the pin metric to include only capacity not already reflected there.
Useful? React with 👍 / 👎.
CI clippy (sm_80, -D warnings) flags or_insert_with(SaveCursor::new); the local pre-commit clippy only lints default-members, where qwen3 is built as a capped-lint dependency, so this only surfaces in CI. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>
|
收尾:glm52 编译回归通过(本地,NCCL 2.30.4, |
1 participant
概要
qwen3 的前缀恢复 / KV offload 全路径迁移到
openinfer-kv-store,crate 不再依赖openinfer-kv-cache/openinfer-kv-offload;vLLM P/D 兼容层并入本 PR 删除(issue #828 延期决策表中的切割项全部落地)。+1163 / −1794,42 文件。「single source of orchestration」落地:resolve 编排从 qwen3 手写的 executor 三相状态机(prefetch RemoteFetch/Loading/Committed)收敛为 kv-store 的
resolve_prefix,scheduler 侧一个ResolveHub统一 plain 与 LoRA-control 两条 submit 通道。关键改动
Option<OffloadEngine>→Option<Arc<KvStore>>;prefetch 状态机(begin/drain/wait_kv_prefetch、remote_fetch.rs、schedulerloading队列)整体删除,无开关硬切;seal/retire/flush 直调 store(flush-on-finish 屏障 =flush_saves_then)。executor 构建 PegaflowHost + 逐层 ArenaSpec(几何与 kv-offloadfrom_buffer逐字段对齐)。store.resolve_prefix上 store runtime(CacheScope{lora},TokenSink 作 CancelProbe,echo 跳过);KvPrefixRAII hold 护队列期,首个 PrefillRequestResult 释放;admission 一本账——resolved_prefix_blocks(hit_tokens/block_size)抵扣 need、pinned_blocks(0)扣 available,双重计数堵死。--kv-pd-*flags 全删。kernels 的PagedKvPageFirst是无关 LayoutTag,保留。kv_offload_cpu_hit重写为新两幕(execute_prefill+retire→seal → flush_offload_saves → evict_cached_blocks → resolve_prefix 取回 → prefill 自动 match,REGRET/MEAN 容差不变);scheduler 单测 Fake → 真 tier-less KvStore。验证
本地(5070 Ti):
cargo check --release -p openinfer-qwen3 --all-targets净cargo test --release -p openinfer-qwen3 --lib61/61(基线 69 − 随状态机删除的 remote_fetch 8 个)cargo test --release -p openinfer-kv-store全绿;cargo tree -p openinfer-qwen3无 kv-cache/kv-offload(dev 侧也零)rg vllm_compat|VllmBlockHasher|miss_wait|page_first源码零命中GPU gate(5090,Qwen3-4B)在本分支跟进:build + hf_golden_gate + prefix_cache + kv_offload_cpu_hit + glm52 编译回归。
Ref #828(映射表/延期决策)、#824。