#![allow(unsafe_code)] // Required for CUDA kernel launches

//! GPU-accelerated IQN (Implicit Quantile Network) dual-head trainer.
//!
//! 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 head.
//!
//! ## Architecture
//!
//! For each τ ∈ (0,1):
//!   cos_feat  = [cos(π·1·τ), ..., cos(π·D·τ)]
//!   embed     = ReLU(W_embed × cos_feat + b_embed)
//!   combined  = h_s2 ⊙ embed
//!   q_branch  = W_bd × combined + b_bd    (per branch)
//!
//! ## cuBLAS GEMM pipeline
//!
//! Forward and backward matmuls use batched cublasLtMatmul (TF32 FAST_TF32).
//! Element-wise operations (ReLU, sigmoid, Huber loss) use precompiled cubin
//! kernels from `iqn_dual_head_kernel.cu`. Zero atomicAdd.
//!
//! ## Integration
//!
//! Called after the main DQN training step in `FusedTrainingCtx::run_full_step()`.
//! Reads `save_h_s2` (online trunk activation) from `GpuDqnTrainer` and computes
//! target h_s2 from `next_states_buf` using the target DQN trunk weights.

use std::sync::Arc;

use cudarc::cublaslt::result as cublaslt_result;
use cudarc::cublaslt::sys as cublaslt_sys;
use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtrMut, LaunchConfig, PushKernelArg};
use tracing::info;

use crate::MLError;
use super::gpu_weights::DuelingWeightSet;
use super::shared_cublas_handle::PerStreamCublasHandles;

// ---------------------------------------------------------------------------
// Configuration
// ---------------------------------------------------------------------------

/// Plan 4 Task 3 (E.3): fixed-τ multi-quantile heads.
///
/// Replaces the legacy QR-DQN midpoint sampling τ_i = (2i+1)/(2N) with five
/// well-known risk-aware quantile points sorted ascending: 5th / 25th /
/// 50th (median) / 75th / 95th. The kernel's `IQN_NUM_QUANTILES` macro is
/// fixed to 5 in lockstep so this Rust-side table fully describes the
/// per-step τ axis. The constructor's `clone_htod` broadcasts these 5
/// values across all batch rows (`[B, 5]`) once at init; both online and
/// target IQN forwards plus the CVaR cold path read this static buffer
/// directly — no per-step kernel sampling required.
///
/// Median (τ=0.50) lives at index 2; off-median positions feed the four
/// new ISV diagnostic slots (ISV[99..103) Q_p05/p25/p75/p95 EMAs). The
/// kernel's inference path (`iqn_forward_kernel`) latches the median Q
/// per action via `if (t == IQN_MEDIAN_INDEX) q_acc[a] = q_val;`, so any
/// re-ordering of this table requires updating that constant (= 2) too.
pub const FIXED_TAUS: [f32; 5] = [0.05, 0.25, 0.50, 0.75, 0.95];
/// Index of the median quantile (τ=0.50) in `FIXED_TAUS`. Mirrors
/// `IQN_MEDIAN_INDEX` in `iqn_forward_kernel`. Used by IQR and any other
/// consumer that needs the median position; the spec § E.3 ordering pins
/// this at 2 (sorted ascending {0.05, 0.25, 0.50, 0.75, 0.95}).
pub const FIXED_TAUS_MEDIAN_INDEX: usize = 2;

/// Configuration for the GPU IQN dual-head trainer.
#[derive(Debug, Clone)]
pub struct GpuIqnConfig {
    /// Hidden dimension = SHARED_H2 from DQN trunk (default 256).
    pub hidden_dim: usize,
    /// Cosine embedding dimension (default 64).
    pub embed_dim: usize,
    /// Number of quantile samples per forward pass.
    /// Plan 4 Task 3 (E.3): fixed at 5 — see `FIXED_TAUS`. Must match the
    /// kernel-side `IQN_NUM_QUANTILES` define; mismatch is a build error
    /// because the kernel sizes register arrays from the macro.
    pub num_quantiles: usize,
    /// Huber threshold for quantile loss (default 1.0).
    pub kappa: f32,
    /// Input state dimension (tensor-core-aligned, e.g. 48 or 56).
    pub state_dim: usize,
    /// First shared hidden layer width from DQN trunk (default 256).
    pub shared_h1: usize,
    /// Branch 0 (direction) action count.
    pub branch_0_size: usize,
    /// Branch 1 (magnitude) action count.
    pub branch_1_size: usize,
    /// Branch 2 (order) action count.
    pub branch_2_size: usize,
    /// Branch 3 (urgency) action count.
    pub branch_3_size: usize,
    /// Training batch size (must match DQN trainer).
    pub batch_size: usize,
    /// Learning rate for IQN Adam optimizer.
    pub lr: f32,
    /// Adam beta1 (first moment decay).
    pub beta1: f32,
    /// Adam beta2 (second moment decay).
    pub beta2: f32,
    /// Adam epsilon (numerical stability).
    pub epsilon: f32,
    // SP4 Layer B: `weight_decay` field removed; per-step AdamW rate is
    // sourced from ISV[WD_RATE[Iqn=3]] and threaded into
    // `execute_training_pipeline` as `weight_decay_value`.
    /// Maximum gradient L2 norm for clipping.
    pub max_grad_norm: f32,
    /// Discount factor for Bellman projection.
    pub gamma: f32,
}

impl Default for GpuIqnConfig {
    fn default() -> Self {
        Self {
            hidden_dim: 256,
            embed_dim: 64,
            // Plan 4 Task 3 (E.3): fixed-τ {0.05, 0.25, 0.50, 0.75, 0.95}.
            // Was 32 (random U(0,1) sampling, QR-DQN-equivalent midpoints).
            num_quantiles: FIXED_TAUS.len(),
            kappa: 1.0,
            state_dim: 72,
            shared_h1: 256,
            branch_0_size: 4,
            branch_1_size: 3,
            branch_2_size: 3,
            branch_3_size: 3,
            batch_size: 256,
            lr: 3e-4,
            beta1: 0.9,
            beta2: 0.999,
            epsilon: 1e-8,
            // SP4 Layer B: weight_decay field removed; sourced from ISV at runtime.
            max_grad_norm: 1.0,
            gamma: 0.99,
        }
    }
}

impl GpuIqnConfig {
    /// Total branch actions (sum of all 4 branches).
    fn total_branch_actions(&self) -> usize {
        self.branch_0_size + self.branch_1_size + self.branch_2_size + self.branch_3_size
    }

    /// Total IQN parameters (embedding + per-branch output layers).
    fn total_params(&self) -> usize {
        let h = self.hidden_dim;
        let d = self.embed_dim;
        let b0 = self.branch_0_size;
        let b1 = self.branch_1_size;
        let b2 = self.branch_2_size;
        let b3 = self.branch_3_size;
        // W_embed[H, D] + b_embed[H] + W_b0..b3[Bi, H] + b_b0..b3[Bi]
        h * d + h + (b0 + b1 + b2 + b3) * (h + 1)
    }
}

// ---------------------------------------------------------------------------
// IqnGemmDesc — pre-cached cublasLt descriptors for a single GEMM shape
// ---------------------------------------------------------------------------

/// Pre-created cublasLt descriptors + heuristic algo for a fixed IQN GEMM shape.
///
/// Same pattern as `CachedGemmDesc` in `batched_forward.rs` / `batched_backward.rs`.
/// Created once at init time, reused every training step — zero per-call overhead.
struct IqnGemmDesc {
    matmul_desc: cublaslt_sys::cublasLtMatmulDesc_t,
    a_layout: cublaslt_sys::cublasLtMatrixLayout_t,
    b_layout: cublaslt_sys::cublasLtMatrixLayout_t,
    c_layout: cublaslt_sys::cublasLtMatrixLayout_t,
    d_layout: cublaslt_sys::cublasLtMatrixLayout_t,
    algo: cublaslt_sys::cublasLtMatmulAlgo_t,
}

unsafe impl Send for IqnGemmDesc {}
unsafe impl Sync for IqnGemmDesc {}

impl Drop for IqnGemmDesc {
    fn drop(&mut self) {
        unsafe {
            let _ = cublaslt_result::destroy_matrix_layout(self.d_layout);
            let _ = cublaslt_result::destroy_matrix_layout(self.c_layout);
            let _ = cublaslt_result::destroy_matrix_layout(self.b_layout);
            let _ = cublaslt_result::destroy_matrix_layout(self.a_layout);
            let _ = cublaslt_result::destroy_matmul_desc(self.matmul_desc);
        }
    }
}

// ---------------------------------------------------------------------------
// GpuIqnHead
// ---------------------------------------------------------------------------

/// GPU-accelerated IQN dual-head trainer.
///
/// Owns online + target IQN weight buffers, Adam state, compiled CUDA kernels,
/// cuBLAS GEMM descriptors, and all intermediate activation/loss buffers.
/// Operates entirely on GPU with zero atomicAdd.
#[allow(missing_debug_implementations, non_snake_case)]
pub struct GpuIqnHead {
    config: GpuIqnConfig,
    stream: Arc<CudaStream>,
    shared_handle: Arc<PerStreamCublasHandles>,

    // ── Compiled element-wise kernels ──────────────────────────────
    // Plan 4 Task 2c.3b: trunk_forward_kernel deleted (was orphaned —
    // loaded but never invoked; cuBLAS path is the actual implementation).
    loss_reduce_kernel: CudaFunction,
    grad_norm_phase1_kernel: CudaFunction,
    grad_norm_phase2_kernel: CudaFunction,
    adam_kernel: CudaFunction,
    forward_kernel: CudaFunction,
    ema_kernel: CudaFunction,
    decode_actions_kernel: CudaFunction,
    // Plan 4 Task 3 (E.3): `sample_taus_kernel` deleted. Fixed-τ broadcast
    // at construction replaces per-step Philox sampling. The kernel itself
    // is removed from `iqn_dual_head_kernel.cu` in lockstep — no other
    // consumer existed.

    // ── cuBLAS element-wise cubin kernels (forward + backward) ────
    relu_fwd_kernel: CudaFunction,
    relu_bwd_kernel: CudaFunction,
    hadamard_sigmoid_kernel: CudaFunction,
    hadamard_sigmoid_bwd_kernel: CudaFunction,
    quantile_huber_loss_kernel: CudaFunction,
    bias_add_kernel: CudaFunction,
    bias_grad_reduce_p1_kernel: CudaFunction,
    bias_grad_reduce_p2_kernel: CudaFunction,
    h_s2_tile_kernel: CudaFunction,
    d_h_s2_reduce_kernel: CudaFunction,
    cos_tile_kernel: CudaFunction,

    // ── cuBLAS GEMM descriptors ────────────────────────────────────
    /// Trunk Layer 1: [shared_h1, B] = W_s1^T @ states[state_dim, B].
    /// Plan 4 Task 2c.3b: kept allocated but unused — the cuBLAS fallback
    /// that consumed it was replaced with a hard error in
    /// `execute_training_pipeline` (online + target now share GRN).
    /// Removing the descriptor would require re-flowing
    /// `create_iqn_gemm_desc` initialisation order; left in place to
    /// minimise the cross-cutting diff.
    #[allow(dead_code)]
    trunk_l1_gemm: IqnGemmDesc,
    /// Trunk Layer 2: [hidden_dim, B] = W_s2^T @ h1[shared_h1, B]. See
    /// `trunk_l1_gemm` for why this is retained as dead code.
    #[allow(dead_code)]
    trunk_l2_gemm: IqnGemmDesc,
    /// Embedding forward: [H, B*Q] = W_embed^T[H,D] @ cos_features[D, B*Q]
    gemm_fwd_embed: IqnGemmDesc,
    /// Per-branch forward: [bk, B*Q] = W_bk^T[bk,H] @ combined[H, B*Q]
    gemm_fwd_branch: [IqnGemmDesc; 4],
    /// Embedding backward dW: [H, D] = d_embed[H, B*Q] @ cos^T[B*Q, D]
    gemm_bwd_dw_embed: IqnGemmDesc,
    /// Per-branch backward dW: [bk, H] = dq_bk[bk, B*Q] @ combined^T[B*Q, H]
    gemm_bwd_dw_branch: [IqnGemmDesc; 4],
    /// Per-branch backward dX: [H, B*Q] = W_bk[H,bk] @ dq_bk[bk, B*Q]
    gemm_bwd_dx_branch: [IqnGemmDesc; 4],
    /// Embedding backward dX: [D, B*Q] = W_embed[D,H] @ d_embed[H, B*Q]
    /// (Not needed since cos_features are fixed — no gradient through them.
    ///  But we need d_embed → d_combined via hadamard_sigmoid_bwd.)

    // ── IQN online parameters (flat f32 on GPU) ──────────────────────
    online_params: CudaSlice<f32>,
    // ── IQN target parameters (for Bellman target) ───────────────────
    target_params: CudaSlice<f32>,

    // ── Adam optimizer state ─────────────────────────────────────────
    m_buf: CudaSlice<f32>,
    v_buf: CudaSlice<f32>,
    grad_buf: CudaSlice<f32>,
    grad_norm_buf: CudaSlice<f32>,
    /// Per-block partial sums for deterministic two-phase grad norm [num_norm_blocks].
    grad_norm_block_sums: CudaSlice<f32>,
    /// IQN gradient w.r.t. shared trunk h_s2 [B, hidden_dim].
    /// Written by backward, added to C51's trunk gradient.
    d_h_s2_buf: CudaSlice<f32>,

    /// Per-action IQR (Q75 - Q25) computed from quantile estimates.
    /// Used as exploration bonus in action selection to counter magnitude collapse.
    /// Shape: [total_branch_actions] f32 (e.g. 12 = 3+3+3+3).
    iqr_buf: CudaSlice<f32>,

    // ── Trunk cuBLAS scratch buffer ─────────────────────────────────
    /// Intermediate h1 for trunk cuBLAS path [shared_h1, B]. Plan 4 Task
    /// 2c.3b: kept allocated but unused — see `trunk_l1_gemm`.
    #[allow(dead_code)]
    trunk_h1_scratch: CudaSlice<f32>,

    // ── cuBLAS intermediate buffers (forward) ───────────────────────
    /// Tiled h_s2 [H, B*Q] — h_s2 columns replicated Q times
    h_s2_tiled_buf: CudaSlice<f32>,
    /// Tiled cos_features [D, B*Q] — cos columns replicated B times
    cos_features_tiled_buf: CudaSlice<f32>,
    /// Embedding GEMM output before ReLU [H, B*Q]
    embed_pre_buf: CudaSlice<f32>,
    /// Post-ReLU embedding [H, B*Q]
    embed_buf: CudaSlice<f32>,
    /// h_s2_tiled * sigmoid(embed) [H, B*Q]
    combined_buf: CudaSlice<f32>,
    /// Per-branch forward logits [bk, B*Q] (4 branches, concatenated as [TBA, B*Q])
    branch_logits_buf: CudaSlice<f32>,
    /// Target branch logits [TBA, B*Q]
    target_branch_logits_buf: CudaSlice<f32>,

    // ── cuBLAS intermediate buffers (backward) ──────────────────────
    /// Gradient w.r.t. branch logits [TBA, B*Q] (from quantile Huber loss)
    d_branch_logits_buf: CudaSlice<f32>,
    /// Accumulated dX from all branches [H, B*Q]
    d_combined_buf: CudaSlice<f32>,
    /// d_embed from hadamard_sigmoid_bwd [H, B*Q]
    d_embed_buf: CudaSlice<f32>,
    /// d_embed_pre after ReLU backward [H, B*Q]
    d_embed_pre_buf: CudaSlice<f32>,
    /// d_h_s2_tiled from hadamard_sigmoid_bwd [H, B*Q]
    d_h_s2_tiled_buf: CudaSlice<f32>,
    /// Scratch for 2-phase bias_grad_reduce: [ceil(BQ/256), max_out_dim]
    bias_grad_partials_buf: CudaSlice<f32>,
    /// Number of phase-1 blocks along the BQ axis (ceil(BQ/256)).
    bias_grad_num_blocks: usize,

