#![allow(unsafe_code)]

//! cuBLAS batched forward pass for the DQN trainer (f32/TF32).
//!
//! Replaces the 1-warp-per-sample fused kernel with cuBLAS matrix multiplications
//! that process the entire batch in a single GEMM call per layer. This raises GPU
//! occupancy from ~1.56% to >60% on H100 by eliminating the per-sample block
//! bottleneck caused by 134 KB shared memory usage.
//!
//! ## TF32 tensor core path (cublasLtMatmul)
//!
//! All forward GEMMs use `cublasLtMatmul` with F32 inputs and
//! `CUBLAS_COMPUTE_32F_FAST_TF32` internal accumulation. cublasLtMatmul passes
//! workspace+stream per-call (no handle state), eliminating the
//! `cublasSetStream` workspace conflict that hangs CUDA Graph mega-capture on
//! H100. On H100 TF32 yields ~500 TFLOPS (vs ~60 for pure f32 CUDA cores).
//!
//! F32 bias+relu kernels handle post-GEMM operations.
//!
//! ## Layout convention
//!
//! All tensors use row-major (C-style) layout:
//! - States:  [B, SD]  — B rows, SD columns
//! - Weights: [out_dim, in_dim]  — stored as GOFF_* in the flat params_buf
//!
//! cuBLAS is column-major, so we use the standard row-major trick:
//! For `C[B, N] = A[B, K] @ W[N, K]^T`:
//!   - cuBLAS sees A as col-major [K, B]   (A^T)
//!   - cuBLAS sees W as col-major [K, N]   (W^T)
//!   - We call: sgemm(CUBLAS_OP_T, CUBLAS_OP_N, N, B, K, 1.0, W, K, A, K, 0.0, C, N)
//!   - Result C is col-major [N, B] which == row-major C[B, N]  ✓
//!
//! ## Network forward order (online and target — same layers, different weights)
//!
//! ```text
//! h_s1[B, SH1] = ReLU(states[B, SD] @ W_s1[SH1, SD]^T + b_s1[SH1])
//! h_s2[B, SH2] = ReLU(h_s1[B, SH1] @ W_s2[SH2, SH1]^T + b_s2[SH2])
//!
//! h_v[B, VH]   = ReLU(h_s2[B, SH2] @ W_v1[VH, SH2]^T + b_v1[VH])
//! v_logits[B, NA] = h_v[B, VH] @ W_v2[NA, VH]^T + b_v2[NA]
//!
//! For each branch d in {0, 1, 2, 3}:
//!   h_bd[B, AH]            = ReLU(h_s2[B, SH2] @ W_bdk_fc[AH, SH2]^T + b_bdk_fc[AH])
//!   adv_logits_d[B, BD*NA] = h_bd[B, AH] @ W_bdk_out[BD*NA, AH]^T + b_bdk_out[BD*NA]
//! ```
//!
//! The C51 distributional loss (log_softmax + cross-entropy with Bellman projection)
//! is computed by a separate NVRTC kernel that reads these activation buffers.
//!
//! ## CUDA Graph compatibility
//!
//! `cublasLtMatmul` is capturable in CUDA 12+ stream-capture mode. Unlike
//! `cublasGemmEx`, it passes workspace+stream per-call, so no `cublasSetStream`
//! state management is needed during multi-stream branch dispatch. Handle
//! creation is NOT capturable and must happen before capture begins.

use std::collections::HashMap;
use std::sync::Arc;
use std::mem::ManuallyDrop;

use cudarc::cublaslt::sys as cublaslt_sys;
use cudarc::cublaslt::result as cublaslt_result;
use cudarc::driver::{CudaEvent, CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg};

use crate::MLError;

// ── Cached GEMM descriptor ──────────────────────────────────────────────────

/// Pre-created cublasLt descriptors + heuristic algo for a fixed GEMM shape.
///
/// Created once at init time for each unique (M, N, K, ldb) shape in the forward
/// pass. Eliminates per-call `cublasLtMatmulDescCreate`, `cublasLtMatrixLayoutCreate`,
/// `cublasLtMatmulPreferenceCreate`, and `cublasLtMatmulAlgoGetHeuristic` overhead
/// (22+ GEMMs per training step).
///
/// The algo heuristic runs ONCE at init with `requestedAlgoCount=3` for better
/// algorithm selection, vs the per-call path which only tried 1.
struct CachedGemmDesc {
    matmul_desc: cublaslt_sys::cublasLtMatmulDesc_t,
    a_layout: cublaslt_sys::cublasLtMatrixLayout_t,
    b_layout: cublaslt_sys::cublasLtMatrixLayout_t,
    c_layout: cublaslt_sys::cublasLtMatrixLayout_t,
    d_layout: cublaslt_sys::cublasLtMatrixLayout_t,
    algo: cublaslt_sys::cublasLtMatmulAlgo_t,
}

/// SAFETY: cublasLt descriptors are opaque pointers owned exclusively by
/// `CublasGemmSet`, which is single-threaded in practice (owned by `GpuDqnTrainer`).
unsafe impl Send for CachedGemmDesc {}
unsafe impl Sync for CachedGemmDesc {}

impl Drop for CachedGemmDesc {
    fn drop(&mut self) {
        unsafe {
            let _ = cublaslt_result::destroy_matrix_layout(self.d_layout);
            let _ = cublaslt_result::destroy_matrix_layout(self.c_layout);
            let _ = cublaslt_result::destroy_matrix_layout(self.b_layout);
            let _ = cublaslt_result::destroy_matrix_layout(self.a_layout);
            let _ = cublaslt_result::destroy_matmul_desc(self.matmul_desc);
            // algo is a value type (cublasLtMatmulAlgo_t), no destroy needed.
            // matmul_pref is destroyed after heuristic query at init time.
        }
    }
}

/// Key for the cached descriptor map: (n=out_dim, b=batch, k=in_dim, ldb).
/// All forward GEMMs use TRANSA=T, TRANSB=N so transpose flags are implicit.
type FwdGemmKey = (usize, usize, usize, usize);

// ── Batched forward context ───────────────────────────────────────────────────

/// cuBLAS-based batched forward GEMM set for the branching dueling DQN.
///
/// Holds an `Arc<PerStreamCublasHandles>` (cuBLAS + cublasLt handles, workspaces,
/// bias kernels) and the forward-specific resources: branch streams, per-branch
/// workspaces, cached GEMM descriptors, and network dimensions.
///
/// The handle is set to the trainer's stream so all GEMMs execute in order.
#[allow(missing_debug_implementations)] // handle is not Debug
pub struct CublasGemmSet {
    /// Shared cuBLAS + cublasLt handle (owned via Arc, shared with backward pass
    /// and experience collector for deterministic single-handle training).
    handle: Arc<super::shared_cublas_handle::PerStreamCublasHandles>,

    // ── Network dimensions (baked at construction) ──
    batch_size: usize,
    state_dim: usize,
    /// pad128(state_dim): CUTLASS K-tile alignment for the B-matrix (input states).
    /// Used as `ldb` for the first-layer GemmEx to prevent 128-element OOB reads.
    pub(crate) state_dim_padded: usize,
    /// First shared layer input dimension. Equals state_dim unless bottleneck is active,
    /// where it is bottleneck_dim + portfolio_dim (< state_dim). Used for h_s1 GEMM K.
    s1_input_dim: usize,
    /// Leading dimension for s1 input buffer. pad128(s1_input_dim) when bottleneck
    /// is active, pad128(state_dim) otherwise. MUST be K-tile aligned for
    /// cuBLAS graph-compatible algorithm selection on Hopper.
    pub(crate) s1_ldb: usize,
    shared_h1: usize,
    shared_h2: usize,
    value_h: usize,
    adv_h: usize,
    num_atoms: usize,
    branch_0_size: usize,
    branch_1_size: usize,
    branch_2_size: usize,
    branch_3_size: usize,

    /// Magnitude branch (d==1) FC input dimension: shared_h2 + branch_0_size
    /// (concat of h_s2 and per-direction Q_dir). Direction-conditioned, so
    /// the trailing-stride width derives from `branch_0_size` (4 in production:
    /// S/H/L/F). Used when mag_concat_ptr != 0.
    mag_concat_dim: usize,

    /// Order/urgency branch (d∈{2,3}) FC input dimension: shared_h2 + 3
    /// (concat of vsn_masked and 3 OFI features per branch). OFI-conditioned,
    /// so the trailing-stride width is fixed at 3 (matches `concat_ofi_features`
    /// kernel writes). Used when ord_concat_ptr/urg_concat_ptr != 0.
    ofi_concat_dim: usize,

    // ── Multi-stream branch dispatch ──
    /// 4 forked CUDA streams for parallel advantage branch execution.
    /// Each branch (exposure, order, urgency, ...) submits GEMMs to its own stream,
    /// then the main stream joins all four before consuming the logits.
    branch_streams: [Arc<CudaStream>; 4],

    /// Per-branch cublasLt workspace buffers — eliminates workspace contention
    /// when 4 branch streams run GEMMs in parallel. Each 32 MB.
    /// Main stream GEMMs (trunk, value head) use `lt_workspace_ptr` (shared).
    _branch_workspace_bufs: [CudaSlice<u8>; 4],
    branch_workspace_ptrs: [u64; 4],

    // ── Pre-allocated CUDA events for fork-join synchronization ──
    /// Reused across forward passes — avoids cuEventCreate/Destroy per call.
    trunk_done_event: CudaEvent,
    branch_done_events: [CudaEvent; 4],

    // ── Cached GEMM descriptors (created once at init, reused per-call) ──
    /// Map from (n, batch, k, ldb) → pre-created descriptors + algo.
    /// All forward GEMMs share TRANSA=T, TRANSB=N.
    gemm_cache: HashMap<FwdGemmKey, CachedGemmDesc>,
    /// Cached descriptors with RELU_BIAS epilogue (for hidden layers).
    /// Bias pointer is set dynamically per-call. Same key as gemm_cache.
    gemm_cache_relu_bias: HashMap<FwdGemmKey, CachedGemmDesc>,
    /// Cached descriptors with BIAS epilogue (for linear output layers — value
    /// logits, advantage logits, GRN Linear_a/Linear_b, VSN linear2, GLU value/gate).
    /// Bias pointer is set dynamically per-call. Same key as `gemm_cache`.
    /// Phase H Site 3: fuses GEMM + bias-add into one cublasLtMatmul call,
    /// eliminating the standalone `add_bias_f32_kernel` launch (one less kernel
    /// + memory round-trip per BIAS-only output layer).
    gemm_cache_bias: HashMap<FwdGemmKey, CachedGemmDesc>,
    /// Cached descriptors with RELU_AUX_BIAS epilogue (writes the dReLU bit-mask
    /// for backward fusion via DRELU_BGRAD). Currently populated only for the
    /// online value-FC layer's shape — the only site whose backward uses
    /// `EPILOGUE_DRELU_BGRAD` (Phase H Site 2). Bias and AUX pointers are set
    /// dynamically per-call.
    gemm_cache_relu_aux_bias: HashMap<FwdGemmKey, CachedGemmDesc>,

    // ── Value-FC dReLU bit-mask aux buffer (online network only) ──
    /// Bit-packed mask buffer for the value-FC layer's ReLU activation. Written
    /// by the online forward via `CUBLASLT_EPILOGUE_RELU_AUX_BIAS` (1 bit per
    /// activation element). Read by `CublasBackwardSet::backward_fc_layer_drelu_bgrad`
    /// via `CUBLASLT_EPILOGUE_DRELU_BGRAD` to gate dy and emit `b_v1` grad in
    /// one fused cublasLt call (saves the standalone `relu_mask` + `bias_grad`
    /// reduce launches at this site). Allocated as `CudaSlice<u8>` with size
    /// `value_h_aux_ld * batch_size / 8` bytes; the leading dim
    /// `value_fc_relu_mask_aux_ld` is rounded up to the nearest 128-bit (16-byte)
    /// boundary per cuBLAS alignment requirements for ReLU mask aux buffers.
    _value_fc_relu_mask_buf: CudaSlice<u8>,
    /// Raw device pointer for the value-FC ReLU mask aux buffer.
    /// Accessed externally via the `value_fc_relu_mask_ptr()` accessor.
    value_fc_relu_mask_ptr_raw: u64,
    /// AUX leading dim for the value-FC ReLU mask, in elements (= bits).
    /// Equals `value_h` rounded up to a multiple of 128 for cuBLAS alignment.
    /// Accessed externally via the `value_fc_relu_mask_aux_ld()` accessor.
    value_fc_relu_mask_aux_ld_raw: usize,
    /// Trainer-registered "main online save_h_v" pointer. The AUX_BIAS path
    /// is taken ONLY when `h_v_ptr` argument matches this — otherwise the
    /// forward falls back to plain `RELU_BIAS` so auxiliary forwards
    /// (causal-intervention, IQN, DDQN, etc., which target scratch h_v
    /// buffers) don't clobber the mask the backward reads. Wired post-
    /// construction via `wire_value_fc_save_h_v`.
    value_fc_save_h_v_marker: u64,

    // ── VSN bottleneck + GLU gating (wired post-construction via wire_vsn_glu) ──
    /// Variable Selection Network bottleneck kernel (SH2→R=16→SH2).
    vsn_kernel: Option<CudaFunction>,
    /// KAN spline gate kernel: output = kan_spline(gate_pre, coeffs, resid_w) * value.
    kan_gate_combine_kernel: Option<CudaFunction>,
    /// Strided scatter: vsn_masked [B, SH2] → first SH2 cols of mag_concat [B, SH2+3].
    strided_scatter_kernel: Option<CudaFunction>,
    /// VSN masked output scratch [B, SH2].
    vsn_masked_ptr: u64,
    /// Gate pre-activation scratch [B, AH] per branch.
    glu_gate_pre_ptrs: [u64; 4],
    /// Value path scratch [B, AH] per branch (holds bias-added GEMM output before GLU).
    glu_value_ptrs: [u64; 4],
    /// VSN bottleneck rank (R=16).
    vsn_rank: usize,
    /// Order branch (d=2) OFI concat buffer ptr [B, SH2+3] (0 if not wired).
    ord_concat_ptr: u64,
    /// Urgency branch (d=3) OFI concat buffer ptr [B, SH2+3] (0 if not wired).
    urg_concat_ptr: u64,

    // ── GRN trunk (Plan 4 Task 2c.3b — encoder forward swap) ────────────
    /// h_s1 GRN block (online network, save-for-backward enabled).
    /// Owns its 5 save-for-backward buffers + 2 LN partial-reduction
    /// scratch buffers + 8 kernel handles. SD → SH1 with `Linear_residual`
    /// projection (dimensions differ).
    grn_h_s1_online: super::gpu_grn::GrnBlock,
    /// h_s2 GRN block (online network, save-for-backward enabled).
    /// SH1 → SH2 with identity residual (SH1 == SH2 in current spec).
    grn_h_s2_online: super::gpu_grn::GrnBlock,
    /// h_s1 GRN block (target network, save buffers allocated but unused).
    /// Target forward is inference-only; the dedicated instance prevents
    /// the online network's save state from being clobbered when the
    /// target forward runs after the online forward in the same step.
    grn_h_s1_target: super::gpu_grn::GrnBlock,
    /// h_s2 GRN block (target network, save buffers allocated but unused).
    grn_h_s2_target: super::gpu_grn::GrnBlock,

    // ── GRN scratch buffers (per-block, sized at construction) ─────────
    // Owned as raw `CudaSlice<u8>` (byte buffers) because the encoder is
    // graph-captured and reads them via raw u64 pointers. The lifetime is
    // bound to `Self`, ensuring the GPU allocations outlive every forward
    // launch.

    /// h_s1 Linear_a output [B, SH1] — also holds post-ELU activation
    /// after the in-place ELU launch. Online instance.
    _grn_h_s1_linear_a_buf: CudaSlice<f32>,
    grn_h_s1_linear_a_ptr: u64,
    /// h_s1 Linear_b output [B, 2*SH1]. Online instance.
    _grn_h_s1_linear_b_buf: CudaSlice<f32>,
    grn_h_s1_linear_b_ptr: u64,
    /// h_s1 Linear_residual output [B, SH1]. Online instance.
    _grn_h_s1_residual_buf: CudaSlice<f32>,
    grn_h_s1_residual_ptr: u64,
    /// h_s1 GLU output [B, SH1]. Online instance.
    _grn_h_s1_glu_buf: CudaSlice<f32>,
    grn_h_s1_glu_ptr: u64,
    /// h_s2 Linear_a output [B, SH2] — also post-ELU. Online instance.
    _grn_h_s2_linear_a_buf: CudaSlice<f32>,
    grn_h_s2_linear_a_ptr: u64,
    /// h_s2 Linear_b output [B, 2*SH2]. Online instance.
    _grn_h_s2_linear_b_buf: CudaSlice<f32>,
    grn_h_s2_linear_b_ptr: u64,
    /// h_s2 GLU output [B, SH2]. Online instance.
    _grn_h_s2_glu_buf: CudaSlice<f32>,
    grn_h_s2_glu_ptr: u64,

    /// h_s1 Linear_a output [B, SH1] — target instance scratch.
    _grn_h_s1_linear_a_buf_tg: CudaSlice<f32>,
    grn_h_s1_linear_a_ptr_tg: u64,
    /// h_s1 Linear_b output [B, 2*SH1] — target instance scratch.
    _grn_h_s1_linear_b_buf_tg: CudaSlice<f32>,
    grn_h_s1_linear_b_ptr_tg: u64,
    /// h_s1 Linear_residual output [B, SH1] — target instance scratch.
    _grn_h_s1_residual_buf_tg: CudaSlice<f32>,
    grn_h_s1_residual_ptr_tg: u64,
    /// h_s1 GLU output [B, SH1] — target instance scratch.
    _grn_h_s1_glu_buf_tg: CudaSlice<f32>,
    grn_h_s1_glu_ptr_tg: u64,
    /// h_s2 Linear_a output [B, SH2] — target instance scratch.
    _grn_h_s2_linear_a_buf_tg: CudaSlice<f32>,
    grn_h_s2_linear_a_ptr_tg: u64,
    /// h_s2 Linear_b output [B, 2*SH2] — target instance scratch.
    _grn_h_s2_linear_b_buf_tg: CudaSlice<f32>,
    grn_h_s2_linear_b_ptr_tg: u64,
    /// h_s2 GLU output [B, SH2] — target instance scratch.
    _grn_h_s2_glu_buf_tg: CudaSlice<f32>,
    grn_h_s2_glu_ptr_tg: u64,

    // ── Element-wise SAXPY kernel (Plan 4 Task 2c.3c.2) ──────────────────
    /// `dqn_saxpy_f32_kernel(y, x, alpha, n)` — computes `y[i] += alpha * x[i]`.
    /// Used by GRN trunk backward to accumulate the h_s2 identity-residual
    /// gradient (`d_h_s1 += d_pre_ln_h_s2`) after Linear_a's backward
    /// overwrites `d_h_s1`.
    saxpy_kernel: CudaFunction,

    // ── Plan 4 Task 1B-iii: VSN feature-selection forward orchestrator ──
    /// `vsn_softmax_and_gate_forward(state_in, logits, B, state_dim_padded,
    /// num_groups, group_begins, group_ends, mask_out, gated_state_out)` —
    /// numerically-stable softmax over `[B, num_groups]` logits, per-feature
    /// multiplicative gate of `state_in` rows by the resulting mask. Loaded
    /// from `VSN_FEATURE_SELECTION_CUBIN` (1B-i kernel landing). Producer of
    /// `vsn_mask_buf` consumed both by `vsn_mask_ema_update` (1B-i ISV kernel)
    /// and by 1B-iv's softmax-Jacobian backward.
    vsn_softmax_and_gate_fwd_kernel: CudaFunction,
    /// Plan 4 Task 1B-iv consumer: `vsn_softmax_and_gate_backward(
    /// d_gated_state, state, mask, B, state_dim_padded, num_groups,
    /// group_begins, group_ends, d_logits, d_state)` implements the FULL
    /// softmax-over-groups Jacobian. Wired into the backward chain as the
    /// final step of `launch_cublas_backward_to` (post-bottleneck Linear bwd
    /// + portfolio passthrough assembly), writing the 24 VSN dW/dB slots at
    /// grad_buf indices [95..119).
    vsn_softmax_and_gate_bwd_kernel: CudaFunction,
    /// Plan 4 Task 1B-iii: scatter handle for column-g writes of per-group
    /// `Linear_2_g` outputs into the assembled `vsn_logits_buf [B, num_groups]`.
    /// Reuses the existing `strided_scatter` kernel from `experience_kernels.cu`
    /// (already wired for VSN bottleneck via `wire_vsn_glu`); the trainer
    /// passes the same kernel handle through `wire_vsn_scatter` after
    /// construction. Kept Optional so non-trainer consumers (e.g. the
    /// experience collector's CublasGemmSet instance) that never call
    /// `vsn_forward` don't need to wire it. `vsn_forward` requires it to be
    /// `Some`; a `debug_assert!` flags the absence loudly.
    vsn_logit_scatter_kernel: Option<CudaFunction>,
}

/// Backward-compatible type alias — existing callers (`GpuDqnTrainer`,
/// `GpuExperienceCollector`, etc.) keep using `CublasForward` until Task 4
/// wires the shared handle through all constructors.
pub type CublasForward = CublasGemmSet;

