From f9e55a45d7c91b4fe83460b595d4164f83c06bcf Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 26 Mar 2026 02:54:19 +0100 Subject: [PATCH] fix: parameterize IQN_EMBED_DIM in all 9 IQN CUDA kernels MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- crates/ml/src/cuda_pipeline/gpu_iqn_head.rs | 13 +++++ .../src/cuda_pipeline/iqn_dual_head_kernel.cu | 57 +++++++++++-------- 2 files changed, 46 insertions(+), 24 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs index 986ebb3c2..e57518160 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs @@ -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}")))?; } diff --git a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu index afacff86e..ae5d2db53 100644 --- a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu @@ -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;