perf: presize paged pool per prefill span and eval grown slabs eagerly

src/lib/mlxcel-core/src/cache/paged.rs

@@ -461,6 +461,11 @@ pub struct PagedBlockPool {
     /// evicts against; the scheduler sets it from the configured / estimated KV
     /// byte budget divided by the per-block byte size.
     block_budget: Option<usize>,
+    /// Number of [`Self::grow_pool`] reallocations that actually grew a slab
+    /// (across all layers). Each one copies an entire layer tensor, so this is
+    /// the observable cost #224's presize exists to avoid; tests pin it and
+    /// observability can surface it.
+    grow_events: u64,
 }
 
 /// Number of physical block rows the pool tensor grows by when a freshly
@@ -494,9 +499,17 @@ impl PagedBlockPool {
             free_rows: vec![Vec::new(); num_layers],
             next_row: vec![0; num_layers],
             block_budget: None,
+            grow_events: 0,
         }
     }
 
+    /// Number of slab-copy pool growths since construction (see
+    /// [`Self::grow_pool`]). A long prefill should presize instead of
+    /// accumulating these.
+    pub fn pool_grow_events(&self) -> u64 {
+        self.grow_events
+    }
+
     pub fn layout(&self) -> &PagedKvLayout {
         &self.layout
     }
@@ -1175,6 +1188,22 @@ impl PagedBlockPool {
         // Walk every block the new tokens span and write its slice.
         let first_block = (first_abs / block_size) as usize;
         let last_block = ((last_abs - 1) / block_size) as usize;
+
+        // Size the layer's pool once for the whole span before any block
+        // write. Without this, assign_row grows the slab in
+        // POOL_GROW_CHUNK_BLOCKS steps and every step reallocates and copies
+        // the entire layer tensor, transiently holding each step's old+new
+        // slab pair in one lazy graph on long prefills (#224).
+        self.presize_for_span(
+            state,
+            layer_idx,
+            first_block,
+            last_block,
+            n_kv_heads,
+            head_dim,
+            k_dtype,
+        )?;
+
         for block_index in first_block..=last_block {
             let block_start_abs = block_index as i32 * block_size;
             let slot_start = (first_abs.max(block_start_abs) - block_start_abs) as usize;
@@ -1210,6 +1239,60 @@ impl PagedBlockPool {
         Ok(())
     }
 
+    /// Ensure the layer's pool tensors can hold every span block in
+    /// `first_block..=last_block` of `state` without incremental growth.
+    ///
+    /// Counts the span blocks that still need a physical row, nets out the
+    /// reusable free rows, and either allocates the pool at the right size
+    /// (first write to the layer) or grows it once. [`Self::assign_row`]
+    /// keeps its grow fallback for rows minted outside the span (e.g. a
+    /// copy-on-write fork), so undershooting here stays correct.
+    #[allow(clippy::too_many_arguments)]
+    fn presize_for_span(
+        &mut self,
+        state: &PagedSequenceState,
+        layer_idx: usize,
+        first_block: usize,
+        last_block: usize,
+        n_kv_heads: i32,
+        head_dim: i32,
+        dtype: i32,
+    ) -> Result<(), String> {
+        let mut unassigned = 0usize;
+        for block_index in first_block..=last_block {
+            let block_id = state.layers[layer_idx].block_ids[block_index];
+            if !self.block_rows[layer_idx].contains_key(&block_id) {
+                unassigned += 1;
+            }
+        }
+        let minted = unassigned.saturating_sub(self.free_rows[layer_idx].len());
+        if minted == 0 {
+            return Ok(());
+        }
+        let target = self.next_row[layer_idx].saturating_add(minted);
+        match self.pool_meta[layer_idx] {
+            Some(meta) if target <= meta.capacity_blocks => Ok(()),
+            Some(_) => self.grow_pool(layer_idx, target),
+            None => {
+                let capacity = target
+                    .div_ceil(POOL_GROW_CHUNK_BLOCKS)
+                    .saturating_mul(POOL_GROW_CHUNK_BLOCKS)
+                    .max(POOL_GROW_CHUNK_BLOCKS);
+                let block_size = self.layout.block_size as i32;
+                let shape = [capacity as i32, block_size, n_kv_heads, head_dim];
+                self.pool_k[layer_idx] = Some(ffi::zeros(&shape, dtype));
+                self.pool_v[layer_idx] = Some(ffi::zeros(&shape, dtype));
+                self.pool_meta[layer_idx] = Some(PagedPoolMeta {
+                    n_kv_heads,
+                    head_dim,
+                    dtype,
+                    capacity_blocks: capacity,
+                });
+                Ok(())
+            }
+        }
+    }
+
     /// Copy the current contents of `src_block_id` into a freshly acquired block
     /// on the same layer and return the new block id. Used by [`write_prefill`]
     /// to fork a shared partial tail block before mutating it (copy-on-write).
@@ -1628,6 +1711,22 @@ impl PagedBlockPool {
             capacity_blocks: new_capacity,
             ..meta
         });
+        self.grow_events += 1;
+        // Materialize the grown copies immediately. The old slabs then return
+        // to the MLX buffer cache before the next layer grows, so a
+        // multi-layer growth episode transiently holds one layer's old+new
+        // pair instead of accumulating every layer's pair in a single lazy
+        // graph (#224).
+        ffi::eval(
+            self.pool_k[layer_idx]
+                .as_ref()
+                .expect("pool_k present after grow"),
+        );
+        ffi::eval(
+            self.pool_v[layer_idx]
+                .as_ref()
+                .expect("pool_v present after grow"),
+        );
         Ok(())
     }
 

