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

//! GpuGrn — Gated Residual Network block wrapper (Plan 4 Task 2c.2).
//!
//! Additive Rust wrapper around the 8 kernels landed in `grn_kernel.cu`
//! (Plan 4 Task 2c.1, commit `68197c2c2`). One `GrnBlock` instance owns
//! the per-block save-for-backward state buffers and the `CudaFunction`
//! handles for all 8 GRN kernels. Forward and backward dispatcher methods
//! sequence the custom kernels; the cuBLAS GEMMs (`Linear_a`, `Linear_b`,
//! optional `Linear_residual`) are called by the caller outside this
//! wrapper, matching the pattern of `gpu_tlob.rs` / `gpu_attention.rs`.
//!
//! NO PRODUCTION CALLERS in this commit — `GrnBlock` is dead code until
//! Task 2c.3+4 wires it into the trunk encoder. Module compiles, kernels
//! load via the existing `CudaFunction` infrastructure, wrapper sits as
//! additive scaffolding.
//!
//! # GRN block formula (per spec §4.E.2)
//!
//!   GRN(x) = LayerNorm(P(x) + GLU(Linear_b(ELU(Linear_a(x)))))
//!
//! where `P(x) = Linear_residual(x)` if `dim_in != dim_out`, else identity.
//!
//! # Forward composition (caller sequences cuBLAS + this wrapper)
//!
//! 1. cuBLAS:  `linear_a_out = Linear_a(x)`                  → [B, hidden]
//! 2. kernel:  `grn_elu_inplace(linear_a_out)`               (post-ELU written
//!             back into the same buffer — wrapper then DtoD-copies the
//!             post-activation into `state.elu_post` for backward, since
//!             `grn_elu_backward` consumes the POST-activation per its
//!             docstring (it recovers `f'(x_pre)` from `x_post` via the
//!             identity `exp(x_pre) = x_post + 1` for `x_post < 0`).
//! 3. cuBLAS:  `linear_b_out = Linear_b(elu_out)`            → [B, 2*hidden]
//!             Wrapper DtoD-copies into `state.linear_b_out` (consumed by
//!             `grn_glu_backward` as `in_2h_saved`).
//! 4. kernel:  `grn_glu_forward(linear_b_out, glu_out, state.sigmoid_b)`
//!                                                            → [B, hidden]
//! 5. cuBLAS (h_s1 only): `linear_residual_out = Linear_residual(x)`
//!             Else (h_s2): residual is `x` directly (caller passes ptr).
//! 6. kernel:  `grn_residual_layernorm_forward(glu_out, residual, gamma,
//!             beta, y, state.ln_mean, state.ln_rstd, state.ln_normed)`
//!                                                            → [B, hidden]
//!
//! # Backward composition
//!
//! 1. kernel:  `grn_layernorm_backward_dx(d_y, ln_normed, ln_rstd, gamma)`
//!                                                            → d_pre_LN
//! 2. kernel:  `grn_layernorm_backward_dgamma_dbeta_p1(d_y, ln_normed)`
//!             then `_p2(...)` finalises d_gamma, d_beta. No atomicAdd.
//! 3. Residual split: `d_glu_out = d_pre_LN`, `d_residual = d_pre_LN`.
//!             Pure pointer aliasing — caller routes the same buffer to
//!             both downstream consumers (Linear_b backward GEMM and the
//!             residual-projection backward / direct accumulation).
//! 4. kernel:  `grn_glu_backward(d_glu_out, state.linear_b_out,
//!             state.sigmoid_b)` → d_linear_b_out [B, 2*hidden]
//! 5. cuBLAS:  Linear_b backward — `d_W_b`, `d_b_b`, `d_elu_out` (caller).
//! 6. kernel:  `grn_elu_backward(d_elu_out, state.elu_post)` → d_linear_a_out
//! 7. cuBLAS:  Linear_a backward — `d_W_a`, `d_b_a`, `d_x_from_a` (caller).
//! 8. cuBLAS (h_s1 only): Linear_residual backward — `d_W_r`, `d_b_r`,
//!             `d_x_from_residual`. Caller accumulates
//!             `d_x = d_x_from_a + d_x_from_residual`. Else (h_s2),
//!             `d_x = d_x_from_a + d_residual` (residual gradient flows
//!             directly back through the identity branch).
//!
//! # Pearls applied
//!
//! - `pearl_cold_path_no_exception_to_gpu_drives.md` — every reduction
//!   stays GPU-side; this wrapper does zero CPU/DtoH work.
//! - `feedback_no_atomicadd.md` — the dgamma/dbeta backward is two-phase
//!   partial-reduce + final-reduce, not atomicAdd.
//! - `feedback_no_feature_flags.md` — no enable_/use_ booleans; the only
//!   axis of behavioural variation is `has_residual_projection`, which is
//!   a structural property of the block (h_s1 vs h_s2), not a toggle.
//! - `feedback_wire_everything_up.md` — this wrapper is intentionally
//!   built as a pre-wire scaffold for Task 2c.3+4, mirroring the same
//!   commit-by-commit ordering used for the kernel landing in 2c.1. The
//!   2c.3+4 wire-up is the obligatory follow-up that retires the dead-
//!   code status.

use std::sync::Arc;

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

use crate::MLError;

// ── Cubin embedding ─────────────────────────────────────────────────────────

use super::gpu_dqn_trainer::GRN_CUBIN;

/// One GRN block (h_s1 SD→SH1 or h_s2 SH1→SH2). Owns 5 save-for-backward
/// state buffers + 2 partial-reduction scratch buffers + 8 kernel handles.
#[allow(missing_debug_implementations)]
pub(crate) struct GrnBlock {
    stream: Arc<CudaStream>,
    /// Hidden dim H of this block's GRN output (= LN width = GLU output width).
    hidden_dim: usize,
    /// Maximum batch size used to size the state buffers.
    batch_size: usize,
    /// True for h_s1 (SD → SH1, requires `Linear_residual`); false for h_s2
    /// (SH1 → SH2, residual is identity since the spec sets SH1 == SH2 for
    /// the second GRN). Used by 2c.3+4 to decide whether the caller must
    /// run the residual-projection cuBLAS GEMM.
    has_residual_projection: bool,

