gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

style: format Rust workspace

Browse Source
This commit is contained in:
Gahow Wang
2026-06-18 18:11:58 +08:00
parent 013465fc06
commit 531cd3fe08
57 changed files with 4045 additions and 1204 deletions
Download Patch File Download Diff File Expand all files Collapse all files

248
crates/xserv-model/src/paged_kv_cache.rs
View File

@@ -29,7 +29,10 @@ impl BlockAllocator {
for b in (1..total_blocks).rev() {
free_stack.push(b as u32);
}
Self { free_stack, total: total_blocks }
Self {
free_stack,
total: total_blocks,
}
}
pub fn alloc(&mut self) -> Option<u32> {
@@ -136,8 +139,14 @@ impl PagedKVCache {
device: u32,
) -> Self {
Self::new_tp(
config, config.num_kv_heads(), total_blocks, cpu_total_blocks,
max_seqs, max_blocks_per_seq, dtype, device,
config,
config.num_kv_heads(),
total_blocks,
cpu_total_blocks,
max_seqs,
max_blocks_per_seq,
dtype,
device,
)
}
@@ -155,7 +164,10 @@ impl PagedKVCache {
dtype: DType,
device: u32,
) -> Self {
assert!(total_blocks >= 2, "need at least 2 blocks (one is sentinel)");
assert!(
total_blocks >= 2,
"need at least 2 blocks (one is sentinel)"
);
let num_layers = config.num_layers();
let head_dim = config.head_dim();
let elem_size = dtype.size_bytes();
@@ -179,11 +191,17 @@ impl PagedKVCache {
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"));
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 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>())
@@ -220,22 +238,49 @@ impl PagedKVCache {
}
}
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 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 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)
self.seq_states[slot]
.as_ref()
.map(|s| s.seq_len)
.unwrap_or(0)
}
pub fn is_slot_free(&self, slot: usize) -> bool {
@@ -280,7 +325,11 @@ impl PagedKVCache {
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 }
if needed_total > cur {
needed_total - cur
} else {
0
}
}
/// Pre-allocate enough physical blocks in `slot` to cover positions
@@ -290,8 +339,14 @@ impl PagedKVCache {
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");
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);
}
}
@@ -318,7 +373,9 @@ impl PagedKVCache {
num_tokens: usize,
start_pos: usize,
) {
if num_tokens == 0 { return; }
if num_tokens == 0 {
return;
}
// Make sure blocks exist for the target range.
self.ensure_capacity(slot, start_pos + num_tokens);
@@ -328,15 +385,21 @@ impl PagedKVCache {
// Stage block_ids on the GPU. Pool-allocated so this is essentially
// free after the first call (same bucket every step).
let block_ids: Vec<i32> = self.seq_states[slot].as_ref().unwrap()
.block_ids.iter().map(|&b| b as i32).collect();
let block_ids: Vec<i32> = self.seq_states[slot]
.as_ref()
.unwrap()
.block_ids
.iter()
.map(|&b| b as i32)
.collect();
let bytes = block_ids.len() * std::mem::size_of::<i32>();
let mut block_ids_gpu = xserv_cuda::allocator::cached_alloc(bytes)
.expect("alloc append block_ids");
let block_ids_bytes = unsafe {
std::slice::from_raw_parts(block_ids.as_ptr() as *const u8, bytes)
};
block_ids_gpu.copy_from_host(block_ids_bytes).expect("upload block_ids");
let mut block_ids_gpu =
xserv_cuda::allocator::cached_alloc(bytes).expect("alloc append block_ids");
let block_ids_bytes =
unsafe { std::slice::from_raw_parts(block_ids.as_ptr() as *const u8, bytes) };
block_ids_gpu
.copy_from_host(block_ids_bytes)
.expect("upload block_ids");
let k_src = k_new.data_ptr() as *const std::ffi::c_void;
let v_src = v_new.data_ptr() as *const std::ffi::c_void;
@@ -345,10 +408,16 @@ impl PagedKVCache {
unsafe {
xserv_kernels::reshape_and_cache_bf16(
k_src, v_src,
k_pool_ptr, v_pool_ptr,
k_src,
v_src,
k_pool_ptr,
v_pool_ptr,
block_ids_gpu.as_ptr() as *const i32,
num_tokens, nkv, hd, start_pos, bs,
num_tokens,
nkv,
hd,
start_pos,
bs,
xserv_cuda::current_stream_raw(),
);
}
@@ -378,7 +447,9 @@ impl PagedKVCache {
v_new: &Tensor,
batch: usize,
) {
if batch == 0 { return; }
if batch == 0 {
return;
}
let nkv = self.num_kv_heads;
let hd = self.head_dim;
debug_assert_eq!(k_new.shape(), &[batch, nkv, hd]);
@@ -393,10 +464,17 @@ impl PagedKVCache {
unsafe {
xserv_kernels::reshape_and_cache_batched_bf16(
k_src, v_src,
k_pool_ptr, v_pool_ptr,
bt_ptr, cl_ptr,
batch, nkv, hd, BLOCK_SIZE, self.max_blocks_per_seq,
k_src,
v_src,
k_pool_ptr,
v_pool_ptr,
bt_ptr,
cl_ptr,
batch,
nkv,
hd,
BLOCK_SIZE,
self.max_blocks_per_seq,
xserv_cuda::current_stream_raw(),
);
}
@@ -447,7 +525,10 @@ impl PagedKVCache {
/// 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");
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;
@@ -456,7 +537,9 @@ impl PagedKVCache {
*cl = 0;
}
for (i, &slot) in slots.iter().enumerate() {
let s = self.seq_states[slot].as_ref().expect("unregistered slot in active batch");
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;
@@ -515,8 +598,12 @@ impl PagedKVCache {
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();
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;
}
@@ -529,16 +616,26 @@ impl PagedKVCache {
// ----- 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 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))
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)
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).
@@ -554,11 +651,17 @@ impl PagedKVCache {
/// 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 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"); }
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"))
@@ -570,10 +673,18 @@ impl PagedKVCache {
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)
.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)
.copy_to_host_at(
&mut self.cpu_v_pools[layer].as_mut_slice()[c_off..c_off + bb],
g_off,
bb,
)
.unwrap();
}
}
@@ -589,11 +700,17 @@ impl PagedKVCache {
/// 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 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"); }
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"))
@@ -605,10 +722,18 @@ impl PagedKVCache {
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)
.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)
.copy_from_host_at(
&self.cpu_v_pools[layer].as_slice()[c_off..c_off + bb],
g_off,
bb,
)
.unwrap();
}
}
@@ -623,7 +748,12 @@ impl PagedKVCache {
}
}
unsafe fn tensor_from_owned_buf(buf: GpuBuffer, shape: &[usize], dtype: DType, device: u32) -> Tensor {
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;