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

310
crates/xserv-model/src/bin/xserv-chat.rs
View File

@@ -1,10 +1,13 @@
use std::io::{self, IsTerminal, Read, Write};
use std::path::PathBuf;
use std::sync::{mpsc, Arc};
use std::sync::{Arc, mpsc};
use std::thread;
use xserv_model::{GraphedGptOssDecoder, loader, sample, sample_greedy_penalized, GptOss, ModelConfig, PagedKVCache, Qwen3, SamplingParams, BLOCK_SIZE};
use xserv_model::{
BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, PagedKVCache, Qwen3, SamplingParams,
loader, sample, sample_greedy_penalized,
};
use xserv_tensor::{DType, Device};
use xserv_tokenizer::Tokenizer;
@@ -14,13 +17,24 @@ enum ChatModel {
}
impl ChatModel {
fn forward_prefill_paged(&self, tokens: &[u32], slot: usize, cache: &mut PagedKVCache) -> xserv_tensor::Tensor {
fn forward_prefill_paged(
&self,
tokens: &[u32],
slot: usize,
cache: &mut PagedKVCache,
) -> xserv_tensor::Tensor {
match self {
ChatModel::Qwen3(m) => m.forward_prefill_paged(tokens, slot, cache),
ChatModel::GptOss(m) => m.forward_prefill_paged(tokens, slot, cache),
}
}
fn forward_decode_paged(&self, tokens: &[u32], positions: &[usize], slots: &[usize], cache: &mut PagedKVCache) -> xserv_tensor::Tensor {
fn forward_decode_paged(
&self,
tokens: &[u32],
positions: &[usize],
slots: &[usize],
cache: &mut PagedKVCache,
) -> xserv_tensor::Tensor {
match self {
ChatModel::Qwen3(m) => m.forward_decode_paged(tokens, positions, slots, cache),
ChatModel::GptOss(m) => m.forward_decode_paged(tokens, positions, slots, cache),
@@ -33,8 +47,15 @@ impl ChatModel {
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>,
},
}
struct TpHandle {
@@ -56,7 +77,8 @@ impl TpHandle {
}
fn tp_worker_loop(
rank: usize, world: usize,
rank: usize,
world: usize,
id: xserv_distributed::UniqueId,
model_dir: std::path::PathBuf,
config: ModelConfig,
@@ -64,29 +86,68 @@ fn tp_worker_loop(
cmd_rx: mpsc::Receiver<TpCommand>,
ack_tx: mpsc::Sender<()>,
) {
let tp = Arc::new(xserv_distributed::TpContext::init(rank, world, id, rank as u32));
let tp = Arc::new(xserv_distributed::TpContext::init(
rank,
world,
id,
rank as u32,
));
let weights = loader::load_model_dir(&model_dir, Device::Cpu);
let model = if config.is_moe() {
ChatModel::GptOss(GptOss::from_weights_tp(config.clone(), weights, rank, world, rank as u32, Some(tp)))
ChatModel::GptOss(GptOss::from_weights_tp(
config.clone(),
weights,
rank,
world,
rank as u32,
Some(tp),
))
} else {
ChatModel::Qwen3(Qwen3::from_weights_tp(config.clone(), weights, rank, world, rank as u32, Some(tp)))
ChatModel::Qwen3(Qwen3::from_weights_tp(
config.clone(),
weights,
rank,
world,
rank as u32,
Some(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 mut cache = PagedKVCache::new_tp(
&config, local_kv, total_blocks, 0, 1, max_blocks_per_seq, DType::BF16, rank as u32,
&config,
local_kv,
total_blocks,
0,
1,
max_blocks_per_seq,
DType::BF16,
rank as u32,
);
let mut decoder = GraphedGptOssDecoder::new();
while let Ok(cmd) = cmd_rx.recv() {
match cmd {
TpCommand::Register(slot) => { let _ = cache.register_sequence(slot); }
TpCommand::Register(slot) => {
let _ = cache.register_sequence(slot);
}
TpCommand::Free(slot) => cache.free_sequence(slot),
TpCommand::Prefill { tokens, slot } => {
let _ = model.forward_prefill_paged(&tokens, slot, &mut cache);
}
TpCommand::Decode { tokens, positions, slots } => {
let _ = chat_decode(&model, &mut decoder, &tokens, &positions, &slots, &mut cache);
TpCommand::Decode {
tokens,
positions,
slots,
} => {
let _ = chat_decode(
&model,
&mut decoder,
&tokens,
&positions,
&slots,
&mut cache,
);
}
}
let _ = ack_tx.send(());
@@ -221,7 +282,13 @@ fn read_line_edited(prompt: &str) -> Line {
}
b => {
// UTF-8 multi-byte: read the continuation bytes for this char.
let extra = if b >= 0xF0 { 3 } else if b >= 0xE0 { 2 } else { 1 };
let extra = if b >= 0xF0 {
3
} else if b >= 0xE0 {
2
} else {
1
};
let mut bytes = vec![b];
let mut cont = [0u8; 1];
let mut ok = true;
@@ -275,7 +342,8 @@ fn main() {
if world > 1 {
assert!(
config.num_kv_heads() % world == 0,
"num_kv_heads {} not divisible by tp {world}", config.num_kv_heads()
"num_kv_heads {} not divisible by tp {world}",
config.num_kv_heads()
);
}
@@ -290,7 +358,16 @@ fn main() {
let model_dir = opts.model_dir.clone();
let config = config.clone();
thread::spawn(move || {
tp_worker_loop(rank, world, id, model_dir, config, max_seq_len, ctx_rx, ack_tx);
tp_worker_loop(
rank,
world,
id,
model_dir,
config,
max_seq_len,
ctx_rx,
ack_tx,
);
});
}
eprintln!("Loading weights (tp={world})...");
@@ -298,14 +375,37 @@ fn main() {
let weights = loader::load_model_dir(&opts.model_dir, Device::Cpu);
eprintln!("Loaded {} tensors", weights.len());
let m = if is_moe {
ChatModel::GptOss(GptOss::from_weights_tp(config.clone(), weights, 0, world, 0, Some(tp)))
ChatModel::GptOss(GptOss::from_weights_tp(
config.clone(),
weights,
0,
world,
0,
Some(tp),
))
} else {
ChatModel::Qwen3(Qwen3::from_weights_tp(config.clone(), weights, 0, world, 0, Some(tp)))
ChatModel::Qwen3(Qwen3::from_weights_tp(
config.clone(),
weights,
0,
world,
0,
Some(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 c = PagedKVCache::new_tp(&config, local_kv, total_blocks, 0, 1, max_blocks_per_seq, DType::BF16, 0);
let c = PagedKVCache::new_tp(
&config,
local_kv,
total_blocks,
0,
1,
max_blocks_per_seq,
DType::BF16,
0,
);
let h = TpHandle { cmd_txs, ack_rx };
(m, c, Some(h))
} else {
@@ -323,7 +423,10 @@ fn main() {
let tokenizer = Tokenizer::from_file(&opts.model_dir.join("tokenizer.json"));
let mut decoder = GraphedGptOssDecoder::new();
if let Some(h) = &tp_handle { h.send(TpCommand::Register(SLOT)); h.wait(); }
if let Some(h) = &tp_handle {
h.send(TpCommand::Register(SLOT));
h.wait();
}
cache.register_sequence(SLOT).expect("register chat slot");
let use_color = opts.color && io::stdout().is_terminal();
@@ -365,11 +468,8 @@ fn main() {
if is_moe {
// Harmony multi-turn: re-render the whole conversation (prior
// analysis dropped) and re-prefill into a freshly cleared slot.
let prompt = build_conversation_gpt_oss(
opts.system_prompt.as_deref(),
&moe_history,
input,
);
let prompt =
build_conversation_gpt_oss(opts.system_prompt.as_deref(), &moe_history, input);
let prompt_tokens = tokenizer.encode(&prompt);
if prompt_tokens.is_empty() {
continue;
@@ -386,8 +486,17 @@ fn main() {
print!("assistant> ");
io::stdout().flush().unwrap();
let (_finish, answer) = generate_with_paged_cache(
&model, &mut decoder, &mut cache, &tokenizer, &prompt_tokens, &opts.sampling,
max_new_tokens, use_color, &tp_handle, is_moe, opts.enable_thinking,
&model,
&mut decoder,
&mut cache,
&tokenizer,
&prompt_tokens,
&opts.sampling,
max_new_tokens,
use_color,
&tp_handle,
is_moe,
opts.enable_thinking,
);
moe_history.push((input.to_string(), answer));
println!();
@@ -436,10 +545,24 @@ fn main() {
);
match finish {
Finish::Stop { token_id } => {
append_after_stop(&model, &mut cache, &tokenizer, max_seq_len, token_id, &tp_handle);
append_after_stop(
&model,
&mut cache,
&tokenizer,
max_seq_len,
token_id,
&tp_handle,
);
}
Finish::Length => {
append_text_to_cache(&model, &mut cache, &tokenizer, max_seq_len, "<|im_end|>\n", &tp_handle);
append_text_to_cache(
&model,
&mut cache,
&tokenizer,
max_seq_len,
"<|im_end|>\n",
&tp_handle,
);
}
}
println!();
@@ -448,9 +571,15 @@ fn main() {
/// Free and re-register the single chat KV slot (clears all cached context).
fn reset_slot(cache: &mut PagedKVCache, tp: &Option<TpHandle>) {
if let Some(h) = tp { h.send(TpCommand::Free(SLOT)); h.wait(); }
if let Some(h) = tp {
h.send(TpCommand::Free(SLOT));
h.wait();
}
cache.free_sequence(SLOT);
if let Some(h) = tp { h.send(TpCommand::Register(SLOT)); h.wait(); }
if let Some(h) = tp {
h.send(TpCommand::Register(SLOT));
h.wait();
}
cache.register_sequence(SLOT).expect("register chat slot");
}
@@ -588,7 +717,15 @@ fn new_paged_cache(config: &ModelConfig, max_seq_len: usize) -> PagedKVCache {
let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
let total_blocks = (max_blocks_per_seq + 1).max(2);
// Single-slot interactive CLI: no swap pool (cpu_total_blocks = 0).
PagedKVCache::new(config, total_blocks, 0, 1, max_blocks_per_seq, DType::BF16, 0)
PagedKVCache::new(
config,
total_blocks,
0,
1,
max_blocks_per_seq,
DType::BF16,
0,
)
}
fn build_turn_prompt(
@@ -668,7 +805,10 @@ fn build_conversation_gpt_oss(
/// civil-calendar conversion (same algorithm the server uses for strftime_now).
fn today_ymd() -> String {
use std::time::{SystemTime, UNIX_EPOCH};
let secs = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs();
let secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
let z = (secs / 86400) as i64 + 719468;
let era = (if z >= 0 { z } else { z - 146096 }) / 146097;
let doe = z - era * 146097;
@@ -709,12 +849,32 @@ fn generate_with_paged_cache(
is_moe: bool,
enable_thinking: bool,
) -> (Finish, String) {
let harmony_end_id = if is_moe { tokenizer.special_token_id("<|end|>") } else { None };
let harmony_channel_id = if is_moe { tokenizer.special_token_id("<|channel|>") } else { None };
let harmony_message_id = if is_moe { tokenizer.special_token_id("<|message|>") } else { None };
let harmony_end_id = if is_moe {
tokenizer.special_token_id("<|end|>")
} else {
None
};
let harmony_channel_id = if is_moe {
tokenizer.special_token_id("<|channel|>")
} else {
None
};
let harmony_message_id = if is_moe {
tokenizer.special_token_id("<|message|>")
} else {
None
};
let harmony_special: Vec<u32> = if is_moe {
["<|channel|>", "<|start|>", "<|end|>", "<|message|>", "<|return|>"]
.iter().filter_map(|s| tokenizer.special_token_id(s)).collect()
[
"<|channel|>",
"<|start|>",
"<|end|>",
"<|message|>",
"<|return|>",
]
.iter()
.filter_map(|s| tokenizer.special_token_id(s))
.collect()
} else {
Vec::new()
};
@@ -722,18 +882,29 @@ fn generate_with_paged_cache(
// "analysis" channel is rendered as thinking (gray). After <|channel|>
// we read the channel name tokens until <|message|>.
#[derive(PartialEq, Clone, Copy)]
enum HarmonyState { Normal, ReadingChannel, InAnalysis, InFinal }
let mut hstate = if is_moe { HarmonyState::InFinal } else { HarmonyState::Normal };
enum HarmonyState {
Normal,
ReadingChannel,
InAnalysis,
InFinal,
}
let mut hstate = if is_moe {
HarmonyState::InFinal
} else {
HarmonyState::Normal
};
// Off by default. A repetition penalty over a harmony stream penalizes the
// control tokens (<|channel|>, <|message|>, <|start|>) that MUST repeat to
// open the final channel — so a non-1.0 default makes gpt-oss stop right
// after the analysis block, before emitting any answer. Opt in via the env
// var if you want it for plain (non-harmony) generation.
let rep_penalty: f32 = std::env::var("XSERV_REP_PENALTY").ok()
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()
let rep_window: usize = std::env::var("XSERV_REP_WINDOW")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(512);
let mut history: Vec<u32> = Vec::new();
@@ -747,9 +918,16 @@ fn generate_with_paged_cache(
}
};
if let Some(h) = tp { h.send(TpCommand::Prefill { tokens: prompt_tokens.to_vec(), slot: SLOT }); }
if let Some(h) = tp {
h.send(TpCommand::Prefill {
tokens: prompt_tokens.to_vec(),
slot: SLOT,
});
}
let logits = model.forward_prefill_paged(prompt_tokens, SLOT, cache);
if let Some(h) = tp { h.wait(); }
if let Some(h) = tp {
h.wait();
}
let mut next = pick(&logits, sampling, &history);
let mut decode_buffer = Vec::new();
let mut in_thinking = false;
@@ -762,9 +940,17 @@ fn generate_with_paged_cache(
for _ in 0..max_tokens {
let position = cache.seq_len(SLOT);
if let Some(h) = tp { h.send(TpCommand::Decode { tokens: vec![next], positions: vec![position], slots: vec![SLOT] }); }
if let Some(h) = tp {
h.send(TpCommand::Decode {
tokens: vec![next],
positions: vec![position],
slots: vec![SLOT],
});
}
let logits = chat_decode(model, decoder, &[next], &[position], &[SLOT], cache);
if let Some(h) = tp { h.wait(); }
if let Some(h) = tp {
h.wait();
}
if tokenizer.is_eos(next) {
print_stream_text(
&tokenizer.flush_decode_stream(&mut decode_buffer),
@@ -775,7 +961,10 @@ fn generate_with_paged_cache(
print_stream_text("\n</think>\n\n", true, use_color);
}
io::stdout().flush().unwrap();
return (Finish::Stop { token_id: next }, tokenizer.decode(&answer_ids));
return (
Finish::Stop { token_id: next },
tokenizer.decode(&answer_ids),
);
}
if harmony_end_id == Some(next) {
// <|end|> closes current segment; if in final channel, we're done
@@ -786,7 +975,10 @@ fn generate_with_paged_cache(
);
if hstate == HarmonyState::InFinal {
io::stdout().flush().unwrap();
return (Finish::Stop { token_id: next }, tokenizer.decode(&answer_ids));
return (
Finish::Stop { token_id: next },
tokenizer.decode(&answer_ids),
);
}
// Closing a thinking (analysis/commentary) channel: emit the </think>
// marker so it renders like Qwen3's thinking block.
@@ -842,7 +1034,13 @@ fn generate_with_paged_cache(
// Analysis channel = the model's reasoning. With --think, show it as a
// thinking block (gray if color); otherwise suppress it (answer only).
if show_thinking {
print_generated_token(tokenizer, next, &mut decode_buffer, &mut in_thinking, use_color);
print_generated_token(
tokenizer,
next,
&mut decode_buffer,
&mut in_thinking,
use_color,
);
io::stdout().flush().unwrap();
}
next = pick(&logits, sampling, &history);
@@ -904,9 +1102,16 @@ fn append_text_to_cache(
if tokens.is_empty() || cache.seq_len(SLOT) + tokens.len() > max_seq_len {
return;
}
if let Some(h) = tp { h.send(TpCommand::Prefill { tokens: tokens.clone(), slot: SLOT }); }
if let Some(h) = tp {
h.send(TpCommand::Prefill {
tokens: tokens.clone(),
slot: SLOT,
});
}
let _ = model.forward_prefill_paged(&tokens, SLOT, cache);
if let Some(h) = tp { h.wait(); }
if let Some(h) = tp {
h.wait();
}
}
fn print_generated_token(
@@ -952,4 +1157,3 @@ fn print_stream_text(text: &str, in_thinking: bool, use_color: bool) {
print!("{text}");
}
}