impl CublasGemmSet {
    /// Create a new cuBLAS forward GEMM set backed by a shared handle.
    ///
    /// The shared handle owns the cuBLAS + cublasLt handles, workspaces, and
    /// bias kernels. This constructor only creates forward-specific resources:
    /// branch streams, per-branch workspaces, cached GEMM descriptors, events,
    /// and stores the network dimensions.
    pub fn new(
        shared: Arc<super::shared_cublas_handle::PerStreamCublasHandles>,
        batch_size: usize,
        state_dim: usize,
        shared_h1: usize,
        shared_h2: usize,
        value_h: usize,
        adv_h: usize,
        num_atoms: usize,
        branch_0_size: usize,
        branch_1_size: usize,
        branch_2_size: usize,
        branch_3_size: usize,
        s1_input_dim: usize,
    ) -> Result<Self, MLError> {
        let stream = &shared.stream;

        // ── Fork 4 branch streams for parallel advantage head dispatch ──
        let branch_streams = [
            stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 0: {e}")))?,
            stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 1: {e}")))?,
            stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 2: {e}")))?,
            stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 3: {e}")))?,
        ];

        // Pre-register each branch stream in the per-stream cublasLt registry
        // (Option C). Must happen BEFORE any CUDA Graph capture — first-access
        // workspace allocation would otherwise fail inside capture with
        // `CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED`.
        for bs in &branch_streams {
            shared.pre_register_stream(bs)?;
        }

        // ── Allocate per-branch cublasLt workspace buffers (32 MB each) ──
        // Eliminates workspace contention when 4 branch streams run GEMMs in parallel.
        let branch_ws_size: usize = 32 * 1024 * 1024;
        let branch_ws_buf_0 = stream.alloc_zeros::<u8>(branch_ws_size)
            .map_err(|e| MLError::ModelError(format!("branch workspace 0 alloc: {e}")))?;
        let branch_ws_buf_1 = stream.alloc_zeros::<u8>(branch_ws_size)
            .map_err(|e| MLError::ModelError(format!("branch workspace 1 alloc: {e}")))?;
        let branch_ws_buf_2 = stream.alloc_zeros::<u8>(branch_ws_size)
            .map_err(|e| MLError::ModelError(format!("branch workspace 2 alloc: {e}")))?;
        let branch_ws_buf_3 = stream.alloc_zeros::<u8>(branch_ws_size)
            .map_err(|e| MLError::ModelError(format!("branch workspace 3 alloc: {e}")))?;
        let branch_workspace_ptrs = [
            branch_ws_buf_0.raw_ptr(),
            branch_ws_buf_1.raw_ptr(),
            branch_ws_buf_2.raw_ptr(),
            branch_ws_buf_3.raw_ptr(),
        ];
        let branch_workspace_bufs = [branch_ws_buf_0, branch_ws_buf_1, branch_ws_buf_2, branch_ws_buf_3];

        // ── Pre-allocate CUDA events for fork-join (reused across all forward passes) ──
        let ctx = stream.context();
        let trunk_done_event = ctx.new_event(None)
            .map_err(|e| MLError::DeviceError(format!("alloc trunk_done_event: {e}")))?;
        let branch_done_events = [
            ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("alloc branch_done_event 0: {e}")))?,
            ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("alloc branch_done_event 1: {e}")))?,
            ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("alloc branch_done_event 2: {e}")))?,
            ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("alloc branch_done_event 3: {e}")))?,
        ];

        // ── Pre-create GEMM descriptors for all unique forward shapes ──
        // All forward GEMMs: TRANSA=T, TRANSB=N, f32 types, FAST_TF32 compute.
        // Shapes keyed by (n=out_dim, b=batch, k=in_dim, ldb).
        let state_dim_padded = (state_dim + 127) & !127;
        // s1_ldb: when bottleneck active, bn_concat has no K-tile padding.
        let s1_ldb = if s1_input_dim < state_dim {
            s1_input_dim
        } else {
            state_dim_padded
        };
        let branch_sizes = [branch_0_size, branch_1_size, branch_2_size, branch_3_size];

        let lt_raw_handle = shared.lt_handle.0;
        let lt_ws_size = shared.lt_workspace_size;

        let mut unique_shapes: Vec<FwdGemmKey> = Vec::new();
        // h_s1 GRN Linear_a / Linear_residual: M=shared_h1, N=batch, K=s1_input_dim, ldb=s1_ldb
        unique_shapes.push((shared_h1, batch_size, s1_input_dim, s1_ldb));
        // h_s1 GRN Linear_b: M=2*shared_h1, N=batch, K=shared_h1, ldb=shared_h1
        unique_shapes.push((2 * shared_h1, batch_size, shared_h1, shared_h1));
        // h_s2 GRN Linear_a: M=shared_h2, N=batch, K=shared_h1, ldb=shared_h1
        unique_shapes.push((shared_h2, batch_size, shared_h1, shared_h1));
        // h_s2 GRN Linear_b: M=2*shared_h2, N=batch, K=shared_h2, ldb=shared_h2
        unique_shapes.push((2 * shared_h2, batch_size, shared_h2, shared_h2));
        // h_v: M=value_h, N=batch, K=shared_h2, ldb=shared_h2
        unique_shapes.push((value_h, batch_size, shared_h2, shared_h2));
        // v_logits: M=num_atoms, N=batch, K=value_h, ldb=value_h
        unique_shapes.push((num_atoms, batch_size, value_h, value_h));
        // h_bd (×4): M=adv_h, N=batch, K=shared_h2, ldb=shared_h2
        // (Direction d==0 uses K=shared_h2+1 post-SP14 — see direction-Q wider shape below.)
        unique_shapes.push((adv_h, batch_size, shared_h2, shared_h2));
        // h_bd magnitude (d==1) wider input (direction-conditioned):
        //   M=adv_h, N=batch, K=shared_h2+branch_0_size, ldb=shared_h2+branch_0_size
        unique_shapes.push((adv_h, batch_size, shared_h2 + branch_0_size, shared_h2 + branch_0_size));
        // h_bd direction (d==0) wider input (SP14 EGF aux→Q wire):
        //   M=adv_h, N=batch, K=shared_h2+1, ldb=shared_h2+1 — extra column carries
        //   `aux_softmax_diff = up_prob - down_prob ∈ [-1, +1]` from the aux next-bar
        //   classifier head. Per SP14 plan §B.9, the direction Q-head's first FC
        //   consumes `sp14_dir_qaux_concat_scratch [B, SH2+1]` instead of raw h_s2.
        unique_shapes.push((adv_h, batch_size, shared_h2 + 1, shared_h2 + 1));
        // h_bd order/urgency (d∈{2,3}) wider input (OFI-conditioned, 3 features per branch):
        //   M=adv_h, N=batch, K=shared_h2+3, ldb=shared_h2+3
        unique_shapes.push((adv_h, batch_size, shared_h2 + 3, shared_h2 + 3));
        // adv_logits (×4): M=branch_k*num_atoms, N=batch, K=adv_h, ldb=adv_h
        for &bs in &branch_sizes {
            unique_shapes.push((bs * num_atoms, batch_size, adv_h, adv_h));
        }
        unique_shapes.sort();
        unique_shapes.dedup();

        let mut gemm_cache = HashMap::new();
        for &(n, b, k, ldb) in &unique_shapes {
            let desc = create_cached_fwd_gemm_desc(lt_raw_handle, n, b, k, ldb, lt_ws_size)?;
            gemm_cache.insert((n, b, k, ldb), desc);
        }

        // Create RELU_BIAS epilogue variants for hidden layers.
        // These fuse GEMM + bias-add + ReLU into one kernel, eliminating separate launches.
        const VSN_HIDDEN_RB: usize = super::gpu_dqn_trainer::VSN_HIDDEN_DIM;
        let mut relu_bias_shapes: Vec<FwdGemmKey> = vec![
            (shared_h1, batch_size, s1_input_dim, s1_ldb), // h_s1
            (shared_h2, batch_size, shared_h1, shared_h1),        // h_s2
            (value_h, batch_size, shared_h2, shared_h2),          // h_v
            (adv_h, batch_size, shared_h2, shared_h2),            // h_bd (×4) — d ∈ {0..3} legacy / VSN-disabled path
            (adv_h, batch_size, shared_h2 + branch_0_size, shared_h2 + branch_0_size), // h_bd magnitude wider input (direction-conditioned)
            (adv_h, batch_size, shared_h2 + 1, shared_h2 + 1), // h_bd direction wider input (SP14 EGF aux→Q wire — d==0 first FC)
            (adv_h, batch_size, shared_h2 + 3, shared_h2 + 3), // h_bd order/urgency wider input (OFI-conditioned)
        ];
        // VSN feature-selection Linear_1 (×6 groups, group_dim varies).
        // Phase H Site 3 — fuses per-group GEMM + bias + ReLU into one kernel.
        // ldb is the full state_dim_padded so cuBLAS reads the per-group slice
        // via pointer offset on the same column-major view as `vsn_forward`.
        for &(gb, ge) in &ml_core::state_layout::FEATURE_GROUP_RANGES {
            let group_dim = ge - gb;
            relu_bias_shapes.push((VSN_HIDDEN_RB, batch_size, group_dim, state_dim_padded));
        }
        relu_bias_shapes.sort();
        relu_bias_shapes.dedup();
        let mut gemm_cache_relu_bias = HashMap::new();
        for &(n, b, k, ldb) in &relu_bias_shapes {
            match create_cached_fwd_gemm_desc_relu_bias(lt_raw_handle, n, b, k, ldb, lt_ws_size) {
                Ok(desc) => { gemm_cache_relu_bias.insert((n, b, k, ldb), desc); }
                Err(e) => {
                    tracing::warn!("RELU_BIAS epilogue not available for ({n},{b},{k},{ldb}): {e}");
                    // Fall back to separate kernels — no entry in cache
                }
            }
        }

        // ── Phase H Site 3: BIAS-only epilogue variants for linear output layers ──
        //
        // Fuses GEMM + bias-add into a single cublasLtMatmul call (no ReLU,
        // because these layers are either output logits or feed into a separate
        // gating activation downstream — KAN/GLU/softmax). Each fused descriptor
        // eliminates one `add_bias_f32_kernel` launch + one memory round-trip
        // per call.
        //
        // Shapes covered:
        //   - VSN linear2 (group logit producer): (1, B, VSN_HIDDEN, VSN_HIDDEN)
        //   - GRN Linear_a / Linear_b (h_s1, h_s2 trunks)
        //   - Value-head v_logits: (num_atoms, B, value_h, value_h)
        //   - Branch adv_logits (×4 shapes — branch_size_d * num_atoms output dim)
        //   - Branch h_bd / GLU value & gate (no-ReLU when GLU-gated downstream):
        //     same K dims as RELU_BIAS but BIAS-only epilogue for the GLU path.
        const VSN_HIDDEN: usize = super::gpu_dqn_trainer::VSN_HIDDEN_DIM;
        let mut bias_shapes: Vec<FwdGemmKey> = vec![
            // VSN linear2 — projects [B, VSN_HIDDEN] → [B, 1] per group.
            (1, batch_size, VSN_HIDDEN, VSN_HIDDEN),
            // GRN h_s1 Linear_a / Linear_residual share K=s1_input_dim, ldb=s1_ldb.
            // (Linear_residual has no bias — only Linear_a uses the BIAS epilogue.)
            (shared_h1, batch_size, s1_input_dim, s1_ldb),
            // GRN h_s1 Linear_b: M=2*shared_h1, K=shared_h1.
            (2 * shared_h1, batch_size, shared_h1, shared_h1),
            // GRN h_s2 Linear_a: M=shared_h2, K=shared_h1.
            (shared_h2, batch_size, shared_h1, shared_h1),
            // GRN h_s2 Linear_b: M=2*shared_h2, K=shared_h2.
            (2 * shared_h2, batch_size, shared_h2, shared_h2),
            // Value-head v_logits: M=num_atoms, K=value_h.
            (num_atoms, batch_size, value_h, value_h),
            // Branch h_bd / GLU value & gate — BIAS only (GLU/KAN provides gating):
            (adv_h, batch_size, shared_h2, shared_h2),
            (adv_h, batch_size, shared_h2 + branch_0_size, shared_h2 + branch_0_size),
            (adv_h, batch_size, shared_h2 + 1, shared_h2 + 1), // SP14 EGF: direction (d==0) wider input — VSN/GLU value+gate GEMMs
            (adv_h, batch_size, shared_h2 + 3, shared_h2 + 3),
        ];
        // Branch adv_logits (×4 distinct branch sizes): M=bs*num_atoms, K=adv_h.
        for &bs in &branch_sizes {
            bias_shapes.push((bs * num_atoms, batch_size, adv_h, adv_h));
        }
        bias_shapes.sort();
        bias_shapes.dedup();
        let mut gemm_cache_bias = HashMap::new();
        for &(n, b, k, ldb) in &bias_shapes {
            match create_cached_fwd_gemm_desc_bias(lt_raw_handle, n, b, k, ldb, lt_ws_size) {
                Ok(desc) => { gemm_cache_bias.insert((n, b, k, ldb), desc); }
                Err(e) => {
                    tracing::warn!("BIAS epilogue not available for ({n},{b},{k},{ldb}): {e}");
                    // Fall back to separate kernels — no entry in cache
                }
            }
        }

        // ── Phase H Site 2: RELU_AUX_BIAS variant for the value-FC layer ──
        //
        // Writes the dReLU bit-mask aux buffer alongside the post-ReLU output.
        // The mask is consumed by `CublasBackwardSet::backward_fc_layer_drelu_bgrad`
        // via `EPILOGUE_DRELU_BGRAD`, fusing the standalone `relu_mask` +
        // `bias_grad` reduce kernels into the value-output dX GEMM (saves
        // 2 launches per online backward pass).
        //
        // AUX_LD is in elements (= bits for ReLU mask). cuBLAS requires the
        // AUX_LD to be a multiple of 128 bits (16 bytes). value_h is normally
        // 128 (already aligned); the rounded form keeps the contract regardless.
        let value_fc_relu_mask_aux_ld: usize = (value_h + 127) & !127;
        let mut gemm_cache_relu_aux_bias = HashMap::new();
        let aux_bias_shape: FwdGemmKey = (value_h, batch_size, shared_h2, shared_h2);
        match create_cached_fwd_gemm_desc_relu_aux_bias(
            lt_raw_handle,
            aux_bias_shape.0, aux_bias_shape.1, aux_bias_shape.2, aux_bias_shape.3,
            value_fc_relu_mask_aux_ld,
            lt_ws_size,
        ) {
            Ok(desc) => { gemm_cache_relu_aux_bias.insert(aux_bias_shape, desc); }
            Err(e) => {
                tracing::warn!(
                    "RELU_AUX_BIAS epilogue not available for value-FC ({},{},{},{}): {e}",
                    aux_bias_shape.0, aux_bias_shape.1, aux_bias_shape.2, aux_bias_shape.3,
                );
                // Fall back to RELU_BIAS path — backward will use the standalone
                // relu_mask + bias_grad sequence at this site (slow path).
            }
        }

        // Allocate the bit-packed mask buffer. Stride per column = aux_ld bits =
        // aux_ld/8 bytes; total bytes = aux_ld/8 * batch_size. The mask only
        // covers the online value-FC forward (target/ensemble paths still use
        // the plain RELU_BIAS epilogue and the manual `relu_mask` kernel).
        let value_fc_mask_bytes = (value_fc_relu_mask_aux_ld / 8) * batch_size;
        let value_fc_relu_mask_buf = stream.alloc_zeros::<u8>(value_fc_mask_bytes)
            .map_err(|e| MLError::ModelError(format!("value-FC ReLU mask aux buf alloc [{value_fc_mask_bytes} bytes]: {e}")))?;
        let value_fc_relu_mask_ptr = value_fc_relu_mask_buf.raw_ptr();

        // ── GRN trunk allocation (Plan 4 Task 2c.3b) ────────────────
        // Allocate two GrnBlock instances per network (online + target) so
        // the target forward never clobbers the online's save-for-backward
        // state. Each instance owns its 5 save buffers + 2 LN partial-
        // reduction scratch buffers. h_s1 needs `Linear_residual`
        // (SD != SH1); h_s2 uses identity residual (SH1 == SH2).
        let stream_for_grn = Arc::clone(&shared.stream);
        let grn_h_s1_online = super::gpu_grn::GrnBlock::new(
            Arc::clone(&stream_for_grn), shared_h1, batch_size, /* has_residual_projection = */ true,
        )?;
        let grn_h_s2_online = super::gpu_grn::GrnBlock::new(
            Arc::clone(&stream_for_grn), shared_h2, batch_size, /* has_residual_projection = */ false,
        )?;
        let grn_h_s1_target = super::gpu_grn::GrnBlock::new(
            Arc::clone(&stream_for_grn), shared_h1, batch_size, /* has_residual_projection = */ true,
        )?;
        let grn_h_s2_target = super::gpu_grn::GrnBlock::new(
            Arc::clone(&stream_for_grn), shared_h2, batch_size, /* has_residual_projection = */ false,
        )?;

        // ── GRN scratch buffers (online + target) ────────────────────
        let alloc_grn = |n: usize, label: &str| -> Result<CudaSlice<f32>, MLError> {
            stream_for_grn
                .alloc_zeros::<f32>(n)
                .map_err(|e| MLError::ModelError(format!("grn scratch {label}: {e}")))
        };

        let h_s1_linear_a_buf = alloc_grn(batch_size * shared_h1,         "h_s1_linear_a")?;
        let h_s1_linear_b_buf = alloc_grn(batch_size * 2 * shared_h1,     "h_s1_linear_b")?;
        let h_s1_residual_buf = alloc_grn(batch_size * shared_h1,         "h_s1_residual")?;
        let h_s1_glu_buf      = alloc_grn(batch_size * shared_h1,         "h_s1_glu")?;
        let h_s2_linear_a_buf = alloc_grn(batch_size * shared_h2,         "h_s2_linear_a")?;
        let h_s2_linear_b_buf = alloc_grn(batch_size * 2 * shared_h2,     "h_s2_linear_b")?;
        let h_s2_glu_buf      = alloc_grn(batch_size * shared_h2,         "h_s2_glu")?;

        let h_s1_linear_a_buf_tg = alloc_grn(batch_size * shared_h1,      "h_s1_linear_a_tg")?;
        let h_s1_linear_b_buf_tg = alloc_grn(batch_size * 2 * shared_h1,  "h_s1_linear_b_tg")?;
        let h_s1_residual_buf_tg = alloc_grn(batch_size * shared_h1,      "h_s1_residual_tg")?;
        let h_s1_glu_buf_tg      = alloc_grn(batch_size * shared_h1,      "h_s1_glu_tg")?;
        let h_s2_linear_a_buf_tg = alloc_grn(batch_size * shared_h2,      "h_s2_linear_a_tg")?;
        let h_s2_linear_b_buf_tg = alloc_grn(batch_size * 2 * shared_h2,  "h_s2_linear_b_tg")?;
        let h_s2_glu_buf_tg      = alloc_grn(batch_size * shared_h2,      "h_s2_glu_tg")?;

        // ── Load SAXPY kernel (Plan 4 Task 2c.3c.2) ──────────────────
        // `dqn_saxpy_f32_kernel(y, x, alpha, n)` is shared with the
        // experience collector's IQR/ensemble-variance Q-bonus paths.
        // Loaded here so the GRN trunk backward (h_s2 identity residual
        // accumulation) can call `saxpy_inplace` without plumbing another
        // kernel handle through the trainer's wire-up sequence.
        let saxpy_kernel = {
            let saxpy_module = stream_for_grn
                .context()
                .load_cubin(super::gpu_dqn_trainer::DQN_UTILITY_CUBIN.to_vec())
                .map_err(|e| MLError::ModelError(format!("saxpy cubin load: {e}")))?;
            saxpy_module
                .load_function("dqn_saxpy_f32_kernel")
                .map_err(|e| MLError::ModelError(format!("dqn_saxpy_f32_kernel load: {e}")))?
        };

        // ── Plan 4 Task 1B-iii: load VSN forward + backward kernels ─────
        // Forward kernel `vsn_softmax_and_gate_forward` is consumed by the
        // `vsn_forward` orchestrator below; backward kernel
        // `vsn_softmax_and_gate_backward` is loaded here so the field is in
        // place for 1B-iv (which only needs to wire the consumer call site).
        // The cubin module is reused across both function lookups — one cubin
        // load, two kernel handles. Mirrors the `h_s2_rms_ema_update` pattern
        // from `gpu_dqn_trainer.rs`.
        let (vsn_softmax_and_gate_fwd_kernel, vsn_softmax_and_gate_bwd_kernel) = {
            let module = stream_for_grn
                .context()
                .load_cubin(super::gpu_dqn_trainer::VSN_FEATURE_SELECTION_CUBIN.to_vec())
                .map_err(|e| MLError::ModelError(format!("vsn_feature_selection cubin load: {e}")))?;
            let fwd = module
                .load_function("vsn_softmax_and_gate_forward")
                .map_err(|e| MLError::ModelError(format!("vsn_softmax_and_gate_forward load: {e}")))?;
            let bwd = module
                .load_function("vsn_softmax_and_gate_backward")
                .map_err(|e| MLError::ModelError(format!("vsn_softmax_and_gate_backward load: {e}")))?;
            (fwd, bwd)
        };

        let h_s1_linear_a_ptr     = h_s1_linear_a_buf.raw_ptr();
        let h_s1_linear_b_ptr     = h_s1_linear_b_buf.raw_ptr();
        let h_s1_residual_ptr     = h_s1_residual_buf.raw_ptr();
        let h_s1_glu_ptr          = h_s1_glu_buf.raw_ptr();
        let h_s2_linear_a_ptr     = h_s2_linear_a_buf.raw_ptr();
        let h_s2_linear_b_ptr     = h_s2_linear_b_buf.raw_ptr();
        let h_s2_glu_ptr          = h_s2_glu_buf.raw_ptr();
        let h_s1_linear_a_ptr_tg  = h_s1_linear_a_buf_tg.raw_ptr();
        let h_s1_linear_b_ptr_tg  = h_s1_linear_b_buf_tg.raw_ptr();
        let h_s1_residual_ptr_tg  = h_s1_residual_buf_tg.raw_ptr();
        let h_s1_glu_ptr_tg       = h_s1_glu_buf_tg.raw_ptr();
        let h_s2_linear_a_ptr_tg  = h_s2_linear_a_buf_tg.raw_ptr();
        let h_s2_linear_b_ptr_tg  = h_s2_linear_b_buf_tg.raw_ptr();
        let h_s2_glu_ptr_tg       = h_s2_glu_buf_tg.raw_ptr();

        Ok(Self {
            handle: shared,
            // Dimensions
            batch_size,
            state_dim,
            state_dim_padded,
            s1_input_dim,
            s1_ldb,
            shared_h1,
            shared_h2,
            value_h,
            adv_h,
            num_atoms,
            branch_0_size,
            branch_1_size,
            branch_2_size,
            branch_3_size,
            mag_concat_dim: shared_h2 + branch_0_size,
            ofi_concat_dim: shared_h2 + 3,
            branch_streams,
            _branch_workspace_bufs: branch_workspace_bufs,
            branch_workspace_ptrs,
            trunk_done_event,
            branch_done_events,
            gemm_cache,
            gemm_cache_relu_bias,
            gemm_cache_bias,
            gemm_cache_relu_aux_bias,
            _value_fc_relu_mask_buf: value_fc_relu_mask_buf,
            value_fc_relu_mask_ptr_raw: value_fc_relu_mask_ptr,
            value_fc_relu_mask_aux_ld_raw: value_fc_relu_mask_aux_ld,
            value_fc_save_h_v_marker: 0u64,  // wired post-construction by the trainer
            // VSN/KAN: not wired until wire_vsn_glu() is called
            vsn_kernel: None,
            kan_gate_combine_kernel: None,
            strided_scatter_kernel: None,
            vsn_masked_ptr: 0,
            glu_gate_pre_ptrs: [0; 4],
            glu_value_ptrs: [0; 4],
            vsn_rank: 0,
            ord_concat_ptr: 0,
            urg_concat_ptr: 0,

            grn_h_s1_online,
            grn_h_s2_online,
            grn_h_s1_target,
            grn_h_s2_target,

            _grn_h_s1_linear_a_buf: h_s1_linear_a_buf,
            grn_h_s1_linear_a_ptr: h_s1_linear_a_ptr,
            _grn_h_s1_linear_b_buf: h_s1_linear_b_buf,
            grn_h_s1_linear_b_ptr: h_s1_linear_b_ptr,
            _grn_h_s1_residual_buf: h_s1_residual_buf,
            grn_h_s1_residual_ptr: h_s1_residual_ptr,
            _grn_h_s1_glu_buf: h_s1_glu_buf,
            grn_h_s1_glu_ptr: h_s1_glu_ptr,
            _grn_h_s2_linear_a_buf: h_s2_linear_a_buf,
            grn_h_s2_linear_a_ptr: h_s2_linear_a_ptr,
            _grn_h_s2_linear_b_buf: h_s2_linear_b_buf,
            grn_h_s2_linear_b_ptr: h_s2_linear_b_ptr,
            _grn_h_s2_glu_buf: h_s2_glu_buf,
            grn_h_s2_glu_ptr: h_s2_glu_ptr,

            _grn_h_s1_linear_a_buf_tg: h_s1_linear_a_buf_tg,
            grn_h_s1_linear_a_ptr_tg: h_s1_linear_a_ptr_tg,
            _grn_h_s1_linear_b_buf_tg: h_s1_linear_b_buf_tg,
            grn_h_s1_linear_b_ptr_tg: h_s1_linear_b_ptr_tg,
            _grn_h_s1_residual_buf_tg: h_s1_residual_buf_tg,
            grn_h_s1_residual_ptr_tg: h_s1_residual_ptr_tg,
            _grn_h_s1_glu_buf_tg: h_s1_glu_buf_tg,
            grn_h_s1_glu_ptr_tg: h_s1_glu_ptr_tg,
            _grn_h_s2_linear_a_buf_tg: h_s2_linear_a_buf_tg,
            grn_h_s2_linear_a_ptr_tg: h_s2_linear_a_ptr_tg,
            _grn_h_s2_linear_b_buf_tg: h_s2_linear_b_buf_tg,
            grn_h_s2_linear_b_ptr_tg: h_s2_linear_b_ptr_tg,
            _grn_h_s2_glu_buf_tg: h_s2_glu_buf_tg,
            grn_h_s2_glu_ptr_tg: h_s2_glu_ptr_tg,

            saxpy_kernel,

            // ── Plan 4 Task 1B-iii: VSN forward orchestrator handles ────
            vsn_softmax_and_gate_fwd_kernel,
            vsn_softmax_and_gate_bwd_kernel,
            // `vsn_logit_scatter_kernel` is wired post-construction by the
            // trainer via `wire_vsn_scatter` (mirrors `wire_vsn_glu`'s
            // post-construction pattern). The experience collector instance
            // never calls `vsn_forward` and therefore does not need it.
            vsn_logit_scatter_kernel: None,
        })
    }

    /// Wire VSN bottleneck + GLU gating kernels and buffer pointers.
    ///
    /// Called once after construction from `GpuDqnTrainer`. After this call,
    /// all 3 forward functions insert VSN→GEMM→bias→GLU in the branch loop
    /// instead of the old GEMM→bias→ReLU path.
    ///
    /// For the experience collector (which doesn't train), this is not called —
    /// VSN/GLU fields stay `None` and forward falls through to the existing
    /// ReLU path.
    #[allow(clippy::too_many_arguments)]
    pub fn wire_vsn_glu(
        &mut self,
        vsn_kernel: CudaFunction,
        kan_gate_combine_kernel: CudaFunction,
        strided_scatter_kernel: CudaFunction,
        vsn_masked_ptr: u64,
        glu_gate_pre_ptrs: [u64; 4],
        glu_value_ptrs: [u64; 4],
        vsn_rank: usize,
    ) {
        self.vsn_kernel = Some(vsn_kernel);
        self.kan_gate_combine_kernel = Some(kan_gate_combine_kernel);
        self.strided_scatter_kernel = Some(strided_scatter_kernel);
        self.vsn_masked_ptr = vsn_masked_ptr;
        self.glu_gate_pre_ptrs = glu_gate_pre_ptrs;
        self.glu_value_ptrs = glu_value_ptrs;
        self.vsn_rank = vsn_rank;
    }

    /// Wire OFI concat buffer pointers for order (d=2) and urgency (d=3) branches.
    ///
    /// Called once after construction from `GpuDqnTrainer`, after `wire_vsn_glu`.
    /// These pointers point to [B, SH2+3] buffers pre-built by `launch_concat_ofi`.
    pub fn wire_ofi_concat(&mut self, ord_concat_ptr: u64, urg_concat_ptr: u64) {
        self.ord_concat_ptr = ord_concat_ptr;
        self.urg_concat_ptr = urg_concat_ptr;
    }

    /// Plan 4 Task 1B-iii: wire the strided_scatter kernel handle used by
    /// `vsn_forward` to write each per-group `Linear_2_g` output (a tight
    /// `[B]` buffer) into column `g` of the assembled `vsn_logits_buf [B,
    /// num_groups]`. The trainer loads `strided_scatter` from
    /// `experience_kernels.cubin` once at construction (already wired for
    /// VSN bottleneck via `wire_vsn_glu`); the same handle is reused here.
    /// Called once post-construction from `GpuDqnTrainer`. Required for
    /// `vsn_forward` to function — `vsn_forward` `debug_assert!`s the
    /// presence of this handle.
    pub fn wire_vsn_scatter(&mut self, strided_scatter_kernel: CudaFunction) {
        self.vsn_logit_scatter_kernel = Some(strided_scatter_kernel);
    }

    /// Rebind the cuBLAS handle to a (potentially new) stream.
    ///
    // Accessors for diagnostic test
    pub(crate) fn adv_h(&self) -> usize { self.adv_h }
    pub(crate) fn batch_size(&self) -> usize { self.batch_size }
    pub(crate) fn shared_h2(&self) -> usize { self.shared_h2 }

    /// Online h_s2 GRN block accessor (read-only) — exposes the wrapper's
    /// save-for-backward state ptrs (elu_post / linear_b_out / sigmoid_b /
    /// ln_rstd / ln_normed) for the trainer's GRN-stage NaN diagnostic
    /// instrumentation (P, post-bkdx5 — slots 16-20 of nan_flags_buf).
    /// h_s2 is the bottleneck stage where save_h_s2 NaN was observed at
    /// F1 ~step 1745 in smoke-test-bkdx5; instrumenting only h_s2 covers
    /// the failing stage without doubling the per-step kernel-launch cost.
    pub(crate) fn grn_h_s2_online(&self) -> &super::gpu_grn::GrnBlock {
        &self.grn_h_s2_online
    }

    /// Online h_s2 Linear_b output buffer pointer (Linear_b output, [B, 2*SH2]).
    /// Same buffer that `state.linear_b_out` is DtoD-copied from inside
    /// `forward_raw_phase2` — exposed directly so the trainer can NaN-check
    /// the live cuBLAS output without depending on the save-DtoD landing.
    pub(crate) fn grn_h_s2_linear_b_ptr(&self) -> u64 {
        self.grn_h_s2_linear_b_ptr
    }

    /// Diagnostic: test cublasLtMatmul with raw cudaMalloc buffers.
    /// Uses the SAME lt_handle and workspace as the real forward.
    pub(crate) fn test_lt_matmul_raw(
        &self,
        stream: &CudaStream,
        w_ptr: u64, a_ptr: u64, c_ptr: u64,
        m: usize, n: usize, k: usize,
    ) -> Result<(), MLError> {
        self.lt_matmul(stream, w_ptr, a_ptr, c_ptr, m, n, k, k, self.handle.lt_workspace_ptr, self.handle.lt_workspace_size, "DIAG_raw")
    }

    /// Per-branch cuBLAS workspace pointers (shared with backward pass).
    pub fn branch_workspace_ptrs(&self) -> [u64; 4] {
        self.branch_workspace_ptrs
    }

    /// Element-wise SAXPY: `y[i] += alpha * x[i]` for `i = 0..n`.
    ///
    /// Plan 4 Task 2c.3c.2: provides the gradient-accumulation primitive the
    /// h_s2 GRN block's identity-residual backward needs. Reuses the existing
    /// `dqn_saxpy_f32_kernel` (loaded once at construction) instead of
    /// introducing a new cuBLAS legacy-handle stream-bind path.
    ///
    /// Both buffers live on the device; `n` is the f32 element count.
    /// Launch config: 256 threads/block, `ceil(n/256)` blocks (grid-stride
    /// not needed — kernel is already bounded by `if (i < n)`).
    pub fn saxpy_inplace(
        &self,
        stream: &Arc<CudaStream>,
        n: usize,
        alpha: f32,
        x_ptr: u64,
        y_ptr: u64,
    ) -> Result<(), MLError> {
        if n == 0 {
            return Ok(());
        }
        let blocks = ((n + 255) / 256) as u32;
        let cfg = LaunchConfig {
            grid_dim: (blocks, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };
        let n_i32 = n as i32;
        unsafe {
            stream
                .launch_builder(&self.saxpy_kernel)
                .arg(&y_ptr)
                .arg(&x_ptr)
                .arg(&alpha)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("saxpy_inplace launch: {e}")))?;
        }
        Ok(())
    }

    // ══════════════════════════════════════════════════════════════════════════
    // Plan 4 Task 1B-iii: Variable Selection Network forward orchestrator.
    // ══════════════════════════════════════════════════════════════════════════

    /// VSN feature-selection forward — per spec §4.E.1.
    ///
    /// Sequence (per group g in [0, SL_NUM_FEATURE_GROUPS)):
    ///   1. cuBLAS `Linear_1[g]`: state_in[B, group_dim_g] @ w1_g^T → h1_g[B, 16]
    ///      (uses `sgemm_f32_ldb` with `state_dim_padded` as ldb; the per-group
    ///       slice is selected via pointer offset `gb * sizeof(f32)` — cuBLAS
    ///       reads col-major [state_dim_padded, B] starting at column 0 of the
    ///       slice).
    ///   2. Fused bias-add + ReLU on h1_g via `add_bias_relu_f32_kernel`
    ///      (writes h1_g back as the post-ReLU activation; if save_h1_g_ptrs
    ///      is `Some`, we DtoD-copy into the supplied save buffer for 1B-iv
    ///      backward — but DtoD-copy is unnecessary because we ALLOCATE h1_g
    ///      AS the save buffer; see `save_h1_g_ptrs` doc below).
    ///   3. cuBLAS `Linear_2[g]`: h1_g[B, 16] @ w2_g^T → linear2_scratch[B, 1]
    ///   4. Bias-add (`+b2_g[0]`) on linear2_scratch.
    ///   5. Scatter linear2_scratch[B] → column `g` of `logits_buf[B, num_groups]`
    ///      via the existing `strided_scatter` kernel: src_stride=1,
    ///      dst_stride=num_groups, dst pointer offset by `g * sizeof(f32)`.
    ///
    /// After the per-group loop, launch `vsn_softmax_and_gate_forward`
    /// (1 block per sample, 256 threads, `num_groups*sizeof(f32)` shmem).
    /// The kernel reads logits + state_in, writes mask_out + gated_state_out.
    ///
    /// `w_ptrs` provides the 24 VSN tensors at indices `[95..119)` per the
    /// Plan 4 Task 1B-ii layout. Group g's tensors live at:
    ///   `[95 + 4*g + 0]` = vsn_w1_g{g}  (size `[VSN_HIDDEN_DIM, group_dim_g]`)
    ///   `[95 + 4*g + 1]` = vsn_b1_g{g}  (size `[VSN_HIDDEN_DIM]`)
    ///   `[95 + 4*g + 2]` = vsn_w2_g{g}  (size `[1, VSN_HIDDEN_DIM]`, flat-stored)
    ///   `[95 + 4*g + 3]` = vsn_b2_g{g}  (size `[1]`)
    ///
    /// `save_h1_g_ptrs`: 6 device-buffer pointers `[B, VSN_HIDDEN_DIM]` each.
    /// We write the per-group h1 (post-ReLU) DIRECTLY into these buffers —
    /// they ARE the save-for-backward storage on the online pass. For target /
    /// DDQN passes (no backward), the caller passes a single shared scratch
    /// pointer for all 6 groups (each Linear_1 overwrites the previous).
    ///
    /// `linear2_scratch_ptr`: shared `[B]` scratch overwritten per group; the
    /// scatter at step 5 reads it before the next group's Linear_2 writes it.
    /// Sharing across passes (online/target/ddqn) is safe because each pass
    /// runs the per-group loop sequentially before launching the softmax+gate
    /// kernel.
    ///
    /// `logits_buf_ptr`: `[B, num_groups]` — must be SAVED on the online pass
    /// (1B-iv backward needs it indirectly via `mask`); throwaway scratch on
    /// target/DDQN.
    ///
    /// `mask_out_ptr`: `[B, num_groups]` — saved on the online pass (read by
    /// `vsn_mask_ema_update` ISV producer + by 1B-iv softmax-Jacobian backward);
    /// throwaway scratch on target/DDQN.
    #[allow(clippy::too_many_arguments)]
    pub fn vsn_forward(
        &self,
        stream: &Arc<CudaStream>,
        state_in_ptr: u64,
        gated_state_out_ptr: u64,
        w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
        logits_buf_ptr: u64,
        mask_out_ptr: u64,
        save_h1_g_ptrs: &[u64; ml_core::state_layout::SL_NUM_FEATURE_GROUPS],
        linear2_scratch_ptr: u64,
        group_begins_dev: u64,
        group_ends_dev: u64,
        batch_size: usize,
    ) -> Result<(), MLError> {
        const VSN_HIDDEN: usize = super::gpu_dqn_trainer::VSN_HIDDEN_DIM;
        const NUM_GROUPS: usize = ml_core::state_layout::SL_NUM_FEATURE_GROUPS;
        const VSN_TENSORS_BASE: usize = 95;

        // Scatter kernel must have been wired by the trainer via `wire_vsn_scatter`.
        // Silent-skip on `None` would mask a misconfiguration; `debug_assert!` flags
        // it in dev builds and release inherits the loud `expect()` panic if the
        // invariant ever breaks.
        debug_assert!(
            self.vsn_logit_scatter_kernel.is_some(),
            "vsn_forward: vsn_logit_scatter_kernel must be wired via `wire_vsn_scatter` \
             before any forward call (trainer-side `wire_vsn_scatter` invocation \
             missing for this CublasGemmSet instance)"
        );
        let scatter_kernel = self
            .vsn_logit_scatter_kernel
            .as_ref()
            .expect("vsn_logit_scatter_kernel must be wired before vsn_forward");

        let group_ranges = ml_core::state_layout::FEATURE_GROUP_RANGES;
        let state_dim_padded = self.state_dim_padded;
        let f32_size = std::mem::size_of::<f32>() as u64;

        // ── Per-group MLP: Linear_1 → ReLU → Linear_2 → bias → scatter ──
        for g in 0..NUM_GROUPS {
            let (gb, ge) = group_ranges[g];
            let group_dim = ge - gb;
            debug_assert!(
                group_dim > 0,
                "vsn_forward: feature group {g} has zero dim — group ranges invariant violated"
            );

            let w1_ptr = w_ptrs[VSN_TENSORS_BASE + 4 * g];
            let b1_ptr = w_ptrs[VSN_TENSORS_BASE + 4 * g + 1];
            let w2_ptr = w_ptrs[VSN_TENSORS_BASE + 4 * g + 2];
            let b2_ptr = w_ptrs[VSN_TENSORS_BASE + 4 * g + 3];
            let h1_ptr = save_h1_g_ptrs[g];
            let label_l1 = ["vsn_l1_g0", "vsn_l1_g1", "vsn_l1_g2", "vsn_l1_g3", "vsn_l1_g4", "vsn_l1_g5"][g];
            let label_l2 = ["vsn_l2_g0", "vsn_l2_g1", "vsn_l2_g2", "vsn_l2_g3", "vsn_l2_g4", "vsn_l2_g5"][g];

            // Step 1+2: Linear_1[g] + bias + ReLU. Phase H Site 3: fused via
            // cublasLt EPILOGUE_RELU_BIAS — h1_ptr ends up holding the POST-ReLU
            // activation (same contract as the legacy two-kernel path). The
            // post-ReLU buffer is what 1B-iv's backward needs as the "saved h1"
            // (relu_mask_standalone takes the post-ReLU activation as the gating
            // signal). No separate DtoD save copy needed.
            //
            // Pointer offset selects the per-group column slice; ldb stays at
            // state_dim_padded so cuBLAS reads the original [state_dim_padded, B]
            // col-major view starting at the group's first feature.
            let state_slice_ptr = state_in_ptr + (gb as u64) * f32_size;
            let vsn_l1_ws = self.handle.lt_workspace_ptr;
            let vsn_l1_wss = self.handle.lt_workspace_size;
            if self.sgemm_f32_fused_relu_bias(
                stream,
                w1_ptr,
                state_slice_ptr,
                h1_ptr,
                b1_ptr,
                VSN_HIDDEN, batch_size, group_dim, state_dim_padded,
                vsn_l1_ws, vsn_l1_wss,
                label_l1,
            ).is_err() {
                self.sgemm_f32_ldb(
                    stream,
                    w1_ptr,
                    state_slice_ptr,
                    h1_ptr,
                    VSN_HIDDEN, batch_size, group_dim,
                    state_dim_padded,
                    label_l1,
                )?;
                self.launch_add_bias_relu_f32_raw(stream, h1_ptr, b1_ptr, VSN_HIDDEN, batch_size)?;
            }

            // Step 3+4: Linear_2[g] + bias_add (no activation — Linear_2 produces a
            // logit). Phase H Site 3: fused GEMM+bias via cublasLt EPILOGUE_BIAS;
            // falls back to separate GEMM + add_bias on missing-cache.
            let vsn_ws = self.handle.lt_workspace_ptr;
            let vsn_wss = self.handle.lt_workspace_size;
            if self.sgemm_f32_fused_bias(
                stream,
                w2_ptr,
                h1_ptr,
                linear2_scratch_ptr,
                b2_ptr,
                1, batch_size, VSN_HIDDEN, VSN_HIDDEN,
                vsn_ws, vsn_wss,
                label_l2,
            ).is_err() {
                self.sgemm_f32(
                    stream,
                    w2_ptr,
                    h1_ptr,
                    linear2_scratch_ptr,
                    1, batch_size, VSN_HIDDEN,
                    label_l2,
                )?;
                self.launch_add_bias_f32_raw(stream, linear2_scratch_ptr, b2_ptr, 1, batch_size)?;
            }

            // Step 5: scatter linear2_scratch[B] → col g of logits_buf[B, num_groups].
            // Reuses `strided_scatter` from experience_kernels.cu:
            //   dst[b * dst_stride + i] = src[b * src_stride + i]   for i in [0, src_stride)
            // With src_stride=1, dst_stride=num_groups, and dst offset by g*sizeof(f32):
            //   dst_ptr[b * num_groups + 0] in shifted view = logits_buf[b * num_groups + g]
            // No atomicAdd; one thread per sample, fully data-parallel.
            let src_stride: i32 = 1;
            let dst_stride: i32 = NUM_GROUPS as i32;
            let total: i32 = batch_size as i32; // = B * src_stride
            let blocks = ((batch_size as u32 + 255) / 256).max(1);
            let dst_col_ptr = logits_buf_ptr + (g as u64) * f32_size;
            unsafe {
                stream
                    .launch_builder(scatter_kernel)
                    .arg(&linear2_scratch_ptr)
                    .arg(&dst_col_ptr)
                    .arg(&src_stride)
                    .arg(&dst_stride)
                    .arg(&total)
                    .launch(LaunchConfig {
                        grid_dim: (blocks, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| MLError::ModelError(format!(
                        "vsn_forward strided_scatter group {g}: {e}"
                    )))?;
            }
        }

        // ── Softmax-over-groups + per-feature gate-multiply ──
        // 1 block per sample, 256 threads. Shared memory = num_groups floats
        // (24 bytes for num_groups=6).
        let b_i = batch_size as i32;
        let sd_padded_i = state_dim_padded as i32;
        let num_groups_i = NUM_GROUPS as i32;
        let smem_bytes = (NUM_GROUPS * std::mem::size_of::<f32>()) as u32;
        unsafe {
            stream
                .launch_builder(&self.vsn_softmax_and_gate_fwd_kernel)
                .arg(&state_in_ptr)
                .arg(&logits_buf_ptr)
                .arg(&b_i)
                .arg(&sd_padded_i)
                .arg(&num_groups_i)
                .arg(&group_begins_dev)
                .arg(&group_ends_dev)
                .arg(&mask_out_ptr)
                .arg(&gated_state_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (batch_size as u32, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: smem_bytes,
                })
                .map_err(|e| MLError::ModelError(format!("vsn_softmax_and_gate_forward: {e}")))?;
        }

        Ok(())
    }

    // ══════════════════════════════════════════════════════════════════════════
    // Forward pass (online network — saves activations for backward)
    // ══════════════════════════════════════════════════════════════════════════

    /// Run the online network forward pass using cuBLAS SGEMM.
    ///
    /// Writes into the trainer's pre-allocated activation save buffers:
    /// - `save_h_s1`: [B, SH1]     — shared trunk layer 1 activations
    /// - `save_h_s2`: [B, SH2]     — shared trunk layer 2 activations
    /// - `save_h_v`:  [B, VH]      — value head hidden activations
    /// - `save_h_b0`: [B, AH]      — branch 0 hidden activations
    /// - `save_h_b1`: [B, AH]      — branch 1 hidden activations
    /// - `save_h_b2`: [B, AH]      — branch 2 hidden activations
    ///
    /// Value logits  → written to `v_logits_buf`  [B, NA]
    /// Branch logits → written to `b_logits_buf`  [B, (B0+B1+B2)*NA]
    ///
    /// Weight pointers are raw u64 into the flat F32 `params_buf` at GOFF_* offsets.
    #[allow(clippy::too_many_arguments)]
    /// Run only the trunk encoder (h_s1 + h_s2) — GRN composition.
    ///
    /// Plan 4 Task 2c.3b: swaps the legacy `Linear→ReLU→Linear→ReLU`
    /// trunk GEMMs for two GRN blocks. Each block runs:
    ///
    ///   GRN(x) = LayerNorm(P(x) + GLU(Linear_b(ELU(Linear_a(x)))))
    ///
    /// where `Linear_a`, `Linear_b`, and (for h_s1) `Linear_residual` are
    /// cuBLAS GEMMs and ELU + GLU + residual+LN are dispatched via
    /// `GrnBlock::forward_raw` from `gpu_grn.rs`.
    ///
    /// Param-tensor layout (post Task 2c.3a reshuffle):
    /// h_s1 GRN — [0]=w_a, [1]=b_a, [2]=w_b, [3]=b_b, [4]=w_residual,
    ///            [5]=gamma, [6]=beta.
    /// h_s2 GRN — [7]=w_a, [8]=b_a, [9]=w_b, [10]=b_b, [11]=gamma,
    ///            [12]=beta. (Identity residual; no Linear_residual.)
    ///
    /// Save-for-backward state lives in the wrapper's `grn_h_s1_online` /
    /// `grn_h_s2_online` instances. The save-buffer DtoD copies are
    /// always issued here (this is the online path) so the GRN backward
    /// in 2c.3c will see consistent state.
    pub fn encoder_forward_only(
        &mut self,
        stream: &Arc<CudaStream>,
        states_ptr: u64,
        w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
        h_s1_ptr: u64,
        h_s2_ptr: u64,
    ) -> Result<(), MLError> {
        let b = self.batch_size;

        // Snapshot the scratch pointers so we don't borrow `self` while
        // also calling `self.sgemm_f32_*` and `self.grn_h_s1_online.*`.
        let h_s1_la_ptr  = self.grn_h_s1_linear_a_ptr;
        let h_s1_lb_ptr  = self.grn_h_s1_linear_b_ptr;
        let h_s1_res_ptr = self.grn_h_s1_residual_ptr;
        let h_s1_glu_ptr = self.grn_h_s1_glu_ptr;
        let h_s2_la_ptr  = self.grn_h_s2_linear_a_ptr;
        let h_s2_lb_ptr  = self.grn_h_s2_linear_b_ptr;
        let h_s2_glu_ptr = self.grn_h_s2_glu_ptr;
        let sh1 = self.shared_h1;
        let sh2 = self.shared_h2;
        let s1_in = self.s1_input_dim;
        let s1_ldb = self.s1_ldb;
        let ws = self.handle.lt_workspace_ptr;
        let wss = self.handle.lt_workspace_size;

        // ── h_s1 GRN block ───────────────────────────────────────────
        // 1. Linear_a: states @ w_a^T + b_a → linear_a_out [B, SH1]
        if self.sgemm_f32_fused_bias(stream, w_ptrs[0], states_ptr, h_s1_la_ptr, w_ptrs[1],
            sh1, b, s1_in, s1_ldb, ws, wss, "grn_h_s1_linear_a").is_err()
        {
            self.sgemm_f32_ldb(stream, w_ptrs[0], states_ptr, h_s1_la_ptr,
                sh1, b, s1_in, s1_ldb, "grn_h_s1_linear_a")?;
            self.launch_add_bias_f32_raw(stream, h_s1_la_ptr, w_ptrs[1], sh1, b)?;
        }
        // 2. Linear_b: linear_a_out (post-ELU happens inside GrnBlock) — but
        //    Linear_b reads pre-ELU? No: per spec, Linear_b reads ELU(Linear_a).
        //    The wrapper does ELU in-place on linear_a_out FIRST. Reorder:
        //    a) ELU happens in GrnBlock::forward_raw step 1.
        //    Therefore we must run Linear_b AFTER ELU. Move Linear_b GEMM
        //    into the dispatcher itself? Cleaner: do ELU here too via a
        //    micro split — the wrapper's forward_raw expects pre-ELU
        //    linear_a_out as input. So:
        //      - run Linear_a + bias here
        //      - call wrapper.elu_then_save (the wrapper's first 2 steps)
        //      - run Linear_b (input = post-ELU = linear_a_out) + bias
        //      - run Linear_residual (input = states) + nothing (no bias)
        //      - call wrapper.glu_then_ln (the wrapper's last 2 steps)
        //    But splitting forward_raw across the cuBLAS fence would
        //    duplicate the dispatcher logic. Simpler alternative: the
        //    wrapper's contract IS pre-ELU input → it does ELU inside.
        //    So we need Linear_b's input to be the linear_a_out buffer
        //    AFTER the wrapper has run ELU on it. That means we must
        //    issue Linear_b AFTER calling forward_raw, but forward_raw
        //    also runs GLU/LN which expect Linear_b output already.
        //
        // Resolution: split the wrapper into two phases. Run the cuBLAS
        // GEMMs in dispatcher order (Linear_a → ELU → Linear_b →
        // Linear_residual → GLU+LN). Easiest: inline the kernel calls
        // here by reaching into `grn_h_s1_online` / `grn_h_s2_online`
        // directly via dedicated dispatcher methods that segment the
        // forward_raw work.
        //
        // To keep the wrapper API simple, we issue each kernel through
        // `GrnBlock::forward_raw` after Linear_a has run, but only AFTER
        // we have also run Linear_b (because forward_raw needs linear_b
        // ready). Therefore: run Linear_a + bias, then run a "phase-1"
        // wrapper call that just does ELU + DtoD save, then Linear_b +
        // bias, then a "phase-2" wrapper call that does
        // GLU + residual+LN. That requires two new wrapper methods
        // (`forward_raw_phase1`, `forward_raw_phase2`).
        //
        // Concrete impl below: avoid both wrapper splits by manually
        // issuing the kernel sequence here through references obtained
        // via dedicated helpers. We use one method on GrnBlock that does
        // the full sequence, and we run the cuBLAS GEMMs of the dispatcher
        // in matching order. The trick: feed Linear_b a buffer AFTER ELU
        // has been applied. Linear_a output buffer holds pre-ELU; we
        // can NOT simply call forward_raw with linear_a_out (still pre-ELU)
        // and linear_b_out (also still containing zeros) — that would run
        // ELU in-place on linear_a_out (now post-ELU), then immediately
        // try to GLU on stale `linear_b_out_ptr` (zero-init).
        //
        // Therefore the correct order is:
        //   1. Linear_a + bias  → linear_a_out [pre-ELU]
        //   2. wrapper.elu_save → in-place ELU + DtoD save
        //   3. Linear_b + bias  → linear_b_out [post-Linear_b, pre-GLU]
        //   4. Linear_residual + (no bias) → residual_proj
        //   5. wrapper.glu_ln_save → GLU + LN+residual + saves
        //
        // Step 2 and 5 are the "phase 1" and "phase 2" of the wrapper.
        // We add `forward_raw_phase1` and `forward_raw_phase2` to GrnBlock
        // — small dispatchers that each launch one or two of the existing
        // 5 forward kernels. No new kernels. Pearl-clean.
        //
        // (Code follows; see gpu_grn.rs for the two new methods.)

        self.grn_h_s1_online.forward_raw_phase1(stream, h_s1_la_ptr, b, /* save = */ true)?;

        // 3. Linear_b: linear_a_out (now post-ELU) @ w_b^T + b_b → linear_b_out [B, 2*SH1]
        if self.sgemm_f32_fused_bias(stream, w_ptrs[2], h_s1_la_ptr, h_s1_lb_ptr, w_ptrs[3],
            2 * sh1, b, sh1, sh1, ws, wss, "grn_h_s1_linear_b").is_err()
        {
            self.sgemm_f32(stream, w_ptrs[2], h_s1_la_ptr, h_s1_lb_ptr,
                2 * sh1, b, sh1, "grn_h_s1_linear_b")?;
            self.launch_add_bias_f32_raw(stream, h_s1_lb_ptr, w_ptrs[3], 2 * sh1, b)?;
        }

        // 4. Linear_residual: states @ w_residual^T → residual_proj [B, SH1]
        //    (no bias term; Linear_residual is a pure projection.)
        self.sgemm_f32_ldb(stream, w_ptrs[4], states_ptr, h_s1_res_ptr,
            sh1, b, s1_in, s1_ldb, "grn_h_s1_linear_residual")?;

        // 5. Wrapper phase 2: GLU forward + residual+LN forward → h_s1_ptr
        self.grn_h_s1_online.forward_raw_phase2(
            stream,
            h_s1_lb_ptr,
            h_s1_res_ptr,
            w_ptrs[5],   // gamma_h_s1
            w_ptrs[6],   // beta_h_s1
            h_s1_glu_ptr,
            h_s1_ptr,
            b,
            /* save = */ true,
        )?;

        // ── h_s2 GRN block (identity residual, h_s1_ptr is the residual) ──
        // 1. Linear_a: h_s1 @ w_a^T + b_a → linear_a_out [B, SH2]
        if self.sgemm_f32_fused_bias(stream, w_ptrs[7], h_s1_ptr, h_s2_la_ptr, w_ptrs[8],
            sh2, b, sh1, sh1, ws, wss, "grn_h_s2_linear_a").is_err()
        {
            self.sgemm_f32(stream, w_ptrs[7], h_s1_ptr, h_s2_la_ptr,
                sh2, b, sh1, "grn_h_s2_linear_a")?;
            self.launch_add_bias_f32_raw(stream, h_s2_la_ptr, w_ptrs[8], sh2, b)?;
        }

        // 2. Wrapper phase 1: ELU in-place + save.
        self.grn_h_s2_online.forward_raw_phase1(stream, h_s2_la_ptr, b, /* save = */ true)?;

        // 3. Linear_b: linear_a_out (post-ELU) @ w_b^T + b_b → linear_b_out [B, 2*SH2]
        if self.sgemm_f32_fused_bias(stream, w_ptrs[9], h_s2_la_ptr, h_s2_lb_ptr, w_ptrs[10],
            2 * sh2, b, sh2, sh2, ws, wss, "grn_h_s2_linear_b").is_err()
        {
            self.sgemm_f32(stream, w_ptrs[9], h_s2_la_ptr, h_s2_lb_ptr,
                2 * sh2, b, sh2, "grn_h_s2_linear_b")?;
            self.launch_add_bias_f32_raw(stream, h_s2_lb_ptr, w_ptrs[10], 2 * sh2, b)?;
        }

        // 4. (no Linear_residual — h_s1_ptr is the residual directly).
        // 5. Wrapper phase 2: GLU + residual+LN with h_s1_ptr as residual → h_s2_ptr
        self.grn_h_s2_online.forward_raw_phase2(
            stream,
            h_s2_lb_ptr,
            h_s1_ptr,    // identity residual
            w_ptrs[11],  // gamma_h_s2
            w_ptrs[12],  // beta_h_s2
            h_s2_glu_ptr,
            h_s2_ptr,
            b,
            /* save = */ true,
        )?;

        Ok(())
    }

    /// Run the target trunk encoder (h_s1 + h_s2) — GRN composition.
    ///
    /// Plan 4 Task 2c.3b: mirrors `encoder_forward_only` but uses the
    /// dedicated `grn_h_s1_target` / `grn_h_s2_target` instances and
    /// target-side scratch buffers so it never clobbers the online's
    /// save-for-backward state. `save_for_backward = false` because the
    /// GRN backward only runs against the online network (target is
    /// inference-only).
    ///
    /// `tg_w_ptrs` indexes into the target weight buffer (same flat
    /// layout as `params_buf`).
    pub fn target_encoder_forward_only(
        &mut self,
        stream: &Arc<CudaStream>,
        states_ptr: u64,
        tg_w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
        h_s1_ptr: u64,
        h_s2_ptr: u64,
    ) -> Result<(), MLError> {
        let b = self.batch_size;

        let h_s1_la_ptr  = self.grn_h_s1_linear_a_ptr_tg;
        let h_s1_lb_ptr  = self.grn_h_s1_linear_b_ptr_tg;
        let h_s1_res_ptr = self.grn_h_s1_residual_ptr_tg;
        let h_s1_glu_ptr = self.grn_h_s1_glu_ptr_tg;
        let h_s2_la_ptr  = self.grn_h_s2_linear_a_ptr_tg;
        let h_s2_lb_ptr  = self.grn_h_s2_linear_b_ptr_tg;
        let h_s2_glu_ptr = self.grn_h_s2_glu_ptr_tg;
        let sh1 = self.shared_h1;
        let sh2 = self.shared_h2;
        let s1_in = self.s1_input_dim;
        let s1_ldb = self.s1_ldb;
        let ws = self.handle.lt_workspace_ptr;
        let wss = self.handle.lt_workspace_size;

        // ── h_s1 target GRN block ────────────────────────────────────
        if self.sgemm_f32_fused_bias(stream, tg_w_ptrs[0], states_ptr, h_s1_la_ptr, tg_w_ptrs[1],
            sh1, b, s1_in, s1_ldb, ws, wss, "tg_grn_h_s1_linear_a").is_err()
        {
            self.sgemm_f32_ldb(stream, tg_w_ptrs[0], states_ptr, h_s1_la_ptr,
                sh1, b, s1_in, s1_ldb, "tg_grn_h_s1_linear_a")?;
            self.launch_add_bias_f32_raw(stream, h_s1_la_ptr, tg_w_ptrs[1], sh1, b)?;
        }
        self.grn_h_s1_target.forward_raw_phase1(stream, h_s1_la_ptr, b, /* save = */ false)?;
        if self.sgemm_f32_fused_bias(stream, tg_w_ptrs[2], h_s1_la_ptr, h_s1_lb_ptr, tg_w_ptrs[3],
            2 * sh1, b, sh1, sh1, ws, wss, "tg_grn_h_s1_linear_b").is_err()
        {
            self.sgemm_f32(stream, tg_w_ptrs[2], h_s1_la_ptr, h_s1_lb_ptr,
                2 * sh1, b, sh1, "tg_grn_h_s1_linear_b")?;
            self.launch_add_bias_f32_raw(stream, h_s1_lb_ptr, tg_w_ptrs[3], 2 * sh1, b)?;
        }
        self.sgemm_f32_ldb(stream, tg_w_ptrs[4], states_ptr, h_s1_res_ptr,
            sh1, b, s1_in, s1_ldb, "tg_grn_h_s1_linear_residual")?;
        self.grn_h_s1_target.forward_raw_phase2(
            stream,
            h_s1_lb_ptr,
            h_s1_res_ptr,
            tg_w_ptrs[5],
            tg_w_ptrs[6],
            h_s1_glu_ptr,
            h_s1_ptr,
            b,
            /* save = */ false,
        )?;

        // ── h_s2 target GRN block (identity residual) ────────────────
        if self.sgemm_f32_fused_bias(stream, tg_w_ptrs[7], h_s1_ptr, h_s2_la_ptr, tg_w_ptrs[8],
            sh2, b, sh1, sh1, ws, wss, "tg_grn_h_s2_linear_a").is_err()
        {
            self.sgemm_f32(stream, tg_w_ptrs[7], h_s1_ptr, h_s2_la_ptr,
                sh2, b, sh1, "tg_grn_h_s2_linear_a")?;
            self.launch_add_bias_f32_raw(stream, h_s2_la_ptr, tg_w_ptrs[8], sh2, b)?;
        }
        self.grn_h_s2_target.forward_raw_phase1(stream, h_s2_la_ptr, b, /* save = */ false)?;
        if self.sgemm_f32_fused_bias(stream, tg_w_ptrs[9], h_s2_la_ptr, h_s2_lb_ptr, tg_w_ptrs[10],
            2 * sh2, b, sh2, sh2, ws, wss, "tg_grn_h_s2_linear_b").is_err()
        {
            self.sgemm_f32(stream, tg_w_ptrs[9], h_s2_la_ptr, h_s2_lb_ptr,
                2 * sh2, b, sh2, "tg_grn_h_s2_linear_b")?;
            self.launch_add_bias_f32_raw(stream, h_s2_lb_ptr, tg_w_ptrs[10], 2 * sh2, b)?;
        }
        self.grn_h_s2_target.forward_raw_phase2(
            stream,
            h_s2_lb_ptr,
            h_s1_ptr,
            tg_w_ptrs[11],
            tg_w_ptrs[12],
            h_s2_glu_ptr,
            h_s2_ptr,
            b,
            /* save = */ false,
        )?;

        Ok(())
    }

    /// Run the GRN trunk encoder backward chain (h_s2 GRN → h_s1 GRN).
    ///
    /// Plan 4 Task 2c.3c.4 — replaces the legacy `Linear→ReLU→Linear`
    /// trunk-backward GEMM sequence in `BatchedBackward::backward_full`,
    /// `apply_iqn_trunk_gradient`, and `apply_ensemble_diversity_backward`
    /// with the GRN backward composition that mirrors `encoder_forward_only`.
    ///
    /// # Sequence
    ///
    /// Composition (reversed from forward — see `encoder_forward_only`):
    /// ```text
    ///   d_h_s2_ptr
    ///     │
    ///     ├─ h_s2 GRN backward ──────────────────────────────
    ///     │   1. backward_raw_phase1: LN_bwd + LN_dgdb + GLU_bwd
    ///     │      ─→ d_pre_ln_h_s2 [B, SH2]    (residual aliasing)
    ///     │      ─→ d_linear_b_out [B, 2*SH2]
    ///     │      ─→ d_gamma_h_s2, d_beta_h_s2 (into grad_buf [11], [12])
    ///     │   2. cuBLAS Linear_b: dW into grad_buf[9],[10]; dX into d_elu_out
    ///     │   3. backward_raw_phase2: ELU_bwd → d_linear_a_out
    ///     │   4. cuBLAS Linear_a: dW into grad_buf[7],[8]; dX → d_h_s1_ptr
    ///     │   5. saxpy_inplace identity residual: d_h_s1 += d_pre_ln_h_s2
    ///     │
    ///     └─ d_h_s1_ptr
    ///         │
    ///         ├─ h_s1 GRN backward ──────────────────────────
    ///         │   6. backward_raw_phase1 → d_pre_ln_h_s1, d_linear_b_out
    ///         │      → d_gamma_h_s1, d_beta_h_s1 (grad_buf [5], [6])
    ///         │   7. cuBLAS Linear_b: dW into grad_buf[2],[3]; dX → d_elu_out
    ///         │   8. backward_raw_phase2: ELU_bwd → d_linear_a_out
    ///         │   9. cuBLAS Linear_a (lda): dW into grad_buf[0],[1];
    ///         │      dX → s1_dx_output (or skip if 0)
    ///         │  10. cuBLAS Linear_residual (no_bias_lda): dW into grad_buf[4]
    ///         │  11. (optional) Linear_residual dX accumulation when bottleneck
    ///         │      writes input gradient into s1_dx_output (beta=1.0)
    ///         │
    ///         └─ s1_dx_output (when bottleneck) — gradient w.r.t. bn_concat
    /// ```
    ///
    /// # Param-tensor offsets (post Task 2c.3a, 13 GRN trunk tensors at [0..13)):
    ///   [0..7)   h_s1: w_a, b_a, w_b, b_b, w_residual, gamma, beta
    ///   [7..13)  h_s2: w_a, b_a, w_b, b_b, gamma, beta (no Linear_residual)
    ///
    /// # Args
    /// * `stream` — main stream that issued the GRN forward (save buffers
    ///   are graph-safe to read on the same stream).
    /// * `cublas_backward` — provides cuBLAS dW/dX backward GEMMs.
    /// * `states_ptr` — h_s1 input buffer pointer (`states_buf` or `bn_concat`).
    /// * `w_ptrs` — pointer table indexing into the f32 master weight buffer.
    /// * `grad_buf_base` — base address of the accumulator the per-tensor
    ///   gradients write into. Each tensor gets its own padded offset via
    ///   `padded_byte_offset(param_sizes, idx)`. The h_s1/h_s2 GRN tensors
    ///   live at indices [0..13). For the auxiliary IQN/ensemble paths the
    ///   caller passes `iqn_trunk_m` (the trunk-only scratch); for the main
    ///   and CQL paths the caller passes the full `grad_buf` / `cql_grad_scratch`.
    /// * `param_sizes` — flat tensor sizes (output of `compute_param_sizes`).
    /// * `d_h_s2_ptr` — [B, SH2] gradient flowing into h_s2 GRN output.
    ///   Read-only (Phase-1 doesn't write back into it).
    /// * `d_h_s1_ptr` — [B, SH1] writable scratch. The h_s2 backward writes
    ///   `d_h_s1 = Linear_a^T(d_h_s2_la_out)` here (beta=0 — overwrite),
    ///   then SAXPY accumulates `d_pre_ln_h_s2` (identity residual). The
    ///   final value is the gradient flowing into h_s1's GRN output.
    /// * `s1_dx_output` — [B, s1_input_dim] writable input-gradient slot.
    ///   Pass `0` to skip dX (top-level path with no trainable input).
    ///   Pass non-zero (e.g. `bn_d_concat_buf`) for the bottleneck case;
    ///   Linear_a writes (beta=0) and Linear_residual accumulates (beta=1).
    /// * Eight scratch buffers — see `bw_grn_h_s{1,2}_d_*` on the trainer.
    ///   Sized [B, hidden] for `d_pre_ln`, `d_elu_out`, `d_linear_a_out`
    ///   and [B, 2*hidden] for `d_linear_b_out`. Live caller-side because
    ///   the auxiliary trainer paths reuse them across multiple chains.
    /// * `save_h_s1_ptr` — h_s1 GRN output saved during forward (input to h_s2
    ///   Linear_a — bypasses the wrapper's `linear_a_post_elu` save which is
    ///   for h_s2's own ELU input).
    #[allow(clippy::too_many_arguments)]
    pub fn encoder_backward_chain(
        &mut self,
        stream: &Arc<CudaStream>,
        cublas_backward: &super::batched_backward::CublasBackwardSet,
        states_ptr: u64,
        w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
        grad_buf_base: u64,
        param_sizes: &[usize],
        d_h_s2_ptr: u64,
        d_h_s1_ptr: u64,
        s1_dx_output: u64,
        grn_h_s1_d_pre_ln: u64,
        grn_h_s1_d_linear_b_out: u64,
        grn_h_s1_d_elu_out: u64,
        grn_h_s1_d_linear_a_out: u64,
        grn_h_s2_d_pre_ln: u64,
        grn_h_s2_d_linear_b_out: u64,
        grn_h_s2_d_elu_out: u64,
        grn_h_s2_d_linear_a_out: u64,
        save_h_s1_ptr: u64,
    ) -> Result<(), MLError> {
        use super::gpu_dqn_trainer::padded_byte_offset;

        let b = self.batch_size;
        let sh1 = self.shared_h1;
        let sh2 = self.shared_h2;
        let s1_in = self.s1_input_dim;
        let s1_lda = self.s1_ldb;

        // Linear_a saved-input pointers — same buffers the forward filled
        // (post-ELU after `forward_raw_phase1`'s in-place ELU launch).
        let h_s1_la_ptr = self.grn_h_s1_linear_a_ptr;
        let h_s2_la_ptr = self.grn_h_s2_linear_a_ptr;

        // ── h_s2 GRN backward ────────────────────────────────────────
        // 1. Phase 1: LN_bwd + LN_dgdb + GLU_bwd.
        //    d_h_s2_ptr → grn_h_s2_d_pre_ln (gated+residual aliasing)
        //               → grn_h_s2_d_linear_b_out
        //               → grad_buf[11] (d_gamma_h_s2), grad_buf[12] (d_beta_h_s2)
        self.grn_h_s2_online.backward_raw_phase1(
            stream,
            d_h_s2_ptr,
            w_ptrs[11],                                            // gamma_h_s2
            b,
            grn_h_s2_d_pre_ln,
            grn_h_s2_d_linear_b_out,
            grad_buf_base + padded_byte_offset(param_sizes, 11),   // d_gamma_h_s2
            grad_buf_base + padded_byte_offset(param_sizes, 12),   // d_beta_h_s2
        )?;

        // 2. cuBLAS Linear_b backward (h_s2):
        //      dy = grn_h_s2_d_linear_b_out [B, 2*SH2]
        //      x  = h_s2_la_ptr (post-ELU, kept in-place by phase1)
        //      dW into goff[9], dB into goff[10], dX → grn_h_s2_d_elu_out (beta=0)
        cublas_backward.launch_dw_only(
            stream,
            grn_h_s2_d_linear_b_out,
            h_s2_la_ptr,
            grad_buf_base + padded_byte_offset(param_sizes, 9),    // d_w_b_h_s2
            grad_buf_base + padded_byte_offset(param_sizes, 10),   // d_b_b_h_s2
            2 * sh2, sh2, b,
        )?;
        cublas_backward.launch_dx_only(
            stream,
            grn_h_s2_d_linear_b_out,
            w_ptrs[9],
            grn_h_s2_d_elu_out,
            2 * sh2, sh2, b, 0.0_f32,
        )?;

        // 3. Phase 2: ELU_bwd → grn_h_s2_d_linear_a_out.
        self.grn_h_s2_online.backward_raw_phase2(
            stream,
            grn_h_s2_d_elu_out,
            b,
            grn_h_s2_d_linear_a_out,
        )?;

        // 4. cuBLAS Linear_a backward (h_s2):
        //      dy = grn_h_s2_d_linear_a_out [B, SH2]
        //      x  = save_h_s1_ptr (h_s1 GRN output, fed Linear_a in forward)
        //      dW into goff[7], dB into goff[8], dX → d_h_s1_ptr (beta=0 — overwrites)
        cublas_backward.launch_dw_only(
            stream,
            grn_h_s2_d_linear_a_out,
            save_h_s1_ptr,
            grad_buf_base + padded_byte_offset(param_sizes, 7),    // d_w_a_h_s2
            grad_buf_base + padded_byte_offset(param_sizes, 8),    // d_b_a_h_s2
            sh2, sh1, b,
        )?;
        cublas_backward.launch_dx_only(
            stream,
            grn_h_s2_d_linear_a_out,
            w_ptrs[7],
            d_h_s1_ptr,
            sh2, sh1, b, 0.0_f32,                                  // overwrite d_h_s1_ptr
        )?;

        // 5. h_s2 identity residual: d_h_s1 += d_pre_ln_h_s2 (alpha = 1.0).
        //    The h_s2 GRN block uses h_s1 as its residual (no `Linear_residual`),
        //    so the LN-input gradient flows straight into d_h_s1 via SAXPY.
        self.saxpy_inplace(
            stream, b * sh1, 1.0_f32,
            grn_h_s2_d_pre_ln,
            d_h_s1_ptr,
        )?;

        // ── h_s1 GRN backward ────────────────────────────────────────
        // 6. Phase 1: LN_bwd + LN_dgdb + GLU_bwd.
        //    d_h_s1_ptr → grn_h_s1_d_pre_ln, grn_h_s1_d_linear_b_out
        //               → grad_buf[5] (d_gamma_h_s1), grad_buf[6] (d_beta_h_s1)
        self.grn_h_s1_online.backward_raw_phase1(
            stream,
            d_h_s1_ptr,
            w_ptrs[5],                                             // gamma_h_s1
            b,
            grn_h_s1_d_pre_ln,
            grn_h_s1_d_linear_b_out,
            grad_buf_base + padded_byte_offset(param_sizes, 5),    // d_gamma_h_s1
            grad_buf_base + padded_byte_offset(param_sizes, 6),    // d_beta_h_s1
        )?;

        // 7. cuBLAS Linear_b backward (h_s1):
        //      dy = grn_h_s1_d_linear_b_out [B, 2*SH1]
        //      x  = h_s1_la_ptr (post-ELU)
        //      dW into goff[2], dB into goff[3], dX → grn_h_s1_d_elu_out (beta=0)
        cublas_backward.launch_dw_only(
            stream,
            grn_h_s1_d_linear_b_out,
            h_s1_la_ptr,
            grad_buf_base + padded_byte_offset(param_sizes, 2),    // d_w_b_h_s1
            grad_buf_base + padded_byte_offset(param_sizes, 3),    // d_b_b_h_s1
            2 * sh1, sh1, b,
        )?;
        cublas_backward.launch_dx_only(
            stream,
            grn_h_s1_d_linear_b_out,
            w_ptrs[2],
            grn_h_s1_d_elu_out,
            2 * sh1, sh1, b, 0.0_f32,
        )?;

        // 8. Phase 2: ELU_bwd → grn_h_s1_d_linear_a_out.
        self.grn_h_s1_online.backward_raw_phase2(
            stream,
            grn_h_s1_d_elu_out,
            b,
            grn_h_s1_d_linear_a_out,
        )?;

        // 9. cuBLAS Linear_a backward (h_s1, padded states stride):
        //      dy = grn_h_s1_d_linear_a_out [B, SH1]
        //      x  = states_ptr [B, s1_in], lda = s1_lda (padded row stride)
        //      dW into goff[0], dB into goff[1], dX → s1_dx_output (beta=0).
        //      `s1_dx_output == 0` skips the dX GEMM entirely (no bottleneck).
        cublas_backward.backward_fc_layer_lda(
            stream,
            grn_h_s1_d_linear_a_out,
            states_ptr,
            w_ptrs[0],
            grad_buf_base + padded_byte_offset(param_sizes, 0),    // d_w_a_h_s1
            grad_buf_base + padded_byte_offset(param_sizes, 1),    // d_b_a_h_s1
            s1_dx_output,                                          // 0 ⇒ skip dX
            sh1, s1_in, s1_lda, b,
        )?;

        // 10. cuBLAS Linear_residual backward (h_s1, NO BIAS, padded stride):
        //       dy = grn_h_s1_d_pre_ln [B, SH1]   (LN-input gradient)
        //       x  = states_ptr [B, s1_in] with x_lda = s1_lda
        //       dW into goff[4]; no bias term for `Linear_residual`.
        cublas_backward.launch_dw_only_no_bias_lda(
            stream,
            grn_h_s1_d_pre_ln,
            states_ptr,
            grad_buf_base + padded_byte_offset(param_sizes, 4),    // d_w_residual_h_s1
            sh1, s1_in, s1_lda, b,
        )?;

        // 11. Linear_residual dX accumulation (only when bottleneck active):
        //       s1_dx_output += grn_h_s1_d_pre_ln @ w_residual^T (beta=1.0).
        //     The Linear_residual forward used `sgemm_f32_ldb` with the same
        //     `s1_lda` row stride, so the backward dX must also write at that
        //     stride — `backward_fc_layer_lda` did so for the Linear_a path,
        //     and we want this contribution to land at the same offsets.
        //     Reuse the lda-aware backward_fc_layer entry to write dX with
        //     beta=1 and skip dW/dB (already handled in step 10) — but the
        //     existing helpers don't expose a "dX-only with lda + beta" form.
        //     The plain `launch_dx_only` writes without lda; with bottleneck
        //     active, `s1_lda == s1_in` (no padding) so the two coincide.
        //     Without bottleneck, `s1_dx_output == 0` and we skip entirely.
        //     `debug_assert_eq!` guards the assumption.
        if s1_dx_output != 0 {
            debug_assert_eq!(
                s1_lda, s1_in,
                "encoder_backward_chain: Linear_residual dX accumulation needs \
                 s1_lda == s1_in_dim when bottleneck is active (cf. forward \
                 sgemm_f32_ldb with s1_ldb=s1_input_dim in the bottleneck branch)"
            );
            cublas_backward.launch_dx_only(
                stream,
                grn_h_s1_d_pre_ln,
                w_ptrs[4],
                s1_dx_output,
                sh1, s1_in, b, 1.0_f32,                            // accumulate
            )?;
        }

        Ok(())
    }

    /// Plan 4 Task 1B-iv: VSN feature-selection backward.
    ///
    /// Consumes the gradient flowing into the VSN-gated state (`d_gated_state`,
    /// assembled by the trainer from the bottleneck Linear backward dX into
    /// the market slice + the bottleneck-concat portfolio passthrough into
    /// `[market_dim..STATE_DIM]` + zero pad into `[STATE_DIM..STATE_DIM_PADDED]`)
    /// and writes per-group MLP weight + bias gradients into 24 grad_buf slots
    /// at indices `[95..119)` (the layout from 1B-ii).
    ///
    /// Sequence (all GPU; no atomicAdd, no DtoH):
    ///   1. `vsn_softmax_and_gate_backward` kernel:
    ///        d_gated_state, state, mask → d_logits[B, num_groups],
    ///                                     d_state[B, state_dim_padded] (discarded)
    ///   2. For each group g in [0, num_groups):
    ///      a. `strided_gather` extracts column g of d_logits[B, num_groups]
    ///         into `d_logit_scratch[B]` (contiguous tight buffer).
    ///      b. cuBLAS Linear_2 backward (`launch_dw_only` + `launch_dx_only`):
    ///           dY = d_logit_scratch[B, 1] (contiguous after gather),
    ///           X  = save_h1_g[g] [B, VSN_HIDDEN]
    ///           dW + dB into grad_buf[95+4g+2 / +3]; dX → d_h1_scratch[B, VSN_HIDDEN]
    ///      c. `relu_mask` on d_h1_scratch using save_h1_g[g] (post-ReLU activation
    ///         from forward — `relu_mask`'s "activation" arg is post-ReLU per kernel
    ///         contract).
    ///      d. cuBLAS Linear_1 backward via `backward_fc_layer_lda`:
    ///           dY = d_h1_scratch[B, VSN_HIDDEN],
    ///           X  = state[B, group_dim_g] sliced from state[B, state_dim_padded]
    ///                with row stride state_dim_padded
    ///           dW + dB into grad_buf[95+4g+0 / +1]; no dX (`d_state` was already
    ///           computed by step 1, and the value is discarded since the state
    ///           buffer is not trainable).
    ///
    /// `state_ptr` and `mask_ptr` are the saved buffers from the matching online
    /// `vsn_forward` call (the trainer keeps them in `states_buf` and
    /// `vsn_mask_buf` respectively). `save_h1_g_ptrs` are the 6 per-group
    /// post-ReLU saves the online forward populated.
    ///
    /// `d_logit_scratch_ptr` (tight `[B]`) and `d_h1_scratch_ptr`
    /// (`[B, VSN_HIDDEN_DIM]`) are reused across the 6 groups since the loop
    /// runs sequentially on the main stream.
    #[allow(clippy::too_many_arguments)]
    pub fn vsn_backward(
        &mut self,
        stream: &Arc<CudaStream>,
        cublas_backward: &super::batched_backward::CublasBackwardSet,
        d_gated_state_ptr: u64,                  // INPUT [B, STATE_DIM_PADDED]
        state_ptr: u64,                           // INPUT [B, STATE_DIM_PADDED] saved
        mask_ptr: u64,                            // INPUT [B, num_groups] saved
        save_h1_g_ptrs: &[u64; ml_core::state_layout::SL_NUM_FEATURE_GROUPS],
        w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
        grad_buf_base: u64,
        param_sizes: &[usize],
        d_logits_ptr: u64,                        // SCRATCH [B, num_groups]
        d_state_ptr: u64,                         // SCRATCH [B, STATE_DIM_PADDED] (kernel writes; discarded)
        d_logit_scratch_ptr: u64,                 // SCRATCH [B] (per-group, reused)
        d_h1_scratch_ptr: u64,                    // SCRATCH [B, VSN_HIDDEN_DIM] (per-group, reused)
        group_begins_dev: u64,
        group_ends_dev: u64,
        gather_kernel: &CudaFunction,             // strided_gather (column extraction)
        batch_size: usize,
    ) -> Result<(), MLError> {
        use super::gpu_dqn_trainer::{padded_byte_offset, VSN_HIDDEN_DIM};
        const VSN_TENSORS_BASE: usize = 95;
        const NUM_GROUPS: usize = ml_core::state_layout::SL_NUM_FEATURE_GROUPS;

        let group_ranges = ml_core::state_layout::FEATURE_GROUP_RANGES;
        let state_dim_padded = self.state_dim_padded;
        let f32_size = std::mem::size_of::<f32>() as u64;

        // ── Step 1: softmax-and-gate backward ────────────────────────────────
        // Writes d_logits[B, num_groups] (consumed by per-group Linear_2 bwd
        // via gather) and d_state[B, state_dim_padded] (discarded — state buf
        // not trainable). Block size 256, grid size B; shared mem layout
        // (per kernel comment): 2*num_groups + block floats.
        {
            let b_i32 = batch_size as i32;
            let state_dim_padded_i32 = state_dim_padded as i32;
            let num_groups_i32 = NUM_GROUPS as i32;
            let smem_floats = 2 * NUM_GROUPS + 256;
            let smem_bytes = (smem_floats * std::mem::size_of::<f32>()) as u32;
            unsafe {
                stream
                    .launch_builder(&self.vsn_softmax_and_gate_bwd_kernel)
                    .arg(&d_gated_state_ptr)
                    .arg(&state_ptr)
                    .arg(&mask_ptr)
                    .arg(&b_i32)
                    .arg(&state_dim_padded_i32)
                    .arg(&num_groups_i32)
                    .arg(&group_begins_dev)
                    .arg(&group_ends_dev)
                    .arg(&d_logits_ptr)
                    .arg(&d_state_ptr)
                    .launch(LaunchConfig {
                        grid_dim: (batch_size as u32, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: smem_bytes,
                    })
                    .map_err(|e| MLError::ModelError(format!(
                        "vsn_softmax_and_gate_backward: {e}"
                    )))?;
            }
        }

        // ── Step 2: per-group Linear_2 backward + ReLU mask + Linear_1 bwd ───
        for g in 0..NUM_GROUPS {
            let (gb, ge) = group_ranges[g];
            let group_dim = ge - gb;
            debug_assert!(
                group_dim > 0,
                "vsn_backward: feature group {g} has zero dim — group ranges invariant violated"
            );

            // (a) Gather column g of d_logits[B, num_groups] → d_logit_scratch[B].
            // Column-extraction via strided_gather: dst_stride=1, col=g,
            // src_lda=num_groups, total = B * dst_stride = B.
            let src_lda_i32 = NUM_GROUPS as i32;
            let col_i32 = g as i32;
            let dst_stride_i32: i32 = 1;
            let total_i32 = batch_size as i32;
            let blocks = ((batch_size as u32 + 255) / 256).max(1);
            unsafe {
                stream
                    .launch_builder(gather_kernel)
                    .arg(&d_logits_ptr)
                    .arg(&d_logit_scratch_ptr)
                    .arg(&src_lda_i32)
                    .arg(&col_i32)
                    .arg(&dst_stride_i32)
                    .arg(&total_i32)
                    .launch(LaunchConfig {
                        grid_dim: (blocks, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| MLError::ModelError(format!(
                        "vsn_backward strided_gather group {g}: {e}"
                    )))?;
            }

            // (b) Linear_2 backward: dY = d_logit_scratch[B, 1] (tight),
            //     X = save_h1_g[g] [B, VSN_HIDDEN] (tight),
            //     dW [1, VSN_HIDDEN] into grad_buf[95+4g+2],
            //     dB [1] into grad_buf[95+4g+3],
            //     dX [B, VSN_HIDDEN] into d_h1_scratch (overwrite, beta=0).
            let goff_w2 = padded_byte_offset(param_sizes, VSN_TENSORS_BASE + 4 * g + 2);
            let goff_b2 = padded_byte_offset(param_sizes, VSN_TENSORS_BASE + 4 * g + 3);
            cublas_backward.launch_dw_only(
                stream,
                d_logit_scratch_ptr,                   // dY [B, 1]
                save_h1_g_ptrs[g],                     // X  [B, VSN_HIDDEN]
                grad_buf_base + goff_w2,               // dW [1, VSN_HIDDEN]
                grad_buf_base + goff_b2,               // dB [1]
                1, VSN_HIDDEN_DIM, batch_size,
            )?;
            cublas_backward.launch_dx_only(
                stream,
                d_logit_scratch_ptr,                   // dY [B, 1]
                w_ptrs[VSN_TENSORS_BASE + 4 * g + 2],  // W2 [1, VSN_HIDDEN]
                d_h1_scratch_ptr,                      // dX [B, VSN_HIDDEN]
                1, VSN_HIDDEN_DIM, batch_size, 0.0_f32,
            )?;

            // (c) ReLU mask on d_h1_scratch using saved post-ReLU h1.
            cublas_backward.relu_mask(
                stream,
                d_h1_scratch_ptr,
                save_h1_g_ptrs[g],
                batch_size * VSN_HIDDEN_DIM,
            )?;

            // (d) Linear_1 backward (state slice with x_lda = state_dim_padded):
            //     dY = d_h1_scratch [B, VSN_HIDDEN] (tight),
            //     X  = state + gb*sizeof(f32) (group-slice column offset),
            //          x_lda = state_dim_padded (row stride of full state row),
            //     dW [VSN_HIDDEN, group_dim_g] into grad_buf[95+4g+0],
            //     dB [VSN_HIDDEN] into grad_buf[95+4g+1],
            //     no dX (d_state was computed in step 1; state buf untrainable).
            let goff_w1 = padded_byte_offset(param_sizes, VSN_TENSORS_BASE + 4 * g + 0);
            let goff_b1 = padded_byte_offset(param_sizes, VSN_TENSORS_BASE + 4 * g + 1);
            let state_slice_ptr = state_ptr + (gb as u64) * f32_size;
            cublas_backward.backward_fc_layer_lda(
                stream,
                d_h1_scratch_ptr,                      // dY [B, VSN_HIDDEN]
                state_slice_ptr,                       // X  [B, group_dim] @ stride=state_dim_padded
                w_ptrs[VSN_TENSORS_BASE + 4 * g + 0],  // W1 [VSN_HIDDEN, group_dim] (unused for dX-skip path)
                grad_buf_base + goff_w1,               // dW [VSN_HIDDEN, group_dim]
                grad_buf_base + goff_b1,               // dB [VSN_HIDDEN]
                0,                                      // dX skipped (state buf untrainable)
                VSN_HIDDEN_DIM,
                group_dim,
                state_dim_padded,                      // x_lda
                batch_size,
            )?;
        }

        Ok(())
    }

    /// Run a single advantage branch's decoder (FC head + adv-logits) on the
    /// caller-supplied stream, sequentially.
    ///
    /// Plan 4 Task 4 (E.4) attachment point for the per-branch `ValueDecoder`
    /// Rust wrapper. **Additive** to `forward_online_raw`: this helper runs
    /// the same GEMM sequence the per-branch loop body in `forward_online_raw`
    /// runs (when `use_vsn_glu == false`), but on a single stream — no fork-
    /// join, no event handshake. It is intended for callers that drive the
    /// branches one at a time (e.g. validation paths, future per-branch
    /// experiments). Production training stays on the existing multi-stream
    /// `forward_online_raw` path until/unless a follow-up explicitly
    /// migrates it.
    ///
    /// `branch_idx ∈ [0, 4)` selects which branch's weights/buffers to use.
    /// `h_s2_ptr` is the encoder output (read-only). For `branch_idx == 1`
    /// (magnitude), pass the pre-built mag_concat pointer to use the wider
    /// `[B, SH2 + branch_0_size]` input (direction-conditioned); pass 0 for
    /// the legacy `[B, SH2]` input. For `branch_idx == 0` (direction),
    /// pass `dir_qaux_concat_ptr` (SP14 EGF) to use `[B, SH2 + 1]`; pass 0
    /// for the legacy `[B, SH2]` input.
    /// `branch_h_ptr` is the per-branch hidden activation buffer; the
    /// `adv_logits_ptr` is the per-branch slice into the flat
    /// `b_logits_buf` at byte offset
    /// `sum_{j<branch_idx}(B * branch_size_j * num_atoms * 4)`.
    #[allow(clippy::too_many_arguments)]
    pub fn decoder_forward_only(
        &self,
        stream: &Arc<CudaStream>,
        branch_idx: usize,
        h_s2_ptr: u64,
        mag_concat_ptr: u64,
        dir_qaux_concat_ptr: u64,
        w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
        branch_h_ptr: u64,
        adv_logits_ptr: u64,
    ) -> Result<(), MLError> {
        if branch_idx >= 4 {
            return Err(MLError::ModelError(format!(
                "decoder_forward_only: branch_idx {branch_idx} out of range [0, 4)"
            )));
        }
        let b = self.batch_size;
        let ws = self.handle.lt_workspace_ptr;
        let wss = self.handle.lt_workspace_size;
        let na = self.num_atoms;
        // Plan 4 Task 2c.3a: branch base indices +9 (b0fc=17, b1fc=21, b2fc=25, b3fc=29).
        let branch_w_base = [17_usize, 21, 25, 29];
        let w_fc_idx = branch_w_base[branch_idx];
        let n_d = self.branch_q_dim_internal(branch_idx);

        let use_vsn_glu = self.vsn_kernel.is_some();

        if use_vsn_glu {
            // VSN bottleneck + GLU gating path (sequential — main stream).
            self.launch_vsn_glu_branch(
                stream, branch_idx, h_s2_ptr, mag_concat_ptr, dir_qaux_concat_ptr, w_ptrs,
                branch_h_ptr,
                ws, wss,
                false, "decoder",
            )?;
        } else {
            // Legacy GEMM+bias+ReLU FC head.
            let (fc_input, fc_k) = if branch_idx == 0 && dir_qaux_concat_ptr != 0 {
                // SP14 EGF: direction Q-head reads [h_s2 ; aux_softmax_diff].
                (dir_qaux_concat_ptr, self.shared_h2 + 1)
            } else if branch_idx == 1 && mag_concat_ptr != 0 {
                (mag_concat_ptr, self.mag_concat_dim)
            } else {
                (h_s2_ptr, self.shared_h2)
            };
            if self.sgemm_f32_fused_relu_bias(stream, w_ptrs[w_fc_idx], fc_input, branch_h_ptr, w_ptrs[w_fc_idx + 1],
                self.adv_h, b, fc_k, fc_k, ws, wss, "decoder_h_bd").is_err()
            {
                self.sgemm_f32(stream, w_ptrs[w_fc_idx], fc_input, branch_h_ptr, self.adv_h, b, fc_k, "decoder_h_bd")?;
                self.launch_add_bias_relu_f32_raw(stream, branch_h_ptr, w_ptrs[w_fc_idx + 1], self.adv_h, b)?;
            }
        }

        // Output GEMM + bias (linear): adv_logits = h_bd @ W_out^T + b_out
        if self.sgemm_f32_fused_bias(stream, w_ptrs[w_fc_idx + 2], branch_h_ptr, adv_logits_ptr, w_ptrs[w_fc_idx + 3],
            n_d * na, b, self.adv_h, self.adv_h, ws, wss, "decoder_adv_logits").is_err()
        {
            self.sgemm_f32(stream, w_ptrs[w_fc_idx + 2], branch_h_ptr, adv_logits_ptr, n_d * na, b, self.adv_h, "decoder_adv_logits")?;
            self.launch_add_bias_f32_raw(stream, adv_logits_ptr, w_ptrs[w_fc_idx + 3], n_d * na, b)?;
        }

        Ok(())
    }

    /// Phase H Site 2 — register the trainer's "main online save_h_v" pointer.
    /// The forward only takes the `RELU_AUX_BIAS` path (writing the dReLU mask)
    /// when `h_v_ptr` argument matches this registered marker. Auxiliary
    /// forwards that target scratch h_v buffers (causal-intervention, IQN,
    /// DDQN, ensemble) get the plain `RELU_BIAS` path, so the mask consumed
    /// by backward isn't clobbered between the main forward and backward.
    pub fn wire_value_fc_save_h_v(&mut self, save_h_v_ptr: u64) {
        self.value_fc_save_h_v_marker = save_h_v_ptr;
    }

    /// Phase H Site 2 — raw u64 pointer to the value-FC dReLU mask aux buffer.
    /// Read by `CublasBackwardSet::backward_fc_layer_drelu_bgrad` via
    /// `EPILOGUE_AUX_POINTER`. Always non-zero (buffer is allocated at
    /// construction); the backward only consumes it when its DRELU_BGRAD
    /// descriptor is present (cached at init).
    pub fn value_fc_relu_mask_ptr(&self) -> u64 {
        self.value_fc_relu_mask_ptr_raw
    }

    /// Phase H Site 2 — AUX leading dimension (in elements/bits) for the
    /// value-FC ReLU mask. Matches the `EPILOGUE_AUX_LD` baked into both the
    /// forward `RELU_AUX_BIAS` descriptor and the backward `DRELU_BGRAD`
    /// descriptor.
    pub fn value_fc_relu_mask_aux_ld(&self) -> usize {
        self.value_fc_relu_mask_aux_ld_raw
    }

    /// Phase H Site 2 — whether the forward path actually wrote a valid
    /// dReLU bit-mask in the previous online forward pass. Returns `true`
    /// iff the cached `RELU_AUX_BIAS` descriptor for the value-FC shape is
    /// available — i.e. the forward took the fused path, not a fallback.
    pub fn value_fc_relu_mask_available(&self) -> bool {
        let key: FwdGemmKey = (self.value_h, self.batch_size, self.shared_h2, self.shared_h2);
        self.gemm_cache_relu_aux_bias.contains_key(&key)
    }

    /// Per-branch advantage dimension (action count for that branch — pre
    /// num_atoms multiplication). Internal so `ValueDecoder` can report
    /// `q_per_action_dim` without exposing the four `branch_k_size` fields.
    pub(crate) fn branch_q_dim_internal(&self, branch_idx: usize) -> usize {
        match branch_idx {
            0 => self.branch_0_size,
            1 => self.branch_1_size,
            2 => self.branch_2_size,
            3 => self.branch_3_size,
            _ => 0,
        }
    }

    /// Distributional atoms per branch action (C51 support size).
    pub(crate) fn num_atoms(&self) -> usize { self.num_atoms }

    /// Graph-safe forward: takes pre-resolved u64 buffer pointers (no device_ptr calls).
    /// Use CachedPtrs from GpuDqnTrainer for all buffer addresses.
    ///
    /// `dir_qaux_concat_ptr` — SP14 Layer B (Earned Gradient Flow): when non-zero,
    /// the direction Q-head's first FC SGEMM (d == 0) consumes this `[B, SH2 + 1]`
    /// scratch buffer (last column = `aux_softmax_diff = up_prob - down_prob`)
    /// instead of `h_s2_ptr` directly. Caller must launch
    /// `dir_concat_qaux_kernel` AFTER aux head forward AND BEFORE this call,
    /// on the same stream — the launch order enforces freshness per
    /// `pearl_canary_input_freshness_launch_order`. Pass 0 for paths that should
    /// keep the legacy `[B, SH2]` direct-h_s2 input (causal intervention,
    /// experience collector, DDQN argmax pass — they have no aux dependency).
    #[allow(clippy::too_many_arguments)]
    pub fn forward_online_raw(
        &mut self,
        stream: &Arc<CudaStream>,
        states_ptr: u64,
        w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
        h_s1_ptr: u64, h_s2_ptr: u64, h_v_ptr: u64,
        h_b0_ptr: u64, h_b1_ptr: u64, h_b2_ptr: u64, h_b3_ptr: u64,
        v_logits_ptr: u64, b_logits_ptr: u64,
        mag_concat_ptr: u64,
        dir_qaux_concat_ptr: u64,
    ) -> Result<(), MLError> {
        let b = self.batch_size;
        let ws = self.handle.lt_workspace_ptr;
        let wss = self.handle.lt_workspace_size;

        // Trunk encoder (h_s1, h_s2) — extracted to `encoder_forward_only` so
        // the new `StateEncoder` Rust wrapper (Plan 4 Task 4 / E.4) can reuse
        // the identical dispatch sequence. Behaviorally a no-op refactor.
        // Plan 4 Task 2c.3a: encoder_forward_only panics — runtime is
        // intentionally trapped here until 2c.3b swaps in GRN forward.
        self.encoder_forward_only(stream, states_ptr, w_ptrs, h_s1_ptr, h_s2_ptr)?;

        // Plan 4 Task 2c.3a: value-head index +9 (w_v1=13, b_v1=14, w_v2=15, b_v2=16).
        //
        // Phase H Site 2: take the RELU_AUX_BIAS path ONLY when the caller's
        // `h_v_ptr` matches the trainer-registered "main save_h_v" buffer
        // (`value_fc_save_h_v_marker`). Auxiliary forwards (causal-intervention,
        // IQN, DDQN, ensemble) target scratch h_v buffers; running the AUX path
        // for those would clobber the dReLU mask the backward reads. Those
        // callers fall through to RELU_BIAS / manual fallback — backward
        // detects the stale state via `value_fc_relu_mask_available()` (false
        // until a marker-matching forward runs).
        let take_aux_path = h_v_ptr != 0
            && h_v_ptr == self.value_fc_save_h_v_marker;
        let aux_failed = if take_aux_path {
            self.sgemm_f32_fused_relu_aux_bias(stream, w_ptrs[13], h_s2_ptr, h_v_ptr, w_ptrs[14],
                self.value_fc_relu_mask_ptr_raw,
                self.value_h, b, self.shared_h2, self.shared_h2, ws, wss, "h_v").is_err()
        } else {
            true // skip AUX, go straight to RELU_BIAS fallback
        };
        if aux_failed {
            // No-AUX path: plain RELU_BIAS, then manual kernels. Backward will
            // see no mask (because forward did not write one this step) and
            // use the standalone relu_mask + bias_grad sequence — i.e., when
            // this branch fires for the marker-matching forward, backward
            // falls back to the slow path automatically via the trainer's
            // `value_fc_relu_mask_available()` gate.
            if self.sgemm_f32_fused_relu_bias(stream, w_ptrs[13], h_s2_ptr, h_v_ptr, w_ptrs[14],
                self.value_h, b, self.shared_h2, self.shared_h2, ws, wss, "h_v").is_err()
            {
                self.sgemm_f32(stream, w_ptrs[13], h_s2_ptr, h_v_ptr, self.value_h, b, self.shared_h2, "h_v")?;
                self.launch_add_bias_relu_f32_raw(stream, h_v_ptr, w_ptrs[14], self.value_h, b)?;
            }
        }

        // Output layer: cublasLtMatmul writes f32 C-matrix (no f32 truncation overflow)
        if self.sgemm_f32_fused_bias(stream, w_ptrs[15], h_v_ptr, v_logits_ptr, w_ptrs[16],
            self.num_atoms, b, self.value_h, self.value_h, ws, wss, "v_logits").is_err()
        {
            self.sgemm_f32(stream, w_ptrs[15], h_v_ptr, v_logits_ptr, self.num_atoms, b, self.value_h, "v_logits")?;
            self.launch_add_bias_f32_raw(stream, v_logits_ptr, w_ptrs[16], self.num_atoms, b)?;
        }

        let branch_sizes = [self.branch_0_size, self.branch_1_size, self.branch_2_size, self.branch_3_size];
        let branch_h_ptrs = [h_b0_ptr, h_b1_ptr, h_b2_ptr, h_b3_ptr];
        // Plan 4 Task 2c.3a: branch base indices +9 (b0fc=17, b1fc=21, b2fc=25, b3fc=29).
        let branch_w_base = [17_usize, 21, 25, 29];
        let na = self.num_atoms;

        // Multi-stream only when branch hidden buffers are distinct (no aliasing).
        // Pass 3 (Double DQN on next_states) reuses a single scratch for all 4
        // branches — must stay sequential to avoid concurrent writes.
        let distinct_branches = h_b0_ptr != h_b1_ptr && h_b1_ptr != h_b2_ptr && h_b0_ptr != h_b2_ptr
            && h_b0_ptr != h_b3_ptr && h_b1_ptr != h_b3_ptr && h_b2_ptr != h_b3_ptr;

        let use_vsn_glu = self.vsn_kernel.is_some();

        if distinct_branches {
            // Multi-stream branch dispatch: submit each branch's GEMMs to a
            // dedicated CUDA stream so all 4 advantage heads execute in parallel.
            // Safe because branches read h_s2 (shared trunk, read-only) and write
            // to non-overlapping regions of branch_h_ptrs and b_logits_buf.
            //
            // Works inside CUDA Graph capture (CUDA 12+): cuStreamWaitEvent
            // on a non-captured stream makes it join the capture graph, and the
            // fork/join pattern is recorded as graph dependencies.
            //
            // cublasLtMatmul takes stream per-call — no cublasSetStream needed.
            // This eliminates the workspace state conflict that hangs CUDA Graph
            // mega-capture on H100.

            // Record trunk completion on the main stream (pre-allocated event).
            self.trunk_done_event.record(stream)
                .map_err(|e| MLError::ModelError(format!("trunk event record: {e}")))?;

            let mut logit_byte_offset: u64 = 0;
            for d in 0..4 {
                let bs = &self.branch_streams[d];
                let n_d = branch_sizes[d];
                let w_fc_idx = branch_w_base[d];

                // Branch stream waits for trunk completion.
                bs.wait(&self.trunk_done_event)
                    .map_err(|e| MLError::ModelError(format!("branch {d} wait trunk: {e}")))?;

                if use_vsn_glu {
                    // VSN bottleneck + GLU gating path
                    self.launch_vsn_glu_branch(
                        bs, d, h_s2_ptr, mag_concat_ptr, dir_qaux_concat_ptr, w_ptrs,
                        branch_h_ptrs[d],
                        self.branch_workspace_ptrs[d], self.handle.lt_workspace_size,
                        true, "on",
                    )?;
                } else {
                    // Legacy GEMM+bias+ReLU path (experience collector, unconditioned)
                    let bws = self.branch_workspace_ptrs[d];
                    let bwss = self.handle.lt_workspace_size;
                    let (fc_input, fc_k) = if d == 0 && dir_qaux_concat_ptr != 0 {
                        // SP14 EGF: direction Q-head reads [h_s2 ; aux_softmax_diff].
                        (dir_qaux_concat_ptr, self.shared_h2 + 1)
                    } else if d == 1 && mag_concat_ptr != 0 {
                        (mag_concat_ptr, self.mag_concat_dim)
                    } else {
                        (h_s2_ptr, self.shared_h2)
                    };
                    if self.sgemm_f32_fused_relu_bias(bs, w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1],
                        self.adv_h, b, fc_k, fc_k, bws, bwss, "h_bd").is_err()
                    {
                        self.sgemm_f32_branch(bs, w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], self.adv_h, b, fc_k, d, "h_bd")?;
                        self.launch_add_bias_relu_f32_raw(bs, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], self.adv_h, b)?;
                    }
                }

                let adv_out_ptr = b_logits_ptr + logit_byte_offset;
                let out_bws = self.branch_workspace_ptrs[d];
                let out_bwss = self.handle.lt_workspace_size;
                if self.sgemm_f32_fused_bias(bs, w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, w_ptrs[w_fc_idx + 3],
                    n_d * na, b, self.adv_h, self.adv_h, out_bws, out_bwss, "adv_logits").is_err()
                {
                    self.sgemm_f32_branch(bs, w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, n_d * na, b, self.adv_h, d, "adv_logits")?;
                    self.launch_add_bias_f32_raw(bs, adv_out_ptr, w_ptrs[w_fc_idx + 3], n_d * na, b)?;
                }

                logit_byte_offset += (b * n_d * na * std::mem::size_of::<f32>()) as u64;
            }

            // Join: main stream waits for all 4 branches to complete (pre-allocated events).
            for d in 0..4 {
                self.branch_done_events[d].record(&self.branch_streams[d])
                    .map_err(|e| MLError::ModelError(format!("branch {d} done event: {e}")))?;
                stream.wait(&self.branch_done_events[d])
                    .map_err(|e| MLError::ModelError(format!("main wait branch {d}: {e}")))?;
            }
        } else {
            // Sequential fallback: branch hidden buffers alias (scratch reuse).
            let mut logit_byte_offset: u64 = 0;
            for d in 0..4 {
                let n_d = branch_sizes[d];
                let w_fc_idx = branch_w_base[d];

                if use_vsn_glu {
                    // VSN bottleneck + GLU gating path (sequential)
                    self.launch_vsn_glu_branch(
                        stream, d, h_s2_ptr, mag_concat_ptr, dir_qaux_concat_ptr, w_ptrs,
                        branch_h_ptrs[d],
                        self.handle.lt_workspace_ptr, self.handle.lt_workspace_size,
                        false, "on_seq",
                    )?;
                } else {
                    // Legacy GEMM+bias+ReLU path — try fused epilogue first.
                    let (fc_input, fc_k) = if d == 0 && dir_qaux_concat_ptr != 0 {
                        // SP14 EGF: direction Q-head reads [h_s2 ; aux_softmax_diff].
                        (dir_qaux_concat_ptr, self.shared_h2 + 1)
                    } else if d == 1 && mag_concat_ptr != 0 {
                        (mag_concat_ptr, self.mag_concat_dim)
                    } else {
                        (h_s2_ptr, self.shared_h2)
                    };
                    if self.sgemm_f32_fused_relu_bias(stream, w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1],
                        self.adv_h, b, fc_k, fc_k, ws, wss, "h_bd").is_err()
                    {
                        self.sgemm_f32(stream, w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], self.adv_h, b, fc_k, "h_bd")?;
                        self.launch_add_bias_relu_f32_raw(stream, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], self.adv_h, b)?;
                    }
                }

                let adv_out_ptr = b_logits_ptr + logit_byte_offset;
                if self.sgemm_f32_fused_bias(stream, w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, w_ptrs[w_fc_idx + 3],
                    n_d * na, b, self.adv_h, self.adv_h, ws, wss, "adv_logits").is_err()
                {
                    self.sgemm_f32(stream, w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, n_d * na, b, self.adv_h, "adv_logits")?;
                    self.launch_add_bias_f32_raw(stream, adv_out_ptr, w_ptrs[w_fc_idx + 3], n_d * na, b)?;
                }

                logit_byte_offset += (b * n_d * na * std::mem::size_of::<f32>()) as u64;
            }
        }

        Ok(())
    }

    /// #30 Pure F32 forward pass for experience collector.
    /// All weights, activations, and logits are f32. No f32 truncation.
    /// Used during experience collection for maximum Q-value precision.
    ///
    /// States must be f32 [B, state_dim_padded]. Weights are f32 [total_params].
    /// All hidden activations are f32. Output logits are f32 (same as f32 path).
    #[allow(clippy::too_many_arguments)]
    pub fn forward_online_f32(
        &mut self,
        stream: &Arc<CudaStream>,
        states_f32_ptr: u64,              // F32 states [B, state_dim_padded]
        w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS], // F32 weight pointers
        h_s1_ptr: u64, h_s2_ptr: u64, h_v_ptr: u64, // F32 activation buffers
        h_b0_ptr: u64, h_b1_ptr: u64, h_b2_ptr: u64, h_b3_ptr: u64,
        v_logits_ptr: u64, b_logits_ptr: u64,        // F32 output logits
        mag_concat_ptr: u64,
        dir_qaux_concat_ptr: u64,                    // SP14 EGF direction Q-head input concat (0 = unused).
    ) -> Result<(), MLError> {
        let b = self.batch_size;
        let ws = self.handle.lt_workspace_ptr;
        let wss = self.handle.lt_workspace_size;

        // Plan 4 Task 2c.3b: trunk encoder via GRN composition (online
        // network — uses the online GrnBlocks' save buffers). The
        // experience collector calls this from outside the training step;
        // its save state is unused (no backward against the collector
        // path) but is still written for layout consistency. The next
        // online encoder forward overwrites it.
        self.encoder_forward_only(stream, states_f32_ptr, w_ptrs, h_s1_ptr, h_s2_ptr)?;

        // Plan 4 Task 2c.3a: value-head index +9 (w_v1=13, b_v1=14).
        if self.sgemm_f32_fused_relu_bias(stream, w_ptrs[13], h_s2_ptr, h_v_ptr, w_ptrs[14],
            self.value_h, b, self.shared_h2, self.shared_h2, ws, wss, "f32_h_v").is_err()
        {
            self.sgemm_f32(stream, w_ptrs[13], h_s2_ptr, h_v_ptr, self.value_h, b, self.shared_h2, "f32_h_v")?;
            self.launch_add_bias_relu_f32_raw(stream, h_v_ptr, w_ptrs[14], self.value_h, b)?;
        }

        // Value output: linear (no ReLU) — w_v2=15, b_v2=16.
        if self.sgemm_f32_fused_bias(stream, w_ptrs[15], h_v_ptr, v_logits_ptr, w_ptrs[16],
            self.num_atoms, b, self.value_h, self.value_h, ws, wss, "f32_v_logits").is_err()
        {
            self.sgemm_f32(stream, w_ptrs[15], h_v_ptr, v_logits_ptr, self.num_atoms, b, self.value_h, "f32_v_logits")?;
            self.launch_add_bias_f32_raw(stream, v_logits_ptr, w_ptrs[16], self.num_atoms, b)?;
        }

        let branch_sizes = [self.branch_0_size, self.branch_1_size, self.branch_2_size, self.branch_3_size];
        let branch_h_ptrs = [h_b0_ptr, h_b1_ptr, h_b2_ptr, h_b3_ptr];
        // Plan 4 Task 2c.3a: branch base indices +9 (b0fc=17, b1fc=21, b2fc=25, b3fc=29).
        let branch_w_base = [17_usize, 21, 25, 29];
        let na = self.num_atoms;

        let distinct_branches = h_b0_ptr != h_b1_ptr && h_b1_ptr != h_b2_ptr && h_b0_ptr != h_b2_ptr
            && h_b0_ptr != h_b3_ptr && h_b1_ptr != h_b3_ptr && h_b2_ptr != h_b3_ptr;

        let use_vsn_glu = self.vsn_kernel.is_some();

        if distinct_branches {
            // Multi-stream branch dispatch — cublasLtMatmul takes stream per-call.
            self.trunk_done_event.record(stream)
                .map_err(|e| MLError::ModelError(format!("f32 trunk event record: {e}")))?;

            let mut logit_byte_offset: u64 = 0;
            for d in 0..4 {
                let bs = &self.branch_streams[d];
                let n_d = branch_sizes[d];
                let w_fc_idx = branch_w_base[d];

                bs.wait(&self.trunk_done_event)
                    .map_err(|e| MLError::ModelError(format!("f32 branch {d} wait trunk: {e}")))?;

                if use_vsn_glu {
                    self.launch_vsn_glu_branch(
                        bs, d, h_s2_ptr, mag_concat_ptr, dir_qaux_concat_ptr, w_ptrs,
                        branch_h_ptrs[d],
                        self.branch_workspace_ptrs[d], self.handle.lt_workspace_size,
                        true, "f32",
                    )?;
                } else {
                    let bws = self.branch_workspace_ptrs[d];
                    let bwss = self.handle.lt_workspace_size;
                    let (fc_input, fc_k) = if d == 0 && dir_qaux_concat_ptr != 0 {
                        // SP14 EGF: direction Q-head reads [h_s2 ; aux_softmax_diff].
                        (dir_qaux_concat_ptr, self.shared_h2 + 1)
                    } else if d == 1 && mag_concat_ptr != 0 {
                        (mag_concat_ptr, self.mag_concat_dim)
                    } else {
                        (h_s2_ptr, self.shared_h2)
                    };
                    if self.sgemm_f32_fused_relu_bias(bs, w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1],
                        self.adv_h, b, fc_k, fc_k, bws, bwss, "f32_h_bd").is_err()
                    {
                        self.sgemm_f32_branch(bs, w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], self.adv_h, b, fc_k, d, "f32_h_bd")?;
                        self.launch_add_bias_relu_f32_raw(bs, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], self.adv_h, b)?;
                    }
                }

                let adv_out_ptr = b_logits_ptr + logit_byte_offset;
                let out_bws = self.branch_workspace_ptrs[d];
                let out_bwss = self.handle.lt_workspace_size;
                if self.sgemm_f32_fused_bias(bs, w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, w_ptrs[w_fc_idx + 3],
                    n_d * na, b, self.adv_h, self.adv_h, out_bws, out_bwss, "f32_adv_logits").is_err()
                {
                    self.sgemm_f32_branch(bs, w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, n_d * na, b, self.adv_h, d, "f32_adv_logits")?;
                    self.launch_add_bias_f32_raw(bs, adv_out_ptr, w_ptrs[w_fc_idx + 3], n_d * na, b)?;
                }

                logit_byte_offset += (b * n_d * na * std::mem::size_of::<f32>()) as u64;
            }

            for d in 0..4 {
                self.branch_done_events[d].record(&self.branch_streams[d])
                    .map_err(|e| MLError::ModelError(format!("f32 branch {d} done event: {e}")))?;
                stream.wait(&self.branch_done_events[d])
                    .map_err(|e| MLError::ModelError(format!("f32 main wait branch {d}: {e}")))?;
            }
        } else {
            // Sequential fallback: branch hidden buffers alias.
            let mut logit_byte_offset: u64 = 0;
            for d in 0..4 {
                let n_d = branch_sizes[d];
                let w_fc_idx = branch_w_base[d];

                if use_vsn_glu {
                    self.launch_vsn_glu_branch(
                        stream, d, h_s2_ptr, mag_concat_ptr, dir_qaux_concat_ptr, w_ptrs,
                        branch_h_ptrs[d],
                        self.handle.lt_workspace_ptr, self.handle.lt_workspace_size,
                        false, "f32_seq",
                    )?;
                } else {
                    let (fc_input, fc_k) = if d == 0 && dir_qaux_concat_ptr != 0 {
                        // SP14 EGF: direction Q-head reads [h_s2 ; aux_softmax_diff].
                        (dir_qaux_concat_ptr, self.shared_h2 + 1)
                    } else if d == 1 && mag_concat_ptr != 0 {
                        (mag_concat_ptr, self.mag_concat_dim)
                    } else {
                        (h_s2_ptr, self.shared_h2)
                    };
                    if self.sgemm_f32_fused_relu_bias(stream, w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1],
                        self.adv_h, b, fc_k, fc_k, ws, wss, "f32_h_bd").is_err()
                    {
                        self.sgemm_f32(stream, w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], self.adv_h, b, fc_k, "f32_h_bd")?;
                        self.launch_add_bias_relu_f32_raw(stream, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], self.adv_h, b)?;
                    }
                }

                let adv_out_ptr = b_logits_ptr + logit_byte_offset;
                if self.sgemm_f32_fused_bias(stream, w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, w_ptrs[w_fc_idx + 3],
                    n_d * na, b, self.adv_h, self.adv_h, ws, wss, "f32_adv_logits").is_err()
                {
                    self.sgemm_f32(stream, w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, n_d * na, b, self.adv_h, "f32_adv_logits")?;
                    self.launch_add_bias_f32_raw(stream, adv_out_ptr, w_ptrs[w_fc_idx + 3], n_d * na, b)?;
                }

                logit_byte_offset += (b * n_d * na * std::mem::size_of::<f32>()) as u64;
            }
        }

        Ok(())
    }

    /// Forward pass using CudaSlice references (calls device_ptr — NOT graph-safe).
    /// Use forward_online_raw with CachedPtrs for graph-captured code.
    pub fn forward_online(
        &mut self,
        stream: &Arc<CudaStream>,
        states_buf: &CudaSlice<f32>,
        w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
        save_h_s1: &CudaSlice<f32>,
        save_h_s2: &CudaSlice<f32>,
        save_h_v: &CudaSlice<f32>,
        save_h_b0: &CudaSlice<f32>,
        save_h_b1: &CudaSlice<f32>,
        save_h_b2: &CudaSlice<f32>,
        save_h_b3: &CudaSlice<f32>,
        v_logits_buf: &CudaSlice<f32>,
        b_logits_buf: &CudaSlice<f32>,
    ) -> Result<(), MLError> {
        self.forward_online_raw(
            stream,
            raw_f32_ptr(states_buf, stream),
            w_ptrs,
            raw_f32_ptr(save_h_s1, stream),
            raw_f32_ptr(save_h_s2, stream),
            raw_f32_ptr(save_h_v, stream),
            raw_f32_ptr(save_h_b0, stream),
            raw_f32_ptr(save_h_b1, stream),
            raw_f32_ptr(save_h_b2, stream),
            raw_f32_ptr(save_h_b3, stream),
            raw_f32_ptr(v_logits_buf, stream),
            raw_f32_ptr(b_logits_buf, stream),
            0u64, // mag_concat_ptr: not used in CudaSlice wrapper path
            0u64, // dir_qaux_concat_ptr: SP14 EGF disabled in CudaSlice wrapper path
        )
    }

    /// Run value head forward only: h_s2 → W_v1 → ReLU → W_v2 → v_logits.
    ///
    /// Used by ensemble heads (1..K-1) to compute per-head value logits
    /// from shared trunk activations (save_h_s2).
    ///
    /// Takes raw u64 weight pointers (w_v1, b_v1, w_v2, b_v2) to support
    /// different weight sets per ensemble head.
    #[allow(clippy::too_many_arguments)]
    pub fn forward_value_head(
        &self,
        stream: &Arc<CudaStream>,
        h_s2_ptr: u64,          // [B, SH2] trunk activation (shared)
        w_v1: u64,              // [VH, SH2] value head layer 1 weights
        b_v1: u64,              // [VH] bias
        w_v2: u64,              // [NA, VH] value head layer 2 weights
        b_v2: u64,              // [NA] bias
        h_v_scratch: u64,       // [B, VH] scratch for hidden activation
        v_logits_out: u64,      // [B, NA] output logits
        batch_size: usize,
    ) -> Result<(), MLError> {
        let ws = self.handle.lt_workspace_ptr;
        let wss = self.handle.lt_workspace_size;

        // h_v = ReLU(h_s2 @ W_v1^T + b_v1) — fused epilogue when available
        if self.sgemm_f32_fused_relu_bias(stream, w_v1, h_s2_ptr, h_v_scratch, b_v1,
            self.value_h, batch_size, self.shared_h2, self.shared_h2, ws, wss, "ens_h_v").is_err()
        {
            self.sgemm_f32(
                stream, w_v1, h_s2_ptr, h_v_scratch,
                self.value_h, batch_size, self.shared_h2, "ens_h_v",
            )?;
            self.launch_add_bias_relu_f32_raw(stream, h_v_scratch, b_v1, self.value_h, batch_size)?;
        }

        // v_logits = h_v @ W_v2^T + b_v2 (linear output — no ReLU)
        if self.sgemm_f32_fused_bias(stream, w_v2, h_v_scratch, v_logits_out, b_v2,
            self.num_atoms, batch_size, self.value_h, self.value_h, ws, wss, "ens_v_logits").is_err()
        {
            self.sgemm_f32(
                stream, w_v2, h_v_scratch, v_logits_out,
                self.num_atoms, batch_size, self.value_h, "ens_v_logits",
            )?;
            self.launch_add_bias_f32_raw(stream, v_logits_out, b_v2, self.num_atoms, batch_size)?;
        }

        Ok(())
    }

    /// Run the target network forward pass (inference only — no activation saves).
    ///
    /// Same GEMM sequence as `forward_online` but writes to separate output buffers
    /// and does NOT save h_s2 (it's not needed for the backward pass through target).
    ///
    /// Uses multi-stream parallelism for the 4 branch heads (same pattern as
    /// `forward_online_raw`). Each branch gets its own hidden buffer (`h_b{0..3}_ptr`)
    /// and runs on a dedicated branch stream, giving ~4x overlap on the branch GEMMs.
    ///
    /// Outputs:
    /// - `tg_h_s2_buf`: [B, SH2]       — target trunk output (needed for IQN reuse)
    /// - `tg_v_logits_buf`: [B, NA]    — target value head logits
    /// - `tg_b_logits_buf`: [B, (B0+B1+B2+B3)*NA] — target branch logits
    ///
    /// The caller also needs to run the target network on next_states for DDQN.
    /// We write tg_h_s2_buf here so the caller can inspect it if needed.
    #[allow(clippy::too_many_arguments)]
    /// Graph-safe target forward: takes pre-resolved u64 buffer pointers.
    #[allow(clippy::too_many_arguments)]
    pub fn forward_target_raw(
        &mut self,
        stream: &Arc<CudaStream>,
        states_ptr: u64,
        tg_w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
        h_s1_ptr: u64, h_s2_ptr: u64, h_v_ptr: u64,
        h_b0_ptr: u64, h_b1_ptr: u64, h_b2_ptr: u64, h_b3_ptr: u64,
        v_logits_ptr: u64, b_logits_ptr: u64,
        mag_concat_ptr: u64,
        dir_qaux_concat_ptr: u64,                    // SP14 EGF direction Q-head input concat (0 = unused).
    ) -> Result<(), MLError> {
        let b = self.batch_size;
        let ws = self.handle.lt_workspace_ptr;
        let wss = self.handle.lt_workspace_size;

        // ── Trunk encoder (target GRN — separate save buffers from online) ──
        self.target_encoder_forward_only(stream, states_ptr, tg_w_ptrs, h_s1_ptr, h_s2_ptr)?;

        // h_v[B, VH] — Plan 4 Task 2c.3a: value-head index +9 (w_v1=13, b_v1=14).
        if self.sgemm_f32_fused_relu_bias(stream, tg_w_ptrs[13], h_s2_ptr, h_v_ptr, tg_w_ptrs[14],
            self.value_h, b, self.shared_h2, self.shared_h2, ws, wss, "tg_h_v").is_err()
        {
            self.sgemm_f32(stream, tg_w_ptrs[13], h_s2_ptr, h_v_ptr,
                self.value_h, b, self.shared_h2, "tg_h_v")?;
            self.launch_add_bias_relu_f32_raw(stream, h_v_ptr, tg_w_ptrs[14], self.value_h, b)?;
        }

        // v_logits[B, NA] — output layer: no ReLU (linear output) (w_v2=15, b_v2=16).
        if self.sgemm_f32_fused_bias(stream, tg_w_ptrs[15], h_v_ptr, v_logits_ptr, tg_w_ptrs[16],
            self.num_atoms, b, self.value_h, self.value_h, ws, wss, "tg_v_logits").is_err()
        {
            self.sgemm_f32(stream, tg_w_ptrs[15], h_v_ptr, v_logits_ptr,
                self.num_atoms, b, self.value_h, "tg_v_logits")?;
            self.launch_add_bias_f32_raw(stream, v_logits_ptr, tg_w_ptrs[16], self.num_atoms, b)?;
        }

        // Branch heads — multi-stream parallel dispatch (4 distinct hidden buffers)
        let branch_sizes = [self.branch_0_size, self.branch_1_size, self.branch_2_size, self.branch_3_size];
        let branch_h_ptrs = [h_b0_ptr, h_b1_ptr, h_b2_ptr, h_b3_ptr];
        // Plan 4 Task 2c.3a: branch base indices +9 (b0fc=17, b1fc=21, b2fc=25, b3fc=29).
        let branch_w_base = [17_usize, 21, 25, 29];
        let na = self.num_atoms;

        // Record trunk+value completion — branch streams wait on this.
        self.trunk_done_event.record(stream)
            .map_err(|e| MLError::ModelError(format!("tg trunk event record: {e}")))?;

        let use_vsn_glu = self.vsn_kernel.is_some();

        let mut logit_byte_offset: u64 = 0;
        for d in 0..4 {
            let bs = &self.branch_streams[d];
            let n_d = branch_sizes[d];
            let w_fc_idx = branch_w_base[d];

            // Branch stream waits for trunk completion.
            bs.wait(&self.trunk_done_event)
                .map_err(|e| MLError::ModelError(format!("tg branch {d} wait trunk: {e}")))?;

            if use_vsn_glu {
                // VSN bottleneck + GLU gating path
                self.launch_vsn_glu_branch(
                    bs, d, h_s2_ptr, mag_concat_ptr, dir_qaux_concat_ptr, tg_w_ptrs,
                    branch_h_ptrs[d],
                    self.branch_workspace_ptrs[d], self.handle.lt_workspace_size,
                    true, "tg",
                )?;
            } else {
                // GEMM+bias+ReLU — try fused epilogue, fall back to separate kernels
                let bws = self.branch_workspace_ptrs[d];
                let bwss = self.handle.lt_workspace_size;
                let (fc_input, fc_k) = if d == 0 && dir_qaux_concat_ptr != 0 {
                    // SP14 EGF: target direction Q-head reads [h_s2 ; aux_softmax_diff].
                    (dir_qaux_concat_ptr, self.shared_h2 + 1)
                } else if d == 1 && mag_concat_ptr != 0 {
                    (mag_concat_ptr, self.mag_concat_dim)
                } else {
                    (h_s2_ptr, self.shared_h2)
                };
                if self.sgemm_f32_fused_relu_bias(bs, tg_w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], tg_w_ptrs[w_fc_idx + 1],
                    self.adv_h, b, fc_k, fc_k, bws, bwss, "tg_h_bd").is_err()
                {
                    self.sgemm_f32_branch(bs, tg_w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d],
                        self.adv_h, b, fc_k, d, "tg_h_bd")?;
                    self.launch_add_bias_relu_f32_raw(bs, branch_h_ptrs[d], tg_w_ptrs[w_fc_idx + 1], self.adv_h, b)?;
                }
            }

            // Output layer: f32 cublasLtMatmul + f32 bias — per-branch workspace.
            let adv_out_ptr = b_logits_ptr + logit_byte_offset;
            let out_bws = self.branch_workspace_ptrs[d];
            let out_bwss = self.handle.lt_workspace_size;
            if self.sgemm_f32_fused_bias(bs, tg_w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, tg_w_ptrs[w_fc_idx + 3],
                n_d * na, b, self.adv_h, self.adv_h, out_bws, out_bwss, "tg_adv_logits").is_err()
            {
                self.sgemm_f32_branch(bs, tg_w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr,
                    n_d * na, b, self.adv_h, d, "tg_adv_logits")?;
                self.launch_add_bias_f32_raw(bs, adv_out_ptr, tg_w_ptrs[w_fc_idx + 3], n_d * na, b)?;
            }

            logit_byte_offset += (b * n_d * na * std::mem::size_of::<f32>()) as u64;
        }

        // Join: main stream waits for all 4 target branches to complete.
        for d in 0..4 {
            self.branch_done_events[d].record(&self.branch_streams[d])
                .map_err(|e| MLError::ModelError(format!("tg branch {d} done event: {e}")))?;
            stream.wait(&self.branch_done_events[d])
                .map_err(|e| MLError::ModelError(format!("tg main wait branch {d}: {e}")))?;
        }

        Ok(())
    }

    /// Target forward using CudaSlice references (calls device_ptr — NOT graph-safe).
    #[allow(clippy::too_many_arguments)]
    pub fn forward_target(
        &mut self,
        stream: &Arc<CudaStream>,
        next_states_buf: &CudaSlice<f32>,
        tg_w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
        tg_h_s1_scratch: &CudaSlice<f32>,
        tg_h_s2_buf: &CudaSlice<f32>,
        tg_h_v_scratch: &CudaSlice<f32>,
        tg_h_b0_scratch: &CudaSlice<f32>,
        tg_h_b1_scratch: &CudaSlice<f32>,
        tg_h_b2_scratch: &CudaSlice<f32>,
        tg_h_b3_scratch: &CudaSlice<f32>,
        tg_v_logits_buf: &CudaSlice<f32>,
        tg_b_logits_buf: &CudaSlice<f32>,
    ) -> Result<(), MLError> {
        self.forward_target_raw(
            stream,
            raw_f32_ptr(next_states_buf, stream),
            tg_w_ptrs,
            raw_f32_ptr(tg_h_s1_scratch, stream),
            raw_f32_ptr(tg_h_s2_buf, stream),
            raw_f32_ptr(tg_h_v_scratch, stream),
            raw_f32_ptr(tg_h_b0_scratch, stream),
            raw_f32_ptr(tg_h_b1_scratch, stream),
            raw_f32_ptr(tg_h_b2_scratch, stream),
            raw_f32_ptr(tg_h_b3_scratch, stream),
            raw_f32_ptr(tg_v_logits_buf, stream),
            raw_f32_ptr(tg_b_logits_buf, stream),
            0u64, // mag_concat_ptr: not used in CudaSlice wrapper path
            0u64, // dir_qaux_concat_ptr: SP14 EGF disabled in CudaSlice wrapper path
        )
    }

    // ══════════════════════════════════════════════════════════════════════════
    // VSN bottleneck + GLU gating — branch forward path
    // ══════════════════════════════════════════════════════════════════════════

    /// Execute one branch forward with VSN bottleneck + GLU gating:
    ///   1. VSN: variable_select_bottleneck(h_s2) → vsn_masked [B, SH2]
    ///   2. For d==0 with dir_qaux_concat (SP14 EGF): vsn_input =
    ///      [vsn_masked; aux_softmax_diff] (last column was written by
    ///      `dir_concat_qaux_kernel` BEFORE this call; we scatter
    ///      vsn_masked into the first SH2 cols here, overwriting the
    ///      stale h_s2 prefix the kernel wrote — same pattern as d==1/2/3).
    ///      For d==1 with mag_concat: vsn_input = [vsn_masked; Q_dir] (already in mag_concat_ptr)
    ///      For others: vsn_input = vsn_masked
    ///   3. Value GEMM: W_bdf @ vsn_input → glu_value [B, AH], add bias (no ReLU)
    ///   4. Gate  GEMM: W_gate @ vsn_input → glu_gate_pre [B, AH], add gate bias
    ///   5. KAN gate: h_bd = kan_spline(gate_pre, coeffs, resid_w) * value
    ///
    /// `ws_ptr`/`ws_size`: cublasLt workspace for this stream (branch or main).
    /// `is_branch_stream`: if true, uses `sgemm_f32_branch`; if false, `sgemm_f32`.
    #[allow(clippy::too_many_arguments)]
    fn launch_vsn_glu_branch(
        &self,
        stream: &Arc<CudaStream>,
        d: usize,                    // branch index 0..3
        h_s2_ptr: u64,              // [B, SH2] trunk output
        mag_concat_ptr: u64,        // [B, SH2+branch_0_size] magnitude concat (0 if unused)
        dir_qaux_concat_ptr: u64,   // [B, SH2+1] direction-aux concat (SP14 EGF; 0 if unused)
        w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
        branch_h_ptr: u64,          // [B, AH] output (save_h_bd)
        _ws_ptr: u64,
        _ws_size: usize,
        is_branch_stream: bool,
        _label_prefix: &str,
    ) -> Result<(), MLError> {
        let vsn_k = self.vsn_kernel.as_ref().unwrap();
        let kan_k = self.kan_gate_combine_kernel.as_ref().unwrap();
        let b = self.batch_size;
        let sh2 = self.shared_h2;
        let ah = self.adv_h;
        let r = self.vsn_rank;

        // ── 1. VSN bottleneck: h_s2 [B, SH2] → vsn_masked [B, SH2] ──
        let w_vsn1_idx = 26 + d * 2;       // W_vsn1_d
        let w_vsn2_idx = 26 + d * 2 + 1;   // W_vsn2_d
        let smem = (sh2 + r) * std::mem::size_of::<f32>();  // h_local[SH2] + proj[R], no mask buf (sigmoid is inline)
        unsafe {
            stream
                .launch_builder(vsn_k)
                .arg(&h_s2_ptr)
                .arg(&self.vsn_masked_ptr)
                .arg(&w_ptrs[w_vsn1_idx])
                .arg(&w_ptrs[w_vsn2_idx])
                .arg(&(b as i32))
                .arg(&(sh2 as i32))
                .arg(&(r as i32))
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: smem as u32,
                })
                .map_err(|e| MLError::ModelError(format!(
                    "{_label_prefix} vsn_bottleneck branch {d}: {e}"
                )))?;
        }

        // ── 2. Rebuild concat from VSN-masked output ──
        //
        // d==1 (magnitude): Scatter vsn_masked → first SH2 cols of mag_concat
        //   [B, SH2+branch_0_size]. Q_dir (last branch_0_size cols) was written
        //   by the pre-VSN launch_mag_concat_from call.
        //
        // d==2 (order) / d==3 (urgency): Scatter vsn_masked → first SH2 cols of
        //   ord_concat / urg_concat [B, SH2+3]. OFI (last 3 cols) was written
        //   by the pre-VSN launch_concat_ofi call.
        //
        // Note: dst_stride differs between mag and ord/urg branches when
        // branch_0_size != 3 (production: 4 vs 3).
        // This is one kernel launch (B threads, negligible cost).
        let scatter_k = self.strided_scatter_kernel.as_ref().unwrap();
        let total_scatter = (b * sh2) as i32;
        let src_stride_scatter = sh2 as i32;
        let mag_dst_stride = self.mag_concat_dim as i32;
        let ofi_dst_stride = self.ofi_concat_dim as i32;
        let scatter_blocks = ((total_scatter as u32 + 255) / 256).max(1);

        let dir_qaux_dst_stride = (sh2 + 1) as i32;

        let (vsn_input, fc_k) = if d == 0 && dir_qaux_concat_ptr != 0 {
            // SP14 EGF: scatter vsn_masked into first SH2 cols of
            // sp14_dir_qaux_concat_scratch [B, SH2+1]. The trailing aux_softmax_diff
            // column at offset SH2 was already written by the pre-forward
            // `dir_concat_qaux_kernel` launch — same one-step-lag pattern as
            // mag_concat. dst_stride = SH2 + 1.
            unsafe {
                stream
                    .launch_builder(scatter_k)
                    .arg(&self.vsn_masked_ptr)
                    .arg(&dir_qaux_concat_ptr)
                    .arg(&src_stride_scatter)
                    .arg(&dir_qaux_dst_stride)
                    .arg(&total_scatter)
                    .launch(LaunchConfig {
                        grid_dim: (scatter_blocks, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| MLError::ModelError(format!(
                        "strided_scatter vsn→dir_qaux_concat: {e}"
                    )))?;
            }
            (dir_qaux_concat_ptr, sh2 + 1)
        } else if d == 1 && mag_concat_ptr != 0 {
            unsafe {
                stream
                    .launch_builder(scatter_k)
                    .arg(&self.vsn_masked_ptr)
                    .arg(&mag_concat_ptr)
                    .arg(&src_stride_scatter)
                    .arg(&mag_dst_stride)
                    .arg(&total_scatter)
                    .launch(LaunchConfig {
                        grid_dim: (scatter_blocks, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| MLError::ModelError(format!(
                        "strided_scatter vsn→mag_concat: {e}"
                    )))?;
            }
            (mag_concat_ptr, self.mag_concat_dim)
        } else if d == 2 && self.ord_concat_ptr != 0 {
            unsafe {
                stream
                    .launch_builder(scatter_k)
                    .arg(&self.vsn_masked_ptr)
                    .arg(&self.ord_concat_ptr)
                    .arg(&src_stride_scatter)
                    .arg(&ofi_dst_stride)
                    .arg(&total_scatter)
                    .launch(LaunchConfig {
                        grid_dim: (scatter_blocks, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| MLError::ModelError(format!(
                        "strided_scatter vsn→ord_concat: {e}"
                    )))?;
            }
            (self.ord_concat_ptr, self.ofi_concat_dim)
        } else if d == 3 && self.urg_concat_ptr != 0 {
            unsafe {
                stream
                    .launch_builder(scatter_k)
                    .arg(&self.vsn_masked_ptr)
                    .arg(&self.urg_concat_ptr)
                    .arg(&src_stride_scatter)
                    .arg(&ofi_dst_stride)
                    .arg(&total_scatter)
                    .launch(LaunchConfig {
                        grid_dim: (scatter_blocks, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| MLError::ModelError(format!(
                        "strided_scatter vsn→urg_concat: {e}"
                    )))?;
            }
            (self.urg_concat_ptr, self.ofi_concat_dim)
        } else {
            (self.vsn_masked_ptr, sh2)
        };

        // ── 3. Value GEMM: W_bdf @ vsn_input → glu_value [B, AH] + bias (no ReLU) ──
        // Plan 4 Task 2c.3a: branch base indices +9 (b0fc=17, b1fc=21, b2fc=25, b3fc=29).
        let w_fc_base = [17_usize, 21, 25, 29];
        let w_fc_idx = w_fc_base[d];
        let value_ptr = self.glu_value_ptrs[d];

        // Phase H Site 3: fused GEMM + bias (no ReLU — GLU provides the gating).
        // Workspace selected per-stream: branch streams use the per-branch
        // workspace to avoid contention with concurrent branch dispatches.
        let glu_ws = if is_branch_stream { self.branch_workspace_ptrs[d] } else { self.handle.lt_workspace_ptr };
        let glu_wss = self.handle.lt_workspace_size;
        if self.sgemm_f32_fused_bias(stream, w_ptrs[w_fc_idx], vsn_input, value_ptr, w_ptrs[w_fc_idx + 1],
            ah, b, fc_k, fc_k, glu_ws, glu_wss, _label_prefix).is_err()
        {
            if is_branch_stream {
                self.sgemm_f32_branch(stream, w_ptrs[w_fc_idx], vsn_input, value_ptr, ah, b, fc_k, d, _label_prefix)?;
            } else {
                self.sgemm_f32(stream, w_ptrs[w_fc_idx], vsn_input, value_ptr, ah, b, fc_k, _label_prefix)?;
            }
            self.launch_add_bias_f32_raw(stream, value_ptr, w_ptrs[w_fc_idx + 1], ah, b)?;
        }

        // ── 4. Gate GEMM: W_gate @ vsn_input → glu_gate_pre [B, AH] + gate bias ──
        // Plan 4 Task 2c.3a: gate base index 34 +9 → 43.
        let w_gate_idx = 43 + d * 2;     // W_gate_d
        let b_gate_idx = 43 + d * 2 + 1; // b_gate_d
        let gate_ptr = self.glu_gate_pre_ptrs[d];

        if self.sgemm_f32_fused_bias(stream, w_ptrs[w_gate_idx], vsn_input, gate_ptr, w_ptrs[b_gate_idx],
            ah, b, fc_k, fc_k, glu_ws, glu_wss, _label_prefix).is_err()
        {
            if is_branch_stream {
                self.sgemm_f32_branch(stream, w_ptrs[w_gate_idx], vsn_input, gate_ptr, ah, b, fc_k, d, _label_prefix)?;
            } else {
                self.sgemm_f32(stream, w_ptrs[w_gate_idx], vsn_input, gate_ptr, ah, b, fc_k, _label_prefix)?;
            }
            self.launch_add_bias_f32_raw(stream, gate_ptr, w_ptrs[b_gate_idx], ah, b)?;
        }

        // ── 5. KAN gate: h_bd = kan_spline(gate_pre, coeffs, resid_w) * value ──
        // Plan 4 Task 2c.3a: KAN coeff/resid base index 42 +9 → 51.
        let kan_coeff_idx = 51 + d * 2;       // kan_coeff_d
        let kan_resid_idx = 51 + d * 2 + 1;   // kan_resid_d
        let n_total = (b * ah) as i32;
        let num_neurons = ah as i32;
        let blocks = ((b * ah + 255) / 256) as u32;
        unsafe {
            stream
                .launch_builder(kan_k)
                .arg(&branch_h_ptr)
                .arg(&gate_ptr)
                .arg(&value_ptr)
                .arg(&w_ptrs[kan_coeff_idx])
                .arg(&w_ptrs[kan_resid_idx])
                .arg(&n_total)
                .arg(&num_neurons)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!(
                    "{_label_prefix} kan_gate_combine branch {d}: {e}"
                )))?;
        }

        Ok(())
    }

    // ══════════════════════════════════════════════════════════════════════════
    // cublasLtMatmul — the ONLY forward GEMM (f32 + TF32 via CUBLAS_COMPUTE_32F_FAST_TF32)
    // ══════════════════════════════════════════════════════════════════════════

    /// F32 GEMM with TF32 tensor core acceleration via cublasLtMatmul.
    /// `C[B, N] = A[B, K] @ W[N, K]^T` — all f32 types.
    ///
    /// cublasLtMatmul passes workspace+stream per-call (no handle state),
    /// eliminating the `cublasSetStream` workspace conflict that hangs CUDA
    /// Graph mega-capture on H100. CUBLAS_COMPUTE_32F_FAST_TF32 gives ~500
    /// TFLOPS on H100 (vs ~60 for pure f32 CUDA cores).
    ///
    /// Uses cached descriptors (pre-created at init) when available — zero
    /// per-call descriptor creation overhead.
    pub(crate) fn sgemm_f32(
        &self,
        stream: &CudaStream,
        w_f32_ptr: u64,   // F32 weights [out_dim, in_dim]
        a_f32_ptr: u64,   // F32 input [B, in_dim]
        c_f32_ptr: u64,   // F32 output [B, out_dim]
        n: usize,         // out_dim
        b: usize,         // batch
        k: usize,         // in_dim
        _label: &str,
    ) -> Result<(), MLError> {
        self.lt_matmul(stream, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, k, self.handle.lt_workspace_ptr, self.handle.lt_workspace_size, _label)
    }

    /// F32 GEMM on a branch stream with per-branch workspace.
    /// Eliminates workspace contention when 4 branches run in parallel.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn sgemm_f32_branch(
        &self,
        stream: &CudaStream,
        w_f32_ptr: u64,
        a_f32_ptr: u64,
        c_f32_ptr: u64,
        n: usize,
        b: usize,
        k: usize,
        branch_idx: usize,
        _label: &str,
    ) -> Result<(), MLError> {
        self.lt_matmul(stream, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, k, self.branch_workspace_ptrs[branch_idx], self.handle.lt_workspace_size, _label)
    }

    /// F32 GEMM with TF32 + explicit `ldb` (for padded states buffer).
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn sgemm_f32_ldb(
        &self,
        stream: &CudaStream,
        w_f32_ptr: u64,
        a_f32_ptr: u64,
        c_f32_ptr: u64,
        n: usize,
        b: usize,
        k: usize,
        ldb: usize,
        _label: &str,
    ) -> Result<(), MLError> {
        self.lt_matmul(stream, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, ldb, self.handle.lt_workspace_ptr, self.handle.lt_workspace_size, _label)
    }

    /// Inner cublasLtMatmul dispatch: uses pre-cached descriptors when available,
    /// falls back to inline descriptor creation for uncached shapes.
    ///
    /// Layout (row-major trick):
    ///   cuBLAS sees A(W) as col-major [K, N] → TRANSA=T
    ///   cuBLAS sees B(input) as col-major [K, B] → TRANSB=N
    ///   Result C is col-major [N, B] == row-major [B, N]
    #[allow(clippy::too_many_arguments)]
    fn lt_matmul(
        &self,
        stream: &CudaStream,
        w_f32_ptr: u64,
        a_f32_ptr: u64,
        c_f32_ptr: u64,
        n: usize,
        b: usize,
        k: usize,
        ldb: usize,
        ws_ptr: u64,
        ws_size: usize,
        _label: &str,
    ) -> Result<(), MLError> {
        let key: FwdGemmKey = (n, b, k, ldb);
        if let Some(cached) = self.gemm_cache.get(&key) {
            // ── Fast path: use pre-created descriptors + algo ──
            self.lt_matmul_cached(stream, cached, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, ldb, ws_ptr, ws_size, _label)
        } else {
            // ── Slow path: inline descriptor creation (backward compat for uncached shapes) ──
            self.lt_matmul_uncached(stream, w_f32_ptr, a_f32_ptr, c_f32_ptr, n, b, k, ldb, ws_ptr, ws_size, _label)
        }
    }

    /// Fast path: execute cublasLtMatmul with pre-cached descriptors.
    /// Zero per-call descriptor creation overhead.
    #[allow(clippy::too_many_arguments)]
    fn lt_matmul_cached(
        &self,
        stream: &CudaStream,
        cached: &CachedGemmDesc,
        w_f32_ptr: u64,
        a_f32_ptr: u64,
        c_f32_ptr: u64,
        n: usize,
        b: usize,
        k: usize,
        ldb: usize,
        ws_ptr: u64,
        ws_size: usize,
        _label: &str,
    ) -> Result<(), MLError> {
        let alpha = 1.0_f32;
        let beta = 0.0_f32;

        // Per-stream lt_handle lookup (Option C determinism fix).
        // Each CUDA stream gets its own cublasLt handle — sharing a handle
        // across concurrent streams lets cuBLAS pick different internal
        // implementations (per NVIDIA cuBLAS §2.1.4).
        let lt_handle = self.handle.lt_handle_for(stream)?;

        unsafe {
            let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t;

            tracing::debug!(
                handle = ?lt_handle,
                matmul_desc = ?cached.matmul_desc,
                w_ptr = w_f32_ptr,
                a_ptr = a_f32_ptr,
                c_ptr = c_f32_ptr,
                a_layout = ?cached.a_layout,
                b_layout = ?cached.b_layout,
                c_layout = ?cached.c_layout,
                d_layout = ?cached.d_layout,
                ws_ptr = ws_ptr,
                ws_size = ws_size,
                stream = ?cu_stream,
                m = n, n_batch = b, k = k, ldb = ldb,
                "lt_matmul CACHED FFI args for {_label}"
            );

            let matmul_status = cublaslt_sys::cublasLtMatmul(
                lt_handle,
                cached.matmul_desc,
                &alpha as *const f32 as *const std::ffi::c_void,
                w_f32_ptr as *const std::ffi::c_void,
                cached.a_layout,
                a_f32_ptr as *const std::ffi::c_void,
                cached.b_layout,
                &beta as *const f32 as *const std::ffi::c_void,
                c_f32_ptr as *const std::ffi::c_void,
                cached.c_layout,
                c_f32_ptr as *mut std::ffi::c_void,
                cached.d_layout,
                &cached.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t,
                ws_ptr as *mut std::ffi::c_void,
                ws_size,
                cu_stream,
            );

            if matmul_status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
                let e = cublaslt_result::CublasError(matmul_status);
                tracing::error!(
                    m=n, n_batch=b, k=k, ldb=ldb,
                    ws_provided=ws_size,
                    ws_ptr=ws_ptr,
                    ?e,
                    "cublasLtMatmul CACHED FAILED for {_label}"
                );
                return Err(MLError::ModelError(format!("cublasLtMatmul cached {_label} (m={n},n={b},k={k},ldb={ldb}): {e:?}")));
            }
        }
        Ok(())
    }

    /// Fused GEMM + bias + ReLU via cublasLt epilogue.
    /// Uses cached descriptor with RELU_BIAS epilogue. Sets bias pointer per-call.
    /// Falls back to separate GEMM + add_bias_relu if epilogue not available.
    pub(crate) fn sgemm_f32_fused_relu_bias(
        &self,
        stream: &CudaStream,
        w_ptr: u64,
        input_ptr: u64,
        output_ptr: u64,
        bias_ptr: u64,
        n: usize,    // out_dim
        b: usize,    // batch
        k: usize,    // in_dim
        ldb: usize,
        ws_ptr: u64,
        ws_size: usize,
        label: &str,
    ) -> Result<(), MLError> {
        let key: FwdGemmKey = (n, b, k, ldb);
        if let Some(cached) = self.gemm_cache_relu_bias.get(&key) {
            // Set bias pointer on the descriptor (lightweight CPU write)
            unsafe {
                cublaslt_result::set_matmul_desc_attribute(
                    cached.matmul_desc,
                    cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_BIAS_POINTER,
                    &bias_ptr as *const u64 as *const std::ffi::c_void,
                    std::mem::size_of::<u64>(),
                ).map_err(|e| MLError::ModelError(format!("set bias ptr {label}: {e:?}")))?;
            }

            // Launch fused GEMM+bias+ReLU
            let alpha = 1.0_f32;
            let beta = 0.0_f32;
            // Per-stream lt_handle lookup (Option C determinism fix).
            let lt_handle = self.handle.lt_handle_for(stream)?;
            unsafe {
                let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t;
                let status = cublaslt_sys::cublasLtMatmul(
                    lt_handle,
                    cached.matmul_desc,
                    &alpha as *const f32 as *const std::ffi::c_void,
                    w_ptr as *const std::ffi::c_void,
                    cached.a_layout,
                    input_ptr as *const std::ffi::c_void,
                    cached.b_layout,
                    &beta as *const f32 as *const std::ffi::c_void,
                    output_ptr as *const std::ffi::c_void,
                    cached.c_layout,
                    output_ptr as *mut std::ffi::c_void,
                    cached.d_layout,
                    &cached.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t,
                    ws_ptr as *mut std::ffi::c_void,
                    ws_size,
                    cu_stream,
                );
                if status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
                    let e = cublaslt_result::CublasError(status);
                    tracing::error!(m=n, n_batch=b, k=k, ?e, "cublasLtMatmul FUSED RELU_BIAS FAILED for {label}");
                    return Err(MLError::ModelError(format!("cublasLtMatmul fused {label}: {e:?}")));
                }
            }
            Ok(())
        } else {
            // No cached epilogue descriptor — caller should use separate GEMM + bias_relu
            Err(MLError::ModelError(format!("no RELU_BIAS epilogue for {label} (n={n},b={b},k={k})")))
        }
    }

    /// Fused GEMM + bias via cublasLt epilogue (no ReLU).
    /// Uses cached descriptor with BIAS epilogue. Sets bias pointer per-call.
    /// Falls back to an `Err` if the shape is not pre-cached at init, signalling
    /// the caller to use the slower separate-kernel path (`sgemm_f32` +
    /// `launch_add_bias_f32_raw`).
    ///
    /// Phase H Site 3: collapses GEMM + add_bias into one cublasLtMatmul call,
    /// eliminating the standalone `add_bias_f32_kernel` launch and one memory
    /// round-trip per BIAS-only output layer (value logits, advantage logits,
    /// GRN Linear_a/b, VSN linear2, GLU value/gate).
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn sgemm_f32_fused_bias(
        &self,
        stream: &CudaStream,
        w_ptr: u64,
        input_ptr: u64,
        output_ptr: u64,
        bias_ptr: u64,
        n: usize,    // out_dim
        b: usize,    // batch
        k: usize,    // in_dim
        ldb: usize,
        ws_ptr: u64,
        ws_size: usize,
        label: &str,
    ) -> Result<(), MLError> {
        let key: FwdGemmKey = (n, b, k, ldb);
        if let Some(cached) = self.gemm_cache_bias.get(&key) {
            // Set bias pointer on the descriptor (lightweight CPU write).
            unsafe {
                cublaslt_result::set_matmul_desc_attribute(
                    cached.matmul_desc,
                    cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_BIAS_POINTER,
                    &bias_ptr as *const u64 as *const std::ffi::c_void,
                    std::mem::size_of::<u64>(),
                ).map_err(|e| MLError::ModelError(format!("set bias ptr {label}: {e:?}")))?;
            }

            // Launch fused GEMM+bias.
            let alpha = 1.0_f32;
            let beta = 0.0_f32;
            let lt_handle = self.handle.lt_handle_for(stream)?;
            unsafe {
                let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t;
                let status = cublaslt_sys::cublasLtMatmul(
                    lt_handle,
                    cached.matmul_desc,
                    &alpha as *const f32 as *const std::ffi::c_void,
                    w_ptr as *const std::ffi::c_void,
                    cached.a_layout,
                    input_ptr as *const std::ffi::c_void,
                    cached.b_layout,
                    &beta as *const f32 as *const std::ffi::c_void,
                    output_ptr as *const std::ffi::c_void,
                    cached.c_layout,
                    output_ptr as *mut std::ffi::c_void,
                    cached.d_layout,
                    &cached.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t,
                    ws_ptr as *mut std::ffi::c_void,
                    ws_size,
                    cu_stream,
                );
                if status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
                    let e = cublaslt_result::CublasError(status);
                    tracing::error!(m=n, n_batch=b, k=k, ?e, "cublasLtMatmul FUSED BIAS FAILED for {label}");
                    return Err(MLError::ModelError(format!("cublasLtMatmul fused-bias {label}: {e:?}")));
                }
            }
            Ok(())
        } else {
            // No cached epilogue descriptor — caller should use separate GEMM + bias.
            Err(MLError::ModelError(format!("no BIAS epilogue for {label} (n={n},b={b},k={k},ldb={ldb})")))
        }
    }

    /// Fused GEMM + bias + ReLU + dReLU bit-mask aux output via cublasLt epilogue.
    ///
    /// Same as `sgemm_f32_fused_relu_bias` but uses `EPILOGUE_RELU_AUX_BIAS` —
    /// in addition to writing the post-ReLU output, cublasLt writes a 1-bit
    /// per-element mask of the dReLU pattern to the AUX buffer. The mask is
    /// consumed by `CublasBackwardSet::backward_fc_layer_drelu_bgrad` to fuse
    /// the standalone ReLU-derivative + bias-grad reduce into the dX GEMM
    /// (Phase H Site 2).
    ///
    /// Bias and AUX pointers are set on the cached descriptor per-call
    /// (lightweight CPU writes — same pattern as `sgemm_f32_fused_relu_bias`).
    /// Falls back to an `Err` if the shape is not pre-cached, signalling the
    /// caller to use the slower separate-kernel path.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn sgemm_f32_fused_relu_aux_bias(
        &self,
        stream: &CudaStream,
        w_ptr: u64,
        input_ptr: u64,
        output_ptr: u64,
        bias_ptr: u64,
        aux_ptr: u64,
        n: usize,    // out_dim
        b: usize,    // batch
        k: usize,    // in_dim
        ldb: usize,
        ws_ptr: u64,
        ws_size: usize,
        label: &str,
    ) -> Result<(), MLError> {
        let key: FwdGemmKey = (n, b, k, ldb);
        if let Some(cached) = self.gemm_cache_relu_aux_bias.get(&key) {
            // Set bias + AUX pointers on the descriptor (lightweight CPU writes).
            unsafe {
                cublaslt_result::set_matmul_desc_attribute(
                    cached.matmul_desc,
                    cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_BIAS_POINTER,
                    &bias_ptr as *const u64 as *const std::ffi::c_void,
                    std::mem::size_of::<u64>(),
                ).map_err(|e| MLError::ModelError(format!("set bias ptr aux {label}: {e:?}")))?;
                cublaslt_result::set_matmul_desc_attribute(
                    cached.matmul_desc,
                    cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_POINTER,
                    &aux_ptr as *const u64 as *const std::ffi::c_void,
                    std::mem::size_of::<u64>(),
                ).map_err(|e| MLError::ModelError(format!("set aux ptr {label}: {e:?}")))?;
            }

            let alpha = 1.0_f32;
            let beta = 0.0_f32;
            let lt_handle = self.handle.lt_handle_for(stream)?;
            unsafe {
                let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t;
                let status = cublaslt_sys::cublasLtMatmul(
                    lt_handle,
                    cached.matmul_desc,
                    &alpha as *const f32 as *const std::ffi::c_void,
                    w_ptr as *const std::ffi::c_void,
                    cached.a_layout,
                    input_ptr as *const std::ffi::c_void,
                    cached.b_layout,
                    &beta as *const f32 as *const std::ffi::c_void,
                    output_ptr as *const std::ffi::c_void,
                    cached.c_layout,
                    output_ptr as *mut std::ffi::c_void,
                    cached.d_layout,
                    &cached.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t,
                    ws_ptr as *mut std::ffi::c_void,
                    ws_size,
                    cu_stream,
                );
                if status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
                    let e = cublaslt_result::CublasError(status);
                    tracing::error!(m=n, n_batch=b, k=k, ?e, "cublasLtMatmul FUSED RELU_AUX_BIAS FAILED for {label}");
                    return Err(MLError::ModelError(format!("cublasLtMatmul fused aux {label}: {e:?}")));
                }
            }
            Ok(())
        } else {
            Err(MLError::ModelError(format!("no RELU_AUX_BIAS epilogue for {label} (n={n},b={b},k={k})")))
        }
    }

    /// Slow path: inline descriptor creation for uncached GEMM shapes.
    /// Used for diagnostic calls and any shape not pre-cached at init.
    #[allow(clippy::too_many_arguments)]
    fn lt_matmul_uncached(
        &self,
        stream: &CudaStream,
        w_f32_ptr: u64,
        a_f32_ptr: u64,
        c_f32_ptr: u64,
        n: usize,
        b: usize,
        k: usize,
        ldb: usize,
        ws_ptr: u64,
        ws_size: usize,
        _label: &str,
    ) -> Result<(), MLError> {
        let alpha = 1.0_f32;
        let beta = 0.0_f32;
        let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
        let compute_type = cublaslt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32;

        unsafe {
            // ── Matmul descriptor ──
            let matmul_desc = cublaslt_result::create_matmul_desc(compute_type, f32_type)
                .map_err(|e| MLError::ModelError(format!("cublasLtMatmulDescCreate {_label}: {e:?}")))?;

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

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

            // ── Matrix layouts ──
            let a_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, k as i64)
                .map_err(|e| MLError::ModelError(format!("A layout {_label}: {e:?}")))?;
            let b_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, b as u64, ldb as i64)
                .map_err(|e| MLError::ModelError(format!("B layout {_label}: {e:?}")))?;
            let c_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64)
                .map_err(|e| MLError::ModelError(format!("C layout {_label}: {e:?}")))?;
            let d_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64)
                .map_err(|e| MLError::ModelError(format!("D layout {_label}: {e:?}")))?;

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

            // ── Deterministic algorithm selection (Option B fix) ──
            // Replaces timing-based `cublasLtMatmulAlgoGetHeuristic` with a
            // hardware-stable AlgoGetIds → AlgoInit → AlgoCheck loop.
            // TF32 compute type preserved.
            let shape = super::cublas_algo_deterministic::ShapeKey::new(
                transa, transb,
                n as i32, b as i32, k as i32,
                k as i32, ldb as i32, n as i32,
                ws_size,
            );
            let heuristic = super::cublas_algo_deterministic::get_matmul_algo_f32_tf32(
                lt_handle,
                matmul_desc,
                a_layout,
                b_layout,
                c_layout,
                d_layout,
                shape,
            );
            if let Err(ref e) = heuristic {
                tracing::error!(m=n, n_batch=b, k=k, ldb=ldb, ws=ws_size, ?e, "cublasLt deterministic algo selection FAILED for {_label}");
            }
            let heuristic = heuristic.map_err(|e| MLError::ModelError(format!("algo deterministic {_label} (m={n},n={b},k={k},ldb={ldb}): {e}")))?;

            // ── Execute matmul ──
            let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t;

            tracing::debug!(
                handle = ?lt_handle,
                matmul_desc = ?matmul_desc,
                w_ptr = w_f32_ptr,
                a_ptr = a_f32_ptr,
                c_ptr = c_f32_ptr,
                a_layout = ?a_layout,
                b_layout = ?b_layout,
                c_layout = ?c_layout,
                d_layout = ?d_layout,
                ws_ptr = ws_ptr,
                ws_size = ws_size,
                stream = ?cu_stream,
                m = n, n_batch = b, k = k, ldb = ldb,
                "lt_matmul UNCACHED FFI args for {_label}"
            );

            let matmul_status = cublaslt_sys::cublasLtMatmul(
                lt_handle,
                matmul_desc,
                &alpha as *const f32 as *const std::ffi::c_void,
                w_f32_ptr as *const std::ffi::c_void,
                a_layout,
                a_f32_ptr as *const std::ffi::c_void,
                b_layout,
                &beta as *const f32 as *const std::ffi::c_void,
                c_f32_ptr as *const std::ffi::c_void,
                c_layout,
                c_f32_ptr as *mut std::ffi::c_void,
                d_layout,
                &heuristic.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t,
                ws_ptr as *mut std::ffi::c_void,
                ws_size,
                cu_stream,
            );
            let matmul_result = if matmul_status == cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
                Ok(())
            } else {
                Err(cublaslt_result::CublasError(matmul_status))
            };
            if let Err(ref e) = matmul_result {
                tracing::error!(
                    m=n, n_batch=b, k=k, ldb=ldb,
                    ws_provided=ws_size,
                    ws_needed=heuristic.workspaceSize,
                    ws_ptr=ws_ptr,
                    ?e,
                    "cublasLtMatmul UNCACHED FAILED for {_label}"
                );
            }
            matmul_result.map_err(|e| MLError::ModelError(format!("cublasLtMatmul {_label} (m={n},n={b},k={k},ldb={ldb},ws={ws_size},needed={}): {e:?}", heuristic.workspaceSize)))?;

            // ── Cleanup ALL descriptors ──
            let _ = cublaslt_result::destroy_matrix_layout(d_layout);
            let _ = cublaslt_result::destroy_matrix_layout(c_layout);
            let _ = cublaslt_result::destroy_matrix_layout(b_layout);
            let _ = cublaslt_result::destroy_matrix_layout(a_layout);
            let _ = cublaslt_result::destroy_matmul_desc(matmul_desc);
        }
        Ok(())
    }

    // ══════════════════════════════════════════════════════════════════════════
    // Bias + activation kernel launchers (all f32)
    // ══════════════════════════════════════════════════════════════════════════

    // ── Bias + activation kernel launchers ──
    pub(crate) fn launch_add_bias_relu_f32_raw(
        &self,
        stream: &Arc<CudaStream>,
        out_ptr: u64,      // F32
        bias_ptr: u64,     // F32
        out_dim: usize,
        batch: usize,
    ) -> Result<(), MLError> {
        let total = (batch * out_dim) as i32;
        let out_dim_i32 = out_dim as i32;
        let blocks = ((batch * out_dim + 255) / 256) as u32;
        unsafe {
            stream
                .launch_builder(&self.handle.add_bias_relu_f32_kernel)
                .arg(&out_ptr)
                .arg(&bias_ptr)
                .arg(&out_dim_i32)
                .arg(&total)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("add_bias_relu_f32: {e}")))?;
        }
        Ok(())
    }

    /// F32 bias (no activation) for logit output layers.
    pub(crate) fn launch_add_bias_f32_raw(
        &self,
        stream: &Arc<CudaStream>,
        out_ptr: u64,
        bias_ptr: u64,
        out_dim: usize,
        batch: usize,
    ) -> Result<(), MLError> {
        let total = (batch * out_dim) as i32;
        let out_dim_i32 = out_dim as i32;
        let blocks = ((batch * out_dim + 255) / 256) as u32;

        unsafe {
            stream
                .launch_builder(&self.handle.add_bias_f32_kernel)
                .arg(&out_ptr)
                .arg(&bias_ptr)
                .arg(&out_dim_i32)
                .arg(&total)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("add_bias_f32 kernel: {e}")))?;
        }

        Ok(())
    }

}

