/**
 * IQN (Implicit Quantile Network) Dual-Head Kernel
 *
 * Mixed-precision: reads BF16 trunk activations/weights, ALL arithmetic in float.
 *
 * Auxiliary quantile head that shares the DQN trunk (h_s2) but trains its own
 * weights independently. Provides quantile-based Q-values for risk-sensitive
 * action selection (CVaR) alongside the primary C51 distributional head.
 *
 * Architecture per sample:
 *   For each τ_i ∈ (0,1), i = 1..N:
 *     cos_feat  = [cos(π·1·τ), ..., cos(π·D·τ)]          [IQN_EMBED_DIM]
 *     embed     = ReLU(W_embed × cos_feat + b_embed)      [hidden_dim]
 *     combined  = h_s2 ⊙ embed                             [hidden_dim]
 *     per-branch: q_d = W_bd × combined + b_bd             [n_d] scalars
 *
 * Loss: Quantile Huber loss (Dabney et al., 2018b)
 *   For each predicted τ_i and target τ'_j:
 *     δ_ij = target_θ_j - predicted_θ_i
 *     ρ_τi(δ) = |τ_i - 𝟙{δ<0}| × L_κ(δ)
 *   L_iqn = (1/N) Σ_i (1/N') Σ_j ρ_τi(δ_ij)
 *
 * Launch config:
 *   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)
 *
 * Runtime parameters (passed to each kernel):
 *   shared_h1   -- first shared hidden layer width (was compile-time IQN_SHARED_H1)
 *   hidden_dim  -- hidden dim = SHARED_H2 from DQN trunk (was compile-time IQN_HIDDEN)
 *
 * Compile-time defines (still used for array sizing and non-hidden constants):
 *   IQN_EMBED_DIM     -- cosine embedding dimension (default 64)
 *   IQN_NUM_QUANTILES -- τ samples per forward pass (default 32)
 *   IQN_KAPPA         -- Huber threshold for quantile loss (default 1.0)
 *   Branch sizes (b0, b1, b2, b3) are runtime kernel parameters.
 */

/* ── Compile-time defaults ──────────────────────────────────────────── */
/* IQN_HIDDEN and IQN_SHARED_H1 are NOW runtime kernel parameters.
 * The #define is kept ONLY for register array sizing (max upper bound).
 * Kernels use the runtime hidden_dim / shared_h1 params for all logic. */
/* IQN_HIDDEN sizes register arrays (IQN_DIST_MAX). Must be >= the maximum
 * runtime hidden_dim from hyperopt (hidden_dim_base upper bound = 512,
 * clamped to 2048 in from_continuous). Using 2048 to cover all cases.
 * Runtime loops use the actual hidden_dim param — no perf impact. */
#ifndef IQN_HIDDEN
#define IQN_HIDDEN 2048
#endif
#ifndef IQN_EMBED_DIM
#define IQN_EMBED_DIM 64
#endif
/* Plan 4 Task 3 (E.3): IQN_NUM_QUANTILES 32 → 5. Fixed-τ multi-quantile
 * heads {0.05, 0.25, 0.50, 0.75, 0.95} replace random U(0,1) sampling.
 * Per-step quantile axis shrinks 32→5 (6.4× smaller); the 5 fixed
 * quantile heads are interpretable / consistent across batches. */
#ifndef IQN_NUM_QUANTILES
#define IQN_NUM_QUANTILES 5
#endif
#ifndef IQN_KAPPA
#define IQN_KAPPA 1.0f
#endif
/* Branch sizes (b0_size, b1_size, b2_size, b3_size) are runtime kernel parameters.
 * No compile-time BRANCH_*_SIZE defines — all offsets computed at runtime. */

/* ── Weight layout (flat f32 buffer) ─────────────────────────────────── */
/* W_embed [hidden_dim, embed_dim], b_embed [hidden_dim]
 * W_b0 [b0_size, hidden_dim], b_b0 [b0_size]
 * W_b1 [b1_size, hidden_dim], b_b1 [b1_size]
 * W_b2 [b2_size, hidden_dim], b_b2 [b2_size]
 * W_b3 [b3_size, hidden_dim], b_b3 [b3_size]
 *
 * Offsets computed at runtime by iqn_compute_offsets(). */

/* Block size for IQN forward/backward/trunk kernels (8 warps = 256 threads). */
#define IQN_BLOCK_SIZE 256

/* 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
#endif
#ifndef IQN_SHARED_H1
#define IQN_SHARED_H1 256
#endif

/* Max shared_h1 for register array sizing */
#define IQN_SHARED_H1_MAX (((IQN_SHARED_H1) + IQN_BLOCK_SIZE - 1) / IQN_BLOCK_SIZE)

/* ── Device helpers ──────────────────────────────────────────────────── */

/** Warp-reduce sum across all 32 lanes (butterfly) — f32. */
__device__ __forceinline__ float iqn_warp_sum(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
    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);
    /* L_κ(δ) = 0.5δ² if |δ|≤κ, κ(|δ| - 0.5κ) otherwise */
    float huber = (abs_delta <= kappa)
        ? 0.5f * delta * delta
        : kappa * (abs_delta - 0.5f * kappa);
    /* Asymmetric weight: |τ - 𝟙{δ<0}| */
    float indicator = (delta < 0.0f) ? 1.0f : 0.0f;
    float weight = fabsf(tau - indicator);
    return weight * huber;
}

/** Gradient of quantile Huber loss element w.r.t. predicted value.
 *  dρ_τ/d(pred) = -|τ - 𝟙{δ<0}| × dL_κ/dδ
 *  where dL_κ/dδ = δ if |δ|≤κ, κ·sign(δ) otherwise */
__device__ __forceinline__ float quantile_huber_grad(float tau, float delta, float kappa) {
    float abs_delta = fabsf(delta);
    float huber_grad = (abs_delta <= kappa)
        ? delta
        : copysignf(kappa, delta);
    float indicator = (delta < 0.0f) ? 1.0f : 0.0f;
    float weight = fabsf(tau - indicator);
    /* Gradient w.r.t. predicted = -weight × huber_grad (negative because δ = target - pred) */
    return -weight * huber_grad;
}

/** Compute runtime weight offsets from hidden_dim, embed_dim, and 4 branch sizes.
 *  Returns 10 offsets:
 *  [0]=W_EMBED, [1]=B_EMBED, [2]=W_B0, [3]=B_B0, [4]=W_B1, [5]=B_B1,
 *  [6]=W_B2, [7]=B_B2, [8]=W_B3, [9]=B_B3 */
__device__ __forceinline__ void iqn_compute_offsets(
    int hidden_dim, int embed_dim,
    int b0, int b1, int b2, int b3,
    int offsets[10])
{
    offsets[0] = 0;                                                      /* W_EMBED */
    offsets[1] = offsets[0] + hidden_dim * embed_dim;                    /* B_EMBED */
    offsets[2] = offsets[1] + hidden_dim;                                /* W_B0 */
    offsets[3] = offsets[2] + b0 * hidden_dim;                           /* B_B0 */
    offsets[4] = offsets[3] + b0;                                        /* W_B1 */
    offsets[5] = offsets[4] + b1 * hidden_dim;                           /* B_B1 */
    offsets[6] = offsets[5] + b1;                                        /* W_B2 */
    offsets[7] = offsets[6] + b2 * hidden_dim;                           /* B_B2 */
    offsets[8] = offsets[7] + b2;                                        /* W_B3 */
    offsets[9] = offsets[8] + b3 * hidden_dim;                           /* B_B3 */
}

