#![allow(unsafe_code)]

//! cuBLAS SGEMM-based batched backward pass for the DQN trainer.
//!
//! Replaces the 1-warp-per-sample `dqn_backward_kernel` with cuBLAS matrix
//! multiplications that process the entire batch in a single GEMM call per
//! layer.  This eliminates `atomicAdd` serialization on the 288K gradient
//! buffer and raises GPU occupancy from 1.56% to >60% on H100.
//!
//! ## TF32 tensor core path (cublasLtMatmul)
//!
//! All backward 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. TF32 tensor cores are activated via the CUBLAS_TF32_TENSOR_OP_MATH
//! math mode set on the cuBLAS handle at construction.
//!
//! ## Layout convention
//!
//! All tensors use row-major (C-style) layout, same as the forward pass:
//! - `dY`:  [B, out_dim]  — upstream gradient
//! - `X`:   [B, in_dim]   — saved activation (input to this layer)
//! - `W`:   [out_dim, in_dim] — weight matrix (row-major)
//! - `dW`:  [out_dim, in_dim] — weight gradient accumulator
//!
//! ## cuBLAS row-major → col-major mapping
//!
//! ### Weight gradient  `dW[out, in] = dY^T[out, B] @ X[B, in]`
//!
//! cuBLAS works column-major.  We observe the dual:
//! - `dW'[in, out]`  col-major  = `X'[in, B]` col-major  @  `dY'^T[B, out]`
//! - Row-major `dW[out, in]` shares the same memory as col-major `dW'[in, out]`.
//! - Row-major `X[B, in]`   shares memory with col-major `X'[in, B]`.
//! - Row-major `dY[B, out]` shares memory with col-major `dY'[out, B]`.
//!
//! Therefore:
//! ```text
//! sgemm(N, T, m=in_dim, n=out_dim, k=B,
//!        1.0, X, lda=in_dim,
//!        dY,  ldb=out_dim,
//!        1.0, dW, ldc=in_dim)
//! ```
//! The `beta=1.0` means ACCUMULATE into `dW` (not overwrite), which is required
//! since multiple layers share the same flat `grad_buf`.
//!
//! ### Upstream gradient  `dX[B, in] = dY[B, out] @ W[out, in]`
//!
//! ```text
//! sgemm(N, N, m=in_dim, n=B, k=out_dim,
//!        1.0, W,  lda=in_dim,
//!        dY,  ldb=out_dim,
//!        0.0, dX, ldc=in_dim)
//! ```
//! Here `beta=0.0` because `dX` is a fresh scratch buffer per layer.
//!
//! ## Bias gradient
//!
//! `db[out] = sum_b(dY[b, out])` — a row-wise reduction.
//! Implemented as a 2-phase shared-memory reduction (`bias_grad_reduce_f32_p1/p2`):
//! Phase 1 does block-level parallel reduce over the batch dimension,
//! Phase 2 sums block partials and accumulates into `db[j]`.
//!
//! ## ReLU derivative mask
//!
//! After computing `dX = W @ dY`, gate by saved pre-ReLU activation:
//! `dX[i] *= (activation[i] > 0)`.
//! Implemented as `relu_mask_kernel` (1 thread per element).
//!
//! ## Phase status
//!
//! This module provides the building blocks for Phase 2 Task 2.  The full
//! The cuBLAS backward path consumes `dL/d_logits` produced by the
//! standalone C51 loss kernel (`c51_loss_batched` in
//! `c51_loss_kernel.cu`) and is wired via
//! `GpuDqnTrainer::launch_cublas_backward()`, which is called from the
//! fused training step. The earlier 1-warp-per-sample
//! `dqn_backward_kernel` (atomicAdd path) has been removed.

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

use cudarc::cublas::sys as cublas_sys;
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;
use super::gpu_dqn_trainer::{GpuDqnTrainConfig, compute_param_sizes, padded_byte_offset};

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

/// Pre-created cublasLt descriptors + heuristic algo for a fixed backward GEMM shape.
///
/// Backward GEMMs have variable TRANSA/TRANSB (N/T for dW, N/N for dX), so we
/// key the cache by the full shape tuple including transpose flags and leading dims.
struct CachedBwdGemmDesc {
    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
/// `CublasBackward`, which is single-threaded in practice.
unsafe impl Send for CachedBwdGemmDesc {}
unsafe impl Sync for CachedBwdGemmDesc {}

impl Drop for CachedBwdGemmDesc {
    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);
        }
    }
}

/// Key for the backward GEMM cache: (transa, transb, m, n, k, lda, ldb, ldc).
/// Beta is a runtime parameter passed to cublasLtMatmul, not part of descriptors.
type BwdGemmKey = (i32, i32, i32, i32, i32, i32, i32, i32);

// ── CublasBackwardSet context ─────────────────────────────────────────────────

/// cuBLAS-based batched backward pass context for the branching dueling DQN.
///
/// Holds a shared `Arc<PerStreamCublasHandles>` (cuBLAS + cublasLt handles,
/// workspaces) plus two NVRTC kernels for operations cuBLAS cannot express
/// directly (bias gradient reduction and ReLU derivative masking).
///
/// ## Usage
///
/// 1. Call `backward_fc_layer()` in reverse layer order, passing the upstream
///    gradient `dY`, the saved forward activation `X`, the weight matrix `W`,
///    and gradient accumulator offsets into the flat `grad_buf`.
/// 2. After each weight layer (except the first), call `relu_mask()` to gate
///    the computed upstream gradient by the saved pre-ReLU activation.
/// 3. For the final full backward, call `backward_full()` which orchestrates
///    all layers automatically.
#[allow(missing_debug_implementations)] // handle is not Debug
pub struct CublasBackwardSet {
    /// Shared cuBLAS + cublasLt handle (owned via Arc, shared with forward pass
    /// and experience collector for deterministic single-handle training).
    handle: Arc<super::shared_cublas_handle::PerStreamCublasHandles>,

    /// `relu_mask_kernel(dx, activation, n)` — element-wise ReLU derivative gate.
    /// Both dx and activation are f32 (entire training pipeline is f32).
    relu_mask_kernel: CudaFunction,

    /// Phase 1: block-level shared-memory reduce over batch dimension.
    bias_grad_p1_kernel: CudaFunction,
    /// Phase 2: final reduce across block partials + accumulate into db.
    bias_grad_p2_kernel: CudaFunction,
    /// Partials scratch [num_blocks, max_out_dim] — shared across all launch_bias_grad calls.
    bias_grad_partials_buf: CudaSlice<f32>,
    /// Number of blocks in the phase-1 grid x-dimension: ceil(batch_size / 256).
    bias_grad_num_blocks: u32,

    /// `kan_gate_backward(d_output, gate_pre, value, spline_coeff, residual_w,
    ///   d_gate_pre, d_value, d_coeff_per_elem, d_coeff_span,
    ///   d_residual_per_elem, n, adv_h)` —
    /// KAN spline backward pass — per-element scatter stage. Writes per-element
    /// basis contributions without atomicAdd.
    kan_gate_backward_kernel: CudaFunction,

    /// `kan_grad_reduce_p1(d_coeff_per_elem, d_coeff_span, d_residual_per_elem,
    ///   partials, n, adv_h, total_params)` —
    /// Phase 1: partial block sums for KAN spline/residual gradients.
    /// Grid=(ceil(B/256), total_params), Block=256, Shared=256*sizeof(float).
    kan_grad_reduce_p1_kernel: CudaFunction,

    /// `kan_grad_reduce_p2(partials, d_spline_coeff, d_residual_w,
    ///   num_blocks, adv_h, total_params)` —
    /// Phase 2: warp-shuffle final reduction for KAN spline/residual gradients.
    /// Grid=(total_params), Block=256.
    kan_grad_reduce_p2_kernel: CudaFunction,

    // ── KAN backward intermediate buffers (allocated once, reused per-call) ──
    /// Per-element B-spline basis contributions [B*AH, 4] — f32.
    kan_d_coeff_per_elem: CudaSlice<f32>,
    /// Per-element span index [B*AH] — i32.
    kan_d_coeff_span: CudaSlice<i32>,
    /// Per-element residual_w gradient contribution [B*AH] — f32.
    kan_d_residual_per_elem: CudaSlice<f32>,
    /// Partials scratch [ceil(B/256), total_params] — f32 (trunk kan reduction).
    kan_reduce_partials: CudaSlice<f32>,

    // ── Network dimensions (baked at construction) ──
    batch_size: usize,
    state_dim: usize,
    state_dim_padded: usize,
    /// First shared layer input dimension (= bottleneck_dim + portfolio_dim when
    /// bottleneck is active, otherwise = state_dim). Used for s1 backward.
    s1_input_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,

    /// Pre-computed parameter sizes for all weight tensors.
    /// Used to compute `padded_byte_offset` for gradient buffer offsets.
    param_sizes: [usize; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],

    // ── Multi-stream branch dispatch ──
    /// 4 forked CUDA streams for parallel advantage branch backward execution.
    /// Each branch submits dW/db GEMMs and KAN backward to its own stream,
    /// then the main stream joins all four before the dX (Step 5) accumulation.
    branch_streams: [Arc<CudaStream>; 4],

    /// Per-branch cublasLt workspace buffers — eliminates workspace contention
    /// when 4 branch streams run GEMMs in parallel. Each 32 MB.
    // Branch workspaces reused from forward pass (same graph, sequential).
    // No separate allocation — saves 128MB VRAM.
    branch_workspace_ptrs: [u64; 4],

    /// Reused CUDA events for fork-join synchronization.
    trunk_done_event: CudaEvent,
    branch_done_events: [CudaEvent; 4],

    // ── Per-branch scratch buffers (safe for parallel execution) ──
    /// Per-branch GLU value gradient scratch [B, AH] each — avoids cross-branch contention.
    branch_d_glu_value: [CudaSlice<f32>; 4],
    /// Per-branch GLU gate_pre gradient scratch [B, AH] each.
    branch_d_glu_gate: [CudaSlice<f32>; 4],

    // ── Per-branch KAN backward intermediate buffers ──
    /// Per-element B-spline basis contributions [B*AH, 4] per branch — f32.
    branch_kan_d_coeff_per_elem: [CudaSlice<f32>; 4],
    /// Per-element span index [B*AH] per branch — i32.
    branch_kan_d_coeff_span: [CudaSlice<i32>; 4],
    /// Per-element residual_w gradient contribution [B*AH] per branch — f32.
    branch_kan_d_residual_per_elem: [CudaSlice<f32>; 4],
    /// Partials scratch [ceil(B/256), total_params] per branch — f32 (2-phase KAN reduce).
    branch_kan_reduce_partials: [CudaSlice<f32>; 4],

    // ── Cached GEMM descriptors (created once at init, reused per-call) ──
    /// Map from (transa, transb, m, n, k, lda, ldb, ldc) → pre-created descriptors + algo.
    gemm_cache: HashMap<BwdGemmKey, CachedBwdGemmDesc>,

    /// Phase H Site 2 — cached descriptor for the value-output dX GEMM with
    /// `EPILOGUE_DRELU_BGRAD`. Same shape key as the regular value-output dX
    /// in `gemm_cache`, but a separate descriptor because epilogue is part of
    /// the descriptor (and the deterministic-algo cache keys on epilogue).
    /// `Option`: absent when the matching forward `RELU_AUX_BIAS` algo is
    /// unavailable (cuBLAS heuristic mismatch) — backward then falls back to
    /// the standalone `relu_mask` + `bias_grad` reduce sequence.
    gemm_value_dx_drelu_bgrad: Option<CachedBwdGemmDesc>,
    /// Phase H Site 2 — AUX leading dim (in elements/bits) for the value-output
    /// dX GEMM's DRELU_BGRAD epilogue. Mirrors the forward's
    /// `value_fc_relu_mask_aux_ld` so the same bit-mask is consumed correctly.
    gemm_value_dx_drelu_bgrad_aux_ld: usize,
}

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

