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

perf: IQN kernel — block 32→256, precompute cosines, parallel loss

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-06 23:20:11 +02:00
parent d4b484e30b
commit 6ff45466f9
2 changed files with 243 additions and 164 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -167,6 +167,9 @@ pub struct GpuIqnHead {
online_taus: CudaSlice<f32>,
/// Pre-sampled τ values for target network [B, N]
target_taus: CudaSlice<f32>,
/// Precomputed cosine features [N, embed_dim] — cos(π·(d+1)·τ_i)
/// Fixed at construction (τ are midpoints). Eliminates 16.7M cosf() per step.
cos_features: CudaSlice<f32>,
/// Branch actions decoded from flat actions [B, 3]
branch_actions: CudaSlice<i32>,
/// 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<f32> = (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)

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

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