rustyverse/rustytorch

Fork 0

Files

T

redclawsystems 4d88dc0584 Initial commit

2026-03-04 00:08:42 +00:00

24 KiB

Raw Blame History

CUDA Backend & API Unification Implementation Guide

Purpose: Complete guide for wiring CUDA GPU kernels and API unification on a CUDA-capable node.

Prerequisites:

CUDA 12.0+ installed
cudarc 0.18.2 (already in workspace)
RTX 5090 ready (compute capability 9.0)

Overview

Phase	Task	Status	Estimated LOC
1.1	Wire element-wise CUDA kernels	Pending	~100
1.2	Use cuBLAS for GEMM	Pending	~80
1.3	Wire reduction CUDA kernels	Pending	~100
1.4	Integrate Flash Attention	Pending	~150
3.1	Unified error types	Pending	~50
3.2	Backend parity test suite	Pending	~200

Total estimated changes: ~680 LOC

Phase 1: CUDA Backend Kernel Wiring

Current Problem

All CUDA operations use host fallback - data is copied to CPU, computed there, and copied back:

// Current pattern in all ops (DEFEATS GPU ACCELERATION):
let a_host = device.stream().clone_dtoh(&*lhs.data)?;  // GPU -> CPU
let b_host = device.stream().clone_dtoh(&*rhs.data)?;  // GPU -> CPU
// ... CPU computation ...
device.stream().clone_htod(&result)?                    // CPU -> GPU

Existing CUDA Kernels (Ready to Wire)

Kernel File	Location	Functions
`element_wise.cu`	`rtx-tensor/src/cuda_kernels/`	add, mul, div, max, sin, cos
`vector_add.cu`	`rtx-tensor/src/cuda_kernels/`	optimized vector addition
`reduction.cu`	`rtx-cfd/src/kernels/cuda/`	sum, mean, max, min
`rtx-flash-attention`	`crates/training/rtx-flash-attention/`	Flash Attention CUDA impl

Task 1.1: Wire Element-Wise CUDA Kernels

File: crates/core/rtx-backend-cuda/src/ops/basic.rs

Current Implementation (Host Fallback)

pub fn add(lhs: &CudaTensorPrimitive<D>, rhs: &CudaTensorPrimitive<D>) -> Result<...> {
    let a_host = device.stream().clone_dtoh(&*lhs.data)?;
    let b_host = device.stream().clone_dtoh(&*rhs.data)?;
    let result: Vec<f32> = a_host.iter().zip(&b_host).map(|(a, b)| a + b).collect();
    device.stream().clone_htod(&result)?
}

Required Implementation (GPU Kernel)

Add kernel compilation to device.rs:

// In crates/core/rtx-backend-cuda/src/device.rs
use cudarc::driver::{CudaFunction, LaunchConfig};
use std::collections::HashMap;
use parking_lot::RwLock;

pub struct CudaDeviceWrapper {
    // ... existing fields ...
    kernel_cache: RwLock<HashMap<String, CudaFunction>>,
}

impl CudaDeviceWrapper {
    /// Get or compile a CUDA kernel
    pub fn get_or_compile_kernel(&self, name: &str) -> Result<CudaFunction, CudaBackendError> {
        // Check cache first
        if let Some(kernel) = self.kernel_cache.read().get(name) {
            return Ok(kernel.clone());
        }

        // Compile kernel from PTX
        let ptx = match name {
            "element_wise_add" => include_str!("../../../rtx-tensor/src/cuda_kernels/element_wise.ptx"),
            "element_wise_mul" => include_str!("../../../rtx-tensor/src/cuda_kernels/element_wise.ptx"),
            // ... other kernels
            _ => return Err(CudaBackendError::KernelNotFound(name.to_string())),
        };

        let module = self.device.load_ptx(ptx.into(), name, &[name])?;
        let kernel = module.get_func(name)?;

        self.kernel_cache.write().insert(name.to_string(), kernel.clone());
        Ok(kernel)
    }
}

Update basic.rs to use kernels:

// In crates/core/rtx-backend-cuda/src/ops/basic.rs
pub fn add<const D: usize>(
    lhs: &CudaTensorPrimitive<D>,
    rhs: &CudaTensorPrimitive<D>,
) -> Result<CudaTensorPrimitive<D>, CudaBackendError> {
    let device = &lhs.device;
    let numel = lhs.numel();

    // Allocate output
    let output = device.alloc_zeros::<f32>(numel)?;

    // Get compiled kernel
    let kernel = device.get_or_compile_kernel("element_wise_add")?;

    // Launch configuration
    let block_size = 256;
    let grid_size = (numel + block_size - 1) / block_size;
    let config = LaunchConfig {
        grid_dim: (grid_size as u32, 1, 1),
        block_dim: (block_size as u32, 1, 1),
        shared_mem_bytes: 0,
    };

    // Launch kernel
    unsafe {
        kernel.launch(config, (&lhs.data, &rhs.data, &output, numel as u32))?;
    }

    Ok(CudaTensorPrimitive::new(output, lhs.shape, device.clone()))
}

// Similar for: sub, mul, div, neg, exp, log, sin, cos, sqrt, abs

Operations to Update

Function	Kernel Name	Notes
`add`	`element_wise_add`	Binary op
`sub`	`element_wise_sub`	Binary op
`mul`	`element_wise_mul`	Binary op
`div`	`element_wise_div`	Binary op, handle div-by-zero
`neg`	`element_wise_neg`	Unary op
`exp`	`element_wise_exp`	Unary op
`log`	`element_wise_log`	Unary op
`sin`	`element_wise_sin`	Unary op
`cos`	`element_wise_cos`	Unary op
`sqrt`	`element_wise_sqrt`	Unary op
`abs`	`element_wise_abs`	Unary op

Task 1.2: Use cuBLAS for GEMM

File: crates/core/rtx-backend-cuda/src/ops/gemm.rs

Current Implementation (Host Fallback - O(n³) CPU loop)

pub fn matmul(...) -> Result<...> {
    let a_host = device.stream().clone_dtoh(&*lhs.data)?;
    let b_host = device.stream().clone_dtoh(&*rhs.data)?;
    // Naive triple-nested loop on CPU
    for i in 0..m {
        for j in 0..n {
            for k in 0..inner {
                result[i * n + j] += a[i * inner + k] * b[k * n + j];
            }
        }
    }
}

Required Implementation (cuBLAS)

Add cuBLAS to Cargo.toml:

# Already available via cudarc, just need to enable:
[dependencies]
cudarc = { workspace = true, features = ["cublas"] }

Update gemm.rs:

// In crates/core/rtx-backend-cuda/src/ops/gemm.rs
use cudarc::cublas::{CudaBlas, GemmConfig};
use cudarc::cublas::sys::cublasOperation_t;

pub fn matmul<const D: usize>(
    lhs: &CudaTensorPrimitive<D>,
    rhs: &CudaTensorPrimitive<D>,
) -> Result<CudaTensorPrimitive<D>, CudaBackendError> {
    let device = &lhs.device;

    // Get dimensions [batch..., M, K] x [batch..., K, N] -> [batch..., M, N]
    let lhs_shape = &lhs.shape;
    let rhs_shape = &rhs.shape;

    let m = lhs_shape[D - 2];
    let k = lhs_shape[D - 1];
    let n = rhs_shape[D - 1];

    // Allocate output
    let mut output_shape = lhs_shape.clone();
    output_shape[D - 1] = n;
    let numel: usize = output_shape.iter().product();
    let output = device.alloc_zeros::<f32>(numel)?;

    // Create cuBLAS handle (cached in device)
    let blas = device.get_cublas_handle()?;

    // For 2D matrices:
    if D == 2 {
        unsafe {
            blas.gemm(
                cublasOperation_t::CUBLAS_OP_N,  // op(B)
                cublasOperation_t::CUBLAS_OP_N,  // op(A)
                n as i32,                         // N
                m as i32,                         // M
                k as i32,                         // K
                &1.0f32,                          // alpha
                rhs.data.as_ptr(),                // B
                n as i32,                         // ldb
                lhs.data.as_ptr(),                // A
                k as i32,                         // lda
                &0.0f32,                          // beta
                output.as_mut_ptr(),              // C
                n as i32,                         // ldc
            )?;
        }
    } else {
        // Batched GEMM for higher dimensions
        let batch_size: usize = lhs_shape[..D-2].iter().product();
        unsafe {
            blas.gemm_strided_batched(
                cublasOperation_t::CUBLAS_OP_N,
                cublasOperation_t::CUBLAS_OP_N,
                n as i32, m as i32, k as i32,
                &1.0f32,
                rhs.data.as_ptr(), n as i32, (k * n) as i64,
                lhs.data.as_ptr(), k as i32, (m * k) as i64,
                &0.0f32,
                output.as_mut_ptr(), n as i32, (m * n) as i64,
                batch_size as i32,
            )?;
        }
    }

    Ok(CudaTensorPrimitive::new(output, output_shape, device.clone()))
}

