gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(ml-supervised): resolve all 104 compile errors — clean build

Browse Source 
- mamba/mod.rs: ~90 errors fixed — _candle suffixed functions replaced,
  operator overloads→free functions, autograd→pseudo-gradients,
  checkpoint→JSON serialization
- gpu_tensor.rs: added gpu_eye, gpu_cat_dim0, gpu_stack_tensors
- TFT/SSD: unused imports cleaned, type mismatches fixed
- ml-core: GpuTensor algebra methods (17 new), cuda_compat.rs deleted,
  GpuVarStore::vars/all_vars/linear_xavier added

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-18 07:26:56 +01:00
parent dd62f3fcfd
commit 6d8ba0708c
14 changed files with 1059 additions and 535 deletions
Download Patch File Download Diff File Expand all files Collapse all files

565
crates/ml-core/src/cuda_autograd/gpu_tensor.rs
View File

@@ -4,11 +4,31 @@
//! wrapper — no autograd graph, no dtype dispatch, no device polymorphism.
//! All data is F32 on a single CUDA device.
use std::mem::ManuallyDrop;
use std::sync::Arc;
use cudarc::driver::{CudaSlice, CudaStream};
use cudarc::cublas::CudaBlas;
use cudarc::cublas::sys::cublasOperation_t;
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut};
use crate::MLError;
use crate::native_types::NativeDType;
// ── Pointer helpers (same pattern as linear.rs) ──────────────────────────
/// Extract raw CUDA device pointer from a CudaSlice (read-only).
fn raw_ptr(slice: &CudaSlice<f32>, stream: &CudaStream) -> u64 {
let (ptr, guard) = slice.device_ptr(stream);
let _no_drop = ManuallyDrop::new(guard);
ptr
}
/// Extract raw mutable CUDA device pointer from a CudaSlice.
fn raw_ptr_mut(slice: &mut CudaSlice<f32>, stream: &CudaStream) -> u64 {
let (ptr, guard) = slice.device_ptr_mut(stream);
let _no_drop = ManuallyDrop::new(guard);
ptr
}
/// Helper trait for scalar values that can be converted to f32.
pub trait ScalarValue {
@@ -288,6 +308,549 @@ impl GpuTensor {
))
}
// -----------------------------------------------------------------------
// Tensor algebra: element-wise ops, reductions, slicing.
// Most ops download to host, compute, re-upload. Marked for future CUDA kernels.
// -----------------------------------------------------------------------
/// Clone this tensor by copying GPU data to a new allocation.
///
/// Uses host-side roundtrip (same as `clone_gpu_tensor` in linear.rs).
/// TODO: CUDA kernel — use `cuMemcpyDtoDAsync` for zero-host clone.
pub fn gpu_clone(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
let n = self.numel();
let mut host = vec![0.0_f32; n];
stream.memcpy_dtoh(&self.data, &mut host).map_err(|e| {
MLError::ModelError(format!("gpu_clone DtoH: {e}"))
})?;
Self::from_host(&host, self.shape.clone(), stream)
}
/// No-op dtype cast — `GpuTensor` is always F32. Returns a clone.
pub fn to_dtype(&self, _dtype: NativeDType, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
self.gpu_clone(stream)
}
/// No-op detach — no autograd tape. Returns a clone.
pub fn detach(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
self.gpu_clone(stream)
}
/// Create a zero-filled tensor with the same shape as `self`.
pub fn zeros_like(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
Self::zeros(&self.shape, stream)
}
/// Create a tensor filled with a constant value.
pub fn full(shape: &[usize], value: f32, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
let n: usize = shape.iter().product();
let host = vec![value; n];
Self::from_host(&host, shape.to_vec(), stream)
}
/// Create a 1-D tensor with values in `[start, end)` stepping by `step`.
pub fn arange(start: f32, end: f32, step: f32, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
if step == 0.0 {
return Err(MLError::InvalidInput("arange: step cannot be zero".to_string()));
}
let n = ((end - start) / step).ceil().max(0.0) as usize;
let mut host = Vec::with_capacity(n);
let mut v = start;
for _ in 0..n {
host.push(v);
v += step;
}
let len = host.len();
Self::from_host(&host, vec![len], stream)
}
/// Element-wise addition. Shapes must match exactly.
/// TODO: CUDA kernel
pub fn add(&self, other: &Self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
if self.shape != other.shape {
return Err(MLError::DimensionMismatch {
expected: self.numel(),
actual: other.numel(),
});
}
let a = self.to_host(stream)?;
let b = other.to_host(stream)?;
let c: Vec<f32> = a.iter().zip(b.iter()).map(|(x, y)| x + y).collect();
Self::from_host(&c, self.shape.clone(), stream)
}
/// Element-wise subtraction. Shapes must match exactly.
/// TODO: CUDA kernel
pub fn sub(&self, other: &Self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
if self.shape != other.shape {
return Err(MLError::DimensionMismatch {
expected: self.numel(),
actual: other.numel(),
});
}
let a = self.to_host(stream)?;
let b = other.to_host(stream)?;
let c: Vec<f32> = a.iter().zip(b.iter()).map(|(x, y)| x - y).collect();
Self::from_host(&c, self.shape.clone(), stream)
}
/// Element-wise multiply. Shapes must match exactly.
/// TODO: CUDA kernel
pub fn mul(&self, other: &Self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
if self.shape != other.shape {
return Err(MLError::DimensionMismatch {
expected: self.numel(),
actual: other.numel(),
});
}
let a = self.to_host(stream)?;
let b = other.to_host(stream)?;
let c: Vec<f32> = a.iter().zip(b.iter()).map(|(x, y)| x * y).collect();
Self::from_host(&c, self.shape.clone(), stream)
}
/// Element-wise multiply with broadcasting.
///
/// Supports the common cases:
/// - Same shape: plain element-wise multiply.
/// - Scalar (numel=1) * tensor: broadcast the scalar.
/// - `[1, N]` * `[M, N]` or `[M, N]` * `[1, N]`: broadcast rows.
/// TODO: CUDA kernel
pub fn broadcast_mul(&self, other: &Self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
// Fast path: identical shapes
if self.shape == other.shape {
return self.mul(other, stream);
}
// Scalar broadcast
if other.numel() == 1 {
let scalar = other.to_host(stream)?;
let s = scalar.first().copied().unwrap_or(0.0);
let a = self.to_host(stream)?;
let c: Vec<f32> = a.iter().map(|x| x * s).collect();
return Self::from_host(&c, self.shape.clone(), stream);
}
if self.numel() == 1 {
let scalar = self.to_host(stream)?;
let s = scalar.first().copied().unwrap_or(0.0);
let b = other.to_host(stream)?;
let c: Vec<f32> = b.iter().map(|x| x * s).collect();
return Self::from_host(&c, other.shape.clone(), stream);
}
// Row broadcast: [1, N] * [M, N]
if self.ndim() == 2 && other.ndim() == 2 {
let (m_a, n_a) = (
self.shape.first().copied().unwrap_or(0),
self.shape.get(1).copied().unwrap_or(0),
);
let (m_b, n_b) = (
other.shape.first().copied().unwrap_or(0),
other.shape.get(1).copied().unwrap_or(0),
);
if n_a == n_b {
let a = self.to_host(stream)?;
let b = other.to_host(stream)?;
if m_a == 1 {
let c: Vec<f32> = (0..m_b * n_b)
.map(|i| {
let col = i % n_b;
b.get(i).copied().unwrap_or(0.0) * a.get(col).copied().unwrap_or(0.0)
})
.collect();
return Self::from_host(&c, other.shape.clone(), stream);
}
if m_b == 1 {
let c: Vec<f32> = (0..m_a * n_a)
.map(|i| {
let col = i % n_a;
a.get(i).copied().unwrap_or(0.0) * b.get(col).copied().unwrap_or(0.0)
})
.collect();
return Self::from_host(&c, self.shape.clone(), stream);
}
}
}
Err(MLError::ModelError(format!(
"broadcast_mul: incompatible shapes {:?} vs {:?}",
self.shape, other.shape
)))
}
/// Element-wise divide with broadcasting.
///
/// Supports the same broadcast patterns as `broadcast_mul`.
/// TODO: CUDA kernel
pub fn broadcast_div(&self, other: &Self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
// Fast path: identical shapes
if self.shape == other.shape {
let a = self.to_host(stream)?;
let b = other.to_host(stream)?;
let c: Vec<f32> = a.iter().zip(b.iter()).map(|(x, y)| x / y).collect();
return Self::from_host(&c, self.shape.clone(), stream);
}
// Scalar broadcast (divisor is scalar)
if other.numel() == 1 {
let scalar = other.to_host(stream)?;
let s = scalar.first().copied().unwrap_or(1.0);
let a = self.to_host(stream)?;
let c: Vec<f32> = a.iter().map(|x| x / s).collect();
return Self::from_host(&c, self.shape.clone(), stream);
}
if self.numel() == 1 {
let scalar = self.to_host(stream)?;
let s = scalar.first().copied().unwrap_or(0.0);
let b = other.to_host(stream)?;
let c: Vec<f32> = b.iter().map(|x| s / x).collect();
return Self::from_host(&c, other.shape.clone(), stream);
}
// Row broadcast for 2-D
if self.ndim() == 2 && other.ndim() == 2 {
let (m_a, n_a) = (
self.shape.first().copied().unwrap_or(0),
self.shape.get(1).copied().unwrap_or(0),
);
let (m_b, n_b) = (
other.shape.first().copied().unwrap_or(0),
other.shape.get(1).copied().unwrap_or(0),
);
if n_a == n_b {
let a = self.to_host(stream)?;
let b = other.to_host(stream)?;
if m_b == 1 {
// [M, N] / [1, N]
let c: Vec<f32> = (0..m_a * n_a)
.map(|i| {
let col = i % n_a;
a.get(i).copied().unwrap_or(0.0) / b.get(col).copied().unwrap_or(1.0)
})
.collect();
return Self::from_host(&c, self.shape.clone(), stream);
}
if m_a == 1 {
// [1, N] / [M, N]
let c: Vec<f32> = (0..m_b * n_b)
.map(|i| {
let col = i % n_b;
a.get(col).copied().unwrap_or(0.0) / b.get(i).copied().unwrap_or(1.0)
})
.collect();
return Self::from_host(&c, other.shape.clone(), stream);
}
}
}
Err(MLError::ModelError(format!(
"broadcast_div: incompatible shapes {:?} vs {:?}",
self.shape, other.shape
)))
}
/// Matrix multiplication via cuBLAS sgemm.
///
/// For 2-D tensors `[M, K]` @ `[K, N]` -> `[M, N]`.
/// For 1-D `[K]` @ `[K, N]` -> `[N]` (treated as `[1, K]` matmul, then squeeze).
pub fn matmul(
&self,
other: &Self,
cublas: &CudaBlas,
stream: &Arc<CudaStream>,
) -> Result<Self, MLError> {
let (m, k_a) = match self.ndim() {
1 => (1_usize, self.shape.first().copied().unwrap_or(0)),
2 => (
self.shape.first().copied().unwrap_or(0),
self.shape.get(1).copied().unwrap_or(0),
),
_ => {
return Err(MLError::ModelError(format!(
"matmul: lhs must be 1-D or 2-D, got {:?}",
self.shape
)));
}
};
let (k_b, n) = match other.ndim() {
2 => (
other.shape.first().copied().unwrap_or(0),
other.shape.get(1).copied().unwrap_or(0),
),
_ => {
return Err(MLError::ModelError(format!(
"matmul: rhs must be 2-D, got {:?}",
other.shape
)));
}
};
if k_a != k_b {
return Err(MLError::DimensionMismatch {
expected: k_a,
actual: k_b,
});
}
let mut c = Self::zeros(&[m, n], stream)?;
// cuBLAS column-major: C_col[N, M] = B_col[N, K] @ A_col[K, M]
// Row-major A[M,K] stored as col-major A_col[K,M] (transposed)
// So: transA=N (A_col is already [K,M]), transB=N for B_col[N,K]
// Wait — we have row-major data but cuBLAS expects col-major.
// Row-major C[M,N] = A[M,K] @ B[K,N]
// In col-major view: C_col[N,M] = B_col[N,K] @ A_col[K,M]
// transA=N, transB=N, m=N, n=M, k=K, lda=N, ldb=K, ldc=N
let a_ptr = raw_ptr(&self.data, stream);
let b_ptr = raw_ptr(&other.data, stream);
let c_ptr = raw_ptr_mut(&mut c.data, stream);
unsafe {
cudarc::cublas::result::sgemm(
*cublas.handle(),
cublasOperation_t::CUBLAS_OP_N, // transA (B in row-major)
cublasOperation_t::CUBLAS_OP_N, // transB (A in row-major)
n as i32, // m (cols of C)
m as i32, // n (rows of C)
k_a as i32, // k
&1.0_f32 as *const f32,
b_ptr as *const f32, // A in cuBLAS = B in row-major
n as i32, // lda = N
a_ptr as *const f32, // B in cuBLAS = A in row-major
k_a as i32, // ldb = K
&0.0_f32 as *const f32,
c_ptr as *mut f32,
n as i32, // ldc = N
)
.map_err(|e| MLError::ModelError(format!("cuBLAS sgemm matmul: {e:?}")))?;
}
// If lhs was 1-D, squeeze the leading dim
if self.ndim() == 1 {
c = c.reshape(vec![n])?;
}
Ok(c)
}
/// Add a dimension of size 1 at position `dim`.
pub fn unsqueeze(&self, dim: usize, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
if dim > self.ndim() {
return Err(MLError::ModelError(format!(
"unsqueeze: dim {dim} > ndim {}",
self.ndim()
)));
}
let mut new_shape = self.shape.clone();
new_shape.insert(dim, 1);
// Data is the same, just reshape
let cloned = self.gpu_clone(stream)?;
Ok(Self {
data: cloned.data,
shape: new_shape,
})
}
/// Remove a dimension of size 1 at position `dim`.
pub fn squeeze(&self, dim: usize, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
if dim >= self.ndim() {
return Err(MLError::ModelError(format!(
"squeeze: dim {dim} >= ndim {}",
self.ndim()
)));
}
if self.shape.get(dim).copied().unwrap_or(0) != 1 {
return Err(MLError::ModelError(format!(
"squeeze: dim {dim} has size {}, expected 1",
self.shape.get(dim).copied().unwrap_or(0)
)));
}
let mut new_shape = self.shape.clone();
new_shape.remove(dim);
if new_shape.is_empty() {
new_shape.push(1);
}
let cloned = self.gpu_clone(stream)?;
Ok(Self {
data: cloned.data,
shape: new_shape,
})
}
/// Slice along a dimension: extract `[start..start+len)` along `dim`.
///
/// Only dim=0 is currently optimized; other dims use host roundtrip.
/// TODO: CUDA kernel for general narrow.
pub fn narrow(
&self,
dim: usize,
start: usize,
len: usize,
stream: &Arc<CudaStream>,
) -> Result<Self, MLError> {
if dim >= self.ndim() {
return Err(MLError::ModelError(format!(
"narrow: dim {dim} >= ndim {}",
self.ndim()
)));
}
let dim_size = self.shape.get(dim).copied().unwrap_or(0);
if start + len > dim_size {
return Err(MLError::ModelError(format!(
"narrow: start({start}) + len({len}) > dim_size({dim_size})"
)));
}
let host = self.to_host(stream)?;
// Compute strides for row-major layout
let ndim = self.ndim();
let mut strides = vec![1_usize; ndim];
for i in (0..ndim.saturating_sub(1)).rev() {
let next = strides.get(i + 1).copied().unwrap_or(1);
let sh = self.shape.get(i + 1).copied().unwrap_or(1);
if let Some(s) = strides.get_mut(i) {
*s = next * sh;
}
}
let mut new_shape = self.shape.clone();
if let Some(d) = new_shape.get_mut(dim) {
*d = len;
}
let new_numel: usize = new_shape.iter().product();
let mut out = Vec::with_capacity(new_numel);
let stride = strides.get(dim).copied().unwrap_or(1);
// For dim=0 this is a contiguous slice
if dim == 0 {
let offset = start * stride;
let end = offset + len * stride;
for i in offset..end {
out.push(host.get(i).copied().unwrap_or(0.0));
}
} else {
// General case: iterate over all output elements
let mut idx = vec![0_usize; ndim];
for _ in 0..new_numel {
// Compute source flat index
let mut src_flat = 0;
for d in 0..ndim {
let coord = if d == dim {
idx.get(d).copied().unwrap_or(0) + start
} else {
idx.get(d).copied().unwrap_or(0)
};
src_flat += coord * strides.get(d).copied().unwrap_or(1);
}
out.push(host.get(src_flat).copied().unwrap_or(0.0));
// Increment multi-index
for d in (0..ndim).rev() {
if let Some(v) = idx.get_mut(d) {
*v += 1;
let limit = new_shape.get(d).copied().unwrap_or(1);
if *v < limit {
break;
}
*v = 0;
}
}
}
}
Self::from_host(&out, new_shape, stream)
}
/// Argmax along a dimension. Returns indices as `Vec<u32>`.
///
/// For a 2-D tensor `[M, N]` with `dim=1`, returns `M` indices in `[0..N)`.
/// TODO: CUDA kernel
pub fn argmax(&self, dim: usize, stream: &Arc<CudaStream>) -> Result<Vec<u32>, MLError> {
if dim >= self.ndim() {
return Err(MLError::ModelError(format!(
"argmax: dim {dim} >= ndim {}",
self.ndim()
)));
}
let host = self.to_host(stream)?;
// For 1-D, return single argmax
if self.ndim() == 1 {
let mut best_idx = 0_u32;
let mut best_val = f32::NEG_INFINITY;
for (i, v) in host.iter().enumerate() {
if *v > best_val {
best_val = *v;
best_idx = i as u32;
}
}
return Ok(vec![best_idx]);
}
// For 2-D with dim=1: argmax per row
if self.ndim() == 2 && dim == 1 {
let rows = self.shape.first().copied().unwrap_or(0);
let cols = self.shape.get(1).copied().unwrap_or(0);
let mut indices = Vec::with_capacity(rows);
for r in 0..rows {
let mut best_idx = 0_u32;
let mut best_val = f32::NEG_INFINITY;
for c in 0..cols {
let v = host.get(r * cols + c).copied().unwrap_or(f32::NEG_INFINITY);
if v > best_val {
best_val = v;
best_idx = c as u32;
}
}
indices.push(best_idx);
}
return Ok(indices);
}
// For 2-D with dim=0: argmax per column
if self.ndim() == 2 && dim == 0 {
let rows = self.shape.first().copied().unwrap_or(0);
let cols = self.shape.get(1).copied().unwrap_or(0);
let mut indices = Vec::with_capacity(cols);
for c in 0..cols {
let mut best_idx = 0_u32;
let mut best_val = f32::NEG_INFINITY;
for r in 0..rows {
let v = host.get(r * cols + c).copied().unwrap_or(f32::NEG_INFINITY);
if v > best_val {
best_val = v;
best_idx = r as u32;
}
}
indices.push(best_idx);
}
return Ok(indices);
}
Err(MLError::ModelError(format!(
"argmax: unsupported ndim={} dim={dim}",
self.ndim()
)))
}
/// Extract a single scalar from a 1-element tensor.
pub fn to_scalar(&self, stream: &Arc<CudaStream>) -> Result<f32, MLError> {
if self.numel() != 1 {
return Err(MLError::ModelError(format!(
"to_scalar: expected 1 element, got {}",
self.numel()
)));
}
let host = self.to_host(stream)?;
Ok(host.first().copied().unwrap_or(0.0))
}
/// Mean of all elements.
/// TODO: CUDA kernel
pub fn mean_all(&self, stream: &Arc<CudaStream>) -> Result<f32, MLError> {
let host = self.to_host(stream)?;
let n = host.len();
if n == 0 {
return Err(MLError::ModelError("mean_all: empty tensor".to_string()));
}
let sum: f64 = host.iter().map(|&v| v as f64).sum();
Ok((sum / n as f64) as f32)
}
/// Reshape without copying — returns error if total elements differ.
pub fn reshape(self, new_shape: Vec<usize>) -> Result<Self, MLError> {
let new_numel: usize = new_shape.iter().product();

31
crates/ml-core/src/cuda_autograd/var_store.rs
View File

@@ -132,6 +132,25 @@ impl GpuVarStore {
self.params.keys().map(|s| s.as_str()).collect()
}
/// List all parameter names (alias for `param_names`).
pub fn vars(&self) -> Vec<&str> {
self.param_names()
}
/// Return all parameters as `GpuTensor` values (cloned to new allocations).
pub fn all_vars(&self) -> Result<Vec<GpuTensor>, MLError> {
let mut out = Vec::with_capacity(self.params.len());
for (name, param) in &self.params {
let mut host = vec![0.0_f32; param.data.len()];
self.stream.memcpy_dtoh(&param.data, &mut host).map_err(|e| {
MLError::ModelError(format!("all_vars DtoH '{name}': {e}"))
})?;
let t = GpuTensor::from_host(&host, param.shape.clone(), &self.stream)?;
out.push(t);
}
Ok(out)
}
// ── Layer builders ────────────────────────────────────────────────
/// Create a linear layer with Xavier initialization and register its
@@ -161,6 +180,18 @@ impl GpuVarStore {
})
}
/// Alias for `linear()` — Xavier-initialized linear layer.
///
/// Provided for call sites that want the initialization scheme explicit in the name.
pub fn linear_xavier(
&mut self,
prefix: &str,
in_dim: usize,
out_dim: usize,
) -> Result<GpuLinear, MLError> {
self.linear(prefix, in_dim, out_dim)
}
/// Create a linear layer with near-zero Xavier initialization.
///
/// For distributional output heads where uniform softmax at init is desired.