impl CublasBackwardSet {
    /// Create a new cuBLAS backward context backed by a shared handle.
    ///
    /// The shared handle owns the cuBLAS + cublasLt handles and workspaces.
    /// This constructor only creates backward-specific resources: the
    /// `relu_mask_kernel`, `bias_grad_reduce_f32_p1/p2` (compiled from the
    /// precompiled cubin), and cached GEMM descriptors.
    pub fn new(
        shared: Arc<super::shared_cublas_handle::PerStreamCublasHandles>,
        config: &GpuDqnTrainConfig,
        kan_gate_backward_kernel: CudaFunction,
        kan_grad_reduce_p1_kernel: CudaFunction,
        kan_grad_reduce_p2_kernel: CudaFunction,
        fwd_branch_workspace_ptrs: [u64; 4],
    ) -> Result<Self, MLError> {
        let stream = &shared.stream;
        let lt_raw_handle = shared.lt_handle.0;
        let lt_ws_size = shared.lt_workspace_size;

        // ── Compile helper kernels ──────────────────────────────────
        let (relu_mask_kernel, bias_grad_p1_kernel, bias_grad_p2_kernel) =
            compile_backward_kernels(stream)?;

        let param_sizes = compute_param_sizes(config);

        // ── Pre-create backward GEMM descriptors for all unique shapes ──
        let batch = config.batch_size;
        let sh1 = config.shared_h1;
        let sh2 = config.shared_h2;
        let vh = config.value_h;
        let ah = config.adv_h;
        let na = config.num_atoms;
        let _sd = ml_core::state_layout::STATE_DIM;
        let sd_pad = (ml_core::state_layout::STATE_DIM + 127) & !127;
        // SP15 Wave 4.1b: the bottleneck-on s1_input_dim grows by +1 to carry
        // the dd_pct column appended by `bn_tanh_concat_dd_kernel`. Mirrors
        // `compute_param_sizes` so the backward gemm cache shape keys match
        // forward's (and the encoder_backward_chain dX width sees the same
        // 103-column buffer).
        let s1_input_dim = if config.bottleneck_dim > 0 {
            config.bottleneck_dim
                + ml_core::state_layout::STATE_DIM.saturating_sub(config.market_dim)
                + 1
        } else {
            ml_core::state_layout::STATE_DIM
        };
        let branch_sizes = [config.branch_0_size, config.branch_1_size, config.branch_2_size, config.branch_3_size];

        let op_n = cublas_sys::cublasOperation_t::CUBLAS_OP_N as i32;
        let op_t = cublas_sys::cublasOperation_t::CUBLAS_OP_T as i32;

        // Collect all unique (transa, transb, m, n, k, lda, ldb, ldc) shapes.
        let mut unique_shapes: Vec<BwdGemmKey> = Vec::new();

        // For each FC layer with (out_dim, in_dim), backward produces:
        //   dW: (N, T, in_dim, out_dim, batch, in_dim, out_dim, in_dim)
        //   dX: (N, N, in_dim, batch, out_dim, in_dim, out_dim, in_dim)
        // For lda variants (shared layer 1 with padded input):
        //   dW: (N, T, in_dim, out_dim, batch, x_lda, out_dim, in_dim)
        //   dX: (N, N, in_dim, batch, out_dim, in_dim, out_dim, x_lda)

        let add_fc_shapes = |shapes: &mut Vec<BwdGemmKey>, out_dim: usize, in_dim: usize| {
            // dW: sgemm(N, T, in_dim, out_dim, batch, in_dim, out_dim, in_dim)
            shapes.push((op_n, op_t, in_dim as i32, out_dim as i32, batch as i32, in_dim as i32, out_dim as i32, in_dim as i32));
            // dX: sgemm(N, N, in_dim, batch, out_dim, in_dim, out_dim, in_dim)
            shapes.push((op_n, op_n, in_dim as i32, batch as i32, out_dim as i32, in_dim as i32, out_dim as i32, in_dim as i32));
        };

        // Branch output layers: (branch_k*NA, adv_h)
        for &bs in &branch_sizes {
            add_fc_shapes(&mut unique_shapes, bs * na, ah);
        }
        // Branch FC layers: (adv_h, shared_h2)
        add_fc_shapes(&mut unique_shapes, ah, sh2);
        // Branch 1 FC (magnitude) wider input (direction-conditioned): (adv_h, shared_h2 + branch_0_size)
        add_fc_shapes(&mut unique_shapes, ah, sh2 + config.branch_0_size);
        // Branches 2/3 FC (order/urgency) wider input (OFI-conditioned, 3 features per branch): (adv_h, shared_h2 + 3)
        add_fc_shapes(&mut unique_shapes, ah, sh2 + 3);
        // Value output: (NA, value_h)
        add_fc_shapes(&mut unique_shapes, na, vh);
        // Value FC: (value_h, shared_h2)
        add_fc_shapes(&mut unique_shapes, vh, sh2);
        // Shared layer 2: (shared_h2, shared_h1)
        add_fc_shapes(&mut unique_shapes, sh2, sh1);
        // Shared layer 1 with padded input: (shared_h1, state_dim) — uses x_lda=state_dim_padded
        // When bottleneck is active, s1_input_dim < state_dim (e.g. 54 vs 80).
        // Weight W_s1 is [SH1, s1_input_dim], input X has stride sd_pad (no bottleneck) or s1_input_dim (bottleneck).
        let s1d = s1_input_dim;
        let s1_x_lda = if config.bottleneck_dim > 0 { s1_input_dim } else { sd_pad };
        // dW: (N, T, s1_input_dim, sh1, batch, s1_x_lda, sh1, s1_input_dim)
        unique_shapes.push((op_n, op_t, s1d as i32, sh1 as i32, batch as i32, s1_x_lda as i32, sh1 as i32, s1d as i32));
        // dX: (N, N, s1_input_dim, batch, sh1, s1_input_dim, sh1, s1_x_lda) — only when s1_dx_output != 0
        unique_shapes.push((op_n, op_n, s1d as i32, batch as i32, sh1 as i32, s1d as i32, sh1 as i32, s1_x_lda as i32));

        unique_shapes.sort();
        unique_shapes.dedup();

        let mut gemm_cache = HashMap::new();
        for &key in &unique_shapes {
            let (transa_i32, transb_i32, m, n, k, lda, ldb, ldc) = key;
            let desc = create_cached_bwd_gemm_desc(lt_raw_handle, transa_i32, transb_i32, m, n, k, lda, ldb, ldc, lt_ws_size)?;
            gemm_cache.insert(key, desc);
        }

        // ── Phase H Site 2: DRELU_BGRAD descriptor for value-output dX ──
        //
        // The value-output backward dX GEMM is:
        //   dX[VH, B] = W_v2[VH, NA] @ dY[NA, B]
        //   sgemm(N, N, VH, B, NA, 1.0, w_v2, VH, dy, NA, 0.0, dx, VH)
        //
        // Applying `EPILOGUE_DRELU_BGRAD` to this GEMM:
        //   1. Gates D in-place by the dReLU bit-mask read from EPILOGUE_AUX_POINTER
        //      (mask was written by the value-FC forward via RELU_AUX_BIAS).
        //   2. Reduces D along the column (N=B) dim into BIAS_POINTER → vector
        //      of length M = VH = b_v1's gradient shape.
        //
        // After this single GEMM, both the standalone `relu_mask` and
        // `launch_bias_grad` at the value-FC site are redundant.
        //
        // AUX_LD must match the forward's AUX_LD (in elements/bits). We
        // recompute it here independently (the forward and backward agree
        // on `(value_h + 127) & !127` to keep them in sync without coupling).
        let value_dx_aux_ld: usize = (vh + 127) & !127;
        let gemm_value_dx_drelu_bgrad = match create_cached_bwd_gemm_desc_drelu_bgrad(
            lt_raw_handle,
            op_n, op_n,
            vh as i32, batch as i32, na as i32,
            vh as i32, na as i32, vh as i32,
            value_dx_aux_ld,
            lt_ws_size,
        ) {
            Ok(d) => Some(d),
            Err(e) => {
                tracing::warn!(
                    "DRELU_BGRAD descriptor unavailable for value-output dX (vh={vh}, b={batch}, na={na}): {e}. \
                     Falling back to standalone relu_mask + bias_grad."
                );
                None
            }
        };

        // ── KAN backward intermediate buffers ─────────────────────────
        let kan_n = config.batch_size * config.adv_h;
        let kan_d_coeff_per_elem = stream.alloc_zeros::<f32>(kan_n * 4)
            .map_err(|e| MLError::ModelError(format!("kan_d_coeff_per_elem alloc [{kan_n}*4]: {e}")))?;
        let kan_d_coeff_span = stream.alloc_zeros::<i32>(kan_n)
            .map_err(|e| MLError::ModelError(format!("kan_d_coeff_span alloc [{kan_n}]: {e}")))?;
        let kan_d_residual_per_elem = stream.alloc_zeros::<f32>(kan_n)
            .map_err(|e| MLError::ModelError(format!("kan_d_residual_per_elem alloc [{kan_n}]: {e}")))?;

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

        // Pre-register each backward branch stream in the per-stream cublasLt
        // registry (Option C). Must happen BEFORE CUDA Graph capture.
        for bs in &branch_streams {
            shared.pre_register_stream(bs)?;
        }

        // Reuse the FORWARD pass's branch workspace pointers — forward and backward
        // are in the same child graph (sequential), so workspaces never conflict.
        // Saves 128MB VRAM (4 × 32MB) that was causing OOM on H100.
        let branch_workspace_ptrs = fwd_branch_workspace_ptrs;

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

        // ── Per-branch GLU scratch buffers ──────────────────────────
        let b_ah = config.batch_size * config.adv_h;
        let alloc_f32 = |n: usize, label: &str| -> Result<CudaSlice<f32>, MLError> {
            stream.alloc_zeros::<f32>(n)
                .map_err(|e| MLError::ModelError(format!("{label} alloc [{n}]: {e}")))
        };
        let branch_d_glu_value = [
            alloc_f32(b_ah, "branch_d_glu_value_0")?,
            alloc_f32(b_ah, "branch_d_glu_value_1")?,
            alloc_f32(b_ah, "branch_d_glu_value_2")?,
            alloc_f32(b_ah, "branch_d_glu_value_3")?,
        ];
        let branch_d_glu_gate = [
            alloc_f32(b_ah, "branch_d_glu_gate_0")?,
            alloc_f32(b_ah, "branch_d_glu_gate_1")?,
            alloc_f32(b_ah, "branch_d_glu_gate_2")?,
            alloc_f32(b_ah, "branch_d_glu_gate_3")?,
        ];

        // ── Per-branch KAN backward intermediate buffers ────────────
        let alloc_i32 = |n: usize, label: &str| -> Result<CudaSlice<i32>, MLError> {
            stream.alloc_zeros::<i32>(n)
                .map_err(|e| MLError::ModelError(format!("{label} alloc [{n}]: {e}")))
        };
        let branch_kan_d_coeff_per_elem = [
            alloc_f32(kan_n * 4, "branch_kan_d_coeff_pe_0")?,
            alloc_f32(kan_n * 4, "branch_kan_d_coeff_pe_1")?,
            alloc_f32(kan_n * 4, "branch_kan_d_coeff_pe_2")?,
            alloc_f32(kan_n * 4, "branch_kan_d_coeff_pe_3")?,
        ];
        let branch_kan_d_coeff_span = [
            alloc_i32(kan_n, "branch_kan_d_coeff_span_0")?,
            alloc_i32(kan_n, "branch_kan_d_coeff_span_1")?,
            alloc_i32(kan_n, "branch_kan_d_coeff_span_2")?,
            alloc_i32(kan_n, "branch_kan_d_coeff_span_3")?,
        ];
        let branch_kan_d_residual_per_elem = [
            alloc_f32(kan_n, "branch_kan_d_resid_pe_0")?,
            alloc_f32(kan_n, "branch_kan_d_resid_pe_1")?,
            alloc_f32(kan_n, "branch_kan_d_resid_pe_2")?,
            alloc_f32(kan_n, "branch_kan_d_resid_pe_3")?,
        ];

        // ── KAN 2-phase reduce partials buffers ────────────────────────
        // total_params = adv_h * 9 (8 spline coeffs + 1 residual per neuron)
        // num_blocks = ceil(batch_size / 256)
        let kan_total_params = config.adv_h * 9;
        let kan_num_blocks = (config.batch_size + 255) / 256;
        let kan_partials_size = kan_num_blocks * kan_total_params;
        let kan_reduce_partials = alloc_f32(kan_partials_size, "kan_reduce_partials")?;
        let branch_kan_reduce_partials = [
            alloc_f32(kan_partials_size, "branch_kan_reduce_partials_0")?,
            alloc_f32(kan_partials_size, "branch_kan_reduce_partials_1")?,
            alloc_f32(kan_partials_size, "branch_kan_reduce_partials_2")?,
            alloc_f32(kan_partials_size, "branch_kan_reduce_partials_3")?,
        ];

        // ── Bias grad 2-phase reduce partials buffer ──────────────────
        // num_blocks = ceil(batch_size / 256), max_out_dim = max across all FC layers.
        let bias_grad_num_blocks = ((batch + 255) / 256) as u32;
        // Safe upper bound: total_params covers any single layer's out_dim. Buffer is tiny (num_blocks × max_od × 4).
        let max_out_dim = (ml_core::state_layout::STATE_DIM).max(sh1).max(sh2).max(vh).max(ah).max(na)
            .max(config.num_atoms + 32).max(config.batch_size)  // some layers reduce over batch-sized dims
            .max(512);  // absolute floor — partials buffer is <64KB even at 512
        let bias_grad_partials_buf = alloc_f32(
            bias_grad_num_blocks as usize * max_out_dim,
            "bias_grad_partials",
        )?;

        Ok(Self {
            handle: shared,
            relu_mask_kernel,
            bias_grad_p1_kernel,
            bias_grad_p2_kernel,
            bias_grad_partials_buf,
            bias_grad_num_blocks,
            kan_gate_backward_kernel,
            kan_grad_reduce_p1_kernel,
            kan_grad_reduce_p2_kernel,
            kan_d_coeff_per_elem,
            kan_d_coeff_span,
            kan_d_residual_per_elem,
            kan_reduce_partials,
            batch_size: config.batch_size,
            state_dim: ml_core::state_layout::STATE_DIM,
            state_dim_padded: sd_pad,
            s1_input_dim,
            shared_h1: config.shared_h1,
            shared_h2: config.shared_h2,
            value_h: config.value_h,
            adv_h: config.adv_h,
            num_atoms: config.num_atoms,
            branch_0_size: config.branch_0_size,
            branch_1_size: config.branch_1_size,
            branch_2_size: config.branch_2_size,
            branch_3_size: config.branch_3_size,
            branch_streams,
            branch_workspace_ptrs,
            trunk_done_event,
            branch_done_events,
            branch_d_glu_value,
            branch_d_glu_gate,
            branch_kan_d_coeff_per_elem,
            branch_kan_d_coeff_span,
            branch_kan_d_residual_per_elem,
            branch_kan_reduce_partials,
            param_sizes,
            gemm_cache,
            gemm_value_dx_drelu_bgrad,
            gemm_value_dx_drelu_bgrad_aux_ld: value_dx_aux_ld,
        })
    }

