diff --git a/crates/xserv-model/src/qwen3.rs b/crates/xserv-model/src/qwen3.rs index c012dcd..1900c64 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 { @@ -113,9 +118,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,