phase19: MoE support — gpt-oss-20b end-to-end inference with TP=2
Add Mixture-of-Experts support for the gpt-oss-20b model (20.9B params, 32 experts × top-4 routing). Key additions: - ModelConfig: MoE fields (num_local_experts, layer_types, sliding_window, attention_bias, explicit head_dim, rope_scaling, swiglu_limit) - YaRN RoPE: RopeCache::new_yarn() with correct frequency interpolation and attention_scaling = 0.1*ln(factor)+1 - Custom GLU kernel: gpt_oss_glu_bf16 (clamped sigmoid gate activation) - Paged attention with sinks + sliding window kernel variant - GptOss model struct with expert-parallel TP (split 32 experts across ranks) - bench-gpt-oss binary for TP inference benchmarking Verified on dash5 with 2x RTX 5090: 63.6 tok/s decode, ~160ms TTFT. Model generates topically-coherent output (needs chat template for quality). Known issues: - Custom GEMV kernel produces NaN with small N (workaround: pad to M=2) - Prefill doesn't use attention sinks (uses standard flash attention) - Output quality requires chat template formatting Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
231
crates/xserv-model/src/bin/bench-gpt-oss.rs
Normal file
231
crates/xserv-model/src/bin/bench-gpt-oss.rs
Normal file
@@ -0,0 +1,231 @@
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use xserv_distributed::{TpContext, UniqueId, get_unique_id};
|
||||
use xserv_model::{loader, GptOss, ModelConfig, PagedKVCache, BLOCK_SIZE};
|
||||
use xserv_tensor::{DType, Device};
|
||||
use xserv_tokenizer::Tokenizer;
|
||||
|
||||
fn main() {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
if args.len() < 2 {
|
||||
eprintln!("Usage: bench-gpt-oss <model-dir> [--max-tokens N] [--tp N]");
|
||||
std::process::exit(1);
|
||||
}
|
||||
|
||||
let model_dir = PathBuf::from(&args[1]);
|
||||
let max_tokens: usize = get_arg(&args, "--max-tokens").unwrap_or(32);
|
||||
let world: usize = get_arg(&args, "--tp").unwrap_or(2);
|
||||
|
||||
let config = ModelConfig::from_file(&model_dir.join("config.json"));
|
||||
let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
|
||||
|
||||
eprintln!(
|
||||
"gpt-oss-20b: layers={}, hidden={}, heads={}/{} kv, experts={}, top_k={}, vocab={}",
|
||||
config.num_layers(), config.hidden(), config.num_heads(),
|
||||
config.num_kv_heads(), config.num_experts(), config.experts_per_token(),
|
||||
config.vocab_size
|
||||
);
|
||||
eprintln!("TP world={world}, max_tokens={max_tokens}");
|
||||
|
||||
let max_seq_len: usize = 2048;
|
||||
let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
|
||||
|
||||
// TP setup
|
||||
let uid = get_unique_id();
|
||||
let local_kv = config.num_kv_heads() / world;
|
||||
|
||||
// Spawn worker threads for ranks 1..world
|
||||
let mut worker_handles = Vec::new();
|
||||
let mut worker_txs = Vec::new();
|
||||
for rank in 1..world {
|
||||
let (tx, rx) = std::sync::mpsc::channel::<WorkerCmd>();
|
||||
let (ack_tx, ack_rx) = std::sync::mpsc::channel::<()>();
|
||||
let cfg = config.clone();
|
||||
let md = model_dir.clone();
|
||||
let uid_copy = uid;
|
||||
worker_handles.push((
|
||||
std::thread::spawn(move || {
|
||||
worker_loop(rank, world, uid_copy, md, cfg, max_seq_len, rx, ack_tx);
|
||||
}),
|
||||
ack_rx,
|
||||
));
|
||||
worker_txs.push(tx);
|
||||
}
|
||||
|
||||
// Rank 0 setup
|
||||
xserv_cuda::device::set_device(0).unwrap();
|
||||
let tp0 = Arc::new(TpContext::init(0, world, uid, 0));
|
||||
eprintln!("[rank 0] Loading weights...");
|
||||
let weights = loader::load_model_dir(&model_dir, Device::Cpu);
|
||||
eprintln!("[rank 0] Loaded {} tensors, building model...", weights.len());
|
||||
let model = GptOss::from_weights_tp(config.clone(), weights, 0, world, 0, Some(tp0));
|
||||
let total_blocks = max_blocks_per_seq + 64;
|
||||
let mut cache = PagedKVCache::new_tp(
|
||||
&config, local_kv, total_blocks, 0, 4, max_blocks_per_seq, DType::BF16, 0,
|
||||
);
|
||||
eprintln!("[rank 0] Ready.");
|
||||
|
||||
// Prompt
|
||||
let prompt = "What is the meaning of life?";
|
||||
let token_ids = tokenizer.encode(prompt);
|
||||
eprintln!("Prompt ({} tokens): {prompt}", token_ids.len());
|
||||
|
||||
// Register sequence
|
||||
let slot = 0;
|
||||
cache.register_sequence(slot).unwrap();
|
||||
broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Register(slot));
|
||||
|
||||
// Prefill
|
||||
let t0 = Instant::now();
|
||||
broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Prefill {
|
||||
tokens: token_ids.clone(), slot,
|
||||
});
|
||||
let logits = model.forward_prefill_paged(&token_ids, slot, &mut cache);
|
||||
wait_workers(&worker_handles);
|
||||
let ttft = t0.elapsed();
|
||||
|
||||
let mut next = sample_greedy_last(&logits);
|
||||
let mut output_tokens = vec![next];
|
||||
|
||||
eprintln!("TTFT: {:.1}ms", ttft.as_secs_f64() * 1000.0);
|
||||
print!("{prompt}");
|
||||
|
||||
// Decode
|
||||
let decode_start = Instant::now();
|
||||
for _ in 1..max_tokens {
|
||||
let text = tokenizer.decode(&[next]);
|
||||
print!("{text}");
|
||||
|
||||
if tokenizer.eos_token_id() == Some(next) { break; }
|
||||
|
||||
let pos = cache.seq_len(slot);
|
||||
broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Decode {
|
||||
tokens: vec![next], positions: vec![pos], slots: vec![slot],
|
||||
});
|
||||
let logits = model.forward_decode_paged(&[next], &[pos], &[slot], &mut cache);
|
||||
wait_workers(&worker_handles);
|
||||
|
||||
next = sample_greedy_last(&logits);
|
||||
output_tokens.push(next);
|
||||
}
|
||||
let decode_elapsed = decode_start.elapsed();
|
||||
println!();
|
||||
|
||||
let gen_tokens = output_tokens.len();
|
||||
let full_text = tokenizer.decode(&output_tokens);
|
||||
eprintln!("\nGenerated text: {full_text}");
|
||||
eprintln!("Token IDs: {:?}", &output_tokens[..output_tokens.len().min(20)]);
|
||||
let tpot = if gen_tokens > 1 {
|
||||
decode_elapsed.as_secs_f64() * 1000.0 / (gen_tokens - 1) as f64
|
||||
} else { 0.0 };
|
||||
let tok_s = if gen_tokens > 1 {
|
||||
(gen_tokens - 1) as f64 / decode_elapsed.as_secs_f64()
|
||||
} else { 0.0 };
|
||||
|
||||
eprintln!("\n--- Performance ---");
|
||||
eprintln!("Generated: {} tokens", gen_tokens);
|
||||
eprintln!("TTFT: {:.1}ms", ttft.as_secs_f64() * 1000.0);
|
||||
eprintln!("TPOT: {:.1}ms", tpot);
|
||||
eprintln!("Throughput: {:.1} tok/s", tok_s);
|
||||
|
||||
// Cleanup
|
||||
broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Shutdown);
|
||||
for (h, _) in worker_handles {
|
||||
h.join().unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
// --- Worker infrastructure ---
|
||||
|
||||
#[derive(Clone)]
|
||||
enum WorkerCmd {
|
||||
Register(usize),
|
||||
Prefill { tokens: Vec<u32>, slot: usize },
|
||||
Decode { tokens: Vec<u32>, positions: Vec<usize>, slots: Vec<usize> },
|
||||
Shutdown,
|
||||
}
|
||||
|
||||
fn worker_loop(
|
||||
rank: usize,
|
||||
world: usize,
|
||||
uid: UniqueId,
|
||||
model_dir: PathBuf,
|
||||
config: ModelConfig,
|
||||
max_seq_len: usize,
|
||||
rx: std::sync::mpsc::Receiver<WorkerCmd>,
|
||||
ack_tx: std::sync::mpsc::Sender<()>,
|
||||
) {
|
||||
xserv_cuda::device::set_device(rank as u32).unwrap();
|
||||
let tp = Arc::new(TpContext::init(rank, world, uid, rank as u32));
|
||||
eprintln!("[rank {rank}] Loading weights...");
|
||||
let weights = loader::load_model_dir(&model_dir, Device::Cpu);
|
||||
let model = GptOss::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 + 64;
|
||||
let mut cache = PagedKVCache::new_tp(
|
||||
&config, local_kv, total_blocks, 0, 4, max_blocks_per_seq, DType::BF16, rank as u32,
|
||||
);
|
||||
eprintln!("[rank {rank}] Ready.");
|
||||
ack_tx.send(()).unwrap();
|
||||
|
||||
while let Ok(cmd) = rx.recv() {
|
||||
match cmd {
|
||||
WorkerCmd::Register(slot) => {
|
||||
let _ = cache.register_sequence(slot);
|
||||
}
|
||||
WorkerCmd::Prefill { tokens, slot } => {
|
||||
let _ = model.forward_prefill_paged(&tokens, slot, &mut cache);
|
||||
}
|
||||
WorkerCmd::Decode { tokens, positions, slots } => {
|
||||
let _ = model.forward_decode_paged(&tokens, &positions, &slots, &mut cache);
|
||||
}
|
||||
WorkerCmd::Shutdown => break,
|
||||
}
|
||||
ack_tx.send(()).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
fn broadcast_cmd(
|
||||
txs: &[std::sync::mpsc::Sender<WorkerCmd>],
|
||||
_handles: &[(std::thread::JoinHandle<()>, std::sync::mpsc::Receiver<()>)],
|
||||
cmd: WorkerCmd,
|
||||
) {
|
||||
for tx in txs {
|
||||
tx.send(cmd.clone()).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
fn wait_workers(handles: &[(std::thread::JoinHandle<()>, std::sync::mpsc::Receiver<()>)]) {
|
||||
for (_, rx) in handles {
|
||||
rx.recv().unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_greedy_last(logits: &xserv_tensor::Tensor) -> u32 {
|
||||
use half::bf16;
|
||||
assert_eq!(logits.ndim(), 2);
|
||||
let logits_cpu = logits.to_device(Device::Cpu);
|
||||
let vocab_size = logits.shape()[1];
|
||||
let seq_len = logits.shape()[0];
|
||||
let data = logits_cpu.as_slice::<bf16>();
|
||||
let last = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
|
||||
|
||||
|
||||
last.iter().enumerate()
|
||||
.max_by(|a, b| {
|
||||
let af = a.1.to_f32();
|
||||
let bf = b.1.to_f32();
|
||||
af.partial_cmp(&bf).unwrap_or(std::cmp::Ordering::Equal)
|
||||
})
|
||||
.map(|(i, _)| i as u32).unwrap()
|
||||
}
|
||||
|
||||
fn get_arg<T: std::str::FromStr>(args: &[String], flag: &str) -> Option<T> {
|
||||
args.iter()
|
||||
.position(|a| a == flag)
|
||||
.and_then(|i| args.get(i + 1))
|
||||
.and_then(|s| s.parse().ok())
|
||||
}
|
||||
@@ -1,6 +1,6 @@
|
||||
use std::io::{self, Write};
|
||||
use std::path::PathBuf;
|
||||
use xserv_model::{loader, KVCache, ModelConfig};
|
||||
use xserv_model::{loader, KVCache, ModelConfig, PagedKVCache, BLOCK_SIZE};
|
||||
use xserv_tensor::{DType, Device};
|
||||
use xserv_tokenizer::Tokenizer;
|
||||
|
||||
@@ -36,14 +36,18 @@ fn main() {
|
||||
eprintln!("Loaded {} tensors", weights.len());
|
||||
|
||||
let is_qwen3 = model_type.contains("qwen");
|
||||
let dtype = if is_qwen3 { DType::BF16 } else { DType::F32 };
|
||||
let is_gpt_oss = model_type.contains("gpt_oss");
|
||||
let dtype = if is_qwen3 || is_gpt_oss { DType::BF16 } else { DType::F32 };
|
||||
|
||||
// Build model
|
||||
enum Model {
|
||||
GPT2(xserv_model::GPT2),
|
||||
Qwen3(xserv_model::Qwen3),
|
||||
GptOss(xserv_model::GptOss),
|
||||
}
|
||||
let model = if is_qwen3 {
|
||||
let model = if is_gpt_oss {
|
||||
Model::GptOss(xserv_model::GptOss::from_weights(config.clone(), weights))
|
||||
} else if is_qwen3 {
|
||||
Model::Qwen3(xserv_model::Qwen3::from_weights(config.clone(), weights))
|
||||
} else {
|
||||
Model::GPT2(xserv_model::GPT2::from_weights(config.clone(), weights))
|
||||
@@ -62,40 +66,92 @@ fn main() {
|
||||
if input == "quit" || input == "exit" { break; }
|
||||
|
||||
let token_ids = tokenizer.encode(input);
|
||||
let kv_heads = if is_qwen3 { config.num_kv_heads() } else { config.num_heads() };
|
||||
let mut cache = KVCache::new(
|
||||
config.num_layers(), kv_heads, config.head_dim(), dtype, Device::Cuda(0),
|
||||
);
|
||||
|
||||
// Prefill + decode
|
||||
let logits = match &model {
|
||||
Model::GPT2(m) => m.forward_with_cache(&token_ids, &mut cache),
|
||||
Model::Qwen3(m) => m.forward_with_cache(&token_ids, &mut cache),
|
||||
};
|
||||
let mut next = match &model {
|
||||
Model::GPT2(_) => xserv_model::gpt2::sample_greedy(&logits),
|
||||
Model::Qwen3(_) => xserv_model::qwen3::sample_greedy(&logits),
|
||||
};
|
||||
if is_gpt_oss {
|
||||
// GptOss uses paged KV cache
|
||||
let max_seq = 2048;
|
||||
let max_blocks_per_seq = (max_seq + BLOCK_SIZE - 1) / BLOCK_SIZE;
|
||||
let total_blocks = max_blocks_per_seq + 64;
|
||||
let mut paged_cache = PagedKVCache::new(
|
||||
&config, total_blocks, 0, 4, max_blocks_per_seq, DType::BF16, 0,
|
||||
);
|
||||
let slot = 0;
|
||||
paged_cache.register_sequence(slot).expect("register slot");
|
||||
|
||||
print!("{input}");
|
||||
io::stdout().flush().unwrap();
|
||||
let model = match &model { Model::GptOss(m) => m, _ => unreachable!() };
|
||||
let logits = model.forward_prefill_paged(&token_ids, slot, &mut paged_cache);
|
||||
let mut next = sample_greedy_last(&logits);
|
||||
|
||||
for _ in 0..max_tokens {
|
||||
let text = tokenizer.decode(&[next]);
|
||||
print!("{text}");
|
||||
print!("{input}");
|
||||
io::stdout().flush().unwrap();
|
||||
|
||||
if tokenizer.eos_token_id() == Some(next) { break; }
|
||||
for _ in 0..max_tokens {
|
||||
let text = tokenizer.decode(&[next]);
|
||||
print!("{text}");
|
||||
io::stdout().flush().unwrap();
|
||||
|
||||
if tokenizer.eos_token_id() == Some(next) { break; }
|
||||
|
||||
let pos = paged_cache.seq_len(slot);
|
||||
let logits = model.forward_decode_paged(
|
||||
&[next], &[pos], &[slot], &mut paged_cache,
|
||||
);
|
||||
next = sample_greedy_last(&logits);
|
||||
}
|
||||
println!();
|
||||
paged_cache.free_sequence(slot);
|
||||
} else {
|
||||
let kv_heads = if is_qwen3 { config.num_kv_heads() } else { config.num_heads() };
|
||||
let mut cache = KVCache::new(
|
||||
config.num_layers(), kv_heads, config.head_dim(), dtype, Device::Cuda(0),
|
||||
);
|
||||
|
||||
let logits = match &model {
|
||||
Model::GPT2(m) => m.forward_with_cache(&[next], &mut cache),
|
||||
Model::Qwen3(m) => m.forward_with_cache(&[next], &mut cache),
|
||||
Model::GPT2(m) => m.forward_with_cache(&token_ids, &mut cache),
|
||||
Model::Qwen3(m) => m.forward_with_cache(&token_ids, &mut cache),
|
||||
Model::GptOss(_) => unreachable!(),
|
||||
};
|
||||
next = match &model {
|
||||
let mut next = match &model {
|
||||
Model::GPT2(_) => xserv_model::gpt2::sample_greedy(&logits),
|
||||
Model::Qwen3(_) => xserv_model::qwen3::sample_greedy(&logits),
|
||||
Model::GptOss(_) => unreachable!(),
|
||||
};
|
||||
|
||||
print!("{input}");
|
||||
io::stdout().flush().unwrap();
|
||||
|
||||
for _ in 0..max_tokens {
|
||||
let text = tokenizer.decode(&[next]);
|
||||
print!("{text}");
|
||||
io::stdout().flush().unwrap();
|
||||
|
||||
if tokenizer.eos_token_id() == Some(next) { break; }
|
||||
|
||||
let logits = match &model {
|
||||
Model::GPT2(m) => m.forward_with_cache(&[next], &mut cache),
|
||||
Model::Qwen3(m) => m.forward_with_cache(&[next], &mut cache),
|
||||
Model::GptOss(_) => unreachable!(),
|
||||
};
|
||||
next = match &model {
|
||||
Model::GPT2(_) => xserv_model::gpt2::sample_greedy(&logits),
|
||||
Model::Qwen3(_) => xserv_model::qwen3::sample_greedy(&logits),
|
||||
Model::GptOss(_) => unreachable!(),
|
||||
};
|
||||
}
|
||||
println!();
|
||||
}
|
||||
println!();
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_greedy_last(logits: &xserv_tensor::Tensor) -> u32 {
|
||||
use half::bf16;
|
||||
assert_eq!(logits.ndim(), 2);
|
||||
let logits_cpu = logits.to_device(Device::Cpu);
|
||||
let vocab_size = logits.shape()[1];
|
||||
let seq_len = logits.shape()[0];
|
||||
let data = logits_cpu.as_slice::<bf16>();
|
||||
let last = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
|
||||
last.iter().enumerate()
|
||||
.max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
|
||||
.map(|(i, _)| i as u32).unwrap()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user