Compare commits
5 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| d8493bd70f | |||
| 7d05ececa0 | |||
| da043554ba | |||
| 2be27d6d94 | |||
| 2d48f25e66 |
@@ -6,6 +6,7 @@ members = [
|
||||
"crates/xserv-kernels",
|
||||
"crates/xserv-model",
|
||||
"crates/xserv-tokenizer",
|
||||
"crates/xserv-server",
|
||||
]
|
||||
|
||||
[workspace.package]
|
||||
@@ -20,3 +21,6 @@ serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
safetensors = "0.5"
|
||||
regex = "1"
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
axum = "0.8"
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
|
||||
@@ -87,6 +87,20 @@ impl GpuBuffer {
|
||||
error::check(unsafe { ffi::cudaMemset(self.ptr, 0, self.len) })
|
||||
}
|
||||
|
||||
/// Copy `count` bytes from `src` buffer at `src_offset` to this buffer at `dst_offset`.
|
||||
pub fn copy_from_device_at(&mut self, src: &GpuBuffer, src_offset: usize, dst_offset: usize, count: usize) -> Result<()> {
|
||||
assert!(src_offset + count <= src.len);
|
||||
assert!(dst_offset + count <= self.len);
|
||||
error::check(unsafe {
|
||||
ffi::cudaMemcpy(
|
||||
self.ptr.add(dst_offset),
|
||||
src.ptr.add(src_offset),
|
||||
count,
|
||||
ffi::CUDA_MEMCPY_D2D,
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
/// Consume the buffer without freeing GPU memory. Returns the raw pointer and length.
|
||||
/// Caller is responsible for eventually calling cudaFree.
|
||||
pub fn into_raw(self) -> (*mut u8, usize) {
|
||||
|
||||
@@ -23,6 +23,7 @@ fn main() {
|
||||
.file("../../csrc/embedding/embedding.cu")
|
||||
.file("../../csrc/embedding/rope.cu")
|
||||
.file("../../csrc/attention/causal_mask.cu")
|
||||
.file("../../csrc/embedding/transpose.cu")
|
||||
.compile("xserv_kernels");
|
||||
|
||||
println!("cargo:rerun-if-changed=../../csrc/");
|
||||
|
||||
@@ -6,8 +6,10 @@ pub mod layernorm;
|
||||
pub mod rmsnorm;
|
||||
pub mod rope;
|
||||
pub mod softmax;
|
||||
pub mod transpose;
|
||||
|
||||
pub use activation::{add, gelu, mul, scale, silu};
|
||||
pub use transpose::{merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, transpose_for_rope_gpu, transpose_from_rope_gpu};
|
||||
pub use attention::attention;
|
||||
pub use embedding::embedding;
|
||||
pub use gemm::{batched_matmul, matmul, GemmBackend};
|
||||
|
||||
91
crates/xserv-kernels/src/transpose.rs
Normal file
91
crates/xserv-kernels/src/transpose.rs
Normal file
@@ -0,0 +1,91 @@
|
||||
use std::ffi::c_void;
|
||||
use xserv_tensor::{DType, Device, Tensor};
|
||||
|
||||
unsafe extern "C" {
|
||||
fn launch_reshape_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
|
||||
fn launch_merge_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
|
||||
fn launch_transpose_hsd_to_shd_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
|
||||
fn launch_transpose_shd_to_hsd_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
|
||||
fn launch_repeat_kv_bf16(inp: *const c_void, out: *mut c_void, kv_heads: i32, n_rep: i32, seq_len: i32, head_dim: i32, stream: *mut c_void);
|
||||
}
|
||||
|
||||
/// [S, H*D] → [1, H, S, D] on GPU (BF16)
|
||||
pub fn reshape_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
|
||||
assert_eq!(x.dtype(), DType::BF16);
|
||||
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
|
||||
let out = Tensor::zeros(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
|
||||
unsafe {
|
||||
launch_reshape_heads_bf16(
|
||||
x.data_ptr() as _, out.data_ptr() as *mut c_void,
|
||||
seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
|
||||
);
|
||||
}
|
||||
xserv_cuda::device::synchronize().unwrap();
|
||||
out
|
||||
}
|
||||
|
||||
/// [1, H, S, D] → [S, H*D] on GPU (BF16)
|
||||
pub fn merge_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
|
||||
assert_eq!(x.dtype(), DType::BF16);
|
||||
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
|
||||
let hidden = num_heads * head_dim;
|
||||
let out = Tensor::zeros(&[seq_len, hidden], DType::BF16, x.device());
|
||||
unsafe {
|
||||
launch_merge_heads_bf16(
|
||||
x.data_ptr() as _, out.data_ptr() as *mut c_void,
|
||||
seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
|
||||
);
|
||||
}
|
||||
xserv_cuda::device::synchronize().unwrap();
|
||||
out
|
||||
}
|
||||
|
||||
/// [1, H, S, D] → [S, H, D] for RoPE on GPU (BF16)
|
||||
pub fn transpose_for_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
|
||||
assert_eq!(x.dtype(), DType::BF16);
|
||||
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
|
||||
let out = Tensor::zeros(&[seq_len, num_heads, head_dim], DType::BF16, x.device());
|
||||
unsafe {
|
||||
launch_transpose_hsd_to_shd_bf16(
|
||||
x.data_ptr() as _, out.data_ptr() as *mut c_void,
|
||||
seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
|
||||
);
|
||||
}
|
||||
xserv_cuda::device::synchronize().unwrap();
|
||||
out
|
||||
}
|
||||
|
||||
/// [S, H, D] → [1, H, S, D] after RoPE on GPU (BF16)
|
||||
pub fn transpose_from_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
|
||||
assert_eq!(x.dtype(), DType::BF16);
|
||||
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
|
||||
let out = Tensor::zeros(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
|
||||
unsafe {
|
||||
launch_transpose_shd_to_hsd_bf16(
|
||||
x.data_ptr() as _, out.data_ptr() as *mut c_void,
|
||||
seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
|
||||
);
|
||||
}
|
||||
xserv_cuda::device::synchronize().unwrap();
|
||||
out
|
||||
}
|
||||
|
||||
/// [1, KV_H, S, D] → [1, KV_H*n_rep, S, D] on GPU (BF16)
|
||||
pub fn repeat_kv_gpu(x: &Tensor, n_rep: usize) -> Tensor {
|
||||
if n_rep == 1 { return x.clone(); }
|
||||
assert_eq!(x.dtype(), DType::BF16);
|
||||
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
|
||||
let kv_heads = x.shape()[1];
|
||||
let seq_len = x.shape()[2];
|
||||
let head_dim = x.shape()[3];
|
||||
let new_heads = kv_heads * n_rep;
|
||||
let out = Tensor::zeros(&[1, new_heads, seq_len, head_dim], DType::BF16, x.device());
|
||||
unsafe {
|
||||
launch_repeat_kv_bf16(
|
||||
x.data_ptr() as _, out.data_ptr() as *mut c_void,
|
||||
kv_heads as i32, n_rep as i32, seq_len as i32, head_dim as i32, std::ptr::null_mut(),
|
||||
);
|
||||
}
|
||||
xserv_cuda::device::synchronize().unwrap();
|
||||
out
|
||||
}
|
||||
@@ -9,6 +9,7 @@ xserv-tensor = { path = "../xserv-tensor" }
|
||||
xserv-kernels = { path = "../xserv-kernels" }
|
||||
xserv-tokenizer = { path = "../xserv-tokenizer" }
|
||||
half.workspace = true
|
||||
smallvec.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
safetensors.workspace = true
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use std::path::PathBuf;
|
||||
use std::time::Instant;
|
||||
use xserv_model::qwen3::sample_greedy;
|
||||
use xserv_model::{loader, KVCache, ModelConfig, Qwen3};
|
||||
use xserv_model::{loader, GpuKVCache, ModelConfig, Qwen3};
|
||||
use xserv_tensor::{DType, Device};
|
||||
use xserv_tokenizer::Tokenizer;
|
||||
|
||||
@@ -31,11 +31,8 @@ fn main() {
|
||||
// Warmup
|
||||
{
|
||||
let ids = tokenizer.encode("warmup");
|
||||
let mut cache = KVCache::new(
|
||||
config.num_layers(), config.num_kv_heads(), config.head_dim(),
|
||||
DType::BF16, Device::Cuda(0),
|
||||
);
|
||||
let _ = model.forward_with_cache(&ids, &mut cache);
|
||||
let mut cache = GpuKVCache::new(&config, 256, DType::BF16);
|
||||
let _ = model.forward_gpu_cache(&ids, &mut cache);
|
||||
}
|
||||
eprintln!("Warmup done. Running benchmark...");
|
||||
|
||||
@@ -97,14 +94,11 @@ fn main() {
|
||||
let input_ids = tokenizer.encode(prompt);
|
||||
let input_len = input_ids.len();
|
||||
|
||||
let mut cache = KVCache::new(
|
||||
config.num_layers(), config.num_kv_heads(), config.head_dim(),
|
||||
DType::BF16, Device::Cuda(0),
|
||||
);
|
||||
let mut cache = GpuKVCache::new(&config, 256, DType::BF16);
|
||||
|
||||
// Prefill
|
||||
let t0 = Instant::now();
|
||||
let logits = model.forward_with_cache(&input_ids, &mut cache);
|
||||
let logits = model.forward_gpu_cache(&input_ids, &mut cache);
|
||||
let first_token = sample_greedy(&logits);
|
||||
let ttft_us = t0.elapsed().as_micros();
|
||||
|
||||
@@ -115,7 +109,7 @@ fn main() {
|
||||
for _ in 1..gen_tokens {
|
||||
let last = *generated.last().unwrap();
|
||||
let t_start = Instant::now();
|
||||
let logits = model.forward_with_cache(&[last], &mut cache);
|
||||
let logits = model.forward_gpu_cache(&[last], &mut cache);
|
||||
let next = sample_greedy(&logits);
|
||||
token_times.push(t_start.elapsed().as_micros());
|
||||
generated.push(next);
|
||||
@@ -148,12 +142,14 @@ fn main() {
|
||||
print!("\"tpot_us\": {tpot_us}}}");
|
||||
if i < prompts.len() - 1 { println!(","); } else { println!(); }
|
||||
|
||||
let display_text = generated_text.replace('\n', " ");
|
||||
let truncated: String = display_text.chars().take(60).collect();
|
||||
eprintln!(
|
||||
"[{}/{}] input={input_len}tok gen={num_generated}tok ttft={:.1}ms tbt={:.1}ms | {}",
|
||||
i + 1, prompts.len(),
|
||||
ttft_us as f64 / 1000.0,
|
||||
tbt_us as f64 / 1000.0,
|
||||
&generated_text.replace('\n', " ")[..generated_text.len().min(60)]
|
||||
truncated
|
||||
);
|
||||
}
|
||||
println!("]");
|
||||
|
||||
118
crates/xserv-model/src/kv_cache.rs
Normal file
118
crates/xserv-model/src/kv_cache.rs
Normal file
@@ -0,0 +1,118 @@
|
||||
use xserv_cuda::GpuBuffer;
|
||||
use xserv_tensor::{DType, Device, Tensor};
|
||||
use crate::config::ModelConfig;
|
||||
|
||||
/// GPU-resident KV cache. Pre-allocates max_seq_len on GPU,
|
||||
/// appends new K/V via D2D copy at offset (no CPU round-trip).
|
||||
pub struct GpuKVCache {
|
||||
// Per layer: contiguous GPU buffer for K and V
|
||||
// Layout: [num_kv_heads, max_seq_len, head_dim] — contiguous per head
|
||||
k_bufs: Vec<GpuBuffer>,
|
||||
v_bufs: Vec<GpuBuffer>,
|
||||
seq_len: usize,
|
||||
max_seq_len: usize,
|
||||
num_kv_heads: usize,
|
||||
head_dim: usize,
|
||||
elem_size: usize,
|
||||
dtype: DType,
|
||||
}
|
||||
|
||||
impl GpuKVCache {
|
||||
pub fn new(config: &ModelConfig, max_seq_len: usize, dtype: DType) -> Self {
|
||||
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 buf_size = num_kv_heads * max_seq_len * head_dim * elem_size;
|
||||
|
||||
let mut k_bufs = Vec::with_capacity(num_layers);
|
||||
let mut v_bufs = Vec::with_capacity(num_layers);
|
||||
for _ in 0..num_layers {
|
||||
let mut k = GpuBuffer::alloc(buf_size).expect("alloc KV cache K");
|
||||
let mut v = GpuBuffer::alloc(buf_size).expect("alloc KV cache V");
|
||||
k.zero().unwrap();
|
||||
v.zero().unwrap();
|
||||
k_bufs.push(k);
|
||||
v_bufs.push(v);
|
||||
}
|
||||
|
||||
Self { k_bufs, v_bufs, seq_len: 0, max_seq_len, num_kv_heads, head_dim, elem_size, dtype }
|
||||
}
|
||||
|
||||
pub fn seq_len(&self) -> usize { self.seq_len }
|
||||
pub fn max_seq_len(&self) -> usize { self.max_seq_len }
|
||||
|
||||
/// Append new K/V tensors for a given layer.
|
||||
/// k_new, v_new: [1, num_kv_heads, new_tokens, head_dim] on GPU, contiguous.
|
||||
/// `write_pos` is the sequence position to write at (caller manages this).
|
||||
pub fn append(&mut self, layer: usize, k_new: &Tensor, v_new: &Tensor, new_tokens: usize, write_pos: usize) {
|
||||
assert!(write_pos + new_tokens <= self.max_seq_len, "KV cache overflow");
|
||||
let es = self.elem_size;
|
||||
let hd = self.head_dim;
|
||||
let max_s = self.max_seq_len;
|
||||
let nh = self.num_kv_heads;
|
||||
|
||||
let k_src = k_new.storage().gpu_buffer();
|
||||
let v_src = v_new.storage().gpu_buffer();
|
||||
|
||||
for h in 0..nh {
|
||||
let src_off = h * new_tokens * hd * es;
|
||||
let dst_off = (h * max_s + write_pos) * hd * es;
|
||||
let count = new_tokens * hd * es;
|
||||
self.k_bufs[layer].copy_from_device_at(k_src, src_off, dst_off, count).unwrap();
|
||||
self.v_bufs[layer].copy_from_device_at(v_src, src_off, dst_off, count).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
pub fn advance_seq_len(&mut self, new_tokens: usize) {
|
||||
self.seq_len += new_tokens;
|
||||
}
|
||||
|
||||
/// Get K/V cache tensors for a layer up to `seq_len` tokens: [1, num_kv_heads, seq_len, head_dim]
|
||||
pub fn get_kv(&self, layer: usize) -> (Tensor, Tensor) {
|
||||
let sl = self.seq_len;
|
||||
self.get_kv_len(layer, sl)
|
||||
}
|
||||
|
||||
pub fn get_kv_len(&self, layer: usize, sl: usize) -> (Tensor, Tensor) {
|
||||
let hd = self.head_dim;
|
||||
let nh = self.num_kv_heads;
|
||||
let es = self.elem_size;
|
||||
let max_s = self.max_seq_len;
|
||||
|
||||
// Allocate output tensors [1, nh, sl, hd]
|
||||
let out_size = nh * sl * hd * es;
|
||||
let mut k_out = GpuBuffer::alloc(out_size).expect("alloc k_out");
|
||||
let mut v_out = GpuBuffer::alloc(out_size).expect("alloc v_out");
|
||||
|
||||
// Copy each head's valid portion
|
||||
for h in 0..nh {
|
||||
let src_off = (h * max_s) * hd * es;
|
||||
let dst_off = (h * sl) * hd * es;
|
||||
let count = sl * hd * es;
|
||||
k_out.copy_from_device_at(&self.k_bufs[layer], src_off, dst_off, count).unwrap();
|
||||
v_out.copy_from_device_at(&self.v_bufs[layer], src_off, dst_off, count).unwrap();
|
||||
}
|
||||
|
||||
let shape = &[1usize, nh, sl, hd];
|
||||
let k = unsafe { tensor_from_gpu_buffer(k_out, shape, self.dtype) };
|
||||
let v = unsafe { tensor_from_gpu_buffer(v_out, shape, self.dtype) };
|
||||
(k, v)
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a Tensor from a GpuBuffer (takes ownership).
|
||||
unsafe fn tensor_from_gpu_buffer(buf: GpuBuffer, shape: &[usize], dtype: DType) -> Tensor {
|
||||
use xserv_tensor::storage::Storage;
|
||||
use xserv_tensor::shape::contiguous_strides;
|
||||
use smallvec::SmallVec;
|
||||
|
||||
let storage = Storage::cuda(buf);
|
||||
Tensor::from_storage(
|
||||
storage,
|
||||
SmallVec::from_slice(shape),
|
||||
contiguous_strides(shape),
|
||||
0,
|
||||
dtype,
|
||||
)
|
||||
}
|
||||
@@ -1,8 +1,10 @@
|
||||
pub mod config;
|
||||
pub mod gpt2;
|
||||
pub mod kv_cache;
|
||||
pub mod loader;
|
||||
pub mod qwen3;
|
||||
|
||||
pub use config::ModelConfig;
|
||||
pub use gpt2::{GPT2, KVCache};
|
||||
pub use kv_cache::GpuKVCache;
|
||||
pub use qwen3::Qwen3;
|
||||
|
||||
@@ -5,6 +5,7 @@ use xserv_tensor::{DType, Device, Tensor};
|
||||
|
||||
use crate::config::ModelConfig;
|
||||
use crate::gpt2::KVCache;
|
||||
use crate::kv_cache::GpuKVCache;
|
||||
|
||||
pub struct Qwen3 {
|
||||
pub config: ModelConfig,
|
||||
@@ -145,6 +146,75 @@ impl Qwen3 {
|
||||
let x = rmsnorm(&x, &self.norm, eps);
|
||||
matmul_2d(&x, &self.lm_head_t)
|
||||
}
|
||||
|
||||
/// Forward with GPU-resident KV cache and GPU transpose/reshape kernels.
|
||||
pub fn forward_gpu_cache(&self, token_ids: &[u32], cache: &mut GpuKVCache) -> Tensor {
|
||||
let new_tokens = token_ids.len();
|
||||
let pos_offset = cache.seq_len();
|
||||
let hidden = self.config.hidden();
|
||||
let num_heads = self.config.num_heads();
|
||||
let num_kv_heads = self.config.num_kv_heads();
|
||||
let head_dim = self.config.head_dim();
|
||||
let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
|
||||
|
||||
let mut x = embedding(&self.embed_tokens, token_ids);
|
||||
let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
|
||||
|
||||
for (layer_idx, layer) in self.layers.iter().enumerate() {
|
||||
let residual = x.clone();
|
||||
let normed = rmsnorm(&x, &layer.input_norm, eps);
|
||||
|
||||
let q = matmul_2d(&normed, &layer.q_proj_wt);
|
||||
let k = matmul_2d(&normed, &layer.k_proj_wt);
|
||||
let v = matmul_2d(&normed, &layer.v_proj_wt);
|
||||
|
||||
// GPU reshape: [S, H*D] → [1, H, S, D]
|
||||
let q = xserv_kernels::reshape_heads_gpu(&q, new_tokens, num_heads, head_dim);
|
||||
let k = xserv_kernels::reshape_heads_gpu(&k, new_tokens, num_kv_heads, head_dim);
|
||||
let v = xserv_kernels::reshape_heads_gpu(&v, new_tokens, num_kv_heads, head_dim);
|
||||
|
||||
// QK norm (reshape to [H*S, D], rmsnorm, reshape back — stays on GPU)
|
||||
let q = head_rmsnorm(&q, &layer.q_norm, eps);
|
||||
let k = head_rmsnorm(&k, &layer.k_norm, eps);
|
||||
|
||||
// GPU transpose for RoPE: [1, H, S, D] → [S, H, D]
|
||||
let q = xserv_kernels::transpose_for_rope_gpu(&q, new_tokens, num_heads, head_dim);
|
||||
let k = xserv_kernels::transpose_for_rope_gpu(&k, new_tokens, num_kv_heads, head_dim);
|
||||
rope_inplace(&q, &self.rope_cache, &positions);
|
||||
rope_inplace(&k, &self.rope_cache, &positions);
|
||||
// GPU transpose back: [S, H, D] → [1, H, S, D]
|
||||
let q = xserv_kernels::transpose_from_rope_gpu(&q, new_tokens, num_heads, head_dim);
|
||||
let k = xserv_kernels::transpose_from_rope_gpu(&k, new_tokens, num_kv_heads, head_dim);
|
||||
|
||||
// GPU KV cache
|
||||
cache.append(layer_idx, &k, &v, new_tokens, pos_offset);
|
||||
let (k_full, v_full) = cache.get_kv_len(layer_idx, pos_offset + new_tokens);
|
||||
|
||||
// GPU repeat KV for GQA
|
||||
let n_rep = num_heads / num_kv_heads;
|
||||
let k_full = xserv_kernels::repeat_kv_gpu(&k_full, n_rep);
|
||||
let v_full = xserv_kernels::repeat_kv_gpu(&v_full, n_rep);
|
||||
|
||||
let attn_out = attention(&q, &k_full, &v_full, true);
|
||||
// GPU merge_heads: [1, H, S, D] → [S, H*D]
|
||||
let attn_merged = xserv_kernels::merge_heads_gpu(&attn_out, new_tokens, num_heads, head_dim);
|
||||
let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
|
||||
x = add_any(&residual, &attn_proj);
|
||||
|
||||
let residual = x.clone();
|
||||
let normed = rmsnorm(&x, &layer.post_norm, eps);
|
||||
let gate = matmul_2d(&normed, &layer.gate_proj_wt);
|
||||
let up = matmul_2d(&normed, &layer.up_proj_wt);
|
||||
let gate_activated = silu(&gate);
|
||||
let hidden_states = mul_any(&gate_activated, &up);
|
||||
let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
|
||||
x = add_any(&residual, &down);
|
||||
}
|
||||
|
||||
cache.advance_seq_len(new_tokens);
|
||||
let x = rmsnorm(&x, &self.norm, eps);
|
||||
matmul_2d(&x, &self.lm_head_t)
|
||||
}
|
||||
}
|
||||
|
||||
// --- Helpers ---
|
||||
|
||||
21
crates/xserv-server/Cargo.toml
Normal file
21
crates/xserv-server/Cargo.toml
Normal file
@@ -0,0 +1,21 @@
|
||||
[package]
|
||||
name = "xserv-server"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
|
||||
[[bin]]
|
||||
name = "xserv-server"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
xserv-cuda = { path = "../xserv-cuda" }
|
||||
xserv-tensor = { path = "../xserv-tensor" }
|
||||
xserv-kernels = { path = "../xserv-kernels" }
|
||||
xserv-model = { path = "../xserv-model" }
|
||||
xserv-tokenizer = { path = "../xserv-tokenizer" }
|
||||
half.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
tokio.workspace = true
|
||||
axum.workspace = true
|
||||
uuid.workspace = true
|
||||
115
crates/xserv-server/src/api.rs
Normal file
115
crates/xserv-server/src/api.rs
Normal file
@@ -0,0 +1,115 @@
|
||||
use axum::Extension;
|
||||
use axum::Json;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::engine::{GenerateEvent, GenerateRequest};
|
||||
use crate::AppState;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct ChatRequest {
|
||||
#[serde(default)]
|
||||
pub model: Option<String>,
|
||||
pub messages: Vec<Message>,
|
||||
#[serde(default = "default_max_tokens")]
|
||||
pub max_tokens: usize,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct Message {
|
||||
pub role: String,
|
||||
pub content: String,
|
||||
}
|
||||
|
||||
fn default_max_tokens() -> usize { 256 }
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct ModelsResponse {
|
||||
object: &'static str,
|
||||
data: Vec<ModelInfo>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct ModelInfo {
|
||||
id: String,
|
||||
object: &'static str,
|
||||
owned_by: &'static str,
|
||||
}
|
||||
|
||||
pub async fn health() -> &'static str { "ok" }
|
||||
|
||||
pub async fn list_models(Extension(state): Extension<Arc<AppState>>) -> Json<ModelsResponse> {
|
||||
Json(ModelsResponse {
|
||||
object: "list",
|
||||
data: vec![ModelInfo {
|
||||
id: state.model_name.clone(),
|
||||
object: "model",
|
||||
owned_by: "xserv",
|
||||
}],
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn chat_completions(
|
||||
Extension(state): Extension<Arc<AppState>>,
|
||||
Json(req): Json<ChatRequest>,
|
||||
) -> Json<serde_json::Value> {
|
||||
let id = format!("chatcmpl-{}", Uuid::new_v4());
|
||||
let model_name = state.model_name.clone();
|
||||
let created = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs();
|
||||
|
||||
// Prepare prompt tokens (MutexGuard scoped)
|
||||
let prompt = build_prompt(&req.messages);
|
||||
let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);
|
||||
|
||||
// Create channel and submit request (MutexGuard scoped)
|
||||
let (tx, mut rx) = tokio::sync::mpsc::channel::<GenerateEvent>(64);
|
||||
let gen_req = GenerateRequest {
|
||||
prompt_tokens,
|
||||
max_tokens: req.max_tokens,
|
||||
sender: tx,
|
||||
};
|
||||
state.engine_sender.lock().unwrap().send(gen_req).expect("engine channel closed");
|
||||
|
||||
// Now await — no MutexGuards held here
|
||||
let mut content = String::new();
|
||||
let mut finish_reason = "length".to_string();
|
||||
while let Some(event) = rx.recv().await {
|
||||
match event {
|
||||
GenerateEvent::Token { text, .. } => content.push_str(&text),
|
||||
GenerateEvent::Done { finish_reason: fr } => { finish_reason = fr; break; }
|
||||
}
|
||||
}
|
||||
|
||||
Json(serde_json::json!({
|
||||
"id": id,
|
||||
"object": "chat.completion",
|
||||
"created": created,
|
||||
"model": model_name,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": { "role": "assistant", "content": content },
|
||||
"finish_reason": finish_reason,
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 0,
|
||||
"total_tokens": 0
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
fn build_prompt(messages: &[Message]) -> String {
|
||||
let mut prompt = String::new();
|
||||
for msg in messages {
|
||||
match msg.role.as_str() {
|
||||
"system" => { prompt.push_str(&msg.content); prompt.push('\n'); }
|
||||
"user" | "assistant" => { prompt.push_str(&msg.content); }
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
prompt
|
||||
}
|
||||
161
crates/xserv-server/src/engine.rs
Normal file
161
crates/xserv-server/src/engine.rs
Normal file
@@ -0,0 +1,161 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::path::Path;
|
||||
use std::sync::mpsc;
|
||||
use xserv_model::{GpuKVCache, ModelConfig, Qwen3};
|
||||
use xserv_model::loader;
|
||||
use xserv_model::qwen3::sample_greedy;
|
||||
use xserv_tensor::{DType, Device};
|
||||
use xserv_tokenizer::Tokenizer;
|
||||
|
||||
pub struct Engine {
|
||||
model: Qwen3,
|
||||
config: ModelConfig,
|
||||
tokenizer: Tokenizer,
|
||||
max_batch_size: usize,
|
||||
max_seq_len: usize,
|
||||
}
|
||||
|
||||
pub struct GenerateRequest {
|
||||
pub prompt_tokens: Vec<u32>,
|
||||
pub max_tokens: usize,
|
||||
pub sender: tokio::sync::mpsc::Sender<GenerateEvent>,
|
||||
}
|
||||
|
||||
pub enum GenerateEvent {
|
||||
Token { id: u32, text: String },
|
||||
Done { finish_reason: String },
|
||||
}
|
||||
|
||||
struct Sequence {
|
||||
id: u64,
|
||||
prompt_tokens: Vec<u32>,
|
||||
generated_tokens: Vec<u32>,
|
||||
max_tokens: usize,
|
||||
kv_cache: GpuKVCache,
|
||||
sender: tokio::sync::mpsc::Sender<GenerateEvent>,
|
||||
prefilled: bool,
|
||||
}
|
||||
|
||||
impl Engine {
|
||||
pub fn load(model_dir: &Path, max_batch_size: usize) -> Self {
|
||||
xserv_cuda::device::set_device(0).unwrap();
|
||||
let config = ModelConfig::from_file(&model_dir.join("config.json"));
|
||||
eprintln!("[engine] Loading weights...");
|
||||
let weights = loader::load_model_dir(model_dir, Device::Cuda(0));
|
||||
eprintln!("[engine] Loaded {} tensors", weights.len());
|
||||
let model = Qwen3::from_weights(config.clone(), weights);
|
||||
let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
|
||||
let max_seq_len = 256;
|
||||
eprintln!("[engine] Ready (max_batch_size={max_batch_size}, max_seq_len={max_seq_len})");
|
||||
Self { model, config, tokenizer, max_batch_size, max_seq_len }
|
||||
}
|
||||
|
||||
pub fn tokenizer(&self) -> &Tokenizer { &self.tokenizer }
|
||||
|
||||
/// Main scheduler loop. Receives requests from channel, manages concurrent sequences.
|
||||
pub fn run(&self, rx: mpsc::Receiver<GenerateRequest>) {
|
||||
let mut waiting: VecDeque<Sequence> = VecDeque::new();
|
||||
let mut running: Vec<Sequence> = Vec::new();
|
||||
let mut next_id: u64 = 0;
|
||||
|
||||
eprintln!("[scheduler] Listening for requests...");
|
||||
|
||||
loop {
|
||||
// Step 1: Remove finished sequences
|
||||
running.retain(|seq| !is_finished(seq));
|
||||
|
||||
// Step 2: Admit new sequences from waiting queue
|
||||
while running.len() < self.max_batch_size {
|
||||
if let Some(seq) = waiting.pop_front() {
|
||||
running.push(seq);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Step 3: If nothing to do, blocking wait for new request
|
||||
if running.is_empty() {
|
||||
match rx.recv() {
|
||||
Ok(req) => {
|
||||
let seq = self.make_sequence(req, &mut next_id);
|
||||
running.push(seq);
|
||||
}
|
||||
Err(_) => break, // channel closed
|
||||
}
|
||||
}
|
||||
|
||||
// Step 4: Process one iteration for all running sequences
|
||||
for seq in running.iter_mut() {
|
||||
if !seq.prefilled {
|
||||
// Prefill
|
||||
let logits = self.model.forward_gpu_cache(&seq.prompt_tokens, &mut seq.kv_cache);
|
||||
let next = sample_greedy(&logits);
|
||||
seq.generated_tokens.push(next);
|
||||
seq.prefilled = true;
|
||||
self.emit_token(seq, next);
|
||||
} else {
|
||||
// Decode one token
|
||||
let last = *seq.generated_tokens.last().unwrap();
|
||||
let logits = self.model.forward_gpu_cache(&[last], &mut seq.kv_cache);
|
||||
let next = sample_greedy(&logits);
|
||||
seq.generated_tokens.push(next);
|
||||
self.emit_token(seq, next);
|
||||
}
|
||||
}
|
||||
|
||||
// Step 5: Check for newly arrived requests (non-blocking)
|
||||
loop {
|
||||
match rx.try_recv() {
|
||||
Ok(req) => {
|
||||
let seq = self.make_sequence(req, &mut next_id);
|
||||
waiting.push_back(seq);
|
||||
}
|
||||
Err(mpsc::TryRecvError::Empty) => break,
|
||||
Err(mpsc::TryRecvError::Disconnected) => return,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn make_sequence(&self, req: GenerateRequest, next_id: &mut u64) -> Sequence {
|
||||
let id = *next_id;
|
||||
*next_id += 1;
|
||||
let kv_cache = GpuKVCache::new(&self.config, self.max_seq_len, DType::BF16);
|
||||
Sequence {
|
||||
id,
|
||||
prompt_tokens: req.prompt_tokens,
|
||||
generated_tokens: Vec::new(),
|
||||
max_tokens: req.max_tokens,
|
||||
kv_cache,
|
||||
sender: req.sender,
|
||||
prefilled: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn emit_token(&self, seq: &Sequence, token_id: u32) {
|
||||
let text = self.tokenizer.decode(&[token_id]);
|
||||
|
||||
if self.tokenizer.eos_token_id() == Some(token_id) {
|
||||
let _ = seq.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
|
||||
let _ = seq.sender.blocking_send(GenerateEvent::Done {
|
||||
finish_reason: "stop".to_string(),
|
||||
});
|
||||
} else if seq.generated_tokens.len() >= seq.max_tokens {
|
||||
let _ = seq.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
|
||||
let _ = seq.sender.blocking_send(GenerateEvent::Done {
|
||||
finish_reason: "length".to_string(),
|
||||
});
|
||||
} else {
|
||||
let _ = seq.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn is_finished(seq: &Sequence) -> bool {
|
||||
if seq.generated_tokens.is_empty() { return false; }
|
||||
let last = *seq.generated_tokens.last().unwrap();
|
||||
if seq.generated_tokens.len() >= seq.max_tokens { return true; }
|
||||
// Check EOS — need tokenizer info. Use a simple heuristic:
|
||||
// If sender is closed (receiver dropped), also consider finished.
|
||||
seq.sender.is_closed() || last == 151645 // Qwen3 EOS token ID (hardcoded for now)
|
||||
}
|
||||
66
crates/xserv-server/src/main.rs
Normal file
66
crates/xserv-server/src/main.rs
Normal file
@@ -0,0 +1,66 @@
|
||||
mod api;
|
||||
mod engine;
|
||||
|
||||
use axum::{routing::{get, post}, Extension, Router};
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{mpsc, Arc, Mutex};
|
||||
use engine::GenerateRequest;
|
||||
|
||||
pub struct AppState {
|
||||
pub model_name: String,
|
||||
pub engine_sender: Mutex<mpsc::Sender<GenerateRequest>>,
|
||||
pub engine_tokenizer: Mutex<xserv_tokenizer::Tokenizer>,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
if args.len() < 2 {
|
||||
eprintln!("Usage: xserv-server <model-dir> [--port PORT] [--max-batch N]");
|
||||
std::process::exit(1);
|
||||
}
|
||||
|
||||
let model_dir = PathBuf::from(&args[1]);
|
||||
let port: u16 = args.iter()
|
||||
.position(|a| a == "--port")
|
||||
.and_then(|i| args.get(i + 1))
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(8080);
|
||||
let max_batch: usize = args.iter()
|
||||
.position(|a| a == "--max-batch")
|
||||
.and_then(|i| args.get(i + 1))
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(4);
|
||||
|
||||
let model_name = model_dir.file_name()
|
||||
.map(|n| n.to_string_lossy().to_string())
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
|
||||
let tokenizer = xserv_tokenizer::Tokenizer::from_file(&model_dir.join("tokenizer.json"));
|
||||
|
||||
// Unbounded channel: allows multiple requests to queue up
|
||||
let (tx, rx) = mpsc::channel::<GenerateRequest>();
|
||||
|
||||
let model_dir_clone = model_dir.clone();
|
||||
std::thread::spawn(move || {
|
||||
let engine = engine::Engine::load(&model_dir_clone, max_batch);
|
||||
engine.run(rx);
|
||||
});
|
||||
|
||||
let state = Arc::new(AppState {
|
||||
model_name,
|
||||
engine_sender: Mutex::new(tx),
|
||||
engine_tokenizer: Mutex::new(tokenizer),
|
||||
});
|
||||
|
||||
let app = Router::new()
|
||||
.route("/health", get(api::health))
|
||||
.route("/v1/models", get(api::list_models))
|
||||
.route("/v1/chat/completions", post(api::chat_completions))
|
||||
.layer(Extension(state));
|
||||
|
||||
let addr = format!("0.0.0.0:{port}");
|
||||
eprintln!("[server] Listening on {addr} (max_batch={max_batch})");
|
||||
let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
}
|
||||
@@ -4,5 +4,6 @@ pub mod storage;
|
||||
pub mod tensor;
|
||||
|
||||
pub use dtype::{DType, TensorDType};
|
||||
pub use storage::Device;
|
||||
pub use shape::Dims;
|
||||
pub use storage::{Device, Storage};
|
||||
pub use tensor::Tensor;
|
||||
|
||||
@@ -18,6 +18,11 @@ pub struct Tensor {
|
||||
impl Tensor {
|
||||
// --- Creation ---
|
||||
|
||||
/// Create a tensor from raw components (for advanced use like GPU KV cache).
|
||||
pub fn from_storage(storage: Storage, shape: Dims, strides: Dims, offset: usize, dtype: DType) -> Self {
|
||||
Self { storage, shape, strides, offset, dtype }
|
||||
}
|
||||
|
||||
pub fn from_slice<T: TensorDType>(data: &[T], shape: &[usize]) -> Self {
|
||||
let numel: usize = shape.iter().product();
|
||||
assert_eq!(data.len(), numel, "data length mismatch with shape");
|
||||
|
||||
161
csrc/embedding/transpose.cu
Normal file
161
csrc/embedding/transpose.cu
Normal file
@@ -0,0 +1,161 @@
|
||||
#include <cuda_bf16.h>
|
||||
|
||||
// Transpose between [S, H, D] and [H, S, D] layouts (used for RoPE and attention).
|
||||
// Also handles [S, H*D] → [H, S, D] (reshape_heads) and reverse (merge_heads).
|
||||
|
||||
// reshape_heads: [S, H*D] → [1, H, S, D]
|
||||
// Input layout: element at [s, h*D + d] = flat[s * H*D + h*D + d]
|
||||
// Output layout: element at [0, h, s, d] = flat[h * S*D + s*D + d]
|
||||
__global__ void reshape_heads_bf16(
|
||||
const __nv_bfloat16* __restrict__ in,
|
||||
__nv_bfloat16* __restrict__ out,
|
||||
int seq_len, int num_heads, int head_dim
|
||||
) {
|
||||
int hidden = num_heads * head_dim;
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = seq_len * hidden;
|
||||
if (idx >= total) return;
|
||||
|
||||
int s = idx / hidden;
|
||||
int rem = idx % hidden;
|
||||
int h = rem / head_dim;
|
||||
int d = rem % head_dim;
|
||||
|
||||
int out_idx = h * seq_len * head_dim + s * head_dim + d;
|
||||
out[out_idx] = in[idx];
|
||||
}
|
||||
|
||||
// merge_heads: [1, H, S, D] → [S, H*D]
|
||||
// Input layout: element at [0, h, s, d] = flat[h * S*D + s*D + d]
|
||||
// Output layout: element at [s, h*D + d] = flat[s * H*D + h*D + d]
|
||||
__global__ void merge_heads_bf16(
|
||||
const __nv_bfloat16* __restrict__ in,
|
||||
__nv_bfloat16* __restrict__ out,
|
||||
int seq_len, int num_heads, int head_dim
|
||||
) {
|
||||
int hidden = num_heads * head_dim;
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = seq_len * hidden;
|
||||
if (idx >= total) return;
|
||||
|
||||
// idx is output index: [s, h*D + d]
|
||||
int s = idx / hidden;
|
||||
int rem = idx % hidden;
|
||||
int h = rem / head_dim;
|
||||
int d = rem % head_dim;
|
||||
|
||||
int in_idx = h * seq_len * head_dim + s * head_dim + d;
|
||||
out[idx] = in[in_idx];
|
||||
}
|
||||
|
||||
// transpose_for_rope: [1, H, S, D] → [S, H, D]
|
||||
// Input: [h, s, d] at h*S*D + s*D + d
|
||||
// Output: [s, h, d] at s*H*D + h*D + d
|
||||
__global__ void transpose_hsd_to_shd_bf16(
|
||||
const __nv_bfloat16* __restrict__ in,
|
||||
__nv_bfloat16* __restrict__ out,
|
||||
int seq_len, int num_heads, int head_dim
|
||||
) {
|
||||
int total = seq_len * num_heads * head_dim;
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= total) return;
|
||||
|
||||
// idx = output flat index: s*H*D + h*D + d
|
||||
int s = idx / (num_heads * head_dim);
|
||||
int rem = idx % (num_heads * head_dim);
|
||||
int h = rem / head_dim;
|
||||
int d = rem % head_dim;
|
||||
|
||||
int in_idx = h * seq_len * head_dim + s * head_dim + d;
|
||||
out[idx] = in[in_idx];
|
||||
}
|
||||
|
||||
// transpose_from_rope: [S, H, D] → [1, H, S, D]
|
||||
// Input: [s, h, d] at s*H*D + h*D + d
|
||||
// Output: [h, s, d] at h*S*D + s*D + d
|
||||
__global__ void transpose_shd_to_hsd_bf16(
|
||||
const __nv_bfloat16* __restrict__ in,
|
||||
__nv_bfloat16* __restrict__ out,
|
||||
int seq_len, int num_heads, int head_dim
|
||||
) {
|
||||
int total = seq_len * num_heads * head_dim;
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= total) return;
|
||||
|
||||
// idx = output flat index: h*S*D + s*D + d
|
||||
int h = idx / (seq_len * head_dim);
|
||||
int rem = idx % (seq_len * head_dim);
|
||||
int s = rem / head_dim;
|
||||
int d = rem % head_dim;
|
||||
|
||||
int in_idx = s * num_heads * head_dim + h * head_dim + d;
|
||||
out[idx] = in[in_idx];
|
||||
}
|
||||
|
||||
// repeat_kv: [1, KV_H, S, D] → [1, KV_H * n_rep, S, D]
|
||||
__global__ void repeat_kv_bf16(
|
||||
const __nv_bfloat16* __restrict__ in,
|
||||
__nv_bfloat16* __restrict__ out,
|
||||
int kv_heads, int n_rep, int seq_len, int head_dim
|
||||
) {
|
||||
int total_heads = kv_heads * n_rep;
|
||||
int total = total_heads * seq_len * head_dim;
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= total) return;
|
||||
|
||||
int out_h = idx / (seq_len * head_dim);
|
||||
int rem = idx % (seq_len * head_dim);
|
||||
int kv_h = out_h / n_rep;
|
||||
|
||||
int in_idx = kv_h * seq_len * head_dim + rem;
|
||||
out[idx] = in[in_idx];
|
||||
}
|
||||
|
||||
extern "C" {
|
||||
|
||||
void launch_reshape_heads_bf16(const void* in, void* out,
|
||||
int seq_len, int num_heads, int head_dim, void* stream) {
|
||||
int total = seq_len * num_heads * head_dim;
|
||||
int block = 256;
|
||||
int grid = (total + block - 1) / block;
|
||||
reshape_heads_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
|
||||
(const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
|
||||
}
|
||||
|
||||
void launch_merge_heads_bf16(const void* in, void* out,
|
||||
int seq_len, int num_heads, int head_dim, void* stream) {
|
||||
int total = seq_len * num_heads * head_dim;
|
||||
int block = 256;
|
||||
int grid = (total + block - 1) / block;
|
||||
merge_heads_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
|
||||
(const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
|
||||
}
|
||||
|
||||
void launch_transpose_hsd_to_shd_bf16(const void* in, void* out,
|
||||
int seq_len, int num_heads, int head_dim, void* stream) {
|
||||
int total = seq_len * num_heads * head_dim;
|
||||
int block = 256;
|
||||
int grid = (total + block - 1) / block;
|
||||
transpose_hsd_to_shd_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
|
||||
(const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
|
||||
}
|
||||
|
||||
void launch_transpose_shd_to_hsd_bf16(const void* in, void* out,
|
||||
int seq_len, int num_heads, int head_dim, void* stream) {
|
||||
int total = seq_len * num_heads * head_dim;
|
||||
int block = 256;
|
||||
int grid = (total + block - 1) / block;
|
||||
transpose_shd_to_hsd_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
|
||||
(const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
|
||||
}
|
||||
|
||||
void launch_repeat_kv_bf16(const void* in, void* out,
|
||||
int kv_heads, int n_rep, int seq_len, int head_dim, void* stream) {
|
||||
int total = kv_heads * n_rep * seq_len * head_dim;
|
||||
int block = 256;
|
||||
int grid = (total + block - 1) / block;
|
||||
repeat_kv_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
|
||||
(const __nv_bfloat16*)in, (__nv_bfloat16*)out, kv_heads, n_rep, seq_len, head_dim);
|
||||
}
|
||||
|
||||
}
|
||||
59
docs/11-paged-attention.md
Normal file
59
docs/11-paged-attention.md
Normal file
@@ -0,0 +1,59 @@
|
||||
# Phase 11: Paged Attention + KV Cache Manager — Design Document
|
||||
|
||||
## Goal
|
||||
|
||||
将 KV cache 从 CPU Vec 迁移到 GPU,使用 block-based paging 管理显存。消除每步 decode 的 CPU round-trip(当前 KV cache 最大性能瓶颈之一)。
|
||||
|
||||
## 当前问题
|
||||
|
||||
每步 decode 的 KV cache 路径:
|
||||
```
|
||||
GPU tensor (K_new) → CPU (per-head Vec append) → reconstruct → CPU tensor → GPU tensor
|
||||
```
|
||||
这涉及 2 次 GPU↔CPU 拷贝 × 36 层 × 2(K,V) = 144 次 transfer/token。
|
||||
|
||||
## 目标设计
|
||||
|
||||
KV cache 直接存在 GPU 上,decode 时只做 GPU→GPU append:
|
||||
```
|
||||
GPU tensor (K_new) → GPU KV cache (in-place append, no CPU)
|
||||
```
|
||||
|
||||
## 实现方案
|
||||
|
||||
### GPU KV Cache(简化版,非 paged)
|
||||
|
||||
先实现连续分配的 GPU KV cache(预分配 max_seq_len),消除 CPU round-trip。Paged allocation 留待后续优化。
|
||||
|
||||
```rust
|
||||
pub struct GpuKVCache {
|
||||
// 预分配: [num_layers, 2, num_kv_heads, max_seq_len, head_dim] on GPU
|
||||
k_caches: Vec<Tensor>, // per layer: [1, num_kv_heads, max_seq_len, head_dim]
|
||||
v_caches: Vec<Tensor>,
|
||||
seq_len: usize, // 当前已填充的长度
|
||||
max_seq_len: usize,
|
||||
}
|
||||
```
|
||||
|
||||
### Append 操作
|
||||
|
||||
用 cudaMemcpy D2D 将新 K/V 写入 cache 的正确偏移位置:
|
||||
```
|
||||
k_cache[layer][0, :, seq_len:seq_len+new, :] = k_new[0, :, :, :]
|
||||
```
|
||||
|
||||
### 读取操作
|
||||
|
||||
不需要拷贝——直接用 view/slice 返回 [0, :, 0:seq_len, :] 的 GPU tensor。
|
||||
|
||||
## 需要的新功能
|
||||
|
||||
1. Tensor slice 支持(view into sub-range of a dimension)
|
||||
2. GPU D2D copy at offset(写入 cache 指定位置)
|
||||
3. 去掉 Qwen3/GPT-2 forward 中的 CPU round-trip KV cache 路径
|
||||
|
||||
## Test Plan
|
||||
|
||||
- [ ] GPU KV cache 输出与 CPU KV cache bit-identical
|
||||
- [ ] Benchmark: TBT 应显著降低(消除 144 次 CPU round-trip)
|
||||
- [ ] 50-prompt correctness re-validation
|
||||
153
docs/12-continuous-batching.md
Normal file
153
docs/12-continuous-batching.md
Normal file
@@ -0,0 +1,153 @@
|
||||
# Phase 12: Continuous Batching + Request Scheduler — Design Document
|
||||
|
||||
## Goal
|
||||
|
||||
实现 iteration-level 请求调度,支持多个请求并发生成 token。核心能力:同时发 N 个请求,N 个请求同时产出 token,新请求可以在 mid-generation 加入 batch。
|
||||
|
||||
## 为什么需要 Continuous Batching
|
||||
|
||||
**当前问题(串行)**:
|
||||
```
|
||||
时间 → [req1 prefill][req1 decode x 100][req2 prefill][req2 decode x 50]...
|
||||
GPU利用: ████████████████████████████████████████████████████████████████████
|
||||
req2 等了 100 个 token 的时间才开始
|
||||
```
|
||||
|
||||
**目标(continuous batching)**:
|
||||
```
|
||||
时间 → [req1+req2 prefill][req1+req2 decode][req1 done, req3 加入][req2+req3 decode]...
|
||||
GPU利用: ████████████████████████████████████████████████████████████████████
|
||||
req2 和 req1 同时推理,req3 在 req1 完成后立即加入
|
||||
```
|
||||
|
||||
## 核心设计
|
||||
|
||||
### 数据结构
|
||||
|
||||
```rust
|
||||
pub struct Sequence {
|
||||
pub id: u64,
|
||||
pub prompt_tokens: Vec<u32>,
|
||||
pub generated_tokens: Vec<u32>,
|
||||
pub status: SeqStatus,
|
||||
pub max_tokens: usize,
|
||||
pub kv_cache: GpuKVCache, // 每个 seq 独立的 KV cache
|
||||
pub output_tx: mpsc::Sender<GenerateEvent>,
|
||||
}
|
||||
|
||||
pub enum SeqStatus {
|
||||
Waiting, // 在队列中等待被 admit
|
||||
Running, // 正在参与 batch forward
|
||||
Finished, // EOS 或 max_tokens 达到
|
||||
}
|
||||
|
||||
pub struct Scheduler {
|
||||
waiting: VecDeque<Sequence>,
|
||||
running: Vec<Sequence>,
|
||||
max_batch_size: usize, // 最大并发请求数
|
||||
next_seq_id: u64,
|
||||
}
|
||||
```
|
||||
|
||||
### 调度循环(Engine 主循环)
|
||||
|
||||
```rust
|
||||
loop {
|
||||
// Step 1: 回收已完成的 sequence
|
||||
running.retain(|seq| seq.status != Finished);
|
||||
|
||||
// Step 2: Admit 新请求(如果 running < max_batch_size)
|
||||
while running.len() < max_batch_size {
|
||||
if let Some(seq) = waiting.pop_front() {
|
||||
running.push(seq);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if running.is_empty() {
|
||||
// 没有任何工作,等待新请求
|
||||
let new_req = request_rx.recv(); // blocking wait
|
||||
waiting.push_back(new_req);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Step 3: 分类 — 哪些需要 prefill,哪些需要 decode
|
||||
let to_prefill: 新加入的 seq(generated_tokens 为空)
|
||||
let to_decode: 已在运行的 seq
|
||||
|
||||
// Step 4: 执行
|
||||
for seq in to_prefill {
|
||||
// Prefill: 完整 prompt 一次 forward
|
||||
model.forward_gpu_cache(&seq.prompt_tokens, &mut seq.kv_cache);
|
||||
seq.status = Running;
|
||||
}
|
||||
|
||||
// Decode: 每个 seq 独立做一步(当前不做 batch forward,留待优化)
|
||||
for seq in to_decode {
|
||||
let last_token = seq.last_generated_token();
|
||||
let logits = model.forward_gpu_cache(&[last_token], &mut seq.kv_cache);
|
||||
let next = sample_greedy(&logits);
|
||||
seq.generated_tokens.push(next);
|
||||
// 发送 token 给客户端
|
||||
seq.output_tx.blocking_send(Token { id: next, text: decode(next) });
|
||||
// 检查完成
|
||||
if next == eos || seq.generated_tokens.len() >= seq.max_tokens {
|
||||
seq.output_tx.blocking_send(Done);
|
||||
seq.status = Finished;
|
||||
}
|
||||
}
|
||||
|
||||
// Step 5: 检查是否有新请求到达(non-blocking)
|
||||
while let Ok(new_req) = request_rx.try_recv() {
|
||||
waiting.push_back(new_req);
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### 关键设计决策
|
||||
|
||||
1. **每个 seq 独立 KV cache**:当前不做 batch forward(需要对齐 seq_len),而是每个 seq 独立调用 model.forward_gpu_cache。未来优化为 batched forward。
|
||||
|
||||
2. **Prefill 和 Decode 混合**:新加入的 seq 先 prefill(一次 forward),然后下一轮加入 decode batch。
|
||||
|
||||
3. **Non-blocking request receive**:decode 循环中用 `try_recv()` 检查新请求,不阻塞推理。
|
||||
|
||||
4. **max_batch_size**:受限于 GPU 显存(每个 seq 的 KV cache 占用)。Qwen3-8B 单卡 32GB,每个 seq 的 KV cache 约 256 tokens × 8 heads × 128 dim × 2(KV) × 2B = 1MB。可以并发 ~100 seq。实际受限于推理速度。
|
||||
|
||||
## 与 Phase 13 (HTTP API) 的接口
|
||||
|
||||
```
|
||||
HTTP Handler Engine Thread
|
||||
│ │
|
||||
│ ──── GenerateRequest ────────► │
|
||||
│ (prompt_tokens, max_tokens, │
|
||||
│ output_tx) │
|
||||
│ │
|
||||
│ ◄──── GenerateEvent (Token/Done) ──── │
|
||||
│ (via tokio::sync::mpsc) │
|
||||
│ │
|
||||
```
|
||||
|
||||
多个 HTTP handler 可以同时提交请求。Engine 线程内部通过 Scheduler 管理并发。
|
||||
|
||||
## 验收测试
|
||||
|
||||
必须通过以下测试才算 Phase 12 完成:
|
||||
|
||||
1. **并发 3 请求测试**:同时发 3 个请求,验证 3 个请求同时产出 token(不是串行等待)
|
||||
2. **吞吐量测试**:并发请求的总 token 吞吐量应接近单请求(因为单个 seq 的 decode 是串行的)
|
||||
3. **动态加入测试**:先发 1 个请求开始生成,过 2 秒再发第 2 个,验证第 2 个立即开始(不等第 1 个完成)
|
||||
4. **正确性测试**:并发请求的输出内容应与单独跑每个请求一致
|
||||
|
||||
## 实现计划
|
||||
|
||||
1. 重构 Engine:从 `while recv → generate` 改为 scheduler loop
|
||||
2. 每个 Sequence 持有独立的 GpuKVCache
|
||||
3. 调度循环实现 admit + prefill + decode + finish
|
||||
4. HTTP API 侧改为 unbounded channel(允许多请求同时提交)
|
||||
5. 编写并发测试脚本
|
||||
|
||||
## 当前状态
|
||||
|
||||
**未实现**。当前是 FIFO 串行,一次只处理一个请求。本文档是实现的设计规格。
|
||||
133
docs/13-http-api.md
Normal file
133
docs/13-http-api.md
Normal file
@@ -0,0 +1,133 @@
|
||||
# Phase 13: HTTP API + Streaming — Design Document (Milestone ③)
|
||||
|
||||
## Goal
|
||||
|
||||
提供 OpenAI 兼容的 HTTP API,让 xserv 可以作为一个 serving 后端被任何 OpenAI SDK 调用。
|
||||
|
||||
## 职责划分
|
||||
|
||||
| 组件 | 职责 |
|
||||
|------|------|
|
||||
| Phase 12 (Scheduler/Engine) | 模型推理 + 请求调度 + token 生成循环 |
|
||||
| **Phase 13 (HTTP API)** | HTTP 请求解析 → 内部格式 → 提交给 engine → 从 channel 接收 token → 编码为 HTTP 响应 |
|
||||
|
||||
Phase 13 不关心模型如何推理,只负责 HTTP 协议层。
|
||||
|
||||
## 技术栈
|
||||
|
||||
- **HTTP framework**: axum 0.8
|
||||
- **Async runtime**: tokio
|
||||
- **Serialization**: serde_json
|
||||
- **Channel**: tokio::sync::mpsc (API ↔ Engine)
|
||||
|
||||
## API 端点
|
||||
|
||||
```
|
||||
GET /health → "ok"
|
||||
GET /v1/models → {"data": [{"id": "qwen3-8b", ...}]}
|
||||
POST /v1/chat/completions → JSON response (non-streaming)
|
||||
POST /v1/chat/completions → SSE stream (streaming, TODO)
|
||||
```
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
Client
|
||||
│ HTTP POST /v1/chat/completions
|
||||
▼
|
||||
┌──────────────────────────────┐
|
||||
│ axum handler │
|
||||
│ 1. Deserialize ChatRequest │
|
||||
│ 2. Build prompt text │
|
||||
│ 3. Tokenize (Mutex<Tokenizer>)│
|
||||
│ 4. Create mpsc channel │
|
||||
│ 5. Submit GenerateRequest │
|
||||
│ 6. await tokens from rx │
|
||||
│ 7. Build JSON response │
|
||||
└──────────────────────────────┘
|
||||
│ GenerateRequest via SyncSender
|
||||
▼
|
||||
┌──────────────────────────────┐
|
||||
│ Engine thread (Phase 12) │
|
||||
│ - recv() request │
|
||||
│ - model.forward_gpu_cache() │
|
||||
│ - blocking_send() tokens │
|
||||
└──────────────────────────────┘
|
||||
```
|
||||
|
||||
## OpenAI 兼容格式
|
||||
|
||||
### Request
|
||||
```json
|
||||
{
|
||||
"model": "qwen3-8b",
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are helpful."},
|
||||
{"role": "user", "content": "Hello"}
|
||||
],
|
||||
"max_tokens": 256,
|
||||
"stream": false
|
||||
}
|
||||
```
|
||||
|
||||
### Response (non-streaming)
|
||||
```json
|
||||
{
|
||||
"id": "chatcmpl-xxx",
|
||||
"object": "chat.completion",
|
||||
"created": 1234567890,
|
||||
"model": "qwen3-8b",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "Hi there!"},
|
||||
"finish_reason": "stop"
|
||||
}],
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}
|
||||
}
|
||||
```
|
||||
|
||||
### SSE Streaming (TODO)
|
||||
```
|
||||
data: {"choices":[{"delta":{"content":"Hi"}}]}
|
||||
|
||||
data: {"choices":[{"delta":{},"finish_reason":"stop"}]}
|
||||
|
||||
data: [DONE]
|
||||
```
|
||||
|
||||
## 当前实现状态
|
||||
|
||||
- [x] `/health` — 健康检查
|
||||
- [x] `/v1/models` — 模型列表
|
||||
- [x] `/v1/chat/completions` (non-streaming) — JSON response
|
||||
- [ ] `/v1/chat/completions` (streaming) — SSE
|
||||
- [ ] 完整的 `usage` 统计 (token 计数)
|
||||
- [ ] 错误处理 (400 for bad request, etc.)
|
||||
- [ ] 多轮对话 chat template
|
||||
|
||||
## Key Design Decisions
|
||||
|
||||
1. **Extension vs State**: 用 `axum::Extension<Arc<AppState>>` 而不是 `Router::with_state`,因为 `SyncSender` 不是 `Sync`(需要 Mutex 包装)。
|
||||
|
||||
2. **Engine 在独立 thread**: GPU 同步操作 block 线程,不能放在 tokio runtime 中。
|
||||
|
||||
3. **tokio::sync::mpsc 做 token 传输**: Engine (std thread) 用 `blocking_send()`,API (async) 用 `.recv().await`。跨 async/sync 边界通信。
|
||||
|
||||
## Test Plan
|
||||
|
||||
- [x] curl /health → "ok"
|
||||
- [x] curl /v1/models → JSON model list
|
||||
- [x] curl /v1/chat/completions → JSON with generated text
|
||||
- [ ] Python OpenAI SDK 兼容性测试
|
||||
- [ ] SSE streaming 测试
|
||||
- [ ] 多轮对话测试
|
||||
|
||||
## Takeaways
|
||||
|
||||
1. **axum 0.8 的 Handler trait 对 Send 很严格**:async fn 返回的 Future 必须是 Send。`std::sync::MutexGuard` 不是 Send,必须确保它不活过 await point(用 scope 或显式 drop)。
|
||||
|
||||
2. **std::sync::mpsc::SyncSender 不是 Sync**:不能直接放在 `Arc<T>` 中被多个 async task 共享。解决方案:`Mutex<SyncSender>` 或换用 `tokio::sync::mpsc::Sender`(是 Sync 的)。
|
||||
|
||||
3. **非 streaming 更简单,先跑通再加 SSE**:SSE streaming 涉及 `Stream` trait、lifetime 问题和复杂的类型推导。先用 collect-all-then-respond 跑通 E2E,streaming 作为增量优化。
|
||||
|
||||
4. **Engine 加载时间 ~20s(Qwen3-8B)**:需要在 server 启动后等 engine ready 才接受请求,否则请求会 hang 在 channel send 上。当前靠 sync_channel(1) 的背压天然处理。
|
||||
107
tools/test_concurrent.py
Normal file
107
tools/test_concurrent.py
Normal file
@@ -0,0 +1,107 @@
|
||||
"""
|
||||
Test concurrent request handling.
|
||||
Sends N requests simultaneously, verifies they all produce tokens concurrently.
|
||||
|
||||
Usage: python3 tools/test_concurrent.py <server_url> [num_requests]
|
||||
"""
|
||||
import sys
|
||||
import time
|
||||
import json
|
||||
import threading
|
||||
import urllib.request
|
||||
import urllib.error
|
||||
|
||||
|
||||
def send_request(url, prompt, max_tokens, results, idx):
|
||||
"""Send a chat completion request and record timing."""
|
||||
body = json.dumps({
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"max_tokens": max_tokens,
|
||||
}).encode()
|
||||
|
||||
req = urllib.request.Request(
|
||||
f"{url}/v1/chat/completions",
|
||||
data=body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
|
||||
t0 = time.time()
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=120) as resp:
|
||||
data = json.loads(resp.read())
|
||||
t1 = time.time()
|
||||
content = data["choices"][0]["message"]["content"]
|
||||
results[idx] = {
|
||||
"status": "ok",
|
||||
"content": content,
|
||||
"duration_s": t1 - t0,
|
||||
"finish_reason": data["choices"][0]["finish_reason"],
|
||||
}
|
||||
except Exception as e:
|
||||
t1 = time.time()
|
||||
results[idx] = {"status": "error", "error": str(e), "duration_s": t1 - t0}
|
||||
|
||||
|
||||
def main():
|
||||
url = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:9090"
|
||||
n = int(sys.argv[2]) if len(sys.argv) > 2 else 3
|
||||
max_tokens = 10
|
||||
|
||||
prompts = [
|
||||
"What is the capital of France?",
|
||||
"Tell me about quantum computing",
|
||||
"How do airplanes fly?",
|
||||
"What is machine learning?",
|
||||
"Explain gravity in simple terms",
|
||||
][:n]
|
||||
|
||||
print(f"Sending {n} concurrent requests to {url} (max_tokens={max_tokens})")
|
||||
print("=" * 70)
|
||||
|
||||
results = [None] * n
|
||||
threads = []
|
||||
|
||||
t_start = time.time()
|
||||
for i, prompt in enumerate(prompts):
|
||||
t = threading.Thread(target=send_request, args=(url, prompt, max_tokens, results, i))
|
||||
threads.append(t)
|
||||
t.start()
|
||||
|
||||
for t in threads:
|
||||
t.join()
|
||||
t_total = time.time() - t_start
|
||||
|
||||
print(f"\n{'#':>2} {'Status':>6} {'Duration':>8} {'Content':<50}")
|
||||
print("-" * 70)
|
||||
for i, r in enumerate(results):
|
||||
if r["status"] == "ok":
|
||||
content_short = r["content"].replace("\n", " ")[:48]
|
||||
print(f"{i+1:>2} {'OK':>6} {r['duration_s']:>6.1f}s {content_short}")
|
||||
else:
|
||||
print(f"{i+1:>2} {'FAIL':>6} {r['duration_s']:>6.1f}s {r['error'][:48]}")
|
||||
|
||||
print("=" * 70)
|
||||
print(f"Total wall time: {t_total:.1f}s")
|
||||
|
||||
# Analyze concurrency
|
||||
durations = [r["duration_s"] for r in results if r["status"] == "ok"]
|
||||
if len(durations) >= 2:
|
||||
sequential_estimate = sum(durations)
|
||||
actual_wall = t_total
|
||||
concurrency_ratio = sequential_estimate / actual_wall if actual_wall > 0 else 0
|
||||
|
||||
print(f"Sum of individual durations: {sequential_estimate:.1f}s")
|
||||
print(f"Actual wall time: {actual_wall:.1f}s")
|
||||
print(f"Concurrency ratio: {concurrency_ratio:.2f}x")
|
||||
|
||||
if concurrency_ratio > 1.5:
|
||||
print("✓ CONCURRENT: requests are being processed in parallel")
|
||||
else:
|
||||
print("✗ SERIAL: requests appear to be processed sequentially")
|
||||
|
||||
all_ok = all(r["status"] == "ok" for r in results)
|
||||
print(f"\nAll requests succeeded: {all_ok}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user