#![allow(unsafe_code)]

//! Fused CUDA DQN training — forward + loss + backward + Adam in 4 kernel launches.
//!
//! Replaces 2,100+ Candle kernel dispatches per training batch with 4 fused CUDA
//! kernels captured in a CUDA Graph for zero-overhead replay.
//!
//! ## Architecture
//!
//! All tensors are pre-allocated at construction (fixed shapes for CUDA Graph
//! compatibility). Weight pointers come from existing `DuelingWeightSet` +
//! `BranchingWeightSet` objects (no duplication). The only per-batch host work
//! is copying batch data into pre-allocated input buffers and reading back
//! the scalar loss + td_errors.
//!
//! ## CUDA Graphs (split: forward + optimizer)
//!
//! The training pipeline is captured into TWO CUDA Graphs on the first
//! `train_step()` call:
//!
//! - **`graph_forward`**: zero → cuBLAS forward → C51 loss → C51 grad → cuBLAS backward
//! - **`graph_adam`**: grad_norm → Adam → unflatten (20 d2d copies)
//!
//! Between the two graph replays, external code can inject auxiliary gradients
//! (IQN trunk, attention, ensemble) into `grad_buf` via SAXPY. The Adam graph
//! then sees the combined C51 + auxiliary gradients.
//!
//! Only the Adam step counter (`t_buf`) and batch input data need updating
//! before each replay — both happen outside the captured graphs.
//!
//! ## Kernel phases
//!
//! Forward and backward passes are handled by cuBLAS (`CublasForward` /
//! `CublasBackward`). The two NVRTC utility kernels from `dqn_utility_kernels.cu`
//! handle the optimizer step:
//!
//! 1. **Grad norm** (`dqn_grad_norm_kernel`): Gradient L2 norm via warp + block
//!    reduction + atomicAdd into a single output float.
//! 2. **Adam update** (`dqn_adam_update_kernel`): Reads completed norm, clips
//!    gradients, applies AdamW update over flattened parameter view.
//!
//! ## Parameter layout
//!
//! The flat parameter/gradient/moment buffers use the same layout as the CUDA
//! `GOFF_*` defines: 20 weight tensors concatenated in order (w_s1, b_s1, w_s2,
//! b_s2, ..., w_b2out, b_b2out). See `compute_param_sizes()`.

use std::sync::Arc;

use cudarc::driver::{
    CudaFunction, CudaGraph, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg,
};
use cudarc::nvrtc::Ptx;
use tracing::info;

use crate::MLError;
use super::gpu_attention::GpuAttention;
use super::gpu_weights::{DuelingWeightSet, BranchingWeightSet};
use super::batched_forward::{CublasForward, f32_weight_ptrs};
use super::batched_backward::{CublasBackward, alloc_backward_scratch, raw_f32_ptr as bw_raw_ptr};

// ── Event tracking RAII guard ────────────────────────────────────────────────

/// RAII guard that disables cudarc event tracking on creation and
/// re-enables it on drop. Prevents early-return bugs where tracking
/// is left permanently disabled.
struct EventTrackingGuard<'a> {
    ctx: &'a cudarc::driver::CudaContext,
}

impl<'a> EventTrackingGuard<'a> {
    fn new(ctx: &'a cudarc::driver::CudaContext) -> Self {
        unsafe { ctx.disable_event_tracking(); }
        Self { ctx }
    }
}

impl Drop for EventTrackingGuard<'_> {
    fn drop(&mut self) {
        unsafe { self.ctx.enable_event_tracking(); }
        let _ = self.ctx.check_err();
    }
}

// ── CUDA Graph wrapper ──────────────────────────────────────────────────────

/// Wrapper to send `CudaGraph` across thread boundaries.
///
/// The graph is only launched on the same stream/context that created it.
/// The trainer is not shared across threads in practice.
struct SendSyncGraph(CudaGraph);

// Safety: CudaGraph is bound to a specific CUDA context. The trainer
// that owns it is always used from the thread that created the context.
// No concurrent access occurs.
unsafe impl Send for SendSyncGraph {}
// Safety: same reasoning — single-owner, no concurrent launch calls.
unsafe impl Sync for SendSyncGraph {}

// ── Configuration ───────────────────────────────────────────────────────────

/// Network dimensions and training hyperparameters for the fused CUDA trainer.
#[derive(Debug, Clone)]
pub struct GpuDqnTrainConfig {
    /// Input state dimension (e.g., 48 or 56 with OFI).
    pub state_dim: usize,
    /// First shared hidden layer width (default: 256).
    pub shared_h1: usize,
    /// Second shared hidden layer width (default: 256).
    pub shared_h2: usize,
    /// Value head hidden width (default: 128).
    pub value_h: usize,
    /// Advantage/branch head hidden width (default: 128).
    pub adv_h: usize,
    /// Number of C51 distributional atoms (default: 51).
    pub num_atoms: usize,
    /// C51 minimum support value (default: -50.0).
    pub v_min: f32,
    /// C51 maximum support value (default: 50.0).
    pub v_max: f32,
    /// Branch 0 (exposure) action count (default: 5).
    pub branch_0_size: usize,
    /// Branch 1 (order) action count (default: 3).
    pub branch_1_size: usize,
    /// Branch 2 (urgency) action count (default: 3).
    pub branch_2_size: usize,
    /// Training batch size (fixed for CUDA Graph capture).
    pub batch_size: usize,
    /// Discount factor for Bellman projection.
    pub gamma: f32,
    /// Learning rate for Adam optimizer.
    pub lr: f32,
    /// Adam β1 (first moment decay).
    pub beta1: f32,
    /// Adam β2 (second moment decay).
    pub beta2: f32,
    /// Adam ε (numerical stability).
    pub epsilon: f32,
    /// Decoupled weight decay (AdamW).
    pub weight_decay: f32,
    /// Maximum gradient L2 norm for clipping.
    pub max_grad_norm: f32,
    /// Spectral norm clipping target (σ_max). Constrains ||W||_σ ≤ σ_max.
    /// Default 3.0 (permits natural Xavier-init scaling, prevents explosion).
    /// Lower values = tighter constraint = more Q-value stability but less capacity.
    pub spectral_norm_sigma_max: f32,
    /// N-step returns (default: 1). When > 1, the experience collector pre-computes
    /// R_n = sum(gamma^i * r_{t+i}) and the C51 Bellman uses gamma^n.
    pub n_steps: usize,
    /// IQN quantile loss weight relative to C51 (0.0 = C51 only)
    pub iqn_lambda: f32,
    /// Number of IQN quantile samples per forward pass (0 = IQN disabled)
    pub iqn_num_quantiles: usize,
    /// Cosine embedding dimension for IQN τ encoding
    pub iqn_embedding_dim: usize,
    /// Huber threshold for quantile loss
    pub iqn_kappa: f32,
    /// Entropy regularization coefficient for C51 distributions.
    /// Adds -coeff * H(p) to the loss, preventing distribution collapse to a single atom.
    /// 0.0 = disabled. Typical: 0.01-0.05.
    pub entropy_coefficient: f32,
    /// Number of epochs to use MSE loss on expected Q-values before switching to C51.
    /// MSE provides stronger gradient signal during early training.
    pub c51_warmup_epochs: usize,
    /// CQL regularization strength (0.0 = disabled, 0.1 = mild, 1.0 = full offline-RL).
    pub cql_alpha: f32,
}

impl Default for GpuDqnTrainConfig {
    fn default() -> Self {
        Self {
            state_dim: 48,   // Actual aligned feature dim (40 market + 3 portfolio = 43, padded to 48)
            shared_h1: 256,
            shared_h2: 256,
            value_h: 128,
            adv_h: 128,
            num_atoms: 51,
            // v_min/v_max computed from reward_scale=10, gamma=0.95:
            // Q* ~ 10 / (1-0.95) * 1.2 = 240, clamped to [20, 300] → ±240
            v_min: -(10.0_f32 / (1.0 - 0.95) * 1.2).clamp(20.0, 300.0),
            v_max: (10.0_f32 / (1.0 - 0.95) * 1.2).clamp(20.0, 300.0),
            branch_0_size: 9,
            branch_1_size: 3,
            branch_2_size: 3,
            batch_size: 64,   // Smoke test default (production overrides via from_hyperparams)
            gamma: 0.99,
            n_steps: 1,
            lr: 3e-5,     // Conservative default for stable training
            beta1: 0.9,
            beta2: 0.999,
            epsilon: 1e-8,
            weight_decay: 1e-4,  // Standard L2 regularization strength
            max_grad_norm: 10.0,
            spectral_norm_sigma_max: 3.0,
            iqn_lambda: 0.25,
            iqn_num_quantiles: 0,
            iqn_embedding_dim: 64,
            iqn_kappa: 1.0,
            entropy_coefficient: 0.001, // Must be ≤ reward magnitude (~0.001). Old 0.01 was 10x reward → entropy dominated learning.
            c51_warmup_epochs: 5,
            cql_alpha: 0.1,
        }
    }
}

/// Loss mode for the fused DQN trainer.
///
/// During warmup, MSE on expected Q-values provides stronger gradient signal
/// than C51 cross-entropy between nearly-identical distributions. After warmup,
/// C51 cross-entropy refines the full value distribution shape.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LossMode {
    /// MSE on expected Q-values -- strong gradient signal for warmup.
    Mse,
    /// C51 cross-entropy on full distribution -- refines value distribution shape.
    C51,
}

/// Result of a single fused training step.
#[derive(Debug, Clone)]
pub struct FusedTrainResult {
    /// Batch-mean weighted loss (for logging).
    pub total_loss: f32,
    /// Per-sample TD errors for PER priority update.
    pub td_errors: Vec<f32>,
    /// Pre-clip gradient L2 norm (for monitoring).
    pub grad_norm: f32,
}

/// Q-value statistics computed entirely on GPU (5 scalars, 20-byte readback).
#[derive(Debug, Clone, Copy)]
pub struct QValueStatsResult {
    /// Average of per-sample max Q-values
    pub avg_max_q: f64,
    /// Global minimum Q-value
    pub q_min: f32,
    /// Global maximum Q-value
    pub q_max: f32,
    /// Mean of all Q-values
    pub q_mean: f32,
    /// Variance of all Q-values
    pub q_variance: f32,
}

/// Scalar-only result from the fused GPU training path.
///
/// TD errors stay on GPU (`td_errors_buf`) — no readback. PER priority update
/// is done by a dedicated CUDA kernel that reads `td_errors_buf` directly.
/// Only 8 bytes DtoH per step (loss + grad_norm).
#[derive(Debug, Clone, Copy)]
pub struct FusedTrainScalars {
    /// Batch-mean weighted loss (for logging).
    pub total_loss: f32,
    /// Pre-clip gradient L2 norm (for monitoring).
    pub grad_norm: f32,
}

// ── Parameter layout ────────────────────────────────────────────────────────
//
// Matches GOFF_* / PARAM_* defines in dqn_training_kernel.cu exactly.
// 20 weight tensors in order: w_s1, b_s1, w_s2, b_s2, w_v1, b_v1, w_v2, b_v2,
// w_b0fc, b_b0fc, w_b0out, b_b0out, w_b1fc, b_b1fc, w_b1out, b_b1out,
// w_b2fc, b_b2fc, w_b2out, b_b2out.

/// Compute the size (element count) of each of the 20 weight tensors.
pub(crate) fn compute_param_sizes(cfg: &GpuDqnTrainConfig) -> [usize; 20] {
    [
        cfg.shared_h1 * cfg.state_dim,                      // w_s1
        cfg.shared_h1,                                       // b_s1
        cfg.shared_h2 * cfg.shared_h1,                      // w_s2
        cfg.shared_h2,                                       // b_s2
        cfg.value_h * cfg.shared_h2,                        // w_v1
        cfg.value_h,                                         // b_v1
        cfg.num_atoms * cfg.value_h,                        // w_v2
        cfg.num_atoms,                                       // b_v2
        cfg.adv_h * cfg.shared_h2,                          // w_b0fc
        cfg.adv_h,                                           // b_b0fc
        cfg.branch_0_size * cfg.num_atoms * cfg.adv_h,     // w_b0out
        cfg.branch_0_size * cfg.num_atoms,                  // b_b0out
        cfg.adv_h * cfg.shared_h2,                          // w_b1fc
        cfg.adv_h,                                           // b_b1fc
        cfg.branch_1_size * cfg.num_atoms * cfg.adv_h,     // w_b1out
        cfg.branch_1_size * cfg.num_atoms,                  // b_b1out
        cfg.adv_h * cfg.shared_h2,                          // w_b2fc
        cfg.adv_h,                                           // b_b2fc
        cfg.branch_2_size * cfg.num_atoms * cfg.adv_h,     // w_b2out
        cfg.branch_2_size * cfg.num_atoms,                  // b_b2out
    ]
}

pub(crate) fn compute_total_params(cfg: &GpuDqnTrainConfig) -> usize {
    compute_param_sizes(cfg).iter().sum()
}

// ── Main struct ─────────────────────────────────────────────────────────────

/// Fused CUDA DQN trainer — replaces Candle dispatch chain with 4 kernel launches.
///
/// All GPU buffers are pre-allocated at construction. The trainer borrows
/// weight sets from the caller (no weight duplication). Only batch input
/// data is uploaded per step; outputs (loss, td_errors, grad_norm) are downloaded.
///
/// On the first `train_step()`, the kernel sequence is captured into two CUDA
/// Graphs (`graph_forward` and `graph_adam`). Subsequent calls replay both
/// graphs for zero kernel-launch overhead, with an injection point between
/// them for auxiliary gradients (IQN, ensemble, attention).
#[allow(missing_debug_implementations)] // CudaSlice does not implement Debug
pub struct GpuDqnTrainer {
    config: GpuDqnTrainConfig,
    stream: Arc<CudaStream>,

    // ── Compiled kernels ────────────────────────────────────────────
    // Dead kernels (forward_loss_kernel, forward_only_kernel, backward_kernel)
    // removed — replaced by cuBLAS forward/backward. See Task 2 for method cleanup.
    grad_norm_kernel: CudaFunction,
    adam_update_kernel: CudaFunction,
    ema_kernel: CudaFunction,
    per_update_kernel: CudaFunction,
    f32_to_bf16_kernel: CudaFunction,
    bf16_to_f32_kernel: CudaFunction,
    saxpy_kernel: CudaFunction,
    zero_kernel: CudaFunction,
    regime_scale_kernel: CudaFunction,
    shrink_perturb_kernel: CudaFunction,
    relu_mask_kernel: CudaFunction,
    spectral_norm_kernel: CudaFunction,
    /// Spectral norm left singular vectors: [out_dim] per weight matrix (W_s1, W_s2)
    spec_u_s1: CudaSlice<f32>,
    spec_v_s1: CudaSlice<f32>,
    spec_u_s2: CudaSlice<f32>,
    spec_v_s2: CudaSlice<f32>,
    /// IQN trunk gradient Adam first moment [trunk_params] (separate from C51 Adam).
    iqn_trunk_m: CudaSlice<f32>,
    /// IQN trunk gradient Adam second moment [trunk_params].
    iqn_trunk_v: CudaSlice<f32>,
    /// IQN trunk Adam step counter.
    iqn_trunk_adam_step: i32,
    /// IQN trunk grad norm scratch [1].
    iqn_trunk_grad_norm: CudaSlice<f32>,
    /// IQN trunk Adam step counter on device [1].
    iqn_trunk_t_buf: CudaSlice<i32>,
    /// Number of trunk parameters (w_s1 + b_s1 + w_s2 + b_s2).
    trunk_param_count: usize,

    // ── Flat BF16 weight mirrors (forward kernel reads BF16 for tensor core throughput) ──
    // Single contiguous BF16 buffer per network (online + target), same GOFF_* layout
    // as the F32 flat buffers. One f32_to_bf16_kernel launch converts the entire buffer
    // instead of 20 per-tensor launches. Forward kernels read at precomputed byte offsets.
    bf16_params_buf: CudaSlice<u16>,         // [TOTAL_PARAMS] flat online BF16
    bf16_target_params_buf: CudaSlice<u16>,  // [TOTAL_PARAMS] flat target BF16
    bf16_mirrors_initialized: bool,
    /// Precomputed byte offsets into flat BF16 buffers for each of the 20 weight tensors.
    /// Layout matches GOFF_* order: w_s1, b_s1, w_s2, b_s2, ..., w_bu2, b_bu2.
    /// Each offset is in bytes (element offset * sizeof(u16)).
    /// Segments are padded to even element counts for 4-byte aligned short2 loads.
    bf16_goff_byte_offsets: [u64; 20],
    /// Total u16 elements in padded BF16 buffers (>= total_params due to alignment padding).
    bf16_padded_total: usize,

    // ── Batch input buffers (uploaded per step) ─────────────────────
    states_buf: CudaSlice<f32>,       // [B, STATE_DIM]
    next_states_buf: CudaSlice<f32>,  // [B, STATE_DIM]
    actions_buf: CudaSlice<i32>,      // [B]
    rewards_buf: CudaSlice<f32>,      // [B]
    dones_buf: CudaSlice<f32>,        // [B]
    is_weights_buf: CudaSlice<f32>,   // [B]

    // ── Activation save buffers (forward → backward) ────────────────
    save_h_s1: CudaSlice<f32>,        // [B, SHARED_H1]
    save_h_s2: CudaSlice<f32>,        // [B, SHARED_H2]
    save_h_v: CudaSlice<f32>,         // [B, VALUE_H]
    save_h_b0: CudaSlice<f32>,        // [B, ADV_H]
    save_h_b1: CudaSlice<f32>,        // [B, ADV_H]
    save_h_b2: CudaSlice<f32>,        // [B, ADV_H]
    save_current_lp: CudaSlice<f32>,  // [B, NUM_BRANCHES(3), NUM_ATOMS]
    save_projected: CudaSlice<f32>,   // [B, NUM_BRANCHES(3), NUM_ATOMS]

    // ── Forward output buffers ──────────────────────────────────────
    per_sample_loss_buf: CudaSlice<f32>,  // [B]
    td_errors_buf: CudaSlice<f32>,        // [B]
    total_loss_buf: CudaSlice<f32>,       // [1]

    // ── Forward-only Q-value output ─────────────────────────────────
    q_out_buf: CudaSlice<f32>,            // [B, TOTAL_ACTIONS(11)]

    // ── Backward / Adam buffers ─────────────────────────────────────
    grad_buf: CudaSlice<f32>,         // [TOTAL_PARAMS] gradient accumulator
    params_buf: CudaSlice<f32>,       // [TOTAL_PARAMS] flat online parameters
    target_params_buf: CudaSlice<f32>, // [TOTAL_PARAMS] flat target parameters (EMA)
    m_buf: CudaSlice<f32>,            // [TOTAL_PARAMS] Adam first moment
    v_buf: CudaSlice<f32>,            // [TOTAL_PARAMS] Adam second moment
    grad_norm_buf: CudaSlice<f32>,    // [1] pre-clip gradient L2 norm

    // ── Adam step counter on device (CUDA Graph cannot bake scalars) ─
    t_buf: CudaSlice<i32>,            // [1] current Adam step

    // ── Training state ──────────────────────────────────────────────
    adam_step: i32,
    total_params: usize,
    params_initialized: bool,
    target_params_initialized: bool,

    // ── CUDA Graphs (split: forward+backward, then optimizer) ──────
    // Graph A: zero → cuBLAS forward → C51 loss → C51 grad → cuBLAS backward
    // Graph B: grad_norm → Adam → unflatten
    // Between A and B: external code can ADD auxiliary gradients to grad_buf
    // (IQN trunk, ensemble heads) — single Adam sees combined gradient.
    graph_forward: Option<SendSyncGraph>,
    graph_adam: Option<SendSyncGraph>,

    // ── Consolidated transfer buffers ─────────────────────────────
    /// Single staging buffer for batch upload consolidation.
    /// Layout: [states(B*SD) | next_states(B*SD) | actions_as_f32(B) | rewards(B) | dones(B) | is_weights(B)]
    /// One HtoD transfer replaces 6 separate PCIe round-trips.
    upload_staging_buf: CudaSlice<f32>,
    /// Total element count of the staging buffer.
    upload_staging_len: usize,

    /// Single readback buffer for loss download consolidation.
    /// Layout: [total_loss(1) | grad_norm(1) | td_errors(B)]
    /// One DtoH transfer replaces 3 separate PCIe round-trips.
    readback_buf: CudaSlice<f32>,
    /// Pre-allocated host-side readback vec (avoids per-step allocation).
    readback_host: Vec<f32>,

    /// Pre-allocated host-side staging vec for batch upload (avoids per-step malloc).
    /// Reused via `.clear()` + `.extend_from_slice()` each step.
    upload_staging_host: Vec<f32>,

    // ── PER priority update (GPU-native) ─────────────────────────────
    /// Running max priority across all batches in an epoch (atomicMax in kernel).
    /// Read back once per epoch via `batch_max_readback_and_reset()`.
    batch_max_buf: CudaSlice<f32>,

    /// Scalar-only readback buffer: [total_loss(1) | grad_norm(1)] = 8 bytes.
    /// Used by `execute_train_scalars_only()` to avoid reading back td_errors.
    scalars_readback_buf: CudaSlice<f32>,
    scalars_readback_host: [f32; 2],

    // ── BF16 batch staging buffers ──────────────────────────────────
    // Pre-allocated CudaSlice<u16> buffers for DtoD copy of BF16 tensors
    // from GpuBatch (states, next_states only — stored in BF16 ring buffer).
    // The bf16_to_f32_kernel converts these to the F32 trainer input buffers.
    // Rewards, dones, weights are F32 in GpuBatch and use direct DtoD to F32 bufs.
    bf16_states_buf: CudaSlice<u16>,      // [B * STATE_DIM]
    bf16_next_states_buf: CudaSlice<u16>, // [B * STATE_DIM]

    // ── cuBLAS batched forward (Phase 2 — high-occupancy path) ───────
    /// cuBLAS forward context (handle + bias kernels).
    cublas_forward: CublasForward,

    // Scratch buffers for cuBLAS target network forward pass.
    // Online activations reuse the existing save_h_* buffers.
    // Target only needs h_s2 (for DDQN) + scratch for intermediate layers.
    /// Target trunk layer 1 scratch: [B, SH1]
    tg_h_s1_scratch: CudaSlice<f32>,
    /// Target trunk layer 2 output — kept for Double-DQN action selection: [B, SH2]
    tg_h_s2_buf: CudaSlice<f32>,
    /// Target value head hidden scratch: [B, VH]
    tg_h_v_scratch: CudaSlice<f32>,
    /// Target branch head hidden scratch (reused for all 3 branches): [B, AH]
    tg_h_b_scratch: CudaSlice<f32>,
    /// Target value head logits: [B, NA]
    tg_v_logits_buf: CudaSlice<f32>,
    /// Target branch logits: [B, (B0+B1+B2)*NA]
    tg_b_logits_buf: CudaSlice<f32>,

    /// Online value head logits (from cuBLAS pass on current states): [B, NA]
    on_v_logits_buf: CudaSlice<f32>,
    /// Online branch logits (from cuBLAS pass on current states): [B, (B0+B1+B2)*NA]
    on_b_logits_buf: CudaSlice<f32>,

    /// Online-next value logits (Double DQN action selector on next_states): [B, NA]
    on_next_v_logits_buf: CudaSlice<f32>,
    /// Online-next branch logits (Double DQN on next_states): [B, (B0+B1+B2)*NA]
    on_next_b_logits_buf: CudaSlice<f32>,
    /// Online-next scratch (reused): h_s1[B,SH1], h_s2[B,SH2], h_v[B,VH], h_b[B,AH]
    on_next_h_s1_scratch: CudaSlice<f32>,
    on_next_h_s2_scratch: CudaSlice<f32>,
    on_next_h_v_scratch: CudaSlice<f32>,
    on_next_h_b_scratch: CudaSlice<f32>,

    /// Compiled C51 loss kernel (standalone, replaces the fused forward+loss kernel).
    c51_loss_kernel: CudaFunction,
    /// C51 loss gradient kernel: computes dL/d_logits for cuBLAS backward.
    c51_grad_kernel: CudaFunction,
    /// MSE loss kernel on expected Q-values (warmup before C51).
    mse_loss_kernel: CudaFunction,
    /// MSE loss gradient kernel through softmax expectation (warmup before C51).
    mse_grad_kernel: CudaFunction,
    /// Current loss mode (MSE warmup vs C51 distributional).
    loss_mode: LossMode,
    /// C51 blend factor: 0.0 = pure MSE, 1.0 = pure C51.
    /// Ramps linearly over warmup epochs for smooth transition.
    c51_alpha: f32,
    /// Loss mode at the time of the last CUDA Graph capture.
    /// When this differs from `loss_mode`, the graph must be recaptured.
    last_captured_loss_mode: Option<LossMode>,
    /// Scratch buffers for blended loss (MSE grad stored here, then blended into d_value/d_adv)
    d_value_logits_mse: CudaSlice<f32>,
    d_adv_logits_mse: CudaSlice<f32>,
    /// Gradient w.r.t. value logits: [B, NA]
    d_value_logits_buf: CudaSlice<f32>,
    /// Gradient w.r.t. branch logits: [B, (B0+B1+B2)*NA]
    d_adv_logits_buf: CudaSlice<f32>,

    /// Precomputed F32 weight byte offsets into params_buf.
    /// Same GOFF_* layout as BF16, but for the F32 cuBLAS path.
    f32_goff_byte_offsets: [u64; 20],

    // ── cuBLAS batched backward (Phase 2 Task 2) ──────────────────────
    /// cuBLAS backward context (handle + ReLU mask + bias grad kernels).
    cublas_backward: CublasBackward,

