Cuda Graph Example

[jafioti]

Hi! Very exciting project. Is it possible to use these kernels inside cuda graphs? I've gotten llama working end-to-end with cuTile-rs kernels (and ordinary kernels for rope + scatter because of an iota bug) but I cant figure out how to put them all in a cuda graph. Currently I'm bottlenecked to ~75ms per forward pass because of the overhead of launching each kernel one-by-one. An example showing how to use cuTile-rs kernels inside a cuda graph (preferably a mixed cuda graph with normal kernels in there as well) would be very useful.

[elibol]

I'm working on adding support for CUDA graphs. I'll include examples as soon as it is available.

Can you share your code? I might be able to help improve your impl in the meantime (as well as to make sure we understand and cover your use case).

[jafioti]

The full code is a bit large but I generated a minimal example:

/*
 * Minimal LLaMA-style inference loop in cuTile Rust.
 *
 * Demonstrates the performance bottleneck when running a multi-layer
 * transformer: each layer requires ~15 kernel launches, and `sync_on()`
 * synchronizes the stream after every one. For a 22-layer model this means
 * ~330 stream syncs per token, dominating the runtime.
 *
 * Request for the cuTile team:
 *   1. A way to launch kernels without per-launch stream sync (the
 *      `DeviceOperation::execute()` path works but requires `unsafe`).
 *   2. CUDA graph capture support — capture the full forward pass into a
 *      graph on the first call, then replay it with `cuGraphLaunch`.
 *      The challenge is that some kernel arguments change per token
 *      (e.g. position, kv_padded), so we'd need graph-node parameter
 *      updates or parameter indirection via GPU buffers.
 *   3. Non-owning Tensor views (zero-copy reshape). Currently
 *      `DeviceBox::from_raw_parts` takes ownership, so creating a view
 *      of existing GPU memory triggers a double-free on drop. A
 *      `BorrowedTensor` or `TensorView` type would eliminate many
 *      GPU-to-GPU memcpy operations in the inner loop.
 *
 */

use cuda_async::device_operation::DeviceOperation;
use cuda_core::CudaContext;
use cutile::api;
use cutile::error::Error;
use cutile::tensor::{IntoPartition, Partition, Tensor, ToHostVec};
use cutile::tile_kernel::TileKernel;
use std::sync::Arc;
use std::time::Instant;

// ═══════════════════════════════════════════════════════════════════════════════
// Kernels — a minimal transformer layer: RMSNorm → MatVec → Add
// ═══════════════════════════════════════════════════════════════════════════════

#[cutile::module]
mod kernels {
    use cutile::core::*;

    /// RMS normalization: out = rms_norm(x) * w
    #[cutile::entry()]
    pub fn rms_norm<const D: i32, const BS: i32>(
        x: &Tensor<f32, { [-1, D] }>,
        w: &Tensor<f32, { [D] }>,
        out: &mut Tensor<f32, { [1, D] }>,
        eps: f32,
    ) {
        let shape: Shape<{ [1, BS] }> = const_shape![1, BS];
        let tiles: i32 = D / BS;
        let pid: (i32, i32, i32) = get_tile_block_id();
        let row = pid.0;
        let x_part: Partition<f32, { [1, BS] }> = x.partition(shape);
        let mut rms: Tile<f32, { [1, BS] }> = constant(0.0, shape);
        for j in 0i32..tiles {
            let t: Tile<f32, { [1, BS] }> = x_part.load([row, j]);
            rms = rms + t * t;
        }
        let s: Tile<f32, { [1] }> = reduce_sum(rms, 1i32);
        let s: f32 = tile_to_scalar(s.reshape(const_shape![]));
        let n: f32 = convert_scalar(D);
        let inv: f32 = 1.0f32 / (s / n + eps);
        let inv_tile: Tile<f32, { [] }> = sqrt(scalar_to_tile(inv), "negative_inf");
        let inv: f32 = tile_to_scalar(inv_tile);
        let scale: Tile<f32, { [1, BS] }> = inv.broadcast(shape);
        let w_part: Partition<f32, { [BS] }> = w.partition(const_shape![BS]);
        let mut out_part: PartitionMut<f32, { [1, BS] }> =
            unsafe { out.partition_mut(shape) };
        for j in 0i32..tiles {
            let t: Tile<f32, { [1, BS] }> = x_part.load([row, j]);
            let tw: Tile<f32, { [1, BS] }> = w_part.load([j]).reshape(shape);
            unsafe { out_part.store(t * scale * tw, [0i32, j]) };
        }
    }

