feat(qwen35): add joint prefix cache

docs/models/qwen35/tp-implementation.md

@@ -2,7 +2,7 @@
 
 > **TL;DR:** Qwen3.5 TP Phase 1 is implemented as correctness-first eager dense TP: TP2 worker/scheduler execution, short/long HF logits gates, scheduler e2e, and real OpenAI-compatible HTTP serving smoke pass. The branch is rebased onto current `main` with the newer engine, sampling, config, and golden-fixture contracts; remaining TP work is tracked as follow-up, not a Phase 1 claim.
 >
-> **Last touched:** 2026-07
+> **Last touched:** 2026-08
 
 ## Scope
 
@@ -39,6 +39,24 @@ Not implemented in Phase 1:
 - Prefix-cache or recurrent-state snapshot support.
 - Performance claims.
 
+## Post-Phase 1 Follow-up: RequestKv and Joint Prefix Cache
+
+This follow-up unifies the TP and single-GPU request lifecycle around `RequestKv` and adds joint full-attention KV plus recurrent/conv prefix reuse.
+
+1. **Unified KV lifecycle**
+   - The controller uses one `KvCacheManager` for prefill/decode scheduling and commit.
+   - Immutable `KvView`s carry the logical page ids to every worker; each rank writes its local KV shard into its own `KvBuffer`.
+2. **TP capacity and layout**
+   - Startup validates identical KV geometry and snapshot-slot counts across ranks.
+   - The logical pool is capped by the smallest rank-local physical capacity.
+3. **Joint recurrent snapshots**
+   - Key, pin, and LRU metadata is centralized; recurrent/conv tensors remain rank-local.
+   - Publication reserves one common slot, saves it on every rank, verifies the committed boundary, and only then publishes the key.
+   - Restore follows the same all-rank rule and reports a hit only after every rank confirms the selected boundary.
+4. **Opt-in budget**
+   - `--qwen35-prefix-cache-mib` reserves the snapshot budget independently on each rank.
+   - `0` keeps cold serving.
+
 ## Important Fixes
 
 ### Gated q projection layout

kvbm/kvbm-logical/src/integrations/scheduled.rs