// Also update: bmm (batched matmul)

Add cuBLAS handle caching to device.rs:

// In device.rs
use cudarc::cublas::CudaBlas;
use once_cell::sync::OnceCell;

pub struct CudaDeviceWrapper {
    // ... existing fields ...
    cublas: OnceCell<CudaBlas>,
}

impl CudaDeviceWrapper {
    pub fn get_cublas_handle(&self) -> Result<&CudaBlas, CudaBackendError> {
        self.cublas.get_or_try_init(|| {
            CudaBlas::new(self.stream.clone())
                .map_err(|e| CudaBackendError::CuBlas(e.to_string()))
        })
    }
}

Task 1.3: Wire Reduction CUDA Kernels

File: crates/core/rtx-backend-cuda/src/ops/reduction.rs

Current Implementation (Host Fallback)

pub fn sum(tensor: &CudaTensorPrimitive<D>) -> Result<...> {
    let data = device.stream().clone_dtoh(&*tensor.data)?;
    let sum: f32 = data.iter().sum();
    // Copy single value back to GPU
}

Required Implementation (GPU Kernels)

// In crates/core/rtx-backend-cuda/src/ops/reduction.rs

/// Parallel reduction sum using GPU kernel
pub fn sum<const D: usize>(
    tensor: &CudaTensorPrimitive<D>,
) -> Result<CudaTensorPrimitive<0>, CudaBackendError> {
    let device = &tensor.device;
    let numel = tensor.numel();

    // Two-phase reduction: first to blocks, then final reduction
    let block_size = 256;
    let grid_size = (numel + block_size - 1) / block_size;

    // Allocate intermediate results (one per block)
    let block_results = device.alloc_zeros::<f32>(grid_size)?;

    // Phase 1: Reduce within blocks
    let kernel = device.get_or_compile_kernel("reduce_sum_phase1")?;
    let config = LaunchConfig {
        grid_dim: (grid_size as u32, 1, 1),
        block_dim: (block_size as u32, 1, 1),
        shared_mem_bytes: block_size * std::mem::size_of::<f32>() as u32,
    };

    unsafe {
        kernel.launch(config, (&tensor.data, &block_results, numel as u32))?;
    }

    // Phase 2: Reduce block results to single value
    if grid_size > 1 {
        // Recursive reduction or final kernel
        let final_result = device.alloc_zeros::<f32>(1)?;
        let final_kernel = device.get_or_compile_kernel("reduce_sum_final")?;
        let final_config = LaunchConfig {
            grid_dim: (1, 1, 1),
            block_dim: (block_size.min(grid_size) as u32, 1, 1),
            shared_mem_bytes: grid_size * std::mem::size_of::<f32>() as u32,
        };

        unsafe {
            final_kernel.launch(final_config, (&block_results, &final_result, grid_size as u32))?;
        }

        Ok(CudaTensorPrimitive::new(final_result, [], device.clone()))
    } else {
        Ok(CudaTensorPrimitive::new(block_results, [], device.clone()))
    }
}