    // Scratch buffers for inter-layer gradient propagation during cuBLAS backward.
    // Sized for the widest layer at each point in the network.
    /// Accumulated d_h_s2 from value head + all 3 branch heads: [B, SH2]
    bw_d_h_s2: CudaSlice<f32>,
    /// Upstream gradient for shared layer 1: [B, SH1]
    bw_d_h_s1: CudaSlice<f32>,
    /// Upstream gradient for value FC: [B, VH]
    bw_d_h_v: CudaSlice<f32>,
    /// Branch FC upstream gradients (one per branch): [B, AH]
    bw_d_h_b0: CudaSlice<f32>,
    bw_d_h_b1: CudaSlice<f32>,
    bw_d_h_b2: CudaSlice<f32>,

    // ── Expected Q-value kernel (ad-hoc validation, not captured in CUDA Graph) ─
    /// Converts C51 value+advantage logits → expected Q-values (validation path).
    /// Input: on_v_logits_buf [B, NA], on_b_logits_buf [B, (B0+B1+B2)*NA]
    /// Output: q_out_buf [B, B0+B1+B2]
    expected_q_kernel: CudaFunction,
    /// GPU reduction: q_out_buf → 5 scalars [avg_max_q, q_min, q_max, q_mean, q_var]
    q_stats_kernel: CudaFunction,
    /// GPU buffer for Q-value statistics [5 floats]
    q_stats_buf: CudaSlice<f32>,

    // ── CQL (Conservative Q-Learning) penalty kernel ─────────────────
    /// Computes CQL logit gradients: dCQL/d_value_logits and dCQL/d_adv_logits.
    /// Only used when `config.use_cql == true && config.cql_alpha > 0`.
    cql_logit_grad_kernel: Option<CudaFunction>,
    /// CQL scratch: value logit gradients [B, NA]
    cql_d_value_logits: CudaSlice<f32>,
    /// CQL scratch: advantage logit gradients [B, (B0+B1+B2)*NA]
    cql_d_adv_logits: CudaSlice<f32>,
}

impl Drop for GpuDqnTrainer {
    fn drop(&mut self) {
        // Synchronize stream and destroy graphs BEFORE CudaSlice fields drop.
        unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }
        self.graph_forward = None;
        self.graph_adam = None;
    }
}

impl GpuDqnTrainer {
    /// Configured batch size (fixed at construction for CUDA Graph compatibility).
    pub fn batch_size(&self) -> usize {
        self.config.batch_size
    }

    /// Reference to saved h_s2 activations from the last forward pass.
    ///
    /// Shape: `[B, SHARED_H2]` — the shared trunk output for online states.
    /// Valid after `train_step()` or `forward_loss()` has been called.
    /// Used by the IQN dual-head to read trunk activations without re-computation.
    pub fn save_h_s2(&self) -> &CudaSlice<f32> {
        &self.save_h_s2
    }

    /// Backward gradient w.r.t. h_s2 (trunk activation).
    pub fn bw_d_h_s2_buf(&self) -> &CudaSlice<f32> {
        &self.bw_d_h_s2
    }

    /// Shared trunk hidden layer 2 dimension.
    pub fn shared_h2(&self) -> usize {
        self.config.shared_h2
    }

    /// Reference to the states buffer on GPU.
    ///
    /// Shape: `[B, STATE_DIM]` — contains the batch's states after `upload_batch_gpu()`.
    /// Used by IQL value network to read states without Candle tensor intermediaries.
    pub fn states_buf(&self) -> &CudaSlice<f32> {
        &self.states_buf
    }

    /// Reference to the next_states buffer on GPU.
    ///
    /// Shape: `[B, STATE_DIM]` — contains the batch's next states after `train_step()`.
    /// Used by the IQN dual-head to compute target h_s2 via a trunk forward kernel.
    pub fn next_states_buf(&self) -> &CudaSlice<f32> {
        &self.next_states_buf
    }

    /// Reference to the actions buffer on GPU.
    ///
    /// Shape: `[B]` i32 — flat actions from the batch. Valid after `train_step()` or
    /// `train_step_gpu()`. Used by IQN to decode flat → branch actions on GPU.
    pub fn actions_buf(&self) -> &CudaSlice<i32> {
        &self.actions_buf
    }

    /// Reference to the rewards buffer on GPU.
    ///
    /// Shape: `[B]` f32 — rewards from the batch. Valid after `train_step()` or
    /// `train_step_gpu()`. Used by IQN for Bellman target computation on GPU.
    pub fn rewards_buf(&self) -> &CudaSlice<f32> {
        &self.rewards_buf
    }

    /// Reference to the dones buffer on GPU.
    ///
    /// Shape: `[B]` f32 — done flags (0.0/1.0). Valid after `train_step()` or
    /// `train_step_gpu()`. Used by IQN for Bellman target masking on GPU.
    pub fn dones_buf(&self) -> &CudaSlice<f32> {
        &self.dones_buf
    }

    /// Reference to the training config.
    pub fn config(&self) -> &GpuDqnTrainConfig {
        &self.config
    }

    /// Raw device pointer to `grad_buf` for auxiliary gradient injection between graph phases.
    ///
    /// Between `graph_forward` replay (which populates `grad_buf` with C51 gradients)
    /// and `graph_adam` replay (which reads `grad_buf` for Adam), external code can ADD
    /// auxiliary gradients (IQN trunk, ensemble heads) to `grad_buf` via SAXPY.
    ///
    /// Caller must use `EventTrackingGuard` and operate on the same stream.
    pub fn grad_buf_ptr(&self) -> u64 {
        raw_device_ptr(&self.grad_buf, &self.stream)
    }

    /// Apply IQN auxiliary gradient to the shared trunk via SAXPY into `grad_buf`.
    ///
    /// After `graph_forward` replay (C51 forward → backward) and IQN backward
    /// have both completed, this method backpropagates IQN's `d_h_s2` through
    /// the trunk's two FC layers and ADDs the resulting gradients to `grad_buf`
    /// (which already contains C51's gradients from `graph_forward`).
    ///
    /// The caller then replays `graph_adam` which sees the combined C51+IQN
    /// gradient and applies a single Adam update. This eliminates the momentum
    /// mismatch from having a separate IQN Adam optimizer.
    ///
    /// Steps:
    /// 1. Zero scratch buffer (`iqn_trunk_m`, repurposed as scratch)
    /// 2. cuBLAS backward through trunk → trunk gradients in scratch
    /// 3. SAXPY: `grad_buf[trunk] += iqn_lambda * scratch[trunk]`
    pub fn apply_iqn_trunk_gradient(
        &mut self,
        iqn_d_h_s2: &CudaSlice<f32>,
        _online_dueling: &mut DuelingWeightSet,
    ) -> Result<(), MLError> {
        let b = self.config.batch_size;
        let sd = self.config.state_dim;
        let sh1 = self.config.shared_h1;
        let sh2 = self.config.shared_h2;
        let f32_size = std::mem::size_of::<f32>();

        // Sync stream to ensure graph_forward replay completed before we touch buffers.
        unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }
        let _ = self.stream.context().check_err();

        // Disable event tracking for all buffer pointer extractions in this method.
        let _evt_guard = EventTrackingGuard::new(self.stream.context());

        // Trunk gradient element counts
        let w_s1_n = sh1 * sd;
        let b_s1_n = sh1;
        let w_s2_n = sh2 * sh1;
        let b_s2_n = sh2;
        let trunk_grad_total = w_s1_n + b_s1_n + w_s2_n + b_s2_n;

        // ── 1. Zero the scratch buffer (repurposed iqn_trunk_m) ───────────
        // cuBLAS backward_fc_layer accumulates (beta=1.0), so scratch must be zeroed.
        // iqn_trunk_m is [trunk_param_count] — same size as the trunk portion of grad_buf.
        {
            let scratch_base = raw_device_ptr(&self.iqn_trunk_m, &self.stream);
            let n_bytes = trunk_grad_total * f32_size;
            unsafe {
                cudarc::driver::result::memset_d8_async(
                    scratch_base, 0u8, n_bytes, self.stream.cu_stream()
                ).map_err(|e| MLError::ModelError(format!("memset iqn scratch: {e}")))?;
            }
        }