    // ── Per-step buffers ─────────────────────────────────────────────
    /// Per-branch online τ values [4][B, N] — mapped pinned (cuMemHostAlloc DEVICEMAP).
    /// Host writes via host_ptr; IQN kernel reads via dev_ptr. No DtoD copy.
    /// Indexed by `active_branch_idx` at each kernel launch site.
    online_taus_per_branch: [super::mapped_pinned::MappedF32Buffer; 4],
    /// Per-branch target τ values [4][B, N] — mapped pinned.
    target_taus_per_branch: [super::mapped_pinned::MappedF32Buffer; 4],
    /// Per-branch precomputed cosine features [4][D, N] — mapped pinned.
    /// cos_feat[dim + q*D] = cos(π·(d+1)·τ_q). Column-major.
    cos_features_per_branch: [super::mapped_pinned::MappedF32Buffer; 4],
    /// Index of the branch whose τ buffers are currently active for the IQN
    /// forward/backward pass. Set by `activate_branch_taus(branch_idx)`,
    /// reset to 0 by `deactivate_branch_taus`. Kernel launch sites index
    /// the per-branch arrays with this value.
    active_branch_idx: usize,
    /// Branch actions decoded from flat actions [B, 4]
    branch_actions: CudaSlice<i32>,
    /// Target h_s2 computed from next_states + target trunk weights [B, H]
    target_h_s2: CudaSlice<f32>,
    /// Pre-computed target h_s2 from DQN Pass 2 (avoids redundant trunk forward).
    /// Set via `set_cached_target_h_s2()`, consumed (taken) by `execute_training_pipeline`.
    cached_target_h_s2_ptr: Option<u64>,

    /// SP4 Task A14: Pearl C engagement-counter buffer pointers, shared
    /// from the main DQN trainer's mapped-pinned buffers via
    /// `set_pearl_c_buffers`. `nan_flags_buf` is the shared `[i32; 55]`
    /// sticky-flag store; the IQN Adam writes `[51]=1` on overflow.
    /// `clamp_engage_per_block_buf` is the shared `[i32; 2048]` per-block
    /// counter buffer indexed by `engage_buf_offset = ParamGroup::Iqn.idx()
    /// × MAX_BLOCKS_PER_ADAM`. Both default to 0 → kernel skips the write
    /// (compatible with smoke tests / standalone `train_iqn_step_gpu`).
    pearl_c_nan_flags_ptr: u64,
    pearl_c_engage_buf_ptr: u64,
    /// Persistent rewards buffer [B] f32
    rewards_buf: CudaSlice<f32>,
    /// Persistent dones buffer [B] f32
    dones_buf: CudaSlice<f32>,

    // ── Activation save buffers (forward → backward) ─────────────────
    /// Q values from forward pass — used by compute_iqr cold path [TBA, B*Q]
    save_q_online: CudaSlice<f32>,

    // ── Loss buffers ─────────────────────────────────────────────────
    per_sample_loss: CudaSlice<f32>, // [B]
    total_loss: CudaSlice<f32>,      // [1] device buffer for reduce kernel output
    total_loss_pinned: *mut f32,     // pinned host for zero-copy readback
    total_loss_dev_ptr: u64,         // device pointer to pinned total_loss

    // ── Training state ───────────────────────────────────────────────
    adam_step: i32,
    t_pinned: *mut i32,
    t_dev_ptr: u64,
    tau_pinned: *mut f32, // pinned+device-mapped host pointer to tau scalar [1]
    tau_dev_ptr: u64,     // device pointer to same physical page — passed to EMA kernel
    // Plan 4 Task 3 (E.3): `rng_step` (Philox seed counter for per-step τ
    // sampling) deleted alongside the sampler kernel. Fixed-τ broadcast
    // at construction has no per-step state.
    total_params: usize,
}

impl GpuIqnHead {
    /// Create a new GPU IQN dual-head trainer.
    ///
    /// Uses the shared cuBLAS handle for all cublasLtMatmul operations.
    /// Compiles element-wise CUDA kernels, creates GEMM descriptors,
    /// initializes online + target weights with Xavier/Glorot, and
    /// pre-allocates all GPU buffers.
    pub fn new(
        shared_handle: Arc<PerStreamCublasHandles>,
        config: GpuIqnConfig,
    ) -> Result<Self, MLError> {
        let stream = Arc::clone(&shared_handle.stream);
        let total_params = config.total_params();
        let b = config.batch_size;
        let n = config.num_quantiles;
        let h = config.hidden_dim;
        let d = config.embed_dim;
        let tba = config.total_branch_actions();
        let bq = b * n; // B*Q — number of columns in tiled matrices
        let lt_handle = shared_handle.lt_handle.0;
        let lt_ws_size = shared_handle.lt_workspace_size;

        // Load element-wise kernels + retained kernels from precompiled cubin
        let kernels = load_iqn_kernels(&stream, &config)?;

        // ── Create GEMM descriptors ──────────────────────────────────────
        let state_dim = ml_core::state_layout::STATE_DIM;
        let shared_h1 = config.shared_h1;

        // Trunk Layer 1: [shared_h1, B] = W_s1^T @ states[state_dim, B]
        // W_s1 physical [state_dim, shared_h1], TRANSA=T → logical [shared_h1, state_dim]
        // states physical [state_dim, B] (row-major [B, state_dim]), TRANSB=N
        let trunk_l1_gemm = create_iqn_gemm_desc(
            lt_handle, 1, 0,
            shared_h1, b, state_dim,
            state_dim, state_dim, shared_h1,
            lt_ws_size, "trunk_l1",
        )?;

        // Trunk Layer 2: [hidden_dim, B] = W_s2^T @ h1[shared_h1, B]
        // W_s2 physical [shared_h1, hidden_dim], TRANSA=T → logical [hidden_dim, shared_h1]
        // h1 physical [shared_h1, B], TRANSB=N
        let trunk_l2_gemm = create_iqn_gemm_desc(
            lt_handle, 1, 0,
            h, b, shared_h1,
            shared_h1, shared_h1, h,
            lt_ws_size, "trunk_l2",
        )?;

        // Forward embedding: [H, B*Q] = W_embed^T[H,D] @ cos_features_tiled[D, B*Q]
        // TRANSA=T, TRANSB=N: A = W_embed stored [D, H] -> logical [H, D]
        let gemm_fwd_embed = create_iqn_gemm_desc(
            lt_handle, 1, 0, h, bq, d, d, d, h, lt_ws_size, "fwd_embed",
        )?;

        // Forward branches: each outputs separate [bk, B*Q] col-major block.
        // Layout in branch_logits_buf: branch_0[b0*bq] | branch_1[b1*bq] | ...
        // [bk, B*Q] = W_bk^T[bk,H] @ combined[H, B*Q]
        let branch_sizes = [
            config.branch_0_size, config.branch_1_size,
            config.branch_2_size, config.branch_3_size,
        ];
        let gemm_fwd_branch = [
            create_iqn_gemm_desc(lt_handle, 1, 0, branch_sizes[0], bq, h, h, h, branch_sizes[0], lt_ws_size, "fwd_b0")?,
            create_iqn_gemm_desc(lt_handle, 1, 0, branch_sizes[1], bq, h, h, h, branch_sizes[1], lt_ws_size, "fwd_b1")?,
            create_iqn_gemm_desc(lt_handle, 1, 0, branch_sizes[2], bq, h, h, h, branch_sizes[2], lt_ws_size, "fwd_b2")?,
            create_iqn_gemm_desc(lt_handle, 1, 0, branch_sizes[3], bq, h, h, h, branch_sizes[3], lt_ws_size, "fwd_b3")?,
        ];

        // Backward dW_embed: [H, D] = d_embed[H, B*Q] @ cos^T[B*Q, D]
        // TRANSA=N, TRANSB=T: C[H,D] = A[H,B*Q] @ B^T[D,B*Q]^T
        let gemm_bwd_dw_embed = create_iqn_gemm_desc(
            lt_handle, 0, 1, h, d, bq, h, d, h, lt_ws_size, "bwd_dw_embed",
        )?;

        // Backward dW_bk: [bk, H] = dq_bk[bk, B*Q] @ combined^T[B*Q, H]
        // dq_bk is separate contiguous [bk, B*Q] block, lda=bk
        let gemm_bwd_dw_branch = [
            create_iqn_gemm_desc(lt_handle, 0, 1, branch_sizes[0], h, bq, branch_sizes[0], h, branch_sizes[0], lt_ws_size, "bwd_dw_b0")?,
            create_iqn_gemm_desc(lt_handle, 0, 1, branch_sizes[1], h, bq, branch_sizes[1], h, branch_sizes[1], lt_ws_size, "bwd_dw_b1")?,
            create_iqn_gemm_desc(lt_handle, 0, 1, branch_sizes[2], h, bq, branch_sizes[2], h, branch_sizes[2], lt_ws_size, "bwd_dw_b2")?,
            create_iqn_gemm_desc(lt_handle, 0, 1, branch_sizes[3], h, bq, branch_sizes[3], h, branch_sizes[3], lt_ws_size, "bwd_dw_b3")?,
        ];

        // Backward dX_bk: [H, B*Q] = W_bk[H,bk] @ dq_bk[bk, B*Q]
        // W_bk physical [H, bk] with lda=H, TRANSA=N → logical [H, bk]
        // dq_bk is separate [bk, B*Q] block, ldb=bk
        let gemm_bwd_dx_branch = [
            create_iqn_gemm_desc(lt_handle, 0, 0, h, bq, branch_sizes[0], h, branch_sizes[0], h, lt_ws_size, "bwd_dx_b0")?,
            create_iqn_gemm_desc(lt_handle, 0, 0, h, bq, branch_sizes[1], h, branch_sizes[1], h, lt_ws_size, "bwd_dx_b1")?,
            create_iqn_gemm_desc(lt_handle, 0, 0, h, bq, branch_sizes[2], h, branch_sizes[2], h, lt_ws_size, "bwd_dx_b2")?,
            create_iqn_gemm_desc(lt_handle, 0, 0, h, bq, branch_sizes[3], h, branch_sizes[3], h, lt_ws_size, "bwd_dx_b3")?,
        ];

        // ── Allocate and initialize weights (Xavier/Glorot) ─────────────
        let online_params = init_iqn_xavier_weights(&stream, &config)?;
        let target_params = clone_cuda_slice(&stream, &online_params, total_params)?;

        // Adam state (zeroed) — with cuBLAS tile padding
        let cublas_pad = 32 * config.hidden_dim;
        let m_buf = alloc_f32(&stream, total_params + cublas_pad, "iqn_m")?;
        let v_buf = alloc_f32(&stream, total_params + cublas_pad, "iqn_v")?;
        let grad_buf = alloc_f32(&stream, total_params + cublas_pad, "iqn_grad")?;
        let grad_norm_buf = alloc_f32(&stream, 1, "iqn_grad_norm")?;
        let norm_blocks = (total_params + 255) / 256;
        let grad_norm_block_sums = alloc_f32(&stream, norm_blocks, "iqn_grad_norm_block_sums")?;
        let d_h_s2_buf = alloc_f32(&stream, b * h, "iqn_d_h_s2")?;

        // Trunk cuBLAS scratch
        let trunk_h1_scratch = alloc_f32(&stream, shared_h1 * b, "iqn_trunk_h1")?;

        // ── cuBLAS intermediate buffers ─────────────────────────────────
        let h_s2_tiled_buf = alloc_f32(&stream, h * bq, "iqn_h_s2_tiled")?;
        let cos_features_tiled_buf = alloc_f32(&stream, d * bq, "iqn_cos_tiled")?;
        let embed_pre_buf = alloc_f32(&stream, h * bq, "iqn_embed_pre")?;
        let embed_buf = alloc_f32(&stream, h * bq, "iqn_embed")?;
        let combined_buf = alloc_f32(&stream, h * bq, "iqn_combined")?;
        let branch_logits_buf = alloc_f32(&stream, tba * bq, "iqn_branch_logits")?;
        let target_branch_logits_buf = alloc_f32(&stream, tba * bq, "iqn_target_logits")?;

        // Backward intermediates
        let d_branch_logits_buf = alloc_f32(&stream, tba * bq, "iqn_d_branch_logits")?;
        let d_combined_buf = alloc_f32(&stream, h * bq, "iqn_d_combined")?;
        let d_embed_buf = alloc_f32(&stream, h * bq, "iqn_d_embed")?;
        let d_embed_pre_buf = alloc_f32(&stream, h * bq, "iqn_d_embed_pre")?;
        let d_h_s2_tiled_buf = alloc_f32(&stream, h * bq, "iqn_d_h_s2_tiled")?;

        // 2-phase bias_grad_reduce scratch: [ceil(BQ/256), max_out_dim]
        // max_out_dim = h (hidden_dim, used for embed layer — the largest bias)
        let bias_grad_num_blocks = (bq + 255) / 256;
        let bias_grad_partials_buf = alloc_f32(&stream, bias_grad_num_blocks * h, "iqn_bias_grad_partials")?;

        // ── Per-step buffers ────────────────────────────────────────────
        // Plan 4 Task 3 (E.3): fixed-τ {0.05, 0.25, 0.50, 0.75, 0.95}.
        // All 4 per-branch slabs are initialised from FIXED_TAUS at construction.
        // refresh_taus_for_branch / refresh_taus_from_isv write updated values
        // directly to host_ptr (no copy, no DtoD) per
        // feedback_no_htod_htoh_only_mapped_pinned.
        assert_eq!(
            n, FIXED_TAUS.len(),
            "GpuIqnConfig::num_quantiles ({n}) must equal FIXED_TAUS.len() \
             ({}) — kernel-side IQN_NUM_QUANTILES is sized to FIXED_TAUS \
             in lockstep (Plan 4 Task 3, E.3)",
            FIXED_TAUS.len(),
        );
        let fixed: &[f32] = &FIXED_TAUS;

        // Precompute cosine features [D, N] from FIXED_TAUS (col-major):
        // cos_feat[dim + q*D] = cos(π·(d+1)·τ_q)
        let mut cos_feat_host = vec![0.0_f32; d * n];
        for q in 0..n {
            let tau_q = fixed[q];
            for dim in 0..d {
                cos_feat_host[dim + q * d] =
                    (std::f32::consts::PI * ((dim + 1) as f32) * tau_q).cos();
            }
        }
        // Tile taus [B, N]: each row is FIXED_TAUS.
        let mut tiled_fixed = Vec::with_capacity(b * n);
        for _ in 0..b {
            tiled_fixed.extend_from_slice(fixed);
        }

        // Allocate and initialise 4 per-branch mapped-pinned slab triplets.
        // MappedF32Buffer::new zero-inits; write_from_slice populates with FIXED_TAUS data.
        let make_taus_buf = |len: usize| -> Result<super::mapped_pinned::MappedF32Buffer, MLError> {
            unsafe { super::mapped_pinned::MappedF32Buffer::new(len) }
                .map_err(|e| MLError::ModelError(format!("IQN tau buf alloc ({len} f32): {e}")))
        };

        let tau_buf_len = b * n;
        let cos_buf_len = d * n;
        let online_taus_b0 = { let buf = make_taus_buf(tau_buf_len)?; buf.write_from_slice(&tiled_fixed); buf };
        let online_taus_b1 = { let buf = make_taus_buf(tau_buf_len)?; buf.write_from_slice(&tiled_fixed); buf };
        let online_taus_b2 = { let buf = make_taus_buf(tau_buf_len)?; buf.write_from_slice(&tiled_fixed); buf };
        let online_taus_b3 = { let buf = make_taus_buf(tau_buf_len)?; buf.write_from_slice(&tiled_fixed); buf };
        let target_taus_b0 = { let buf = make_taus_buf(tau_buf_len)?; buf.write_from_slice(&tiled_fixed); buf };
        let target_taus_b1 = { let buf = make_taus_buf(tau_buf_len)?; buf.write_from_slice(&tiled_fixed); buf };
        let target_taus_b2 = { let buf = make_taus_buf(tau_buf_len)?; buf.write_from_slice(&tiled_fixed); buf };
        let target_taus_b3 = { let buf = make_taus_buf(tau_buf_len)?; buf.write_from_slice(&tiled_fixed); buf };
        let cos_features_b0 = { let buf = make_taus_buf(cos_buf_len)?; buf.write_from_slice(&cos_feat_host); buf };
        let cos_features_b1 = { let buf = make_taus_buf(cos_buf_len)?; buf.write_from_slice(&cos_feat_host); buf };
        let cos_features_b2 = { let buf = make_taus_buf(cos_buf_len)?; buf.write_from_slice(&cos_feat_host); buf };
        let cos_features_b3 = { let buf = make_taus_buf(cos_buf_len)?; buf.write_from_slice(&cos_feat_host); buf };

        let branch_actions = stream.alloc_zeros::<i32>(b * 4).map_err(|e| {
            MLError::ModelError(format!("IQN alloc branch_actions: {e}"))
        })?;
        let target_h_s2 = alloc_f32(&stream, b * h, "iqn_target_h_s2")?;

        // Activation saves (for compute_iqr cold path)
        let save_q_online = alloc_f32(&stream, tba * bq, "iqn_save_q_online")?;

        // Loss buffers
        let per_sample_loss = alloc_f32(&stream, b, "iqn_per_sample_loss")?;
        let total_loss = alloc_f32(&stream, 1, "iqn_total_loss")?;
        // Pinned device-mapped for zero-copy readback (no cuStreamSynchronize per step)
        let total_loss_pinned: *mut f32 = unsafe {
            let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP;
            cudarc::driver::result::malloc_host(std::mem::size_of::<f32>(), flags)
                .map_err(|e| MLError::ModelError(format!("pinned iqn_total_loss alloc: {e}")))?
                as *mut f32
        };
        unsafe { *total_loss_pinned = 0.0; }
        let total_loss_dev_ptr = unsafe {
            let mut dp = 0u64;
            cudarc::driver::sys::cuMemHostGetDevicePointer_v2(&mut dp as *mut u64, total_loss_pinned.cast(), 0);
            dp
        };

        // Persistent rewards/dones buffers
        let rewards_buf = stream.alloc_zeros::<f32>(b)
            .map_err(|e| MLError::ModelError(format!("GpuIqn alloc iqn_rewards ({b} f32): {e}")))?;
        let dones_buf = stream.alloc_zeros::<f32>(b)
            .map_err(|e| MLError::ModelError(format!("GpuIqn alloc iqn_dones ({b} f32): {e}")))?;

        // IQR buffer
        let iqr_buf = alloc_f32(&stream, tba, "iqn_iqr")?;

        // SP6 Pearl 5: allocate per-branch tau/cos buffer triplets (same size as main buffers).
        let vram_bytes = (total_params * 6 + h * bq * 8 + tba * bq * 4
            + d * bq + b * h * 2 + b * 4 + tba + b * 2 + 2) * 4;

        info!(
            hidden_dim = h,
            embed_dim = d,
            num_quantiles = n,
            kappa = config.kappa,
            total_params,
            total_branch_actions = tba,
            vram_kb = vram_bytes / 1024,
            "GpuIqnHead initialized: cuBLAS GEMMs + element-wise kernels, Xavier weights"
        );

        let t_pinned: *mut i32 = unsafe {
            let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP;
            cudarc::driver::result::malloc_host(std::mem::size_of::<i32>(), flags)
                .map_err(|e| MLError::ModelError(format!("iqn pinned t alloc: {e}")))?
                as *mut i32
        };
        unsafe { *t_pinned = 0; }
        let t_dev_ptr = unsafe {
            let mut dp = 0u64;
            cudarc::driver::sys::cuMemHostGetDevicePointer_v2(&mut dp as *mut u64, t_pinned.cast(), 0);
            dp
        };
        // Allocate tau as pinned+device-mapped so the EMA kernel reads it directly
        // from host-visible memory without any PCIe copy on the per-step hot path.
        let tau_pinned: *mut f32 = unsafe {
            let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP;
            cudarc::driver::result::malloc_host(std::mem::size_of::<f32>(), flags)
                .map_err(|e| MLError::ModelError(format!("iqn tau_pinned alloc: {e}")))?
                as *mut f32
        };
        unsafe { *tau_pinned = 0.0_f32; }
        let tau_dev_ptr = unsafe {
            let mut dp = 0u64;
            cudarc::driver::sys::cuMemHostGetDevicePointer_v2(
                &mut dp as *mut u64,
                tau_pinned.cast(),
                0,
            );
            dp
        };

        Ok(Self {
            config,
            stream,
            shared_handle,
            loss_reduce_kernel: kernels.loss_reduce,
            grad_norm_phase1_kernel: kernels.gnorm_p1,
            grad_norm_phase2_kernel: kernels.gnorm_p2,
            adam_kernel: kernels.adam,
            forward_kernel: kernels.fwd_only,
            ema_kernel: kernels.ema,
            decode_actions_kernel: kernels.decode_act,
            relu_fwd_kernel: kernels.relu_fwd,
            relu_bwd_kernel: kernels.relu_bwd,
            hadamard_sigmoid_kernel: kernels.hadamard_sigmoid,
            hadamard_sigmoid_bwd_kernel: kernels.hadamard_sigmoid_bwd,
            quantile_huber_loss_kernel: kernels.quantile_huber_loss,
            bias_add_kernel: kernels.bias_add,
            bias_grad_reduce_p1_kernel: kernels.bias_grad_reduce_p1,
            bias_grad_reduce_p2_kernel: kernels.bias_grad_reduce_p2,
            h_s2_tile_kernel: kernels.h_s2_tile,
            d_h_s2_reduce_kernel: kernels.d_h_s2_reduce,
            cos_tile_kernel: kernels.cos_tile,
            trunk_l1_gemm,
            trunk_l2_gemm,
            gemm_fwd_embed,
            gemm_fwd_branch,
            gemm_bwd_dw_embed,
            gemm_bwd_dw_branch,
            gemm_bwd_dx_branch,
            online_params,
            target_params,
            m_buf,
            v_buf,
            grad_buf,
            grad_norm_buf,
            grad_norm_block_sums,
            d_h_s2_buf,
            iqr_buf,
            trunk_h1_scratch,
            h_s2_tiled_buf,
            cos_features_tiled_buf,
            embed_pre_buf,
            embed_buf,
            combined_buf,
            branch_logits_buf,
            target_branch_logits_buf,
            d_branch_logits_buf,
            d_combined_buf,
            d_embed_buf,
            d_embed_pre_buf,
            d_h_s2_tiled_buf,
            bias_grad_partials_buf,
            bias_grad_num_blocks,
            online_taus_per_branch: [online_taus_b0, online_taus_b1, online_taus_b2, online_taus_b3],
            target_taus_per_branch: [target_taus_b0, target_taus_b1, target_taus_b2, target_taus_b3],
            cos_features_per_branch: [cos_features_b0, cos_features_b1, cos_features_b2, cos_features_b3],
            active_branch_idx: 0,
            branch_actions,
            target_h_s2,
            cached_target_h_s2_ptr: None,
            pearl_c_nan_flags_ptr: 0,
            pearl_c_engage_buf_ptr: 0,
            rewards_buf,
            dones_buf,
            save_q_online,
            per_sample_loss,
            total_loss,
            total_loss_pinned,
            total_loss_dev_ptr,
            adam_step: 0,
            t_pinned,
            t_dev_ptr,
            tau_pinned,
            tau_dev_ptr,
            total_params,
        })
    }