// ── Raw device pointer helpers ────────────────────────────────────────────────

/// Extract raw F32 device pointer from a `CudaSlice<f32>`.
fn raw_f32_ptr(slice: &CudaSlice<f32>, stream: &Arc<CudaStream>) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = ManuallyDrop::new(guard);
    ptr
}

// ── Cached GEMM descriptor factory ──────────────────────────────────────────

/// Create a `CachedGemmDesc` for a forward GEMM shape (TRANSA=T, TRANSB=N, f32, FAST_TF32).
///
/// Runs the algorithm heuristic ONCE at init time with `requestedAlgoCount=3`
/// for better algorithm selection than the per-call path (which only tried 1).
fn create_cached_fwd_gemm_desc(
    lt_handle: cublaslt_sys::cublasLtHandle_t,
    n: usize,       // out_dim (M in cuBLAS)
    b: usize,       // batch (N in cuBLAS)
    k: usize,       // in_dim
    ldb: usize,     // leading dimension of B
    ws_size: usize,  // workspace size for heuristic
) -> Result<CachedGemmDesc, MLError> {
    let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
    let compute_type = cublaslt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32;

    unsafe {
        // ── Matmul descriptor ──
        let matmul_desc = cublaslt_result::create_matmul_desc(compute_type, f32_type)
            .map_err(|e| MLError::ModelError(format!("cached MatmulDescCreate (n={n},b={b},k={k}): {e:?}")))?;

        let transa: i32 = 1; // CUBLAS_OP_T
        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSA,
            &transa as *const i32 as *const std::ffi::c_void,
            std::mem::size_of::<i32>(),
        ).map_err(|e| MLError::ModelError(format!("cached set TRANSA (n={n},b={b},k={k}): {e:?}")))?;

        let transb: i32 = 0; // CUBLAS_OP_N
        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSB,
            &transb as *const i32 as *const std::ffi::c_void,
            std::mem::size_of::<i32>(),
        ).map_err(|e| MLError::ModelError(format!("cached set TRANSB (n={n},b={b},k={k}): {e:?}")))?;

        // ── Matrix layouts (physical storage, before transpose) ──
        // A = W: col-major [K, N] with ld=K. TRANSA=T gives logical [N, K].
        let a_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, k as i64)
            .map_err(|e| MLError::ModelError(format!("cached A layout (n={n},b={b},k={k}): {e:?}")))?;
        // B = input: col-major [K, B] with ld=ldb. TRANSB=N keeps [K, B].
        let b_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, b as u64, ldb as i64)
            .map_err(|e| MLError::ModelError(format!("cached B layout (n={n},b={b},k={k}): {e:?}")))?;
        // C = output: col-major [N, B] with ld=N.
        let c_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64)
            .map_err(|e| MLError::ModelError(format!("cached C layout (n={n},b={b},k={k}): {e:?}")))?;
        // D layout (separate descriptor, same dimensions as C)
        let d_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64)
            .map_err(|e| MLError::ModelError(format!("cached D layout (n={n},b={b},k={k}): {e:?}")))?;

        // ── Deterministic algorithm selection (Option B fix) ──
        // Replaces timing-based `cublasLtMatmulAlgoGetHeuristic` with a
        // hardware-stable AlgoGetIds → AlgoInit → AlgoCheck loop. TF32
        // compute type preserved. AlgoCheck validates the selected algo
        // against the descriptor + layouts — the workspace filter continues
        // to weed out split-K/stream-K variants that exceed workspace.
        let shape = super::cublas_algo_deterministic::ShapeKey::new(
            transa, transb,
            n as i32, b as i32, k as i32,
            k as i32, ldb as i32, n as i32,
            ws_size,
        );
        let heuristic = super::cublas_algo_deterministic::get_matmul_algo_f32_tf32(
            lt_handle, matmul_desc,
            a_layout, b_layout, c_layout, d_layout,
            shape,
        );
        let heuristic = match heuristic {
            Ok(h) => h,
            Err(e) => {
                let _ = cublaslt_result::destroy_matrix_layout(d_layout);
                let _ = cublaslt_result::destroy_matrix_layout(c_layout);
                let _ = cublaslt_result::destroy_matrix_layout(b_layout);
                let _ = cublaslt_result::destroy_matrix_layout(a_layout);
                let _ = cublaslt_result::destroy_matmul_desc(matmul_desc);
                return Err(MLError::ModelError(format!(
                    "cached fwd deterministic algo (n={n},b={b},k={k},ldb={ldb}): {e}"
                )));
            }
        };
        let algo = heuristic.algo;
        tracing::info!(
            n, b, k, ldb,
            ws_needed = heuristic.workspaceSize,
            "cached fwd GEMM desc created (deterministic algorithm)"
        );

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