        // ── 2. Copy IQN d_h_s2 → bw_d_h_s2 scratch ──────────────────────
        {
            let n_bytes = b * sh2 * f32_size;
            let src = raw_device_ptr(iqn_d_h_s2, &self.stream);
            let dst = raw_device_ptr(&self.bw_d_h_s2, &self.stream);
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst, src, n_bytes, self.stream.cu_stream()
                ).map_err(|e| MLError::ModelError(format!("IQN d_h_s2 DtoD: {e}")))?;
            }
        }

        // ── 3. ReLU mask: bw_d_h_s2 *= (save_h_s2 > 0) ──────────────────
        {
            let d_ptr = raw_device_ptr(&self.bw_d_h_s2, &self.stream);
            let act_ptr = raw_device_ptr(&self.save_h_s2, &self.stream);
            let n_relu = (b * sh2) as i32;
            let blocks = ((b * sh2 + 255) / 256) as u32;
            unsafe {
                self.stream
                    .launch_builder(&self.relu_mask_kernel)
                    .arg(&d_ptr)
                    .arg(&act_ptr)
                    .arg(&n_relu)
                    .launch(LaunchConfig {
                        grid_dim: (blocks, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| MLError::ModelError(format!("relu_mask h_s2: {e}")))?;
            }
        }

        // ── 4. Backward FC layer 2: h_s1 → h_s2 (into SCRATCH) ──────────
        // Computes dW_s2, db_s2 into scratch (iqn_trunk_m), d_h_s1 into bw_d_h_s1.
        {
            let dy = bw_raw_ptr(&self.bw_d_h_s2, &self.stream);
            let x = bw_raw_ptr(&self.save_h_s1, &self.stream);
            let param_sizes = compute_param_sizes(&self.config);
            let w_ptrs = f32_weight_ptrs(&self.params_buf, &param_sizes, &self.stream);
            let w = w_ptrs[2]; // w_s2

            let scratch_base = raw_device_ptr(&self.iqn_trunk_m, &self.stream);
            let f32_u = f32_size as u64;
            let dw = scratch_base + (w_s1_n + b_s1_n) as u64 * f32_u; // goff_w_s2 in scratch
            let db = dw + w_s2_n as u64 * f32_u;                       // goff_b_s2 in scratch

            let dx = bw_raw_ptr(&self.bw_d_h_s1, &self.stream);

            self.cublas_backward.backward_fc_layer(
                &self.stream, dy, x, w, dw, db, dx, sh2, sh1, b,
            )?;
        }

        // ── 5. ReLU mask: bw_d_h_s1 *= (save_h_s1 > 0) ──────────────────
        {
            let d_ptr = raw_device_ptr(&self.bw_d_h_s1, &self.stream);
            let act_ptr = raw_device_ptr(&self.save_h_s1, &self.stream);
            let n_relu = (b * sh1) as i32;
            let blocks = ((b * sh1 + 255) / 256) as u32;
            unsafe {
                self.stream
                    .launch_builder(&self.relu_mask_kernel)
                    .arg(&d_ptr)
                    .arg(&act_ptr)
                    .arg(&n_relu)
                    .launch(LaunchConfig {
                        grid_dim: (blocks, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| MLError::ModelError(format!("relu_mask h_s1: {e}")))?;
            }
        }

        // ── 6. Backward FC layer 1: states → h_s1 (into SCRATCH) ────────
        // Computes dW_s1, db_s1 into scratch. dx=0 (skip input gradient).
        {
            let dy = bw_raw_ptr(&self.bw_d_h_s1, &self.stream);
            let x = bw_raw_ptr(&self.states_buf, &self.stream);
            let param_sizes = compute_param_sizes(&self.config);
            let w_ptrs = f32_weight_ptrs(&self.params_buf, &param_sizes, &self.stream);
            let w = w_ptrs[0]; // w_s1

            let scratch_base = raw_device_ptr(&self.iqn_trunk_m, &self.stream);
            let f32_u = f32_size as u64;
            let dw = scratch_base;                              // goff_w_s1 in scratch
            let db = scratch_base + w_s1_n as u64 * f32_u;     // goff_b_s1 in scratch

            self.cublas_backward.backward_fc_layer(
                &self.stream, dy, x, w, dw, db, 0, sh1, sd, b,
            )?;
        }

        // ── 7. SAXPY: grad_buf[trunk] += iqn_lambda * scratch[trunk] ─────
        // Adds scaled IQN trunk gradient to C51's gradient already in grad_buf.
        // graph_adam will then see the combined gradient and apply single Adam.
        {
            let grad_ptr = raw_device_ptr(&self.grad_buf, &self.stream);
            let scratch_ptr = raw_device_ptr(&self.iqn_trunk_m, &self.stream);
            let scale = self.config.iqn_lambda;
            let n_i32 = trunk_grad_total as i32;
            let blocks = ((trunk_grad_total + 255) / 256) as u32;
            unsafe {
                self.stream
                    .launch_builder(&self.saxpy_kernel)
                    .arg(&grad_ptr)
                    .arg(&scratch_ptr)
                    .arg(&scale)
                    .arg(&n_i32)
                    .launch(LaunchConfig {
                        grid_dim: (blocks, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| MLError::ModelError(format!("IQN trunk SAXPY: {e}")))?;
            }
        }

        // _evt_guard drops here → re-enables event tracking

        Ok(())
    }

    /// Apply ensemble diversity gradient through the value head into the trunk.
    ///
    /// Takes `d_logits` [B, NA] (the KL gradient w.r.t. head-0 value logits
    /// produced by the ensemble KL gradient kernel) and backpropagates through:
    ///   1. Value output layer:  d_logits [B, NA] × W_v2^T → d_h_v [B, VH]
    ///   2. ReLU mask:           d_h_v *= (save_h_v > 0)
    ///   3. Value FC layer:      d_h_v [B, VH]  × W_v1^T → d_h_s2 [B, SH2]
    ///   4. Trunk backward (shared layers 2 → 1) into scratch buffer
    ///   5. SAXPY: grad_buf[trunk] += scale * scratch[trunk]
    ///
    /// This mirrors `apply_iqn_trunk_gradient` but starts from value logits
    /// instead of h_s2. The `scale` parameter is `ensemble_diversity_weight`.
    ///
    /// All computation is GPU-only — no CPU readback in the gradient path.
    #[allow(clippy::too_many_arguments)]
    pub fn apply_ensemble_diversity_backward(
        &mut self,
        d_logits_ptr: u64,
        scale: f32,
    ) -> Result<(), MLError> {
        let b = self.config.batch_size;
        let sd = self.config.state_dim;
        let sh1 = self.config.shared_h1;
        let sh2 = self.config.shared_h2;
        let vh = self.config.value_h;
        let na = self.config.num_atoms;
        let f32_size = std::mem::size_of::<f32>();

        // Sync stream to ensure graph_forward replay completed before we touch buffers.
        unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }
        let _ = self.stream.context().check_err();

        let _evt_guard = EventTrackingGuard::new(self.stream.context());

        // Trunk gradient element counts (for the scratch → SAXPY step)
        let w_s1_n = sh1 * sd;
        let b_s1_n = sh1;
        let w_s2_n = sh2 * sh1;
        let b_s2_n = sh2;
        let trunk_grad_total = w_s1_n + b_s1_n + w_s2_n + b_s2_n;

        // ── 1. Zero the scratch buffer (iqn_trunk_m) ───────────────────────
        {
            let scratch_base = raw_device_ptr(&self.iqn_trunk_m, &self.stream);
            let n_bytes = trunk_grad_total * f32_size;
            unsafe {
                cudarc::driver::result::memset_d8_async(
                    scratch_base, 0u8, n_bytes, self.stream.cu_stream()
                ).map_err(|e| MLError::ModelError(format!("memset ens scratch: {e}")))?;
            }
        }

        // ── 2. Backward value output layer: d_logits → d_h_v ───────────────
        // d_logits [B, NA] × W_v2^T [VH, NA] → d_h_v [B, VH]
        // Only upstream gradient (dX) is needed — skip dW/db for value head.
        {
            let param_sizes = compute_param_sizes(&self.config);
            let w_ptrs = f32_weight_ptrs(&self.params_buf, &param_sizes, &self.stream);
            let w_v2 = w_ptrs[6]; // W_v2 [NA, VH]

            let dx = raw_device_ptr(&self.bw_d_h_v, &self.stream);

            // launch_dx_only: computes only dX = dY @ W^T (no dW/db)
            self.cublas_backward.launch_dx_only(
                &self.stream, d_logits_ptr, w_v2, dx,
                na, vh, b, 0.0,
            )?;
        }

        // ── 3. ReLU mask: d_h_v *= (save_h_v > 0) ─────────────────────────
        {
            let d_ptr = raw_device_ptr(&self.bw_d_h_v, &self.stream);
            let act_ptr = raw_device_ptr(&self.save_h_v, &self.stream);
            let n_relu = (b * vh) as i32;
            let blocks = ((b * vh + 255) / 256) as u32;
            unsafe {
                self.stream
                    .launch_builder(&self.relu_mask_kernel)
                    .arg(&d_ptr)
                    .arg(&act_ptr)
                    .arg(&n_relu)
                    .launch(LaunchConfig {
                        grid_dim: (blocks, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| MLError::ModelError(format!("relu_mask ens_h_v: {e}")))?;
            }
        }

        // ── 4. Backward value FC layer: d_h_v → d_h_s2 ────────────────────
        // d_h_v [B, VH] × W_v1^T [SH2, VH] → d_h_s2 [B, SH2]
        // Only upstream gradient (dX) is needed — skip dW/db for value head.
        {
            let param_sizes = compute_param_sizes(&self.config);
            let w_ptrs = f32_weight_ptrs(&self.params_buf, &param_sizes, &self.stream);
            let w_v1 = w_ptrs[4]; // W_v1 [VH, SH2]

            let dy = raw_device_ptr(&self.bw_d_h_v, &self.stream);
            let dx = raw_device_ptr(&self.bw_d_h_s2, &self.stream);

            // launch_dx_only: computes only dX = dY @ W^T (no dW/db)
            self.cublas_backward.launch_dx_only(
                &self.stream, dy, w_v1, dx,
                vh, sh2, b, 0.0,
            )?;
        }

        // ── 5. ReLU mask: d_h_s2 *= (save_h_s2 > 0) ──────────────────────
        {
            let d_ptr = raw_device_ptr(&self.bw_d_h_s2, &self.stream);
            let act_ptr = raw_device_ptr(&self.save_h_s2, &self.stream);
            let n_relu = (b * sh2) as i32;
            let blocks = ((b * sh2 + 255) / 256) as u32;
            unsafe {
                self.stream
                    .launch_builder(&self.relu_mask_kernel)
                    .arg(&d_ptr)
                    .arg(&act_ptr)
                    .arg(&n_relu)
                    .launch(LaunchConfig {
                        grid_dim: (blocks, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| MLError::ModelError(format!("relu_mask ens_h_s2: {e}")))?;
            }
        }

        // ── 6. Backward FC layer 2: h_s1 → h_s2 (into SCRATCH) ────────────
        {
            let dy = bw_raw_ptr(&self.bw_d_h_s2, &self.stream);
            let x = bw_raw_ptr(&self.save_h_s1, &self.stream);
            let param_sizes = compute_param_sizes(&self.config);
            let w_ptrs = f32_weight_ptrs(&self.params_buf, &param_sizes, &self.stream);
            let w = w_ptrs[2]; // w_s2

            let scratch_base = raw_device_ptr(&self.iqn_trunk_m, &self.stream);
            let f32_u = f32_size as u64;
            let dw = scratch_base + (w_s1_n + b_s1_n) as u64 * f32_u;
            let db = dw + w_s2_n as u64 * f32_u;

            let dx = bw_raw_ptr(&self.bw_d_h_s1, &self.stream);

            self.cublas_backward.backward_fc_layer(
                &self.stream, dy, x, w, dw, db, dx, sh2, sh1, b,
            )?;
        }

        // ── 7. ReLU mask: d_h_s1 *= (save_h_s1 > 0) ──────────────────────
        {
            let d_ptr = raw_device_ptr(&self.bw_d_h_s1, &self.stream);
            let act_ptr = raw_device_ptr(&self.save_h_s1, &self.stream);
            let n_relu = (b * sh1) as i32;
            let blocks = ((b * sh1 + 255) / 256) as u32;
            unsafe {
                self.stream
                    .launch_builder(&self.relu_mask_kernel)
                    .arg(&d_ptr)
                    .arg(&act_ptr)
                    .arg(&n_relu)
                    .launch(LaunchConfig {
                        grid_dim: (blocks, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| MLError::ModelError(format!("relu_mask ens_h_s1: {e}")))?;
            }
        }

        // ── 8. Backward FC layer 1: states → h_s1 (into SCRATCH) ──────────
        {
            let dy = bw_raw_ptr(&self.bw_d_h_s1, &self.stream);
            let x = bw_raw_ptr(&self.states_buf, &self.stream);
            let param_sizes = compute_param_sizes(&self.config);
            let w_ptrs = f32_weight_ptrs(&self.params_buf, &param_sizes, &self.stream);
            let w = w_ptrs[0]; // w_s1

            let scratch_base = raw_device_ptr(&self.iqn_trunk_m, &self.stream);
            let f32_u = f32_size as u64;
            let dw = scratch_base;
            let db = scratch_base + w_s1_n as u64 * f32_u;

            self.cublas_backward.backward_fc_layer(
                &self.stream, dy, x, w, dw, db, 0, sh1, sd, b,
            )?;
        }

        // ── 9. SAXPY: grad_buf[trunk] += scale * scratch[trunk] ────────────
        {
            let grad_ptr = raw_device_ptr(&self.grad_buf, &self.stream);
            let scratch_ptr = raw_device_ptr(&self.iqn_trunk_m, &self.stream);
            let n_i32 = trunk_grad_total as i32;
            let blocks = ((trunk_grad_total + 255) / 256) as u32;
            unsafe {
                self.stream
                    .launch_builder(&self.saxpy_kernel)
                    .arg(&grad_ptr)
                    .arg(&scratch_ptr)
                    .arg(&scale)
                    .arg(&n_i32)
                    .launch(LaunchConfig {
                        grid_dim: (blocks, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| MLError::ModelError(format!("ensemble trunk SAXPY: {e}")))?;
            }
        }

        Ok(())
    }

    /// Apply regime-adaptive scaling to td_errors before PER priority update.
    ///
    /// Computes mean ADX/CUSUM from the batch states as the "current regime" proxy,
    /// then scales td_errors by Gaussian similarity: samples from similar regimes get
    /// higher priority, adapting replay to the current market conditions.
    ///
    /// GPU-only: reads from states_buf, writes to td_errors_buf.
    /// GPU-only regime-adaptive PER scaling. Reads target ADX/CUSUM from
    /// the first sample in states_buf directly in the kernel — zero CPU readback.
    pub fn regime_scale_td_errors(&self) -> Result<(), MLError> {
        let bs = self.config.batch_size as i32;
        let blocks = ((self.config.batch_size + 255) / 256) as u32;

        let _evt_guard = EventTrackingGuard::new(self.stream.context());

        let td_ptr = raw_device_ptr(&self.td_errors_buf, &self.stream);
        let states_ptr = raw_device_ptr(&self.states_buf, &self.stream);

        unsafe {
            self.stream
                .launch_builder(&self.regime_scale_kernel)
                .arg(&td_ptr)
                .arg(&states_ptr)
                .arg(&bs)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("regime_scale_kernel: {e}")))?;
        }

        Ok(())
    }

    /// Apply CQL (Conservative Q-Learning) gradient injection.
    ///
    /// Computes the CQL penalty gradient w.r.t. the C51 logits and runs a
    /// second cuBLAS backward pass to produce parameter gradients, which are
    /// accumulated into `grad_buf` (beta=1.0) on top of the C51 gradients.
    ///
    /// Called between `graph_forward` and `graph_adam` in `FusedTrainingCtx`.
    /// Returns `true` if CQL was applied, `false` if skipped (disabled or no kernel).
    pub fn apply_cql_gradient(
        &mut self,
    ) -> Result<bool, MLError> {
        let cql_kernel = match &self.cql_logit_grad_kernel {
            Some(k) => k.clone(),
            None => return Ok(false),
        };
        if self.config.cql_alpha <= 0.0 {
            return Ok(false);
        }

        let b = self.config.batch_size;
        let na = self.config.num_atoms;
        let b0 = self.config.branch_0_size;
        let b1 = self.config.branch_1_size;
        let b2 = self.config.branch_2_size;
        let total_actions = b0 + b1 + b2;
        let _total_branch_atoms = total_actions * na;

        let _evt_guard = EventTrackingGuard::new(self.stream.context());

        // Step 1: Compute CQL logit gradients via kernel
        let v_logits_ptr = raw_device_ptr(&self.on_v_logits_buf, &self.stream);
        let adv_logits_ptr = raw_device_ptr(&self.on_b_logits_buf, &self.stream);
        let actions_ptr = raw_device_ptr_i32(&self.actions_buf, &self.stream);
        let d_v_ptr = raw_device_ptr(&self.cql_d_value_logits, &self.stream);
        let d_adv_ptr = raw_device_ptr(&self.cql_d_adv_logits, &self.stream);

        let cql_alpha = self.config.cql_alpha;
        let n_i32 = b as i32;
        let na_i32 = na as i32;
        let b0_i32 = b0 as i32;
        let b1_i32 = b1 as i32;
        let b2_i32 = b2 as i32;
        let v_min = self.config.v_min;
        let v_max = self.config.v_max;

        let blocks = ((b + 255) / 256) as u32;
        unsafe {
            self.stream
                .launch_builder(&cql_kernel)
                .arg(&v_logits_ptr)
                .arg(&adv_logits_ptr)
                .arg(&actions_ptr)
                .arg(&d_v_ptr)
                .arg(&d_adv_ptr)
                .arg(&cql_alpha)
                .arg(&n_i32)
                .arg(&na_i32)
                .arg(&b0_i32)
                .arg(&b1_i32)
                .arg(&b2_i32)
                .arg(&v_min)
                .arg(&v_max)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("cql_logit_grad_kernel: {e}")))?;
        }

        // Step 2: SAXPY — add CQL logit gradients to d_value_logits_buf and d_adv_logits_buf
        // These are the same buffers that already hold C51's logit gradients.
        // After this, when we run cuBLAS backward, the backward pass sees
        // combined C51+CQL logit gradients and produces combined parameter gradients.
        // However, backward_full uses beta=1.0 accumulation, so calling it again
        // would ADD to grad_buf (doubling the C51 contribution).
        //
        // Instead, we SAXPY the CQL logit gradients into the d_logits buffers and
        // then rely on the fact that the NEXT graph_forward will use these updated
        // logit gradients. But that's for the next step, not this one.
        //
        // For the current step, the correct approach is to add CQL gradients
        // directly to grad_buf using a separate cuBLAS backward pass with
        // cql_d_value_logits and cql_d_adv_logits as input.
        // This is the same pattern as apply_iqn_trunk_gradient.

        // Extract weight pointers for cuBLAS backward
        let param_sizes = compute_param_sizes(&self.config);
        let w_ptrs = f32_weight_ptrs(&self.params_buf, &param_sizes, &self.stream);

        // Construct d_adv_logits pointers per branch
        let f32_size = std::mem::size_of::<f32>();
        let d_adv_base = d_adv_ptr;
        let d_adv_ptrs = [
            d_adv_base,
            d_adv_base + (b0 * na * f32_size) as u64,
            d_adv_base + ((b0 + b1) * na * f32_size) as u64,
        ];

        // Saved activations from the forward pass (still valid)
        let states_ptr_fw = raw_device_ptr(&self.states_buf, &self.stream);
        let h_s1_ptr = raw_device_ptr(&self.save_h_s1, &self.stream);
        let h_s2_ptr = raw_device_ptr(&self.save_h_s2, &self.stream);
        let h_v_ptr = raw_device_ptr(&self.save_h_v, &self.stream);
        let h_b0_ptr = raw_device_ptr(&self.save_h_b0, &self.stream);
        let h_b1_ptr = raw_device_ptr(&self.save_h_b1, &self.stream);
        let h_b2_ptr = raw_device_ptr(&self.save_h_b2, &self.stream);

        // Scratch buffers for inter-layer gradients (can reuse bw_d_h_* buffers)
        let scratch_d_h_s2 = raw_device_ptr(&self.bw_d_h_s2, &self.stream);
        let scratch_d_h_s1 = raw_device_ptr(&self.bw_d_h_s1, &self.stream);
        let scratch_d_h_v = raw_device_ptr(&self.bw_d_h_v, &self.stream);
        let scratch_d_h_b0 = raw_device_ptr(&self.bw_d_h_b0, &self.stream);
        let scratch_d_h_b1 = raw_device_ptr(&self.bw_d_h_b1, &self.stream);
        let scratch_d_h_b2 = raw_device_ptr(&self.bw_d_h_b2, &self.stream);

        // Run full backward pass with CQL logit gradients.
        // This ACCUMULATES into grad_buf (beta=1.0 in cuBLAS SGEMM),
        // adding CQL parameter gradients on top of C51's.
        self.cublas_backward.backward_full(
            &self.stream,
            d_v_ptr,
            &d_adv_ptrs,
            states_ptr_fw,
            h_s1_ptr, h_s2_ptr, h_v_ptr,
            &[h_b0_ptr, h_b1_ptr, h_b2_ptr],
            &w_ptrs,
            raw_device_ptr(&self.grad_buf, &self.stream),
            scratch_d_h_s2, scratch_d_h_s1, scratch_d_h_v,
            &[scratch_d_h_b0, scratch_d_h_b1, scratch_d_h_b2],
        ).map_err(|e| MLError::ModelError(format!("CQL backward_full: {e}")))?;

        Ok(true)
    }

    /// Whether CQL is enabled and the kernel was compiled.
    pub fn has_cql(&self) -> bool {
        self.cql_logit_grad_kernel.is_some() && self.config.cql_alpha > 0.0
    }

    /// Apply spectral normalization to trunk weight matrices W_s1 and W_s2.
    ///
    /// One step of power iteration per call (standard practice — single step
    /// per training step converges in practice). Constrains ||W||_σ ≤ 1.0,
    /// bounding the network's Lipschitz constant and preventing Q-value explosion.
    ///
    /// Runs outside the CUDA Graph on `params_buf` and `DuelingWeightSet`.
    pub fn apply_spectral_norm(
        &mut self,
        online_dueling: &mut DuelingWeightSet,
    ) -> Result<(), MLError> {
        let sh1 = self.config.shared_h1 as i32;
        let sh2 = self.config.shared_h2 as i32;
        let sd = self.config.state_dim as i32;
        let sigma_max = self.config.spectral_norm_sigma_max;

        let _evt_guard = EventTrackingGuard::new(self.stream.context());

        // Spectral norm on W_s1 [shared_h1, state_dim]
        {
            let (w_ptr, _g) = online_dueling.w_s1.device_ptr(&self.stream);
            let (u_ptr, _g2) = self.spec_u_s1.device_ptr(&self.stream);
            let (v_ptr, _g3) = self.spec_v_s1.device_ptr(&self.stream);
            unsafe {
                self.stream
                    .launch_builder(&self.spectral_norm_kernel)
                    .arg(&w_ptr).arg(&u_ptr).arg(&v_ptr)
                    .arg(&sh1).arg(&sd).arg(&sigma_max)
                    .launch(LaunchConfig {
                        grid_dim: (1, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 512 * 4,
                    })
                    .map_err(|e| MLError::ModelError(format!("spectral_norm W_s1: {e}")))?;
            }
        }

        // Spectral norm on W_s2 [shared_h2, shared_h1]
        {
            let (w_ptr, _g) = online_dueling.w_s2.device_ptr(&self.stream);
            let (u_ptr, _g2) = self.spec_u_s2.device_ptr(&self.stream);
            let (v_ptr, _g3) = self.spec_v_s2.device_ptr(&self.stream);
            unsafe {
                self.stream
                    .launch_builder(&self.spectral_norm_kernel)
                    .arg(&w_ptr).arg(&u_ptr).arg(&v_ptr)
                    .arg(&sh2).arg(&sh1).arg(&sigma_max)
                    .launch(LaunchConfig {
                        grid_dim: (1, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 512 * 4,
                    })
                    .map_err(|e| MLError::ModelError(format!("spectral_norm W_s2: {e}")))?;
            }
        }

        // Sync scaled weights back to params_buf (trunk portion) so the CUDA
        // Graph's Adam sees the spectrally normalized weights. Without this,
        // the graph reads stale unscaled weights from params_buf.
        {
            let f32_sz = std::mem::size_of::<f32>();
            let sh1_u = self.config.shared_h1;
            let sh2_u = self.config.shared_h2;
            let sd_u = self.config.state_dim;
            let params_base = raw_device_ptr(&self.params_buf, &self.stream);

            // w_s1 → params_buf[goff_w_s1]
            let w_s1_src = raw_device_ptr(&online_dueling.w_s1, &self.stream);
            let w_s1_n = sh1_u * sd_u * f32_sz;
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    params_base, w_s1_src, w_s1_n, self.stream.cu_stream()
                ).map_err(|e| MLError::ModelError(format!("spectral sync w_s1: {e}")))?;
            }

            // w_s2 → params_buf[goff_w_s2]
            let w_s2_src = raw_device_ptr(&online_dueling.w_s2, &self.stream);
            // GOFF layout: w_s1(sh1*sd), b_s1(sh1), w_s2(sh2*sh1), b_s2(sh2)
            let goff_w_s2_bytes = (sh1_u * sd_u + sh1_u) * f32_sz;
            let w_s2_n = sh2_u * sh1_u * f32_sz;
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    params_base + goff_w_s2_bytes as u64,
                    w_s2_src, w_s2_n, self.stream.cu_stream()
                ).map_err(|e| MLError::ModelError(format!("spectral sync w_s2: {e}")))?;
            }
        }

        Ok(())
    }

    /// Apply GPU multi-head feature attention to `save_h_s2` (post-graph).
    ///
    /// Runs 4-head self-attention over the trunk output `h_s2 [B, SHARED_H2]`
    /// and writes the attended representation back into `save_h_s2`. The next
    /// CUDA Graph replay then uses the attended features as its trunk activation.
    ///
    /// This is a 1-step lag operation: attention modifies the buffer AFTER the
    /// current graph replay has already consumed the unmodified `save_h_s2`.
    /// The lag is acceptable — equivalent to async multi-task SGD.
    ///
    /// Must be called outside the CUDA Graph. Uses `EventTrackingGuard` for all
    /// buffer pointer extractions to avoid `CUDA_ERROR_INVALID_VALUE` from stale
    /// write events recorded during graph capture.
    pub fn apply_attention_forward(
        &mut self,
        attention: &mut GpuAttention,
        batch_size: usize,
    ) -> Result<(), MLError> {
        let f32_size = std::mem::size_of::<f32>();
        let sh2 = self.config.shared_h2;
        let n_bytes = batch_size * sh2 * f32_size;

        // Sync stream to ensure graph replay completed before touching save_h_s2.
        unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }
        let _ = self.stream.context().check_err();

        // Disable event tracking for all buffer access in this method.
        // After CUDA Graph capture, cudarc's device_ptr() records write events
        // that fail with CUDA_ERROR_INVALID_VALUE on stale captured graph events.
        let _evt_guard = EventTrackingGuard::new(self.stream.context());

        // Run attention forward: save_h_s2 → attention output buffer
        let attn_output = attention.forward(&self.save_h_s2, batch_size)?;

        // DtoD copy: attention output → save_h_s2 (in-place update for next graph replay)
        let src_ptr = raw_device_ptr(attn_output, &self.stream);
        let dst_ptr = raw_device_ptr(&self.save_h_s2, &self.stream);
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr, src_ptr, n_bytes, self.stream.cu_stream()
            ).map_err(|e| MLError::ModelError(format!("attention→save_h_s2 DtoD: {e}")))?;
        }

        Ok(())
    }

    /// Shrink-and-Perturb: prevent loss of plasticity between walk-forward windows.
    ///
    /// `params[i] = alpha * params[i] + (1 - alpha) * noise[i]`
    ///
    /// Where noise is drawn from N(0, sigma). This maintains the network's ability
    /// to learn new patterns across market regime shifts (Ash & Adams 2020).
    ///
    /// Called at walk-forward window boundaries (not per-step).
    /// Also invalidates the CUDA Graph since weights change outside the graph.
    pub fn shrink_and_perturb(
        &mut self,
        alpha: f32,
        sigma: f32,
    ) -> Result<(), MLError> {
        let n = self.total_params;
        let n_i32 = n as i32;
        let blocks = ((n + 255) / 256) as u32;
        let seed = (self.adam_step as u32).wrapping_mul(2654435761);

        // GPU-native: single kernel generates noise + blends in one pass.
        // params[i] = alpha * params[i] + (1-alpha) * N(0, sigma)
        // No CPU noise generation, no HtoD upload.
        let _evt_guard = EventTrackingGuard::new(self.stream.context());
        let params_ptr = raw_device_ptr(&self.params_buf, &self.stream);
        unsafe {
            self.stream
                .launch_builder(&self.shrink_perturb_kernel)
                .arg(&params_ptr)
                .arg(&alpha)
                .arg(&sigma)
                .arg(&n_i32)
                .arg(&seed)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("shrink_perturb_kernel: {e}")))?;
        }
        drop(_evt_guard);

        // Invalidate both CUDA Graphs — weights changed outside captured ops
        self.graph_forward = None;
        self.graph_adam = None;

        info!(alpha, sigma, total_params = n, "Shrink-and-Perturb applied (GPU-native, zero CPU)");
        Ok(())
    }

    /// Create a new fused DQN trainer with pre-allocated GPU buffers.
    ///
    /// Compiles all 4 kernels (forward+loss, backward, grad_norm, adam_update)
    /// from a single NVRTC compilation. All buffers are allocated once (fixed
    /// shapes for CUDA Graph compatibility).
    pub fn new(
        stream: Arc<CudaStream>,
        config: GpuDqnTrainConfig,
    ) -> Result<Self, MLError> {
        let b = config.batch_size;
        let total_params = compute_total_params(&config);

        // Event tracking: kept ENABLED during normal ops (buffer allocation, DtoD).
        // DISABLED only during CUDA Graph capture (in capture_training_graphs).
        // Global disable pollutes the context for other tests/models.

        // Stack size for the training kernel: ~46KB/thread with NUM_ATOMS=51,
        // ~12KB/thread with NUM_ATOMS=11. DIST_SIZE(SCRATCH1_DIM=256) * 4 = 11KB
        // per array. Stack is set once in DQNTrainer::new() (64KB for all kernels).

        // ── Compile 4 utility kernels (grad_norm, adam_update, BF16 converters) ─
        let (grad_norm_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel) =
            compile_training_kernels(&stream, &config)?;

        // ── Compile EMA kernel (standalone — not captured in CUDA Graph) ──
        let ema_kernel = compile_ema_kernel(&stream)?;
        let relu_mask_standalone = compile_relu_mask_standalone(&stream)?;

        // ── Compile PER priority update kernel (standalone) ──
        let per_update_kernel = compile_per_update_kernel(&stream)?;

        // ── Allocate batch input buffers ────────────────────────────
        let states_buf = alloc_f32(&stream, b * config.state_dim, "states")?;
        let next_states_buf = alloc_f32(&stream, b * config.state_dim, "next_states")?;
        let actions_buf = alloc_i32(&stream, b, "actions")?;
        let rewards_buf = alloc_f32(&stream, b, "rewards")?;
        let dones_buf = alloc_f32(&stream, b, "dones")?;
        let is_weights_buf = alloc_f32(&stream, b, "is_weights")?;

        // ── Allocate activation save buffers ────────────────────────
        let num_branches = 3;
        let save_h_s1 = alloc_f32(&stream, b * config.shared_h1, "save_h_s1")?;
        let save_h_s2 = alloc_f32(&stream, b * config.shared_h2, "save_h_s2")?;
        let save_h_v = alloc_f32(&stream, b * config.value_h, "save_h_v")?;
        let save_h_b0 = alloc_f32(&stream, b * config.adv_h, "save_h_b0")?;
        let save_h_b1 = alloc_f32(&stream, b * config.adv_h, "save_h_b1")?;
        let save_h_b2 = alloc_f32(&stream, b * config.adv_h, "save_h_b2")?;
        let save_current_lp = alloc_f32(
            &stream,
            b * num_branches * config.num_atoms,
            "save_current_lp",
        )?;
        let save_projected = alloc_f32(
            &stream,
            b * num_branches * config.num_atoms,
            "save_projected",
        )?;

        // ── Allocate forward output buffers ─────────────────────────
        let per_sample_loss_buf = alloc_f32(&stream, b, "per_sample_loss")?;
        let td_errors_buf = alloc_f32(&stream, b, "td_errors")?;
        let total_loss_buf = alloc_f32(&stream, 1, "total_loss")?;
        let total_actions = config.branch_0_size + config.branch_1_size + config.branch_2_size;
        let q_out_buf = alloc_f32(&stream, b * total_actions, "q_out")?;

        // ── Allocate backward / Adam buffers ────────────────────────
        let grad_buf = alloc_f32(&stream, total_params, "grad_buf")?;
        let params_buf = alloc_f32(&stream, total_params, "params_buf")?;
        let target_params_buf = alloc_f32(&stream, total_params, "target_params_buf")?;
        let m_buf = alloc_f32(&stream, total_params, "adam_m")?;
        let v_buf = alloc_f32(&stream, total_params, "adam_v")?;
        let grad_norm_buf = alloc_f32(&stream, 1, "grad_norm")?;
        let t_buf = alloc_i32(&stream, 1, "adam_t")?;

        // ── Allocate flat BF16 weight mirror buffers ──────────────────
        // One contiguous BF16 buffer per network. Forward kernels read at
        // precomputed GOFF byte offsets. Each segment is padded to even element
        // count so that every base pointer is 4-byte aligned — required by
        // cooperative_load_tile_bf16 which vectorizes as short2 (4-byte loads).
        let param_sizes = compute_param_sizes(&config);
        let mut bf16_goff_byte_offsets = [0_u64; 20];
        let bf16_padded_total: usize;
        {
            let mut offset = 0_u64;
            for i in 0..20 {
                bf16_goff_byte_offsets[i] = offset;
                // Pad to even element count → 4-byte aligned short2 loads
                let padded = ((param_sizes[i] + 1) & !1) as u64;
                offset += padded * (std::mem::size_of::<u16>() as u64);
            }
            bf16_padded_total = (offset / std::mem::size_of::<u16>() as u64) as usize;
        }
        let bf16_params_buf = alloc_u16(&stream, bf16_padded_total, "bf16_params")?;
        let bf16_target_params_buf = alloc_u16(&stream, bf16_padded_total, "bf16_target_params")?;

        // ── Allocate consolidated transfer buffers ─────────────────
        // Upload staging: states + next_states + actions(as f32) + rewards + dones + is_weights
        let upload_staging_len = b * config.state_dim * 2 + b * 4; // 2*B*SD + 4*B
        let upload_staging_buf = alloc_f32(&stream, upload_staging_len, "upload_staging")?;

        // Readback: total_loss(1) + grad_norm(1) + td_errors(B)
        let readback_len = 2 + b;
        let readback_buf = alloc_f32(&stream, readback_len, "readback")?;
        let readback_host = vec![0.0_f32; readback_len];

        // Pre-allocate host staging vec for upload_batch (reused each step)
        let upload_staging_host = vec![0.0_f32; upload_staging_len];

        // PER priority update: batch max accumulator + scalar-only readback
        let mut batch_max_buf = alloc_f32(&stream, 1, "batch_max")?;
        stream.memset_zeros(&mut batch_max_buf)
            .map_err(|e| MLError::ModelError(format!("zero batch_max: {e}")))?;
        let scalars_readback_buf = alloc_f32(&stream, 2, "scalars_readback")?;
        let scalars_readback_host = [0.0_f32; 2];

        // ── Allocate BF16 batch staging buffers ──────────────────────
        // DtoD from GpuBatch's BF16 Candle tensors (states, next_states only)
        // → these staging buffers, then bf16_to_f32_kernel → F32 trainer bufs.
        // Rewards, dones, weights are F32 in GpuBatch and skip BF16 staging.
        let bf16_states_buf = alloc_u16(&stream, b * config.state_dim, "bf16_states")?;
        let bf16_next_states_buf = alloc_u16(&stream, b * config.state_dim, "bf16_next_states")?;

        // ── cuBLAS target network scratch buffers ──────────────────
        // Target forward uses these instead of save_h_* (no grad saves needed).
        let tg_h_s1_scratch = alloc_f32(&stream, b * config.shared_h1, "tg_h_s1")?;
        let tg_h_s2_buf     = alloc_f32(&stream, b * config.shared_h2, "tg_h_s2")?;
        let tg_h_v_scratch  = alloc_f32(&stream, b * config.value_h,   "tg_h_v")?;
        let tg_h_b_scratch  = alloc_f32(&stream, b * config.adv_h,     "tg_h_b")?;
        let tg_v_logits_buf = alloc_f32(&stream, b * config.num_atoms, "tg_v_logits")?;
        let total_branch_atoms =
            config.branch_0_size * config.num_atoms
            + config.branch_1_size * config.num_atoms
            + config.branch_2_size * config.num_atoms;
        let tg_b_logits_buf = alloc_f32(&stream, b * total_branch_atoms, "tg_b_logits")?;

        // ── cuBLAS online logit buffers ────────────────────────────
        let on_v_logits_buf = alloc_f32(&stream, b * config.num_atoms, "on_v_logits")?;
        let on_b_logits_buf = alloc_f32(&stream, b * total_branch_atoms, "on_b_logits")?;

        // ── cuBLAS online-on-next-states buffers (Double DQN action selector) ──
        let on_next_v_logits_buf = alloc_f32(&stream, b * config.num_atoms, "on_next_v_logits")?;
        let on_next_b_logits_buf = alloc_f32(&stream, b * total_branch_atoms, "on_next_b_logits")?;
        let on_next_h_s1_scratch = alloc_f32(&stream, b * config.shared_h1, "on_next_h_s1")?;
        let on_next_h_s2_scratch = alloc_f32(&stream, b * config.shared_h2, "on_next_h_s2")?;
        let on_next_h_v_scratch = alloc_f32(&stream, b * config.value_h, "on_next_h_v")?;
        let on_next_h_b_scratch = alloc_f32(&stream, b * config.adv_h, "on_next_h_b")?;

        // ── Compile standalone C51 loss + gradient kernels (required) ─
        let c51_loss_kernel = compile_c51_loss_kernel(&stream, &config)?;
        let c51_grad_kernel = compile_c51_grad_kernel(&stream, &config)?;
        info!("GpuDqnTrainer: c51_loss + c51_grad kernels compiled");

        // ── Compile MSE loss + gradient kernels (warmup before C51) ─
        let mse_loss_kernel = compile_mse_loss_kernel(&stream, &config)?;
        let mse_grad_kernel = compile_mse_grad_kernel(&stream, &config)?;
        info!("GpuDqnTrainer: mse_loss + mse_grad kernels compiled");

        // ── Compile expected Q-value + stats kernels (validation, not in CUDA Graph) ─
        let expected_q_kernel = compile_expected_q_kernel(&stream)?;
        let q_stats_kernel = compile_q_stats_kernel(&stream)?;
        let q_stats_buf = alloc_f32(&stream, 5, "q_stats")?;
        info!("GpuDqnTrainer: expected_q + q_stats kernels compiled");

        // ── Compile CQL penalty kernel (if enabled) ──────────────────────
        let cql_logit_grad_kernel = if config.cql_alpha > 0.0 {
            match compile_cql_logit_grad_kernel(&stream) {
                Ok(k) => {
                    info!(cql_alpha = config.cql_alpha, "GpuDqnTrainer: CQL logit gradient kernel compiled");
                    Some(k)
                }
                Err(e) => {
                    tracing::warn!("CQL kernel compilation failed (non-fatal): {e}");
                    None
                }
            }
        } else {
            None
        };
        let cql_d_value_logits = alloc_f32(&stream, b * config.num_atoms, "cql_d_value_logits")?;
        let cql_d_adv_logits = alloc_f32(&stream, b * total_branch_atoms, "cql_d_adv_logits")?;

        // ── Gradient output buffers for cuBLAS backward ──────────────
        let d_value_logits_buf = alloc_f32(&stream, b * config.num_atoms, "d_value_logits")?;
        let d_adv_logits_buf = alloc_f32(&stream, b * total_branch_atoms, "d_adv_logits")?;
        // Scratch buffers for blended MSE+C51 loss (MSE grad stored here, then blended)
        let d_value_logits_mse = alloc_f32(&stream, b * config.num_atoms, "d_value_logits_mse")?;
        let d_adv_logits_mse = alloc_f32(&stream, b * total_branch_atoms, "d_adv_logits_mse")?;

        // ── Spectral normalization singular vectors ─────────────────
        // Initialize with random unit vectors for proper power iteration convergence.
        // Uniform 1/√n converges but slowly (aligned with no singular direction).
        // Random vectors almost surely have nonzero projection onto the leading
        // singular direction, giving fast convergence (typically 3-5 iterations).
        let mut rng_state = 0x5DEECE66Du64.wrapping_add(config.batch_size as u64);
        let rand_unit = |n: usize, rng: &mut u64| -> Vec<f32> {
            let mut v = Vec::with_capacity(n);
            for _ in 0..n {
                *rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
                let u = (*rng >> 33) as f32 / (1u64 << 31) as f32 - 0.5;
                v.push(u);
            }
            let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt().max(1e-8);
            for x in v.iter_mut() { *x /= norm; }
            v
        };
        let init_u_s1 = rand_unit(config.shared_h1, &mut rng_state);
        let init_v_s1 = rand_unit(config.state_dim, &mut rng_state);
        let init_u_s2 = rand_unit(config.shared_h2, &mut rng_state);
        let init_v_s2 = rand_unit(config.shared_h1, &mut rng_state);

        // IQN trunk Adam state (separate from C51 Adam — prevents momentum mismatch)
        let trunk_params = config.shared_h1 * config.state_dim + config.shared_h1
            + config.shared_h2 * config.shared_h1 + config.shared_h2;
        let iqn_trunk_m = alloc_f32(&stream, trunk_params, "iqn_trunk_m")?;
        let iqn_trunk_v = alloc_f32(&stream, trunk_params, "iqn_trunk_v")?;
        let iqn_trunk_grad_norm = alloc_f32(&stream, 1, "iqn_trunk_grad_norm")?;
        let iqn_trunk_t_buf = alloc_i32(&stream, 1, "iqn_trunk_t_buf")?;
        let mut spec_u_s1 = alloc_f32(&stream, config.shared_h1, "spec_u_s1")?;
        let mut spec_v_s1 = alloc_f32(&stream, config.state_dim, "spec_v_s1")?;
        let mut spec_u_s2 = alloc_f32(&stream, config.shared_h2, "spec_u_s2")?;
        let mut spec_v_s2 = alloc_f32(&stream, config.shared_h1, "spec_v_s2")?;
        stream.memcpy_htod(&init_u_s1, &mut spec_u_s1)
            .map_err(|e| MLError::ModelError(format!("spec_u_s1 init: {e}")))?;
        stream.memcpy_htod(&init_v_s1, &mut spec_v_s1)
            .map_err(|e| MLError::ModelError(format!("spec_v_s1 init: {e}")))?;
        stream.memcpy_htod(&init_u_s2, &mut spec_u_s2)
            .map_err(|e| MLError::ModelError(format!("spec_u_s2 init: {e}")))?;
        stream.memcpy_htod(&init_v_s2, &mut spec_v_s2)
            .map_err(|e| MLError::ModelError(format!("spec_v_s2 init: {e}")))?;

        // ── Precompute F32 GOFF byte offsets (for cuBLAS weight pointers) ──
        let mut f32_goff_byte_offsets = [0_u64; 20];
        {
            let mut offset = 0_u64;
            for i in 0..20 {
                f32_goff_byte_offsets[i] = offset;
                offset += (param_sizes[i] * std::mem::size_of::<f32>()) as u64;
            }
        }

        // ── Initialize cuBLAS forward context (required) ─────────────
        let cublas_forward = CublasForward::new(
            &stream,
            config.batch_size,
            config.state_dim,
            config.shared_h1,
            config.shared_h2,
            config.value_h,
            config.adv_h,
            config.num_atoms,
            config.branch_0_size,
            config.branch_1_size,
            config.branch_2_size,
        )?;
        info!("GpuDqnTrainer: cuBLAS batched forward initialized");

        // ── Initialize cuBLAS backward context (required) ──────────
        let cublas_backward = CublasBackward::new(&stream, &config)?;
        info!("GpuDqnTrainer: cuBLAS batched backward initialized");

        // ── Backward scratch buffers ────────────────────────────────
        // Pre-allocate inter-layer gradient buffers for the cuBLAS backward
        // pass.  These are separate from the activation saves used in forward.
        let (bw_d_h_s2, bw_d_h_s1, bw_d_h_v, bw_d_h_b0, bw_d_h_b1, bw_d_h_b2) =
            alloc_backward_scratch(&stream, &config)
                .map_err(|e| MLError::ModelError(format!("backward scratch alloc: {e}")))?;

        let batch_bytes = (b * config.state_dim * 2
            + b * 4
            + b * config.shared_h1
            + b * config.shared_h2
            + b * config.value_h
            + b * config.adv_h * 3
            + b * num_branches * config.num_atoms * 2
            + b * 2 + 1)
            * std::mem::size_of::<f32>();
        let optim_bytes = total_params * 5 * std::mem::size_of::<f32>(); // grad + params + target_params + m + v

        info!(
            batch_size = b,
            state_dim = config.state_dim,
            total_params,
            batch_alloc_mb = batch_bytes as f64 / (1024.0 * 1024.0),
            optim_alloc_mb = optim_bytes as f64 / (1024.0 * 1024.0),
            "GpuDqnTrainer: all buffers allocated, 4 kernels compiled"
        );

        let initial_loss_mode = if config.c51_warmup_epochs > 0 { LossMode::Mse } else { LossMode::C51 };
        let initial_c51_alpha = if config.c51_warmup_epochs > 0 { 0.0 } else { 1.0 };

        Ok(Self {
            config,
            stream,
            grad_norm_kernel,
            adam_update_kernel,
            ema_kernel,
            per_update_kernel,
            f32_to_bf16_kernel,
            bf16_to_f32_kernel,
            saxpy_kernel,
            zero_kernel,
            regime_scale_kernel,
            shrink_perturb_kernel: shrink_perturb,
            relu_mask_kernel: relu_mask_standalone,
            spectral_norm_kernel,
            spec_u_s1,
            spec_v_s1,
            spec_u_s2,
            spec_v_s2,
            iqn_trunk_m,
            iqn_trunk_v,
            iqn_trunk_adam_step: 0,
            iqn_trunk_grad_norm,
            iqn_trunk_t_buf,
            trunk_param_count: trunk_params,
            bf16_params_buf,
            bf16_target_params_buf,
            bf16_mirrors_initialized: false,
            bf16_goff_byte_offsets,
            bf16_padded_total,
            states_buf,
            next_states_buf,
            actions_buf,
            rewards_buf,
            dones_buf,
            is_weights_buf,
            save_h_s1,
            save_h_s2,
            save_h_v,
            save_h_b0,
            save_h_b1,
            save_h_b2,
            save_current_lp,
            save_projected,
            per_sample_loss_buf,
            td_errors_buf,
            total_loss_buf,
            q_out_buf,
            grad_buf,
            params_buf,
            target_params_buf,
            m_buf,
            v_buf,
            grad_norm_buf,
            t_buf,
            adam_step: 0,
            total_params,
            params_initialized: false,
            target_params_initialized: false,
            graph_forward: None,
            graph_adam: None,
            upload_staging_buf,
            upload_staging_len,
            readback_buf,
            readback_host,
            upload_staging_host,
            batch_max_buf,
            scalars_readback_buf,
            scalars_readback_host,
            bf16_states_buf,
            bf16_next_states_buf,
            cublas_forward,
            tg_h_s1_scratch,
            tg_h_s2_buf,
            tg_h_v_scratch,
            tg_h_b_scratch,
            tg_v_logits_buf,
            tg_b_logits_buf,
            on_v_logits_buf,
            on_b_logits_buf,
            on_next_v_logits_buf,
            on_next_b_logits_buf,
            on_next_h_s1_scratch,
            on_next_h_s2_scratch,
            on_next_h_v_scratch,
            on_next_h_b_scratch,
            c51_loss_kernel,
            c51_grad_kernel,
            mse_loss_kernel,
            mse_grad_kernel,
            loss_mode: initial_loss_mode,
            c51_alpha: initial_c51_alpha,
            last_captured_loss_mode: None,
            d_value_logits_buf,
            d_adv_logits_buf,
            d_value_logits_mse,
            d_adv_logits_mse,
            f32_goff_byte_offsets,
            cublas_backward,
            bw_d_h_s2,
            bw_d_h_s1,
            bw_d_h_v,
            bw_d_h_b0,
            bw_d_h_b1,
            bw_d_h_b2,
            expected_q_kernel,
            q_stats_kernel,
            q_stats_buf,
            cql_logit_grad_kernel,
            cql_d_value_logits,
            cql_d_adv_logits,
        })
    }

    /// Reference to the bf16→f32 conversion kernel.
    ///
    /// Exposed so that callers (e.g., `FusedTrainingCtx`) can convert BF16
    /// GpuBatch tensors to F32 CudaSlice buffers without creating Candle
    /// temporary tensors.
    pub fn bf16_to_f32_kernel(&self) -> &CudaFunction {
        &self.bf16_to_f32_kernel
    }

    /// Reference to the trainer's forked CudaStream.
    pub fn stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }

    // ═══════════════════════════════════════════════════════════════════
    // BF16 weight mirror management
    // ═══════════════════════════════════════════════════════════════════

    /// Ensure flat BF16 weight mirrors are synced from F32 originals.
    ///
    /// Called lazily on first `train_step()` or `forward_loss()`. Performs
    /// initial F32 → BF16 conversion via 2 kernel launches (1 online + 1 target)
    /// over the flat parameter buffers (same GOFF_* layout).
    fn ensure_bf16_mirrors(
        &mut self,
        _online_d: &DuelingWeightSet,
        _online_b: &BranchingWeightSet,
        _target_d: &DuelingWeightSet,
        _target_b: &BranchingWeightSet,
    ) -> Result<(), MLError> {
        if self.bf16_mirrors_initialized {
            return Ok(());
        }

        // Per-segment F32 → BF16 at padded offsets (4-byte aligned short2 loads)
        self.launch_bf16_convert_online()?;
        self.launch_bf16_convert_target()?;

        self.bf16_mirrors_initialized = true;

        info!(
            total_params = self.total_params,
            bf16_padded_total = self.bf16_padded_total,
            "GpuDqnTrainer: BF16 weight mirrors synced (20 seg launches × 2 networks)"
        );
        Ok(())
    }

    /// Per-segment F32 → BF16 conversion: f32_buf → bf16_buf at padded offsets.
    ///
    /// 20 small kernel launches (one per weight tensor). Each segment is written
    /// at its padded BF16 byte offset so all base pointers are 4-byte aligned
    /// for short2 vectorized loads in the forward kernel.
    fn launch_segmented_bf16_convert(
        &self,
        f32_buf: &CudaSlice<f32>,
        bf16_buf: &CudaSlice<u16>,
    ) -> Result<(), MLError> {
        let param_sizes = compute_param_sizes(&self.config);
        let f32_base = raw_device_ptr(f32_buf, &self.stream);
        let bf16_base = raw_device_ptr_u16(bf16_buf, &self.stream);

        let mut f32_byte_offset = 0_u64;
        for i in 0..20 {
            let n = param_sizes[i];
            if n == 0 { continue; }
            let f32_ptr = f32_base + f32_byte_offset;
            let bf16_ptr = bf16_base + self.bf16_goff_byte_offsets[i];
            let n_i32 = n as i32;
            let blocks = ((n + 255) / 256) as u32;
            unsafe {
                self.stream
                    .launch_builder(&self.f32_to_bf16_kernel)
                    .arg(&f32_ptr)
                    .arg(&bf16_ptr)
                    .arg(&n_i32)
                    .launch(LaunchConfig {
                        grid_dim: (blocks, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| MLError::ModelError(format!("f32_to_bf16 seg[{i}]: {e}")))?;
            }
            f32_byte_offset += (n as u64) * (std::mem::size_of::<f32>() as u64);
        }
        Ok(())
    }

    /// Segmented F32 → BF16: params_buf → bf16_params_buf (online weights).
    fn launch_bf16_convert_online(&self) -> Result<(), MLError> {
        self.launch_segmented_bf16_convert(&self.params_buf, &self.bf16_params_buf)
    }

    /// Segmented F32 → BF16: target_params_buf → bf16_target_params_buf (target weights).
    fn launch_bf16_convert_target(&self) -> Result<(), MLError> {
        self.launch_segmented_bf16_convert(&self.target_params_buf, &self.bf16_target_params_buf)
    }

    /// Compute 20 raw device pointers into a flat BF16 buffer at GOFF_* offsets.
    ///
    /// Returns `[ptr_w_s1, ptr_b_s1, ptr_w_s2, ..., ptr_w_bu2, ptr_b_bu2]`.
    /// Each pointer is the base of the flat buffer + the precomputed byte offset.
    fn bf16_weight_ptrs(&self, bf16_buf: &CudaSlice<u16>) -> [u64; 20] {
        let base = raw_device_ptr_u16(bf16_buf, &self.stream);
        let mut ptrs = [0_u64; 20];
        for i in 0..20 {
            ptrs[i] = base + self.bf16_goff_byte_offsets[i];
        }
        ptrs
    }

    /// Sync online BF16 mirrors from the flat F32 `params_buf`.
    ///
    /// Called inside the CUDA Graph after `unflatten_online_weights()` so that
    /// the next graph replay's forward kernel reads updated BF16 weights.
    /// 20 per-segment launches with padded offsets for 4-byte BF16 alignment.
    fn sync_online_bf16(
        &self,
        _online_d: &DuelingWeightSet,
        _online_b: &BranchingWeightSet,
    ) -> Result<(), MLError> {
        self.launch_bf16_convert_online()
    }

    /// Sync target BF16 mirrors from the flat F32 `target_params_buf`.
    ///
    /// Called after `target_ema_update()` so that the next forward kernel
    /// reads updated BF16 target weights. 20 per-segment launches with padded offsets.
    fn sync_target_bf16(&self) -> Result<(), MLError> {
        self.launch_bf16_convert_target()
    }

    // ═══════════════════════════════════════════════════════════════════
    // Full training step (CUDA Graph — capture once, replay many)
    // ═══════════════════════════════════════════════════════════════════

    /// Run a complete fused training step: forward+loss → backward → Adam.
    ///
    /// On the first call (or after `invalidate_training_graph()`), captures two
    /// CUDA Graphs: `graph_forward` (zero→backward) and `graph_adam` (grad_norm→unflatten).
    /// Subsequent calls replay both graphs with zero kernel-launch overhead.
    ///
    /// Per-step host work (OUTSIDE the graphs):
    /// - Upload batch data (host pointers change each step)
    /// - Update Adam step counter `t_buf`
    /// - Download results (loss, td_errors, grad_norm)
    ///
    /// Captured in `graph_forward`:
    /// - Zero accumulators (memset_zeros)
    /// - cuBLAS forward (3 passes)
    /// - C51 loss + gradient kernels
    /// - cuBLAS backward
    ///
    /// Captured in `graph_adam`:
    /// - Gradient norm reduction
    /// - Adam update kernel
    /// - Unflatten d2d copies (params_buf → individual weight tensors)
    #[allow(clippy::too_many_arguments)]
    pub fn train_step(
        &mut self,
        states: &[f32],
        next_states: &[f32],
        actions: &[i32],
        rewards: &[f32],
        dones: &[f32],
        is_weights: &[f32],
        online_dueling: &DuelingWeightSet,
        online_branching: &BranchingWeightSet,
        _target_dueling: &DuelingWeightSet,
        _target_branching: &BranchingWeightSet,
    ) -> Result<FusedTrainResult, MLError> {
        let b = self.config.batch_size;
        let sd = self.config.state_dim;

        // ── Validate ────────────────────────────────────────────────
        if states.len() != b * sd {
            return Err(MLError::ModelError(format!(
                "states length {} != batch_size({b}) × state_dim({sd})",
                states.len()
            )));
        }
        if actions.len() != b {
            return Err(MLError::ModelError(format!(
                "actions length {} != batch_size({b})",
                actions.len()
            )));
        }

        // ── First call: flatten weights ──────────────────────────────
        if !self.params_initialized {
            self.flatten_online_weights(online_dueling, online_branching)?;
            self.params_initialized = true;
        }

        // ── Upload batch data (OUTSIDE graph — host pointers change) ─
        self.upload_batch(states, next_states, actions, rewards, dones, is_weights)?;

        self.execute_train_and_readback(online_dueling, online_branching)
    }

    /// GPU-direct training step — no CPU roundtrip.
    ///
    /// Takes Candle GPU tensors from `GpuBatch` and copies them to trainer
    /// buffers via DtoD memcpy (skipping the GPU→CPU `.to_vec1()` → CPU→GPU
    /// `memcpy_htod` roundtrip of `train_step()`).
    ///
    /// Same CUDA Graph execution and readback as `train_step()`.
    #[allow(clippy::too_many_arguments)]
    pub fn train_step_gpu(
        &mut self,
        gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch,
        online_dueling: &DuelingWeightSet,
        online_branching: &BranchingWeightSet,
        _target_dueling: &DuelingWeightSet,
        _target_branching: &BranchingWeightSet,
    ) -> Result<FusedTrainScalars, MLError> {
        // ── First call: flatten weights ──────────────────────────────
        if !self.params_initialized {
            self.flatten_online_weights(online_dueling, online_branching)?;
            self.params_initialized = true;
        }

        // ── GPU-direct upload (DtoD, no CPU staging) ─────────────────
        self.upload_batch_gpu(gpu_batch)?;

        // Upload + capture + replay forward ONLY.
        // Adam is deferred to replay_adam_and_readback() so the caller can
        // inject IQN/attention/ensemble gradients into grad_buf first.
        self.adam_step += 1;
        self.stream
            .memcpy_htod(&[self.adam_step], &mut self.t_buf)
            .map_err(|e| MLError::ModelError(format!("HtoD adam_step: {e}")))?;
        if self.graph_forward.is_none() {
            self.capture_training_graphs(online_dueling, online_branching)?;
        }
        self.replay_forward()?;

        // Return placeholder — real scalars come from replay_adam_and_readback()
        Ok(FusedTrainScalars { total_loss: 0.0, grad_norm: 0.0 })
    }

    /// Replay graph_adam + readback scalars. Call AFTER injecting auxiliary
    /// gradients into grad_buf (IQN, attention, ensemble).
    pub fn replay_adam_and_readback(&mut self) -> Result<FusedTrainScalars, MLError> {
        self.replay_adam()?;

        // Sync + disable event tracking for readback (same pattern as all post-graph ops)
        unsafe {
            cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream());
        }
        let _evt_guard = EventTrackingGuard::new(self.stream.context());
        let mut loss_host = [0.0_f32; 1];
        let mut norm_host = [0.0_f32; 1];
        unsafe {
            cudarc::driver::sys::cuMemcpyDtoH_v2(
                loss_host.as_mut_ptr().cast(),
                raw_device_ptr(&self.total_loss_buf, &self.stream), 4,
            );
            cudarc::driver::sys::cuMemcpyDtoH_v2(
                norm_host.as_mut_ptr().cast(),
                raw_device_ptr(&self.grad_norm_buf, &self.stream), 4, // grad_norm_buf, NOT scalars_readback_buf
            );
        }
        Ok(FusedTrainScalars {
            total_loss: loss_host[0],
            grad_norm: norm_host[0].sqrt(), // sqrt: grad_norm_buf stores sum_of_squares
        })
    }

    /// Common post-upload execution: Adam step → CUDA Graph → readback.
    ///
    /// Shared by `train_step()` (CPU upload path) and `train_step_gpu()`
    /// (GPU-direct DtoD path). Everything after batch data lands in the
    /// trainer's CudaSlice buffers is identical.
    ///
    /// Graph split:
    /// - `graph_forward`: zero → cuBLAS forward → C51 loss → C51 grad → cuBLAS backward
    /// - (injection point: external code adds IQN/attention/ensemble gradients to grad_buf)
    /// - `graph_adam`: grad_norm → Adam → unflatten
    fn execute_train_and_readback(
        &mut self,
        online_dueling: &DuelingWeightSet,
        online_branching: &BranchingWeightSet,
    ) -> Result<FusedTrainResult, MLError> {
        let b = self.config.batch_size;

        // ── Update Adam step counter on device (OUTSIDE graph) ───────
        self.adam_step += 1;
        self.stream
            .memcpy_htod(&[self.adam_step], &mut self.t_buf)
            .map_err(|e| MLError::ModelError(format!("HtoD adam_step: {e}")))?;

        // ── Execute training step (split graph: forward then adam) ──
        if self.graph_forward.is_none() {
            self.capture_training_graphs(online_dueling, online_branching)?;
        }
        self.replay_forward()?;
        // --- INJECTION POINT: caller should inject auxiliary gradients here ---
        self.replay_adam()?;

        // ── Consolidate readback: gather 3 GPU buffers → 1 readback buf ──
        // Layout: [total_loss(1) | grad_norm(1) | td_errors(B)]
        // D2D copies are async on the same stream — no host sync until the
        // single memcpy_dtoh at the end.
        let readback_base = raw_device_ptr(&self.readback_buf, &self.stream);
        let f32_bytes = std::mem::size_of::<f32>();

        // total_loss → readback[0]
        let loss_src = raw_device_ptr(&self.total_loss_buf, &self.stream);
        dtod_copy(readback_base, loss_src, f32_bytes, &self.stream, 0, "readback_gather")?;

        // grad_norm → readback[1]
        let norm_src = raw_device_ptr(&self.grad_norm_buf, &self.stream);
        dtod_copy(readback_base + f32_bytes as u64, norm_src, f32_bytes, &self.stream, 1, "readback_gather")?;

        // td_errors → readback[2..2+B]
        let td_src = raw_device_ptr(&self.td_errors_buf, &self.stream);
        let td_bytes = b * f32_bytes;
        dtod_copy(readback_base + (2 * f32_bytes) as u64, td_src, td_bytes, &self.stream, 2, "readback_gather")?;

        // ── Single DtoH transfer ──────────────────────────────────────
        self.stream
            .memcpy_dtoh(&self.readback_buf, &mut self.readback_host) // gpu-exit: loss+grad+td_errors
            .map_err(|e| MLError::ModelError(format!("DtoH readback: {e}")))?;

        // ── Unpack on CPU ─────────────────────────────────────────────
        let total_loss = self.readback_host[0];
        let grad_norm_sq = self.readback_host[1];
        let td_errors = self.readback_host[2..2 + b].to_vec(); // cpu-side host-buf slice

        Ok(FusedTrainResult {
            total_loss,
            td_errors,
            grad_norm: grad_norm_sq.sqrt(),
        })
    }

    /// Replay graph_forward only (zero → forward → loss → grad → backward).
    /// After this call, grad_buf contains C51 gradients. The caller can inject
    /// auxiliary gradients (IQN, attention, ensemble) before calling replay_adam().
    pub fn replay_forward(&self) -> Result<(), MLError> {
        if let Some(ref graph) = self.graph_forward {
            graph.0.launch().map_err(|e| {
                MLError::ModelError(format!("CUDA graph_forward replay: {e}"))
            })?;
        }
        Ok(())
    }

    /// Replay graph_adam only (grad_norm → Adam → unflatten).
    /// Must be called AFTER replay_forward() and any auxiliary gradient injection.
    pub fn replay_adam(&self) -> Result<(), MLError> {
        if let Some(ref graph) = self.graph_adam {
            graph.0.launch().map_err(|e| {
                MLError::ModelError(format!("CUDA graph_adam replay: {e}"))
            })?;
        }
        Ok(())
    }

    /// Scalar-only readback path: only total_loss + grad_norm (8 bytes DtoH).
    ///
    /// TD errors stay on GPU in `td_errors_buf` — PER priority update is done
    /// by `update_priorities_cuda()` which reads them directly via device pointer.
    fn execute_train_scalars_only(
        &mut self,
        online_dueling: &DuelingWeightSet,
        online_branching: &BranchingWeightSet,
    ) -> Result<FusedTrainScalars, MLError> {
        // ── Update Adam step counter on device (OUTSIDE graph) ───────
        self.adam_step += 1;
        self.stream
            .memcpy_htod(&[self.adam_step], &mut self.t_buf)
            .map_err(|e| MLError::ModelError(format!("HtoD adam_step: {e}")))?;

        // ── Execute training step (split graph: forward then adam) ──
        if self.graph_forward.is_none() {
            self.capture_training_graphs(online_dueling, online_branching)?;
        }
        self.replay_forward()?;
        // --- INJECTION POINT: caller should inject auxiliary gradients here ---
        self.replay_adam()?;

        // ── Raw sync + readback (bypasses stale events from graph capture) ──
        unsafe {
            cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream());
        }
        let mut loss_host = [0.0_f32; 1];
        let mut norm_host = [0.0_f32; 1];
        unsafe {
            cudarc::driver::sys::cuMemcpyDtoH_v2(
                loss_host.as_mut_ptr().cast(),
                raw_device_ptr(&self.total_loss_buf, &self.stream), 4,
            );
            cudarc::driver::sys::cuMemcpyDtoH_v2(
                norm_host.as_mut_ptr().cast(),
                raw_device_ptr(&self.grad_norm_buf, &self.stream), 4,
            );
        } // gpu-exit: 2 scalar readbacks (8 bytes)
        self.scalars_readback_host = [loss_host[0], norm_host[0]];

        Ok(FusedTrainScalars {
            total_loss: self.scalars_readback_host[0],
            grad_norm: self.scalars_readback_host[1].sqrt(),
        })
    }

    /// Launch GPU-native PER priority update kernel.
    ///
    /// Reads `td_errors_buf` (populated by training graph) and scatter-writes
    /// new priorities directly to the replay buffer's priority array.
    /// Accumulates batch max via atomicMax for epoch-boundary flush.
    ///
    /// Zero CPU readback, single CUDA kernel launch (replaces 7-8 Candle ops).
    ///
    /// # Arguments
    /// * `indices` — CudaSlice<u32> [B] buffer indices from PER sampling
    /// * `priorities` — CudaSlice<f32> [capacity] priority array in GpuReplayBuffer
    /// * `alpha` — PER alpha exponent
    /// * `epsilon` — PER epsilon floor
    pub fn update_priorities_cuda(
        &mut self,
        indices: &CudaSlice<u32>,
        priorities: &CudaSlice<f32>,
        alpha: f32,
        epsilon: f32,
    ) -> Result<(), MLError> {
        let b = self.config.batch_size as i32;
        let capacity = priorities.len() as i32;
        let blocks = ((self.config.batch_size + 255) / 256) as u32;
        let launch_cfg = LaunchConfig {
            grid_dim: (blocks, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };

        unsafe {
            self.stream
                .launch_builder(&self.per_update_kernel)
                .arg(&self.td_errors_buf)
                .arg(indices)
                .arg(priorities)
                .arg(&self.batch_max_buf)
                .arg(&alpha)
                .arg(&epsilon)
                .arg(&b)
                .arg(&capacity)
                .launch(launch_cfg)
                .map_err(|e| MLError::ModelError(format!("per_update_priorities_kernel: {e}")))?;
        }

        Ok(())
    }

    /// Read back the accumulated batch max priority and reset for next epoch.
    ///
    /// Single scalar DtoH (4 bytes) — called once per epoch, not per step.
    /// Returns 0.0 if no priorities were updated this epoch.
    pub fn batch_max_readback_and_reset(&mut self) -> Result<f32, MLError> {
        let mut host = [0.0_f32; 1];
        self.stream
            .memcpy_dtoh(&self.batch_max_buf, &mut host) // gpu-exit: 4B scalar batch_max
            .map_err(|e| MLError::ModelError(format!("DtoH batch_max: {e}")))?;
        // Reset to 0 for next epoch
        self.stream
            .memset_zeros(&mut self.batch_max_buf)
            .map_err(|e| MLError::ModelError(format!("zero batch_max: {e}")))?;
        Ok(host[0])
    }

    /// Reference to the td_errors buffer on GPU.
    ///
    /// Shape: `[B]` f32 — per-sample TD errors from the last training step.
    /// Valid after `train_step_gpu()` or `train_step()` has been called.
    /// Used by the PER priority update kernel to avoid DtoH readback.
    pub fn td_errors_buf(&self) -> &CudaSlice<f32> {
        &self.td_errors_buf
    }

    /// Reference to the online value logits buffer from the last cuBLAS forward.
    ///
    /// Shape: `[B, num_atoms]`. Used by ensemble heads to get head-0 C51 logits
    /// for diversity loss computation without re-running the forward pass.
    pub fn on_v_logits_buf(&self) -> &CudaSlice<f32> {
        &self.on_v_logits_buf
    }

    /// Reference to the target h_v scratch buffer (reused for ensemble head intermediate).
    ///
    /// Shape: `[B, VALUE_H]`. Safe to reuse between the CUDA Graph step and the next
    /// train_step call (not used during the post-graph ensemble forward phase).
    pub fn tg_h_v_scratch_ptr(&self) -> &CudaSlice<f32> {
        &self.tg_h_v_scratch
    }

    /// Run value head forward for an ensemble extra head.
    ///
    /// h_s2 → W_v1 → ReLU → W_v2 → v_logits
    ///
    /// Uses shared trunk activations (save_h_s2) with per-head value weights.
    /// Writes v_logits to `out_ptr` (offset into ensemble_logits_buf).
    pub fn forward_value_head_for_ensemble(
        &self,
        head_w_v1: u64, head_b_v1: u64,
        head_w_v2: u64, head_b_v2: u64,
        v_logits_out: u64,
        batch_size: usize,
    ) -> Result<(), MLError> {
        let h_s2_ptr = raw_device_ptr(&self.save_h_s2, &self.stream);
        let h_v_ptr = raw_device_ptr(&self.tg_h_v_scratch, &self.stream);
        self.cublas_forward.forward_value_head(
            &self.stream,
            h_s2_ptr, head_w_v1, head_b_v1, head_w_v2, head_b_v2,
            h_v_ptr, v_logits_out, batch_size,
        )
    }

    /// Total number of per-branch actions (BRANCH_0 + BRANCH_1 + BRANCH_2).
    // NOTE: forward_loss(), forward_only_q(), launch_forward_only(), launch_forward_loss(),
    // and launch_backward() were removed here. They depended on forward_loss_kernel,
    // forward_only_kernel, and backward_kernel which are dead code (replaced by cuBLAS).
    // See Task 2 for full cleanup context.
    pub fn total_actions(&self) -> usize {
        self.config.branch_0_size + self.config.branch_1_size + self.config.branch_2_size
    }

    /// Reference to the Q-value output buffer from the last `forward_only_q()` call.
    ///
    /// Shape: `[B, total_actions]`. Only valid after `forward_only_q()`.
    pub fn q_out_buf(&self) -> &CudaSlice<f32> {
        &self.q_out_buf
    }

    /// Compute expected Q-values for a batch of states using cuBLAS forward.
    ///
    /// Runs the online network forward pass (cuBLAS SGEMM) and converts
    /// C51 distributional logits to expected Q-values via the `compute_expected_q` kernel.
    /// Does NOT interact with any Candle types.
    ///
    /// This is an ad-hoc forward pass for validation — it is NOT captured in the CUDA Graph.
    /// It reuses the existing `on_v_logits_buf` and `on_b_logits_buf` scratch buffers,
    /// and writes expected Q-values into `q_out_buf`.
    ///
    /// Input:  `states` — `CudaSlice<f32>` of shape `[batch_size, state_dim]`
    /// Output: reference to `q_out_buf` — `[batch_size, total_actions(11)]`
    pub fn compute_q_values(
        &self,
        states: &CudaSlice<f32>,
        batch_size: usize,
    ) -> Result<&CudaSlice<f32>, MLError> {
        if batch_size > self.config.batch_size {
            return Err(MLError::ModelError(format!(
                "compute_q_values: batch_size {batch_size} exceeds trainer batch_size {}",
                self.config.batch_size
            )));
        }

        // Step 1: cuBLAS forward pass — online network on the provided states.
        // Writes logits into on_v_logits_buf [batch_size, NA] and on_b_logits_buf [batch_size, (B0+B1+B2)*NA].
        // Activation buffers (save_h_*) are used as scratch here (no backward will use them).
        let param_sizes = compute_param_sizes(&self.config);
        let on_w_ptrs = f32_weight_ptrs(&self.params_buf, &param_sizes, &self.stream);

        self.cublas_forward.forward_online(
            &self.stream,
            states,
            &on_w_ptrs,
            &self.save_h_s1,
            &self.save_h_s2,
            &self.save_h_v,
            &self.save_h_b0,
            &self.save_h_b1,
            &self.save_h_b2,
            &self.on_v_logits_buf,
            &self.on_b_logits_buf,
        )?;

        // Step 2: compute_expected_q kernel — logits → expected Q-values.
        // Writes into q_out_buf [batch_size, total_actions].
        let n = batch_size as i32;
        let na = self.config.num_atoms as i32;
        let b0 = self.config.branch_0_size as i32;
        let b1 = self.config.branch_1_size as i32;
        let b2 = self.config.branch_2_size as i32;
        let v_min = self.config.v_min;
        let v_max = self.config.v_max;

        let total_threads = batch_size;
        let block_dim = 256_u32;
        let grid_dim = ((total_threads as u32 + block_dim - 1) / block_dim).max(1);

        unsafe {
            self.stream
                .launch_builder(&self.expected_q_kernel)
                .arg(&self.on_v_logits_buf)
                .arg(&self.on_b_logits_buf)
                .arg(&self.q_out_buf)
                .arg(&n)
                .arg(&na)
                .arg(&b0)
                .arg(&b1)
                .arg(&b2)
                .arg(&v_min)
                .arg(&v_max)
                .launch(LaunchConfig {
                    grid_dim: (grid_dim, 1, 1),
                    block_dim: (block_dim, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("compute_expected_q kernel: {e}")))?;
        }

        Ok(&self.q_out_buf)
    }

    /// Compute Q-value statistics entirely on GPU — zero CPU reduction.
    ///
    /// Runs cuBLAS forward + expected_q + q_stats_kernel → returns 5 scalars:
    /// `(avg_max_q, q_min, q_max, q_mean, q_variance)`.
    /// Only 20 bytes (5 f32) are read back to host.
    pub fn compute_q_stats(
        &mut self,
        states: &CudaSlice<f32>,
        batch_size: usize,
    ) -> Result<QValueStatsResult, MLError> {
        self.compute_q_values(states, batch_size)?;

        // Zero stats buf then launch reduction
        self.stream.memset_zeros(&mut self.q_stats_buf)
            .map_err(|e| MLError::ModelError(format!("zero q_stats: {e}")))?;

        let total_actions = self.total_actions() as i32;
        let n = batch_size as i32;

        unsafe {
            self.stream
                .launch_builder(&self.q_stats_kernel)
                .arg(&self.q_out_buf)
                .arg(&self.q_stats_buf)
                .arg(&n)
                .arg(&total_actions)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("q_stats_kernel: {e}")))?;
        }

        // Single 20-byte readback: [avg_max_q, q_min, q_max, q_mean, q_var]
        let mut host = [0.0_f32; 5];
        self.stream.memcpy_dtoh(&self.q_stats_buf, &mut host)
            .map_err(|e| MLError::ModelError(format!("q_stats DtoH: {e}")))?;

        Ok(QValueStatsResult {
            avg_max_q: host[0] as f64,
            q_min: host[1],
            q_max: host[2],
            q_mean: host[3],
            q_variance: host[4],
        })
    }

    // ═══════════════════════════════════════════════════════════════════
    // CUDA Graph capture and invalidation
    // ═══════════════════════════════════════════════════════════════════

    /// Capture two CUDA Graphs: forward (zero→backward) and adam (grad_norm→unflatten).
    ///
    /// Called on the first `train_step()` or after `invalidate_training_graph()`.
    /// The split allows external code to inject auxiliary gradients (IQN, attention,
    /// ensemble) into `grad_buf` between the two graph replays.
    ///
    /// Graph A (`graph_forward`): zero → cuBLAS forward → C51 loss → C51 grad → cuBLAS backward
    /// Graph B (`graph_adam`):    grad_norm → Adam → unflatten (20 d2d copies)
    fn capture_training_graphs(
        &mut self,
        online_d: &DuelingWeightSet,
        online_b: &BranchingWeightSet,
    ) -> Result<(), MLError> {
        // Synchronize the stream before capture to ensure all pending work
        // (BF16 mirror sync, batch upload, adam_step memcpy) is complete.
        self.stream.synchronize()
            .map_err(|e| MLError::ModelError(format!("stream sync before capture: {e}")))?;

        // Disable event tracking during capture — cudarc's device_ptr records
        // CudaEvents which are DISALLOWED inside CUDA Graph capture.
        unsafe { self.stream.context().disable_event_tracking(); }

        // ── Capture graph_forward ──────────────────────────────────────
        let begin_result = self.stream.begin_capture(
            cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL,
        );
        if let Err(e) = begin_result {
            unsafe { self.stream.context().enable_event_tracking(); }
            let _ = self.stream.context().check_err();
            return Err(MLError::ModelError(format!("CUDA graph_forward begin_capture: {e}")));
        }

        let submit_fwd_result = self.submit_forward_ops();

        let graph_fwd_result = self.stream.end_capture(
            cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
        );

        // Check forward submission
        if let Err(e) = submit_fwd_result {
            unsafe { self.stream.context().enable_event_tracking(); }
            let _ = self.stream.context().check_err();
            return Err(e);
        }

        let graph_fwd = graph_fwd_result
            .map_err(|e| {
                unsafe { self.stream.context().enable_event_tracking(); }
                let _ = self.stream.context().check_err();
                MLError::ModelError(format!("CUDA graph_forward end_capture: {e}"))
            })?
            .ok_or_else(|| {
                unsafe { self.stream.context().enable_event_tracking(); }
                let _ = self.stream.context().check_err();
                MLError::ModelError(
                    "CUDA graph_forward capture returned None — stream may not support capture".into()
                )
            })?;

        // ── Capture graph_adam ──────────────────────────────────────────
        let begin_result = self.stream.begin_capture(
            cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL,
        );
        if let Err(e) = begin_result {
            unsafe { self.stream.context().enable_event_tracking(); }
            let _ = self.stream.context().check_err();
            return Err(MLError::ModelError(format!("CUDA graph_adam begin_capture: {e}")));
        }

        let submit_adam_result = self.submit_adam_ops(online_d, online_b);

        let graph_adam_result = self.stream.end_capture(
            cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
        );

        // Re-enable event tracking after both captures and drain stale errors.
        unsafe { self.stream.context().enable_event_tracking(); }
        let _ = self.stream.context().check_err();

        // Propagate adam submission error
        submit_adam_result?;

        let graph_adam = graph_adam_result
            .map_err(|e| MLError::ModelError(format!("CUDA graph_adam end_capture: {e}")))?
            .ok_or_else(|| MLError::ModelError(
                "CUDA graph_adam capture returned None — stream may not support capture".into()
            ))?;

        // Launch both graphs on first capture (warm-up). On subsequent steps,
        // only graph_forward is replayed by train_step_gpu(). The caller then
        // injects auxiliary gradients and calls replay_adam_and_readback().
        graph_fwd.launch().map_err(|e| {
            MLError::ModelError(format!("CUDA graph_forward first launch: {e}"))
        })?;
        graph_adam.launch().map_err(|e| {
            MLError::ModelError(format!("CUDA graph_adam first launch: {e}"))
        })?;

        info!(
            "GpuDqnTrainer: 2 CUDA graphs captured and launched \
             (graph_forward: 5 memsets + forward + loss + grad + backward; \
              graph_adam: grad_norm + adam + 20 d2d unflatten)"
        );
        self.graph_forward = Some(SendSyncGraph(graph_fwd));
        self.graph_adam = Some(SendSyncGraph(graph_adam));
        self.last_captured_loss_mode = Some(self.loss_mode);
        Ok(())
    }

    /// Submit the forward phase ops to the stream (captured into graph_forward).
    ///
    /// Steps: zero accumulators → cuBLAS forward → C51 loss → C51 grad → cuBLAS backward.
    /// After this graph, `grad_buf` contains C51's gradients.
    fn submit_forward_ops(&mut self) -> Result<(), MLError> {
        // ── Zero accumulators (capturable: memset_zeros uses cuMemsetD32Async) ─
        self.stream
            .memset_zeros(&mut self.total_loss_buf)
            .map_err(|e| MLError::ModelError(format!("zero total_loss: {e}")))?;
        self.stream
            .memset_zeros(&mut self.grad_buf)
            .map_err(|e| MLError::ModelError(format!("zero grad_buf: {e}")))?;
        // Zero C51 gradient output buffers (atomicAdd accumulates)
        self.stream
            .memset_zeros(&mut self.d_value_logits_buf)
            .map_err(|e| MLError::ModelError(format!("zero d_value_logits: {e}")))?;
        self.stream
            .memset_zeros(&mut self.d_adv_logits_buf)
            .map_err(|e| MLError::ModelError(format!("zero d_adv_logits: {e}")))?;

        // ── 1. Forward (cuBLAS SGEMM, 3 passes) ──────────────────
        self.launch_cublas_forward()?;

        // ── 2+3. Loss + gradient (blended MSE + C51 via c51_alpha ramp) ─
        // Always run BOTH loss paths. Blend gradients:
        //   grad = (1-α) * MSE_grad + α * C51_grad
        // When α=0 (warmup): pure MSE. When α=1 (converged): pure C51.
        // The kernel sequence is FIXED (both always run) → CUDA Graph compatible.
        //
        // MSE writes to d_value_logits_mse / d_adv_logits_mse (scratch)
        // C51 writes to d_value_logits_buf / d_adv_logits_buf (main)
        // Then SAXPY blends: main = α * main + (1-α) * scratch

        // MSE path → scratch buffers
        self.stream.memset_zeros(&mut self.d_value_logits_mse)
            .map_err(|e| MLError::ModelError(format!("zero d_value_mse: {e}")))?;
        self.stream.memset_zeros(&mut self.d_adv_logits_mse)
            .map_err(|e| MLError::ModelError(format!("zero d_adv_mse: {e}")))?;
        self.launch_mse_loss()?;
        self.launch_mse_grad_to_scratch()?;

        // C51 path → main buffers (already zeroed above)
        self.launch_c51_loss()?;
        self.launch_c51_grad()?;

        // Blend: main = α * C51 + (1-α) * MSE
        // Step 1: main *= α (scale C51 grad)
        // Step 2: main += (1-α) * scratch (add MSE grad)
        {
            let alpha = self.c51_alpha;
            let na = self.config.num_atoms;
            let b = self.config.batch_size;
            let b0 = self.config.branch_0_size;
            let b1 = self.config.branch_1_size;
            let b2 = self.config.branch_2_size;
            let n_val = (b * na) as i32;
            let n_adv = (b * (b0 + b1 + b2) * na) as i32;
            let scale_c51 = alpha - 1.0; // SAXPY: y += (α-1)*y → y *= α
            let scale_mse = 1.0 - alpha;
            let val_blocks = ((b * na + 255) / 256) as u32;
            let adv_blocks = ((b * (b0 + b1 + b2) * na + 255) / 256) as u32;
            let cfg_val = LaunchConfig { grid_dim: (val_blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
            let cfg_adv = LaunchConfig { grid_dim: (adv_blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };

            // Use raw device pointers to avoid &mut/& borrow conflict on same buffer.
            let val_ptr = raw_device_ptr(&self.d_value_logits_buf, &self.stream);
            let val_mse_ptr = raw_device_ptr(&self.d_value_logits_mse, &self.stream);
            let adv_ptr = raw_device_ptr(&self.d_adv_logits_buf, &self.stream);
            let adv_mse_ptr = raw_device_ptr(&self.d_adv_logits_mse, &self.stream);

            unsafe {
                // d_value += (α-1) * d_value → d_value *= α
                self.stream.launch_builder(&self.saxpy_kernel)
                    .arg(&val_ptr).arg(&val_ptr)
                    .arg(&scale_c51).arg(&n_val)
                    .launch(cfg_val).map_err(|e| MLError::ModelError(format!("blend c51 val: {e}")))?;
                // d_value += (1-α) * mse_scratch
                self.stream.launch_builder(&self.saxpy_kernel)
                    .arg(&val_ptr).arg(&val_mse_ptr)
                    .arg(&scale_mse).arg(&n_val)
                    .launch(cfg_val).map_err(|e| MLError::ModelError(format!("blend mse val: {e}")))?;
                // d_adv += (α-1) * d_adv → d_adv *= α
                self.stream.launch_builder(&self.saxpy_kernel)
                    .arg(&adv_ptr).arg(&adv_ptr)
                    .arg(&scale_c51).arg(&n_adv)
                    .launch(cfg_adv).map_err(|e| MLError::ModelError(format!("blend c51 adv: {e}")))?;
                // d_adv += (1-α) * mse_scratch
                self.stream.launch_builder(&self.saxpy_kernel)
                    .arg(&adv_ptr).arg(&adv_mse_ptr)
                    .arg(&scale_mse).arg(&n_adv)
                    .launch(cfg_adv).map_err(|e| MLError::ModelError(format!("blend mse adv: {e}")))?;
            }
        }

        // ── 4. Backward (cuBLAS SGEMM, chain rule through layers) ─
        self.launch_cublas_backward()?;

        Ok(())
    }

    /// Submit the optimizer phase ops to the stream (captured into graph_adam).
    ///
    /// Steps: zero grad_norm → grad_norm → Adam → unflatten.
    /// Reads `grad_buf` which may contain combined C51 + IQN + ensemble gradients.
    fn submit_adam_ops(
        &mut self,
        online_d: &DuelingWeightSet,
        online_b: &BranchingWeightSet,
    ) -> Result<(), MLError> {
        // ── Zero grad_norm accumulator ─────────────────────────────
        self.stream
            .memset_zeros(&mut self.grad_norm_buf)
            .map_err(|e| MLError::ModelError(format!("zero grad_norm: {e}")))?;

        // ── 5. Gradient norm ───────────────────────────────────────
        self.launch_grad_norm()?;

        // ── 6. Adam update ────────────────────────────────────────
        self.launch_adam_update()?;

        // ── 7. Unflatten: params_buf → individual F32 weight tensors ─
        self.unflatten_online_weights(online_d, online_b)?;

        Ok(())
    }

    /// Discard the cached CUDA Graphs and force re-flatten + re-capture.
    ///
    /// Call after:
    /// - Target network EMA update (target weight pointers may change)
    /// - Learning rate schedule change (lr is baked into the graph)
    /// - Any external weight modification
    ///
    /// The next `train_step()` will re-capture fresh graphs.
    pub fn invalidate_training_graph(&mut self) {
        self.graph_forward = None;
        self.graph_adam = None;
        self.last_captured_loss_mode = None;
        self.params_initialized = false;
        // Reset BF16 mirrors flag — they'll be re-synced on next train_step.
        // Padded BF16 buffers are pre-allocated (no reallocation needed); only the
        // content needs refreshing via 20 per-segment f32_to_bf16 launches each.
        self.bf16_mirrors_initialized = false;
    }

    /// Switch the loss mode (MSE warmup vs C51 distributional).
    ///
    /// With blended loss, both MSE and C51 kernels run every step
    /// (CUDA Graph kernel sequence is fixed). This method sets `c51_alpha`
    /// accordingly: Mse → 0.0 (pure MSE), C51 → 1.0 (pure C51).
    /// No graph invalidation needed — alpha is read from host memory each step.
    pub fn set_loss_mode(&mut self, mode: LossMode) {
        if self.loss_mode != mode {
            info!("Loss mode switched to {:?}", mode);
            self.loss_mode = mode;
            match mode {
                LossMode::Mse => self.c51_alpha = 0.0,
                LossMode::C51 => self.c51_alpha = 1.0,
            }
        }
    }

    /// Set the C51 blend factor directly for gradual ramp.
    ///
    /// `alpha = 0.0`: pure MSE loss (warmup).
    /// `alpha = 1.0`: pure C51 distributional loss (converged).
    /// Intermediate values blend: `grad = alpha*C51 + (1-alpha)*MSE`.
    ///
    /// No CUDA Graph invalidation — the kernel sequence is always the same
    /// (both MSE and C51 run every step); only the SAXPY scale factors change.
    pub fn set_c51_alpha(&mut self, alpha: f32) {
        self.c51_alpha = alpha.clamp(0.0, 1.0);
    }

    /// Current loss mode.
    pub fn loss_mode(&self) -> LossMode {
        self.loss_mode
    }

    /// Invalidate cached state after external weight modifications.
    ///
    /// Call after target network EMA update or any manual weight change.
    /// Forces re-flatten on the next `train_step()`.
    pub fn invalidate_params(&mut self) {
        self.invalidate_training_graph();
    }

    // ═══════════════════════════════════════════════════════════════════
    // Batch upload helper
    // ═══════════════════════════════════════════════════════════════════

    /// Upload batch data to pre-allocated GPU buffers via consolidated staging.
    ///
    /// Packs all 6 arrays into a single contiguous host buffer, performs one
    /// `memcpy_htod` to the GPU staging buffer, then scatters to individual
    /// buffers via `memcpy_dtod_async`. This turns 6 PCIe round-trips into 1,
    /// saving ~12-30us per batch on H100 PCIe Gen5.
    ///
    /// Layout: [states(B*SD) | next_states(B*SD) | actions_as_f32(B) | rewards(B) | dones(B) | is_weights(B)]
    fn upload_batch(
        &mut self,
        states: &[f32],
        next_states: &[f32],
        actions: &[i32],
        rewards: &[f32],
        dones: &[f32],
        is_weights: &[f32],
    ) -> Result<(), MLError> {
        let b = self.config.batch_size;
        let sd = self.config.state_dim;

        // ── Pack all data into pre-allocated host buffer (zero malloc) ──
        self.upload_staging_host.clear();
        self.upload_staging_host.extend_from_slice(states);          // B * SD f32s
        self.upload_staging_host.extend_from_slice(next_states);     // B * SD f32s
        // Reinterpret i32 actions as f32 bits (same 4 bytes, no conversion)
        for &a in actions {
            self.upload_staging_host.push(f32::from_bits(a as u32));
        }
        self.upload_staging_host.extend_from_slice(rewards);         // B f32s
        self.upload_staging_host.extend_from_slice(dones);           // B f32s
        self.upload_staging_host.extend_from_slice(is_weights);      // B f32s

        // ── Single HtoD transfer ──────────────────────────────────
        self.stream
            .memcpy_htod(&self.upload_staging_host, &mut self.upload_staging_buf)
            .map_err(|e| MLError::ModelError(format!("HtoD staging: {e}")))?;

        // ── Scatter from staging to individual buffers via DtoD ───
        let staging_base = raw_device_ptr(&self.upload_staging_buf, &self.stream);
        let f32_size = std::mem::size_of::<f32>();
        let mut byte_offset: u64 = 0;

        // states: B * SD elements
        let states_bytes = b * sd * f32_size;
        let states_dst = raw_device_ptr(&self.states_buf, &self.stream);
        dtod_copy(states_dst, staging_base + byte_offset, states_bytes, &self.stream, 0, "upload_scatter")?;
        byte_offset += states_bytes as u64;

        // next_states: B * SD elements
        let next_states_bytes = b * sd * f32_size;
        let next_states_dst = raw_device_ptr(&self.next_states_buf, &self.stream);
        dtod_copy(next_states_dst, staging_base + byte_offset, next_states_bytes, &self.stream, 1, "upload_scatter")?;
        byte_offset += next_states_bytes as u64;

        // actions: B elements (reinterpreted as f32 in staging, copy raw bytes to i32 buf)
        let actions_bytes = b * f32_size; // i32 and f32 are both 4 bytes
        let actions_dst = raw_device_ptr_i32(&self.actions_buf, &self.stream);
        dtod_copy(actions_dst, staging_base + byte_offset, actions_bytes, &self.stream, 2, "upload_scatter")?;
        byte_offset += actions_bytes as u64;

        // rewards: B elements
        let rewards_bytes = b * f32_size;
        let rewards_dst = raw_device_ptr(&self.rewards_buf, &self.stream);
        dtod_copy(rewards_dst, staging_base + byte_offset, rewards_bytes, &self.stream, 3, "upload_scatter")?;
        byte_offset += rewards_bytes as u64;

        // dones: B elements
        let dones_bytes = b * f32_size;
        let dones_dst = raw_device_ptr(&self.dones_buf, &self.stream);
        dtod_copy(dones_dst, staging_base + byte_offset, dones_bytes, &self.stream, 4, "upload_scatter")?;
        byte_offset += dones_bytes as u64;

        // is_weights: B elements
        let is_weights_bytes = b * f32_size;
        let is_weights_dst = raw_device_ptr(&self.is_weights_buf, &self.stream);
        dtod_copy(is_weights_dst, staging_base + byte_offset, is_weights_bytes, &self.stream, 5, "upload_scatter")?;

        Ok(())
    }

    /// Upload batch data from GPU tensors — pure CudaSlice, zero Candle temporaries.
    ///
    /// GpuBatch layout from GpuReplayBuffer:
    /// - states/next_states: BF16 (stored in BF16 ring buffer)
    /// - rewards/dones/weights: F32 (stored in F32 ring buffer / computed as F32)
    /// - actions: U32 (stored in U32 ring buffer)
    ///
    /// For BF16 tensors (states, next_states): DtoD → BF16 staging → bf16_to_f32 kernel → F32 buf.
    /// For F32 tensors (rewards, dones, weights): DtoD → F32 buf directly.
    /// For U32 actions: DtoD → I32 buf (bit-compatible for values in [0, 44]).
    ///
    /// No Candle `to_dtype()` temporaries — eliminates Drop/event conflicts with forked stream.
    fn upload_batch_gpu(
        &mut self,
        gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch,
    ) -> Result<(), MLError> {
        let b = self.config.batch_size;
        let sd = self.config.state_dim;

        // GpuBatch stores GpuTensor (CudaSlice<f32>) for all fields.
        // States may be BF16-typed GpuTensors whose underlying data is u16-reinterpreted f32.
        // Check shape to determine if we have BF16 data that needs conversion.
        let states_data = gpu_batch.states.data();
        let next_states_data = gpu_batch.next_states.data();

        // If states are stored as BF16 (u16 packed in f32 storage with half the elements),
        // use the bf16 path. Otherwise copy f32 directly.
        let states_numel = gpu_batch.states.numel();
        if states_numel == b * sd {
            // F32 states -- direct DtoD copy to states_buf
            dtod_from_slice_f32(states_data, &self.states_buf, b * sd, &self.stream, "states")?;
            dtod_from_slice_f32(next_states_data, &self.next_states_buf, b * sd, &self.stream, "next_states")?;
        } else {
            // BF16 states -- reinterpret f32 storage as u16 and convert
            let states_u16: &CudaSlice<u16> = unsafe {
                &*(states_data as *const CudaSlice<f32> as *const CudaSlice<u16>)
            };
            let next_states_u16: &CudaSlice<u16> = unsafe {
                &*(next_states_data as *const CudaSlice<f32> as *const CudaSlice<u16>)
            };
            dtod_from_bf16(states_u16, &self.bf16_states_buf, b * sd, &self.stream, "states")?;
            dtod_from_bf16(next_states_u16, &self.bf16_next_states_buf, b * sd, &self.stream, "next_states")?;
            launch_bf16_to_f32(&self.bf16_to_f32_kernel, &self.bf16_states_buf, &self.states_buf, b * sd, &self.stream, "states")?;
            launch_bf16_to_f32(&self.bf16_to_f32_kernel, &self.bf16_next_states_buf, &self.next_states_buf, b * sd, &self.stream, "next_states")?;
        }

        // Rewards, dones, IS-weights: already F32 -- direct DtoD
        dtod_from_slice_f32(gpu_batch.rewards.data(), &self.rewards_buf, b, &self.stream, "rewards")?;
        dtod_from_slice_f32(gpu_batch.dones.data(), &self.dones_buf, b, &self.stream, "dones")?;
        dtod_from_slice_f32(gpu_batch.weights.data(), &self.is_weights_buf, b, &self.stream, "is_weights")?;

        // Actions: GpuTensor f32 -> I32 buf. Reinterpret f32 as u32 then DtoD as i32.
        let actions_f32 = gpu_batch.actions.data();
        let actions_u32: &CudaSlice<u32> = unsafe {
            &*(actions_f32 as *const CudaSlice<f32> as *const CudaSlice<u32>)
        };
        dtod_from_u32_to_i32(actions_u32, &self.actions_buf, b, &self.stream, "actions")?;

        Ok(())
    }

    // ═══════════════════════════════════════════════════════════════════
    // Kernel launch methods
    // ═══════════════════════════════════════════════════════════════════

    /// Launch the cuBLAS batched forward pass (online + target networks).
    ///
    /// Replaces `launch_forward_loss` for the forward phase only.
    /// The C51 distributional loss is still computed by the existing NVRTC kernel
    /// which reads the cuBLAS output buffers (`on_v_logits_buf`, `on_b_logits_buf`,
    /// `tg_v_logits_buf`, `tg_b_logits_buf`) together with the saved activations.
    ///
    /// This raises GPU occupancy from ~1.56% to >60% on H100 by replacing the
    /// 1-warp-per-sample shared-memory-bound kernel with cuBLAS DGEMM which
    /// uses tensor cores and efficiently tiles across the full SM.
    ///
    /// Returns `Ok(false)` if cuBLAS is not initialized (caller falls back to BF16 kernel).
    fn launch_cublas_forward(&self) -> Result<(), MLError> {
        let cublas = &self.cublas_forward;

        // Compute 20 raw F32 device pointers (online + target) into flat param buffers.
        let param_sizes = compute_param_sizes(&self.config);
        let on_w_ptrs = f32_weight_ptrs(&self.params_buf, &param_sizes, &self.stream);
        let tg_w_ptrs = f32_weight_ptrs(&self.target_params_buf, &param_sizes, &self.stream);

        // ── Pass 1: Online network forward on STATES ──────────────────
        // Writes into save_h_s1, save_h_s2, save_h_v, save_h_b0/b1/b2,
        // on_v_logits_buf, on_b_logits_buf.
        cublas.forward_online(
            &self.stream,
            &self.states_buf,
            &on_w_ptrs,
            &self.save_h_s1,
            &self.save_h_s2,
            &self.save_h_v,
            &self.save_h_b0,
            &self.save_h_b1,
            &self.save_h_b2,
            &self.on_v_logits_buf,
            &self.on_b_logits_buf,
        )?;

        // ── Pass 2: Target network forward on NEXT_STATES ────────────
        // Writes into tg_h_s2_buf (kept for DDQN), scratch bufs, tg_v_logits_buf, tg_b_logits_buf.
        cublas.forward_target(
            &self.stream,
            &self.next_states_buf,
            &tg_w_ptrs,
            &self.tg_h_s1_scratch,
            &self.tg_h_s2_buf,
            &self.tg_h_v_scratch,
            &self.tg_h_b_scratch,
            &self.tg_v_logits_buf,
            &self.tg_b_logits_buf,
        )?;

        // ── Pass 3: Online network forward on NEXT_STATES (Double DQN) ─
        // DDQN action selection: online network on next_states picks the greedy
        // action; target distribution for that action is the Bellman target.
        cublas.forward_online(
            &self.stream,
            &self.next_states_buf,
            &on_w_ptrs,
            &self.on_next_h_s1_scratch,
            &self.on_next_h_s2_scratch,
            &self.on_next_h_v_scratch,
            &self.on_next_h_b_scratch,
            &self.on_next_h_b_scratch, // reuse for b1 (scratch only)
            &self.on_next_h_b_scratch, // reuse for b2 (scratch only)
            &self.on_next_v_logits_buf,
            &self.on_next_b_logits_buf,
        )?;

        Ok(())
    }

    /// Launch the standalone C51 loss kernel on pre-computed cuBLAS logit outputs.
    ///
    /// Reads: on_v_logits, on_b_logits, tg_v_logits, tg_b_logits,
    ///        on_next_v_logits, on_next_b_logits (Double DQN).
    /// Writes: per_sample_loss, td_errors, total_loss, save_current_lp, save_projected.
    fn launch_c51_loss(&self) -> Result<(), MLError> {
        let kernel = &self.c51_loss_kernel;

        let b = self.config.batch_size;
        let na = self.config.num_atoms;
        let b0 = self.config.branch_0_size;
        let b1 = self.config.branch_1_size;
        let b2 = self.config.branch_2_size;

        // Extract raw pointers for online branch logits (contiguous [B, (B0+B1+B2)*NA])
        let on_b_base = raw_device_ptr(&self.on_b_logits_buf, &self.stream);
        let on_b0_ptr = on_b_base;
        let on_b1_ptr = on_b_base + (b * b0 * na * 4) as u64;
        let on_b2_ptr = on_b1_ptr + (b * b1 * na * 4) as u64;

        // Target branch logits
        let tg_b_base = raw_device_ptr(&self.tg_b_logits_buf, &self.stream);
        let tg_b0_ptr = tg_b_base;
        let tg_b1_ptr = tg_b_base + (b * b0 * na * 4) as u64;
        let tg_b2_ptr = tg_b1_ptr + (b * b1 * na * 4) as u64;

        // Online-next branch logits (Double DQN action selector)
        let on_next_b_base = raw_device_ptr(&self.on_next_b_logits_buf, &self.stream);
        let on_next_b0_ptr = on_next_b_base;
        let on_next_b1_ptr = on_next_b_base + (b * b0 * na * 4) as u64;
        let on_next_b2_ptr = on_next_b1_ptr + (b * b1 * na * 4) as u64;

        // N-step returns: use gamma^n for the Bellman projection.
        // The experience collector pre-computes R_n = sum(gamma^i * r_i).
        let gamma = self.config.gamma.powi(self.config.n_steps as i32);
        let batch_i32 = b as i32;

        // Shared memory: SHMEM_TOTAL_FLOATS * 4 bytes
        // 51+51+255+51+51+51+8 = 518 floats = 2072 bytes
        let max_branch = b0.max(b1).max(b2);
        let shmem_floats = na + na + max_branch * na + na + na + na + 8;
        let shmem_bytes = (shmem_floats * std::mem::size_of::<f32>()) as u32;

        unsafe {
            self.stream
                .launch_builder(kernel)
                // ── Online on current states (4) ──
                .arg(&self.on_v_logits_buf)
                .arg(&on_b0_ptr)
                .arg(&on_b1_ptr)
                .arg(&on_b2_ptr)
                // ── Target on next_states (4) ──
                .arg(&self.tg_v_logits_buf)
                .arg(&tg_b0_ptr)
                .arg(&tg_b1_ptr)
                .arg(&tg_b2_ptr)
                // ── Online-next on next_states (Double DQN, 4) ──
                .arg(&self.on_next_v_logits_buf)
                .arg(&on_next_b0_ptr)
                .arg(&on_next_b1_ptr)
                .arg(&on_next_b2_ptr)
                // ── Batch data (4) ──
                .arg(&self.actions_buf)
                .arg(&self.rewards_buf)
                .arg(&self.dones_buf)
                .arg(&self.is_weights_buf)
                // ── Outputs (3) ──
                .arg(&self.per_sample_loss_buf)
                .arg(&self.td_errors_buf)
                .arg(&self.total_loss_buf)
                // ── Saved for backward (2) ──
                .arg(&self.save_current_lp)
                .arg(&self.save_projected)
                // ── Config (2) ──
                .arg(&gamma)
                .arg(&batch_i32)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: shmem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!("c51_loss_batched launch: {e}")))?;
        }

        Ok(())
    }

    /// Launch C51 loss gradient kernel: save_current_lp, save_projected → dL/d_logits.
    ///
    /// Computes dL/d_combined = is_weights * (exp(current_lp) - projected) per branch,
    /// then routes through dueling to produce d_value_logits and d_adv_logits.
    fn launch_c51_grad(&self) -> Result<(), MLError> {
        let b = self.config.batch_size;
        let na = self.config.num_atoms;
        let b0 = self.config.branch_0_size;
        let b1 = self.config.branch_1_size;
        let b2 = self.config.branch_2_size;
        let total_branch_atoms = (b0 + b1 + b2) * na;

        let batch_i32 = b as i32;
        let na_i32 = na as i32;
        let b0_i32 = b0 as i32;
        let b1_i32 = b1 as i32;
        let b2_i32 = b2 as i32;
        let total_branch_atoms_i32 = total_branch_atoms as i32;

        let blocks = ((b * 3 * na + 255) / 256) as u32;

        unsafe {
            self.stream
                .launch_builder(&self.c51_grad_kernel)
                .arg(&self.save_current_lp)
                .arg(&self.save_projected)
                .arg(&self.is_weights_buf)
                .arg(&self.actions_buf)
                .arg(&self.d_value_logits_buf)
                .arg(&self.d_adv_logits_buf)
                .arg(&batch_i32)
                .arg(&na_i32)
                .arg(&b0_i32)
                .arg(&b1_i32)
                .arg(&b2_i32)
                .arg(&total_branch_atoms_i32)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("c51_grad_kernel launch: {e}")))?;
        }
        Ok(())
    }

    /// Launch the MSE loss kernel on expected Q-values (warmup before C51).
    ///
    /// Computes MSE loss between online E[Q] and Bellman target E[Q].
    /// Writes: per_sample_loss, td_errors, total_loss, save_current_lp (online softmax probs),
    /// save_projected (repurposed: stores online E[Q] per branch for the grad kernel).
    fn launch_mse_loss(&self) -> Result<(), MLError> {
        let kernel = &self.mse_loss_kernel;

        let b = self.config.batch_size;
        let na = self.config.num_atoms;
        let b0 = self.config.branch_0_size;
        let b1 = self.config.branch_1_size;
        let b2 = self.config.branch_2_size;

        // Extract raw pointers for online branch logits (contiguous [B, (B0+B1+B2)*NA])
        let on_b_base = raw_device_ptr(&self.on_b_logits_buf, &self.stream);
        let on_b0_ptr = on_b_base;
        let on_b1_ptr = on_b_base + (b * b0 * na * 4) as u64;
        let on_b2_ptr = on_b1_ptr + (b * b1 * na * 4) as u64;

        // Target branch logits
        let tg_b_base = raw_device_ptr(&self.tg_b_logits_buf, &self.stream);
        let tg_b0_ptr = tg_b_base;
        let tg_b1_ptr = tg_b_base + (b * b0 * na * 4) as u64;
        let tg_b2_ptr = tg_b1_ptr + (b * b1 * na * 4) as u64;

        // Online-next branch logits (Double DQN action selector)
        let on_next_b_base = raw_device_ptr(&self.on_next_b_logits_buf, &self.stream);
        let on_next_b0_ptr = on_next_b_base;
        let on_next_b1_ptr = on_next_b_base + (b * b0 * na * 4) as u64;
        let on_next_b2_ptr = on_next_b1_ptr + (b * b1 * na * 4) as u64;

        // N-step returns: use gamma^n for the Bellman projection.
        let gamma = self.config.gamma.powi(self.config.n_steps as i32);
        let batch_i32 = b as i32;

        // Shared memory: same layout as C51 kernel (support + val + adv + lp + proj + current_lp + reduce)
        let max_branch = b0.max(b1).max(b2);
        let shmem_floats = na + na + max_branch * na + na + na + na + 8;
        let shmem_bytes = (shmem_floats * std::mem::size_of::<f32>()) as u32;

        unsafe {
            self.stream
                .launch_builder(kernel)
                // ── Online on current states (4) ──
                .arg(&self.on_v_logits_buf)
                .arg(&on_b0_ptr)
                .arg(&on_b1_ptr)
                .arg(&on_b2_ptr)
                // ── Target on next_states (4) ──
                .arg(&self.tg_v_logits_buf)
                .arg(&tg_b0_ptr)
                .arg(&tg_b1_ptr)
                .arg(&tg_b2_ptr)
                // ── Online-next on next_states (Double DQN, 4) ──
                .arg(&self.on_next_v_logits_buf)
                .arg(&on_next_b0_ptr)
                .arg(&on_next_b1_ptr)
                .arg(&on_next_b2_ptr)
                // ── Batch data (4) ──
                .arg(&self.actions_buf)
                .arg(&self.rewards_buf)
                .arg(&self.dones_buf)
                .arg(&self.is_weights_buf)
                // ── Outputs (3) ──
                .arg(&self.per_sample_loss_buf)
                .arg(&self.td_errors_buf)
                .arg(&self.total_loss_buf)
                // ── Saved for backward (2) ── repurposed: save_current_lp = softmax probs,
                // save_projected = per-branch E[Q] values (3 floats per sample per branch)
                .arg(&self.save_current_lp)
                .arg(&self.save_projected)
                // ── Config (2) ──
                .arg(&gamma)
                .arg(&batch_i32)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: shmem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!("mse_loss_batched launch: {e}")))?;
        }

        Ok(())
    }

    /// Launch MSE loss gradient kernel through softmax expectation.
    ///
    /// Reads save_current_lp (softmax probs) and save_projected (expected Q + td_error),
    /// computes d_logit_j = td_error * is_weight * p_j * (z_j - E[Q]) per branch,
    /// then routes through dueling to produce d_value_logits and d_adv_logits.
    fn launch_mse_grad(&self) -> Result<(), MLError> {
        self.launch_mse_grad_inner(&self.d_value_logits_buf, &self.d_adv_logits_buf)
    }

    /// Launch MSE grad to scratch buffers (for blended MSE+C51 loss path).
    ///
    /// Identical to `launch_mse_grad()` but writes to `d_value_logits_mse` /
    /// `d_adv_logits_mse` instead of the main gradient buffers. The caller then
    /// SAXPY-blends scratch into main: main = alpha*C51 + (1-alpha)*MSE.
    fn launch_mse_grad_to_scratch(&self) -> Result<(), MLError> {
        self.launch_mse_grad_inner(&self.d_value_logits_mse, &self.d_adv_logits_mse)
    }

    /// Shared implementation for MSE gradient kernel launch.
    ///
    /// Writes gradient outputs to the provided destination buffers.
    fn launch_mse_grad_inner(
        &self,
        d_value_dst: &CudaSlice<f32>,
        d_adv_dst: &CudaSlice<f32>,
    ) -> Result<(), MLError> {
        let b = self.config.batch_size;
        let na = self.config.num_atoms;
        let b0 = self.config.branch_0_size;
        let b1 = self.config.branch_1_size;
        let b2 = self.config.branch_2_size;
        let total_branch_atoms = (b0 + b1 + b2) * na;

        let batch_i32 = b as i32;
        let na_i32 = na as i32;
        let b0_i32 = b0 as i32;
        let b1_i32 = b1 as i32;
        let b2_i32 = b2 as i32;
        let total_branch_atoms_i32 = total_branch_atoms as i32;

        let blocks = ((b * 3 * na + 255) / 256) as u32;

        unsafe {
            self.stream
                .launch_builder(&self.mse_grad_kernel)
                .arg(&self.save_current_lp)
                .arg(&self.save_projected)
                .arg(&self.is_weights_buf)
                .arg(&self.actions_buf)
                .arg(d_value_dst)
                .arg(d_adv_dst)
                .arg(&batch_i32)
                .arg(&na_i32)
                .arg(&b0_i32)
                .arg(&b1_i32)
                .arg(&b2_i32)
                .arg(&total_branch_atoms_i32)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("mse_grad_kernel launch: {e}")))?;
        }
        Ok(())
    }

    /// cuBLAS SGEMM backward pass: chain rule through all layers.
    ///
    /// Reads dL/d_logits from `d_value_logits_buf` and `d_adv_logits_buf`
    /// (populated by `launch_c51_grad` or `launch_mse_grad`), propagates gradients
    /// through all layers using cuBLAS GEMM, and accumulates into `grad_buf`.
    fn launch_cublas_backward(&self) -> Result<(), MLError> {
        let bw = &self.cublas_backward;

        let param_sizes = compute_param_sizes(&self.config);
        let w_ptrs = f32_weight_ptrs(&self.params_buf, &param_sizes, &self.stream);

        let grad_base = bw_raw_ptr(&self.grad_buf, &self.stream);
        let states_ptr = bw_raw_ptr(&self.states_buf, &self.stream);
        let h_s1_ptr = bw_raw_ptr(&self.save_h_s1, &self.stream);
        let h_s2_ptr = bw_raw_ptr(&self.save_h_s2, &self.stream);
        let h_v_ptr  = bw_raw_ptr(&self.save_h_v,  &self.stream);
        let h_b0_ptr = bw_raw_ptr(&self.save_h_b0, &self.stream);
        let h_b1_ptr = bw_raw_ptr(&self.save_h_b1, &self.stream);
        let h_b2_ptr = bw_raw_ptr(&self.save_h_b2, &self.stream);

        let d_h_s2_ptr = bw_raw_ptr(&self.bw_d_h_s2, &self.stream);
        let d_h_s1_ptr = bw_raw_ptr(&self.bw_d_h_s1, &self.stream);
        let d_h_v_ptr  = bw_raw_ptr(&self.bw_d_h_v,  &self.stream);
        let d_h_b0_ptr = bw_raw_ptr(&self.bw_d_h_b0, &self.stream);
        let d_h_b1_ptr = bw_raw_ptr(&self.bw_d_h_b1, &self.stream);
        let d_h_b2_ptr = bw_raw_ptr(&self.bw_d_h_b2, &self.stream);

        // dL/d_logits from c51_grad_kernel
        let d_value_logits_ptr = bw_raw_ptr(&self.d_value_logits_buf, &self.stream);
        let d_adv_logits_ptr = bw_raw_ptr(&self.d_adv_logits_buf, &self.stream);

        let na = self.config.num_atoms;
        let f32_size = std::mem::size_of::<f32>() as u64;
        let d_adv0 = d_adv_logits_ptr;
        let d_adv1 = d_adv0 + (self.config.batch_size * self.config.branch_0_size * na) as u64 * f32_size;
        let d_adv2 = d_adv1 + (self.config.batch_size * self.config.branch_1_size * na) as u64 * f32_size;

        bw.backward_full(
            &self.stream,
            d_value_logits_ptr,
            &[d_adv0, d_adv1, d_adv2],
            states_ptr,
            h_s1_ptr,
            h_s2_ptr,
            h_v_ptr,
            &[h_b0_ptr, h_b1_ptr, h_b2_ptr],
            &w_ptrs,
            grad_base,
            d_h_s2_ptr,
            d_h_s1_ptr,
            d_h_v_ptr,
            &[d_h_b0_ptr, d_h_b1_ptr, d_h_b2_ptr],
        )?;

        Ok(())
    }

    /// Launch the gradient L2 norm reduction kernel.
    ///
    /// Argument order matches `dqn_grad_norm_kernel` in `dqn_training_kernel.cu`:
    /// grads, out_grad_norm, total_params = 3 args.
    ///
    /// `grad_norm_buf` must be zeroed before this call (done in `submit_training_ops`).
    fn launch_grad_norm(&self) -> Result<(), MLError> {
        let tp = self.total_params as i32;
        let blocks = ((self.total_params + 255) / 256) as u32;

        let launch_cfg = LaunchConfig {
            grid_dim: (blocks, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0, // uses static __shared__ warp_sums[8]
        };

        // Safety: argument order matches the extern "C" kernel signature exactly.
        // grad_buf has size = total_params; grad_norm_buf has size = 1.
        unsafe {
            self.stream
                .launch_builder(&self.grad_norm_kernel)
                .arg(&self.grad_buf)
                .arg(&self.grad_norm_buf)
                .arg(&tp)
                .launch(launch_cfg)
                .map_err(|e| {
                    MLError::ModelError(format!("dqn_grad_norm_kernel launch: {e}"))
                })?;
        }

        Ok(())
    }

    /// Launch the Adam update kernel with correct gradient clipping.
    ///
    /// Argument order matches `dqn_adam_update_kernel` in `dqn_training_kernel.cu`:
    /// params, grads, m, v, grad_norm_sq, lr, beta1, beta2, epsilon,
    /// weight_decay, max_grad_norm, t_ptr, total_params = 13 args.
    ///
    /// Must be launched AFTER `launch_grad_norm` so that `grad_norm_buf`
    /// contains the completed sum of squares (no race condition).
    /// `t_ptr` is a device buffer (not scalar) so the CUDA Graph can be
    /// replayed with an updated step counter.
    fn launch_adam_update(&self) -> Result<(), MLError> {
        let tp = self.total_params as i32;
        let blocks = ((self.total_params + 255) / 256) as u32;

        let launch_cfg = LaunchConfig {
            grid_dim: (blocks, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };

        let lr = self.config.lr;
        let beta1 = self.config.beta1;
        let beta2 = self.config.beta2;
        let epsilon = self.config.epsilon;
        let weight_decay = self.config.weight_decay;
        let max_grad_norm = self.config.max_grad_norm;

        // Safety: argument order matches the extern "C" kernel signature exactly.
        // All buffers are pre-allocated with size = total_params.
        // grad_norm_buf contains the completed norm from launch_grad_norm().
        // t_buf is a device pointer (kernel reads *t_ptr) for CUDA Graph compatibility.
        unsafe {
            self.stream
                .launch_builder(&self.adam_update_kernel)
                .arg(&self.params_buf)
                .arg(&self.grad_buf)
                .arg(&self.m_buf)
                .arg(&self.v_buf)
                .arg(&self.grad_norm_buf)
                .arg(&lr)
                .arg(&beta1)
                .arg(&beta2)
                .arg(&epsilon)
                .arg(&weight_decay)
                .arg(&max_grad_norm)
                .arg(&self.t_buf)   // device pointer — not baked scalar
                .arg(&tp)
                .launch(launch_cfg)
                .map_err(|e| {
                    MLError::ModelError(format!("dqn_adam_update_kernel launch: {e}"))
                })?;
        }

        Ok(())
    }

    // ═══════════════════════════════════════════════════════════════════
    // Weight flattening / unflattening (GPU d2d — zero host roundtrip)
    // ═══════════════════════════════════════════════════════════════════

    /// Copy 20 individual weight tensors into the flat `params_buf`.
    ///
    /// Order matches the GOFF_* layout in `dqn_training_kernel.cu`.
    /// Pure device-to-device copies via `cuMemcpyDtoDAsync` — zero host roundtrip.
    fn flatten_online_weights(
        &self,
        online_d: &DuelingWeightSet,
        online_b: &BranchingWeightSet,
    ) -> Result<(), MLError> {
        let sizes = compute_param_sizes(&self.config);
        let dst_base = raw_device_ptr(&self.params_buf, &self.stream);

        // Ordered to match GOFF_* layout
        let slices: [&CudaSlice<f32>; 20] = [
            &online_d.w_s1, &online_d.b_s1,
            &online_d.w_s2, &online_d.b_s2,
            &online_d.w_v1, &online_d.b_v1,
            &online_d.w_v2, &online_d.b_v2,
            &online_d.w_a1, &online_d.b_a1,     // branch 0 = advantage slot
            &online_d.w_a2, &online_d.b_a2,
            &online_b.w_bo1, &online_b.b_bo1,   // branch 1
            &online_b.w_bo2, &online_b.b_bo2,
            &online_b.w_bu1, &online_b.b_bu1,   // branch 2
            &online_b.w_bu2, &online_b.b_bu2,
        ];

        let weight_names = [
            "w_s1", "b_s1", "w_s2", "b_s2",
            "w_v1", "b_v1", "w_v2", "b_v2",
            "w_a1", "b_a1", "w_a2", "b_a2",
            "w_bo1", "b_bo1", "w_bo2", "b_bo2",
            "w_bu1", "b_bu1", "w_bu2", "b_bu2",
        ];

        let mut byte_offset: u64 = 0;
        for (i, slice) in slices.iter().enumerate() {
            let actual = slice.len();
            if actual != sizes[i] {
                return Err(MLError::ModelError(format!(
                    "flatten[{}] ({}) size mismatch: CudaSlice has {} elems, expected {}",
                    i, weight_names[i], actual, sizes[i],
                )));
            }
            let num_bytes = sizes[i] * std::mem::size_of::<f32>();
            let src = raw_device_ptr(slice, &self.stream);
            dtod_copy(dst_base + byte_offset, src, num_bytes, &self.stream, i, "flatten")?;
            byte_offset += num_bytes as u64;
        }

        Ok(())
    }

    /// Copy flat `params_buf` back to 20 individual weight tensors.
    ///
    /// Called after each Adam step to sync the canonical weight storage.
    /// Pure device-to-device copies via `cuMemcpyDtoDAsync` — zero host roundtrip.
    ///
    /// During CUDA Graph capture, these d2d copies are recorded and replayed.
    /// The device pointers are stable (CudaSlice allocations don't move).
    fn unflatten_online_weights(
        &self,
        online_d: &DuelingWeightSet,
        online_b: &BranchingWeightSet,
    ) -> Result<(), MLError> {
        let sizes = compute_param_sizes(&self.config);
        let src_base = raw_device_ptr(&self.params_buf, &self.stream);

        // Must match GOFF_* order exactly
        let slices: [&CudaSlice<f32>; 20] = [
            &online_d.w_s1, &online_d.b_s1,
            &online_d.w_s2, &online_d.b_s2,
            &online_d.w_v1, &online_d.b_v1,
            &online_d.w_v2, &online_d.b_v2,
            &online_d.w_a1, &online_d.b_a1,
            &online_d.w_a2, &online_d.b_a2,
            &online_b.w_bo1, &online_b.b_bo1,
            &online_b.w_bo2, &online_b.b_bo2,
            &online_b.w_bu1, &online_b.b_bu1,
            &online_b.w_bu2, &online_b.b_bu2,
        ];

        let mut byte_offset: u64 = 0;
        for (i, slice) in slices.iter().enumerate() {
            let num_bytes = sizes[i] * std::mem::size_of::<f32>();
            let dst = raw_device_ptr(slice, &self.stream);
            dtod_copy(dst, src_base + byte_offset, num_bytes, &self.stream, i, "unflatten")?;
            byte_offset += num_bytes as u64;
        }

        Ok(())
    }

    // ═══════════════════════════════════════════════════════════════════
    // GPU-native Polyak EMA target update
    // ═══════════════════════════════════════════════════════════════════

    /// Copy 20 individual target weight tensors into the flat `target_params_buf`.
    ///
    /// Same GOFF_* layout as `flatten_online_weights()`.
    /// Pure device-to-device copies -- zero host roundtrip.
    fn flatten_target_weights(
        &self,
        target_d: &DuelingWeightSet,
        target_b: &BranchingWeightSet,
    ) -> Result<(), MLError> {
        let sizes = compute_param_sizes(&self.config);
        let dst_base = raw_device_ptr(&self.target_params_buf, &self.stream);

        let slices: [&CudaSlice<f32>; 20] = [
            &target_d.w_s1, &target_d.b_s1,
            &target_d.w_s2, &target_d.b_s2,
            &target_d.w_v1, &target_d.b_v1,
            &target_d.w_v2, &target_d.b_v2,
            &target_d.w_a1, &target_d.b_a1,
            &target_d.w_a2, &target_d.b_a2,
            &target_b.w_bo1, &target_b.b_bo1,
            &target_b.w_bo2, &target_b.b_bo2,
            &target_b.w_bu1, &target_b.b_bu1,
            &target_b.w_bu2, &target_b.b_bu2,
        ];

        // Sync stream to catch any deferred errors from graph replay
        let sync_r = unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()) };
        if sync_r != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
            return Err(MLError::ModelError(format!("flatten_target pre-sync: {sync_r:?}")));
        }
        let mut byte_offset: u64 = 0;
        for (i, slice) in slices.iter().enumerate() {
            let num_bytes = sizes[i] * std::mem::size_of::<f32>();
            let src = raw_device_ptr(slice, &self.stream);
            let dst = dst_base + byte_offset;
            dtod_copy(dst, src, num_bytes, &self.stream, i, "flatten_target")?;
            byte_offset += num_bytes as u64;
        }

        Ok(())
    }

    /// Copy flat `target_params_buf` back to 20 individual target weight tensors.
    ///
    /// Called after the fused EMA kernel to scatter updated flat target weights
    /// back into the individual `DuelingWeightSet` + `BranchingWeightSet` tensors.
    /// Pure device-to-device copies -- zero host roundtrip.
    fn unflatten_target_weights(
        &self,
        target_d: &DuelingWeightSet,
        target_b: &BranchingWeightSet,
    ) -> Result<(), MLError> {
        let sizes = compute_param_sizes(&self.config);
        let src_base = raw_device_ptr(&self.target_params_buf, &self.stream);

        let slices: [&CudaSlice<f32>; 20] = [
            &target_d.w_s1, &target_d.b_s1,
            &target_d.w_s2, &target_d.b_s2,
            &target_d.w_v1, &target_d.b_v1,
            &target_d.w_v2, &target_d.b_v2,
            &target_d.w_a1, &target_d.b_a1,
            &target_d.w_a2, &target_d.b_a2,
            &target_b.w_bo1, &target_b.b_bo1,
            &target_b.w_bo2, &target_b.b_bo2,
            &target_b.w_bu1, &target_b.b_bu1,
            &target_b.w_bu2, &target_b.b_bu2,
        ];

        let mut byte_offset: u64 = 0;
        for (i, slice) in slices.iter().enumerate() {
            let num_bytes = sizes[i] * std::mem::size_of::<f32>();
            let dst = raw_device_ptr(slice, &self.stream);
            dtod_copy(dst, src_base + byte_offset, num_bytes, &self.stream, i, "unflatten_target")?;
            byte_offset += num_bytes as u64;
        }

        Ok(())
    }

    // ═══════════════════════════════════════════════════════════════════
    // GPU-native Polyak EMA target update (fused single kernel)
    // ═══════════════════════════════════════════════════════════════════

    /// GPU-native Polyak EMA: `target[i] = (1-tau)*target[i] + tau*online[i]`
    ///
    /// Fused single-kernel update over flat parameter buffers. On first call,
    /// flattens target weights into `target_params_buf` (same GOFF_* layout as
    /// `params_buf`). Then launches ONE EMA kernel over the entire flat buffer
    /// instead of 20 per-tensor launches. After the kernel, scatters the updated
    /// flat target weights back to individual tensors via `unflatten_target_weights()`.
    ///
    /// Runs OUTSIDE the captured CUDA Graph -- device pointers are stable so
    /// the graph stays valid.
    pub fn target_ema_update(
        &mut self,
        _online_d: &DuelingWeightSet,
        _online_b: &BranchingWeightSet,
        target_d: &mut DuelingWeightSet,
        target_b: &mut BranchingWeightSet,
        tau: f32,
    ) -> Result<(), MLError> {
        // Sync stream and clear any stale errors from graph capture phase.
        // cudarc stores errors from cuStreamWaitEvent on disabled events during
        // graph capture. check_err() consumes them so bind_to_thread() succeeds.
        unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }
        let _ = self.stream.context().check_err();

        // First call: flatten target weights into flat buffer (20 DtoD copies, once only).
        // Subsequent calls reuse the flat buffer which is kept in sync by unflatten at end.
        if !self.target_params_initialized {
            self.flatten_target_weights(target_d, target_b)?;
            self.target_params_initialized = true;
        }

        // Single fused EMA kernel over all parameters at once.
        // params_buf contains the latest online weights (kept in sync by
        // unflatten_online_weights after each Adam step).
        let n = self.total_params as i32;
        let blocks = ((self.total_params + 255) / 256) as u32;
        let launch_cfg = LaunchConfig {
            grid_dim: (blocks, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };

        let t_ptr = raw_device_ptr(&self.target_params_buf, &self.stream);
        let o_ptr = raw_device_ptr(&self.params_buf, &self.stream);
        unsafe {
            self.stream
                .launch_builder(&self.ema_kernel)
                .arg(&t_ptr)
                .arg(&o_ptr)
                .arg(&tau)
                .arg(&n)
                .launch(launch_cfg)
                .map_err(|e| {
                    MLError::ModelError(format!("dqn_ema_kernel flat launch: {e}"))
                })?;
        }

        // Scatter flat target buffer back to individual weight tensors
        self.unflatten_target_weights(target_d, target_b)?;

        // Sync target BF16 mirrors from updated F32 target weights
        // 20 per-segment launches: target_params_buf → bf16_target_params_buf
        self.sync_target_bf16()?;

        Ok(())
    }
}

