Gahow Wang
2be27d6d94
New CUDA kernels (csrc/embedding/transpose.cu): - reshape_heads_bf16: [S, H*D] → [1, H, S, D] - merge_heads_bf16: [1, H, S, D] → [S, H*D] - transpose_hsd_to_shd_bf16: [1, H, S, D] → [S, H, D] (for RoPE) - transpose_shd_to_hsd_bf16: [S, H, D] → [1, H, S, D] (from RoPE) - repeat_kv_bf16: [1, KV_H, S, D] → [1, KV_H*n_rep, S, D] Rust wrappers (xserv-kernels/src/transpose.rs): - reshape_heads_gpu, merge_heads_gpu, transpose_for/from_rope_gpu, repeat_kv_gpu Qwen3 forward_gpu_cache now uses all GPU kernels — zero CPU data round-trips. Result: 50/50 self-consistent, 3-5% faster (TBT 142→137ms) Remaining bottleneck: ~900 device::synchronize() calls + 252 cuBLAS handle creations per token (Phase 15 targets) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
92 lines
4.0 KiB
Rust
92 lines
4.0 KiB
Rust
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);
|
|
}
|
|
|
|
/// [S, H*D] → [1, H, S, D] on GPU (BF16)
|
|
pub fn reshape_heads_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::zeros(&[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, std::ptr::null_mut(),
|
|
);
|
|
}
|
|
xserv_cuda::device::synchronize().unwrap();
|
|
out
|
|
}
|
|
|
|
/// [1, H, S, D] → [S, H*D] on GPU (BF16)
|
|
pub fn merge_heads_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 hidden = num_heads * head_dim;
|
|
let out = Tensor::zeros(&[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, std::ptr::null_mut(),
|
|
);
|
|
}
|
|
xserv_cuda::device::synchronize().unwrap();
|
|
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 {
|
|
assert_eq!(x.dtype(), DType::BF16);
|
|
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
|
|
let out = Tensor::zeros(&[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, std::ptr::null_mut(),
|
|
);
|
|
}
|
|
xserv_cuda::device::synchronize().unwrap();
|
|
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 {
|
|
assert_eq!(x.dtype(), DType::BF16);
|
|
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
|
|
let out = Tensor::zeros(&[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, std::ptr::null_mut(),
|
|
);
|
|
}
|
|
xserv_cuda::device::synchronize().unwrap();
|
|
out
|
|
}
|
|
|
|
/// [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(); }
|
|
assert_eq!(x.dtype(), DType::BF16);
|
|
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
|
|
let kv_heads = x.shape()[1];
|
|
let seq_len = x.shape()[2];
|
|
let head_dim = x.shape()[3];
|
|
let new_heads = kv_heads * n_rep;
|
|
let out = Tensor::zeros(&[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, std::ptr::null_mut(),
|
|
);
|
|
}
|
|
xserv_cuda::device::synchronize().unwrap();
|
|
out
|
|
}
|