#![allow(unsafe_code)] // Required for f32->u32 CudaSlice reinterpret cast in PER index path
//! Fused CUDA Training Module
//!
//! High-performance H100-optimized training path that replaces 2,100+ Candle kernel
//! dispatches per batch with 3 fused CUDA kernels captured in a CUDA Graph:
//!
//! 1. **Forward + Loss kernel** -- shared trunk, 3 branching advantage heads, C51 distributional loss
//! 2. **Backward kernel** -- full backprop through all layers with gradient clipping
//! 3. **Adam optimizer kernel** -- fused parameter update with weight decay
//!
//! After the first batch, all 3 kernels are captured into a CUDA Graph and replayed
//! on subsequent steps with zero launch overhead.
//!
//! ## Weight Sync Architecture (GPU-native EMA)
//!
//! MlDevice pointers in `CudaSlice` buffers are stable across in-place updates, so
//! the CUDA Graph remains valid across Polyak EMA target updates:
//!
//! 1. **After fused step**: online `CudaSlice` weights are already updated by the Adam kernel
//! 2. **GPU Polyak EMA**: `target[i] = (1-tau)*target[i] + tau*online[i]` via EMA kernel
//! 3. **No reverse/forward sync per step** -- flat buffer sync deferred to epoch boundary
//!
//! This eliminates ~120 D2D copies + 120 Candle flatten/contiguous ops per step
//! (~13 GB wasted PCIe per epoch).

use std::sync::Arc;

use anyhow::Result;
use ml_core::device::MlDevice;
use ml_core::state_layout;
use tracing::info;

use cudarc::driver::DevicePtr;
use cudarc::driver::sys as cuda_sys;

use crate::cuda_pipeline::mapped_pinned::{MappedF32Buffer, MappedI32Buffer};
use crate::cuda_pipeline::gpu_attention::{GpuAttention, GpuAttentionConfig};
use crate::cuda_pipeline::gpu_tlob::GpuTlob;
use crate::cuda_pipeline::gpu_dqn_trainer::{
    GpuDqnTrainConfig, GpuDqnTrainer, TAU_EFF_INDEX,
};
use crate::cuda_pipeline::sp5_isv_slots::{
    BUDGET_IQN_BASE, BUDGET_ENS_BASE,
    ADAM_BETA1_BASE, ADAM_BETA2_BASE, ADAM_EPS_BASE,
    IQN_TAU_BASE,
};
use crate::cuda_pipeline::gpu_her::{GpuHer, GpuHerConfig, parse_her_strategy};
use crate::cuda_pipeline::gpu_iql_trainer::{GpuIqlConfig, GpuIqlTrainer};
use crate::cuda_pipeline::gpu_iqn_head::{GpuIqnConfig, GpuIqnHead};
use crate::cuda_pipeline::gpu_weights::{
    self, BranchingWeightSet, DuelingWeightSet,
};
use super::config::DQNAgentType;
use super::DQNHyperparameters;

/// Lightweight result from fused GPU training — CPU scalars only, no GPU
/// allocations or HtoD copies.  The training guard reads loss/grad_norm
/// directly from the GPU-resident buffers via `loss_gpu_buf()` /
/// `grad_norm_gpu_buf()`, so uploading scalars back to GPU is unnecessary
/// (and on H100 CUDA 13 it causes a sync-stall hang).
pub(crate) struct FusedStepResult {
    pub total_loss: f32,
    pub grad_norm: f32,
}

// ── Mapped-pinned staging helpers (HtoD elimination) ─────────────────────────
// Per `feedback_no_htod_htoh_only_mapped_pinned.md`, mapped pinned
// (cuMemHostAlloc DEVICEMAP) is the only allowed CPU↔GPU path. These helpers
// stage CPU data through a transient mapped pinned buffer and DtoD-copy into
// a GPU-resident `CudaSlice<T>`, preserving the existing kernel-arg surface
// (e.g., `her_source_indices_gpu` consumed by gpu_her.rs out of scope).

fn staging_upload_i32(
    stream: &Arc<cudarc::driver::CudaStream>,
    n: usize,
    data: &[i32],
    label: &str,
) -> Result<cudarc::driver::CudaSlice<i32>> {
    use cudarc::driver::DevicePtrMut;
    assert!(data.len() <= n, "{label}: data.len {} > buf.len {n}", data.len());
    let mut buf = stream.alloc_zeros::<i32>(n)
        .map_err(|e| anyhow::anyhow!("{label} alloc: {e}"))?;
    let staging = unsafe { MappedI32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("{label} staging alloc: {e}"))?;
    staging.write_from_slice(data);
    let n_bytes = n * std::mem::size_of::<i32>();
    {
        let (dst_ptr, _g) = buf.device_ptr_mut(stream);
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr, staging.dev_ptr, n_bytes, stream.cu_stream(),
            ).map_err(|e| anyhow::anyhow!("{label} dtod: {e}"))?;
        }
    }
    stream.synchronize()
        .map_err(|e| anyhow::anyhow!("{label} dtod sync: {e}"))?;
    Ok(buf)
}

fn staging_upload_f32(
    stream: &Arc<cudarc::driver::CudaStream>,
    n: usize,
    data: &[f32],
    label: &str,
) -> Result<cudarc::driver::CudaSlice<f32>> {
    use cudarc::driver::DevicePtrMut;
    assert!(data.len() <= n, "{label}: data.len {} > buf.len {n}", data.len());
    let mut buf = stream.alloc_zeros::<f32>(n)
        .map_err(|e| anyhow::anyhow!("{label} alloc: {e}"))?;
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("{label} staging alloc: {e}"))?;
    staging.write_from_slice(data);
    let n_bytes = n * std::mem::size_of::<f32>();
    {
        let (dst_ptr, _g) = buf.device_ptr_mut(stream);
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr, staging.dev_ptr, n_bytes, stream.cu_stream(),
            ).map_err(|e| anyhow::anyhow!("{label} dtod: {e}"))?;
        }
    }
    stream.synchronize()
        .map_err(|e| anyhow::anyhow!("{label} dtod sync: {e}"))?;
    Ok(buf)
}

/// Child sub-graph: owns both the CUgraph topology and an instantiated CUgraphExec.
///
/// The exec comes directly from cudarc's `end_capture` instantiation — we do NOT
/// clone + re-instantiate because `cuGraphClone` corrupts cublasLtMatmul HMMA tile
/// scheduling state on Hopper (sm_90), causing graph replay to hang.
struct ChildGraph {
    graph: cuda_sys::CUgraph,
    exec: cuda_sys::CUgraphExec,
}

impl ChildGraph {
    /// Create an additional exec from an existing graph (for eval replays).
    /// Does NOT clone — instantiates directly from the original graph.
    fn instantiate_extra_exec(graph: cuda_sys::CUgraph) -> Result<cuda_sys::CUgraphExec> {
        let mut exec: cuda_sys::CUgraphExec = std::ptr::null_mut();
        let result = unsafe { cuda_sys::cuGraphInstantiateWithFlags(&mut exec, graph, 0) };
        if result != cuda_sys::cudaError_enum::CUDA_SUCCESS {
            anyhow::bail!("cuGraphInstantiateWithFlags failed: {result:?}");
        }
        Ok(exec)
    }

    fn launch(&self, stream: cuda_sys::CUstream) -> Result<()> {
        let result = unsafe { cuda_sys::cuGraphLaunch(self.exec, stream) };
        if result != cuda_sys::cudaError_enum::CUDA_SUCCESS {
            anyhow::bail!("cuGraphLaunch failed: {result:?}");
        }
        Ok(())
    }
}

impl Drop for ChildGraph {
    fn drop(&mut self) {
        if !self.exec.is_null() {
            unsafe { cuda_sys::cuGraphExecDestroy(self.exec); }
        }
        if !self.graph.is_null() {
            unsafe { cuda_sys::cuGraphDestroy(self.graph); }
        }
    }
}

// Safety: The graph is bound to a CUDA context. FusedTrainingCtx is always used
// from the thread that created the context. No concurrent access occurs.
unsafe impl Send for ChildGraph {}
unsafe impl Sync for ChildGraph {}

/// Composed parent graph -- single `cuGraphLaunch` per training step.
///
/// Contains 7 child sub-graphs (spectral, forward, ddqn, aux, post_aux, adam, maintenance)
/// added sequentially via `cuGraphAddChildGraphNode`.
struct ParentGraph {
    exec: cuda_sys::CUgraphExec,
    graph: cuda_sys::CUgraph,
}

impl ParentGraph {
    /// Launch the composed parent graph on the given stream.
    fn launch(&self, stream: cuda_sys::CUstream) -> Result<()> {
        let result = unsafe { cuda_sys::cuGraphLaunch(self.exec, stream) };
        if result != cuda_sys::cudaError_enum::CUDA_SUCCESS {
            anyhow::bail!("cuGraphLaunch failed: {result:?}");
        }
        Ok(())
    }
}

impl Drop for ParentGraph {
    fn drop(&mut self) {
        unsafe {
            if !self.exec.is_null() {
                cuda_sys::cuGraphExecDestroy(self.exec);
            }
            if !self.graph.is_null() {
                cuda_sys::cuGraphDestroy(self.graph);
            }
        }
    }
}

// Safety: same as ChildGraph -- single-owner, same CUDA context.
unsafe impl Send for ParentGraph {}
unsafe impl Sync for ParentGraph {}

/// Per-component gradient norm budget fractions for auxiliary objectives.
/// C51 gets whatever remains: `1.0 - sum(active_auxiliary_budgets)`.
/// Default: IQN=40%, CQL=10%, ENS=5%, C51=45%.
/// Old hardcoded: IQN=75% — starved C51 directional learning (grad_norm frozen at 0.655).
pub(crate) struct GradBudget {
    pub cql: f32,
    pub iqn: f32,
    pub ens: f32,
}

/// Fused CUDA training context -- owns the `GpuDqnTrainer` and extracted weight sets.
///
/// Weight sets are extracted from the Candle flat buffer at initialization via direct
/// DtoD copy (no Candle flatten/cast/contiguous -- Vars are guaranteed F32
/// contiguous by `BranchingDuelingQNetwork::ensure_f32_contiguous()`).
///
/// Per-step training operates entirely on GPU `CudaSlice` buffers: the Adam
/// kernel updates online weights, then the EMA kernel blends online into target
/// weights in-place. The `CudaSlice` weights are the source of truth during
/// training -- the flat buffer is never accessed per-step.
///
/// flat buffer reverse sync is deferred to epoch boundary for checkpointing and
/// Candle forward (experience collector, inference). MlDevice pointers are stable
/// across in-place updates, so the CUDA Graph stays valid.
// PhaseEvents removed: single parent graph launch makes per-child event profiling
// irrelevant. Wall-clock epoch timing at epoch boundary is sufficient.

pub(crate) struct FusedTrainingCtx {
    trainer: GpuDqnTrainer,
    online_dueling: DuelingWeightSet,
    online_branching: BranchingWeightSet,
    target_dueling: DuelingWeightSet,
    target_branching: BranchingWeightSet,
    /// Backing GPU allocations for online/target weight sets.
    /// Must outlive the weight sets (pointers reference this memory).
    /// `None` when weight sets are zero-copy views into `params_buf`.
    #[allow(dead_code)]
    _weight_backings: Option<(
        gpu_weights::DuelingWeightBacking, gpu_weights::BranchingWeightBacking,
        gpu_weights::DuelingWeightBacking, gpu_weights::BranchingWeightBacking,
    )>,
    /// Backing GPU allocations for ensemble extra head weight sets.
    #[allow(dead_code)]
    _ensemble_backings: Vec<(gpu_weights::DuelingWeightBacking, gpu_weights::BranchingWeightBacking)>,
    stream: Arc<cudarc::driver::CudaStream>,
    /// Batch size at creation time -- must match `current_batch_size` to reuse CUDA Graph.
    batch_size: usize,
    /// Steps since last flat buffer sync (deferred to epoch boundary).
    steps_since_varmap_sync: usize,
    /// GPU-resident copy of best-epoch params. Pure DtoD — no Candle, no safetensors.
    /// Saved on new best epoch, restored before evaluation.
    best_params_snapshot: Option<cudarc::driver::CudaSlice<f32>>,
    /// Event recorded after restore_best_params writes complete on the trainer's stream.
    /// The evaluator waits on this event on its own stream (cuStreamWaitEvent) to ensure
    /// cross-stream visibility without CPU sync.
    restore_event: Option<cudarc::driver::CudaEvent>,
    /// Last observed C51 raw gradient norm (before budget clip). Updated every 1000 steps.
    last_c51_raw_norm: f32,
    /// Last observed combined gradient norm (from Adam readback).
    last_combined_norm: f32,
    last_loss: f32,
    loss_accum: f64,
    grad_accum: f64,
    accum_steps: u64,
    /// GPU HER relabeler -- initialized when `her_ratio > 0.0`.
    /// When active, `run_full_step()` splits the batch into normal + HER portions,
    /// runs the relabel kernel on the HER portion, and merges before training.
    pub(crate) gpu_her: Option<GpuHer>,
    /// Pre-computed HER source indices on GPU — avoids per-step Vec + HtoD.
    /// Contains `[normal_count, normal_count+1, ..., batch_size-1]`.
    pub(crate) her_source_indices_gpu: Option<cudarc::driver::CudaSlice<i32>>,
    /// Pre-allocated HER reward ones on GPU — avoids per-step vec![1.0; N] + HtoD.
    pub(crate) her_reward_ones_gpu: Option<cudarc::driver::CudaSlice<f32>>,
    /// Compiled HER in-place relabel kernel (loaded from dqn_utility cubin).
    pub(crate) her_inplace_kernel: Option<cudarc::driver::CudaFunction>,
    /// GPU IQL value trainer (τ=high, default 0.7) — always active.
    /// Trains V(s) with expectile regression after each DQN training step.
    pub(crate) gpu_iql: GpuIqlTrainer,
    /// GPU IQL value trainer (τ=low, derived as 1-τ_high) — always active.
    /// Provides the low-expectile V(s) for expectile gap uncertainty estimation.
    pub(crate) gpu_iql_low: GpuIqlTrainer,
    /// Trains quantile head that shares the DQN trunk (h_s2) independently.
    /// Combined loss: L_total = L_c51 + iqn_lambda × L_iqn.
    pub(crate) gpu_iqn: Option<GpuIqnHead>,
    /// CVaR scales buffer (kept alive so device pointer remains valid).
    cvar_scales_buf: Option<cudarc::driver::CudaSlice<f32>>,
    /// GPU multi-head feature attention over h_s2 (post-graph, 1-step lag).
    /// When Some, applied after EMA update, before IQN, each training step.
    pub(crate) gpu_attention: Option<GpuAttention>,
    /// TLOB attention module: OFI[32]→TLOB[16], trained end-to-end via DQN reward.
    /// Forward runs after state gather (before cuBLAS forward). Backward runs after
    /// trunk backward. D.8 Plan 2 Task 6C.
    ///
    /// Always present — `GpuTlob::new` is unconditional in the trainer ctor.
    /// Previously wrapped in `Option<_>` to graceful-degrade on cuBLAS-Lt
    /// heuristic-not-found; that was hiding the real API-choice bug. The
    /// classic `cublasSgemm_v2` path now serves all 8 TLOB GEMMs successfully
    /// for the tiny attention shapes (M=16, K∈{16,32}). Per
    /// `feedback_no_hiding.md` + `feedback_wire_everything_up.md`: wire it
    /// or delete it; never optional.
    pub(crate) gpu_tlob: GpuTlob,
    // ── Phase 2: Aux multi-stream ──
    /// Dedicated CUDA stream for IQN aux training (parallel with attn_stream after IQL).
    iqn_stream: Option<Arc<cudarc::driver::CudaStream>>,
    /// Dedicated CUDA stream for Attention aux training (parallel with iqn_stream after IQL).
    attn_stream: Option<Arc<cudarc::driver::CudaStream>>,
    /// cuBLAS workspace for IQN stream (32MB — matches PerStreamCublasHandles lt_workspace).
    iqn_workspace: Option<cudarc::driver::CudaSlice<u8>>,
    /// cuBLAS workspace for Attention stream (32MB — matches PerStreamCublasHandles lt_workspace).
    attn_workspace: Option<cudarc::driver::CudaSlice<u8>>,
    /// Event recorded on main stream after IQL completes; IQN/Attention streams wait on it.
    iql_done_event: Option<cudarc::driver::CudaEvent>,
    /// Event recorded on iqn_stream after IQN completes; main stream waits before Adam.
    iqn_done_event: Option<cudarc::driver::CudaEvent>,
    /// Event recorded on attn_stream after Attention completes; main stream waits before Adam.
    attn_done_event: Option<cudarc::driver::CudaEvent>,
    /// Ensemble extra heads (heads 1..K-1). Head 0 lives inside the CUDA Graph.
    /// Each element is an independent (DuelingWeightSet, BranchingWeightSet) pair
    /// that shares the same trunk but has perturbed value/advantage weights.
    /// Empty when ensemble_count <= 1.
    pub(crate) ensemble_extra_heads: Vec<(DuelingWeightSet, BranchingWeightSet)>,
    /// Ensemble diversity weight λ — scales the KL loss relative to C51 loss.
    /// Default 0.01. Zero when ensemble_count <= 1.
    pub(crate) ensemble_diversity_weight: f32,
    /// Compiled ensemble aggregate kernel (Q-value mean/variance across heads).
    /// None when ensemble_count <= 1.
    pub(crate) ensemble_aggregate_kernel: Option<cudarc::driver::CudaFunction>,
    /// Pre-allocated buffer: [K * B * num_atoms] f32 for per-head value logits.
    /// Used by the KL gradient kernel to assemble logits from all K heads.
    /// F32 because output layer GemmEx writes f32 (no f32 truncation overflow).
    /// None when ensemble_count <= 1.
    pub(crate) ensemble_logits_buf: Option<cudarc::driver::CudaSlice<f32>>,
    /// Pre-allocated buffer: [B * num_atoms] for mean Q-values across K heads.
    /// None when ensemble_count <= 1.
    pub(crate) ensemble_mean_q_buf: Option<cudarc::driver::CudaSlice<f32>>,
    /// Pre-allocated buffer: [B * num_atoms] for Q-value variance across K heads.
    /// None when ensemble_count <= 1.
    pub(crate) ensemble_var_q_buf: Option<cudarc::driver::CudaSlice<f32>>,
    /// Compiled KL gradient kernel for diversity gradient flow.
    pub(crate) ensemble_kl_grad_kernel: Option<cudarc::driver::CudaFunction>,
    /// Pre-allocated buffer: [B * num_atoms] for diversity gradient on head 0 logits.
    pub(crate) ensemble_d_logits_buf: Option<cudarc::driver::CudaSlice<f32>>,
    /// #32 Gradient Vaccine: pending validation batch for next training step.
    /// Set by the training loop, consumed by `run_full_step()`.
    pub(crate) pending_vaccine_batch: Option<crate::dqn::replay_buffer_type::GpuBatch>,
    /// Batch count for logging (wall-clock).
    phase_batch_count: usize,
    /// Child sub-graphs (captured independently, composed into parent).
    /// Children are raw CUgraph handles -- NOT individually instantiated.
    per_sample_child: Option<ChildGraph>,
    counters_child: Option<ChildGraph>,
    spectral_child: Option<ChildGraph>,
    forward_child: Option<ChildGraph>,
    ddqn_child: Option<ChildGraph>,
    aux_child: Option<ChildGraph>,
    post_aux_child: Option<ChildGraph>,
    adam_grad_child: Option<ChildGraph>,
    adam_child: Option<ChildGraph>,
    maintenance_child: Option<ChildGraph>,
    iql_modulate_child: Option<ChildGraph>,
    per_priority_child: Option<ChildGraph>,
    /// PER RNG step counter (pinned device-mapped, GPU-side increment).
    per_rng_step_dev_ptr: u64,
    /// Composed parent -- single launch replays all children sequentially.
    parent_graph: Option<ParentGraph>,
    /// Eval-only forward graph exec for deterministic Q-value replay.
    /// Separately instantiated from forward_child for eval isolation.
    eval_forward_exec: Option<cuda_sys::CUgraphExec>,
    /// Host-side tau for graph replay. Graph captures HtoD from this address.
    tau_host: f32,
    /// v8: PopArt reward normalization enabled (false by default).
    pub(crate) popart_enabled: bool,
    /// v8: Use robust PopArt (median/IQR) instead of Welford mean/var.
    pub(crate) popart_robust: bool,
    /// Cached median from sorted epoch rewards (for robust PopArt).
    pub(crate) cached_median: f32,
    /// Cached IQR from sorted epoch rewards (for robust PopArt).
    pub(crate) cached_iqr: f32,
    /// PopArt running variance from previous epoch (for tau reset detection).
    prev_popart_var: f32,
    /// v8: Hindsight relabeling fraction (0.0 = disabled).
    pub(crate) hindsight_fraction: f64,
    /// v8: Curriculum learning enabled (false by default).
    pub(crate) curriculum_enabled: bool,
}

impl Drop for FusedTrainingCtx {
    fn drop(&mut self) {
        // Sync stream before fields drop to prevent stale CUDA events.
        let _ = self.stream.context().check_err();
        // Destroy eval forward exec (not owned by ParentGraph or ChildGraph).
        if let Some(exec) = self.eval_forward_exec.take() {
            unsafe { cuda_sys::cuGraphExecDestroy(exec); }
        }
        // ParentGraph + ChildGraph Drop impls handle their own cleanup.
    }
}