// ── Compilation ─────────────────────────────────────────────────────────────

/// Compile the 4 live utility kernels from a single NVRTC compilation.
///
/// The source is `common_device_functions.cuh` (provides `f32_to_bf16_kernel` and
/// `bf16_to_f32_kernel`) concatenated with `dqn_utility_kernels.cu` (provides
/// `dqn_grad_norm_kernel` and `dqn_adam_update_kernel`).
///
/// Returns `(grad_norm, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, relu_mask, spectral_norm)`.
fn compile_training_kernels(
    stream: &Arc<CudaStream>,
    config: &GpuDqnTrainConfig,
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
    // Inject the three defines required by common_device_functions.cuh (#error guards).
    // The utility kernels (grad_norm, adam_update, f32_to_bf16, bf16_to_f32) take all
    // sizes as runtime arguments — no shmem tiling or network-dimension #defines needed.
    let dim_overrides = format!(
        "#define STATE_DIM {state_dim}\n\
         #define MARKET_DIM 42\n\
         #define PORTFOLIO_DIM 8\n",
        state_dim = config.state_dim,
    );

    let common_src = include_str!("common_device_functions.cuh");
    let kernel_src = include_str!("dqn_utility_kernels.cu");
    let full_source = format!("{dim_overrides}\n{common_src}\n{kernel_src}");

    info!(
        state_dim = config.state_dim,
        total_params = compute_total_params(config),
        "GpuDqnTrainer: compiling utility kernels (grad_norm + adam_update + BF16 converters)"
    );

    let context = stream.context();
    let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context)
        .map_err(|e| {
            MLError::ModelError(format!("dqn_utility_kernels compilation failed: {e}"))
        })?;
    let module = context.load_module(ptx).map_err(|e| {
        MLError::ModelError(format!("dqn_training module load: {e}"))
    })?;

    let grad_norm = module
        .load_function("dqn_grad_norm_kernel")
        .map_err(|e| {
            MLError::ModelError(format!("dqn_grad_norm_kernel load: {e}"))
        })?;
    let adam_update = module
        .load_function("dqn_adam_update_kernel")
        .map_err(|e| {
            MLError::ModelError(format!("dqn_adam_update_kernel load: {e}"))
        })?;
    let f32_to_bf16 = module
        .load_function("f32_to_bf16_kernel")
        .map_err(|e| {
            MLError::ModelError(format!("f32_to_bf16_kernel load: {e}"))
        })?;
    let bf16_to_f32 = module
        .load_function("bf16_to_f32_kernel")
        .map_err(|e| {
            MLError::ModelError(format!("bf16_to_f32_kernel load: {e}"))
        })?;

    let saxpy = module
        .load_function("dqn_saxpy_kernel")
        .map_err(|e| {
            MLError::ModelError(format!("dqn_saxpy_kernel load: {e}"))
        })?;
    let zero = module
        .load_function("dqn_zero_kernel")
        .map_err(|e| {
            MLError::ModelError(format!("dqn_zero_kernel load: {e}"))
        })?;

    let regime_scale = module
        .load_function("dqn_regime_scale_kernel")
        .map_err(|e| {
            MLError::ModelError(format!("dqn_regime_scale_kernel load: {e}"))
        })?;

    let shrink_perturb = module
        .load_function("dqn_shrink_perturb_kernel")
        .map_err(|e| {
            MLError::ModelError(format!("dqn_shrink_perturb_kernel load: {e}"))
        })?;
    let _relu_mask_from_module = module
        .load_function("dqn_relu_mask_kernel")
        .map_err(|e| {
            MLError::ModelError(format!("dqn_relu_mask_kernel load: {e}"))
        })?;
    let spectral_norm = module
        .load_function("spectral_norm_kernel")
        .map_err(|e| {
            MLError::ModelError(format!("spectral_norm_kernel load: {e}"))
        })?;

    info!("GpuDqnTrainer: 10 utility kernels compiled and loaded");
    Ok((grad_norm, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm))
}