src/lib/mlxcel-core/src/cache/paged_pool_tests.rs

@@ -1210,3 +1210,102 @@ fn trim_to_logical_start_empties_the_layer() {
     assert_eq!(layer.block_ids.len(), 0);
     assert!(pool.gather_visible(&state, 0).unwrap().is_none());
 }
+
+// ---------------------------------------------------------------------------
+// 12. write_prefill presize (#224): one allocation for a multi-chunk span
+// ---------------------------------------------------------------------------
+
+#[test]
+fn write_prefill_presizes_pool_in_one_step_for_multi_chunk_span() {
+    // 100 blocks of 4 slots = 400 tokens, spanning four 32-block grow chunks.
+    // presize_for_span must allocate the layer pool once at ceil(100/32)*32 =
+    // 128 rows instead of growing 32 -> 64 -> 96 -> 128 with a full slab copy
+    // at each step.
+    let block_size = 4usize;
+    let total = 400i32;
+    let mut pool = fp16_pool(block_size, 1);
+    let mut state = PagedSequenceState::new(pool.layout());
+
+    let k_prefill = make_block(100.0, total);
+    let v_prefill = make_block(500.0, total);
+    pool.write_prefill(&mut state, 0, &k_prefill, &v_prefill)
+        .unwrap();
+
+    assert_eq!(state.layer(0).unwrap().block_ids.len(), 100);
+    assert_eq!(state.layer(0).unwrap().len, total as usize);
+
+    // Capacity must be the single presized target: 128 rows for K and V, in
+    // FP32 (make_block writes f32), 4 bytes per element.
+    let expected_capacity_rows = 128usize;
+    let expected_bytes = 2 * expected_capacity_rows * block_size * (H as usize) * (D as usize) * 4;
+    assert_eq!(pool.pool_tensor_bytes(), expected_bytes);
+
+    // The regression #224 fixes: presize allocates the pool at the final size
+    // directly, so NO slab-copy growth happened. The old incremental path
+    // converged to the same 128 rows via three grow_pool reallocations, so
+    // this assertion (not the final capacity above) is what pins presize.
+    assert_eq!(pool.pool_grow_events(), 0);
+
+    // Round-trip stays byte-identical to the dense reference.
+    let (gk, gv) = pool
+        .gather_visible(&state, 0)
+        .unwrap()
+        .expect("gather must return data");
+    let dense_k = dense_reference(&[(make_block(100.0, total), 0)], total, 0, total);
+    let dense_v = dense_reference(&[(make_block(500.0, total), 0)], total, 0, total);
+    assert_eq!(flatten_fp32(&gk), flatten_fp32(&dense_k));
+    assert_eq!(flatten_fp32(&gv), flatten_fp32(&dense_v));
+}
+
+#[test]
+fn write_prefill_after_presize_grows_incrementally_and_round_trips() {
+    // First prefill presizes to 32 rows (8 blocks used). A second prefill
+    // pushing past the presized capacity must take the grow_pool path (now
+    // with eager eval) and stay byte-identical.
+    let block_size = 4usize;
+    let first = 32i32; // 8 blocks
+    let second = 160i32; // 40 more blocks -> 48 total, beyond the 32-row chunk
+    let mut pool = fp16_pool(block_size, 1);
+    let mut state = PagedSequenceState::new(pool.layout());
+
+    let k1 = make_block(100.0, first);
+    let v1 = make_block(500.0, first);
+    pool.write_prefill(&mut state, 0, &k1, &v1).unwrap();
+    let bytes_after_first = pool.pool_tensor_bytes();
+
+    let k2 = make_block(200.0, second);
+    let v2 = make_block(600.0, second);
+    pool.write_prefill(&mut state, 0, &k2, &v2).unwrap();
+    assert!(pool.pool_tensor_bytes() > bytes_after_first);
+    assert_eq!(state.layer(0).unwrap().len, (first + second) as usize);
+
+    // The second span needed 40 more rows past the presized 32: exactly one
+    // grow_pool reallocation (32 -> 64), not one per 32-row chunk.
+    assert_eq!(pool.pool_grow_events(), 1);
+
+    let total = first + second;
+    let (gk, gv) = pool
+        .gather_visible(&state, 0)
+        .unwrap()
+        .expect("gather must return data");
+    let dense_k = dense_reference(
+        &[
+            (make_block(100.0, first), 0),
+            (make_block(200.0, second), first as usize),
+        ],
+        total,
+        0,
+        total,
+    );
+    let dense_v = dense_reference(
+        &[
+            (make_block(500.0, first), 0),
+            (make_block(600.0, second), first as usize),
+        ],
+        total,
+        0,
+        total,
+    );
+    assert_eq!(flatten_fp32(&gk), flatten_fp32(&dense_k));
+    assert_eq!(flatten_fp32(&gv), flatten_fp32(&dense_v));
+}