#![allow(unsafe_code)]

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

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

use cudarc::cublas::result as cublas_result;
use cudarc::cublas::sys as cublas_sys;
use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg};

use crate::MLError;

// ── cuBLAS handle wrapper ─────────────────────────────────────────────────────

/// Wrapper to make `cublasHandle_t` Send + Sync.
///
/// cuBLAS handles are not thread-safe by default, but `CublasForward` is
/// owned exclusively by `GpuDqnTrainer` which is single-threaded in practice.
struct SendSyncCublasHandle(cublas_sys::cublasHandle_t);

unsafe impl Send for SendSyncCublasHandle {}
unsafe impl Sync for SendSyncCublasHandle {}

impl Drop for SendSyncCublasHandle {
    fn drop(&mut self) {
        unsafe {
            let _ = cublas_result::destroy_handle(self.0);
        }
    }
}

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

/// cuBLAS-based batched forward pass context for the branching dueling DQN.
///
/// Holds the cuBLAS handle and references to the trainer's pre-allocated
/// activation buffers. GEMM output lands directly in those buffers.
///
/// The handle is set to the trainer's stream so all GEMMs execute in order.
#[allow(missing_debug_implementations)] // handle is not Debug
pub struct CublasForward {
    handle: SendSyncCublasHandle,
    /// cuBLAS workspace buffer — must outlive the handle for CUDA Graph replay.
    _workspace_buf: CudaSlice<u8>,

    // ── Compiled bias+activation kernels (loaded from the DQN module) ──
    /// Kernel: `add_bias_relu(output, bias, out_dim, total_elements)`
    /// Grid: ceil(B * out_dim / 256), Block: 256.
    add_bias_relu_kernel: CudaFunction,

    /// Kernel: `add_bias(output, bias, out_dim, total_elements)`
    /// Same shape as add_bias_relu but no clamping (used for final logit layers).
    add_bias_kernel: CudaFunction,

    // ── Network dimensions (baked at construction) ──
    batch_size: usize,
    state_dim: usize,
    shared_h1: usize,
    shared_h2: usize,
    value_h: usize,
    adv_h: usize,
    num_atoms: usize,
    branch_0_size: usize,
    branch_1_size: usize,
    branch_2_size: usize,
}

impl CublasForward {
    /// Create a new cuBLAS forward context.
    ///
    /// Compiles the `add_bias_relu` and `add_bias` kernels from inline NVRTC source.
    /// Binds the cuBLAS handle to the provided stream.
    pub fn new(
        stream: &Arc<CudaStream>,
        batch_size: usize,
        state_dim: usize,
        shared_h1: usize,
        shared_h2: usize,
        value_h: usize,
        adv_h: usize,
        num_atoms: usize,
        branch_0_size: usize,
        branch_1_size: usize,
        branch_2_size: usize,
    ) -> Result<Self, MLError> {
        // ── Create cuBLAS handle ────────────────────────────────────────
        let raw_handle = cublas_result::create_handle()
            .map_err(|e| MLError::ModelError(format!("cublasCreate: {e:?}")))?;

        // Bind handle to stream — ALL GEMMs will execute on this stream.
        // SAFETY: cudarc::driver::sys and cudarc::cublas::sys both typedef
        // CUstream to *mut CUstream_st, which is the same underlying type.
        // The cast is safe as long as the pointer value is preserved.
        unsafe {
            let cu_stream = stream.cu_stream() as *mut cublas_sys::CUstream_st;
            cublas_result::set_stream(raw_handle, cu_stream)
                .map_err(|e| MLError::ModelError(format!("cublasSetStream: {e:?}")))?;
        }

        // Allocate explicit cuBLAS workspace for CUDA Graph compatibility.
        // Without this, cuBLAS uses an internal workspace that becomes invalid
        // during graph replay — causing zero gradients on H100 (SM_90) where
        // HMMA algorithms require workspace for batch_size > ~128.
        let workspace_size: usize = 4 * 1024 * 1024; // 4 MB
        let workspace_buf = stream.alloc_zeros::<u8>(workspace_size)
            .map_err(|e| MLError::ModelError(format!("cuBLAS workspace alloc: {e}")))?;
        unsafe {
            let (dev_ptr, _guard) = workspace_buf.device_ptr(stream);
            let ws_ptr = dev_ptr as *mut std::ffi::c_void;
            let status = cublas_sys::cublasSetWorkspace_v2(raw_handle, ws_ptr, workspace_size);
            if status != cublas_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
                return Err(MLError::ModelError(format!("cublasSetWorkspace_v2 failed: {status:?}")));
            }
        }