/// Compile the standalone Polyak EMA kernel.
///
/// `ema_kernel(target, online, tau, n)`: `target[i] = (1-tau)*target[i] + tau*online[i]`
///
/// This kernel is NOT captured in the CUDA Graph — it runs after graph replay
/// to blend online weights into the target network in-place. Device pointers
/// are stable across calls, so the training graph remains valid.
fn compile_ema_kernel(stream: &Arc<CudaStream>) -> Result<CudaFunction, MLError> {
    let src = r#"
extern "C" __global__
void dqn_ema_kernel(float* __restrict__ target,
                    const float* __restrict__ online,
                    float tau,
                    int n)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        target[i] = (1.0f - tau) * target[i] + tau * online[i];
    }
}
"#;
    let context = stream.context();
    let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context)
        .map_err(|e| MLError::ModelError(format!("dqn_ema_kernel compilation: {e}")))?;
    let module = context.load_module(ptx).map_err(|e| {
        MLError::ModelError(format!("dqn_ema module load: {e}"))
    })?;
    module.load_function("dqn_ema_kernel").map_err(|e| {
        MLError::ModelError(format!("dqn_ema_kernel load: {e}"))
    })
}

/// Compile standalone ReLU mask kernel for IQN trunk gradient.
///
/// Compiled as a separate module (like EMA/PER) rather than as part of the
/// training utility module. This avoids potential interference from
/// `common_device_functions.cuh` defines during cubin compilation.
fn compile_relu_mask_standalone(stream: &Arc<CudaStream>) -> Result<CudaFunction, MLError> {
    let src = r#"
extern "C" __global__
void relu_mask_standalone(float* __restrict__ dx,
                          const float* __restrict__ activation,
                          int n)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    if (activation[i] <= 0.0f) dx[i] = 0.0f;
}
"#;
    let context = stream.context();
    let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context)
        .map_err(|e| MLError::ModelError(format!("relu_mask_standalone compilation: {e}")))?;
    let module = context.load_module(ptx).map_err(|e| {
        MLError::ModelError(format!("relu_mask_standalone module load: {e}"))
    })?;
    module.load_function("relu_mask_standalone").map_err(|e| {
        MLError::ModelError(format!("relu_mask_standalone load: {e}"))
    })
}