    /// Matrix-vector multiply: out = x @ w^T
    /// x: (1, K), w: (N, K), out: (N,)
    #[cutile::entry()]
    pub fn matvec<const BN: i32, const BK: i32, const K: i32>(
        out: &mut Tensor<f32, { [BN] }>,
        x: &Tensor<f32, { [-1, K] }>,
        w: &Tensor<f32, { [-1, K] }>,
    ) {
        let pid: (i32, i32, i32) = get_tile_block_id();
        let x_part = x.partition(const_shape![1, BK]);
        let w_part = w.partition(const_shape![BN, BK]);
        let mut acc = out.load().reshape(const_shape![BN, 1]);
        let transpose: Array<{ [1, 0] }> = Array::<{ [1, 0] }> {
            dims: &[1i32, 0i32],
        };
        for k in 0i32..(K / BK) {
            let tx = x_part.load([0i32, k]).reshape(const_shape![1, BK]);
            let tw = w_part.load([pid.0, k]);
            let txt: Tile<f32, { [BK, 1] }> = permute(tx, transpose);
            acc = mma(tw, txt, acc);
        }
        out.store(acc.reshape(const_shape![BN]));
    }

    /// Element-wise add: out = a + b
    #[cutile::entry()]
    pub fn add<const B: i32>(
        a: &Tensor<f32, { [-1] }>,
        b: &Tensor<f32, { [-1] }>,
        out: &mut Tensor<f32, { [B] }>,
    ) {
        let ta: Tile<f32, { [B] }> = load_tile_like_1d(a, out);
        let tb: Tile<f32, { [B] }> = load_tile_like_1d(b, out);
        out.store(ta + tb);
    }
}

use crate::kernels::{add_sync, matvec_sync, rms_norm_sync};

// ═══════════════════════════════════════════════════════════════════════════════
// Fake "model" with random weights — same shapes as a real LLaMA layer
// ═══════════════════════════════════════════════════════════════════════════════

struct Layer {
    norm_w: Arc<Tensor<f32>>,  // (D,)
    wq: Arc<Tensor<f32>>,     // (D, D)
    wo: Arc<Tensor<f32>>,     // (D, D)
}

