/**
 * DQN utility kernels — gradient norm reduction and AdamW update.
 *
 * These two kernels form the optimizer step in the DQN training pipeline.
 * They operate on the flat parameter / gradient buffers produced by the
 * cuBLAS backward pass and are compiled together with
 * common_device_functions.cuh (which also provides helper functions
 * needed by the same NVRTC module).
 *
 * No #include of common_device_functions.cuh is needed here — the build
 * function in gpu_dqn_trainer.rs prepends it via string concatenation
 * before passing the combined source to NVRTC.
 *
 * Extracted from dqn_training_kernel.cu (lines 1207-1293).
 */

/* ══════════════════════════════════════════════════════════════════════
 * GRADIENT NORM KERNEL (Phase 3a)
 *
 * Computes gradient L2 norm (sum of squares) via warp + block reduction
 * via warp + block reduction. Must be followed by
 * dqn_adam_update_kernel which reads the completed norm.
 *
 * Launch config: grid=(ceil(TOTAL_PARAMS/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

/* Phase 1: Each block reduces its chunk of grads to a single partial sum.
 * Writes block_sums[blockIdx.x] — NO atomicAdd, NO memset needed.
 * Launch: grid=(num_blocks), block=(256). */
extern "C" __global__ void dqn_grad_norm_kernel(
    const float* __restrict__ grads,       /* [TOTAL_PARAMS] */
    float* __restrict__ block_sums,        /* [num_blocks] partial sums */
    int total_params
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    float g_f = (idx < total_params) ? grads[idx] : 0.0f;
    float g2 = g_f * g_f;

    /* Warp-level reduction */
    for (int offset = 16; offset > 0; offset >>= 1)
        g2 += __shfl_xor_sync(0xFFFFFFFF, g2, offset);

    /* Block-level reduction via shared memory */
    __shared__ float warp_sums[8];
    int warp_id = threadIdx.x / 32;
    int warp_lane = threadIdx.x % 32;
    if (warp_lane == 0) warp_sums[warp_id] = g2;
    __syncthreads();

    if (warp_id == 0) {
        float val = (warp_lane < blockDim.x / 32) ? warp_sums[warp_lane] : 0.0f;
        for (int offset = 16; offset > 0; offset >>= 1)
            val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
        if (warp_lane == 0)
            block_sums[blockIdx.x] = val;
    }
}

/* Phase 2: Reduce block_sums[0..num_blocks] → sqrt → f32 + bf16 output.
 * Single block, 256 threads. Each thread sums a stripe, then warp+block reduce.
 * Launch: grid=(1), block=(256). */
