536 lines
16 KiB
Rust
536 lines
16 KiB
Rust
use std::ffi::c_void;
|
|
use xserv_tensor::{DType, Tensor};
|
|
|
|
use crate::gemm::{CublasHandle, cublas_handle};
|
|
|
|
unsafe extern "C" {
|
|
fn launch_moe_topk_softmax_bf16(
|
|
router_logits: *const c_void,
|
|
topk_ids: *mut c_void,
|
|
topk_weights: *mut c_void,
|
|
num_tokens: i32,
|
|
num_experts: i32,
|
|
top_k: i32,
|
|
stream: *mut c_void,
|
|
);
|
|
fn launch_moe_replicate_bf16(
|
|
x: *const c_void,
|
|
x_rep: *mut c_void,
|
|
num_tokens: i32,
|
|
hidden: i32,
|
|
local_experts: i32,
|
|
stream: *mut c_void,
|
|
);
|
|
fn launch_moe_bias_add_3d_bf16(
|
|
x: *mut c_void,
|
|
bias: *const c_void,
|
|
batch: i32,
|
|
num_tokens: i32,
|
|
dim: i32,
|
|
stream: *mut c_void,
|
|
);
|
|
fn launch_moe_weighted_sum_bf16(
|
|
expert_out: *const c_void,
|
|
topk_ids: *const c_void,
|
|
topk_weights: *const c_void,
|
|
out: *mut c_void,
|
|
num_tokens: i32,
|
|
hidden: i32,
|
|
top_k: i32,
|
|
expert_start: i32,
|
|
local_experts: i32,
|
|
stream: *mut c_void,
|
|
);
|
|
|
|
fn launch_moe_sparse_gemv_bf16_bf16(
|
|
x: *const c_void,
|
|
w: *const c_void,
|
|
topk_ids: *const c_void,
|
|
y: *mut c_void,
|
|
num_tokens: i32,
|
|
n: i32,
|
|
k: i32,
|
|
top_k: i32,
|
|
expert_start: i32,
|
|
local_experts: i32,
|
|
x_per_slot: i32,
|
|
stream: *mut c_void,
|
|
);
|
|
|
|
fn launch_moe_sparse_gemv_fp8_bf16(
|
|
x: *const c_void,
|
|
w: *const c_void,
|
|
w_scales: *const c_void,
|
|
bias: *const c_void,
|
|
topk_ids: *const c_void,
|
|
y: *mut c_void,
|
|
num_tokens: i32,
|
|
n: i32,
|
|
k: i32,
|
|
top_k: i32,
|
|
expert_start: i32,
|
|
local_experts: i32,
|
|
x_per_slot: i32,
|
|
stream: *mut c_void,
|
|
);
|
|
fn launch_moe_sparse_gemv_mxfp4_bf16(
|
|
x: *const c_void,
|
|
w_packed: *const c_void,
|
|
w_scales: *const c_void,
|
|
bias: *const c_void,
|
|
topk_ids: *const c_void,
|
|
y: *mut c_void,
|
|
num_tokens: i32,
|
|
n: i32,
|
|
k: i32,
|
|
top_k: i32,
|
|
expert_start: i32,
|
|
local_experts: i32,
|
|
x_per_slot: i32,
|
|
stream: *mut c_void,
|
|
);
|
|
fn launch_moe_weighted_sum_sparse_bf16(
|
|
down: *const c_void,
|
|
topk_ids: *const c_void,
|
|
topk_weights: *const c_void,
|
|
out: *mut c_void,
|
|
num_tokens: i32,
|
|
hidden: i32,
|
|
top_k: i32,
|
|
expert_start: i32,
|
|
local_experts: i32,
|
|
stream: *mut c_void,
|
|
);
|
|
|
|
fn cublasGemmStridedBatchedEx(
|
|
handle: CublasHandle,
|
|
transa: i32,
|
|
transb: i32,
|
|
m: i32,
|
|
n: i32,
|
|
k: i32,
|
|
alpha: *const c_void,
|
|
a: *const c_void,
|
|
a_type: i32,
|
|
lda: i32,
|
|
stride_a: i64,
|
|
b: *const c_void,
|
|
b_type: i32,
|
|
ldb: i32,
|
|
stride_b: i64,
|
|
beta: *const c_void,
|
|
c: *mut c_void,
|
|
c_type: i32,
|
|
ldc: i32,
|
|
stride_c: i64,
|
|
batch_count: i32,
|
|
compute_type: i32,
|
|
algo: i32,
|
|
) -> i32;
|
|
|
|
fn cublasSetStream_v2(handle: CublasHandle, stream: *mut c_void) -> i32;
|
|
}
|
|
|
|
const CUDA_R_16BF: i32 = 14;
|
|
const CUBLAS_COMPUTE_32F: i32 = 68;
|
|
const CUBLAS_GEMM_DEFAULT: i32 = -1;
|
|
|
|
/// GPU top-k selection + softmax over router logits.
|
|
///
|
|
/// Input: router_logits [num_tokens, num_experts] BF16 on GPU
|
|
/// Output: (topk_ids [num_tokens, top_k] i32, topk_weights [num_tokens, top_k] f32)
|
|
pub fn moe_topk_softmax(
|
|
router_logits: &Tensor,
|
|
num_experts: usize,
|
|
top_k: usize,
|
|
) -> (Tensor, Tensor) {
|
|
assert_eq!(router_logits.ndim(), 2);
|
|
assert_eq!(router_logits.dtype(), DType::BF16);
|
|
assert!(router_logits.is_contiguous());
|
|
let num_tokens = router_logits.shape()[0];
|
|
assert_eq!(router_logits.shape()[1], num_experts);
|
|
|
|
// NOTE: topk_ids actually holds i32 expert indices; DType has no I32, so
|
|
// this is a raw 4-byte buffer mislabeled F32. Never read it as floats —
|
|
// all consumers (weighted-sum / sparse GEMV kernels) cast to int*.
|
|
let topk_ids = Tensor::empty(&[num_tokens, top_k], DType::F32, router_logits.device());
|
|
let topk_weights = Tensor::empty(&[num_tokens, top_k], DType::F32, router_logits.device());
|
|
|
|
unsafe {
|
|
launch_moe_topk_softmax_bf16(
|
|
router_logits.data_ptr() as *const c_void,
|
|
topk_ids.data_ptr() as *mut c_void,
|
|
topk_weights.data_ptr() as *mut c_void,
|
|
num_tokens as i32,
|
|
num_experts as i32,
|
|
top_k as i32,
|
|
xserv_cuda::current_stream_raw(),
|
|
);
|
|
}
|
|
|
|
(topk_ids, topk_weights)
|
|
}
|
|
|
|
/// Replicate x [num_tokens, hidden] → [local_experts, num_tokens, hidden].
|
|
pub fn moe_replicate(x: &Tensor, local_experts: usize) -> Tensor {
|
|
assert_eq!(x.ndim(), 2);
|
|
assert_eq!(x.dtype(), DType::BF16);
|
|
assert!(x.is_contiguous());
|
|
let num_tokens = x.shape()[0];
|
|
let hidden = x.shape()[1];
|
|
let out = Tensor::empty(
|
|
&[local_experts, num_tokens, hidden],
|
|
DType::BF16,
|
|
x.device(),
|
|
);
|
|
|
|
unsafe {
|
|
launch_moe_replicate_bf16(
|
|
x.data_ptr() as *const c_void,
|
|
out.data_ptr() as *mut c_void,
|
|
num_tokens as i32,
|
|
hidden as i32,
|
|
local_experts as i32,
|
|
xserv_cuda::current_stream_raw(),
|
|
);
|
|
}
|
|
|
|
out
|
|
}
|
|
|
|
/// In-place 3D bias add: x [batch, num_tokens, dim] += bias [batch, dim].
|
|
pub fn moe_bias_add_3d(x: &Tensor, bias: &Tensor) {
|
|
assert_eq!(x.ndim(), 3);
|
|
assert_eq!(bias.ndim(), 2);
|
|
assert_eq!(x.dtype(), DType::BF16);
|
|
assert!(x.is_contiguous());
|
|
let batch = x.shape()[0];
|
|
let num_tokens = x.shape()[1];
|
|
let dim = x.shape()[2];
|
|
assert_eq!(bias.shape(), &[batch, dim]);
|
|
|
|
unsafe {
|
|
launch_moe_bias_add_3d_bf16(
|
|
x.data_ptr() as *mut c_void,
|
|
bias.data_ptr() as *const c_void,
|
|
batch as i32,
|
|
num_tokens as i32,
|
|
dim as i32,
|
|
xserv_cuda::current_stream_raw(),
|
|
);
|
|
}
|
|
}
|
|
|
|
/// Weighted sum of expert outputs → [num_tokens, hidden].
|
|
///
|
|
/// expert_out: [local_experts, num_tokens, hidden] BF16
|
|
/// topk_ids: [num_tokens, top_k] i32 (global expert indices)
|
|
/// topk_weights: [num_tokens, top_k] f32
|
|
pub fn moe_weighted_sum(
|
|
expert_out: &Tensor,
|
|
topk_ids: &Tensor,
|
|
topk_weights: &Tensor,
|
|
expert_start: usize,
|
|
local_experts: usize,
|
|
top_k: usize,
|
|
) -> Tensor {
|
|
assert_eq!(expert_out.ndim(), 3);
|
|
assert_eq!(expert_out.dtype(), DType::BF16);
|
|
let num_tokens = expert_out.shape()[1];
|
|
let hidden = expert_out.shape()[2];
|
|
|
|
let out = Tensor::empty(&[num_tokens, hidden], DType::BF16, expert_out.device());
|
|
|
|
unsafe {
|
|
launch_moe_weighted_sum_bf16(
|
|
expert_out.data_ptr() as *const c_void,
|
|
topk_ids.data_ptr() as *const c_void,
|
|
topk_weights.data_ptr() as *const c_void,
|
|
out.data_ptr() as *mut c_void,
|
|
num_tokens as i32,
|
|
hidden as i32,
|
|
top_k as i32,
|
|
expert_start as i32,
|
|
local_experts as i32,
|
|
xserv_cuda::current_stream_raw(),
|
|
);
|
|
}
|
|
|
|
out
|
|
}
|
|
|
|
/// Sparse MoE GEMV (BF16 weights): compute only routed experts.
|
|
/// x: [num_tokens, K] or [num_tokens * top_k, K], w_t: [local_experts, N, K].
|
|
pub fn moe_sparse_gemv_bf16(
|
|
x: &Tensor,
|
|
w_t: &Tensor,
|
|
topk_ids: &Tensor,
|
|
top_k: usize,
|
|
expert_start: usize,
|
|
local_experts: usize,
|
|
x_per_slot: bool,
|
|
) -> Tensor {
|
|
assert_eq!(x.ndim(), 2);
|
|
assert_eq!(w_t.ndim(), 3);
|
|
assert_eq!(x.dtype(), DType::BF16);
|
|
assert_eq!(w_t.dtype(), DType::BF16);
|
|
assert!(x.is_contiguous() && w_t.is_contiguous());
|
|
let local = w_t.shape()[0];
|
|
let n = w_t.shape()[1];
|
|
let k = w_t.shape()[2];
|
|
assert_eq!(local, local_experts);
|
|
assert_eq!(x.shape()[1], k);
|
|
let num_tokens = if x_per_slot {
|
|
x.shape()[0] / top_k
|
|
} else {
|
|
x.shape()[0]
|
|
};
|
|
let y = Tensor::empty(&[num_tokens, top_k, n], DType::BF16, x.device());
|
|
unsafe {
|
|
launch_moe_sparse_gemv_bf16_bf16(
|
|
x.data_ptr() as *const c_void,
|
|
w_t.data_ptr() as *const c_void,
|
|
topk_ids.data_ptr() as *const c_void,
|
|
y.data_ptr() as *mut c_void,
|
|
num_tokens as i32,
|
|
n as i32,
|
|
k as i32,
|
|
top_k as i32,
|
|
expert_start as i32,
|
|
local_experts as i32,
|
|
if x_per_slot { 1 } else { 0 },
|
|
xserv_cuda::current_stream_raw(),
|
|
);
|
|
}
|
|
y
|
|
}
|
|
|
|
/// Sparse MoE GEMV (FP8 W8A16): compute only the routed experts.
|
|
///
|
|
/// x: [num_tokens, K] BF16 (x_per_slot=false, gate_up) or
|
|
/// [num_tokens * top_k, K] BF16 (x_per_slot=true, down)
|
|
/// w_fp8_t: [local_experts, N, K] FP8E4M3 (transposed weight layout)
|
|
/// w_scales: [local_experts] F32 per-expert scalar scales
|
|
/// bias: [local_experts, N] BF16 (fused into the epilogue)
|
|
/// topk_ids: [num_tokens, top_k] i32 global expert ids (GPU)
|
|
///
|
|
/// Returns y [num_tokens, top_k, N] BF16. Slots routed to experts NOT
|
|
/// owned by this rank are left UNWRITTEN (uninitialized memory) — the
|
|
/// consumer must skip them (see moe_weighted_sum_sparse).
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn moe_sparse_gemv_fp8(
|
|
x: &Tensor,
|
|
w_fp8_t: &Tensor,
|
|
w_scales: &Tensor,
|
|
bias: &Tensor,
|
|
topk_ids: &Tensor,
|
|
num_tokens: usize,
|
|
top_k: usize,
|
|
expert_start: usize,
|
|
local_experts: usize,
|
|
x_per_slot: bool,
|
|
) -> Tensor {
|
|
assert_eq!(x.dtype(), DType::BF16);
|
|
assert!(x.is_contiguous());
|
|
assert_eq!(w_fp8_t.dtype(), DType::FP8E4M3);
|
|
let n = w_fp8_t.shape()[1];
|
|
let k = w_fp8_t.shape()[2];
|
|
// The kernel reads weights as uint4 (16 FP8 values per lane) and would
|
|
// silently skip a K%16 tail.
|
|
assert_eq!(k % 16, 0, "sparse FP8 GEMV requires K % 16 == 0, got {k}");
|
|
assert_eq!(x.shape()[x.ndim() - 1], k);
|
|
assert_eq!(
|
|
x.shape()[0],
|
|
if x_per_slot {
|
|
num_tokens * top_k
|
|
} else {
|
|
num_tokens
|
|
}
|
|
);
|
|
|
|
let y = Tensor::empty(&[num_tokens, top_k, n], DType::BF16, x.device());
|
|
unsafe {
|
|
launch_moe_sparse_gemv_fp8_bf16(
|
|
x.data_ptr() as *const c_void,
|
|
w_fp8_t.data_ptr() as *const c_void,
|
|
w_scales.data_ptr() as *const c_void,
|
|
bias.data_ptr() as *const c_void,
|
|
topk_ids.data_ptr() as *const c_void,
|
|
y.data_ptr() as *mut c_void,
|
|
num_tokens as i32,
|
|
n as i32,
|
|
k as i32,
|
|
top_k as i32,
|
|
expert_start as i32,
|
|
local_experts as i32,
|
|
x_per_slot as i32,
|
|
xserv_cuda::current_stream_raw(),
|
|
);
|
|
}
|
|
y
|
|
}
|
|
|
|
/// Sparse MoE GEMV (MXFP4 W4A16): same contract as moe_sparse_gemv_fp8,
|
|
/// with packed 4-bit weights [E, N, K/2] + UE8M0 block scales [E, N, K/32].
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn moe_sparse_gemv_mxfp4(
|
|
x: &Tensor,
|
|
w_packed: &Tensor,
|
|
w_scales: &Tensor,
|
|
bias: &Tensor,
|
|
topk_ids: &Tensor,
|
|
num_tokens: usize,
|
|
top_k: usize,
|
|
n: usize,
|
|
k: usize,
|
|
expert_start: usize,
|
|
local_experts: usize,
|
|
x_per_slot: bool,
|
|
) -> Tensor {
|
|
assert_eq!(x.dtype(), DType::BF16);
|
|
assert!(x.is_contiguous());
|
|
// 32-element MXFP4 blocks, read as uint4 (32 nibbles) per lane.
|
|
assert_eq!(k % 32, 0, "sparse MXFP4 GEMV requires K % 32 == 0, got {k}");
|
|
assert_eq!(x.shape()[x.ndim() - 1], k);
|
|
assert_eq!(
|
|
x.shape()[0],
|
|
if x_per_slot {
|
|
num_tokens * top_k
|
|
} else {
|
|
num_tokens
|
|
}
|
|
);
|
|
|
|
let y = Tensor::empty(&[num_tokens, top_k, n], DType::BF16, x.device());
|
|
unsafe {
|
|
launch_moe_sparse_gemv_mxfp4_bf16(
|
|
x.data_ptr() as *const c_void,
|
|
w_packed.data_ptr() as *const c_void,
|
|
w_scales.data_ptr() as *const c_void,
|
|
bias.data_ptr() as *const c_void,
|
|
topk_ids.data_ptr() as *const c_void,
|
|
y.data_ptr() as *mut c_void,
|
|
num_tokens as i32,
|
|
n as i32,
|
|
k as i32,
|
|
top_k as i32,
|
|
expert_start as i32,
|
|
local_experts as i32,
|
|
x_per_slot as i32,
|
|
xserv_cuda::current_stream_raw(),
|
|
);
|
|
}
|
|
y
|
|
}
|
|
|
|
/// Weighted sum over the slot axis of the sparse GEMV output.
|
|
///
|
|
/// down: [num_tokens, top_k, hidden] BF16 (non-local slots uninitialized
|
|
/// and skipped, never multiplied by zero — NaN * 0 = NaN).
|
|
pub fn moe_weighted_sum_sparse(
|
|
down: &Tensor,
|
|
topk_ids: &Tensor,
|
|
topk_weights: &Tensor,
|
|
expert_start: usize,
|
|
local_experts: usize,
|
|
) -> Tensor {
|
|
assert_eq!(down.ndim(), 3);
|
|
assert_eq!(down.dtype(), DType::BF16);
|
|
let num_tokens = down.shape()[0];
|
|
let top_k = down.shape()[1];
|
|
let hidden = down.shape()[2];
|
|
|
|
let out = Tensor::empty(&[num_tokens, hidden], DType::BF16, down.device());
|
|
unsafe {
|
|
launch_moe_weighted_sum_sparse_bf16(
|
|
down.data_ptr() as *const c_void,
|
|
topk_ids.data_ptr() as *const c_void,
|
|
topk_weights.data_ptr() as *const c_void,
|
|
out.data_ptr() as *mut c_void,
|
|
num_tokens as i32,
|
|
hidden as i32,
|
|
top_k as i32,
|
|
expert_start as i32,
|
|
local_experts as i32,
|
|
xserv_cuda::current_stream_raw(),
|
|
);
|
|
}
|
|
out
|
|
}
|
|
|
|
/// Strided batched GEMM for MoE expert forward.
|
|
/// C[b] = A[b] @ B[b] for b in 0..batch
|
|
///
|
|
/// A: [batch, M, K] BF16 contiguous
|
|
/// B: [batch, K, N] BF16 contiguous
|
|
/// Returns C: [batch, M, N] BF16
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn batched_gemm_strided(a: &Tensor, b: &Tensor) -> Tensor {
|
|
assert_eq!(a.ndim(), 3);
|
|
assert_eq!(b.ndim(), 3);
|
|
assert_eq!(a.dtype(), DType::BF16);
|
|
assert_eq!(b.dtype(), DType::BF16);
|
|
assert!(a.is_contiguous() && b.is_contiguous());
|
|
assert_eq!(a.shape()[0], b.shape()[0]);
|
|
assert_eq!(a.shape()[2], b.shape()[1]);
|
|
|
|
let batch = a.shape()[0];
|
|
let m = a.shape()[1];
|
|
let k = a.shape()[2];
|
|
let n = b.shape()[2];
|
|
|
|
let c = Tensor::empty(&[batch, m, n], DType::BF16, a.device());
|
|
|
|
let alpha: f32 = 1.0;
|
|
let beta: f32 = 0.0;
|
|
|
|
// cuBLAS column-major: we compute C^T = B^T @ A^T
|
|
// A is [batch, M, K] row-major → A^T is [K, M] col-major, lda=K
|
|
// B is [batch, K, N] row-major → B^T is [N, K] col-major, ldb=N? No...
|
|
//
|
|
// Actually for row-major: A[M,K] in memory = col-major A^T[K,M] with lda=K.
|
|
// So we call cublasGemmStridedBatchedEx with:
|
|
// transa=N, transb=N
|
|
// m=N, n=M, k=K (because cuBLAS sees col-major)
|
|
// A_cublas = B_row (pointer), lda=N
|
|
// B_cublas = A_row (pointer), ldb=K
|
|
// C_cublas = C_row (pointer), ldc=N
|
|
|
|
let stride_a = (m * k) as i64;
|
|
let stride_b = (k * n) as i64;
|
|
let stride_c = (m * n) as i64;
|
|
|
|
let handle = cublas_handle();
|
|
unsafe {
|
|
cublasSetStream_v2(handle, xserv_cuda::current_stream_raw());
|
|
let status = cublasGemmStridedBatchedEx(
|
|
handle,
|
|
0,
|
|
0, // CUBLAS_OP_N, CUBLAS_OP_N
|
|
n as i32,
|
|
m as i32,
|
|
k as i32,
|
|
&alpha as *const f32 as *const c_void,
|
|
b.data_ptr() as *const c_void,
|
|
CUDA_R_16BF,
|
|
n as i32,
|
|
stride_b,
|
|
a.data_ptr() as *const c_void,
|
|
CUDA_R_16BF,
|
|
k as i32,
|
|
stride_a,
|
|
&beta as *const f32 as *const c_void,
|
|
c.data_ptr() as *mut c_void,
|
|
CUDA_R_16BF,
|
|
n as i32,
|
|
stride_c,
|
|
batch as i32,
|
|
CUBLAS_COMPUTE_32F,
|
|
CUBLAS_GEMM_DEFAULT,
|
|
);
|
|
assert_eq!(status, 0, "cublasGemmStridedBatchedEx failed: {status}");
|
|
}
|
|
|
|
c
|
|
}
|