gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

style: format Rust workspace

Browse Source
This commit is contained in:
Gahow Wang
2026-06-18 18:11:58 +08:00
parent 013465fc06
commit 531cd3fe08
57 changed files with 4045 additions and 1204 deletions
Download Patch File Download Diff File Expand all files Collapse all files

295
crates/xserv-model/src/qwen3.rs
View File

@@ -1,5 +1,5 @@
use std::collections::HashMap;
use half::bf16;
use std::collections::HashMap;
use xserv_kernels::*;
use xserv_tensor::{Device, Tensor};
@@ -13,7 +13,7 @@ pub struct Qwen3 {
embed_tokens: Tensor,
layers: Vec<Qwen3Block>,
norm: Tensor,
lm_head_t: Tensor, // precomputed transpose
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.
@@ -28,22 +28,29 @@ pub struct Qwen3 {
}
struct Qwen3Block {
input_norm: Tensor, // [hidden]
input_norm: Tensor, // [hidden]
qkv_proj_wt: Tensor, // FUSED: [hidden, (H+2*KV)*D] — Q|K|V columns
q_dim: usize, // num_heads * head_dim (Q slice boundary)
kv_dim: usize, // num_kv_heads * head_dim (K/V slice size)
o_proj_wt: Tensor, // TRANSPOSED: [num_heads*head_dim, hidden]
q_norm: Tensor, // [head_dim]
k_norm: Tensor, // [head_dim]
post_norm: Tensor, // [hidden]
gate_up_proj_wt: Tensor, // FUSED: [hidden, 2*intermediate]
down_proj_wt: Tensor, // TRANSPOSED: [intermediate, hidden]
o_proj_wt: Tensor, // TRANSPOSED: [num_heads*head_dim, hidden]
q_norm: Tensor, // [head_dim]
k_norm: Tensor, // [head_dim]
post_norm: Tensor, // [hidden]
gate_up_proj_wt: Tensor, // FUSED: [hidden, 2*intermediate]
down_proj_wt: Tensor, // TRANSPOSED: [intermediate, hidden]
}
impl Qwen3Block {
fn q_proj_wt(&self) -> Tensor { self.qkv_proj_wt.narrow(1, 0, self.q_dim) }
fn k_proj_wt(&self) -> Tensor { self.qkv_proj_wt.narrow(1, self.q_dim, self.kv_dim) }
fn v_proj_wt(&self) -> Tensor { self.qkv_proj_wt.narrow(1, self.q_dim + self.kv_dim, self.kv_dim) }
fn q_proj_wt(&self) -> Tensor {
self.qkv_proj_wt.narrow(1, 0, self.q_dim)
}
fn k_proj_wt(&self) -> Tensor {
self.qkv_proj_wt.narrow(1, self.q_dim, self.kv_dim)
}
fn v_proj_wt(&self) -> Tensor {
self.qkv_proj_wt
.narrow(1, self.q_dim + self.kv_dim, self.kv_dim)
}
fn gate_proj_wt(&self) -> Tensor {
let half = self.gate_up_proj_wt.shape()[1] / 2;
self.gate_up_proj_wt.narrow(1, 0, half)
@@ -80,18 +87,31 @@ impl Qwen3 {
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}"))
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() };
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 row = |t: Tensor| -> Tensor {
shard_cols(&t, rank, world)
.to_device(dev)
.transpose(0, 1)
.contiguous()
};
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 lm_head_t = repl(take(&mut w, "lm_head.weight"))
.transpose(0, 1)
.contiguous();
let rope_cache = RopeCache::new(
config.max_seq_len(),
@@ -102,7 +122,10 @@ impl Qwen3 {
let num_layers = config.num_layers();
let mut layers = Vec::with_capacity(num_layers);
if rank == 0 {
eprintln!("Loading+sharding weights for {} layers (world={world})...", num_layers);
eprintln!(
"Loading+sharding weights for {} layers (world={world})...",
num_layers
);
}
for i in 0..num_layers {
let p = format!("model.layers.{i}");
@@ -126,7 +149,10 @@ impl Qwen3 {
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"))),
post_norm: repl(take(
&mut w,
&format!("{p}.post_attention_layernorm.weight"),
)),
gate_up_proj_wt,
down_proj_wt: row(take(&mut w, &format!("{p}.mlp.down_proj.weight"))),
});
@@ -165,7 +191,10 @@ impl Qwen3 {
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}");
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;
@@ -173,16 +202,29 @@ impl Qwen3 {
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}"))
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 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(),
@@ -217,7 +259,10 @@ impl Qwen3 {
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"))),
post_norm: repl(take(
&mut w,
&format!("{p}.post_attention_layernorm.weight"),
)),
gate_up_proj_wt,
down_proj_wt: wt(take(&mut w, &format!("{p}.mlp.down_proj.weight"))),
});
@@ -252,8 +297,12 @@ impl Qwen3 {
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 }
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
@@ -276,7 +325,9 @@ impl Qwen3 {
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();
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();
@@ -285,7 +336,9 @@ impl Qwen3 {
let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
let q = qkv.narrow(1, 0, layer.q_dim).contiguous();
let k = qkv.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
let v = qkv.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim).contiguous();
let v = qkv
.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
.contiguous();
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);
@@ -305,10 +358,12 @@ impl Qwen3 {
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_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 (normed, x_new) =
xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let residual = x_new.clone();
let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
@@ -356,7 +411,9 @@ impl Qwen3 {
let qkv_all = matmul_2d(&normed, &layer.qkv_proj_wt);
let q_all = qkv_all.narrow(1, 0, layer.q_dim).contiguous();
let k_all = qkv_all.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
let v_all = qkv_all.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim).contiguous();
let v_all = qkv_all
.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
.contiguous();
let mut q_rows: Vec<Tensor> = Vec::with_capacity(batch);
for b in 0..batch {
@@ -394,14 +451,23 @@ impl Qwen3 {
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,
&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 (normed, x_new) =
xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let residual = x_new.clone();
let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
@@ -441,7 +507,9 @@ impl Qwen3 {
let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
let mut x = embedding(&self.embed_tokens, token_ids);
let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
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();
@@ -450,7 +518,9 @@ impl Qwen3 {
let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
let q = qkv.narrow(1, 0, layer.q_dim).contiguous();
let k = qkv.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
let v = qkv.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim).contiguous();
let v = qkv
.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
.contiguous();
let q = reshape_heads(&q, new_tokens, num_heads, head_dim);
let k = reshape_heads(&k, new_tokens, num_kv_heads, head_dim);
@@ -531,7 +601,9 @@ impl Qwen3 {
let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
let q_all = qkv.narrow(1, 0, layer.q_dim).contiguous();
let k_all = qkv.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
let v_all = qkv.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim).contiguous();
let v_all = qkv
.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
.contiguous();
// Per-sequence: reshape, qk-norm, RoPE, KV cache, attention, merge
let mut attn_outputs: Vec<Tensor> = Vec::with_capacity(batch);
@@ -583,7 +655,8 @@ impl Qwen3 {
let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
// Fused add + rmsnorm
let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let (normed, x_new) =
xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let residual = x_new.clone();
let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
@@ -662,13 +735,15 @@ impl Qwen3 {
let qkv = matmul_2d(&normed, &layer.qkv_proj_wt); // [B, (H+2*KV)*D]
let q_dim = num_heads * head_dim;
let kv_dim = num_kv_heads * head_dim;
let q_all = qkv.narrow(1, 0, q_dim); // [B, H*D] (view)
let k_all = qkv.narrow(1, q_dim, kv_dim); // [B, KV*D] (view)
let q_all = qkv.narrow(1, 0, q_dim); // [B, H*D] (view)
let k_all = qkv.narrow(1, q_dim, kv_dim); // [B, KV*D] (view)
let v_all = qkv.narrow(1, q_dim + kv_dim, kv_dim);
// Per-head RMSNorm on contiguous copies (narrow views are strided).
let q_flat = q_all.contiguous().reshape(&[batch * num_heads, head_dim]);
let k_flat = k_all.contiguous().reshape(&[batch * num_kv_heads, head_dim]);
let k_flat = k_all
.contiguous()
.reshape(&[batch * num_kv_heads, head_dim]);
let q_normed = rmsnorm(&q_flat, &layer.q_norm, eps);
let k_normed = rmsnorm(&k_flat, &layer.k_norm, eps);
@@ -688,8 +763,16 @@ impl Qwen3 {
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,
&q_4d,
k_pool_ptr,
v_pool_ptr,
bt_ptr,
cl_ptr,
batch,
num_heads,
num_kv_heads,
head_dim,
max_blocks,
);
// attn_out shape [B, H, 1, D] is contiguous-equivalent to [B, H*D].
@@ -697,7 +780,8 @@ impl Qwen3 {
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 (normed, x_new) =
xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let residual = x_new.clone();
// Fused gate+up projection: one GEMV instead of two.
@@ -743,7 +827,9 @@ impl Qwen3 {
paged_cache.advance_seq_len(slot, new_tokens);
let mut x = embedding(&self.embed_tokens, token_ids);
let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
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();
@@ -752,7 +838,9 @@ impl Qwen3 {
let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
let q = qkv.narrow(1, 0, layer.q_dim).contiguous();
let k = qkv.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
let v = qkv.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim).contiguous();
let v = qkv
.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
.contiguous();
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);
@@ -773,11 +861,13 @@ impl Qwen3 {
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_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);
let (normed, x_new) = xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let (normed, x_new) =
xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let residual = x_new.clone();
let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
@@ -805,7 +895,9 @@ impl Qwen3 {
let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
let mut x = embedding(&self.embed_tokens, token_ids);
let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
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();
@@ -814,7 +906,9 @@ impl Qwen3 {
let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
let q = qkv.narrow(1, 0, layer.q_dim).contiguous();
let k = qkv.narrow(1, layer.q_dim, layer.kv_dim).contiguous();
let v = qkv.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim).contiguous();
let v = qkv
.narrow(1, layer.q_dim + layer.kv_dim, layer.kv_dim)
.contiguous();
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);
@@ -834,10 +928,12 @@ impl Qwen3 {
let (k_full, v_full) = cache.get_kv_len(layer_idx, pos_offset + new_tokens);
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_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 (normed, x_new) =
xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let residual = x_new.clone();
let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
@@ -856,28 +952,33 @@ impl Qwen3 {
/// Extract weight pointers for CUDA Graph capture.
pub fn layer_weight_ptrs(&self) -> Vec<crate::decode_graph::LayerWeightPtrs> {
self.layers.iter().map(|l| crate::decode_graph::LayerWeightPtrs {
input_norm: l.input_norm.data_ptr() as *const std::ffi::c_void,
q_proj_wt: l.q_proj_wt().data_ptr() as *const std::ffi::c_void,
k_proj_wt: l.k_proj_wt().data_ptr() as *const std::ffi::c_void,
v_proj_wt: l.v_proj_wt().data_ptr() as *const std::ffi::c_void,
o_proj_wt: l.o_proj_wt.data_ptr() as *const std::ffi::c_void,
q_norm: l.q_norm.data_ptr() as *const std::ffi::c_void,
k_norm: l.k_norm.data_ptr() as *const std::ffi::c_void,
post_norm: l.post_norm.data_ptr() as *const std::ffi::c_void,
gate_proj_wt: l.gate_proj_wt().data_ptr() as *const std::ffi::c_void,
up_proj_wt: l.up_proj_wt().data_ptr() as *const std::ffi::c_void,
down_proj_wt: l.down_proj_wt.data_ptr() as *const std::ffi::c_void,
}).collect()
self.layers
.iter()
.map(|l| crate::decode_graph::LayerWeightPtrs {
input_norm: l.input_norm.data_ptr() as *const std::ffi::c_void,
q_proj_wt: l.q_proj_wt().data_ptr() as *const std::ffi::c_void,
k_proj_wt: l.k_proj_wt().data_ptr() as *const std::ffi::c_void,
v_proj_wt: l.v_proj_wt().data_ptr() as *const std::ffi::c_void,
o_proj_wt: l.o_proj_wt.data_ptr() as *const std::ffi::c_void,
q_norm: l.q_norm.data_ptr() as *const std::ffi::c_void,
k_norm: l.k_norm.data_ptr() as *const std::ffi::c_void,
post_norm: l.post_norm.data_ptr() as *const std::ffi::c_void,
gate_proj_wt: l.gate_proj_wt().data_ptr() as *const std::ffi::c_void,
up_proj_wt: l.up_proj_wt().data_ptr() as *const std::ffi::c_void,
down_proj_wt: l.down_proj_wt.data_ptr() as *const std::ffi::c_void,
})
.collect()
}
/// Get pointers needed for CUDA Graph capture.
pub fn graph_capture_ptrs(&self) -> (
*const std::ffi::c_void, // norm weight
*const std::ffi::c_void, // lm_head_t
*const std::ffi::c_void, // embed_tokens
*const std::ffi::c_void, // rope cos
*const std::ffi::c_void, // rope sin
pub fn graph_capture_ptrs(
&self,
) -> (
*const std::ffi::c_void, // norm weight
*const std::ffi::c_void, // lm_head_t
*const std::ffi::c_void, // embed_tokens
*const std::ffi::c_void, // rope cos
*const std::ffi::c_void, // rope sin
) {
(
self.norm.data_ptr() as *const std::ffi::c_void,
@@ -895,11 +996,16 @@ impl Qwen3 {
/// (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(); }
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}");
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>();
@@ -911,11 +1017,16 @@ fn shard_rows(t: &Tensor, rank: usize, world: usize) -> Tensor {
/// 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(); }
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}");
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);
@@ -1009,7 +1120,9 @@ fn transpose_from_rope(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: u
}
fn repeat_kv(x: &Tensor, n_rep: usize) -> Tensor {
if n_rep == 1 { return x.clone(); }
if n_rep == 1 {
return x.clone();
}
let kv_heads = x.shape()[1];
let seq_len = x.shape()[2];
let head_dim = x.shape()[3];
@@ -1065,11 +1178,16 @@ fn concat_rows(rows: &[Tensor]) -> Tensor {
let src_buf = row.storage().gpu_buffer();
let src_offset = row.offset() * elem_size;
let dst_offset = b * row_bytes;
out_buf.copy_from_device_at(src_buf, src_offset, dst_offset, row_bytes).unwrap();
out_buf
.copy_from_device_at(src_buf, src_offset, dst_offset, row_bytes)
.unwrap();
}
// Wrap in a Tensor
let device_id = match device { Device::Cuda(id) => id, _ => panic!("expected CUDA device") };
let device_id = match device {
Device::Cuda(id) => id,
_ => panic!("expected CUDA device"),
};
unsafe {
crate::kv_cache::tensor_from_gpu_buffer_pub(out_buf, &[batch, cols], dtype, device_id)
}
@@ -1082,12 +1200,15 @@ fn cat_cols(tensors: &[&Tensor]) -> Tensor {
let dtype = tensors[0].dtype();
let device = tensors[0].device();
let elem = dtype.size_bytes();
let total_cols: usize = tensors.iter().map(|t| {
assert_eq!(t.ndim(), 2);
assert_eq!(t.shape()[0], rows);
assert!(t.is_contiguous());
t.shape()[1]
}).sum();
let total_cols: usize = tensors
.iter()
.map(|t| {
assert_eq!(t.ndim(), 2);
assert_eq!(t.shape()[0], rows);
assert!(t.is_contiguous());
t.shape()[1]
})
.sum();
let out = Tensor::empty(&[rows, total_cols], dtype, device);
let dst_base = out.data_ptr() as *mut u8;
for r in 0..rows {
@@ -1126,7 +1247,9 @@ pub fn sample_greedy(logits: &Tensor) -> u32 {
let seq_len = logits.shape()[0];
let data = logits_cpu.as_slice::<bf16>();
let last = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
last.iter().enumerate()
last.iter()
.enumerate()
.max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
.map(|(i, _)| i as u32).unwrap()
.map(|(i, _)| i as u32)
.unwrap()
}