/// Compile the standalone PER priority update kernel.
///
/// `per_update_priorities_kernel(td_errors, indices, priorities, batch_max, alpha, epsilon, B)`:
/// Computes `|td|^alpha + eps`, scatter-writes to priorities, atomicMax for batch max.
/// Single kernel launch replaces 7-8 Candle tensor ops per training step.
fn compile_per_update_kernel(stream: &Arc<CudaStream>) -> Result<CudaFunction, MLError> {
    let src = r#"
extern "C" __global__
void per_update_priorities_kernel(
    const float* __restrict__ td_errors,
    const unsigned int* __restrict__ indices,
    float* __restrict__ priorities,
    float* __restrict__ batch_max,
    float alpha,
    float epsilon,
    int batch_size,
    int capacity)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    float td = fabsf(td_errors[i]);
    float new_prio = powf(td, alpha) + epsilon;

    // Clamp to [epsilon, 1e6] — same bounds as GpuReplayBuffer::update_priorities_gpu
    if (new_prio < epsilon) new_prio = epsilon;
    if (new_prio > 1e6f) new_prio = 1e6f;

    // Scatter-write priority at buffer index (bounds-checked)
    unsigned int idx = indices[i];
    if (idx < (unsigned int)capacity)
        priorities[idx] = new_prio;

    // atomicMax for batch max: IEEE 754 positive floats preserve int ordering
    int ival = __float_as_int(new_prio);
    atomicMax((int*)batch_max, ival);
}
"#;
    let context = stream.context();
    let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context)
        .map_err(|e| MLError::ModelError(format!("per_update_priorities_kernel compilation: {e}")))?;
    let module = context.load_module(ptx).map_err(|e| {
        MLError::ModelError(format!("per_update module load: {e}"))
    })?;
    module.load_function("per_update_priorities_kernel").map_err(|e| {
        MLError::ModelError(format!("per_update_priorities_kernel load: {e}"))
    })
}

