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

44
crates/xserv-server/src/api.rs
View File

@@ -72,7 +72,10 @@ impl ChatTemplate {
let source = std::fs::read_to_string(&jinja_path)
.unwrap_or_else(|e| panic!("failed to read {}: {e}", jinja_path.display()));
eprintln!("[chat-template] loaded from {}", jinja_path.display());
return Self { source, model_type: model_type.to_string() };
return Self {
source,
model_type: model_type.to_string(),
};
}
// 2. Try tokenizer_config.json → chat_template field
@@ -82,7 +85,10 @@ impl ChatTemplate {
if let Ok(v) = serde_json::from_str::<serde_json::Value>(&data) {
if let Some(ct) = v.get("chat_template").and_then(|v| v.as_str()) {
eprintln!("[chat-template] loaded from tokenizer_config.json");
return Self { source: ct.to_string(), model_type: model_type.to_string() };
return Self {
source: ct.to_string(),
model_type: model_type.to_string(),
};
}
}
}
@@ -90,7 +96,10 @@ impl ChatTemplate {
// 3. No template found — use empty source, will fall back to hardcoded
eprintln!("[chat-template] no Jinja template found, using hardcoded fallback");
Self { source: String::new(), model_type: model_type.to_string() }
Self {
source: String::new(),
model_type: model_type.to_string(),
}
}
pub fn render(&self, messages: &[Message]) -> String {
@@ -206,7 +215,10 @@ fn build_prompt_gpt_oss(messages: &[Message]) -> String {
prompt.push_str("<|start|>system<|message|>");
prompt.push_str("You are ChatGPT, a large language model trained by OpenAI.\n");
prompt.push_str("Knowledge cutoff: 2024-06\n");
prompt.push_str(&format!("Current date: {}\n\n", strftime_now("%Y-%m-%d".to_string())));
prompt.push_str(&format!(
"Current date: {}\n\n",
strftime_now("%Y-%m-%d".to_string())
));
prompt.push_str("Reasoning: low\n\n");
prompt.push_str("# Valid channels: analysis, commentary, final. Channel must be included for every message.");
prompt.push_str("<|end|>");
@@ -334,13 +346,11 @@ async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
"completion_tokens": completion_token_count,
"total_tokens": prompt_token_count + completion_token_count
}
})).into_response()
}))
.into_response()
}
fn chat_stream(
state: Arc<AppState>,
req: ChatRequest,
) -> Response {
fn chat_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
let id = format!("chatcmpl-{}", Uuid::new_v4());
let model_name = state.model_name.clone();
let created = unix_timestamp();
@@ -356,7 +366,8 @@ fn chat_stream(
if prompt_tokens.len() >= max_seq_len {
return bad_request(format!(
"prompt is {} tokens, exceeds max_seq_len {}",
prompt_tokens.len(), max_seq_len
prompt_tokens.len(),
max_seq_len
));
}
let max_tokens = req.max_tokens.min(max_seq_len - prompt_tokens.len());
@@ -413,7 +424,9 @@ fn chat_stream(
}
});
Sse::new(ReceiverStream::new(sse_rx)).keep_alive(KeepAlive::default()).into_response()
Sse::new(ReceiverStream::new(sse_rx))
.keep_alive(KeepAlive::default())
.into_response()
}
fn validate_request(req: &ChatRequest, model_name: &str) -> Option<Response> {
@@ -436,8 +449,13 @@ fn validate_request(req: &ChatRequest, model_name: &str) -> Option<Response> {
/// prior handler panicked) and returns a clean 503 instead of panicking when the
/// engine thread is gone, so one dead engine doesn't cascade into every request.
fn submit_to_engine(state: &AppState, req: GenerateRequest) -> Result<(), Response> {
let sender = state.engine_sender.lock().unwrap_or_else(|e| e.into_inner());
sender.send(req).map_err(|_| service_unavailable("inference engine is not available"))
let sender = state
.engine_sender
.lock()
.unwrap_or_else(|e| e.into_inner());
sender
.send(req)
.map_err(|_| service_unavailable("inference engine is not available"))
}
fn service_unavailable(message: impl Into<String>) -> Response {

105
crates/xserv-server/src/engine.rs
View File

@@ -1,10 +1,10 @@
use std::collections::VecDeque;
use std::path::Path;
use std::sync::mpsc;
use std::sync::Once;
use std::sync::mpsc;
use std::time::Instant;
use xserv_model::{ModelConfig, PagedKVCache, Qwen3, SamplingParams, sample, BLOCK_SIZE};
use xserv_model::loader;
use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, SamplingParams, sample};
use xserv_tensor::{DType, Device};
use xserv_tokenizer::Tokenizer;
@@ -109,12 +109,23 @@ impl Engine {
(total_blocks * bytes_per_block) as f64 / 1e9,
info.free_memory as f64 / 1e9,
);
Self { model, config, tokenizer, max_batch_size, max_seq_len, paged_cache }
Self {
model,
config,
tokenizer,
max_batch_size,
max_seq_len,
paged_cache,
}
}
pub fn tokenizer(&self) -> &Tokenizer { &self.tokenizer }
pub fn tokenizer(&self) -> &Tokenizer {
&self.tokenizer
}
pub fn max_seq_len(&self) -> usize { self.max_seq_len }
pub fn max_seq_len(&self) -> usize {
self.max_seq_len
}
/// Main scheduler loop. Receives requests from channel, manages concurrent sequences.
///
@@ -134,7 +145,8 @@ impl Engine {
loop {
// Step 1: Remove finished sequences and return their slots.
let finished_slots: Vec<usize> = running.iter()
let finished_slots: Vec<usize> = running
.iter()
.filter(|s| is_finished(s))
.filter_map(|s| s.seq_slot)
.collect();
@@ -147,10 +159,16 @@ impl Engine {
// room (oldest first). They resume decoding from where they paused.
while running.len() < self.max_batch_size && !swapped.is_empty() {
let slot = swapped[0].seq_slot.expect("swapped slot");
if !self.paged_cache.can_swap_in(slot) { break; }
if !self.paged_cache.can_swap_in(slot) {
break;
}
self.paged_cache.swap_in(slot).expect("swap_in");
let seq = swapped.remove(0);
eprintln!("[scheduler] swapped in seq {} ({} blocks)", seq.id, self.paged_cache.block_count(slot));
eprintln!(
"[scheduler] swapped in seq {} ({} blocks)",
seq.id,
self.paged_cache.block_count(slot)
);
running.push(seq);
}
@@ -161,14 +179,22 @@ impl Engine {
let mut avail = self.paged_cache.free_blocks();
let decode_reserve = running.len();
while running.len() < self.max_batch_size {
let Some(front) = waiting.front() else { break; };
let Some(front) = waiting.front() else {
break;
};
let prompt_blocks = front.prompt_tokens.len().div_ceil(BLOCK_SIZE).max(1);
if avail < prompt_blocks + decode_reserve { break; }
if avail < prompt_blocks + decode_reserve {
break;
}
let free_slot = (0..self.paged_cache.max_seqs())
.find(|&s| self.paged_cache.is_slot_free(s));
let Some(slot) = free_slot else { break; };
let Some(slot) = free_slot else {
break;
};
let mut seq = waiting.pop_front().unwrap();
self.paged_cache.register_sequence(slot).expect("register paged slot");
self.paged_cache
.register_sequence(slot)
.expect("register paged slot");
seq.seq_slot = Some(slot);
running.push(seq);
avail -= prompt_blocks; // projected free after this seq prefills
@@ -199,7 +225,9 @@ impl Engine {
if !seq.prefilled {
let slot = seq.seq_slot.expect("slot");
let logits = self.model.forward_prefill_paged(
&seq.prompt_tokens, slot, &mut self.paged_cache,
&seq.prompt_tokens,
slot,
&mut self.paged_cache,
);
let next = sample(&logits, &seq.sampling);
seq.generated_tokens.push(next);
@@ -219,13 +247,18 @@ impl Engine {
&& !newly_prefilled.contains(&running[p].id)
&& running[p].seq_slot.is_some()
});
let Some(pos) = victim else { break; };
let Some(pos) = victim else {
break;
};
let seq = running.remove(pos);
let slot = seq.seq_slot.unwrap();
if self.paged_cache.can_swap_out(slot) {
let nblocks = self.paged_cache.block_count(slot);
self.paged_cache.swap_out(slot).expect("swap_out");
eprintln!("[scheduler] preempt: swapped out seq {} ({nblocks} blocks) to host", seq.id);
eprintln!(
"[scheduler] preempt: swapped out seq {} ({nblocks} blocks) to host",
seq.id
);
swapped.push(seq);
needed = decode_block_need(&self.paged_cache, &running, &newly_prefilled);
} else {
@@ -235,7 +268,9 @@ impl Engine {
}
// Step 5c: Batched paged decode for the surviving prefilled sequences.
let decode_indices: Vec<usize> = running.iter().enumerate()
let decode_indices: Vec<usize> = running
.iter()
.enumerate()
.filter(|(_, s)| s.prefilled && !newly_prefilled.contains(&s.id))
.map(|(i, _)| i)
.collect();
@@ -246,25 +281,32 @@ impl Engine {
eprintln!("[scheduler] paged decode active");
});
let tokens: Vec<u32> = decode_indices.iter()
let tokens: Vec<u32> = decode_indices
.iter()
.map(|&i| *running[i].generated_tokens.last().unwrap())
.collect();
let positions: Vec<usize> = decode_indices.iter()
let positions: Vec<usize> = decode_indices
.iter()
.map(|&i| self.paged_cache.seq_len(running[i].seq_slot.unwrap()))
.collect();
let slots: Vec<usize> = decode_indices.iter()
let slots: Vec<usize> = decode_indices
.iter()
.map(|&i| running[i].seq_slot.unwrap())
.collect();
let logits = self.model.forward_decode_paged(
&tokens, &positions, &slots, &mut self.paged_cache,
&tokens,
&positions,
&slots,
&mut self.paged_cache,
);
// Fast path: every active sequence is greedy → run argmax on
// the GPU and only D2H the chosen token ids (a few bytes per
// sequence) instead of the full [B, vocab_size] BF16 logits
// (~1.2 MB for B=4, Qwen3 vocab=152K).
let all_greedy = decode_indices.iter()
let all_greedy = decode_indices
.iter()
.all(|&i| running[i].sampling.temperature == 0.0);
if all_greedy {
let next_ids = xserv_kernels::argmax_bf16_to_host(&logits);
@@ -285,11 +327,15 @@ impl Engine {
let row_start = j * vocab_size;
let row_logits = &data[row_start..row_start + vocab_size];
let next = if running[i].sampling.temperature == 0.0 {
row_logits.iter().enumerate()
row_logits
.iter()
.enumerate()
.max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
.map(|(idx, _)| idx as u32).unwrap()
.map(|(idx, _)| idx as u32)
.unwrap()
} else {
let row_tensor = xserv_tensor::Tensor::from_slice(row_logits, &[1, vocab_size]);
let row_tensor =
xserv_tensor::Tensor::from_slice(row_logits, &[1, vocab_size]);
sample(&row_tensor, &running[i].sampling)
};
running[i].generated_tokens.push(next);
@@ -334,7 +380,8 @@ impl Engine {
/// Total additional GPU blocks the next decode step needs across all
/// currently-decoding (prefilled, not just-prefilled) sequences.
fn decode_block_need(paged: &PagedKVCache, running: &[Sequence], newly_prefilled: &[u64]) -> usize {
running.iter()
running
.iter()
.filter(|s| s.prefilled && !newly_prefilled.contains(&s.id))
.filter_map(|s| s.seq_slot)
.map(|slot| paged.additional_blocks_needed(slot, 1))
@@ -372,8 +419,12 @@ fn send_token_if_nonempty(seq: &Sequence, text: String) {
}
fn is_finished(seq: &Sequence) -> bool {
if seq.generated_tokens.is_empty() { return false; }
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; }
if seq.generated_tokens.len() >= seq.max_tokens {
return true;
}
seq.sender.is_closed() || seq.eos_token_id == Some(last)
}

45
crates/xserv-server/src/main.rs
View File

@@ -3,10 +3,13 @@ mod engine;
mod pp_engine;
mod tp_engine;
use axum::{routing::{get, post}, Extension, Router};
use std::path::PathBuf;
use std::sync::{mpsc, Arc, Mutex};
use axum::{
Extension, Router,
routing::{get, post},
};
use engine::GenerateRequest;
use std::path::PathBuf;
use std::sync::{Arc, Mutex, mpsc};
use xserv_model::ModelConfig;
pub struct AppState {
@@ -21,40 +24,48 @@ pub struct AppState {
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] [--max-seq-len N] [--swap-space-gb N] [--tp N] [--pp N]");
eprintln!(
"Usage: xserv-server <model-dir> [--port PORT] [--max-batch N] [--max-seq-len N] [--swap-space-gb N] [--tp N] [--pp N]"
);
std::process::exit(1);
}
let model_dir = PathBuf::from(&args[1]);
let port: u16 = args.iter()
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()
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)
.max(1);
let requested_max_seq_len: usize = args.iter()
let requested_max_seq_len: usize = args
.iter()
.position(|a| a == "--max-seq-len")
.and_then(|i| args.get(i + 1))
.and_then(|s| s.parse().ok())
.unwrap_or(2048)
.max(1);
let swap_space_gb: usize = args.iter()
let swap_space_gb: usize = args
.iter()
.position(|a| a == "--swap-space-gb")
.and_then(|i| args.get(i + 1))
.and_then(|s| s.parse().ok())
.unwrap_or(8);
let tp: usize = args.iter()
let tp: usize = args
.iter()
.position(|a| a == "--tp")
.and_then(|i| args.get(i + 1))
.and_then(|s| s.parse().ok())
.unwrap_or(1)
.max(1);
let pp: usize = args.iter()
let pp: usize = args
.iter()
.position(|a| a == "--pp")
.and_then(|i| args.get(i + 1))
.and_then(|s| s.parse().ok())
@@ -69,7 +80,9 @@ async fn main() {
// tp=1 (single-rank world) so quantized models can serve on one GPU.
let is_gpt_oss = model_config.model_type.as_deref() == Some("gpt_oss");
if pp > 1 && is_gpt_oss {
eprintln!("gpt-oss is not supported by the pipeline-parallel engine (Qwen3 only); use --tp instead");
eprintln!(
"gpt-oss is not supported by the pipeline-parallel engine (Qwen3 only); use --tp instead"
);
std::process::exit(1);
}
let model_max_seq_len = model_config.max_seq_len();
@@ -84,7 +97,8 @@ async fn main() {
);
}
let model_name = model_dir.file_name()
let model_name = model_dir
.file_name()
.map(|n| n.to_string_lossy().to_string())
.unwrap_or_else(|| "unknown".to_string());
@@ -99,7 +113,12 @@ async fn main() {
// Pipeline-parallel path: stage-0 coordinator + worker stage threads.
pp_engine::run_pp(&model_dir_clone, pp, max_seq_len, rx);
} else if tp <= 1 && !is_gpt_oss {
let mut engine = engine::Engine::load_with_swap(&model_dir_clone, max_batch, max_seq_len, swap_space_gb);
let mut engine = engine::Engine::load_with_swap(
&model_dir_clone,
max_batch,
max_seq_len,
swap_space_gb,
);
engine.run(rx);
} else {
// Tensor-parallel path: rank-0 coordinator + worker rank threads.

100
crates/xserv-server/src/pp_engine.rs
View File

@@ -15,15 +15,15 @@
use std::ffi::c_void;
use std::path::{Path, PathBuf};
use std::sync::mpsc;
use std::sync::Arc;
use std::sync::mpsc;
use std::thread;
use half::bf16;
use xserv_distributed::{PpContext, UniqueId};
use xserv_model::loader;
use xserv_model::sampling::SamplingParams;
use xserv_model::{sample, ModelConfig, PagedKVCache, Qwen3, BLOCK_SIZE};
use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, sample};
use xserv_tensor::{DType, Device, Tensor};
use xserv_tokenizer::Tokenizer;
@@ -38,9 +38,16 @@ enum PpCommand {
Free(usize),
/// Receive `[n_tokens, hidden]` from the previous stage, run this stage's
/// layers; if last stage, sample with `sampling` and return the token.
Prefill { n_tokens: usize, slot: usize, sampling: SamplingParams },
Prefill {
n_tokens: usize,
slot: usize,
sampling: SamplingParams,
},
/// Receive `[1, hidden]`, run this stage's layers; last stage samples.
Decode { slot: usize, sampling: SamplingParams },
Decode {
slot: usize,
sampling: SamplingParams,
},
Shutdown,
}
@@ -76,9 +83,21 @@ fn build_stage(
let max_blocks_per_seq = max_seq_len.div_ceil(BLOCK_SIZE);
let total_blocks = max_blocks_per_seq + 8; // v1 serial: one active sequence
let cache = PagedKVCache::new(
&stage_config, total_blocks, 0, 4, max_blocks_per_seq, DType::BF16, device,
&stage_config,
total_blocks,
0,
4,
max_blocks_per_seq,
DType::BF16,
device,
);
StageCtx { model, cache, pp, hidden: config.hidden(), device }
StageCtx {
model,
cache,
pp,
hidden: config.hidden(),
device,
}
}
/// Allocate a zeroed `[n, hidden]` device tensor and receive into it from `peer`.
@@ -110,7 +129,15 @@ fn worker_loop(
ack_tx: mpsc::Sender<()>,
token_tx: mpsc::Sender<u32>,
) {
let mut sc = build_stage(&model_dir, &config, stage, world, stage as u32, max_seq_len, id);
let mut sc = build_stage(
&model_dir,
&config,
stage,
world,
stage as u32,
max_seq_len,
id,
);
let is_last = stage == world - 1;
let prev = stage - 1;
let next = stage + 1;
@@ -125,7 +152,11 @@ fn worker_loop(
sc.cache.free_sequence(slot);
let _ = ack_tx.send(());
}
PpCommand::Prefill { n_tokens, slot, sampling } => {
PpCommand::Prefill {
n_tokens,
slot,
sampling,
} => {
let x = recv_hidden(&sc, n_tokens, prev);
let x = sc.model.forward_layers_prefill(x, slot, &mut sc.cache);
if is_last {
@@ -155,7 +186,12 @@ fn worker_loop(
/// Run the PP coordinator (stage 0) on the calling thread. Spawns worker stages
/// 1..world and consumes generation requests from `rx`.
pub fn run_pp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Receiver<GenerateRequest>) {
pub fn run_pp(
model_dir: &Path,
world: usize,
max_seq_len: usize,
rx: mpsc::Receiver<GenerateRequest>,
) {
assert!(world >= 2, "run_pp requires world >= 2");
let config = ModelConfig::from_file(&model_dir.join("config.json"));
assert!(
@@ -179,7 +215,17 @@ pub fn run_pp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
let model_dir = model_dir.to_path_buf();
let config = config.clone();
thread::spawn(move || {
worker_loop(stage, world, id, model_dir, config, max_seq_len, ctx_rx, ack_tx, token_tx);
worker_loop(
stage,
world,
id,
model_dir,
config,
max_seq_len,
ctx_rx,
ack_tx,
token_tx,
);
});
}
@@ -207,11 +253,14 @@ pub fn run_pp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
wait_acks(&ack_rx);
// Prefill: embed prompt, run stage-0 layers, push hidden into the pipe.
broadcast(&cmd_txs, PpCommand::Prefill {
n_tokens: req.prompt_tokens.len(),
slot,
sampling: req.sampling.clone(),
});
broadcast(
&cmd_txs,
PpCommand::Prefill {
n_tokens: req.prompt_tokens.len(),
slot,
sampling: req.sampling.clone(),
},
);
let x = sc.model.embed(&req.prompt_tokens);
let x = sc.model.forward_layers_prefill(x, slot, &mut sc.cache);
send_hidden(&sc, &x, next_peer);
@@ -228,7 +277,13 @@ pub fn run_pp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
if generated >= req.max_tokens {
break "length";
}
broadcast(&cmd_txs, PpCommand::Decode { slot, sampling: req.sampling.clone() });
broadcast(
&cmd_txs,
PpCommand::Decode {
slot,
sampling: req.sampling.clone(),
},
);
let x = sc.model.embed(&[next]);
let x = sc.model.forward_layers_decode(x, &[slot], &mut sc.cache);
send_hidden(&sc, &x, next_peer);
@@ -239,9 +294,14 @@ pub fn run_pp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
let tail = tokenizer.flush_decode_stream(&mut decode_buf);
if !tail.is_empty() {
let _ = req.sender.blocking_send(GenerateEvent::Token { id: next, text: tail });
let _ = req.sender.blocking_send(GenerateEvent::Token {
id: next,
text: tail,
});
}
let _ = req.sender.blocking_send(GenerateEvent::Done { finish_reason: finish.to_string() });
let _ = req.sender.blocking_send(GenerateEvent::Done {
finish_reason: finish.to_string(),
});
broadcast(&cmd_txs, PpCommand::Free(slot));
sc.cache.free_sequence(slot);
@@ -258,6 +318,8 @@ fn emit_text(tokenizer: &Tokenizer, req: &GenerateRequest, token_id: u32, buf: &
}
let text = tokenizer.decode_token_stream(token_id, buf);
if !text.is_empty() {
let _ = req.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
let _ = req
.sender
.blocking_send(GenerateEvent::Token { id: token_id, text });
}
}

154
crates/xserv-server/src/tp_engine.rs
View File

@@ -13,13 +13,16 @@
//! work; the single-GPU `Engine` still handles TP=1.
use std::path::{Path, PathBuf};
use std::sync::mpsc;
use std::sync::Arc;
use std::sync::mpsc;
use std::thread;
use xserv_distributed::{TpContext, UniqueId};
use xserv_model::loader;
use xserv_model::{sample, sample_greedy_penalized, GptOss, GraphedGptOssDecoder, ModelConfig, PagedKVCache, Qwen3, BLOCK_SIZE};
use xserv_model::{
BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, PagedKVCache, Qwen3, sample,
sample_greedy_penalized,
};
use xserv_tensor::{DType, Device, Tensor};
use xserv_tokenizer::Tokenizer;
@@ -29,8 +32,15 @@ use crate::engine::{GenerateEvent, GenerateRequest};
enum TpCommand {
Register(usize),
Free(usize),
Prefill { tokens: Vec<u32>, slot: usize },
Decode { tokens: Vec<u32>, positions: Vec<usize>, slots: Vec<usize> },
Prefill {
tokens: Vec<u32>,
slot: usize,
},
Decode {
tokens: Vec<u32>,
positions: Vec<usize>,
slots: Vec<usize>,
},
Shutdown,
}
@@ -40,14 +50,25 @@ enum TpModel {
}
impl TpModel {
fn forward_prefill_paged(&self, tokens: &[u32], slot: usize, cache: &mut PagedKVCache) -> Tensor {
fn forward_prefill_paged(
&self,
tokens: &[u32],
slot: usize,
cache: &mut PagedKVCache,
) -> Tensor {
match self {
TpModel::Qwen3(m) => m.forward_prefill_paged(tokens, slot, cache),
TpModel::GptOss(m) => m.forward_prefill_paged(tokens, slot, cache),
}
}
fn forward_decode_paged(&self, tokens: &[u32], positions: &[usize], slots: &[usize], cache: &mut PagedKVCache) -> Tensor {
fn forward_decode_paged(
&self,
tokens: &[u32],
positions: &[usize],
slots: &[usize],
cache: &mut PagedKVCache,
) -> Tensor {
match self {
TpModel::Qwen3(m) => m.forward_decode_paged(tokens, positions, slots, cache),
TpModel::GptOss(m) => m.forward_decode_paged(tokens, positions, slots, cache),
@@ -65,8 +86,12 @@ struct RankCtx {
/// (lazy capture, replay thereafter); everything else runs eager.
fn rank_decode(rc: &mut RankCtx, tokens: &[u32], positions: &[usize], slots: &[usize]) -> Tensor {
match &rc.model {
TpModel::GptOss(m) => rc.decoder.decode(m, tokens, positions, slots, &mut rc.cache),
TpModel::Qwen3(_) => rc.model.forward_decode_paged(tokens, positions, slots, &mut rc.cache),
TpModel::GptOss(m) => rc
.decoder
.decode(m, tokens, positions, slots, &mut rc.cache),
TpModel::Qwen3(_) => rc
.model
.forward_decode_paged(tokens, positions, slots, &mut rc.cache),
}
}
@@ -81,17 +106,42 @@ fn build_rank(
) -> RankCtx {
let weights = loader::load_model_dir(model_dir, Device::Cpu);
let model = if config.is_moe() {
TpModel::GptOss(GptOss::from_weights_tp(config.clone(), weights, rank, world, device, tp))
TpModel::GptOss(GptOss::from_weights_tp(
config.clone(),
weights,
rank,
world,
device,
tp,
))
} else {
TpModel::Qwen3(Qwen3::from_weights_tp(config.clone(), weights, rank, world, device, tp))
TpModel::Qwen3(Qwen3::from_weights_tp(
config.clone(),
weights,
rank,
world,
device,
tp,
))
};
let local_kv = config.num_kv_heads() / world;
let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
let total_blocks = max_blocks_per_seq + 8;
let cache = PagedKVCache::new_tp(
config, local_kv, total_blocks, 0, 4, max_blocks_per_seq, DType::BF16, device,
config,
local_kv,
total_blocks,
0,
4,
max_blocks_per_seq,
DType::BF16,
device,
);
RankCtx { model, cache, decoder: GraphedGptOssDecoder::new() }
RankCtx {
model,
cache,
decoder: GraphedGptOssDecoder::new(),
}
}
fn worker_loop(
@@ -105,7 +155,15 @@ fn worker_loop(
ack_tx: mpsc::Sender<()>,
) {
let tp = Arc::new(TpContext::init(rank, world, id, rank as u32));
let mut rc = build_rank(&model_dir, &config, rank, world, rank as u32, max_seq_len, Some(tp));
let mut rc = build_rank(
&model_dir,
&config,
rank,
world,
rank as u32,
max_seq_len,
Some(tp),
);
while let Ok(cmd) = cmd_rx.recv() {
match cmd {
TpCommand::Register(slot) => {
@@ -115,7 +173,11 @@ fn worker_loop(
TpCommand::Prefill { tokens, slot } => {
let _ = rc.model.forward_prefill_paged(&tokens, slot, &mut rc.cache);
}
TpCommand::Decode { tokens, positions, slots } => {
TpCommand::Decode {
tokens,
positions,
slots,
} => {
let _ = rank_decode(&mut rc, &tokens, &positions, &slots);
}
TpCommand::Shutdown => {
@@ -129,7 +191,12 @@ fn worker_loop(
/// Run the TP coordinator (rank 0) on the calling thread. Spawns worker ranks
/// internally and consumes generation requests from `rx`.
pub fn run_tp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Receiver<GenerateRequest>) {
pub fn run_tp(
model_dir: &Path,
world: usize,
max_seq_len: usize,
rx: mpsc::Receiver<GenerateRequest>,
) {
// world=1 is a valid single-rank configuration (gpt-oss has no
// single-GPU engine path; NCCL init and all_reduce no-op at world=1).
assert!(world >= 1, "run_tp requires world >= 1");
@@ -152,7 +219,16 @@ pub fn run_tp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
let model_dir = model_dir.to_path_buf();
let config = config.clone();
thread::spawn(move || {
worker_loop(rank, world, id, model_dir, config, max_seq_len, ctx_rx, ack_tx);
worker_loop(
rank,
world,
id,
model_dir,
config,
max_seq_len,
ctx_rx,
ack_tx,
);
});
}
@@ -165,10 +241,14 @@ pub fn run_tp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
// models loop under pure greedy when numerics diverge from the reference).
// Off by default; XSERV_REP_PENALTY>1 enables it over the last
// XSERV_REP_WINDOW generated tokens. Applied only on the greedy path.
let rep_penalty: f32 = std::env::var("XSERV_REP_PENALTY").ok()
.and_then(|s| s.parse().ok()).unwrap_or(1.0);
let rep_window: usize = std::env::var("XSERV_REP_WINDOW").ok()
.and_then(|s| s.parse().ok()).unwrap_or(128);
let rep_penalty: f32 = std::env::var("XSERV_REP_PENALTY")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(1.0);
let rep_window: usize = std::env::var("XSERV_REP_WINDOW")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(128);
let pick = |logits: &Tensor, sp: &xserv_model::SamplingParams, history: &[u32]| -> u32 {
if rep_penalty > 1.0 && sp.temperature == 0.0 {
let start = history.len().saturating_sub(rep_window);
@@ -197,8 +277,16 @@ pub fn run_tp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
wait_acks(&ack_rx);
// Prefill.
broadcast(&cmd_txs, TpCommand::Prefill { tokens: req.prompt_tokens.clone(), slot });
let logits = rc.model.forward_prefill_paged(&req.prompt_tokens, slot, &mut rc.cache);
broadcast(
&cmd_txs,
TpCommand::Prefill {
tokens: req.prompt_tokens.clone(),
slot,
},
);
let logits = rc
.model
.forward_prefill_paged(&req.prompt_tokens, slot, &mut rc.cache);
wait_acks(&ack_rx);
let mut gen_ids: Vec<u32> = Vec::new();
let mut next = pick(&logits, &req.sampling, &gen_ids);
@@ -216,7 +304,14 @@ pub fn run_tp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
break "length";
}
let pos = rc.cache.seq_len(slot);
broadcast(&cmd_txs, TpCommand::Decode { tokens: vec![next], positions: vec![pos], slots: vec![slot] });
broadcast(
&cmd_txs,
TpCommand::Decode {
tokens: vec![next],
positions: vec![pos],
slots: vec![slot],
},
);
let logits = rank_decode(&mut rc, &[next], &[pos], &[slot]);
wait_acks(&ack_rx);
next = pick(&logits, &req.sampling, &gen_ids);
@@ -227,9 +322,14 @@ pub fn run_tp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
let tail = tokenizer.flush_decode_stream(&mut decode_buf);
if !tail.is_empty() {
let _ = req.sender.blocking_send(GenerateEvent::Token { id: next, text: tail });
let _ = req.sender.blocking_send(GenerateEvent::Token {
id: next,
text: tail,
});
}
let _ = req.sender.blocking_send(GenerateEvent::Done { finish_reason: finish.to_string() });
let _ = req.sender.blocking_send(GenerateEvent::Done {
finish_reason: finish.to_string(),
});
broadcast(&cmd_txs, TpCommand::Free(slot));
rc.cache.free_sequence(slot);
@@ -246,6 +346,8 @@ fn emit_text(tokenizer: &Tokenizer, req: &GenerateRequest, token_id: u32, buf: &
}
let text = tokenizer.decode_token_stream(token_id, buf);
if !text.is_empty() {
let _ = req.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
let _ = req
.sender
.blocking_send(GenerateEvent::Token { id: token_id, text });
}
}