extern "C" __global__ void dqn_grad_norm_finalize(
    const float* __restrict__ block_sums,  /* [num_blocks] from phase 1 */
    float* __restrict__ out_norm,          /* [1] f32 L2 norm (single buffer, no duplicates) */
    int num_blocks
) {
    int tid = threadIdx.x;
    float sum = 0.0f;
    for (int i = tid; i < num_blocks; i += blockDim.x)
        sum += block_sums[i];

    /* Warp reduce */
    for (int offset = 16; offset > 0; offset >>= 1)
        sum += __shfl_xor_sync(0xFFFFFFFF, sum, offset);

    __shared__ float warp_sums[8];
    int warp_id = tid / 32;
    int warp_lane = tid % 32;
    if (warp_lane == 0) warp_sums[warp_id] = sum;
    __syncthreads();

    if (warp_id == 0) {
        float val = (warp_lane < blockDim.x / 32) ? warp_sums[warp_lane] : 0.0f;
        for (int offset = 16; offset > 0; offset >>= 1)
            val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
        if (warp_lane == 0) {
            out_norm[0] = sqrtf(fmaxf(val, 0.0f));
        }
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * ADAM UPDATE KERNEL (Phase 3b)
 *
 * Reads the COMPLETED gradient L2 norm from dqn_grad_norm_kernel,
 * applies gradient clipping, then AdamW parameter update.
 * Trivially parallel: 1 thread per parameter element.
 *
 * SP3 Mech 9 (2026-04-30): post-Adam ISV-driven weight clamp. After the
 * Adam + AdamW + L1-proximal step writes p_val, |p_val| is clamped to
 * `weight_clamp_max_abs` (computed host-side from
 * ISV[Q_ABS_REF_INDEX].max(1.0) × 100 — see `update_adaptive_clip` and
 * the Adam launch sites in `gpu_dqn_trainer.rs`). Targets cuBLAS sgemm
 * f32-accumulator overflow at `iqn_trunk_m` / `bn_d_concat_buf` GEMMs
 * with finite inputs but extreme-magnitude weights (root cause of the
 * F1 step-3540 NaN observed at slots 26+32 with slots 27+35 unflagged).
 * Bound is one OOM below the slot 44-45 diagnostic threshold
 * (1e3 × Q_ABS_REF) so the regression sentinel keeps 10× firing
 * headroom. Disabled when `weight_clamp_max_abs == 0.0` (debug
 * fast-path; production never sets 0).
 *
 * SP4 Task A14 (2026-04-30): Pearl C engagement counter. Each thread
 * sets `local_engage = 1` when its parameter overshoots the bound;
 * block-wide warp+block tree-reduce (mirrors `dqn_grad_norm_kernel`
 * pattern, no atomicAdd per `feedback_no_atomicadd`); single block-leader
 * writes per-block count to
 * `clamp_engage_per_block_buf[engage_buf_offset + blockIdx.x]`. Host
 * sums across blocks → engagement_rate; rate-deficit Wiener-EMA detects
 * post-clamp feedback-loop saturation. `engage_buf_offset < 0` skips
 * the counter write (used by aux trainers outside the SP4 8-group
 * taxonomy — DT, OFI embed, denoise, recursive_conf, sel).
 * `nan_flags_buf[diag_slot] = 1` is a sticky idempotent write capturing
 * "any thread overshot the bound this step" for HEALTH_DIAG.
 *
 * Launch config: grid=(ceil(TOTAL_PARAMS/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_adam_update_kernel(
    float* __restrict__ params,                   /* [TOTAL_PARAMS] f32 master weights */
    const float* __restrict__ grads,              /* [TOTAL_PARAMS] accumulated gradients (f32) */
    float* __restrict__ m,                        /* [TOTAL_PARAMS] first moment (f32) */
    float* __restrict__ v,                        /* [TOTAL_PARAMS] second moment (f32) */
    const float* __restrict__ grad_norm_f32,      /* [1] L2 norm from grad_norm_finalize kernel */
    const float* __restrict__ lr_ptr,             /* [1] learning rate (pinned device-mapped for CUDA Graph) */
    float beta1,
    float beta2,
    float epsilon,
    float weight_decay,
    const float* __restrict__ max_grad_norm_buf, /* [1] adaptive clip norm (device buffer for CUDA Graph) */
    const int* __restrict__ t_ptr,        /* Adam step counter on device (for CUDA Graph) */
    int total_params,
    const float* __restrict__ weight_decay_mask,  /* [TOTAL_PARAMS] 1.0=decay, 0.0=no decay */
    int l1_end,                                    /* param index where L1 stops (end of w_s1) */
    float l1_lambda,                               /* L1 penalty (1e-3). 0=disabled */
    float weight_clamp_max_abs,                    /* SP4 Layer B: ISV-driven |p_val| bound (`ISV[WEIGHT_BOUND[group]].max(EPS_CLAMP_FLOOR)`). 0=disabled. */
    /* SP4 Task A14: Pearl C engagement counter. */
    int* __restrict__ nan_flags_buf,               /* sticky overflow flag store; nan_flags_buf[diag_slot]=1 on first thread that overshoots */
    int diag_slot,                                 /* index into nan_flags_buf (one per Adam kernel) */
    int* __restrict__ clamp_engage_per_block_buf,  /* [SP4_PARAM_GROUP_COUNT × MAX_BLOCKS_PER_ADAM] per-block engagement counts */
    int engage_buf_offset,                         /* per-group offset (group*MAX_BLOCKS); -1 skips writeback (aux trainers) */
    /* SP21 T2.2 Phase 4 (2026-05-11) — per-launch LR multiplier from
     * E4 (`enrichment::compute_branch_lr_scale`). Each branch sub-launch
     * passes its own `BRANCH_LR_SCALE_*_INDEX` ISV value here, scaling
     * the effective Adam LR for that branch's params. Non-branch
     * sub-launches (Trunk, Value, Trunk-extras, IQN, Attn, Curiosity)
     * pass 1.0 (no-op). The kernel multiplies `*lr_ptr * lr_scale`
     * before applying the Adam update; combined with the existing
     * `lr_ptr` (host-controlled global LR schedule), this gives
     * `effective_lr = global_lr × branch_scale`. Cold-start sentinel
     * (ISV reads 0.0 before E4 emits) is handled HOST-side by the
     * launcher (floor at 1.0) — kernel sees only validated scales. */
    float lr_scale
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    /* Out-of-range threads still participate in the block reduce (with
     * local_engage=0) so the block-wide tree-reduce stays correct. */
    bool active = (idx < total_params);

    int local_engage = 0;
    /* SP21 Phase 4: scale effective LR by per-branch multiplier. */
    float lr = active ? (*lr_ptr * lr_scale) : 0.0f;
    int t = active ? *t_ptr : 1;
    float g_f = active ? grads[idx] : 0.0f;

    /* Skip NaN/Inf gradients — leave param, momentum, variance unchanged.
     * Non-deterministic NaN can arise from IQN/CQL auxiliary gradient
     * injection during early training with extreme quantile estimates. */
    bool grad_finite = active && isfinite(g_f);

    if (grad_finite) {
        /* Clip using the COMPLETED L2 norm (single buffer, written by grad_norm_finalize).
         * max_grad_norm read from pinned device-mapped buffer (EMA-based adaptive clip). */
        float max_grad_norm = *max_grad_norm_buf;
        float norm_f = *grad_norm_f32; /* Already L2 norm (sqrt applied by finalize) */
        float clip_scale_f = (norm_f > max_grad_norm) ? (max_grad_norm / norm_f) : 1.0f;
        float clipped_g_f = g_f * clip_scale_f;

        /* Adam bias correction: native float — no bf16 rounding issues. */
        float beta1_t = fmaxf(1.0f - powf(beta1, (float)t), 1e-8f);
        float beta2_t = fmaxf(1.0f - powf(beta2, (float)t), 1e-8f);

        /* Adam moment updates in native float — full precision, no bf16 conversion. */
        float m_i = beta1 * m[idx] + (1.0f - beta1) * clipped_g_f;
        float v_i = beta2 * v[idx] + (1.0f - beta2) * clipped_g_f * clipped_g_f;
        m[idx] = m_i;
        v[idx] = v_i;

        float m_hat = m_i / beta1_t;
        float v_hat = v_i / beta2_t;

        /* AdamW weight decay (decoupled) — masked: trunk+value=1.0, branches=0.0. */
        float p_val = params[idx];
        float wd = weight_decay * weight_decay_mask[idx];
        p_val = p_val - lr * (m_hat / (sqrtf(v_hat) + epsilon) + wd * p_val);

        /* G8: L1 proximal step on w_s1 — automatic feature selection */
        if (idx < l1_end && l1_lambda > 0.0f) {
            float sign_val = (p_val > 0.0f) ? 1.0f : -1.0f;
            p_val = sign_val * fmaxf(fabsf(p_val) - lr * l1_lambda, 0.0f);
        }

        /* SP4 Layer B + SP4 Task A14: post-Adam weight clamp + Pearl C
         * engagement counter. ISV-driven bound
         * (`ISV[WEIGHT_BOUND[group]].max(EPS_CLAMP_FLOOR)`) computed host-side
         * via `read_group_adam_bounds` and passed in. Targets cuBLAS
         * sgemm f32 accumulator overflow at slots 26+32 — Wiener-EMA over
         * p99(|params|) keeps the bound just above legitimate weight
         * magnitudes; the diagnostic retains regression-sentinel
         * headroom. Disabled when bound==0 (debug fast-path; production
         * never sets 0).
         *
         * `local_engage = 1` when |p_val| would overshoot — captured into
         * the per-block tree-reduce below. The sticky `nan_flags_buf` write
         * is idempotent (always writes 1, never 0) so no race exists. */
        if (weight_clamp_max_abs > 0.0f) {
            if (fabsf(p_val) > weight_clamp_max_abs) {
                if (nan_flags_buf != nullptr && diag_slot >= 0) {
                    nan_flags_buf[diag_slot] = 1;
                }
                local_engage = 1;
            }
            p_val = fminf(fmaxf(p_val, -weight_clamp_max_abs), weight_clamp_max_abs);
        }

        params[idx] = p_val;
    }

    /* SP4 Task A14: warp + block tree-reduce of `local_engage`. Mirrors
     * `dqn_grad_norm_kernel`'s pattern — no atomicAdd. blockDim.x=256 ⇒ 8
     * warps; warp_sums fits in __shared__[8]. */
    int e = local_engage;
    for (int offset = 16; offset > 0; offset >>= 1)
        e += __shfl_xor_sync(0xFFFFFFFF, e, offset);

    __shared__ int warp_engage[8];
    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;
        }
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * SAXPY KERNEL
 *
 * y[i] += alpha * x[i]   for i = 0..n-1
 *
 * Used by IQN trunk gradient to apply auxiliary SGD correction to the
 * shared trunk weights. Runs outside the CUDA Graph.
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_saxpy_kernel(
    float* __restrict__ y,
    const float* __restrict__ x,
    float alpha,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) y[i] = y[i] + alpha * x[i];
}

/* ══════════════════════════════════════════════════════════════════════
 * F32 SAXPY KERNEL
 *
 * y[i] += alpha * x[i]   for i = 0..n-1
 *
 * Float variant for f32 d_logits blending (MSE+C51 gradient mix).
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_saxpy_f32_kernel(
    float* __restrict__ y,
    const float* __restrict__ x,
    float alpha,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) y[i] = y[i] + alpha * x[i];
}

/* ══════════════════════════════════════════════════════════════════════
 * DISTILLATION SAXPY KERNEL (D1/N1) — fully GPU-native, zero CPU I/O
 *
 * grad[i] += alpha × (params[i] − best_snap[i])
 *
 * where alpha is computed ON-GPU from the ISV learning_health signal:
 *
 *     alpha = DISTILL_PER_STEP_SCALE × (1 − health)
 *
 * With DISTILL_PER_STEP_SCALE = 0.1 / 250 = 4e-4, this amortizes the
 * spec's intended per-epoch pull (0.1 × (1 − health)) across the typical
 * 250 training steps per epoch. Cumulative effect over an epoch:
 *   1 − (1 − alpha)^250  ≈  alpha × 250  ≈  0.1 × (1 − health)
 * matching the spec to within first-order accuracy. Adam's inverse-rms
 * scaling absorbs any step-count drift, so this doesn't need dynamic
 * adjustment per-run.
 *
 * No CPU-side writes to device memory anywhere. All inputs flow via:
 *   - grad, params, best_snap: existing device buffers
 *   - isv_signals: pinned device-mapped, written by GPU ISV-update kernel
 *
 * The SAXPY is a no-op (zero net effect) when best_snap == params —
 * which is the initial state before any real snapshot is DtoD'd in.
 * The caller (epoch-boundary snapshot capture) DtoD-copies the best
 * ring snapshot into the stable `best_snap` buffer, keeping the pointer
 * baked into the graph while varying its contents between replays.
 *
 * ISV[12] = learning_health ∈ [0, 1]. Clamped defensively in case of
 * NaN/out-of-range writes. ISV pointer may be null during pre-ISV
 * warmup; the null-guard falls back to health=0.5 (no-pull equilibrium).
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

#define DISTILL_PER_STEP_SCALE 4.0e-4f  /* = 0.1 / 250 = spec_alpha / typical_steps_per_epoch */
#define DISTILL_ISV_HEALTH_INDEX 12

extern "C" __global__ void dqn_distill_saxpy_kernel(
    float* __restrict__ grad,
    const float* __restrict__ params,
    const float* __restrict__ best_snap,
    const float* __restrict__ isv_signals,  /* [13] pinned device-mapped; NULL → fallback */
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;

    float health = (isv_signals != nullptr)
        ? fminf(1.0f, fmaxf(0.0f, isv_signals[DISTILL_ISV_HEALTH_INDEX]))
        : 0.5f;
    float alpha = DISTILL_PER_STEP_SCALE * (1.0f - health);

    grad[i] = grad[i] + alpha * (params[i] - best_snap[i]);
}

/* ══════════════════════════════════════════════════════════════════════
 * F32 IN-PLACE SCALE KERNEL
 *
 * y[i] *= alpha   for i = 0..n-1
 *
 * Single-pointer variant — no aliasing possible, no __restrict__ needed.
 * Used for gradient scaling where SAXPY y+=α*x with x==y violates
 * __restrict__ and causes undefined behavior (NaN from register
 * corruption under NVCC's no-alias optimizations).
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_scale_f32_kernel(
    float* y,
    float alpha,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        /* NaN-safe: IEEE 754 says 0*NaN=NaN. When c51_alpha=0 (MSE warmup),
         * C51 grad buffer can contain NaN. Scaling by 0 must yield 0. */
        y[i] = (alpha == 0.0f) ? 0.0f : y[i] * alpha;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * CLIPPED SAXPY KERNEL
 *
 * y[i] += alpha * clip(x_norm, max_norm) * x[i]
 *
 * Reads the pre-computed L2 norm-squared of x from x_norm_sq[0],
 * computes clip_scale = min(1, max_norm / ||x||), and applies:
 *   y[i] += alpha * clip_scale * x[i]
 *
 * This prevents any auxiliary gradient source (IQN, CQL, ensemble)
 * from dominating the combined gradient buffer. Each component is
 * independently clipped BEFORE accumulation.
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_clipped_saxpy_kernel(
    float* __restrict__ y,                         /* [n] f32 gradient accumulator (grad_buf) */
    const float* __restrict__ x,                   /* [n] f32 auxiliary gradient scratch */
    float alpha,
    float max_norm,
    const float* __restrict__ x_norm,      /* [1] L2 norm from grad_norm_finalize */
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;

    float norm_f = *x_norm; /* Already L2 norm (sqrt applied by finalize) */
    float clip_f = (norm_f > max_norm) ? (max_norm / norm_f) : 1.0f;
    y[i] = y[i] + alpha * clip_f * x[i];
}

/* ══════════════════════════════════════════════════════════════════════
 * IN-PLACE GRADIENT CLIPPING KERNEL
 *
 * grads[i] *= min(1, max_norm / ||grads||)
 *
 * Reads the pre-computed L2 norm-squared from grad_norm_sq[0].
 * Used to clip the combined C51+CQL gradient BEFORE IQN/ensemble
 * injection, preventing gradient interference cascade.
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_clip_grad_kernel(
    float* __restrict__ grads,                     /* [TOTAL_PARAMS] f32 gradient accumulator */
    const float* __restrict__ grad_norm_ptr,  /* bf16 L2 norm from finalize kernel */
    float max_norm,
    int total_params
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total_params) return;

    float norm_f = (float)(*grad_norm_ptr);  /* bf16 L2 norm, already sqrt'd */
    if (norm_f > max_norm) {
        float scale = max_norm / norm_f;
        grads[idx] = grads[idx] * scale;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * ZERO KERNEL
 *
 * buf[i] = 0.0f   for i = 0..n-1
 *
 * Zeroes a section of the gradient buffer before IQN trunk backward
 * accumulation (backward_fc_layer uses beta=1.0 for grad accumulation).
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

/* ══════════════════════════════════════════════════════════════════════
 * NaN DETECTION KERNEL
 *
 * Scans a float buffer for NaN/Inf values. Writes 1 to flags[flag_idx]
 * if ANY non-finite value is found. Used for diagnostic instrumentation
 * between CUDA graph replays to pinpoint NaN source.
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_nan_check_f32(
    const float* __restrict__ buf,
    int n,
    int* __restrict__ flags,
    int flag_idx
) {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {
        if (!isfinite(buf[i])) {
            flags[flag_idx] = 1;
            return;
        }
    }
}

/* F32 variant for checking float buffers (cuBLAS outputs, weight shadows) */
extern "C" __global__ void dqn_nan_check_f32_b(
    const float* __restrict__ buf,
    int n,
    int* __restrict__ flags,
    int flag_idx
) {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {
        float v = (float)buf[i];
        if (!isfinite(v)) {
            flags[flag_idx] = 1;
            return;
        }
    }
}

extern "C" __global__ void dqn_zero_kernel(
    float* __restrict__ buf,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) buf[i] = 0.0f;
}

/* ══════════════════════════════════════════════════════════════════════
 * IN-PLACE FINITE CLAMP KERNEL — SP1 Phase C surgical fix infrastructure.
 *
 * For each element of `buf`:
 *   - If !isfinite(v) (NaN or ±Inf): set to 0.0f.
 *   - Else clamp finite |v| to ±max_abs (ISV-driven bound supplied by host).
 *
 * Used as pre-call sanitisation on cuBLAS GEMM backward inputs (slots 27 +
 * 24/25 per `docs/dqn-backward-nan-audit.md`) and as post-call defence-in-
 * depth on outputs (slots 26 + 32) where extreme intermediate-product
 * accumulator overflow has been observed (smoke-test-xvzgk @ commit
 * f139a63ee). max_abs is sourced from existing ISV slots:
 *   - 1e6 × ISV[H_S2_RMS_EMA_INDEX] for the IQN trunk path
 *   - 1e6 × ISV[Q_DIR_ABS_REF_INDEX]  for the Bottleneck Linear path
 * with a 1e3 floor applied host-side when the ISV slot is uninitialised
 * (Invariant 1 ε carve-out — prevents zero-clamping a hot buffer when the
 * ISV producer hasn't run yet).
 *
 * Launch config: grid=(ceil(n/256)), block=(256). In-place — no extra
 * device memory and no DtoD/HtoH copies (per
 * `feedback_no_htod_htoh_only_mapped_pinned`).
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dqn_clamp_finite_f32_kernel(
    float* __restrict__ buf,
    int n,
    float max_abs
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) return;
    float v = buf[idx];
    /* NaN/Inf → 0; finite values clamped to ±max_abs */
    buf[idx] = isfinite(v) ? fmaxf(fminf(v, max_abs), -max_abs) : 0.0f;
}

/* ══════════════════════════════════════════════════════════════════════
 * REGIME-ADAPTIVE PER SCALING KERNEL
 *
 * Scales td_errors by regime similarity before PER priority update.
 * Samples from similar market regimes (ADX/CUSUM) get higher priority.
 *
 * Uses STATE_DIM from common_device_functions.cuh for state stride.
 *
 * Launch config: grid=(ceil(batch_size/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_regime_scale_kernel(
    float* __restrict__ td_errors,          /* [B] scaled in-place */
    const float* __restrict__ states,       /* [B, state_dim] */
    int batch_size,
    int state_dim                           /* runtime state dimension */
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    /* Read target regime from first sample (GPU-only, no CPU readback).
     * ADX at feature index 40, CUSUM at index 41. */
    float target_adx   = states[0 * state_dim + 40];
    float target_cusum = states[0 * state_dim + 41];

    float sample_adx   = states[i * state_dim + 40];
    float sample_cusum = states[i * state_dim + 41];

    float adx_diff   = sample_adx   - target_adx;
    float cusum_diff = sample_cusum - target_cusum;
    float dist_sq = adx_diff * adx_diff + cusum_diff * cusum_diff;

    float temperature = 0.3f;
    float sim = expf(-(dist_sq) / (temperature * temperature + 1e-8f));
    /* Clamp to [0.5, 2.0]: never fully suppress, max 2x upweight */
    float scale = 0.5f + 1.5f * sim;

    td_errors[i] = td_errors[i] * scale;
}

