1047 lines
40 KiB
Rust
1047 lines
40 KiB
Rust
use half::bf16;
|
||
use std::collections::HashMap;
|
||
use std::ffi::c_void;
|
||
use xserv_kernels::*;
|
||
use xserv_tensor::{Device, Tensor};
|
||
|
||
use crate::config::ModelConfig;
|
||
use crate::paged_kv_cache::PagedKVCache;
|
||
|
||
pub struct GptOss {
|
||
pub config: ModelConfig,
|
||
embed_tokens: Tensor,
|
||
layers: Vec<GptOssBlock>,
|
||
norm: Tensor,
|
||
norm_bias: Option<Tensor>,
|
||
lm_head_t: Tensor,
|
||
rope_cache: RopeCache,
|
||
tp: Option<std::sync::Arc<xserv_distributed::TpContext>>,
|
||
local_num_heads: usize,
|
||
local_num_kv_heads: usize,
|
||
has_norm_bias: bool,
|
||
}
|
||
|
||
struct GptOssBlock {
|
||
input_norm: Tensor,
|
||
input_norm_bias: Option<Tensor>,
|
||
// Attention (with bias)
|
||
q_proj_wt: Tensor,
|
||
q_proj_bias: Tensor,
|
||
k_proj_wt: Tensor,
|
||
k_proj_bias: Tensor,
|
||
v_proj_wt: Tensor,
|
||
v_proj_bias: Tensor,
|
||
o_proj_wt: Tensor,
|
||
o_proj_bias: Tensor,
|
||
sinks: Tensor,
|
||
#[allow(dead_code)]
|
||
is_sliding: bool,
|
||
window_size: usize,
|
||
// MoE MLP
|
||
post_norm: Tensor,
|
||
post_norm_bias: Option<Tensor>,
|
||
router_wt: Tensor,
|
||
router_bias: Tensor,
|
||
// 3D expert weights for batched GEMM (contiguous on GPU)
|
||
expert_gate_up_wt: Tensor, // [local_experts, hidden, 2*inter] BF16
|
||
expert_gate_up_bias: Tensor, // [local_experts, 2*inter]
|
||
expert_down_wt: Tensor, // [local_experts, inter, hidden] BF16
|
||
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
|
||
// 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)>,
|
||
expert_down_mxfp4: Option<(Tensor, Tensor)>,
|
||
local_experts: usize,
|
||
// Activation params
|
||
glu_alpha: f32,
|
||
glu_limit: f32,
|
||
}
|
||
|
||
impl GptOss {
|
||
pub fn from_weights(config: ModelConfig, w: HashMap<String, Tensor>) -> Self {
|
||
Self::from_weights_tp(config, w, 0, 1, 0, None)
|
||
}
|
||
|
||
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}"))
|
||
};
|
||
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()
|
||
};
|
||
// 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()
|
||
};
|
||
// Bias sharding helpers
|
||
let col_bias = |t: Tensor| -> Tensor { shard_1d(&t, rank, world).to_device(dev) };
|
||
let repl_bias = |t: Tensor| -> Tensor { t.to_device(dev) };
|
||
|
||
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 head_dim = config.head_dim();
|
||
let rope_theta = config.rope_theta.unwrap_or(150000.0);
|
||
let max_seq_len = config.max_seq_len().min(8192); // cap for memory
|
||
|
||
let rope_cache = if let Some(ref rs) = config.rope_scaling {
|
||
if rs.rope_type.as_deref() == Some("yarn") {
|
||
RopeCache::new_yarn(
|
||
max_seq_len,
|
||
head_dim,
|
||
rope_theta,
|
||
rs.factor.unwrap_or(1.0),
|
||
rs.original_max_position_embeddings.unwrap_or(4096),
|
||
rs.beta_fast.unwrap_or(32.0),
|
||
rs.beta_slow.unwrap_or(1.0),
|
||
)
|
||
} else {
|
||
RopeCache::new(max_seq_len, head_dim, rope_theta as f32)
|
||
}
|
||
} else {
|
||
RopeCache::new(max_seq_len, head_dim, rope_theta as f32)
|
||
};
|
||
|
||
let num_layers = config.num_layers();
|
||
let num_experts = config.num_experts();
|
||
let glu_alpha = config.geglu_alpha();
|
||
let glu_limit = config.swiglu_limit.unwrap_or(7.0) as f32;
|
||
|
||
let mut layers = Vec::with_capacity(num_layers);
|
||
if rank == 0 {
|
||
eprintln!(
|
||
"Loading gpt-oss weights: {} layers, {} experts, world={world}, glu_alpha={glu_alpha}...",
|
||
num_layers, num_experts
|
||
);
|
||
}
|
||
|
||
for i in 0..num_layers {
|
||
let p = format!("model.layers.{i}");
|
||
|
||
// Attention weights — column-parallel for Q/K/V, row-parallel for O
|
||
let q_proj_wt = col(take(&mut w, &format!("{p}.self_attn.q_proj.weight")));
|
||
let q_proj_bias = col_bias(take(&mut w, &format!("{p}.self_attn.q_proj.bias")));
|
||
let k_proj_wt = col(take(&mut w, &format!("{p}.self_attn.k_proj.weight")));
|
||
let k_proj_bias = col_bias(take(&mut w, &format!("{p}.self_attn.k_proj.bias")));
|
||
let v_proj_wt = col(take(&mut w, &format!("{p}.self_attn.v_proj.weight")));
|
||
let v_proj_bias = col_bias(take(&mut w, &format!("{p}.self_attn.v_proj.bias")));
|
||
let o_proj_wt = row(take(&mut w, &format!("{p}.self_attn.o_proj.weight")));
|
||
let o_proj_bias = repl_bias(take(&mut w, &format!("{p}.self_attn.o_proj.bias")));
|
||
|
||
// Sinks: shard per-head across TP ranks
|
||
let sinks_full = take(&mut w, &format!("{p}.self_attn.sinks"));
|
||
let sinks = shard_1d(&sinks_full, rank, world).to_device(dev);
|
||
|
||
let is_sliding = config.is_sliding_layer(i);
|
||
let window_size = if is_sliding { config.window_size() } else { 0 };
|
||
|
||
// MoE weights — router replicated, experts split across TP ranks
|
||
let router_wt_raw = take(&mut w, &format!("{p}.mlp.router.weight"));
|
||
let router_wt = router_wt_raw.to_device(dev).transpose(0, 1).contiguous();
|
||
let router_bias = repl_bias(take(&mut w, &format!("{p}.mlp.router.bias")));
|
||
|
||
// Expert weights: [num_experts, hidden, 2*inter] — stored as 3D tensors
|
||
// Expert parallelism: rank owns experts [rank*E/world .. (rank+1)*E/world)
|
||
let gate_up_3d = take(&mut w, &format!("{p}.mlp.experts.gate_up_proj"));
|
||
let gate_up_bias_2d = take(&mut w, &format!("{p}.mlp.experts.gate_up_proj_bias"));
|
||
let down_3d = take(&mut w, &format!("{p}.mlp.experts.down_proj"));
|
||
let down_bias_2d = take(&mut w, &format!("{p}.mlp.experts.down_proj_bias"));
|
||
|
||
// FP8 scale tensors (present only in FP8-quantized models)
|
||
let gate_up_scale = w.remove(&format!("{p}.mlp.experts.gate_up_proj_scale"));
|
||
let down_scale = w.remove(&format!("{p}.mlp.experts.down_proj_scale"));
|
||
|
||
let local_experts = num_experts / world;
|
||
let expert_start = rank * local_experts;
|
||
|
||
// 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_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]
|
||
};
|
||
|
||
// Slice the rank's range of experts as contiguous 3D tensors on GPU
|
||
let expert_gate_up_wt;
|
||
let expert_down_wt;
|
||
let expert_gate_up_fp8;
|
||
let expert_gate_up_scale_gpu;
|
||
let expert_down_fp8;
|
||
let expert_down_scale_gpu;
|
||
|
||
if is_mxfp4 {
|
||
// MXFP4 W4A16: weights already [E, N, K] packed ([E, N, K/2] bytes)
|
||
// + 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);
|
||
expert_gate_up_mxfp4 = Some((gu_packed, gu_scl));
|
||
expert_down_mxfp4 = Some((dn_packed, dn_scl));
|
||
expert_gate_up_fp8 = None;
|
||
expert_gate_up_scale_gpu = None;
|
||
expert_down_fp8 = None;
|
||
expert_down_scale_gpu = None;
|
||
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 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),
|
||
);
|
||
// 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));
|
||
// 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_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);
|
||
|
||
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));
|
||
|
||
layers.push(GptOssBlock {
|
||
input_norm,
|
||
input_norm_bias,
|
||
q_proj_wt,
|
||
q_proj_bias,
|
||
k_proj_wt,
|
||
k_proj_bias,
|
||
v_proj_wt,
|
||
v_proj_bias,
|
||
o_proj_wt,
|
||
o_proj_bias,
|
||
sinks,
|
||
is_sliding,
|
||
window_size,
|
||
post_norm,
|
||
post_norm_bias,
|
||
router_wt,
|
||
router_bias,
|
||
expert_gate_up_wt,
|
||
expert_gate_up_bias,
|
||
expert_down_wt,
|
||
expert_down_bias,
|
||
expert_gate_up_fp8,
|
||
expert_gate_up_scale: expert_gate_up_scale_gpu,
|
||
expert_down_fp8,
|
||
expert_down_scale: expert_down_scale_gpu,
|
||
expert_gate_up_mxfp4,
|
||
expert_down_mxfp4,
|
||
local_experts,
|
||
glu_alpha,
|
||
glu_limit,
|
||
});
|
||
}
|
||
|
||
let local_num_heads = config.num_heads() / world;
|
||
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);
|
||
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)"
|
||
);
|
||
}
|
||
if is_mxfp4 {
|
||
eprintln!(
|
||
"gpt-oss: MXFP4 quantized expert weights detected (W4A16 fused-GEMV mode)"
|
||
);
|
||
}
|
||
}
|
||
|
||
// Warn about unused weights that the model didn't consume
|
||
if rank == 0 && !w.is_empty() {
|
||
eprintln!("WARNING: {} unused weight(s) in model:", w.len());
|
||
let mut keys: Vec<_> = w.keys().collect();
|
||
keys.sort();
|
||
for k in &keys {
|
||
eprintln!(" {k}");
|
||
}
|
||
}
|
||
|
||
Self {
|
||
config,
|
||
embed_tokens,
|
||
layers,
|
||
norm,
|
||
norm_bias,
|
||
lm_head_t,
|
||
rope_cache,
|
||
tp,
|
||
local_num_heads,
|
||
local_num_kv_heads,
|
||
has_norm_bias,
|
||
}
|
||
}
|
||
|
||
#[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 c_void;
|
||
tp.all_reduce_sum_bf16_ptr(ptr, t.numel());
|
||
}
|
||
}
|
||
}
|
||
|
||
#[inline]
|
||
fn norm(x: &Tensor, weight: &Tensor, bias: &Option<Tensor>, eps: f32) -> Tensor {
|
||
match bias {
|
||
Some(b) => layernorm(x, weight, b, eps),
|
||
None => rmsnorm(x, weight, eps),
|
||
}
|
||
}
|
||
|
||
#[inline]
|
||
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);
|
||
let normed = layernorm(&sum, weight, b, eps);
|
||
(normed, sum)
|
||
}
|
||
None => add_rmsnorm(x, residual, weight, eps),
|
||
}
|
||
}
|
||
|
||
fn norm_eps(&self) -> f32 {
|
||
if self.has_norm_bias {
|
||
self.config.ln_eps()
|
||
} else {
|
||
self.config.rms_norm_eps.unwrap_or(1e-5) as f32
|
||
}
|
||
}
|
||
|
||
/// Paged decode: process one token per sequence using paged KV cache.
|
||
pub fn forward_decode_paged(
|
||
&self,
|
||
tokens: &[u32],
|
||
positions: &[usize],
|
||
seq_slots: &[usize],
|
||
paged_cache: &mut PagedKVCache,
|
||
) -> Tensor {
|
||
let batch = tokens.len();
|
||
assert_eq!(positions.len(), batch);
|
||
assert_eq!(seq_slots.len(), batch);
|
||
assert!(batch > 0);
|
||
|
||
self.decode_prepare(positions, seq_slots, paged_cache);
|
||
|
||
// Upload token ids + positions, then run the pure-GPU core.
|
||
let ids_gpu = upload_u32(tokens);
|
||
let positions_u32: Vec<u32> = positions.iter().map(|&p| p as u32).collect();
|
||
let pos_gpu = upload_u32(&positions_u32);
|
||
let logits = self.decode_core(
|
||
ids_gpu.as_ptr() as *const c_void,
|
||
pos_gpu.as_ptr() as *const c_void,
|
||
batch,
|
||
paged_cache,
|
||
);
|
||
|
||
for &slot in seq_slots {
|
||
paged_cache.advance_seq_len(slot, 1);
|
||
}
|
||
logits
|
||
}
|
||
|
||
/// Host-side per-step cache bookkeeping: block allocation + uploading
|
||
/// block tables / context lens to their (stable-address) GPU buffers.
|
||
/// Runs OUTSIDE the CUDA-graph captured region.
|
||
pub fn decode_prepare(
|
||
&self,
|
||
positions: &[usize],
|
||
seq_slots: &[usize],
|
||
paged_cache: &mut PagedKVCache,
|
||
) {
|
||
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);
|
||
}
|
||
|
||
/// The pure-GPU decode step: embedding → 24 layers → final norm → logits.
|
||
/// Token ids and positions are read from device buffers; every other input
|
||
/// (weights, KV pools, block table, context lens) has a stable address —
|
||
/// which is exactly what makes this region CUDA-graph capturable.
|
||
pub fn decode_core(
|
||
&self,
|
||
ids_gpu: *const c_void,
|
||
pos_gpu: *const c_void,
|
||
batch: usize,
|
||
paged_cache: &mut PagedKVCache,
|
||
) -> Tensor {
|
||
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.norm_eps();
|
||
|
||
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();
|
||
|
||
let mut x = embedding_device_ids(&self.embed_tokens, ids_gpu, batch);
|
||
|
||
for (layer_idx, layer) in self.layers.iter().enumerate() {
|
||
let residual = x.clone();
|
||
let normed = Self::norm(&x, &layer.input_norm, &layer.input_norm_bias, eps);
|
||
|
||
// Q/K/V projections with bias
|
||
let q_all = add_bias(&matmul_2d(&normed, &layer.q_proj_wt), &layer.q_proj_bias);
|
||
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]);
|
||
|
||
// RoPE (no QK-norm for gpt-oss)
|
||
rope_inplace_device_pos(&q_3d, &self.rope_cache, pos_gpu);
|
||
rope_inplace_device_pos(&k_3d, &self.rope_cache, pos_gpu);
|
||
|
||
let v_3d = v_all.reshape(&[batch, num_kv_heads, head_dim]);
|
||
|
||
// KV cache scatter
|
||
paged_cache.append_tokens_batched(layer_idx, &k_3d, &v_3d, batch);
|
||
|
||
// Paged attention with sinks + sliding window
|
||
let q_4d = q_3d.reshape(&[batch, num_heads, 1, head_dim]);
|
||
let k_pool_ptr = paged_cache.k_pool(layer_idx).as_ptr() as *const c_void;
|
||
let v_pool_ptr = paged_cache.v_pool(layer_idx).as_ptr() as *const c_void;
|
||
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,
|
||
sinks_ptr,
|
||
batch,
|
||
num_heads,
|
||
num_kv_heads,
|
||
head_dim,
|
||
max_blocks,
|
||
layer.window_size,
|
||
);
|
||
|
||
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);
|
||
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 residual = x_new;
|
||
let normed = normed.contiguous();
|
||
|
||
// MoE MLP
|
||
let moe_out = self.moe_forward(&normed, layer, batch);
|
||
x = xserv_kernels::add(&residual, &moe_out);
|
||
}
|
||
|
||
let x = Self::norm(&x, &self.norm, &self.norm_bias, eps);
|
||
matmul_2d(&x, &self.lm_head_t)
|
||
}
|
||
|
||
/// Paged prefill: process full prompt tokens.
|
||
pub fn forward_prefill_paged(
|
||
&self,
|
||
token_ids: &[u32],
|
||
slot: usize,
|
||
paged_cache: &mut PagedKVCache,
|
||
) -> Tensor {
|
||
let new_tokens = token_ids.len();
|
||
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.norm_eps();
|
||
|
||
paged_cache.ensure_capacity(slot, pos_offset + new_tokens);
|
||
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();
|
||
|
||
for (layer_idx, layer) in self.layers.iter().enumerate() {
|
||
let residual = x.clone();
|
||
let normed = Self::norm(&x, &layer.input_norm, &layer.input_norm_bias, eps);
|
||
|
||
let q = add_bias(&matmul_2d(&normed, &layer.q_proj_wt), &layer.q_proj_bias);
|
||
let k = add_bias(&matmul_2d(&normed, &layer.k_proj_wt), &layer.k_proj_bias);
|
||
let v = add_bias(&matmul_2d(&normed, &layer.v_proj_wt), &layer.v_proj_bias);
|
||
|
||
let q = reshape_heads_gpu(&q, new_tokens, num_heads, head_dim);
|
||
let k = reshape_heads_gpu(&k, new_tokens, num_kv_heads, head_dim);
|
||
let v = reshape_heads_gpu(&v, new_tokens, num_kv_heads, head_dim);
|
||
|
||
// RoPE
|
||
let q = transpose_for_rope_gpu(&q, new_tokens, num_heads, head_dim);
|
||
let k = 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 = transpose_from_rope_gpu(&q, new_tokens, num_heads, head_dim);
|
||
let k = transpose_from_rope_gpu(&k, new_tokens, num_kv_heads, head_dim);
|
||
|
||
// KV cache
|
||
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);
|
||
|
||
// 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_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 residual = x_new;
|
||
|
||
// MoE MLP
|
||
let moe_out = self.moe_forward(&normed, layer, new_tokens);
|
||
x = xserv_kernels::add(&residual, &moe_out);
|
||
}
|
||
|
||
let x = Self::norm(&x, &self.norm, &self.norm_bias, eps);
|
||
matmul_2d(&x, &self.lm_head_t)
|
||
}
|
||
|
||
/// MoE forward pass — fully on GPU via batched GEMM.
|
||
///
|
||
/// Each rank owns `local_experts` experts. The input is replicated across all
|
||
/// local experts, processed via two batched cuBLAS GEMMs (gate_up and down),
|
||
/// and the selected experts' outputs are weighted-summed on GPU. Non-selected
|
||
/// experts contribute zero (via the routing weights), so no scatter/gather is
|
||
/// needed. AllReduce sums partial results across TP ranks.
|
||
fn moe_forward(&self, x: &Tensor, layer: &GptOssBlock, num_tokens: usize) -> Tensor {
|
||
let num_experts = self.config.num_experts();
|
||
let top_k = self.config.experts_per_token();
|
||
let rank = self.tp.as_ref().map(|tp| tp.rank).unwrap_or(0);
|
||
let local_experts = layer.local_experts;
|
||
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);
|
||
|
||
// 2. GPU top-k + softmax
|
||
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
|
||
// GEMVs read ~top_k/num_experts of the bytes, which dominates decode
|
||
// (memory-bound). Dense reads each weight once for ALL tokens, so it
|
||
// wins back at num_tokens ≈ local_experts / E[local hits] ≈ 8.
|
||
const SPARSE_MAX_TOKENS: usize = 8;
|
||
let quantized = layer.expert_gate_up_fp8.is_some() || layer.expert_gate_up_mxfp4.is_some();
|
||
if num_tokens <= SPARSE_MAX_TOKENS && quantized && !dense_moe_forced() {
|
||
let gate_up = if let Some((ref packed, ref scales)) = layer.expert_gate_up_mxfp4 {
|
||
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,
|
||
)
|
||
} else {
|
||
xserv_kernels::moe::moe_sparse_gemv_fp8(
|
||
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,
|
||
)
|
||
};
|
||
|
||
// GLU over all slots. Non-local slots hold unwritten memory; they
|
||
// are never consumed (the down GEMV and the weighted sum both skip
|
||
// slots whose expert this rank does not own).
|
||
let inter2 = gate_up.shape()[2];
|
||
let gate_up_flat = gate_up.reshape(&[num_tokens * top_k, inter2]);
|
||
let activated = gpt_oss_glu(&gate_up_flat, layer.glu_alpha, layer.glu_limit);
|
||
|
||
let down = if let Some((ref packed, ref scales)) = layer.expert_down_mxfp4 {
|
||
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,
|
||
)
|
||
} else {
|
||
xserv_kernels::moe::moe_sparse_gemv_fp8(
|
||
&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,
|
||
)
|
||
};
|
||
|
||
let moe_out = xserv_kernels::moe::moe_weighted_sum_sparse(
|
||
&down,
|
||
&topk_ids,
|
||
&topk_weights,
|
||
expert_start,
|
||
local_experts,
|
||
);
|
||
self.all_reduce(&moe_out);
|
||
return moe_out;
|
||
}
|
||
|
||
// 3. Replicate input: [tokens, hidden] → [local_experts, tokens, hidden]
|
||
let x_rep = xserv_kernels::moe::moe_replicate(x, local_experts);
|
||
|
||
// 4. Batched GEMM gate_up: [E, tokens, hidden] @ [E, hidden, 2*inter] → [E, tokens, 2*inter]
|
||
let gate_up = if let Some((ref packed, ref scales)) = layer.expert_gate_up_mxfp4 {
|
||
// MXFP4 W4A16: decode (M=1) uses the fused 4-bit dequant GEMV; prefill
|
||
// dequantizes to BF16 then reuses the batched GEMM.
|
||
let n = packed.shape()[1];
|
||
let k = packed.shape()[2] * 2;
|
||
if num_tokens == 1 {
|
||
let x2 = x_rep.reshape(&[local_experts, k]);
|
||
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,
|
||
);
|
||
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);
|
||
xserv_kernels::quantization::batched_gemm_fp8(
|
||
&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)
|
||
};
|
||
|
||
// 5. Bias add: gate_up += expert_gate_up_bias (in-place)
|
||
xserv_kernels::moe::moe_bias_add_3d(&gate_up, &layer.expert_gate_up_bias);
|
||
|
||
// 6. GLU activation: treat [E * tokens, 2*inter] → [E * tokens, inter]
|
||
let inter2 = gate_up.shape()[2];
|
||
let flat_rows = local_experts * num_tokens;
|
||
let gate_up_flat = gate_up.reshape(&[flat_rows, inter2]);
|
||
let activated_flat = gpt_oss_glu(&gate_up_flat, layer.glu_alpha, layer.glu_limit);
|
||
let inter = inter2 / 2;
|
||
let activated = activated_flat.reshape(&[local_experts, num_tokens, inter]);
|
||
|
||
// 7. Batched GEMM down: [E, tokens, inter] @ [E, inter, hidden] → [E, tokens, hidden]
|
||
let down = if let Some((ref packed, ref scales)) = layer.expert_down_mxfp4 {
|
||
let n = packed.shape()[1];
|
||
let k = packed.shape()[2] * 2;
|
||
if num_tokens == 1 {
|
||
let a2 = activated.reshape(&[local_experts, k]);
|
||
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,
|
||
);
|
||
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);
|
||
xserv_kernels::quantization::batched_gemm_fp8(
|
||
&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)
|
||
};
|
||
|
||
// 8. Bias add: down += expert_down_bias (in-place)
|
||
xserv_kernels::moe::moe_bias_add_3d(&down, &layer.expert_down_bias);
|
||
|
||
// 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,
|
||
);
|
||
|
||
self.all_reduce(&moe_out);
|
||
moe_out
|
||
}
|
||
}
|
||
|
||
// --- Helpers ---
|
||
|
||
/// 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 mut buf = xserv_cuda::allocator::cached_alloc(bytes.len()).expect("alloc u32 upload");
|
||
buf.copy_from_host(bytes).unwrap();
|
||
buf
|
||
}
|
||
|
||
/// XSERV_DENSE_MOE=1 forces the dense all-expert path (A/B benchmarking).
|
||
fn dense_moe_forced() -> bool {
|
||
static FORCED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
|
||
*FORCED.get_or_init(|| std::env::var("XSERV_DENSE_MOE").is_ok_and(|v| v != "0"))
|
||
}
|
||
|
||
fn matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
|
||
assert_eq!(a.ndim(), 2);
|
||
assert_eq!(b.ndim(), 2);
|
||
matmul(a, b, GemmBackend::CuBlas)
|
||
}
|
||
|
||
/// Add bias to a 2D tensor: [rows, cols] + [cols] → [rows, cols].
|
||
/// Single GPU broadcast kernel — the old rows>1 path tiled the bias on the
|
||
/// CPU (D2H + host loop + H2D) on every call, 96×/prefill in the hot path.
|
||
fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
|
||
let x_c = x.contiguous();
|
||
xserv_kernels::bias_add_2d(&x_c, bias)
|
||
}
|
||
|
||
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);
|
||
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])
|
||
}
|
||
|
||
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);
|
||
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 shard_1d(t: &Tensor, rank: usize, world: usize) -> Tensor {
|
||
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}"
|
||
);
|
||
let local = total / world;
|
||
let host = t.to_device(Device::Cpu);
|
||
let data = host.as_slice::<bf16>();
|
||
let start = rank * local;
|
||
let shard = data[start..start + local].to_vec();
|
||
Tensor::from_slice(&shard, &[local])
|
||
}
|
||
|
||
/// Transpose the inner two dimensions of a [batch, rows, cols] tensor → [batch, cols, rows].
|
||
/// Works on raw bytes (any dtype). CPU-only.
|
||
fn transpose_3d_inner_raw(t: &Tensor, batch: usize, rows: usize, cols: usize) -> Tensor {
|
||
assert_eq!(t.ndim(), 3);
|
||
assert_eq!(t.shape(), &[batch, rows, cols]);
|
||
let host = t.to_device(Device::Cpu);
|
||
let es = t.dtype().size_bytes();
|
||
let raw = host.as_raw_bytes();
|
||
let mut out = vec![0u8; batch * cols * rows * es];
|
||
for b in 0..batch {
|
||
for r in 0..rows {
|
||
for c in 0..cols {
|
||
let src_off = (b * rows * cols + r * cols + c) * es;
|
||
let dst_off = (b * cols * rows + c * rows + r) * es;
|
||
out[dst_off..dst_off + es].copy_from_slice(&raw[src_off..src_off + es]);
|
||
}
|
||
}
|
||
}
|
||
Tensor::from_raw_bytes(&out, &[batch, cols, rows], t.dtype())
|
||
}
|
||
|
||
/// 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 {
|
||
assert_eq!(t.ndim(), 3);
|
||
let host = t.to_device(Device::Cpu);
|
||
let elem_size = t.dtype().size_bytes();
|
||
let raw = host.as_raw_bytes();
|
||
let stride = rows * cols * elem_size;
|
||
let offset = start * stride;
|
||
let slice = &raw[offset..offset + count * stride];
|
||
Tensor::from_raw_bytes(slice, &[count, rows, cols], t.dtype())
|
||
}
|
||
|
||
/// Slice scale tensor [num_experts] F32 → [count] starting at `start`.
|
||
fn slice_scale_range(t: &Tensor, start: usize, count: usize) -> Tensor {
|
||
assert_eq!(t.ndim(), 1);
|
||
assert_eq!(t.dtype(), xserv_tensor::DType::F32);
|
||
let host = t.to_device(Device::Cpu);
|
||
let data = host.as_slice::<f32>();
|
||
let slice = data[start..start + count].to_vec();
|
||
Tensor::from_slice(&slice, &[count])
|
||
}
|
||
|
||
/// 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 {
|
||
assert_eq!(t.ndim(), 3);
|
||
let host = t.to_device(Device::Cpu);
|
||
let data = host.as_slice::<bf16>();
|
||
let stride = rows * cols;
|
||
let offset = start * stride;
|
||
let slice = data[offset..offset + count * stride].to_vec();
|
||
Tensor::from_slice(&slice, &[count, rows, cols])
|
||
}
|
||
|
||
/// Extract experts [start..start+count) from a [num_experts, dim] 2D tensor
|
||
fn slice_expert_range_2d(t: &Tensor, start: usize, count: usize, dim: usize) -> Tensor {
|
||
assert_eq!(t.ndim(), 2);
|
||
let host = t.to_device(Device::Cpu);
|
||
let data = host.as_slice::<bf16>();
|
||
let offset = start * dim;
|
||
let slice = data[offset..offset + count * dim].to_vec();
|
||
Tensor::from_slice(&slice, &[count, dim])
|
||
}
|