kernels: reshape_and_cache, GPU argmax, single-launch GEMV
Three new CUDA kernels and one rewrite: - reshape_and_cache: scatter K/V into paged pool in a single kernel per layer, replacing the Rust-side per-token per-head cudaMemcpy loop. Includes both single-sequence (prefill) and batched (decode) variants. - argmax: GPU-side BF16 argmax with warp-shuffle reduction. Greedy decode now only D2H-transfers B×4 bytes (token ids) instead of the full [B, vocab] logits tensor. - GEMV rewrite: fused zero-init inside the K-split kernel eliminates the cudaMemsetAsync call, reducing launches from 3 to 2 per GEMV. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
@@ -21,12 +21,14 @@ fn main() {
|
||||
.file("../../csrc/normalization/layernorm.cu")
|
||||
.file("../../csrc/activation/activations.cu")
|
||||
.file("../../csrc/reduce/softmax.cu")
|
||||
.file("../../csrc/reduce/argmax.cu")
|
||||
.file("../../csrc/embedding/embedding.cu")
|
||||
.file("../../csrc/embedding/rope.cu")
|
||||
.file("../../csrc/attention/causal_mask.cu")
|
||||
.file("../../csrc/embedding/transpose.cu")
|
||||
.file("../../csrc/attention/flash_attention.cu")
|
||||
.file("../../csrc/attention/paged_attention.cu")
|
||||
.file("../../csrc/attention/reshape_and_cache.cu")
|
||||
.compile("xserv_kernels");
|
||||
|
||||
println!("cargo:rerun-if-changed=../../csrc/");
|
||||
|
||||
65
crates/xserv-kernels/src/argmax.rs
Normal file
65
crates/xserv-kernels/src/argmax.rs
Normal file
@@ -0,0 +1,65 @@
|
||||
use std::ffi::c_void;
|
||||
use xserv_tensor::{DType, Device, Tensor};
|
||||
|
||||
unsafe extern "C" {
|
||||
fn launch_argmax_bf16(logits: *const c_void, out_idx: *mut c_void,
|
||||
rows: i32, cols: i32, stream: *mut c_void);
|
||||
}
|
||||
|
||||
/// GPU argmax over the last dim of a [rows, cols] BF16 tensor.
|
||||
///
|
||||
/// Returns a host `Vec<u32>` of length `rows`. Internally:
|
||||
/// - launches one kernel that writes [rows] i32 indices on device
|
||||
/// - D2H copies just `rows * 4` bytes (vs `rows * cols * 2` for the
|
||||
/// "copy logits to CPU then argmax" path it replaces)
|
||||
///
|
||||
/// This is the greedy-decode hot path: avoids touching the full
|
||||
/// [B, vocab] logits buffer on the host every step.
|
||||
pub fn argmax_bf16_to_host(logits: &Tensor) -> Vec<u32> {
|
||||
assert_eq!(logits.ndim(), 2, "argmax expects a 2D [rows, cols] tensor");
|
||||
assert_eq!(logits.dtype(), DType::BF16, "argmax kernel is BF16-only");
|
||||
assert!(logits.is_contiguous(), "argmax requires contiguous input");
|
||||
assert!(matches!(logits.device(), Device::Cuda(_)), "argmax requires GPU input");
|
||||
|
||||
let rows = logits.shape()[0];
|
||||
let cols = logits.shape()[1];
|
||||
assert!(rows <= i32::MAX as usize);
|
||||
assert!(cols <= i32::MAX as usize);
|
||||
|
||||
// Output buffer: rows * i32. Pooled allocator so this is essentially free
|
||||
// after the first call.
|
||||
let bytes = rows * std::mem::size_of::<i32>();
|
||||
let mut out = xserv_cuda::allocator::cached_alloc(bytes).expect("argmax out alloc");
|
||||
|
||||
unsafe {
|
||||
launch_argmax_bf16(
|
||||
logits.data_ptr() as *const c_void,
|
||||
out.as_mut_ptr() as *mut c_void,
|
||||
rows as i32, cols as i32,
|
||||
std::ptr::null_mut(),
|
||||
);
|
||||
}
|
||||
|
||||
let mut host_bytes = vec![0u8; bytes];
|
||||
out.copy_to_host(&mut host_bytes).expect("argmax D2H");
|
||||
drop(out); // returned to pool
|
||||
|
||||
let host_i32: &[i32] = unsafe {
|
||||
std::slice::from_raw_parts(host_bytes.as_ptr() as *const i32, rows)
|
||||
};
|
||||
host_i32.iter().map(|&v| v as u32).collect()
|
||||
}
|
||||
|
||||
/// Convenience: argmax of a single row [1, cols] (or [cols] reshaped to [1, cols]).
|
||||
pub fn argmax_bf16_single(logits: &Tensor) -> u32 {
|
||||
let cols = *logits.shape().last().unwrap();
|
||||
let rows = logits.numel() / cols;
|
||||
assert_eq!(rows, 1, "argmax_bf16_single requires a single row");
|
||||
let view = if logits.ndim() == 2 {
|
||||
logits.clone()
|
||||
} else {
|
||||
logits.reshape(&[1, cols])
|
||||
};
|
||||
argmax_bf16_to_host(&view)[0]
|
||||
}
|
||||
|
||||
@@ -33,6 +33,85 @@ unsafe extern "C" {
|
||||
head_dim: i32, max_blocks_per_seq: i32,
|
||||
scale: f32, stream: *mut c_void,
|
||||
);
|
||||
fn launch_reshape_and_cache_bf16(
|
||||
k_src: *const c_void, v_src: *const c_void,
|
||||
k_pool: *mut c_void, v_pool: *mut c_void,
|
||||
block_ids: *const c_void,
|
||||
num_tokens: i32, num_heads: i32,
|
||||
head_dim: i32, start_pos: i32, block_size: i32,
|
||||
stream: *mut c_void,
|
||||
);
|
||||
fn launch_reshape_and_cache_batched_bf16(
|
||||
k_src: *const c_void, v_src: *const c_void,
|
||||
k_pool: *mut c_void, v_pool: *mut c_void,
|
||||
block_tables: *const c_void, kv_lens: *const c_void,
|
||||
batch: i32, num_heads: i32,
|
||||
head_dim: i32, block_size: i32, max_blocks_per_seq: i32,
|
||||
stream: *mut c_void,
|
||||
);
|
||||
}
|
||||
|
||||
/// Scatter `[num_kv_heads, num_tokens, head_dim]` BF16 K/V into a paged
|
||||
/// pool for a single sequence whose block table lives at `block_ids_gpu`
|
||||
/// (int32, on device).
|
||||
///
|
||||
/// `k_pool_ptr`/`v_pool_ptr` point to one layer's pool, of logical shape
|
||||
/// `[num_blocks_total, num_kv_heads, block_size, head_dim]`.
|
||||
///
|
||||
/// All pointers must be on the same GPU as the launching context.
|
||||
///
|
||||
/// # Safety
|
||||
/// Pointers must be valid GPU pointers with the documented layouts.
|
||||
/// `block_ids_gpu` must contain at least `(start_pos + num_tokens + block_size - 1) / block_size`
|
||||
/// valid physical block ids.
|
||||
pub unsafe fn reshape_and_cache_bf16(
|
||||
k_src: *const c_void, v_src: *const c_void,
|
||||
k_pool_ptr: *mut c_void, v_pool_ptr: *mut c_void,
|
||||
block_ids_gpu: *const i32,
|
||||
num_tokens: usize, num_heads: usize,
|
||||
head_dim: usize, start_pos: usize, block_size: usize,
|
||||
stream: *mut c_void,
|
||||
) {
|
||||
unsafe {
|
||||
launch_reshape_and_cache_bf16(
|
||||
k_src, v_src,
|
||||
k_pool_ptr, v_pool_ptr,
|
||||
block_ids_gpu as *const c_void,
|
||||
num_tokens as i32, num_heads as i32,
|
||||
head_dim as i32, start_pos as i32, block_size as i32,
|
||||
stream,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Batched scatter for the multi-sequence decode step. Reads
|
||||
/// `block_tables` (`[batch, max_blocks_per_seq]` int32 — same buffer the
|
||||
/// paged-attention kernel reads) and `kv_lens` (`[batch]` int32, current
|
||||
/// seq_len + 1 — i.e., the index of the just-written token + 1) so the
|
||||
/// caller doesn't need a separate per-step upload of block ids.
|
||||
///
|
||||
/// # Safety
|
||||
/// All pointers must be on the same GPU. `block_tables` and `kv_lens` must
|
||||
/// already be synced to the device for the active batch.
|
||||
pub unsafe fn reshape_and_cache_batched_bf16(
|
||||
k_src: *const c_void, v_src: *const c_void,
|
||||
k_pool_ptr: *mut c_void, v_pool_ptr: *mut c_void,
|
||||
block_tables_gpu: *const i32, kv_lens_gpu: *const i32,
|
||||
batch: usize, num_heads: usize,
|
||||
head_dim: usize, block_size: usize, max_blocks_per_seq: usize,
|
||||
stream: *mut c_void,
|
||||
) {
|
||||
unsafe {
|
||||
launch_reshape_and_cache_batched_bf16(
|
||||
k_src, v_src,
|
||||
k_pool_ptr, v_pool_ptr,
|
||||
block_tables_gpu as *const c_void,
|
||||
kv_lens_gpu as *const c_void,
|
||||
batch as i32, num_heads as i32,
|
||||
head_dim as i32, block_size as i32, max_blocks_per_seq as i32,
|
||||
stream,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_causal_mask(scores: &Tensor, offset: usize) {
|
||||
|
||||
@@ -1,8 +1,16 @@
|
||||
use std::cell::RefCell;
|
||||
use std::ffi::c_void;
|
||||
use xserv_cuda::error::{self, Result};
|
||||
use xserv_cuda::GpuBuffer;
|
||||
use xserv_tensor::{DType, Device, Tensor};
|
||||
|
||||
const CUBLAS_WORKSPACE_BYTES: usize = 32 * 1024 * 1024;
|
||||
|
||||
// GEMV: single-kernel, no FP32 temp buffer needed
|
||||
unsafe extern "C" {
|
||||
fn launch_gemv_bf16(x: *const c_void, w: *const c_void, y_bf16: *mut c_void, y_fp32_buf: *mut c_void, k: i32, n: i32, stream: *mut c_void);
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum GemmBackend {
|
||||
Naive,
|
||||
@@ -16,7 +24,6 @@ unsafe extern "C" {
|
||||
fn launch_gemm_naive_bf16(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
|
||||
fn launch_gemm_tiled_f32(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
|
||||
fn launch_gemm_tiled_bf16(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
|
||||
fn launch_gemv_bf16(x: *const c_void, w: *const c_void, y_bf16: *mut c_void, y_fp32_buf: *mut c_void, k: i32, n: i32, stream: *mut c_void);
|
||||
}
|
||||
|
||||
// --- FFI: cuBLAS ---
|
||||
@@ -36,6 +43,7 @@ unsafe extern "C" {
|
||||
fn cublasCreate_v2(handle: *mut CublasHandle) -> i32;
|
||||
fn cublasDestroy_v2(handle: CublasHandle) -> i32;
|
||||
fn cublasSetStream_v2(handle: CublasHandle, stream: *mut c_void) -> i32;
|
||||
fn cublasSetWorkspace_v2(handle: CublasHandle, workspace: *mut c_void, size: usize) -> i32;
|
||||
fn cublasGemmEx(
|
||||
handle: CublasHandle,
|
||||
transa: i32, transb: i32,
|
||||
@@ -65,13 +73,25 @@ unsafe extern "C" {
|
||||
|
||||
pub struct CublasContext {
|
||||
handle: CublasHandle,
|
||||
/// Dedicated 32 MiB workspace owned by this handle. Held to keep the GPU
|
||||
/// buffer alive for the lifetime of the handle; cuBLAS reads/writes into
|
||||
/// it during GEMM. Dropped after `cublasDestroy_v2` so cuBLAS can't touch
|
||||
/// freed memory.
|
||||
_workspace: Option<GpuBuffer>,
|
||||
}
|
||||
|
||||
impl CublasContext {
|
||||
pub fn new() -> Result<Self> {
|
||||
let mut handle = std::ptr::null_mut();
|
||||
error::check(unsafe { cublasCreate_v2(&mut handle) })?;
|
||||
Ok(Self { handle })
|
||||
// Attach a per-handle workspace. cublasSetWorkspace requires the
|
||||
// pointer to remain valid until destroy or until a new workspace is
|
||||
// set, so we keep the GpuBuffer in this struct.
|
||||
let mut workspace = GpuBuffer::alloc(CUBLAS_WORKSPACE_BYTES)?;
|
||||
error::check(unsafe {
|
||||
cublasSetWorkspace_v2(handle, workspace.as_mut_ptr() as *mut c_void, CUBLAS_WORKSPACE_BYTES)
|
||||
})?;
|
||||
Ok(Self { handle, _workspace: Some(workspace) })
|
||||
}
|
||||
}
|
||||
|
||||
@@ -80,6 +100,7 @@ impl Drop for CublasContext {
|
||||
if !self.handle.is_null() {
|
||||
unsafe { cublasDestroy_v2(self.handle) };
|
||||
}
|
||||
// _workspace drops here, after cublasDestroy_v2 has released the handle.
|
||||
}
|
||||
}
|
||||
|
||||
@@ -152,7 +173,6 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
|
||||
}
|
||||
}
|
||||
GemmBackend::CuBlas => {
|
||||
// Fast path: custom GEMV for M=1 BF16 (bandwidth-optimal decode)
|
||||
if m == 1 && dtype == DType::BF16 {
|
||||
let mut fp32_buf = xserv_cuda::allocator::cached_alloc(n * 4).unwrap();
|
||||
unsafe {
|
||||
@@ -163,11 +183,7 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
|
||||
null_stream,
|
||||
);
|
||||
}
|
||||
// fp32_buf returned to caching allocator pool on drop
|
||||
} else {
|
||||
// cuBLAS uses column-major, but we have row-major tensors.
|
||||
// Trick: compute C^T = B^T @ A^T, which gives us C in row-major.
|
||||
// cuBLAS sees our row-major data as column-major transposed.
|
||||
let alpha = 1.0f32;
|
||||
let beta = 0.0f32;
|
||||
|
||||
@@ -179,19 +195,17 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
|
||||
|
||||
with_cublas(|handle| unsafe {
|
||||
cublasSetStream_v2(handle, null_stream);
|
||||
// Row-major trick: swap A/B and transpose flags
|
||||
// C(row-major) = A @ B <=> C^T(col-major) = B^T @ A^T
|
||||
error::check(cublasGemmEx(
|
||||
handle,
|
||||
CUBLAS_OP_N, CUBLAS_OP_N,
|
||||
n as i32, m as i32, k as i32,
|
||||
&alpha as *const f32 as *const c_void,
|
||||
b_ptr, b_type, n as i32, // B as col-major = B^T
|
||||
a_ptr, a_type, k as i32, // A as col-major = A^T
|
||||
b_ptr, b_type, n as i32,
|
||||
a_ptr, a_type, k as i32,
|
||||
&beta as *const f32 as *const c_void,
|
||||
c_ptr, c_type, n as i32, // C as col-major = C^T
|
||||
c_ptr, c_type, n as i32,
|
||||
CUBLAS_COMPUTE_32F,
|
||||
-1, // default algo
|
||||
-1,
|
||||
)).expect("cuBLAS GEMM failed");
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
pub mod activation;
|
||||
pub mod argmax;
|
||||
pub mod attention;
|
||||
pub mod dispatch;
|
||||
pub mod embedding;
|
||||
@@ -10,8 +11,9 @@ pub mod softmax;
|
||||
pub mod transpose;
|
||||
|
||||
pub use activation::{add, gelu, mul, scale, silu, silu_mul};
|
||||
pub use argmax::{argmax_bf16_single, argmax_bf16_to_host};
|
||||
pub use transpose::{merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, strided_to_contiguous_gpu, transpose_for_rope_gpu, transpose_from_rope_gpu};
|
||||
pub use attention::{attention, decode_attention, flash_attention, paged_decode_attention};
|
||||
pub use attention::{attention, decode_attention, flash_attention, paged_decode_attention, reshape_and_cache_bf16, reshape_and_cache_batched_bf16};
|
||||
pub use embedding::embedding;
|
||||
pub use gemm::{batched_matmul, matmul, GemmBackend};
|
||||
pub use layernorm::layernorm;
|
||||
|
||||
Reference in New Issue
Block a user