docs: TP=1/2/4 xserv vs llama.cpp benchmark results

AIME 2025 + GSM8K at TP=1/2/4. Quality on par across engines/TP. Opposite
perf scaling: xserv TPOT improves with TP (21->17->15ms) while llama.cpp
row-split regresses over PCIe (10->19->20ms), crossing over so xserv is faster
at TP=4. Includes the clean same-path bench-tp scaling (58/76/86 tok/s).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Cargo.toml

@@ -7,6 +7,7 @@ members = [
     "crates/xserv-model",
     "crates/xserv-tokenizer",
     "crates/xserv-server",
+    "crates/xserv-distributed",
 ]
 
 [workspace.package]

crates/xserv-distributed/Cargo.toml

@@ -0,0 +1,8 @@
+[package]
+name = "xserv-distributed"
+version.workspace = true
+edition.workspace = true
+
+[dependencies]
+xserv-cuda = { path = "../xserv-cuda" }
+half.workspace = true

crates/xserv-distributed/build.rs

@@ -0,0 +1,13 @@
+use std::env;
+
+fn main() {
+    let cuda_path = env::var("CUDA_HOME")
+        .or_else(|_| env::var("CUDA_PATH"))
+        .unwrap_or_else(|_| "/usr/local/cuda".to_string());
+
+    println!("cargo:rustc-link-search=native={cuda_path}/lib64");
+    // NCCL is typically installed as a system library.
+    println!("cargo:rustc-link-search=native=/usr/lib/x86_64-linux-gnu");
+    println!("cargo:rustc-link-lib=dylib=nccl");
+    println!("cargo:rustc-link-lib=dylib=cudart");
+}

crates/xserv-distributed/src/ffi.rs

@@ -0,0 +1,65 @@
+//! Minimal NCCL FFI bindings (hand-written, like the CUDA bindings).
+//! Only the collectives we need for tensor parallelism.
+
+use std::ffi::c_void;
+use std::os::raw::c_char;
+use xserv_cuda::ffi::CudaStream;
+
+/// Opaque NCCL communicator handle (`ncclComm_t`).
+pub type NcclComm = *mut c_void;
+
+/// `ncclUniqueId` is a 128-byte opaque blob shared from rank 0 to all ranks.
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct NcclUniqueId {
+    pub internal: [c_char; 128],
+}
+
+impl Default for NcclUniqueId {
+    fn default() -> Self {
+        Self { internal: [0; 128] }
+    }
+}
+
+// ncclDataType_t (subset)
+pub const NCCL_FLOAT32: i32 = 7;
+pub const NCCL_BF16: i32 = 9;
+
+// ncclRedOp_t
+pub const NCCL_SUM: i32 = 0;
+
+// ncclResult_t
+pub const NCCL_SUCCESS: i32 = 0;
+
+unsafe extern "C" {
+    pub fn ncclGetUniqueId(uid: *mut NcclUniqueId) -> i32;
+    // ncclUniqueId is passed BY VALUE (a 128-byte struct) per the NCCL ABI.
+    pub fn ncclCommInitRank(comm: *mut NcclComm, nranks: i32, commid: NcclUniqueId, rank: i32) -> i32;
+    pub fn ncclCommDestroy(comm: NcclComm) -> i32;
+    pub fn ncclAllReduce(
+        sendbuff: *const c_void,
+        recvbuff: *mut c_void,
+        count: usize,
+        datatype: i32,
+        op: i32,
+        comm: NcclComm,
+        stream: CudaStream,
+    ) -> i32;
+    pub fn ncclGroupStart() -> i32;
+    pub fn ncclGroupEnd() -> i32;
+    pub fn ncclGetErrorString(result: i32) -> *const c_char;
+}
+
+pub fn err_string(result: i32) -> String {
+    unsafe {
+        let p = ncclGetErrorString(result);
+        if p.is_null() {
+            return format!("nccl error {result}");
+        }
+        std::ffi::CStr::from_ptr(p).to_string_lossy().into_owned()
+    }
+}
+
+pub fn check(result: i32, what: &str) {
+    assert_eq!(result, NCCL_SUCCESS, "{what} failed: {}", err_string(result));
+}

crates/xserv-distributed/src/lib.rs

@@ -0,0 +1,97 @@
+//! Tensor-parallel primitives for xserv.
+//!
+//! Process model: one OS thread per TP rank, each bound to one GPU. NCCL is
+//! used for the collective (AllReduce); a hand-rolled P2P AllReduce may replace
+//! it later as a learning exercise (see docs/17-tensor-parallelism.md).
+
+pub mod ffi;
+
+use std::ffi::c_void;
+
+use ffi::{NcclComm, NcclUniqueId};
+use xserv_cuda::device;
+use xserv_cuda::GpuBuffer;
+
+pub use ffi::NcclUniqueId as UniqueId;
+
+/// The CUDA "null" (default) stream. The model's kernels and cuBLAS calls run
+/// on it, so issuing NCCL on the same stream keeps AllReduce correctly ordered
+/// after the producing matmul and before the consuming kernel — no extra sync.
+const NULL_STREAM: xserv_cuda::ffi::CudaStream = std::ptr::null_mut();
+
+/// Generate a unique id on one rank (typically rank 0) and broadcast the bytes
+/// to all ranks out-of-band (e.g. via a shared variable across threads).
+pub fn get_unique_id() -> NcclUniqueId {
+    let mut id = NcclUniqueId::default();
+    ffi::check(unsafe { ffi::ncclGetUniqueId(&mut id) }, "ncclGetUniqueId");
+    id
+}
+
+/// Per-rank tensor-parallel context: NCCL communicator + a dedicated stream.
+pub struct TpContext {
+    pub rank: usize,
+    pub world: usize,
+    pub device: u32,
+    comm: NcclComm,
+}
+
+// The NCCL communicator is owned by exactly one rank thread.
+unsafe impl Send for TpContext {}
+
+impl TpContext {
+    /// Initialize this rank. Must be called from the thread that will own this
+    /// rank's GPU work; binds the thread to `device` first. All ranks must call
+    /// this concurrently with the same `id` and `world`.
+    pub fn init(rank: usize, world: usize, id: NcclUniqueId, device: u32) -> Self {
+        device::set_device(device).expect("set_device");
+        let mut comm: NcclComm = std::ptr::null_mut();
+        // Wrap the concurrent inits in a group so they rendezvous without deadlock.
+        ffi::check(unsafe { ffi::ncclGroupStart() }, "ncclGroupStart(init)");
+        ffi::check(
+            unsafe { ffi::ncclCommInitRank(&mut comm, world as i32, id, rank as i32) },
+            "ncclCommInitRank",
+        );
+        ffi::check(unsafe { ffi::ncclGroupEnd() }, "ncclGroupEnd(init)");
+        Self { rank, world, device, comm }
+    }
+
+    /// In-place AllReduce(sum) over `count` BF16 elements in `buf`.
+    pub fn all_reduce_sum_bf16(&self, buf: &mut GpuBuffer, count: usize) {
+        self.all_reduce_sum_bf16_ptr(buf.as_mut_ptr() as *mut c_void, count);
+    }
+
+    /// In-place AllReduce(sum) directly on a device pointer (`count` BF16 elems),
+    /// issued on the null stream so it is ordered with the model's kernels.
+    /// Asynchronous: a later sync (e.g. the D2H logits copy) completes it.
+    ///
+    /// # Safety
+    /// `ptr` must point to at least `count` BF16 elements of valid device memory
+    /// on this rank's device. The reduction is in-place (send == recv).
+    pub fn all_reduce_sum_bf16_ptr(&self, ptr: *mut c_void, count: usize) {
+        if self.world == 1 {
+            return; // nothing to reduce
+        }
+        ffi::check(
+            unsafe {
+                ffi::ncclAllReduce(
+                    ptr as *const c_void,
+                    ptr,
+                    count,
+                    ffi::NCCL_BF16,
+                    ffi::NCCL_SUM,
+                    self.comm,
+                    NULL_STREAM,
+                )
+            },
+            "ncclAllReduce",
+        );
+    }
+}
+
+impl Drop for TpContext {
+    fn drop(&mut self) {
+        if !self.comm.is_null() {
+            unsafe { ffi::ncclCommDestroy(self.comm) };
+        }
+    }
+}