impl FusedTrainingCtx {
    /// Create a new fused training context.
    ///
    /// Extracts online + target weight sets from the DQN's Candle VarMaps into
    /// `CudaSlice` buffers, compiles the 3 fused kernels + EMA kernel, and
    /// allocates all pre-allocated buffers for CUDA Graph capture.
    ///
    /// Only valid when `matches!(device, MlDevice::Cuda { .. }) && agent.is_using_branching()`.
    pub(crate) fn new(
        _device: &MlDevice,
        agent: &DQNAgentType,
        hyperparams: &DQNHyperparameters,
        batch_size: usize,
        stream: Arc<cudarc::driver::CudaStream>,
    ) -> Result<Self> {
        let dqn = agent.primary_dqn();

        dqn.branching_q_network.as_ref().ok_or_else(|| {
            anyhow::anyhow!("Fused CUDA training requires branching Q-network")
        })?;
        dqn.branching_target_network.as_ref().ok_or_else(|| {
            anyhow::anyhow!("Fused CUDA training requires branching target network")
        })?;

        // Use the caller-provided forked CudaStream. All GPU components (experience
        // collector, fused trainer, portfolio sim, monitoring) share this single
        // stream via Arc::clone — no more dual-stream event tracking conflicts.

        // Resolve gradient clip norm ONCE — single source of truth for all sub-trainers.
        // Was: main=10.0, IQN=1.0, IQL=1.0 → inconsistent 10× mismatch.
        let resolved_grad_norm = hyperparams.gradient_clip_norm.unwrap_or(10.0);

        // Build config from DQN network dimensions
        let (shared_h1, shared_h2, value_h, adv_h) = agent.network_dims();
        let config = GpuDqnTrainConfig {
            state_dim: ml_core::state_layout::STATE_DIM,
            shared_h1,
            shared_h2,
            value_h,
            adv_h,
            num_atoms: dqn.config.num_atoms,
            v_min: dqn.config.v_min,
            v_max: dqn.config.v_max,
            branch_0_size: dqn.config.num_actions,
            branch_1_size: dqn.config.num_order_types,
            branch_2_size: dqn.config.num_urgency_levels,
            branch_3_size: 3, // urgency — fixed at 3 (4-branch refactor)
            batch_size,
            // Hyperparams are now f32 native (post f64→f32 migration); pass-through without cast.
            // Adam lr/eps/decay still derive from f64 hyperparams → f32 at optimizer boundary.
            gamma: hyperparams.gamma,
            n_steps: hyperparams.n_steps,
            lr: hyperparams.learning_rate as f32,
            beta1: 0.9,
            beta2: 0.999,
            epsilon: hyperparams.adam_epsilon as f32,
            // SP4 Layer B: `weight_decay` field removed from GpuDqnTrainConfig
            // and DQNHyperparameters; per-group rates flow from ISV at runtime.
            max_grad_norm: resolved_grad_norm as f32,
            spectral_norm_sigma_max: hyperparams.spectral_norm_sigma_max,
            spectral_decoupling_lambda: hyperparams.spectral_decoupling_lambda,
            iqn_lambda: hyperparams.iqn_lambda,
            // Plan 4 Task 3 (E.3): IQN τ schedule is `FIXED_TAUS` (len 5),
            // pinned in lockstep with the kernel-side `IQN_NUM_QUANTILES`
            // macro. Embedding dim (64) and Huber κ (1.0) are kernel-fixed.
            iqn_num_quantiles: if hyperparams.iqn_lambda > 0.0 {
                crate::cuda_pipeline::gpu_iqn_head::FIXED_TAUS.len()
            } else { 0 },
            iqn_embedding_dim: 64,
            iqn_kappa: 1.0,
            entropy_coefficient: dqn.config.entropy_coefficient as f32,
            c51_warmup_epochs: hyperparams.c51_warmup_epochs,
            cql_alpha: hyperparams.cql_alpha as f32,
            curiosity_q_penalty_lambda: hyperparams.curiosity_q_penalty_lambda,
            asymmetric_dd_weight: hyperparams.asymmetric_dd_weight,
            ensemble_disagreement_penalty: hyperparams.ensemble_disagreement_penalty,
            stochastic_depth_prob: hyperparams.stochastic_depth_prob,
            pruning_epoch: hyperparams.pruning_epoch,
            pruning_fraction: hyperparams.pruning_fraction,
            causal_weight: f64::from(hyperparams.causal_weight),
            causal_intervention_interval: hyperparams.causal_intervention_interval,
            bottleneck_dim: hyperparams.bottleneck_dim,
            market_dim: 42, // Always 42 base market features — OFI features bypass bottleneck via portfolio_dim
            total_epochs: hyperparams.epochs,
            replay_seed_steps: hyperparams.replay_seed_steps,
            // MoE adaptive load-balance λ controller (2026-04-27).
            moe_lambda_floor: hyperparams.moe_lambda_floor.unwrap_or(0.01),
            moe_lambda_max_extra: hyperparams.moe_lambda_max_extra.unwrap_or(0.09),
            moe_entropy_target_frac: hyperparams.moe_entropy_target_frac.unwrap_or(0.7),
        };

        // Create weight set pointer views AFTER GpuDqnTrainer is constructed below.
        // Weight sets point directly into params_buf — zero copies.

        // Permanently disable cudarc event tracking BEFORE any CudaSlice allocation.
        // This ensures all buffers have read=None, write=None, preventing SyncOnDrop
        // from calling cuEventRecord which poisons error state during graph capture.
        let _ = stream.context().check_err();
        unsafe { stream.context().disable_event_tracking(); }

        // Compute param sizes BEFORE moving config into trainer.
        let param_sizes = crate::cuda_pipeline::gpu_dqn_trainer::compute_param_sizes(&config);

        // Create the fused trainer (compiles kernels, allocates buffers)
        let mut trainer = GpuDqnTrainer::new(stream.clone(), config)
            .map_err(|e| anyhow::anyhow!("GpuDqnTrainer init: {e}"))?;

        // Xavier-initialize both params_buf and target_params_buf BEFORE creating
        // weight set pointer views. Uses compute_param_sizes() for correct
        // bottleneck-aware dimensions (s1_input_dim, not state_dim).
        trainer.xavier_init_params_buf()
            .map_err(|e| anyhow::anyhow!("xavier_init_params_buf: {e}"))?;

        // Create weight set pointer views directly into params_buf — zero copies.
        // params_buf is Xavier-initialized by xavier_init_params_buf() above.
        let online_dueling = gpu_weights::DuelingWeightSet::from_flat_buffer(
            trainer.params_buf_ptr(), &param_sizes,
        );
        let online_branching = gpu_weights::BranchingWeightSet::from_flat_buffer(
            trainer.params_buf_ptr(), &param_sizes,
        );
        let target_dueling = gpu_weights::DuelingWeightSet::from_flat_buffer(
            trainer.target_params_buf_ptr(), &param_sizes,
        );
        let target_branching = gpu_weights::BranchingWeightSet::from_flat_buffer(
            trainer.target_params_buf_ptr(), &param_sizes,
        );

        // Initialize GPU HER if her_ratio > 0
        let gpu_her = if hyperparams.her_ratio > 0.0 {
            let her_config = GpuHerConfig {
                goal_dim: 1, // Single scalar goal (target return)
                her_ratio: f64::from(hyperparams.her_ratio),
                strategy: parse_her_strategy(&hyperparams.her_strategy),
                goal_threshold: 0.01,
                batch_size,
                state_dim: ml_core::state_layout::STATE_DIM,
            };
            match GpuHer::new(stream.clone(), her_config) {
                Ok(her) => {
                    info!(
                        her_ratio = hyperparams.her_ratio,
                        her_strategy = %hyperparams.her_strategy,
                        "GPU HER initialized: relabel kernel compiled"
                    );
                    Some(her)
                }
                Err(e) => {
                    tracing::warn!("GPU HER init failed, continuing without HER: {e}");
                    None
                }
            }
        } else {
            None
        };

        // Pre-compute HER source indices and reward ones on GPU.
        // Mapped pinned staging + DtoD copy — no HtoD per
        // `feedback_no_htod_htoh_only_mapped_pinned.md`.
        // Destination types remain `CudaSlice<{i32,f32}>` so the downstream
        // `relabel_batch_with_strategy` kernel-arg surface in gpu_her.rs is
        // unchanged (out of scope per task).
        let (her_source_indices_gpu, her_reward_ones_gpu) = if gpu_her.is_some() {
            let her_batch_size = (batch_size as f64 * f64::from(hyperparams.her_ratio)) as usize;
            let normal_count = batch_size.saturating_sub(her_batch_size);
            // Source indices: [normal_count, normal_count+1, ..., batch_size-1]
            let source_idx: Vec<i32> = (normal_count..batch_size).map(|i| i as i32).collect();
            let src_gpu = staging_upload_i32(&stream, her_batch_size, &source_idx, "HER_source_indices")?;
            // Reward ones: [1.0; her_batch_size]
            let ones: Vec<f32> = vec![1.0_f32; her_batch_size];
            let ones_gpu = staging_upload_f32(&stream, her_batch_size, &ones, "HER_reward_ones")?;
            (Some(src_gpu), Some(ones_gpu))
        } else {
            (None, None)
        };

        // Initialize dual IQL value trainers (mandatory — always active)
        let iql_config = GpuIqlConfig {
            expectile_tau: hyperparams.iql_expectile_tau,
            advantage_temperature: hyperparams.iql_advantage_temperature,
            value_hidden_dim: 128,
            state_dim: ml_core::state_layout::STATE_DIM,
            batch_size,
            lr: hyperparams.learning_rate as f32,
            max_grad_norm: resolved_grad_norm as f32,
            num_atoms: dqn.config.num_atoms,
            total_actions: dqn.config.num_actions + dqn.config.num_order_types
                + dqn.config.num_urgency_levels + 3,
            branch_sizes: [
                dqn.config.num_actions,
                dqn.config.num_order_types,
                dqn.config.num_urgency_levels,
                3,
            ],
            gamma: hyperparams.gamma as f32,
            ..GpuIqlConfig::default()
        };
        let shared_cublas_for_iql = Arc::clone(trainer.shared_cublas());
        let gpu_iql = GpuIqlTrainer::new(Arc::clone(&shared_cublas_for_iql), iql_config.clone())
            .map_err(|e| anyhow::anyhow!("GPU IQL (high-tau) init: {e}"))?;

        let iql_low_config = GpuIqlConfig {
            expectile_tau: 1.0 - hyperparams.iql_expectile_tau,
            ..iql_config
        };
        let gpu_iql_low = GpuIqlTrainer::new(shared_cublas_for_iql, iql_low_config)
            .map_err(|e| anyhow::anyhow!("GPU IQL (low-tau) init: {e}"))?;

        // Initialize GPU IQN dual-head when iqn_lambda > 0 (CVaR risk sizing).
        //
        // Plan 4 Task 3 (E.3): `num_quantiles` is pinned to `FIXED_TAUS.len()`
        // (= 5). The kernel-side `IQN_NUM_QUANTILES` macro sizes register
        // arrays in lockstep with FIXED_TAUS; the embedding dim (64) and
        // Huber threshold κ (1.0) are kernel-fixed in the same lockstep.
        // `iqn_lambda > 0` enables the dual-head; `iqn_lambda = 0.0` disables
        // the IQN aux contribution to the loss without removing the head.
        let gpu_iqn = if hyperparams.iqn_lambda > 0.0 {
            let iqn_config = GpuIqnConfig {
                hidden_dim: shared_h2,
                embed_dim: 64,
                num_quantiles: crate::cuda_pipeline::gpu_iqn_head::FIXED_TAUS.len(),
                kappa: 1.0,
                state_dim: ml_core::state_layout::STATE_DIM,
                shared_h1,
                branch_0_size: dqn.config.num_actions,
                branch_1_size: dqn.config.num_order_types,
                branch_2_size: dqn.config.num_urgency_levels,
                branch_3_size: 3, // urgency — fixed at 3 (4-branch refactor)
                batch_size,
                lr: hyperparams.learning_rate as f32,
                max_grad_norm: resolved_grad_norm as f32,
                gamma: hyperparams.gamma as f32,
                ..GpuIqnConfig::default()
            };
            match GpuIqnHead::new(Arc::clone(trainer.shared_cublas()), iqn_config) {
                Ok(iqn) => {
                    info!(
                        iqn_lambda = hyperparams.iqn_lambda,
                        num_quantiles = hyperparams.num_quantiles,
                        embed_dim = 64,
                        "GPU IQN dual-head initialized: 7 kernels compiled"
                    );
                    Some(iqn)
                }
                Err(e) => {
                    tracing::warn!("GPU IQN init failed, continuing without IQN: {e}");
                    None
                }
            }
        } else {
            None
        };

        // SP2: populate fused NaN-check (ptr, len) tables via host-side write
        // through the trainer's mapped-pinned `nan_check_buf_ptrs` /
        // `nan_check_buf_lens` (no HtoD copy per
        // `feedback_no_htod_htoh_only_mapped_pinned`). Slot pointers are
        // stable across the trainer's lifetime — populate ONCE here, after
        // `gpu_iqn` has been constructed so its slot 27/28 accessors resolve
        // correctly. When IQN is inactive (`iqn_lambda == 0.0` →
        // `gpu_iqn = None`), slots 27/28 receive null entries which the fused
        // kernel skips via its null-pointer guard.
        let iqn_q_for_meta = gpu_iqn
            .as_ref()
            .map(|_| crate::cuda_pipeline::gpu_iqn_head::FIXED_TAUS.len());
        trainer.populate_nan_check_meta(
            batch_size,
            shared_h2,
            gpu_iqn.as_ref().map(|h| h.d_h_s2_raw_ptr()),
            gpu_iqn.as_ref().map(|h| h.d_branch_logits_buf_ptr()),
            iqn_q_for_meta,
            // SP3 Mech 5: IQN Adam m/v ptr+len for slots 39/43. None when
            // IQN inactive — the kernel's null-pointer guard skips those
            // slots, symmetric with slots 27/28.
            gpu_iqn.as_ref().map(|h| h.adam_m_ptr()),
            gpu_iqn.as_ref().map(|h| h.adam_m_len()),
            gpu_iqn.as_ref().map(|h| h.adam_v_ptr()),
            gpu_iqn.as_ref().map(|h| h.adam_v_len()),
            // SP3 close-out v2: IQN online_params ptr+len for slot 48 — the
            // diagnostic blind spot that hid the F1 step-3540 NaN. Same
            // None=inactive convention.
            gpu_iqn.as_ref().map(|h| h.online_params_ptr()),
            gpu_iqn.as_ref().map(|h| h.online_params_len()),
        ).map_err(|e| anyhow::anyhow!("populate_nan_check_meta: {e}"))?;

        // GPU attention — always active. 4-head self-attention on h_s2.
        let attn_config = GpuAttentionConfig {
            state_dim: shared_h2,
            num_heads: 4,
            batch_size,
        };
        let gpu_attention = match GpuAttention::new(Arc::clone(trainer.shared_cublas()), attn_config) {
            Ok(attn) => {
                info!(
                    hidden_dim = shared_h2,
                    num_heads = 4,
                    batch_size,
                    "GPU attention initialized: 4-head feature attention on h_s2"
                );
                Some(attn)
            }
            Err(e) => {
                tracing::warn!("GPU attention init failed, continuing without: {e}");
                None
            }
        };

        // D.8 TLOB: initialize attention module over OFI features.
        // No graceful-degrade: TLOB is on the production training path. If init
        // fails the training cannot proceed correctly — surface the error to
        // the caller rather than silently dropping the OFI→TLOB attention head
        // (which previously hid the cuBLAS-Lt heuristic failure for tiny
        // attention shapes; now resolved via classic cublasSgemm_v2). Per
        // `feedback_no_hiding.md`: wire up or surface, never suppress.
        let gpu_tlob = GpuTlob::new(Arc::clone(trainer.shared_cublas()), batch_size)
            .map_err(|e| anyhow::anyhow!("GpuTlob init: {e}"))?;

        // Phase 2: Allocate aux multi-stream infrastructure (streams, workspaces, events).
        // Only allocate when IQN or Attention is present — parallel dispatch needs at least
        // one aux trainer to be worthwhile.
        let (iqn_stream, attn_stream, iqn_workspace, attn_workspace, iql_done_event, iqn_done_event, attn_done_event) =
            if gpu_iqn.is_some() || gpu_attention.is_some() {
                let iqn_s = stream.context().new_stream()
                    .map_err(|e| anyhow::anyhow!("iqn_stream: {e}"))?;
                let attn_s = stream.context().new_stream()
                    .map_err(|e| anyhow::anyhow!("attn_stream: {e}"))?;
                let iqn_ws = stream.alloc_zeros::<u8>(32 * 1024 * 1024)
                    .map_err(|e| anyhow::anyhow!("iqn workspace: {e}"))?;
                let attn_ws = stream.alloc_zeros::<u8>(32 * 1024 * 1024)
                    .map_err(|e| anyhow::anyhow!("attn workspace: {e}"))?;
                let iql_ev = stream.context().new_event(None)
                    .map_err(|e| anyhow::anyhow!("iql_done_event: {e}"))?;
                let iqn_ev = stream.context().new_event(None)
                    .map_err(|e| anyhow::anyhow!("iqn_done_event: {e}"))?;
                let attn_ev = stream.context().new_event(None)
                    .map_err(|e| anyhow::anyhow!("attn_done_event: {e}"))?;

                // Pre-register iqn_stream and attn_stream in the per-stream
                // cublasLt registry (Option C). Must happen BEFORE CUDA Graph
                // capture — the first-access cold path inside `lt_for` calls
                // `cudaMalloc`, which fails during capture.
                let shared = trainer.shared_cublas();
                shared.pre_register_stream(&iqn_s)
                    .map_err(|e| anyhow::anyhow!("pre_register iqn_stream: {e:?}"))?;
                shared.pre_register_stream(&attn_s)
                    .map_err(|e| anyhow::anyhow!("pre_register attn_stream: {e:?}"))?;

                info!("Phase 2 aux multi-stream: 2 streams + 64MB workspace + 3 events allocated");
                (Some(iqn_s), Some(attn_s), Some(iqn_ws), Some(attn_ws), Some(iql_ev), Some(iqn_ev), Some(attn_ev))
            } else {
                (None, None, None, None, None, None, None)
            };

        // Initialize ensemble extra heads (heads 1..K-1) when ensemble_count > 1.
        // Head 0 is head_0 = online_dueling + online_branching (already inside CUDA Graph).
        // Heads 1..K-1 run OUTSIDE the CUDA Graph and share the trunk activations (h_s2).
        let k = hyperparams.ensemble_count.max(1);
        let (
            ensemble_extra_heads,
            ensemble_aggregate_kernel,
            ensemble_logits_buf,
            ensemble_mean_q_buf,
            ensemble_var_q_buf,
            ensemble_kl_grad_kernel,
            ensemble_d_logits_buf,
            ens_backings_vec,
        ) = if k > 1 {
            use crate::cuda_pipeline::gpu_dqn_trainer::compile_ensemble_kernels;

            let agg_kernel =
                compile_ensemble_kernels(&stream, ml_core::state_layout::STATE_DIM)
                    .map_err(|e| anyhow::anyhow!("Ensemble kernels compile: {e}"))?;

            // Allocate K-1 extra head weight sets via DtoD clone of head 0 + small noise.
            // DtoD: head_k ← online_dueling + online_branching weights copied on GPU.
            // Each head has its own backing storage + pointer view.
            let mut extra_heads: Vec<(DuelingWeightSet, BranchingWeightSet)> =
                Vec::with_capacity(k - 1);
            let mut ens_backings: Vec<(gpu_weights::DuelingWeightBacking, gpu_weights::BranchingWeightBacking)> =
                Vec::with_capacity(k - 1);
            for head_idx in 1..k {
                let (d_backing, head_dueling) = clone_dueling_weights(&online_dueling, &stream, head_idx)
                    .map_err(|e| anyhow::anyhow!("Clone ensemble dueling head {head_idx}: {e}"))?;
                let (b_backing, head_branching) = clone_branching_weights(&online_branching, &stream, head_idx)
                    .map_err(|e| anyhow::anyhow!("Clone ensemble branching head {head_idx}: {e}"))?;
                extra_heads.push((head_dueling, head_branching));
                ens_backings.push((d_backing, b_backing));
            }

            // Pre-allocate ensemble buffers.
            // ensemble_logits_buf: [K * B * num_atoms] — all K heads' value logits
            let na = dqn.config.num_atoms;
            // total_actions removed — ensemble buffers use num_atoms, not total_actions.
            let logits_buf = stream.alloc_zeros::<f32>(k * batch_size * na)
                .map_err(|e| anyhow::anyhow!("Alloc ensemble_logits_buf f32: {e}"))?;
            // mean_q_buf and var_q_buf: [B * num_atoms] — aggregate kernel outputs
            // per-atom mean/variance across K heads (NOT per-action).
            let mean_q_buf = stream.alloc_zeros::<f32>(batch_size * na)
                .map_err(|e| anyhow::anyhow!("Alloc ensemble_mean_q_buf: {e}"))?;
            let var_q_buf = stream.alloc_zeros::<f32>(batch_size * na)
                .map_err(|e| anyhow::anyhow!("Alloc ensemble_var_q_buf: {e}"))?;

            info!(
                ensemble_count = k,
                extra_heads = k - 1,
                "Ensemble multi-head initialized: {k} heads, KL diversity weight={diversity_weight}",
                diversity_weight = hyperparams.ensemble_diversity_weight,
            );

            // Load KL gradient kernel from precompiled cubin (ZERO runtime nvcc)
            let kl_grad_kernel = {
                static ENSEMBLE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/ensemble_kernels.cubin"));
                let module = stream.context().load_cubin(ENSEMBLE_CUBIN.to_vec())
                    .map_err(|e| anyhow::anyhow!("ensemble_kl_grad cubin load: {e}"))?;
                module.load_function("ensemble_kl_gradient_kernel")
                    .map_err(|e| anyhow::anyhow!("ensemble_kl_gradient_kernel load: {e}"))?
            };

            // d_logits buffer for diversity gradient [B * num_atoms]
            let d_logits_buf = stream.alloc_zeros::<f32>(batch_size * dqn.config.num_atoms)
                .map_err(|e| anyhow::anyhow!("alloc ensemble_d_logits: {e}"))?;

            (
                extra_heads,
                Some(agg_kernel),
                Some(logits_buf),
                Some(mean_q_buf),
                Some(var_q_buf),
                Some(kl_grad_kernel),
                Some(d_logits_buf),
                ens_backings,
            )
        } else {
            (Vec::new(), None, None, None, None, None, None, Vec::new())
        };

        // Wire IQL buffer pointers to the DQN trainer for C51 kernel launches.
        trainer.set_per_sample_support_ptr(gpu_iql.per_sample_support_ptr());
        trainer.set_branch_scales_ptr(gpu_iql.branch_scales_ptr());

        info!(
            batch_size,
            her_enabled = gpu_her.is_some(),
            iql_enabled = true,
            iqn_enabled = gpu_iqn.is_some(),
            attention_enabled = gpu_attention.is_some(),
            ensemble_heads = k,
            "Fused CUDA training initialized: 4 kernels + EMA compiled, \
             ~291K params, CUDA Graph will capture on first step"
        );

        Ok(Self {
            trainer,
            online_dueling,
            online_branching,
            target_dueling,
            target_branching,
            _weight_backings: None, // zero-copy: params_buf owns the memory
            _ensemble_backings: ens_backings_vec,
            stream,
            batch_size,
            steps_since_varmap_sync: 0,
            best_params_snapshot: None,
            restore_event: None,
            last_c51_raw_norm: 0.0,
            last_combined_norm: 0.0,
            last_loss: 0.0,
            loss_accum: 0.0,
            grad_accum: 0.0,
            accum_steps: 0,
            gpu_her,
            her_source_indices_gpu,
            her_reward_ones_gpu,
            her_inplace_kernel: None, // kernel lives in trainer.her_inplace_kernel
            gpu_iql,
            gpu_iql_low,
            gpu_iqn,
            cvar_scales_buf: None,
            gpu_attention,
            gpu_tlob,
            iqn_stream,
            attn_stream,
            iqn_workspace,
            attn_workspace,
            iql_done_event,
            iqn_done_event,
            attn_done_event,
            ensemble_extra_heads,
            ensemble_diversity_weight: hyperparams.ensemble_diversity_weight as f32,
            ensemble_aggregate_kernel,
            ensemble_logits_buf,
            ensemble_mean_q_buf,
            ensemble_var_q_buf,
            ensemble_kl_grad_kernel,
            ensemble_d_logits_buf,
            pending_vaccine_batch: None,
            phase_batch_count: 0,
            per_sample_child: None,
            counters_child: None,
            spectral_child: None,
            forward_child: None,
            ddqn_child: None,
            aux_child: None,
            post_aux_child: None,
            adam_grad_child: None,
            adam_child: None,
            maintenance_child: None,
            iql_modulate_child: None,
            per_priority_child: None,
            per_rng_step_dev_ptr: 0,
            parent_graph: None,
            eval_forward_exec: None,
            tau_host: 0.0,
            popart_enabled: hyperparams.popart_enabled,
            popart_robust: hyperparams.popart_robust,
            cached_median: 0.0,
            cached_iqr: 0.0,
            prev_popart_var: 0.0,
            hindsight_fraction: f64::from(hyperparams.hindsight_fraction),
            curriculum_enabled: hyperparams.curriculum_enabled,
        })
    }

    /// Batch size this context was created for.
    /// #20 Delegate to GPU trainer's pruning mask computation.
    pub(crate) fn compute_pruning_mask(&mut self, epoch: usize) -> Result<()> {
        self.trainer.compute_pruning_mask(epoch)
            .map_err(|e| anyhow::anyhow!("pruning: {e}"))
    }

    pub(crate) fn batch_size(&self) -> usize {
        self.batch_size
    }

    /// Reset per-fold state. Keeps GPU infrastructure + compiled kernels.
    /// Invalidates child graphs and parent (fold data range changed).
    pub(crate) fn reset_for_fold(&mut self) -> anyhow::Result<()> {
        self.steps_since_varmap_sync = 0;
        self.last_combined_norm = 0.0;
        self.last_loss = 0.0;
        // Destroy parent graph and children -- re-capture on next step.
        self.parent_graph = None;
        self.per_sample_child = None;
        self.counters_child = None;
        self.spectral_child = None;
        self.forward_child = None;
        self.ddqn_child = None;
        self.aux_child = None;
        self.post_aux_child = None;
        self.adam_grad_child = None;
        self.adam_child = None;
        self.maintenance_child = None;
        self.iql_modulate_child = None;
        self.per_priority_child = None;
        // Phase 2 aux streams, workspaces, and events persist across folds.
        // They are reusable GPU resources — no reset or reallocation needed.
        if let Some(exec) = self.eval_forward_exec.take() {
            unsafe { cuda_sys::cuGraphExecDestroy(exec); }
        }
        // Shrink-and-perturb at fold boundary -- reduce overfit to previous fold.
        // alpha=0.8: keep 80% of learned representations, add 20% random noise.
        // Without this, the greedy policy from fold 1 produces val_Sharpe=-33
        // on fold 2 data and takes many epochs to recover.
        let sp_alpha = 0.8_f32;
        let sp_sigma = 0.01_f32;
        if let Err(e) = self.trainer.shrink_and_perturb(sp_alpha, sp_sigma) {
            tracing::warn!("Fold-boundary shrink-and-perturb failed (non-fatal): {e}");
        } else {
            tracing::info!(alpha = sp_alpha, sigma = sp_sigma, "Fold-boundary shrink-and-perturb applied");
        }
        // Hard-copy the shrink-and-perturb'd online weights into target params.
        // Without this, target_params_buf retains end-of-previous-fold weights
        // while online was just modified — the Bellman target would use stale
        // weights against perturbed online predictions, producing an outsized
        // TD error gap in the first fold-N+1 steps. Polyak averaging alone is
        // too slow (tau=0.005) to close that gap before the oversized gradients
        // compound through Adam. This pairs with reset_adam_state below.
        self.trainer.sync_target_from_online()?;
        // Reset Adam optimizer state — each fold starts with fresh momentum.
        // Stale momentum from a previous fold causes weight explosion → Q-value -1e30.
        self.trainer.reset_adam_state()?;

        // The IQN head carries its OWN Adam state (m_buf/v_buf/adam_step on
        // `GpuIqnHead`) — separate from `GpuDqnTrainer::reset_adam_state`,
        // which only zeroes the legacy `iqn_trunk_*` buffers. Without this
        // reset the IQN optimizer enters fold N+1 with fold N's momentum,
        // producing oversized Adam steps in the first epochs whose effect
        // compounds through the IQN backward pass into runaway gradients.
        if let Some(iqn) = self.gpu_iqn.as_mut() {
            // ┌─────────────────────────────────────────────────────────────┐
            // │ DO NOT DELETE — fold-1 Q-drift regression trip-wire (#84)   │
            // └─────────────────────────────────────────────────────────────┘
            //
            // Hard-sync IQN target ← online at the same fold boundary as the
            // main DQN's `sync_target_from_online` above. Without this, the
            // IQN target buffer retains end-of-prior-fold weights while the
            // online side is otherwise advancing through fold N+1's first
            // gradient steps; Polyak EMA at τ=0.005 closes ~0.5%/step, far
            // too slow before the inflated TD-error compounds geometrically
            // (observed: Q ≈ 0.8 → 2.2 → 4.0 → 10.7 → fp32 overflow past
            // atom support inside fold 1). Pairs with the IQN Adam reset
            // immediately below — both are fold-boundary contract consumers
            // that must migrate together (see `feedback_no_partial_refactor`).
            //
            // If you remove the `iqn.sync_target_from_online()` call below,
            // fold-1 Q-drift returns and the model becomes unsafe for live
            // trading (Kelly cap + runaway Q → oversized positions). The
            // root-cause fix landed in commit `7c19b5903` (issue #84). A
            // GPU runtime regression test exercises the contract end-to-
            // end — see
            // `cuda_pipeline::gpu_iqn_head::tests::iqn_sync_target_from_online_makes_target_equal_online`.
            // It runs on the L40S smoke validation pool (carries
            // `#[ignore = "gpu"]`); a stub body or a sync against the
            // wrong buffer/stream fails the bit-for-bit equality assertion.
            // Don't silence the test — restore the call.
            if let Err(e) = iqn.sync_target_from_online() {
                tracing::warn!("Fold-boundary IQN target sync failed (non-fatal): {e}");
            } else {
                tracing::info!("Fold-boundary IQN target hard-synced from online");
            }
            if let Err(e) = iqn.reset_adam_state() {
                tracing::warn!("Fold-boundary IQN Adam reset failed (non-fatal): {e}");
            } else {
                tracing::info!("Fold-boundary IQN Adam state reset");
            }
        }

        // TLOB and IQL each carry their own Adam state separate from the
        // main DQN's `m_buf` / `v_buf` (and from the IQN head's). The
        // fold-boundary contract that resets the main DQN Adam must also
        // reset every auxiliary optimizer, otherwise fold N's momentum
        // drives oversized first-epoch updates in fold N+1's expectile +
        // attention backward paths and feeds into the trunk gradient.
        // Per `feedback_no_partial_refactor.md` — same shared contract,
        // every consumer migrates together.
        if let Err(e) = self.gpu_tlob.reset_adam_state() {
            tracing::warn!("Fold-boundary TLOB Adam reset failed (non-fatal): {e}");
        } else {
            tracing::info!("Fold-boundary TLOB Adam state reset");
        }
        if let Err(e) = self.gpu_iql.reset_adam_state() {
            tracing::warn!("Fold-boundary IQL Adam reset failed (non-fatal): {e}");
        } else {
            tracing::info!("Fold-boundary IQL Adam state reset");
        }
        if let Err(e) = self.gpu_iql_low.reset_adam_state() {
            tracing::warn!("Fold-boundary IQL-low Adam reset failed (non-fatal): {e}");
        } else {
            tracing::info!("Fold-boundary IQL-low Adam state reset");
        }

        // SP3 Mech 4: 4-head attention Adam reset (fold-boundary).
        // GpuAttention owns its own (attn_m, attn_v, attn_adam_step) — the
        // SP3 Mech 4 audit (Task B5) found this optimizer was the lone
        // omission from the IQN/TLOB/IQL/main-DQN reset chain. The module
        // initializes with `Option<GpuAttention>` purely so init failure
        // is non-fatal; when present (which is the production default per
        // the constructor's "always active" comment) it trains every step
        // alongside the main DQN, so its momentum must zero at fold
        // boundary too. Same partial-refactor contract as IQN/TLOB/IQL —
        // see `feedback_no_partial_refactor.md`.
        if let Some(ref mut attn) = self.gpu_attention {
            if let Err(e) = attn.reset_adam_state() {
                tracing::warn!("Fold-boundary attention Adam reset failed (non-fatal): {e}");
            } else {
                tracing::info!("Fold-boundary attention Adam state reset");
            }
        }

        // SP3 Mech 8 (slow_ema fold-boundary reset, 2026-04-29) was
        // tested here and reverted 2026-04-30 after smoke-test-rxhjh
        // showed F1 NaN @ 2300 — WORSE than the 3720-step ceiling Mech 6
        // alone produced. Persisting `grad_norm_slow_ema` across folds
        // (anchored at F0's smaller grad scale) provides unintentional
        // F1 protection by tightening Mech 6's upper_bound during F1
        // ramp-up. Resetting it loosened the bound and accelerated Adam
        // saturation. SP3 Mech 9 (post-Adam ISV-driven weight clamp in
        // `dqn_adam_update_kernel`) addresses cross-fold Adam pathology
        // by bounding parameter magnitudes directly rather than tuning
        // the clip anchor.
        //
        // FoldReset state (eval_q_mean_ema, eval_q_std_ema, isv_v_range_slots,
        // isv_learning_health, isv_sharpe_ema, isv_q_means) is now handled by
        // StateResetRegistry-driven dispatch in DQNTrainer::reset_for_fold
        // (trainer/mod.rs). reset_eval_v_range_state is called via the
        // FusedTrainingCtx::reset_eval_v_range_state wrapper from there.
        // Clear replay buffer — stale experiences from previous fold have wrong
        // reward distributions (different data window, different policy).
        // NOTE: the replay buffer clear is handled by DQNTrainer::reset_for_fold()
        // (trainer/mod.rs) which calls self.clear_replay_buffer().await BEFORE
        // invoking fused.reset_for_fold(). GpuDqnTrainer does not own the replay
        // buffer, so we cannot clear it here directly.

        // PopArt fold-boundary policy:
        //   prev_popart_var: reset (it's used for tau-change detection — a fresh
        //     fold counts as a "change point" anyway, no signal lost).
        //   cached_median / cached_iqr: CARRY FORWARD from the previous fold.
        //     Resetting them to 0 forced the first epoch of fold N+1 to fall
        //     through to the Welford GPU path (`if cached_iqr > 0.0` check),
        //     and the Welford running stats from fold N — combined with fold
        //     N+1's slightly different reward distribution post-S&P + adversarial
        //     regime — produced reward normalisations far outside the C51
        //     atom-support [-50, 50]. This drove the categorical loss to 10⁵×
        //     readings (cosmetic only, val Sharpe stayed healthy) and on rare
        //     timing-sensitive paths overflowed to NaN.
        //     The reward distribution between adjacent walk-forward folds is
        //     similar — same instrument, same regime mix — so fold N's
        //     final-epoch median/IQR is a sensible bridge for fold N+1's
        //     first epoch. By the end of fold N+1's epoch 1, fresh stats are
        //     re-computed and replace the carried values via
        //     `set_robust_popart_stats`.
        if self.popart_enabled {
            self.prev_popart_var = 0.0;
        }

        // Reset NaN diagnostic flags so previous fold's NaN events don't taint
        // the next fold's diagnostic. Per-step checks inside the captured graph
        // (`submit_post_aux_ops`) accumulate flags across steps within a fold so
        // the FIRST buffer to go NaN remains visible at host-readback time.
        self.trainer.reset_nan_flags()
            .map_err(|e| anyhow::anyhow!("nan_flags reset at fold boundary: {e}"))?;
        Ok(())
    }

    /// Reset the adaptive C51 v_range EMAs and the pinned range buffer.
    ///
    /// Called from DQNTrainer::reset_named_state for the "eval_q_mean_ema",
    /// "eval_q_std_ema", and "isv_v_range_slots" FoldReset entries.
    /// These three states form a single logical unit — the EMA state that
    /// drives the ISV v-range slots. Separating them in the registry allows
    /// future tasks to migrate them independently; for now one call resets all.
    pub(crate) fn reset_eval_v_range_state(&mut self) {
        self.trainer.reset_eval_v_range_state();
    }

    /// Reset the IQN readiness gauge (`iqn_loss_initial`, `iqn_loss_ema`,
    /// `iqn_readiness` scalar + pinned slot).
    ///
    /// Called from DQNTrainer::reset_named_state for the
    /// "isv_iqn_loss_initial" / "isv_iqn_loss_ema" / "isv_iqn_readiness"
    /// FoldReset entries. The three names map to the same logical reset
    /// (the streaming-improvement gauge's three coupled scalars).
    pub(crate) fn reset_iqn_readiness_state(&mut self) {
        self.trainer.reset_iqn_readiness_state();
    }

    /// Steps since last flat buffer sync.
    pub(crate) fn steps_since_varmap_sync(&self) -> usize {
        self.steps_since_varmap_sync
    }

    /// Per-component gradient norm diagnostics for epoch-level Prometheus reporting.
    /// Returns (c51_raw_norm, combined_norm).
    pub(crate) fn grad_norm_diagnostics(&self) -> (f32, f32) {
        (self.last_c51_raw_norm, self.last_combined_norm)
    }

    /// Read and reset epoch-level loss/grad_norm averages (accumulated every step).
    pub(crate) fn read_and_reset_accumulators(&mut self) -> (f64, f64) {
        if self.accum_steps == 0 { return (0.0, 0.0); }
        let avg_loss = self.loss_accum / self.accum_steps as f64;
        let avg_grad = self.grad_accum / self.accum_steps as f64;
        self.loss_accum = 0.0;
        self.grad_accum = 0.0;
        self.accum_steps = 0;
        (avg_loss, avg_grad)
    }

    /// Last readback scalars from pinned device-mapped memory (no GPU sync).
    pub(crate) fn last_readback_scalars(&self) -> crate::cuda_pipeline::gpu_dqn_trainer::FusedTrainScalars {
        crate::cuda_pipeline::gpu_dqn_trainer::FusedTrainScalars {
            total_loss: self.last_loss,
            grad_norm: self.last_combined_norm,
        }
    }

