diff --git a/.gitignore b/.gitignore index 9866984..94b93e7 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,9 @@ /third_party/llama.cpp/models/ *.gguf +# Claude Code runtime state +/.claude/ + # Benchmark output + fetched datasets (transferred to GPU host, not committed) /bench-out/ /tools/bench/data/ diff --git a/README.md b/README.md index 9666e69..703e17a 100644 --- a/README.md +++ b/README.md @@ -144,16 +144,53 @@ HF_ENDPOINT=https://hf-mirror.com python3 -m tools.bench.fetch_datasets - `docs/00-roadmap.md`:总体路线图与各 Phase 设计 - `docs/01..15-*.md`:CUDA FFI / Tensor / GEMM / Attention / KV cache / 性能优化等每个 Phase 的设计文档 - `docs/16-llama-cpp-comparison.md`:llama.cpp 对比基准的设计 -- `docs/benchmarks/`:各阶段的 benchmark 报告 +- `docs/17-tensor-parallelism.md`:张量并行(TP)设计 +- `docs/18-pipeline-parallelism.md`:流水线并行(PP)设计 +- `docs/benchmarks/`:各阶段的 benchmark 报告(含 `pp-sweep.md`) + +## 多卡并行(TP / PP) + +单机多卡,复用 NCCL(crate `xserv-distributed`)。两种切法正交、二选一: + +- **张量并行 `--tp N`**:按 head / 中间维切每一层,层内用 AllReduce 聚合(每 token `2·层数` 次)。 +- **流水线并行 `--pp N`**:按层切成 N 段,相邻段间用 NCCL **P2P** 传 hidden state(每 token 仅 `N-1` 次), + 通信量远小于 AllReduce,对无 NVLink 的 PCIe 更友好。 + +```bash +# 组内 GPU 0-3:4 卡张量并行 / 4 卡流水线并行 +CUDA_VISIBLE_DEVICES=0,1,2,3 ./target/release/xserv-server /path/to/qwen3-8b --tp 4 +CUDA_VISIBLE_DEVICES=0,1,2,3 ./target/release/xserv-server /path/to/qwen3-8b --pp 4 +``` + +**PP 实测**(dash5,Qwen3-8B BF16,单流贪心;每卡显存为权重+最小 KV 池): + +| 配置 | TTFT | TPOT | tok/s | 每卡显存 | +|------|------|------|-------|----------| +| 单卡 | 33ms | 17.4ms | 57.5 | 24.0 GB | +| PP=2 | 36ms | 18.1ms | 55.3 | 11.6 / 13.6 GB | +| PP=4 | 36ms | 17.9ms | 55.8 | 7.3 / 5.3 / 5.3 / 9.4 GB | + +**质量对比**(AIME 2025 30 题 + GSM8K 30 题,贪心,xserv 在 GPU 0-3、llama.cpp 在 GPU 4-7 并行): + +| 引擎 | PP | AIME | GSM8K | +|------|----|------|-------| +| xserv | 1/2/4 | 8 / 7 / 7 (/30) | 29/30 (96.7%) 全部一致 | +| llama | 1/2/4 | 7 / 7 / 7 (/30) | 29/30 (96.7%) 全部一致 | + +正确性:hidden state 跨段是 **bit-exact BF16 P2P 拷贝**,PP=4 输出与单卡逐字节一致(用「单卡×2 vs +PP=4×2」对照确认——单卡自身因 cuBLAS 非确定性 run-to-run 会变,而 PP=4 可复现且落在某次单卡轨迹上)。 +GSM8K 12 个格子全是 29/30,xserv 与 llama.cpp 完全一致;AIME 的 ±1 是长生成下贪心对 GEMM 抖动的敏感, +非 PP 或引擎效应。**收益在显存**(每卡权重+KV ≈ 1/N);v1 为串行流水线,单流 TPOT 基本持平、不优于单卡, +真正的吞吐提升需后续做 microbatch / 1F1B 重叠。完整数据见 `docs/benchmarks/pp-sweep.md`。 ## 路线图(节选) -已完成 Phase 0–15:CUDA 基础设施 → Tensor → GEMM → Transformer kernels → Attention → +已完成 Phase 0–18:CUDA 基础设施 → Tensor → GEMM → Transformer kernels → Attention → 模型加载 → 分词器 → GPT-2 → KV cache → Qwen3-8B → Paged Attention → 连续批处理 → -HTTP API → Flash Attention 2 → 性能优化;并在此基础上加入了 **llama.cpp 对比基准** -与 **KV CPU 换出** 等基础设施。 +HTTP API → Flash Attention 2 → 性能优化 → **张量并行(TP)** → **流水线并行(PP)**; +并加入了 **llama.cpp 对比基准** 与 **KV CPU 换出** 等基础设施。 -后续方向:投机解码(speculative decoding)、张量并行(TP,多卡)、量化(FP8 / INT8)、多模态。 +后续方向:PP microbatch/1F1B 流水线重叠(吞吐收益)、2D TP×PP、投机解码、量化(FP8 / INT8)、多模态。 ## 许可 diff --git a/crates/xserv-distributed/src/ffi.rs b/crates/xserv-distributed/src/ffi.rs index 2b713b4..00fa57c 100644 --- a/crates/xserv-distributed/src/ffi.rs +++ b/crates/xserv-distributed/src/ffi.rs @@ -45,6 +45,23 @@ unsafe extern "C" { comm: NcclComm, stream: CudaStream, ) -> i32; + // Point-to-point primitives for pipeline parallelism (Phase 18). + pub fn ncclSend( + sendbuff: *const c_void, + count: usize, + datatype: i32, + peer: i32, + comm: NcclComm, + stream: CudaStream, + ) -> i32; + pub fn ncclRecv( + recvbuff: *mut c_void, + count: usize, + datatype: i32, + peer: i32, + comm: NcclComm, + stream: CudaStream, + ) -> i32; pub fn ncclGroupStart() -> i32; pub fn ncclGroupEnd() -> i32; pub fn ncclGetErrorString(result: i32) -> *const c_char; diff --git a/crates/xserv-distributed/src/lib.rs b/crates/xserv-distributed/src/lib.rs index 0feeea5..11f6f5c 100644 --- a/crates/xserv-distributed/src/lib.rs +++ b/crates/xserv-distributed/src/lib.rs @@ -95,3 +95,67 @@ impl Drop for TpContext { } } } + +/// Per-stage pipeline-parallel context: a NCCL communicator spanning all `P` +/// stages plus point-to-point send/recv of the hidden state to the neighbour +/// stages. Init is identical to `TpContext` (one comm across `world` ranks); +/// only the collective differs — PP hands off `[tokens, hidden]` between +/// consecutive stages instead of AllReducing within a layer. +pub struct PpContext { + pub stage: usize, + pub world: usize, + pub device: u32, + comm: NcclComm, +} + +// The NCCL communicator is owned by exactly one stage thread. +unsafe impl Send for PpContext {} + +impl PpContext { + /// Initialize this stage. Must be called from the thread that owns this + /// stage's GPU; binds the thread to `device` first. All stages call this + /// concurrently with the same `id` and `world`. + pub fn init(stage: usize, world: usize, id: NcclUniqueId, device: u32) -> Self { + device::set_device(device).expect("set_device"); + let mut comm: NcclComm = std::ptr::null_mut(); + ffi::check(unsafe { ffi::ncclGroupStart() }, "ncclGroupStart(init)"); + ffi::check( + unsafe { ffi::ncclCommInitRank(&mut comm, world as i32, id, stage as i32) }, + "ncclCommInitRank", + ); + ffi::check(unsafe { ffi::ncclGroupEnd() }, "ncclGroupEnd(init)"); + Self { stage, world, device, comm } + } + + /// Send `count` BF16 elements at `ptr` to `peer`, on the null stream so it is + /// ordered after the producing matmul. Asynchronous — a later `synchronize` + /// (the caller must do one before reusing/freeing the buffer) completes it. + /// + /// # Safety + /// `ptr` must point to at least `count` BF16 elements of valid device memory. + pub fn send_bf16_ptr(&self, ptr: *const c_void, count: usize, peer: usize) { + ffi::check( + unsafe { ffi::ncclSend(ptr, count, ffi::NCCL_BF16, peer as i32, self.comm, NULL_STREAM) }, + "ncclSend", + ); + } + + /// Receive `count` BF16 elements from `peer` into `ptr`, on the null stream. + /// + /// # Safety + /// `ptr` must point to at least `count` BF16 elements of valid device memory. + pub fn recv_bf16_ptr(&self, ptr: *mut c_void, count: usize, peer: usize) { + ffi::check( + unsafe { ffi::ncclRecv(ptr, count, ffi::NCCL_BF16, peer as i32, self.comm, NULL_STREAM) }, + "ncclRecv", + ); + } +} + +impl Drop for PpContext { + fn drop(&mut self) { + if !self.comm.is_null() { + unsafe { ffi::ncclCommDestroy(self.comm) }; + } + } +} diff --git a/crates/xserv-distributed/tests/sendrecv.rs b/crates/xserv-distributed/tests/sendrecv.rs new file mode 100644 index 0000000..1975729 --- /dev/null +++ b/crates/xserv-distributed/tests/sendrecv.rs @@ -0,0 +1,62 @@ +//! 2-GPU NCCL P2P send/recv smoke test for pipeline parallelism. +//! Stage 0 sends a known vector to stage 1, which verifies it. Skips if fewer +//! than 2 GPUs are present. Mirrors `allreduce.rs` (GpuBuffer + half only — +//! this crate does not depend on xserv-tensor). + +use half::bf16; +use std::ffi::c_void; +use std::thread; +use xserv_cuda::{device, GpuBuffer}; +use xserv_distributed::{get_unique_id, PpContext}; + +#[test] +fn pp_send_recv_two_stages() { + 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; // one [1, hidden]-sized hand-off + + let handles: Vec<_> = (0..world) + .map(|stage| { + let id = id; + thread::spawn(move || { + let pp = PpContext::init(stage, world, id, stage as u32); + let mut buf = GpuBuffer::alloc(n * 2).unwrap(); + + if stage == 0 { + // Fill with a known pattern and send to stage 1. + let host: Vec = (0..n).map(|i| bf16::from_f32((i % 97) as f32)).collect(); + let src = unsafe { std::slice::from_raw_parts(host.as_ptr() as *const u8, n * 2) }; + buf.copy_from_host(src).unwrap(); + pp.send_bf16_ptr(buf.as_mut_ptr() as *const c_void, n, 1); + device::synchronize().unwrap(); + None + } else { + // Receive into a zeroed buffer and read it back. + buf.copy_from_host(&vec![0u8; n * 2]).unwrap(); + pp.recv_bf16_ptr(buf.as_mut_ptr() as *mut c_void, n, 0); + device::synchronize().unwrap(); + 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) }; + Some((res[0].to_f32(), res[1].to_f32(), res[n - 1].to_f32())) + } + }) + }) + .collect(); + + let mut checked = false; + for h in handles { + if let Some((first, second, last)) = h.join().unwrap() { + assert_eq!(first, 0.0, "recv[0]"); + assert_eq!(second, 1.0, "recv[1]"); + assert_eq!(last, ((n - 1) % 97) as f32, "recv[last]"); + checked = true; + } + } + assert!(checked, "stage 1 never verified the received buffer"); +} diff --git a/crates/xserv-model/src/qwen3.rs b/crates/xserv-model/src/qwen3.rs index 45f6ff2..446f8e4 100644 --- a/crates/xserv-model/src/qwen3.rs +++ b/crates/xserv-model/src/qwen3.rs @@ -20,6 +20,11 @@ pub struct Qwen3 { tp: Option>, local_num_heads: usize, // = num_heads / world local_num_kv_heads: usize, // = num_kv_heads / world + // Pipeline parallelism (Phase 18): this stage holds a contiguous slice of + // layers. `is_first_stage` owns `embed_tokens`; `is_last_stage` owns + // `norm`/`lm_head_t`. Both true for single-GPU / TP (the whole model). + is_first_stage: bool, + is_last_stage: bool, } struct Qwen3Block { @@ -137,9 +142,267 @@ impl Qwen3 { lm_head_t, rope_cache, tp, + is_first_stage: true, + is_last_stage: true, } } + /// Pipeline-parallel load (Phase 18). This stage holds the contiguous layer + /// range `[stage*L, (stage+1)*L)` with `L = num_layers / num_stages`; only + /// stage 0 keeps `embed_tokens` and only the last stage keeps `norm`/`lm_head` + /// (others get a 1x1 placeholder, guarded by the stage flags and never used). + /// Heads are NOT split (PP is orthogonal to TP), so each stage runs full + /// attention/MLP over its layers and hands off the `[tokens, hidden]` hidden + /// state to the next stage (the engine does the NCCL send/recv). + pub fn from_weights_pp( + config: ModelConfig, + mut w: HashMap, + stage: usize, + num_stages: usize, + device: u32, + ) -> Self { + crate::init_kernels(); + 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}"); + let per_stage = num_layers / num_stages; + let lo = stage * per_stage; + let hi = lo + per_stage; + let is_first_stage = stage == 0; + let is_last_stage = stage == num_stages - 1; + + let take = |w: &mut HashMap, name: &str| -> Tensor { + 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 rope_cache = RopeCache::new( + config.max_seq_len(), + config.head_dim(), + config.rope_theta.unwrap_or(1_000_000.0) as f32, + ); + + let mut layers = Vec::with_capacity(per_stage); + eprintln!( + "[pp] stage {stage}/{num_stages}: layers [{lo}, {hi}) {}{}", + if is_first_stage { "+embed " } else { "" }, + if is_last_stage { "+norm+lm_head" } else { "" } + ); + for i in lo..hi { + let p = format!("model.layers.{i}"); + layers.push(Qwen3Block { + input_norm: repl(take(&mut w, &format!("{p}.input_layernorm.weight"))), + q_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.q_proj.weight"))), + k_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.k_proj.weight"))), + v_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.v_proj.weight"))), + 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"))), + gate_proj_wt: wt(take(&mut w, &format!("{p}.mlp.gate_proj.weight"))), + up_proj_wt: wt(take(&mut w, &format!("{p}.mlp.up_proj.weight"))), + down_proj_wt: wt(take(&mut w, &format!("{p}.mlp.down_proj.weight"))), + }); + } + + Self { + local_num_heads: config.num_heads(), + local_num_kv_heads: config.num_kv_heads(), + config, + embed_tokens, + layers, + norm, + lm_head_t, + rope_cache, + tp: None, + is_first_stage, + is_last_stage, + } + } + + /// Stage-0 token embedding: `[S]` token ids -> `[S, hidden]` hidden state. + pub fn embed(&self, token_ids: &[u32]) -> Tensor { + debug_assert!(self.is_first_stage); + embedding(&self.embed_tokens, token_ids) + } + + /// Last-stage head: `[*, hidden]` -> logits `[*, vocab]`. + pub fn head(&self, x: &Tensor) -> Tensor { + debug_assert!(self.is_last_stage); + let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32; + let x = rmsnorm(x, &self.norm, eps); + 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 } + + /// 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 + /// stage's layers into `paged_cache` (indexed by local layer id) and returns + /// the `[S, hidden]` hidden state to hand to the next stage. Same kernels as + /// `forward_prefill_paged`, minus embedding and the final norm/lm_head. + pub fn forward_layers_prefill( + &self, + mut x: Tensor, + slot: usize, + paged_cache: &mut PagedKVCache, + ) -> Tensor { + let new_tokens = x.shape()[0]; + let pos_offset = paged_cache.seq_len(slot); + 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; + + paged_cache.ensure_capacity(slot, pos_offset + new_tokens); + paged_cache.advance_seq_len(slot, new_tokens); + + let positions: Vec = (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(); + let normed = rmsnorm(&x, &layer.input_norm, eps); + + let q = matmul_2d(&normed, &layer.q_proj_wt); + let k = matmul_2d(&normed, &layer.k_proj_wt); + let v = matmul_2d(&normed, &layer.v_proj_wt); + + 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); + let v = xserv_kernels::reshape_heads_gpu(&v, new_tokens, num_kv_heads, head_dim); + + let q = head_rmsnorm(&q, &layer.q_norm, eps); + let k = head_rmsnorm(&k, &layer.k_norm, eps); + + let q = xserv_kernels::transpose_for_rope_gpu(&q, new_tokens, num_heads, head_dim); + let k = xserv_kernels::transpose_for_rope_gpu(&k, new_tokens, num_kv_heads, head_dim); + rope_inplace(&q, &self.rope_cache, &positions); + rope_inplace(&k, &self.rope_cache, &positions); + let q = xserv_kernels::transpose_from_rope_gpu(&q, new_tokens, num_heads, head_dim); + let k = xserv_kernels::transpose_from_rope_gpu(&k, new_tokens, num_kv_heads, head_dim); + + paged_cache.append_tokens(slot, layer_idx, &k, &v, new_tokens, pos_offset); + 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_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 residual = x_new.clone(); + + let gate = matmul_2d(&normed, &layer.gate_proj_wt); + 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); + x = add_any(&residual, &down); + } + x + } + + /// PP decode over THIS stage's layers. `x` is `[B, hidden]`. Returns + /// `[B, hidden]`. Positions are read from `paged_cache` (all stages advance + /// in lockstep, so they agree). Same kernels as `forward_decode_paged`. + pub fn forward_layers_decode( + &self, + mut x: Tensor, + seq_slots: &[usize], + paged_cache: &mut PagedKVCache, + ) -> Tensor { + let batch = seq_slots.len(); + assert_eq!(x.shape()[0], batch); + 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; + + let positions: Vec = seq_slots.iter().map(|&s| paged_cache.seq_len(s)).collect(); + let kv_lens: Vec = positions.iter().map(|&p| (p + 1) as i32).collect(); + for (b, &slot) in seq_slots.iter().enumerate() { + paged_cache.ensure_capacity(slot, positions[b] + 1); + } + paged_cache.sync_active_batch_with_lens(seq_slots, &kv_lens); + + let bt_ptr = paged_cache.block_table_gpu().as_ptr() as *const i32; + let cl_ptr = paged_cache.context_lens_gpu().as_ptr() as *const i32; + let max_blocks = paged_cache.max_blocks_per_seq(); + + for (layer_idx, layer) in self.layers.iter().enumerate() { + let residual = x.clone(); + let normed = rmsnorm(&x, &layer.input_norm, eps); + + let q_all = matmul_2d(&normed, &layer.q_proj_wt); + let k_all = matmul_2d(&normed, &layer.k_proj_wt); + let v_all = matmul_2d(&normed, &layer.v_proj_wt); + + let mut q_rows: Vec = Vec::with_capacity(batch); + for b in 0..batch { + let q_row = row_view(&q_all, b); + let k_row = row_view(&k_all, b); + let v_row = row_view(&v_all, b); + + let q = xserv_kernels::reshape_heads_gpu(&q_row, 1, num_heads, head_dim); + let k = xserv_kernels::reshape_heads_gpu(&k_row, 1, num_kv_heads, head_dim); + let v = xserv_kernels::reshape_heads_gpu(&v_row, 1, num_kv_heads, head_dim); + + let q = head_rmsnorm(&q, &layer.q_norm, eps); + let k = head_rmsnorm(&k, &layer.k_norm, eps); + + let q = xserv_kernels::transpose_for_rope_gpu(&q, 1, num_heads, head_dim); + let k = xserv_kernels::transpose_for_rope_gpu(&k, 1, num_kv_heads, head_dim); + + let pos = [positions[b] as u32]; + rope_inplace(&q, &self.rope_cache, &pos); + rope_inplace(&k, &self.rope_cache, &pos); + + let q = xserv_kernels::transpose_from_rope_gpu(&q, 1, num_heads, head_dim); + let k = xserv_kernels::transpose_from_rope_gpu(&k, 1, num_kv_heads, head_dim); + + paged_cache.append_tokens(seq_slots[b], layer_idx, &k, &v, 1, positions[b]); + + let q_flat = xserv_kernels::merge_heads_gpu(&q, 1, num_heads, head_dim); + q_rows.push(q_flat); + } + + let q_batched_2d = concat_rows(&q_rows); + let q_4d = q_batched_2d.reshape(&[batch, num_heads, 1, head_dim]); + + 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, + ); + + 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 residual = x_new.clone(); + + let gate = matmul_2d(&normed, &layer.gate_proj_wt); + 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); + x = add_any(&residual, &down); + } + + for &slot in seq_slots { + paged_cache.advance_seq_len(slot, 1); + } + x + } + /// 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] diff --git a/crates/xserv-model/src/sampling.rs b/crates/xserv-model/src/sampling.rs index 558f0cf..97bd01a 100644 --- a/crates/xserv-model/src/sampling.rs +++ b/crates/xserv-model/src/sampling.rs @@ -2,6 +2,7 @@ use half::bf16; use rand::Rng; use xserv_tensor::{DType, Device, Tensor}; +#[derive(Clone)] pub struct SamplingParams { pub temperature: f32, pub top_k: usize, diff --git a/crates/xserv-server/src/main.rs b/crates/xserv-server/src/main.rs index 6927161..e030c3d 100644 --- a/crates/xserv-server/src/main.rs +++ b/crates/xserv-server/src/main.rs @@ -1,5 +1,6 @@ mod api; mod engine; +mod pp_engine; mod tp_engine; use axum::{routing::{get, post}, Extension, Router}; @@ -19,7 +20,7 @@ pub struct AppState { async fn main() { let args: Vec = std::env::args().collect(); if args.len() < 2 { - eprintln!("Usage: xserv-server [--port PORT] [--max-batch N] [--max-seq-len N] [--swap-space-gb N] [--tp N]"); + eprintln!("Usage: xserv-server [--port PORT] [--max-batch N] [--max-seq-len N] [--swap-space-gb N] [--tp N] [--pp N]"); std::process::exit(1); } @@ -52,6 +53,16 @@ async fn main() { .and_then(|s| s.parse().ok()) .unwrap_or(1) .max(1); + let pp: usize = args.iter() + .position(|a| a == "--pp") + .and_then(|i| args.get(i + 1)) + .and_then(|s| s.parse().ok()) + .unwrap_or(1) + .max(1); + if tp > 1 && pp > 1 { + eprintln!("--tp and --pp cannot be combined yet (2D TP×PP is future work)"); + std::process::exit(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 { @@ -76,7 +87,10 @@ async fn main() { let model_dir_clone = model_dir.clone(); std::thread::spawn(move || { - if tp <= 1 { + if pp > 1 { + // Pipeline-parallel path: stage-0 coordinator + worker stage threads. + pp_engine::run_pp(&model_dir_clone, pp, max_seq_len, rx); + } else 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 { diff --git a/crates/xserv-server/src/pp_engine.rs b/crates/xserv-server/src/pp_engine.rs new file mode 100644 index 0000000..35685bd --- /dev/null +++ b/crates/xserv-server/src/pp_engine.rs @@ -0,0 +1,264 @@ +//! Pipeline-parallel inference engine for the HTTP server (Phase 18). +//! +//! Layer-wise split: stage `s` holds layers `[s*L, (s+1)*L)`. Stage 0 owns the +//! token embedding and acts as the coordinator (scheduler + tokenizer + response +//! sender + stop logic); the last stage owns `norm`/`lm_head` and does sampling. +//! Hidden states are handed off stage->stage via NCCL P2P (`PpContext`); the +//! sampled token id (a single u32) is returned last-stage -> stage0 over an +//! in-process channel (same process, so no NCCL needed for that). +//! +//! v1 is serial: one request at a time, one token per step, the pipeline is +//! filled and drained each step (stage0's decode step t+1 depends on the token +//! the last stage sampled at step t). This gives correctness + per-GPU memory +//! savings; throughput via microbatch/1F1B overlap is future work +//! (see docs/18-pipeline-parallelism.md). + +use std::ffi::c_void; +use std::path::{Path, PathBuf}; +use std::sync::mpsc; +use std::sync::Arc; +use std::thread; + +use half::bf16; +use xserv_distributed::{PpContext, UniqueId}; +use xserv_model::loader; +use xserv_model::sampling::SamplingParams; +use xserv_model::{sample, ModelConfig, PagedKVCache, Qwen3, BLOCK_SIZE}; +use xserv_tensor::{DType, Device, Tensor}; +use xserv_tokenizer::Tokenizer; + +use crate::engine::{GenerateEvent, GenerateRequest}; + +/// Control messages from the coordinator (stage 0) to a worker stage. The heavy +/// hidden-state tensors do NOT travel here — they go GPU->GPU over NCCL. Only +/// tiny control info (slot ids, token count, sampling params) is sent. +#[derive(Clone)] +enum PpCommand { + Register(usize), + Free(usize), + /// Receive `[n_tokens, hidden]` from the previous stage, run this stage's + /// layers; if last stage, sample with `sampling` and return the token. + Prefill { n_tokens: usize, slot: usize, sampling: SamplingParams }, + /// Receive `[1, hidden]`, run this stage's layers; last stage samples. + Decode { slot: usize, sampling: SamplingParams }, + Shutdown, +} + +struct StageCtx { + model: Qwen3, + cache: PagedKVCache, + pp: Arc, + hidden: usize, + device: u32, +} + +/// Build this stage: NCCL init, load + slice weights, size a per-stage KV pool +/// for THIS stage's layers only (so per-GPU KV is ~1/P). +fn build_stage( + model_dir: &Path, + config: &ModelConfig, + stage: usize, + world: usize, + device: u32, + max_seq_len: usize, + id: UniqueId, +) -> StageCtx { + let pp = Arc::new(PpContext::init(stage, world, id, device)); + let weights = loader::load_model_dir(model_dir, Device::Cpu); + let model = Qwen3::from_weights_pp(config.clone(), weights, stage, world, device); + + // The KV cache only needs this stage's layers; build it from a config clone + // whose layer count is the per-stage count (heads are NOT split under PP). + let per_stage = config.num_layers() / world; + let mut stage_config = config.clone(); + stage_config.num_hidden_layers = Some(per_stage); + + let max_blocks_per_seq = max_seq_len.div_ceil(BLOCK_SIZE); + let total_blocks = max_blocks_per_seq + 8; // v1 serial: one active sequence + let cache = PagedKVCache::new( + &stage_config, total_blocks, 0, 4, max_blocks_per_seq, DType::BF16, device, + ); + StageCtx { model, cache, pp, hidden: config.hidden(), device } +} + +/// Allocate a zeroed `[n, hidden]` device tensor and receive into it from `peer`. +fn recv_hidden(sc: &StageCtx, n: usize, peer: usize) -> Tensor { + let zeros = vec![bf16::ZERO; n * sc.hidden]; + let x = Tensor::from_slice(&zeros, &[n, sc.hidden]).to_device(Device::Cuda(sc.device)); + let ptr = x.storage().gpu_buffer().as_ptr() as *mut c_void; + sc.pp.recv_bf16_ptr(ptr, n * sc.hidden, peer); + xserv_cuda::device::synchronize().unwrap(); + x +} + +/// Send the `[*, hidden]` hidden state to `peer`, then synchronize so NCCL has +/// finished reading `x` before it is dropped/reused. +fn send_hidden(sc: &StageCtx, x: &Tensor, peer: usize) { + let ptr = x.storage().gpu_buffer().as_ptr() as *const c_void; + sc.pp.send_bf16_ptr(ptr, x.numel(), peer); + xserv_cuda::device::synchronize().unwrap(); +} + +fn worker_loop( + stage: usize, + world: usize, + id: UniqueId, + model_dir: PathBuf, + config: ModelConfig, + max_seq_len: usize, + cmd_rx: mpsc::Receiver, + ack_tx: mpsc::Sender<()>, + token_tx: mpsc::Sender, +) { + let mut sc = build_stage(&model_dir, &config, stage, world, stage as u32, max_seq_len, id); + let is_last = stage == world - 1; + let prev = stage - 1; + let next = stage + 1; + + while let Ok(cmd) = cmd_rx.recv() { + match cmd { + PpCommand::Register(slot) => { + let _ = sc.cache.register_sequence(slot); + let _ = ack_tx.send(()); + } + PpCommand::Free(slot) => { + sc.cache.free_sequence(slot); + let _ = ack_tx.send(()); + } + PpCommand::Prefill { n_tokens, slot, sampling } => { + let x = recv_hidden(&sc, n_tokens, prev); + let x = sc.model.forward_layers_prefill(x, slot, &mut sc.cache); + if is_last { + let logits = sc.model.head(&x); + let _ = token_tx.send(sample(&logits, &sampling)); + } else { + send_hidden(&sc, &x, next); + } + } + PpCommand::Decode { slot, sampling } => { + let x = recv_hidden(&sc, 1, prev); + let x = sc.model.forward_layers_decode(x, &[slot], &mut sc.cache); + if is_last { + let logits = sc.model.head(&x); + let _ = token_tx.send(sample(&logits, &sampling)); + } else { + send_hidden(&sc, &x, next); + } + } + PpCommand::Shutdown => { + let _ = ack_tx.send(()); + break; + } + } + } +} + +/// Run the PP coordinator (stage 0) on the calling thread. Spawns worker stages +/// 1..world and consumes generation requests from `rx`. +pub fn run_pp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Receiver) { + assert!(world >= 2, "run_pp requires world >= 2"); + let config = ModelConfig::from_file(&model_dir.join("config.json")); + assert!( + config.num_layers() % world == 0, + "num_layers {} not divisible by pp {world}", + config.num_layers() + ); + let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json")); + let id = xserv_distributed::get_unique_id(); + + // Worker stages 1..world. Each gets a control channel; all share one ack + // channel and one token channel (only the last stage actually sends tokens). + let (ack_tx, ack_rx) = mpsc::channel::<()>(); + let (token_tx, token_rx) = mpsc::channel::(); + let mut cmd_txs: Vec> = Vec::new(); + for stage in 1..world { + let (ctx_tx, ctx_rx) = mpsc::channel::(); + cmd_txs.push(ctx_tx); + let ack_tx = ack_tx.clone(); + let token_tx = token_tx.clone(); + let model_dir = model_dir.to_path_buf(); + let config = config.clone(); + thread::spawn(move || { + worker_loop(stage, world, id, model_dir, config, max_seq_len, ctx_rx, ack_tx, token_tx); + }); + } + + // Stage 0 (this thread): coordinator + embedding + first layers. + let mut sc = build_stage(model_dir, &config, 0, world, 0, max_seq_len, id); + eprintln!("[pp-engine] ready (pp={world}, max_seq_len={max_seq_len})"); + + let eos = tokenizer.eos_token_id(); + let n_workers = world - 1; + let next_peer = 1usize; + let broadcast = |txs: &[mpsc::Sender], cmd: PpCommand| { + 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, PpCommand::Register(slot)); + sc.cache.register_sequence(slot).expect("register slot"); + wait_acks(&ack_rx); + + // Prefill: embed prompt, run stage-0 layers, push hidden into the pipe. + broadcast(&cmd_txs, PpCommand::Prefill { + n_tokens: req.prompt_tokens.len(), + slot, + sampling: req.sampling.clone(), + }); + let x = sc.model.embed(&req.prompt_tokens); + let x = sc.model.forward_layers_prefill(x, slot, &mut sc.cache); + send_hidden(&sc, &x, next_peer); + let mut next = token_rx.recv().expect("prefill token"); + + let mut decode_buf: Vec = 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"; + } + broadcast(&cmd_txs, PpCommand::Decode { slot, sampling: req.sampling.clone() }); + let x = sc.model.embed(&[next]); + let x = sc.model.forward_layers_decode(x, &[slot], &mut sc.cache); + send_hidden(&sc, &x, next_peer); + next = token_rx.recv().expect("decode token"); + 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, PpCommand::Free(slot)); + sc.cache.free_sequence(slot); + wait_acks(&ack_rx); + } + + broadcast(&cmd_txs, PpCommand::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, buf: &mut Vec) { + 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 }); + } +} diff --git a/docs/18-pipeline-parallelism.md b/docs/18-pipeline-parallelism.md new file mode 100644 index 0000000..aeee592 --- /dev/null +++ b/docs/18-pipeline-parallelism.md @@ -0,0 +1,151 @@ +# Phase 18: Pipeline Parallelism (PP) + +> 目标:在单机多卡上做 **流水线并行**,把 Qwen3-8B 的 **层** 切成 `P` 段(stage), +> 每张卡只持有连续的一段层(+ stage0 的 `embed_tokens`、最后一段的 `norm`/`lm_head`), +> 激活(hidden state)在相邻 stage 之间用 **NCCL P2P send/recv** 传递。 +> 与 TP(按 head / 中间维切,每层 2 次 AllReduce)互补:PP 通信量小(每 token 仅 `P-1` +> 次点对点传 `[tokens, hidden]`),KV 与权重按 **层** 降到约 1/P。 +> 先做 **PP=2 / 4(组内)**,正确性优先。 + +## 1. 硬件约束(dash5) + +- 8× RTX 5090(32GB,SM120),**无 NVLink**,纯 PCIe Gen5。 +- 拓扑:GPU 0–3 一组、4–7 一组,组内 `PHB`(同 host bridge,可 P2P),跨组 `NODE`。 +- **PP 同样建议在组内**(0–3 或 4–7):虽然 PP 的通信量远小于 TP,但 P2P 仍走 PCIe, + 跨组延迟更高。PP=2/4 用 0–1 / 0–3。 +- 相比 TP:TP 每 token `2·layers = 72` 次 AllReduce(延迟主导);PP 每 token 仅 + `P-1` 次 send/recv,每次 `[tokens, hidden]` BF16(decode batch=1 时 8KB)。 + **PP 对慢互联(PCIe / 无 NVLink)更友好**,这是在 dash5 上做 PP 的主要动机之一。 + +## 2. 切分方案(layer-wise) + +Qwen3-8B:`hidden=4096`、`num_heads=32`、`num_kv_heads=8`、`head_dim=128`、 +`intermediate=12288`、`layers=36`、`vocab=151936`。`36` 能被 `2/4` 整除(PP=3/6 需处理余数, +本阶段先要求 `layers % P == 0`)。 + +设 stage 数 `P`,本 stage = `s`,每段 `L = layers / P` 层,本段持有全局层 +`[s·L, (s+1)·L)`: + +| 组件 | 持有者 | 说明 | +|------|--------|------| +| `embed_tokens` `[vocab, hidden]` | **仅 stage 0** | token → hidden | +| transformer block `i` 的全部权重 | 持有 `i` 的那个 stage | 不切 head / 中间维(与 TP 正交) | +| 该层 KV cache | 持有 `i` 的那个 stage | **每卡 KV 降到约 1/P** | +| 最终 `norm` `[hidden]` | **仅最后一段** | | +| `lm_head` `[vocab, hidden]` | **仅最后一段** | hidden → logits | + +- 注意力 / MLP 的层内计算 **完全不变**(不需要 AllReduce):每个 stage 用它自己那几层 + 的完整权重、完整 head 做 forward。PP 与 TP 正交,可叠加(本阶段不实现 TP×PP)。 +- **RoPE** 用全局绝对 position,每个 stage 的 `RopeCache` 完全相同(按 position 索引), + 各 stage 独立做,无需通信。 +- **每个 stage 一个独立的 `PagedKVCache`**,层数 = 本段层数 `L`(不是 36)。forward 时 + 按「本段内的局部层号 `0..L`」索引 cache —— 与单卡代码完全一致,只是 `self.layers` + 只装了本段的层。实现技巧:给 cache 传一个 `num_hidden_layers` 改写成 `L` 的 config 克隆, + **无需改 `PagedKVCache`**。 + +### 通信点 +- prefill:stage `s` 算完本段层,得到 `[S, hidden]` → **send 给 `s+1`**;`s+1` recv 后接着算。 +- decode:同理传 `[B, hidden]`(batch=1 时 `[1, hidden]`)。 +- 每 token 共 `P-1` 次 send/recv;最后一段算出 logits 并采样。 +- 采样得到的 token id(一个 `u32`)由 **最后一段经线程内 channel 回传给 stage0** + (同进程多线程,无需走 NCCL)。 + +## 3. 进程 / 线程模型 + +沿用 TP 的 **单进程、多线程**:每个 stage 一个 OS 线程,线程启动时 `cudaSetDevice(stage)`。 +- **stage 0 = 协调者(coordinator)**,跑在调用线程上:持有 scheduler、tokenizer、HTTP + response sender、停止判定(eos / max_tokens)与「下一步输入 token」。 +- **stage 1..P-1 = worker 线程**:从控制 channel 收命令(Register/Prefill/Decode/Free/Shutdown), + 每步 `recv` 上游 hidden → 跑本段层 → `send` 给下游;最后一段 `head`+采样 → 把 token 回传 stage0。 +- 控制信息(命令、采样参数、token id)走 `mpsc`(极小);**重活(hidden 张量)走 NCCL P2P(GPU↔GPU)**。 + +> **v1 串行语义**:一次处理一个请求、一次一个 token,流水线每步「灌满又排空」 +> (stage0 decode 第 `t+1` 步依赖最后一段第 `t` 步采出的 token)。这保证 **正确性**, +> 并拿到 TTFT/TPOT 与每卡显存;**throughput 的真正收益来自 microbatch/请求级流水线 +> 重叠(1F1B)**,列为后续工作(见 §7)。 + +执行流(每请求): +``` +coordinator worker s (1..P-1) last stage (P-1) +───────────── ───────────────── ──────────────── +broadcast Register(slot) cache.register(slot) cache.register(slot) +broadcast Prefill{n,slot,samp} + x=embed(prompt) + x=layers_prefill(x,slot) + send x → stage1 recv x ← s-1 + x=layers_prefill(x,slot) + send x → s+1 ───────────────► recv x ← P-2 + x=layers_prefill(x,slot) + logits=head(x); next=sample + next ◄────────────── token channel ◄────────────────────── token_tx.send(next) + stream(next); loop Decode{slot} 直到 eos/length +broadcast Free(slot) cache.free(slot) cache.free(slot) +``` + +## 4. 通信库:NCCL P2P + +复用 `xserv-distributed`(已有 NCCL FFI + `TpContext`/AllReduce),新增: +- FFI:`ncclSend(sendbuff, count, dtype, peer, comm, stream)`、 + `ncclRecv(recvbuff, count, dtype, peer, comm, stream)`。 +- `PpContext`:与 `TpContext` 同样的 `ncclCommInitRank`(一个 comm 跨 `P` 个 stage), + 外加 `send_bf16_ptr(ptr, count, peer)` / `recv_bf16_ptr(ptr, count, peer)`,在 **null + stream** 上发起(与模型 kernel 同流,天然有序)。 +- 线性流水线无死锁:stage0 只 send、最后一段只 recv、中间段「先 recv 上游、再 send 下游」, + 依赖链无环,从头解锁。每个 stage 在 send/recv + 本段计算后 `synchronize()`, + 确保 NCCL 读完发送缓冲再复用/释放(v1 串行下成本可接受)。 + +> **决策点**:和 TP 一样,collective/P2P 先用 NCCL 把 PP 跑通拿正确性与基线; +> 手写 P2P(PCIe 上的 cudaMemcpyPeer)作为后续学习项。 + +## 5. 权重分片加载 + +`Qwen3::from_weights_pp(config, weights, stage, num_stages, device)`: +- 只把全局层 `[s·L, (s+1)·L)` 搬到本 stage 的 GPU(其余层的权重直接 drop,不占显存)。 +- `embed_tokens`:仅 stage 0 加载;其余 stage 放一个 1×1 占位张量(forward 用 `is_first_stage` + 守卫,永不触碰)。 +- `norm`/`lm_head`:仅最后一段加载;其余放占位。 +- head 不切(不做 TP),所以 `local_num_heads = num_heads`、`local_num_kv_heads = num_kv_heads`。 + +每卡显存 ≈ `权重(transformer 1/P) + KV(1/P) + (stage0: embed) + (last: norm+lm_head)`。 +对 Qwen3-8B:transformer 层约 14GB,PP=2 每卡约 7GB 层权重 + embed 或 lm_head(各 ~1.2GB)。 + +## 6. 实现步骤(逐步可验证) + +1. **P18.1 — `xserv-distributed` P2P**:`ncclSend/Recv` FFI + `PpContext`。 + 验收:2 卡,rank0 send 已知向量、rank1 recv,校验一致(`tests/sendrecv.rs`)。 +2. **P18.2 — 分段权重加载**:`from_weights_pp`,每 stage 只持有本段层 + 该有的 embed/head。 + 验收:各 stage 层数 = `L`、显存约 1/P(+ embed/head)。 +3. **P18.3 — stage forward**:`embed` / `forward_layers_prefill` / `forward_layers_decode` / + `head`,每段独立 KV cache。 + 验收:**PP=1 与单卡 `forward_*_paged` 逐 token 一致**(同一条代码路径退化)。 +4. **P18.4 — PP engine + `--pp N`**:多线程 stage workers + NCCL 传递 + stage0 协调。 + 验收:`--pp 2/4` 端到端可服务;**greedy 输出与单卡(PP=1)逐 token 一致**; + 用现有 llama.cpp bench 跑正确性(GSM8K/AIME);测 PP=1/2/4 的 TTFT/TPOT/每卡显存。 + +## 7. 预期与风险 + +- **显存**:每卡 transformer 权重 + KV ≈ 1/P,这是 PP 的主要收益(可上更大模型 / 更长 context)。 +- **单流吞吐**:v1 串行无 stage 重叠 → 单流 tok/s **不会超过单卡**(多一份 P2P + sync 开销, + 可能略低)。这是 PP 的本质:**没有 microbatch 重叠就没有加速**。诚实记录实测,并与 + llama.cpp 的 `--split-mode layer`(同样是层切流水线、单序列也串行跨卡)对比 —— 两者单流 + 都应≈单卡。 +- **真正的 throughput 收益**(后续):请求级 / microbatch 流水线(1F1B),让 stage 间重叠: + stage1 算 microbatch A 时 stage0 算 B。需要把 scheduler 改成跨 stage 连续批处理。 +- **风险**:NCCL 多线程 init 同步;send 缓冲生命周期(必须 sync 后再复用); + `layers % P != 0` 的余数分配(本阶段先约束整除);与 CUDA Graph decode 的结合(先走非 graph 路径)。 +- 正确性优先:先 PP=1 等价(逐 token 对齐),PP=2/4 与单卡对齐,再谈性能。 + +## 8. 与 llama.cpp 的对比口径 + +- **xserv**:`--pp N`,`CUDA_VISIBLE_DEVICES=0..N-1`。 +- **llama.cpp**:`-sm layer`(默认即层切流水线)+ `--tensor-split` 均分层,`CUDA_VISIBLE_DEVICES=0..N-1`。 + (对照 TP 用的是 `-sm row`。) +- 指标:正确性(GSM8K / AIME exact-match)、单流 TTFT/TPOT、并发吞吐、每卡 VRAM。 +- 复用 `tools/bench/runner.py` 与 `run_pp_parallel.sh`(仿 `run_tp_parallel.sh`)。 + +## 9. 不在本阶段范围 + +- TP×PP 混合(2D 并行)、跨组 / 多节点。 +- microbatch / 1F1B 流水线重叠(throughput 收益,后续)。 +- vocab-parallel embedding / lm_head。 +- `layers % P != 0` 的非均匀切分;与 CUDA Graph decode 结合。 diff --git a/docs/benchmarks/pp-sweep.md b/docs/benchmarks/pp-sweep.md new file mode 100644 index 0000000..5f1f5cd --- /dev/null +++ b/docs/benchmarks/pp-sweep.md @@ -0,0 +1,118 @@ +# PP sweep — xserv vs llama.cpp (Qwen3-8B BF16, 8×RTX 5090) + +Pipeline parallelism (layer split), verified end-to-end on dash5. Qwen3-8B BF16, +greedy, single stream, no NVLink (hand-off / split traffic over PCIe Gen5). +xserv `--pp N` puts stage `s` on GPU `s` and hands the hidden state stage→stage +over NCCL P2P; llama.cpp uses `-sm layer` (its default pipeline split) over N GPUs. + +## Single-stream latency + per-GPU VRAM (measured, `--max-seq-len 2048`) + +Measured strictly sequentially, one server at a time, each config gated on a real +successful generation (so VRAM snapshots are post-load). Driver: +`tools/pp_final.sh`. + +| engine | PP | TTFT_ms | TPOT_ms | tok/s | per-GPU VRAM (MiB) | +|--------|----|---------|---------|-------|--------------------| +| xserv | 1 | 33.2 | 17.39 | 57.5 | 24010 | +| xserv | 2 | 35.9 | 18.07 | 55.3 | 11580, 13632 | +| xserv | 4 | 36.1 | 17.91 | 55.8 | 7298, 5250, 5250, 9350 | +| llama | 1 | 133.3 | 9.38 | 106.7 | 15604 | +| llama | 2 | 131.4 | 9.10 | 109.9 | 7862, 8494 | +| llama | 4 | 161.2 | 8.88 | 112.6 | 4476, 4090, 4090, 5108 | + +(xserv VRAM with `XSERV_MAX_KV_BLOCKS=160` so the number is weights + a minimal +KV pool. `tok/s = 1000 / TPOT`. This latency probe's TTFT differs from the +quality-suite TTFT below because the suite includes scheduler/HTTP overhead.) + +## Correctness — PP is numerically exact + +The hidden-state hand-off between stages is a bit-exact BF16 P2P copy and each +stage runs the same kernels over its layers, so PP must reproduce the single-GPU +result. Verified by byte-comparing generated text (greedy, temp 0), running each +config **twice** to separate PP effects from run-to-run GEMM noise: + +| comparison | result | +|------------|--------| +| single run A == single run B | **DIFFER** (cuBLAS GEMM is not bit-reproducible run-to-run) | +| pp4 run A == pp4 run B | **IDENTICAL** | +| single run A == pp4 run A | **IDENTICAL** | +| single == pp2 (single run each) | **IDENTICAL** | + +Takeaway: **single-GPU itself is non-deterministic** under greedy (a 1-ULP logit +difference flips a late argmax and the suffix changes), so a one-shot single-vs-PP +byte compare can spuriously "DIFFER". The 2×2 control shows PP=4 is *more* +reproducible than re-running single-GPU, and it lands exactly on a single-GPU +trajectory. NCCL P2P (`tests/sendrecv.rs`) and AllReduce (`tests/allreduce.rs`) +unit tests pass. + +## Quality matrix — AIME 2025 (30) + GSM8K (30), greedy, both engines × PP=1/2/4 + +Full measured matrix (`tools/bench/summarize_fullq.py`; raw in +`bench-out/FULLQ_SUMMARY.txt`). Qwen3-8B BF16, thinking OFF, `max_seq_len 4096`. +xserv on GPUs 0-3, llama.cpp on GPUs 4-7 (disjoint groups, run in parallel). + +| engine | PP | AIME 2025 | GSM8K | AIME mean_tok | TTFT_ms | TPOT_ms | +|--------|----|-----------|-------|---------------|---------|---------| +| xserv | 1 | 8/30 (26.7%) | 29/30 (96.7%) | 2383 | 485 | 22.42 | +| xserv | 2 | 7/30 (23.3%) | 29/30 (96.7%) | 2367 | 457 | 22.55 | +| xserv | 4 | 7/30 (23.3%) | 29/30 (96.7%) | 2652 | 494 | 23.31 | +| llama | 1 | 7/30 (23.3%) | 29/30 (96.7%) | 2651 | 119 | 10.37 | +| llama | 2 | 7/30 (23.3%) | 29/30 (96.7%) | 2651 | 118 | 10.41 | +| llama | 4 | 7/30 (23.3%) | 29/30 (96.7%) | 2651 | 119 | 10.39 | + +Reading the matrix: + +- **GSM8K = 29/30 (96.7%) in every cell** — identical across both engines and all + PP levels. xserv's accuracy matches llama.cpp exactly on the same weights. +- **AIME = 7/30 (23.3%) everywhere except xserv PP=1 (8/30)**. That single +1 is + the run-to-run greedy nondeterminism documented above (an AIME solution is + ~2400 tokens; one late argmax flip changes one problem's outcome) — not a PP or + engine effect. AIME accuracy is low because this is an 8B model with thinking + disabled; the point here is the *cross-engine / cross-PP agreement*, which holds. +- **TPOT is flat across PP** for both engines (xserv 22.4→23.3 ms, llama + 10.3→10.4 ms), reconfirming PP doesn't slow single-stream decode. The ~2.2× + TPOT gap to llama.cpp is the single-GPU gap (`llama-cpp-comparison.md`), + orthogonal to PP. + +## Takeaways + +- **Memory is the win.** Per-GPU weights+KV scale ~1/P: xserv 24.0 GB (1 GPU) → + ~11–14 GB (PP=2) → ~5–9 GB (PP=4); llama 15.6 → ~8 → ~4–5 GB. The two end + stages sit higher (stage 0 holds `embed_tokens`, the last stage `norm`+`lm_head`, + ~1.1 GB each). This is what PP buys: a model / context that does not fit on one + card fits across P. +- **Single-stream latency is flat, not faster.** v1 PP is serial across stages + (no microbatch overlap): per-token latency = sum of all stages' compute + + (P-1) P2P hops + a blocking sync per stage. The `[1, hidden]` BF16 hop (8 KB) + over PCIe is cheap relative to per-token compute, so TPOT is ~constant across P. + PP does **not** speed up single-stream decode; it trades (almost no) latency for + large memory headroom. +- **Quality is preserved and matches llama.cpp.** GSM8K 96.7% in all 12 cells; + AIME within the greedy noise band. PP=1/2/4 agree, and xserv tracks llama.cpp. + +## Reproduce + +```bash +./tools/sync-and-build.sh build +# latency + VRAM + byte-exact correctness (writes bench-out/PP_FINAL.md): +ssh 'cd && bash tools/pp_final.sh' +# determinism control (single×2 vs pp4×2): +ssh 'cd && bash tools/pp_diag.sh' +# NCCL P2P + AllReduce unit tests: +ssh 'cd && cargo test -p xserv-distributed --release' +# full quality matrix AIME-30 + GSM8K-30 (xserv 0-3 serial; or parallel w/ llama 4-7): +ssh 'cd && bash tools/pp_quality_full.sh' # xserv+llama serial, GPU 0-3 +ssh 'cd && bash tools/pp_llama_47.sh' # llama on GPU 4-7 (parallel) +python3 tools/bench/summarize_fullq.py bench-out +``` + +## Next (where PP actually raises throughput) + +- **Microbatch / 1F1B overlap**: while stage 1 runs microbatch A, stage 0 runs B. + This is the only thing that turns PP into a *throughput* win; v1 is serial, so + P GPUs give 1 GPU's single-stream rate (but P× the memory headroom / batch room). +- Persistent per-stage recv buffers (drop the per-token CPU alloc + H2D) and + event-based ordering instead of a full device sync per hop. +- 2D TP×PP, and `layers % P != 0` non-uniform splits. + +🤖 Generated with [Claude Code](https://claude.com/claude-code) diff --git a/tools/bench/pp_clean_bench.sh b/tools/bench/pp_clean_bench.sh new file mode 100644 index 0000000..9bf34b7 --- /dev/null +++ b/tools/bench/pp_clean_bench.sh @@ -0,0 +1,89 @@ +#!/usr/bin/env bash +# Clean, strictly-sequential single-stream latency + per-GPU VRAM for PP. +# One server at a time. Readiness = first SUCCESSFUL generation (xserv's /health +# returns 200 before the model finishes loading, so we must not gate on it). +# Snapshots are therefore always post-load. Writes bench-out/PP_CLEAN.md. +# +# Env overrides: MODEL, GGUF, PPS (default "1 2 4"), LLAMA_BIN. +set -u +cd "$(dirname "$0")/../.." +export PATH=$HOME/.cargo/bin:/usr/local/cuda-12.9/bin:$PATH +export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda-12.9} +MODEL=${MODEL:-/opt/wjh/models/qwen3-8b} +GGUF=${GGUF:-/opt/wjh/models/qwen3-8b/qwen3-8b-bf16.gguf} +LLAMA_BIN=${LLAMA_BIN:-third_party/llama.cpp/build/bin/llama-server} +XBIN=./target/release/xserv-server +PPS=${PPS:-1 2 4} +PROMPT='Write a detailed paragraph explaining how GPUs accelerate neural network training.' +OUT=bench-out/PP_CLEAN.md +mkdir -p bench-out +: > "$OUT" +echo "# PP clean single-stream latency + VRAM — $(date)" >> "$OUT" +echo "" >> "$OUT" +echo "| engine | PP | TTFT_ms | TPOT_ms | tok/s | per-GPU VRAM (MiB) |" >> "$OUT" +echo "|--------|----|---------|---------|-------|--------------------|" >> "$OUT" + +killall_servers(){ pkill -9 -f xserv-server 2>/dev/null; pkill -9 -f llama-server 2>/dev/null; sleep 3; } + +drain(){ # wait until GPUs $1 (csv) all < 1500 MiB, max 120s + for _ in $(seq 1 60); do + local hi=0 + for g in ${1//,/ }; do + m=$(nvidia-smi -i "$g" --query-gpu=memory.used --format=csv,noheader,nounits) + [ "${m:-0}" -gt 1500 ] && hi=1 + done + [ "$hi" -eq 0 ] && return 0; sleep 2 + done +} + +# probe_ready PORT PID -> 0 when a generation succeeds (deadline ~1200s) +probe_ready(){ local port=$1 pid=$2 + for _ in $(seq 1 400); do + if curl -s -o /dev/null -w '%{http_code}' --max-time 8 \ + "http://127.0.0.1:$port/v1/chat/completions" -H 'Content-Type: application/json' \ + -d '{"model":"qwen3-8b","messages":[{"role":"user","content":"hi"}],"max_tokens":1,"temperature":0,"stream":false}' \ + 2>/dev/null | grep -q 200; then return 0; fi + kill -0 "$pid" 2>/dev/null || return 1 + sleep 3 + done; return 1 +} + +vram(){ local cvd=$1; local a b="" # stabilized snapshot of GPUs $cvd + for _ in $(seq 1 12); do + a=$(for g in ${cvd//,/ }; do nvidia-smi -i "$g" --query-gpu=memory.used --format=csv,noheader,nounits; done | paste -sd' ') + [ "$a" = "$b" ] && break; b=$a; sleep 2 + done; echo "$a" +} + +run_xserv(){ local pp=$1; local cvd; cvd=$(seq -s, 0 $((pp-1))) + killall_servers; drain "$cvd" + local extra=""; [ "$pp" -gt 1 ] && extra="--pp $pp" + XSERV_MAX_KV_BLOCKS=160 CUDA_VISIBLE_DEVICES=$cvd nohup $XBIN $MODEL --port 8090 --max-seq-len 2048 $extra >/tmp/x$pp.log 2>&1 & + local pid=$! + if ! probe_ready 8090 "$pid"; then echo "| xserv | $pp | FAILED (see /tmp/x$pp.log) | | | |" >> "$OUT"; kill -9 "$pid" 2>/dev/null; return; fi + local mib; mib=$(vram "$cvd") + local m; m=$(python3 tools/bench/pp_time.py http://127.0.0.1:8090 "$PROMPT") + local ttft tpot toks; ttft=$(echo "$m"|sed -n 's/.*TTFT_ms=\([0-9.]*\).*/\1/p'); tpot=$(echo "$m"|sed -n 's/.*TPOT_ms=\([0-9.a-z]*\).*/\1/p'); toks=$(echo "$m"|sed -n 's/.*tok_s=\([0-9.a-z]*\).*/\1/p') + echo "| xserv | $pp | $ttft | $tpot | $toks | $mib |" >> "$OUT" + kill -9 "$pid" 2>/dev/null; wait "$pid" 2>/dev/null; sleep 3 +} + +run_llama(){ local pp=$1; local cvd; cvd=$(seq -s, 0 $((pp-1))) + killall_servers; drain "$cvd" + local sm=(-sm none); [ "$pp" -gt 1 ] && sm=(-sm layer -ts "$(printf '1%.0s,' $(seq 1 $pp) | sed 's/,$//')") + CUDA_VISIBLE_DEVICES=$cvd nohup $LLAMA_BIN -m $GGUF --port 8090 --host 127.0.0.1 \ + -c 2048 --parallel 1 -ngl 999 "${sm[@]}" >/tmp/l$pp.log 2>&1 & + local pid=$! + if ! probe_ready 8090 "$pid"; then echo "| llama | $pp | FAILED (see /tmp/l$pp.log) | | | |" >> "$OUT"; kill -9 "$pid" 2>/dev/null; return; fi + local mib; mib=$(vram "$cvd") + local m; m=$(python3 tools/bench/pp_time.py http://127.0.0.1:8090 "$PROMPT") + local ttft tpot toks; ttft=$(echo "$m"|sed -n 's/.*TTFT_ms=\([0-9.]*\).*/\1/p'); tpot=$(echo "$m"|sed -n 's/.*TPOT_ms=\([0-9.a-z]*\).*/\1/p'); toks=$(echo "$m"|sed -n 's/.*tok_s=\([0-9.a-z]*\).*/\1/p') + echo "| llama | $pp | $ttft | $tpot | $toks | $mib |" >> "$OUT" + kill -9 "$pid" 2>/dev/null; wait "$pid" 2>/dev/null; sleep 3 +} + +for pp in $PPS; do run_xserv "$pp"; done +for pp in $PPS; do run_llama "$pp"; done +killall_servers +echo "" >> "$OUT" +echo "PP_CLEAN_DONE" >> "$OUT" diff --git a/tools/bench/pp_time.py b/tools/bench/pp_time.py new file mode 100644 index 0000000..4d0ad74 --- /dev/null +++ b/tools/bench/pp_time.py @@ -0,0 +1,44 @@ +"""Tiny single-stream latency probe over the OpenAI HTTP API. + +Usage: python3 pp_time.py BASE_URL "PROMPT" +Prints: TTFT_ms=.. TPOT_ms=.. tok_full=.. tok_s=.. + +TTFT ~ wall time of a max_tokens=1 request (prefill + 1 token). +TPOT ~ (t_full - t_1) / (tokens_full - tokens_1), using the server's reported +completion_tokens so it is exact even if generation stops early. +""" +import json +import sys +import time +import urllib.request + +base = sys.argv[1].rstrip("/") +prompt = sys.argv[2] + + +def req(max_tokens): + body = json.dumps({ + "model": "qwen3-8b", + "messages": [{"role": "user", "content": prompt}], + "max_tokens": max_tokens, + "temperature": 0, + "stream": False, + }).encode() + r = urllib.request.Request(base + "/v1/chat/completions", body, + {"Content-Type": "application/json"}) + t = time.time() + d = json.load(urllib.request.urlopen(r, timeout=600)) + dt = time.time() - t + ct = d.get("usage", {}).get("completion_tokens") + return dt, ct + + +t1, c1 = req(1) +tF, cF = req(160) +ttft = t1 * 1000.0 +denom = (cF - c1) if (cF and c1 and cF > c1) else None +if denom: + tpot = (tF - t1) / denom * 1000.0 + print(f"TTFT_ms={ttft:.1f} TPOT_ms={tpot:.2f} tok_full={cF} tok_s={1000.0/tpot:.1f}") +else: + print(f"TTFT_ms={ttft:.1f} TPOT_ms=nan tok_full={cF} tok_s=nan") diff --git a/tools/bench/run_pp_parallel.sh b/tools/bench/run_pp_parallel.sh new file mode 100644 index 0000000..1a43f44 --- /dev/null +++ b/tools/bench/run_pp_parallel.sh @@ -0,0 +1,42 @@ +#!/usr/bin/env bash +# Run the PP=1/2/4 sweep with xserv and llama.cpp CONCURRENTLY on disjoint GPU +# groups: xserv (--pp) on GPUs 0..N-1, llama.cpp (-sm layer) on GPUs 4..4+N-1. +# The 8x5090 box is grouped 0-3 / 4-7 (PHB intra-group), so each engine's P2P +# stays intra-group and the two engines never contend for a GPU. +# +# xserv splits layers across N GPUs and hands off hidden states via NCCL P2P; +# llama.cpp's default `-sm layer` does the analogous layer-wise split. +# +# Run from the repo root on the GPU host. Produces bench-out/pp{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:-20}" +MAXSEQ="${MAXSEQ:-2048}" +PPS="${PPS:-1 2 4}" +TASKS="${TASKS:-gsm8k}" + +for PP in $PPS; do + LD=$(seq -s, 4 $((3 + PP))) # llama GPUs: 4 / 4,5 / 4,5,6,7 + echo "##### PP=$PP (xserv GPU 0..$((PP-1)) || llama GPU $LD) #####" + rm -rf "bench-out/pp$PP-xserv" "bench-out/pp$PP-llama" + + python3 -u -m tools.bench.runner --systems xserv --pp "$PP" \ + --xserv-bin ./target/release/xserv-server --xserv-model "$MODEL" \ + --suite quality --quality-tasks "$TASKS" --quality-limit "$LIMIT" \ + --max-batch 1 --max-seq-len "$MAXSEQ" \ + --out-dir "bench-out/pp$PP-xserv" > "/tmp/pp$PP-xserv.log" 2>&1 & + XP=$! + + python3 -u -m tools.bench.runner --systems llama.cpp --pp "$PP" --llama-devices "$LD" \ + --llama-bin third_party/llama.cpp/build/bin/llama-server --llama-gguf "$GGUF" \ + --suite quality --quality-tasks "$TASKS" --quality-limit "$LIMIT" \ + --max-batch 1 --max-seq-len "$MAXSEQ" \ + --out-dir "bench-out/pp$PP-llama" > "/tmp/pp$PP-llama.log" 2>&1 & + LP=$! + + wait "$XP" "$LP" + echo "PP=$PP done" +done +echo ALL_DONE diff --git a/tools/bench/runner.py b/tools/bench/runner.py index f9d4380..5bc1eab 100644 --- a/tools/bench/runner.py +++ b/tools/bench/runner.py @@ -72,6 +72,9 @@ def parse_args() -> argparse.Namespace: 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("--pp", type=int, default=1, + help="Pipeline-parallel degree for BOTH engines (xserv --pp N; " + "llama.cpp --split-mode layer 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 " @@ -113,7 +116,7 @@ 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, tp=args.tp, + max_batch=args.max_batch, max_seq_len=args.max_seq_len, tp=args.tp, pp=args.pp, ), health_path="/health", ready_timeout_s=1200.0, @@ -140,10 +143,10 @@ def build_endpoints(args) -> list[SystemEndpoint]: # 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)] + devs = [d.strip() for d in args.llama_devices.split(",") if d.strip()][: max(args.tp, args.pp, 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))} + elif args.tp > 1 or args.pp > 1: + llama_env = {"CUDA_VISIBLE_DEVICES": ",".join(str(d) for d in range(max(args.tp, args.pp)))} else: llama_env = {} eps.append(SystemEndpoint( @@ -152,7 +155,7 @@ def build_endpoints(args) -> list[SystemEndpoint]: 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, tp=args.tp, + n_parallel=args.max_batch, ctx_per_slot=args.max_seq_len, tp=args.tp, pp=args.pp, ), launch_env=llama_env, # llama-server's health endpoint also returns 200 only when model is loaded. diff --git a/tools/bench/servers.py b/tools/bench/servers.py index bfdf4ee..e07d8de 100644 --- a/tools/bench/servers.py +++ b/tools/bench/servers.py @@ -114,6 +114,7 @@ def xserv_launch_cmd( max_batch: int, max_seq_len: int, tp: int = 1, + pp: int = 1, ) -> list[str]: cmd = [ bin_path, @@ -122,7 +123,9 @@ def xserv_launch_cmd( "--max-batch", str(max_batch), "--max-seq-len", str(max_seq_len), ] - if tp > 1: + if pp > 1: + cmd += ["--pp", str(pp)] # xserv binds stage s -> GPU s internally + elif tp > 1: cmd += ["--tp", str(tp)] # xserv binds rank r -> GPU r internally return cmd @@ -136,6 +139,7 @@ def llama_cpp_launch_cmd( ctx_per_slot: int, n_gpu_layers: int = 99, tp: int = 1, + pp: 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 @@ -153,7 +157,10 @@ 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: + if pp > 1: + # Pipeline / layer split across the visible GPUs (llama.cpp default). + cmd += ["--split-mode", "layer", "-ts", ",".join(["1"] * pp)] + elif 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). diff --git a/tools/bench/summarize_fullq.py b/tools/bench/summarize_fullq.py new file mode 100644 index 0000000..2c4673f --- /dev/null +++ b/tools/bench/summarize_fullq.py @@ -0,0 +1,17 @@ +"""Summarize the full quality matrix: bench-out/fullq-{xserv,llama}-pp{1,2,4}. +Prints one row per (engine, pp, task) with accuracy + latency.""" +import glob, json, os, sys +base = sys.argv[1] if len(sys.argv) > 1 else "bench-out" +print("%-6s %-3s %-9s %-8s %6s %9s %9s %10s" % + ("engine","PP","task","correct","acc%","mean_tok","TTFT_ms","TPOT_ms")) +for eng in ("xserv","llama"): + for pp in (1,2,4): + files = sorted(glob.glob(os.path.join(base, f"fullq-{eng}-pp{pp}", "comparison-*.json"))) + if not files: + print(f"{eng:<6} {pp:<3} (no results)"); continue + d = json.load(open(files[-1])) + for r in d.get("quality",{}).get("summary",[]): + print("%-6s %-3d %-9s %-8s %5.1f%% %9.0f %9.1f %10.2f" % ( + eng, pp, r["task"], f'{r["n_correct"]}/{r["n_total"]}', + r["accuracy"]*100, r.get("mean_completion_tokens",0), + r.get("mean_ttft_ms",0), r.get("mean_tpot_ms",0))) diff --git a/tools/bench/summarize_pp.py b/tools/bench/summarize_pp.py new file mode 100644 index 0000000..8f6517c --- /dev/null +++ b/tools/bench/summarize_pp.py @@ -0,0 +1,24 @@ +"""Summarize the concurrent PP sweep: bench-out/pp{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 pp in (1, 2, 4): + for sysname in ("xserv", "llama"): + files = sorted(glob.glob(os.path.join(base, f"pp{pp}-{sysname}", "comparison-*.json"))) + if not files: + continue + d = json.load(open(files[-1])) + for r in d["quality"]["summary"]: + rows.append((pp, 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" % + ("PP", "engine", "task", "correct", "acc%", "mean_tok", "TTFT_ms", "TPOT_ms", "wall_s")) +for (pp, 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" % + (pp, s, task, f"{nc}/{nt}", acc, tok, ttft, tpot, wall)) diff --git a/tools/pp_diag.sh b/tools/pp_diag.sh new file mode 100644 index 0000000..901b4d6 --- /dev/null +++ b/tools/pp_diag.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +# Diagnose pp4 divergence: run single x2 and pp4 x2, same prompt, compare all. +set -u +cd /opt/wjh/projects/xserv +export PATH=$HOME/.cargo/bin:/usr/local/cuda-12.9/bin:$PATH +MODEL=/opt/wjh/models/qwen3-8b; XBIN=./target/release/xserv-server +P='Explain what a transformer is in machine learning, in 3 sentences.' +D=bench-out/PP_DIAG.md; : > "$D" +kall(){ pkill -9 -f xserv-server 2>/dev/null; sleep 3; } +ready(){ for _ in $(seq 1 400); do [ "$(curl -s -o /dev/null -w '%{http_code}' --max-time 8 http://127.0.0.1:8090/v1/chat/completions -H 'Content-Type: application/json' -d '{"model":"qwen3-8b","messages":[{"role":"user","content":"hi"}],"max_tokens":1,"temperature":0,"stream":false}' 2>/dev/null)" = 200 ] && return 0; kill -0 $1 2>/dev/null||return 1; sleep 3; done; return 1; } +run(){ local out=$1 cvd=$2; shift 2 + kall + CUDA_VISIBLE_DEVICES=$cvd nohup $XBIN $MODEL --port 8090 --max-seq-len 2048 "$@" >/tmp/d.log 2>&1 & + local pid=$!; ready $pid || { echo "FAIL" >"$out"; kill -9 $pid 2>/dev/null; return; } + curl -s --max-time 200 http://127.0.0.1:8090/v1/chat/completions -H 'Content-Type: application/json' \ + -d "{\"model\":\"qwen3-8b\",\"messages\":[{\"role\":\"user\",\"content\":\"$P\"}],\"max_tokens\":128,\"temperature\":0,\"stream\":false}" \ + | python3 -c 'import sys,json;print(json.load(sys.stdin)["choices"][0]["message"]["content"])' > "$out" 2>/dev/null + kill -9 $pid 2>/dev/null; wait $pid 2>/dev/null; sleep 3 +} +run /tmp/s_a.txt 0 +run /tmp/s_b.txt 0 +run /tmp/p4_a.txt 0,1,2,3 --pp 4 +run /tmp/p4_b.txt 0,1,2,3 --pp 4 +echo "single_A==single_B: $(cmp -s /tmp/s_a.txt /tmp/s_b.txt && echo IDENTICAL || echo DIFFER)" | tee -a "$D" +echo "pp4_A==pp4_B: $(cmp -s /tmp/p4_a.txt /tmp/p4_b.txt && echo IDENTICAL || echo DIFFER)" | tee -a "$D" +echo "single_A==pp4_A: $(cmp -s /tmp/s_a.txt /tmp/p4_a.txt && echo IDENTICAL || echo DIFFER)" | tee -a "$D" +echo "--- first diff offset single_A vs pp4_A ---" | tee -a "$D" +cmp /tmp/s_a.txt /tmp/p4_a.txt 2>&1 | tee -a "$D" +echo "--- lengths (chars) ---" | tee -a "$D" +wc -c /tmp/s_a.txt /tmp/s_b.txt /tmp/p4_a.txt /tmp/p4_b.txt | tee -a "$D" +echo "PP_DIAG_DONE" >> "$D" diff --git a/tools/pp_final.sh b/tools/pp_final.sh new file mode 100644 index 0000000..af78fdc --- /dev/null +++ b/tools/pp_final.sh @@ -0,0 +1,72 @@ +#!/usr/bin/env bash +# Definitive PP measurement, strictly sequential, with generated text captured +# for a real correctness byte-compare. Writes bench-out/PP_FINAL.md and per-config +# text files. One server at a time; readiness gated on a real generation. +set -u +cd /opt/wjh/projects/xserv +export PATH=$HOME/.cargo/bin:/usr/local/cuda-12.9/bin:$PATH +export CUDA_HOME=/usr/local/cuda-12.9 +MODEL=/opt/wjh/models/qwen3-8b +GGUF=/opt/wjh/models/qwen3-8b/qwen3-8b-bf16.gguf +LBIN=third_party/llama.cpp/build/bin/llama-server +XBIN=./target/release/xserv-server +PROMPT='Explain what a transformer is in machine learning, in 3 sentences.' +R=bench-out/PP_FINAL.md +: > "$R" +log(){ echo "$@" >> "$R"; } + +kall(){ pkill -9 -f xserv-server 2>/dev/null; pkill -9 -f llama-server 2>/dev/null; sleep 3; } +drain(){ for _ in $(seq 1 90); do hi=0; for g in ${1//,/ }; do m=$(nvidia-smi -i "$g" --query-gpu=memory.used --format=csv,noheader,nounits); [ "${m:-0}" -gt 1500 ] && hi=1; done; [ "$hi" = 0 ] && return 0; sleep 2; done; } +# gen PORT MAXTOK -> echoes JSON; http code in $GCODE +gen(){ GCODE=$(curl -s -o /tmp/resp.json -w '%{http_code}' --max-time 300 \ + "http://127.0.0.1:$1/v1/chat/completions" -H 'Content-Type: application/json' \ + -d "{\"model\":\"qwen3-8b\",\"messages\":[{\"role\":\"user\",\"content\":\"$PROMPT\"}],\"max_tokens\":$2,\"temperature\":0,\"stream\":false}"); cat /tmp/resp.json; } +ready(){ local port=$1 pid=$2; for _ in $(seq 1 400); do + c=$(curl -s -o /dev/null -w '%{http_code}' --max-time 8 "http://127.0.0.1:$port/v1/chat/completions" -H 'Content-Type: application/json' -d '{"model":"qwen3-8b","messages":[{"role":"user","content":"hi"}],"max_tokens":1,"temperature":0,"stream":false}' 2>/dev/null) + [ "$c" = 200 ] && return 0; kill -0 "$pid" 2>/dev/null || return 1; sleep 3; done; return 1; } +snap(){ for g in ${1//,/ }; do nvidia-smi -i "$g" --query-gpu=memory.used --format=csv,noheader,nounits; done | paste -sd' '; } +# latency: TTFT from 1-tok, TPOT from 96-tok using server completion_tokens +lat(){ local port=$1 + local t0 t1 c1 cF tF + t0=$(date +%s.%N); gen "$port" 1 >/tmp/g1.json; t1=$(date +%s.%N) + c1=$(python3 -c 'import json;print(json.load(open("/tmp/g1.json"))["usage"]["completion_tokens"])' 2>/dev/null || echo 1) + local ta tb; ta=$(date +%s.%N); gen "$port" 96 >/tmp/gF.json; tb=$(date +%s.%N) + cF=$(python3 -c 'import json;print(json.load(open("/tmp/gF.json"))["usage"]["completion_tokens"])' 2>/dev/null || echo 0) + python3 -c " +ttft=($t1-$t0)*1000 +d=$cF-$c1 +print('TTFT_ms=%.1f TPOT_ms=%.2f tok_s=%.1f tokF=$cF'%(ttft,(($tb-$ta)-($t1-$t0))/d*1000 if d>0 else float('nan'),(1000.0/((($tb-$ta)-($t1-$t0))/d*1000)) if d>0 else float('nan')))" +} + +xserv(){ local pp=$1 cvd; cvd=$(seq -s, 0 $((pp-1))); kall; drain "$cvd" + local ex=""; [ "$pp" -gt 1 ] && ex="--pp $pp" + XSERV_MAX_KV_BLOCKS=160 CUDA_VISIBLE_DEVICES=$cvd nohup $XBIN $MODEL --port 8090 --max-seq-len 2048 $ex >/tmp/xf$pp.log 2>&1 & + local pid=$!; if ! ready 8090 $pid; then log "xserv pp=$pp: FAILED"; kill -9 $pid 2>/dev/null; return; fi + local mib; mib=$(snap "$cvd") + gen 8090 64 | python3 -c 'import sys,json;print(json.load(sys.stdin)["choices"][0]["message"]["content"])' > /tmp/xtext_$pp.txt 2>/dev/null + local L; L=$(lat 8090) + log "xserv pp=$pp | VRAM=$mib MiB | $L" + kill -9 $pid 2>/dev/null; wait $pid 2>/dev/null; sleep 3 +} +llama(){ local pp=$1 cvd; cvd=$(seq -s, 0 $((pp-1))); kall; drain "$cvd" + local sm; if [ "$pp" -gt 1 ]; then sm="-sm layer -ts $(printf '1%.0s,' $(seq 1 $pp)|sed 's/,$//')"; else sm="-sm none"; fi + CUDA_VISIBLE_DEVICES=$cvd nohup $LBIN -m $GGUF --port 8090 --host 127.0.0.1 -c 2048 --parallel 1 -ngl 999 $sm >/tmp/lf$pp.log 2>&1 & + local pid=$!; if ! ready 8090 $pid; then log "llama pp=$pp: FAILED"; kill -9 $pid 2>/dev/null; return; fi + local mib; mib=$(snap "$cvd"); local L; L=$(lat 8090) + log "llama pp=$pp | VRAM=$mib MiB | $L" + kill -9 $pid 2>/dev/null; wait $pid 2>/dev/null; sleep 3 +} + +log "# PP FINAL — $(date)" +for pp in 1 2 4; do xserv $pp; done +log "" +log "## correctness (xserv greedy, byte compare of generated text)" +log "single==pp2: $(cmp -s /tmp/xtext_1.txt /tmp/xtext_2.txt && echo IDENTICAL || echo DIFFER)" +log "single==pp4: $(cmp -s /tmp/xtext_1.txt /tmp/xtext_4.txt && echo IDENTICAL || echo DIFFER)" +log "single_text: $(head -c 200 /tmp/xtext_1.txt)" +log "pp2_text: $(head -c 200 /tmp/xtext_2.txt)" +log "" +for pp in 1 2 4; do llama $pp; done +kall +log "" +log "PP_FINAL_DONE" diff --git a/tools/pp_llama_47.sh b/tools/pp_llama_47.sh new file mode 100644 index 0000000..f30ef33 --- /dev/null +++ b/tools/pp_llama_47.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# llama.cpp PP=1/2/4 quality (aime2025+gsm8k, 30 each) on physical GPUs 4-7, +# parallel with the xserv matrix on 0-3. Pass --llama-devices so the runner pins +# CUDA_VISIBLE_DEVICES to 4.. (it otherwise forces 0..N-1). Distinct port + dirs. +set -u +cd /opt/wjh/projects/xserv +export PATH=$HOME/.cargo/bin:/usr/local/cuda-12.9/bin:$PATH +export CUDA_HOME=/usr/local/cuda-12.9 +GGUF=/opt/wjh/models/qwen3-8b/qwen3-8b-bf16.gguf +LBIN=third_party/llama.cpp/build/bin/llama-server +PROG=bench-out/LLAMA47_PROGRESS.md +: > "$PROG"; echo "# llama on GPU 4-7 — $(date)" >> "$PROG" +for pp in 1 2 4; do + dev=$(seq -s, 4 $((3+pp))) + out=bench-out/fullq-llama-pp$pp; rm -rf "$out" + echo "=== START llama pp=$pp dev=$dev $(date +%H:%M:%S) ===" >> "$PROG" + pkill -9 -f "llama-server.*18181" 2>/dev/null; sleep 2 + python3 -u -m tools.bench.runner --systems llama.cpp --pp "$pp" --llama-devices "$dev" \ + --llama-bin "$LBIN" --llama-gguf "$GGUF" --llama-port 18181 \ + --suite quality --quality-tasks aime2025,gsm8k --quality-limit 30 \ + --max-batch 1 --max-seq-len 4096 --out-dir "$out" >/tmp/fql-$pp.log 2>&1 + echo "=== END llama pp=$pp rc=$? $(date +%H:%M:%S) $(ls $out/comparison-*.json 2>/dev/null|wc -l) json ===" >> "$PROG" +done +pkill -9 -f "llama-server.*18181" 2>/dev/null +echo "LLAMA47_DONE" >> "$PROG" diff --git a/tools/pp_quality_full.sh b/tools/pp_quality_full.sh new file mode 100644 index 0000000..dab6ea9 --- /dev/null +++ b/tools/pp_quality_full.sh @@ -0,0 +1,54 @@ +#!/usr/bin/env bash +# FULL quality matrix, strictly sequential (one server at a time, same GPU group +# 0..N-1, no concurrency). Both engines x PP=1/2/4 x {aime2025, gsm8k}. +# Each (engine,pp) invocation runs runner.py once (it does start->both tasks->stop). +# Writes bench-out/fullq--pp/comparison-*.json ; summarized at the end. +set -u +cd /opt/wjh/projects/xserv +export PATH=$HOME/.cargo/bin:/usr/local/cuda-12.9/bin:$PATH +export CUDA_HOME=/usr/local/cuda-12.9 +MODEL=/opt/wjh/models/qwen3-8b +GGUF=/opt/wjh/models/qwen3-8b/qwen3-8b-bf16.gguf +XBIN=./target/release/xserv-server +LBIN=third_party/llama.cpp/build/bin/llama-server +AIME_LIMIT=${AIME_LIMIT:-30} +GSM_LIMIT=${GSM_LIMIT:-20} +MAXSEQ=${MAXSEQ:-4096} +PROG=bench-out/FULLQ_PROGRESS.md +: > "$PROG" +echo "# full quality matrix — $(date)" >> "$PROG" + +kall(){ pkill -9 -f xserv-server 2>/dev/null; pkill -9 -f llama-server 2>/dev/null; pkill -9 -f runner.py 2>/dev/null; sleep 4; } +drain(){ for _ in $(seq 1 90); do hi=0; for g in $(seq 0 $1); do m=$(nvidia-smi -i "$g" --query-gpu=memory.used --format=csv,noheader,nounits); [ "${m:-0}" -gt 1500 ] && hi=1; done; [ "$hi" = 0 ] && return 0; sleep 2; done; } + +run_one(){ # $1 engine $2 pp + local eng=$1 pp=$2 dev; dev=$(seq -s, 0 $((pp-1))) + kall; drain $((pp-1)) + local out=bench-out/fullq-$eng-pp$pp + rm -rf "$out" + echo "=== START $eng pp=$pp on GPU $dev $(date +%H:%M:%S) ===" >> "$PROG" + if [ "$eng" = xserv ]; then + python3 -u -m tools.bench.runner --systems xserv --pp "$pp" \ + --xserv-bin "$XBIN" --xserv-model "$MODEL" \ + --suite quality --quality-tasks aime2025,gsm8k --quality-limit 30 \ + --max-batch 1 --max-seq-len "$MAXSEQ" \ + --out-dir "$out" >/tmp/fq-$eng-$pp.log 2>&1 + else + python3 -u -m tools.bench.runner --systems llama.cpp --pp "$pp" \ + --llama-bin "$LBIN" --llama-gguf "$GGUF" \ + --suite quality --quality-tasks aime2025,gsm8k --quality-limit 30 \ + --max-batch 1 --max-seq-len "$MAXSEQ" \ + --out-dir "$out" >/tmp/fq-$eng-$pp.log 2>&1 + fi + echo "=== END $eng pp=$pp rc=$? $(date +%H:%M:%S) $(ls $out/comparison-*.json 2>/dev/null | wc -l) json ===" >> "$PROG" +} + +# aime2025 has 30 problems; runner uses one --quality-limit for ALL tasks, so we +# pass max(limits) and rely on the datasets' own sizes (gsm8k.json may be larger, +# but we cap with --quality-limit). To keep gsm8k at 20 and aime at 30 we run the +# matrix with --quality-limit 30 (aime full; gsm8k uses first 30 -> report shows n_total). +for eng in xserv llama; do + for pp in 1 2 4; do run_one "$eng" "$pp"; done +done +kall +echo "FULLQ_DONE" >> "$PROG" diff --git a/tools/pp_verify.sh b/tools/pp_verify.sh new file mode 100644 index 0000000..bd4c9d0 --- /dev/null +++ b/tools/pp_verify.sh @@ -0,0 +1,90 @@ +#!/usr/bin/env bash +# One-shot pipeline-parallel (PP) verification + benchmark for Qwen3-8B. +# Run on the GPU host from the repo root. Writes bench-out/PP_RESULTS.md. +# +# 1. NCCL P2P send/recv + AllReduce unit tests +# 2. correctness: greedy (temp=0) output single == --pp 2 == --pp 4 (byte compare) +# 3. per-GPU VRAM (health-gated; weights + a minimal KV pool, ~1/P per card) +# 4. quality+latency sweep vs llama.cpp (-sm layer), gsm8k +# +# Env: MODEL, GGUF, LIMIT (problems), PPS (e.g. "1 2 4") may be overridden. +set -u +cd "$(dirname "$0")/.." +export PATH=$HOME/.cargo/bin:/usr/local/cuda-12.9/bin:$PATH +export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda-12.9} +MODEL=${MODEL:-/opt/wjh/models/qwen3-8b} +GGUF=${GGUF:-/opt/wjh/models/qwen3-8b/qwen3-8b-bf16.gguf} +LIMIT=${LIMIT:-20} +PPS=${PPS:-1 2 4} +BIN=./target/release/xserv-server +R=bench-out/PP_RESULTS.md +mkdir -p bench-out +: > "$R" +log(){ echo "$@" | tee -a "$R"; } + +pkill -9 -f xserv-server 2>/dev/null; pkill -9 -f llama-server 2>/dev/null; sleep 3 + +log "# PP verification — $(date)" + +# ---- 1. NCCL P2P + AllReduce unit tests ---- +log ""; log "## 1. NCCL P2P + AllReduce test" +cargo test -p xserv-distributed --release -- --test-threads=1 >/tmp/pp_t.log 2>&1 +log " cargo test exit=$?" +grep -hE "test result|pp_send_recv|allreduce_two_gpu" /tmp/pp_t.log | sed 's/^/ /' | tee -a "$R" + +# wait_ready PORT PID -> 0 when a real generation succeeds (xserv's /health +# returns 200 before the model is loaded, so gate on a generation, not /health). +wait_ready(){ local port=$1 pid=$2 + for _ in $(seq 1 400); do + curl -s -o /dev/null -w '%{http_code}' --max-time 8 \ + "http://127.0.0.1:$port/v1/chat/completions" -H 'Content-Type: application/json' \ + -d '{"model":"qwen3-8b","messages":[{"role":"user","content":"hi"}],"max_tokens":1,"temperature":0,"stream":false}' \ + 2>/dev/null | grep -q 200 && return 0 + kill -0 "$pid" 2>/dev/null || return 1 + sleep 3 + done; return 1 +} + +# ---- 2. correctness ---- +PROMPT='Explain what a transformer is in machine learning, in 3 sentences.' +gen(){ local port=$1 cvd=$2; shift 2 + CUDA_VISIBLE_DEVICES=$cvd nohup $BIN $MODEL --port $port --max-seq-len 2048 "$@" >/tmp/pp_s$port.log 2>&1 & + local pid=$! + wait_ready "$port" "$pid" || { echo "(server $port failed)"; kill -9 "$pid" 2>/dev/null; return; } + curl -s --max-time 200 "http://127.0.0.1:$port/v1/chat/completions" -H 'Content-Type: application/json' \ + -d "{\"model\":\"qwen3-8b\",\"messages\":[{\"role\":\"user\",\"content\":\"$PROMPT\"}],\"max_tokens\":64,\"temperature\":0,\"stream\":false}" \ + | python3 -c 'import sys,json;print(json.load(sys.stdin)["choices"][0]["message"]["content"])' 2>/dev/null + kill -9 "$pid" 2>/dev/null; wait "$pid" 2>/dev/null; sleep 3 +} +gen 8091 0 > /tmp/o_single.txt +gen 8092 0,1 --pp 2 > /tmp/o_pp2.txt +gen 8093 0,1,2,3 --pp 4 > /tmp/o_pp4.txt +log ""; log "## 2. Correctness (greedy temp=0, byte compare)" +log " single==pp2: $(cmp -s /tmp/o_single.txt /tmp/o_pp2.txt && echo IDENTICAL || echo DIFFER)" +log " single==pp4: $(cmp -s /tmp/o_single.txt /tmp/o_pp4.txt && echo IDENTICAL || echo DIFFER)" +log " single text: $(head -c 160 /tmp/o_single.txt)" + +# ---- 3. per-GPU VRAM (health-gated, KV pool capped so all configs comparable) ---- +log ""; log "## 3. Per-GPU VRAM (XSERV_MAX_KV_BLOCKS=160; weights + minimal KV)" +snap(){ nvidia-smi -i "$1" --query-gpu=memory.used --format=csv,noheader,nounits | paste -sd' '; } +vram(){ local label=$1 cvd=$2 port=$3; shift 3 + XSERV_MAX_KV_BLOCKS=160 CUDA_VISIBLE_DEVICES=$cvd nohup $BIN $MODEL --port $port --max-seq-len 2048 "$@" >/tmp/pp_v$port.log 2>&1 & + local pid=$! + wait_ready "$port" "$pid" || { log " $label: server failed"; kill -9 "$pid" 2>/dev/null; return; } + curl -s --max-time 120 "http://127.0.0.1:$port/v1/chat/completions" -H 'Content-Type: application/json' \ + -d '{"model":"qwen3-8b","messages":[{"role":"user","content":"hi"}],"max_tokens":8,"temperature":0,"stream":false}' >/dev/null + local a b=""; for _ in $(seq 1 12); do a=$(snap "$cvd"); [ "$a" = "$b" ] && break; b=$a; sleep 2; done + log " $label ($cvd): $a MiB" + kill -9 "$pid" 2>/dev/null; wait "$pid" 2>/dev/null; sleep 5 +} +vram single 0 8094 +vram pp2 0,1 8095 --pp 2 +vram pp4 0,1,2,3 8096 --pp 4 + +# ---- 4. sweep vs llama.cpp ---- +log ""; log "## 4. Sweep (gsm8k $LIMIT, xserv --pp 0..N-1 vs llama -sm layer 4..)" +PPS="$PPS" LIMIT="$LIMIT" TASKS=gsm8k bash tools/bench/run_pp_parallel.sh >/tmp/pp_sweep.log 2>&1 +log '```' +python3 tools/bench/summarize_pp.py bench-out >> "$R" 2>&1 +log '```' +log ""; log "PP_VERIFY_DONE"