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

359
crates/xserv-model/src/gpt_oss.rs
View File

@@ -1,6 +1,6 @@
use half::bf16;
use std::collections::HashMap;
use std::ffi::c_void;
use half::bf16;
use xserv_kernels::*;
use xserv_tensor::{Device, Tensor};
@@ -49,10 +49,10 @@ struct GptOssBlock {
expert_down_bias: Tensor, // [local_experts, hidden]
// FP8 quantized expert weights (Some when running FP8 W8A8)
// Transposed layout [E, N, K] for cuBLASLt FP8 (Blackwell requires transA=T)
expert_gate_up_fp8: Option<Tensor>, // [local_experts, 2*inter, hidden] FP8E4M3
expert_gate_up_scale: Option<Tensor>,// [local_experts] F32
expert_down_fp8: Option<Tensor>, // [local_experts, hidden, inter] FP8E4M3
expert_down_scale: Option<Tensor>, // [local_experts] F32
expert_gate_up_fp8: Option<Tensor>, // [local_experts, 2*inter, hidden] FP8E4M3
expert_gate_up_scale: Option<Tensor>, // [local_experts] F32
expert_down_fp8: Option<Tensor>, // [local_experts, hidden, inter] FP8E4M3
expert_down_scale: Option<Tensor>, // [local_experts] F32
// MXFP4 W4A16 expert weights (Some when running 4-bit weight-only).
// (packed [E, N, K/2] u8, scales [E, N, K/32] u8) in [E, N, K] layout.
expert_gate_up_mxfp4: Option<(Tensor, Tensor)>,
@@ -79,16 +79,23 @@ impl GptOss {
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}"))
};
let repl = |t: Tensor| -> Tensor { t.to_device(dev) };
// column-parallel: shard rows of [out, in], transpose → [in, out/world]
let col = |t: Tensor| -> Tensor {
shard_rows(&t, rank, world).to_device(dev).transpose(0, 1).contiguous()
shard_rows(&t, rank, world)
.to_device(dev)
.transpose(0, 1)
.contiguous()
};
// row-parallel: shard cols of [out, in], transpose → [in/world, out]
let row = |t: Tensor| -> Tensor {
shard_cols(&t, rank, world).to_device(dev).transpose(0, 1).contiguous()
shard_cols(&t, rank, world)
.to_device(dev)
.transpose(0, 1)
.contiguous()
};
// Bias sharding helpers
let col_bias = |t: Tensor| -> Tensor { shard_1d(&t, rank, world).to_device(dev) };
@@ -97,7 +104,9 @@ impl GptOss {
let embed_tokens = repl(take(&mut w, "model.embed_tokens.weight"));
let norm = repl(take(&mut w, "model.norm.weight"));
let norm_bias = w.remove("model.norm.bias").map(|t| repl(t));
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 head_dim = config.head_dim();
let rope_theta = config.rope_theta.unwrap_or(150000.0);
@@ -176,15 +185,30 @@ impl GptOss {
// MXFP4 stores 4-bit weights in an FP8E4M3 byte container (same dtype
// as FP8), so distinguish by the scale rank: FP8 scale is 1-D [E],
// MXFP4 scale is 3-D [E, N, K/32].
let is_mxfp4 = gate_up_scale.as_ref().map(|s| s.ndim() == 3).unwrap_or(false);
let is_mxfp4 = gate_up_scale
.as_ref()
.map(|s| s.ndim() == 3)
.unwrap_or(false);
let is_fp8 = !is_mxfp4 && gate_up_3d.dtype() == xserv_tensor::DType::FP8E4M3;
let mut expert_gate_up_mxfp4: Option<(Tensor, Tensor)> = None;
let mut expert_down_mxfp4: Option<(Tensor, Tensor)> = None;
let inter2 = if is_mxfp4 { gate_up_3d.shape()[1] } else { gate_up_3d.shape()[2] }; // 2*inter (N)
let hidden = if is_mxfp4 { gate_up_3d.shape()[2] * 2 } else { gate_up_3d.shape()[1] };
let inter = if is_mxfp4 { down_3d.shape()[2] * 2 } else { down_3d.shape()[1] };
let inter2 = if is_mxfp4 {
gate_up_3d.shape()[1]
} else {
gate_up_3d.shape()[2]
}; // 2*inter (N)
let hidden = if is_mxfp4 {
gate_up_3d.shape()[2] * 2
} else {
gate_up_3d.shape()[1]
};
let inter = if is_mxfp4 {
down_3d.shape()[2] * 2
} else {
down_3d.shape()[1]
};
// Slice the rank's range of experts as contiguous 3D tensors on GPU
let expert_gate_up_wt;
@@ -199,10 +223,38 @@ impl GptOss {
// + scales [E, N, K/32]. Slice this rank's experts (raw bytes).
let gu_s = gate_up_scale.expect("MXFP4 model missing gate_up_proj_scale");
let d_s = down_scale.expect("MXFP4 model missing down_proj_scale");
let gu_packed = slice_expert_range_3d_raw(&gate_up_3d, expert_start, local_experts, inter2, hidden / 2).to_device(dev);
let gu_scl = slice_expert_range_3d_raw(&gu_s, expert_start, local_experts, inter2, hidden / 32).to_device(dev);
let dn_packed = slice_expert_range_3d_raw(&down_3d, expert_start, local_experts, hidden, inter / 2).to_device(dev);
let dn_scl = slice_expert_range_3d_raw(&d_s, expert_start, local_experts, hidden, inter / 32).to_device(dev);
let gu_packed = slice_expert_range_3d_raw(
&gate_up_3d,
expert_start,
local_experts,
inter2,
hidden / 2,
)
.to_device(dev);
let gu_scl = slice_expert_range_3d_raw(
&gu_s,
expert_start,
local_experts,
inter2,
hidden / 32,
)
.to_device(dev);
let dn_packed = slice_expert_range_3d_raw(
&down_3d,
expert_start,
local_experts,
hidden,
inter / 2,
)
.to_device(dev);
let dn_scl = slice_expert_range_3d_raw(
&d_s,
expert_start,
local_experts,
hidden,
inter / 32,
)
.to_device(dev);
expert_gate_up_mxfp4 = Some((gu_packed, gu_scl));
expert_down_mxfp4 = Some((dn_packed, dn_scl));
expert_gate_up_fp8 = None;
@@ -214,36 +266,65 @@ impl GptOss {
} else if is_fp8 {
// FP8 W8A8 path: load and TRANSPOSE weights for cuBLASLt (requires transA=T on Blackwell).
// Original: [E, K, N] → Transposed: [E, N, K]
let gu_sliced = slice_expert_range_3d_raw(&gate_up_3d, expert_start, local_experts, hidden, inter2);
let dn_sliced = slice_expert_range_3d_raw(&down_3d, expert_start, local_experts, inter, hidden);
expert_gate_up_fp8 = Some(transpose_3d_inner_raw(&gu_sliced, local_experts, hidden, inter2).to_device(dev));
expert_down_fp8 = Some(transpose_3d_inner_raw(&dn_sliced, local_experts, inter, hidden).to_device(dev));
let gu_sliced = slice_expert_range_3d_raw(
&gate_up_3d,
expert_start,
local_experts,
hidden,
inter2,
);
let dn_sliced =
slice_expert_range_3d_raw(&down_3d, expert_start, local_experts, inter, hidden);
expert_gate_up_fp8 = Some(
transpose_3d_inner_raw(&gu_sliced, local_experts, hidden, inter2)
.to_device(dev),
);
expert_down_fp8 = Some(
transpose_3d_inner_raw(&dn_sliced, local_experts, inter, hidden).to_device(dev),
);
// Scales: [num_experts] F32 → slice to [local_experts]
let gu_s = gate_up_scale.expect("FP8 model missing gate_up_proj_scale");
let d_s = down_scale.expect("FP8 model missing down_proj_scale");
expert_gate_up_scale_gpu = Some(slice_scale_range(&gu_s, expert_start, local_experts).to_device(dev));
expert_down_scale_gpu = Some(slice_scale_range(&d_s, expert_start, local_experts).to_device(dev));
expert_gate_up_scale_gpu =
Some(slice_scale_range(&gu_s, expert_start, local_experts).to_device(dev));
expert_down_scale_gpu =
Some(slice_scale_range(&d_s, expert_start, local_experts).to_device(dev));
// Dummy BF16 tensors (never read in FP8 path)
expert_gate_up_wt = Tensor::empty(&[1, 1, 1], xserv_tensor::DType::BF16, dev);
expert_down_wt = Tensor::empty(&[1, 1, 1], xserv_tensor::DType::BF16, dev);
} else {
// BF16 path: existing behavior
expert_gate_up_wt = slice_expert_range_3d(&gate_up_3d, expert_start, local_experts, hidden, inter2).to_device(dev);
expert_down_wt = slice_expert_range_3d(&down_3d, expert_start, local_experts, inter, hidden).to_device(dev);
expert_gate_up_wt =
slice_expert_range_3d(&gate_up_3d, expert_start, local_experts, hidden, inter2)
.to_device(dev);
expert_down_wt =
slice_expert_range_3d(&down_3d, expert_start, local_experts, inter, hidden)
.to_device(dev);
expert_gate_up_fp8 = None;
expert_gate_up_scale_gpu = None;
expert_down_fp8 = None;
expert_down_scale_gpu = None;
}
let expert_gate_up_bias = slice_expert_range_2d(&gate_up_bias_2d, expert_start, local_experts, inter2).to_device(dev);
let expert_down_bias = slice_expert_range_2d(&down_bias_2d, expert_start, local_experts, hidden).to_device(dev);
let expert_gate_up_bias =
slice_expert_range_2d(&gate_up_bias_2d, expert_start, local_experts, inter2)
.to_device(dev);
let expert_down_bias =
slice_expert_range_2d(&down_bias_2d, expert_start, local_experts, hidden)
.to_device(dev);
xserv_cuda::allocator::cached_trim();
let input_norm = repl(take(&mut w, &format!("{p}.input_layernorm.weight")));
let input_norm_bias = w.remove(&format!("{p}.input_layernorm.bias")).map(|t| repl(t));
let post_norm = repl(take(&mut w, &format!("{p}.post_attention_layernorm.weight")));
let post_norm_bias = w.remove(&format!("{p}.post_attention_layernorm.bias")).map(|t| repl(t));
let input_norm_bias = w
.remove(&format!("{p}.input_layernorm.bias"))
.map(|t| repl(t));
let post_norm = repl(take(
&mut w,
&format!("{p}.post_attention_layernorm.weight"),
));
let post_norm_bias = w
.remove(&format!("{p}.post_attention_layernorm.bias"))
.map(|t| repl(t));
layers.push(GptOssBlock {
input_norm,
@@ -283,17 +364,27 @@ impl GptOss {
let local_num_kv_heads = config.num_kv_heads() / world;
let has_norm_bias = norm_bias.is_some();
let is_fp8 = layers.first().map(|l| l.expert_gate_up_fp8.is_some()).unwrap_or(false);
let is_mxfp4 = layers.first().map(|l| l.expert_gate_up_mxfp4.is_some()).unwrap_or(false);
let is_fp8 = layers
.first()
.map(|l| l.expert_gate_up_fp8.is_some())
.unwrap_or(false);
let is_mxfp4 = layers
.first()
.map(|l| l.expert_gate_up_mxfp4.is_some())
.unwrap_or(false);
if rank == 0 {
if has_norm_bias {
eprintln!("gpt-oss: detected LayerNorm bias — using LayerNorm instead of RMSNorm");
}
if is_fp8 {
eprintln!("gpt-oss: FP8 E4M3 quantized expert weights detected (W8A8 cuBLASLt mode)");
eprintln!(
"gpt-oss: FP8 E4M3 quantized expert weights detected (W8A8 cuBLASLt mode)"
);
}
if is_mxfp4 {
eprintln!("gpt-oss: MXFP4 quantized expert weights detected (W4A16 fused-GEMV mode)");
eprintln!(
"gpt-oss: MXFP4 quantized expert weights detected (W4A16 fused-GEMV mode)"
);
}
}
@@ -341,7 +432,13 @@ impl GptOss {
}
#[inline]
fn add_norm(x: &Tensor, residual: &Tensor, weight: &Tensor, bias: &Option<Tensor>, eps: f32) -> (Tensor, Tensor) {
fn add_norm(
x: &Tensor,
residual: &Tensor,
weight: &Tensor,
bias: &Option<Tensor>,
eps: f32,
) -> (Tensor, Tensor) {
match bias {
Some(b) => {
let sum = xserv_kernels::add(x, residual);
@@ -439,7 +536,6 @@ impl GptOss {
let k_all = add_bias(&matmul_2d(&normed, &layer.k_proj_wt), &layer.k_proj_bias);
let v_all = add_bias(&matmul_2d(&normed, &layer.v_proj_wt), &layer.v_proj_bias);
// Reshape for RoPE: [B, H*D] → [B, H, D]
let q_3d = q_all.reshape(&[batch, num_heads, head_dim]);
let k_3d = k_all.reshape(&[batch, num_kv_heads, head_dim]);
@@ -460,9 +556,17 @@ impl GptOss {
let sinks_ptr = layer.sinks.data_ptr() as *const c_void;
let attn_out = paged_decode_attention_sinks(
&q_4d, k_pool_ptr, v_pool_ptr, bt_ptr, cl_ptr,
&q_4d,
k_pool_ptr,
v_pool_ptr,
bt_ptr,
cl_ptr,
sinks_ptr,
batch, num_heads, num_kv_heads, head_dim, max_blocks,
batch,
num_heads,
num_kv_heads,
head_dim,
max_blocks,
layer.window_size,
);
@@ -471,9 +575,14 @@ impl GptOss {
self.all_reduce(&attn_proj);
let attn_proj = add_bias(&attn_proj, &layer.o_proj_bias);
// Residual + post-norm
let (normed, x_new) = Self::add_norm(&attn_proj, &residual, &layer.post_norm, &layer.post_norm_bias, eps);
let (normed, x_new) = Self::add_norm(
&attn_proj,
&residual,
&layer.post_norm,
&layer.post_norm_bias,
eps,
);
let residual = x_new;
let normed = normed.contiguous();
@@ -505,7 +614,9 @@ impl GptOss {
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();
@@ -532,14 +643,21 @@ impl GptOss {
let (k_full, v_full) = paged_cache.gather_kv_contiguous(slot, layer_idx);
// Flash attention with gpt-oss sinks + (per-layer) sliding window.
let attn_out = flash_attention_sinks(&q, &k_full, &v_full, &layer.sinks, layer.window_size);
let attn_out =
flash_attention_sinks(&q, &k_full, &v_full, &layer.sinks, layer.window_size);
let attn_merged = 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 attn_proj = add_bias(&attn_proj, &layer.o_proj_bias);
let (normed, x_new) = Self::add_norm(&attn_proj, &residual, &layer.post_norm, &layer.post_norm_bias, eps);
let (normed, x_new) = Self::add_norm(
&attn_proj,
&residual,
&layer.post_norm,
&layer.post_norm_bias,
eps,
);
let residual = x_new;
// MoE MLP
@@ -566,15 +684,11 @@ impl GptOss {
let expert_start = rank * local_experts;
// 1. Router: [tokens, hidden] @ [hidden, num_experts] + bias → [tokens, num_experts]
let router_logits = add_bias(
&matmul_2d(x, &layer.router_wt),
&layer.router_bias,
);
let router_logits = add_bias(&matmul_2d(x, &layer.router_wt), &layer.router_bias);
// 2. GPU top-k + softmax
let (topk_ids, topk_weights) = xserv_kernels::moe::moe_topk_softmax(
&router_logits, num_experts, top_k,
);
let (topk_ids, topk_weights) =
xserv_kernels::moe::moe_topk_softmax(&router_logits, num_experts, top_k);
// Sparse decode path: compute ONLY the routed experts. The dense path
// below reads every local expert's weights per forward; the sparse
@@ -588,15 +702,31 @@ impl GptOss {
let n = packed.shape()[1];
let k = packed.shape()[2] * 2;
xserv_kernels::moe::moe_sparse_gemv_mxfp4(
x, packed, scales, &layer.expert_gate_up_bias, &topk_ids,
num_tokens, top_k, n, k, expert_start, local_experts, false,
x,
packed,
scales,
&layer.expert_gate_up_bias,
&topk_ids,
num_tokens,
top_k,
n,
k,
expert_start,
local_experts,
false,
)
} else {
xserv_kernels::moe::moe_sparse_gemv_fp8(
x, layer.expert_gate_up_fp8.as_ref().unwrap(),
x,
layer.expert_gate_up_fp8.as_ref().unwrap(),
layer.expert_gate_up_scale.as_ref().unwrap(),
&layer.expert_gate_up_bias, &topk_ids,
num_tokens, top_k, expert_start, local_experts, false,
&layer.expert_gate_up_bias,
&topk_ids,
num_tokens,
top_k,
expert_start,
local_experts,
false,
)
};
@@ -611,20 +741,40 @@ impl GptOss {
let n = packed.shape()[1];
let k = packed.shape()[2] * 2;
xserv_kernels::moe::moe_sparse_gemv_mxfp4(
&activated, packed, scales, &layer.expert_down_bias, &topk_ids,
num_tokens, top_k, n, k, expert_start, local_experts, true,
&activated,
packed,
scales,
&layer.expert_down_bias,
&topk_ids,
num_tokens,
top_k,
n,
k,
expert_start,
local_experts,
true,
)
} else {
xserv_kernels::moe::moe_sparse_gemv_fp8(
&activated, layer.expert_down_fp8.as_ref().unwrap(),
&activated,
layer.expert_down_fp8.as_ref().unwrap(),
layer.expert_down_scale.as_ref().unwrap(),
&layer.expert_down_bias, &topk_ids,
num_tokens, top_k, expert_start, local_experts, true,
&layer.expert_down_bias,
&topk_ids,
num_tokens,
top_k,
expert_start,
local_experts,
true,
)
};
let moe_out = xserv_kernels::moe::moe_weighted_sum_sparse(
&down, &topk_ids, &topk_weights, expert_start, local_experts,
&down,
&topk_ids,
&topk_weights,
expert_start,
local_experts,
);
self.all_reduce(&moe_out);
return moe_out;
@@ -644,14 +794,24 @@ impl GptOss {
xserv_kernels::quantization::batched_gemv_mxfp4(&x2, packed, scales, n, k)
.reshape(&[local_experts, 1, n])
} else {
let w_bf16 = xserv_kernels::quantization::dequant_mxfp4_to_bf16_t(packed, scales, local_experts, n, k);
let w_bf16 = xserv_kernels::quantization::dequant_mxfp4_to_bf16_t(
packed,
scales,
local_experts,
n,
k,
);
xserv_kernels::moe::batched_gemm_strided(&x_rep, &w_bf16)
}
} else if let Some(ref wt_fp8_t) = layer.expert_gate_up_fp8 {
// W8A8: quantize activations with per-expert scalar scale, use cuBLASLt FP8 GEMM
let (x_fp8, x_scales) = xserv_kernels::quantization::quantize_bf16_to_fp8_rowwise(&x_rep);
let (x_fp8, x_scales) =
xserv_kernels::quantization::quantize_bf16_to_fp8_rowwise(&x_rep);
xserv_kernels::quantization::batched_gemm_fp8(
&x_fp8, &x_scales, wt_fp8_t, layer.expert_gate_up_scale.as_ref().unwrap(),
&x_fp8,
&x_scales,
wt_fp8_t,
layer.expert_gate_up_scale.as_ref().unwrap(),
)
} else {
xserv_kernels::moe::batched_gemm_strided(&x_rep, &layer.expert_gate_up_wt)
@@ -677,14 +837,24 @@ impl GptOss {
xserv_kernels::quantization::batched_gemv_mxfp4(&a2, packed, scales, n, k)
.reshape(&[local_experts, 1, n])
} else {
let w_bf16 = xserv_kernels::quantization::dequant_mxfp4_to_bf16_t(packed, scales, local_experts, n, k);
let w_bf16 = xserv_kernels::quantization::dequant_mxfp4_to_bf16_t(
packed,
scales,
local_experts,
n,
k,
);
xserv_kernels::moe::batched_gemm_strided(&activated, &w_bf16)
}
} else if let Some(ref wt_fp8) = layer.expert_down_fp8 {
// W8A8: quantize post-GLU activations to FP8, use cuBLASLt FP8 GEMM
let (act_fp8, act_scales) = xserv_kernels::quantization::quantize_bf16_to_fp8_rowwise(&activated);
let (act_fp8, act_scales) =
xserv_kernels::quantization::quantize_bf16_to_fp8_rowwise(&activated);
xserv_kernels::quantization::batched_gemm_fp8(
&act_fp8, &act_scales, wt_fp8, layer.expert_down_scale.as_ref().unwrap(),
&act_fp8,
&act_scales,
wt_fp8,
layer.expert_down_scale.as_ref().unwrap(),
)
} else {
xserv_kernels::moe::batched_gemm_strided(&activated, &layer.expert_down_wt)
@@ -695,8 +865,12 @@ impl GptOss {
// 9. Weighted sum across experts → [tokens, hidden]
let moe_out = xserv_kernels::moe::moe_weighted_sum(
&down, &topk_ids, &topk_weights,
expert_start, local_experts, top_k,
&down,
&topk_ids,
&topk_weights,
expert_start,
local_experts,
top_k,
);
self.all_reduce(&moe_out);
@@ -708,9 +882,7 @@ impl GptOss {
/// Upload a u32 slice to a pooled GPU buffer (synchronous H2D).
fn upload_u32(vals: &[u32]) -> xserv_cuda::GpuBuffer {
let bytes = unsafe {
std::slice::from_raw_parts(vals.as_ptr() as *const u8, vals.len() * 4)
};
let bytes = unsafe { std::slice::from_raw_parts(vals.as_ptr() as *const u8, vals.len() * 4) };
let mut buf = xserv_cuda::allocator::cached_alloc(bytes.len()).expect("alloc u32 upload");
buf.copy_from_host(bytes).unwrap();
buf
@@ -737,11 +909,16 @@ fn add_bias(x: &Tensor, bias: &Tensor) -> 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);
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>();
@@ -751,11 +928,16 @@ fn shard_rows(t: &Tensor, rank: usize, world: usize) -> Tensor {
}
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);
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);
@@ -769,11 +951,16 @@ fn shard_cols(t: &Tensor, rank: usize, world: usize) -> Tensor {
}
fn shard_1d(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(), 1);
let total = shape[0];
assert!(total % world == 0, "dim {total} not divisible by world {world}");
assert!(
total % world == 0,
"dim {total} not divisible by world {world}"
);
let local = total / world;
let host = t.to_device(Device::Cpu);
let data = host.as_slice::<bf16>();
@@ -804,7 +991,13 @@ fn transpose_3d_inner_raw(t: &Tensor, batch: usize, rows: usize, cols: usize) ->
}
/// Extract experts [start..start+count) from a [num_experts, rows, cols] 3D tensor (any dtype, raw bytes).
fn slice_expert_range_3d_raw(t: &Tensor, start: usize, count: usize, rows: usize, cols: usize) -> Tensor {
fn slice_expert_range_3d_raw(
t: &Tensor,
start: usize,
count: usize,
rows: usize,
cols: usize,
) -> Tensor {
assert_eq!(t.ndim(), 3);
let host = t.to_device(Device::Cpu);
let elem_size = t.dtype().size_bytes();
@@ -826,7 +1019,13 @@ fn slice_scale_range(t: &Tensor, start: usize, count: usize) -> Tensor {
}
/// Extract experts [start..start+count) from a [num_experts, rows, cols] 3D tensor
fn slice_expert_range_3d(t: &Tensor, start: usize, count: usize, rows: usize, cols: usize) -> Tensor {
fn slice_expert_range_3d(
t: &Tensor,
start: usize,
count: usize,
rows: usize,
cols: usize,
) -> Tensor {
assert_eq!(t.ndim(), 3);
let host = t.to_device(Device::Cpu);
let data = host.as_slice::<bf16>();