/// Compile the standalone C51 distributional loss kernel.
///
/// This kernel operates on pre-computed forward pass outputs (value + advantage logits)
/// from cuBLAS, computing the C51 cross-entropy loss with Bellman projection.
/// Grid: (batch_size, 1, 1), Block: (256, 1, 1), Shmem: ~2 KB/block.
fn compile_c51_loss_kernel(
    stream: &Arc<CudaStream>,
    config: &GpuDqnTrainConfig,
) -> Result<CudaFunction, MLError> {
    let defines = format!(
        "#define NUM_ATOMS {na}\n\
         #define V_MIN ({v_min:.1}f)\n\
         #define V_MAX ({v_max:.1}f)\n\
         #define BRANCH_0_SIZE {b0}\n\
         #define BRANCH_1_SIZE {b1}\n\
         #define BRANCH_2_SIZE {b2}\n",
        na = config.num_atoms,
        v_min = config.v_min,
        v_max = config.v_max,
        b0 = config.branch_0_size,
        b1 = config.branch_1_size,
        b2 = config.branch_2_size,
    );
    let kernel_src = include_str!("c51_loss_kernel.cu");
    let full_source = format!("{defines}\n{kernel_src}");

    let context = stream.context();
    let ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context)
        .map_err(|e| MLError::ModelError(format!("c51_loss_kernel compilation: {e}")))?;
    let module = context.load_module(ptx)
        .map_err(|e| MLError::ModelError(format!("c51_loss module load: {e}")))?;
    module.load_function("c51_loss_batched")
        .map_err(|e| MLError::ModelError(format!("c51_loss_batched load: {e}")))
}

/// Compile the C51 loss gradient kernel.
///
/// Computes dL/d_logits from save_current_lp and save_projected, then routes
/// through the dueling architecture to produce d_value_logits and d_adv_logits.
///
/// dL/d_combined[b,d,j] = is_weights[b] * (exp(current_lp[b,d,j]) - projected[b,d,j])
/// d_value[b,j] = sum_d dL/d_combined[b,d,j]
/// d_adv[b,d,a,j] = dL/d_combined[b,d,j] * (delta(a,a_d) - 1/A_d)
fn compile_c51_grad_kernel(
    stream: &Arc<CudaStream>,
    _config: &GpuDqnTrainConfig,
) -> Result<CudaFunction, MLError> {
    let entropy_coeff = _config.entropy_coefficient;
    let src = format!(r#"
extern "C" __global__ void c51_grad_kernel(
    const float* __restrict__ current_lp,      // [B, 3, NA]
    const float* __restrict__ projected,        // [B, 3, NA]
    const float* __restrict__ is_weights,       // [B]
    const int*   __restrict__ actions,          // [B] factored
    float* __restrict__ d_value_logits,         // [B, NA]
    float* __restrict__ d_adv_logits,           // [B, (B0+B1+B2)*NA]
    int batch_size,
    int num_atoms,
    int b0_size, int b1_size, int b2_size,
    int total_branch_atoms)
{{
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int total_grad_elems = batch_size * 3 * num_atoms;
    if (tid >= total_grad_elems) return;

    int j  = tid % num_atoms;
    int d  = (tid / num_atoms) % 3;
    int b  = tid / (3 * num_atoms);

    float isw = is_weights[b];
    float lp  = current_lp[tid];
    float proj = projected[tid];

    /* Cross-entropy gradient: d/d_logits(-Σ proj * lp) = exp(lp) - proj */
    float d_combined = isw * (expf(lp) - proj);

    /* Entropy regularization: d/d_logits(-coeff * H(p)) = coeff * (1 + lp)
     * where H(p) = -Σ p * log(p) and p = exp(lp) (log-probs).
     * This prevents C51 distributions from collapsing to a single atom.
     * CRITICAL: clamp lp to [-10, 0] to prevent gradient explosion from
     * near-zero probability atoms (lp=-50 → gradient term = -0.49 per atom). */
    float entropy_coeff = {entropy_coeff:.6}f;
    if (entropy_coeff > 0.0f) {{
        float lp_clamped = fmaxf(lp, -10.0f); /* floor: p ≥ exp(-10) ≈ 4.5e-5 */
        d_combined += entropy_coeff * (1.0f + lp_clamped);
    }}

    // Route through dueling: d_value[b,j] += d_combined
    atomicAdd(&d_value_logits[b * num_atoms + j], d_combined);

    // Factored action decode
    int factored = actions[b];
    int max_action = b0_size * b1_size * b2_size;
    if (factored < 0 || factored >= max_action) factored = 0;
    int branch_sizes[3];
    branch_sizes[0] = b0_size;
    branch_sizes[1] = b1_size;
    branch_sizes[2] = b2_size;

    int a_d;
    if (d == 0) a_d = factored / (b1_size * b2_size);
    else if (d == 1) a_d = (factored / b2_size) % b1_size;
    else a_d = factored % b2_size;

    int A_d = branch_sizes[d];
    float inv_A = 1.0f / (float)A_d;

    // Compute per-branch adv offset
    int branch_offset = 0;
    for (int dd = 0; dd < d; dd++)
        branch_offset += branch_sizes[dd] * num_atoms;

    // d_adv[b, d, a, j] = d_combined * (delta(a, a_d) - 1/A_d)
    for (int a = 0; a < A_d; a++) {{
        float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A);
        int adv_idx = b * total_branch_atoms + branch_offset + a * num_atoms + j;
        atomicAdd(&d_adv_logits[adv_idx], d_combined * dueling_grad);
    }}
}}
"#);

    let context = stream.context();
    let ptx = crate::cuda_pipeline::compile_ptx_for_device(&src, &context)
        .map_err(|e| MLError::ModelError(format!("c51_grad_kernel compilation: {e}")))?;
    let module = context.load_module(ptx)
        .map_err(|e| MLError::ModelError(format!("c51_grad module load: {e}")))?;
    module.load_function("c51_grad_kernel")
        .map_err(|e| MLError::ModelError(format!("c51_grad_kernel load: {e}")))
}

/// Compile the MSE loss kernel on expected Q-values (warmup before C51).
///
/// Same signature as `c51_loss_batched` but computes:
///   1. Online E[Q] = sum_j softmax(dueling_logits[j]) * z_j
///   2. Target E[Q] via Double DQN (online selects, target evaluates)
///   3. Bellman: target_Q = reward + gamma * (1-done) * target_E[Q]
///   4. MSE loss: 0.5 * (online_E[Q] - target_Q)^2 per branch, averaged
///
/// Saves softmax probs in save_current_lp (for gradient kernel).
/// Saves [td_error, E_Q, 0...] per branch per sample in save_projected.
fn compile_mse_loss_kernel(
    stream: &Arc<CudaStream>,
    config: &GpuDqnTrainConfig,
) -> Result<CudaFunction, MLError> {
    let defines = format!(
        "#define NUM_ATOMS {na}\n\
         #define V_MIN ({v_min:.1}f)\n\
         #define V_MAX ({v_max:.1}f)\n\
         #define BRANCH_0_SIZE {b0}\n\
         #define BRANCH_1_SIZE {b1}\n\
         #define BRANCH_2_SIZE {b2}\n",
        na = config.num_atoms,
        v_min = config.v_min,
        v_max = config.v_max,
        b0 = config.branch_0_size,
        b1 = config.branch_1_size,
        b2 = config.branch_2_size,
    );
    let kernel_src = include_str!("mse_loss_kernel.cu");
    let full_source = format!("{defines}\n{kernel_src}");

    let context = stream.context();
    let ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context)
        .map_err(|e| MLError::ModelError(format!("mse_loss_kernel compilation: {e}")))?;
    let module = context.load_module(ptx)
        .map_err(|e| MLError::ModelError(format!("mse_loss module load: {e}")))?;
    module.load_function("mse_loss_batched")
        .map_err(|e| MLError::ModelError(format!("mse_loss_batched load: {e}")))
}

