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

192
crates/xserv-model/src/bin/bench-gpt-oss.rs
View File

@@ -3,7 +3,7 @@ use std::sync::Arc;
use std::time::Instant;
use xserv_distributed::{TpContext, UniqueId, get_unique_id};
use xserv_model::{loader, GptOss, GraphedGptOssDecoder, ModelConfig, PagedKVCache, BLOCK_SIZE};
use xserv_model::{BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, PagedKVCache, loader};
use xserv_tensor::{DType, Device};
use xserv_tokenizer::Tokenizer;
@@ -23,8 +23,12 @@ fn main() {
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.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}");
@@ -59,17 +63,29 @@ fn main() {
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());
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,
&config,
local_kv,
total_blocks,
0,
4,
max_blocks_per_seq,
DType::BF16,
0,
);
eprintln!("[rank 0] Ready.");
// Prompt
let prompt_arg = get_arg::<String>(&args, "--prompt");
let prompt = prompt_arg.as_deref().unwrap_or("What is the meaning of life?");
let prompt = prompt_arg
.as_deref()
.unwrap_or("What is the meaning of life?");
let token_ids = tokenizer.encode(prompt);
eprintln!("Prompt ({} tokens): {prompt}", token_ids.len());
@@ -83,11 +99,21 @@ fn main() {
// (oracle) next token. Removes free-running compounding so it isolates
// whether per-position logits agree with the llama.cpp trajectory.
if let Some(forced) = get_arg::<String>(&args, "--forced") {
let forced_ids: Vec<u32> = forced.split(',').filter_map(|s| s.trim().parse().ok()).collect();
let forced_ids: Vec<u32> = forced
.split(',')
.filter_map(|s| s.trim().parse().ok())
.collect();
let mut seq = token_ids.clone();
seq.extend_from_slice(&forced_ids);
// Workers must run the same prefill in lockstep (TP AllReduces match up).
broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Prefill { tokens: seq.clone(), slot });
broadcast_cmd(
&worker_txs,
&worker_handles,
WorkerCmd::Prefill {
tokens: seq.clone(),
slot,
},
);
let logits = model.forward_prefill_paged(&seq, slot, &mut cache);
wait_workers(&worker_handles);
let logits_cpu = logits.to_device(Device::Cpu);
@@ -99,19 +125,31 @@ fn main() {
// position i predicts seq[i+1]; we check the forced region
for i in (plen - 1)..(seq.len() - 1) {
let row = &data[i * vocab..(i + 1) * vocab];
let argmax = row.iter().enumerate()
let argmax = row
.iter()
.enumerate()
.max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
.map(|(j, _)| j as u32).unwrap();
.map(|(j, _)| j as u32)
.unwrap();
let expected = seq[i + 1];
let ok = argmax == expected;
if ok { matches += 1; }
if ok {
matches += 1;
}
total += 1;
eprintln!("pos {i}: xserv_argmax={argmax} oracle={expected} {}", if ok {"OK"} else {"DIFF"});
eprintln!(
"pos {i}: xserv_argmax={argmax} oracle={expected} {}",
if ok { "OK" } else { "DIFF" }
);
}
eprintln!("\nTeacher-forced top-1 agreement: {matches}/{total} = {:.1}%",
100.0 * matches as f64 / total as f64);
eprintln!(
"\nTeacher-forced top-1 agreement: {matches}/{total} = {:.1}%",
100.0 * matches as f64 / total as f64
);
broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Shutdown);
for (h, _) in worker_handles { h.join().unwrap(); }
for (h, _) in worker_handles {
h.join().unwrap();
}
return;
}
@@ -120,8 +158,18 @@ fn main() {
// per-position top-1 agreement bucketed by position. Localizes long-context
// decode degradation (which prefill teacher-forcing cannot see).
if let Some(forced) = get_arg::<String>(&args, "--forced-decode") {
let forced_ids: Vec<u32> = forced.split(',').filter_map(|s| s.trim().parse().ok()).collect();
broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Prefill { tokens: token_ids.clone(), slot });
let forced_ids: Vec<u32> = forced
.split(',')
.filter_map(|s| s.trim().parse().ok())
.collect();
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 mut pred = sample_greedy_last(&logits); // prediction for forced[0]
@@ -133,34 +181,55 @@ fn main() {
matches += ok as usize;
total += 1;
let b = i / bucket;
if buckets.len() <= b { buckets.push((0, 0)); }
if buckets.len() <= b {
buckets.push((0, 0));
}
buckets[b].0 += ok as usize;
buckets[b].1 += 1;
// Teacher-force: feed the oracle token through the decode path.
let pos = cache.seq_len(slot);
broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Decode {
tokens: vec![f], positions: vec![pos], slots: vec![slot],
});
broadcast_cmd(
&worker_txs,
&worker_handles,
WorkerCmd::Decode {
tokens: vec![f],
positions: vec![pos],
slots: vec![slot],
},
);
let logits = model.forward_decode_paged(&[f], &[pos], &[slot], &mut cache);
wait_workers(&worker_handles);
pred = sample_greedy_last(&logits);
}
eprintln!("Teacher-forced DECODE agreement: {matches}/{total} = {:.1}%",
100.0 * matches as f64 / total as f64);
eprintln!(
"Teacher-forced DECODE agreement: {matches}/{total} = {:.1}%",
100.0 * matches as f64 / total as f64
);
for (b, (m, t)) in buckets.iter().enumerate() {
eprintln!(" pos[{:>4}..{:<4}]: {m:>3}/{t:<3} = {:.0}%",
b * bucket, b * bucket + t, 100.0 * (*m as f64) / (*t as f64));
eprintln!(
" pos[{:>4}..{:<4}]: {m:>3}/{t:<3} = {:.0}%",
b * bucket,
b * bucket + t,
100.0 * (*m as f64) / (*t as f64)
);
}
broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Shutdown);
for (h, _) in worker_handles { h.join().unwrap(); }
for (h, _) in worker_handles {
h.join().unwrap();
}
return;
}
// Prefill
let t0 = Instant::now();
broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Prefill {
tokens: token_ids.clone(), slot,
});
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();
@@ -178,12 +247,20 @@ fn main() {
let text = tokenizer.decode(&[next]);
print!("{text}");
if tokenizer.eos_token_id() == Some(next) { break; }
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],
});
broadcast_cmd(
&worker_txs,
&worker_handles,
WorkerCmd::Decode {
tokens: vec![next],
positions: vec![pos],
slots: vec![slot],
},
);
let logits = decoder.decode(&model, &[next], &[pos], &[slot], &mut cache);
wait_workers(&worker_handles);
@@ -196,13 +273,20 @@ fn main() {
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)]);
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 };
} else {
0.0
};
let tok_s = if gen_tokens > 1 {
(gen_tokens - 1) as f64 / decode_elapsed.as_secs_f64()
} else { 0.0 };
} else {
0.0
};
eprintln!("\n--- Performance ---");
eprintln!("Generated: {} tokens", gen_tokens);
@@ -222,8 +306,15 @@ fn main() {
#[derive(Clone)]
enum WorkerCmd {
Register(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,
}
@@ -241,12 +332,20 @@ fn worker_loop(
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 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,
&config,
local_kv,
total_blocks,
0,
4,
max_blocks_per_seq,
DType::BF16,
rank as u32,
);
eprintln!("[rank {rank}] Ready.");
ack_tx.send(()).unwrap();
@@ -260,7 +359,11 @@ fn worker_loop(
WorkerCmd::Prefill { tokens, slot } => {
let _ = model.forward_prefill_paged(&tokens, slot, &mut cache);
}
WorkerCmd::Decode { tokens, positions, slots } => {
WorkerCmd::Decode {
tokens,
positions,
slots,
} => {
let _ = decoder.decode(&model, &tokens, &positions, &slots, &mut cache);
}
WorkerCmd::Shutdown => break,
@@ -299,14 +402,15 @@ fn sample_greedy_last(logits: &xserv_tensor::Tensor) -> u32 {
let data = logits_cpu.as_slice::<bf16>();
let last = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
last.iter().enumerate()
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()
.map(|(i, _)| i as u32)
.unwrap()
}
fn get_arg<T: std::str::FromStr>(args: &[String], flag: &str) -> Option<T> {