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

fix: IQL kernels decode factored actions into branch indices — eliminates OOB

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-14 00:19:26 +02:00
parent 2d704eaa07
commit d3d9c4a145
3 changed files with 57 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs
View File

@@ -540,11 +540,15 @@ impl GpuIqlTrainer {
&mut self,
q_out_buf: &CudaSlice<f32>,
actions_buf: &CudaSlice<i32>,
total_actions: usize,
_total_actions: usize,
) -> Result<&CudaSlice<f32>, 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<f32>,
actions_buf: &CudaSlice<i32>,
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),

45
crates/ml/src/cuda_pipeline/iql_value_kernel.cu
View File

@@ -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);

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

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