rustyverse/rustytorch
2
0
Fork 0
Code Issues Pull Requests Actions 28 Packages Projects Releases Wiki Activity
Files
3fd4fb7b8c56f93c1a38ee270f4c5870e4affc5e
rustytorch/wire_cuda.md
T
redclawsystems 4d88dc0584 Initial commit
2026-03-04 00:08:42 +00:00

812 lines
24 KiB
Markdown
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:
```rust
// 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)
```rust
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)
1. **Add kernel compilation to device.rs:**
```rust
// 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)
}
}
```
2. **Update basic.rs to use kernels:**
```rust
// 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)
```rust
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)
1. **Add cuBLAS to Cargo.toml:**
```toml
# Already available via cudarc, just need to enable:
[dependencies]
cudarc = { workspace = true, features = ["cublas"] }
```
2. **Update gemm.rs:**
```rust
// 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)
```
3. **Add cuBLAS handle caching to device.rs:**
```rust
// 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)
```rust
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)
```rust
// 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)
```rust
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
1. **Add dependency to Cargo.toml:**
```toml
[dependencies]
rtx-flash-attention = { workspace = true }
```
2. **Update attention.rs:**
```rust
// 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`
```rust
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`
```rust
//! 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
```bash
# 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
```bash
# 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
Reference in New Issue View Git Blame Copy Permalink