diff --git a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs index e87fdbe7a..108213f59 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs @@ -137,8 +137,8 @@ impl GpuIqlConfig { fn total_params(&self) -> usize { let h = self.value_hidden_dim; let sd = self.state_dim; - // w1[H*SD] + b1[H] + w2[H*H] + b2[H] + w3[H] + b3[1] - h * sd + h + h * h + h + h + 1 + // w1[H*SD] + b1[H] + w2[H*H] + b2[H] + w3[H*2] + b3[2] + h * sd + h + h * h + h + h * 2 + 2 } } @@ -207,7 +207,8 @@ pub struct GpuIqlTrainer { adv_sigma_ema_kernel: CudaFunction, // ── V network parameters (flat f32 on GPU) ───────────────────── - // Layout: W1[H*SD] + b1[H] + W2[H*H] + b2[H] + W3[H] + b3[1] + // Layout: W1[H*SD] + b1[H] + W2[H*H] + b2[H] + W3[H*2] + b3[2] + // W3 is the 2-output head: row 0 = V_short weights, row 1 = V_long weights. params_buf: CudaSlice, // ── Adam optimizer state ──────────────────────────────────────── @@ -224,7 +225,7 @@ pub struct GpuIqlTrainer { h2_buf: CudaSlice, // [H, B] post-SiLU layer 2 // ── cuBLAS backward intermediate buffers ──────────────────────── - dv_buf: CudaSlice, // [1, B] d_expectile_loss output + dv_buf: CudaSlice, // [2, B] d_expectile_loss for both heads dh2_buf: CudaSlice, // [H, B] dh2 (pre silu_bwd) dh2_pre_buf: CudaSlice, // [H, B] dh2_pre (post silu_bwd) dh1_buf: CudaSlice, // [H, B] dh1 @@ -282,14 +283,17 @@ impl GpuIqlTrainer { 1, 0, "iql_fwd_h1")?; // TRANSA=T, TRANSB=N let gemm_fwd_h2 = create_iql_gemm_desc(lt_handle, h, b, h, h, lt_ws_size, 1, 0, "iql_fwd_h2")?; - let gemm_fwd_v = create_iql_gemm_desc(lt_handle, 1, b, h, h, lt_ws_size, + let gemm_fwd_v = create_iql_gemm_desc(lt_handle, 2, b, h, h, lt_ws_size, 1, 0, "iql_fwd_v")?; // Create 5 backward GEMM descriptors - // dW3 = dv[1,B] @ h2^T[B,H] -> [1,H] TRANSA=N, TRANSB=T - let gemm_bwd_dw3 = create_iql_gemm_desc_nt(lt_handle, 1, h, b, 1, h, lt_ws_size, "iql_bwd_dw3")?; - // dh2 = W3[H,1] @ dv[1,B] -> [H,B] TRANSA=N, TRANSB=N - let gemm_bwd_dh2 = create_iql_gemm_desc(lt_handle, h, b, 1, 1, lt_ws_size, + // dW3 = dv[2,B] @ h2^T[B,H] -> [2,H] TRANSA=N, TRANSB=T + let gemm_bwd_dw3 = create_iql_gemm_desc_nt(lt_handle, 2, h, b, 2, h, lt_ws_size, "iql_bwd_dw3")?; + // dh2 = W3[H,2]^T @ dv[2,B] -> [H,B] (K=2: W3 has 2 output rows) + // W3 is [H, 2] col-major (stored as [2, H] because TRANSA=T for fwd). + // Backward: dh2[H,B] = W3^T[H,2]^T @ dv[2,B] where the GEMM is + // M=H, N=B, K=2 with TRANSA=N (A=W3 physical [H,2] col-major, ld=H). + let gemm_bwd_dh2 = create_iql_gemm_desc(lt_handle, h, b, 2, 2, lt_ws_size, 0, 0, "iql_bwd_dh2")?; // dW2 = dh2_pre[H,B] @ h1^T[B,H] -> [H,H] let gemm_bwd_dw2 = create_iql_gemm_desc_nt(lt_handle, h, h, b, h, h, lt_ws_size, "iql_bwd_dw2")?; @@ -344,7 +348,7 @@ impl GpuIqlTrainer { let h1_buf = alloc_f32(&stream, h * b, "iql_h1")?; let h2_pre_buf = alloc_f32(&stream, h * b, "iql_h2_pre")?; let h2_buf = alloc_f32(&stream, h * b, "iql_h2")?; - let dv_buf = alloc_f32(&stream, b, "iql_dv")?; + let dv_buf = alloc_f32(&stream, b * 2, "iql_dv")?; let dh2_buf = alloc_f32(&stream, h * b, "iql_dh2")?; let dh2_pre_buf = alloc_f32(&stream, h * b, "iql_dh2_pre")?; let dh1_buf = alloc_f32(&stream, h * b, "iql_dh1")?; @@ -364,8 +368,8 @@ impl GpuIqlTrainer { let grad_norm_partials = alloc_f32(&stream, grad_norm_blocks, "iql_grad_norm_partials")?; // Allocate output buffers - let v_out_buf = alloc_f32(&stream, b, "iql_v_out")?; - let loss_buf = alloc_f32(&stream, b, "iql_loss")?; + let v_out_buf = alloc_f32(&stream, b * 2, "iql_v_out")?; + let loss_buf = alloc_f32(&stream, b * 2, "iql_loss")?; let total_loss_buf = alloc_f32(&stream, 1, "iql_total_loss")?; let q_taken_buf = alloc_f32(&stream, b, "iql_q_taken")?; let advantage_weights_buf = alloc_f32(&stream, b, "iql_adv_weights")?; @@ -402,11 +406,11 @@ impl GpuIqlTrainer { } super::htod_f32(&stream, &default_support, &mut per_sample_support_buf)?; - let cublas_fwd_vram = h * b * 4 * 4 + b * 4; // h1_pre+h1+h2_pre+h2 + dv - let cublas_bwd_vram = h * b * 4 * 4; // dh2+dh2_pre+dh1+dh1_pre - let vram_bytes = total_params * 4 * 4 // params + m + v + grad + let cublas_fwd_vram = h * b * 4 * 4 + b * 2 * 4; // h1_pre+h1+h2_pre+h2 + dv[B*2] + let cublas_bwd_vram = h * b * 4 * 4; // dh2+dh2_pre+dh1+dh1_pre + let vram_bytes = total_params * 4 * 4 // params + m + v + grad + cublas_fwd_vram + cublas_bwd_vram - + (b * 3 + 2) * 4; // v_out + loss + total_loss + adv_weights + grad_norm + + (b * 4 + 2) * 4; // v_out[B*2] + loss[B*2] + total_loss + adv_weights + grad_norm info!( state_dim = ml_core::state_layout::STATE_DIM, @@ -542,7 +546,7 @@ impl GpuIqlTrainer { let w2_off = b1_off + h; let b2_off = w2_off + h * h; let w3_off = b2_off + h; - let b3_off = w3_off + h; + let b3_off = w3_off + h * 2; let f32_sz = std::mem::size_of::(); // Raw pointers into params_buf @@ -639,7 +643,7 @@ impl GpuIqlTrainer { .map_err(|e| MLError::ModelError(format!("IQL silu_fwd h2: {e}")))?; } - // 3. v_pre[1,B] = W3^T @ h2 -> v_out_buf + // 3. v_pre[2,B] = W3^T @ h2 -> v_out_buf (2 heads: V_short row 0, V_long row 1) unsafe { cublaslt_sys::cublasLtMatmul( lt_handle, @@ -660,23 +664,25 @@ impl GpuIqlTrainer { cu_stream, ); } - // Add output bias (b3, scalar broadcast) - let n_v = b as i32; - let one_i32 = 1_i32; - let v_blocks = ((b + 255) / 256) as u32; + // Add output bias (b3, 2-element broadcast across B*2 elements) + let n_v2 = (b * 2) as i32; + let two_i32 = 2_i32; + let v_blocks2 = ((b * 2 + 255) / 256) as u32; unsafe { self.stream.launch_builder(&self.bias_add_kernel) .arg(&self.v_out_buf) .arg(&b3_ptr) - .arg(&one_i32) - .arg(&n_v) - .launch(LaunchConfig { grid_dim: (v_blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) + .arg(&two_i32) + .arg(&n_v2) + .launch(LaunchConfig { grid_dim: (v_blocks2, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) .map_err(|e| MLError::ModelError(format!("IQL bias_add v: {e}")))?; } - // 4. Expectile loss + dv: combined kernel - // loss[B] = |tau - 1(u<0)| * u^2 where u = q_taken - v_out - // dv[B] = d(loss)/d(v_out) + // 4. Expectile loss + dv: both heads, same Q-target per sample. + // loss[B*2] = |tau - 1(u<0)| * u^2 for each (sample, head) pair + // dv[B*2] = -2*w*u + // iql_expectile_loss indexes q_taken[b] = q_taken[idx / num_heads] + let v_blocks = ((b + 255) / 256) as u32; unsafe { self.stream.launch_builder(&self.expectile_loss_kernel) .arg(&self.v_out_buf) @@ -685,7 +691,8 @@ impl GpuIqlTrainer { .arg(&mut self.loss_buf) .arg(&mut self.dv_buf) .arg(&batch_size_i32) - .launch(LaunchConfig { grid_dim: (v_blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) + .arg(&two_i32) + .launch(LaunchConfig { grid_dim: (v_blocks2, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) .map_err(|e| MLError::ModelError(format!("IQL expectile_loss: {e}")))?; } @@ -703,7 +710,7 @@ impl GpuIqlTrainer { let inv_batch: f32 = 1.0 / b as f32; let beta_zero: f32 = 0.0; - // dW3[1,H] = dv[1,B] @ h2^T[B,H] (alpha=1/B for mean reduction) + // dW3[2,H] = dv[2,B] @ h2^T[B,H] (alpha=1/B for mean reduction) let dw3_ptr = self.grad_buf.raw_ptr() + (w3_off * f32_sz) as u64; unsafe { cublaslt_sys::cublasLtMatmul( @@ -726,7 +733,8 @@ impl GpuIqlTrainer { ); } - // db3: sum(dv) / B -- scalar bias gradient via 2-phase reduce + // db3[2]: sum columns of dv[2,B] per output head via 2-phase reduce + // Grid y-dim=2 so each bias element gets its own reduce lane. let db3_ptr = self.grad_buf.raw_ptr() + (b3_off * f32_sz) as u64; let num_blocks = self.bias_grad_num_blocks; let partials_ptr = self.bias_grad_partials_buf.raw_ptr(); @@ -734,9 +742,9 @@ impl GpuIqlTrainer { self.stream.launch_builder(&self.bias_grad_reduce_p1_kernel) .arg(&self.dv_buf) .arg(&partials_ptr) - .arg(&one_i32) + .arg(&two_i32) .arg(&batch_size_i32) - .launch(LaunchConfig { grid_dim: (num_blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 256 * 4 }) + .launch(LaunchConfig { grid_dim: (num_blocks, 2, 1), block_dim: (256, 1, 1), shared_mem_bytes: 256 * 4 }) .map_err(|e| MLError::ModelError(format!("IQL bias_grad_reduce_p1 b3: {e}")))?; } let num_blocks_i32 = num_blocks as i32; @@ -744,14 +752,14 @@ impl GpuIqlTrainer { self.stream.launch_builder(&self.bias_grad_reduce_p2_kernel) .arg(&partials_ptr) .arg(&db3_ptr) - .arg(&one_i32) + .arg(&two_i32) .arg(&num_blocks_i32) .arg(&inv_batch) .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) .map_err(|e| MLError::ModelError(format!("IQL bias_grad_reduce_p2 b3: {e}")))?; } - // dh2[H,B] = W3[H,1] @ dv[1,B] + // dh2[H,B] = W3^T[H,2] @ dv[2,B] (K=2 now) unsafe { cublaslt_sys::cublasLtMatmul( lt_handle, @@ -910,13 +918,14 @@ impl GpuIqlTrainer { // ── Loss reduce + grad norm + Adam (unchanged) ── - // 4. Loss reduce (deterministic sequential sum) + // 4. Loss reduce (deterministic sequential sum over B*2 elements) unsafe { self.stream .launch_builder(&self.loss_reduce_kernel) .arg(&self.loss_buf) .arg(&mut self.total_loss_buf) .arg(&batch_size_i32) + .arg(&two_i32) .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), @@ -1441,14 +1450,14 @@ fn init_xavier_weights( } // b2 stays zero (w2_end .. w2_end + h) - // Output layer: w3[H], b3[1] + // Output layer: w3[H*2] (2 output heads: V_short row 0, V_long row 1), b3[2] let w3_start = w2_end + h; - let limit3 = (6.0_f64 / (h + 1) as f64).sqrt() as f32; - let w3_end = w3_start + h; + let limit3 = (6.0_f64 / (h + 2) as f64).sqrt() as f32; + let w3_end = w3_start + h * 2; for w in &mut weights[w3_start..w3_end] { *w = rng.gen_range(-limit3..limit3); } - // b3 stays zero + // b3[2] stays zero // Upload to GPU let mut params_buf = alloc_f32(stream, total, "iql_params")?; diff --git a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu index 334dce9d4..df29c6934 100644 --- a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu @@ -33,8 +33,8 @@ #define V_B1_SIZE (VALUE_HIDDEN_DIM) #define V_W2_SIZE (VALUE_HIDDEN_DIM * VALUE_HIDDEN_DIM) #define V_B2_SIZE (VALUE_HIDDEN_DIM) -#define V_W3_SIZE (VALUE_HIDDEN_DIM) /* output layer: 1 x H */ -#define V_B3_SIZE (1) +#define V_W3_SIZE (VALUE_HIDDEN_DIM * 2) /* output layer: 2 x H (V_short + V_long) */ +#define V_B3_SIZE (2) /* Runtime offset computation helper -- state_dim is a kernel parameter. * W1 size = VALUE_HIDDEN_DIM * state_dim (varies with state_dim). */ @@ -98,15 +98,20 @@ __device__ __forceinline__ float iql_block_sum( */ extern "C" __global__ void iql_loss_reduce( - const float* __restrict__ loss_out, /* [B] per-sample loss */ + const float* __restrict__ loss_out, /* [B * num_heads] per-element loss */ float* __restrict__ total_loss, /* [1] output */ - int batch_size + int batch_size, + int num_heads ) { float sum = 0.0f; - for (int b = 0; b < batch_size; b++) { - sum += loss_out[b]; + int total = batch_size * num_heads; + for (int i = 0; i < total; i++) { + sum += loss_out[i]; } + /* Normalise by B so the loss scale is consistent with the single-head case. + * Summing over num_heads and dividing by B gives the mean over samples; + * each head contributes equally. */ total_loss[0] = sum / (float)batch_size; } @@ -330,15 +335,25 @@ void iql_forward_kernel( } __syncthreads(); - /* Output: dot product + bias */ - float v_acc = 0.0f; + /* Output: two heads (V_short, V_long) — dot products with rows 0 and 1 + * of W3[2, H] stored row-major in params. v_out layout [2, B] col-major: + * v_out[0*B + sample] = V_short, v_out[1*B + sample] = V_long. + * Consumers sum the two: V(s) = v_out[sample*2+0] + v_out[sample*2+1] + * using the [B*2] stride-2 view produced by the cuBLAS forward path. */ + float v0_acc = 0.0f; + float v1_acc = 0.0f; for (int k = tid; k < VALUE_HIDDEN_DIM; k += 256) { - v_acc += w3[k] * sh2[k]; + v0_acc += w3[k] * sh2[k]; + v1_acc += w3[VALUE_HIDDEN_DIM + k] * sh2[k]; } - v_acc = iql_block_sum(v_acc, warp_sums); + v0_acc = iql_block_sum(v0_acc, warp_sums); + __syncthreads(); + v1_acc = iql_block_sum(v1_acc, warp_sums); if (tid == 0) { - v_out[sample] = v_acc + b3[0]; + /* col-major [2, B]: row r, col b stored at r + b*2 */ + v_out[sample * 2 + 0] = v0_acc + b3[0]; + v_out[sample * 2 + 1] = v1_acc + b3[1]; } } @@ -358,7 +373,7 @@ extern "C" __global__ void iql_compute_advantage_weights( const float* __restrict__ q_out, /* [B, b0+b1+b2+b3] branch Q-values */ const int* __restrict__ actions, /* [B] factored action indices */ - const float* __restrict__ v_out, /* [B] V(s) from IQL */ + const float* __restrict__ v_out, /* [B*2] V(s) two heads col-major [2,B] */ const float* __restrict__ readiness_buf,/* [1] CV-based readiness */ float* __restrict__ adv_weights, /* [B] output weights */ float beta, @@ -378,7 +393,8 @@ void iql_compute_advantage_weights( int mag = a % b1; a /= b1; int dir = a; float q_taken = row[dir] + row[b0 + mag] + row[b0 + b1 + ord] + row[b0 + b1 + b2 + urg]; - float adv = q_taken - v_out[b]; + /* V(s) = V_short + V_long; col-major [2, B]: element (head, b) at head + b*2 */ + float adv = q_taken - (v_out[b * 2 + 0] + v_out[b * 2 + 1]); float raw_w = expf(beta * adv); float clamped_w = fminf(fmaxf(raw_w, 0.01f), 100.0f); /* Blend: readiness=0 → neutral weight 1.0, readiness=1 → IQL weight */ @@ -571,7 +587,7 @@ void iql_adv_variance_reduce( */ extern "C" __global__ void iql_compute_per_sample_support( - const float* __restrict__ v_out, /* [B] */ + const float* __restrict__ v_out, /* [B*2] two heads col-major [2,B] */ const float* __restrict__ q_out, /* [B, total_actions] */ float* __restrict__ per_sample_support, /* [B, 4, 3] */ const float* __restrict__ readiness_buf, /* [1] CV-based readiness */ @@ -586,7 +602,8 @@ void iql_compute_per_sample_support( if (b >= batch_size) return; float r = readiness_buf[0]; - float v = v_out[b]; + /* V(s) = V_short + V_long; col-major [2, B]: element (head, b) at head + b*2 */ + float v = v_out[b * 2 + 0] + v_out[b * 2 + 1]; const float* q = q_out + b * total_actions; float spread = 0.0f; @@ -730,7 +747,7 @@ void iql_support_floor( extern "C" __global__ void iql_per_branch_advantage( const float* __restrict__ q_out, /* [B, total_actions] */ - const float* __restrict__ v_out, /* [B] */ + const float* __restrict__ v_out, /* [B*2] two heads col-major [2,B] */ const int* __restrict__ actions, /* [B] factored action indices */ float* __restrict__ branch_scales, /* [B*4] */ const float* __restrict__ readiness_buf, /* [1] CV-based readiness */ @@ -754,7 +771,8 @@ void iql_per_branch_advantage( int b = blockIdx.x * blockDim.x + threadIdx.x; if (b >= batch_size) return; - float v = v_out[b]; + /* V(s) = V_short + V_long; col-major [2, B]: element (head, b) at head + b*2 */ + float v = v_out[b * 2 + 0] + v_out[b * 2 + 1]; const float* q = q_out + b * total_actions; int factored = actions[b]; @@ -936,10 +954,21 @@ void iql_silu_bwd( } } -/* iql_expectile_loss — Expectile loss + gradient (fused) - * loss[i] = w*(Q[i]-V[i])^2, w = tau if u>=0 else (1-tau) - * dv[i] = -2*w*(Q[i]-V[i]) - * Launch: grid=(ceil(B/256),1,1), block=(256,1,1) +/* iql_expectile_loss — Expectile loss + gradient (fused), horizon-decomposed. + * + * Operates over B*num_heads elements. For idx in [0, B*num_heads): + * b = idx / num_heads (sample index) + * v_out[idx] = V_head for that (sample, head) pair + * q_taken[b] = same Q-target for both heads of the same sample + * + * loss[idx] = w*(Q[b]-V[idx])^2, w = tau if u>=0 else (1-tau) + * dv[idx] = -2*w*(Q[b]-V[idx]) + * + * For the 2-head case (num_heads=2): v_out is [B*2] col-major [2,B]; + * both heads regress the same expectile target. The sum V(s)=V_short+V_long + * is consumed by downstream kernels. + * + * Launch: grid=(ceil(B*num_heads/256),1,1), block=(256,1,1) */ extern "C" __global__ void iql_expectile_loss( @@ -948,12 +977,14 @@ void iql_expectile_loss( float tau, float* __restrict__ loss_out, float* __restrict__ dv_out, - int B + int B, + int num_heads ) { int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx < B) { + if (idx < B * num_heads) { + int b = idx / num_heads; float v = v_out[idx]; - float q = q_taken[idx]; + float q = q_taken[b]; float u = q - v; float weight = (u >= 0.0f) ? tau : (1.0f - tau); loss_out[idx] = weight * u * u; diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index e86a9794c..4324d1787 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -234,6 +234,8 @@ Updated after Task 6 cleanup (2026-04-24): 5 confirmed-orphan files deleted, 3 O Plan 1 Tasks 12/15/16 + pre-allocation (2026-04-24): No new modules added. Changes are ISV slot allocation + consumer migration only. Task 15 confirmed no-op (`IQL_BRANCH_SCALE_FLOOR_INDEX` already serves conviction-floor role). Tasks 12 and 16 migrate `cql_alpha` and plan-threshold consumers from config fields / hardcoded literals to ISV slots. 8 new ISV slots allocated ([39..47)); fingerprint tail moves from [37..39) to [47..49); `ISV_TOTAL_DIM` 39 → 49. `GpuDqnTrainConfig` gains `total_epochs` field (written to `TOTAL_EPOCHS_INDEX` at construction). `write_isv_signal_at` bound extended from `ISV_DIM` to `ISV_TOTAL_DIM` to allow writes beyond slot 22. +Plan 2 Task 6B D.3 (2026-04-24): IQL value head widened from 1 to 2 outputs (V_short + V_long). `v_out_buf` shape `[B]` → `[B*2]`. `gemm_fwd_v` M=1→2, `gemm_bwd_dw3` M=1→2, `gemm_bwd_dh2` K=1→2. `W3` param block `[H*1]` → `[H*2]`, `b3` `[1]` → `[2]`. `total_params` += H+1. `iql_expectile_loss` kernel extended with `num_heads` argument. 4 consumer kernels in `iql_value_kernel.cu` updated to read `v_out[b*2+0] + v_out[b*2+1]`. Checkpoint compat break — retrain required. + | Classification | Count | |---|---| | Wired | 75 |