/* ══════════════════════════════════════════════════════════════════════
 * SHRINK-AND-PERTURB KERNEL (GPU-native noise generation)
 *
 * params[i] = alpha * params[i] + (1-alpha) * sigma * noise_i
 *
 * Params in [skip_start, skip_end) are EXCLUDED from perturbation (kept
 * intact). This protects ALL 4 branch head weights (direction, magnitude,
 * order, urgency). Branch heads encode the learned policy — perturbing
 * direction kills Q-gap conviction → forces Flat → magnitude diversity
 * collapses. Only the shared trunk + value head + bottleneck are perturbed.
 *
 * Generates Gaussian noise via Box-Muller transform using per-element
 * LCG PRNG seeded from the element index + step counter. No CPU noise
 * generation or HtoD upload needed.
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_shrink_perturb_kernel(
    float* __restrict__ params,
    float alpha,
    float sigma,
    int n,
    unsigned int seed,           /* changes per call to produce different noise */
    int skip_start,              /* first param index to skip (magnitude branch) */
    int skip_end                 /* one past last param index to skip */
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;

    /* Skip magnitude branch head weights — preserve learned diversity */
    if (i >= skip_start && i < skip_end) return;

    /* Per-element LCG PRNG: hash(seed, index) → uniform random */
    unsigned int state = seed ^ (unsigned int)(i * 2654435761u);
    state = state * 1664525u + 1013904223u;
    float u1 = (float)(state >> 8) / 16777216.0f;  /* (0, 1) */
    state = state * 1664525u + 1013904223u;
    float u2 = (float)(state >> 8) / 16777216.0f;

    /* Box-Muller: uniform → Gaussian N(0, sigma) */
    u1 = fmaxf(u1, 1e-7f);  /* prevent log(0) */
    float noise = sigma * sqrtf(-2.0f * logf(u1)) * cosf(u2 * 6.283185307f);

    /* Shrink-and-Perturb: blend old weights with noise */
    params[i] = alpha * params[i] + (1.0f - alpha) * noise;
}

