1160 lines
36 KiB
Rust
1160 lines
36 KiB
Rust
use std::io::{self, IsTerminal, Read, Write};
|
|
use std::path::PathBuf;
|
|
|
|
use std::sync::{Arc, mpsc};
|
|
use std::thread;
|
|
|
|
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;
|
|
|
|
enum ChatModel {
|
|
Qwen3(Qwen3),
|
|
GptOss(GptOss),
|
|
}
|
|
|
|
impl ChatModel {
|
|
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 {
|
|
match self {
|
|
ChatModel::Qwen3(m) => m.forward_decode_paged(tokens, positions, slots, cache),
|
|
ChatModel::GptOss(m) => m.forward_decode_paged(tokens, positions, slots, cache),
|
|
}
|
|
}
|
|
}
|
|
|
|
// TP worker infrastructure (reused from tp_engine pattern)
|
|
#[derive(Clone)]
|
|
enum TpCommand {
|
|
Register(usize),
|
|
Free(usize),
|
|
Prefill {
|
|
tokens: Vec<u32>,
|
|
slot: usize,
|
|
},
|
|
Decode {
|
|
tokens: Vec<u32>,
|
|
positions: Vec<usize>,
|
|
slots: Vec<usize>,
|
|
},
|
|
}
|
|
|
|
struct TpHandle {
|
|
cmd_txs: Vec<mpsc::Sender<TpCommand>>,
|
|
ack_rx: mpsc::Receiver<()>,
|
|
}
|
|
|
|
impl TpHandle {
|
|
fn send(&self, cmd: TpCommand) {
|
|
for tx in &self.cmd_txs {
|
|
tx.send(cmd.clone()).ok();
|
|
}
|
|
}
|
|
fn wait(&self) {
|
|
for _ in 0..self.cmd_txs.len() {
|
|
self.ack_rx.recv().ok();
|
|
}
|
|
}
|
|
}
|
|
|
|
fn tp_worker_loop(
|
|
rank: usize,
|
|
world: usize,
|
|
id: xserv_distributed::UniqueId,
|
|
model_dir: std::path::PathBuf,
|
|
config: ModelConfig,
|
|
max_seq_len: usize,
|
|
cmd_rx: mpsc::Receiver<TpCommand>,
|
|
ack_tx: mpsc::Sender<()>,
|
|
) {
|
|
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),
|
|
))
|
|
} else {
|
|
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,
|
|
);
|
|
let mut decoder = GraphedGptOssDecoder::new();
|
|
while let Ok(cmd) = cmd_rx.recv() {
|
|
match cmd {
|
|
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,
|
|
);
|
|
}
|
|
}
|
|
let _ = ack_tx.send(());
|
|
}
|
|
}
|
|
|
|
const SLOT: usize = 0;
|
|
|
|
struct CliOptions {
|
|
model_dir: PathBuf,
|
|
max_tokens: usize,
|
|
max_seq_len: usize,
|
|
tp: usize,
|
|
sampling: SamplingParams,
|
|
system_prompt: Option<String>,
|
|
enable_thinking: bool,
|
|
color: bool,
|
|
}
|
|
|
|
enum Finish {
|
|
Stop { token_id: u32 },
|
|
Length,
|
|
}
|
|
|
|
enum Line {
|
|
Text(String),
|
|
Eof,
|
|
}
|
|
|
|
/// RAII terminal raw-mode guard. Disables canonical mode + echo (keeps output
|
|
/// post-processing and signals), so we read keystrokes ourselves and edit the
|
|
/// line UTF-8-aware. Restores the original termios on drop.
|
|
struct RawMode {
|
|
orig: libc::termios,
|
|
}
|
|
|
|
impl RawMode {
|
|
fn enable() -> Option<Self> {
|
|
unsafe {
|
|
let mut orig: libc::termios = std::mem::zeroed();
|
|
if libc::tcgetattr(libc::STDIN_FILENO, &mut orig) != 0 {
|
|
return None;
|
|
}
|
|
let mut raw = orig;
|
|
raw.c_lflag &= !(libc::ICANON | libc::ECHO);
|
|
raw.c_cc[libc::VMIN as usize] = 1;
|
|
raw.c_cc[libc::VTIME as usize] = 0;
|
|
if libc::tcsetattr(libc::STDIN_FILENO, libc::TCSANOW, &raw) != 0 {
|
|
return None;
|
|
}
|
|
Some(RawMode { orig })
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Drop for RawMode {
|
|
fn drop(&mut self) {
|
|
unsafe {
|
|
libc::tcsetattr(libc::STDIN_FILENO, libc::TCSANOW, &self.orig);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Read one line with UTF-8/CJK-aware editing. In a TTY this enters raw mode and
|
|
/// handles keystrokes so Backspace deletes a whole character (not a byte), and
|
|
/// multi-byte input (汉字/日本語/한글) renders correctly. Non-TTY (piped) input
|
|
/// falls back to a plain cooked read.
|
|
fn read_line_edited(prompt: &str) -> Line {
|
|
let cooked = || -> Line {
|
|
print!("{prompt}");
|
|
io::stdout().flush().ok();
|
|
let mut s = String::new();
|
|
match io::stdin().read_line(&mut s) {
|
|
Ok(0) | Err(_) => Line::Eof,
|
|
Ok(_) => Line::Text(s),
|
|
}
|
|
};
|
|
|
|
if !io::stdin().is_terminal() {
|
|
return cooked();
|
|
}
|
|
let Some(_raw) = RawMode::enable() else {
|
|
return cooked();
|
|
};
|
|
|
|
// Single-line editor: on every edit, rewrite the whole line so the terminal
|
|
// renders correct (incl. double-width CJK) glyphs; \x1b[K clears leftovers.
|
|
let redraw = |buf: &str| {
|
|
print!("\r{prompt}{buf}\x1b[K");
|
|
io::stdout().flush().ok();
|
|
};
|
|
|
|
let mut buf = String::new();
|
|
redraw(&buf);
|
|
let mut stdin = io::stdin().lock();
|
|
let mut byte = [0u8; 1];
|
|
|
|
loop {
|
|
if stdin.read(&mut byte).unwrap_or(0) == 0 {
|
|
// EOF on the stream.
|
|
if buf.is_empty() {
|
|
return Line::Eof;
|
|
}
|
|
break;
|
|
}
|
|
match byte[0] {
|
|
b'\r' | b'\n' => {
|
|
println!();
|
|
break;
|
|
}
|
|
0x7f | 0x08 => {
|
|
// Backspace: drop one whole char (String::pop is char-aware).
|
|
buf.pop();
|
|
redraw(&buf);
|
|
}
|
|
0x04 => {
|
|
// Ctrl-D: EOF only when the line is empty.
|
|
if buf.is_empty() {
|
|
return Line::Eof;
|
|
}
|
|
}
|
|
0x1b => {
|
|
// Escape sequence (arrows, etc.): consume and ignore the 2 bytes
|
|
// of a typical CSI sequence so they don't land in the buffer.
|
|
let mut seq = [0u8; 2];
|
|
let _ = stdin.read(&mut seq);
|
|
}
|
|
b if b < 0x20 => { /* other control bytes: ignore */ }
|
|
b if b < 0x80 => {
|
|
buf.push(b as char);
|
|
redraw(&buf);
|
|
}
|
|
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 mut bytes = vec![b];
|
|
let mut cont = [0u8; 1];
|
|
let mut ok = true;
|
|
for _ in 0..extra {
|
|
if stdin.read(&mut cont).unwrap_or(0) == 0 {
|
|
ok = false;
|
|
break;
|
|
}
|
|
bytes.push(cont[0]);
|
|
}
|
|
if ok {
|
|
if let Ok(s) = std::str::from_utf8(&bytes) {
|
|
buf.push_str(s);
|
|
redraw(&buf);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Line::Text(buf)
|
|
}
|
|
|
|
fn main() {
|
|
let opts = parse_args();
|
|
|
|
xserv_cuda::device::set_device(0).unwrap();
|
|
let info = xserv_cuda::device::device_info(0).unwrap();
|
|
eprintln!(
|
|
"GPU: {} ({} MB free)",
|
|
info.name,
|
|
info.free_memory / 1024 / 1024
|
|
);
|
|
|
|
let config = ModelConfig::from_file(&opts.model_dir.join("config.json"));
|
|
let model_type = config.model_type.as_deref().unwrap_or("unknown");
|
|
let is_moe = config.is_moe();
|
|
|
|
let max_seq_len = opts.max_seq_len.min(config.max_seq_len()).max(1);
|
|
eprintln!(
|
|
"Model: {model_type}{}, layers={}, hidden={}, heads={}/{} kv, vocab={}, max_seq_len={}",
|
|
if is_moe { " (MoE)" } else { "" },
|
|
config.num_layers(),
|
|
config.hidden(),
|
|
config.num_heads(),
|
|
config.num_kv_heads(),
|
|
config.vocab_size,
|
|
max_seq_len
|
|
);
|
|
|
|
let world = opts.tp;
|
|
if world > 1 {
|
|
assert!(
|
|
config.num_kv_heads() % world == 0,
|
|
"num_kv_heads {} not divisible by tp {world}",
|
|
config.num_kv_heads()
|
|
);
|
|
}
|
|
|
|
let (model, mut cache, tp_handle) = if world > 1 {
|
|
let id = xserv_distributed::get_unique_id();
|
|
let (ack_tx, ack_rx) = mpsc::channel::<()>();
|
|
let mut cmd_txs = Vec::new();
|
|
for rank in 1..world {
|
|
let (ctx_tx, ctx_rx) = mpsc::channel::<TpCommand>();
|
|
cmd_txs.push(ctx_tx);
|
|
let ack_tx = ack_tx.clone();
|
|
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,
|
|
);
|
|
});
|
|
}
|
|
eprintln!("Loading weights (tp={world})...");
|
|
let tp = Arc::new(xserv_distributed::TpContext::init(0, world, id, 0));
|
|
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),
|
|
))
|
|
} else {
|
|
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 h = TpHandle { cmd_txs, ack_rx };
|
|
(m, c, Some(h))
|
|
} else {
|
|
eprintln!("Loading weights...");
|
|
let weights = loader::load_model_dir(&opts.model_dir, Device::Cuda(0));
|
|
eprintln!("Loaded {} tensors", weights.len());
|
|
let m = if is_moe {
|
|
ChatModel::GptOss(GptOss::from_weights(config.clone(), weights))
|
|
} else {
|
|
ChatModel::Qwen3(Qwen3::from_weights(config.clone(), weights))
|
|
};
|
|
let c = new_paged_cache(&config, max_seq_len);
|
|
(m, c, None)
|
|
};
|
|
|
|
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();
|
|
}
|
|
cache.register_sequence(SLOT).expect("register chat slot");
|
|
let use_color = opts.color && io::stdout().is_terminal();
|
|
|
|
eprintln!("Ready (paged KV cache, tp={world}).");
|
|
eprintln!("Commands: /exit, /quit, /clear\n");
|
|
|
|
// gpt-oss multi-turn history of (user, assistant-final) text. Harmony
|
|
// requires re-rendering the conversation each turn with prior analysis
|
|
// dropped, so the moe path re-prefills from this rather than reusing an
|
|
// incremental KV cache (which would accumulate CoT + <|return|> and collapse
|
|
// at longer context). Qwen3 ignores this and keeps the incremental cache.
|
|
let mut moe_history: Vec<(String, String)> = Vec::new();
|
|
|
|
loop {
|
|
let line = match read_line_edited("user> ") {
|
|
Line::Eof => break,
|
|
Line::Text(s) => s,
|
|
};
|
|
let input = line.trim();
|
|
if input.is_empty() {
|
|
continue;
|
|
}
|
|
|
|
match input {
|
|
"/exit" | "/quit" | "exit" | "quit" => break,
|
|
"/clear" => {
|
|
reset_slot(&mut cache, &tp_handle);
|
|
moe_history.clear();
|
|
eprintln!("history and KV cache cleared");
|
|
continue;
|
|
}
|
|
"/help" => {
|
|
print_help();
|
|
continue;
|
|
}
|
|
_ => {}
|
|
}
|
|
|
|
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_tokens = tokenizer.encode(&prompt);
|
|
if prompt_tokens.is_empty() {
|
|
continue;
|
|
}
|
|
if prompt_tokens.len() >= max_seq_len {
|
|
eprintln!(
|
|
"context full: conversation needs {} tokens >= max_seq_len {max_seq_len}; use /clear",
|
|
prompt_tokens.len()
|
|
);
|
|
continue;
|
|
}
|
|
let max_new_tokens = opts.max_tokens.min(max_seq_len - prompt_tokens.len());
|
|
reset_slot(&mut cache, &tp_handle);
|
|
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,
|
|
);
|
|
moe_history.push((input.to_string(), answer));
|
|
println!();
|
|
continue;
|
|
}
|
|
|
|
// Qwen3: incremental KV cache — only the new turn is prefilled and the
|
|
// assistant's tokens stay cached for the next turn.
|
|
let include_system = cache.seq_len(SLOT) == 0;
|
|
let prompt = build_turn_prompt(
|
|
opts.system_prompt.as_deref(),
|
|
include_system,
|
|
input,
|
|
opts.enable_thinking,
|
|
);
|
|
let prompt_tokens = tokenizer.encode(&prompt);
|
|
if prompt_tokens.is_empty() {
|
|
continue;
|
|
}
|
|
|
|
let used = cache.seq_len(SLOT);
|
|
let remaining = max_seq_len.saturating_sub(used);
|
|
if prompt_tokens.len() >= remaining {
|
|
eprintln!(
|
|
"context full: {used}/{max_seq_len} tokens used, new turn needs {} tokens; use /clear",
|
|
prompt_tokens.len()
|
|
);
|
|
continue;
|
|
}
|
|
let max_new_tokens = opts.max_tokens.min(remaining - prompt_tokens.len());
|
|
|
|
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,
|
|
);
|
|
match finish {
|
|
Finish::Stop { token_id } => {
|
|
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,
|
|
);
|
|
}
|
|
}
|
|
println!();
|
|
}
|
|
}
|
|
|
|
/// 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();
|
|
}
|
|
cache.free_sequence(SLOT);
|
|
if let Some(h) = tp {
|
|
h.send(TpCommand::Register(SLOT));
|
|
h.wait();
|
|
}
|
|
cache.register_sequence(SLOT).expect("register chat slot");
|
|
}
|
|
|
|
fn parse_args() -> CliOptions {
|
|
let args: Vec<String> = std::env::args().skip(1).collect();
|
|
if args.is_empty() || args.iter().any(|a| a == "--help" || a == "-h") {
|
|
print_usage_and_exit(0);
|
|
}
|
|
|
|
let mut model_dir = None;
|
|
let mut max_tokens = 256usize;
|
|
let mut max_seq_len = 2048usize;
|
|
let mut tp = 1usize;
|
|
let mut temperature = 0.0f32;
|
|
let mut top_k = 0usize;
|
|
let mut top_p = 1.0f32;
|
|
let mut system_prompt = None;
|
|
let mut enable_thinking = false;
|
|
let mut color = true;
|
|
|
|
let mut i = 0;
|
|
while i < args.len() {
|
|
match args[i].as_str() {
|
|
"-m" | "--model" => {
|
|
i += 1;
|
|
model_dir = args.get(i).map(PathBuf::from);
|
|
}
|
|
"--max-tokens" => {
|
|
i += 1;
|
|
max_tokens = parse_value(&args, i, "--max-tokens");
|
|
}
|
|
"--max-seq-len" => {
|
|
i += 1;
|
|
max_seq_len = parse_value(&args, i, "--max-seq-len");
|
|
}
|
|
"--tp" => {
|
|
i += 1;
|
|
tp = parse_value(&args, i, "--tp");
|
|
}
|
|
"--temperature" => {
|
|
i += 1;
|
|
temperature = parse_value(&args, i, "--temperature");
|
|
}
|
|
"--top-k" => {
|
|
i += 1;
|
|
top_k = parse_value(&args, i, "--top-k");
|
|
}
|
|
"--top-p" => {
|
|
i += 1;
|
|
top_p = parse_value(&args, i, "--top-p");
|
|
}
|
|
"--system" => {
|
|
i += 1;
|
|
system_prompt = args.get(i).cloned();
|
|
if system_prompt.is_none() {
|
|
eprintln!("missing value for --system");
|
|
std::process::exit(2);
|
|
}
|
|
}
|
|
"--think" => {
|
|
enable_thinking = true;
|
|
}
|
|
"--no-color" => {
|
|
color = false;
|
|
}
|
|
arg if arg.starts_with('-') => {
|
|
eprintln!("unknown option: {arg}");
|
|
print_usage_and_exit(2);
|
|
}
|
|
arg => {
|
|
if model_dir.is_some() {
|
|
eprintln!("unexpected extra argument: {arg}");
|
|
print_usage_and_exit(2);
|
|
}
|
|
model_dir = Some(PathBuf::from(arg));
|
|
}
|
|
}
|
|
i += 1;
|
|
}
|
|
|
|
CliOptions {
|
|
model_dir: model_dir.unwrap_or_else(|| {
|
|
eprintln!("missing model directory");
|
|
print_usage_and_exit(2);
|
|
}),
|
|
max_tokens: max_tokens.max(1),
|
|
max_seq_len: max_seq_len.max(1),
|
|
tp: tp.max(1),
|
|
sampling: SamplingParams {
|
|
temperature,
|
|
top_k,
|
|
top_p,
|
|
},
|
|
system_prompt,
|
|
enable_thinking,
|
|
color,
|
|
}
|
|
}
|
|
|
|
fn parse_value<T: std::str::FromStr>(args: &[String], i: usize, name: &str) -> T {
|
|
args.get(i).and_then(|s| s.parse().ok()).unwrap_or_else(|| {
|
|
eprintln!("invalid or missing value for {name}");
|
|
std::process::exit(2);
|
|
})
|
|
}
|
|
|
|
fn print_usage_and_exit(code: i32) -> ! {
|
|
eprintln!(
|
|
"Usage: xserv-chat <model-dir> [options]\n\
|
|
\n\
|
|
Options:\n\
|
|
\t-m, --model DIR Model directory\n\
|
|
\t--max-tokens N Max generated tokens per turn (default: 256)\n\
|
|
\t--max-seq-len N Persistent KV context length (default: 2048)\n\
|
|
\t--tp N Tensor parallelism degree (default: 1)\n\
|
|
\t--temperature F Sampling temperature, 0 = greedy (default: 0)\n\
|
|
\t--top-k N Top-k sampling, 0 = disabled (default: 0)\n\
|
|
\t--top-p F Top-p sampling (default: 1.0)\n\
|
|
\t--system TEXT System prompt for the first turn after start or /clear\n\
|
|
\t--think Let Qwen3 emit thinking; rendered gray on terminals\n\
|
|
\t--no-color Disable ANSI color for thinking output\n\
|
|
\t-h, --help Show this help"
|
|
);
|
|
std::process::exit(code);
|
|
}
|
|
|
|
fn print_help() {
|
|
eprintln!("Commands:");
|
|
eprintln!(" /clear clear chat history and free/recreate the paged KV slot");
|
|
eprintln!(" /exit quit");
|
|
eprintln!(" /quit quit");
|
|
}
|
|
|
|
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,
|
|
)
|
|
}
|
|
|
|
fn build_turn_prompt(
|
|
system: Option<&str>,
|
|
include_system: bool,
|
|
user_input: &str,
|
|
enable_thinking: bool,
|
|
) -> String {
|
|
let mut prompt = String::new();
|
|
if include_system {
|
|
if let Some(system) = system {
|
|
if !system.trim().is_empty() {
|
|
prompt.push_str("<|im_start|>system\n");
|
|
prompt.push_str(system.trim());
|
|
prompt.push_str("<|im_end|>\n");
|
|
}
|
|
}
|
|
}
|
|
prompt.push_str("<|im_start|>user\n");
|
|
prompt.push_str(user_input);
|
|
prompt.push_str("<|im_end|>\n");
|
|
prompt.push_str("<|im_start|>assistant\n");
|
|
if !enable_thinking {
|
|
prompt.push_str("<think>\n\n</think>\n\n");
|
|
}
|
|
prompt
|
|
}
|
|
|
|
/// Render the full gpt-oss harmony conversation for re-prefill. gpt-oss was
|
|
/// trained on this exact system-message structure (identity / knowledge cutoff
|
|
/// / current date / Reasoning level / channels — see the model's
|
|
/// chat_template.jinja `build_system_message`). A hand-rolled substitute puts
|
|
/// the model out of distribution and destabilizes channel selection.
|
|
///
|
|
/// Harmony multi-turn drops prior chain-of-thought: past assistant messages are
|
|
/// rendered as completed `final` channels ending in `<|end|>` (not the
|
|
/// `<|return|>` stop token). Keeping the analysis + `<|return|>` of every turn
|
|
/// in context — as an incremental KV cache does — is out of distribution and
|
|
/// makes the model collapse at longer context. "Reasoning: low" keeps the
|
|
/// analysis channel short for an interactive chat.
|
|
fn build_conversation_gpt_oss(
|
|
system: Option<&str>,
|
|
history: &[(String, String)],
|
|
current_user: &str,
|
|
) -> String {
|
|
let mut prompt = String::new();
|
|
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", today_ymd()));
|
|
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|>");
|
|
if let Some(sys) = system {
|
|
if !sys.trim().is_empty() {
|
|
prompt.push_str("<|start|>developer<|message|># Instructions\n\n");
|
|
prompt.push_str(sys.trim());
|
|
prompt.push_str("<|end|>");
|
|
}
|
|
}
|
|
for (user, assistant) in history {
|
|
prompt.push_str("<|start|>user<|message|>");
|
|
prompt.push_str(user);
|
|
prompt.push_str("<|end|>");
|
|
prompt.push_str("<|start|>assistant<|channel|>final<|message|>");
|
|
prompt.push_str(assistant.trim());
|
|
prompt.push_str("<|end|>");
|
|
}
|
|
prompt.push_str("<|start|>user<|message|>");
|
|
prompt.push_str(current_user);
|
|
prompt.push_str("<|end|>");
|
|
prompt.push_str("<|start|>assistant");
|
|
prompt
|
|
}
|
|
|
|
/// Current UTC date as "YYYY-MM-DD" for the harmony system message. Rata Die
|
|
/// 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 z = (secs / 86400) as i64 + 719468;
|
|
let era = (if z >= 0 { z } else { z - 146096 }) / 146097;
|
|
let doe = z - era * 146097;
|
|
let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
|
|
let y = yoe + era * 400;
|
|
let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
|
|
let mp = (5 * doy + 2) / 153;
|
|
let d = doy - (153 * mp + 2) / 5 + 1;
|
|
let m = if mp < 10 { mp + 3 } else { mp - 9 };
|
|
let y = if m <= 2 { y + 1 } else { y };
|
|
format!("{y:04}-{m:02}-{d:02}")
|
|
}
|
|
|
|
fn chat_decode(
|
|
model: &ChatModel,
|
|
decoder: &mut GraphedGptOssDecoder,
|
|
tokens: &[u32],
|
|
positions: &[usize],
|
|
slots: &[usize],
|
|
cache: &mut PagedKVCache,
|
|
) -> xserv_tensor::Tensor {
|
|
match model {
|
|
ChatModel::GptOss(m) => decoder.decode(m, tokens, positions, slots, cache),
|
|
ChatModel::Qwen3(_) => model.forward_decode_paged(tokens, positions, slots, cache),
|
|
}
|
|
}
|
|
|
|
fn generate_with_paged_cache(
|
|
model: &ChatModel,
|
|
decoder: &mut GraphedGptOssDecoder,
|
|
cache: &mut PagedKVCache,
|
|
tokenizer: &Tokenizer,
|
|
prompt_tokens: &[u32],
|
|
sampling: &SamplingParams,
|
|
max_tokens: usize,
|
|
use_color: bool,
|
|
tp: &Option<TpHandle>,
|
|
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_special: Vec<u32> = if is_moe {
|
|
[
|
|
"<|channel|>",
|
|
"<|start|>",
|
|
"<|end|>",
|
|
"<|message|>",
|
|
"<|return|>",
|
|
]
|
|
.iter()
|
|
.filter_map(|s| tokenizer.special_token_id(s))
|
|
.collect()
|
|
} else {
|
|
Vec::new()
|
|
};
|
|
// Harmony channel state: "final" channel text is printed normally,
|
|
// "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
|
|
};
|
|
|
|
// 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()
|
|
.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(512);
|
|
let mut history: Vec<u32> = Vec::new();
|
|
|
|
let pick = |logits: &xserv_tensor::Tensor, sp: &SamplingParams, hist: &[u32]| -> u32 {
|
|
if rep_penalty > 1.0 && sp.temperature == 0.0 {
|
|
let start = hist.len().saturating_sub(rep_window);
|
|
sample_greedy_penalized(logits, &hist[start..], rep_penalty)
|
|
} else {
|
|
sample(logits, sp)
|
|
}
|
|
};
|
|
|
|
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();
|
|
}
|
|
let mut next = pick(&logits, sampling, &history);
|
|
let mut decode_buffer = Vec::new();
|
|
let mut in_thinking = false;
|
|
let show_thinking = is_moe && enable_thinking;
|
|
// Visible answer tokens, returned for multi-turn history. For moe this is
|
|
// the final-channel content only (analysis is suppressed/gray); for Qwen3
|
|
// it is everything printed. The caller decodes these into the assistant
|
|
// message it re-renders into the next prompt.
|
|
let mut answer_ids: Vec<u32> = Vec::new();
|
|
|
|
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],
|
|
});
|
|
}
|
|
let logits = chat_decode(model, decoder, &[next], &[position], &[SLOT], cache);
|
|
if let Some(h) = tp {
|
|
h.wait();
|
|
}
|
|
if tokenizer.is_eos(next) {
|
|
print_stream_text(
|
|
&tokenizer.flush_decode_stream(&mut decode_buffer),
|
|
in_thinking,
|
|
use_color,
|
|
);
|
|
if show_thinking && in_thinking {
|
|
print_stream_text("\n</think>\n\n", true, use_color);
|
|
}
|
|
io::stdout().flush().unwrap();
|
|
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
|
|
print_stream_text(
|
|
&tokenizer.flush_decode_stream(&mut decode_buffer),
|
|
in_thinking,
|
|
use_color,
|
|
);
|
|
if hstate == HarmonyState::InFinal {
|
|
io::stdout().flush().unwrap();
|
|
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.
|
|
if show_thinking && hstate == HarmonyState::InAnalysis {
|
|
print_stream_text("\n</think>\n\n", true, use_color);
|
|
in_thinking = false;
|
|
}
|
|
hstate = HarmonyState::Normal;
|
|
next = pick(&logits, sampling, &history);
|
|
continue;
|
|
}
|
|
|
|
history.push(next);
|
|
|
|
// Harmony channel routing state machine
|
|
if harmony_channel_id == Some(next) {
|
|
decode_buffer.clear();
|
|
hstate = HarmonyState::ReadingChannel;
|
|
next = pick(&logits, sampling, &history);
|
|
continue;
|
|
}
|
|
if harmony_message_id == Some(next) {
|
|
if hstate == HarmonyState::ReadingChannel {
|
|
// Channel name was accumulated but we don't need to parse it —
|
|
// we just check via the channel_name buffer below
|
|
}
|
|
decode_buffer.clear();
|
|
next = pick(&logits, sampling, &history);
|
|
continue;
|
|
}
|
|
if hstate == HarmonyState::ReadingChannel {
|
|
// Reading channel name tokens (e.g. "final", "analysis")
|
|
let tok_text = tokenizer.decode(&[next]);
|
|
if tok_text.contains("final") {
|
|
hstate = HarmonyState::InFinal;
|
|
in_thinking = false;
|
|
} else {
|
|
hstate = HarmonyState::InAnalysis;
|
|
// Open a Qwen3-style thinking block for the analysis channel.
|
|
if show_thinking {
|
|
print_stream_text("<think>\n", true, use_color);
|
|
in_thinking = true;
|
|
}
|
|
}
|
|
next = pick(&logits, sampling, &history);
|
|
continue;
|
|
}
|
|
if harmony_special.contains(&next) {
|
|
next = pick(&logits, sampling, &history);
|
|
continue;
|
|
}
|
|
if hstate == HarmonyState::InAnalysis {
|
|
// 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,
|
|
);
|
|
io::stdout().flush().unwrap();
|
|
}
|
|
next = pick(&logits, sampling, &history);
|
|
continue;
|
|
}
|
|
if is_moe && hstate != HarmonyState::InFinal {
|
|
// Between harmony messages (after a channel's <|end|>, before the
|
|
// next <|channel|>): the model emits a role header like "assistant".
|
|
// That's structural, not user-visible content — suppress it. Only
|
|
// for moe/harmony; non-moe (Qwen3) stays in Normal and prints here.
|
|
next = pick(&logits, sampling, &history);
|
|
continue;
|
|
}
|
|
|
|
answer_ids.push(next);
|
|
print_generated_token(
|
|
tokenizer,
|
|
next,
|
|
&mut decode_buffer,
|
|
&mut in_thinking,
|
|
use_color,
|
|
);
|
|
io::stdout().flush().unwrap();
|
|
next = pick(&logits, sampling, &history);
|
|
}
|
|
|
|
print_stream_text(
|
|
&tokenizer.flush_decode_stream(&mut decode_buffer),
|
|
in_thinking,
|
|
use_color,
|
|
);
|
|
if show_thinking && in_thinking {
|
|
print_stream_text("\n</think>\n\n", true, use_color);
|
|
}
|
|
io::stdout().flush().unwrap();
|
|
(Finish::Length, tokenizer.decode(&answer_ids))
|
|
}
|
|
|
|
fn append_after_stop(
|
|
model: &ChatModel,
|
|
cache: &mut PagedKVCache,
|
|
tokenizer: &Tokenizer,
|
|
max_seq_len: usize,
|
|
_stop_token_id: u32,
|
|
tp: &Option<TpHandle>,
|
|
) {
|
|
append_text_to_cache(model, cache, tokenizer, max_seq_len, "\n", tp);
|
|
}
|
|
|
|
fn append_text_to_cache(
|
|
model: &ChatModel,
|
|
cache: &mut PagedKVCache,
|
|
tokenizer: &Tokenizer,
|
|
max_seq_len: usize,
|
|
text: &str,
|
|
tp: &Option<TpHandle>,
|
|
) {
|
|
let tokens = tokenizer.encode(text);
|
|
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,
|
|
});
|
|
}
|
|
let _ = model.forward_prefill_paged(&tokens, SLOT, cache);
|
|
if let Some(h) = tp {
|
|
h.wait();
|
|
}
|
|
}
|
|
|
|
fn print_generated_token(
|
|
tokenizer: &Tokenizer,
|
|
token_id: u32,
|
|
decode_buffer: &mut Vec<u8>,
|
|
in_thinking: &mut bool,
|
|
use_color: bool,
|
|
) {
|
|
if tokenizer.special_token_id("<think>") == Some(token_id) {
|
|
print_stream_text(
|
|
&tokenizer.flush_decode_stream(decode_buffer),
|
|
*in_thinking,
|
|
use_color,
|
|
);
|
|
*in_thinking = true;
|
|
print_stream_text("<think>", true, use_color);
|
|
return;
|
|
}
|
|
|
|
if tokenizer.special_token_id("</think>") == Some(token_id) {
|
|
print_stream_text(
|
|
&tokenizer.flush_decode_stream(decode_buffer),
|
|
*in_thinking,
|
|
use_color,
|
|
);
|
|
print_stream_text("</think>", true, use_color);
|
|
*in_thinking = false;
|
|
return;
|
|
}
|
|
|
|
let text = tokenizer.decode_token_stream(token_id, decode_buffer);
|
|
print_stream_text(&text, *in_thinking, use_color);
|
|
}
|
|
|
|
fn print_stream_text(text: &str, in_thinking: bool, use_color: bool) {
|
|
if text.is_empty() {
|
|
return;
|
|
}
|
|
if in_thinking && use_color {
|
|
print!("\x1b[90m{text}\x1b[0m");
|
|
} else {
|
|
print!("{text}");
|
|
}
|
|
}
|