diff --git a/src/lib/mlxcel-core/src/cache/paged.rs b/src/lib/mlxcel-core/src/cache/paged.rs index 93cb1075..79d8bb87 100644 --- a/src/lib/mlxcel-core/src/cache/paged.rs +++ b/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, + /// 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(()) } diff --git a/src/lib/mlxcel-core/src/cache/paged_pool_tests.rs b/src/lib/mlxcel-core/src/cache/paged_pool_tests.rs index ca115f06..f74f9bef 100644 --- a/src/lib/mlxcel-core/src/cache/paged_pool_tests.rs +++ b/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)); +}