/* ── Forward + Loss Kernel ──────────────────────────────────────────── */
/**
 * Per-sample IQN forward pass + quantile Huber loss.
 *
 * Reads h_s2 (shared trunk activation, bf16 from DQN trunk), converts to f32
 * at boundary. All internal computation in f32.
 *
 * τ values are pre-sampled and uploaded via the `taus` buffer (f32).
 *
 * One warp (32 threads) per sample. Quantiles are processed sequentially
 * to conserve registers — each τ produces tba (= b0+b1+b2+b3) values.
 *
 * Saves per-quantile intermediate activations for backward pass.
 */
extern "C" __global__
void iqn_forward_loss_kernel(
    /* Inputs (bf16 from DQN trunk) */
    const float* __restrict__ h_s2,           /* [B, hidden_dim] trunk activations (bf16) */
    const float* __restrict__ target_h_s2,    /* [B, hidden_dim] target trunk activations (bf16) */
    /* Inputs (f32) */
    const float* __restrict__ taus,            /* [B, N] online τ samples ∈ (0,1) */
    const float* __restrict__ target_taus,     /* [B, N] target τ samples */
    const int*   __restrict__ actions,         /* [B, 4] branch actions (direction, magnitude, order, urgency) */
    const float* __restrict__ rewards,         /* [B] */
    const float* __restrict__ dones,           /* [B] (0.0/1.0) */
    float gamma,
    /* 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) */
    /* Activation saves for backward (f32) */
    float* __restrict__ save_embed,            /* [B, N, hidden_dim] post-ReLU embeddings */
    float* __restrict__ save_combined,         /* [B, N, hidden_dim] h_s2 ⊙ embed */
    float* __restrict__ save_q_online,         /* [B, N, tba] quantile Q-values */
    float* __restrict__ save_q_target,         /* [B, N, tba] 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 embed_dim,                             /* runtime: cosine embedding dimension (was IQN_EMBED_DIM) */
    int b0_size,                               /* runtime: branch 0 (direction) action count */
    int b1_size,                               /* runtime: branch 1 (magnitude) action count */
    int b2_size,                               /* runtime: branch 2 (order) action count */
    int b3_size                                /* runtime: branch 3 (urgency) action count */
)
{
    int sample = blockIdx.x;
    if (sample >= batch_size) return;
    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 tba = b0_size + b1_size + b2_size + b3_size;
    int off[10];
    iqn_compute_offsets(hidden_dim, embed_dim, b0_size, b1_size, b2_size, b3_size, off);

    /* ── Pointers into this sample's data ── */
    const float* my_h_s2_bf16 = h_s2 + sample * hidden_dim;
    const float* my_target_h_s2_bf16 = target_h_s2 + sample * hidden_dim;
    const float* my_taus = taus + sample * IQN_NUM_QUANTILES;
    const float* my_target_taus = target_taus + sample * IQN_NUM_QUANTILES;

    int a0 = actions[sample * 4 + 0];  /* direction action */
    int a1 = actions[sample * 4 + 1];  /* magnitude action */
    int a2 = actions[sample * 4 + 2];  /* order action */
    int a3 = actions[sample * 4 + 3];  /* urgency action */
    float reward = rewards[sample];
    float done   = dones[sample];

    /* Weight pointers (f32, same for all samples) */
    const float* w_embed = online_params + off[0];
    const float* b_embed = online_params + off[1];
    const float* w_b0 = online_params + off[2];
    const float* b_b0 = online_params + off[3];
    const float* w_b1 = online_params + off[4];
    const float* b_b1 = online_params + off[5];
    const float* w_b2 = online_params + off[6];
    const float* b_b2 = online_params + off[7];
    const float* w_b3 = online_params + off[8];
    const float* b_b3 = online_params + off[9];

    const float* tw_embed = target_params + off[0];
    const float* tb_embed = target_params + off[1];
    const float* tw_b0 = target_params + off[2];
    const float* tb_b0 = target_params + off[3];
    const float* tw_b1 = target_params + off[4];
    const float* tb_b1 = target_params + off[5];
    const float* tw_b2 = target_params + off[6];
    const float* tb_b2 = target_params + off[7];
    const float* tw_b3 = target_params + off[8];
    const float* tb_b3 = target_params + off[9];

    /* ── Load h_s2 into distributed registers (bf16 → f32 at boundary) ── */
    float h_dist[IQN_DIST_MAX];
    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 = 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++) {
        /* Precomputed cosine features: cos_features[t * embed_dim + d] */
        const float* my_cos = cos_features + t * embed_dim;

        /* 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 = 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++) {
                acc += w_row[d] * my_cos[d];
            }
            /* ReLU */
            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 = 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 = 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 */
        int q_save_offset = (sample * IQN_NUM_QUANTILES + t) * tba;

        /* Branch 0 (direction) */
        float q0_taken = 0.0f;
        for (int a = 0; a < b0_size; a++) {
            float partial = 0.0f;
            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;
            }
        }
        /* Broadcast q0_taken from thread 0 */
        if (tid == 0) shmem_reduce[0] = q0_taken;
        __syncthreads();
        q0_taken = shmem_reduce[0];

        /* Branch 1 (magnitude) */
        float q1_taken = 0.0f;
        for (int a = 0; a < b1_size; a++) {
            float partial = 0.0f;
            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 + b0_size + a] = acc;
                if (a == a1) q1_taken = acc;
            }
        }
        if (tid == 0) shmem_reduce[0] = q1_taken;
        __syncthreads();
        q1_taken = shmem_reduce[0];

        /* Branch 2 (order) */
        float q2_taken = 0.0f;
        for (int a = 0; a < b2_size; a++) {
            float partial = 0.0f;
            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 + b0_size + b1_size + a] = acc;
                if (a == a2) q2_taken = acc;
            }
        }
        if (tid == 0) shmem_reduce[0] = q2_taken;
        __syncthreads();
        q2_taken = shmem_reduce[0];

        /* Branch 3 (urgency) */
        float q3_taken = 0.0f;
        for (int a = 0; a < b3_size; a++) {
            float partial = 0.0f;
            const float* w_row = w_b3 + 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_b3[a] + iqn_block_sum(partial, shmem_reduce);
            if (tid == 0) {
                save_q_online[q_save_offset + b0_size + b1_size + b2_size + a] = acc;
                if (a == a3) q3_taken = acc;
            }
        }
        if (tid == 0) shmem_reduce[0] = q3_taken;
        __syncthreads();
        q3_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 + q3_taken;
    }

    /* ── Process TARGET quantiles ── */
    float target_q_a[IQN_NUM_QUANTILES];

    for (int t = 0; t < IQN_NUM_QUANTILES; t++) {
        /* Target uses same fixed τ midpoints, same precomputed cosine features */
        const float* my_cos = cos_features + t * embed_dim;

        /* Cosine embedding + linear (target weights, precomputed cos_features) */
        float embed_dist[IQN_DIST_MAX];
        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++) {
                acc += w_row[d] * my_cos[d];
            }
            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 = 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, q1, q2, q3;

        /* Branch 0 */
        {
            float partial = 0.0f;
            const float* w_row = tw_b0 + a0 * hidden_dim;
            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 partial = 0.0f;
            const float* w_row = tw_b1 + a1 * hidden_dim;
            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 partial = 0.0f;
            const float* w_row = tw_b2 + a2 * hidden_dim;
            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);
        }
        /* Branch 3 */
        {
            float partial = 0.0f;
            const float* w_row = tw_b3 + a3 * hidden_dim;
            for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
                partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE];
            q3 = tb_b3[a3] + iqn_block_sum(partial, shmem_reduce);
        }

        /* Bellman target: r + γ(1-done) × Q_target */
        float raw_target = q0 + q1 + q2 + q3;
        target_q_a[t] = reward + gamma * (1.0f - done) * raw_target;

        /* Save target Q for backward */
        int q_save_offset = (sample * IQN_NUM_QUANTILES + t) * tba;
        if (tid == 0) {
            save_q_target[q_save_offset + a0] = target_q_a[t];
            save_q_target[q_save_offset + b0_size + a1] = target_q_a[t];
            save_q_target[q_save_offset + b0_size + b1_size + a2] = target_q_a[t];
            save_q_target[q_save_offset + b0_size + b1_size + b2_size + a3] = target_q_a[t];
        }
    }

    /* ── 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_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;
        /* total_loss reduced by separate deterministic iqn_loss_reduce kernel.
         * No atomicAdd — fully deterministic training. */
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * IQN BACKWARD — PER-SAMPLE PHASE (zero atomicAdd)
 *
 * Computes:
 *   1. dL_dq[B, N] — quantile Huber gradient scalar per (sample, quantile)
 *   2. d_h_s2_out[B, hidden_dim] — trunk gradient (register accumulation)
 *
 * Weight gradients are NOT computed here — they are computed by the
 * separate iqn_weight_grad_reduce kernel using the saved dL_dq values
 * and the already-saved forward activations. This eliminates all
 * cross-sample atomicAdd non-determinism.
 *
 * Launch: grid=(batch_size), block=(256), shmem=8*sizeof(float).
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__
void iqn_backward_per_sample(
    const float* __restrict__ h_s2,            /* [B, hidden_dim] f32 trunk activations */
    const float* __restrict__ save_embed,      /* [B, N, hidden_dim] */
    const float* __restrict__ save_q_online,   /* [B, N, tba] */
    const float* __restrict__ save_q_target,   /* [B, N, tba] */
    const float* __restrict__ taus,            /* [B, N] */
    const int*   __restrict__ actions,         /* [B, 4] */
    const float* __restrict__ online_params,   /* IQN weights (for dL_dcomb computation) */
    /* Outputs */
    float* __restrict__ save_dL_dq,            /* [B, N] — saved for weight_grad_reduce */
    float* __restrict__ d_h_s2_out,            /* [B, hidden_dim] — plain write, no atomicAdd */
    int batch_size,
    int shared_h1,
    int hidden_dim,
    int embed_dim,
    int b0_size, int b1_size, int b2_size, int b3_size
)
{
    int sample = blockIdx.x;
    if (sample >= batch_size) return;
    int tid = threadIdx.x;

    __shared__ float shmem_reduce[IQN_BLOCK_SIZE / 32];

    int tba = b0_size + b1_size + b2_size + b3_size;
    int off[10];
    iqn_compute_offsets(hidden_dim, embed_dim, b0_size, b1_size, b2_size, b3_size, off);

    const float* my_h_s2 = h_s2 + sample * hidden_dim;
    const float* my_taus = taus + sample * IQN_NUM_QUANTILES;
    int a0 = actions[sample * 4 + 0];
    int a1 = actions[sample * 4 + 1];
    int a2 = actions[sample * 4 + 2];
    int a3 = actions[sample * 4 + 3];

    const float* w_b0 = online_params + off[2];
    const float* w_b1 = online_params + off[4];
    const float* w_b2 = online_params + off[6];
    const float* w_b3 = online_params + off[8];

    /* Load h_s2 into distributed registers */
    float h_dist[IQN_DIST_MAX];
    for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
        h_dist[h / IQN_BLOCK_SIZE] = my_h_s2[h];

    /* Pre-compute online and target Q-values for taken actions */
    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) * tba;
        if (tid == 0) {
            online_q[t] = save_q_online[qoff + a0]
                        + save_q_online[qoff + b0_size + a1]
                        + save_q_online[qoff + b0_size + b1_size + a2]
                        + save_q_online[qoff + b0_size + b1_size + b2_size + a3];
            target_q[t] = save_q_target[qoff + a0];
        }
    }
    /* Broadcast Q-values from thread 0 to all threads */
    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];
    }

    float kappa = IQN_KAPPA;
    float inv_n_sq = 1.0f / ((float)IQN_NUM_QUANTILES * (float)IQN_NUM_QUANTILES);
    float inv_batch = 1.0f / (float)batch_size;

    /* Register accumulator for d_h_s2 — replaces atomicAdd */
    float d_h_s2_acc[IQN_DIST_MAX];
    for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
        d_h_s2_acc[h / IQN_BLOCK_SIZE] = 0.0f;

    for (int ti = 0; ti < IQN_NUM_QUANTILES; ti++) {
        float tau_i = my_taus[ti];

        /* Compute dL/dq scalar via block reduction */
        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);
        }
        float dL_dq = iqn_block_sum(my_dL_dq, shmem_reduce) * inv_n_sq;

        /* Save dL_dq for deterministic weight gradient computation */
        if (tid == 0)
            save_dL_dq[sample * IQN_NUM_QUANTILES + ti] = dL_dq;

        /* Load saved embed for ReLU mask and trunk gradient */
        int save_offset = (sample * IQN_NUM_QUANTILES + ti) * hidden_dim;
        float embed_dist[IQN_DIST_MAX];
        for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
            embed_dist[h / IQN_BLOCK_SIZE] = save_embed[save_offset + h];

        /* Compute dL/d(combined) for trunk gradient only (no weight writes) */
        for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) {
            float dL_dcomb = (w_b0[a0 * hidden_dim + h]
                            + w_b1[a1 * hidden_dim + h]
                            + w_b2[a2 * hidden_dim + h]
                            + w_b3[a3 * hidden_dim + h]) * dL_dq;

            /* Accumulate trunk gradient: dL/d(h_s2) += dL_dcomb ⊙ embed */
            float d_trunk = inv_batch * dL_dcomb * embed_dist[h / IQN_BLOCK_SIZE];
            if (isfinite(d_trunk))
                d_h_s2_acc[h / IQN_BLOCK_SIZE] += d_trunk;
        }
        /* Weight gradient writes REMOVED — handled by iqn_weight_grad_reduce */
    }

    /* Write d_h_s2_out ONCE per (sample, h) — no atomicAdd */
    for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
        d_h_s2_out[sample * hidden_dim + h] = d_h_s2_acc[h / IQN_BLOCK_SIZE];
}