    // ══════════════════════════════════════════════════════════════════════════
    // Pure f32 SGEMM (the ONLY backward GEMM — no f32 anywhere)
    // ══════════════════════════════════════════════════════════════════════════

    /// F32 GEMM with TF32 tensor core acceleration via cublasLtMatmul.
    ///
    /// cublasLtMatmul passes workspace+stream per-call (no handle state),
    /// eliminating the `cublasSetStream` workspace conflict that hangs
    /// CUDA Graph mega-capture on H100.
    ///
    /// Uses cached descriptors when the shape is pre-cached at init time.
    /// Falls back to inline descriptor creation for uncached shapes.
    ///
    /// Uses the per-stream workspace from the registry (Option C). For side
    /// streams (branch/iqn/attn) the registry has provisioned an isolated
    /// workspace on first access.
    fn sgemm_f32(
        &self,
        stream: &CudaStream,
        transa: cublas_sys::cublasOperation_t,
        transb: cublas_sys::cublasOperation_t,
        m: i32, n: i32, k: i32,
        alpha: f32,
        a: u64, lda: i32,
        b_ptr: u64, ldb: i32,
        beta: f32,
        c: u64, ldc: i32,
        label: &str,
    ) -> Result<(), MLError> {
        let (_, ws_ptr, _) = self.handle.lt_for(stream)?;
        self.sgemm_f32_ws(stream, transa, transb, m, n, k, alpha, a, lda, b_ptr, ldb, beta, c, ldc, ws_ptr, label)
    }

    /// F32 GEMM with explicit workspace pointer.
    ///
    /// Branch streams MUST use per-branch workspace pointers to avoid
    /// contention with the shared workspace used by the main stream.
    #[allow(clippy::too_many_arguments)]
    fn sgemm_f32_ws(
        &self,
        stream: &CudaStream,
        transa: cublas_sys::cublasOperation_t,
        transb: cublas_sys::cublasOperation_t,
        m: i32, n: i32, k: i32,
        alpha: f32,
        a: u64, lda: i32,
        b_ptr: u64, ldb: i32,
        beta: f32,
        c: u64, ldc: i32,
        ws_ptr: u64,
        label: &str,
    ) -> Result<(), MLError> {
        let key: BwdGemmKey = (transa as i32, transb as i32, m, n, k, lda, ldb, ldc);
        if let Some(cached) = self.gemm_cache.get(&key) {
            // ── Fast path: use pre-created descriptors + algo ──
            self.sgemm_f32_cached_ws(stream, cached, alpha, a, b_ptr, beta, c, m, n, k, lda, ldb, ldc, ws_ptr, label)
        } else {
            // ── Slow path: inline descriptor creation ──
            self.sgemm_f32_uncached_ws(stream, transa, transb, m, n, k, alpha, a, lda, b_ptr, ldb, beta, c, ldc, ws_ptr, label)
        }
    }

    /// Fast path: execute backward GEMM with pre-cached descriptors (per-stream workspace).
    #[allow(clippy::too_many_arguments)]
    fn sgemm_f32_cached(
        &self,
        stream: &CudaStream,
        cached: &CachedBwdGemmDesc,
        alpha: f32,
        a: u64,
        b_ptr: u64,
        beta: f32,
        c: u64,
        m: i32, n: i32, k: i32,
        lda: i32, ldb: i32, ldc: i32,
        label: &str,
    ) -> Result<(), MLError> {
        let (_, ws_ptr, _) = self.handle.lt_for(stream)?;
        self.sgemm_f32_cached_ws(stream, cached, alpha, a, b_ptr, beta, c, m, n, k, lda, ldb, ldc, ws_ptr, label)
    }

    /// Fast path: execute backward GEMM with pre-cached descriptors and explicit workspace.
    #[allow(clippy::too_many_arguments)]
    fn sgemm_f32_cached_ws(
        &self,
        stream: &CudaStream,
        cached: &CachedBwdGemmDesc,
        alpha: f32,
        a: u64,
        b_ptr: u64,
        beta: f32,
        c: u64,
        m: i32, n: i32, k: i32,
        lda: i32, ldb: i32, ldc: i32,
        ws_ptr: u64,
        label: &str,
    ) -> Result<(), MLError> {
        // Per-stream lt_handle lookup (Option C determinism fix).
        let lt_handle = self.handle.lt_handle_for(stream)?;
        let lt_workspace_size = self.handle.lt_workspace_size;

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

            let matmul_status = cublaslt_sys::cublasLtMatmul(
                lt_handle,
                cached.matmul_desc,
                &alpha as *const f32 as *const std::ffi::c_void,
                a as *const std::ffi::c_void,
                cached.a_layout,
                b_ptr as *const std::ffi::c_void,
                cached.b_layout,
                &beta as *const f32 as *const std::ffi::c_void,
                c as *const std::ffi::c_void,
                cached.c_layout,
                c 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,
                lt_workspace_size,
                cu_stream,
            );

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

    /// Slow path: inline descriptor creation for uncached backward GEMM shapes (per-stream workspace).
    #[allow(clippy::too_many_arguments)]
    fn sgemm_f32_uncached(
        &self,
        stream: &CudaStream,
        transa: cublas_sys::cublasOperation_t,
        transb: cublas_sys::cublasOperation_t,
        m: i32, n: i32, k: i32,
        alpha: f32,
        a: u64, lda: i32,
        b_ptr: u64, ldb: i32,
        beta: f32,
        c: u64, ldc: i32,
        label: &str,
    ) -> Result<(), MLError> {
        let (_, ws_ptr, _) = self.handle.lt_for(stream)?;
        self.sgemm_f32_uncached_ws(stream, transa, transb, m, n, k, alpha, a, lda, b_ptr, ldb, beta, c, ldc, ws_ptr, label)
    }

    /// Slow path: inline descriptor creation with explicit workspace pointer.
    #[allow(clippy::too_many_arguments)]
    fn sgemm_f32_uncached_ws(
        &self,
        stream: &CudaStream,
        transa: cublas_sys::cublasOperation_t,
        transb: cublas_sys::cublasOperation_t,
        m: i32, n: i32, k: i32,
        alpha: f32,
        a: u64, lda: i32,
        b_ptr: u64, ldb: i32,
        beta: f32,
        c: u64, ldc: i32,
        ws_ptr: u64,
        label: &str,
    ) -> Result<(), MLError> {
        let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
        // FAST_TF32: explicit TF32 tensor core path. Required on cuBLAS 13.0+ (SM90)
        // where COMPUTE_32F alone returns NOT_SUPPORTED for certain algorithm heuristics.
        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 bw {label}: {e:?}")))?;

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

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

            // ── Matrix layouts ──
            let a_layout = if transa == cublas_sys::cublasOperation_t::CUBLAS_OP_N {
                cublaslt_result::create_matrix_layout(f32_type, m as u64, k as u64, lda as i64)
            } else {
                cublaslt_result::create_matrix_layout(f32_type, k as u64, m as u64, lda as i64)
            }.map_err(|e| MLError::ModelError(format!("A layout bw {label}: {e:?}")))?;

            let b_layout = if transb == cublas_sys::cublasOperation_t::CUBLAS_OP_N {
                cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, ldb as i64)
            } else {
                cublaslt_result::create_matrix_layout(f32_type, n as u64, k as u64, ldb as i64)
            }.map_err(|e| MLError::ModelError(format!("B layout bw {label}: {e:?}")))?;

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

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

            // Per-stream lt_handle (Option C determinism fix). Used for both
            // deterministic algo selection and the matmul dispatch itself.
            let lt_handle = self.handle.lt_handle_for(stream)?;
            let lt_workspace_size = self.handle.lt_workspace_size;

            // ── Deterministic algorithm selection (Option B fix) ──
            // Replaces timing-based `cublasLtMatmulAlgoGetHeuristic` with a
            // hardware-stable AlgoGetIds → AlgoInit → AlgoCheck loop. TF32
            // compute type preserved. See cublas_algo_deterministic.rs.
            let shape = super::cublas_algo_deterministic::ShapeKey::new(
                transa_i32, transb_i32, m, n, k, lda, ldb, ldc,
                lt_workspace_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=m, n=n, k=k, lda=lda, ldb=ldb, ldc=ldc, ws=lt_workspace_size, ?e, "cublasLt deterministic algo selection FAILED for bw {label}");
            }
            let heuristic = heuristic.map_err(|e| MLError::ModelError(format!("algo deterministic bw {label} (m={m},n={n},k={k},lda={lda},ldb={ldb}): {e}")))?;

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

            let matmul_result = cublaslt_result::matmul(
                lt_handle,
                matmul_desc,
                &alpha as *const f32 as *const std::ffi::c_void,
                &beta as *const f32 as *const std::ffi::c_void,
                a as *const std::ffi::c_void,
                a_layout,
                b_ptr as *const std::ffi::c_void,
                b_layout,
                c as *const std::ffi::c_void,
                c_layout,
                c as *mut std::ffi::c_void,
                d_layout,
                &heuristic.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t,
                ws_ptr as *mut std::ffi::c_void,
                lt_workspace_size,
                cu_stream,
            );
            if let Err(ref e) = matmul_result {
                tracing::error!(
                    m=m, n=n, k=k, lda=lda, ldb=ldb, ldc=ldc,
                    ws_provided=lt_workspace_size,
                    ws_needed=heuristic.workspaceSize,
                    ws_ptr=ws_ptr,
                    ?e,
                    "cublasLtMatmul UNCACHED FAILED for bw {label}"
                );
            }
            matmul_result.map_err(|e| MLError::ModelError(format!("cublasLtMatmul bw {label} (m={m},n={n},k={k},lda={lda},ldb={ldb},ws={},needed={}): {e:?}", lt_workspace_size, heuristic.workspaceSize)))?;

