KAIO

Rust-native GPU kernel authoring framework.

Write GPU compute kernels in Rust, compile to PTX, run on NVIDIA GPUs. No CUDA C++, no Python, no CUDA toolkit required. Windows + Linux.

Quick Start

use kaio::prelude::*;

#[gpu_kernel(block_size = 256)]
fn saxpy(x: *const [f32], y: *mut [f32], alpha: f32, n: u32) {
    let idx = thread_idx_x() + block_idx_x() * block_dim_x();
    if idx < n {
        y[idx] = alpha * x[idx] + y[idx];
    }
}

fn main() -> Result<()> {
    let device = KaioDevice::new(0)?;
    let n = 1024u32;

    let x = device.alloc_from(&vec![1.0f32; n as usize])?;
    let mut y = device.alloc_from(&vec![2.0f32; n as usize])?;

    saxpy::launch(&device, &x, &mut y, 2.5f32, n)?;

    let result = y.to_host(&device)?;
    println!("result: {:?}", &result[..8]);
    // prints: result: [4.5, 4.5, 4.5, 4.5, 4.5, 4.5, 4.5, 4.5]
    Ok(())
}

Requires an NVIDIA GPU with driver installed. No CUDA toolkit needed.

When to Use KAIO

KAIO is not a replacement for ML frameworks like Candle or Burn. It is the layer you use when you need more control than they provide.

• Write custom GPU kernels your framework doesn't support
• Control GPU memory explicitly — deterministic VRAM, buffer reuse
• Ship GPU binaries without Python or Triton in your dependency chain
• Run on Windows (Triton is Linux-only)
• Prototype GPU code in Rust without learning CUDA C++

Features

Architecture

Limitations

• NVIDIA only (SM 7.0+) — no AMD, no Intel
• Performance is size-dependent: tensor-core matmul meets or beats cuBLAS sgemm at 4096³ on RTX 4090 (worst-of-10: 107% sync / 115% async, precision caveats disclosed) but lags heavily at small sizes; the scalar path tops out at ~31% of cuBLAS — see the repository's docs/performance.md
• DSL subset of Rust — no closures, traits, generics, or &&/||
• FlashAttention requires d_k <= 256
• No autograd in the core crates (the kaio-candle bridge provides backward for the matmul TC family and FlashAttention); no multi-GPU
• API may change

Status

Phase 9 complete — tensor-core matmul family (f16 + bf16, sync + cp.async, auto-tuned), quantized INT8/INT4 matmul + fused QKV projections, FlashAttention forward + backward, candle bridge (kaio-candle), ldmatrix IR primitive. v0.5.0.

See the repository for full documentation, runnable examples, copy-paste patterns, and development logs.

License

Licensed under either of Apache-2.0 or MIT at your option.