kaio-candle

Candle bridge for KAIO — CustomOp bindings that let you call KAIO's tensor-core GPU kernels directly on candle_core::Tensor.

Ships twelve ops: matmul_tc, matmul_tc_bf16, matmul_tc_async, matmul_tc_bf16_async, matmul_int4, matmul_int8, attention_tc, attention_tc_causal, attention_flash, attention_flash_causal, qkv_project_int8, qkv_project_int4. All four matmul TC variants (f16 + bf16, sync + async) support backward (autograd) via the forward-reuse pattern — no new PTX in either precision. FlashAttention (attention_flash + attention_flash_causal) supports backward through dedicated backward PTX kernels, preserving the no-O(seq²)-memory profile through the backward pass. attention_tc and quantized ops are forward-only.

Status — v0.2.0

v0.2.0 adds the bf16 matmul family (matmul_tc_bf16, matmul_tc_bf16_async, forward + backward) and FlashAttention backward (plain + causal) to the v0.1 surface.

All ops are bit-exact verified against direct kaio-ops calls with the same input bits.

Why a separate crate?

kaio-candle is not a member of the main KAIO workspace. cudarc rejects dynamic-loading + dynamic-linking as simultaneously active features:

• Main KAIO defaults to dynamic-loading — no CUDA toolkit required to build. Host tests pass on bare GitHub runners.
• candle-core with its cuda feature activates dynamic-linking — it links against libcuda at compile time.

Cargo unions features across a workspace build, so including kaio-candle in the main workspace would force every main-workspace build to also carry candle's dynamic-linking, breaking no-CUDA CI. The standalone crate keeps the two worlds apart.

Consumers who already build candle with the cuda feature see no new system requirement beyond what candle itself needs.

Build

cd kaio-candle
cargo build --features cuda