            // ── Cleanup 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(())
    }

    // ══════════════════════════════════════════════════════════════════════════
    // Public building blocks
    // ══════════════════════════════════════════════════════════════════════════

    /// Backward through one fully-connected layer using cuBLAS SGEMM.
    ///
    /// Computes:
    /// - `dW += dY^T @ X`   (weight gradient, accumulated with beta=1.0)
    /// - `db += sum_b(dY)`  (bias gradient, accumulated via atomicAdd)
    /// - `dX  = dY  @ W^T`  (upstream gradient, overwritten with beta=0.0)
    ///
    /// This does NOT apply the ReLU derivative mask; call `relu_mask()` after
    /// if the layer had a ReLU activation.
    ///
    /// ## Arguments
    ///
    /// - `dy`      : raw u64 device pointer to `dY[B, out_dim]` (upstream gradient)
    /// - `x`       : raw u64 device pointer to `X[B, in_dim]`  (saved input activation)
    /// - `w`       : raw u64 device pointer to `W[out_dim, in_dim]` (weight matrix)
    /// - `dw`      : raw u64 device pointer to `dW[out_dim, in_dim]` in `grad_buf`
    /// - `db`      : raw u64 device pointer to `db[out_dim]` in `grad_buf`
    /// - `dx`      : raw u64 device pointer to `dX[B, in_dim]` scratch (may be 0
    ///               if upstream gradient is not needed, e.g. input layer)
    /// - `out_dim` : number of neurons in this layer's output
    /// - `in_dim`  : number of neurons in this layer's input
    /// - `batch`   : batch size B
    #[allow(clippy::too_many_arguments)]
    pub fn backward_fc_layer(
        &self,
        stream: &Arc<CudaStream>,
        dy: u64,       // f32
        x: u64,        // f32 (saved activations)
        w: u64,        // f32 (master weights)
        dw: u64,       // f32 (grad_buf)
        db: u64,       // f32 (grad_buf)
        dx: u64,       // f32 (scratch) — 0 means skip
        out_dim: usize,
        in_dim: usize,
        batch: usize,
    ) -> Result<(), MLError> {
        // ── Weight gradient: dW[out, in] += dY^T @ X ──
        self.sgemm_f32(
            stream,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            in_dim as i32, out_dim as i32, batch as i32,
            1.0, x, in_dim as i32,
            dy, out_dim as i32,
            1.0, dw, in_dim as i32,
            "fc_dW",
        )?;

        self.launch_bias_grad(stream, dy, db, out_dim, batch)?;

        // ── Upstream gradient: dX[B, in] = dY[B, out] @ W[out, in] ──
        if dx != 0 {
            self.sgemm_f32(
                stream,
                cublas_sys::cublasOperation_t::CUBLAS_OP_N,
                cublas_sys::cublasOperation_t::CUBLAS_OP_N,
                in_dim as i32, batch as i32, out_dim as i32,
                1.0, w, in_dim as i32,
                dy, out_dim as i32,
                0.0, dx, in_dim as i32,
                "fc_dX",
            )?;
        }

        Ok(())
    }

    /// Like `backward_fc_layer` but with a custom leading dimension for X
    /// (padded row stride, e.g. pad128(state_dim) for CUTLASS K-tile).
    pub fn backward_fc_layer_lda(
        &self,
        stream: &Arc<CudaStream>,
        dy: u64, x: u64, w: u64, dw: u64, db: u64, dx: u64,
        out_dim: usize, in_dim: usize, x_lda: usize, batch: usize,
    ) -> Result<(), MLError> {
        self.sgemm_f32(
            stream,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            in_dim as i32, out_dim as i32, batch as i32,
            1.0, x, x_lda as i32,
            dy, out_dim as i32,
            1.0, dw, in_dim as i32,
            "fc_dW_lda",
        )?;

        self.launch_bias_grad(stream, dy, db, out_dim, batch)?;

        if dx != 0 {
            self.sgemm_f32(
                stream,
                cublas_sys::cublasOperation_t::CUBLAS_OP_N,
                cublas_sys::cublasOperation_t::CUBLAS_OP_N,
                in_dim as i32, batch as i32, out_dim as i32,
                1.0, w, in_dim as i32,
                dy, out_dim as i32,
                0.0, dx, x_lda as i32,
                "fc_dX_lda",
            )?;
        }

        Ok(())
    }

    /// Phase H Site 2 — value-output backward with `EPILOGUE_DRELU_BGRAD` on dX.
    ///
    /// Same contract as `backward_fc_layer` for the dW + db (output-layer
    /// bias gradient) computations, but the dX GEMM uses `DRELU_BGRAD`:
    /// gates the output by the dReLU bit-mask saved by the value-FC forward
    /// AND reduces the gated output along N (=batch) into `next_layer_db`.
    ///
    /// After this single GEMM:
    ///   - `dx` (= `scratch_d_h_v`) holds the dReLU-gated upstream gradient
    ///     ready for the value-FC dW/dX (no separate `relu_mask` needed).
    ///   - `next_layer_db` (= `b_v1` slot in grad_buf) holds the value-FC
    ///     bias gradient (no separate `bias_grad` reduce needed).
    ///
    /// Falls back to the regular `backward_fc_layer` + caller-side `relu_mask`
    /// + `bias_grad` sequence when the cached DRELU_BGRAD descriptor is
    /// unavailable (returns `Ok(false)` to signal the caller to do the
    /// fallback work).
    #[allow(clippy::too_many_arguments)]
    pub fn backward_fc_layer_drelu_bgrad(
        &self,
        stream: &Arc<CudaStream>,
        dy: u64,             // [B, out_dim] f32 — upstream gradient
        x: u64,              // [B, in_dim] f32 — saved input (post-activation of THIS layer)
        w: u64,              // [out_dim, in_dim] f32 — master weights
        dw: u64,             // grad_buf — output dW
        db: u64,             // grad_buf — output bias gradient FOR THIS LAYER (the output layer)
        dx: u64,             // [B, in_dim] f32 — output upstream gradient (gated by drelu)
        aux_mask: u64,       // dReLU bit-mask written by value-FC forward (RELU_AUX_BIAS)
        next_layer_db: u64,  // grad_buf — bias gradient FOR THE NEXT LAYER (value-FC)
        out_dim: usize,
        in_dim: usize,
        batch: usize,
    ) -> Result<bool, MLError> {
        let cached = match self.gemm_value_dx_drelu_bgrad.as_ref() {
            Some(c) => c,
            None => return Ok(false), // signal fallback to caller
        };

        // ── dW + db (output-layer bias) — same as backward_fc_layer ──
        self.sgemm_f32(
            stream,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            in_dim as i32, out_dim as i32, batch as i32,
            1.0, x, in_dim as i32,
            dy, out_dim as i32,
            1.0, dw, in_dim as i32,
            "fc_dW_drelu_bgrad",
        )?;
        self.launch_bias_grad(stream, dy, db, out_dim, batch)?;

        // ── Fused dX + drelu + bgrad → next_layer_db ──
        //
        // sgemm(N, N, in_dim=VH, batch=B, out_dim=NA, 1.0, w, VH, dy, NA, 0.0, dx, VH)
        // with EPILOGUE_DRELU_BGRAD: gates D by aux_mask (in-place) and reduces
        // gated D along N=batch dim into next_layer_db (length M=in_dim=VH).
        //
        // Set per-call AUX_POINTER (mask) + BIAS_POINTER (next-layer bgrad
        // destination) on the cached descriptor before launching.
        unsafe {
            cublaslt_result::set_matmul_desc_attribute(
                cached.matmul_desc,
                cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_POINTER,
                &aux_mask as *const u64 as *const std::ffi::c_void,
                std::mem::size_of::<u64>(),
            ).map_err(|e| MLError::ModelError(format!("set drelu_bgrad AUX_POINTER: {e:?}")))?;
            cublaslt_result::set_matmul_desc_attribute(
                cached.matmul_desc,
                cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_BIAS_POINTER,
                &next_layer_db as *const u64 as *const std::ffi::c_void,
                std::mem::size_of::<u64>(),
            ).map_err(|e| MLError::ModelError(format!("set drelu_bgrad BIAS_POINTER: {e:?}")))?;
        }

        let alpha = 1.0_f32;
        let beta = 0.0_f32;
        let lt_handle = self.handle.lt_handle_for(stream)?;
        let lt_ws = self.handle.lt_workspace_size;
        let (_, ws_ptr, _) = self.handle.lt_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 as *const std::ffi::c_void,
                cached.a_layout,
                dy as *const std::ffi::c_void,
                cached.b_layout,
                &beta as *const f32 as *const std::ffi::c_void,
                dx as *const std::ffi::c_void,
                cached.c_layout,
                dx 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,
                lt_ws,
                cu_stream,
            );
            if status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
                let e = cublaslt_result::CublasError(status);
                tracing::error!(
                    m = in_dim, n = batch, k = out_dim,
                    aux_ld = self.gemm_value_dx_drelu_bgrad_aux_ld,
                    ?e,
                    "cublasLtMatmul DRELU_BGRAD FAILED for value-FC dX"
                );
                return Err(MLError::ModelError(format!(
                    "cublasLtMatmul drelu_bgrad value-FC dX (m={in_dim},n={batch},k={out_dim}): {e:?}"
                )));
            }
        }

        Ok(true)
    }

    /// Apply the ReLU derivative mask in-place: `dx[i] *= (activation[i] > 0)`.
    ///
    /// Called after `backward_fc_layer()` for layers that used ReLU activation.
    /// The `activation` argument is the saved post-ReLU value from the forward
    /// pass (i.e. `save_h_*`).  For ReLU, `activation[i] > 0` iff the pre-ReLU
    /// value was positive, so we can use the post-ReLU value directly.
    ///
    /// Grid: `ceil(n / 256)`, Block: 256.
    pub fn relu_mask(
        &self,
        stream: &Arc<CudaStream>,
        dx: u64,
        activation: u64,
        n: usize,
    ) -> Result<(), MLError> {
        let n_i32 = n as i32;
        let blocks = ((n + 255) / 256) as u32;

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

        Ok(())
    }

    /// KAN spline gate backward: split upstream gradient `d_output` into `d_value`
    /// and `d_gate_pre`, and deterministically reduce coefficient/residual gradients.
    ///
    /// Two-kernel approach (no atomicAdd):
    ///   1. `kan_gate_backward` — per-element scatter: computes d_value, d_gate_pre,
    ///      and writes per-element basis contributions to intermediate buffers.
    ///   2. `kan_grad_reduce` — deterministic reduction: one thread per output
    ///      parameter accumulates from the intermediate buffers.
    ///
    /// ```text
    /// gate = clamp(sum_k coeff[k]*B_k(pre) + resid_w*pre, 0, 1)
    /// d_value[i]    = d_output[i] * gate
    /// d_gate_pre[i] = d_output[i] * value[i] * clamp_mask * (dB/dx + resid_w)
    /// d_coeff[k]    = sum_i { d_output[i] * value[i] * clamp_mask * B_k(pre) }
    /// d_resid_w     = sum_i { d_output[i] * value[i] * clamp_mask * pre }
    /// ```
    #[allow(clippy::too_many_arguments)]
    pub fn launch_kan_gate_backward(
        &self,
        stream: &Arc<CudaStream>,
        d_output: u64,        // [B, AH] upstream gradient from branch output backward
        gate_pre: u64,        // [B, AH] saved pre-activation gate input
        value: u64,           // [B, AH] saved value path output
        spline_coeff: u64,    // [AH, 8] current spline coefficients (from params_buf)
        residual_w: u64,      // [AH] current residual weights (from params_buf)
        d_gate_pre: u64,      // [B, AH] output: gradient for gate path (upstream to W_gate)
        d_value: u64,         // [B, AH] output: gradient for value path (upstream to W_bdf)
        d_spline_coeff: u64,  // [AH, 8] output: coeff gradients (in grad_buf)
        d_residual_w: u64,    // [AH] output: residual weight gradients (in grad_buf)
        n: usize,
        adv_h: usize,
    ) -> Result<(), MLError> {
        let n_i32 = n as i32;
        let adv_h_i32 = adv_h as i32;
        let blocks = ((n + 255) / 256) as u32;

        // Extract raw device pointers for the intermediate buffers.
        let d_coeff_pe_ptr = raw_f32_ptr(&self.kan_d_coeff_per_elem, stream);
        let d_span_ptr = {
            let (ptr, guard) = self.kan_d_coeff_span.device_ptr(stream);
            let _no_drop = ManuallyDrop::new(guard);
            ptr
        };
        let d_resid_pe_ptr = raw_f32_ptr(&self.kan_d_residual_per_elem, stream);

        // ── Stage 1: per-element scatter (no atomicAdd) ──
        unsafe {
            stream
                .launch_builder(&self.kan_gate_backward_kernel)
                .arg(&d_output)
                .arg(&gate_pre)
                .arg(&value)
                .arg(&spline_coeff)
                .arg(&residual_w)
                .arg(&d_gate_pre)
                .arg(&d_value)
                .arg(&d_coeff_pe_ptr)
                .arg(&d_span_ptr)
                .arg(&d_resid_pe_ptr)
                .arg(&n_i32)
                .arg(&adv_h_i32)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("kan_gate_backward_kernel: {e}")))?;
        }

        // ── Stage 2: 2-phase warp-parallel reduction ──
        // total_params = adv_h * 8 (coefficients) + adv_h (residual) = adv_h * 9
        let batch_size = n / adv_h;
        let total_params = adv_h * 9;
        let p1_blocks_x = ((batch_size + 255) / 256) as u32;
        let total_params_i32 = total_params as i32;
        let p1_blocks_x_i32 = p1_blocks_x as i32;
        let partials_ptr = raw_f32_ptr(&self.kan_reduce_partials, stream);

        unsafe {
            // Phase 1: partial block sums — grid=(ceil(B/256), total_params)
            stream
                .launch_builder(&self.kan_grad_reduce_p1_kernel)
                .arg(&d_coeff_pe_ptr)
                .arg(&d_span_ptr)
                .arg(&d_resid_pe_ptr)
                .arg(&partials_ptr)
                .arg(&n_i32)
                .arg(&adv_h_i32)
                .arg(&total_params_i32)
                .launch(LaunchConfig {
                    grid_dim: (p1_blocks_x, total_params as u32, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 256 * 4,
                })
                .map_err(|e| MLError::ModelError(format!("kan_grad_reduce_p1_kernel: {e}")))?;

            // Phase 2: final reduction — grid=(total_params)
            stream
                .launch_builder(&self.kan_grad_reduce_p2_kernel)
                .arg(&partials_ptr)
                .arg(&d_spline_coeff)
                .arg(&d_residual_w)
                .arg(&p1_blocks_x_i32)
                .arg(&adv_h_i32)
                .arg(&total_params_i32)
                .launch(LaunchConfig {
                    grid_dim: (total_params as u32, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("kan_grad_reduce_p2_kernel: {e}")))?;
        }

        Ok(())
    }

    // ══════════════════════════════════════════════════════════════════════════
    // Per-branch backward (Steps 1-4: dW/db only, no dX)
    // ══════════════════════════════════════════════════════════════════════════

    /// KAN spline gate backward using per-branch intermediate buffers.
    ///
    /// Identical to `launch_kan_gate_backward` but uses branch-specific
    /// `d_coeff_per_elem`, `d_coeff_span`, and `d_residual_per_elem` buffers,
    /// making it safe for concurrent execution on branch streams.
    #[allow(clippy::too_many_arguments)]
    fn launch_kan_gate_backward_branch(
        &self,
        stream: &Arc<CudaStream>,
        d: usize,             // branch index 0..3
        d_output: u64,        // [B, AH]
        gate_pre: u64,        // [B, AH]
        value: u64,           // [B, AH]
        spline_coeff: u64,    // [AH, 8]
        residual_w: u64,      // [AH]
        d_gate_pre: u64,      // [B, AH] output
        d_value: u64,         // [B, AH] output
        d_spline_coeff: u64,  // [AH, 8] output (in grad_buf)
        d_residual_w: u64,    // [AH] output (in grad_buf)
        n: usize,
        adv_h: usize,
    ) -> Result<(), MLError> {
        let n_i32 = n as i32;
        let adv_h_i32 = adv_h as i32;
        let blocks = ((n + 255) / 256) as u32;

        // Use per-branch intermediate buffers (no cross-branch contention).
        let d_coeff_pe_ptr = raw_f32_ptr(&self.branch_kan_d_coeff_per_elem[d], stream);
        let d_span_ptr = {
            let (ptr, guard) = self.branch_kan_d_coeff_span[d].device_ptr(stream);
            let _no_drop = ManuallyDrop::new(guard);
            ptr
        };
        let d_resid_pe_ptr = raw_f32_ptr(&self.branch_kan_d_residual_per_elem[d], stream);

        // ── Stage 1: per-element scatter (no atomicAdd) ──
        unsafe {
            stream
                .launch_builder(&self.kan_gate_backward_kernel)
                .arg(&d_output)
                .arg(&gate_pre)
                .arg(&value)
                .arg(&spline_coeff)
                .arg(&residual_w)
                .arg(&d_gate_pre)
                .arg(&d_value)
                .arg(&d_coeff_pe_ptr)
                .arg(&d_span_ptr)
                .arg(&d_resid_pe_ptr)
                .arg(&n_i32)
                .arg(&adv_h_i32)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("kan_gate_backward_kernel branch {d}: {e}")))?;
        }

        // ── Stage 2: 2-phase warp-parallel reduction ──
        let batch_size = n / adv_h;
        let total_params = adv_h * 9;
        let p1_blocks_x = ((batch_size + 255) / 256) as u32;
        let total_params_i32 = total_params as i32;
        let p1_blocks_x_i32 = p1_blocks_x as i32;
        let partials_ptr = raw_f32_ptr(&self.branch_kan_reduce_partials[d], stream);

        unsafe {
            // Phase 1: partial block sums — grid=(ceil(B/256), total_params)
            stream
                .launch_builder(&self.kan_grad_reduce_p1_kernel)
                .arg(&d_coeff_pe_ptr)
                .arg(&d_span_ptr)
                .arg(&d_resid_pe_ptr)
                .arg(&partials_ptr)
                .arg(&n_i32)
                .arg(&adv_h_i32)
                .arg(&total_params_i32)
                .launch(LaunchConfig {
                    grid_dim: (p1_blocks_x, total_params as u32, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 256 * 4,
                })
                .map_err(|e| MLError::ModelError(format!("kan_grad_reduce_p1_kernel branch {d}: {e}")))?;

            // Phase 2: final reduction — grid=(total_params)
            stream
                .launch_builder(&self.kan_grad_reduce_p2_kernel)
                .arg(&partials_ptr)
                .arg(&d_spline_coeff)
                .arg(&d_residual_w)
                .arg(&p1_blocks_x_i32)
                .arg(&adv_h_i32)
                .arg(&total_params_i32)
                .launch(LaunchConfig {
                    grid_dim: (total_params as u32, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("kan_grad_reduce_p2_kernel branch {d}: {e}")))?;
        }

        Ok(())
    }

    /// Backward Steps 1-4 for a single branch on a branch stream.
    ///
    /// Computes dW/db for the branch output layer, KAN spline backward, and
    /// value/gate path weight gradients. All writes go to non-overlapping
    /// regions of `grad_buf`, making this safe for parallel execution.
    ///
    /// Does NOT compute Step 5 (dX to d_h_s2 or d_concat). That runs on the
    /// main stream after all 4 branches join.
    #[allow(clippy::too_many_arguments)]
    fn backward_branch_dw(
        &self,
        stream: &Arc<CudaStream>,
        d: usize,
        ws_ptr: u64,
        // Per-branch data
        d_adv_logits_d: u64,   // [B, n_d]
        save_h_b_d: u64,      // [B, AH] (KAN gate output)
        scratch_d_h_b_d: u64, // [B, AH] scratch for d_h_bd
        glu_gate_pre_d: u64,  // [B, AH] saved gate pre-activation
        glu_value_d: u64,     // [B, AH] saved value path output
        // Weight pointers (read-only)
        w_out: u64,           // W_bdk_out [n_d, AH]
        w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
        // Grad buf offsets
        grad_buf_base: u64,
        goff_w_bout_d: u64,
        goff_b_bout_d: u64,
        goff_w_bfc_d: u64,
        goff_b_bfc_d: u64,
        goff_w_gate_d: u64,
        goff_b_gate_d: u64,
        goff_kan_coeff_d: u64,
        goff_kan_resid_d: u64,
        // FC input for dW (branches 1,2,3 may use concat buffers)
        fc_input: u64,
        fc_in_dim: usize,
        // Branch output dim
        n_d: usize,
        b: usize,
    ) -> Result<(), MLError> {
        // ── Step 1: Branch output layer (no activation — logit layer) ──
        // dW += dY^T @ X,  db += sum(dY),  dX = dY @ W^T → scratch_d_h_b[d]
        self.backward_fc_layer_ws(
            stream,
            d_adv_logits_d,
            save_h_b_d,
            w_out,
            grad_buf_base + goff_w_bout_d,
            grad_buf_base + goff_b_bout_d,
            scratch_d_h_b_d,
            n_d,
            self.adv_h,
            b,
            ws_ptr,
        )?;

        // Per-branch GLU scratch pointers
        let d_glu_value_ptr = raw_f32_ptr(&self.branch_d_glu_value[d], stream);
        let d_glu_gate_ptr = raw_f32_ptr(&self.branch_d_glu_gate[d], stream);

        // ── Step 2: KAN spline backward ──
        let kan_coeff_idx = 42 + d * 2;
        let kan_resid_idx = 42 + d * 2 + 1;
        self.launch_kan_gate_backward_branch(
            stream,
            d,
            scratch_d_h_b_d,
            glu_gate_pre_d,
            glu_value_d,
            w_ptrs[kan_coeff_idx],
            w_ptrs[kan_resid_idx],
            d_glu_gate_ptr,
            d_glu_value_ptr,
            grad_buf_base + goff_kan_coeff_d,
            grad_buf_base + goff_kan_resid_d,
            b * self.adv_h,
            self.adv_h,
        )?;

        // ── Step 3: Value path weight gradients ──
        self.launch_dw_only_ws(
            stream,
            d_glu_value_ptr,
            fc_input,
            grad_buf_base + goff_w_bfc_d,
            grad_buf_base + goff_b_bfc_d,
            self.adv_h,
            fc_in_dim,
            b,
            ws_ptr,
        )?;

        // ── Step 4: Gate path weight gradients ──
        self.launch_dw_only_ws(
            stream,
            d_glu_gate_ptr,
            fc_input,
            grad_buf_base + goff_w_gate_d,
            grad_buf_base + goff_b_gate_d,
            self.adv_h,
            fc_in_dim,
            b,
            ws_ptr,
        )?;

        Ok(())
    }

    /// Backward Step 5 for a single branch on the main stream.
    ///
    /// Computes dX (upstream gradient to VSN input / d_h_s2 / concat buffers).
    /// Must run on the main stream AFTER all branches finish Steps 1-4 so that
    /// d_h_s2 accumulation with beta is ordered correctly.
    #[allow(clippy::too_many_arguments)]
    fn backward_branch_dx(
        &self,
        stream: &Arc<CudaStream>,
        d: usize,
        w_fc: u64,
        w_gate: u64,
        scratch_d_h_s2: u64,
        mag_concat_ptr: u64,
        d_mag_concat_ptr: u64,
        ord_concat_ptr: u64,
        d_ord_concat_ptr: u64,
        urg_concat_ptr: u64,
        d_urg_concat_ptr: u64,
        // SP14 B.10 (2026-05-05): when non-zero, direction branch (d==0)
        // dX writes to this `[B, SH2+1]` buffer with K = SH2+1 (matching
        // the widened `w_b0fc [adv_h, SH2+1]` weight from B.8). When 0,
        // falls back to legacy direct-write into `scratch_d_h_s2 [B, SH2]`.
        d_dir_qaux_concat_ptr: u64,
        b: usize,
    ) -> Result<(), MLError> {
        // Read per-branch GLU scratch (produced by Steps 1-4 on branch stream,
        // now readable by main stream after join).
        let d_glu_value_ptr = raw_f32_ptr(&self.branch_d_glu_value[d], stream);
        let d_glu_gate_ptr = raw_f32_ptr(&self.branch_d_glu_gate[d], stream);

        if d == 0 && d_dir_qaux_concat_ptr != 0 {
            // SP14 B.10: direction Q-head with EGF wire — dX writes to
            // d_dir_qaux_concat [B, SH2+1]. K = SH2+1 because `w_b0fc` is
            // `[adv_h, SH2+1]` post-B.8 (last col = aux_softmax_diff weight).
            // After this returns, the caller scales col SH2 by ISV[393] then
            // accumulates first SH2 cols into scratch_d_h_s2 (with beta=0
            // because d==0 is the first branch's contribution to d_h_s2).
            self.launch_dx_only(
                stream,
                d_glu_value_ptr,
                w_fc,
                d_dir_qaux_concat_ptr,
                self.adv_h,
                self.shared_h2 + 1,
                b,
                0.0_f32,
            )?;
            self.launch_dx_only(
                stream,
                d_glu_gate_ptr,
                w_gate,
                d_dir_qaux_concat_ptr,
                self.adv_h,
                self.shared_h2 + 1,
                b,
                1.0_f32,
            )?;
        } else if d == 1 && mag_concat_ptr != 0 {
            // Magnitude (direction-conditioned): dX writes to d_mag_concat [B, SH2+branch_0_size].
            self.launch_dx_only(
                stream,
                d_glu_value_ptr,
                w_fc,
                d_mag_concat_ptr,
                self.adv_h,
                self.shared_h2 + self.branch_0_size,
                b,
                0.0_f32,
            )?;
            self.launch_dx_only(
                stream,
                d_glu_gate_ptr,
                w_gate,
                d_mag_concat_ptr,
                self.adv_h,
                self.shared_h2 + self.branch_0_size,
                b,
                1.0_f32,
            )?;
        } else if d == 2 && ord_concat_ptr != 0 {
            // Order: dX writes to d_ord_concat [B, SH2+3].
            self.launch_dx_only(
                stream,
                d_glu_value_ptr,
                w_fc,
                d_ord_concat_ptr,
                self.adv_h,
                self.shared_h2 + 3,
                b,
                0.0_f32,
            )?;
            self.launch_dx_only(
                stream,
                d_glu_gate_ptr,
                w_gate,
                d_ord_concat_ptr,
                self.adv_h,
                self.shared_h2 + 3,
                b,
                1.0_f32,
            )?;
        } else if d == 3 && urg_concat_ptr != 0 {
            // Urgency: dX writes to d_urg_concat [B, SH2+3].
            self.launch_dx_only(
                stream,
                d_glu_value_ptr,
                w_fc,
                d_urg_concat_ptr,
                self.adv_h,
                self.shared_h2 + 3,
                b,
                0.0_f32,
            )?;
            self.launch_dx_only(
                stream,
                d_glu_gate_ptr,
                w_gate,
                d_urg_concat_ptr,
                self.adv_h,
                self.shared_h2 + 3,
                b,
                1.0_f32,
            )?;
        } else {
            // Direction (d==0) or any branch with no concat: accumulate into d_h_s2 directly.
            let beta_s2 = if d == 0 { 0.0_f32 } else { 1.0_f32 };
            self.launch_dx_only(
                stream,
                d_glu_value_ptr,
                w_fc,
                scratch_d_h_s2,
                self.adv_h,
                self.shared_h2,
                b,
                beta_s2,
            )?;
            self.launch_dx_only(
                stream,
                d_glu_gate_ptr,
                w_gate,
                scratch_d_h_s2,
                self.adv_h,
                self.shared_h2,
                b,
                1.0_f32,
            )?;
        }

        Ok(())
    }

    /// Workspace-aware `backward_fc_layer`: dW + db + dX using explicit workspace.
    #[allow(clippy::too_many_arguments)]
    fn backward_fc_layer_ws(
        &self,
        stream: &Arc<CudaStream>,
        dy: u64, x: u64, w: u64, dw: u64, db: u64, dx: u64,
        out_dim: usize, in_dim: usize, batch: usize,
        ws_ptr: u64,
    ) -> Result<(), MLError> {
        // dW += dY^T @ X
        self.sgemm_f32_ws(
            stream,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            in_dim as i32, out_dim as i32, batch as i32,
            1.0, x, in_dim as i32,
            dy, out_dim as i32,
            1.0, dw, in_dim as i32,
            ws_ptr,
            "fc_dW_ws",
        )?;

        self.launch_bias_grad(stream, dy, db, out_dim, batch)?;

        // dX = dY @ W^T
        if dx != 0 {
            self.sgemm_f32_ws(
                stream,
                cublas_sys::cublasOperation_t::CUBLAS_OP_N,
                cublas_sys::cublasOperation_t::CUBLAS_OP_N,
                in_dim as i32, batch as i32, out_dim as i32,
                1.0, w, in_dim as i32,
                dy, out_dim as i32,
                0.0, dx, in_dim as i32,
                ws_ptr,
                "fc_dX_ws",
            )?;
        }

        Ok(())
    }

    /// Workspace-aware `launch_dw_only`: dW + db using explicit workspace.
    #[allow(clippy::too_many_arguments)]
    fn launch_dw_only_ws(
        &self,
        stream: &Arc<CudaStream>,
        dy: u64, x: u64, dw: u64, db: u64,
        out_dim: usize, in_dim: usize, batch: usize,
        ws_ptr: u64,
    ) -> Result<(), MLError> {
        self.sgemm_f32_ws(
            stream,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            in_dim as i32, out_dim as i32, batch as i32,
            1.0, x, in_dim as i32,
            dy, out_dim as i32,
            1.0, dw, in_dim as i32,
            ws_ptr,
            "dW_only_ws",
        )?;

        self.launch_bias_grad(stream, dy, db, out_dim, batch)
    }

    // ══════════════════════════════════════════════════════════════════════════
    // Full backward pass
    // ══════════════════════════════════════════════════════════════════════════

    /// Full backward pass through the branching dueling network via cuBLAS SGEMM.
    ///
    /// Consumes pre-computed `dL/d_logits` for each head emitted by
    /// `c51_loss_batched` (`c51_loss_kernel.cu`), and is invoked by
    /// `GpuDqnTrainer::launch_cublas_backward` during the fused training
    /// step.
    ///
    /// ## Backward order (reverse of forward)
    ///
    /// For each branch d in {0, 1, 2}:
    ///   1. Branch output layer:  `d_adv_logits_d` → `h_bd`
    ///   2. Branch FC layer:      `h_bd` gradient  → `h_s2` (accumulated)
    ///
    /// Value head:
    ///   3. Value output layer:  `d_value_logits` → `h_v`
    ///   4. Value FC layer:      `h_v` gradient   → `h_s2` (accumulated)
    ///
    /// Shared trunk (using accumulated `d_h_s2`):
    ///   5. Shared layer 2: `d_h_s2` → `h_s1`
    ///   6. Shared layer 1: `d_h_s1` → (states — skip, not trainable)
    ///
    /// ## Arguments
    ///
    /// - `d_value_logits`   : `[B, NA]` gradient from C51 loss for the value head
    /// - `d_adv_logits`     : four `[B, branch_k * NA]` gradients for each branch
    /// - `states`           : `[B, SD]` raw input states (for last shared layer dW)
    /// - `save_h_s1`        : `[B, SH1]` saved shared layer 1 ReLU output
    /// - `save_h_s2`        : `[B, SH2]` saved shared layer 2 ReLU output
    /// - `save_h_v`         : `[B, VH]` saved value FC ReLU output
    /// - `save_h_b`         : four `[B, AH]` saved GLU output (post-gating activation)
    /// - `w_ptrs`           : 50-element array of raw F32 device pointers into
    ///                        `params_buf` (online weights, read-only)
    /// - `grad_buf`         : flat gradient accumulator `[TOTAL_PARAMS]` — must be
    ///                        zeroed before this call
    /// - `scratch_dh`       : scratch buffers for layer-to-layer gradient passing
    /// - `glu_gate_pre`     : four `[B, AH]` saved pre-sigmoid gate activations
    /// - `glu_value`        : four `[B, AH]` saved value path outputs (pre-gating)
    /// - `scratch_d_glu_value`: `[B, AH]` scratch for GLU value gradient
    /// - `scratch_d_glu_gate` : `[B, AH]` scratch for GLU gate_pre gradient
    #[allow(clippy::too_many_arguments, unreachable_code, unused_variables)]
    pub fn backward_full(
        &self,
        stream: &Arc<CudaStream>,
        // dL/d_logits from C51 loss kernel
        d_value_logits: u64,          // [B, NA]
        d_adv_logits: &[u64; 4],      // [B, branch_k * NA] each
        // Saved forward activations
        states: u64,                  // [B, SD]
        save_h_s1: u64,               // [B, SH1]
        save_h_s2: u64,               // [B, SH2]
        save_h_v: u64,                // [B, VH]
        save_h_b: &[u64; 4],          // [B, AH] each — GLU output
        // Online weight pointers (read-only for W^T computation)
        w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
        // Flat gradient accumulator (must be zeroed by caller)
        grad_buf_base: u64,
        // Scratch buffers for inter-layer gradients (f32)
        scratch_d_h_s2: u64,          // [B, SH2] — f32, accumulated from value + all 4 branches
        scratch_d_h_s1: u64,          // [B, SH1] — f32
        scratch_d_h_v: u64,           // [B, VH]  — f32
        scratch_d_h_b: &[u64; 4],     // [B, AH] each — f32
        // #31 Bottleneck: optional f32 dX output for shared layer 1.
        // When non-zero, computes d_loss/d_bn_concat (upstream gradient for bottleneck).
        // When 0, dX is not computed (states not trainable — no bottleneck).
        s1_dx_output: u64,
        // Magnitude branch conditioning: saved forward concat and dX output
        // (direction-conditioned: SH2 + branch_0_size cols).
        mag_concat_ptr: u64,          // [B, SH2+branch_0_size] saved forward concat (for dW)
        d_mag_concat_ptr: u64,        // [B, SH2+branch_0_size] dX output for branch 1 (caller accumulates)
        // Order branch OFI conditioning: saved forward concat and dX output
        ord_concat_ptr: u64,          // [B, SH2+3] saved forward concat (for dW)
        d_ord_concat_ptr: u64,        // [B, SH2+3] dX output for branch 2 (caller accumulates)
        // Urgency branch OFI conditioning: saved forward concat and dX output
        urg_concat_ptr: u64,          // [B, SH2+3] saved forward concat (for dW)
        d_urg_concat_ptr: u64,        // [B, SH2+3] dX output for branch 3 (caller accumulates)
        // GLU saved forward activations for backward pass
        glu_gate_pre: &[u64; 4],      // [B, AH] each — saved pre-sigmoid gate activations
        glu_value: &[u64; 4],         // [B, AH] each — saved value path outputs
        // GLU backward scratch buffers (kept for backward compatibility; per-branch
        // versions in branch_d_glu_value/branch_d_glu_gate are used for parallel dispatch)
        _scratch_d_glu_value: u64,    // [B, AH] — superseded by per-branch scratch
        _scratch_d_glu_gate: u64,     // [B, AH] — superseded by per-branch scratch
        // Phase H Site 2 — value-FC dReLU bit-mask aux buffer written by the
        // online forward via `RELU_AUX_BIAS`. Consumed via `DRELU_BGRAD` on the
        // value-output dX GEMM (fuses the standalone `relu_mask` + `bias_grad`
        // reduce into one cublasLt call, saving 2 launches per backward). When
        // 0 OR when the cached DRELU_BGRAD descriptor is unavailable, the
        // backward path falls back to the standalone-kernel sequence.
        value_fc_relu_mask_ptr: u64,
        // SP14 Layer B Task B.10 (2026-05-05): direction-Q-head Earned
        // Gradient Flow wire. When non-zero, the direction branch (`d == 0`)
        // backward operates on the wider `[B, SH2 + 1]` concat (last col =
        // `aux_softmax_diff`), matching the SP14 B.8/B.9 forward consumer
        // pattern. dW reads `dir_qaux_concat_ptr` (saved forward concat,
        // re-built by the caller via `launch_sp14_dir_concat_qaux` right
        // before this `backward_full`), and dX writes to `d_dir_qaux_concat_ptr`
        // with `K = SH2 + 1`. The caller then accumulates the first SH2
        // columns into `scratch_d_h_s2` via `strided_accumulate`
        // (mirroring the magnitude branch pattern); column SH2 is dropped.
        // SP14 Layer C Phase C.1 (2026-05-08) deleted the
        // `sp14_scale_wire_col_kernel` α-gate that previously scaled
        // column SH2 — Coupling B is gone, the wire column flows
        // unscaled and is simply not propagated downstream.
        // When 0 (e.g. CudaSlice wrapper path), `d == 0` falls back to the
        // legacy direct-write path into `scratch_d_h_s2 [B, SH2]`.
        dir_qaux_concat_ptr: u64,     // [B, SH2+1] saved forward concat (for dW)
        d_dir_qaux_concat_ptr: u64,   // [B, SH2+1] dX output (caller scales col SH2 + accumulates first SH2 cols)
    ) -> Result<(), MLError> {
        // Plan 4 Task 2c.3c.4: trunk backward swap.
        //
        // After Task 2c.3a reshuffled the 4 legacy Linear→ReLU→Linear trunk
        // tensors into 13 GRN tensors at indices [0..13), running the legacy
        // backward GEMM sequence would write garbage into GRN-shaped slots.
        // Task 2c.3c.4 removes the legacy trunk-backward block from this
        // function entirely; the trainer is now responsible for invoking
        // `BatchedForward::encoder_backward_chain` (which mirrors the GRN
        // forward composition) AFTER `backward_full` returns.
        //
        // `backward_full` therefore writes:
        //   - branch FC + value-head gradients into `grad_buf_base`
        //   - `scratch_d_h_s2` populated with the value-head + branch dX,
        //     accumulated for consumption by `encoder_backward_chain`
        // The trainer's downstream block runs:
        //   - magnitude/order/urgency concat dX accumulation into d_h_s2
        //   - `encoder_backward_chain` (writes h_s2 + h_s1 GRN dW/dB/dgamma
        //     /dbeta and produces the bottleneck dX in `s1_dx_output`)
        //   - bottleneck backward (chain rule through tanh + dW)
        let b = self.batch_size;
        let na = self.num_atoms;

        // ── Precomputed grad_buf offsets (Plan 4 Task 2c.3a — GRN reshuffle) ──
        //
        // Layout (post-2c.3a, 95 tensors total — 13 GRN trunk tensors at
        // indices [0..13), then everything that used to start at index 4
        // shifts +9 to start at index 13):
        //   [0..7)   h_s1 GRN: w_a, b_a, w_b, b_b, w_residual, gamma, beta
        //   [7..13)  h_s2 GRN: w_a, b_a, w_b, b_b, gamma, beta (no Linear_residual)
        //   [13..17) value head (w_v1, b_v1, w_v2, b_v2)
        //   [17..33) 4 advantage branches × {fc, fc_bias, out, out_bias}
        //   [33..35) bottleneck
        //   [35..43) VSN
        //   [43..51) GLU gates
        //   [51..59) KAN
        //   etc.
        //
        // Trunk-tensor offsets (indices [0..13)) are written by
        // `BatchedForward::encoder_backward_chain` invoked by the trainer.
        // The branches and value head still write here; the value-head dX
        // accumulates into `scratch_d_h_s2` for the trainer's chain call.
        let goff_w_v1:    u64 = padded_byte_offset(&self.param_sizes, 13);
        let goff_b_v1:    u64 = padded_byte_offset(&self.param_sizes, 14);
        let goff_w_v2:    u64 = padded_byte_offset(&self.param_sizes, 15);
        let goff_b_v2:    u64 = padded_byte_offset(&self.param_sizes, 16);
        let goff_w_b0fc:  u64 = padded_byte_offset(&self.param_sizes, 17);
        let goff_b_b0fc:  u64 = padded_byte_offset(&self.param_sizes, 18);
        let goff_w_b0out: u64 = padded_byte_offset(&self.param_sizes, 19);
        let goff_b_b0out: u64 = padded_byte_offset(&self.param_sizes, 20);
        let goff_w_b1fc:  u64 = padded_byte_offset(&self.param_sizes, 21);
        let goff_b_b1fc:  u64 = padded_byte_offset(&self.param_sizes, 22);
        let goff_w_b1out: u64 = padded_byte_offset(&self.param_sizes, 23);
        let goff_b_b1out: u64 = padded_byte_offset(&self.param_sizes, 24);
        let goff_w_b2fc:  u64 = padded_byte_offset(&self.param_sizes, 25);
        let goff_b_b2fc:  u64 = padded_byte_offset(&self.param_sizes, 26);
        let goff_w_b2out: u64 = padded_byte_offset(&self.param_sizes, 27);
        let goff_b_b2out: u64 = padded_byte_offset(&self.param_sizes, 28);
        let goff_w_b3fc:  u64 = padded_byte_offset(&self.param_sizes, 29);
        let goff_b_b3fc:  u64 = padded_byte_offset(&self.param_sizes, 30);
        let goff_w_b3out: u64 = padded_byte_offset(&self.param_sizes, 31);
        let goff_b_b3out: u64 = padded_byte_offset(&self.param_sizes, 32);

        let goff_w_bfc  = [goff_w_b0fc,  goff_w_b1fc,  goff_w_b2fc,  goff_w_b3fc];
        let goff_b_bfc  = [goff_b_b0fc,  goff_b_b1fc,  goff_b_b2fc,  goff_b_b3fc];
        let goff_w_bout = [goff_w_b0out, goff_w_b1out, goff_w_b2out, goff_w_b3out];
        let goff_b_bout = [goff_b_b0out, goff_b_b1out, goff_b_b2out, goff_b_b3out];

        let branch_n = [
            self.branch_0_size * na,
            self.branch_1_size * na,
            self.branch_2_size * na,
            self.branch_3_size * na,
        ];

        // w_ptrs indices (Plan 4 Task 2c.3a — GRN reshuffle, all post-h_s2 +9):
        //   [0..7)   h_s1 GRN tensors (LeakyReLU/ReLU consumers panic-gated)
        //   [7..13)  h_s2 GRN tensors
        //   [13..17) value head (w_v1, b_v1, w_v2, b_v2)
        //   [17..21) branch 0 (direction)
        //   [21..25) branch 1 (magnitude)
        //   [25..29) branch 2 (order)
        //   [29..33) branch 3 (urgency)
        //   [43..51) GLU gates
        //   [51..59) KAN coeff/resid pairs
        let w_bfc_idx  = [17_usize, 21, 25, 29];
        let w_bout_idx = [19_usize, 23, 27, 31];
        // GLU gate weight indices (post-2c.3a): w_gate_d at 43+d*2, b_gate_d at 43+d*2+1
        let w_gate_idx = [43_usize, 45, 47, 49];

        // GLU gate gradient offsets (use padded_byte_offset for correctness)
        let goff_w_gate = [
            padded_byte_offset(&self.param_sizes, 43),
            padded_byte_offset(&self.param_sizes, 45),
            padded_byte_offset(&self.param_sizes, 47),
            padded_byte_offset(&self.param_sizes, 49),
        ];
        let goff_b_gate = [
            padded_byte_offset(&self.param_sizes, 44),
            padded_byte_offset(&self.param_sizes, 46),
            padded_byte_offset(&self.param_sizes, 48),
            padded_byte_offset(&self.param_sizes, 50),
        ];

        // KAN spline coefficient gradient offsets (in grad_buf)
        let goff_kan_coeff = [
            padded_byte_offset(&self.param_sizes, 51),
            padded_byte_offset(&self.param_sizes, 53),
            padded_byte_offset(&self.param_sizes, 55),
            padded_byte_offset(&self.param_sizes, 57),
        ];
        let goff_kan_resid = [
            padded_byte_offset(&self.param_sizes, 52),
            padded_byte_offset(&self.param_sizes, 54),
            padded_byte_offset(&self.param_sizes, 56),
            padded_byte_offset(&self.param_sizes, 58),
        ];

        // ══════════════════════════════════════════════════════════════════
        // BRANCHES: parallel dW/db (Steps 1-4), sequential dX (Step 5)
        //
        // Steps 1-4 write to non-overlapping regions of grad_buf and
        // per-branch scratch buffers, so they run safely in parallel on
        // branch_streams[0..3]. Step 5 (dX) may accumulate into the
        // shared scratch_d_h_s2 with beta ordering, so it runs
        // sequentially on the main stream after all branches join.
        // ══════════════════════════════════════════════════════════════════

        // ── Fork: record trunk completion, dispatch Steps 1-4 on branch streams ──
        self.trunk_done_event.record(stream)
            .map_err(|e| MLError::DeviceError(format!("bw trunk_done_event record: {e}")))?;

        for d in 0..4 {
            let bs = &self.branch_streams[d];
            bs.wait(&self.trunk_done_event)
                .map_err(|e| MLError::DeviceError(format!("bw branch stream {d} wait trunk: {e}")))?;

            let n_d = branch_n[d];
            let w_out = w_ptrs[w_bout_idx[d]];

            // FC input for dW (each branch may use a wider concat buffer).
            // Direction (d==0): SP14 B.10 EGF wire — `[h_s2 ; aux_softmax_diff]`
            //   when `dir_qaux_concat_ptr != 0` (SH2 + 1 cols).
            // Magnitude (d==1): direction-conditioned (SH2 + branch_0_size).
            // Order/urgency (d==2,3): OFI-conditioned (SH2 + 3).
            // Otherwise: raw `save_h_s2` (SH2 cols), the legacy path.
            let (fc_input, fc_in_dim) = if d == 0 && dir_qaux_concat_ptr != 0 {
                (dir_qaux_concat_ptr, self.shared_h2 + 1)
            } else if d == 1 && mag_concat_ptr != 0 {
                (mag_concat_ptr, self.shared_h2 + self.branch_0_size)
            } else if d == 2 && ord_concat_ptr != 0 {
                (ord_concat_ptr, self.shared_h2 + 3)
            } else if d == 3 && urg_concat_ptr != 0 {
                (urg_concat_ptr, self.shared_h2 + 3)
            } else {
                (save_h_s2, self.shared_h2)
            };

            self.backward_branch_dw(
                bs,
                d,
                self.branch_workspace_ptrs[d],
                d_adv_logits[d],
                save_h_b[d],
                scratch_d_h_b[d],
                glu_gate_pre[d],
                glu_value[d],
                w_out,
                w_ptrs,
                grad_buf_base,
                goff_w_bout[d],
                goff_b_bout[d],
                goff_w_bfc[d],
                goff_b_bfc[d],
                goff_w_gate[d],
                goff_b_gate[d],
                goff_kan_coeff[d],
                goff_kan_resid[d],
                fc_input,
                fc_in_dim,
                n_d,
                b,
            )?;

            self.branch_done_events[d].record(bs)
                .map_err(|e| MLError::DeviceError(format!("bw branch_done_event {d} record: {e}")))?;
        }

        // ── Join: main stream waits for all 4 branches to finish Steps 1-4 ──
        for d in 0..4 {
            stream.wait(&self.branch_done_events[d])
                .map_err(|e| MLError::DeviceError(format!("bw main stream wait branch {d}: {e}")))?;
        }

        // ── Step 5: dX accumulation runs sequentially on main stream ──
        // d_h_s2 requires ordered beta accumulation (d==0 overwrites, d>0 accumulates).
        // Concat branches (mag/ord/urg) write to separate buffers, so ordering
        // among them does not matter — but they must complete before the trunk
        // backward reads d_h_s2.
        for d in 0..4 {
            let w_fc  = w_ptrs[w_bfc_idx[d]];
            let w_gate = w_ptrs[w_gate_idx[d]];

            self.backward_branch_dx(
                stream,
                d,
                w_fc,
                w_gate,
                scratch_d_h_s2,
                mag_concat_ptr,
                d_mag_concat_ptr,
                ord_concat_ptr,
                d_ord_concat_ptr,
                urg_concat_ptr,
                d_urg_concat_ptr,
                // SP14 B.10: route direction (d==0) dX to the wire-aware
                // [B, SH2+1] buffer when the EGF wire is enabled.
                d_dir_qaux_concat_ptr,
                b,
            )?;
        }

        // ══════════════════════════════════════════════════════════════════
        // VALUE HEAD
        // ══════════════════════════════════════════════════════════════════

        // ── Value output + value-FC backward (DRELU_BGRAD-fused when available) ──
        //
        // Phase H Site 2: when the cached `EPILOGUE_DRELU_BGRAD` descriptor
        // is present AND the forward wrote a valid mask, fuse the standalone
        // `relu_mask` (on scratch_d_h_v) and the value-FC bias-gradient reduce
        // into the value-output dX GEMM. Saves 2 kernel launches per backward
        // pass × 2 forward passes (target + online) per training step.
        //
        // After the fused path:
        //   - scratch_d_h_v contains the dReLU-gated upstream gradient
        //   - grad_buf[goff_b_v1] holds the value-FC bias gradient
        //   - the subsequent value-FC dW uses `launch_dw_only_no_bias`
        //     (skips the now-redundant bias-grad reduce)
        //
        // Fallback (when `value_fc_relu_mask_ptr == 0` OR DRELU_BGRAD unavailable):
        //   regular backward_fc_layer + relu_mask + launch_dw_only sequence.
        let drelu_bgrad_used = if value_fc_relu_mask_ptr != 0 {
            self.backward_fc_layer_drelu_bgrad(
                stream,
                d_value_logits,
                save_h_v,
                w_ptrs[15],                // W_v2 [NA, VH]
                grad_buf_base + goff_w_v2,
                grad_buf_base + goff_b_v2,
                scratch_d_h_v,             // dX [B, VH] — gated by drelu in-place
                value_fc_relu_mask_ptr,    // dReLU mask aux from forward
                grad_buf_base + goff_b_v1, // value-FC bias gradient (next layer)
                na,
                self.value_h,
                b,
            )?
        } else {
            false
        };

        if !drelu_bgrad_used {
            // Slow path: standalone relu_mask + bias_grad reduce.
            // Plan 4 Task 2c.3a: W_v2 index +9 (was 6, now 15).
            self.backward_fc_layer(
                stream,
                d_value_logits,
                save_h_v,
                w_ptrs[15],                // W_v2 [NA, VH]
                grad_buf_base + goff_w_v2,
                grad_buf_base + goff_b_v2,
                scratch_d_h_v,             // dX [B, VH]
                na,
                self.value_h,
                b,
            )?;

            // ── Value FC layer (ReLU) ─────────────────────────────────────
            // ReLU mask on f32 dX_v
            self.relu_mask(stream, scratch_d_h_v, save_h_v, b * self.value_h)?;

            // dW for value FC: dW[VH, SH2] += dY^T @ h_s2 (with bias-grad reduce)
            self.launch_dw_only(
                stream,
                scratch_d_h_v,             // dY [B, VH] — f32
                save_h_s2,
                grad_buf_base + goff_w_v1,
                grad_buf_base + goff_b_v1,
                self.value_h,
                self.shared_h2,
                b,
            )?;
        } else {
            // Fast path: DRELU_BGRAD already wrote scratch_d_h_v (gated) and
            // grad_buf[goff_b_v1] (bias gradient). Only dW remains for value-FC.
            self.launch_dw_only_no_bias(
                stream,
                scratch_d_h_v,             // dY [B, VH] — f32 (already drelu-gated)
                save_h_s2,
                grad_buf_base + goff_w_v1,
                self.value_h,
                self.shared_h2,
                b,
            )?;
        }

        // dX for value FC: accumulated into f32 scratch_d_h_s2 (beta=1.0 — branches already wrote)
        // Plan 4 Task 2c.3a: W_v1 index +9 (was 4, now 13).
        self.launch_dx_only(
            stream,
            scratch_d_h_v,             // dY [B, VH] — f32
            w_ptrs[13],                // W_v1 [VH, SH2]
            scratch_d_h_s2,
            self.value_h,
            self.shared_h2,
            b,
            1.0_f32,                   // accumulate on top of branch contributions
        )?;

        // ══════════════════════════════════════════════════════════════════
        // SHARED TRUNK (Plan 4 Task 2c.3c.4 — moved out)
        // ══════════════════════════════════════════════════════════════════
        //
        // Trunk backward (ReLU + Linear×2) was the legacy chain matching the
        // pre-GRN trunk forward. Task 2c.3a migrated the trunk to GRN and
        // 2c.3c.4 moved the backward chain into
        // `BatchedForward::encoder_backward_chain` so it owns the cuBLAS +
        // GRN-kernel orchestration symmetrically with the forward path. The
        // trainer invokes `encoder_backward_chain` immediately after this
        // function returns, passing in `scratch_d_h_s2` as the chain's
        // entry-point gradient.
        //
        // Inputs to this function still include `save_h_s1`, `scratch_d_h_s1`,
        // `states`, and `s1_dx_output` because the auxiliary IQN/ensemble
        // paths (which pre-existed `encoder_backward_chain`) still rely on
        // the same signature shape — the trainer threads them straight
        // through to `encoder_backward_chain`. This function intentionally
        // does not consume them; their purpose is documentation of the
        // contract between branch-decoder backward and the trunk encoder
        // backward that runs immediately afterwards.
        let _ = (save_h_s1, scratch_d_h_s1, states, s1_dx_output);
        let _ = goff_b_b3out;     // suppress dead_code for last computed offset

        Ok(())
    }

    // ══════════════════════════════════════════════════════════════════════════
    // Private SGEMM helpers
    // ══════════════════════════════════════════════════════════════════════════

    /// Compute only weight gradient + bias gradient (no upstream dX).
    ///
    /// Used for branch FC layers where the upstream dX is computed separately
    /// via `launch_dx_only` so multiple branches can accumulate into the
    /// shared `d_h_s2` scratch buffer with controllable beta.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn launch_dw_only(
        &self,
        stream: &Arc<CudaStream>,
        dy: u64,
        x: u64,
        dw: u64,
        db: u64,
        out_dim: usize,
        in_dim: usize,
        batch: usize,
    ) -> Result<(), MLError> {
        self.sgemm_f32(
            stream,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            in_dim as i32, out_dim as i32, batch as i32,
            1.0, x, in_dim as i32,
            dy, out_dim as i32,
            1.0, dw, in_dim as i32,
            "dW_only",
        )?;

        self.launch_bias_grad(stream, dy, db, out_dim, batch)
    }

    /// Compute ONLY weight gradient (no bias gradient, no upstream dX).
    ///
    /// Variant of `launch_dw_only` for FC layers that have NO bias term — e.g.
    /// the GRN block's `Linear_residual` projection (Plan 4 Task 2c.3c.2).
    /// `launch_dw_only` would segfault here because the 2-phase bias-grad
    /// kernel dereferences a NULL `db` pointer; this method skips that launch
    /// entirely and only runs the cuBLAS dW GEMM.
    ///
    /// Same convention as `launch_dw_only`:
    ///   `dW[out, in] += dY^T[out, B] @ X[B, in]`
    /// using the row-major → col-major dual:
    ///   `sgemm(N, T, in_dim, out_dim, batch, 1.0, X, in_dim, dY, out_dim, 1.0, dW, in_dim)`.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn launch_dw_only_no_bias(
        &self,
        stream: &Arc<CudaStream>,
        dy: u64,
        x: u64,
        dw: u64,
        out_dim: usize,
        in_dim: usize,
        batch: usize,
    ) -> Result<(), MLError> {
        self.sgemm_f32(
            stream,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            in_dim as i32, out_dim as i32, batch as i32,
            1.0, x, in_dim as i32,
            dy, out_dim as i32,
            1.0, dw, in_dim as i32,
            "dW_only_no_bias",
        )
    }

    /// Compute ONLY weight gradient (no bias gradient, no upstream dX) with a
    /// custom leading dimension for X (Plan 4 Task 2c.3c.4).
    ///
    /// Matches `launch_dw_only_no_bias` but threads `x_lda` into cuBLAS — needed
    /// for the GRN h_s1 `Linear_residual` backward, whose input is the padded
    /// `states_buf` (row stride = `pad128(state_dim)`) or the bottleneck
    /// `bn_concat` (row stride = `s1_input_dim`). Bias gradient is skipped
    /// because `Linear_residual` is bias-less; mixing this into `launch_dw_only`
    /// would dereference a NULL `db` pointer in the bias-grad kernel.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn launch_dw_only_no_bias_lda(
        &self,
        stream: &Arc<CudaStream>,
        dy: u64,
        x: u64,
        dw: u64,
        out_dim: usize,
        in_dim: usize,
        x_lda: usize,
        batch: usize,
    ) -> Result<(), MLError> {
        self.sgemm_f32(
            stream,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            in_dim as i32, out_dim as i32, batch as i32,
            1.0, x, x_lda as i32,
            dy, out_dim as i32,
            1.0, dw, in_dim as i32,
            "dW_only_no_bias_lda",
        )
    }

    /// Compute only upstream gradient dX = dY @ W^T, with configurable beta.
    ///
    /// `beta=0.0` overwrites dX; `beta=1.0` accumulates into dX.  Used to merge
    /// contributions from multiple branches into the shared `d_h_s2` buffer.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn launch_dx_only(
        &self,
        stream: &Arc<CudaStream>,
        dy: u64,
        w: u64,
        dx: u64,
        out_dim: usize,
        in_dim: usize,
        batch: usize,
        beta: f32,
    ) -> Result<(), MLError> {
        self.sgemm_f32(
            stream,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            in_dim as i32, batch as i32, out_dim as i32,
            1.0, w, in_dim as i32,
            dy, out_dim as i32,
            beta, dx, in_dim as i32,
            "dX_only",
        )
    }

    /// Compute only upstream gradient dX = dY @ W^T with a custom leading
    /// dimension for X (Plan 4 Task 1B-iv).
    ///
    /// Mirrors `launch_dx_only` but threads `x_lda` into cuBLAS — needed for
    /// the bottleneck Linear backward when its dX must be written into a
    /// slice of a wider buffer (`d_vsn_gated_state[B, state_dim_padded]`'s
    /// market slice `[0..market_dim]` with row stride `state_dim_padded`).
    /// `beta=0.0` overwrites the destination slice; the surrounding
    /// `vsn_d_gated_state_portfolio_pad_kernel` zero-initialises the market
    /// slice so callers can rely on `dx` being fully defined regardless of
    /// kernel/GEMM ordering.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn launch_dx_only_lda(
        &self,
        stream: &Arc<CudaStream>,
        dy: u64,
        w: u64,
        dx: u64,
        out_dim: usize,
        in_dim: usize,
        x_lda: usize,
        batch: usize,
        beta: f32,
    ) -> Result<(), MLError> {
        self.sgemm_f32(
            stream,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            in_dim as i32, batch as i32, out_dim as i32,
            1.0, w, in_dim as i32,
            dy, out_dim as i32,
            beta, dx, x_lda as i32,
            "dX_only_lda",
        )
    }

    /// Launch 2-phase bias gradient reduction: db[j] += sum_b(dY[b * out_dim + j]).
    ///
    /// Phase 1: block-level shared-memory reduce into `bias_grad_partials_buf`.
    /// Phase 2: final reduce across block partials + accumulate into db.
    fn launch_bias_grad(
        &self,
        stream: &Arc<CudaStream>,
        dy: u64,
        db: u64,
        out_dim: usize,
        batch: usize,
    ) -> Result<(), MLError> {
        let out_dim_i32 = out_dim as i32;
        let batch_i32 = batch as i32;
        let num_blocks = self.bias_grad_num_blocks;
        let partials_ptr = raw_f32_ptr(&self.bias_grad_partials_buf, stream);

        // Phase 1: block-level reduce
        unsafe {
            stream
                .launch_builder(&self.bias_grad_p1_kernel)
                .arg(&dy)
                .arg(&partials_ptr)
                .arg(&out_dim_i32)
                .arg(&batch_i32)
                .launch(LaunchConfig {
                    grid_dim: (num_blocks, out_dim as u32, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 256 * 4,
                })
                .map_err(|e| MLError::ModelError(format!("bias_grad_reduce_f32_p1: {e}")))?;
        }

        // Phase 2: final reduce + accumulate
        let num_blocks_i32 = num_blocks as i32;
        let p2_blocks = ((out_dim + 255) / 256) as u32;
        unsafe {
            stream
                .launch_builder(&self.bias_grad_p2_kernel)
                .arg(&partials_ptr)
                .arg(&db)
                .arg(&out_dim_i32)
                .arg(&num_blocks_i32)
                .launch(LaunchConfig {
                    grid_dim: (p2_blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("bias_grad_reduce_f32_p2: {e}")))?;
        }

        Ok(())
    }
}

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

/// Create a `CachedBwdGemmDesc` for a backward GEMM shape with given transpose flags.
///
/// Runs the algorithm heuristic ONCE at init time with `requestedAlgoCount=3`
/// for better algorithm selection.
#[allow(clippy::too_many_arguments)]
fn create_cached_bwd_gemm_desc(
    lt_handle: cublaslt_sys::cublasLtHandle_t,
    transa_i32: i32,
    transb_i32: i32,
    m: i32, n: i32, k: i32,
    lda: i32, ldb: i32, ldc: i32,
    ws_size: usize,
) -> Result<CachedBwdGemmDesc, 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 bw MatmulDescCreate (m={m},n={n},k={k}): {e:?}")))?;

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

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

        // Matrix layouts — match the logic in sgemm_f32_uncached.
        let transa_is_n = transa_i32 == 0; // CUBLAS_OP_N
        let transb_is_n = transb_i32 == 0;

        let a_layout = if transa_is_n {
            cublaslt_result::create_matrix_layout(f32_type, m as u64, k as u64, lda as i64)
        } else {
            cublaslt_result::create_matrix_layout(f32_type, k as u64, m as u64, lda as i64)
        }.map_err(|e| MLError::ModelError(format!("cached bw A layout (m={m},n={n},k={k}): {e:?}")))?;

        let b_layout = if transb_is_n {
            cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, ldb as i64)
        } else {
            cublaslt_result::create_matrix_layout(f32_type, n as u64, k as u64, ldb as i64)
        }.map_err(|e| MLError::ModelError(format!("cached bw B layout (m={m},n={n},k={k}): {e:?}")))?;

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

        let d_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, ldc as i64)
            .map_err(|e| MLError::ModelError(format!("cached bw D layout (m={m},n={n},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. The original graph-capture safety
        // concern (split-K/stream-K) is still handled — AlgoCheck validates
        // the selected algo against the current descriptor + layouts.
        let shape = super::cublas_algo_deterministic::ShapeKey::new(
            transa_i32, transb_i32, m, n, k, lda, ldb, ldc, 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 bw deterministic algo (transa={transa_i32},transb={transb_i32},m={m},n={n},k={k}): {e}"
                )));
            }
        };

        let algo = heuristic.algo;

        tracing::info!(
            transa = transa_i32, transb = transb_i32,
            m = m, n = n, k = k, lda = lda, ldb = ldb, ldc = ldc,
            ws_needed = heuristic.workspaceSize,
            "cached bwd GEMM desc created (deterministic algorithm)"
        );

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

