From a797e61d96b1cc8688c515eefa355382ae7b79cf Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 27 Mar 2026 23:00:45 +0100 Subject: [PATCH] =?UTF-8?q?WIP(bf16):=20atomic=20BF16=20conversion=20?= =?UTF-8?q?=E2=80=94=2036/37=20CUDA=20kernels,=20all=20Rust=20types?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CUDA kernels: 36 of 37 compiled with __nv_bfloat16* (dt_kernels.cu remaining) Rust types: ALL CudaSlice → CudaSlice across ml + ml-core Build: all kernels now compiled with common_device_functions.cuh (BF16 helpers) BF16 math wrappers: bf16_sqrt, bf16_log, bf16_exp, bf16_pow, bf16_fabs, bf16_fmax, bf16_fmin, bf16_cos, bf16_zero, bf16_one, bf16(), atomicAddBF16 NOT YET COMPILING — ml-core boundary errors (Vec → Vec) and dt_kernels.cu float*__nv_bfloat16 ambiguity remain. Co-Authored-By: Claude Opus 4.6 (1M context) --- Cargo.lock | 1 + crates/ml-core/Cargo.toml | 3 +- .../ml-core/src/cuda_autograd/activations.rs | 32 +-- .../ml-core/src/cuda_autograd/elementwise.rs | 146 ++++++------ .../ml-core/src/cuda_autograd/gpu_tensor.rs | 38 +-- crates/ml-core/src/cuda_autograd/init.rs | 24 +- crates/ml-core/src/cuda_autograd/linear.rs | 6 +- crates/ml-core/src/cuda_autograd/mod.rs | 4 +- crates/ml-core/src/cuda_autograd/optimizer.rs | 10 +- .../ml-core/src/cuda_autograd/reductions.rs | 46 ++-- .../ml-core/src/cuda_autograd/stream_ops.rs | 62 ++--- crates/ml-core/src/cuda_autograd/var_store.rs | 12 +- crates/ml/build.rs | 32 +-- .../attention_backward_kernel.cu | 118 +++++----- .../ml/src/cuda_pipeline/attention_kernel.cu | 52 ++--- .../src/cuda_pipeline/backtest_env_kernel.cu | 10 +- .../backtest_forward_ppo_kernel.cu | 16 +- .../backtest_forward_supervised_kernel.cu | 4 +- .../cuda_pipeline/backtest_gather_kernel.cu | 38 +-- .../cuda_pipeline/backtest_metrics_kernel.cu | 10 +- .../ml/src/cuda_pipeline/backward_kernels.cu | 23 +- .../ml/src/cuda_pipeline/batched_backward.rs | 18 +- .../ml/src/cuda_pipeline/batched_forward.rs | 50 ++-- crates/ml/src/cuda_pipeline/bias_kernels.cu | 14 +- .../ml/src/cuda_pipeline/c51_loss_kernel.cu | 14 +- .../cuda_pipeline/common_device_functions.cuh | 81 ++++++- .../ml/src/cuda_pipeline/cql_grad_kernel.cu | 26 +-- .../curiosity_training_kernel.cu | 93 ++++---- .../src/cuda_pipeline/decision_transformer.rs | 74 +++--- .../src/cuda_pipeline/dqn_utility_kernels.cu | 108 ++++----- crates/ml/src/cuda_pipeline/dt_kernels.cu | 148 ++++++------ crates/ml/src/cuda_pipeline/ema_kernel.cu | 8 +- .../ml/src/cuda_pipeline/ensemble_kernels.cu | 58 ++--- .../cuda_pipeline/epsilon_greedy_kernel.cu | 52 ++--- .../src/cuda_pipeline/experience_kernels.cu | 140 +++++------ .../src/cuda_pipeline/gpu_action_selector.rs | 18 +- crates/ml/src/cuda_pipeline/gpu_attention.rs | 38 +-- .../cuda_pipeline/gpu_backtest_evaluator.rs | 124 +++++----- .../cuda_pipeline/gpu_curiosity_trainer.rs | 62 ++--- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 218 +++++++++--------- .../cuda_pipeline/gpu_experience_collector.rs | 120 +++++----- crates/ml/src/cuda_pipeline/gpu_her.rs | 38 +-- .../ml/src/cuda_pipeline/gpu_iql_trainer.rs | 34 +-- crates/ml/src/cuda_pipeline/gpu_iqn_head.rs | 68 +++--- crates/ml/src/cuda_pipeline/gpu_monitoring.rs | 6 +- crates/ml/src/cuda_pipeline/gpu_portfolio.rs | 32 +-- .../ml/src/cuda_pipeline/gpu_ppo_collector.rs | 50 ++-- crates/ml/src/cuda_pipeline/gpu_statistics.rs | 10 +- .../src/cuda_pipeline/gpu_training_guard.rs | 18 +- .../ml/src/cuda_pipeline/gpu_walk_forward.rs | 10 +- crates/ml/src/cuda_pipeline/gpu_weights.rs | 144 ++++++------ .../src/cuda_pipeline/her_relabel_kernel.cu | 26 +-- .../ml/src/cuda_pipeline/iql_value_kernel.cu | 208 ++++++++--------- .../ml/src/cuda_pipeline/iqn_cvar_kernel.cu | 6 +- .../src/cuda_pipeline/iqn_dual_head_kernel.cu | 200 ++++++++-------- crates/ml/src/cuda_pipeline/mod.rs | 64 ++--- .../ml/src/cuda_pipeline/monitoring_kernel.cu | 4 +- .../ml/src/cuda_pipeline/mse_loss_kernel.cu | 14 +- crates/ml/src/cuda_pipeline/nstep_kernel.cu | 20 +- .../ml/src/cuda_pipeline/per_update_kernel.cu | 16 +- .../cuda_pipeline/ppo_experience_kernel.cu | 86 +++---- crates/ml/src/cuda_pipeline/q_stats_kernel.cu | 12 +- .../ml/src/cuda_pipeline/relu_mask_kernel.cu | 6 +- crates/ml/src/cuda_pipeline/signal_adapter.rs | 36 +-- .../cuda_pipeline/signal_adapter_kernel.cu | 40 ++-- .../ml/src/cuda_pipeline/statistics_kernel.cu | 12 +- .../src/cuda_pipeline/trade_stats_kernel.cu | 4 +- .../cuda_pipeline/training_guard_kernel.cu | 34 +-- crates/ml/src/hyperopt/adapters/ppo.rs | 6 +- crates/ml/src/trainers/dqn/config.rs | 6 +- crates/ml/src/trainers/dqn/fused_training.rs | 40 ++-- .../trainers/dqn/smoke_tests/gpu_residency.rs | 10 +- .../dqn/smoke_tests/gradient_budget.rs | 6 +- .../trainers/dqn/smoke_tests/performance.rs | 2 +- .../dqn/smoke_tests/training_stability.rs | 4 +- crates/ml/src/trainers/dqn/trainer/action.rs | 4 +- crates/ml/src/trainers/dqn/trainer/metrics.rs | 6 +- crates/ml/src/trainers/dqn/trainer/mod.rs | 4 +- .../ml/src/trainers/dqn/trainer/train_step.rs | 4 +- .../src/trainers/dqn/trainer/training_loop.rs | 6 +- crates/ml/src/trainers/ppo.rs | 4 +- 81 files changed, 1776 insertions(+), 1675 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 80f1541fc..f33122180 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6228,6 +6228,7 @@ dependencies = [ "cudarc", "dashmap 6.1.0", "futures", + "half", "libc", "mimalloc", "ndarray", diff --git a/crates/ml-core/Cargo.toml b/crates/ml-core/Cargo.toml index e87d3ba5d..17bce14bd 100644 --- a/crates/ml-core/Cargo.toml +++ b/crates/ml-core/Cargo.toml @@ -49,7 +49,8 @@ uuid.workspace = true parking_lot = { version = "0.12", features = ["hardware-lock-elision"] } ndarray = { workspace = true, features = ["rayon"] } sha2 = "0.10" -cudarc = { version = "0.19", optional = true, default-features = false, features = ["driver", "nvrtc", "cublas", "dynamic-linking", "std", "cuda-version-from-build-system"] } +cudarc = { version = "0.19", optional = true, default-features = false, features = ["driver", "cublas", "dynamic-linking", "std", "cuda-version-from-build-system", "f16"] } +half = { workspace = true } rand.workspace = true rand_distr.workspace = true toml.workspace = true diff --git a/crates/ml-core/src/cuda_autograd/activations.rs b/crates/ml-core/src/cuda_autograd/activations.rs index b280abd13..e2a918f88 100644 --- a/crates/ml-core/src/cuda_autograd/activations.rs +++ b/crates/ml-core/src/cuda_autograd/activations.rs @@ -332,14 +332,14 @@ impl ActivationKernels { /// Use this when the backward pass is not needed (inference / PPO forward). pub fn relu_fwd_raw( &self, - x: &CudaSlice, + x: &CudaSlice, len: usize, stream: &Arc, - ) -> Result, MLError> { - let y = stream.alloc_zeros::(len).map_err(|e| { + ) -> Result, MLError> { + let y = stream.alloc_zeros::(len).map_err(|e| { MLError::ModelError(format!("relu_fwd_raw alloc: {e}")) })?; - let mask = stream.alloc_zeros::(len).map_err(|e| { + let mask = stream.alloc_zeros::(len).map_err(|e| { MLError::ModelError(format!("relu_fwd_raw mask alloc: {e}")) })?; let n_i32 = len as i32; @@ -360,16 +360,16 @@ impl ActivationKernels { /// ReLU backward on raw `CudaSlice` pairs. Returns `grad_input`. pub fn relu_bwd_raw( &self, - grad_output: &CudaSlice, - input: &CudaSlice, + grad_output: &CudaSlice, + input: &CudaSlice, len: usize, stream: &Arc, - ) -> Result, MLError> { + ) -> Result, MLError> { // Forward to get the mask - let mask = stream.alloc_zeros::(len).map_err(|e| { + let mask = stream.alloc_zeros::(len).map_err(|e| { MLError::ModelError(format!("relu_bwd_raw mask alloc: {e}")) })?; - let y_scratch = stream.alloc_zeros::(len).map_err(|e| { + let y_scratch = stream.alloc_zeros::(len).map_err(|e| { MLError::ModelError(format!("relu_bwd_raw y_scratch alloc: {e}")) })?; let n_i32 = len as i32; @@ -385,7 +385,7 @@ impl ActivationKernels { .map_err(|e| MLError::ModelError(format!("relu_bwd_raw fwd: {e}")))?; } // Backward using the mask - let dx = stream.alloc_zeros::(len).map_err(|e| { + let dx = stream.alloc_zeros::(len).map_err(|e| { MLError::ModelError(format!("relu_bwd_raw dx alloc: {e}")) })?; unsafe { @@ -404,11 +404,11 @@ impl ActivationKernels { /// Tanh forward on raw `CudaSlice`. Returns output only (no saved tensor). pub fn tanh_fwd_raw( &self, - x: &CudaSlice, + x: &CudaSlice, len: usize, stream: &Arc, - ) -> Result, MLError> { - let y = stream.alloc_zeros::(len).map_err(|e| { + ) -> Result, MLError> { + let y = stream.alloc_zeros::(len).map_err(|e| { MLError::ModelError(format!("tanh_fwd_raw alloc: {e}")) })?; let n_i32 = len as i32; @@ -428,11 +428,11 @@ impl ActivationKernels { /// Sigmoid forward on raw `CudaSlice`. Returns output only (no saved tensor). pub fn sigmoid_fwd_raw( &self, - x: &CudaSlice, + x: &CudaSlice, len: usize, stream: &Arc, - ) -> Result, MLError> { - let y = stream.alloc_zeros::(len).map_err(|e| { + ) -> Result, MLError> { + let y = stream.alloc_zeros::(len).map_err(|e| { MLError::ModelError(format!("sigmoid_fwd_raw alloc: {e}")) })?; let n_i32 = len as i32; diff --git a/crates/ml-core/src/cuda_autograd/elementwise.rs b/crates/ml-core/src/cuda_autograd/elementwise.rs index 4bc779e2e..f88c441a7 100644 --- a/crates/ml-core/src/cuda_autograd/elementwise.rs +++ b/crates/ml-core/src/cuda_autograd/elementwise.rs @@ -85,12 +85,12 @@ impl ElementwiseKernels { /// `op`: 0=add, 1=sub, 2=mul, 3=div, 4=min, 5=max pub fn binary( &self, - a: &CudaSlice, - b: &CudaSlice, + a: &CudaSlice, + b: &CudaSlice, n: usize, op: i32, - ) -> Result, MLError> { - let out = self.stream.alloc_zeros::(n).map_err(|e| { + ) -> Result, MLError> { + let out = self.stream.alloc_zeros::(n).map_err(|e| { MLError::ModelError(format!("binary alloc: {e}")) })?; let n_i32 = n as i32; @@ -98,8 +98,8 @@ impl ElementwiseKernels { // SAFETY: Kernel arguments match the CUDA function signature exactly: // (const float* a, const float* b, float* out, int n, int op) - // a and b are valid CudaSlice of at least n elements. - // out is a freshly allocated CudaSlice of n elements. + // a and b are valid CudaSlice of at least n elements. + // out is a freshly allocated CudaSlice of n elements. unsafe { self.stream .launch_builder(&self.binary_fn) @@ -122,13 +122,13 @@ impl ElementwiseKernels { /// 17=sigmoid, 18=tanh pub fn unary( &self, - x: &CudaSlice, + x: &CudaSlice, n: usize, op: i32, param1: f32, param2: f32, - ) -> Result, MLError> { - let out = self.stream.alloc_zeros::(n).map_err(|e| { + ) -> Result, MLError> { + let out = self.stream.alloc_zeros::(n).map_err(|e| { MLError::ModelError(format!("unary alloc: {e}")) })?; let n_i32 = n as i32; @@ -136,8 +136,8 @@ impl ElementwiseKernels { // SAFETY: Kernel arguments match the CUDA function signature exactly: // (const float* x, float* out, int n, int op, float param1, float param2) - // x is a valid CudaSlice of at least n elements. - // out is a freshly allocated CudaSlice of n elements. + // x is a valid CudaSlice of at least n elements. + // out is a freshly allocated CudaSlice of n elements. unsafe { self.stream .launch_builder(&self.unary_fn) @@ -160,13 +160,13 @@ impl ElementwiseKernels { /// Otherwise: out[i] = tensor[i] op scalar pub fn broadcast_scalar( &self, - tensor: &CudaSlice, - scalar: &CudaSlice, + tensor: &CudaSlice, + scalar: &CudaSlice, n: usize, op: i32, scalar_is_lhs: i32, - ) -> Result, MLError> { - let out = self.stream.alloc_zeros::(n).map_err(|e| { + ) -> Result, MLError> { + let out = self.stream.alloc_zeros::(n).map_err(|e| { MLError::ModelError(format!("broadcast_scalar alloc: {e}")) })?; let n_i32 = n as i32; @@ -197,14 +197,14 @@ impl ElementwiseKernels { /// `row_is_lhs`: if true, out[i] = row[i%cols] op matrix[i]; else matrix[i] op row[i%cols] pub fn broadcast_row( &self, - matrix: &CudaSlice, - row: &CudaSlice, + matrix: &CudaSlice, + row: &CudaSlice, total: usize, cols: usize, op: i32, row_is_lhs: i32, - ) -> Result, MLError> { - let out = self.stream.alloc_zeros::(total).map_err(|e| { + ) -> Result, MLError> { + let out = self.stream.alloc_zeros::(total).map_err(|e| { MLError::ModelError(format!("broadcast_row alloc: {e}")) })?; let total_i32 = total as i32; @@ -237,14 +237,14 @@ impl ElementwiseKernels { /// `col_is_lhs`: if true, `out[i] = col_vec[i/cols] op matrix[i]`; else `matrix[i] op col_vec[i/cols]` pub fn broadcast_col( &self, - matrix: &CudaSlice, - col_vec: &CudaSlice, + matrix: &CudaSlice, + col_vec: &CudaSlice, total: usize, cols: usize, op: i32, col_is_lhs: i32, - ) -> Result, MLError> { - let out = self.stream.alloc_zeros::(total).map_err(|e| { + ) -> Result, MLError> { + let out = self.stream.alloc_zeros::(total).map_err(|e| { MLError::ModelError(format!("broadcast_col alloc: {e}")) })?; let total_i32 = total as i32; @@ -277,14 +277,14 @@ impl ElementwiseKernels { /// `tgt_strides` is a GPU-resident u32 buffer of length `ndim`. pub fn expand( &self, - src: &CudaSlice, + src: &CudaSlice, src_strides_dev: &CudaSlice, src_shape_dev: &CudaSlice, tgt_strides_dev: &CudaSlice, total: usize, ndim: usize, - ) -> Result, MLError> { - let out = self.stream.alloc_zeros::(total).map_err(|e| { + ) -> Result, MLError> { + let out = self.stream.alloc_zeros::(total).map_err(|e| { MLError::ModelError(format!("expand alloc: {e}")) })?; let total_i32 = total as i32; @@ -315,12 +315,12 @@ impl ElementwiseKernels { /// Transpose a 2-D matrix [rows, cols] -> [cols, rows]. pub fn transpose_2d( &self, - src: &CudaSlice, + src: &CudaSlice, rows: usize, cols: usize, - ) -> Result, MLError> { + ) -> Result, MLError> { let total = rows * cols; - let out = self.stream.alloc_zeros::(total).map_err(|e| { + let out = self.stream.alloc_zeros::(total).map_err(|e| { MLError::ModelError(format!("transpose alloc: {e}")) })?; let total_i32 = total as i32; @@ -330,7 +330,7 @@ impl ElementwiseKernels { // SAFETY: Kernel arguments match the CUDA function signature exactly: // (const float* src, float* dst, int total, int rows, int cols) - // src and dst are both valid CudaSlice of `total` elements. + // src and dst are both valid CudaSlice of `total` elements. unsafe { self.stream .launch_builder(&self.transpose_2d_fn) @@ -351,15 +351,15 @@ impl ElementwiseKernels { /// `indices_dev` is a GPU-resident u32 buffer. pub fn gather( &self, - src: &CudaSlice, + src: &CudaSlice, indices_dev: &CudaSlice, num_indices: usize, inner_size: usize, outer_count: usize, src_dim_size: usize, - ) -> Result, MLError> { + ) -> Result, MLError> { let total = outer_count * num_indices * inner_size; - let out = self.stream.alloc_zeros::(total).map_err(|e| { + let out = self.stream.alloc_zeros::(total).map_err(|e| { MLError::ModelError(format!("gather alloc: {e}")) })?; let total_i32 = total as i32; @@ -393,18 +393,18 @@ impl ElementwiseKernels { /// /// For each row `i`, picks the element at column `indices[i]` and writes /// it to `out[i]`. `indices` is a GPU-resident `i32` buffer of length `rows`. - /// Output is a flat `CudaSlice` of length `rows`. + /// Output is a flat `CudaSlice` of length `rows`. /// /// This eliminates the need to download actions to CPU for per-row /// index selection from log-softmax matrices in PPO/policy gradient methods. pub fn gather_rows( &self, - src: &CudaSlice, + src: &CudaSlice, indices: &CudaSlice, rows: usize, cols: usize, - ) -> Result, MLError> { - let out = self.stream.alloc_zeros::(rows).map_err(|e| { + ) -> Result, MLError> { + let out = self.stream.alloc_zeros::(rows).map_err(|e| { MLError::ModelError(format!("gather_rows alloc: {e}")) })?; let rows_i32 = rows as i32; @@ -413,9 +413,9 @@ impl ElementwiseKernels { // SAFETY: Kernel arguments match the CUDA function signature exactly: // (const float* src, const int* indices, float* out, int rows, int cols) - // src is a valid CudaSlice of at least rows*cols elements. + // src is a valid CudaSlice of at least rows*cols elements. // indices is a valid CudaSlice of at least rows elements. - // out is a freshly allocated CudaSlice of rows elements. + // out is a freshly allocated CudaSlice of rows elements. unsafe { self.stream .launch_builder(&self.gather_rows_fn) @@ -436,12 +436,12 @@ impl ElementwiseKernels { /// (common for action indices stored as unsigned). pub fn gather_rows_u32( &self, - src: &CudaSlice, + src: &CudaSlice, indices: &CudaSlice, rows: usize, cols: usize, - ) -> Result, MLError> { - let out = self.stream.alloc_zeros::(rows).map_err(|e| { + ) -> Result, MLError> { + let out = self.stream.alloc_zeros::(rows).map_err(|e| { MLError::ModelError(format!("gather_rows_u32 alloc: {e}")) })?; let rows_i32 = rows as i32; @@ -450,9 +450,9 @@ impl ElementwiseKernels { // SAFETY: Kernel arguments match the CUDA function signature exactly: // (const float* src, const unsigned int* indices, float* out, int rows, int cols) - // src is a valid CudaSlice of at least rows*cols elements. + // src is a valid CudaSlice of at least rows*cols elements. // indices is a valid CudaSlice of at least rows elements. - // out is a freshly allocated CudaSlice of rows elements. + // out is a freshly allocated CudaSlice of rows elements. unsafe { self.stream .launch_builder(&self.gather_rows_u32_fn) @@ -470,82 +470,82 @@ impl ElementwiseKernels { // ── Convenience wrappers for unary ops 5-10 ─────────────────────── /// Element-wise absolute value (`|x|`). Op 5. - pub fn abs(&self, x: &CudaSlice, n: usize) -> Result, MLError> { + pub fn abs(&self, x: &CudaSlice, n: usize) -> Result, MLError> { self.unary(x, n, 5, 0.0, 0.0) } /// Element-wise negation (`-x`). Op 6. - pub fn neg(&self, x: &CudaSlice, n: usize) -> Result, MLError> { + pub fn neg(&self, x: &CudaSlice, n: usize) -> Result, MLError> { self.unary(x, n, 6, 0.0, 0.0) } /// Element-wise exponential (`e^x`). Op 7. - pub fn exp(&self, x: &CudaSlice, n: usize) -> Result, MLError> { + pub fn exp(&self, x: &CudaSlice, n: usize) -> Result, MLError> { self.unary(x, n, 7, 0.0, 0.0) } /// Element-wise natural logarithm (`ln(x)`). Op 8. - pub fn log(&self, x: &CudaSlice, n: usize) -> Result, MLError> { + pub fn log(&self, x: &CudaSlice, n: usize) -> Result, MLError> { self.unary(x, n, 8, 0.0, 0.0) } /// Element-wise less-or-equal comparison. Returns 1.0 if `x <= threshold`, else 0.0. Op 9. - pub fn le(&self, x: &CudaSlice, threshold: f32, n: usize) -> Result, MLError> { + pub fn le(&self, x: &CudaSlice, threshold: f32, n: usize) -> Result, MLError> { self.unary(x, n, 9, threshold, 0.0) } /// Element-wise affine transform: `x * mul + add`. Op 10. - pub fn affine(&self, x: &CudaSlice, mul: f32, add: f32, n: usize) -> Result, MLError> { + pub fn affine(&self, x: &CudaSlice, mul: f32, add: f32, n: usize) -> Result, MLError> { self.unary(x, n, 10, mul, add) } /// Element-wise sine (`sin(x)`). Op 11. - pub fn sin(&self, x: &CudaSlice, n: usize) -> Result, MLError> { + pub fn sin(&self, x: &CudaSlice, n: usize) -> Result, MLError> { self.unary(x, n, 11, 0.0, 0.0) } /// Element-wise cosine (`cos(x)`). Op 12. - pub fn cos(&self, x: &CudaSlice, n: usize) -> Result, MLError> { + pub fn cos(&self, x: &CudaSlice, n: usize) -> Result, MLError> { self.unary(x, n, 12, 0.0, 0.0) } /// Element-wise square root (`sqrt(x)`). Op 13. - pub fn sqrt(&self, x: &CudaSlice, n: usize) -> Result, MLError> { + pub fn sqrt(&self, x: &CudaSlice, n: usize) -> Result, MLError> { self.unary(x, n, 13, 0.0, 0.0) } /// Element-wise SiLU (`x * sigmoid(x)`). Op 14. - pub fn silu(&self, x: &CudaSlice, n: usize) -> Result, MLError> { + pub fn silu(&self, x: &CudaSlice, n: usize) -> Result, MLError> { self.unary(x, n, 14, 0.0, 0.0) } /// Element-wise ELU (`x if x>=0, alpha*(exp(x)-1) otherwise`). Op 15. - pub fn elu(&self, x: &CudaSlice, alpha: f32, n: usize) -> Result, MLError> { + pub fn elu(&self, x: &CudaSlice, alpha: f32, n: usize) -> Result, MLError> { self.unary(x, n, 15, alpha, 0.0) } /// Element-wise reciprocal (`1/x`, safe: returns 1e10 for `|x| < 1e-10`). Op 16. - pub fn recip(&self, x: &CudaSlice, n: usize) -> Result, MLError> { + pub fn recip(&self, x: &CudaSlice, n: usize) -> Result, MLError> { self.unary(x, n, 16, 0.0, 0.0) } /// Element-wise sigmoid (`1 / (1 + exp(-x))`). Op 17. - pub fn sigmoid(&self, x: &CudaSlice, n: usize) -> Result, MLError> { + pub fn sigmoid(&self, x: &CudaSlice, n: usize) -> Result, MLError> { self.unary(x, n, 17, 0.0, 0.0) } /// Element-wise tanh. Op 18. - pub fn tanh(&self, x: &CudaSlice, n: usize) -> Result, MLError> { + pub fn tanh(&self, x: &CudaSlice, n: usize) -> Result, MLError> { self.unary(x, n, 18, 0.0, 0.0) } /// Element-wise minimum of two same-shape buffers. - pub fn min(&self, a: &CudaSlice, b: &CudaSlice, n: usize) -> Result, MLError> { + pub fn min(&self, a: &CudaSlice, b: &CudaSlice, n: usize) -> Result, MLError> { self.binary(a, b, n, 4) } /// Element-wise maximum of two same-shape buffers. - pub fn max(&self, a: &CudaSlice, b: &CudaSlice, n: usize) -> Result, MLError> { + pub fn max(&self, a: &CudaSlice, b: &CudaSlice, n: usize) -> Result, MLError> { self.binary(a, b, n, 5) } @@ -887,8 +887,8 @@ mod tests { let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); let a_host = vec![1.0_f32, 2.0, 3.0, 4.0]; let b_host = vec![10.0_f32, 20.0, 30.0, 40.0]; - let mut a = stream.alloc_zeros::(4).unwrap(); - let mut b = stream.alloc_zeros::(4).unwrap(); + let mut a = stream.alloc_zeros::(4).unwrap(); + let mut b = stream.alloc_zeros::(4).unwrap(); stream.memcpy_htod(&a_host, &mut a).unwrap(); stream.memcpy_htod(&b_host, &mut b).unwrap(); @@ -907,8 +907,8 @@ mod tests { let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); let a_host = vec![10.0_f32, 20.0, 30.0]; let b_host = vec![1.0_f32, 2.0, 3.0]; - let mut a = stream.alloc_zeros::(3).unwrap(); - let mut b = stream.alloc_zeros::(3).unwrap(); + let mut a = stream.alloc_zeros::(3).unwrap(); + let mut b = stream.alloc_zeros::(3).unwrap(); stream.memcpy_htod(&a_host, &mut a).unwrap(); stream.memcpy_htod(&b_host, &mut b).unwrap(); @@ -926,8 +926,8 @@ mod tests { let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); let a_host = vec![2.0_f32, 3.0, 4.0]; let b_host = vec![5.0_f32, 6.0, 7.0]; - let mut a = stream.alloc_zeros::(3).unwrap(); - let mut b = stream.alloc_zeros::(3).unwrap(); + let mut a = stream.alloc_zeros::(3).unwrap(); + let mut b = stream.alloc_zeros::(3).unwrap(); stream.memcpy_htod(&a_host, &mut a).unwrap(); stream.memcpy_htod(&b_host, &mut b).unwrap(); @@ -944,7 +944,7 @@ mod tests { let stream = make_stream(); let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); let x_host = vec![-2.0_f32, -1.0, 0.0, 1.0, 2.0]; - let mut x = stream.alloc_zeros::(5).unwrap(); + let mut x = stream.alloc_zeros::(5).unwrap(); stream.memcpy_htod(&x_host, &mut x).unwrap(); let out = kernels.unary(&x, 5, 3, 0.0, 0.0).unwrap(); // relu @@ -962,7 +962,7 @@ mod tests { let stream = make_stream(); let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); let x_host = vec![1.0_f32, 2.0, 3.0, -4.0]; - let mut x = stream.alloc_zeros::(4).unwrap(); + let mut x = stream.alloc_zeros::(4).unwrap(); stream.memcpy_htod(&x_host, &mut x).unwrap(); let out = kernels.unary(&x, 4, 1, 0.0, 0.0).unwrap(); // sqr @@ -979,7 +979,7 @@ mod tests { let stream = make_stream(); let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); let x_host = vec![-5.0_f32, -1.0, 0.5, 2.0, 10.0]; - let mut x = stream.alloc_zeros::(5).unwrap(); + let mut x = stream.alloc_zeros::(5).unwrap(); stream.memcpy_htod(&x_host, &mut x).unwrap(); let out = kernels.unary(&x, 5, 4, -2.0, 3.0).unwrap(); // clamp(-2, 3) @@ -998,7 +998,7 @@ mod tests { let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); // 2x3 matrix: [[1,2,3],[4,5,6]] let x_host = vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]; - let mut x = stream.alloc_zeros::(6).unwrap(); + let mut x = stream.alloc_zeros::(6).unwrap(); stream.memcpy_htod(&x_host, &mut x).unwrap(); let out = kernels.transpose_2d(&x, 2, 3).unwrap(); @@ -1019,8 +1019,8 @@ mod tests { let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); let tensor_host = vec![1.0_f32, 2.0, 3.0, 4.0]; let scalar_host = vec![10.0_f32]; - let mut tensor = stream.alloc_zeros::(4).unwrap(); - let mut scalar = stream.alloc_zeros::(1).unwrap(); + let mut tensor = stream.alloc_zeros::(4).unwrap(); + let mut scalar = stream.alloc_zeros::(1).unwrap(); stream.memcpy_htod(&tensor_host, &mut tensor).unwrap(); stream.memcpy_htod(&scalar_host, &mut scalar).unwrap(); @@ -1039,7 +1039,7 @@ mod tests { let kernels = ElementwiseKernels::new(&stream).expect("compile kernels"); // 4x3 src: row0=[0,1,2], row1=[3,4,5], row2=[6,7,8], row3=[9,10,11] let src_host: Vec = (0..12).map(|i| i as f32).collect(); - let mut src = stream.alloc_zeros::(12).unwrap(); + let mut src = stream.alloc_zeros::(12).unwrap(); stream.memcpy_htod(&src_host, &mut src).unwrap(); // Select rows 2, 0 diff --git a/crates/ml-core/src/cuda_autograd/gpu_tensor.rs b/crates/ml-core/src/cuda_autograd/gpu_tensor.rs index ba0e7506b..f0fddc6a9 100644 --- a/crates/ml-core/src/cuda_autograd/gpu_tensor.rs +++ b/crates/ml-core/src/cuda_autograd/gpu_tensor.rs @@ -1,4 +1,4 @@ -//! GPU-resident tensor wrapping `CudaSlice` with shape metadata. +//! GPU-resident tensor wrapping `CudaSlice` with shape metadata. //! //! Replaces `candle_core::Tensor` for the CUDA autograd system. This is a thin //! wrapper — no autograd graph, no dtype dispatch, no device polymorphism. @@ -19,14 +19,14 @@ 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, stream: &CudaStream) -> u64 { +fn raw_ptr(slice: &CudaSlice, 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, stream: &CudaStream) -> u64 { +fn raw_ptr_mut(slice: &mut CudaSlice, stream: &CudaStream) -> u64 { let (ptr, guard) = slice.device_ptr_mut(stream); let _no_drop = ManuallyDrop::new(guard); ptr @@ -51,19 +51,19 @@ impl ScalarValue for &[f32] { /// A GPU-resident f32 tensor with shape information. /// -/// This is a lightweight wrapper around `CudaSlice` that carries shape +/// This is a lightweight wrapper around `CudaSlice` that carries shape /// metadata. Unlike Candle's `Tensor`, it has no autograd graph, no reference /// counting of operations, and no device abstraction — it is always on GPU. /// /// ## Ownership /// -/// The inner `CudaSlice` is reference-counted (cudarc 0.19). Cloning a +/// The inner `CudaSlice` is reference-counted (cudarc 0.19). Cloning a /// `GpuTensor` shares the underlying GPU allocation — no device-side copy. /// The GPU memory is freed when the last reference is dropped. #[derive(Clone)] pub struct GpuTensor { /// Raw GPU buffer. Element count = product of `shape`. - pub(crate) data: CudaSlice, + pub(crate) data: CudaSlice, /// Tensor dimensions, e.g. `[batch, features]` or `[out_dim, in_dim]`. pub(crate) shape: Vec, } @@ -74,7 +74,7 @@ impl GpuTensor { /// # Errors /// /// Returns an error if the product of `shape` does not match `data.len()`. - pub fn new(data: CudaSlice, shape: Vec) -> Result { + pub fn new(data: CudaSlice, shape: Vec) -> Result { let expected: usize = shape.iter().product(); let actual = data.len(); if expected != actual { @@ -83,16 +83,16 @@ impl GpuTensor { Ok(Self { data, shape }) } - /// Access the underlying `CudaSlice` for direct CUDA kernel operations. + /// Access the underlying `CudaSlice` for direct CUDA kernel operations. /// /// Prefer higher-level methods (add, sub, affine, argmax, etc.) when available. /// Use this only when interfacing with raw CUDA kernels (elementwise, reductions). - pub fn cuda_data(&self) -> &CudaSlice { + pub fn cuda_data(&self) -> &CudaSlice { &self.data } /// Wrap an existing `CudaSlice` as a 1-D tensor (no shape validation needed). - pub fn from_flat(data: CudaSlice) -> Self { + pub fn from_flat(data: CudaSlice) -> Self { let len = data.len(); Self { data, @@ -103,7 +103,7 @@ impl GpuTensor { /// Allocate a zero-filled tensor of the given shape on `stream`. pub fn zeros(shape: &[usize], stream: &Arc) -> Result { let n: usize = shape.iter().product(); - let data = stream.alloc_zeros::(n).map_err(|e| { + let data = stream.alloc_zeros::(n).map_err(|e| { MLError::ModelError(format!("GpuTensor::zeros({shape:?}): {e}")) })?; Ok(Self { @@ -125,7 +125,7 @@ impl GpuTensor { actual: host.len(), }); } - let mut data = stream.alloc_zeros::(expected).map_err(|e| { + let mut data = stream.alloc_zeros::(expected).map_err(|e| { MLError::ModelError(format!("GpuTensor::from_host alloc: {e}")) })?; stream.memcpy_htod(host, &mut data).map_err(|e| { @@ -160,12 +160,12 @@ impl GpuTensor { } /// Borrow the underlying `CudaSlice`. - pub fn data(&self) -> &CudaSlice { + pub fn data(&self) -> &CudaSlice { &self.data } /// Mutably borrow the underlying `CudaSlice`. - pub fn data_mut(&mut self) -> &mut CudaSlice { + pub fn data_mut(&mut self) -> &mut CudaSlice { &mut self.data } @@ -173,7 +173,7 @@ impl GpuTensor { /// /// Useful for converting back to other GPU buffer types (e.g. `CudaVec`) /// without a GPU-side copy. - pub fn into_parts(self) -> (CudaSlice, Vec) { + pub fn into_parts(self) -> (CudaSlice, Vec) { (self.data, self.shape) } @@ -241,7 +241,7 @@ impl GpuTensor { let total_elems: usize = new_shape.iter().product(); // GPU-native: allocate output and DtoD copy each tensor's data - let mut output = stream.alloc_zeros::(total_elems).map_err(|e| { + let mut output = stream.alloc_zeros::(total_elems).map_err(|e| { MLError::ModelError(format!("cat: alloc {total_elems} f32: {e}")) })?; @@ -302,7 +302,7 @@ impl GpuTensor { let elems_per: usize = inner_shape.iter().product(); let total = batch * elems_per; - let mut output = stream.alloc_zeros::(total).map_err(|e| { + let mut output = stream.alloc_zeros::(total).map_err(|e| { MLError::ModelError(format!("stack: alloc {total} f32: {e}")) })?; @@ -335,7 +335,7 @@ impl GpuTensor { /// Uses async device-to-device memcpy — zero host download. pub fn gpu_clone(&self, stream: &Arc) -> Result { let n = self.numel(); - let dst = stream.alloc_zeros::(n).map_err(|e| { + let dst = stream.alloc_zeros::(n).map_err(|e| { MLError::ModelError(format!("gpu_clone alloc: {e}")) })?; let num_bytes = n * std::mem::size_of::(); @@ -749,7 +749,7 @@ impl GpuTensor { if dim == 0 { let offset = start * stride; let src_view = self.data.slice(offset..offset + new_numel); - let dst = stream.alloc_zeros::(new_numel).map_err(|e| { + let dst = stream.alloc_zeros::(new_numel).map_err(|e| { MLError::ModelError(format!("narrow alloc: {e}")) })?; let num_bytes = new_numel * std::mem::size_of::(); diff --git a/crates/ml-core/src/cuda_autograd/init.rs b/crates/ml-core/src/cuda_autograd/init.rs index 1aeeade4e..639029610 100644 --- a/crates/ml-core/src/cuda_autograd/init.rs +++ b/crates/ml-core/src/cuda_autograd/init.rs @@ -11,7 +11,7 @@ use cudarc::driver::{CudaSlice, CudaStream}; use crate::MLError; -/// Initialize a `CudaSlice` with Xavier uniform values. +/// Initialize a `CudaSlice` with Xavier uniform values. /// /// Samples from U(-limit, limit) where `limit = sqrt(6 / (fan_in + fan_out))`. /// Shape: `fan_out * fan_in` elements (row-major, matching `[out, in]` convention). @@ -19,14 +19,14 @@ pub fn xavier_uniform( fan_in: usize, fan_out: usize, stream: &Arc, -) -> Result, MLError> { +) -> Result, MLError> { let n = fan_out * fan_in; let limit = (6.0_f64 / (fan_in + fan_out) as f64).sqrt(); let host = generate_uniform(n, -limit, limit); upload_to_gpu(&host, stream) } -/// Initialize a `CudaSlice` with Kaiming uniform values (for ReLU networks). +/// Initialize a `CudaSlice` with Kaiming uniform values (for ReLU networks). /// /// Samples from U(-limit, limit) where `limit = sqrt(6 / fan_in)`. /// Shape: `fan_out * fan_in` elements. @@ -34,7 +34,7 @@ pub fn kaiming_uniform( fan_in: usize, fan_out: usize, stream: &Arc, -) -> Result, MLError> { +) -> Result, MLError> { let n = fan_out * fan_in; let limit = (6.0_f64 / fan_in as f64).sqrt(); let host = generate_uniform(n, -limit, limit); @@ -46,23 +46,23 @@ pub fn bias_uniform( fan_in: usize, fan_out: usize, stream: &Arc, -) -> Result, MLError> { +) -> Result, MLError> { let bound = 1.0 / (fan_in as f64).sqrt(); let host = generate_uniform(fan_out, -bound, bound); upload_to_gpu(&host, stream) } -/// Initialize a `CudaSlice` with zeros. +/// Initialize a `CudaSlice` with zeros. pub fn zeros( n: usize, stream: &Arc, -) -> Result, MLError> { - stream.alloc_zeros::(n).map_err(|e| { +) -> Result, MLError> { + stream.alloc_zeros::(n).map_err(|e| { MLError::ModelError(format!("init::zeros({n}): {e}")) }) } -/// Initialize a `CudaSlice` with near-zero Xavier values. +/// Initialize a `CudaSlice` with near-zero Xavier values. /// /// 0.01x Xavier scale — useful for distributional output layers where /// near-uniform softmax is desired at initialization. @@ -70,7 +70,7 @@ pub fn near_zero_xavier( fan_in: usize, fan_out: usize, stream: &Arc, -) -> Result, MLError> { +) -> Result, MLError> { let n = fan_out * fan_in; let limit = 0.01 * (6.0_f64 / (fan_in + fan_out) as f64).sqrt(); let host = generate_uniform(n, -limit, limit); @@ -126,8 +126,8 @@ fn generate_uniform(n: usize, lo: f64, hi: f64) -> Vec { pub fn upload_to_gpu( host: &[f32], stream: &Arc, -) -> Result, MLError> { - let mut buf = stream.alloc_zeros::(host.len()).map_err(|e| { +) -> Result, MLError> { + let mut buf = stream.alloc_zeros::(host.len()).map_err(|e| { MLError::ModelError(format!("init upload alloc: {e}")) })?; stream.memcpy_htod(host, &mut buf).map_err(|e| { diff --git a/crates/ml-core/src/cuda_autograd/linear.rs b/crates/ml-core/src/cuda_autograd/linear.rs index 7c13d8fbb..4372e00b1 100644 --- a/crates/ml-core/src/cuda_autograd/linear.rs +++ b/crates/ml-core/src/cuda_autograd/linear.rs @@ -24,14 +24,14 @@ use super::var_store::GpuVarStore; // ── Pointer helpers (same pattern as ml-supervised/gpu_tensor.rs) ──────── /// Extract raw CUDA device pointer from a CudaSlice (read-only). -fn raw_ptr(slice: &CudaSlice, stream: &CudaStream) -> u64 { +fn raw_ptr(slice: &CudaSlice, 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, stream: &CudaStream) -> u64 { +fn raw_ptr_mut(slice: &mut CudaSlice, stream: &CudaStream) -> u64 { let (ptr, guard) = slice.device_ptr_mut(stream); let _no_drop = ManuallyDrop::new(guard); ptr @@ -342,7 +342,7 @@ void reduce_sum_axis0_kernel(const float* __restrict__ x, /// Add bias to a 2D tensor: Y[b, j] += bias[j]. fn add_bias_2d( y: &mut GpuTensor, - bias: &CudaSlice, + bias: &CudaSlice, rows: usize, cols: usize, stream: &Arc, diff --git a/crates/ml-core/src/cuda_autograd/mod.rs b/crates/ml-core/src/cuda_autograd/mod.rs index 590e027df..aca4cebe7 100644 --- a/crates/ml-core/src/cuda_autograd/mod.rs +++ b/crates/ml-core/src/cuda_autograd/mod.rs @@ -7,7 +7,7 @@ //! //! ## Architecture //! -//! - [`GpuTensor`]: Thin wrapper around `CudaSlice` with shape metadata. +//! - [`GpuTensor`]: Thin wrapper around `CudaSlice` with shape metadata. //! - [`GpuVarStore`]: Named parameter store replacing `VarMap`/`VarBuilder`. //! - [`GpuLinear`]: Dense layer using cuBLAS sgemm for forward, manual backward. //! - [`GpuAdamW`]: Fused AdamW optimizer running entirely on GPU. @@ -23,7 +23,7 @@ //! This is more memory-efficient and CUDA Graph-friendly. //! //! 2. **Flat parameter buffer**: `GpuVarStore` can optionally flatten all parameters -//! into a contiguous `CudaSlice` for single-kernel Adam updates (matching +//! into a contiguous `CudaSlice` for single-kernel Adam updates (matching //! the DQN fused trainer pattern). //! //! 3. **cuBLAS for matmul**: Linear forward/backward use cuBLAS sgemm rather than diff --git a/crates/ml-core/src/cuda_autograd/optimizer.rs b/crates/ml-core/src/cuda_autograd/optimizer.rs index 4dfb5121f..d0cd12b10 100644 --- a/crates/ml-core/src/cuda_autograd/optimizer.rs +++ b/crates/ml-core/src/cuda_autograd/optimizer.rs @@ -55,9 +55,9 @@ impl Default for AdamWConfig { /// Per-parameter optimizer state (first and second moment estimates). struct ParamState { /// First moment (moving average of gradients). - m: CudaSlice, + m: CudaSlice, /// Second moment (moving average of squared gradients). - v: CudaSlice, + v: CudaSlice, } /// GPU-native AdamW optimizer. @@ -184,10 +184,10 @@ impl GpuAdamW { // Lazily allocate moment buffers if !self.states.contains_key(name) { - let m = self.stream.alloc_zeros::(n).map_err(|e| { + let m = self.stream.alloc_zeros::(n).map_err(|e| { MLError::ModelError(format!("AdamW alloc m '{name}': {e}")) })?; - let v = self.stream.alloc_zeros::(n).map_err(|e| { + let v = self.stream.alloc_zeros::(n).map_err(|e| { MLError::ModelError(format!("AdamW alloc v '{name}': {e}")) })?; self.states.insert(name.clone(), ParamState { m, v }); @@ -240,7 +240,7 @@ impl GpuAdamW { gradients: &BTreeMap, ) -> Result<(f32, f32), MLError> { // Accumulate squared norms per parameter into a scalar on GPU - let norm_sq_buf = self.stream.alloc_zeros::(1).map_err(|e| { + let norm_sq_buf = self.stream.alloc_zeros::(1).map_err(|e| { MLError::ModelError(format!("AdamW norm_sq alloc: {e}")) })?; diff --git a/crates/ml-core/src/cuda_autograd/reductions.rs b/crates/ml-core/src/cuda_autograd/reductions.rs index 52153740b..21bf9e7dc 100644 --- a/crates/ml-core/src/cuda_autograd/reductions.rs +++ b/crates/ml-core/src/cuda_autograd/reductions.rs @@ -78,13 +78,13 @@ impl ReductionKernels { /// The kernel performs a tree-reduction in shared memory (warp-level then block-level) /// with atomic writes to a 5-element global result buffer: [min, max, sum, sum_sq, count]. /// Mean and variance are derived on the CPU from the 5 raw values (20 bytes transferred). - pub fn stats(&self, data: &CudaSlice, n: usize) -> Result { + pub fn stats(&self, data: &CudaSlice, n: usize) -> Result { if n == 0 { return Err(MLError::InvalidInput("stats: empty buffer".to_string())); } // result[5] = [min, max, sum, sum_sq, count] - let mut result = self.stream.alloc_zeros::(5).map_err(|e| { + let mut result = self.stream.alloc_zeros::(5).map_err(|e| { MLError::ModelError(format!("stats result alloc: {e}")) })?; @@ -109,8 +109,8 @@ impl ReductionKernels { // SAFETY: Kernel arguments match the CUDA function signature exactly: // (const float* data, float* result, int n) - // data is a valid CudaSlice of at least n elements. - // result is a valid CudaSlice of 5 elements, pre-initialized. + // data is a valid CudaSlice of at least n elements. + // result is a valid CudaSlice of 5 elements, pre-initialized. unsafe { self.stream .launch_builder(&self.fused_stats_fn) @@ -151,13 +151,13 @@ impl ReductionKernels { /// Argmax of a flat buffer -- returns index as u32. /// /// Single kernel launch with shared-memory reduction, single 4-byte DtoH readback. - pub fn argmax(&self, data: &CudaSlice, n: usize) -> Result { + pub fn argmax(&self, data: &CudaSlice, n: usize) -> Result { if n == 0 { return Err(MLError::InvalidInput("argmax: empty buffer".to_string())); } // Output: [best_val, best_idx] as float2 (index packed into float bits) - let mut result = self.stream.alloc_zeros::(2).map_err(|e| { + let mut result = self.stream.alloc_zeros::(2).map_err(|e| { MLError::ModelError(format!("argmax result alloc: {e}")) })?; let init_vals = [f32::NEG_INFINITY, 0.0_f32]; // val, idx @@ -179,8 +179,8 @@ impl ReductionKernels { // SAFETY: Kernel arguments match the CUDA function signature exactly: // (const float* data, float* result, int n) - // data is a valid CudaSlice of at least n elements. - // result is a valid CudaSlice of 2 elements, pre-initialized. + // data is a valid CudaSlice of at least n elements. + // result is a valid CudaSlice of 2 elements, pre-initialized. unsafe { self.stream .launch_builder(&self.argmax_flat_fn) @@ -203,7 +203,7 @@ impl ReductionKernels { /// Argmax per row for 2D data [rows, cols] -- returns `Vec` of length `rows`. /// /// Each block processes one row. Single DtoH readback of `rows * 4` bytes. - pub fn argmax_rows(&self, data: &CudaSlice, rows: usize, cols: usize) -> Result, MLError> { + pub fn argmax_rows(&self, data: &CudaSlice, rows: usize, cols: usize) -> Result, MLError> { if rows == 0 || cols == 0 { return Err(MLError::InvalidInput("argmax_rows: empty data".to_string())); } @@ -225,7 +225,7 @@ impl ReductionKernels { // SAFETY: Kernel arguments match the CUDA function signature exactly: // (const float* data, unsigned int* result, int rows, int cols) - // data is a valid CudaSlice of at least rows*cols elements. + // data is a valid CudaSlice of at least rows*cols elements. // result is a valid CudaSlice of at least rows elements. unsafe { self.stream @@ -249,12 +249,12 @@ impl ReductionKernels { /// Sum of all elements -- single scalar readback. /// /// Tree-reduction in shared memory with atomic add to global accumulator. - pub fn sum(&self, data: &CudaSlice, n: usize) -> Result { + pub fn sum(&self, data: &CudaSlice, n: usize) -> Result { if n == 0 { return Err(MLError::InvalidInput("sum: empty buffer".to_string())); } - let mut result = self.stream.alloc_zeros::(1).map_err(|e| { + let mut result = self.stream.alloc_zeros::(1).map_err(|e| { MLError::ModelError(format!("sum result alloc: {e}")) })?; self.stream.memset_zeros(&mut result).map_err(|e| { @@ -275,8 +275,8 @@ impl ReductionKernels { // SAFETY: Kernel arguments match the CUDA function signature exactly: // (const float* data, float* result, int n) - // data is a valid CudaSlice of at least n elements. - // result is a valid CudaSlice of 1 element, zero-initialized. + // data is a valid CudaSlice of at least n elements. + // result is a valid CudaSlice of 1 element, zero-initialized. unsafe { self.stream .launch_builder(&self.sum_reduce_fn) @@ -298,12 +298,12 @@ impl ReductionKernels { /// Per-column sum for 2D data [rows, cols] -- returns `Vec` of length `cols`. /// /// One block per column, shared-memory tree reduction along the row axis. - pub fn col_sums(&self, data: &CudaSlice, rows: usize, cols: usize) -> Result, MLError> { + pub fn col_sums(&self, data: &CudaSlice, rows: usize, cols: usize) -> Result, MLError> { if rows == 0 || cols == 0 { return Err(MLError::InvalidInput("col_sums: empty data".to_string())); } - let mut result = self.stream.alloc_zeros::(cols).map_err(|e| { + let mut result = self.stream.alloc_zeros::(cols).map_err(|e| { MLError::ModelError(format!("col_sums result alloc: {e}")) })?; self.stream.memset_zeros(&mut result).map_err(|e| { @@ -323,8 +323,8 @@ impl ReductionKernels { // SAFETY: Kernel arguments match the CUDA function signature exactly: // (const float* data, float* result, int rows, int cols) - // data is a valid CudaSlice of at least rows*cols elements. - // result is a valid CudaSlice of at least cols elements. + // data is a valid CudaSlice of at least rows*cols elements. + // result is a valid CudaSlice of at least cols elements. unsafe { self.stream .launch_builder(&self.col_sum_fn) @@ -721,7 +721,7 @@ mod tests { let kernels = ReductionKernels::new(&stream).unwrap(); let data: Vec = (1..=100).map(|i| i as f32).collect(); - let mut gpu_data = stream.alloc_zeros::(100).unwrap(); + let mut gpu_data = stream.alloc_zeros::(100).unwrap(); stream.memcpy_htod(&data, &mut gpu_data).unwrap(); let stats = kernels.stats(&gpu_data, 100).unwrap(); @@ -743,7 +743,7 @@ mod tests { if let Some(v) = data.get_mut(42) { *v = 999.0; } - let mut gpu_data = stream.alloc_zeros::(1000).unwrap(); + let mut gpu_data = stream.alloc_zeros::(1000).unwrap(); stream.memcpy_htod(&data, &mut gpu_data).unwrap(); let idx = kernels.argmax(&gpu_data, 1000).unwrap(); @@ -763,7 +763,7 @@ mod tests { 1.0, 2.0, 3.0, 4.0, 8.0, // max at col 4 1.0, 2.0, 7.0, 3.0, 4.0, // max at col 2 ]; - let mut gpu_data = stream.alloc_zeros::(20).unwrap(); + let mut gpu_data = stream.alloc_zeros::(20).unwrap(); stream.memcpy_htod(&data, &mut gpu_data).unwrap(); let indices = kernels.argmax_rows(&gpu_data, 4, 5).unwrap(); @@ -776,7 +776,7 @@ mod tests { let kernels = ReductionKernels::new(&stream).unwrap(); let data: Vec = (1..=100).map(|i| i as f32).collect(); - let mut gpu_data = stream.alloc_zeros::(100).unwrap(); + let mut gpu_data = stream.alloc_zeros::(100).unwrap(); stream.memcpy_htod(&data, &mut gpu_data).unwrap(); let total = kernels.sum(&gpu_data, 100).unwrap(); @@ -794,7 +794,7 @@ mod tests { 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, ]; - let mut gpu_data = stream.alloc_zeros::(12).unwrap(); + let mut gpu_data = stream.alloc_zeros::(12).unwrap(); stream.memcpy_htod(&data, &mut gpu_data).unwrap(); let sums = kernels.col_sums(&gpu_data, 3, 4).unwrap(); diff --git a/crates/ml-core/src/cuda_autograd/stream_ops.rs b/crates/ml-core/src/cuda_autograd/stream_ops.rs index 466c3eccf..823ccf6ab 100644 --- a/crates/ml-core/src/cuda_autograd/stream_ops.rs +++ b/crates/ml-core/src/cuda_autograd/stream_ops.rs @@ -22,14 +22,14 @@ use crate::MLError; // ── Pointer helpers ─────────────────────────────────────────────────────── /// Extract raw CUDA device pointer from a `CudaSlice` (read-only). -fn raw_ptr(slice: &CudaSlice, stream: &CudaStream) -> u64 { +fn raw_ptr(slice: &CudaSlice, 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, stream: &CudaStream) -> u64 { +fn raw_ptr_mut(slice: &mut CudaSlice, stream: &CudaStream) -> u64 { let (ptr, guard) = slice.device_ptr_mut(stream); let _no_drop = ManuallyDrop::new(guard); ptr @@ -45,8 +45,8 @@ fn ew(stream: &Arc) -> Result, - dst: &mut CudaSlice, + src: &CudaSlice, + dst: &mut CudaSlice, src_offset: usize, dst_offset: usize, count: usize, @@ -85,7 +85,7 @@ fn dtod_copy( #[derive(Clone)] pub struct StreamTensor { /// Contiguous f32 data on device. Length = product of `shape`. - pub data: CudaSlice, + pub data: CudaSlice, /// Row-major shape, e.g. `[batch, features]`. pub shape: Vec, /// CUDA stream handle. @@ -111,7 +111,7 @@ impl StreamTensor { pub fn zeros(shape: &[usize], stream: &Arc) -> Result { let n: usize = shape.iter().product(); let data = stream - .alloc_zeros::(n) + .alloc_zeros::(n) .map_err(|e| MLError::DeviceError(format!("StreamTensor zeros alloc: {e}")))?; Ok(Self { data, @@ -201,14 +201,14 @@ impl StreamTensor { /// GPU-native linear layer: `y = x @ W^T + b`. /// -/// Weights are stored as `CudaSlice` and matrix multiplication is +/// Weights are stored as `CudaSlice` and matrix multiplication is /// dispatched via cuBLAS `sgemm`. Carries its own stream handle. #[derive(Debug, Clone)] pub struct StreamLinear { /// Weight matrix: shape (`out_features`, `in_features`) -- row-major. - pub weight: CudaSlice, + pub weight: CudaSlice, /// Bias vector: shape (`out_features`,). - pub bias: Option>, + pub bias: Option>, pub in_features: usize, pub out_features: usize, stream: Arc, @@ -248,8 +248,8 @@ impl StreamLinear { /// Create from existing weight/bias slices (e.g. loaded from checkpoint). pub fn from_slices( - weight: CudaSlice, - bias: Option>, + weight: CudaSlice, + bias: Option>, in_features: usize, out_features: usize, stream: &Arc, @@ -289,7 +289,7 @@ impl StreamLinear { let mut out = self .stream - .alloc_zeros::(m * n) + .alloc_zeros::(m * n) .map_err(|e| MLError::DeviceError(format!("StreamLinear output alloc: {e}")))?; let blas = cudarc::cublas::CudaBlas::new(Arc::clone(&self.stream)) @@ -369,7 +369,7 @@ impl StreamLinear { /// Helper: wrap a `CudaSlice` result from a kernel into a `StreamTensor`. #[inline] -fn wrap(data: CudaSlice, shape: &[usize], stream: &Arc) -> StreamTensor { +fn wrap(data: CudaSlice, shape: &[usize], stream: &Arc) -> StreamTensor { StreamTensor { data, shape: shape.to_vec(), @@ -526,7 +526,7 @@ pub fn gpu_cat_dim1(a: &StreamTensor, b: &StreamTensor) -> Result(total) + let mut out = a.stream.alloc_zeros::(total) .map_err(|e| MLError::DeviceError(format!("cat_dim1 alloc: {e}")))?; for r in 0..rows { @@ -556,7 +556,7 @@ pub fn gpu_cat_dim1(a: &StreamTensor, b: &StreamTensor) -> Result(total) + let mut out = a.stream.alloc_zeros::(total) .map_err(|e| MLError::DeviceError(format!("cat_dim1_3d alloc: {e}")))?; for bi in 0..batch { @@ -593,7 +593,7 @@ pub fn gpu_matmul(a: &StreamTensor, b: &StreamTensor) -> Result(m * n) + .alloc_zeros::(m * n) .map_err(|e| MLError::DeviceError(format!("matmul output alloc: {e}")))?; let blas = cudarc::cublas::CudaBlas::new(Arc::clone(&a.stream)) @@ -682,7 +682,7 @@ pub fn gpu_select_dim1(x: &StreamTensor, idx: usize) -> Result(out_elems) + let mut out = x.stream.alloc_zeros::(out_elems) .map_err(|e| MLError::DeviceError(format!("select_dim1 alloc: {e}")))?; for bi in 0..batch { @@ -731,7 +731,7 @@ pub fn gpu_scalar_sub(scalar: f32, x: &StreamTensor) -> Result) -> Result { let n: usize = shape.iter().product(); // Create zeros then use affine(mul=0, add=value) to fill - let zeros = stream.alloc_zeros::(n) + let zeros = stream.alloc_zeros::(n) .map_err(|e| MLError::DeviceError(format!("gpu_full alloc: {e}")))?; let k = ew(stream)?; let out = k.affine(&zeros, 0.0, value, n)?; @@ -760,7 +760,7 @@ pub fn gpu_gather_dim0(source: &StreamTensor, indices: &StreamTensor) -> Result< // Download only the indices (small), keep source on GPU let idx_host = indices.to_vec()?; // gpu-exit: small index array for gather let out_elems = n * s_cols; - let mut out = source.stream.alloc_zeros::(out_elems) + let mut out = source.stream.alloc_zeros::(out_elems) .map_err(|e| MLError::DeviceError(format!("gather_dim0 alloc: {e}")))?; // For each row in output, copy the corresponding source row via DtoD @@ -799,7 +799,7 @@ pub fn gpu_unsqueeze_broadcast(x: &StreamTensor, new_cols: usize) -> Result out[i] = x[i/cols] * 1.0 = x[i/cols] - let ones = x.stream.alloc_zeros::(total) + let ones = x.stream.alloc_zeros::(total) .map_err(|e| MLError::DeviceError(format!("unsqueeze_broadcast ones alloc: {e}")))?; // Fill ones using affine(mul=0, add=1) let k = ew(&x.stream)?; @@ -854,7 +854,7 @@ pub fn gpu_sum_dim(x: &StreamTensor, dim: usize) -> Result(cols) + let mut row_buf = x.stream.alloc_zeros::(cols) .map_err(|e| MLError::DeviceError(format!("sum_dim row alloc: {e}")))?; dtod_copy(&x.stream, &x.data, &mut row_buf, r * cols, 0, cols)?; let s = red.sum(&row_buf, cols)?; @@ -930,7 +930,7 @@ pub fn gpu_narrow_2d(x: &StreamTensor, dim: usize, start: usize, len: usize) -> } // Contiguous slice: just copy from offset let out_elems = len * cols; - let mut out = x.stream.alloc_zeros::(out_elems) + let mut out = x.stream.alloc_zeros::(out_elems) .map_err(|e| MLError::DeviceError(format!("narrow_2d alloc: {e}")))?; dtod_copy(&x.stream, &x.data, &mut out, start * cols, 0, out_elems)?; Ok(StreamTensor { @@ -944,7 +944,7 @@ pub fn gpu_narrow_2d(x: &StreamTensor, dim: usize, start: usize, len: usize) -> } // Non-contiguous: copy per-row let out_elems = rows * len; - let mut out = x.stream.alloc_zeros::(out_elems) + let mut out = x.stream.alloc_zeros::(out_elems) .map_err(|e| MLError::DeviceError(format!("narrow_2d_dim1 alloc: {e}")))?; for r in 0..rows { dtod_copy(&x.stream, &x.data, &mut out, r * cols + start, r * len, len)?; @@ -974,12 +974,12 @@ pub fn gpu_softmax(x: &StreamTensor) -> Result { // Per-row softmax: for each row, subtract max, exp, divide by sum. // Process row by row to use contiguous reduction. let total = rows * cols; - let mut out = x.stream.alloc_zeros::(total) + let mut out = x.stream.alloc_zeros::(total) .map_err(|e| MLError::DeviceError(format!("softmax alloc: {e}")))?; for r in 0..rows { // Extract row - let mut row_buf = x.stream.alloc_zeros::(cols) + let mut row_buf = x.stream.alloc_zeros::(cols) .map_err(|e| MLError::DeviceError(format!("softmax row alloc: {e}")))?; dtod_copy(&x.stream, &x.data, &mut row_buf, r * cols, 0, cols)?; @@ -1032,12 +1032,12 @@ pub fn gpu_layer_norm( let red = super::reductions::ReductionKernels::new(&x.stream)?; let total = rows * cols; - let mut out = x.stream.alloc_zeros::(total) + let mut out = x.stream.alloc_zeros::(total) .map_err(|e| MLError::DeviceError(format!("layer_norm alloc: {e}")))?; for r in 0..rows { // Extract row - let mut row_buf = x.stream.alloc_zeros::(cols) + let mut row_buf = x.stream.alloc_zeros::(cols) .map_err(|e| MLError::DeviceError(format!("layer_norm row alloc: {e}")))?; dtod_copy(&x.stream, &x.data, &mut row_buf, r * cols, 0, cols)?; @@ -1077,7 +1077,7 @@ pub fn gpu_mean_keepdim_last(x: &StreamTensor) -> Result let mut result_host = Vec::with_capacity(rows); for r in 0..rows { - let mut row_buf = x.stream.alloc_zeros::(cols) + let mut row_buf = x.stream.alloc_zeros::(cols) .map_err(|e| MLError::DeviceError(format!("mean_keepdim row alloc: {e}")))?; dtod_copy(&x.stream, &x.data, &mut row_buf, r * cols, 0, cols)?; let stats = red.stats(&row_buf, cols)?; @@ -1134,7 +1134,7 @@ pub fn gpu_stack_2d(tensors: &[StreamTensor]) -> Result { let elem_per_tensor = rows * cols; let total = n * elem_per_tensor; - let mut out = first.stream.alloc_zeros::(total) + let mut out = first.stream.alloc_zeros::(total) .map_err(|e| MLError::DeviceError(format!("stack_2d alloc: {e}")))?; for (i, t) in tensors.iter().enumerate() { @@ -1256,7 +1256,7 @@ pub fn gpu_cat_dim0(tensors: &[StreamTensor], stream: &Arc) -> Resul total_batch += t.dim(0)?; } - let mut out = stream.alloc_zeros::(total_elems) + let mut out = stream.alloc_zeros::(total_elems) .map_err(|e| MLError::DeviceError(format!("cat_dim0 alloc: {e}")))?; let mut offset = 0_usize; @@ -1290,7 +1290,7 @@ pub fn gpu_stack_tensors(tensors: &[StreamTensor], stream: &Arc) -> let elem_count = first.numel(); let total = tensors.len() * elem_count; - let mut out = stream.alloc_zeros::(total) + let mut out = stream.alloc_zeros::(total) .map_err(|e| MLError::DeviceError(format!("stack_tensors alloc: {e}")))?; for (i, t) in tensors.iter().enumerate() { diff --git a/crates/ml-core/src/cuda_autograd/var_store.rs b/crates/ml-core/src/cuda_autograd/var_store.rs index c11cf7bc2..a1e04105e 100644 --- a/crates/ml-core/src/cuda_autograd/var_store.rs +++ b/crates/ml-core/src/cuda_autograd/var_store.rs @@ -1,6 +1,6 @@ //! GPU variable store — replaces Candle's `VarMap`/`VarBuilder`. //! -//! Stores named weight tensors as `CudaSlice` on GPU. Provides: +//! Stores named weight tensors as `CudaSlice` on GPU. Provides: //! - Named parameter registration and lookup //! - Linear layer builder (returns [`GpuLinear`]) //! - Flat parameter view for single-kernel optimizer updates @@ -19,7 +19,7 @@ use super::linear::GpuLinear; /// A named GPU parameter (weight or bias). pub struct GpuParam { /// The GPU buffer holding this parameter's data. - pub data: CudaSlice, + pub data: CudaSlice, /// Shape of the parameter tensor. pub shape: Vec, } @@ -32,14 +32,14 @@ impl std::fmt::Debug for GpuParam { /// GPU variable store replacing Candle's `VarMap`. /// -/// Stores all trainable parameters as named `CudaSlice` buffers. +/// Stores all trainable parameters as named `CudaSlice` buffers. /// Parameters are ordered by insertion order (via `BTreeMap` on name). /// /// ## Flat parameter view /// /// For optimizers that operate on a single contiguous buffer (like the fused /// DQN Adam kernel), call [`flatten()`](GpuVarStore::flatten) to get all -/// parameters concatenated into one `CudaSlice`. After optimizer update, +/// parameters concatenated into one `CudaSlice`. After optimizer update, /// call [`unflatten()`](GpuVarStore::unflatten) to scatter back. pub struct GpuVarStore { /// Named parameters, ordered by name for deterministic iteration. @@ -106,7 +106,7 @@ impl GpuVarStore { pub fn register>( &mut self, name: S, - data: CudaSlice, + data: CudaSlice, shape: Vec, ) -> Result<(), MLError> { let name = name.into(); @@ -275,7 +275,7 @@ impl GpuVarStore { pub fn flatten(&self) -> Result<(GpuTensor, BTreeMap), MLError> { let total = self.total_elements(); if total == 0 { - let buf = self.stream.alloc_zeros::(1).map_err(|e| { + let buf = self.stream.alloc_zeros::(1).map_err(|e| { MLError::ModelError(format!("flatten alloc: {e}")) })?; return Ok((GpuTensor::from_flat(buf), BTreeMap::new())); diff --git a/crates/ml/build.rs b/crates/ml/build.rs index ae13ed023..5122d0949 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -66,17 +66,9 @@ fn main() { "monitoring_kernel.cu", "nstep_kernel.cu", "ppo_experience_kernel.cu", - // bias_kernels needs common header for f32_to_bf16_kernel (used by CublasForward) "bias_kernels.cu", - ]; - - // Standalone kernels: no common header needed. - // experience_kernels.cu has its own inline helpers and explicitly does NOT - // use common_device_functions.cuh. The remaining standalone kernels are - // small utility kernels that have no #include dependencies. - let standalone_kernels = [ + // Formerly standalone — now need common header for BF16 types "experience_kernels.cu", - // Inline kernels extracted from gpu_dqn_trainer.rs "ema_kernel.cu", "relu_mask_kernel.cu", "per_update_kernel.cu", @@ -85,12 +77,16 @@ fn main() { "expected_q_kernel.cu", "q_stats_kernel.cu", "cql_grad_kernel.cu", - // Inline kernels extracted from other modules "trade_stats_kernel.cu", "backward_kernels.cu", "iqn_cvar_kernel.cu", ]; + // Standalone kernels: no common header needed. + // experience_kernels.cu has its own inline helpers and explicitly does NOT + // use common_device_functions.cuh. The remaining standalone kernels are + // small utility kernels that have no #include dependencies. + // Compile kernels that need common header prepended for kernel_name in &kernels_with_common { compile_kernel( @@ -103,20 +99,8 @@ fn main() { ); } - // Compile standalone kernels (no common header) - for kernel_name in &standalone_kernels { - compile_kernel( - &nvcc, - kernel_dir, - kernel_name, - &arch, - &out_dir, - None, - ); - } - - eprintln!(" Precompiled {} CUDA kernels ({arch}) — ZERO runtime compilation", - kernels_with_common.len() + standalone_kernels.len()); + eprintln!(" Precompiled {} CUDA kernels ({arch}) — ZERO runtime compilation, all BF16", + kernels_with_common.len()); } /// Compile a single .cu kernel file to a .cubin via nvcc. diff --git a/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu b/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu index 51bd0f3e5..c5a9b8d8a 100644 --- a/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu +++ b/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu @@ -49,9 +49,9 @@ * @param B batch size */ extern "C" __global__ void attention_backward_kernel( - const float* __restrict__ d_output, - const float* __restrict__ states_input, - const float* __restrict__ params, + const __nv_bfloat16* __restrict__ d_output, + const __nv_bfloat16* __restrict__ states_input, + const __nv_bfloat16* __restrict__ params, float* d_input, float* d_params, int B, @@ -66,20 +66,20 @@ extern "C" __global__ void attention_backward_kernel( if (sample >= B) return; int tid = threadIdx.x; /* 0-31 warp lane */ - const float* x = states_input + sample * D; - const float* dy = d_output + sample * D; + const __nv_bfloat16* x = states_input + sample * D; + const __nv_bfloat16* dy = d_output + sample * D; float* dx = d_input + sample * D; /* Weight offsets in flat params buffer (same layout as forward) */ - const float* W_Q = params; - const float* W_K = W_Q + D * D; - const float* W_V = W_K + D * D; - const float* b_Q = W_V + D * D; - const float* b_K = b_Q + D; - const float* b_V = b_K + D; - const float* W_O = b_V + D; - const float* b_O = W_O + D * D; - const float* ln_gamma = b_O + D; + const __nv_bfloat16* W_Q = params; + const __nv_bfloat16* W_K = W_Q + D * D; + const __nv_bfloat16* W_V = W_K + D * D; + const __nv_bfloat16* b_Q = W_V + D * D; + const __nv_bfloat16* b_K = b_Q + D; + const __nv_bfloat16* b_V = b_K + D; + const __nv_bfloat16* W_O = b_V + D; + const __nv_bfloat16* b_O = W_O + D * D; + const __nv_bfloat16* ln_gamma = b_O + D; /* ln_beta = ln_gamma + D; (not needed in backward computation) */ /* Gradient offsets in d_params (same layout) */ @@ -112,12 +112,12 @@ extern "C" __global__ void attention_backward_kernel( /* Compute Q, K for attention scores */ for (int f = tid; f < Dh; f += 32) { int gf = head_start + f; - float q = b_Q[gf]; - float k = b_K[gf]; + float q = __bfloat162float(b_Q[gf]); + float k = __bfloat162float(b_K[gf]); for (int j = 0; j < D; j++) { - float xj = x[j]; - q += xj * W_Q[j * D + gf]; - k += xj * W_K[j * D + gf]; + float xj = __bfloat162float(x[j]); + q += xj * __bfloat162float(W_Q[j * D + gf]); + k += xj * __bfloat162float(W_K[j * D + gf]); } concat[gf] = q * k / sqrtf((float)Dh); } @@ -147,22 +147,22 @@ extern "C" __global__ void attention_backward_kernel( for (int f = tid; f < Dh; f += 32) { float attn = concat[head_start + f] * inv_sum; int gf = head_start + f; - float v = b_V[gf]; + float v = __bfloat162float(b_V[gf]); for (int j = 0; j < D; j++) - v += x[j] * W_V[j * D + gf]; + v += __bfloat162float(x[j]) * __bfloat162float(W_V[j * D + gf]); concat[gf] = attn * v; } __syncthreads(); } /* Now concat[0..D) is the recomputed multi-head attention output. - * Compute pre_ln[f] = x[f] + sum_j(concat[j] * W_O[j*D+f]) + b_O[f] + * Compute pre_ln[f] = __bfloat162float(x[f]) + sum_j(concat[j] * __bfloat162float(W_O[j*D+f])) + __bfloat162float(b_O[f]) * Store pre_ln in d_proj temporarily (we'll overwrite with d_projection later). */ for (int f = tid; f < D; f += 32) { - float proj = b_O[f]; + float proj = __bfloat162float(b_O[f]); for (int j = 0; j < D; j++) - proj += concat[j] * W_O[j * D + f]; - d_proj[f] = x[f] + proj; /* pre_ln = states + projection */ + proj += concat[j] * __bfloat162float(W_O[j * D + f]); + d_proj[f] = __bfloat162float(x[f]) + proj; /* pre_ln = states + projection */ } __syncthreads(); @@ -193,7 +193,7 @@ extern "C" __global__ void attention_backward_kernel( float inv_std = 1.0f / sqrtf(var + 1e-5f); /* LN backward: - * d_prenorm[f] = (1/D) * inv_std * (D * dy[f]*gamma[f] + * d_prenorm[f] = (1/D) * inv_std * (D * __bfloat162float(dy[f])*gamma[f] * - sum(dy*gamma) - x_hat[f] * sum(dy*gamma*x_hat)) * where x_hat = (pre_ln - mean) * inv_std */ @@ -202,7 +202,7 @@ extern "C" __global__ void attention_backward_kernel( float sum_dy_gamma_xhat = 0.0f; for (int f = tid; f < D; f += 32) { float x_hat = (d_proj[f] - mean) * inv_std; /* d_proj holds pre_ln */ - float dy_g = dy[f] * ln_gamma[f]; + float dy_g = __bfloat162float(dy[f]) * __bfloat162float(ln_gamma[f]); sum_dy_gamma += dy_g; sum_dy_gamma_xhat += dy_g * x_hat; } @@ -217,7 +217,7 @@ extern "C" __global__ void attention_backward_kernel( * Overwrite d_proj with d_prenorm (we're done with pre_ln). */ for (int f = tid; f < D; f += 32) { float x_hat = (d_proj[f] - mean) * inv_std; - float dy_g = dy[f] * ln_gamma[f]; + float dy_g = __bfloat162float(dy[f]) * __bfloat162float(ln_gamma[f]); float d_prenorm = inv_std / (float)D * ((float)D * dy_g - sum_dy_gamma - x_hat * sum_dy_gamma_xhat); @@ -225,8 +225,8 @@ extern "C" __global__ void attention_backward_kernel( d_proj[f] = d_prenorm; /* overwrite: now d_proj holds d_prenorm */ /* LN parameter gradients */ - atomicAdd(&d_ln_gamma[f], dy[f] * x_hat); - atomicAdd(&d_ln_beta[f], dy[f]); + atomicAdd(&d_ln_gamma[f], __bfloat162float(dy[f]) * x_hat); + atomicAdd(&d_ln_beta[f], __bfloat162float(dy[f])); } __syncthreads(); @@ -248,8 +248,8 @@ extern "C" __global__ void attention_backward_kernel( /* =================================================================== * Step 4: Output projection backward * - * Forward: projection[f] = sum_j(concat[j] * W_O[j*D + f]) + b_O[f] - * d_concat[j] = sum_f(d_proj[f] * W_O[j*D + f]) + * Forward: projection[f] = sum_j(concat[j] * __bfloat162float(W_O[j*D + f])) + __bfloat162float(b_O[f]) + * d_concat[j] = sum_f(d_proj[f] * __bfloat162float(W_O[j*D + f])) * dW_O[j*D+f] += concat[j] * d_proj[f] * db_O[f] += d_proj[f] * @@ -280,7 +280,7 @@ extern "C" __global__ void attention_backward_kernel( for (int j = tid; j < D; j += 32) { float d_c = 0.0f; for (int f = 0; f < D; f++) - d_c += d_proj[f] * W_O[j * D + f]; + d_c += d_proj[f] * __bfloat162float(W_O[j * D + f]); /* Temporarily store in register; we need to sync before overwriting d_proj */ concat[j] = d_c; /* repurpose concat as d_concat temporarily */ } @@ -311,12 +311,12 @@ extern "C" __global__ void attention_backward_kernel( * Use concat[head_start..] as scratch for attn values. */ for (int f = tid; f < Dh; f += 32) { int gf = head_start + f; - float q = b_Q[gf]; - float k = b_K[gf]; + float q = __bfloat162float(b_Q[gf]); + float k = __bfloat162float(b_K[gf]); for (int j = 0; j < D; j++) { - float xj = x[j]; - q += xj * W_Q[j * D + gf]; - k += xj * W_K[j * D + gf]; + float xj = __bfloat162float(x[j]); + q += xj * __bfloat162float(W_Q[j * D + gf]); + k += xj * __bfloat162float(W_K[j * D + gf]); } concat[gf] = q * k / sqrtf((float)Dh); } @@ -351,9 +351,9 @@ extern "C" __global__ void attention_backward_kernel( float attn = concat[gf] * inv_sum; /* Recompute V */ - float v = b_V[gf]; + float v = __bfloat162float(b_V[gf]); for (int j = 0; j < D; j++) - v += x[j] * W_V[j * D + gf]; + v += __bfloat162float(x[j]) * __bfloat162float(W_V[j * D + gf]); float d_attn = d_proj[gf] * v; /* d_proj holds d_concat */ sum_da_a += d_attn * attn; @@ -368,14 +368,14 @@ extern "C" __global__ void attention_backward_kernel( float attn = concat[gf] * inv_sum; /* Recompute Q, K, V */ - float q = b_Q[gf]; - float k = b_K[gf]; - float v = b_V[gf]; + float q = __bfloat162float(b_Q[gf]); + float k = __bfloat162float(b_K[gf]); + float v = __bfloat162float(b_V[gf]); for (int j = 0; j < D; j++) { - float xj = x[j]; - q += xj * W_Q[j * D + gf]; - k += xj * W_K[j * D + gf]; - v += xj * W_V[j * D + gf]; + float xj = __bfloat162float(x[j]); + q += xj * __bfloat162float(W_Q[j * D + gf]); + k += xj * __bfloat162float(W_K[j * D + gf]); + v += xj * __bfloat162float(W_V[j * D + gf]); } float d_attn_v = d_proj[gf]; /* d_concat[gf] */ @@ -396,9 +396,9 @@ extern "C" __global__ void attention_backward_kernel( atomicAdd(&db_K[gf], d_k); atomicAdd(&db_V[gf], d_v); - /* -- Weight gradients: dW[j*D+gf] += x[j] * d_proj -- */ + /* -- Weight gradients: dW[j*D+gf] += __bfloat162float(x[j]) * d_proj -- */ for (int j = 0; j < D; j++) { - float xj = x[j]; + float xj = __bfloat162float(x[j]); atomicAdd(&dW_Q[j * D + gf], xj * d_q); atomicAdd(&dW_K[j * D + gf], xj * d_k); atomicAdd(&dW_V[j * D + gf], xj * d_v); @@ -406,9 +406,9 @@ extern "C" __global__ void attention_backward_kernel( /* -- Input gradients: dx[j] += d_q*W_Q + d_k*W_K + d_v*W_V -- */ for (int j = 0; j < D; j++) { - float grad_j = d_q * W_Q[j * D + gf] - + d_k * W_K[j * D + gf] - + d_v * W_V[j * D + gf]; + float grad_j = d_q * __bfloat162float(W_Q[j * D + gf]) + + d_k * __bfloat162float(W_K[j * D + gf]) + + d_v * __bfloat162float(W_V[j * D + gf]); atomicAdd(&dx[j], grad_j); } } @@ -425,8 +425,8 @@ extern "C" __global__ void attention_backward_kernel( * Launch: grid=(ceil(total_params/256), 1, 1), block=(256, 1, 1) * ===================================================================== */ extern "C" __global__ void attn_grad_norm_kernel( - const float* __restrict__ grads, - float* __restrict__ norm_out, /* [1] */ + const __nv_bfloat16* __restrict__ grads, + __nv_bfloat16* __restrict__ norm_out, /* [1] */ int total_params ) { int tid = blockIdx.x * blockDim.x + threadIdx.x; @@ -460,11 +460,11 @@ extern "C" __global__ void attn_grad_norm_kernel( * Launch: grid=(ceil(total_params/256), 1, 1), block=(256, 1, 1) * ===================================================================== */ extern "C" __global__ void attn_adam_kernel( - float* __restrict__ params, - float* __restrict__ grads, - float* __restrict__ m, - float* __restrict__ v, - const float* __restrict__ norm, /* [1] gradient L2 norm */ + __nv_bfloat16* __restrict__ params, + __nv_bfloat16* __restrict__ grads, + __nv_bfloat16* __restrict__ m, + __nv_bfloat16* __restrict__ v, + const __nv_bfloat16* __restrict__ norm, /* [1] gradient L2 norm */ float lr, float beta1, float beta2, float eps, float weight_decay, float max_grad_norm, int adam_t, diff --git a/crates/ml/src/cuda_pipeline/attention_kernel.cu b/crates/ml/src/cuda_pipeline/attention_kernel.cu index 51c517b79..6a79a0790 100644 --- a/crates/ml/src/cuda_pipeline/attention_kernel.cu +++ b/crates/ml/src/cuda_pipeline/attention_kernel.cu @@ -60,8 +60,8 @@ * @param B batch size */ extern "C" __global__ void multihead_feature_attention( - const float* __restrict__ states, - const float* __restrict__ params, + const __nv_bfloat16* __restrict__ states, + const __nv_bfloat16* __restrict__ params, float* output, int B, int state_dim, /* runtime state dimension (replaces ATTN_STATE_DIM) */ @@ -75,20 +75,20 @@ extern "C" __global__ void multihead_feature_attention( if (sample >= B) return; int tid = threadIdx.x; /* 0-31 warp lane */ - const float* x = states + sample * D; + const __nv_bfloat16* x = states + sample * D; float* out = output + sample * D; /* Weight offsets in flat params buffer */ - const float* W_Q = params; /* [D, D] */ - const float* W_K = W_Q + D * D; /* [D, D] */ - const float* W_V = W_K + D * D; /* [D, D] */ - const float* b_Q = W_V + D * D; /* [D] */ - const float* b_K = b_Q + D; /* [D] */ - const float* b_V = b_K + D; /* [D] */ - const float* W_O = b_V + D; /* [D, D] */ - const float* b_O = W_O + D * D; /* [D] */ - const float* ln_gamma = b_O + D; /* [D] */ - const float* ln_beta = ln_gamma + D; /* [D] */ + const __nv_bfloat16* W_Q = params; /* [D, D] */ + const __nv_bfloat16* W_K = W_Q + D * D; /* [D, D] */ + const __nv_bfloat16* W_V = W_K + D * D; /* [D, D] */ + const __nv_bfloat16* b_Q = W_V + D * D; /* [D] */ + const __nv_bfloat16* b_K = b_Q + D; /* [D] */ + const __nv_bfloat16* b_V = b_K + D; /* [D] */ + const __nv_bfloat16* W_O = b_V + D; /* [D, D] */ + const __nv_bfloat16* b_O = W_O + D * D; /* [D] */ + const __nv_bfloat16* ln_gamma = b_O + D; /* [D] */ + const __nv_bfloat16* ln_beta = ln_gamma + D; /* [D] */ /* ── Step 1: Compute Q, K, V projections ── */ /* Each thread handles D/32 ≈ 2-3 features */ @@ -107,14 +107,14 @@ extern "C" __global__ void multihead_feature_attention( /* Compute Q_h, K_h, V_h for features this thread handles */ for (int f = tid; f < Dh; f += 32) { int global_f = head_start + f; - float q = b_Q[global_f]; - float k = b_K[global_f]; - float v = b_V[global_f]; + float q = __bfloat162float(b_Q[global_f]); + float k = __bfloat162float(b_K[global_f]); + float v = __bfloat162float(b_V[global_f]); for (int j = 0; j < D; j++) { - float xj = x[j]; - q += xj * W_Q[j * D + global_f]; - k += xj * W_K[j * D + global_f]; - v += xj * W_V[j * D + global_f]; + float xj = __bfloat162float(x[j]); + q += xj * __bfloat162float(W_Q[j * D + global_f]); + k += xj * __bfloat162float(W_K[j * D + global_f]); + v += xj * __bfloat162float(W_V[j * D + global_f]); } /* Attention score: Q[f] * K[f] / sqrt(Dh) */ @@ -151,7 +151,7 @@ extern "C" __global__ void multihead_feature_attention( /* Recompute V[f] */ int global_f = head_start + f; - float v = b_V[global_f]; + float v = __bfloat162float(b_V[global_f]); for (int j = 0; j < D; j++) { v += x[j] * W_V[j * D + global_f]; } @@ -163,18 +163,18 @@ extern "C" __global__ void multihead_feature_attention( /* ── Step 2: Output projection + residual ── */ /* out[f] = x[f] + sum_j(concat[j] * W_O[j * D + f]) + b_O[f] */ for (int f = tid; f < D; f += 32) { - float proj = b_O[f]; + float proj = __bfloat162float(b_O[f]); for (int j = 0; j < D; j++) { - proj += shared_concat[j] * W_O[j * D + f]; + proj += shared_concat[j] * __bfloat162float(W_O[j * D + f]); } - out[f] = x[f] + proj; /* residual connection */ + out[f] = __bfloat162float(x[f]) + proj; /* residual connection */ } __syncthreads(); /* ── Step 3: Layer normalization ── */ /* mean = avg(out[f]) */ float local_mean = 0.0f; - for (int f = tid; f < D; f += 32) local_mean += out[f]; + for (int f = tid; f < D; f += 32) local_mean += out[f]; /* out is float* - OK */ for (int mask = 16; mask >= 1; mask >>= 1) local_mean += __shfl_xor_sync(0xFFFFFFFF, local_mean, mask); float mean = __shfl_sync(0xFFFFFFFF, local_mean / (float)D, 0); @@ -191,7 +191,7 @@ extern "C" __global__ void multihead_feature_attention( /* Normalize + affine */ for (int f = tid; f < D; f += 32) { - out[f] = ln_gamma[f] * (out[f] - mean) * inv_std + ln_beta[f]; + out[f] = __bfloat162float(ln_gamma[f]) * (out[f] - mean) * inv_std + __bfloat162float(ln_beta[f]); } } diff --git a/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu index 78bf651c4..1bd682f88 100644 --- a/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu +++ b/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu @@ -43,7 +43,7 @@ __device__ void handle_capital_floor_breach( extern "C" __global__ void backtest_env_step( // Market data (read-only, uploaded once) - const float* __restrict__ prices, // [n_windows * max_len * 4] (OHLC) + const __nv_bfloat16* __restrict__ prices, // [n_windows * max_len * 4] (OHLC) const int* __restrict__ window_lens, // [n_windows] // Actions from model for current step @@ -97,10 +97,10 @@ extern "C" __global__ void backtest_env_step( // Read current prices int price_base = (w * max_len + current_step) * 4; - float open = prices[price_base + 0]; - float high = prices[price_base + 1]; - float low = prices[price_base + 2]; - float close = prices[price_base + 3]; + float open = __bfloat162float(prices[price_base + 0]); + float high = __bfloat162float(prices[price_base + 1]); + float low = __bfloat162float(prices[price_base + 2]); + float close = __bfloat162float(prices[price_base + 3]); // Read portfolio state from shared memory tile float value = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 0]; diff --git a/crates/ml/src/cuda_pipeline/backtest_forward_ppo_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_forward_ppo_kernel.cu index 1d9386eed..4dcd9a98a 100644 --- a/crates/ml/src/cuda_pipeline/backtest_forward_ppo_kernel.cu +++ b/crates/ml/src/cuda_pipeline/backtest_forward_ppo_kernel.cu @@ -27,14 +27,14 @@ #define ACTIONS_PER_BIN 9 /* 3 order × 3 urgency */ extern "C" __global__ void backtest_forward_ppo_kernel( - const float* __restrict__ states, /* [N, state_dim] */ + const __nv_bfloat16* __restrict__ states, /* [N, state_dim] */ /* Actor weights (6 pointers) */ - const float* __restrict__ pw1, /* [ACTOR_H1, state_dim] = [128, S_DIM] */ - const float* __restrict__ pb1, /* [ACTOR_H1] = [128] */ - const float* __restrict__ pw2, /* [ACTOR_H2, ACTOR_H1] = [64, 128] */ - const float* __restrict__ pb2, /* [ACTOR_H2] = [64] */ - const float* __restrict__ pw3, /* [num_actions, ACTOR_H2] = [45, 64] */ - const float* __restrict__ pb3, /* [num_actions] = [45] */ + const __nv_bfloat16* __restrict__ pw1, /* [ACTOR_H1, state_dim] = [128, S_DIM] */ + const __nv_bfloat16* __restrict__ pb1, /* [ACTOR_H1] = [128] */ + const __nv_bfloat16* __restrict__ pw2, /* [ACTOR_H2, ACTOR_H1] = [64, 128] */ + const __nv_bfloat16* __restrict__ pb2, /* [ACTOR_H2] = [64] */ + const __nv_bfloat16* __restrict__ pw3, /* [num_actions, ACTOR_H2] = [45, 64] */ + const __nv_bfloat16* __restrict__ pb3, /* [num_actions] = [45] */ /* Output */ int* __restrict__ out_actions, /* [N] greedy exposure action (0-4) */ int N, @@ -48,7 +48,7 @@ extern "C" __global__ void backtest_forward_ppo_kernel( float state[MAX_STATE_DIM]; int base = window_id * state_dim; for (int i = 0; i < state_dim; i++) { - state[i] = states[base + i]; + state[i] = __bfloat162float(states[base + i]); } /* ---- Actor forward: state → logits ---- */ diff --git a/crates/ml/src/cuda_pipeline/backtest_forward_supervised_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_forward_supervised_kernel.cu index 812d65287..3459adc37 100644 --- a/crates/ml/src/cuda_pipeline/backtest_forward_supervised_kernel.cu +++ b/crates/ml/src/cuda_pipeline/backtest_forward_supervised_kernel.cu @@ -19,7 +19,7 @@ */ extern "C" __global__ void signal_to_action_kernel( - const float* __restrict__ predictions, /* [N] or [N,1] scalar signals */ + const __nv_bfloat16* __restrict__ predictions, /* [N] or [N,1] scalar signals */ int* __restrict__ out_actions, /* [N] action indices 0-4 */ float high_threshold_bps, float low_threshold_bps, @@ -28,7 +28,7 @@ extern "C" __global__ void signal_to_action_kernel( int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= N) return; - float pred = predictions[idx]; + float pred = __bfloat162float(predictions[idx]); int action; if (pred < -high_threshold_bps) { diff --git a/crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu index 99b36f094..9c7b52e66 100644 --- a/crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu +++ b/crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu @@ -32,8 +32,8 @@ // Per window: return, volatility, volume_trend, momentum extern "C" __global__ void gather_states( - const float* __restrict__ features, // [n_windows * max_len * feat_dim] - const float* __restrict__ portfolio, // [n_windows * 8] + const __nv_bfloat16* __restrict__ features, // [n_windows * max_len * feat_dim] + const __nv_bfloat16* __restrict__ portfolio, // [n_windows * 8] float* states_out, // [n_windows * state_dim] int n_windows, int max_len, @@ -55,7 +55,7 @@ extern "C" __global__ void gather_states( int ps = w * 8; if (w < n_windows) { for (int i = 0; i < 8; i++) { - shmem_portfolio[local_tid * 8 + i] = __ldg(&portfolio[ps + i]); + shmem_portfolio[local_tid * 8 + i] = __bfloat162float(__ldg(&portfolio[ps + i])); } } __syncthreads(); @@ -79,13 +79,13 @@ extern "C" __global__ void gather_states( // ── 1. Market features [0 .. market_dim) ── // int i = 0; for (; i + 3 < market_dim; i += 4) { - states_out[out_base + i] = __ldg(&features[feat_base + i]); - states_out[out_base + i + 1] = __ldg(&features[feat_base + i + 1]); - states_out[out_base + i + 2] = __ldg(&features[feat_base + i + 2]); - states_out[out_base + i + 3] = __ldg(&features[feat_base + i + 3]); + states_out[out_base + i] = __bfloat162float(__ldg(&features[feat_base + i])); + states_out[out_base + i + 1] = __bfloat162float(__ldg(&features[feat_base + i + 1])); + states_out[out_base + i + 2] = __bfloat162float(__ldg(&features[feat_base + i + 2])); + states_out[out_base + i + 3] = __bfloat162float(__ldg(&features[feat_base + i + 3])); } for (; i < market_dim; i++) { - states_out[out_base + i] = __ldg(&features[feat_base + i]); + states_out[out_base + i] = __bfloat162float(__ldg(&features[feat_base + i])); } // ── 2. Portfolio features [market_dim .. market_dim+8) ── // @@ -100,7 +100,7 @@ extern "C" __global__ void gather_states( // Unrealized PnL: position * (current_close - entry_price) // Use feature[0] as close price proxy (same as training kernel) - float close_now = __ldg(&features[feat_base]); + float close_now = __bfloat162float(__ldg(&features[feat_base])); float unrealized_pnl = (entry_price > 0.0f && position != 0.0f) ? position * (close_now - entry_price) : 0.0f; @@ -147,8 +147,8 @@ extern "C" __global__ void gather_states( if (past_step >= 0 && slot + 3 < state_dim) { int now_feat_off = (window_bar_base + current_step) * feat_dim; int past_feat_off = (window_bar_base + past_step) * feat_dim; - float close_cur = __ldg(&features[now_feat_off]); - float close_past = __ldg(&features[past_feat_off]); + float close_cur = __bfloat162float(__ldg(&features[now_feat_off])); + float close_past = __bfloat162float(__ldg(&features[past_feat_off])); // Return over N bars float ret = (close_past > 0.0f) ? (close_cur - close_past) / close_past : 0.0f; @@ -162,12 +162,12 @@ extern "C" __global__ void gather_states( int vol_count = 0; for (int j = past_step; j <= current_step; j++) { int j_off = (window_bar_base + j) * feat_dim; - float v = __ldg(&features[j_off]); + float v = __bfloat162float(__ldg(&features[j_off])); if (v > max_val) max_val = v; if (v < min_val) min_val = v; // Volume proxy: feature index 4 if (feat_dim > 4) { - vol_sum += __ldg(&features[j_off + 4]); + vol_sum += __bfloat162float(__ldg(&features[j_off + 4])); vol_count++; } } @@ -177,7 +177,7 @@ extern "C" __global__ void gather_states( // Volume trend: current vs average float avg_vol = (vol_count > 0) ? vol_sum / (float)vol_count : 1.0f; - float cur_vol = (feat_dim > 4) ? __ldg(&features[now_feat_off + 4]) : 1.0f; + float cur_vol = (feat_dim > 4) ? __bfloat162float(__ldg(&features[now_feat_off + 4])) : 1.0f; float vol_ratio = (avg_vol > 0.0f) ? cur_vol / avg_vol : 1.0f; states_out[out_base + slot + 2] = fmaxf(0.0f, fminf(5.0f, vol_ratio)); @@ -202,13 +202,13 @@ extern "C" __global__ void gather_states( int ofi_out_base = market_dim + 8 + 16; // OFI features are stored after market features in the feature vector for (i = 0; i + 3 < ofi_dim; i += 4) { - states_out[out_base + ofi_out_base + i] = __ldg(&features[feat_base + market_dim + i]); - states_out[out_base + ofi_out_base + i + 1] = __ldg(&features[feat_base + market_dim + i + 1]); - states_out[out_base + ofi_out_base + i + 2] = __ldg(&features[feat_base + market_dim + i + 2]); - states_out[out_base + ofi_out_base + i + 3] = __ldg(&features[feat_base + market_dim + i + 3]); + states_out[out_base + ofi_out_base + i] = __bfloat162float(__ldg(&features[feat_base + market_dim + i])); + states_out[out_base + ofi_out_base + i + 1] = __bfloat162float(__ldg(&features[feat_base + market_dim + i + 1])); + states_out[out_base + ofi_out_base + i + 2] = __bfloat162float(__ldg(&features[feat_base + market_dim + i + 2])); + states_out[out_base + ofi_out_base + i + 3] = __bfloat162float(__ldg(&features[feat_base + market_dim + i + 3])); } for (; i < ofi_dim; i++) { - states_out[out_base + ofi_out_base + i] = __ldg(&features[feat_base + market_dim + i]); + states_out[out_base + ofi_out_base + i] = __bfloat162float(__ldg(&features[feat_base + market_dim + i])); } } diff --git a/crates/ml/src/cuda_pipeline/backtest_metrics_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_metrics_kernel.cu index 9e9ce9712..d3564eb5a 100644 --- a/crates/ml/src/cuda_pipeline/backtest_metrics_kernel.cu +++ b/crates/ml/src/cuda_pipeline/backtest_metrics_kernel.cu @@ -18,8 +18,8 @@ // [13] unique_action_count (popcount of 9-bit mask — how many exposure levels 0..8 were used) extern "C" __global__ void compute_backtest_metrics( - const float* __restrict__ step_returns, // [n_windows * max_len] - const float* __restrict__ portfolio_state, // [n_windows * 8] + const __nv_bfloat16* __restrict__ step_returns, // [n_windows * max_len] + const __nv_bfloat16* __restrict__ portfolio_state, // [n_windows * 8] const int* __restrict__ window_lens, // [n_windows] const int* __restrict__ actions_history, // [n_windows * max_len] for trade counting float* metrics_out, // [n_windows * 14] @@ -86,7 +86,7 @@ extern "C" __global__ void compute_backtest_metrics( if (chunk_end > wlen) chunk_end = wlen; for (int i = chunk_start; i < chunk_end; i++) { - float r = step_returns[base + i]; + float r = __bfloat162float(step_returns[base + i]); // Sharpe/Sortino use arithmetic mean of per-bar returns local_sum += r; local_sq += r * r; @@ -216,7 +216,7 @@ extern "C" __global__ void compute_backtest_metrics( float exact_peak = 1.0f; float exact_max_dd = 0.0f; for (int i = 0; i < wlen; i++) { - float r = step_returns[base + i]; + float r = __bfloat162float(step_returns[base + i]); exact_equity *= (1.0f + r); exact_peak = fmaxf(exact_peak, exact_equity); float dd = (exact_peak - exact_equity) / fmaxf(exact_peak, 1e-8f); @@ -338,7 +338,7 @@ extern "C" __global__ void compute_backtest_metrics( int sample_step = (wlen + sort_len - 1) / sort_len; // ceil(wlen / sort_len) for (int i = tid; i < sort_len; i += stride) { int src = i * sample_step; - s_sorted[i] = (src < wlen) ? step_returns[base + src] : 0.0f; + s_sorted[i] = (src < wlen) ? __bfloat162float(step_returns[base + src]) : 0.0f; } // Compute next power-of-two for bitonic sort diff --git a/crates/ml/src/cuda_pipeline/backward_kernels.cu b/crates/ml/src/cuda_pipeline/backward_kernels.cu index 9a50bf06e..7c08f3f72 100644 --- a/crates/ml/src/cuda_pipeline/backward_kernels.cu +++ b/crates/ml/src/cuda_pipeline/backward_kernels.cu @@ -10,18 +10,18 @@ */ extern "C" __global__ void relu_mask_kernel( - float* __restrict__ dy, - const float* __restrict__ activation, + __nv_bfloat16* __restrict__ dy, + const __nv_bfloat16* __restrict__ activation, int n) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= n) return; - if (activation[i] <= 0.0f) dy[i] = 0.0f; + if (__bfloat162float(activation[i]) <= 0.0f) dy[i] = __float2bfloat16(0.0f); } extern "C" __global__ void bias_grad_reduce_kernel( - const float* __restrict__ dy, - float* __restrict__ db, + const __nv_bfloat16* __restrict__ dy, + __nv_bfloat16* __restrict__ db, int out_dim, int batch_size) { @@ -29,7 +29,16 @@ extern "C" __global__ void bias_grad_reduce_kernel( if (j >= out_dim) return; float sum = 0.0f; for (int b = 0; b < batch_size; b++) { - sum += dy[b * out_dim + j]; + sum += __bfloat162float(dy[b * out_dim + j]); } - atomicAdd(&db[j], sum); + /* BF16 atomicAdd via CAS loop */ + unsigned short* addr = (unsigned short*)&db[j]; + unsigned short old_bits = *addr; + unsigned short assumed; + do { + assumed = old_bits; + __nv_bfloat16 old_bf = *(__nv_bfloat16*)&assumed; + __nv_bfloat16 new_bf = __float2bfloat16(__bfloat162float(old_bf) + sum); + old_bits = atomicCAS(addr, assumed, *(unsigned short*)&new_bf); + } while (old_bits != assumed); } diff --git a/crates/ml/src/cuda_pipeline/batched_backward.rs b/crates/ml/src/cuda_pipeline/batched_backward.rs index 35818a4e5..829445fa5 100644 --- a/crates/ml/src/cuda_pipeline/batched_backward.rs +++ b/crates/ml/src/cuda_pipeline/batched_backward.rs @@ -816,7 +816,7 @@ fn compile_backward_kernels( // private functions. /// Extract raw F32 device pointer from a CudaSlice (read-only). -pub(crate) fn raw_f32_ptr(slice: &CudaSlice, stream: &Arc) -> u64 { +pub(crate) fn raw_f32_ptr(slice: &CudaSlice, stream: &Arc) -> u64 { let (ptr, guard) = slice.device_ptr(stream); let _no_drop = ManuallyDrop::new(guard); ptr @@ -832,16 +832,16 @@ pub fn alloc_backward_scratch( stream: &Arc, config: &GpuDqnTrainConfig, ) -> Result<( - CudaSlice, // d_h_s2 [B, SH2] - CudaSlice, // d_h_s1 [B, SH1] - CudaSlice, // d_h_v [B, VH] - CudaSlice, // d_h_b0 [B, AH] - CudaSlice, // d_h_b1 [B, AH] - CudaSlice, // d_h_b2 [B, AH] + CudaSlice, // d_h_s2 [B, SH2] + CudaSlice, // d_h_s1 [B, SH1] + CudaSlice, // d_h_v [B, VH] + CudaSlice, // d_h_b0 [B, AH] + CudaSlice, // d_h_b1 [B, AH] + CudaSlice, // d_h_b2 [B, AH] ), MLError> { let b = config.batch_size; - let alloc = |n: usize| -> Result, MLError> { - stream.alloc_zeros::(n) + let alloc = |n: usize| -> Result, MLError> { + stream.alloc_zeros::(n) .map_err(|e| MLError::ModelError(format!("backward scratch alloc [{n}]: {e}"))) }; diff --git a/crates/ml/src/cuda_pipeline/batched_forward.rs b/crates/ml/src/cuda_pipeline/batched_forward.rs index 3be500d57..800db37c6 100644 --- a/crates/ml/src/cuda_pipeline/batched_forward.rs +++ b/crates/ml/src/cuda_pipeline/batched_forward.rs @@ -120,7 +120,7 @@ pub struct CublasForward { f32_to_bf16_kernel: CudaFunction, // ── BF16 activation buffers for tensor core forward path ───────── - // Internal BF16 buffers needed because forward_online takes &CudaSlice + // Internal BF16 buffers needed because forward_online takes &CudaSlice // and we cannot change that signature without cascading to callers. // The gemmex_bf16 path reads/writes these instead. @@ -322,19 +322,19 @@ impl CublasForward { &self, stream: &Arc, // ── Inputs ────────────────────────────────────────────────────── - states_buf: &CudaSlice, // [B, SD] + states_buf: &CudaSlice, // [B, SD] // ── F32 weight pointers (raw u64 into flat params_buf) ───────── w_ptrs: &[u64; 20], // [w_s1, b_s1, w_s2, b_s2, ...] // ── Activation save buffers (written by this function) ───────── - save_h_s1: &CudaSlice, // [B, SH1] - save_h_s2: &CudaSlice, // [B, SH2] - save_h_v: &CudaSlice, // [B, VH] - save_h_b0: &CudaSlice, // [B, AH] - save_h_b1: &CudaSlice, // [B, AH] - save_h_b2: &CudaSlice, // [B, AH] + save_h_s1: &CudaSlice, // [B, SH1] + save_h_s2: &CudaSlice, // [B, SH2] + save_h_v: &CudaSlice, // [B, VH] + save_h_b0: &CudaSlice, // [B, AH] + save_h_b1: &CudaSlice, // [B, AH] + save_h_b2: &CudaSlice, // [B, AH] // ── Logit output buffers ──────────────────────────────────────── - v_logits_buf: &CudaSlice, // [B, NA] — value head logits - b_logits_buf: &CudaSlice, // [B, (B0+B1+B2)*NA] — branch logits + v_logits_buf: &CudaSlice, // [B, NA] — value head logits + b_logits_buf: &CudaSlice, // [B, (B0+B1+B2)*NA] — branch logits ) -> Result<(), MLError> { let b = self.batch_size; @@ -465,7 +465,7 @@ impl CublasForward { &self, stream: &Arc, // ── Inputs ────────────────────────────────────────────────────── - states_f32: &CudaSlice, // [B, SD] F32 from experience collector + states_f32: &CudaSlice, // [B, SD] F32 from experience collector // ── BF16 weight pointers (raw u64 into flat bf16_params_buf) ── bf16_w_ptrs: &[u64; 20], // [w_s1, b_s1, w_s2, b_s2, ...] ) -> Result<(), MLError> { @@ -615,7 +615,7 @@ impl CublasForward { pub fn forward_target_bf16( &self, stream: &Arc, - next_states_f32: &CudaSlice, // [B, SD] F32 from experience collector + next_states_f32: &CudaSlice, // [B, SD] F32 from experience collector bf16_w_ptrs: &[u64; 20], // BF16 target weight pointers ) -> Result<(), MLError> { let b = self.batch_size; @@ -787,17 +787,17 @@ impl CublasForward { &self, stream: &Arc, // ── Inputs ────────────────────────────────────────────────────── - next_states_buf: &CudaSlice, // [B, SD] + next_states_buf: &CudaSlice, // [B, SD] // ── F32 weight pointers (raw u64 into flat target_params_buf) ── tg_w_ptrs: &[u64; 20], // ── Scratch activation buffers (not saved for backward) ──────── - tg_h_s1_scratch: &CudaSlice, // [B, SH1] scratch - tg_h_s2_buf: &CudaSlice, // [B, SH2] kept for Double-DQN - tg_h_v_scratch: &CudaSlice, // [B, VH] scratch - tg_h_b_scratch: &CudaSlice, // [B, AH] scratch (reused for all 3 branches) + tg_h_s1_scratch: &CudaSlice, // [B, SH1] scratch + tg_h_s2_buf: &CudaSlice, // [B, SH2] kept for Double-DQN + tg_h_v_scratch: &CudaSlice, // [B, VH] scratch + tg_h_b_scratch: &CudaSlice, // [B, AH] scratch (reused for all 3 branches) // ── Output buffers ────────────────────────────────────────────── - tg_v_logits_buf: &CudaSlice, // [B, NA] - tg_b_logits_buf: &CudaSlice, // [B, (B0+B1+B2)*NA] + tg_v_logits_buf: &CudaSlice, // [B, NA] + tg_b_logits_buf: &CudaSlice, // [B, (B0+B1+B2)*NA] ) -> Result<(), MLError> { let b = self.batch_size; @@ -898,7 +898,7 @@ impl CublasForward { /// Arguments: /// - `w_ptr` : raw device pointer to W[N, K] (row-major, stride K) /// - `a_ptr` : raw device pointer to A[B, K] (row-major, stride K) - /// - `c_buf` : output CudaSlice [B * N] (written in col-major [N, B]) + /// - `c_buf` : output CudaSlice [B * N] (written in col-major [N, B]) /// - `n` : output cols (out_dim) /// - `b` : batch size (rows of A / cols of output in col-major) /// - `k` : inner dim (in_dim = cols of W = cols of A) @@ -1035,7 +1035,7 @@ impl CublasForward { fn launch_add_bias_relu( &self, stream: &Arc, - output: &CudaSlice, + output: &CudaSlice, bias_ptr: u64, out_dim: usize, batch: usize, @@ -1101,7 +1101,7 @@ impl CublasForward { fn launch_add_bias( &self, stream: &Arc, - output: &CudaSlice, + output: &CudaSlice, bias_ptr: u64, out_dim: usize, batch: usize, @@ -1244,14 +1244,14 @@ impl CublasForward { // ── Raw device pointer helpers ──────────────────────────────────────────────── /// Extract raw F32 device pointer (read-only). -fn raw_f32_ptr(slice: &CudaSlice, stream: &Arc) -> u64 { +fn raw_f32_ptr(slice: &CudaSlice, stream: &Arc) -> u64 { let (ptr, guard) = slice.device_ptr(stream); let _no_drop = ManuallyDrop::new(guard); ptr } /// Extract raw F32 device pointer (mutable). -fn raw_f32_ptr_mut(slice: &CudaSlice, stream: &Arc) -> u64 { +fn raw_f32_ptr_mut(slice: &CudaSlice, stream: &Arc) -> u64 { let (ptr, guard) = slice.device_ptr(stream); let _no_drop = ManuallyDrop::new(guard); ptr @@ -1315,7 +1315,7 @@ fn compile_bias_kernels( /// /// Returns 20 raw u64 device pointers (base + byte_offset for each tensor). pub fn f32_weight_ptrs( - params_buf: &CudaSlice, + params_buf: &CudaSlice, param_sizes: &[usize; 20], stream: &Arc, ) -> [u64; 20] { diff --git a/crates/ml/src/cuda_pipeline/bias_kernels.cu b/crates/ml/src/cuda_pipeline/bias_kernels.cu index c1fbe001e..022685ae6 100644 --- a/crates/ml/src/cuda_pipeline/bias_kernels.cu +++ b/crates/ml/src/cuda_pipeline/bias_kernels.cu @@ -9,26 +9,26 @@ */ extern "C" __global__ void add_bias_relu_kernel( - float* __restrict__ output, - const float* __restrict__ bias, + __nv_bfloat16* __restrict__ output, + const __nv_bfloat16* __restrict__ bias, int out_dim, int total_elements) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= total_elements) return; - float val = output[i] + bias[i % out_dim]; - output[i] = (val > 0.0f) ? val : 0.0f; + float val = __bfloat162float(output[i]) + __bfloat162float(bias[i % out_dim]); + output[i] = __float2bfloat16((val > 0.0f) ? val : 0.0f); } extern "C" __global__ void add_bias_kernel( - float* __restrict__ output, - const float* __restrict__ bias, + __nv_bfloat16* __restrict__ output, + const __nv_bfloat16* __restrict__ bias, int out_dim, int total_elements) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= total_elements) return; - output[i] += bias[i % out_dim]; + output[i] = __float2bfloat16(__bfloat162float(output[i]) + __bfloat162float(bias[i % out_dim])); } /* ------------------------------------------------------------------ */ diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index ec92e814c..735ba73b0 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -291,7 +291,7 @@ extern "C" __global__ void c51_loss_batched( /* ── Outputs ──────────────────────────────────────────────────── */ __nv_bfloat16* __restrict__ per_sample_loss, /* [B] IS-weighted loss per sample */ __nv_bfloat16* __restrict__ td_errors, /* [B] unweighted, for PER priority update */ - float* __restrict__ total_loss, /* [1] batch mean loss (atomicAdd) — stays F32, no BF16 atomicAdd */ + __nv_bfloat16* __restrict__ total_loss, /* [1] batch mean loss (atomicAdd) — stays F32, no BF16 atomicAdd */ /* ── Saved tensors for backward pass ─────────────────────────── */ __nv_bfloat16* __restrict__ save_current_lp, /* [B, NUM_BRANCHES, num_atoms] */ @@ -600,6 +600,16 @@ extern "C" __global__ void c51_loss_batched( float weighted_loss = clamped_ce * is_weight; per_sample_loss[sample_id] = __float2bfloat16(weighted_loss); td_errors[sample_id] = __float2bfloat16(clamped_ce); /* PER sees clamped too */ - atomicAdd(total_loss, weighted_loss / (float)batch_size); + /* BF16 atomicAdd via CAS loop */ + { + unsigned short* addr = (unsigned short*)total_loss; + unsigned short old_bits = *addr, assumed; + do { + assumed = old_bits; + __nv_bfloat16 old_bf = *(__nv_bfloat16*)&assumed; + __nv_bfloat16 new_bf = __float2bfloat16(__bfloat162float(old_bf) + weighted_loss / (float)batch_size); + old_bits = atomicCAS(addr, assumed, *(unsigned short*)&new_bf); + } while (old_bits != assumed); + } } } diff --git a/crates/ml/src/cuda_pipeline/common_device_functions.cuh b/crates/ml/src/cuda_pipeline/common_device_functions.cuh index 75fd65e08..49e8f7aa1 100644 --- a/crates/ml/src/cuda_pipeline/common_device_functions.cuh +++ b/crates/ml/src/cuda_pipeline/common_device_functions.cuh @@ -13,10 +13,52 @@ * This matches NVIDIA's native tensor core accumulate mode. */ #include +/* ── BF16 native math wrappers ─────────────────────────────────────── */ +/* Thin wrappers around F32 transcendentals for BF16 arguments. */ +/* The cast is hidden inside — kernel code reads as pure BF16. */ +/* Arithmetic (+, -, *, /, >, <) uses native __nv_bfloat16 ops (SM80+)*/ + +__device__ __forceinline__ __nv_bfloat16 bf16_sqrt(__nv_bfloat16 x) { + return __float2bfloat16(sqrtf(__bfloat162float(x))); +} +__device__ __forceinline__ __nv_bfloat16 bf16_log(__nv_bfloat16 x) { + return __float2bfloat16(logf(__bfloat162float(x))); +} +__device__ __forceinline__ __nv_bfloat16 bf16_exp(__nv_bfloat16 x) { + return __float2bfloat16(expf(__bfloat162float(x))); +} +__device__ __forceinline__ __nv_bfloat16 bf16_pow(__nv_bfloat16 x, __nv_bfloat16 p) { + return __float2bfloat16(powf(__bfloat162float(x), __bfloat162float(p))); +} +__device__ __forceinline__ __nv_bfloat16 bf16_fabs(__nv_bfloat16 x) { + return __float2bfloat16(fabsf(__bfloat162float(x))); +} +__device__ __forceinline__ __nv_bfloat16 bf16_fmax(__nv_bfloat16 a, __nv_bfloat16 b) { + return __float2bfloat16(fmaxf(__bfloat162float(a), __bfloat162float(b))); +} +__device__ __forceinline__ __nv_bfloat16 bf16_fmin(__nv_bfloat16 a, __nv_bfloat16 b) { + return __float2bfloat16(fminf(__bfloat162float(a), __bfloat162float(b))); +} +__device__ __forceinline__ __nv_bfloat16 bf16_cos(__nv_bfloat16 x) { + return __float2bfloat16(cosf(__bfloat162float(x))); +} + +/** BF16 zero constant */ +__device__ __forceinline__ __nv_bfloat16 bf16_zero() { + return __float2bfloat16(0.0f); +} +/** BF16 one constant */ +__device__ __forceinline__ __nv_bfloat16 bf16_one() { + return __float2bfloat16(1.0f); +} +/** BF16 from float scalar (for host-passed params like lr, gamma) */ +__device__ __forceinline__ __nv_bfloat16 bf16(float x) { + return __float2bfloat16(x); +} + /** BF16 leaky ReLU — operates on raw bf16, returns bf16. */ __device__ __forceinline__ __nv_bfloat16 leaky_relu_bf16(__nv_bfloat16 x) { - float xf = __bfloat162float(x); - return __float2bfloat16((xf >= 0.0f) ? xf : 0.01f * xf); + return (x > bf16_zero()) ? x : bf16(0.01f) * x; } /** Convert F32 to BF16 (truncation, matching PyTorch default). */ @@ -24,6 +66,20 @@ __device__ __forceinline__ __nv_bfloat16 f32_to_bf16(float x) { return __float2bfloat16(x); } +/** BF16 atomicAdd via CAS loop (no native BF16 atomicAdd on any SM). */ +__device__ __forceinline__ void atomicAddBF16(__nv_bfloat16* addr, __nv_bfloat16 val) { + unsigned short* addr_as_us = (unsigned short*)addr; + unsigned short old = *addr_as_us; + unsigned short assumed; + do { + assumed = old; + __nv_bfloat16 sum = __float2bfloat16( + __bfloat162float(*(__nv_bfloat16*)&assumed) + __bfloat162float(val) + ); + old = atomicCAS(addr_as_us, assumed, *(unsigned short*)&sum); + } while (assumed != old); +} + /* ------------------------------------------------------------------ */ /* Constants */ /* ------------------------------------------------------------------ */ @@ -154,6 +210,27 @@ __device__ void matvec_leaky_relu( } } + +/** BF16 overload: reads weights/bias from __nv_bfloat16, input/output in float. */ +__device__ void matvec_leaky_relu( + const __nv_bfloat16* __restrict__ W, + const __nv_bfloat16* __restrict__ b, + const float* input, + float* output, + int in_dim, + int out_dim, + int activate +) { + for (int j = 0; j < out_dim; j++) { + float acc = __bfloat162float(b[j]); + const __nv_bfloat16* row = W + j * in_dim; + for (int i = 0; i < in_dim; i++) { + acc += __bfloat162float(row[i]) * input[i]; + } + output[j] = activate ? leaky_relu(acc) : acc; + } +} + /* Maximum dimension for noise scratch arrays. * Overridable via NVRTC #define injection to match actual network widths. */ #ifndef NOISY_MAX_DIM diff --git a/crates/ml/src/cuda_pipeline/cql_grad_kernel.cu b/crates/ml/src/cuda_pipeline/cql_grad_kernel.cu index 5e58c035f..0f6210a12 100644 --- a/crates/ml/src/cuda_pipeline/cql_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/cql_grad_kernel.cu @@ -15,8 +15,8 @@ */ extern "C" __global__ void cql_logit_grad_kernel( - const float* __restrict__ v_logits, // [N, num_atoms] - const float* __restrict__ adv_logits, // [N, total_actions * num_atoms] + const __nv_bfloat16* __restrict__ v_logits, // [N, num_atoms] + const __nv_bfloat16* __restrict__ adv_logits, // [N, total_actions * num_atoms] const int* __restrict__ actions, // [N] factored action indices (0-44) float* __restrict__ d_v_logits, // [N, num_atoms] output float* __restrict__ d_adv_logits, // [N, total_actions * num_atoms] output @@ -52,7 +52,7 @@ extern "C" __global__ void cql_logit_grad_kernel( branch_sizes[1] = b1_size; branch_sizes[2] = b2_size; - const float* val = v_logits + (long long)i * num_atoms; + const __nv_bfloat16* val = v_logits + (long long)i * num_atoms; // Zero value logit gradient accumulator float d_val_accum[256]; // max num_atoms (C51: 51, generous buffer) @@ -69,31 +69,31 @@ extern "C" __global__ void cql_logit_grad_kernel( // Compute mean advantage per atom (for dueling centering) float mean_adv_sum = 0.0f; for (int aa = 0; aa < bd; aa++) { - const float* adv_aa = adv_logits + (long long)i * total_actions * num_atoms + const __nv_bfloat16* adv_aa = adv_logits + (long long)i * total_actions * num_atoms + (long long)(adv_offset + aa) * num_atoms; - for (int j = 0; j < num_atoms; j++) mean_adv_sum += adv_aa[j]; + for (int j = 0; j < num_atoms; j++) mean_adv_sum += __bfloat162float(adv_aa[j]); } float mean_adv_per_atom = mean_adv_sum / (float)(bd * num_atoms); for (int a = 0; a < bd; a++) { - const float* adv = adv_logits + (long long)i * total_actions * num_atoms + const __nv_bfloat16* adv = adv_logits + (long long)i * total_actions * num_atoms + (long long)(adv_offset + a) * num_atoms; // softmax over atoms + expected value float max_logit = -1e30f; for (int j = 0; j < num_atoms; j++) { - float c = val[j] + adv[j] - mean_adv_per_atom; + float c = __bfloat162float(val[j]) + __bfloat162float(adv[j]) - mean_adv_per_atom; if (c > max_logit) max_logit = c; } float sum_exp = 0.0f; for (int j = 0; j < num_atoms; j++) { - sum_exp += expf(val[j] + adv[j] - mean_adv_per_atom - max_logit); + sum_exp += expf(__bfloat162float(val[j]) + __bfloat162float(adv[j]) - mean_adv_per_atom - max_logit); } float log_sum = logf(sum_exp + 1e-8f) + max_logit; float eq = 0.0f; for (int j = 0; j < num_atoms; j++) { - float p = expf(val[j] + adv[j] - mean_adv_per_atom - log_sum); + float p = expf(__bfloat162float(val[j]) + __bfloat162float(adv[j]) - mean_adv_per_atom - log_sum); float z = v_min + (float)j * dz; eq += p * z; } @@ -117,7 +117,7 @@ extern "C" __global__ void cql_logit_grad_kernel( // Step 3: Chain rule through softmax-expectation to get d_logits for (int a = 0; a < bd; a++) { - const float* adv = adv_logits + (long long)i * total_actions * num_atoms + const __nv_bfloat16* adv = adv_logits + (long long)i * total_actions * num_atoms + (long long)(adv_offset + a) * num_atoms; float* d_adv = d_adv_logits + (long long)i * total_actions * num_atoms + (long long)(adv_offset + a) * num_atoms; @@ -125,19 +125,19 @@ extern "C" __global__ void cql_logit_grad_kernel( // Recompute p[j] for this action float max_logit = -1e30f; for (int j = 0; j < num_atoms; j++) { - float c = val[j] + adv[j] - mean_adv_per_atom; + float c = __bfloat162float(val[j]) + __bfloat162float(adv[j]) - mean_adv_per_atom; if (c > max_logit) max_logit = c; } float sum_exp = 0.0f; for (int j = 0; j < num_atoms; j++) { - sum_exp += expf(val[j] + adv[j] - mean_adv_per_atom - max_logit); + sum_exp += expf(__bfloat162float(val[j]) + __bfloat162float(adv[j]) - mean_adv_per_atom - max_logit); } float log_sum = logf(sum_exp + 1e-8f) + max_logit; float eq = q_values_branch[a]; for (int j = 0; j < num_atoms; j++) { - float p = expf(val[j] + adv[j] - mean_adv_per_atom - log_sum); + float p = expf(__bfloat162float(val[j]) + __bfloat162float(adv[j]) - mean_adv_per_atom - log_sum); float z = v_min + (float)j * dz; // d_combined_logit[j] = d_cql_dq[a] * p * (z - Q) float d_combined = d_cql_dq[a] * p * (z - eq); diff --git a/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu b/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu index a2a2b3c3d..aa038d2df 100644 --- a/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu +++ b/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu @@ -1,3 +1,4 @@ +#include /** * GPU-resident curiosity forward model training kernel. * @@ -25,8 +26,8 @@ * negligible for a tiny auxiliary model. */ extern "C" __global__ void curiosity_shift_states( - const float* __restrict__ states, - float* __restrict__ next_states, + const __nv_bfloat16* __restrict__ states, + __nv_bfloat16* __restrict__ next_states, int N, int state_dim ) { @@ -64,32 +65,32 @@ __device__ __forceinline__ float warp_sum_cur(float val) { * to the same weight indices, so warp reduction before atomicAdd is valid. */ extern "C" __global__ void curiosity_forward_backward( - const float* __restrict__ states, /* [N, state_dim] */ + const __nv_bfloat16* __restrict__ states, /* [N, state_dim] */ const int* __restrict__ actions, /* [N] */ - const float* __restrict__ next_states, /* [N, state_dim] */ - const float* __restrict__ w1, /* [CUR_HIDDEN, CUR_INPUT] */ - const float* __restrict__ b1, /* [CUR_HIDDEN] */ - const float* __restrict__ w2, /* [CUR_OUTPUT, CUR_HIDDEN] */ - const float* __restrict__ b2, /* [CUR_OUTPUT] */ - float* __restrict__ grad_w1, /* [CUR_HIDDEN, CUR_INPUT] */ - float* __restrict__ grad_b1, /* [CUR_HIDDEN] */ - float* __restrict__ grad_w2, /* [CUR_OUTPUT, CUR_HIDDEN] */ - float* __restrict__ grad_b2, /* [CUR_OUTPUT] */ + const __nv_bfloat16* __restrict__ next_states, /* [N, state_dim] */ + const __nv_bfloat16* __restrict__ w1, /* [CUR_HIDDEN, CUR_INPUT] */ + const __nv_bfloat16* __restrict__ b1, /* [CUR_HIDDEN] */ + const __nv_bfloat16* __restrict__ w2, /* [CUR_OUTPUT, CUR_HIDDEN] */ + const __nv_bfloat16* __restrict__ b2, /* [CUR_OUTPUT] */ + __nv_bfloat16* __restrict__ grad_w1, /* [CUR_HIDDEN, CUR_INPUT] */ + __nv_bfloat16* __restrict__ grad_b1, /* [CUR_HIDDEN] */ + __nv_bfloat16* __restrict__ grad_w2, /* [CUR_OUTPUT, CUR_HIDDEN] */ + __nv_bfloat16* __restrict__ grad_b2, /* [CUR_OUTPUT] */ int N, int state_dim ) { int tid = blockIdx.x * blockDim.x + threadIdx.x; if (tid >= N) return; - const float* state = states + tid * state_dim; + const __nv_bfloat16* state_bf = states + tid * state_dim; int action_idx = actions[tid]; - const float* next_state = next_states + tid * state_dim; + const __nv_bfloat16* next_state_bf = next_states + tid * state_dim; /* ---- Forward pass ---- */ /* Build input: first MARKET_DIM state features + 3-class action one-hot */ float input[CUR_INPUT]; - for (int i = 0; i < MARKET_DIM; i++) input[i] = state[i]; + for (int i = 0; i < MARKET_DIM; i++) input[i] = __bfloat162float(state_bf[i]); input[MARKET_DIM + 0] = 0.0f; input[MARKET_DIM + 1] = 0.0f; input[MARKET_DIM + 2] = 0.0f; @@ -182,10 +183,10 @@ extern "C" __global__ void curiosity_forward_backward( * Gradients are divided by batch_size to compute mean gradient. */ extern "C" __global__ void curiosity_adam_step( - float* __restrict__ params, /* weights/biases to update */ - const float* __restrict__ grads, /* accumulated gradients */ - float* __restrict__ m, /* first moment */ - float* __restrict__ v, /* second moment */ + __nv_bfloat16* __restrict__ params, /* weights/biases to update */ + const __nv_bfloat16* __restrict__ grads, /* accumulated gradients */ + __nv_bfloat16* __restrict__ m, /* first moment */ + __nv_bfloat16* __restrict__ v, /* second moment */ int num_params, int batch_size, /* divide grads by this for mean */ float lr, @@ -215,10 +216,10 @@ extern "C" __global__ void curiosity_adam_step( * Eliminates 3 kernel launch overheads (~150 µs/step on H100). */ extern "C" __global__ void curiosity_adam_step_fused( - float* __restrict__ p0, const float* __restrict__ g0, float* __restrict__ m0, float* __restrict__ v0, int n0, - float* __restrict__ p1, const float* __restrict__ g1, float* __restrict__ m1, float* __restrict__ v1, int n1, - float* __restrict__ p2, const float* __restrict__ g2, float* __restrict__ m2, float* __restrict__ v2, int n2, - float* __restrict__ p3, const float* __restrict__ g3, float* __restrict__ m3, float* __restrict__ v3, int n3, + __nv_bfloat16* __restrict__ p0, const __nv_bfloat16* __restrict__ g0, __nv_bfloat16* __restrict__ m0, __nv_bfloat16* __restrict__ v0, int n0, + __nv_bfloat16* __restrict__ p1, const __nv_bfloat16* __restrict__ g1, __nv_bfloat16* __restrict__ m1, __nv_bfloat16* __restrict__ v1, int n1, + __nv_bfloat16* __restrict__ p2, const __nv_bfloat16* __restrict__ g2, __nv_bfloat16* __restrict__ m2, __nv_bfloat16* __restrict__ v2, int n2, + __nv_bfloat16* __restrict__ p3, const __nv_bfloat16* __restrict__ g3, __nv_bfloat16* __restrict__ m3, __nv_bfloat16* __restrict__ v3, int n3, int batch_size, float lr, float beta1, float beta2, float eps, int step ) { @@ -227,7 +228,7 @@ extern "C" __global__ void curiosity_adam_step_fused( if (idx >= total) return; /* Determine which param group and local offset */ - float* p; const float* g; float* mm; float* vv; +__nv_bfloat16* p; const __nv_bfloat16* g; __nv_bfloat16* mm; __nv_bfloat16* vv; int local_i; if (idx < n0) { p = p0; g = g0; mm = m0; vv = v0; local_i = idx; @@ -269,29 +270,29 @@ extern "C" __global__ void curiosity_adam_step_fused( * per training step (~30-50 µs on H100 at high training frequency). */ extern "C" __global__ void curiosity_fused_zero_fwd_bwd_adam( - const float* __restrict__ states, /* [N, state_dim] */ + const __nv_bfloat16* __restrict__ states, /* [N, state_dim] */ const int* __restrict__ actions, /* [N] */ - const float* __restrict__ next_states, /* [N, state_dim] */ + const __nv_bfloat16* __restrict__ next_states, /* [N, state_dim] */ /* Weights (updated in-place by Adam in phase 3) */ - float* __restrict__ w1, /* [CUR_HIDDEN, CUR_INPUT] */ - float* __restrict__ b1, /* [CUR_HIDDEN] */ - float* __restrict__ w2, /* [CUR_OUTPUT, CUR_HIDDEN] */ - float* __restrict__ b2, /* [CUR_OUTPUT] */ + __nv_bfloat16* __restrict__ w1, /* [CUR_HIDDEN, CUR_INPUT] */ + __nv_bfloat16* __restrict__ b1, /* [CUR_HIDDEN] */ + __nv_bfloat16* __restrict__ w2, /* [CUR_OUTPUT, CUR_HIDDEN] */ + __nv_bfloat16* __restrict__ b2, /* [CUR_OUTPUT] */ /* Gradient accumulators (zeroed in phase 1, accumulated in phase 2) */ - float* __restrict__ grad_w1, /* [CUR_HIDDEN, CUR_INPUT] */ - float* __restrict__ grad_b1, /* [CUR_HIDDEN] */ - float* __restrict__ grad_w2, /* [CUR_OUTPUT, CUR_HIDDEN] */ - float* __restrict__ grad_b2, /* [CUR_OUTPUT] */ + __nv_bfloat16* __restrict__ grad_w1, /* [CUR_HIDDEN, CUR_INPUT] */ + __nv_bfloat16* __restrict__ grad_b1, /* [CUR_HIDDEN] */ + __nv_bfloat16* __restrict__ grad_w2, /* [CUR_OUTPUT, CUR_HIDDEN] */ + __nv_bfloat16* __restrict__ grad_b2, /* [CUR_OUTPUT] */ /* Adam first moment */ - float* __restrict__ adam_m_w1, /* [CUR_HIDDEN * CUR_INPUT] */ - float* __restrict__ adam_m_b1, /* [CUR_HIDDEN] */ - float* __restrict__ adam_m_w2, /* [CUR_OUTPUT * CUR_HIDDEN] */ - float* __restrict__ adam_m_b2, /* [CUR_OUTPUT] */ + __nv_bfloat16* __restrict__ adam_m_w1, /* [CUR_HIDDEN * CUR_INPUT] */ + __nv_bfloat16* __restrict__ adam_m_b1, /* [CUR_HIDDEN] */ + __nv_bfloat16* __restrict__ adam_m_w2, /* [CUR_OUTPUT * CUR_HIDDEN] */ + __nv_bfloat16* __restrict__ adam_m_b2, /* [CUR_OUTPUT] */ /* Adam second moment */ - float* __restrict__ adam_v_w1, /* [CUR_HIDDEN * CUR_INPUT] */ - float* __restrict__ adam_v_b1, /* [CUR_HIDDEN] */ - float* __restrict__ adam_v_w2, /* [CUR_OUTPUT * CUR_HIDDEN] */ - float* __restrict__ adam_v_b2, /* [CUR_OUTPUT] */ + __nv_bfloat16* __restrict__ adam_v_w1, /* [CUR_HIDDEN * CUR_INPUT] */ + __nv_bfloat16* __restrict__ adam_v_b1, /* [CUR_HIDDEN] */ + __nv_bfloat16* __restrict__ adam_v_w2, /* [CUR_OUTPUT * CUR_HIDDEN] */ + __nv_bfloat16* __restrict__ adam_v_b2, /* [CUR_OUTPUT] */ /* Atomic block-arrival counter (must be zeroed before launch) */ int* __restrict__ block_counter, /* Scalar parameters */ @@ -320,15 +321,15 @@ extern "C" __global__ void curiosity_fused_zero_fwd_bwd_adam( /* PHASE 1: Forward + backward pass (one sample per thread) */ /* ================================================================ */ if (tid < N) { - const float* state = states + tid * state_dim; + const __nv_bfloat16* state_bf2 = states + tid * state_dim; int action_idx = actions[tid]; - const float* next_state = next_states + tid * state_dim; + const __nv_bfloat16* next_state_bf2 = next_states + tid * state_dim; /* ---- Forward pass ---- */ /* Build input: first MARKET_DIM state features + 3-class action one-hot */ float input[CUR_INPUT]; - for (int i = 0; i < MARKET_DIM; i++) input[i] = state[i]; + for (int i = 0; i < MARKET_DIM; i++) input[i] = __bfloat162float(state_bf[i]); input[MARKET_DIM + 0] = 0.0f; input[MARKET_DIM + 1] = 0.0f; input[MARKET_DIM + 2] = 0.0f; @@ -437,7 +438,7 @@ extern "C" __global__ void curiosity_fused_zero_fwd_bwd_adam( for (int i = block_tid; i < total_grad_elems; i += block_size) { /* Determine which param group and local offset */ - float* p; float* g; float* mm; float* vv; +__nv_bfloat16* p; __nv_bfloat16* g; __nv_bfloat16* mm; __nv_bfloat16* vv; int local_i; if (i < w1_len) { p = w1; g = grad_w1; mm = adam_m_w1; vv = adam_v_w1; local_i = i; diff --git a/crates/ml/src/cuda_pipeline/decision_transformer.rs b/crates/ml/src/cuda_pipeline/decision_transformer.rs index 63a9f7748..55792d3b0 100644 --- a/crates/ml/src/cuda_pipeline/decision_transformer.rs +++ b/crates/ml/src/cuda_pipeline/decision_transformer.rs @@ -218,13 +218,13 @@ struct LayerOffsets { // Extract the u64 device pointer and prevent the SyncOnDrop guard from // recording an event. Safe because all operations happen on the same stream. -fn raw_ptr_f32(slice: &CudaSlice, stream: &CudaStream) -> u64 { +fn raw_ptr_f32(slice: &CudaSlice, stream: &CudaStream) -> u64 { let (ptr, guard) = slice.device_ptr(stream); let _no_drop = std::mem::ManuallyDrop::new(guard); ptr } -fn raw_ptr_f32_mut(slice: &mut CudaSlice, stream: &CudaStream) -> u64 { +fn raw_ptr_f32_mut(slice: &mut CudaSlice, stream: &CudaStream) -> u64 { let (ptr, guard) = slice.device_ptr_mut(stream); let _no_drop = std::mem::ManuallyDrop::new(guard); ptr @@ -299,39 +299,39 @@ fn compile_dt_kernels( struct DtScratch { /// Embedded tokens: [B, T, E] - embedded: CudaSlice, + embedded: CudaSlice, /// Q, K, V projections: [B, T, E] each - q_buf: CudaSlice, - k_buf: CudaSlice, - v_buf: CudaSlice, + q_buf: CudaSlice, + k_buf: CudaSlice, + v_buf: CudaSlice, /// Attention output (pre-LN): [B, T, E] - attn_out: CudaSlice, + attn_out: CudaSlice, /// Post-LN1 output: [B, T, E] - ln1_out: CudaSlice, + ln1_out: CudaSlice, /// FFN output (post-LN2): [B, T, E] - ffn_out: CudaSlice, + ffn_out: CudaSlice, /// LN2 temp output: [B, T, E] - ln2_out: CudaSlice, + ln2_out: CudaSlice, /// Action logits: [B*T, A] - logits: CudaSlice, + logits: CudaSlice, /// Per-sample loss: [B*T] - per_sample_loss: CudaSlice, + per_sample_loss: CudaSlice, /// Scalar total loss: [1] - total_loss: CudaSlice, + total_loss: CudaSlice, /// Gradient of logits: [B*T, A] - d_logits: CudaSlice, + d_logits: CudaSlice, /// Gradient buffer for layer backward: [B, T, E] - d_layer: CudaSlice, + d_layer: CudaSlice, /// Second gradient buffer for embed backward input grad: [B, T, input_dim] - d_embed_input: CudaSlice, + d_embed_input: CudaSlice, /// Gradient accumulator for all parameters: [total_params] - d_params: CudaSlice, + d_params: CudaSlice, /// Adam first moment: [total_params] - adam_m: CudaSlice, + adam_m: CudaSlice, /// Adam second moment: [total_params] - adam_v: CudaSlice, + adam_v: CudaSlice, /// Grad norm scratch: [1] - grad_norm: CudaSlice, + grad_norm: CudaSlice, /// Adam step counter on device: [1] adam_t: CudaSlice, } @@ -346,8 +346,8 @@ fn alloc_dt_scratch( let bt_input = config.batch_size * config.context_len * config.input_dim(); let total_params = config.total_params(); - let alloc_f = |n: usize, name: &str| -> Result, MLError> { - stream.alloc_zeros::(n) + let alloc_f = |n: usize, name: &str| -> Result, MLError> { + stream.alloc_zeros::(n) .map_err(|e| MLError::ModelError(format!("DT scratch {name} ({n} f32): {e}"))) }; @@ -408,11 +408,11 @@ pub struct DecisionTransformer { config: DecisionTransformerConfig, stream: Arc, /// Flat parameter buffer (all weights + biases). - params: CudaSlice, + params: CudaSlice, /// Context buffer: [batch_size, context_len, state_dim + 2]. - context_buf: CudaSlice, + context_buf: CudaSlice, /// Output buffer: [batch_size, num_actions] action logits. - action_logits: CudaSlice, + action_logits: CudaSlice, /// Compiled DT kernels (lazy init on first pretrain_step). kernels: Option, /// Adam optimizer kernels (grad_norm + adam_update). @@ -464,14 +464,14 @@ impl DecisionTransformer { } } - let mut params = stream.alloc_zeros::(total_params) + let mut params = stream.alloc_zeros::(total_params) .map_err(|e| MLError::ModelError(format!("DT params alloc: {e}")))?; stream.memcpy_htod(&host_params, &mut params) .map_err(|e| MLError::ModelError(format!("DT params upload: {e}")))?; - let context_buf = stream.alloc_zeros::(context_size) + let context_buf = stream.alloc_zeros::(context_size) .map_err(|e| MLError::ModelError(format!("DT context alloc: {e}")))?; - let action_logits = stream.alloc_zeros::(output_size) + let action_logits = stream.alloc_zeros::(output_size) .map_err(|e| MLError::ModelError(format!("DT action_logits alloc: {e}")))?; let vram_mb = config.vram_estimate() as f64 / (1024.0 * 1024.0); @@ -524,7 +524,7 @@ impl DecisionTransformer { /// Mean cross-entropy loss for this batch. pub fn pretrain_step( &mut self, - trajectories: &CudaSlice, + trajectories: &CudaSlice, target_actions: &CudaSlice, batch_size: usize, traj_elem_offset: usize, @@ -911,7 +911,7 @@ impl DecisionTransformer { } /// Get the action logits buffer. - pub fn action_logits(&self) -> &CudaSlice { + pub fn action_logits(&self) -> &CudaSlice { &self.action_logits } @@ -935,11 +935,11 @@ impl DecisionTransformer { /// - `num_complete_batches` is how many full batches of `batch_size` episodes are available pub fn build_dt_trajectories( &mut self, - features_gpu: &CudaSlice, - targets_gpu: &CudaSlice, + features_gpu: &CudaSlice, + targets_gpu: &CudaSlice, num_bars: usize, gamma: f64, - ) -> Result<(CudaSlice, CudaSlice, usize), MLError> { + ) -> Result<(CudaSlice, CudaSlice, usize), MLError> { self.ensure_compiled()?; let stream = &self.stream; @@ -978,7 +978,7 @@ impl DecisionTransformer { ); // ── Step 1: Compute rewards and expert actions on GPU ── - let mut rewards_gpu = stream.alloc_zeros::(num_bars) + let mut rewards_gpu = stream.alloc_zeros::(num_bars) .map_err(|e| MLError::ModelError(format!("DT rewards alloc: {e}")))?; let mut actions_all_gpu = stream.alloc_zeros::(num_bars) .map_err(|e| MLError::ModelError(format!("DT actions alloc: {e}")))?; @@ -1025,9 +1025,9 @@ impl DecisionTransformer { // Instead of copying, we'll compute RTG per-episode using the global rewards array. // Build a temporary [num_episodes, context_len] reward slice and RTG. let ep_total = num_episodes * context_len; - let mut ep_rewards_gpu = stream.alloc_zeros::(ep_total) + let mut ep_rewards_gpu = stream.alloc_zeros::(ep_total) .map_err(|e| MLError::ModelError(format!("DT ep_rewards alloc: {e}")))?; - let mut rtg_gpu = stream.alloc_zeros::(ep_total) + let mut rtg_gpu = stream.alloc_zeros::(ep_total) .map_err(|e| MLError::ModelError(format!("DT rtg alloc: {e}")))?; // Copy per-episode reward windows from the global reward array. @@ -1091,7 +1091,7 @@ impl DecisionTransformer { // ── Step 4: Build trajectory tensor on GPU ── let traj_total = num_episodes * context_len * input_dim; - let mut trajectories_gpu = stream.alloc_zeros::(traj_total) + let mut trajectories_gpu = stream.alloc_zeros::(traj_total) .map_err(|e| MLError::ModelError(format!("DT trajectories alloc: {e}")))?; let mut target_actions_gpu = stream.alloc_zeros::(ep_total) .map_err(|e| MLError::ModelError(format!("DT target_actions alloc: {e}")))?; diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index d6cd113a9..93d0213c4 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -25,12 +25,12 @@ * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dqn_grad_norm_kernel( - const float* __restrict__ grads, /* [TOTAL_PARAMS] accumulated gradients */ - float* __restrict__ out_grad_norm, /* [1] output: sum of squares */ + const __nv_bfloat16* __restrict__ grads, /* [TOTAL_PARAMS] accumulated gradients */ + __nv_bfloat16* __restrict__ out_grad_norm, /* [1] output: sum of squares */ int total_params ) { int idx = blockIdx.x * blockDim.x + threadIdx.x; - float g = (idx < total_params) ? grads[idx] : 0.0f; + float g = (idx < total_params) ? __bfloat162float(grads[idx]) : 0.0f; float g2 = g * g; /* Warp-level reduction via shuffle (no shared memory) */ @@ -56,7 +56,7 @@ extern "C" __global__ void dqn_grad_norm_kernel( for (int offset = 16; offset > 0; offset >>= 1) val += __shfl_xor_sync(0xFFFFFFFF, val, offset); if (warp_lane == 0) - atomicAdd(out_grad_norm, val); + atomicAddBF16(out_grad_norm, __float2bfloat16(val)); } } @@ -71,11 +71,11 @@ extern "C" __global__ void dqn_grad_norm_kernel( * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dqn_adam_update_kernel( - float* __restrict__ params, /* [TOTAL_PARAMS] flattened parameters */ - const float* __restrict__ grads, /* [TOTAL_PARAMS] accumulated gradients */ - float* __restrict__ m, /* [TOTAL_PARAMS] first moment (Adam) */ - float* __restrict__ v, /* [TOTAL_PARAMS] second moment (Adam) */ - const float* __restrict__ grad_norm_sq, /* [1] completed sum of squares */ + __nv_bfloat16* __restrict__ params, /* [TOTAL_PARAMS] flattened parameters */ + const __nv_bfloat16* __restrict__ grads, /* [TOTAL_PARAMS] accumulated gradients */ + __nv_bfloat16* __restrict__ m, /* [TOTAL_PARAMS] first moment (Adam) */ + __nv_bfloat16* __restrict__ v, /* [TOTAL_PARAMS] second moment (Adam) */ + const __nv_bfloat16* __restrict__ grad_norm_sq, /* [1] completed sum of squares */ float lr, float beta1, float beta2, @@ -89,10 +89,10 @@ extern "C" __global__ void dqn_adam_update_kernel( if (idx >= total_params) return; int t = *t_ptr; - float g = grads[idx]; + float g = __bfloat162float(grads[idx]); /* Clip using the COMPLETED norm (no race — separate kernel launch) */ - float norm = sqrtf(*grad_norm_sq + 1e-12f); + float norm = sqrtf(__bfloat162float(*grad_norm_sq) + 1e-12f); float clip_scale = (norm > max_grad_norm) ? (max_grad_norm / norm) : 1.0f; float clipped_g = g * clip_scale; @@ -100,16 +100,18 @@ extern "C" __global__ void dqn_adam_update_kernel( float beta1_t = 1.0f - powf(beta1, (float)t); float beta2_t = 1.0f - powf(beta2, (float)t); - float m_i = beta1 * m[idx] + (1.0f - beta1) * clipped_g; - float v_i = beta2 * v[idx] + (1.0f - beta2) * clipped_g * clipped_g; - m[idx] = m_i; - v[idx] = v_i; + float m_i = beta1 * __bfloat162float(m[idx]) + (1.0f - beta1) * clipped_g; + float v_i = beta2 * __bfloat162float(v[idx]) + (1.0f - beta2) * clipped_g * clipped_g; + m[idx] = __float2bfloat16(m_i); + v[idx] = __float2bfloat16(v_i); float m_hat = m_i / beta1_t; float v_hat = v_i / beta2_t; /* AdamW weight decay (decoupled) */ - params[idx] -= lr * (m_hat / (sqrtf(v_hat) + epsilon) + weight_decay * params[idx]); + float p_val = __bfloat162float(params[idx]); + p_val -= lr * (m_hat / (sqrtf(v_hat) + epsilon) + weight_decay * p_val); + params[idx] = __float2bfloat16(p_val); } /* ══════════════════════════════════════════════════════════════════════ @@ -124,13 +126,13 @@ extern "C" __global__ void dqn_adam_update_kernel( * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dqn_saxpy_kernel( - float* __restrict__ y, - const float* __restrict__ x, + __nv_bfloat16* __restrict__ y, + const __nv_bfloat16* __restrict__ x, float alpha, int n ) { int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) y[i] += alpha * x[i]; + if (i < n) y[i] = __float2bfloat16(__bfloat162float(y[i]) + alpha * __bfloat162float(x[i])); } /* ══════════════════════════════════════════════════════════════════════ @@ -150,19 +152,19 @@ extern "C" __global__ void dqn_saxpy_kernel( * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dqn_clipped_saxpy_kernel( - float* __restrict__ y, - const float* __restrict__ x, + __nv_bfloat16* __restrict__ y, + const __nv_bfloat16* __restrict__ x, float alpha, float max_norm, - const float* __restrict__ x_norm_sq, + const __nv_bfloat16* __restrict__ x_norm_sq, int n ) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= n) return; - float norm = sqrtf(*x_norm_sq + 1e-12f); + float norm = sqrtf(__bfloat162float(*x_norm_sq) + 1e-12f); float clip = (norm > max_norm) ? (max_norm / norm) : 1.0f; - y[i] += alpha * clip * x[i]; + y[i] = __float2bfloat16(__bfloat162float(y[i]) + alpha * clip * __bfloat162float(x[i])); } /* ══════════════════════════════════════════════════════════════════════ @@ -178,15 +180,15 @@ extern "C" __global__ void dqn_clipped_saxpy_kernel( * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dqn_clip_grad_kernel( - float* __restrict__ grads, - const float* __restrict__ grad_norm_sq, + __nv_bfloat16* __restrict__ grads, + const __nv_bfloat16* __restrict__ grad_norm_sq, float max_norm, int total_params ) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= total_params) return; - float norm = sqrtf(*grad_norm_sq + 1e-12f); + float norm = sqrtf(__bfloat162float(*grad_norm_sq) + 1e-12f); if (norm > max_norm) { grads[idx] *= max_norm / norm; } @@ -204,11 +206,11 @@ extern "C" __global__ void dqn_clip_grad_kernel( * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dqn_zero_kernel( - float* __restrict__ buf, + __nv_bfloat16* __restrict__ buf, int n ) { int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) buf[i] = 0.0f; + if (i < n) buf[i] = __float2bfloat16(0.0f); } /* ══════════════════════════════════════════════════════════════════════ @@ -223,8 +225,8 @@ extern "C" __global__ void dqn_zero_kernel( * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dqn_regime_scale_kernel( - float* __restrict__ td_errors, /* [B] scaled in-place */ - const float* __restrict__ states, /* [B, state_dim] */ + __nv_bfloat16* __restrict__ td_errors, /* [B] scaled in-place */ + const __nv_bfloat16* __restrict__ states, /* [B, state_dim] */ int batch_size, int state_dim /* runtime state dimension */ ) { @@ -233,11 +235,11 @@ extern "C" __global__ void dqn_regime_scale_kernel( /* Read target regime from first sample (GPU-only, no CPU readback). * ADX at feature index 40, CUSUM at index 41. */ - float target_adx = states[0 * state_dim + 40]; - float target_cusum = states[0 * state_dim + 41]; + float target_adx = __bfloat162float(states[0 * state_dim + 40]); + float target_cusum = __bfloat162float(states[0 * state_dim + 41]); - float sample_adx = states[i * state_dim + 40]; - float sample_cusum = states[i * state_dim + 41]; + float sample_adx = __bfloat162float(states[i * state_dim + 40]); + float sample_cusum = __bfloat162float(states[i * state_dim + 41]); float adx_diff = sample_adx - target_adx; float cusum_diff = sample_cusum - target_cusum; @@ -248,7 +250,7 @@ extern "C" __global__ void dqn_regime_scale_kernel( /* Clamp to [0.5, 2.0]: never fully suppress, max 2x upweight */ float scale = 0.5f + 1.5f * sim; - td_errors[i] *= scale; + td_errors[i] = __float2bfloat16(__bfloat162float(td_errors[i]) * scale); } /* ══════════════════════════════════════════════════════════════════════ @@ -264,7 +266,7 @@ extern "C" __global__ void dqn_regime_scale_kernel( * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dqn_shrink_perturb_kernel( - float* __restrict__ params, + __nv_bfloat16* __restrict__ params, float alpha, float sigma, int n, @@ -285,7 +287,7 @@ extern "C" __global__ void dqn_shrink_perturb_kernel( float noise = sigma * sqrtf(-2.0f * logf(u1)) * cosf(6.283185307f * u2); /* Shrink-and-Perturb: blend old weights with noise */ - params[i] = alpha * params[i] + (1.0f - alpha) * noise; + params[i] = __float2bfloat16(alpha * __bfloat162float(params[i]) + (1.0f - alpha) * noise); } /* ══════════════════════════════════════════════════════════════════════ @@ -297,13 +299,13 @@ extern "C" __global__ void dqn_shrink_perturb_kernel( * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dqn_relu_mask_kernel( - float* __restrict__ dx, - const float* __restrict__ activation, + __nv_bfloat16* __restrict__ dx, + const __nv_bfloat16* __restrict__ activation, int n ) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= n) return; - if (activation[i] <= 0.0f) dx[i] = 0.0f; + if (__bfloat162float(activation[i]) <= 0.0f) dx[i] = __float2bfloat16(0.0f); } /* ══════════════════════════════════════════════════════════════════════ @@ -323,9 +325,9 @@ extern "C" __global__ void dqn_relu_mask_kernel( * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void spectral_norm_kernel( - float* __restrict__ W, /* [out_dim, in_dim] weight matrix — scaled in-place */ - float* __restrict__ u, /* [out_dim] left singular vector (persistent) */ - float* __restrict__ v, /* [in_dim] right singular vector (persistent) */ + __nv_bfloat16* __restrict__ W, /* [out_dim, in_dim] weight matrix — scaled in-place */ + __nv_bfloat16* __restrict__ u, /* [out_dim] left singular vector (persistent) */ + __nv_bfloat16* __restrict__ v, /* [in_dim] right singular vector (persistent) */ int out_dim, int in_dim, float sigma_max /* clip sigma above this (typically 1.0) */ @@ -340,15 +342,15 @@ extern "C" __global__ void spectral_norm_kernel( for (int col = tid; col < in_dim; col += bd) { float val = 0.0f; for (int row = 0; row < out_dim; row++) - val += W[row * in_dim + col] * u[row]; - v[col] = val; /* unnormalized — normalize below */ + val += __bfloat162float(W[row * in_dim + col]) * __bfloat162float(u[row]); + v[col] = __float2bfloat16(val); /* unnormalized — normalize below */ } __syncthreads(); /* ── Normalize v_new: ||v||₂ via parallel reduction ── */ float local_v2 = 0.0f; for (int col = tid; col < in_dim; col += bd) - local_v2 += v[col] * v[col]; + { float vc = __bfloat162float(v[col]); local_v2 += vc * vc; } shmem[tid] = local_v2; __syncthreads(); for (int s = bd / 2; s > 0; s >>= 1) { @@ -357,22 +359,22 @@ extern "C" __global__ void spectral_norm_kernel( } float v_norm = sqrtf(shmem[0] + 1e-12f); for (int col = tid; col < in_dim; col += bd) - v[col] /= v_norm; + v[col] = __float2bfloat16(__bfloat162float(v[col]) / v_norm); __syncthreads(); /* ── Step 2: u_new = W v_new ── (strided for dims > blockDim) */ for (int row = tid; row < out_dim; row += bd) { float val = 0.0f; for (int col = 0; col < in_dim; col++) - val += W[row * in_dim + col] * v[col]; - u[row] = val; /* unnormalized — sigma = ||u_new|| */ + val += __bfloat162float(W[row * in_dim + col]) * __bfloat162float(v[col]); + u[row] = __float2bfloat16(val); /* unnormalized — sigma = ||u_new|| */ } __syncthreads(); /* ── Sigma = ||W v_new|| = ||u_new|| (before normalizing) ── */ float local_u2 = 0.0f; for (int row = tid; row < out_dim; row += bd) - local_u2 += u[row] * u[row]; + { float ur = __bfloat162float(u[row]); local_u2 += ur * ur; } shmem[tid] = local_u2; __syncthreads(); for (int s = bd / 2; s > 0; s >>= 1) { @@ -383,7 +385,7 @@ extern "C" __global__ void spectral_norm_kernel( /* Normalize u_new */ for (int row = tid; row < out_dim; row += bd) - u[row] /= (sigma + 1e-12f); + u[row] = __float2bfloat16(__bfloat162float(u[row]) / (sigma + 1e-12f)); __syncthreads(); /* ── Step 3: Scale W if sigma > sigma_max ── */ @@ -391,6 +393,6 @@ extern "C" __global__ void spectral_norm_kernel( float scale = (sigma > sigma_max) ? (sigma_max / sigma) : 1.0f; if (scale < 1.0f) { for (int i = tid; i < n_total; i += bd) - W[i] *= scale; + W[i] = __float2bfloat16(__bfloat162float(W[i]) * scale); } } diff --git a/crates/ml/src/cuda_pipeline/dt_kernels.cu b/crates/ml/src/cuda_pipeline/dt_kernels.cu index c12586dac..20792a02b 100644 --- a/crates/ml/src/cuda_pipeline/dt_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dt_kernels.cu @@ -35,16 +35,16 @@ * Linear projection of each timestep's token (return-to-go, state, action) * plus learned positional encoding. * - * out[b][t][d] = sum_j(traj[b][t][j] * W_embed[j][d]) + b_embed[d] + pos_embed[t][d] + * out[b][t][d] = sum_j(__bfloat162float(traj[b])[t][j] * W_embed[j][d]) + b_embed[d] + pos_embed[t][d] * * Grid: (B*T, 1, 1) * Block: (min(E, 256), 1, 1) * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dt_embed_kernel( - const float* __restrict__ trajectories, /* [B, T, input_dim] */ - const float* __restrict__ W_embed, /* [input_dim, embed_dim] */ - const float* __restrict__ b_embed, /* [embed_dim] */ - const float* __restrict__ pos_embed, /* [T, embed_dim] */ + const __nv_bfloat16* __restrict__ trajectories, /* [B, T, input_dim] */ + const __nv_bfloat16* __restrict__ W_embed, /* [input_dim, embed_dim] */ + const __nv_bfloat16* __restrict__ b_embed, /* [embed_dim] */ + const __nv_bfloat16* __restrict__ pos_embed, /* [T, embed_dim] */ float* __restrict__ output, /* [B, T, embed_dim] */ int B, int T, int input_dim, int embed_dim ) { @@ -55,13 +55,13 @@ extern "C" __global__ void dt_embed_kernel( int b = bt / T; int t = bt % T; - const float* traj = trajectories + (b * T + t) * input_dim; + const __nv_bfloat16* traj = trajectories + (b * T + t) * input_dim; float* out = output + (b * T + t) * embed_dim; for (int dd = d; dd < embed_dim; dd += blockDim.x) { float val = b_embed[dd] + pos_embed[t * embed_dim + dd]; for (int j = 0; j < input_dim; j++) { - val += traj[j] * W_embed[j * embed_dim + dd]; + val += __bfloat162float(traj[j]) * W_embed[j * embed_dim + dd]; } out[dd] = val; } @@ -77,10 +77,10 @@ extern "C" __global__ void dt_embed_kernel( * Block: (min(E, 256), 1, 1) * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dt_qkv_projection_kernel( - const float* __restrict__ input, /* [B, T, E] */ - const float* __restrict__ W_Q, /* [E, E] */ - const float* __restrict__ W_K, /* [E, E] */ - const float* __restrict__ W_V, /* [E, E] */ + const __nv_bfloat16* __restrict__ input, /* [B, T, E] */ + const __nv_bfloat16* __restrict__ W_Q, /* [E, E] */ + const __nv_bfloat16* __restrict__ W_K, /* [E, E] */ + const __nv_bfloat16* __restrict__ W_V, /* [E, E] */ float* __restrict__ Q_out, /* [B, T, E] */ float* __restrict__ K_out, /* [B, T, E] */ float* __restrict__ V_out, /* [B, T, E] */ @@ -90,12 +90,12 @@ extern "C" __global__ void dt_qkv_projection_kernel( if (bt >= B * T) return; int d = threadIdx.x; - const float* x = input + bt * E; + const __nv_bfloat16* x = input + bt * E; for (int dd = d; dd < E; dd += blockDim.x) { float q = 0.0f, k = 0.0f, v = 0.0f; for (int j = 0; j < E; j++) { - float xj = x[j]; + float xj = __bfloat162float(x[j]); q += xj * W_Q[j * E + dd]; k += xj * W_K[j * E + dd]; v += xj * W_V[j * E + dd]; @@ -123,11 +123,11 @@ extern "C" __global__ void dt_qkv_projection_kernel( * final[b][t][d] = sum_h( out_h[t][dh] * W_O[h*Dh + dh, d] ) * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dt_causal_attention_kernel( - const float* __restrict__ Q, /* [B, T, E] */ - const float* __restrict__ K, /* [B, T, E] */ - const float* __restrict__ V, /* [B, T, E] */ - const float* __restrict__ W_O, /* [E, E] output projection */ - const float* __restrict__ b_O, /* [E] output bias */ + const __nv_bfloat16* __restrict__ Q, /* [B, T, E] */ + const __nv_bfloat16* __restrict__ K, /* [B, T, E] */ + const __nv_bfloat16* __restrict__ V, /* [B, T, E] */ + const __nv_bfloat16* __restrict__ W_O, /* [E, E] output projection */ + const __nv_bfloat16* __restrict__ b_O, /* [E] output bias */ float* __restrict__ output, /* [B, T, E] */ int B, int T, int E, int num_heads ) { @@ -141,9 +141,9 @@ extern "C" __global__ void dt_causal_attention_kernel( int head_off = h * Dh; /* Pointers for this batch element */ - const float* Q_b = Q + b * T * E; - const float* K_b = K + b * T * E; - const float* V_b = V + b * T * E; + const __nv_bfloat16* Q_b = Q + b * T * E; + const __nv_bfloat16* K_b = K + b * T * E; + const __nv_bfloat16* V_b = V + b * T * E; /* Shared memory for K and V of this head: [T, Dh] each */ extern __shared__ float shmem[]; @@ -152,8 +152,8 @@ extern "C" __global__ void dt_causal_attention_kernel( /* Load K and V for this head into shared memory */ for (int dd = 0; dd < Dh; dd++) { - sh_K[t_idx * Dh + dd] = K_b[t_idx * E + head_off + dd]; - sh_V[t_idx * Dh + dd] = V_b[t_idx * E + head_off + dd]; + sh_K[t_idx * Dh + dd] = __bfloat162float(K_b[t_idx * E + head_off + dd]); + sh_V[t_idx * Dh + dd] = __bfloat162float(V_b[t_idx * E + head_off + dd]); } __syncthreads(); @@ -165,7 +165,7 @@ extern "C" __global__ void dt_causal_attention_kernel( for (int j = 0; j <= t_idx; j++) { float score = 0.0f; for (int dd = 0; dd < Dh; dd++) { - score += Q_b[t_idx * E + head_off + dd] * sh_K[j * Dh + dd]; + score += __bfloat162float(Q_b[t_idx * E + head_off + dd]) * sh_K[j * Dh + dd]; } score *= inv_sqrt_dh; if (score > max_score) max_score = score; @@ -179,7 +179,7 @@ extern "C" __global__ void dt_causal_attention_kernel( for (int j = 0; j <= t_idx; j++) { float score = 0.0f; for (int dd = 0; dd < Dh; dd++) { - score += Q_b[t_idx * E + head_off + dd] * sh_K[j * Dh + dd]; + score += __bfloat162float(Q_b[t_idx * E + head_off + dd]) * sh_K[j * Dh + dd]; } score = expf(score * inv_sqrt_dh - max_score); sum_exp += score; @@ -239,9 +239,9 @@ extern "C" __global__ void dt_causal_attention_kernel( * Block: (min(E, 256), 1, 1) * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dt_layernorm_kernel( - const float* __restrict__ input, /* [B*T, E] */ - const float* __restrict__ gamma, /* [E] */ - const float* __restrict__ beta, /* [E] */ + const __nv_bfloat16* __restrict__ input, /* [B*T, E] */ + const __nv_bfloat16* __restrict__ gamma, /* [E] */ + const __nv_bfloat16* __restrict__ beta, /* [E] */ float* __restrict__ output, /* [B*T, E] */ int N, /* B*T */ int E @@ -250,13 +250,13 @@ extern "C" __global__ void dt_layernorm_kernel( if (n >= N) return; int tid = threadIdx.x; - const float* x = input + n * E; + const __nv_bfloat16* x = input + n * E; float* out = output + n * E; /* Compute mean via block-parallel reduction */ float local_sum = 0.0f; for (int d = tid; d < E; d += blockDim.x) { - local_sum += x[d]; + local_sum += __bfloat162float(x[d]); } /* Warp-level reduction */ @@ -287,7 +287,7 @@ extern "C" __global__ void dt_layernorm_kernel( /* Compute variance */ float local_var = 0.0f; for (int d = tid; d < E; d += blockDim.x) { - float diff = x[d] - mean; + float diff = __bfloat162float(x[d]) - mean; local_var += diff * diff; } @@ -312,7 +312,7 @@ extern "C" __global__ void dt_layernorm_kernel( /* Normalize + affine */ for (int d = tid; d < E; d += blockDim.x) { - out[d] = gamma[d] * (x[d] - mean) * inv_std + beta[d]; + out[d] = gamma[d] * (__bfloat162float(x[d]) - mean) * inv_std + beta[d]; } } @@ -328,11 +328,11 @@ extern "C" __global__ void dt_layernorm_kernel( * Block: (min(E, 256), 1, 1) * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dt_ffn_kernel( - const float* __restrict__ input, /* [B*T, E] */ - const float* __restrict__ W1, /* [E, 4E] */ - const float* __restrict__ b1, /* [4E] */ - const float* __restrict__ W2, /* [4E, E] */ - const float* __restrict__ b2, /* [E] */ + const __nv_bfloat16* __restrict__ input, /* [B*T, E] */ + const __nv_bfloat16* __restrict__ W1, /* [E, 4E] */ + const __nv_bfloat16* __restrict__ b1, /* [4E] */ + const __nv_bfloat16* __restrict__ W2, /* [4E, E] */ + const __nv_bfloat16* __restrict__ b2, /* [E] */ float* __restrict__ output, /* [B*T, E] */ int N, /* B*T */ int E @@ -342,7 +342,7 @@ extern "C" __global__ void dt_ffn_kernel( int tid = threadIdx.x; int FFN = E * 4; - const float* x = input + n * E; + const __nv_bfloat16* x = input + n * E; float* out = output + n * E; /* Each thread computes its output elements via stride loop. @@ -354,7 +354,7 @@ extern "C" __global__ void dt_ffn_kernel( for (int ff = tid; ff < FFN; ff += blockDim.x) { float h = b1[ff]; for (int j = 0; j < E; j++) { - h += x[j] * W1[j * FFN + ff]; + h += __bfloat162float(x[j]) * W1[j * FFN + ff]; } /* GELU approximation */ float x3 = h * h * h; @@ -366,7 +366,7 @@ extern "C" __global__ void dt_ffn_kernel( /* Step 2: output = hidden @ W2 + b2 + input (residual) */ for (int d = tid; d < E; d += blockDim.x) { - float val = b2[d] + x[d]; /* bias + residual */ + float val = b2[d] + __bfloat162float(x[d]); /* bias + residual */ for (int ff = 0; ff < FFN; ff++) { val += sh_hidden[ff] * W2[ff * E + d]; } @@ -385,9 +385,9 @@ extern "C" __global__ void dt_ffn_kernel( * Block: (min(A, 256), 1, 1) * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dt_action_head_kernel( - const float* __restrict__ input, /* [B*T, E] */ - const float* __restrict__ W_head, /* [E, A] */ - const float* __restrict__ b_head, /* [A] */ + const __nv_bfloat16* __restrict__ input, /* [B*T, E] */ + const __nv_bfloat16* __restrict__ W_head, /* [E, A] */ + const __nv_bfloat16* __restrict__ b_head, /* [A] */ float* __restrict__ logits, /* [B*T, A] */ int N, /* B*T */ int E, @@ -397,13 +397,13 @@ extern "C" __global__ void dt_action_head_kernel( if (n >= N) return; int tid = threadIdx.x; - const float* x = input + n * E; + const __nv_bfloat16* x = input + n * E; float* out = logits + n * A; for (int a = tid; a < A; a += blockDim.x) { float val = b_head[a]; for (int d = 0; d < E; d++) { - val += x[d] * W_head[d * A + a]; + val += __bfloat162float(x[d]) * W_head[d * A + a]; } out[a] = val; } @@ -419,7 +419,7 @@ extern "C" __global__ void dt_action_head_kernel( * Block: (min(A, 256), 1, 1) or (32, 1, 1) for warp reduction * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dt_cross_entropy_kernel( - const float* __restrict__ logits, /* [B*T, A] */ + const __nv_bfloat16* __restrict__ logits, /* [B*T, A] */ const int* __restrict__ targets, /* [B*T] */ float* __restrict__ per_sample_loss, /* [B*T] */ float* __restrict__ total_loss, /* [1] atomicAdd */ @@ -429,23 +429,23 @@ extern "C" __global__ void dt_cross_entropy_kernel( int n = blockIdx.x; if (n >= N) return; - const float* lg = logits + n * A; + const __nv_bfloat16* lg = logits + n * A; int target = targets[n]; /* Find max for numerical stability */ float max_val = -1e30f; for (int a = 0; a < A; a++) { - if (lg[a] > max_val) max_val = lg[a]; + if (__bfloat162float(lg[a]) > max_val) max_val = __bfloat162float(lg[a]); } /* log-sum-exp */ float sum_exp = 0.0f; for (int a = 0; a < A; a++) { - sum_exp += expf(lg[a] - max_val); + sum_exp += expf(__bfloat162float(lg[a]) - max_val); } float log_sum_exp = max_val + logf(sum_exp + 1e-8f); - float target_logit = (target >= 0 && target < A) ? lg[target] : 0.0f; + float target_logit = (target >= 0 && target < A) ? __bfloat162float(lg[target]) : 0.0f; float loss = log_sum_exp - target_logit; /* Clamp to prevent NaN propagation */ @@ -464,7 +464,7 @@ extern "C" __global__ void dt_cross_entropy_kernel( * Block: (min(A, 256), 1, 1) * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dt_ce_backward_kernel( - const float* __restrict__ logits, /* [B*T, A] */ + const __nv_bfloat16* __restrict__ logits, /* [B*T, A] */ const int* __restrict__ targets, /* [B*T] */ float* __restrict__ d_logits, /* [B*T, A] */ int N, /* B*T */ @@ -474,26 +474,26 @@ extern "C" __global__ void dt_ce_backward_kernel( if (n >= N) return; int tid = threadIdx.x; - const float* lg = logits + n * A; + const __nv_bfloat16* lg = logits + n * A; float* d_lg = d_logits + n * A; int target = targets[n]; /* Numerically stable softmax */ float max_val = -1e30f; for (int a = 0; a < A; a++) { - if (lg[a] > max_val) max_val = lg[a]; + if (__bfloat162float(lg[a]) > max_val) max_val = __bfloat162float(lg[a]); } float sum_exp = 0.0f; for (int a = 0; a < A; a++) { - sum_exp += expf(lg[a] - max_val); + sum_exp += expf(__bfloat162float(lg[a]) - max_val); } float inv_sum = 1.0f / (sum_exp + 1e-8f); /* d_logits = softmax - one_hot, scaled by 1/N for mean loss */ float scale = 1.0f / (float)N; for (int a = tid; a < A; a += blockDim.x) { - float softmax_a = expf(lg[a] - max_val) * inv_sum; + float softmax_a = expf(__bfloat162float(lg[a]) - max_val) * inv_sum; float one_hot = (a == target) ? 1.0f : 0.0f; d_lg[a] = (softmax_a - one_hot) * scale; } @@ -513,9 +513,9 @@ extern "C" __global__ void dt_ce_backward_kernel( * Note: weight gradients use atomicAdd across samples. * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dt_linear_backward_kernel( - const float* __restrict__ d_output, /* [N, O] */ - const float* __restrict__ input, /* [N, I] */ - const float* __restrict__ W, /* [I, O] */ + const __nv_bfloat16* __restrict__ d_output, /* [N, O] */ + const __nv_bfloat16* __restrict__ input, /* [N, I] */ + const __nv_bfloat16* __restrict__ W, /* [I, O] */ float* __restrict__ dW, /* [I, O] atomicAdd */ float* __restrict__ db, /* [O] atomicAdd */ float* __restrict__ d_input, /* [N, I] */ @@ -527,30 +527,30 @@ extern "C" __global__ void dt_linear_backward_kernel( if (n >= N) return; int tid = threadIdx.x; - const float* dout = d_output + n * O; - const float* x = input + n * I; + const __nv_bfloat16* dout = d_output + n * O; + const __nv_bfloat16* x = input + n * I; float* dx = d_input + n * I; /* Compute d_input[n][i] = sum_o(d_output[n][o] * W[i][o]) */ for (int i = tid; i < I; i += blockDim.x) { float val = 0.0f; for (int o = 0; o < O; o++) { - val += dout[o] * W[i * O + o]; + val += __bfloat162float(dout[o]) * W[i * O + o]; } dx[i] = val; } /* Accumulate weight + bias gradients via atomicAdd */ for (int i = tid; i < I; i += blockDim.x) { - float xi = x[i]; + float xi = __bfloat162float(x[i]); for (int o = 0; o < O; o++) { - atomicAdd(&dW[i * O + o], xi * dout[o]); + atomicAdd(&dW[i * O + o], xi * __bfloat162float(dout[o])); } } /* Bias gradient: only one thread per sample contributes */ for (int o = tid; o < O; o += blockDim.x) { - atomicAdd(&db[o], dout[o]); + atomicAdd(&db[o], __bfloat162float(dout[o])); } } @@ -563,7 +563,7 @@ extern "C" __global__ void dt_linear_backward_kernel( * Block: (256, 1, 1) * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dt_zero_kernel( - float* __restrict__ buf, + __nv_bfloat16* __restrict__ buf, int n ) { int i = blockIdx.x * blockDim.x + threadIdx.x; @@ -581,8 +581,8 @@ extern "C" __global__ void dt_zero_kernel( * Block: (256, 1, 1) * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dt_residual_add_kernel( - float* __restrict__ output, - const float* __restrict__ residual, + __nv_bfloat16* __restrict__ output, + const __nv_bfloat16* __restrict__ residual, int n ) { int i = blockIdx.x * blockDim.x + threadIdx.x; @@ -599,7 +599,7 @@ extern "C" __global__ void dt_residual_add_kernel( * Block: (1, 1, 1) * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dt_return_to_go_kernel( - const float* __restrict__ rewards, /* [num_episodes, episode_len] */ + const __nv_bfloat16* __restrict__ rewards, /* [num_episodes, episode_len] */ float* __restrict__ returns_to_go, /* [num_episodes, episode_len] */ float gamma, int episode_len, @@ -608,7 +608,7 @@ extern "C" __global__ void dt_return_to_go_kernel( int ep = blockIdx.x; if (ep >= num_episodes) return; - const float* r = rewards + ep * episode_len; + const __nv_bfloat16* r = rewards + ep * episode_len; float* rtg = returns_to_go + ep * episode_len; /* Backward scan: R_{T-1} = r_{T-1}, R_t = r_t + gamma * R_{t+1} */ @@ -636,8 +636,8 @@ extern "C" __global__ void dt_return_to_go_kernel( * Block: (min(T, 256), 1, 1) — one thread per timestep within episode * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dt_build_trajectories_kernel( - const float* __restrict__ features, /* [num_bars, feat_dim] */ - const float* __restrict__ returns_to_go, /* [num_episodes, T] */ + const __nv_bfloat16* __restrict__ features, /* [num_bars, feat_dim] */ + const __nv_bfloat16* __restrict__ returns_to_go, /* [num_episodes, T] */ const int* __restrict__ actions, /* [num_bars] — global per-bar actions */ const int* __restrict__ episode_start_indices, /* [num_episodes] */ float* __restrict__ trajectories, /* [num_episodes, T, input_dim] */ @@ -664,9 +664,9 @@ extern "C" __global__ void dt_build_trajectories_kernel( out[0] = returns_to_go[ep * T + t]; /* Slots 1..feat_dim: state features */ - const float* feat = features + bar_idx * feat_dim; + const __nv_bfloat16* feat = features + bar_idx * feat_dim; for (int f = 0; f < feat_dim; f++) { - out[1 + f] = feat[f]; + out[1 + f] = __bfloat162float(feat[f]); } /* Slot state_dim+1: action as float (state_dim = feat_dim) */ @@ -688,7 +688,7 @@ extern "C" __global__ void dt_build_trajectories_kernel( * Block: (256, 1, 1) * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void dt_compute_rewards_actions_kernel( - const float* __restrict__ targets, /* [num_bars, 4] — O,H,L,C per bar */ + const __nv_bfloat16* __restrict__ targets, /* [num_bars, 4] — O,H,L,C per bar */ float* __restrict__ rewards, /* [num_bars] */ int* __restrict__ actions, /* [num_bars] */ int num_bars, diff --git a/crates/ml/src/cuda_pipeline/ema_kernel.cu b/crates/ml/src/cuda_pipeline/ema_kernel.cu index e5759d53e..70953db98 100644 --- a/crates/ml/src/cuda_pipeline/ema_kernel.cu +++ b/crates/ml/src/cuda_pipeline/ema_kernel.cu @@ -7,13 +7,15 @@ */ extern "C" __global__ -void dqn_ema_kernel(float* __restrict__ target, - const float* __restrict__ online, +void dqn_ema_kernel(__nv_bfloat16* __restrict__ target, + const __nv_bfloat16* __restrict__ online, float tau, int n) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i < n) { - target[i] = (1.0f - tau) * target[i] + tau * online[i]; + float t = __bfloat162float(target[i]); + float o = __bfloat162float(online[i]); + target[i] = __float2bfloat16((1.0f - tau) * t + tau * o); } } diff --git a/crates/ml/src/cuda_pipeline/ensemble_kernels.cu b/crates/ml/src/cuda_pipeline/ensemble_kernels.cu index 09d18aa8e..3f7428193 100644 --- a/crates/ml/src/cuda_pipeline/ensemble_kernels.cu +++ b/crates/ml/src/cuda_pipeline/ensemble_kernels.cu @@ -1,3 +1,5 @@ +#include + /** * Ensemble Q-Network kernels — aggregate and diversity loss. * @@ -27,9 +29,9 @@ * Launch config: grid=(ceil(B*num_actions/256), 1, 1), block=(256, 1, 1). * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void ensemble_aggregate_kernel( - const float* __restrict__ head_q_values, /* [K * B * num_actions] flat */ - float* __restrict__ mean_q, /* [B * num_actions] output: mean */ - float* __restrict__ var_q, /* [B * num_actions] output: variance */ + const __nv_bfloat16* __restrict__ head_q_values, /* [K * B * num_actions] flat */ + __nv_bfloat16* __restrict__ mean_q, /* [B * num_actions] output: mean */ + __nv_bfloat16* __restrict__ var_q, /* [B * num_actions] output: variance */ int K, int B, int num_actions @@ -40,15 +42,15 @@ extern "C" __global__ void ensemble_aggregate_kernel( float sum = 0.0f, sum_sq = 0.0f; for (int k = 0; k < K; k++) { - float q = head_q_values[(long long)k * total + idx]; + float q = __bfloat162float(head_q_values[(long long)k * total + idx]); sum += q; sum_sq += q * q; } float mean = sum / (float)K; - mean_q[idx] = mean; + mean_q[idx] = __float2bfloat16(mean); /* Var = E[X^2] - E[X]^2, clamped to 0 for numerical stability */ float v = sum_sq / (float)K - mean * mean; - var_q[idx] = (v > 0.0f) ? v : 0.0f; + var_q[idx] = __float2bfloat16((v > 0.0f) ? v : 0.0f); } /* ══════════════════════════════════════════════════════════════════════ @@ -70,8 +72,8 @@ extern "C" __global__ void ensemble_aggregate_kernel( * Launch config: grid=(ceil(B*num_pairs/256), 1, 1), block=(256, 1, 1). * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void ensemble_diversity_kernel( - const float* __restrict__ head_logits, /* [K * B * num_atoms] */ - float* __restrict__ diversity_loss, /* [1] output: accumulated KL */ + const __nv_bfloat16* __restrict__ head_logits, /* [K * B * num_atoms] */ + __nv_bfloat16* __restrict__ diversity_loss, /* [1] output: accumulated KL */ int K, int B, int num_atoms @@ -101,26 +103,26 @@ extern "C" __global__ void ensemble_diversity_kernel( } /* Pointers to head i and head j logits for this sample */ - const float* logits_i = head_logits + (long long)hi * B * num_atoms + sample * num_atoms; - const float* logits_j = head_logits + (long long)hj * B * num_atoms + sample * num_atoms; + const __nv_bfloat16* logits_i = head_logits + (long long)hi * B * num_atoms + sample * num_atoms; + const __nv_bfloat16* logits_j = head_logits + (long long)hj * B * num_atoms + sample * num_atoms; /* Softmax of logits_i */ - float max_i = logits_i[0]; - for (int a = 1; a < num_atoms; a++) max_i = fmaxf(max_i, logits_i[a]); + float max_i = __bfloat162float(logits_i[0]); + for (int a = 1; a < num_atoms; a++) max_i = fmaxf(max_i, __bfloat162float(logits_i[a])); float sum_i = 0.0f; - for (int a = 0; a < num_atoms; a++) sum_i += expf(logits_i[a] - max_i); + for (int a = 0; a < num_atoms; a++) sum_i += expf(__bfloat162float(logits_i[a]) - max_i); /* Softmax of logits_j */ - float max_j = logits_j[0]; - for (int a = 1; a < num_atoms; a++) max_j = fmaxf(max_j, logits_j[a]); + float max_j = __bfloat162float(logits_j[0]); + for (int a = 1; a < num_atoms; a++) max_j = fmaxf(max_j, __bfloat162float(logits_j[a])); float sum_j = 0.0f; - for (int a = 0; a < num_atoms; a++) sum_j += expf(logits_j[a] - max_j); + for (int a = 0; a < num_atoms; a++) sum_j += expf(__bfloat162float(logits_j[a]) - max_j); /* KL(p_i || p_j) = sum_a p_i[a] * log(p_i[a] / p_j[a]) */ float kl = 0.0f; for (int a = 0; a < num_atoms; a++) { - float pi = expf(logits_i[a] - max_i) / sum_i; - float pj = expf(logits_j[a] - max_j) / sum_j; + float pi = expf(__bfloat162float(logits_i[a]) - max_i) / sum_i; + float pj = expf(__bfloat162float(logits_j[a]) - max_j) / sum_j; /* Symmetrized KL (Jensen–Shannon style): (KL(i||j) + KL(j||i)) / 2 */ float log_ratio_ij = logf(pi / (pj + 1e-8f) + 1e-8f); float log_ratio_ji = logf(pj / (pi + 1e-8f) + 1e-8f); @@ -147,7 +149,7 @@ extern "C" __global__ void ensemble_diversity_kernel( for (int offset = 16; offset > 0; offset >>= 1) val += __shfl_xor_sync(0xFFFFFFFF, val, offset); if (warp_lane == 0) - atomicAdd(diversity_loss, val); + atomicAddBF16(diversity_loss, __float2bfloat16(val)); } } @@ -167,8 +169,8 @@ extern "C" __global__ void ensemble_diversity_kernel( * ══════════════════════════════════════════════════════════════════════ */ extern "C" __global__ void ensemble_kl_gradient_kernel( - const float* __restrict__ head_logits, /* [K * B * NA] */ - float* __restrict__ d_logits_0, /* [B * NA] output: diversity gradient for head 0 */ + const __nv_bfloat16* __restrict__ head_logits, /* [K * B * NA] */ + __nv_bfloat16* __restrict__ d_logits_0, /* [B * NA] output: diversity gradient for head 0 */ float diversity_weight, int K, int B, int NA ) { @@ -181,31 +183,31 @@ extern "C" __global__ void ensemble_kl_gradient_kernel( /* Softmax of head 0 at (b, a) — numerically stable */ float max0 = -1e30f; for (int j = 0; j < NA; j++) { - float v = head_logits[0 * stride + b * NA + j]; + float v = __bfloat162float(head_logits[0 * stride + b * NA + j]); if (v > max0) max0 = v; } float sum0 = 0.0f; for (int j = 0; j < NA; j++) - sum0 += expf(head_logits[0 * stride + b * NA + j] - max0); - float p0 = expf(head_logits[0 * stride + b * NA + a] - max0) / (sum0 + 1e-8f); + sum0 += expf(__bfloat162float(head_logits[0 * stride + b * NA + j]) - max0); + float p0 = expf(__bfloat162float(head_logits[0 * stride + b * NA + a]) - max0) / (sum0 + 1e-8f); /* Average gradient from all other heads: mean_k(p0 - pk) */ float grad = 0.0f; for (int k = 1; k < K; k++) { float maxk = -1e30f; for (int j = 0; j < NA; j++) { - float v = head_logits[k * stride + b * NA + j]; + float v = __bfloat162float(head_logits[k * stride + b * NA + j]); if (v > maxk) maxk = v; } float sumk = 0.0f; for (int j = 0; j < NA; j++) - sumk += expf(head_logits[k * stride + b * NA + j] - maxk); - float pk = expf(head_logits[k * stride + b * NA + a] - maxk) / (sumk + 1e-8f); + sumk += expf(__bfloat162float(head_logits[k * stride + b * NA + j]) - maxk); + float pk = expf(__bfloat162float(head_logits[k * stride + b * NA + a]) - maxk) / (sumk + 1e-8f); grad += (p0 - pk); } grad /= (float)(K - 1); /* NEGATIVE: maximize divergence (subtract from loss) */ - d_logits_0[b * NA + a] = -diversity_weight * grad; + d_logits_0[b * NA + a] = __float2bfloat16(-diversity_weight * grad); } diff --git a/crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu b/crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu index 43f212c20..11596820b 100644 --- a/crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu +++ b/crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu @@ -18,7 +18,7 @@ * Used when fill simulation is not needed (e.g. inference, CPU routing fallback). */ extern "C" __global__ void epsilon_greedy_select( - const float* __restrict__ q_values, /* [batch_size, num_actions] */ + const __nv_bfloat16* __restrict__ q_values, /* [batch_size, num_actions] */ unsigned int* rng_states, /* [batch_size] — persistent LCG state */ unsigned int* actions_out, /* [batch_size] — output action indices */ const float epsilon, @@ -39,8 +39,8 @@ extern "C" __global__ void epsilon_greedy_select( if (action >= (unsigned int)num_actions) action = (unsigned int)(num_actions - 1); } else { /* Greedy: argmax over Q-values for this batch element */ - const float* q = q_values + idx * num_actions; - float best_q = q[0]; + const __nv_bfloat16* q = q_values + idx * num_actions; + __nv_bfloat16 best_q = q[0]; action = 0; for (int a = 1; a < num_actions; a++) { if (q[a] > best_q) { @@ -50,8 +50,8 @@ extern "C" __global__ void epsilon_greedy_select( } /* Q-gap conviction filter: force flat when conviction is low */ if (q_gap_threshold > 0.0f && num_actions > 2) { - float q_flat = q[num_actions / 2]; /* center = Flat */ - if (best_q - q_flat < q_gap_threshold) { + float q_flat_f = __bfloat162float(q[num_actions / 2]); /* center = Flat */ + if (__bfloat162float(best_q) - q_flat_f < q_gap_threshold) { action = 2; } } @@ -81,7 +81,7 @@ extern "C" __global__ void epsilon_greedy_select( * fill_mask_out — 1 if filled, 0 if not (optional: NULL to skip) */ extern "C" __global__ void epsilon_greedy_routed( - const float* __restrict__ q_values, /* [batch_size, num_actions] */ + const __nv_bfloat16* __restrict__ q_values, /* [batch_size, num_actions] */ unsigned int* rng_states, /* [batch_size] */ unsigned int* actions_out, /* [batch_size] — post-fill exposure */ int* fill_mask_out, /* [batch_size] — 1=filled, 0=not (nullable) */ @@ -113,8 +113,8 @@ extern "C" __global__ void epsilon_greedy_routed( exposure_idx = (int)(gpu_random(&rng) * (float)num_actions); if (exposure_idx >= num_actions) exposure_idx = num_actions - 1; } else { - const float* q = q_values + idx * num_actions; - float best_q = q[0]; + const __nv_bfloat16* q = q_values + idx * num_actions; + __nv_bfloat16 best_q = q[0]; exposure_idx = 0; for (int a = 1; a < num_actions; a++) { if (q[a] > best_q) { @@ -125,7 +125,7 @@ extern "C" __global__ void epsilon_greedy_routed( /* Q-gap conviction filter */ if (q_gap_threshold > 0.0f && num_actions > 2) { float q_flat = q[num_actions / 2]; - if (best_q - q_flat < q_gap_threshold) { + if (__bfloat162float(best_q) - q_flat < q_gap_threshold) { exposure_idx = num_actions / 2; } } @@ -171,17 +171,17 @@ extern "C" __global__ void epsilon_greedy_routed( * Launch: grid=(ceil(batch_size/256),1,1), block=(256,1,1). */ extern "C" __global__ void branching_action_select( - const float* __restrict__ q_exposure, /* [batch_size, 5] */ - const float* __restrict__ q_order, /* [batch_size, 3] */ - const float* __restrict__ q_urgency, /* [batch_size, 3] */ + const __nv_bfloat16* __restrict__ q_exposure, /* [batch_size, 5] */ + const __nv_bfloat16* __restrict__ q_order, /* [batch_size, 3] */ + const __nv_bfloat16* __restrict__ q_urgency, /* [batch_size, 3] */ unsigned int* rng_states, /* [batch_size] */ unsigned int* actions_out, /* [batch_size] — factored index 0-44 */ const float epsilon, const int batch_size, const float q_gap_threshold, /* min Q-gap for trade entry (0.0 = disabled) */ - const float* __restrict__ bonus_exposure, /* [5] UCB count bonus per exposure action (NULL = disabled) */ - const float* __restrict__ bonus_order, /* [3] UCB count bonus per order action (NULL = disabled) */ - const float* __restrict__ bonus_urgency /* [3] UCB count bonus per urgency action (NULL = disabled) */ + const __nv_bfloat16* __restrict__ bonus_exposure, /* [5] UCB count bonus per exposure action (NULL = disabled) */ + const __nv_bfloat16* __restrict__ bonus_order, /* [3] UCB count bonus per order action (NULL = disabled) */ + const __nv_bfloat16* __restrict__ bonus_urgency /* [3] UCB count bonus per urgency action (NULL = disabled) */ ) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= batch_size) return; @@ -195,11 +195,11 @@ extern "C" __global__ void branching_action_select( exposure = (int)(gpu_random(&rng) * 5.0f); if (exposure >= 5) exposure = 4; } else { - const float* qe = q_exposure + idx * 5; - float best = qe[0] + ((bonus_exposure != NULL) ? bonus_exposure[0] : 0.0f); + const __nv_bfloat16* qe = q_exposure + idx * 5; + float best = __bfloat162float(qe[0]) + ((bonus_exposure != NULL) ? __bfloat162float(bonus_exposure[0]) : 0.0f); exposure = 0; for (int a = 1; a < 5; a++) { - float q_plus_bonus = qe[a] + ((bonus_exposure != NULL) ? bonus_exposure[a] : 0.0f); + float q_plus_bonus = __bfloat162float(qe[a]) + ((bonus_exposure != NULL) ? __bfloat162float(bonus_exposure[a]) : 0.0f); if (q_plus_bonus > best) { best = q_plus_bonus; exposure = a; } } /* Q-gap filter: force flat when conviction is low. @@ -207,8 +207,8 @@ extern "C" __global__ void branching_action_select( * Compare best Q (without bonus) against Q(Flat) — conviction is * about Q-value quality, not exploration bonus. */ if (q_gap_threshold > 0.0f) { - float q_flat = qe[2]; /* index 2 = Flat (center of 5-action branching) */ - float q_best_raw = qe[exposure]; /* raw Q without bonus */ + float q_flat = __bfloat162float(qe[2]); /* index 2 = Flat (center of 5-action branching) */ + float q_best_raw = __bfloat162float(qe[exposure]); /* raw Q without bonus */ if (q_best_raw - q_flat < q_gap_threshold) { exposure = 2; } @@ -222,11 +222,11 @@ extern "C" __global__ void branching_action_select( order = (int)(gpu_random(&rng) * 3.0f); if (order >= 3) order = 2; } else { - const float* qo = q_order + idx * 3; - float best = qo[0] + ((bonus_order != NULL) ? bonus_order[0] : 0.0f); + const __nv_bfloat16* qo = q_order + idx * 3; + float best = __bfloat162float(qo[0]) + ((bonus_order != NULL) ? __bfloat162float(bonus_order[0]) : 0.0f); order = 0; for (int a = 1; a < 3; a++) { - float q_plus_bonus = qo[a] + ((bonus_order != NULL) ? bonus_order[a] : 0.0f); + float q_plus_bonus = __bfloat162float(qo[a]) + ((bonus_order != NULL) ? __bfloat162float(bonus_order[a]) : 0.0f); if (q_plus_bonus > best) { best = q_plus_bonus; order = a; } } } @@ -238,11 +238,11 @@ extern "C" __global__ void branching_action_select( urgency = (int)(gpu_random(&rng) * 3.0f); if (urgency >= 3) urgency = 2; } else { - const float* qu = q_urgency + idx * 3; - float best = qu[0] + ((bonus_urgency != NULL) ? bonus_urgency[0] : 0.0f); + const __nv_bfloat16* qu = q_urgency + idx * 3; + float best = __bfloat162float(qu[0]) + ((bonus_urgency != NULL) ? __bfloat162float(bonus_urgency[0]) : 0.0f); urgency = 0; for (int a = 1; a < 3; a++) { - float q_plus_bonus = qu[a] + ((bonus_urgency != NULL) ? bonus_urgency[a] : 0.0f); + float q_plus_bonus = __bfloat162float(qu[a]) + ((bonus_urgency != NULL) ? __bfloat162float(bonus_urgency[a]) : 0.0f); if (q_plus_bonus > best) { best = q_plus_bonus; urgency = a; } } } diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 1c1aa7d29..ab1ce5397 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -131,10 +131,10 @@ __device__ __forceinline__ int argmax_n(const float* arr, int n) { * @param market_dim number of market features per bar */ extern "C" __global__ void experience_state_gather( - const float* __restrict__ market_features, + const __nv_bfloat16* __restrict__ market_features, const int* __restrict__ episode_starts, const int* __restrict__ current_timesteps, - const float* __restrict__ portfolio_states, + const __nv_bfloat16* __restrict__ portfolio_states, float* batch_states, int N, int total_bars, @@ -156,9 +156,9 @@ extern "C" __global__ void experience_state_gather( } /* -- Market features: [0 .. market_dim) -- */ - const float* mf_row = market_features + (long long)bar_idx * market_dim; + const __nv_bfloat16* mf_row = market_features + (long long)bar_idx * market_dim; for (int k = 0; k < market_dim; k++) - out[k] = mf_row[k]; + out[k] = __bfloat162float(mf_row[k]); /* -- Portfolio features: [market_dim .. market_dim+8) -- * @@ -178,16 +178,16 @@ extern "C" __global__ void experience_state_gather( * approximation here (max_pos not available in gather kernel). * The network learns the actual scale via the position feature. */ - const float* ps = portfolio_states + (long long)i * PORTFOLIO_STRIDE; - float position = ps[0]; - float cash = ps[1]; - float portfolio_value = ps[2]; - float peak_equity = ps[7]; - float prev_equity = ps[9]; - float hold_time = ps[10]; - float realized_pnl = ps[11]; - float entry_price = ps[12]; - float trade_start_pnl = ps[13]; + const __nv_bfloat16* ps = portfolio_states + (long long)i * PORTFOLIO_STRIDE; + float position = __bfloat162float(ps[0]); + float cash = __bfloat162float(ps[1]); + float portfolio_value = __bfloat162float(ps[2]); + float peak_equity = __bfloat162float(ps[7]); + float prev_equity = __bfloat162float(ps[9]); + float hold_time = __bfloat162float(ps[10]); + float realized_pnl = __bfloat162float(ps[11]); + float entry_price = __bfloat162float(ps[12]); + float trade_start_pnl = __bfloat162float(ps[13]); float equity = (portfolio_value > 1.0f) ? portfolio_value : 1.0f; float drawdown = (peak_equity > 1.0f) @@ -246,10 +246,10 @@ extern "C" __global__ void experience_state_gather( if (past_idx >= 0 && slot + 3 < state_dim) { /* Close prices: feature[0] of each bar is typically close or a return. * Use raw market_features — index 0 is the first feature per bar. */ - const float* now_row = market_features + (long long)bar_idx * market_dim; - const float* past_row = market_features + (long long)past_idx * market_dim; - float close_now = now_row[0]; - float close_past = past_row[0]; + const __nv_bfloat16* now_row = market_features + (long long)bar_idx * market_dim; + const __nv_bfloat16* past_row = market_features + (long long)past_idx * market_dim; + float close_now = __bfloat162float(now_row[0]); + float close_past = __bfloat162float(past_row[0]); /* Return over N bars */ float ret = (close_past > 0.0f) ? (close_now - close_past) / close_past : 0.0f; @@ -277,7 +277,7 @@ extern "C" __global__ void experience_state_gather( /* Volume trend: current vs average */ float avg_vol = (vol_count > 0) ? vol_sum / (float)vol_count : 1.0f; - float cur_vol = (market_dim > 4) ? now_row[4] : 1.0f; + float cur_vol = (market_dim > 4) ? __bfloat162float(now_row[4]) : 1.0f; float vol_ratio = (avg_vol > 0.0f) ? cur_vol / avg_vol : 1.0f; out[slot + 2] = fmaxf(0.0f, fminf(5.0f, vol_ratio)); /* clamp 0-5x */ @@ -348,7 +348,7 @@ extern "C" __global__ void experience_action_select( int b1_size, int b2_size, float q_gap_threshold - ,const float* __restrict__ portfolio_states, /* [N, 20] read-only */ + ,const __nv_bfloat16* __restrict__ portfolio_states, /* [N, 20] read-only */ int min_hold_bars, float max_position ) { @@ -509,7 +509,7 @@ extern "C" __global__ void experience_action_select( * @param margin_pct initial margin as fraction of notional (0.06 = 6%) */ extern "C" __global__ void experience_env_step( - const float* __restrict__ targets, + const __nv_bfloat16* __restrict__ targets, const int* __restrict__ episode_starts, int* current_timesteps, const int* __restrict__ actions, @@ -518,11 +518,11 @@ extern "C" __global__ void experience_env_step( int* out_actions, float* out_rewards, float* out_dones, - const float* __restrict__ batch_states, + const __nv_bfloat16* __restrict__ batch_states, float max_position, float tx_cost_multiplier, float loss_aversion, - const float* __restrict__ features, + const __nv_bfloat16* __restrict__ features, int market_dim, int L, int N, @@ -532,8 +532,8 @@ extern "C" __global__ void experience_env_step( int b1_size, int b2_size, int current_t, - const float* __restrict__ cvar_scales, /* [N] or NULL — CVaR position scaling */ - const float* __restrict__ q_gaps, /* [N] or NULL — Q-gap conviction scaling */ + const __nv_bfloat16* __restrict__ cvar_scales, /* [N] or NULL — CVaR position scaling */ + const __nv_bfloat16* __restrict__ q_gaps, /* [N] or NULL — Q-gap conviction scaling */ float* raw_returns_out, /* [N, L] output: true per-bar portfolio return (unshapen) */ int min_hold_bars, /* minimum bars to hold before exiting or reversing */ float spread_cost, /* bid-ask spread cost per unit (matches backtest) */ @@ -550,10 +550,10 @@ extern "C" __global__ void experience_env_step( long long out_off = (long long)i * L + current_t; /* ---- Write current state to replay buffer ---- */ - const float* src = batch_states + (long long)i * state_dim; + const __nv_bfloat16* src = batch_states + (long long)i * state_dim; float* dst = out_states + out_off * state_dim; for (int k = 0; k < state_dim; k++) - dst[k] = src[k]; + dst[k] = __bfloat162float(src[k]); /* ---- Write action ---- */ int action_idx = actions[i]; @@ -571,9 +571,9 @@ extern "C" __global__ void experience_env_step( /* Target layout: [preproc_close, preproc_next, raw_close, raw_next] * [0:1] = preprocessed (log-return normalized) — for Q-network * [2:3] = raw dollar prices — for portfolio simulation P&L + tx costs */ - const float* tgt = targets + (long long)bar_idx * 4; - float raw_close = tgt[2]; - float raw_next = tgt[3]; + const __nv_bfloat16* tgt = targets + (long long)bar_idx * 4; + float raw_close = __bfloat162float(tgt[2]); + float raw_next = __bfloat162float(tgt[3]); /* Guard against degenerate prices from data gaps. */ if (raw_close <= 0.0f) raw_close = 1.0f; @@ -581,23 +581,23 @@ extern "C" __global__ void experience_env_step( /* ---- Read full portfolio state (PORTFOLIO_STRIDE=20) ---- */ float* ps = portfolio_states + (long long)i * PORTFOLIO_STRIDE; - float position = ps[0]; - float cash = ps[1]; + float position = __bfloat162float(ps[0]); + float cash = __bfloat162float(ps[1]); /* ps[2] = portfolio_value (updated at end) */ - /* ps[3:6] reserved (unused by reward v6) */ - float peak_equity = ps[7]; - float flat_counter = ps[8]; - float prev_equity = ps[9]; - float hold_time = ps[10]; + /* __bfloat162float(ps[3:6]) reserved (unused by reward v6) */ + float peak_equity = __bfloat162float(ps[7]); + float flat_counter = __bfloat162float(ps[8]); + float prev_equity = __bfloat162float(ps[9]); + float hold_time = __bfloat162float(ps[10]); /* ps[11] = realized_pnl (cumulative, updated at end) */ - float entry_price = ps[12]; /* price when trade was entered */ - float trade_start_pnl = ps[13]; /* realized_pnl snapshot at trade entry */ - float win_count = ps[14]; /* Kelly: number of profitable trade exits */ - float loss_count = ps[15]; /* Kelly: number of losing trade exits */ - float sum_wins = ps[16]; /* Kelly: cumulative profit from winners */ - float sum_losses = ps[17]; /* Kelly: cumulative |loss| from losers */ - float sum_returns = ps[18]; /* Kelly: cumulative net returns (for μ) */ - float sum_sq_returns = ps[19]; /* Kelly: cumulative squared returns (for σ²) */ + float entry_price = __bfloat162float(ps[12]); /* price when trade was entered */ + float trade_start_pnl = __bfloat162float(ps[13]); /* realized_pnl snapshot at trade entry */ + float win_count = __bfloat162float(ps[14]); /* Kelly: number of profitable trade exits */ + float loss_count = __bfloat162float(ps[15]); /* Kelly: number of losing trade exits */ + float sum_wins = __bfloat162float(ps[16]); /* Kelly: cumulative profit from winners */ + float sum_losses = __bfloat162float(ps[17]); /* Kelly: cumulative |loss| from losers */ + float sum_returns = __bfloat162float(ps[18]); /* Kelly: cumulative net returns (for μ) */ + float sum_sq_returns = __bfloat162float(ps[19]); /* Kelly: cumulative squared returns (for σ²) */ /* Pre-trade capital floor: skip trade execution on blown accounts. * When the floor triggers, write done=1 AND reset the portfolio to @@ -605,7 +605,7 @@ extern "C" __global__ void experience_env_step( * the reset, the blown portfolio persists and every subsequent step * hits the floor again (producing fake 92% MaxDD from stuck episodes). */ { - float portfolio_val = ps[2]; + float portfolio_val = __bfloat162float(ps[2]); if (check_capital_floor(portfolio_val, peak_equity)) { out_rewards[out_off] = -10.0f; out_dones[out_off] = 1.0f; @@ -664,12 +664,12 @@ extern "C" __global__ void experience_env_step( } /* Kelly criterion: scale by optimal fraction after sufficient samples. - * win_count/loss_count already declared earlier from ps[14]/ps[15]. */ + * win_count/loss_count already declared earlier from __bfloat162float(ps[14])/__bfloat162float(ps[15]). */ { float total_trades = win_count + loss_count; if (total_trades >= 20.0f) { - float sum_wins_val = ps[16]; - float sum_losses_val = ps[17]; + float sum_wins_val = __bfloat162float(ps[16]); + float sum_losses_val = __bfloat162float(ps[17]); float avg_win = sum_wins_val / fmaxf(win_count, 1.0f); float avg_loss = sum_losses_val / fmaxf(loss_count, 1.0f); float win_rate = win_count / total_trades; @@ -691,7 +691,7 @@ extern "C" __global__ void experience_env_step( * margin ≈ 6% of notional (CME ES initial margin ~$15K per contract). */ { float margin_per_contract = raw_close * contract_multiplier * margin_pct; - float portfolio_val = ps[2]; + float portfolio_val = __bfloat162float(ps[2]); target_position = apply_margin_cap(target_position, portfolio_val, margin_per_contract, peak_equity); } @@ -716,7 +716,7 @@ extern "C" __global__ void experience_env_step( float pre_trade_position = position; /* Sign of position BEFORE trade: -1 (short), 0 (flat), +1 (long) */ - int prev_sign = (ps[0] > 0.001f) ? 1 : ((ps[0] < -0.001f) ? -1 : 0); + int prev_sign = (__bfloat162float(ps[0]) > 0.001f) ? 1 : ((__bfloat162float(ps[0]) < -0.001f) ? -1 : 0); /* ════════════════════════════════════════════════════════════════════ * DYNAMIC TRAILING STOP — regime-adaptive, locks in profits. @@ -754,14 +754,14 @@ extern "C" __global__ void experience_env_step( * violations, but epsilon exploration can still slip through. * ════════════════════════════════════════════════════════════════════ */ int is_last_bar = (bar_idx >= total_bars - 1) ? 1 : 0; - float enforced_position = enforce_hold(ps[0], target_position, hold_time, min_hold_bars, is_last_bar); + float enforced_position = enforce_hold(__bfloat162float(ps[0]), target_position, hold_time, min_hold_bars, is_last_bar); int hold_violation = (fabsf(enforced_position - target_position) > 0.001f); if (hold_violation) { target_position = enforced_position; /* keep old position */ /* Fix the stored action to match held exposure (prevent action aliasing) */ - float prev_exposure_frac = ps[0] / (max_position > 0.0f ? max_position : 1.0f); + float prev_exposure_frac = __bfloat162float(ps[0]) / (max_position > 0.0f ? max_position : 1.0f); int held_exposure = (int)roundf((prev_exposure_frac + 1.0f) * 0.5f * (float)(b0_size - 1)); if (held_exposure < 0) held_exposure = 0; if (held_exposure >= b0_size) held_exposure = b0_size - 1; @@ -800,14 +800,14 @@ extern "C" __global__ void experience_env_step( if (entering_trade) { entry_price = raw_close; - trade_start_pnl = ps[11]; + trade_start_pnl = __bfloat162float(ps[11]); } /* On reversal: close old segment, open new one. * MUST run AFTER hold enforcement — if hold guard cancelled the reversal, * reversing_trade is 0 and this block correctly does nothing. */ if (reversing_trade) { - float closing_pnl = ps[11] + old_pos_pnl - trade_start_pnl; + float closing_pnl = __bfloat162float(ps[11]) + old_pos_pnl - trade_start_pnl; reversal_return = closing_pnl / (prev_equity > 1.0f ? prev_equity : 1.0f); /* Kelly stats for completed segment */ @@ -823,7 +823,7 @@ extern "C" __global__ void experience_env_step( /* Reset for new segment */ entry_price = raw_close; - trade_start_pnl = ps[11] + old_pos_pnl; + trade_start_pnl = __bfloat162float(ps[11]) + old_pos_pnl; } /* Save hold_time BEFORE reset — sparse reward needs the pre-reset value @@ -871,7 +871,7 @@ extern "C" __global__ void experience_env_step( /* ---- Sparse: trade completion reward (vol-normalized) ---- */ if (segment_complete && segment_hold_time > 0.0f) { - float segment_pnl = ps[11] + raw_pnl - trade_start_pnl; + float segment_pnl = __bfloat162float(ps[11]) + raw_pnl - trade_start_pnl; float segment_return = segment_pnl / (prev_equity > 1.0f ? prev_equity : 1.0f); /* Kelly statistics for EXIT (reversal Kelly already done above) */ @@ -973,7 +973,7 @@ extern "C" __global__ void experience_env_step( ps[0] = position; ps[1] = cash; ps[2] = new_portfolio_value; - /* ps[3:6] reserved — reward v6 does not use DSR/PnL EMA */ + /* __bfloat162float(ps[3:6]) reserved — reward v6 does not use DSR/PnL EMA */ ps[7] = peak_equity; ps[8] = flat_counter; ps[9] = new_portfolio_value; /* prev_equity = current equity for next step */ @@ -982,9 +982,9 @@ extern "C" __global__ void experience_env_step( * (saved before position was overwritten at ps[0] = position). * For non-reversal bars, raw_pnl (= current position's PnL) is correct. */ if (reversing_trade) { - ps[11] = ps[11] + old_pos_pnl; /* old position's PnL (saved at line ~700) */ + ps[11] = __bfloat162float(ps[11]) + old_pos_pnl; /* old position's PnL (saved at line ~700) */ } else { - ps[11] = ps[11] + raw_pnl; /* normal: current position's PnL */ + ps[11] = __bfloat162float(ps[11]) + raw_pnl; /* normal: current position's PnL */ } ps[12] = entry_price; /* preserved across bars of a trade */ ps[13] = trade_start_pnl; /* realized_pnl snapshot at trade entry */ @@ -1067,7 +1067,7 @@ __device__ __forceinline__ float action_to_tx_cost(int action_idx) { } extern "C" __global__ void portfolio_sim_kernel( - const float* __restrict__ targets, + const __nv_bfloat16* __restrict__ targets, const int* __restrict__ actions, float* portfolio_state, float* portfolio_out, @@ -1241,8 +1241,8 @@ extern "C" __global__ void portfolio_sim_kernel( * @param v_max C51 maximum support */ extern "C" __global__ void compute_expected_q( - const float* __restrict__ v_logits, - const float* __restrict__ b_logits, + const __nv_bfloat16* __restrict__ v_logits, + const __nv_bfloat16* __restrict__ b_logits, float* q_values, int N, int num_atoms, @@ -1259,7 +1259,7 @@ extern "C" __global__ void compute_expected_q( float dz = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f; /* Per-sample value logits */ - const float* v_row = v_logits + (long long)i * num_atoms; + const __nv_bfloat16* v_row = v_logits + (long long)i * num_atoms; const float* b_row = b_logits + (long long)i * total_actions * num_atoms; /* Compute advantage mean across all actions for dueling subtraction */ @@ -1276,19 +1276,19 @@ extern "C" __global__ void compute_expected_q( /* Softmax + expectation over atoms */ float max_logit = -1e30f; for (int z = 0; z < num_atoms; z++) { - float logit = v_row[z] + adv_a[z]; + float logit = __bfloat162float(v_row[z]) + adv_a[z]; if (logit > max_logit) max_logit = logit; } float sum_exp = 0.0f; for (int z = 0; z < num_atoms; z++) { - float logit = v_row[z] + adv_a[z]; + float logit = __bfloat162float(v_row[z]) + adv_a[z]; sum_exp += expf(logit - max_logit); } float expected_q = 0.0f; for (int z = 0; z < num_atoms; z++) { - float logit = v_row[z] + adv_a[z]; + float logit = __bfloat162float(v_row[z]) + adv_a[z]; float prob = expf(logit - max_logit) / sum_exp; float z_val = v_min + (float)z * dz; expected_q += prob * z_val; @@ -1320,9 +1320,9 @@ extern "C" __global__ void compute_expected_q( */ extern "C" __global__ void expert_action_override( int* out_actions, /* [N] Q-network actions (overwritten) */ - const float* __restrict__ targets, /* [total_bars, 4] OHLCV price data */ - const float* __restrict__ features, /* [total_bars, MARKET_DIM] market features */ - const float* __restrict__ portfolio_states, /* [N, PORTFOLIO_STRIDE] per-episode state */ + const __nv_bfloat16* __restrict__ targets, /* [total_bars, 4] OHLCV price data */ + const __nv_bfloat16* __restrict__ features, /* [total_bars, MARKET_DIM] market features */ + const __nv_bfloat16* __restrict__ portfolio_states, /* [N, PORTFOLIO_STRIDE] per-episode state */ const int* __restrict__ episode_starts, /* [N] bar offset per episode */ const int* __restrict__ current_timesteps, /* [N] current step in episode */ unsigned int* rng_states, /* [N] per-episode RNG (updated) */ diff --git a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs index b358417ee..10c8dd9a9 100644 --- a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs @@ -2,7 +2,7 @@ //! GPU-fused epsilon-greedy action selection -- pure cudarc. //! -//! All inputs are `CudaSlice` Q-values, all outputs are `CudaSlice` action +//! All inputs are `CudaSlice` Q-values, all outputs are `CudaSlice` action //! indices. The selector stores an `Arc`. //! //! Pure cudarc implementation with no Candle dependency. @@ -41,9 +41,9 @@ pub struct GpuActionSelector { bonus_exposure_ptr: u64, bonus_order_ptr: u64, bonus_urgency_ptr: u64, - bonus_exposure_buf: Option>, - bonus_order_buf: Option>, - bonus_urgency_buf: Option>, + bonus_exposure_buf: Option>, + bonus_order_buf: Option>, + bonus_urgency_buf: Option>, } impl GpuActionSelector { @@ -102,7 +102,7 @@ impl GpuActionSelector { Ok(()) } - pub fn select_actions(&mut self, q_values: &CudaSlice, epsilon: f32, batch_size: usize, num_actions: usize) -> Result, MLError> { + pub fn select_actions(&mut self, q_values: &CudaSlice, epsilon: f32, batch_size: usize, num_actions: usize) -> Result, MLError> { if batch_size == 0 { return self.stream.alloc_zeros::(0).map_err(|e| MLError::ModelError(format!("alloc empty: {e}"))); } if batch_size > self.max_batch_size { return Err(MLError::ModelError(format!("batch_size {batch_size} exceeds max_batch_size {}", self.max_batch_size))); } let config = launch_config_1d(batch_size); @@ -119,7 +119,7 @@ impl GpuActionSelector { } #[allow(clippy::too_many_arguments)] - pub fn select_actions_routed(&mut self, q_values: &CudaSlice, epsilon: f32, batch_size: usize, num_actions: usize, + pub fn select_actions_routed(&mut self, q_values: &CudaSlice, epsilon: f32, batch_size: usize, num_actions: usize, step_offset: i32, spread: f32, median_spread: f32, volatility: f32, median_vol: f32, spread_bps: f32, ioc_fill_prob: f32, limit_fill_min: f32, limit_fill_max: f32, spread_cost_frac: f32, spread_capture_frac: f32, @@ -141,8 +141,8 @@ impl GpuActionSelector { self.copy_actions_out(batch_size) } - pub fn select_actions_branching(&mut self, q_exposure: &CudaSlice, q_order: &CudaSlice, - q_urgency: &CudaSlice, epsilon: f32, batch_size: usize, + pub fn select_actions_branching(&mut self, q_exposure: &CudaSlice, q_order: &CudaSlice, + q_urgency: &CudaSlice, epsilon: f32, batch_size: usize, ) -> Result, MLError> { if batch_size == 0 { return self.stream.alloc_zeros::(0).map_err(|e| MLError::ModelError(format!("alloc empty: {e}"))); } if batch_size > self.max_batch_size { return Err(MLError::ModelError(format!("batch_size {} exceeds max {}", batch_size, self.max_batch_size))); } @@ -242,7 +242,7 @@ mod tests { let mut selector = GpuActionSelector::new(stream.clone(), batch_size, 12345).expect("init"); // Allocate random Q-values on GPU via host upload let q_host: Vec = (0..batch_size * num_actions).map(|i| (i as f32) * 0.1 - 8.0).collect(); - let mut q_buf = stream.alloc_zeros::(batch_size * num_actions).expect("alloc q_values"); + let mut q_buf = stream.alloc_zeros::(batch_size * num_actions).expect("alloc q_values"); stream.memcpy_htod(&q_host, &mut q_buf).expect("upload q_values"); let greedy_actions = selector.select_actions(&q_buf, 0.0, batch_size, num_actions).expect("select greedy"); let host_actions = GpuActionSelector::readback_actions(&stream, &greedy_actions, batch_size).expect("readback"); diff --git a/crates/ml/src/cuda_pipeline/gpu_attention.rs b/crates/ml/src/cuda_pipeline/gpu_attention.rs index aef78d669..bb0ec0d5a 100644 --- a/crates/ml/src/cuda_pipeline/gpu_attention.rs +++ b/crates/ml/src/cuda_pipeline/gpu_attention.rs @@ -62,23 +62,23 @@ pub struct GpuAttention { grad_norm_kernel: CudaFunction, adam_kernel: CudaFunction, /// Attention parameters (Xavier-initialized, trainable). - params: CudaSlice, + params: CudaSlice, /// Output buffer for attended states [batch_size, state_dim]. - output_buf: CudaSlice, + output_buf: CudaSlice, /// Saved input states from forward pass (for backward recomputation). - saved_input: CudaSlice, + saved_input: CudaSlice, /// Scratch buffer for input gradient [batch_size, state_dim]. /// Written by backward kernel via atomicAdd. Not propagated further /// (attention is the first layer — no upstream to backprop into). - d_input_scratch: CudaSlice, + d_input_scratch: CudaSlice, /// Gradient accumulator for attention parameters [total_params]. - d_params: CudaSlice, + d_params: CudaSlice, /// Gradient norm output buffer [1]. - grad_norm_buf: CudaSlice, + grad_norm_buf: CudaSlice, /// Adam first moment (m) for attention parameters [total_params]. - attn_m: CudaSlice, + attn_m: CudaSlice, /// Adam second moment (v) for attention parameters [total_params]. - attn_v: CudaSlice, + attn_v: CudaSlice, /// Adam step counter (1-indexed). attn_adam_step: i32, /// Total attention parameters count. @@ -144,28 +144,28 @@ impl GpuAttention { } // beta stays 0 - let mut params = stream.alloc_zeros::(total_params) + let mut params = stream.alloc_zeros::(total_params) .map_err(|e| MLError::ModelError(format!("attention params alloc: {e}")))?; stream.memcpy_htod(&host_params, &mut params) .map_err(|e| MLError::ModelError(format!("attention params upload: {e}")))?; - let output_buf = stream.alloc_zeros::(b * d) + let output_buf = stream.alloc_zeros::(b * d) .map_err(|e| MLError::ModelError(format!("attention output alloc: {e}")))?; // Saved state buffer for backward pass (input states) - let saved_input = stream.alloc_zeros::(b * d) + let saved_input = stream.alloc_zeros::(b * d) .map_err(|e| MLError::ModelError(format!("attention saved_input alloc: {e}")))?; // Gradient and optimizer buffers - let d_input_scratch = stream.alloc_zeros::(b * d) + let d_input_scratch = stream.alloc_zeros::(b * d) .map_err(|e| MLError::ModelError(format!("attention d_input_scratch alloc: {e}")))?; - let d_params = stream.alloc_zeros::(total_params) + let d_params = stream.alloc_zeros::(total_params) .map_err(|e| MLError::ModelError(format!("attention d_params alloc: {e}")))?; - let grad_norm_buf = stream.alloc_zeros::(1) + let grad_norm_buf = stream.alloc_zeros::(1) .map_err(|e| MLError::ModelError(format!("attention grad_norm alloc: {e}")))?; - let attn_m = stream.alloc_zeros::(total_params) + let attn_m = stream.alloc_zeros::(total_params) .map_err(|e| MLError::ModelError(format!("attention adam_m alloc: {e}")))?; - let attn_v = stream.alloc_zeros::(total_params) + let attn_v = stream.alloc_zeros::(total_params) .map_err(|e| MLError::ModelError(format!("attention adam_v alloc: {e}")))?; // VRAM: params(4) + d_params + m + v + output + saved_input + d_input_scratch + grad_norm @@ -204,7 +204,7 @@ impl GpuAttention { /// /// Saves input states and pre-layernorm output for the backward pass. /// Returns a reference to the output buffer (attended states). - pub fn forward(&mut self, states: &CudaSlice, batch_size: usize) -> Result<&CudaSlice, MLError> { + pub fn forward(&mut self, states: &CudaSlice, batch_size: usize) -> Result<&CudaSlice, MLError> { let b = batch_size; let d = self.config.state_dim; @@ -266,7 +266,7 @@ impl GpuAttention { /// further since attention is the first layer). pub fn backward( &mut self, - d_output: &CudaSlice, + d_output: &CudaSlice, batch_size: usize, ) -> Result<(), MLError> { let b = batch_size; @@ -403,7 +403,7 @@ impl GpuAttention { } /// Get the output buffer reference. - pub fn output(&self) -> &CudaSlice { + pub fn output(&self) -> &CudaSlice { &self.output_buf } } diff --git a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs index 7d594c2b4..c6f832d3f 100644 --- a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs +++ b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs @@ -249,23 +249,23 @@ pub struct GpuBacktestEvaluator { gather_kernel: CudaFunction, // Uploaded data (read-only; persists across the step loop) - prices_buf: CudaSlice, // [n_windows * max_len * 4] - features_buf: CudaSlice, // [n_windows * max_len * feat_dim] + prices_buf: CudaSlice, // [n_windows * max_len * 4] + features_buf: CudaSlice, // [n_windows * max_len * feat_dim] window_lens_buf: CudaSlice, // [n_windows] // Mutable state (written by kernels each step) - portfolio_buf: CudaSlice, // [n_windows * 8] - step_rewards_buf: CudaSlice, // [n_windows] - step_returns_buf: CudaSlice, // [n_windows * max_len] + portfolio_buf: CudaSlice, // [n_windows * 8] + step_rewards_buf: CudaSlice, // [n_windows] + step_returns_buf: CudaSlice, // [n_windows * max_len] done_buf: CudaSlice, // [n_windows] actions_buf: CudaSlice, // [n_windows] actions_history_buf: CudaSlice, // [n_windows * max_len] // Gather kernel output buffer (overwritten every step) - states_buf: CudaSlice, // [n_windows * state_dim] (8-aligned) + states_buf: CudaSlice, // [n_windows * state_dim] (8-aligned) // Output buffer (written by metrics kernel) - metrics_buf: CudaSlice, // [n_windows * 14] + metrics_buf: CudaSlice, // [n_windows * 14] // Dimensions and config n_windows: usize, @@ -289,20 +289,20 @@ pub struct GpuBacktestEvaluator { /// Flat F32 weight buffer for cuBLAS (20 tensors concatenated). /// Created on first `evaluate_dqn()` call. - cublas_params_flat: Option>, + cublas_params_flat: Option>, /// cuBLAS activation scratch buffers (created on first call). - cublas_h_s1: Option>, - cublas_h_s2: Option>, - cublas_h_v: Option>, - cublas_h_b0: Option>, - cublas_h_b1: Option>, - cublas_h_b2: Option>, - cublas_v_logits: Option>, - cublas_b_logits: Option>, + cublas_h_s1: Option>, + cublas_h_s2: Option>, + cublas_h_v: Option>, + cublas_h_b0: Option>, + cublas_h_b1: Option>, + cublas_h_b2: Option>, + cublas_v_logits: Option>, + cublas_b_logits: Option>, /// Q-values buffer for C51 expected-Q computation. - cublas_q_values: Option>, + cublas_q_values: Option>, /// `compute_expected_q` kernel (loaded from experience_kernels.cu). expected_q_kernel: Option, @@ -323,10 +323,10 @@ pub struct GpuBacktestEvaluator { /// Q-gap output buffer for `experience_action_select` kernel (per-step path). /// [n_windows] f32 on GPU — values are unused by backtest but the kernel writes them. - q_gaps_buf: Option>, + q_gaps_buf: Option>, /// Chunked Q-gap output buffer (chunked path: [n_windows * CHUNK_SIZE]). - chunked_q_gaps_buf: Option>, + chunked_q_gaps_buf: Option>, /// Cached CUDA graph for `evaluate_dqn_graphed()` step loop. /// @@ -351,20 +351,20 @@ pub struct GpuBacktestEvaluator { chunked_cublas_forward: Option, /// Chunked activation scratch buffers (sized for n_windows * CHUNK_SIZE). - chunked_h_s1: Option>, - chunked_h_s2: Option>, - chunked_h_v: Option>, - chunked_h_b0: Option>, - chunked_h_b1: Option>, - chunked_h_b2: Option>, - chunked_v_logits: Option>, - chunked_b_logits: Option>, + chunked_h_s1: Option>, + chunked_h_s2: Option>, + chunked_h_v: Option>, + chunked_h_b0: Option>, + chunked_h_b1: Option>, + chunked_h_b2: Option>, + chunked_v_logits: Option>, + chunked_b_logits: Option>, /// Q-values buffer for chunked C51 expected-Q (n_windows * CHUNK_SIZE rows). - chunked_q_values: Option>, + chunked_q_values: Option>, /// Batched states buffer for chunked gather: [n_windows * CHUNK_SIZE, state_dim]. - chunked_states_buf: Option>, + chunked_states_buf: Option>, /// Batched forward actions buffer for chunked action select: [n_windows * CHUNK_SIZE]. chunked_actions_buf: Option>, @@ -518,10 +518,10 @@ impl GpuBacktestEvaluator { .map_err(|e| MLError::ModelError(format!("portfolio alloc: {e}")))?; let step_rewards_buf = stream - .alloc_zeros::(n_windows) + .alloc_zeros::(n_windows) .map_err(|e| MLError::ModelError(format!("step_rewards alloc: {e}")))?; let step_returns_buf = stream - .alloc_zeros::(n_windows * max_len) + .alloc_zeros::(n_windows * max_len) .map_err(|e| MLError::ModelError(format!("step_returns alloc: {e}")))?; let done_buf = stream .alloc_zeros::(n_windows) @@ -533,7 +533,7 @@ impl GpuBacktestEvaluator { .alloc_zeros::(n_windows * max_len) .map_err(|e| MLError::ModelError(format!("actions_history alloc: {e}")))?; let metrics_buf = stream - .alloc_zeros::(n_windows * 14) + .alloc_zeros::(n_windows * 14) .map_err(|e| MLError::ModelError(format!("metrics alloc: {e}")))?; // State layout matches training experience_state_gather: @@ -543,7 +543,7 @@ impl GpuBacktestEvaluator { const PORTFOLIO_AND_MTF_DIM: usize = 24; // 8 portfolio + 16 multi-timeframe let state_dim = (feature_dim + PORTFOLIO_AND_MTF_DIM + 7) & !7; let states_buf = stream - .alloc_zeros::(n_windows * state_dim) + .alloc_zeros::(n_windows * state_dim) .map_err(|e| MLError::ModelError(format!("states_buf alloc: {e}")))?; // Pre-allocate forward kernel output buffer (used by evaluate_dqn path) @@ -669,7 +669,7 @@ impl GpuBacktestEvaluator { step: usize, portfolio_dim: usize, _device: &(), - ) -> Result, MLError> { + ) -> Result, MLError> { if portfolio_dim != self.portfolio_dim { return Err(MLError::ConfigError(format!( "gather_states: portfolio_dim={portfolio_dim} != expected {}", @@ -718,9 +718,9 @@ impl GpuBacktestEvaluator { .map_err(|e| MLError::ModelError(format!("gather_states launch step {step}: {e}")))?; } - // Return a view of the states buffer as a CudaSlice. + // Return a view of the states buffer as a CudaSlice. let n_elems = self.n_windows * state_dim; - let dst = self.stream.alloc_zeros::(n_elems) + let dst = self.stream.alloc_zeros::(n_elems) .map_err(|e| MLError::ModelError(format!("alloc states step {step}: {e}")))?; let src_view = self.states_buf.slice(..n_elems); { @@ -752,7 +752,7 @@ impl GpuBacktestEvaluator { portfolio_dim: usize, ) -> Result, MLError> where - F: Fn(&CudaSlice, usize, usize) -> Result, MLError>, + F: Fn(&CudaSlice, usize, usize) -> Result, MLError>, { let _nvtx = NvtxRange::new("backtest_evaluate"); @@ -1253,7 +1253,7 @@ impl GpuBacktestEvaluator { /// replacing the candle `signal_to_action_scores` + `argmax` pipeline. /// /// # Arguments - /// * `forward_fn` — model forward: `&CudaSlice[N * state_dim]` → `CudaSlice[N]` predictions + /// * `forward_fn` — model forward: `&CudaSlice[N * state_dim]` → `CudaSlice[N]` predictions /// * `high_threshold_bps` — strong signal threshold (actions 0 and 4) /// * `low_threshold_bps` — mild signal threshold (actions 1 and 3) /// * `portfolio_dim` — must be 24 (8 portfolio + 16 multi-timeframe) @@ -1265,7 +1265,7 @@ impl GpuBacktestEvaluator { portfolio_dim: usize, ) -> Result, MLError> where - F: Fn(&CudaSlice, usize, usize) -> Result, MLError>, + F: Fn(&CudaSlice, usize, usize) -> Result, MLError>, { if portfolio_dim != self.portfolio_dim { return Err(MLError::ConfigError(format!( @@ -1368,27 +1368,27 @@ impl GpuBacktestEvaluator { let param_sizes = compute_backtest_param_sizes(dqn_cfg, self.state_dim); let total_params: usize = param_sizes.iter().sum(); let params_flat = self.stream - .alloc_zeros::(total_params) + .alloc_zeros::(total_params) .map_err(|e| MLError::ModelError(format!("alloc cublas_params_flat: {e}")))?; // Allocate activation scratch buffers - let h_s1 = self.stream.alloc_zeros::(n * dqn_cfg.shared_h1) + let h_s1 = self.stream.alloc_zeros::(n * dqn_cfg.shared_h1) .map_err(|e| MLError::ModelError(format!("alloc h_s1: {e}")))?; - let h_s2 = self.stream.alloc_zeros::(n * dqn_cfg.shared_h2) + let h_s2 = self.stream.alloc_zeros::(n * dqn_cfg.shared_h2) .map_err(|e| MLError::ModelError(format!("alloc h_s2: {e}")))?; - let h_v = self.stream.alloc_zeros::(n * dqn_cfg.value_h) + let h_v = self.stream.alloc_zeros::(n * dqn_cfg.value_h) .map_err(|e| MLError::ModelError(format!("alloc h_v: {e}")))?; - let h_b0 = self.stream.alloc_zeros::(n * dqn_cfg.adv_h) + let h_b0 = self.stream.alloc_zeros::(n * dqn_cfg.adv_h) .map_err(|e| MLError::ModelError(format!("alloc h_b0: {e}")))?; - let h_b1 = self.stream.alloc_zeros::(n * dqn_cfg.adv_h) + let h_b1 = self.stream.alloc_zeros::(n * dqn_cfg.adv_h) .map_err(|e| MLError::ModelError(format!("alloc h_b1: {e}")))?; - let h_b2 = self.stream.alloc_zeros::(n * dqn_cfg.adv_h) + let h_b2 = self.stream.alloc_zeros::(n * dqn_cfg.adv_h) .map_err(|e| MLError::ModelError(format!("alloc h_b2: {e}")))?; - let v_logits = self.stream.alloc_zeros::(n * na) + let v_logits = self.stream.alloc_zeros::(n * na) .map_err(|e| MLError::ModelError(format!("alloc v_logits: {e}")))?; - let b_logits = self.stream.alloc_zeros::(n * total_branch_atoms) + let b_logits = self.stream.alloc_zeros::(n * total_branch_atoms) .map_err(|e| MLError::ModelError(format!("alloc b_logits: {e}")))?; - let q_values = self.stream.alloc_zeros::(n * total_actions) + let q_values = self.stream.alloc_zeros::(n * total_actions) .map_err(|e| MLError::ModelError(format!("alloc q_values: {e}")))?; // Compile experience kernels for compute_expected_q + action_select @@ -1418,7 +1418,7 @@ impl GpuBacktestEvaluator { // Q-gap output buffer: same sizing as rng_states — chunked action_select // writes batch = n * chunk_len entries per call. - let q_gaps_buf = self.stream.alloc_zeros::(cn) + let q_gaps_buf = self.stream.alloc_zeros::(cn) .map_err(|e| MLError::ModelError(format!("alloc q_gaps_buf: {e}")))?; let chunked_cublas = CublasForward::new( @@ -1435,26 +1435,26 @@ impl GpuBacktestEvaluator { dqn_cfg.branch_2_size, )?; - let ch_h_s1 = self.stream.alloc_zeros::(cn * dqn_cfg.shared_h1) + let ch_h_s1 = self.stream.alloc_zeros::(cn * dqn_cfg.shared_h1) .map_err(|e| MLError::ModelError(format!("alloc chunked h_s1: {e}")))?; - let ch_h_s2 = self.stream.alloc_zeros::(cn * dqn_cfg.shared_h2) + let ch_h_s2 = self.stream.alloc_zeros::(cn * dqn_cfg.shared_h2) .map_err(|e| MLError::ModelError(format!("alloc chunked h_s2: {e}")))?; - let ch_h_v = self.stream.alloc_zeros::(cn * dqn_cfg.value_h) + let ch_h_v = self.stream.alloc_zeros::(cn * dqn_cfg.value_h) .map_err(|e| MLError::ModelError(format!("alloc chunked h_v: {e}")))?; - let ch_h_b0 = self.stream.alloc_zeros::(cn * dqn_cfg.adv_h) + let ch_h_b0 = self.stream.alloc_zeros::(cn * dqn_cfg.adv_h) .map_err(|e| MLError::ModelError(format!("alloc chunked h_b0: {e}")))?; - let ch_h_b1 = self.stream.alloc_zeros::(cn * dqn_cfg.adv_h) + let ch_h_b1 = self.stream.alloc_zeros::(cn * dqn_cfg.adv_h) .map_err(|e| MLError::ModelError(format!("alloc chunked h_b1: {e}")))?; - let ch_h_b2 = self.stream.alloc_zeros::(cn * dqn_cfg.adv_h) + let ch_h_b2 = self.stream.alloc_zeros::(cn * dqn_cfg.adv_h) .map_err(|e| MLError::ModelError(format!("alloc chunked h_b2: {e}")))?; - let ch_v_logits = self.stream.alloc_zeros::(cn * na) + let ch_v_logits = self.stream.alloc_zeros::(cn * na) .map_err(|e| MLError::ModelError(format!("alloc chunked v_logits: {e}")))?; - let ch_b_logits = self.stream.alloc_zeros::(cn * total_branch_atoms) + let ch_b_logits = self.stream.alloc_zeros::(cn * total_branch_atoms) .map_err(|e| MLError::ModelError(format!("alloc chunked b_logits: {e}")))?; - let ch_q_values = self.stream.alloc_zeros::(cn * total_actions) + let ch_q_values = self.stream.alloc_zeros::(cn * total_actions) .map_err(|e| MLError::ModelError(format!("alloc chunked q_values: {e}")))?; - let ch_states_buf = self.stream.alloc_zeros::(cn * self.state_dim) + let ch_states_buf = self.stream.alloc_zeros::(cn * self.state_dim) .map_err(|e| MLError::ModelError(format!("alloc chunked states_buf: {e}")))?; let ch_actions_buf = self.stream.alloc_zeros::(cn) .map_err(|e| MLError::ModelError(format!("alloc chunked actions_buf: {e}")))?; @@ -1462,7 +1462,7 @@ impl GpuBacktestEvaluator { let ch_rng_seeds: Vec = (0..cn).map(|_| fastrand::u32(..)).collect(); let ch_rng_states = self.stream.clone_htod(&ch_rng_seeds) .map_err(|e| MLError::ModelError(format!("alloc chunked rng_states: {e}")))?; - let ch_q_gaps = self.stream.alloc_zeros::(cn) + let ch_q_gaps = self.stream.alloc_zeros::(cn) .map_err(|e| MLError::ModelError(format!("alloc chunked q_gaps: {e}")))?; let chunked_mem_mb = ( @@ -1813,7 +1813,7 @@ fn compute_backtest_param_sizes(cfg: &DqnBacktestConfig, state_dim: usize) -> [u } /// Extract raw F32 device pointer from a CudaSlice. -fn raw_f32_ptr(slice: &CudaSlice, stream: &Arc) -> u64 { +fn raw_f32_ptr(slice: &CudaSlice, stream: &Arc) -> u64 { let (ptr, guard) = slice.device_ptr(stream); let _no_drop = ManuallyDrop::new(guard); ptr diff --git a/crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs index 0494e38cb..df36b57b4 100644 --- a/crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs @@ -77,25 +77,25 @@ pub struct GpuCuriosityTrainer { fused_zero_fwd_bwd_adam_func: CudaFunction, // Gradient buffers - grad_w1: CudaSlice, // [CUR_W1_LEN] - grad_b1: CudaSlice, // [CUR_B1_LEN] - grad_w2: CudaSlice, // [CUR_W2_LEN] - grad_b2: CudaSlice, // [CUR_B2_LEN] + grad_w1: CudaSlice, // [CUR_W1_LEN] + grad_b1: CudaSlice, // [CUR_B1_LEN] + grad_w2: CudaSlice, // [CUR_W2_LEN] + grad_b2: CudaSlice, // [CUR_B2_LEN] // Adam first moment (per-param-group) - adam_m_w1: CudaSlice, // [CUR_W1_LEN] - adam_m_b1: CudaSlice, // [CUR_B1_LEN] - adam_m_w2: CudaSlice, // [CUR_W2_LEN] - adam_m_b2: CudaSlice, // [CUR_B2_LEN] + adam_m_w1: CudaSlice, // [CUR_W1_LEN] + adam_m_b1: CudaSlice, // [CUR_B1_LEN] + adam_m_w2: CudaSlice, // [CUR_W2_LEN] + adam_m_b2: CudaSlice, // [CUR_B2_LEN] // Adam second moment (per-param-group) - adam_v_w1: CudaSlice, // [CUR_W1_LEN] - adam_v_b1: CudaSlice, // [CUR_B1_LEN] - adam_v_w2: CudaSlice, // [CUR_W2_LEN] - adam_v_b2: CudaSlice, // [CUR_B2_LEN] + adam_v_w1: CudaSlice, // [CUR_W1_LEN] + adam_v_b1: CudaSlice, // [CUR_B1_LEN] + adam_v_w2: CudaSlice, // [CUR_W2_LEN] + adam_v_b2: CudaSlice, // [CUR_B2_LEN] // Shifted next_states buffer - next_states_buf: CudaSlice, + next_states_buf: CudaSlice, /// Atomic block-arrival counter for fused kernel grid-wide sync (single i32 on GPU). block_counter: CudaSlice, @@ -118,11 +118,11 @@ pub struct GpuCuriosityTrainer { fn launch_adam_step( stream: &CudaStream, adam_func: &CudaFunction, - params: &mut CudaSlice, - grads: &CudaSlice, + params: &mut CudaSlice, + grads: &CudaSlice, num_params: usize, - m: &mut CudaSlice, - v: &mut CudaSlice, + m: &mut CudaSlice, + v: &mut CudaSlice, batch_size: usize, step: i32, ) -> Result<(), MLError> { @@ -204,50 +204,50 @@ impl GpuCuriosityTrainer { })?; // ---- Allocate gradient buffers ---- - let grad_w1 = stream.alloc_zeros::(CUR_W1_LEN).map_err(|e| { + let grad_w1 = stream.alloc_zeros::(CUR_W1_LEN).map_err(|e| { MLError::ModelError(format!("alloc grad_w1: {e}")) })?; - let grad_b1 = stream.alloc_zeros::(CUR_B1_LEN).map_err(|e| { + let grad_b1 = stream.alloc_zeros::(CUR_B1_LEN).map_err(|e| { MLError::ModelError(format!("alloc grad_b1: {e}")) })?; - let grad_w2 = stream.alloc_zeros::(CUR_W2_LEN).map_err(|e| { + let grad_w2 = stream.alloc_zeros::(CUR_W2_LEN).map_err(|e| { MLError::ModelError(format!("alloc grad_w2: {e}")) })?; - let grad_b2 = stream.alloc_zeros::(CUR_B2_LEN).map_err(|e| { + let grad_b2 = stream.alloc_zeros::(CUR_B2_LEN).map_err(|e| { MLError::ModelError(format!("alloc grad_b2: {e}")) })?; // ---- Allocate Adam first moment buffers ---- - let adam_m_w1 = stream.alloc_zeros::(CUR_W1_LEN).map_err(|e| { + let adam_m_w1 = stream.alloc_zeros::(CUR_W1_LEN).map_err(|e| { MLError::ModelError(format!("alloc adam_m_w1: {e}")) })?; - let adam_m_b1 = stream.alloc_zeros::(CUR_B1_LEN).map_err(|e| { + let adam_m_b1 = stream.alloc_zeros::(CUR_B1_LEN).map_err(|e| { MLError::ModelError(format!("alloc adam_m_b1: {e}")) })?; - let adam_m_w2 = stream.alloc_zeros::(CUR_W2_LEN).map_err(|e| { + let adam_m_w2 = stream.alloc_zeros::(CUR_W2_LEN).map_err(|e| { MLError::ModelError(format!("alloc adam_m_w2: {e}")) })?; - let adam_m_b2 = stream.alloc_zeros::(CUR_B2_LEN).map_err(|e| { + let adam_m_b2 = stream.alloc_zeros::(CUR_B2_LEN).map_err(|e| { MLError::ModelError(format!("alloc adam_m_b2: {e}")) })?; // ---- Allocate Adam second moment buffers ---- - let adam_v_w1 = stream.alloc_zeros::(CUR_W1_LEN).map_err(|e| { + let adam_v_w1 = stream.alloc_zeros::(CUR_W1_LEN).map_err(|e| { MLError::ModelError(format!("alloc adam_v_w1: {e}")) })?; - let adam_v_b1 = stream.alloc_zeros::(CUR_B1_LEN).map_err(|e| { + let adam_v_b1 = stream.alloc_zeros::(CUR_B1_LEN).map_err(|e| { MLError::ModelError(format!("alloc adam_v_b1: {e}")) })?; - let adam_v_w2 = stream.alloc_zeros::(CUR_W2_LEN).map_err(|e| { + let adam_v_w2 = stream.alloc_zeros::(CUR_W2_LEN).map_err(|e| { MLError::ModelError(format!("alloc adam_v_w2: {e}")) })?; - let adam_v_b2 = stream.alloc_zeros::(CUR_B2_LEN).map_err(|e| { + let adam_v_b2 = stream.alloc_zeros::(CUR_B2_LEN).map_err(|e| { MLError::ModelError(format!("alloc adam_v_b2: {e}")) })?; // ---- Allocate shifted next_states buffer ---- let next_states_buf = stream - .alloc_zeros::(max_samples * state_dim) + .alloc_zeros::(max_samples * state_dim) .map_err(|e| { MLError::ModelError(format!("alloc next_states_buf: {e}")) })?; @@ -310,7 +310,7 @@ impl GpuCuriosityTrainer { pub fn train_on_collector_buffers( &mut self, weights: &mut CuriosityWeightSet, - states: &CudaSlice, + states: &CudaSlice, actions: &CudaSlice, n_samples: usize, ) -> Result<(), MLError> { diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 55317ba52..5939cf921 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -383,27 +383,27 @@ pub struct GpuDqnTrainer { clipped_saxpy_kernel: CudaFunction, clip_grad_kernel: CudaFunction, /// Spectral norm left singular vectors: [out_dim] per weight matrix (W_s1, W_s2) - spec_u_s1: CudaSlice, - spec_v_s1: CudaSlice, - spec_u_s2: CudaSlice, - spec_v_s2: CudaSlice, + spec_u_s1: CudaSlice, + spec_v_s1: CudaSlice, + spec_u_s2: CudaSlice, + spec_v_s2: CudaSlice, // Head spectral norm vectors (power iteration) - spec_u_v1: CudaSlice, spec_v_v1: CudaSlice, - spec_u_v2: CudaSlice, spec_v_v2: CudaSlice, - spec_u_a1: CudaSlice, spec_v_a1: CudaSlice, - spec_u_a2: CudaSlice, spec_v_a2: CudaSlice, - spec_u_bo1: CudaSlice, spec_v_bo1: CudaSlice, - spec_u_bo2: CudaSlice, spec_v_bo2: CudaSlice, - spec_u_bu1: CudaSlice, spec_v_bu1: CudaSlice, - spec_u_bu2: CudaSlice, spec_v_bu2: CudaSlice, + spec_u_v1: CudaSlice, spec_v_v1: CudaSlice, + spec_u_v2: CudaSlice, spec_v_v2: CudaSlice, + spec_u_a1: CudaSlice, spec_v_a1: CudaSlice, + spec_u_a2: CudaSlice, spec_v_a2: CudaSlice, + spec_u_bo1: CudaSlice, spec_v_bo1: CudaSlice, + spec_u_bo2: CudaSlice, spec_v_bo2: CudaSlice, + spec_u_bu1: CudaSlice, spec_v_bu1: CudaSlice, + spec_u_bu2: CudaSlice, spec_v_bu2: CudaSlice, /// IQN trunk gradient Adam first moment [trunk_params] (separate from C51 Adam). - iqn_trunk_m: CudaSlice, + iqn_trunk_m: CudaSlice, /// IQN trunk gradient Adam second moment [trunk_params]. - iqn_trunk_v: CudaSlice, + iqn_trunk_v: CudaSlice, /// IQN trunk Adam step counter. iqn_trunk_adam_step: i32, /// IQN trunk grad norm scratch [1]. - iqn_trunk_grad_norm: CudaSlice, + iqn_trunk_grad_norm: CudaSlice, /// IQN trunk Adam step counter on device [1]. iqn_trunk_t_buf: CudaSlice, /// Number of trunk parameters (w_s1 + b_s1 + w_s2 + b_s2). @@ -425,39 +425,39 @@ pub struct GpuDqnTrainer { bf16_padded_total: usize, // ── Batch input buffers (uploaded per step) ───────────────────── - states_buf: CudaSlice, // [B, STATE_DIM] - next_states_buf: CudaSlice, // [B, STATE_DIM] + states_buf: CudaSlice, // [B, STATE_DIM] + next_states_buf: CudaSlice, // [B, STATE_DIM] actions_buf: CudaSlice, // [B] - rewards_buf: CudaSlice, // [B] - dones_buf: CudaSlice, // [B] - is_weights_buf: CudaSlice, // [B] + rewards_buf: CudaSlice, // [B] + dones_buf: CudaSlice, // [B] + is_weights_buf: CudaSlice, // [B] // ── Activation save buffers (forward → backward) ──────────────── - save_h_s1: CudaSlice, // [B, SHARED_H1] - save_h_s2: CudaSlice, // [B, SHARED_H2] - save_h_v: CudaSlice, // [B, VALUE_H] - save_h_b0: CudaSlice, // [B, ADV_H] - save_h_b1: CudaSlice, // [B, ADV_H] - save_h_b2: CudaSlice, // [B, ADV_H] - save_current_lp: CudaSlice, // [B, NUM_BRANCHES(3), NUM_ATOMS] - save_projected: CudaSlice, // [B, NUM_BRANCHES(3), NUM_ATOMS] + save_h_s1: CudaSlice, // [B, SHARED_H1] + save_h_s2: CudaSlice, // [B, SHARED_H2] + save_h_v: CudaSlice, // [B, VALUE_H] + save_h_b0: CudaSlice, // [B, ADV_H] + save_h_b1: CudaSlice, // [B, ADV_H] + save_h_b2: CudaSlice, // [B, ADV_H] + save_current_lp: CudaSlice, // [B, NUM_BRANCHES(3), NUM_ATOMS] + save_projected: CudaSlice, // [B, NUM_BRANCHES(3), NUM_ATOMS] // ── Forward output buffers ────────────────────────────────────── - per_sample_loss_buf: CudaSlice, // [B] - td_errors_buf: CudaSlice, // [B] - total_loss_buf: CudaSlice, // [1] + per_sample_loss_buf: CudaSlice, // [B] + td_errors_buf: CudaSlice, // [B] + total_loss_buf: CudaSlice, // [1] // ── Forward-only Q-value output ───────────────────────────────── - q_out_buf: CudaSlice, // [B, TOTAL_ACTIONS(11)] + q_out_buf: CudaSlice, // [B, TOTAL_ACTIONS(11)] // ── Backward / Adam buffers ───────────────────────────────────── - grad_buf: CudaSlice, // [TOTAL_PARAMS] gradient accumulator - params_buf: CudaSlice, // [TOTAL_PARAMS] flat online parameters - target_params_buf: CudaSlice, // [TOTAL_PARAMS] flat target parameters (EMA) - m_buf: CudaSlice, // [TOTAL_PARAMS] Adam first moment - v_buf: CudaSlice, // [TOTAL_PARAMS] Adam second moment - grad_norm_buf: CudaSlice, // [1] pre-clip gradient L2 norm - cql_grad_scratch: CudaSlice, // [TOTAL_PARAMS] CQL gradient isolation buffer + grad_buf: CudaSlice, // [TOTAL_PARAMS] gradient accumulator + params_buf: CudaSlice, // [TOTAL_PARAMS] flat online parameters + target_params_buf: CudaSlice, // [TOTAL_PARAMS] flat target parameters (EMA) + m_buf: CudaSlice, // [TOTAL_PARAMS] Adam first moment + v_buf: CudaSlice, // [TOTAL_PARAMS] Adam second moment + grad_norm_buf: CudaSlice, // [1] pre-clip gradient L2 norm + cql_grad_scratch: CudaSlice, // [TOTAL_PARAMS] CQL gradient isolation buffer // ── Adam step counter on device (CUDA Graph cannot bake scalars) ─ t_buf: CudaSlice, // [1] current Adam step @@ -480,14 +480,14 @@ pub struct GpuDqnTrainer { /// Single staging buffer for batch upload consolidation. /// Layout: [states(B*SD) | next_states(B*SD) | actions_as_f32(B) | rewards(B) | dones(B) | is_weights(B)] /// One HtoD transfer replaces 6 separate PCIe round-trips. - upload_staging_buf: CudaSlice, + upload_staging_buf: CudaSlice, /// Total element count of the staging buffer. upload_staging_len: usize, /// Single readback buffer for loss download consolidation. /// Layout: [total_loss(1) | grad_norm(1) | td_errors(B)] /// One DtoH transfer replaces 3 separate PCIe round-trips. - readback_buf: CudaSlice, + readback_buf: CudaSlice, /// Pre-allocated host-side readback vec (avoids per-step allocation). readback_host: Vec, @@ -498,11 +498,11 @@ pub struct GpuDqnTrainer { // ── PER priority update (GPU-native) ───────────────────────────── /// Running max priority across all batches in an epoch (atomicMax in kernel). /// Read back once per epoch via `batch_max_readback_and_reset()`. - batch_max_buf: CudaSlice, + batch_max_buf: CudaSlice, /// Scalar-only readback buffer: [total_loss(1) | grad_norm(1)] = 8 bytes. /// Used by `execute_train_scalars_only()` to avoid reading back td_errors. - scalars_readback_buf: CudaSlice, + scalars_readback_buf: CudaSlice, scalars_readback_host: [f32; 2], // ── BF16 batch staging buffers ────────────────────────────────── @@ -521,32 +521,32 @@ pub struct GpuDqnTrainer { // Online activations reuse the existing save_h_* buffers. // Target only needs h_s2 (for DDQN) + scratch for intermediate layers. /// Target trunk layer 1 scratch: [B, SH1] - tg_h_s1_scratch: CudaSlice, + tg_h_s1_scratch: CudaSlice, /// Target trunk layer 2 output — kept for Double-DQN action selection: [B, SH2] - tg_h_s2_buf: CudaSlice, + tg_h_s2_buf: CudaSlice, /// Target value head hidden scratch: [B, VH] - tg_h_v_scratch: CudaSlice, + tg_h_v_scratch: CudaSlice, /// Target branch head hidden scratch (reused for all 3 branches): [B, AH] - tg_h_b_scratch: CudaSlice, + tg_h_b_scratch: CudaSlice, /// Target value head logits: [B, NA] - tg_v_logits_buf: CudaSlice, + tg_v_logits_buf: CudaSlice, /// Target branch logits: [B, (B0+B1+B2)*NA] - tg_b_logits_buf: CudaSlice, + tg_b_logits_buf: CudaSlice, /// Online value head logits (from cuBLAS pass on current states): [B, NA] - on_v_logits_buf: CudaSlice, + on_v_logits_buf: CudaSlice, /// Online branch logits (from cuBLAS pass on current states): [B, (B0+B1+B2)*NA] - on_b_logits_buf: CudaSlice, + on_b_logits_buf: CudaSlice, /// Online-next value logits (Double DQN action selector on next_states): [B, NA] - on_next_v_logits_buf: CudaSlice, + on_next_v_logits_buf: CudaSlice, /// Online-next branch logits (Double DQN on next_states): [B, (B0+B1+B2)*NA] - on_next_b_logits_buf: CudaSlice, + on_next_b_logits_buf: CudaSlice, /// Online-next scratch (reused): h_s1[B,SH1], h_s2[B,SH2], h_v[B,VH], h_b[B,AH] - on_next_h_s1_scratch: CudaSlice, - on_next_h_s2_scratch: CudaSlice, - on_next_h_v_scratch: CudaSlice, - on_next_h_b_scratch: CudaSlice, + on_next_h_s1_scratch: CudaSlice, + on_next_h_s2_scratch: CudaSlice, + on_next_h_v_scratch: CudaSlice, + on_next_h_b_scratch: CudaSlice, /// Compiled C51 loss kernel (standalone, replaces the fused forward+loss kernel). c51_loss_kernel: CudaFunction, @@ -565,12 +565,12 @@ pub struct GpuDqnTrainer { /// When this differs from `loss_mode`, the graph must be recaptured. last_captured_loss_mode: Option, /// Scratch buffers for blended loss (MSE grad stored here, then blended into d_value/d_adv) - d_value_logits_mse: CudaSlice, - d_adv_logits_mse: CudaSlice, + d_value_logits_mse: CudaSlice, + d_adv_logits_mse: CudaSlice, /// Gradient w.r.t. value logits: [B, NA] - d_value_logits_buf: CudaSlice, + d_value_logits_buf: CudaSlice, /// Gradient w.r.t. branch logits: [B, (B0+B1+B2)*NA] - d_adv_logits_buf: CudaSlice, + d_adv_logits_buf: CudaSlice, /// Precomputed F32 weight byte offsets into params_buf. /// Same GOFF_* layout as BF16, but for the F32 cuBLAS path. @@ -583,15 +583,15 @@ pub struct GpuDqnTrainer { // Scratch buffers for inter-layer gradient propagation during cuBLAS backward. // Sized for the widest layer at each point in the network. /// Accumulated d_h_s2 from value head + all 3 branch heads: [B, SH2] - bw_d_h_s2: CudaSlice, + bw_d_h_s2: CudaSlice, /// Upstream gradient for shared layer 1: [B, SH1] - bw_d_h_s1: CudaSlice, + bw_d_h_s1: CudaSlice, /// Upstream gradient for value FC: [B, VH] - bw_d_h_v: CudaSlice, + bw_d_h_v: CudaSlice, /// Branch FC upstream gradients (one per branch): [B, AH] - bw_d_h_b0: CudaSlice, - bw_d_h_b1: CudaSlice, - bw_d_h_b2: CudaSlice, + bw_d_h_b0: CudaSlice, + bw_d_h_b1: CudaSlice, + bw_d_h_b2: CudaSlice, // ── Expected Q-value kernel (ad-hoc validation, not captured in CUDA Graph) ─ /// Converts C51 value+advantage logits → expected Q-values (validation path). @@ -601,16 +601,16 @@ pub struct GpuDqnTrainer { /// GPU reduction: q_out_buf → 5 scalars [avg_max_q, q_min, q_max, q_mean, q_var] q_stats_kernel: CudaFunction, /// GPU buffer for Q-value statistics [5 floats] - q_stats_buf: CudaSlice, + q_stats_buf: CudaSlice, // ── CQL (Conservative Q-Learning) penalty kernel ───────────────── /// Computes CQL logit gradients: dCQL/d_value_logits and dCQL/d_adv_logits. /// Only used when `config.use_cql == true && config.cql_alpha > 0`. cql_logit_grad_kernel: Option, /// CQL scratch: value logit gradients [B, NA] - cql_d_value_logits: CudaSlice, + cql_d_value_logits: CudaSlice, /// CQL scratch: advantage logit gradients [B, (B0+B1+B2)*NA] - cql_d_adv_logits: CudaSlice, + cql_d_adv_logits: CudaSlice, } impl Drop for GpuDqnTrainer { @@ -633,12 +633,12 @@ impl GpuDqnTrainer { /// Shape: `[B, SHARED_H2]` — the shared trunk output for online states. /// Valid after `train_step()` or `forward_loss()` has been called. /// Used by the IQN dual-head to read trunk activations without re-computation. - pub fn save_h_s2(&self) -> &CudaSlice { + pub fn save_h_s2(&self) -> &CudaSlice { &self.save_h_s2 } /// Backward gradient w.r.t. h_s2 (trunk activation). - pub fn bw_d_h_s2_buf(&self) -> &CudaSlice { + pub fn bw_d_h_s2_buf(&self) -> &CudaSlice { &self.bw_d_h_s2 } @@ -651,7 +651,7 @@ impl GpuDqnTrainer { /// /// Shape: `[B, STATE_DIM]` — contains the batch's states after `upload_batch_gpu()`. /// Used by IQL value network to read states without Candle tensor intermediaries. - pub fn states_buf(&self) -> &CudaSlice { + pub fn states_buf(&self) -> &CudaSlice { &self.states_buf } @@ -659,7 +659,7 @@ impl GpuDqnTrainer { /// /// Shape: `[B, STATE_DIM]` — contains the batch's next states after `train_step()`. /// Used by the IQN dual-head to compute target h_s2 via a trunk forward kernel. - pub fn next_states_buf(&self) -> &CudaSlice { + pub fn next_states_buf(&self) -> &CudaSlice { &self.next_states_buf } @@ -675,7 +675,7 @@ impl GpuDqnTrainer { /// /// Shape: `[B]` f32 — rewards from the batch. Valid after `train_step()` or /// `train_step_gpu()`. Used by IQN for Bellman target computation on GPU. - pub fn rewards_buf(&self) -> &CudaSlice { + pub fn rewards_buf(&self) -> &CudaSlice { &self.rewards_buf } @@ -683,7 +683,7 @@ impl GpuDqnTrainer { /// /// Shape: `[B]` f32 — done flags (0.0/1.0). Valid after `train_step()` or /// `train_step_gpu()`. Used by IQN for Bellman target masking on GPU. - pub fn dones_buf(&self) -> &CudaSlice { + pub fn dones_buf(&self) -> &CudaSlice { &self.dones_buf } @@ -705,13 +705,13 @@ impl GpuDqnTrainer { /// Mutable reference to grad_buf for test injection. #[cfg(test)] - pub fn grad_buf_mut(&mut self) -> &mut CudaSlice { + pub fn grad_buf_mut(&mut self) -> &mut CudaSlice { &mut self.grad_buf } /// Mutable reference to CQL gradient scratch for test injection. #[cfg(test)] - pub fn cql_grad_scratch_mut(&mut self) -> &mut CudaSlice { + pub fn cql_grad_scratch_mut(&mut self) -> &mut CudaSlice { &mut self.cql_grad_scratch } @@ -737,7 +737,7 @@ impl GpuDqnTrainer { /// 3. SAXPY: `grad_buf[trunk] += iqn_lambda * scratch[trunk]` pub fn apply_iqn_trunk_gradient( &mut self, - iqn_d_h_s2: &CudaSlice, + iqn_d_h_s2: &CudaSlice, _online_dueling: &mut DuelingWeightSet, ) -> Result<(), MLError> { let b = self.config.batch_size; @@ -2239,7 +2239,7 @@ impl GpuDqnTrainer { /// for short2 vectorized loads in the forward kernel. fn launch_segmented_bf16_convert( &self, - f32_buf: &CudaSlice, + f32_buf: &CudaSlice, bf16_buf: &CudaSlice, ) -> Result<(), MLError> { let param_sizes = compute_param_sizes(&self.config); @@ -2609,13 +2609,13 @@ impl GpuDqnTrainer { /// /// # Arguments /// * `indices` — CudaSlice [B] buffer indices from PER sampling - /// * `priorities` — CudaSlice [capacity] priority array in GpuReplayBuffer + /// * `priorities` — CudaSlice [capacity] priority array in GpuReplayBuffer /// * `alpha` — PER alpha exponent /// * `epsilon` — PER epsilon floor pub fn update_priorities_cuda( &mut self, indices: &CudaSlice, - priorities: &CudaSlice, + priorities: &CudaSlice, alpha: f32, epsilon: f32, ) -> Result<(), MLError> { @@ -2667,7 +2667,7 @@ impl GpuDqnTrainer { /// Shape: `[B]` f32 — per-sample TD errors from the last training step. /// Valid after `train_step_gpu()` or `train_step()` has been called. /// Used by the PER priority update kernel to avoid DtoH readback. - pub fn td_errors_buf(&self) -> &CudaSlice { + pub fn td_errors_buf(&self) -> &CudaSlice { &self.td_errors_buf } @@ -2675,7 +2675,7 @@ impl GpuDqnTrainer { /// /// Shape: `[B, num_atoms]`. Used by ensemble heads to get head-0 C51 logits /// for diversity loss computation without re-running the forward pass. - pub fn on_v_logits_buf(&self) -> &CudaSlice { + pub fn on_v_logits_buf(&self) -> &CudaSlice { &self.on_v_logits_buf } @@ -2683,7 +2683,7 @@ impl GpuDqnTrainer { /// /// Shape: `[B, VALUE_H]`. Safe to reuse between the CUDA Graph step and the next /// train_step call (not used during the post-graph ensemble forward phase). - pub fn tg_h_v_scratch_ptr(&self) -> &CudaSlice { + pub fn tg_h_v_scratch_ptr(&self) -> &CudaSlice { &self.tg_h_v_scratch } @@ -2721,7 +2721,7 @@ impl GpuDqnTrainer { /// Reference to the Q-value output buffer from the last `forward_only_q()` call. /// /// Shape: `[B, total_actions]`. Only valid after `forward_only_q()`. - pub fn q_out_buf(&self) -> &CudaSlice { + pub fn q_out_buf(&self) -> &CudaSlice { &self.q_out_buf } @@ -2735,13 +2735,13 @@ impl GpuDqnTrainer { /// It reuses the existing `on_v_logits_buf` and `on_b_logits_buf` scratch buffers, /// and writes expected Q-values into `q_out_buf`. /// - /// Input: `states` — `CudaSlice` of shape `[batch_size, state_dim]` + /// Input: `states` — `CudaSlice` of shape `[batch_size, state_dim]` /// Output: reference to `q_out_buf` — `[batch_size, total_actions(11)]` pub fn compute_q_values( &self, - states: &CudaSlice, + states: &CudaSlice, batch_size: usize, - ) -> Result<&CudaSlice, MLError> { + ) -> Result<&CudaSlice, MLError> { if batch_size > self.config.batch_size { return Err(MLError::ModelError(format!( "compute_q_values: batch_size {batch_size} exceeds trainer batch_size {}", @@ -2814,7 +2814,7 @@ impl GpuDqnTrainer { /// Only 20 bytes (5 f32) are read back to host. pub fn compute_q_stats( &mut self, - states: &CudaSlice, + states: &CudaSlice, batch_size: usize, ) -> Result { self.compute_q_values(states, batch_size)?; @@ -3258,7 +3258,7 @@ impl GpuDqnTrainer { let b = self.config.batch_size; let sd = self.config.state_dim; - // GpuBatch stores GpuTensor (CudaSlice) for all fields. + // GpuBatch stores GpuTensor (CudaSlice) for all fields. // States may be BF16-typed GpuTensors whose underlying data is u16-reinterpreted f32. // Check shape to determine if we have BF16 data that needs conversion. let states_data = gpu_batch.states.data(); @@ -3274,10 +3274,10 @@ impl GpuDqnTrainer { } else { // BF16 states -- reinterpret f32 storage as u16 and convert let states_u16: &CudaSlice = unsafe { - &*(states_data as *const CudaSlice as *const CudaSlice) + &*(states_data as *const CudaSlice as *const CudaSlice) }; let next_states_u16: &CudaSlice = unsafe { - &*(next_states_data as *const CudaSlice as *const CudaSlice) + &*(next_states_data as *const CudaSlice as *const CudaSlice) }; dtod_from_bf16(states_u16, &self.bf16_states_buf, b * sd, &self.stream, "states")?; dtod_from_bf16(next_states_u16, &self.bf16_next_states_buf, b * sd, &self.stream, "next_states")?; @@ -3293,7 +3293,7 @@ impl GpuDqnTrainer { // Actions: GpuTensor f32 -> I32 buf. Reinterpret f32 as u32 then DtoD as i32. let actions_f32 = gpu_batch.actions.data(); let actions_u32: &CudaSlice = unsafe { - &*(actions_f32 as *const CudaSlice as *const CudaSlice) + &*(actions_f32 as *const CudaSlice as *const CudaSlice) }; dtod_from_u32_to_i32(actions_u32, &self.actions_buf, b, &self.stream, "actions")?; @@ -3642,8 +3642,8 @@ impl GpuDqnTrainer { /// Writes gradient outputs to the provided destination buffers. fn launch_mse_grad_inner( &self, - d_value_dst: &CudaSlice, - d_adv_dst: &CudaSlice, + d_value_dst: &CudaSlice, + d_adv_dst: &CudaSlice, ) -> Result<(), MLError> { let b = self.config.batch_size; let na = self.config.num_atoms; @@ -3863,7 +3863,7 @@ impl GpuDqnTrainer { let dst_base = raw_device_ptr(&self.params_buf, &self.stream); // Ordered to match GOFF_* layout - let slices: [&CudaSlice; 20] = [ + let slices: [&CudaSlice; 20] = [ &online_d.w_s1, &online_d.b_s1, &online_d.w_s2, &online_d.b_s2, &online_d.w_v1, &online_d.b_v1, @@ -3918,7 +3918,7 @@ impl GpuDqnTrainer { let src_base = raw_device_ptr(&self.params_buf, &self.stream); // Must match GOFF_* order exactly - let slices: [&CudaSlice; 20] = [ + let slices: [&CudaSlice; 20] = [ &online_d.w_s1, &online_d.b_s1, &online_d.w_s2, &online_d.b_s2, &online_d.w_v1, &online_d.b_v1, @@ -3958,7 +3958,7 @@ impl GpuDqnTrainer { let sizes = compute_param_sizes(&self.config); let dst_base = raw_device_ptr(&self.target_params_buf, &self.stream); - let slices: [&CudaSlice; 20] = [ + let slices: [&CudaSlice; 20] = [ &target_d.w_s1, &target_d.b_s1, &target_d.w_s2, &target_d.b_s2, &target_d.w_v1, &target_d.b_v1, @@ -4001,7 +4001,7 @@ impl GpuDqnTrainer { let sizes = compute_param_sizes(&self.config); let src_base = raw_device_ptr(&self.target_params_buf, &self.stream); - let slices: [&CudaSlice; 20] = [ + let slices: [&CudaSlice; 20] = [ &target_d.w_s1, &target_d.b_s1, &target_d.w_s2, &target_d.b_s2, &target_d.w_v1, &target_d.b_v1, @@ -4321,7 +4321,7 @@ pub(crate) fn query_max_shmem_bytes() -> usize { /// /// Same pattern as `raw_device_ptr` in `gpu_weights.rs` — wraps the SyncOnDrop /// guard in ManuallyDrop to skip read event recording (safe on same stream). -pub(crate) fn raw_device_ptr(slice: &CudaSlice, stream: &CudaStream) -> u64 { +pub(crate) fn raw_device_ptr(slice: &CudaSlice, stream: &CudaStream) -> u64 { let (ptr, guard) = slice.device_ptr(stream); let _no_drop = std::mem::ManuallyDrop::new(guard); ptr @@ -4371,10 +4371,10 @@ pub(crate) fn dtod_copy( // ── Candle tensor → CudaSlice DtoD copy helpers ──────────────────────────── -/// DtoD copy from a CudaSlice to a pre-allocated CudaSlice. +/// DtoD copy from a CudaSlice to a pre-allocated CudaSlice. fn dtod_from_slice_f32( - src: &CudaSlice, - dst: &CudaSlice, + src: &CudaSlice, + dst: &CudaSlice, num_elements: usize, stream: &Arc, ctx: &str, @@ -4405,8 +4405,8 @@ fn alloc_f32( stream: &Arc, n: usize, name: &str, -) -> Result, MLError> { - stream.alloc_zeros::(n).map_err(|e| { +) -> Result, MLError> { + stream.alloc_zeros::(n).map_err(|e| { MLError::ModelError(format!("GpuDqnTrainer alloc {name} ({n} f32): {e}")) }) } @@ -4469,14 +4469,14 @@ pub(crate) fn compile_ensemble_kernels( Ok((aggregate, diversity)) } -/// Launch the bf16_to_f32 CUDA kernel: converts BF16 CudaSlice → F32 CudaSlice. +/// Launch the bf16_to_f32 CUDA kernel: converts BF16 CudaSlice → F32 CudaSlice. /// /// Grid: ceil(n/256), Block: 256. Kernel signature: `bf16_to_f32_kernel(src, dst, n)`. /// Zero Candle tensor involvement — pure cudarc kernel launch. pub(crate) fn launch_bf16_to_f32( kernel: &CudaFunction, src: &CudaSlice, - dst: &CudaSlice, + dst: &CudaSlice, n: usize, stream: &Arc, ctx: &str, @@ -4488,7 +4488,7 @@ pub(crate) fn launch_bf16_to_f32( shared_mem_bytes: 0, }; - // Safety: src is CudaSlice holding BF16 bit patterns, dst is CudaSlice, + // Safety: src is CudaSlice holding BF16 bit patterns, dst is CudaSlice, // both are valid GPU allocations on the same context. Kernel reads src[0..n] and // writes dst[0..n]. unsafe { diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index e78fa6422..f3261c118 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -329,15 +329,15 @@ impl ExperienceCollectorConfig { #[allow(missing_debug_implementations)] pub struct GpuExperienceBatch { /// States `[total * state_dim]` on GPU (f32, row-major) - pub states: CudaSlice, + pub states: CudaSlice, /// Next-states `[total * state_dim]` on GPU (f32, episode-aware shift) - pub next_states: CudaSlice, + pub next_states: CudaSlice, /// Action indices `[total]` on GPU (i32, values 0..4) pub actions: CudaSlice, /// Rewards `[total]` on GPU (f32) - pub rewards: CudaSlice, + pub rewards: CudaSlice, /// Done flags `[total]` on GPU (i32, 0 or 1) - pub dones: CudaSlice, + pub dones: CudaSlice, /// Episode index per transition `[total]` on GPU (i32). /// /// `episode_ids[i] = i / timesteps_per_episode`. @@ -393,7 +393,7 @@ pub struct GpuExperienceCollector { // ── Flat F32 online weight buffer for cuBLAS ──────────────────── /// Flattened online network parameters [total_params] — single DtoD copy for sync. - online_params_flat: CudaSlice, + online_params_flat: CudaSlice, /// Total number of F32 parameters in the flat buffer. total_params: usize, /// Per-tensor sizes for computing weight pointers. @@ -417,26 +417,26 @@ pub struct GpuExperienceCollector { // ── Per-timestep batch buffers (reused each timestep) ─────────── /// Assembled state batch for cuBLAS: [N, state_dim]. - batch_states: CudaSlice, + batch_states: CudaSlice, /// Selected actions per episode: [N]. batch_actions: CudaSlice, /// Q-gap conviction per episode [N]: output by action_select, input to env_step. - q_gaps_buf: CudaSlice, + q_gaps_buf: CudaSlice, /// Kelly optimal fraction (per-epoch scalar from CPU). /// Current timestep counter per episode: [N]. current_timesteps: CudaSlice, /// Q-values output from expected Q or direct logits: [N, total_branch_actions]. - q_values: CudaSlice, + q_values: CudaSlice, // ── cuBLAS activation scratch (reused each timestep) ──────────── - exp_h_s1: CudaSlice, // [N, shared_h1] - exp_h_s2: CudaSlice, // [N, shared_h2] - exp_h_v: CudaSlice, // [N, value_h] - exp_h_b0: CudaSlice, // [N, adv_h] - exp_h_b1: CudaSlice, // [N, adv_h] - exp_h_b2: CudaSlice, // [N, adv_h] - exp_v_logits: CudaSlice, // [N, num_atoms] - exp_b_logits: CudaSlice, // [N, total_branch_atoms] + exp_h_s1: CudaSlice, // [N, shared_h1] + exp_h_s2: CudaSlice, // [N, shared_h2] + exp_h_v: CudaSlice, // [N, value_h] + exp_h_b0: CudaSlice, // [N, adv_h] + exp_h_b1: CudaSlice, // [N, adv_h] + exp_h_b2: CudaSlice, // [N, adv_h] + exp_v_logits: CudaSlice, // [N, num_atoms] + exp_b_logits: CudaSlice, // [N, total_branch_atoms] // ── Legacy weight sets (for sync_online_weights/sync_target_weights compat) ── online_weights: DuelingWeightSet, @@ -448,26 +448,26 @@ pub struct GpuExperienceCollector { target_branching: BranchingWeightSet, /// GPU buffer for trade_stats_reduce output: [TRADE_STATS_FLOATS] = 6 floats. - trade_stats_buf: CudaSlice, + trade_stats_buf: CudaSlice, // Per-episode state buffers [alloc_episodes, ...] - portfolio_states: CudaSlice, // [alloc_episodes * PORTFOLIO_STRIDE] (20 floats per episode) + portfolio_states: CudaSlice, // [alloc_episodes * PORTFOLIO_STRIDE] (20 floats per episode) rng_states: CudaSlice, // [alloc_episodes] episode_starts_buf: CudaSlice,// [alloc_episodes] // Output buffers [alloc_episodes * alloc_timesteps, ...] - states_out: CudaSlice, // [alloc_episodes * alloc_timesteps * STATE_DIM] + states_out: CudaSlice, // [alloc_episodes * alloc_timesteps * STATE_DIM] actions_out: CudaSlice, // [alloc_episodes * alloc_timesteps] - rewards_out: CudaSlice, // [alloc_episodes * alloc_timesteps] - done_out: CudaSlice, // [alloc_episodes * alloc_timesteps] + rewards_out: CudaSlice, // [alloc_episodes * alloc_timesteps] + done_out: CudaSlice, // [alloc_episodes * alloc_timesteps] /// Raw per-bar portfolio returns (unshaped) for accurate Sharpe/MaxDD/Sortino. /// True fractional return: (equity_t - equity_{t-1}) / equity_{t-1}. - raw_returns_out: CudaSlice, // [alloc_episodes * alloc_timesteps] + raw_returns_out: CudaSlice, // [alloc_episodes * alloc_timesteps] /// Running reward statistics for cross-epoch normalization (Welford's online algorithm) // Persistent epoch state — survives across kernel launches. - epoch_state: CudaSlice, // [8] + epoch_state: CudaSlice, // [8] /// Bitfield: bit 0 = reset portfolio, bit 1 = reset DSR, bit 2 = reset vol EMA reset_flags: u32, @@ -479,7 +479,7 @@ pub struct GpuExperienceCollector { expert_ratio: f32, // OFI features pre-uploaded to GPU [total_bars * OFI_DIM] (None when OFI disabled) - ofi_gpu: Option>, + ofi_gpu: Option>, ofi_dim: usize, // GPU-resident curiosity forward model trainer (None if curiosity disabled) @@ -608,7 +608,7 @@ impl GpuExperienceCollector { // Allocate flat online parameter buffer let online_params_flat = stream - .alloc_zeros::(total_params) + .alloc_zeros::(total_params) .map_err(|e| MLError::ModelError(format!("alloc online_params_flat: {e}")))?; info!( @@ -675,7 +675,7 @@ impl GpuExperienceCollector { // ── Step 6: Allocate per-episode buffers ──────────────────────── // Portfolio states for experience kernels: [N, PORTFOLIO_STRIDE=20] let mut portfolio_states = stream - .alloc_zeros::(alloc_episodes * PORTFOLIO_STRIDE) + .alloc_zeros::(alloc_episodes * PORTFOLIO_STRIDE) .map_err(|e| MLError::ModelError(format!("alloc portfolio_states: {e}")))?; // Initialize portfolio state (20 floats per episode) @@ -729,34 +729,34 @@ impl GpuExperienceCollector { // ── Step 7: Allocate output buffers ───────────────────────────── let total_output = alloc_episodes * alloc_timesteps; let states_out = stream - .alloc_zeros::(total_output * state_dim) + .alloc_zeros::(total_output * state_dim) .map_err(|e| MLError::ModelError(format!("alloc states_out: {e}")))?; let actions_out = stream .alloc_zeros::(total_output) .map_err(|e| MLError::ModelError(format!("alloc actions_out: {e}")))?; let rewards_out = stream - .alloc_zeros::(total_output) + .alloc_zeros::(total_output) .map_err(|e| MLError::ModelError(format!("alloc rewards_out: {e}")))?; let done_out = stream - .alloc_zeros::(total_output) + .alloc_zeros::(total_output) .map_err(|e| MLError::ModelError(format!("alloc done_out: {e}")))?; let raw_returns_out = stream - .alloc_zeros::(total_output) + .alloc_zeros::(total_output) .map_err(|e| MLError::ModelError(format!("alloc raw_returns_out: {e}")))?; let trade_stats_buf = stream - .alloc_zeros::(TRADE_STATS_FLOATS) + .alloc_zeros::(TRADE_STATS_FLOATS) .map_err(|e| MLError::ModelError(format!("alloc trade_stats_buf: {e}")))?; // ── Step 8: Allocate per-timestep batch buffers ───────────────── let n = alloc_episodes; let batch_states = stream - .alloc_zeros::(n * state_dim) + .alloc_zeros::(n * state_dim) .map_err(|e| MLError::ModelError(format!("alloc batch_states: {e}")))?; let batch_actions = stream .alloc_zeros::(n) .map_err(|e| MLError::ModelError(format!("alloc batch_actions: {e}")))?; let q_gaps_buf = stream - .alloc_zeros::(n) + .alloc_zeros::(n) .map_err(|e| MLError::ModelError(format!("alloc q_gaps_buf: {e}")))?; let current_timesteps = stream .alloc_zeros::(n) @@ -764,33 +764,33 @@ impl GpuExperienceCollector { let total_branch_atoms = (branch_sizes[0] + branch_sizes[1] + branch_sizes[2]) * num_atoms; let q_values = stream - .alloc_zeros::(n * total_branch_actions) + .alloc_zeros::(n * total_branch_actions) .map_err(|e| MLError::ModelError(format!("alloc q_values: {e}")))?; // cuBLAS activation scratch buffers let exp_h_s1 = stream - .alloc_zeros::(n * shared_h1) + .alloc_zeros::(n * shared_h1) .map_err(|e| MLError::ModelError(format!("alloc exp_h_s1: {e}")))?; let exp_h_s2 = stream - .alloc_zeros::(n * shared_h2) + .alloc_zeros::(n * shared_h2) .map_err(|e| MLError::ModelError(format!("alloc exp_h_s2: {e}")))?; let exp_h_v = stream - .alloc_zeros::(n * value_h) + .alloc_zeros::(n * value_h) .map_err(|e| MLError::ModelError(format!("alloc exp_h_v: {e}")))?; let exp_h_b0 = stream - .alloc_zeros::(n * adv_h) + .alloc_zeros::(n * adv_h) .map_err(|e| MLError::ModelError(format!("alloc exp_h_b0: {e}")))?; let exp_h_b1 = stream - .alloc_zeros::(n * adv_h) + .alloc_zeros::(n * adv_h) .map_err(|e| MLError::ModelError(format!("alloc exp_h_b1: {e}")))?; let exp_h_b2 = stream - .alloc_zeros::(n * adv_h) + .alloc_zeros::(n * adv_h) .map_err(|e| MLError::ModelError(format!("alloc exp_h_b2: {e}")))?; let exp_v_logits = stream - .alloc_zeros::(n * num_atoms) + .alloc_zeros::(n * num_atoms) .map_err(|e| MLError::ModelError(format!("alloc exp_v_logits: {e}")))?; let exp_b_logits = stream - .alloc_zeros::(n * total_branch_atoms) + .alloc_zeros::(n * total_branch_atoms) .map_err(|e| MLError::ModelError(format!("alloc exp_b_logits: {e}")))?; // ── Step 9: Initialize curiosity trainer ──────────────────────── @@ -818,7 +818,7 @@ impl GpuExperienceCollector { let pinned_dones = unsafe { PinnedHostBuf::::new(total_output)? }; // OFI placeholder - let ofi_placeholder = stream.alloc_zeros::(1) + let ofi_placeholder = stream.alloc_zeros::(1) .map_err(|e| MLError::ModelError(format!("alloc OFI placeholder: {e}"))) .ok(); @@ -1034,8 +1034,8 @@ impl GpuExperienceCollector { /// these buffers to ops that run on a different CUDA stream. pub fn collect_experiences_gpu( &mut self, - market_features_buf: &CudaSlice, - targets_buf: &CudaSlice, + market_features_buf: &CudaSlice, + targets_buf: &CudaSlice, episode_starts: &[i32], config: &ExperienceCollectorConfig, ) -> Result { @@ -1135,8 +1135,8 @@ impl GpuExperienceCollector { /// 5. `experience_env_step` — portfolio simulation, reward, done, advance timestep fn launch_timestep_loop( &mut self, - market_features_buf: &CudaSlice, - targets_buf: &CudaSlice, + market_features_buf: &CudaSlice, + targets_buf: &CudaSlice, episode_starts: &[i32], config: &ExperienceCollectorConfig, ) -> Result<(usize, usize), MLError> { @@ -1237,7 +1237,7 @@ impl GpuExperienceCollector { // ── 3. Convert logits → expected Q-values ─────────────────── // When num_atoms == 1, b_logits is already [N, 11] Q-values. // When num_atoms > 1, we need softmax + expectation over z-support. - let q_source: &CudaSlice = if self.num_atoms > 1 { + let q_source: &CudaSlice = if self.num_atoms > 1 { let na = self.num_atoms as i32; let v_min_f = self.v_min; let v_max_f = self.v_max; @@ -1429,7 +1429,7 @@ impl GpuExperienceCollector { /// /// Use this when the caller already has a contiguous F32 parameter buffer /// (e.g. from `GpuDqnTrainer`). Avoids the per-tensor extraction overhead. - pub fn sync_weights_flat(&mut self, params_buf: &CudaSlice) -> Result<(), MLError> { + pub fn sync_weights_flat(&mut self, params_buf: &CudaSlice) -> Result<(), MLError> { let num_bytes = self.total_params * std::mem::size_of::(); let src = raw_device_ptr(params_buf, &self.stream); let dst = raw_device_ptr(&self.online_params_flat, &self.stream); @@ -1545,7 +1545,7 @@ impl GpuExperienceCollector { } /// Get a reference to the epoch state GPU buffer (for monitoring kernel). - pub fn epoch_state_gpu(&self) -> &CudaSlice { + pub fn epoch_state_gpu(&self) -> &CudaSlice { &self.epoch_state } @@ -1564,7 +1564,7 @@ impl GpuExperienceCollector { } /// Get a reference to the rewards output GPU buffer. - pub fn rewards_gpu(&self) -> &CudaSlice { + pub fn rewards_gpu(&self) -> &CudaSlice { &self.rewards_out } @@ -1613,7 +1613,7 @@ impl GpuExperienceCollector { if ofi_flat.is_empty() { return Ok(()); } - let mut gpu_buf = self.stream.alloc_zeros::(ofi_flat.len()) + let mut gpu_buf = self.stream.alloc_zeros::(ofi_flat.len()) .map_err(|e| MLError::ModelError(format!("alloc OFI GPU buffer: {e}")))?; self.stream.memcpy_htod(ofi_flat, &mut gpu_buf) .map_err(|e| MLError::ModelError(format!("upload OFI features to GPU: {e}")))?; @@ -1765,14 +1765,14 @@ fn fill_episode_ids_gpu( // Pure cudarc DtoD helpers // --------------------------------------------------------------------------- -/// Allocate a fresh `CudaSlice` and DtoD-copy `n_elems` floats from `src`. +/// Allocate a fresh `CudaSlice` and DtoD-copy `n_elems` floats from `src`. fn dtod_clone_f32( stream: &Arc, - src: &CudaSlice, + src: &CudaSlice, n_elems: usize, label: &str, -) -> Result, MLError> { - let mut dst = stream.alloc_zeros::(n_elems) +) -> Result, MLError> { + let mut dst = stream.alloc_zeros::(n_elems) .map_err(|e| MLError::ModelError(format!("alloc {label} f32[{n_elems}]: {e}")))?; let src_view = src.slice(..n_elems); stream.memcpy_dtod(&src_view, &mut dst) @@ -1801,12 +1801,12 @@ fn dtod_clone_i32( /// When `t + shift >= timesteps`, duplicates the last available state. fn build_next_states_dtod( stream: &Arc, - states: &CudaSlice, + states: &CudaSlice, n_episodes: usize, timesteps: usize, sd: usize, shift: usize, -) -> Result, MLError> { +) -> Result, MLError> { let total = n_episodes * timesteps; let total_elems = total * sd; let shift = shift.max(1).min(timesteps); // clamp shift to valid range @@ -1815,7 +1815,7 @@ fn build_next_states_dtod( return dtod_clone_f32(stream, states, total_elems, "next_states_trivial"); } - let mut dst = stream.alloc_zeros::(total_elems) + let mut dst = stream.alloc_zeros::(total_elems) .map_err(|e| MLError::ModelError(format!("alloc next_states f32[{total_elems}]: {e}")))?; let shifted_count = timesteps - shift; @@ -1844,7 +1844,7 @@ fn build_next_states_dtod( } /// Extract raw F32 device pointer (read-only). -fn raw_f32_ptr(slice: &CudaSlice, stream: &Arc) -> u64 { +fn raw_f32_ptr(slice: &CudaSlice, stream: &Arc) -> u64 { let (ptr, guard) = slice.device_ptr(stream); let _no_drop = ManuallyDrop::new(guard); ptr diff --git a/crates/ml/src/cuda_pipeline/gpu_her.rs b/crates/ml/src/cuda_pipeline/gpu_her.rs index 063b188a0..df4277d33 100644 --- a/crates/ml/src/cuda_pipeline/gpu_her.rs +++ b/crates/ml/src/cuda_pipeline/gpu_her.rs @@ -101,15 +101,15 @@ impl GpuHerConfig { #[allow(missing_debug_implementations)] pub struct HerBatch { /// Relabeled states `[her_batch_size, state_dim]` -- f32 on device. - pub states: CudaSlice, + pub states: CudaSlice, /// Relabeled next_states `[her_batch_size, state_dim]` -- f32 on device. - pub next_states: CudaSlice, + pub next_states: CudaSlice, /// Actions (copied unchanged) `[her_batch_size]` -- i32 on device. pub actions: CudaSlice, /// Recomputed rewards `[her_batch_size]` -- f32 on device. - pub rewards: CudaSlice, + pub rewards: CudaSlice, /// Dones (copied unchanged) `[her_batch_size]` -- f32 on device. - pub dones: CudaSlice, + pub dones: CudaSlice, /// Number of relabeled samples in this batch. pub batch_size: usize, } @@ -129,11 +129,11 @@ pub struct GpuHer { pub config: GpuHerConfig, // Staging buffers for relabeled experiences (pre-allocated at max HER batch size) - out_states: CudaSlice, - out_next_states: CudaSlice, + out_states: CudaSlice, + out_next_states: CudaSlice, out_actions: CudaSlice, - out_rewards: CudaSlice, - out_dones: CudaSlice, + out_rewards: CudaSlice, + out_dones: CudaSlice, // Index buffers (uploaded each step) source_indices: CudaSlice, @@ -248,11 +248,11 @@ impl GpuHer { /// * `her_batch_size` - Number of experiences to relabel pub fn relabel_batch( &mut self, - buffer_states: &CudaSlice, - buffer_next_states: &CudaSlice, + buffer_states: &CudaSlice, + buffer_next_states: &CudaSlice, buffer_actions: &CudaSlice, - buffer_rewards: &CudaSlice, - buffer_dones: &CudaSlice, + buffer_rewards: &CudaSlice, + buffer_dones: &CudaSlice, source_idx_host: &[i32], donor_idx_host: &[i32], buffer_size: usize, @@ -512,8 +512,8 @@ fn alloc_f32( stream: &Arc, n: usize, name: &str, -) -> Result, MLError> { - stream.alloc_zeros::(n).map_err(|e| { +) -> Result, MLError> { + stream.alloc_zeros::(n).map_err(|e| { MLError::ModelError(format!("GpuHer alloc {name} ({n} f32): {e}")) }) } @@ -528,12 +528,12 @@ fn alloc_i32( }) } -/// Extract raw CUDA device pointer from a CudaSlice. +/// Extract raw CUDA device pointer from a CudaSlice. /// /// Same pattern as `raw_device_ptr` in `gpu_dqn_trainer.rs` -- wraps the /// `SyncOnDrop` guard in `ManuallyDrop` to skip read event recording (safe /// on same stream, same context). -fn raw_ptr_f32(slice: &CudaSlice, stream: &CudaStream) -> u64 { +fn raw_ptr_f32(slice: &CudaSlice, stream: &CudaStream) -> u64 { let (ptr, guard) = slice.device_ptr(stream); let _no_drop = std::mem::ManuallyDrop::new(guard); ptr @@ -546,12 +546,12 @@ fn raw_ptr_i32(slice: &CudaSlice, stream: &CudaStream) -> u64 { ptr } -/// Clone a sub-range of a CudaSlice via async D2D memcpy. +/// Clone a sub-range of a CudaSlice via async D2D memcpy. fn slice_clone_f32( stream: &Arc, - src: &CudaSlice, + src: &CudaSlice, n: usize, -) -> Result, MLError> { +) -> Result, MLError> { let dst = alloc_f32(stream, n, "her_clone_f32")?; let num_bytes = n * std::mem::size_of::(); let src_ptr = raw_ptr_f32(src, stream); diff --git a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs index 2ee1cfdd5..dfadc8e73 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs @@ -111,24 +111,24 @@ pub struct GpuIqlTrainer { forward_kernel: CudaFunction, // ── V network parameters (flat f32 on GPU) ───────────────────── - params_buf: CudaSlice, + params_buf: CudaSlice, // ── Adam optimizer state ──────────────────────────────────────── - m_buf: CudaSlice, // first moment - v_buf: CudaSlice, // second moment - grad_buf: CudaSlice, // gradient accumulator - grad_norm_buf: CudaSlice, // [1] gradient L2 norm + m_buf: CudaSlice, // first moment + v_buf: CudaSlice, // second moment + grad_buf: CudaSlice, // gradient accumulator + grad_norm_buf: CudaSlice, // [1] gradient L2 norm // ── Activation save buffers (forward -> backward) ─────────────── - save_pre1: CudaSlice, // [B, H] pre-activation layer 1 - save_pre2: CudaSlice, // [B, H] pre-activation layer 2 - save_h1: CudaSlice, // [B, H] post-activation layer 1 - save_h2: CudaSlice, // [B, H] post-activation layer 2 + save_pre1: CudaSlice, // [B, H] pre-activation layer 1 + save_pre2: CudaSlice, // [B, H] pre-activation layer 2 + save_h1: CudaSlice, // [B, H] post-activation layer 1 + save_h2: CudaSlice, // [B, H] post-activation layer 2 // ── Output buffers ────────────────────────────────────────────── - v_out_buf: CudaSlice, // [B] V(s) predictions - loss_buf: CudaSlice, // [B] per-sample loss - total_loss_buf: CudaSlice, // [1] batch-mean loss + v_out_buf: CudaSlice, // [B] V(s) predictions + loss_buf: CudaSlice, // [B] per-sample loss + total_loss_buf: CudaSlice, // [1] batch-mean loss // ── Training state ────────────────────────────────────────────── adam_step: i32, @@ -234,8 +234,8 @@ impl GpuIqlTrainer { /// tensor involvement. pub fn train_value_step( &mut self, - states_f32: &CudaSlice, - q_values_f32: &CudaSlice, + states_f32: &CudaSlice, + q_values_f32: &CudaSlice, ) -> Result { let b = self.config.batch_size; @@ -433,7 +433,7 @@ fn compile_iql_kernels( fn init_xavier_weights( stream: &Arc, config: &GpuIqlConfig, -) -> Result, MLError> { +) -> Result, MLError> { use rand::Rng; let total = config.total_params(); @@ -486,8 +486,8 @@ fn alloc_f32( stream: &Arc, n: usize, name: &str, -) -> Result, MLError> { - stream.alloc_zeros::(n).map_err(|e| { +) -> Result, MLError> { + stream.alloc_zeros::(n).map_err(|e| { MLError::ModelError(format!("GpuIql alloc {name} ({n} f32): {e}")) }) } diff --git a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs index 9c5550574..e90294730 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs @@ -150,42 +150,42 @@ pub struct GpuIqnHead { sample_taus_kernel: CudaFunction, // ── IQN online parameters (flat f32 on GPU) ────────────────────── - online_params: CudaSlice, + online_params: CudaSlice, // ── IQN target parameters (for Bellman target) ─────────────────── - target_params: CudaSlice, + target_params: CudaSlice, // ── Adam optimizer state ───────────────────────────────────────── - m_buf: CudaSlice, - v_buf: CudaSlice, - grad_buf: CudaSlice, - grad_norm_buf: CudaSlice, + m_buf: CudaSlice, + v_buf: CudaSlice, + grad_buf: CudaSlice, + grad_norm_buf: CudaSlice, /// IQN gradient w.r.t. shared trunk h_s2 [B, hidden_dim]. /// Accumulated during backward, added to C51's trunk gradient. - d_h_s2_buf: CudaSlice, + d_h_s2_buf: CudaSlice, // ── Per-step buffers ───────────────────────────────────────────── /// Pre-sampled τ values for online network [B, N] - online_taus: CudaSlice, + online_taus: CudaSlice, /// Pre-sampled τ values for target network [B, N] - target_taus: CudaSlice, + target_taus: CudaSlice, /// Branch actions decoded from flat actions [B, 3] branch_actions: CudaSlice, /// Target h_s2 computed from next_states + target trunk weights [B, H] - target_h_s2: CudaSlice, + target_h_s2: CudaSlice, /// Persistent rewards buffer [B] — avoids per-step temp allocs - rewards_buf: CudaSlice, + rewards_buf: CudaSlice, /// Persistent dones buffer [B] — avoids per-step temp allocs - dones_buf: CudaSlice, + dones_buf: CudaSlice, // ── Activation save buffers (forward → backward) ───────────────── - save_embed: CudaSlice, // [B, N, H] - save_combined: CudaSlice, // [B, N, H] - save_q_online: CudaSlice, // [B, N, TOTAL_BRANCH_ACTIONS] - save_q_target: CudaSlice, // [B, N, TOTAL_BRANCH_ACTIONS] + save_embed: CudaSlice, // [B, N, H] + save_combined: CudaSlice, // [B, N, H] + save_q_online: CudaSlice, // [B, N, TOTAL_BRANCH_ACTIONS] + save_q_target: CudaSlice, // [B, N, TOTAL_BRANCH_ACTIONS] // ── Loss buffers ───────────────────────────────────────────────── - per_sample_loss: CudaSlice, // [B] - total_loss: CudaSlice, // [1] + per_sample_loss: CudaSlice, // [B] + total_loss: CudaSlice, // [1] // ── Training state ─────────────────────────────────────────────── adam_step: i32, @@ -330,12 +330,12 @@ impl GpuIqnHead { /// Returns the batch-mean quantile Huber loss. pub fn train_iqn_step_gpu( &mut self, - online_h_s2: &CudaSlice, - next_states_buf: &CudaSlice, + online_h_s2: &CudaSlice, + next_states_buf: &CudaSlice, target_dueling: &DuelingWeightSet, dqn_actions_buf: &CudaSlice, - dqn_rewards_buf: &CudaSlice, - dqn_dones_buf: &CudaSlice, + dqn_rewards_buf: &CudaSlice, + dqn_dones_buf: &CudaSlice, ) -> Result { let b = self.config.batch_size; let shared_h1_i32 = self.config.shared_h1 as i32; @@ -389,8 +389,8 @@ impl GpuIqnHead { /// on GPU (by either `train_iqn_step_gpu` or the legacy HtoD path). fn execute_training_pipeline( &mut self, - online_h_s2: &CudaSlice, - next_states_buf: &CudaSlice, + online_h_s2: &CudaSlice, + next_states_buf: &CudaSlice, target_dueling: &DuelingWeightSet, ) -> Result { let b = self.config.batch_size; @@ -616,14 +616,14 @@ impl GpuIqnHead { } /// Per-sample IQN loss buffer reference (for PER weighting). - pub fn per_sample_loss(&self) -> &CudaSlice { + pub fn per_sample_loss(&self) -> &CudaSlice { &self.per_sample_loss } /// IQN gradient w.r.t. shared trunk h_s2 [B, hidden_dim]. /// After backward, this contains dL_iqn/d(h_s2) — add to C51's trunk gradient /// so the trunk receives bounded IQN Huber gradients alongside C51. - pub fn d_h_s2(&self) -> &CudaSlice { + pub fn d_h_s2(&self) -> &CudaSlice { &self.d_h_s2_buf } @@ -649,11 +649,11 @@ impl GpuIqnHead { /// The scale is applied to target_position in the env_step kernel. pub fn compute_cvar_scales( &mut self, - h_s2: &CudaSlice, + h_s2: &CudaSlice, actions: &CudaSlice, batch_size: usize, alpha: f32, - ) -> Result, MLError> { + ) -> Result, MLError> { let b = batch_size; let n = self.config.num_quantiles; let tba = self.config.total_branch_actions(); @@ -709,7 +709,7 @@ impl GpuIqnHead { // 3. Compute CVaR on GPU — ZERO CPU readback. // Inline CUDA kernel: one thread per sample, insertion-sort 32 quantiles // in registers, compute CVaR as mean of lowest alpha_count values. - let mut scales_buf = self.stream.alloc_zeros::(b) + let mut scales_buf = self.stream.alloc_zeros::(b) .map_err(|e| MLError::ModelError(format!("IQN CVaR alloc scales: {e}")))?; let alpha_f = alpha; let n_i32 = n as i32; @@ -874,7 +874,7 @@ fn compile_iqn_kernels( fn init_iqn_xavier_weights( stream: &Arc, config: &GpuIqnConfig, -) -> Result, MLError> { +) -> Result, MLError> { use rand::Rng; let total = config.total_params(); @@ -927,8 +927,8 @@ fn alloc_f32( stream: &Arc, n: usize, name: &str, -) -> Result, MLError> { - stream.alloc_zeros::(n).map_err(|e| { +) -> Result, MLError> { + stream.alloc_zeros::(n).map_err(|e| { MLError::ModelError(format!("GpuIqn alloc {name} ({n} f32): {e}")) }) } @@ -936,9 +936,9 @@ fn alloc_f32( /// Clone a CudaSlice by copying device → host → device. fn clone_cuda_slice( stream: &Arc, - src: &CudaSlice, + src: &CudaSlice, n: usize, -) -> Result, MLError> { +) -> Result, MLError> { let mut host = vec![0.0_f32; n]; stream.memcpy_dtoh(src, &mut host) // gpu-exit: weight clone for target network .map_err(|e| MLError::ModelError(format!("IQN clone D→H: {e}")))?; diff --git a/crates/ml/src/cuda_pipeline/gpu_monitoring.rs b/crates/ml/src/cuda_pipeline/gpu_monitoring.rs index ae23dce4f..bc3c1d547 100644 --- a/crates/ml/src/cuda_pipeline/gpu_monitoring.rs +++ b/crates/ml/src/cuda_pipeline/gpu_monitoring.rs @@ -29,7 +29,7 @@ pub struct MonitoringSummary { pub struct GpuMonitoringReducer { stream: Arc, kernel_func: CudaFunction, - summary_buf: CudaSlice, // [12] + summary_buf: CudaSlice, // [12] } impl GpuMonitoringReducer { @@ -40,7 +40,7 @@ impl GpuMonitoringReducer { .map_err(|e| MLError::ModelError(format!("monitoring module load: {e}")))?; let kernel_func = module.load_function("monitoring_reduce") .map_err(|e| MLError::ModelError(format!("monitoring_reduce load: {e}")))?; - let summary_buf = stream.alloc_zeros::(16) // 5 stats + 9 action counts + 1 total + 1 pad + let summary_buf = stream.alloc_zeros::(16) // 5 stats + 9 action counts + 1 total + 1 pad .map_err(|e| MLError::ModelError(format!("monitoring summary alloc: {e}")))?; Ok(Self { stream: Arc::clone(stream), kernel_func, summary_buf }) @@ -50,7 +50,7 @@ impl GpuMonitoringReducer { /// Does NOT synchronize — caller must sync before reading result. pub fn reduce( &mut self, - rewards: &CudaSlice, + rewards: &CudaSlice, actions: &CudaSlice, n: usize, ) -> Result<(), MLError> { diff --git a/crates/ml/src/cuda_pipeline/gpu_portfolio.rs b/crates/ml/src/cuda_pipeline/gpu_portfolio.rs index aeae6f9ec..b94395704 100644 --- a/crates/ml/src/cuda_pipeline/gpu_portfolio.rs +++ b/crates/ml/src/cuda_pipeline/gpu_portfolio.rs @@ -36,9 +36,9 @@ pub struct GpuPortfolioSimulator { // GPU buffers (persistent across batches) actions_buf: CudaSlice, - portfolio_state_buf: CudaSlice, - portfolio_out_buf: CudaSlice, - rewards_out_buf: CudaSlice, + portfolio_state_buf: CudaSlice, + portfolio_out_buf: CudaSlice, + rewards_out_buf: CudaSlice, done_out_buf: CudaSlice, // Config @@ -54,9 +54,9 @@ pub struct GpuPortfolioSimulator { #[allow(missing_debug_implementations)] pub struct GpuPortfolioSimResult { /// Normalized portfolio features [batch_size * 3] on GPU - pub portfolio_features: CudaSlice, + pub portfolio_features: CudaSlice, /// Raw PnL rewards [batch_size] on GPU - pub rewards: CudaSlice, + pub rewards: CudaSlice, /// Episode boundary flags [batch_size] on GPU (1=done, 0=continue) pub done_flags: CudaSlice, /// Number of bars processed @@ -98,13 +98,13 @@ impl GpuPortfolioSimulator { .alloc_zeros::(MAX_BATCH_SIZE) .map_err(|e| MLError::ModelError(format!("Failed to alloc actions buffer: {e}")))?; let mut portfolio_state_buf = stream - .alloc_zeros::(PORTFOLIO_STATE_SIZE) + .alloc_zeros::(PORTFOLIO_STATE_SIZE) .map_err(|e| MLError::ModelError(format!("Failed to alloc portfolio state: {e}")))?; let portfolio_out_buf = stream - .alloc_zeros::(MAX_BATCH_SIZE * 3) + .alloc_zeros::(MAX_BATCH_SIZE * 3) .map_err(|e| MLError::ModelError(format!("Failed to alloc portfolio output: {e}")))?; let rewards_out_buf = stream - .alloc_zeros::(MAX_BATCH_SIZE) + .alloc_zeros::(MAX_BATCH_SIZE) .map_err(|e| MLError::ModelError(format!("Failed to alloc rewards output: {e}")))?; let done_out_buf = stream .alloc_zeros::(MAX_BATCH_SIZE) @@ -150,16 +150,16 @@ impl GpuPortfolioSimulator { /// directly by the next GPU stage (state builder, forward pass, etc.). pub fn simulate_batch_gpu( &mut self, - targets_buf: &CudaSlice, + targets_buf: &CudaSlice, actions: &[i32], batch_start: usize, ) -> Result { let batch_size = actions.len(); if batch_size == 0 { return Ok(GpuPortfolioSimResult { - portfolio_features: self.stream.alloc_zeros::(0) + portfolio_features: self.stream.alloc_zeros::(0) .map_err(|e| MLError::ModelError(format!("empty alloc: {e}")))?, - rewards: self.stream.alloc_zeros::(0) + rewards: self.stream.alloc_zeros::(0) .map_err(|e| MLError::ModelError(format!("empty alloc: {e}")))?, done_flags: self.stream.alloc_zeros::(0) .map_err(|e| MLError::ModelError(format!("empty alloc: {e}")))?, @@ -228,11 +228,11 @@ impl GpuPortfolioSimulator { /// DtoD clone the first `count` f32 elements of `src`. #[allow(unsafe_code)] fn dtod_clone_range_f32( - src: &CudaSlice, + src: &CudaSlice, count: usize, stream: &Arc, - ) -> Result, MLError> { - let mut dst = stream.alloc_zeros::(count) + ) -> Result, MLError> { + let mut dst = stream.alloc_zeros::(count) .map_err(|e| MLError::ModelError(format!("dtod_clone alloc: {e}")))?; { let nbytes = count * std::mem::size_of::(); @@ -282,9 +282,9 @@ impl GpuPortfolioSimulator { /// Get reference to the GPU-resident portfolio state buffer (for diagnostics). /// - /// Returns the raw `CudaSlice` of shape `[PORTFOLIO_STATE_SIZE]` on GPU. + /// Returns the raw `CudaSlice` of shape `[PORTFOLIO_STATE_SIZE]` on GPU. /// Use `memcpy_dtoh` explicitly at checkpoint boundaries only. - pub fn portfolio_state_buf(&self) -> &CudaSlice { + pub fn portfolio_state_buf(&self) -> &CudaSlice { &self.portfolio_state_buf } } diff --git a/crates/ml/src/cuda_pipeline/gpu_ppo_collector.rs b/crates/ml/src/cuda_pipeline/gpu_ppo_collector.rs index 5b257f3d1..534179bd3 100644 --- a/crates/ml/src/cuda_pipeline/gpu_ppo_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_ppo_collector.rs @@ -118,15 +118,15 @@ impl PpoCollectorConfig { #[allow(missing_debug_implementations)] pub struct PpoExperienceBatch { /// Flattened state vectors `[N * L * STATE_DIM]` -- GPU-resident - pub states: CudaSlice, + pub states: CudaSlice, /// Selected action indices `[N * L]` -- GPU-resident pub actions: CudaSlice, /// Log-probabilities of selected actions `[N * L]` -- GPU-resident - pub log_probs: CudaSlice, + pub log_probs: CudaSlice, /// GAE advantages `[N * L]` -- GPU-resident - pub advantages: CudaSlice, + pub advantages: CudaSlice, /// Discounted returns (value targets) `[N * L]` -- GPU-resident - pub returns: CudaSlice, + pub returns: CudaSlice, /// Episode-done flags `[N * L]` -- GPU-resident pub done_flags: CudaSlice, /// Number of episodes in this batch @@ -232,24 +232,24 @@ pub struct GpuPpoExperienceCollector { curiosity_weights: CuriosityWeightSet, // Per-episode state buffers [N, ...] - portfolio_states: CudaSlice, // [MAX_EPISODES * PORTFOLIO_STATE_SIZE] - barrier_states: CudaSlice, // [MAX_EPISODES * BARRIER_STATE_SIZE] + portfolio_states: CudaSlice, // [MAX_EPISODES * PORTFOLIO_STATE_SIZE] + barrier_states: CudaSlice, // [MAX_EPISODES * BARRIER_STATE_SIZE] diversity_windows: CudaSlice, // [MAX_EPISODES * DIVERSITY_WINDOW] diversity_metas: CudaSlice, // [MAX_EPISODES * 2] - barrier_config: CudaSlice, // [3] + barrier_config: CudaSlice, // [3] rng_states: CudaSlice, // [MAX_EPISODES] episode_starts_buf: CudaSlice,// [MAX_EPISODES] // Output buffers [MAX_EPISODES * MAX_TIMESTEPS, ...] - states_out: CudaSlice, // [MAX_EPISODES * MAX_TIMESTEPS * STATE_DIM] + states_out: CudaSlice, // [MAX_EPISODES * MAX_TIMESTEPS * STATE_DIM] actions_out: CudaSlice, // [MAX_EPISODES * MAX_TIMESTEPS] - log_probs_out: CudaSlice, // [MAX_EPISODES * MAX_TIMESTEPS] - advantages_out: CudaSlice, // [MAX_EPISODES * MAX_TIMESTEPS] - returns_out: CudaSlice, // [MAX_EPISODES * MAX_TIMESTEPS] + log_probs_out: CudaSlice, // [MAX_EPISODES * MAX_TIMESTEPS] + advantages_out: CudaSlice, // [MAX_EPISODES * MAX_TIMESTEPS] + returns_out: CudaSlice, // [MAX_EPISODES * MAX_TIMESTEPS] dones_out: CudaSlice, // [MAX_EPISODES * MAX_TIMESTEPS] // Consolidated readback staging: D2D gather then single DtoH per dtype - readback_f32_staging: CudaSlice, + readback_f32_staging: CudaSlice, readback_f32_host: Vec, readback_i32_staging: CudaSlice, readback_i32_host: Vec, @@ -306,12 +306,12 @@ impl GpuPpoExperienceCollector { // ---- Step 3: Allocate per-episode buffers ---- let portfolio_states = stream - .alloc_zeros::(MAX_EPISODES * PORTFOLIO_STATE_SIZE) + .alloc_zeros::(MAX_EPISODES * PORTFOLIO_STATE_SIZE) .map_err(|e| { MLError::ModelError(format!("Failed to alloc portfolio_states: {e}")) })?; let barrier_states = stream - .alloc_zeros::(MAX_EPISODES * BARRIER_STATE_SIZE) + .alloc_zeros::(MAX_EPISODES * BARRIER_STATE_SIZE) .map_err(|e| { MLError::ModelError(format!("Failed to alloc barrier_states: {e}")) })?; @@ -326,7 +326,7 @@ impl GpuPpoExperienceCollector { MLError::ModelError(format!("Failed to alloc diversity_metas: {e}")) })?; let barrier_config = stream - .alloc_zeros::(3) + .alloc_zeros::(3) .map_err(|e| { MLError::ModelError(format!("Failed to alloc barrier_config: {e}")) })?; @@ -375,7 +375,7 @@ impl GpuPpoExperienceCollector { // ---- Step 6: Allocate output buffers ---- let total_output = MAX_EPISODES * MAX_TIMESTEPS; let states_out = stream - .alloc_zeros::(total_output * STATE_DIM) + .alloc_zeros::(total_output * STATE_DIM) .map_err(|e| { MLError::ModelError(format!("Failed to alloc states_out: {e}")) })?; @@ -385,17 +385,17 @@ impl GpuPpoExperienceCollector { MLError::ModelError(format!("Failed to alloc actions_out: {e}")) })?; let log_probs_out = stream - .alloc_zeros::(total_output) + .alloc_zeros::(total_output) .map_err(|e| { MLError::ModelError(format!("Failed to alloc log_probs_out: {e}")) })?; let advantages_out = stream - .alloc_zeros::(total_output) + .alloc_zeros::(total_output) .map_err(|e| { MLError::ModelError(format!("Failed to alloc advantages_out: {e}")) })?; let returns_out = stream - .alloc_zeros::(total_output) + .alloc_zeros::(total_output) .map_err(|e| { MLError::ModelError(format!("Failed to alloc returns_out: {e}")) })?; @@ -409,7 +409,7 @@ impl GpuPpoExperienceCollector { // f32 staging: states (N*L*STATE_DIM) + log_probs (N*L) + advantages (N*L) + returns (N*L) let f32_staging_len = total_output * STATE_DIM + 3 * total_output; let readback_f32_staging = stream - .alloc_zeros::(f32_staging_len) + .alloc_zeros::(f32_staging_len) .map_err(|e| { MLError::ModelError(format!("Failed to alloc readback_f32_staging: {e}")) })?; @@ -486,8 +486,8 @@ impl GpuPpoExperienceCollector { /// * `config` - Kernel configuration parameters pub fn collect_experiences( &mut self, - market_features_buf: &CudaSlice, - targets_buf: &CudaSlice, + market_features_buf: &CudaSlice, + targets_buf: &CudaSlice, episode_starts: &[i32], config: &PpoCollectorConfig, ) -> Result { @@ -668,13 +668,13 @@ impl GpuPpoExperienceCollector { /// D2D copy `count` f32 elements from `src` into a fresh CudaSlice (GPU-to-GPU, no CPU). fn d2d_copy_f32( &self, - src: &CudaSlice, + src: &CudaSlice, count: usize, label: &str, - ) -> Result, MLError> { + ) -> Result, MLError> { let mut dst = self .stream - .alloc_zeros::(count) + .alloc_zeros::(count) .map_err(|e| MLError::ModelError(format!("alloc {label}: {e}")))?; let src_view = src.slice(..count); let nbytes = count * std::mem::size_of::(); diff --git a/crates/ml/src/cuda_pipeline/gpu_statistics.rs b/crates/ml/src/cuda_pipeline/gpu_statistics.rs index f4d02e7b2..f96ad59d9 100644 --- a/crates/ml/src/cuda_pipeline/gpu_statistics.rs +++ b/crates/ml/src/cuda_pipeline/gpu_statistics.rs @@ -35,7 +35,7 @@ pub struct BatchStatistics { #[allow(missing_debug_implementations)] pub struct GpuStatistics { kernel_func: CudaFunction, - output_buf: CudaSlice, + output_buf: CudaSlice, stream: Arc, } @@ -52,7 +52,7 @@ impl GpuStatistics { MLError::ModelError(format!("batch_statistics function load: {e}")) })?; - let output_buf = stream.alloc_zeros::(10).map_err(|e| { + let output_buf = stream.alloc_zeros::(10).map_err(|e| { MLError::ModelError(format!("alloc output_buf: {e}")) })?; @@ -65,9 +65,9 @@ impl GpuStatistics { /// drawdown are computed on the host from raw sums for correctness. pub fn compute( &mut self, - rewards: &CudaSlice, - portfolio_values: &CudaSlice, - positions: &CudaSlice, + rewards: &CudaSlice, + portfolio_values: &CudaSlice, + positions: &CudaSlice, n: usize, ) -> Result { let _nvtx = NvtxRange::new("gpu_statistics"); diff --git a/crates/ml/src/cuda_pipeline/gpu_training_guard.rs b/crates/ml/src/cuda_pipeline/gpu_training_guard.rs index 785bf67f6..c4946a9fd 100644 --- a/crates/ml/src/cuda_pipeline/gpu_training_guard.rs +++ b/crates/ml/src/cuda_pipeline/gpu_training_guard.rs @@ -191,7 +191,7 @@ pub struct GpuTrainingGuard { /// Accumulator buffer (3 floats: loss_sum, grad_norm_sum, step_count). /// Stays in device memory -- kernel does read-modify-write. /// Single `memcpy_dtoh` at boundary only (not per-step). - acc_buf: CudaSlice, + acc_buf: CudaSlice, /// Welford running-mean accumulator for Q-value estimation (CPU-side, no GPU tensor). q_count: usize, @@ -225,7 +225,7 @@ impl GpuTrainingGuard { .map_err(|e| MLError::ModelError(format!("qvalue_divergence_check load: {e}")))?; // Allocate device accumulator buffer (stays on-device, read at epoch boundary) - let acc_buf = stream.alloc_zeros::(3).map_err(|e| { + let acc_buf = stream.alloc_zeros::(3).map_err(|e| { MLError::ModelError(format!("alloc acc_buf: {e}")) })?; @@ -264,15 +264,15 @@ impl GpuTrainingGuard { /// Run the guard check and accumulate kernels for one training step. /// - /// `loss_gpu` and `grad_norm_gpu` must be F32 `CudaSlice` scalars (1 element each). + /// `loss_gpu` and `grad_norm_gpu` must be F32 `CudaSlice` scalars (1 element each). /// /// Returns the safety flags and scalar values from the *previous* step /// (one-step delay due to double-buffering). On the very first call, /// returns safe defaults (no halts, zero loss/grad_norm). pub fn check_and_accumulate( &mut self, - loss_gpu: &CudaSlice, - grad_norm_gpu: &CudaSlice, + loss_gpu: &CudaSlice, + grad_norm_gpu: &CudaSlice, clip_threshold: f32, collapse_threshold: f32, warmup: bool, @@ -365,10 +365,10 @@ impl GpuTrainingGuard { /// Compute Q-value statistics over a full batch via `qvalue_stats_reduce`. /// - /// `q_values` must be a contiguous F32 `CudaSlice` with `batch_size * num_actions` elements. + /// `q_values` must be a contiguous F32 `CudaSlice` with `batch_size * num_actions` elements. pub fn qvalue_stats( &mut self, - q_values: &CudaSlice, + q_values: &CudaSlice, batch_size: usize, num_actions: usize, ) -> Result { @@ -410,10 +410,10 @@ impl GpuTrainingGuard { /// Check a single sample for Q-value divergence via `qvalue_divergence_check`. /// - /// `q_values` must be a contiguous F32 `CudaSlice` with `num_actions` elements. + /// `q_values` must be a contiguous F32 `CudaSlice` with `num_actions` elements. pub fn qvalue_divergence( &mut self, - q_values: &CudaSlice, + q_values: &CudaSlice, num_actions: usize, threshold: f32, ) -> Result { diff --git a/crates/ml/src/cuda_pipeline/gpu_walk_forward.rs b/crates/ml/src/cuda_pipeline/gpu_walk_forward.rs index 9477ed4de..26ca40c32 100644 --- a/crates/ml/src/cuda_pipeline/gpu_walk_forward.rs +++ b/crates/ml/src/cuda_pipeline/gpu_walk_forward.rs @@ -158,9 +158,9 @@ impl GpuWalkForwardConfig { /// GPU-resident walk-forward data store. /// -/// Holds the entire dataset on GPU as cudarc `CudaSlice` buffers. +/// Holds the entire dataset on GPU as cudarc `CudaSlice` buffers. /// Folds are expressed as index ranges into these buffers — no re-upload needed. -// CudaSlice doesn't implement Debug, so manual impl. +// CudaSlice doesn't implement Debug, so manual impl. impl std::fmt::Debug for GpuWalkForwardData { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("GpuWalkForwardData") @@ -176,11 +176,11 @@ impl std::fmt::Debug for GpuWalkForwardData { pub struct GpuWalkForwardData { /// Market features [total_bars * FEATURE_DIM] (FEATURE_DIM = 42). - pub features: CudaSlice, + pub features: CudaSlice, /// Target prices [total_bars * TARGET_DIM] (TARGET_DIM = 4). - pub targets: CudaSlice, + pub targets: CudaSlice, /// OFI features [total_bars * 8] (optional, from MBP-10). - pub ofi: Option>, + pub ofi: Option>, /// Total number of bars in the dataset. pub total_bars: usize, /// Feature dimension (42). diff --git a/crates/ml/src/cuda_pipeline/gpu_weights.rs b/crates/ml/src/cuda_pipeline/gpu_weights.rs index f6b01c60e..b193b3b9b 100644 --- a/crates/ml/src/cuda_pipeline/gpu_weights.rs +++ b/crates/ml/src/cuda_pipeline/gpu_weights.rs @@ -1,7 +1,7 @@ //! Weight extraction and GPU upload for DQN experience collection kernel. //! //! Extracts neural network weights from `GpuVarStore` objects and copies -//! them as flat `CudaSlice` GPU buffers for the CUDA experience +//! them as flat `CudaSlice` GPU buffers for the CUDA experience //! collection kernels (`experience_kernels.cu` + cuBLAS). //! //! Two weight sets are managed: @@ -11,7 +11,7 @@ //! Weights are stored row-major `[out_features, in_features]`. The CUDA //! kernel reads them in the same layout. //! -//! All weights are `CudaSlice` on GPU. Extraction and sync use +//! All weights are `CudaSlice` on GPU. Extraction and sync use //! device-to-device copy (zero CPU roundtrip). // CUDA FFI module — memcpy_dtod_async requires unsafe by design. #![allow(unsafe_code)] @@ -108,20 +108,20 @@ const PPO_CRITIC_WEIGHT_NAMES: [&str; 12] = [ #[allow(missing_debug_implementations)] // CudaSlice does not implement Debug pub struct DuelingWeightSet { // Shared layers - pub w_s1: CudaSlice, // shared_0.weight [SHARED_H1, STATE_DIM] - pub b_s1: CudaSlice, // shared_0.bias [SHARED_H1] - pub w_s2: CudaSlice, // shared_1.weight [SHARED_H2, SHARED_H1] - pub b_s2: CudaSlice, // shared_1.bias [SHARED_H2] + pub w_s1: CudaSlice, // shared_0.weight [SHARED_H1, STATE_DIM] + pub b_s1: CudaSlice, // shared_0.bias [SHARED_H1] + pub w_s2: CudaSlice, // shared_1.weight [SHARED_H2, SHARED_H1] + pub b_s2: CudaSlice, // shared_1.bias [SHARED_H2] // Value head (distributional: NUM_ATOMS outputs, not 1) - pub w_v1: CudaSlice, // value_fc.weight [VALUE_H, SHARED_H2] - pub b_v1: CudaSlice, // value_fc.bias [VALUE_H] - pub w_v2: CudaSlice, // value_out.weight [NUM_ATOMS, VALUE_H] - pub b_v2: CudaSlice, // value_out.bias [NUM_ATOMS] + pub w_v1: CudaSlice, // value_fc.weight [VALUE_H, SHARED_H2] + pub b_v1: CudaSlice, // value_fc.bias [VALUE_H] + pub w_v2: CudaSlice, // value_out.weight [NUM_ATOMS, VALUE_H] + pub b_v2: CudaSlice, // value_out.bias [NUM_ATOMS] // Exposure head / branch 0 (distributional: BRANCH_0_SIZE*NUM_ATOMS outputs) - pub w_a1: CudaSlice, // branch_0_fc.weight [ADV_H, SHARED_H2] - pub b_a1: CudaSlice, // branch_0_fc.bias [ADV_H] - pub w_a2: CudaSlice, // branch_0_out.weight [BRANCH_0_SIZE*NUM_ATOMS, ADV_H] - pub b_a2: CudaSlice, // branch_0_out.bias [BRANCH_0_SIZE*NUM_ATOMS] + pub w_a1: CudaSlice, // branch_0_fc.weight [ADV_H, SHARED_H2] + pub b_a1: CudaSlice, // branch_0_fc.bias [ADV_H] + pub w_a2: CudaSlice, // branch_0_out.weight [BRANCH_0_SIZE*NUM_ATOMS, ADV_H] + pub b_a2: CudaSlice, // branch_0_out.bias [BRANCH_0_SIZE*NUM_ATOMS] } /// GPU buffers holding the 2 extra advantage head weight sets for Branching DQN. @@ -140,15 +140,15 @@ pub struct DuelingWeightSet { #[allow(missing_debug_implementations)] // CudaSlice does not implement Debug pub struct BranchingWeightSet { // Order head (branch 1) - pub w_bo1: CudaSlice, // branch_1_fc.weight [ADV_H, SHARED_H2] - pub b_bo1: CudaSlice, // branch_1_fc.bias [ADV_H] - pub w_bo2: CudaSlice, // branch_1_out.weight [BRANCH_1_SIZE*NUM_ATOMS, ADV_H] - pub b_bo2: CudaSlice, // branch_1_out.bias [BRANCH_1_SIZE*NUM_ATOMS] + pub w_bo1: CudaSlice, // branch_1_fc.weight [ADV_H, SHARED_H2] + pub b_bo1: CudaSlice, // branch_1_fc.bias [ADV_H] + pub w_bo2: CudaSlice, // branch_1_out.weight [BRANCH_1_SIZE*NUM_ATOMS, ADV_H] + pub b_bo2: CudaSlice, // branch_1_out.bias [BRANCH_1_SIZE*NUM_ATOMS] // Urgency head (branch 2) - pub w_bu1: CudaSlice, // branch_2_fc.weight [ADV_H, SHARED_H2] - pub b_bu1: CudaSlice, // branch_2_fc.bias [ADV_H] - pub w_bu2: CudaSlice, // branch_2_out.weight [BRANCH_2_SIZE*NUM_ATOMS, ADV_H] - pub b_bu2: CudaSlice, // branch_2_out.bias [BRANCH_2_SIZE*NUM_ATOMS] + pub w_bu1: CudaSlice, // branch_2_fc.weight [ADV_H, SHARED_H2] + pub b_bu1: CudaSlice, // branch_2_fc.bias [ADV_H] + pub w_bu2: CudaSlice, // branch_2_out.weight [BRANCH_2_SIZE*NUM_ATOMS, ADV_H] + pub b_bu2: CudaSlice, // branch_2_out.bias [BRANCH_2_SIZE*NUM_ATOMS] } /// GPU buffers holding all 4 weight tensors of the Curiosity Forward Model. @@ -166,10 +166,10 @@ const CUR_OUTPUT_DIM: usize = CUR_MARKET_DIM; // 42 #[allow(missing_debug_implementations)] // CudaSlice does not implement Debug pub struct CuriosityWeightSet { - pub w1: CudaSlice, // fc1.weight [CUR_HIDDEN, CUR_INPUT] = [128, 45] - pub b1: CudaSlice, // fc1.bias [CUR_HIDDEN] = [128] - pub w2: CudaSlice, // fc2.weight [CUR_OUTPUT, CUR_HIDDEN] = [42, 128] - pub b2: CudaSlice, // fc2.bias [CUR_OUTPUT] = [42] + pub w1: CudaSlice, // fc1.weight [CUR_HIDDEN, CUR_INPUT] = [128, 45] + pub b1: CudaSlice, // fc1.bias [CUR_HIDDEN] = [128] + pub w2: CudaSlice, // fc2.weight [CUR_OUTPUT, CUR_HIDDEN] = [42, 128] + pub b2: CudaSlice, // fc2.bias [CUR_OUTPUT] = [42] } impl CuriosityWeightSet { @@ -181,16 +181,16 @@ impl CuriosityWeightSet { /// `curiosity_scale` is set to 0.0 in the config). pub fn zeros(stream: &Arc) -> Result { Ok(Self { - w1: stream.alloc_zeros::(CUR_HIDDEN_DIM * CUR_INPUT_DIM).map_err(|e| { + w1: stream.alloc_zeros::(CUR_HIDDEN_DIM * CUR_INPUT_DIM).map_err(|e| { MLError::ModelError(format!("Failed to alloc curiosity w1 zeros: {e}")) })?, - b1: stream.alloc_zeros::(CUR_HIDDEN_DIM).map_err(|e| { + b1: stream.alloc_zeros::(CUR_HIDDEN_DIM).map_err(|e| { MLError::ModelError(format!("Failed to alloc curiosity b1 zeros: {e}")) })?, - w2: stream.alloc_zeros::(CUR_OUTPUT_DIM * CUR_HIDDEN_DIM).map_err(|e| { + w2: stream.alloc_zeros::(CUR_OUTPUT_DIM * CUR_HIDDEN_DIM).map_err(|e| { MLError::ModelError(format!("Failed to alloc curiosity w2 zeros: {e}")) })?, - b2: stream.alloc_zeros::(CUR_OUTPUT_DIM).map_err(|e| { + b2: stream.alloc_zeros::(CUR_OUTPUT_DIM).map_err(|e| { MLError::ModelError(format!("Failed to alloc curiosity b2 zeros: {e}")) })?, }) @@ -205,12 +205,12 @@ impl CuriosityWeightSet { /// - Output: `[45, 64]`, `[45]` #[allow(missing_debug_implementations)] // CudaSlice does not implement Debug pub struct PpoActorWeightSet { - pub pw1: CudaSlice, // policy_layer_0.weight [128, 54] - pub pb1: CudaSlice, // policy_layer_0.bias [128] - pub pw2: CudaSlice, // policy_layer_1.weight [64, 128] - pub pb2: CudaSlice, // policy_layer_1.bias [64] - pub pw3: CudaSlice, // policy_output.weight [45, 64] - pub pb3: CudaSlice, // policy_output.bias [45] + pub pw1: CudaSlice, // policy_layer_0.weight [128, 54] + pub pb1: CudaSlice, // policy_layer_0.bias [128] + pub pw2: CudaSlice, // policy_layer_1.weight [64, 128] + pub pb2: CudaSlice, // policy_layer_1.bias [64] + pub pw3: CudaSlice, // policy_output.weight [45, 64] + pub pb3: CudaSlice, // policy_output.bias [45] } /// GPU buffers holding all 12 weight tensors of the PPO Critic (ValueNetwork). @@ -224,18 +224,18 @@ pub struct PpoActorWeightSet { /// - Output: `[1, 64]`, `[1]` #[allow(missing_debug_implementations)] // CudaSlice does not implement Debug pub struct PpoCriticWeightSet { - pub vw1: CudaSlice, // value_layer_0.weight [512, 54] - pub vb1: CudaSlice, // value_layer_0.bias [512] - pub vw2: CudaSlice, // value_layer_1.weight [384, 512] - pub vb2: CudaSlice, // value_layer_1.bias [384] - pub vw3: CudaSlice, // value_layer_2.weight [256, 384] - pub vb3: CudaSlice, // value_layer_2.bias [256] - pub vw4: CudaSlice, // value_layer_3.weight [128, 256] - pub vb4: CudaSlice, // value_layer_3.bias [128] - pub vw5: CudaSlice, // value_layer_4.weight [64, 128] - pub vb5: CudaSlice, // value_layer_4.bias [64] - pub vw6: CudaSlice, // value_output.weight [1, 64] - pub vb6: CudaSlice, // value_output.bias [1] + pub vw1: CudaSlice, // value_layer_0.weight [512, 54] + pub vb1: CudaSlice, // value_layer_0.bias [512] + pub vw2: CudaSlice, // value_layer_1.weight [384, 512] + pub vb2: CudaSlice, // value_layer_1.bias [384] + pub vw3: CudaSlice, // value_layer_2.weight [256, 384] + pub vb3: CudaSlice, // value_layer_2.bias [256] + pub vw4: CudaSlice, // value_layer_3.weight [128, 256] + pub vb4: CudaSlice, // value_layer_3.bias [128] + pub vw5: CudaSlice, // value_layer_4.weight [64, 128] + pub vb5: CudaSlice, // value_layer_4.bias [64] + pub vw6: CudaSlice, // value_output.weight [1, 64] + pub vb6: CudaSlice, // value_output.bias [1] } /// RMSNorm gamma weight names for distributional dueling network (4 total). @@ -255,10 +255,10 @@ const RMSNORM_WEIGHT_NAMES: [&str; 4] = [ /// - advantage_rmsnorm:[128] (after advantage FC) #[allow(missing_debug_implementations)] pub struct RmsNormWeightSet { - pub gamma_s0: CudaSlice, // shared_rmsnorm_0.weight [shared_h1] - pub gamma_s1: CudaSlice, // shared_rmsnorm_1.weight [shared_h2] - pub gamma_v: CudaSlice, // value_rmsnorm.weight [value_h] - pub gamma_a: CudaSlice, // advantage_rmsnorm.weight [adv_h] + pub gamma_s0: CudaSlice, // shared_rmsnorm_0.weight [shared_h1] + pub gamma_s1: CudaSlice, // shared_rmsnorm_1.weight [shared_h2] + pub gamma_v: CudaSlice, // value_rmsnorm.weight [value_h] + pub gamma_a: CudaSlice, // advantage_rmsnorm.weight [adv_h] } impl RmsNormWeightSet { @@ -271,10 +271,10 @@ impl RmsNormWeightSet { network_dims: (usize, usize, usize, usize), ) -> Result { let (shared_h1, shared_h2, value_h, adv_h) = network_dims; - let alloc_ones = |size: usize| -> Result, MLError> { + let alloc_ones = |size: usize| -> Result, MLError> { let data = vec![1.0_f32; size]; let mut buf = stream - .alloc_zeros::(size) + .alloc_zeros::(size) .map_err(|e| MLError::ModelError(format!("Alloc RMSNorm ones: {e}")))?; stream .memcpy_htod(&data, &mut buf) @@ -372,7 +372,7 @@ unsafe impl DeviceRepr for KernelWeightPack {} /// - The CudaSlice is owned by the caller and outlives the kernel launch /// - The kernel launches on the same stream, so CUDA ordering is guaranteed /// - `SyncOnDrop::Record` only records events for cross-stream sync (not needed) -fn raw_device_ptr(slice: &CudaSlice, stream: &CudaStream) -> u64 { +fn raw_device_ptr(slice: &CudaSlice, stream: &CudaStream) -> u64 { let (ptr, guard) = slice.device_ptr(stream); // Keep guard alive without running its Drop (skip read event recording). // Named binding prevents `let_underscore_must_use`; ManuallyDrop prevents drop. @@ -509,7 +509,7 @@ pub struct CuriosityWeightSetBf16 { /// Allocate a BF16 buffer matching an F32 buffer's element count. fn alloc_bf16_mirror( - f32_buf: &CudaSlice, + f32_buf: &CudaSlice, stream: &Arc, ) -> Result, MLError> { let n = f32_buf.len(); @@ -525,7 +525,7 @@ fn alloc_bf16_mirror( /// /// `func`: pre-compiled `CudaFunction` for `f32_to_bf16_kernel`. pub fn convert_f32_to_bf16( - src: &CudaSlice, + src: &CudaSlice, dst: &mut CudaSlice, func: &cudarc::driver::CudaFunction, stream: &Arc, @@ -789,14 +789,14 @@ fn extract_one( vars: &GpuVarStore, name: &str, stream: &Arc, -) -> Result, MLError> { +) -> Result, MLError> { let param = vars .get(name) .ok_or_else(|| MLError::ModelError(format!("Missing weight: {name}")))?; let n_elems = param.data.len(); let mut buf = stream - .alloc_zeros::(n_elems) + .alloc_zeros::(n_elems) .map_err(|e| MLError::ModelError(format!("Alloc {name}: {e}")))?; stream.memcpy_dtod(¶m.data, &mut buf) @@ -810,7 +810,7 @@ fn extract_one( fn sync_one( vars: &GpuVarStore, name: &str, - buf: &mut CudaSlice, + buf: &mut CudaSlice, stream: &Arc, ) -> Result<(), MLError> { let param = vars @@ -828,7 +828,7 @@ fn sync_one( // --------------------------------------------------------------------------- /// Extract all 12 Dueling Q-Network weight tensors from a `GpuVarStore` and -/// upload them to GPU as flat `CudaSlice` buffers. +/// upload them to GPU as flat `CudaSlice` buffers. /// /// The `GpuVarStore` must contain tensor keys matching `DuelingQNetwork`: /// `shared_0.weight`, `shared_0.bias`, `shared_1.weight`, `shared_1.bias`, @@ -840,7 +840,7 @@ pub fn extract_dueling_weights( stream: &Arc, ) -> Result { - let extract = |name: &str| -> Result, MLError> { + let extract = |name: &str| -> Result, MLError> { extract_one(vars, name, stream) }; @@ -887,7 +887,7 @@ pub fn extract_dueling_weights( } /// Extract all 8 Branching DQN extra head weight tensors from a `GpuVarStore` -/// and upload them to GPU as flat `CudaSlice` buffers. +/// and upload them to GPU as flat `CudaSlice` buffers. /// /// The `GpuVarStore` must contain tensor keys matching `BranchingDuelingQNetwork`: /// `branch_1_fc.weight`, `branch_1_fc.bias`, `branch_1_out.weight`, `branch_1_out.bias`, @@ -897,7 +897,7 @@ pub fn extract_branching_weights( stream: &Arc, ) -> Result { - let extract = |name: &str| -> Result, MLError> { + let extract = |name: &str| -> Result, MLError> { extract_one(vars, name, stream) }; @@ -949,7 +949,7 @@ pub fn extract_dueling_weights_branching( stream: &Arc, ) -> Result { - let extract = |name: &str| -> Result, MLError> { + let extract = |name: &str| -> Result, MLError> { extract_one(vars, name, stream) }; @@ -1004,7 +1004,7 @@ pub fn sync_dueling_weights_branching( } /// Extract all 4 Curiosity Forward Model weight tensors from a `GpuVarStore` -/// and upload them to GPU as flat `CudaSlice` buffers. +/// and upload them to GPU as flat `CudaSlice` buffers. /// /// The `GpuVarStore` must contain tensor keys matching `ForwardDynamicsModel`: /// `fc1.weight`, `fc1.bias`, `fc2.weight`, `fc2.bias`. @@ -1013,7 +1013,7 @@ pub fn extract_curiosity_weights( stream: &Arc, ) -> Result { - let extract = |name: &str| -> Result, MLError> { + let extract = |name: &str| -> Result, MLError> { extract_one(vars, name, stream) }; @@ -1098,7 +1098,7 @@ pub fn sync_branching_weights( fn reverse_sync_one( vars: &GpuVarStore, name: &str, - src: &CudaSlice, + src: &CudaSlice, stream: &Arc, ) -> Result<(), MLError> { let param = vars @@ -1244,7 +1244,7 @@ pub fn sync_curiosity_weights( } /// Extract all 6 PPO Actor (PolicyNetwork) weight tensors from a `GpuVarStore` -/// and upload them to GPU as flat `CudaSlice` buffers. +/// and upload them to GPU as flat `CudaSlice` buffers. /// /// The `GpuVarStore` must contain tensor keys matching `PolicyNetwork`: /// `policy_layer_0.weight`, `policy_layer_0.bias`, `policy_layer_1.weight`, @@ -1254,7 +1254,7 @@ pub fn extract_ppo_actor_weights( stream: &Arc, ) -> Result { - let extract = |name: &str| -> Result, MLError> { + let extract = |name: &str| -> Result, MLError> { extract_one(vars, name, stream) }; @@ -1289,7 +1289,7 @@ pub fn extract_ppo_actor_weights( } /// Extract all 12 PPO Critic (ValueNetwork) weight tensors from a `GpuVarStore` -/// and upload them to GPU as flat `CudaSlice` buffers. +/// and upload them to GPU as flat `CudaSlice` buffers. /// /// The `GpuVarStore` must contain tensor keys matching `ValueNetwork`: /// `value_layer_0.weight`, `value_layer_0.bias`, `value_layer_1.weight`, @@ -1301,7 +1301,7 @@ pub fn extract_ppo_critic_weights( stream: &Arc, ) -> Result { - let extract = |name: &str| -> Result, MLError> { + let extract = |name: &str| -> Result, MLError> { extract_one(vars, name, stream) }; diff --git a/crates/ml/src/cuda_pipeline/her_relabel_kernel.cu b/crates/ml/src/cuda_pipeline/her_relabel_kernel.cu index e10ec4c65..6c546df8f 100644 --- a/crates/ml/src/cuda_pipeline/her_relabel_kernel.cu +++ b/crates/ml/src/cuda_pipeline/her_relabel_kernel.cu @@ -27,19 +27,19 @@ extern "C" __global__ void her_relabel_kernel( - const float* __restrict__ states, /* [capacity, state_dim] */ - const float* __restrict__ next_states, /* [capacity, state_dim] */ + const __nv_bfloat16* __restrict__ states, /* [capacity, state_dim] */ + const __nv_bfloat16* __restrict__ next_states, /* [capacity, state_dim] */ const int* __restrict__ actions, /* [capacity] */ - const float* __restrict__ rewards, /* [capacity] */ - const float* __restrict__ dones, /* [capacity] */ + const __nv_bfloat16* __restrict__ rewards, /* [capacity] */ + const __nv_bfloat16* __restrict__ dones, /* [capacity] */ const int* __restrict__ source_indices,/* [her_batch_size] */ const int* __restrict__ donor_indices, /* [her_batch_size] */ int buffer_size, - float* __restrict__ out_states, /* [her_batch_size, state_dim] */ - float* __restrict__ out_next_states, /* [her_batch_size, state_dim] */ + __nv_bfloat16* __restrict__ out_states, /* [her_batch_size, state_dim] */ + __nv_bfloat16* __restrict__ out_next_states, /* [her_batch_size, state_dim] */ int* __restrict__ out_actions, /* [her_batch_size] */ - float* __restrict__ out_rewards, /* [her_batch_size] */ - float* __restrict__ out_dones, /* [her_batch_size] */ + __nv_bfloat16* __restrict__ out_rewards, /* [her_batch_size] */ + __nv_bfloat16* __restrict__ out_dones, /* [her_batch_size] */ int state_dim /* runtime state dimension */ ) { @@ -60,7 +60,7 @@ void her_relabel_kernel( /* 2. Replace goal portion (first GOAL_DIM elements) with donor's achieved goal. * The donor's achieved goal is extracted from the donor's next_state. */ for (int d = lane_id; d < GOAL_DIM; d += 32) { - float achieved = next_states[donor_idx * state_dim + d]; + __nv_bfloat16 achieved = next_states[donor_idx * state_dim + d]; out_states[sample_idx * state_dim + d] = achieved; out_next_states[sample_idx * state_dim + d] = achieved; } @@ -68,21 +68,21 @@ void her_relabel_kernel( /* 3. Scalar work: copy action/done, recompute sparse reward (lane 0 only). */ if (lane_id == 0) { out_actions[sample_idx] = actions[src_idx]; - out_dones[sample_idx] = dones[src_idx]; + out_dones[sample_idx] = dones[src_idx]; /* BF16 copy - OK */ /* Goal distance: L2 between source's achieved goal and donor's achieved goal. * Source achieved = next_states[src_idx, 0..GOAL_DIM] * Donor achieved = next_states[donor_idx, 0..GOAL_DIM] */ float dist_sq = 0.0f; for (int d = 0; d < GOAL_DIM; d++) { - float src_achieved = next_states[src_idx * state_dim + d]; - float donor_achieved = next_states[donor_idx * state_dim + d]; + float src_achieved = __bfloat162float(next_states[src_idx * state_dim + d]); + float donor_achieved = __bfloat162float(next_states[donor_idx * state_dim + d]); float diff = src_achieved - donor_achieved; dist_sq += diff * diff; } float dist = sqrtf(dist_sq); /* Sparse reward: +1.0 if goal achieved, -0.01 otherwise */ - out_rewards[sample_idx] = (dist < GOAL_THRESHOLD) ? 1.0f : -0.01f; + out_rewards[sample_idx] = __float2bfloat16((dist < GOAL_THRESHOLD) ? 1.0f : -0.01f); } } diff --git a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu index 56a6ccbba..f84691627 100644 --- a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu @@ -90,16 +90,16 @@ __device__ __forceinline__ float silu_grad(float x) { */ extern "C" __global__ void iql_forward_loss_kernel( - const float* __restrict__ states, - const float* __restrict__ q_values, - const float* __restrict__ params, - float* __restrict__ v_out, - float* __restrict__ loss_out, - float* __restrict__ total_loss, - float* __restrict__ save_pre1, - float* __restrict__ save_pre2, - float* __restrict__ save_h1, - float* __restrict__ save_h2, + const __nv_bfloat16* __restrict__ states, + const __nv_bfloat16* __restrict__ q_values, + const __nv_bfloat16* __restrict__ params, + __nv_bfloat16* __restrict__ v_out, + __nv_bfloat16* __restrict__ loss_out, + __nv_bfloat16* __restrict__ total_loss, + __nv_bfloat16* __restrict__ save_pre1, + __nv_bfloat16* __restrict__ save_pre2, + __nv_bfloat16* __restrict__ save_h1, + __nv_bfloat16* __restrict__ save_h2, int batch_size, int state_dim /* runtime state dimension */ ) @@ -112,41 +112,41 @@ void iql_forward_loss_kernel( int off_w1, off_b1, off_w2, off_b2, off_w3, off_b3; iql_compute_offsets(state_dim, &off_w1, &off_b1, &off_w2, &off_b2, &off_w3, &off_b3); - const float* w1 = params + off_w1; - const float* b1 = params + off_b1; - const float* w2 = params + off_w2; - const float* b2 = params + off_b2; - const float* w3 = params + off_w3; - const float* b3 = params + off_b3; + const __nv_bfloat16* w1 = params + off_w1; + const __nv_bfloat16* b1 = params + off_b1; + const __nv_bfloat16* w2 = params + off_w2; + const __nv_bfloat16* b2 = params + off_b2; + const __nv_bfloat16* w3 = params + off_w3; + const __nv_bfloat16* b3 = params + off_b3; - const float* x = states + sample * state_dim; + const __nv_bfloat16* x = states + sample * state_dim; /* ---- Layer 1: Linear(state_dim -> H) + SiLU ---- */ /* Each lane computes a subset of H outputs via warp reduction */ - float* pre1_ptr = save_pre1 + sample * VALUE_HIDDEN_DIM; - float* h1_ptr = save_h1 + sample * VALUE_HIDDEN_DIM; + __nv_bfloat16* pre1_ptr = save_pre1 + sample * VALUE_HIDDEN_DIM; + __nv_bfloat16* h1_ptr = save_h1 + sample * VALUE_HIDDEN_DIM; for (int j = lane; j < VALUE_HIDDEN_DIM; j += 32) { - float acc = b1[j]; + float acc = __bfloat162float(b1[j]); for (int k = 0; k < state_dim; k++) { - acc += w1[j * state_dim + k] * x[k]; + acc += __bfloat162float(w1[j * state_dim + k]) * __bfloat162float(x[k]); } - pre1_ptr[j] = acc; - h1_ptr[j] = silu(acc); + pre1_ptr[j] = __float2bfloat16(acc); + h1_ptr[j] = __float2bfloat16(silu(acc)); } __syncwarp(); /* ---- Layer 2: Linear(H -> H) + SiLU ---- */ - float* pre2_ptr = save_pre2 + sample * VALUE_HIDDEN_DIM; - float* h2_ptr = save_h2 + sample * VALUE_HIDDEN_DIM; + __nv_bfloat16* pre2_ptr = save_pre2 + sample * VALUE_HIDDEN_DIM; + __nv_bfloat16* h2_ptr = save_h2 + sample * VALUE_HIDDEN_DIM; for (int j = lane; j < VALUE_HIDDEN_DIM; j += 32) { - float acc = b2[j]; + float acc = __bfloat162float(b2[j]); for (int k = 0; k < VALUE_HIDDEN_DIM; k++) { - acc += w2[j * VALUE_HIDDEN_DIM + k] * h1_ptr[k]; + acc += __bfloat162float(w2[j * VALUE_HIDDEN_DIM + k]) * __bfloat162float(h1_ptr[k]); } - pre2_ptr[j] = acc; - h2_ptr[j] = silu(acc); + pre2_ptr[j] = __float2bfloat16(acc); + h2_ptr[j] = __float2bfloat16(silu(acc)); } __syncwarp(); @@ -154,7 +154,7 @@ void iql_forward_loss_kernel( /* Warp-reduce the dot product across lanes */ float v_acc = 0.0f; for (int k = lane; k < VALUE_HIDDEN_DIM; k += 32) { - v_acc += w3[k] * h2_ptr[k]; + v_acc += __bfloat162float(w3[k]) * __bfloat162float(h2_ptr[k]); } /* Warp reduce */ for (int offset = 16; offset > 0; offset >>= 1) { @@ -162,17 +162,17 @@ void iql_forward_loss_kernel( } /* Lane 0 has the full sum */ if (lane == 0) { - float v_val = v_acc + b3[0]; - v_out[sample] = v_val; + float v_val = v_acc + __bfloat162float(b3[0]); + v_out[sample] = __float2bfloat16(v_val); /* Expectile loss: L_tau(u) = |tau - 1(u<0)| * u^2 */ - float u = q_values[sample] - v_val; + float u = __bfloat162float(q_values[sample]) - v_val; float weight = (u >= 0.0f) ? IQL_EXPECTILE_TAU : (1.0f - IQL_EXPECTILE_TAU); float sample_loss = weight * u * u; - loss_out[sample] = sample_loss; + loss_out[sample] = __float2bfloat16(sample_loss); /* Accumulate batch-mean loss */ - atomicAdd(total_loss, sample_loss / (float)batch_size); + atomicAddBF16(total_loss, __float2bfloat16(sample_loss / (float)batch_size)); } } @@ -194,15 +194,15 @@ void iql_forward_loss_kernel( */ extern "C" __global__ void iql_backward_kernel( - const float* __restrict__ states, - const float* __restrict__ q_values, - const float* __restrict__ v_out, - const float* __restrict__ params, - const float* __restrict__ save_pre1, - const float* __restrict__ save_pre2, - const float* __restrict__ save_h1, - const float* __restrict__ save_h2, - float* __restrict__ grads, + const __nv_bfloat16* __restrict__ states, + const __nv_bfloat16* __restrict__ q_values, + const __nv_bfloat16* __restrict__ v_out, + const __nv_bfloat16* __restrict__ params, + const __nv_bfloat16* __restrict__ save_pre1, + const __nv_bfloat16* __restrict__ save_pre2, + const __nv_bfloat16* __restrict__ save_h1, + const __nv_bfloat16* __restrict__ save_h2, + __nv_bfloat16* __restrict__ grads, int batch_size, int state_dim /* runtime state dimension */ ) @@ -215,25 +215,25 @@ void iql_backward_kernel( int off_w1, off_b1, off_w2, off_b2, off_w3, off_b3; iql_compute_offsets(state_dim, &off_w1, &off_b1, &off_w2, &off_b2, &off_w3, &off_b3); - const float* w1 = params + off_w1; - const float* w2 = params + off_w2; - const float* w3 = params + off_w3; + const __nv_bfloat16* w1 = params + off_w1; + const __nv_bfloat16* w2 = params + off_w2; + const __nv_bfloat16* w3 = params + off_w3; - float* gw1 = grads + off_w1; - float* gb1 = grads + off_b1; - float* gw2 = grads + off_w2; - float* gb2 = grads + off_b2; - float* gw3 = grads + off_w3; - float* gb3 = grads + off_b3; + __nv_bfloat16* gw1 = grads + off_w1; + __nv_bfloat16* gb1 = grads + off_b1; + __nv_bfloat16* gw2 = grads + off_w2; + __nv_bfloat16* gb2 = grads + off_b2; + __nv_bfloat16* gw3 = grads + off_w3; + __nv_bfloat16* gb3 = grads + off_b3; - const float* x = states + sample * state_dim; - const float* h1 = save_h1 + sample * VALUE_HIDDEN_DIM; - const float* h2 = save_h2 + sample * VALUE_HIDDEN_DIM; - const float* pre1 = save_pre1 + sample * VALUE_HIDDEN_DIM; - const float* pre2 = save_pre2 + sample * VALUE_HIDDEN_DIM; + const __nv_bfloat16* x = states + sample * state_dim; + const __nv_bfloat16* h1 = save_h1 + sample * VALUE_HIDDEN_DIM; + const __nv_bfloat16* h2 = save_h2 + sample * VALUE_HIDDEN_DIM; + const __nv_bfloat16* pre1 = save_pre1 + sample * VALUE_HIDDEN_DIM; + const __nv_bfloat16* pre2 = save_pre2 + sample * VALUE_HIDDEN_DIM; - float v_val = v_out[sample]; - float q_val = q_values[sample]; + float v_val = __bfloat162float(v_out[sample]); + float q_val = __bfloat162float(q_values[sample]); /* dL/dV = -2 * weight * (Q - V) / batch_size */ float u = q_val - v_val; @@ -243,25 +243,25 @@ void iql_backward_kernel( /* ---- Output layer gradient: dL/dw3, dL/db3 ---- */ /* dV/dw3[k] = h2[k], dV/db3 = 1 */ for (int k = lane; k < VALUE_HIDDEN_DIM; k += 32) { - atomicAdd(&gw3[k], dldv * h2[k]); + atomicAddBF16(&gw3[k], __float2bfloat16(dldv * __bfloat162float(h2[k]))); } if (lane == 0) { - atomicAdd(&gb3[0], dldv); + atomicAddBF16(&gb3[0], __float2bfloat16(dldv)); } /* ---- Backprop through layer 2 ---- */ /* dL/dh2[k] = dldv * w3[k] */ /* dL/dpre2[k] = dL/dh2[k] * silu'(pre2[k]) */ for (int j = lane; j < VALUE_HIDDEN_DIM; j += 32) { - float dh2_j = dldv * w3[j]; - float dpre2_j = dh2_j * silu_grad(pre2[j]); + float dh2_j = dldv * __bfloat162float(w3[j]); + float dpre2_j = dh2_j * silu_grad(__bfloat162float(pre2[j])); /* dL/db2[j] = dpre2_j */ - atomicAdd(&gb2[j], dpre2_j); + atomicAddBF16(&gb2[j], __float2bfloat16(dpre2_j)); /* dL/dw2[j, k] = dpre2_j * h1[k] */ for (int k = 0; k < VALUE_HIDDEN_DIM; k++) { - atomicAdd(&gw2[j * VALUE_HIDDEN_DIM + k], dpre2_j * h1[k]); + atomicAddBF16(&gw2[j * VALUE_HIDDEN_DIM + k], __float2bfloat16(dpre2_j * __bfloat162float(h1[k]))); } } __syncwarp(); @@ -271,18 +271,18 @@ void iql_backward_kernel( for (int k = lane; k < VALUE_HIDDEN_DIM; k += 32) { float dh1_k = 0.0f; for (int j = 0; j < VALUE_HIDDEN_DIM; j++) { - float dh2_j = dldv * w3[j]; - float dpre2_j = dh2_j * silu_grad(pre2[j]); + float dh2_j = dldv * __bfloat162float(w3[j]); + float dpre2_j = dh2_j * silu_grad(__bfloat162float(pre2[j])); dh1_k += dpre2_j * w2[j * VALUE_HIDDEN_DIM + k]; } - float dpre1_k = dh1_k * silu_grad(pre1[k]); + float dpre1_k = dh1_k * silu_grad(__bfloat162float(pre1[k])); /* dL/db1[k] = dpre1_k */ - atomicAdd(&gb1[k], dpre1_k); + atomicAddBF16(&gb1[k], __float2bfloat16(dpre1_k)); /* dL/dw1[k, d] = dpre1_k * x[d] */ for (int d = 0; d < state_dim; d++) { - atomicAdd(&gw1[k * state_dim + d], dpre1_k * x[d]); + atomicAddBF16(&gw1[k * state_dim + d], __float2bfloat16(dpre1_k * __bfloat162float(x[d]))); } } } @@ -301,11 +301,11 @@ void iql_backward_kernel( */ extern "C" __global__ void iql_adam_kernel( - float* __restrict__ params, - float* __restrict__ grads, - float* __restrict__ m, /* first moment */ - float* __restrict__ v, /* second moment */ - const float* __restrict__ grad_norm_buf, + __nv_bfloat16* __restrict__ params, + __nv_bfloat16* __restrict__ grads, + __nv_bfloat16* __restrict__ m, /* first moment */ + __nv_bfloat16* __restrict__ v, /* second moment */ + const __nv_bfloat16* __restrict__ grad_norm_buf, float lr, float beta1, float beta2, @@ -320,17 +320,17 @@ void iql_adam_kernel( if (i >= total_params) return; /* Gradient clipping */ - float gn = grad_norm_buf[0]; + float gn = __bfloat162float(grad_norm_buf[0]); float clip_scale = (gn > max_grad_norm && gn > 0.0f) ? (max_grad_norm / gn) : 1.0f; - float g = grads[i] * clip_scale; + float g = __bfloat162float(grads[i]) * clip_scale; /* AdamW: decoupled weight decay */ - float p = params[i]; + float p = __bfloat162float(params[i]); p -= lr * weight_decay * p; /* Adam moment updates */ - float m_i = beta1 * m[i] + (1.0f - beta1) * g; - float v_i = beta2 * v[i] + (1.0f - beta2) * g * g; + float m_i = beta1 * __bfloat162float(m[i]) + (1.0f - beta1) * g; + float v_i = beta2 * __bfloat162float(v[i]) + (1.0f - beta2) * g * g; /* Bias correction */ float m_hat = m_i / (1.0f - powf(beta1, (float)t)); @@ -339,10 +339,10 @@ void iql_adam_kernel( /* Parameter update */ p -= lr * m_hat / (sqrtf(v_hat) + eps); - params[i] = p; - m[i] = m_i; - v[i] = v_i; - grads[i] = 0.0f; /* zero gradient for next step */ + params[i] = __float2bfloat16(p); + m[i] = __float2bfloat16(m_i); + v[i] = __float2bfloat16(v_i); + grads[i] = __float2bfloat16(0.0f); /* zero gradient for next step */ } /* ------------------------------------------------------------------ */ @@ -357,8 +357,8 @@ void iql_adam_kernel( */ extern "C" __global__ void iql_grad_norm_kernel( - const float* __restrict__ grads, - float* __restrict__ grad_norm_buf, + const __nv_bfloat16* __restrict__ grads, + __nv_bfloat16* __restrict__ grad_norm_buf, int total_params ) { @@ -368,7 +368,7 @@ void iql_grad_norm_kernel( int i = blockIdx.x * blockDim.x + threadIdx.x; /* Load and square */ - float val = (i < total_params) ? grads[i] : 0.0f; + float val = (i < total_params) ? __bfloat162float(grads[i]) : 0.0f; sdata[tid] = val * val; __syncthreads(); @@ -397,9 +397,9 @@ void iql_grad_norm_kernel( */ extern "C" __global__ void iql_forward_kernel( - const float* __restrict__ states, - const float* __restrict__ params, - float* __restrict__ v_out, + const __nv_bfloat16* __restrict__ states, + const __nv_bfloat16* __restrict__ params, + __nv_bfloat16* __restrict__ v_out, int batch_size, int state_dim /* runtime state dimension */ ) @@ -412,14 +412,14 @@ void iql_forward_kernel( int off_w1, off_b1, off_w2, off_b2, off_w3, off_b3; iql_compute_offsets(state_dim, &off_w1, &off_b1, &off_w2, &off_b2, &off_w3, &off_b3); - const float* w1 = params + off_w1; - const float* b1 = params + off_b1; - const float* w2 = params + off_w2; - const float* b2 = params + off_b2; - const float* w3 = params + off_w3; - const float* b3 = params + off_b3; + const __nv_bfloat16* w1 = params + off_w1; + const __nv_bfloat16* b1 = params + off_b1; + const __nv_bfloat16* w2 = params + off_w2; + const __nv_bfloat16* b2 = params + off_b2; + const __nv_bfloat16* w3 = params + off_w3; + const __nv_bfloat16* b3 = params + off_b3; - const float* x = states + sample * state_dim; + const __nv_bfloat16* x = states + sample * state_dim; /* Shared memory for intermediate activations (reused per layer) */ __shared__ float sh1[VALUE_HIDDEN_DIM]; @@ -427,9 +427,9 @@ void iql_forward_kernel( /* Layer 1: SiLU */ for (int j = lane; j < VALUE_HIDDEN_DIM; j += 32) { - float acc = b1[j]; + float acc = __bfloat162float(b1[j]); for (int k = 0; k < state_dim; k++) { - acc += w1[j * state_dim + k] * x[k]; + acc += __bfloat162float(w1[j * state_dim + k]) * __bfloat162float(x[k]); } sh1[j] = silu(acc); } @@ -437,9 +437,9 @@ void iql_forward_kernel( /* Layer 2: SiLU */ for (int j = lane; j < VALUE_HIDDEN_DIM; j += 32) { - float acc = b2[j]; + float acc = __bfloat162float(b2[j]); for (int k = 0; k < VALUE_HIDDEN_DIM; k++) { - acc += w2[j * VALUE_HIDDEN_DIM + k] * sh1[k]; + acc += __bfloat162float(w2[j * VALUE_HIDDEN_DIM + k]) * sh1[k]; } sh2[j] = silu(acc); } @@ -448,12 +448,12 @@ void iql_forward_kernel( /* Output: dot product + bias */ float v_acc = 0.0f; for (int k = lane; k < VALUE_HIDDEN_DIM; k += 32) { - v_acc += w3[k] * sh2[k]; + v_acc += __bfloat162float(w3[k]) * sh2[k]; } for (int offset = 16; offset > 0; offset >>= 1) { v_acc += __shfl_down_sync(0xFFFFFFFF, v_acc, offset); } if (lane == 0) { - v_out[sample] = v_acc + b3[0]; + v_out[sample] = __float2bfloat16(v_acc + __bfloat162float(b3[0])); } } diff --git a/crates/ml/src/cuda_pipeline/iqn_cvar_kernel.cu b/crates/ml/src/cuda_pipeline/iqn_cvar_kernel.cu index e757f153b..649c10b7b 100644 --- a/crates/ml/src/cuda_pipeline/iqn_cvar_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iqn_cvar_kernel.cu @@ -11,7 +11,7 @@ */ extern "C" __global__ void iqn_cvar_kernel( - const float* __restrict__ q_values, // [B, N_TAU, TBA] + const __nv_bfloat16* __restrict__ q_values, // [B, N_TAU, TBA] const int* __restrict__ actions, // [B] factored action indices float* scales_out, // [B] output CVaR scales int B, int N_TAU, int TBA, int B0_SIZE, int B1B2, float ALPHA @@ -35,7 +35,7 @@ extern "C" __global__ void iqn_cvar_kernel( // Fast path: just find the minimum float min_val = 1e30f; for (int t = 0; t < N_TAU; t++) { - float v = q_values[i * N_TAU * TBA + t * TBA + exposure_idx]; + float v = __bfloat162float(q_values[i * N_TAU * TBA + t * TBA + exposure_idx]); if (v < min_val) min_val = v; } cvar_sum = min_val; @@ -45,7 +45,7 @@ extern "C" __global__ void iqn_cvar_kernel( float sorted[8]; // max alpha_count for alpha=0.25, N_TAU=32 int sorted_len = 0; for (int t = 0; t < N_TAU; t++) { - float v = q_values[i * N_TAU * TBA + t * TBA + exposure_idx]; + float v = __bfloat162float(q_values[i * N_TAU * TBA + t * TBA + exposure_idx]); if (sorted_len < alpha_count) { // Insert into sorted array int pos = sorted_len; diff --git a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu index 8348b0250..777cc2821 100644 --- a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu @@ -185,25 +185,25 @@ __device__ __forceinline__ void iqn_compute_offsets(int hidden_dim, int embed_di extern "C" __global__ void iqn_forward_loss_kernel( /* Inputs */ - const float* __restrict__ h_s2, /* [B, hidden_dim] trunk activations */ - const float* __restrict__ target_h_s2, /* [B, hidden_dim] target trunk activations */ - const float* __restrict__ taus, /* [B, N] online τ samples ∈ (0,1) */ - const float* __restrict__ target_taus, /* [B, N] target τ samples */ + const __nv_bfloat16* __restrict__ h_s2, /* [B, hidden_dim] trunk activations */ + const __nv_bfloat16* __restrict__ target_h_s2, /* [B, hidden_dim] target trunk activations */ + const __nv_bfloat16* __restrict__ taus, /* [B, N] online τ samples ∈ (0,1) */ + const __nv_bfloat16* __restrict__ target_taus, /* [B, N] target τ samples */ const int* __restrict__ actions, /* [B, 3] branch actions (exposure, order, urgency) */ - const float* __restrict__ rewards, /* [B] */ - const float* __restrict__ dones, /* [B] */ + const __nv_bfloat16* __restrict__ rewards, /* [B] */ + const __nv_bfloat16* __restrict__ dones, /* [B] */ float gamma, /* Weights */ - const float* __restrict__ online_params, /* online IQN weights */ - const float* __restrict__ target_params, /* target IQN weights */ + const __nv_bfloat16* __restrict__ online_params, /* online IQN weights */ + const __nv_bfloat16* __restrict__ target_params, /* target IQN weights */ /* Outputs */ - float* __restrict__ per_sample_loss, /* [B] per-sample loss (for PER) */ - float* __restrict__ total_loss, /* [1] batch-mean loss (atomicAdd) */ + __nv_bfloat16* __restrict__ per_sample_loss, /* [B] per-sample loss (for PER) */ + __nv_bfloat16* __restrict__ total_loss, /* [1] batch-mean loss (atomicAdd) */ /* Activation saves for backward */ - float* __restrict__ save_embed, /* [B, N, hidden_dim] post-ReLU embeddings */ - float* __restrict__ save_combined, /* [B, N, hidden_dim] h_s2 ⊙ embed */ - float* __restrict__ save_q_online, /* [B, N, TOTAL_BRANCH_ACTIONS] quantile Q-values */ - float* __restrict__ save_q_target, /* [B, N, TOTAL_BRANCH_ACTIONS] target Q-values */ + __nv_bfloat16* __restrict__ save_embed, /* [B, N, hidden_dim] post-ReLU embeddings */ + __nv_bfloat16* __restrict__ save_combined, /* [B, N, hidden_dim] h_s2 ⊙ embed */ + __nv_bfloat16* __restrict__ save_q_online, /* [B, N, TOTAL_BRANCH_ACTIONS] quantile Q-values */ + __nv_bfloat16* __restrict__ save_q_target, /* [B, N, TOTAL_BRANCH_ACTIONS] target Q-values */ int batch_size, int shared_h1, /* runtime: first hidden layer width (unused here, for consistency) */ int hidden_dim, /* runtime: IQN hidden dim = SHARED_H2 */ @@ -219,10 +219,10 @@ void iqn_forward_loss_kernel( iqn_compute_offsets(hidden_dim, embed_dim, off); /* ── Pointers into this sample's data ── */ - const float* my_h_s2 = h_s2 + sample * hidden_dim; - const float* my_target_h_s2 = target_h_s2 + sample * hidden_dim; - const float* my_taus = taus + sample * IQN_NUM_QUANTILES; - const float* my_target_taus = target_taus + sample * IQN_NUM_QUANTILES; + const __nv_bfloat16* my_h_s2 = h_s2 + sample * hidden_dim; + const __nv_bfloat16* my_target_h_s2 = target_h_s2 + sample * hidden_dim; + const __nv_bfloat16* my_taus = taus + sample * IQN_NUM_QUANTILES; + const __nv_bfloat16* my_target_taus = target_taus + sample * IQN_NUM_QUANTILES; int a0 = actions[sample * 3 + 0]; /* exposure action */ int a1 = actions[sample * 3 + 1]; /* order action */ @@ -231,46 +231,46 @@ void iqn_forward_loss_kernel( float done = dones[sample]; /* Weight pointers (same for all samples) */ - const float* w_embed = online_params + off[0]; - const float* b_embed = online_params + off[1]; - const float* w_b0 = online_params + off[2]; - const float* b_b0 = online_params + off[3]; - const float* w_b1 = online_params + off[4]; - const float* b_b1 = online_params + off[5]; - const float* w_b2 = online_params + off[6]; - const float* b_b2 = online_params + off[7]; + const __nv_bfloat16* w_embed = online_params + off[0]; + const __nv_bfloat16* b_embed = online_params + off[1]; + const __nv_bfloat16* w_b0 = online_params + off[2]; + const __nv_bfloat16* b_b0 = online_params + off[3]; + const __nv_bfloat16* w_b1 = online_params + off[4]; + const __nv_bfloat16* b_b1 = online_params + off[5]; + const __nv_bfloat16* w_b2 = online_params + off[6]; + const __nv_bfloat16* b_b2 = online_params + off[7]; - const float* tw_embed = target_params + off[0]; - const float* tb_embed = target_params + off[1]; - const float* tw_b0 = target_params + off[2]; - const float* tb_b0 = target_params + off[3]; - const float* tw_b1 = target_params + off[4]; - const float* tb_b1 = target_params + off[5]; - const float* tw_b2 = target_params + off[6]; - const float* tb_b2 = target_params + off[7]; + const __nv_bfloat16* tw_embed = target_params + off[0]; + const __nv_bfloat16* tb_embed = target_params + off[1]; + const __nv_bfloat16* tw_b0 = target_params + off[2]; + const __nv_bfloat16* tb_b0 = target_params + off[3]; + const __nv_bfloat16* tw_b1 = target_params + off[4]; + const __nv_bfloat16* tb_b1 = target_params + off[5]; + const __nv_bfloat16* tw_b2 = target_params + off[6]; + const __nv_bfloat16* tb_b2 = target_params + off[7]; /* ── Load h_s2 into distributed registers ── */ float h_dist[IQN_DIST_MAX]; for (int d = lane; d < hidden_dim; d += 32) - h_dist[d / 32] = my_h_s2[d]; + h_dist[d / 32] = __bfloat162float(my_h_s2[d]); float h_target_dist[IQN_DIST_MAX]; for (int d = lane; d < hidden_dim; d += 32) - h_target_dist[d / 32] = my_target_h_s2[d]; + h_target_dist[d / 32] = __bfloat162float(my_target_h_s2[d]); /* ── Process each ONLINE quantile ── */ /* For each τ_i, compute quantile Q-values for the taken actions */ float online_q_a[IQN_NUM_QUANTILES]; /* Q(s, a_taken, τ_i) per branch sum */ for (int t = 0; t < IQN_NUM_QUANTILES; t++) { - float tau = my_taus[t]; + float tau = __bfloat162float(my_taus[t]); /* 1. Cosine features: cos(π·(d+1)·τ) for d=0..embed_dim-1 */ - /* 2. Linear embedding: embed[h] = Σ_d W_embed[h,d]·cos_feat[d] + b_embed[h] */ + /* 2. Linear embedding: embed[h] = Σ_d W_embed[h,d]·cos_feat[d] + __bfloat162float(b_embed[h]) */ /* Process via warp-cooperative matvec: each lane handles hidden_dim/32 outputs */ float embed_dist[IQN_DIST_MAX]; for (int h = lane; h < hidden_dim; h += 32) { - float acc = b_embed[h]; + float acc = __bfloat162float(b_embed[h]); const float* w_row = w_embed + h * embed_dim; for (int d = 0; d < embed_dim; d++) { float cos_val = cosf(3.14159265f * (float)(d + 1) * tau); @@ -298,7 +298,7 @@ void iqn_forward_loss_kernel( /* Branch 0 (exposure) */ float q0_taken = 0.0f; for (int a = 0; a < BRANCH_0_SIZE; a++) { - float acc = b_b0[a]; + float acc = __bfloat162float(b_b0[a]); const float* w_row = w_b0 + a * hidden_dim; float partial = 0.0f; for (int h = lane; h < hidden_dim; h += 32) @@ -314,7 +314,7 @@ void iqn_forward_loss_kernel( /* Branch 1 (order) */ float q1_taken = 0.0f; for (int a = 0; a < BRANCH_1_SIZE; a++) { - float acc = b_b1[a]; + float acc = __bfloat162float(b_b1[a]); const float* w_row = w_b1 + a * hidden_dim; float partial = 0.0f; for (int h = lane; h < hidden_dim; h += 32) @@ -330,7 +330,7 @@ void iqn_forward_loss_kernel( /* Branch 2 (urgency) */ float q2_taken = 0.0f; for (int a = 0; a < BRANCH_2_SIZE; a++) { - float acc = b_b2[a]; + float acc = __bfloat162float(b_b2[a]); const float* w_row = w_b2 + a * hidden_dim; float partial = 0.0f; for (int h = lane; h < hidden_dim; h += 32) @@ -351,12 +351,12 @@ void iqn_forward_loss_kernel( float target_q_a[IQN_NUM_QUANTILES]; for (int t = 0; t < IQN_NUM_QUANTILES; t++) { - float tau = my_target_taus[t]; + float tau = __bfloat162float(my_target_taus[t]); /* Cosine embedding + linear (target weights) */ float embed_dist[IQN_DIST_MAX]; for (int h = lane; h < hidden_dim; h += 32) { - float acc = tb_embed[h]; + float acc = __bfloat162float(tb_embed[h]); const float* w_row = tw_embed + h * embed_dim; for (int d = 0; d < embed_dim; d++) { float cos_val = cosf(3.14159265f * (float)(d + 1) * tau); @@ -375,7 +375,7 @@ void iqn_forward_loss_kernel( /* Branch 0 */ { - float acc = tb_b0[a0]; + float acc = __bfloat162float(tb_b0[a0]); float partial = 0.0f; const float* w_row = tw_b0 + a0 * hidden_dim; for (int h = lane; h < hidden_dim; h += 32) @@ -384,7 +384,7 @@ void iqn_forward_loss_kernel( } /* Branch 1 */ { - float acc = tb_b1[a1]; + float acc = __bfloat162float(tb_b1[a1]); float partial = 0.0f; const float* w_row = tw_b1 + a1 * hidden_dim; for (int h = lane; h < hidden_dim; h += 32) @@ -393,7 +393,7 @@ void iqn_forward_loss_kernel( } /* Branch 2 */ { - float acc = tb_b2[a2]; + float acc = __bfloat162float(tb_b2[a2]); float partial = 0.0f; const float* w_row = tw_b2 + a2 * hidden_dim; for (int h = lane; h < hidden_dim; h += 32) @@ -418,7 +418,7 @@ void iqn_forward_loss_kernel( float sample_loss = 0.0f; if (lane == 0) { for (int i = 0; i < IQN_NUM_QUANTILES; i++) { - float tau_i = my_taus[i]; + float tau_i = __bfloat162float(my_taus[i]); float inner_sum = 0.0f; for (int j = 0; j < IQN_NUM_QUANTILES; j++) { float delta = target_q_a[j] - online_q_a[i]; @@ -448,18 +448,18 @@ void iqn_forward_loss_kernel( extern "C" __global__ void iqn_backward_kernel( /* Saved activations (from forward pass) */ - const float* __restrict__ h_s2, /* [B, hidden_dim] */ - const float* __restrict__ save_embed, /* [B, N, hidden_dim] */ - const float* __restrict__ save_combined, /* [B, N, hidden_dim] */ - const float* __restrict__ save_q_online, /* [B, N, TOTAL_BRANCH_ACTIONS] */ - const float* __restrict__ save_q_target, /* [B, N, TOTAL_BRANCH_ACTIONS] */ - const float* __restrict__ taus, /* [B, N] */ + const __nv_bfloat16* __restrict__ h_s2, /* [B, hidden_dim] */ + const __nv_bfloat16* __restrict__ save_embed, /* [B, N, hidden_dim] */ + const __nv_bfloat16* __restrict__ save_combined, /* [B, N, hidden_dim] */ + const __nv_bfloat16* __restrict__ save_q_online, /* [B, N, TOTAL_BRANCH_ACTIONS] */ + const __nv_bfloat16* __restrict__ save_q_target, /* [B, N, TOTAL_BRANCH_ACTIONS] */ + const __nv_bfloat16* __restrict__ taus, /* [B, N] */ const int* __restrict__ actions, /* [B, 3] */ /* Weights (for backward through linear layers) */ - const float* __restrict__ online_params, + const __nv_bfloat16* __restrict__ online_params, /* Gradient output */ - float* __restrict__ grad_buf, /* accumulated gradients */ - float* __restrict__ d_h_s2_out, /* [B, hidden_dim] dL/d(h_s2) for trunk gradient */ + __nv_bfloat16* __restrict__ grad_buf, /* accumulated gradients */ + __nv_bfloat16* __restrict__ d_h_s2_out, /* [B, hidden_dim] dL/d(h_s2) for trunk gradient */ int batch_size, int shared_h1, /* runtime: first hidden layer width (unused here) */ int hidden_dim, /* runtime: IQN hidden dim = SHARED_H2 */ @@ -474,21 +474,21 @@ void iqn_backward_kernel( int off[8]; iqn_compute_offsets(hidden_dim, embed_dim, off); - const float* my_h_s2 = h_s2 + sample * hidden_dim; - const float* my_taus = taus + sample * IQN_NUM_QUANTILES; + const __nv_bfloat16* my_h_s2 = h_s2 + sample * hidden_dim; + const __nv_bfloat16* my_taus = taus + sample * IQN_NUM_QUANTILES; int a0 = actions[sample * 3 + 0]; int a1 = actions[sample * 3 + 1]; int a2 = actions[sample * 3 + 2]; - const float* w_embed = online_params + off[0]; - const float* w_b0 = online_params + off[2]; - const float* w_b1 = online_params + off[4]; - const float* w_b2 = online_params + off[6]; + const __nv_bfloat16* w_embed = online_params + off[0]; + const __nv_bfloat16* w_b0 = online_params + off[2]; + const __nv_bfloat16* w_b1 = online_params + off[4]; + const __nv_bfloat16* w_b2 = online_params + off[6]; /* Load h_s2 into distributed registers */ float h_dist[IQN_DIST_MAX]; for (int h = lane; h < hidden_dim; h += 32) - h_dist[h / 32] = my_h_s2[h]; + h_dist[h / 32] = __bfloat162float(my_h_s2[h]); /* Pre-compute online and target Q-values for taken actions across all quantiles */ /* (reading from saved buffers) */ @@ -510,7 +510,7 @@ void iqn_backward_kernel( float inv_n_sq = 1.0f / ((float)IQN_NUM_QUANTILES * (float)IQN_NUM_QUANTILES); for (int ti = 0; ti < IQN_NUM_QUANTILES; ti++) { - float tau_i = my_taus[ti]; + float tau_i = __bfloat162float(my_taus[ti]); int save_offset = (sample * IQN_NUM_QUANTILES + ti) * hidden_dim; /* Compute dL/dq_online for this quantile: @@ -541,7 +541,7 @@ void iqn_backward_kernel( /* Branch 0: accumulate gradient */ for (int h = lane; h < hidden_dim; h += 32) { - dL_dcomb[h / 32] += w_b0[a0 * hidden_dim + h] * dL_dq; + dL_dcomb[h / 32] += __bfloat162float(w_b0[a0 * hidden_dim + h]) * dL_dq; /* dL/dW_b0[a0, h] += dL/dq × combined[h] */ atomicAdd(&grad_buf[off[2] + a0 * hidden_dim + h], dL_dq * comb_dist[h / 32]); @@ -551,7 +551,7 @@ void iqn_backward_kernel( /* Branch 1 */ for (int h = lane; h < hidden_dim; h += 32) { - dL_dcomb[h / 32] += w_b1[a1 * hidden_dim + h] * dL_dq; + dL_dcomb[h / 32] += __bfloat162float(w_b1[a1 * hidden_dim + h]) * dL_dq; atomicAdd(&grad_buf[off[4] + a1 * hidden_dim + h], dL_dq * comb_dist[h / 32]); } @@ -560,7 +560,7 @@ void iqn_backward_kernel( /* Branch 2 */ for (int h = lane; h < hidden_dim; h += 32) { - dL_dcomb[h / 32] += w_b2[a2 * hidden_dim + h] * dL_dq; + dL_dcomb[h / 32] += __bfloat162float(w_b2[a2 * hidden_dim + h]) * dL_dq; atomicAdd(&grad_buf[off[6] + a2 * hidden_dim + h], dL_dq * comb_dist[h / 32]); } @@ -588,10 +588,10 @@ void iqn_backward_kernel( } /* ── Gradient through embedding linear layer ── */ - /* pre_relu[h] = Σ_d W_embed[h,d]·cos(π(d+1)τ) + b_embed[h] + /* pre_relu[h] = Σ_d W_embed[h,d]·cos(π(d+1)τ) + __bfloat162float(b_embed[h]) * dL/dW_embed[h,d] += dL/d(pre_relu)[h] × cos(π(d+1)τ) * dL/db_embed[h] += dL/d(pre_relu)[h] */ - float tau = my_taus[ti]; + float tau = __bfloat162float(my_taus[ti]); for (int h = lane; h < hidden_dim; h += 32) { float dL_dpre = dL_dembed[h / 32]; /* Bias gradient */ @@ -609,8 +609,8 @@ void iqn_backward_kernel( /* ── Gradient Norm Kernel ───────────────────────────────────────────── */ extern "C" __global__ void iqn_grad_norm_kernel( - const float* __restrict__ grads, - float* __restrict__ norm_out, /* [1] */ + const __nv_bfloat16* __restrict__ grads, + __nv_bfloat16* __restrict__ norm_out, /* [1] */ int total_params, int shared_h1, /* runtime: unused, for consistent interface */ int hidden_dim, /* runtime: unused, for consistent interface */ @@ -645,11 +645,11 @@ void iqn_grad_norm_kernel( /* ── Adam Update Kernel ─────────────────────────────────────────────── */ extern "C" __global__ void iqn_adam_kernel( - float* __restrict__ params, - float* __restrict__ grads, - float* __restrict__ m, /* first moment */ - float* __restrict__ v, /* second moment */ - const float* __restrict__ norm, /* [1] gradient L2 norm */ + __nv_bfloat16* __restrict__ params, + __nv_bfloat16* __restrict__ grads, + __nv_bfloat16* __restrict__ m, /* first moment */ + __nv_bfloat16* __restrict__ v, /* second moment */ + const __nv_bfloat16* __restrict__ norm, /* [1] gradient L2 norm */ float lr, float beta1, float beta2, float eps, float weight_decay, float max_grad_norm, int adam_t, /* step counter (1-indexed) */ @@ -666,7 +666,7 @@ void iqn_adam_kernel( float grad_norm = norm[0]; float clip_scale = (grad_norm > max_grad_norm && grad_norm > 0.0f) ? max_grad_norm / grad_norm : 1.0f; - float g = grads[tid] * clip_scale; + float g = __bfloat162float(grads[tid]) * clip_scale; /* AdamW update */ float mi = beta1 * m[tid] + (1.0f - beta1) * g; @@ -684,7 +684,7 @@ void iqn_adam_kernel( params[tid] = p; /* Zero gradient for next step */ - grads[tid] = 0.0f; + grads[tid] = __float2bfloat16(0.0f); } /* ── Forward-Only Kernel (Inference) ────────────────────────────────── */ @@ -698,10 +698,10 @@ void iqn_adam_kernel( */ extern "C" __global__ void iqn_forward_kernel( - const float* __restrict__ h_s2, /* [B, hidden_dim] */ - const float* __restrict__ taus, /* [B, N] */ - const float* __restrict__ params, /* IQN weights */ - float* __restrict__ expected_q, /* [B, TOTAL_BRANCH_ACTIONS] */ + const __nv_bfloat16* __restrict__ h_s2, /* [B, hidden_dim] */ + const __nv_bfloat16* __restrict__ taus, /* [B, N] */ + const __nv_bfloat16* __restrict__ params, /* IQN weights */ + __nv_bfloat16* __restrict__ expected_q, /* [B, TOTAL_BRANCH_ACTIONS] */ int batch_size, int shared_h1, /* runtime: unused here, for consistency */ int hidden_dim, /* runtime: IQN hidden dim = SHARED_H2 */ @@ -716,8 +716,8 @@ void iqn_forward_kernel( int off[8]; iqn_compute_offsets(hidden_dim, embed_dim, off); - const float* my_h_s2 = h_s2 + sample * hidden_dim; - const float* my_taus = taus + sample * IQN_NUM_QUANTILES; + const __nv_bfloat16* my_h_s2 = h_s2 + sample * hidden_dim; + const __nv_bfloat16* my_taus = taus + sample * IQN_NUM_QUANTILES; const float* w_embed = params + off[0]; const float* b_embed = params + off[1]; const float* w_b0 = params + off[2]; @@ -730,7 +730,7 @@ void iqn_forward_kernel( /* Load h_s2 */ float h_dist[IQN_DIST_MAX]; for (int h = lane; h < hidden_dim; h += 32) - h_dist[h / 32] = my_h_s2[h]; + h_dist[h / 32] = __bfloat162float(my_h_s2[h]); /* Accumulate Q-values across quantiles (for mean) */ float q_acc[TOTAL_BRANCH_ACTIONS]; @@ -738,12 +738,12 @@ void iqn_forward_kernel( q_acc[a] = 0.0f; for (int t = 0; t < IQN_NUM_QUANTILES; t++) { - float tau = my_taus[t]; + float tau = __bfloat162float(my_taus[t]); /* Cosine embedding + linear + ReLU */ float embed_dist[IQN_DIST_MAX]; for (int h = lane; h < hidden_dim; h += 32) { - float acc = b_embed[h]; + float acc = __bfloat162float(b_embed[h]); const float* w_row = w_embed + h * embed_dim; for (int d = 0; d < embed_dim; d++) { float cos_val = cosf(3.14159265f * (float)(d + 1) * tau); @@ -764,7 +764,7 @@ void iqn_forward_kernel( const float* w_row = w_b0 + a * hidden_dim; for (int h = lane; h < hidden_dim; h += 32) partial += w_row[h] * comb_dist[h / 32]; - float q_val = iqn_warp_sum(partial) + b_b0[a]; + float q_val = iqn_warp_sum(partial) + __bfloat162float(b_b0[a]); q_acc[a] += q_val; } /* Branch 1 */ @@ -773,7 +773,7 @@ void iqn_forward_kernel( const float* w_row = w_b1 + a * hidden_dim; for (int h = lane; h < hidden_dim; h += 32) partial += w_row[h] * comb_dist[h / 32]; - float q_val = iqn_warp_sum(partial) + b_b1[a]; + float q_val = iqn_warp_sum(partial) + __bfloat162float(b_b1[a]); q_acc[BRANCH_0_SIZE + a] += q_val; } /* Branch 2 */ @@ -782,7 +782,7 @@ void iqn_forward_kernel( const float* w_row = w_b2 + a * hidden_dim; for (int h = lane; h < hidden_dim; h += 32) partial += w_row[h] * comb_dist[h / 32]; - float q_val = iqn_warp_sum(partial) + b_b2[a]; + float q_val = iqn_warp_sum(partial) + __bfloat162float(b_b2[a]); q_acc[BRANCH_0_SIZE + BRANCH_1_SIZE + a] += q_val; } } @@ -808,12 +808,12 @@ void iqn_forward_kernel( */ extern "C" __global__ void iqn_trunk_forward_kernel( - const float* __restrict__ states, /* [B, state_dim] */ - const float* __restrict__ w_s1, /* [shared_h1, state_dim] */ - const float* __restrict__ b_s1, /* [shared_h1] */ - const float* __restrict__ w_s2, /* [hidden_dim, shared_h1] */ - const float* __restrict__ b_s2, /* [hidden_dim] */ - float* __restrict__ h_s2_out, /* [B, hidden_dim] */ + const __nv_bfloat16* __restrict__ states, /* [B, state_dim] */ + const __nv_bfloat16* __restrict__ w_s1, /* [shared_h1, state_dim] */ + const __nv_bfloat16* __restrict__ b_s1, /* [shared_h1] */ + const __nv_bfloat16* __restrict__ w_s2, /* [hidden_dim, shared_h1] */ + const __nv_bfloat16* __restrict__ b_s2, /* [hidden_dim] */ + __nv_bfloat16* __restrict__ h_s2_out, /* [B, hidden_dim] */ int batch_size, int state_dim, /* runtime state dimension (replaces IQN_STATE_DIM) */ int shared_h1, /* runtime: first hidden layer width */ @@ -857,8 +857,8 @@ void iqn_trunk_forward_kernel( */ extern "C" __global__ void iqn_ema_kernel( - float* __restrict__ target, - const float* __restrict__ online, + __nv_bfloat16* __restrict__ target, + const __nv_bfloat16* __restrict__ online, float tau, int n, int shared_h1, /* runtime: unused, for consistent interface */ @@ -937,7 +937,7 @@ unsigned int iqn_philox_single(unsigned int counter, unsigned int key) { extern "C" __global__ void iqn_sample_taus_kernel( - float* __restrict__ taus, + __nv_bfloat16* __restrict__ taus, unsigned int seed, int total, int shared_h1, /* runtime: unused, for consistent interface */ diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index 0e1365448..578e9fbed 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -5,7 +5,7 @@ //! pre-uploaded buffers instead of creating per-step copies. //! //! All modules operate on raw `CudaSlice` via cudarc. -//! `DqnGpuData` and `PpoGpuData` hold GPU-resident `CudaSlice` buffers. +//! `DqnGpuData` and `PpoGpuData` hold GPU-resident `CudaSlice` buffers. use std::sync::Arc; use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut}; @@ -102,15 +102,15 @@ pub fn optimal_launch_dims(n_items: u32, max_threads_per_block: u32) -> (u32, u3 /// The fused experience collector kernel (4490 lines, branching+C51+NoisyNets) // compile_ptx_for_device re-export removed: all kernels now use precompiled cubins. -/// Clone a `CudaSlice` via DtoD memcpy. +/// Clone a `CudaSlice` via DtoD memcpy. /// /// The caller owns the returned slice independently of the source. pub fn clone_cuda_slice_f32( - src: &CudaSlice, + src: &CudaSlice, stream: &Arc, -) -> Result, crate::MLError> { +) -> Result, crate::MLError> { let n = src.len(); - let mut dst = stream.alloc_zeros::(n).map_err(|e| { + let mut dst = stream.alloc_zeros::(n).map_err(|e| { crate::MLError::ModelError(format!("clone_cuda_slice_f32 alloc: {e}")) })?; { @@ -153,16 +153,16 @@ pub fn clone_cuda_slice_u32( /// Pre-uploaded GPU training data for DQN trainer. /// /// Holds market features [N * 42] and target prices [N * 4] as GPU-resident -/// `CudaSlice` buffers. Portfolio features (3 dims) are computed per-bar +/// `CudaSlice` buffers. Portfolio features (3 dims) are computed per-bar /// by the trainer and concatenated on-device. OFI features (8 dims from MBP-10 /// order book) are optionally uploaded and concatenated to produce [N, 51] states. pub struct DqnGpuData { /// Market features [num_bars * 42] on GPU (f32, row-major) - pub features: CudaSlice, + pub features: CudaSlice, /// Target prices [num_bars * 4] on GPU (f32, row-major) - pub targets: CudaSlice, + pub targets: CudaSlice, /// OFI features [num_bars * 8] on GPU (f32), from MBP-10 data - pub ofi_features: Option>, + pub ofi_features: Option>, /// Number of training bars pub num_bars: usize, /// Feature dimension (42) @@ -287,12 +287,12 @@ impl DqnGpuData { /// D2D subrange copy: extract `count` f32 elements at `offset` into a new `CudaSlice`. pub(crate) fn d2d_subrange( - src: &CudaSlice, + src: &CudaSlice, offset: usize, count: usize, stream: &Arc, - ) -> Result, MLError> { - let mut dst = stream.alloc_zeros::(count).map_err(|e| { + ) -> Result, MLError> { + let mut dst = stream.alloc_zeros::(count).map_err(|e| { MLError::ModelError(format!("d2d_subrange alloc: {e}")) })?; { @@ -311,7 +311,7 @@ impl DqnGpuData { } /// Get market features for a single bar as a [42] `CudaSlice` (D2D copy from GPU buffer). - pub fn bar_features(&self, bar_idx: usize, stream: &Arc) -> Result, MLError> { + pub fn bar_features(&self, bar_idx: usize, stream: &Arc) -> Result, MLError> { if bar_idx >= self.num_bars { return Err(MLError::ModelError(format!("Bar {bar_idx} out of range (num_bars={})", self.num_bars))); } @@ -319,13 +319,13 @@ impl DqnGpuData { } /// Get market features for a batch of bars as a [batch_size * 42] `CudaSlice`. - pub fn batch_features(&self, start: usize, count: usize, stream: &Arc) -> Result, MLError> { + pub fn batch_features(&self, start: usize, count: usize, stream: &Arc) -> Result, MLError> { let count = count.min(self.num_bars.saturating_sub(start)); Self::d2d_subrange(&self.features, start * self.feature_dim, count * self.feature_dim, stream) } /// Get target prices for a single bar as a [4] `CudaSlice`. - pub fn bar_targets(&self, bar_idx: usize, stream: &Arc) -> Result, MLError> { + pub fn bar_targets(&self, bar_idx: usize, stream: &Arc) -> Result, MLError> { if bar_idx >= self.num_bars { return Err(MLError::ModelError(format!("Bar {bar_idx} target out of range (num_bars={})", self.num_bars))); } @@ -333,16 +333,16 @@ impl DqnGpuData { } /// Get target prices for a batch as [batch_size * 4] `CudaSlice`. - pub fn batch_targets(&self, start: usize, count: usize, stream: &Arc) -> Result, MLError> { + pub fn batch_targets(&self, start: usize, count: usize, stream: &Arc) -> Result, MLError> { let count = count.min(self.num_bars.saturating_sub(start)); Self::d2d_subrange(&self.targets, start * 4, count * 4, stream) } - /// Get target values for a bar as a GPU-resident [4] `CudaSlice`. + /// Get target values for a bar as a GPU-resident [4] `CudaSlice`. /// /// Tensor order: [current_close_preproc, next_close_preproc, current_close_raw, next_close_raw]. /// Stays on GPU. Callers that need CPU scalars should extract at their own boundary. - pub fn bar_target_values(&self, bar_idx: usize, stream: &Arc) -> Result, MLError> { + pub fn bar_target_values(&self, bar_idx: usize, stream: &Arc) -> Result, MLError> { self.bar_targets(bar_idx, stream) } @@ -355,13 +355,13 @@ impl DqnGpuData { bar_idx: usize, portfolio_features: &[f32; 3], stream: &Arc, - ) -> Result, MLError> { + ) -> Result, MLError> { let ofi_dim = if self.ofi_features.is_some() { 8 } else { 0 }; let raw_dim = self.feature_dim + 3 + ofi_dim; let final_dim = self.aligned_state_dim.unwrap_or(raw_dim); // Allocate output buffer (zeros for padding region) - let mut dst = stream.alloc_zeros::(final_dim) + let mut dst = stream.alloc_zeros::(final_dim) .map_err(|e| MLError::ModelError(format!("build_state alloc: {e}")))?; // DtoD copy: market features -> dst[0..feature_dim] @@ -390,7 +390,7 @@ impl DqnGpuData { count: usize, portfolio_features: &[f32; 3], stream: &Arc, - ) -> Result, MLError> { + ) -> Result, MLError> { let count = count.min(self.num_bars.saturating_sub(start)); if count == 0 { return Err(MLError::ModelError("Empty batch for state construction".to_owned())); @@ -400,7 +400,7 @@ impl DqnGpuData { let final_dim = self.aligned_state_dim.unwrap_or(raw_dim); // Allocate output buffer (zeros for padding region) - let mut dst = stream.alloc_zeros::(count * final_dim) + let mut dst = stream.alloc_zeros::(count * final_dim) .map_err(|e| MLError::ModelError(format!("build_batch alloc: {e}")))?; // Upload portfolio features once (3 scalars) @@ -447,8 +447,8 @@ impl DqnGpuData { /// DtoD copy `count` f32 elements from `src` (offset 0) into `dst` at element offset. #[allow(unsafe_code)] fn dtod_copy_into( - src: &CudaSlice, - dst: &mut CudaSlice, + src: &CudaSlice, + dst: &mut CudaSlice, dst_offset_elems: usize, count: usize, stream: &Arc, @@ -462,9 +462,9 @@ impl DqnGpuData { /// buffer. The source is sliced via `CudaSlice::slice()` for safety. #[allow(unsafe_code)] fn dtod_copy_into_at_offset( - src: &CudaSlice, + src: &CudaSlice, src_offset_elems: usize, - dst: &mut CudaSlice, + dst: &mut CudaSlice, dst_offset_elems: usize, count: usize, stream: &Arc, @@ -493,7 +493,7 @@ impl DqnGpuData { #[allow(unsafe_code)] fn htod_copy_into( src: &[f32], - dst: &mut CudaSlice, + dst: &mut CudaSlice, dst_offset_elems: usize, stream: &Arc, ) -> Result<(), MLError> { @@ -608,10 +608,10 @@ impl GpuBufferPool { /// Pre-uploaded GPU training data for PPO trainer. /// -/// Holds all rollout states as a single [N * state_dim] `CudaSlice`. +/// Holds all rollout states as a single [N * state_dim] `CudaSlice`. pub struct PpoGpuData { /// All states [num_steps * state_dim] on GPU (f32, row-major) - pub states: CudaSlice, + pub states: CudaSlice, /// Number of training steps pub num_steps: usize, /// State dimension @@ -628,7 +628,7 @@ impl std::fmt::Debug for PpoGpuData { } impl PpoGpuData { - /// Upload PPO market data to GPU as a single contiguous `CudaSlice`. + /// Upload PPO market data to GPU as a single contiguous `CudaSlice`. pub fn upload(market_data: &[Vec], stream: &Arc) -> Result { let num_steps = market_data.len(); if num_steps == 0 { @@ -674,7 +674,7 @@ impl PpoGpuData { } /// Get a single step's state as a [state_dim] `CudaSlice` (D2D copy from GPU buffer). - pub fn step_state(&self, step_idx: usize, stream: &Arc) -> Result, MLError> { + pub fn step_state(&self, step_idx: usize, stream: &Arc) -> Result, MLError> { if step_idx >= self.num_steps { return Err(MLError::ModelError(format!("Step {step_idx} out of range (num_steps={})", self.num_steps))); } @@ -682,7 +682,7 @@ impl PpoGpuData { } /// Get a batch of states as [batch_size * state_dim] `CudaSlice`. - pub fn batch_states(&self, start: usize, count: usize, stream: &Arc) -> Result, MLError> { + pub fn batch_states(&self, start: usize, count: usize, stream: &Arc) -> Result, MLError> { let count = count.min(self.num_steps.saturating_sub(start)); DqnGpuData::d2d_subrange(&self.states, start * self.state_dim, count * self.state_dim, stream) } diff --git a/crates/ml/src/cuda_pipeline/monitoring_kernel.cu b/crates/ml/src/cuda_pipeline/monitoring_kernel.cu index 67e8cf753..560f1d0ec 100644 --- a/crates/ml/src/cuda_pipeline/monitoring_kernel.cu +++ b/crates/ml/src/cuda_pipeline/monitoring_kernel.cu @@ -24,7 +24,7 @@ __device__ float atomicMax_float(float* addr, float val) { // Reduce per-experience rewards and actions into a compact summary. // One block, parallel reduction across N elements. extern "C" __global__ void monitoring_reduce( - const float* __restrict__ rewards, // [N] + const __nv_bfloat16* __restrict__ rewards, // [N] const int* __restrict__ actions, // [N] float* summary, // [16]: mean, std, min, max, sharpe, counts[9], total, _pad int N, @@ -55,7 +55,7 @@ extern "C" __global__ void monitoring_reduce( int local_counts[9] = {0, 0, 0, 0, 0, 0, 0, 0, 0}; for (int i = tid; i < N; i += stride) { - float r = rewards[i]; + float r = __bfloat162float(rewards[i]); local_sum += r; local_sq += r * r; local_min = fminf(local_min, r); diff --git a/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu b/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu index 12a3e5be4..9f784ee51 100644 --- a/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu @@ -150,7 +150,7 @@ extern "C" __global__ void mse_loss_batched( /* ── Outputs ──────────────────────────────────────────────────── */ __nv_bfloat16* __restrict__ per_sample_loss, /* [B] IS-weighted loss per sample */ __nv_bfloat16* __restrict__ td_errors, /* [B] unweighted, for PER priority update */ - float* __restrict__ total_loss, /* [1] batch mean loss (atomicAdd) — stays F32, no BF16 atomicAdd */ + __nv_bfloat16* __restrict__ total_loss, /* [1] batch mean loss (atomicAdd) — stays F32, no BF16 atomicAdd */ /* ── Saved tensors for backward pass ─────────────────────────── */ __nv_bfloat16* __restrict__ save_current_lp, /* [B, NUM_BRANCHES, num_atoms] online probs */ @@ -385,6 +385,16 @@ extern "C" __global__ void mse_loss_batched( float weighted_loss = avg_mse * is_weight; per_sample_loss[sample_id] = __float2bfloat16(weighted_loss); td_errors[sample_id] = __float2bfloat16(avg_td); - atomicAdd(total_loss, weighted_loss / (float)batch_size); + /* BF16 atomicAdd via CAS loop */ + { + unsigned short* addr = (unsigned short*)total_loss; + unsigned short old_bits = *addr, assumed; + do { + assumed = old_bits; + __nv_bfloat16 old_bf = *(__nv_bfloat16*)&assumed; + __nv_bfloat16 new_bf = __float2bfloat16(__bfloat162float(old_bf) + weighted_loss / (float)batch_size); + old_bits = atomicCAS(addr, assumed, *(unsigned short*)&new_bf); + } while (old_bits != assumed); + } } } diff --git a/crates/ml/src/cuda_pipeline/nstep_kernel.cu b/crates/ml/src/cuda_pipeline/nstep_kernel.cu index bb5e303eb..509875ec0 100644 --- a/crates/ml/src/cuda_pipeline/nstep_kernel.cu +++ b/crates/ml/src/cuda_pipeline/nstep_kernel.cu @@ -14,10 +14,10 @@ */ extern "C" __global__ void nstep_accumulate_kernel( - const float* __restrict__ raw_rewards, /* [N * L] original 1-step rewards */ - const float* __restrict__ raw_dones, /* [N * L] original 1-step dones */ - float* __restrict__ out_rewards, /* [N * L] overwritten with R_n */ - float* __restrict__ out_dones, /* [N * L] overwritten with done_n */ + const __nv_bfloat16* __restrict__ raw_rewards, /* [N * L] original 1-step rewards */ + const __nv_bfloat16* __restrict__ raw_dones, /* [N * L] original 1-step dones */ + __nv_bfloat16* __restrict__ out_rewards, /* [N * L] overwritten with R_n */ + __nv_bfloat16* __restrict__ out_dones, /* [N * L] overwritten with done_n */ float gamma, int n_steps, int L, /* timesteps per episode */ @@ -38,8 +38,8 @@ extern "C" __global__ void nstep_accumulate_kernel( if (step >= L) break; int off = ep * L + step; - float r_i = raw_rewards[off]; - float d_i = raw_dones[off]; + float r_i = __bfloat162float(raw_rewards[off]); + float d_i = __bfloat162float(raw_dones[off]); R_n += gamma_pow * r_i; gamma_pow *= gamma; @@ -51,8 +51,8 @@ extern "C" __global__ void nstep_accumulate_kernel( } int base = ep * L + t; - out_rewards[base] = R_n; - out_dones[base] = any_done; + out_rewards[base] = __float2bfloat16(R_n); + out_dones[base] = __float2bfloat16(any_done); } /* ══════════════════════════════════════════════════════════════════════ @@ -72,11 +72,11 @@ extern "C" __global__ void nstep_accumulate_kernel( * which has STATE_DIM as a compile-time constant for tensor core alignment. */ extern "C" __global__ void reward_normalize_kernel( - float* __restrict__ rewards, + __nv_bfloat16* __restrict__ rewards, float mean, float inv_std, int n ) { int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) rewards[i] = (rewards[i] - mean) * inv_std; + if (i < n) rewards[i] = __float2bfloat16((__bfloat162float(rewards[i]) - mean) * inv_std); } diff --git a/crates/ml/src/cuda_pipeline/per_update_kernel.cu b/crates/ml/src/cuda_pipeline/per_update_kernel.cu index ff50adcc9..1ab6b77e4 100644 --- a/crates/ml/src/cuda_pipeline/per_update_kernel.cu +++ b/crates/ml/src/cuda_pipeline/per_update_kernel.cu @@ -9,10 +9,10 @@ extern "C" __global__ void per_update_priorities_kernel( - const float* __restrict__ td_errors, + const __nv_bfloat16* __restrict__ td_errors, const unsigned int* __restrict__ indices, - float* __restrict__ priorities, - float* __restrict__ batch_max, + __nv_bfloat16* __restrict__ priorities, + __nv_bfloat16* __restrict__ batch_max, float alpha, float epsilon, int batch_size, @@ -21,7 +21,7 @@ void per_update_priorities_kernel( int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= batch_size) return; - float td = fabsf(td_errors[i]); + float td = fabsf(__bfloat162float(td_errors[i])); float new_prio = powf(td, alpha) + epsilon; // Clamp to [epsilon, 1e6] — same bounds as GpuReplayBuffer::update_priorities_gpu @@ -31,9 +31,9 @@ void per_update_priorities_kernel( // Scatter-write priority at buffer index (bounds-checked) unsigned int idx = indices[i]; if (idx < (unsigned int)capacity) - priorities[idx] = new_prio; + priorities[idx] = __float2bfloat16(new_prio); - // atomicMax for batch max: IEEE 754 positive floats preserve int ordering - int ival = __float_as_int(new_prio); - atomicMax((int*)batch_max, ival); + // atomicMax for batch max: store as float bits in the BF16 pair + // batch_max is a single scalar; use float* alias for atomicMax + atomicMax((int*)batch_max, __float_as_int(new_prio)); } diff --git a/crates/ml/src/cuda_pipeline/ppo_experience_kernel.cu b/crates/ml/src/cuda_pipeline/ppo_experience_kernel.cu index d0076ad45..32908bded 100644 --- a/crates/ml/src/cuda_pipeline/ppo_experience_kernel.cu +++ b/crates/ml/src/cuda_pipeline/ppo_experience_kernel.cu @@ -49,12 +49,12 @@ */ __device__ void ppo_actor_forward( const float* state, - const float* __restrict__ pw1, /* [ACTOR_H1, STATE_DIM] */ - const float* __restrict__ pb1, /* [ACTOR_H1] */ - const float* __restrict__ pw2, /* [ACTOR_H2, ACTOR_H1] */ - const float* __restrict__ pb2, /* [ACTOR_H2] */ - const float* __restrict__ pw3, /* [NUM_ACTIONS, ACTOR_H2] */ - const float* __restrict__ pb3, /* [NUM_ACTIONS] */ + const __nv_bfloat16* __restrict__ pw1, /* [ACTOR_H1, STATE_DIM] */ + const __nv_bfloat16* __restrict__ pb1, /* [ACTOR_H1] */ + const __nv_bfloat16* __restrict__ pw2, /* [ACTOR_H2, ACTOR_H1] */ + const __nv_bfloat16* __restrict__ pb2, /* [ACTOR_H2] */ + const __nv_bfloat16* __restrict__ pw3, /* [NUM_ACTIONS, ACTOR_H2] */ + const __nv_bfloat16* __restrict__ pb3, /* [NUM_ACTIONS] */ float* h1, /* [ACTOR_H1] scratch */ float* h2, /* [ACTOR_H2] scratch */ float* logits /* [NUM_ACTIONS] output */ @@ -145,18 +145,18 @@ __device__ void softmax_sample( */ __device__ float ppo_critic_forward( const float* state, - const float* __restrict__ vw1, /* [CRITIC_H1, STATE_DIM] = [512, 54] */ - const float* __restrict__ vb1, /* [CRITIC_H1] = [512] */ - const float* __restrict__ vw2, /* [CRITIC_H2, CRITIC_H1] = [384, 512] */ - const float* __restrict__ vb2, /* [CRITIC_H2] = [384] */ - const float* __restrict__ vw3, /* [CRITIC_H3, CRITIC_H2] = [256, 384] */ - const float* __restrict__ vb3, /* [CRITIC_H3] = [256] */ - const float* __restrict__ vw4, /* [CRITIC_H4, CRITIC_H3] = [128, 256] */ - const float* __restrict__ vb4, /* [CRITIC_H4] = [128] */ - const float* __restrict__ vw5, /* [CRITIC_H5, CRITIC_H4] = [64, 128] */ - const float* __restrict__ vb5, /* [CRITIC_H5] = [64] */ - const float* __restrict__ vw6, /* [1, CRITIC_H5] = [1, 64] */ - const float* __restrict__ vb6, /* [1] */ + const __nv_bfloat16* __restrict__ vw1, /* [CRITIC_H1, STATE_DIM] = [512, 54] */ + const __nv_bfloat16* __restrict__ vb1, /* [CRITIC_H1] = [512] */ + const __nv_bfloat16* __restrict__ vw2, /* [CRITIC_H2, CRITIC_H1] = [384, 512] */ + const __nv_bfloat16* __restrict__ vb2, /* [CRITIC_H2] = [384] */ + const __nv_bfloat16* __restrict__ vw3, /* [CRITIC_H3, CRITIC_H2] = [256, 384] */ + const __nv_bfloat16* __restrict__ vb3, /* [CRITIC_H3] = [256] */ + const __nv_bfloat16* __restrict__ vw4, /* [CRITIC_H4, CRITIC_H3] = [128, 256] */ + const __nv_bfloat16* __restrict__ vb4, /* [CRITIC_H4] = [128] */ + const __nv_bfloat16* __restrict__ vw5, /* [CRITIC_H5, CRITIC_H4] = [64, 128] */ + const __nv_bfloat16* __restrict__ vb5, /* [CRITIC_H5] = [64] */ + const __nv_bfloat16* __restrict__ vw6, /* [1, CRITIC_H5] = [1, 64] */ + const __nv_bfloat16* __restrict__ vb6, /* [1] */ float* scratch_a, /* [CRITIC_H1] = [512] ping-pong A */ float* scratch_b /* [CRITIC_H1] = [512] ping-pong B */ ) { @@ -240,37 +240,37 @@ __device__ void compute_gae_backward( */ extern "C" __global__ void ppo_full_experience_kernel( /* ---- Market data ---- */ - const float* __restrict__ market_features, /* [total_bars, MARKET_DIM] */ - const float* __restrict__ targets, /* [total_bars, 4] */ + const __nv_bfloat16* __restrict__ market_features, /* [total_bars, MARKET_DIM] */ + const __nv_bfloat16* __restrict__ targets, /* [total_bars, 4] */ const int* __restrict__ episode_starts, /* [N] */ /* ---- Actor weights (6 pointers) ---- */ - const float* __restrict__ pw1, /* [ACTOR_H1, STATE_DIM] = [128, 54] */ - const float* __restrict__ pb1, /* [ACTOR_H1] = [128] */ - const float* __restrict__ pw2, /* [ACTOR_H2, ACTOR_H1] = [64, 128] */ - const float* __restrict__ pb2, /* [ACTOR_H2] = [64] */ - const float* __restrict__ pw3, /* [NUM_ACTIONS, ACTOR_H2] = [45, 64] */ - const float* __restrict__ pb3, /* [NUM_ACTIONS] = [45] */ + const __nv_bfloat16* __restrict__ pw1, /* [ACTOR_H1, STATE_DIM] = [128, 54] */ + const __nv_bfloat16* __restrict__ pb1, /* [ACTOR_H1] = [128] */ + const __nv_bfloat16* __restrict__ pw2, /* [ACTOR_H2, ACTOR_H1] = [64, 128] */ + const __nv_bfloat16* __restrict__ pb2, /* [ACTOR_H2] = [64] */ + const __nv_bfloat16* __restrict__ pw3, /* [NUM_ACTIONS, ACTOR_H2] = [45, 64] */ + const __nv_bfloat16* __restrict__ pb3, /* [NUM_ACTIONS] = [45] */ /* ---- Critic weights (12 pointers) ---- */ - const float* __restrict__ vw1, /* [CRITIC_H1, STATE_DIM] = [512, 54] */ - const float* __restrict__ vb1, /* [CRITIC_H1] = [512] */ - const float* __restrict__ vw2, /* [CRITIC_H2, CRITIC_H1] = [384, 512] */ - const float* __restrict__ vb2, /* [CRITIC_H2] = [384] */ - const float* __restrict__ vw3, /* [CRITIC_H3, CRITIC_H2] = [256, 384] */ - const float* __restrict__ vb3, /* [CRITIC_H3] = [256] */ - const float* __restrict__ vw4, /* [CRITIC_H4, CRITIC_H3] = [128, 256] */ - const float* __restrict__ vb4, /* [CRITIC_H4] = [128] */ - const float* __restrict__ vw5, /* [CRITIC_H5, CRITIC_H4] = [64, 128] */ - const float* __restrict__ vb5, /* [CRITIC_H5] = [64] */ - const float* __restrict__ vw6, /* [1, CRITIC_H5] = [1, 64] */ - const float* __restrict__ vb6, /* [1] */ + const __nv_bfloat16* __restrict__ vw1, /* [CRITIC_H1, STATE_DIM] = [512, 54] */ + const __nv_bfloat16* __restrict__ vb1, /* [CRITIC_H1] = [512] */ + const __nv_bfloat16* __restrict__ vw2, /* [CRITIC_H2, CRITIC_H1] = [384, 512] */ + const __nv_bfloat16* __restrict__ vb2, /* [CRITIC_H2] = [384] */ + const __nv_bfloat16* __restrict__ vw3, /* [CRITIC_H3, CRITIC_H2] = [256, 384] */ + const __nv_bfloat16* __restrict__ vb3, /* [CRITIC_H3] = [256] */ + const __nv_bfloat16* __restrict__ vw4, /* [CRITIC_H4, CRITIC_H3] = [128, 256] */ + const __nv_bfloat16* __restrict__ vb4, /* [CRITIC_H4] = [128] */ + const __nv_bfloat16* __restrict__ vw5, /* [CRITIC_H5, CRITIC_H4] = [64, 128] */ + const __nv_bfloat16* __restrict__ vb5, /* [CRITIC_H5] = [64] */ + const __nv_bfloat16* __restrict__ vw6, /* [1, CRITIC_H5] = [1, 64] */ + const __nv_bfloat16* __restrict__ vb6, /* [1] */ /* ---- Curiosity model weights (4 pointers) ---- */ - const float* __restrict__ cur_w1, /* [CUR_HIDDEN, CUR_INPUT] */ - const float* __restrict__ cur_b1, /* [CUR_HIDDEN] */ - const float* __restrict__ cur_w2, /* [CUR_OUTPUT, CUR_HIDDEN] */ - const float* __restrict__ cur_b2, /* [CUR_OUTPUT] */ + const __nv_bfloat16* __restrict__ cur_w1, /* [CUR_HIDDEN, CUR_INPUT] */ + const __nv_bfloat16* __restrict__ cur_b1, /* [CUR_HIDDEN] */ + const __nv_bfloat16* __restrict__ cur_w2, /* [CUR_OUTPUT, CUR_HIDDEN] */ + const __nv_bfloat16* __restrict__ cur_b2, /* [CUR_OUTPUT] */ /* ---- Per-episode mutable state arrays ---- */ float* portfolio_states, /* [N, PORTFOLIO_STATE_SIZE] */ @@ -279,7 +279,7 @@ extern "C" __global__ void ppo_full_experience_kernel( int* diversity_metas, /* [N, 2] */ /* ---- Barrier config (shared) ---- */ - const float* __restrict__ barrier_config, /* [3]: profit_mult, loss_mult, max_bars */ + const __nv_bfloat16* __restrict__ barrier_config, /* [3]: profit_mult, loss_mult, max_bars */ /* ---- Scalar configs ---- */ float max_position, diff --git a/crates/ml/src/cuda_pipeline/q_stats_kernel.cu b/crates/ml/src/cuda_pipeline/q_stats_kernel.cu index fac96a783..8a9776b49 100644 --- a/crates/ml/src/cuda_pipeline/q_stats_kernel.cu +++ b/crates/ml/src/cuda_pipeline/q_stats_kernel.cu @@ -11,7 +11,7 @@ extern "C" __global__ void q_stats_reduce( const __nv_bfloat16* __restrict__ q_values, // [N, total_actions] - float* __restrict__ out, // [5]: avg_max_q, q_min, q_max, q_mean, q_var — stays F32 (monitoring scalars) + __nv_bfloat16* __restrict__ out, // [5]: avg_max_q, q_min, q_max, q_mean, q_var — stays F32 (monitoring scalars) int N, int total_actions) { @@ -46,9 +46,9 @@ extern "C" __global__ void q_stats_reduce( float variance = (total > 0) ? var_sum / (float)total : 0.0f; float avg_max = (N > 0) ? sum_max_q / (float)N : 0.0f; - out[0] = avg_max; - out[1] = global_min; - out[2] = global_max; - out[3] = mean; - out[4] = variance; + out[0] = __float2bfloat16(avg_max); + out[1] = __float2bfloat16(global_min); + out[2] = __float2bfloat16(global_max); + out[3] = __float2bfloat16(mean); + out[4] = __float2bfloat16(variance); } diff --git a/crates/ml/src/cuda_pipeline/relu_mask_kernel.cu b/crates/ml/src/cuda_pipeline/relu_mask_kernel.cu index fc3ad8eed..4e121b6c5 100644 --- a/crates/ml/src/cuda_pipeline/relu_mask_kernel.cu +++ b/crates/ml/src/cuda_pipeline/relu_mask_kernel.cu @@ -7,11 +7,11 @@ */ extern "C" __global__ -void relu_mask_standalone(float* __restrict__ dx, - const float* __restrict__ activation, +void relu_mask_standalone(__nv_bfloat16* __restrict__ dx, + const __nv_bfloat16* __restrict__ activation, int n) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= n) return; - if (activation[i] <= 0.0f) dx[i] = 0.0f; + if (__bfloat162float(activation[i]) <= 0.0f) dx[i] = __float2bfloat16(0.0f); } diff --git a/crates/ml/src/cuda_pipeline/signal_adapter.rs b/crates/ml/src/cuda_pipeline/signal_adapter.rs index 539f1a944..0d0644b18 100644 --- a/crates/ml/src/cuda_pipeline/signal_adapter.rs +++ b/crates/ml/src/cuda_pipeline/signal_adapter.rs @@ -2,7 +2,7 @@ //! Signal adapter utilities for converting between model outputs and action scores. //! -//! Pure cudarc `CudaSlice` operations -- zero Candle Tensor usage. +//! Pure cudarc `CudaSlice` operations -- zero Candle Tensor usage. //! Three CUDA kernels bridge the gap between model predictions and the 5-action //! DQN exposure space (Short100, Short50, Flat, Long50, Long100). @@ -55,20 +55,20 @@ fn load_kernels(context: &Arc) -> Result { /// * `stream` - CUDA stream for kernel launch and allocation /// /// # Returns -/// GPU-resident `CudaSlice` of shape `[batch * 5]`. +/// GPU-resident `CudaSlice` of shape `[batch * 5]`. /// /// # Errors /// Returns `MLError::ModelError` on kernel compilation or launch failure. pub fn ppo_to_exposure_scores( - probs: &CudaSlice, + probs: &CudaSlice, batch: usize, stream: &Arc, -) -> Result, MLError> { +) -> Result, MLError> { let context = stream.context(); let kernels = load_kernels(&context)?; let mut out = stream - .alloc_zeros::(batch * 5) + .alloc_zeros::(batch * 5) .map_err(|e| MLError::ModelError(format!("alloc exposure scores: {e}")))?; let batch_i32 = batch as i32; @@ -111,22 +111,22 @@ pub fn ppo_to_exposure_scores( /// * `stream` - CUDA stream for kernel launch and allocation /// /// # Returns -/// GPU-resident `CudaSlice` of shape `[batch * 5]` (one-hot rows). +/// GPU-resident `CudaSlice` of shape `[batch * 5]` (one-hot rows). /// /// # Errors /// Returns `MLError::ModelError` on kernel compilation or launch failure. pub fn signal_to_action_scores( - predictions: &CudaSlice, + predictions: &CudaSlice, batch: usize, high_threshold_bps: f32, low_threshold_bps: f32, stream: &Arc, -) -> Result, MLError> { +) -> Result, MLError> { let context = stream.context(); let kernels = load_kernels(&context)?; let mut out = stream - .alloc_zeros::(batch * 5) + .alloc_zeros::(batch * 5) .map_err(|e| MLError::ModelError(format!("alloc action scores: {e}")))?; let batch_i32 = batch as i32; @@ -167,17 +167,17 @@ pub fn signal_to_action_scores( /// * `stream` - CUDA stream for kernel launch and allocation /// /// # Returns -/// GPU-resident `CudaSlice` of shape `[batch]`. +/// GPU-resident `CudaSlice` of shape `[batch]`. /// /// # Errors /// Returns `MLError::ModelError` if dimensions are invalid or kernel fails. pub fn tft_quantile_to_signal( - quantiles: &CudaSlice, + quantiles: &CudaSlice, batch: usize, horizon: usize, num_quantiles: usize, stream: &Arc, -) -> Result, MLError> { +) -> Result, MLError> { if horizon < 1 { return Err(MLError::ModelError( "tft_quantile_to_signal: horizon must be >= 1".to_owned(), @@ -193,7 +193,7 @@ pub fn tft_quantile_to_signal( let kernels = load_kernels(&context)?; let mut out = stream - .alloc_zeros::(batch) + .alloc_zeros::(batch) .map_err(|e| MLError::ModelError(format!("alloc tft signal: {e}")))?; let batch_i32 = batch as i32; @@ -262,7 +262,7 @@ mod tests { let stream = cuda_stream(); let batch = 4; let uniform = vec![1.0_f32 / 45.0; batch * 45]; - let mut probs_buf = stream.alloc_zeros::(batch * 45).unwrap(); + let mut probs_buf = stream.alloc_zeros::(batch * 45).unwrap(); stream.memcpy_htod(&uniform, &mut probs_buf).unwrap(); let scores = ppo_to_exposure_scores(&probs_buf, batch, &stream).unwrap(); @@ -295,7 +295,7 @@ mod tests { } } } - let mut probs_buf = stream.alloc_zeros::(batch * 45).unwrap(); + let mut probs_buf = stream.alloc_zeros::(batch * 45).unwrap(); stream.memcpy_htod(&raw, &mut probs_buf).unwrap(); let scores = ppo_to_exposure_scores(&probs_buf, batch, &stream).unwrap(); @@ -336,7 +336,7 @@ mod tests { fn run_signal_test(pred_val: f32, high: f32, low: f32) -> Vec { let stream = cuda_stream(); - let mut pred_buf = stream.alloc_zeros::(1).unwrap(); + let mut pred_buf = stream.alloc_zeros::(1).unwrap(); stream.memcpy_htod(&[pred_val], &mut pred_buf).unwrap(); let scores = signal_to_action_scores(&pred_buf, 1, high, low, &stream).unwrap(); let mut host = vec![0.0_f32; 5]; @@ -385,7 +385,7 @@ mod tests { -1.0, 0.5, 2.0, // batch 0: median = 0.5 -3.0, -1.5, 0.0, // batch 1: median = -1.5 ]; - let mut q_buf = stream.alloc_zeros::(data.len()).unwrap(); + let mut q_buf = stream.alloc_zeros::(data.len()).unwrap(); stream.memcpy_htod(&data, &mut q_buf).unwrap(); let signal = tft_quantile_to_signal(&q_buf, 2, 1, 3, &stream).unwrap(); @@ -408,7 +408,7 @@ mod tests { #[test] fn test_tft_quantile_rejects_bad_dims() { let stream = cuda_stream(); - let mut q_buf = stream.alloc_zeros::(6).unwrap(); + let mut q_buf = stream.alloc_zeros::(6).unwrap(); stream.memcpy_htod(&[0.0_f32; 6], &mut q_buf).unwrap(); // horizon=0 should fail diff --git a/crates/ml/src/cuda_pipeline/signal_adapter_kernel.cu b/crates/ml/src/cuda_pipeline/signal_adapter_kernel.cu index da7939414..c29a0fd2b 100644 --- a/crates/ml/src/cuda_pipeline/signal_adapter_kernel.cu +++ b/crates/ml/src/cuda_pipeline/signal_adapter_kernel.cu @@ -24,8 +24,8 @@ */ extern "C" __global__ void ppo_to_exposure_scores_kernel( - const float* __restrict__ probs, /* [batch * 45] softmax probs */ - float* __restrict__ out_scores, /* [batch * 5] exposure scores */ + const __nv_bfloat16* __restrict__ probs, /* [batch * 45] softmax probs */ + __nv_bfloat16* __restrict__ out_scores, /* [batch * 5] exposure scores */ int batch ) { int b = blockIdx.x * blockDim.x + threadIdx.x; @@ -39,15 +39,15 @@ extern "C" __global__ void ppo_to_exposure_scores_kernel( float sum = 0.0f; int offset = base + e * 9; for (int j = 0; j < 9; ++j) { - sum += probs[offset + j]; + sum += __bfloat162float(probs[offset + j]); } - out_scores[b * 5 + e] = sum; + out_scores[b * 5 + e] = __float2bfloat16(sum); } } extern "C" __global__ void signal_to_action_scores_kernel( - const float* __restrict__ predictions, /* [batch] scalar signals */ - float* __restrict__ out_scores, /* [batch * 5] one-hot scores */ + const __nv_bfloat16* __restrict__ predictions, /* [batch] scalar signals */ + __nv_bfloat16* __restrict__ out_scores, /* [batch * 5] one-hot scores */ float high_threshold_bps, float low_threshold_bps, int batch @@ -55,32 +55,34 @@ extern "C" __global__ void signal_to_action_scores_kernel( int b = blockIdx.x * blockDim.x + threadIdx.x; if (b >= batch) return; - float pred = predictions[b]; + float pred = __bfloat162float(predictions[b]); int out_base = b * 5; /* Zero all 5 scores, then set the matching bucket to 1.0 */ - out_scores[out_base + 0] = 0.0f; - out_scores[out_base + 1] = 0.0f; - out_scores[out_base + 2] = 0.0f; - out_scores[out_base + 3] = 0.0f; - out_scores[out_base + 4] = 0.0f; + __nv_bfloat16 zero_bf = __float2bfloat16(0.0f); + __nv_bfloat16 one_bf = __float2bfloat16(1.0f); + out_scores[out_base + 0] = zero_bf; + out_scores[out_base + 1] = zero_bf; + out_scores[out_base + 2] = zero_bf; + out_scores[out_base + 3] = zero_bf; + out_scores[out_base + 4] = zero_bf; if (pred < -high_threshold_bps) { - out_scores[out_base + 0] = 1.0f; /* Short100 */ + out_scores[out_base + 0] = one_bf; /* Short100 */ } else if (pred < -low_threshold_bps) { - out_scores[out_base + 1] = 1.0f; /* Short50 */ + out_scores[out_base + 1] = one_bf; /* Short50 */ } else if (pred > high_threshold_bps) { - out_scores[out_base + 4] = 1.0f; /* Long100 */ + out_scores[out_base + 4] = one_bf; /* Long100 */ } else if (pred > low_threshold_bps) { - out_scores[out_base + 3] = 1.0f; /* Long50 */ + out_scores[out_base + 3] = one_bf; /* Long50 */ } else { - out_scores[out_base + 2] = 1.0f; /* Flat */ + out_scores[out_base + 2] = one_bf; /* Flat */ } } extern "C" __global__ void tft_quantile_extract_kernel( - const float* __restrict__ quantiles, /* [batch * horizon * num_q] */ - float* __restrict__ out_signal, /* [batch] median values */ + const __nv_bfloat16* __restrict__ quantiles, /* [batch * horizon * num_q] */ + __nv_bfloat16* __restrict__ out_signal, /* [batch] median values */ int horizon, int num_quantiles, int batch diff --git a/crates/ml/src/cuda_pipeline/statistics_kernel.cu b/crates/ml/src/cuda_pipeline/statistics_kernel.cu index e93f8ba85..745578822 100644 --- a/crates/ml/src/cuda_pipeline/statistics_kernel.cu +++ b/crates/ml/src/cuda_pipeline/statistics_kernel.cu @@ -21,9 +21,9 @@ * Grid-stride loop handles N > 256. */ extern "C" __global__ void batch_statistics( - const float* __restrict__ rewards, - const float* __restrict__ portfolio_values, - const float* __restrict__ positions, + const __nv_bfloat16* __restrict__ rewards, + const __nv_bfloat16* __restrict__ portfolio_values, + const __nv_bfloat16* __restrict__ positions, float* output_stats, int N ) { @@ -47,16 +47,16 @@ extern "C" __global__ void batch_statistics( float local_pos_sq_sum = 0.0f; for (int i = tid; i < N; i += blockDim.x) { - float r = rewards[i]; + float r = __bfloat162float(rewards[i]); local_sum += r; local_sq_sum += r * r; if (r > 0.0f) local_wins++; - float pv = portfolio_values[i]; + float pv = __bfloat162float(portfolio_values[i]); if (pv < local_min_val) local_min_val = pv; if (pv > local_max_val) local_max_val = pv; - float pos = fabsf(positions[i]); + float pos = fabsf(__bfloat162float(positions[i])); local_pos_sum += pos; local_pos_sq_sum += pos * pos; } diff --git a/crates/ml/src/cuda_pipeline/trade_stats_kernel.cu b/crates/ml/src/cuda_pipeline/trade_stats_kernel.cu index 91dcc33cf..ea5440ac8 100644 --- a/crates/ml/src/cuda_pipeline/trade_stats_kernel.cu +++ b/crates/ml/src/cuda_pipeline/trade_stats_kernel.cu @@ -12,7 +12,7 @@ #define PORTFOLIO_STRIDE 20 extern "C" __global__ void trade_stats_reduce( - const float* __restrict__ portfolio_states, // [N * PORTFOLIO_STRIDE] + const __nv_bfloat16* __restrict__ portfolio_states, // [N * PORTFOLIO_STRIDE] float* out, // [6]: sums of ps[14..19] int N ) { @@ -33,7 +33,7 @@ extern "C" __global__ void trade_stats_reduce( for (int i = tid; i < N; i += stride) { int base = i * PORTFOLIO_STRIDE; for (int k = 0; k < 6; k++) { - local_sums[k] += portfolio_states[base + 14 + k]; + local_sums[k] += __bfloat162float(portfolio_states[base + 14 + k]); } } diff --git a/crates/ml/src/cuda_pipeline/training_guard_kernel.cu b/crates/ml/src/cuda_pipeline/training_guard_kernel.cu index 0e31113af..fb0442e62 100644 --- a/crates/ml/src/cuda_pipeline/training_guard_kernel.cu +++ b/crates/ml/src/cuda_pipeline/training_guard_kernel.cu @@ -42,16 +42,16 @@ /* [2] step_count (stored as float) */ /* ------------------------------------------------------------------ */ extern "C" __global__ void training_guard_check_and_accumulate( - const float* __restrict__ loss_scalar, /* GPU-resident scalar (1 float) */ - const float* __restrict__ grad_norm_scalar, /* GPU-resident scalar (1 float) */ + const __nv_bfloat16* __restrict__ loss_scalar, /* GPU-resident scalar (1 float) */ + const __nv_bfloat16* __restrict__ grad_norm_scalar, /* GPU-resident scalar (1 float) */ float* output, /* pinned host buffer (7 floats) */ float* acc_buf, /* device accumulator (3 floats) */ float clip_threshold, /* e.g. 1e6f */ float collapse_threshold, /* e.g. 1e-7f */ int warmup /* 1 = in warmup, skip collapse */ ) { - float loss = loss_scalar[0]; - float grad_norm = grad_norm_scalar[0]; + float loss = __bfloat162float(loss_scalar[0]); + float grad_norm = __bfloat162float(grad_norm_scalar[0]); /* ── Guard check ── */ @@ -100,15 +100,15 @@ extern "C" __global__ void training_guard_check_and_accumulate( to the fused kernel's logic. New code should use the fused version. */ extern "C" __global__ void training_guard_check( - const float* __restrict__ loss_scalar, - const float* __restrict__ grad_norm_scalar, + const __nv_bfloat16* __restrict__ loss_scalar, + const __nv_bfloat16* __restrict__ grad_norm_scalar, float* output, float clip_threshold, float collapse_threshold, int warmup ) { - float loss = loss_scalar[0]; - float grad_norm = grad_norm_scalar[0]; + float loss = __bfloat162float(loss_scalar[0]); + float grad_norm = __bfloat162float(grad_norm_scalar[0]); int loss_nan = isnan(loss) || isinf(loss); int grad_nan = isnan(grad_norm) || isinf(grad_norm); int halt_nan = (loss_nan || grad_nan) ? 1 : 0; @@ -129,12 +129,12 @@ extern "C" __global__ void training_guard_check( } extern "C" __global__ void training_guard_accumulate( - const float* __restrict__ loss_scalar, - const float* __restrict__ grad_norm_scalar, + const __nv_bfloat16* __restrict__ loss_scalar, + const __nv_bfloat16* __restrict__ grad_norm_scalar, float* acc_buf ) { - float loss = loss_scalar[0]; - float grad_norm = grad_norm_scalar[0]; + float loss = __bfloat162float(loss_scalar[0]); + float grad_norm = __bfloat162float(grad_norm_scalar[0]); if (!isnan(loss) && !isinf(loss)) { acc_buf[0] += loss; } @@ -160,7 +160,7 @@ extern "C" __global__ void training_guard_accumulate( /* [3] mean_of_all — mean of ALL Q-values (collapse detection) */ /* ------------------------------------------------------------------ */ extern "C" __global__ void qvalue_stats_reduce( - const float* __restrict__ q_values, /* [batch_size * num_actions] */ + const __nv_bfloat16* __restrict__ q_values, /* [batch_size * num_actions] */ float* output, /* pinned host buffer (4 floats) */ int batch_size, int num_actions @@ -179,12 +179,12 @@ extern "C" __global__ void qvalue_stats_reduce( /* Grid-stride loop: each thread processes multiple samples */ for (int i = tid; i < batch_size; i += blockDim.x) { - const float* row = q_values + i * num_actions; + const __nv_bfloat16* row = q_values + i * num_actions; /* Find max Q across actions for this sample */ float sample_max = -1e30f; for (int a = 0; a < num_actions; a++) { - float qv = row[a]; + float qv = __bfloat162float(row[a]); if (!isnan(qv) && !isinf(qv)) { if (qv > sample_max) sample_max = qv; local_all += qv; @@ -266,7 +266,7 @@ extern "C" __global__ void qvalue_stats_reduce( /* [4] divergence_detected — 1.0 if |q_min| or |q_max| > threshold */ /* ------------------------------------------------------------------ */ extern "C" __global__ void qvalue_divergence_check( - const float* __restrict__ q_values, /* [num_actions] for one sample */ + const __nv_bfloat16* __restrict__ q_values, /* [num_actions] for one sample */ float* output, /* pinned host buffer (5 floats) */ int num_actions, float divergence_threshold /* e.g. 1e4f */ @@ -278,7 +278,7 @@ extern "C" __global__ void qvalue_divergence_check( int valid = 0; for (int a = 0; a < num_actions; a++) { - float qv = q_values[a]; + float qv = __bfloat162float(q_values[a]); if (isnan(qv) || isinf(qv)) continue; if (qv < q_min) q_min = qv; diff --git a/crates/ml/src/hyperopt/adapters/ppo.rs b/crates/ml/src/hyperopt/adapters/ppo.rs index 0f8251e92..17f0446f5 100644 --- a/crates/ml/src/hyperopt/adapters/ppo.rs +++ b/crates/ml/src/hyperopt/adapters/ppo.rs @@ -328,9 +328,9 @@ pub struct PPOTrainer { /// Spread width in ticks (default 1.0) spread_ticks: f64, /// GPU market features [num_bars * 42] — uploaded once, reused across trials - features_cuda: Option>, + features_cuda: Option>, /// GPU market targets [num_bars * 4] — uploaded once, reused across trials - targets_cuda: Option>, + targets_cuda: Option>, /// GPU PPO experience collector — initialized on first trial, weights synced per trial gpu_collector: Option, /// Number of bars in the uploaded market data @@ -1305,7 +1305,7 @@ impl PPOTrainer { // Forward function: PPO actor -> softmax -> ppo_to_exposure_scores -> [B, 5] // Uses host-side action_probabilities API, then maps to exposure scores on GPU. let metrics = evaluator.evaluate( - &|states_gpu: &CudaSlice, batch_size: usize, state_dim: usize| -> Result, MLError> { + &|states_gpu: &CudaSlice, batch_size: usize, state_dim: usize| -> Result, MLError> { // Download states to host for PPO forward pass let states_host = stream.clone_dtoh(states_gpu).map_err(|e| { // gpu-exit: PPO actor forward MLError::ModelError(format!("PPO backtest state download: {e}")) diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 2dd76f63a..1307f6971 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -420,7 +420,7 @@ impl DQNAgentType { let s_slice = states.data(); let ns_slice = next_states.data(); // actions are stored as f32 in GpuTensor but insert_batch wants u32. - // Reinterpret the CudaSlice as CudaSlice (bit-compatible for small ints). + // Reinterpret the CudaSlice as CudaSlice (bit-compatible for small ints). let a_f32_slice = actions.data(); let r_slice = rewards.data(); let d_slice = dones.data(); @@ -434,14 +434,14 @@ impl DQNAgentType { // Safety: f32 and u32 are same size. Action values [0..44] are // integers stored as f32 then bitwise reinterpreted as u32. let a_slice: &cudarc::driver::CudaSlice = unsafe { - &*(a_f32_slice as *const cudarc::driver::CudaSlice + &*(a_f32_slice as *const cudarc::driver::CudaSlice as *const cudarc::driver::CudaSlice) }; gpu_buf.gpu.insert_batch(s_slice, ns_slice, a_slice, r_slice, d_slice, batch_size) } Self::RegimeConditional(agent) => { let a_slice: &cudarc::driver::CudaSlice = unsafe { - &*(a_f32_slice as *const cudarc::driver::CudaSlice + &*(a_f32_slice as *const cudarc::driver::CudaSlice as *const cudarc::driver::CudaSlice) }; macro_rules! insert_head { diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 045f516ff..7bd1d9fed 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -93,7 +93,7 @@ pub(crate) struct FusedTrainingCtx { /// Combined loss: L_total = L_c51 + iqn_lambda × L_iqn. pub(crate) gpu_iqn: Option, /// CVaR scales buffer (kept alive so device pointer remains valid). - cvar_scales_buf: Option>, + cvar_scales_buf: Option>, /// GPU multi-head feature attention over h_s2 (post-graph, 1-step lag). /// When Some, applied after EMA update, before IQN, each training step. pub(crate) gpu_attention: Option, @@ -114,20 +114,20 @@ pub(crate) struct FusedTrainingCtx { /// Pre-allocated buffer: [K * B * num_atoms] for per-head value logits. /// Used to assemble logits from all K heads before diversity kernel. /// None when ensemble_count <= 1. - pub(crate) ensemble_logits_buf: Option>, + pub(crate) ensemble_logits_buf: Option>, /// Pre-allocated buffer: [B * total_actions] for mean Q-values across K heads. /// None when ensemble_count <= 1. - pub(crate) ensemble_mean_q_buf: Option>, + pub(crate) ensemble_mean_q_buf: Option>, /// Pre-allocated buffer: [B * total_actions] for Q-value variance across K heads. /// None when ensemble_count <= 1. - pub(crate) ensemble_var_q_buf: Option>, + pub(crate) ensemble_var_q_buf: Option>, /// Pre-allocated buffer: [1] for accumulated diversity loss scalar. /// None when ensemble_count <= 1. - pub(crate) ensemble_diversity_loss_buf: Option>, + pub(crate) ensemble_diversity_loss_buf: Option>, /// Compiled KL gradient kernel for diversity gradient flow. pub(crate) ensemble_kl_grad_kernel: Option, /// Pre-allocated buffer: [B * num_atoms] for diversity gradient on head 0 logits. - pub(crate) ensemble_d_logits_buf: Option>, + pub(crate) ensemble_d_logits_buf: Option>, } impl Drop for FusedTrainingCtx { @@ -381,13 +381,13 @@ impl FusedTrainingCtx { // ensemble_logits_buf: [K * B * num_atoms] — all K heads' value logits let na = dqn.config.num_atoms; let total_actions = dqn.config.num_actions + dqn.config.num_order_types + dqn.config.num_urgency_levels; - let logits_buf = stream.alloc_zeros::(k * batch_size * na) + let logits_buf = stream.alloc_zeros::(k * batch_size * na) .map_err(|e| anyhow::anyhow!("Alloc ensemble_logits_buf: {e}"))?; - let mean_q_buf = stream.alloc_zeros::(batch_size * total_actions) + let mean_q_buf = stream.alloc_zeros::(batch_size * total_actions) .map_err(|e| anyhow::anyhow!("Alloc ensemble_mean_q_buf: {e}"))?; - let var_q_buf = stream.alloc_zeros::(batch_size * total_actions) + let var_q_buf = stream.alloc_zeros::(batch_size * total_actions) .map_err(|e| anyhow::anyhow!("Alloc ensemble_var_q_buf: {e}"))?; - let div_loss_buf = stream.alloc_zeros::(1) + let div_loss_buf = stream.alloc_zeros::(1) .map_err(|e| anyhow::anyhow!("Alloc ensemble_diversity_loss_buf: {e}"))?; info!( @@ -408,7 +408,7 @@ impl FusedTrainingCtx { }; // d_logits buffer for diversity gradient [B * num_atoms] - let d_logits_buf = stream.alloc_zeros::(batch_size * dqn.config.num_atoms) + let d_logits_buf = stream.alloc_zeros::(batch_size * dqn.config.num_atoms) .map_err(|e| anyhow::anyhow!("alloc ensemble_d_logits: {e}"))?; ( @@ -806,7 +806,7 @@ impl FusedTrainingCtx { // Reinterpret the f32 CudaSlice as u32 to pass to the PER kernel. let idx_f32 = effective_gpu.indices.data(); let idx_u32: &cudarc::driver::CudaSlice = unsafe { - &*(idx_f32 as *const cudarc::driver::CudaSlice + &*(idx_f32 as *const cudarc::driver::CudaSlice as *const cudarc::driver::CudaSlice) }; self.trainer.update_priorities_cuda( @@ -887,7 +887,7 @@ impl FusedTrainingCtx { let head0_ptr = raw_device_ptr(self.trainer.on_v_logits_buf(), &self.stream); let dst_ptr = raw_device_ptr(logits_buf, &self.stream); let n_bytes = b * na * f32_size; - // Safety: both are valid CudaSlice on the same context. Byte sizes match. + // Safety: both are valid CudaSlice on the same context. Byte sizes match. unsafe { cudarc::driver::result::memcpy_dtod_async( dst_ptr, head0_ptr, n_bytes, self.stream.cu_stream() @@ -1135,7 +1135,7 @@ impl FusedTrainingCtx { /// Compute Q-value statistics entirely on GPU — 20-byte readback (5 scalars). pub(crate) fn compute_q_stats( &mut self, - states: &cudarc::driver::CudaSlice, + states: &cudarc::driver::CudaSlice, batch_size: usize, ) -> Result { self.trainer.compute_q_stats(states, batch_size) @@ -1148,9 +1148,9 @@ impl FusedTrainingCtx { /// No DtoH transfer — caller decides how to consume the CudaSlice. pub(crate) fn compute_q_values( &mut self, - states: &cudarc::driver::CudaSlice, + states: &cudarc::driver::CudaSlice, batch_size: usize, - ) -> Result<&cudarc::driver::CudaSlice> { + ) -> Result<&cudarc::driver::CudaSlice> { self.trainer.compute_q_values(states, batch_size) .map_err(|e| anyhow::anyhow!("compute_q_values: {e}")) } @@ -1444,9 +1444,9 @@ fn clone_dueling_weights( stream: &Arc, head_idx: usize, ) -> Result { - let clone_f32 = |slice: &cudarc::driver::CudaSlice| -> Result> { + let clone_f32 = |slice: &cudarc::driver::CudaSlice| -> Result> { let n = slice.len(); - let dst = stream.alloc_zeros::(n) + let dst = stream.alloc_zeros::(n) .map_err(|e| anyhow::anyhow!("Clone alloc {n}xf32: {e}"))?; let (src_ptr, src_guard) = slice.device_ptr(stream); let (dst_ptr, dst_guard) = dst.device_ptr(stream); @@ -1489,9 +1489,9 @@ fn clone_branching_weights( stream: &Arc, head_idx: usize, ) -> Result { - let clone_f32 = |slice: &cudarc::driver::CudaSlice| -> Result> { + let clone_f32 = |slice: &cudarc::driver::CudaSlice| -> Result> { let n = slice.len(); - let dst = stream.alloc_zeros::(n) + let dst = stream.alloc_zeros::(n) .map_err(|e| anyhow::anyhow!("Clone alloc {n}xf32: {e}"))?; let (src_ptr, src_guard) = slice.device_ptr(stream); let (dst_ptr, dst_guard) = dst.device_ptr(stream); diff --git a/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs b/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs index f23431bfd..3848c1289 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs @@ -32,7 +32,7 @@ fn smoke_stream() -> Arc { } /// Create random f32 CudaSlice on GPU (host-generated, uploaded). -fn random_cuda_f32(n: usize, stream: &Arc) -> anyhow::Result> { +fn random_cuda_f32(n: usize, stream: &Arc) -> anyhow::Result> { use rand::Rng; let mut rng = rand::thread_rng(); let host: Vec = (0..n).map(|_| rng.gen_range(-1.0_f32..1.0)).collect(); @@ -47,8 +47,8 @@ fn zeros_cuda_u32(n: usize, stream: &Arc) -> anyhow: } /// Create zeroed f32 CudaSlice on GPU. -fn zeros_cuda_f32(n: usize, stream: &Arc) -> anyhow::Result> { - let slice = stream.alloc_zeros::(n).map_err(|e| anyhow::anyhow!("alloc f32: {e}"))?; +fn zeros_cuda_f32(n: usize, stream: &Arc) -> anyhow::Result> { + let slice = stream.alloc_zeros::(n).map_err(|e| anyhow::anyhow!("alloc f32: {e}"))?; Ok(slice) } @@ -68,8 +68,8 @@ fn insert_random_batch( Ok(()) } -/// Download a CudaSlice to host for assertion. -fn dtoh_f32(slice: &cudarc::driver::CudaSlice, stream: &Arc) -> anyhow::Result> { +/// Download a CudaSlice to host for assertion. +fn dtoh_f32(slice: &cudarc::driver::CudaSlice, stream: &Arc) -> anyhow::Result> { let n = slice.len(); let mut host = vec![0.0_f32; n]; stream.memcpy_dtoh(slice, &mut host).map_err(|e| anyhow::anyhow!("dtoh f32: {e}"))?; // test-only readback diff --git a/crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs b/crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs index 0688f7000..28c050ef0 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs @@ -36,7 +36,7 @@ fn test_config() -> GpuDqnTrainConfig { /// needs non-degenerate matrices for power iteration to converge). fn alloc_dueling(stream: &Arc, cfg: &GpuDqnTrainConfig) -> DuelingWeightSet { let na = cfg.num_atoms; - let alloc = |n: usize| -> cudarc::driver::CudaSlice { + let alloc = |n: usize| -> cudarc::driver::CudaSlice { // Fill with 1.0 (non-zero, non-degenerate for spectral norm) stream.clone_htod(&vec![0.1_f32; n]).expect("alloc dueling weight") }; @@ -58,7 +58,7 @@ fn alloc_dueling(stream: &Arc, cfg: &GpuDqnTrainConf fn alloc_branching(stream: &Arc, cfg: &GpuDqnTrainConfig) -> BranchingWeightSet { let na = cfg.num_atoms; - let alloc = |n: usize| -> cudarc::driver::CudaSlice { + let alloc = |n: usize| -> cudarc::driver::CudaSlice { stream.clone_htod(&vec![0.1_f32; n]).expect("alloc branching weight") }; BranchingWeightSet { @@ -188,7 +188,7 @@ fn test_spectral_norm_constrains_operator_norm() -> anyhow::Result<()> { let mut trainer = GpuDqnTrainer::new(stream.clone(), cfg.clone())?; // Create weights with large values — spectral norm should constrain them - let alloc_large = |n: usize| -> cudarc::driver::CudaSlice { + let alloc_large = |n: usize| -> cudarc::driver::CudaSlice { stream.clone_htod(&vec![5.0_f32; n]).expect("alloc large weight") }; let na = cfg.num_atoms; diff --git a/crates/ml/src/trainers/dqn/smoke_tests/performance.rs b/crates/ml/src/trainers/dqn/smoke_tests/performance.rs index 0d88cc643..ebbed8790 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/performance.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/performance.rs @@ -87,7 +87,7 @@ fn test_per_sample_latency() -> anyhow::Result<()> { let next_states = stream.clone_htod(&host_states).map_err(|e| anyhow::anyhow!("{e}"))?; let actions = stream.alloc_zeros::(batch_insert).map_err(|e| anyhow::anyhow!("{e}"))?; let rewards = stream.clone_htod(&host_rewards).map_err(|e| anyhow::anyhow!("{e}"))?; - let dones = stream.alloc_zeros::(batch_insert).map_err(|e| anyhow::anyhow!("{e}"))?; + let dones = stream.alloc_zeros::(batch_insert).map_err(|e| anyhow::anyhow!("{e}"))?; buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, batch_insert)?; } assert_eq!(buf.len(), fill_count); diff --git a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs index 4ed248d32..a4caf4aa0 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs @@ -110,7 +110,7 @@ fn test_per_weights_valid() -> anyhow::Result<()> { let next_states = stream.clone_htod(&vec![0.0_f32; 45]).map_err(|e| anyhow::anyhow!("{e}"))?; let actions = stream.alloc_zeros::(1).map_err(|e| anyhow::anyhow!("{e}"))?; let rewards = stream.clone_htod(&[0.5_f32]).map_err(|e| anyhow::anyhow!("{e}"))?; - let dones = stream.alloc_zeros::(1).map_err(|e| anyhow::anyhow!("{e}"))?; + let dones = stream.alloc_zeros::(1).map_err(|e| anyhow::anyhow!("{e}"))?; buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, 1)?; } @@ -154,7 +154,7 @@ fn test_per_indices_valid() -> anyhow::Result<()> { let next_states = stream.clone_htod(&vec![0.0_f32; 45]).map_err(|e| anyhow::anyhow!("{e}"))?; let actions = stream.alloc_zeros::(1).map_err(|e| anyhow::anyhow!("{e}"))?; let rewards = stream.clone_htod(&[0.5_f32]).map_err(|e| anyhow::anyhow!("{e}"))?; - let dones = stream.alloc_zeros::(1).map_err(|e| anyhow::anyhow!("{e}"))?; + let dones = stream.alloc_zeros::(1).map_err(|e| anyhow::anyhow!("{e}"))?; buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, 1)?; } diff --git a/crates/ml/src/trainers/dqn/trainer/action.rs b/crates/ml/src/trainers/dqn/trainer/action.rs index 800335cba..1e4fc4072 100644 --- a/crates/ml/src/trainers/dqn/trainer/action.rs +++ b/crates/ml/src/trainers/dqn/trainer/action.rs @@ -10,8 +10,8 @@ use crate::dqn::TradingState; use crate::dqn::order_router::OrderRouter; use ml_core::fill_simulator::FillResult; -/// Extract the raw `&CudaSlice` from a `GpuTensor` reference. -/// GpuTensor stores data as `CudaSlice` directly — no conversion needed. +/// Extract the raw `&CudaSlice` from a `GpuTensor` reference. +/// GpuTensor stores data as `CudaSlice` directly — no conversion needed. macro_rules! extract_cuda_f32 { ($tensor:expr) => {{ $tensor.data() diff --git a/crates/ml/src/trainers/dqn/trainer/metrics.rs b/crates/ml/src/trainers/dqn/trainer/metrics.rs index 40afcf874..5e4466806 100644 --- a/crates/ml/src/trainers/dqn/trainer/metrics.rs +++ b/crates/ml/src/trainers/dqn/trainer/metrics.rs @@ -254,7 +254,7 @@ impl DQNTrainer { Err(_) => return None, }; - // GPU PER path: states are already a CudaSlice + // GPU PER path: states are already a CudaSlice let gpu_batch = batch_sample.gpu_batch.as_ref()?; let states_data = gpu_batch.states.data(); @@ -389,7 +389,7 @@ impl DQNTrainer { /// Compute validation loss (negative Sharpe proxy) using cuBLAS forward — zero GpuTensor. /// /// Builds the state tensor on CPU (features + portfolio broadcast + OFI + alignment - /// padding), uploads it once as a raw `CudaSlice`, runs cuBLAS forward via + /// padding), uploads it once as a raw `CudaSlice`, runs cuBLAS forward via /// `fused_ctx`, then downloads Q-values and performs all action selection and Sharpe /// computation on CPU. This avoids GpuTensor elementwise kernels (add, broadcast_mul, /// cat, etc.) that cache `CudaFunction` handles on the shared CUDA primary context, @@ -505,7 +505,7 @@ impl DQNTrainer { state_flat.extend(std::iter::repeat(0.0_f32).take(pad_cols)); } - // Upload state to GPU as raw CudaSlice — no GpuTensor, no elementwise kernels. + // Upload state to GPU as raw CudaSlice — no GpuTensor, no elementwise kernels. let state_gpu = stream.clone_htod(&state_flat) .map_err(|e| anyhow::anyhow!("val state HtoD: {e}"))?; diff --git a/crates/ml/src/trainers/dqn/trainer/mod.rs b/crates/ml/src/trainers/dqn/trainer/mod.rs index e81acba84..d4c0b181c 100644 --- a/crates/ml/src/trainers/dqn/trainer/mod.rs +++ b/crates/ml/src/trainers/dqn/trainer/mod.rs @@ -219,10 +219,10 @@ pub struct DQNTrainer { pub(crate) gpu_data: Option, /// Raw cudarc targets buffer for CUDA kernel (parallel to raw features buffer) - pub(crate) targets_raw_cuda: Option>, + pub(crate) targets_raw_cuda: Option>, /// Raw cudarc features buffer for CUDA experience kernel [num_bars * 42] - pub(crate) features_raw_cuda: Option>, + pub(crate) features_raw_cuda: Option>, /// GPU experience collector for zero-roundtrip CUDA kernel (Phase 2b) pub(crate) gpu_experience_collector: Option, diff --git a/crates/ml/src/trainers/dqn/trainer/train_step.rs b/crates/ml/src/trainers/dqn/trainer/train_step.rs index cfba37f63..95b1ff766 100644 --- a/crates/ml/src/trainers/dqn/trainer/train_step.rs +++ b/crates/ml/src/trainers/dqn/trainer/train_step.rs @@ -177,7 +177,7 @@ impl DQNTrainer { .sample(sample_size) .map_err(|e| anyhow::anyhow!("Q-est sample: {e}"))?; - // GPU PER path: states are already a CudaSlice + // GPU PER path: states are already a CudaSlice let gpu_batch = batch_sample.gpu_batch.as_ref() .ok_or_else(|| anyhow::anyhow!("GPU PER must be active — gpu_batch is None"))?; let states_data = gpu_batch.states.data(); @@ -471,7 +471,7 @@ impl DQNTrainer { .sample(sample_size) .map_err(|e| anyhow::anyhow!("Q-est sample: {e}"))?; - // GPU PER path: states are already a CudaSlice + // GPU PER path: states are already a CudaSlice let gpu_batch = batch_sample.gpu_batch.as_ref() .ok_or_else(|| anyhow::anyhow!("GPU PER must be active — gpu_batch is None"))?; let states_data = gpu_batch.states.data(); diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index a4525962c..15bc021e1 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1119,7 +1119,7 @@ impl DQNTrainer { &*(&gpu_batch.actions as *const CudaSlice as *const CudaSlice) }; - // dones: already CudaSlice (kernel outputs float 0.0/1.0) + // dones: already CudaSlice (kernel outputs float 0.0/1.0) let agent = self.agent.read().await; let mut gpu_buf = agent.memory().as_gpu_buffer() .ok_or_else(|| anyhow::anyhow!("GPU PER buffer required"))?; @@ -1668,8 +1668,8 @@ impl DQNTrainer { let state_dim_val = flat_states.len() / valid_indices.len(); let truncated_states = &flat_states[..batch_size_refresh * state_dim_val]; - // Upload states to GPU as CudaSlice - let mut states_gpu = refresh_stream.alloc_zeros::(batch_size_refresh * state_dim_val) + // Upload states to GPU as CudaSlice + let mut states_gpu = refresh_stream.alloc_zeros::(batch_size_refresh * state_dim_val) .map_err(|e| anyhow::anyhow!("PER refresh alloc: {e}"))?; refresh_stream.memcpy_htod(truncated_states, &mut states_gpu) .map_err(|e| anyhow::anyhow!("PER refresh HtoD: {e}"))?; diff --git a/crates/ml/src/trainers/ppo.rs b/crates/ml/src/trainers/ppo.rs index 665145268..ef852bd03 100644 --- a/crates/ml/src/trainers/ppo.rs +++ b/crates/ml/src/trainers/ppo.rs @@ -221,9 +221,9 @@ pub struct PpoTrainer { explained_variance_history: Arc>>, gpu_ppo_collector: Option, /// Raw cudarc features buffer for GPU experience kernel [num_bars * 51] - features_raw_cuda: Option>, + features_raw_cuda: Option>, /// Raw cudarc targets buffer for GPU experience kernel [num_bars * 4] - targets_raw_cuda: Option>, + targets_raw_cuda: Option>, /// Number of bars in the raw data buffers (needed to configure kernel) raw_data_num_bars: usize, }