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

23
crates/xserv-model/src/config.rs
View File

@@ -88,23 +88,33 @@ impl ModelConfig {
}
pub fn hidden(&self) -> usize {
self.hidden_size.or(self.n_embd).expect("hidden_size or n_embd required")
self.hidden_size
.or(self.n_embd)
.expect("hidden_size or n_embd required")
}
pub fn num_heads(&self) -> usize {
self.num_attention_heads.or(self.n_head).expect("num_attention_heads or n_head required")
self.num_attention_heads
.or(self.n_head)
.expect("num_attention_heads or n_head required")
}
pub fn num_layers(&self) -> usize {
self.num_hidden_layers.or(self.n_layer).expect("num_hidden_layers or n_layer required")
self.num_hidden_layers
.or(self.n_layer)
.expect("num_hidden_layers or n_layer required")
}
pub fn max_seq_len(&self) -> usize {
self.max_position_embeddings.or(self.n_positions).unwrap_or(2048)
self.max_position_embeddings
.or(self.n_positions)
.unwrap_or(2048)
}
pub fn ffn_hidden(&self) -> usize {
self.intermediate_size.or(self.n_inner).unwrap_or(self.hidden() * 4)
self.intermediate_size
.or(self.n_inner)
.unwrap_or(self.hidden() * 4)
}
pub fn num_kv_heads(&self) -> usize {
@@ -112,7 +122,8 @@ impl ModelConfig {
}
pub fn head_dim(&self) -> usize {
self.explicit_head_dim.unwrap_or_else(|| self.hidden() / self.num_heads())
self.explicit_head_dim
.unwrap_or_else(|| self.hidden() / self.num_heads())
}
pub fn ln_eps(&self) -> f32 {

338
crates/xserv-model/src/decode_graph.rs
View File

@@ -18,19 +18,19 @@ use crate::kv_cache::GpuKVCache;
/// All buffers have stable GPU addresses for CUDA Graph replay.
struct DecodeBuffers {
// Hidden-size buffers: [1, hidden]
x: GpuBuffer, // running hidden state
normed: GpuBuffer, // rmsnorm output
attn_out: GpuBuffer, // attention output [1, num_heads, 1, head_dim]
attn_merged: GpuBuffer, // merge_heads output [1, hidden]
o_proj: GpuBuffer, // O projection output [1, hidden]
normed2: GpuBuffer, // post-attn norm output [1, hidden]
sum_out: GpuBuffer, // add_rmsnorm sum output [1, hidden]
down: GpuBuffer, // down projection output [1, hidden]
x: GpuBuffer, // running hidden state
normed: GpuBuffer, // rmsnorm output
attn_out: GpuBuffer, // attention output [1, num_heads, 1, head_dim]
attn_merged: GpuBuffer, // merge_heads output [1, hidden]
o_proj: GpuBuffer, // O projection output [1, hidden]
normed2: GpuBuffer, // post-attn norm output [1, hidden]
sum_out: GpuBuffer, // add_rmsnorm sum output [1, hidden]
down: GpuBuffer, // down projection output [1, hidden]
// QKV projection outputs
q_proj: GpuBuffer, // [1, num_heads * head_dim]
k_proj: GpuBuffer, // [1, num_kv_heads * head_dim]
v_proj: GpuBuffer, // [1, num_kv_heads * head_dim]
q_proj: GpuBuffer, // [1, num_heads * head_dim]
k_proj: GpuBuffer, // [1, num_kv_heads * head_dim]
v_proj: GpuBuffer, // [1, num_kv_heads * head_dim]
// Reshaped: [1, H, 1, D]
q_reshaped: GpuBuffer,
@@ -50,23 +50,23 @@ struct DecodeBuffers {
k_final: GpuBuffer,
// FFN intermediates
gate: GpuBuffer, // [1, intermediate]
up: GpuBuffer, // [1, intermediate]
silu_out: GpuBuffer, // [1, intermediate]
gate: GpuBuffer, // [1, intermediate]
up: GpuBuffer, // [1, intermediate]
silu_out: GpuBuffer, // [1, intermediate]
// GEMV fp32 accumulators (separate per output dimension)
fp32_hidden: GpuBuffer, // for hidden-sized GEMV outputs
fp32_q: GpuBuffer, // for Q projection
fp32_kv: GpuBuffer, // for K/V projection
fp32_intermediate: GpuBuffer,// for gate/up projections
fp32_vocab: GpuBuffer, // for lm_head
fp32_hidden: GpuBuffer, // for hidden-sized GEMV outputs
fp32_q: GpuBuffer, // for Q projection
fp32_kv: GpuBuffer, // for K/V projection
fp32_intermediate: GpuBuffer, // for gate/up projections
fp32_vocab: GpuBuffer, // for lm_head
// Token ID and position (GPU-resident, updated before replay)
token_id_gpu: GpuBuffer, // 4 bytes (u32)
position_gpu: GpuBuffer, // 4 bytes (u32)
token_id_gpu: GpuBuffer, // 4 bytes (u32)
position_gpu: GpuBuffer, // 4 bytes (u32)
// Final output
logits: GpuBuffer, // [1, vocab_size]
logits: GpuBuffer, // [1, vocab_size]
}
pub struct DecodeGraphState {
@@ -199,127 +199,296 @@ impl DecodeGraphState {
let cublas = cublas_handle();
// Set cuBLAS to use our stream
unsafe { dispatch::set_cublas_stream(cublas, s); }
unsafe {
dispatch::set_cublas_stream(cublas, s);
}
for (l, lw) in layers.iter().enumerate() {
// === Pre-attention graph ===
self.pre_attn_graphs[l].begin_capture(&self.stream).expect("begin pre-attn capture");
self.pre_attn_graphs[l]
.begin_capture(&self.stream)
.expect("begin pre-attn capture");
unsafe {
// RMSNorm
dispatch::rmsnorm_bf16(
self.buffers.x.as_ptr() as _, lw.input_norm, self.buffers.normed.as_mut_ptr() as _,
1, h, eps, s,
self.buffers.x.as_ptr() as _,
lw.input_norm,
self.buffers.normed.as_mut_ptr() as _,
1,
h,
eps,
s,
);
// Q projection (GEMV)
dispatch::gemv_bf16(
self.buffers.normed.as_ptr() as _, lw.q_proj_wt, self.buffers.q_proj.as_mut_ptr() as _,
self.buffers.normed.as_ptr() as _,
lw.q_proj_wt,
self.buffers.q_proj.as_mut_ptr() as _,
self.buffers.fp32_q.as_mut_ptr() as _,
h, nh * hd, s,
h,
nh * hd,
s,
);
// K projection (GEMV)
dispatch::gemv_bf16(
self.buffers.normed.as_ptr() as _, lw.k_proj_wt, self.buffers.k_proj.as_mut_ptr() as _,
self.buffers.normed.as_ptr() as _,
lw.k_proj_wt,
self.buffers.k_proj.as_mut_ptr() as _,
self.buffers.fp32_kv.as_mut_ptr() as _,
h, nkv * hd, s,
h,
nkv * hd,
s,
);
// V projection (GEMV)
dispatch::gemv_bf16(
self.buffers.normed.as_ptr() as _, lw.v_proj_wt, self.buffers.v_proj.as_mut_ptr() as _,
self.buffers.normed.as_ptr() as _,
lw.v_proj_wt,
self.buffers.v_proj.as_mut_ptr() as _,
self.buffers.fp32_kv.as_mut_ptr() as _,
h, nkv * hd, s,
h,
nkv * hd,
s,
);
// Reshape heads: [1, H*D] -> [1, H, 1, D]
dispatch::reshape_heads_bf16(self.buffers.q_proj.as_ptr() as _, self.buffers.q_reshaped.as_mut_ptr() as _, 1, nh, hd, s);
dispatch::reshape_heads_bf16(self.buffers.k_proj.as_ptr() as _, self.buffers.k_reshaped.as_mut_ptr() as _, 1, nkv, hd, s);
dispatch::reshape_heads_bf16(self.buffers.v_proj.as_ptr() as _, self.buffers.v_reshaped.as_mut_ptr() as _, 1, nkv, hd, s);
dispatch::reshape_heads_bf16(
self.buffers.q_proj.as_ptr() as _,
self.buffers.q_reshaped.as_mut_ptr() as _,
1,
nh,
hd,
s,
);
dispatch::reshape_heads_bf16(
self.buffers.k_proj.as_ptr() as _,
self.buffers.k_reshaped.as_mut_ptr() as _,
1,
nkv,
hd,
s,
);
dispatch::reshape_heads_bf16(
self.buffers.v_proj.as_ptr() as _,
self.buffers.v_reshaped.as_mut_ptr() as _,
1,
nkv,
hd,
s,
);
// QK norm (head-level rmsnorm: treat [1,H,1,D] as [H, D])
dispatch::rmsnorm_bf16(self.buffers.q_reshaped.as_ptr() as _, lw.q_norm, self.buffers.q_normed.as_mut_ptr() as _, nh, hd, eps, s);
dispatch::rmsnorm_bf16(self.buffers.k_reshaped.as_ptr() as _, lw.k_norm, self.buffers.k_normed.as_mut_ptr() as _, nkv, hd, eps, s);
dispatch::rmsnorm_bf16(
self.buffers.q_reshaped.as_ptr() as _,
lw.q_norm,
self.buffers.q_normed.as_mut_ptr() as _,
nh,
hd,
eps,
s,
);
dispatch::rmsnorm_bf16(
self.buffers.k_reshaped.as_ptr() as _,
lw.k_norm,
self.buffers.k_normed.as_mut_ptr() as _,
nkv,
hd,
eps,
s,
);
// Transpose for RoPE: [1,H,1,D] -> [1,H,D]
dispatch::transpose_hsd_to_shd_bf16(self.buffers.q_normed.as_ptr() as _, self.buffers.q_rope.as_mut_ptr() as _, 1, nh, hd, s);
dispatch::transpose_hsd_to_shd_bf16(self.buffers.k_normed.as_ptr() as _, self.buffers.k_rope.as_mut_ptr() as _, 1, nkv, hd, s);
dispatch::transpose_hsd_to_shd_bf16(
self.buffers.q_normed.as_ptr() as _,
self.buffers.q_rope.as_mut_ptr() as _,
1,
nh,
hd,
s,
);
dispatch::transpose_hsd_to_shd_bf16(
self.buffers.k_normed.as_ptr() as _,
self.buffers.k_rope.as_mut_ptr() as _,
1,
nkv,
hd,
s,
);
// RoPE (in-place, reads position_gpu)
dispatch::rope_bf16(self.buffers.q_rope.as_mut_ptr() as _, rope_cos, rope_sin, self.buffers.position_gpu.as_ptr() as _, 1, nh, hd, s);
dispatch::rope_bf16(self.buffers.k_rope.as_mut_ptr() as _, rope_cos, rope_sin, self.buffers.position_gpu.as_ptr() as _, 1, nkv, hd, s);
dispatch::rope_bf16(
self.buffers.q_rope.as_mut_ptr() as _,
rope_cos,
rope_sin,
self.buffers.position_gpu.as_ptr() as _,
1,
nh,
hd,
s,
);
dispatch::rope_bf16(
self.buffers.k_rope.as_mut_ptr() as _,
rope_cos,
rope_sin,
self.buffers.position_gpu.as_ptr() as _,
1,
nkv,
hd,
s,
);
// Transpose back: [1,H,D] -> [1,H,1,D]
dispatch::transpose_shd_to_hsd_bf16(self.buffers.q_rope.as_ptr() as _, self.buffers.q_final.as_mut_ptr() as _, 1, nh, hd, s);
dispatch::transpose_shd_to_hsd_bf16(self.buffers.k_rope.as_ptr() as _, self.buffers.k_final.as_mut_ptr() as _, 1, nkv, hd, s);
dispatch::transpose_shd_to_hsd_bf16(
self.buffers.q_rope.as_ptr() as _,
self.buffers.q_final.as_mut_ptr() as _,
1,
nh,
hd,
s,
);
dispatch::transpose_shd_to_hsd_bf16(
self.buffers.k_rope.as_ptr() as _,
self.buffers.k_final.as_mut_ptr() as _,
1,
nkv,
hd,
s,
);
}
self.pre_attn_graphs[l].end_capture(&self.stream).expect("end pre-attn capture");
self.pre_attn_graphs[l]
.end_capture(&self.stream)
.expect("end pre-attn capture");
// === Post-attention graph ===
self.post_attn_graphs[l].begin_capture(&self.stream).expect("begin post-attn capture");
self.post_attn_graphs[l]
.begin_capture(&self.stream)
.expect("begin post-attn capture");
unsafe {
// Merge heads: [1,H,1,D] -> [1, hidden]
// attn_out is written by ungraphed attention
dispatch::merge_heads_bf16(self.buffers.attn_out.as_ptr() as _, self.buffers.attn_merged.as_mut_ptr() as _, 1, nh, hd, s);
dispatch::merge_heads_bf16(
self.buffers.attn_out.as_ptr() as _,
self.buffers.attn_merged.as_mut_ptr() as _,
1,
nh,
hd,
s,
);
// O projection
dispatch::gemv_bf16(
self.buffers.attn_merged.as_ptr() as _, lw.o_proj_wt, self.buffers.o_proj.as_mut_ptr() as _,
self.buffers.attn_merged.as_ptr() as _,
lw.o_proj_wt,
self.buffers.o_proj.as_mut_ptr() as _,
self.buffers.fp32_hidden.as_mut_ptr() as _,
nh * hd, h, s,
nh * hd,
h,
s,
);
// Fused Add+RMSNorm: normed2 = rmsnorm(o_proj + x), sum_out = o_proj + x
dispatch::add_rmsnorm_bf16(
self.buffers.o_proj.as_ptr() as _, self.buffers.x.as_ptr() as _, lw.post_norm,
self.buffers.normed2.as_mut_ptr() as _, self.buffers.sum_out.as_mut_ptr() as _,
1, h, eps, s,
self.buffers.o_proj.as_ptr() as _,
self.buffers.x.as_ptr() as _,
lw.post_norm,
self.buffers.normed2.as_mut_ptr() as _,
self.buffers.sum_out.as_mut_ptr() as _,
1,
h,
eps,
s,
);
// Gate projection
dispatch::gemv_bf16(
self.buffers.normed2.as_ptr() as _, lw.gate_proj_wt, self.buffers.gate.as_mut_ptr() as _,
self.buffers.normed2.as_ptr() as _,
lw.gate_proj_wt,
self.buffers.gate.as_mut_ptr() as _,
self.buffers.fp32_intermediate.as_mut_ptr() as _,
h, inter, s,
h,
inter,
s,
);
// Up projection
dispatch::gemv_bf16(
self.buffers.normed2.as_ptr() as _, lw.up_proj_wt, self.buffers.up.as_mut_ptr() as _,
self.buffers.normed2.as_ptr() as _,
lw.up_proj_wt,
self.buffers.up.as_mut_ptr() as _,
self.buffers.fp32_intermediate.as_mut_ptr() as _,
h, inter, s,
h,
inter,
s,
);
// Fused SiLU x Mul
dispatch::silu_mul_bf16(self.buffers.gate.as_ptr() as _, self.buffers.up.as_ptr() as _, self.buffers.silu_out.as_mut_ptr() as _, inter, s);
dispatch::silu_mul_bf16(
self.buffers.gate.as_ptr() as _,
self.buffers.up.as_ptr() as _,
self.buffers.silu_out.as_mut_ptr() as _,
inter,
s,
);
// Down projection
dispatch::gemv_bf16(
self.buffers.silu_out.as_ptr() as _, lw.down_proj_wt, self.buffers.down.as_mut_ptr() as _,
self.buffers.silu_out.as_ptr() as _,
lw.down_proj_wt,
self.buffers.down.as_mut_ptr() as _,
self.buffers.fp32_hidden.as_mut_ptr() as _,
inter, h, s,
inter,
h,
s,
);
// x = sum_out + down (residual connection for next layer)
dispatch::add_bf16(self.buffers.sum_out.as_ptr() as _, self.buffers.down.as_ptr() as _, self.buffers.x.as_mut_ptr() as _, h, s);
dispatch::add_bf16(
self.buffers.sum_out.as_ptr() as _,
self.buffers.down.as_ptr() as _,
self.buffers.x.as_mut_ptr() as _,
h,
s,
);
}
self.post_attn_graphs[l].end_capture(&self.stream).expect("end post-attn capture");
self.post_attn_graphs[l]
.end_capture(&self.stream)
.expect("end post-attn capture");
}
// === Final graph: norm + lm_head ===
self.final_graph.begin_capture(&self.stream).expect("begin final capture");
self.final_graph
.begin_capture(&self.stream)
.expect("begin final capture");
unsafe {
dispatch::rmsnorm_bf16(self.buffers.x.as_ptr() as _, norm_weight, self.buffers.normed.as_mut_ptr() as _, 1, h, eps, s);
dispatch::rmsnorm_bf16(
self.buffers.x.as_ptr() as _,
norm_weight,
self.buffers.normed.as_mut_ptr() as _,
1,
h,
eps,
s,
);
dispatch::gemv_bf16(
self.buffers.normed.as_ptr() as _, lm_head_wt, self.buffers.logits.as_mut_ptr() as _,
self.buffers.normed.as_ptr() as _,
lm_head_wt,
self.buffers.logits.as_mut_ptr() as _,
self.buffers.fp32_vocab.as_mut_ptr() as _,
h, vocab, s,
h,
vocab,
s,
);
}
self.final_graph.end_capture(&self.stream).expect("end final capture");
self.final_graph
.end_capture(&self.stream)
.expect("end final capture");
// Reset cuBLAS back to null stream
unsafe { dispatch::set_cublas_stream(cublas, std::ptr::null_mut()); }
unsafe {
dispatch::set_cublas_stream(cublas, std::ptr::null_mut());
}
self.captured = true;
}
@@ -343,8 +512,14 @@ impl DecodeGraphState {
let es = 2usize; // BF16
// Upload token ID and position to fixed GPU buffers
self.buffers.token_id_gpu.copy_from_host(&token_id.to_le_bytes()).unwrap();
self.buffers.position_gpu.copy_from_host(&position.to_le_bytes()).unwrap();
self.buffers
.token_id_gpu
.copy_from_host(&token_id.to_le_bytes())
.unwrap();
self.buffers
.position_gpu
.copy_from_host(&position.to_le_bytes())
.unwrap();
// Embedding (outside graph since token_id changes each step)
unsafe {
@@ -352,13 +527,18 @@ impl DecodeGraphState {
embed_table,
self.buffers.token_id_gpu.as_ptr() as _,
self.buffers.x.as_mut_ptr() as _,
1, hidden_size, vocab_size, s,
1,
hidden_size,
vocab_size,
s,
);
}
for l in 0..self.num_layers {
// Pre-attention graph (norm + QKV + reshape + QK-norm + RoPE)
self.pre_attn_graphs[l].launch(&self.stream).expect("launch pre-attn graph");
self.pre_attn_graphs[l]
.launch(&self.stream)
.expect("launch pre-attn graph");
// Ungraphed: KV cache append
// k_final shape: [1, num_kv_heads, 1, head_dim] (after RoPE pipeline)
@@ -402,9 +582,13 @@ impl DecodeGraphState {
k_full.data_ptr() as _,
v_full.data_ptr() as _,
self.buffers.attn_out.as_mut_ptr() as _,
1, nh as i32, nkv as i32,
kv_len, hd as i32,
scale, s,
1,
nh as i32,
nkv as i32,
kv_len,
hd as i32,
scale,
s,
);
}
@@ -412,11 +596,15 @@ impl DecodeGraphState {
self.stream.synchronize().expect("sync before post-attn");
// Post-attention graph (merge + O-proj + add_rmsnorm + FFN + residual)
self.post_attn_graphs[l].launch(&self.stream).expect("launch post-attn graph");
self.post_attn_graphs[l]
.launch(&self.stream)
.expect("launch post-attn graph");
}
// Final graph (norm + lm_head)
self.final_graph.launch(&self.stream).expect("launch final graph");
self.final_graph
.launch(&self.stream)
.expect("launch final graph");
// Sync to ensure logits are ready
self.stream.synchronize().expect("sync after decode");

109
crates/xserv-model/src/gpt2.rs
View File

@@ -31,7 +31,7 @@ struct GPT2Block {
pub struct KVCache {
// Per layer, per head: raw bytes (works for both f32 and bf16)
k: Vec<Vec<Vec<u8>>>, // [num_layers][num_heads][seq_len * head_dim * elem_size]
k: Vec<Vec<Vec<u8>>>, // [num_layers][num_heads][seq_len * head_dim * elem_size]
v: Vec<Vec<Vec<u8>>>,
len: usize,
num_heads: usize,
@@ -42,7 +42,13 @@ pub struct KVCache {
}
impl KVCache {
pub fn new(num_layers: usize, num_heads: usize, head_dim: usize, dtype: DType, device: Device) -> Self {
pub fn new(
num_layers: usize,
num_heads: usize,
head_dim: usize,
dtype: DType,
device: Device,
) -> Self {
Self {
k: (0..num_layers).map(|_| vec![vec![]; num_heads]).collect(),
v: (0..num_layers).map(|_| vec![vec![]; num_heads]).collect(),
@@ -55,10 +61,18 @@ impl KVCache {
}
}
pub fn seq_len(&self) -> usize { self.len }
pub fn seq_len(&self) -> usize {
self.len
}
/// Append from a CPU tensor with shape [1, H, new_tokens, D].
pub fn append_kv_tensor(&mut self, layer: usize, k_cpu: &Tensor, v_cpu: &Tensor, new_tokens: usize) {
pub fn append_kv_tensor(
&mut self,
layer: usize,
k_cpu: &Tensor,
v_cpu: &Tensor,
new_tokens: usize,
) {
let hd = self.head_dim;
let es = self.elem_size;
let k_bytes = k_cpu.storage().as_cpu_bytes();
@@ -118,7 +132,8 @@ impl GPT2 {
pub fn from_weights(config: ModelConfig, mut w: HashMap<String, Tensor>) -> Self {
crate::init_kernels();
let take = |w: &mut HashMap<String, Tensor>, name: &str| -> Tensor {
w.remove(name).unwrap_or_else(|| panic!("missing weight: {name}"))
w.remove(name)
.unwrap_or_else(|| panic!("missing weight: {name}"))
};
let wte = take(&mut w, "wte.weight");
@@ -147,7 +162,15 @@ impl GPT2 {
});
}
Self { config, wte, wpe, layers, ln_f_g, ln_f_b, lm_head }
Self {
config,
wte,
wpe,
layers,
ln_f_g,
ln_f_b,
lm_head,
}
}
/// Full forward pass without KV cache (for testing / correctness comparison).
@@ -179,14 +202,22 @@ impl GPT2 {
let head_dim = self.config.head_dim();
let tok_emb = embedding(&self.wte, token_ids);
let pos_ids: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
let pos_ids: Vec<u32> = (pos_offset..pos_offset + new_tokens)
.map(|p| p as u32)
.collect();
let pos_emb = embedding(&self.wpe, &pos_ids);
let mut x = add_tensors(&tok_emb, &pos_emb);
for (layer_idx, layer) in self.layers.iter().enumerate() {
x = self.transformer_block(
layer, &x, Some((cache, layer_idx)),
pos_offset, new_tokens, num_heads, head_dim, hidden,
layer,
&x,
Some((cache, layer_idx)),
pos_offset,
new_tokens,
num_heads,
head_dim,
hidden,
);
}
@@ -238,7 +269,11 @@ impl GPT2 {
fn linear(x: &Tensor, weight: &Tensor, bias: Option<&Tensor>) -> Tensor {
let out = matmul_2d(x, weight);
if let Some(b) = bias { add_bias(&out, b) } else { out }
if let Some(b) = bias {
add_bias(&out, b)
} else {
out
}
}
fn matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
@@ -277,7 +312,12 @@ fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
}
}
fn split_qkv(qkv: &Tensor, num_heads: usize, head_dim: usize, seq_len: usize) -> (Tensor, Tensor, Tensor) {
fn split_qkv(
qkv: &Tensor,
num_heads: usize,
head_dim: usize,
seq_len: usize,
) -> (Tensor, Tensor, Tensor) {
let hidden = num_heads * head_dim;
let qkv_cpu = qkv.to_device(Device::Cpu);
let device = qkv.device();
@@ -294,14 +334,21 @@ fn split_qkv(qkv: &Tensor, num_heads: usize, head_dim: usize, seq_len: usize) ->
for h in 0..num_heads {
let src_off = h * head_dim;
let dst_off = (h * seq_len + s) * head_dim;
q_data[dst_off..dst_off + head_dim].copy_from_slice(&row[src_off..src_off + head_dim]);
k_data[dst_off..dst_off + head_dim].copy_from_slice(&row[hidden + src_off..hidden + src_off + head_dim]);
v_data[dst_off..dst_off + head_dim].copy_from_slice(&row[2 * hidden + src_off..2 * hidden + src_off + head_dim]);
q_data[dst_off..dst_off + head_dim]
.copy_from_slice(&row[src_off..src_off + head_dim]);
k_data[dst_off..dst_off + head_dim]
.copy_from_slice(&row[hidden + src_off..hidden + src_off + head_dim]);
v_data[dst_off..dst_off + head_dim].copy_from_slice(
&row[2 * hidden + src_off..2 * hidden + src_off + head_dim],
);
}
}
let q = Tensor::from_slice(&q_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
let k = Tensor::from_slice(&k_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
let v = Tensor::from_slice(&v_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
let q =
Tensor::from_slice(&q_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
let k =
Tensor::from_slice(&k_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
let v =
Tensor::from_slice(&v_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
(q, k, v)
}
DType::BF16 => {
@@ -314,14 +361,21 @@ fn split_qkv(qkv: &Tensor, num_heads: usize, head_dim: usize, seq_len: usize) ->
for h in 0..num_heads {
let src_off = h * head_dim;
let dst_off = (h * seq_len + s) * head_dim;
q_data[dst_off..dst_off + head_dim].copy_from_slice(&row[src_off..src_off + head_dim]);
k_data[dst_off..dst_off + head_dim].copy_from_slice(&row[hidden + src_off..hidden + src_off + head_dim]);
v_data[dst_off..dst_off + head_dim].copy_from_slice(&row[2 * hidden + src_off..2 * hidden + src_off + head_dim]);
q_data[dst_off..dst_off + head_dim]
.copy_from_slice(&row[src_off..src_off + head_dim]);
k_data[dst_off..dst_off + head_dim]
.copy_from_slice(&row[hidden + src_off..hidden + src_off + head_dim]);
v_data[dst_off..dst_off + head_dim].copy_from_slice(
&row[2 * hidden + src_off..2 * hidden + src_off + head_dim],
);
}
}
let q = Tensor::from_slice(&q_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
let k = Tensor::from_slice(&k_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
let v = Tensor::from_slice(&v_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
let q =
Tensor::from_slice(&q_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
let k =
Tensor::from_slice(&k_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
let v =
Tensor::from_slice(&v_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
(q, k, v)
}
_ => panic!("unsupported dtype {:?} in split_qkv", dtype),
@@ -343,7 +397,8 @@ fn merge_heads(x: &Tensor, seq_len: usize, hidden: usize) -> Tensor {
for h in 0..num_heads {
let src_off = (h * seq_len + s) * head_dim;
let dst_off = s * hidden + h * head_dim;
out[dst_off..dst_off + head_dim].copy_from_slice(&src[src_off..src_off + head_dim]);
out[dst_off..dst_off + head_dim]
.copy_from_slice(&src[src_off..src_off + head_dim]);
}
}
Tensor::from_slice(&out, &[seq_len, hidden]).to_device(device)
@@ -355,7 +410,8 @@ fn merge_heads(x: &Tensor, seq_len: usize, hidden: usize) -> Tensor {
for h in 0..num_heads {
let src_off = (h * seq_len + s) * head_dim;
let dst_off = s * hidden + h * head_dim;
out[dst_off..dst_off + head_dim].copy_from_slice(&src[src_off..src_off + head_dim]);
out[dst_off..dst_off + head_dim]
.copy_from_slice(&src[src_off..src_off + head_dim]);
}
}
Tensor::from_slice(&out, &[seq_len, hidden]).to_device(device)
@@ -372,7 +428,8 @@ pub fn sample_greedy(logits: &Tensor) -> u32 {
let vocab_size = logits.shape()[1];
let seq_len = logits.shape()[0];
let last_row = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
last_row.iter()
last_row
.iter()
.enumerate()
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
.map(|(idx, _)| idx as u32)

359
crates/xserv-model/src/gpt_oss.rs
View File

@@ -1,6 +1,6 @@
use half::bf16;
use std::collections::HashMap;
use std::ffi::c_void;
use half::bf16;
use xserv_kernels::*;
use xserv_tensor::{Device, Tensor};
@@ -49,10 +49,10 @@ struct GptOssBlock {
expert_down_bias: Tensor, // [local_experts, hidden]
// FP8 quantized expert weights (Some when running FP8 W8A8)
// Transposed layout [E, N, K] for cuBLASLt FP8 (Blackwell requires transA=T)
expert_gate_up_fp8: Option<Tensor>, // [local_experts, 2*inter, hidden] FP8E4M3
expert_gate_up_scale: Option<Tensor>,// [local_experts] F32
expert_down_fp8: Option<Tensor>, // [local_experts, hidden, inter] FP8E4M3
expert_down_scale: Option<Tensor>, // [local_experts] F32
expert_gate_up_fp8: Option<Tensor>, // [local_experts, 2*inter, hidden] FP8E4M3
expert_gate_up_scale: Option<Tensor>, // [local_experts] F32
expert_down_fp8: Option<Tensor>, // [local_experts, hidden, inter] FP8E4M3
expert_down_scale: Option<Tensor>, // [local_experts] F32
// MXFP4 W4A16 expert weights (Some when running 4-bit weight-only).
// (packed [E, N, K/2] u8, scales [E, N, K/32] u8) in [E, N, K] layout.
expert_gate_up_mxfp4: Option<(Tensor, Tensor)>,
@@ -79,16 +79,23 @@ impl GptOss {
crate::init_kernels();
let dev = Device::Cuda(device);
let take = |w: &mut HashMap<String, Tensor>, name: &str| -> Tensor {
w.remove(name).unwrap_or_else(|| panic!("missing weight: {name}"))
w.remove(name)
.unwrap_or_else(|| panic!("missing weight: {name}"))
};
let repl = |t: Tensor| -> Tensor { t.to_device(dev) };
// column-parallel: shard rows of [out, in], transpose → [in, out/world]
let col = |t: Tensor| -> Tensor {
shard_rows(&t, rank, world).to_device(dev).transpose(0, 1).contiguous()
shard_rows(&t, rank, world)
.to_device(dev)
.transpose(0, 1)
.contiguous()
};
// row-parallel: shard cols of [out, in], transpose → [in/world, out]
let row = |t: Tensor| -> Tensor {
shard_cols(&t, rank, world).to_device(dev).transpose(0, 1).contiguous()
shard_cols(&t, rank, world)
.to_device(dev)
.transpose(0, 1)
.contiguous()
};
// Bias sharding helpers
let col_bias = |t: Tensor| -> Tensor { shard_1d(&t, rank, world).to_device(dev) };
@@ -97,7 +104,9 @@ impl GptOss {
let embed_tokens = repl(take(&mut w, "model.embed_tokens.weight"));
let norm = repl(take(&mut w, "model.norm.weight"));
let norm_bias = w.remove("model.norm.bias").map(|t| repl(t));
let lm_head_t = repl(take(&mut w, "lm_head.weight")).transpose(0, 1).contiguous();
let lm_head_t = repl(take(&mut w, "lm_head.weight"))
.transpose(0, 1)
.contiguous();
let head_dim = config.head_dim();
let rope_theta = config.rope_theta.unwrap_or(150000.0);
@@ -176,15 +185,30 @@ impl GptOss {
// MXFP4 stores 4-bit weights in an FP8E4M3 byte container (same dtype
// as FP8), so distinguish by the scale rank: FP8 scale is 1-D [E],
// MXFP4 scale is 3-D [E, N, K/32].
let is_mxfp4 = gate_up_scale.as_ref().map(|s| s.ndim() == 3).unwrap_or(false);
let is_mxfp4 = gate_up_scale
.as_ref()
.map(|s| s.ndim() == 3)
.unwrap_or(false);
let is_fp8 = !is_mxfp4 && gate_up_3d.dtype() == xserv_tensor::DType::FP8E4M3;
let mut expert_gate_up_mxfp4: Option<(Tensor, Tensor)> = None;
let mut expert_down_mxfp4: Option<(Tensor, Tensor)> = None;
let inter2 = if is_mxfp4 { gate_up_3d.shape()[1] } else { gate_up_3d.shape()[2] }; // 2*inter (N)
let hidden = if is_mxfp4 { gate_up_3d.shape()[2] * 2 } else { gate_up_3d.shape()[1] };
let inter = if is_mxfp4 { down_3d.shape()[2] * 2 } else { down_3d.shape()[1] };
let inter2 = if is_mxfp4 {
gate_up_3d.shape()[1]
} else {
gate_up_3d.shape()[2]
}; // 2*inter (N)
let hidden = if is_mxfp4 {
gate_up_3d.shape()[2] * 2
} else {
gate_up_3d.shape()[1]
};
let inter = if is_mxfp4 {
down_3d.shape()[2] * 2
} else {
down_3d.shape()[1]
};
// Slice the rank's range of experts as contiguous 3D tensors on GPU
let expert_gate_up_wt;
@@ -199,10 +223,38 @@ impl GptOss {
// + scales [E, N, K/32]. Slice this rank's experts (raw bytes).
let gu_s = gate_up_scale.expect("MXFP4 model missing gate_up_proj_scale");
let d_s = down_scale.expect("MXFP4 model missing down_proj_scale");
let gu_packed = slice_expert_range_3d_raw(&gate_up_3d, expert_start, local_experts, inter2, hidden / 2).to_device(dev);
let gu_scl = slice_expert_range_3d_raw(&gu_s, expert_start, local_experts, inter2, hidden / 32).to_device(dev);
let dn_packed = slice_expert_range_3d_raw(&down_3d, expert_start, local_experts, hidden, inter / 2).to_device(dev);
let dn_scl = slice_expert_range_3d_raw(&d_s, expert_start, local_experts, hidden, inter / 32).to_device(dev);
let gu_packed = slice_expert_range_3d_raw(
&gate_up_3d,
expert_start,
local_experts,
inter2,
hidden / 2,
)
.to_device(dev);
let gu_scl = slice_expert_range_3d_raw(
&gu_s,
expert_start,
local_experts,
inter2,
hidden / 32,
)
.to_device(dev);
let dn_packed = slice_expert_range_3d_raw(
&down_3d,
expert_start,
local_experts,
hidden,
inter / 2,
)
.to_device(dev);
let dn_scl = slice_expert_range_3d_raw(
&d_s,
expert_start,
local_experts,
hidden,
inter / 32,
)
.to_device(dev);
expert_gate_up_mxfp4 = Some((gu_packed, gu_scl));
expert_down_mxfp4 = Some((dn_packed, dn_scl));
expert_gate_up_fp8 = None;
@@ -214,36 +266,65 @@ impl GptOss {
} else if is_fp8 {
// FP8 W8A8 path: load and TRANSPOSE weights for cuBLASLt (requires transA=T on Blackwell).
// Original: [E, K, N] → Transposed: [E, N, K]
let gu_sliced = slice_expert_range_3d_raw(&gate_up_3d, expert_start, local_experts, hidden, inter2);
let dn_sliced = slice_expert_range_3d_raw(&down_3d, expert_start, local_experts, inter, hidden);
expert_gate_up_fp8 = Some(transpose_3d_inner_raw(&gu_sliced, local_experts, hidden, inter2).to_device(dev));
expert_down_fp8 = Some(transpose_3d_inner_raw(&dn_sliced, local_experts, inter, hidden).to_device(dev));
let gu_sliced = slice_expert_range_3d_raw(
&gate_up_3d,
expert_start,
local_experts,
hidden,
inter2,
);
let dn_sliced =
slice_expert_range_3d_raw(&down_3d, expert_start, local_experts, inter, hidden);
expert_gate_up_fp8 = Some(
transpose_3d_inner_raw(&gu_sliced, local_experts, hidden, inter2)
.to_device(dev),
);
expert_down_fp8 = Some(
transpose_3d_inner_raw(&dn_sliced, local_experts, inter, hidden).to_device(dev),
);
// Scales: [num_experts] F32 → slice to [local_experts]
let gu_s = gate_up_scale.expect("FP8 model missing gate_up_proj_scale");
let d_s = down_scale.expect("FP8 model missing down_proj_scale");
expert_gate_up_scale_gpu = Some(slice_scale_range(&gu_s, expert_start, local_experts).to_device(dev));
expert_down_scale_gpu = Some(slice_scale_range(&d_s, expert_start, local_experts).to_device(dev));
expert_gate_up_scale_gpu =
Some(slice_scale_range(&gu_s, expert_start, local_experts).to_device(dev));
expert_down_scale_gpu =
Some(slice_scale_range(&d_s, expert_start, local_experts).to_device(dev));
// Dummy BF16 tensors (never read in FP8 path)
expert_gate_up_wt = Tensor::empty(&[1, 1, 1], xserv_tensor::DType::BF16, dev);
expert_down_wt = Tensor::empty(&[1, 1, 1], xserv_tensor::DType::BF16, dev);
} else {
// BF16 path: existing behavior
expert_gate_up_wt = slice_expert_range_3d(&gate_up_3d, expert_start, local_experts, hidden, inter2).to_device(dev);
expert_down_wt = slice_expert_range_3d(&down_3d, expert_start, local_experts, inter, hidden).to_device(dev);
expert_gate_up_wt =
slice_expert_range_3d(&gate_up_3d, expert_start, local_experts, hidden, inter2)
.to_device(dev);
expert_down_wt =
slice_expert_range_3d(&down_3d, expert_start, local_experts, inter, hidden)
.to_device(dev);
expert_gate_up_fp8 = None;
expert_gate_up_scale_gpu = None;
expert_down_fp8 = None;
expert_down_scale_gpu = None;
}
let expert_gate_up_bias = slice_expert_range_2d(&gate_up_bias_2d, expert_start, local_experts, inter2).to_device(dev);
let expert_down_bias = slice_expert_range_2d(&down_bias_2d, expert_start, local_experts, hidden).to_device(dev);
let expert_gate_up_bias =
slice_expert_range_2d(&gate_up_bias_2d, expert_start, local_experts, inter2)
.to_device(dev);
let expert_down_bias =
slice_expert_range_2d(&down_bias_2d, expert_start, local_experts, hidden)
.to_device(dev);
xserv_cuda::allocator::cached_trim();
let input_norm = repl(take(&mut w, &format!("{p}.input_layernorm.weight")));
let input_norm_bias = w.remove(&format!("{p}.input_layernorm.bias")).map(|t| repl(t));
let post_norm = repl(take(&mut w, &format!("{p}.post_attention_layernorm.weight")));
let post_norm_bias = w.remove(&format!("{p}.post_attention_layernorm.bias")).map(|t| repl(t));
let input_norm_bias = w
.remove(&format!("{p}.input_layernorm.bias"))
.map(|t| repl(t));
let post_norm = repl(take(
&mut w,
&format!("{p}.post_attention_layernorm.weight"),
));
let post_norm_bias = w
.remove(&format!("{p}.post_attention_layernorm.bias"))
.map(|t| repl(t));
layers.push(GptOssBlock {
input_norm,
@@ -283,17 +364,27 @@ impl GptOss {
let local_num_kv_heads = config.num_kv_heads() / world;
let has_norm_bias = norm_bias.is_some();
let is_fp8 = layers.first().map(|l| l.expert_gate_up_fp8.is_some()).unwrap_or(false);
let is_mxfp4 = layers.first().map(|l| l.expert_gate_up_mxfp4.is_some()).unwrap_or(false);
let is_fp8 = layers
.first()
.map(|l| l.expert_gate_up_fp8.is_some())
.unwrap_or(false);
let is_mxfp4 = layers
.first()
.map(|l| l.expert_gate_up_mxfp4.is_some())
.unwrap_or(false);
if rank == 0 {
if has_norm_bias {
eprintln!("gpt-oss: detected LayerNorm bias — using LayerNorm instead of RMSNorm");
}
if is_fp8 {
eprintln!("gpt-oss: FP8 E4M3 quantized expert weights detected (W8A8 cuBLASLt mode)");
eprintln!(
"gpt-oss: FP8 E4M3 quantized expert weights detected (W8A8 cuBLASLt mode)"
);
}
if is_mxfp4 {
eprintln!("gpt-oss: MXFP4 quantized expert weights detected (W4A16 fused-GEMV mode)");
eprintln!(
"gpt-oss: MXFP4 quantized expert weights detected (W4A16 fused-GEMV mode)"
);
}
}
@@ -341,7 +432,13 @@ impl GptOss {
}
#[inline]
fn add_norm(x: &Tensor, residual: &Tensor, weight: &Tensor, bias: &Option<Tensor>, eps: f32) -> (Tensor, Tensor) {
fn add_norm(
x: &Tensor,
residual: &Tensor,
weight: &Tensor,
bias: &Option<Tensor>,
eps: f32,
) -> (Tensor, Tensor) {
match bias {
Some(b) => {
let sum = xserv_kernels::add(x, residual);
@@ -439,7 +536,6 @@ impl GptOss {
let k_all = add_bias(&matmul_2d(&normed, &layer.k_proj_wt), &layer.k_proj_bias);
let v_all = add_bias(&matmul_2d(&normed, &layer.v_proj_wt), &layer.v_proj_bias);
// Reshape for RoPE: [B, H*D] → [B, H, D]
let q_3d = q_all.reshape(&[batch, num_heads, head_dim]);
let k_3d = k_all.reshape(&[batch, num_kv_heads, head_dim]);
@@ -460,9 +556,17 @@ impl GptOss {
let sinks_ptr = layer.sinks.data_ptr() as *const c_void;
let attn_out = paged_decode_attention_sinks(
&q_4d, k_pool_ptr, v_pool_ptr, bt_ptr, cl_ptr,
&q_4d,
k_pool_ptr,
v_pool_ptr,
bt_ptr,
cl_ptr,
sinks_ptr,
batch, num_heads, num_kv_heads, head_dim, max_blocks,
batch,
num_heads,
num_kv_heads,
head_dim,
max_blocks,
layer.window_size,
);
@@ -471,9 +575,14 @@ impl GptOss {
self.all_reduce(&attn_proj);
let attn_proj = add_bias(&attn_proj, &layer.o_proj_bias);
// Residual + post-norm
let (normed, x_new) = Self::add_norm(&attn_proj, &residual, &layer.post_norm, &layer.post_norm_bias, eps);
let (normed, x_new) = Self::add_norm(
&attn_proj,
&residual,
&layer.post_norm,
&layer.post_norm_bias,
eps,
);
let residual = x_new;
let normed = normed.contiguous();
@@ -505,7 +614,9 @@ impl GptOss {
paged_cache.advance_seq_len(slot, new_tokens);
let mut x = embedding(&self.embed_tokens, token_ids);
let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens)
.map(|p| p as u32)
.collect();
for (layer_idx, layer) in self.layers.iter().enumerate() {
let residual = x.clone();
@@ -532,14 +643,21 @@ impl GptOss {
let (k_full, v_full) = paged_cache.gather_kv_contiguous(slot, layer_idx);
// Flash attention with gpt-oss sinks + (per-layer) sliding window.
let attn_out = flash_attention_sinks(&q, &k_full, &v_full, &layer.sinks, layer.window_size);
let attn_out =
flash_attention_sinks(&q, &k_full, &v_full, &layer.sinks, layer.window_size);
let attn_merged = merge_heads_gpu(&attn_out, new_tokens, num_heads, head_dim);
let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
self.all_reduce(&attn_proj);
let attn_proj = add_bias(&attn_proj, &layer.o_proj_bias);
let (normed, x_new) = Self::add_norm(&attn_proj, &residual, &layer.post_norm, &layer.post_norm_bias, eps);
let (normed, x_new) = Self::add_norm(
&attn_proj,
&residual,
&layer.post_norm,
&layer.post_norm_bias,
eps,
);
let residual = x_new;
// MoE MLP
@@ -566,15 +684,11 @@ impl GptOss {
let expert_start = rank * local_experts;
// 1. Router: [tokens, hidden] @ [hidden, num_experts] + bias → [tokens, num_experts]
let router_logits = add_bias(
&matmul_2d(x, &layer.router_wt),
&layer.router_bias,
);
let router_logits = add_bias(&matmul_2d(x, &layer.router_wt), &layer.router_bias);
// 2. GPU top-k + softmax
let (topk_ids, topk_weights) = xserv_kernels::moe::moe_topk_softmax(
&router_logits, num_experts, top_k,
);
let (topk_ids, topk_weights) =
xserv_kernels::moe::moe_topk_softmax(&router_logits, num_experts, top_k);
// Sparse decode path: compute ONLY the routed experts. The dense path
// below reads every local expert's weights per forward; the sparse
@@ -588,15 +702,31 @@ impl GptOss {
let n = packed.shape()[1];
let k = packed.shape()[2] * 2;
xserv_kernels::moe::moe_sparse_gemv_mxfp4(
x, packed, scales, &layer.expert_gate_up_bias, &topk_ids,
num_tokens, top_k, n, k, expert_start, local_experts, false,
x,
packed,
scales,
&layer.expert_gate_up_bias,
&topk_ids,
num_tokens,
top_k,
n,
k,
expert_start,
local_experts,
false,
)
} else {
xserv_kernels::moe::moe_sparse_gemv_fp8(
x, layer.expert_gate_up_fp8.as_ref().unwrap(),
x,
layer.expert_gate_up_fp8.as_ref().unwrap(),
layer.expert_gate_up_scale.as_ref().unwrap(),
&layer.expert_gate_up_bias, &topk_ids,
num_tokens, top_k, expert_start, local_experts, false,
&layer.expert_gate_up_bias,
&topk_ids,
num_tokens,
top_k,
expert_start,
local_experts,
false,
)
};
@@ -611,20 +741,40 @@ impl GptOss {
let n = packed.shape()[1];
let k = packed.shape()[2] * 2;
xserv_kernels::moe::moe_sparse_gemv_mxfp4(
&activated, packed, scales, &layer.expert_down_bias, &topk_ids,
num_tokens, top_k, n, k, expert_start, local_experts, true,
&activated,
packed,
scales,
&layer.expert_down_bias,
&topk_ids,
num_tokens,
top_k,
n,
k,
expert_start,
local_experts,
true,
)
} else {
xserv_kernels::moe::moe_sparse_gemv_fp8(
&activated, layer.expert_down_fp8.as_ref().unwrap(),
&activated,
layer.expert_down_fp8.as_ref().unwrap(),
layer.expert_down_scale.as_ref().unwrap(),
&layer.expert_down_bias, &topk_ids,
num_tokens, top_k, expert_start, local_experts, true,
&layer.expert_down_bias,
&topk_ids,
num_tokens,
top_k,
expert_start,
local_experts,
true,
)
};
let moe_out = xserv_kernels::moe::moe_weighted_sum_sparse(
&down, &topk_ids, &topk_weights, expert_start, local_experts,
&down,
&topk_ids,
&topk_weights,
expert_start,
local_experts,
);
self.all_reduce(&moe_out);
return moe_out;
@@ -644,14 +794,24 @@ impl GptOss {
xserv_kernels::quantization::batched_gemv_mxfp4(&x2, packed, scales, n, k)
.reshape(&[local_experts, 1, n])
} else {
let w_bf16 = xserv_kernels::quantization::dequant_mxfp4_to_bf16_t(packed, scales, local_experts, n, k);
let w_bf16 = xserv_kernels::quantization::dequant_mxfp4_to_bf16_t(
packed,
scales,
local_experts,
n,
k,
);
xserv_kernels::moe::batched_gemm_strided(&x_rep, &w_bf16)
}
} else if let Some(ref wt_fp8_t) = layer.expert_gate_up_fp8 {
// W8A8: quantize activations with per-expert scalar scale, use cuBLASLt FP8 GEMM
let (x_fp8, x_scales) = xserv_kernels::quantization::quantize_bf16_to_fp8_rowwise(&x_rep);
let (x_fp8, x_scales) =
xserv_kernels::quantization::quantize_bf16_to_fp8_rowwise(&x_rep);
xserv_kernels::quantization::batched_gemm_fp8(
&x_fp8, &x_scales, wt_fp8_t, layer.expert_gate_up_scale.as_ref().unwrap(),
&x_fp8,
&x_scales,
wt_fp8_t,
layer.expert_gate_up_scale.as_ref().unwrap(),
)
} else {
xserv_kernels::moe::batched_gemm_strided(&x_rep, &layer.expert_gate_up_wt)
@@ -677,14 +837,24 @@ impl GptOss {
xserv_kernels::quantization::batched_gemv_mxfp4(&a2, packed, scales, n, k)
.reshape(&[local_experts, 1, n])
} else {
let w_bf16 = xserv_kernels::quantization::dequant_mxfp4_to_bf16_t(packed, scales, local_experts, n, k);
let w_bf16 = xserv_kernels::quantization::dequant_mxfp4_to_bf16_t(
packed,
scales,
local_experts,
n,
k,
);
xserv_kernels::moe::batched_gemm_strided(&activated, &w_bf16)
}
} else if let Some(ref wt_fp8) = layer.expert_down_fp8 {
// W8A8: quantize post-GLU activations to FP8, use cuBLASLt FP8 GEMM
let (act_fp8, act_scales) = xserv_kernels::quantization::quantize_bf16_to_fp8_rowwise(&activated);
let (act_fp8, act_scales) =
xserv_kernels::quantization::quantize_bf16_to_fp8_rowwise(&activated);
xserv_kernels::quantization::batched_gemm_fp8(
&act_fp8, &act_scales, wt_fp8, layer.expert_down_scale.as_ref().unwrap(),
&act_fp8,
&act_scales,
wt_fp8,
layer.expert_down_scale.as_ref().unwrap(),
)
} else {
xserv_kernels::moe::batched_gemm_strided(&activated, &layer.expert_down_wt)
@@ -695,8 +865,12 @@ impl GptOss {
// 9. Weighted sum across experts → [tokens, hidden]
let moe_out = xserv_kernels::moe::moe_weighted_sum(
&down, &topk_ids, &topk_weights,
expert_start, local_experts, top_k,
&down,
&topk_ids,
&topk_weights,
expert_start,
local_experts,
top_k,
);
self.all_reduce(&moe_out);
@@ -708,9 +882,7 @@ impl GptOss {
/// Upload a u32 slice to a pooled GPU buffer (synchronous H2D).
fn upload_u32(vals: &[u32]) -> xserv_cuda::GpuBuffer {
let bytes = unsafe {
std::slice::from_raw_parts(vals.as_ptr() as *const u8, vals.len() * 4)
};
let bytes = unsafe { std::slice::from_raw_parts(vals.as_ptr() as *const u8, vals.len() * 4) };
let mut buf = xserv_cuda::allocator::cached_alloc(bytes.len()).expect("alloc u32 upload");
buf.copy_from_host(bytes).unwrap();
buf
@@ -737,11 +909,16 @@ fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
}
fn shard_rows(t: &Tensor, rank: usize, world: usize) -> Tensor {
if world == 1 { return t.clone(); }
if world == 1 {
return t.clone();
}
let shape = t.shape();
assert_eq!(shape.len(), 2);
let (rows, cols) = (shape[0], shape[1]);
assert!(rows % world == 0, "rows {rows} not divisible by world {world}");
assert!(
rows % world == 0,
"rows {rows} not divisible by world {world}"
);
let local = rows / world;
let host = t.to_device(Device::Cpu);
let data = host.as_slice::<bf16>();
@@ -751,11 +928,16 @@ fn shard_rows(t: &Tensor, rank: usize, world: usize) -> Tensor {
}
fn shard_cols(t: &Tensor, rank: usize, world: usize) -> Tensor {
if world == 1 { return t.clone(); }
if world == 1 {
return t.clone();
}
let shape = t.shape();
assert_eq!(shape.len(), 2);
let (rows, cols) = (shape[0], shape[1]);
assert!(cols % world == 0, "cols {cols} not divisible by world {world}");
assert!(
cols % world == 0,
"cols {cols} not divisible by world {world}"
);
let local = cols / world;
let c0 = rank * local;
let host = t.to_device(Device::Cpu);
@@ -769,11 +951,16 @@ fn shard_cols(t: &Tensor, rank: usize, world: usize) -> Tensor {
}
fn shard_1d(t: &Tensor, rank: usize, world: usize) -> Tensor {
if world == 1 { return t.clone(); }
if world == 1 {
return t.clone();
}
let shape = t.shape();
assert_eq!(shape.len(), 1);
let total = shape[0];
assert!(total % world == 0, "dim {total} not divisible by world {world}");
assert!(
total % world == 0,
"dim {total} not divisible by world {world}"
);
let local = total / world;
let host = t.to_device(Device::Cpu);
let data = host.as_slice::<bf16>();
@@ -804,7 +991,13 @@ fn transpose_3d_inner_raw(t: &Tensor, batch: usize, rows: usize, cols: usize) ->
}
/// Extract experts [start..start+count) from a [num_experts, rows, cols] 3D tensor (any dtype, raw bytes).
fn slice_expert_range_3d_raw(t: &Tensor, start: usize, count: usize, rows: usize, cols: usize) -> Tensor {
fn slice_expert_range_3d_raw(
t: &Tensor,
start: usize,
count: usize,
rows: usize,
cols: usize,
) -> Tensor {
assert_eq!(t.ndim(), 3);
let host = t.to_device(Device::Cpu);
let elem_size = t.dtype().size_bytes();
@@ -826,7 +1019,13 @@ fn slice_scale_range(t: &Tensor, start: usize, count: usize) -> Tensor {
}
/// Extract experts [start..start+count) from a [num_experts, rows, cols] 3D tensor
fn slice_expert_range_3d(t: &Tensor, start: usize, count: usize, rows: usize, cols: usize) -> Tensor {
fn slice_expert_range_3d(
t: &Tensor,
start: usize,
count: usize,
rows: usize,
cols: usize,
) -> Tensor {
assert_eq!(t.ndim(), 3);
let host = t.to_device(Device::Cpu);
let data = host.as_slice::<bf16>();

39
crates/xserv-model/src/gpt_oss_graph.rs
View File

@@ -59,7 +59,9 @@ impl GptOssDecodeGraph {
model.decode_prepare(&[position], &[slot], cache);
ids_buf.copy_from_host(&token.to_le_bytes()).unwrap();
pos_buf.copy_from_host(&(position as u32).to_le_bytes()).unwrap();
pos_buf
.copy_from_host(&(position as u32).to_le_bytes())
.unwrap();
// Retained warmup: run the exact step once eagerly with the quarantine
// ON. Freed intermediates are held back instead of recycled, so the
@@ -88,21 +90,32 @@ impl GptOssDecodeGraph {
let logits;
{
let _guard = xserv_cuda::stream::push_stream(&stream);
graph.begin_capture(&stream).expect("begin decode-graph capture");
graph
.begin_capture(&stream)
.expect("begin decode-graph capture");
logits = model.decode_core(
ids_buf.as_ptr() as *const c_void,
pos_buf.as_ptr() as *const c_void,
1,
cache,
);
graph.end_capture(&stream).expect("end decode-graph capture");
graph
.end_capture(&stream)
.expect("end decode-graph capture");
}
let arena = allocator::end_retain();
graph.launch(&stream).expect("first decode-graph replay");
cache.advance_seq_len(slot, 1);
Self { stream, graph, ids_buf, pos_buf, logits, _arena: arena }
Self {
stream,
graph,
ids_buf,
pos_buf,
logits,
_arena: arena,
}
}
/// Run one decode step by replaying the captured graph.
@@ -116,8 +129,12 @@ impl GptOssDecodeGraph {
) -> Tensor {
model.decode_prepare(&[position], &[slot], cache);
self.ids_buf.copy_from_host(&token.to_le_bytes()).unwrap();
self.pos_buf.copy_from_host(&(position as u32).to_le_bytes()).unwrap();
self.graph.launch(&self.stream).expect("decode-graph replay");
self.pos_buf
.copy_from_host(&(position as u32).to_le_bytes())
.unwrap();
self.graph
.launch(&self.stream)
.expect("decode-graph replay");
cache.advance_seq_len(slot, 1);
// Shallow clone: the caller reads these logits before the next replay
// rewrites the underlying buffer.
@@ -137,8 +154,14 @@ pub struct GraphedGptOssDecoder {
impl GraphedGptOssDecoder {
pub fn new() -> Self {
let enabled = std::env::var("XSERV_DECODE_GRAPH").map(|v| v != "0").unwrap_or(true);
Self { graph: None, eager_steps: 0, enabled }
let enabled = std::env::var("XSERV_DECODE_GRAPH")
.map(|v| v != "0")
.unwrap_or(true);
Self {
graph: None,
eager_steps: 0,
enabled,
}
}
pub fn decode(

99
crates/xserv-model/src/kv_cache.rs
View File

@@ -1,6 +1,6 @@
use crate::config::ModelConfig;
use xserv_cuda::GpuBuffer;
use xserv_tensor::{DType, Tensor};
use crate::config::ModelConfig;
/// GPU-resident KV cache. Pre-allocates max_seq_len on GPU,
/// appends new K/V via D2D copy at offset (no CPU round-trip).
@@ -46,17 +46,43 @@ impl GpuKVCache {
v_staging.push(GpuBuffer::alloc(buf_size).expect("alloc KV staging V"));
}
Self { k_bufs, v_bufs, k_staging, v_staging, seq_len: 0, max_seq_len, num_kv_heads, head_dim, elem_size, dtype, device }
Self {
k_bufs,
v_bufs,
k_staging,
v_staging,
seq_len: 0,
max_seq_len,
num_kv_heads,
head_dim,
elem_size,
dtype,
device,
}
}
pub fn seq_len(&self) -> usize { self.seq_len }
pub fn max_seq_len(&self) -> usize { self.max_seq_len }
pub fn seq_len(&self) -> usize {
self.seq_len
}
pub fn max_seq_len(&self) -> usize {
self.max_seq_len
}
/// Append new K/V tensors for a given layer.
/// k_new, v_new: [1, num_kv_heads, new_tokens, head_dim] on GPU, contiguous.
/// `write_pos` is the sequence position to write at (caller manages this).
pub fn append(&mut self, layer: usize, k_new: &Tensor, v_new: &Tensor, new_tokens: usize, write_pos: usize) {
assert!(write_pos + new_tokens <= self.max_seq_len, "KV cache overflow");
pub fn append(
&mut self,
layer: usize,
k_new: &Tensor,
v_new: &Tensor,
new_tokens: usize,
write_pos: usize,
) {
assert!(
write_pos + new_tokens <= self.max_seq_len,
"KV cache overflow"
);
let es = self.elem_size;
let hd = self.head_dim;
let max_s = self.max_seq_len;
@@ -69,14 +95,23 @@ impl GpuKVCache {
let src_off = h * new_tokens * hd * es;
let dst_off = (h * max_s + write_pos) * hd * es;
let count = new_tokens * hd * es;
self.k_bufs[layer].copy_from_device_at(k_src, src_off, dst_off, count).unwrap();
self.v_bufs[layer].copy_from_device_at(v_src, src_off, dst_off, count).unwrap();
self.k_bufs[layer]
.copy_from_device_at(k_src, src_off, dst_off, count)
.unwrap();
self.v_bufs[layer]
.copy_from_device_at(v_src, src_off, dst_off, count)
.unwrap();
}
}
pub fn advance_seq_len(&mut self, new_tokens: usize) {
self.seq_len += new_tokens;
assert!(self.seq_len <= self.max_seq_len, "KV cache seq_len ({}) exceeds max_seq_len ({})", self.seq_len, self.max_seq_len);
assert!(
self.seq_len <= self.max_seq_len,
"KV cache seq_len ({}) exceeds max_seq_len ({})",
self.seq_len,
self.max_seq_len
);
}
/// Get K/V cache tensors for a layer up to `seq_len` tokens: [1, num_kv_heads, seq_len, head_dim]
@@ -86,7 +121,11 @@ impl GpuKVCache {
}
pub fn get_kv_len(&mut self, layer: usize, sl: usize) -> (Tensor, Tensor) {
assert!(sl <= self.max_seq_len, "get_kv_len: sl ({sl}) exceeds max_seq_len ({})", self.max_seq_len);
assert!(
sl <= self.max_seq_len,
"get_kv_len: sl ({sl}) exceeds max_seq_len ({})",
self.max_seq_len
);
let hd = self.head_dim;
let nh = self.num_kv_heads;
let es = self.elem_size;
@@ -104,8 +143,12 @@ impl GpuKVCache {
let src_off = (h * max_s) * hd * es;
let dst_off = (h * sl) * hd * es;
let count = sl * hd * es;
k_stg.copy_from_device_at(k_buf, src_off, dst_off, count).unwrap();
v_stg.copy_from_device_at(v_buf, src_off, dst_off, count).unwrap();
k_stg
.copy_from_device_at(k_buf, src_off, dst_off, count)
.unwrap();
v_stg
.copy_from_device_at(v_buf, src_off, dst_off, count)
.unwrap();
}
// Grab raw pointers before dropping the mutable borrows
let k_ptr = k_stg.as_mut_ptr();
@@ -117,20 +160,35 @@ impl GpuKVCache {
// get_kv_len call overwrites the staging buffer).
let shape = &[1usize, nh, sl, hd];
let k = unsafe {
tensor_from_gpu_buffer(GpuBuffer::borrow_raw(k_ptr, out_size), shape, self.dtype, self.device)
tensor_from_gpu_buffer(
GpuBuffer::borrow_raw(k_ptr, out_size),
shape,
self.dtype,
self.device,
)
};
let v = unsafe {
tensor_from_gpu_buffer(GpuBuffer::borrow_raw(v_ptr, out_size), shape, self.dtype, self.device)
tensor_from_gpu_buffer(
GpuBuffer::borrow_raw(v_ptr, out_size),
shape,
self.dtype,
self.device,
)
};
(k, v)
}
}
/// Create a Tensor from a GpuBuffer (takes ownership).
unsafe fn tensor_from_gpu_buffer(buf: GpuBuffer, shape: &[usize], dtype: DType, device: u32) -> Tensor {
use xserv_tensor::storage::Storage;
use xserv_tensor::shape::contiguous_strides;
unsafe fn tensor_from_gpu_buffer(
buf: GpuBuffer,
shape: &[usize],
dtype: DType,
device: u32,
) -> Tensor {
use smallvec::SmallVec;
use xserv_tensor::shape::contiguous_strides;
use xserv_tensor::storage::Storage;
let storage = Storage::cuda(buf, device);
Tensor::from_storage(
@@ -146,6 +204,11 @@ unsafe fn tensor_from_gpu_buffer(buf: GpuBuffer, shape: &[usize], dtype: DType,
///
/// # Safety
/// `buf` must be a valid GPU allocation with at least `product(shape) * dtype.size_bytes()` bytes.
pub unsafe fn tensor_from_gpu_buffer_pub(buf: GpuBuffer, shape: &[usize], dtype: DType, device: u32) -> Tensor {
pub unsafe fn tensor_from_gpu_buffer_pub(
buf: GpuBuffer,
shape: &[usize],
dtype: DType,
device: u32,
) -> Tensor {
tensor_from_gpu_buffer(buf, shape, dtype, device)
}

4
crates/xserv-model/src/lib.rs
View File

@@ -11,11 +11,11 @@ pub mod sampling;
pub use config::ModelConfig;
pub use decode_graph::{DecodeGraphState, LayerWeightPtrs};
pub use gpt2::{GPT2, KVCache};
pub use gpt_oss::GptOss;
pub use gpt_oss_graph::{GptOssDecodeGraph, GraphedGptOssDecoder};
pub use gpt2::{GPT2, KVCache};
pub use kv_cache::GpuKVCache;
pub use paged_kv_cache::{BlockAllocator, Location, PagedKVCache, BLOCK_SIZE};
pub use paged_kv_cache::{BLOCK_SIZE, BlockAllocator, Location, PagedKVCache};
pub use qwen3::Qwen3;
pub use sampling::{SamplingParams, sample, sample_greedy_penalized};

14
crates/xserv-model/src/loader.rs
View File

@@ -5,8 +5,8 @@ use std::path::Path;
use xserv_tensor::{DType, Device, Tensor};
pub fn load_safetensors(path: &Path, device: Device) -> HashMap<String, Tensor> {
let data = std::fs::read(path)
.unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
let data =
std::fs::read(path).unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
let st = SafeTensors::deserialize(&data)
.unwrap_or_else(|e| panic!("failed to parse safetensors {}: {e}", path.display()));
@@ -60,7 +60,11 @@ pub fn load_model_dir(dir: &Path, device: Device) -> HashMap<String, Tensor> {
all_tensors.extend(tensors);
}
assert!(!all_tensors.is_empty(), "no safetensors files found in {}", dir.display());
assert!(
!all_tensors.is_empty(),
"no safetensors files found in {}",
dir.display()
);
all_tensors
}
@@ -84,8 +88,6 @@ fn make_tensor(raw_bytes: &[u8], shape: &[usize], dtype: DType) -> Tensor {
};
Tensor::from_slice(bfs, shape)
}
DType::FP8E4M3 => {
Tensor::from_raw_bytes(raw_bytes, shape, DType::FP8E4M3)
}
DType::FP8E4M3 => Tensor::from_raw_bytes(raw_bytes, shape, DType::FP8E4M3),
}
}

248
crates/xserv-model/src/paged_kv_cache.rs
View File

@@ -29,7 +29,10 @@ impl BlockAllocator {
for b in (1..total_blocks).rev() {
free_stack.push(b as u32);
}
Self { free_stack, total: total_blocks }
Self {
free_stack,
total: total_blocks,
}
}
pub fn alloc(&mut self) -> Option<u32> {
@@ -136,8 +139,14 @@ impl PagedKVCache {
device: u32,
) -> Self {
Self::new_tp(
config, config.num_kv_heads(), total_blocks, cpu_total_blocks,
max_seqs, max_blocks_per_seq, dtype, device,
config,
config.num_kv_heads(),
total_blocks,
cpu_total_blocks,
max_seqs,
max_blocks_per_seq,
dtype,
device,
)
}
@@ -155,7 +164,10 @@ impl PagedKVCache {
dtype: DType,
device: u32,
) -> Self {
assert!(total_blocks >= 2, "need at least 2 blocks (one is sentinel)");
assert!(
total_blocks >= 2,
"need at least 2 blocks (one is sentinel)"
);
let num_layers = config.num_layers();
let head_dim = config.head_dim();
let elem_size = dtype.size_bytes();
@@ -179,11 +191,17 @@ impl PagedKVCache {
if cpu_total_blocks >= 2 {
let cpu_pool_bytes = cpu_total_blocks * block_bytes;
for _ in 0..num_layers {
cpu_k_pools.push(PinnedBuffer::alloc(cpu_pool_bytes).expect("alloc CPU K swap pool"));
cpu_v_pools.push(PinnedBuffer::alloc(cpu_pool_bytes).expect("alloc CPU V swap pool"));
cpu_k_pools
.push(PinnedBuffer::alloc(cpu_pool_bytes).expect("alloc CPU K swap pool"));
cpu_v_pools
.push(PinnedBuffer::alloc(cpu_pool_bytes).expect("alloc CPU V swap pool"));
}
}
let cpu_allocator = BlockAllocator::new(if cpu_total_blocks >= 2 { cpu_total_blocks } else { 0 });
let cpu_allocator = BlockAllocator::new(if cpu_total_blocks >= 2 {
cpu_total_blocks
} else {
0
});
let block_table_gpu =
GpuBuffer::alloc(max_seqs * max_blocks_per_seq * std::mem::size_of::<i32>())
@@ -220,22 +238,49 @@ impl PagedKVCache {
}
}
pub fn num_layers(&self) -> usize { self.num_layers }
pub fn num_kv_heads(&self) -> usize { self.num_kv_heads }
pub fn head_dim(&self) -> usize { self.head_dim }
pub fn dtype(&self) -> DType { self.dtype }
pub fn max_seqs(&self) -> usize { self.max_seqs }
pub fn max_blocks_per_seq(&self) -> usize { self.max_blocks_per_seq }
pub fn free_blocks(&self) -> usize { self.allocator.free_count() }
pub fn total_blocks(&self) -> usize { self.allocator.total() }
pub fn num_layers(&self) -> usize {
self.num_layers
}
pub fn num_kv_heads(&self) -> usize {
self.num_kv_heads
}
pub fn head_dim(&self) -> usize {
self.head_dim
}
pub fn dtype(&self) -> DType {
self.dtype
}
pub fn max_seqs(&self) -> usize {
self.max_seqs
}
pub fn max_blocks_per_seq(&self) -> usize {
self.max_blocks_per_seq
}
pub fn free_blocks(&self) -> usize {
self.allocator.free_count()
}
pub fn total_blocks(&self) -> usize {
self.allocator.total()
}
pub fn k_pool(&self, layer: usize) -> &GpuBuffer { &self.k_pools[layer] }
pub fn v_pool(&self, layer: usize) -> &GpuBuffer { &self.v_pools[layer] }
pub fn block_table_gpu(&self) -> &GpuBuffer { &self.block_table_gpu }
pub fn context_lens_gpu(&self) -> &GpuBuffer { &self.context_lens_gpu }
pub fn k_pool(&self, layer: usize) -> &GpuBuffer {
&self.k_pools[layer]
}
pub fn v_pool(&self, layer: usize) -> &GpuBuffer {
&self.v_pools[layer]
}
pub fn block_table_gpu(&self) -> &GpuBuffer {
&self.block_table_gpu
}
pub fn context_lens_gpu(&self) -> &GpuBuffer {
&self.context_lens_gpu
}
pub fn seq_len(&self, slot: usize) -> usize {
self.seq_states[slot].as_ref().map(|s| s.seq_len).unwrap_or(0)
self.seq_states[slot]
.as_ref()
.map(|s| s.seq_len)
.unwrap_or(0)
}
pub fn is_slot_free(&self, slot: usize) -> bool {
@@ -280,7 +325,11 @@ impl PagedKVCache {
let state = self.seq_states[slot].as_ref().expect("unregistered slot");
let cur = state.block_ids.len();
let needed_total = (state.seq_len + new_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE;
if needed_total > cur { needed_total - cur } else { 0 }
if needed_total > cur {
needed_total - cur
} else {
0
}
}
/// Pre-allocate enough physical blocks in `slot` to cover positions
@@ -290,8 +339,14 @@ impl PagedKVCache {
let state = self.seq_states[slot].as_mut().expect("unregistered slot");
let needed_total = (end_pos + BLOCK_SIZE - 1) / BLOCK_SIZE;
while state.block_ids.len() < needed_total {
let b = self.allocator.alloc().expect("out of blocks (caller must check)");
assert!(state.block_ids.len() < self.max_blocks_per_seq, "block table overflow");
let b = self
.allocator
.alloc()
.expect("out of blocks (caller must check)");
assert!(
state.block_ids.len() < self.max_blocks_per_seq,
"block table overflow"
);
state.block_ids.push(b);
}
}
@@ -318,7 +373,9 @@ impl PagedKVCache {
num_tokens: usize,
start_pos: usize,
) {
if num_tokens == 0 { return; }
if num_tokens == 0 {
return;
}
// Make sure blocks exist for the target range.
self.ensure_capacity(slot, start_pos + num_tokens);
@@ -328,15 +385,21 @@ impl PagedKVCache {
// Stage block_ids on the GPU. Pool-allocated so this is essentially
// free after the first call (same bucket every step).
let block_ids: Vec<i32> = self.seq_states[slot].as_ref().unwrap()
.block_ids.iter().map(|&b| b as i32).collect();
let block_ids: Vec<i32> = self.seq_states[slot]
.as_ref()
.unwrap()
.block_ids
.iter()
.map(|&b| b as i32)
.collect();
let bytes = block_ids.len() * std::mem::size_of::<i32>();
let mut block_ids_gpu = xserv_cuda::allocator::cached_alloc(bytes)
.expect("alloc append block_ids");
let block_ids_bytes = unsafe {
std::slice::from_raw_parts(block_ids.as_ptr() as *const u8, bytes)
};
block_ids_gpu.copy_from_host(block_ids_bytes).expect("upload block_ids");
let mut block_ids_gpu =
xserv_cuda::allocator::cached_alloc(bytes).expect("alloc append block_ids");
let block_ids_bytes =
unsafe { std::slice::from_raw_parts(block_ids.as_ptr() as *const u8, bytes) };
block_ids_gpu
.copy_from_host(block_ids_bytes)
.expect("upload block_ids");
let k_src = k_new.data_ptr() as *const std::ffi::c_void;
let v_src = v_new.data_ptr() as *const std::ffi::c_void;
@@ -345,10 +408,16 @@ impl PagedKVCache {
unsafe {
xserv_kernels::reshape_and_cache_bf16(
k_src, v_src,
k_pool_ptr, v_pool_ptr,
k_src,
v_src,
k_pool_ptr,
v_pool_ptr,
block_ids_gpu.as_ptr() as *const i32,
num_tokens, nkv, hd, start_pos, bs,
num_tokens,
nkv,
hd,
start_pos,
bs,
xserv_cuda::current_stream_raw(),
);
}
@@ -378,7 +447,9 @@ impl PagedKVCache {
v_new: &Tensor,
batch: usize,
) {
if batch == 0 { return; }
if batch == 0 {
return;
}
let nkv = self.num_kv_heads;
let hd = self.head_dim;
debug_assert_eq!(k_new.shape(), &[batch, nkv, hd]);
@@ -393,10 +464,17 @@ impl PagedKVCache {
unsafe {
xserv_kernels::reshape_and_cache_batched_bf16(
k_src, v_src,
k_pool_ptr, v_pool_ptr,
bt_ptr, cl_ptr,
batch, nkv, hd, BLOCK_SIZE, self.max_blocks_per_seq,
k_src,
v_src,
k_pool_ptr,
v_pool_ptr,
bt_ptr,
cl_ptr,
batch,
nkv,
hd,
BLOCK_SIZE,
self.max_blocks_per_seq,
xserv_cuda::current_stream_raw(),
);
}
@@ -447,7 +525,10 @@ impl PagedKVCache {
/// before advance_seq_len has run).
pub fn sync_active_batch_with_lens(&mut self, slots: &[usize], kv_lens: &[i32]) {
assert_eq!(slots.len(), kv_lens.len());
assert!(slots.len() <= self.max_seqs, "active batch exceeds max_seqs");
assert!(
slots.len() <= self.max_seqs,
"active batch exceeds max_seqs"
);
let stride = self.max_blocks_per_seq;
for row in &mut self.block_table_host {
*row = 0;
@@ -456,7 +537,9 @@ impl PagedKVCache {
*cl = 0;
}
for (i, &slot) in slots.iter().enumerate() {
let s = self.seq_states[slot].as_ref().expect("unregistered slot in active batch");
let s = self.seq_states[slot]
.as_ref()
.expect("unregistered slot in active batch");
let row = &mut self.block_table_host[i * stride..(i + 1) * stride];
for (j, b) in s.block_ids.iter().enumerate() {
row[j] = *b as i32;
@@ -515,8 +598,12 @@ impl PagedKVCache {
let src_off = ((phys * nkv + h) * bs + slot_in_blk) * hd * es;
let dst_off = (h * sl + p) * hd * es;
let count = chunk * hd * es;
k_dst.copy_from_device_at(k_pool, src_off, dst_off, count).unwrap();
v_dst.copy_from_device_at(v_pool, src_off, dst_off, count).unwrap();
k_dst
.copy_from_device_at(k_pool, src_off, dst_off, count)
.unwrap();
v_dst
.copy_from_device_at(v_pool, src_off, dst_off, count)
.unwrap();
}
p += chunk;
}
@@ -529,16 +616,26 @@ impl PagedKVCache {
// ----- Swapping (vLLM-style preemption to pinned host memory) -----
pub fn free_cpu_blocks(&self) -> usize { self.cpu_allocator.free_count() }
pub fn swap_enabled(&self) -> bool { !self.cpu_k_pools.is_empty() }
pub fn free_cpu_blocks(&self) -> usize {
self.cpu_allocator.free_count()
}
pub fn swap_enabled(&self) -> bool {
!self.cpu_k_pools.is_empty()
}
pub fn is_swapped(&self, slot: usize) -> bool {
matches!(self.seq_states[slot].as_ref().map(|s| s.location), Some(Location::Cpu))
matches!(
self.seq_states[slot].as_ref().map(|s| s.location),
Some(Location::Cpu)
)
}
/// Number of physical blocks currently held by `slot` (in either pool).
pub fn block_count(&self, slot: usize) -> usize {
self.seq_states[slot].as_ref().map(|s| s.block_ids.len()).unwrap_or(0)
self.seq_states[slot]
.as_ref()
.map(|s| s.block_ids.len())
.unwrap_or(0)
}
/// Whether a swapped sequence at `slot` can be brought back (enough free GPU blocks).
@@ -554,11 +651,17 @@ impl PagedKVCache {
/// Evict `slot`'s KV from GPU to pinned host memory and free its GPU blocks.
/// The slot stays registered (location = Cpu); the sequence is paused.
pub fn swap_out(&mut self, slot: usize) -> Result<(), &'static str> {
let state = self.seq_states[slot].as_ref().ok_or("swap_out: empty slot")?;
if state.location == Location::Cpu { return Ok(()); }
let state = self.seq_states[slot]
.as_ref()
.ok_or("swap_out: empty slot")?;
if state.location == Location::Cpu {
return Ok(());
}
let gpu_ids = state.block_ids.clone();
let n = gpu_ids.len();
if !self.cpu_allocator.can_alloc(n) { return Err("swap_out: CPU pool full"); }
if !self.cpu_allocator.can_alloc(n) {
return Err("swap_out: CPU pool full");
}
let cpu_ids: Vec<u32> = (0..n)
.map(|_| self.cpu_allocator.alloc().expect("checked can_alloc"))
@@ -570,10 +673,18 @@ impl PagedKVCache {
let g_off = gpu_ids[i] as usize * bb;
let c_off = cpu_ids[i] as usize * bb;
self.k_pools[layer]
.copy_to_host_at(&mut self.cpu_k_pools[layer].as_mut_slice()[c_off..c_off + bb], g_off, bb)
.copy_to_host_at(
&mut self.cpu_k_pools[layer].as_mut_slice()[c_off..c_off + bb],
g_off,
bb,
)
.unwrap();
self.v_pools[layer]
.copy_to_host_at(&mut self.cpu_v_pools[layer].as_mut_slice()[c_off..c_off + bb], g_off, bb)
.copy_to_host_at(
&mut self.cpu_v_pools[layer].as_mut_slice()[c_off..c_off + bb],
g_off,
bb,
)
.unwrap();
}
}
@@ -589,11 +700,17 @@ impl PagedKVCache {
/// Bring `slot`'s KV back from host to GPU and free its CPU blocks.
pub fn swap_in(&mut self, slot: usize) -> Result<(), &'static str> {
let state = self.seq_states[slot].as_ref().ok_or("swap_in: empty slot")?;
if state.location == Location::Gpu { return Ok(()); }
let state = self.seq_states[slot]
.as_ref()
.ok_or("swap_in: empty slot")?;
if state.location == Location::Gpu {
return Ok(());
}
let cpu_ids = state.block_ids.clone();
let n = cpu_ids.len();
if !self.allocator.can_alloc(n) { return Err("swap_in: GPU pool full"); }
if !self.allocator.can_alloc(n) {
return Err("swap_in: GPU pool full");
}
let gpu_ids: Vec<u32> = (0..n)
.map(|_| self.allocator.alloc().expect("checked can_alloc"))
@@ -605,10 +722,18 @@ impl PagedKVCache {
let g_off = gpu_ids[i] as usize * bb;
let c_off = cpu_ids[i] as usize * bb;
self.k_pools[layer]
.copy_from_host_at(&self.cpu_k_pools[layer].as_slice()[c_off..c_off + bb], g_off, bb)
.copy_from_host_at(
&self.cpu_k_pools[layer].as_slice()[c_off..c_off + bb],
g_off,
bb,
)
.unwrap();
self.v_pools[layer]
.copy_from_host_at(&self.cpu_v_pools[layer].as_slice()[c_off..c_off + bb], g_off, bb)
.copy_from_host_at(
&self.cpu_v_pools[layer].as_slice()[c_off..c_off + bb],
g_off,
bb,
)
.unwrap();
}
}
@@ -623,7 +748,12 @@ impl PagedKVCache {
}
}
unsafe fn tensor_from_owned_buf(buf: GpuBuffer, shape: &[usize], dtype: DType, device: u32) -> Tensor {
unsafe fn tensor_from_owned_buf(
buf: GpuBuffer,
shape: &[usize],
dtype: DType,
device: u32,
) -> Tensor {
use smallvec::SmallVec;
use xserv_tensor::shape::contiguous_strides;
use xserv_tensor::storage::Storage;

295
crates/xserv-model/src/qwen3.rs
View File

@@ -1,5 +1,5 @@
use std::collections::HashMap;
use half::bf16;
use std::collections::HashMap;
use xserv_kernels::*;
use xserv_tensor::{Device, Tensor};
@@ -13,7 +13,7 @@ pub struct Qwen3 {
embed_tokens: Tensor,
layers: Vec<Qwen3Block>,
norm: Tensor,
lm_head_t: Tensor, // precomputed transpose
lm_head_t: Tensor, // precomputed transpose
rope_cache: RopeCache,
// Tensor parallelism. `tp` is None (or world==1) for single-GPU; otherwise
// this rank holds 1/world of the heads and AllReduces after o_proj/down_proj.
@@ -28,22 +28,29 @@ pub struct Qwen3 {
}
struct Qwen3Block {
input_norm: Tensor, // [hidden]
input_norm: Tensor, // [hidden]
qkv_proj_wt: Tensor, // FUSED: [hidden, (H+2*KV)*D] — Q|K|V columns
q_dim: usize, // num_heads * head_dim (Q slice boundary)
kv_dim: usize, // num_kv_heads * head_dim (K/V slice size)
o_proj_wt: Tensor, // TRANSPOSED: [num_heads*head_dim, hidden]
q_norm: Tensor, // [head_dim]
k_norm: Tensor, // [head_dim]
post_norm: Tensor, // [hidden]
gate_up_proj_wt: Tensor, // FUSED: [hidden, 2*intermediate]
down_proj_wt: Tensor, // TRANSPOSED: [intermediate, hidden]
o_proj_wt: Tensor, // TRANSPOSED: [num_heads*head_dim, hidden]
q_norm: Tensor, // [head_dim]
k_norm: Tensor, // [head_dim]
post_norm: Tensor, // [hidden]
gate_up_proj_wt: Tensor, // FUSED: [hidden, 2*intermediate]
down_proj_wt: Tensor, // TRANSPOSED: [intermediate, hidden]
}
impl Qwen3Block {
fn q_proj_wt(&self) -> Tensor { self.qkv_proj_wt.narrow(1, 0, self.q_dim) }
fn k_proj_wt(&self) -> Tensor { self.qkv_proj_wt.narrow(1, self.q_dim, self.kv_dim) }
fn v_proj_wt(&self) -> Tensor { self.qkv_proj_wt.narrow(1, self.q_dim + self.kv_dim, self.kv_dim) }
fn q_proj_wt(&self) -> Tensor {
self.qkv_proj_wt.narrow(1, 0, self.q_dim)
}
fn k_proj_wt(&self) -> Tensor {
self.qkv_proj_wt.narrow(1, self.q_dim, self.kv_dim)
}
fn v_proj_wt(&self) -> Tensor {
self.qkv_proj_wt
.narrow(1, self.q_dim + self.kv_dim, self.kv_dim)
}
fn gate_proj_wt(&self) -> Tensor {
let half = self.gate_up_proj_wt.shape()[1] / 2;
self.gate_up_proj_wt.narrow(1, 0, half)
@@ -80,18 +87,31 @@ impl Qwen3 {
crate::init_kernels();
let dev = Device::Cuda(device);
let take = |w: &mut HashMap<String, Tensor>, name: &str| -> Tensor {
w.remove(name).unwrap_or_else(|| panic!("missing weight: {name}"))
w.remove(name)
.unwrap_or_else(|| panic!("missing weight: {name}"))
};
// Replicated weight: upload whole to this rank's device.
let repl = |t: Tensor| -> Tensor { t.to_device(dev) };
// column-parallel: keep this rank's rows of [out, in], upload, transpose → [in, out/world].
let col = |t: Tensor| -> Tensor { shard_rows(&t, rank, world).to_device(dev).transpose(0, 1).contiguous() };
let col = |t: Tensor| -> Tensor {
shard_rows(&t, rank, world)
.to_device(dev)
.transpose(0, 1)
.contiguous()
};
// row-parallel: keep this rank's cols of [out, in], upload, transpose → [in/world, out].
let row = |t: Tensor| -> Tensor { shard_cols(&t, rank, world).to_device(dev).transpose(0, 1).contiguous() };
let row = |t: Tensor| -> Tensor {
shard_cols(&t, rank, world)
.to_device(dev)
.transpose(0, 1)
.contiguous()
};
let embed_tokens = repl(take(&mut w, "model.embed_tokens.weight"));
let norm = repl(take(&mut w, "model.norm.weight"));
let lm_head_t = repl(take(&mut w, "lm_head.weight")).transpose(0, 1).contiguous();
let lm_head_t = repl(take(&mut w, "lm_head.weight"))
.transpose(0, 1)
.contiguous();
let rope_cache = RopeCache::new(
config.max_seq_len(),
@@ -102,7 +122,10 @@ impl Qwen3 {
let num_layers = config.num_layers();
let mut layers = Vec::with_capacity(num_layers);
if rank == 0 {
eprintln!("Loading+sharding weights for {} layers (world={world})...", num_layers);
eprintln!(
"Loading+sharding weights for {} layers (world={world})...",
num_layers
);
}
for i in 0..num_layers {
let p = format!("model.layers.{i}");
@@ -126,7 +149,10 @@ impl Qwen3 {
o_proj_wt: row(take(&mut w, &format!("{p}.self_attn.o_proj.weight"))),
q_norm: repl(take(&mut w, &format!("{p}.self_attn.q_norm.weight"))),
k_norm: repl(take(&mut w, &format!("{p}.self_attn.k_norm.weight"))),
post_norm: repl(take(&mut w, &format!("{p}.post_attention_layernorm.weight"))),
post_norm: repl(take(
&mut w,
&format!("{p}.post_attention_layernorm.weight"),
)),
gate_up_proj_wt,
down_proj_wt: row(take(&mut w, &format!("{p}.mlp.down_proj.weight"))),
});
@@ -165,7 +191,10 @@ impl Qwen3 {
let dev = Device::Cuda(device);
assert!(num_stages >= 1);
let num_layers = config.num_layers();
assert!(num_layers % num_stages == 0, "num_layers {num_layers} not divisible by pp {num_stages}");
assert!(
num_layers % num_stages == 0,
"num_layers {num_layers} not divisible by pp {num_stages}"
);
let per_stage = num_layers / num_stages;
let lo = stage * per_stage;
let hi = lo + per_stage;
@@ -173,16 +202,29 @@ impl Qwen3 {
let is_last_stage = stage == num_stages - 1;
let take = |w: &mut HashMap<String, Tensor>, name: &str| -> Tensor {
w.remove(name).unwrap_or_else(|| panic!("missing weight: {name}"))
w.remove(name)
.unwrap_or_else(|| panic!("missing weight: {name}"))
};
let repl = |t: Tensor| -> Tensor { t.to_device(dev) };
// Pre-transpose like the TP path's `col`/`row` do for world==1 (no shard).
let wt = |t: Tensor| -> Tensor { t.to_device(dev).transpose(0, 1).contiguous() };
let placeholder = || Tensor::from_slice(&[bf16::ZERO], &[1, 1]).to_device(dev);
let embed_tokens = if is_first_stage { repl(take(&mut w, "model.embed_tokens.weight")) } else { placeholder() };
let norm = if is_last_stage { repl(take(&mut w, "model.norm.weight")) } else { placeholder() };
let lm_head_t = if is_last_stage { wt(take(&mut w, "lm_head.weight")) } else { placeholder() };
let embed_tokens = if is_first_stage {
repl(take(&mut w, "model.embed_tokens.weight"))
} else {
placeholder()
};
let norm = if is_last_stage {
repl(take(&mut w, "model.norm.weight"))
} else {
placeholder()
};
let lm_head_t = if is_last_stage {
wt(take(&mut w, "lm_head.weight"))
} else {
placeholder()
};
let rope_cache = RopeCache::new(
config.max_seq_len(),
@@ -217,7 +259,10 @@ impl Qwen3 {
o_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.o_proj.weight"))),
q_norm: repl(take(&mut w, &format!("{p}.self_attn.q_norm.weight"))),
k_norm: repl(take(&mut w, &format!("{p}.self_attn.k_norm.weight"))),
post_norm: repl(take(&mut w, &format!("{p}.post_attention_layernorm.weight"))),
post_norm: repl(take(
&mut w,
&format!("{p}.post_attention_layernorm.weight"),
)),
gate_up_proj_wt,
down_proj_wt: wt(take(&mut w, &format!("{p}.mlp.down_proj.weight"))),
});
@@ -252,8 +297,12 @@ impl Qwen3 {
matmul_2d(&x, &self.lm_head_t)
}
pub fn pp_is_first(&self) -> bool { self.is_first_stage }
pub fn pp_is_last(&self) -> bool { self.is_last_stage }
pub fn pp_is_first(&self) -> bool {
self.is_first_stage
}
pub fn pp_is_last(&self) -> bool {
self.is_last_stage
}
/// PP prefill over THIS stage's layers. `x` is `[S, hidden]` (stage 0: from
/// `embed`; otherwise received from the previous stage). Writes K/V for this
@@ -276,7 +325,9 @@ impl Qwen3 {
paged_cache.ensure_capacity(slot, pos_offset + new_tokens);
paged_cache.advance_seq_len(slot, new_tokens);
let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens)
.map(|p| p as u32)
.collect();
for (layer_idx, layer) in self.layers.iter().enumerate() {
let residual = x.clone();
@@ -285,7 +336,9 @@ impl Qwen3 {
let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
let q = qkv.narrow(1, 0, layer.q_dim).contiguous();
let k = qkv.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
let v = qkv.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim).contiguous();
let v = qkv
.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
.contiguous();
let q = xserv_kernels::reshape_heads_gpu(&q, new_tokens, num_heads, head_dim);
let k = xserv_kernels::reshape_heads_gpu(&k, new_tokens, num_kv_heads, head_dim);
@@ -305,10 +358,12 @@ impl Qwen3 {
let (k_full, v_full) = paged_cache.gather_kv_contiguous(slot, layer_idx);
let attn_out = flash_attention(&q, &k_full, &v_full, true);
let attn_merged = xserv_kernels::merge_heads_gpu(&attn_out, new_tokens, num_heads, head_dim);
let attn_merged =
xserv_kernels::merge_heads_gpu(&attn_out, new_tokens, num_heads, head_dim);
let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let (normed, x_new) =
xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let residual = x_new.clone();
let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
@@ -356,7 +411,9 @@ impl Qwen3 {
let qkv_all = matmul_2d(&normed, &layer.qkv_proj_wt);
let q_all = qkv_all.narrow(1, 0, layer.q_dim).contiguous();
let k_all = qkv_all.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
let v_all = qkv_all.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim).contiguous();
let v_all = qkv_all
.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
.contiguous();
let mut q_rows: Vec<Tensor> = Vec::with_capacity(batch);
for b in 0..batch {
@@ -394,14 +451,23 @@ impl Qwen3 {
let v_pool_ptr = paged_cache.v_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
let attn_out = xserv_kernels::paged_decode_attention(
&q_4d, k_pool_ptr, v_pool_ptr, bt_ptr, cl_ptr,
batch, num_heads, num_kv_heads, head_dim, max_blocks,
&q_4d,
k_pool_ptr,
v_pool_ptr,
bt_ptr,
cl_ptr,
batch,
num_heads,
num_kv_heads,
head_dim,
max_blocks,
);
let attn_merged = attn_out.reshape(&[batch, num_heads * head_dim]);
let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let (normed, x_new) =
xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let residual = x_new.clone();
let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
@@ -441,7 +507,9 @@ impl Qwen3 {
let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
let mut x = embedding(&self.embed_tokens, token_ids);
let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens)
.map(|p| p as u32)
.collect();
for (layer_idx, layer) in self.layers.iter().enumerate() {
let residual = x.clone();
@@ -450,7 +518,9 @@ impl Qwen3 {
let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
let q = qkv.narrow(1, 0, layer.q_dim).contiguous();
let k = qkv.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
let v = qkv.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim).contiguous();
let v = qkv
.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
.contiguous();
let q = reshape_heads(&q, new_tokens, num_heads, head_dim);
let k = reshape_heads(&k, new_tokens, num_kv_heads, head_dim);
@@ -531,7 +601,9 @@ impl Qwen3 {
let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
let q_all = qkv.narrow(1, 0, layer.q_dim).contiguous();
let k_all = qkv.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
let v_all = qkv.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim).contiguous();
let v_all = qkv
.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
.contiguous();
// Per-sequence: reshape, qk-norm, RoPE, KV cache, attention, merge
let mut attn_outputs: Vec<Tensor> = Vec::with_capacity(batch);
@@ -583,7 +655,8 @@ impl Qwen3 {
let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
// Fused add + rmsnorm
let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let (normed, x_new) =
xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let residual = x_new.clone();
let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
@@ -662,13 +735,15 @@ impl Qwen3 {
let qkv = matmul_2d(&normed, &layer.qkv_proj_wt); // [B, (H+2*KV)*D]
let q_dim = num_heads * head_dim;
let kv_dim = num_kv_heads * head_dim;
let q_all = qkv.narrow(1, 0, q_dim); // [B, H*D] (view)
let k_all = qkv.narrow(1, q_dim, kv_dim); // [B, KV*D] (view)
let q_all = qkv.narrow(1, 0, q_dim); // [B, H*D] (view)
let k_all = qkv.narrow(1, q_dim, kv_dim); // [B, KV*D] (view)
let v_all = qkv.narrow(1, q_dim + kv_dim, kv_dim);
// Per-head RMSNorm on contiguous copies (narrow views are strided).
let q_flat = q_all.contiguous().reshape(&[batch * num_heads, head_dim]);
let k_flat = k_all.contiguous().reshape(&[batch * num_kv_heads, head_dim]);
let k_flat = k_all
.contiguous()
.reshape(&[batch * num_kv_heads, head_dim]);
let q_normed = rmsnorm(&q_flat, &layer.q_norm, eps);
let k_normed = rmsnorm(&k_flat, &layer.k_norm, eps);
@@ -688,8 +763,16 @@ impl Qwen3 {
let k_pool_ptr = paged_cache.k_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
let v_pool_ptr = paged_cache.v_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
let attn_out = xserv_kernels::paged_decode_attention(
&q_4d, k_pool_ptr, v_pool_ptr, bt_ptr, cl_ptr,
batch, num_heads, num_kv_heads, head_dim, max_blocks,
&q_4d,
k_pool_ptr,
v_pool_ptr,
bt_ptr,
cl_ptr,
batch,
num_heads,
num_kv_heads,
head_dim,
max_blocks,
);
// attn_out shape [B, H, 1, D] is contiguous-equivalent to [B, H*D].
@@ -697,7 +780,8 @@ impl Qwen3 {
let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
self.all_reduce(&attn_proj); // TP: sum partial attention outputs
let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let (normed, x_new) =
xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let residual = x_new.clone();
// Fused gate+up projection: one GEMV instead of two.
@@ -743,7 +827,9 @@ impl Qwen3 {
paged_cache.advance_seq_len(slot, new_tokens);
let mut x = embedding(&self.embed_tokens, token_ids);
let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens)
.map(|p| p as u32)
.collect();
for (layer_idx, layer) in self.layers.iter().enumerate() {
let residual = x.clone();
@@ -752,7 +838,9 @@ impl Qwen3 {
let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
let q = qkv.narrow(1, 0, layer.q_dim).contiguous();
let k = qkv.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
let v = qkv.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim).contiguous();
let v = qkv
.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
.contiguous();
let q = xserv_kernels::reshape_heads_gpu(&q, new_tokens, num_heads, head_dim);
let k = xserv_kernels::reshape_heads_gpu(&k, new_tokens, num_kv_heads, head_dim);
@@ -773,11 +861,13 @@ impl Qwen3 {
let (k_full, v_full) = paged_cache.gather_kv_contiguous(slot, layer_idx);
let attn_out = flash_attention(&q, &k_full, &v_full, true);
let attn_merged = xserv_kernels::merge_heads_gpu(&attn_out, new_tokens, num_heads, head_dim);
let attn_merged =
xserv_kernels::merge_heads_gpu(&attn_out, new_tokens, num_heads, head_dim);
let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
self.all_reduce(&attn_proj);
let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let (normed, x_new) =
xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let residual = x_new.clone();
let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
@@ -805,7 +895,9 @@ impl Qwen3 {
let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
let mut x = embedding(&self.embed_tokens, token_ids);
let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens)
.map(|p| p as u32)
.collect();
for (layer_idx, layer) in self.layers.iter().enumerate() {
let residual = x.clone();
@@ -814,7 +906,9 @@ impl Qwen3 {
let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
let q = qkv.narrow(1, 0, layer.q_dim).contiguous();
let k = qkv.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
let v = qkv.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim).contiguous();
let v = qkv
.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
.contiguous();
let q = xserv_kernels::reshape_heads_gpu(&q, new_tokens, num_heads, head_dim);
let k = xserv_kernels::reshape_heads_gpu(&k, new_tokens, num_kv_heads, head_dim);
@@ -834,10 +928,12 @@ impl Qwen3 {
let (k_full, v_full) = cache.get_kv_len(layer_idx, pos_offset + new_tokens);
let attn_out = flash_attention(&q, &k_full, &v_full, true);
let attn_merged = xserv_kernels::merge_heads_gpu(&attn_out, new_tokens, num_heads, head_dim);
let attn_merged =
xserv_kernels::merge_heads_gpu(&attn_out, new_tokens, num_heads, head_dim);
let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let (normed, x_new) =
xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let residual = x_new.clone();
let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
@@ -856,28 +952,33 @@ impl Qwen3 {
/// Extract weight pointers for CUDA Graph capture.
pub fn layer_weight_ptrs(&self) -> Vec<crate::decode_graph::LayerWeightPtrs> {
self.layers.iter().map(|l| crate::decode_graph::LayerWeightPtrs {
input_norm: l.input_norm.data_ptr() as *const std::ffi::c_void,
q_proj_wt: l.q_proj_wt().data_ptr() as *const std::ffi::c_void,
k_proj_wt: l.k_proj_wt().data_ptr() as *const std::ffi::c_void,
v_proj_wt: l.v_proj_wt().data_ptr() as *const std::ffi::c_void,
o_proj_wt: l.o_proj_wt.data_ptr() as *const std::ffi::c_void,
q_norm: l.q_norm.data_ptr() as *const std::ffi::c_void,
k_norm: l.k_norm.data_ptr() as *const std::ffi::c_void,
post_norm: l.post_norm.data_ptr() as *const std::ffi::c_void,
gate_proj_wt: l.gate_proj_wt().data_ptr() as *const std::ffi::c_void,
up_proj_wt: l.up_proj_wt().data_ptr() as *const std::ffi::c_void,
down_proj_wt: l.down_proj_wt.data_ptr() as *const std::ffi::c_void,
}).collect()
self.layers
.iter()
.map(|l| crate::decode_graph::LayerWeightPtrs {
input_norm: l.input_norm.data_ptr() as *const std::ffi::c_void,
q_proj_wt: l.q_proj_wt().data_ptr() as *const std::ffi::c_void,
k_proj_wt: l.k_proj_wt().data_ptr() as *const std::ffi::c_void,
v_proj_wt: l.v_proj_wt().data_ptr() as *const std::ffi::c_void,
o_proj_wt: l.o_proj_wt.data_ptr() as *const std::ffi::c_void,
q_norm: l.q_norm.data_ptr() as *const std::ffi::c_void,
k_norm: l.k_norm.data_ptr() as *const std::ffi::c_void,
post_norm: l.post_norm.data_ptr() as *const std::ffi::c_void,
gate_proj_wt: l.gate_proj_wt().data_ptr() as *const std::ffi::c_void,
up_proj_wt: l.up_proj_wt().data_ptr() as *const std::ffi::c_void,
down_proj_wt: l.down_proj_wt.data_ptr() as *const std::ffi::c_void,
})
.collect()
}
/// Get pointers needed for CUDA Graph capture.
pub fn graph_capture_ptrs(&self) -> (
*const std::ffi::c_void, // norm weight
*const std::ffi::c_void, // lm_head_t
*const std::ffi::c_void, // embed_tokens
*const std::ffi::c_void, // rope cos
*const std::ffi::c_void, // rope sin
pub fn graph_capture_ptrs(
&self,
) -> (
*const std::ffi::c_void, // norm weight
*const std::ffi::c_void, // lm_head_t
*const std::ffi::c_void, // embed_tokens
*const std::ffi::c_void, // rope cos
*const std::ffi::c_void, // rope sin
) {
(
self.norm.data_ptr() as *const std::ffi::c_void,
@@ -895,11 +996,16 @@ impl Qwen3 {
/// (column-parallel split: split the OUTPUT dim). `world==1` returns the whole.
/// Input must be a contiguous CPU (or device) BF16 tensor.
fn shard_rows(t: &Tensor, rank: usize, world: usize) -> Tensor {
if world == 1 { return t.clone(); }
if world == 1 {
return t.clone();
}
let shape = t.shape();
assert_eq!(shape.len(), 2, "shard_rows expects 2D weight");
let (rows, cols) = (shape[0], shape[1]);
assert!(rows % world == 0, "rows {rows} not divisible by world {world}");
assert!(
rows % world == 0,
"rows {rows} not divisible by world {world}"
);
let local = rows / world;
let host = t.to_device(Device::Cpu);
let data = host.as_slice::<bf16>();
@@ -911,11 +1017,16 @@ fn shard_rows(t: &Tensor, rank: usize, world: usize) -> Tensor {
/// Keep this rank's column-block of a 2D `[rows, cols]` BF16 tensor (row-parallel
/// split: split the INPUT dim). Strided copy. `world==1` returns the whole.
fn shard_cols(t: &Tensor, rank: usize, world: usize) -> Tensor {
if world == 1 { return t.clone(); }
if world == 1 {
return t.clone();
}
let shape = t.shape();
assert_eq!(shape.len(), 2, "shard_cols expects 2D weight");
let (rows, cols) = (shape[0], shape[1]);
assert!(cols % world == 0, "cols {cols} not divisible by world {world}");
assert!(
cols % world == 0,
"cols {cols} not divisible by world {world}"
);
let local = cols / world;
let c0 = rank * local;
let host = t.to_device(Device::Cpu);
@@ -1009,7 +1120,9 @@ fn transpose_from_rope(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: u
}
fn repeat_kv(x: &Tensor, n_rep: usize) -> Tensor {
if n_rep == 1 { return x.clone(); }
if n_rep == 1 {
return x.clone();
}
let kv_heads = x.shape()[1];
let seq_len = x.shape()[2];
let head_dim = x.shape()[3];
@@ -1065,11 +1178,16 @@ fn concat_rows(rows: &[Tensor]) -> Tensor {
let src_buf = row.storage().gpu_buffer();
let src_offset = row.offset() * elem_size;
let dst_offset = b * row_bytes;
out_buf.copy_from_device_at(src_buf, src_offset, dst_offset, row_bytes).unwrap();
out_buf
.copy_from_device_at(src_buf, src_offset, dst_offset, row_bytes)
.unwrap();
}
// Wrap in a Tensor
let device_id = match device { Device::Cuda(id) => id, _ => panic!("expected CUDA device") };
let device_id = match device {
Device::Cuda(id) => id,
_ => panic!("expected CUDA device"),
};
unsafe {
crate::kv_cache::tensor_from_gpu_buffer_pub(out_buf, &[batch, cols], dtype, device_id)
}
@@ -1082,12 +1200,15 @@ fn cat_cols(tensors: &[&Tensor]) -> Tensor {
let dtype = tensors[0].dtype();
let device = tensors[0].device();
let elem = dtype.size_bytes();
let total_cols: usize = tensors.iter().map(|t| {
assert_eq!(t.ndim(), 2);
assert_eq!(t.shape()[0], rows);
assert!(t.is_contiguous());
t.shape()[1]
}).sum();
let total_cols: usize = tensors
.iter()
.map(|t| {
assert_eq!(t.ndim(), 2);
assert_eq!(t.shape()[0], rows);
assert!(t.is_contiguous());
t.shape()[1]
})
.sum();
let out = Tensor::empty(&[rows, total_cols], dtype, device);
let dst_base = out.data_ptr() as *mut u8;
for r in 0..rows {
@@ -1126,7 +1247,9 @@ pub fn sample_greedy(logits: &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()
}

17
crates/xserv-model/src/sampling.rs
View File

@@ -11,7 +11,11 @@ pub struct SamplingParams {
impl Default for SamplingParams {
fn default() -> Self {
Self { temperature: 0.0, top_k: 0, top_p: 1.0 }
Self {
temperature: 0.0,
top_k: 0,
top_p: 1.0,
}
}
}
@@ -134,9 +138,14 @@ pub fn sample_greedy_penalized(logits: &Tensor, recent: &[u32], penalty: f32) ->
let seq_len = logits.shape()[0];
let logits_cpu = logits.to_device(Device::Cpu);
let mut last_row: Vec<f32> = match logits.dtype() {
DType::F32 => logits_cpu.as_slice::<f32>()[(seq_len - 1) * vocab_size..seq_len * vocab_size].to_vec(),
DType::BF16 => logits_cpu.as_slice::<bf16>()[(seq_len - 1) * vocab_size..seq_len * vocab_size]
.iter().map(|v| v.to_f32()).collect(),
DType::F32 => {
logits_cpu.as_slice::<f32>()[(seq_len - 1) * vocab_size..seq_len * vocab_size].to_vec()
}
DType::BF16 => logits_cpu.as_slice::<bf16>()
[(seq_len - 1) * vocab_size..seq_len * vocab_size]
.iter()
.map(|v| v.to_f32())
.collect(),
_ => panic!("unsupported dtype for sampling: {:?}", logits.dtype()),
};
if penalty > 1.0 {