/// Phase H Site 2 — create a cached backward GEMM descriptor with
/// `EPILOGUE_DRELU_BGRAD`.
///
/// Same shape contract as `create_cached_bwd_gemm_desc` (TRANSA/TRANSB and
/// M/N/K/lda/ldb/ldc passed in), but additionally stamps:
///   - EPILOGUE = DRELU_BGRAD
///   - BIAS_DATA_TYPE = F32 (for the bgrad output)
///   - EPILOGUE_AUX_LD = `aux_ld_elements` (matches forward's RELU_AUX_BIAS)
///
/// AUX and BIAS pointers are set per-call via `set_matmul_desc_attribute` —
/// matches the pattern used by `sgemm_f32_fused_relu_aux_bias` on the
/// forward side.
#[allow(clippy::too_many_arguments)]
fn create_cached_bwd_gemm_desc_drelu_bgrad(
    lt_handle: cublaslt_sys::cublasLtHandle_t,
    transa_i32: i32,
    transb_i32: i32,
    m: i32, n: i32, k: i32,
    lda: i32, ldb: i32, ldc: i32,
    aux_ld_elements: usize,
    ws_size: usize,
) -> Result<CachedBwdGemmDesc, 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!("drelu_bgrad MatmulDescCreate (m={m},n={n},k={k}): {e:?}")))?;

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

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

        // EPILOGUE = DRELU_BGRAD: gate output by mask + reduce along N for bgrad.
        let epilogue: i32 = cublaslt_sys::cublasLtEpilogue_t::CUBLASLT_EPILOGUE_DRELU_BGRAD 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!("drelu_bgrad set EPILOGUE: {e:?}")))?;

        // bgrad output 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!("drelu_bgrad set BIAS_DATA_TYPE: {e:?}")))?;

        // AUX_LD (in elements/bits) — matches forward's RELU_AUX_BIAS aux_ld.
        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!("drelu_bgrad set AUX_LD: {e:?}")))?;

        // Matrix layouts — same convention as create_cached_bwd_gemm_desc.
        let transa_is_n = transa_i32 == 0;
        let transb_is_n = transb_i32 == 0;

        let a_layout = if transa_is_n {
            cublaslt_result::create_matrix_layout(f32_type, m as u64, k as u64, lda as i64)
        } else {
            cublaslt_result::create_matrix_layout(f32_type, k as u64, m as u64, lda as i64)
        }.map_err(|e| MLError::ModelError(format!("drelu_bgrad A layout: {e:?}")))?;

        let b_layout = if transb_is_n {
            cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, ldb as i64)
        } else {
            cublaslt_result::create_matrix_layout(f32_type, n as u64, k as u64, ldb as i64)
        }.map_err(|e| MLError::ModelError(format!("drelu_bgrad B layout: {e:?}")))?;

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

        // Deterministic algo selection — keys on the new epilogue value.
        let shape = super::cublas_algo_deterministic::ShapeKey::with_epilogue(
            transa_i32, transb_i32, m, n, k, lda, ldb, ldc,
            cublaslt_sys::cublasLtEpilogue_t::CUBLASLT_EPILOGUE_DRELU_BGRAD,
            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!(
                    "drelu_bgrad deterministic algo (m={m},n={n},k={k},lda={lda},ldb={ldb},ldc={ldc}): {e}"
                )));
            }
        };
        let algo = heuristic.algo;

        tracing::info!(
            transa = transa_i32, transb = transb_i32,
            m = m, n = n, k = k, lda = lda, ldb = ldb, ldc = ldc,
            aux_ld = aux_ld_elements,
            ws_needed = heuristic.workspaceSize,
            "cached bwd GEMM+DRELU_BGRAD desc created (deterministic algorithm)"
        );

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