    /// Prepare IQN buffers: decode flat actions and DtoD copy rewards/dones.
    ///
    /// Runs on the IQN head's own stream (the main training stream). Called
    /// Launch bias_add + relu_fwd for IQN trunk cuBLAS GEMMs (in-place on data_ptr).
    /// Plan 4 Task 2c.3b: dead code — the cuBLAS fallback that called
    /// this is gone; online + target share `target_encoder_forward_only`.
    #[allow(dead_code)]
    fn launch_trunk_bias_relu(
        &self,
        stream: &CudaStream,
        data_ptr: u64,
        bias_ptr: u64,
        out_dim: usize,
        batch: usize,
    ) -> Result<(), MLError> {
        let total = out_dim * batch;
        let out_dim_i32 = out_dim as i32;
        let total_i32 = total as i32;
        let elem_cfg = LaunchConfig::for_num_elems(total as u32);
        unsafe {
            stream.launch_builder(&self.bias_add_kernel)
                .arg(&data_ptr).arg(&bias_ptr).arg(&out_dim_i32).arg(&total_i32)
                .launch(elem_cfg)
                .map_err(|e| MLError::ModelError(format!("trunk bias_add: {e}")))?;
            stream.launch_builder(&self.relu_fwd_kernel)
                .arg(&data_ptr).arg(&data_ptr).arg(&total_i32)
                .launch(elem_cfg)
                .map_err(|e| MLError::ModelError(format!("trunk relu_fwd: {e}")))?;
        }
        Ok(())
    }

    /// Prepare IQN buffers: decode flat actions and DtoD copy rewards/dones.
    ///
    /// Runs on the IQN head's own stream (the main training stream). Called
    /// before `execute_training_pipeline` so that a dedicated override stream
    /// can safely read `branch_actions`, `rewards_buf`, and `dones_buf`.
    pub fn prepare_buffers(
        &mut self,
        dqn_actions_buf: &CudaSlice<i32>,
        dqn_rewards_buf: &CudaSlice<f32>,
        dqn_dones_buf: &CudaSlice<f32>,
    ) -> Result<(), MLError> {
        let b = self.config.batch_size;
        let shared_h1_i32 = self.config.shared_h1 as i32;
        let hidden_dim_i32 = self.config.hidden_dim as i32;
        let embed_dim_i32 = self.config.embed_dim as i32;
        let b0_i32 = self.config.branch_0_size as i32;
        let b1_i32 = self.config.branch_1_size as i32;
        let b2_i32 = self.config.branch_2_size as i32;
        let b3_i32 = self.config.branch_3_size as i32;

        let decode_blocks = (b + 255) / 256;
        let decode_config = LaunchConfig {
            grid_dim: (decode_blocks as u32, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };
        let batch_size_i32 = b as i32;

        unsafe {
            self.stream
                .launch_builder(&self.decode_actions_kernel)
                .arg(dqn_actions_buf)
                .arg(&mut self.branch_actions)
                .arg(&batch_size_i32)
                .arg(&shared_h1_i32)
                .arg(&hidden_dim_i32)
                .arg(&embed_dim_i32)
                .arg(&b0_i32)
                .arg(&b1_i32)
                .arg(&b2_i32)
                .arg(&b3_i32)
                .launch(decode_config)
                .map_err(|e| MLError::ModelError(format!("IQN decode_actions kernel: {e}")))?;
        }

        use super::gpu_dqn_trainer::dtod_copy_async;
        let f32_bytes = std::mem::size_of::<f32>();
        let r_src = dqn_rewards_buf.raw_ptr();
        let r_dst = self.rewards_buf.raw_ptr();
        dtod_copy_async(r_dst, r_src, b * f32_bytes, &self.stream, 0, "iqn_rewards_dtod")?;

        let d_src = dqn_dones_buf.raw_ptr();
        let d_dst = self.dones_buf.raw_ptr();
        dtod_copy_async(d_dst, d_src, b * f32_bytes, &self.stream, 0, "iqn_dones_dtod")?;

        Ok(())
    }

    /// Run one GPU-native IQN training step.
    ///
    /// All data sourced from DQN trainer's GPU buffers via DtoD copies.
    /// No CPU arrays, no HtoD uploads for batch data.
    ///
    /// Full pipeline:
    /// 1. Decode flat actions → branch actions on GPU
    /// 2. DtoD copy rewards/dones from DQN trainer buffers
    /// 3. Compute target h_s2 (or use cached)
    /// 4. Forward (cuBLAS GEMMs + element-wise kernels)
    /// 5. Loss (quantile Huber)
    /// 6. Backward (cuBLAS GEMMs + element-wise kernels)
    /// 7. Grad norm + Adam
    ///
    /// Returns the batch-mean quantile Huber loss.
    pub fn train_iqn_step_gpu(
        &mut self,
        online_h_s2: &CudaSlice<f32>,
        next_states_buf: &CudaSlice<f32>,
        target_dueling: &DuelingWeightSet,
        dqn_actions_buf: &CudaSlice<i32>,
        dqn_rewards_buf: &CudaSlice<f32>,
        dqn_dones_buf: &CudaSlice<f32>,
        /* SP4 Layer B (Mech 9): ISV-derived |p| bound for `iqn_adam_kernel`
         * — caller reads `ISV[WEIGHT_BOUND[Iqn=3]]` and passes it through.
         * 0=disabled (debug only); production callers always set a non-zero
         * bound. */
        weight_clamp_max_abs: f32,
        /* SP4 Layer B (Mech 9): ISV-derived AdamW weight_decay rate. */
        weight_decay_value: f32,
        /* SP5 Pearl 4: ISV-driven Adam β1/β2/ε for ParamGroup::Iqn (index 3).
         * Caller reads ISV[ADAM_BETA1_BASE+3], ISV[ADAM_BETA2_BASE+3],
         * ISV[ADAM_EPS_BASE+3] and threads through here. */
        iqn_beta1: f32,
        iqn_beta2: f32,
        iqn_epsilon: f32,
    ) -> Result<f32, MLError> {
        // 1-2. Decode actions + DtoD copy rewards/dones
        self.prepare_buffers(dqn_actions_buf, dqn_rewards_buf, dqn_dones_buf)?;

        // 3-9. Execute core training pipeline
        self.execute_training_pipeline(
            online_h_s2,
            next_states_buf,
            target_dueling,
            None,
            None,
            weight_clamp_max_abs,
            weight_decay_value,
            iqn_beta1,
            iqn_beta2,
            iqn_epsilon,
        )
    }

    /// Set a pre-computed target h_s2 buffer from the DQN trainer's Pass 2.
    /// When set, `execute_training_pipeline` skips the redundant trunk forward kernel.
    /// The pointer must remain valid until the next `train_iqn_step_gpu` call.
    pub fn set_cached_target_h_s2(&mut self, ptr: u64) {
        self.cached_target_h_s2_ptr = Some(ptr);
    }

    /// SP4 Task A14: thread the main DQN trainer's mapped-pinned Pearl C
    /// buffers down to the IQN Adam launch. `nan_flags` is the shared
    /// `[i32; SP4_NAN_FLAGS_END]` sticky-flag buffer; `engage_buf` is the
    /// shared `[i32; SP4_ENGAGE_BUF_LEN]` per-block engagement-count buffer.
    /// Pass 0 for either pointer to disable Pearl C for this IQN trainer
    /// (smoke tests / standalone runs).
    pub fn set_pearl_c_buffers(&mut self, nan_flags: u64, engage_buf: u64) {
        self.pearl_c_nan_flags_ptr = nan_flags;
        self.pearl_c_engage_buf_ptr = engage_buf;
    }

    /// Core IQN training pipeline: trunk forward → cuBLAS forward → loss → backward → Adam.
    ///
    /// Assumes `branch_actions`, `rewards_buf`, `dones_buf` are already populated.
    ///
    /// `override_stream` and `override_workspace` allow a dedicated CUDA stream and
    /// cuBLAS workspace to be used instead of the IQN head's defaults.  Pass `None`
    /// for both to preserve the existing behaviour (used by all current callers).
    pub fn execute_training_pipeline(
        &mut self,
        online_h_s2: &CudaSlice<f32>,
        next_states_buf: &CudaSlice<f32>,
        target_dueling: &DuelingWeightSet,
        override_stream: Option<&Arc<CudaStream>>,
        override_workspace: Option<(u64, usize)>,
        /* SP4 Layer B (Mech 9): ISV-derived |p| bound for the trailing
         * `iqn_adam_kernel` arg. Caller reads `ISV[WEIGHT_BOUND[Iqn=3]]`
         * with EPS_CLAMP_FLOOR ε. 0=disabled (debug only). */
        weight_clamp_max_abs: f32,
        /* SP4 Layer B (Mech 9): ISV-derived AdamW weight_decay rate.
         * Caller reads `ISV[WD_RATE[Iqn=3]]` with `.max(0.0)` non-negativity. */
        weight_decay_value: f32,
        /* SP5 Layer B (Pearl 4): ISV-derived Adam β1/β2/ε for ParamGroup::Iqn=3.
         * Caller reads ISV[ADAM_BETA1_BASE+3], ISV[ADAM_BETA2_BASE+3], ISV[ADAM_EPS_BASE+3]
         * and applies the Invariant 1 floors (β1≥0.5, β2≥0.9, ε≥1e-12). */
        iqn_beta1: f32,
        iqn_beta2: f32,
        iqn_epsilon: f32,
    ) -> Result<f32, MLError> {
        let effective_stream = override_stream.unwrap_or(&self.stream);
        let (lt_ws_ptr, lt_ws_size) = override_workspace
            .unwrap_or((self.shared_handle.lt_workspace_ptr, self.shared_handle.lt_workspace_size));
        let b = self.config.batch_size;
        let n = self.config.num_quantiles;
        let h = self.config.hidden_dim;
        let d = self.config.embed_dim;
        let bq = b * n;
        let tba = self.config.total_branch_actions();
        let branch_sizes = [
            self.config.branch_0_size, self.config.branch_1_size,
            self.config.branch_2_size, self.config.branch_3_size,
        ];
        let shared_h1_i32 = self.config.shared_h1 as i32;
        let hidden_dim_i32 = self.config.hidden_dim as i32;
        let embed_dim_i32 = self.config.embed_dim as i32;
        let b0_i32 = self.config.branch_0_size as i32;
        let b1_i32 = self.config.branch_1_size as i32;
        let b2_i32 = self.config.branch_2_size as i32;
        let b3_i32 = self.config.branch_3_size as i32;

        // ── Step 1: Target h_s2 (cached from DQN's target trunk forward) ──
        //
        // Plan 4 Task 2c.3b: the cuBLAS fallback that ran a legacy
        // `Linear→ReLU→Linear` on `target_dueling` weights is gone — it
        // was duplicate code that produced ReLU activations while the
        // primary DQN target forward (now `target_encoder_forward_only`)
        // produces LayerNorm activations. Online and target trunks now
        // share ONE implementation: `BatchedForward::target_encoder_forward_only`,
        // which writes `tg_h_s2_buf`. The trainer pre-populates that
        // pointer here via `set_cached_target_h_s2` before each step;
        // a missing cache is now a hard error because there is no
        // fallback that would produce semantically-equivalent activations.
        let cached_ptr = self.cached_target_h_s2_ptr.take().ok_or_else(|| {
            MLError::ModelError(
                "IQN execute_training_pipeline: cached_target_h_s2_ptr is None. \
                 Plan 4 Task 2c.3b: the legacy cuBLAS fallback was deleted \
                 because online and target trunks must share the GRN forward \
                 implementation. The caller must call set_cached_target_h_s2 \
                 with the buffer written by BatchedForward::target_encoder_forward_only \
                 before invoking this method."
                    .to_string(),
            )
        })?;
        {
            let n_bytes = b * h * std::mem::size_of::<f32>();
            let dst = self.target_h_s2.raw_ptr();
            unsafe {
                let r = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
                    dst, cached_ptr, n_bytes, effective_stream.cu_stream(),
                );
                if r != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
                    return Err(MLError::ModelError(format!("IQN cached h_s2 DtoD: {r:?}")));
                }
            }
            // Suppress warnings for the now-unused fallback inputs.
            let _ = (target_dueling, next_states_buf, lt_ws_ptr, lt_ws_size);
        }

