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

189
crates/xserv-kernels/src/activation.rs
View File

@@ -6,78 +6,189 @@ unsafe extern "C" {
fn launch_gelu_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_silu_f32(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_silu_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_scale_f32(x: *const c_void, out: *mut c_void, scale: f32, n: i32, stream: *mut c_void);
fn launch_scale_bf16(x: *const c_void, out: *mut c_void, scale: f32, n: i32, stream: *mut c_void);
fn launch_add_f32(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_add_bf16(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_mul_f32(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_mul_bf16(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_silu_mul_bf16(gate: *const c_void, up: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_gpt_oss_glu_bf16(gate_up: *const c_void, out: *mut c_void, n_elements: i32,
alpha: f32, limit: f32, stream: *mut c_void);
fn launch_bias_add_2d_bf16(x: *const c_void, bias: *const c_void, out: *mut c_void,
rows: i32, cols: i32, stream: *mut c_void);
fn launch_scale_f32(
x: *const c_void,
out: *mut c_void,
scale: f32,
n: i32,
stream: *mut c_void,
);
fn launch_scale_bf16(
x: *const c_void,
out: *mut c_void,
scale: f32,
n: i32,
stream: *mut c_void,
);
fn launch_add_f32(
a: *const c_void,
b: *const c_void,
out: *mut c_void,
n: i32,
stream: *mut c_void,
);
fn launch_add_bf16(
a: *const c_void,
b: *const c_void,
out: *mut c_void,
n: i32,
stream: *mut c_void,
);
fn launch_mul_f32(
a: *const c_void,
b: *const c_void,
out: *mut c_void,
n: i32,
stream: *mut c_void,
);
fn launch_mul_bf16(
a: *const c_void,
b: *const c_void,
out: *mut c_void,
n: i32,
stream: *mut c_void,
);
fn launch_silu_mul_bf16(
gate: *const c_void,
up: *const c_void,
out: *mut c_void,
n: i32,
stream: *mut c_void,
);
fn launch_gpt_oss_glu_bf16(
gate_up: *const c_void,
out: *mut c_void,
n_elements: i32,
alpha: f32,
limit: f32,
stream: *mut c_void,
);
fn launch_bias_add_2d_bf16(
x: *const c_void,
bias: *const c_void,
out: *mut c_void,
rows: i32,
cols: i32,
stream: *mut c_void,
);
}
fn dispatch_unary(x: &Tensor, f32_fn: unsafe extern "C" fn(*const c_void, *mut c_void, i32, *mut c_void),
bf16_fn: unsafe extern "C" fn(*const c_void, *mut c_void, i32, *mut c_void)) -> Tensor {
fn dispatch_unary(
x: &Tensor,
f32_fn: unsafe extern "C" fn(*const c_void, *mut c_void, i32, *mut c_void),
bf16_fn: unsafe extern "C" fn(*const c_void, *mut c_void, i32, *mut c_void),
) -> Tensor {
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
let out = Tensor::empty(x.shape(), x.dtype(), x.device());
let n = x.numel();
assert!(n <= i32::MAX as usize, "tensor too large for i32 kernel param ({n} elements)");
assert!(
n <= i32::MAX as usize,
"tensor too large for i32 kernel param ({n} elements)"
);
let n = n as i32;
unsafe {
match x.dtype() {
DType::F32 => f32_fn(x.data_ptr() as _, out.data_ptr() as *mut c_void, n, xserv_cuda::current_stream_raw()),
DType::BF16 => bf16_fn(x.data_ptr() as _, out.data_ptr() as *mut c_void, n, xserv_cuda::current_stream_raw()),
DType::F32 => f32_fn(
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
n,
xserv_cuda::current_stream_raw(),
),
DType::BF16 => bf16_fn(
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
n,
xserv_cuda::current_stream_raw(),
),
_ => panic!("unsupported dtype"),
}
}
out
}
fn dispatch_binary(a: &Tensor, b: &Tensor,
f32_fn: unsafe extern "C" fn(*const c_void, *const c_void, *mut c_void, i32, *mut c_void),
bf16_fn: unsafe extern "C" fn(*const c_void, *const c_void, *mut c_void, i32, *mut c_void)) -> Tensor {
fn dispatch_binary(
a: &Tensor,
b: &Tensor,
f32_fn: unsafe extern "C" fn(*const c_void, *const c_void, *mut c_void, i32, *mut c_void),
bf16_fn: unsafe extern "C" fn(*const c_void, *const c_void, *mut c_void, i32, *mut c_void),
) -> Tensor {
assert_eq!(a.shape(), b.shape());
assert!(a.is_contiguous() && b.is_contiguous());
assert!(matches!(a.device(), Device::Cuda(_)));
assert_eq!(a.dtype(), b.dtype());
let out = Tensor::empty(a.shape(), a.dtype(), a.device());
let n = a.numel();
assert!(n <= i32::MAX as usize, "tensor too large for i32 kernel param ({n} elements)");
assert!(
n <= i32::MAX as usize,
"tensor too large for i32 kernel param ({n} elements)"
);
let n = n as i32;
unsafe {
match a.dtype() {
DType::F32 => f32_fn(a.data_ptr() as _, b.data_ptr() as _, out.data_ptr() as *mut c_void, n, xserv_cuda::current_stream_raw()),
DType::BF16 => bf16_fn(a.data_ptr() as _, b.data_ptr() as _, out.data_ptr() as *mut c_void, n, xserv_cuda::current_stream_raw()),
DType::F32 => f32_fn(
a.data_ptr() as _,
b.data_ptr() as _,
out.data_ptr() as *mut c_void,
n,
xserv_cuda::current_stream_raw(),
),
DType::BF16 => bf16_fn(
a.data_ptr() as _,
b.data_ptr() as _,
out.data_ptr() as *mut c_void,
n,
xserv_cuda::current_stream_raw(),
),
_ => panic!("unsupported dtype"),
}
}
out
}
pub fn gelu(x: &Tensor) -> Tensor { dispatch_unary(x, launch_gelu_f32, launch_gelu_bf16) }
pub fn silu(x: &Tensor) -> Tensor { dispatch_unary(x, launch_silu_f32, launch_silu_bf16) }
pub fn gelu(x: &Tensor) -> Tensor {
dispatch_unary(x, launch_gelu_f32, launch_gelu_bf16)
}
pub fn silu(x: &Tensor) -> Tensor {
dispatch_unary(x, launch_silu_f32, launch_silu_bf16)
}
pub fn scale(x: &Tensor, scale_val: f32) -> Tensor {
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
let out = Tensor::empty(x.shape(), x.dtype(), x.device());
let n = x.numel();
assert!(n <= i32::MAX as usize, "tensor too large for i32 kernel param ({n} elements)");
assert!(
n <= i32::MAX as usize,
"tensor too large for i32 kernel param ({n} elements)"
);
let n = n as i32;
unsafe {
match x.dtype() {
DType::F32 => launch_scale_f32(x.data_ptr() as _, out.data_ptr() as *mut c_void, scale_val, n, xserv_cuda::current_stream_raw()),
DType::BF16 => launch_scale_bf16(x.data_ptr() as _, out.data_ptr() as *mut c_void, scale_val, n, xserv_cuda::current_stream_raw()),
DType::F32 => launch_scale_f32(
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
scale_val,
n,
xserv_cuda::current_stream_raw(),
),
DType::BF16 => launch_scale_bf16(
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
scale_val,
n,
xserv_cuda::current_stream_raw(),
),
_ => panic!("unsupported dtype for scale"),
}
}
out
}
pub fn add(a: &Tensor, b: &Tensor) -> Tensor { dispatch_binary(a, b, launch_add_f32, launch_add_bf16) }
pub fn mul(a: &Tensor, b: &Tensor) -> Tensor { dispatch_binary(a, b, launch_mul_f32, launch_mul_bf16) }
pub fn add(a: &Tensor, b: &Tensor) -> Tensor {
dispatch_binary(a, b, launch_add_f32, launch_add_bf16)
}
pub fn mul(a: &Tensor, b: &Tensor) -> Tensor {
dispatch_binary(a, b, launch_mul_f32, launch_mul_bf16)
}
/// Row-broadcast bias add: out[r, c] = x[r, c] + bias[c] (BF16 only).
pub fn bias_add_2d(x: &Tensor, bias: &Tensor) -> Tensor {
@@ -89,13 +200,22 @@ pub fn bias_add_2d(x: &Tensor, bias: &Tensor) -> Tensor {
assert!(matches!(x.device(), Device::Cuda(_)));
let rows = x.shape()[0];
let cols = x.shape()[1];
assert_eq!(bias.shape()[0], cols, "bias size {} != cols {cols}", bias.shape()[0]);
assert_eq!(
bias.shape()[0],
cols,
"bias size {} != cols {cols}",
bias.shape()[0]
);
assert!(rows * cols <= i32::MAX as usize);
let out = Tensor::empty(&[rows, cols], DType::BF16, x.device());
unsafe {
launch_bias_add_2d_bf16(
x.data_ptr() as _, bias.data_ptr() as _, out.data_ptr() as *mut c_void,
rows as i32, cols as i32, xserv_cuda::current_stream_raw(),
x.data_ptr() as _,
bias.data_ptr() as _,
out.data_ptr() as *mut c_void,
rows as i32,
cols as i32,
xserv_cuda::current_stream_raw(),
);
}
out
@@ -110,7 +230,10 @@ pub fn silu_mul(gate: &Tensor, up: &Tensor) -> Tensor {
assert_eq!(gate.dtype(), DType::BF16, "silu_mul requires BF16");
let out = Tensor::empty(gate.shape(), gate.dtype(), gate.device());
let n = gate.numel();
assert!(n <= i32::MAX as usize, "tensor too large for i32 kernel param ({n} elements)");
assert!(
n <= i32::MAX as usize,
"tensor too large for i32 kernel param ({n} elements)"
);
let n = n as i32;
unsafe {
launch_silu_mul_bf16(