/// Create a cached GEMM descriptor with BIAS epilogue (no ReLU).
/// The bias pointer is set dynamically per-call via `set_matmul_desc_attribute`.
/// This fuses GEMM + bias-add into a single cublasLtMatmul kernel, eliminating
/// the separate `add_bias_f32_kernel` launch (Phase H Site 3).
fn create_cached_fwd_gemm_desc_bias(
    lt_handle: cublaslt_sys::cublasLtHandle_t,
    n: usize, b: usize, k: usize, ldb: usize, ws_size: usize,
) -> Result<CachedGemmDesc, MLError> {
    let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
    let compute_type = cublaslt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32;

    unsafe {
        let matmul_desc = cublaslt_result::create_matmul_desc(compute_type, f32_type)
            .map_err(|e| MLError::ModelError(format!("cached+bias MatmulDescCreate (n={n},b={b},k={k}): {e:?}")))?;

        let transa: i32 = 1; // CUBLAS_OP_T
        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSA,
            &transa as *const i32 as *const std::ffi::c_void,
            std::mem::size_of::<i32>(),
        ).map_err(|e| MLError::ModelError(format!("cached+bias set TRANSA: {e:?}")))?;

        let transb: i32 = 0; // CUBLAS_OP_N
        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSB,
            &transb as *const i32 as *const std::ffi::c_void,
            std::mem::size_of::<i32>(),
        ).map_err(|e| MLError::ModelError(format!("cached+bias set TRANSB: {e:?}")))?;

        // Set BIAS epilogue — fuses bias add into the matmul kernel (no ReLU).
        let epilogue: i32 = cublaslt_sys::cublasLtEpilogue_t::CUBLASLT_EPILOGUE_BIAS as i32;
        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_EPILOGUE,
            &epilogue as *const i32 as *const std::ffi::c_void,
            std::mem::size_of::<i32>(),
        ).map_err(|e| MLError::ModelError(format!("cached+bias set EPILOGUE: {e:?}")))?;

        // Set bias data type to F32.
        let bias_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE,
            &bias_type as *const _ as *const std::ffi::c_void,
            std::mem::size_of::<cublaslt_sys::cudaDataType_t>(),
        ).map_err(|e| MLError::ModelError(format!("cached+bias set BIAS_DATA_TYPE: {e:?}")))?;

        let a_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, k as i64)
            .map_err(|e| MLError::ModelError(format!("cached+bias A layout: {e:?}")))?;
        let b_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, b as u64, ldb as i64)
            .map_err(|e| MLError::ModelError(format!("cached+bias B layout: {e:?}")))?;
        let c_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64)
            .map_err(|e| MLError::ModelError(format!("cached+bias C layout: {e:?}")))?;
        let d_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64)
            .map_err(|e| MLError::ModelError(format!("cached+bias D layout: {e:?}")))?;

        // Deterministic algorithm selection — keys on the BIAS epilogue value
        // so the cache produces an algo selection valid for BIAS (independent
        // from DEFAULT and RELU_BIAS entries for the same shape).
        let shape = super::cublas_algo_deterministic::ShapeKey::with_epilogue(
            transa, transb,
            n as i32, b as i32, k as i32,
            k as i32, ldb as i32, n as i32,
            cublaslt_sys::cublasLtEpilogue_t::CUBLASLT_EPILOGUE_BIAS,
            ws_size,
        );
        let heuristic = super::cublas_algo_deterministic::get_matmul_algo_f32_tf32(
            lt_handle, matmul_desc,
            a_layout, b_layout, c_layout, d_layout,
            shape,
        );
        let heuristic = match heuristic {
            Ok(h) => h,
            Err(e) => {
                let _ = cublaslt_result::destroy_matrix_layout(d_layout);
                let _ = cublaslt_result::destroy_matrix_layout(c_layout);
                let _ = cublaslt_result::destroy_matrix_layout(b_layout);
                let _ = cublaslt_result::destroy_matrix_layout(a_layout);
                let _ = cublaslt_result::destroy_matmul_desc(matmul_desc);
                return Err(MLError::ModelError(format!(
                    "cached+bias deterministic algo (n={n},b={b},k={k},ldb={ldb}): {e}"
                )));
            }
        };
        let algo = heuristic.algo;
        tracing::info!(
            n, b, k, ldb,
            ws_needed = heuristic.workspaceSize,
            "cached fwd GEMM+BIAS desc created (deterministic algorithm)"
        );

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

