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> {

51
crates/xserv-model/src/bin/bench-gpt2.rs
View File

@@ -1,7 +1,7 @@
use std::path::PathBuf;
use std::time::Instant;
use xserv_model::gpt2::{sample_greedy, KVCache};
use xserv_model::{loader, GPT2, ModelConfig};
use xserv_model::gpt2::{KVCache, sample_greedy};
use xserv_model::{GPT2, ModelConfig, loader};
use xserv_tensor::Device;
use xserv_tokenizer::Tokenizer;
@@ -104,9 +104,15 @@ fn main() {
let tbt_us = if !token_times_us.is_empty() {
token_times_us.iter().sum::<u128>() / token_times_us.len() as u128
} else { 0 };
} else {
0
};
let total_gen_us: u128 = ttft_us + token_times_us.iter().sum::<u128>();
let tpot_us = if num_generated > 0 { total_gen_us / num_generated as u128 } else { 0 };
let tpot_us = if num_generated > 0 {
total_gen_us / num_generated as u128
} else {
0
};
let gen_text_escaped = generated_text
.replace('\\', "\\\\")
@@ -124,11 +130,16 @@ fn main() {
print!("\"ttft_us\": {ttft_us}, ");
print!("\"tbt_us\": {tbt_us}, ");
print!("\"tpot_us\": {tpot_us}}}");
if i < prompts.len() - 1 { println!(","); } else { println!(); }
if i < prompts.len() - 1 {
println!(",");
} else {
println!();
}
eprintln!(
"[{}/{}] input={input_len}tok gen={num_generated}tok ttft={:.1}ms tbt={:.1}ms | {}",
i + 1, prompts.len(),
i + 1,
prompts.len(),
ttft_us as f64 / 1000.0,
tbt_us as f64 / 1000.0,
&generated_text.replace('\n', " ")[..generated_text.len().min(60)]
@@ -138,12 +149,18 @@ fn main() {
}
fn generate_with_cache(
model: &GPT2, config: &ModelConfig, tokenizer: &Tokenizer,
input_ids: &[u32], gen_tokens: usize,
model: &GPT2,
config: &ModelConfig,
tokenizer: &Tokenizer,
input_ids: &[u32],
gen_tokens: usize,
) -> (Vec<u32>, u128, Vec<u128>) {
let mut cache = KVCache::new(
config.num_layers(), config.num_heads(), config.head_dim(),
xserv_tensor::DType::F32, Device::Cuda(0),
config.num_layers(),
config.num_heads(),
config.head_dim(),
xserv_tensor::DType::F32,
Device::Cuda(0),
);
// Prefill
@@ -163,15 +180,19 @@ fn generate_with_cache(
let next = sample_greedy(&logits);
token_times.push(t_start.elapsed().as_micros());
generated.push(next);
if tokenizer.eos_token_id() == Some(next) { break; }
if tokenizer.eos_token_id() == Some(next) {
break;
}
}
(generated, ttft_us, token_times)
}
fn generate_no_cache(
model: &GPT2, tokenizer: &Tokenizer,
input_ids: &[u32], gen_tokens: usize,
model: &GPT2,
tokenizer: &Tokenizer,
input_ids: &[u32],
gen_tokens: usize,
) -> (Vec<u32>, u128, Vec<u128>) {
let mut all_ids = input_ids.to_vec();
@@ -191,7 +212,9 @@ fn generate_no_cache(
token_times.push(t_start.elapsed().as_micros());
all_ids.push(next);
generated.push(next);
if tokenizer.eos_token_id() == Some(next) { break; }
if tokenizer.eos_token_id() == Some(next) {
break;
}
}
(generated, ttft_us, token_times)

52
crates/xserv-model/src/bin/bench-qwen3.rs
View File

@@ -1,7 +1,7 @@
use std::path::PathBuf;
use std::time::Instant;
use xserv_model::qwen3::sample_greedy;
use xserv_model::{loader, DecodeGraphState, GpuKVCache, ModelConfig, Qwen3};
use xserv_model::{DecodeGraphState, GpuKVCache, ModelConfig, Qwen3, loader};
use xserv_tensor::{DType, Device};
use xserv_tokenizer::Tokenizer;
@@ -139,18 +139,35 @@ fn main() {
} else {
// Replay captured graphs
let pos = cache.seq_len() as u32;
graph.execute(last, pos, &mut cache, &layer_ptrs, embed, config.vocab_size as i32, config.hidden() as i32);
graph.execute(
last,
pos,
&mut cache,
&layer_ptrs,
embed,
config.vocab_size as i32,
config.hidden() as i32,
);
cache.advance_seq_len(1);
// Read logits from graph buffer
let vocab_size = config.vocab_size;
let mut logits_bytes = vec![0u8; vocab_size * 2];
graph.logits_buffer().copy_to_host(&mut logits_bytes).unwrap();
graph
.logits_buffer()
.copy_to_host(&mut logits_bytes)
.unwrap();
let logits_data: &[half::bf16] = unsafe {
std::slice::from_raw_parts(logits_bytes.as_ptr() as *const half::bf16, vocab_size)
std::slice::from_raw_parts(
logits_bytes.as_ptr() as *const half::bf16,
vocab_size,
)
};
logits_data.iter().enumerate()
logits_data
.iter()
.enumerate()
.max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
.map(|(idx, _)| idx as u32).unwrap()
.map(|(idx, _)| idx as u32)
.unwrap()
}
} else {
let logits = model.forward_gpu_cache(&[last], &mut cache);
@@ -159,16 +176,24 @@ fn main() {
token_times.push(t_start.elapsed().as_micros());
generated.push(next);
if tokenizer.eos_token_id() == Some(next) { break; }
if tokenizer.eos_token_id() == Some(next) {
break;
}
}
let num_generated = generated.len();
let generated_text = tokenizer.decode(&generated);
let tbt_us = if !token_times.is_empty() {
token_times.iter().sum::<u128>() / token_times.len() as u128
} else { 0 };
} else {
0
};
let total_gen_us: u128 = ttft_us + token_times.iter().sum::<u128>();
let tpot_us = if num_generated > 0 { total_gen_us / num_generated as u128 } else { 0 };
let tpot_us = if num_generated > 0 {
total_gen_us / num_generated as u128
} else {
0
};
let gen_text_escaped = generated_text
.replace('\\', "\\\\")
@@ -186,13 +211,18 @@ fn main() {
print!("\"ttft_us\": {ttft_us}, ");
print!("\"tbt_us\": {tbt_us}, ");
print!("\"tpot_us\": {tpot_us}}}");
if i < prompts.len() - 1 { println!(","); } else { println!(); }
if i < prompts.len() - 1 {
println!(",");
} else {
println!();
}
let display_text = generated_text.replace('\n', " ");
let truncated: String = display_text.chars().take(60).collect();
eprintln!(
"[{}/{}] input={input_len}tok gen={num_generated}tok ttft={:.1}ms tbt={:.1}ms | {}",
i + 1, prompts.len(),
i + 1,
prompts.len(),
ttft_us as f64 / 1000.0,
tbt_us as f64 / 1000.0,
truncated

78
crates/xserv-model/src/bin/bench-tp.rs
View File

@@ -18,7 +18,7 @@ use std::thread;
use std::time::Instant;
use xserv_model::qwen3::sample_greedy;
use xserv_model::{loader, ModelConfig, PagedKVCache, Qwen3, BLOCK_SIZE};
use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, loader};
use xserv_tensor::{DType, Device};
use xserv_tokenizer::Tokenizer;
@@ -35,8 +35,13 @@ fn main() {
std::process::exit(1);
}
let model_dir = PathBuf::from(&args[1]);
let world: usize = arg(&args, "--tp").and_then(|s| s.parse().ok()).unwrap_or(1).max(1);
let gen_tokens: usize = arg(&args, "--gen-tokens").and_then(|s| s.parse().ok()).unwrap_or(64);
let world: usize = arg(&args, "--tp")
.and_then(|s| s.parse().ok())
.unwrap_or(1)
.max(1);
let gen_tokens: usize = arg(&args, "--gen-tokens")
.and_then(|s| s.parse().ok())
.unwrap_or(64);
let devices: Vec<u32> = match arg(&args, "--devices") {
Some(s) => s.split(',').filter_map(|d| d.trim().parse().ok()).collect(),
None => (0..world as u32).collect(),
@@ -67,7 +72,11 @@ fn main() {
// Tensors are not Send (their Storage holds a raw GPU pointer), so each rank
// thread loads its own CPU copy of the weights and shards in-thread. Loading
// is not part of the timed region.
let id = if world > 1 { Some(xserv_distributed::get_unique_id()) } else { None };
let id = if world > 1 {
Some(xserv_distributed::get_unique_id())
} else {
None
};
let handles: Vec<_> = (0..world)
.map(|rank| {
@@ -76,7 +85,9 @@ fn main() {
let prompt_ids = prompt_ids.clone();
let device = devices[rank];
thread::spawn(move || {
run_rank(rank, world, device, id, config, model_dir, prompt_ids, gen_tokens, eos)
run_rank(
rank, world, device, id, config, model_dir, prompt_ids, gen_tokens, eos,
)
})
})
.collect();
@@ -91,7 +102,10 @@ fn main() {
let results = rank0.expect("rank 0 produced no results");
println!("\n=== TP={world} (devices {devices:?}) — Qwen3 E2E benchmark ===");
println!("{:<45} {:>10} {:>12} {:>8}", "prompt", "TTFT(ms)", "decode tok/s", "gen");
println!(
"{:<45} {:>10} {:>12} {:>8}",
"prompt", "TTFT(ms)", "decode tok/s", "gen"
);
let mut tps_sum = 0.0;
for (i, r) in results.iter().enumerate() {
let text = tokenizer.decode(&r.gen_ids).replace('\n', " ");
@@ -99,16 +113,29 @@ fn main() {
let p: String = prompts[i].chars().take(43).collect();
println!(
"{:<45} {:>10.1} {:>12.1} {:>8} | {}",
p, r.ttft_ms, r.decode_tok_s, r.gen_ids.len(), short
p,
r.ttft_ms,
r.decode_tok_s,
r.gen_ids.len(),
short
);
tps_sum += r.decode_tok_s;
}
println!("--- mean decode throughput: {:.1} tok/s ---", tps_sum / results.len() as f64);
println!(
"--- mean decode throughput: {:.1} tok/s ---",
tps_sum / results.len() as f64
);
// Machine-readable line for cross-TP correctness diffing (rank 0 token ids).
let all_ids: Vec<String> = results
.iter()
.map(|r| r.gen_ids.iter().map(|x| x.to_string()).collect::<Vec<_>>().join(","))
.map(|r| {
r.gen_ids
.iter()
.map(|x| x.to_string())
.collect::<Vec<_>>()
.join(",")
})
.collect();
println!("CORRECTNESS_IDS tp={world} {}", all_ids.join(" | "));
}
@@ -126,7 +153,12 @@ fn run_rank(
) -> Option<Vec<PromptResult>> {
// Bind this thread to its GPU and set up the TP communicator.
let tp = if world > 1 {
Some(Arc::new(xserv_distributed::TpContext::init(rank, world, id.unwrap(), device)))
Some(Arc::new(xserv_distributed::TpContext::init(
rank,
world,
id.unwrap(),
device,
)))
} else {
xserv_cuda::device::set_device(device).unwrap();
None
@@ -142,7 +174,14 @@ fn run_rank(
let max_blocks_per_seq = max_seq.div_ceil(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, device,
&config,
local_kv,
total_blocks,
0,
1,
max_blocks_per_seq,
DType::BF16,
device,
);
// Warmup (init kernels / allocator / NCCL channels) — not timed.
@@ -177,12 +216,20 @@ fn run_rank(
steps += 1;
}
let decode_s = t1.elapsed().as_secs_f64();
let decode_tok_s = if steps > 0 && decode_s > 0.0 { steps as f64 / decode_s } else { 0.0 };
let decode_tok_s = if steps > 0 && decode_s > 0.0 {
steps as f64 / decode_s
} else {
0.0
};
cache.free_sequence(0);
if rank == 0 {
out.push(PromptResult { gen_ids: generated, ttft_ms, decode_tok_s });
out.push(PromptResult {
gen_ids: generated,
ttft_ms,
decode_tok_s,
});
}
}
@@ -190,5 +237,8 @@ fn run_rank(
}
fn arg<'a>(args: &'a [String], flag: &str) -> Option<&'a str> {
args.iter().position(|a| a == flag).and_then(|i| args.get(i + 1)).map(|s| s.as_str())
args.iter()
.position(|a| a == flag)
.and_then(|i| args.get(i + 1))
.map(|s| s.as_str())
}

15
crates/xserv-model/src/bin/dump-logits.rs
View File

@@ -1,8 +1,8 @@
use half::bf16;
use std::path::PathBuf;
use xserv_model::{loader, KVCache, ModelConfig, Qwen3};
use xserv_model::{KVCache, ModelConfig, Qwen3, loader};
use xserv_tensor::{DType, Device};
use xserv_tokenizer::Tokenizer;
use half::bf16;
fn main() {
let args: Vec<String> = std::env::args().collect();
@@ -20,8 +20,11 @@ fn main() {
eprintln!("Token IDs: {token_ids:?}");
let mut cache = KVCache::new(
config.num_layers(), config.num_kv_heads(), config.head_dim(),
DType::BF16, Device::Cuda(0),
config.num_layers(),
config.num_kv_heads(),
config.head_dim(),
DType::BF16,
Device::Cuda(0),
);
let logits = model.forward_with_cache(&token_ids, &mut cache);
let logits_cpu = logits.to_device(Device::Cpu);
@@ -31,7 +34,9 @@ fn main() {
// Print top-20 logits for the last position
let last_row = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
let mut indexed: Vec<(usize, f32)> = last_row.iter().enumerate()
let mut indexed: Vec<(usize, f32)> = last_row
.iter()
.enumerate()
.map(|(i, v)| (i, v.to_f32()))
.collect();
indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

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}");
}
}

76
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, PagedKVCache, BLOCK_SIZE};
use xserv_model::{BLOCK_SIZE, KVCache, ModelConfig, PagedKVCache, loader};
use xserv_tensor::{DType, Device};
use xserv_tokenizer::Tokenizer;
@@ -21,14 +21,21 @@ fn main() {
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);
eprintln!(
"GPU: {} ({} MB free)",
info.name,
info.free_memory / 1024 / 1024
);
let config = ModelConfig::from_file(&model_dir.join("config.json"));
let model_type = config.model_type.as_deref().unwrap_or("unknown");
eprintln!(
"Model: {model_type}, layers={}, hidden={}, heads={}/{} kv, vocab={}",
config.num_layers(), config.hidden(), config.num_heads(),
config.num_kv_heads(), config.vocab_size
config.num_layers(),
config.hidden(),
config.num_heads(),
config.num_kv_heads(),
config.vocab_size
);
eprintln!("Loading weights...");
@@ -37,7 +44,11 @@ fn main() {
let is_qwen3 = model_type.contains("qwen");
let is_gpt_oss = model_type.contains("gpt_oss");
let dtype = if is_qwen3 || is_gpt_oss { DType::BF16 } else { DType::F32 };
let dtype = if is_qwen3 || is_gpt_oss {
DType::BF16
} else {
DType::F32
};
// Build model
enum Model {
@@ -60,10 +71,16 @@ fn main() {
print!("xserv> ");
io::stdout().flush().unwrap();
let mut input = String::new();
if io::stdin().read_line(&mut input).unwrap() == 0 { break; }
if io::stdin().read_line(&mut input).unwrap() == 0 {
break;
}
let input = input.trim();
if input.is_empty() { continue; }
if input == "quit" || input == "exit" { break; }
if input.is_empty() {
continue;
}
if input == "quit" || input == "exit" {
break;
}
let token_ids = tokenizer.encode(input);
@@ -73,12 +90,21 @@ fn main() {
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,
&config,
total_blocks,
0,
4,
max_blocks_per_seq,
DType::BF16,
0,
);
let slot = 0;
paged_cache.register_sequence(slot).expect("register slot");
let model = match &model { Model::GptOss(m) => m, _ => unreachable!() };
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);
@@ -90,20 +116,28 @@ fn main() {
print!("{text}");
io::stdout().flush().unwrap();
if tokenizer.eos_token_id() == Some(next) { break; }
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,
);
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 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),
config.num_layers(),
kv_heads,
config.head_dim(),
dtype,
Device::Cuda(0),
);
let logits = match &model {
@@ -125,7 +159,9 @@ fn main() {
print!("{text}");
io::stdout().flush().unwrap();
if tokenizer.eos_token_id() == Some(next) { break; }
if tokenizer.eos_token_id() == Some(next) {
break;
}
let logits = match &model {
Model::GPT2(m) => m.forward_with_cache(&[next], &mut cache),
@@ -151,7 +187,9 @@ fn sample_greedy_last(logits: &xserv_tensor::Tensor) -> u32 {
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()
last.iter()
.enumerate()
.max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
.map(|(i, _)| i as u32).unwrap()
.map(|(i, _)| i as u32)
.unwrap()
}