fn main() -> Result<(), Error> {
    let ctx = CudaContext::new(0)?;
    let stream = ctx.new_stream()?;

    // Model config (TinyLlama-like)
    let d = 2048usize;
    let n_layers = 22usize;
    let block = 128usize;
    let bn = 16usize;
    let bk = 16usize;
    let eps = 1e-5f32;

    println!("Allocating fake model: d={d}, layers={n_layers}");

    // Create random weights
    let mut layers = Vec::new();
    for _ in 0..n_layers {
        layers.push(Layer {
            norm_w: api::randn(0.0f32, 0.02, [d]).sync_on(&stream)?.into(),
            wq: api::randn(0.0f32, 0.02, [d, d]).sync_on(&stream)?.into(),
            wo: api::randn(0.0f32, 0.02, [d, d]).sync_on(&stream)?.into(),
        });
    }

    let rms_gen = vec![(d as i32).to_string(), (block as i32).to_string()];
    let mv_gen = vec![
        (bn as i32).to_string(),
        (bk as i32).to_string(),
        d.to_string(),
    ];
    let add_gen = vec![(block as i32).to_string()];

    // ── Warmup: one pass to compile all kernels ─────────────────────────
    println!("Warming up (compiling kernels)...");
    {
        let h: Arc<Tensor<f32>> = api::randn(0.0f32, 1.0, [d]).sync_on(&stream)?.into();
        let h_2d: Arc<Tensor<f32>> = h
            .copy_sync(&stream)?
            .reshape([1, d])
            .into();
        let layer = &layers[0];

        let rms_out: Partition<Tensor<f32>> = api::zeros([1, d])
            .sync_on(&stream)?
            .partition([1, d as i32]);
        let _ = rms_norm_sync(h_2d.clone(), layer.norm_w.clone(), rms_out, eps)
            .generics(rms_gen.clone())
            .sync_on(&stream)?;

        let q_out: Partition<Tensor<f32>> = api::zeros([d])
            .sync_on(&stream)?
            .partition([bn as i32]);
        let _ = matvec_sync(q_out, h_2d.clone(), layer.wq.clone())
            .generics(mv_gen.clone())
            .sync_on(&stream)?;

        let add_out: Partition<Tensor<f32>> = api::zeros([d])
            .sync_on(&stream)?
            .partition([block as i32]);
        let _ = add_sync(h.clone(), h.clone(), add_out)
            .generics(add_gen.clone())
            .sync_on(&stream)?;
    }
    println!("Warmup done.");

    // ── Benchmark: sync_on path (current) ───────────────────────────────
    let n_tokens = 20;
    println!("\nBenchmarking {n_tokens} tokens with sync_on (current path)...");

    let start = Instant::now();
    for _tok in 0..n_tokens {
        let mut h: Arc<Tensor<f32>> = api::randn(0.0f32, 1.0, [d])
            .sync_on(&stream)?.into();

        for layer in &layers {
            // RMSNorm
            let h_2d: Arc<Tensor<f32>> = h.copy_sync(&stream)?.reshape([1, d]).into();
            let rms_out = api::zeros([1, d]).sync_on(&stream)?.partition([1, d as i32]);
            let (_, _, rms_out, _) = rms_norm_sync(
                h_2d, layer.norm_w.clone(), rms_out, eps,
            )
            .generics(rms_gen.clone())
            .sync_on(&stream)?;
            let norm: Arc<Tensor<f32>> = rms_out.unpartition().reshape([1, d]).into();

            // MatVec (Q projection)
            let q_out = api::zeros([d]).sync_on(&stream)?.partition([bn as i32]);
            let (q_out, _, _) = matvec_sync(
                q_out, norm.clone(), layer.wq.clone(),
            )
            .generics(mv_gen.clone())
            .sync_on(&stream)?;

            // MatVec (O projection) — pretend attention output = Q
            let q_2d: Arc<Tensor<f32>> = q_out.unpartition().reshape([1, d]).into();
            let o_out = api::zeros([d]).sync_on(&stream)?.partition([bn as i32]);
            let (o_out, _, _) = matvec_sync(
                o_out, q_2d, layer.wo.clone(),
            )
            .generics(mv_gen.clone())
            .sync_on(&stream)?;

            // Residual add
            let o_flat: Arc<Tensor<f32>> = o_out.unpartition().into();
            let add_out = api::zeros([d]).sync_on(&stream)?.partition([block as i32]);
            let (_, _, add_out) = add_sync(
                h.clone(), o_flat, add_out,
            )
            .generics(add_gen.clone())
            .sync_on(&stream)?;
            h = add_out.unpartition().into();
        }

        // "Logits" — just to consume h
        let _ = h.to_host_vec().sync_on(&stream)?;
    }
    let elapsed_sync = start.elapsed();
    let tps_sync = n_tokens as f64 / elapsed_sync.as_secs_f64();
    println!(
        "sync_on path: {n_tokens} tokens in {:.2}s = {tps_sync:.1} tok/s",
        elapsed_sync.as_secs_f64()
    );

    // ── Benchmark: execute path (no per-kernel sync) ────────────────────
    println!("\nBenchmarking {n_tokens} tokens with execute (async path)...");

    let start = Instant::now();
    let exec_ctx = cuda_async::device_operation::ExecutionContext::new(stream.clone());
    for _tok in 0..n_tokens {
        let mut h: Arc<Tensor<f32>> = Arc::new(unsafe {
            api::randn(0.0f32, 1.0, [d]).execute(&exec_ctx)
                .expect("randn failed")
        });

        for layer in &layers {
            // RMSNorm
            let h_2d: Arc<Tensor<f32>> = Arc::new(unsafe {
                api::copy(&h).execute(&exec_ctx).expect("copy").reshape([1, d])
            });
            let rms_out = unsafe {
                api::zeros([1, d]).execute(&exec_ctx).expect("zeros")
            }.partition([1, d as i32]);
            let (_, _, rms_out, _) = unsafe {
                rms_norm_sync(h_2d, layer.norm_w.clone(), rms_out, eps)
                    .generics(rms_gen.clone())
                    .execute(&exec_ctx)
                    .expect("rms_norm")
            };
            let norm: Arc<Tensor<f32>> = Arc::new(
                rms_out.unpartition().reshape([1, d])
            );

            // MatVec (Q)
            let q_out = unsafe {
                api::zeros([d]).execute(&exec_ctx).expect("zeros")
            }.partition([bn as i32]);
            let (q_out, _, _) = unsafe {
                matvec_sync(q_out, norm.clone(), layer.wq.clone())
                    .generics(mv_gen.clone())
                    .execute(&exec_ctx)
                    .expect("matvec")
            };

            // MatVec (O)
            let q_2d: Arc<Tensor<f32>> = Arc::new(unsafe {
                api::copy(&Arc::new(q_out.unpartition()))
                    .execute(&exec_ctx).expect("copy")
                    .reshape([1, d])
            });
            let o_out = unsafe {
                api::zeros([d]).execute(&exec_ctx).expect("zeros")
            }.partition([bn as i32]);
            let (o_out, _, _) = unsafe {
                matvec_sync(o_out, q_2d, layer.wo.clone())
                    .generics(mv_gen.clone())
                    .execute(&exec_ctx)
                    .expect("matvec")
            };

            // Add
            let o_flat: Arc<Tensor<f32>> = o_out.unpartition().into();
            let add_out = unsafe {
                api::zeros([d]).execute(&exec_ctx).expect("zeros")
            }.partition([block as i32]);
            let (_, _, add_out) = unsafe {
                add_sync(h.clone(), o_flat, add_out)
                    .generics(add_gen.clone())
                    .execute(&exec_ctx)
                    .expect("add")
            };
            h = add_out.unpartition().into();
        }

        // Single sync at end of token
        stream.synchronize().expect("sync");
        let _ = h.to_host_vec().sync_on(&stream).expect("download");
    }
    let elapsed_exec = start.elapsed();
    let tps_exec = n_tokens as f64 / elapsed_exec.as_secs_f64();
    println!(
        "execute path: {n_tokens} tokens in {:.2}s = {tps_exec:.1} tok/s",
        elapsed_exec.as_secs_f64()
    );

    println!("\nSpeedup: {:.2}x", tps_exec / tps_sync);
    println!("\n--- Bottleneck analysis ---");
    println!("Per token: {} layers × 4 kernel launches = {} total", n_layers, n_layers * 4);
    println!("sync_on path: {} stream.synchronize() calls per token", n_layers * 4);
    println!("execute path: 1 stream.synchronize() call per token");
    println!("\nWith a full LLaMA forward pass (~15 kernels/layer × 22 layers = 330 launches),");
    println!("the sync_on overhead dominates at ~10μs per sync × 330 = 3.3ms/token.");
    println!("\nRequested features:");
    println!("  1. CUDA graph capture/replay for the forward pass");
    println!("  2. Non-owning TensorView (zero-copy reshape without DeviceBox ownership)");
    println!("  3. Pre-allocated tensor pool (reuse buffers across tokens)");

    Ok(())
}