/// Create a cached GEMM descriptor with RELU_BIAS epilogue.
/// The bias pointer is set dynamically per-call via set_matmul_desc_attribute.
/// This fuses GEMM + bias-add + ReLU into a single cublasLtMatmul kernel,
/// eliminating separate add_bias_relu kernel launches.
fn create_cached_fwd_gemm_desc_relu_bias(
    lt_handle: cublaslt_sys::cublasLtHandle_t,
    n: usize, b: usize, k: usize, ldb: usize, ws_size: usize,
) -> Result<CachedGemmDesc, MLError> {
    let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
    let compute_type = cublaslt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32;

    unsafe {
        let matmul_desc = cublaslt_result::create_matmul_desc(compute_type, f32_type)
            .map_err(|e| MLError::ModelError(format!("cached+relu MatmulDescCreate (n={n},b={b},k={k}): {e:?}")))?;

        let transa: i32 = 1; // CUBLAS_OP_T
        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSA,
            &transa as *const i32 as *const std::ffi::c_void,
            std::mem::size_of::<i32>(),
        ).map_err(|e| MLError::ModelError(format!("cached+relu set TRANSA: {e:?}")))?;

        let transb: i32 = 0; // CUBLAS_OP_N
        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSB,
            &transb as *const i32 as *const std::ffi::c_void,
            std::mem::size_of::<i32>(),
        ).map_err(|e| MLError::ModelError(format!("cached+relu set TRANSB: {e:?}")))?;

        // Set RELU_BIAS epilogue — fuses bias add + ReLU into the matmul kernel
        let epilogue: i32 = cublaslt_sys::cublasLtEpilogue_t::CUBLASLT_EPILOGUE_RELU_BIAS as i32;
        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_EPILOGUE,
            &epilogue as *const i32 as *const std::ffi::c_void,
            std::mem::size_of::<i32>(),
        ).map_err(|e| MLError::ModelError(format!("cached+relu set EPILOGUE: {e:?}")))?;

        // Set bias data type to F32
        let bias_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE,
            &bias_type as *const _ as *const std::ffi::c_void,
            std::mem::size_of::<cublaslt_sys::cudaDataType_t>(),
        ).map_err(|e| MLError::ModelError(format!("cached+relu set BIAS_DATA_TYPE: {e:?}")))?;

        let a_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, k as i64)
            .map_err(|e| MLError::ModelError(format!("cached+relu A layout: {e:?}")))?;
        let b_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, b as u64, ldb as i64)
            .map_err(|e| MLError::ModelError(format!("cached+relu B layout: {e:?}")))?;
        let c_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64)
            .map_err(|e| MLError::ModelError(format!("cached+relu C layout: {e:?}")))?;
        let d_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64)
            .map_err(|e| MLError::ModelError(format!("cached+relu D layout: {e:?}")))?;

        // ── Deterministic algorithm selection (Option B fix) ──
        // Replaces timing-based `cublasLtMatmulAlgoGetHeuristic` with a
        // hardware-stable AlgoGetIds → AlgoInit → AlgoCheck loop. TF32
        // compute type preserved. RELU_BIAS epilogue flows through the
        // matmul descriptor, so AlgoCheck validates epilogue compatibility.
        let shape = super::cublas_algo_deterministic::ShapeKey::with_epilogue(
            transa, transb,
            n as i32, b as i32, k as i32,
            k as i32, ldb as i32, n as i32,
            cublaslt_sys::cublasLtEpilogue_t::CUBLASLT_EPILOGUE_RELU_BIAS,
            ws_size,
        );
        let heuristic = super::cublas_algo_deterministic::get_matmul_algo_f32_tf32(
            lt_handle, matmul_desc,
            a_layout, b_layout, c_layout, d_layout,
            shape,
        );
        let heuristic = match heuristic {
            Ok(h) => h,
            Err(e) => {
                let _ = cublaslt_result::destroy_matrix_layout(d_layout);
                let _ = cublaslt_result::destroy_matrix_layout(c_layout);
                let _ = cublaslt_result::destroy_matrix_layout(b_layout);
                let _ = cublaslt_result::destroy_matrix_layout(a_layout);
                let _ = cublaslt_result::destroy_matmul_desc(matmul_desc);
                return Err(MLError::ModelError(format!(
                    "cached+relu deterministic algo (n={n},b={b},k={k},ldb={ldb}): {e}"
                )));
            }
        };
        let algo = heuristic.algo;
        tracing::info!(
            n, b, k, ldb,
            ws_needed = heuristic.workspaceSize,
            "cached fwd GEMM+RELU_BIAS desc created (deterministic algorithm)"
        );

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