/* ══════════════════════════════════════════════════════════════════════
 * IQN DETERMINISTIC WEIGHT GRADIENT REDUCTION
 *
 * Computes IQN weight gradients from saved dL_dq[B, N] and forward
 * activations. One thread per weight parameter — loops over all samples
 * and quantiles with fixed summation order. Zero atomicAdd.
 *
 * For embedding weights: recomputes dL_dpre from dL_dq via chain rule
 * (dL_dcomb → dL_dembed → ReLU mask → dL_dpre), then outer-products
 * with cosine features.
 *
 * For branch weights: accumulates dL_dq × combined for matching actions.
 *
 * Launch: grid=(ceil(total_params/256)), block=(256).
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__
void iqn_weight_grad_reduce(
    const float* __restrict__ save_dL_dq,      /* [B, N] quantile gradient scalars */
    const float* __restrict__ h_s2,            /* [B, hidden_dim] trunk activations */
    const float* __restrict__ save_embed,      /* [B, N, hidden_dim] post-ReLU embeddings */
    const float* __restrict__ save_combined,   /* [B, N, hidden_dim] h_s2 ⊙ embed */
    const float* __restrict__ cos_features,    /* [N, embed_dim] precomputed cosines */
    const int*   __restrict__ actions,         /* [B, 4] branch actions */
    const float* __restrict__ online_params,   /* current weights (for dL_dcomb chain rule) */
    float* __restrict__ grad_buf,              /* [total_params] output */
    int batch_size,
    int hidden_dim,
    int embed_dim,
    int b0_size, int b1_size, int b2_size, int b3_size,
    int total_params
)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total_params) return;

    int off[10];
    iqn_compute_offsets(hidden_dim, embed_dim, b0_size, b1_size, b2_size, b3_size, off);
    int total_computed = off[9] + b3_size;
    if (idx >= total_computed) { grad_buf[idx] = 0.0f; return; }

    float inv_batch = 1.0f / (float)batch_size;
    float acc = 0.0f;
    int N = IQN_NUM_QUANTILES;

    if (idx < off[1]) {
        /* ── W_EMBED[h, d] ── */
        int local = idx - off[0];
        int h = local / embed_dim;
        int d = local % embed_dim;

        for (int b = 0; b < batch_size; b++) {
            int a0 = actions[b * 4 + 0];
            int a1 = actions[b * 4 + 1];
            int a2 = actions[b * 4 + 2];
            int a3 = actions[b * 4 + 3];
            float h_s2_bh = h_s2[b * hidden_dim + h];
            float w_sum = online_params[off[2] + a0 * hidden_dim + h]
                        + online_params[off[4] + a1 * hidden_dim + h]
                        + online_params[off[6] + a2 * hidden_dim + h]
                        + online_params[off[8] + a3 * hidden_dim + h];

            for (int ti = 0; ti < N; ti++) {
                float dq = save_dL_dq[b * N + ti];
                float embed_bth = save_embed[(b * N + ti) * hidden_dim + h];
                /* Chain rule: dL_dcomb × h_s2 × ReLU_mask = dL_dpre */
                float dL_dpre = dq * w_sum * h_s2_bh * ((embed_bth > 0.0f) ? 1.0f : 0.0f);
                if (isfinite(dL_dpre))
                    acc += dL_dpre * cos_features[ti * embed_dim + d];
            }
        }
    } else if (idx < off[2]) {
        /* ── B_EMBED[h] ── */
        int h = idx - off[1];

        for (int b = 0; b < batch_size; b++) {
            int a0 = actions[b * 4 + 0];
            int a1 = actions[b * 4 + 1];
            int a2 = actions[b * 4 + 2];
            int a3 = actions[b * 4 + 3];
            float h_s2_bh = h_s2[b * hidden_dim + h];
            float w_sum = online_params[off[2] + a0 * hidden_dim + h]
                        + online_params[off[4] + a1 * hidden_dim + h]
                        + online_params[off[6] + a2 * hidden_dim + h]
                        + online_params[off[8] + a3 * hidden_dim + h];

            for (int ti = 0; ti < N; ti++) {
                float dq = save_dL_dq[b * N + ti];
                float embed_bth = save_embed[(b * N + ti) * hidden_dim + h];
                float dL_dpre = dq * w_sum * h_s2_bh * ((embed_bth > 0.0f) ? 1.0f : 0.0f);
                if (isfinite(dL_dpre))
                    acc += dL_dpre;
            }
        }
    } else {
        /* ── Branch weight or bias (W_Bi or B_Bi) ── */
        int bsizes[4] = {b0_size, b1_size, b2_size, b3_size};
        int branch_idx = -1;
        int is_bias = 0;
        int local_offset = 0;

        for (int i = 0; i < 4; i++) {
            int w_start = off[2 + i * 2];
            int b_start = off[3 + i * 2];
            int w_end = b_start;
            int b_end = (i < 3) ? off[4 + i * 2] : total_computed;

            if (idx >= w_start && idx < w_end) {
                branch_idx = i;
                is_bias = 0;
                local_offset = idx - w_start;
                break;
            }
            if (idx >= b_start && idx < b_end) {
                branch_idx = i;
                is_bias = 1;
                local_offset = idx - b_start;
                break;
            }
        }

        if (branch_idx < 0) { grad_buf[idx] = 0.0f; return; }

        if (is_bias) {
            /* B_Bi[a] = inv_batch × Σ_{b: action_b == a} Σ_ti dL_dq[b, ti] */
            int a = local_offset;
            for (int b = 0; b < batch_size; b++) {
                if (actions[b * 4 + branch_idx] != a) continue;
                for (int ti = 0; ti < N; ti++)
                    acc += save_dL_dq[b * N + ti];
            }
        } else {
            /* W_Bi[a, h] = inv_batch × Σ_{b: action_b == a} Σ_ti dL_dq[b, ti] × combined[b, ti, h] */
            int a = local_offset / hidden_dim;
            int h = local_offset % hidden_dim;
            for (int b = 0; b < batch_size; b++) {
                if (actions[b * 4 + branch_idx] != a) continue;
                for (int ti = 0; ti < N; ti++) {
                    float dq = save_dL_dq[b * N + ti];
                    float comb = save_combined[(b * N + ti) * hidden_dim + h];
                    acc += dq * comb;
                }
            }
        }
    }

    grad_buf[idx] = acc * inv_batch;
}

