style: format Rust workspace

crates/xserv-model/src/bin/bench-tp.rs

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