Merge branch 'phase18-pipeline-parallelism': pipeline-parallel inference

Adds --pp N for layer-wise pipeline parallelism via NCCL P2P send/recv.
Each stage holds layers [s*L, (s+1)*L), stage 0 owns embedding, last
stage owns norm/lm_head. v1 serial (one request at a time) — correctness
+ per-GPU memory savings (~1/N). Refactors model to unfused QKV/gate_up
projections and removes unused kernels (argmax, reshape_and_cache).

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;

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) };
+        }
+    }
+}

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<bf16> = (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");
+}

crates/xserv-model/src/qwen3.rs

@@ -20,6 +20,11 @@ pub struct Qwen3 {
     tp: Option<std::sync::Arc<xserv_distributed::TpContext>>,
     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<String, Tensor>,
+        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<String, Tensor>, 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<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
+
+        for (layer_idx, layer) in self.layers.iter().enumerate() {
+            let residual = x.clone();
+            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<usize> = seq_slots.iter().map(|&s| paged_cache.seq_len(s)).collect();
+        let kv_lens: Vec<i32> = 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<Tensor> = 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]

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,

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<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] [--tp N]");
+        eprintln!("Usage: xserv-server <model-dir> [--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 {

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<PpContext>,
+    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<PpCommand>,
+    ack_tx: mpsc::Sender<()>,
+    token_tx: mpsc::Sender<u32>,
+) {
+    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<GenerateRequest>) {
+    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::<u32>();
+    let mut cmd_txs: Vec<mpsc::Sender<PpCommand>> = Vec::new();
+    for stage in 1..world {
+        let (ctx_tx, ctx_rx) = mpsc::channel::<PpCommand>();
+        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<PpCommand>], 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<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";
+            }
+            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<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 });
+    }
+}