/* ══════════════════════════════════════════════════════════════════════
 * IQN DETERMINISTIC LOSS REDUCTION
 *
 * Sums per_sample_loss[0..batch_size-1] / batch_size into total_loss[0].
 * Single-thread sequential — fully deterministic (fixed summation order).
 *
 * Launch: grid=(1), block=(1).
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void iqn_loss_reduce(
    const float* __restrict__ per_sample_loss,
    float* __restrict__ total_loss,
    int batch_size
) {
    float sum = 0.0f;
    for (int i = 0; i < batch_size; i++)
        sum += per_sample_loss[i];
    total_loss[0] = sum / (float)batch_size;
}

/* ══════════════════════════════════════════════════════════════════════
 * IQN DETERMINISTIC GRADIENT NORM — TWO-PHASE REDUCTION
 *
 * Phase 1: each block computes partial sum-of-squares → block_sums[blockIdx.x]
 *   NO atomicAdd. One output per block.
 *
 * Phase 2: single block sums all partial sums → norm_out[0]
 *   Sequential (grid=1, block=1) for full determinism.
 *
 * The Adam kernel reads norm_out[0] as raw sum-of-squares (takes sqrt internally).
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__
void iqn_grad_norm_phase1(
    const float* __restrict__ grads,
    float* __restrict__ block_sums,    /* [num_blocks] — one partial sum per block */
    int total_params
)
{
    int tid = blockIdx.x * blockDim.x + threadIdx.x;

    float partial = 0.0f;
    for (int i = tid; i < total_params; i += gridDim.x * blockDim.x) {
        float g = grads[i];
        partial += g * g;
    }

    /* Block-level reduction via shared memory */
    __shared__ float shared[256];
    shared[threadIdx.x] = partial;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s)
            shared[threadIdx.x] += shared[threadIdx.x + s];
        __syncthreads();
    }

    /* One output per block — no atomicAdd */
    if (threadIdx.x == 0)
        block_sums[blockIdx.x] = shared[0];
}