// ── Kernel compilation ────────────────────────────────────────────────────────

/// Compile `relu_mask_kernel` and `bias_grad_reduce_f32_p1/p2` from precompiled cubin.
///
/// - `relu_mask_kernel`: 1 thread per element, gating dx by saved activation.
/// - `bias_grad_reduce_f32_p1`: phase 1 — block-level shared-memory reduce over batch.
/// - `bias_grad_reduce_f32_p2`: phase 2 — final reduce across block partials + accumulate.
///
/// Precompiled backward kernels cubin (build.rs — ZERO runtime nvcc).
static BACKWARD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/backward_kernels.cubin"));

fn compile_backward_kernels(
    stream: &Arc<CudaStream>,
) -> Result<(CudaFunction, CudaFunction, CudaFunction), MLError> {
    let context = stream.context();
    let module = context
        .load_cubin(BACKWARD_CUBIN.to_vec())
        .map_err(|e| MLError::ModelError(format!("backward cubin load: {e}")))?;

    let relu_mask = module
        .load_function("relu_mask_kernel")
        .map_err(|e| MLError::ModelError(format!("relu_mask_kernel load: {e}")))?;
    let bias_grad_p1 = module
        .load_function("bias_grad_reduce_f32_p1")
        .map_err(|e| MLError::ModelError(format!("bias_grad_reduce_f32_p1 load: {e}")))?;
    let bias_grad_p2 = module
        .load_function("bias_grad_reduce_f32_p2")
        .map_err(|e| MLError::ModelError(format!("bias_grad_reduce_f32_p2 load: {e}")))?;

    Ok((relu_mask, bias_grad_p1, bias_grad_p2))
}

