From 2e06d86e4335c563a43ffca732598e4cfb245aa0 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 4 Apr 2026 11:18:29 +0200 Subject: [PATCH] fix: replace aliased SAXPY with in-place scale kernel (fixes __restrict__ NaN bug) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The C51+MSE gradient blending passed the same pointer as both y (output) and x (input) to dqn_saxpy_f32_kernel. That kernel has __restrict__ on both parameters, promising the compiler they don't alias. When they DO alias, NVCC's no-alias optimizations produce undefined behavior (NaN from register corruption). Added dqn_scale_f32_kernel (y[i] *= alpha) — single-pointer variant with no __restrict__, no aliasing possible. The two self-aliased SAXPY calls are replaced with scale+SAXPY pairs using distinct pointers. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../src/cuda_pipeline/dqn_utility_kernels.cu | 22 +++++++++++ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 37 ++++++++++--------- 2 files changed, 42 insertions(+), 17 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index 3ba373c53..e26c90966 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -165,6 +165,28 @@ extern "C" __global__ void dqn_saxpy_f32_kernel( if (i < n) y[i] = y[i] + alpha * x[i]; } +/* ══════════════════════════════════════════════════════════════════════ + * F32 IN-PLACE SCALE KERNEL + * + * y[i] *= alpha for i = 0..n-1 + * + * Single-pointer variant — no aliasing possible, no __restrict__ needed. + * Used for gradient scaling where SAXPY y+=α*x with x==y violates + * __restrict__ and causes undefined behavior (NaN from register + * corruption under NVCC's no-alias optimizations). + * + * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1). + * ══════════════════════════════════════════════════════════════════════ */ + +extern "C" __global__ void dqn_scale_f32_kernel( + float* y, + float alpha, + int n +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) y[i] *= alpha; +} + /* ══════════════════════════════════════════════════════════════════════ * CLIPPED SAXPY KERNEL * diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 655714c8c..16e31ad9e 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -497,6 +497,7 @@ pub struct GpuDqnTrainer { bf16_to_f32_kernel: CudaFunction, saxpy_kernel: CudaFunction, saxpy_f32_kernel: CudaFunction, + scale_f32_kernel: CudaFunction, zero_kernel: CudaFunction, regime_scale_kernel: CudaFunction, shrink_perturb_kernel: CudaFunction, @@ -2166,7 +2167,7 @@ impl GpuDqnTrainer { // per array. Stack is set once in DQNTrainer::new() (64KB for all kernels). // ── Compile 4 utility kernels (grad_norm, adam_update, BF16 converters) ─ - let (grad_norm_kernel, grad_norm_finalize_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clipped_saxpy_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, causal_intervene_kernel_fn, causal_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel) = + let (grad_norm_kernel, grad_norm_finalize_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clipped_saxpy_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, causal_intervene_kernel_fn, causal_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel) = compile_training_kernels(&stream, &config)?; // Separate grad_norm instance for non-graph launches (clip_grad_buf_inplace). @@ -2666,6 +2667,7 @@ impl GpuDqnTrainer { bf16_to_f32_kernel, saxpy_kernel, saxpy_f32_kernel, + scale_f32_kernel, zero_kernel, regime_scale_kernel, shrink_perturb_kernel: shrink_perturb, @@ -4093,7 +4095,6 @@ impl GpuDqnTrainer { let b2 = self.config.branch_2_size; let n_val = (b * na) as i32; let n_adv = (b * (b0 + b1 + b2) * na) as i32; - let scale_c51 = alpha - 1.0; // SAXPY: y += (α-1)*y → y *= α let scale_mse = 1.0 - alpha; let val_blocks = ((b * na + 255) / 256) as u32; let adv_blocks = ((b * (b0 + b1 + b2) * na + 255) / 256) as u32; @@ -4107,22 +4108,22 @@ impl GpuDqnTrainer { let adv_mse_ptr = self.d_adv_logits_mse.raw_ptr(); unsafe { - // d_value += (α-1) * d_value → d_value *= α (f32 SAXPY) - self.stream.launch_builder(&self.saxpy_f32_kernel) - .arg(&val_ptr).arg(&val_ptr) - .arg(&scale_c51).arg(&n_val) - .launch(cfg_val).map_err(|e| MLError::ModelError(format!("blend c51 val: {e}")))?; - // d_value += (1-α) * mse_scratch + // d_value *= α (in-place scale — single pointer, no aliasing) + self.stream.launch_builder(&self.scale_f32_kernel) + .arg(&val_ptr) + .arg(&alpha).arg(&n_val) + .launch(cfg_val).map_err(|e| MLError::ModelError(format!("scale c51 val: {e}")))?; + // d_value += (1-α) * mse_scratch (SAXPY with distinct pointers — safe) self.stream.launch_builder(&self.saxpy_f32_kernel) .arg(&val_ptr).arg(&val_mse_ptr) .arg(&scale_mse).arg(&n_val) .launch(cfg_val).map_err(|e| MLError::ModelError(format!("blend mse val: {e}")))?; - // d_adv += (α-1) * d_adv → d_adv *= α - self.stream.launch_builder(&self.saxpy_f32_kernel) - .arg(&adv_ptr).arg(&adv_ptr) - .arg(&scale_c51).arg(&n_adv) - .launch(cfg_adv).map_err(|e| MLError::ModelError(format!("blend c51 adv: {e}")))?; - // d_adv += (1-α) * mse_scratch + // d_adv *= α (in-place scale — single pointer, no aliasing) + self.stream.launch_builder(&self.scale_f32_kernel) + .arg(&adv_ptr) + .arg(&alpha).arg(&n_adv) + .launch(cfg_adv).map_err(|e| MLError::ModelError(format!("scale c51 adv: {e}")))?; + // d_adv += (1-α) * mse_scratch (SAXPY with distinct pointers — safe) self.stream.launch_builder(&self.saxpy_f32_kernel) .arg(&adv_ptr).arg(&adv_mse_ptr) .arg(&scale_mse).arg(&n_adv) @@ -5567,7 +5568,7 @@ impl GpuDqnTrainer { fn compile_training_kernels( stream: &Arc, config: &GpuDqnTrainConfig, -) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { +) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { info!( state_dim = config.state_dim, total_params = compute_total_params(config), @@ -5593,6 +5594,8 @@ fn compile_training_kernels( .map_err(|e| MLError::ModelError(format!("dqn_saxpy_kernel load: {e}")))?; let saxpy_f32 = module.load_function("dqn_saxpy_f32_kernel") .map_err(|e| MLError::ModelError(format!("dqn_saxpy_f32_kernel load: {e}")))?; + let scale_f32 = module.load_function("dqn_scale_f32_kernel") + .map_err(|e| MLError::ModelError(format!("dqn_scale_f32_kernel load: {e}")))?; let zero = module.load_function("dqn_zero_kernel") .map_err(|e| MLError::ModelError(format!("dqn_zero_kernel load: {e}")))?; let regime_scale = module.load_function("dqn_regime_scale_kernel") @@ -5635,8 +5638,8 @@ fn compile_training_kernels( let her_inplace = module.load_function("her_inplace_relabel") .map_err(|e| MLError::ModelError(format!("her_inplace_relabel load: {e}")))?; - info!("GpuDqnTrainer: 27 utility kernels loaded from precompiled cubin"); - Ok((grad_norm, grad_norm_finalize, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clipped_saxpy, clip_grad, pad_states, saxpy_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, causal_intervene, causal_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace)) + info!("GpuDqnTrainer: 28 utility kernels loaded from precompiled cubin"); + Ok((grad_norm, grad_norm_finalize, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clipped_saxpy, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, causal_intervene, causal_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace)) } /// Load the standalone Polyak EMA kernel from precompiled cubin.