crates/xserv-distributed/tests/allreduce.rs

@@ -0,0 +1,50 @@
+//! 2-GPU AllReduce smoke test. Skips if fewer than 2 GPUs are present.
+
+use half::bf16;
+use std::thread;
+use xserv_cuda::{device, GpuBuffer};
+use xserv_distributed::{get_unique_id, TpContext};
+
+#[test]
+fn allreduce_two_gpu_sum() {
+    let world = 2usize;
+    if device::device_count().unwrap_or(0) < world as i32 {
+        eprintln!("skip: need >= {world} GPUs");
+        return;
+    }
+
+    let id = get_unique_id();
+    let n = 4096usize;
+
+    let handles: Vec<_> = (0..world)
+        .map(|rank| {
+            let id = id;
+            thread::spawn(move || {
+                let ctx = TpContext::init(rank, world, id, rank as u32);
+
+                // Rank r fills its buffer with (r + 1).
+                let val = bf16::from_f32((rank + 1) as f32);
+                let host = vec![val; n];
+                let src = unsafe {
+                    std::slice::from_raw_parts(host.as_ptr() as *const u8, n * 2)
+                };
+                let mut buf = GpuBuffer::alloc(n * 2).unwrap();
+                buf.copy_from_host(src).unwrap();
+
+                ctx.all_reduce_sum_bf16(&mut buf, n);
+
+                let mut out = vec![0u8; n * 2];
+                buf.copy_to_host(&mut out).unwrap();
+                let res = unsafe { std::slice::from_raw_parts(out.as_ptr() as *const bf16, n) };
+                (res[0].to_f32(), res[n - 1].to_f32())
+            })
+        })
+        .collect();
+
+    // sum over ranks of (r+1) = 1 + 2 = 3
+    for h in handles {
+        let (first, last) = h.join().unwrap();
+        assert_eq!(first, 3.0, "AllReduce(sum) first element");
+        assert_eq!(last, 3.0, "AllReduce(sum) last element");
+    }
+}

crates/xserv-model/Cargo.toml

@@ -8,6 +8,7 @@ xserv-cuda = { path = "../xserv-cuda" }
 xserv-tensor = { path = "../xserv-tensor" }
 xserv-kernels = { path = "../xserv-kernels" }
 xserv-tokenizer = { path = "../xserv-tokenizer" }
+xserv-distributed = { path = "../xserv-distributed" }
 half.workspace = true
 smallvec.workspace = true
 serde.workspace = true

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

@@ -0,0 +1,194 @@
+//! Tensor-parallel E2E benchmark for Qwen3.
+//!
+//! Spawns one thread per TP rank (each bound to one GPU), loads the sharded
+//! model, and runs greedy autoregressive generation. Because lm_head is
+//! replicated and the post-AllReduce hidden state is identical on every rank,
+//! all ranks compute identical logits and pick the same greedy token — so the
+//! rank threads stay in lockstep via the per-layer AllReduces without any
+//! token broadcast. Rank 0 records output + timings.
+//!
+//! Usage: bench-tp <model-dir> [--tp N] [--gen-tokens N] [--devices 0,1,2,3]
+//!
+//! Run with --tp 1 / 2 / 4 and compare the printed text (correctness) and
+//! tok/s (performance).
+
+use std::path::PathBuf;
+use std::sync::Arc;
+use std::thread;
+use std::time::Instant;
+
+use xserv_model::qwen3::sample_greedy;
+use xserv_model::{loader, ModelConfig, PagedKVCache, Qwen3, BLOCK_SIZE};
+use xserv_tensor::{DType, Device};
+use xserv_tokenizer::Tokenizer;
+
+struct PromptResult {
+    gen_ids: Vec<u32>,
+    ttft_ms: f64,
+    decode_tok_s: f64,
+}
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    if args.len() < 2 {
+        eprintln!("Usage: bench-tp <model-dir> [--tp N] [--gen-tokens N] [--devices 0,1,2,3]");
+        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 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(),
+    };
+    assert_eq!(devices.len(), world, "--devices count must equal --tp");
+
+    let config = ModelConfig::from_file(&model_dir.join("config.json"));
+    assert!(
+        config.num_kv_heads() % world == 0,
+        "num_kv_heads {} not divisible by tp {world}",
+        config.num_kv_heads()
+    );
+    let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
+    let eos = tokenizer.eos_token_id();
+
+    let prompts: Vec<&str> = vec![
+        "The capital of France is",
+        "Explain photosynthesis in one sentence.",
+        "Write a haiku about the ocean.",
+        "List three uses of a hammer.",
+        "What is the speed of light?",
+        "Describe the water cycle briefly.",
+        "Who wrote Romeo and Juliet?",
+        "Translate 'good morning' into Spanish.",
+    ];
+    let prompt_ids: Vec<Vec<u32>> = prompts.iter().map(|p| tokenizer.encode(p)).collect();
+
+    // 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 handles: Vec<_> = (0..world)
+        .map(|rank| {
+            let model_dir = model_dir.clone();
+            let config = config.clone();
+            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)
+            })
+        })
+        .collect();
+
+    let mut rank0: Option<Vec<PromptResult>> = None;
+    for (rank, h) in handles.into_iter().enumerate() {
+        let r = h.join().expect("rank thread panicked");
+        if rank == 0 {
+            rank0 = r;
+        }
+    }
+
+    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");
+    let mut tps_sum = 0.0;
+    for (i, r) in results.iter().enumerate() {
+        let text = tokenizer.decode(&r.gen_ids).replace('\n', " ");
+        let short: String = text.chars().take(50).collect();
+        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
+        );
+        tps_sum += r.decode_tok_s;
+    }
+    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(","))
+        .collect();
+    println!("CORRECTNESS_IDS tp={world} {}", all_ids.join(" | "));
+}
+
+fn run_rank(
+    rank: usize,
+    world: usize,
+    device: u32,
+    id: Option<xserv_distributed::UniqueId>,
+    config: ModelConfig,
+    model_dir: PathBuf,
+    prompt_ids: Vec<Vec<u32>>,
+    gen_tokens: usize,
+    eos: Option<u32>,
+) -> 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)))
+    } else {
+        xserv_cuda::device::set_device(device).unwrap();
+        None
+    };
+
+    // Load this rank's own CPU copy of the weights and shard in-thread.
+    let weights = loader::load_model_dir(&model_dir, Device::Cpu);
+    let model = Qwen3::from_weights_tp(config.clone(), weights, rank, world, device, tp.clone());
+
+    // Per-rank paged KV cache holds only this rank's local KV heads.
+    let local_kv = config.num_kv_heads() / world;
+    let max_seq = 2048usize;
+    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,
+    );
+
+    // Warmup (init kernels / allocator / NCCL channels) — not timed.
+    cache.register_sequence(0).unwrap();
+    let _ = model.forward_prefill_paged(&[1u32, 2, 3], 0, &mut cache);
+    cache.free_sequence(0);
+
+    let mut out = Vec::new();
+    for ids in &prompt_ids {
+        cache.register_sequence(0).unwrap();
+
+        // Prefill (TTFT).
+        let t0 = Instant::now();
+        let logits = model.forward_prefill_paged(ids, 0, &mut cache);
+        let first = sample_greedy(&logits);
+        let ttft_ms = t0.elapsed().as_secs_f64() * 1000.0;
+
+        let mut generated = vec![first];
+
+        // Decode.
+        let t1 = Instant::now();
+        let mut steps = 0usize;
+        for _ in 1..gen_tokens {
+            let last = *generated.last().unwrap();
+            if eos == Some(last) {
+                break;
+            }
+            let pos = cache.seq_len(0);
+            let logits = model.forward_decode_paged(&[last], &[pos], &[0], &mut cache);
+            let next = sample_greedy(&logits);
+            generated.push(next);
+            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 };
+
+        cache.free_sequence(0);
+
+        if rank == 0 {
+            out.push(PromptResult { gen_ids: generated, ttft_ms, decode_tok_s });
+        }
+    }
+
+    if rank == 0 { Some(out) } else { None }
+}
+
+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())
+}