9
crates/ml-core/src/cuda_compat.rs
View File

@@ -1,9 +0,0 @@
//! CUDA-compatible operations (legacy shim — most functionality moved to cuda_autograd).
//!
//! This module previously provided manual implementations of operations missing
//! in Candle's CUDA kernels. With the Candle elimination, these operations are
//! now available through `cuda_autograd::ActivationKernels` (sigmoid, etc.) and
//! `cuda_autograd::GpuLayerNorm` (layer normalization).
//!
//! The module is kept as an empty placeholder so that downstream `use crate::cuda_compat`
//! statements do not fail. Callers should migrate to the cuda_autograd equivalents.

1
crates/ml-core/src/lib.rs
View File

@@ -131,7 +131,6 @@ pub mod traits;
pub mod optimizers;
pub mod gradient_accumulation;
pub mod gradient_utils;
pub mod cuda_compat;
pub mod tensor_ops;
pub mod gpu;
pub mod safety;

69
crates/ml-supervised/src/gpu_tensor.rs
View File

@@ -1136,6 +1136,75 @@ pub fn gpu_broadcast_add_col(a: &GpuTensor, b: &GpuTensor) -> Result<GpuTensor,
GpuTensor::from_vec(result, &[rows, cols], &a.stream)
}
/// Identity matrix (2D, square).
pub fn gpu_eye(n: usize, stream: &Arc<CudaStream>) -> Result<GpuTensor, MLError> {
let mut data = vec![0.0_f32; n * n];
for i in 0..n {
if let Some(slot) = data.get_mut(i * n + i) {
*slot = 1.0;
}
}
GpuTensor::from_vec(data, &[n, n], stream)
}
/// Concatenate tensors along dimension 0 (batch dimension).
///
/// All tensors must have the same shape except for dimension 0.
/// Result shape: [sum(dim0), dim1, dim2, ...].
pub fn gpu_cat_dim0(tensors: &[GpuTensor], stream: &Arc<CudaStream>) -> Result<GpuTensor, MLError> {
if tensors.is_empty() {
return Err(MLError::InvalidInput("gpu_cat_dim0: empty list".to_owned()));
}
let first = tensors.first().ok_or_else(|| MLError::InvalidInput("empty".to_owned()))?;
let tail_shape: Vec<usize> = first.shape.iter().skip(1).copied().collect();
let tail_elems: usize = tail_shape.iter().product::<usize>().max(1);
let mut result = Vec::new();
let mut total_batch = 0_usize;
for t in tensors {
let t_tail: Vec<usize> = t.shape.iter().skip(1).copied().collect();
if t_tail != tail_shape {
return Err(MLError::DimensionMismatch {
expected: tail_elems,
actual: t_tail.iter().product::<usize>().max(1),
});
}
let host = t.to_vec()?;
result.extend_from_slice(&host);
total_batch += t.dim(0)?;
}
let mut out_shape = vec![total_batch];
out_shape.extend_from_slice(&tail_shape);
GpuTensor::from_vec(result, &out_shape, stream)
}
/// Stack tensors along a new leading dimension.
///
/// All tensors must have the same shape.
/// Result shape: [N, ...original_shape].
pub fn gpu_stack_tensors(tensors: &[GpuTensor], stream: &Arc<CudaStream>) -> Result<GpuTensor, MLError> {
if tensors.is_empty() {
return Err(MLError::InvalidInput("gpu_stack_tensors: empty list".to_owned()));
}
let first = tensors.first().ok_or_else(|| MLError::InvalidInput("empty".to_owned()))?;
let elem_shape = first.shape.clone();
let mut result = Vec::new();
for t in tensors {
if t.shape != elem_shape {
return Err(MLError::DimensionMismatch {
expected: first.numel(),
actual: t.numel(),
});
}
let host = t.to_vec()?;
result.extend_from_slice(&host);
}
let mut out_shape = vec![tensors.len()];
out_shape.extend_from_slice(&elem_shape);
GpuTensor::from_vec(result, &out_shape, stream)
}
/// Ones tensor.
pub fn gpu_ones(shape: &[usize], stream: &Arc<CudaStream>) -> Result<GpuTensor, MLError> {
let n: usize = shape.iter().product();

1
crates/ml-supervised/src/lib.rs
View File

@@ -16,7 +16,6 @@
#![allow(clippy::single_char_lifetime_names)] // 'a is idiomatic Rust
// Re-export shared types from ml-core
pub use ml_core::cuda_compat;
pub use ml_core::xavier_init;
// GPU tensor abstraction (cudarc CudaSlice + cuBLAS)

3
crates/ml-supervised/src/mamba/loss.rs
View File

@@ -6,9 +6,6 @@
//!
//! Uses local `GpuTensor` (cudarc-backed) instead of candle Tensor.
use std::sync::Arc;
use cudarc::driver::CudaStream;
use ml_core::{MLError, MLResult};
use crate::gpu_tensor::GpuTensor;

840
crates/ml-supervised/src/mamba/mod.rs
View File

File diff suppressed because it is too large Load Diff

2
crates/ml-supervised/src/mamba/scan_algorithms.rs
View File

@@ -10,7 +10,7 @@ use std::time::Instant;
use cudarc::driver::CudaStream;
use tracing::{debug, instrument};
use crate::gpu_tensor::{gpu_add, gpu_mul, gpu_scale, GpuTensor};
use crate::gpu_tensor::GpuTensor;
use crate::liquid::FixedPoint;
use ml_core::MLError;

8
crates/ml-supervised/src/mamba/ssd_layer.rs
View File

@@ -25,8 +25,8 @@ use tracing::instrument;
use super::{Mamba2Config, Mamba2State};
use crate::gpu_tensor::{
gpu_add, gpu_full, gpu_layer_norm, gpu_matmul, gpu_mul, gpu_ones,
gpu_relu, gpu_scale, gpu_sigmoid, gpu_softmax, gpu_transpose,
gpu_add, gpu_layer_norm, gpu_matmul, gpu_mul, gpu_ones,
gpu_scale, gpu_sigmoid, gpu_transpose,
GpuLinear, GpuTensor,
};
use ml_core::MLError;
@@ -196,7 +196,7 @@ impl SSDLayer {
fn state_space_transform(
&mut self,
input: &GpuTensor,
state: &mut Mamba2State,
_state: &mut Mamba2State,
) -> Result<GpuTensor, MLError> {
// Project input to state space
let state_input = self.state_projection.forward(input)?;
@@ -233,7 +233,7 @@ impl SSDLayer {
qkv.shape
)));
}
let rows = qkv.dim(0)?;
let _rows = qkv.dim(0)?;
let total_cols = qkv.dim(1)?;
let expected_cols = 3 * single_head_size;
if total_cols != expected_cols {

12
crates/ml-supervised/src/tft/mod.rs
View File

@@ -28,7 +28,7 @@ use std::time::{Instant, SystemTime};
use cudarc::driver::{CudaContext, CudaStream};
use crate::gpu_tensor::{
gpu_add, gpu_matmul, gpu_mean_all, gpu_scale, gpu_sqr, gpu_sub,
gpu_add,
GpuLinear, GpuTensor,
};
@@ -768,10 +768,9 @@ impl TemporalFusionTransformer {
let predictions = self.forward(&static_tensor, &hist_tensor, &fut_tensor)?;
// Compute quantile loss as scalar
let loss_tensor = self
let loss_val = self
.quantile_outputs
.quantile_loss(&predictions, &target_tensor)?;
let loss_val = gpu_mean_all(&loss_tensor)?;
if loss_val.is_finite() {
epoch_loss += loss_val as f64;
@@ -818,22 +817,21 @@ impl TemporalFusionTransformer {
let target_tensor = self.array_to_gpu_1d(targets)?;
let predictions = self.forward(&static_tensor, &hist_tensor, &fut_tensor)?;
let loss_tensor = self
let loss_val = self
.quantile_outputs
.quantile_loss(&predictions, &target_tensor)?;
let loss_val = gpu_mean_all(&loss_tensor)?;
total_loss += loss_val as f64;
}
Ok(total_loss / validation_data.len() as f64)
}
/// Compute quantile loss for training
/// Compute quantile loss for training (returns scalar f32)
pub fn compute_quantile_loss(
&self,
predictions: &GpuTensor,
targets: &GpuTensor,
) -> Result<GpuTensor, MLError> {
) -> Result<f32, MLError> {
self.quantile_outputs.quantile_loss(predictions, targets)
}

4
crates/ml-supervised/src/tft/quantile_outputs.rs
View File

@@ -9,8 +9,8 @@ use cudarc::driver::CudaStream;
use ml_core::MLError;
use crate::gpu_tensor::{
gpu_add, gpu_exp, gpu_log, gpu_add_scalar, gpu_mean_all, gpu_sub,
gpu_mul, gpu_scale, gpu_abs, GpuLinear, GpuTensor,
gpu_add, gpu_exp, gpu_log, gpu_add_scalar,
GpuLinear, GpuTensor,
};
/// Quantile output layer for uncertainty estimation

7
crates/ml-supervised/src/tft/temporal_attention.rs
View File

@@ -11,7 +11,7 @@ use cudarc::driver::CudaStream;
use tracing::instrument;
use crate::gpu_tensor::{
gpu_add, gpu_cat_dim1, gpu_layer_norm, gpu_matmul, gpu_scale, gpu_softmax, gpu_transpose,
gpu_add, gpu_layer_norm, gpu_matmul, gpu_scale, gpu_softmax, gpu_transpose,
GpuLinear, GpuTensor,
};
use ml_core::MLError;
@@ -98,6 +98,7 @@ pub struct TemporalSelfAttention {
layer_norm: CudaLayerNorm,
head_dim: usize,
last_attention_weights: RwLock<HashMap<String, f64>>,
#[allow(dead_code)]
stream: Arc<CudaStream>,
}
@@ -172,8 +173,8 @@ impl TemporalSelfAttention {
/// Core 2D attention forward: [N, hidden_dim] -> [N, hidden_dim]
fn forward_2d(&self, x: &GpuTensor) -> Result<GpuTensor, MLError> {
let n = x.dim(0)?;
let hidden_dim = x.dim(1)?;
let _n = x.dim(0)?;
let _hidden_dim = x.dim(1)?;
// QKV projection: [N, hidden_dim] -> [N, 3 * hidden_dim]
let qkv = self.qkv_projection.forward(x)?;

42
docs/superpowers/plans/2026-03-18-gpu-dedup-consolidation.md Normal file
View File

@@ -0,0 +1,42 @@
# GPU Implementation Deduplication Plan
> Post-candle-elimination cleanup. Consolidate duplicate GPU tensor, linear, activation, and optimizer implementations.
**Goal:** Single canonical implementation per GPU primitive in ml-core, with ml-ppo CudaLinear kept for its distinct ownership model.
---
## Phase 1: Delete ml-supervised GpuTensor/GpuLinear (CRITICAL)
- DELETE `crates/ml-supervised/src/gpu_tensor.rs` GpuTensor + GpuLinear definitions
- Re-export `ml_core::cuda_autograd::{GpuTensor, GpuLinear}` in ml-supervised
- Update all ml-supervised imports
- ~200 LOC reduction
## Phase 2: Consolidate activations into ml-core
- MERGE `crates/ml-ppo/src/cuda_nn/activations.rs` CUDA kernels into ml-core
- DELETE `crates/ml-ppo/src/cuda_nn/activations.rs`
- Add thin wrapper in ml-ppo returning CudaVec for PPO compat
- ~400 LOC reduction
## Phase 3: Merge AdamW kernels
- MERGE `crates/ml-ppo/src/cuda_nn/adam.rs` kernel source into ml-core GpuAdamW
- Keep wrapper in ml-ppo for flat param group interface
- ~200 LOC reduction
## Phase 4: GPU-ify host-side workarounds
Priority order:
1. `gpu_clone()` — cudaMemcpyDtoDAsync (trivial, HIGH frequency)
2. `add/sub/mul()` — 3 element-wise CUDA kernels (~10 lines each)
3. `broadcast_mul/div()` — strided CUDA kernel
4. `softmax` — standard CUDA softmax kernel for attention
5. `narrow(dim>0)` — 2D copy kernel
6. `argmax/mean_all` — CUB reductions (LOW frequency)
## Phase 5: Consolidate pointer helpers
- Define `raw_ptr()`/`raw_ptr_mut()` once in ml-core, re-export
- Delete 3 duplicate definitions