/// Reduction along specific dimensions
pub fn sum_dim<const D: usize>(
    tensor: &CudaTensorPrimitive<D>,
    dims: &[usize],
    keepdim: bool,
) -> Result<CudaTensorPrimitive<D>, CudaBackendError> {
    // Use strided reduction kernel
    let device = &tensor.device;
    let kernel = device.get_or_compile_kernel("reduce_sum_dim")?;

    // Calculate output shape
    let mut output_shape = tensor.shape.clone();
    for &dim in dims {
        output_shape[dim] = if keepdim { 1 } else { output_shape[dim] };
    }

    // ... kernel launch with dimension info ...
    todo!("Implement strided reduction")
}

// Similar for: mean, max, min, prod

Reduction Operations to Implement

Function	Kernel	Notes
`sum`	`reduce_sum`	Two-phase parallel reduction
`sum_dim`	`reduce_sum_dim`	Strided reduction along axis
`mean`	`reduce_mean`	sum / count
`mean_dim`	`reduce_mean_dim`	Strided mean
`max`	`reduce_max`	Parallel max with index tracking
`min`	`reduce_min`	Parallel min with index tracking
`prod`	`reduce_prod`	Parallel product

Task 1.4: Integrate Flash Attention

File: crates/core/rtx-backend-cuda/src/ops/attention.rs

Current Implementation (Host Fallback Reference)

pub fn flash_attention(...) -> Result<...> {
    let q_data = device.stream().clone_dtoh(&*query.data)?;
    let k_data = device.stream().clone_dtoh(&*key.data)?;
    let v_data = device.stream().clone_dtoh(&*value.data)?;
    // ... O(N²) CPU attention computation ...
}

Required Implementation

Add dependency to Cargo.toml:

[dependencies]
rtx-flash-attention = { workspace = true }

Update attention.rs:

// In crates/core/rtx-backend-cuda/src/ops/attention.rs
use rtx_flash_attention::{FlashAttention, FlashAttentionConfig};
use rtx_tensor::{Tensor, Device};

/// Convert CudaTensorPrimitive to rtx_tensor::Tensor
fn primitive_to_tensor<const D: usize>(primitive: &CudaTensorPrimitive<D>) -> Tensor {
    let data = primitive.to_vec();  // Copy to host temporarily
    let shape: Vec<usize> = primitive.shape.to_vec();
    let device = Device::Cuda(primitive.device.index());
    Tensor::from_slice(&data, &shape, &device)
        .expect("Failed to create Tensor from CudaTensorPrimitive")
}

/// Convert rtx_tensor::Tensor back to CudaTensorPrimitive
fn tensor_to_primitive_4d(
    tensor: &Tensor,
    device: &CudaDeviceWrapper,
) -> CudaTensorPrimitive<4> {
    let data = tensor.to_cpu().expect("Failed to copy tensor to CPU");
    let shape_vec = tensor.shape().dims();
    let shape: [usize; 4] = [shape_vec[0], shape_vec[1], shape_vec[2], shape_vec[3]];

    let cuda_data = device.alloc_copy(&data).expect("Failed to copy to GPU");
    CudaTensorPrimitive::new(cuda_data, shape, device.clone())
}

/// Flash Attention using optimized CUDA kernels
pub fn flash_attention(
    query: &CudaTensorPrimitive<4>,
    key: &CudaTensorPrimitive<4>,
    value: &CudaTensorPrimitive<4>,
    mask: Option<&CudaTensorPrimitive<4>>,
    scale: f32,
    causal: bool,
) -> Result<CudaTensorPrimitive<4>, CudaBackendError> {
    // Try optimized path
    match flash_attention_optimized(query, key, value, scale, causal) {
        Ok(output) => {
            if mask.is_some() && !causal {
                // Fall back for explicit masks
                return flash_attention_reference(query, key, value, mask, scale, causal);
            }
            Ok(output)
        }
        Err(_) => flash_attention_reference(query, key, value, mask, scale, causal),
    }
}

