From 6ff45466f985aa34b01d61d3ce5897917bb7f538 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 6 Apr 2026 23:20:11 +0200 Subject: [PATCH] =?UTF-8?q?perf:=20IQN=20kernel=20=E2=80=94=20block=2032?= =?UTF-8?q?=E2=86=92256,=20precompute=20cosines,=20parallel=20loss?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three optimizations to IQN forward+loss and backward kernels: 1. Block size 32→256 (8 warps): occupancy 3.1%→25% on H100. Inner hidden_dim loops now stride by 256 instead of 32. Block-level reduction via shared memory replaces warp-only shuffle. 2. Precomputed cosine features [N, embed_dim]: eliminates 16.7M cosf() calls per step. cos(π·(d+1)·τ_i) computed once at construction (τ are fixed QR-DQN midpoints). 3. Quantile Huber loss distributed across 256 threads: was single-lane serial (32×32=1024 iterations on lane 0). Now each thread handles ~4 pairs. Expected impact: IQN from ~80ms to ~10-15ms per step (occupancy + cosine elimination + parallel loss). Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/cuda_pipeline/gpu_iqn_head.rs | 54 ++- .../src/cuda_pipeline/iqn_dual_head_kernel.cu | 353 ++++++++++-------- 2 files changed, 243 insertions(+), 164 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs index 2c2c1d50b..a7a22db93 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs @@ -167,6 +167,9 @@ pub struct GpuIqnHead { online_taus: CudaSlice, /// Pre-sampled τ values for target network [B, N] target_taus: CudaSlice, + /// Precomputed cosine features [N, embed_dim] — cos(π·(d+1)·τ_i) + /// Fixed at construction (τ are midpoints). Eliminates 16.7M cosf() per step. + cos_features: CudaSlice, /// Branch actions decoded from flat actions [B, 3] branch_actions: CudaSlice, /// Target h_s2 computed from next_states + target trunk weights [B, H] @@ -229,14 +232,33 @@ impl GpuIqnHead { let d_h_s2_buf = alloc_f32(&stream, b * h, "iqn_d_h_s2")?; // Per-step buffers — fixed τ midpoints (QR-DQN style). - // τ_i = (2i - 1) / (2N) for i = 1..N. Deterministic → CUDA Graph compatible. + // τ_i = (2i + 1) / (2N) for i = 0..N-1. Deterministic → CUDA Graph compatible. // No per-step random sampling. IQN cosine embedding still operates on these. let online_taus; let target_taus; + let cos_features; { + let d = config.embed_dim; let fixed: Vec = (0..n) .map(|i| (2.0 * (i as f32) + 1.0) / (2.0 * n as f32)) .collect(); + + // Precompute cosine features: cos_features[i * embed_dim + d] = cos(π·(d+1)·τ_i) + // These are constant across all training steps (τ are fixed midpoints). + // Eliminates ~16.7M cosf() calls per step in the CUDA kernels. + let mut cos_feat_host = Vec::with_capacity(n * d); + for i in 0..n { + let tau_i = fixed[i]; + for dim in 0..d { + cos_feat_host.push( + (std::f32::consts::PI * ((dim + 1) as f32) * tau_i).cos() + ); + } + } + cos_features = stream.clone_htod(&cos_feat_host).map_err(|e| { + MLError::ModelError(format!("IQN htod cos_features ({} f32): {e}", n * d)) + })?; + let mut tiled = Vec::with_capacity(b * n); for _ in 0..b { tiled.extend_from_slice(&fixed); @@ -309,6 +331,7 @@ impl GpuIqnHead { d_h_s2_buf, online_taus, target_taus, + cos_features, branch_actions, target_h_s2, rewards_buf, @@ -412,11 +435,11 @@ impl GpuIqnHead { 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 + // 4. Compute target h_s2 via trunk forward kernel (256 threads per block) let shmem_bytes = self.config.shared_h1 * 4; let trunk_config = LaunchConfig { grid_dim: (b as u32, 1, 1), - block_dim: (32, 1, 1), + block_dim: (256, 1, 1), shared_mem_bytes: shmem_bytes as u32, }; @@ -455,11 +478,13 @@ impl GpuIqnHead { let batch_size_i32 = b as i32; let gamma = self.config.gamma; - // 6. Forward + loss kernel: one warp per sample + // 6. Forward + loss kernel: 256 threads per sample (8 warps) + // Shared memory: 8 floats for cross-warp block reduction + let fwd_shmem = (256 / 32) * 4; // 8 warps × sizeof(f32) let fwd_config = LaunchConfig { grid_dim: (b as u32, 1, 1), - block_dim: (32, 1, 1), - shared_mem_bytes: 0, + block_dim: (256, 1, 1), + shared_mem_bytes: fwd_shmem as u32, }; // Safety: all buffer pointers are valid GPU allocations on the same context @@ -477,6 +502,7 @@ impl GpuIqnHead { .arg(&gamma) .arg(&self.online_params) .arg(&self.target_params) + .arg(&self.cos_features) .arg(&mut self.per_sample_loss) .arg(&mut self.total_loss) .arg(&mut self.save_embed) @@ -491,11 +517,12 @@ impl GpuIqnHead { .map_err(|e| MLError::ModelError(format!("IQN forward+loss kernel: {e}")))?; } - // 7. Backward kernel: one warp per sample + // 7. Backward kernel: 256 threads per sample (8 warps) + let bwd_shmem = (256 / 32) * 4; // 8 warps × sizeof(f32) let bwd_config = LaunchConfig { grid_dim: (b as u32, 1, 1), - block_dim: (32, 1, 1), - shared_mem_bytes: 0, + block_dim: (256, 1, 1), + shared_mem_bytes: bwd_shmem as u32, }; // Safety: all pointers valid, save buffers written by forward kernel on same stream. @@ -510,6 +537,7 @@ impl GpuIqnHead { .arg(&self.online_taus) .arg(&self.branch_actions) .arg(&self.online_params) + .arg(&self.cos_features) .arg(&mut self.grad_buf) .arg(&mut self.d_h_s2_buf) // IQN trunk gradient output .arg(&batch_size_i32) @@ -734,11 +762,12 @@ impl GpuIqnHead { .map_err(|e| MLError::ModelError(format!("IQN CVaR sample_taus: {e}")))?; } - // 2. Run IQN forward-only kernel + // 2. Run IQN forward-only kernel (256 threads per sample) + let fwd_inf_shmem = (256 / 32) * 4; // 8 warps × sizeof(f32) let fwd_config = LaunchConfig { grid_dim: (b as u32, 1, 1), - block_dim: (32, 1, 1), - shared_mem_bytes: 0, + block_dim: (256, 1, 1), + shared_mem_bytes: fwd_inf_shmem as u32, }; let batch_i32 = b as i32; unsafe { @@ -747,6 +776,7 @@ impl GpuIqnHead { .arg(h_s2) .arg(&self.online_taus) .arg(&self.online_params) + .arg(&self.cos_features) .arg(&mut self.save_q_online) .arg(&batch_i32) .arg(&shared_h1_i32) 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 5a24fda43..c340a5a67 100644 --- a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu @@ -21,8 +21,8 @@ * L_iqn = (1/N) Σ_i (1/N') Σ_j ρ_τi(δ_ij) * * Launch config: - * Forward + loss: grid=(batch_size, 1, 1), block=(32, 1, 1) - * Backward: grid=(batch_size, 1, 1), block=(32, 1, 1) + * Forward + loss: grid=(batch_size, 1, 1), block=(256, 1, 1), shmem=8*4 bytes + * Backward: grid=(batch_size, 1, 1), block=(256, 1, 1), shmem=8*4 bytes * Grad norm: grid=(ceil(total_params/256), 1, 1), block=(256, 1, 1) * Adam: grid=(ceil(total_params/256), 1, 1), block=(256, 1, 1) * @@ -104,11 +104,14 @@ #define IQN_OFF_W_B2 (IQN_OFF_B_B1 + IQN_B_B1_SIZE) #define IQN_OFF_B_B2 (IQN_OFF_W_B2 + IQN_W_B2_SIZE) -/* Max distributed array size: ceil(IQN_HIDDEN / 32) — used for register arrays */ -#define IQN_DIST_MAX (((IQN_HIDDEN) + 31) / 32) +/* Block size for IQN forward/backward/trunk kernels (8 warps = 256 threads). */ +#define IQN_BLOCK_SIZE 256 -/* Distributed array size: ceil(dim / 32) elements per lane */ -#define IQN_DIST(dim) (((dim) + 31) / 32) +/* Max distributed array size: ceil(IQN_HIDDEN / IQN_BLOCK_SIZE) — used for register arrays */ +#define IQN_DIST_MAX (((IQN_HIDDEN) + IQN_BLOCK_SIZE - 1) / IQN_BLOCK_SIZE) + +/* Distributed array size: ceil(dim / IQN_BLOCK_SIZE) elements per thread */ +#define IQN_DIST(dim) (((dim) + IQN_BLOCK_SIZE - 1) / IQN_BLOCK_SIZE) #ifndef IQN_STATE_DIM #define IQN_STATE_DIM 48 @@ -118,7 +121,7 @@ #endif /* Max shared_h1 for register array sizing */ -#define IQN_SHARED_H1_MAX (((IQN_SHARED_H1) + 31) / 32) +#define IQN_SHARED_H1_MAX (((IQN_SHARED_H1) + IQN_BLOCK_SIZE - 1) / IQN_BLOCK_SIZE) /* ── Device helpers ──────────────────────────────────────────────────── */ @@ -129,6 +132,32 @@ __device__ __forceinline__ float iqn_warp_sum(float val) { return val; } +/** Block-level reduce sum for IQN_BLOCK_SIZE (256) threads. + * First reduces within each warp via shuffle, then reduces + * across warps via shared memory. Returns the sum on ALL threads + * (via broadcast). Requires shmem_reduce[8] in shared memory. */ +__device__ __forceinline__ float iqn_block_sum(float val, float* shmem_reduce) { + int lane = threadIdx.x & 31; + int warp_id = threadIdx.x >> 5; + /* Intra-warp reduction */ + val = iqn_warp_sum(val); + /* Lane 0 of each warp writes to shared memory */ + if (lane == 0) + shmem_reduce[warp_id] = val; + __syncthreads(); + /* First warp reduces across all 8 warp partial sums */ + float result = 0.0f; + if (warp_id == 0) { + result = (lane < (IQN_BLOCK_SIZE / 32)) ? shmem_reduce[lane] : 0.0f; + result = iqn_warp_sum(result); + } + /* Broadcast result from thread 0 to all threads via shared memory */ + if (threadIdx.x == 0) + shmem_reduce[0] = result; + __syncthreads(); + return shmem_reduce[0]; +} + /** Quantile Huber loss element: ρ_τ(δ) = |τ - 𝟙{δ<0}| × L_κ(δ) */ __device__ __forceinline__ float quantile_huber_element(float tau, float delta, float kappa) { float abs_delta = fabsf(delta); @@ -198,6 +227,8 @@ void iqn_forward_loss_kernel( /* Weights (f32) */ const float* __restrict__ online_params, /* online IQN weights */ const float* __restrict__ target_params, /* target IQN weights */ + /* Precomputed cosine features (Optimization C) */ + const float* __restrict__ cos_features, /* [N, embed_dim] — cos(π·(d+1)·τ_i) */ /* Outputs (f32) */ float* __restrict__ per_sample_loss, /* [B] per-sample loss (for PER) */ float* __restrict__ total_loss, /* [1] batch-mean loss (atomicAdd) */ @@ -214,7 +245,10 @@ void iqn_forward_loss_kernel( { int sample = blockIdx.x; if (sample >= batch_size) return; - int lane = threadIdx.x; + int tid = threadIdx.x; + + /* Shared memory for block-level reductions (8 warps × 1 float) */ + __shared__ float shmem_reduce[IQN_BLOCK_SIZE / 32]; /* ── Compute runtime weight offsets ── */ int off[8]; @@ -253,45 +287,44 @@ void iqn_forward_loss_kernel( /* ── Load h_s2 into distributed registers (bf16 → f32 at boundary) ── */ float h_dist[IQN_DIST_MAX]; - for (int d = lane; d < hidden_dim; d += 32) - h_dist[d / 32] = (float)my_h_s2_bf16[d]; + for (int d = tid; d < hidden_dim; d += IQN_BLOCK_SIZE) + h_dist[d / IQN_BLOCK_SIZE] = (float)my_h_s2_bf16[d]; float h_target_dist[IQN_DIST_MAX]; - for (int d = lane; d < hidden_dim; d += 32) - h_target_dist[d / 32] = (float)my_target_h_s2_bf16[d]; + for (int d = tid; d < hidden_dim; d += IQN_BLOCK_SIZE) + h_target_dist[d / IQN_BLOCK_SIZE] = (float)my_target_h_s2_bf16[d]; /* ── Process each ONLINE quantile ── */ /* For each τ_i, compute quantile Q-values for the taken actions */ float online_q_a[IQN_NUM_QUANTILES]; /* Q(s, a_taken, τ_i) per branch sum */ for (int t = 0; t < IQN_NUM_QUANTILES; t++) { - float tau = my_taus[t]; + /* Precomputed cosine features: cos_features[t * embed_dim + d] */ + const float* my_cos = cos_features + t * embed_dim; - /* 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 */ + /* 1+2. Linear embedding using precomputed cosine features: + * embed[h] = ReLU(Σ_d W_embed[h,d]·cos_feat[d] + b_embed[h]) */ float embed_dist[IQN_DIST_MAX]; - for (int h = lane; h < hidden_dim; h += 32) { + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { float acc = b_embed[h]; 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; + acc += w_row[d] * my_cos[d]; } /* ReLU */ - embed_dist[h / 32] = fmaxf(acc, 0.0f); + embed_dist[h / IQN_BLOCK_SIZE] = fmaxf(acc, 0.0f); } /* 3. Element-wise product: combined = h_s2 ⊙ embed */ float comb_dist[IQN_DIST_MAX]; - for (int h = lane; h < hidden_dim; h += 32) - comb_dist[h / 32] = h_dist[h / 32] * embed_dist[h / 32]; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + comb_dist[h / IQN_BLOCK_SIZE] = h_dist[h / IQN_BLOCK_SIZE] * embed_dist[h / IQN_BLOCK_SIZE]; /* Save activations for backward pass */ int save_offset = (sample * IQN_NUM_QUANTILES + t) * hidden_dim; - for (int h = lane; h < hidden_dim; h += 32) { - save_embed[save_offset + h] = embed_dist[h / 32]; - save_combined[save_offset + h] = comb_dist[h / 32]; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { + save_embed[save_offset + h] = embed_dist[h / IQN_BLOCK_SIZE]; + save_combined[save_offset + h] = comb_dist[h / IQN_BLOCK_SIZE]; } /* 4. Per-branch output: q_d = W_bd × combined + b_bd */ @@ -300,50 +333,54 @@ void iqn_forward_loss_kernel( /* Branch 0 (exposure) */ float q0_taken = 0.0f; for (int a = 0; a < BRANCH_0_SIZE; a++) { - float acc = b_b0[a]; - const float* w_row = w_b0 + a * hidden_dim; float partial = 0.0f; - for (int h = lane; h < hidden_dim; h += 32) - partial += w_row[h] * comb_dist[h / 32]; - acc += iqn_warp_sum(partial); - if (lane == 0) { + const float* w_row = w_b0 + a * hidden_dim; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; + float acc = b_b0[a] + iqn_block_sum(partial, shmem_reduce); + if (tid == 0) { save_q_online[q_save_offset + a] = acc; if (a == a0) q0_taken = acc; } - if (a == a0) q0_taken = __shfl_sync(0xFFFFFFFF, q0_taken, 0); } + /* Broadcast q0_taken from thread 0 */ + if (tid == 0) shmem_reduce[0] = q0_taken; + __syncthreads(); + q0_taken = shmem_reduce[0]; /* Branch 1 (order) */ float q1_taken = 0.0f; for (int a = 0; a < BRANCH_1_SIZE; a++) { - float acc = b_b1[a]; - const float* w_row = w_b1 + a * hidden_dim; float partial = 0.0f; - for (int h = lane; h < hidden_dim; h += 32) - partial += w_row[h] * comb_dist[h / 32]; - acc += iqn_warp_sum(partial); - if (lane == 0) { + const float* w_row = w_b1 + a * hidden_dim; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; + float acc = b_b1[a] + iqn_block_sum(partial, shmem_reduce); + if (tid == 0) { save_q_online[q_save_offset + BRANCH_0_SIZE + a] = acc; if (a == a1) q1_taken = acc; } - if (a == a1) q1_taken = __shfl_sync(0xFFFFFFFF, q1_taken, 0); } + if (tid == 0) shmem_reduce[0] = q1_taken; + __syncthreads(); + q1_taken = shmem_reduce[0]; /* Branch 2 (urgency) */ float q2_taken = 0.0f; for (int a = 0; a < BRANCH_2_SIZE; a++) { - float acc = b_b2[a]; - const float* w_row = w_b2 + a * hidden_dim; float partial = 0.0f; - for (int h = lane; h < hidden_dim; h += 32) - partial += w_row[h] * comb_dist[h / 32]; - acc += iqn_warp_sum(partial); - if (lane == 0) { + const float* w_row = w_b2 + a * hidden_dim; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; + float acc = b_b2[a] + iqn_block_sum(partial, shmem_reduce); + if (tid == 0) { save_q_online[q_save_offset + BRANCH_0_SIZE + BRANCH_1_SIZE + a] = acc; if (a == a2) q2_taken = acc; } - if (a == a2) q2_taken = __shfl_sync(0xFFFFFFFF, q2_taken, 0); } + if (tid == 0) shmem_reduce[0] = q2_taken; + __syncthreads(); + q2_taken = shmem_reduce[0]; /* Sum of taken-action Q-values across branches (for loss computation) */ online_q_a[t] = q0_taken + q1_taken + q2_taken; @@ -353,54 +390,51 @@ void iqn_forward_loss_kernel( float target_q_a[IQN_NUM_QUANTILES]; for (int t = 0; t < IQN_NUM_QUANTILES; t++) { - float tau = my_target_taus[t]; + /* Target uses same fixed τ midpoints, same precomputed cosine features */ + const float* my_cos = cos_features + t * embed_dim; - /* Cosine embedding + linear (target weights) */ + /* Cosine embedding + linear (target weights, precomputed cos_features) */ float embed_dist[IQN_DIST_MAX]; - for (int h = lane; h < hidden_dim; h += 32) { + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { float acc = tb_embed[h]; 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; + acc += w_row[d] * my_cos[d]; } - embed_dist[h / 32] = fmaxf(acc, 0.0f); + embed_dist[h / IQN_BLOCK_SIZE] = fmaxf(acc, 0.0f); } /* Element-wise product with target h_s2 */ float comb_dist[IQN_DIST_MAX]; - for (int h = lane; h < hidden_dim; h += 32) - comb_dist[h / 32] = h_target_dist[h / 32] * embed_dist[h / 32]; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + comb_dist[h / IQN_BLOCK_SIZE] = h_target_dist[h / IQN_BLOCK_SIZE] * embed_dist[h / IQN_BLOCK_SIZE]; /* Per-branch Q-values for taken actions */ - float q0 = 0.0f, q1 = 0.0f, q2 = 0.0f; + float q0, q1, q2; /* Branch 0 */ { - float acc = tb_b0[a0]; float partial = 0.0f; const float* w_row = tw_b0 + a0 * hidden_dim; - for (int h = lane; h < hidden_dim; h += 32) - partial += w_row[h] * comb_dist[h / 32]; - q0 = acc + iqn_warp_sum(partial); + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; + q0 = tb_b0[a0] + iqn_block_sum(partial, shmem_reduce); } /* Branch 1 */ { - float acc = tb_b1[a1]; float partial = 0.0f; const float* w_row = tw_b1 + a1 * hidden_dim; - for (int h = lane; h < hidden_dim; h += 32) - partial += w_row[h] * comb_dist[h / 32]; - q1 = acc + iqn_warp_sum(partial); + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; + q1 = tb_b1[a1] + iqn_block_sum(partial, shmem_reduce); } /* Branch 2 */ { - float acc = tb_b2[a2]; float partial = 0.0f; const float* w_row = tw_b2 + a2 * hidden_dim; - for (int h = lane; h < hidden_dim; h += 32) - partial += w_row[h] * comb_dist[h / 32]; - q2 = acc + iqn_warp_sum(partial); + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; + q2 = tb_b2[a2] + iqn_block_sum(partial, shmem_reduce); } /* Bellman target: r + γ(1-done) × Q_target */ @@ -409,7 +443,7 @@ void iqn_forward_loss_kernel( /* Save target Q for backward */ int q_save_offset = (sample * IQN_NUM_QUANTILES + t) * TOTAL_BRANCH_ACTIONS; - if (lane == 0) { + if (tid == 0) { save_q_target[q_save_offset + a0] = target_q_a[t]; save_q_target[q_save_offset + BRANCH_0_SIZE + a1] = target_q_a[t]; save_q_target[q_save_offset + BRANCH_0_SIZE + BRANCH_1_SIZE + a2] = target_q_a[t]; @@ -417,21 +451,23 @@ void iqn_forward_loss_kernel( } /* ── Quantile Huber Loss: (1/N) Σ_i (1/N') Σ_j ρ_τi(δ_ij) ── */ + /* Distribute the N×N' pairs across all 256 threads. + * Total pairs = IQN_NUM_QUANTILES × IQN_NUM_QUANTILES = 32×32 = 1024. + * 256 threads → 4 pairs per thread. */ float kappa = IQN_KAPPA; - float inv_nq = 1.0f / (float)IQN_NUM_QUANTILES; - float sample_loss = 0.0f; - if (lane == 0) { - for (int i = 0; i < IQN_NUM_QUANTILES; i++) { - float tau_i = my_taus[i]; - float inner_sum = 0.0f; - for (int j = 0; j < IQN_NUM_QUANTILES; j++) { - float delta = target_q_a[j] - online_q_a[i]; - inner_sum += quantile_huber_element(tau_i, delta, kappa); - } - sample_loss += inner_sum * inv_nq; - } - sample_loss *= inv_nq; + float inv_nq_sq = 1.0f / ((float)IQN_NUM_QUANTILES * (float)IQN_NUM_QUANTILES); + int total_pairs = IQN_NUM_QUANTILES * IQN_NUM_QUANTILES; + float my_loss = 0.0f; + for (int p = tid; p < total_pairs; p += IQN_BLOCK_SIZE) { + int i = p / IQN_NUM_QUANTILES; + int j = p % IQN_NUM_QUANTILES; + float tau_i = my_taus[i]; + float delta = target_q_a[j] - online_q_a[i]; + my_loss += quantile_huber_element(tau_i, delta, kappa); + } + float sample_loss = iqn_block_sum(my_loss, shmem_reduce) * inv_nq_sq; + if (tid == 0) { per_sample_loss[sample] = sample_loss; atomicAdd(total_loss, sample_loss * (1.0f / (float)batch_size)); } @@ -461,6 +497,8 @@ void iqn_backward_kernel( const int* __restrict__ actions, /* [B, 3] */ /* Weights (f32, for backward through linear layers) */ const float* __restrict__ online_params, + /* Precomputed cosine features (Optimization C) */ + const float* __restrict__ cos_features, /* [N, embed_dim] — cos(π·(d+1)·τ_i) */ /* Gradient output (f32) */ float* __restrict__ grad_buf, /* accumulated gradients */ float* __restrict__ d_h_s2_out, /* [B, hidden_dim] dL/d(h_s2) for trunk gradient */ @@ -472,7 +510,10 @@ void iqn_backward_kernel( { int sample = blockIdx.x; if (sample >= batch_size) return; - int lane = threadIdx.x; + int tid = threadIdx.x; + + /* Shared memory for block-level reductions */ + __shared__ float shmem_reduce[IQN_BLOCK_SIZE / 32]; /* ── Compute runtime weight offsets ── */ int off[8]; @@ -491,23 +532,31 @@ void iqn_backward_kernel( /* Load h_s2 into distributed registers (bf16 → f32 at boundary) */ float h_dist[IQN_DIST_MAX]; - for (int h = lane; h < hidden_dim; h += 32) - h_dist[h / 32] = (float)my_h_s2_bf16[h]; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + h_dist[h / IQN_BLOCK_SIZE] = (float)my_h_s2_bf16[h]; /* Pre-compute online and target Q-values for taken actions across all quantiles */ - /* (reading from saved f32 buffers) */ + /* (reading from saved f32 buffers — only thread 0 reads, then broadcast) */ float online_q[IQN_NUM_QUANTILES]; float target_q[IQN_NUM_QUANTILES]; for (int t = 0; t < IQN_NUM_QUANTILES; t++) { int qoff = (sample * IQN_NUM_QUANTILES + t) * TOTAL_BRANCH_ACTIONS; - if (lane == 0) { + if (tid == 0) { online_q[t] = save_q_online[qoff + a0] + save_q_online[qoff + BRANCH_0_SIZE + a1] + save_q_online[qoff + BRANCH_0_SIZE + BRANCH_1_SIZE + a2]; target_q[t] = save_q_target[qoff + a0]; /* target_q already contains Bellman target */ } - online_q[t] = __shfl_sync(0xFFFFFFFF, online_q[t], 0); - target_q[t] = __shfl_sync(0xFFFFFFFF, target_q[t], 0); + } + /* Broadcast all Q-values from thread 0 to all threads via shared memory */ + /* Process in chunks fitting the shared memory */ + for (int t = 0; t < IQN_NUM_QUANTILES; t++) { + if (tid == 0) shmem_reduce[0] = online_q[t]; + __syncthreads(); + online_q[t] = shmem_reduce[0]; + if (tid == 0) shmem_reduce[0] = target_q[t]; + __syncthreads(); + target_q[t] = shmem_reduce[0]; } /* Process each quantile τ_i */ @@ -518,58 +567,55 @@ void iqn_backward_kernel( float tau_i = my_taus[ti]; int save_offset = (sample * IQN_NUM_QUANTILES + ti) * hidden_dim; - /* Compute dL/dq_online for this quantile: + /* Compute dL/dq_online for this quantile — distribute across threads: * dL/dq_i = (1/N²) Σ_j dρ_τi(δ_ij)/dq_i */ - float dL_dq = 0.0f; - if (lane == 0) { - for (int tj = 0; tj < IQN_NUM_QUANTILES; tj++) { - float delta = target_q[tj] - online_q[ti]; - dL_dq += quantile_huber_grad(tau_i, delta, kappa); - } - dL_dq *= inv_n_sq; + float my_dL_dq = 0.0f; + for (int tj = tid; tj < IQN_NUM_QUANTILES; tj += IQN_BLOCK_SIZE) { + float delta = target_q[tj] - online_q[ti]; + my_dL_dq += quantile_huber_grad(tau_i, delta, kappa); } - dL_dq = __shfl_sync(0xFFFFFFFF, dL_dq, 0); + float dL_dq = iqn_block_sum(my_dL_dq, shmem_reduce) * inv_n_sq; /* Load saved activations */ float embed_dist[IQN_DIST_MAX]; float comb_dist[IQN_DIST_MAX]; - for (int h = lane; h < hidden_dim; h += 32) { - embed_dist[h / 32] = save_embed[save_offset + h]; - comb_dist[h / 32] = save_combined[save_offset + h]; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { + embed_dist[h / IQN_BLOCK_SIZE] = save_embed[save_offset + h]; + comb_dist[h / IQN_BLOCK_SIZE] = save_combined[save_offset + h]; } /* ── Gradient through branch output layers (taken action only) ── */ /* dL/d(combined) = Σ_d W_bd[a_d, :] × dL/dq (only for taken actions) */ float dL_dcomb[IQN_DIST_MAX]; - for (int h = lane; h < hidden_dim; h += 32) - dL_dcomb[h / 32] = 0.0f; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + dL_dcomb[h / IQN_BLOCK_SIZE] = 0.0f; /* Branch 0: accumulate gradient */ - for (int h = lane; h < hidden_dim; h += 32) { - dL_dcomb[h / 32] += w_b0[a0 * hidden_dim + h] * dL_dq; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { + dL_dcomb[h / IQN_BLOCK_SIZE] += w_b0[a0 * hidden_dim + h] * dL_dq; /* dL/dW_b0[a0, h] += dL/dq × combined[h] */ atomicAdd(&grad_buf[off[2] + a0 * hidden_dim + h], - dL_dq * comb_dist[h / 32]); + dL_dq * comb_dist[h / IQN_BLOCK_SIZE]); } - if (lane == 0) + if (tid == 0) atomicAdd(&grad_buf[off[3] + a0], dL_dq); /* Branch 1 */ - for (int h = lane; h < hidden_dim; h += 32) { - dL_dcomb[h / 32] += w_b1[a1 * hidden_dim + h] * dL_dq; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { + dL_dcomb[h / IQN_BLOCK_SIZE] += w_b1[a1 * hidden_dim + h] * dL_dq; atomicAdd(&grad_buf[off[4] + a1 * hidden_dim + h], - dL_dq * comb_dist[h / 32]); + dL_dq * comb_dist[h / IQN_BLOCK_SIZE]); } - if (lane == 0) + if (tid == 0) atomicAdd(&grad_buf[off[5] + a1], dL_dq); /* Branch 2 */ - for (int h = lane; h < hidden_dim; h += 32) { - dL_dcomb[h / 32] += w_b2[a2 * hidden_dim + h] * dL_dq; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { + dL_dcomb[h / IQN_BLOCK_SIZE] += w_b2[a2 * hidden_dim + h] * dL_dq; atomicAdd(&grad_buf[off[6] + a2 * hidden_dim + h], - dL_dq * comb_dist[h / 32]); + dL_dq * comb_dist[h / IQN_BLOCK_SIZE]); } - if (lane == 0) + if (tid == 0) atomicAdd(&grad_buf[off[7] + a2], dL_dq); /* ── Gradient through element-wise product ── */ @@ -577,37 +623,36 @@ void iqn_backward_kernel( * dL/d(embed) = dL/d(combined) ⊙ h_s2 * dL/d(h_s2) = dL/d(combined) ⊙ embed — NOW COMPUTED for trunk gradient */ float dL_dembed[IQN_DIST_MAX]; - for (int h = lane; h < hidden_dim; h += 32) { - dL_dembed[h / 32] = dL_dcomb[h / 32] * h_dist[h / 32]; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { + dL_dembed[h / IQN_BLOCK_SIZE] = dL_dcomb[h / IQN_BLOCK_SIZE] * h_dist[h / IQN_BLOCK_SIZE]; /* Accumulate dL/d(h_s2) across all quantiles. * This flows IQN's bounded Huber gradient to the shared trunk. */ - float d_trunk = dL_dcomb[h / 32] * embed_dist[h / 32]; + float d_trunk = dL_dcomb[h / IQN_BLOCK_SIZE] * embed_dist[h / IQN_BLOCK_SIZE]; if (isfinite(d_trunk)) atomicAdd(&d_h_s2_out[sample * hidden_dim + h], d_trunk); } /* ── Gradient through ReLU ── */ /* embed = ReLU(pre_relu) → dL/d(pre_relu) = dL/d(embed) × 𝟙{embed > 0} */ - for (int h = lane; h < hidden_dim; h += 32) { - if (!(embed_dist[h / 32] > 0.0f)) - dL_dembed[h / 32] = 0.0f; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { + if (!(embed_dist[h / IQN_BLOCK_SIZE] > 0.0f)) + dL_dembed[h / IQN_BLOCK_SIZE] = 0.0f; } /* ── Gradient through embedding linear layer ── */ /* pre_relu[h] = Σ_d W_embed[h,d]·cos(π(d+1)τ) + b_embed[h] * dL/dW_embed[h,d] += dL/d(pre_relu)[h] × cos(π(d+1)τ) * dL/db_embed[h] += dL/d(pre_relu)[h] */ - float tau = my_taus[ti]; - for (int h = lane; h < hidden_dim; h += 32) { - float dL_dpre = dL_dembed[h / 32]; + const float* my_cos = cos_features + ti * embed_dim; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { + float dL_dpre = dL_dembed[h / IQN_BLOCK_SIZE]; if (!isfinite(dL_dpre)) continue; /* Bias gradient */ atomicAdd(&grad_buf[off[1] + h], dL_dpre); - /* Weight gradient: outer product with cosine features */ + /* Weight gradient: outer product with precomputed cosine features */ for (int d = 0; d < embed_dim; d++) { - float cos_val = cosf(3.14159265f * (float)(d + 1) * tau); atomicAdd(&grad_buf[off[0] + h * embed_dim + d], - dL_dpre * cos_val); + dL_dpre * my_cos[d]); } } } @@ -717,6 +762,7 @@ void iqn_forward_kernel( const __nv_bfloat16* __restrict__ h_s2, /* [B, hidden_dim] (bf16 from DQN trunk) */ const float* __restrict__ taus, /* [B, N] (f32) */ const float* __restrict__ params, /* IQN weights (f32) */ + const float* __restrict__ cos_features, /* [N, embed_dim] precomputed cosine features */ float* __restrict__ expected_q, /* [B, TOTAL_BRANCH_ACTIONS] (f32) */ int batch_size, int shared_h1, /* runtime: unused here, for consistency */ @@ -726,7 +772,10 @@ void iqn_forward_kernel( { int sample = blockIdx.x; if (sample >= batch_size) return; - int lane = threadIdx.x; + int tid = threadIdx.x; + + /* Shared memory for block-level reductions */ + __shared__ float shmem_reduce[IQN_BLOCK_SIZE / 32]; /* ── Compute runtime weight offsets ── */ int off[8]; @@ -745,8 +794,8 @@ void iqn_forward_kernel( /* Load h_s2 (bf16 → f32 at boundary) */ float h_dist[IQN_DIST_MAX]; - for (int h = lane; h < hidden_dim; h += 32) - h_dist[h / 32] = (float)my_h_s2_bf16[h]; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + h_dist[h / IQN_BLOCK_SIZE] = (float)my_h_s2_bf16[h]; /* Accumulate Q-values across quantiles (for mean) */ float q_acc[TOTAL_BRANCH_ACTIONS]; @@ -754,58 +803,58 @@ void iqn_forward_kernel( q_acc[a] = 0.0f; for (int t = 0; t < IQN_NUM_QUANTILES; t++) { - float tau = my_taus[t]; + /* Precomputed cosine features for this quantile */ + const float* my_cos = cos_features + t * embed_dim; - /* Cosine embedding + linear + ReLU */ + /* Cosine embedding + linear + ReLU (using precomputed cos_features) */ float embed_dist[IQN_DIST_MAX]; - for (int h = lane; h < hidden_dim; h += 32) { + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { float acc = b_embed[h]; 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; + acc += w_row[d] * my_cos[d]; } - embed_dist[h / 32] = fmaxf(acc, 0.0f); + embed_dist[h / IQN_BLOCK_SIZE] = fmaxf(acc, 0.0f); } /* Element-wise product */ float comb_dist[IQN_DIST_MAX]; - for (int h = lane; h < hidden_dim; h += 32) - comb_dist[h / 32] = h_dist[h / 32] * embed_dist[h / 32]; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + comb_dist[h / IQN_BLOCK_SIZE] = h_dist[h / IQN_BLOCK_SIZE] * embed_dist[h / IQN_BLOCK_SIZE]; /* Branch outputs */ /* Branch 0 */ for (int a = 0; a < BRANCH_0_SIZE; a++) { float partial = 0.0f; const float* w_row = w_b0 + a * hidden_dim; - for (int h = lane; h < hidden_dim; h += 32) - partial += w_row[h] * comb_dist[h / 32]; - float q_val = iqn_warp_sum(partial) + b_b0[a]; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; + float q_val = iqn_block_sum(partial, shmem_reduce) + b_b0[a]; q_acc[a] += q_val; } /* Branch 1 */ for (int a = 0; a < BRANCH_1_SIZE; a++) { float partial = 0.0f; const float* w_row = w_b1 + a * hidden_dim; - for (int h = lane; h < hidden_dim; h += 32) - partial += w_row[h] * comb_dist[h / 32]; - float q_val = iqn_warp_sum(partial) + b_b1[a]; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; + float q_val = iqn_block_sum(partial, shmem_reduce) + b_b1[a]; q_acc[BRANCH_0_SIZE + a] += q_val; } /* Branch 2 */ for (int a = 0; a < BRANCH_2_SIZE; a++) { float partial = 0.0f; const float* w_row = w_b2 + a * hidden_dim; - for (int h = lane; h < hidden_dim; h += 32) - partial += w_row[h] * comb_dist[h / 32]; - float q_val = iqn_warp_sum(partial) + b_b2[a]; + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) + partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE]; + float q_val = iqn_block_sum(partial, shmem_reduce) + b_b2[a]; q_acc[BRANCH_0_SIZE + BRANCH_1_SIZE + a] += q_val; } } /* Write mean Q-values (averaged over quantiles) */ float inv_n = 1.0f / (float)IQN_NUM_QUANTILES; - if (lane == 0) { + if (tid == 0) { for (int a = 0; a < TOTAL_BRANCH_ACTIONS; a++) expected_q[sample * TOTAL_BRANCH_ACTIONS + a] = q_acc[a] * inv_n; } @@ -818,7 +867,7 @@ void iqn_forward_kernel( * state → LeakyReLU(W_s1 × state + b_s1) → LeakyReLU(W_s2 × h1 + b_s2) → h_s2 * * Used to compute target_h_s2 from next_states + target DQN weights. - * One warp per sample. Uses shared memory for h1 intermediate. + * One block (256 threads) per sample. Uses shared memory for h1 intermediate. * * Inputs: bf16 states and DQN trunk weights (from external DQN system). * Output: f32 h_s2_out (consumed by IQN forward/backward kernels as f32). @@ -842,7 +891,7 @@ void iqn_trunk_forward_kernel( { int sample = blockIdx.x; if (sample >= batch_size) return; - int lane = threadIdx.x; + int tid = threadIdx.x; extern __shared__ float shmem[]; float* h1 = shmem; /* [shared_h1] */ @@ -854,17 +903,17 @@ void iqn_trunk_forward_kernel( const __nv_bfloat16* my_state = states + sample * padded_sd; /* Layer 1: h1 = leaky_relu(W_s1 @ state + b_s1) — bf16 inputs, f32 compute */ - for (int h = lane; h < shared_h1; h += 32) { + for (int h = tid; h < shared_h1; h += IQN_BLOCK_SIZE) { float acc = (float)b_s1[h]; const __nv_bfloat16* w_row = w_s1 + h * state_dim; for (int d = 0; d < state_dim; d++) acc += (float)w_row[d] * (float)my_state[d]; h1[h] = (acc > 0.0f) ? acc : 0.01f * acc; } - __syncwarp(); + __syncthreads(); /* Layer 2: h2 = leaky_relu(W_s2 @ h1 + b_s2) — bf16 weights, f32 h1 and output */ - for (int h = lane; h < hidden_dim; h += 32) { + for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) { float acc = (float)b_s2[h]; const __nv_bfloat16* w_row = w_s2 + h * shared_h1; for (int d = 0; d < shared_h1; d++)