    /// Plan C T11 follow-up — F1 ep2 explosion diagnostic (2026-04-29).
    /// Loss decomposition snapshot for the per-step FOLD_EXPLOSION_DIAG
    /// block. Returns `(c51_loss, mse_loss, iqn_loss)` from already-pinned
    /// device-mapped scalars (no GPU sync, no DtoH).
    ///
    /// * `c51_loss` ← `total_loss_pinned` (the C51 KL loss; same value the
    ///   blended-loss formula treats as the C51 component).
    /// * `mse_loss` ← `mse_loss_pinned` (the MSE warmup loss).
    /// * `iqn_loss` ← `gpu_iqn.read_total_loss()` if present; `NaN` otherwise.
    ///
    /// CQL and ensemble losses are not pinned-readback — they live in
    /// transient device buffers consumed inside the training graph and are
    /// not exposed for monitoring. Their absence from the diagnostic is
    /// acknowledged: the goal is to identify which of {c51, iqn, mse} is
    /// the runaway driver; if none of those three explode and the total
    /// still does, it implicates the unexposed CQL/ens path.
    /// Diagnostic-only; remove with the FOLD_EXPLOSION_DIAG block once root
    /// cause is identified.
    pub(crate) fn explosion_diag_loss_components(&self) -> (f32, f32, f32) {
        let c51_loss = self.trainer.c51_loss_pinned_value();
        let mse_loss = self.trainer.mse_loss_pinned_value();
        let iqn_loss = self.gpu_iqn
            .as_ref()
            .and_then(|iqn| iqn.read_total_loss().ok())
            .unwrap_or(f32::NAN);
        (c51_loss, mse_loss, iqn_loss)
    }

