model: tensor-parallel Qwen3 (sharded weights + AllReduce)

from_weights_tp shards each rank's weights (column-split q/k/v/gate/up,
row-split o/down; replicate norms/embed/lm_head) and the paged forward uses
local head counts + AllReduces after o_proj and down_proj. PagedKVCache::new_tp
sizes the pool for the rank's local KV heads (KV is sharded too). TP=1 is the
identity path. New bench-tp binary runs E2E multi-GPU generation per TP degree.

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

crates/xserv-model/src/qwen3.rs

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