gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

phase19: MoE support — gpt-oss-20b end-to-end inference with TP=2

Browse Source 
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
2026-05-30 15:18:01 +08:00
parent 46bfb59f30
commit 9ad91a4a92
12 changed files with 1390 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

108
crates/xserv-model/src/bin/xserv-cli.rs
View File

@@ -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()
}