    // ── Save-for-backward state ───────────────────────────────────────
    /// Post-ELU activation: `ELU(Linear_a(x))`. Consumed by
    /// `grn_elu_backward` — see kernel docstring: backward takes the
    /// POST-activation, not pre-, recovers `f'(x_pre)` via the identity
    /// `exp(x_pre) = x_post + 1` for `x_post < 0`.
    elu_post: CudaSlice<f32>,           // [B, H]
    /// `Linear_b` output (pre-GLU split). Consumed by `grn_glu_backward`
    /// as `in_2h_saved`; backward needs `chunk_a` (the first H columns)
    /// to compute `d_chunk_b = d_out * chunk_a * σ * (1-σ)`.
    linear_b_out: CudaSlice<f32>,       // [B, 2*H]
    /// Sigmoid of GLU's gate half. Consumed by `grn_glu_backward`.
    sigmoid_b: CudaSlice<f32>,          // [B, H]
    /// LN forward saved per-sample mean. Consumed by LN backward dx.
    ln_mean: CudaSlice<f32>,            // [B]
    /// LN forward saved per-sample reciprocal-stddev. Consumed by LN
    /// backward dx.
    ln_rstd: CudaSlice<f32>,            // [B]
    /// LN forward saved normalized values (pre-affine). Consumed by both
    /// LN backward kernels.
    ln_normed: CudaSlice<f32>,          // [B, H]

    // ── Partial-reduction scratch (no atomicAdd in dgamma/dbeta) ─────
    /// Per-block partials for d_gamma reduction across batch.
    /// Layout: `[num_blocks_b, H]` row-major. Phase 1 writes per-block
    /// partials; phase 2 sums across `num_blocks_b` into d_gamma[H].
    dgamma_partials: CudaSlice<f32>,    // [num_blocks_b, H]
    /// Per-block partials for d_beta reduction across batch.
    dbeta_partials: CudaSlice<f32>,     // [num_blocks_b, H]
    /// Number of phase-1 blocks across batch (`(B + 255) / 256`).
    num_blocks_b: u32,

    // ── Kernel handles ───────────────────────────────────────────────
    elu_inplace_kernel: CudaFunction,
    glu_forward_kernel: CudaFunction,
    residual_layernorm_forward_kernel: CudaFunction,
    layernorm_backward_dx_kernel: CudaFunction,
    layernorm_backward_dgamma_dbeta_p1_kernel: CudaFunction,
    layernorm_backward_dgamma_dbeta_p2_kernel: CudaFunction,
    glu_backward_kernel: CudaFunction,
    elu_backward_kernel: CudaFunction,
}

impl GrnBlock {
    /// Allocate state buffers and load all 8 GRN kernel handles from the
    /// precompiled cubin.
    ///
    /// `hidden_dim` is the GRN output width (LN width). For h_s1 this is
    /// SH1; for h_s2 this is SH2 (== SH1 in the current spec). `batch_size`
    /// is the maximum batch the trainer will dispatch through this block.
    /// `has_residual_projection` is `true` for h_s1 (input dim SD differs
    /// from output dim SH1, so a `Linear_residual` GEMM is needed) and
    /// `false` for h_s2 (identity residual).
    pub(crate) fn new(
        stream: Arc<CudaStream>,
        hidden_dim: usize,
        batch_size: usize,
        has_residual_projection: bool,
    ) -> Result<Self, MLError> {
        let h = hidden_dim;
        let b = batch_size;

        // ── Load cubin module + 8 kernel functions ───────────────────
        let context = stream.context();
        let module = context.load_cubin(GRN_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("grn cubin load: {e}")))?;
        let elu_inplace_kernel = module.load_function("grn_elu_inplace")
            .map_err(|e| MLError::ModelError(format!("grn_elu_inplace load: {e}")))?;
        let glu_forward_kernel = module.load_function("grn_glu_forward")
            .map_err(|e| MLError::ModelError(format!("grn_glu_forward load: {e}")))?;
        let residual_layernorm_forward_kernel = module.load_function("grn_residual_layernorm_forward")
            .map_err(|e| MLError::ModelError(format!("grn_residual_layernorm_forward load: {e}")))?;
        let layernorm_backward_dx_kernel = module.load_function("grn_layernorm_backward_dx")
            .map_err(|e| MLError::ModelError(format!("grn_layernorm_backward_dx load: {e}")))?;
        let layernorm_backward_dgamma_dbeta_p1_kernel = module.load_function("grn_layernorm_backward_dgamma_dbeta_p1")
            .map_err(|e| MLError::ModelError(format!("grn_layernorm_backward_dgamma_dbeta_p1 load: {e}")))?;
        let layernorm_backward_dgamma_dbeta_p2_kernel = module.load_function("grn_layernorm_backward_dgamma_dbeta_p2")
            .map_err(|e| MLError::ModelError(format!("grn_layernorm_backward_dgamma_dbeta_p2 load: {e}")))?;
        let glu_backward_kernel = module.load_function("grn_glu_backward")
            .map_err(|e| MLError::ModelError(format!("grn_glu_backward load: {e}")))?;
        let elu_backward_kernel = module.load_function("grn_elu_backward")
            .map_err(|e| MLError::ModelError(format!("grn_elu_backward load: {e}")))?;

        // ── Allocate save-for-backward state ─────────────────────────
        let alloc = |n: usize, label: &str| -> Result<CudaSlice<f32>, MLError> {
            stream.alloc_zeros::<f32>(n)
                .map_err(|e| MLError::ModelError(format!("grn {label} alloc: {e}")))
        };

        let elu_post     = alloc(b * h,         "elu_post")?;
        let linear_b_out = alloc(b * 2 * h,     "linear_b_out")?;
        let sigmoid_b    = alloc(b * h,         "sigmoid_b")?;
        let ln_mean      = alloc(b,             "ln_mean")?;
        let ln_rstd      = alloc(b,             "ln_rstd")?;
        let ln_normed    = alloc(b * h,         "ln_normed")?;

        // Phase-1 batch reduction grid: matches the kernel's launch shape
        //     grid=(num_blocks_b, H), block=256.
        // Each phase-1 block reduces a 256-stripe of B for one hidden dim.
        let num_blocks_b = ((b + 255) / 256) as u32;
        let dgamma_partials = alloc(num_blocks_b as usize * h, "dgamma_partials")?;
        let dbeta_partials  = alloc(num_blocks_b as usize * h, "dbeta_partials")?;

        info!(
            hidden_dim = h,
            batch_size = b,
            has_residual_projection = has_residual_projection,
            num_blocks_b = num_blocks_b,
            state_bytes = (b * h * 3 + b * 2 * h + b * 2 + num_blocks_b as usize * h * 2) * std::mem::size_of::<f32>(),
            "GrnBlock initialised: 8 kernels loaded, 5 save-for-backward + 2 partial-reduce buffers allocated"
        );

