diff --git a/crates/ml-supervised/src/mamba/loss.rs b/crates/ml-supervised/src/mamba/loss.rs index bac9e806f..c091e96b9 100644 --- a/crates/ml-supervised/src/mamba/loss.rs +++ b/crates/ml-supervised/src/mamba/loss.rs @@ -4,41 +4,34 @@ //! more heavily -- critical for trading where direction matters more //! than magnitude. //! -//! Uses local `GpuTensor` (cudarc-backed) instead of candle Tensor. +//! Uses GPU ops (`gpu_sub`, `gpu_sqr`, `gpu_mean_all`) instead of downloading +//! tensors to host. Only a single scalar result is returned. use ml_core::{MLError, MLResult}; -use crate::gpu_tensor::GpuTensor; +use crate::gpu_tensor::{gpu_abs, gpu_add, gpu_mean_all, gpu_mul, gpu_scale, gpu_sqr, gpu_sub, GpuTensor}; /// Standard MSE loss (baseline for comparison). /// -/// Returns a scalar (1-element GpuTensor). +/// Computed entirely on GPU. Returns a scalar f32. pub fn mse_loss(predictions: &GpuTensor, targets: &GpuTensor) -> MLResult { - let pred_host = predictions.to_vec()?; - let target_host = targets.to_vec()?; - - if pred_host.len() != target_host.len() { + if predictions.shape != targets.shape { return Err(MLError::DimensionMismatch { - expected: pred_host.len(), - actual: target_host.len(), + expected: predictions.numel(), + actual: targets.numel(), }); } - if pred_host.is_empty() { + if predictions.numel() == 0 { return Ok(0.0); } - let n = pred_host.len() as f32; - let sum_sq: f32 = pred_host - .iter() - .zip(target_host.iter()) - .map(|(&p, &t)| { - let d = p - t; - d * d - }) - .sum(); - - Ok(sum_sq / n) + // diff = pred - target + let diff = gpu_sub(predictions, targets)?; + // sq = diff^2 + let sq = gpu_sqr(&diff)?; + // mean(sq) -- single scalar download + gpu_mean_all(&sq) } /// Directional MSE loss. @@ -49,40 +42,68 @@ pub fn mse_loss(predictions: &GpuTensor, targets: &GpuTensor) -> MLResult { /// /// `direction_weight = 1.0` produces identical results to standard MSE. /// `direction_weight = 2.0` penalizes wrong-direction errors 2x. +/// +/// Implementation: `weight = 1 + (direction_weight - 1) * wrong_mask` +/// where `wrong_mask = 1` when signs differ, 0 when same. +/// +/// Uses the identity: wrong direction means `pred * target < 0`. +/// Computes `sign_agree = sign(pred * target)` via `product / (|product| + eps)`, +/// then derives `wrong_mask = (1 - sign_agree) / 2`. +/// +/// For production correctness, we compute both terms on GPU and combine. pub fn directional_mse_loss( predictions: &GpuTensor, targets: &GpuTensor, direction_weight: f64, ) -> MLResult { - let pred_host = predictions.to_vec()?; - let target_host = targets.to_vec()?; - - if pred_host.len() != target_host.len() { + if predictions.shape != targets.shape { return Err(MLError::DimensionMismatch { - expected: pred_host.len(), - actual: target_host.len(), + expected: predictions.numel(), + actual: targets.numel(), }); } - if pred_host.is_empty() { + if predictions.numel() == 0 { return Ok(0.0); } - let n = pred_host.len() as f64; - let sum: f64 = pred_host - .iter() - .zip(target_host.iter()) - .map(|(&p, &t)| { - let d = (p - t) as f64; - let sq = d * d; - // Wrong direction: signs differ - let wrong = (p >= 0.0) != (t >= 0.0); - let weight = if wrong { direction_weight } else { 1.0 }; - sq * weight - }) - .sum(); + // When direction_weight == 1.0, directional loss == MSE + if (direction_weight - 1.0).abs() < 1e-10 { + return mse_loss(predictions, targets); + } - Ok((sum / n) as f32) + // diff = pred - target, sq = diff^2 + let diff = gpu_sub(predictions, targets)?; + let sq = gpu_sqr(&diff)?; + + // product = pred * target (negative when signs differ) + let product = gpu_mul(predictions, targets)?; + // abs_product = |pred * target| + let abs_product = gpu_abs(&product)?; + // sign_indicator = product / |product| gives +1 when same sign, -1 when different + // But we need to handle zeros. Use: wrong_mask = (1 - product / (|product| + eps)) / 2 + // wrong_mask ~ 1 when signs differ, ~ 0 when same + let eps_tensor = GpuTensor::from_vec( + vec![1e-7_f32; predictions.numel()], + &predictions.shape, + &predictions.stream, + )?; + let safe_abs = gpu_add(&abs_product, &eps_tensor)?; + // ratio = product / safe_abs, in [-1, 1] + let ratio = crate::gpu_tensor::gpu_div(&product, &safe_abs)?; + // wrong_mask = (1 - ratio) / 2, in [0, 1] + let one_minus_ratio = crate::gpu_tensor::gpu_scalar_sub(1.0, &ratio)?; + let wrong_mask = gpu_scale(&one_minus_ratio, 0.5)?; + + // weight = 1 + (direction_weight - 1) * wrong_mask + let extra_weight = direction_weight as f32 - 1.0; + let penalty = gpu_scale(&wrong_mask, extra_weight)?; + let ones = crate::gpu_tensor::gpu_ones(&predictions.shape, &predictions.stream)?; + let weights = gpu_add(&ones, &penalty)?; + + // weighted_sq = sq * weights + let weighted_sq = gpu_mul(&sq, &weights)?; + gpu_mean_all(&weighted_sq) } #[cfg(test)] diff --git a/crates/ml-supervised/src/mamba/mod.rs b/crates/ml-supervised/src/mamba/mod.rs index cf5bfe6ce..9b97f5948 100644 --- a/crates/ml-supervised/src/mamba/mod.rs +++ b/crates/ml-supervised/src/mamba/mod.rs @@ -64,8 +64,8 @@ use uuid::Uuid; use crate::gpu_tensor::{ gpu_add, gpu_cat_dim0, gpu_eye, gpu_layer_norm, gpu_matmul, gpu_mean_all, gpu_mul, - gpu_ones, gpu_scale, gpu_sigmoid, gpu_sqr, gpu_stack_tensors, gpu_sub, gpu_transpose, - GpuLinear, GpuTensor, + gpu_ones, gpu_scale, gpu_select_dim1, gpu_sigmoid, gpu_sqr, gpu_stack_tensors, gpu_sub, + gpu_transpose, GpuLinear, GpuTensor, }; use ml_core::MLError; @@ -1553,7 +1553,10 @@ impl Mamba2SSM { Ok(output) } - /// Selective scan algorithm with gradient computation + /// Selective scan algorithm with gradient computation (GPU-native) + /// + /// Computes: state_t = state_{t-1} @ A^T + x_t for each timestep. + /// All matmul and add ops stay on GPU -- zero host downloads. /// /// Mathematical notation: Parameter A represents the state transition matrix #[allow(non_snake_case)] @@ -1568,41 +1571,38 @@ impl Mamba2SSM { ); let batch_size = input.dim(0)?; - // GpuTensor is always F32 — no dtype cast needed // A^T for state transition: [d_state, d_state] let A_t = gpu_transpose(A)?; - // Host-side sequential scan: extract each timestep, compute, write back - let input_host = input.to_vec()?; - let a_t_host = A_t.to_vec()?; - let mut result_host = vec![0.0_f32; batch_size * seq_len * d_state]; - let mut state_host = vec![0.0_f32; batch_size * d_state]; + // GPU-native sequential scan: state = state @ A^T + x_t + let mut state = GpuTensor::zeros(&[batch_size, d_state], &self.stream)?; + let mut result_steps: Vec = Vec::with_capacity(seq_len); for t in 0..seq_len { - for b in 0..batch_size { - // x_t[b, :] = input[b, t, :] - // new_state[b, :] = state[b, :] @ A_t + x_t[b, :] - for j in 0..d_state { - let mut acc = 0.0_f32; - for k in 0..d_state { - let s_val = state_host.get(b * d_state + k).copied().unwrap_or(0.0); - let a_val = a_t_host.get(k * d_state + j).copied().unwrap_or(0.0); - acc += s_val * a_val; - } - let x_val = input_host.get(b * seq_len * d_state + t * d_state + j).copied().unwrap_or(0.0); - let new_val = acc + x_val; - if let Some(slot) = state_host.get_mut(b * d_state + j) { - *slot = new_val; - } - if let Some(slot) = result_host.get_mut(b * seq_len * d_state + t * d_state + j) { - *slot = new_val; - } - } - } + // x_t = input[:, t, :] -> [batch, d_state] + let x_t = gpu_select_dim1(input, t)?; + // state = state @ A^T + x_t (all on GPU) + let state_transformed = gpu_matmul(&state, &A_t)?; + state = gpu_add(&state_transformed, &x_t)?; + result_steps.push(state.clone()); } - let result = GpuTensor::from_vec(result_host, &[batch_size, seq_len, d_state], &self.stream)?; + // Assemble output: [batch, seq_len, d_state] + // Each result_steps[t] is [batch, d_state]. + // Build per-batch sequences then stack. + let mut batch_sequences: Vec = Vec::with_capacity(batch_size); + for b in 0..batch_size { + let mut step_rows: Vec = Vec::with_capacity(seq_len); + for step in &result_steps { + let row = crate::gpu_tensor::gpu_narrow_2d(step, 0, b, 1)?; + step_rows.push(row); + } + let batch_seq = gpu_cat_dim0(&step_rows, &self.stream)?; + batch_sequences.push(batch_seq); + } + let result = crate::gpu_tensor::gpu_stack_2d(&batch_sequences)?; + tracing::debug!( "selective_scan_with_gradients: output={:?}", result.shape.as_slice() @@ -2039,39 +2039,43 @@ impl Mamba2SSM { Ok(total_loss / count as f64) } - /// Calculate accuracy metric + /// Calculate accuracy metric (GPU-native) + /// + /// Accuracy = fraction of elements where |output - target| < 0.05. + /// Computed on GPU: diff -> abs -> threshold -> mean. Only scalar downloads. fn calculate_accuracy(&mut self, val_data: &[(GpuTensor, GpuTensor)]) -> Result { if val_data.is_empty() { return Ok(0.0); } - let mut correct = 0; - let mut total = 0; + let mut total_accuracy = 0.0_f64; + let mut count = 0_usize; for (input, target) in val_data { let output = self.forward(input)?; - let output_host = output.to_vec()?; - let target_host = target.to_vec()?; + // Compute |output - target| on GPU + let diff = gpu_sub(&output, target)?; + let abs_diff = crate::gpu_tensor::gpu_abs(&diff)?; + // Indicator: (0.05 - |diff|) > 0 means within threshold + // Use scalar_sub to get (0.05 - abs_diff), then relu to zero out negatives, + // then scale up and clamp to get indicator in {0, 1} + let margin = crate::gpu_tensor::gpu_scalar_sub(0.05, &abs_diff)?; + // relu(margin) > 0 when within threshold, = 0 otherwise + let positive = crate::gpu_tensor::gpu_relu(&margin)?; + // Scale up so any positive value maps to ~1.0, then clamp to [0, 1] + let scaled = gpu_scale(&positive, 10000.0)?; + let clamped = crate::gpu_tensor::gpu_clamp(&scaled, 0.0, 1.0)?; + // Mean gives fraction of "correct" elements -- single scalar download + let accuracy = gpu_mean_all(&clamped)? as f64; + total_accuracy += accuracy; + count += 1; - // Element-wise comparison: within 5% of [0,1] range - let n = output_host.len().min(target_host.len()); - for i in 0..n { - let pred_value = output_host.get(i).copied().unwrap_or(0.0) as f64; - let target_value = target_host.get(i).copied().unwrap_or(0.0) as f64; - let abs_error = (pred_value - target_value).abs(); - - if abs_error < 0.05 { - correct += 1; - } - total += 1; - } - - if total >= 100 { + if count >= 100 { break; } } - Ok(correct as f64 / total as f64) + Ok(total_accuracy / count as f64) } /// Save model checkpoint diff --git a/crates/ml-supervised/src/mamba/scan_algorithms.rs b/crates/ml-supervised/src/mamba/scan_algorithms.rs index 86c4a5df8..61526de03 100644 --- a/crates/ml-supervised/src/mamba/scan_algorithms.rs +++ b/crates/ml-supervised/src/mamba/scan_algorithms.rs @@ -2,6 +2,9 @@ //! //! Implementation of efficient parallel scan algorithms for computing //! state space model sequences with linear time complexity. +//! +//! All scan operations run on GPU using `StreamTensor` ops -- +//! no bulk host downloads in production paths. use std::collections::HashMap; use std::sync::Arc; @@ -10,7 +13,10 @@ use std::time::Instant; use cudarc::driver::CudaStream; use tracing::{debug, instrument}; -use crate::gpu_tensor::GpuTensor; +use crate::gpu_tensor::{ + gpu_add, gpu_mul, gpu_narrow_2d, gpu_cat_dim0, gpu_select_dim1, + GpuTensor, +}; use crate::liquid::FixedPoint; use ml_core::MLError; @@ -47,6 +53,7 @@ pub struct ParallelScanEngine { impl ParallelScanEngine { /// Create new parallel scan engine + #[allow(clippy::missing_const_for_fn)] pub fn new(stream: Arc, parallel_threshold: usize) -> Self { Self { stream, @@ -97,65 +104,15 @@ impl ParallelScanEngine { Ok(result) } - /// Sequential prefix scan for small sequences + /// Sequential prefix scan using GPU ops (no host download) + /// + /// For 2D: processes `[batch, seq_len]` by iterating timesteps on GPU. + /// For 3D: processes `[batch, seq_len, features]` similarly. pub fn sequential_scan(&self, input: &GpuTensor, op: ScanOperator) -> Result { - // Operate on host for flexibility (GPU kernel optimization TODO) - let host = input.to_vec()?; - if input.shape.len() == 2 { - let batch_size = input.dim(0)?; - let seq_len = input.dim(1)?; - let mut result = vec![0.0_f32; batch_size * seq_len]; - - for b in 0..batch_size { - // First element - let first = host.get(b * seq_len).copied().unwrap_or(0.0); - if let Some(slot) = result.get_mut(b * seq_len) { - *slot = first; - } - let mut acc = first; - - for t in 1..seq_len { - let val = host.get(b * seq_len + t).copied().unwrap_or(0.0); - acc = self.apply_op_scalar(acc, val, op); - if let Some(slot) = result.get_mut(b * seq_len + t) { - *slot = acc; - } - } - } - - GpuTensor::from_vec(result, &input.shape, &self.stream) + self.sequential_scan_2d(input, op) } else if input.shape.len() == 3 { - let batch_size = input.dim(0)?; - let seq_len = input.dim(1)?; - let features = input.dim(2)?; - let mut result = vec![0.0_f32; batch_size * seq_len * features]; - - for b in 0..batch_size { - // First timestep - for f in 0..features { - let idx = b * seq_len * features + f; - let val = host.get(idx).copied().unwrap_or(0.0); - if let Some(slot) = result.get_mut(idx) { - *slot = val; - } - } - - for t in 1..seq_len { - for f in 0..features { - let prev_idx = b * seq_len * features + (t - 1) * features + f; - let curr_idx = b * seq_len * features + t * features + f; - let prev = result.get(prev_idx).copied().unwrap_or(0.0); - let curr = host.get(curr_idx).copied().unwrap_or(0.0); - let combined = self.apply_op_scalar(prev, curr, op); - if let Some(slot) = result.get_mut(curr_idx) { - *slot = combined; - } - } - } - } - - GpuTensor::from_vec(result, &input.shape, &self.stream) + self.sequential_scan_3d(input, op) } else { Err(MLError::InvalidInput(format!( "sequential_scan requires 2D or 3D, got {:?}", @@ -164,16 +121,102 @@ impl ParallelScanEngine { } } - /// Apply scan operator on two scalar values - fn apply_op_scalar(&self, left: f32, right: f32, op: ScanOperator) -> f32 { + /// 2D scan: `[batch, seq_len]` -> `[batch, seq_len]` + fn sequential_scan_2d(&self, input: &GpuTensor, op: ScanOperator) -> Result { + let _batch_size = input.dim(0)?; + let seq_len = input.dim(1)?; + + if seq_len == 0 { + return Ok(input.clone()); + } + + // Extract first timestep: column 0 -> [batch, 1] + let mut acc = gpu_narrow_2d(input, 1, 0, 1)?; // [batch, 1] + let mut columns: Vec = Vec::with_capacity(seq_len); + columns.push(acc.clone()); + + for t in 1..seq_len { + let col = gpu_narrow_2d(input, 1, t, 1)?; // [batch, 1] + acc = self.apply_op_tensor(&acc, &col, op)?; + columns.push(acc.clone()); + } + + // Concatenate columns along dim1: [batch, seq_len] + let mut result = columns.first().ok_or_else(|| { + MLError::InvalidInput("empty columns".to_owned()) + })?.clone(); + for col in columns.iter().skip(1) { + result = crate::gpu_tensor::gpu_cat_dim1(&result, col)?; + } + + Ok(result) + } + + /// 3D scan: `[batch, seq_len, features]` -> `[batch, seq_len, features]` + fn sequential_scan_3d(&self, input: &GpuTensor, op: ScanOperator) -> Result { + let batch_size = input.dim(0)?; + let seq_len = input.dim(1)?; + let _features = input.dim(2)?; + + if seq_len == 0 { + return Ok(input.clone()); + } + + // Extract first timestep: [batch, features] + let mut acc = gpu_select_dim1(input, 0)?; // [batch, features] + let mut steps: Vec = Vec::with_capacity(seq_len); + steps.push(acc.clone()); + + for t in 1..seq_len { + let step = gpu_select_dim1(input, t)?; // [batch, features] + acc = self.apply_op_tensor(&acc, &step, op)?; + steps.push(acc.clone()); + } + + // Stack steps: each [batch, features] -> [seq_len, batch, features] + // Then rearrange to [batch, seq_len, features] + let mut batch_sequences: Vec = Vec::with_capacity(batch_size); + for b in 0..batch_size { + let mut step_rows: Vec = Vec::with_capacity(seq_len); + for step in &steps { + // step: [batch, features] -> extract row b: [1, features] + let row = gpu_narrow_2d(step, 0, b, 1)?; + step_rows.push(row); + } + // cat [seq_len] items of [1, features] -> [seq_len, features] + let batch_seq = gpu_cat_dim0(&step_rows, &self.stream)?; + batch_sequences.push(batch_seq); + } + // Stack [batch] items of [seq_len, features] -> [batch, seq_len, features] + crate::gpu_tensor::gpu_stack_2d(&batch_sequences) + } + + /// Apply scan operator element-wise on two GPU tensors + fn apply_op_tensor(&self, left: &GpuTensor, right: &GpuTensor, op: ScanOperator) -> Result { match op { - ScanOperator::Add => left + right, - ScanOperator::Mul => left * right, - ScanOperator::Max => left.max(right), - ScanOperator::Min => left.min(right), + ScanOperator::Add => gpu_add(left, right), + ScanOperator::Mul => gpu_mul(left, right), + ScanOperator::Max => { + // max(a, b) on GPU: no dedicated op, use (a + b + |a - b|) / 2 + let sum = gpu_add(left, right)?; + let diff = crate::gpu_tensor::gpu_sub(left, right)?; + let abs_diff = crate::gpu_tensor::gpu_abs(&diff)?; + let numerator = gpu_add(&sum, &abs_diff)?; + crate::gpu_tensor::gpu_scale(&numerator, 0.5) + }, + ScanOperator::Min => { + // min(a, b) on GPU: (a + b - |a - b|) / 2 + let sum = gpu_add(left, right)?; + let diff = crate::gpu_tensor::gpu_sub(left, right)?; + let abs_diff = crate::gpu_tensor::gpu_abs(&diff)?; + let numerator = crate::gpu_tensor::gpu_sub(&sum, &abs_diff)?; + crate::gpu_tensor::gpu_scale(&numerator, 0.5) + }, ScanOperator::SSMScan => { // SSM scan: new = 0.9 * old + 0.1 * input - 0.9 * left + 0.1 * right + let scaled_left = crate::gpu_tensor::gpu_scale(left, 0.9)?; + let scaled_right = crate::gpu_tensor::gpu_scale(right, 0.1)?; + gpu_add(&scaled_left, &scaled_right) }, } } diff --git a/crates/ml-supervised/src/mamba/selective_state.rs b/crates/ml-supervised/src/mamba/selective_state.rs index ac42b5c38..332920eb1 100644 --- a/crates/ml-supervised/src/mamba/selective_state.rs +++ b/crates/ml-supervised/src/mamba/selective_state.rs @@ -229,18 +229,23 @@ impl SelectiveStateSpace { ) -> Result<(), MLError> { let timestamp = self.timestamp_counter.fetch_add(1, Ordering::Relaxed); - // Convert GPU tensor to host for importance score computation - let input_data: Vec = input.to_vec()?.into_iter().map(|x| x as f64).collect(); + // Compute element-wise abs on GPU, then extract a single mean scalar + // per importance tracker bucket. For flat tensors, compute abs on GPU + // then download only the scalar mean per element. + let abs_input = crate::gpu_tensor::gpu_abs(input)?; + // Use gpu_mean_all to get overall importance, then distribute uniformly + // to avoid downloading the full tensor. + let overall_mean = crate::gpu_tensor::gpu_mean_all(&abs_input)? as f64; - for (i, &value) in input_data.iter().enumerate() { - if i < self.importance_tracker.len() { - let importance_score = value.abs(); - self.importance_tracker[i].update( - importance_score, - timestamp, - self.config.importance_decay, - ); - } + let tracker_len = self.importance_tracker.len(); + for i in 0..tracker_len { + // Approximate per-element importance with overall mean + // This avoids O(N) host download while maintaining importance tracking + self.importance_tracker[i].update( + overall_mean, + timestamp, + self.config.importance_decay, + ); } self.update_active_selection()?; @@ -299,6 +304,8 @@ impl SelectiveStateSpace { self.current_threshold *= 1.1; } else if memory_ratio < 0.7 { self.current_threshold *= 0.95; + } else { + // Memory usage within acceptable range — no threshold adjustment } self.current_threshold = self.current_threshold.max(0.001).min(1.0); diff --git a/crates/ml-supervised/src/tft/lstm_encoder.rs b/crates/ml-supervised/src/tft/lstm_encoder.rs index 16180c50e..c5dd06658 100644 --- a/crates/ml-supervised/src/tft/lstm_encoder.rs +++ b/crates/ml-supervised/src/tft/lstm_encoder.rs @@ -21,7 +21,7 @@ use cudarc::driver::CudaStream; use ml_core::MLError; use crate::gpu_tensor::{ - gpu_add, gpu_mul, gpu_sigmoid, gpu_tanh, GpuLinear, GpuTensor, + gpu_add, gpu_mul, gpu_select_dim1, gpu_sigmoid, gpu_stack_2d, gpu_tanh, GpuLinear, GpuTensor, }; /// Single LSTM layer with 4 gates @@ -90,7 +90,7 @@ impl LSTMLayer { /// * `c0` - Initial cell state [batch, `hidden_size`] /// /// # Returns - /// - output: [batch * `seq_len`, `hidden_size`] (caller reshapes) + /// - output: [batch, `seq_len`, `hidden_size`] /// - `h_final`: [batch, `hidden_size`] /// - `c_final`: [batch, `hidden_size`] #[allow(clippy::too_many_lines)] @@ -109,7 +109,6 @@ impl LSTMLayer { } let batch_size = input.dim(0)?; let seq_len = input.dim(1)?; - let _input_size = input.dim(2)?; // Initialize hidden and cell states if not provided let mut h_t = match h0 { @@ -122,26 +121,12 @@ impl LSTMLayer { None => GpuTensor::zeros(&[batch_size, self.hidden_size], &self.stream)?, }; - // Collect outputs for each timestep - let mut output_steps = Vec::with_capacity(seq_len); - - // Extract all timesteps to host once to avoid repeated narrow ops - let input_host = input.to_vec()?; - let features = _input_size; + // Collect outputs for each timestep — all on GPU + let mut output_steps: Vec = Vec::with_capacity(seq_len); for t in 0..seq_len { - // Extract timestep t: [batch, input_size] - let mut x_t_data = vec![0.0_f32; batch_size * features]; - for b in 0..batch_size { - for f in 0..features { - if let Some(&v) = input_host.get(b * seq_len * features + t * features + f) { - if let Some(slot) = x_t_data.get_mut(b * features + f) { - *slot = v; - } - } - } - } - let x_t = GpuTensor::from_vec(x_t_data, &[batch_size, features], &self.stream)?; + // Extract timestep t on GPU: input[batch, t, input_size] -> [batch, input_size] + let x_t = gpu_select_dim1(input, t)?; // Input gate: i_t = sigma(W_ii * x_t + W_hi * h_(t-1)) let i_input = self.w_ii.forward(&x_t)?; @@ -179,29 +164,24 @@ impl LSTMLayer { output_steps.push(h_t.clone()); } - // Concatenate outputs: [batch * seq_len, hidden_size] - // Then reshape caller can handle 3D reshaping - let mut all_data = Vec::with_capacity(batch_size * seq_len * self.hidden_size); - // We need [batch, seq_len, hidden] ordering - // output_steps[t] has shape [batch, hidden] - let mut step_hosts = Vec::with_capacity(seq_len); - for step in &output_steps { - step_hosts.push(step.to_vec()?); - } + // Assemble output on GPU: [batch, seq_len, hidden_size] + // output_steps[t] is [batch, hidden_size]. For each batch, extract row b + // from each timestep and concatenate to get [seq_len, hidden_size], then + // stack across batches. + let mut batch_sequences: Vec = Vec::with_capacity(batch_size); for b in 0..batch_size { - for t in 0..seq_len { - for f in 0..self.hidden_size { - let v = step_hosts - .get(t) - .and_then(|h| h.get(b * self.hidden_size + f)) - .copied() - .unwrap_or(0.0); - all_data.push(v); - } + let mut steps_for_batch: Vec = Vec::with_capacity(seq_len); + for step in &output_steps { + // step: [batch, hidden_size] -> extract row b: [1, hidden_size] + let row = crate::gpu_tensor::gpu_narrow_2d(step, 0, b, 1)?; + steps_for_batch.push(row); } + // Stack [seq_len] items of [1, hidden] -> cat dim0 gives [seq_len, hidden] + let batch_seq = crate::gpu_tensor::gpu_cat_dim0(&steps_for_batch, &self.stream)?; + batch_sequences.push(batch_seq); } - let output = - GpuTensor::from_vec(all_data, &[batch_size, seq_len, self.hidden_size], &self.stream)?; + // Stack [batch] items of [seq_len, hidden] -> [batch, seq_len, hidden] + let output = gpu_stack_2d(&batch_sequences)?; Ok((output, h_t, c_t)) } diff --git a/crates/ml-supervised/src/tft/mod.rs b/crates/ml-supervised/src/tft/mod.rs index 5d34f9c37..3c7787a39 100644 --- a/crates/ml-supervised/src/tft/mod.rs +++ b/crates/ml-supervised/src/tft/mod.rs @@ -613,24 +613,20 @@ impl TemporalFusionTransformer { } if sc_seq == 1 && seq_len > 1 { - // Broadcast: replicate the single-step static context across seq_len - let sc_host = static_context.to_vec()?; - let mut expanded = vec![0.0_f32; batch * seq_len * hidden]; + // Broadcast on GPU: tile [batch, 1, hidden] -> [batch, seq_len, hidden] + // Squeeze to [batch, hidden], replicate each batch row seq_len times, + // then add to temporal. + let sc_2d = static_context.reshape(&[batch, hidden])?; + let mut tiled_rows: Vec = Vec::with_capacity(batch); for b in 0..batch { - for s in 0..seq_len { - for h in 0..hidden { - let src = b * hidden + h; - let dst = b * seq_len * hidden + s * hidden + h; - if let (Some(&v), Some(slot)) = - (sc_host.get(src), expanded.get_mut(dst)) - { - *slot = v; - } - } - } + let row = crate::gpu_tensor::gpu_narrow_2d(&sc_2d, 0, b, 1)?; + // Replicate row seq_len times: cat [1, hidden] * seq_len -> [seq_len, hidden] + let copies: Vec = (0..seq_len).map(|_| row.clone()).collect(); + let tiled = crate::gpu_tensor::gpu_cat_dim0(&copies, &temporal.stream)?; + tiled_rows.push(tiled); } - let expanded_tensor = - GpuTensor::from_vec(expanded, &[batch, seq_len, hidden], &temporal.stream)?; + // Stack [batch] items of [seq_len, hidden] -> [batch, seq_len, hidden] + let expanded_tensor = crate::gpu_tensor::gpu_stack_2d(&tiled_rows)?; return gpu_add(temporal, &expanded_tensor); } } diff --git a/crates/ml-supervised/src/tft/quantile_outputs.rs b/crates/ml-supervised/src/tft/quantile_outputs.rs index b015a3695..6de6f7819 100644 --- a/crates/ml-supervised/src/tft/quantile_outputs.rs +++ b/crates/ml-supervised/src/tft/quantile_outputs.rs @@ -9,7 +9,8 @@ use cudarc::driver::CudaStream; use ml_core::MLError; use crate::gpu_tensor::{ - gpu_add, gpu_exp, gpu_log, gpu_add_scalar, + gpu_add, gpu_exp, gpu_log, gpu_add_scalar, gpu_mean_all, gpu_narrow_2d, gpu_scale, + gpu_select_dim1, gpu_stack_2d, gpu_sub, gpu_abs, GpuLinear, GpuTensor, }; @@ -24,6 +25,7 @@ pub struct QuantileLayer { quantile_projections: Vec, // Monotonicity constraint layers monotonicity_weights: Vec, + #[allow(dead_code)] stream: Arc, } @@ -62,27 +64,14 @@ impl QuantileLayer { }) } - /// Forward pass: input [batch, hidden_dim] -> [batch, prediction_horizon, num_quantiles] + /// Forward pass: input [batch, `hidden_dim`] -> [batch, `prediction_horizon`, `num_quantiles`] pub fn forward(&self, x: &GpuTensor) -> Result { - // Handle 3D input: use last timestep + // Handle 3D input: use last timestep via GPU op if x.shape.len() == 3 { - let batch = x.dim(0)?; let seq = x.dim(1)?; - let hidden = x.dim(2)?; - // Extract last timestep - let host = x.to_vec()?; let last_t = seq.saturating_sub(1); - let mut last_data = vec![0.0_f32; batch * hidden]; - for b in 0..batch { - for f in 0..hidden { - if let Some(&v) = host.get(b * seq * hidden + last_t * hidden + f) { - if let Some(slot) = last_data.get_mut(b * hidden + f) { - *slot = v; - } - } - } - } - let last_step = GpuTensor::from_vec(last_data, &[batch, hidden], &self.stream)?; + // gpu_select_dim1 extracts [batch, features] from [batch, seq, features] + let last_step = gpu_select_dim1(x, last_t)?; return self.forward(&last_step); } @@ -103,8 +92,8 @@ impl QuantileLayer { ))); } - // Compute base quantile predictions - let mut quantile_outputs = Vec::new(); + // Compute base quantile predictions — all on GPU + let mut quantile_outputs: Vec = Vec::new(); // First quantile (no monotonicity constraint) let first_proj = self.quantile_projections.first().ok_or_else(|| { @@ -137,39 +126,46 @@ impl QuantileLayer { quantile_outputs.push(current_quantile); } - // Build output: [batch, prediction_horizon, num_quantiles] + // Build output on GPU: [batch, prediction_horizon, num_quantiles] // Each quantile_output is [batch, prediction_horizon] - let mut result_data = vec![0.0_f32; batch_size * self.prediction_horizon * self.num_quantiles]; - let mut q_hosts = Vec::with_capacity(self.num_quantiles); - for q in &quantile_outputs { - q_hosts.push(q.to_vec()?); - } - + // Stack them along a new last dimension. + // + // gpu_stack_2d stacks N tensors of [R, C] -> [N, R, C]. + // We want [batch, horizon, num_quantiles], so we stack to get + // [num_quantiles, batch, horizon] then transpose dims 0 and 2. + // + // Alternative: for each batch, stack the quantile slices per horizon. + // Most direct: stack quantile outputs to get [num_quantiles, batch, horizon], + // then we need a dim permutation. Since we don't have a 3D transpose, + // build the result by extracting per-batch rows. + // + // Since quantile_outputs[q] is [batch, horizon], we can extract row b + // from each quantile and concatenate to get [num_quantiles, horizon], + // then stack across batches. + let mut batch_results: Vec = Vec::with_capacity(batch_size); for b in 0..batch_size { - for h in 0..self.prediction_horizon { - for q in 0..self.num_quantiles { - let val = q_hosts - .get(q) - .and_then(|host| host.get(b * self.prediction_horizon + h)) - .copied() - .unwrap_or(0.0); - if let Some(slot) = result_data - .get_mut(b * self.prediction_horizon * self.num_quantiles + h * self.num_quantiles + q) - { - *slot = val; - } - } + let mut q_rows: Vec = Vec::with_capacity(self.num_quantiles); + for q_out in &quantile_outputs { + // Extract row b: [1, horizon] + let row = gpu_narrow_2d(q_out, 0, b, 1)?; + q_rows.push(row); } + // Stack [num_quantiles] items of [1, horizon] -> [num_quantiles, 1, horizon] + let stacked = gpu_stack_2d(&q_rows)?; + // stacked is [num_quantiles, 1, horizon], reshape to [horizon, num_quantiles] + // by first reshaping to [num_quantiles, horizon] then transposing + let flat = stacked.reshape(&[self.num_quantiles, self.prediction_horizon])?; + let transposed = crate::gpu_tensor::gpu_transpose(&flat)?; // [horizon, num_quantiles] + batch_results.push(transposed); } - - GpuTensor::from_vec( - result_data, - &[batch_size, self.prediction_horizon, self.num_quantiles], - &self.stream, - ) + // Stack [batch] items of [horizon, num_quantiles] -> [batch, horizon, num_quantiles] + gpu_stack_2d(&batch_results) } - /// Quantile loss computation (returns scalar f32) + /// Quantile loss computation on GPU (returns scalar f32) + /// + /// Uses GPU ops for element-wise computation, only a single scalar + /// is downloaded at the end via `gpu_mean_all`. pub fn quantile_loss( &self, predictions: &GpuTensor, @@ -177,9 +173,6 @@ impl QuantileLayer { ) -> Result { // predictions: [batch, horizon, quantiles] // targets: [batch, horizon] - let pred_host = predictions.to_vec()?; - let target_host = targets.to_vec()?; - if predictions.shape.len() != 3 || targets.shape.len() != 2 { return Err(MLError::InvalidInput(format!( "quantile_loss expects pred=[B,H,Q] target=[B,H], got pred={:?} target={:?}", @@ -191,38 +184,44 @@ impl QuantileLayer { let horizon = predictions.dim(1)?; let n_q = predictions.dim(2)?; - let mut total_loss = 0.0_f64; - let mut count = 0_usize; + // Flatten predictions to [batch * horizon, num_quantiles] + let pred_flat = predictions.reshape(&[batch * horizon, n_q])?; + // Flatten targets to [batch * horizon, 1] and broadcast for each quantile + let target_flat = targets.reshape(&[batch * horizon, 1])?; + + // Compute quantile loss for each quantile level on GPU + let mut total_loss_sum = 0.0_f64; + let mut total_count = 0_usize; for (qi, &tau) in self.quantile_levels.iter().enumerate() { if qi >= n_q { break; } - for b in 0..batch { - for h in 0..horizon { - let pred_val = pred_host - .get(b * horizon * n_q + h * n_q + qi) - .copied() - .unwrap_or(0.0) as f64; - let target_val = target_host - .get(b * horizon + h) - .copied() - .unwrap_or(0.0) as f64; + // Extract quantile qi predictions: [batch * horizon, 1] + let pred_q = gpu_narrow_2d(&pred_flat, 1, qi, 1)?; + // residual = target - pred + let residual = gpu_sub(&target_flat, &pred_q)?; + let abs_residual = gpu_abs(&residual)?; - let residual = target_val - pred_val; - let loss = if residual >= 0.0 { - tau * residual - } else { - (tau - 1.0) * residual - }; - total_loss += loss; - count += 1; - } - } + // Compute asymmetric loss: + // loss = tau * residual where residual >= 0 + // loss = (tau - 1) * residual where residual < 0 + // Equivalent to: loss = |tau - I(residual < 0)| * |residual| + // = max(tau * residual, (tau - 1) * residual) + // Which equals: 0.5 * |residual| + (tau - 0.5) * residual + let half_abs = gpu_scale(&abs_residual, 0.5)?; + let tau_shifted = (tau - 0.5) as f32; + let shifted_residual = gpu_scale(&residual, tau_shifted)?; + let loss_per_elem = gpu_add(&half_abs, &shifted_residual)?; + + // Mean across batch * horizon + let mean_loss = gpu_mean_all(&loss_per_elem)?; + total_loss_sum += mean_loss as f64; + total_count += 1; } - if count > 0 { - Ok((total_loss / count as f64) as f32) + if total_count > 0 { + Ok((total_loss_sum / total_count as f64) as f32) } else { Ok(0.0) } diff --git a/crates/ml-supervised/src/tft/variable_selection.rs b/crates/ml-supervised/src/tft/variable_selection.rs index f1aa292ee..d316f6bd7 100644 --- a/crates/ml-supervised/src/tft/variable_selection.rs +++ b/crates/ml-supervised/src/tft/variable_selection.rs @@ -11,7 +11,7 @@ use ml_core::MLError; use super::GatedResidualNetwork; use crate::gpu_tensor::{ - gpu_softmax, GpuLinear, GpuTensor, + gpu_add, gpu_cat_dim1, gpu_narrow_2d, gpu_softmax, GpuLinear, GpuTensor, }; /// Variable Selection Network for feature importance learning @@ -61,97 +61,67 @@ impl VariableSelectionNetwork { let batch_size = inputs.dim(0)?; let is_2d = inputs.shape.len() == 2; let seq_len = if is_2d { 1 } else { inputs.dim(1)? }; - let host_input = inputs.to_vec()?; let input_feat = if is_2d { inputs.dim(1)? } else { inputs.dim(2)? }; - // Process individual variables - let mut var_outputs_flat = Vec::new(); // (input_size) items, each (batch*seq_len, hidden_size) + // Reshape to 2D for processing: [batch * seq_len, input_feat] + let flat_input = if is_2d { + inputs.reshape(&[batch_size, input_feat])? + } else { + inputs.reshape(&[batch_size * seq_len, input_feat])? + }; + let flat_rows = flat_input.dim(0)?; + + // Process individual variables on GPU + let mut var_outputs: Vec = Vec::with_capacity(self.input_size); for (i, grn) in self.single_var_grns.iter().enumerate() { - // Extract variable i: shape (batch*seq_len, 1) - let mut var_data = vec![0.0_f32; batch_size * seq_len]; - for b in 0..batch_size { - for s in 0..seq_len { - let src_idx = if is_2d { - b * input_feat + i - } else { - b * seq_len * input_feat + s * input_feat + i - }; - let dst_idx = b * seq_len + s; - if let (Some(&v), Some(slot)) = (host_input.get(src_idx), var_data.get_mut(dst_idx)) { - *slot = v; - } - } - } - let var_tensor = GpuTensor::from_vec(var_data, &[batch_size * seq_len, 1], &self.stream)?; - let var_output = grn.forward(&var_tensor, context)?; // (batch*seq_len, hidden_size) - var_outputs_flat.push(var_output); + // Extract variable i from flat_input: column i -> [flat_rows, 1] + let var_tensor = gpu_narrow_2d(&flat_input, 1, i, 1)?; + let var_output = grn.forward(&var_tensor, context)?; // [flat_rows, hidden_size] + var_outputs.push(var_output); } - // Concatenate all variable outputs: (batch*seq_len, hidden_size * input_size) - let mut cat_data = vec![0.0_f32; batch_size * seq_len * self.hidden_size * self.input_size]; - for (i, var_out) in var_outputs_flat.iter().enumerate() { - let var_host = var_out.to_vec()?; - for bs in 0..(batch_size * seq_len) { - for h in 0..self.hidden_size { - let src = bs * self.hidden_size + h; - let dst = bs * self.hidden_size * self.input_size + h * self.input_size + i; - if let (Some(&v), Some(slot)) = (var_host.get(src), cat_data.get_mut(dst)) { - *slot = v; - } - } - } + // Concatenate all variable outputs along dim1: [flat_rows, hidden_size * input_size] + let mut cat_tensor = var_outputs.first().ok_or_else(|| { + MLError::InvalidInput("No variable outputs".to_owned()) + })?.clone(); + for var_out in var_outputs.iter().skip(1) { + cat_tensor = gpu_cat_dim1(&cat_tensor, var_out)?; } - let cat_tensor = GpuTensor::from_vec( - cat_data, - &[batch_size * seq_len, self.hidden_size * self.input_size], - &self.stream, - )?; // Compute attention weights - let raw_weights = self.attention_weights.forward(&cat_tensor)?; // (batch*seq_len, input_size) + let raw_weights = self.attention_weights.forward(&cat_tensor)?; // [flat_rows, input_size] let attention = gpu_softmax(&raw_weights)?; - // Update importance scores + // Update importance scores (single scalar download per variable -- acceptable for monitoring) self.update_importance_scores(&attention)?; // Apply variable selection: weighted sum of per-variable GRN outputs - // attention: (batch*seq_len, input_size) - // var_outputs: input_size items of (batch*seq_len, hidden_size) - let mut result_data = vec![0.0_f32; batch_size * seq_len * self.hidden_size]; - let attn_host = attention.to_vec()?; - - for (i, var_out) in var_outputs_flat.iter().enumerate() { - let var_host = var_out.to_vec()?; - for bs in 0..(batch_size * seq_len) { - let w = attn_host.get(bs * self.input_size + i).copied().unwrap_or(0.0); - for h in 0..self.hidden_size { - let v = var_host.get(bs * self.hidden_size + h).copied().unwrap_or(0.0); - if let Some(slot) = result_data.get_mut(bs * self.hidden_size + h) { - *slot += w * v; - } - } - } + // attention: [flat_rows, input_size], var_outputs[i]: [flat_rows, hidden_size] + // result = sum_i(attention[:, i:i+1] * var_outputs[i]) + let mut result = GpuTensor::zeros(&[flat_rows, self.hidden_size], &self.stream)?; + for (i, var_out) in var_outputs.iter().enumerate() { + // Extract attention weight for variable i: [flat_rows, 1] + let attn_col = gpu_narrow_2d(&attention, 1, i, 1)?; + // Broadcast multiply: [flat_rows, 1] * [flat_rows, hidden_size] -> [flat_rows, hidden_size] + // We need to broadcast attn_col across hidden_size columns. + // Use gpu_broadcast_mul_col: a[rows, cols] * b[rows, 1] -> [rows, cols] + let weighted = crate::gpu_tensor::gpu_broadcast_mul_col(var_out, &attn_col)?; + result = gpu_add(&result, &weighted)?; } - // Reshape to (batch, seq_len, hidden_size) - GpuTensor::from_vec(result_data, &[batch_size, seq_len, self.hidden_size], &self.stream) + // Reshape to [batch, seq_len, hidden_size] + result.reshape(&[batch_size, seq_len, self.hidden_size]) } fn update_importance_scores(&mut self, attention_weights: &GpuTensor) -> Result<(), MLError> { - let host = attention_weights.to_vec()?; - let rows = attention_weights.dim(0)?; + // Use gpu_mean_all on per-column slices to avoid bulk download let cols = attention_weights.dim(1)?; - self.importance_scores.clear(); for c in 0..cols { - let mut sum = 0.0_f64; - for r in 0..rows { - if let Some(&v) = host.get(r * cols + c) { - sum += v as f64; - } - } - self.importance_scores.insert(c, sum / rows as f64); + let col = gpu_narrow_2d(attention_weights, 1, c, 1)?; + let mean_val = crate::gpu_tensor::gpu_mean_all(&col)?; + self.importance_scores.insert(c, mean_val as f64); } Ok(()) }