// ── Raw device pointer helpers ────────────────────────────────────────────────
//
// Mirrors the helpers in batched_forward.rs.  Declared here so this module
// can be compiled independently without `pub use`-ing the forward module's
// private functions.

/// Extract raw device pointer from an f32 CudaSlice (read-only).
pub(crate) 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
}

// ── Scratch buffer allocation helper ─────────────────────────────────────────

/// Allocate zero-initialised scratch buffers for the cuBLAS backward pass.
///
/// Returns `(d_h_s2, d_h_s1, d_h_v, d_h_b0, d_h_b1, d_h_b2, d_h_b3)`.
/// All f32 — the entire training pipeline is f32.
pub fn alloc_backward_scratch(
    stream: &Arc<CudaStream>,
    config: &GpuDqnTrainConfig,
) -> Result<(
    CudaSlice<f32>,         // d_h_s2 [B, SH2]
    CudaSlice<f32>,         // d_h_s1 [B, SH1]
    CudaSlice<f32>,         // d_h_v  [B, VH]
    CudaSlice<f32>,         // d_h_b0 [B, AH]
    CudaSlice<f32>,         // d_h_b1 [B, AH]
    CudaSlice<f32>,         // d_h_b2 [B, AH]
    CudaSlice<f32>,         // d_h_b3 [B, AH]
    CudaSlice<f32>,         // d_glu_value [B, AH] — GLU value path gradient scratch
    CudaSlice<f32>,         // d_glu_gate  [B, AH] — GLU gate_pre gradient scratch
), MLError> {
    let b = config.batch_size;
    let alloc_f32 = |n: usize| -> Result<CudaSlice<f32>, MLError> {
        stream.alloc_zeros::<f32>(n)
            .map_err(|e| MLError::ModelError(format!("backward scratch f32 alloc [{n}]: {e}")))
    };

    Ok((
        alloc_f32(b * config.shared_h2)?,
        alloc_f32(b * config.shared_h1)?,
        alloc_f32(b * config.value_h)?,
        alloc_f32(b * config.adv_h)?,
        alloc_f32(b * config.adv_h)?,
        alloc_f32(b * config.adv_h)?,
        alloc_f32(b * config.adv_h)?,
        alloc_f32(b * config.adv_h)?,  // GLU d_value scratch
        alloc_f32(b * config.adv_h)?,  // GLU d_gate_pre scratch
    ))
}
