From e328bcdb138c25e129b5db202463597204bcda7a Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 15 Apr 2026 09:19:23 +0200 Subject: [PATCH] =?UTF-8?q?feat:=20diffusion=20Q-refinement=20=E2=80=94=20?= =?UTF-8?q?3-step=20denoiser=20conditioned=20on=20Var[Q]=20(2,700=20params?= =?UTF-8?q?)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 156 +++++++++++++++++- crates/ml/src/trainers/dqn/fused_training.rs | 7 + 2 files changed, 157 insertions(+), 6 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index ab6778b81..ab139043e 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -393,11 +393,11 @@ pub(crate) fn compute_param_sizes(cfg: &GpuDqnTrainConfig) -> [usize; NUM_WEIGHT cfg.adv_h, // [13] b_b1fc cfg.branch_1_size * cfg.num_atoms * cfg.adv_h, // [14] w_b1out cfg.branch_1_size * cfg.num_atoms, // [15] b_b1out - cfg.adv_h * cfg.shared_h2, // [16] w_b2fc + cfg.adv_h * (cfg.shared_h2 + 3), // [16] w_b2fc (OFI-conditioned) cfg.adv_h, // [17] b_b2fc cfg.branch_2_size * cfg.num_atoms * cfg.adv_h, // [18] w_b2out cfg.branch_2_size * cfg.num_atoms, // [19] b_b2out - cfg.adv_h * cfg.shared_h2, // [20] w_b3fc + cfg.adv_h * (cfg.shared_h2 + 3), // [20] w_b3fc (OFI-conditioned) cfg.adv_h, // [21] b_b3fc cfg.branch_3_size * cfg.num_atoms * cfg.adv_h, // [22] w_b3out cfg.branch_3_size * cfg.num_atoms, // [23] b_b3out @@ -418,9 +418,9 @@ pub(crate) fn compute_param_sizes(cfg: &GpuDqnTrainConfig) -> [usize; NUM_WEIGHT cfg.adv_h, // [35] b_gate_0 cfg.adv_h * (cfg.shared_h2 + 3), // [36] w_gate_1 [AH, SH2+3] cfg.adv_h, // [37] b_gate_1 - cfg.adv_h * cfg.shared_h2, // [38] w_gate_2 + cfg.adv_h * (cfg.shared_h2 + 3), // [38] w_gate_2 (OFI-conditioned) cfg.adv_h, // [39] b_gate_2 - cfg.adv_h * cfg.shared_h2, // [40] w_gate_3 + cfg.adv_h * (cfg.shared_h2 + 3), // [40] w_gate_3 (OFI-conditioned) cfg.adv_h, // [41] b_gate_3 // ── KAN spline coefficients (8 bases per neuron) + residual weight ── cfg.adv_h * 8, // [42] kan_coeff_0 [AH, 8] @@ -488,6 +488,8 @@ struct CachedPtrs { save_h_b3: u64, mag_concat_buf: u64, d_mag_concat_buf: u64, + pub(crate) ord_concat_buf: u64, + pub(crate) urg_concat_buf: u64, bw_d_h_s1: u64, bw_d_h_s2: u64, bw_d_h_v: u64, @@ -1043,6 +1045,24 @@ pub struct GpuDqnTrainer { graph_adam_step: i32, /// branch_graph_message_pass kernel handle. graph_msg_kernel: CudaFunction, + + // ── Diffusion Q-refinement (3-step denoiser conditioned on Var[Q]) ─ + /// Flat params for 3 denoising steps: 3 × (W1[24,24]+b1[24]+W2[12,24]+b2[12]) = 3 × 900 = 2700. + denoise_params: CudaSlice, + /// Adam first moment for denoise_params [2700]. + denoise_adam_m: CudaSlice, + /// Adam second moment for denoise_params [2700]. + denoise_adam_v: CudaSlice, + /// Adam step counter for denoise optimizer. + denoise_adam_step: i32, + /// Ping-pong buffer A for iterative denoising [B, 12]. + denoise_buf_a: CudaSlice, + /// Ping-pong buffer B for iterative denoising [B, 12]. + denoise_buf_b: CudaSlice, + /// Var[Q] buffer for training path [B, 12] — populated by compute_expected_q. + q_var_buf_trainer: CudaSlice, + /// q_denoise_step kernel handle. + q_denoise_kernel: CudaFunction, } impl GpuDqnTrainer { @@ -1361,6 +1381,71 @@ impl GpuDqnTrainer { Ok(()) } + /// Run 3-step diffusion denoising on q_coord_buf [batch_size, 12] conditioned on Var[Q]. + /// + /// Each step: input = [Q_prev; sqrt(Var[Q]) * schedule] → FC(24→24)-SiLU-FC(24→12) → Q_next. + /// Ping-pong: a→b (k=1), b→a (k=2), a→b (k=3). Result (in buf_b) copied back to q_coord_buf. + /// Must be called AFTER launch_graph_message_pass. Reads q_var_buf_trainer (populated by + /// compute_expected_q with the training path's variance output). + pub(crate) fn launch_q_denoise(&self, batch_size: usize) -> Result<(), MLError> { + let b = batch_size as i32; + let blocks = ((batch_size as u32 + 255) / 256).max(1); + let copy_bytes = batch_size * 12 * std::mem::size_of::(); + let q_coord_ptr = self.q_coord_buf.raw_ptr(); + let buf_a = self.denoise_buf_a.raw_ptr(); + let buf_b = self.denoise_buf_b.raw_ptr(); + let var_ptr = self.q_var_buf_trainer.raw_ptr(); + let params_ptr = self.denoise_params.raw_ptr(); + + // Copy q_coord_buf → denoise_buf_a as the denoiser starting point. + unsafe { + cudarc::driver::result::memcpy_dtod_async(buf_a, q_coord_ptr, copy_bytes, self.stream.cu_stream()) + .map_err(|e| MLError::ModelError(format!("denoise init copy: {e}")))?; + } + + // Layout per step: W1[24,24]=576 + b1[24]=24 + W2[12,24]=288 + b2[12]=12 = 900 floats. + const STEP_PARAMS: usize = 900; + // 3 ping-pong steps: k=1(a→b), k=2(b→a), k=3(a→b). After step 3, result is in buf_b. + for k in 0..3_i32 { + let (in_ptr, out_ptr) = if k % 2 == 0 { (buf_a, buf_b) } else { (buf_b, buf_a) }; + let step_off = k as usize * STEP_PARAMS; + let w1_ptr = params_ptr + (step_off * std::mem::size_of::()) as u64; + let b1_ptr = w1_ptr + (576 * std::mem::size_of::()) as u64; + let w2_ptr = b1_ptr + (24 * std::mem::size_of::()) as u64; + let b2_ptr = w2_ptr + (288 * std::mem::size_of::()) as u64; + let step_k = k + 1_i32; + let total_steps = 3_i32; + unsafe { + self.stream + .launch_builder(&self.q_denoise_kernel) + .arg(&in_ptr) + .arg(&out_ptr) + .arg(&var_ptr) + .arg(&w1_ptr) + .arg(&b1_ptr) + .arg(&w2_ptr) + .arg(&b2_ptr) + .arg(&b) + .arg(&step_k) + .arg(&total_steps) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("q_denoise_step k={}: {e}", k + 1)))?; + } + } + + // After 3 steps (k=0,1,2), the last output is buf_b. Copy buf_b → q_coord_buf. + unsafe { + cudarc::driver::result::memcpy_dtod_async(q_coord_ptr, buf_b, copy_bytes, self.stream.cu_stream()) + .map_err(|e| MLError::ModelError(format!("denoise final copy: {e}")))?; + } + + Ok(()) + } + /// Center Q-values in q_out_buf by subtracting their global mean. /// /// Two-phase GPU operation (no CPU roundtrip): @@ -3455,6 +3540,54 @@ impl GpuDqnTrainer { super::htod_f32(&stream, &graph_params_host, &mut graph_params)?; info!("GpuDqnTrainer: branch_graph_message_pass kernel loaded (60 params, near-identity W_msg init)"); + // ── Diffusion Q-refinement ──────────────────────────────────────────────── + // 3 steps × (W1[24,24]=576 + b1[24]=24 + W2[12,24]=288 + b2[12]=12) = 3 × 900 = 2700 floats. + const DENOISE_PARAMS_PER_STEP: usize = 576 + 24 + 288 + 12; // 900 + const DENOISE_TOTAL_PARAMS: usize = 3 * DENOISE_PARAMS_PER_STEP; // 2700 + let denoise_adam_m = stream.alloc_zeros::(DENOISE_TOTAL_PARAMS) + .map_err(|e| MLError::ModelError(format!("alloc denoise_adam_m: {e}")))?; + let denoise_adam_v = stream.alloc_zeros::(DENOISE_TOTAL_PARAMS) + .map_err(|e| MLError::ModelError(format!("alloc denoise_adam_v: {e}")))?; + let denoise_buf_a = alloc_f32(&stream, b * 12, "denoise_buf_a")?; + let denoise_buf_b = alloc_f32(&stream, b * 12, "denoise_buf_b")?; + let q_var_buf_trainer = alloc_f32(&stream, b * 12, "q_var_buf_trainer")?; + // Xavier uniform init for W1 and W2; biases stay zero. + let mut denoise_params_host = vec![0.0_f32; DENOISE_TOTAL_PARAMS]; + { + let mut rng = 0xDEAD_C0DE_u64.wrapping_add(DENOISE_TOTAL_PARAMS as u64); + let mut off = 0usize; + for _step in 0..3 { + // W1 [24, 24]: fan_in=24, fan_out=24 → limit = sqrt(6/48) + let w1_limit = (6.0_f64 / (24 + 24) as f64).sqrt() as f32; + for w in &mut denoise_params_host[off..off + 576] { + rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); + let u = (rng >> 33) as f32 / (1u64 << 31) as f32 - 0.5; + *w = u * 2.0 * w1_limit; + } + off += 576; + // b1 [24]: stays zero + off += 24; + // W2 [12, 24]: fan_in=24, fan_out=12 → limit = sqrt(6/36) + let w2_limit = (6.0_f64 / (12 + 24) as f64).sqrt() as f32; + for w in &mut denoise_params_host[off..off + 288] { + rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); + let u = (rng >> 33) as f32 / (1u64 << 31) as f32 - 0.5; + *w = u * 2.0 * w2_limit; + } + off += 288; + // b2 [12]: stays zero + off += 12; + } + } + let mut denoise_params = stream.alloc_zeros::(DENOISE_TOTAL_PARAMS) + .map_err(|e| MLError::ModelError(format!("alloc denoise_params: {e}")))?; + super::htod_f32(&stream, &denoise_params_host, &mut denoise_params)?; + let cpbi_module_denoise = stream.context().load_cubin(EXPECTED_Q_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("cpbi cubin (denoise): {e}")))?; + let q_denoise_kernel = cpbi_module_denoise.load_function("q_denoise_step") + .map_err(|e| MLError::ModelError(format!("q_denoise_step load: {e}")))?; + info!("GpuDqnTrainer: q_denoise_step kernel loaded (2700 params, 3-step diffusion denoiser)"); + // v8: Pessimistic Q-value initialization — shift value head bias to -0.1 // Pessimistic Q-init REMOVED — incompatible with per-sample support. // Xavier init gives near-zero value head output, correct for adaptive C51. @@ -3727,6 +3860,14 @@ impl GpuDqnTrainer { graph_adam_v, graph_adam_step: 0, graph_msg_kernel, + denoise_params, + denoise_adam_m, + denoise_adam_v, + denoise_adam_step: 0, + denoise_buf_a, + denoise_buf_b, + q_var_buf_trainer, + q_denoise_kernel, }) } @@ -4944,8 +5085,8 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("atom_stats_buf memset: {e}")))?; let atom_stats_ptr = self.atom_stats_buf.raw_ptr(); let support_ptr = self.per_sample_support_ptr; + let q_var_ptr = self.q_var_buf_trainer.raw_ptr(); unsafe { - let null_q_var = 0u64; // No variance output needed for Q-stats path self.stream .launch_builder(&self.expected_q_kernel) .arg(&on_v_ptr) @@ -4959,7 +5100,7 @@ impl GpuDqnTrainer { .arg(&b3) .arg(&support_ptr) .arg(&atom_stats_ptr) - .arg(&null_q_var) + .arg(&q_var_ptr) .launch(LaunchConfig { grid_dim: (eq_blocks, 1, 1), block_dim: (256, 1, 1), @@ -4978,6 +5119,9 @@ impl GpuDqnTrainer { // Graph message passing: structural coordination across branches (in-place on q_coord_buf). self.launch_graph_message_pass(batch_size)?; + // Diffusion Q-refinement: 3-step denoiser conditioned on Var[Q]. + self.launch_q_denoise(batch_size)?; + // Stats reduction on q_out_buf. let num_atoms = na; unsafe { diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 5fb07c15e..481cc702e 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1990,6 +1990,13 @@ impl FusedTrainingCtx { .map_err(|e| anyhow::anyhow!("launch_graph_message_pass: {e}")) } + /// Run 3-step diffusion denoising on q_coord_buf conditioned on Var[Q]. + /// Must be called AFTER launch_graph_message_pass. + pub(crate) fn launch_q_denoise(&self, batch_size: usize) -> Result<()> { + self.trainer.launch_q_denoise(batch_size) + .map_err(|e| anyhow::anyhow!("launch_q_denoise: {e}")) + } + /// Center Q-values in q_out_buf by subtracting their global mean. /// Prevents bootstrapping drift. Must be called after compute_expected_q, /// before launch_q_attention.