        // ── Step 2: Tile h_s2 [H,B] → h_s2_tiled [H, B*Q] ─────────────
        let elem_config = |total: usize| LaunchConfig {
            grid_dim: (((total + 255) / 256) as u32, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };

        let h_i32 = h as i32;
        let b_i32 = b as i32;
        let n_i32 = n as i32;
        let d_i32 = d as i32;

        // Tile online h_s2
        unsafe {
            effective_stream
                .launch_builder(&self.h_s2_tile_kernel)
                .arg(online_h_s2)
                .arg(&mut self.h_s2_tiled_buf)
                .arg(&h_i32)
                .arg(&b_i32)
                .arg(&n_i32)
                .launch(elem_config(h * bq))
                .map_err(|e| MLError::ModelError(format!("IQN h_s2_tile: {e}")))?;
        }

        // Tile cos_features [D, N] → [D, B*Q]
        unsafe {
            effective_stream
                .launch_builder(&self.cos_tile_kernel)
                .arg(&self.cos_features_per_branch[self.active_branch_idx].dev_ptr)
                .arg(&mut self.cos_features_tiled_buf)
                .arg(&d_i32)
                .arg(&b_i32)
                .arg(&n_i32)
                .launch(elem_config(d * bq))
                .map_err(|e| MLError::ModelError(format!("IQN cos_tile: {e}")))?;
        }

        // Weight layout in online_params:
        // W_embed[H, D] | b_embed[H] | W_b0[b0, H] | b_b0[b0] | ... | W_b3[b3, H] | b_b3[b3]
        let params_ptr = self.online_params.raw_ptr();
        let f32_sz = std::mem::size_of::<f32>() as u64;
        let w_embed_ptr = params_ptr;
        let b_embed_ptr = params_ptr + (h * d) as u64 * f32_sz;
        let mut branch_w_ptrs = [0u64; 4];
        let mut branch_b_ptrs = [0u64; 4];
        {
            let mut off = h * d + h; // past W_embed + b_embed
            for i in 0..4 {
                branch_w_ptrs[i] = params_ptr + off as u64 * f32_sz;
                off += branch_sizes[i] * h;
                branch_b_ptrs[i] = params_ptr + off as u64 * f32_sz;
                off += branch_sizes[i];
            }
        }

        // Per-stream lt_handle (Option C determinism fix).
        let lt_handle_raw = self.shared_handle.lt_handle_for(effective_stream)?;

        // ── Step 3: Forward embedding GEMM ──────────────────────────────
        // embed_pre[H, B*Q] = W_embed^T @ cos_features_tiled
        iqn_lt_matmul(
            effective_stream, lt_handle_raw, &self.gemm_fwd_embed,
            1.0, 0.0,
            w_embed_ptr, self.cos_features_tiled_buf.raw_ptr(),
            self.embed_pre_buf.raw_ptr(),
            lt_ws_ptr, lt_ws_size, "iqn_fwd_embed",
        )?;

        // Step 3b: Bias add + ReLU → embed_buf
        let hbq = h * bq;
        let hbq_i32 = hbq as i32;
        unsafe {
            effective_stream
                .launch_builder(&self.bias_add_kernel)
                .arg(&mut self.embed_pre_buf)
                .arg(&b_embed_ptr)
                .arg(&h_i32)
                .arg(&hbq_i32)
                .launch(elem_config(hbq))
                .map_err(|e| MLError::ModelError(format!("IQN bias_add embed: {e}")))?;
        }
        unsafe {
            effective_stream
                .launch_builder(&self.relu_fwd_kernel)
                .arg(&self.embed_pre_buf)
                .arg(&mut self.embed_buf)
                .arg(&hbq_i32)
                .launch(elem_config(hbq))
                .map_err(|e| MLError::ModelError(format!("IQN relu_fwd embed: {e}")))?;
        }

        // ── Step 4: Hadamard product: combined = h_s2_tiled * sigmoid(embed) ──
        unsafe {
            effective_stream
                .launch_builder(&self.hadamard_sigmoid_kernel)
                .arg(&self.h_s2_tiled_buf)
                .arg(&self.embed_buf)
                .arg(&mut self.combined_buf)
                .arg(&hbq_i32)
                .launch(elem_config(hbq))
                .map_err(|e| MLError::ModelError(format!("IQN hadamard_sigmoid: {e}")))?;
        }

        // ── Step 5: Per-branch forward GEMMs ────────────────────────────
        // Concatenated output in branch_logits_buf [TBA, B*Q]
        {
            let mut row_off = 0usize;
            for i in 0..4 {
                let bs = branch_sizes[i];
                let out_ptr = self.branch_logits_buf.raw_ptr() + (row_off * bq) as u64 * f32_sz;
                iqn_lt_matmul(
                    effective_stream, lt_handle_raw, &self.gemm_fwd_branch[i],
                    1.0, 0.0,
                    branch_w_ptrs[i], self.combined_buf.raw_ptr(),
                    out_ptr,
                    lt_ws_ptr, lt_ws_size,
                    &format!("iqn_fwd_branch_{i}"),
                )?;
                // Bias add for this branch
                let bs_i32 = bs as i32;
                let total_i32 = (bs * bq) as i32;
                unsafe {
                    effective_stream
                        .launch_builder(&self.bias_add_kernel)
                        .arg(&out_ptr)
                        .arg(&branch_b_ptrs[i])
                        .arg(&bs_i32)
                        .arg(&total_i32)
                        .launch(elem_config(bs * bq))
                        .map_err(|e| MLError::ModelError(format!("IQN bias_add branch_{i}: {e}")))?;
                }
                row_off += bs;
            }
        }

        // Copy online logits → save_q_online for compute_iqr
        {
            let n_bytes = tba * bq * std::mem::size_of::<f32>();
            let src = self.branch_logits_buf.raw_ptr();
            let dst = self.save_q_online.raw_ptr();
            unsafe {
                let r = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(dst, src, n_bytes, effective_stream.cu_stream());
                if r != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
                    return Err(MLError::ModelError(format!("IQN save_q_online DtoD: {r:?}")));
                }
            }
        }

        // ── Step 5b: Target forward pass (same GEMMs with target params) ──
        {
            let tgt_ptr = self.target_params.raw_ptr();
            let tw_embed_ptr = tgt_ptr;
            let tb_embed_ptr = tgt_ptr + (h * d) as u64 * f32_sz;

            // Tile target_h_s2 into h_s2_tiled_buf (reuse buffer)
            unsafe {
                effective_stream
                    .launch_builder(&self.h_s2_tile_kernel)
                    .arg(&self.target_h_s2)
                    .arg(&mut self.h_s2_tiled_buf)
                    .arg(&h_i32)
                    .arg(&b_i32)
                    .arg(&n_i32)
                    .launch(elem_config(h * bq))
                    .map_err(|e| MLError::ModelError(format!("IQN target h_s2_tile: {e}")))?;
            }

            // Target embedding forward
            iqn_lt_matmul(
                effective_stream, lt_handle_raw, &self.gemm_fwd_embed,
                1.0, 0.0,
                tw_embed_ptr, self.cos_features_tiled_buf.raw_ptr(),
                self.embed_pre_buf.raw_ptr(), // reuse for target
                lt_ws_ptr, lt_ws_size, "iqn_tgt_fwd_embed",
            )?;
            unsafe {
                effective_stream
                    .launch_builder(&self.bias_add_kernel)
                    .arg(&mut self.embed_pre_buf)
                    .arg(&tb_embed_ptr)
                    .arg(&h_i32)
                    .arg(&hbq_i32)
                    .launch(elem_config(hbq))
                    .map_err(|e| MLError::ModelError(format!("IQN tgt bias_add embed: {e}")))?;
            }
            // ReLU into a scratch — reuse embed_buf (we saved online logits already)
            unsafe {
                effective_stream
                    .launch_builder(&self.relu_fwd_kernel)
                    .arg(&self.embed_pre_buf)
                    .arg(&mut self.embed_buf)
                    .arg(&hbq_i32)
                    .launch(elem_config(hbq))
                    .map_err(|e| MLError::ModelError(format!("IQN tgt relu_fwd: {e}")))?;
            }
            // Target combined
            unsafe {
                effective_stream
                    .launch_builder(&self.hadamard_sigmoid_kernel)
                    .arg(&self.h_s2_tiled_buf)
                    .arg(&self.embed_buf)
                    .arg(&mut self.combined_buf) // reuse combined_buf for target
                    .arg(&hbq_i32)
                    .launch(elem_config(hbq))
                    .map_err(|e| MLError::ModelError(format!("IQN tgt hadamard_sigmoid: {e}")))?;
            }

            // Target branch GEMMs
            let mut tgt_off = h * d + h;
            let mut row_off = 0usize;
            for i in 0..4 {
                let bs = branch_sizes[i];
                let tw_ptr = tgt_ptr + tgt_off as u64 * f32_sz;
                tgt_off += bs * h;
                let tb_ptr = tgt_ptr + tgt_off as u64 * f32_sz;
                tgt_off += bs;

                let out_ptr = self.target_branch_logits_buf.raw_ptr() + (row_off * bq) as u64 * f32_sz;
                iqn_lt_matmul(
                    effective_stream, lt_handle_raw, &self.gemm_fwd_branch[i],
                    1.0, 0.0,
                    tw_ptr, self.combined_buf.raw_ptr(),
                    out_ptr,
                    lt_ws_ptr, lt_ws_size,
                    &format!("iqn_tgt_fwd_branch_{i}"),
                )?;
                let bs_i32 = bs as i32;
                let total_i32 = (bs * bq) as i32;
                unsafe {
                    effective_stream
                        .launch_builder(&self.bias_add_kernel)
                        .arg(&out_ptr)
                        .arg(&tb_ptr)
                        .arg(&bs_i32)
                        .arg(&total_i32)
                        .launch(elem_config(bs * bq))
                        .map_err(|e| MLError::ModelError(format!("IQN tgt bias_add branch_{i}: {e}")))?;
                }
                row_off += bs;
            }
        }

        // ── Step 6: Quantile Huber loss + dq gradients ──────────────────
        // Zero d_branch_logits_buf — raw cuMemsetD8Async, no cudarc overhead during graph capture
        unsafe {
            cudarc::driver::sys::cuMemsetD32Async(
                self.d_branch_logits_buf.raw_ptr(), 0,
                self.d_branch_logits_buf.len(),
                effective_stream.cu_stream(),
            );
        }

        let kappa = self.config.kappa;
        let gamma = self.config.gamma;
        let batch_size_i32 = b as i32;