fn flash_attention_optimized(
    query: &CudaTensorPrimitive<4>,
    key: &CudaTensorPrimitive<4>,
    value: &CudaTensorPrimitive<4>,
    scale: f32,
    causal: bool,
) -> Result<CudaTensorPrimitive<4>, rtx_flash_attention::FlashError> {
    let config = FlashAttentionConfig::default()
        .with_softmax_scale(scale)
        .with_causal(causal);

    let attn = FlashAttention::new(config)?;

    let q_tensor = primitive_to_tensor(query);
    let k_tensor = primitive_to_tensor(key);
    let v_tensor = primitive_to_tensor(value);

    let (output, _lse) = attn.forward(&q_tensor, &k_tensor, &v_tensor)?;

    Ok(tensor_to_primitive_4d(&output, &query.device))
}

// Keep reference implementation for fallback
fn flash_attention_reference(...) -> Result<...> {
    // ... existing CPU implementation ...
}

Phase 3: API Unification

Task 3.1: Unified Error Types

File: crates/core/rtx-backend/src/error.rs

use thiserror::Error;

/// Unified backend error type for all GPU backends
#[derive(Error, Debug)]
pub enum BackendError {
    #[error("Device not found: index {0}")]
    DeviceNotFound(usize),

    #[error("Memory allocation failed: requested {requested} bytes, available {available}")]
    MemoryAllocation { requested: usize, available: usize },

    #[error("Kernel compilation failed: {0}")]
    KernelCompilation(String),

    #[error("Kernel execution failed: {0}")]
    KernelExecution(String),

    #[error("Synchronization failed")]
    Synchronization,

    #[error("Shape mismatch: {0}")]
    ShapeMismatch(String),

    #[error("Invalid operation: {0}")]
    InvalidOperation(String),