/// Create a cached GEMM descriptor with `RELU_AUX_BIAS` epilogue (Phase H Site 2).
///
/// Same as `create_cached_fwd_gemm_desc_relu_bias` but additionally writes the
/// dReLU bit-mask to an auxiliary buffer for backward fusion via DRELU_BGRAD.
/// The bias and AUX pointers are set per-call via `set_matmul_desc_attribute`;
/// the AUX leading dim is baked here at descriptor creation (immutable per shape).
///
/// AUX_LD is in elements (= bits for ReLU mask) and must be a multiple of
/// 128 bits per cuBLAS alignment requirements.
fn create_cached_fwd_gemm_desc_relu_aux_bias(
    lt_handle: cublaslt_sys::cublasLtHandle_t,
    n: usize, b: usize, k: usize, ldb: usize,
    aux_ld_elements: usize,
    ws_size: usize,
) -> Result<CachedGemmDesc, MLError> {
    let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
    let compute_type = cublaslt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32;

    unsafe {
        let matmul_desc = cublaslt_result::create_matmul_desc(compute_type, f32_type)
            .map_err(|e| MLError::ModelError(format!("cached+relu_aux MatmulDescCreate (n={n},b={b},k={k}): {e:?}")))?;

        let transa: i32 = 1; // CUBLAS_OP_T
        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSA,
            &transa as *const i32 as *const std::ffi::c_void,
            std::mem::size_of::<i32>(),
        ).map_err(|e| MLError::ModelError(format!("cached+relu_aux set TRANSA: {e:?}")))?;

        let transb: i32 = 0; // CUBLAS_OP_N
        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSB,
            &transb as *const i32 as *const std::ffi::c_void,
            std::mem::size_of::<i32>(),
        ).map_err(|e| MLError::ModelError(format!("cached+relu_aux set TRANSB: {e:?}")))?;

        // Set RELU_AUX_BIAS epilogue — fuses bias + ReLU into matmul AND writes
        // the dReLU bit-mask to the aux buffer for backward DRELU_BGRAD fusion.
        let epilogue: i32 = cublaslt_sys::cublasLtEpilogue_t::CUBLASLT_EPILOGUE_RELU_AUX_BIAS as i32;
        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_EPILOGUE,
            &epilogue as *const i32 as *const std::ffi::c_void,
            std::mem::size_of::<i32>(),
        ).map_err(|e| MLError::ModelError(format!("cached+relu_aux set EPILOGUE: {e:?}")))?;

        // Bias data type = F32
        let bias_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE,
            &bias_type as *const _ as *const std::ffi::c_void,
            std::mem::size_of::<cublaslt_sys::cudaDataType_t>(),
        ).map_err(|e| MLError::ModelError(format!("cached+relu_aux set BIAS_DATA_TYPE: {e:?}")))?;

        // AUX leading dimension (in elements/bits for ReLU mask). Baked at
        // creation; AUX pointer is set per-call.
        let aux_ld_i64: i64 = aux_ld_elements as i64;
        cublaslt_result::set_matmul_desc_attribute(
            matmul_desc,
            cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_LD,
            &aux_ld_i64 as *const i64 as *const std::ffi::c_void,
            std::mem::size_of::<i64>(),
        ).map_err(|e| MLError::ModelError(format!("cached+relu_aux set AUX_LD: {e:?}")))?;

        let a_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, k as i64)
            .map_err(|e| MLError::ModelError(format!("cached+relu_aux A layout: {e:?}")))?;
        let b_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, b as u64, ldb as i64)
            .map_err(|e| MLError::ModelError(format!("cached+relu_aux B layout: {e:?}")))?;
        let c_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64)
            .map_err(|e| MLError::ModelError(format!("cached+relu_aux C layout: {e:?}")))?;
        let d_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64)
            .map_err(|e| MLError::ModelError(format!("cached+relu_aux D layout: {e:?}")))?;

        // Deterministic algorithm selection — keys on the new epilogue value
        // so the cache produces an algo selection valid for RELU_AUX_BIAS
        // (independent from the RELU_BIAS / DEFAULT entries for the same shape).
        let shape = super::cublas_algo_deterministic::ShapeKey::with_epilogue(
            transa, transb,
            n as i32, b as i32, k as i32,
            k as i32, ldb as i32, n as i32,
            cublaslt_sys::cublasLtEpilogue_t::CUBLASLT_EPILOGUE_RELU_AUX_BIAS,
            ws_size,
        );
        let heuristic = super::cublas_algo_deterministic::get_matmul_algo_f32_tf32(
            lt_handle, matmul_desc,
            a_layout, b_layout, c_layout, d_layout,
            shape,
        );
        let heuristic = match heuristic {
            Ok(h) => h,
            Err(e) => {
                let _ = cublaslt_result::destroy_matrix_layout(d_layout);
                let _ = cublaslt_result::destroy_matrix_layout(c_layout);
                let _ = cublaslt_result::destroy_matrix_layout(b_layout);
                let _ = cublaslt_result::destroy_matrix_layout(a_layout);
                let _ = cublaslt_result::destroy_matmul_desc(matmul_desc);
                return Err(MLError::ModelError(format!(
                    "cached+relu_aux deterministic algo (n={n},b={b},k={k},ldb={ldb}): {e}"
                )));
            }
        };
        let algo = heuristic.algo;
        tracing::info!(
            n, b, k, ldb,
            aux_ld = aux_ld_elements,
            ws_needed = heuristic.workspaceSize,
            "cached fwd GEMM+RELU_AUX_BIAS desc created (deterministic algorithm)"
        );

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