        Ok(Self {
            stream,
            hidden_dim: h,
            batch_size: b,
            has_residual_projection,
            elu_post,
            linear_b_out,
            sigmoid_b,
            ln_mean,
            ln_rstd,
            ln_normed,
            dgamma_partials,
            dbeta_partials,
            num_blocks_b,
            elu_inplace_kernel,
            glu_forward_kernel,
            residual_layernorm_forward_kernel,
            layernorm_backward_dx_kernel,
            layernorm_backward_dgamma_dbeta_p1_kernel,
            layernorm_backward_dgamma_dbeta_p2_kernel,
            glu_backward_kernel,
            elu_backward_kernel,
        })
    }

    /// Whether this block needs a `Linear_residual` projection (h_s1) or
    /// can use identity residual (h_s2). Set at construction by 2c.3+4.
    #[inline]
    pub(crate) fn has_residual_projection(&self) -> bool {
        self.has_residual_projection
    }

    /// Hidden dim H (= GRN output width = LN width).
    #[inline]
    pub(crate) fn hidden_dim(&self) -> usize {
        self.hidden_dim
    }

    /// Maximum batch size this block was sized for.
    #[inline]
    pub(crate) fn max_batch_size(&self) -> usize {
        self.batch_size
    }

    /// Raw device pointer to `state.elu_post` (post-ELU activation, [B, H]).
    /// Populated by `forward_raw` / `forward_raw_phase1` when
    /// `save_for_backward = true`. Used by the trainer's GRN-stage NaN
    /// diagnostic instrumentation (P, post-bkdx5 — slot 16/17 of nan_flags).
    #[inline]
    pub(crate) fn elu_post_ptr(&self) -> u64 {
        self.elu_post.raw_ptr()
    }

    /// Raw device pointer to `state.linear_b_out` (Linear_b output pre-GLU,
    /// [B, 2*H]). Populated by `forward_raw` / `forward_raw_phase2` when
    /// `save_for_backward = true`. Used by the trainer's GRN-stage NaN
    /// diagnostic instrumentation (P, post-bkdx5).
    #[inline]
    pub(crate) fn linear_b_out_ptr(&self) -> u64 {
        self.linear_b_out.raw_ptr()
    }

    /// Raw device pointer to `state.sigmoid_b` (sigmoid of GLU gate half,
    /// [B, H]). Populated inside `grn_glu_forward`. Used by the trainer's
    /// GRN-stage NaN diagnostic instrumentation (P, post-bkdx5).
    #[inline]
    pub(crate) fn sigmoid_b_ptr(&self) -> u64 {
        self.sigmoid_b.raw_ptr()
    }

    /// Raw device pointer to `state.ln_rstd` (LN reciprocal-stddev, [B]).
    /// Populated inside `grn_residual_layernorm_forward`. Used by the
    /// trainer's GRN-stage NaN diagnostic instrumentation (P, post-bkdx5):
    /// if forward variance hits zero, rstd → ∞ → downstream NaN.
    #[inline]
    pub(crate) fn ln_rstd_ptr(&self) -> u64 {
        self.ln_rstd.raw_ptr()
    }

    /// Raw device pointer to `state.ln_normed` (LN normalized values pre-
    /// affine, [B, H]). Populated inside `grn_residual_layernorm_forward`.
    /// Used by the trainer's GRN-stage NaN diagnostic instrumentation
    /// (P, post-bkdx5).
    #[inline]
    pub(crate) fn ln_normed_ptr(&self) -> u64 {
        self.ln_normed.raw_ptr()
    }

    /// Forward pass — sequences the 3 GRN forward kernels (CudaSlice variant).
    ///
    /// Prefer `forward_raw` for graph-captured callers that operate on raw
    /// `u64` device pointers (avoids cudarc's `device_ptr` event-tracking
    /// machinery, which is incompatible with CUDA Graph capture).
    ///
    ///   1. `grn_elu_inplace(linear_a_out)` — in-place ELU.
    ///   2. DtoD copy `linear_a_out` (now post-ELU) → `state.elu_post`
    ///      (saves post-activation for `grn_elu_backward`).
    ///   3. DtoD copy caller's `linear_b_out_in` → `state.linear_b_out`
    ///      (saves pre-split for `grn_glu_backward`).
    ///   4. `grn_glu_forward(linear_b_out, B, H, glu_out, state.sigmoid_b)`.
    ///   5. `grn_residual_layernorm_forward(glu_out, residual, gamma, beta,
    ///         B, H, y, state.ln_mean, state.ln_rstd, state.ln_normed)`.
    ///
    /// cuBLAS GEMMs (`Linear_a`, `Linear_b`, optional `Linear_residual`)
    /// are issued by the caller before calling this method. Buffers are
    /// passed as `&mut CudaSlice<f32>` (or `&CudaSlice<f32>` for read-only
    /// inputs) so cudarc tracks their lifetimes against the launch.
    ///
    /// # Args
    /// * `linear_a_out` — caller-owned [B, H] buffer holding `Linear_a(x)
    ///   + b_a` on entry. Mutated in-place to `ELU(Linear_a(x) + b_a)`.
    /// * `linear_b_out_in` — caller-owned [B, 2*H] buffer holding
    ///   `Linear_b(ELU_out) + b_b`. The wrapper saves a DtoD copy.
    /// * `residual` — [B, H] residual input. For h_s1 this is the output
    ///   of `Linear_residual(x)`; for h_s2 this is `x` directly.
    /// * `gamma`, `beta` — LN affine params [H] each.
    /// * `glu_out` — caller-owned [B, H] scratch for GLU output.
    /// * `y` — caller-owned [B, H] output of the GRN block.
    /// * `batch_size` — runtime batch count (must satisfy
    ///   `batch_size <= self.batch_size`).
    #[allow(dead_code)]
    pub(crate) fn forward(
        &mut self,
        linear_a_out: &mut CudaSlice<f32>,
        linear_b_out_in: &CudaSlice<f32>,
        residual: &CudaSlice<f32>,
        gamma: &CudaSlice<f32>,
        beta: &CudaSlice<f32>,
        glu_out: &mut CudaSlice<f32>,
        y: &mut CudaSlice<f32>,
        batch_size: usize,
    ) -> Result<(), MLError> {
        if batch_size > self.batch_size {
            return Err(MLError::ModelError(format!(
                "grn forward: batch_size {} exceeds wrapper capacity {}",
                batch_size, self.batch_size,
            )));
        }
        let h = self.hidden_dim;
        let b = batch_size;
        let h_i32 = h as i32;
        let b_i32 = b as i32;
        let n_elem_h = (b * h) as i32;     // post-ELU and per-element kernels
        let blocks_bh = ((b * h + 255) / 256) as u32;

        // ── 1. ELU in-place on linear_a_out ──────────────────────────
        unsafe {
            self.stream
                .launch_builder(&self.elu_inplace_kernel)
                .arg(&mut *linear_a_out)
                .arg(&n_elem_h)
                .launch(LaunchConfig {
                    grid_dim: (blocks_bh, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("grn_elu_inplace: {e}")))?;
        }

        // ── 2. DtoD save: post-ELU → state.elu_post ──────────────────
        {
            let nbytes = b * h * std::mem::size_of::<f32>();
            let (src_ptr, _sg) = linear_a_out.device_ptr(&self.stream);
            let (dst_ptr, _dg) = self.elu_post.device_ptr_mut(&self.stream);
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, src_ptr, nbytes, self.stream.cu_stream(),
                ).map_err(|e| MLError::ModelError(format!("grn save elu_post: {e}")))?;
            }
        }

        // ── 3. DtoD save: linear_b_out_in → state.linear_b_out ───────
        {
            let nbytes = b * 2 * h * std::mem::size_of::<f32>();
            let (src_ptr, _sg) = linear_b_out_in.device_ptr(&self.stream);
            let (dst_ptr, _dg) = self.linear_b_out.device_ptr_mut(&self.stream);
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, src_ptr, nbytes, self.stream.cu_stream(),
                ).map_err(|e| MLError::ModelError(format!("grn save linear_b_out: {e}")))?;
            }
        }

        // ── 4. GLU forward ───────────────────────────────────────────
        unsafe {
            self.stream
                .launch_builder(&self.glu_forward_kernel)
                .arg(linear_b_out_in)
                .arg(&b_i32)
                .arg(&h_i32)
                .arg(&mut *glu_out)
                .arg(&mut self.sigmoid_b)
                .launch(LaunchConfig {
                    grid_dim: (blocks_bh, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("grn_glu_forward: {e}")))?;
        }

        // ── 5. Residual + LayerNorm fused forward ────────────────────
        // Grid = B blocks, each block reduces one row. Block dim = min(H, 256)
        // — matches kernel expectation (warp_sums sized for up to 8 warps).
        let ln_block_dim = if h < 256 { h.max(32) as u32 } else { 256 };
        unsafe {
            self.stream
                .launch_builder(&self.residual_layernorm_forward_kernel)
                .arg(&mut *glu_out)
                .arg(residual)
                .arg(gamma)
                .arg(beta)
                .arg(&b_i32)
                .arg(&h_i32)
                .arg(y)
                .arg(&mut self.ln_mean)
                .arg(&mut self.ln_rstd)
                .arg(&mut self.ln_normed)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (ln_block_dim, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("grn_residual_layernorm_forward: {e}")))?;
        }

        Ok(())
    }

    /// Graph-safe forward pass: takes raw `u64` device pointers instead of
    /// `&mut CudaSlice<f32>`. Equivalent to `forward()` but bypasses
    /// cudarc's `device_ptr` event-tracking machinery so the launches can
    /// be captured into a CUDA Graph (Plan 4 Task 2c.3b).
    ///
    /// The kernel sequence is identical to `forward`. The save-for-backward
    /// DtoD copies (`elu_post`, `linear_b_out`) are gated on
    /// `save_for_backward` — set to `false` for target-network forwards
    /// where the saved state is never consumed (the GRN backward only runs
    /// against the online network).
    ///
    /// # Args (all u64 device pointers)
    /// * `linear_a_out_ptr` — caller-owned [B, H] buffer holding
    ///   `Linear_a(x) + b_a` on entry. Mutated in-place to
    ///   `ELU(Linear_a(x) + b_a)`.
    /// * `linear_b_out_ptr` — caller-owned [B, 2*H] holding
    ///   `Linear_b(ELU_out) + b_b`. Read-only (DtoD-copied into
    ///   `state.linear_b_out` when saving).
    /// * `residual_ptr` — caller-owned [B, H] residual input.
    /// * `gamma_ptr`, `beta_ptr` — LN affine params [H] each.
    /// * `glu_out_ptr` — caller-owned [B, H] scratch for GLU output.
    /// * `y_ptr` — caller-owned [B, H] output of the GRN block.
    /// * `batch_size` — runtime batch count.
    /// * `save_for_backward` — when true, DtoD-copy post-ELU and Linear_b
    ///   output into the wrapper's save buffers; when false skip the
    ///   copies (target-network case where backward is never run).
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn forward_raw(
        &mut self,
        stream: &Arc<CudaStream>,
        linear_a_out_ptr: u64,
        linear_b_out_ptr: u64,
        residual_ptr: u64,
        gamma_ptr: u64,
        beta_ptr: u64,
        glu_out_ptr: u64,
        y_ptr: u64,
        batch_size: usize,
        save_for_backward: bool,
    ) -> Result<(), MLError> {
        if batch_size > self.batch_size {
            return Err(MLError::ModelError(format!(
                "grn forward_raw: batch_size {} exceeds wrapper capacity {}",
                batch_size, self.batch_size,
            )));
        }
        let h = self.hidden_dim;
        let b = batch_size;
        let h_i32 = h as i32;
        let b_i32 = b as i32;
        let n_elem_h = (b * h) as i32;
        let blocks_bh = ((b * h + 255) / 256) as u32;

        // ── 1. ELU in-place on linear_a_out (raw u64) ────────────────
        unsafe {
            stream
                .launch_builder(&self.elu_inplace_kernel)
                .arg(&linear_a_out_ptr)
                .arg(&n_elem_h)
                .launch(LaunchConfig {
                    grid_dim: (blocks_bh, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("grn_elu_inplace (raw): {e}")))?;
        }

        // ── 2/3. Save-for-backward DtoD copies (skipped for target) ──
        if save_for_backward {
            // post-ELU → state.elu_post
            let nbytes_h = b * h * std::mem::size_of::<f32>();
            let (dst_ptr, _dg) = self.elu_post.device_ptr_mut(stream);
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, linear_a_out_ptr, nbytes_h, stream.cu_stream(),
                ).map_err(|e| MLError::ModelError(format!("grn save elu_post (raw): {e}")))?;
            }
            // linear_b_out → state.linear_b_out
            let nbytes_2h = b * 2 * h * std::mem::size_of::<f32>();
            let (dst2_ptr, _dg2) = self.linear_b_out.device_ptr_mut(stream);
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst2_ptr, linear_b_out_ptr, nbytes_2h, stream.cu_stream(),
                ).map_err(|e| MLError::ModelError(format!("grn save linear_b_out (raw): {e}")))?;
            }
        }

        // ── 4. GLU forward (raw u64) ─────────────────────────────────
        unsafe {
            stream
                .launch_builder(&self.glu_forward_kernel)
                .arg(&linear_b_out_ptr)
                .arg(&b_i32)
                .arg(&h_i32)
                .arg(&glu_out_ptr)
                .arg(&mut self.sigmoid_b)
                .launch(LaunchConfig {
                    grid_dim: (blocks_bh, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("grn_glu_forward (raw): {e}")))?;
        }

        // ── 5. Residual + LayerNorm fused forward (raw u64) ──────────
        let ln_block_dim = if h < 256 { h.max(32) as u32 } else { 256 };
        unsafe {
            stream
                .launch_builder(&self.residual_layernorm_forward_kernel)
                .arg(&glu_out_ptr)
                .arg(&residual_ptr)
                .arg(&gamma_ptr)
                .arg(&beta_ptr)
                .arg(&b_i32)
                .arg(&h_i32)
                .arg(&y_ptr)
                .arg(&mut self.ln_mean)
                .arg(&mut self.ln_rstd)
                .arg(&mut self.ln_normed)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (ln_block_dim, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("grn_residual_layernorm_forward (raw): {e}")))?;
        }

        Ok(())
    }

    /// Phase-1 forward (raw u64): in-place ELU on `linear_a_out_ptr` +
    /// optional DtoD save into `state.elu_post`. Splits `forward_raw`
    /// across the cuBLAS Linear_b GEMM — the encoder dispatcher must
    /// run Linear_a → ELU → Linear_b → Linear_residual → GLU+LN, so the
    /// wrapper needs a "before Linear_b" entry point and an "after
    /// Linear_residual" entry point.
    ///
    /// Set `save_for_backward = false` to skip the DtoD copy (target
    /// network forward never runs backward against this instance).
    pub(crate) fn forward_raw_phase1(
        &mut self,
        stream: &Arc<CudaStream>,
        linear_a_out_ptr: u64,
        batch_size: usize,
        save_for_backward: bool,
    ) -> Result<(), MLError> {
        if batch_size > self.batch_size {
            return Err(MLError::ModelError(format!(
                "grn forward_raw_phase1: batch_size {} exceeds wrapper capacity {}",
                batch_size, self.batch_size,
            )));
        }
        let h = self.hidden_dim;
        let b = batch_size;
        let n_elem_h = (b * h) as i32;
        let blocks_bh = ((b * h + 255) / 256) as u32;

        // ── 1. ELU in-place on linear_a_out (raw u64) ────────────────
        unsafe {
            stream
                .launch_builder(&self.elu_inplace_kernel)
                .arg(&linear_a_out_ptr)
                .arg(&n_elem_h)
                .launch(LaunchConfig {
                    grid_dim: (blocks_bh, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("grn_elu_inplace (raw p1): {e}")))?;
        }

        // ── 2. Save-for-backward DtoD copy: post-ELU → state.elu_post ─
        if save_for_backward {
            let nbytes_h = b * h * std::mem::size_of::<f32>();
            let (dst_ptr, _dg) = self.elu_post.device_ptr_mut(stream);
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, linear_a_out_ptr, nbytes_h, stream.cu_stream(),
                ).map_err(|e| MLError::ModelError(format!("grn save elu_post (raw p1): {e}")))?;
            }
        }

        Ok(())
    }

    /// Phase-2 forward (raw u64): DtoD-save `linear_b_out_ptr` into
    /// `state.linear_b_out` (when saving), then GLU forward, then
    /// residual+LN forward, writing `y_ptr`.
    ///
    /// The caller must have already issued cuBLAS Linear_b on the
    /// post-ELU `linear_a_out` and (for h_s1) Linear_residual on the
    /// input `x` to fill `linear_b_out_ptr` and `residual_ptr`
    /// respectively.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn forward_raw_phase2(
        &mut self,
        stream: &Arc<CudaStream>,
        linear_b_out_ptr: u64,
        residual_ptr: u64,
        gamma_ptr: u64,
        beta_ptr: u64,
        glu_out_ptr: u64,
        y_ptr: u64,
        batch_size: usize,
        save_for_backward: bool,
    ) -> Result<(), MLError> {
        if batch_size > self.batch_size {
            return Err(MLError::ModelError(format!(
                "grn forward_raw_phase2: batch_size {} exceeds wrapper capacity {}",
                batch_size, self.batch_size,
            )));
        }
        let h = self.hidden_dim;
        let b = batch_size;
        let h_i32 = h as i32;
        let b_i32 = b as i32;
        let blocks_bh = ((b * h + 255) / 256) as u32;

        // ── DtoD save: linear_b_out → state.linear_b_out ─────────────
        if save_for_backward {
            let nbytes_2h = b * 2 * h * std::mem::size_of::<f32>();
            let (dst_ptr, _dg) = self.linear_b_out.device_ptr_mut(stream);
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, linear_b_out_ptr, nbytes_2h, stream.cu_stream(),
                ).map_err(|e| MLError::ModelError(format!("grn save linear_b_out (raw p2): {e}")))?;
            }
        }

        // ── GLU forward ──────────────────────────────────────────────
        unsafe {
            stream
                .launch_builder(&self.glu_forward_kernel)
                .arg(&linear_b_out_ptr)
                .arg(&b_i32)
                .arg(&h_i32)
                .arg(&glu_out_ptr)
                .arg(&mut self.sigmoid_b)
                .launch(LaunchConfig {
                    grid_dim: (blocks_bh, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("grn_glu_forward (raw p2): {e}")))?;
        }

        // ── Residual + LayerNorm fused forward ───────────────────────
        let ln_block_dim = if h < 256 { h.max(32) as u32 } else { 256 };
        unsafe {
            stream
                .launch_builder(&self.residual_layernorm_forward_kernel)
                .arg(&glu_out_ptr)
                .arg(&residual_ptr)
                .arg(&gamma_ptr)
                .arg(&beta_ptr)
                .arg(&b_i32)
                .arg(&h_i32)
                .arg(&y_ptr)
                .arg(&mut self.ln_mean)
                .arg(&mut self.ln_rstd)
                .arg(&mut self.ln_normed)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (ln_block_dim, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("grn_residual_layernorm_forward (raw p2): {e}")))?;
        }

        Ok(())
    }

    /// Backward pass — sequences the 5 GRN backward kernels (LN dx, LN
    /// dgamma/dbeta two-phase, GLU bwd, ELU bwd):
    ///
    ///   1. `grn_layernorm_backward_dx(d_y, ln_normed, ln_rstd, gamma)`
    ///      → `d_pre_LN`. This same buffer feeds BOTH the gated branch
    ///      and the residual branch (residual split is just pointer
    ///      aliasing, no kernel).
    ///   2. `grn_layernorm_backward_dgamma_dbeta_p1(d_y, ln_normed)`
    ///      → `dgamma_partials`, `dbeta_partials`.
    ///   3. `grn_layernorm_backward_dgamma_dbeta_p2(...)` → `d_gamma`,
    ///      `d_beta`. No atomicAdd — phase 1 writes per-block stripes,
    ///      phase 2 reduces across blocks.
    ///   4. `grn_glu_backward(d_glu_out, state.linear_b_out,
    ///      state.sigmoid_b)` → `d_linear_b_out` [B, 2*H].
    ///   5. Caller runs `Linear_b` backward GEMMs → `d_elu_out` [B, H].
    ///   6. `grn_elu_backward(d_elu_out, state.elu_post)` → `d_linear_a_out`.
    ///
    /// The caller is responsible for:
    ///   - The cuBLAS Linear_a / Linear_b / Linear_residual backward GEMMs
    ///     between steps 4 and 6 (and after step 6).
    ///   - Treating `d_pre_LN` as the gradient flowing into BOTH the gated
    ///     branch (passed as `d_glu_out` to step 4) and the residual
    ///     branch (consumed by `Linear_residual` backward for h_s1 or
    ///     accumulated directly into `d_x` for h_s2).
    ///
    /// # Args
    /// * `d_y` — gradient of the GRN output [B, H], coming from above.
    /// * `gamma` — LN gamma parameter [H], same buffer as forward.
    /// * `d_pre_LN` — caller-owned [B, H] output: gradient flowing into
    ///   the (gated + residual) sum prior to LN. Used as both `d_glu_out`
    ///   and `d_residual` by the caller via aliasing.
    /// * `d_linear_b_out` — caller-owned [B, 2*H] output: gradient w.r.t.
    ///   the pre-GLU `Linear_b` output, feeds Linear_b backward GEMM.
    /// * `d_elu_out` — caller-owned [B, H] input: gradient w.r.t. the
    ///   post-ELU activation, written by caller's Linear_b backward GEMM.
    /// * `d_linear_a_out` — caller-owned [B, H] output: gradient w.r.t.
    ///   the pre-ELU activation, feeds Linear_a backward GEMM.
    /// * `d_gamma`, `d_beta` — caller-owned [H] outputs: LN parameter
    ///   gradients.
    /// * `batch_size` — runtime batch count.
    #[allow(dead_code)]
    pub(crate) fn backward(
        &mut self,
        d_y: &CudaSlice<f32>,
        gamma: &CudaSlice<f32>,
        d_pre_ln: &mut CudaSlice<f32>,
        d_linear_b_out: &mut CudaSlice<f32>,
        d_elu_out: &CudaSlice<f32>,
        d_linear_a_out: &mut CudaSlice<f32>,
        d_gamma: &mut CudaSlice<f32>,
        d_beta: &mut CudaSlice<f32>,
        batch_size: usize,
    ) -> Result<(), MLError> {
        if batch_size > self.batch_size {
            return Err(MLError::ModelError(format!(
                "grn backward: batch_size {} exceeds wrapper capacity {}",
                batch_size, self.batch_size,
            )));
        }
        let h = self.hidden_dim;
        let b = batch_size;
        let h_i32 = h as i32;
        let b_i32 = b as i32;
        let num_blocks_b_runtime = ((b + 255) / 256) as u32;
        if num_blocks_b_runtime > self.num_blocks_b {
            // Should be impossible given the batch_size check above, but guard
            // so a future signature change can't silently undersize the partials.
            return Err(MLError::ModelError(format!(
                "grn backward: runtime num_blocks_b {} exceeds preallocated {}",
                num_blocks_b_runtime, self.num_blocks_b,
            )));
        }

        // ── 1. LN backward dx ────────────────────────────────────────
        // Grid = B (one block per row); block dim = min(H, 256), matching
        // the forward LN kernel's launch shape.
        let ln_block_dim = if h < 256 { h.max(32) as u32 } else { 256 };
        unsafe {
            self.stream
                .launch_builder(&self.layernorm_backward_dx_kernel)
                .arg(d_y)
                .arg(&self.ln_normed)
                .arg(&self.ln_rstd)
                .arg(gamma)
                .arg(&b_i32)
                .arg(&h_i32)
                .arg(&mut *d_pre_ln)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (ln_block_dim, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("grn_layernorm_backward_dx: {e}")))?;
        }

        // ── 2. LN backward dgamma/dbeta phase 1 (per-block partials) ─
        // Kernel uses `extern __shared__ float sdata[]` sized
        // 2 × blockDim.x floats (one buffer for gamma partials, one for
        // beta partials, reduced in lockstep).
        let smem_bytes_p1 = (2 * 256 * std::mem::size_of::<f32>()) as u32;
        unsafe {
            self.stream
                .launch_builder(&self.layernorm_backward_dgamma_dbeta_p1_kernel)
                .arg(d_y)
                .arg(&self.ln_normed)
                .arg(&b_i32)
                .arg(&h_i32)
                .arg(&mut self.dgamma_partials)
                .arg(&mut self.dbeta_partials)
                .launch(LaunchConfig {
                    grid_dim: (num_blocks_b_runtime, h as u32, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: smem_bytes_p1,
                })
                .map_err(|e| MLError::ModelError(format!("grn_layernorm_backward_dgamma_dbeta_p1: {e}")))?;
        }

        // ── 3. LN backward dgamma/dbeta phase 2 (final reduce) ───────
        let num_blocks_b_i32 = num_blocks_b_runtime as i32;
        let p2_blocks = ((h + 255) / 256) as u32;
        unsafe {
            self.stream
                .launch_builder(&self.layernorm_backward_dgamma_dbeta_p2_kernel)
                .arg(&self.dgamma_partials)
                .arg(&self.dbeta_partials)
                .arg(&h_i32)
                .arg(&num_blocks_b_i32)
                .arg(d_gamma)
                .arg(d_beta)
                .launch(LaunchConfig {
                    grid_dim: (p2_blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("grn_layernorm_backward_dgamma_dbeta_p2: {e}")))?;
        }

        // ── 4. GLU backward ──────────────────────────────────────────
        // Residual split: d_pre_ln carries d_(gated + residual). The kernel
        // consumes it as d_glu_out — caller is responsible for ALSO routing
        // the same d_pre_ln to the residual-branch backward (Linear_residual
        // for h_s1 or direct accumulation for h_s2). The aliasing is the
        // outer wiring's concern; here we just feed d_pre_ln in.
        let blocks_bh = ((b * h + 255) / 256) as u32;
        unsafe {
            self.stream
                .launch_builder(&self.glu_backward_kernel)
                .arg(&*d_pre_ln)             // d_out_h: gradient into GLU output (read-only)
                .arg(&self.linear_b_out)     // in_2h_saved: pre-GLU Linear_b output
                .arg(&self.sigmoid_b)
                .arg(&b_i32)
                .arg(&h_i32)
                .arg(&mut *d_linear_b_out)
                .launch(LaunchConfig {
                    grid_dim: (blocks_bh, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("grn_glu_backward: {e}")))?;
        }

        // ── 5. (caller runs Linear_b backward GEMMs here, producing
        //        d_elu_out from d_linear_b_out) ────────────────────────

        // ── 6. ELU backward — d_x = d_out * f'(x_pre) where f' is
        //      recovered from the post-activation x_saved via the kernel's
        //      identity (x_saved >= 0 → 1; x_saved < 0 → x_saved + 1). ──
        let n_elem_h = (b * h) as i32;
        let blocks_n = ((b * h + 255) / 256) as u32;
        unsafe {
            self.stream
                .launch_builder(&self.elu_backward_kernel)
                .arg(d_elu_out)
                .arg(&self.elu_post)
                .arg(&n_elem_h)
                .arg(d_linear_a_out)
                .launch(LaunchConfig {
                    grid_dim: (blocks_n, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("grn_elu_backward: {e}")))?;
        }

        Ok(())
    }

    /// Plan 4 Task 2c.3c.1: backward phase 1 (raw u64) — LN backward + GLU
    /// backward. Mirrors `forward_raw_phase1`/`forward_raw_phase2` for the
    /// same composition reason as the forward split: the encoder's cuBLAS
    /// `Linear_b` backward GEMM must run BETWEEN GLU backward and ELU
    /// backward (it produces the `d_elu_out` gradient that ELU backward
    /// consumes). Splitting the monolithic `backward()` at that boundary
    /// lets the trunk-encoder dispatcher (Task 2c.3c.4) compose:
    ///
    ///   phase1: d_y → d_linear_b_out + d_pre_LN (residual aliasing)
    ///   caller: cuBLAS Linear_b backward — d_linear_b_out → d_elu_out
    ///                                                     + d_W_b + d_b_b
    ///   phase2: d_elu_out → d_linear_a_out
    ///   caller: cuBLAS Linear_a backward — d_linear_a_out → d_x_from_a
    ///                                                     + d_W_a + d_b_a
    ///   caller (h_s1 only): cuBLAS Linear_residual backward —
    ///     d_pre_LN → d_x_from_residual + d_W_residual
    ///   caller: d_x = d_x_from_a + d_x_from_residual (h_s1)
    ///         OR d_x = d_x_from_a + d_pre_LN          (h_s2, identity residual)
    ///
    /// Phase 1 sequences four kernels:
    ///   1. `grn_layernorm_backward_dx`            → `d_pre_LN_ptr`
    ///   2. `grn_layernorm_backward_dgamma_dbeta_p1` (per-block partials)
    ///   3. `grn_layernorm_backward_dgamma_dbeta_p2` (final reduce → `d_gamma`,
    ///      `d_beta`; no atomicAdd)
    ///   4. `grn_glu_backward(d_pre_LN, state.linear_b_out, state.sigmoid_b)`
    ///      → `d_linear_b_out_ptr`
    ///
    /// **Residual-split contract**: `d_pre_LN_ptr` is the gradient flowing
    /// into the (gated + residual) sum prior to LN — i.e. `d_glu_out`
    /// **and** `d_residual` share the same buffer. Phase 1 writes it once;
    /// the caller consumes it from BOTH downstream paths via implicit
    /// pointer aliasing (no kernel, no copy). This mirrors the GRN forward
    /// where the residual branch is also a pure pointer pass.
    ///
    /// `save_for_backward` state buffers (`elu_post`, `linear_b_out`,
    /// `sigmoid_b`, `ln_mean`, `ln_rstd`, `ln_normed`) must have been
    /// populated by a preceding `forward_raw` / `forward_raw_phase{1,2}`
    /// with `save_for_backward = true` on this same `GrnBlock` instance.
    ///
    /// # Args (all u64 device pointers)
    /// * `d_y_ptr` — [B, H] gradient of the GRN output, coming from above.
    /// * `gamma_ptr` — [H] LN scale (same buffer as forward).
    /// * `batch_size` — runtime batch count.
    /// * `d_pre_ln_ptr` — [B, H] OUTPUT: gradient flowing into (gated +
    ///   residual) sum pre-LN. Used as both `d_glu_out` (consumed inside
    ///   this method by `grn_glu_backward`) AND `d_residual` (consumed by
    ///   the caller's residual-branch backward).
    /// * `d_linear_b_out_ptr` — [B, 2*H] OUTPUT: gradient w.r.t. pre-GLU
    ///   `Linear_b` output, feeds the caller's Linear_b backward GEMM.
    /// * `d_gamma_ptr` — [H] OUTPUT: LN gamma gradient (no atomicAdd —
    ///   two-phase per-block reduction).
    /// * `d_beta_ptr` — [H] OUTPUT: LN beta gradient.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn backward_raw_phase1(
        &mut self,
        stream: &Arc<CudaStream>,
        d_y_ptr: u64,
        gamma_ptr: u64,
        batch_size: usize,
        d_pre_ln_ptr: u64,
        d_linear_b_out_ptr: u64,
        d_gamma_ptr: u64,
        d_beta_ptr: u64,
    ) -> Result<(), MLError> {
        if batch_size > self.batch_size {
            return Err(MLError::ModelError(format!(
                "grn backward_raw_phase1: batch_size {} exceeds wrapper capacity {}",
                batch_size, self.batch_size,
            )));
        }
        let h = self.hidden_dim;
        let b = batch_size;
        let h_i32 = h as i32;
        let b_i32 = b as i32;
        let num_blocks_b_runtime = ((b + 255) / 256) as u32;
        if num_blocks_b_runtime > self.num_blocks_b {
            return Err(MLError::ModelError(format!(
                "grn backward_raw_phase1: runtime num_blocks_b {} exceeds preallocated {}",
                num_blocks_b_runtime, self.num_blocks_b,
            )));
        }

        // ── 1. LN backward dx (raw u64) ──────────────────────────────
        let ln_block_dim = if h < 256 { h.max(32) as u32 } else { 256 };
        unsafe {
            stream
                .launch_builder(&self.layernorm_backward_dx_kernel)
                .arg(&d_y_ptr)
                .arg(&self.ln_normed)
                .arg(&self.ln_rstd)
                .arg(&gamma_ptr)
                .arg(&b_i32)
                .arg(&h_i32)
                .arg(&d_pre_ln_ptr)
                .launch(LaunchConfig {
                    grid_dim: (b as u32, 1, 1),
                    block_dim: (ln_block_dim, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("grn_layernorm_backward_dx (raw p1): {e}")))?;
        }

        // ── 2. LN backward dgamma/dbeta phase 1 (per-block partials) ─
        let smem_bytes_p1 = (2 * 256 * std::mem::size_of::<f32>()) as u32;
        unsafe {
            stream
                .launch_builder(&self.layernorm_backward_dgamma_dbeta_p1_kernel)
                .arg(&d_y_ptr)
                .arg(&self.ln_normed)
                .arg(&b_i32)
                .arg(&h_i32)
                .arg(&mut self.dgamma_partials)
                .arg(&mut self.dbeta_partials)
                .launch(LaunchConfig {
                    grid_dim: (num_blocks_b_runtime, h as u32, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: smem_bytes_p1,
                })
                .map_err(|e| MLError::ModelError(format!("grn_layernorm_backward_dgamma_dbeta_p1 (raw p1): {e}")))?;
        }

        // ── 3. LN backward dgamma/dbeta phase 2 (final reduce) ───────
        let num_blocks_b_i32 = num_blocks_b_runtime as i32;
        let p2_blocks = ((h + 255) / 256) as u32;
        unsafe {
            stream
                .launch_builder(&self.layernorm_backward_dgamma_dbeta_p2_kernel)
                .arg(&self.dgamma_partials)
                .arg(&self.dbeta_partials)
                .arg(&h_i32)
                .arg(&num_blocks_b_i32)
                .arg(&d_gamma_ptr)
                .arg(&d_beta_ptr)
                .launch(LaunchConfig {
                    grid_dim: (p2_blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("grn_layernorm_backward_dgamma_dbeta_p2 (raw p1): {e}")))?;
        }

        // ── 4. GLU backward — consumes d_pre_ln_ptr as d_glu_out ─────
        // Residual aliasing: caller will ALSO route the same d_pre_ln_ptr
        // to the residual-branch backward (Linear_residual for h_s1, or
        // direct accumulation for h_s2). The kernel only reads d_glu_out
        // (it doesn't write into d_pre_ln_ptr), so the aliasing is safe.
        let blocks_bh = ((b * h + 255) / 256) as u32;
        unsafe {
            stream
                .launch_builder(&self.glu_backward_kernel)
                .arg(&d_pre_ln_ptr)              // d_out_h: read-only (gradient into GLU out)
                .arg(&self.linear_b_out)         // in_2h_saved: pre-GLU Linear_b output
                .arg(&self.sigmoid_b)
                .arg(&b_i32)
                .arg(&h_i32)
                .arg(&d_linear_b_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (blocks_bh, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("grn_glu_backward (raw p1): {e}")))?;
        }

        Ok(())
    }

    /// Plan 4 Task 2c.3c.1: backward phase 2 (raw u64) — ELU backward.
    ///
    /// Runs only `grn_elu_backward`. Takes `d_elu_out_ptr` (computed by the
    /// caller's cuBLAS Linear_b backward GEMM, which sits between phase 1
    /// and phase 2) and writes `d_linear_a_out_ptr` for the caller's
    /// subsequent Linear_a backward GEMM.
    ///
    /// `state.elu_post` (populated by `forward_raw_phase1` with
    /// `save_for_backward = true`) is consumed here — the kernel recovers
    /// `f'(x_pre)` from the POST-activation via the identity
    /// `exp(x_pre) = x_post + 1` for `x_post < 0`, per the kernel's
    /// docstring.
    ///
    /// # Args (all u64 device pointers)
    /// * `d_elu_out_ptr` — [B, H] INPUT: gradient w.r.t. post-ELU
    ///   activation, written by the caller's Linear_b backward GEMM.
    /// * `batch_size` — runtime batch count.
    /// * `d_linear_a_out_ptr` — [B, H] OUTPUT: gradient w.r.t. pre-ELU
    ///   activation (= `Linear_a(x) + b_a`), feeds the caller's Linear_a
    ///   backward GEMM.
    pub(crate) fn backward_raw_phase2(
        &mut self,
        stream: &Arc<CudaStream>,
        d_elu_out_ptr: u64,
        batch_size: usize,
        d_linear_a_out_ptr: u64,
    ) -> Result<(), MLError> {
        if batch_size > self.batch_size {
            return Err(MLError::ModelError(format!(
                "grn backward_raw_phase2: batch_size {} exceeds wrapper capacity {}",
                batch_size, self.batch_size,
            )));
        }
        let h = self.hidden_dim;
        let b = batch_size;
        let n_elem_h = (b * h) as i32;
        let blocks_n = ((b * h + 255) / 256) as u32;

        unsafe {
            stream
                .launch_builder(&self.elu_backward_kernel)
                .arg(&d_elu_out_ptr)
                .arg(&self.elu_post)
                .arg(&n_elem_h)
                .arg(&d_linear_a_out_ptr)
                .launch(LaunchConfig {
                    grid_dim: (blocks_n, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("grn_elu_backward (raw p2): {e}")))?;
        }

        Ok(())
    }
}

// `GrnBlock` holds an `Arc<CudaStream>` and `CudaFunction` handles, both of
// which are `Send + Sync` in cudarc. The save buffers are `CudaSlice<f32>`
// which is `Send` but not `Sync`; the wrapper itself takes `&mut self` for
// any stateful method, so it is safely `Send` for cross-thread ownership
// transfer (the trainer pattern). Matching `gpu_tlob.rs`'s explicit impl.
unsafe impl Send for GrnBlock {}