extern "C" __global__
void iqn_grad_norm_phase2(
    const float* __restrict__ block_sums,
    float* __restrict__ norm_out,      /* [1] sum-of-squares output */
    int num_blocks
)
{
    /* Sequential sum — deterministic. Grid=(1), Block=(1). */
    float sum = 0.0f;
    for (int i = 0; i < num_blocks; i++)
        sum += block_sums[i];
    norm_out[0] = sum;
}

/* ── Adam Update Kernel ─────────────────────────────────────────────── */
/* SP4 Layer B: trailing `weight_clamp_max_abs` arg matches the
 * pattern wired into all five Adam kernels (dqn/iqn/iql/attn/curiosity).
 * Bound = `ISV[WEIGHT_BOUND[ParamGroup::Iqn]].max(EPS_CLAMP_FLOOR)` computed
 * host-side via `read_group_adam_bounds`. Targets the IQN-internal weight
 * buffer that the slot-26 cuBLAS GEMM reads — `apply_iqn_trunk_gradient`
 * saw NaN at slot 26 with finite inputs because `online_params` here
 * drifted via this Adam, never reached by the single-kernel Mech 9 in
 * `dqn_adam_update_kernel`. Disabled when bound==0.
 *
 * SP4 Task A14 (2026-04-30): Pearl C engagement counter via warp+block
 * tree-reduce mirroring `dqn_grad_norm_kernel` (no atomicAdd). Single
 * block-leader writes `clamp_engage_per_block_buf[engage_buf_offset +
 * blockIdx.x]`. `engage_buf_offset < 0` skips the writeback (defensive
 * default for callers that disable Pearl C).
 */
extern "C" __global__
void iqn_adam_kernel(
    float* __restrict__ params,
    float* __restrict__ grads,
    float* __restrict__ m,          /* first moment */
    float* __restrict__ v,          /* second moment */
    const float* __restrict__ norm_sq, /* [1] gradient sum-of-squares (NOT sqrt'd) */
    float lr, float beta1, float beta2, float eps,
    float weight_decay, float max_grad_norm,
    const int* __restrict__ adam_t_buf, /* [1] step counter on device */
    int total_params,
    int shared_h1,                 /* runtime: unused, for consistent interface */
    int hidden_dim,                /* runtime: unused, for consistent interface */
    int embed_dim,                 /* runtime: unused, for consistent interface */
    int b0_size,                   /* runtime: unused, for consistent interface */
    int b1_size,                   /* runtime: unused, for consistent interface */
    int b2_size,                   /* runtime: unused, for consistent interface */
    int b3_size,                   /* runtime: unused, for consistent interface */
    float weight_clamp_max_abs,    /* SP4 Layer B: ISV-driven |p| bound (`ISV[WEIGHT_BOUND[Iqn]].max(EPS_CLAMP_FLOOR)`). 0=disabled. */
    /* SP4 Task A14: Pearl C engagement counter. */
    int* __restrict__ nan_flags_buf,              /* sticky overflow flag */
    int diag_slot,                                /* nan_flags_buf index for this Adam kernel */
    int* __restrict__ clamp_engage_per_block_buf, /* [SP4_PARAM_GROUP_COUNT × MAX_BLOCKS_PER_ADAM] */
    int engage_buf_offset                          /* group*MAX_BLOCKS; -1 to skip */
)
{
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    bool active = (tid < total_params);
    int local_engage = 0;

    if (active) {
        float g = grads[tid];

        /* Skip NaN/Inf gradients — leave param, momentum, variance unchanged.
         * Still fall through to the block reduce so the engagement counter
         * sees this thread's local_engage=0. */
        if (!isfinite(g)) {
            grads[tid] = 0.0f;
        } else {
            /* Gradient clipping: norm_sq[0] is raw sum-of-squares from grad_norm kernel */
            float grad_norm_val = sqrtf(fmaxf(norm_sq[0], 0.0f));
            float clip_scale = (grad_norm_val > max_grad_norm && grad_norm_val > 0.0f)
                ? max_grad_norm / grad_norm_val : 1.0f;
            g *= clip_scale;

            /* AdamW update */
            float mi = beta1 * m[tid] + (1.0f - beta1) * g;
            float vi = beta2 * v[tid] + (1.0f - beta2) * g * g;
            m[tid] = mi;
            v[tid] = vi;

            /* Bias correction */
            int adam_t = adam_t_buf[0];
            float m_hat = mi / (1.0f - powf(beta1, (float)adam_t));
            float v_hat = vi / (1.0f - powf(beta2, (float)adam_t));

            /* Parameter update with decoupled weight decay */
            float p = params[tid];
            p = p - lr * (m_hat / (sqrtf(v_hat) + eps) + weight_decay * p);

            /* SP4 Layer B + SP4 Task A14: post-Adam |p| clamp +
             * engagement-counter increment. Bound is host-supplied
             * (`ISV[WEIGHT_BOUND[Iqn]].max(EPS_CLAMP_FLOOR)`). The
             * ISV producer (Wiener-EMA over p99(|params|)) keeps the
             * bound naturally above legitimate weight magnitudes;
             * `EPS_CLAMP_FLOOR=1.0` is a cold-start floor only. */
            if (weight_clamp_max_abs > 0.0f) {
                if (fabsf(p) > weight_clamp_max_abs) {
                    if (nan_flags_buf != nullptr && diag_slot >= 0) {
                        nan_flags_buf[diag_slot] = 1;
                    }
                    local_engage = 1;
                }
                p = fminf(fmaxf(p, -weight_clamp_max_abs), weight_clamp_max_abs);
            }

            params[tid] = p;

            /* Zero gradient for next step */
            grads[tid] = 0.0f;
        }
    }

    /* SP4 Task A14: warp + block tree-reduce of `local_engage`. */
    int e = local_engage;
    for (int offset = 16; offset > 0; offset >>= 1)
        e += __shfl_xor_sync(0xFFFFFFFF, e, offset);

    __shared__ int warp_engage[8]; /* blockDim.x = 256 ⇒ 8 warps */
    int warp_id   = threadIdx.x / 32;
    int warp_lane = threadIdx.x % 32;
    if (warp_lane == 0) warp_engage[warp_id] = e;
    __syncthreads();

    if (warp_id == 0) {
        int val = (warp_lane < blockDim.x / 32) ? warp_engage[warp_lane] : 0;
        for (int offset = 16; offset > 0; offset >>= 1)
            val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
        if (warp_lane == 0 && engage_buf_offset >= 0 && clamp_engage_per_block_buf != nullptr) {
            clamp_engage_per_block_buf[engage_buf_offset + blockIdx.x] = val;
        }
    }
}