    /// Plan C T11 follow-up — F1 ep2 explosion diagnostic (2026-04-29).
    /// Atom positions for sample 0, directions 0 (Short) and 2 (Long) from
    /// the per-sample C51 support tile [B, 4, 3] stride-12. Layout per
    /// `gpu_iql_trainer::per_sample_support_ptr` doc: each per-sample
    /// per-branch triple is `(v_min, v_max, delta_z)`.
    ///
    /// Returns `[v_min_d0, v_max_d0, delta_z_d0, v_min_d2, v_max_d2, delta_z_d2]`.
    /// On any failure the array is filled with NaN — diagnostic non-fatal.
    ///
    /// Cold-path DtoH: a single 12-element f32 transfer per fired step,
    /// gated by the explosion-trigger predicate (so it runs at most a
    /// handful of times across an entire fold's ~3700 steps). Per
    /// `feedback_no_htod_htoh_only_mapped_pinned.md` the canonical zero-copy
    /// path applies on the hot path; this is a cold-path one-shot read of
    /// 48 bytes from a regular (non-pinned) `CudaSlice<f32>` and uses the
    /// existing `dtoh_f32` helper. Diagnostic-only; remove with the
    /// FOLD_EXPLOSION_DIAG block once root cause is identified.
    pub(crate) fn explosion_diag_atom_range(&self) -> [f32; 6] {
        let mut buf = [f32::NAN; 6];
        let support_ptr = self.gpu_iql.per_sample_support_ptr();
        if support_ptr == 0 {
            return buf;
        }
        // Read first 12 elements (sample 0, all 4 dirs × 3 floats), then
        // pick the dir-0 and dir-2 triples. Reading 12 instead of 6 keeps
        // the read contiguous from the buffer base — simpler than a
        // mid-buffer offset and the extra 24 bytes are inconsequential
        // on a cold-path one-shot.
        let mut sample0: [f32; 12] = [0.0; 12];
        // Build a view of the first 12 elements via the underlying CudaSlice.
        // Use raw cuMemcpyDtoH_v2 so we don't depend on cudarc's slice API
        // (mirrors `dtoh_f32` but we only have the raw pointer here).
        #[allow(unsafe_code)]
        unsafe {
            cudarc::driver::sys::cuStreamSynchronize(
                self.trainer.stream_for_diag().cu_stream(),
            );
            let rc = cudarc::driver::sys::cuMemcpyDtoH_v2(
                sample0.as_mut_ptr().cast(),
                support_ptr,
                12 * std::mem::size_of::<f32>(),
            );
            if rc != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
                return buf;
            }
        }
        // Layout: [d0_vmin, d0_vmax, d0_dz, d1_..., d2_vmin, d2_vmax, d2_dz, d3_...]
        buf[0] = sample0[0];   // dir 0 v_min
        buf[1] = sample0[1];   // dir 0 v_max
        buf[2] = sample0[2];   // dir 0 delta_z
        buf[3] = sample0[6];   // dir 2 v_min
        buf[4] = sample0[7];   // dir 2 v_max
        buf[5] = sample0[8];   // dir 2 delta_z
        buf
    }

    /// Plan C T11 follow-up — F1 ep2 explosion diagnostic (2026-04-29).
    /// Q-value range snapshot via the existing GPU reduction kernel.
    /// Returns `(q_min, q_mean, q_max)`. On reduction failure returns
    /// `(NaN, NaN, NaN)` — diagnostic non-fatal. Cold-path: launches the
    /// reduce + DtoH the 28-byte stats struct, gated by the explosion
    /// trigger so it runs at most a handful of times per fold.
    /// Diagnostic-only; remove with the FOLD_EXPLOSION_DIAG block once
    /// root cause is identified.
    pub(crate) fn explosion_diag_q_range(&mut self) -> (f32, f32, f32) {
        match self.trainer.reduce_current_q_stats() {
            Ok(s) => (s.q_min, s.q_mean, s.q_max),
            Err(_) => (f32::NAN, f32::NAN, f32::NAN),
        }
    }

    /// G12 predictive_coding loss readback (sync DtoH; epoch-boundary only).
    /// Used by HEALTH_DIAG `gems [g12_predictive=X]`. G6 + G10 were measured
    /// sub-noise via V7 methodology and their readbacks deleted.
    pub(crate) fn read_gem_g12_loss(&self) -> f32 {
        self.trainer.predictive_loss_value()
    }

    /// Task 0.4 — per-action-branch grad-norm ratio (magnitude / direction)
    /// for H4 detection. Reads grad_buf to a pinned host slot and computes
    /// L2 norm per branch slice. Returns 0.0 on failure (non-fatal — diag only).
    pub(crate) fn grad_ratio_mag_dir(&mut self) -> f32 {
        self.trainer.grad_ratio_mag_dir().unwrap_or(0.0)
    }

    /// Task 2.0 confirmation — absolute per-branch grad L2 norms, returning
    /// `(dir, mag)`. Complements `grad_ratio_mag_dir` — the ratio collapses
    /// to 0.0 when `dir < 1e-9`, masking whether dir is starved (tiny) or
    /// mag is zero. Raw norms disambiguate. Zeros on accessor failure
    /// (non-fatal — diag only), matching the Task 0.5 / Task 2.0 pattern.
    pub(crate) fn grad_norms_dir_mag_abs(&mut self) -> (f32, f32) {
        match self.trainer.grad_norms_dir_mag_abs() {
            Ok(v) => v,
            Err(e) => {
                tracing::warn!("grad_norms_dir_mag_abs failed: {e}");
                (0.0, 0.0)
            }
        }
    }

    /// Task 2.0 — refresh the cached per-component grad-norm arrays from the
    /// pinned result slot populated by the in-graph reduction kernel.
    /// Epoch-boundary cold path; stream-syncs once. Returns `Err` on
    /// pinned-buffer absence or sync failure — the caller is expected to
    /// `tracing::warn!` + fall back to stale values (same pattern as
    /// `update_q_mag_means_cached`).
    pub(crate) fn refresh_grad_component_norms(&mut self) -> Result<(), crate::MLError> {
        self.trainer.refresh_grad_component_norms()
    }

    /// Task 2.0 — IQN magnitude/direction grad-norm ratio. Index 0 in the
    /// `grad_component_norms_*` caches (populated by
    /// `refresh_grad_component_norms`). Index order (9 slots, Task 2.0
    /// extension):
    /// `[iqn, cql, cql_sx, c51, c51_bs, distill, rec, pred, ens]`.
    pub(crate) fn grad_mag_iqn_ratio(&self) -> f32 {
        let m = self.trainer.grad_component_norms_mag()[0];
        let d = self.trainer.grad_component_norms_dir()[0];
        if d > 1e-9 { m / d } else { 0.0 }
    }
    /// Task 2.0 — CQL magnitude/direction grad-norm ratio (index 1; spans
    /// the snapshot→reduction pair that wraps `apply_cql_gradient` +
    /// `apply_cql_saxpy` together).
    pub(crate) fn grad_mag_cql_ratio(&self) -> f32 {
        let m = self.trainer.grad_component_norms_mag()[1];
        let d = self.trainer.grad_component_norms_dir()[1];
        if d > 1e-9 { m / d } else { 0.0 }
    }
    /// Task 2.0 extension — CQL-SAXPY magnitude/direction grad-norm ratio
    /// (index 2; isolates only `apply_cql_saxpy` — the budget-scaled add of
    /// `cql_grad_scratch` into `grad_buf`).
    pub(crate) fn grad_mag_cql_sx_ratio(&self) -> f32 {
        let m = self.trainer.grad_component_norms_mag()[2];
        let d = self.trainer.grad_component_norms_dir()[2];
        if d > 1e-9 { m / d } else { 0.0 }
    }
    /// Task 2.0 — C51 magnitude/direction grad-norm ratio (index 3; pre-
    /// budget-scale — same semantics as first-pass Task 2.0).
    pub(crate) fn grad_mag_c51_ratio(&self) -> f32 {
        let m = self.trainer.grad_component_norms_mag()[3];
        let d = self.trainer.grad_component_norms_dir()[3];
        if d > 1e-9 { m / d } else { 0.0 }
    }
    /// Task 2.0 extension — C51-budget-scale magnitude/direction delta ratio
    /// (index 4; isolates the multiplicative amplitude change applied by
    /// `apply_c51_budget_scale`).
    pub(crate) fn grad_mag_c51_bs_ratio(&self) -> f32 {
        let m = self.trainer.grad_component_norms_mag()[4];
        let d = self.trainer.grad_component_norms_dir()[4];
        if d > 1e-9 { m / d } else { 0.0 }
    }
    /// Task 2.0 extension — Distillation magnitude/direction grad-norm ratio
    /// (index 5; isolates `apply_distillation_gradient`).
    pub(crate) fn grad_mag_distill_ratio(&self) -> f32 {
        let m = self.trainer.grad_component_norms_mag()[5];
        let d = self.trainer.grad_component_norms_dir()[5];
        if d > 1e-9 { m / d } else { 0.0 }
    }
    /// Task 2.0 extension — Recursive-confidence magnitude/direction grad-
    /// norm ratio (index 6; isolates `launch_recursive_confidence_backward`).
    pub(crate) fn grad_mag_rec_ratio(&self) -> f32 {
        let m = self.trainer.grad_component_norms_mag()[6];
        let d = self.trainer.grad_component_norms_dir()[6];
        if d > 1e-9 { m / d } else { 0.0 }
    }
    /// Task 2.0 extension — Predictive-coding magnitude/direction grad-norm
    /// ratio (index 7; isolates `compute_predictive_coding_loss`).
    pub(crate) fn grad_mag_pred_ratio(&self) -> f32 {
        let m = self.trainer.grad_component_norms_mag()[7];
        let d = self.trainer.grad_component_norms_dir()[7];
        if d > 1e-9 { m / d } else { 0.0 }
    }
    /// Task 2.0 — Ensemble magnitude/direction grad-norm ratio (index 8).
    pub(crate) fn grad_mag_ens_ratio(&self) -> f32 {
        let m = self.trainer.grad_component_norms_mag()[8];
        let d = self.trainer.grad_component_norms_dir()[8];
        if d > 1e-9 { m / d } else { 0.0 }
    }

    /// Task 2.Z diagnostic — absolute trunk-slice L2 norms per component
    /// populated by `refresh_grad_component_norms`. Index order matches
    /// the mag/dir accessors:
    /// `[iqn, cql, cql_sx, c51, c51_bs, distill, rec, pred, ens]`.
    ///
    /// Values report the L2 norm of the additive delta each component's
    /// backward / scaling op applied to trunk tensors 0..4 (w_s1, b_s1,
    /// w_s2, b_s2). Expected non-zero on `[iqn]` (via
    /// `apply_iqn_trunk_gradient`) and `[ens]` (via
    /// `apply_ensemble_diversity_backward`) whenever their readiness /
    /// scale factors are non-zero; expected near-zero on all others
    /// (they write exclusively to branch heads 8..16).
    pub(crate) fn grad_trunk_norms_by_component(&self) -> [f32; 9] {
        self.trainer.grad_component_norms_trunk()
    }
    /// Task 2.Z diagnostic — per-component absolute trunk L2 norm accessors.
    pub(crate) fn grad_trunk_iqn_abs(&self)     -> f32 { self.trainer.grad_component_norms_trunk()[0] }
    pub(crate) fn grad_trunk_cql_abs(&self)     -> f32 { self.trainer.grad_component_norms_trunk()[1] }
    pub(crate) fn grad_trunk_cql_sx_abs(&self)  -> f32 { self.trainer.grad_component_norms_trunk()[2] }
    pub(crate) fn grad_trunk_c51_abs(&self)     -> f32 { self.trainer.grad_component_norms_trunk()[3] }
    pub(crate) fn grad_trunk_c51_bs_abs(&self)  -> f32 { self.trainer.grad_component_norms_trunk()[4] }
    pub(crate) fn grad_trunk_distill_abs(&self) -> f32 { self.trainer.grad_component_norms_trunk()[5] }
    pub(crate) fn grad_trunk_rec_abs(&self)     -> f32 { self.trainer.grad_component_norms_trunk()[6] }
    pub(crate) fn grad_trunk_pred_abs(&self)    -> f32 { self.trainer.grad_component_norms_trunk()[7] }
    pub(crate) fn grad_trunk_ens_abs(&self)     -> f32 { self.trainer.grad_component_norms_trunk()[8] }

    // Legacy `per_branch_target_drift()` and `per_branch_vsn_mean()` removed
    // 2026-04-25 per pearl_cold_path_no_exception_to_gpu_drives.md. Replaced by
    // GPU kernels writing to ISV slots; HEALTH_DIAG reads via read_isv_signal_at:
    //   - vsn_mag/vsn_dir: ISV[VSN_MAG_EMA_INDEX=87], ISV[88]
    //   - drift_mag/drift_dir: ISV[TARGET_DRIFT_MAG_EMA_INDEX=92], ISV[93]

    /// Plan 4 Task 5 Mode A (E.5 attention-focus): forwarding accessors for the
    /// GPU-side reduction kernel. No host work, no DtoH.
    pub(crate) fn vsn_branch_slice_indices(&self) -> (usize, usize, usize, usize) {
        self.trainer.vsn_branch_slice_indices()
    }
    pub(crate) fn mamba2_h_enriched_device(&self, batch_size: usize) -> (u64, usize) {
        self.trainer.mamba2_h_enriched_device(batch_size)
    }
    pub(crate) fn params_buf_device(&self) -> u64 {
        self.trainer.params_buf_device_ptr()
    }
    /// Plan 4 follow-up: target-drift kernel accessors.
    pub(crate) fn target_params_buf_device(&self) -> u64 {
        self.trainer.target_params_buf_device_ptr()
    }
    pub(crate) fn branch_param_slice_indices(&self) -> (usize, usize, usize, usize) {
        self.trainer.branch_param_slice_indices()
    }

    /// Task 0.3 — per-magnitude-bucket Q-value mean.
    ///
    /// Returns `[q_full, q_half, q_quarter]` read from the cached values
    /// populated by `update_q_mag_means_cached` (cold-path, epoch boundary).
    /// Callers should invoke `update_q_mag_means_cached` before reading.
    pub(crate) fn q_magnitude_bucket_means(&self) -> [f32; 3] {
        self.trainer.q_magnitude_bucket_means()
    }

    /// Task 0.3 — refresh the cached magnitude-branch Q-value means.
    ///
    /// Cold-path dtoh of `q_out_buf` contents. Should be called at HEALTH_DIAG
    /// emit time (epoch boundary) immediately before `q_magnitude_bucket_means`.
    pub(crate) fn update_q_mag_means_cached(&mut self) -> Result<(), crate::MLError> {
        self.trainer.update_q_mag_means_cached()
    }

    /// Flush the last in-flight async readback from the training step.
    /// Called once at epoch end to retrieve the final step's scalars.
    pub(crate) fn flush_readback(&mut self) -> Result<crate::cuda_pipeline::gpu_dqn_trainer::FusedTrainScalars> {
        self.trainer.flush_readback()
            .map_err(|e| anyhow::anyhow!("flush_readback: {e}"))
    }

    /// Run one full fused training step -- GPU-only, no CPU paths.
    ///
    /// Single parent graph launch replays all 11 children:
    ///   counters -> spectral -> forward -> ddqn -> aux -> post_aux ->
    ///   adam_grad -> adam -> maintenance -> iql_modulate -> per_priority
    ///
    /// Step 0 runs ungraphed + captures all children + composes parent.
    /// Steps 1+ replay the parent graph via single `cuGraphLaunch`.
    ///
    /// After graph launch: only host-side pinned scalar reads (nanoseconds, no GPU ops).
    pub(crate) fn run_full_step(
        &mut self,
        agent: &mut DQNAgentType,
        _device: &MlDevice,
    ) -> Result<FusedStepResult> {
        let cu_stream = self.stream.cu_stream();

        // ── Graph launch or ungraphed step 0 ──────────────────────────
        if let Some(ref pg) = self.parent_graph {
            // Steps 1+: single parent graph launch replays all children.
            tracing::debug!("parent graph launch START");
            pg.launch(cu_stream)?;
            tracing::debug!("parent graph launch DONE (async — GPU working)");
        } else {
            // Ungraphed step 0: run everything, then capture + compose.
            // Ensure params are initialized for the first step.
            if !self.trainer.params_initialized {
                self.trainer.flatten_online_weights(
                    &self.online_dueling, &self.online_branching,
                ).map_err(|e| anyhow::anyhow!("flatten weights: {e}"))?;
                self.trainer.params_initialized = true;
            }
            if !self.trainer.target_params_initialized {
                let n_bytes = self.trainer.total_params() * std::mem::size_of::<f32>();
                let src = self.trainer.params_buf_ptr();
                let dst = self.trainer.target_params_buf_ptr();
                unsafe {
                    cudarc::driver::result::memcpy_dtod_async(dst, src, n_bytes, cu_stream)
                        .map_err(|e| anyhow::anyhow!("init target DtoD: {e}"))?;
                }
                self.trainer.target_params_initialized = true;
            }

            // Phase -1: PER sampling (writes directly to trainer buffers via gather_padded)
            let ptrs = agent.primary_dqn_mut().memory.gpu.sample_proportional(self.batch_size)
                .map_err(|e| anyhow::anyhow!("PER sample step 0: {e}"))?;
            let gpu_batch = crate::dqn::replay_buffer_type::GpuBatch {
                states_ptr: ptrs.states_ptr,
                actions_ptr: ptrs.actions_ptr,
                rewards_ptr: ptrs.rewards_ptr,
                next_states_ptr: ptrs.next_states_ptr,
                dones_ptr: ptrs.dones_ptr,
                weights_ptr: ptrs.weights_ptr,
                indices_ptr: ptrs.indices_ptr,
                episode_ids_ptr: ptrs.episode_ids_ptr,
                batch_size: self.batch_size,
            };

            // Phase 0a: PopArt normalize rewards in-place.
            // Rescales raw rewards to unit variance so C51 atom support [-50,50]
            // covers the actual reward range. Without this, rewards in [-1,1] use
            // only 2% of the atom range → near-zero C51 gradient.
            if self.popart_robust && self.cached_iqr > 0.0 {
                // Robust path: median/IQR from sorted epoch rewards (no running stats).
                self.trainer.normalize_rewards_robust(self.batch_size, self.cached_median, self.cached_iqr)
                    .map_err(|e| anyhow::anyhow!("PopArt robust normalize: {e}"))?;
            } else {
                // Welford path: running mean/var with warmup.
                self.trainer.normalize_rewards_popart_inplace(self.batch_size)
                    .map_err(|e| anyhow::anyhow!("PopArt normalize: {e}"))?;
            }

            // Phase 0b: Counter increments + stochastic depth mask (ungraphed)
            self.submit_counters_ops()?;

            // HER donor computation (outside graph -- indices vary per step)
            if let Some(ref mut her) = self.gpu_her {
                use crate::cuda_pipeline::gpu_her::HerGpuStrategy;
                let batch_size = self.batch_size;
                match her.config.strategy {
                    HerGpuStrategy::Random => {
                        her.generate_random_donors_gpu(batch_size)
                            .map_err(|e| anyhow::anyhow!("HER random donors GPU: {e}"))?;
                    }
                    HerGpuStrategy::Future | HerGpuStrategy::Final => {
                        let source_indices = self.her_source_indices_gpu.as_ref()
                            .ok_or_else(|| anyhow::anyhow!("HER source indices not pre-computed"))?;
                        let her_batch_size = her.config.her_batch_size();
                        her.relabel_batch_with_strategy(
                            gpu_batch.episode_ids_ptr, source_indices,
                            1, source_indices.len(), her_batch_size,
                        ).map_err(|e| anyhow::anyhow!("HER donor selection: {e}"))?;
                    }
                }
            }

            // Phase 1: Spectral norm
            self.trainer.apply_spectral_norm(&self.online_dueling, &self.online_branching)
                .map_err(|e| anyhow::anyhow!("spectral norm: {e}"))?;

            // D.8 TLOB forward: OFI[32]→TLOB[16] written into states_buf before trunk forward.
            {
                let states = self.trainer.states_buf_mut();
                self.gpu_tlob
                    .forward(states, self.batch_size)
                    .map_err(|e| anyhow::anyhow!("TLOB forward: {e}"))?;
            }

            // Phase 2: Forward (cuBLAS trunk + ISV + loss + backward)
            self.trainer.submit_forward_ops_main()
                .map_err(|e| anyhow::anyhow!("forward ops: {e}"))?;

            // Phase 3: DDQN Pass 3
            self.trainer.submit_forward_ops_ddqn()
                .map_err(|e| anyhow::anyhow!("ddqn ops: {e}"))?;

            // Phase 4: Aux ops (HER + EMA + IQL + IQN + attention + CQL)
            self.submit_aux_ops(agent, &gpu_batch)?;

            // Phase 5b: Post-aux ops (selectivity + denoise + Q-stats + risk_sgd)
            self.trainer.submit_post_aux_ops(self.batch_size)
                .map_err(|e| anyhow::anyhow!("post_aux ops: {e}"))?;

            // SP1 Phase B (Task 4) + SP2 (Task A4): single fused NaN-check
            // launch covers all 12 backward-path slots (24-35) in 12 blocks.
            // Slots 33/34/35 fire inline during backward orchestration (see
            // `launch_cublas_backward_to` and `apply_iqn_trunk_gradient` in
            // `gpu_dqn_trainer.rs`). Permanent diagnostic infrastructure;
            // flags accumulate within fold and reset on fold boundary via
            // `reset_nan_flags()`. Captures into the same `post_aux` child
            // graph as `submit_post_aux_ops` so the GPU reduce kernels share
            // launch overhead with the existing pre/post forward checks.
            //
            // SP2: IQN ptrs no longer per-step args — `populate_nan_check_meta`
            // (called once at construction) baked Option<u64> nullity into
            // the fused kernel's metadata buffer entries; null/zero entries
            // are skipped by the kernel's null-pointer guard.
            self.trainer.run_nan_checks_post_backward()
                .map_err(|e| anyhow::anyhow!("post_backward nan checks: {e}"))?;

            // Phase 6: TLOB bwd + Mamba2 bwd + OFI embed bwd + pruning + grad_norm + Adam + ISV update
            // TLOB backward: reads TLOB-slice gradient from bn_d_concat_buf, updates W_Q/K/V/O.
            {
                let concat_dim = self.trainer.bn_concat_dim();
                // TLOB starts at SL_TLOB_START=74 in the state. In the concat buffer (which is
                // [bottleneck_dim | state[market_dim..STATE_DIM]]), TLOB offset =
                //   bottleneck_dim + (TLOB_START - market_dim)
                let tlob_concat_off = self.trainer.config().bottleneck_dim
                    + (state_layout::TLOB_START
                        .saturating_sub(self.trainer.config().market_dim));
                let d_concat = self.trainer.bn_d_concat_buf() as *const _ as *mut _;
                // SAFETY: bn_d_concat_buf is not borrowed mutably elsewhere during tlob.backward.
                let d_concat_ref = unsafe { &*d_concat };
                self.gpu_tlob
                    .backward(d_concat_ref, self.batch_size, concat_dim, tlob_concat_off)
                    .map_err(|e| anyhow::anyhow!("TLOB backward: {e}"))?;
                // SP4 Layer B (Mech 9): TLOB shares `attn_adam_kernel`
                // (cubin-shared with GpuAttention). TLOB params co-live with
                // attention in the Attn (group 6) param-group taxonomy, so
                // the same ISV-driven WEIGHT_BOUND[Attn] / WD_RATE[Attn]
                // contract applies. `read_group_adam_bounds` enforces the
                // EPS_CLAMP_FLOOR=1.0 ε for cold-start and AdamW
                // non-negativity.
                let (tlob_weight_clamp_max_abs, tlob_weight_decay_value) =
                    self.trainer.read_group_adam_bounds(crate::cuda_pipeline::ParamGroup::Attn);
                // SP5 Layer B (Pearl 4): TLOB shares ParamGroup::Attn=6.
                let tlob_g = crate::cuda_pipeline::ParamGroup::Attn.idx();
                let tlob_beta1   = self.trainer.read_isv_signal_at(ADAM_BETA1_BASE + tlob_g).clamp(0.85_f32, 0.95_f32);
                let tlob_beta2   = self.trainer.read_isv_signal_at(ADAM_BETA2_BASE + tlob_g).clamp(0.99_f32, 0.9995_f32);
                let tlob_epsilon = self.trainer.read_isv_signal_at(ADAM_EPS_BASE   + tlob_g).clamp(1e-10_f32, 1e-6_f32);
                self.gpu_tlob
                    .adam_step(
                        self.trainer.config().lr,
                        self.trainer.config().max_grad_norm,
                        tlob_weight_clamp_max_abs,
                        tlob_weight_decay_value,
                        tlob_beta1,
                        tlob_beta2,
                        tlob_epsilon,
                    )
                    .map_err(|e| anyhow::anyhow!("TLOB Adam: {e}"))?;
            }
            self.mamba2_backward(self.batch_size)?;
            self.step_mamba2_adam()?;
            self.launch_ofi_embed_backward(self.batch_size)?;
            self.step_ofi_embed_adam()?;
            self.trainer.apply_pruning_mask()
                .map_err(|e| anyhow::anyhow!("pruning mask: {e}"))?;
            // Adaptive per-branch gradient-norm balancer — caps each
            // branch's weight-gradient L2 norm at
            // `num_branches × median_branch_norm`.  Must run AFTER all
            // aux grad writes complete and BEFORE `compute_grad_norm_for_adam`
            // so the global clip + Adam update both observe the rebalanced
            // gradient.  See `launch_branch_grad_balance` for the full
            // architectural rationale (num_branches = 4 is the factored
            // action-axis count, median is a per-step reference — no
            // tuned knobs).
            self.trainer.launch_branch_grad_balance()
                .map_err(|e| anyhow::anyhow!("branch_grad_balance: {e}"))?;
            self.trainer.compute_grad_norm_for_adam()
                .map_err(|e| anyhow::anyhow!("grad_norm: {e}"))?;
            self.trainer.submit_adam_ops(&self.online_dueling, &self.online_branching)
                .map_err(|e| anyhow::anyhow!("adam ops: {e}"))?;
            // Plan C Phase 2 follow-up F (2026-04-29): per-step bin_means_reduce
            // producers feed ISV slots [13..16] (magnitude) and [17..21]
            // (direction) via the captured `update_isv_signals` below. Without
            // these per-step launches the producers fire only at the once-per-
            // epoch `reduce_current_q_stats` cadence, leaving ISV[16] and
            // ISV[21] stuck near the bootstrap-EMA value (e.g. 0.05 after a
            // single α=0.05 update from raw≈1.0). The Q-drift kill criterion's
            // adaptive floor `max(0.5, 3 × max(ISV[16], ISV[21]))` then never
            // graduates above the 0.5 cold-start hard-floor and the criterion
            // fires on legitimate cold-start growth (smoke-test-n9xzr).
            // Reads `q_out_buf` populated by this step's `forward_child` (which
            // runs before `adam_update_child` per the parent graph order at
            // ~line 1304) and writes to mapped-pinned scratch consumed by the
            // immediately-following `update_isv_signals` on the same stream —
            // no extra sync, captures cleanly into the adam_update child.
            self.trainer.launch_q_mag_bin_means_reduce(self.batch_size)
                .map_err(|e| anyhow::anyhow!("q_mag_bin_means_reduce: {e}"))?;
            self.trainer.launch_q_dir_bin_means_reduce(self.batch_size)
                .map_err(|e| anyhow::anyhow!("q_dir_bin_means_reduce: {e}"))?;
            self.trainer.update_isv_signals()
                .map_err(|e| anyhow::anyhow!("ISV update: {e}"))?;

            // Phase 7: Maintenance ops (causal + vaccine)
            self.trainer.submit_maintenance_ops()
                .map_err(|e| anyhow::anyhow!("maintenance ops: {e}"))?;

            // Phase 8: IQL modulate + PER priority update (ungraphed)
            self.submit_iql_modulate_ops(agent, &gpu_batch)?;
            self.submit_per_priority_ops(agent)?;

            // Capture all children + compose parent after ungraphed step succeeds.
            if let Err(e) = self.capture_training_graph(agent, &gpu_batch) {
                tracing::warn!("parent graph capture failed (continuing ungraphed): {e}");
            }
        }

        // ── Outside graph: kernel-free host ops only ─────────────────
        self.pending_vaccine_batch = None;

        // SP4 Task A15: Pearl C post-Adam engagement-rate check.
        // Reads per-block engagement counts from the mapped-pinned buffer
        // (zero-copy host access), reduces → engagement_rate, applies
        // Pearls A+D to the rate-deficit time series.
        //
        // SP4 Layer B (2026-05-01): the main DQN Adam now splits into 3
        // sub-launches (Trunk/Value/Branches) — each writes per-block
        // counts at its own group offset, so Pearl C is now per-group:
        //   group 0 (DqnTrunk)    — DQN trunk Adam sub-launch.
        //   group 1 (DqnValue)    — DQN value-head Adam sub-launch.
        //   group 2 (DqnBranches) — DQN branch-head Adam sub-launch.
        //   group 3 (Iqn)         — IQN dual-head trainer.
        //   group 4 (IqlHigh)     — IQL high-tau trainer.
        //   group 5 (IqlLow)      — IQL low-tau trainer.
        //   group 6 (Attn)        — 4-head attention (TLOB shares cubin
        //                           but disables Pearl C).
        //   group 7 (Curiosity)   — handled in `train_curiosity_gpu`,
        //                           not here (separate caller path).
        let (trunk_n, value_n, branches_n) = self.trainer.dqn_group_param_counts();
        if let Err(e) = self.trainer.pearl_c_post_adam_engagement_check(
            crate::cuda_pipeline::ParamGroup::DqnTrunk,
            trunk_n,
        ) {
            tracing::warn!("SP4 Pearl C check (DQN trunk) failed (non-fatal): {e}");
        }
        if let Err(e) = self.trainer.pearl_c_post_adam_engagement_check(
            crate::cuda_pipeline::ParamGroup::DqnValue,
            value_n,
        ) {
            tracing::warn!("SP4 Pearl C check (DQN value) failed (non-fatal): {e}");
        }
        if let Err(e) = self.trainer.pearl_c_post_adam_engagement_check(
            crate::cuda_pipeline::ParamGroup::DqnBranches,
            branches_n,
        ) {
            tracing::warn!("SP4 Pearl C check (DQN branches) failed (non-fatal): {e}");
        }
        if let Some(ref iqn) = self.gpu_iqn {
            let n = iqn.total_params();
            if let Err(e) = self.trainer.pearl_c_post_adam_engagement_check(
                crate::cuda_pipeline::ParamGroup::Iqn,
                n,
            ) {
                tracing::warn!("SP4 Pearl C check (IQN) failed (non-fatal): {e}");
            }
        }
        let iql_high_n = self.gpu_iql.total_params();
        if let Err(e) = self.trainer.pearl_c_post_adam_engagement_check(
            crate::cuda_pipeline::ParamGroup::IqlHigh,
            iql_high_n,
        ) {
            tracing::warn!("SP4 Pearl C check (IQL high) failed (non-fatal): {e}");
        }
        let iql_low_n = self.gpu_iql_low.total_params();
        if let Err(e) = self.trainer.pearl_c_post_adam_engagement_check(
            crate::cuda_pipeline::ParamGroup::IqlLow,
            iql_low_n,
        ) {
            tracing::warn!("SP4 Pearl C check (IQL low) failed (non-fatal): {e}");
        }
        if let Some(ref attn) = self.gpu_attention {
            let n = attn.total_params();
            if let Err(e) = self.trainer.pearl_c_post_adam_engagement_check(
                crate::cuda_pipeline::ParamGroup::Attn,
                n,
            ) {
                tracing::warn!("SP4 Pearl C check (Attn) failed (non-fatal): {e}");
            }
        }

        // Adaptive IQN lambda: read loss OUTSIDE graph (pinned read, no kernel).
        if let Some(ref iqn) = self.gpu_iqn {
            if let Ok(loss) = iqn.read_total_loss() {
                self.trainer.update_iqn_readiness(loss);
            }
        }

        // Readback training scalars from pinned device-mapped memory (no kernel).
        // Values are from the PREVIOUS step (double-buffered via GPU→host latency).
        let fused_result = self.trainer.readback_training_scalars()
            .map_err(|e| anyhow::anyhow!("readback scalars: {e}"))?;
        // Accumulate for epoch-level metrics (every step, not just guard steps)
        self.loss_accum += fused_result.total_loss as f64;
        self.grad_accum += fused_result.grad_norm as f64;
        self.accum_steps += 1;

        // Host-only warmup counter increment.
        self.increment_warmup();

        // Host-only bookkeeping (training_steps++, PER beta annealing).
        agent.fused_post_step_gpu_bookkeeping()
            .map_err(|e| anyhow::anyhow!("Fused GPU bookkeeping: {e}"))?;

        self.steps_since_varmap_sync += 1;
        self.last_combined_norm = fused_result.grad_norm;
        self.last_loss = fused_result.total_loss;
        self.phase_batch_count += 1;

        tracing::debug!(
            step = self.steps_since_varmap_sync,
            graphed = self.parent_graph.is_some(),
            "fused step complete"
        );

        Ok(FusedStepResult {
            total_loss: fused_result.total_loss,
            grad_norm: fused_result.grad_norm,
        })
    }

    /// Log epoch-level timing summary (wall-clock, no CUDA events).
    ///
    /// PhaseEvents removed: single parent graph launch makes per-child event
    /// profiling irrelevant. Batch count is logged for throughput measurement.
    pub(crate) fn log_phase_timing(&mut self) {
        if self.phase_batch_count == 0 { return; }
        tracing::info!(
            batches = self.phase_batch_count,
            graphed = self.parent_graph.is_some(),
            "Epoch training: {} batches (parent graph: {})",
            self.phase_batch_count,
            if self.parent_graph.is_some() { "single launch" } else { "ungraphed" },
        );
        self.phase_batch_count = 0;
    }

    /// Submit all auxiliary ops (captured into aux_child sub-graph).
    ///
    /// These ops are deterministic per-step (same kernel topology) and suitable
    /// for CUDA graph capture. Causal + vaccine are in maintenance_child.
    /// Selectivity + denoise + risk_sgd + multi-horizon are in post_aux_child.
    ///
    /// Ops submitted: HER relabel, C51 clip, EMA target, attention, IQL, IQN, CQL, regime_scale.
    fn submit_aux_ops(
        &mut self,
        agent: &mut DQNAgentType,
        gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch,
    ) -> Result<()> {
        // HER in-place relabeling kernel (donor indices already pre-computed).
        if let Some(ref her) = self.gpu_her {
            let her_batch_size = her.config.her_batch_size();
            let state_dim = ml_core::state_layout::STATE_DIM;
            let goal_dim = her.config.goal_dim.min(state_dim);
            let normal_count = self.batch_size.saturating_sub(her_batch_size);

            let src_next_ptr = gpu_batch.next_states_ptr;
            let donor_ptr = her.donor_indices.raw_ptr();
            self.trainer.launch_her_inplace_relabel(
                src_next_ptr, donor_ptr,
                normal_count, her_batch_size, goal_dim, state_dim,
            ).map_err(|e| anyhow::anyhow!("HER in-place relabel kernel: {e}"))?;
        }

        // Task 2.0 — measure C51 gradient contribution to branch 0 (direction)
        // and branch 1 (magnitude). C51's backward ran inside
        // `submit_forward_ops_main` (Phase 2); `grad_buf` was memset-zeroed
        // at its start. At this point `grad_buf` holds exactly the C51/MSE
        // blended backward — no IQN/CQL/Ens contributions yet. Reduction
        // kernel runs BEFORE `apply_c51_budget_scale` so we capture the raw
        // backward contribution rather than the budget-scaled amplitude.
        // Uses `grad_zero_ref_buf` as the "snapshot" (permanent zeros), so
        // kernel computes `‖grad_buf − 0‖` on branch 0+1 slices.
        self.trainer.grad_decomp_launch_c51()
            .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp C51: {e}"))?;

        // NOTE: Per-component gradient budget clip REMOVED (old explicit clip → constant grad norm).
        // B4/G5: Budgets are now applied as SCALE factors on individual SAXPY contributions,
        // not as norm clips. This preserves gradient magnitude variation while steering
        // component balance. Safety is still handled by the single global clip in Adam.

        // SP6 Pearl 2 + SP7 (2026-05-03): launches the on-device dispatch
        // kernel populating `lb_budget_effective_buf` with the resolved CQL/C51
        // trunk + per-branch corrections; returns IQN and ENS scalars as
        // before (prompt-scoped — host-branch dispatch only applied to CQL/C51).
        let (iqn_branch, _ens_branch, _iqn_trunk, _ens_trunk) =
            self.compute_adaptive_budgets();

        // F8/G5: Scale C51 contribution in grad_buf by c51_budget.
        // C51 backward already wrote into grad_buf during submit_forward_ops_main (Phase 2).
        // This scale MUST run before CQL/IQN SAXPYs below so only C51's portion is scaled.
        // Order: c51_budget × C51_grad (here), then cql_budget × CQL_grad, iqn_budget × IQN_grad.
        //
        // Task 2.0 extension — snapshot BEFORE the multiplicative scale so the
        // post-scale reduction measures the amplitude change on branch 0+1.
        // When `c51_budget ≈ 1.0` the scale is a no-op and delta = 0.0.
        self.trainer.grad_decomp_snapshot_c51_bs()
            .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_c51_bs: {e}"))?;
        // SP6 Pearl 2 + SP7 (2026-05-03): trunk/value first (mean budget), then
        // per-branch correction sub-launches. Both kernels now read `alpha`
        // from `lb_budget_effective_buf` at execution time — capture-safe.
        // Net effect: branch slice ends up scaled by trunk × correction[b]
        // = effective[b], identical to the prior host-resolved semantics but
        // with values written by the dispatch kernel against the runtime ISV
        // state (not frozen at capture-time step 0). The "skip when
        // correction ≈ 1.0" optimization is dropped — the value lives on-
        // device and can't be branched on host-side without reintroducing
        // the same capture-freeze bug. Cost: 4 cheap kernel launches with
        // alpha=1.0 in the steady state (multiply by 1.0 is a numeric no-op).
        let c51_trunk_dev_ptr = self.trainer.lb_budget_c51_trunk_dev_ptr();
        self.trainer.apply_c51_budget_scale(c51_trunk_dev_ptr)
            .map_err(|e| anyhow::anyhow!("c51_budget_scale_trunk: {e}"))?;
        for b in 0..4_usize {
            let correction_dev_ptr = self.trainer.lb_budget_c51_correction_dev_ptr(b);
            self.trainer.apply_c51_budget_scale_branch(b, correction_dev_ptr)
                .map_err(|e| anyhow::anyhow!("c51_budget_scale_branch[{b}]: {e}"))?;
        }
        self.trainer.grad_decomp_launch_c51_bs()
            .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp c51_bs: {e}"))?;

        // DIAGNOSTIC: sync between each aux op to find hanging kernel

        // EMA target update — tau read from ISV[TAU_EFF_INDEX] (GPU-computed).
        {
            let tau = self.trainer.read_isv_signal_at(TAU_EFF_INDEX);
            // Tau was computed by tau_update kernel at epoch boundary from
            // ISV[EPOCH_IDX, TOTAL_EPOCHS, LEARNING_HEALTH] — cosine schedule
            // with health-coupled floor. No CPU-side computation here.
            self.trainer.target_ema_update(
                &self.online_dueling, &self.online_branching,
                &self.target_dueling, &self.target_branching,
                tau,
            ).map_err(|e| anyhow::anyhow!("GPU EMA target update: {e}"))?;
        }


        // ── IQL pipeline: gather → train both → advantages → support → branch scales → gap → epsilon ──
        // IQL runs first on the main stream; IQN + Attention fork in parallel after it.
        {
            // 1. Populate q_out_buf from current logits.
            self.trainer.populate_q_out(self.batch_size)
                .map_err(|e| anyhow::anyhow!("IQL populate_q_out: {e}"))?;

            // 2. Gather Q(s, a_taken) for both IQL trainers.
            let q_out = self.trainer.q_out_buf();
            let actions = self.trainer.actions_buf();
            let total_actions = self.trainer.total_actions();
            self.gpu_iql.gather_q_taken(q_out, actions, total_actions)
                .map_err(|e| anyhow::anyhow!("IQL high gather: {e}"))?;
            let q_out = self.trainer.q_out_buf();
            let actions = self.trainer.actions_buf();
            self.gpu_iql_low.gather_q_taken(q_out, actions, total_actions)
                .map_err(|e| anyhow::anyhow!("IQL low gather: {e}"))?;

            // 3. Train V(s) — both tau=high and tau=low.
            //
            // SP4 Layer B (Mech 9): ISV-driven post-Adam |p| clamp for
            // `iql_adam_kernel`. Per-group bounds: IqlHigh (group 4) and
            // IqlLow (group 5) each read `ISV[WEIGHT_BOUND[group]]` and
            // `ISV[WD_RATE[group]]` via `read_group_adam_bounds`, which
            // applies the EPS_CLAMP_FLOOR=1.0 ε for cold-start ISV and
            // the AdamW non-negativity contract.
            let (iql_high_clamp, iql_high_wd) =
                self.trainer.read_group_adam_bounds(crate::cuda_pipeline::ParamGroup::IqlHigh);
            let (iql_low_clamp, iql_low_wd) =
                self.trainer.read_group_adam_bounds(crate::cuda_pipeline::ParamGroup::IqlLow);
            // SP5 Pearl 4: per-group Adam β1/β2/ε for IqlHigh (group 4) and IqlLow (group 5).
            // ISV[ADAM_BETA1_BASE+4=230], ISV[ADAM_BETA2_BASE+4=238], ISV[ADAM_EPS_BASE+4=246].
            // ISV[ADAM_BETA1_BASE+5=231], ISV[ADAM_BETA2_BASE+5=239], ISV[ADAM_EPS_BASE+5=247].
            // Clamps match SP4 producer envelope: β1∈[0.85,0.95], β2∈[0.99,0.9995], ε∈[1e-10,1e-6].
            let iql_high_g = crate::cuda_pipeline::ParamGroup::IqlHigh.idx();
            let iql_high_beta1   = self.trainer.read_isv_signal_at(ADAM_BETA1_BASE + iql_high_g).clamp(0.85_f32, 0.95_f32);
            let iql_high_beta2   = self.trainer.read_isv_signal_at(ADAM_BETA2_BASE + iql_high_g).clamp(0.99_f32, 0.9995_f32);
            let iql_high_epsilon = self.trainer.read_isv_signal_at(ADAM_EPS_BASE   + iql_high_g).clamp(1e-10_f32, 1e-6_f32);
            let iql_low_g = crate::cuda_pipeline::ParamGroup::IqlLow.idx();
            let iql_low_beta1   = self.trainer.read_isv_signal_at(ADAM_BETA1_BASE + iql_low_g).clamp(0.85_f32, 0.95_f32);
            let iql_low_beta2   = self.trainer.read_isv_signal_at(ADAM_BETA2_BASE + iql_low_g).clamp(0.99_f32, 0.9995_f32);
            let iql_low_epsilon = self.trainer.read_isv_signal_at(ADAM_EPS_BASE   + iql_low_g).clamp(1e-10_f32, 1e-6_f32);
            let states_f32 = self.trainer.states_buf();
            self.gpu_iql.train_value_step(states_f32, iql_high_clamp, iql_high_wd,
                iql_high_beta1, iql_high_beta2, iql_high_epsilon)
                .map_err(|e| anyhow::anyhow!("IQL high train: {e}"))?;
            let states_f32 = self.trainer.states_buf();
            self.gpu_iql_low.train_value_step(states_f32, iql_low_clamp, iql_low_wd,
                iql_low_beta1, iql_low_beta2, iql_low_epsilon)
                .map_err(|e| anyhow::anyhow!("IQL low train: {e}"))?;

            // 4. Compute advantage weights (from high-tau V(s)).
            let q_out = self.trainer.q_out_buf();
            let actions = self.trainer.actions_buf();
            self.gpu_iql.compute_advantage_weights(q_out, actions, total_actions)
                .map_err(|e| anyhow::anyhow!("IQL advantage_weights: {e}"))?;

            // 5. Update advantage sigma EMA.
            self.gpu_iql.update_adv_sigma()
                .map_err(|e| anyhow::anyhow!("IQL adv_sigma: {e}"))?;

            // 6. Compute per-sample, per-branch C51 support.
            //    Phase 2d: tile is [B, 4, 3]; centres from ISV slots 23..30,
            //    half-width from V(s)-derived Q spread.
            let q_out = self.trainer.q_out_buf();
            let isv_ptr = self.trainer.isv_signals_dev_ptr();
            self.gpu_iql.compute_per_sample_support(q_out, isv_ptr)
                .map_err(|e| anyhow::anyhow!("IQL per_sample_support: {e}"))?;

            // Floor support using Frugal p5 quantile — no silent zero-loss samples.
            self.gpu_iql.apply_support_floor()
                .map_err(|e| anyhow::anyhow!("IQL support_floor: {e}"))?;

            // 7. Compute per-branch gradient scales.
            // Per-sample branch_scales floor (2026-04-23): pass ISV device
            // pointer + slot index so the kernel can apply the ISV-driven
            // floor (slot 36) and prevent per-sample (b, d) gradient
            // collapse on losing branches when IQL readiness → 1.
            let q_out = self.trainer.q_out_buf();
            let actions = self.trainer.actions_buf();
            let isv_ptr_scales = self.trainer.isv_signals_dev_ptr();
            let iql_floor_idx: i32 =
                crate::cuda_pipeline::gpu_dqn_trainer::IQL_BRANCH_SCALE_FLOOR_INDEX as i32;
            self.gpu_iql.compute_branch_scales(q_out, actions, isv_ptr_scales, iql_floor_idx)
                .map_err(|e| anyhow::anyhow!("IQL branch_scales: {e}"))?;

            // 8. Compute expectile gap and per-sample epsilon.
            let v_low_ptr = self.gpu_iql_low.v_out_ptr();
            self.gpu_iql.compute_expectile_gap(v_low_ptr)
                .map_err(|e| anyhow::anyhow!("IQL expectile_gap: {e}"))?;

            let base_epsilon = agent.primary_dqn().config.epsilon_end;
            self.gpu_iql.compute_per_sample_epsilon(base_epsilon)
                .map_err(|e| anyhow::anyhow!("IQL per_sample_epsilon: {e}"))?;
        }

        // ── IQN prep: decode actions + DtoD copies (main stream, before fork) ──
        // Must complete before execute_training_pipeline runs on iqn_stream.
        if let Some(ref mut iqn) = self.gpu_iqn {
            // SP6 Pearl 5: refresh per-branch τ schedules from ISV[250..270).
            // ISV layout: branch b, quantile q → IQN_TAU_BASE + b*5 + q.
            // Each branch gets its own [B,N] tau slab; the main (averaged) buffer
            // is also kept current via refresh_taus_from_isv for CVaR and any
            // consumer that has not migrated to the per-branch path.
            {
                let mut isv_taus = [0.0_f32; 20];
                for i in 0..20 {
                    isv_taus[i] = self.trainer.read_isv_signal_at(IQN_TAU_BASE + i);
                }
                // Populate per-branch slabs first.
                for branch_idx in 0..4_usize {
                    let mut tau5 = [0.0_f32; 5];
                    for q in 0..5_usize {
                        tau5[q] = isv_taus[branch_idx * 5 + q];
                    }
                    iqn.refresh_taus_for_branch(branch_idx, &tau5);
                }
                // Keep the averaged main buffer current for CVaR and backward-compat consumers.
                iqn.refresh_taus_from_isv(&isv_taus);
            }
            // Note: set_cached_target_h_s2 is called inside the per-branch loops
            // below — `execute_training_pipeline` consumes the cache via `.take()`
            // each call, so the single set-before-loop pattern from the pre-SP6
            // single-branch era left branches 1..4 with `None` (silently
            // non-fatally failed). See SP6 Pearl 5 contract.
            let dqn_actions = self.trainer.actions_buf();
            let dqn_rewards = self.trainer.rewards_buf();
            let dqn_dones = self.trainer.dones_buf();
            // train_iqn_step_gpu does decode + DtoD copies on self.stream (main),
            // then calls execute_training_pipeline. We only need the prep part here;
            // execute_training_pipeline will be called separately with the override stream.
            // Prep: decode flat actions → branch_actions, copy rewards/dones.
            iqn.prepare_buffers(dqn_actions, dqn_rewards, dqn_dones)
                .map_err(|e| anyhow::anyhow!("IQN prepare_buffers: {e}"))?;
        }

        // ── Fork-join: IQN + Attention run in parallel on dedicated streams ──
        // Record iql_done_event on main stream as the fork point. Both iqn_stream
        // and attn_stream wait on this event before starting their work.
        let has_iqn = self.gpu_iqn.is_some();
        let has_attn = self.gpu_attention.is_some();
        let parallel = self.iqn_stream.is_some();
        let mut iqn_ok = false;

        // SP6 Pearl 5 / Task 2.0: snapshot grad_buf branch 0+1 slice BEFORE any
        // IQN trunk gradient fires. The parallel and sequential paths both call
        // apply_iqn_trunk_gradient inline (4 passes × budget/4), so the snapshot
        // must precede the fork to correctly isolate the IQN contribution delta.
        self.trainer.grad_decomp_snapshot_iqn()
            .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_iqn: {e}"))?;

        if parallel {
            // Record fork-point event on main stream.
            let iql_ev = self.iql_done_event.as_ref()
                .ok_or_else(|| anyhow::anyhow!("iql_done_event missing"))?;
            iql_ev.record(&self.stream)
                .map_err(|e| anyhow::anyhow!("iql_done_event record: {e}"))?;

            // ── Fork IQN to iqn_stream ──
            // SP6 Pearl 5: run 4 sequential IQN forward passes on iqn_stream, each
            // with its per-branch τ schedule active. After each pass, synchronize to
            // the main stream and apply iqn_budget/4 trunk gradient contribution so
            // all 4 branches' d_h_s2 values are folded into grad_buf equally.
            if has_iqn {
                let iqn_s = self.iqn_stream.as_ref().unwrap();
                iqn_s.wait(iql_ev)
                    .map_err(|e| anyhow::anyhow!("iqn_stream wait iql_done: {e}"))?;

                let iqn_ws = self.iqn_workspace.as_ref().unwrap();
                let ws_override = Some((iqn_ws.raw_ptr(), 32 * 1024 * 1024_usize));
                let iqn_stream_ref = Arc::clone(iqn_s);

                // SP4 Layer B (Mech 9): ISV-driven post-Adam |p| clamp for
                // `iqn_adam_kernel`. IQN params live in group 3
                // (`ParamGroup::Iqn`); `read_group_adam_bounds` returns
                // `ISV[WEIGHT_BOUND[Iqn]]` (with EPS_CLAMP_FLOOR=1.0 ε
                // for cold-start) and `ISV[WD_RATE[Iqn]]` (with AdamW
                // non-negativity).
                let (weight_clamp_max_abs, weight_decay_value) =
                    self.trainer.read_group_adam_bounds(crate::cuda_pipeline::ParamGroup::Iqn);
                // SP5 Layer B (Pearl 4): per-group Adam β1/β2/ε for Iqn=3.
                let iqn_g = crate::cuda_pipeline::ParamGroup::Iqn.idx();
                let iqn_beta1   = self.trainer.read_isv_signal_at(ADAM_BETA1_BASE + iqn_g).clamp(0.85_f32, 0.95_f32);
                let iqn_beta2   = self.trainer.read_isv_signal_at(ADAM_BETA2_BASE + iqn_g).clamp(0.99_f32, 0.9995_f32);
                let iqn_epsilon = self.trainer.read_isv_signal_at(ADAM_EPS_BASE   + iqn_g).clamp(1e-10_f32, 1e-6_f32);

                // Pearl 2 + Pearl 5 cross: use per-branch IQN budget normalized by /4 to
                // preserve total magnitude. Sum across 4 branches = sum(iqn_branch)/4 =
                // iqn_trunk (since iqn_trunk = mean(iqn_branch)). Per-branch differentiation
                // from Pearl 2 preserved; SP5 Layer B gradient magnitude contract maintained.
                // `iqn_branch` is indexed by `branch_idx` inside the loop below.
                let iqn_ev = self.iqn_done_event.as_ref().unwrap();

                for branch_idx in 0..4_usize {
                    let iqn_budget_per_branch = iqn_branch[branch_idx] / 4.0_f32;
                    let online_h_s2 = self.trainer.save_h_s2();
                    let next_states_buf = self.trainer.next_states_buf();
                    let tg_h_s2_ptr = self.trainer.tg_h_s2_ptr();
                    let iqn = self.gpu_iqn.as_mut().unwrap();

                    // execute_training_pipeline consumes the cached ptr (.take()), so
                    // re-set per branch. Same value (target_h_s2 is per-step, not
                    // per-branch) but the contract is "set before each call".
                    iqn.set_cached_target_h_s2(tg_h_s2_ptr);
                    iqn.activate_branch_taus(branch_idx);
                    let pass_ok = match iqn.execute_training_pipeline(
                        online_h_s2, next_states_buf, &self.target_dueling,
                        Some(&iqn_stream_ref), ws_override,
                        weight_clamp_max_abs,
                        weight_decay_value,
                        iqn_beta1,
                        iqn_beta2,
                        iqn_epsilon,
                    ) {
                        Ok(_) => true,
                        Err(e) => {
                            tracing::warn!("IQN parallel branch {branch_idx} step failed (non-fatal): {e}");
                            false
                        }
                    };
                    iqn.deactivate_branch_taus(branch_idx);

                    if pass_ok {
                        // Synchronize iqn_stream → main stream so apply_iqn_trunk_gradient
                        // (cuBLAS, main stream) sees the completed d_h_s2 from this pass.
                        iqn_ev.record(&iqn_stream_ref)
                            .map_err(|e| anyhow::anyhow!("iqn branch{branch_idx} event record: {e}"))?;
                        self.stream.wait(iqn_ev)
                            .map_err(|e| anyhow::anyhow!("main wait iqn branch{branch_idx}: {e}"))?;

                        let d_h_s2_ptr = iqn.d_h_s2_raw_ptr();
                        self.trainer.apply_iqn_trunk_gradient(
                            d_h_s2_ptr, &mut self.online_dueling, iqn_budget_per_branch,
                        ).map_err(|e| anyhow::anyhow!("IQN trunk gradient branch {branch_idx}: {e}"))?;

                        // Re-fork iqn_stream off main so the next branch pass starts
                        // after main has consumed the gradient (prevents use-after-free
                        // of d_h_s2_buf if iqn_stream were to get ahead of main).
                        iql_ev.record(&self.stream)
                            .map_err(|e| anyhow::anyhow!("re-fork iqn branch{branch_idx}: {e}"))?;
                        iqn_stream_ref.wait(iql_ev)
                            .map_err(|e| anyhow::anyhow!("iqn_stream re-fork wait {branch_idx}: {e}"))?;

                        iqn_ok = true;
                    }
                }

                // Record final iqn_stream completion for the join below.
                iqn_ev.record(&iqn_stream_ref)
                    .map_err(|e| anyhow::anyhow!("iqn_done_event record: {e}"))?;
            }

            // ── Fork Attention to attn_stream ──
            if has_attn {
                let attn_s = self.attn_stream.as_ref().unwrap();
                let iql_ev = self.iql_done_event.as_ref().unwrap();
                attn_s.wait(iql_ev)
                    .map_err(|e| anyhow::anyhow!("attn_stream wait iql_done: {e}"))?;

                let attn_ws = self.attn_workspace.as_ref().unwrap();
                let ws_override = Some((attn_ws.raw_ptr(), 32 * 1024 * 1024_usize));
                let attn_stream_ref = Arc::clone(attn_s);
                let batch_size = self.batch_size;
                let sh2 = self.trainer.shared_h2();
                let n_bytes = batch_size * sh2 * std::mem::size_of::<f32>();

                // Copy save_h_s2 → attn_bw_scratch on attn_stream (graph-safe kernel copy).
                let h_s2_ptr = self.trainer.save_h_s2().raw_ptr();
                let scratch_ptr = self.trainer.attn_bw_scratch().raw_ptr();
                self.trainer.graph_safe_copy_f32_on(
                    scratch_ptr, h_s2_ptr, n_bytes, &attn_stream_ref, "attn h_s2→scratch",
                ).map_err(|e| anyhow::anyhow!("{e}"))?;

                let attn = self.gpu_attention.as_mut().unwrap();

                // Forward on attn_stream with widened input: [D,B] states + [10,B] ofi_embed.
                let attn_output_ptr = {
                    let out = attn.forward(
                        self.trainer.attn_bw_scratch(),
                        self.trainer.ofi_embed_output(),
                        batch_size,
                        Some(&attn_stream_ref), ws_override,
                    ).map_err(|e| anyhow::anyhow!("Attention forward: {e}"))?;
                    out.raw_ptr()
                };

                // Copy attention output → save_h_s2 on attn_stream (graph-safe kernel copy).
                let save_dst = self.trainer.save_h_s2().raw_ptr();
                self.trainer.graph_safe_copy_f32_on(
                    save_dst, attn_output_ptr, n_bytes, &attn_stream_ref, "attn output→h_s2",
                ).map_err(|e| anyhow::anyhow!("{e}"))?;

                // Copy bw_d_h_s2 → attn_bw_scratch on attn_stream for backward (graph-safe).
                let d_h_s2_ptr = self.trainer.bw_d_h_s2_buf().raw_ptr();
                self.trainer.graph_safe_copy_f32_on(
                    scratch_ptr, d_h_s2_ptr, n_bytes, &attn_stream_ref, "attn bw_d_h_s2→scratch",
                ).map_err(|e| anyhow::anyhow!("{e}"))?;

                // Backward on attn_stream.
                attn.backward(
                    self.trainer.attn_bw_scratch(), batch_size,
                    Some(&attn_stream_ref), ws_override,
                ).map_err(|e| anyhow::anyhow!("Attention backward: {e}"))?;

                // Adam on attn_stream.
                let lr = self.trainer.config().lr;
                let mgn = self.trainer.config().max_grad_norm;
                // SP4 Layer B (Mech 9): ISV-driven post-Adam |p| clamp for
                // `attn_adam_kernel`. Attention params live in group 6
                // (`ParamGroup::Attn`); `read_group_adam_bounds` returns
                // `ISV[WEIGHT_BOUND[Attn]]` (with EPS_CLAMP_FLOOR=1.0 ε
                // for cold-start) and `ISV[WD_RATE[Attn]]` (with AdamW
                // non-negativity).
                let (attn_weight_clamp_max_abs, attn_weight_decay_value) =
                    self.trainer.read_group_adam_bounds(crate::cuda_pipeline::ParamGroup::Attn);
                // SP5 Layer B (Pearl 4): per-group Adam β1/β2/ε for Attn=6.
                let attn_g = crate::cuda_pipeline::ParamGroup::Attn.idx();
                let attn_beta1   = self.trainer.read_isv_signal_at(ADAM_BETA1_BASE + attn_g).clamp(0.85_f32, 0.95_f32);
                let attn_beta2   = self.trainer.read_isv_signal_at(ADAM_BETA2_BASE + attn_g).clamp(0.99_f32, 0.9995_f32);
                let attn_epsilon = self.trainer.read_isv_signal_at(ADAM_EPS_BASE   + attn_g).clamp(1e-10_f32, 1e-6_f32);
                attn.adam_step(lr, mgn, Some(&attn_stream_ref), ws_override,
                               attn_weight_clamp_max_abs, attn_weight_decay_value,
                               attn_beta1, attn_beta2, attn_epsilon)
                    .map_err(|e| anyhow::anyhow!("Attention Adam: {e}"))?;

                // Record completion on attn_stream.
                let attn_ev = self.attn_done_event.as_ref().unwrap();
                attn_ev.record(self.attn_stream.as_ref().unwrap())
                    .map_err(|e| anyhow::anyhow!("attn_done_event record: {e}"))?;
            }

            // ── Join: main stream waits for both parallel streams ──
            if has_iqn {
                let iqn_ev = self.iqn_done_event.as_ref().unwrap();
                self.stream.wait(iqn_ev)
                    .map_err(|e| anyhow::anyhow!("main wait iqn_done: {e}"))?;
            }
            if has_attn {
                let attn_ev = self.attn_done_event.as_ref().unwrap();
                self.stream.wait(attn_ev)
                    .map_err(|e| anyhow::anyhow!("main wait attn_done: {e}"))?;
            }
        } else {
            // ── Sequential fallback: no dedicated streams ──

            // Attention forward + backward + Adam (sequential on main stream).
            if let Some(ref mut attn) = self.gpu_attention {
                self.trainer.apply_attention_forward(attn, self.batch_size)
                    .map_err(|e| anyhow::anyhow!("Attention forward: {e}"))?;
                let d_h_s2_ref = self.trainer.bw_d_h_s2()
                    .map_err(|e| anyhow::anyhow!("bw_d_h_s2 ref: {e}"))?;
                attn.backward(d_h_s2_ref, self.batch_size, None, None)
                    .map_err(|e| anyhow::anyhow!("Attention backward: {e}"))?;
                let lr = self.trainer.config().lr;
                let mgn = self.trainer.config().max_grad_norm;
                // SP4 Layer B (Mech 9): same ISV-driven bound as parallel arm.
                // Sequential-fallback path.
                let (attn_weight_clamp_max_abs, attn_weight_decay_value) =
                    self.trainer.read_group_adam_bounds(crate::cuda_pipeline::ParamGroup::Attn);
                // SP5 Layer B (Pearl 4): per-group Adam β1/β2/ε for Attn=6.
                let attn_g = crate::cuda_pipeline::ParamGroup::Attn.idx();
                let attn_beta1   = self.trainer.read_isv_signal_at(ADAM_BETA1_BASE + attn_g).clamp(0.85_f32, 0.95_f32);
                let attn_beta2   = self.trainer.read_isv_signal_at(ADAM_BETA2_BASE + attn_g).clamp(0.99_f32, 0.9995_f32);
                let attn_epsilon = self.trainer.read_isv_signal_at(ADAM_EPS_BASE   + attn_g).clamp(1e-10_f32, 1e-6_f32);
                attn.adam_step(lr, mgn, None, None,
                               attn_weight_clamp_max_abs, attn_weight_decay_value,
                               attn_beta1, attn_beta2, attn_epsilon)
                    .map_err(|e| anyhow::anyhow!("Attention Adam: {e}"))?;
            }

            // SP6 Pearl 5: IQN training pipeline — 4 sequential passes on main stream,
            // each using a per-branch τ schedule. ÷4 budget normalization preserves
            // total IQN gradient magnitude at iqn_budget (same as SP5 Layer B contract).
            if has_iqn {
                // SP4 Layer B (Mech 9): same ISV-driven bound as the parallel arm.
                let (weight_clamp_max_abs, weight_decay_value) =
                    self.trainer.read_group_adam_bounds(crate::cuda_pipeline::ParamGroup::Iqn);
                // SP5 Layer B (Pearl 4): per-group Adam β1/β2/ε for Iqn=3.
                let iqn_g = crate::cuda_pipeline::ParamGroup::Iqn.idx();
                let iqn_beta1   = self.trainer.read_isv_signal_at(ADAM_BETA1_BASE + iqn_g).clamp(0.85_f32, 0.95_f32);
                let iqn_beta2   = self.trainer.read_isv_signal_at(ADAM_BETA2_BASE + iqn_g).clamp(0.99_f32, 0.9995_f32);
                let iqn_epsilon = self.trainer.read_isv_signal_at(ADAM_EPS_BASE   + iqn_g).clamp(1e-10_f32, 1e-6_f32);

                for branch_idx in 0..4_usize {
                    // Pearl 2 + Pearl 5 cross: per-branch IQN budget; /4 preserves total
                    // magnitude (sum across 4 branches = iqn_trunk).
                    let iqn_budget_per_branch = iqn_branch[branch_idx] / 4.0_f32;
                    let online_h_s2 = self.trainer.save_h_s2();
                    let next_states_buf = self.trainer.next_states_buf();
                    let tg_h_s2_ptr = self.trainer.tg_h_s2_ptr();
                    let iqn = self.gpu_iqn.as_mut().unwrap();

                    // Mirror parallel-arm fix: execute_training_pipeline consumes the
                    // cached ptr each call, so set per-branch.
                    iqn.set_cached_target_h_s2(tg_h_s2_ptr);
                    iqn.activate_branch_taus(branch_idx);
                    let pass_ok = match iqn.execute_training_pipeline(
                        online_h_s2, next_states_buf, &self.target_dueling,
                        None, None,
                        weight_clamp_max_abs,
                        weight_decay_value,
                        iqn_beta1,
                        iqn_beta2,
                        iqn_epsilon,
                    ) {
                        Ok(_) => true,
                        Err(e) => {
                            tracing::warn!("IQN sequential branch {branch_idx} step failed (non-fatal): {e}");
                            false
                        }
                    };
                    iqn.deactivate_branch_taus(branch_idx);

                    if pass_ok {
                        // Apply this branch's gradient contribution immediately while
                        // d_h_s2_buf still holds this pass's result.
                        let d_h_s2_ptr = iqn.d_h_s2_raw_ptr();
                        self.trainer.apply_iqn_trunk_gradient(
                            d_h_s2_ptr, &mut self.online_dueling, iqn_budget_per_branch,
                        ).map_err(|e| anyhow::anyhow!("IQN trunk gradient branch {branch_idx}: {e}"))?;
                        iqn_ok = true;
                    }
                }
            }
        }

        // ── Post-join IQN processing (main stream) ──
        // SP6 Pearl 5: apply_iqn_trunk_gradient is now called inline per branch
        // (4 × budget/4) inside the parallel/sequential loops above. The snapshot
        // was moved to before the fork. Only EMA + IQR run here post-join.
        if iqn_ok {
            if let Some(ref mut iqn) = self.gpu_iqn {
                // SP6 Pearl 5: apply_iqn_trunk_gradient was moved into the per-branch
                // loop above (4× iqn_budget_per_branch = iqn_budget/4). The Pearl 2
                // single-call form is removed — the per-branch form supersedes it
                // and supplies the iqn_budget/4 normalization that preserves total
                // gradient magnitude. Only EMA + IQR run here post-join.
                let tau = self.trainer.read_isv_signal_at(TAU_EFF_INDEX);
                iqn.target_ema_update(tau)
                    .map_err(|e| anyhow::anyhow!("IQN EMA update: {e}"))?;
            }
        }

        // Task 2.0 — IQN reduction. Runs unconditionally so the pinned slot
        // stays populated even on steps where `iqn_ok=false`. When IQN was
        // skipped the snapshot and current `grad_buf` are identical over the
        // branch 0+1 range → delta norm = 0.0 (honest semantic).
        self.trainer.grad_decomp_launch_iqn()
            .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp IQN: {e}"))?;

        // IQN→PER loss: cast f32 per_sample_loss → f32 td_errors_buf.
        // IQN is now f32 native; PER system expects f32 td_errors.
        if let Some(ref iqn) = self.gpu_iqn {
            let src = iqn.per_sample_loss();
            self.trainer.cast_f32_to_td_errors(src)
                .map_err(|e| anyhow::anyhow!("IQN→PER f32→f32 cast: {e}"))?;
        }


        // CQL gradient + SAXPY (no per-component clip — Adam handles safety).
        // B4/G5: cql_budget scales the SAXPY contribution (0.10×(1−regime)×health).
        //
        // F5/D2: Q-gap barrier gradient weight — tuned + temporal.
        // Base weight raised from 0.05 → 0.20 after local smoke tests showed the
        // 0.05 barrier was too weak to escape the Q-uniform attractor. Also
        // amplified by meta-Q collapse prediction: when meta_q_pred > 0.5 we
        // scale up linearly (up to 2×) so the barrier gets stronger BEFORE
        // collapse fully lands, not after.
        let meta_q_pred = self.trainer.last_meta_q_pred();
        let temporal_amp = 1.0_f32 + (meta_q_pred - 0.5).max(0.0) * 2.0;
        let f5_barrier_weight = 0.20_f32 * temporal_amp;

        // SP7 Path A (2026-05-03): the historical `grad_decomp_launch_cql`
        // (paired with this snapshot) measured `‖grad_buf − snapshot_cql‖`
        // over the CQL gradient path — structurally always 0 because
        // `apply_cql_gradient` writes into `cql_grad_scratch` (a separate
        // buffer), never touching `grad_buf` until `apply_cql_saxpy` fires.
        // The launch is replaced below by `launch_cql_raw_norm`, which reads
        // `cql_grad_scratch` directly and writes a real raw-CQL norm to
        // `grad_decomp_result_pinned[3..6]`. The SP7 loss-balance controller
        // reads slot 3 as a budget-independent CQL reference signal (was
        // deadlocking on the budget-scaled `cql_sx` slot at offset 6).
        //
        // The snapshot itself is now functionally dead (nothing reads its
        // output) but is left in place to preserve the `grad_snapshot_cql`
        // buffer + `grad_decomp_snapshot_cql()` call site semantics —
        // removing the buffer touches the constructor, struct definition,
        // and 9 sibling `grad_snapshot_*` references; that cleanup belongs
        // in a follow-up commit per `feedback_no_partial_refactor` (this
        // commit is focused on Path A wire-up).
        self.trainer.grad_decomp_snapshot_cql()
            .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql: {e}"))?;

        if self.trainer.has_cql() {
            match self.trainer.apply_cql_gradient(f5_barrier_weight) {
                Ok(true) => {
                    // SP7 Path A (2026-05-03) — populate `cql` slot at offset 3
                    // with the raw CQL gradient norm read directly from
                    // `cql_grad_scratch`. Must run AFTER `apply_cql_gradient`
                    // populated the scratch and BEFORE `apply_cql_saxpy`
                    // consumes it. Replaces the unwired
                    // `grad_decomp_launch_cql` snapshot path (which always
                    // reported 0 because the snapshot was on `grad_buf` and
                    // CQL writes to `cql_grad_scratch`).
                    self.trainer.launch_cql_raw_norm()
                        .map_err(|e| anyhow::anyhow!("SP7 Path A cql_raw_norm: {e}"))?;
                    // Task 2.0 extension — isolate `apply_cql_saxpy`. The
                    // preceding `apply_cql_gradient` only writes into
                    // `cql_grad_scratch`, not `grad_buf`, so the cql_sx
                    // delta captures exactly the budget-scaled SAXPY into
                    // `grad_buf`.
                    self.trainer.grad_decomp_snapshot_cql_sx()
                        .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql_sx: {e}"))?;
                    // SP6 Pearl 2 + SP7 (2026-05-03): trunk/value SAXPY first
                    // (mean budget), then per-branch corrections. Both kernels
                    // read `alpha` from `lb_budget_effective_buf` (populated by
                    // the dispatch kernel earlier in `compute_adaptive_budgets`)
                    // — capture-safe, see audit Fix 33.
                    let cql_trunk_dev_ptr = self.trainer.lb_budget_cql_trunk_dev_ptr();
                    self.trainer.apply_cql_saxpy(cql_trunk_dev_ptr)
                        .map_err(|e| anyhow::anyhow!("CQL SAXPY trunk: {e}"))?;
                    for b in 0..4_usize {
                        let correction_dev_ptr = self.trainer.lb_budget_cql_correction_dev_ptr(b);
                        self.trainer.apply_cql_saxpy_branch(b, correction_dev_ptr)
                            .map_err(|e| anyhow::anyhow!("CQL SAXPY branch[{b}]: {e}"))?;
                    }
                    self.trainer.grad_decomp_launch_cql_sx()
                        .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp cql_sx: {e}"))?;
                }
                Ok(false) => {
                    // Keep cql_sx slot populated on skip path — delta = 0.
                    // SP7 Path A: also keep `cql` slot at offset 3 populated.
                    // `cql_grad_scratch` was zeroed at the top of
                    // `apply_cql_gradient` (`cuMemsetD32Async`) before the
                    // early-return, so the raw-norm kernel reads all-zeros
                    // and writes 0.0 to the slot — honest semantic.
                    self.trainer.launch_cql_raw_norm()
                        .map_err(|e| anyhow::anyhow!("SP7 Path A cql_raw_norm (skip): {e}"))?;
                    self.trainer.grad_decomp_snapshot_cql_sx()
                        .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql_sx (skip): {e}"))?;
                    self.trainer.grad_decomp_launch_cql_sx()
                        .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp cql_sx (skip): {e}"))?;
                }
                Err(e) => {
                    tracing::warn!("CQL gradient failed (non-fatal): {e}");
                    // Keep both slots populated on error path. The scratch
                    // may be in an indeterminate state after the partial
                    // backward; the raw-norm kernel still produces a
                    // numerically-finite L2 norm, which is what SP7 needs.
                    self.trainer.launch_cql_raw_norm()
                        .map_err(|e| anyhow::anyhow!("SP7 Path A cql_raw_norm (err): {e}"))?;
                    self.trainer.grad_decomp_snapshot_cql_sx()
                        .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql_sx (err): {e}"))?;
                    self.trainer.grad_decomp_launch_cql_sx()
                        .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp cql_sx (err): {e}"))?;
                }
            }
        } else {
            // Keep cql_sx slot populated on CQL-disabled path — delta = 0.
            // SP7 Path A: scratch is left at its prior contents (no
            // `apply_cql_gradient` call to zero it), so the raw norm reads
            // whatever survived from the last step where CQL fired. In the
            // CQL-permanently-disabled regime the buffer is still zero from
            // construction, so the kernel writes 0.0. In a transient-disable
            // regime the slot stays at the most-recent active value, which
            // is the appropriate "carry forward" semantic for a controller
            // reading a gradient-magnitude reference.
            self.trainer.launch_cql_raw_norm()
                .map_err(|e| anyhow::anyhow!("SP7 Path A cql_raw_norm (no-cql): {e}"))?;
            self.trainer.grad_decomp_snapshot_cql_sx()
                .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql_sx (no-cql): {e}"))?;
            self.trainer.grad_decomp_launch_cql_sx()
                .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp cql_sx (no-cql): {e}"))?;
        }

        // D1/N1: Temporal self-distillation — per-step kernel pull toward the
        // best historical snapshot. Fully GPU-native: alpha is computed
        // in-kernel from ISV[LEARNING_HEALTH_INDEX] (pinned device-mapped),
        // and `distill_best_buf` is a stable device buffer whose contents
        // are DtoD-updated only at epoch boundaries (when a new best lands
        // in the ring). Initial state is a copy of params so `(params − best)`
        // is zero — the kernel is a strict no-op until a real snapshot
        // populates the buffer. Must run here (aux phase) — prior
        // epoch-boundary placement wrote to grad_buf AFTER the last step,
        // and the next step's graph_forward memset zeroed it before Adam
        // could consume it.
        //
        // Task 2.0 extension — isolate `apply_distillation_gradient` delta.
        self.trainer.grad_decomp_snapshot_distill()
            .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_distill: {e}"))?;
        self.trainer.apply_distillation_gradient()
            .map_err(|e| anyhow::anyhow!("distill SAXPY: {e}"))?;
        self.trainer.grad_decomp_launch_distill()
            .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp distill: {e}"))?;


        // Recursive confidence backward: MSE grad into trunk + conf weight gradients.
        // Must run before Adam (which reads grad_buf for the parameter update).
        //
        // Task 2.0 extension — isolate `launch_recursive_confidence_backward`
        // delta.
        self.trainer.grad_decomp_snapshot_rec()
            .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_rec: {e}"))?;
        self.trainer.launch_recursive_confidence_backward(self.batch_size)
            .map_err(|e| anyhow::anyhow!("Recursive confidence backward: {e}"))?;
        self.trainer.grad_decomp_launch_rec()
            .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp rec: {e}"))?;

        // G12: Predictive coding auxiliary loss + backward.
        // Self-supervised temporal smoothness on the enriched trunk h_s2 —
        // adds 2*lambda*(h[i] - h[i+1])/SH2 regularization gradient into
        // bw_d_h_s2. Runs HERE so the gradient lands in bw_d_h_s2 before
        // the trunk backward GEMMs consume it (and before recursive
        // confidence's own trunk gradient — both accumulate via plain +=).
        //
        // Task 2.0 extension — isolate `compute_predictive_coding_loss`
        // delta.
        self.trainer.grad_decomp_snapshot_pred()
            .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_pred: {e}"))?;
        self.trainer.compute_predictive_coding_loss(self.batch_size)
            .map_err(|e| anyhow::anyhow!("Predictive coding (G12): {e}"))?;
        self.trainer.grad_decomp_launch_pred()
            .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp pred: {e}"))?;

        // G6 + G10 measurement (commit 61ab27ff3) found both penalties are
        // dominated by existing mechanisms:
        //   G10 temporal_consistency: 0.000000–0.000002 across 14 epochs.
        //     Cosine similarity between consecutive h_s2 samples almost never
        //     exceeds the 0.95 threshold; spectral_norm already enforces a
        //     global Lipschitz constraint that subsumes per-pair Lipschitz.
        //   G6 branch_independence: ~1e-4 (orders of magnitude smaller than G12).
        //     NoisyNets injects different parameter noise per layer → the 4
        //     advantage branches are already ~99.99% diverse without explicit
        //     cosine-similarity penalty.
        //
        // Per V7-gem methodology: when the existing mechanisms cover the signal,
        // wiring backward adds gradient noise without value. Both calls removed
        // here; the kernels and Rust glue are deleted in a follow-up cleanup.

        // Task 2.0 — snapshot BEFORE ensemble diversity step. Runs
        // unconditionally even when `ensemble_extra_heads.is_empty()` so the
        // post-ens reduction below can run unconditionally (0.0 delta on
        // steps without ensemble diversity — honest semantic).
        self.trainer.grad_decomp_snapshot_ens()
            .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_ens: {e}"))?;

        // Ensemble diversity: heads 1..K-1 forward + KL gradient + trunk backward.
        // Runs inside aux_child graph — gradients land in grad_buf before Adam.
        if !self.ensemble_extra_heads.is_empty() {
            self.run_ensemble_step()
                .map_err(|e| anyhow::anyhow!("Ensemble diversity: {e}"))?;
        }

        // Task 2.0 — ensemble reduction. Delta = grad_buf − snapshot over branch 0+1.
        self.trainer.grad_decomp_launch_ens()
            .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp ens: {e}"))?;

        // Regime-adaptive PER scaling.
        self.trainer.regime_scale_td_errors()
            .map_err(|e| anyhow::anyhow!("Regime PER scaling: {e}"))?;

        Ok(())
    }

    /// Submit counter increments + stochastic depth mask (captured as counters_child).
    ///
    /// Replaces host-side `adam_step_async()` + `pre_replay_state_update()` with a
    /// single GPU kernel that increments all 8 Adam step counters and computes
    /// cosine-annealed tau on-device.
    fn submit_counters_ops(&mut self) -> Result<()> {
        // GPU-side counter increments (replaces adam_step_async + pre_replay_state_update)
        let iql_t = self.gpu_iql.t_dev_ptr();
        let iql_low_t = self.gpu_iql_low.t_dev_ptr();
        let iqn_t = self.gpu_iqn.as_ref().map(|i| i.t_dev_ptr()).unwrap_or(0);
        let attn_t = self.gpu_attention.as_ref().map(|a| a.t_dev_ptr()).unwrap_or(0);
        let rng_step = self.per_rng_step_dev_ptr;
        self.trainer.submit_counter_increments(iql_t, iql_low_t, iqn_t, attn_t, rng_step)
            .map_err(|e| anyhow::anyhow!("counters: {e}"))?;
        // Stochastic depth mask (already a GPU kernel)
        self.trainer.update_stochastic_depth_mask()
            .map_err(|e| anyhow::anyhow!("stochastic depth mask: {e}"))
    }

    /// SP4 Layer A Task A7 fix-up + fix-up #2: build the aux-trainer buffer
    /// descriptor `SP4AuxBuffers` that
    /// `GpuDqnTrainer::launch_sp4_param_group_oracles_all_groups` consumes
    /// for groups 3-7 (IQN, IQL hi/lo, Attn, Curiosity).
    ///
    /// `gpu_iqn` and `gpu_attention` are `Option<_>` purely as init-failure
    /// fallback (graceful degrade) and are otherwise always present in
    /// production; if either is `None`, this helper emits an empty
    /// descriptor for that group so the launcher's `total_count == 0`
    /// short-circuit silently skips the kernel launch.
    ///
    /// **Curiosity (group 7)** lives outside `FusedTrainingCtx` —
    /// `CuriosityWeightSet` is owned by the experience collector and the
    /// matching `GpuCuriosityTrainer` (grad + Adam state) is also there.
    /// Layer B's training-loop caller threads them in via the
    /// `curiosity_weights` and `curiosity_trainer` parameters; both must
    /// be `Some` together (caller responsibility — they live on the same
    /// collector) and the descriptor describes the union of all four
    /// `[w1, b1, w2, b2]` sub-buffers as one logical group. Passing
    /// `None` for either yields an empty curiosity descriptor and the
    /// launcher silently skips group 7.
    #[allow(dead_code)] // Layer B will consume this; lint suppressed until then.
    pub(crate) fn build_sp4_aux_buffers(
        &self,
        curiosity_weights: Option<&crate::cuda_pipeline::gpu_weights::CuriosityWeightSet>,
        curiosity_trainer: Option<&crate::cuda_pipeline::gpu_curiosity_trainer::GpuCuriosityTrainer>,
    ) -> crate::cuda_pipeline::gpu_dqn_trainer::SP4AuxBuffers
    {
        use crate::cuda_pipeline::gpu_dqn_trainer::{SP4AuxBuffers, Sp4ParamGroupBufs, Sp4SubBuffer};

        let iqn = if let Some(iqn) = self.gpu_iqn.as_ref() {
            Sp4ParamGroupBufs::single(
                iqn.online_params_ptr(),
                iqn.online_grad_ptr(),
                iqn.adam_m_ptr(),
                iqn.adam_v_ptr(),
                iqn.online_params_len(),
            )
        } else {
            Sp4ParamGroupBufs::empty()
        };

        let iql_high = Sp4ParamGroupBufs::single(
            self.gpu_iql.params_ptr(),
            self.gpu_iql.grads_ptr(),
            self.gpu_iql.adam_m_ptr(),
            self.gpu_iql.adam_v_ptr(),
            self.gpu_iql.params_len(),
        );

        let iql_low = Sp4ParamGroupBufs::single(
            self.gpu_iql_low.params_ptr(),
            self.gpu_iql_low.grads_ptr(),
            self.gpu_iql_low.adam_m_ptr(),
            self.gpu_iql_low.adam_v_ptr(),
            self.gpu_iql_low.params_len(),
        );

        let attn = if let Some(attn) = self.gpu_attention.as_ref() {
            Sp4ParamGroupBufs::single(
                attn.params_ptr(),
                attn.grads_ptr(),
                attn.adam_m_ptr(),
                attn.adam_v_ptr(),
                attn.params_len(),
            )
        } else {
            Sp4ParamGroupBufs::empty()
        };

        // Curiosity multi-sub-buffer descriptor. Both halves of state
        // (params from CuriosityWeightSet; grad/Adam from GpuCuriosityTrainer)
        // must be supplied — Layer B's caller fetches both from the same
        // GpuExperienceCollector. If either is missing (curiosity disabled
        // at collector init), emit an empty descriptor.
        let curiosity = match (curiosity_weights, curiosity_trainer) {
            (Some(w), Some(t)) => Sp4ParamGroupBufs {
                sub_buffers: vec![
                    Sp4SubBuffer {
                        params_ptr: w.w1_ptr(),
                        grads_ptr:  t.grad_w1_ptr(),
                        adam_m_ptr: t.adam_m_w1_ptr(),
                        adam_v_ptr: t.adam_v_w1_ptr(),
                        count:      w.w1_len(),
                    },
                    Sp4SubBuffer {
                        params_ptr: w.b1_ptr(),
                        grads_ptr:  t.grad_b1_ptr(),
                        adam_m_ptr: t.adam_m_b1_ptr(),
                        adam_v_ptr: t.adam_v_b1_ptr(),
                        count:      w.b1_len(),
                    },
                    Sp4SubBuffer {
                        params_ptr: w.w2_ptr(),
                        grads_ptr:  t.grad_w2_ptr(),
                        adam_m_ptr: t.adam_m_w2_ptr(),
                        adam_v_ptr: t.adam_v_w2_ptr(),
                        count:      w.w2_len(),
                    },
                    Sp4SubBuffer {
                        params_ptr: w.b2_ptr(),
                        grads_ptr:  t.grad_b2_ptr(),
                        adam_m_ptr: t.adam_m_b2_ptr(),
                        adam_v_ptr: t.adam_v_b2_ptr(),
                        count:      w.b2_len(),
                    },
                ],
            },
            _ => Sp4ParamGroupBufs::empty(),
        };

        SP4AuxBuffers { iqn, iql_high, iql_low, attn, curiosity }
    }

    /// Submit IQL advantage-weighted TD error modulation (captured as iql_modulate_child).
    ///
    /// Uses the IQL trainer's own CUfunction from its own cubin. The `write_pos`
    /// and `capacity` scalars are baked at capture time — `write_pos` staleness is
    /// acceptable for a soft modulation signal (not correctness-critical).
    fn submit_iql_modulate_ops(
        &mut self,
        agent: &mut DQNAgentType,
        gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch,
    ) -> Result<()> {
        let write_pos = agent.replay_write_cursor() as i32;
        let capacity = agent.replay_capacity() as i32;
        self.gpu_iql.modulate_td_errors(
            self.trainer.td_errors_buf_mut(),
            gpu_batch.indices_ptr,
            write_pos,
            capacity,
        ).map_err(|e| anyhow::anyhow!("IQL modulate_td_errors: {e}"))
    }

    /// Submit PER priority update (captured as per_priority_child).
    ///
    /// Routes through the replay buffer's seg_tree_update kernel which propagates
    /// the segment tree from leaf to root entirely on GPU.
    fn submit_per_priority_ops(&mut self, agent: &mut DQNAgentType) -> Result<()> {
        agent.update_priorities_from_td(
            self.trainer.td_errors_buf(),
            self.batch_size,
            &self.stream,
        ).map_err(|e| anyhow::anyhow!("PER priority update: {e}"))
    }

    /// Capture a single child sub-graph using cudarc stream capture.
    ///
    /// Takes ownership of cudarc's CudaGraph handles directly via `mem::forget`.
    /// We do NOT use `cuGraphClone` — cloning corrupts cublasLtMatmul HMMA tile
    /// state on Hopper, causing graph replay to hang when temporal ops are present.
    ///
    /// Uses `disable_event_tracking` around capture to prevent cudarc corruption.
    fn capture_child_graph(
        &mut self,
        label: &str,
        ops: impl FnOnce(&mut Self) -> Result<()>,
    ) -> Result<ChildGraph> {
        self.trainer.sync_all_streams()
            .map_err(|e| anyhow::anyhow!("{label} sync: {e}"))?;
        let _ = self.stream.context().check_err();
        unsafe { self.stream.context().disable_event_tracking(); }

        let begin_result = self.stream.begin_capture(
            cuda_sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED,
        );
        if let Err(e) = begin_result {
            unsafe { self.stream.context().enable_event_tracking(); }
            let _ = self.stream.context().check_err();
            anyhow::bail!("{label} begin_capture: {e}");
        }

        let result = ops(self);

        let graph_result = self.stream.end_capture(
            cuda_sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
        );

        unsafe { self.stream.context().enable_event_tracking(); }
        let _ = self.stream.context().check_err();
        result?;

        let cudarc_graph = graph_result
            .map_err(|e| anyhow::anyhow!("{label} end_capture: {e}"))?
            .ok_or_else(|| anyhow::anyhow!("{label} capture returned None"))?;

        // Take ownership of the raw handles from cudarc's CudaGraph.
        // cudarc already instantiated the exec with AUTO_FREE_ON_LAUNCH — use
        // it directly. mem::forget prevents CudaGraph::drop from destroying them.
        let graph = cudarc_graph.cu_graph;
        let exec = cudarc_graph.cu_graph_exec;
        std::mem::forget(cudarc_graph);

        Ok(ChildGraph { graph, exec })
    }

    /// Capture all 12 children and compose into single parent graph.
    ///
    /// Called after ungraphed step 0 succeeds. Each child is captured
    /// independently, then composed into the parent via `cuGraphAddChildGraphNode`.
    ///
    /// Children (sequential):
    ///   per_sample -> counters -> spectral -> forward -> ddqn -> aux -> post_aux ->
    ///   adam_grad -> adam -> maintenance -> iql_modulate -> per_priority
    fn capture_training_graph(
        &mut self,
        agent: &mut DQNAgentType,
        gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch,
    ) -> Result<()> {
        // Capture PER sampling child (prefix scan + binary search + gather).
        let per_sample = self.capture_child_graph("per_sample", |s| {
            agent.primary_dqn_mut().memory.gpu.sample_proportional(s.batch_size)
                .map_err(|e| anyhow::anyhow!("per_sample capture: {e}"))
                .map(|_| ())
        })?;

        // Capture counters child (GPU-side step counters + stochastic depth mask).
        let counters = self.capture_child_graph("counters", |s| {
            s.submit_counters_ops()
        })?;

        // Capture spectral child.
        let spectral = self.capture_child_graph("spectral", |s| {
            s.trainer.apply_spectral_norm(&s.online_dueling, &s.online_branching)
                .map_err(|e| anyhow::anyhow!("{e}"))
        })?;

        // Capture forward child (cuBLAS trunk + ISV + loss + backward).
        let forward = self.capture_child_graph("forward", |s| {
            s.trainer.submit_forward_ops_main()
                .map_err(|e| anyhow::anyhow!("{e}"))
        })?;

        // Capture DDQN child (Pass 3).
        let ddqn = self.capture_child_graph("ddqn", |s| {
            s.trainer.submit_forward_ops_ddqn()
                .map_err(|e| anyhow::anyhow!("{e}"))
        })?;

        // Capture aux child (HER + EMA + IQL + IQN + attention + CQL).
        let aux = self.capture_child_graph("aux", |s| {
            s.submit_aux_ops(agent, gpu_batch)
        })?;

        // Capture post_aux child (selectivity + denoise + Q-stats + risk_sgd
        // + SP1 Phase B post-backward NaN checks). The post-backward NaN
        // checks join the same captured child as the existing pre/post
        // forward checks since both are GPU-side reduce kernels writing
        // into `nan_flags_buf` — keeping them adjacent in the graph
        // minimises launch overhead and makes the slot-name table
        // (in `training_loop.rs`) align with one capture point per group.
        let post_aux = self.capture_child_graph("post_aux", |s| {
            s.trainer.submit_post_aux_ops(s.batch_size)
                .map_err(|e| anyhow::anyhow!("{e}"))?;
            // SP2 (Task A4): IQN ptrs no longer per-step args —
            // `populate_nan_check_meta` (called once at construction)
            // baked Option<u64> nullity into the fused kernel's
            // metadata buffer entries.
            s.trainer.run_nan_checks_post_backward()
                .map_err(|e| anyhow::anyhow!("{e}"))
        })?;

        // adam_grad: mamba2 bwd + ofi_embed bwd + pruning + grad_norm
        // adam_update: Adam + unflatten + ISV
        let adam_grad = self.capture_child_graph("adam_grad", |s| {
            s.mamba2_backward(s.batch_size)?;
            s.step_mamba2_adam()?;
            s.launch_ofi_embed_backward(s.batch_size)?;
            s.step_ofi_embed_adam()?;
            s.trainer.apply_pruning_mask()
                .map_err(|e| anyhow::anyhow!("{e}"))?;
            // Adaptive per-branch gradient-norm balancer — captured inside
            // adam_grad_child so every graph replay rebalances before the
            // global grad_norm reduction.  See `launch_branch_grad_balance`.
            s.trainer.launch_branch_grad_balance()
                .map_err(|e| anyhow::anyhow!("branch_grad_balance: {e}"))?;
            s.trainer.compute_grad_norm_for_adam()
                .map_err(|e| anyhow::anyhow!("{e}"))
        })?;
        let adam = self.capture_child_graph("adam_update", |s| {
            s.trainer.submit_adam_ops(&s.online_dueling, &s.online_branching)
                .map_err(|e| anyhow::anyhow!("{e}"))?;
            // Plan C Phase 2 follow-up F (2026-04-29): see ungraphed step 0
            // path (~line 1467) for the full rationale. Capturing the bin_
            // means_reduce producers in adam_update_child (immediately before
            // `update_isv_signals`) makes the per-step graph self-consistent:
            // every replay reads `q_out_buf` filled by this step's
            // forward_child, writes mapped-pinned scratch, and the captured
            // update_isv_signals immediately consumes it. Without this, the
            // captured `update_isv_signals` reads stale once-per-epoch scratch
            // and ISV[16,21] stay near the cold-start value, defeating the
            // Q-drift kill criterion's adaptive floor.
            s.trainer.launch_q_mag_bin_means_reduce(s.batch_size)
                .map_err(|e| anyhow::anyhow!("{e}"))?;
            s.trainer.launch_q_dir_bin_means_reduce(s.batch_size)
                .map_err(|e| anyhow::anyhow!("{e}"))?;
            s.trainer.update_isv_signals()
                .map_err(|e| anyhow::anyhow!("{e}"))
        })?;

        // Capture maintenance child (causal intervention + vaccine).
        let maintenance = self.capture_child_graph("maintenance", |s| {
            s.trainer.submit_maintenance_ops()
                .map_err(|e| anyhow::anyhow!("{e}"))
        })?;

        // Capture IQL modulate child (advantage-weighted PER modulation).
        let iql_modulate = self.capture_child_graph("iql_modulate", |s| {
            s.submit_iql_modulate_ops(agent, gpu_batch)
        })?;

        // Capture PER priority child (seg_tree update).
        let per_priority = self.capture_child_graph("per_priority", |s| {
            s.submit_per_priority_ops(agent)
        })?;

        // Store children.
        self.per_sample_child = Some(per_sample);
        self.counters_child = Some(counters);
        self.spectral_child = Some(spectral);
        self.forward_child = Some(forward);
        self.ddqn_child = Some(ddqn);
        self.aux_child = Some(aux);
        self.post_aux_child = Some(post_aux);
        self.adam_grad_child = Some(adam_grad);
        self.adam_child = Some(adam);
        self.maintenance_child = Some(maintenance);
        self.iql_modulate_child = Some(iql_modulate);
        self.per_priority_child = Some(per_priority);

        // Compose parent graph from all children.
        self.compose_parent_graph()?;

        // Create eval-only forward exec for deterministic Q-value replay.
        self.create_eval_forward_exec()?;

        info!("13 child graphs captured + parent composed (per_sample + counters + spectral + forward + ddqn + aux + post_aux + adam_grad + adam + maintenance + iql_modulate + per_priority)");
        Ok(())
    }

    /// Compose parent graph from all child sub-graphs using `cuGraphAddChildGraphNode`.
    ///
    /// Each child is added as a sequential dependency:
    ///   per_sample -> counters -> spectral -> forward -> ddqn -> aux -> post_aux ->
    ///   adam_grad -> adam -> maintenance -> iql_modulate -> per_priority
    /// The parent is instantiated once with `cuGraphInstantiateWithFlags`.
    fn compose_parent_graph(&mut self) -> Result<()> {
        // Create empty parent graph.
        let mut parent: cuda_sys::CUgraph = std::ptr::null_mut();
        let r = unsafe { cuda_sys::cuGraphCreate(&mut parent, 0) };
        if r != cuda_sys::cudaError_enum::CUDA_SUCCESS {
            anyhow::bail!("cuGraphCreate failed: {r:?}");
        }

        // Collect all children in sequential order.
        // Optional children are skipped if None.
        let mut children_vec: Vec<(&str, &ChildGraph)> = Vec::with_capacity(13);
        if let Some(ref c) = self.per_sample_child {
            children_vec.push(("per_sample", c));
        }
        if let Some(ref c) = self.counters_child {
            children_vec.push(("counters", c));
        }
        children_vec.push(("spectral", self.spectral_child.as_ref().unwrap()));
        children_vec.push(("forward",  self.forward_child.as_ref().unwrap()));
        children_vec.push(("ddqn",     self.ddqn_child.as_ref().unwrap()));
        children_vec.push(("aux",      self.aux_child.as_ref().unwrap()));
        if let Some(ref c) = self.post_aux_child {
            children_vec.push(("post_aux", c));
        }
        if let Some(ref c) = self.adam_grad_child {
            children_vec.push(("adam_grad", c));
        }
        children_vec.push(("adam", self.adam_child.as_ref().unwrap()));
        if let Some(ref c) = self.maintenance_child {
            children_vec.push(("maintenance", c));
        }
        if let Some(ref c) = self.iql_modulate_child {
            children_vec.push(("iql_modulate", c));
        }
        if let Some(ref c) = self.per_priority_child {
            children_vec.push(("per_priority", c));
        }
        let children = children_vec;

        let mut prev_node: cuda_sys::CUgraphNode = std::ptr::null_mut();
        for (label, child) in &children {
            let mut node: cuda_sys::CUgraphNode = std::ptr::null_mut();
            let (deps, num_deps) = if prev_node.is_null() {
                (std::ptr::null(), 0)
            } else {
                (&prev_node as *const cuda_sys::CUgraphNode, 1)
            };
            let r = unsafe {
                cuda_sys::cuGraphAddChildGraphNode(
                    &mut node,
                    parent,
                    deps,
                    num_deps,
                    child.graph,
                )
            };
            if r != cuda_sys::cudaError_enum::CUDA_SUCCESS {
                unsafe { cuda_sys::cuGraphDestroy(parent); }
                anyhow::bail!("cuGraphAddChildGraphNode({label}) failed: {r:?}");
            }
            prev_node = node;
        }

        // Instantiate parent.
        let mut graph_exec: cuda_sys::CUgraphExec = std::ptr::null_mut();
        let r = unsafe { cuda_sys::cuGraphInstantiateWithFlags(&mut graph_exec, parent, 0) };
        if r != cuda_sys::cudaError_enum::CUDA_SUCCESS {
            unsafe { cuda_sys::cuGraphDestroy(parent); }
            anyhow::bail!("cuGraphInstantiateWithFlags failed: {r:?}");
        }

        self.parent_graph = Some(ParentGraph { exec: graph_exec, graph: parent });
        self.trainer.mark_graphs_captured();

        info!(
            num_children = children.len(),
            "Parent graph composed: single cuGraphLaunch per step"
        );
        Ok(())
    }

    /// Create eval-only forward graph exec for deterministic Q-value replay.
    ///
    /// Create eval-only forward graph exec for deterministic Q-value replay.
    ///
    /// Instantiates directly from the forward child's original CUgraph — no
    /// cuGraphClone. Multiple execs from the same graph is safe (NVIDIA docs).
    fn create_eval_forward_exec(&mut self) -> Result<()> {
        let forward_child = self.forward_child.as_ref()
            .ok_or_else(|| anyhow::anyhow!("forward_child not captured"))?;

        let eval_exec = ChildGraph::instantiate_extra_exec(forward_child.graph)?;

        if let Some(old) = self.eval_forward_exec.take() {
            unsafe { cuda_sys::cuGraphExecDestroy(old); }
        }
        self.eval_forward_exec = Some(eval_exec);

        // Inject into GpuDqnTrainer for replay_forward_for_q_values().
        let trainer_eval_exec = ChildGraph::instantiate_extra_exec(forward_child.graph)?;
        self.trainer.set_eval_forward_exec(trainer_eval_exec);

        Ok(())
    }

    /// Update sub-trainer host-side state before parent graph replay.
    ///
    /// Sub-trainers (IQN, attention) still use graphed HtoD from stable struct
    /// field addresses. Main GpuDqnTrainer uses pinned device-mapped memory
    /// (zero-copy, no HtoD).
    fn pre_replay_state_update(&mut self, _agent: &mut DQNAgentType) {
        // Increment sub-trainer adam_steps so graphed HtoD copies updated values.
        if let Some(ref mut attn) = self.gpu_attention {
            attn.increment_adam_step();
        }
        self.gpu_iql.increment_adam_step();
        self.gpu_iql_low.increment_adam_step();
        if let Some(ref mut iqn) = self.gpu_iqn {
            iqn.increment_adam_step();
            // Update IQN tau from ISV[TAU_EFF_INDEX] (GPU-computed at epoch boundary).
            let tau = self.trainer.read_isv_signal_at(TAU_EFF_INDEX);
            iqn.set_tau_host(tau);
        }
        // Update DQN trainer tau from ISV[TAU_EFF_INDEX].
        {
            let tau = self.trainer.read_isv_signal_at(TAU_EFF_INDEX);
            self.trainer.set_tau_value(tau);
        }
    }

    /// Run one ensemble diversity step (outside CUDA Graph).
    ///
    /// Heads 1..K-1 run their value/advantage head forward on the trunk activations
    /// (save_h_s2) that head 0 already computed inside the CUDA Graph. The KL
    /// divergence across all K head pairs is accumulated as a diversity loss scalar.
    ///
    /// This method is called after the CUDA Graph replay so save_h_s2 is valid.
    /// EventTrackingGuard is used to prevent stale CUDA event errors.
    ///
    /// The diversity gradient flows back through head 0's value head into the
    /// shared trunk via cuBLAS backward + SAXPY (same pattern as IQN trunk gradient).
    /// graph_adam then sees combined C51 + IQN + ensemble diversity gradients.
    fn run_ensemble_step(&mut self) -> Result<()> {
        
        use cudarc::driver::{LaunchConfig, PushKernelArg};

        let k = self.ensemble_extra_heads.len() + 1; // total heads including head 0
        let b = self.batch_size;
        let na = self.trainer.config().num_atoms;

        // No cuStreamSynchronize needed — all ops are on the same stream.
        // CUDA guarantees in-order execution on a single stream.
        // Event tracking disabled for buffer pointer extractions (CUDA Graph compat).
        use crate::cuda_pipeline::gpu_dqn_trainer::EventTrackingGuard;
        let _evt_guard = EventTrackingGuard::new(self.stream.context());

        let logits_buf = match self.ensemble_logits_buf.as_ref() {
            Some(b) => b,
            None => return Ok(()),
        };
        // ── 1. Copy head 0 value logits (on_v_logits_buf, f32) → logits_buf[0..B*na] ──
        // Head 0's logits are already computed by the CUDA Graph.
        {
            let head0_ptr = self.trainer.on_v_logits_buf().raw_ptr();
            let dst_ptr = logits_buf.raw_ptr();
            let n_bytes = b * na * std::mem::size_of::<f32>();
            // Graph-safe copy: run_ensemble_step is captured in aux child graph.
            self.trainer.graph_safe_copy_f32(dst_ptr, head0_ptr, n_bytes, "ensemble head0 logits")
                .map_err(|e| anyhow::anyhow!("{e}"))?;
        }

        // ── 2. Run heads 1..K-1 forward (value head only on save_h_s2) ──
        // Each extra head uses its own value weights (w_v1, b_v1, w_v2, b_v2)
        // on the shared trunk activation save_h_s2 (no trunk re-forward needed).
        // We compute: h_v_k = ReLU(save_h_s2 @ W_v1_k^T + b_v1_k)
        //             logits_k = h_v_k @ W_v2_k^T + b_v2_k
        // Using raw cuBLAS-style SGEMM (reuse saxpy + bias patterns).
        // For simplicity, use a minimal CUDA SGEMM via cuBLAS forward (shared with trainer).
        // Since we don't have a direct cuBLAS handle here, use the saxpy kernel
        // to accumulate bias and the gemm scratch from trainer.
        //
        // PRACTICAL SIMPLIFICATION: Copy save_h_s2 into the trainer's scratch buffers,
        // then launch head-k value layers via cuBLAS gemm reuse.
        // The head-k weights are in extra_heads[k-1].0 (DuelingWeightSet: w_v1, b_v1, w_v2, b_v2).
        //
        for (head_idx, (head_dueling, _head_branching)) in
            self.ensemble_extra_heads.iter().enumerate()
        {
            let k_idx = head_idx + 1; // head 0 is already copied

            // Full cuBLAS value head forward for each ensemble head.
            // h_s2 → W_v1_k → ReLU → W_v2_k → v_logits_k
            // Uses shared trunk activations (save_h_s2) with per-head value weights.
            let w_v1_ptr = head_dueling.w_v1;
            let b_v1_ptr = head_dueling.b_v1;
            let w_v2_ptr = head_dueling.w_v2;
            let b_v2_ptr = head_dueling.b_v2;
            let dst_k_ptr = logits_buf.raw_ptr()
                + (k_idx * b * na * std::mem::size_of::<f32>()) as u64;

            self.trainer.forward_value_head_for_ensemble(
                w_v1_ptr, b_v1_ptr, w_v2_ptr, b_v2_ptr,
                dst_k_ptr, b,
            ).map_err(|e| anyhow::anyhow!("Ensemble head{k_idx} value forward: {e}"))?;
        }

        // ── 3. Aggregate kernel: compute per-atom mean/variance across K heads ──
        // Populates ensemble_mean_q_buf and ensemble_var_q_buf [B * na] f32.
        // These buffers are used for exploration bonus at the epoch boundary.
        // The aggregate kernel treats logits as [K * B * na] and outputs mean/variance.
        if let (Some(ref agg_kernel), Some(ref mean_q_buf), Some(ref mut var_q_buf)) = (
            self.ensemble_aggregate_kernel.as_ref(),
            self.ensemble_mean_q_buf.as_ref(),
            self.ensemble_var_q_buf.as_mut(),
        ) {
            let total_ba = b * na;
            let agg_blocks = ((total_ba + 255) / 256) as u32;
            if agg_blocks > 0 {
                let logits_ptr = logits_buf.raw_ptr();
                let mean_ptr = mean_q_buf.raw_ptr();
                let var_ptr = var_q_buf.raw_ptr();
                let k_i32 = k as i32;
                let b_i32 = b as i32;
                let na_i32 = na as i32;
                unsafe {
                    self.stream
                        .launch_builder(agg_kernel)
                        .arg(&logits_ptr)
                        .arg(&mean_ptr)
                        .arg(&var_ptr)
                        .arg(&k_i32)
                        .arg(&b_i32)
                        .arg(&na_i32)
                        .launch(LaunchConfig {
                            grid_dim: (agg_blocks, 1, 1),
                            block_dim: (256, 1, 1),
                            shared_mem_bytes: 0,
                        })
                        .map_err(|e| anyhow::anyhow!("Ensemble aggregate kernel: {e}"))?;
                }
            }
        }

        // ── 4. Compute KL gradient and SAXPY into grad_buf ──────────────
        // The KL gradient encourages head 0 to disagree with heads 1..K-1.
        // d_logits = -diversity_weight * mean_k( softmax(logits_0) - softmax(logits_k) )
        // This is added to grad_buf's value logits section (d_value_logits offset).
        // The graph_adam's single Adam then sees: C51 grad + IQN grad + diversity grad.
        if let (Some(ref kl_kernel), Some(ref d_logits_buf)) =
            (&self.ensemble_kl_grad_kernel, &self.ensemble_d_logits_buf)
        {
            let dw = self.ensemble_diversity_weight;
            let k_i32 = k as i32;
            let b_i32 = b as i32;
            let na_i32 = na as i32;
            let logits_ptr = logits_buf.raw_ptr();
            let d_logits_ptr = d_logits_buf.raw_ptr();

            unsafe {
                self.stream
                    .launch_builder(kl_kernel)
                    .arg(&logits_ptr)
                    .arg(&d_logits_ptr)
                    .arg(&dw)
                    .arg(&k_i32)
                    .arg(&b_i32)
                    .arg(&na_i32)
                    .launch(LaunchConfig {
                        grid_dim: (b as u32, 1, 1),
                        block_dim: (na.min(256) as u32, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| anyhow::anyhow!("ensemble_kl_gradient_kernel: {e}"))?;
            }

            // ── 6. Backward through value head → trunk, SAXPY into grad_buf ──
            // d_logits_buf [B, NA] → backward W_v2 → d_h_v → ReLU mask →
            // backward W_v1 → d_h_s2 → backward trunk (layers 2,1) →
            // SAXPY(diversity_weight) into grad_buf's trunk section.
            // graph_adam then sees: C51 + IQN + ensemble diversity in one update.
            self.trainer.apply_ensemble_diversity_backward(
                d_logits_ptr, dw,
            ).map_err(|e| anyhow::anyhow!("Ensemble diversity backward: {e}"))?;

            tracing::debug!(
                "Ensemble: diversity gradient injected into trunk via value head backward"
            );
        }

        Ok(())
    }

    /// Actions buffer from the last training step.
    pub(crate) fn actions_buf(&self) -> &cudarc::driver::CudaSlice<i32> {
        self.trainer.actions_buf()
    }

    /// Per-sample risk budget [B] — sigmoid output from learned risk branch.
    pub(crate) fn risk_budget_buf(&self) -> &cudarc::driver::CudaSlice<f32> {
        self.trainer.risk_budget_buf()
    }

    /// Per-sample CVaR alpha [B] — derived from risk budget for tail-risk optimization.
    pub(crate) fn cvar_alpha_buf(&self) -> &cudarc::driver::CudaSlice<f32> {
        self.trainer.cvar_alpha_buf()
    }

    /// Per-sample commitment lambda [B] — derived from risk budget for position sizing.
    pub(crate) fn commit_lambda_buf(&self) -> &cudarc::driver::CudaSlice<f32> {
        self.trainer.commit_lambda_buf()
    }

    /// Compute CVaR scales from IQN head and return the device pointer.
    /// Returns 0 (NULL) if IQN is not active.
    pub(crate) fn compute_cvar_device_ptr(
        &mut self,
        actions: &cudarc::driver::CudaSlice<i32>,
        batch_size: usize,
        alpha: f32,
    ) -> Result<u64> {
        let iqn = match self.gpu_iqn.as_mut() {
            Some(iqn) => iqn,
            None => return Ok(0), // No IQN → no CVaR scaling
        };
        let h_s2 = self.trainer.save_h_s2();
        let scales = iqn.compute_cvar_scales(h_s2, actions, batch_size, alpha)
            .map_err(|e| anyhow::anyhow!("CVaR scales: {e}"))?;
        // Store the CudaSlice to keep it alive (prevent deallocation)
        self.cvar_scales_buf = Some(scales);
        // Get the device pointer from the stored buffer
        let stored = self.cvar_scales_buf.as_ref().ok_or_else(|| anyhow::anyhow!("BUG: cvar_scales_buf just set"))?;
        let (ptr, _guard) = stored.device_ptr(&self.stream);
        Ok(ptr)
    }

    /// Device pointer to IQN per-action IQR buffer.
    /// Returns `None` if IQN is not active, `Some(ptr)` otherwise.
    /// The pointer is stable (CudaSlice owned by GpuIqnHead).
    pub(crate) fn iqn_iqr_ptr(&self) -> Option<u64> {
        self.gpu_iqn.as_ref().map(|iqn| iqn.iqr_buf_ptr())
    }

    /// Recompute IQN per-action IQR from the last forward pass's quantile estimates.
    ///
    /// Epoch-boundary call only — involves a synchronous DtoH of quantile data followed
    /// by a HtoD upload of the resulting IQR vector, both of which are forbidden on the
    /// per-step hot path. The result populates `iqr_buf` which `set_iqr_bonus` passes
    /// to the experience collector as a device pointer.
    ///
    /// Returns `Ok(())` if IQN is inactive (nothing to compute).
    pub(crate) fn refresh_iqn_iqr(&mut self) -> anyhow::Result<()> {
        if let Some(ref mut iqn) = self.gpu_iqn {
            iqn.compute_iqr()
                .map_err(|e| anyhow::anyhow!("refresh_iqn_iqr: {e}"))?;
        }
        Ok(())
    }

    /// Plan 4 Task 3 (E.3): launch the IQN multi-quantile diagnostic EMA
    /// kernel. Reads `gpu_iqn.save_q_online` and EMAs the four off-median
    /// fixed quantiles into ISV[99..103). No-op when IQN is inactive (no
    /// quantile surface to read). α is derived adaptively from per-slot
    /// Pearls A+D Wiener state — see `sp4_wiener_ema::pearls_ad_update`.
    pub(crate) fn launch_iqn_quantile_ema(&self) -> Result<(), crate::MLError> {
        let iqn = match self.gpu_iqn.as_ref() {
            Some(iqn) => iqn,
            None => return Ok(()),
        };
        self.trainer.launch_iqn_quantile_ema(
            iqn.save_q_online_ptr(),
            iqn.tba(),
            iqn.num_quantiles(),
        )
    }

    /// Device pointer to ISV signals [8] pinned buffer for adaptive hold enforcement.
    /// Returns the dev_ptr (u64) that the experience collector passes to env_step.
    pub(crate) fn isv_signals_dev_ptr(&self) -> u64 {
        self.trainer.isv_signals_dev_ptr()
    }

    /// Raw device pointer to the ensemble Q-value variance buffer and its length.
    /// The buffer contains per-atom variance across K ensemble heads: [B * num_atoms] f32.
    /// Populated after each training step by the aggregate kernel.
    /// Returns None when ensemble_count <= 1 (no ensemble active).
    pub(crate) fn ensemble_var_ptr(&self) -> Option<(u64, usize)> {
        self.ensemble_var_q_buf.as_ref().map(|buf| (buf.raw_ptr(), buf.len()))
    }

    /// Sync CudaSlice weights back to flat buffer (deferred -- called at epoch boundary).
    ///
    /// Copies:
    /// - Online `CudaSlice -> flat buffer` (reverse sync: Adam-updated weights)
    /// - Target `CudaSlice -> flat buffer` (reverse sync: EMA-blended weights)
    ///
    /// This keeps the Candle flat buffer in sync for checkpointing, Q-value estimation
    /// via `forward()`, and any non-fused code paths.
    pub(crate) fn sync_to_varmap(&mut self, _agent: &mut DQNAgentType) -> Result<()> {
        // No-op: VarStore has been removed. The flat buffer (params_buf) is the
        // single source of truth. Checkpoint save reads directly from params_buf.
        // The old reverse sync copied CudaSlice -> flat buffer for checkpoint
        // compatibility, which is no longer needed.

        // seg_tree_update propagates max_priority internally — no batch_max flush needed.

        self.steps_since_varmap_sync = 0;
        Ok(())
    }

    /// Shrink-and-Perturb: prevent loss of plasticity at walk-forward boundaries.
    ///
    /// `params[i] = alpha * params[i] + (1-alpha) * N(0, sigma)`
    ///
    /// Call between walk-forward windows to maintain the network's ability to
    /// adapt to new market regimes. Invalidates the CUDA Graph.
    pub(crate) fn shrink_and_perturb(&mut self, alpha: f32, sigma: f32) -> Result<()> {
        self.trainer.shrink_and_perturb(alpha, sigma)
            .map_err(|e| anyhow::anyhow!("Shrink-and-Perturb: {e}"))
    }

    /// Update cached median and IQR for robust PopArt normalization.
    /// Called once per epoch after experience collection sorts rewards.
    pub(crate) fn set_robust_popart_stats(&mut self, median: f32, iqr: f32) {
        self.cached_median = median;
        self.cached_iqr = iqr;
    }

    /// Read ISV regime signals (regime_stability, volatility) from pinned memory.
    pub(crate) fn read_isv_regime(&self) -> (f32, f32) {
        self.trainer.read_isv_regime()
    }

    /// D8/N8: Forward (health, regime_stability) read through to caller.
    pub(crate) fn read_isv_regime_via_trainer(&self) -> (f32, f32) {
        self.trainer.read_isv_health_and_regime()
    }

    // ── A3: LearningHealth signal writers / readers ───────────────────────

    /// Read a single f32 value from the ISV signals buffer at given index (pinned DtoH).
    pub(crate) fn read_isv_signal_at(&self, index: usize) -> f32 {
        self.trainer.read_isv_signal_at(index)
    }

    /// Write a single f32 value to ISV signals buffer at given index (pinned HtoD).
    pub(crate) fn write_isv_signal_at(&self, index: usize, value: f32) {
        self.trainer.write_isv_signal_at(index, value);
    }

    /// Read atom utilization from the last Q-stats readback (pinned pointer offset 6).
    pub(crate) fn read_atom_utilization(&self) -> f32 {
        self.trainer.read_atom_utilization()
    }

    /// F2: Read aggregated Adam m-state sample for vector cosine grad_consistency.
    pub(crate) fn read_adam_m_flat_sample(&self, max_per_component: usize) -> anyhow::Result<Vec<f32>> {
        self.trainer.read_adam_m_flat_sample(max_per_component)
            .map_err(|e| anyhow::anyhow!("{e}"))
    }

    /// Last effective cql_alpha from the most recent apply_cql_gradient() call.
    pub(crate) fn last_cql_alpha_eff(&self) -> f32 {
        self.trainer.last_cql_alpha_eff
    }

    /// B3/G4: Last health-scaled Expected SARSA tau factor (1.0 at full health, 6.0 at collapse).
    pub(crate) fn last_sarsa_tau_factor(&self) -> f32 {
        self.trainer.last_sarsa_tau_factor
    }

    /// SP6 Pearl 2 + SP7 (2026-05-03): Compute adaptive gradient budgets.
    ///
    /// Returns `(iqn_branch, ens_branch, iqn_trunk, ens_trunk)`. The CQL and
    /// C51 budgets are no longer returned as scalars — they live on-device in
    /// `lb_budget_effective_buf`, populated by `launch_lb_budget_dispatch`
    /// (called inline below). Callers consume them via the trainer's
    /// `lb_budget_{cql,c51}_{trunk,correction}_dev_ptr` accessors and pass
    /// the device pointers to the SAXPY/scale kernels, which read `alpha`
    /// from device memory at runtime — capture-safe.
    ///
    /// Bug fixed (commit landing this change): `cql_active` / `c51_active`
    /// were resolved via host-side `if/else` here, then passed by value to
    /// `apply_c51_budget_scale` / `apply_cql_saxpy`. Inside the captured
    /// `aux_child` graph, the `if/else` ran ONCE at capture time (step 0 of
    /// each fold, when the activation flag is the FoldReset sentinel `0.0`)
    /// → bootstrap branch resolved → literal `0.02` / `0.05` baked into the
    /// captured kernel-arg buffers, freezing the budget for ~1000 graph
    /// replays per fold regardless of what the controller wrote at runtime.
    /// Same disease as SP4 host-side EMA elimination 2026-05-01; see
    /// `feedback_no_cpu_compute_strict.md` and audit Fix 33.
    ///
    /// IQN keeps a structural BASE_IQN=0.11 floor (reference budget,
    /// Invariant 1 carve-out — the controller's ratio compute would
    /// degenerate if IQN reached 0). ENS keeps the prior numeric-equality
    /// bootstrap because no controller drives BUDGET_ENS_BASE in SP7
    /// (Pearl 2 stopped writing it in T6). Both IQN and ENS still resolve
    /// host-side here; the capture-time freeze applies but is
    /// out-of-scope for this fix per prompt scope.
    pub(crate) fn compute_adaptive_budgets(
        &mut self,
    ) -> ([f32; 4], [f32; 4], f32, f32) {
        // SP7 (2026-05-03): launch the on-device dispatch kernel that resolves
        // `(active >= 0.5) ? max(controller, EPS_DIV) : bootstrap` for CQL and
        // C51 across all 4 branches, computes the trunk mean and per-branch
        // correction, and writes them to `lb_budget_effective_buf`. Captured
        // into the same `aux_child` graph as the SAXPY/scale callers below;
        // every replay sees fresh runtime ISV state.
        if let Err(e) = self.trainer.launch_lb_budget_dispatch() {
            tracing::error!("SP7 lb_budget_dispatch launch failed: {e}");
        }

        let mut iqn = [0.0_f32; 4];
        let mut ens = [0.0_f32; 4];
        // IQN/ENS bootstrap constants. ENS_BOOTSTRAP_BUDGET kept at 0.02 to
        // preserve prior numeric-equality semantics (no controller drives it
        // in SP7 T6+).
        const SP7_EPS_DIV: f32 = 1e-8;
        const ENS_BOOTSTRAP_BUDGET: f32 = 0.02;
        const BASE_IQN: f32 = 0.11;
        let budget_or_bootstrap = |raw: f32, bootstrap: f32| -> f32 {
            if raw < SP7_EPS_DIV { bootstrap } else { raw }
        };
        for b in 0..4_usize {
            iqn[b] = self.read_isv_signal_at(BUDGET_IQN_BASE + b).max(BASE_IQN);
            ens[b] = budget_or_bootstrap(self.read_isv_signal_at(BUDGET_ENS_BASE + b),
                                         ENS_BOOTSTRAP_BUDGET);
        }
        // Trunk/value use mean of 4 branch budgets (D3 decision in SP6 spec).
        let iqn_trunk = iqn.iter().sum::<f32>() / 4.0_f32;
        let ens_trunk = ens.iter().sum::<f32>() / 4.0_f32;
        // Cache IQN/ENS host-side scalars for HEALTH_DIAG. CQL/C51 are read
        // from `lb_budget_effective_buf` at HEALTH_DIAG emit time (see
        // `training_loop.rs` HEALTH_DIAG block) — same memory the kernel
        // populated, no DtoH copy.
        self.trainer.last_iqn_budget_per_branch = iqn;
        self.trainer.last_ens_budget_per_branch = ens;
        self.trainer.last_iqn_budget_eff = iqn_trunk;
        self.trainer.last_ens_budget_eff = ens_trunk;
        (iqn, ens, iqn_trunk, ens_trunk)
    }

    /// SP5 Pearl 2 (flatness-modulated): last per-step IQN budget (mean of 4
    /// per-branch values). BASE_IQN=0.11 floor is enforced in
    /// `compute_adaptive_budgets` — IQN is the reference budget used by the
    /// SP7 loss-balance controller's ratio compute and cannot be driven to 0.
    pub(crate) fn last_iqn_budget_eff(&self) -> f32 { self.trainer.last_iqn_budget_eff }
    /// SP7 (2026-05-03): last per-step CQL budget — mean of the GPU-resolved
    /// per-branch effective budgets in `lb_budget_effective_buf`. Each
    /// per-branch value dispatches on `ISV[LB_CQL_ACTIVE_BASE+b]` on-device:
    /// `CQL_BOOTSTRAP_BUDGET=0.02` while the controller has yet to fire,
    /// controller verbatim (clamped at SP7_EPS_DIV) once active. Caller MUST
    /// have synced the producing stream (HEALTH_DIAG path runs at end of
    /// epoch — already synced).
    pub(crate) fn last_cql_budget_eff(&self) -> f32 {
        let (cql, _c51, _t_cql, _t_c51) = self.trainer.read_lb_budget_effective();
        cql.iter().sum::<f32>() / 4.0_f32
    }
    /// SP7 (2026-05-03): last per-step C51 budget — mean of the GPU-resolved
    /// per-branch effective budgets in `lb_budget_effective_buf`. Same dispatch
    /// semantics as `last_cql_budget_eff` against
    /// `ISV[LB_C51_ACTIVE_BASE+b]` with `C51_BOOTSTRAP_BUDGET=0.05`.
    pub(crate) fn last_c51_budget_eff(&self) -> f32 {
        let (_cql, c51, _t_cql, _t_c51) = self.trainer.read_lb_budget_effective();
        c51.iter().sum::<f32>() / 4.0_f32
    }
    /// SP7 (2026-05-03): per-branch CQL effective budget read from
    /// `lb_budget_effective_buf`. Used by HEALTH_DIAG.
    pub(crate) fn last_cql_budget_per_branch(&self) -> [f32; 4] {
        let (cql, _c51, _t_cql, _t_c51) = self.trainer.read_lb_budget_effective();
        cql
    }
    /// SP7 (2026-05-03): per-branch C51 effective budget read from
    /// `lb_budget_effective_buf`. Used by HEALTH_DIAG.
    pub(crate) fn last_c51_budget_per_branch(&self) -> [f32; 4] {
        let (_cql, c51, _t_cql, _t_c51) = self.trainer.read_lb_budget_effective();
        c51
    }
    /// SP7 (2026-05-03): last per-step ensemble budget (mean of 4 per-branch
    /// values). ENS has no active controller driver in SP7 (Pearl 2 no longer
    /// writes BUDGET_ENS_BASE; loss-balance controller covers CQL/C51 only).
    /// Sentinel-aware bootstrap ENS_BOOTSTRAP_BUDGET=0.02 applied in
    /// `compute_adaptive_budgets` on cold start / fold boundary.
    pub(crate) fn last_ens_budget_eff(&self) -> f32 { self.trainer.last_ens_budget_eff }

    /// F3: Spectral gap — sigma_1 / sigma_2 via rolling Gram + power iteration.
    pub(crate) fn compute_q_spectral_gap(&mut self) -> f32 {
        self.trainer.compute_q_spectral_gap()
    }

    /// D1/N1: Conditionally snapshot current weights at epoch boundary.
    /// Snapshots if health >= 0.7 and q_gap beats the worst stored entry.
    pub(crate) fn maybe_snapshot_qnet(&mut self, health: f32, q_gap: f32, epoch: u32) -> anyhow::Result<bool> {
        self.trainer.maybe_snapshot_params(health, q_gap, epoch)
            .map_err(|e| anyhow::anyhow!("snapshot: {e}"))
    }


    /// D1/N1: Whether distillation was active in the most recent epoch boundary.
    pub(crate) fn last_distill_active(&self) -> bool {
        self.trainer.last_distill_active
    }

    /// F4/F5: Forward meta-Q collapse prediction from the host trainer into the
    /// GPU trainer so aux-op kernels can amplify their weights based on predicted collapse.
    pub(crate) fn set_meta_q_pred(&mut self, pred: f32) {
        self.trainer.set_meta_q_pred(pred);
    }

    pub(crate) fn read_popart_variance(&self) -> f32 {
        if !self.popart_enabled { return 0.0; }
        let mut var = [0.0_f32];
        if self.trainer.read_popart_var(&mut var).is_ok() {
            var[0]
        } else {
            0.0
        }
    }

    pub(crate) fn should_reset_tau(&mut self) -> bool {
        let current_var = self.read_popart_variance();
        if self.prev_popart_var > 0.0 && current_var > 0.0 {
            let ratio = current_var / self.prev_popart_var;
            let shifted = ratio > 2.0 || ratio < 0.5;
            self.prev_popart_var = current_var;
            shifted
        } else {
            self.prev_popart_var = current_var;
            false
        }
    }

    /// Compute expected Q-values for a batch of states using cuBLAS forward (zero Candle).
    ///
    /// Delegates to `GpuDqnTrainer::compute_q_values` which runs the online network
    /// cuBLAS SGEMM forward pass and converts C51 distributional logits to expected
    /// Q-values via the `compute_expected_q` kernel.
    ///
    /// The DtoH readback is intentionally small: `batch_size × total_actions` floats
    /// (e.g. 10 × 11 = 110 floats = 440 bytes) — epoch-boundary only.
    /// Total number of per-branch actions (B0+B1+B2).
    pub(crate) fn total_actions(&self) -> usize {
        self.trainer.total_actions()
    }

    /// Set curiosity weight pointers from the experience collector's `CuriosityWeightSet`.
    /// Must be called once after the experience collector is initialized.
    /// Invalidates CUDA Graphs so the curiosity inference kernel is captured in the next replay.
    pub(crate) fn set_curiosity_weights(
        &mut self,
        weights: &crate::cuda_pipeline::gpu_weights::CuriosityWeightSet,
    ) {
        self.trainer.set_curiosity_weights(&weights.w1, &weights.b1, &weights.w2, &weights.b2);
    }

    /// SP4 Task A14: device pointer to the trainer's `nan_flags_buf`.
    /// Used by external aux paths (`GpuExperienceCollector` for the
    /// curiosity trainer) to share the Pearl C sticky-flag store.
    pub(crate) fn nan_flags_buf_ptr(&self) -> u64 {
        self.trainer.nan_flags_buf_ptr()
    }

    /// SP4 Task A14: device pointer to the trainer's
    /// `clamp_engage_per_block_buf` (mapped pinned). Used by external
    /// aux paths to share the Pearl C per-block engagement-counter buf.
    pub(crate) fn clamp_engage_per_block_buf_dev_ptr(&self) -> u64 {
        self.trainer.clamp_engage_per_block_buf_dev_ptr()
    }

    /// SP4 Task A13.5 (2026-05-01): mapped-pinned host pointer for
    /// `wiener_state_buf`. Used by `init_gpu_experience_collector`
    /// to wire `set_reward_component_pearls_buffers` on the
    /// experience collector (cross-boundary Pearls A+D wiring).
    pub(crate) fn wiener_state_buf_host_ptr(&self) -> *mut f32 {
        self.trainer.wiener_state_buf_host_ptr()
    }

    /// SP4 GPU-Pearls refactor (2026-05-01): mapped-pinned device pointer
    /// for `wiener_state_buf`. The GPU `apply_pearls_ad_kernel` invoked
    /// from `launch_reward_component_ema_inplace` reads/writes Wiener
    /// state via this pointer.
    pub(crate) fn wiener_state_buf_dev_ptr(&self) -> u64 {
        self.trainer.wiener_state_buf_dev_ptr()
    }

    /// SP4 Task A13.5 (2026-05-01): mapped-pinned device pointer for
    /// `producer_step_scratch_buf` (used by the Pearls-A+D-retrofitted
    /// `reward_component_ema` kernel as its scratch destination).
    pub(crate) fn producer_step_scratch_buf_dev_ptr(&self) -> u64 {
        self.trainer.producer_step_scratch_buf_dev_ptr()
    }

    /// SP4 Task A13.5 (2026-05-01): mapped-pinned host pointer for
    /// `producer_step_scratch_buf` (used by `GpuExperienceCollector`
    /// to read the 6 reward-component step observations after kernel
    /// + sync).
    pub(crate) fn producer_step_scratch_buf_host_ptr(&self) -> *mut f32 {
        self.trainer.producer_step_scratch_buf_host_ptr()
    }

    /// SP4 Task A13.5 (2026-05-01): mapped-pinned host pointer for
    /// `isv_signals_pinned` (used by `GpuExperienceCollector` to
    /// read/write ISV[REWARD_POPART_EMA_INDEX..+6) during Pearls A+D
    /// host-side update).
    pub(crate) fn isv_signals_pinned_ptr(&self) -> *mut f32 {
        self.trainer.isv_signals_pinned_ptr()
    }

    /// SP4 Layer A Task A13.5 (2026-05-01): wire the main DQN trainer's
    /// mapped-pinned Pearls A+D buffers into the experience collector
    /// for `launch_reward_component_ema_inplace`. Mirrors the A14/A15
    /// `wire_aux_trainer_pearl_c_buffers` precedent — idempotent, safe
    /// to call multiple times. Per `feedback_wire_everything_up.md`,
    /// the trainer-side orphan `launch_reward_component_ema` was
    /// deleted; this is the canonical (and only) wiring path now.
    pub(crate) fn wire_reward_component_pearls_buffers(
        &mut self,
        collector: &mut crate::cuda_pipeline::gpu_experience_collector::GpuExperienceCollector,
    ) {
        let wiener_dev   = self.trainer.wiener_state_buf_dev_ptr();
        let wiener_host  = self.trainer.wiener_state_buf_host_ptr();
        let scratch_dev  = self.trainer.producer_step_scratch_buf_dev_ptr();
        let scratch_host = self.trainer.producer_step_scratch_buf_host_ptr();
        let isv_dev      = self.trainer.isv_signals_dev_ptr();
        let isv_pinned   = self.trainer.isv_signals_pinned_ptr();
        collector.set_reward_component_pearls_buffers(
            wiener_dev,
            wiener_host,
            scratch_dev,
            scratch_host,
            isv_dev,
            isv_pinned,
        );
    }

    /// SP4 Task A14: wire the main DQN trainer's mapped-pinned Pearl C
    /// buffers into all aux trainers owned by this `FusedTrainingCtx`.
    /// Idempotent — safe to call multiple times. Curiosity is wired
    /// separately because it lives in `GpuExperienceCollector`, not
    /// here; that path goes through
    /// `GpuExperienceCollector::set_curiosity_pearl_c_buffers`.
    pub(crate) fn wire_aux_trainer_pearl_c_buffers(&mut self) {
        let nan_flags_ptr = self.trainer.nan_flags_buf_ptr();
        let engage_buf_ptr = self.trainer.clamp_engage_per_block_buf_dev_ptr();

        // IQN head (group 3 = ParamGroup::Iqn).
        if let Some(ref mut iqn) = self.gpu_iqn {
            iqn.set_pearl_c_buffers(nan_flags_ptr, engage_buf_ptr);
        }
        // IQL high-tau trainer (group 4 = ParamGroup::IqlHigh) — default
        // group on construction; explicitly re-set for clarity.
        self.gpu_iql.set_pearl_c_buffers(nan_flags_ptr, engage_buf_ptr);
        self.gpu_iql.set_pearl_c_param_group(
            crate::cuda_pipeline::ParamGroup::IqlHigh,
        );
        // IQL low-tau trainer (group 5 = ParamGroup::IqlLow).
        self.gpu_iql_low.set_pearl_c_buffers(nan_flags_ptr, engage_buf_ptr);
        self.gpu_iql_low.set_pearl_c_param_group(
            crate::cuda_pipeline::ParamGroup::IqlLow,
        );
        // Attention head (group 6 = ParamGroup::Attn). TLOB shares the
        // attn_adam_kernel cubin but disables Pearl C at the launch site
        // (Layer A); we do not wire it here.
        if let Some(ref mut attn) = self.gpu_attention {
            attn.set_pearl_c_buffers(nan_flags_ptr, engage_buf_ptr);
        }
    }

    /// Compute Q-value statistics from raw f32 states pointer (cold path).
    pub(crate) fn compute_q_stats_from_ptr(
        &mut self,
        states_f32_ptr: u64,
        batch_size: usize,
    ) -> Result<crate::cuda_pipeline::gpu_dqn_trainer::QValueStatsResult> {
        let dst = self.trainer.states_buf_ptr();
        self.trainer.launch_pad_states(dst, states_f32_ptr, batch_size)
            .map_err(|e| anyhow::anyhow!("q_stats pad: {e}"))?;
        // compute_q_stats reads from internal states_buf directly
        self.trainer.compute_q_stats_internal(batch_size)
            .map_err(|e| anyhow::anyhow!("compute_q_stats: {e}"))
    }

    /// Compute Q-value statistics entirely on GPU — 28-byte readback (7 scalars).
    pub(crate) fn compute_q_stats(
        &mut self,
        states: &cudarc::driver::CudaSlice<f32>,
        batch_size: usize,
    ) -> Result<crate::cuda_pipeline::gpu_dqn_trainer::QValueStatsResult> {
        self.trainer.compute_q_stats(states, batch_size)
            .map_err(|e| anyhow::anyhow!("compute_q_stats: {e}"))
    }

    /// Reduce Q-stats from the current training batch (no extra forward pass).
    /// Returns PREVIOUS call's results (async one-step delay).
    pub(crate) fn reduce_current_q_stats(
        &mut self,
    ) -> Result<crate::cuda_pipeline::gpu_dqn_trainer::QValueStatsResult> {
        self.trainer.reduce_current_q_stats()
            .map_err(|e| anyhow::anyhow!("reduce_current_q_stats: {e}"))
    }

    /// Calibrate homeostatic targets — driven by model readiness, not epoch.
    pub(crate) fn calibrate_homeostatic_targets(&mut self) {
        self.trainer.calibrate_homeostatic_targets();
    }

    /// Run cuBLAS forward pass and return reference to GPU-resident Q-values.
    ///
    /// Output shape: `[batch_size, total_actions]` in `q_out_buf`.
    /// No DtoH transfer — caller decides how to consume the CudaSlice.
    /// Deterministic Q-values via CUDA Graph replay.
    pub(crate) fn replay_forward_for_q_values(
        &mut self,
        batch_size: usize,
    ) -> Result<&cudarc::driver::CudaSlice<f32>> {
        self.trainer.replay_forward_for_q_values(batch_size)
            .map_err(|e| anyhow::anyhow!("replay_forward_for_q_values: {e}"))
    }

    /// Pad states from raw pointer into trainer's states_buf.
    pub(crate) fn launch_pad_states(&mut self, dst: u64, src: u64, batch_size: usize) -> Result<()> {
        self.trainer.launch_pad_states(dst, src, batch_size)
            .map_err(|e| anyhow::anyhow!("launch_pad_states: {e}"))
    }

    /// Task 2.X diagnostic — read-only borrow of the inner `GpuDqnTrainer`.
    /// Used by `DQNTrainer::last_isv_branch_q_spreads` to surface ISV
    /// slots [13]/[14] to smoke tests without widening the field's
    /// visibility. Do NOT add write paths through this accessor — the
    /// trainer owns its CUDA context and mutation must go through
    /// properly ordered submission paths.
    pub(crate) fn trainer(&self) -> &GpuDqnTrainer {
        &self.trainer
    }

    /// Plan C Phase 2 follow-up K (2026-04-29): mut accessor for the
    /// fold-boundary fast-grad-norm-EMA reset. The reset is host-side
    /// (mapped-pinned scalar write) so it does NOT touch the captured
    /// CUDA graphs — same caveat as `set_lr` / `set_aux_weight`. Use
    /// only for plain mapped-pinned scalar mutations and other host-side
    /// state on the trainer; do NOT use this to launch graph-affecting
    /// kernels (those must go through properly-ordered submission paths).
    pub(crate) fn trainer_mut(&mut self) -> &mut GpuDqnTrainer {
        &mut self.trainer
    }

    /// Raw pointer to trainer's states buffer.
    pub(crate) fn trainer_states_buf_ptr(&self) -> u64 {
        self.trainer.states_buf_ptr()
    }

    /// Raw pointer to trainer's next_states buffer.
    pub(crate) fn trainer_next_states_buf_ptr(&self) -> u64 {
        self.trainer.next_states_buf_ptr()
    }

    /// Raw pointer to trainer's actions buffer.
    pub(crate) fn trainer_actions_buf_ptr(&self) -> u64 {
        self.trainer.actions_buf_ptr()
    }

    /// Raw pointer to trainer's rewards buffer.
    pub(crate) fn trainer_rewards_buf_ptr(&self) -> u64 {
        self.trainer.rewards_buf_ptr()
    }

    /// Raw pointer to trainer's dones buffer.
    pub(crate) fn trainer_dones_buf_ptr(&self) -> u64 {
        self.trainer.dones_buf_ptr()
    }

    /// Raw pointer to trainer's IS weights buffer.
    pub(crate) fn trainer_is_weights_buf_ptr(&self) -> u64 {
        self.trainer.is_weights_buf_ptr()
    }

    /// SP13 Layer B Commit B1.1b (2026-05-05): raw pointer to trainer's
    /// `aux_nb_label_buf` for direct-to-trainer gather. Wraps the inner
    /// `GpuDqnTrainer::aux_nb_label_buf_ptr()` accessor — same delegation
    /// pattern as the 6 sibling getters above. Caller (`set_trainer_buffers`
    /// in `gpu_replay_buffer.rs`) pipes this into the PER sampler so the
    /// per-step `gather_i32_scalar` writes -1/0/1 labels straight into the
    /// trainer's i32 buffer the K=2 softmax CE consumer reads.
    pub(crate) fn trainer_aux_sign_labels_buf_ptr(&self) -> u64 {
        self.trainer.aux_nb_label_buf_ptr()
    }

    /// Padded state dimension for direct-to-trainer gather.
    pub(crate) fn trainer_state_dim_padded(&self) -> usize {
        self.trainer.state_dim_padded()
    }

    /// Set the PER RNG step counter pointer (pinned device-mapped).
    pub(crate) fn set_per_rng_step_dev_ptr(&mut self, ptr: u64) {
        self.per_rng_step_dev_ptr = ptr;
    }

    /// Run cross-branch Q attention: q_out_buf → q_coord_buf [batch_size, 12].
    pub(crate) fn launch_q_attention(&self, batch_size: usize) -> Result<()> {
        self.trainer.launch_q_attention(batch_size)
            .map_err(|e| anyhow::anyhow!("launch_q_attention: {e}"))
    }

    /// Run cross-branch graph message passing on q_coord_buf [batch_size, 12] in-place.
    /// Must be called AFTER launch_q_attention.
    pub(crate) fn launch_graph_message_pass(&self, batch_size: usize) -> Result<()> {
        self.trainer.launch_graph_message_pass(batch_size)
            .map_err(|e| anyhow::anyhow!("launch_graph_message_pass: {e}"))
    }

    /// Run 3-step diffusion denoising on q_coord_buf conditioned on Var[Q].
    /// Must be called AFTER launch_graph_message_pass.
    pub(crate) fn launch_q_denoise(&self, batch_size: usize) -> Result<()> {
        self.trainer.launch_q_denoise(batch_size)
            .map_err(|e| anyhow::anyhow!("launch_q_denoise: {e}"))
    }

    /// Center Q-values in q_out_buf by subtracting their global mean.
    /// Prevents bootstrapping drift. Must be called after compute_expected_q,
    /// before launch_q_attention.
    pub(crate) fn launch_q_mean_center(&self, batch_size: usize) -> Result<()> {
        self.trainer.launch_q_mean_center(batch_size)
            .map_err(|e| anyhow::anyhow!("launch_q_mean_center: {e}"))
    }

    /// Mamba2 temporal scan: update rolling history and enrich h_s2.
    pub(crate) fn mamba2_step(&self, batch_size: usize) -> Result<()> {
        self.trainer.mamba2_step(batch_size)
            .map_err(|e| anyhow::anyhow!("mamba2_step: {e}"))
    }

    /// Mamba2 BPTT: compute gradients for W_A, W_B, W_C through K=8 scan steps.
    pub(crate) fn mamba2_backward(&mut self, batch_size: usize) -> Result<()> {
        self.trainer.mamba2_backward(batch_size)
            .map_err(|e| anyhow::anyhow!("mamba2_backward: {e}"))
    }

    /// Standalone SGD step for Mamba2 parameters.
    pub(crate) fn step_mamba2_adam(&mut self) -> Result<()> {
        self.trainer.step_mamba2_adam()
            .map_err(|e| anyhow::anyhow!("step_mamba2_adam: {e}"))
    }

    /// OFI embed MLP backward: accumulate Mamba2 + attention gradients, compute dW/db.
    pub(crate) fn launch_ofi_embed_backward(&mut self, batch_size: usize) -> Result<()> {
        // Get attention gradient pointer if attention is present
        let d_ofi_attn_ptr = self.gpu_attention.as_ref()
            .map(|attn| attn.d_ofi_embed_attn().raw_ptr());
        self.trainer.launch_ofi_embed_backward(batch_size, d_ofi_attn_ptr)
            .map_err(|e| anyhow::anyhow!("launch_ofi_embed_backward: {e}"))
    }

    /// Adam optimizer step for OFI embed MLP parameters (LR=1e-4, clip=1.0).
    pub(crate) fn step_ofi_embed_adam(&mut self) -> Result<()> {
        self.trainer.step_ofi_embed_adam()
            .map_err(|e| anyhow::anyhow!("step_ofi_embed_adam: {e}"))
    }

    /// Gentle weight decay for risk branch parameters (MVP training signal).
    pub(crate) fn step_risk_sgd(&mut self) -> Result<()> {
        self.trainer.step_risk_sgd()
            .map_err(|e| anyhow::anyhow!("step_risk_sgd: {e}"))
    }

    /// Increment new component LR warmup counter.
    pub(crate) fn increment_warmup(&mut self) {
        self.trainer.increment_warmup();
    }

    /// Multi-horizon value forward: 5-bar and 20-bar value heads.
    /// Degenerate blend (copies 1-bar logits) until cuBLAS GEMMs are wired in Task 9c.
    pub(crate) fn multi_horizon_value_forward(&self, batch_size: usize) -> Result<()> {
        self.trainer.multi_horizon_value_forward(batch_size)
            .map_err(|e| anyhow::anyhow!("multi_horizon_value_forward: {e}"))
    }

    /// Selectivity gate forward: sigmoid(dot(W_sel, h_s2) + b_sel) per sample.
    pub(crate) fn launch_selectivity_forward(&self, batch_size: usize) -> Result<()> {
        self.trainer.launch_selectivity_forward(batch_size)
            .map_err(|e| anyhow::anyhow!("launch_selectivity_forward: {e}"))
    }

    /// Selectivity gate backward: BCE gradient on (sel, per_sample_loss / mean_loss).
    pub(crate) fn launch_selectivity_backward(&mut self, batch_size: usize, mean_loss: f32) -> Result<()> {
        self.trainer.launch_selectivity_backward(batch_size, mean_loss)
            .map_err(|e| anyhow::anyhow!("launch_selectivity_backward: {e}"))
    }

    /// Adam update for selectivity gate parameters (LR=1e-4, clip=1.0).
    pub(crate) fn step_selectivity_adam(&mut self) -> Result<()> {
        self.trainer.step_selectivity_adam()
            .map_err(|e| anyhow::anyhow!("step_selectivity_adam: {e}"))
    }

    /// Compute target network Q-values [B, 12] for denoiser training.
    pub(crate) fn compute_denoise_target_q(&self, batch_size: usize) -> Result<()> {
        self.trainer.compute_denoise_target_q(batch_size)
            .map_err(|e| anyhow::anyhow!("compute_denoise_target_q: {e}"))
    }

    /// Backward pass through diffusion denoiser (MSE vs target Q-values).
    pub(crate) fn launch_q_denoise_backward(&mut self, batch_size: usize) -> Result<()> {
        self.trainer.launch_q_denoise_backward(batch_size)
            .map_err(|e| anyhow::anyhow!("launch_q_denoise_backward: {e}"))
    }

    /// Adam update for denoise_params (LR=1e-4, clip=1.0).
    pub(crate) fn step_denoise_adam(&mut self) -> Result<()> {
        self.trainer.step_denoise_adam()
            .map_err(|e| anyhow::anyhow!("step_denoise_adam: {e}"))
    }

    /// Set C51 blend factor for gradual MSE→C51 ramp.
    ///
    /// `alpha = 0.0`: pure MSE (warmup). `alpha = 1.0`: pure C51 (converged).
    /// Alpha is a scalar arg baked at graph capture time.
    pub(crate) fn set_c51_alpha(&mut self, alpha: f32) {
        self.trainer.set_c51_alpha(alpha);
    }

    /// Plan 4 Task 6 Commit B: thin delegate for the aux-loss weight setter.
    /// `aux_weight` is the per-step ISV-driven scalar that scales aux head
    /// gradient SAXPYs into `grad_buf[119..127)` and `bw_d_h_s2`. See
    /// `GpuDqnTrainer::set_aux_weight` for the [0.05, 0.3] clamp rationale.
    pub(crate) fn set_aux_weight(&mut self, aux_weight: f32) {
        self.trainer.set_aux_weight(aux_weight);
    }

    /// Return a reference to the GPU-resident loss scalar (f32) for the training guard.
    ///
    /// During warmup (c51_alpha near 0), C51 loss may be NaN from random logits.
    /// The training guard reads this pointer for NaN detection, so we must return
    /// `mse_loss_buf` during warmup to avoid false NaN halts.
    /// Once C51 is active (alpha > 0), returns `total_loss_buf` (C51 accumulator).
    pub(crate) fn online_dueling_ref(&self) -> &DuelingWeightSet {
        &self.online_dueling
    }

    pub(crate) fn online_branching_ref(&self) -> &BranchingWeightSet {
        &self.online_branching
    }

    /// Wait for the restore_best_params writes on a different stream.
    /// Uses cuStreamWaitEvent — GPU-side only, no CPU stall.
    pub(crate) fn wait_restore_on_stream(&self, other_stream: &std::sync::Arc<cudarc::driver::CudaStream>) -> Result<()> {
        if let Some(ref event) = self.restore_event {
            unsafe {
                cudarc::driver::sys::cuStreamWaitEvent(
                    other_stream.cu_stream(),
                    event.cu_event(),
                    0,
                );
            }
        }
        Ok(())
    }

    /// Raw pointer + length of the best-epoch flat f32 params buffer.
    /// The evaluator can DtoD copy from this on ITS OWN stream — no
    /// cross-stream visibility issues. Returns None if no best snapshot saved.
    pub(crate) fn best_params_flat(&self) -> Option<(u64, usize)> {
        self.best_params_snapshot.as_ref().map(|s| (s.raw_ptr(), s.len()))
    }

    /// Raw device pointer to the flat f32 online params buffer (for DtoD sync).
    pub(crate) fn params_flat_ptr(&self) -> u64 {
        self.trainer.params().raw_ptr()
    }

    /// Read one sample's state vector from the GPU states_buf (DtoH, 96 floats).
    /// Diagnostic use only -- negligible cost. Returns the padded state vector.
    pub(crate) fn read_state_sample(&self, sample_idx: usize) -> anyhow::Result<Vec<f32>> {
        let sd = ml_core::state_layout::STATE_DIM;
        let sd_padded = (sd + 127) & !127;
        let states = self.trainer.states_buf();
        let host: Vec<f32> = self.stream.clone_dtoh(states)
            .map_err(|e| anyhow::anyhow!("states DtoH for diagnostic: {e}"))?;
        let offset = sample_idx * sd_padded;
        if offset + sd > host.len() {
            anyhow::bail!("sample_idx {} out of range (host len={}, sd={})", sample_idx, host.len(), sd);
        }
        Ok(host[offset..offset + sd].to_vec())
    }

    /// Compute a simple checksum of the flat params buffer (GPU→CPU readback).
    /// For determinism diagnostics only — synchronizes the stream.
    pub(crate) fn params_checksum(&self) -> anyhow::Result<u64> {
        let n = self.trainer.total_params();
        let params = self.trainer.params();
        let host: Vec<f32> = self.stream.clone_dtoh(params)
            .map_err(|e| anyhow::anyhow!("params DtoH for checksum: {e}"))?;
        let mut hash: u64 = 0xcbf29ce484222325; // FNV-1a
        for &v in &host[..n] {
            let bits = v.to_bits();
            hash ^= bits as u64;
            hash = hash.wrapping_mul(0x100000001b3);
        }
        Ok(hash)
    }

    /// Access the shared cuBLAS handle for passing to new CublasGemmSet instances
    /// (e.g., the GPU experience collector).
    pub(crate) fn shared_cublas(&self) -> &Arc<crate::cuda_pipeline::shared_cublas_handle::PerStreamCublasHandles> {
        self.trainer.shared_cublas()
    }

    /// Per-sample support pointer from IQL trainer.
    pub(crate) fn per_sample_support_ptr(&self) -> u64 { self.gpu_iql.per_sample_support_ptr() }
    /// Fixed wide v_range pointer for eval compute_expected_q (not C51 loss).
    pub(crate) fn eval_v_range_ptr(&self) -> u64 { self.trainer.eval_v_range_ptr() }
    /// Return current eval v_range as [v_min, v_max] via synchronous DtoH.
    /// Reads the same buffer that update_eval_v_range wrote.
    pub(crate) fn eval_v_range(&self) -> [f32; 2] {
        let mut vr = [0.0_f32; 2];
        unsafe {
            cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream());
            cudarc::driver::sys::cuMemcpyDtoH_v2(
                vr.as_mut_ptr().cast(),
                self.trainer.eval_v_range_ptr(),
                2 * std::mem::size_of::<f32>(),
            );
        }
        vr
    }
    /// Update eval v_range from observed per-branch Q-value statistics.
    /// Spec 2026-04-23: writes 8 ISV slots (centre + half-width per branch)
    /// and, for backward compatibility, the legacy scalar `eval_v_range_pinned`
    /// from the direction-branch bounds.
    pub(crate) fn update_eval_v_range(
        &mut self,
        per_branch_stats: &crate::cuda_pipeline::gpu_dqn_trainer::PerBranchQValueStats,
        per_branch_q_gaps: [f32; 4],
    ) {
        self.trainer.update_eval_v_range(per_branch_stats, per_branch_q_gaps);
    }

    /// Reduce per-branch Q-stats from the current `q_out_buf` — 4 × 7 floats
    /// covering `[avg_max_q, q_min, q_max, q_mean, q_var, atom_entropy(=0),
    /// atom_utilization(=0)]` per branch. Feeds `update_eval_v_range` for the
    /// ISV v-range per-branch EMA bus.
    pub(crate) fn reduce_current_q_stats_per_branch(&mut self)
        -> Result<crate::cuda_pipeline::gpu_dqn_trainer::PerBranchQValueStats, crate::MLError>
    {
        self.trainer.reduce_current_q_stats_per_branch()
    }
    /// Per-sample epsilon from IQL expectile gap.
    pub(crate) fn per_sample_epsilon_ptr(&self) -> u64 { self.gpu_iql.per_sample_epsilon_ptr() }
    pub(crate) fn num_atoms(&self) -> usize { self.trainer.config().num_atoms }
    /// Plan 1 Task 9: GPU-driven atom position recompute (all 4 branches, single launch).
    pub(crate) fn launch_atoms_update(&self) -> anyhow::Result<()> {
        self.trainer.launch_atoms_update()
            .map_err(|e| anyhow::anyhow!("launch_atoms_update: {e}"))
    }
    pub(crate) fn update_adaptive_clip(&mut self, grad_norm: f32) { self.trainer.update_adaptive_clip(grad_norm); }
    pub(crate) fn adaptive_clip_value(&self) -> f32 { self.trainer.adaptive_clip_value() }
    /// Plan C Phase 2 follow-up K (2026-04-29): override the active
    /// adaptive-clip pinned slot with a warmup-factor-scaled value. Mirrors
    /// `set_lr`'s zero-graph-recapture contract — the kernel reads
    /// `*adaptive_clip_pinned` at replay time, so writing the pinned slot
    /// takes effect on the next captured graph replay without invalidation.
    /// The companion `update_adaptive_clip` sets the EMA-derived baseline
    /// (`grad_norm_ema × 2`, floored at 1.0); this setter then scales that
    /// baseline by the warmup-factor consumer's `0.1 + 0.9 × factor`.
    pub(crate) fn set_active_clip(&mut self, clip: f32) { self.trainer.set_active_clip(clip); }
    /// Raw pointer to plan_params_buf [B, 6] for trade plan integration.
    pub(crate) fn plan_params_buf_ptr(&self) -> u64 {
        self.trainer.plan_params_buf_ptr()
    }

    /// Raw device pointer to the flat f32 target params buffer (for DtoD sync).
    pub(crate) fn target_params_flat_ptr(&self) -> u64 {
        self.trainer.target_params().raw_ptr()
    }

    /// Total number of parameters (same for online and target).
    pub(crate) fn total_param_count(&self) -> usize {
        self.trainer.total_params()
    }

    // ── Trajectory backtracking accessors ──────────────────────────────

    /// Download online params from GPU to host (synchronous DtoH).
    pub(crate) fn download_params(&self, dst: &mut [f32]) -> anyhow::Result<()> {
        self.trainer.download_params(dst)
            .map_err(|e| anyhow::anyhow!("download params: {e}"))
    }

    /// Upload online params from host to GPU (HtoD).
    pub(crate) fn upload_params(&mut self, src: &[f32]) -> anyhow::Result<()> {
        self.trainer.upload_params(src)
            .map_err(|e| anyhow::anyhow!("upload params: {e}"))
    }

    /// Download target params from GPU to host (synchronous DtoH).
    pub(crate) fn download_target_params(&self, dst: &mut [f32]) -> anyhow::Result<()> {
        self.trainer.download_target_params(dst)
            .map_err(|e| anyhow::anyhow!("download target params: {e}"))
    }

    /// Upload target params from host to GPU (HtoD).
    pub(crate) fn upload_target_params(&mut self, src: &[f32]) -> anyhow::Result<()> {
        self.trainer.upload_target_params(src)
            .map_err(|e| anyhow::anyhow!("upload target params: {e}"))
    }

    /// Clone params_buf to a new CudaSlice via DtoD — checkpoint stays on GPU, zero CPU.
    pub(crate) fn clone_params_gpu(&self) -> anyhow::Result<cudarc::driver::CudaSlice<f32>> {
        self.trainer.clone_params_gpu()
            .map_err(|e| anyhow::anyhow!("clone params gpu: {e}"))
    }

    /// Clone target_params_buf to a new CudaSlice via DtoD — checkpoint stays on GPU, zero CPU.
    pub(crate) fn clone_target_params_gpu(&self) -> anyhow::Result<cudarc::driver::CudaSlice<f32>> {
        self.trainer.clone_target_params_gpu()
            .map_err(|e| anyhow::anyhow!("clone target params gpu: {e}"))
    }

    /// Restore params_buf from a GPU checkpoint via DtoD — zero CPU involvement.
    pub(crate) fn restore_params_from_gpu(&mut self, src: &cudarc::driver::CudaSlice<f32>) -> anyhow::Result<()> {
        self.trainer.restore_params_from_gpu(src)
            .map_err(|e| anyhow::anyhow!("restore params from gpu: {e}"))
    }

    /// Restore target_params_buf from a GPU checkpoint via DtoD — zero CPU involvement.
    pub(crate) fn restore_target_params_from_gpu(&mut self, src: &cudarc::driver::CudaSlice<f32>) -> anyhow::Result<()> {
        self.trainer.restore_target_params_from_gpu(src)
            .map_err(|e| anyhow::anyhow!("restore target params from gpu: {e}"))
    }

    /// Reset Adam optimizer state (m=0, v=0, step=0).
    pub(crate) fn reset_adam_state(&mut self) -> anyhow::Result<()> {
        self.trainer.reset_adam_state()
            .map_err(|e| anyhow::anyhow!("reset adam: {e}"))
    }

    /// Reset Adam momentum (m/v) for branch weights only (indices 8-41).
    /// Trunk momentum (indices 0-7) preserved. Lightweight plateau-breaking fix.
    pub(crate) fn reset_branch_adam_momentum(&mut self) -> anyhow::Result<()> {
        self.trainer.reset_branch_adam_momentum()
            .map_err(|e| anyhow::anyhow!("reset branch adam: {e}"))
    }

    /// Scale Adam momentum buffer (m_buf) by `scale`. Used for warm restart.
    pub(crate) fn scale_adam_momentum(&mut self, scale: f32) -> anyhow::Result<()> {
        self.trainer.scale_adam_momentum(scale)
            .map_err(|e| anyhow::anyhow!("scale adam momentum: {e}"))
    }

    /// Set learning rate. Invalidates parent graph for re-capture on next step.
    pub(crate) fn set_lr(&mut self, lr: f32) {
        self.trainer.set_lr(lr);
    }

    /// Read adam_step counter.
    pub(crate) fn adam_step_val(&self) -> i32 { self.trainer.adam_step_val() }

    /// Read eval_q_mean_ema.
    pub(crate) fn eval_q_mean_ema(&self) -> f32 { self.trainer.eval_q_mean_ema() }

    /// Read eval_q_std_ema.
    pub(crate) fn eval_q_std_ema(&self) -> f32 { self.trainer.eval_q_std_ema() }

    /// Read per_branch_q_gap_ema.
    pub(crate) fn per_branch_q_gap_ema(&self) -> [f32; 4] { self.trainer.per_branch_q_gap_ema() }
    pub(crate) fn get_per_branch_q_gaps(&self) -> [f32; 4] { self.trainer.get_per_branch_q_gaps() }
    pub(crate) fn utilization_ema(&self) -> f32 { self.trainer.utilization_ema() }

    /// G9: Set epoch seed for regime dropout PRNG.
    pub(crate) fn set_regime_dropout_seed(&mut self, seed: i32) {
        self.trainer.set_regime_dropout_seed(seed);
    }

    /// G4: Set adaptive gamma for C51 Bellman projection.
    /// Plan 1 Task 13: delegate GPU tau update to trainer.
    pub(crate) fn launch_tau_update(&self, tau_base: f32, tau_final: f32)
        -> anyhow::Result<()>
    {
        self.trainer.launch_tau_update(tau_base, tau_final)
            .map_err(|e| anyhow::anyhow!("launch_tau_update: {e}"))
    }

    /// Plan 1 Task 14: delegate GPU epsilon update to trainer.
    pub(crate) fn launch_epsilon_update(&self, eps_start: f32, eps_end: f32)
        -> anyhow::Result<()>
    {
        self.trainer.launch_epsilon_update(eps_start, eps_end)
            .map_err(|e| anyhow::anyhow!("launch_epsilon_update: {e}"))
    }

    /// Plan 2 Task 3 D.2: delegate GPU per-branch gamma update to trainer.
    pub(crate) fn launch_per_branch_gamma_update(&self)
        -> anyhow::Result<()>
    {
        self.trainer.launch_per_branch_gamma_update()
            .map_err(|e| anyhow::anyhow!("launch_per_branch_gamma_update: {e}"))
    }

    /// Plan 1 Task 11: delegate GPU kelly_cap update to trainer.
    pub(crate) fn launch_kelly_cap_update(
        &self,
        portfolio_state_dev_ptr: u64,
        n_envs: i32,
    ) -> anyhow::Result<()> {
        self.trainer.launch_kelly_cap_update(portfolio_state_dev_ptr, n_envs)
            .map_err(|e| anyhow::anyhow!("launch_kelly_cap_update: {e}"))
    }

    /// Plan 2 Task 1 C.1: GPU-driven per-branch Q P5/P95 percentile EMA update.
    /// Delegates to `GpuDqnTrainer::launch_q_quantile_reduce`. Cold-path per-epoch.
    pub(crate) fn launch_q_quantile_reduce(&self, ema_alpha: f32) -> anyhow::Result<()> {
        self.trainer.launch_q_quantile_reduce(ema_alpha)
            .map_err(|e| anyhow::anyhow!("launch_q_quantile_reduce: {e}"))
    }

    /// G5: Set the EMA threshold for epistemic variance gating.
    pub(crate) fn set_var_ema(&self, val: f32) {
        self.trainer.set_var_ema(val);
    }

    /// Current IQN readiness [0, 1]: 0=exploring, 1=exploiting.
    pub(crate) fn iqn_readiness(&self) -> f32 { self.trainer.iqn_readiness() }

    /// Device pointer to IQN readiness [1] pinned — for plan enforcement gating.
    pub(crate) fn iqn_readiness_dev_ptr(&self) -> u64 { self.trainer.iqn_readiness_dev_ptr() }


    /// Set eval_q_mean_ema (trajectory backtracking restore).
    pub(crate) fn set_eval_q_mean_ema(&mut self, v: f32) { self.trainer.set_eval_q_mean_ema(v); }

    /// Set eval_q_std_ema (trajectory backtracking restore).
    pub(crate) fn set_eval_q_std_ema(&mut self, v: f32) { self.trainer.set_eval_q_std_ema(v); }

    /// Set per_branch_q_gap_ema (trajectory backtracking restore).
    pub(crate) fn set_per_branch_q_gap_ema(&mut self, v: [f32; 4]) { self.trainer.set_per_branch_q_gap_ema(v); }

    // ── End trajectory backtracking accessors ──────────────────────────

    /// Raw pointer + length of the current flat f32 params buffer.
    pub(crate) fn params_flat(&self) -> (u64, usize) {
        let p = self.trainer.params();
        (p.raw_ptr(), p.len())
    }

    /// Snapshot current params as the best model. Pure async DtoD — no sync, no Candle.
    /// Called when a new best epoch is found during training.
    pub(crate) fn save_best_params(&mut self) -> Result<()> {
        let src = self.trainer.params();
        let n = src.len();
        let n_bytes = n * std::mem::size_of::<f32>();
        let src_ptr = src.raw_ptr();

        // Allocate or reuse snapshot buffer
        if self.best_params_snapshot.is_none() || self.best_params_snapshot.as_ref().unwrap().len() != n {
            let buf = self.stream.alloc_zeros::<f32>(n)
                .map_err(|e| anyhow::anyhow!("save_best_params alloc: {e}"))?;
            self.best_params_snapshot = Some(buf);
        }
        let dst_ptr = self.best_params_snapshot.as_ref().unwrap().raw_ptr();

        // Async DtoD — no sync
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr, src_ptr, n_bytes, self.stream.cu_stream(),
            ).map_err(|e| anyhow::anyhow!("save_best_params DtoD: {e}"))?;
        }
        Ok(())
    }

    /// Restore best-epoch params into online weights. DtoD + unflatten + sync.
    /// Called before evaluation to ensure the evaluator sees the best model.
    /// Syncs the stream to guarantee writes are globally visible before the
    /// evaluator (which may use a different stream) reads them.
    pub(crate) fn restore_best_params(&mut self) -> Result<()> {
        let snapshot = self.best_params_snapshot.as_ref()
            .ok_or_else(|| anyhow::anyhow!("No best params snapshot saved"))?;
        let n_bytes = snapshot.len() * std::mem::size_of::<f32>();
        let src_ptr = snapshot.raw_ptr();
        let dst_ptr = self.trainer.params_mut().raw_ptr();
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr, src_ptr, n_bytes, self.stream.cu_stream(),
            ).map_err(|e| anyhow::anyhow!("restore_best_params DtoD: {e}"))?;
        }
        // Unflatten: params_buf → individual weight tensors
        self.trainer.unflatten_online_weights(&self.online_dueling, &self.online_branching)
            .map_err(|e| anyhow::anyhow!("restore_best_params unflatten: {e}"))?;
        // Record event on trainer's stream — the evaluator waits for this event
        // on its own stream before reading the weights. No CPU sync — purely
        // GPU-side inter-stream dependency via cuStreamWaitEvent.
        let event = self.stream.record_event(None)
            .map_err(|e| anyhow::anyhow!("restore_best_params event record: {e}"))?;
        self.restore_event = Some(event);
        Ok(())
    }

    /// Raw device pointer to the active loss scalar (pinned device-mapped).
    /// Returns MSE dev_ptr when c51_alpha ≈ 0, otherwise C51 total_loss dev_ptr.
    pub(crate) fn loss_gpu_ptr(&self) -> u64 {
        if self.trainer.c51_alpha() < 1e-6 {
            self.trainer.mse_loss_dev_ptr()
        } else {
            self.trainer.total_loss_dev_ptr()
        }
    }

    /// Raw device pointer to grad_norm scalar (pinned device-mapped).
    /// Written by the CUDA graph's grad_norm kernel — valid after `replay_adam()`.
    pub(crate) fn grad_norm_gpu_ptr(&self) -> u64 {
        self.trainer.grad_norm_dev_ptr()
    }

    /// Read NaN detection flags (synchronizes stream). Used by training guard on NaN halt.
    /// Index → buffer mapping is documented in
    /// `GpuDqnTrainer::run_nan_checks_post_forward`.
    /// SP1 Phase B: returns 50 flags. Slots 24-35 cover backward-path
    /// checks; slots 36-47 SP3 Mech 5 ISV-threshold; slot 48 SP3 close-out v2
    /// IQN online_params; slot 49 SP3 Mech 10 post-trunk h_s2 sentinel.
    pub(crate) fn read_nan_flags(&self) -> Result<[i32; 50], crate::MLError> {
        self.trainer.read_nan_flags()
    }

    /// Get CUDA stream reference for DtoH transfers after compute_q_values.
    pub(crate) fn stream(&self) -> &Arc<cudarc::driver::CudaStream> {
        &self.stream
    }
}