@@ -436,11 +436,24 @@ impl<T: BlockMetadata> SchedulableSequence<T> {
     pub fn match_and_add_prefix(
         &mut self,
         manager: &BlockManager<T>,
+    ) -> Result<usize, ScheduleError> {
+        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<T>,
+        requested_max_blocks: usize,
     ) -> Result<usize, ScheduleError> {
         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)

pegainfer-kv-cache/src/manager.rs

@@ -38,6 +38,21 @@ impl KvCacheManager {
         Ok(Self { pool, buffer })
     }
 
+    /// Pair an existing physical KV buffer with a new logical block pool.
+    ///
+    /// `num_blocks` may be smaller than the physical allocation. Tensor-parallel
+    /// executors use this to choose the minimum common logical capacity across
+    /// rank-local buffers while keeping one shared page-id namespace.
+    pub fn from_buffer(buffer: KvBuffer, num_blocks: usize) -> anyhow::Result<Self> {
+        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(

pegainfer-kv-cache/src/pool.rs

@@ -210,6 +210,7 @@ impl BlockPool {
             seq_hashes,
             gpu_hit,
             cacheable,
+            block_size: self.block_size,
             held: gpu_guard,
         }
     }
@@ -253,6 +254,8 @@ pub struct PrefixProbe {
     gpu_hit: usize,
     /// Reuse cap: blocks past this are never matched (the final chunk forwards).
     cacheable: usize,
+    /// Tokens represented by one complete KV block.
+    block_size: usize,
     /// Strong refs keeping matched/loaded blocks resident until prefill.
     held: Vec<ImmutableBlock<()>>,
 }
@@ -271,6 +274,37 @@ impl PrefixProbe {
         self.held.len()
     }
 
+    /// Complete prefix blocks eligible for request reuse.
+    ///
+    /// This is capped by the final-token rule even if a caller extended the
+    /// probe with additional loaded blocks.
+    pub fn reusable_blocks(&self) -> usize {
+        self.held.len().min(self.cacheable)
+    }
+
+    /// Returns the lineage hash identifying the complete reusable prefix ending
+    /// at `boundary_tokens`.
+    ///
+    /// `boundary_tokens` is measured in tokens. For example,
+    /// with a 16-token block size, `boundary_hash(32)` identifies the token
+    /// prefix `[0, 32)`. The returned hash covers the full prefix lineage,
+    /// rather than only the contents of the final block.
+    ///
+    /// Returns `None` when the boundary is zero, is not block-aligned, or
+    /// exceeds the reusable prefix.
+    pub fn boundary_hash(&self, boundary_tokens: usize) -> Option<[u8; 16]> {
+        if boundary_tokens == 0 || !boundary_tokens.is_multiple_of(self.block_size) {
+            return None;
+        }
+        let block_count = boundary_tokens / self.block_size;
+        if block_count > self.reusable_blocks() {
+            return None;
+        }
+        self.seq_hashes
+            .get(block_count - 1)
+            .map(sequence_hash_bytes)
+    }
+
     /// Content hashes to query the CPU tier with: the blocks past the GPU hit,
     /// capped at the reuse boundary. Empty when the GPU hit already covers
     /// every reusable block (nothing to load — prefill normally).
@@ -357,17 +391,53 @@ 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<usize> {
+        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<usize> {
         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.
         self.emitted_blocks = self.seq.assigned_blocks();
         Ok(blocks * self.seq.block_size())
     }
 
+    /// Returns the lineage hash identifying the registered prefix ending at
+    /// `boundary_tokens`.
+    ///
+    /// `boundary_tokens` must be a non-zero multiple of the KV `block_size`,
+    /// and be no more than the number of blocks already registered by this request;
+    /// otherwise this method returns `None`.
+    pub fn registered_boundary_hash(&self, boundary_tokens: usize) -> Option<[u8; 16]> {
+        let block_size = self.seq.block_size();
+        if boundary_tokens == 0 || !boundary_tokens.is_multiple_of(block_size) {
+            return None;
+        }
+        let block_count = boundary_tokens / block_size;
+        if block_count > self.seq.assigned_blocks() {
+            return None;
+        }
+        self.seq
+            .inner()
+            .sequence()
+            .all_sequence_hashes()
+            .get(block_count - 1)
+            .map(sequence_hash_bytes)
+    }
+
     // ── Scheduling (allocates blocks) ──────────────────────────────────
 
     pub fn schedule_prefill(
@@ -760,6 +830,57 @@ mod tests {
         );
     }
 
+    #[test]
+    fn probe_and_attach_can_select_a_shorter_exact_boundary() {
+        let pool = BlockPool::new(16, 32).unwrap();
+        let prompt = (0..80u32).collect::<Vec<_>>();
+
+        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::<Vec<_>>();
+        let mut seed = pool.new_request(prompt.clone(), 4, None);
+        seed.schedule_prefill(64, &pool).expect("seed schedule");
+        seed.apply_prefill(9000, &pool).expect("seed apply");
+        seed.release().expect("seed release");
+
+        let probe = pool.probe_prefix(prompt, None);
+        assert_eq!(probe.gpu_hit_blocks(), 4);
+        assert_eq!(
+            probe.reusable_blocks(),
+            3,
+            "one prompt token must remain uncached"
+        );
+        assert!(probe.boundary_hash(0).is_none());
+        assert!(probe.boundary_hash(47).is_none());
+        assert!(probe.boundary_hash(48).is_some());
+        assert!(probe.boundary_hash(64).is_none());
+    }
+
     #[test]
     fn request_reports_the_lifetime_capacity_it_was_created_with() {
         let pool = BlockPool::new(16, 8).unwrap();

pegainfer-qwen35/Cargo.toml

@@ -11,6 +11,7 @@ cudarc = { workspace = true }
 half = { workspace = true }
 log = { workspace = true }
 pegainfer-core = { workspace = true }
+pegainfer-kv-cache = { workspace = true }
 pegainfer-kernels = { workspace = true }
 pegainfer-sample = { workspace = true }
 rand = { workspace = true }
@@ -52,6 +53,10 @@ required-features = ["qwen35"]
 name = "chunked_prefill"
 required-features = ["qwen35"]
 
+[[test]]
+name = "prefix_cache"
+required-features = ["qwen35"]
+
 [[test]]
 name = "serving_tp2"
 required-features = ["qwen35"]