model: paged KV cache with CPU swap pool, decode graph, qwen3 updates
- paged_kv_cache: new block-paged KV cache; adds a pinned-host swap pool with
a second BlockAllocator, per-sequence Location {Gpu,Cpu}, and lossless
swap_out/swap_in (block-granular D2H/H2D) for vLLM-style preemption.
bytes_per_block helper exposes per-block cost for VRAM-based sizing.
- decode_graph: CUDA-graph decode path.
- qwen3/gpt2/kv_cache: paged prefill/decode forward + related updates.
- tokenizer/bins: BPE updates, new xserv-chat CLI, bench-qwen3 tweaks.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
569
crates/xserv-model/src/paged_kv_cache.rs
Normal file
569
crates/xserv-model/src/paged_kv_cache.rs
Normal file
@@ -0,0 +1,569 @@
|
||||
//! Paged KV cache: vLLM-style block-based KV cache with O(1) allocation
|
||||
//! and indirection via per-sequence block tables.
|
||||
//!
|
||||
//! Physical layout per layer:
|
||||
//! K pool: [total_blocks, num_kv_heads, BLOCK_SIZE, head_dim] BF16
|
||||
//! V pool: same
|
||||
//!
|
||||
//! Logical view per sequence: a list of physical block ids. Token at logical
|
||||
//! position p lives in block_ids[p / BLOCK_SIZE] at slot (p % BLOCK_SIZE).
|
||||
|
||||
use crate::config::ModelConfig;
|
||||
use xserv_cuda::{GpuBuffer, PinnedBuffer};
|
||||
use xserv_tensor::{DType, Tensor};
|
||||
|
||||
pub const BLOCK_SIZE: usize = 16;
|
||||
|
||||
/// Stack-based block allocator: O(1) alloc/free.
|
||||
pub struct BlockAllocator {
|
||||
free_stack: Vec<u32>,
|
||||
total: usize,
|
||||
}
|
||||
|
||||
impl BlockAllocator {
|
||||
pub fn new(total_blocks: usize) -> Self {
|
||||
// Reserve block 0 as a sentinel "null" block (never allocated).
|
||||
// Free list contains [total-1, total-2, ..., 1] so pop returns 1 first.
|
||||
// total_blocks==0 means "disabled" (e.g. swap off): empty free list.
|
||||
let mut free_stack = Vec::with_capacity(total_blocks.saturating_sub(1));
|
||||
for b in (1..total_blocks).rev() {
|
||||
free_stack.push(b as u32);
|
||||
}
|
||||
Self { free_stack, total: total_blocks }
|
||||
}
|
||||
|
||||
pub fn alloc(&mut self) -> Option<u32> {
|
||||
self.free_stack.pop()
|
||||
}
|
||||
|
||||
pub fn free(&mut self, block: u32) {
|
||||
debug_assert!((block as usize) < self.total && block != 0);
|
||||
self.free_stack.push(block);
|
||||
}
|
||||
|
||||
pub fn free_count(&self) -> usize {
|
||||
self.free_stack.len()
|
||||
}
|
||||
|
||||
pub fn total(&self) -> usize {
|
||||
self.total
|
||||
}
|
||||
|
||||
pub fn can_alloc(&self, n: usize) -> bool {
|
||||
self.free_stack.len() >= n
|
||||
}
|
||||
}
|
||||
|
||||
/// Where a sequence's KV blocks currently live.
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
|
||||
pub enum Location {
|
||||
Gpu,
|
||||
Cpu,
|
||||
}
|
||||
|
||||
/// Per-sequence state held in the cache.
|
||||
#[derive(Clone)]
|
||||
pub struct SeqState {
|
||||
/// Block ids into the GPU pool when `location == Gpu`, or into the CPU
|
||||
/// (pinned host) pool when `location == Cpu`.
|
||||
pub block_ids: Vec<u32>,
|
||||
pub seq_len: usize,
|
||||
pub location: Location,
|
||||
}
|
||||
|
||||
pub struct PagedKVCache {
|
||||
// [layer]: GpuBuffer of size total_blocks * nkv * BLOCK_SIZE * hd * elem_size
|
||||
k_pools: Vec<GpuBuffer>,
|
||||
v_pools: Vec<GpuBuffer>,
|
||||
|
||||
// CPU (pinned host) swap pools, same per-layer layout as the GPU pools but
|
||||
// sized for `cpu_total_blocks`. Empty when swap is disabled.
|
||||
cpu_k_pools: Vec<PinnedBuffer>,
|
||||
cpu_v_pools: Vec<PinnedBuffer>,
|
||||
cpu_allocator: BlockAllocator,
|
||||
|
||||
// Bytes occupied by one block within a single layer pool:
|
||||
// num_kv_heads * BLOCK_SIZE * head_dim * elem_size.
|
||||
block_bytes: usize,
|
||||
|
||||
allocator: BlockAllocator,
|
||||
seq_states: Vec<Option<SeqState>>,
|
||||
|
||||
// GPU-resident per-sequence metadata. Uploaded each step via sync_to_gpu().
|
||||
// block_table_gpu: i32 [max_seqs, max_blocks_per_seq]
|
||||
// context_lens_gpu: i32 [max_seqs]
|
||||
block_table_gpu: GpuBuffer,
|
||||
context_lens_gpu: GpuBuffer,
|
||||
// Host-side staging mirroring the GPU buffers above.
|
||||
block_table_host: Vec<i32>,
|
||||
context_lens_host: Vec<i32>,
|
||||
|
||||
// Config
|
||||
num_layers: usize,
|
||||
num_kv_heads: usize,
|
||||
head_dim: usize,
|
||||
elem_size: usize,
|
||||
dtype: DType,
|
||||
device: u32,
|
||||
max_seqs: usize,
|
||||
max_blocks_per_seq: usize,
|
||||
}
|
||||
|
||||
impl PagedKVCache {
|
||||
/// Bytes occupied by all KV blocks for ONE physical block across the whole
|
||||
/// model (both K and V, all layers). Use this to size pools against VRAM.
|
||||
pub fn bytes_per_block(config: &ModelConfig, dtype: DType) -> usize {
|
||||
2 * config.num_layers()
|
||||
* config.num_kv_heads()
|
||||
* BLOCK_SIZE
|
||||
* config.head_dim()
|
||||
* dtype.size_bytes()
|
||||
}
|
||||
|
||||
/// Create a new paged cache.
|
||||
/// - `total_blocks`: total number of physical GPU blocks across all sequences.
|
||||
/// - `cpu_total_blocks`: physical blocks in the pinned-host swap pool (0 = swap off).
|
||||
/// - `max_seqs`: max number of concurrent sequences (slots), incl. swapped.
|
||||
/// - `max_blocks_per_seq`: capacity of the block table per slot
|
||||
/// (must be >= ceil(max_seq_len / BLOCK_SIZE)).
|
||||
pub fn new(
|
||||
config: &ModelConfig,
|
||||
total_blocks: usize,
|
||||
cpu_total_blocks: usize,
|
||||
max_seqs: usize,
|
||||
max_blocks_per_seq: usize,
|
||||
dtype: DType,
|
||||
device: u32,
|
||||
) -> Self {
|
||||
assert!(total_blocks >= 2, "need at least 2 blocks (one is sentinel)");
|
||||
let num_layers = config.num_layers();
|
||||
let num_kv_heads = config.num_kv_heads();
|
||||
let head_dim = config.head_dim();
|
||||
let elem_size = dtype.size_bytes();
|
||||
let block_bytes = num_kv_heads * BLOCK_SIZE * head_dim * elem_size;
|
||||
let pool_bytes = total_blocks * block_bytes;
|
||||
|
||||
let mut k_pools = Vec::with_capacity(num_layers);
|
||||
let mut v_pools = Vec::with_capacity(num_layers);
|
||||
for _ in 0..num_layers {
|
||||
let mut k = GpuBuffer::alloc(pool_bytes).expect("alloc paged K pool");
|
||||
let mut v = GpuBuffer::alloc(pool_bytes).expect("alloc paged V pool");
|
||||
k.zero().unwrap();
|
||||
v.zero().unwrap();
|
||||
k_pools.push(k);
|
||||
v_pools.push(v);
|
||||
}
|
||||
|
||||
// Pinned-host swap pools (one per layer, mirroring the GPU layout).
|
||||
let mut cpu_k_pools = Vec::new();
|
||||
let mut cpu_v_pools = Vec::new();
|
||||
if cpu_total_blocks >= 2 {
|
||||
let cpu_pool_bytes = cpu_total_blocks * block_bytes;
|
||||
for _ in 0..num_layers {
|
||||
cpu_k_pools.push(PinnedBuffer::alloc(cpu_pool_bytes).expect("alloc CPU K swap pool"));
|
||||
cpu_v_pools.push(PinnedBuffer::alloc(cpu_pool_bytes).expect("alloc CPU V swap pool"));
|
||||
}
|
||||
}
|
||||
let cpu_allocator = BlockAllocator::new(if cpu_total_blocks >= 2 { cpu_total_blocks } else { 0 });
|
||||
|
||||
let block_table_gpu =
|
||||
GpuBuffer::alloc(max_seqs * max_blocks_per_seq * std::mem::size_of::<i32>())
|
||||
.expect("alloc block table");
|
||||
let context_lens_gpu =
|
||||
GpuBuffer::alloc(max_seqs * std::mem::size_of::<i32>()).expect("alloc context lens");
|
||||
|
||||
let block_table_host = vec![0i32; max_seqs * max_blocks_per_seq];
|
||||
let context_lens_host = vec![0i32; max_seqs];
|
||||
|
||||
let seq_states = (0..max_seqs).map(|_| None).collect();
|
||||
|
||||
Self {
|
||||
k_pools,
|
||||
v_pools,
|
||||
cpu_k_pools,
|
||||
cpu_v_pools,
|
||||
cpu_allocator,
|
||||
block_bytes,
|
||||
allocator: BlockAllocator::new(total_blocks),
|
||||
seq_states,
|
||||
block_table_gpu,
|
||||
context_lens_gpu,
|
||||
block_table_host,
|
||||
context_lens_host,
|
||||
num_layers,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
elem_size,
|
||||
dtype,
|
||||
device,
|
||||
max_seqs,
|
||||
max_blocks_per_seq,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn num_layers(&self) -> usize { self.num_layers }
|
||||
pub fn num_kv_heads(&self) -> usize { self.num_kv_heads }
|
||||
pub fn head_dim(&self) -> usize { self.head_dim }
|
||||
pub fn dtype(&self) -> DType { self.dtype }
|
||||
pub fn max_seqs(&self) -> usize { self.max_seqs }
|
||||
pub fn max_blocks_per_seq(&self) -> usize { self.max_blocks_per_seq }
|
||||
pub fn free_blocks(&self) -> usize { self.allocator.free_count() }
|
||||
pub fn total_blocks(&self) -> usize { self.allocator.total() }
|
||||
|
||||
pub fn k_pool(&self, layer: usize) -> &GpuBuffer { &self.k_pools[layer] }
|
||||
pub fn v_pool(&self, layer: usize) -> &GpuBuffer { &self.v_pools[layer] }
|
||||
pub fn block_table_gpu(&self) -> &GpuBuffer { &self.block_table_gpu }
|
||||
pub fn context_lens_gpu(&self) -> &GpuBuffer { &self.context_lens_gpu }
|
||||
|
||||
pub fn seq_len(&self, slot: usize) -> usize {
|
||||
self.seq_states[slot].as_ref().map(|s| s.seq_len).unwrap_or(0)
|
||||
}
|
||||
|
||||
pub fn is_slot_free(&self, slot: usize) -> bool {
|
||||
self.seq_states[slot].is_none()
|
||||
}
|
||||
|
||||
/// Register a new sequence at `slot`. Allocates the first block.
|
||||
/// Returns Err(()) if no slot or no blocks are available.
|
||||
pub fn register_sequence(&mut self, slot: usize) -> Result<(), &'static str> {
|
||||
if slot >= self.max_seqs {
|
||||
return Err("slot out of range");
|
||||
}
|
||||
if self.seq_states[slot].is_some() {
|
||||
return Err("slot already in use");
|
||||
}
|
||||
let block = self.allocator.alloc().ok_or("out of blocks")?;
|
||||
self.seq_states[slot] = Some(SeqState {
|
||||
block_ids: vec![block],
|
||||
seq_len: 0,
|
||||
location: Location::Gpu,
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Free all blocks for `slot` and clear the slot. Frees from whichever pool
|
||||
/// (GPU or CPU) the sequence currently lives in.
|
||||
pub fn free_sequence(&mut self, slot: usize) {
|
||||
if let Some(state) = self.seq_states[slot].take() {
|
||||
let alloc = match state.location {
|
||||
Location::Gpu => &mut self.allocator,
|
||||
Location::Cpu => &mut self.cpu_allocator,
|
||||
};
|
||||
for b in state.block_ids {
|
||||
alloc.free(b);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Number of blocks needed to hold `seq_len + new_tokens` tokens, beyond
|
||||
/// what is currently allocated for `slot`.
|
||||
pub fn additional_blocks_needed(&self, slot: usize, new_tokens: usize) -> usize {
|
||||
let state = self.seq_states[slot].as_ref().expect("unregistered slot");
|
||||
let cur = state.block_ids.len();
|
||||
let needed_total = (state.seq_len + new_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE;
|
||||
if needed_total > cur { needed_total - cur } else { 0 }
|
||||
}
|
||||
|
||||
/// Pre-allocate enough physical blocks in `slot` to cover positions
|
||||
/// `[0, end_pos)`. Call once before the per-layer append loop so that
|
||||
/// every layer's append uses the same block table.
|
||||
pub fn ensure_capacity(&mut self, slot: usize, end_pos: usize) {
|
||||
let state = self.seq_states[slot].as_mut().expect("unregistered slot");
|
||||
let needed_total = (end_pos + BLOCK_SIZE - 1) / BLOCK_SIZE;
|
||||
while state.block_ids.len() < needed_total {
|
||||
let b = self.allocator.alloc().expect("out of blocks (caller must check)");
|
||||
assert!(state.block_ids.len() < self.max_blocks_per_seq, "block table overflow");
|
||||
state.block_ids.push(b);
|
||||
}
|
||||
}
|
||||
|
||||
/// Append `num_tokens` of K/V into the paged pool for `slot` at logical
|
||||
/// position `start_pos`. Caller must have called `ensure_capacity(slot, start_pos + num_tokens)`
|
||||
/// first (or accept that this method may also extend block list).
|
||||
/// Does NOT touch `seq_len`. Call `advance_seq_len(slot, num_tokens)` after
|
||||
/// every layer has been written.
|
||||
///
|
||||
/// `k_new`, `v_new`: GPU tensors with logical shape
|
||||
/// [1, num_kv_heads, num_tokens, head_dim]
|
||||
/// stored contiguously (head-major, then tokens, then dim).
|
||||
pub fn append_tokens(
|
||||
&mut self,
|
||||
slot: usize,
|
||||
layer: usize,
|
||||
k_new: &Tensor,
|
||||
v_new: &Tensor,
|
||||
num_tokens: usize,
|
||||
start_pos: usize,
|
||||
) {
|
||||
if num_tokens == 0 { return; }
|
||||
// Make sure blocks exist for the target range.
|
||||
self.ensure_capacity(slot, start_pos + num_tokens);
|
||||
|
||||
let block_ids = self.seq_states[slot].as_ref().unwrap().block_ids.clone();
|
||||
|
||||
let nkv = self.num_kv_heads;
|
||||
let hd = self.head_dim;
|
||||
let es = self.elem_size;
|
||||
let bs = BLOCK_SIZE;
|
||||
|
||||
let k_src = k_new.storage().gpu_buffer();
|
||||
let v_src = v_new.storage().gpu_buffer();
|
||||
|
||||
let k_pool = &mut self.k_pools[layer];
|
||||
let v_pool = &mut self.v_pools[layer];
|
||||
|
||||
let mut t = 0usize;
|
||||
while t < num_tokens {
|
||||
let p = start_pos + t;
|
||||
let logical_blk = p / bs;
|
||||
let slot_in_blk = p % bs;
|
||||
let chunk = (bs - slot_in_blk).min(num_tokens - t);
|
||||
let phys = block_ids[logical_blk] as usize;
|
||||
|
||||
for h in 0..nkv {
|
||||
let src_off = (h * num_tokens + t) * hd * es;
|
||||
let dst_off = ((phys * nkv + h) * bs + slot_in_blk) * hd * es;
|
||||
let count = chunk * hd * es;
|
||||
k_pool.copy_from_device_at(k_src, src_off, dst_off, count).unwrap();
|
||||
v_pool.copy_from_device_at(v_src, src_off, dst_off, count).unwrap();
|
||||
}
|
||||
|
||||
t += chunk;
|
||||
}
|
||||
}
|
||||
|
||||
/// Advance the logical seq_len after append_tokens for ALL layers has completed.
|
||||
pub fn advance_seq_len(&mut self, slot: usize, num_tokens: usize) {
|
||||
let state = self.seq_states[slot].as_mut().expect("unregistered slot");
|
||||
state.seq_len += num_tokens;
|
||||
}
|
||||
|
||||
/// Refresh the host-side block table + context lens from `seq_states`,
|
||||
/// then upload to GPU. Call once per decode step before the paged kernel.
|
||||
pub fn sync_to_gpu(&mut self) {
|
||||
let stride = self.max_blocks_per_seq;
|
||||
for slot in 0..self.max_seqs {
|
||||
let row = &mut self.block_table_host[slot * stride..(slot + 1) * stride];
|
||||
row.fill(0);
|
||||
let len = match &self.seq_states[slot] {
|
||||
Some(s) => {
|
||||
for (i, b) in s.block_ids.iter().enumerate() {
|
||||
row[i] = *b as i32;
|
||||
}
|
||||
s.seq_len as i32
|
||||
}
|
||||
None => 0,
|
||||
};
|
||||
self.context_lens_host[slot] = len;
|
||||
}
|
||||
|
||||
self.upload_metadata();
|
||||
}
|
||||
|
||||
/// Pack the given active slots into rows 0..slots.len() of block_table_gpu
|
||||
/// and context_lens_gpu, then upload. Used by paged decode where the kernel
|
||||
/// iterates over `batch` active sequences in order.
|
||||
pub fn sync_active_batch_to_gpu(&mut self, slots: &[usize]) {
|
||||
let lens: Vec<i32> = slots
|
||||
.iter()
|
||||
.map(|&s| self.seq_states[s].as_ref().unwrap().seq_len as i32)
|
||||
.collect();
|
||||
self.sync_active_batch_with_lens(slots, &lens);
|
||||
}
|
||||
|
||||
/// Like sync_active_batch_to_gpu but uses caller-supplied kv_lens (number
|
||||
/// of valid K/V tokens to attend over per active row). Useful when the
|
||||
/// kv_len for the current step differs from the cached seq_len (e.g.
|
||||
/// before advance_seq_len has run).
|
||||
pub fn sync_active_batch_with_lens(&mut self, slots: &[usize], kv_lens: &[i32]) {
|
||||
assert_eq!(slots.len(), kv_lens.len());
|
||||
assert!(slots.len() <= self.max_seqs, "active batch exceeds max_seqs");
|
||||
let stride = self.max_blocks_per_seq;
|
||||
for row in &mut self.block_table_host {
|
||||
*row = 0;
|
||||
}
|
||||
for cl in &mut self.context_lens_host {
|
||||
*cl = 0;
|
||||
}
|
||||
for (i, &slot) in slots.iter().enumerate() {
|
||||
let s = self.seq_states[slot].as_ref().expect("unregistered slot in active batch");
|
||||
let row = &mut self.block_table_host[i * stride..(i + 1) * stride];
|
||||
for (j, b) in s.block_ids.iter().enumerate() {
|
||||
row[j] = *b as i32;
|
||||
}
|
||||
self.context_lens_host[i] = kv_lens[i];
|
||||
}
|
||||
self.upload_metadata();
|
||||
}
|
||||
|
||||
fn upload_metadata(&mut self) {
|
||||
let bt_bytes = unsafe {
|
||||
std::slice::from_raw_parts(
|
||||
self.block_table_host.as_ptr() as *const u8,
|
||||
self.block_table_host.len() * std::mem::size_of::<i32>(),
|
||||
)
|
||||
};
|
||||
self.block_table_gpu.copy_from_host(bt_bytes).unwrap();
|
||||
|
||||
let cl_bytes = unsafe {
|
||||
std::slice::from_raw_parts(
|
||||
self.context_lens_host.as_ptr() as *const u8,
|
||||
self.context_lens_host.len() * std::mem::size_of::<i32>(),
|
||||
)
|
||||
};
|
||||
self.context_lens_gpu.copy_from_host(cl_bytes).unwrap();
|
||||
}
|
||||
|
||||
/// Materialize a contiguous K/V tensor for a sequence at `layer`, shaped
|
||||
/// [1, num_kv_heads, seq_len, head_dim]. Used for prefill, where Flash
|
||||
/// Attention 2 expects contiguous K/V.
|
||||
///
|
||||
/// Allocates from the cached allocator; the returned Tensors own their storage.
|
||||
pub fn gather_kv_contiguous(&self, slot: usize, layer: usize) -> (Tensor, Tensor) {
|
||||
let state = self.seq_states[slot].as_ref().expect("unregistered slot");
|
||||
let sl = state.seq_len;
|
||||
let nkv = self.num_kv_heads;
|
||||
let hd = self.head_dim;
|
||||
let es = self.elem_size;
|
||||
let bs = BLOCK_SIZE;
|
||||
|
||||
let out_bytes = nkv * sl * hd * es;
|
||||
let mut k_dst = xserv_cuda::allocator::cached_alloc(out_bytes).expect("alloc gather K");
|
||||
let mut v_dst = xserv_cuda::allocator::cached_alloc(out_bytes).expect("alloc gather V");
|
||||
|
||||
let k_pool = &self.k_pools[layer];
|
||||
let v_pool = &self.v_pools[layer];
|
||||
|
||||
let mut p = 0usize;
|
||||
while p < sl {
|
||||
let logical_blk = p / bs;
|
||||
let slot_in_blk = p % bs;
|
||||
let chunk = (bs - slot_in_blk).min(sl - p);
|
||||
let phys = state.block_ids[logical_blk] as usize;
|
||||
|
||||
for h in 0..nkv {
|
||||
let src_off = ((phys * nkv + h) * bs + slot_in_blk) * hd * es;
|
||||
let dst_off = (h * sl + p) * hd * es;
|
||||
let count = chunk * hd * es;
|
||||
k_dst.copy_from_device_at(k_pool, src_off, dst_off, count).unwrap();
|
||||
v_dst.copy_from_device_at(v_pool, src_off, dst_off, count).unwrap();
|
||||
}
|
||||
p += chunk;
|
||||
}
|
||||
|
||||
let shape = &[1usize, nkv, sl, hd];
|
||||
let k = unsafe { tensor_from_owned_buf(k_dst, shape, self.dtype, self.device) };
|
||||
let v = unsafe { tensor_from_owned_buf(v_dst, shape, self.dtype, self.device) };
|
||||
(k, v)
|
||||
}
|
||||
|
||||
// ----- Swapping (vLLM-style preemption to pinned host memory) -----
|
||||
|
||||
pub fn free_cpu_blocks(&self) -> usize { self.cpu_allocator.free_count() }
|
||||
pub fn swap_enabled(&self) -> bool { !self.cpu_k_pools.is_empty() }
|
||||
|
||||
pub fn is_swapped(&self, slot: usize) -> bool {
|
||||
matches!(self.seq_states[slot].as_ref().map(|s| s.location), Some(Location::Cpu))
|
||||
}
|
||||
|
||||
/// Number of physical blocks currently held by `slot` (in either pool).
|
||||
pub fn block_count(&self, slot: usize) -> usize {
|
||||
self.seq_states[slot].as_ref().map(|s| s.block_ids.len()).unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Whether a swapped sequence at `slot` can be brought back (enough free GPU blocks).
|
||||
pub fn can_swap_in(&self, slot: usize) -> bool {
|
||||
self.allocator.can_alloc(self.block_count(slot))
|
||||
}
|
||||
|
||||
/// Whether the GPU sequence at `slot` can be evicted (enough free CPU blocks).
|
||||
pub fn can_swap_out(&self, slot: usize) -> bool {
|
||||
self.cpu_allocator.can_alloc(self.block_count(slot))
|
||||
}
|
||||
|
||||
/// Evict `slot`'s KV from GPU to pinned host memory and free its GPU blocks.
|
||||
/// The slot stays registered (location = Cpu); the sequence is paused.
|
||||
pub fn swap_out(&mut self, slot: usize) -> Result<(), &'static str> {
|
||||
let state = self.seq_states[slot].as_ref().ok_or("swap_out: empty slot")?;
|
||||
if state.location == Location::Cpu { return Ok(()); }
|
||||
let gpu_ids = state.block_ids.clone();
|
||||
let n = gpu_ids.len();
|
||||
if !self.cpu_allocator.can_alloc(n) { return Err("swap_out: CPU pool full"); }
|
||||
|
||||
let cpu_ids: Vec<u32> = (0..n)
|
||||
.map(|_| self.cpu_allocator.alloc().expect("checked can_alloc"))
|
||||
.collect();
|
||||
|
||||
let bb = self.block_bytes;
|
||||
for layer in 0..self.num_layers {
|
||||
for i in 0..n {
|
||||
let g_off = gpu_ids[i] as usize * bb;
|
||||
let c_off = cpu_ids[i] as usize * bb;
|
||||
self.k_pools[layer]
|
||||
.copy_to_host_at(&mut self.cpu_k_pools[layer].as_mut_slice()[c_off..c_off + bb], g_off, bb)
|
||||
.unwrap();
|
||||
self.v_pools[layer]
|
||||
.copy_to_host_at(&mut self.cpu_v_pools[layer].as_mut_slice()[c_off..c_off + bb], g_off, bb)
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
for b in gpu_ids {
|
||||
self.allocator.free(b);
|
||||
}
|
||||
let state = self.seq_states[slot].as_mut().unwrap();
|
||||
state.block_ids = cpu_ids;
|
||||
state.location = Location::Cpu;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Bring `slot`'s KV back from host to GPU and free its CPU blocks.
|
||||
pub fn swap_in(&mut self, slot: usize) -> Result<(), &'static str> {
|
||||
let state = self.seq_states[slot].as_ref().ok_or("swap_in: empty slot")?;
|
||||
if state.location == Location::Gpu { return Ok(()); }
|
||||
let cpu_ids = state.block_ids.clone();
|
||||
let n = cpu_ids.len();
|
||||
if !self.allocator.can_alloc(n) { return Err("swap_in: GPU pool full"); }
|
||||
|
||||
let gpu_ids: Vec<u32> = (0..n)
|
||||
.map(|_| self.allocator.alloc().expect("checked can_alloc"))
|
||||
.collect();
|
||||
|
||||
let bb = self.block_bytes;
|
||||
for layer in 0..self.num_layers {
|
||||
for i in 0..n {
|
||||
let g_off = gpu_ids[i] as usize * bb;
|
||||
let c_off = cpu_ids[i] as usize * bb;
|
||||
self.k_pools[layer]
|
||||
.copy_from_host_at(&self.cpu_k_pools[layer].as_slice()[c_off..c_off + bb], g_off, bb)
|
||||
.unwrap();
|
||||
self.v_pools[layer]
|
||||
.copy_from_host_at(&self.cpu_v_pools[layer].as_slice()[c_off..c_off + bb], g_off, bb)
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
for b in cpu_ids {
|
||||
self.cpu_allocator.free(b);
|
||||
}
|
||||
let state = self.seq_states[slot].as_mut().unwrap();
|
||||
state.block_ids = gpu_ids;
|
||||
state.location = Location::Gpu;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn tensor_from_owned_buf(buf: GpuBuffer, shape: &[usize], dtype: DType, device: u32) -> Tensor {
|
||||
use smallvec::SmallVec;
|
||||
use xserv_tensor::shape::contiguous_strides;
|
||||
use xserv_tensor::storage::Storage;
|
||||
|
||||
let storage = Storage::cuda(buf, device);
|
||||
Tensor::from_storage(
|
||||
storage,
|
||||
SmallVec::from_slice(shape),
|
||||
contiguous_strides(shape),
|
||||
0,
|
||||
dtype,
|
||||
)
|
||||
}
|
||||
Reference in New Issue
Block a user