[elibol]

Just want to drop a note and let you know that I haven't forgotten about this :)

I ran some tests using your example code, and extended it with an unsafe CUDA graphs prototype. There's a bit more performance testing I need to do to make sure we're taking full advantage of the perf benefits offered by CUDA graphs. I'm not seeing a significant difference between just submitting work to a stream and syncing after each token (similar to your ExecutionContext example, but optimized to avoid device-wide syncs) vs. doing the same thing using CUDA graphs.

Right now, I'm working on getting out documentation that hopefully answers some of your questions. I should have that out today/tomorrow.

Once I have docs out, I'm going to switch back to CUDA graphs and begin working out the safe API that covers your use case. My sense is that we'd like to have an API that leans toward greater flexibility vs. rigid static checks. I'll post a proposed API once I have the impl worked out. Would be great if we could iterate on that together before I make more wide-spread changes.

In the meantime, I'm happy to add that as an unsafe example to the repo if you think it's something you could use in the short-term.

Let me know if that works for you! If your needs are more urgent, I can focus on getting an unsafe API + example out ASAP.

[jafioti]

@elibol Thank you! I really appreciate it. My use case isn't super urgent, so if you dont think an unsafe api is likely to stay then I wouldnt bother. I'd be happy to iterate on a safe API for cuda graphs!