/* ── Forward-Only Kernel (Inference) ────────────────────────────────── */
/**
 * IQN forward pass for inference — no loss, no activation saves.
 *
 * Reads bf16 h_s2 from DQN trunk, converts to f32 at boundary.
 * All params and taus are f32. Output is f32.
 *
 * Plan 4 Task 3 (E.3): with `IQN_NUM_QUANTILES = 5` and FIXED_TAUS
 * = {0.05, 0.25, 0.50, 0.75, 0.95} (sorted ascending, populated from the
 * Rust constructor's HtoD upload), the median quantile lives at index 2
 * (τ = 0.50). Action ranking now uses the median Q rather than the mean
 * over all quantiles — the median is the risk-neutral central estimate
 * that the C51 mode/argmax-of-expected aggregation already provides at
 * its head, and using it here keeps the IQN inference path consistent
 * with the spec's 5-quantile decomposition.
 *
 * Outputs: expected_q [B, tba] — median (τ=0.50) Q per action.
 */
extern "C" __global__
void iqn_forward_kernel(
    const float* __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, tba] (f32) */
    int batch_size,
    int shared_h1,                               /* runtime: unused here, for consistency */
    int hidden_dim,                              /* runtime: IQN hidden dim = SHARED_H2 */
    int embed_dim,                               /* runtime: cosine embedding dimension (was IQN_EMBED_DIM) */
    int b0_size,                                 /* runtime: branch 0 (direction) action count */
    int b1_size,                                 /* runtime: branch 1 (magnitude) action count */
    int b2_size,                                 /* runtime: branch 2 (order) action count */
    int b3_size                                  /* runtime: branch 3 (urgency) action count */
)
{
    int sample = blockIdx.x;
    if (sample >= batch_size) return;
    int tid = threadIdx.x;

    /* Shared memory for block-level reductions */
    __shared__ float shmem_reduce[IQN_BLOCK_SIZE / 32];

    /* ── Compute runtime weight offsets ── */
    int tba = b0_size + b1_size + b2_size + b3_size;
    int off[10];
    iqn_compute_offsets(hidden_dim, embed_dim, b0_size, b1_size, b2_size, b3_size, off);

    const float* my_h_s2_bf16 = h_s2 + sample * hidden_dim;
    const float* my_taus = taus + sample * IQN_NUM_QUANTILES;
    const float* w_embed = params + off[0];
    const float* b_embed = params + off[1];
    const float* w_b0 = params + off[2];
    const float* b_b0 = params + off[3];
    const float* w_b1 = params + off[4];
    const float* b_b1 = params + off[5];
    const float* w_b2 = params + off[6];
    const float* b_b2 = params + off[7];
    const float* w_b3 = params + off[8];
    const float* b_b3 = params + off[9];

    /* Load h_s2 (bf16 → f32 at boundary) */
    float h_dist[IQN_DIST_MAX];
    for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
        h_dist[h / IQN_BLOCK_SIZE] = (float)my_h_s2_bf16[h];

    /* Plan 4 Task 3 (E.3): hold the MEDIAN (τ=0.50, index 2) Q per action.
     * The fixed-τ broadcast guarantees τ_2 = 0.50 across all samples; the
     * mean over the asymmetric set {0.05, 0.25, 0.50, 0.75, 0.95} would
     * mix tail-risk into the action-ranking estimate. Median is the risk-
     * neutral central tendency consistent with the spec's 5-quantile
     * decomposition (other 4 quantiles surface as ISV diagnostics).
     * Max tba = 4+3+3+3 = 13 — fits easily in registers. */
    float q_acc[16]; /* sized >= max tba */
    for (int a = 0; a < tba; a++)
        q_acc[a] = 0.0f;

    for (int t = 0; t < IQN_NUM_QUANTILES; t++) {
        /* Precomputed cosine features for this quantile */
        const float* my_cos = cos_features + t * embed_dim;

        /* Cosine embedding + linear + ReLU (using precomputed cos_features) */
        float embed_dist[IQN_DIST_MAX];
        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++) {
                acc += w_row[d] * my_cos[d];
            }
            embed_dist[h / IQN_BLOCK_SIZE] = fmaxf(acc, 0.0f);
        }

        /* Element-wise product */
        float comb_dist[IQN_DIST_MAX];
        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];

        /* Plan 4 Task 3 (E.3): only the median quantile (τ=0.50, index 2)
         * feeds expected_q. Off-median quantiles still execute their GEMMs
         * (the trunk is fixed; per-quantile cost is the W_b·combined dot
         * product) so that the IQN's per-step Q surface is fully sampled
         * for downstream compute_iqr / CVaR consumers — but we discard the
         * non-median values here. The 4-fold reduction in cost vs the
         * legacy 32-quantile mean comes from the kernel-wide
         * IQN_NUM_QUANTILES drop (32 → 5), not from skipping non-medians.
         * IQN_MEDIAN_INDEX = 2 is anchored to the FIXED_TAUS table layout
         * `[0.05, 0.25, 0.50, 0.75, 0.95]` populated by the constructor's
         * HtoD; if that ordering ever changes the seed-fingerprint catches
         * the mismatch at checkpoint load before this kernel executes. */
        const int IQN_MEDIAN_INDEX = 2;

        /* Branch 0 */
        for (int a = 0; a < b0_size; a++) {
            float partial = 0.0f;
            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 q_val = iqn_block_sum(partial, shmem_reduce) + b_b0[a];
            if (t == IQN_MEDIAN_INDEX) q_acc[a] = q_val;
        }
        /* Branch 1 */
        for (int a = 0; a < b1_size; a++) {
            float partial = 0.0f;
            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 q_val = iqn_block_sum(partial, shmem_reduce) + b_b1[a];
            if (t == IQN_MEDIAN_INDEX) q_acc[b0_size + a] = q_val;
        }
        /* Branch 2 */
        for (int a = 0; a < b2_size; a++) {
            float partial = 0.0f;
            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 q_val = iqn_block_sum(partial, shmem_reduce) + b_b2[a];
            if (t == IQN_MEDIAN_INDEX) q_acc[b0_size + b1_size + a] = q_val;
        }
        /* Branch 3 */
        for (int a = 0; a < b3_size; a++) {
            float partial = 0.0f;
            const float* w_row = w_b3 + 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 q_val = iqn_block_sum(partial, shmem_reduce) + b_b3[a];
            if (t == IQN_MEDIAN_INDEX) q_acc[b0_size + b1_size + b2_size + a] = q_val;
        }
    }

    /* Plan 4 Task 3 (E.3): write median Q (τ=0.50). No mean-over-quantiles
     * division — q_acc already holds the median quantile values per action. */
    if (tid == 0) {
        for (int a = 0; a < tba; a++)
            expected_q[sample * tba + a] = q_acc[a];
    }
}

