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

feat: widen w_b1fc [AH, SH2+3] + allocate mag_concat buffers + launch methods

Browse Source 
param_sizes[12] grows by 3*adv_h for direction conditioning.
mag_concat_buf [B, SH2+3] and d_mag_concat_buf allocated.
mag_concat_qdir and strided_accumulate kernels loaded.
launch_mag_concat and accumulate_d_h_s2_from_concat methods added.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-14 22:37:17 +02:00
parent abb0795759
commit cf009ad49e
1 changed files with 97 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -389,7 +389,7 @@ pub(crate) fn compute_param_sizes(cfg: &GpuDqnTrainConfig) -> [usize; NUM_WEIGHT
cfg.adv_h, // [9] b_b0fc
cfg.branch_0_size * cfg.num_atoms * cfg.adv_h, // [10] w_b0out
cfg.branch_0_size * cfg.num_atoms, // [11] b_b0out
cfg.adv_h * cfg.shared_h2, // [12] w_b1fc
cfg.adv_h * (cfg.shared_h2 + 3), // [12] w_b1fc (direction-conditioned: SH2+3 input)
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
@@ -459,6 +459,8 @@ struct CachedPtrs {
save_h_b1: u64,
save_h_b2: u64,
save_h_b3: u64,
mag_concat_buf: u64,
d_mag_concat_buf: u64,
bw_d_h_s1: u64,
bw_d_h_s2: u64,
bw_d_h_v: u64,
@@ -588,6 +590,16 @@ pub struct GpuDqnTrainer {
save_h_b1: CudaSlice<f32>, // [B, ADV_H]
save_h_b2: CudaSlice<f32>, // [B, ADV_H]
save_h_b3: CudaSlice<f32>, // [B, ADV_H]
/// Magnitude branch concat input [B, SH2+3]: [h_s2; Q_dir] for direction conditioning.
mag_concat_buf: CudaSlice<f32>,
/// Backward scratch for magnitude concat dX [B, SH2+3].
d_mag_concat_buf: CudaSlice<f32>,
/// Kernel: mag_concat_qdir — builds [h_s2; Q_dir] concat.
mag_concat_kernel: CudaFunction,
/// Kernel: strided_accumulate — extracts d_h_s2 from d_mag_concat in backward.
strided_accumulate_kernel: CudaFunction,
save_current_lp: CudaSlice<f32>, // [B, NUM_BRANCHES(4), NUM_ATOMS]
save_projected: CudaSlice<f32>, // [B, NUM_BRANCHES(4), NUM_ATOMS]
@@ -1083,6 +1095,71 @@ impl GpuDqnTrainer {
self.branch_scales_ptr = ptr;
}
/// Build [h_s2; Q_dir] concat buffer for magnitude branch conditioning.
/// Must be called AFTER direction branch forward (branch 0 logits are ready).
pub(crate) fn launch_mag_concat(&self, v_logits_ptr: u64, b_logits_ptr: u64) -> Result<(), MLError> {
let b = self.config.batch_size;
let sh2 = self.config.shared_h2 as i32;
let na = self.config.num_atoms as i32;
let b0 = self.config.branch_0_size as i32;
let n = b as i32;
let blocks = ((b as u32 + 255) / 256).max(1);
let h_s2_ptr = self.ptrs.save_h_s2;
let concat_ptr = self.ptrs.mag_concat_buf;
let support_ptr = self.per_sample_support_ptr;
unsafe {
self.stream
.launch_builder(&self.mag_concat_kernel)
.arg(&h_s2_ptr)
.arg(&v_logits_ptr)
.arg(&b_logits_ptr)
.arg(&concat_ptr)
.arg(&n)
.arg(&sh2)
.arg(&na)
.arg(&b0)
.arg(&support_ptr)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("mag_concat_qdir: {e}")))?;
}
Ok(())
}
/// Accumulate first SH2 columns of d_mag_concat [B, SH2+3] into d_h_s2 [B, SH2].
pub(crate) fn accumulate_d_h_s2_from_concat(
&self,
d_concat_ptr: u64,
d_h_s2_ptr: u64,
batch: usize,
beta: f32,
) -> Result<(), MLError> {
let total = (batch * self.config.shared_h2) as i32;
let blocks = ((total as u32 + 255) / 256).max(1);
let sh2 = self.config.shared_h2 as i32;
let sh2_plus_3 = (self.config.shared_h2 + 3) as i32;
unsafe {
self.stream
.launch_builder(&self.strided_accumulate_kernel)
.arg(&d_concat_ptr)
.arg(&d_h_s2_ptr)
.arg(&sh2)
.arg(&sh2_plus_3)
.arg(&total)
.arg(&beta)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("strided_accumulate: {e}")))?;
}
Ok(())
}
/// Update IQN readiness from the current IQN loss.
/// Readiness ramps 0→1 as loss decreases from initial value.
/// Uses adaptive-rate EMA (same pattern as eval_v_range).
@@ -2370,6 +2447,9 @@ impl GpuDqnTrainer {
let save_h_b1 = alloc_f32(&stream, b * config.adv_h + kt, "save_h_b1")?;
let save_h_b2 = alloc_f32(&stream, b * config.adv_h + kt, "save_h_b2")?;
let save_h_b3 = alloc_f32(&stream, b * config.adv_h + kt, "save_h_b3")?;
let mag_concat_dim = config.shared_h2 + 3;
let mag_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "mag_concat_buf")?;
let d_mag_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "d_mag_concat_buf")?;
let save_current_lp = alloc_f32(
&stream,
b * num_branches * config.num_atoms,
@@ -2591,6 +2671,16 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("alloc atom_stats: {e}")))?;
info!("GpuDqnTrainer: expected_q + q_stats kernels compiled");
// ── Load mag_concat and strided_accumulate from experience_kernels cubin ──
let exp_module_for_mag = stream.context()
.load_cubin(EXPECTED_Q_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("exp cubin for mag_concat: {e}")))?;
let mag_concat_kernel = exp_module_for_mag.load_function("mag_concat_qdir")
.map_err(|e| MLError::ModelError(format!("mag_concat_qdir load: {e}")))?;
let strided_accumulate_kernel = exp_module_for_mag.load_function("strided_accumulate")
.map_err(|e| MLError::ModelError(format!("strided_accumulate load: {e}")))?;
info!("GpuDqnTrainer: mag_concat_qdir + strided_accumulate kernels loaded");
// ── Compile CQL penalty kernel (if enabled) ──────────────────────
let cql_logit_grad_kernel = if config.cql_alpha > 0.0 {
match compile_cql_logit_grad_kernel(&stream) {
@@ -2908,6 +2998,8 @@ impl GpuDqnTrainer {
save_h_b1: save_h_b1.raw_ptr(),
save_h_b2: save_h_b2.raw_ptr(),
save_h_b3: save_h_b3.raw_ptr(),
mag_concat_buf: mag_concat_buf.raw_ptr(),
d_mag_concat_buf: d_mag_concat_buf.raw_ptr(),
bw_d_h_s1: bw_d_h_s1.raw_ptr(),
bw_d_h_s2: bw_d_h_s2.raw_ptr(),
bw_d_h_v: bw_d_h_v.raw_ptr(),
@@ -3060,6 +3152,10 @@ impl GpuDqnTrainer {
save_h_b1,
save_h_b2,
save_h_b3,
mag_concat_buf,
d_mag_concat_buf,
mag_concat_kernel,
strided_accumulate_kernel,
save_current_lp,
save_projected,
per_sample_loss_buf,