The cuda feature is required for any actual bridge functionality. Without it, kaio-candle is an empty shell (matches candle-core's own opt-in cuda pattern) — attempting to call kaio_candle::matmul_tc(...) surfaces a "function not found" compile error pointing at the missing feature.

Build requirements with cuda:

• CUDA toolkit (candle-core's cudarc feature uses dynamic-linking).
• NVIDIA GPU with SM 8.0 or newer (Ampere, Ada, Hopper).

Quickstart

# Cargo.toml
[dependencies]
kaio-candle = { version = "0.2", features = ["cuda"] }
kaio = "0.5"
candle-core = { version = "0.10", features = ["cuda"] }
half = "2"

use std::sync::Arc;
use candle_core::{Device, Tensor};
use half::f16;
use kaio::prelude::KaioDevice;

fn main() -> anyhow::Result<()> {
    let candle_dev = Device::new_cuda(0)?;
    let kaio_dev = Arc::new(KaioDevice::new(0)?);

    let m = 128usize;
    let k = 128usize;
    let n = 128usize;

    let a_host: Vec<f16> = (0..m * k).map(|i| f16::from_f32((i % 17) as f32 * 0.01)).collect();
    let b_host: Vec<f16> = (0..k * n).map(|i| f16::from_f32((i % 13) as f32 * 0.02)).collect();

    let a = Tensor::from_vec(a_host, (m, k), &candle_dev)?;
    let b = Tensor::from_vec(b_host, (k, n), &candle_dev)?;

    // f16[m,k] x f16[k,n] -> f32[m,n]
    let c = kaio_candle::matmul_tc(&kaio_dev, &a, &b)?;
    println!("output shape: {:?}", c.shape().dims());

    Ok(())
}

Three runnable examples ship in examples/:

cd kaio-candle
cargo run --release --features cuda --example matmul_tc_candle
cargo run --release --features cuda --example matmul_int8_candle
cargo run --release --features cuda --example attention_tc_candle

Op surface

matmul_int4 is GPTQ-style: group_size=128 is locked in by the kaio-ops kernel contract. K must be a multiple of 128, weights are packed 8 INT4 values per u32, one f16 scale per group of 128 elements.

matmul_int8 is W8A8 symmetric quant. Candle has no DType::I8, so the convention is DType::U8 tensors whose bytes are interpreted as signed INT8 (-128..=127) by the kernel. The bridge reinterprets the storage via a same-layout transmute. scale is a scalar f32 applied in the accumulator; a typical realistic value is max_abs / 127.

attention_tc uses a shared-memory scores buffer capped at seq_k ≤ 384. attention_flash has no seq cap (online softmax — no materialized score matrix) but is strictly single-head self-attention: Q, K, V must all be [seq_len, d_k] with d_k ≤ 256; cross-attention shapes are rejected with a pointer back to attention_tc.

qkv_project_int8 and qkv_project_int4 are direct-call functions (not CustomOpN — candle's trait maxes at 3 inputs and single output). They return (Tensor, Tensor, Tensor) with DType::F16 output because the fused kernel performs the f32→f16 conversion internally as part of the projection fusion. Gradient-tracked inputs are rejected with a loud error requiring .detach() — these ops are forward-only.

Backward support

Numerically approximate (matmul TC backward): The backward implementation downcasts the f32 upstream gradient to f16 (or bf16) to reuse the existing tensor-core forward kernels, and casts the output gradients back to the input dtype to satisfy candle's dtype-matching constraint. This is an initial autograd integration proving the bwd() bridge pattern, not a final mixed-precision training stack. The FlashAttention backward has no such round-trip — it is f32 end-to-end.

Memory (matmul TC backward): The backward pass materializes transposed tensors in VRAM (.t()?.contiguous() = allocation + copy). Peak backward memory is approximately 2–3x the forward input size. Designed for integration testing and light training, not high-throughput training loops where allocator overhead matters.

Recompute (FlashAttention backward): candle's CustomOp3 has no fwd→bwd saved-intermediate channel, so each backward call re-runs the stats-saving forward once to recover the per-row logsumexp before launching the backward kernels (the forward is deterministic; recomputed stats are bit-identical to saved ones). Direct kaio_ops::attention_flash_with_stats + attention_flash_bwd callers can hold the stats buffer themselves and skip the recompute. Measured cost of both tiers is in docs/performance.md.

Device lifetime

The Arc<kaio::prelude::KaioDevice> you construct and pass to kaio-candle wrapper functions is independent of the candle_core::Device you use for your tensors. Both retain the same CUDA primary context via cuDevicePrimaryCtxRetain; neither owns the other. Drop order between them is unconstrained.

Every wrapper call checks that the KAIO device and candle device share the same CUDA ordinal; a mismatch is a loud error.

Candle version policy

kaio-candle = 0.2 pins candle-core = "=0.10.2" exactly (unchanged from the 0.1.x line — candle 0.10.2 remains the current release). This is deliberate:

• The CustomOp2 / CustomOp3 surface has changed between candle minor versions in the past.
• cudarc feature conventions change with candle releases.

We re-pin kaio-candle against each new candle minor release. Use kaio-candle 0.1.x–0.2.x with candle-core 0.10.x; the next kaio-candle minor after a candle minor bump will target whichever candle release is current when it publishes.

A weekly GitHub Actions workflow (.github/workflows/candle-head.yml) builds kaio-candle against candle-core's git main branch once per Monday. If this badge goes red for more than two consecutive weeks, either the pin moves to the new candle minor or this section documents the divergence.

Known limitations (v0.2)

• Non-contiguous tensors rejected. Call .contiguous()? upstream.
• Non-zero storage offset rejected (e.g. from .narrow(...) / .slice(...)). Call .contiguous()? to compact.
• Rank-2 only. Multi-head attention callers must reshape [heads, seq, d] to [heads * seq, d] or call per-head with rank-2 slices. Wrappers error with a concrete reshape hint for higher-rank inputs.
• CUDA Graph capture partially unblocked. Event-based sync (Sprint 7.4c) removes the prior cuCtxSynchronize blocker. However, full CUDA Graph capture requires non-default streams on both the candle and KAIO sides, which is not yet verified.
• f32 output (CustomOp ops) / f16 output (direct-call ops). matmul_tc, matmul_tc_bf16 (and the async siblings), matmul_int4, matmul_int8, attention_tc, and attention_flash return f32 matching the kaio-ops accumulator. qkv_project_int{4,8} return f16 because the fused kernel converts internally.
• No CPU fallback. cpu_fwd returns a loud error rather than silently routing to candle.matmul(). KAIO's value is GPU-specific PTX; a silent CPU fallback would mask every perf claim.
• Bench numbers vs direct-call gap. Each bridge call issues event-based stream sync (two join() calls per op). This replaced the heavier cuCtxSynchronize fencing used during early bridge development but still allocates a transient CudaEvent per call. KAIO's published %-of-cuBLAS numbers are measured via direct kaio-ops calls, not through the bridge.

License

Dual-licensed under MIT or Apache-2.0, at your option.