rustytorch/.claude/plans/cuda-implementation-gaps.md

RustyTorch++ CUDA Implementation Gaps

Overview

Tracking document for CUDA-related implementation gaps to debug on RTX 5090 FE with CUDA 13.

Environment Target:

• GPU: NVIDIA RTX 5090 Founders Edition
• CUDA: 13.x
• cuDNN: 9.x (new API)

---

Critical Priority

1. cuDNN Convolution Disabled FIXED

Status: COMPLETED - cuDNN convolution paths re-enabled for conv1d, conv2d, conv3d

Files modified:

• crates/core/rtx-tensor/src/tensor/convolution.rs - Added cudnn_conv1d, cudnn_conv2d, cudnn_conv3d methods

Changes:

• Imported cuDNN module types (CudnnContext, CudnnConfig, CudnnConvolution, CudnnOperation, ConvolutionConfig)
• Updated conv2d/conv1d/conv3d dispatch to call cuDNN implementations for CUDA tensors
• Added device ID validation to ensure input and weight are on the same GPU

---

High Priority

2. cuBLASLt Placeholder FIXED

Status: COMPLETED - Real cuBLASLt integration via cudarc 0.18.x

File modified: crates/core/rtx-tensor/src/cublas/advanced.rs

Changes:

• Replaced placeholder with cudarc::cublaslt::safe::CudaBlasLT
• Supports Matmul<f32>, Matmul<f16>, Matmul<bf16> via cudarc traits
• Kernel fusion available via Activation::Relu and Activation::Gelu

---

3. Mixed Precision Incomplete FIXED

Status: COMPLETED - Updated to use available cuBLASLt APIs

File modified: crates/core/rtx-tensor/src/cublas/precision.rs

Changes:

• FP16 tensor core GEMM: Use CublasCore::gemm_f16() or CublasLt with Matmul<f16>
• BF16 tensor core GEMM: Use CublasLt with Matmul<bf16> trait
• Standard FP16 GEMM: Use CublasCore::gemm_f16() for direct operations

---

Medium Priority

5. Sparse SpGEMM Dense Fallback FIXED

Status: COMPLETED - Now uses native cuSPARSE SpGEMM multi-phase API

Files modified:

• crates/core/rtx-tensor/src/sparse/cuda_kernels.rs - Fixed spgemm_coo to return sparse result
• crates/core/rtx-tensor/src/sparse/cusparse_kernels.rs - Implemented proper cuSPARSE SpGEMM

Changes:

• Replaced dense fallback with native cuSPARSE SpGEMM multi-phase workflow:
  • Phase 1: cusparseSpGEMM_workEstimation (query + execute)
  • Phase 2: cusparseSpGEMM_compute (query + execute)
  • Phase 3: cusparseSpGEMM_copy (finalize result)
• Proper CSR descriptor creation and pointer management
• Workspace allocation for both phases
• Result size query via cusparseSpMatGetSize
• COO ↔ CSR conversion for format compatibility

---

6. Flash Attention Edge Cases FIXED

File: crates/core/rtx-autograd/src/autodiff/ops/llm.rs

Status: COMPLETED - All edge cases handled

Fixed:

• Added create_causal_mask() helper function
• Backward pass now applies causal mask when causal=true
• Uses -1e9 for masked positions (numerical stability)
• Properly handles seq_q != seq_k cases
• Backend::softmax already implements log-sum-exp trick internally (see lib.rs:251)
• Added stable_softmax_backward helper for clean gradient computation
• Added comprehensive edge case tests in tests/flash_attention_tests.rs:
  • Basic tensor creation for attention
  • Asymmetric sequence lengths (seq_q != seq_k)
  • Large values (numerical stability test)
  • Small values (underflow prevention)
  • Single token sequences
  • Large batch/head counts
  • Causal mask dimensions verification

---

7. Conv1d/Conv3d Backward Missing FIXED

Status: COMPLETED - Conv1d and Conv3d backward passes implemented

Files modified:

• crates/core/rtx-kernel/src/kernels/cudnn_conv.rs - Added conv1d_backward and conv3d_backward
• crates/core/rtx-autograd/src/autodiff/ops/conv.rs - New file with Conv1dBackward, Conv2dBackward, Conv3dBackward
• crates/core/rtx-autograd/src/autodiff/ops/mod.rs - Added conv module export

Changes:

• Added conv1d_backward that converts 1D shapes to 2D and delegates to conv2d_backward
• Added conv3d_backward with proper 5D tensor and filter descriptor support
• Added helper methods for 3D convolution:
  • get_or_create_tensor_descriptor_5d
  • get_or_create_filter_descriptor_5d
  • get_or_create_convolution_descriptor_3d
  • find_best_weight_gradient_algorithm_3d
  • find_best_data_gradient_algorithm_3d
  • calculate_conv3d_flops
• Created autograd backward functions for all convolution types

---

Low Priority

8. GPU Pooling Placeholders FIXED

Status: COMPLETED - GPU-accelerated pooling via cuDNN