        // ── Compile bias kernels ────────────────────────────────────────
        let (add_bias_relu_kernel, add_bias_kernel) =
            compile_bias_kernels(stream)?;

        Ok(Self {
            handle: SendSyncCublasHandle(raw_handle),
            _workspace_buf: workspace_buf,
            add_bias_relu_kernel,
            add_bias_kernel,
            batch_size,
            state_dim,
            shared_h1,
            shared_h2,
            value_h,
            adv_h,
            num_atoms,
            branch_0_size,
            branch_1_size,
            branch_2_size,
        })
    }

    /// Rebind the cuBLAS handle to a (potentially new) stream.
    ///
    /// Must be called before CUDA Graph capture when the stream changes.
    pub fn set_stream(&self, stream: &CudaStream) -> Result<(), MLError> {
        unsafe {
            let cu_stream = stream.cu_stream() as *mut cublas_sys::CUstream_st;
            cublas_result::set_stream(self.handle.0, cu_stream)
                .map_err(|e| MLError::ModelError(format!("cublasSetStream: {e:?}")))
        }
    }

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

    /// Run the online network forward pass using cuBLAS SGEMM.
    ///
    /// Writes into the trainer's pre-allocated activation save buffers:
    /// - `save_h_s1`: [B, SH1]     — shared trunk layer 1 activations
    /// - `save_h_s2`: [B, SH2]     — shared trunk layer 2 activations
    /// - `save_h_v`:  [B, VH]      — value head hidden activations
    /// - `save_h_b0`: [B, AH]      — branch 0 hidden activations
    /// - `save_h_b1`: [B, AH]      — branch 1 hidden activations
    /// - `save_h_b2`: [B, AH]      — branch 2 hidden activations
    ///
    /// Value logits  → written to `v_logits_buf`  [B, NA]
    /// Branch logits → written to `b_logits_buf`  [B, (B0+B1+B2)*NA]
    ///
    /// Weight pointers are raw u64 into the flat F32 `params_buf` at GOFF_* offsets.
    #[allow(clippy::too_many_arguments)]
    pub fn forward_online(
        &self,
        stream: &Arc<CudaStream>,
        // ── Inputs ──────────────────────────────────────────────────────
        states_buf: &CudaSlice<f32>,          // [B, SD]
        // ── F32 weight pointers (raw u64 into flat params_buf) ─────────
        w_ptrs: &[u64; 20],                    // [w_s1, b_s1, w_s2, b_s2, ...]
        // ── Activation save buffers (written by this function) ─────────
        save_h_s1: &CudaSlice<f32>,           // [B, SH1]
        save_h_s2: &CudaSlice<f32>,           // [B, SH2]
        save_h_v: &CudaSlice<f32>,            // [B, VH]
        save_h_b0: &CudaSlice<f32>,           // [B, AH]
        save_h_b1: &CudaSlice<f32>,           // [B, AH]
        save_h_b2: &CudaSlice<f32>,           // [B, AH]
        // ── Logit output buffers ────────────────────────────────────────
        v_logits_buf: &CudaSlice<f32>,        // [B, NA]       — value head logits
        b_logits_buf: &CudaSlice<f32>,        // [B, (B0+B1+B2)*NA] — branch logits
    ) -> Result<(), MLError> {
        let b = self.batch_size;

        // ── Shared trunk layer 1: h_s1[B, SH1] = ReLU(states @ W_s1^T + b_s1) ──
        self.sgemm_layer(
            stream,
            w_ptrs[0], // W_s1[SH1, SD]
            raw_f32_ptr(states_buf, stream),
            raw_f32_ptr_mut(save_h_s1, stream),
            self.shared_h1, b, self.state_dim,
            "h_s1",
        )?;
        self.launch_add_bias_relu(stream, save_h_s1, w_ptrs[1], self.shared_h1, b)?; // b_s1

        // ── Shared trunk layer 2: h_s2[B, SH2] = ReLU(h_s1 @ W_s2^T + b_s2) ──
        self.sgemm_layer(
            stream,
            w_ptrs[2], // W_s2[SH2, SH1]
            raw_f32_ptr(save_h_s1, stream),
            raw_f32_ptr_mut(save_h_s2, stream),
            self.shared_h2, b, self.shared_h1,
            "h_s2",
        )?;
        self.launch_add_bias_relu(stream, save_h_s2, w_ptrs[3], self.shared_h2, b)?; // b_s2

        // ── Value head layer 1: h_v[B, VH] = ReLU(h_s2 @ W_v1^T + b_v1) ──
        self.sgemm_layer(
            stream,
            w_ptrs[4], // W_v1[VH, SH2]
            raw_f32_ptr(save_h_s2, stream),
            raw_f32_ptr_mut(save_h_v, stream),
            self.value_h, b, self.shared_h2,
            "h_v",
        )?;
        self.launch_add_bias_relu(stream, save_h_v, w_ptrs[5], self.value_h, b)?; // b_v1

        // ── Value head layer 2: v_logits[B, NA] = h_v @ W_v2^T + b_v2 ──
        self.sgemm_layer(
            stream,
            w_ptrs[6], // W_v2[NA, VH]
            raw_f32_ptr(save_h_v, stream),
            raw_f32_ptr_mut(v_logits_buf, stream),
            self.num_atoms, b, self.value_h,
            "v_logits",
        )?;
        self.launch_add_bias(stream, v_logits_buf, w_ptrs[7], self.num_atoms, b)?; // b_v2

        // ── Branch heads ─────────────────────────────────────────────────────
        let branch_sizes = [self.branch_0_size, self.branch_1_size, self.branch_2_size];
        let branch_save_h = [save_h_b0, save_h_b1, save_h_b2];
        // Weight indices: [w_b0fc=8, b_b0fc=9, w_b0out=10, b_b0out=11,
        //                  w_b1fc=12, b_b1fc=13, w_b1out=14, b_b1out=15,
        //                  w_b2fc=16, b_b2fc=17, w_b2out=18, b_b2out=19]
        let branch_w_base = [8_usize, 12, 16];
        let b_logits_ptr = raw_f32_ptr_mut(b_logits_buf, stream);
        let na = self.num_atoms;

        let mut logit_byte_offset: u64 = 0;
        for d in 0..3 {
            let n_d = branch_sizes[d];
            let w_fc_idx = branch_w_base[d];
            let b_fc_idx = w_fc_idx + 1;
            let w_out_idx = w_fc_idx + 2;
            let b_out_idx = w_fc_idx + 3;

            // h_bd[B, AH] = ReLU(h_s2 @ W_bdk_fc^T + b_bdk_fc)
            self.sgemm_layer(
                stream,
                w_ptrs[w_fc_idx], // W_bdk_fc[AH, SH2]
                raw_f32_ptr(save_h_s2, stream),
                raw_f32_ptr_mut(branch_save_h[d], stream),
                self.adv_h, b, self.shared_h2,
                "h_bd",
            )?;
            self.launch_add_bias_relu(stream, branch_save_h[d], w_ptrs[b_fc_idx], self.adv_h, b)?;

            // adv_logits_d[B, n_d*NA] = h_bd @ W_bdk_out^T + b_bdk_out
            // Output pointer: offset into b_logits_buf by (sum of prior branch atoms) * B * 4 bytes
            let adv_out_ptr = b_logits_ptr + logit_byte_offset;
            self.sgemm_layer_raw(
                stream,
                w_ptrs[w_out_idx], // W_bdk_out[n_d*NA, AH]
                raw_f32_ptr(branch_save_h[d], stream),
                adv_out_ptr,
                n_d * na, b, self.adv_h,
                "adv_logits",
            )?;
            // Bias for adv logits: b_logits_buf at current offset as a slice
            // We use raw add_bias on the pointer + offset
            self.launch_add_bias_raw(stream, adv_out_ptr, w_ptrs[b_out_idx], n_d * na, b)?;

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

        Ok(())
    }

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

        // v_logits = h_v @ W_v2^T + b_v2
        self.sgemm_layer_raw(
            stream, w_v2, h_v_scratch, v_logits_out,
            self.num_atoms, batch_size, self.value_h, "ens_v_logits",
        )?;
        self.launch_add_bias_raw(stream, v_logits_out, b_v2, self.num_atoms, batch_size)?;

        Ok(())
    }

    /// Run the target network forward pass (inference only — no activation saves).
    ///
    /// Same GEMM sequence as `forward_online` but writes to separate output buffers
    /// and does NOT save h_s2 (it's not needed for the backward pass through target).
    ///
    /// Outputs:
    /// - `tg_h_s2_buf`: [B, SH2]       — target trunk output (needed for Bellman target computation)
    /// - `tg_h_v_buf`:  [B, VH]        — target value head hidden (scratch, reused)
    /// - `tg_v_logits_buf`: [B, NA]    — target value head logits
    /// - `tg_b_logits_buf`: [B, (B0+B1+B2)*NA] — target branch logits
    ///
    /// The caller also needs to run the target network on next_states for DDQN.
    /// We write tg_h_s2_buf here so the caller can inspect it if needed.
    #[allow(clippy::too_many_arguments)]
    pub fn forward_target(
        &self,
        stream: &Arc<CudaStream>,
        // ── Inputs ──────────────────────────────────────────────────────
        next_states_buf: &CudaSlice<f32>,     // [B, SD]
        // ── F32 weight pointers (raw u64 into flat target_params_buf) ──
        tg_w_ptrs: &[u64; 20],
        // ── Scratch activation buffers (not saved for backward) ────────
        tg_h_s1_scratch: &CudaSlice<f32>,    // [B, SH1]   scratch
        tg_h_s2_buf: &CudaSlice<f32>,        // [B, SH2]   kept for Double-DQN
        tg_h_v_scratch: &CudaSlice<f32>,     // [B, VH]    scratch
        tg_h_b_scratch: &CudaSlice<f32>,     // [B, AH]    scratch (reused for all 3 branches)
        // ── Output buffers ──────────────────────────────────────────────
        tg_v_logits_buf: &CudaSlice<f32>,    // [B, NA]
        tg_b_logits_buf: &CudaSlice<f32>,    // [B, (B0+B1+B2)*NA]
    ) -> Result<(), MLError> {
        let b = self.batch_size;

        // h_s1[B, SH1]
        self.sgemm_layer(
            stream,
            tg_w_ptrs[0],
            raw_f32_ptr(next_states_buf, stream),
            raw_f32_ptr_mut(tg_h_s1_scratch, stream),
            self.shared_h1, b, self.state_dim,
            "tg_h_s1",
        )?;
        self.launch_add_bias_relu(stream, tg_h_s1_scratch, tg_w_ptrs[1], self.shared_h1, b)?;

        // h_s2[B, SH2]
        self.sgemm_layer(
            stream,
            tg_w_ptrs[2],
            raw_f32_ptr(tg_h_s1_scratch, stream),
            raw_f32_ptr_mut(tg_h_s2_buf, stream),
            self.shared_h2, b, self.shared_h1,
            "tg_h_s2",
        )?;
        self.launch_add_bias_relu(stream, tg_h_s2_buf, tg_w_ptrs[3], self.shared_h2, b)?;

        // h_v[B, VH]
        self.sgemm_layer(
            stream,
            tg_w_ptrs[4],
            raw_f32_ptr(tg_h_s2_buf, stream),
            raw_f32_ptr_mut(tg_h_v_scratch, stream),
            self.value_h, b, self.shared_h2,
            "tg_h_v",
        )?;
        self.launch_add_bias_relu(stream, tg_h_v_scratch, tg_w_ptrs[5], self.value_h, b)?;

        // v_logits[B, NA]
        self.sgemm_layer(
            stream,
            tg_w_ptrs[6],
            raw_f32_ptr(tg_h_v_scratch, stream),
            raw_f32_ptr_mut(tg_v_logits_buf, stream),
            self.num_atoms, b, self.value_h,
            "tg_v_logits",
        )?;
        self.launch_add_bias(stream, tg_v_logits_buf, tg_w_ptrs[7], self.num_atoms, b)?;

        // Branch heads — reuse tg_h_b_scratch for all 3 (sequential, no aliasing)
        let branch_sizes = [self.branch_0_size, self.branch_1_size, self.branch_2_size];
        let branch_w_base = [8_usize, 12, 16];
        let b_logits_ptr = raw_f32_ptr_mut(tg_b_logits_buf, stream);
        let na = self.num_atoms;
        let mut logit_byte_offset: u64 = 0;

        for d in 0..3 {
            let n_d = branch_sizes[d];
            let w_fc_idx = branch_w_base[d];

            self.sgemm_layer(
                stream,
                tg_w_ptrs[w_fc_idx],
                raw_f32_ptr(tg_h_s2_buf, stream),
                raw_f32_ptr_mut(tg_h_b_scratch, stream),
                self.adv_h, b, self.shared_h2,
                "tg_h_bd",
            )?;
            self.launch_add_bias_relu(stream, tg_h_b_scratch, tg_w_ptrs[w_fc_idx + 1], self.adv_h, b)?;

            let adv_out_ptr = b_logits_ptr + logit_byte_offset;
            self.sgemm_layer_raw(
                stream,
                tg_w_ptrs[w_fc_idx + 2],
                raw_f32_ptr(tg_h_b_scratch, stream),
                adv_out_ptr,
                n_d * na, b, self.adv_h,
                "tg_adv_logits",
            )?;
            self.launch_add_bias_raw(stream, adv_out_ptr, tg_w_ptrs[w_fc_idx + 3], n_d * na, b)?;

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

        Ok(())
    }

    // ══════════════════════════════════════════════════════════════════════════
    // cuBLAS SGEMM helpers
    // ══════════════════════════════════════════════════════════════════════════

    /// Single cuBLAS SGEMM: C[B, N] = A[B, K] @ W[N, K]^T  (row-major).
    ///
    /// Call convention (column-major cuBLAS, row-major data):
    ///
    /// sgemm(CUBLAS_OP_T, CUBLAS_OP_N, N, B, K, 1.0, W, K, A, K, 0.0, C, N)
    ///
    /// This outputs C in col-major [N, B] which equals row-major C[B, N]. ✓
    ///
    /// Arguments:
    /// - `w_ptr`  : raw device pointer to W[N, K] (row-major, stride K)
    /// - `a_ptr`  : raw device pointer to A[B, K] (row-major, stride K)
    /// - `c_buf`  : output CudaSlice<f32> [B * N] (written in col-major [N, B])
    /// - `n`      : output cols (out_dim)
    /// - `b`      : batch size (rows of A / cols of output in col-major)
    /// - `k`      : inner dim (in_dim = cols of W = cols of A)
    #[allow(clippy::too_many_arguments)]
    fn sgemm_layer(
        &self,
        _stream: &Arc<CudaStream>,
        w_ptr: u64,
        a_ptr: u64,
        c_ptr: u64,
        n: usize,
        b: usize,
        k: usize,
        _label: &str,
    ) -> Result<(), MLError> {
        let alpha = 1.0_f32;
        let beta = 0.0_f32;

        // SAFETY: w_ptr, a_ptr, c_ptr are valid CUDA device pointers in the same
        // context. Sizes are computed from config values that were validated at
        // construction. The cuBLAS handle is bound to the correct stream.
        unsafe {
            cublas_result::sgemm(
                self.handle.0,
                cublas_sys::cublasOperation_t::CUBLAS_OP_T,  // transa: transpose W
                cublas_sys::cublasOperation_t::CUBLAS_OP_N,  // transb: keep A as-is
                n as i32,      // m (rows of op(W) = N = out_dim)
                b as i32,      // n (cols of op(A) = B = batch)
                k as i32,      // k (inner dim)
                &alpha,
                w_ptr as *const f32,
                k as i32,      // lda: leading dim of W (before transpose) = K
                a_ptr as *const f32,
                k as i32,      // ldb: leading dim of A = K
                &beta,
                c_ptr as *mut f32,
                n as i32,      // ldc: leading dim of C = N
            )
            .map_err(|e| MLError::ModelError(format!("cublasSgemm {_label}: {e:?}")))?;
        }

        Ok(())
    }

    /// Same as `sgemm_layer` but takes raw output pointer directly.
    ///
    /// Used when writing into a sub-region of a larger buffer (branch logits).
    #[allow(clippy::too_many_arguments)]
    fn sgemm_layer_raw(
        &self,
        _stream: &Arc<CudaStream>,
        w_ptr: u64,
        a_ptr: u64,
        c_ptr: u64,
        n: usize,
        b: usize,
        k: usize,
        _label: &str,
    ) -> Result<(), MLError> {
        self.sgemm_layer(_stream, w_ptr, a_ptr, c_ptr, n, b, k, _label)
    }

    // ══════════════════════════════════════════════════════════════════════════
    // Bias + activation kernel launchers
    // ══════════════════════════════════════════════════════════════════════════

    /// Launch `add_bias_relu` over an activation buffer.
    ///
    /// `output[i] = max(0, output[i] + bias[i % out_dim])`
    ///
    /// Grid: ceil(B * out_dim / 256), Block: 256.
    fn launch_add_bias_relu(
        &self,
        stream: &Arc<CudaStream>,
        output: &CudaSlice<f32>,
        bias_ptr: u64,
        out_dim: usize,
        batch: usize,
    ) -> Result<(), MLError> {
        let total = (batch * out_dim) as i32;
        let out_dim_i32 = out_dim as i32;
        let blocks = ((batch * out_dim + 255) / 256) as u32;

        let out_ptr = raw_f32_ptr_mut(output, stream);

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

        Ok(())
    }

    /// Launch `add_bias_relu` with a raw output pointer (sub-buffer support).
    fn launch_add_bias_relu_raw(
        &self,
        stream: &Arc<CudaStream>,
        out_ptr: u64,
        bias_ptr: u64,
        out_dim: usize,
        batch: usize,
    ) -> Result<(), MLError> {
        let total = (batch * out_dim) as i32;
        let out_dim_i32 = out_dim as i32;
        let blocks = ((batch * out_dim + 255) / 256) as u32;

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

        Ok(())
    }

    /// Launch `add_bias` (no activation) over an output buffer.
    ///
    /// `output[i] = output[i] + bias[i % out_dim]`
    fn launch_add_bias(
        &self,
        stream: &Arc<CudaStream>,
        output: &CudaSlice<f32>,
        bias_ptr: u64,
        out_dim: usize,
        batch: usize,
    ) -> Result<(), MLError> {
        let out_ptr = raw_f32_ptr_mut(output, stream);
        self.launch_add_bias_raw(stream, out_ptr, bias_ptr, out_dim, batch)
    }

    /// Launch `add_bias` with a raw output pointer (sub-buffer support).
    fn launch_add_bias_raw(
        &self,
        stream: &Arc<CudaStream>,
        out_ptr: u64,
        bias_ptr: u64,
        out_dim: usize,
        batch: usize,
    ) -> Result<(), MLError> {
        let total = (batch * out_dim) as i32;
        let out_dim_i32 = out_dim as i32;
        let blocks = ((batch * out_dim + 255) / 256) as u32;

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

        Ok(())
    }
}

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