impl crate::cuda_pipeline::q_value_provider::QValueProvider for FusedTrainingCtx {
    fn max_batch_size(&self) -> usize {
        self.batch_size
    }

    fn total_actions(&self) -> usize {
        self.trainer.total_actions()
    }

    fn padded_state_dim(&self) -> usize {
        // pad128: round up to 128-byte alignment for cuBLAS
        ml_core::state_layout::STATE_DIM_PADDED
    }

    fn compute_q_values_to(
        &mut self,
        states_ptr: u64,
        batch_size: usize,
        q_out_ptr: u64,
    ) -> Result<(), crate::MLError> {
        let max_b = self.batch_size;
        let tba = self.trainer.total_actions();
        let f32_size = std::mem::size_of::<f32>();

        for offset in (0..batch_size).step_by(max_b) {
            let sub_b = (batch_size - offset).min(max_b);
            let src_ptr = states_ptr + (offset * self.padded_state_dim() * f32_size) as u64;

            // Write states to trainer's states_buf, then run graphed forward
            let dst = self.trainer.states_buf_ptr();
            self.trainer.launch_pad_states(dst, src_ptr, sub_b)
                .map_err(|e| crate::MLError::ModelError(format!("QValueProvider pad: {e}")))?;
            self.trainer.replay_forward_for_q_values(sub_b)
                .map_err(|e| crate::MLError::ModelError(format!("QValueProvider: {e}")))?;

            // DtoD copy: trainer q_out_buf → caller's output at offset
            let q_src = self.trainer.q_out_buf().raw_ptr();
            let q_dst = q_out_ptr + (offset * tba * f32_size) as u64;
            let copy_bytes = sub_b * tba * f32_size;
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    q_dst, q_src, copy_bytes, self.stream.cu_stream(),
                ).map_err(|e| crate::MLError::ModelError(format!("QValueProvider DtoD: {e}")))?;
            }
        }
        Ok(())
    }

    fn compute_q_and_b_logits_to(
        &mut self,
        states_ptr: u64,
        batch_size: usize,
        q_out_ptr: u64,
        b_logits_out_ptr: u64,
    ) -> Result<(), crate::MLError> {
        // Mirror of `compute_q_values_to`, but additionally DtoD-copies the
        // trainer's `on_b_logits_buf` (branch-major
        // `[sub_b, total_actions * num_atoms]`) into the caller's output at
        // the same per-iteration offset. Plan C T2 amendment buffer feed for
        // `experience_action_select` Thompson direction sampling
        // (gpu_backtest_evaluator val pipeline, Plan C T4, 2026-04-29).
        let max_b = self.batch_size;
        let tba = self.trainer.total_actions();
        let na = self.trainer.config().num_atoms;
        let row_b_logits = tba * na; // total_branch_atoms
        let f32_size = std::mem::size_of::<f32>();

        for offset in (0..batch_size).step_by(max_b) {
            let sub_b = (batch_size - offset).min(max_b);
            let src_ptr = states_ptr + (offset * self.padded_state_dim() * f32_size) as u64;

            let dst = self.trainer.states_buf_ptr();
            self.trainer.launch_pad_states(dst, src_ptr, sub_b)
                .map_err(|e| crate::MLError::ModelError(format!("QValueProvider pad: {e}")))?;
            self.trainer.replay_forward_for_q_values(sub_b)
                .map_err(|e| crate::MLError::ModelError(format!("QValueProvider: {e}")))?;

            // DtoD: q_out_buf → caller's q_out at offset.
            let q_src = self.trainer.q_out_buf().raw_ptr();
            let q_dst = q_out_ptr + (offset * tba * f32_size) as u64;
            let q_bytes = sub_b * tba * f32_size;
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    q_dst, q_src, q_bytes, self.stream.cu_stream(),
                ).map_err(|e| crate::MLError::ModelError(format!("QValueProvider DtoD q: {e}")))?;
            }

            // DtoD: on_b_logits_buf → caller's b_logits at offset.
            // Branch-major layout `[B, total_actions * num_atoms]`; the
            // direction-branch slice (Branch 0) lives at the start of each
            // sample's row, matching the kernel's
            // `b_logits_dir + i * b0 * n_atoms` indexing.
            let b_src = self.trainer.on_b_logits_buf_ptr();
            let b_dst = b_logits_out_ptr + (offset * row_b_logits * f32_size) as u64;
            let b_bytes = sub_b * row_b_logits * f32_size;
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    b_dst, b_src, b_bytes, self.stream.cu_stream(),
                ).map_err(|e| crate::MLError::ModelError(format!("QValueProvider DtoD b_logits: {e}")))?;
            }
        }
        Ok(())
    }

    fn total_branch_atoms(&self) -> usize {
        self.trainer.total_actions() * self.trainer.config().num_atoms
    }

    fn num_atoms(&self) -> usize {
        self.trainer.config().num_atoms
    }

    fn per_sample_support_ptr(&self) -> u64 {
        self.trainer.per_sample_support_ptr_get()
    }

    fn atom_positions_ptr(&self) -> u64 {
        self.trainer.atom_positions_buf_ptr()
    }

    fn stream_handle(&self) -> u64 {
        self.stream.cu_stream() as u64
    }

    /// Run the trade-plan MLP on `save_h_s2` (populated by the most recent
    /// `compute_q_values_to` call) and DtoD-copy the resulting plan_params
    /// to the caller's buffer.
    ///
    /// Val-Flat-collapse fix #5 (task #94, 2026-04-24). The backtest
    /// evaluator calls this after each chunk's forward pass so that
    /// `plan_isv` state-vector positions [98..105) carry real plan signals
    /// matching the training distribution, instead of the zero-fill that
    /// was biasing val argmax toward Flat/Hold on 99.99% of bars.
    ///
    /// Batches are processed in the trainer's internal `batch_size`
    /// chunks (same as `compute_q_values_to`) so that the plan MLP sees
    /// the same h_s2 layout as training.
    fn compute_plan_params(
        &mut self,
        plan_out_ptr: u64,
        batch_size: usize,
    ) -> Result<(), crate::MLError> {
        let max_b = self.batch_size;
        let f32_size = std::mem::size_of::<f32>();
        let plan_width = 6_usize;

        for offset in (0..batch_size).step_by(max_b) {
            let sub_b = (batch_size - offset).min(max_b);

            // Run plan MLP on trainer's current save_h_s2 (populated by
            // the preceding compute_q_values_to call for this offset).
            self.trainer.launch_trade_plan_forward(sub_b)
                .map_err(|e| crate::MLError::ModelError(
                    format!("QValueProvider plan_params: {e}")
                ))?;

            // DtoD copy plan_params_buf → caller's output at offset.
            let plan_src = self.trainer.plan_params_buf_ptr();
            let plan_dst = plan_out_ptr + (offset * plan_width * f32_size) as u64;
            let copy_bytes = sub_b * plan_width * f32_size;
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    plan_dst, plan_src, copy_bytes, self.stream.cu_stream(),
                ).map_err(|e| crate::MLError::ModelError(
                    format!("QValueProvider plan_params DtoD: {e}")
                ))?;
            }
        }
        Ok(())
    }
}