/* ── Trunk Forward Kernel ──────────────────────────────────────────── */
/* ── iqn_trunk_forward_kernel deleted (Plan 4 Task 2c.3b) ────────────
 * The kernel was orphaned — `IqnHead::trunk_forward_kernel` loaded it
 * but never invoked it. The actual IQN target trunk runs through
 * cuBLAS `iqn_lt_matmul` calls in
 * `gpu_iqn_head.rs::execute_training_pipeline`; that path is now
 * replaced with `target_encoder_forward_only` so online and target use
 * one trunk implementation (the GRN forward). See audit row #11 in
 * `docs/h_s2_consumers_audit.md` and the audit-corrections commit
 * `4402dbaf3` for the orphan classification. */

/* ── EMA Target Update Kernel ──────────────────────────────────────── */
/**
 * Polyak EMA: target[i] = (1 - tau) * target[i] + tau * online[i]
 * All f32.
 */
extern "C" __global__
void iqn_ema_kernel(
    float* __restrict__ target,
    const float* __restrict__ online,
    const float* __restrict__ tau_buf,
    int n,
    int shared_h1,
    int hidden_dim,
    int embed_dim,
    int b0_size,
    int b1_size,
    int b2_size,
    int b3_size
)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float tau = tau_buf[0];
        target[i] = (1.0f - tau) * target[i] + tau * online[i];
    }
}

/* ── Flat action → branch action decode kernel ─────────────────────────
 *
 * Decodes flat action index ∈ [0, b0×b1×b2×b3) into 4 branch indices.
 * action = dir*b1*b2*b3 + mag*b2*b3 + order*b3 + urgency
 * Grid: (ceil(batch_size/256), 1, 1), Block: (256, 1, 1)
 *
 * Input:  flat_actions[B]   -- i32 flat action indices
 * Output: branch_actions[B*4] -- i32 [direction, magnitude, order, urgency] per sample
 */
extern "C" __global__
void iqn_decode_actions_kernel(
    const int* __restrict__ flat_actions,
    int* __restrict__ branch_actions,
    int batch_size,
    int shared_h1,                 /* runtime: unused, for consistent interface */
    int hidden_dim,                /* runtime: unused, for consistent interface */
    int embed_dim,                 /* runtime: unused, for consistent interface */
    int b0_size,                   /* runtime: branch 0 (direction) action count */
    int b1_size,                   /* runtime: branch 1 (magnitude) action count */
    int b2_size,                   /* runtime: branch 2 (order) action count */
    int b3_size                    /* runtime: branch 3 (urgency) action count */
)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    int flat = flat_actions[i];
    /* Bounds check: clamp to valid action range [0, total_actions) */
    int max_action = b0_size * b1_size * b2_size * b3_size;
    if (flat < 0 || flat >= max_action) flat = 0;
    int s1 = b1_size * b2_size * b3_size;
    int s2 = b2_size * b3_size;
    int direction = flat / s1;
    int remainder = flat % s1;
    int magnitude = remainder / s2;
    remainder     = remainder % s2;
    int order     = remainder / b3_size;
    int urgency   = remainder % b3_size;

    branch_actions[i * 4 + 0] = direction;
    branch_actions[i * 4 + 1] = magnitude;
    branch_actions[i * 4 + 2] = order;
    branch_actions[i * 4 + 3] = urgency;
}

/* Plan 4 Task 3 (E.3): iqn_sample_taus_kernel deleted.
 *
 * The legacy Philox-driven random τ sampler is obsolete under the fixed-τ
 * multi-quantile design. The Rust constructor uploads FIXED_TAUS =
 * {0.05, 0.25, 0.50, 0.75, 0.95} into `online_taus` / `target_taus` once
 * via clone_htod (broadcast B times to fill the [B, 5] buffer); both the
 * online and target IQN forwards and the CVaR cold path read this static
 * device buffer directly. No per-step kernel launch is needed.
 *
 * Pre-edit consumers: only `compute_cvar_scales` invoked this kernel; that
 * call site is removed in lockstep. No other Rust-side or kernel-side
 * consumer existed (verified by grep on `iqn_sample_taus_kernel`).
 */

/* ================================================================== */
/* Element-wise kernels for cuBLAS-based IQN forward/backward         */
/* ================================================================== */

/* ReLU forward: out = max(x, 0) */
extern "C" __global__
void iqn_relu_fwd(
    const float* __restrict__ x,
    float* __restrict__ out,
    int N
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        out[idx] = fmaxf(x[idx], 0.0f);
    }
}

/* ReLU backward: d_x = (x > 0) ? d_out : 0 */
extern "C" __global__
void iqn_relu_bwd(
    const float* __restrict__ x,       /* saved pre-activation */
    const float* __restrict__ d_out,   /* upstream gradient */
    float* __restrict__ d_x,           /* output gradient */
    int N
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        d_x[idx] = (x[idx] > 0.0f) ? d_out[idx] : 0.0f;
    }
}

/* Hadamard product with sigmoid: out = h_s2_tiled * sigmoid(embed) */
extern "C" __global__
void iqn_hadamard_sigmoid(
    const float* __restrict__ h_s2_tiled,  /* [H, B*Q] */
    const float* __restrict__ embed,       /* [H, B*Q] post-ReLU embedding */
    float* __restrict__ out,               /* [H, B*Q] combined */
    int N                                  /* H * B * Q */
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        float sig = 1.0f / (1.0f + expf(-embed[idx]));
        out[idx] = h_s2_tiled[idx] * sig;
    }
}

/* Hadamard-sigmoid backward:
   d_embed = d_combined * h_s2_tiled * sigmoid'(embed_pre)
   d_h_s2_tiled = d_combined * sigmoid(embed_post)
   embed_post is post-ReLU, so sigmoid acts on the ReLU output. */
extern "C" __global__
void iqn_hadamard_sigmoid_bwd(
    const float* __restrict__ h_s2_tiled,  /* [H, B*Q] */
    const float* __restrict__ embed,       /* [H, B*Q] post-ReLU embedding */
    const float* __restrict__ d_combined,  /* [H, B*Q] upstream gradient */
    float* __restrict__ d_embed,           /* [H, B*Q] gradient w.r.t. embed */
    float* __restrict__ d_h_s2_tiled,      /* [H, B*Q] gradient w.r.t. h_s2 */
    int N
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        float sig = 1.0f / (1.0f + expf(-embed[idx]));
        float dsig = sig * (1.0f - sig);
        d_embed[idx] = d_combined[idx] * h_s2_tiled[idx] * dsig;
        d_h_s2_tiled[idx] = d_combined[idx] * sig;
    }
}

/* Bias add: x[idx] += bias[idx % out_dim] */
extern "C" __global__
void iqn_bias_add(
    float* __restrict__ x,
    const float* __restrict__ bias,
    int out_dim,
    int N
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        x[idx] += bias[idx % out_dim];
    }
}

/* Bias gradient reduce — Phase 1: partial block sums.
   Grid: (ceil(BQ/256), out_dim, 1), Block: (256, 1, 1), Shared: 256*sizeof(float).
   d_pre is [out_dim, BQ] column-major.
   partials is [num_blocks_x, out_dim] — one partial sum per (block, neuron). */
extern "C" __global__
void iqn_bias_grad_reduce_p1(
    const float* __restrict__ d_pre,
    float* __restrict__ partials,
    int out_dim,
    int BQ
) {
    extern __shared__ float sdata[];
    int j   = blockIdx.y;              /* which output neuron */
    int tid = threadIdx.x;
    int gid = blockIdx.x * blockDim.x + tid;

    float val = (gid < BQ) ? d_pre[j + (long)gid * out_dim] : 0.0f;
    sdata[tid] = val;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) sdata[tid] += sdata[tid + s];
        __syncthreads();
    }

    if (tid == 0) {
        partials[blockIdx.x * out_dim + j] = sdata[0];
    }
}