crates/xserv-model/src/paged_kv_cache.rs

@@ -134,10 +134,29 @@ impl PagedKVCache {
         max_blocks_per_seq: usize,
         dtype: DType,
         device: u32,
+    ) -> Self {
+        Self::new_tp(
+            config, config.num_kv_heads(), total_blocks, cpu_total_blocks,
+            max_seqs, max_blocks_per_seq, dtype, device,
+        )
+    }
+
+    /// Like `new`, but with an explicit `num_kv_heads` — under tensor parallelism
+    /// each rank only stores its `num_kv_heads / world` heads, so the pool is
+    /// sized for the local head count, not the model's full count.
+    #[allow(clippy::too_many_arguments)]
+    pub fn new_tp(
+        config: &ModelConfig,
+        num_kv_heads: usize,
+        total_blocks: usize,
+        cpu_total_blocks: usize,
+        max_seqs: usize,
+        max_blocks_per_seq: usize,
+        dtype: DType,
+        device: u32,
     ) -> Self {
         assert!(total_blocks >= 2, "need at least 2 blocks (one is sentinel)");
         let num_layers = config.num_layers();
-        let num_kv_heads = config.num_kv_heads();
         let head_dim = config.head_dim();
         let elem_size = dtype.size_bytes();
         let block_bytes = num_kv_heads * BLOCK_SIZE * head_dim * elem_size;

crates/xserv-model/src/qwen3.rs

@@ -15,6 +15,11 @@ pub struct Qwen3 {
     norm: Tensor,
     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.
+    tp: Option<std::sync::Arc<xserv_distributed::TpContext>>,
+    local_num_heads: usize,    // = num_heads / world
+    local_num_kv_heads: usize, // = num_kv_heads / world
 }
 
 struct Qwen3Block {
@@ -32,15 +37,43 @@ struct Qwen3Block {
 }
 
 impl Qwen3 {
-    pub fn from_weights(config: ModelConfig, mut w: HashMap<String, Tensor>) -> Self {
+    /// Single-GPU load (weights already on the target GPU). Equivalent to
+    /// `from_weights_tp(.., rank=0, world=1, device=0, tp=None)`.
+    pub fn from_weights(config: ModelConfig, w: HashMap<String, Tensor>) -> Self {
+        Self::from_weights_tp(config, w, 0, 1, 0, None)
+    }
+
+    /// Tensor-parallel load. `w` may live on CPU or any device; each weight is
+    /// sharded for `rank`/`world`, uploaded to `device`, and transposed.
+    /// `world==1` shards are identity, so this is also the single-GPU path.
+    ///
+    /// Split scheme (Megatron-style):
+    /// - column-parallel (split output): q/k/v/gate/up  → shard rows of `[out,in]`
+    /// - row-parallel    (split input):  o/down         → shard cols of `[out,in]`
+    /// - replicated: norms, embed_tokens, lm_head
+    pub fn from_weights_tp(
+        config: ModelConfig,
+        mut w: HashMap<String, Tensor>,
+        rank: usize,
+        world: usize,
+        device: u32,
+        tp: Option<std::sync::Arc<xserv_distributed::TpContext>>,
+    ) -> Self {
         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}"))
         };
+        // 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() };
+        // 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 embed_tokens = take(&mut w, "model.embed_tokens.weight");
-        let norm = take(&mut w, "model.norm.weight");
-        let lm_head_raw = take(&mut w, "lm_head.weight");
+        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 rope_cache = RopeCache::new(
             config.max_seq_len(),
@@ -48,33 +81,51 @@ impl Qwen3 {
             config.rope_theta.unwrap_or(1_000_000.0) as f32,
         );
 
-        // Precompute transposed weights: [out, in] → [in, out] so we can do x @ wt directly
-        let transpose_w = |t: Tensor| -> Tensor {
-            t.transpose(0, 1).contiguous()
-        };
-
         let num_layers = config.num_layers();
         let mut layers = Vec::with_capacity(num_layers);
-        eprintln!("Transposing weights for {} layers...", num_layers);
+        if rank == 0 {
+            eprintln!("Loading+sharding weights for {} layers (world={world})...", num_layers);
+        }
         for i in 0..num_layers {
             let p = format!("model.layers.{i}");
             layers.push(Qwen3Block {
-                input_norm: take(&mut w, &format!("{p}.input_layernorm.weight")),
-                q_proj_wt: transpose_w(take(&mut w, &format!("{p}.self_attn.q_proj.weight"))),
-                k_proj_wt: transpose_w(take(&mut w, &format!("{p}.self_attn.k_proj.weight"))),
-                v_proj_wt: transpose_w(take(&mut w, &format!("{p}.self_attn.v_proj.weight"))),
-                o_proj_wt: transpose_w(take(&mut w, &format!("{p}.self_attn.o_proj.weight"))),
-                q_norm: take(&mut w, &format!("{p}.self_attn.q_norm.weight")),
-                k_norm: take(&mut w, &format!("{p}.self_attn.k_norm.weight")),
-                post_norm: take(&mut w, &format!("{p}.post_attention_layernorm.weight")),
-                gate_proj_wt: transpose_w(take(&mut w, &format!("{p}.mlp.gate_proj.weight"))),
-                up_proj_wt: transpose_w(take(&mut w, &format!("{p}.mlp.up_proj.weight"))),
-                down_proj_wt: transpose_w(take(&mut w, &format!("{p}.mlp.down_proj.weight"))),
+                input_norm: repl(take(&mut w, &format!("{p}.input_layernorm.weight"))),
+                q_proj_wt: col(take(&mut w, &format!("{p}.self_attn.q_proj.weight"))),
+                k_proj_wt: col(take(&mut w, &format!("{p}.self_attn.k_proj.weight"))),
+                v_proj_wt: col(take(&mut w, &format!("{p}.self_attn.v_proj.weight"))),
+                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"))),
+                gate_proj_wt: col(take(&mut w, &format!("{p}.mlp.gate_proj.weight"))),
+                up_proj_wt: col(take(&mut w, &format!("{p}.mlp.up_proj.weight"))),
+                down_proj_wt: row(take(&mut w, &format!("{p}.mlp.down_proj.weight"))),
             });
         }
 
-        let lm_head_t = transpose_w(lm_head_raw);
-        Self { config, embed_tokens, layers, norm, lm_head_t, rope_cache }
+        Self {
+            local_num_heads: config.num_heads() / world,
+            local_num_kv_heads: config.num_kv_heads() / world,
+            config,
+            embed_tokens,
+            layers,
+            norm,
+            lm_head_t,
+            rope_cache,
+            tp,
+        }
+    }
+
+    /// In-place AllReduce(sum) of a partial `[*, hidden]` BF16 activation across
+    /// TP ranks (no-op when not tensor-parallel). Used after o_proj and down_proj.
+    #[inline]
+    fn all_reduce(&self, t: &Tensor) {
+        if let Some(tp) = &self.tp {
+            if tp.world > 1 {
+                let ptr = t.storage().gpu_buffer().as_ptr() as *mut std::ffi::c_void;
+                tp.all_reduce_sum_bf16_ptr(ptr, t.numel());
+            }
+        }
     }
 
     pub fn forward_with_cache(&self, token_ids: &[u32], cache: &mut KVCache) -> Tensor {
@@ -273,8 +324,9 @@ impl Qwen3 {
         assert_eq!(seq_slots.len(), batch);
         assert!(batch > 0);
 
-        let num_heads = self.config.num_heads();
-        let num_kv_heads = self.config.num_kv_heads();
+        // TP: this rank owns a slice of the heads (local_* == full when world==1).
+        let num_heads = self.local_num_heads;
+        let num_kv_heads = self.local_num_kv_heads;
         let head_dim = self.config.head_dim();
         let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
 
@@ -356,6 +408,7 @@ impl Qwen3 {
             // Plain reshape is a view; merge_heads_gpu would incorrectly swap B<->H.
             let attn_merged = attn_out.reshape(&[batch, num_heads * head_dim]);
             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 residual = x_new.clone();
@@ -364,6 +417,7 @@ impl Qwen3 {
             let up = matmul_2d(&normed, &layer.up_proj_wt);
             let hidden_states = xserv_kernels::silu_mul(&gate, &up);
             let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
+            self.all_reduce(&down); // TP: sum partial MLP outputs
             x = add_any(&residual, &down);
         }
 
@@ -387,8 +441,9 @@ impl Qwen3 {
     ) -> Tensor {
         let new_tokens = token_ids.len();
         let pos_offset = paged_cache.seq_len(slot);
-        let num_heads = self.config.num_heads();
-        let num_kv_heads = self.config.num_kv_heads();
+        // TP: this rank owns a slice of the heads (local_* == full when world==1).
+        let num_heads = self.local_num_heads;
+        let num_kv_heads = self.local_num_kv_heads;
         let head_dim = self.config.head_dim();
         let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
 
@@ -431,6 +486,7 @@ impl Qwen3 {
 
             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); // TP: sum partial attention outputs
 
             let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
             let residual = x_new.clone();
@@ -439,6 +495,7 @@ impl Qwen3 {
             let up = matmul_2d(&normed, &layer.up_proj_wt);
             let hidden_states = xserv_kernels::silu_mul(&gate, &up);
             let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
+            self.all_reduce(&down); // TP: sum partial MLP outputs
             x = add_any(&residual, &down);
         }
 
@@ -549,6 +606,43 @@ impl Qwen3 {
 
 // --- Helpers ---
 
+/// Keep this rank's contiguous row-block of a 2D `[rows, cols]` BF16 tensor
+/// (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(); }
+    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}");
+    let local = rows / world;
+    let host = t.to_device(Device::Cpu);
+    let data = host.as_slice::<bf16>();
+    let start = rank * local * cols;
+    let shard = data[start..start + local * cols].to_vec();
+    Tensor::from_slice(&shard, &[local, cols])
+}
+
+/// 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(); }
+    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}");
+    let local = cols / world;
+    let c0 = rank * local;
+    let host = t.to_device(Device::Cpu);
+    let data = host.as_slice::<bf16>();
+    let mut shard = Vec::with_capacity(rows * local);
+    for r in 0..rows {
+        let base = r * cols + c0;
+        shard.extend_from_slice(&data[base..base + local]);
+    }
+    Tensor::from_slice(&shard, &[rows, local])
+}
+
 fn matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
     assert_eq!(a.ndim(), 2);
     assert_eq!(b.ndim(), 2);

crates/xserv-server/Cargo.toml

@@ -13,6 +13,7 @@ xserv-tensor = { path = "../xserv-tensor" }
 xserv-kernels = { path = "../xserv-kernels" }
 xserv-model = { path = "../xserv-model" }
 xserv-tokenizer = { path = "../xserv-tokenizer" }
+xserv-distributed = { path = "../xserv-distributed" }
 half.workspace = true
 serde.workspace = true
 serde_json.workspace = true

crates/xserv-server/src/main.rs

@@ -1,5 +1,6 @@
 mod api;
 mod engine;
+mod tp_engine;
 
 use axum::{routing::{get, post}, Extension, Router};
 use std::path::PathBuf;
@@ -18,7 +19,7 @@ pub struct AppState {
 async fn main() {
     let args: Vec<String> = std::env::args().collect();
     if args.len() < 2 {
-        eprintln!("Usage: xserv-server <model-dir> [--port PORT] [--max-batch N] [--max-seq-len N] [--swap-space-gb N]");
+        eprintln!("Usage: xserv-server <model-dir> [--port PORT] [--max-batch N] [--max-seq-len N] [--swap-space-gb N] [--tp N]");
         std::process::exit(1);
     }
 
@@ -45,6 +46,12 @@ async fn main() {
         .and_then(|i| args.get(i + 1))
         .and_then(|s| s.parse().ok())
         .unwrap_or(8);
+    let tp: usize = args.iter()
+        .position(|a| a == "--tp")
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(1)
+        .max(1);
     let model_config = ModelConfig::from_file(&model_dir.join("config.json"));
     let model_max_seq_len = model_config.max_seq_len();
     if model_max_seq_len == 0 {
@@ -69,8 +76,13 @@ async fn main() {
 
     let model_dir_clone = model_dir.clone();
     std::thread::spawn(move || {
-        let mut engine = engine::Engine::load_with_swap(&model_dir_clone, max_batch, max_seq_len, swap_space_gb);
-        engine.run(rx);
+        if tp <= 1 {
+            let mut engine = engine::Engine::load_with_swap(&model_dir_clone, max_batch, max_seq_len, swap_space_gb);
+            engine.run(rx);
+        } else {
+            // Tensor-parallel path: rank-0 coordinator + worker rank threads.
+            tp_engine::run_tp(&model_dir_clone, tp, max_seq_len, rx);
+        }
     });
 
     let state = Arc::new(AppState {

crates/xserv-server/src/tp_engine.rs

@@ -0,0 +1,195 @@
+//! Tensor-parallel inference engine for the HTTP server.
+//!
+//! Serial coordinator model: one rank-0 coordinator thread (the caller) drives
+//! generation and owns the scheduler; ranks 1..world are worker threads. For
+//! each step the coordinator broadcasts a command (Register/Prefill/Decode/Free)
+//! to the workers and runs the same op on its own shard; the per-layer NCCL
+//! AllReduces keep all ranks in lockstep. Only the coordinator samples — the
+//! chosen token is carried in the next Decode command, so this is correct for
+//! both greedy and stochastic sampling.
+//!
+//! Requests are processed one at a time (sufficient for the quality benchmark,
+//! which issues serial requests). Continuous batching across ranks is future
+//! work; the single-GPU `Engine` still handles TP=1.
+
+use std::path::{Path, PathBuf};
+use std::sync::mpsc;
+use std::sync::Arc;
+use std::thread;
+
+use xserv_distributed::{TpContext, UniqueId};
+use xserv_model::loader;
+use xserv_model::{sample, ModelConfig, PagedKVCache, Qwen3, BLOCK_SIZE};
+use xserv_tensor::{DType, Device};
+use xserv_tokenizer::Tokenizer;
+
+use crate::engine::{GenerateEvent, GenerateRequest};
+
+#[derive(Clone)]
+enum TpCommand {
+    Register(usize),
+    Free(usize),
+    Prefill { tokens: Vec<u32>, slot: usize },
+    Decode { tokens: Vec<u32>, positions: Vec<usize>, slots: Vec<usize> },
+    Shutdown,
+}
+
+struct RankCtx {
+    model: Qwen3,
+    cache: PagedKVCache,
+}
+
+fn build_rank(
+    model_dir: &Path,
+    config: &ModelConfig,
+    rank: usize,
+    world: usize,
+    device: u32,
+    max_seq_len: usize,
+    tp: Option<Arc<TpContext>>,
+) -> RankCtx {
+    let weights = loader::load_model_dir(model_dir, Device::Cpu);
+    let model = Qwen3::from_weights_tp(config.clone(), weights, rank, world, device, tp);
+    let local_kv = config.num_kv_heads() / world;
+    let max_blocks_per_seq = max_seq_len.div_ceil(BLOCK_SIZE);
+    let total_blocks = max_blocks_per_seq + 8;
+    let cache = PagedKVCache::new_tp(
+        config, local_kv, total_blocks, 0, 4, max_blocks_per_seq, DType::BF16, device,
+    );
+    RankCtx { model, cache }
+}
+
+fn worker_loop(
+    rank: usize,
+    world: usize,
+    id: UniqueId,
+    model_dir: PathBuf,
+    config: ModelConfig,
+    max_seq_len: usize,
+    cmd_rx: mpsc::Receiver<TpCommand>,
+    ack_tx: mpsc::Sender<()>,
+) {
+    let tp = Arc::new(TpContext::init(rank, world, id, rank as u32));
+    let mut rc = build_rank(&model_dir, &config, rank, world, rank as u32, max_seq_len, Some(tp));
+    while let Ok(cmd) = cmd_rx.recv() {
+        match cmd {
+            TpCommand::Register(slot) => {
+                let _ = rc.cache.register_sequence(slot);
+            }
+            TpCommand::Free(slot) => rc.cache.free_sequence(slot),
+            TpCommand::Prefill { tokens, slot } => {
+                let _ = rc.model.forward_prefill_paged(&tokens, slot, &mut rc.cache);
+            }
+            TpCommand::Decode { tokens, positions, slots } => {
+                let _ = rc.model.forward_decode_paged(&tokens, &positions, &slots, &mut rc.cache);
+            }
+            TpCommand::Shutdown => {
+                let _ = ack_tx.send(());
+                break;
+            }
+        }
+        let _ = ack_tx.send(());
+    }
+}
+
+/// Run the TP coordinator (rank 0) on the calling thread. Spawns worker ranks
+/// internally and consumes generation requests from `rx`.
+pub fn run_tp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Receiver<GenerateRequest>) {
+    assert!(world >= 2, "run_tp requires world >= 2");
+    let config = ModelConfig::from_file(&model_dir.join("config.json"));
+    assert!(
+        config.num_kv_heads() % world == 0,
+        "num_kv_heads {} not divisible by tp {world}",
+        config.num_kv_heads()
+    );
+    let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
+    let id = xserv_distributed::get_unique_id();
+
+    // Spawn worker ranks 1..world.
+    let (ack_tx, ack_rx) = mpsc::channel::<()>();
+    let mut cmd_txs: Vec<mpsc::Sender<TpCommand>> = Vec::new();
+    for rank in 1..world {
+        let (ctx_tx, ctx_rx) = mpsc::channel::<TpCommand>();
+        cmd_txs.push(ctx_tx);
+        let ack_tx = ack_tx.clone();
+        let model_dir = model_dir.to_path_buf();
+        let config = config.clone();
+        thread::spawn(move || {
+            worker_loop(rank, world, id, model_dir, config, max_seq_len, ctx_rx, ack_tx);
+        });
+    }
+
+    // Rank 0 (this thread).
+    let tp = Arc::new(TpContext::init(0, world, id, 0));
+    let mut rc = build_rank(model_dir, &config, 0, world, 0, max_seq_len, Some(tp));
+    eprintln!("[tp-engine] ready (tp={world}, max_seq_len={max_seq_len})");
+
+    let eos = tokenizer.eos_token_id();
+    let n_workers = world - 1;
+    let broadcast = |txs: &[mpsc::Sender<TpCommand>], cmd: TpCommand| {
+        for t in txs {
+            let _ = t.send(cmd.clone());
+        }
+    };
+    let wait_acks = |rx: &mpsc::Receiver<()>| {
+        for _ in 0..n_workers {
+            let _ = rx.recv();
+        }
+    };
+
+    let slot = 0usize;
+    while let Ok(req) = rx.recv() {
+        broadcast(&cmd_txs, TpCommand::Register(slot));
+        rc.cache.register_sequence(slot).expect("register slot");
+        wait_acks(&ack_rx);
+
+        // Prefill.
+        broadcast(&cmd_txs, TpCommand::Prefill { tokens: req.prompt_tokens.clone(), slot });
+        let logits = rc.model.forward_prefill_paged(&req.prompt_tokens, slot, &mut rc.cache);
+        wait_acks(&ack_rx);
+        let mut next = sample(&logits, &req.sampling);
+
+        let mut decode_buf: Vec<u8> = Vec::new();
+        let mut generated = 1usize;
+        emit_text(&tokenizer, &req, next, eos, &mut decode_buf);
+
+        let finish = loop {
+            if eos == Some(next) {
+                break "stop";
+            }
+            if generated >= req.max_tokens {
+                break "length";
+            }
+            let pos = rc.cache.seq_len(slot);
+            broadcast(&cmd_txs, TpCommand::Decode { tokens: vec![next], positions: vec![pos], slots: vec![slot] });
+            let logits = rc.model.forward_decode_paged(&[next], &[pos], &[slot], &mut rc.cache);
+            wait_acks(&ack_rx);
+            next = sample(&logits, &req.sampling);
+            generated += 1;
+            emit_text(&tokenizer, &req, next, eos, &mut decode_buf);
+        };
+
+        let tail = tokenizer.flush_decode_stream(&mut decode_buf);
+        if !tail.is_empty() {
+            let _ = req.sender.blocking_send(GenerateEvent::Token { id: next, text: tail });
+        }
+        let _ = req.sender.blocking_send(GenerateEvent::Done { finish_reason: finish.to_string() });
+
+        broadcast(&cmd_txs, TpCommand::Free(slot));
+        rc.cache.free_sequence(slot);
+        wait_acks(&ack_rx);
+    }
+
+    broadcast(&cmd_txs, TpCommand::Shutdown);
+}
+
+/// Stream a token's decoded text to the client (EOS contributes no text).
+fn emit_text(tokenizer: &Tokenizer, req: &GenerateRequest, token_id: u32, eos: Option<u32>, buf: &mut Vec<u8>) {
+    if eos == Some(token_id) {
+        return;
+    }
+    let text = tokenizer.decode_token_stream(token_id, buf);
+    if !text.is_empty() {
+        let _ = req.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
+    }
+}

docs/17-tensor-parallelism.md

@@ -0,0 +1,122 @@
+# Phase 17: Tensor Parallelism (TP)
+
+> 目标：在单机多卡上做 **张量并行**，把 Qwen3-8B 的权重、计算和 KV cache 按
+> head / 中间维切分到 TP 个 GPU 上，用 AllReduce 聚合，降低单卡显存压力并提升吞吐。
+> 先做 **TP=2 / 4（组内）**，跳过投机解码（原 Phase 16）。
+
+## 1. 硬件约束（dash5）
+
+- 8× RTX 5090（32GB，SM120），**无 NVLink**，纯 PCIe Gen5。
+- 拓扑：GPU 0–3 一组、4–7 一组，组内 `PHB`（同 host bridge，可 P2P），跨组 `NODE`。
+- **TP 必须在组内**（0–3 或 4–7），否则 AllReduce 走跨组 PCIe，延迟更高。
+- AllReduce 带宽受限于 PCIe（~单向 64GB/s），远低于 NVLink；通信会是 decode 的主要开销。
+
+## 2. 切分方案（Megatron-style）
+
+Qwen3-8B：`hidden=4096`、`num_heads=32`、`num_kv_heads=8`、`head_dim=128`、
+`intermediate=12288`、`layers=36`、`vocab=151936`。TP=2/4/8 都能整除 32/8/12288。
+
+每个 transformer block 的切分（设 world size = `T`，本 rank = `r`）：
+
+### Attention（column → row）
+| 权重 | 原 shape (已转置) | 切分 | 每 rank shape |
+|------|-------------------|------|---------------|
+| `q_proj_wt` | `[hidden, num_heads·head_dim]` | column（按 Q head） | `[hidden, (num_heads/T)·head_dim]` |
+| `k_proj_wt` | `[hidden, num_kv_heads·head_dim]` | column（按 KV head） | `[hidden, (num_kv_heads/T)·head_dim]` |
+| `v_proj_wt` | 同上 | column | 同上 |
+| `o_proj_wt` | `[num_heads·head_dim, hidden]` | **row** | `[(num_heads/T)·head_dim, hidden]` |
+
+- 每个 rank 只算自己的 `num_heads/T` 个 Q head 和对应的 `num_kv_heads/T` 个 KV head；
+  GQA 的 `n_rep = num_heads/num_kv_heads = 4` 在每个 rank 内保持不变。
+- `q_norm`/`k_norm`（`[head_dim]`）逐 head 应用，**复制**到每个 rank。
+- RoPE 逐 head、按 position 应用，每个 rank 独立做。
+- **KV cache 也切分**：每个 rank 的 paged KV 只存自己的 `num_kv_heads/T` 个 head
+  → 每卡 KV 显存降为 1/T（TP 的一大收益）。
+- `attn = merge_heads(...) @ o_proj_wt` 得到**部分** `[T_tok, hidden]` → **AllReduce(sum)** → 完整。
+
+### MLP / SwiGLU（column → row）
+| 权重 | 原 shape | 切分 | 每 rank shape |
+|------|----------|------|---------------|
+| `gate_proj_wt` | `[hidden, intermediate]` | column | `[hidden, intermediate/T]` |
+| `up_proj_wt` | 同上 | column | 同上 |
+| `down_proj_wt` | `[intermediate, hidden]` | **row** | `[intermediate/T, hidden]` |
+
+- `silu(gate)*up` 在切分后的 `[T_tok, intermediate/T]` 上逐元素做，无需通信。
+- `down = (...) @ down_proj_wt` 得到部分 `[T_tok, hidden]` → **AllReduce(sum)** → 完整。
+
+### 复制（不切分）
+- 所有 RMSNorm 权重（`input_norm`/`post_norm`/最终 `norm`）：每个 rank 在 AllReduce 后
+  拿到完整 hidden，本地用复制的权重归一化。
+- **第一版**：`embed_tokens` 和 `lm_head` 复制（各 ~1.2GB）。
+  后续优化：vocab-parallel embedding（local lookup + AllReduce）、column-parallel lm_head + AllGather。
+
+### 通信点
+每层 **2 次 AllReduce**（o_proj 后、down_proj 后）→ 每生成 1 token 共 `2·36 = 72` 次。
+decode 时每次 AllReduce 张量是 `[batch, 4096]` BF16（单 token batch=1 时 8KB），**延迟主导**。
+
+## 3. 进程 / 线程模型
+
+**单进程、多线程**：每个 TP rank 一个 OS 线程，线程启动时 `cudaSetDevice(rank_device)` 并绑定。
+
+选择理由：
+- xserv 的 caching allocator 是 `thread_local`，每线程独立池 → 天然契合「一线程一卡一池」。
+- CUDA context 隐式按 device/thread 管理；线程内只 set 一次 device、不再切换即可。
+- HTTP server / 调度器仍在主线程（rank 0 协调），无需多进程 IPC，改动最小。
+- 单机 8 卡足够；多进程（torchrun 式）留待真正跨节点时再说。
+
+执行流：主线程调度器准备一个 step 的输入（tokens/positions/slots），广播给 `T` 个 rank 线程；
+每个 rank 线程跑自己的分片 forward（含层内 AllReduce），rank 0 拿到完整 logits 后采样。
+用 barrier / channel 同步每个 step。
+
+## 4. 通信库：NCCL
+
+用 **NCCL**（dash5 已装：`/usr/lib/x86_64-linux-gnu/libnccl.so.2`，`/usr/include/nccl.h`）。
+
+- 新建 crate `xserv-distributed`：NCCL FFI（`ncclGetUniqueId`、`ncclCommInitRank`、
+  `ncclAllReduce`、`ncclGroupStart/End`）+ `TpContext`（rank/world/comm/stream）+
+  `all_reduce_sum(&mut GpuBuffer)` 原语。
+- NCCL 多线程模式：主线程生成 `ncclUniqueId`，各 rank 线程用 `ncclCommInitRank(comm, world, id, rank)`
+  初始化（需 `ncclGroupStart/End` 包裹并发 init）。
+- AllReduce 用 BF16（`ncclBfloat16`）+ `ncclSum`，在每个 rank 自己的 stream 上。
+
+> **决策点**：collective 用 NCCL 还是自己手写 P2P ring/tree AllReduce？
+> 本项目是「从零构建」，但 collective 属于基础设施（类比我们也用 cuBLAS 作为可切换后端）。
+> 推荐先用 NCCL 把 TP 跑通、拿到正确性与加速比，后续可选做手写 AllReduce 作为学习项。
+
+## 5. 权重分片加载
+
+每个 rank 只加载/保留自己的分片，省显存：
+- column-parallel 权重：按输出维切片取本 rank 的 `[*, dim/T]` 段。
+- row-parallel 权重：按输入维切片取本 rank 的 `[dim/T, *]` 段。
+- 复制权重（norm/embed/lm_head）：每个 rank 各留一份。
+
+实现：`loader` 读 safetensors（mmap）时按 rank 只搬运需要的切片到该 rank 的 GPU；
+`Qwen3::from_weights_tp(config, weights, rank, world)` 在转置/切分时按 rank 取段。
+
+## 6. 实现步骤（逐步可验证）
+
+1. **P17.1 — `xserv-distributed` 基础**：NCCL FFI + `TpContext` + `all_reduce_sum`。
+   验收：2 卡各放一个已知向量，AllReduce 后两卡结果都等于和。
+2. **P17.2 — 分片权重加载**：`from_weights_tp`，每 rank 只持有自己的分片。
+   验收：各 rank 权重 shape 正确、显存占用约为 1/T（+ 复制项）。
+3. **P17.3 — TP forward**：rank 内 attention/MLP + 层内 AllReduce。
+   验收：**TP=2 的 logits 与 TP=1 在 BF16 容差内一致**（top-1 一致，top-5 重合）。
+4. **P17.4 — 接入 engine/server**：`--tp N`，多线程 rank workers + rank0 调度。
+   验收：`--tp 2` 端到端可服务；用现有 llama.cpp bench 跑正确性；
+   测 TP=2 vs TP=1 的吞吐 / TTFT / 每卡显存。
+
+## 7. 预期与风险
+
+- **显存**：每卡权重 + KV 降到约 1/T（embed/lm_head 暂复制）。TP=2 时单卡 ~8GB 权重 + 更大 KV/并发空间。
+- **吞吐**：PCIe AllReduce 延迟会吃掉部分收益；decode 是延迟敏感的，72 次小 AllReduce/token
+  可能让 TP=2 的单流 tok/s **不一定线性提升**，但能跑更大 batch / 更长 context。先测实测数。
+- **风险**：NCCL 多线程初始化的同步、每 rank stream 与现有 kernel stream 的协调、
+  KV cache 按 rank 切 head 后 paged kernel 的 head 维参数要用 per-rank 值。
+- 正确性优先：先 TP=1 等价（logits 对齐），再谈性能。
+
+## 8. 不在本阶段范围
+
+- 跨组 TP=8、Pipeline Parallelism、多节点。
+- vocab-parallel embedding / lm_head（先复制）。
+- 手写 AllReduce（NCCL 跑通后可选）。
+- 与 CUDA Graph decode 的结合（先走非 graph 路径）。

docs/benchmarks/tensor-parallelism.md

@@ -0,0 +1,73 @@
+# Benchmark: Tensor Parallelism (TP=1/2/4) — xserv vs llama.cpp
+
+**Setup.** Qwen3-8B BF16 on 8× RTX 5090 (PCIe Gen5, **no NVLink**; GPUs grouped
+0-3 / 4-7 by PHB). Both engines driven over the same OpenAI HTTP harness, same
+scorers, thinking-off, greedy (temp 0), `max_tokens` 2048. Datasets: **AIME
+2025** (30) + **GSM8K** (30). The two engines run **concurrently on disjoint
+groups** — xserv on GPU 0..N-1, llama.cpp (`--split-mode row`) on GPU 4..4+N-1
+(`tools/bench/run_tp_parallel.sh`).
+
+## Correctness — on par across engines and TP
+
+| TP | task | xserv | llama.cpp |
+|----|------|-------|-----------|
+| 1 | AIME 2025 | 16.7% (5/30) | 13.3% (4/30) |
+| 1 | GSM8K | 96.7% (29/30) | 96.7% (29/30) |
+| 2 | AIME 2025 | 13.3% (4/30) | 13.3% (4/30) |
+| 2 | GSM8K | 93.3% (28/30) | 96.7% (29/30) |
+| 4 | AIME 2025 | 16.7% (5/30) | 13.3% (4/30) |
+| 4 | GSM8K | 96.7% (29/30) | 96.7% (29/30) |
+
+Within ±1 problem everywhere — TP changes nothing about quality on either
+engine, and the two engines agree. (AIME is low for both: Qwen3-8B thinking-off,
+capped at 2048 tokens.)
+
+## Performance — TPOT (ms/token, lower is better)
+
+| TP | xserv AIME / GSM8K | llama.cpp AIME / GSM8K |
+|----|--------------------|------------------------|
+| 1 | 21.0 / 17.8 | **10.4 / 10.3** |
+| 2 | 17.2 / 13.9 | 19.0 / 18.9 |
+| 4 | **15.2 / 12.1** | 20.2 / 20.2 |
+
+**Opposite TP scaling, with a crossover:**
+
+- **xserv TP scales positively**: TPOT 21.0 → 17.2 → 15.2 ms (AIME),
+  17.8 → 13.9 → 12.1 ms (GSM8K) — TP=4 is ~1.4–1.5× faster than TP=1. GPU 0-3
+  all ~82% utilized. (Sublinear because of the 72 PCIe AllReduces/token.)
+- **llama.cpp row-split regresses**: TPOT 10.4 → 19.0 → 20.2 ms — TP=1 is its
+  best; TP=2/4 nearly double the latency. GPU 4-7 only ~24% utilized
+  (communication-bound). Row-split's per-layer cross-GPU traffic over PCIe
+  without NVLink dominates.
+- **Crossover**: llama.cpp is ~2× faster at TP=1, still ahead at TP=2, and xserv
+  is ~1.3× faster at TP=4 (15.2 vs 20.2 ms AIME). AIME-30 wall clock: xserv
+  1046 → 846 → 730 s (falling), llama.cpp 520 → 952 → 1012 s (rising).
+
+On a NVLink-less PCIe box, **xserv's TP is a genuine win and llama.cpp's
+tensor-split is counterproductive** — exactly what the topology predicts.
+
+### Clean same-path xserv scaling (bench-tp)
+
+The HTTP numbers above mix engines (xserv TP=1 uses the production continuous-
+batching engine; TP≥2 uses the serial TP coordinator). The single-stream,
+same-code-path scaling from `bench-tp` (greedy, 8 prompts × 64 tokens):
+
+| TP | xserv decode tok/s | speedup | TTFT |
+|----|--------------------|---------|------|
+| 1 | 58.5 | 1.00× | 18.0 ms |
+| 2 | 75.7 | 1.29× | 13.4 ms |
+| 4 | 86.1 | 1.47× | 11.5 ms |
+
+## Caveats
+
+- xserv **TP=1 uses the production `Engine`**, TP≥2 the serial `tp_engine`
+  coordinator — different per-token paths, so the HTTP TP=1→2 step has an engine
+  confound. The clean same-path scaling (bench-tp, above) confirms the trend.
+- xserv **TTFT is weaker** on long AIME prompts (~460–500 ms vs llama ~100–190 ms)
+  — prefill is a known optimization target.
+- llama.cpp uses `--split-mode row` (its tensor-parallel mode); the default
+  `layer` split only memory-splits, without parallel compute.
+- The TP HTTP server processes requests **serially** (sufficient for this serial
+  quality benchmark); continuous-batching TP is future work.
+
+Raw artifacts: `bench-out/tp{1,2,4}-{xserv,llama}/comparison-*.{md,json}`.

tools/bench/run_tp_parallel.sh

@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+# Run the TP=1/2/4 quality sweep with xserv and llama.cpp CONCURRENTLY on
+# disjoint GPU groups: xserv on GPUs 0..N-1, llama.cpp on GPUs 4..4+N-1.
+# The 8x5090 box is grouped 0-3 / 4-7 (PHB intra-group), so each engine's TP
+# comm stays intra-group and the two engines never contend for a GPU.
+#
+# Run from the repo root on the GPU host. Produces bench-out/tp{1,2,4}-{xserv,llama}.
+
+set -u
+MODEL="${MODEL:-/opt/wjh/models/qwen3-8b}"
+GGUF="${GGUF:-/opt/wjh/models/qwen3-8b/qwen3-8b-bf16.gguf}"
+LIMIT="${LIMIT:-30}"
+MAXSEQ="${MAXSEQ:-2048}"
+TPS="${TPS:-1 2 4}"
+
+for TP in $TPS; do
+    LD=$(seq -s, 4 $((3 + TP)))   # llama GPUs: 4 / 4,5 / 4,5,6,7
+    echo "##### TP=$TP  (xserv GPU 0..$((TP-1))  ||  llama GPU $LD) #####"
+    rm -rf "bench-out/tp$TP-xserv" "bench-out/tp$TP-llama"
+
+    python3 -u -m tools.bench.runner --systems xserv --tp "$TP" \
+        --xserv-bin ./target/release/xserv-server --xserv-model "$MODEL" \
+        --suite quality --quality-tasks aime2025,gsm8k --quality-limit "$LIMIT" \
+        --max-batch 1 --max-seq-len "$MAXSEQ" \
+        --out-dir "bench-out/tp$TP-xserv" > "/tmp/tp$TP-xserv.log" 2>&1 &
+    XP=$!
+
+    python3 -u -m tools.bench.runner --systems llama.cpp --tp "$TP" --llama-devices "$LD" \
+        --llama-bin third_party/llama.cpp/build/bin/llama-server --llama-gguf "$GGUF" \
+        --suite quality --quality-tasks aime2025,gsm8k --quality-limit "$LIMIT" \
+        --max-batch 1 --max-seq-len "$MAXSEQ" \
+        --out-dir "bench-out/tp$TP-llama" > "/tmp/tp$TP-llama.log" 2>&1 &
+    LP=$!
+
+    wait "$XP" "$LP"
+    echo "TP=$TP done (xserv exit=$? )"
+done
+echo ALL_DONE

tools/bench/runner.py

@@ -69,6 +69,13 @@ def parse_args() -> argparse.Namespace:
     p.add_argument("--max-seq-len", type=int, default=8192)
     p.add_argument("--systems", default="xserv,llama.cpp",
                    help="Comma-separated subset to run, e.g. 'xserv' to skip llama.cpp")
+    p.add_argument("--tp", type=int, default=1,
+                   help="Tensor-parallel degree for BOTH engines (xserv --tp N; "
+                        "llama.cpp --split-mode row over the first N GPUs).")
+    p.add_argument("--llama-devices", default=None,
+                   help="Comma list of GPU ordinals for llama.cpp (first --tp used). "
+                        "Lets llama run on a disjoint GPU group (e.g. 4,5,6,7) so it "
+                        "can run concurrently with xserv on 0..N-1.")
     p.add_argument("--enable-thinking", action="store_true",
                    help="Enable Qwen3 thinking on llama.cpp. Default OFF to match "
                         "xserv, which hardcodes thinking off in its prompt builder.")
@@ -106,10 +113,10 @@ def build_endpoints(args) -> list[SystemEndpoint]:
                 model_id=args.xserv_model_id,
                 launch_cmd=xserv_launch_cmd(
                     args.xserv_bin, model_dir, args.xserv_port,
-                    max_batch=args.max_batch, max_seq_len=args.max_seq_len,
+                    max_batch=args.max_batch, max_seq_len=args.max_seq_len, tp=args.tp,
                 ),
                 health_path="/health",
-                ready_timeout_s=900.0,
+                ready_timeout_s=1200.0,
             ))
 
     # Match xserv's hardcoded thinking-OFF mode unless explicitly overridden.
@@ -128,17 +135,29 @@ def build_endpoints(args) -> list[SystemEndpoint]:
             gguf = args.llama_gguf or os.environ.get("LLAMA_GGUF")
             if not gguf:
                 raise SystemExit("--llama-gguf or LLAMA_GGUF required (or pass --llama-base-url)")
+            # Pick the GPUs llama.cpp runs on. Default is the first `tp` GPUs;
+            # pass --llama-devices to place it on a disjoint group (e.g. 4,5,6,7)
+            # so it can run concurrently with xserv on 0..N-1. --split-mode row
+            # then tensor-parallel-splits across exactly these devices.
+            if args.llama_devices:
+                devs = [d.strip() for d in args.llama_devices.split(",") if d.strip()][: max(args.tp, 1)]
+                llama_env = {"CUDA_VISIBLE_DEVICES": ",".join(devs)}
+            elif args.tp > 1:
+                llama_env = {"CUDA_VISIBLE_DEVICES": ",".join(str(d) for d in range(args.tp))}
+            else:
+                llama_env = {}
             eps.append(SystemEndpoint(
                 name=SYSTEM_LLAMA_CPP,
                 base_url=f"http://127.0.0.1:{args.llama_port}",
                 model_id=args.llama_model_id,
                 launch_cmd=llama_cpp_launch_cmd(
                     args.llama_bin, gguf, args.llama_port,
-                    n_parallel=args.max_batch, ctx_per_slot=args.max_seq_len,
+                    n_parallel=args.max_batch, ctx_per_slot=args.max_seq_len, tp=args.tp,
                 ),
+                launch_env=llama_env,
                 # llama-server's health endpoint also returns 200 only when model is loaded.
                 health_path="/health",
-                ready_timeout_s=900.0,
+                ready_timeout_s=1200.0,
                 extra_body=llama_extra_body,
             ))
     return eps

tools/bench/servers.py

@@ -113,14 +113,18 @@ def xserv_launch_cmd(
     *,
     max_batch: int,
     max_seq_len: int,
+    tp: int = 1,
 ) -> list[str]:
-    return [
+    cmd = [
         bin_path,
         model_dir,
         "--port", str(port),
         "--max-batch", str(max_batch),
         "--max-seq-len", str(max_seq_len),
     ]
+    if tp > 1:
+        cmd += ["--tp", str(tp)]  # xserv binds rank r -> GPU r internally
+    return cmd
 
 
 def llama_cpp_launch_cmd(
@@ -131,13 +135,14 @@ def llama_cpp_launch_cmd(
     n_parallel: int,
     ctx_per_slot: int,
     n_gpu_layers: int = 99,
+    tp: int = 1,
 ) -> list[str]:
     # llama.cpp DIVIDES total -c across --parallel slots: per-slot context is
     # n_ctx / n_parallel. xserv gives each sequence the full max_seq_len, so to
     # match we must set total -c = ctx_per_slot * n_parallel. Getting this wrong
     # silently truncates long generations (e.g. AIME) on llama.cpp's side.
     total_ctx = ctx_per_slot * n_parallel
-    return [
+    cmd = [
         bin_path,
         "-m", gguf_path,
         "--port", str(port),
@@ -148,3 +153,9 @@ def llama_cpp_launch_cmd(
         # NOTE: do NOT pass --log-disable; its startup log reports per-slot
         # n_ctx, which is exactly the diagnostic that catches ctx misconfig.
     ]
+    if tp > 1:
+        # Tensor-parallel split across the visible GPUs (caller restricts the
+        # set via CUDA_VISIBLE_DEVICES in launch_env). Row-split is llama.cpp's
+        # tensor-parallel mode (vs the default layer/pipeline split).
+        cmd += ["--split-mode", "row"]
+    return cmd

tools/bench/summarize_tp.py

@@ -0,0 +1,24 @@
+"""Summarize the concurrent TP sweep: bench-out/tp{1,2,4}-{xserv,llama}."""
+import glob
+import json
+import os
+import sys
+
+base = sys.argv[1] if len(sys.argv) > 1 else "bench-out"
+rows = []
+for tp in (1, 2, 4):
+    for sysname in ("xserv", "llama"):
+        files = sorted(glob.glob(os.path.join(base, f"tp{tp}-{sysname}", "comparison-*.json")))
+        if not files:
+            continue
+        d = json.load(open(files[-1]))
+        for r in d["quality"]["summary"]:
+            rows.append((tp, sysname, r["task"], r["n_correct"], r["n_total"],
+                         r["accuracy"] * 100, r["mean_completion_tokens"],
+                         r["mean_ttft_ms"], r["mean_tpot_ms"], r["wall_s"]))
+
+print("%-3s %-7s %-9s %-9s %7s %9s %9s %10s %9s" %
+      ("TP", "engine", "task", "correct", "acc%", "mean_tok", "TTFT_ms", "TPOT_ms", "wall_s"))
+for (tp, s, task, nc, nt, acc, tok, ttft, tpot, wall) in rows:
+    print("%-3d %-7s %-9s %-9s %6.1f%% %9.0f %9.1f %10.2f %9.0f" %
+          (tp, s, task, f"{nc}/{nt}", acc, tok, ttft, tpot, wall))