Gahow Wang
4c3f914459
CUDA layer for the paged-KV + swap work: - csrc: new paged_attention.cu plus updates across attention/gemm/norm/ activation/embedding/reduce kernels and common.cuh. - xserv-kernels: new dispatch module and kernel-binding updates. - xserv-cuda: cudaMallocHost/FreeHost bindings + PinnedBuffer (host swap pool backing) and offset-aware D2H/H2D copies used to move KV blocks between the GPU pool and pinned host memory. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
74 lines
2.5 KiB
Rust
74 lines
2.5 KiB
Rust
use std::ffi::c_void;
|
|
use std::os::raw::c_char;
|
|
|
|
pub type CudaStream = *mut c_void;
|
|
pub type CudaEvent = *mut c_void;
|
|
pub type CudaGraph = *mut c_void;
|
|
pub type CudaGraphExec = *mut c_void;
|
|
|
|
pub const CUDA_MEMCPY_H2D: i32 = 1;
|
|
pub const CUDA_MEMCPY_D2H: i32 = 2;
|
|
pub const CUDA_MEMCPY_D2D: i32 = 3;
|
|
|
|
pub const CUDA_SUCCESS: i32 = 0;
|
|
pub const CUDA_ERROR_OUT_OF_MEMORY: i32 = 2;
|
|
|
|
/// cudaStreamCaptureMode::cudaStreamCaptureModeGlobal
|
|
pub const CUDA_STREAM_CAPTURE_MODE_GLOBAL: i32 = 0;
|
|
|
|
unsafe extern "C" {
|
|
// --- Device ---
|
|
pub fn cudaGetDeviceCount(count: *mut i32) -> i32;
|
|
pub fn cudaSetDevice(device: i32) -> i32;
|
|
pub fn cudaGetDevice(device: *mut i32) -> i32;
|
|
/// Takes a raw pointer; caller provides a heap buffer large enough for any CUDA version.
|
|
pub fn cudaGetDeviceProperties(prop: *mut u8, device: i32) -> i32;
|
|
pub fn cudaDeviceSynchronize() -> i32;
|
|
|
|
// --- Memory ---
|
|
pub fn cudaMalloc(devptr: *mut *mut u8, size: usize) -> i32;
|
|
pub fn cudaFree(devptr: *mut u8) -> i32;
|
|
pub fn cudaMallocHost(ptr: *mut *mut u8, size: usize) -> i32;
|
|
pub fn cudaFreeHost(ptr: *mut u8) -> i32;
|
|
pub fn cudaMemcpy(dst: *mut u8, src: *const u8, count: usize, kind: i32) -> i32;
|
|
pub fn cudaMemcpyAsync(
|
|
dst: *mut u8,
|
|
src: *const u8,
|
|
count: usize,
|
|
kind: i32,
|
|
stream: CudaStream,
|
|
) -> i32;
|
|
pub fn cudaMemset(devptr: *mut u8, value: i32, count: usize) -> i32;
|
|
pub fn cudaMemsetAsync(devptr: *mut u8, value: i32, count: usize, stream: CudaStream) -> i32;
|
|
|
|
// --- Stream ---
|
|
pub fn cudaStreamCreate(stream: *mut CudaStream) -> i32;
|
|
pub fn cudaStreamDestroy(stream: CudaStream) -> i32;
|
|
pub fn cudaStreamSynchronize(stream: CudaStream) -> i32;
|
|
|
|
// --- Error ---
|
|
pub fn cudaGetLastError() -> i32;
|
|
pub fn cudaGetErrorString(error: i32) -> *const c_char;
|
|
|
|
// --- CUDA Graphs ---
|
|
pub fn cudaStreamBeginCapture(stream: CudaStream, mode: i32) -> i32;
|
|
pub fn cudaStreamEndCapture(stream: CudaStream, graph: *mut CudaGraph) -> i32;
|
|
pub fn cudaGraphInstantiate(
|
|
graph_exec: *mut CudaGraphExec,
|
|
graph: CudaGraph,
|
|
flags: u64,
|
|
) -> i32;
|
|
pub fn cudaGraphLaunch(graph_exec: CudaGraphExec, stream: CudaStream) -> i32;
|
|
pub fn cudaGraphDestroy(graph: CudaGraph) -> i32;
|
|
pub fn cudaGraphExecDestroy(graph_exec: CudaGraphExec) -> i32;
|
|
|
|
// --- Our test kernel ---
|
|
pub fn launch_vecadd_f32(
|
|
a: *const f32,
|
|
b: *const f32,
|
|
c: *mut f32,
|
|
n: i32,
|
|
stream: CudaStream,
|
|
);
|
|
}
|