/* Bias gradient reduce — Phase 2: final reduction + scale.
   Grid: (out_dim, 1, 1), Block: (256, 1, 1).
   partials is [num_blocks, out_dim]. */
extern "C" __global__
void iqn_bias_grad_reduce_p2(
    const float* __restrict__ partials,
    float* __restrict__ d_bias,
    int num_blocks,
    int out_dim,
    float inv_bq
) {
    int j   = blockIdx.x;
    int tid = threadIdx.x;

    float sum = 0.0f;
    for (int i = tid; i < num_blocks; i += blockDim.x) {
        sum += partials[i * out_dim + j];
    }

    /* Warp-level reduction */
    for (int offset = 16; offset > 0; offset >>= 1) {
        sum += __shfl_down_sync(0xFFFFFFFF, sum, offset);
    }

    if (tid == 0) {
        d_bias[j] = sum * inv_bq;
    }
}

/* Tile h_s2 [H, B] into h_s2_tiled [H, B*Q] by repeating each column Q times.
   Column-major: h_s2[h + b*H] -> h_s2_tiled[h + (b*Q + q)*H] */
extern "C" __global__
void iqn_h_s2_tile(
    const float* __restrict__ h_s2,        /* [H, B] */
    float* __restrict__ h_s2_tiled,        /* [H, B*Q] */
    int H, int B, int Q
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = H * B * Q;
    if (idx < total) {
        int h = idx % H;
        int bq = idx / H;
        int b = bq / Q;
        h_s2_tiled[idx] = h_s2[h + b * H];
    }
}

/* Reduce d_h_s2_tiled [H, B*Q] back to d_h_s2 [H, B] by summing Q quantiles.
   d_h_s2[h + b*H] = sum_{q=0}^{Q-1} d_h_s2_tiled[h + (b*Q+q)*H] */
extern "C" __global__
void iqn_d_h_s2_reduce(
    const float* __restrict__ d_h_s2_tiled,  /* [H, B*Q] */
    float* __restrict__ d_h_s2,              /* [H, B] */
    int H, int B, int Q
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = H * B;
    if (idx < total) {
        int h = idx % H;
        int b = idx / H;
        float sum = 0.0f;
        for (int q = 0; q < Q; q++) {
            sum += d_h_s2_tiled[h + (b * Q + q) * H];
        }
        d_h_s2[idx] = sum;
    }
}

/* Quantile Huber loss + per-branch dq gradients.
   One block per sample. Threads cooperate over quantiles/actions.
   Layout: q_online/q_target are [TBA, B*Q] col-major.
   actions: [B, 4] row-major (branch actions per sample).

   For each sample s and quantile q, computes:
     For each branch b, the taken action a_b:
       delta = q_online[offset_b + a_b, s*Q + q] - gamma * q_target[offset_b + a_b, s*Q + q]
       Quantile Huber loss with tau = taus[s*Q + q]
       d_q_online = (1/Q) * d_huber * |tau - I(delta<0)|

   per_sample_loss[s] = (1/Q) * sum over q of sum over branches of huber_loss */
extern "C" __global__
void iqn_quantile_huber_loss(
    const float* __restrict__ q_online,    /* [TBA, B*Q] */
    const float* __restrict__ q_target,    /* [TBA, B*Q] */
    const float* __restrict__ taus,        /* [B*Q] (fixed midpoints) */
    const int* __restrict__ actions,       /* [B, 4] */
    float* __restrict__ per_sample_loss,   /* [B] */
    float* __restrict__ d_q_online,        /* [TBA, B*Q] zero-init required */
    int B, int Q, int b0, int b1, int b2, int b3,
    float kappa, float gamma
) {
    int s = blockIdx.x;
    if (s >= B) return;
    int tid = threadIdx.x;
    int TBA = b0 + b1 + b2 + b3;

    /* Decode branch actions for this sample */
    int a0 = actions[s * 4 + 0];
    int a1 = actions[s * 4 + 1];
    int a2 = actions[s * 4 + 2];
    int a3 = actions[s * 4 + 3];

    /* Branch offsets in the TBA dimension */
    int off0 = 0, off1 = b0, off2 = b0 + b1, off3 = b0 + b1 + b2;
    int branch_actions[4] = {a0, a1, a2, a3};
    int branch_offsets[4] = {off0, off1, off2, off3};

    float sample_loss = 0.0f;

    /* Each thread processes a subset of quantiles */
    for (int q = tid; q < Q; q += blockDim.x) {
        int col = s * Q + q;
        float tau_q = taus[col];
        float q_loss = 0.0f;

        for (int br = 0; br < 4; br++) {
            int row = branch_offsets[br] + branch_actions[br];
            int idx = row + col * TBA;  /* col-major index */

            float online_val = q_online[idx];
            float target_val = q_target[idx];
            float delta = online_val - gamma * target_val;

            /* Huber loss */
            float abs_delta = fabsf(delta);
            float huber;
            if (abs_delta <= kappa) {
                huber = 0.5f * delta * delta;
            } else {
                huber = kappa * (abs_delta - 0.5f * kappa);
            }

            /* Quantile weight */
            float indicator = (delta < 0.0f) ? 1.0f : 0.0f;
            float qw = fabsf(tau_q - indicator);
            q_loss += qw * huber;

            /* Gradient of quantile Huber w.r.t. q_online */
            float d_huber;
            if (abs_delta <= kappa) {
                d_huber = delta;
            } else {
                d_huber = (delta > 0.0f) ? kappa : -kappa;
            }
            d_q_online[idx] = qw * d_huber / (float)Q;
        }
        sample_loss += q_loss;
    }

    /* Warp reduce then block reduce for sample_loss */
    sample_loss /= (float)Q;

    /* Warp-level reduction */
    for (int offset = 16; offset > 0; offset >>= 1) {
        sample_loss += __shfl_down_sync(0xFFFFFFFF, sample_loss, offset);
    }

    /* Block-level reduction via shared memory */
    __shared__ float warp_sums[8]; /* max 256 threads = 8 warps */
    int warp_id = tid / 32;
    int lane_id = tid % 32;
    if (lane_id == 0) {
        warp_sums[warp_id] = sample_loss;
    }
    __syncthreads();

    if (tid == 0) {
        float total = 0.0f;
        int num_warps = (blockDim.x + 31) / 32;
        for (int w = 0; w < num_warps; w++) {
            total += warp_sums[w];
        }
        per_sample_loss[s] = total;
    }
}

/* Tile cos_features [D, Q] into [D, B*Q] by repeating B times.
   cos_features_src[d + q*D] -> cos_features_tiled[d + (b*Q+q)*D] */
extern "C" __global__
void iqn_cos_tile(
    const float* __restrict__ cos_src,     /* [D, Q] */
    float* __restrict__ cos_tiled,         /* [D, B*Q] */
    int D, int B, int Q
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = D * B * Q;
    if (idx < total) {
        int d = idx % D;
        int bq = idx / D;
        int q = bq % Q;
        cos_tiled[idx] = cos_src[d + q * D];
    }
}
