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

model: paged KV cache with CPU swap pool, decode graph, qwen3 updates

Browse Source 
- 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:
Gahow Wang
2026-05-28 19:58:54 +08:00
parent 4c3f914459
commit d52baa0006
9 changed files with 1896 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

101
crates/xserv-model/src/gpt2.rs
View File

@@ -280,45 +280,88 @@ fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
fn split_qkv(qkv: &Tensor, num_heads: usize, head_dim: usize, seq_len: usize) -> (Tensor, Tensor, Tensor) {
let hidden = num_heads * head_dim;
let qkv_cpu = qkv.to_device(Device::Cpu);
let data = qkv_cpu.as_slice::<f32>();
let mut q_data = vec![0.0f32; num_heads * seq_len * head_dim];
let mut k_data = vec![0.0f32; num_heads * seq_len * head_dim];
let mut v_data = vec![0.0f32; num_heads * seq_len * head_dim];
for s in 0..seq_len {
let row = &data[s * 3 * hidden..(s + 1) * 3 * hidden];
for h in 0..num_heads {
let src_off = h * head_dim;
let dst_off = (h * seq_len + s) * head_dim;
q_data[dst_off..dst_off + head_dim].copy_from_slice(&row[src_off..src_off + head_dim]);
k_data[dst_off..dst_off + head_dim].copy_from_slice(&row[hidden + src_off..hidden + src_off + head_dim]);
v_data[dst_off..dst_off + head_dim].copy_from_slice(&row[2 * hidden + src_off..2 * hidden + src_off + head_dim]);
}
}
let device = qkv.device();
let q = Tensor::from_slice(&q_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
let k = Tensor::from_slice(&k_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
let v = Tensor::from_slice(&v_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
(q, k, v)
let dtype = qkv.dtype();
match dtype {
DType::F32 => {
let data = qkv_cpu.as_slice::<f32>();
let mut q_data = vec![0.0f32; num_heads * seq_len * head_dim];
let mut k_data = vec![0.0f32; num_heads * seq_len * head_dim];
let mut v_data = vec![0.0f32; num_heads * seq_len * head_dim];
for s in 0..seq_len {
let row = &data[s * 3 * hidden..(s + 1) * 3 * hidden];
for h in 0..num_heads {
let src_off = h * head_dim;
let dst_off = (h * seq_len + s) * head_dim;
q_data[dst_off..dst_off + head_dim].copy_from_slice(&row[src_off..src_off + head_dim]);
k_data[dst_off..dst_off + head_dim].copy_from_slice(&row[hidden + src_off..hidden + src_off + head_dim]);
v_data[dst_off..dst_off + head_dim].copy_from_slice(&row[2 * hidden + src_off..2 * hidden + src_off + head_dim]);
}
}
let q = Tensor::from_slice(&q_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
let k = Tensor::from_slice(&k_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
let v = Tensor::from_slice(&v_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
(q, k, v)
}
DType::BF16 => {
let data = qkv_cpu.as_slice::<half::bf16>();
let mut q_data = vec![half::bf16::ZERO; num_heads * seq_len * head_dim];
let mut k_data = vec![half::bf16::ZERO; num_heads * seq_len * head_dim];
let mut v_data = vec![half::bf16::ZERO; num_heads * seq_len * head_dim];
for s in 0..seq_len {
let row = &data[s * 3 * hidden..(s + 1) * 3 * hidden];
for h in 0..num_heads {
let src_off = h * head_dim;
let dst_off = (h * seq_len + s) * head_dim;
q_data[dst_off..dst_off + head_dim].copy_from_slice(&row[src_off..src_off + head_dim]);
k_data[dst_off..dst_off + head_dim].copy_from_slice(&row[hidden + src_off..hidden + src_off + head_dim]);
v_data[dst_off..dst_off + head_dim].copy_from_slice(&row[2 * hidden + src_off..2 * hidden + src_off + head_dim]);
}
}
let q = Tensor::from_slice(&q_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
let k = Tensor::from_slice(&k_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
let v = Tensor::from_slice(&v_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
(q, k, v)
}
_ => panic!("unsupported dtype {:?} in split_qkv", dtype),
}
}
fn merge_heads(x: &Tensor, seq_len: usize, hidden: usize) -> Tensor {
let num_heads = x.shape()[1];
let head_dim = x.shape()[3];
let x_cpu = x.to_device(Device::Cpu);
let src = x_cpu.as_slice::<f32>();
let device = x.device();
let dtype = x.dtype();
let mut out = vec![0.0f32; seq_len * hidden];
for s in 0..seq_len {
for h in 0..num_heads {
let src_off = (h * seq_len + s) * head_dim;
let dst_off = s * hidden + h * head_dim;
out[dst_off..dst_off + head_dim].copy_from_slice(&src[src_off..src_off + head_dim]);
match dtype {
DType::F32 => {
let src = x_cpu.as_slice::<f32>();
let mut out = vec![0.0f32; seq_len * hidden];
for s in 0..seq_len {
for h in 0..num_heads {
let src_off = (h * seq_len + s) * head_dim;
let dst_off = s * hidden + h * head_dim;
out[dst_off..dst_off + head_dim].copy_from_slice(&src[src_off..src_off + head_dim]);
}
}
Tensor::from_slice(&out, &[seq_len, hidden]).to_device(device)
}
DType::BF16 => {
let src = x_cpu.as_slice::<half::bf16>();
let mut out = vec![half::bf16::ZERO; seq_len * hidden];
for s in 0..seq_len {
for h in 0..num_heads {
let src_off = (h * seq_len + s) * head_dim;
let dst_off = s * hidden + h * head_dim;
out[dst_off..dst_off + head_dim].copy_from_slice(&src[src_off..src_off + head_dim]);
}
}
Tensor::from_slice(&out, &[seq_len, hidden]).to_device(device)
}
_ => panic!("unsupported dtype {:?} in merge_heads", dtype),
}
Tensor::from_slice(&out, &[seq_len, hidden]).to_device(x.device())
}
/// Greedy sampling: return the argmax token ID from the last position's logits.