        unsafe {
            effective_stream
                .launch_builder(&self.quantile_huber_loss_kernel)
                .arg(&self.branch_logits_buf)
                .arg(&self.target_branch_logits_buf)
                .arg(&self.online_taus_per_branch[self.active_branch_idx].dev_ptr)
                .arg(&self.branch_actions)
                .arg(&mut self.per_sample_loss)
                .arg(&mut self.d_branch_logits_buf)
                .arg(&batch_size_i32)
                .arg(&n_i32)
                .arg(&b0_i32)
                .arg(&b1_i32)
                .arg(&b2_i32)
                .arg(&b3_i32)
                .arg(&kappa)
                .arg(&gamma)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 8 * std::mem::size_of::<f32>() as u32,
                })
                .map_err(|e| MLError::ModelError(format!("IQN quantile_huber_loss: {e}")))?;
        }

        // Loss reduce → pinned device-mapped buffer (zero-copy readback, no sync)
        unsafe {
            cudarc::driver::sys::cuMemsetD32Async(
                self.total_loss_dev_ptr, 0, 1, effective_stream.cu_stream(),
            );
            let loss_ptr = self.total_loss_dev_ptr;
            effective_stream
                .launch_builder(&self.loss_reduce_kernel)
                .arg(&self.per_sample_loss)
                .arg(&loss_ptr)
                .arg(&batch_size_i32)
                .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 })
                .map_err(|e| MLError::ModelError(format!("IQN loss_reduce: {e}")))?;
        }

        // ── Step 7: Backward pass (cuBLAS GEMMs) ────────────────────────
        // We need online h_s2_tiled, embed_buf, combined_buf, embed_pre_buf
        // from the ONLINE forward pass. But we overwrote them with target.
        // Re-tile online h_s2 and redo online embedding forward to recover saves.
        unsafe {
            effective_stream
                .launch_builder(&self.h_s2_tile_kernel)
                .arg(online_h_s2)
                .arg(&mut self.h_s2_tiled_buf)
                .arg(&h_i32)
                .arg(&b_i32)
                .arg(&n_i32)
                .launch(elem_config(h * bq))
                .map_err(|e| MLError::ModelError(format!("IQN bwd h_s2_tile: {e}")))?;
        }
        // Redo online embedding forward to recover embed_pre_buf + embed_buf + combined_buf
        iqn_lt_matmul(
            effective_stream, lt_handle_raw, &self.gemm_fwd_embed,
            1.0, 0.0,
            w_embed_ptr, self.cos_features_tiled_buf.raw_ptr(),
            self.embed_pre_buf.raw_ptr(),
            lt_ws_ptr, lt_ws_size, "iqn_bwd_recover_embed",
        )?;
        unsafe {
            effective_stream
                .launch_builder(&self.bias_add_kernel)
                .arg(&mut self.embed_pre_buf)
                .arg(&b_embed_ptr)
                .arg(&h_i32)
                .arg(&hbq_i32)
                .launch(elem_config(hbq))
                .map_err(|e| MLError::ModelError(format!("IQN bwd recover bias_add: {e}")))?;
        }
        unsafe {
            effective_stream
                .launch_builder(&self.relu_fwd_kernel)
                .arg(&self.embed_pre_buf)
                .arg(&mut self.embed_buf)
                .arg(&hbq_i32)
                .launch(elem_config(hbq))
                .map_err(|e| MLError::ModelError(format!("IQN bwd recover relu: {e}")))?;
        }
        unsafe {
            effective_stream
                .launch_builder(&self.hadamard_sigmoid_kernel)
                .arg(&self.h_s2_tiled_buf)
                .arg(&self.embed_buf)
                .arg(&mut self.combined_buf)
                .arg(&hbq_i32)
                .launch(elem_config(hbq))
                .map_err(|e| MLError::ModelError(format!("IQN bwd recover hadamard: {e}")))?;
        }

        // Now backward through branches:
        let inv_bq = 1.0_f32 / bq as f32;

        // Zero grad_buf + d_combined_buf — raw cuMemsetD8Async, no cudarc overhead
        unsafe {
            cudarc::driver::sys::cuMemsetD32Async(
                self.grad_buf.raw_ptr(), 0, self.grad_buf.len(),
                effective_stream.cu_stream(),
            );
            cudarc::driver::sys::cuMemsetD32Async(
                self.d_combined_buf.raw_ptr(), 0, self.d_combined_buf.len(),
                effective_stream.cu_stream(),
            );
        }

        {
            let mut row_off = 0usize;
            let mut grad_off = h * d + h; // skip W_embed + b_embed in grad_buf
            for i in 0..4 {
                let bs = branch_sizes[i];
                let dq_ptr = self.d_branch_logits_buf.raw_ptr() + (row_off * bq) as u64 * f32_sz;

                // dW_bk = dq_bk @ combined^T (alpha=inv_bq for mean reduction)
                let dw_ptr = self.grad_buf.raw_ptr() + grad_off as u64 * f32_sz;
                iqn_lt_matmul(
                    effective_stream, lt_handle_raw, &self.gemm_bwd_dw_branch[i],
                    inv_bq, 0.0,
                    dq_ptr, self.combined_buf.raw_ptr(),
                    dw_ptr,
                    lt_ws_ptr, lt_ws_size,
                    &format!("iqn_bwd_dw_branch_{i}"),
                )?;
                grad_off += bs * h;

                // db_bk: 2-phase bias gradient reduce
                // Phase 1: grid(ceil(BQ/256), bs, 1), block(256) — partial sums per block
                // Phase 2: grid(bs, 1, 1), block(256) — final reduce + scale
                let db_ptr = self.grad_buf.raw_ptr() + grad_off as u64 * f32_sz;
                let bs_i32 = bs as i32;
                let bq_i32 = bq as i32;
                let p1_blocks = self.bias_grad_num_blocks as u32;
                let num_blocks_i32 = self.bias_grad_num_blocks as i32;
                let partials_ptr = self.bias_grad_partials_buf.raw_ptr();
                unsafe {
                    effective_stream
                        .launch_builder(&self.bias_grad_reduce_p1_kernel)
                        .arg(&dq_ptr)
                        .arg(&partials_ptr)
                        .arg(&bs_i32)
                        .arg(&bq_i32)
                        .launch(LaunchConfig {
                            grid_dim: (p1_blocks, bs as u32, 1),
                            block_dim: (256, 1, 1),
                            shared_mem_bytes: 256 * 4,
                        })
                        .map_err(|e| MLError::ModelError(format!("IQN bias_grad_reduce_p1 branch_{i}: {e}")))?;
                    effective_stream
                        .launch_builder(&self.bias_grad_reduce_p2_kernel)
                        .arg(&partials_ptr)
                        .arg(&db_ptr)
                        .arg(&num_blocks_i32)
                        .arg(&bs_i32)
                        .arg(&inv_bq)
                        .launch(LaunchConfig {
                            grid_dim: (bs as u32, 1, 1),
                            block_dim: (256, 1, 1),
                            shared_mem_bytes: 0,
                        })
                        .map_err(|e| MLError::ModelError(format!("IQN bias_grad_reduce_p2 branch_{i}: {e}")))?;
                }
                grad_off += bs;

                // d_combined += W_bk @ dq_bk (beta=1.0 for accumulation)
                iqn_lt_matmul(
                    effective_stream, lt_handle_raw, &self.gemm_bwd_dx_branch[i],
                    1.0, 1.0, // beta=1.0 to accumulate across branches
                    branch_w_ptrs[i], dq_ptr,
                    self.d_combined_buf.raw_ptr(),
                    lt_ws_ptr, lt_ws_size,
                    &format!("iqn_bwd_dx_branch_{i}"),
                )?;

                row_off += bs;
            }
        }

        // Step 7b: Backward through hadamard_sigmoid → d_embed, d_h_s2_tiled
        unsafe {
            effective_stream
                .launch_builder(&self.hadamard_sigmoid_bwd_kernel)
                .arg(&self.h_s2_tiled_buf)
                .arg(&self.embed_buf)
                .arg(&self.d_combined_buf)
                .arg(&mut self.d_embed_buf)
                .arg(&mut self.d_h_s2_tiled_buf)
                .arg(&hbq_i32)
                .launch(elem_config(hbq))
                .map_err(|e| MLError::ModelError(format!("IQN hadamard_sigmoid_bwd: {e}")))?;
        }

        // Step 7c: Backward through ReLU
        unsafe {
            effective_stream
                .launch_builder(&self.relu_bwd_kernel)
                .arg(&self.embed_pre_buf)
                .arg(&self.d_embed_buf)
                .arg(&mut self.d_embed_pre_buf)
                .arg(&hbq_i32)
                .launch(elem_config(hbq))
                .map_err(|e| MLError::ModelError(format!("IQN relu_bwd: {e}")))?;
        }

        // Step 7d: dW_embed = d_embed_pre @ cos_features_tiled^T (alpha=inv_bq)
        let dw_embed_ptr = self.grad_buf.raw_ptr(); // offset 0 in grad_buf
        iqn_lt_matmul(
            effective_stream, lt_handle_raw, &self.gemm_bwd_dw_embed,
            inv_bq, 0.0,
            self.d_embed_pre_buf.raw_ptr(), self.cos_features_tiled_buf.raw_ptr(),
            dw_embed_ptr,
            lt_ws_ptr, lt_ws_size, "iqn_bwd_dw_embed",
        )?;

        // Step 7e: db_embed: 2-phase bias gradient reduce
        // Phase 1: grid(ceil(BQ/256), h, 1), block(256) — partial sums per block
        // Phase 2: grid(h, 1, 1), block(256) — final reduce + scale
        let db_embed_ptr = self.grad_buf.raw_ptr() + (h * d) as u64 * f32_sz;
        let bq_i32 = bq as i32;
        let p1_blocks = self.bias_grad_num_blocks as u32;
        let num_blocks_i32 = self.bias_grad_num_blocks as i32;
        let partials_ptr = self.bias_grad_partials_buf.raw_ptr();
        unsafe {
            effective_stream
                .launch_builder(&self.bias_grad_reduce_p1_kernel)
                .arg(&self.d_embed_pre_buf.raw_ptr())
                .arg(&partials_ptr)
                .arg(&h_i32)
                .arg(&bq_i32)
                .launch(LaunchConfig {
                    grid_dim: (p1_blocks, h as u32, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 256 * 4,
                })
                .map_err(|e| MLError::ModelError(format!("IQN bias_grad_reduce_p1 embed: {e}")))?;
            effective_stream
                .launch_builder(&self.bias_grad_reduce_p2_kernel)
                .arg(&partials_ptr)
                .arg(&db_embed_ptr)
                .arg(&num_blocks_i32)
                .arg(&h_i32)
                .arg(&inv_bq)
                .launch(LaunchConfig {
                    grid_dim: (h as u32, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("IQN bias_grad_reduce_p2 embed: {e}")))?;
        }

        // Step 7f: Reduce d_h_s2_tiled [H, B*Q] → d_h_s2 [H, B]
        unsafe {
            effective_stream
                .launch_builder(&self.d_h_s2_reduce_kernel)
                .arg(&self.d_h_s2_tiled_buf)
                .arg(&mut self.d_h_s2_buf)
                .arg(&h_i32)
                .arg(&b_i32)
                .arg(&n_i32)
                .launch(elem_config(h * b))
                .map_err(|e| MLError::ModelError(format!("IQN d_h_s2_reduce: {e}")))?;
        }

        // ── Step 8: Grad norm (two-phase deterministic) ─────────────────
        let total_params_i32 = self.total_params as i32;
        let norm_blocks = (self.total_params + 255) / 256;

        unsafe {
            effective_stream
                .launch_builder(&self.grad_norm_phase1_kernel)
                .arg(&self.grad_buf)
                .arg(&mut self.grad_norm_block_sums)
                .arg(&total_params_i32)
                .launch(LaunchConfig {
                    grid_dim: (norm_blocks as u32, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("IQN grad_norm phase1: {e}")))?;
        }
        let norm_blocks_i32 = norm_blocks as i32;
        unsafe {
            effective_stream
                .launch_builder(&self.grad_norm_phase2_kernel)
                .arg(&self.grad_norm_block_sums)
                .arg(&mut self.grad_norm_buf)
                .arg(&norm_blocks_i32)
                .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 })
                .map_err(|e| MLError::ModelError(format!("IQN grad_norm phase2: {e}")))?;
        }

        // ── Step 9: Adam update ─────────────────────────────────────────
        self.adam_step += 1;
        unsafe { *self.t_pinned = self.adam_step; }
        let adam_blocks = (self.total_params + 255) / 256;
        let adam_config = LaunchConfig {
            grid_dim: (adam_blocks as u32, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };

        let lr = self.config.lr;
        // SP5 Layer B (Pearl 4): per-group Adam β1/β2/ε threaded from
        // FusedTrainingCtx via ISV[ADAM_BETA{1,2}_BASE+3] / ISV[ADAM_EPS_BASE+3].
        let beta1 = iqn_beta1;
        let beta2 = iqn_beta2;
        let eps   = iqn_epsilon;
        // SP4 Layer B (Mech 9): AdamW weight_decay sourced from
        // ISV[WD_RATE[Iqn=3]] by the caller (`FusedTrainingCtx`) and
        // threaded through `execute_training_pipeline` alongside
        // `weight_clamp_max_abs`. `.max(0.0)` upstream enforces the
        // AdamW non-negativity contract.
        let wd = weight_decay_value;
        let mgn = self.config.max_grad_norm;
        let t_ptr = self.t_dev_ptr;

        // SP4 Task A14: Pearl C engagement-counter args. IQN params live in
        // ParamGroup::Iqn (= group 3); offset = 3 × MAX_BLOCKS_PER_ADAM. The
        // caller (`FusedTrainingCtx`) plumbs through `nan_flags_dev_ptr` +
        // `clamp_engage_dev_ptr` from the main DQN trainer's mapped-pinned
        // buffers via `set_pearl_c_buffers`.
        let nan_flags_ptr = self.pearl_c_nan_flags_ptr;
        let diag_slot: i32 = crate::cuda_pipeline::SP4_IQN_ADAM_DIAG_SLOT;
        let engage_buf_ptr = self.pearl_c_engage_buf_ptr;
        let engage_buf_offset: i32 = if engage_buf_ptr == 0 {
            crate::cuda_pipeline::SP4_ENGAGE_OFFSET_DISABLED
        } else {
            (crate::cuda_pipeline::ParamGroup::Iqn.idx()
                * crate::cuda_pipeline::MAX_BLOCKS_PER_ADAM) as i32
        };

        unsafe {
            effective_stream
                .launch_builder(&self.adam_kernel)
                .arg(&mut self.online_params)
                .arg(&mut self.grad_buf)
                .arg(&mut self.m_buf)
                .arg(&mut self.v_buf)
                .arg(&self.grad_norm_buf)
                .arg(&lr)
                .arg(&beta1)
                .arg(&beta2)
                .arg(&eps)
                .arg(&wd)
                .arg(&mgn)
                .arg(&t_ptr)
                .arg(&total_params_i32)
                .arg(&shared_h1_i32)
                .arg(&hidden_dim_i32)
                .arg(&embed_dim_i32)
                .arg(&b0_i32)
                .arg(&b1_i32)
                .arg(&b2_i32)
                .arg(&b3_i32)
                .arg(&weight_clamp_max_abs) // SP3 Mech 9 (close-out v2): post-Adam |p| clamp
                .arg(&nan_flags_ptr)        // SP4 A14: Pearl C sticky engage flag
                .arg(&diag_slot)            // SP4 A14: nan_flags_buf slot
                .arg(&engage_buf_ptr)       // SP4 A14: per-block engage buffer
                .arg(&engage_buf_offset)    // SP4 A14: group offset (Iqn=3)
                .launch(adam_config)
                .map_err(|e| MLError::ModelError(format!("IQN adam kernel: {e}")))?;
        }

        // Loss stays on GPU — no per-step readback.
        Ok(0.0)
    }

    /// Hard-copy online IQN params into target params via DtoD.
    ///
    /// Mirrors `GpuDqnTrainer::sync_target_from_online`: at fold boundaries the
    /// main DQN trunk is shrink-and-perturb'd and Adam reset, but the IQN
    /// target buffer still holds end-of-prior-fold weights. Polyak EMA at
    /// τ=0.005/step is far too slow to close that gap before the inflated
    /// Bellman TD error compounds through the IQN backward pass — this drives
    /// the geometric Q-drift observed in fold-1 (Q ≈ 0.8 → 2.2 → 4.0 → 10.7
    /// → fp32 overflow past atom support). Hard-syncing target ← online at
    /// the same fold boundary as the main DQN closes the gap in a single
    /// DtoD copy. No CPU staging — pure on-GPU.
    pub fn sync_target_from_online(&mut self) -> Result<(), MLError> {
        let n_bytes = self.total_params * std::mem::size_of::<f32>();
        let src = self.online_params.raw_ptr();
        let dst = self.target_params.raw_ptr();
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(dst, src, n_bytes, self.stream.cu_stream())
                .map_err(|e| MLError::ModelError(format!("IQN sync_target_from_online: {e}")))?;
        }
        Ok(())
    }

    /// Test-only borrow of the online IQN parameter buffer.
    ///
    /// Only compiled under `cfg(test)` so this does not widen the public
    /// API. Used by the GPU runtime test for `sync_target_from_online`
    /// (see `tests::iqn_sync_target_from_online_makes_target_equal_online`)
    /// to issue DtoD reads through `MappedF32Buffer`'s `dev_ptr` and
    /// verify the buffer contents post-sync.
    #[cfg(test)]
    pub(crate) fn online_params_slice(&self) -> &CudaSlice<f32> {
        &self.online_params
    }

    /// Test-only borrow of the target IQN parameter buffer.
    /// See `online_params_slice` for the rationale.
    #[cfg(test)]
    pub(crate) fn target_params_slice(&self) -> &CudaSlice<f32> {
        &self.target_params
    }

    /// Test-only accessor for the parameter count (`config.total_params()`
    /// cached at construction).
    #[cfg(test)]
    pub(crate) fn total_params_for_test(&self) -> usize {
        self.total_params
    }

    /// SP4 Task A15: parameter count accessor for the host-side Pearl C
    /// engagement-rate check. Returns `config.total_params()` cached at
    /// construction.
    pub fn total_params(&self) -> usize {
        self.total_params
    }

    /// Test-only borrow of the owning CUDA stream — the same stream every
    /// async DtoD copy on this head is queued against. The test issues its
    /// fills and read-backs on this stream, then synchronises before
    /// asserting on the mapped-pinned host pointer.
    #[cfg(test)]
    pub(crate) fn stream_for_test(&self) -> &Arc<CudaStream> {
        &self.stream
    }

    /// Update target IQN weights via Polyak EMA.
    ///
    /// `target[i] = (1 - tau) * target[i] + tau * online[i]`
    pub fn target_ema_update(&mut self, tau: f32) -> Result<(), MLError> {
        let n = self.total_params;
        // Write tau directly into the pinned+device-mapped page.  The GPU EMA kernel
        // reads via tau_dev_ptr which aliases the same physical page — no PCIe copy.
        unsafe { *self.tau_pinned = tau; }
        let blocks = (n + 255) / 256;
        let config = LaunchConfig {
            grid_dim: (blocks as u32, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };

        let n_i32 = n as i32;
        let shared_h1_i32 = self.config.shared_h1 as i32;
        let hidden_dim_i32 = self.config.hidden_dim as i32;
        let embed_dim_i32 = self.config.embed_dim as i32;
        let b0_i32 = self.config.branch_0_size as i32;
        let b1_i32 = self.config.branch_1_size as i32;
        let b2_i32 = self.config.branch_2_size as i32;
        let b3_i32 = self.config.branch_3_size as i32;
        let tau_ptr = self.tau_dev_ptr;

        unsafe {
            self.stream
                .launch_builder(&self.ema_kernel)
                .arg(&mut self.target_params)
                .arg(&self.online_params)
                .arg(&tau_ptr)
                .arg(&n_i32)
                .arg(&shared_h1_i32)
                .arg(&hidden_dim_i32)
                .arg(&embed_dim_i32)
                .arg(&b0_i32)
                .arg(&b1_i32)
                .arg(&b2_i32)
                .arg(&b3_i32)
                .launch(config)
                .map_err(|e| MLError::ModelError(format!("IQN EMA kernel: {e}")))?;
        }

        Ok(())
    }

    /// Current Adam step count.
    pub fn adam_step(&self) -> i32 {
        self.adam_step
    }

    /// Increment the host-side adam step counter without launching kernels.
    /// Used before graph_aux replay so the captured HtoD copies the updated value.
    pub fn increment_adam_step(&mut self) {
        self.adam_step += 1;
    }

    /// Zero the Adam first- and second-moment buffers and the step counter.
    /// Called at fold boundaries to stop stale fold-0 momentum from driving
    /// fold-1 parameter updates into a divergent regime — the fold-1 reward
    /// distribution differs from fold-0, so reusing the previous fold's
    /// m/v estimates produces oversized updates in the first training steps,
    /// which then snowball through the loss kernel's backward pass.
    pub fn reset_adam_state(&mut self) -> Result<(), MLError> {
        self.stream.memset_zeros(&mut self.m_buf)
            .map_err(|e| MLError::ModelError(format!("reset iqn m_buf: {e}")))?;
        self.stream.memset_zeros(&mut self.v_buf)
            .map_err(|e| MLError::ModelError(format!("reset iqn v_buf: {e}")))?;
        self.adam_step = 0;
        // t_pinned is device-mapped pinned host memory; writing the host slot
        // propagates immediately to the GPU via the pinned mapping, matching
        // the pattern used elsewhere in this module.
        unsafe { *self.t_pinned = 0; }
        Ok(())
    }

    /// Device pointer for the Adam step counter (pinned device-mapped).
    /// Used by `submit_counter_increments` to atomically increment all counters on GPU.
    pub fn t_dev_ptr(&self) -> u64 {
        self.t_dev_ptr
    }

    /// SP3 slot 39: raw device pointer to the IQN Adam first-moment buffer.
    /// Used by the SP3 Mech 5 fused threshold-check kernel (Task B6) to inspect
    /// IQN Adam state alongside the main DQN Adam buffers (slots 36-38, 40-42).
    /// Length matches `total_params + cublas_pad` from `Self::new` — see
    /// `adam_m_len` for the elementcount.
    pub fn adam_m_ptr(&self) -> u64 {
        self.m_buf.raw_ptr()
    }

    /// Element count of the IQN Adam first-moment buffer.
    pub fn adam_m_len(&self) -> usize {
        self.m_buf.len()
    }

    /// SP3 slot 43: raw device pointer to the IQN Adam second-moment buffer.
    /// Companion to `adam_m_ptr` — same length and layout, sharded `m`/`v`
    /// pair feeding the fused threshold-check kernel.
    pub fn adam_v_ptr(&self) -> u64 {
        self.v_buf.raw_ptr()
    }

    /// Element count of the IQN Adam second-moment buffer.
    pub fn adam_v_len(&self) -> usize {
        self.v_buf.len()
    }

    /// SP3 close-out v2 / slot 48: raw device pointer to the IQN online
    /// parameter buffer. Companion to `adam_m_ptr` / `adam_v_ptr` —
    /// captures the IQN-internal weights that drive the slot-26 cuBLAS
    /// GEMM (`apply_iqn_trunk_gradient`). Wired into the fused
    /// threshold-check kernel as the diagnostic blind spot identified in
    /// the F1 step-3540 NaN trace: slots 44-45 only cover the main DQN
    /// trunk + heads, leaving IQN's separate param buffer unmonitored.
    pub fn online_params_ptr(&self) -> u64 {
        self.online_params.raw_ptr()
    }

    /// Element count of the IQN online parameter buffer.
    /// Length = `total_params + cublas_pad` (last tensor can be overread
    /// by 32-element cuBLAS tiles — see `init_iqn_xavier_weights`).
    pub fn online_params_len(&self) -> usize {
        self.online_params.len()
    }

    /// SP4 Task A7 fix-up: raw device pointer to the IQN reduced gradient
    /// buffer. Companion to `online_params_ptr`/`adam_m_ptr`/`adam_v_ptr`,
    /// completing the four-pointer set required by the param-group oracle
    /// kernel (Pass D `Σ w·g / Σ w²` weight-decay rate observation). Length
    /// matches `online_params_len` (Adam state + grads mirror the params
    /// shape including the cuBLAS overread pad).
    pub fn online_grad_ptr(&self) -> u64 {
        self.grad_buf.raw_ptr()
    }

    /// Element count of the IQN reduced gradient buffer. Mirrors
    /// `online_params_len` — the gradient is allocated alongside the params
    /// in `Self::new` with the same `total_params + cublas_pad` count.
    pub fn online_grad_len(&self) -> usize {
        self.grad_buf.len()
    }

    /// Write tau into the pinned+device-mapped page for graph replay.
    /// The EMA kernel reads via `tau_dev_ptr` which aliases the same physical page.
    pub fn set_tau_host(&mut self, tau: f32) {
        unsafe { *self.tau_pinned = tau; }
    }

    /// Per-sample IQN loss buffer reference (for PER weighting).
    pub fn per_sample_loss(&self) -> &CudaSlice<f32> {
        &self.per_sample_loss
    }

    /// Read IQN total loss from pinned device-mapped memory (zero-copy, no sync).
    /// Returns the PREVIOUS step's value (one-step lag due to async pipeline).
    pub fn read_total_loss(&self) -> Result<f32, MLError> {
        Ok(unsafe { *self.total_loss_pinned })
    }

    /// IQN gradient w.r.t. shared trunk h_s2 [B, hidden_dim].
    /// After backward, this contains dL_iqn/d(h_s2) — add to C51's trunk gradient
    /// so the trunk receives bounded IQN Huber gradients alongside C51.
    pub fn d_h_s2(&self) -> &CudaSlice<f32> {
        &self.d_h_s2_buf
    }

    /// Raw device pointer to the pinned+device-mapped tau scalar.
    /// The GPU EMA kernel reads tau through this pointer on the per-step hot path.
    pub fn tau_buf_ptr(&self) -> u64 { self.tau_dev_ptr }

    pub fn d_h_s2_raw_ptr(&self) -> u64 {
        self.d_h_s2_buf.raw_ptr()
    }

    /// SP1 Phase B (slot 28): production IQN backward output buffer
    /// `d_branch_logits_buf [TBA, B*Q]` (f32) — written by
    /// `iqn_quantile_huber_loss` (`iqn_dual_head_kernel.cu:1346-1413`,
    /// `d_q_online[idx] = qw * d_huber / Q` with no inline `isfinite` guard).
    /// Wired into the NaN-check infrastructure by `run_nan_checks_post_backward`
    /// (Task 4) via the `FusedTrainingCtx` call site (the IQN head is owned by
    /// `FusedTrainingCtx`, not `GpuDqnTrainer`, so the pointer is forwarded as a
    /// `u64` argument to the trainer-side check method — same pattern as
    /// `apply_iqn_trunk_gradient(iqn_d_h_s2_ptr: u64, ...)` at
    /// `gpu_dqn_trainer.rs:6843`). Per-slot semantics in
    /// `docs/dqn-backward-nan-audit.md` per-slot accessor table (slot 28 row).
    /// NOTE: distinct from the kernel-internal `save_dL_dq` name used inside
    /// the dormant `iqn_backward_per_sample` path (not loaded in production).
    pub(crate) fn d_branch_logits_buf_ptr(&self) -> u64 {
        self.d_branch_logits_buf.raw_ptr()
    }

    /// Read IQN loss from GPU (call sparingly — synchronizes the stream).
    /// Use for epoch-end logging, NOT per-step monitoring.
    pub fn read_loss(&self) -> Result<f32, MLError> {
        let host = self.stream.clone_dtoh(&self.total_loss)
            .map_err(|e| MLError::ModelError(format!("IQN read_loss dtoh: {e}")))?;
        Ok(host[0])
    }

    /// Total branch actions (sum across all 4 branches).
    fn total_branch_actions(&self) -> usize {
        self.config.branch_0_size + self.config.branch_1_size
            + self.config.branch_2_size + self.config.branch_3_size
    }

    /// Compute IQR per action from the last forward pass's quantile estimates.
    ///
    /// IQR = quantile_75 − quantile_25 for each action, drawn from the last
    /// sample in the training batch. Plan 4 Task 3 (E.3): with `FIXED_TAUS`
    /// = `[0.05, 0.25, 0.50, 0.75, 0.95]`, the 25th and 75th percentiles
    /// live at slot indices 1 and 3 respectively (= `n_q / 4` and
    /// `(3*n_q) / 4` for `n_q = 5` — the integer-truncation arithmetic
    /// happens to land on the right slots, kept for parity with the prior
    /// 32-quantile QR-DQN-midpoint behaviour).
    ///
    /// Call after `train_iqn_step_gpu()` which populates `save_q_online`.
    /// This is a cold-path operation (epoch boundary) -- host-side readback is acceptable.
    pub fn compute_iqr(&mut self) -> Result<(), MLError> {
        let n_q = self.config.num_quantiles;
        let n_a = self.total_branch_actions();

        // Plan 4 Task 3 (E.3): with FIXED_TAUS = [0.05, 0.25, 0.50, 0.75, 0.95],
        // idx_25 = 1 and idx_75 = 3. Computed as `n_q/4` / `3*n_q/4` so the
        // formula remains valid if FIXED_TAUS were ever extended to a longer
        // grid that still placed Q25/Q75 at the canonical positions; with
        // n_q=5 these evaluate to 1 and 3 by integer truncation.
        let idx_25 = n_q / 4;
        let idx_75 = (3 * n_q) / 4;
        debug_assert!(idx_25 < n_q && idx_75 < n_q && idx_25 < idx_75,
            "IQR indices out of order: idx_25={idx_25}, idx_75={idx_75}, n_q={n_q}");

        // save_q_online is [TBA, B*Q] col-major.
        // For last sample (b=B-1), quantile q: column = (B-1)*Q + q, row = action
        // Index = row + col * TBA = action + ((B-1)*Q + q) * TBA
        let last_sample_col_start = (self.config.batch_size - 1) * n_q;

        // Read all quantiles for the last sample's actions
        let mut host_q = vec![0.0_f32; n_q * n_a];
        // We need columns [last_sample_col_start .. last_sample_col_start + n_q]
        // Each column has TBA elements. Total = n_q * TBA floats.
        let src_ptr = self.save_q_online.raw_ptr()
            + (last_sample_col_start * n_a * std::mem::size_of::<f32>()) as u64;
        unsafe {
            let r = cudarc::driver::sys::cuMemcpyDtoH_v2(
                host_q.as_mut_ptr() as *mut std::ffi::c_void,
                src_ptr,
                (n_q * n_a * std::mem::size_of::<f32>()) as usize,
            );
            if r != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
                return Err(MLError::ModelError(format!("IQR DtoH failed: {r:?}")));
            }
        }

        // host_q layout: col-major [TBA, n_q] = host_q[action + q * TBA]
        let mut iqr = vec![0.0_f32; n_a];
        for a in 0..n_a {
            let q25 = host_q[a + idx_25 * n_a];
            let q75 = host_q[a + idx_75 * n_a];
            iqr[a] = (q75 - q25).abs();
        }

        // Upload IQR to GPU
        unsafe {
            let r = cudarc::driver::sys::cuMemcpyHtoDAsync_v2(
                self.iqr_buf.raw_ptr(),
                iqr.as_ptr() as *const std::ffi::c_void,
                n_a * std::mem::size_of::<f32>(),
                self.stream.cu_stream(),
            );
            if r != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
                return Err(MLError::ModelError(format!("IQR HtoD failed: {r:?}")));
            }
        }

        Ok(())
    }

    /// Raw device pointer to IQR buffer for passing to action selector.
    pub fn iqr_buf_ptr(&self) -> u64 {
        self.iqr_buf.raw_ptr()
    }

    /// Raw device pointer to the per-quantile online Q surface
    /// `save_q_online [TBA, B*Q]` (col-major) populated by
    /// `execute_training_pipeline`. Plan 4 Task 3 (E.3) consumer:
    /// `iqn_quantile_ema_update` GPU kernel, which reduces (B × TBA)
    /// per off-median quantile into ISV[99..103). The buffer's lifetime
    /// matches `GpuIqnHead`'s; the pointer is stable across training steps.
    pub fn save_q_online_ptr(&self) -> u64 {
        self.save_q_online.raw_ptr()
    }

    /// Total branch actions (TBA) — sum across all 4 branches. Public
    /// accessor for Plan 4 Task 3 (E.3) callers that need to size the
    /// `iqn_quantile_ema_update` reduction without re-reading the config
    /// (which is private to this struct).
    pub fn tba(&self) -> usize {
        self.total_branch_actions()
    }

    /// Number of quantiles (= `FIXED_TAUS.len()` = 5 under Plan 4 Task 3).
    pub fn num_quantiles(&self) -> usize {
        self.config.num_quantiles
    }

    /// SP5 Pearl 5: Update τ schedule from ISV[IQN_TAU_BASE=250..270).
    ///
    /// Reads 4 branches × 5 quantiles from ISV, averages across branches per
    /// quantile slot, applies Invariant-1 cold-start floors (if ISV reads 0,
    /// fall back to `FIXED_TAUS[q]`), then writes the averaged τ schedule and
    /// recomputed cosine features to all 4 per-branch mapped-pinned slabs via
    /// `write_from_slice` (direct host_ptr write — no DtoD copy).
    ///
    /// Must be called once per epoch from `FusedTrainingCtx` before the IQN
    /// training pipeline runs.
    pub fn refresh_taus_from_isv(
        &mut self,
        isv_per_branch_taus: &[f32; 20], // ISV[250..270): branch b, quantile q at [b*5+q]
    ) {
        let n = self.config.num_quantiles;   // 5
        let b = self.config.batch_size;      // batch dimension for tiling
        let d = self.config.embed_dim;       // 64

        // Average 4 branches per quantile, apply cold-start floor from FIXED_TAUS.
        let mut taus = [0.0_f32; 5];
        for q in 0..n {
            let mean = (isv_per_branch_taus[q]
                + isv_per_branch_taus[5 + q]
                + isv_per_branch_taus[10 + q]
                + isv_per_branch_taus[15 + q]) / 4.0_f32;
            // Invariant-1 cold-start floor: if ISV hasn't been populated yet
            // (all four branches read 0.0), retain FIXED_TAUS[q] to avoid
            // collapsing the quantile distribution to 0.
            taus[q] = if mean < 1e-6_f32 { FIXED_TAUS[q] } else { mean.min(1.0_f32) };
        }

        // Recompute cosine features [D, N]: cos_feat[dim + q*D] = cos(π·(d+1)·τ_q)
        let mut cos_feat_host = vec![0.0_f32; d * n];
        for q in 0..n {
            let tau_q = taus[q];
            for dim in 0..d {
                cos_feat_host[dim + q * d] =
                    (std::f32::consts::PI * ((dim + 1) as f32) * tau_q).cos();
            }
        }

        // Tile taus [B, N] — each row is the same N quantile values repeated B times.
        let mut tiled = Vec::with_capacity(b * n);
        for _ in 0..b {
            tiled.extend_from_slice(&taus[..n]);
        }

        // Write the averaged tau schedule to all 4 per-branch mapped-pinned slabs.
        // Direct host_ptr writes — no DtoD copy, no upload. The kernel reads via
        // dev_ptr of the same mapped-pinned allocation after the next stream sync
        // (feedback_no_htod_htoh_only_mapped_pinned).
        for b_idx in 0..4 {
            self.online_taus_per_branch[b_idx].write_from_slice(&tiled);
            self.target_taus_per_branch[b_idx].write_from_slice(&tiled);
            self.cos_features_per_branch[b_idx].write_from_slice(&cos_feat_host);
        }
    }

    /// SP6 Pearl 5: Refresh the τ tensors for a single branch from its 5-quantile
    /// schedule `tau5` (from ISV[IQN_TAU_BASE + branch_idx*5 .. +5]).
    ///
    /// Writes directly to `online_taus_per_branch[branch_idx]`,
    /// `target_taus_per_branch[branch_idx]`, and `cos_features_per_branch[branch_idx]`
    /// via `write_from_slice` (mapped-pinned host_ptr — no DtoD copy).
    /// Called 4 times per epoch refresh (once per branch) from
    /// `FusedTrainingCtx::run_full_step`.
    ///
    /// Cold-start floor: if a tau value is below 1e-6 (ISV not yet populated),
    /// uses `FIXED_TAUS[q]` to preserve the interpretable quantile distribution.
    pub fn refresh_taus_for_branch(
        &mut self,
        branch_idx: usize,
        tau5: &[f32; 5],
    ) {
        debug_assert!(branch_idx < 4, "refresh_taus_for_branch: branch_idx {branch_idx} >= 4");
        let n = self.config.num_quantiles;   // 5
        let b = self.config.batch_size;
        let d = self.config.embed_dim;

        // Apply cold-start floor: if ISV slot is zero (not yet populated),
        // fall back to FIXED_TAUS[q] so we don't collapse the quantile distribution.
        let mut taus = [0.0_f32; 5];
        for q in 0..n {
            taus[q] = if tau5[q] < 1e-6_f32 { FIXED_TAUS[q] } else { tau5[q].min(1.0_f32) };
        }

        // Recompute cosine features [D, N]: cos_feat[dim + q*D] = cos(π·(d+1)·τ_q)
        let mut cos_feat_host = vec![0.0_f32; d * n];
        for q in 0..n {
            let tau_q = taus[q];
            for dim in 0..d {
                cos_feat_host[dim + q * d] =
                    (std::f32::consts::PI * ((dim + 1) as f32) * tau_q).cos();
            }
        }

        // Tile taus [B, N].
        let mut tiled = Vec::with_capacity(b * n);
        for _ in 0..b {
            tiled.extend_from_slice(&taus[..n]);
        }

        // Write directly to mapped-pinned host_ptr — no DtoD copy, no upload.
        // The IQN kernel reads via dev_ptr of the same mapped-pinned allocation
        // after the stream sync that precedes kernel launch
        // (feedback_no_htod_htoh_only_mapped_pinned).
        self.online_taus_per_branch[branch_idx].write_from_slice(&tiled);
        self.target_taus_per_branch[branch_idx].write_from_slice(&tiled);
        self.cos_features_per_branch[branch_idx].write_from_slice(&cos_feat_host);
    }

    /// SP6 Pearl 5: Set `active_branch_idx` so all subsequent IQN kernel launches
    /// use `branch_idx`'s per-branch mapped-pinned τ / cos_features slabs.
    ///
    /// Must be paired with `deactivate_branch_taus(branch_idx)` after the forward pass.
    pub fn activate_branch_taus(&mut self, branch_idx: usize) {
        debug_assert!(branch_idx < 4, "activate_branch_taus: branch_idx {branch_idx} >= 4");
        self.active_branch_idx = branch_idx;
    }

    /// SP6 Pearl 5: Reset `active_branch_idx` to 0 after a per-branch IQN forward pass.
    ///
    /// Symmetric inverse of `activate_branch_taus(branch_idx)`.
    pub fn deactivate_branch_taus(&mut self, branch_idx: usize) {
        debug_assert!(branch_idx < 4, "deactivate_branch_taus: branch_idx {branch_idx} >= 4");
        self.active_branch_idx = 0;
    }

    /// Compute CVaR-based position scaling from IQN quantiles.
    ///
    /// Runs the IQN forward-only kernel on `h_s2` (trunk activation),
    /// then computes CVaR at `alpha` (e.g. 0.05 = 5th percentile) for
    /// each sample's selected exposure action.
    ///
    /// Returns a GPU buffer `[batch_size]` with scaling factors in [0.25, 1.0].
    pub fn compute_cvar_scales(
        &mut self,
        h_s2: &CudaSlice<f32>,
        actions: &CudaSlice<i32>,
        batch_size: usize,
        alpha: f32,
    ) -> Result<CudaSlice<f32>, MLError> {
        let b = batch_size;
        let n = self.config.num_quantiles;
        let tba = self.config.total_branch_actions();
        let shared_h1_i32 = self.config.shared_h1 as i32;
        let hidden_dim_i32 = self.config.hidden_dim as i32;
        let embed_dim_i32 = self.config.embed_dim as i32;

        // Plan 4 Task 3 (E.3): per-step Philox τ sampling deleted. The
        // constructor populated all per-branch mapped-pinned τ slabs with
        // FIXED_TAUS broadcast across all B rows; CVaR's quantile estimates
        // are drawn from the same fixed quantile levels {0.05, 0.25, 0.50,
        // 0.75, 0.95} as the training pipeline, making the CVaR estimate
        // stable across steps and consistent with the spec's
        // interpretable-quantile contract.
        let b0_i32 = self.config.branch_0_size as i32;
        let b1_i32 = self.config.branch_1_size as i32;
        let b2_i32 = self.config.branch_2_size as i32;
        let b3_i32 = self.config.branch_3_size as i32;

        // 2. Run IQN forward-only kernel (256 threads per sample)
        let fwd_inf_shmem = (256 / 32) * 4;
        let fwd_config = LaunchConfig {
            grid_dim: (b as u32, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: fwd_inf_shmem as u32,
        };
        let batch_i32 = b as i32;
        unsafe {
            self.stream
                .launch_builder(&self.forward_kernel)
                .arg(h_s2)
                .arg(&self.online_taus_per_branch[self.active_branch_idx].dev_ptr)
                .arg(&self.online_params)
                .arg(&self.cos_features_per_branch[self.active_branch_idx].dev_ptr)
                .arg(&mut self.save_q_online)
                .arg(&batch_i32)
                .arg(&shared_h1_i32)
                .arg(&hidden_dim_i32)
                .arg(&embed_dim_i32)
                .arg(&b0_i32)
                .arg(&b1_i32)
                .arg(&b2_i32)
                .arg(&b3_i32)
                .launch(fwd_config)
                .map_err(|e| MLError::ModelError(format!("IQN CVaR forward: {e}")))?;
        }

        // 3. Compute CVaR on GPU — ZERO CPU readback.
        let mut scales_buf = self.stream.alloc_zeros::<f32>(b)
            .map_err(|e| MLError::ModelError(format!("IQN CVaR alloc scales: {e}")))?;
        let alpha_f = alpha;
        let n_i32 = n as i32;
        let tba_i32 = tba as i32;
        let b0_i32 = self.config.branch_0_size as i32;
        let b1b2 = (self.config.branch_1_size * self.config.branch_2_size) as i32;
        let cvar_blocks = (b + 255) / 256;
        let cvar_config = LaunchConfig {
            grid_dim: (cvar_blocks as u32, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };
        // Load from precompiled cubin (ZERO runtime nvcc)
        static CVAR_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/iqn_cvar_kernel.cubin"));
        let context = self.stream.context();
        let module = context.load_cubin(CVAR_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("CVaR cubin load: {e}")))?;
        let cvar_kernel = module.load_function("iqn_cvar_kernel")
            .map_err(|e| MLError::ModelError(format!("CVaR kernel load: {e}")))?;

        unsafe {
            self.stream
                .launch_builder(&cvar_kernel)
                .arg(&self.save_q_online)
                .arg(actions)
                .arg(&mut scales_buf)
                .arg(&batch_i32)
                .arg(&n_i32)
                .arg(&tba_i32)
                .arg(&b0_i32)
                .arg(&b1b2)
                .arg(&alpha_f)
                .launch(cvar_config)
                .map_err(|e| MLError::ModelError(format!("CVaR kernel launch: {e}")))?;
        }

        Ok(scales_buf)
    }
}