/// Deep-copy the 12 weight tensors of a DuelingWeightSet from pointer sources
/// into newly-allocated CudaSlice buffers. Returns both the backing storage and
/// a pointer-view weight set.
///
/// Used for ensemble extra heads that need independent weight copies.
fn clone_dueling_weights(
    src: &DuelingWeightSet,
    stream: &Arc<cudarc::driver::CudaStream>,
    head_idx: usize,
) -> Result<(gpu_weights::DuelingWeightBacking, DuelingWeightSet)> {
    use crate::cuda_pipeline::gpu_dqn_trainer::dtod_copy_async;

    let clone_from_ptr = |ptr: u64, len: usize| -> Result<cudarc::driver::CudaSlice<f32>> {
        let dst = stream.alloc_zeros::<f32>(len)
            .map_err(|e| anyhow::anyhow!("Clone alloc {len}xf32: {e}"))?;
        let n_bytes = len * std::mem::size_of::<f32>();
        dtod_copy_async(dst.raw_ptr(), ptr, n_bytes, stream, 0, "clone_dueling")
            .map_err(|e| anyhow::anyhow!("DtoD clone: {e}"))?;
        Ok(dst)
    };

    let backing = gpu_weights::DuelingWeightBacking {
        slices: [
            clone_from_ptr(src.w_s1, src.w_s1_len)?,
            clone_from_ptr(src.b_s1, src.b_s1_len)?,
            clone_from_ptr(src.w_s2, src.w_s2_len)?,
            clone_from_ptr(src.b_s2, src.b_s2_len)?,
            clone_from_ptr(src.w_v1, src.w_v1_len)?,
            clone_from_ptr(src.b_v1, src.b_v1_len)?,
            clone_from_ptr(src.w_v2, src.w_v2_len)?,
            clone_from_ptr(src.b_v2, src.b_v2_len)?,
            clone_from_ptr(src.w_a1, src.w_a1_len)?,
            clone_from_ptr(src.b_a1, src.b_a1_len)?,
            clone_from_ptr(src.w_a2, src.w_a2_len)?,
            clone_from_ptr(src.b_a2, src.b_a2_len)?,
        ],
    };
    let ws = backing.weight_set();

    let _ = head_idx; // used for future per-head noise seeding
    unsafe { cudarc::driver::sys::cuStreamSynchronize(stream.cu_stream()); }

    Ok((backing, ws))
}