/// Compile the MSE loss gradient kernel through softmax expectation.
///
/// For each sample [b], branch [d], atom [j]:
///   d_logit_j = td_error * is_weight * p_j * (z_j - E[Q])
/// where p_j = softmax prob (from save_current_lp), z_j = support value,
/// E[Q] = expected Q (from save_projected), td_error (from save_projected).
///
/// Routes through dueling: same pattern as c51_grad_kernel.
fn compile_mse_grad_kernel(
    stream: &Arc<CudaStream>,
    config: &GpuDqnTrainConfig,
) -> Result<CudaFunction, MLError> {
    let src = format!(r#"
extern "C" __global__ void mse_grad_kernel(
    const float* __restrict__ save_probs,       // [B, 3, NA] softmax probs
    const float* __restrict__ save_eq_td,       // [B, 3, NA] layout: [td_error, E_Q, 0, ...]
    const float* __restrict__ is_weights,       // [B]
    const int*   __restrict__ actions,          // [B] factored
    float* __restrict__ d_value_logits,         // [B, NA]
    float* __restrict__ d_adv_logits,           // [B, (B0+B1+B2)*NA]
    int batch_size,
    int num_atoms,
    int b0_size, int b1_size, int b2_size,
    int total_branch_atoms)
{{
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int total_grad_elems = batch_size * 3 * num_atoms;
    if (tid >= total_grad_elems) return;

    int j  = tid % num_atoms;
    int d  = (tid / num_atoms) % 3;
    int b  = tid / (3 * num_atoms);

    float isw = is_weights[b];
    float p_j = save_probs[tid];  /* softmax prob for atom j */

    /* Read td_error and E[Q] from save_eq_td.
     * Layout: for each [b, d, :], element [0] = td_error, element [1] = E[Q] */
    int base = (b * 3 + d) * num_atoms;
    float td_error = save_eq_td[base + 0];
    float e_q = save_eq_td[base + 1];

    /* C51 support: z_j = v_min + j * delta_z */
    float v_min = {v_min:.1}f;
    float v_max = {v_max:.1}f;
    float delta_z = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f;
    float z_j = v_min + (float)j * delta_z;

    /* Gradient of MSE loss through softmax expectation:
     * L = 0.5 * td^2 where td = E[Q] - target_Q and E[Q] = sum_j(p_j * z_j)
     * dL/d_logit_j = td * d(E[Q])/d_logit_j
     *              = td * p_j * (z_j - E[Q])
     * (softmax Jacobian: dE[Q]/d_logit_j = p_j * (z_j - sum_k(p_k * z_k)) = p_j * (z_j - E[Q]))
     */
    float d_combined = isw * td_error * p_j * (z_j - e_q);

    // Route through dueling: d_value[b,j] += d_combined
    atomicAdd(&d_value_logits[b * num_atoms + j], d_combined);

    // Factored action decode
    int factored = actions[b];
    int max_action = b0_size * b1_size * b2_size;
    if (factored < 0 || factored >= max_action) factored = 0;
    int branch_sizes[3];
    branch_sizes[0] = b0_size;
    branch_sizes[1] = b1_size;
    branch_sizes[2] = b2_size;

    int a_d;
    if (d == 0) a_d = factored / (b1_size * b2_size);
    else if (d == 1) a_d = (factored / b2_size) % b1_size;
    else a_d = factored % b2_size;

    int A_d = branch_sizes[d];
    float inv_A = 1.0f / (float)A_d;

    // Compute per-branch adv offset
    int branch_offset = 0;
    for (int dd = 0; dd < d; dd++)
        branch_offset += branch_sizes[dd] * num_atoms;

    // d_adv[b, d, a, j] = d_combined * (delta(a, a_d) - 1/A_d)
    for (int a = 0; a < A_d; a++) {{
        float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A);
        int adv_idx = b * total_branch_atoms + branch_offset + a * num_atoms + j;
        atomicAdd(&d_adv_logits[adv_idx], d_combined * dueling_grad);
    }}
}}
"#,
    v_min = config.v_min,
    v_max = config.v_max,
    );

    let context = stream.context();
    let ptx = crate::cuda_pipeline::compile_ptx_for_device(&src, &context)
        .map_err(|e| MLError::ModelError(format!("mse_grad_kernel compilation: {e}")))?;
    let module = context.load_module(ptx)
        .map_err(|e| MLError::ModelError(format!("mse_grad module load: {e}")))?;
    module.load_function("mse_grad_kernel")
        .map_err(|e| MLError::ModelError(format!("mse_grad_kernel load: {e}")))
}

/// Compile the `compute_expected_q` NVRTC kernel.
///
/// Converts C51 distributional logits (value + advantage, branching dueling)
/// to expected Q-values for the ad-hoc validation forward pass.
/// One thread per sample; iterates over branches and atoms.
fn compile_expected_q_kernel(
    stream: &Arc<CudaStream>,
) -> Result<CudaFunction, MLError> {
    let src = r#"
extern "C" __global__ void compute_expected_q(
    const float* __restrict__ v_logits,   // [N, num_atoms]
    const float* __restrict__ b_logits,   // [N, (b0+b1+b2)*num_atoms]
    float* __restrict__ q_values,         // [N, b0+b1+b2]
    int N, int num_atoms,
    int b0_size, int b1_size, int b2_size,
    float v_min, float v_max)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;

    int total_actions = b0_size + b1_size + b2_size;
    float dz = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f;

    // Value logits for this sample: [num_atoms]
    const float* val = v_logits + (long long)i * num_atoms;

    int branch_sizes[3];
    branch_sizes[0] = b0_size;
    branch_sizes[1] = b1_size;
    branch_sizes[2] = b2_size;

    int adv_offset = 0;
    int q_offset = 0;
    for (int d = 0; d < 3; d++) {
        int bd = branch_sizes[d];
        for (int a = 0; a < bd; a++) {
            const float* adv = b_logits + (long long)i * total_actions * num_atoms
                             + (long long)(adv_offset + a) * num_atoms;

            // Compute mean advantage logit sum for this branch (for dueling centering)
            float mean_adv_sum = 0.0f;
            for (int aa = 0; aa < bd; aa++) {
                const float* adv_aa = b_logits + (long long)i * total_actions * num_atoms
                                    + (long long)(adv_offset + aa) * num_atoms;
                for (int j = 0; j < num_atoms; j++) {
                    mean_adv_sum += adv_aa[j];
                }
            }
            float mean_adv_per_atom = mean_adv_sum / (float)(bd * num_atoms);

            // Numerically stable log_softmax over combined = val[j] + adv[j] - mean_adv_per_atom
            float max_logit = -1e30f;
            for (int j = 0; j < num_atoms; j++) {
                float combined = val[j] + adv[j] - mean_adv_per_atom;
                if (combined > max_logit) max_logit = combined;
            }
            float sum_exp = 0.0f;
            for (int j = 0; j < num_atoms; j++) {
                float combined = val[j] + adv[j] - mean_adv_per_atom;
                sum_exp += expf(combined - max_logit);
            }
            float log_sum = logf(sum_exp + 1e-8f) + max_logit;

            // Expected Q = sum_j( softmax_j * z_j )
            float eq = 0.0f;
            for (int j = 0; j < num_atoms; j++) {
                float combined = val[j] + adv[j] - mean_adv_per_atom;
                float p = expf(combined - log_sum);
                float z = v_min + (float)j * dz;
                eq += p * z;
            }

            q_values[(long long)i * total_actions + q_offset + a] = eq;
        }
        adv_offset += bd;
        q_offset += bd;
    }
}
"#;

    let context = stream.context();
    let ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context)
        .map_err(|e| MLError::ModelError(format!("compute_expected_q compilation: {e}")))?;
    let module = context.load_module(ptx)
        .map_err(|e| MLError::ModelError(format!("compute_expected_q module load: {e}")))?;
    module.load_function("compute_expected_q")
        .map_err(|e| MLError::ModelError(format!("compute_expected_q load: {e}")))
}

/// Compile the Q-value statistics reduction kernel.
///
/// Single-block kernel: computes [avg_max_q, q_min, q_max, q_mean, q_variance]
/// from the q_out_buf [N, total_actions]. Grid: (1,1,1), Block: (256,1,1).
fn compile_q_stats_kernel(
    stream: &Arc<CudaStream>,
) -> Result<CudaFunction, MLError> {
    let src = r#"
extern "C" __global__ void q_stats_reduce(
    const float* __restrict__ q_values,  // [N, total_actions]
    float* __restrict__ out,             // [5]: avg_max_q, q_min, q_max, q_mean, q_var
    int N,
    int total_actions)
{
    // Single block — thread 0 does the reduction sequentially.
    // N is tiny (<=10 samples), total_actions=11. Total work: 110 elements.
    if (threadIdx.x != 0) return;

    float global_min = 1e30f;
    float global_max = -1e30f;
    float global_sum = 0.0f;
    float sum_max_q = 0.0f;
    int total = N * total_actions;

    for (int i = 0; i < N; i++) {
        float row_max = -1e30f;
        for (int a = 0; a < total_actions; a++) {
            float v = q_values[i * total_actions + a];
            if (v < global_min) global_min = v;
            if (v > global_max) global_max = v;
            if (v > row_max) row_max = v;
            global_sum += v;
        }
        sum_max_q += row_max;
    }

    float mean = (total > 0) ? global_sum / (float)total : 0.0f;
    float var_sum = 0.0f;
    for (int i = 0; i < total; i++) {
        float d = q_values[i] - mean;
        var_sum += d * d;
    }
    float variance = (total > 0) ? var_sum / (float)total : 0.0f;
    float avg_max = (N > 0) ? sum_max_q / (float)N : 0.0f;

    out[0] = avg_max;
    out[1] = global_min;
    out[2] = global_max;
    out[3] = mean;
    out[4] = variance;
}
"#;

    let context = stream.context();
    let ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context)
        .map_err(|e| MLError::ModelError(format!("q_stats_reduce compilation: {e}")))?;
    let module = context.load_module(ptx)
        .map_err(|e| MLError::ModelError(format!("q_stats module load: {e}")))?;
    module.load_function("q_stats_reduce")
        .map_err(|e| MLError::ModelError(format!("q_stats_reduce load: {e}")))
}

// ── Shared memory sizing ────────────────────────────────────────────────────

/// Query the hardware's max shared memory per block via cuDeviceGetAttribute.
/// Returns usable bytes (floored at 48KB default).
///
/// Uses `CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN` to get the
/// true hardware ceiling. Uses 80% to leave headroom for kernel state.
/// Falls back to 48KB if the query fails.
pub(crate) fn query_max_shmem_bytes() -> usize {
    use cudarc::driver::sys::{cuDeviceGetAttribute, CUdevice_attribute};
    let mut max_shmem: i32 = 49152;
    let result = unsafe {
        cuDeviceGetAttribute(
            &mut max_shmem,
            CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN,
            0,
        )
    };
    if result != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
        tracing::warn!("cuDeviceGetAttribute(MAX_SHARED_MEMORY_PER_BLOCK_OPTIN) failed: {result:?}");
    }
    // A100 reports 164KB, H100 reports 228KB — these genuinely support opt-in.
    // Consumer GPUs (RTX 3050: 100KB, RTX 4090: 100KB) report > 48KB but
    // fail with CUDA_ERROR_ILLEGAL_ADDRESS on dynamic shmem > 48KB with
    // some driver versions. Use 48KB safe default for < 164KB.
    let usable = if max_shmem >= 164_000 {
        (max_shmem as usize) * 80 / 100
    } else {
        49152
    };
    usable.max(49152)
}

// ── Device-to-device copy helpers ────────────────────────────────────────────

/// Extract raw CUDA device pointer (CUdeviceptr = u64) from a CudaSlice.
///
/// Same pattern as `raw_device_ptr` in `gpu_weights.rs` — wraps the SyncOnDrop
/// guard in ManuallyDrop to skip read event recording (safe on same stream).
pub(crate) fn raw_device_ptr(slice: &CudaSlice<f32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = std::mem::ManuallyDrop::new(guard);
    ptr
}

/// Extract raw CUDA device pointer from an i32 CudaSlice.
///
/// Same pattern as `raw_device_ptr` but for `CudaSlice<i32>` (actions buffer).
fn raw_device_ptr_i32(slice: &CudaSlice<i32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = std::mem::ManuallyDrop::new(guard);
    ptr
}

/// Extract raw CUDA device pointer from a `CudaSlice<u32>`.
fn raw_device_ptr_u32(slice: &CudaSlice<u32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = std::mem::ManuallyDrop::new(guard);
    ptr
}

/// Extract raw CUDA device pointer from a `CudaSlice<u16>` (BF16 weight buffers).
fn raw_device_ptr_u16(slice: &CudaSlice<u16>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = std::mem::ManuallyDrop::new(guard);
    ptr
}

/// Async device-to-device memcpy with error context.
pub(crate) fn dtod_copy(
    dst: u64,
    src: u64,
    num_bytes: usize,
    stream: &Arc<CudaStream>,
    idx: usize,
    op: &str,
) -> Result<(), MLError> {
    // Safety: caller guarantees src/dst are valid device pointers within
    // the same CUDA context, num_bytes ≤ allocation size of both regions,
    // and stream belongs to the same context.
    unsafe {
        cudarc::driver::result::memcpy_dtod_async(dst, src, num_bytes, stream.cu_stream())
            .map_err(|e| MLError::ModelError(format!("DtoD {op}[{idx}]: {e}")))?;
    }
    Ok(())
}

// ── Candle tensor → CudaSlice DtoD copy helpers ────────────────────────────

/// DtoD copy from a CudaSlice<f32> to a pre-allocated CudaSlice<f32>.
fn dtod_from_slice_f32(
    src: &CudaSlice<f32>,
    dst: &CudaSlice<f32>,
    num_elements: usize,
    stream: &Arc<CudaStream>,
    ctx: &str,
) -> Result<(), MLError> {
    let src_ptr = raw_device_ptr(src, stream);
    let dst_ptr = raw_device_ptr(dst, stream);
    dtod_copy(dst_ptr, src_ptr, num_elements * std::mem::size_of::<f32>(), stream, 0, ctx)
}

/// DtoD copy from a CudaSlice<u32> to a CudaSlice<i32>.
///
/// U32 and I32 are bit-compatible for action indices in [0, 44].
fn dtod_from_u32_to_i32(
    src: &CudaSlice<u32>,
    dst: &CudaSlice<i32>,
    num_elements: usize,
    stream: &Arc<CudaStream>,
    ctx: &str,
) -> Result<(), MLError> {
    let src_ptr = raw_device_ptr_u32(src, stream);
    let dst_ptr = raw_device_ptr_i32(dst, stream);
    dtod_copy(dst_ptr, src_ptr, num_elements * std::mem::size_of::<u32>(), stream, 0, ctx)
}

// ── Allocation helpers ──────────────────────────────────────────────────────

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!("GpuDqnTrainer alloc {name} ({n} f32): {e}"))
    })
}

fn alloc_i32(
    stream: &Arc<CudaStream>,
    n: usize,
    name: &str,
) -> Result<CudaSlice<i32>, MLError> {
    stream.alloc_zeros::<i32>(n).map_err(|e| {
        MLError::ModelError(format!("GpuDqnTrainer alloc {name} ({n} i32): {e}"))
    })
}

fn alloc_u16(
    stream: &Arc<CudaStream>,
    n: usize,
    name: &str,
) -> Result<CudaSlice<u16>, MLError> {
    stream.alloc_zeros::<u16>(n).map_err(|e| {
        MLError::ModelError(format!("GpuDqnTrainer alloc {name} ({n} u16): {e}"))
    })
}

/// DtoD copy from a CudaSlice<u16> (BF16 bit pattern) to a pre-allocated CudaSlice<u16>.
fn dtod_from_bf16(
    src: &CudaSlice<u16>,
    dst: &CudaSlice<u16>,
    num_elements: usize,
    stream: &Arc<CudaStream>,
    ctx: &str,
) -> Result<(), MLError> {
    let (src_ptr, src_guard) = src.device_ptr(stream);
    let _no_drop = std::mem::ManuallyDrop::new(src_guard);
    let (dst_ptr, dst_guard) = dst.device_ptr(stream);
    let _no_drop2 = std::mem::ManuallyDrop::new(dst_guard);
    dtod_copy(dst_ptr, src_ptr, num_elements * std::mem::size_of::<u16>(), stream, 0, ctx)
}

/// Compile ensemble aggregate + diversity kernels from `ensemble_kernels.cu`.
///
/// Returns `(ensemble_aggregate_kernel, ensemble_diversity_kernel)`.
/// Both kernels take all sizes as runtime arguments — no network-dim #defines needed.
pub(crate) fn compile_ensemble_kernels(
    stream: &Arc<CudaStream>,
    state_dim: usize,
) -> Result<(CudaFunction, CudaFunction), MLError> {
    let dim_overrides = format!(
        "#define STATE_DIM {state_dim}\n\
         #define MARKET_DIM 42\n\
         #define PORTFOLIO_DIM 8\n",
    );
    let common_src = include_str!("common_device_functions.cuh");
    let kernel_src = include_str!("ensemble_kernels.cu");
    let full_source = format!("{dim_overrides}\n{common_src}\n{kernel_src}");

    let context = stream.context();
    let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context)
        .map_err(|e| MLError::ModelError(format!("ensemble_kernels compilation: {e}")))?;
    let module = context.load_module(ptx)
        .map_err(|e| MLError::ModelError(format!("ensemble_kernels module load: {e}")))?;

    let aggregate = module
        .load_function("ensemble_aggregate_kernel")
        .map_err(|e| MLError::ModelError(format!("ensemble_aggregate_kernel load: {e}")))?;
    let diversity = module
        .load_function("ensemble_diversity_kernel")
        .map_err(|e| MLError::ModelError(format!("ensemble_diversity_kernel load: {e}")))?;

    Ok((aggregate, diversity))
}

/// Launch the bf16_to_f32 CUDA kernel: converts BF16 CudaSlice<u16> → F32 CudaSlice<f32>.
///
/// Grid: ceil(n/256), Block: 256. Kernel signature: `bf16_to_f32_kernel(src, dst, n)`.
/// Zero Candle tensor involvement — pure cudarc kernel launch.
pub(crate) fn launch_bf16_to_f32(
    kernel: &CudaFunction,
    src: &CudaSlice<u16>,
    dst: &CudaSlice<f32>,
    n: usize,
    stream: &Arc<CudaStream>,
    ctx: &str,
) -> Result<(), MLError> {
    let n_i32 = n as i32;
    let launch_cfg = LaunchConfig {
        grid_dim: (((n + 255) / 256) as u32, 1, 1),
        block_dim: (256, 1, 1),
        shared_mem_bytes: 0,
    };

    // Safety: src is CudaSlice<u16> holding BF16 bit patterns, dst is CudaSlice<f32>,
    // both are valid GPU allocations on the same context. Kernel reads src[0..n] and
    // writes dst[0..n].
    unsafe {
        stream
            .launch_builder(kernel)
            .arg(src)
            .arg(dst)
            .arg(&n_i32)
            .launch(launch_cfg)
            .map_err(|e| MLError::ModelError(format!("bf16_to_f32_kernel {ctx}: {e}")))?;
    }
    Ok(())
}

/// Compile the CQL (Conservative Q-Learning) logit gradient kernel.
///
/// Computes dCQL/d_value_logits and dCQL/d_adv_logits for the Branching Dueling
/// C51 architecture. The CQL penalty per branch is:
///   `penalty = alpha * (logsumexp_a(Q(a)) - Q(a_taken))`
///
/// The gradient flows through the softmax expectation:
///   `dCQL/d_logit[a,j] = alpha * dCQL_dQ[a] * p[j] * (z[j] - Q[a])`
/// where `dCQL_dQ[a] = softmax_Q[a] - I(a == a_taken)` and `p[j] = softmax(combined_logits)[j]`.
///
/// One thread per sample. Iterates over the 3 branches (exposure, order, urgency).
fn compile_cql_logit_grad_kernel(
    stream: &Arc<CudaStream>,
) -> Result<CudaFunction, MLError> {
    let src = r#"
extern "C" __global__ void cql_logit_grad_kernel(
    const float* __restrict__ v_logits,      // [N, num_atoms]
    const float* __restrict__ adv_logits,    // [N, total_actions * num_atoms]
    const int*   __restrict__ actions,       // [N] factored action indices (0-44)
    float*       __restrict__ d_v_logits,    // [N, num_atoms] output
    float*       __restrict__ d_adv_logits,  // [N, total_actions * num_atoms] output
    float cql_alpha,
    int N, int num_atoms,
    int b0_size, int b1_size, int b2_size,
    float v_min, float v_max)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;

    int total_actions = b0_size + b1_size + b2_size;
    float dz = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f;

    // Decode factored action: action = a0 * b1 * b2 + a1 * b2 + a2
    int action = actions[i];
    int a2_taken = action % b2_size;
    int a1_taken = (action / b2_size) % b1_size;
    int a0_taken = action / (b1_size * b2_size);

    // Clamp to valid range (safety)
    if (a0_taken >= b0_size) a0_taken = b0_size - 1;
    if (a1_taken >= b1_size) a1_taken = b1_size - 1;
    if (a2_taken >= b2_size) a2_taken = b2_size - 1;

    int branch_taken[3];
    branch_taken[0] = a0_taken;
    branch_taken[1] = a1_taken;
    branch_taken[2] = a2_taken;

    int branch_sizes[3];
    branch_sizes[0] = b0_size;
    branch_sizes[1] = b1_size;
    branch_sizes[2] = b2_size;

    const float* val = v_logits + (long long)i * num_atoms;

    // Zero value logit gradient accumulator
    float d_val_accum[256]; // max num_atoms (C51: 51, generous buffer)
    for (int j = 0; j < num_atoms && j < 256; j++) d_val_accum[j] = 0.0f;

    int adv_offset = 0;
    for (int d = 0; d < 3; d++) {
        int bd = branch_sizes[d];
        int a_taken_d = branch_taken[d];

        // Step 1: Compute expected Q for each action in this branch
        // using dueling: combined_logit[j] = val[j] + adv[a,j] - mean_adv
        float q_values_branch[9]; // max branch size = 5 (exposure), 3+3=6 total per branch

        // Compute mean advantage per atom (for dueling centering)
        float mean_adv_sum = 0.0f;
        for (int aa = 0; aa < bd; aa++) {
            const float* adv_aa = adv_logits + (long long)i * total_actions * num_atoms
                                + (long long)(adv_offset + aa) * num_atoms;
            for (int j = 0; j < num_atoms; j++) mean_adv_sum += adv_aa[j];
        }
        float mean_adv_per_atom = mean_adv_sum / (float)(bd * num_atoms);

        for (int a = 0; a < bd; a++) {
            const float* adv = adv_logits + (long long)i * total_actions * num_atoms
                             + (long long)(adv_offset + a) * num_atoms;

            // softmax over atoms + expected value
            float max_logit = -1e30f;
            for (int j = 0; j < num_atoms; j++) {
                float c = val[j] + adv[j] - mean_adv_per_atom;
                if (c > max_logit) max_logit = c;
            }
            float sum_exp = 0.0f;
            for (int j = 0; j < num_atoms; j++) {
                sum_exp += expf(val[j] + adv[j] - mean_adv_per_atom - max_logit);
            }
            float log_sum = logf(sum_exp + 1e-8f) + max_logit;

            float eq = 0.0f;
            for (int j = 0; j < num_atoms; j++) {
                float p = expf(val[j] + adv[j] - mean_adv_per_atom - log_sum);
                float z = v_min + (float)j * dz;
                eq += p * z;
            }
            q_values_branch[a] = eq;
        }

        // Step 2: CQL penalty gradient w.r.t. Q-values
        // dCQL/dQ[a] = softmax_Q[a] - I(a == a_taken)
        float max_q = q_values_branch[0];
        for (int a = 1; a < bd; a++)
            if (q_values_branch[a] > max_q) max_q = q_values_branch[a];
        float sum_exp_q = 0.0f;
        for (int a = 0; a < bd; a++)
            sum_exp_q += expf(q_values_branch[a] - max_q);

        float d_cql_dq[9];
        for (int a = 0; a < bd; a++) {
            float softmax_q = expf(q_values_branch[a] - max_q) / (sum_exp_q + 1e-8f);
            d_cql_dq[a] = cql_alpha * (softmax_q - ((a == a_taken_d) ? 1.0f : 0.0f));
        }

        // Step 3: Chain rule through softmax-expectation to get d_logits
        for (int a = 0; a < bd; a++) {
            const float* adv = adv_logits + (long long)i * total_actions * num_atoms
                             + (long long)(adv_offset + a) * num_atoms;
            float* d_adv = d_adv_logits + (long long)i * total_actions * num_atoms
                         + (long long)(adv_offset + a) * num_atoms;

            // Recompute p[j] for this action
            float max_logit = -1e30f;
            for (int j = 0; j < num_atoms; j++) {
                float c = val[j] + adv[j] - mean_adv_per_atom;
                if (c > max_logit) max_logit = c;
            }
            float sum_exp = 0.0f;
            for (int j = 0; j < num_atoms; j++) {
                sum_exp += expf(val[j] + adv[j] - mean_adv_per_atom - max_logit);
            }
            float log_sum = logf(sum_exp + 1e-8f) + max_logit;

            float eq = q_values_branch[a];

            for (int j = 0; j < num_atoms; j++) {
                float p = expf(val[j] + adv[j] - mean_adv_per_atom - log_sum);
                float z = v_min + (float)j * dz;
                // d_combined_logit[j] = d_cql_dq[a] * p * (z - Q)
                float d_combined = d_cql_dq[a] * p * (z - eq);

                // Split combined gradient to adv and val
                d_adv[j] = d_combined;
                if (j < 256) d_val_accum[j] += d_combined;
            }
        }
        adv_offset += bd;
    }

    // Write accumulated value logit gradient (summed across all branches and actions)
    float* d_val = d_v_logits + (long long)i * num_atoms;
    for (int j = 0; j < num_atoms && j < 256; j++) {
        d_val[j] = d_val_accum[j];
    }
}
"#;

    let context = stream.context();
    let ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context)
        .map_err(|e| MLError::ModelError(format!("cql_logit_grad_kernel compilation: {e}")))?;
    let module = context.load_module(ptx)
        .map_err(|e| MLError::ModelError(format!("cql_logit_grad module load: {e}")))?;
    module.load_function("cql_logit_grad_kernel")
        .map_err(|e| MLError::ModelError(format!("cql_logit_grad_kernel load: {e}")))
}