// ---------------------------------------------------------------------------
// cuBLAS GEMM helpers
// ---------------------------------------------------------------------------

/// Execute a single cublasLtMatmul with pre-cached IQN GEMM descriptors.
fn iqn_lt_matmul(
    stream: &CudaStream,
    lt_handle: cublaslt_sys::cublasLtHandle_t,
    desc: &IqnGemmDesc,
    alpha: f32,
    beta: f32,
    a_ptr: u64,
    b_ptr: u64,
    c_ptr: u64,
    ws_ptr: u64,
    ws_size: usize,
    label: &str,
) -> Result<(), MLError> {
    unsafe {
        let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t;
        let status = cublaslt_sys::cublasLtMatmul(
            lt_handle,
            desc.matmul_desc,
            &alpha as *const f32 as *const std::ffi::c_void,
            a_ptr as *const std::ffi::c_void,
            desc.a_layout,
            b_ptr as *const std::ffi::c_void,
            desc.b_layout,
            &beta as *const f32 as *const std::ffi::c_void,
            c_ptr as *const std::ffi::c_void,
            desc.c_layout,
            c_ptr as *mut std::ffi::c_void,
            desc.d_layout,
            &desc.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t,
            ws_ptr as *mut std::ffi::c_void,
            ws_size,
            cu_stream,
        );
        if status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
            let e = cublaslt_result::CublasError(status);
            return Err(MLError::ModelError(format!("cublasLtMatmul IQN {label}: {e:?}")));
        }
    }
    Ok(())
}