/// Deep-copy a BranchingWeightSet on GPU via DtoD memcpy.
fn clone_branching_weights(
    src: &BranchingWeightSet,
    stream: &Arc<cudarc::driver::CudaStream>,
    head_idx: usize,
) -> Result<(gpu_weights::BranchingWeightBacking, BranchingWeightSet)> {
    use crate::cuda_pipeline::gpu_dqn_trainer::dtod_copy_async;

    let clone_from_ptr = |ptr: u64, len: usize| -> Result<cudarc::driver::CudaSlice<f32>> {
        let dst = stream.alloc_zeros::<f32>(len)
            .map_err(|e| anyhow::anyhow!("Clone alloc {len}xf32: {e}"))?;
        let n_bytes = len * std::mem::size_of::<f32>();
        dtod_copy_async(dst.raw_ptr(), ptr, n_bytes, stream, 0, "clone_branching")
            .map_err(|e| anyhow::anyhow!("DtoD clone: {e}"))?;
        Ok(dst)
    };

    let backing = gpu_weights::BranchingWeightBacking {
        slices: [
            clone_from_ptr(src.w_bo1, src.w_bo1_len)?,
            clone_from_ptr(src.b_bo1, src.b_bo1_len)?,
            clone_from_ptr(src.w_bo2, src.w_bo2_len)?,
            clone_from_ptr(src.b_bo2, src.b_bo2_len)?,
            clone_from_ptr(src.w_bu1, src.w_bu1_len)?,
            clone_from_ptr(src.b_bu1, src.b_bu1_len)?,
            clone_from_ptr(src.w_bu2, src.w_bu2_len)?,
            clone_from_ptr(src.b_bu2, src.b_bu2_len)?,
            clone_from_ptr(src.w_bg1, src.w_bg1_len)?,
            clone_from_ptr(src.b_bg1, src.b_bg1_len)?,
            clone_from_ptr(src.w_bg2, src.w_bg2_len)?,
            clone_from_ptr(src.b_bg2, src.b_bg2_len)?,
        ],
    };
    let ws = backing.weight_set();

    let _ = head_idx;
    unsafe { cudarc::driver::sys::cuStreamSynchronize(stream.cu_stream()); }

    Ok((backing, ws))
}