/* ══════════════════════════════════════════════════════════════════════
 * RELU MASK KERNEL (for IQN trunk gradient — separate from backward module)
 *
 * dx[i] *= (activation[i] > 0.0f)
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_relu_mask_kernel(
    float* __restrict__ dx,
    const float* __restrict__ activation,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    if (activation[i] <= 0.0f) dx[i] = 0.0f;
}

/* ══════════════════════════════════════════════════════════════════════
 * BATCHED SPECTRAL NORM — all weight matrices in a single launch
 *
 * Each block processes one weight matrix independently. 10 blocks = 10
 * matrices, all running in parallel on H100's 132 SMs.
 *
 * Descriptor layout (per matrix, 6 values):
 *   [0] = W pointer (device), [1] = u pointer, [2] = v pointer,
 *   [3] = out_dim, [4] = in_dim, [5] = (unused, alignment pad)
 *
 * Launch: grid=(num_matrices,1,1), block=(256,1,1), shmem=2048
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void spectral_norm_batched(
    const unsigned long long* __restrict__ descriptors,  /* [num_matrices, 6] */
    float sigma_max,
    int num_matrices)
{
    int mat_idx = blockIdx.x;
    if (mat_idx >= num_matrices) return;

    /* Unpack descriptor for this matrix */
    const unsigned long long* desc = &descriptors[mat_idx * 6];
    float* W = (float*)desc[0];
    float* u = (float*)desc[1];
    float* v = (float*)desc[2];
    int out_dim = (int)desc[3];
    int in_dim  = (int)desc[4];

    __shared__ float shmem[256];
    int tid = threadIdx.x;
    int bd = blockDim.x;
    int n_total = out_dim * in_dim;

    /* 3 power iterations for robust sigma estimation.
     * 1 iteration underestimates sigma when Adam makes large weight updates,
     * creating a feedback loop: undetected large sigma → unconstrained backward
     * amplification → gradient spike → large Adam update → repeat.
     * 3 iterations converge to within 0.1% of true sigma for typical weight matrices
     * (cost: 4 extra GEMV per matrix, negligible vs cuBLAS forward). */
    float sigma = 0.0f;
    for (int iter = 0; iter < 3; iter++) {
        /* v_new = W^T u */
        for (int col = tid; col < in_dim; col += bd) {
            float val = 0.0f;
            for (int row = 0; row < out_dim; row++)
                val += W[row * in_dim + col] * u[row];
            v[col] = val;
        }
        __syncthreads();

        /* Normalize v_new */
        float local_v2 = 0.0f;
        for (int col = tid; col < in_dim; col += bd)
            { float vc = v[col]; local_v2 += vc * vc; }
        shmem[tid] = local_v2;
        __syncthreads();
        for (int s = bd / 2; s > 0; s >>= 1) {
            if (tid < s) shmem[tid] += shmem[tid + s];
            __syncthreads();
        }
        float v_norm = sqrtf(shmem[0] + 1e-12f);
        for (int col = tid; col < in_dim; col += bd)
            v[col] /= v_norm;
        __syncthreads();

        /* u_new = W v_new */
        for (int row = tid; row < out_dim; row += bd) {
            float val = 0.0f;
            for (int col = 0; col < in_dim; col++)
                val += W[row * in_dim + col] * v[col];
            u[row] = val;
        }
        __syncthreads();

        /* Sigma = ||u_new|| */
        float local_u2 = 0.0f;
        for (int row = tid; row < out_dim; row += bd)
            { float ur = u[row]; local_u2 += ur * ur; }
        shmem[tid] = local_u2;
        __syncthreads();
        for (int s = bd / 2; s > 0; s >>= 1) {
            if (tid < s) shmem[tid] += shmem[tid + s];
            __syncthreads();
        }
        sigma = sqrtf(shmem[0] + 1e-12f);

        /* Normalize u_new */
        for (int row = tid; row < out_dim; row += bd)
            u[row] /= (sigma + 1e-12f);
        __syncthreads();
    }

    /* Scale W if sigma > sigma_max */
    if (sigma < 1e-6f) return;
    float scale = (sigma > sigma_max) ? (sigma_max / sigma) : 1.0f;
    if (scale < 1.0f) {
        for (int i = tid; i < n_total; i += bd)
            W[i] *= scale;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * PAD STATES KERNEL (CUTLASS K-tile alignment)
 *
 * Scatters contiguous [batch, sd] BF16 states into a padded
 * [batch, padded_sd] layout with zero-filled columns [sd..padded_sd).
 *
 * CUTLASS uses 128-element K-tiles in GemmEx.  When state_dim (K) is
 * not a multiple of 128, the kernel reads past row boundaries.  Padding
 * each row to pad128(state_dim) with zeros eliminates the OOB reads.
 *
 * Launch config: grid=(ceil(batch * padded_sd / 256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void pad_states_kernel(
    float* __restrict__ dst,               /* [batch, padded_sd] — padded f32 output */
    const float* __restrict__ src,         /* [batch, sd]        — contiguous f32 input (PER stores f32) */
    int batch,
    int sd,
    int padded_sd
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= batch * padded_sd) return;
    int b = idx / padded_sd;
    int j = idx % padded_sd;
    dst[idx] = (j < sd) ? src[b * sd + j] : 0.0f;
}

/* ══════════════════════════════════════════════════════════════════════
 * BOTTLENECK TANH + CONCAT KERNEL (#31)
 *
 * Applies tanh activation to bottleneck hidden [B, bn_dim], then
 * concatenates with portfolio features from the state vector to produce
 * the input for h_s1: [B, bn_dim + portfolio_dim].
 *
 * States layout: [0..market_dim) = market features, [market_dim..state_dim) = portfolio
 * Portfolio features are copied AS-IS from the states buffer.
 *
 * One thread per output element in the concatenated buffer.
 * Grid: ceil(B * concat_dim / 256), Block: 256.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void bn_tanh_concat_kernel(
    float* __restrict__ bn_hidden,     /* [B, bn_dim] — IN/OUT: tanh applied in-place */
    const float* __restrict__ states,  /* [B, state_dim_padded] — source for portfolio features */
    float* __restrict__ concat_out,    /* [B, concat_dim] — output: [bn_tanh | portfolio] */
    int batch_size,
    int bn_dim,
    int market_dim,
    int state_dim_padded,
    int concat_dim   /* = bn_dim + (state_dim - market_dim) = bn_dim + portfolio_dim */
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * concat_dim;
    if (idx >= total) return;

    int b = idx / concat_dim;
    int j = idx % concat_dim;

    if (j < bn_dim) {
        /* Bottleneck dimension: apply tanh and copy to output */
        float val = (float)bn_hidden[b * bn_dim + j];
        val = tanhf(val);
        bn_hidden[b * bn_dim + j] = (val);  /* write back tanh'd value */
        concat_out[idx] = (val);
    } else {
        /* Portfolio dimension: copy from states[b, market_dim + (j - bn_dim)] */
        int portfolio_idx = market_dim + (j - bn_dim);
        concat_out[idx] = states[(long long)b * state_dim_padded + portfolio_idx];
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * FEATURE IMPORTANCE KERNEL (#8)
 *
 * Computes per-feature importance as |gradient × activation| summed
 * over the batch. Features with high importance actually influence
 * the Q-value; features with low importance are noise.
 *
 * Reads the first-layer gradient (dW_s1 from grad_buf) and the
 * first-layer input (states from states_buf). Each feature k's
 * importance = sum_b sum_h |dW_s1[h,k] * states[b,k]|.
 *
 * Grid: ceil(market_dim / 256), Block: 256.  One thread per feature.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void feature_importance_kernel(
    float* __restrict__ importance,           /* [market_dim] output */
    const float* __restrict__ grad_w_s1,      /* [shared_h1, state_dim] f32 gradient */
    const float* __restrict__ states, /* [batch_size, state_dim_padded] bf16 */
    int batch_size,
    int shared_h1,
    int state_dim,
    int state_dim_padded,
    int market_dim
) {
    int k = blockIdx.x * blockDim.x + threadIdx.x;
    if (k >= market_dim) return;

    /* Sum |dW[h,k]| over all hidden units — measures gradient magnitude for feature k */
    float grad_sum = 0.0f;
    for (int h = 0; h < shared_h1; h++) {
        grad_sum += fabsf(grad_w_s1[h * state_dim + k]);
    }

    /* Sum |states[b,k]| over batch — measures activation magnitude for feature k */
    float act_sum = 0.0f;
    for (int b = 0; b < batch_size; b++) {
        act_sum += fabsf((float)states[(long long)b * state_dim_padded + k]);
    }

    /* Importance = gradient magnitude × activation magnitude */
    importance[k] = grad_sum * act_sum;
}

/* ══════════════════════════════════════════════════════════════════════
 * LOTTERY TICKET PRUNING KERNEL (#20)
 *
 * Applies a binary weight mask after each Adam optimizer step.
 * mask[i] = 0 → weight is pruned (forced to zero permanently).
 * mask[i] = 1 → weight survives (trained normally).
 *
 * The mask is computed ONCE at a pruning epoch (e.g. epoch 50):
 * sort weights by magnitude, set bottom 70% to 0 in mask.
 * After that, the mask is applied every training step.
 *
 * Grid: ceil(total_params / 256), Block: 256.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void lottery_ticket_mask(
    float* __restrict__ params_f32,           /* [total_params] f32 master weights */
    float* __restrict__ params_bf16,  /* [total_params] bf16 shadow weights */
    const float* __restrict__ mask,           /* [total_params] binary mask (0 or 1) */
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    params_f32[i]  *= mask[i];
    params_bf16[i] = (params_f32[i]);
}

/**
 * Compute the pruning mask: sort by magnitude, bottom fraction → 0.
 * Each thread checks if |weight| < threshold and writes 0 or 1.
 * Threshold is pre-computed on GPU via a separate reduction.
 *
 * Grid: ceil(total_params / 256), Block: 256.
 */
extern "C" __global__ void lottery_ticket_compute_mask(
    float* __restrict__ mask,                 /* [total_params] output: 0 or 1 */
    const float* __restrict__ params_f32,     /* [total_params] current weights */
    float magnitude_threshold,                /* weights below this magnitude are pruned */
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    float mag = fabsf(params_f32[i]);
    mask[i] = (mag >= magnitude_threshold) ? 1.0f : 0.0f;
}

/* ══════════════════════════════════════════════════════════════════════
 * CAUSAL INTERVENTION KERNEL (#34)
 *
 * Copies states to scratch, zeros feature k, all on GPU.
 * Used for causal sensitivity analysis — measures how Q-values change
 * when a specific feature is set to 0 (do-calculus intervention).
 *
 * Grid: ceil(B / 256), Block: 256.  One thread per batch sample.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void causal_intervene_feature(
    const float* __restrict__ states,   /* [B, state_dim_padded] source */
    float* __restrict__ scratch,         /* [B, state_dim_padded] output */
    int feature_k,                               /* which feature to zero */
    int state_dim_padded,                        /* row stride */
    int batch_size
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;

    long long base = (long long)b * state_dim_padded;
    for (int j = 0; j < state_dim_padded; j++) {
        scratch[base + j] = (j == feature_k) ? 0.0f : states[base + j];
    }
}

/**
 * Compute |Q_original - Q_intervened|² summed over batch for a single feature.
 * Reads two f32 logit buffers, reduces to a single scalar per feature.
 *
 * Grid: (1, 1, 1), Block: (256, 1, 1). Called once per (feature_k, pass).
 *
 * Determinism: tid 0 is the only writer, sensitivity_out[feature_k] is
 * unique per kernel launch (feature_k loop variable), and the caller
 * memset_zeros the buffer before the loop. Plain += is sequential within
 * one block + serialized across kernel launches on the same stream — no
 * atomic needed. Result is bit-stable across runs.
 */
extern "C" __global__ void causal_q_delta_reduce(
    const float* __restrict__ q_original,    /* [B, num_atoms] */
    const float* __restrict__ q_intervened,  /* [B, num_atoms] */
    float* __restrict__ sensitivity_out,     /* [market_dim] — accumulator (zeroed by caller) */
    int feature_k,
    int total_elements                        /* B * num_atoms */
) {
    __shared__ float s_sum[256];
    int tid = threadIdx.x;
    float local_sum = 0.0f;

    for (int i = tid; i < total_elements; i += 256) {
        float d = q_original[i] - q_intervened[i];
        local_sum += d * d;
    }

    s_sum[tid] = local_sum;
    __syncthreads();

    for (int s = 128; s > 0; s >>= 1) {
        if (tid < s) s_sum[tid] += s_sum[tid + s];
        __syncthreads();
    }

    if (tid == 0) {
        sensitivity_out[feature_k] += s_sum[0] / (float)total_elements;
    }
}

/**
 * Reduce causal_sensitivity_buf[0..n_features] to a single mean scalar.
 * Writes result to out[0]. Single block, 256 threads.
 */
extern "C" __global__ void causal_mean_reduce(
    const float* __restrict__ sensitivity,
    float* __restrict__ out,
    int n_features)
{
    __shared__ float s_sum[256];
    int tid = threadIdx.x;
    float val = (tid < n_features) ? sensitivity[tid] : 0.0f;
    s_sum[tid] = val;
    __syncthreads();

    for (int s = 128; s > 0; s >>= 1) {
        if (tid < s) s_sum[tid] += s_sum[tid + s];
        __syncthreads();
    }

    if (tid == 0) {
        out[0] = s_sum[0] / (float)n_features;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * GRADIENT REVERSAL KERNEL (#9 — DANN)
 *
 * Negates a section of the gradient buffer to implement gradient reversal
 * for Domain Adversarial Neural Networks. When the discriminator's gradient
 * flows back to the shared trunk, it's NEGATED — forcing the trunk to
 * learn features that the discriminator CANNOT use to predict the regime.
 *
 * grad[i] = -grad[i]   for i in [offset, offset + count)
 *
 * Grid: ceil(count / 256), Block: 256.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void gradient_reversal(
    float* __restrict__ grad_buf,
    int offset,
    int count
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= count) return;
    grad_buf[offset + i] = -grad_buf[offset + i];
}

/* ══════════════════════════════════════════════════════════════════════
 * CROSS-FOLD CONSISTENCY KERNEL (#29)
 *
 * Compares actions between two forward passes (different folds) on the
 * same states. Outputs the fraction of samples where actions agree.
 *
 * agreement = mean(argmax(Q_fold1) == argmax(Q_fold2))
 *
 * Grid: (1, 1, 1), Block: (256, 1, 1).  Single-block reduction.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void cross_fold_consistency(
    const int* __restrict__ actions_fold1,    /* [B] argmax actions from fold 1 */
    const int* __restrict__ actions_fold2,    /* [B] argmax actions from fold 2 */
    float* __restrict__ consistency_out,      /* [1] output: agreement fraction */
    int batch_size
) {
    __shared__ int s_agree[256];
    int tid = threadIdx.x;
    int local_agree = 0;

    for (int i = tid; i < batch_size; i += 256) {
        if (actions_fold1[i] == actions_fold2[i]) local_agree++;
    }

    s_agree[tid] = local_agree;
    __syncthreads();

    for (int s = 128; s > 0; s >>= 1) {
        if (tid < s) s_agree[tid] += s_agree[tid + s];
        __syncthreads();
    }

    if (tid == 0) {
        consistency_out[0] = (float)s_agree[0] / (float)batch_size;
    }
}

/* phantom_liquidity_gbm DELETED (Fix 30 Ambiguous-A, 2026-05-02).
 *
 * The synthetic-data overlay kernel was dead code — zero callers in the
 * production crate, in `services/`, and in `crates/`. The audit grep
 * (`grep -rn "phantom_liquidity_gbm" crates/ml/src/ services/ crates/`)
 * returned only the kernel definition itself; no `load_function`, no
 * `launch_builder`, no Rust-side launcher.
 *
 * Per Fix 29 row #18 (`dqn_utility_kernels.cu:1089-1093`), the writer's
 * downstream consumer contract was unverifiable because no consumer
 * existed. The kernel wrote `log_ret` (synthetic GBM log-return) to
 * `market_features[bar*market_dim + 0]` AND `+3`. Post Bug 1 (commit
 * `5a5dd0fed`) the writer for `market_features[..+0]` is z-normed
 * log-return, not raw log-return — a hypothetical re-wire of this
 * kernel into the production path would have introduced exactly the
 * Bug-1 contract drift the audit was triaging.
 *
 * Resolution per `feedback_no_hiding` + `feedback_wire_everything_up`:
 * delete the kernel rather than leave it as a contract-ambiguous orphan
 * that future re-wirers might trip over. Re-introducing GBM-overlay
 * augmentation is a separate design with its own writer/consumer
 * contract decision (raw vs z-normed log-return; column 0 vs column 3
 * post the `data_loading.rs` Bug-1 normalisation pivot) and should be
 * landed with caller wiring in the same commit per
 * `feedback_wire_everything_up`.
 */

/* ══════════════════════════════════════════════════════════════════════
 * PRICE-LEVEL INVARIANCE KERNEL (#12)
 *
 * Creates a shifted copy of the states batch by adding a random delta
 * to price-derived features [0..3]. The Q-values on original and shifted
 * states should match — if they don't, the model has learned absolute
 * price levels (which never generalize).
 *
 * Grid: ceil(B / 256), Block: 256.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void price_level_shift(
    const float* __restrict__ states,  /* [B, state_dim_padded] */
    float* __restrict__ shifted,       /* [B, state_dim_padded] output */
    unsigned int* __restrict__ rng_state,      /* [1] LCG state */
    int state_dim_padded,
    int batch_size
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;

    /* Generate a random shift from LCG (same shift for all features in one sample) */
    unsigned int rng = rng_state[0] + (unsigned int)b * 2654435761u;
    rng = rng * 1664525u + 1013904223u;
    float u = (float)(rng & 0x00FFFFFFu) / 16777216.0f;
    float delta = (u - 0.5f) * 2.0f; /* shift in [-1, +1] */

    long long base = (long long)b * state_dim_padded;
    for (int k = 0; k < state_dim_padded; k++) {
        float val = (float)states[base + k];
        if (k < 4) {
            val += delta; /* shift price-derived features only */
        }
        shifted[base + k] = (val);
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * SHARPNESS PERTURBATION KERNEL (#28)
 *
 * Perturbs weights by epsilon in a random direction for flatness measurement.
 * For SAM (Sharpness-Aware Minimization) model selection.
 *
 * params_perturbed[i] = params[i] + epsilon * direction[i]
 * where direction is pre-generated random unit vector on GPU.
 *
 * Grid: ceil(n/256), Block: 256.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void sharpness_perturb_weights(
    float* __restrict__ perturbed,           /* [total_params] output */
    const float* __restrict__ original,      /* [total_params] original weights */
    unsigned int* __restrict__ rng_state,    /* [1] LCG state */
    float epsilon,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;

    /* Generate random direction component via LCG */
    unsigned int rng = rng_state[0] + (unsigned int)i * 2654435761u;
    rng = rng * 1664525u + 1013904223u;
    float u1 = (float)(rng & 0x00FFFFFFu) / 16777216.0f + 1e-6f;
    rng = rng * 1664525u + 1013904223u;
    float u2 = (float)(rng & 0x00FFFFFFu) / 16777216.0f;
    float gaussian = sqrtf(-2.0f * logf(u1)) * cosf(6.2831853f * u2);

    perturbed[i] = original[i] + epsilon * gaussian;
}

/* ══════════════════════════════════════════════════════════════════════
 * GRADIENT VACCINE KERNELS (#32)
 *
 * Two kernels that implement the gradient projection for the vaccine:
 *   1. gradient_dot_and_norm: computes dot(g_train, g_val) and |g_val|^2
 *      via warp + block reduction.
 *   2. gradient_project: conditionally projects g_train when dot < 0.
 *      g_train -= (dot / norm_sq) * g_val
 *
 * The dot and norm_sq scalars are written to GPU buffers (2 floats)
 * and read by the projection kernel — no CPU roundtrip.
 * ══════════════════════════════════════════════════════════════════════ */

/* Phase 1: each block writes partial dot + norm to block_results.
 * block_results[blockIdx.x * 2 + 0] = partial dot(g_train, g_val)
 * block_results[blockIdx.x * 2 + 1] = partial |g_val|²
 * NO atomicAdd. Launch: grid=(num_blocks), block=(256). */
extern "C" __global__ void gradient_dot_and_norm(
    const float* __restrict__ g_train,      /* [total_params] f32 */
    const float* __restrict__ g_val,        /* [total_params] f32 */
    float* __restrict__ block_results,      /* [num_blocks * 2] */
    int n
) {
    __shared__ float s_dot[256];
    __shared__ float s_norm[256];
    int tid = threadIdx.x;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    float local_dot = 0.0f, local_norm = 0.0f;

    for (int i = idx; i < n; i += blockDim.x * gridDim.x) {
        float t = g_train[i], v = g_val[i];
        local_dot += t * v;
        local_norm += v * v;
    }

    s_dot[tid] = local_dot;
    s_norm[tid] = local_norm;
    __syncthreads();
    for (int s = 128; s > 0; s >>= 1) {
        if (tid < s) {
            s_dot[tid] += s_dot[tid + s];
            s_norm[tid] += s_norm[tid + s];
        }
        __syncthreads();
    }
    if (tid == 0) {
        block_results[blockIdx.x * 2 + 0] = s_dot[0];
        block_results[blockIdx.x * 2 + 1] = s_norm[0];
    }
}

/* Phase 2: reduce block partials → result[0]=dot, result[1]=norm_sq.
 * Launch: grid=(1), block=(256). */
extern "C" __global__ void gradient_dot_norm_finalize(
    const float* __restrict__ block_results, /* [num_blocks * 2] */
    float* __restrict__ result,              /* [2]: dot, norm_sq */
    int num_blocks
) {
    int tid = threadIdx.x;
    float sum_dot = 0.0f, sum_norm = 0.0f;
    for (int i = tid; i < num_blocks; i += blockDim.x) {
        sum_dot += block_results[i * 2 + 0];
        sum_norm += block_results[i * 2 + 1];
    }
    __shared__ float s_dot[256];
    __shared__ float s_norm[256];
    s_dot[tid] = sum_dot;
    s_norm[tid] = sum_norm;
    __syncthreads();
    for (int s = 128; s > 0; s >>= 1) {
        if (tid < s) {
            s_dot[tid] += s_dot[tid + s];
            s_norm[tid] += s_norm[tid + s];
        }
        __syncthreads();
    }
    if (tid == 0) {
        result[0] = s_dot[0];
        result[1] = s_norm[0];
    }
}

extern "C" __global__ void gradient_project(
    float* __restrict__ g_train,        /* [total_params] f32 — MODIFIED in-place */
    const float* __restrict__ g_val,    /* [total_params] f32 */
    const float* __restrict__ dot_norm, /* [2]: dot_norm[0]=dot, dot_norm[1]=norm_sq */
    int n
) {
    float dot = dot_norm[0];
    float norm_sq = dot_norm[1];

    /* Only project when gradients DISAGREE (dot < 0).
     * When they agree (dot >= 0), both directions benefit generalization. */
    if (dot >= 0.0f) return;

    /* Projection scale: g_train -= (dot / norm_sq) * g_val */
    if (norm_sq < 1e-12f) return; /* g_val is zero — nothing to project */
    float scale = dot / norm_sq;

    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        g_train[idx] -= scale * g_val[idx];
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * BOTTLENECK TANH BACKWARD KERNEL (#31)
 *
 * Computes d_loss/d_bn from d_loss/d_bn_concat (f32) via tanh derivative:
 *   d_bn[b, j] = d_concat[b, j] * (1 - tanh_val[b, j]^2)
 *
 * Only the first bn_dim columns of d_concat are used (portfolio columns skipped).
 * d_concat has stride concat_dim, d_bn has stride bn_dim.
 *
 * tanh_val is read from bn_hidden (bf16, already has tanh applied from forward).
 *
 * Grid: ceil(B * bn_dim / 256), Block: 256.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void bn_tanh_backward_kernel(
    float* __restrict__ d_bn,                     /* [B, bn_dim] f32 output */
    const float* __restrict__ d_concat,            /* [B, concat_dim] f32 input */
    const float* __restrict__ bn_hidden,   /* [B, bn_dim] bf16 (tanh values from forward) */
    int batch_size,
    int bn_dim,
    int concat_dim
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * bn_dim;
    if (idx >= total) return;

    int b = idx / bn_dim;
    int j = idx % bn_dim;

    float tanh_val = (float)bn_hidden[b * bn_dim + j];
    float d_upstream = d_concat[b * concat_dim + j];
    /* tanh'(x) = 1 - tanh(x)^2 */
    d_bn[idx] = d_upstream * (1.0f - tanh_val * tanh_val);
}

/* ══════════════════════════════════════════════════════════════════════
 * VSN-AWARE BOTTLENECK CONCAT BACKWARD: PORTFOLIO PASSTHROUGH (Plan 4 Task 1B-iv)
 *
 * Inverse of the portfolio-passthrough path inside `bn_tanh_concat_kernel`:
 *   forward (concat):   concat_out[b, bn_dim + k] = states[b, market_dim + k]
 *                       for k in [0, portfolio_dim)
 *   backward (here):    d_gated_state[b, market_dim + k] = d_concat[b, bn_dim + k]
 *
 * The market slice of `d_gated_state[b, 0..market_dim]` is filled by the
 * cuBLAS bottleneck Linear backward (`launch_dx_only_lda`, with row stride
 * `state_dim_padded`) AFTER this kernel runs. Padding region
 * `[state_dim..state_dim_padded)` is zero-filled here (no-op when
 * STATE_DIM == STATE_DIM_PADDED, but we honor the contract for callers
 * that pad the state).
 *
 * One thread per (batch, padded-state-index) pair handles either the
 * portfolio slice or the padding tail; the market slice is left untouched
 * (cuBLAS overwrites it with beta=0).
 *
 * Grid: ceil(B*state_dim_padded/256), Block: 256.
 * No atomicAdd, deterministic.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void vsn_d_gated_state_portfolio_pad_kernel(
    float* __restrict__ d_gated_state,    /* [B, state_dim_padded] f32 OUT */
    const float* __restrict__ d_concat,    /* [B, concat_dim] f32 IN */
    int batch_size,
    int market_dim,
    int bn_dim,
    int concat_dim,
    int state_dim,
    int state_dim_padded
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * state_dim_padded;
    if (idx >= total) return;

    int b = idx / state_dim_padded;
    int i = idx % state_dim_padded;

    if (i < market_dim) {
        /* Market slice: zero-initialised here so the buffer is fully defined
         * even before cuBLAS launch_dx_only_lda overwrites it (beta=0 GEMM)
         * with the bottleneck Linear backward dX. Removes a write-order
         * dependency between this kernel and the cuBLAS launch. */
        d_gated_state[idx] = 0.0f;
    } else if (i < state_dim) {
        /* Portfolio passthrough: d_gated_state[b, market_dim + k]
         *                     = d_concat    [b, bn_dim + k]
         * where k = i - market_dim (∈ [0, portfolio_dim)). */
        int k = i - market_dim;
        d_gated_state[idx] = d_concat[(long long)b * concat_dim + bn_dim + k];
    } else {
        /* Padding tail [state_dim, state_dim_padded): zero. */
        d_gated_state[idx] = 0.0f;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * BOTTLENECK WEIGHT GRADIENT BIAS REDUCE (#31)
 *
 * Computes db_bn[j] = sum_b(d_bn[b, j]) for j = 0..bn_dim-1.
 * Each thread handles one output neuron, sums over batch dimension.
 * Deterministic: each thread writes exactly one output — no atomicAdd.
 *
 * Grid: ceil(bn_dim / 256), Block: 256.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void bn_bias_grad_kernel(
    float* __restrict__ grad_bias,       /* [bn_dim] f32 */
    const float* __restrict__ d_bn,      /* [B, bn_dim] f32 */
    int batch_size,
    int bn_dim
) {
    int j = blockIdx.x * blockDim.x + threadIdx.x;
    if (j >= bn_dim) return;

    float sum = 0.0f;
    for (int b = 0; b < batch_size; b++) {
        sum += d_bn[b * bn_dim + j];
    }
    grad_bias[j] += sum;
}

/* ══════════════════════════════════════════════════════════════════════
 * BOTTLENECK TANH + CONCAT F32 KERNEL (#30 + #31)
 *
 * Same as bn_tanh_concat_kernel but with all-f32 I/O.
 * Used by the f32 experience collector forward path.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void bn_tanh_concat_f32_kernel(
    float* __restrict__ bn_hidden,       /* [B, bn_dim] — IN/OUT: tanh applied in-place */
    const float* __restrict__ states,    /* [B, state_dim_padded] — source for portfolio features */
    float* __restrict__ concat_out,      /* [B, concat_dim] — output: [bn_tanh | portfolio] */
    int batch_size,
    int bn_dim,
    int market_dim,
    int state_dim_padded,
    int concat_dim
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * concat_dim;
    if (idx >= total) return;

    int b = idx / concat_dim;
    int j = idx % concat_dim;

    if (j < bn_dim) {
        float val = bn_hidden[b * bn_dim + j];
        val = tanhf(val);
        bn_hidden[b * bn_dim + j] = val;
        concat_out[idx] = val;
    } else {
        int portfolio_idx = market_dim + (j - bn_dim);
        concat_out[idx] = states[(long long)b * state_dim_padded + portfolio_idx];
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * SP15 Wave 4.1a (2026-05-06): BOTTLENECK TANH + CONCAT + DD_PCT KERNEL
 *
 * Bottleneck-aware extension of `bn_tanh_concat_kernel`. Same forward
 * computation (tanh on bn_hidden in-place + portfolio passthrough), with
 * one additional output column appended that broadcasts
 * `isv[DD_PCT_INDEX=406]` across the batch dimension.
 *
 * The standalone `dd_pct_concat_kernel.cu` from Phase 1.5 was bottleneck-
 * incompatible — it operated on raw `[B, state_dim_padded=128]` state but
 * production trunk consumes `[B, s1_input_dim] = [B, bn_dim + portfolio_dim]`
 * (= 102 with current bottleneck_dim=16, market_dim=42, STATE_DIM=128) post-
 * bottleneck. This kernel folds the dd_pct append into the same launch as
 * the bn_tanh + portfolio concat — eval-time policy SEES drawdown context
 * via state on every forward pass, with the trunk consuming the dd_pct
 * column as the (s1_input_dim+1)th feature once Wave 4.1b lands the
 * consumer migration (s1_input_dim 102→103 propagation, GRN w_s1 reshape,
 * 3 forward + 3 backward call-site updates).
 *
 * Kernel contract (mirrors `bn_tanh_concat_kernel` with one extra column):
 *   Inputs:  bn_hidden [B, bn_dim]              — IN/OUT: tanh in-place
 *            states    [B, state_dim_padded]    — source for portfolio
 *            isv[DD_PCT_INDEX=406]              — dd_pct ∈ [0, 1] broadcast
 *   Output:  concat_out [B, concat_dd_dim]
 *            where concat_dd_dim = bn_dim + portfolio_dim + 1
 *            concat_out[b, 0..bn_dim)          = tanh(bn_hidden[b, 0..bn_dim))
 *            concat_out[b, bn_dim..bn_dim+pd)  = states[b, market_dim..market_dim+pd)
 *            concat_out[b, bn_dim+pd]          = isv[DD_PCT_INDEX]
 *
 * Caller must supply `concat_dd_dim = concat_dim + 1` (i.e. the OUTPUT
 * total columns, including the dd_pct append). Pure scatter-copy + one
 * broadcast column; no atomicAdd; deterministic.
 *
 * Wave 4.1a transient state: this kernel + its launcher are currently
 * ORPHAN — production callers continue to use `bn_tanh_concat_kernel` /
 * `bn_tanh_concat_f32_kernel` until Wave 4.1b's atomic refactor flips
 * the 4 launch sites + grows `bn_concat_buf` / `tg_bn_concat_buf` /
 * `on_next_bn_concat_buf` / experience_collector's `bn_concat` by `+B`
 * floats and propagates `s1_input_dim+1` through GRN w_s1 + cuBLAS
 * gemm_caches. Acceptable per the documented Wave 3a/3b split precedent
 * (kernel + launcher + oracle test verify in isolation first).
 *
 * Grid: ceil(B * concat_dd_dim / 256), Block: 256.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void bn_tanh_concat_dd_kernel(
    float* __restrict__ bn_hidden,       /* [B, bn_dim] — IN/OUT: tanh applied in-place */
    const float* __restrict__ states,    /* [B, state_dim_padded] — source for portfolio features */
    const float* __restrict__ isv,       /* [ISV_LEN] — bus, reads slot 406 (DD_PCT_INDEX) */
    float* __restrict__ concat_out,      /* [B, concat_dd_dim] — output: [bn_tanh | portfolio | dd_pct] */
    int batch_size,
    int bn_dim,
    int market_dim,
    int state_dim_padded,
    int concat_dd_dim                    /* = bn_dim + portfolio_dim + 1 */
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * concat_dd_dim;
    if (idx >= total) return;

    int b = idx / concat_dd_dim;
    int j = idx % concat_dd_dim;

    /* Last column = isv[DD_PCT_INDEX] broadcast across batch.
     * DD_PCT_INDEX=406 is a compile-time constant matching the Rust
     * `sp15_isv_slots::DD_PCT_INDEX` and the deleted standalone
     * `dd_pct_concat_kernel`'s contract. */
    if (j == concat_dd_dim - 1) {
        concat_out[idx] = isv[406];
        return;
    }

    if (j < bn_dim) {
        /* Bottleneck dimension: apply tanh and copy to output. */
        float val = bn_hidden[b * bn_dim + j];
        val = tanhf(val);
        bn_hidden[b * bn_dim + j] = val;
        concat_out[idx] = val;
    } else {
        /* Portfolio dimension: copy from states[b, market_dim + (j - bn_dim)]. */
        int portfolio_idx = market_dim + (j - bn_dim);
        concat_out[idx] = states[(long long)b * state_dim_padded + portfolio_idx];
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * STOCHASTIC DEPTH RNG KERNEL (#21)
 *
 * Generates per-layer drop/keep scales ENTIRELY on GPU using LCG RNG.
 * Replaces CPU rand::thread_rng() + memcpy_htod that ran every training step.
 *
 * For each of 3 layers: draw uniform [0,1), compare with drop_prob.
 * If uniform > p: scale = 1/(1-p) (keep, expected-value correction).
 * If uniform <= p: scale = 0 (drop this layer).
 *
 * Single thread — 3 layers is trivially fast, no parallelism needed.
 * Grid: (1, 1, 1), Block: (1, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void stochastic_depth_rng(
    float* __restrict__ scale_buf,         /* [3] output: per-layer scales */
    unsigned int* __restrict__ rng_state,   /* [1] step counter (read-write) */
    float drop_prob                         /* probability of dropping each layer */
) {
    /* Stateless Philox: deterministic per (step, layer). No accumulated LCG state. */
    unsigned int step = rng_state[0];
    float keep_scale = 1.0f / (1.0f - drop_prob);

    for (int layer = 0; layer < 3; layer++) {
        unsigned int key = step;
        unsigned int ctr = (unsigned int)(layer * 37 + 7777);
        ctr ^= key * 0x9E3779B9u;
        ctr *= 0x85ebca6bu;
        ctr ^= ctr >> 13;
        ctr *= 0xc2b2ae35u;
        ctr ^= ctr >> 16;
        float u = (float)(ctr & 0x00FFFFFFu) / 16777216.0f;
        scale_buf[layer] = (u > drop_prob) ? keep_scale : 0.0f;
    }
    rng_state[0] = step + 1; /* monotonic counter, not LCG */
}

/* ══════════════════════════════════════════════════════════════════════
 * STOCHASTIC DEPTH SCALE KERNEL (#21)
 *
 * Scales a bf16 hidden activation buffer by a per-layer scalar.
 * Used for stochastic depth: each hidden layer is either kept (scaled
 * by 1/(1-p) for expected-value correction) or dropped (scaled by 0).
 *
 * The scale buffer [3] is written by stochastic_depth_rng BEFORE each
 * CUDA Graph replay. The graph reads the updated values — this is legal
 * because CUDA Graphs capture buffer addresses, not contents.
 *
 * buf[i] *= scale_buf[layer_idx]   for i = 0..n-1
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void stochastic_depth_scale(
    float* __restrict__ buf,
    const float* __restrict__ scale_buf,
    int layer_idx,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    float scale = scale_buf[layer_idx];
    float val = (float)buf[i] * scale;
    buf[i] = (val);
}

/* ══════════════════════════════════════════════════════════════════════
 * F32 → I32 CAST KERNEL
 *
 * Converts action indices stored as f32 (from GpuTensor) to i32
 * (required by loss/grad kernels). Eliminates the synchronous
 * GPU→CPU→GPU round-trip in upload_batch_gpu that was a per-step
 * serialization bottleneck.
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void f32_to_i32_kernel(
    const float* __restrict__ src,
    int* __restrict__ dst,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    dst[i] = (int)(src[i]);
}

/* ══════════════════════════════════════════════════════════════════════
 * HER IN-PLACE GOAL RELABEL KERNEL
 *
 * Applies HER goal relabeling directly into the trainer's padded
 * staging buffers. For each HER sample i (i = 0..her_batch_size-1):
 *
 *   dst_states[(offset + i) * dst_stride + d] =
 *       src_next_states[donors[i] * src_stride + d]   for d < goal_dim
 *
 * Also sets rewards[offset + i] = 1.0 (goal achieved by construction).
 *
 * This eliminates the intermediate GpuBatch allocation that was
 * previously needed for HER relabeling (17+ cuMemAlloc per step).
 *
 * Launch config: grid=(her_batch_size, 1, 1), block=(32, 1, 1).
 * One warp per relabeled sample — coalesced column writes.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void her_inplace_relabel(
    float* __restrict__ dst_states,       /* [batch_size, dst_stride] padded bf16 */
    float* __restrict__ dst_next_states,  /* [batch_size, dst_stride] padded bf16 */
    float*         __restrict__ dst_rewards,       /* [batch_size] f32 */
    const float*   __restrict__ src_next_states,   /* [batch_size, src_stride] f32 (PER stores f32) */
    const int*     __restrict__ donor_indices,     /* [her_batch_size] i32 */
    int offset,       /* normal_count: first HER row in dst */
    int goal_dim,     /* number of goal columns to replace */
    int src_stride,   /* unpadded state_dim in src (f32 elements) */
    int dst_stride,   /* padded state_dim in dst (bf16 elements) */
    int state_dim     /* min(src_stride, dst_stride) for full copy */
) {
    int i = blockIdx.x;     /* HER sample index */
    int lane = threadIdx.x; /* warp lane 0-31 */
    int donor = donor_indices[i];
    int dst_row = offset + i;

    /* Replace goal columns (first goal_dim elements) with donor's achieved goal.
     * Source is f32 (PER buffer stores full-precision states). Convert to bf16
     * for the padded staging buffer. */
    for (int d = lane; d < goal_dim; d += 32) {
        float achieved = src_next_states[donor * src_stride + d];
        float achieved_bf16 = (achieved);
        dst_states[dst_row * dst_stride + d] = achieved_bf16;
        dst_next_states[dst_row * dst_stride + d] = achieved_bf16;
    }

    /* Set reward = 1.0 (only lane 0) */
    if (lane == 0) {
        dst_rewards[dst_row] = 1.0f;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * BUILD NEXT STATES KERNEL
 *
 * Replaces the CPU loop in build_next_states_f32() that launched ~795K
 * individual memcpy_dtod calls (one per episode per fill timestep).
 * Single kernel launch handles all episodes in parallel.
 *
 * For each element (ep, t, d):
 *   if t < shifted_count: dst[ep,t,d] = src[ep, t+shift, d]
 *   else:                 dst[ep,t,d] = src[ep, timesteps-1, d]  (last state)
 *
 * Launch config: grid=(ceil(total_elems/256), 1, 1), block=(256, 1, 1).
 * total_elems = n_episodes * timesteps * state_dim
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void build_next_states_kernel(
    float* __restrict__ dst,
    const float* __restrict__ src,
    int n_episodes,
    int timesteps,
    int sd,
    int shift
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = n_episodes * timesteps * sd;
    if (idx >= total) return;

    int d  = idx % sd;
    int t  = (idx / sd) % timesteps;
    int ep = idx / (timesteps * sd);

    int ep_offset = ep * timesteps * sd;
    int shifted_count = timesteps - shift;

    if (t < shifted_count) {
        dst[idx] = src[ep_offset + (t + shift) * sd + d];
    } else {
        dst[idx] = src[ep_offset + (timesteps - 1) * sd + d];
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * REGIME CLASSIFICATION KERNEL
 *
 * Classifies each bar into a market regime based on pre-computed
 * ADX (feature index 40) and CUSUM (feature index 41).
 *   ADX > 0.25        → 0 (Trending)
 *   |CUSUM| > 0.7     → 2 (Volatile)
 *   else              → 1 (Ranging)
 *
 * Launch config: grid=(ceil(total_bars/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void regime_classify_kernel(
    const float* __restrict__ features,  /* [total_bars, feat_stride] f32 */
    unsigned char* __restrict__ regimes, /* [total_bars] u8 output */
    int total_bars,
    int feat_stride                      /* 42 */
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= total_bars) return;
    float adx   = features[i * feat_stride + 40];
    float cusum = features[i * feat_stride + 41];
    unsigned char regime;
    if (adx > 0.25f)             regime = 0;  /* Trending */
    else if (fabsf(cusum) > 0.7f) regime = 2;  /* Volatile */
    else                          regime = 1;  /* Ranging */
    regimes[i] = regime;
}

/* ══════════════════════════════════════════════════════════════════════
 * REGIME PREFIX SUM KERNEL
 *
 * Computes per-regime cumulative counts for O(1) range queries.
 * Output: prefix[r * (total_bars+1) + i] = count of regime r in bars [0, i).
 *
 * Simple sequential scan — runs once per training start, not latency-critical.
 * For 1.1M bars this takes <1ms on any GPU.
 *
 * Launch config: grid=(1, 1, 1), block=(3, 1, 1).  One thread per regime.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void regime_prefix_sum_kernel(
    const unsigned char* __restrict__ regimes, /* [total_bars] */
    unsigned int* __restrict__ prefix,         /* [3, total_bars+1] output */
    int total_bars
) {
    int r = threadIdx.x;  /* regime index: 0, 1, or 2 */
    if (r >= 3) return;
    int stride = total_bars + 1;
    prefix[r * stride + 0] = 0;
    unsigned int count = 0;
    for (int i = 0; i < total_bars; i++) {
        if (regimes[i] == r) count++;
        prefix[r * stride + (i + 1)] = count;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * EXPOSURE AUXILIARY GRADIENT KERNEL (v7.1)
 *
 * Breaks the i%3 Q-value degeneracy in branching DQN by giving each
 * exposure output a UNIQUE gradient via cross-entropy against the
 * ══════════════════════════════════════════════════════════════════════ */

/* ══════════════════════════════════════════════════════════════════════
 * POPART REWARD NORMALIZATION KERNEL (v8 D2)
 *
 * Normalizes rewards in-place using running mean/variance (Welford).
 * Updates statistics from the current batch, then normalizes.
 * Eliminates manual v_min/v_max tuning for C51.
 *
 * Launch: grid=(1), block=(256). Single block for atomic stat update.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void popart_normalize_kernel(
    float* __restrict__ rewards,
    float* __restrict__ mean_buf,
    float* __restrict__ var_buf,
    float* __restrict__ count_buf,
    int N, int warmup_count
) {
    __shared__ float s_sum;
    __shared__ float s_sq_sum;
    __shared__ float s_mu;
    __shared__ float s_sigma;

    int tid = threadIdx.x;

    /* Phase 1: compute batch sum and sum-of-squares via reduction */
    float local_sum = 0.0f;
    float local_sq = 0.0f;
    for (int i = tid; i < N; i += blockDim.x) {
        float r = rewards[i];
        local_sum += r;
        local_sq += r * r;
    }

    /* Warp reduction */
    for (int offset = 16; offset > 0; offset >>= 1) {
        local_sum += __shfl_xor_sync(0xFFFFFFFF, local_sum, offset);
        local_sq += __shfl_xor_sync(0xFFFFFFFF, local_sq, offset);
    }

    __shared__ float warp_sums[8];
    __shared__ float warp_sqs[8];
    int warp_id = tid / 32;
    int lane = tid % 32;
    if (lane == 0) { warp_sums[warp_id] = local_sum; warp_sqs[warp_id] = local_sq; }
    __syncthreads();

    if (warp_id == 0) {
        float val = (lane < blockDim.x / 32) ? warp_sums[lane] : 0.0f;
        float sq = (lane < blockDim.x / 32) ? warp_sqs[lane] : 0.0f;
        for (int offset = 16; offset > 0; offset >>= 1) {
            val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
            sq += __shfl_xor_sync(0xFFFFFFFF, sq, offset);
        }
        if (lane == 0) { s_sum = val; s_sq_sum = sq; }
    }
    __syncthreads();

    /* Phase 2: thread 0 updates running statistics */
    if (tid == 0) {
        float old_count = *count_buf;
        float new_count = old_count + (float)N;
        float batch_mean = s_sum / fmaxf((float)N, 1.0f);
        float new_mean = (old_count > 0.0f)
            ? (*mean_buf * old_count + s_sum) / new_count
            : batch_mean;
        float new_var = (new_count > 1.0f)
            ? (s_sq_sum + *var_buf * old_count) / new_count - new_mean * new_mean
            : 1.0f;
        /* Variance floor: 1e-8 (was 0.0001 which capped sigma at 0.01,
         * preventing unit-variance normalization when raw rewards are O(0.001).
         * With rewards of 0.001 and sigma_floor=0.01, normalized rewards were
         * only O(0.1) → Bellman shift of 0.17 atoms → near-zero C51 gradient.
         * 1e-8 allows sigma as low as 1e-4, producing O(10) normalized rewards
         * which shift the C51 distribution by ~17 atoms — full gradient signal. */
        new_var = fmaxf(new_var, 1e-8f);
        *mean_buf = new_mean;
        *var_buf = new_var;
        *count_buf = new_count;
        s_mu = new_mean;
        s_sigma = sqrtf(new_var);
    }
    __syncthreads();

    /* Phase 3: normalize rewards in-place */
    float mu = s_mu;
    float sigma = fmaxf(s_sigma, 1e-4f);
    float total = *count_buf;

    if (total >= (float)warmup_count) {
        for (int i = tid; i < N; i += blockDim.x) {
            rewards[i] = (rewards[i] - mu) / sigma;
        }
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * ROBUST POPART REWARD NORMALIZATION KERNEL
 *
 * Normalizes rewards in-place using pre-computed median and IQR
 * (interquartile range). More robust than Welford mean/var for
 * heavy-tailed bimodal distributions (many ±0.1 micro-rewards +
 * few ±5.0 trade exits).
 *
 * Median/IQR are computed once per epoch from the full sorted
 * experience collection rewards (tiny 3-float DtoH), then passed
 * as scalars to this kernel for every training batch.
 *
 * Launch: grid=ceil(N/256), block=256.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void popart_normalize_robust(
    float* __restrict__ rewards,   /* [N] in-place normalization */
    float median,                  /* pre-computed from sorted rewards */
    float iqr,                     /* interquartile range (Q75 - Q25) */
    int N
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;
    /* Permanent floor on iqr: 1e-4. Was 1e-6 — too low to defend against
     * pathological cold-start iqr (e.g. fold-boundary reset or genuine
     * data corruption producing degenerate quantiles). At iqr=1e-6 a
     * trade-exit reward of 5.0 normalises to 5e6, far outside the C51
     * atom support [-50, 50], producing 10^5× categorical loss readings
     * and rare overflow→NaN paths. 1e-4 caps the worst-case
     * normalised reward at 5e4 / floor=1e-4 = unchanged — but only
     * when actual iqr is truly tiny; healthy iqr (~1e-3 to 1e-1) is
     * unaffected. Per `pearl_blend_formulas_must_have_permanent_floor.md`. */
    float safe_iqr = fmaxf(iqr, 1e-4f);
    rewards[i] = (rewards[i] - median) / safe_iqr;
}

// ============================================================================
// SP2 + SP3 Mech 5: Fused multi-buffer NaN + threshold check
//
// SP2 (Task A4) initial scope: replaces 8 individual `dqn_nan_check_f32`
// launches in run_nan_checks_post_backward (slots 24-30 + 32) with a single
// launch. Each block (blockIdx.x = 0..N_SLOTS) scans its assigned buffer
// end-to-end, sets the corresponding flag if any non-finite value found.
// buf_ptrs[i] = nullptr → block i no-ops (deferred slot).
//
// SP3 Mech 5 (Task B6) extension: grid extended from 12 → 24 blocks. Slots
// 0-11 (= absolute flag indices 24-35) keep NaN-only check (threshold = +inf
// trivially false). Slots 12-25 (= absolute 36-49) add SP4-driven magnitude
// threshold checks against backward-path Adam state, weights, target_q,
// atom-positions, IQN online_params, and post-trunk h_s2 buffers. Each
// threshold reads ISV[SP4_*_BOUND[group/branch]] directly — no magnitude
// multiplier, no per-slot threshold buffer, no per-step HtoD copies.
//
// SP3 close-out v2 extension: grid extended 24 → 25 blocks. Relative slot 24
// (= absolute 48) covers IQN's separate online_params buffer — the F1 step-
// 3540 NaN was traced to this buffer drifting through `iqn_adam_kernel` (its
// own param + Adam state, never reached by the single-kernel Mech 9 in
// `dqn_adam_update_kernel`). SP4 Layer B reads ISV[WEIGHT_BOUND[Iqn=3]] for
// this slot.
//
// SP3 Mech 10 extension: grid extended 25 → 26 blocks. Relative slot 25
// (= absolute 49) covers post-trunk h_s2 activation max-abs — the regression
// sentinel for Mech 10's ISV-driven activation clamp. SP4 Layer B reads
// ISV[H_S2_BOUND_INDEX=169] for this slot (no magnitude multiplier; the
// p99 producer kernel `h_s2_p99_update` already tracks the upper bound that
// Mech 10 clamps to). Mirrors the Mech 5 family pattern.
//
// SP4 Layer B (2026-05-01): per-slot threshold sourced directly from the
// SP4 ISV bound for that slot. No magnitude multipliers; the diagnostic
// fires when |buf[i]| ≥ ISV[bound_idx] (i.e. when a buffer overshoots its
// own clamp threshold). Mech 5's regression-sentinel headroom is now
// supplied by Pearl D's adaptive α tracking the underlying p99 directly.
//
// Slot-index → ISV index mapping (relative slot ≥ 12; absolute = slot + 24):
//   12-15 (abs 36-39): Adam m max-abs ≥ ISV[ADAM_M_BOUND[group_idx]]
//                       group_idx = abs_slot - 36 (groups 0..3)
//   16-19 (abs 40-43): Adam v max-abs ≥ ISV[ADAM_V_BOUND[group_idx]]
//                       group_idx = abs_slot - 40 (groups 0..3)
//   20-21 (abs 44-45): Weight max-abs  ≥ ISV[WEIGHT_BOUND[group_idx]]
//                       group_idx = abs_slot - 44 (groups 0..1 — trunk/value)
//   22    (abs 46)   : target_q post-clip ≥ ISV[TARGET_Q_BOUND]
//   23    (abs 47)   : atom-position span ≥ ISV[ATOM_POS_BOUND[branch=0]]
//                       (any branch is representative since the 4-branch
//                        diagnostic checks the full positions buffer
//                        union; branch 0 = direction)
//   24    (abs 48)   : IQN online_params max-abs ≥ ISV[WEIGHT_BOUND[Iqn=3]]
//   25    (abs 49)   : h_s2 post-clamp max-abs ≥ ISV[H_S2_BOUND]
//
// SP4 ISV slot constants (mirrored from `sp4_isv_slots.rs`):
//   ATOM_POS_BOUND_BASE = 132
//   WEIGHT_BOUND_BASE   = 136
//   ADAM_M_BOUND_BASE   = 144
//   ADAM_V_BOUND_BASE   = 152
//   TARGET_Q_BOUND_INDEX = 131
//   H_S2_BOUND_INDEX     = 169
//
// Invariants preserved from dqn_nan_check_f32:
//   - sticky-flag semantics (writes 1 only; never clears)
//   - block-local reduce (no atomicAdd)
//   - graph-capture safe (pure kernel launch, no host ops)
// ============================================================================

#ifndef SP4_TARGET_Q_BOUND_INDEX
#define SP4_TARGET_Q_BOUND_INDEX 131
#endif
#ifndef SP4_ATOM_POS_BOUND_BASE
#define SP4_ATOM_POS_BOUND_BASE 132
#endif
#ifndef SP4_WEIGHT_BOUND_BASE
#define SP4_WEIGHT_BOUND_BASE 136
#endif
#ifndef SP4_ADAM_M_BOUND_BASE
#define SP4_ADAM_M_BOUND_BASE 144
#endif
#ifndef SP4_ADAM_V_BOUND_BASE
#define SP4_ADAM_V_BOUND_BASE 152
#endif
#ifndef SP4_H_S2_BOUND_INDEX
#define SP4_H_S2_BOUND_INDEX 169
#endif

extern "C" __global__ void dqn_nan_check_fused_f32_kernel(
    const float* const* buf_ptrs,        // [N_SLOTS] (26 entries: 12 NaN-only + 14 threshold)
    const int* buf_lens,                 // [N_SLOTS]
    const float* __restrict__ isv_signals, // SP4 Layer B: per-slot threshold source (mapped-pinned)
    int base_flag_idx,                   // first slot index (24 for backward checks)
    int* nan_flags_buf
) {
    const int slot = blockIdx.x;
    const float* buf = buf_ptrs[slot];
    const int n = buf_lens[slot];
    if (buf == nullptr || n == 0) return;  // deferred slot or unused entry

    // SP4 Layer B: per-slot threshold by relative slot index. Slots 0-11
    // (= absolute 24-35) are NaN-only checks (threshold = +inf so the
    // magnitude check is trivially false). Slots 12-25 (= absolute 36-49)
    // read directly from the SP4 ISV bound for that slot. The bound is
    // floored at 1.0 to handle pre-first-observation (Pearl A sentinel)
    // ISV reads as 0 — same cold-start convention as the Rust consumers.
    float thr = INFINITY;
    if (slot >= 12 && slot < 16) {
        // slots 36-39: Adam m, groups 0..3
        const int group_idx = slot - 12;  // abs_slot - 36
        thr = fmaxf(isv_signals[SP4_ADAM_M_BOUND_BASE + group_idx], 1.0f);
    } else if (slot >= 16 && slot < 20) {
        // slots 40-43: Adam v, groups 0..3
        const int group_idx = slot - 16;  // abs_slot - 40
        thr = fmaxf(isv_signals[SP4_ADAM_V_BOUND_BASE + group_idx], 1.0f);
    } else if (slot >= 20 && slot < 22) {
        // slots 44-45: weight slices, groups 0..1 (trunk / value)
        const int group_idx = slot - 20;  // abs_slot - 44
        thr = fmaxf(isv_signals[SP4_WEIGHT_BOUND_BASE + group_idx], 1.0f);
    } else if (slot == 22) {
        // slot 46: target_q post-clip
        thr = fmaxf(isv_signals[SP4_TARGET_Q_BOUND_INDEX], 1.0f);
    } else if (slot == 23) {
        // slot 47: atom-positions (branch 0 = direction; positions buffer
        // union diagnostic — any branch's bound is representative).
        thr = fmaxf(isv_signals[SP4_ATOM_POS_BOUND_BASE], 1.0f);
    } else if (slot == 24) {
        // slot 48 (SP3 close-out v2): IQN online_params — uses the IQN
        // (group 3) weight bound.
        thr = fmaxf(isv_signals[SP4_WEIGHT_BOUND_BASE + 3], 1.0f);
    } else if (slot == 25) {
        // slot 49 (SP3 Mech 10): post-trunk h_s2 max-abs sentinel.
        thr = fmaxf(isv_signals[SP4_H_S2_BOUND_INDEX], 1.0f);
    }

    int flag_set = 0;
    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        const float v = buf[i];
        if (!isfinite(v)) { flag_set = 1; }
        else if (fabsf(v) >= thr) { flag_set = 1; }
    }
    flag_set = __syncthreads_or(flag_set);
    if (threadIdx.x == 0 && flag_set) {
        nan_flags_buf[base_flag_idx + slot] = 1;
    }
}