// ── Compute weight pointers from flat params_buf ────────────────────────────

/// Compute the 20 raw device pointers into a flat params_buf at GOFF_* offsets.
///
/// The flat buffer layout matches `compute_param_sizes()`:
/// [w_s1, b_s1, w_s2, b_s2, w_v1, b_v1, w_v2, b_v2,
///  w_b0fc, b_b0fc, w_b0out, b_b0out,
///  w_b1fc, b_b1fc, w_b1out, b_b1out,
///  w_b2fc, b_b2fc, w_b2out, b_b2out]
/// Compute raw F32 device pointers from a `CudaSlice<f32>` params buffer.
pub fn f32_weight_ptrs(
    params_buf: &CudaSlice<f32>,
    param_sizes: &[usize; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
    stream: &Arc<CudaStream>,
) -> [u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS] {
    let base = {
        let (ptr, guard) = params_buf.device_ptr(stream);
        let _no_drop = ManuallyDrop::new(guard);
        ptr
    };
    f32_weight_ptrs_from_base(base, param_sizes)
}

/// #30 Compute raw F32 device pointers from a pre-resolved base pointer.
/// Compute per-tensor device pointers from a flat f32 params buffer.
///
/// Each tensor start is **16-byte aligned** (4 floats) for cublasLtMatmul
/// FAST_TF32 tensor core requirements. Tensors with non-aligned element
/// counts (e.g. bias of 51 atoms) get implicit padding to the next
/// 4-element boundary. The padding elements exist in the buffer but are
/// never read by kernels (bias-add uses the correct count, not the padded size).
pub fn f32_weight_ptrs_from_base(
    base: u64,
    param_sizes: &[usize; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
) -> [u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS] {
    let n = super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS;
    let mut ptrs = [0_u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS];
    let mut byte_offset: u64 = 0;
    for i in 0..n {
        ptrs[i] = base + byte_offset;
        byte_offset += (super::gpu_dqn_trainer::align4(param_sizes[i]) * std::mem::size_of::<f32>()) as u64;
    }
    ptrs
}