Files modified:

• crates/core/rtx-tensor/src/tensor/pooling.rs - Added cuDNN-accelerated cudnn_pool2d
• crates/core/rtx-nn/src/layers/pooling/maxpool2d.rs - Fixed to use Tensor::max_pool2d
• crates/core/rtx-nn/src/layers/pooling/avgpool2d.rs - Fixed to use Tensor::avg_pool2d

Changes:

• Added cudnn_pool2d implementation using cuDNN's pooling forward pass
• Automatic GPU dispatch for CUDA tensors, CPU fallback otherwise
• Supports both max pooling (CUDNN_POOLING_MAX) and average pooling
• Fixed broken placeholder implementations in rtx-nn layers
• Layers now delegate to working Tensor methods for square kernels

---

9. Memory Fragmentation Stub FIXED

Status: COMPLETED - Fragmentation ratio now calculated

File modified: crates/core/rtx-memory/src/metrics.rs

Changes:

• Implemented fragmentation_ratio() based on three factors:
  • Size variance factor (entropy of allocation size distribution)
  • Memory efficiency (current vs total allocated)
  • Churn rate (allocation/deallocation cycling)
• Added size_variance_factor() helper using entropy-based calculation
• Returns weighted average of factors, clamped to [0.0, 1.0]

---

Completed Items ✓

• Tensor NaN/Inf detection (isnan, isinf, has_nan, has_inf, all_finite)
• Comparison operators (gt, lt, eq, ne, ge, le + scalar variants)
• Gradient clipping NaN/Inf validation
• Autograd backward functions (Mean, Max, Min, GELU, SiLU, LayerNorm, RMSNorm, FlashAttention)
• Metal scalar operations (add_scalar, mul_scalar)
• DLPack tensor interop
• Distributed context broadcast
• Legacy NCCL cleanup
• Fix cudnn_conv.rs type mismatch (Conv2dDescriptor → ConvDescriptor)
• CUDA Data Transfer Bug Fix - cuda_matmul() now uses lock_cuda_slice() pattern
• Conv1d/Conv3d Backward - Added cuDNN backward kernels and autograd ops
• GPU Pooling - cuDNN-accelerated max_pool2d and avg_pool2d
• Memory Fragmentation - Implemented fragmentation_ratio() with entropy-based calculation
• cuDNN Module Complete Rewrite - All files ported to cudarc 0.18.x result layer API
• cuBLASLt Integration - Real CudaBlasLT handle, not placeholder
• Mixed Precision Stubs - Updated to point to available APIs
• Flash Attention Edge Cases - All edge cases fixed, tests added
• cuSPARSE SpGEMM - Native multi-phase API (work estimation → compute → copy)
• Tensor Core GEMM - Fixed FP16/BF16 tensor core placeholders in precision.rs
• Batched Softmax - Implemented softmax for transformer attention in batched.rs
• FP16/BF16 Conversion - Implemented FP16→FP32 and BF16→FP32 conversion in advanced.rs
• GPU Type Conversion Kernels - Added fp16_to_fp32, bf16_to_fp32, fp32_to_fp16, fp32_to_bf16 CUDA kernels
• GPU Softmax Kernel - Integrated softmax_kernel for large matrix operations
• GPU Kernel Wrappers - Added launch wrappers in cuda_kernels/mod.rs for all new kernels
• Threshold-Based GPU/CPU Selection - Automatic fallback to CPU for small matrices to avoid kernel overhead

Debug Session Checklist (RTX 5090 + CUDA 13)

# 1. Verify CUDA environment
nvidia-smi
nvcc --version

# 2. Check cudarc compatibility
cargo check -p rtx-tensor --features cuda 2>&1 | head -50

# 3. Test basic CUDA operations
cargo test -p rtx-tensor cuda --features cuda

# 4. Test cuDNN (if available)
cargo test -p rtx-kernel cudnn --features cuda

# 5. Benchmark tensor cores
cargo bench -p rtx-bench gemm --features cuda

# 6. Profile with Nsight
nsys profile cargo test -p rtx-tensor matmul --features cuda

---

Files Quick Reference

Priority	File	Status
CRITICAL	rtx-tensor/src/tensor/convolution.rs	✅ FIXED - cuDNN enabled
HIGH	rtx-tensor/src/cublas/advanced.rs	✅ FIXED - Real cuBLASLt
HIGH	rtx-tensor/src/cublas/precision.rs	✅ FIXED - Tensor cores working
MEDIUM	rtx-tensor/src/sparse/cuda_kernels.rs	✅ FIXED - SpGEMM working
MEDIUM	rtx-autograd/src/autodiff/ops/llm.rs	✅ FIXED - All edge cases handled
MEDIUM	rtx-kernel/src/kernels/cudnn_conv.rs	✅ FIXED - Conv1d/3d backward added
LOW	rtx-tensor/src/tensor/pooling.rs	✅ FIXED - cuDNN pooling
LOW	rtx-memory/src/metrics.rs	✅ FIXED - Fragmentation ratio