style: format Rust workspace
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user