/// Extract raw F32 device pointer (read-only).
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
}

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

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

/// Compile `add_bias_relu` and `add_bias` kernels from inline NVRTC source.
///
/// These kernels are tiny (1 FMA per thread) and launch after each GEMM to:
/// 1. Broadcast the bias across the batch dimension.
/// 2. Apply ReLU (for hidden layers) or no activation (for logit layers).
///
/// They are NOT captured in the CUDA Graph — they run inline in the stream.
/// For CUDA Graph compatibility, we'd add them to the captured region, but
/// the small launch overhead (< 1 µs per kernel) is negligible vs. GEMM time.
fn compile_bias_kernels(
    stream: &Arc<CudaStream>,
) -> Result<(CudaFunction, CudaFunction), MLError> {
    // IMPORTANT: these are standalone kernels, not requiring common_device_functions.cuh.
    // They must be compiled without the STATE_DIM / MARKET_DIM guards.
    let src = r#"
extern "C" __global__ void add_bias_relu_kernel(
    float* __restrict__ output,
    const float* __restrict__ bias,
    int out_dim,
    int total_elements)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= total_elements) return;
    float val = output[i] + bias[i % out_dim];
    output[i] = (val > 0.0f) ? val : 0.0f;
}

extern "C" __global__ void add_bias_kernel(
    float* __restrict__ output,
    const float* __restrict__ bias,
    int out_dim,
    int total_elements)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= total_elements) return;
    output[i] += bias[i % out_dim];
}
"#;

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

    let add_bias_relu = module
        .load_function("add_bias_relu_kernel")
        .map_err(|e| MLError::ModelError(format!("add_bias_relu_kernel load: {e}")))?;
    let add_bias = module
        .load_function("add_bias_kernel")
        .map_err(|e| MLError::ModelError(format!("add_bias_kernel load: {e}")))?;

    Ok((add_bias_relu, add_bias))
}

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

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

    let mut ptrs = [0_u64; 20];
    let mut byte_offset: u64 = 0;
    for i in 0..20 {
        ptrs[i] = base + byte_offset;
        byte_offset += (param_sizes[i] * std::mem::size_of::<f32>()) as u64;
    }
    ptrs
}
