From d3d9c4a14511c74efdd7fae47746ce4500bec250 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 14 Apr 2026 00:19:26 +0200 Subject: [PATCH] =?UTF-8?q?fix:=20IQL=20kernels=20decode=20factored=20acti?= =?UTF-8?q?ons=20into=20branch=20indices=20=E2=80=94=20eliminates=20OOB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit iql_gather_q_taken and iql_compute_advantage_weights indexed q_out with factored action (0..80) into a [B,12] buffer — massive OOB. compute-sanitizer found 2586 errors. Now decodes factored action (dir*b1*b2*b3 + mag*b2*b3 + ord*b3 + urg) into 4 branch indices and sums per-branch Q-values. Also fixes total_actions config (was product 81, now sum 12). Sanitizer: 0 errors. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_iql_trainer.rs | 26 ++++++++--- .../ml/src/cuda_pipeline/iql_value_kernel.cu | 45 ++++++++++++++----- crates/ml/src/trainers/dqn/fused_training.rs | 4 +- 3 files changed, 57 insertions(+), 18 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs index 97a19787f..ca5ca0506 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs @@ -540,11 +540,15 @@ impl GpuIqlTrainer { &mut self, q_out_buf: &CudaSlice, actions_buf: &CudaSlice, - total_actions: usize, + _total_actions: usize, ) -> Result<&CudaSlice, MLError> { let b = self.config.batch_size; let batch_size_i32 = b as i32; - let total_actions_i32 = total_actions as i32; + let bs = &self.config.branch_sizes; + let b0 = bs[0] as i32; + let b1 = bs[1] as i32; + let b2 = bs[2] as i32; + let b3 = bs[3] as i32; let blocks = (b + 255) / 256; unsafe { @@ -554,7 +558,10 @@ impl GpuIqlTrainer { .arg(actions_buf) .arg(&mut self.q_taken_buf) .arg(&batch_size_i32) - .arg(&total_actions_i32) + .arg(&b0) + .arg(&b1) + .arg(&b2) + .arg(&b3) .launch(LaunchConfig { grid_dim: (blocks as u32, 1, 1), block_dim: (256, 1, 1), @@ -582,12 +589,16 @@ impl GpuIqlTrainer { &mut self, q_out_buf: &CudaSlice, actions_buf: &CudaSlice, - total_actions: usize, + _total_actions: usize, ) -> Result<(), MLError> { let b = self.config.batch_size; let beta = self.config.advantage_temperature; let batch_size_i32 = b as i32; - let total_actions_i32 = total_actions as i32; + let bs = &self.config.branch_sizes; + let b0 = bs[0] as i32; + let b1 = bs[1] as i32; + let b2 = bs[2] as i32; + let b3 = bs[3] as i32; let blocks = (b + 255) / 256; unsafe { @@ -600,7 +611,10 @@ impl GpuIqlTrainer { .arg(&mut self.advantage_weights_buf) .arg(&beta) .arg(&batch_size_i32) - .arg(&total_actions_i32) + .arg(&b0) + .arg(&b1) + .arg(&b2) + .arg(&b3) .launch(LaunchConfig { grid_dim: (blocks as u32, 1, 1), block_dim: (256, 1, 1), diff --git a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu index 24e560ec5..9073364e0 100644 --- a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu @@ -468,23 +468,41 @@ void iql_adam_kernel( /* Gather Q(s, a_taken) from q_out_buf */ /* ------------------------------------------------------------------ */ /** - * Extract Q-value for the taken action from the full Q-value buffer. - * q_taken[b] = q_out[b * total_actions + actions[b]] + * Extract Q-value for the taken factored action from branch Q-values. + * + * q_out is [B, b0+b1+b2+b3] — per-branch expected Q-values. + * actions[b] is a factored action index: dir*b1*b2*b3 + mag*b2*b3 + ord*b3 + urg. + * Q(s,a) = q_out[b, dir_offset + dir] + q_out[b, mag_offset + mag] + * + q_out[b, ord_offset + ord] + q_out[b, urg_offset + urg] * * Launch: grid=ceil(B/256), block=256. */ extern "C" __global__ void iql_gather_q_taken( - const float* __restrict__ q_out, /* [B, total_actions] */ - const int* __restrict__ actions, /* [B] taken action indices */ + const float* __restrict__ q_out, /* [B, b0+b1+b2+b3] */ + const int* __restrict__ actions, /* [B] factored action indices */ float* __restrict__ q_taken, /* [B] output */ int batch_size, - int total_actions + int b0, int b1, int b2, int b3 ) { int b = blockIdx.x * blockDim.x + threadIdx.x; if (b >= batch_size) return; - q_taken[b] = q_out[b * total_actions + actions[b]]; + + int a = actions[b]; + int total_per_row = b0 + b1 + b2 + b3; + const float* row = q_out + b * total_per_row; + + /* Decode factored action: a = dir*(b1*b2*b3) + mag*(b2*b3) + ord*b3 + urg */ + int urg = a % b3; a /= b3; + int ord = a % b2; a /= b2; + int mag = a % b1; a /= b1; + int dir = a; + + q_taken[b] = row[dir] + + row[b0 + mag] + + row[b0 + b1 + ord] + + row[b0 + b1 + b2 + urg]; } /* ------------------------------------------------------------------ */ @@ -574,21 +592,28 @@ void iql_forward_kernel( */ extern "C" __global__ void iql_compute_advantage_weights( - const float* __restrict__ q_out, /* [B, total_actions] */ - const int* __restrict__ actions, /* [B] taken action indices */ + 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__ readiness_buf,/* [1] CV-based readiness */ float* __restrict__ adv_weights, /* [B] output weights */ float beta, int batch_size, - int total_actions + int b0, int b1, int b2, int b3 ) { int b = blockIdx.x * blockDim.x + threadIdx.x; if (b >= batch_size) return; float r = readiness_buf[0]; - float q_taken = q_out[b * total_actions + actions[b]]; + int total_per_row = b0 + b1 + b2 + b3; + const float* row = q_out + b * total_per_row; + int a = actions[b]; + int urg = a % b3; a /= b3; + int ord = a % b2; a /= b2; + 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]; float raw_w = expf(beta * adv); float clamped_w = fminf(fmaxf(raw_w, 0.01f), 100.0f); diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 58f658811..7e96854a0 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -479,8 +479,8 @@ impl FusedTrainingCtx { lr: hyperparams.learning_rate as f32, max_grad_norm: resolved_grad_norm as f32, num_atoms: dqn.config.num_atoms, - total_actions: dqn.config.num_actions * dqn.config.num_order_types - * dqn.config.num_urgency_levels * 3, + total_actions: dqn.config.num_actions + dqn.config.num_order_types + + dqn.config.num_urgency_levels + 3, branch_sizes: [ dqn.config.num_actions, dqn.config.num_order_types,