diff --git a/Cargo.lock b/Cargo.lock index bc37d4c36..9d9f96e3c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3872,6 +3872,7 @@ dependencies = [ "pegainfer-core", "pegainfer-frontend", "pegainfer-kernels", + "pegainfer-kv-cache", "pegainfer-sample", "rand 0.10.1", "reqwest 0.12.28", diff --git a/docs/models/qwen35/prefix-cache.md b/docs/models/qwen35/prefix-cache.md index 84572ee46..8795fddfd 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. KV stays in kvbm's reclaimable inactive pool, while `SnapshotCache` owns snapshot 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-09 ## 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 are rejected before cache lookup 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 `SnapshotCache` 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,135 +121,154 @@ 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, + snapshots: SnapshotCache, + stats: PrefixCacheStats, } -struct RecurrentSnapshotCache { - slots: Vec, - index: HashMap, - free: Vec, - lru: LruList, +struct SnapshotCache { + entries: HashMap>, + free_slots: Vec, + slot_count: usize, + clock: u64, } -struct SnapshotSlot { - state: RecurrentState, - key: Option, - pin_count: usize, +struct SnapshotEntry { + recurrent_slot: usize, + last_used: AtomicU64, } -#[derive(Clone, Hash, Eq, PartialEq)] -struct SnapshotKey { - sequence_hash: SequenceHash, +struct SnapshotGuard { boundary: usize, + entry: Arc, +} + +struct SnapshotReservation { + 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. +`SnapshotEntry` owns one recurrent-state slot; a `SnapshotGuard` pins it by holding another `Arc`. `BlockPool` owns KV, with `PrefixProbe` and the request-local attachment keeping matched blocks alive through lookup and attach. TP publishes only after every worker saves the same slot at the same boundary, and reports a hit only after every worker confirms the restored 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_snapshot`; if it returns a reservation, copy the complete `RecurrentState` into that slot. +8. Publish the `SnapshotEntry` 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 recurrent snapshots; `BlockPool` owns KV lifetime: -- KV candidates are held by strong immutable-block guards from probe until exact-boundary attachment or abandonment. -- 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. +- With caching enabled, released KV enters kvbm's inactive pool, where it remains reusable and counts toward `available_blocks()`. +- With caching disabled, released KV returns directly to the free pool. +- Snapshot slots are immutable while indexed and can be evicted only when no `SnapshotGuard` exists. +- `SnapshotGuard` protects a snapshot 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: +Snapshot eviction reuses its recurrent-state slot; KV reclamation remains independent. An aborted replacement returns the unpublished slot to the free list. Lookup can continue to the next shorter published boundary. -- 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. +LRU selects an unpinned snapshot victim. Correctness depends on pinning and joint validation, not on LRU ordering. -This cannot produce a partial hit: only `Qwen35PrefixCache` can declare a hit, and it checks both resources on every lookup. +## TP behavior -LRU is the first victim policy for unpinned snapshot slots. Correctness depends on pinning and joint validation, not on LRU itself. +TP uses one logical cache decision and rank-local physical storage: + +1. The controller owns the only `KvCacheManager`, `RequestKv` map, `SnapshotCache`, 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. + +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 are rejected before 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. +- A snapshot may outlive its KV blocks; lookup treats missing KV as a cache miss. +- Disabling the feature preserves cold-serving behavior. ## Implementation order -Implementation is a follow-on to this design and should land in this order: - -1. Move Qwen3.5 KV management to `BlockPool`/`RequestKv` while keeping the existing KV memory layout, admission behavior, direct-paged kernels, and accuracy gates. -2. Add the fixed GPU snapshot allocator, config-derived slot sizing, pin guards, and unpinned LRU eviction. -3. Make scheduler chunk planning boundary-aware and publish snapshots after committed whole-model windows. -4. Add exact-boundary KV lookup/attach and expose the single `Qwen35PrefixCache::restore_prefix` operation. -5. Add metrics for joint hit length, KV-only fallback, snapshot miss, skipped insertion, eviction, and restore latency. -6. Run correctness, pressure, and warm-TTFT gates before enabling the feature by default. +1. Move Qwen3.5 KV management to `KvCacheManager` and `RequestKv`. +2. Add fixed-budget recurrent snapshot storage and joint prefix entry management. +3. Publish snapshots at committed 256-token prefill boundaries. +4. Add joint KV/snapshot lookup, restore, pinning, and LRU eviction for TP1 and TP. +5. Validate correctness, serving behavior, and cold/warm performance before default enablement. ## Validation @@ -264,9 +281,50 @@ The implementation acceptance surface should include: - 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. +- retained cache metrics for joint hits, hit length, misses, insertions, evictions, and slot occupancy. + +The cold/warm measurements below guide future stride changes. They must not be replaced by the old RTX 4090 CPU-transfer estimates, which measured a deferred design and a different snapshot shape. + +## 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** + +- **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 +333,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 d6fc6ab79..48e410b23 100644 --- a/docs/models/qwen35/tp-implementation.md +++ b/docs/models/qwen35/tp-implementation.md @@ -40,6 +40,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 40c8fcb60..51df71587 100644 --- a/pegainfer-kv-cache/src/manager.rs +++ b/pegainfer-kv-cache/src/manager.rs @@ -65,6 +65,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 99e37dc14..b3f23a987 100644 --- a/pegainfer-kv-cache/src/pool.rs +++ b/pegainfer-kv-cache/src/pool.rs @@ -208,6 +208,7 @@ impl BlockPool { seq_hashes, gpu_hit, cacheable, + block_size: self.block_size, held: gpu_guard, } } @@ -251,6 +252,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>, } @@ -269,6 +272,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). @@ -349,10 +383,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. @@ -360,6 +407,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( @@ -736,6 +806,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()); + } + fn complete_non_retained_speculative_request( pool: &BlockPool, prompt: &[u32], diff --git a/pegainfer-qwen35/Cargo.toml b/pegainfer-qwen35/Cargo.toml index f308934bf..c279b4f4e 100644 --- a/pegainfer-qwen35/Cargo.toml +++ b/pegainfer-qwen35/Cargo.toml @@ -13,6 +13,7 @@ half = { workspace = true } log = { workspace = true } pegainfer-core = { workspace = true } pegainfer-frontend = { workspace = true } +pegainfer-kv-cache = { workspace = true } pegainfer-kernels = { workspace = true } pegainfer-sample = { workspace = true } rand = { workspace = true } @@ -54,6 +55,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 b41da41d4..e75a2d7de 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::tensor::HiddenStates; use pegainfer_frontend::sampler::SamplingParams; +use pegainfer_kv_cache::KvBuffer; +use pegainfer_kv_cache::KvView; use super::batch_decode_graph::BATCH_BUCKETS; use super::batch_decode_graph::BatchDecodeGraphState; @@ -264,7 +265,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, @@ -274,7 +276,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" @@ -286,12 +288,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); } @@ -304,8 +311,7 @@ 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)?; // When this GQA group has no compiled batch-decode kernel, run full // attention through the paged-prefill kernel with a per-step plan. @@ -317,7 +323,7 @@ impl Qwen35Model { } else { let start_positions: Vec = positions.iter().map(|&p| p as usize).collect(); Some(self.one_token_paged_plan( - &kv_refs, + views, &start_positions, self.geometry.local_num_attention_heads(), self.geometry.local_num_key_value_heads(), @@ -325,10 +331,15 @@ impl Qwen35Model { )?) }; - 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, prefill_attn_plan.as_ref(), @@ -354,12 +365,20 @@ 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, graph_use: DecodeGraphUse, ) -> Result<()> { let padded_bs = bucket_for(token_ids.len()); - self.batch_decode_graph_padded(token_ids, kv_states, graph_state, graph_use, padded_bs) + self.batch_decode_graph_padded( + token_ids, + views, + kv_buffer, + graph_state, + graph_use, + padded_bs, + ) } /// `batch_decode_graph` with the bucket chosen by the caller instead of @@ -371,14 +390,15 @@ impl Qwen35Model { pub(crate) fn batch_decode_graph_padded( &self, token_ids: &[u32], - kv_states: &mut [&mut KvState], + views: &[KvView], + kv_buffer: &KvBuffer, graph_state: &mut BatchDecodeGraphState, graph_use: DecodeGraphUse, padded_bs: usize, ) -> 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 {}", @@ -397,15 +417,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); } graph_state.linear_pointer_tables.validate_for( @@ -414,14 +431,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); } @@ -440,13 +460,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. @@ -456,7 +480,7 @@ impl Qwen35Model { let result = match graph_use { DecodeGraphUse::Serve => graphs[bucket_idx].run_or_capture(&self.ctx, || { self.batch_decode_kernels_graph( - kv_buffer, + kv_buffer.buffer(), &layout, padded_bs, None, @@ -467,7 +491,7 @@ impl Qwen35Model { }), DecodeGraphUse::CaptureOnly => graphs[bucket_idx].capture_only(&self.ctx, || { self.batch_decode_kernels_graph( - kv_buffer, + kv_buffer.buffer(), &layout, padded_bs, None, @@ -489,7 +513,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(); @@ -500,13 +525,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); @@ -516,45 +543,37 @@ 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 - ) - })?; + .memcpy_htod(&positions_i32, &mut bufs.positions_d)?; - let kv_refs: Vec<&KvState> = kv_states.iter().map(|s| &**s).collect(); let plan = self.one_token_paged_plan( - &kv_refs, + views, &start_positions, self.geometry.local_num_attention_heads(), self.geometry.local_num_key_value_heads(), "hybrid decode", )?; - 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, + )?; anyhow::ensure!( - layout.num_kv_heads == self.config.num_key_value_heads + layout.num_kv_heads == self.geometry.local_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={}", + "hybrid decode KV layout mismatch bs={bs}: layout kv_heads={}, head_dim={}; local kv_heads={}, config head_dim={}", layout.num_kv_heads, layout.head_dim, - self.config.num_key_value_heads, + self.geometry.local_num_key_value_heads(), self.config.head_dim ); self.batch_decode_batched_hybrid_kernels( - kv_buffer, + kv_buffer.buffer(), &layout, &plan, bs, @@ -570,15 +589,18 @@ impl Qwen35Model { /// the hd256 FFI takes no override. fn one_token_paged_plan( &self, - kv_refs: &[&KvState], + views: &[KvView], start_positions: &[usize], num_q_heads: usize, num_kv_heads: usize, label: &str, ) -> Result { - let bs = kv_refs.len(); - let page_indices: Vec> = kv_refs.iter().map(|kv| kv.page_indices_i32()).collect(); - let last_page_lens: Vec = kv_refs.iter().map(|kv| kv.last_page_len()).collect(); + let bs = views.len(); + 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]; ops::PrefillPagedPlan::from_raw_batch_with_cta_tile_q( &self.ctx, diff --git a/pegainfer-qwen35/src/batch_decode_graph.rs b/pegainfer-qwen35/src/batch_decode_graph.rs index e68348ea1..4269e1144 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, geometry: LocalGeometry, - 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,17 +113,9 @@ 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}")) } /// D2D move slot `from` into slot `to`, leaving `to` as the canonical diff --git a/pegainfer-qwen35/src/decode_buffers.rs b/pegainfer-qwen35/src/decode_buffers.rs index cea677445..3f5ea9d39 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::LocalGeometry; @@ -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 10803a293..f4312e0ff 100644 --- a/pegainfer-qwen35/src/executor.rs +++ b/pegainfer-qwen35/src/executor.rs @@ -7,14 +7,16 @@ use std::collections::HashSet; use anyhow::Result; -use pegainfer_core::kv_pool::KvState; use pegainfer_core::tensor::HiddenStates; use pegainfer_frontend::engine::TokenLogprob; use pegainfer_frontend::sampler::SamplingParams; +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,11 +169,34 @@ 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() @@ -176,9 +209,18 @@ impl Qwen35Executor { }) .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(); @@ -188,8 +230,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, top_k)| { pegainfer_sample::token_logprob_from_row(row, first_token, *top_k) }); @@ -231,14 +274,39 @@ 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( + 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, - &mut kv_refs, + &views, + self.kv_cache.buffer(), &mut self.graph_state, crate::batch_decode::DecodeGraphUse::Serve, - )?; + ) { + 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(); @@ -258,6 +326,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, top_k)| { pegainfer_sample::token_logprob_from_row(row, token, *top_k) }); @@ -281,9 +350,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!( @@ -323,7 +394,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 3da029e66..425dce5f5 100644 --- a/pegainfer-qwen35/src/lib.rs +++ b/pegainfer-qwen35/src/lib.rs @@ -17,6 +17,7 @@ pub mod model_line; mod ops; mod prefill; pub mod prefill_buffers; +mod prefix_cache; pub(crate) mod recurrent; pub(crate) mod recurrent_state; mod scheduler; @@ -113,6 +114,7 @@ pub fn start_engine( max_batch, max_prefill_tokens, Qwen35SchedulerPolicy::Off, + 0, ) } @@ -127,9 +129,29 @@ pub struct Qwen35LaunchOptions { cuda_graph: bool, max_batch: usize, max_prefill_tokens: usize, + prefix_cache_mib: usize, } impl Qwen35LaunchOptions { + /// Configure the runtime and optional per-rank snapshot budget. + pub fn new( + device_ordinal: usize, + tp_size: usize, + cuda_graph: bool, + max_batch: usize, + max_prefill_tokens: usize, + prefix_cache_mib: usize, + ) -> Self { + Self { + device_ordinal, + tp_size, + cuda_graph, + max_batch, + max_prefill_tokens, + prefix_cache_mib, + } + } + fn device_ordinals(&self) -> Result> { anyhow::ensure!(self.tp_size >= 1, "Qwen3.5 tp_size must be >= 1"); Ok(if self.tp_size == 1 { @@ -161,6 +183,7 @@ pub fn launch_with_options_policy_and_overlap( options.max_prefill_tokens, scheduler_policy, decode_overlap, + options.prefix_cache_mib, ) } @@ -176,6 +199,7 @@ pub fn start_engine_with_capacity( max_batch, max_prefill_tokens, Qwen35SchedulerPolicy::Off, + 0, ) } @@ -185,6 +209,7 @@ pub(crate) fn start_engine_with_capacity_and_policy( max_batch: usize, max_prefill_tokens: usize, scheduler_policy: Qwen35SchedulerPolicy, + prefix_cache_mib: usize, ) -> Result { start_engine_with_capacity_policy_and_overlap( model_path, @@ -193,6 +218,7 @@ pub(crate) fn start_engine_with_capacity_and_policy( max_prefill_tokens, scheduler_policy, Qwen35DecodeOverlap::Off, + prefix_cache_mib, ) } @@ -203,6 +229,7 @@ pub fn start_engine_with_capacity_policy_and_overlap( max_prefill_tokens: usize, scheduler_policy: Qwen35SchedulerPolicy, decode_overlap: Qwen35DecodeOverlap, + prefix_cache_mib: usize, ) -> Result { anyhow::ensure!( (1..=MAX_DECODE_BATCH).contains(&max_batch), @@ -220,6 +247,9 @@ pub fn start_engine_with_capacity_policy_and_overlap( "Qwen3.5 --decode-overlap=stream currently requires --max-batch <= {MAX_SHARED_SM_DECODE_BATCH}; larger decode buckets still fall back to the prefill GEMM handle" ); } + 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 { anyhow::ensure!( decode_overlap == Qwen35DecodeOverlap::Off, @@ -240,6 +270,7 @@ pub fn start_engine_with_capacity_policy_and_overlap( max_batch, max_prefill_tokens, enable_cuda_graph, + prefix_snapshot_bytes, ); } @@ -260,7 +291,12 @@ pub fn start_engine_with_capacity_policy_and_overlap( 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, @@ -293,6 +329,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(); @@ -318,6 +355,7 @@ mod tests { 1, 1, Qwen35SchedulerPolicy::Auto, + 0, ) .err() .expect("scheduler policy validation should reject TP launch") @@ -342,6 +380,7 @@ mod tests { 1, Qwen35SchedulerPolicy::Off, Qwen35DecodeOverlap::SharedSm, + 0, ) .err() .expect("decode-overlap validation should reject TP launch") @@ -365,6 +404,7 @@ mod tests { 1, Qwen35SchedulerPolicy::Off, Qwen35DecodeOverlap::SharedSm, + 0, ) .err() .expect("decode-overlap validation should reject unsafe decode bucket") diff --git a/pegainfer-qwen35/src/model_line.rs b/pegainfer-qwen35/src/model_line.rs index 0788a0dfd..ecd4385fd 100644 --- a/pegainfer-qwen35/src/model_line.rs +++ b/pegainfer-qwen35/src/model_line.rs @@ -21,6 +21,9 @@ pub struct Qwen35Line; // Qwen3.5-exclusive CLI flags. #[derive(ClapArgs)] struct Qwen35Cli { + /// Per-rank GPU budget in MiB for joint prefix snapshots; zero disables reuse. + #[arg(long, default_value_t = 0)] + qwen35_prefix_cache_mib: usize, /// Decode-batch capacity, 1..=64. Qwen3.5 internally rounds allocation to /// the next graph bucket but admits only this many scheduler slots; defaults /// to 64. @@ -92,6 +95,7 @@ impl ModelLine for Qwen35Line { "device_ordinal", "tp_size", "cuda_graph", + "no_prefix_cache", "max_prefill_tokens", "decode_overlap", "decode_sm_pct", @@ -105,6 +109,11 @@ impl ModelLine for Qwen35Line { ) -> Result<(), CliError> { let cli = cli(ctx); let decode_overlap = resolve_decode_overlap(ctx.shared.decode_overlap)?; + if cli.qwen35_prefix_cache_mib > 0 && ctx.shared.no_prefix_cache { + return Err(CliError::rule( + "--qwen35-prefix-cache-mib and --no-prefix-cache are contradictory", + )); + } if let Some(max_batch) = cli.max_batch { if !(1..=crate::MAX_DECODE_BATCH).contains(&max_batch) { return Err(CliError::rule(format!( @@ -142,6 +151,7 @@ impl ModelLine for Qwen35Line { crate::launch_with_options_policy_and_overlap( ctx.model_path, Qwen35LaunchOptions { + prefix_cache_mib: cli.qwen35_prefix_cache_mib, device_ordinal: ctx.shared.device_ordinal, tp_size: ctx.shared.tp_size, cuda_graph: ctx.shared.cuda_graph, @@ -178,6 +188,33 @@ mod tests { MODEL_LINE.validate(&ctx, &provided) } + #[test] + fn accepts_prefix_cache_on_tp1_and_tp2() { + validate_argv(&["pegainfer", "--qwen35-prefix-cache-mib", "128"]).unwrap(); + validate_argv(&[ + "pegainfer", + "--tp-size", + "2", + "--cuda-graph=false", + "--qwen35-prefix-cache-mib", + "128", + ]) + .unwrap(); + } + + #[test] + fn rejects_contradictory_prefix_cache_flags() { + let error = validate_argv(&[ + "pegainfer", + "--qwen35-prefix-cache-mib", + "128", + "--no-prefix-cache", + ]) + .unwrap_err() + .to_string(); + assert!(error.contains("contradictory"), "{error}"); + } + #[test] fn probe_accepts_qwen35_identity() { let json = serde_json::json!({ diff --git a/pegainfer-qwen35/src/prefill.rs b/pegainfer-qwen35/src/prefill.rs index 6fa2555d7..00ef9fd37 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.geometry.local_num_attention_heads(), + self.geometry.local_num_key_value_heads(), + 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,65 +183,39 @@ 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, )?; - - // 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 = + let gdr_chunkwise_scratch = GdrChunkwiseScratch35::new(&self.ctx, c, self.geometry, seq_len)?; + Ok((hidden_batch, gdr_chunkwise_scratch)) + } - // 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 geom = self.geometry; - let prefill_plan = PrefillPagedPlan::new( - &self.ctx, - &kv_desc, - base_pos, - seq_len, - geom.local_num_attention_heads(), - geom.local_num_key_value_heads(), - 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; @@ -200,14 +228,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) @@ -223,8 +250,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; @@ -250,8 +276,7 @@ impl Qwen35Model { attn, &normed_batch, full_idx, - kv_state, - prefill_plan, + kv, attn_out_dim, seq_len, )?, @@ -289,8 +314,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 { @@ -305,16 +329,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; @@ -329,8 +351,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, @@ -360,24 +382,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, @@ -399,8 +415,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..860595ac1 --- /dev/null +++ b/pegainfer-qwen35/src/prefix_cache.rs @@ -0,0 +1,658 @@ +//! Joint full-attention KV and recurrent/conv prefix cache for Qwen3.5. + +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use anyhow::Result; +use pegainfer_core::tensor::DeviceContext; +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::config::LocalGeometry; +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, +} + +/// One recurrent snapshot indexed by a reusable prefix boundary. +struct SnapshotEntry { + /// Identically numbered physical recurrent snapshot on every rank. + recurrent_slot: usize, + /// Logical timestamp used to select an unpinned LRU victim. + last_used: AtomicU64, +} + +/// RAII pin on a recurrent snapshot during restore. +pub(crate) struct SnapshotGuard { + /// Token boundary represented by the pinned entry. + boundary: usize, + /// The directory's extra strong references are active restore pins. + entry: Arc, +} + +impl SnapshotGuard { + pub(crate) fn boundary(&self) -> usize { + self.boundary + } + + pub(crate) fn recurrent_slot(&self) -> usize { + self.entry.recurrent_slot + } +} + +/// Mutable prefix-cache state shared by every execution rank. +struct SnapshotCache { + /// Published boundary key to its recurrent snapshot 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 SnapshotCache { + 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) -> Option { + let last_used = self.tick(); + let entry = self.entries.get(&key)?; + entry.last_used.store(last_used, Ordering::Relaxed); + Some(SnapshotGuard { + boundary: key.boundary_tokens, + entry: Arc::clone(entry), + }) + } + + /// 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(&key) { + entry.last_used.store(last_used, Ordering::Relaxed); + 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)| Arc::strong_count(entry) == 1) + .min_by_key(|(_, entry)| entry.last_used.load(Ordering::Relaxed)) + .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(SnapshotReservation { + recurrent_slot: slot, + key, + replaced: evicted, + }) + } + + /// Publish after every rank has enqueued its recurrent-state copy. + fn publish(&mut self, reservation: SnapshotReservation) { + let last_used = self.tick(); + let previous = self.entries.insert( + reservation.key, + Arc::new(SnapshotEntry { + recurrent_slot: reservation.recurrent_slot, + last_used: AtomicU64::new(last_used), + }), + ); + debug_assert!(previous.is_none()); + } + + /// Return an unpublished slot after a physical copy failure. + fn abort(&mut self, reservation: SnapshotReservation) { + 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 SnapshotReservation { + recurrent_slot: usize, + key: PrefixBoundaryKey, + /// Whether this insertion replaced an unpinned snapshot entry. + replaced: bool, +} + +impl SnapshotReservation { + 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, + geometry: LocalGeometry, + slot_count: usize, + ) -> Result { + let mut slots = Vec::with_capacity(slot_count); + for _ in 0..slot_count { + slots.push(RecurrentState::new(ctx, config, geometry)?); + } + 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, and LRU state for reusable entries. + snapshots: SnapshotCache, + /// 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, + snapshots: SnapshotCache::new(snapshot_slots), + 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.snapshots.capacity() > 0 + } + + /// Total number of preallocated recurrent snapshot slots. + pub(crate) fn snapshot_slots(&self) -> usize { + self.snapshots.capacity() + } + + /// Number of snapshot slots that currently have a published key. + pub(crate) fn snapshot_occupancy(&self) -> usize { + self.snapshots.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 restore + /// is enqueued (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; + 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.snapshots.lookup(key) 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 enqueued the recurrent-state copy. + /// + /// `begin_request` performs lookup and KV attachment. The physical + /// recurrent-state copy remains executor-specific: single-GPU execution + /// copies from the local store, while TP coordinates all workers. This + /// validates the positions and releases the guard once every copy is + /// enqueued in stream order. + pub(crate) fn finish_restore( + &mut self, + request: &RequestKv, + guard: SnapshotGuard, + 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; + // End the cache pin after restore enqueue and position validation. + 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. The returned boundary can be passed to + /// [`Self::reserve_snapshot`], which accepts only non-zero multiples of + /// [`SNAPSHOT_STRIDE_TOKENS`]. + pub(crate) fn apply_prefill( + &self, + request: &mut RequestKv, + first_token: Option, + ) -> Result { + let applied = if let Some(first_token) = first_token { + request.apply_prefill(first_token, self.kv.pool()) + } else { + request.apply_prefill_chunk(self.kv.pool()) + }; + // Retain registered KV only when prefix caching is enabled. + if !self.enabled() { + request.mark_blocks_reset_on_release(); + } + applied?; + let boundary = request.kv_position(); + Ok(boundary) + } + + /// Reserve a recurrent-state slot for an eligible applied boundary. + /// + /// Returns a reservation when the boundary is eligible and a free or + /// unpinned slot is available. The caller copies recurrent state into that + /// slot on every rank, then publishes or aborts the reservation. Duplicate + /// boundaries, disabled caching, ineligible boundaries, and fully pinned + /// capacity return `None`. + pub(crate) fn reserve_snapshot( + &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.snapshots.reserve(key)) + } + + /// Publish after rank-local snapshot copies succeed. `reserve_snapshot` + /// already validated registered KV at the boundary. + pub(crate) fn publish_snapshot(&mut self, reservation: SnapshotReservation) { + self.stats.inserts += 1; + if reservation.was_replacement() { + self.stats.evictions += 1; + } + self.snapshots.publish(reservation); + } + + /// Abort a prefix reservation after any rank-local copy failure. + pub(crate) fn abort_snapshot(&mut self, reservation: SnapshotReservation) { + self.snapshots.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<()> { + let applied = request.apply_decode(token, self.kv.pool()); + if !self.enabled() { + request.mark_blocks_reset_on_release(); + } + applied?; + 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. + pub(crate) fn release_request(&self, request: &mut RequestKv) -> Result<()> { + if !self.enabled() { + 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 pegainfer_kv_cache::BlockPool; + + use super::PrefixBoundaryKey; + use super::SNAPSHOT_STRIDE_TOKENS; + use super::SnapshotCache; + 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 = SnapshotCache::new(1); + let reservation = state + .reserve(key(1)) + .expect("empty state must reserve a write"); + assert!(state.lookup(key(1)).is_none()); + state.publish(reservation); + assert!(state.lookup(key(1)).is_some()); + } + + #[test] + fn aborted_write_does_not_expose_partial_snapshot() { + let mut state = SnapshotCache::new(1); + let reservation = state + .reserve(key(1)) + .expect("empty state must reserve a write"); + state.abort(reservation); + assert!(state.lookup(key(1)).is_none()); + assert!(state.reserve(key(2)).is_some()); + } + + #[test] + fn duplicate_refreshes_lru_without_allocating_a_slot() { + let mut state = SnapshotCache::new(1); + let reservation = state + .reserve(key(1)) + .expect("empty directory must reserve a write"); + state.publish(reservation); + 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 = SnapshotCache::new(2); + for tag in [1, 2] { + let reservation = state + .reserve(key(tag)) + .expect("state should have a free slot"); + state.publish(reservation); + } + drop(state.lookup(key(1)).expect("key 1 should be present")); + let reservation = state + .reserve(key(3)) + .expect("an unpinned LRU victim should be available"); + state.publish(reservation); + assert!(state.lookup(key(1)).is_some()); + assert!(state.lookup(key(2)).is_none()); + assert!(state.lookup(key(3)).is_some()); + } + + #[test] + fn all_pinned_snapshots_make_insertion_a_soft_miss() { + let mut state = SnapshotCache::new(1); + let reservation = state + .reserve(key(1)) + .expect("empty state must reserve a write"); + state.publish(reservation); + let guard = state.lookup(key(1)).expect("key 1 should be present"); + assert!(state.reserve(key(2)).is_none()); + drop(guard); + assert!(state.reserve(key(2)).is_some()); + } + + #[test] + fn released_prefix_kv_is_reclaimable_without_snapshot_eviction() { + let pool = BlockPool::new(16, 32).expect("block pool"); + let baseline = pool.available_blocks(); + let prompt = vec![7; SNAPSHOT_STRIDE_TOKENS + 16]; + let mut request = pool.new_request(prompt.clone(), 0, None); + request + .schedule_prefill(prompt.len(), &pool) + .expect("schedule prefill"); + request.apply_prefill_chunk(&pool).expect("apply prefill"); + let cached_boundary = request + .registered_boundary_hash(SNAPSHOT_STRIDE_TOKENS) + .expect("full snapshot boundary is registered"); + + let mut state = SnapshotCache::new(1); + let key = PrefixBoundaryKey { + sequence_hash: cached_boundary, + boundary_tokens: SNAPSHOT_STRIDE_TOKENS, + }; + let reservation = state + .reserve(key) + .expect("empty state must reserve a write"); + state.publish(reservation); + request.release().expect("release request"); + + // Released KV remains reusable and counts as available capacity. + assert_eq!(pool.available_blocks(), baseline); + + let mut warm = pool.new_request(prompt, 0, None); + assert_eq!( + warm.match_and_add_prefix(&pool).expect("match inactive KV"), + SNAPSHOT_STRIDE_TOKENS + ); + warm.release().expect("release warm request"); + + // A full-pool cold reservation reclaims KV without evicting the snapshot. + let cold_prompt = vec![9; baseline * pool.block_size()]; + let mut cold = pool.new_request(cold_prompt.clone(), 0, None); + cold.schedule_prefill(cold_prompt.len(), &pool) + .expect("inactive KV must satisfy cold allocation"); + cold.revert_schedule().expect("revert cold reservation"); + cold.release().expect("release cold request"); + + assert!(state.lookup(key).is_some()); + } +} diff --git a/pegainfer-qwen35/src/recurrent_state.rs b/pegainfer-qwen35/src/recurrent_state.rs index 1b97505d1..6219f18a2 100644 --- a/pegainfer-qwen35/src/recurrent_state.rs +++ b/pegainfer-qwen35/src/recurrent_state.rs @@ -84,6 +84,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/backend.rs b/pegainfer-qwen35/src/scheduler/backend.rs index 406c56307..13c412518 100644 --- a/pegainfer-qwen35/src/scheduler/backend.rs +++ b/pegainfer-qwen35/src/scheduler/backend.rs @@ -2,9 +2,12 @@ use super::*; use crate::logprobs::LogprobSnapshot; +use crate::tp_executor::TpBeginRequestError; pub(super) struct SingleGpuBackend { - model: Qwen35Model, + pub(super) model: Qwen35Model, + pub(super) kv_cache: Qwen35PrefixCache, + recurrent_store: RecurrentStateStore, graph_state: BatchDecodeGraphState, prefill_stream: Option>, } @@ -16,6 +19,12 @@ pub(super) enum SchedulerBackend { Tp(TpSchedulerBackend), } +impl Drop for SchedulerBackend { + fn drop(&mut self) { + self.log_prefix_cache_stats(); + } +} + pub(super) struct AsyncPrefillOutput { logits: Option, done: CudaEvent, @@ -68,37 +77,6 @@ pub(super) fn fatal_cuda_lifecycle(message: &str) -> ! { std::process::abort(); } -/// Borrowed single-GPU prefill inputs, in request order: token windows, KV -/// states, recurrent states. -type SinglePrefillViews<'a> = ( - Vec<&'a [u32]>, - &'a mut Vec, - Vec<&'a mut RecurrentState>, -); - -/// Borrow a single-GPU prefill chunk as the model's prefill inputs. -fn single_prefill_views(chunk: &mut ScheduledChunk) -> Result> { - let windows = 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"); - }; - Ok((windows, kvs, recs.iter_mut().collect())) -} - -/// Borrow the active decode batch as the model's decode inputs, in slot order: -/// last tokens and KV states. -fn single_decode_views(active: &mut [ActiveRequest35]) -> (Vec, Vec<&mut KvState>) { - let tokens = active.iter().map(|r| r.last_token).collect(); - let kvs = active - .iter_mut() - .map(|r| match &mut r.backend_state { - ActiveBackendState::Single { kv, .. } => kv, - ActiveBackendState::Tp { .. } => panic!("single-GPU decode received TP active state"), - }) - .collect(); - (tokens, kvs) -} - /// Pair each sampled token with its host logprob row, where one was requested. fn attached_logprobs( cpu_logits: Vec>, @@ -123,15 +101,156 @@ pub(super) struct TpSchedulerBackend { pub(super) pending_compaction: Option, } +pub(super) enum AdmissionError { + Recoverable(anyhow::Error), + Fatal(anyhow::Error), +} + +impl From for AdmissionError { + fn from(error: TpBeginRequestError) -> Self { + match error { + TpBeginRequestError::Recoverable(error) => Self::Recoverable(error), + TpBeginRequestError::Fatal(error) => Self::Fatal(error), + } + } +} + impl SingleGpuBackend { + fn schedule_prefill_views( + &self, + kvs: &mut [Box], + windows: &[Vec], + ) -> Result> { + debug_assert_eq!(kvs.len(), windows.len()); + for (scheduled, (kv, window)) in kvs.iter_mut().zip(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).map(Box::as_mut), + ); + return Err(error); + } + } + Ok(kvs + .iter() + .zip(windows) + .map(|(kv, window)| self.kv_cache.prefill_view(kv, window.len())) + .collect()) + } + + fn schedule_decode_views(&self, active: &mut [ActiveRequest35]) -> Result> { + for (scheduled, request) in active.iter_mut().enumerate() { + let ActiveBackendState::Single { kv, .. } = &mut request.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); + } + } + Ok(active + .iter() + .map(|request| match &request.backend_state { + ActiveBackendState::Single { kv, .. } => self.kv_cache.decode_view(kv), + ActiveBackendState::Tp { .. } => { + panic!("single-GPU decode received TP active state") + } + }) + .collect()) + } + + pub(super) 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_snapshot(kv, boundary)? { + if let Err(error) = self.recurrent_store.save( + self.model.device_ctx(), + reservation.recurrent_slot(), + rec, + ) { + self.kv_cache.abort_snapshot(reservation); + return Err(error); + } + self.kv_cache.publish_snapshot(reservation); + } + } + Ok(()) + } + + pub(super) 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(()) + } + + pub(super) 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={}, occupancy={}/{}", + stats.joint_hits, + stats.joint_hit_tokens, + stats.kv_only_fallbacks, + stats.snapshot_misses, + stats.inserts, + stats.evictions, + cache.snapshot_occupancy(), + cache.snapshot_slots(), + ); + } + pub(super) fn new( model: Qwen35Model, max_batch: usize, decode_overlap: Qwen35DecodeOverlap, ) -> 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.geometry, + 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)?; + let graph_state = model.create_batch_decode_graph_state_with_capacity( + graph_capacity, + kv_cache.pool().total_blocks(), + kv_cache.pool().padding_block_id(), + )?; let prefill_stream = match decode_overlap { Qwen35DecodeOverlap::Off => None, Qwen35DecodeOverlap::SharedSm => Some( @@ -144,6 +263,8 @@ impl SingleGpuBackend { }; Ok(Self { model, + kv_cache, + recurrent_store, graph_state, prefill_stream, }) @@ -163,23 +284,63 @@ impl SingleGpuBackend { } pub(super) fn page_size(&self) -> usize { - self.model.kv_pool().layout().page_size + self.kv_cache.pool().block_size() } pub(super) fn available_pages(&self) -> usize { - self.model.kv_pool().available_pages() + self.kv_cache.pool().available_blocks() } pub(super) fn capacity_pages_for_requests(&self) -> usize { - self.model.kv_pool().capacity_pages().saturating_sub(1) + self.kv_cache.pool().max_request_blocks() } pub(super) fn max_position_embeddings(&self) -> usize { self.model.config().max_position_embeddings } - pub(super) fn alloc_kv(&self) -> KvState { - self.model.alloc_kv() + pub(super) fn alloc_prefill_state( + &mut self, + req: &SchedulerRequest, + ) -> Result<(PrefillBackendState, usize), AdmissionError> { + let mut rec = self + .alloc_recurrent() + .map_err(AdmissionError::Recoverable)?; + let (mut kv, restore) = self + .kv_cache + .begin_request( + &req.prompt_tokens, + req.max_tokens, + req.lora_adapter.as_deref(), + true, + ) + .map_err(AdmissionError::Recoverable)?; + 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(AdmissionError::Recoverable(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(AdmissionError::Recoverable(error)); + } + } + } else { + 0 + }; + Ok(( + PrefillBackendState::Single { + kv: Box::new(kv), + rec, + }, + cached_tokens, + )) } pub(super) fn alloc_recurrent(&self) -> Result { @@ -191,8 +352,22 @@ impl SingleGpuBackend { } pub(super) fn batch_prefill_logits(&self, chunk: &mut ScheduledChunk) -> Result { - let (windows, kvs, mut recs) = single_prefill_views(chunk)?; - self.model.batch_prefill_logits(&windows, kvs, &mut recs) + 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"); + }; + let views = self.schedule_prefill_views(kvs, &chunk.windows)?; + let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); + 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.iter_mut().map(Box::as_mut)); + } + result } pub(super) fn overlap_enabled(&self) -> bool { @@ -214,12 +389,18 @@ impl SingleGpuBackend { .join(&self.model.device_ctx().stream) .map_err(|err| anyhow::anyhow!("join Qwen3.5 prefill stream: {err}"))?; - let (windows, kvs, mut recs) = single_prefill_views(chunk)?; + 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 async prefill received TP chunk state"); + }; + let views = self.schedule_prefill_views(kvs, &chunk.windows)?; + let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); let logits = match self.model.batch_prefill_logits_on_stream( Arc::clone(&prefill_stream), - &windows, - kvs, - &mut recs, + &window_refs, + &views, + self.kv_cache.buffer(), + &mut rec_refs, ) { Ok(logits) => logits, Err(err) => { @@ -228,6 +409,7 @@ impl SingleGpuBackend { "Qwen3.5 async prefill failed ({err}); stream drain failed: {sync_err}" )); } + revert_scheduled_requests(&self.kv_cache, kvs.iter_mut().map(Box::as_mut)); return Err(err); } }; @@ -255,26 +437,56 @@ impl SingleGpuBackend { chunk: &mut ScheduledChunk, active: &mut [ActiveRequest35], ) -> Result { - let (windows, kvs, mut recs) = single_prefill_views(chunk)?; - let (decode_tokens, mut decode_kvs) = single_decode_views(active); - self.model.unified_step( - &windows, - kvs, - &mut recs, + 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 unified step received TP chunk state"); + }; + let prefill_views = self.schedule_prefill_views(kvs, &chunk.windows)?; + let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); + let decode_tokens: Vec = active.iter().map(|r| r.last_token).collect(); + let decode_views = match self.schedule_decode_views(active) { + Ok(views) => views, + Err(error) => { + revert_scheduled_requests(&self.kv_cache, kvs.iter_mut().map(Box::as_mut)); + return Err(error); + } + }; + let result = self.model.unified_step( + &window_refs, + &prefill_views, + &mut rec_refs, &decode_tokens, - &mut decode_kvs, + &decode_views, + self.kv_cache.buffer(), &mut self.graph_state, - ) + ); + if result.is_err() { + revert_scheduled_requests(&self.kv_cache, kvs.iter_mut().map(Box::as_mut)); + revert_scheduled_requests( + &self.kv_cache, + active.iter_mut().filter_map(active_request_kv), + ); + } + result } pub(super) fn decode_graph(&mut self, active: &mut [ActiveRequest35]) -> Result<()> { - let (tokens, mut kvs) = single_decode_views(active); - self.model.batch_decode_graph( - &tokens, - &mut kvs, + let token_ids: Vec = active.iter().map(|r| r.last_token).collect(); + let views = self.schedule_decode_views(active)?; + let result = self.model.batch_decode_graph( + &token_ids, + &views, + self.kv_cache.buffer(), &mut self.graph_state, crate::batch_decode::DecodeGraphUse::Serve, - ) + ); + if result.is_err() { + revert_scheduled_requests( + &self.kv_cache, + active.iter_mut().filter_map(active_request_kv), + ); + } + result } pub(super) fn sample_prefill_logits( @@ -389,19 +601,39 @@ impl SingleGpuBackend { } impl TpSchedulerBackend { + pub(super) fn alloc_prefill_state( + &mut self, + req: &SchedulerRequest, + ) -> Result<(PrefillBackendState, usize), AdmissionError> { + 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(), + true, + ) + .map_err(AdmissionError::from)?; + Ok((PrefillBackendState::Tp { request_id }, cached_tokens)) + } + pub(super) fn new( model_path: &str, device_ordinals: &[usize], max_batch: usize, max_prefill_tokens: usize, enable_cuda_graph: bool, + prefix_snapshot_bytes: usize, ) -> Result { - let executor = Qwen35TpExecutor::from_runtime_with_limits( + let executor = Qwen35TpExecutor::from_runtime_with_limits_and_prefix( model_path, enable_cuda_graph, device_ordinals, max_batch, max_prefill_tokens, + prefix_snapshot_bytes, )?; Ok(Self { executor, @@ -438,24 +670,14 @@ impl TpSchedulerBackend { pub(super) fn available_pages( &self, - active: &[ActiveRequest35], - prefilling: &[PrefillingRequest35], + _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)) + self.executor.available_pages() } pub(super) fn execute_prefill_chunk( - &self, + &mut self, chunk: &ScheduledChunk, sample_seed: u64, ) -> Result>> { @@ -463,23 +685,27 @@ impl TpSchedulerBackend { let result = self .executor .execute_prefill_chunks_with_seed(&items, sample_seed)?; - align_prefill_results(chunk, &result) - .map_err(|err| self.executor.poison_artifact_contract("prefill", &err)) + align_prefill_results(chunk, &result).map_err(|err| { + self.executor + .poison_after_mutation("prefill artifacts", &err) + }) } pub(super) fn execute_decode( - &self, + &mut self, active: &[ActiveRequest35], sample_seed: u64, ) -> Result> { let items = tp_decode_items(active)?; let result = self.executor.execute_decode_items(&items, sample_seed)?; - align_decode_results(active, &result) - .map_err(|err| self.executor.poison_artifact_contract("decode", &err)) + align_decode_results(active, &result).map_err(|err| { + self.executor + .poison_after_mutation("decode artifacts", &err) + }) } pub(super) fn execute_unified( - &self, + &mut self, chunk: &ScheduledChunk, active: &[ActiveRequest35], decode_sample_seed: u64, @@ -494,17 +720,17 @@ impl TpSchedulerBackend { let result = self.executor.execute_unified(&plan)?; let prefill = align_prefill_results(chunk, &result.prefill).map_err(|err| { self.executor - .poison_artifact_contract("unified prefill", &err) + .poison_after_mutation("unified prefill artifacts", &err) })?; let decode = align_decode_results(active, &result.decode).map_err(|err| { self.executor - .poison_artifact_contract("unified decode", &err) + .poison_after_mutation("unified decode artifacts", &err) })?; Ok(AlignedUnifiedArtifacts { prefill, decode }) } pub(super) fn drop_request( - &self, + &mut self, request_id: RequestId, expectation: DropExpectation, ) -> Result<()> { @@ -546,6 +772,25 @@ impl TpSchedulerBackend { } impl SchedulerBackend { + pub(super) 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(), + } + } + + pub(super) 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, + } + } + pub(super) fn max_batch(&self) -> usize { match self { Self::Single(backend) => backend.max_batch(), @@ -585,15 +830,13 @@ impl SchedulerBackend { } } - pub(super) fn alloc_prefill_state(&mut self) -> Result { + pub(super) fn alloc_prefill_state( + &mut self, + req: &SchedulerRequest, + ) -> Result<(PrefillBackendState, usize), AdmissionError> { 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), } } diff --git a/pegainfer-qwen35/src/scheduler/mod.rs b/pegainfer-qwen35/src/scheduler/mod.rs index ff726d1b4..5fb7433d8 100644 --- a/pegainfer-qwen35/src/scheduler/mod.rs +++ b/pegainfer-qwen35/src/scheduler/mod.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) +//! - `RecurrentState` alongside content-hashed `RequestKv` (hybrid attention) //! - `BatchDecodeGraphState` for CUDA Graph batch decode (stable-address slots) mod backend; @@ -25,7 +25,6 @@ use cudarc::driver::sys; use log::debug; use log::info; use log::warn; -use pegainfer_core::kv_pool::KvState; use pegainfer_core::tensor::HiddenStates; use pegainfer_frontend::engine::EngineHandle as SchedulerHandle; use pegainfer_frontend::engine::FinishReason; @@ -38,6 +37,9 @@ use pegainfer_frontend::engine::TokenLogprob; use pegainfer_frontend::engine::TokenSink; use pegainfer_frontend::engine::panic_message; use pegainfer_frontend::sampler::SamplingParams; +use pegainfer_kv_cache::KvCacheManager; +use pegainfer_kv_cache::KvView; +use pegainfer_kv_cache::RequestKv; use rand::SeedableRng; use rand::rngs::StdRng; use tokio::sync::mpsc; @@ -53,7 +55,6 @@ use self::plan::RejectReason; use self::plan::admit_pending_requests; use self::plan::choose_prefill_budget; use self::plan::compaction_after_retire; -use self::plan::max_kv_tokens; use self::plan::plan_prefill_chunks; use self::plan::prefilling_future_pages; use self::plan::slot_for_new_request; @@ -67,6 +68,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::DropExpectation; use crate::tp_executor::Qwen35TpExecutor; @@ -107,7 +110,7 @@ struct PrefillingRequest35 { enum ActiveBackendState { Single { - kv: KvState, + kv: Box, /// Index into `BatchDecodeGraphState.slot_states`. graph_slot_idx: usize, }, @@ -120,8 +123,13 @@ enum ActiveBackendState { } enum PrefillBackendState { - Single { kv: KvState, rec: RecurrentState }, - Tp { request_id: RequestId }, + Single { + kv: Box, + rec: RecurrentState, + }, + Tp { + request_id: RequestId, + }, } struct TerminalRequest { @@ -360,15 +368,15 @@ pub(crate) fn start_with_capacity_and_policy( ); // 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 backend = SingleGpuBackend::new(model, max_batch, decode_overlap)?; + let total_blocks = backend.capacity_pages_for_requests(); let kv_total_blocks = total_blocks as u64; - let block_size = model.kv_pool().layout().page_size; + let block_size = backend.page_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, decode_overlap)?; let (submit_tx, submit_rx) = mpsc::unbounded_channel(); let (startup_tx, startup_rx) = std_mpsc::channel(); @@ -426,6 +434,7 @@ pub(crate) fn start_tp_with_capacity( max_batch: usize, max_prefill_tokens: usize, enable_cuda_graph: bool, + prefix_snapshot_bytes: usize, ) -> Result { assert!( max_prefill_tokens > 0, @@ -437,6 +446,7 @@ pub(crate) fn start_tp_with_capacity( max_batch, max_prefill_tokens, enable_cuda_graph, + prefix_snapshot_bytes, )?; let servable = servable_len( backend.max_position_embeddings(), @@ -475,15 +485,6 @@ pub(crate) fn start_tp_with_capacity( ) } -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 servable_len(max_context: usize, max_pages: usize, page_size: usize) -> u32 { max_context .min(max_pages.saturating_mul(page_size)) @@ -637,11 +638,8 @@ where "request pruned before scheduling: request_id={:?} phase=prefill cursor={}", removed.req.request_id, removed.cursor ); - let expectation = if removed.cursor == 0 { - DropExpectation::MustBeAbsent - } else { - DropExpectation::MustExist - }; + // Admission now creates rank-local recurrent state, even on a cold miss. + let expectation = DropExpectation::MustExist; if let Err(err) = backend.drop_prefill_state(&removed.backend_state, expectation) { return Err(FatalSchedulerError::new(err.to_string()).with_request(removed)); } @@ -912,31 +910,58 @@ fn scheduler_loop( } // 6. Move freshly admitted prompts into the chunked-prefill queue. - for req in admission.pending { + let mut admitted = admission.pending.into_iter(); + while let Some(req) = admitted.next() { debug!( "request admitted: request_id={:?} prompt_len={} max_tokens={}", req.request_id, 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, - }), - Err(e) => { - warn!("failed to allocate recurrent state for new request: {e}"); + match backend.alloc_prefill_state(&req) { + Ok((backend_state, cached_tokens)) => { + let scheduled_at_unix_s = unix_now_s(); + let _ = 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, + }); + prefilling.push(PrefillingRequest35 { + backend_state, + cursor: cached_tokens, + step_chunk: 0, + req, + }); + } + Err(AdmissionError::Recoverable(error)) => { + warn!("failed to admit new request: {error}"); let _ = req.token_tx.send(TokenEvent::Error { - message: e.to_string(), + message: error.to_string(), prompt_tokens: req.prompt_tokens.len(), completion_tokens: 0, }); } + Err(AdmissionError::Fatal(error)) => { + let kv_total_blocks = backend.capacity_pages_for_requests() as u64; + let failure = FatalSchedulerError::new(error.to_string()) + .with_request(req) + .with_requests(admitted); + terminal_scheduler_shutdown( + &mut submit_rx, + &load_tx, + kv_total_blocks, + active, + prefilling, + Vec::new(), + admission.deferred, + inflight_prefill.take(), + failure, + ); + return; + } } } - deferred = admission.deferred; // 7. Choose this tick's prefill budget, take that chunk off the front of @@ -963,7 +988,11 @@ fn scheduler_loop( &prefill_queue, decode_overlap, ); - 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. @@ -1044,7 +1073,8 @@ fn send_rejection(req: &SchedulerRequest, reason: RejectReason) { req.max_tokens ), RejectReason::KvBudget => { - let max_request_tokens = max_kv_tokens(req.prompt_tokens.len(), req.max_tokens); + let max_request_tokens = + plan::request_lifetime_tokens(req.prompt_tokens.len(), req.max_tokens); format!( "request requires more KV pages than this model instance can provide: prompt_tokens={}, max_request_tokens={max_request_tokens}", req.prompt_tokens.len() @@ -1083,7 +1113,14 @@ fn prefill_batch( }; let prefill_sample_seed = rand::RngExt::random(rng); match single.sample_prefill_logits(&chunk.reqs, &logits, prefill_sample_seed) { - Ok((tokens, logprobs)) => PrefillStepArtifacts::Single { tokens, logprobs }, + Ok((tokens, logprobs)) => { + if let Err(error) = single.apply_prefill(&mut chunk, &tokens) { + return Err( + FatalSchedulerError::new(error.to_string()).with_requests(chunk.reqs) + ); + } + PrefillStepArtifacts::Single { tokens, logprobs } + } Err(e) => { warn!("prefill sampling failed: {e}"); fail_chunk(chunk, &e.to_string()); @@ -1141,7 +1178,7 @@ fn finish_async_prefill( inflight: InflightPrefill, ) -> std::result::Result<(), FatalSchedulerError> { let InflightPrefill { - chunk, + mut chunk, output, sample_seed, } = inflight; @@ -1157,6 +1194,9 @@ fn finish_async_prefill( return Ok(()); } }; + if let Err(error) = single.apply_prefill(&mut chunk, &tokens) { + return Err(FatalSchedulerError::new(error.to_string()).with_requests(chunk.reqs)); + } let artifacts = PrefillStepArtifacts::Single { tokens, logprobs }; promote_or_requeue(single, active, prefilling, chunk, &artifacts) } @@ -1229,8 +1269,10 @@ fn unified_step_sched( // 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, decode_seed)?; + if output.decoded + && let Err(failure) = process_decode_logits(backend, active, decode_seed) + { + return Err(failure.with_requests(chunk.reqs)); } let prefill_logits = output @@ -1246,6 +1288,9 @@ fn unified_step_sched( return Ok(()); } }; + if let Err(error) = backend.apply_prefill(&mut chunk, &tokens) { + return Err(FatalSchedulerError::new(error.to_string()).with_requests(chunk.reqs)); + } let prefill = PrefillStepArtifacts::Single { tokens, logprobs }; promote_or_requeue(backend, active, prefilling, chunk, &prefill) } @@ -1313,6 +1358,11 @@ fn decode_step_with_seed( }, }; + if let SchedulerBackend::Single(single) = backend { + single + .apply_decode(active, &tokens) + .map_err(|e| FatalSchedulerError::new(e.to_string()))?; + } dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec) } @@ -1338,6 +1388,9 @@ fn process_decode_logits( } }; + backend + .apply_decode(active, &tokens) + .map_err(|e| FatalSchedulerError::new(e.to_string()))?; dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec) } @@ -1562,7 +1615,10 @@ struct InflightPrefill { enum ScheduledChunkBackendState { Single { - kvs: Vec, + // Keep request allocations stable while an async prefill owns the chunk; + // failure cleanup can then move each request back through one ownership path. + #[allow(clippy::vec_box)] + kvs: Vec>, recs: Vec, }, Tp { @@ -1622,6 +1678,7 @@ impl From> for ScheduledChunk { fn take_prefill_chunks( prefilling: &mut Vec, prefill_budget: usize, + snapshot_stride: Option, ) -> Vec { let remaining: Vec = prefilling .iter() @@ -1630,7 +1687,7 @@ 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 } @@ -1917,3 +1974,31 @@ fn split_scheduled_backend_state( #[cfg(test)] mod tests; + +fn active_request_kv(request: &mut ActiveRequest35) -> Option<&mut RequestKv> { + match &mut request.backend_state { + ActiveBackendState::Single { kv, .. } => Some(kv), + ActiveBackendState::Tp { .. } => None, + } +} +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}"); + } + } +} +fn unix_now_s() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0.0, |duration| duration.as_secs_f64()) +} +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) + }) +} diff --git a/pegainfer-qwen35/src/scheduler/plan.rs b/pegainfer-qwen35/src/scheduler/plan.rs index 86bc66545..143c79102 100644 --- a/pegainfer-qwen35/src/scheduler/plan.rs +++ b/pegainfer-qwen35/src/scheduler/plan.rs @@ -124,7 +124,7 @@ pub(super) fn admit_pending_requests( continue; } - let request_pages = pages_needed(max_kv_tokens(prompt_len, max_tokens), page_size); + let request_pages = request_lifetime_pages(prompt_len, max_tokens, page_size); if request_pages > max_request_pages { rejected.push((req, RejectReason::KvBudget)); continue; @@ -151,12 +151,19 @@ fn pages_needed(token_count: usize, page_size: usize) -> usize { token_count.div_ceil(page_size) } -// Prefill samples the first output token but does not append it to KV. A -// generated token occupies KV only when it is fed as the next decode input. -// Therefore N returned completion tokens occupy at most N - 1 generated-token -// KV slots. -pub(super) fn max_kv_tokens(prompt_len: usize, max_tokens: usize) -> usize { - prompt_len.saturating_add(max_tokens.saturating_sub(1)) +// One-token completions finish after prefill, before schedule_decode provisions +// a dangling generation block. Multi-token requests can draw that extra block, +// so admission reserves prompt + max_tokens for their lifetime peak. +fn request_lifetime_pages(prompt_len: usize, max_tokens: usize, page_size: usize) -> usize { + pages_needed(request_lifetime_tokens(prompt_len, max_tokens), page_size) +} + +pub(super) fn request_lifetime_tokens(prompt_len: usize, max_tokens: usize) -> usize { + if max_tokens <= 1 { + prompt_len + } else { + prompt_len.saturating_add(max_tokens) + } } fn current_active_tokens(req: ActiveKvBudget) -> usize { @@ -168,7 +175,7 @@ fn active_future_pages(active: &[ActiveKvBudget], page_size: usize) -> usize { active .iter() .map(|req| { - let max_pages = pages_needed(max_kv_tokens(req.prompt_len, req.max_tokens), page_size); + let max_pages = request_lifetime_pages(req.prompt_len, req.max_tokens, page_size); let current_pages = pages_needed(current_active_tokens(*req), page_size); assert!( current_pages <= max_pages, @@ -196,7 +203,7 @@ pub(super) fn prefilling_future_pages(prefilling: &[PrefillKvBudget], page_size: prefilling .iter() .map(|req| { - let max_pages = pages_needed(max_kv_tokens(req.prompt_len, req.max_tokens), page_size); + let max_pages = request_lifetime_pages(req.prompt_len, req.max_tokens, page_size); let current_pages = pages_needed(req.current_tokens, page_size); assert!( current_pages <= max_pages, @@ -244,6 +251,8 @@ pub(super) fn compaction_after_retire( #[cfg(test)] mod tests { + use pegainfer_kv_cache::BlockPool; + use super::*; #[derive(Clone, Debug)] @@ -727,13 +736,13 @@ mod tests { fn admission_counts_pending_generation_budget() { let outcome = admit_pending_requests( vec![ - pending_with_max(1, 16, 17), // 32 KV tokens -> 2 pages - pending(2, 16), // 16 KV tokens -> 1 page + pending_with_max(1, 16, 17), // 33-token peak -> 3 pages + pending(2, 16), // 16-token peak -> 1 page ], &[], 8, 16, - 2, + 3, 8, usize::MAX, |req| req.prompt_len, @@ -755,23 +764,23 @@ mod tests { ActiveKvBudget { prompt_len: 16, generated_count: 1, // current 16 tokens -> 1 page - max_tokens: 33, // max 48 KV tokens -> 3 pages + max_tokens: 33, // 49-token lifetime peak -> 4 pages }, ActiveKvBudget { prompt_len: 16, generated_count: 17, // current 32 tokens -> 2 pages - max_tokens: 17, // max 32 KV tokens -> 2 pages + max_tokens: 17, // 33-token lifetime peak -> 3 pages }, ActiveKvBudget { prompt_len: 9, generated_count: 8, // current 16 tokens -> 1 page - max_tokens: 24, // max 32 KV tokens -> 2 pages + max_tokens: 24, // 33-token lifetime peak -> 3 pages }, ]; assert_eq!( active_future_pages(&active, 16), - 3, + 6, "active admission reserves only future page growth, not pages already held" ); } @@ -781,14 +790,14 @@ mod tests { let active = [ActiveKvBudget { prompt_len: 16, generated_count: 1, // current 16 tokens -> 1 page - max_tokens: 49, // max 64 KV tokens -> 4 pages; future growth = 3 pages + max_tokens: 49, // 65-token lifetime peak -> 5 pages; future growth = 4 pages }]; let outcome = admit_pending_requests( vec![pending(1, 16), pending(2, 16)], &active, 8, 16, - 4, + 5, 8, usize::MAX, |req| req.prompt_len, @@ -808,7 +817,7 @@ mod tests { fn admission_rejects_impossible_request_without_blocking_later_fit() { let outcome = admit_pending_requests( vec![ - pending_with_max(1, 16, 65), // 80 KV tokens -> 5 pages + pending_with_max(1, 16, 65), // 81-token peak -> 6 pages pending(2, 16), ], &[], @@ -833,7 +842,7 @@ mod tests { #[test] fn admission_allows_request_at_single_request_page_cap() { let outcome = admit_pending_requests( - vec![pending_with_max(1, 16, 49)], // 64 KV tokens -> 4 pages + vec![pending_with_max(1, 16, 48)], // 64-token peak -> 4 pages &[], 1, 16, @@ -849,9 +858,48 @@ mod tests { assert!(outcome.rejected.is_empty()); } + fn kvbm_peak_pages(prompt_len: usize, max_tokens: usize, page_size: usize) -> usize { + let pool = BlockPool::new(page_size, 256).expect("test block pool"); + let baseline = pool.available_blocks(); + let mut peak = 0; + let mut kv = pool.new_request(vec![1; prompt_len], max_tokens, None); + + kv.schedule_prefill(prompt_len, &pool) + .expect("schedule prefill"); + peak = peak.max(baseline - pool.available_blocks()); + kv.apply_prefill(100, &pool).expect("apply prefill"); + for step in 1..max_tokens { + kv.schedule_decode(&pool).expect("schedule decode"); + peak = peak.max(baseline - pool.available_blocks()); + kv.apply_decode(100 + step as u32, &pool) + .expect("apply decode"); + } + kv.release().expect("release request KV"); + assert_eq!(pool.available_blocks(), baseline); + peak + } + + #[test] + fn lifetime_pages_cover_request_kv_peak_draw() { + let page_size = 16; + assert_eq!(request_lifetime_pages(16, 17, page_size), 3); + assert_eq!(kvbm_peak_pages(16, 17, page_size), 3); + + for prompt_len in 1..=32 { + for max_tokens in 1..=32 { + let reserved = request_lifetime_pages(prompt_len, max_tokens, page_size); + let peak = kvbm_peak_pages(prompt_len, max_tokens, page_size); + assert!( + reserved >= peak, + "prompt={prompt_len}, max_tokens={max_tokens}: reserved={reserved}, peak={peak}" + ); + } + } + } + #[test] fn one_token_completion_on_page_boundary_uses_only_prompt_page() { - assert_eq!(max_kv_tokens(16, 1), 16); + assert_eq!(request_lifetime_pages(16, 1, 16), 1); let outcome = admit_pending_requests( vec![pending_with_max(1, 16, 1)], &[], @@ -949,13 +997,13 @@ mod tests { #[test] fn prefilling_future_pages_reserves_only_remaining_growth() { - // current 16 tokens -> 1 page; max 48 KV tokens -> 3 pages; future = 2. + // Current 16 tokens use 1 page; the 49-token lifetime peak uses 4. let prefilling = [PrefillKvBudget { current_tokens: 16, prompt_len: 16, max_tokens: 33, }]; - assert_eq!(prefilling_future_pages(&prefilling, 16), 2); + assert_eq!(prefilling_future_pages(&prefilling, 16), 3); } #[test] @@ -964,14 +1012,14 @@ mod tests { PrefillKvBudget { current_tokens: 0, // just admitted, nothing in KV yet -> 0 pages prompt_len: 16, - max_tokens: 17, // max 32 KV tokens -> 2 pages; future = 2 + max_tokens: 17, // 33-token lifetime peak -> 3 pages }, PrefillKvBudget { current_tokens: 32, // 2 pages held prompt_len: 40, - max_tokens: 9, // max 48 KV tokens -> 3 pages; future = 1 + max_tokens: 9, // 49-token lifetime peak -> 4 pages; future = 2 }, ]; - assert_eq!(prefilling_future_pages(&prefilling, 16), 3); + assert_eq!(prefilling_future_pages(&prefilling, 16), 5); } } diff --git a/pegainfer-qwen35/src/scheduler/tests.rs b/pegainfer-qwen35/src/scheduler/tests.rs index 9bc6f1e3f..f887e43b7 100644 --- a/pegainfer-qwen35/src/scheduler/tests.rs +++ b/pegainfer-qwen35/src/scheduler/tests.rs @@ -306,7 +306,7 @@ fn closed_resident_work_is_absent_from_post_prune_load() { assert_eq!(backend.retired_active, vec![RequestId::new(10)]); assert_eq!( backend.dropped_prefilling, - vec![(RequestId::new(12), DropExpectation::MustBeAbsent)] + vec![(RequestId::new(12), DropExpectation::MustExist)] ); } @@ -795,6 +795,16 @@ fn collect_finished_with_timeout( } } +#[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_lifetime_kv_and_context_limits() { let rejection_message = |reason: RejectReason, max_tokens: usize| { @@ -814,10 +824,10 @@ fn send_rejection_reports_lifetime_kv_and_context_limits() { } }; - let kv = rejection_message(RejectReason::KvBudget, 65); + let kv = rejection_message(RejectReason::KvBudget, 49); assert!( - kv.contains("max_request_tokens=80"), - "rejection should report the full lifetime KV request: {kv}" + kv.contains("max_request_tokens=65"), + "rejection should report the lifetime KV peak used by admission: {kv}" ); let context = rejection_message(RejectReason::ContextLength { limit: 32 }, 17); @@ -832,13 +842,18 @@ fn send_rejection_reports_lifetime_kv_and_context_limits() { } #[test] -fn echo_request_is_rejected_before_backend_admission() { - let (echo_tx, mut echo_rx) = TokenSink::standalone(); +fn prompt_logprobs_request_is_rejected_before_backend_admission() { + let (unsupported_tx, mut unsupported_rx) = TokenSink::standalone(); let (regular_tx, mut regular_rx) = TokenSink::standalone(); - let mut echo = test_request_with_shape("unsupported-echo", echo_tx, vec![1, 2, 3], 4); - echo.prompt_logprobs = Some(0); + let mut unsupported = test_request_with_shape( + "unsupported-prompt-logprobs", + unsupported_tx, + vec![1, 2, 3], + 4, + ); + unsupported.prompt_logprobs = Some(0); let regular = test_request("regular", regular_tx); - let mut pending = vec![echo, regular]; + let mut pending = vec![unsupported, regular]; reject_unsupported_prompt_logprobs(&mut pending); @@ -848,7 +863,7 @@ fn echo_request_is_rejected_before_backend_admission() { pending[0].prompt_logprobs.is_none(), "only requests eligible for backend admission may remain" ); - match echo_rx.blocking_recv().map(|(_, event)| event) { + match unsupported_rx.blocking_recv().map(|(_, event)| event) { Some(TokenEvent::Rejected { message, prompt_tokens, @@ -858,7 +873,7 @@ fn echo_request_is_rejected_before_backend_admission() { assert_eq!(prompt_tokens, 3); assert_eq!(completion_tokens, 0); } - event => panic!("expected unsupported echo rejection, got {event:?}"), + event => panic!("expected unsupported prompt-logprobs rejection, got {event:?}"), } assert!(matches!( regular_rx.try_recv(), @@ -884,8 +899,8 @@ fn tp2_scheduler_runs_forced_mixed_steps() { else { return; }; - let handle = - start_tp_with_capacity(&model_path, 42, &[0, 1], 2, 1, false).expect("start TP2 scheduler"); + let handle = start_tp_with_capacity(&model_path, 42, &[0, 1], 2, 1, false, 0) + .expect("start TP2 scheduler"); let (decode_tx, mut decode_rx) = TokenSink::standalone(); let (prefill_tx, mut prefill_rx) = TokenSink::standalone(); diff --git a/pegainfer-qwen35/src/tp_executor.rs b/pegainfer-qwen35/src/tp_executor.rs index e02ace2c3..b898e6a42 100644 --- a/pegainfer-qwen35/src/tp_executor.rs +++ b/pegainfer-qwen35/src/tp_executor.rs @@ -4,6 +4,7 @@ //! and state are sharded per rank; decode rows run as one batched forward per //! rank plus one batched rank-0 sampling pass. +use std::collections::HashMap; use std::collections::HashSet; use std::panic::AssertUnwindSafe; use std::panic::catch_unwind; @@ -19,8 +20,10 @@ use std::thread::{self}; use std::time::Instant; use anyhow::Result; -use pegainfer_core::kv_pool::KvState; use pegainfer_frontend::sampler::SamplingParams; +use pegainfer_kv_cache::KvCacheManager; +use pegainfer_kv_cache::KvView; +use pegainfer_kv_cache::RequestKv; use crate::batch_decode::DecodeGraphUse; use crate::batch_decode_graph::BATCH_BUCKETS; @@ -41,6 +44,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::Qwen35PrefixCache; +use crate::prefix_cache::RecurrentStateStore; +use crate::prefix_cache::SnapshotGuard; use crate::recurrent_state::LinearStatePointerTables; use crate::recurrent_state::RecurrentState; use crate::weights::ModelRuntimeConfig; @@ -55,6 +61,20 @@ const TP_PRECAPTURE_TIMEOUT: std::time::Duration = std::time::Duration::from_sec const TP_RUNTIME_MEMORY_RESERVE_BYTES: usize = 512 * 1024 * 1024; const TRITON_AOT_DEVICE_TABLE_LEN: usize = 16; +#[derive(Debug)] +pub(crate) enum TpBeginRequestError { + Recoverable(anyhow::Error), + Fatal(anyhow::Error), +} + +impl TpBeginRequestError { + pub(crate) fn into_inner(self) -> anyhow::Error { + match self { + Self::Recoverable(error) | Self::Fatal(error) => error, + } + } +} + /// One controller-barriered phase of the TP decode-graph pre-capture sweep. /// /// Capture and launch are separate phases because a captured collective's @@ -88,23 +108,40 @@ pub(crate) struct TpSlotCompaction { #[allow(dead_code)] enum TpWorkerCommand { + RestoreRequest { + request_id: RequestId, + snapshot_slot: Option, + boundary: usize, + start: Arc, + resp: mpsc::Sender, + }, + SaveSnapshot { + request_id: RequestId, + snapshot_slot: usize, + start: Arc, + resp: mpsc::Sender, + }, Ping { resp: mpsc::Sender, }, RunPrefillChunks { chunks: Vec, + kv_views: Vec, sample_seed: u64, start: Arc, resp: mpsc::Sender, }, RunDecodeStep { requests: Vec, + kv_views: Vec, sample_seed: u64, start: Arc, resp: mpsc::Sender, }, RunUnifiedStep { plan: TpUnifiedPlan, + prefill_views: Vec, + decode_views: Vec, start: Arc, resp: mpsc::Sender, }, @@ -142,6 +179,7 @@ enum TpWorkerCommand { #[derive(Debug)] enum TpWorkerReply { + Position(usize), Ack, DropAck { existed: bool, @@ -238,6 +276,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, @@ -360,6 +400,261 @@ pub(crate) struct TpUnifiedResult { } impl Qwen35TpExecutor { + fn schedule_prefill(&mut self, chunks: &[TpPrefillChunkItem]) -> Result> { + anyhow::ensure!( + chunks + .iter() + .all(|c| self.request_kvs.contains_key(&c.request_id)), + "TP prefill requires begin_request" + ); + for (scheduled, chunk) in chunks.iter().enumerate() { + if let Err(error) = self.kv_cache.schedule_prefill( + self.request_kvs.get_mut(&chunk.request_id).unwrap(), + chunk.prompt_tokens.len(), + ) { + for prior in &chunks[..scheduled] { + self.request_kvs + .get_mut(&prior.request_id) + .unwrap() + .revert_schedule()?; + } + return Err(error); + } + } + Ok(chunks + .iter() + .map(|c| { + self.kv_cache + .prefill_view(&self.request_kvs[&c.request_id], c.prompt_tokens.len()) + }) + .collect()) + } + fn schedule_decode(&mut self, requests: &[TpDecodeStepItem]) -> Result> { + anyhow::ensure!( + requests + .iter() + .all(|r| self.request_kvs.contains_key(&r.request_id)), + "TP decode requires begin_request" + ); + for (scheduled, request) in requests.iter().enumerate() { + if let Err(error) = self + .kv_cache + .schedule_decode(self.request_kvs.get_mut(&request.request_id).unwrap()) + { + for prior in &requests[..scheduled] { + self.request_kvs + .get_mut(&prior.request_id) + .unwrap() + .revert_schedule()?; + } + return Err(error); + } + } + Ok(requests + .iter() + .map(|r| self.kv_cache.decode_view(&self.request_kvs[&r.request_id])) + .collect()) + } + + fn revert_requests<'a>(&mut self, request_ids: impl IntoIterator) { + for request_id in request_ids { + if let Some(request) = self.request_kvs.get_mut(request_id) { + if let Err(error) = self.kv_cache.revert_schedule(request) { + log::warn!( + "failed to revert Qwen3.5 TP request {} KV schedule: {error}", + request_id.get() + ); + } + } + } + } + fn apply_prefill_result( + &mut self, + chunks: &[TpPrefillChunkItem], + result: &PrefillResult, + ) -> Result<()> { + for chunk in chunks { + let first_token = if chunk.finish_prefill { + Some( + result + .requests + .iter() + .find(|r| r.request_id == chunk.request_id) + .ok_or_else(|| anyhow::anyhow!("missing final prefill artifact"))? + .first_token, + ) + } else { + None + }; + let kv = self.request_kvs.get_mut(&chunk.request_id).unwrap(); + let boundary = self.kv_cache.apply_prefill(kv, first_token)?; + if let Some(reservation) = self.kv_cache.reserve_snapshot(kv, boundary)? { + match self.broadcast_save_snapshot(chunk.request_id, reservation.recurrent_slot()) { + Ok(positions) if positions.iter().all(|&p| p == boundary) => { + self.kv_cache.publish_snapshot(reservation); + } + result => { + self.kv_cache.abort_snapshot(reservation); + let positions = result?; + anyhow::bail!( + "TP snapshot boundary mismatch: {positions:?}, expected {boundary}" + ); + } + } + } + } + Ok(()) + } + fn apply_decode_result( + &mut self, + requests: &[TpDecodeStepItem], + result: &DecodeResult, + ) -> Result<()> { + for request in requests { + let token = result + .requests + .iter() + .find(|r| r.request_id == request.request_id) + .ok_or_else(|| anyhow::anyhow!("missing decode artifact"))? + .token; + self.kv_cache.apply_decode( + self.request_kvs.get_mut(&request.request_id).unwrap(), + token, + )?; + } + Ok(()) + } + + /// Admit one request. An error leaves no distributed request state behind + /// unless the executor is poisoned and can no longer serve another command. + 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 { + self.poison + .ensure_healthy() + .map_err(TpBeginRequestError::Fatal)?; + if self.request_kvs.contains_key(&request_id) { + return Err(TpBeginRequestError::Recoverable(anyhow::anyhow!( + "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) + .map_err(TpBeginRequestError::Recoverable)?; + let boundary = restore.as_ref().map_or(0, SnapshotGuard::boundary); + let snapshot_slot = restore.as_ref().map(SnapshotGuard::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(TpBeginRequestError::Fatal(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, DropExpectation::MustExist); + let _ = self.kv_cache.release_request(&mut kv); + return Err(TpBeginRequestError::Fatal( + self.poison_after_mutation("request restore", &error), + )); + } + } + } else { + if !positions.iter().all(|&position| position == 0) { + let _ = self.drop_request(request_id, DropExpectation::MustExist); + let _ = self.kv_cache.release_request(&mut kv); + let error = anyhow::anyhow!( + "Qwen3.5 TP cold request restored non-zero positions {positions:?}" + ); + return Err(TpBeginRequestError::Fatal( + self.poison_after_mutation("request restore", &error), + )); + } + 0 + }; + self.request_kvs.insert(request_id, kv); + Ok(cached_tokens) + } + + 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={}, occupancy={}/{}", + self.world_size, + stats.joint_hits, + stats.joint_hit_tokens, + stats.kv_only_fallbacks, + stats.snapshot_misses, + stats.inserts, + stats.evictions, + self.kv_cache.snapshot_occupancy(), + self.kv_cache.snapshot_slots(), + ); + } + + fn broadcast_restore_request( + &self, + request_id: RequestId, + snapshot_slot: Option, + boundary: usize, + ) -> Result> { + let resp_rx = self.dispatch_mutating("RestoreRequest", |start, resp| { + TpWorkerCommand::RestoreRequest { + request_id, + snapshot_slot, + boundary, + start, + resp, + } + })?; + let responses = + recv_runtime_responses(&resp_rx, self.world_size, "RestoreRequest", &self.poison)?; + validate_dispatched_responses( + validate_position_responses(responses, self.world_size), + "RestoreRequest", + &self.poison, + ) + } + + fn broadcast_save_snapshot( + &self, + request_id: RequestId, + snapshot_slot: usize, + ) -> Result> { + let resp_rx = self.dispatch_mutating("SaveSnapshot", |start, resp| { + TpWorkerCommand::SaveSnapshot { + request_id, + snapshot_slot, + start, + resp, + } + })?; + let responses = + recv_runtime_responses(&resp_rx, self.world_size, "SaveSnapshot", &self.poison)?; + validate_dispatched_responses( + validate_position_responses(responses, self.world_size), + "SaveSnapshot", + &self.poison, + ) + } + pub fn from_runtime_with_capacity( model_path: &str, enable_cuda_graph: bool, @@ -381,6 +676,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 { validate_cuda_ordinals(device_ordinals)?; anyhow::ensure!( @@ -402,6 +715,7 @@ impl Qwen35TpExecutor { enable_cuda_graph, tensor_parallel: Some(TensorParallelConfig::try_from((rank, world_size))?), device_ordinal, + prefix_snapshot_bytes, }, )?); } @@ -425,16 +739,28 @@ impl Qwen35TpExecutor { ); }); } - let page_size = first.kv_pool().layout().page_size; + let page_size = first.kv_buffer().layout().page_size; 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_page_size = model.kv_buffer().layout().page_size; 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}" ); - min_capacity_pages = min_capacity_pages.min(model.kv_pool().capacity_pages()); + min_capacity_pages = min_capacity_pages.min(model.kv_buffer().num_blocks()); } + let snapshot_slots = first.prefix_snapshot_slots(); + anyhow::ensure!( + models + .iter() + .all(|m| m.prefix_snapshot_slots() == snapshot_slots), + "TP snapshot slot counts differ" + ); + let kv_cache = Qwen35PrefixCache::new( + KvCacheManager::from_buffer(first.kv_buffer().clone(), min_capacity_pages)?, + snapshot_slots, + )?; + let padding_block_id = kv_cache.pool().padding_block_id(); let capacity_pages_for_requests = min_capacity_pages.saturating_sub(1); let max_position_embeddings = first.config().max_position_embeddings; let eos_token_id = first.config().eos_token_id; @@ -455,6 +781,7 @@ impl Qwen35TpExecutor { max_batch, max_prefill_tokens, graph_enabled, + padding_block_id, nccl_id, Arc::clone(&startup_gate), Arc::clone(&effective_max_batch), @@ -524,6 +851,8 @@ impl Qwen35TpExecutor { let executor = Self { workers, + kv_cache, + request_kvs: HashMap::new(), poison, world_size, max_batch: min_rank_max_batch, @@ -678,18 +1007,42 @@ impl Qwen35TpExecutor { ) } - pub fn execute_prefill(&self, plan: PrefillPlan<'_>) -> Result { - anyhow::ensure!( - !plan.requests.is_empty(), - "Qwen3.5 TP prefill plan requires at least one request" - ); + pub fn execute_prefill(&mut self, plan: PrefillPlan<'_>) -> Result { + self.poison.ensure_healthy()?; let chunks: Vec = plan .requests .iter() .cloned() .map(TpPrefillChunkItem::from) .collect(); - let result = self.execute_prefill_chunks(&chunks)?; + validate_prefill_layout( + &chunks, + self.max_batch, + self.max_position_embeddings, + self.request_kvs.len(), + |request_id| self.request_kvs.contains_key(&request_id), + )?; + + for (index, request) in plan.requests.iter().enumerate() { + if let Err(error) = self.begin_request( + request.request_id, + &request.prompt_tokens, + self.max_position_embeddings - request.prompt_tokens.len(), + None, + false, + ) { + let error = error.into_inner(); + // Once an earlier request commits, a later failure leaves a + // partially admitted plan; the executor must not keep serving. + if index == 0 { + return Err(error); + } + return Err(self.poison_after_mutation("prefill admission", &error)); + } + } + let result = self + .execute_prefill_chunks(&chunks) + .map_err(|error| self.poison_after_mutation("prefill", &error))?; if self.graph_enabled { // Convenience-API slot tracking: every prefill plan item finishes // prefill (TpPrefillChunkItem::from sets finish_prefill), so each @@ -705,12 +1058,12 @@ impl Qwen35TpExecutor { Ok(result) } - 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 { @@ -720,25 +1073,39 @@ impl Qwen35TpExecutor { "Qwen3.5 TP prefill chunk command requires at least one chunk" ); validate_prefill_chunks(chunks)?; + let kv_views = self.schedule_prefill(chunks)?; let chunks = chunks.to_vec(); - let resp_rx = self.dispatch_mutating("prefill chunks", |start, resp| { - TpWorkerCommand::RunPrefillChunks { - chunks: chunks.clone(), - sample_seed, - start, - resp, + let result = (|| { + let resp_rx = self.dispatch_mutating("prefill chunks", |start, resp| { + TpWorkerCommand::RunPrefillChunks { + chunks: chunks.clone(), + kv_views: kv_views.clone(), + sample_seed, + start, + resp, + } + })?; + let responses = + recv_runtime_responses(&resp_rx, self.world_size, "prefill chunks", &self.poison)?; + validate_dispatched_responses( + validate_prefill_responses(responses, self.world_size), + "prefill chunks", + &self.poison, + ) + })(); + let result = match result { + Ok(result) => result, + Err(error) => { + self.revert_requests(chunks.iter().map(|chunk| &chunk.request_id)); + return Err(error); } - })?; - let responses = - recv_runtime_responses(&resp_rx, self.world_size, "prefill chunks", &self.poison)?; - validate_dispatched_responses( - validate_prefill_responses(responses, self.world_size), - "prefill chunks", - &self.poison, - ) + }; + self.apply_prefill_result(&chunks, &result) + .map_err(|e| self.poison_after_mutation("prefill apply", &e))?; + 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" @@ -790,7 +1157,7 @@ impl Qwen35TpExecutor { } pub(crate) fn execute_decode_items( - &self, + &mut self, requests: &[TpDecodeStepItem], sample_seed: u64, ) -> Result { @@ -800,55 +1167,104 @@ impl Qwen35TpExecutor { "Qwen3.5 TP decode plan requires at least one request" ); validate_decode_requests(requests)?; + let kv_views = self.schedule_decode(requests)?; let requests = requests.to_vec(); - let resp_rx = self.dispatch_mutating("decode step", |start, resp| { - TpWorkerCommand::RunDecodeStep { - requests: requests.clone(), - sample_seed, - start, - resp, + let result = (|| { + let resp_rx = self.dispatch_mutating("decode step", |start, resp| { + TpWorkerCommand::RunDecodeStep { + requests: requests.clone(), + kv_views: kv_views.clone(), + sample_seed, + start, + resp, + } + })?; + let responses = + recv_runtime_responses(&resp_rx, self.world_size, "decode step", &self.poison)?; + validate_dispatched_responses( + validate_decode_responses(responses, self.world_size), + "decode step", + &self.poison, + ) + })(); + let result = match result { + Ok(result) => result, + Err(error) => { + self.revert_requests(requests.iter().map(|request| &request.request_id)); + return Err(error); } - })?; - let responses = - recv_runtime_responses(&resp_rx, self.world_size, "decode step", &self.poison)?; - validate_dispatched_responses( - validate_decode_responses(responses, self.world_size), - "decode step", - &self.poison, - ) + }; + self.apply_decode_result(&requests, &result) + .map_err(|e| self.poison_after_mutation("decode apply", &e))?; + Ok(result) } - pub(crate) fn execute_unified(&self, plan: &TpUnifiedPlan) -> Result { + pub(crate) fn execute_unified(&mut self, plan: &TpUnifiedPlan) -> Result { self.poison.ensure_healthy()?; validate_unified_plan(plan, self.max_batch)?; - let resp_rx = self.dispatch_mutating("unified step", |start, resp| { - TpWorkerCommand::RunUnifiedStep { - plan: plan.clone(), - start, - resp, + let prefill_views = self.schedule_prefill(&plan.prefill)?; + let decode_views = match self.schedule_decode(&plan.decode) { + Ok(views) => views, + Err(error) => { + for chunk in &plan.prefill { + let _ = self + .request_kvs + .get_mut(&chunk.request_id) + .unwrap() + .revert_schedule(); + } + return Err(error); } - })?; - let responses = - recv_runtime_responses(&resp_rx, self.world_size, "unified step", &self.poison)?; - validate_dispatched_responses( - validate_unified_responses(responses, self.world_size), - "unified step", - &self.poison, - ) + }; + let result = (|| { + let resp_rx = self.dispatch_mutating("unified step", |start, resp| { + TpWorkerCommand::RunUnifiedStep { + plan: plan.clone(), + prefill_views: prefill_views.clone(), + decode_views: decode_views.clone(), + start, + resp, + } + })?; + let responses = + recv_runtime_responses(&resp_rx, self.world_size, "unified step", &self.poison)?; + validate_dispatched_responses( + validate_unified_responses(responses, self.world_size), + "unified step", + &self.poison, + ) + })(); + let result = match result { + Ok(result) => result, + Err(error) => { + self.revert_requests(plan.prefill.iter().map(|item| &item.request_id)); + self.revert_requests(plan.decode.iter().map(|item| &item.request_id)); + return Err(error); + } + }; + self.apply_prefill_result(&plan.prefill, &result.prefill) + .map_err(|e| self.poison_after_mutation("unified prefill apply", &e))?; + self.apply_decode_result(&plan.decode, &result.decode) + .map_err(|e| self.poison_after_mutation("unified decode apply", &e))?; + Ok(result) } - pub(crate) fn poison_artifact_contract( + pub(crate) fn poison_after_mutation( &self, operation: &'static str, err: &anyhow::Error, ) -> anyhow::Error { let reason = self.poison.poison(format!( - "invalid Qwen3.5 TP {operation} artifact set: {err:#}" + "Qwen3.5 TP {operation} failed after mutation: {err:#}" )); anyhow::anyhow!(reason) } - pub fn drop_request(&self, request_id: RequestId, expectation: DropExpectation) -> Result<()> { + pub fn drop_request( + &mut self, + request_id: RequestId, + expectation: DropExpectation, + ) -> Result<()> { let compaction = self.track_retired_slot(request_id); self.drop_request_with_compaction(request_id, expectation, compaction) } @@ -857,7 +1273,7 @@ impl Qwen35TpExecutor { /// already applied to its own dense-slot bookkeeping. Workers apply the /// move and poison on occupancy mismatch; eager workers ignore it. pub(crate) fn drop_request_with_compaction( - &self, + &mut self, request_id: RequestId, expectation: DropExpectation, compaction: Option, @@ -876,7 +1292,11 @@ impl Qwen35TpExecutor { validate_drop_responses(responses, self.world_size, expectation), "drop request", &self.poison, - ) + )?; + if let Some(mut kv) = self.request_kvs.remove(&request_id) { + self.kv_cache.release_request(&mut kv)?; + } + Ok(()) } /// Convenience-API tracker: swap-remove the retired request and derive the @@ -944,6 +1364,7 @@ impl Qwen35TpExecutor { &self.poison, |start, resp| TpWorkerCommand::RunPrefillChunks { chunks: chunks.clone(), + kv_views: Vec::new(), sample_seed: 0, start, resp, @@ -1164,6 +1585,7 @@ impl TpWorker { max_batch: usize, max_prefill_tokens: usize, graph_enabled: bool, + padding_block_id: i32, nccl_id: cudarc::nccl::safe::Id, startup_gate: Arc, effective_max_batch: Arc, @@ -1189,6 +1611,7 @@ impl TpWorker { max_batch, max_prefill_tokens, graph_enabled, + padding_block_id, ); let prepared = match prepared { Ok((prepared, rank_max_batch)) => { @@ -1259,6 +1682,7 @@ impl Drop for TpWorker { } struct TpWorkerState { + snapshots: RecurrentStateStore, rank: usize, _world_size: usize, max_batch: usize, @@ -1283,10 +1707,12 @@ struct TpWorkerState { } struct TpWorkerPrepared { + snapshots: RecurrentStateStore, rank: usize, world_size: usize, max_batch: usize, model: Qwen35Model, + padding_block_id: i32, decode_buffers: BatchDecodeBuffers35, sample_scratch: pegainfer_sample::SampleScratch, cublas_guard: CublasThreadGuard, @@ -1295,7 +1721,6 @@ struct TpWorkerPrepared { struct TpRequestState { request_id: RequestId, phase: TpRequestPhase, - kv: KvState, /// Prefill-owned recurrent state. Graph mode moves it into the decode slot /// on the request's first decode row (`None` afterwards); the eager path /// keeps it for the request's whole lifetime. @@ -1324,8 +1749,15 @@ impl TpWorkerPrepared { requested_max_batch: usize, max_prefill_tokens: usize, graph_enabled: bool, + padding_block_id: i32, ) -> Result<(Self, usize)> { let cublas_guard = bind_worker_thread(&model)?; + let snapshots = RecurrentStateStore::new( + model.device_ctx(), + model.config(), + model.geometry, + model.prefix_snapshot_slots(), + )?; let (free_bytes, total_bytes) = model .device_ctx() .ctx @@ -1394,7 +1826,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, + model.kv_buffer().num_blocks(), + padding_block_id, + )?; let sample_scratch = pegainfer_sample::SampleScratch::new( model.device_ctx(), model.config().selection_vocab, @@ -1402,10 +1838,12 @@ impl TpWorkerPrepared { )?; Ok(( Self { + snapshots, rank, world_size, max_batch, model, + padding_block_id, decode_buffers, sample_scratch, cublas_guard, @@ -1422,10 +1860,12 @@ impl TpWorkerPrepared { poison: Arc, ) -> Result { let Self { + snapshots, rank, world_size, max_batch, mut model, + padding_block_id, decode_buffers, sample_scratch, cublas_guard, @@ -1454,12 +1894,17 @@ impl TpWorkerPrepared { // cuStreamBeginCapture during the pre-capture sweep. model.tune_decode_gemm_algos()?; let slots = bucket_for(effective_max_batch); - let graph_state = model.create_batch_decode_graph_state_with_capacity(slots)?; + let graph_state = model.create_batch_decode_graph_state_with_capacity( + slots, + model.kv_buffer().num_blocks(), + padding_block_id, + )?; (Some(graph_state), vec![None; slots]) } else { (None, Vec::new()) }; Ok(TpWorkerState { + snapshots, rank, _world_size: world_size, max_batch: effective_max_batch, @@ -1499,15 +1944,100 @@ fn effective_recurrent_capacity( } impl TpWorkerState { + 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(), + self.model.geometry, + )?; + 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 state = TpRequestState { + request_id, + phase: TpRequestPhase::Prefilling, + recurrent: Some(recurrent), + }; + self.requests.push(state); + 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.as_ref().ok_or_else(|| { + anyhow::anyhow!("cannot snapshot a TP request after graph-slot promotion") + })?; + self.snapshots + .save(self.model.device_ctx(), snapshot_slot, recurrent)?; + Ok(TpWorkerReply::Position(recurrent.seq_len)) + } + #[allow(clippy::needless_pass_by_value)] fn run(&mut self, rx: mpsc::Receiver) { while let Ok(command) = rx.recv() { let fatal = match command { + TpWorkerCommand::RestoreRequest { + request_id, + snapshot_slot, + boundary, + start, + resp, + } => { + if start.wait() == TpCommandDecision::Cancel { + false + } else { + let result = self.restore_request(request_id, snapshot_slot, boundary); + self.respond(resp, "restore request", result) + } + } + TpWorkerCommand::SaveSnapshot { + request_id, + snapshot_slot, + start, + resp, + } => { + if start.wait() == TpCommandDecision::Cancel { + false + } else { + let result = self.save_snapshot(request_id, snapshot_slot); + self.respond(resp, "save snapshot", result) + } + } TpWorkerCommand::Ping { resp } => { self.respond(resp, "ping", Ok(TpWorkerReply::Ack)) } TpWorkerCommand::RunPrefillChunks { chunks, + kv_views, sample_seed, start, resp, @@ -1515,12 +2045,13 @@ impl TpWorkerState { if start.wait() == TpCommandDecision::Cancel { false } else { - 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, start, resp, @@ -1528,15 +2059,21 @@ impl TpWorkerState { if start.wait() == TpCommandDecision::Cancel { false } else { - 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 { plan, start, resp } => { + TpWorkerCommand::RunUnifiedStep { + plan, + prefill_views, + decode_views, + start, + resp, + } => { if start.wait() == TpCommandDecision::Cancel { false } else { - let result = self.execute_unified(&plan); + let result = self.execute_unified(&plan, &prefill_views, &decode_views); self.respond(resp, "unified step", result) } } @@ -1630,9 +2167,10 @@ impl TpWorkerState { fn execute_prefill_chunks( &mut self, chunks: &[TpPrefillChunkItem], + kv_views: &[KvView], sample_seed: u64, ) -> Result { - let requests = self.execute_prefill_rows(chunks, sample_seed)?; + let requests = self.execute_prefill_rows(chunks, kv_views, sample_seed)?; if self.rank == 0 { Ok(TpWorkerReply::Prefill(PrefillResult { requests })) } else { @@ -1643,6 +2181,7 @@ impl TpWorkerState { fn execute_prefill_rows( &mut self, chunks: &[TpPrefillChunkItem], + kv_views: &[KvView], sample_seed: u64, ) -> Result> { anyhow::ensure!( @@ -1650,6 +2189,10 @@ impl TpWorkerState { "Qwen3.5 TP prefill chunk command requires at least one chunk" ); validate_prefill_chunks(chunks)?; + anyhow::ensure!( + chunks.len() == kv_views.len(), + "TP prefill view count mismatch" + ); let new_requests = chunks .iter() .filter(|chunk| self.request_index(chunk.request_id).is_none()) @@ -1662,8 +2205,10 @@ impl TpWorkerState { 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 (row_idx, chunk) in chunks.iter().enumerate() { + let state_idx = self + .request_index(chunk.request_id) + .ok_or_else(|| anyhow::anyhow!("TP prefill missing restored request state"))?; let state = &mut self.requests[state_idx]; anyhow::ensure!( state.phase == TpRequestPhase::Prefilling, @@ -1680,8 +2225,9 @@ impl TpWorkerState { ]; let logits = self.model.batch_prefill_logits( &prompt, - std::slice::from_mut(&mut state.kv), + std::slice::from_ref(&kv_views[row_idx]), &mut recurrent_refs, + self.model.kv_buffer(), )?; if chunk.finish_prefill { @@ -1708,6 +2254,7 @@ impl TpWorkerState { fn run_decode_batch( &mut self, requests: &[TpDecodeStepItem], + kv_views: &[KvView], sample_seed: u64, ) -> Result> { let bs = requests.len(); @@ -1715,7 +2262,7 @@ impl TpWorkerState { return Ok(Vec::new()); } if self.graph_state.is_some() { - return self.run_decode_batch_graph(requests, sample_seed); + return self.run_decode_batch_graph(requests, kv_views, sample_seed); } // Resolve the worker state slot of every row in command order. @@ -1737,11 +2284,9 @@ impl TpWorkerState { debug_assert!(row_of_state[state_idx].is_none()); row_of_state[state_idx] = Some(row); } - let mut kv_refs: Vec<&mut KvState> = Vec::with_capacity(bs); let mut recurrent_refs: Vec<&mut RecurrentState> = Vec::with_capacity(bs); for state in states_in_row_order(&mut self.requests, &row_of_state) { - let TpRequestState { kv, recurrent, .. } = state; - kv_refs.push(kv); + let TpRequestState { recurrent, .. } = state; recurrent_refs.push( recurrent .as_mut() @@ -1762,7 +2307,8 @@ impl TpWorkerState { let token_ids: Vec = requests.iter().map(|request| request.token_id).collect(); self.model.batch_decode_eager_logits( &token_ids, - &mut kv_refs, + kv_views, + self.model.kv_buffer(), &mut recurrent_refs, &self.decode_pointer_tables, &mut self.decode_buffers, @@ -1793,6 +2339,7 @@ impl TpWorkerState { fn run_decode_batch_graph( &mut self, requests: &[TpDecodeStepItem], + kv_views: &[KvView], sample_seed: u64, ) -> Result> { let bs = requests.len(); @@ -1849,16 +2396,13 @@ impl TpWorkerState { } } - // KV refs in row (slot) order; page tables stay per-step H2D via - // sync_paged_meta inside batch_decode_graph. - let mut kv_refs: Vec<&mut KvState> = states_in_row_order(&mut self.requests, &row_of_state) - .into_iter() - .map(|state| &mut state.kv) - .collect(); + // KV views arrive in row (slot) order; page tables stay per-step H2D via + // sync_paged_views inside batch_decode_graph. let token_ids: Vec = requests.iter().map(|request| request.token_id).collect(); self.model.batch_decode_graph( &token_ids, - &mut kv_refs, + kv_views, + self.model.kv_buffer(), graph_state, DecodeGraphUse::Replay, )?; @@ -1906,9 +2450,10 @@ impl TpWorkerState { fn execute_decode( &mut self, requests: &[TpDecodeStepItem], + kv_views: &[KvView], sample_seed: u64, ) -> Result { - let requests = self.execute_decode_rows(requests, sample_seed)?; + let requests = self.execute_decode_rows(requests, kv_views, sample_seed)?; if self.rank == 0 { Ok(TpWorkerReply::Decode(DecodeResult { requests })) } else { @@ -1919,6 +2464,7 @@ impl TpWorkerState { fn execute_decode_rows( &mut self, requests: &[TpDecodeStepItem], + kv_views: &[KvView], sample_seed: u64, ) -> Result> { anyhow::ensure!( @@ -1926,6 +2472,10 @@ impl TpWorkerState { "Qwen3.5 TP decode command requires at least one request" ); validate_decode_requests(requests)?; + anyhow::ensure!( + requests.len() == kv_views.len(), + "TP decode view count mismatch" + ); anyhow::ensure!( requests.len() <= self.max_batch, "Qwen3.5 TP decode batch {} exceeds worker capacity {}", @@ -1933,18 +2483,24 @@ impl TpWorkerState { self.max_batch ); - self.run_decode_batch(requests, sample_seed) + self.run_decode_batch(requests, kv_views, sample_seed) } - fn execute_unified(&mut self, plan: &TpUnifiedPlan) -> Result { + fn execute_unified( + &mut self, + plan: &TpUnifiedPlan, + prefill_views: &[KvView], + decode_views: &[KvView], + ) -> Result { validate_unified_worker_state(self, plan)?; // The command order is canonical across ranks. Sampling seeds are // selected by the scheduler in decode-then-prefill order, independent // of this forward order. let prefill_requests = - self.execute_prefill_rows(&plan.prefill, plan.prefill_sample_seed)?; - let decode_requests = self.execute_decode_rows(&plan.decode, plan.decode_sample_seed)?; + self.execute_prefill_rows(&plan.prefill, prefill_views, plan.prefill_sample_seed)?; + let decode_requests = + self.execute_decode_rows(&plan.decode, decode_views, plan.decode_sample_seed)?; if self.rank == 0 { Ok(TpWorkerReply::Unified(TpUnifiedResult { @@ -1960,25 +2516,6 @@ impl TpWorkerState { } } - fn ensure_prefill_state(&mut self, request_id: RequestId) -> Result { - if let Some(idx) = self.request_index(request_id) { - return Ok(idx); - } - let recurrent = RecurrentState::new( - self.model.device_ctx(), - self.model.config(), - self.model.geometry, - )?; - let state = TpRequestState { - request_id, - phase: TpRequestPhase::Prefilling, - kv: self.model.alloc_kv(), - recurrent: Some(recurrent), - }; - self.requests.push(state); - Ok(self.requests.len() - 1) - } - fn request_index(&self, request_id: RequestId) -> Option { self.requests .iter() @@ -2016,7 +2553,7 @@ impl TpWorkerState { /// Capture or launch one bucket with synthetic rows. Outputs are /// discarded; the rows exist only to give the recorded kernels valid /// addresses. One real row (token 0 at position 0 over a freshly - /// allocated one-page KV state) selects nothing — the bucket is passed + /// constructed one-page KV view) selects nothing — the bucket is passed /// explicitly — and every other row is padding on the pool's reserved /// padding page, exactly as when serving. The sweep therefore holds one /// KV page at a time regardless of pool size or bucket. @@ -2030,11 +2567,13 @@ impl TpWorkerState { "Qwen3.5 TP pre-capture bucket {bucket} exceeds {} slots", graph_state.slot_states.len() ); - let mut synthetic_kv = self.model.alloc_kv(); - let mut kv_refs = [&mut synthetic_kv]; + // Startup owns the whole buffer; page 0 is scratch until admission. + graph_state.slot_states[0].seq_len = 0; + let synthetic_view = KvView::new(vec![0], 1, self.model.kv_buffer().layout().page_size); self.model.batch_decode_graph_padded( &[0u32], - &mut kv_refs, + &[synthetic_view], + self.model.kv_buffer(), graph_state, graph_use, bucket, @@ -2173,6 +2712,39 @@ fn validate_prefill_chunks(chunks: &[TpPrefillChunkItem]) -> Result<()> { Ok(()) } +fn validate_prefill_layout( + chunks: &[TpPrefillChunkItem], + max_batch: usize, + max_position_embeddings: usize, + resident_count: usize, + mut request_exists: impl FnMut(RequestId) -> bool, +) -> Result<()> { + anyhow::ensure!( + !chunks.is_empty(), + "Qwen3.5 TP prefill plan requires at least one request" + ); + validate_prefill_chunks(chunks)?; + anyhow::ensure!( + resident_count.saturating_add(chunks.len()) <= max_batch, + "Qwen3.5 TP prefill plan would exceed request capacity {max_batch}" + ); + for chunk in chunks { + anyhow::ensure!( + !request_exists(chunk.request_id), + "Qwen3.5 TP request {} already exists", + chunk.request_id.get() + ); + anyhow::ensure!( + chunk.prompt_tokens.len() < max_position_embeddings, + "Qwen3.5 TP prefill request {} with {} prompt tokens leaves no room in the {}-token context window", + chunk.request_id.get(), + chunk.prompt_tokens.len(), + max_position_embeddings + ); + } + Ok(()) +} + /// Worker request states in decode-row order: `row_of_state[i]` is the row /// `states[i]` occupies in this command, `None` when it is not part of it. fn states_in_row_order<'a>( @@ -2533,6 +3105,7 @@ fn validate_unified_responses( fn reply_name(reply: &TpWorkerReply) -> &'static str { match reply { + TpWorkerReply::Position(_) => "snapshot position", TpWorkerReply::Ack => "acknowledgement", TpWorkerReply::DropAck { .. } => "drop acknowledgement", TpWorkerReply::Prefill(_) => "prefill result", @@ -2564,6 +3137,9 @@ fn wait_for_worker_snapshots( response.rank ); match response.result? { + TpWorkerReply::Position(_) => { + anyhow::bail!("expected worker state, got snapshot position") + } TpWorkerReply::Snapshot(snapshot) => { anyhow::ensure!( snapshot.rank == response.rank, @@ -2660,6 +3236,27 @@ fn bind_worker_thread(model: &Qwen35Model) -> Result { Ok(CublasThreadGuard) } +fn validate_position_responses( + responses: Vec, + world_size: usize, +) -> Result> { + let mut positions = vec![None; world_size]; + for response in responses { + anyhow::ensure!( + response.rank < world_size && positions[response.rank].is_none(), + "invalid or duplicate snapshot response rank" + ); + let TpWorkerReply::Position(position) = response.result? else { + anyhow::bail!("expected snapshot position response"); + }; + positions[response.rank] = Some(position); + } + positions + .into_iter() + .map(|p| p.ok_or_else(|| anyhow::anyhow!("missing snapshot rank response"))) + .collect() +} + #[cfg(test)] mod tests { use super::*; @@ -2784,6 +3381,7 @@ mod tests { None, true, )], + kv_views: Vec::new(), sample_seed: 0, start, resp, @@ -2957,6 +3555,41 @@ mod tests { assert!(err.contains("duplicate")); } + #[test] + fn validates_prefill_layout_before_admission() { + let request = |id, tokens| { + TpPrefillChunkItem::new(RequestId::new(id), vec![9707; tokens], None, true) + }; + + validate_prefill_layout(&[request(1, 2)], 2, 4, 1, |_| false) + .expect("one new request fits the remaining slot and context"); + + let err = validate_prefill_layout(&[], 2, 4, 0, |_| false) + .unwrap_err() + .to_string(); + assert!(err.contains("at least one request")); + + let err = validate_prefill_layout(&[request(1, 2), request(1, 2)], 2, 4, 0, |_| false) + .unwrap_err() + .to_string(); + assert!(err.contains("duplicate")); + + let err = validate_prefill_layout(&[request(2, 2)], 2, 4, 1, |id| id == RequestId::new(2)) + .unwrap_err() + .to_string(); + assert!(err.contains("already exists")); + + let err = validate_prefill_layout(&[request(2, 2)], 1, 4, 1, |_| false) + .unwrap_err() + .to_string(); + assert!(err.contains("capacity")); + + let err = validate_prefill_layout(&[request(2, 4)], 2, 4, 0, |_| false) + .unwrap_err() + .to_string(); + assert!(err.contains("leaves no room")); + } + #[test] fn validates_decode_request_shape() { validate_decode_requests(&[TpDecodeStepItem::new( @@ -3004,8 +3637,9 @@ mod tests { ) else { return; }; - 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"); executor .drop_request(RequestId::new(400), DropExpectation::MustBeAbsent) @@ -3082,24 +3716,70 @@ mod tests { ) else { return; }; - 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"); executor .disconnect_worker_receiver_for_test(1) .expect("disconnect rank-1 worker receiver"); + let error = executor + .begin_request( + RequestId::new(420), + &[151_646, 9707], + executor.max_position_embeddings - 2, + None, + false, + ) + .unwrap_err(); + let TpBeginRequestError::Fatal(error) = error else { + panic!("worker disconnect must be fatal") + }; + let err = error.to_string(); + assert!( + err.contains("failed to dispatch RestoreRequest to TP worker rank 1"), + "unexpected error: {err}" + ); + assert!(executor.ping_all().is_err()); + } + + #[test] + #[ignore = "requires two CUDA devices and Qwen3.5 weights"] + fn tp2_invalid_prefill_plan_does_not_begin_earlier_requests() { + let Some(model_path) = crate::test_fixture::model_path_or_skip( + "tp2_invalid_prefill_plan_does_not_begin_earlier_requests", + ) else { + return; + }; + let mut executor = + Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 2) + .expect("start TP2 executor"); + let request_id = RequestId::new(430); + let duplicate = [ + PrefillStepItem::new(request_id, vec![151_646, 9707], None), + PrefillStepItem::new(request_id, vec![9707], None), + ]; + let err = executor .execute_prefill(PrefillPlan { - requests: &[PrefillStepItem::new( - RequestId::new(420), - vec![151_646, 9707], - None, - )], + requests: &duplicate, }) .unwrap_err() .to_string(); - assert!(err.contains("failed to dispatch prefill chunks to TP worker rank 1")); - assert!(executor.ping_all().is_err()); + assert!(err.contains("duplicate")); + assert!(executor.request_kvs.is_empty()); + assert_workers_empty(&executor); + + executor + .execute_prefill(PrefillPlan { + requests: &[PrefillStepItem::new(request_id, vec![151_646, 9707], None)], + }) + .expect("the rejected request ID remains reusable"); + executor + .drop_request(request_id, DropExpectation::MustExist) + .expect("drop retried request"); + assert!(executor.request_kvs.is_empty()); + assert_workers_empty(&executor); } #[test] @@ -3110,8 +3790,9 @@ mod tests { ) else { return; }; - let executor = Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 2) - .expect("start TP2 executor"); + let mut executor = + Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 2) + .expect("start TP2 executor"); let decode_id = RequestId::new(30); let decode_prefill = executor .execute_prefill(PrefillPlan { @@ -3123,6 +3804,9 @@ mod tests { }) .expect("materialize TP2 decode request"); let prefill_id = RequestId::new(31); + executor + .begin_request(prefill_id, &[151_646, 9707], 8, None, false) + .expect("admit unified prefill"); let unified = executor .execute_unified(&TpUnifiedPlan { prefill: vec![TpPrefillChunkItem::new( @@ -3176,7 +3860,7 @@ mod tests { ) else { return; }; - let executor = Qwen35TpExecutor::from_runtime_with_capacity( + let mut executor = Qwen35TpExecutor::from_runtime_with_capacity( &model_path, false, &[0, 1], @@ -3271,8 +3955,9 @@ mod tests { ) else { return; }; - 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 prompt = vec![151_646, 9707]; let clean_id = RequestId::new(300); diff --git a/pegainfer-qwen35/src/unified_forward.rs b/pegainfer-qwen35/src/unified_forward.rs index e4dfc4560..44fe2c157 100644 --- a/pegainfer-qwen35/src/unified_forward.rs +++ b/pegainfer-qwen35/src/unified_forward.rs @@ -13,10 +13,11 @@ use std::sync::Arc; use anyhow::Result; use cudarc::driver::CudaStream; -use pegainfer_core::kv_pool::KvState; use pegainfer_core::tensor::HiddenStates; use pegainfer_frontend::engine::panic_message; use pegainfer_kernels::tensor::StreamOverrideGuard; +use pegainfer_kv_cache::KvBuffer; +use pegainfer_kv_cache::KvView; use super::batch_decode_graph::BatchDecodeGraphState; use super::recurrent_state::RecurrentState; @@ -36,26 +37,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,14 +67,15 @@ impl Qwen35Model { &mut self, stream: Arc, prompts: &[&[u32]], - kv_states: &mut [KvState], + views: &[KvView], + kv_buffer: &KvBuffer, recurrent_states: &mut [&mut RecurrentState], ) -> Result { let cu_stream = stream.cu_stream(); let original_stream = std::mem::replace(&mut self.ctx.stream, stream); let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { let _stream_override = unsafe { StreamOverrideGuard::activate_prefill(cu_stream) }; - self.batch_prefill_logits(prompts, kv_states, recurrent_states) + self.batch_prefill_logits(prompts, views, recurrent_states, kv_buffer) })); self.ctx.stream = original_stream; @@ -100,10 +101,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!( @@ -117,8 +119,9 @@ impl Qwen35Model { } else { Some(self.batch_prefill_logits( prefill_prompts, - prefill_kv_states, + prefill_views, prefill_recurrent_states, + kv_buffer, )?) }; @@ -128,7 +131,8 @@ impl Qwen35Model { } else { self.batch_decode_graph( decode_tokens, - decode_kv_states, + decode_views, + kv_buffer, graph_state, DecodeGraphUse::Serve, )?; @@ -144,10 +148,11 @@ impl Qwen35Model { #[cfg(test)] mod tests { - use pegainfer_core::kv_pool::KvState; use pegainfer_core::tensor::HiddenStates; + use pegainfer_kv_cache::KvCacheManager; use super::*; + use crate::prefix_cache::Qwen35PrefixCache; fn greedy_sample_batch(model: &Qwen35Model, logits: &HiddenStates, rows: usize) -> Vec { let params = vec![pegainfer_frontend::sampler::SamplingParams::default(); rows]; @@ -161,108 +166,112 @@ 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) = - crate::test_fixture::model_path_or_skip("unified_step_decode_matches_graph_decode") - 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, model.geometry).unwrap(), - RecurrentState::new(&model.ctx, &model.config, model.geometry).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, DecodeGraphUse::Serve) - .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, model.geometry).unwrap(), - RecurrentState::new(&model.ctx, &model.config, model.geometry).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, model.geometry).unwrap(), + RecurrentState::new(&model.ctx, &model.config, model.geometry).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(); + } else { + model + .batch_decode_graph( + &tids, + &views, + cache.buffer(), + &mut gs, + DecodeGraphUse::Serve, + ) .unwrap(); - assert!(output.decoded); - let next = greedy_sample_batch(&model, &gs.buffers.logits, 2); - tokens_a.push(next[0]); - tokens_b.push(next[1]); } - (tokens_a, tokens_b) + 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.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) = + crate::test_fixture::model_path_or_skip("unified_step_decode_matches_graph_decode") + 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 30527c48e..de4a84fde 100644 --- a/pegainfer-qwen35/src/weights.rs +++ b/pegainfer-qwen35/src/weights.rs @@ -17,6 +17,7 @@ use pegainfer_core::weight_loader::WeightPrefetch; use pegainfer_core::weight_loader::deserialize_shards; use pegainfer_core::weight_loader::load_shard_info_fixed; use pegainfer_core::weight_loader::mmap_shards; +use pegainfer_kv_cache::KvBuffer; use safetensors::SafeTensors; use super::config::Config35; @@ -34,6 +35,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 { @@ -42,6 +45,7 @@ impl Default for ModelRuntimeConfig { enable_cuda_graph: true, tensor_parallel: None, device_ordinal: 0, + prefix_snapshot_bytes: 0, } } } @@ -58,8 +62,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. @@ -107,11 +113,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, @@ -256,13 +264,12 @@ 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, geometry.local_num_key_value_heads(), config.head_dim, page_size, - ) - .expect("kv layout geometry"); + ); let bytes_per_page = layout.page_stride * std::mem::size_of::(); let (free_bytes, _total_bytes) = cudarc::driver::result::mem_get_info() .map_err(|e| anyhow::anyhow!("cuMemGetInfo failed: {e}"))?; @@ -277,31 +284,45 @@ impl Qwen35Model { let recurrent_reserve = STATES_PER_DECODE_SLOT * max_batch * super::recurrent_state::bytes_per_request(&config, geometry); + let prefix_snapshot_bytes = runtime.prefix_snapshot_bytes; + let snapshot_bytes_per_slot = super::recurrent_state::bytes_per_request(&config, geometry); + 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, geometry.local_num_key_value_heads(), config.head_dim, @@ -319,7 +340,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, @@ -348,12 +370,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) { @@ -466,13 +488,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, @@ -483,7 +513,8 @@ impl Qwen35Model { &self.ctx, &self.config, self.geometry, - &self.kv_pool, + max_total_pages, + padding_page_id, max_batch, ) } @@ -491,14 +522,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.geometry, 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/e2e_scheduler.rs b/pegainfer-qwen35/tests/e2e_scheduler.rs index 0d8ddebd3..3983b3071 100644 --- a/pegainfer-qwen35/tests/e2e_scheduler.rs +++ b/pegainfer-qwen35/tests/e2e_scheduler.rs @@ -718,6 +718,7 @@ fn test_e2e_qwen35_shared_sm_last_decoder() { 8192, pegainfer_qwen35::Qwen35SchedulerPolicy::Off, pegainfer_qwen35::Qwen35DecodeOverlap::Off, + 0, ) .expect("Failed to start Qwen3.5 default-Off scheduler"); let mut off_rx = submit_repeated_token_request( @@ -757,6 +758,7 @@ fn test_e2e_qwen35_shared_sm_last_decoder() { 8192, pegainfer_qwen35::Qwen35SchedulerPolicy::Auto, pegainfer_qwen35::Qwen35DecodeOverlap::SharedSm, + 0, ) .expect("Failed to start Qwen3.5 auto + shared-SM scheduler"); let mut auto_load = auto_handle @@ -811,6 +813,7 @@ fn test_e2e_qwen35_shared_sm_last_decoder() { 8192, pegainfer_qwen35::Qwen35SchedulerPolicy::Off, pegainfer_qwen35::Qwen35DecodeOverlap::SharedSm, + 0, ) .expect("Failed to start Qwen3.5 shared-SM scheduler"); let mut load = handle diff --git a/pegainfer-qwen35/tests/hf_golden_gate.rs b/pegainfer-qwen35/tests/hf_golden_gate.rs index a932c5037..a1bc86ece 100644 --- a/pegainfer-qwen35/tests/hf_golden_gate.rs +++ b/pegainfer-qwen35/tests/hf_golden_gate.rs @@ -511,7 +511,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)]| { @@ -774,7 +779,7 @@ fn build_tp2_graph_executor(model_path: &str, label: &str) -> Option (Stats, Vec) { assert!( @@ -961,17 +966,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); } @@ -991,10 +996,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" @@ -1020,12 +1025,12 @@ fn pega_logprobs_match_hf_golden_within_qwen35_tolerance_tp2_graph() { report_fixture_shape(&golden); let all: Vec = (0..golden.num_seqs).collect(); - let Some(ex) = build_tp2_graph_executor(&model_path, "TP2 graph") else { + let Some(mut ex) = build_tp2_graph_executor(&model_path, "TP2 graph") else { return; }; - let (stats, fp1) = run_tp(&golden, &ex, &all, false); + let (stats, fp1) = run_tp(&golden, &mut ex, &all, false); report_and_assert("TP2 sequential graph", &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 graph replay must reproduce identical logprobs" @@ -1033,7 +1038,7 @@ fn pega_logprobs_match_hf_golden_within_qwen35_tolerance_tp2_graph() { for n in BUCKET_STRADDLES { if all.len() >= n { - let (batched, _) = run_tp(&golden, &ex, &all[..n], true); + let (batched, _) = run_tp(&golden, &mut ex, &all[..n], true); report_and_assert(&format!("TP2 batched graph ({n} padded)"), &batched); } else { eprintln!( @@ -1045,9 +1050,9 @@ fn pega_logprobs_match_hf_golden_within_qwen35_tolerance_tp2_graph() { if golden.num_seqs >= SLOT_COMPACTION_BATCH && golden.decode_len >= 2 { let (compacted, fp1) = - run_tp_with_slot_compaction(&golden, &ex, &all[..SLOT_COMPACTION_BATCH]); + run_tp_with_slot_compaction(&golden, &mut ex, &all[..SLOT_COMPACTION_BATCH]); report_and_assert("TP2 slot-compaction graph", &compacted); - let (_, fp2) = run_tp_with_slot_compaction(&golden, &ex, &all[..SLOT_COMPACTION_BATCH]); + let (_, fp2) = run_tp_with_slot_compaction(&golden, &mut ex, &all[..SLOT_COMPACTION_BATCH]); assert_eq!( fp1, fp2, "TP2 slot-compaction Qwen3.5 graph 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..e8bddfeec --- /dev/null +++ b/pegainfer-qwen35/tests/prefix_cache.rs @@ -0,0 +1,353 @@ +//! 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_frontend::engine::EngineHandle; +use pegainfer_frontend::engine::FinishReason; +use pegainfer_frontend::engine::GenerateRequest; +use pegainfer_frontend::engine::TokenEvent; +use pegainfer_frontend::engine::TokenLogprob; +use pegainfer_frontend::engine::TokenSink; +use pegainfer_frontend::sampler::SamplingParams; +use pegainfer_qwen35::Qwen35LaunchOptions; +use pegainfer_qwen35::Qwen35SchedulerPolicy; + +mod common; + +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 { + common::model_path_or_skip("prefix_cache") +} + +fn start_engine(model_path: &str, tp_size: usize, prefix_cache_mib: usize) -> EngineHandle { + start_engine_with_graph(model_path, tp_size, prefix_cache_mib, tp_size == 1) +} + +fn start_engine_with_graph( + model_path: &str, + tp_size: usize, + prefix_cache_mib: usize, + cuda_graph: bool, +) -> EngineHandle { + pegainfer_qwen35::launch_with_options_policy_and_overlap( + Path::new(model_path), + Qwen35LaunchOptions::new(0, tp_size, cuda_graph, 2, 1024, prefix_cache_mib), + Qwen35SchedulerPolicy::Off, + pegainfer_qwen35::Qwen35DecodeOverlap::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 submit( + handle: &EngineHandle, + prompt_tokens: Vec, + max_tokens: usize, + logprobs: usize, +) -> pegainfer_frontend::engine::TokenStreamReceiver { + let (token_tx, 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: (logprobs > 0).then_some(logprobs), + prompt_logprobs: None, + }) + .expect("submit failed"); + + rx +} + +fn generate( + handle: &EngineHandle, + prompt_tokens: Vec, + max_tokens: usize, + logprobs: usize, +) -> Generation { + let mut rx = submit(handle, prompt_tokens, max_tokens, logprobs); + 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); + ( + 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); + assert_eq!(cold.cached_tokens, 0); + let warm = generate(&handle, long_prompt.clone(), TRACE_TOKENS, TOP_LOGPROBS); + 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); + 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); + assert_eq!(exact_256.cached_tokens, 0); + + let extended = generate(&handle, long_prompt[..320].to_vec(), 1, 0); + assert_eq!(extended.cached_tokens, 256); +} + +fn restore_during_live_decode( + tp_size: usize, + cuda_graph: bool, + overlap: pegainfer_qwen35::Qwen35DecodeOverlap, +) { + let Some(model_path) = model_path_or_skip() else { + return; + }; + let tokenizer = common::load_tokenizer(&model_path); + let prompt = prompt_tokens( + &tokenizer, + "Joint prefix restore must preserve logits while another request decodes. ", + 576, + ); + let handle = pegainfer_qwen35::launch_with_options_policy_and_overlap( + Path::new(&model_path), + Qwen35LaunchOptions::new(0, tp_size, cuda_graph, 2, 1024, 128), + Qwen35SchedulerPolicy::Off, + overlap, + ) + .unwrap(); + let cold = generate(&handle, prompt.clone(), TRACE_TOKENS, TOP_LOGPROBS); + let mut background = submit(&handle, vec![9707], 1024, 0); + loop { + match background.blocking_recv().map(|(_, event)| event) { + Some(TokenEvent::Token { .. }) => break, + Some(TokenEvent::Scheduled { .. }) => {} + event => panic!("background decode did not start: {event:?}"), + } + } + let warm = generate(&handle, prompt.clone(), TRACE_TOKENS, TOP_LOGPROBS); + assert_eq!(warm.cached_tokens, 512); + assert_trace_close("restore during live decode", &cold, &warm); + let mut background_tokens = 1; + while let Ok((_, event)) = background.try_recv() { + match event { + TokenEvent::Token { .. } => background_tokens += 1, + TokenEvent::Finished { .. } => { + panic!("background finished before the mixed-load probe") + } + TokenEvent::Error { message, .. } | TokenEvent::Rejected { message, .. } => { + panic!("{message}") + } + _ => {} + } + } + assert!(background_tokens < 1024); + drop(background); // Exercise cancellation cleanup followed by another cache hit. + let again = generate(&handle, prompt, TRACE_TOKENS, TOP_LOGPROBS); + assert_eq!(again.cached_tokens, 512); + assert_trace_close("restore after cancellation", &cold, &again); +} + +#[test] +fn async_prefill_prefix_restore_during_live_decode() { + restore_during_live_decode(1, true, pegainfer_qwen35::Qwen35DecodeOverlap::SharedSm); +} + +#[test] +#[ignore = "requires two CUDA devices and Qwen3.5 weights"] +fn tp2_prefix_restore_during_mixed_step() { + restore_during_live_decode(2, false, pegainfer_qwen35::Qwen35DecodeOverlap::Off); +} + +#[test] +#[ignore = "requires two CUDA devices and Qwen3.5 weights"] +fn tp2_graph_prefix_restore_during_mixed_step() { + restore_during_live_decode(2, true, pegainfer_qwen35::Qwen35DecodeOverlap::Off); +}