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

feat: diffusion Q-refinement — 3-step denoiser conditioned on Var[Q] (2,700 params)

Browse Source 
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-15 09:19:23 +02:00
parent cd7c67f541
commit e328bcdb13
2 changed files with 157 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

156
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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<f32>,
/// Adam first moment for denoise_params [2700].
denoise_adam_m: CudaSlice<f32>,
/// Adam second moment for denoise_params [2700].
denoise_adam_v: CudaSlice<f32>,
/// Adam step counter for denoise optimizer.
denoise_adam_step: i32,
/// Ping-pong buffer A for iterative denoising [B, 12].
denoise_buf_a: CudaSlice<f32>,
/// Ping-pong buffer B for iterative denoising [B, 12].
denoise_buf_b: CudaSlice<f32>,
/// Var[Q] buffer for training path [B, 12] — populated by compute_expected_q.
q_var_buf_trainer: CudaSlice<f32>,
/// 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::<f32>();
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::<f32>()) as u64;
let b1_ptr = w1_ptr + (576 * std::mem::size_of::<f32>()) as u64;
let w2_ptr = b1_ptr + (24 * std::mem::size_of::<f32>()) as u64;
let b2_ptr = w2_ptr + (288 * std::mem::size_of::<f32>()) 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::<f32>(DENOISE_TOTAL_PARAMS)
.map_err(|e| MLError::ModelError(format!("alloc denoise_adam_m: {e}")))?;
let denoise_adam_v = stream.alloc_zeros::<f32>(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::<f32>(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 {

7
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -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.