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

fix: parameterize IQN_EMBED_DIM in all 9 IQN CUDA kernels

Browse Source 
IQN_EMBED_DIM (cosine embedding dimension, default 64) converted from
compile-time #define to runtime kernel parameter `embed_dim`. This was
the last remaining hardcoded dimension in the precompiled IQN cubin
that could diverge from GpuIqnConfig.embed_dim at runtime, causing
SIGSEGV from out-of-bounds weight offset computation.

All 9 IQN kernels (forward_loss, backward, grad_norm, adam, forward,
trunk_forward, ema, decode_actions, sample_taus) now receive embed_dim
as the final parameter. The iqn_compute_offsets() device helper uses
the runtime embed_dim for W_EMBED→B_EMBED offset calculation.

The other 7 kernel files in the task spec were verified safe:
- c51/mse_loss_kernel.cu: NVRTC-compiled with config values injected
- attention/attention_backward: already parameterized (state_dim, num_heads)
- dt_kernels.cu: NVRTC-compiled with config values injected
- curiosity/ppo: architecturally constant values matching Rust constants

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-26 02:54:19 +01:00
parent 4ab8ca0162
commit f9e55a45d7
2 changed files with 46 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
View File

@@ -340,6 +340,7 @@ impl GpuIqnHead {
let b = self.config.batch_size;
let shared_h1_i32 = self.config.shared_h1 as i32;
let hidden_dim_i32 = self.config.hidden_dim as i32;
let embed_dim_i32 = self.config.embed_dim as i32;
// Layer 3a: τ values are fixed midpoints, pre-computed in constructor.
// No per-step sampling → deterministic → CUDA Graph compatible.
@@ -362,6 +363,7 @@ impl GpuIqnHead {
.arg(&batch_size_i32)
.arg(&shared_h1_i32)
.arg(&hidden_dim_i32)
.arg(&embed_dim_i32)
.launch(decode_config)
.map_err(|e| MLError::ModelError(format!("IQN decode_actions kernel: {e}")))?;
}
@@ -394,6 +396,7 @@ impl GpuIqnHead {
let b = self.config.batch_size;
let shared_h1_i32 = self.config.shared_h1 as i32;
let hidden_dim_i32 = self.config.hidden_dim as i32;
let embed_dim_i32 = self.config.embed_dim as i32;
// 4. Compute target h_s2 via trunk forward kernel
let shmem_bytes = self.config.shared_h1 * 4;
@@ -420,6 +423,7 @@ impl GpuIqnHead {
.arg(&state_dim_i32)
.arg(&shared_h1_i32)
.arg(&hidden_dim_i32)
.arg(&embed_dim_i32)
.launch(trunk_config)
.map_err(|e| MLError::ModelError(format!("IQN trunk forward kernel: {e}")))?;
}
@@ -468,6 +472,7 @@ impl GpuIqnHead {
.arg(&batch_size_i32)
.arg(&shared_h1_i32)
.arg(&hidden_dim_i32)
.arg(&embed_dim_i32)
.launch(fwd_config)
.map_err(|e| MLError::ModelError(format!("IQN forward+loss kernel: {e}")))?;
}
@@ -496,6 +501,7 @@ impl GpuIqnHead {
.arg(&batch_size_i32)
.arg(&shared_h1_i32)
.arg(&hidden_dim_i32)
.arg(&embed_dim_i32)
.launch(bwd_config)
.map_err(|e| MLError::ModelError(format!("IQN backward kernel: {e}")))?;
}
@@ -518,6 +524,7 @@ impl GpuIqnHead {
.arg(&total_params_i32)
.arg(&shared_h1_i32)
.arg(&hidden_dim_i32)
.arg(&embed_dim_i32)
.launch(norm_config)
.map_err(|e| MLError::ModelError(format!("IQN grad_norm kernel: {e}")))?;
}
@@ -558,6 +565,7 @@ impl GpuIqnHead {
.arg(&total_params_i32)
.arg(&shared_h1_i32)
.arg(&hidden_dim_i32)
.arg(&embed_dim_i32)
.launch(adam_config)
.map_err(|e| MLError::ModelError(format!("IQN adam kernel: {e}")))?;
}
@@ -582,6 +590,7 @@ impl GpuIqnHead {
let n_i32 = n as i32;
let shared_h1_i32 = self.config.shared_h1 as i32;
let hidden_dim_i32 = self.config.hidden_dim as i32;
let embed_dim_i32 = self.config.embed_dim as i32;
// Safety: target_params and online_params have total_params elements each.
unsafe {
@@ -593,6 +602,7 @@ impl GpuIqnHead {
.arg(&n_i32)
.arg(&shared_h1_i32)
.arg(&hidden_dim_i32)
.arg(&embed_dim_i32)
.launch(config)
.map_err(|e| MLError::ModelError(format!("IQN EMA kernel: {e}")))?;
}
@@ -649,6 +659,7 @@ impl GpuIqnHead {
let tba = self.config.total_branch_actions();
let shared_h1_i32 = self.config.shared_h1 as i32;
let hidden_dim_i32 = self.config.hidden_dim as i32;
let embed_dim_i32 = self.config.embed_dim as i32;
// 1. Sample τ values for inference
let total_taus = (b * n) as i32;
@@ -668,6 +679,7 @@ impl GpuIqnHead {
.arg(&rng_step)
.arg(&shared_h1_i32)
.arg(&hidden_dim_i32)
.arg(&embed_dim_i32)
.launch(tau_config)
.map_err(|e| MLError::ModelError(format!("IQN CVaR sample_taus: {e}")))?;
}
@@ -689,6 +701,7 @@ impl GpuIqnHead {
.arg(&batch_i32)
.arg(&shared_h1_i32)
.arg(&hidden_dim_i32)
.arg(&embed_dim_i32)
.launch(fwd_config)
.map_err(|e| MLError::ModelError(format!("IQN CVaR forward: {e}")))?;
}

57
crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu
View File

@@ -150,11 +150,11 @@ __device__ __forceinline__ float quantile_huber_grad(float tau, float delta, flo
return -weight * huber_grad;
}
/** Compute runtime weight offsets from hidden_dim. Returns offsets in an array:
/** Compute runtime weight offsets from hidden_dim and embed_dim. Returns offsets in an array:
* [0]=W_EMBED, [1]=B_EMBED, [2]=W_B0, [3]=B_B0, [4]=W_B1, [5]=B_B1, [6]=W_B2, [7]=B_B2 */
__device__ __forceinline__ void iqn_compute_offsets(int hidden_dim, int offsets[8]) {
__device__ __forceinline__ void iqn_compute_offsets(int hidden_dim, int embed_dim, int offsets[8]) {
offsets[0] = 0; /* W_EMBED */
offsets[1] = offsets[0] + hidden_dim * IQN_EMBED_DIM; /* B_EMBED */
offsets[1] = offsets[0] + hidden_dim * embed_dim; /* B_EMBED */
offsets[2] = offsets[1] + hidden_dim; /* W_B0 */
offsets[3] = offsets[2] + BRANCH_0_SIZE * hidden_dim; /* B_B0 */
offsets[4] = offsets[3] + BRANCH_0_SIZE; /* W_B1 */
@@ -202,7 +202,8 @@ void iqn_forward_loss_kernel(
float* __restrict__ save_q_target, /* [B, N, TOTAL_BRANCH_ACTIONS] target Q-values */
int batch_size,
int shared_h1, /* runtime: first hidden layer width (unused here, for consistency) */
int hidden_dim /* runtime: IQN hidden dim = SHARED_H2 */
int hidden_dim, /* runtime: IQN hidden dim = SHARED_H2 */
int embed_dim /* runtime: cosine embedding dimension (was IQN_EMBED_DIM) */
)
{
int sample = blockIdx.x;
@@ -211,7 +212,7 @@ void iqn_forward_loss_kernel(
/* ── Compute runtime weight offsets ── */
int off[8];
iqn_compute_offsets(hidden_dim, off);
iqn_compute_offsets(hidden_dim, embed_dim, off);
/* ── Pointers into this sample's data ── */
const float* my_h_s2 = h_s2 + sample * hidden_dim;
@@ -260,14 +261,14 @@ void iqn_forward_loss_kernel(
for (int t = 0; t < IQN_NUM_QUANTILES; t++) {
float tau = my_taus[t];
/* 1. Cosine features: cos(π·(d+1)·τ) for d=0..IQN_EMBED_DIM-1 */
/* 1. Cosine features: cos(π·(d+1)·τ) for d=0..embed_dim-1 */
/* 2. Linear embedding: embed[h] = Σ_d W_embed[h,d]·cos_feat[d] + b_embed[h] */
/* Process via warp-cooperative matvec: each lane handles hidden_dim/32 outputs */
float embed_dist[IQN_DIST_MAX];
for (int h = lane; h < hidden_dim; h += 32) {
float acc = b_embed[h];
const float* w_row = w_embed + h * IQN_EMBED_DIM;
for (int d = 0; d < IQN_EMBED_DIM; d++) {
const float* w_row = w_embed + h * embed_dim;
for (int d = 0; d < embed_dim; d++) {
float cos_val = cosf(3.14159265f * (float)(d + 1) * tau);
acc += w_row[d] * cos_val;
}
@@ -352,8 +353,8 @@ void iqn_forward_loss_kernel(
float embed_dist[IQN_DIST_MAX];
for (int h = lane; h < hidden_dim; h += 32) {
float acc = tb_embed[h];
const float* w_row = tw_embed + h * IQN_EMBED_DIM;
for (int d = 0; d < IQN_EMBED_DIM; d++) {
const float* w_row = tw_embed + h * embed_dim;
for (int d = 0; d < embed_dim; d++) {
float cos_val = cosf(3.14159265f * (float)(d + 1) * tau);
acc += w_row[d] * cos_val;
}
@@ -457,7 +458,8 @@ void iqn_backward_kernel(
float* __restrict__ d_h_s2_out, /* [B, hidden_dim] dL/d(h_s2) for trunk gradient */
int batch_size,
int shared_h1, /* runtime: first hidden layer width (unused here) */
int hidden_dim /* runtime: IQN hidden dim = SHARED_H2 */
int hidden_dim, /* runtime: IQN hidden dim = SHARED_H2 */
int embed_dim /* runtime: cosine embedding dimension (was IQN_EMBED_DIM) */
)
{
int sample = blockIdx.x;
@@ -466,7 +468,7 @@ void iqn_backward_kernel(
/* ── Compute runtime weight offsets ── */
int off[8];
iqn_compute_offsets(hidden_dim, off);
iqn_compute_offsets(hidden_dim, embed_dim, off);
const float* my_h_s2 = h_s2 + sample * hidden_dim;
const float* my_taus = taus + sample * IQN_NUM_QUANTILES;
@@ -591,9 +593,9 @@ void iqn_backward_kernel(
/* Bias gradient */
atomicAdd(&grad_buf[off[1] + h], dL_dpre);
/* Weight gradient: outer product with cosine features */
for (int d = 0; d < IQN_EMBED_DIM; d++) {
for (int d = 0; d < embed_dim; d++) {
float cos_val = cosf(3.14159265f * (float)(d + 1) * tau);
atomicAdd(&grad_buf[off[0] + h * IQN_EMBED_DIM + d],
atomicAdd(&grad_buf[off[0] + h * embed_dim + d],
dL_dpre * cos_val);
}
}
@@ -607,7 +609,8 @@ void iqn_grad_norm_kernel(
float* __restrict__ norm_out, /* [1] */
int total_params,
int shared_h1, /* runtime: unused, for consistent interface */
int hidden_dim /* runtime: unused, for consistent interface */
int hidden_dim, /* runtime: unused, for consistent interface */
int embed_dim /* runtime: unused, for consistent interface */
)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
@@ -648,7 +651,8 @@ void iqn_adam_kernel(
int adam_t, /* step counter (1-indexed) */
int total_params,
int shared_h1, /* runtime: unused, for consistent interface */
int hidden_dim /* runtime: unused, for consistent interface */
int hidden_dim, /* runtime: unused, for consistent interface */
int embed_dim /* runtime: unused, for consistent interface */
)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
@@ -696,7 +700,8 @@ void iqn_forward_kernel(
float* __restrict__ expected_q, /* [B, TOTAL_BRANCH_ACTIONS] */
int batch_size,
int shared_h1, /* runtime: unused here, for consistency */
int hidden_dim /* runtime: IQN hidden dim = SHARED_H2 */
int hidden_dim, /* runtime: IQN hidden dim = SHARED_H2 */
int embed_dim /* runtime: cosine embedding dimension (was IQN_EMBED_DIM) */
)
{
int sample = blockIdx.x;
@@ -705,7 +710,7 @@ void iqn_forward_kernel(
/* ── Compute runtime weight offsets ── */
int off[8];
iqn_compute_offsets(hidden_dim, off);
iqn_compute_offsets(hidden_dim, embed_dim, off);
const float* my_h_s2 = h_s2 + sample * hidden_dim;
const float* my_taus = taus + sample * IQN_NUM_QUANTILES;
@@ -735,8 +740,8 @@ void iqn_forward_kernel(
float embed_dist[IQN_DIST_MAX];
for (int h = lane; h < hidden_dim; h += 32) {
float acc = b_embed[h];
const float* w_row = w_embed + h * IQN_EMBED_DIM;
for (int d = 0; d < IQN_EMBED_DIM; d++) {
const float* w_row = w_embed + h * embed_dim;
for (int d = 0; d < embed_dim; d++) {
float cos_val = cosf(3.14159265f * (float)(d + 1) * tau);
acc += w_row[d] * cos_val;
}
@@ -808,7 +813,8 @@ void iqn_trunk_forward_kernel(
int batch_size,
int state_dim, /* runtime state dimension (replaces IQN_STATE_DIM) */
int shared_h1, /* runtime: first hidden layer width */
int hidden_dim /* runtime: IQN hidden dim = SHARED_H2 */
int hidden_dim, /* runtime: IQN hidden dim = SHARED_H2 */
int embed_dim /* runtime: unused here, for consistent interface */
)
{
int sample = blockIdx.x;
@@ -852,7 +858,8 @@ void iqn_ema_kernel(
float tau,
int n,
int shared_h1, /* runtime: unused, for consistent interface */
int hidden_dim /* runtime: unused, for consistent interface */
int hidden_dim, /* runtime: unused, for consistent interface */
int embed_dim /* runtime: unused, for consistent interface */
)
{
int i = blockIdx.x * blockDim.x + threadIdx.x;
@@ -874,7 +881,8 @@ void iqn_decode_actions_kernel(
int* __restrict__ branch_actions,
int batch_size,
int shared_h1, /* runtime: unused, for consistent interface */
int hidden_dim /* runtime: unused, for consistent interface */
int hidden_dim, /* runtime: unused, for consistent interface */
int embed_dim /* runtime: unused, for consistent interface */
)
{
int i = blockIdx.x * blockDim.x + threadIdx.x;
@@ -929,7 +937,8 @@ void iqn_sample_taus_kernel(
unsigned int seed,
int total,
int shared_h1, /* runtime: unused, for consistent interface */
int hidden_dim /* runtime: unused, for consistent interface */
int hidden_dim, /* runtime: unused, for consistent interface */
int embed_dim /* runtime: unused, for consistent interface */
)
{
int i = blockIdx.x * blockDim.x + threadIdx.x;