From 40d56ffdff7147eb429d82938d2e172a02c82b92 Mon Sep 17 00:00:00 2001 From: wangke <364517893@qq.com> Date: Thu, 6 Aug 2026 11:55:32 +0000 Subject: [PATCH] feat(qwen35): add joint prefix cache Unify Qwen3.5 KV and recurrent-state lifecycle across single-GPU and TP serving, accuracy tests, and benchmarks. Add prefix-cache coverage and record the measured performance results. Signed-off-by: wangke <364517893@qq.com> --- Cargo.lock | 1 + docs/models/qwen35/prefix-cache.md | 290 +++++--- docs/models/qwen35/tp-implementation.md | 20 +- .../src/integrations/scheduled.rs | 15 +- pegainfer-kv-cache/src/manager.rs | 15 + pegainfer-kv-cache/src/pool.rs | 125 +++- pegainfer-qwen35/Cargo.toml | 5 + pegainfer-qwen35/src/batch_decode.rs | 153 ++-- pegainfer-qwen35/src/batch_decode_graph.rs | 21 +- pegainfer-qwen35/src/decode_buffers.rs | 19 +- pegainfer-qwen35/src/executor.rs | 118 ++- pegainfer-qwen35/src/lib.rs | 21 +- pegainfer-qwen35/src/prefill.rs | 188 ++--- pegainfer-qwen35/src/prefix_cache.rs | 677 ++++++++++++++++++ pegainfer-qwen35/src/recurrent_state.rs | 20 + pegainfer-qwen35/src/scheduler.rs | 670 ++++++++++++----- pegainfer-qwen35/src/scheduler/tests.rs | 14 +- pegainfer-qwen35/src/tp_executor.rs | 579 +++++++++++++-- pegainfer-qwen35/src/unified_forward.rs | 215 +++--- pegainfer-qwen35/src/weights.rs | 70 +- pegainfer-qwen35/tests/hf_golden_gate.rs | 21 +- pegainfer-qwen35/tests/prefix_cache.rs | 318 ++++++++ pegainfer-server/src/bin/bench_serving/cli.rs | 7 +- .../src/bin/bench_serving/main.rs | 35 +- pegainfer-server/src/config.rs | 55 ++ pegainfer-server/src/main.rs | 1 + 26 files changed, 2982 insertions(+), 691 deletions(-) create mode 100644 pegainfer-qwen35/src/prefix_cache.rs create mode 100644 pegainfer-qwen35/tests/prefix_cache.rs diff --git a/Cargo.lock b/Cargo.lock index b9a6dcaa3..4a581fad0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3693,6 +3693,7 @@ dependencies = [ "log", "pegainfer-core", "pegainfer-kernels", + "pegainfer-kv-cache", "pegainfer-sample", "pegainfer-vllm-frontend", "rand 0.10.1", diff --git a/docs/models/qwen35/prefix-cache.md b/docs/models/qwen35/prefix-cache.md index 84572ee46..a9723157f 100644 --- a/docs/models/qwen35/prefix-cache.md +++ b/docs/models/qwen35/prefix-cache.md @@ -1,78 +1,76 @@ # Qwen3.5-4B prefix cache -> **TL;DR:** A Qwen3.5-4B prefix hit is valid only when full-attention KV and a complete recurrent/conv snapshot exist at the same 256-token boundary. `Qwen35PrefixCache` checks and restores both together, so the scheduler sees either one valid hit or a miss. The first version keeps snapshots on GPU and requires Qwen3.5 KV to move from `KvPool`/`KvState` to the content-hashed `BlockPool`/`RequestKv` cache. +> **TL;DR:** Qwen3.5-4B uses a GPU-only, content-hashed joint prefix cache: full-attention KV is reusable only with a matching complete recurrent/conv snapshot at the same 256-token boundary, otherwise the request is cold. `Qwen35PrefixCache` manages each reusable boundary as one `PrefixEntry`, while `PrefixCacheState` owns lookup, pinning, LRU, and publication. TP1/TP2 correctness and serving tests pass, and a 4,160-token shared prefix cuts warm TTFT by 94.3% (TP1) / 95.1% (TP2). > -> **Last touched:** 2026-07 +> **Last touched:** 2026-08 ## Preparation - **Read**: - `docs/index.md` - identified the current Qwen3.5 roadmap and the related Qwen3 cache and scheduler docs. - - `docs/models/qwen35/roadmap.md` - direct-paged writes and bounded chunked prefill are complete; issue #257 now needs a joint KV/recurrent/conv cache design. + - `docs/models/qwen35/roadmap.md` - direct-paged writes and bounded chunked prefill are complete; issue #257 needs a joint KV/recurrent/conv cache. - Maintainer RFC discussion for issue #257 - narrowed the first version to a GPU snapshot cache with one consistency rule for KV and recurrent state. - `docs/models/qwen3/prefix-cache.md` - provides the existing rules for block hashes, adapter isolation, final-token recompute, and keeping matched KV alive. - - `docs/subsystems/runtime/qwen3-kvbm-integration-spec.md` - describes the content-hashed `BlockPool`/`RequestKv` cache that Qwen3.5 does not yet use. - - `pegainfer-qwen35/src/{scheduler.rs,prefill.rs,prefill_buffers.rs,recurrent.rs,recurrent_state.rs,weights.rs}` - confirmed the current request state flow, valid prefill boundaries, snapshot layout, and GPU memory reservation. - - `pegainfer-core/src/kv_pool.rs` and `pegainfer-kv-cache/src/pool.rs` - confirmed that Qwen3.5 still uses anonymous RAII pages while Qwen3 can register, match, and pin content-hashed blocks. + - `docs/subsystems/runtime/qwen3-kvbm-integration-spec.md` - describes the content-hashed `BlockPool`/`RequestKv` cache adopted by Qwen3.5. + - `pegainfer-qwen35/src/{scheduler.rs,prefill.rs,recurrent_state.rs,weights.rs}` and `pegainfer-kv-cache/src/pool.rs` - confirmed request-state flow, snapshot layout, exact-boundary KV operations, and GPU memory reservation. - **Relevant history**: - The first draft focused on CPU offload but did not say clearly who keeps KV and snapshots consistent. Review narrowed the first version to GPU allocation, lookup, lifetime, and whole-model snapshot creation. - - The first draft also treated the 64-token GDR tile as a correctness boundary. Current resumed-prefill coverage uses 16-token scheduler chunks successfully, so a completed whole-model chunk, not an internal GDR tile, is the state boundary. + - The first draft also treated the 64-token GDR tile as a correctness boundary. Resumed prefill works with smaller scheduler chunks, so the safe boundary is a completed whole-model window, not an internal GDR tile. - **Plan**: - 1. Replace the two-tier proposal with a GPU-only cache that restores KV and recurrent state together. - 2. Base snapshot creation and restore on the current chunked-prefill flow and the future `BlockPool`/`RequestKv` migration. - 3. Define snapshot contents, capacity, publication order, lookup, pinning, eviction, failure behavior, and follow-on validation. + 1. Add exact-boundary, non-mutating KV probe and attach support. + 2. Unify Qwen3.5 execution around `KvCacheManager`/`RequestKv` transactions before enabling reuse. + 3. Add a fixed-budget recurrent/conv snapshot pool, joint restore/publication, TP coordination, observability, and validation. - **Risks / open questions**: - - The current `RequestKv::match_and_add_prefix` immediately changes a new request to use the longest KV-only match. Qwen3.5 instead needs the exact-boundary lookup and attach behavior described below. - - Snapshot copy and cold-prefill costs have not been measured for the current 32-value-head 4B configuration. The initial 256-token interval is therefore a starting policy, not a proven optimum. + - Snapshot copy and cold-prefill costs vary by model shape. The initial 256-token interval is a starting policy, not a proven optimum. + - TP must publish and restore a snapshot only when every rank reports the same token boundary. ## Decisions The first version is deliberately narrow: - GPU-only recurrent snapshot allocator with a fixed load-time byte budget. -- One complete snapshot contains all 24 linear layers' f32 GDR state, bf16 conv state, and the token position. -- Snapshots are published every 256 prompt tokens. Non-aligned request ends are not cached in the first version. +- One complete snapshot contains all linear layers' f32 GDR state, bf16 conv state, and the token position. +- Snapshots are eligible every 256 prompt tokens. Non-aligned request ends apply normally but do not create snapshots. - A reusable boundary must have both registered full-attention KV and a GPU-resident recurrent snapshot for the same token prefix. -- `Qwen35PrefixCache` is the only interface used by the scheduler for prefix creation and restore. +- `Qwen35PrefixCache` is the only interface used by the scheduler for joint prefix creation and restore. - Restored state is copied into the request's own `RecurrentState`; active requests never modify or directly share a cache slot. -- Echo and prompt-logprob requests stay on cold prefill because cached positions would not produce their required logits. +- Echo requests stay on cold prefill because cached positions would not produce their required logits. CPU offload, a second LRU tier, snapshot compression, request-end snapshots, and cross-process transfer are deferred. They must keep the same rule that KV and recurrent state are restored together. -## Current state and prerequisite +## Implemented state flow -Qwen3.5 prefill already has a safe point at which it can create a snapshot. A `PrefillingRequest35` owns: +The scheduler/executor now owns content-hashed request KV and recurrent state together: ```rust -struct PrefillingRequest35 { - req: SchedulerRequest, - kv: KvState, - rec: RecurrentState, - cursor: usize, - step_chunk: usize, +enum PrefillBackendState { + Single { + kv: Box, + rec: RecurrentState, + }, + // TP workers address controller-owned RequestKv by request id. + Tp { request_id: RequestId }, } ``` -For each scheduled window, `prefill_chunk_forward` writes full-attention K/V directly into paged storage and advances all linear layers' recurrent and conv state. When the whole-model call returns successfully: +`RequestKv::schedule_prefill` reserves pages and produces an immutable `KvView`. Prefill and decode kernels write through that view without changing logical KV. After the whole-model call succeeds, the scheduler applies the KV transaction and may publish a recurrent snapshot. After each successful prefill window: ```text -kv.seq_len() == rec.seq_len == cursor + step_chunk +kv.kv_position() == rec.seq_len == cursor + step_chunk ``` -If the prompt is incomplete, the scheduler keeps these states for the next step. If it is complete, it copies `rec` into a stable decode graph slot. Direct-paged prefill and scheduler chunking are therefore no longer blockers. +Single-GPU serving, TP serving, and the low-level accuracy executor all use `KvCacheManager`/`RequestKv`, even when the snapshot budget is zero. With zero budget, release resets registered blocks so the disabled mode remains a true cold control. -The missing prerequisite is content-based KV reuse. Qwen3.5 still uses `pegainfer_core::kv_pool::{KvPool, KvState}`: it allocates and returns pages, but it cannot identify their token content, register completed blocks, or match a new request against them. Qwen3 uses `pegainfer_kv_cache::{BlockPool, RequestKv}`, which provides those operations and keeps matched blocks alive while they are being attached to a request. - -Before prefix reuse can be implemented, Qwen3.5 full-attention KV must move to that cache API while preserving its current page-first memory layout and kernels. Most required operations already exist. Joint lookup adds these requirements: +Joint lookup adds these exact-boundary operations to the shared cache: 1. Probe the longest contiguous registered KV prefix without changing the new request. -2. Keep the probed KV blocks pinned while the snapshot cache is checked. +2. Keep the probed KV blocks pinned while `PrefixCacheState` is checked. 3. Expose complete KV-block boundaries and their canonical `SequenceHash` values in descending order. -4. After a joint boundary is selected, transfer only the blocks through that boundary to the new request and release any longer KV-only tail. -5. Advance the new request's KV position to the selected boundary as part of the same attach operation; a partial attach must not be visible. +4. Attach only the boundary with a matching snapshot; ignore any longer KV-only tail. +5. Advance the new request's KV position as part of the same attachment. 6. Keep at least one prompt token uncached so prefill can produce the first generated token. -The existing `schedule_prefill`, `apply_prefill_chunk`, and `revert_schedule` operations then handle suffix prefill and failed forwards. +`TokenEvent::Scheduled.cached_tokens` reports the selected joint boundary. A KV-only tail is never reported as a hit. ## Valid cache hit @@ -80,24 +78,24 @@ Qwen3.5 stores KV and recurrent snapshots separately, but a request may reuse a 1. Full-attention KV for `[0, N)` is registered and still on GPU. 2. A complete recurrent snapshot for `[0, N)` is still on GPU. -3. Both use the same `SequenceHash`, which includes the token lineage and adapter/LoRA salt. +3. Both use the same `SequenceHash`, which includes token lineage and adapter/LoRA salt. 4. The KV position, snapshot position, recurrent `seq_len`, and scheduler cursor all equal `N`. KV without a matching snapshot is not a Qwen3.5 prefix hit. A snapshot without matching KV is also not a hit. The scheduler never sees either one as partial reuse. ## Snapshot interval -A snapshot boundary is the token position after a scheduled prefill window has completed all 32 layers. The interval for the first slice is: +A snapshot boundary is the token position after a scheduled prefill window has completed all model layers. The current interval is: ```text -SNAPSHOT_STRIDE = 256 tokens +SNAPSHOT_STRIDE_TOKENS = 256 ``` -The stride is a multiple of the 16-token KV block size, so every snapshot key can reference the lineage hash of a complete registered KV block. The scheduler must clamp each request's next window to the next snapshot boundary. With a 900-token prompt, the resulting positions are `256 -> 512 -> 768 -> 900`; only the first three are snapshot candidates. +The stride is a multiple of the 16-token KV block size, so every snapshot key references the lineage hash of a complete registered KV block. The scheduler clamps each request's next window to the next snapshot boundary. With a 900-token prompt, the resulting positions are `256 -> 512 -> 768 -> 900`; only the first three are snapshot candidates. -The GDR implementation internally tiles work in 64-token chunks, but this is not a snapshot correctness constraint. It handles a partial final tile and commits the final recurrent state for arbitrary positive sequence lengths. The existing scheduler-level resumed-prefill gate uses 16-token windows and exercises this behavior. The invariant is therefore "after a successful whole-model window," not `position % 64 == 0`. +The GDR implementation internally tiles work in 64-token chunks, but this is not a snapshot correctness constraint. It handles a partial final tile and applies recurrent state for arbitrary positive sequence lengths. The invariant is “after a successful whole-model window,” not `position % 64 == 0`. -The 256-token interval limits how many large snapshots one prompt creates. Prefixes shorter than 256 tokens intentionally remain cold. This is a starting policy, not a measured optimum; any later interval must still align to complete KV blocks. +The 256-token interval limits how many large snapshots one prompt creates. Prefixes shorter than 256 tokens intentionally remain cold. Any later interval must still align to complete KV blocks. ## Snapshot contents and capacity @@ -107,7 +105,7 @@ Each slot uses the same device layout as request-local `RecurrentState`: - for every linear layer, `conv_state: [linear_attn_qkv_dim, conv_kernel_dim - 1]` bf16; - host metadata recording the exact `seq_len` represented by the slot. -Capacity must be derived from `recurrent_state::bytes_per_request(config)`, not a hard-coded model label. For Qwen3.5-4B (24 linear layers, 16 key heads, 32 value heads, 128x128 state, conv kernel 4), one slot is: +Capacity is derived from `recurrent_state::bytes_per_request(config)`, not a hard-coded model label. For Qwen3.5-4B, one slot is: ```text per-layer GDR state = 32 * 128 * 128 * 4 = 2,097,152 bytes @@ -123,124 +121,150 @@ let bytes_per_slot = bytes_per_request(config); let max_slots = snapshot_budget_bytes / bytes_per_slot; ``` -The reservation participates in the same load-time budget as prefill scratch, recurrent request/decode slots, and KV pages. The current loader reserves two recurrent states per decode capacity slot before sizing KV; snapshot bytes are additional and must also be subtracted before the KV pool is allocated. Snapshot allocation must not opportunistically consume memory that admission assumes belongs to KV. +Snapshot bytes are reserved before KV capacity is finalized, so snapshot allocation cannot consume memory that admission assumes belongs to KV. Zero configured MiB disables prefix reuse. A positive budget smaller than one complete slot is rejected instead of silently acting disabled. -If the configured budget produces zero slots, Qwen3.5 prefix reuse is disabled and serving retains its current cold behavior. +Under TP the configured budget applies to each rank. Every rank must allocate the same number of physical slots. ## Cache ownership and pinning -`Qwen35PrefixCache` owns the existing full-attention KV manager and the recurrent snapshot cache. `KvCacheManager` keeps the logical `BlockPool` and physical GPU `KvBuffer` together: +`Qwen35PrefixCache` owns the full-attention KV manager and the joint-entry directory. Physical snapshot stores are separate so TP can use one metadata decision with identical slot numbers on every rank: ```rust struct Qwen35PrefixCache { kv: KvCacheManager, - snapshots: RecurrentSnapshotCache, - stride: usize, + state: PrefixCacheState, + enabled: bool, + stats: PrefixCacheStats, } -struct RecurrentSnapshotCache { - slots: Vec, - index: HashMap, - free: Vec, - lru: LruList, +struct PrefixCacheState { + entries: HashMap, + free_slots: Vec, + slot_count: usize, + clock: u64, } -struct SnapshotSlot { - state: RecurrentState, - key: Option, - pin_count: usize, +struct PrefixEntry { + recurrent_slot: usize, + kv_lease: Vec, + pin_count: Arc, + last_used: u64, } -#[derive(Clone, Hash, Eq, PartialEq)] -struct SnapshotKey { - sequence_hash: SequenceHash, +struct PrefixGuard { boundary: usize, + recurrent_slot: usize, + pin_count: Arc, + started: Instant, +} + +struct PrefixReservation { + key: PrefixBoundaryKey, + recurrent_slot: usize, + replaced: bool, +} + +struct RecurrentStateStore { + slots: Vec, +} + +struct PrefixBoundaryKey { + sequence_hash: [u8; 16], + boundary_tokens: usize, } ``` -`sequence_hash` is the canonical hash returned by the KV cache for the block ending at `boundary`. It already includes the earlier block lineage and adapter salt, so the snapshot cache does not maintain a second adapter identity. Storing `boundary` explicitly prevents a snapshot from being reused at the wrong token position. +`sequence_hash` is the canonical hash returned by the KV cache for the block ending at `boundary_tokens`. It already includes earlier block lineage and adapter salt. Storing `boundary_tokens` explicitly prevents reuse at the wrong token position. -There is no third stored copy that combines KV and snapshot data. A valid internal hit records the selected boundary and holds both a KV guard and a snapshot guard. `Qwen35PrefixCache` consumes those guards during restore, so neither resource can be evicted in the meantime. The scheduler never sees the separate guards. +There is no third stored tensor copy combining KV and recurrent state. Each `PrefixEntry` owns one physical recurrent-state slot and strong KV guards for every block from token `0` through its boundary. A selected hit returns a `PrefixGuard` that pins the entry until physical restore completes. TP publishes the prefix entry only after every worker saves the same-numbered slot at the same boundary; it reports a hit only after every worker restores and confirms that boundary. ## Creating a snapshot -A snapshot is created after a whole-model window, not inside per-layer GDR scratch. The order is: +A snapshot is created after a whole-model window, not inside per-layer GDR scratch: -1. Clamp the request's scheduled window to the next 256-token boundary or prompt end. -2. `RequestKv::schedule_prefill` reserves the KV blocks for that window. +1. Clamp the scheduled window to the next 256-token boundary or prompt end. +2. Reserve KV blocks with `RequestKv::schedule_prefill`. 3. Run the full model, updating full-attention KV and request-local recurrent/conv state. -4. On failure, revert the KV schedule, fail the request, and publish nothing. -5. On success, commit the KV request state (`apply_prefill_chunk` or final `apply_prefill`) so complete blocks are registered. -6. Assert that committed KV position and `rec.seq_len` equal the candidate boundary. -7. If the boundary is snapshot-eligible, allocate an unpinned slot and D2D-copy the complete `RecurrentState` into it. -8. Publish `SnapshotKey -> SlotId` only after the copy has been successfully enqueued under the scheduler stream's ordering contract. - -The GDR `chunk_state` scratch is per linear-layer call and cannot represent a whole-model snapshot. Conv state is updated separately. Only request-local `RecurrentState` after all layers have finished contains the complete pair required for publication. +4. On failure, revert the KV schedule and publish nothing. +5. On success, apply KV with `apply_prefill_chunk` or final `apply_prefill`. +6. Verify that KV position and `rec.seq_len` equal the candidate boundary. +7. At an eligible boundary, call `reserve_prefix`; if it returns a reservation, copy the complete `RecurrentState` into that slot. +8. Publish the `PrefixEntry` after all rank-local copies succeed; abort the reservation on failure. -Running out of snapshot slots is a soft cache event: skip insertion and continue the request. A CUDA copy failure is an execution error, not a cache-capacity miss; no key is published. +The GDR `chunk_state` scratch is per linear-layer call and cannot represent a whole-model snapshot. Only request-local `RecurrentState` after all layers finish contains the complete recurrent/conv state required for publication. -Insertion of an already-resident key reuses the existing immutable slot and refreshes its LRU position; it does not allocate or copy a duplicate. When replacing an unpinned victim, the cache removes the victim's old index entry before starting the copy. If that copy fails, the slot returns to the free list with no published key. +Running out of snapshot slots is a soft cache event: skip insertion and continue the request. A CUDA copy failure is an execution error and publishes no key. A duplicate key refreshes LRU without copying another immutable snapshot. If replacement copying fails, the reserved slot returns to the free list unpublished. ## Lookup and restore -The scheduler calls one operation: +Restore is two-phase so one directory can coordinate one or many physical ranks: ```rust -let cached_tokens = prefix_cache.restore_prefix( - &mut prefix_request, - &mut request_recurrent, +let (request_kv, restore) = prefix_cache.begin_request(...)?; +// Restore restore.recurrent_slot() on every physical rank. +let cached_tokens = prefix_cache.finish_restore( + &request_kv, + restore, + &rank_positions, )?; ``` -Inside `Qwen35PrefixCache`: +`begin_request` performs the logical lookup and KV attach: -1. Ask the KV cache for the longest registered prefix while keeping the candidate KV blocks alive. -2. Enumerate eligible 256-token boundaries in descending order, subject to the usual rule that at least one prompt token remains to run. -3. Build `SnapshotKey` from the candidate's canonical `SequenceHash` and position. -4. Try to pin the corresponding snapshot slot. -5. The first boundary with both guards becomes the selected hit; if none exists, return `0` without changing request state. +1. Probe the longest registered KV prefix while keeping candidate blocks alive. +2. Enumerate eligible 256-token boundaries from longest to shortest, leaving at least one prompt token to run. +3. Build `PrefixBoundaryKey` from the canonical `SequenceHash` and token position. +4. Pin the corresponding snapshot slot. +5. Select the first boundary with both resources; if none exists, return `0` without changing request state. 6. Attach exactly that KV boundary to request-local `RequestKv`. -7. D2D-copy the immutable snapshot into request-local `RecurrentState`. -8. Verify all positions equal the selected boundary, then set the scheduler cursor and return it as `cached_tokens`. -For example, a 768-token KV match with snapshots at 256 and 512 restores 512 tokens. The scheduler never receives "KV hit 768, snapshot hit 512" as separate facts. +The single-GPU or TP executor then restores `recurrent_slot` into the request-local state. `finish_restore` verifies all positions, records the hit, and releases the pin. + +For example, a 768-token KV match with snapshots at 256 and 512 restores 512 tokens. The scheduler never receives “KV hit 768, snapshot hit 512” as separate facts. -`Qwen35PrefixCache` must acquire both guards before changing the request. If restore then fails, it releases the request KV and snapshot guard and reports an error. It must not expose a partly restored request or treat a restore error as a normal cache miss. +If physical restore or position validation fails after KV attachment, request preparation fails and releases the prepared state. It is not treated as a normal cache miss. -After restore, suffix prefill operates normally on `tokens[cached_tokens..]`. When prefill finishes, the existing copy from request-local recurrent state into the decode graph slot remains unchanged. +After restore, suffix prefill operates on `tokens[cached_tokens..]`. When prefill completes, recurrent state is promoted into the normal decode state. Decode continues to schedule, forward, and apply KV one token at a time; it does not perform another prefix lookup. ## Lifetime and eviction -The KV and snapshot caches keep their own allocation policies, but `Qwen35PrefixCache` decides whether a boundary can be reused: +`Qwen35PrefixCache` owns the joint entry lifetime: -- KV candidates are held by strong immutable-block guards from probe until exact-boundary attachment or abandonment. +- Every request marks its assigned KV blocks to reset on release. Non-aligned prompt tails, decode-generated full blocks, and prefixes without a published snapshot therefore return to the free pool. +- A published snapshot's cache-owned KV lease holds strong immutable-block guards for every leading block through its boundary. Those are the only blocks retained after request release. - Snapshot slots are immutable while indexed and can be evicted only when `pin_count == 0`. -- A snapshot guard is needed only until its D2D copy into request-local state completes. It is not held for the full request lifetime. +- A snapshot guard is needed only until D2D restore and position checks complete. - Restored suffix prefill and decode mutate request-owned state, never the cached slot. - If no free or unpinned snapshot slot exists, insertion is skipped rather than blocking or failing the request. -Eviction does not require synchronized callbacks between the physical pools: +Replacing an entry removes the old `PrefixEntry` during reservation; dropping it releases the KV lease. If the replacement copy aborts, the new slot remains unpublished and returns to the free list. The guarded blocks then reset/free once no active request still references them, so snapshot LRU eviction cannot leave an inactive KV-only tail consuming cache capacity. Lookup can continue to the next shorter published prefix boundary. + +LRU selects an unpinned snapshot victim. Correctness depends on pinning and joint validation, not on LRU ordering. + +## TP behavior -- If KV is evicted first, the snapshot remains indexed but cannot be used. Lookup cannot acquire the KV guard, so snapshot LRU may reclaim it later. -- If the snapshot is evicted first, the KV blocks may remain reusable by the physical pool, but the boundary is KV-only and ineligible for Qwen3.5 restore. -- If either side disappears between candidate discovery and pinning, lookup continues to the next shorter joint boundary. +TP uses one logical cache decision and rank-local physical storage: -This cannot produce a partial hit: only `Qwen35PrefixCache` can declare a hit, and it checks both resources on every lookup. +1. The controller owns the only `KvCacheManager`, `RequestKv` map, `PrefixCacheState`, and LRU state. +2. Startup validates compatible KV geometry and snapshot-slot counts across ranks. +3. Logical capacity is capped by the smallest rank-local physical capacity. +4. The controller broadcasts identical `KvView` page IDs; every worker writes its local KV shard into its own `KvBuffer`. +5. Snapshot insertion reserves one common slot, saves it on every rank, verifies all returned positions, and only then publishes the key. +6. Restore uses the same slot on every rank and reports a hit only after every worker confirms the boundary. -LRU is the first victim policy for unpinned snapshot slots. Correctness depends on pinning and joint validation, not on LRU itself. +This keeps admission, attachment, and eviction deterministic across ranks while leaving recurrent/conv tensors local to each GPU. ## Correctness rules -- `RequestKv::kv_position() == RecurrentState::seq_len == SnapshotKey::boundary` after insertion and restore. +- `RequestKv::kv_position() == RecurrentState::seq_len == PrefixBoundaryKey::boundary_tokens` after insertion and restore. - A snapshot contains both state tensors for every linear layer; GDR-only or conv-only snapshots are invalid. - Snapshot contents are immutable after publication. -- KV and snapshot must use the same canonical `SequenceHash`. - A prefix hit always leaves at least one prompt token uncached so final prefill can emit the first generated token. -- Echo and prompt-logprob requests never use prefix matching in the first slice. -- Allocation pressure or no evictable snapshot slot changes hit rate only, not request output. -- Snapshot insertion failure never converts an otherwise valid cold request into a cache hit. +- Echo requests do not use prefix matching. +- Allocation pressure changes hit rate only, not request output. +- Failed forward or snapshot insertion never publishes a cache key. - A KV-only or snapshot-only boundary is never reported as cached tokens. -- Disabling the feature or configuring zero slots preserves current cold-serving behavior. +- Disabling the feature preserves cold-serving behavior. ## Implementation order @@ -255,18 +279,57 @@ Implementation is a follow-on to this design and should land in this order: ## Validation -The implementation acceptance surface should include: +Completed on an RTX 4090 with fixed-revision local Qwen3.5-4B weights: + +- `cargo test --release -p openinfer-kv-cache`: 19 passed. +- `cargo test --release -p openinfer-qwen35 --features qwen35 --lib`: 78 passed; 6 two-GPU tests remained ignored. +- `cargo test --release -p openinfer-server --features qwen35 config::tests::qwen35`: 9 passed. +- Prefix-cache coverage proved cold `0`, warm `256`, exact 256-token alignment, 512/576 boundary selection, prefix extension, echo bypass, multi-token logprob parity, LRU refresh/eviction, pinned-slot soft misses, cache-owned KV lease release on eviction, KV-only fallback to `0`, and output stability. +- Scheduler e2e and resumed chunked prefill passed. +- Short and long TP1/TP2 HF logits gates passed after `RequestKv` unification. TP2 short sequential mean/p99 delta was `0.0251/0.0983`; long sequential was `0.0215/0.0684`. +- Release clippy passed for the Qwen3.5 library/tests and corresponding server targets. +- TP2 scheduler, HTTP serving, and the full ignored package run passed on physical GPUs 1/2. + +## Performance Result (2026-08-06) + +- **Environment:** local Qwen3.5-4B, RTX 4090s only (GPU 1 for TP1; GPUs 1/2 for TP2). TP1 used CUDA Graphs; TP2 used `--cuda-graph=false`. + +Qwen3.5 reuses the largest 256-token boundary strictly below the prompt length. This table compares cold and warm TTFT p50 for single-token generation (`cache off -> cache on`). + +| Prompt tokens | Cached tokens | TP1 | TP1 reduction | TP2 | TP2 reduction | +| ---: | ---: | ---: | ---: | ---: | ---: | +| 320 | 256 | 29.22 -> 15.56 | 46.7% | 46.18 -> 18.75 | 59.4% | +| 576 | 512 | 46.85 -> 16.92 | 63.9% | 72.97 -> 18.82 | 74.2% | +| 1,088 | 1,024 | 88.75 -> 15.93 | 82.1% | 126.14 -> 19.36 | 84.7% | +| 2,112 | 2,048 | 161.00 -> 16.40 | 89.8% | 230.07 -> 19.82 | 91.4% | +| 4,160 | 4,096 | 308.64 -> 17.71 | 94.3% | 439.37 -> 21.68 | 95.1% | + +This table shows how the same cache hit affects TTFT and end-to-end latency when generating 128 tokens; decode dominates the remaining latency. + +| Prompt tokens | TP1 TTFT p50 off -> on | TP1 E2E p50 off -> on | TP2 TTFT p50 off -> on | TP2 E2E p50 off -> on | +| ---: | ---: | ---: | ---: | ---: | +| 1,088 | 94.58 -> 24.13 | 1,542.02 -> 1,452.78 | 127.05 -> 20.38 | 1,445.78 -> 1,324.97 | +| 2,112 | 166.25 -> 24.23 | 1,700.09 -> 1,535.84 | 231.56 -> 21.89 | 1,640.17 -> 1,421.95 | +| 4,160 | 313.51 -> 24.30 | 2,006.08 -> 1,696.04 | 442.01 -> 24.88 | 2,036.57 -> 1,606.08 | + +This table summarizes behavior under concurrency, mixed load, and decode batching: + +| Additional workload | Result | +| --- | --- | +| TP1 concurrency, 2,112 prompt + 128 output | At concurrency 1/4/8: TTFT p50 `24.67/94.41/152.25 ms`; request throughput `83.11/75.87/69.04 tok/s`; every request hit 2,048 cached tokens. | +| TP1 mixed load | Baseline ITL p50/p99 `12.03/12.25 ms`; mixed ITL `12.03/36.12 ms`. After initial insertion, 4,096-token injections hit 3,840 tokens with `41.42--51.00 ms` prefill and no warnings. | +| Decode TP1 vs TP2 | At context 4,096/batch 1, TP1/TP2 TPOT is `13.08/12.47 ms`; at batch 4 it is `13.74/48.76 ms`. TP2 batch decode needs a separate runtime optimization pass; this run disables CUDA Graphs. | + +**Conclusion** -- cache-management tests for snapshot keys, shorter-boundary fallback, duplicate insertion, pinning, and unpinned eviction; -- scheduler tests for prompts below 256, exactly on boundaries, across multiple boundaries, and with non-aligned tails; -- availability tests for KV-only, snapshot-only, and shorter valid fallback; -- adapter salt isolation; -- mixed cold and warm requests in the same prefill/unified step; -- pool-full behavior proving insertion skip preserves cold output; -- real GPU cold-vs-warm HF logits gates, including resumed suffix prefill and decode-slot promotion; -- retained metrics for snapshot D2D copy time, cold insertion overhead, warm TTFT, joint hit length, and slot occupancy. +- **Core performance gains** + - Warm TTFT improves with prefix length: at 4,160 prompt tokens it falls by 94.3% on TP1 and 95.1% on TP2. + - For 128-token outputs, the same long prompt reduces E2E latency by 15.4% on TP1 and 21.1% on TP2; steady decode TPOT is effectively unchanged. + - The benefit remains under TP1 concurrency and mixed load; warm injections hit 3,840 tokens with 41--51 ms prefill and no warnings. + - TP1 HTTP warm TTFT remains about 19--29 ms for 320--4,160-token prompts with cache enabled. +- **Follow-up** + - TP2 batch-4 decode has high TPOT while CUDA Graphs are disabled; profile and optimize this runtime path separately from prefix cache. -Those measurements determine whether 256 remains the right stride. They must not be replaced by the old RTX 4090 CPU-transfer estimates, which measured a deferred design and a different snapshot shape. ## Deferred work @@ -275,7 +338,6 @@ Those measurements determine whether 256 remains the right stride. They must not - workload-adaptive or per-model snapshot stride; - snapshot compression or reduced-precision state; - cross-worker/P-D transfer of hybrid state; -- integration with speculative rollback state; - sharing snapshot infrastructure across other hybrid model lines. -Each of these must preserve the same logical rule: one reusable boundary restores all model state at one token position. +Each extension must preserve the same logical rule: one reusable boundary restores all model state at one token position. diff --git a/docs/models/qwen35/tp-implementation.md b/docs/models/qwen35/tp-implementation.md index 66d35cd26..599fb8d92 100644 --- a/docs/models/qwen35/tp-implementation.md +++ b/docs/models/qwen35/tp-implementation.md @@ -2,7 +2,7 @@ > **TL;DR:** Qwen3.5 TP Phase 1 is implemented as correctness-first eager dense TP: TP2 worker/scheduler execution, short/long HF logits gates, scheduler e2e, and real OpenAI-compatible HTTP serving smoke pass. The branch is rebased onto current `main` with the newer engine, sampling, config, and golden-fixture contracts; remaining TP work is tracked as follow-up, not a Phase 1 claim. > -> **Last touched:** 2026-07 +> **Last touched:** 2026-08 ## Scope @@ -39,6 +39,24 @@ Not implemented in Phase 1: - Prefix-cache or recurrent-state snapshot support. - Performance claims. +## Post-Phase 1 Follow-up: RequestKv and Joint Prefix Cache + +This follow-up unifies the TP and single-GPU request lifecycle around `RequestKv` and adds joint full-attention KV plus recurrent/conv prefix reuse. + +1. **Unified KV lifecycle** + - The controller uses one `KvCacheManager` for prefill/decode scheduling and commit. + - Immutable `KvView`s carry the logical page ids to every worker; each rank writes its local KV shard into its own `KvBuffer`. +2. **TP capacity and layout** + - Startup validates identical KV geometry and snapshot-slot counts across ranks. + - The logical pool is capped by the smallest rank-local physical capacity. +3. **Joint recurrent snapshots** + - Key, pin, and LRU metadata is centralized; recurrent/conv tensors remain rank-local. + - Publication reserves one common slot, saves it on every rank, verifies the committed boundary, and only then publishes the key. + - Restore follows the same all-rank rule and reports a hit only after every rank confirms the selected boundary. +4. **Opt-in budget** + - `--qwen35-prefix-cache-mib` reserves the snapshot budget independently on each rank. + - `0` keeps cold serving. + ## Important Fixes ### Gated q projection layout diff --git a/kvbm/kvbm-logical/src/integrations/scheduled.rs b/kvbm/kvbm-logical/src/integrations/scheduled.rs index f4bfde5aa..d7f5e690c 100644 --- a/kvbm/kvbm-logical/src/integrations/scheduled.rs +++ b/kvbm/kvbm-logical/src/integrations/scheduled.rs @@ -436,11 +436,24 @@ impl SchedulableSequence { pub fn match_and_add_prefix( &mut self, manager: &BlockManager, + ) -> Result { + self.match_and_add_prefix_up_to(manager, usize::MAX) + } + + /// Match and add at most `requested_max_blocks` prefix blocks. + /// + /// This is useful for hybrid models whose auxiliary state + /// may only be restorable at a boundary shorter than the longest KV hit. + pub fn match_and_add_prefix_up_to( + &mut self, + manager: &BlockManager, + requested_max_blocks: usize, ) -> Result { self.require_idle()?; let bs = self.inner.block_size(); - let max_blocks = self.inner.num_input_tokens().saturating_sub(1) / bs; + let max_blocks = + (self.inner.num_input_tokens().saturating_sub(1) / bs).min(requested_max_blocks); let count = self .inner .match_and_add_prefix(manager, max_blocks) diff --git a/pegainfer-kv-cache/src/manager.rs b/pegainfer-kv-cache/src/manager.rs index 8564b72d9..5f31adbb6 100644 --- a/pegainfer-kv-cache/src/manager.rs +++ b/pegainfer-kv-cache/src/manager.rs @@ -38,6 +38,21 @@ impl KvCacheManager { Ok(Self { pool, buffer }) } + /// Pair an existing physical KV buffer with a new logical block pool. + /// + /// `num_blocks` may be smaller than the physical allocation. Tensor-parallel + /// executors use this to choose the minimum common logical capacity across + /// rank-local buffers while keeping one shared page-id namespace. + pub fn from_buffer(buffer: KvBuffer, num_blocks: usize) -> anyhow::Result { + anyhow::ensure!( + num_blocks <= buffer.num_blocks(), + "logical KV block count {num_blocks} exceeds physical buffer capacity {}", + buffer.num_blocks() + ); + let pool = BlockPool::new(buffer.layout().page_size, num_blocks)?; + Ok(Self { pool, buffer }) + } + /// Like [`new`](Self::new) but the pool emits KV block events; returns the /// receiver to drain. See [`BlockPool::with_events`]. pub fn new_with_events( diff --git a/pegainfer-kv-cache/src/pool.rs b/pegainfer-kv-cache/src/pool.rs index 4848795a7..6f554ce54 100644 --- a/pegainfer-kv-cache/src/pool.rs +++ b/pegainfer-kv-cache/src/pool.rs @@ -210,6 +210,7 @@ impl BlockPool { seq_hashes, gpu_hit, cacheable, + block_size: self.block_size, held: gpu_guard, } } @@ -253,6 +254,8 @@ pub struct PrefixProbe { gpu_hit: usize, /// Reuse cap: blocks past this are never matched (the final chunk forwards). cacheable: usize, + /// Tokens represented by one complete KV block. + block_size: usize, /// Strong refs keeping matched/loaded blocks resident until prefill. held: Vec>, } @@ -271,6 +274,37 @@ impl PrefixProbe { self.held.len() } + /// Complete prefix blocks eligible for request reuse. + /// + /// This is capped by the final-token rule even if a caller extended the + /// probe with additional loaded blocks. + pub fn reusable_blocks(&self) -> usize { + self.held.len().min(self.cacheable) + } + + /// Returns the lineage hash identifying the complete reusable prefix ending + /// at `boundary_tokens`. + /// + /// `boundary_tokens` is measured in tokens. For example, + /// with a 16-token block size, `boundary_hash(32)` identifies the token + /// prefix `[0, 32)`. The returned hash covers the full prefix lineage, + /// rather than only the contents of the final block. + /// + /// Returns `None` when the boundary is zero, is not block-aligned, or + /// exceeds the reusable prefix. + pub fn boundary_hash(&self, boundary_tokens: usize) -> Option<[u8; 16]> { + if boundary_tokens == 0 || !boundary_tokens.is_multiple_of(self.block_size) { + return None; + } + let block_count = boundary_tokens / self.block_size; + if block_count > self.reusable_blocks() { + return None; + } + self.seq_hashes + .get(block_count - 1) + .map(sequence_hash_bytes) + } + /// Content hashes to query the CPU tier with: the blocks past the GPU hit, /// capped at the reuse boundary. Empty when the GPU hit already covers /// every reusable block (nothing to load — prefill normally). @@ -357,10 +391,23 @@ impl RequestKv { /// Matching always leaves at least one prompt token uncached so the /// final prefill chunk can emit the first generated token. pub fn match_and_add_prefix(&mut self, pool: &BlockPool) -> anyhow::Result { + self.match_and_add_prefix_up_to(pool, usize::MAX) + } + + /// Match and attach no more than `max_blocks` of the resident prefix. + /// + /// The underlying sequence still enforces the final-token cap. A caller + /// can hold a [`PrefixProbe`] while invoking this method to ensure the + /// selected blocks remain resident between joint-state lookup and attach. + pub fn match_and_add_prefix_up_to( + &mut self, + pool: &BlockPool, + max_blocks: usize, + ) -> anyhow::Result { let blocks = self .seq - .match_and_add_prefix(&pool.block_manager) - .map_err(|e| anyhow::anyhow!("match_and_add_prefix: {e}"))?; + .match_and_add_prefix_up_to(&pool.block_manager, max_blocks) + .map_err(|e| anyhow::anyhow!("match_and_add_prefix_up_to: {e}"))?; // Prefix-hit blocks are already in the router's tree (whoever first // sealed them stored them, and a GPU hit means they were never evicted), // so the store-event cursor skips them. @@ -368,6 +415,29 @@ impl RequestKv { Ok(blocks * self.seq.block_size()) } + /// Returns the lineage hash identifying the registered prefix ending at + /// `boundary_tokens`. + /// + /// `boundary_tokens` must be a non-zero multiple of the KV `block_size`, + /// and be no more than the number of blocks already registered by this request; + /// otherwise this method returns `None`. + pub fn registered_boundary_hash(&self, boundary_tokens: usize) -> Option<[u8; 16]> { + let block_size = self.seq.block_size(); + if boundary_tokens == 0 || !boundary_tokens.is_multiple_of(block_size) { + return None; + } + let block_count = boundary_tokens / block_size; + if block_count > self.seq.assigned_blocks() { + return None; + } + self.seq + .inner() + .sequence() + .all_sequence_hashes() + .get(block_count - 1) + .map(sequence_hash_bytes) + } + // ── Scheduling (allocates blocks) ────────────────────────────────── pub fn schedule_prefill( @@ -760,6 +830,57 @@ mod tests { ); } + #[test] + fn probe_and_attach_can_select_a_shorter_exact_boundary() { + let pool = BlockPool::new(16, 32).unwrap(); + let prompt = (0..80u32).collect::>(); + + let mut seed = pool.new_request(prompt[..64].to_vec(), 4, None); + seed.schedule_prefill(64, &pool).expect("seed schedule"); + seed.apply_prefill(9000, &pool).expect("seed apply"); + assert_eq!( + seed.registered_boundary_hash(32), + Some(seed.prompt_block_hashes()[1]) + ); + seed.release().expect("seed release"); + + let probe = pool.probe_prefix(prompt.clone(), None); + assert_eq!(probe.gpu_hit_blocks(), 4); + assert_eq!(probe.reusable_blocks(), 4); + let boundary_hash = probe.boundary_hash(32).expect("32-token boundary"); + + let mut warm = pool.new_request(prompt, 4, None); + let matched = warm + .match_and_add_prefix_up_to(&pool, 2) + .expect("exact attach"); + assert_eq!(matched, 32); + assert_eq!(warm.kv_position(), 32); + assert_eq!(warm.prefix_matched_blocks(), 2); + assert_eq!(warm.registered_boundary_hash(32), Some(boundary_hash)); + } + + #[test] + fn probe_boundary_hash_obeys_reusable_cap_and_final_token_rule() { + let pool = BlockPool::new(16, 32).unwrap(); + let prompt = (0..64u32).collect::>(); + let mut seed = pool.new_request(prompt.clone(), 4, None); + seed.schedule_prefill(64, &pool).expect("seed schedule"); + seed.apply_prefill(9000, &pool).expect("seed apply"); + seed.release().expect("seed release"); + + let probe = pool.probe_prefix(prompt, None); + assert_eq!(probe.gpu_hit_blocks(), 4); + assert_eq!( + probe.reusable_blocks(), + 3, + "one prompt token must remain uncached" + ); + assert!(probe.boundary_hash(0).is_none()); + assert!(probe.boundary_hash(47).is_none()); + assert!(probe.boundary_hash(48).is_some()); + assert!(probe.boundary_hash(64).is_none()); + } + #[test] fn request_reports_the_lifetime_capacity_it_was_created_with() { let pool = BlockPool::new(16, 8).unwrap(); diff --git a/pegainfer-qwen35/Cargo.toml b/pegainfer-qwen35/Cargo.toml index 0a7da20b6..869763646 100644 --- a/pegainfer-qwen35/Cargo.toml +++ b/pegainfer-qwen35/Cargo.toml @@ -11,6 +11,7 @@ cudarc = { workspace = true } half = { workspace = true } log = { workspace = true } pegainfer-core = { workspace = true } +pegainfer-kv-cache = { workspace = true } pegainfer-kernels = { workspace = true } pegainfer-sample = { workspace = true } rand = { workspace = true } @@ -52,6 +53,10 @@ required-features = ["qwen35"] name = "chunked_prefill" required-features = ["qwen35"] +[[test]] +name = "prefix_cache" +required-features = ["qwen35"] + [[test]] name = "serving_tp2" required-features = ["qwen35"] diff --git a/pegainfer-qwen35/src/batch_decode.rs b/pegainfer-qwen35/src/batch_decode.rs index 99e65a0b2..a2f6acb69 100644 --- a/pegainfer-qwen35/src/batch_decode.rs +++ b/pegainfer-qwen35/src/batch_decode.rs @@ -7,9 +7,10 @@ use cudarc::driver::CudaSlice; use cudarc::driver::DevicePtr; use cudarc::driver::DevicePtrMut; use pegainfer_core::kv_pool::KvLayout; -use pegainfer_core::kv_pool::KvState; use pegainfer_core::sampler::SamplingParams; use pegainfer_core::tensor::HiddenStates; +use pegainfer_kv_cache::KvBuffer; +use pegainfer_kv_cache::KvView; use super::batch_decode_graph::BATCH_BUCKETS; use super::batch_decode_graph::BatchDecodeGraphState; @@ -245,7 +246,8 @@ impl Qwen35Model { pub(crate) fn batch_decode_eager_logits( &self, token_ids: &[u32], - kv_states: &mut [&mut KvState], + views: &[KvView], + kv_buffer: &KvBuffer, recurrent_states: &mut [&mut RecurrentState], linear_pointer_tables: &LinearStatePointerTables, bufs: &mut BatchDecodeBuffers35, @@ -255,7 +257,7 @@ impl Qwen35Model { bs > 0, "batch_decode_eager_logits requires at least one request" ); - anyhow::ensure!(bs == kv_states.len(), "token_ids / kv_states len mismatch"); + anyhow::ensure!(bs == views.len(), "token_ids / KV views len mismatch"); anyhow::ensure!( bs == recurrent_states.len(), "token_ids / recurrent_states len mismatch" @@ -267,12 +269,17 @@ impl Qwen35Model { ); linear_pointer_tables.validate_for(&self.config, bs, "Qwen3.5 eager decode")?; + // KvView describes the post-step KV extent, so this decode token is + // written at seq_len - 1. Recurrent state must start at that position. let mut positions = Vec::with_capacity(bs); - for (i, kv) in kv_states.iter_mut().enumerate() { - let pos = kv.seq_len(); + for (i, view) in views.iter().enumerate() { + let pos = view.seq_len().saturating_sub(1); + anyhow::ensure!( + recurrent_states[i].seq_len == pos, + "Qwen3.5 eager decode position mismatch at row {i}: recurrent={}, view_pos={pos}", + recurrent_states[i].seq_len + ); self.ensure_rope_cache_covers(pos + 1)?; - kv.ensure_capacity(pos + 1)?; - kv.advance(1); recurrent_states[i].seq_len += 1; positions.push(pos as i32); } @@ -285,13 +292,17 @@ impl Qwen35Model { .stream .memcpy_htod(&positions, &mut bufs.positions_d)?; - let kv_refs: Vec<&KvState> = kv_states.iter().map(|s| &**s).collect(); - bufs.sync_paged_meta(&self.ctx, &kv_refs, bs)?; + bufs.sync_paged_views(&self.ctx, views, bs)?; - let kv_buffer = kv_states[0].buffer(); - let layout = *kv_states[0].layout(); + let cache_layout = kv_buffer.layout(); + let layout = KvLayout::new( + cache_layout.num_layers, + cache_layout.num_kv_heads, + cache_layout.head_dim, + cache_layout.page_size, + ); self.batch_decode_kernels_graph( - kv_buffer, + kv_buffer.buffer(), &layout, bs, &linear_pointer_tables.state_ptrs, @@ -316,12 +327,13 @@ impl Qwen35Model { pub(crate) fn batch_decode_graph( &self, token_ids: &[u32], - kv_states: &mut [&mut KvState], + views: &[KvView], + kv_buffer: &KvBuffer, graph_state: &mut BatchDecodeGraphState, ) -> Result<()> { let bs = token_ids.len(); - anyhow::ensure!(bs > 0, "batch_decode_graph requires at least one request"); - anyhow::ensure!(bs == kv_states.len(), "token_ids / kv_states len mismatch"); + anyhow::ensure!(bs > 0, "batch_decode_graph requires requests"); + anyhow::ensure!(bs == views.len(), "token_ids / KV views len mismatch"); anyhow::ensure!( bs <= graph_state.slot_states.len(), "batch size {bs} exceeds decode capacity {}", @@ -332,15 +344,12 @@ impl Qwen35Model { LOG_UNCOMPILED_DECODE_ROUTE.call_once(|| { let group = self.config.num_attention_heads / self.config.num_key_value_heads; log::info!( - "Qwen3.5 decode GQA group {group} ({} q heads / {} kv heads) has no compiled BatchDecode kernel; batched hybrid eager fallback active, bs_capacity={}", - self.config.num_attention_heads, - self.config.num_key_value_heads, - graph_state.buffers.max_batch_size, + "Qwen3.5 decode GQA group {group} has no compiled BatchDecode kernel; batched hybrid eager fallback active" ); }); // Paged-prefill attention stays eager; verify_graph records that // captured prefill attention under-reads growing decode KV. - return self.batch_decode_batched_hybrid(token_ids, kv_states, graph_state); + return self.batch_decode_batched_hybrid(token_ids, views, kv_buffer, graph_state); } let padded_bs = bucket_for(bs); @@ -350,14 +359,17 @@ impl Qwen35Model { "Qwen3.5 graph decode", )?; - // Advance KV states and collect positions. Slot seq_len is incremented - // on the CPU outside the graph so it never appears inside the capture. + // KvView already includes the page reserved by schedule_decode; model + // execution advances recurrent state but never logical RequestKv state. let mut positions = Vec::with_capacity(bs); - for (i, kv) in kv_states.iter_mut().enumerate() { - let pos = kv.seq_len(); + for (i, view) in views.iter().enumerate() { + let pos = view.seq_len().saturating_sub(1); + anyhow::ensure!( + graph_state.slot_states[i].seq_len == pos, + "Qwen3.5 decode position mismatch at slot {i}: recurrent={}, view_pos={pos}", + graph_state.slot_states[i].seq_len + ); self.ensure_rope_cache_covers(pos + 1)?; - kv.ensure_capacity(pos + 1)?; - kv.advance(1); graph_state.slot_states[i].seq_len += 1; positions.push(pos as i32); } @@ -376,13 +388,17 @@ impl Qwen35Model { .memcpy_htod(&positions, &mut graph_state.buffers.positions_d)?; // H2D: paged KV metadata with padding slots pointing to padding_page_id. - let kv_refs: Vec<&KvState> = kv_states.iter().map(|s| &**s).collect(); graph_state .buffers - .sync_paged_meta(&self.ctx, &kv_refs, padded_bs)?; - - let kv_buffer = kv_states[0].buffer(); - let layout = *kv_states[0].layout(); + .sync_paged_views(&self.ctx, views, padded_bs)?; + + let cache_layout = kv_buffer.layout(); + let layout = KvLayout::new( + cache_layout.num_layers, + cache_layout.num_kv_heads, + cache_layout.head_dim, + cache_layout.page_size, + ); let bucket_idx = BATCH_BUCKETS.iter().position(|&b| b == padded_bs).unwrap(); // Take graphs out of graph_state to avoid split-borrow in the closure. @@ -391,7 +407,7 @@ impl Qwen35Model { let linear_conv_state_ptrs = &graph_state.linear_pointer_tables.conv_state_ptrs; let result = graphs[bucket_idx].run_or_capture(&self.ctx, || { self.batch_decode_kernels_graph( - kv_buffer, + kv_buffer.buffer(), &layout, padded_bs, linear_state_ptrs, @@ -406,7 +422,8 @@ impl Qwen35Model { fn batch_decode_batched_hybrid( &self, token_ids: &[u32], - kv_states: &mut [&mut KvState], + views: &[KvView], + kv_buffer: &KvBuffer, graph_state: &mut BatchDecodeGraphState, ) -> Result<()> { let bs = token_ids.len(); @@ -417,13 +434,15 @@ impl Qwen35Model { )?; let mut positions_i32 = Vec::with_capacity(bs); let mut start_positions = Vec::with_capacity(bs); - for (i, kv) in kv_states.iter_mut().enumerate() { - let pos = kv.seq_len(); + for (i, view) in views.iter().enumerate() { + let pos = view.seq_len().saturating_sub(1); + anyhow::ensure!( + graph_state.slot_states[i].seq_len == pos, + "Qwen3.5 hybrid position mismatch at slot {i}: recurrent={}, view_pos={pos}", + graph_state.slot_states[i].seq_len + ); self.ensure_rope_cache_covers(pos + 1) - .with_context(|| format!("hybrid decode rope cache pos={} slot={i}", pos + 1))?; - kv.ensure_capacity(pos + 1) - .with_context(|| format!("hybrid decode KV capacity pos={} slot={i}", pos + 1))?; - kv.advance(1); + .with_context(|| format!("hybrid decode rope pos={} slot={i}", pos + 1))?; graph_state.slot_states[i].seq_len += 1; positions_i32.push(pos as i32); start_positions.push(pos); @@ -433,26 +452,16 @@ impl Qwen35Model { bufs.set_batch_size(bs); self.ctx .stream - .memcpy_htod(token_ids, &mut bufs.token_ids_d) - .map_err(|e| { - anyhow::anyhow!( - "hybrid decode H2D token_ids bs={bs}, cap={}: {e}", - bufs.max_batch_size - ) - })?; + .memcpy_htod(token_ids, &mut bufs.token_ids_d)?; self.ctx .stream - .memcpy_htod(&positions_i32, &mut bufs.positions_d) - .map_err(|e| { - anyhow::anyhow!( - "hybrid decode H2D positions bs={bs}, cap={}: {e}", - bufs.max_batch_size - ) - })?; - - let page_indices: Vec> = - kv_states.iter().map(|kv| kv.page_indices_i32()).collect(); - let last_page_lens: Vec = kv_states.iter().map(|kv| kv.last_page_len()).collect(); + .memcpy_htod(&positions_i32, &mut bufs.positions_d)?; + + let page_indices = views + .iter() + .map(|view| view.page_indices().to_vec()) + .collect::>(); + let last_page_lens = views.iter().map(KvView::last_page_len).collect::>(); let seq_lens = vec![1usize; bs]; // cta_tile_q 0 = the kernel's own FA2 derivation; the hd256 FFI takes no override. let plan = ops::PrefillPagedPlan::from_raw_batch_with_cta_tile_q( @@ -465,30 +474,16 @@ impl Qwen35Model { self.config.num_key_value_heads, self.config.head_dim, 0, - ) - .with_context(|| { - format!( - "hybrid decode build PrefillPagedPlan bs={bs}, pages={}, heads={}/{}, head_dim={}", - page_indices.iter().map(Vec::len).sum::(), - self.config.num_attention_heads, - self.config.num_key_value_heads, - self.config.head_dim - ) - })?; - - let kv_buffer = kv_states[0].buffer(); - let layout = *kv_states[0].layout(); - anyhow::ensure!( - layout.num_kv_heads == self.config.num_key_value_heads - && layout.head_dim == self.config.head_dim, - "hybrid decode KV layout mismatch bs={bs}: layout kv_heads={}, head_dim={}; config kv_heads={}, head_dim={}", - layout.num_kv_heads, - layout.head_dim, - self.config.num_key_value_heads, - self.config.head_dim + )?; + let cache_layout = kv_buffer.layout(); + let layout = KvLayout::new( + cache_layout.num_layers, + cache_layout.num_kv_heads, + cache_layout.head_dim, + cache_layout.page_size, ); self.batch_decode_batched_hybrid_kernels( - kv_buffer, + kv_buffer.buffer(), &layout, &plan, bs, diff --git a/pegainfer-qwen35/src/batch_decode_graph.rs b/pegainfer-qwen35/src/batch_decode_graph.rs index 27bcbdcab..6af42bc36 100644 --- a/pegainfer-qwen35/src/batch_decode_graph.rs +++ b/pegainfer-qwen35/src/batch_decode_graph.rs @@ -2,7 +2,6 @@ use anyhow::Result; use pegainfer_core::cuda_graph::CudaGraphState; -use pegainfer_core::kv_pool::KvPool; use pegainfer_core::tensor::DeviceContext; use super::config::Config35; @@ -62,12 +61,10 @@ impl BatchDecodeGraphState { ctx: &DeviceContext, config: &Config35, tensor_parallel: TensorParallelConfig, - kv_pool: &KvPool, + max_total_pages: usize, + padding_page_id: i32, max_batch: usize, ) -> Result { - let padding_page_id = kv_pool.padding_page_id(); - let max_total_pages = kv_pool.capacity_pages(); - let buffers = BatchDecodeBuffers35::new( ctx, config, @@ -116,16 +113,8 @@ impl BatchDecodeGraphState { src: &RecurrentState, slot_idx: usize, ) -> Result<()> { - let dst = &mut self.slot_states[slot_idx]; - for (dst_layer, src_layer) in dst.layers.iter_mut().zip(src.layers.iter()) { - ctx.stream - .memcpy_dtod(&src_layer.state, &mut dst_layer.state) - .map_err(|e| anyhow::anyhow!("copy recurrent state to slot {slot_idx}: {e}"))?; - ctx.stream - .memcpy_dtod(&src_layer.conv_state.data, &mut dst_layer.conv_state.data) - .map_err(|e| anyhow::anyhow!("copy conv state to slot {slot_idx}: {e}"))?; - } - dst.seq_len = src.seq_len; - Ok(()) + self.slot_states[slot_idx] + .copy_from(ctx, src) + .map_err(|e| anyhow::anyhow!("copy recurrent state to slot {slot_idx}: {e}")) } } diff --git a/pegainfer-qwen35/src/decode_buffers.rs b/pegainfer-qwen35/src/decode_buffers.rs index 3e5bbfad1..d4d32b477 100644 --- a/pegainfer-qwen35/src/decode_buffers.rs +++ b/pegainfer-qwen35/src/decode_buffers.rs @@ -2,9 +2,9 @@ use anyhow::Result; use cudarc::driver::CudaSlice; -use pegainfer_core::kv_pool::KvState; use pegainfer_core::tensor::DeviceContext; use pegainfer_core::tensor::HiddenStates; +use pegainfer_kv_cache::KvView; use super::config::Config35; use super::config::TensorParallelConfig; @@ -151,15 +151,15 @@ impl BatchDecodeBuffers35 { /// Sync paged attention metadata to GPU. /// - /// `padded_bs` >= `kv_states.len()`: padding slots (if any) point to the + /// `padded_bs` >= `views.len()`: padding slots (if any) point to the /// reserved padding page with seq_len=1 so FlashInfer accesses valid memory. - pub(crate) fn sync_paged_meta( + pub(crate) fn sync_paged_views( &mut self, ctx: &DeviceContext, - kv_states: &[&KvState], + views: &[KvView], padded_bs: usize, ) -> Result<()> { - let real_bs = kv_states.len(); + let real_bs = views.len(); debug_assert!(padded_bs >= real_bs); let mut all_page_indices = Vec::new(); @@ -167,12 +167,11 @@ impl BatchDecodeBuffers35 { let mut last_page_lens = Vec::with_capacity(padded_bs); let mut chunk_sizes = Vec::with_capacity(padded_bs); - for kv in kv_states { - let pages = kv.page_indices_i32(); - all_page_indices.extend_from_slice(&pages); + for view in views { + all_page_indices.extend_from_slice(view.page_indices()); indptr.push(all_page_indices.len() as i32); - last_page_lens.push(kv.last_page_len() as i32); - chunk_sizes.push(kv.seq_len() as i32); + last_page_lens.push(view.last_page_len() as i32); + chunk_sizes.push(view.seq_len() as i32); } // Padding slots: 1 page (the padding page), seq_len=1, last_page_len=1. diff --git a/pegainfer-qwen35/src/executor.rs b/pegainfer-qwen35/src/executor.rs index e7eebb2d4..e878e6bec 100644 --- a/pegainfer-qwen35/src/executor.rs +++ b/pegainfer-qwen35/src/executor.rs @@ -8,13 +8,15 @@ use std::collections::HashSet; use anyhow::Result; use pegainfer_core::engine::TokenLogprob; -use pegainfer_core::kv_pool::KvState; use pegainfer_core::sampler::SamplingParams; use pegainfer_core::tensor::HiddenStates; +use pegainfer_kv_cache::KvCacheManager; +use pegainfer_kv_cache::RequestKv; use crate::batch_decode_graph::BatchDecodeGraphState; use crate::decode_buffers::BatchDecodeBuffers35; use crate::logprobs::snapshot_requested_logprobs; +use crate::prefix_cache::Qwen35PrefixCache; use crate::recurrent_state::RecurrentState; use crate::weights::Qwen35Model; @@ -101,23 +103,31 @@ pub struct DecodeResult { struct ActiveRequest { request_id: RequestId, - kv: KvState, + kv: RequestKv, graph_slot_idx: usize, } pub struct Qwen35Executor { model: Qwen35Model, + kv_cache: Qwen35PrefixCache, graph_state: BatchDecodeGraphState, active: Vec, } impl Qwen35Executor { pub fn from_runtime(model_path: &str, device_ordinal: usize, max_batch: usize) -> Result { - let model = Qwen35Model::from_safetensors(model_path, device_ordinal, max_batch)?; + let model = Qwen35Model::from_safetensors(model_path, device_ordinal, max_batch, 0)?; model.tune_decode_gemm_algos()?; - let graph_state = model.create_batch_decode_graph_state()?; + let manager = + KvCacheManager::from_buffer(model.kv_buffer().clone(), model.kv_buffer().num_blocks())?; + let kv_cache = Qwen35PrefixCache::new(manager, 0)?; + let graph_state = model.create_batch_decode_graph_state( + kv_cache.pool().total_blocks(), + kv_cache.pool().padding_block_id(), + )?; Ok(Self { model, + kv_cache, graph_state, active: Vec::new(), }) @@ -159,20 +169,52 @@ impl Qwen35Executor { .iter() .map(|req| req.prompt_tokens.as_slice()) .collect(); - let mut kv_states: Vec = plan + let mut kv_states: Vec = plan .requests .iter() - .map(|_| self.model.alloc_kv()) + .map(|req| { + self.kv_cache.pool().new_request( + req.prompt_tokens.clone(), + self.model + .config() + .max_position_embeddings + .saturating_sub(req.prompt_tokens.len()), + None, + ) + }) .collect(); + for scheduled in 0..kv_states.len() { + if let Err(error) = self.kv_cache.schedule_prefill( + &mut kv_states[scheduled], + plan.requests[scheduled].prompt_tokens.len(), + ) { + revert_scheduled_requests(&self.kv_cache, kv_states.iter_mut().take(scheduled)); + return Err(error); + } + } + let views = kv_states + .iter() + .zip(plan.requests) + .map(|(kv, req)| self.kv_cache.prefill_view(kv, req.prompt_tokens.len())) + .collect::>(); let mut recurrent_states: Vec = plan .requests .iter() .map(|_| RecurrentState::new(self.model.device_ctx(), self.model.config())) .collect::>()?; let mut recurrent_refs: Vec<&mut RecurrentState> = recurrent_states.iter_mut().collect(); - let logits = - self.model - .batch_prefill_logits(&prompts, &mut kv_states, &mut recurrent_refs)?; + let logits = match self.model.batch_prefill_logits( + &prompts, + &views, + &mut recurrent_refs, + self.kv_cache.buffer(), + ) { + Ok(logits) => logits, + Err(error) => { + revert_scheduled_requests(&self.kv_cache, &mut kv_states); + return Err(error); + } + }; let requested_logprobs: Vec = plan.requests.iter().map(|req| req.logprobs).collect(); let cpu_logits = @@ -181,8 +223,9 @@ impl Qwen35Executor { select_default_tokens_from_logits(&self.model, &logits, &mut self.graph_state.buffers)?; let mut results = Vec::with_capacity(plan.requests.len()); - for (i, (req, kv)) in plan.requests.iter().zip(kv_states).enumerate() { + for (i, (req, mut kv)) in plan.requests.iter().zip(kv_states).enumerate() { let first_token = tokens[i]; + self.kv_cache.apply_prefill(&mut kv, Some(first_token))?; let first_token_logprob = cpu_logits[i].as_ref().and_then(|row| { pegainfer_sample::token_logprob_from_row(row, first_token, req.logprobs) }); @@ -224,10 +267,38 @@ impl Qwen35Executor { } let token_ids: Vec = plan.requests.iter().map(|req| req.token_id).collect(); - let mut kv_refs: Vec<&mut KvState> = - self.active.iter_mut().map(|req| &mut req.kv).collect(); - self.model - .batch_decode_graph(&token_ids, &mut kv_refs, &mut self.graph_state)?; + for scheduled in 0..self.active.len() { + if let Err(error) = self + .kv_cache + .schedule_decode(&mut self.active[scheduled].kv) + { + revert_scheduled_requests( + &self.kv_cache, + self.active + .iter_mut() + .take(scheduled) + .map(|active| &mut active.kv), + ); + return Err(error); + } + } + let views = self + .active + .iter() + .map(|req| self.kv_cache.decode_view(&req.kv)) + .collect::>(); + if let Err(error) = self.model.batch_decode_graph( + &token_ids, + &views, + self.kv_cache.buffer(), + &mut self.graph_state, + ) { + revert_scheduled_requests( + &self.kv_cache, + self.active.iter_mut().map(|active| &mut active.kv), + ); + return Err(error); + } let requested_logprobs: Vec = plan.requests.iter().map(|req| req.logprobs).collect(); let cpu_logits = snapshot_requested_logprobs( @@ -246,6 +317,7 @@ impl Qwen35Executor { let mut results = Vec::with_capacity(plan.requests.len()); for (i, req) in plan.requests.iter().enumerate() { let token = tokens[i]; + self.kv_cache.apply_decode(&mut self.active[i].kv, token)?; let logprob = cpu_logits[i] .as_ref() .and_then(|row| pegainfer_sample::token_logprob_from_row(row, token, req.logprobs)); @@ -269,9 +341,11 @@ impl Qwen35Executor { self.compact_slot(idx) } + /// Remove one active request and keep the dense CUDA Graph slot layout. fn compact_slot(&mut self, idx: usize) -> Result<()> { let last = self.active.len() - 1; - self.active.swap_remove(idx); + let mut removed = self.active.swap_remove(idx); + let release_result = self.kv_cache.release_request(&mut removed.kv); if idx < self.active.len() { anyhow::ensure!( @@ -311,7 +385,19 @@ impl Qwen35Executor { self.graph_state.slot_states[idx].seq_len = self.graph_state.slot_states[last].seq_len; self.active[idx].graph_slot_idx = idx; } - Ok(()) + release_result + } +} + +/// Roll back a set of requests scheduled by the current executor step. +fn revert_scheduled_requests<'a>( + kv_cache: &Qwen35PrefixCache, + requests: impl IntoIterator, +) { + for request in requests { + if let Err(error) = kv_cache.revert_schedule(request) { + log::warn!("failed to revert Qwen3.5 executor KV schedule: {error}"); + } } } diff --git a/pegainfer-qwen35/src/lib.rs b/pegainfer-qwen35/src/lib.rs index 642f8b4b9..8c1b01d49 100644 --- a/pegainfer-qwen35/src/lib.rs +++ b/pegainfer-qwen35/src/lib.rs @@ -14,6 +14,7 @@ mod logprobs; mod ops; mod prefill; pub mod prefill_buffers; +mod prefix_cache; pub(crate) mod recurrent; pub(crate) mod recurrent_state; mod scheduler; @@ -83,6 +84,7 @@ pub fn start_engine( max_batch, max_prefill_tokens, Qwen35SchedulerPolicy::Off, + 0, ) } @@ -96,6 +98,9 @@ pub struct Qwen35LaunchOptions { pub cuda_graph: bool, pub max_batch: usize, pub max_prefill_tokens: usize, + /// Fixed GPU budget for joint recurrent/conv prefix snapshots. + /// Zero keeps Qwen3.5 prefix matching disabled. + pub prefix_cache_mib: usize, } impl Qwen35LaunchOptions { @@ -128,6 +133,7 @@ pub fn launch_with_options_and_policy( options.max_batch, options.max_prefill_tokens, scheduler_policy, + options.prefix_cache_mib, ) } @@ -143,6 +149,7 @@ pub fn start_engine_with_capacity( max_batch, max_prefill_tokens, Qwen35SchedulerPolicy::Off, + 0, ) } @@ -152,6 +159,7 @@ pub fn start_engine_with_capacity_and_policy( max_batch: usize, max_prefill_tokens: usize, scheduler_policy: Qwen35SchedulerPolicy, + prefix_cache_mib: usize, ) -> Result { anyhow::ensure!( (1..=MAX_DECODE_BATCH).contains(&max_batch), @@ -163,6 +171,9 @@ pub fn start_engine_with_capacity_and_policy( seed, .. } = options; + let prefix_snapshot_bytes = prefix_cache_mib + .checked_mul(1024 * 1024) + .ok_or_else(|| anyhow!("Qwen3.5 prefix-cache MiB budget overflows usize"))?; if device_ordinals.len() > 1 { if scheduler_policy == Qwen35SchedulerPolicy::Auto { return Err(anyhow!( @@ -183,6 +194,7 @@ pub fn start_engine_with_capacity_and_policy( &device_ordinals, max_batch, max_prefill_tokens, + prefix_snapshot_bytes, ); } @@ -203,7 +215,12 @@ pub fn start_engine_with_capacity_and_policy( let model_path = model_path .to_str() .ok_or_else(|| anyhow!("model path must be valid UTF-8"))?; - let model = weights::Qwen35Model::from_safetensors(model_path, device_ordinal, max_batch)?; + let model = weights::Qwen35Model::from_safetensors( + model_path, + device_ordinal, + max_batch, + prefix_snapshot_bytes, + )?; scheduler::start_with_capacity_and_policy( model, seed, @@ -232,6 +249,7 @@ mod tests { cuda_graph: false, max_batch: 1, max_prefill_tokens: 1, + prefix_cache_mib: 0, }; let err = options.device_ordinals().unwrap_err().to_string(); @@ -257,6 +275,7 @@ mod tests { 1, 1, Qwen35SchedulerPolicy::Auto, + 0, ) .err() .expect("scheduler policy validation should reject TP launch") diff --git a/pegainfer-qwen35/src/prefill.rs b/pegainfer-qwen35/src/prefill.rs index f641983a9..a1b5aedbb 100644 --- a/pegainfer-qwen35/src/prefill.rs +++ b/pegainfer-qwen35/src/prefill.rs @@ -20,9 +20,10 @@ pub(crate) const SCRATCH_ESTIMATE_SEQ: usize = 20_000; pub(crate) const PREFILL_CHUNK_LEN: usize = SCRATCH_ESTIMATE_SEQ; const HEAD_DIM: usize = 256; -use pegainfer_core::kv_pool::KvState; use pegainfer_core::tensor::DeviceVec; use pegainfer_core::tensor::HiddenStates; +use pegainfer_kv_cache::KvBuffer; +use pegainfer_kv_cache::KvView; use super::prefill_buffers::GdrChunkwiseScratch35; use super::recurrent_state::RecurrentState; @@ -35,6 +36,13 @@ use crate::ffi; use crate::ops; use crate::ops::PrefillPagedPlan; +struct PrefillKvAccess<'a> { + buffer: &'a cudarc::driver::CudaSlice, + layout: pegainfer_core::kv_pool::KvLayout, + plan: &'a PrefillPagedPlan, + base_pos: usize, +} + fn checked_prefill_end_pos( base_pos: usize, seq_len: usize, @@ -54,21 +62,28 @@ impl Qwen35Model { pub(super) fn prefill_last_hidden( &self, token_ids: &[u32], - kv_state: &mut KvState, + full_view: &KvView, + kv_buffer: &KvBuffer, recurrent: &mut RecurrentState, ) -> Result { - let seq_len = token_ids.len(); anyhow::ensure!( - seq_len > 0, - "Qwen3.5 prefill_last_hidden requires at least one token" + !token_ids.is_empty(), + "Qwen3.5 prefill requires at least one token" ); - let c = &self.config; - // Validate the full target range up front (position overflow + RoPE cache // coverage) so an out-of-range prompt is rejected before any chunk mutates // the KV / recurrent state, rather than failing partway through. - let base_pos = kv_state.seq_len(); - let end_pos = checked_prefill_end_pos(base_pos, seq_len, c.max_position_embeddings)?; + let base_pos = recurrent.seq_len; + let end_pos = checked_prefill_end_pos( + base_pos, + token_ids.len(), + self.config.max_position_embeddings, + )?; + anyhow::ensure!( + full_view.seq_len() == end_pos, + "Qwen3.5 prefill view ends at {}, expected {end_pos}", + full_view.seq_len() + ); self.ensure_rope_cache_covers(end_pos)?; // Run prefill in serial chunks of at most `PREFILL_CHUNK_LEN` tokens. Each @@ -76,15 +91,54 @@ impl Qwen35Model { // place, so the next chunk continues from the previous one. This caps the // per-pass GDR scratch (which grows with the pass length) at the budget // reserved at startup, so prompts longer than one chunk prefill without OOM. - let mut hidden_batch: Option = None; + let page_size = kv_buffer.layout().page_size; + let mut hidden_batch = None; + let mut offset = 0usize; for chunk in token_ids.chunks(PREFILL_CHUNK_LEN) { // Free the previous chunk's hidden states before allocating the next // chunk's scratch so peak memory stays within one chunk's reservation. drop(hidden_batch.take()); - hidden_batch = Some(self.prefill_chunk_forward(chunk, kv_state, recurrent)?); + let (chunk_hidden, chunk_scratch) = self.prepare_prefill_chunk(chunk)?; + let chunk_end = base_pos + offset + chunk.len(); + let page_count = chunk_end.div_ceil(page_size); + let chunk_view = KvView::new( + full_view.page_indices()[..page_count].to_vec(), + chunk_end, + page_size, + ); + let page_indices = vec![chunk_view.page_indices().to_vec()]; + let plan = PrefillPagedPlan::from_raw_batch_with_cta_tile_q( + &self.ctx, + &page_indices, + &[chunk_view.last_page_len()], + &[base_pos + offset], + &[chunk.len()], + self.config.local_num_attention_heads(self.tensor_parallel), + self.config.local_num_key_value_heads(self.tensor_parallel), + self.config.head_dim, + 0, + )?; + let access = PrefillKvAccess { + buffer: kv_buffer.buffer(), + layout: pegainfer_core::kv_pool::KvLayout::new( + kv_buffer.layout().num_layers, + kv_buffer.layout().num_kv_heads, + kv_buffer.layout().head_dim, + page_size, + ), + plan: &plan, + base_pos: base_pos + offset, + }; + hidden_batch = Some(self.prefill_chunk_forward( + chunk_hidden, + chunk_scratch, + &access, + recurrent, + )?); + offset += chunk.len(); } // `seq_len > 0` guarantees at least one chunk produced hidden states. - let hidden_batch = hidden_batch.expect("prefill produced no chunk despite seq_len > 0"); + let hidden_batch = hidden_batch.expect("non-empty prefill produced no chunk"); // Last-token logic runs once, on the final chunk's output. ops::extract_vec(&self.ctx, &hidden_batch, hidden_batch.seq_len - 1) @@ -129,64 +183,38 @@ impl Qwen35Model { Ok(logits) } - /// Forward one prefill chunk through all layers, advancing the paged KV state - /// and the linear-attention recurrent/conv state in place. - /// - /// `token_ids.len()` must be in `1..=PREFILL_CHUNK_LEN` so the per-chunk GDR - /// scratch stays within the startup reservation. Returns the chunk's hidden - /// states for every token; only the final chunk's last token feeds the LM head. - fn prefill_chunk_forward( + fn prepare_prefill_chunk( &self, token_ids: &[u32], - kv_state: &mut KvState, - recurrent: &mut RecurrentState, - ) -> Result { + ) -> Result<(HiddenStates, GdrChunkwiseScratch35)> { let seq_len = token_ids.len(); - debug_assert!( - seq_len > 0 && seq_len <= PREFILL_CHUNK_LEN, - "prefill chunk length {seq_len} out of range 1..={PREFILL_CHUNK_LEN}" - ); let c = &self.config; - let base_pos = kv_state.seq_len(); - let end_pos = checked_prefill_end_pos(base_pos, seq_len, c.max_position_embeddings)?; - self.ensure_rope_cache_covers(end_pos)?; // Embeddings for this chunk. let token_ids_gpu = self .ctx .stream .clone_htod(token_ids) - .map_err(|e| anyhow::anyhow!("H2D copy failed: {}", e))?; - - let hidden_dim = c.hidden_size; - let mut hidden_batch = HiddenStates::zeros(&self.ctx, hidden_dim, seq_len)?; + .map_err(|e| anyhow::anyhow!("H2D copy failed: {e}"))?; + let mut hidden_batch = HiddenStates::zeros(&self.ctx, c.hidden_size, seq_len)?; ops::embedding_batch( &self.ctx, &self.embed_tokens, &token_ids_gpu, &mut hidden_batch, )?; + let gdr_chunkwise_scratch = GdrChunkwiseScratch35::new(&self.ctx, c, seq_len)?; + Ok((hidden_batch, gdr_chunkwise_scratch)) + } - // Allocate the chunk scratch before advancing the KV state. It is the - // largest, most allocation-prone buffer here, so failing first leaves - // `kv_state` untouched and the request can be rejected cleanly. - let mut gdr_chunkwise_scratch = GdrChunkwiseScratch35::new(&self.ctx, c, seq_len)?; - - // Advance paged KV state and build this chunk's prefill plan. - kv_state.ensure_capacity(end_pos)?; - kv_state.advance(seq_len); - let kv_desc = kv_state.desc(); - let tp = self.tensor_parallel; - let prefill_plan = PrefillPagedPlan::new( - &self.ctx, - &kv_desc, - base_pos, - seq_len, - c.local_num_attention_heads(tp), - c.local_num_key_value_heads(tp), - c.head_dim, - )?; - + fn prefill_chunk_forward( + &self, + mut hidden_batch: HiddenStates, + mut gdr_chunkwise_scratch: GdrChunkwiseScratch35, + kv: &PrefillKvAccess<'_>, + recurrent: &mut RecurrentState, + ) -> Result { + let seq_len = hidden_batch.seq_len; // Process layers let mut linear_idx = 0usize; let mut full_idx = 0usize; @@ -199,14 +227,13 @@ impl Qwen35Model { &mut gdr_chunkwise_scratch, &mut linear_idx, &mut full_idx, - kv_state, - &prefill_plan, + kv, recurrent, )?; } // Advance recurrent token count for the next chunk / decode step; the - // paged KV position is tracked by `kv_state` (advanced above). + // paged KV position is committed by the caller. recurrent.seq_len += seq_len; Ok(hidden_batch) @@ -222,8 +249,7 @@ impl Qwen35Model { gdr_chunkwise_scratch: &mut GdrChunkwiseScratch35, linear_idx: &mut usize, full_idx: &mut usize, - kv_state: &KvState, - prefill_plan: &PrefillPagedPlan, + kv: &PrefillKvAccess<'_>, recurrent: &mut RecurrentState, ) -> Result { let c = &self.config; @@ -249,8 +275,7 @@ impl Qwen35Model { attn, &normed_batch, full_idx, - kv_state, - prefill_plan, + kv, attn_out_dim, seq_len, )?, @@ -288,8 +313,7 @@ impl Qwen35Model { attn: &FullAttentionLayer, normed_batch: &HiddenStates, full_idx: &mut usize, - kv_state: &KvState, - prefill_plan: &PrefillPagedPlan, + kv: &PrefillKvAccess<'_>, _attn_out_dim: usize, seq_len: usize, ) -> Result { @@ -304,16 +328,14 @@ impl Qwen35Model { let v_batch = ops::gemm(&self.ctx, &attn.v_proj, normed_batch)?; let mut attn_out_batch = HiddenStates::zeros(&self.ctx, attn_out_dim, seq_len)?; - // `kv_state` was advanced by `seq_len` before the layer loop, so the - // base write position for this prefill is `seq_len()` minus this batch. - let base_pos = kv_state.seq_len() - seq_len; + let base_pos = kv.base_pos; let mut q_prepped = HiddenStates::zeros(&self.ctx, attn_out_dim, seq_len)?; let start_pos_cpu: CudaSlice = self .ctx .stream .clone_htod(&[base_pos as i32]) .map_err(|e| anyhow::anyhow!("H2D start_pos failed: {e}"))?; - let layout = kv_state.layout(); + let layout = &kv.layout; let layer_k_off = (*full_idx * layout.layer_stride) as i64; let layer_v_off = layer_k_off + layout.kv_block_len as i64; let stride_page = layout.page_stride as i64; @@ -328,8 +350,8 @@ impl Qwen35Model { let (cos_ptr, _) = self.cos_cache.data.device_ptr(&self.ctx.stream); let (sin_ptr, _) = self.sin_cache.data.device_ptr(&self.ctx.stream); let (qp_ptr, _) = q_prepped.data.device_ptr_mut(&self.ctx.stream); - let (buf_ptr, _) = kv_state.buffer().device_ptr(&self.ctx.stream); - let (pi_ptr, _) = prefill_plan.page_indices_d().device_ptr(&self.ctx.stream); + let (buf_ptr, _) = kv.buffer.device_ptr(&self.ctx.stream); + let (pi_ptr, _) = kv.plan.page_indices_d().device_ptr(&self.ctx.stream); let (sp_ptr, _) = start_pos_cpu.device_ptr(&self.ctx.stream); ffi::prefill_attention_hd256_prep_paged_cuda( qf_ptr as *const ffi::Half, @@ -359,24 +381,18 @@ impl Qwen35Model { // Step 2: Batch prefill paged attention (HD=256). let sm_scale = 1.0f32 / f32::sqrt(HEAD_DIM as f32); { - let (buf_ptr, _gbuf) = kv_state.buffer().device_ptr(&self.ctx.stream); + let (buf_ptr, _gbuf) = kv.buffer.device_ptr(&self.ctx.stream); let (qp_ptr, _gqp) = q_prepped.data.device_ptr(&self.ctx.stream); let (out_ptr, _go) = attn_out_batch.data.device_ptr_mut(&self.ctx.stream); - let (pi_ptr, _gpi) = prefill_plan.page_indices_d().device_ptr(&self.ctx.stream); - let (pip_ptr, _gpip) = prefill_plan.page_indptr_d().device_ptr(&self.ctx.stream); - let (lpl_ptr, _glpl) = prefill_plan.last_page_len_d().device_ptr(&self.ctx.stream); - let (qi_ptr, _gqi) = prefill_plan.q_indptr_d().device_ptr(&self.ctx.stream); - let (ri_ptr, _gri) = prefill_plan - .request_indices_d() - .device_ptr(&self.ctx.stream); - let (qti_ptr, _gqti) = prefill_plan - .qo_tile_indices_d() - .device_ptr(&self.ctx.stream); - let (kti_ptr, _gkti) = prefill_plan - .kv_tile_indices_d() - .device_ptr(&self.ctx.stream); - let (kcs_ptr, _gkcs) = prefill_plan.kv_chunk_size_d().device_ptr(&self.ctx.stream); - let (tnr_ptr, _gtnr) = prefill_plan.total_num_rows_d().device_ptr(&self.ctx.stream); + let (pi_ptr, _gpi) = kv.plan.page_indices_d().device_ptr(&self.ctx.stream); + let (pip_ptr, _gpip) = kv.plan.page_indptr_d().device_ptr(&self.ctx.stream); + let (lpl_ptr, _glpl) = kv.plan.last_page_len_d().device_ptr(&self.ctx.stream); + let (qi_ptr, _gqi) = kv.plan.q_indptr_d().device_ptr(&self.ctx.stream); + let (ri_ptr, _gri) = kv.plan.request_indices_d().device_ptr(&self.ctx.stream); + let (qti_ptr, _gqti) = kv.plan.qo_tile_indices_d().device_ptr(&self.ctx.stream); + let (kti_ptr, _gkti) = kv.plan.kv_tile_indices_d().device_ptr(&self.ctx.stream); + let (kcs_ptr, _gkcs) = kv.plan.kv_chunk_size_d().device_ptr(&self.ctx.stream); + let (tnr_ptr, _gtnr) = kv.plan.total_num_rows_d().device_ptr(&self.ctx.stream); let result = unsafe { ffi::batch_prefill_paged_cuda_hd256( qp_ptr as *const ffi::Half, @@ -398,8 +414,8 @@ impl Qwen35Model { HEAD_DIM as i32, layout.page_size as i32, seq_len as i32, - prefill_plan.batch_size(), - prefill_plan.num_tiles(), + kv.plan.batch_size(), + kv.plan.num_tiles(), stride_page, sm_scale, self.ctx.stream.cu_stream(), diff --git a/pegainfer-qwen35/src/prefix_cache.rs b/pegainfer-qwen35/src/prefix_cache.rs new file mode 100644 index 000000000..70fff76f2 --- /dev/null +++ b/pegainfer-qwen35/src/prefix_cache.rs @@ -0,0 +1,677 @@ +//! Joint full-attention KV and recurrent/conv prefix cache for Qwen3.5. + +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Instant; + +use anyhow::Result; +use pegainfer_core::tensor::DeviceContext; +use pegainfer_kv_cache::KvBlockGuard; +use pegainfer_kv_cache::KvBuffer; +use pegainfer_kv_cache::KvCacheManager; +use pegainfer_kv_cache::KvView; +use pegainfer_kv_cache::RequestKv; + +use crate::config::Config35; +use crate::recurrent_state::RecurrentState; + +/// Token interval at which a complete recurrent/conv snapshot may be cached. +pub(crate) const SNAPSHOT_STRIDE_TOKENS: usize = 256; + +/// Content-addressed identity of one complete hybrid-model prefix boundary. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub(crate) struct PrefixBoundaryKey { + /// Canonical full-attention KV lineage hash for this prefix. + pub(crate) sequence_hash: [u8; 16], + /// Exclusive token position represented by both KV and recurrent state. + pub(crate) boundary_tokens: usize, +} + +/// Cumulative counters for joint KV/recurrent prefix-cache activity. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub(crate) struct PrefixCacheStats { + /// Requests that restored both KV and a recurrent snapshot. + pub(crate) joint_hits: u64, + /// Prompt tokens reused across all joint hits. + pub(crate) joint_hit_tokens: u64, + /// Requests with eligible resident KV but no matching snapshot. + pub(crate) kv_only_fallbacks: u64, + /// Individual eligible boundaries without a matching snapshot. + pub(crate) snapshot_misses: u64, + /// New recurrent snapshots published. + pub(crate) inserts: u64, + /// Published snapshots replaced by LRU insertion. + pub(crate) evictions: u64, + /// Successful joint restore time, including lookup, attach, copy, and checks. + pub(crate) restore_ns: u64, +} + +/// One reusable Qwen3.5 prefix entry. +/// +/// The entry is the ownership boundary for the recurrent snapshot slot +/// and the leading KV blocks. Dropping an entry therefore drops its KV lease as well. +struct PrefixEntry { + /// Identically numbered physical recurrent snapshot on every rank. + recurrent_slot: usize, + /// Strong pins for every KV block through the entry boundary. + #[allow(dead_code)] // ownership is the use: dropping the entry drops the lease + kv_lease: Vec, + /// Active restore guards preventing this entry from being evicted. + pin_count: Arc, + /// Logical timestamp used to select an unpinned LRU victim. + last_used: u64, +} + +/// RAII pin on one prefix cache entry while its recurrent state is being restored. +pub(crate) struct PrefixGuard { + /// Token boundary represented by the pinned entry. + boundary: usize, + /// Rank-local physical snapshot slot selected by the directory. + recurrent_slot: usize, + /// Shared count consulted by insertion before choosing a victim. + pin_count: Arc, + /// Start time used for restore latency accounting. + started: Instant, +} + +impl PrefixGuard { + pub(crate) fn boundary(&self) -> usize { + self.boundary + } + + pub(crate) fn recurrent_slot(&self) -> usize { + self.recurrent_slot + } +} + +impl Drop for PrefixGuard { + fn drop(&mut self) { + let previous = self.pin_count.fetch_sub(1, Ordering::AcqRel); + debug_assert!(previous > 0, "prefix guard pin underflow"); + } +} + +/// Mutable prefix-cache state shared by every execution rank. +struct PrefixCacheState { + /// Published boundary key to its complete joint entry. + entries: HashMap, + /// Unpublished slots available without eviction. + free_slots: Vec, + /// Total number of preallocated physical snapshot slots. + slot_count: usize, + /// Monotonic logical time for LRU ordering. + clock: u64, +} + +impl PrefixCacheState { + fn new(slot_count: usize) -> Self { + Self { + entries: HashMap::with_capacity(slot_count), + free_slots: (0..slot_count).rev().collect(), + slot_count, + clock: 0, + } + } + + /// Number of currently published joint entries. + fn len(&self) -> usize { + self.entries.len() + } + + /// Number of preallocated recurrent-state slots. + fn capacity(&self) -> usize { + self.slot_count + } + + /// Advance the non-zero logical clock used by the LRU policy. + fn tick(&mut self) -> u64 { + self.clock = self.clock.wrapping_add(1).max(1); + self.clock + } + + /// Look up `key`, refresh its LRU timestamp, and pin its entry. + fn lookup(&mut self, key: PrefixBoundaryKey, started: Instant) -> Option { + let last_used = self.tick(); + let entry = self.entries.get_mut(&key)?; + entry.last_used = last_used; + entry.pin_count.fetch_add(1, Ordering::AcqRel); + Some(PrefixGuard { + boundary: key.boundary_tokens, + recurrent_slot: entry.recurrent_slot, + pin_count: Arc::clone(&entry.pin_count), + started, + }) + } + + /// Reserve a free or unpinned LRU slot without publishing the new entry. + fn reserve(&mut self, key: PrefixBoundaryKey) -> Option { + let last_used = self.tick(); + if let Some(entry) = self.entries.get_mut(&key) { + entry.last_used = last_used; + return None; + } + + let (slot, evicted) = if let Some(slot) = self.free_slots.pop() { + (slot, false) + } else { + let (&victim_key, slot) = self + .entries + .iter() + .filter(|(_, entry)| entry.pin_count.load(Ordering::Acquire) == 0) + .min_by_key(|(_, entry)| entry.last_used) + .map(|(key, entry)| (key, entry.recurrent_slot))?; + let evicted = self + .entries + .remove(&victim_key) + .expect("LRU victim must still be present"); + debug_assert_eq!(evicted.recurrent_slot, slot); + (slot, true) + }; + + Some(PrefixReservation { + recurrent_slot: slot, + key, + replaced: evicted, + }) + } + + /// Publish one complete entry only after every rank has enqueued its + /// physical recurrent-state copy. + fn publish(&mut self, reservation: PrefixReservation, kv_lease: Vec) { + let last_used = self.tick(); + let previous = self.entries.insert( + reservation.key, + PrefixEntry { + recurrent_slot: reservation.recurrent_slot, + kv_lease, + pin_count: Arc::new(AtomicUsize::new(0)), + last_used, + }, + ); + debug_assert!(previous.is_none()); + } + + /// Return an unpublished slot after a physical copy failure. + fn abort(&mut self, reservation: PrefixReservation) { + debug_assert!(!self.entries.contains_key(&reservation.key)); + self.free_slots.push(reservation.recurrent_slot); + } +} + +/// One pending central-directory insertion. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct PrefixReservation { + recurrent_slot: usize, + key: PrefixBoundaryKey, + /// Whether this insertion replaced an unpinned joint entry. + replaced: bool, +} + +impl PrefixReservation { + pub(crate) fn recurrent_slot(self) -> usize { + self.recurrent_slot + } + + fn was_replacement(self) -> bool { + self.replaced + } +} + +/// Rank-local physical recurrent/conv snapshot allocations. +pub(crate) struct RecurrentStateStore { + slots: Vec, +} + +impl RecurrentStateStore { + pub(crate) fn new(ctx: &DeviceContext, config: &Config35, slot_count: usize) -> Result { + let mut slots = Vec::with_capacity(slot_count); + for _ in 0..slot_count { + slots.push(RecurrentState::new(ctx, config)?); + } + Ok(Self { slots }) + } + + pub(crate) fn len(&self) -> usize { + self.slots.len() + } + + pub(crate) fn save( + &mut self, + ctx: &DeviceContext, + slot: usize, + src: &RecurrentState, + ) -> Result<()> { + let dst = self + .slots + .get_mut(slot) + .ok_or_else(|| anyhow::anyhow!("snapshot slot {slot} out of range"))?; + dst.copy_from(ctx, src) + } + + pub(crate) fn restore( + &self, + ctx: &DeviceContext, + slot: usize, + dst: &mut RecurrentState, + ) -> Result<()> { + let src = self + .slots + .get(slot) + .ok_or_else(|| anyhow::anyhow!("snapshot slot {slot} out of range"))?; + dst.copy_from(ctx, src) + } +} + +/// The only Qwen3.5 scheduler interface allowed to reconcile paged KV with +/// recurrent/conv state. +pub(crate) struct Qwen35PrefixCache { + /// Logical block pool paired with the full-attention GPU KV buffer. + kv: KvCacheManager, + /// Prefix key, slot, pin, LRU, and KV-lease ownership for reusable entries. + state: PrefixCacheState, + /// Whether at least one joint snapshot can be retained. + enabled: bool, + /// Scheduler-thread-owned cumulative metrics. + stats: PrefixCacheStats, +} + +impl Qwen35PrefixCache { + /// Build the joint coordinator for `snapshot_slots` per-rank allocations. + pub(crate) fn new(kv: KvCacheManager, snapshot_slots: usize) -> Result { + anyhow::ensure!( + SNAPSHOT_STRIDE_TOKENS.is_multiple_of(kv.pool().block_size()), + "Qwen3.5 snapshot stride {SNAPSHOT_STRIDE_TOKENS} must be a multiple of KV block size {}", + kv.pool().block_size() + ); + Ok(Self { + kv, + state: PrefixCacheState::new(snapshot_slots), + enabled: snapshot_slots > 0, + stats: PrefixCacheStats::default(), + }) + } + + /// Logical KV block pool used for request allocation and admission. + pub(crate) fn pool(&self) -> &pegainfer_kv_cache::BlockPool { + self.kv.pool() + } + + /// Rank-0 physical full-attention KV storage indexed by [`Self::pool`]. + pub(crate) fn buffer(&self) -> &KvBuffer { + self.kv.buffer() + } + + /// Whether joint prefix reuse and snapshot publication are enabled. + pub(crate) fn enabled(&self) -> bool { + self.enabled + } + + /// Total number of preallocated recurrent snapshot slots. + pub(crate) fn snapshot_slots(&self) -> usize { + self.state.capacity() + } + + /// Number of snapshot slots that currently have a published key. + pub(crate) fn snapshot_occupancy(&self) -> usize { + self.state.len() + } + + /// Return a point-in-time copy of cumulative cache metrics. + pub(crate) fn stats(&self) -> PrefixCacheStats { + self.stats + } + + /// Create request-local KV state and select the longest joint prefix. + /// + /// If a joint prefix is found, the KV blocks are attached to the request + /// and a guard is returned to prevent eviction until the recurrent state + /// is copied (see `Qwen35PrefixCache::finish_restore`). + /// If no joint prefix is found, the request is still created and returned. + pub(crate) fn begin_request( + &mut self, + prompt_tokens: &[u32], + max_output_tokens: usize, + lora_name: Option<&str>, + allow_match: bool, + ) -> Result<(RequestKv, Option)> { + let mut request = + self.kv + .pool() + .new_request(prompt_tokens.to_vec(), max_output_tokens, lora_name); + if !self.enabled || !allow_match { + return Ok((request, None)); + } + + let probe = self + .kv + .pool() + .probe_prefix(prompt_tokens.to_vec(), lora_name); + let resident_tokens = probe.reusable_blocks() * self.kv.pool().block_size(); + let mut saw_eligible_kv = false; + let started = Instant::now(); + for boundary in eligible_boundaries(resident_tokens, SNAPSHOT_STRIDE_TOKENS) { + saw_eligible_kv = true; + let Some(sequence_hash) = probe.boundary_hash(boundary) else { + continue; + }; + let key = PrefixBoundaryKey { + sequence_hash, + boundary_tokens: boundary, + }; + let Some(guard) = self.state.lookup(key, started) else { + self.stats.snapshot_misses += 1; + continue; + }; + + let max_blocks = boundary / self.kv.pool().block_size(); + let attached = match request.match_and_add_prefix_up_to(self.kv.pool(), max_blocks) { + Ok(attached) => attached, + Err(error) => { + let _ = request.release(); + return Err(error); + } + }; + if attached != boundary { + let _ = request.release(); + anyhow::bail!( + "Qwen3.5 joint prefix attach selected {boundary} tokens but attached {attached}" + ); + } + return Ok((request, Some(guard))); + } + + if saw_eligible_kv { + self.stats.kv_only_fallbacks += 1; + } + Ok((request, None)) + } + + /// Finish a restore after the caller has copied the recurrent state. + /// + /// `begin_request` performs the lookup and KV blocks' attachment, but the physical + /// recurrent-state copy is executor-specific: single-GPU execution copies + /// from the local store, while TP execution coordinates the copy across + /// workers. This method therefore only be called after the recurrent state + /// is copied. It validates the boundary and releases the guard. + pub(crate) fn finish_restore( + &mut self, + request: &RequestKv, + guard: PrefixGuard, + recurrent_positions: &[usize], + ) -> Result { + let boundary = guard.boundary(); + anyhow::ensure!( + request.kv_position() == boundary + && !recurrent_positions.is_empty() + && recurrent_positions + .iter() + .all(|&position| position == boundary), + "Qwen3.5 joint prefix restore position mismatch: kv={}, recurrent={recurrent_positions:?}, boundary={}", + request.kv_position(), + boundary, + ); + self.stats.joint_hits += 1; + self.stats.joint_hit_tokens += boundary as u64; + self.stats.restore_ns = self + .stats + .restore_ns + .saturating_add(guard.started.elapsed().as_nanos() as u64); + // End the cache pin only after every physical restore was checked. + drop(guard); + Ok(boundary) + } + + /// Reserve the KV pages required by the next prefill forward. + pub(crate) fn schedule_prefill(&self, request: &mut RequestKv, tokens: usize) -> Result<()> { + request + .schedule_prefill(tokens, self.kv.pool()) + .map_err(|e| anyhow::anyhow!("Qwen3.5 prefill KV schedule failed: {e}")) + } + + /// Build the exact, immutable KV page-table view for prefill kernels. + #[allow(clippy::unused_self)] // keep KV state transitions behind this facade + pub(crate) fn prefill_view(&self, request: &RequestKv, tokens: usize) -> KvView { + request.prefill_view(tokens) + } + + /// Apply one successful whole-model prefill window. + /// + /// KV is applied for every window. Snapshot publication is attempted + /// only at a non-zero multiple of [`SNAPSHOT_STRIDE_TOKENS`]. A non-aligned + /// prompt tail still applies successfully but does not create a snapshot. + pub(crate) fn apply_prefill( + &self, + request: &mut RequestKv, + first_token: Option, + ) -> Result { + if let Some(first_token) = first_token { + request.apply_prefill(first_token, self.kv.pool())?; + } else { + request.apply_prefill_chunk(self.kv.pool())?; + } + let boundary = request.kv_position(); + Ok(boundary) + } + + /// Reserve a recurrent-state slot for an eligible applied boundary. + /// + /// Return a reservation handle when the boundary is alighed to [`SNAPSHOT_STRIDE_TOKENS`]. + /// The reservation is later performed locally at every rank. + /// Return `None` when the boundary is not aligned or when the cache is disabled. + pub(crate) fn reserve_prefix( + &mut self, + request: &RequestKv, + boundary: usize, + ) -> Result> { + if !self.enabled { + return Ok(None); + } + if boundary == 0 || !boundary.is_multiple_of(SNAPSHOT_STRIDE_TOKENS) { + return Ok(None); + } + let sequence_hash = request + .registered_boundary_hash(boundary) + .ok_or_else(|| anyhow::anyhow!("no registered KV hash at boundary {boundary}"))?; + let key = PrefixBoundaryKey { + sequence_hash, + boundary_tokens: boundary, + }; + Ok(self.state.reserve(key)) + } + + /// Insert a prefix entry to the cache. + /// + /// The caller must have already attached the KV blocks and copied the recurrent state to the reserved slot. + pub(crate) fn publish_prefix(&mut self, request: &RequestKv, reservation: PrefixReservation) { + let block_count = reservation.key.boundary_tokens / self.kv.pool().block_size(); + let mut kv_guards = request.assigned_block_guards(); + assert!( + kv_guards.len() >= block_count, + "Qwen3.5 snapshot boundary {} requires {block_count} KV blocks, request has {}", + reservation.key.boundary_tokens, + kv_guards.len() + ); + kv_guards.truncate(block_count); + self.stats.inserts += 1; + if reservation.was_replacement() { + self.stats.evictions += 1; + } + self.state.publish(reservation, kv_guards); + } + + /// Abort a prefix reservation after any rank-local copy failure. + pub(crate) fn abort_prefix(&mut self, reservation: PrefixReservation) { + self.state.abort(reservation); + } + + /// Reserve KV capacity for the next one-token decode forward. + pub(crate) fn schedule_decode(&self, request: &mut RequestKv) -> Result<()> { + request + .schedule_decode(self.kv.pool()) + .map_err(|e| anyhow::anyhow!("Qwen3.5 decode KV schedule failed: {e}")) + } + + /// Build the exact, immutable KV page-table view for decode kernels. + #[allow(clippy::unused_self)] // keep KV state transitions behind this facade + pub(crate) fn decode_view(&self, request: &RequestKv) -> KvView { + request.decode_view() + } + + /// Apply the KV written by decode and record the newly sampled token. + pub(crate) fn apply_decode(&self, request: &mut RequestKv, token: u32) -> Result<()> { + request.apply_decode(token, self.kv.pool())?; + Ok(()) + } + + /// Roll back pages reserved by a scheduled step that did not apply. + #[allow(clippy::unused_self)] // keep KV state transitions behind this facade + pub(crate) fn revert_schedule(&self, request: &mut RequestKv) -> Result<()> { + request.revert_schedule() + } + + /// Release all request KV. + #[allow(clippy::unused_self)] // keep KV state transitions behind this facade + pub(crate) fn release_request(&self, request: &mut RequestKv) -> Result<()> { + request.mark_blocks_reset_on_release(); + request.release() + } +} + +/// Yield reusable snapshot boundaries from longest to shortest. +fn eligible_boundaries(resident_tokens: usize, stride: usize) -> impl Iterator { + let highest = resident_tokens / stride * stride; + (1..=highest / stride).rev().map(move |n| n * stride) +} + +#[cfg(test)] +mod tests { + use std::time::Instant; + + use pegainfer_kv_cache::BlockPool; + + use super::PrefixBoundaryKey; + use super::PrefixCacheState; + use super::SNAPSHOT_STRIDE_TOKENS; + use super::eligible_boundaries; + + fn key(tag: u8) -> PrefixBoundaryKey { + PrefixBoundaryKey { + sequence_hash: [tag; 16], + boundary_tokens: 256, + } + } + + #[test] + fn joint_boundaries_descend_on_snapshot_stride() { + assert_eq!( + eligible_boundaries(255, 256).collect::>(), + Vec::::new() + ); + assert_eq!(eligible_boundaries(256, 256).collect::>(), [256]); + assert_eq!( + eligible_boundaries(900, 256).collect::>(), + [768, 512, 256] + ); + } + + #[test] + fn state_publishes_only_after_explicit_commit() { + let mut state = PrefixCacheState::new(1); + let reservation = state + .reserve(key(1)) + .expect("empty state must reserve a write"); + assert!(state.lookup(key(1), Instant::now()).is_none()); + state.publish(reservation, Vec::new()); + assert!(state.lookup(key(1), Instant::now()).is_some()); + } + + #[test] + fn aborted_write_does_not_expose_partial_snapshot() { + let mut state = PrefixCacheState::new(1); + let reservation = state + .reserve(key(1)) + .expect("empty state must reserve a write"); + state.abort(reservation); + assert!(state.lookup(key(1), Instant::now()).is_none()); + assert!(state.reserve(key(2)).is_some()); + } + + #[test] + fn duplicate_refreshes_lru_without_allocating_a_slot() { + let mut state = PrefixCacheState::new(1); + let reservation = state + .reserve(key(1)) + .expect("empty directory must reserve a write"); + state.publish(reservation, Vec::new()); + assert!(state.reserve(key(1)).is_none()); + assert_eq!(state.len(), 1); + } + + #[test] + fn lru_evicts_untouched_snapshot_but_preserves_touched_snapshot() { + let mut state = PrefixCacheState::new(2); + for tag in [1, 2] { + let reservation = state + .reserve(key(tag)) + .expect("state should have a free slot"); + state.publish(reservation, Vec::new()); + } + drop( + state + .lookup(key(1), Instant::now()) + .expect("key 1 should be present"), + ); + let reservation = state + .reserve(key(3)) + .expect("an unpinned LRU victim should be available"); + state.publish(reservation, Vec::new()); + assert!(state.lookup(key(1), Instant::now()).is_some()); + assert!(state.lookup(key(2), Instant::now()).is_none()); + assert!(state.lookup(key(3), Instant::now()).is_some()); + } + + #[test] + fn all_pinned_snapshots_make_insertion_a_soft_miss() { + let mut state = PrefixCacheState::new(1); + let reservation = state + .reserve(key(1)) + .expect("empty state must reserve a write"); + state.publish(reservation, Vec::new()); + let guard = state + .lookup(key(1), Instant::now()) + .expect("key 1 should be present"); + assert!(state.reserve(key(2)).is_none()); + drop(guard); + assert!(state.reserve(key(2)).is_some()); + } + + #[test] + fn joint_cache_entry_releases_kv_blocks_on_eviction() { + let pool = BlockPool::new(16, 32).expect("block pool"); + let baseline = pool.available_blocks(); + let mut request = pool.new_request(vec![7; SNAPSHOT_STRIDE_TOKENS], 0, None); + request + .schedule_prefill(SNAPSHOT_STRIDE_TOKENS, &pool) + .expect("schedule prefill"); + request.apply_prefill_chunk(&pool).expect("apply prefill"); + + let mut state = PrefixCacheState::new(1); + let reservation = state + .reserve(key(1)) + .expect("empty state must reserve a write"); + state.publish(reservation, request.assigned_block_guards()); + request.mark_blocks_reset_on_release(); + request.release().expect("release request"); + + let retained_blocks = SNAPSHOT_STRIDE_TOKENS / pool.block_size(); + assert_eq!(pool.available_blocks(), baseline - retained_blocks); + + let reservation = state + .reserve(key(2)) + .expect("full directory must evict its unpinned snapshot"); + assert_eq!(pool.available_blocks(), baseline); + state.abort(reservation); + } +} diff --git a/pegainfer-qwen35/src/recurrent_state.rs b/pegainfer-qwen35/src/recurrent_state.rs index f7279ec12..5200e2925 100644 --- a/pegainfer-qwen35/src/recurrent_state.rs +++ b/pegainfer-qwen35/src/recurrent_state.rs @@ -70,6 +70,26 @@ impl RecurrentState { Ok(Self { layers, seq_len: 0 }) } + + /// Copy one complete target-model recurrent state into this allocation. + pub(crate) fn copy_from(&mut self, ctx: &DeviceContext, src: &Self) -> Result<()> { + anyhow::ensure!( + self.layers.len() == src.layers.len(), + "Qwen3.5 recurrent copy layer mismatch: dst={}, src={}", + self.layers.len(), + src.layers.len() + ); + for (layer_idx, (dst, src)) in self.layers.iter_mut().zip(&src.layers).enumerate() { + ctx.stream + .memcpy_dtod(&src.state, &mut dst.state) + .map_err(|e| anyhow::anyhow!("copy recurrent layer {layer_idx}: {e}"))?; + ctx.stream + .memcpy_dtod(&src.conv_state.data, &mut dst.conv_state.data) + .map_err(|e| anyhow::anyhow!("copy conv state layer {layer_idx}: {e}"))?; + } + self.seq_len = src.seq_len; + Ok(()) + } } impl LinearStatePointerTables { diff --git a/pegainfer-qwen35/src/scheduler.rs b/pegainfer-qwen35/src/scheduler.rs index c016da565..6e1fbda35 100644 --- a/pegainfer-qwen35/src/scheduler.rs +++ b/pegainfer-qwen35/src/scheduler.rs @@ -1,7 +1,7 @@ //! Scheduler for Qwen3.5: dedicated GPU thread that batches concurrent requests. //! //! Mirrors the Qwen3 scheduler but manages: -//! - `RecurrentState` alongside `KvState` (linear attention layers) +//! - controller-owned `RequestKv` plus recurrent state (hybrid attention) //! - `BatchDecodeGraphState` for CUDA Graph batch decode (stable-address slots) mod plan; @@ -27,9 +27,10 @@ use pegainfer_core::engine::TokenEvent; use pegainfer_core::engine::TokenLogprob; use pegainfer_core::engine::TokenSink; use pegainfer_core::engine::panic_message; -use pegainfer_core::kv_pool::KvState; use pegainfer_core::sampler::SamplingParams; use pegainfer_core::tensor::HiddenStates; +use pegainfer_kv_cache::KvCacheManager; +use pegainfer_kv_cache::RequestKv; use rand::SeedableRng; use rand::rngs::StdRng; use tokio::sync::mpsc; @@ -56,6 +57,8 @@ use crate::executor::PrefillRequestResult; use crate::executor::PrefillResult; use crate::executor::RequestId; use crate::logprobs::snapshot_requested_logprobs; +use crate::prefix_cache::Qwen35PrefixCache; +use crate::prefix_cache::RecurrentStateStore; use crate::recurrent_state::RecurrentState; use crate::tp_executor::Qwen35TpExecutor; use crate::tp_executor::TpDecodeStepItem; @@ -93,7 +96,7 @@ struct PrefillingRequest35 { enum ActiveBackendState { Single { - kv: KvState, + kv: Box, /// Index into `BatchDecodeGraphState.slot_states`. graph_slot_idx: usize, }, @@ -103,8 +106,32 @@ enum ActiveBackendState { } enum PrefillBackendState { - Single { kv: KvState, rec: RecurrentState }, - Tp { request_id: RequestId }, + Single { + kv: Box, + rec: RecurrentState, + }, + Tp { + request_id: RequestId, + }, +} + +fn active_request_kv(request: &mut ActiveRequest35) -> Option<&mut RequestKv> { + match &mut request.backend_state { + ActiveBackendState::Single { kv, .. } => Some(kv), + ActiveBackendState::Tp { .. } => None, + } +} + +/// Roll back a set of requests scheduled by the current scheduler step. +fn revert_scheduled_requests<'a>( + kv_cache: &Qwen35PrefixCache, + requests: impl IntoIterator, +) { + for request in requests { + if let Err(error) = kv_cache.revert_schedule(request) { + warn!("failed to revert Qwen3.5 scheduler KV schedule: {error}"); + } + } } pub const DEFAULT_MAX_PREFILL_TOKENS: usize = 1024; @@ -129,6 +156,12 @@ fn itl_debug_mono_us() -> u128 { ORIGIN.get_or_init(Instant::now).elapsed().as_micros() } +fn unix_now_s() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0.0, |duration| duration.as_secs_f64()) +} + // ── Entry point ───────────────────────────────────────────────────────── pub fn start_with_capacity( @@ -157,18 +190,17 @@ pub(crate) fn start_with_capacity_and_policy( max_prefill_tokens > 0, "max_prefill_tokens must be positive: a zero budget can never schedule a prefill chunk" ); + let backend = SingleGpuBackend::new(model, max_batch)?; // Static instance cap for the vLLM bridge's max_model_len. Live admission // still uses the current page budget inside the scheduler loop. - let total_blocks = model.kv_pool().capacity_pages().saturating_sub(1); + let total_blocks = backend.kv_cache.pool().max_request_blocks(); let kv_total_blocks = total_blocks as u64; - let block_size = model.kv_pool().layout().page_size; + let block_size = backend.kv_cache.pool().block_size(); let servable = servable_len( - model.config().max_position_embeddings, + backend.model.config().max_position_embeddings, total_blocks, block_size, ); - let backend = SingleGpuBackend::new(model, max_batch)?; - let (submit_tx, submit_rx) = mpsc::unbounded_channel(); let (startup_tx, startup_rx) = std_mpsc::channel(); let (load_tx, load_rx) = watch::channel(LoadSnapshot { @@ -224,13 +256,19 @@ pub(crate) fn start_tp_with_capacity( device_ordinals: &[usize], max_batch: usize, max_prefill_tokens: usize, + prefix_snapshot_bytes: usize, ) -> Result { assert!( max_prefill_tokens > 0, "max_prefill_tokens must be positive: a zero budget can never schedule a prefill chunk" ); - let backend = - TpSchedulerBackend::new(model_path, device_ordinals, max_batch, max_prefill_tokens)?; + let backend = TpSchedulerBackend::new( + model_path, + device_ordinals, + max_batch, + max_prefill_tokens, + prefix_snapshot_bytes, + )?; let servable = servable_len( backend.max_position_embeddings(), backend.capacity_pages_for_requests(), @@ -270,6 +308,8 @@ pub(crate) fn start_tp_with_capacity( struct SingleGpuBackend { model: Qwen35Model, + kv_cache: Qwen35PrefixCache, + recurrent_store: RecurrentStateStore, graph_state: BatchDecodeGraphState, } @@ -288,9 +328,32 @@ struct TpSchedulerBackend { impl SingleGpuBackend { fn new(model: Qwen35Model, max_batch: usize) -> Result { anyhow::ensure!(max_batch > 0, "Qwen3.5 max_batch must be > 0"); + let manager = + KvCacheManager::from_buffer(model.kv_buffer().clone(), model.kv_buffer().num_blocks())?; + let kv_cache = Qwen35PrefixCache::new(manager, model.prefix_snapshot_slots())?; + let recurrent_store = RecurrentStateStore::new( + model.device_ctx(), + model.config(), + model.prefix_snapshot_slots(), + )?; + debug_assert_eq!(recurrent_store.len(), kv_cache.snapshot_slots()); let graph_capacity = crate::batch_decode_graph::bucket_for(max_batch); - let graph_state = model.create_batch_decode_graph_state_with_capacity(graph_capacity)?; - Ok(Self { model, graph_state }) + let graph_state = model.create_batch_decode_graph_state_with_capacity( + graph_capacity, + kv_cache.pool().total_blocks(), + kv_cache.pool().padding_block_id(), + )?; + info!( + "Qwen3.5 prefix cache: enabled={}, snapshot_slots={}", + kv_cache.enabled(), + kv_cache.snapshot_slots() + ); + Ok(Self { + model, + kv_cache, + recurrent_store, + graph_state, + }) } fn model(&self) -> &Qwen35Model { @@ -307,37 +370,91 @@ impl SingleGpuBackend { } fn page_size(&self) -> usize { - self.model.kv_pool().layout().page_size + self.kv_cache.pool().block_size() } fn available_pages(&self) -> usize { - self.model.kv_pool().available_pages() + self.kv_cache.pool().available_blocks() } fn capacity_pages_for_requests(&self) -> usize { - self.model.kv_pool().capacity_pages().saturating_sub(1) + self.kv_cache.pool().max_request_blocks() } fn max_position_embeddings(&self) -> usize { self.model.config().max_position_embeddings } - fn alloc_kv(&self) -> KvState { - self.model.alloc_kv() - } - fn alloc_recurrent(&self) -> Result { RecurrentState::new(self.model.device_ctx(), self.model.config()) } + fn alloc_prefill_state( + &mut self, + req: &SchedulerRequest, + ) -> Result<(PrefillBackendState, usize)> { + let mut rec = self.alloc_recurrent()?; + let (mut kv, restore) = self.kv_cache.begin_request( + &req.prompt_tokens, + req.max_tokens, + req.lora_adapter.as_deref(), + !req.echo, + )?; + let cached_tokens = if let Some(restore) = restore { + if let Err(error) = self.recurrent_store.restore( + self.model.device_ctx(), + restore.recurrent_slot(), + &mut rec, + ) { + let _ = self.kv_cache.release_request(&mut kv); + return Err(error); + } + match self.kv_cache.finish_restore(&kv, restore, &[rec.seq_len]) { + Ok(tokens) => tokens, + Err(error) => { + let _ = self.kv_cache.release_request(&mut kv); + return Err(error); + } + } + } else { + 0 + }; + Ok(( + PrefillBackendState::Single { + kv: Box::new(kv), + rec, + }, + cached_tokens, + )) + } + fn batch_prefill_logits(&self, chunk: &mut ScheduledChunk) -> Result { let window_refs: Vec<&[u32]> = chunk.windows.iter().map(Vec::as_slice).collect(); let ScheduledChunkBackendState::Single { kvs, recs } = &mut chunk.backend_state else { anyhow::bail!("single-GPU prefill received TP chunk state"); }; + for (scheduled, (kv, window)) in kvs.iter_mut().zip(&chunk.windows).enumerate() { + if let Err(error) = self.kv_cache.schedule_prefill(kv, window.len()) { + revert_scheduled_requests(&self.kv_cache, kvs.iter_mut().take(scheduled)); + return Err(error); + } + } + let views = kvs + .iter() + .zip(&chunk.windows) + .map(|(kv, window)| self.kv_cache.prefill_view(kv, window.len())) + .collect::>(); let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); - self.model - .batch_prefill_logits(&window_refs, kvs, &mut rec_refs) + let result = self.model.batch_prefill_logits( + &window_refs, + &views, + &mut rec_refs, + self.kv_cache.buffer(), + ); + if result.is_err() { + revert_scheduled_requests(&self.kv_cache, kvs); + } + result } fn unified_step( @@ -349,40 +466,172 @@ impl SingleGpuBackend { let ScheduledChunkBackendState::Single { kvs, recs } = &mut chunk.backend_state else { anyhow::bail!("single-GPU unified step received TP chunk state"); }; + for (scheduled_prefills, (kv, window)) in kvs.iter_mut().zip(&chunk.windows).enumerate() { + if let Err(error) = self.kv_cache.schedule_prefill(kv, window.len()) { + revert_scheduled_requests(&self.kv_cache, kvs.iter_mut().take(scheduled_prefills)); + return Err(error); + } + } + let prefill_views = kvs + .iter() + .zip(&chunk.windows) + .map(|(kv, window)| self.kv_cache.prefill_view(kv, window.len())) + .collect::>(); let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); let decode_tokens: Vec = active.iter().map(|r| r.last_token).collect(); - let mut decode_kv_refs: Vec<&mut KvState> = active - .iter_mut() - .map(|r| match &mut r.backend_state { - ActiveBackendState::Single { kv, .. } => kv, + for (scheduled_decodes, req) in active.iter_mut().enumerate() { + let ActiveBackendState::Single { kv, .. } = &mut req.backend_state else { + panic!("single-GPU unified step received TP active state") + }; + if let Err(error) = self.kv_cache.schedule_decode(kv) { + revert_scheduled_requests(&self.kv_cache, kvs); + revert_scheduled_requests( + &self.kv_cache, + active + .iter_mut() + .take(scheduled_decodes) + .filter_map(active_request_kv), + ); + return Err(error); + } + } + let decode_views = active + .iter() + .map(|r| match &r.backend_state { + ActiveBackendState::Single { kv, .. } => self.kv_cache.decode_view(kv), ActiveBackendState::Tp { .. } => { panic!("single-GPU unified step received TP active state") } }) - .collect(); - self.model.unified_step( + .collect::>(); + let result = self.model.unified_step( &window_refs, - kvs, + &prefill_views, &mut rec_refs, &decode_tokens, - &mut decode_kv_refs, + &decode_views, + self.kv_cache.buffer(), &mut self.graph_state, - ) + ); + if result.is_err() { + revert_scheduled_requests(&self.kv_cache, kvs); + revert_scheduled_requests( + &self.kv_cache, + active.iter_mut().filter_map(active_request_kv), + ); + } + result } fn decode_graph(&mut self, active: &mut [ActiveRequest35]) -> Result<()> { let token_ids: Vec = active.iter().map(|r| r.last_token).collect(); - let mut kv_refs: Vec<&mut KvState> = active - .iter_mut() - .map(|r| match &mut r.backend_state { - ActiveBackendState::Single { kv, .. } => kv, - ActiveBackendState::Tp { .. } => { - panic!("single-GPU decode received TP active state") - } + for (scheduled, req) in active.iter_mut().enumerate() { + let ActiveBackendState::Single { kv, .. } = &mut req.backend_state else { + panic!("single-GPU decode received TP active state") + }; + if let Err(error) = self.kv_cache.schedule_decode(kv) { + revert_scheduled_requests( + &self.kv_cache, + active + .iter_mut() + .take(scheduled) + .filter_map(active_request_kv), + ); + return Err(error); + } + } + let views = active + .iter() + .map(|r| match &r.backend_state { + ActiveBackendState::Single { kv, .. } => self.kv_cache.decode_view(kv), + ActiveBackendState::Tp { .. } => panic!("single-GPU decode received TP state"), }) - .collect(); - self.model - .batch_decode_graph(&token_ids, &mut kv_refs, &mut self.graph_state) + .collect::>(); + let result = self.model.batch_decode_graph( + &token_ids, + &views, + self.kv_cache.buffer(), + &mut self.graph_state, + ); + if result.is_err() { + revert_scheduled_requests( + &self.kv_cache, + active.iter_mut().filter_map(active_request_kv), + ); + } + result + } + + fn apply_prefill(&mut self, chunk: &mut ScheduledChunk, tokens: &[u32]) -> Result<()> { + let ScheduledChunkBackendState::Single { kvs, recs } = &mut chunk.backend_state else { + anyhow::bail!("single-GPU commit received TP chunk state") + }; + for (i, (kv, rec)) in kvs.iter_mut().zip(recs.iter()).enumerate() { + let is_final = chunk.ends[i] == chunk.reqs[i].prompt_tokens.len(); + let boundary = self + .kv_cache + .apply_prefill(kv, is_final.then_some(tokens[i]))?; + anyhow::ensure!( + rec.seq_len == boundary, + "Qwen3.5 prefill apply position mismatch: kv={boundary}, recurrent={}", + rec.seq_len + ); + if let Some(reservation) = self.kv_cache.reserve_prefix(kv, boundary)? { + if let Err(error) = self.recurrent_store.save( + self.model.device_ctx(), + reservation.recurrent_slot(), + rec, + ) { + self.kv_cache.abort_prefix(reservation); + return Err(error); + } + self.kv_cache.publish_prefix(kv, reservation); + } + } + Ok(()) + } + + fn apply_decode(&self, active: &mut [ActiveRequest35], tokens: &[u32]) -> Result<()> { + anyhow::ensure!(active.len() == tokens.len(), "decode apply row mismatch"); + for (req, &token) in active.iter_mut().zip(tokens) { + let ActiveBackendState::Single { kv, .. } = &mut req.backend_state else { + anyhow::bail!("single-GPU decode apply received TP state") + }; + self.kv_cache.apply_decode(kv, token)?; + } + Ok(()) + } + + fn fail_active(&self, active: &mut Vec, message: &str) { + for mut req in active.drain(..) { + if let ActiveBackendState::Single { kv, .. } = &mut req.backend_state { + if let Err(error) = self.kv_cache.release_request(kv) { + warn!("failed to release Qwen3.5 request KV: {error}"); + } + } + let _ = req.token_tx.send(TokenEvent::Error { + message: message.to_string(), + prompt_tokens: req.prompt_len, + completion_tokens: req.generated_count, + }); + } + } + + fn log_prefix_cache_stats(&self) { + let cache = &self.kv_cache; + let stats = cache.stats(); + info!( + "Qwen3.5 prefix cache summary: joint_hits={}, hit_tokens={}, kv_only_fallbacks={}, snapshot_misses={}, inserts={}, evictions={}, restore_ms={:.3}, occupancy={}/{}", + stats.joint_hits, + stats.joint_hit_tokens, + stats.kv_only_fallbacks, + stats.snapshot_misses, + stats.inserts, + stats.evictions, + stats.restore_ns as f64 / 1_000_000.0, + cache.snapshot_occupancy(), + cache.snapshot_slots(), + ); } fn sample_prefill_logits( @@ -524,13 +773,15 @@ impl TpSchedulerBackend { device_ordinals: &[usize], max_batch: usize, max_prefill_tokens: usize, + prefix_snapshot_bytes: usize, ) -> Result { - let executor = Qwen35TpExecutor::from_runtime_with_limits( + let executor = Qwen35TpExecutor::from_runtime_with_limits_and_prefix( model_path, false, device_ordinals, max_batch, max_prefill_tokens, + prefix_snapshot_bytes, )?; Ok(Self { executor, @@ -544,6 +795,21 @@ impl TpSchedulerBackend { id } + fn alloc_prefill_state( + &mut self, + req: &SchedulerRequest, + ) -> Result<(PrefillBackendState, usize)> { + let request_id = self.alloc_request_id(); + let cached_tokens = self.executor.begin_request( + request_id, + &req.prompt_tokens, + req.max_tokens, + req.lora_adapter.as_deref(), + !req.echo, + )?; + Ok((PrefillBackendState::Tp { request_id }, cached_tokens)) + } + fn max_batch(&self) -> usize { self.executor.max_batch() } @@ -564,26 +830,12 @@ impl TpSchedulerBackend { self.executor.is_stop_token(token) } - fn available_pages( - &self, - active: &[ActiveRequest35], - prefilling: &[PrefillingRequest35], - ) -> usize { - let page_size = self.page_size(); - let active_pages: usize = active - .iter() - .map(|req| pages_needed(current_active_tokens(req), page_size)) - .sum(); - let prefilling_pages: usize = prefilling - .iter() - .map(|req| pages_needed(req.cursor, page_size)) - .sum(); - self.capacity_pages_for_requests() - .saturating_sub(active_pages.saturating_add(prefilling_pages)) + fn available_pages(&self) -> usize { + self.executor.available_pages() } fn execute_prefill_chunk( - &self, + &mut self, chunk: &ScheduledChunk, sample_seed: u64, ) -> Result<(Vec, Vec>)> { @@ -613,7 +865,7 @@ impl TpSchedulerBackend { } fn execute_decode( - &self, + &mut self, active: &[ActiveRequest35], sample_seed: u64, ) -> Result<(Vec, Vec>)> { @@ -635,7 +887,7 @@ impl TpSchedulerBackend { align_decode_results(active, &result) } - fn drop_request(&self, request_id: RequestId) { + fn drop_request(&mut self, request_id: RequestId) { if let Err(err) = self.executor.drop_request(request_id) { warn!( "failed to drop Qwen3.5 TP worker request {}: {err}", @@ -667,7 +919,10 @@ impl SchedulerBackend { ) -> usize { match self { Self::Single(backend) => backend.available_pages(), - Self::Tp(backend) => backend.available_pages(active, prefilling), + Self::Tp(backend) => { + let _ = (active, prefilling); + backend.available_pages() + } } } @@ -685,15 +940,25 @@ impl SchedulerBackend { } } - fn alloc_prefill_state(&mut self) -> Result { + fn alloc_prefill_state( + &mut self, + req: &SchedulerRequest, + ) -> Result<(PrefillBackendState, usize)> { match self { - Self::Single(backend) => Ok(PrefillBackendState::Single { - kv: backend.alloc_kv(), - rec: backend.alloc_recurrent()?, - }), - Self::Tp(backend) => Ok(PrefillBackendState::Tp { - request_id: backend.alloc_request_id(), - }), + Self::Single(backend) => backend.alloc_prefill_state(req), + Self::Tp(backend) => backend.alloc_prefill_state(req), + } + } + + fn snapshot_stride(&self) -> Option { + match self { + Self::Single(backend) if backend.kv_cache.enabled() => { + Some(crate::prefix_cache::SNAPSHOT_STRIDE_TOKENS) + } + Self::Tp(backend) if backend.executor.prefix_cache_enabled() => { + Some(crate::prefix_cache::SNAPSHOT_STRIDE_TOKENS) + } + Self::Single(_) | Self::Tp(_) => None, } } @@ -707,15 +972,13 @@ impl SchedulerBackend { Self::Tp(backend) => backend.is_stop_token(token), } } -} - -fn current_active_tokens(req: &ActiveRequest35) -> usize { - req.prompt_len - .saturating_add(req.generated_count.saturating_sub(1)) -} -fn pages_needed(token_count: usize, page_size: usize) -> usize { - token_count.div_ceil(page_size) + fn log_prefix_cache_stats(&self) { + match self { + Self::Single(backend) => backend.log_prefix_cache_stats(), + Self::Tp(backend) => backend.executor.log_prefix_cache_stats(), + } + } } fn align_prefill_results( @@ -881,6 +1144,7 @@ fn scheduler_loop( pending.push(req); } else { info!("scheduler: all handles dropped, exiting"); + backend.log_prefix_cache_stats(); return; } while let Ok((req, _kv_prefix)) = submit_rx.try_recv() { @@ -911,16 +1175,25 @@ fn scheduler_loop( .available_pages(&active, &prefilling) .saturating_sub(prefilling_future_pages(&prefilling_budget, page_size)); let decode_batching_slot = max_batch.saturating_sub(prefilling.len()); + // Keep admission's protocol-level prompt + max_tokens limit identical + // to the max_model_len advertised by the handle. The content-hashed + // pool reserves one padding block, so its physical cap can be below + // the model's configured context length. + let max_context_tokens = servable_len( + backend.max_position_embeddings(), + backend.capacity_pages_for_requests(), + page_size, + ) as usize; let admission = admit_pending_requests( pending, &active_budget, decode_batching_slot, page_size, page_budget, - // KvPool capacity includes the CUDA Graph padding page reserved at + // The block pool includes the CUDA Graph padding page reserved at // construction, so a real request can use at most the remaining pages. backend.capacity_pages_for_requests(), - backend.max_position_embeddings(), + max_context_tokens, |req| req.prompt_tokens.len(), |req| req.max_tokens, ); @@ -936,13 +1209,29 @@ fn scheduler_loop( req.prompt_tokens.len(), req.max_tokens ); - match backend.alloc_prefill_state() { - Ok(backend_state) => prefilling.push(PrefillingRequest35 { - backend_state, - cursor: 0, - step_chunk: 0, - req, - }), + match backend.alloc_prefill_state(&req) { + Ok((backend_state, cached_tokens)) => { + let scheduled_at_unix_s = unix_now_s(); + if req + .token_tx + .send(TokenEvent::Scheduled { + queued_at_unix_s: req.queued_at_unix_s.unwrap_or(scheduled_at_unix_s), + scheduled_at_unix_s, + prompt_tokens: req.prompt_tokens.len(), + cached_tokens, + }) + .is_err() + { + backend.drop_prefill_state(backend_state); + continue; + } + prefilling.push(PrefillingRequest35 { + backend_state, + cursor: cached_tokens, + step_chunk: 0, + req, + }); + } Err(e) => { warn!("failed to allocate recurrent state for new request: {e}"); let _ = req.token_tx.send(TokenEvent::Error { @@ -979,7 +1268,11 @@ fn scheduler_loop( &active_decode, &prefill_queue, ); - let scheduled = take_prefill_chunks(&mut prefilling, step_prefill_budget); + let scheduled = take_prefill_chunks( + &mut prefilling, + step_prefill_budget, + backend.snapshot_stride(), + ); // ITL diagnostics (#470): capture the *actual* prefill-chunk token count // and the frozen decode width for this step before the plan consumes the // scheduled set. Off unless PEGAINFER_ITL_DEBUG is set. @@ -1095,25 +1388,34 @@ fn prefill_batch( Ok(v) => v, Err(e) => { warn!("batch prefill failed: {e}"); - fail_chunk(chunk, &e.to_string()); + fail_chunk(single, chunk, &e.to_string()); return; } }; - match single.sample_prefill_logits(&chunk.reqs, &logits, rng) { + let sampled = match single.sample_prefill_logits(&chunk.reqs, &logits, rng) { Ok(v) => v, Err(e) => { warn!("prefill sampling failed: {e}"); - fail_chunk(chunk, &e.to_string()); + if let ScheduledChunkBackendState::Single { kvs, .. } = &mut chunk.backend_state + { + revert_scheduled_requests(&single.kv_cache, kvs); + } + fail_chunk(single, chunk, &e.to_string()); return; } + }; + if let Err(e) = single.apply_prefill(&mut chunk, &sampled.0) { + warn!("prefill KV/snapshot commit failed: {e}"); + fail_chunk(single, chunk, &e.to_string()); + return; } + sampled } SchedulerBackend::Tp(tp) => match tp.execute_prefill_chunk(&chunk, sample_seed) { Ok(v) => v, Err(e) => { warn!("TP prefill chunk failed: {e}"); - drop_tp_chunk_state(tp, &chunk); - fail_chunk(chunk, &e.to_string()); + fail_chunk(backend, chunk, &e.to_string()); return; } }, @@ -1142,7 +1444,7 @@ fn unified_step_sched( completion_tokens: req.generated_count, }); } - fail_chunk(chunk, message); + fail_chunk(backend, chunk, message); return; }; let mut chunk = ScheduledChunk::from(scheduled); @@ -1154,38 +1456,75 @@ fn unified_step_sched( Err(e) => { warn!("unified step failed: {e}"); let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } - fail_chunk(chunk, &message); + backend.fail_active(active, &message); + fail_chunk(backend, chunk, &message); return; } }; - // Process decode results FIRST (it may retire requests and free graph slots - // that promotion then fills densely). - if output.decoded { - process_decode_logits(backend, active, rng); - } - let prefill_logits = output .prefill_logits .as_ref() .expect("scheduled prefill chunk must return prefill logits"); + let decode_sampled = if output.decoded { + match backend.sample_decode_logits(active, rng) { + Ok(sampled) => Some(sampled), + Err(e) => { + warn!("unified decode sampling failed: {e}"); + revert_scheduled_requests( + &backend.kv_cache, + active.iter_mut().filter_map(active_request_kv), + ); + if let ScheduledChunkBackendState::Single { kvs, .. } = &mut chunk.backend_state { + revert_scheduled_requests(&backend.kv_cache, kvs); + } + let message = e.to_string(); + backend.fail_active(active, &message); + fail_chunk(backend, chunk, &message); + return; + } + } + } else { + None + }; let (tokens, logprobs_vec) = match backend.sample_prefill_logits(&chunk.reqs, prefill_logits, rng) { Ok(v) => v, Err(e) => { warn!("unified prefill sampling failed: {e}"); - fail_chunk(chunk, &e.to_string()); + revert_scheduled_requests( + &backend.kv_cache, + active.iter_mut().filter_map(active_request_kv), + ); + if let ScheduledChunkBackendState::Single { kvs, .. } = &mut chunk.backend_state { + revert_scheduled_requests(&backend.kv_cache, kvs); + } + backend.fail_active(active, &e.to_string()); + fail_chunk(backend, chunk, &e.to_string()); return; } }; + if let Some((decode_tokens, _)) = &decode_sampled + && let Err(e) = backend.apply_decode(active, decode_tokens) + { + warn!("unified decode KV apply failed: {e}"); + backend.fail_active(active, &e.to_string()); + fail_chunk(backend, chunk, &e.to_string()); + return; + } + if let Err(e) = backend.apply_prefill(&mut chunk, &tokens) { + warn!("unified prefill KV/snapshot commit failed: {e}"); + backend.fail_active(active, &e.to_string()); + fail_chunk(backend, chunk, &e.to_string()); + return; + } + + // Decode commits and dispatches first; retirements free graph slots that + // the newly-prefilled requests can then occupy densely. + if let Some((decode_tokens, decode_logprobs)) = decode_sampled { + dispatch_decode_tokens(backend, active, &decode_tokens, &decode_logprobs); + } promote_or_requeue(backend, active, prefilling, chunk, &tokens, &logprobs_vec); } @@ -1202,13 +1541,7 @@ fn decode_step( if let Err(e) = single.decode_graph(active) { warn!("batch_decode_graph error: {e}"); let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } + single.fail_active(active, &message); return; } // Snapshot logits to CPU BEFORE sampling (sampling may modify bufs.logits) @@ -1216,14 +1549,12 @@ fn decode_step( Ok(v) => v, Err(e) => { warn!("decode sampling/logprobs error: {e}"); + revert_scheduled_requests( + &single.kv_cache, + active.iter_mut().filter_map(active_request_kv), + ); let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } + single.fail_active(active, &message); return; } } @@ -1249,31 +1580,14 @@ fn decode_step( }, }; - dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec); -} - -/// Process decode logits from unified step: sample, extract logprobs, dispatch. -fn process_decode_logits( - backend: &mut SingleGpuBackend, - active: &mut Vec, - rng: &mut StdRng, -) { - let (tokens, logprobs_vec) = match backend.sample_decode_logits(active, rng) { - Ok(v) => v, - Err(e) => { - warn!("decode sampling/logprobs error: {e}"); + if let SchedulerBackend::Single(single) = backend { + if let Err(e) = single.apply_decode(active, &tokens) { + warn!("decode KV apply failed: {e}"); let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } + single.fail_active(active, &message); return; } - }; - + } dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec); } @@ -1393,7 +1707,12 @@ fn compact_single_slot( idx: usize, ) { let compaction = compaction_after_retire(active.len(), idx); - active.swap_remove(idx); + let mut removed = active.swap_remove(idx); + if let ActiveBackendState::Single { kv, .. } = &mut removed.backend_state { + if let Err(error) = backend.kv_cache.release_request(kv) { + warn!("failed to release Qwen3.5 request KV: {error}"); + } + } if let Some(compaction) = compaction { backend.compact_slot(active, compaction); @@ -1414,7 +1733,7 @@ struct ScheduledChunk { enum ScheduledChunkBackendState { Single { - kvs: Vec, + kvs: Vec, recs: Vec, }, Tp { @@ -1455,7 +1774,7 @@ impl From> for ScheduledChunk { ScheduledChunkBackendState::Single { kvs, recs }, PrefillBackendState::Single { kv, rec }, ) => { - kvs.push(kv); + kvs.push(*kv); recs.push(rec); } ( @@ -1474,6 +1793,7 @@ impl From> for ScheduledChunk { fn take_prefill_chunks( prefilling: &mut Vec, prefill_budget: usize, + snapshot_stride: Option, ) -> Vec { let remaining: Vec = prefilling .iter() @@ -1482,28 +1802,28 @@ fn take_prefill_chunks( let chunks = plan_prefill_chunks(&remaining, prefill_budget); let mut scheduled: Vec = prefilling.drain(0..chunks.len()).collect(); for (p, chunk) in scheduled.iter_mut().zip(&chunks) { - p.step_chunk = *chunk; + p.step_chunk = clamp_prefill_chunk(p.cursor, *chunk, snapshot_stride); } scheduled } +fn clamp_prefill_chunk(cursor: usize, chunk: usize, snapshot_stride: Option) -> usize { + snapshot_stride.map_or(chunk, |stride| { + debug_assert!(stride > 0); + chunk.min(stride - cursor % stride) + }) +} + /// Report a forward/sampling failure to every request in the failed chunk. -fn fail_chunk(chunk: ScheduledChunk, message: &str) { - for req in chunk.reqs { +fn fail_chunk(backend: &mut impl PrefillPromoteBackend, chunk: ScheduledChunk, message: &str) { + let states = split_scheduled_backend_state(chunk.backend_state); + for (req, state) in chunk.reqs.into_iter().zip(states) { let _ = req.token_tx.send(TokenEvent::Error { message: message.to_string(), prompt_tokens: req.prompt_tokens.len(), completion_tokens: 0, }); - } -} - -fn drop_tp_chunk_state(backend: &TpSchedulerBackend, chunk: &ScheduledChunk) { - let ScheduledChunkBackendState::Tp { request_ids } = &chunk.backend_state else { - return; - }; - for &request_id in request_ids { - backend.drop_request(request_id); + backend.drop_prefill_state(state); } } @@ -1654,7 +1974,14 @@ impl PrefillPromoteBackend for SingleGpuBackend { } } - fn drop_prefill_state(&mut self, _state: PrefillBackendState) {} + fn drop_prefill_state(&mut self, state: PrefillBackendState) { + let PrefillBackendState::Single { mut kv, .. } = state else { + panic!("single-GPU drop received TP prefill state"); + }; + if let Err(error) = self.kv_cache.release_request(&mut kv) { + warn!("failed to release Qwen3.5 request KV: {error}"); + } + } } impl PrefillPromoteBackend for SchedulerBackend { @@ -1687,10 +2014,16 @@ impl PrefillPromoteBackend for SchedulerBackend { } fn drop_prefill_state(&mut self, state: PrefillBackendState) { - if let (SchedulerBackend::Tp(backend), PrefillBackendState::Tp { request_id }) = - (self, state) - { - backend.drop_request(request_id); + match (self, state) { + (SchedulerBackend::Single(backend), PrefillBackendState::Single { mut kv, .. }) => { + if let Err(error) = backend.kv_cache.release_request(&mut kv) { + warn!("failed to release Qwen3.5 request KV: {error}"); + } + } + (SchedulerBackend::Tp(backend), PrefillBackendState::Tp { request_id }) => { + backend.drop_request(request_id); + } + _ => panic!("mismatched Qwen3.5 scheduler backend state during drop"), } } } @@ -1702,7 +2035,10 @@ fn split_scheduled_backend_state( ScheduledChunkBackendState::Single { kvs, recs } => kvs .into_iter() .zip(recs) - .map(|(kv, rec)| PrefillBackendState::Single { kv, rec }) + .map(|(kv, rec)| PrefillBackendState::Single { + kv: Box::new(kv), + rec, + }) .collect(), ScheduledChunkBackendState::Tp { request_ids } => request_ids .into_iter() diff --git a/pegainfer-qwen35/src/scheduler/tests.rs b/pegainfer-qwen35/src/scheduler/tests.rs index 9d3b62241..76191b04c 100644 --- a/pegainfer-qwen35/src/scheduler/tests.rs +++ b/pegainfer-qwen35/src/scheduler/tests.rs @@ -5,6 +5,16 @@ use pegainfer_core::engine::EpBackend; use super::*; +#[test] +fn prefix_cache_chunking_stops_at_snapshot_boundaries() { + let stride = Some(crate::prefix_cache::SNAPSHOT_STRIDE_TOKENS); + assert_eq!(clamp_prefill_chunk(0, 900, stride), 256); + assert_eq!(clamp_prefill_chunk(256, 644, stride), 256); + assert_eq!(clamp_prefill_chunk(512, 388, stride), 256); + assert_eq!(clamp_prefill_chunk(768, 132, stride), 132); + assert_eq!(clamp_prefill_chunk(0, 900, None), 900); +} + #[test] fn send_rejection_reports_kv_lifetime_request_tokens() { let (token_tx, mut token_rx) = TokenSink::standalone(); @@ -86,8 +96,8 @@ fn tp_engine_rejects_cuda_graph_before_model_load() { fn tp2_scheduler_chunked_prefill_then_decode_smoke() { let model_path = std::env::var("PEGAINFER_TEST_MODEL_PATH") .unwrap_or_else(|_| "/home/data/mgj/qwen35weights".to_string()); - let handle = - start_tp_with_capacity(&model_path, 42, &[0, 1], 1, 1).expect("start Qwen3.5 TP scheduler"); + let handle = start_tp_with_capacity(&model_path, 42, &[0, 1], 1, 1, 0) + .expect("start Qwen3.5 TP scheduler"); let (token_tx, mut token_rx) = TokenSink::standalone(); handle diff --git a/pegainfer-qwen35/src/tp_executor.rs b/pegainfer-qwen35/src/tp_executor.rs index 3a85fc863..c838c7d81 100644 --- a/pegainfer-qwen35/src/tp_executor.rs +++ b/pegainfer-qwen35/src/tp_executor.rs @@ -3,6 +3,7 @@ //! Phase 1 supports eager dense TP prefill and decode. Unified execution still //! fails closed until the scheduler path can drive ordered eager decode. +use std::collections::HashMap; use std::collections::HashSet; use std::panic::AssertUnwindSafe; use std::panic::catch_unwind; @@ -17,8 +18,10 @@ use std::thread::JoinHandle; use std::thread::{self}; use anyhow::Result; -use pegainfer_core::kv_pool::KvState; use pegainfer_core::sampler::SamplingParams; +use pegainfer_kv_cache::KvCacheManager; +use pegainfer_kv_cache::KvView; +use pegainfer_kv_cache::RequestKv; use crate::batch_decode_graph::MAX_BATCH; use crate::config::TensorParallelConfig; @@ -35,6 +38,9 @@ use crate::executor::RequestId; use crate::logprobs::snapshot_requested_logprobs; use crate::prefill::PREFILL_CHUNK_LEN; use crate::prefill_buffers::GdrChunkwiseScratch35; +use crate::prefix_cache::PrefixGuard; +use crate::prefix_cache::Qwen35PrefixCache; +use crate::prefix_cache::RecurrentStateStore; use crate::recurrent_state::LinearStatePointerTables; use crate::recurrent_state::RecurrentState; use crate::weights::ModelRuntimeConfig; @@ -52,17 +58,30 @@ enum TpWorkerCommand { }, RunPrefillChunks { chunks: Vec, + kv_views: Vec, sample_seed: u64, resp: mpsc::Sender, }, RunDecodeStep { requests: Vec, + kv_views: Vec, sample_seed: u64, resp: mpsc::Sender, }, RunUnifiedStep { resp: mpsc::Sender, }, + RestoreRequest { + request_id: RequestId, + snapshot_slot: Option, + boundary: usize, + resp: mpsc::Sender, + }, + SaveSnapshot { + request_id: RequestId, + snapshot_slot: usize, + resp: mpsc::Sender, + }, DropRequest { request_id: RequestId, resp: mpsc::Sender, @@ -73,6 +92,7 @@ enum TpWorkerCommand { #[derive(Debug)] enum TpWorkerReply { Ack, + Position(usize), Prefill(PrefillResult), Decode(DecodeResult), } @@ -110,6 +130,8 @@ impl TpRuntimePoison { /// TP executor. Rank 0 is the primary worker and returns scheduler-visible /// artifacts; every rank runs the same ordered state-mutating commands. pub struct Qwen35TpExecutor { + kv_cache: Qwen35PrefixCache, + request_kvs: HashMap, workers: Vec, poison: Arc, world_size: usize, @@ -216,6 +238,24 @@ impl Qwen35TpExecutor { device_ordinals: &[usize], max_batch: usize, max_prefill_tokens: usize, + ) -> Result { + Self::from_runtime_with_limits_and_prefix( + model_path, + enable_cuda_graph, + device_ordinals, + max_batch, + max_prefill_tokens, + 0, + ) + } + + pub(crate) fn from_runtime_with_limits_and_prefix( + model_path: &str, + enable_cuda_graph: bool, + device_ordinals: &[usize], + max_batch: usize, + max_prefill_tokens: usize, + prefix_snapshot_bytes: usize, ) -> Result { anyhow::ensure!( device_ordinals.len() > 1, @@ -240,23 +280,40 @@ impl Qwen35TpExecutor { enable_cuda_graph: false, tensor_parallel: Some(TensorParallelConfig { rank, world_size }), device_ordinal, + prefix_snapshot_bytes, }, )?); } let first = models .first() .ok_or_else(|| anyhow::anyhow!("Qwen3.5 TP executor loaded no models"))?; - let page_size = first.kv_pool().layout().page_size; + let first_layout = *first.kv_buffer().layout(); + let page_size = first_layout.page_size; + let snapshot_slots = first.prefix_snapshot_slots(); let mut min_capacity_pages = usize::MAX; for (rank, model) in models.iter().enumerate() { - let rank_page_size = model.kv_pool().layout().page_size; + let rank_layout = model.kv_buffer().layout(); anyhow::ensure!( - rank_page_size == page_size, - "Qwen3.5 TP rank {rank} KV page size {rank_page_size} does not match rank 0 page size {page_size}" + rank_layout.page_size == first_layout.page_size + && rank_layout.num_layers == first_layout.num_layers + && rank_layout.num_kv_heads == first_layout.num_kv_heads + && rank_layout.head_dim == first_layout.head_dim, + "Qwen3.5 TP rank {rank} KV layout {:?} does not match rank 0 layout {:?}", + rank_layout, + first_layout, ); - min_capacity_pages = min_capacity_pages.min(model.kv_pool().capacity_pages()); + anyhow::ensure!( + model.prefix_snapshot_slots() == snapshot_slots, + "Qwen3.5 TP rank {rank} snapshot slots {} do not match rank 0 slots {snapshot_slots}", + model.prefix_snapshot_slots(), + ); + min_capacity_pages = min_capacity_pages.min(model.kv_buffer().num_blocks()); } - let capacity_pages_for_requests = min_capacity_pages.saturating_sub(1); + let manager = + KvCacheManager::from_buffer(models[0].kv_buffer().clone(), min_capacity_pages)?; + let kv_cache = Qwen35PrefixCache::new(manager, snapshot_slots)?; + let capacity_pages_for_requests = kv_cache.pool().max_request_blocks(); + let padding_page_id = kv_cache.pool().padding_block_id(); let max_position_embeddings = first.config().max_position_embeddings; let eos_token_id = first.config().eos_token_id; @@ -275,6 +332,8 @@ impl Qwen35TpExecutor { model, max_batch, max_prefill_tokens, + min_capacity_pages, + padding_page_id, nccl_id, Arc::clone(&startup_gate), Arc::clone(&effective_max_batch), @@ -343,6 +402,8 @@ impl Qwen35TpExecutor { disarm_nccl_startup_watchdog(watchdog_done, watchdog)?; Ok(Self { + kv_cache, + request_kvs: HashMap::new(), workers, poison, world_size, @@ -370,6 +431,31 @@ impl Qwen35TpExecutor { self.capacity_pages_for_requests } + pub(crate) fn available_pages(&self) -> usize { + self.kv_cache.pool().available_blocks() + } + + pub(crate) fn prefix_cache_enabled(&self) -> bool { + self.kv_cache.enabled() + } + + pub(crate) fn log_prefix_cache_stats(&self) { + let stats = self.kv_cache.stats(); + log::info!( + "Qwen3.5 TP prefix cache summary: ranks={}, joint_hits={}, hit_tokens={}, kv_only_fallbacks={}, snapshot_misses={}, inserts={}, evictions={}, restore_ms={:.3}, occupancy={}/{}", + self.world_size, + stats.joint_hits, + stats.joint_hit_tokens, + stats.kv_only_fallbacks, + stats.snapshot_misses, + stats.inserts, + stats.evictions, + stats.restore_ns as f64 / 1_000_000.0, + self.kv_cache.snapshot_occupancy(), + self.kv_cache.snapshot_slots(), + ); + } + pub(crate) fn max_position_embeddings(&self) -> usize { self.max_position_embeddings } @@ -383,26 +469,86 @@ impl Qwen35TpExecutor { self.broadcast_ack(TpWorkerCommandKind::Ping) } - pub fn execute_prefill(&self, plan: PrefillPlan<'_>) -> Result { + /// Create controller-owned RequestKv and restore the same joint boundary on every rank. + pub(crate) fn begin_request( + &mut self, + request_id: RequestId, + prompt_tokens: &[u32], + max_output_tokens: usize, + lora_name: Option<&str>, + allow_match: bool, + ) -> Result { + anyhow::ensure!( + !self.request_kvs.contains_key(&request_id), + "Qwen3.5 TP request {} already exists", + request_id.get() + ); + let (mut kv, restore) = self.kv_cache.begin_request( + prompt_tokens, + max_output_tokens, + lora_name, + allow_match, + )?; + let boundary = restore.as_ref().map_or(0, PrefixGuard::boundary); + let snapshot_slot = restore.as_ref().map(PrefixGuard::recurrent_slot); + let positions = match self.broadcast_restore_request(request_id, snapshot_slot, boundary) { + Ok(positions) => positions, + Err(error) => { + let _ = self.kv_cache.release_request(&mut kv); + return Err(error); + } + }; + let cached_tokens = if let Some(restore) = restore { + match self.kv_cache.finish_restore(&kv, restore, &positions) { + Ok(tokens) => tokens, + Err(error) => { + let _ = self.drop_request(request_id); + let _ = self.kv_cache.release_request(&mut kv); + return Err(error); + } + } + } else { + if !positions.iter().all(|&position| position == 0) { + let _ = self.drop_request(request_id); + let _ = self.kv_cache.release_request(&mut kv); + anyhow::bail!("Qwen3.5 TP cold request restored non-zero positions {positions:?}"); + } + 0 + }; + self.request_kvs.insert(request_id, kv); + Ok(cached_tokens) + } + + pub fn execute_prefill(&mut self, plan: PrefillPlan<'_>) -> Result { anyhow::ensure!( !plan.requests.is_empty(), "Qwen3.5 TP prefill plan requires at least one request" ); - let chunks: Vec = plan - .requests - .iter() - .cloned() - .map(TpPrefillChunkItem::from) - .collect(); + for request in plan.requests { + if !self.request_kvs.contains_key(&request.request_id) { + let max_output = self + .max_position_embeddings + .saturating_sub(request.prompt_tokens.len()); + self.begin_request( + request.request_id, + &request.prompt_tokens, + max_output, + None, + false, + )?; + } + } + let chunks: Vec = + plan.requests.iter().cloned().map(Into::into).collect(); self.execute_prefill_chunks(&chunks) } - fn execute_prefill_chunks(&self, chunks: &[TpPrefillChunkItem]) -> Result { + fn execute_prefill_chunks(&mut self, chunks: &[TpPrefillChunkItem]) -> Result { self.execute_prefill_chunks_with_seed(chunks, 0) } pub(crate) fn execute_prefill_chunks_with_seed( - &self, + &mut self, chunks: &[TpPrefillChunkItem], sample_seed: u64, ) -> Result { @@ -411,6 +557,38 @@ impl Qwen35TpExecutor { !chunks.is_empty(), "Qwen3.5 TP prefill chunk command requires at least one chunk" ); + for scheduled in 0..chunks.len() { + let chunk = &chunks[scheduled]; + let Some(kv) = self.request_kvs.get_mut(&chunk.request_id) else { + anyhow::bail!( + "Qwen3.5 TP prefill request {} has no RequestKv", + chunk.request_id.get() + ); + }; + if let Err(error) = self + .kv_cache + .schedule_prefill(kv, chunk.prompt_tokens.len()) + { + revert_scheduled_requests( + &self.kv_cache, + &mut self.request_kvs, + chunks + .iter() + .take(scheduled) + .map(|previous| previous.request_id), + ); + return Err(error); + } + } + let kv_views = chunks + .iter() + .map(|chunk| { + self.kv_cache.prefill_view( + &self.request_kvs[&chunk.request_id], + chunk.prompt_tokens.len(), + ) + }) + .collect::>(); let chunks = chunks.to_vec(); let (resp_tx, resp_rx) = mpsc::channel(); for worker in &self.workers { @@ -418,16 +596,77 @@ impl Qwen35TpExecutor { worker, TpWorkerCommand::RunPrefillChunks { chunks: chunks.clone(), + kv_views: kv_views.clone(), sample_seed, resp: resp_tx.clone(), }, )?; } drop(resp_tx); - wait_for_prefill(resp_rx, self.workers.len(), &self.poison) + let result = match wait_for_prefill(resp_rx, self.workers.len(), &self.poison) { + Ok(result) => result, + Err(error) => { + revert_scheduled_requests( + &self.kv_cache, + &mut self.request_kvs, + chunks.iter().map(|chunk| chunk.request_id), + ); + return Err(error); + } + }; + for chunk in &chunks { + let first_token = if chunk.finish_prefill { + Some( + result + .requests + .iter() + .find(|result| result.request_id == chunk.request_id) + .ok_or_else(|| { + anyhow::anyhow!( + "Qwen3.5 TP final prefill request {} returned no token", + chunk.request_id.get() + ) + })? + .first_token, + ) + } else { + None + }; + let (boundary, reservation) = { + let kv = self + .request_kvs + .get_mut(&chunk.request_id) + .expect("TP RequestKv exists after forward"); + let boundary = self.kv_cache.apply_prefill(kv, first_token)?; + let reservation = self.kv_cache.reserve_prefix(kv, boundary)?; + (boundary, reservation) + }; + if let Some(reservation) = reservation { + match self.broadcast_save_snapshot(chunk.request_id, reservation.recurrent_slot()) { + Ok(positions) if positions.iter().all(|&position| position == boundary) => { + let kv = self + .request_kvs + .get(&chunk.request_id) + .expect("TP RequestKv exists after snapshot save"); + self.kv_cache.publish_prefix(kv, reservation); + } + Ok(positions) => { + self.kv_cache.abort_prefix(reservation); + anyhow::bail!( + "Qwen3.5 TP snapshot positions {positions:?} do not match boundary {boundary}" + ); + } + Err(error) => { + self.kv_cache.abort_prefix(reservation); + return Err(error); + } + } + } + } + Ok(result) } - pub fn execute_decode(&self, plan: DecodePlan<'_>) -> Result { + pub fn execute_decode(&mut self, plan: DecodePlan<'_>) -> Result { anyhow::ensure!( !plan.requests.is_empty(), "Qwen3.5 TP decode plan requires at least one request" @@ -448,7 +687,7 @@ impl Qwen35TpExecutor { } pub(crate) fn execute_decode_items( - &self, + &mut self, requests: &[TpDecodeStepItem], sample_seed: u64, ) -> Result { @@ -457,6 +696,36 @@ impl Qwen35TpExecutor { !requests.is_empty(), "Qwen3.5 TP decode plan requires at least one request" ); + for scheduled in 0..requests.len() { + let request = &requests[scheduled]; + let kv = self + .request_kvs + .get_mut(&request.request_id) + .ok_or_else(|| { + anyhow::anyhow!( + "Qwen3.5 TP decode request {} has no RequestKv", + request.request_id.get() + ) + })?; + if let Err(error) = self.kv_cache.schedule_decode(kv) { + revert_scheduled_requests( + &self.kv_cache, + &mut self.request_kvs, + requests + .iter() + .take(scheduled) + .map(|previous| previous.request_id), + ); + return Err(error); + } + } + let kv_views = requests + .iter() + .map(|request| { + self.kv_cache + .decode_view(&self.request_kvs[&request.request_id]) + }) + .collect::>(); let requests = requests.to_vec(); let (resp_tx, resp_rx) = mpsc::channel(); for worker in &self.workers { @@ -464,16 +733,35 @@ impl Qwen35TpExecutor { worker, TpWorkerCommand::RunDecodeStep { requests: requests.clone(), + kv_views: kv_views.clone(), sample_seed, resp: resp_tx.clone(), }, )?; } drop(resp_tx); - wait_for_decode(resp_rx, self.workers.len(), &self.poison) + let result = match wait_for_decode(resp_rx, self.workers.len(), &self.poison) { + Ok(result) => result, + Err(error) => { + revert_scheduled_requests( + &self.kv_cache, + &mut self.request_kvs, + requests.iter().map(|request| request.request_id), + ); + return Err(error); + } + }; + for result in &result.requests { + let kv = self + .request_kvs + .get_mut(&result.request_id) + .expect("TP RequestKv exists after decode"); + self.kv_cache.apply_decode(kv, result.token)?; + } + Ok(result) } - pub fn drop_request(&self, request_id: RequestId) -> Result<()> { + pub fn drop_request(&mut self, request_id: RequestId) -> Result<()> { self.poison.ensure_healthy()?; let (resp_tx, resp_rx) = mpsc::channel(); for worker in &self.workers { @@ -486,7 +774,53 @@ impl Qwen35TpExecutor { )?; } drop(resp_tx); - wait_for_acks(resp_rx, self.workers.len(), "drop request", &self.poison) + wait_for_acks(resp_rx, self.workers.len(), "drop request", &self.poison)?; + if let Some(mut kv) = self.request_kvs.remove(&request_id) { + self.kv_cache.release_request(&mut kv)?; + } + Ok(()) + } + + fn broadcast_restore_request( + &self, + request_id: RequestId, + snapshot_slot: Option, + boundary: usize, + ) -> Result> { + let (resp_tx, resp_rx) = mpsc::channel(); + for worker in &self.workers { + self.send_or_poison( + worker, + TpWorkerCommand::RestoreRequest { + request_id, + snapshot_slot, + boundary, + resp: resp_tx.clone(), + }, + )?; + } + drop(resp_tx); + wait_for_positions(resp_rx, self.workers.len(), "restore request", &self.poison) + } + + fn broadcast_save_snapshot( + &self, + request_id: RequestId, + snapshot_slot: usize, + ) -> Result> { + let (resp_tx, resp_rx) = mpsc::channel(); + for worker in &self.workers { + self.send_or_poison( + worker, + TpWorkerCommand::SaveSnapshot { + request_id, + snapshot_slot, + resp: resp_tx.clone(), + }, + )?; + } + drop(resp_tx); + wait_for_positions(resp_rx, self.workers.len(), "save snapshot", &self.poison) } fn broadcast_ack(&self, kind: TpWorkerCommandKind) -> Result<()> { @@ -498,11 +832,13 @@ impl Qwen35TpExecutor { }, TpWorkerCommandKind::RunPrefillChunks => TpWorkerCommand::RunPrefillChunks { chunks: Vec::new(), + kv_views: Vec::new(), sample_seed: 0, resp: resp_tx.clone(), }, TpWorkerCommandKind::RunDecodeStep => TpWorkerCommand::RunDecodeStep { requests: Vec::new(), + kv_views: Vec::new(), sample_seed: 0, resp: resp_tx.clone(), }, @@ -649,6 +985,8 @@ impl TpWorker { model: Qwen35Model, max_batch: usize, max_prefill_tokens: usize, + total_blocks: usize, + padding_block_id: i32, nccl_id: cudarc::nccl::safe::Id, startup_gate: Arc, effective_max_batch: Arc, @@ -673,6 +1011,8 @@ impl TpWorker { model, max_batch, max_prefill_tokens, + total_blocks, + padding_block_id, ); let prepared = match prepared { Ok((prepared, rank_max_batch)) => { @@ -747,6 +1087,7 @@ struct TpWorkerState { _world_size: usize, max_batch: usize, model: Qwen35Model, + snapshots: RecurrentStateStore, requests: Vec, decode_buffers: BatchDecodeBuffers35, sample_scratch: pegainfer_sample::SampleScratch, @@ -759,6 +1100,7 @@ struct TpWorkerPrepared { world_size: usize, max_batch: usize, model: Qwen35Model, + snapshots: RecurrentStateStore, decode_buffers: BatchDecodeBuffers35, sample_scratch: pegainfer_sample::SampleScratch, cublas_guard: CublasThreadGuard, @@ -767,7 +1109,6 @@ struct TpWorkerPrepared { struct TpRequestState { request_id: RequestId, phase: TpRequestPhase, - kv: KvState, recurrent: RecurrentState, linear_pointer_tables: LinearStatePointerTables, } @@ -785,8 +1126,15 @@ impl TpWorkerPrepared { model: Qwen35Model, requested_max_batch: usize, max_prefill_tokens: usize, + total_blocks: usize, + padding_block_id: i32, ) -> Result<(Self, usize)> { let cublas_guard = bind_worker_thread(&model)?; + let snapshots = RecurrentStateStore::new( + model.device_ctx(), + model.config(), + model.prefix_snapshot_slots(), + )?; let (free_bytes, total_bytes) = model .device_ctx() .ctx @@ -821,7 +1169,11 @@ impl TpWorkerPrepared { prefill_scratch_tokens, prefill_scratch_bytes as f64 / 1024.0 / 1024.0, ); - let decode_buffers = model.create_batch_decode_buffers_with_capacity(max_batch)?; + let decode_buffers = model.create_batch_decode_buffers_with_capacity( + max_batch, + total_blocks, + padding_block_id, + )?; let sample_scratch = pegainfer_sample::SampleScratch::new( model.device_ctx(), model.config().selection_vocab, @@ -833,6 +1185,7 @@ impl TpWorkerPrepared { world_size, max_batch, model, + snapshots, decode_buffers, sample_scratch, cublas_guard, @@ -852,6 +1205,7 @@ impl TpWorkerPrepared { world_size, max_batch, mut model, + snapshots, decode_buffers, sample_scratch, cublas_guard, @@ -873,6 +1227,7 @@ impl TpWorkerPrepared { _world_size: world_size, max_batch: effective_max_batch, model, + snapshots, requests: Vec::new(), decode_buffers, sample_scratch, @@ -914,18 +1269,20 @@ impl TpWorkerState { } TpWorkerCommand::RunPrefillChunks { chunks, + kv_views, sample_seed, resp, } => { - let result = self.execute_prefill_chunks(&chunks, sample_seed); + let result = self.execute_prefill_chunks(&chunks, &kv_views, sample_seed); self.respond(resp, "prefill", result) } TpWorkerCommand::RunDecodeStep { requests, + kv_views, sample_seed, resp, } => { - let result = self.execute_decode(&requests, sample_seed); + let result = self.execute_decode(&requests, &kv_views, sample_seed); self.respond(resp, "decode", result) } TpWorkerCommand::RunUnifiedStep { resp } => { @@ -938,6 +1295,23 @@ impl TpWorkerState { )), ) } + TpWorkerCommand::RestoreRequest { + request_id, + snapshot_slot, + boundary, + resp, + } => { + let result = self.restore_request(request_id, snapshot_slot, boundary); + self.respond(resp, "restore request", result) + } + TpWorkerCommand::SaveSnapshot { + request_id, + snapshot_slot, + resp, + } => { + let result = self.save_snapshot(request_id, snapshot_slot); + self.respond(resp, "save snapshot", result) + } TpWorkerCommand::DropRequest { request_id, resp } => { self.drop_request(request_id); self.respond(resp, "drop request", Ok(TpWorkerReply::Ack)) @@ -982,6 +1356,7 @@ impl TpWorkerState { fn execute_prefill_chunks( &mut self, chunks: &[TpPrefillChunkItem], + kv_views: &[KvView], sample_seed: u64, ) -> Result { anyhow::ensure!( @@ -989,20 +1364,20 @@ impl TpWorkerState { "Qwen3.5 TP prefill chunk command requires at least one chunk" ); validate_prefill_chunks(chunks)?; - let new_requests = chunks - .iter() - .filter(|chunk| self.request_index(chunk.request_id).is_none()) - .count(); anyhow::ensure!( - self.requests.len() + new_requests <= self.max_batch, - "Qwen3.5 TP prefill chunks would exceed worker capacity {}", - self.max_batch + chunks.len() == kv_views.len(), + "Qwen3.5 TP prefill chunks / KV views len mismatch" ); let mut primary_results = Vec::new(); let mut final_row_idx = 0usize; - for chunk in chunks { - let state_idx = self.ensure_prefill_state(chunk.request_id)?; + for (chunk, kv_view) in chunks.iter().zip(kv_views) { + let state_idx = self.request_index(chunk.request_id).ok_or_else(|| { + anyhow::anyhow!( + "Qwen3.5 TP prefill request {} has no restored worker state", + chunk.request_id.get() + ) + })?; let state = &mut self.requests[state_idx]; anyhow::ensure!( state.phase == TpRequestPhase::Prefilling, @@ -1014,8 +1389,9 @@ impl TpWorkerState { let mut recurrent_refs = vec![&mut state.recurrent]; let logits = self.model.batch_prefill_logits( &prompt, - std::slice::from_mut(&mut state.kv), + std::slice::from_ref(kv_view), &mut recurrent_refs, + self.model.kv_buffer(), )?; if chunk.finish_prefill { @@ -1072,6 +1448,7 @@ impl TpWorkerState { fn execute_decode( &mut self, requests: &[TpDecodeStepItem], + kv_views: &[KvView], sample_seed: u64, ) -> Result { anyhow::ensure!( @@ -1079,6 +1456,10 @@ impl TpWorkerState { "Qwen3.5 TP decode command requires at least one request" ); validate_decode_requests(requests)?; + anyhow::ensure!( + requests.len() == kv_views.len(), + "Qwen3.5 TP decode requests / KV views len mismatch" + ); anyhow::ensure!( requests.len() <= self.max_batch, "Qwen3.5 TP decode batch {} exceeds worker capacity {}", @@ -1088,7 +1469,7 @@ impl TpWorkerState { let mut primary_results = Vec::with_capacity(if self.rank == 0 { requests.len() } else { 0 }); - for (row_idx, request) in requests.iter().enumerate() { + for (row_idx, (request, kv_view)) in requests.iter().zip(kv_views).enumerate() { let state_idx = self.request_index(request.request_id).ok_or_else(|| { anyhow::anyhow!( "Qwen3.5 TP decode request {} has no worker state", @@ -1103,11 +1484,11 @@ impl TpWorkerState { { let state = &mut self.requests[state_idx]; - let mut kv_refs = [&mut state.kv]; let mut recurrent_refs = [&mut state.recurrent]; self.model.batch_decode_eager_logits( &[request.token_id], - &mut kv_refs, + std::slice::from_ref(kv_view), + self.model.kv_buffer(), &mut recurrent_refs, &state.linear_pointer_tables, &mut self.decode_buffers, @@ -1150,11 +1531,32 @@ impl TpWorkerState { } } - fn ensure_prefill_state(&mut self, request_id: RequestId) -> Result { - if let Some(idx) = self.request_index(request_id) { - return Ok(idx); - } + fn restore_request( + &mut self, + request_id: RequestId, + snapshot_slot: Option, + boundary: usize, + ) -> Result { + anyhow::ensure!( + self.request_index(request_id).is_none(), + "Qwen3.5 TP request {} already has worker state", + request_id.get() + ); + anyhow::ensure!( + self.requests.len() < self.max_batch, + "Qwen3.5 TP restore would exceed worker capacity {}", + self.max_batch + ); let mut recurrent = RecurrentState::new(self.model.device_ctx(), self.model.config())?; + if let Some(slot) = snapshot_slot { + self.snapshots + .restore(self.model.device_ctx(), slot, &mut recurrent)?; + } + anyhow::ensure!( + recurrent.seq_len == boundary, + "Qwen3.5 TP restored recurrent position {} does not match boundary {boundary}", + recurrent.seq_len + ); let linear_pointer_tables = { let mut recurrent_refs = [&mut recurrent]; LinearStatePointerTables::from_recurrent_refs( @@ -1168,12 +1570,28 @@ impl TpWorkerState { let state = TpRequestState { request_id, phase: TpRequestPhase::Prefilling, - kv: self.model.alloc_kv(), recurrent, linear_pointer_tables, }; self.requests.push(state); - Ok(self.requests.len() - 1) + Ok(TpWorkerReply::Position(boundary)) + } + + fn save_snapshot( + &mut self, + request_id: RequestId, + snapshot_slot: usize, + ) -> Result { + let state_idx = self.request_index(request_id).ok_or_else(|| { + anyhow::anyhow!( + "Qwen3.5 TP snapshot request {} has no worker state", + request_id.get() + ) + })?; + let recurrent = &self.requests[state_idx].recurrent; + self.snapshots + .save(self.model.device_ctx(), snapshot_slot, recurrent)?; + Ok(TpWorkerReply::Position(recurrent.seq_len)) } fn request_index(&self, request_id: RequestId) -> Option { @@ -1251,6 +1669,9 @@ fn wait_for_acks( let response = recv_runtime_response(&responses, op_name, poison)?; match response.result? { TpWorkerReply::Ack => {} + TpWorkerReply::Position(_) => { + anyhow::bail!("Qwen3.5 TP {op_name} unexpectedly returned a position") + } TpWorkerReply::Prefill(_) => { anyhow::bail!("Qwen3.5 TP {op_name} unexpectedly returned prefill result") } @@ -1262,6 +1683,26 @@ fn wait_for_acks( Ok(()) } +#[allow(clippy::needless_pass_by_value)] +fn wait_for_positions( + responses: mpsc::Receiver, + expected: usize, + op_name: &'static str, + poison: &TpRuntimePoison, +) -> Result> { + let mut positions = Vec::with_capacity(expected); + for _ in 0..expected { + let response = recv_runtime_response(&responses, op_name, poison)?; + match response.result? { + TpWorkerReply::Position(position) => positions.push(position), + TpWorkerReply::Ack | TpWorkerReply::Prefill(_) | TpWorkerReply::Decode(_) => { + anyhow::bail!("Qwen3.5 TP {op_name} returned an unexpected reply") + } + } + } + Ok(positions) +} + #[allow(clippy::needless_pass_by_value)] fn wait_for_prefill( responses: mpsc::Receiver, @@ -1273,6 +1714,9 @@ fn wait_for_prefill( let response = recv_runtime_response(&responses, "prefill", poison)?; match response.result? { TpWorkerReply::Ack => {} + TpWorkerReply::Position(_) => { + anyhow::bail!("Qwen3.5 TP prefill unexpectedly returned a position") + } TpWorkerReply::Prefill(prefill) => { anyhow::ensure!( response.rank == 0, @@ -1304,6 +1748,9 @@ fn wait_for_decode( let response = recv_runtime_response(&responses, "decode", poison)?; match response.result? { TpWorkerReply::Ack => {} + TpWorkerReply::Position(_) => { + anyhow::bail!("Qwen3.5 TP decode unexpectedly returned a position") + } TpWorkerReply::Decode(decode) => { anyhow::ensure!( response.rank == 0, @@ -1379,6 +1826,25 @@ fn bind_worker_thread(model: &Qwen35Model) -> Result { Ok(CublasThreadGuard) } +/// Roll back controller-side requests scheduled by the current TP step. +fn revert_scheduled_requests( + kv_cache: &Qwen35PrefixCache, + request_kvs: &mut HashMap, + request_ids: impl IntoIterator, +) { + for request_id in request_ids { + let Some(kv) = request_kvs.get_mut(&request_id) else { + continue; + }; + if let Err(error) = kv_cache.revert_schedule(kv) { + log::warn!( + "failed to revert Qwen3.5 TP request {} KV schedule: {error}", + request_id.get() + ); + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -1545,8 +2011,9 @@ mod tests { fn starts_tp2_workers_and_broadcasts_lifecycle_commands() { let model_path = std::env::var("PEGAINFER_TEST_MODEL_PATH") .unwrap_or_else(|_| "/home/data/mgj/qwen35weights".to_string()); - let executor = Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 1) - .expect("start TP2 executor"); + let mut executor = + Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 1) + .expect("start TP2 executor"); assert_eq!(executor.world_size(), 2); assert_eq!(executor.max_batch(), 1); executor.ping_all().expect("ping all workers"); @@ -1575,8 +2042,9 @@ mod tests { fn tp2_prefill_runs_and_returns_primary_result() { let model_path = std::env::var("PEGAINFER_TEST_MODEL_PATH") .unwrap_or_else(|_| "/home/data/mgj/qwen35weights".to_string()); - let executor = Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 1) - .expect("start TP2 executor"); + let mut executor = + Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 1) + .expect("start TP2 executor"); let request_id = RequestId::new(11); let request = PrefillStepItem::new(request_id, vec![151_646, 9707], 0); let result = executor @@ -1596,9 +2064,13 @@ mod tests { fn tp2_chunked_prefill_advances_existing_request_state() { let model_path = std::env::var("PEGAINFER_TEST_MODEL_PATH") .unwrap_or_else(|_| "/home/data/mgj/qwen35weights".to_string()); - let executor = Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 1) - .expect("start TP2 executor"); + let mut executor = + Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 1) + .expect("start TP2 executor"); let request_id = RequestId::new(13); + executor + .begin_request(request_id, &[151_646, 9707], 3, None, false) + .expect("create TP2 RequestKv and recurrent state"); let first = TpPrefillChunkItem::new(request_id, vec![151_646], 0, false); let first_result = executor .execute_prefill_chunks(&[first]) @@ -1622,8 +2094,9 @@ mod tests { fn tp2_decode_runs_after_prefill() { let model_path = std::env::var("PEGAINFER_TEST_MODEL_PATH") .unwrap_or_else(|_| "/home/data/mgj/qwen35weights".to_string()); - let executor = Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 1) - .expect("start TP2 executor"); + let mut executor = + Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 1) + .expect("start TP2 executor"); let request_id = RequestId::new(17); let request = PrefillStepItem::new(request_id, vec![151_646, 9707], 0); let prefill = executor diff --git a/pegainfer-qwen35/src/unified_forward.rs b/pegainfer-qwen35/src/unified_forward.rs index a17420b5e..f679f32a4 100644 --- a/pegainfer-qwen35/src/unified_forward.rs +++ b/pegainfer-qwen35/src/unified_forward.rs @@ -10,8 +10,9 @@ //! compiled GQA groups; eager prefill fallback for uncompiled ones). use anyhow::Result; -use pegainfer_core::kv_pool::KvState; use pegainfer_core::tensor::HiddenStates; +use pegainfer_kv_cache::KvBuffer; +use pegainfer_kv_cache::KvView; use super::batch_decode_graph::BatchDecodeGraphState; use super::recurrent_state::RecurrentState; @@ -30,26 +31,25 @@ impl Qwen35Model { pub(crate) fn batch_prefill_logits( &self, prompts: &[&[u32]], - kv_states: &mut [KvState], + views: &[KvView], recurrent_states: &mut [&mut RecurrentState], + kv_buffer: &KvBuffer, ) -> Result { let n = prompts.len(); - anyhow::ensure!(n > 0, "batch_prefill requires at least one prompt"); - anyhow::ensure!(n == kv_states.len(), "prompts / kv_states len mismatch"); + anyhow::ensure!(n > 0, "batch prefill requires prompts"); + anyhow::ensure!(n == views.len(), "prompts / KV views len mismatch"); anyhow::ensure!( n == recurrent_states.len(), - "prompts / recurrent_states len mismatch" + "prompts / recurrent states len mismatch" ); - let mut last_hiddens = Vec::with_capacity(n); for i in 0..n { - let last_hidden = - self.prefill_last_hidden(prompts[i], &mut kv_states[i], recurrent_states[i])?; - debug_assert_eq!( - last_hidden.len, self.config.hidden_size, - "Qwen3.5 prefill last hidden row must match request {i}" - ); - last_hiddens.push(last_hidden); + last_hiddens.push(self.prefill_last_hidden( + prompts[i], + &views[i], + kv_buffer, + recurrent_states[i], + )?); } self.batch_last_hidden_logits(&last_hiddens) } @@ -67,10 +67,11 @@ impl Qwen35Model { pub(crate) fn unified_step( &self, prefill_prompts: &[&[u32]], - prefill_kv_states: &mut [KvState], + prefill_views: &[KvView], prefill_recurrent_states: &mut [&mut RecurrentState], decode_tokens: &[u32], - decode_kv_states: &mut [&mut KvState], + decode_views: &[KvView], + kv_buffer: &KvBuffer, graph_state: &mut BatchDecodeGraphState, ) -> Result { anyhow::ensure!( @@ -84,8 +85,9 @@ impl Qwen35Model { } else { Some(self.batch_prefill_logits( prefill_prompts, - prefill_kv_states, + prefill_views, prefill_recurrent_states, + kv_buffer, )?) }; @@ -93,7 +95,7 @@ impl Qwen35Model { let decoded = if decode_tokens.is_empty() { false } else { - self.batch_decode_graph(decode_tokens, decode_kv_states, graph_state)?; + self.batch_decode_graph(decode_tokens, decode_views, kv_buffer, graph_state)?; true }; @@ -108,10 +110,11 @@ impl Qwen35Model { mod tests { use std::path::Path; - use pegainfer_core::kv_pool::KvState; use pegainfer_core::tensor::HiddenStates; + use pegainfer_kv_cache::KvCacheManager; use super::*; + use crate::prefix_cache::Qwen35PrefixCache; const MODEL_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../models/Qwen3.5-4B"); @@ -141,106 +144,104 @@ mod tests { .unwrap() } - /// Verify that unified_step decode output matches batch_decode_graph standalone. - #[test] - fn unified_step_decode_matches_graph_decode() { - let Some(model_path) = get_model_path_or_skip() else { - return; - }; - let model = Qwen35Model::from_safetensors(&model_path, 0, 2).unwrap(); - + fn run_decode_path(model: &Qwen35Model, unified: bool) -> (Vec, Vec) { let prompt_a: Vec = vec![9707]; let prompt_b: Vec = vec![3838, 374, 220, 17, 10, 17]; + let prompts = [&prompt_a[..], &prompt_b[..]]; let num_steps = 5; - - // --- Reference: standalone batch_decode_graph --- - let ref_tokens = { - let prompts_ref: Vec<&[u32]> = vec![&prompt_a, &prompt_b]; - let mut kv_states: Vec = vec![model.alloc_kv(), model.alloc_kv()]; - let mut rec_states: Vec = vec![ - RecurrentState::new(&model.ctx, &model.config).unwrap(), - RecurrentState::new(&model.ctx, &model.config).unwrap(), - ]; - let mut rec_refs: Vec<&mut RecurrentState> = rec_states.iter_mut().collect(); - let first_logits = model - .batch_prefill_logits(&prompts_ref, &mut kv_states, &mut rec_refs) - .unwrap(); - let first = greedy_sample_batch(&model, &first_logits, 2); - let first_a = first[0]; - let first_b = first[1]; - - let mut gs = model.create_batch_decode_graph_state().unwrap(); - gs.copy_state_to_slot(&model.ctx, &rec_states[0], 0) + let manager = + KvCacheManager::from_buffer(model.kv_buffer().clone(), model.kv_buffer().num_blocks()) .unwrap(); - gs.copy_state_to_slot(&model.ctx, &rec_states[1], 1) - .unwrap(); - model.ctx.sync().unwrap(); - - let mut tokens_a = vec![first_a]; - let mut tokens_b = vec![first_b]; - let mut kv_refs: Vec<&mut KvState> = kv_states.iter_mut().collect(); - - for _ in 1..num_steps { - let tids = [*tokens_a.last().unwrap(), *tokens_b.last().unwrap()]; - model - .batch_decode_graph(&tids, &mut kv_refs, &mut gs) - .unwrap(); - let next = greedy_sample_batch(&model, &gs.buffers.logits, 2); - tokens_a.push(next[0]); - tokens_b.push(next[1]); - } - (tokens_a, tokens_b) - }; - - // --- unified_step path --- - let unified_tokens = { - let prompts_ref: Vec<&[u32]> = vec![&prompt_a, &prompt_b]; - let mut kv_states: Vec = vec![model.alloc_kv(), model.alloc_kv()]; - let mut rec_states: Vec = vec![ - RecurrentState::new(&model.ctx, &model.config).unwrap(), - RecurrentState::new(&model.ctx, &model.config).unwrap(), - ]; - let mut rec_refs: Vec<&mut RecurrentState> = rec_states.iter_mut().collect(); - - let output = model + let cache = Qwen35PrefixCache::new(manager, 0).unwrap(); + let mut kvs = vec![ + cache.pool().new_request(prompt_a.clone(), num_steps, None), + cache.pool().new_request(prompt_b.clone(), num_steps, None), + ]; + for (kv, prompt) in kvs.iter_mut().zip(prompts) { + cache.schedule_prefill(kv, prompt.len()).unwrap(); + } + let views = kvs + .iter() + .zip(prompts) + .map(|(kv, prompt)| cache.prefill_view(kv, prompt.len())) + .collect::>(); + let mut rec_states = [ + RecurrentState::new(&model.ctx, &model.config).unwrap(), + RecurrentState::new(&model.ctx, &model.config).unwrap(), + ]; + let mut rec_refs: Vec<&mut RecurrentState> = rec_states.iter_mut().collect(); + let mut gs = model + .create_batch_decode_graph_state( + cache.pool().total_blocks(), + cache.pool().padding_block_id(), + ) + .unwrap(); + let first_logits = if unified { + model .unified_step( - &prompts_ref, - &mut kv_states, + &prompts, + &views, &mut rec_refs, &[], - &mut [], - &mut model.create_batch_decode_graph_state().unwrap(), + &[], + cache.buffer(), + &mut gs, ) - .unwrap(); - let prefill_logits = output.prefill_logits.as_ref().unwrap(); - let first = greedy_sample_batch(&model, prefill_logits, 2); - let first_a = first[0]; - let first_b = first[1]; - - // Transfer prefill states to decode graph slots - let mut gs = model.create_batch_decode_graph_state().unwrap(); - gs.copy_state_to_slot(&model.ctx, &rec_states[0], 0) - .unwrap(); - gs.copy_state_to_slot(&model.ctx, &rec_states[1], 1) - .unwrap(); - - let mut kv_refs: Vec<&mut KvState> = kv_states.iter_mut().collect(); - - let mut tokens_a = vec![first_a]; - let mut tokens_b = vec![first_b]; - - for _ in 1..num_steps { - let tids = [*tokens_a.last().unwrap(), *tokens_b.last().unwrap()]; - let output = model - .unified_step(&[], &mut [], &mut [], &tids, &mut kv_refs, &mut gs) + .unwrap() + .prefill_logits + .unwrap() + } else { + model + .batch_prefill_logits(&prompts, &views, &mut rec_refs, cache.buffer()) + .unwrap() + }; + let first = greedy_sample_batch(model, &first_logits, 2); + for (kv, token) in kvs.iter_mut().zip(&first) { + cache.apply_prefill(kv, Some(*token)).unwrap(); + } + gs.copy_state_to_slot(&model.ctx, &rec_states[0], 0) + .unwrap(); + gs.copy_state_to_slot(&model.ctx, &rec_states[1], 1) + .unwrap(); + let mut tokens_a = vec![first[0]]; + let mut tokens_b = vec![first[1]]; + for _ in 1..num_steps { + for kv in &mut kvs { + cache.schedule_decode(kv).unwrap(); + } + let views = kvs + .iter() + .map(|kv| cache.decode_view(kv)) + .collect::>(); + let tids = [*tokens_a.last().unwrap(), *tokens_b.last().unwrap()]; + if unified { + model + .unified_step(&[], &[], &mut [], &tids, &views, cache.buffer(), &mut gs) .unwrap(); - assert!(output.decoded); - let next = greedy_sample_batch(&model, &gs.buffers.logits, 2); - tokens_a.push(next[0]); - tokens_b.push(next[1]); + } else { + model + .batch_decode_graph(&tids, &views, cache.buffer(), &mut gs) + .unwrap(); + } + let next = greedy_sample_batch(model, &gs.buffers.logits, 2); + for (kv, token) in kvs.iter_mut().zip(&next) { + cache.apply_decode(kv, *token).unwrap(); } - (tokens_a, tokens_b) + tokens_a.push(next[0]); + tokens_b.push(next[1]); + } + (tokens_a, tokens_b) + } + + /// Verify that unified_step decode output matches batch_decode_graph standalone. + #[test] + fn unified_step_decode_matches_graph_decode() { + let Some(model_path) = get_model_path_or_skip() else { + return; }; + let model = Qwen35Model::from_safetensors(&model_path, 0, 2, 0).unwrap(); + let ref_tokens = run_decode_path(&model, false); + let unified_tokens = run_decode_path(&model, true); assert_eq!( unified_tokens, ref_tokens, diff --git a/pegainfer-qwen35/src/weights.rs b/pegainfer-qwen35/src/weights.rs index acdf532ac..283c6dd5e 100644 --- a/pegainfer-qwen35/src/weights.rs +++ b/pegainfer-qwen35/src/weights.rs @@ -21,6 +21,7 @@ use pegainfer_core::weight_loader::load_tensor_2d_col_shard; use pegainfer_core::weight_loader::load_tensor_2d_row_shard; use pegainfer_core::weight_loader::mmap_shards; use pegainfer_core::weight_loader::precompute_rope; +use pegainfer_kv_cache::KvBuffer; use safetensors::SafeTensors; use super::config::Config35; @@ -90,6 +91,8 @@ pub(crate) struct ModelRuntimeConfig { pub(crate) enable_cuda_graph: bool, pub(crate) tensor_parallel: Option, pub(crate) device_ordinal: usize, + /// Per-rank GPU budget reserved for complete recurrent/conv snapshots. + pub(crate) prefix_snapshot_bytes: usize, } impl Default for ModelRuntimeConfig { @@ -98,6 +101,7 @@ impl Default for ModelRuntimeConfig { enable_cuda_graph: true, tensor_parallel: None, device_ordinal: 0, + prefix_snapshot_bytes: 0, } } } @@ -114,8 +118,10 @@ pub struct Qwen35Model { // Partial RoPE cache: [max_seq_len * rotary_dim] pub(super) cos_cache: DeviceVec, pub(super) sin_cache: DeviceVec, - /// Shared paged KV pool for full-attention layers. - kv_pool: pegainfer_core::kv_pool::KvPool, + /// Rank-local physical full-attention KV storage. + kv_buffer: KvBuffer, + /// Complete recurrent snapshot slots reserved by the load-time budget. + prefix_snapshot_slots: usize, /// Decode-slot count the recurrent-state reserve was sized for. /// Physical decode capacity actually allocated (recurrent-state slots, /// decode buffers, CUDA-graph slots). Always a `BATCH_BUCKETS` value. @@ -163,11 +169,13 @@ impl Qwen35Model { model_path: &str, device_ordinal: usize, max_batch: usize, + prefix_snapshot_bytes: usize, ) -> Result { Self::from_safetensors_with_runtime_and_capacity( model_path, ModelRuntimeConfig { device_ordinal, + prefix_snapshot_bytes, ..Default::default() }, max_batch, @@ -485,7 +493,7 @@ impl Qwen35Model { // Paged KV pool for the 8 full-attention layers. let page_size = 16usize; let num_full_layers = config.num_full_attention_layers(); - let layout = pegainfer_core::kv_pool::KvLayout::new( + let layout = pegainfer_kv_cache::KvLayout::new( num_full_layers, config.local_num_key_value_heads(tensor_parallel), config.head_dim, @@ -501,31 +509,45 @@ impl Qwen35Model { super::prefill_buffers::GdrChunkwiseScratch35::estimate_bytes(&config, max_prefill_len); let recurrent_reserve = STATES_PER_DECODE_SLOT * max_batch * super::recurrent_state::bytes_per_request(&config); + let prefix_snapshot_bytes = runtime.prefix_snapshot_bytes; + let snapshot_bytes_per_slot = super::recurrent_state::bytes_per_request(&config); + let snapshot_slots = prefix_snapshot_bytes / snapshot_bytes_per_slot; + anyhow::ensure!( + prefix_snapshot_bytes == 0 || snapshot_slots > 0, + "Qwen3.5 prefix-cache budget is {} MiB, but one recurrent/conv snapshot requires {:.3} MiB", + prefix_snapshot_bytes / (1024 * 1024), + snapshot_bytes_per_slot as f64 / 1024.0 / 1024.0, + ); + let snapshot_reserve = snapshot_slots * snapshot_bytes_per_slot; let min_kv_bytes = MIN_KV_PAGES * bytes_per_page; anyhow::ensure!( - free_bytes >= scratch_reserve + recurrent_reserve + min_kv_bytes, + free_bytes >= scratch_reserve + recurrent_reserve + snapshot_reserve + min_kv_bytes, "insufficient device memory for Qwen3.5: {} MB free, but prefill scratch needs {} MB, \ recurrent state needs {} MB ({STATES_PER_DECODE_SLOT} x {max_batch} decode slots), \ - and the minimal KV pool needs {} MB; lower the decode batch capacity (--max-batch) \ + prefix snapshots need {} MB ({} slots), and the minimal KV pool needs {} MB; \ + lower the decode batch capacity (--max-batch) or the prefix-cache budget \ or use a smaller model", free_bytes / (1024 * 1024), scratch_reserve / (1024 * 1024), recurrent_reserve / (1024 * 1024), + snapshot_reserve / (1024 * 1024), + snapshot_slots, min_kv_bytes / (1024 * 1024), ); - let available = free_bytes - scratch_reserve - recurrent_reserve; + let available = free_bytes - scratch_reserve - recurrent_reserve - snapshot_reserve; let kv_budget = (available as f64 * 0.85) as usize; let num_pages = (kv_budget / bytes_per_page).max(MIN_KV_PAGES); let kv_mb = num_pages * bytes_per_page / (1024 * 1024); let scratch_mb = scratch_reserve / (1024 * 1024); let recurrent_mb = recurrent_reserve / (1024 * 1024); + let snapshot_mb = snapshot_reserve / (1024 * 1024); info!( - "Qwen3.5 KV cache: {num_pages} pages ({kv_mb} MB), prefill scratch reserve: {scratch_mb} MB, recurrent-state reserve: {recurrent_mb} MB ({STATES_PER_DECODE_SLOT} x {max_batch} slots), {:.0}% of {:.0} MB free", + "Qwen3.5 KV cache: {num_pages} pages ({kv_mb} MB), prefill scratch reserve: {scratch_mb} MB, recurrent-state reserve: {recurrent_mb} MB ({STATES_PER_DECODE_SLOT} x {max_batch} slots), prefix snapshots: {snapshot_slots} slots ({snapshot_mb} MB), {:.0}% of {:.0} MB free", kv_budget as f64 / free_bytes as f64 * 100.0, free_bytes as f64 / 1024.0 / 1024.0 ); - let kv_pool = pegainfer_core::kv_pool::KvPool::new( - &ctx, + let kv_buffer = KvBuffer::new( + &ctx.stream, num_full_layers, config.local_num_key_value_heads(tensor_parallel), config.head_dim, @@ -543,7 +565,8 @@ impl Qwen35Model { norm, cos_cache, sin_cache, - kv_pool, + kv_buffer, + prefix_snapshot_slots: snapshot_slots, reserved_decode_slots: max_batch, decode_admission_batch, tp_comm: None, @@ -572,12 +595,12 @@ impl Qwen35Model { &self.ctx } - pub(crate) fn alloc_kv(&self) -> pegainfer_core::kv_pool::KvState { - self.kv_pool.alloc() + pub(crate) fn kv_buffer(&self) -> &KvBuffer { + &self.kv_buffer } - pub(crate) fn kv_pool(&self) -> &pegainfer_core::kv_pool::KvPool { - &self.kv_pool + pub(crate) fn prefix_snapshot_slots(&self) -> usize { + self.prefix_snapshot_slots } pub(crate) fn attach_tp_comm(&mut self, comm: Comm) { @@ -708,13 +731,21 @@ impl Qwen35Model { /// Create the CUDA Graph batch decode state at the loaded capacity. pub(crate) fn create_batch_decode_graph_state( &self, + max_total_pages: usize, + padding_page_id: i32, ) -> anyhow::Result { - self.create_batch_decode_graph_state_with_capacity(self.reserved_decode_slots) + self.create_batch_decode_graph_state_with_capacity( + self.reserved_decode_slots, + max_total_pages, + padding_page_id, + ) } pub(crate) fn create_batch_decode_graph_state_with_capacity( &self, max_batch: usize, + max_total_pages: usize, + padding_page_id: i32, ) -> anyhow::Result { anyhow::ensure!( max_batch <= self.reserved_decode_slots, @@ -725,7 +756,8 @@ impl Qwen35Model { &self.ctx, &self.config, self.tensor_parallel, - &self.kv_pool, + max_total_pages, + padding_page_id, max_batch, ) } @@ -733,14 +765,16 @@ impl Qwen35Model { pub(crate) fn create_batch_decode_buffers_with_capacity( &self, max_batch: usize, + max_total_pages: usize, + padding_page_id: i32, ) -> anyhow::Result { super::decode_buffers::BatchDecodeBuffers35::new( &self.ctx, &self.config, self.tensor_parallel, max_batch, - self.kv_pool.capacity_pages(), - self.kv_pool.padding_page_id(), + max_total_pages, + padding_page_id, ) } diff --git a/pegainfer-qwen35/tests/hf_golden_gate.rs b/pegainfer-qwen35/tests/hf_golden_gate.rs index 8f869dfc3..a9ee4d539 100644 --- a/pegainfer-qwen35/tests/hf_golden_gate.rs +++ b/pegainfer-qwen35/tests/hf_golden_gate.rs @@ -525,7 +525,12 @@ fn run(g: &Golden, ex: &mut Qwen35Executor, seqs: &[usize], batched: bool) -> (S (stats, fingerprint) } -fn run_tp(g: &Golden, ex: &Qwen35TpExecutor, seqs: &[usize], batched: bool) -> (Stats, Vec) { +fn run_tp( + g: &Golden, + ex: &mut Qwen35TpExecutor, + seqs: &[usize], + batched: bool, +) -> (Stats, Vec) { let mut stats = Stats::default(); let mut fingerprint = Vec::new(); let mut fold = |stats: &mut Stats, seq, pos, pega: &[(u32, f32)]| { @@ -862,17 +867,17 @@ fn pega_logprobs_match_hf_golden_within_qwen35_tolerance_tp2() { report_fixture_shape(&golden); let all: Vec = (0..golden.num_seqs).collect(); - let ex = build_tp2_executor(&model_path); - let (stats, fp1) = run_tp(&golden, &ex, &all, false); + let mut ex = build_tp2_executor(&model_path); + let (stats, fp1) = run_tp(&golden, &mut ex, &all, false); report_and_assert("TP2 sequential eager", &stats); - let (_, fp2) = run_tp(&golden, &ex, &all, false); + let (_, fp2) = run_tp(&golden, &mut ex, &all, false); assert_eq!( fp1, fp2, "TP2 sequential Qwen3.5 replay must reproduce identical logprobs" ); let batched_n = all.len().min(MAX_EXECUTOR_BATCH); - let (batched, _) = run_tp(&golden, &ex, &all[..batched_n], true); + let (batched, _) = run_tp(&golden, &mut ex, &all[..batched_n], true); report_and_assert("TP2 batched eager", &batched); } @@ -891,10 +896,10 @@ fn pega_logprobs_match_hf_long_golden_within_qwen35_tolerance_tp2() { report_fixture_shape(&golden); let all: Vec = (0..golden.num_seqs).collect(); - let ex = build_tp2_executor(&model_path); - let (stats, fp1) = run_tp(&golden, &ex, &all, false); + let mut ex = build_tp2_executor(&model_path); + let (stats, fp1) = run_tp(&golden, &mut ex, &all, false); report_and_assert("TP2 long sequential eager", &stats); - let (_, fp2) = run_tp(&golden, &ex, &all, false); + let (_, fp2) = run_tp(&golden, &mut ex, &all, false); assert_eq!( fp1, fp2, "TP2 long sequential Qwen3.5 replay must reproduce identical logprobs" diff --git a/pegainfer-qwen35/tests/prefix_cache.rs b/pegainfer-qwen35/tests/prefix_cache.rs new file mode 100644 index 000000000..973416df1 --- /dev/null +++ b/pegainfer-qwen35/tests/prefix_cache.rs @@ -0,0 +1,318 @@ +//! Qwen3.5 joint full-attention KV and recurrent/conv prefix-cache gate. +//! +//! The first request publishes the 256-token boundary. The second identical +//! request must restore both state families at that boundary and report the +//! joint hit through `TokenEvent::Scheduled`. + +use std::path::Path; + +use pegainfer_core::engine::EngineHandle; +use pegainfer_core::engine::FinishReason; +use pegainfer_core::engine::GenerateRequest; +use pegainfer_core::engine::TokenEvent; +use pegainfer_core::engine::TokenLogprob; +use pegainfer_core::engine::TokenSink; +use pegainfer_core::sampler::SamplingParams; +use pegainfer_qwen35::Qwen35LaunchOptions; +use pegainfer_qwen35::Qwen35SchedulerPolicy; + +mod common; + +const MODEL_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../models/Qwen3.5-4B"); +const PREFIX_BOUNDARY: usize = 256; +const PROMPT_TOKENS: usize = 320; +const TRACE_TOKENS: usize = 8; +const TOP_LOGPROBS: usize = 16; +// Qwen3.5-4B uses 49.125 MiB per snapshot, so this is exactly two slots. +const PREFIX_CACHE_MIB: usize = 128; + +fn model_path_or_skip() -> Option { + match std::env::var("pegainfer_TEST_MODEL_PATH") { + Ok(path) => Some(path), + Err(_) if Path::new(MODEL_PATH).join("config.json").exists() => { + Some(MODEL_PATH.to_string()) + } + Err(_) => { + eprintln!( + "skipping qwen35 prefix_cache: {MODEL_PATH}/config.json is missing; set pegainfer_TEST_MODEL_PATH to run it" + ); + None + } + } +} + +fn start_engine(model_path: &str, tp_size: usize, prefix_cache_mib: usize) -> EngineHandle { + pegainfer_qwen35::launch_with_options_and_policy( + Path::new(model_path), + Qwen35LaunchOptions { + device_ordinal: 0, + tp_size, + cuda_graph: tp_size == 1, + max_batch: 2, + max_prefill_tokens: 1024, + prefix_cache_mib, + }, + Qwen35SchedulerPolicy::Off, + ) + .unwrap_or_else(|err| panic!("failed to start Qwen3.5 TP{tp_size} prefix-cache engine: {err}")) +} + +struct Generation { + cached_tokens: usize, + tokens: Vec, + logprobs: Vec>, +} + +fn generate( + handle: &EngineHandle, + prompt_tokens: Vec, + max_tokens: usize, + logprobs: usize, + echo: bool, +) -> Generation { + let (token_tx, mut rx) = TokenSink::standalone(); + handle + .submit(GenerateRequest { + trace_parent: None, + request_id: None, + queued_at_unix_s: None, + data_parallel_rank: None, + prompt_tokens, + params: SamplingParams { + ignore_eos: true, + ..SamplingParams::default() + }, + max_tokens, + lora_adapter: None, + kv_transfer_params: None, + token_tx, + logprobs, + echo, + }) + .expect("submit failed"); + + let mut cached_tokens = None; + let mut generated_tokens = Vec::with_capacity(max_tokens); + let mut generated_logprobs = Vec::with_capacity(max_tokens); + loop { + match rx.blocking_recv().map(|(_, event)| event) { + Some(TokenEvent::Scheduled { + cached_tokens: hit, .. + }) => { + cached_tokens = Some(hit); + } + Some(TokenEvent::Token { id, logprob }) => { + generated_tokens.push(id); + generated_logprobs.push(logprob); + } + Some(TokenEvent::PromptTokens { .. } | TokenEvent::KvTransfer { .. }) => {} + Some(TokenEvent::Finished { finish_reason, .. }) => { + assert_eq!(finish_reason, FinishReason::Length); + return Generation { + cached_tokens: cached_tokens.expect("request did not emit Scheduled"), + tokens: generated_tokens, + logprobs: generated_logprobs, + }; + } + Some(TokenEvent::Error { message, .. }) => panic!("generation failed: {message}"), + Some(TokenEvent::Rejected { message, .. }) => panic!("generation rejected: {message}"), + None => panic!("scheduler channel closed without Finished"), + } + } +} + +fn generate_one(handle: &EngineHandle, prompt_tokens: Vec) -> (usize, u32) { + let result = generate(handle, prompt_tokens, 1, 0, false); + ( + result.cached_tokens, + *result.tokens.first().expect("request emitted no token"), + ) +} + +fn prompt_tokens( + tokenizer: &vllm_text::tokenizer::DynTokenizer, + text: &str, + token_len: usize, +) -> Vec { + let prompt = text.repeat(80); + let mut tokens = tokenizer.encode(&prompt, false).expect("encode failed"); + assert!( + tokens.len() >= token_len, + "test fixture encoded to only {} tokens", + tokens.len() + ); + tokens.truncate(token_len); + tokens +} + +fn assert_trace_close(label: &str, cold: &Generation, warm: &Generation) { + assert_eq!( + cold.tokens, warm.tokens, + "{label}: generated token trace changed" + ); + assert_eq!(cold.logprobs.len(), warm.logprobs.len()); + let mut deltas = Vec::new(); + for (position, (cold_lp, warm_lp)) in cold.logprobs.iter().zip(&warm.logprobs).enumerate() { + let cold_lp = cold_lp + .as_ref() + .unwrap_or_else(|| panic!("{label}: cold position {position} has no logprob")); + let warm_lp = warm_lp + .as_ref() + .unwrap_or_else(|| panic!("{label}: warm position {position} has no logprob")); + let cold_top = cold_lp.top_logprobs[0].1; + let cold_map: std::collections::HashMap = + cold_lp.top_logprobs.iter().copied().collect(); + let warm_argmax = warm_lp.top_logprobs[0].0; + let warm_cold_lp = cold_map.get(&warm_argmax).unwrap_or_else(|| { + panic!("{label}: warm argmax {warm_argmax} missing from cold top-logprobs") + }); + assert!( + cold_top - warm_cold_lp <= 0.20, + "{label}: position {position} argmax regret {:.4} exceeds 0.20", + cold_top - warm_cold_lp + ); + for &(token, warm_value) in warm_lp.top_logprobs.iter().take(8) { + if let Some(cold_value) = cold_map.get(&token) { + deltas.push((warm_value - cold_value).abs()); + } + } + } + assert!(!deltas.is_empty(), "{label}: no top-logprob overlap"); + deltas.sort_by(f32::total_cmp); + let mean = deltas.iter().sum::() / deltas.len() as f32; + let p99 = deltas[((deltas.len() as f64 * 0.99) as usize).min(deltas.len() - 1)]; + eprintln!( + "{label}: {} logprob deltas, mean {mean:.4}, p99 {p99:.4}", + deltas.len() + ); + assert!(mean <= 0.06, "{label}: mean logprob delta {mean:.4} > 0.06"); + assert!(p99 <= 0.20, "{label}: p99 logprob delta {p99:.4} > 0.20"); +} + +#[test] +fn joint_restore_and_unpinned_lru_eviction_preserve_output() { + let Some(model_path) = model_path_or_skip() else { + return; + }; + let tokenizer = common::load_tokenizer(&model_path); + let prompt_a = prompt_tokens( + &tokenizer, + "Alpha prefix exercises full-attention KV plus every recurrent and convolution state. ", + PROMPT_TOKENS, + ); + let prompt_b = prompt_tokens( + &tokenizer, + "Beta prefix is deliberately distinct and occupies a second recurrent snapshot slot. ", + PROMPT_TOKENS, + ); + let prompt_c = prompt_tokens( + &tokenizer, + "Gamma prefix creates pressure and must evict the least-recently-used unpinned snapshot. ", + PROMPT_TOKENS, + ); + + let handle = start_engine(&model_path, 1, PREFIX_CACHE_MIB); + let (cold_cached, cold_token) = generate_one(&handle, prompt_a.clone()); + assert_eq!(cold_cached, 0, "first request must be cold"); + + let (warm_cached, warm_token) = generate_one(&handle, prompt_a.clone()); + assert_eq!( + warm_cached, PREFIX_BOUNDARY, + "the longest jointly committed boundary should be restored" + ); + assert_eq!( + warm_token, cold_token, + "joint restore must preserve greedy output" + ); + + let (beta_cold_cached, beta_token) = generate_one(&handle, prompt_b.clone()); + assert_eq!(beta_cold_cached, 0, "new beta prefix must be cold"); + + let (alpha_touched_cached, _) = generate_one(&handle, prompt_a); + assert_eq!( + alpha_touched_cached, PREFIX_BOUNDARY, + "alpha lookup must refresh its snapshot LRU position" + ); + + let (gamma_cold_cached, _) = generate_one(&handle, prompt_c); + assert_eq!( + gamma_cold_cached, 0, + "new gamma prefix must insert under snapshot pressure" + ); + + let (beta_after_eviction_cached, beta_after_eviction_token) = generate_one(&handle, prompt_b); + assert_eq!( + beta_after_eviction_cached, 0, + "beta KV may remain resident, but its evicted snapshot must force a joint miss" + ); + assert_eq!( + beta_after_eviction_token, beta_token, + "snapshot pressure may change hit rate but must not change output" + ); +} + +#[test] +fn boundary_selection_and_multitoken_restore_preserve_logits() { + let Some(model_path) = model_path_or_skip() else { + return; + }; + let tokenizer = common::load_tokenizer(&model_path); + let long_prompt = prompt_tokens( + &tokenizer, + "Boundary coverage checks exact alignment, prefix extension, and joint recurrent state restore. ", + 576, + ); + let handle = start_engine(&model_path, 1, 512); + + let cold = generate( + &handle, + long_prompt.clone(), + TRACE_TOKENS, + TOP_LOGPROBS, + false, + ); + assert_eq!(cold.cached_tokens, 0); + let warm = generate( + &handle, + long_prompt.clone(), + TRACE_TOKENS, + TOP_LOGPROBS, + false, + ); + assert_eq!(warm.cached_tokens, 512); + assert_trace_close("tp1 576-token restore", &cold, &warm); + + let exact_512 = generate(&handle, long_prompt[..512].to_vec(), 1, 0, false); + assert_eq!( + exact_512.cached_tokens, 256, + "an exactly aligned prompt must retain one token for final prefill" + ); + let exact_256 = generate(&handle, long_prompt[..256].to_vec(), 1, 0, false); + assert_eq!(exact_256.cached_tokens, 0); + let echo = generate(&handle, long_prompt.clone(), 1, 0, true); + assert_eq!(echo.cached_tokens, 0, "echo must bypass prefix reuse"); + + let extended = generate(&handle, long_prompt[..320].to_vec(), 1, 0, false); + assert_eq!(extended.cached_tokens, 256); +} + +#[test] +#[ignore = "requires two CUDA devices and Qwen3.5 weights"] +fn tp2_joint_restore_preserves_output() { + let Some(model_path) = model_path_or_skip() else { + return; + }; + let tokenizer = common::load_tokenizer(&model_path); + let prompt = prompt_tokens( + &tokenizer, + "Tensor parallel prefix reuse restores every rank's recurrent and convolution state. ", + PROMPT_TOKENS, + ); + let handle = start_engine(&model_path, 2, PREFIX_CACHE_MIB); + + let cold = generate(&handle, prompt.clone(), TRACE_TOKENS, TOP_LOGPROBS, false); + assert_eq!(cold.cached_tokens, 0, "first TP2 request must be cold"); + let warm = generate(&handle, prompt, TRACE_TOKENS, TOP_LOGPROBS, false); + assert_eq!(warm.cached_tokens, PREFIX_BOUNDARY); + assert_trace_close("tp2 joint restore", &cold, &warm); +} diff --git a/pegainfer-server/src/bin/bench_serving/cli.rs b/pegainfer-server/src/bin/bench_serving/cli.rs index 9f5efb470..39da76e44 100644 --- a/pegainfer-server/src/bin/bench_serving/cli.rs +++ b/pegainfer-server/src/bin/bench_serving/cli.rs @@ -158,7 +158,7 @@ pub(crate) struct Cli { #[arg(long, default_value_t = false)] pub(crate) cuda_profiler_capture: bool, - /// Tensor-parallel world size for Kimi-K2 + /// Tensor-parallel world size for Kimi-K2 and Qwen3.5. #[arg(long, default_value_t = 1)] pub(crate) tp_size: usize, @@ -182,6 +182,11 @@ pub(crate) struct Cli { #[arg(long, default_value_t = 4)] pub(crate) max_batch: usize, + /// Qwen3.5 fixed GPU budget for recurrent/conv prefix snapshots. + /// Zero keeps the cache disabled. + #[arg(long, default_value_t = 0)] + pub(crate) qwen35_prefix_cache_mib: usize, + /// Qwen3.5 scheduler policy for prefill/decode balancing. #[arg(long, value_enum, default_value_t = CliQwen35SchedulerPolicy::Off)] pub(crate) qwen35_scheduler_policy: CliQwen35SchedulerPolicy, diff --git a/pegainfer-server/src/bin/bench_serving/main.rs b/pegainfer-server/src/bin/bench_serving/main.rs index e8d4514b0..53130280a 100644 --- a/pegainfer-server/src/bin/bench_serving/main.rs +++ b/pegainfer-server/src/bin/bench_serving/main.rs @@ -114,6 +114,23 @@ fn validate_qwen35_max_batch(model_type: ModelType, max_batch: usize) -> Result< Ok(()) } +fn validate_qwen35_parallel(model_type: ModelType, tp_size: usize, cuda_graph: bool) -> Result<()> { + #[cfg(feature = "qwen35")] + if matches!(model_type, ModelType::Qwen35) { + anyhow::ensure!(tp_size > 0, "--tp-size must be positive"); + if tp_size > 1 { + anyhow::ensure!( + !cuda_graph, + "Qwen3.5 TP benchmark requires --cuda-graph=false" + ); + } + return Ok(()); + } + #[cfg(not(feature = "qwen35"))] + let _ = (model_type, tp_size, cuda_graph); + Ok(()) +} + fn dispatch( cli: &Cli, model_type: ModelType, @@ -161,6 +178,7 @@ fn main() -> Result<()> { debug!("Detected model type: {:?}", model_type); validate_qwen35_scheduler_policy(model_type, cli.qwen35_scheduler_policy)?; validate_qwen35_max_batch(model_type, cli.max_batch)?; + validate_qwen35_parallel(model_type, cli.tp_size, cli.cuda_graph)?; let load_start = Instant::now(); // Shared tail for every scheduler-backed model: load the tokenizer, stamp @@ -265,17 +283,16 @@ fn main() -> Result<()> { .max_prefill_tokens .filter(|&v| v > 0) .unwrap_or(pegainfer_qwen35::DEFAULT_MAX_PREFILL_TOKENS); - let handle = pegainfer_qwen35::start_engine_with_capacity_and_policy( + let handle = pegainfer_qwen35::launch_with_options_and_policy( Path::new(&cli.model_path), - EngineLoadOptions { - enable_cuda_graph: cli.cuda_graph, - device_ordinals: vec![0], - parallel_config: None, - ep_backend: EpBackend::Nccl, - seed: command_seed(&cli), + pegainfer_qwen35::Qwen35LaunchOptions { + device_ordinal: 0, + tp_size: cli.tp_size, + cuda_graph: cli.cuda_graph, + max_batch: cli.max_batch, + max_prefill_tokens, + prefix_cache_mib: cli.qwen35_prefix_cache_mib, }, - cli.max_batch, - max_prefill_tokens, cli.qwen35_scheduler_policy.resolve(), )?; finish(handle, cli.cuda_graph) diff --git a/pegainfer-server/src/config.rs b/pegainfer-server/src/config.rs index f62cd8406..713475ab9 100644 --- a/pegainfer-server/src/config.rs +++ b/pegainfer-server/src/config.rs @@ -162,6 +162,11 @@ pub(crate) struct Args { #[arg(long, default_value_t = false)] pub no_prefix_cache: bool, + /// Qwen3.5-only fixed GPU budget for recurrent/conv prefix snapshots. + /// Zero disables Qwen3.5 prefix matching. + #[arg(long, default_value_t = 0)] + pub qwen35_prefix_cache_mib: usize, + /// Speculative drafter model path: Qwen3 DFlash/DSpark decoding, or the /// GLM5.2 DSpark drafter (greedy AND sampled requests speculate; /// per-request accept stats logged). For Qwen3: single-GPU greedy only; @@ -404,8 +409,10 @@ fn consumed_args(model_type: ModelType) -> &'static [&'static str] { "device_ordinal", "tp_size", "cuda_graph", + "no_prefix_cache", "max_prefill_tokens", "max_batch", + "qwen35_prefix_cache_mib", "qwen35_scheduler_policy", ], } @@ -531,6 +538,12 @@ impl Args { "--qwen35-scheduler-policy=auto is single-GPU only; Qwen3.5 TP uses the fixed off policy" ); } + if self.qwen35_prefix_cache_mib > 0 && self.no_prefix_cache { + bail!( + "--qwen35-prefix-cache-mib and --no-prefix-cache are contradictory; \ + use a positive budget to enable Qwen3.5 prefix reuse or zero to disable it" + ); + } } #[cfg(feature = "glm52")] if matches!(model_type, ModelType::Glm52) { @@ -856,6 +869,48 @@ mod tests { ); } + #[cfg(feature = "qwen35")] + #[test] + fn qwen35_accepts_positive_prefix_cache_budget() { + let (args, provided) = + parse_with_provided(&["pegainfer", "--qwen35-prefix-cache-mib", "128"]); + args.validate(ModelType::Qwen35, &provided) + .expect("single-GPU Qwen3.5 should accept a prefix-cache budget"); + assert_eq!(args.qwen35_prefix_cache_mib, 128); + } + + #[cfg(feature = "qwen35")] + #[test] + fn qwen35_rejects_contradictory_prefix_cache_flags() { + let (args, provided) = parse_with_provided(&[ + "pegainfer", + "--qwen35-prefix-cache-mib", + "128", + "--no-prefix-cache", + ]); + let err = args + .validate(ModelType::Qwen35, &provided) + .expect_err("positive Qwen3.5 cache budget must reject --no-prefix-cache") + .to_string(); + assert!(err.contains("contradictory"), "unexpected error: {err}"); + } + + #[cfg(feature = "qwen35")] + #[test] + fn qwen35_accepts_prefix_cache_with_tp() { + let (args, provided) = parse_with_provided(&[ + "pegainfer", + "--tp-size", + "2", + "--cuda-graph=false", + "--qwen35-prefix-cache-mib", + "128", + ]); + args.validate(ModelType::Qwen35, &provided) + .expect("Qwen3.5 TP should accept a prefix-cache budget"); + assert_eq!(args.qwen35_prefix_cache_mib, 128); + } + #[test] fn parses_lora_modules_name_equals_path() { assert_eq!( diff --git a/pegainfer-server/src/main.rs b/pegainfer-server/src/main.rs index 568dd10d3..21667cec1 100644 --- a/pegainfer-server/src/main.rs +++ b/pegainfer-server/src/main.rs @@ -342,6 +342,7 @@ fn load_engine(args: &Args, model_type: ModelType) -> anyhow::Result