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

188
crates/xserv-kernels/src/transpose.rs
View File

@@ -2,19 +2,79 @@ use std::ffi::c_void;
use xserv_tensor::{DType, Device, Tensor};
unsafe extern "C" {
fn launch_reshape_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
fn launch_merge_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
fn launch_transpose_hsd_to_shd_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
fn launch_transpose_shd_to_hsd_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
fn launch_repeat_kv_bf16(inp: *const c_void, out: *mut c_void, kv_heads: i32, n_rep: i32, seq_len: i32, head_dim: i32, stream: *mut c_void);
fn launch_strided_copy_bf16(inp: *const c_void, out: *mut c_void, numel: i32, ndim: i32,
shape0: i32, shape1: i32, shape2: i32, shape3: i32,
in_stride0: i32, in_stride1: i32, in_stride2: i32, in_stride3: i32,
in_offset: i32, stream: *mut c_void);
fn launch_strided_copy_f32(inp: *const c_void, out: *mut c_void, numel: i32, ndim: i32,
shape0: i32, shape1: i32, shape2: i32, shape3: i32,
in_stride0: i32, in_stride1: i32, in_stride2: i32, in_stride3: i32,
in_offset: i32, stream: *mut c_void);
fn launch_reshape_heads_bf16(
inp: *const c_void,
out: *mut c_void,
seq_len: i32,
num_heads: i32,
head_dim: i32,
stream: *mut c_void,
);
fn launch_merge_heads_bf16(
inp: *const c_void,
out: *mut c_void,
seq_len: i32,
num_heads: i32,
head_dim: i32,
stream: *mut c_void,
);
fn launch_transpose_hsd_to_shd_bf16(
inp: *const c_void,
out: *mut c_void,
seq_len: i32,
num_heads: i32,
head_dim: i32,
stream: *mut c_void,
);
fn launch_transpose_shd_to_hsd_bf16(
inp: *const c_void,
out: *mut c_void,
seq_len: i32,
num_heads: i32,
head_dim: i32,
stream: *mut c_void,
);
fn launch_repeat_kv_bf16(
inp: *const c_void,
out: *mut c_void,
kv_heads: i32,
n_rep: i32,
seq_len: i32,
head_dim: i32,
stream: *mut c_void,
);
fn launch_strided_copy_bf16(
inp: *const c_void,
out: *mut c_void,
numel: i32,
ndim: i32,
shape0: i32,
shape1: i32,
shape2: i32,
shape3: i32,
in_stride0: i32,
in_stride1: i32,
in_stride2: i32,
in_stride3: i32,
in_offset: i32,
stream: *mut c_void,
);
fn launch_strided_copy_f32(
inp: *const c_void,
out: *mut c_void,
numel: i32,
ndim: i32,
shape0: i32,
shape1: i32,
shape2: i32,
shape3: i32,
in_stride0: i32,
in_stride1: i32,
in_stride2: i32,
in_stride3: i32,
in_offset: i32,
stream: *mut c_void,
);
}
/// [S, H*D] → [1, H, S, D] on GPU (BF16)
@@ -24,8 +84,12 @@ pub fn reshape_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim:
let out = Tensor::empty(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
unsafe {
launch_reshape_heads_bf16(
x.data_ptr() as _, out.data_ptr() as *mut c_void,
seq_len as i32, num_heads as i32, head_dim as i32, xserv_cuda::current_stream_raw(),
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
seq_len as i32,
num_heads as i32,
head_dim as i32,
xserv_cuda::current_stream_raw(),
);
}
out
@@ -39,36 +103,58 @@ pub fn merge_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: u
let out = Tensor::empty(&[seq_len, hidden], DType::BF16, x.device());
unsafe {
launch_merge_heads_bf16(
x.data_ptr() as _, out.data_ptr() as *mut c_void,
seq_len as i32, num_heads as i32, head_dim as i32, xserv_cuda::current_stream_raw(),
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
seq_len as i32,
num_heads as i32,
head_dim as i32,
xserv_cuda::current_stream_raw(),
);
}
out
}
/// [1, H, S, D] → [S, H, D] for RoPE on GPU (BF16)
pub fn transpose_for_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
pub fn transpose_for_rope_gpu(
x: &Tensor,
seq_len: usize,
num_heads: usize,
head_dim: usize,
) -> Tensor {
assert_eq!(x.dtype(), DType::BF16);
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
let out = Tensor::empty(&[seq_len, num_heads, head_dim], DType::BF16, x.device());
unsafe {
launch_transpose_hsd_to_shd_bf16(
x.data_ptr() as _, out.data_ptr() as *mut c_void,
seq_len as i32, num_heads as i32, head_dim as i32, xserv_cuda::current_stream_raw(),
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
seq_len as i32,
num_heads as i32,
head_dim as i32,
xserv_cuda::current_stream_raw(),
);
}
out
}
/// [S, H, D] → [1, H, S, D] after RoPE on GPU (BF16)
pub fn transpose_from_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
pub fn transpose_from_rope_gpu(
x: &Tensor,
seq_len: usize,
num_heads: usize,
head_dim: usize,
) -> Tensor {
assert_eq!(x.dtype(), DType::BF16);
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
let out = Tensor::empty(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
unsafe {
launch_transpose_shd_to_hsd_bf16(
x.data_ptr() as _, out.data_ptr() as *mut c_void,
seq_len as i32, num_heads as i32, head_dim as i32, xserv_cuda::current_stream_raw(),
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
seq_len as i32,
num_heads as i32,
head_dim as i32,
xserv_cuda::current_stream_raw(),
);
}
out
@@ -76,7 +162,9 @@ pub fn transpose_from_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, hea
/// [1, KV_H, S, D] → [1, KV_H*n_rep, S, D] on GPU (BF16)
pub fn repeat_kv_gpu(x: &Tensor, n_rep: usize) -> Tensor {
if n_rep == 1 { return x.clone(); }
if n_rep == 1 {
return x.clone();
}
assert_eq!(x.dtype(), DType::BF16);
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
let kv_heads = x.shape()[1];
@@ -86,8 +174,13 @@ pub fn repeat_kv_gpu(x: &Tensor, n_rep: usize) -> Tensor {
let out = Tensor::empty(&[1, new_heads, seq_len, head_dim], DType::BF16, x.device());
unsafe {
launch_repeat_kv_bf16(
x.data_ptr() as _, out.data_ptr() as *mut c_void,
kv_heads as i32, n_rep as i32, seq_len as i32, head_dim as i32, xserv_cuda::current_stream_raw(),
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
kv_heads as i32,
n_rep as i32,
seq_len as i32,
head_dim as i32,
xserv_cuda::current_stream_raw(),
);
}
out
@@ -122,20 +215,41 @@ pub fn strided_to_contiguous_gpu(x: &Tensor) -> Tensor {
unsafe {
match x.dtype() {
DType::BF16 => launch_strided_copy_bf16(
storage_ptr as _, out.data_ptr() as *mut c_void,
numel as i32, ndim as i32,
shape4[0], shape4[1], shape4[2], shape4[3],
strides4[0], strides4[1], strides4[2], strides4[3],
in_offset, xserv_cuda::current_stream_raw(),
storage_ptr as _,
out.data_ptr() as *mut c_void,
numel as i32,
ndim as i32,
shape4[0],
shape4[1],
shape4[2],
shape4[3],
strides4[0],
strides4[1],
strides4[2],
strides4[3],
in_offset,
xserv_cuda::current_stream_raw(),
),
DType::F32 => launch_strided_copy_f32(
storage_ptr as _, out.data_ptr() as *mut c_void,
numel as i32, ndim as i32,
shape4[0], shape4[1], shape4[2], shape4[3],
strides4[0], strides4[1], strides4[2], strides4[3],
in_offset, xserv_cuda::current_stream_raw(),
storage_ptr as _,
out.data_ptr() as *mut c_void,
numel as i32,
ndim as i32,
shape4[0],
shape4[1],
shape4[2],
shape4[3],
strides4[0],
strides4[1],
strides4[2],
strides4[3],
in_offset,
xserv_cuda::current_stream_raw(),
),
_ => panic!(
"strided_to_contiguous_gpu: unsupported dtype {:?}",
x.dtype()
),
_ => panic!("strided_to_contiguous_gpu: unsupported dtype {:?}", x.dtype()),
}
}
out