Gahow Wang
d5532ef209
Two optimizations: 1. Tensor::empty() — skip cudaMemset for output tensors All kernel wrappers that fully overwrite their output now use Tensor::empty() instead of Tensor::zeros(). Eliminates ~756 cudaMemset calls per decode step (21 per layer × 36 layers). Improvement: 46.6 → 50.3 tok/s (+8%). 2. CUDA Graph infrastructure (for future use) Added FFI bindings (cudaStreamBeginCapture, cudaGraphInstantiate, cudaGraphLaunch) and RAII CudaGraph wrapper. Not yet used in the forward pass due to variable kv_len, but provides foundation for future graph-based decode optimization. Ablation (dash5, RTX 5090, Qwen3-8B BF16, serial decode): | Optimization | tok/s | vs HF | Roofline | |-------------|-------|-------|----------| | Phase 14 baseline | 12.9 | 36% | 12% | | + Fused kernels | 13.2 | 37% | 12% | | + Batched decode | 13.2 (serial) | 37% | 12% | | + Custom GEMV | 46.6 | 130% | 42% | | + Tensor::empty | 50.3 | 140% | 45% | Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
73 lines
2.4 KiB
Rust
73 lines
2.4 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;
|
|
|
|
// --- 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,
|
|
);
|
|
}
|