    #[cfg(feature = "cuda")]
    #[error("CUDA error: {0}")]
    Cuda(#[from] CudaBackendError),

    #[cfg(feature = "metal")]
    #[error("Metal error: {0}")]
    Metal(#[from] MetalBackendError),
}

/// Result type alias for backend operations
pub type BackendResult<T> = Result<T, BackendError>;

Task 3.2: Backend Parity Test Suite

File: tests/backend_parity_tests.rs

//! Backend parity tests - verify CUDA and Metal produce identical results

use rtx_tensor::{Tensor, Device};

/// Helper to compare tensors with tolerance
fn assert_tensors_close(a: &Tensor, b: &Tensor, rtol: f32, atol: f32) {
    let a_data = a.to_cpu().unwrap();
    let b_data = b.to_cpu().unwrap();

    assert_eq!(a_data.len(), b_data.len(), "Tensor sizes differ");

    for (i, (av, bv)) in a_data.iter().zip(&b_data).enumerate() {
        let diff = (av - bv).abs();
        let threshold = atol + rtol * bv.abs();
        assert!(
            diff <= threshold,
            "Mismatch at index {}: {} vs {} (diff: {}, threshold: {})",
            i, av, bv, diff, threshold
        );
    }
}

/// Create random tensor on specified device
fn random_tensor(shape: &[usize], device: &Device) -> Tensor {
    Tensor::randn(shape, device).unwrap()
}

// ============================================================================
// Element-wise Operations
// ============================================================================

#[test]
fn test_add_parity() {
    let cpu = Device::Cpu;
    let a_cpu = random_tensor(&[128, 256], &cpu);
    let b_cpu = random_tensor(&[128, 256], &cpu);
    let result_cpu = a_cpu.add(&b_cpu).unwrap();

    #[cfg(feature = "cuda")]
    {
        let cuda = Device::Cuda(0);
        let a_cuda = a_cpu.to_device(&cuda).unwrap();
        let b_cuda = b_cpu.to_device(&cuda).unwrap();
        let result_cuda = a_cuda.add(&b_cuda).unwrap();
        assert_tensors_close(&result_cpu, &result_cuda.to_device(&cpu).unwrap(), 1e-5, 1e-6);
    }

    #[cfg(feature = "metal")]
    {
        let metal = Device::Metal(0);
        let a_metal = a_cpu.to_device(&metal).unwrap();
        let b_metal = b_cpu.to_device(&metal).unwrap();
        let result_metal = a_metal.add(&b_metal).unwrap();
        assert_tensors_close(&result_cpu, &result_metal.to_device(&cpu).unwrap(), 1e-5, 1e-6);
    }
}

#[test]
fn test_mul_parity() {
    // Similar to add
}

#[test]
fn test_exp_parity() {
    // Similar pattern for unary ops
}

// ============================================================================
// Matrix Multiplication
// ============================================================================

#[test]
fn test_matmul_parity() {
    let cpu = Device::Cpu;
    let a = random_tensor(&[64, 128], &cpu);
    let b = random_tensor(&[128, 256], &cpu);
    let result_cpu = a.matmul(&b).unwrap();

    #[cfg(feature = "cuda")]
    {
        let cuda = Device::Cuda(0);
        let a_cuda = a.to_device(&cuda).unwrap();
        let b_cuda = b.to_device(&cuda).unwrap();
        let result_cuda = a_cuda.matmul(&b_cuda).unwrap();
        // Slightly looser tolerance for GEMM due to FP accumulation order
        assert_tensors_close(&result_cpu, &result_cuda.to_device(&cpu).unwrap(), 1e-4, 1e-5);
    }

    #[cfg(feature = "metal")]
    {
        let metal = Device::Metal(0);
        let a_metal = a.to_device(&metal).unwrap();
        let b_metal = b.to_device(&metal).unwrap();
        let result_metal = a_metal.matmul(&b_metal).unwrap();
        assert_tensors_close(&result_cpu, &result_metal.to_device(&cpu).unwrap(), 1e-4, 1e-5);
    }
}

#[test]
fn test_batched_matmul_parity() {
    let cpu = Device::Cpu;
    let a = random_tensor(&[8, 64, 128], &cpu);  // [batch, M, K]
    let b = random_tensor(&[8, 128, 256], &cpu); // [batch, K, N]
    let result_cpu = a.bmm(&b).unwrap();

    // ... CUDA and Metal tests ...
}

// ============================================================================
// Reductions
// ============================================================================

#[test]
fn test_sum_parity() {
    let cpu = Device::Cpu;
    let a = random_tensor(&[128, 256], &cpu);
    let result_cpu = a.sum().unwrap();

    #[cfg(feature = "cuda")]
    {
        let cuda = Device::Cuda(0);
        let a_cuda = a.to_device(&cuda).unwrap();
        let result_cuda = a_cuda.sum().unwrap();
        assert_tensors_close(&result_cpu, &result_cuda.to_device(&cpu).unwrap(), 1e-4, 1e-5);
    }
}

#[test]
fn test_mean_dim_parity() {
    let cpu = Device::Cpu;
    let a = random_tensor(&[32, 64, 128], &cpu);
    let result_cpu = a.mean(&[1], true).unwrap();  // Mean over dim 1, keepdim

    // ... CUDA and Metal tests ...
}

// ============================================================================
// Attention
// ============================================================================

#[test]
fn test_flash_attention_parity() {
    let cpu = Device::Cpu;

    // [batch, heads, seq_len, head_dim]
    let q = random_tensor(&[2, 8, 128, 64], &cpu);
    let k = random_tensor(&[2, 8, 128, 64], &cpu);
    let v = random_tensor(&[2, 8, 128, 64], &cpu);

    let scale = 1.0 / (64.0f32).sqrt();
    let result_cpu = flash_attention_reference(&q, &k, &v, scale, true);

    #[cfg(feature = "cuda")]
    {
        let cuda = Device::Cuda(0);
        let q_cuda = q.to_device(&cuda).unwrap();
        let k_cuda = k.to_device(&cuda).unwrap();
        let v_cuda = v.to_device(&cuda).unwrap();

        let result_cuda = flash_attention(&q_cuda, &k_cuda, &v_cuda, None, scale, true).unwrap();
        // Attention has more FP variance, use looser tolerance
        assert_tensors_close(&result_cpu, &result_cuda.to_device(&cpu).unwrap(), 1e-3, 1e-4);
    }
}

// ============================================================================
// Performance Benchmarks
// ============================================================================

#[test]
#[ignore]  // Run with: cargo test benchmark -- --ignored
fn benchmark_matmul_speedup() {
    use std::time::Instant;

    let sizes = [(128, 128), (512, 512), (1024, 1024), (2048, 2048)];

    for (m, n) in sizes {
        let cpu = Device::Cpu;
        let a_cpu = random_tensor(&[m, n], &cpu);
        let b_cpu = random_tensor(&[n, m], &cpu);

        // CPU baseline
        let start = Instant::now();
        for _ in 0..10 {
            let _ = a_cpu.matmul(&b_cpu).unwrap();
        }
        let cpu_time = start.elapsed().as_secs_f64() / 10.0;

        #[cfg(feature = "cuda")]
        {
            let cuda = Device::Cuda(0);
            let a_cuda = a_cpu.to_device(&cuda).unwrap();
            let b_cuda = b_cpu.to_device(&cuda).unwrap();

            // Warmup
            let _ = a_cuda.matmul(&b_cuda).unwrap();

            let start = Instant::now();
            for _ in 0..100 {
                let _ = a_cuda.matmul(&b_cuda).unwrap();
            }
            let cuda_time = start.elapsed().as_secs_f64() / 100.0;

            let speedup = cpu_time / cuda_time;
            println!("MatMul {}x{}: CPU={:.3}ms, CUDA={:.3}ms, Speedup={:.1}x",
                m, n, cpu_time * 1000.0, cuda_time * 1000.0, speedup);

            assert!(speedup > 10.0, "Expected at least 10x speedup for {}x{}", m, n);
        }
    }
}

Verification Checklist

Run these checks after implementation:

Phase 1 Verification

# Build CUDA backend
cargo build -p rtx-backend-cuda --features cuda

# Run basic tests
cargo test -p rtx-backend-cuda --features cuda

# Verify no host fallback in basic.rs
grep -n "clone_dtoh" crates/core/rtx-backend-cuda/src/ops/basic.rs
# Should return nothing (no host fallback)

# Verify cuBLAS in gemm.rs
grep -n "CudaBlas\|gemm" crates/core/rtx-backend-cuda/src/ops/gemm.rs
# Should show cuBLAS usage

# Verify Flash Attention integration
grep -n "FlashAttention" crates/core/rtx-backend-cuda/src/ops/attention.rs
# Should show FlashAttention usage

Phase 3 Verification

# Run parity tests
cargo test -p rtx-backend --features cuda,metal backend_parity

# Run benchmarks
cargo test -p rtx-backend --features cuda benchmark -- --ignored --nocapture

Summary

Task	File	Key Change
1.1	`ops/basic.rs`, `device.rs`	Kernel compilation + launch instead of host fallback
1.2	`ops/gemm.rs`, `device.rs`	cuBLAS handle caching + GEMM calls
1.3	`ops/reduction.rs`	Two-phase parallel reduction kernels
1.4	`ops/attention.rs`, `Cargo.toml`	rtx-flash-attention integration
3.1	`rtx-backend/src/error.rs`	Unified BackendError enum
3.2	`tests/backend_parity_tests.rs`	Cross-backend correctness tests

Expected Speedups:

Element-wise ops: 50-100x (memory bandwidth limited)
MatMul (cuBLAS): 100-500x for large matrices
Reductions: 20-50x
Flash Attention: 10-100x depending on sequence length

24 KiB Raw Blame History

CUDA Backend & API Unification Implementation Guide

Overview

Phase 1: CUDA Backend Kernel Wiring

Current Problem

Existing CUDA Kernels (Ready to Wire)

Task 1.1: Wire Element-Wise CUDA Kernels

Current Implementation (Host Fallback)

Required Implementation (GPU Kernel)

Operations to Update

Task 1.2: Use cuBLAS for GEMM

Current Implementation (Host Fallback - O(n³) CPU loop)

Required Implementation (cuBLAS)

Task 1.3: Wire Reduction CUDA Kernels

Current Implementation (Host Fallback)

Required Implementation (GPU Kernels)

Reduction Operations to Implement

Task 1.4: Integrate Flash Attention

Current Implementation (Host Fallback Reference)

Required Implementation

Phase 3: API Unification

Task 3.1: Unified Error Types

Task 3.2: Backend Parity Test Suite

Verification Checklist

Phase 1 Verification

Phase 3 Verification

Summary

24 KiB

Raw Blame History