/// Create a cached `IqnGemmDesc` for a given GEMM shape.
///
/// Parameters follow cuBLAS column-major convention:
/// - `transa`: 0=OP_N, 1=OP_T
/// - `transb`: 0=OP_N, 1=OP_T
/// - `m, n, k`: logical output dimensions (C is m x n)
/// - `lda, ldb, ldc`: leading dimensions of physical storage
#[allow(clippy::too_many_arguments)]
fn create_iqn_gemm_desc(
    lt_handle: cublaslt_sys::cublasLtHandle_t,
    transa: i32,
    transb: i32,
    m: usize,
    n: usize,
    k: usize,
    lda: usize,
    ldb: usize,
    ldc: usize,
    ws_size: usize,
    label: &str,
) -> Result<IqnGemmDesc, MLError> {
    let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
    let compute_type = cublaslt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32;

    unsafe {
        let matmul_desc = cublaslt_result::create_matmul_desc(compute_type, f32_type)
            .map_err(|e| MLError::ModelError(format!("IQN MatmulDescCreate {label}: {e:?}")))?;

        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSA,
            &transa as *const i32 as *const std::ffi::c_void,
            std::mem::size_of::<i32>(),
        ).map_err(|e| MLError::ModelError(format!("IQN set TRANSA {label}: {e:?}")))?;

        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSB,
            &transb as *const i32 as *const std::ffi::c_void,
            std::mem::size_of::<i32>(),
        ).map_err(|e| MLError::ModelError(format!("IQN set TRANSB {label}: {e:?}")))?;

        // Matrix layouts — physical storage dimensions depend on transpose flags.
        let transa_is_n = transa == 0;
        let transb_is_n = transb == 0;

        let a_layout = if transa_is_n {
            // A stored as [M, K] with ld=lda, OP_N → logical [M, K]
            cublaslt_result::create_matrix_layout(f32_type, m as u64, k as u64, lda as i64)
        } else {
            // A stored as [K, M] with ld=lda, OP_T → logical [M, K]
            cublaslt_result::create_matrix_layout(f32_type, k as u64, m as u64, lda as i64)
        }.map_err(|e| MLError::ModelError(format!("IQN A layout {label}: {e:?}")))?;

        let b_layout = if transb_is_n {
            // B stored as [K, N] with ld=ldb, OP_N → logical [K, N]
            cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, ldb as i64)
        } else {
            // B stored as [N, K] with ld=ldb, OP_T → logical [K, N]
            cublaslt_result::create_matrix_layout(f32_type, n as u64, k as u64, ldb as i64)
        }.map_err(|e| MLError::ModelError(format!("IQN B layout {label}: {e:?}")))?;

        let c_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, ldc as i64)
            .map_err(|e| MLError::ModelError(format!("IQN C layout {label}: {e:?}")))?;

        let d_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, ldc as i64)
            .map_err(|e| MLError::ModelError(format!("IQN D layout {label}: {e:?}")))?;

        // ── Deterministic algorithm selection (Option B fix) ──
        // Replaces timing-based `cublasLtMatmulAlgoGetHeuristic` with a
        // hardware-stable AlgoGetIds → AlgoInit → AlgoCheck loop. TF32
        // compute type preserved.
        let shape = super::cublas_algo_deterministic::ShapeKey::new(
            transa, transb,
            m as i32, n as i32, k as i32,
            lda as i32, ldb as i32, ldc as i32,
            ws_size,
        );
        let heuristic = super::cublas_algo_deterministic::get_matmul_algo_f32_tf32(
            lt_handle, matmul_desc,
            a_layout, b_layout, c_layout, d_layout,
            shape,
        );

        let heuristic = match heuristic {
            Ok(h) => h,
            Err(e) => {
                let _ = cublaslt_result::destroy_matrix_layout(d_layout);
                let _ = cublaslt_result::destroy_matrix_layout(c_layout);
                let _ = cublaslt_result::destroy_matrix_layout(b_layout);
                let _ = cublaslt_result::destroy_matrix_layout(a_layout);
                let _ = cublaslt_result::destroy_matmul_desc(matmul_desc);
                return Err(MLError::ModelError(format!(
                    "IQN deterministic algo {label} (transa={transa},transb={transb},m={m},n={n},k={k}): {e}"
                )));
            }
        };

        let algo = heuristic.algo;
        tracing::info!(
            transa, transb, m, n, k, lda, ldb, ldc,
            ws_needed = heuristic.workspaceSize,
            "IQN GEMM desc created: {label}"
        );

        Ok(IqnGemmDesc {
            matmul_desc, a_layout, b_layout, c_layout, d_layout, algo,
        })
    }
}

// ---------------------------------------------------------------------------
// Kernel loading
// ---------------------------------------------------------------------------

struct IqnKernels {
    // Plan 4 Task 2c.3b: trunk_fwd deleted (orphaned kernel — see kernel
    // file for the deletion comment).
    // Plan 4 Task 3 (E.3): sample_taus deleted (fixed-τ broadcast replaces
    // per-step Philox sampling — see kernel file for the deletion comment).
    loss_reduce: CudaFunction,
    gnorm_p1: CudaFunction,
    gnorm_p2: CudaFunction,
    adam: CudaFunction,
    fwd_only: CudaFunction,
    ema: CudaFunction,
    decode_act: CudaFunction,
    relu_fwd: CudaFunction,
    relu_bwd: CudaFunction,
    hadamard_sigmoid: CudaFunction,
    hadamard_sigmoid_bwd: CudaFunction,
    quantile_huber_loss: CudaFunction,
    bias_add: CudaFunction,
    bias_grad_reduce_p1: CudaFunction,
    bias_grad_reduce_p2: CudaFunction,
    h_s2_tile: CudaFunction,
    d_h_s2_reduce: CudaFunction,
    cos_tile: CudaFunction,
}

/// Load all IQN CUDA kernels from the precompiled cubin.
fn load_iqn_kernels(
    stream: &Arc<CudaStream>,
    config: &GpuIqnConfig,
) -> Result<IqnKernels, MLError> {
    static IQN_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/iqn_dual_head_kernel.cubin"));

    info!(
        hidden = config.hidden_dim,
        embed_dim = config.embed_dim,
        num_quantiles = config.num_quantiles,
        kappa = config.kappa,
        total_params = config.total_params(),
        "GpuIqnHead: loading precompiled IQN kernels (cuBLAS GEMM mode)"
    );

    let context = stream.context();
    let module = context.load_cubin(IQN_CUBIN.to_vec()).map_err(|e| {
        MLError::ModelError(format!("IQN module load: {e}"))
    })?;

    let load = |name: &str| -> Result<CudaFunction, MLError> {
        module.load_function(name).map_err(|e| {
            MLError::ModelError(format!("{name} load: {e}"))
        })
    };

    // Retained kernels (not replaced by cuBLAS)
    // Plan 4 Task 2c.3b: iqn_trunk_forward_kernel deleted; the cuBLAS
    // path in `execute_training_pipeline` is now replaced with a call
    // to `BatchedForward::target_encoder_forward_only` (online + target
    // share one trunk implementation — the GRN forward).
    let loss_reduce = load("iqn_loss_reduce")?;
    let gnorm_p1 = load("iqn_grad_norm_phase1")?;
    let gnorm_p2 = load("iqn_grad_norm_phase2")?;
    let adam = load("iqn_adam_kernel")?;
    let fwd_only = load("iqn_forward_kernel")?;
    let ema = load("iqn_ema_kernel")?;
    let decode_act = load("iqn_decode_actions_kernel")?;
    // Plan 4 Task 3 (E.3): iqn_sample_taus_kernel removed from cubin — no
    // load. Fixed-τ broadcast at construction replaces per-step sampling.

    // New element-wise kernels for cuBLAS pipeline
    let relu_fwd = load("iqn_relu_fwd")?;
    let relu_bwd = load("iqn_relu_bwd")?;
    let hadamard_sigmoid = load("iqn_hadamard_sigmoid")?;
    let hadamard_sigmoid_bwd = load("iqn_hadamard_sigmoid_bwd")?;
    let quantile_huber_loss = load("iqn_quantile_huber_loss")?;
    let bias_add = load("iqn_bias_add")?;
    let bias_grad_reduce_p1 = load("iqn_bias_grad_reduce_p1")?;
    let bias_grad_reduce_p2 = load("iqn_bias_grad_reduce_p2")?;
    let h_s2_tile = load("iqn_h_s2_tile")?;
    let d_h_s2_reduce = load("iqn_d_h_s2_reduce")?;
    let cos_tile = load("iqn_cos_tile")?;

    info!(
        "GpuIqnHead: 18 kernels loaded (7 retained + 11 element-wise; \
         trunk_fwd deleted in Task 2c.3b, sample_taus deleted in Task 3 E.3)"
    );
    Ok(IqnKernels {
        loss_reduce, gnorm_p1, gnorm_p2, adam, fwd_only,
        ema, decode_act,
        relu_fwd, relu_bwd, hadamard_sigmoid, hadamard_sigmoid_bwd,
        quantile_huber_loss, bias_add, bias_grad_reduce_p1, bias_grad_reduce_p2,
        h_s2_tile, d_h_s2_reduce, cos_tile,
    })
}

// ---------------------------------------------------------------------------
// Weight initialization
// ---------------------------------------------------------------------------

/// Initialize IQN weights with Xavier/Glorot uniform distribution — native f32.
fn init_iqn_xavier_weights(
    stream: &Arc<CudaStream>,
    config: &GpuIqnConfig,
) -> Result<CudaSlice<f32>, MLError> {
    use rand::Rng;
    use rand::SeedableRng;
    use rand::rngs::StdRng;

    let total = config.total_params();
    let h = config.hidden_dim;
    let d = config.embed_dim;
    let b0 = config.branch_0_size;
    let b1 = config.branch_1_size;
    let b2 = config.branch_2_size;
    let b3 = config.branch_3_size;
    // cuBLAS tile padding: last tensor can be overread by 32-element tiles
    let cublas_pad = 32 * h;
    let mut weights = vec![0.0_f32; total + cublas_pad];
    let mut rng = StdRng::seed_from_u64(crate::cuda_pipeline::mix_seed(0x1CA_1234));
    let mut offset = 0;

    // W_embed [H, D]
    let limit = (6.0_f64 / (h + d) as f64).sqrt() as f32;
    for w in &mut weights[offset..offset + h * d] {
        *w = rng.gen_range(-limit..limit);
    }
    offset += h * d;
    // b_embed [H] — zero
    offset += h;

    // Branch output layers (all 4 branches)
    for branch_size in [b0, b1, b2, b3] {
        let limit = (6.0_f64 / (branch_size + h) as f64).sqrt() as f32;
        // W_bd [branch_size, H]
        for w in &mut weights[offset..offset + branch_size * h] {
            *w = rng.gen_range(-limit..limit);
        }
        offset += branch_size * h;
        // b_bd [branch_size] — zero
        offset += branch_size;
    }

    debug_assert_eq!(offset, total, "IQN weight layout mismatch");

    // Upload to GPU via mapped-pinned staging (zeros past total_params remain in `weights` Vec)
    let params_buf = {
        let n = weights.len();
        let staging = unsafe { super::mapped_pinned::MappedF32Buffer::new(n) }
            .map_err(|e| MLError::ModelError(format!("IQN online_params staging alloc ({} f32): {e}", n)))?;
        staging.write_from_slice(&weights);
        let mut _dst = stream.alloc_zeros::<f32>(n)
            .map_err(|e| MLError::ModelError(format!("IQN online_params alloc ({} f32): {e}", n)))?;
        if n > 0 {
            let nbytes = n * std::mem::size_of::<f32>();
            #[allow(unsafe_code)]
            unsafe {
                let (dst_ptr, _g) = _dst.device_ptr_mut(stream);
                cudarc::driver::result::memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream())
                    .map_err(|e| MLError::ModelError(format!("IQN online_params DtoD: {e}")))?;
            }
            stream.synchronize().map_err(|e| MLError::ModelError(format!("IQN online_params sync: {e}")))?;
        }
        _dst
    };

    Ok(params_buf)
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Allocate a zeroed f32 buffer on GPU.
fn alloc_f32(
    stream: &Arc<CudaStream>,
    n: usize,
    name: &str,
) -> Result<CudaSlice<f32>, MLError> {
    stream.alloc_zeros::<f32>(n).map_err(|e| {
        MLError::ModelError(format!("GpuIqn alloc {name} ({n} f32): {e}"))
    })
}

// SAFETY: GpuIqnHead is only accessed from the training thread that owns
// the CUDA context.  The raw `t_pinned`, `total_loss_pinned`, and `tau_pinned`
// pointers are each allocated via cuMemHostAlloc(DEVICEMAP) and freed in Drop
// — none are aliased.
unsafe impl Send for GpuIqnHead {}
unsafe impl Sync for GpuIqnHead {}

impl Drop for GpuIqnHead {
    fn drop(&mut self) {
        if !self.t_pinned.is_null() {
            unsafe { cudarc::driver::sys::cuMemFreeHost(self.t_pinned.cast()); }
        }
        if !self.total_loss_pinned.is_null() {
            unsafe { cudarc::driver::sys::cuMemFreeHost(self.total_loss_pinned.cast()); }
        }
        if !self.tau_pinned.is_null() {
            unsafe { cudarc::driver::sys::cuMemFreeHost(self.tau_pinned.cast()); }
        }
    }
}

/// Clone a `CudaSlice<f32>` via a single DtoD copy (no host crossing).
///
/// Replaces the legacy device→host→device round-trip that violated both
/// `feedback_gpu_cpu_roundtrip.md` (no DtoH) and
/// `feedback_no_htod_htoh_only_mapped_pinned.md` (no HtoD). Used at
/// constructor time to seed the IQN target network from the online params.
fn clone_cuda_slice(
    stream: &Arc<CudaStream>,
    src: &CudaSlice<f32>,
    n: usize,
) -> Result<CudaSlice<f32>, MLError> {
    super::clone_cuda_slice_f32(src, stream)
        .map_err(|e| MLError::ModelError(format!("IQN clone DtoD ({n} f32): {e}")))
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    //! GPU runtime regression tests for `GpuIqnHead`.
    //!
    //! All tests carry `#[ignore = "gpu"]` because they require a live CUDA
    //! device. They run on the L40S smoke validation pool — `cargo test
    //! -p ml --lib` on a CPU-only host skips them by default. Trigger
    //! manually with:
    //!     cargo test -p ml --lib gpu_iqn_head -- --ignored --nocapture
    //!
    //! ## CPU↔GPU communication discipline
    //!
    //! Per `feedback_no_htod_htoh_only_mapped_pinned.md`, the only
    //! permitted CPU↔GPU path is `cuMemHostAlloc(DEVICEMAP|PORTABLE)`
    //! mapped pinned memory. Tests are not exempt — these tests fill and
    //! read IQN parameter buffers through `MappedF32Buffer` and use
    //! `cuMemcpyDtoDAsync` between the mapped buffer's `dev_ptr` and the
    //! head's `online_params` / `target_params` `CudaSlice`s. No
    //! `memcpy_dtoh` / `memcpy_htod` anywhere on the test path.
    use super::*;
    use crate::cuda_pipeline::mapped_pinned::MappedF32Buffer;
    use crate::cuda_pipeline::shared_cublas_handle::PerStreamCublasHandles;
    use cudarc::driver::CudaContext;

    /// Fill an entire `CudaSlice<f32>` to a single scalar value via a
    /// mapped-pinned staging buffer and `cuMemcpyDtoDAsync`. No HtoD copy
    /// is issued; the host write through `MappedF32Buffer::host_ptr`
    /// reaches the GPU through the device-mapped alias, then the DtoD
    /// copies the staged values into `dst`.
    ///
    /// Caller must `stream.synchronize()` before treating `dst` as
    /// settled (or before reusing the staging buffer for another value).
    fn fill_slice_via_mapped(
        stream: &Arc<CudaStream>,
        staging: &MappedF32Buffer,
        dst: &CudaSlice<f32>,
        n: usize,
        value: f32,
    ) {
        assert_eq!(staging.len, n, "staging buffer size must match dst");
        // Host-side fill of the pinned mapping; the device sees the same
        // bytes through `staging.dev_ptr`.
        let host_slice: &mut [f32] = unsafe {
            std::slice::from_raw_parts_mut(staging.host_ptr, n)
        };
        host_slice.fill(value);
        let n_bytes = n * std::mem::size_of::<f32>();
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst.raw_ptr(),
                staging.dev_ptr,
                n_bytes,
                stream.cu_stream(),
            )
            .expect("DtoD fill via mapped pinned staging");
        }
    }

    /// Read an entire `CudaSlice<f32>` into a fresh `Vec<f32>` via a
    /// mapped-pinned destination buffer and `cuMemcpyDtoDAsync`. Returns
    /// after `stream.synchronize()` so the host read is coherent.
    fn read_slice_via_mapped(
        stream: &Arc<CudaStream>,
        src: &CudaSlice<f32>,
        n: usize,
    ) -> Vec<f32> {
        let mapped = unsafe { MappedF32Buffer::new(n) }
            .expect("MappedF32Buffer alloc for read-back");
        let n_bytes = n * std::mem::size_of::<f32>();
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                mapped.dev_ptr,
                src.raw_ptr(),
                n_bytes,
                stream.cu_stream(),
            )
            .expect("DtoD read-back into mapped pinned destination");
        }
        stream.synchronize().expect("stream sync after read-back");
        mapped.read_all()
    }

    /// GPU runtime contract test for `GpuIqnHead::sync_target_from_online`.
    ///
    /// The fold-boundary path in `fused_training.rs::reset_for_fold`
    /// depends on `sync_target_from_online` performing a real DtoD copy
    /// from `online_params` into `target_params`. Without it, the IQN
    /// target buffer carries end-of-prior-fold weights into fold N+1 and
    /// Polyak EMA at τ=0.005/step is far too slow to close the resulting
    /// Bellman-target gap before Q drifts geometrically (commit
    /// `7c19b5903`, issue #84).
    ///
    /// This test exercises the contract end-to-end on a real GPU:
    ///   1. Fill `online_params` ← 0.42 (f32) and `target_params` ← 0.99
    ///      via mapped-pinned staging + DtoD. `0.42` and `0.99` are
    ///      arbitrary witnesses (any two distinct fp32 values would do);
    ///      they are NOT tuned constants — the test asserts equality, not
    ///      anything about the magnitudes.
    ///   2. Sanity: read both buffers back via mapped pinned and assert
    ///      they differ pointwise (the staged fill landed).
    ///   3. Call `sync_target_from_online`.
    ///   4. Synchronise the stream so the queued DtoD has retired.
    ///   5. Read both buffers back via mapped pinned and assert
    ///      bit-for-bit equality across all `total_params` slots.
    ///
    /// A stub body (e.g. `Ok(())` returning without copying) or a sync
    /// against the wrong buffer / wrong direction / wrong stream all fail
    /// step 5 — which the previous static `include_str!` regression test
    /// could not detect.
    #[test]
    #[ignore = "gpu"]
    fn iqn_sync_target_from_online_makes_target_equal_online() {
        let ctx = CudaContext::new(0).expect("CUDA context — is GPU available?");
        let stream = ctx.default_stream();
        let shared = Arc::new(
            PerStreamCublasHandles::new(&stream)
                .expect("PerStreamCublasHandles::new"),
        );

        let config = GpuIqnConfig::default();
        let mut iqn = GpuIqnHead::new(Arc::clone(&shared), config)
            .expect("GpuIqnHead::new must succeed on a GPU-equipped host");
        let n = iqn.total_params_for_test();
        let head_stream = Arc::clone(iqn.stream_for_test());

        // Witness values — arbitrary distinct fp32 constants. The test
        // asserts post-sync equality, not anything about their magnitude.
        const ONLINE_FILL: f32 = 0.42;
        const TARGET_FILL: f32 = 0.99;

        // Single mapped-pinned staging buffer reused across the two
        // fills (one DtoD per fill, sync between so the staging memory
        // can be safely overwritten).
        let staging = unsafe { MappedF32Buffer::new(n) }
            .expect("MappedF32Buffer alloc for fills");

        fill_slice_via_mapped(&head_stream, &staging, iqn.online_params_slice(), n, ONLINE_FILL);
        head_stream.synchronize().expect("stream sync after online fill");
        fill_slice_via_mapped(&head_stream, &staging, iqn.target_params_slice(), n, TARGET_FILL);
        head_stream.synchronize().expect("stream sync after target fill");

        // Sanity: the fills actually landed and the buffers differ.
        let online_pre = read_slice_via_mapped(&head_stream, iqn.online_params_slice(), n);
        let target_pre = read_slice_via_mapped(&head_stream, iqn.target_params_slice(), n);
        assert_eq!(online_pre.len(), n);
        assert_eq!(target_pre.len(), n);
        assert!(online_pre.iter().all(|&v| v == ONLINE_FILL),
            "online fill did not propagate (saw {:?}…)", &online_pre[..online_pre.len().min(4)]);
        assert!(target_pre.iter().all(|&v| v == TARGET_FILL),
            "target fill did not propagate (saw {:?}…)", &target_pre[..target_pre.len().min(4)]);
        assert_ne!(online_pre, target_pre,
            "pre-sync buffers must differ — fill values were identical?");

        // Exercise the contract under test.
        iqn.sync_target_from_online()
            .expect("sync_target_from_online must succeed");
        head_stream.synchronize().expect("stream sync after sync_target_from_online");

        // Read both back and assert bit-for-bit equality across the full
        // parameter buffer. A stub body, wrong buffers, wrong direction,
        // or queue against the wrong stream all fail this assertion.
        let online_post = read_slice_via_mapped(&head_stream, iqn.online_params_slice(), n);
        let target_post = read_slice_via_mapped(&head_stream, iqn.target_params_slice(), n);

        // Online must be unchanged.
        assert_eq!(online_post, online_pre,
            "sync_target_from_online must not mutate online_params");
        // Target must now equal online — exact bit equality (DtoD copy is
        // not lossy; we use `f32::to_bits` to make NaN witnesses fail loud
        // if the implementation ever changes).
        assert_eq!(target_post.len(), online_post.len());
        for (i, (t, o)) in target_post.iter().zip(online_post.iter()).enumerate() {
            assert_eq!(t.to_bits(), o.to_bits(),
                "target[{i}] != online[{i}] post-sync: target={t} online={o}");
        }
    }
}
