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

//! GpuTlob — TLOB attention module (D.8 Plan 2 Task 6C).
//!
//! Compresses OFI[SL_OFI_DIM=32] → TLOB[SL_TLOB_DIM=16] using single-head
//! scaled-dot-product attention over the feature dimension.
//!
//! Architecture:
//!   [Q || K || V][3*16,B] = `cublasSgemmStridedBatched`(W_Q/K/V^T, ofi, batch=3)
//!     — single launch in place of three separate Q/K/V SGEMMs (`fwd_qkv`
//!       contiguous strideC=M·B chunks at offsets 0, M·B, 2·M·B in the
//!       fused `proj_qkv_buf` of length 3·M·B).
//!   attn_out[16,B] = SDP(Q, K, V)             (tlob_sdp_forward cubin)
//!   out[16,B] = W_O^T @ attn_out[16,B] → [16,B] (classic cublasSgemm_v2)
//!   written to states_buf at SL_TLOB_START    (tlob_write_states cubin)
//!
//! Backward:
//!   d_out[16,B] extracted from bn_d_concat_buf (tlob_read_grad_from_states)
//!   d_attn_out[16,B] = W_O @ d_out            (classic cublasSgemm_v2)
//!   dW_O[16,16] += attn_out^T @ d_out / B     (classic cublasSgemm_v2)
//!   d_q/k/v[16,B] = SDP backward              (tlob_sdp_backward cubin)
//!   dW_Q/K/V[16,32] += ofi^T @ d_qkv / B      (`cublasSgemmStridedBatched`,
//!                                              batch=3, single launch)
//!
//! Weights are Xavier-initialised (random init, no pretrained checkpoint).
//! Adam step is called after backward — same pattern as GpuAttention.
//!
//! The ISV slot TLOB_REGIME_FOCUS_EMA_INDEX (slot 60) carries the mean of
//! the per-sample max attention weight, updated at per-epoch CPU call.
//! Zero atomicAdd; no CUDA Graph capture for TLOB (runs outside graph on
//! the main stream, following the same pattern as GpuAttention).
//!
//! ## cuBLAS API choice — classic sgemm + strided batched, not cuBLAS-Lt
//!
//! TLOB attention GEMMs have shapes M=TLOB_OUT=16, K∈{16,32}, N=batch.
//! These are too narrow for cuBLAS-Lt's tensor-core-oriented heuristic
//! (`cublasLtMatmulAlgoGetHeuristic` returns "no algo found" for M<32 on
//! TF32 paths). Per NVIDIA best practice, cuBLAS-Lt is the right tool for
//! large/tensor-core-optimised GEMMs at scale (M,K,N ≥ 64–128); classic
//! `cublasSgemm_v2` (and its strided-batched variant) is the right tool
//! for general-purpose any-shape GEMMs. TLOB belongs in the second
//! category. The larger trunk attention in `gpu_attention.rs` (dims ≥ 128)
//! keeps using cuBLAS-Lt where its tensor-core heuristic does pay off.
//!
//! ## Q/K/V fusion (strided batched)
//!
//! The forward/backward Q/K/V GEMMs share operands (ofi for forward, ofi
//! for the dW backward) and only differ in the weight pointer. We dispatch
//! a single `cublasSgemmStridedBatched(batch=3)` instead of three separate
//! `cublasSgemm_v2` launches, cutting kernel-launch overhead 3× without
//! changing the SDP kernel's per-buffer access pattern: the 3 contiguous
//! M·B-float output chunks (Q, K, V) live in a single fused `proj_qkv_buf`
//! and are passed to the SDP kernel as 3 offset device pointers.
//!
//! Strided-batched layout:
//!   Forward:  A=W_Q (strideA=M·K), B=ofi (strideB=0), C=proj_qkv (strideC=M·B)
//!   Backward: A=d_proj_qkv (strideA=M·B), B=ofi (strideB=0), C=d_W_QKV (strideC=M·K)
//!
//! Param flat layout `[W_Q | W_K | W_V | W_O]` is unchanged: W_Q/K/V are
//! consecutive M·K-float blocks, so strideA=M·K reads them in order with
//! no checkpoint format change. This was the deciding factor over Strategy
//! 1 (concatenated weights, single SGEMM with M=3·M=48): the concat-W path
//! requires the SDP kernel to read with stride-3M (col-major [3M, B], ldc=3M),
//! which forces a kernel signature change — incompatible with bit-for-bit
//! equivalence to the prior 3-SGEMM path that we verify in
//! `crates/ml/tests/tlob_qkv_fusion_equivalence.rs`. Strategy 2 keeps the
//! existing `[M, B]` per-projection layout intact, so the SDP kernel and
//! its access pattern are byte-identical pre/post fusion.

use std::sync::Arc;

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

use crate::MLError;
use super::shared_cublas_handle::PerStreamCublasHandles;
use ml_core::state_layout::{OFI_DIM, TLOB_DIM, STATE_DIM_PADDED};

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

static TLOB_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/tlob_kernel.cubin"));

// ── Dimensions ─────────────────────────────────────────────────────────────

/// TLOB Q/K/V projection input width (OFI features).
const TLOB_IN: usize = OFI_DIM;   // 32
/// TLOB Q/K/V projection output width (compressed dim).
const TLOB_OUT: usize = TLOB_DIM; // 16

/// Parameter layout (flat buffer):
///   W_Q[TLOB_OUT, TLOB_IN] + W_K[TLOB_OUT, TLOB_IN] + W_V[TLOB_OUT, TLOB_IN]
///   + W_O[TLOB_OUT, TLOB_OUT]
/// No bias terms — LayerNorm is intentionally omitted for MVP to stay within
/// the 800 LOC budget. The Q/K/V zero-bias simplification is acceptable for
/// D.8 because the OFI features are already zero-mean (MBP-10 delta/ratio feats).
const TLOB_W_QKV_SIZE: usize = 3 * TLOB_OUT * TLOB_IN;  // 3 × 16 × 32 = 1536
const TLOB_W_O_SIZE:   usize = TLOB_OUT * TLOB_OUT;      // 16 × 16 = 256
const TLOB_TOTAL_PARAMS: usize = TLOB_W_QKV_SIZE + TLOB_W_O_SIZE; // 1792

// Offsets within the flat params buffer
const W_Q_OFF: usize = 0;
const W_K_OFF: usize = TLOB_OUT * TLOB_IN;                       // 512
const W_V_OFF: usize = 2 * TLOB_OUT * TLOB_IN;                   // 1024
const W_O_OFF: usize = 3 * TLOB_OUT * TLOB_IN;                   // 1536

// Strided-batched stride constants for the Q/K/V fusion. The flat
// `[W_Q | W_K | W_V]` block is read with `strideA = M·K = 512` floats
// (one weight matrix per batch index). The fused `proj_qkv_buf` /
// `d_proj_qkv_buf` is laid out as 3 contiguous `[M, B]` col-major chunks
// at offsets 0, M·B, 2·M·B (`strideC = M·B`). Both strides are derived,
// not hard-coded — they follow directly from the existing per-projection
// shapes.
const QKV_BATCH: usize = 3;
const W_QKV_STRIDE_FLOATS: usize = TLOB_OUT * TLOB_IN; // M*K per weight matrix

// ── GpuTlob ─────────────────────────────────────────────────────────────────

/// TLOB attention module: OFI[32,B] → compressed features[16,B].
///
/// Weights are random-init (Xavier), trainable end-to-end via DQN reward.
/// No pretrained checkpoint; no ONNX.
#[allow(missing_debug_implementations)]
pub(crate) struct GpuTlob {
    stream: Arc<CudaStream>,
    shared_handle: Arc<PerStreamCublasHandles>,
    batch_size: usize,

    // ── cuBLAS GEMM dispatch ──────────────────────────────────────────
    // TLOB uses classic `cublasSgemm_v2` (not cuBLAS-Lt) because the GEMM
    // shapes — M=TLOB_OUT=16, K∈{16,32}, N=batch — are too narrow for
    // cuBLAS-Lt's tensor-core-oriented algo set. The classic-sgemm path
    // is dispatched on `shared_handle.classic_handle` already bound to
    // `stream`. There are no cached descriptors: classic sgemm takes
    // (transA, transB, m, n, k, lda, ldb, ldc) inline at every call.

    // ── cubin kernels ─────────────────────────────────────────────────
    sdp_fwd_kernel: CudaFunction,
    sdp_bwd_kernel: CudaFunction,
    write_states_kernel: CudaFunction,
    read_grad_kernel: CudaFunction,

    // ── Weight parameters [TLOB_TOTAL_PARAMS] ────────────────────────
    /// Flat: W_Q[512] + W_K[512] + W_V[512] + W_O[256] = 1792 floats.
    params: CudaSlice<f32>,

    // ── Intermediate buffers (forward, saved for backward) ────────────
    /// OFI input in col-major [TLOB_IN=32, B]: extracted from states_buf.
    ofi_col_buf: CudaSlice<f32>,
    /// Fused Q/K/V projections [3 × TLOB_OUT=16, B]: a single allocation
    /// holding three contiguous `[M, B]` col-major chunks at offsets
    /// 0, M·B, 2·M·B. Produced by one `cublasSgemmStridedBatched(batch=3)`
    /// call (replacing three separate Q/K/V SGEMMs). The SDP forward and
    /// backward kernels read each chunk through an offset device pointer,
    /// so their access pattern is byte-identical to the pre-fusion layout.
    proj_qkv_buf: CudaSlice<f32>,
    /// SDP output [TLOB_OUT=16, B] (pre-output-projection), saved for dW_O.
    attn_out_buf: CudaSlice<f32>,
    /// SDP softmax weights [TLOB_OUT=16, B], saved for backward.
    attn_scores_buf: CudaSlice<f32>,
    /// TLOB output [TLOB_OUT=16, B] = W_O^T @ attn_out.
    output_buf: CudaSlice<f32>,

    // ── Gradient buffers (backward) ───────────────────────────────────
    /// Gradient from trunk w.r.t. TLOB output [TLOB_OUT=16, B].
    d_output_buf: CudaSlice<f32>,
    /// Gradient w.r.t. attn_out [TLOB_OUT=16, B].
    d_attn_out_buf: CudaSlice<f32>,
    /// Fused gradients w.r.t. Q/K/V [3 × TLOB_OUT=16, B]: same layout as
    /// `proj_qkv_buf` (three `[M, B]` col-major chunks at offsets 0, M·B,
    /// 2·M·B). The SDP backward kernel writes each chunk through an offset
    /// device pointer; the dW_Q/K/V `cublasSgemmStridedBatched` call then
    /// reads from this fused buffer with `strideA = M·B`.
    d_proj_qkv_buf: CudaSlice<f32>,
    /// Parameter gradients [TLOB_TOTAL_PARAMS].
    d_params: CudaSlice<f32>,

    // ── Adam optimizer state ──────────────────────────────────────────
    adam_m: CudaSlice<f32>,
    adam_v: CudaSlice<f32>,
    adam_step: i32,
    /// Pinned host counter for device-side Adam t (device-mapped).
    t_pinned: *mut i32,
    t_dev_ptr: u64,

    // ── Gradient norm buffers ─────────────────────────────────────────
    grad_norm_buf: CudaSlice<f32>,
    grad_norm_block_sums: CudaSlice<f32>,

    // ── Kernels for grad norm + Adam ──────────────────────────────────
    grad_norm_phase1_kernel: CudaFunction,
    grad_norm_phase2_kernel: CudaFunction,
    adam_kernel: CudaFunction,

    // ── ofi_extract kernel ────────────────────────────────────────────
    /// Extracts OFI slice from states_buf row-major → col-major [32,B].
    ofi_extract_kernel: CudaFunction,

    // ── Bias gradient reduce (for zero-bias architecture, only d_params for W) ──
    bias_grad_reduce_p1_kernel: CudaFunction,
    bias_grad_reduce_p2_kernel: CudaFunction,
    bias_grad_partials_buf: CudaSlice<f32>,
    bias_grad_num_blocks: u32,
}

impl GpuTlob {
    pub(crate) fn new(
        shared_handle: Arc<PerStreamCublasHandles>,
        batch_size: usize,
    ) -> Result<Self, MLError> {
        let stream = Arc::clone(&shared_handle.stream);
        let b = batch_size;

        // Bind the classic cuBLAS handle to `stream` once at construction so
        // every subsequent `cublasSgemm_v2` call inherits the binding. The
        // handle is owned by `shared_handle`; classic_for re-binds before
        // returning, idempotent across calls.
        let _ = shared_handle.classic_for(&stream)?;

        // ── Load cubin module ────────────────────────────────────────
        let context = stream.context();
        let module = context.load_cubin(TLOB_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("tlob cubin load: {e}")))?;
        let sdp_fwd_kernel = module.load_function("tlob_sdp_forward")
            .map_err(|e| MLError::ModelError(format!("tlob_sdp_forward load: {e}")))?;
        let sdp_bwd_kernel = module.load_function("tlob_sdp_backward")
            .map_err(|e| MLError::ModelError(format!("tlob_sdp_backward load: {e}")))?;
        let write_states_kernel = module.load_function("tlob_write_states")
            .map_err(|e| MLError::ModelError(format!("tlob_write_states load: {e}")))?;
        let read_grad_kernel = module.load_function("tlob_read_grad_from_states")
            .map_err(|e| MLError::ModelError(format!("tlob_read_grad_from_states load: {e}")))?;

        // ── Load grad norm + Adam from attention backward cubin ───────
        // Reuse the same kernels as GpuAttention to avoid duplicate cubin compilation.
        static ATTN_BWD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/attention_backward_kernel.cubin"));
        let bwd_module = context.load_cubin(ATTN_BWD_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("attn_bwd cubin for tlob: {e}")))?;
        let grad_norm_phase1_kernel = bwd_module.load_function("attn_grad_norm_phase1")
            .map_err(|e| MLError::ModelError(format!("attn_grad_norm_phase1 load: {e}")))?;
        let grad_norm_phase2_kernel = bwd_module.load_function("attn_grad_norm_phase2")
            .map_err(|e| MLError::ModelError(format!("attn_grad_norm_phase2 load: {e}")))?;
        let adam_kernel = bwd_module.load_function("attn_adam_kernel")
            .map_err(|e| MLError::ModelError(format!("attn_adam_kernel load: {e}")))?;
        let bias_grad_reduce_p1_kernel = bwd_module.load_function("attn_bias_grad_reduce_p1")
            .map_err(|e| MLError::ModelError(format!("attn_bias_grad_reduce_p1 load: {e}")))?;
        let bias_grad_reduce_p2_kernel = bwd_module.load_function("attn_bias_grad_reduce_p2")
            .map_err(|e| MLError::ModelError(format!("attn_bias_grad_reduce_p2 load: {e}")))?;

        // ── ofi_extract kernel (from experience_kernels cubin) ────────
        static EXP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/experience_kernels.cubin"));
        let exp_module = context.load_cubin(EXP_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("exp cubin for tlob: {e}")))?;
        let ofi_extract_kernel = exp_module.load_function("ofi_embed_build_input")
            .map_err(|e| MLError::ModelError(format!("ofi_embed_build_input for tlob: {e}")))?;

        // ── Xavier initialization ─────────────────────────────────────
        let mut host_params = vec![0.0_f32; TLOB_TOTAL_PARAMS];
        let fan_in  = TLOB_IN  as f32;
        let fan_out = TLOB_OUT as f32;
        let xavier_std = (2.0 / (fan_in + fan_out)).sqrt();
        let mut rng = 0xcbf29ce484222325_u64.wrapping_add(b as u64);
        for p in host_params[..TLOB_W_QKV_SIZE].iter_mut() {
            rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
            let u = (rng >> 33) as f32 / (1u64 << 31) as f32;
            *p = (u - 0.5) * 3.46 * xavier_std;
        }
        // W_O starts at zero (identity-like start: output = 0 initially, learns from scratch)

        let mut params = stream.alloc_zeros::<f32>(TLOB_TOTAL_PARAMS)
            .map_err(|e| MLError::ModelError(format!("tlob params alloc: {e}")))?;
        // Mapped-pinned host staging → DtoD copy. Mirrors the canonical
        // `MappedF32Buffer` upload pattern (host_ptr writes; kernel reads
        // dev_ptr; no host-pageable HtoD).  Same shape used by
        // `gpu_attention.rs` / `gpu_iql_trainer.rs` / etc. after the
        // via_pinned cleanup landed in the codebase.
        {
            let n = host_params.len();
            if n > 0 {
                let staging = unsafe { super::mapped_pinned::MappedF32Buffer::new(n) }
                    .map_err(|e| MLError::ModelError(format!("tlob params staging alloc: {e}")))?;
                staging.write_from_slice(&host_params);
                let nbytes = n * std::mem::size_of::<f32>();
                #[allow(unsafe_code)]
                unsafe {
                    let (dst_ptr, _g) = params.device_ptr_mut(&stream);
                    cudarc::driver::result::memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream())
                        .map_err(|e| MLError::ModelError(format!("tlob params DtoD: {e}")))?;
                }
                stream.synchronize().map_err(|e| MLError::ModelError(format!("tlob params sync: {e}")))?;
            }
        }

        // ── Buffer allocation ─────────────────────────────────────────
        let db_out = TLOB_OUT * b;  // 16 × B
        let db_in  = TLOB_IN  * b; // 32 × B

        let alloc = |n: usize, label: &str| -> Result<CudaSlice<f32>, MLError> {
            stream.alloc_zeros::<f32>(n)
                .map_err(|e| MLError::ModelError(format!("tlob {label} alloc: {e}")))
        };

        let ofi_col_buf      = alloc(db_in, "ofi_col")?;
        // Fused Q/K/V buffer: one allocation, three contiguous [M, B]
        // chunks. Replaces the previous three `proj_q/k/v_buf` allocations.
        let proj_qkv_buf     = alloc(QKV_BATCH * db_out, "proj_qkv")?;
        let attn_out_buf     = alloc(db_out, "attn_out")?;
        let attn_scores_buf  = alloc(db_out, "attn_scores")?;
        let output_buf       = alloc(db_out, "output")?;
        let d_output_buf     = alloc(db_out, "d_output")?;
        let d_attn_out_buf   = alloc(db_out, "d_attn_out")?;
        // Fused d_Q/K/V gradient buffer: same layout as `proj_qkv_buf`.
        let d_proj_qkv_buf   = alloc(QKV_BATCH * db_out, "d_proj_qkv")?;
        let d_params         = alloc(TLOB_TOTAL_PARAMS, "d_params")?;
        let adam_m         = alloc(TLOB_TOTAL_PARAMS, "adam_m")?;
        let adam_v         = alloc(TLOB_TOTAL_PARAMS, "adam_v")?;

        let norm_blocks = (TLOB_TOTAL_PARAMS + 255) / 256;
        let grad_norm_buf        = alloc(1, "grad_norm")?;
        let grad_norm_block_sums = alloc(norm_blocks, "grad_norm_sums")?;

        let bias_grad_num_blocks = ((b + 255) / 256) as u32;
        let bias_grad_partials_buf = alloc(
            (bias_grad_num_blocks as usize) * TLOB_OUT,
            "bias_grad_partials",
        )?;

        // ── Adam pinned step counter ──────────────────────────────────
        let mut t_raw: *mut i32 = std::ptr::null_mut();
        unsafe {
            let status = cudarc::driver::sys::cuMemHostAlloc(
                (&mut t_raw) as *mut *mut i32 as *mut *mut std::ffi::c_void,
                std::mem::size_of::<i32>(),
                cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP,
            );
            if status != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
                return Err(MLError::ModelError(format!("tlob: cuMemHostAlloc failed: {status:?}")));
            }
            *t_raw = 0;
        }
        let mut t_dev_raw: u64 = 0;
        unsafe {
            let status = cudarc::driver::sys::cuMemHostGetDevicePointer_v2(
                &mut t_dev_raw,
                t_raw as *mut std::ffi::c_void,
                0,
            );
            if status != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
                cudarc::driver::sys::cuMemFreeHost(t_raw as *mut std::ffi::c_void);
                return Err(MLError::ModelError(format!("tlob: cuMemHostGetDevicePointer failed: {status:?}")));
            }
        }

        info!(
            batch_size = b,
            tlob_in = TLOB_IN,
            tlob_out = TLOB_OUT,
            total_params = TLOB_TOTAL_PARAMS,
            "GpuTlob initialized: OFI[32]→TLOB[16] single-head SDP, Xavier random init, trainable end-to-end"
        );

        Ok(Self {
            stream,
            shared_handle,
            batch_size: b,
            sdp_fwd_kernel,
            sdp_bwd_kernel,
            write_states_kernel,
            read_grad_kernel,
            params,
            ofi_col_buf,
            proj_qkv_buf,
            attn_out_buf,
            attn_scores_buf,
            output_buf,
            d_output_buf,
            d_attn_out_buf,
            d_proj_qkv_buf,
            d_params,
            adam_m,
            adam_v,
            adam_step: 0,
            t_pinned: t_raw,
            t_dev_ptr: t_dev_raw,
            grad_norm_buf,
            grad_norm_block_sums,
            grad_norm_phase1_kernel,
            grad_norm_phase2_kernel,
            adam_kernel,
            ofi_extract_kernel,
            bias_grad_reduce_p1_kernel,
            bias_grad_reduce_p2_kernel,
            bias_grad_partials_buf,
            bias_grad_num_blocks,
        })
    }

    /// TLOB forward pass.
    ///
    /// 1. Extract OFI slice from states_buf → ofi_col_buf [TLOB_IN, B] col-major.
    /// 2. Fused Q/K/V projections via `cublasSgemmStridedBatched(batch=3)`
    ///    into `proj_qkv_buf` (3 contiguous [M, B] chunks, single launch).
    /// 3. SDP (tlob_sdp_forward kernel) — reads Q, K, V via 3 offset device
    ///    pointers into `proj_qkv_buf`.
    /// 4. Output projection via cuBLAS.
    /// 5. Write [TLOB_OUT, B] to states_buf at SL_TLOB_START.
    ///
    /// `states_buf`: row-major [B, state_dim_padded] — both read (OFI) and written (TLOB).
    pub(crate) fn forward(
        &mut self,
        states_buf: &mut CudaSlice<f32>,
        batch_size: usize,
    ) -> Result<(), MLError> {
        let b = batch_size;
        let b_i32 = b as i32;
        let state_dim_padded_i32 = STATE_DIM_PADDED as i32;
        let blocks = ((b + 255) / 256) as u32;
        let f32_sz = std::mem::size_of::<f32>();

        // Bind classic cuBLAS handle to our stream + workspace before any GEMM.
        let classic = self.shared_handle.classic_handle_for(&self.stream)?;

        // ── Step 1: Extract OFI [SL_OFI_START..+SL_OFI_DIM] from states row-major
        // → ofi_col_buf [TLOB_IN=32, B] col-major.
        // Reuse ofi_embed_build_input kernel: reads s[SL_OFI_START..+SL_OFI_DIM] from states.
        let state_dim_i32 = STATE_DIM_PADDED as i32;
        unsafe {
            self.stream
                .launch_builder(&self.ofi_extract_kernel)
                .arg(states_buf as &CudaSlice<f32>)
                .arg(&mut self.ofi_col_buf)
                .arg(&b_i32)
                .arg(&state_dim_i32)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("tlob ofi_extract: {e}")))?;
        }

        let params_base = self.params.raw_ptr();
        let ofi_ptr  = self.ofi_col_buf.raw_ptr();

        // ── Step 2: Fused Q/K/V projections ───────────────────────────
        // Math (per-projection): col-major [TLOB_OUT, B] = op(W)[TLOB_OUT,
        // TLOB_IN] @ op(ofi)[TLOB_IN, B].  W is stored col-major
        // [TLOB_IN, TLOB_OUT] (i.e. row-major [TLOB_OUT, TLOB_IN]), so
        // transA=T reads it as [TLOB_OUT, TLOB_IN] col-major; ofi is
        // already col-major [TLOB_IN, B], transB=N.
        //
        // ── Layout convention (Fix 20, 2026-05-03) ────────────────────
        // W_Q (and W_K, W_V) are stored col-major `[K=TLOB_IN, M=TLOB_OUT]`
        // with lda=K — that is, K-rows × M-cols, "K is the fast axis".
        // Flat index `i` of `params + W_Q_OFF` therefore addresses position
        // `(k = i % K, m = i / K)`. The backward dW_Q SGEMM writes into
        // the same flat slot using the SAME [K, M] layout (see Step 4 of
        // `backward()`), so Adam's element-wise update — `params[i] -=
        // lr * d_params[i] / ...` — sees aligned (k, m) coordinates at
        // every flat index. Pre-Fix-20 the backward wrote col-major
        // `[M, K]` (ldc=M=16), which only agreed with the forward layout
        // at i=0; from i=1 onward, when M ≠ K (TLOB: M=16, K=32), Adam
        // silently updated the wrong weights. The reproduction +
        // regression coverage lives in `tlob_dw_layout_alignment_*` in
        // the inline tests module below.
        //
        // Fused dispatch: `cublasSgemmStridedBatched(batch=3)` over the
        // contiguous `[W_Q | W_K | W_V]` weight block (`strideA = M·K`)
        // and the same ofi (`strideB = 0`), writing into the single
        // `proj_qkv_buf` with `strideC = M·B`. One launch in place of
        // three; cuts per-call cuBLAS heuristic-lookup + kernel-launch
        // overhead 3× without changing the per-projection memory layout.
        let m_qkv = TLOB_OUT as i32;
        let n_qkv = b_i32;
        let k_qkv = TLOB_IN as i32;
        let stride_a_qkv = W_QKV_STRIDE_FLOATS as i64; // M*K floats per weight matrix
        let stride_c_qkv = (TLOB_OUT * b) as i64;      // M*B floats per output chunk
        sgemm_strided_batched(
            classic,
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            m_qkv, n_qkv, k_qkv,
            1.0_f32,
            params_base + (W_Q_OFF * f32_sz) as u64, k_qkv, stride_a_qkv,
            ofi_ptr, k_qkv, 0_i64,
            0.0_f32,
            self.proj_qkv_buf.raw_ptr(), m_qkv, stride_c_qkv,
            QKV_BATCH as i32,
            "fwd_qkv",
        )?;

        // ── Step 3: SDP ───────────────────────────────────────────────
        // Q, K, V live as contiguous [M, B] chunks at offsets 0, M·B,
        // 2·M·B inside `proj_qkv_buf`.  Pass each chunk as an offset
        // device pointer; the SDP kernel signature is unchanged (still
        // takes three `const float*` Q/K/V args), so its access pattern
        // is byte-identical to the pre-fusion path.
        let qkv_base = self.proj_qkv_buf.raw_ptr();
        let chunk_bytes = (TLOB_OUT * b * f32_sz) as u64;
        let proj_q_ptr: u64 = qkv_base;
        let proj_k_ptr: u64 = qkv_base + chunk_bytes;
        let proj_v_ptr: u64 = qkv_base + 2 * chunk_bytes;
        unsafe {
            self.stream
                .launch_builder(&self.sdp_fwd_kernel)
                .arg(&proj_q_ptr)
                .arg(&proj_k_ptr)
                .arg(&proj_v_ptr)
                .arg(&mut self.attn_out_buf)
                .arg(&mut self.attn_scores_buf)
                .arg(&b_i32)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("tlob sdp_fwd: {e}")))?;
        }

        // ── Step 4: Output projection ─────────────────────────────────
        // out = W_O^T @ attn_out: same pattern as Q/K/V but K = TLOB_OUT (W_O is square).
        let k_o = TLOB_OUT as i32;
        sgemm_v2(
            classic,
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            m_qkv, n_qkv, k_o,
            1.0_f32,
            params_base + (W_O_OFF * f32_sz) as u64, k_o,
            self.attn_out_buf.raw_ptr(), k_o,
            0.0_f32,
            self.output_buf.raw_ptr(), m_qkv,
            "fwd_o",
        )?;

        // ── Step 5: Scatter output into states_buf at SL_TLOB_START ──
        unsafe {
            self.stream
                .launch_builder(&self.write_states_kernel)
                .arg(&self.output_buf)
                .arg(states_buf)
                .arg(&b_i32)
                .arg(&state_dim_padded_i32)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("tlob write_states: {e}")))?;
        }
        Ok(())
    }

    /// TLOB backward pass.
    ///
    /// `d_concat`: gradient buffer [B, concat_dim] from trunk backward (bn_d_concat_buf).
    /// `concat_dim`: total concat dim = bottleneck_dim + (STATE_DIM - market_dim).
    /// `tlob_concat_off`: TLOB slot offset within concat row = bottleneck_dim + (TLOB_START - market_dim).
    pub(crate) fn backward(
        &mut self,
        d_concat: &CudaSlice<f32>,
        batch_size: usize,
        concat_dim: usize,
        tlob_concat_off: usize,
    ) -> Result<(), MLError> {
        let b = batch_size;
        let b_i32 = b as i32;
        let blocks = ((b + 255) / 256) as u32;
        let f32_sz = std::mem::size_of::<f32>();
        let inv_batch = 1.0_f32 / b as f32;
        let concat_dim_i32 = concat_dim as i32;
        let tlob_off_i32 = tlob_concat_off as i32;

        // Bind classic cuBLAS handle to our stream + workspace before any GEMM.
        let classic = self.shared_handle.classic_handle_for(&self.stream)?;

        // ── Step 1: Extract TLOB gradient from d_concat ───────────────
        // tlob_read_grad_from_states: d_concat[B, concat_dim] → d_output_buf[TLOB_OUT, B]
        unsafe {
            self.stream
                .launch_builder(&self.read_grad_kernel)
                .arg(d_concat)
                .arg(&mut self.d_output_buf)
                .arg(&b_i32)
                .arg(&concat_dim_i32)
                .arg(&tlob_off_i32)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("tlob read_grad: {e}")))?;
        }

        let params_base = self.params.raw_ptr();
        let grad_base = self.d_params.raw_ptr();

        // ── Step 2: Backward through output projection ────────────────
        // d_attn_out = W_O @ d_out:
        //   col-major [TLOB_OUT, B] = op(W_O)[TLOB_OUT, TLOB_OUT] @ op(d_out)[TLOB_OUT, B]
        // W_O is col-major [TLOB_IN_OUT, TLOB_OUT] but for output, that storage is
        // [TLOB_OUT, TLOB_OUT] (square); we want W_O treated as [TLOB_OUT, TLOB_OUT]
        // — read as ROWS=output of fwd, COLS=input of fwd. transA=N here means
        // "use W_O as-stored": col-major [TLOB_OUT, TLOB_OUT] (the fwd transposed
        // it; in bwd we want the un-transposed view, i.e. transA=N).
        // M=TLOB_OUT, N=B, K=TLOB_OUT.
        let m_o = TLOB_OUT as i32;
        let k_o = TLOB_OUT as i32;
        sgemm_v2(
            classic,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            m_o, b_i32, k_o,
            1.0_f32,
            params_base + (W_O_OFF * f32_sz) as u64, k_o,
            self.d_output_buf.raw_ptr(), k_o,
            0.0_f32,
            self.d_attn_out_buf.raw_ptr(), m_o,
            "bwd_dx_o",
        )?;

        // dW_O = d_out @ attn_out^T   col-major:
        //   [TLOB_OUT, TLOB_OUT] = [TLOB_OUT, B] @ [B, TLOB_OUT]
        //   transA=N (d_out as-stored), transB=T (attn_out -> [B, TLOB_OUT])
        //   M=TLOB_OUT, N=TLOB_OUT, K=B
        //   alpha = 1/B (sdp_bwd does NOT scale d_out, so we scale here)
        sgemm_v2(
            classic,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            m_o, k_o, b_i32,
            inv_batch,
            self.d_output_buf.raw_ptr(), m_o,
            self.attn_out_buf.raw_ptr(), m_o,
            0.0_f32,
            grad_base + (W_O_OFF * f32_sz) as u64, m_o,
            "bwd_dw_o",
        )?;

        // ── Step 3: SDP backward ──────────────────────────────────────
        // Q, K, V (forward saves) and d_Q, d_K, d_V (gradients) all live
        // as contiguous [M, B] chunks at offsets 0, M·B, 2·M·B inside the
        // fused proj_qkv / d_proj_qkv buffers. Pass each chunk via offset
        // device pointer; SDP backward kernel signature is unchanged.
        let f32_sz_u64 = f32_sz as u64;
        let chunk_bytes = (TLOB_OUT * b) as u64 * f32_sz_u64;
        let qkv_fwd_base = self.proj_qkv_buf.raw_ptr();
        let d_qkv_base = self.d_proj_qkv_buf.raw_ptr();
        let proj_q_ptr: u64 = qkv_fwd_base;
        let proj_k_ptr: u64 = qkv_fwd_base + chunk_bytes;
        let proj_v_ptr: u64 = qkv_fwd_base + 2 * chunk_bytes;
        let d_proj_q_ptr: u64 = d_qkv_base;
        let d_proj_k_ptr: u64 = d_qkv_base + chunk_bytes;
        let d_proj_v_ptr: u64 = d_qkv_base + 2 * chunk_bytes;
        unsafe {
            self.stream
                .launch_builder(&self.sdp_bwd_kernel)
                .arg(&self.d_attn_out_buf)
                .arg(&self.attn_scores_buf)
                .arg(&proj_v_ptr)
                .arg(&proj_q_ptr)
                .arg(&proj_k_ptr)
                .arg(&d_proj_q_ptr)
                .arg(&d_proj_k_ptr)
                .arg(&d_proj_v_ptr)
                .arg(&b_i32)
                .arg(&inv_batch)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("tlob sdp_bwd: {e}")))?;
        }

        // ── Step 4: Fused dW_Q/K/V backward ───────────────────────────
        // Layout convention (Fix 20, 2026-05-03):
        //
        //   Forward W_Q is stored col-major `[K, M]` (lda=K=TLOB_IN=32,
        //   shape K-rows × M-cols), set by the forward SGEMM with
        //   `transA=T` reading `params + W_Q_OFF` as that shape. Adam
        //   consumes `params[i]` and `d_params[i]` element-wise — so
        //   `d_params` MUST agree with `params` on the (k, m) coordinates
        //   addressed at each flat index. If we wrote dW_Q col-major
        //   `[M, K]` (the "obvious" `dW = d_proj @ ofi^T` form, ldc=M=16),
        //   then at flat index `i ≠ 0` the mapping would diverge:
        //     params[i]   → (k = i % 32, m = i / 32)   [forward layout]
        //     d_params[i] → (m = i % 16, k = i / 16)   [naive bwd layout]
        //   Identical only at i=0; from i=1 onward they reference different
        //   (k, m) pairs and Adam silently updates the wrong weights.
        //
        // Fix: compute `dW_Q^T = ofi @ d_proj^T` instead, so C lands in
        // col-major `[K, M]` (ldc=K=32) — bit-identical layout to the
        // forward W_Q storage. Math is unchanged (transposing both sides
        // yields the same dW gradient values, just re-laid-out so flat
        // indexing matches `params`).
        //
        // Per-projection math: `dW_X^T[K, M] = ofi[K, B] @ d_proj_X^T[B, M]`
        //   transA=N (ofi as-stored col-major [K, B]),
        //   transB=T (d_proj_X col-major [M, B] → [B, M]),
        //   M_sgemm=TLOB_IN=K, N_sgemm=TLOB_OUT=M, K_sgemm=B.
        //   alpha=1.0 — sdp_bwd already divides by B.
        //
        // Fused dispatch: `cublasSgemmStridedBatched(batch=3)`. The three
        // d_proj_X chunks are now operand B (strideB=M·B); ofi is shared
        // across all three (strideA=0); strideC=M·K is unchanged because
        // each weight slot is the same M·K floats wide regardless of
        // which axis is the leading dim.
        //
        // The inline `tlob_sgemm_parity_with_cpu_reference` test's CPU
        // reference for dW_Q/K/V was updated in lockstep so flat indices
        // line up the same way (forward [K, M] convention).
        let ofi_ptr = self.ofi_col_buf.raw_ptr();
        let m_qkv = TLOB_IN as i32;     // dW^T leading dim = K (forward layout)
        let n_qkv = TLOB_OUT as i32;    // dW^T trailing dim = M
        let k_qkv = b_i32;
        let stride_a_qkv = 0_i64;                       // ofi shared across batch
        let stride_b_qkv = (TLOB_OUT * b) as i64;       // M*B floats per d_proj chunk
        let stride_c_qkv = W_QKV_STRIDE_FLOATS as i64;  // M*K floats per weight matrix
        sgemm_strided_batched(
            classic,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            m_qkv, n_qkv, k_qkv,
            1.0_f32,
            ofi_ptr, m_qkv, stride_a_qkv,
            self.d_proj_qkv_buf.raw_ptr(), n_qkv, stride_b_qkv,
            0.0_f32,
            grad_base + (W_Q_OFF * f32_sz) as u64, m_qkv, stride_c_qkv,
            QKV_BATCH as i32,
            "bwd_dw_qkv",
        )?;

        Ok(())
    }

    /// Adam step: grad norm clip + Adam update on TLOB params.
    ///
    /// SP4 Layer B (Mech 9): `weight_clamp_max_abs` is the ISV-driven |p|
    /// bound for `attn_adam_kernel` (TLOB shares this kernel via the
    /// attention cubin). Caller reads `ISV[WEIGHT_BOUND[Attn=6]]` (TLOB
    /// co-lives with attention in the SP4 group taxonomy). 0=disabled
    /// (debug only). `weight_decay_value` is sourced from
    /// `ISV[WD_RATE[Attn=6]]` with `.max(0.0)`.
    pub(crate) fn adam_step(
        &mut self,
        lr: f32,
        max_grad_norm: f32,
        weight_clamp_max_abs: f32,
        weight_decay_value: f32,
        /* SP5 Layer B (Pearl 4): ISV-derived Adam β1/β2/ε for ParamGroup::Attn=6.
         * TLOB shares the Attn param-group taxonomy with GpuAttention; same ISV slots.
         * Caller applies Invariant 1 floors (β1≥0.5, β2≥0.9, ε≥1e-12). */
        attn_beta1: f32,
        attn_beta2: f32,
        attn_epsilon: f32,
    ) -> Result<(), MLError> {
        let tp = TLOB_TOTAL_PARAMS;
        let tp_i32 = tp as i32;

        let norm_blocks = (tp + 255) / 256;
        unsafe {
            self.stream
                .launch_builder(&self.grad_norm_phase1_kernel)
                .arg(&self.d_params)
                .arg(&mut self.grad_norm_block_sums)
                .arg(&tp_i32)
                .launch(LaunchConfig {
                    grid_dim: (norm_blocks as u32, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("tlob grad_norm phase1: {e}")))?;
        }
        let norm_blocks_i32 = norm_blocks as i32;
        unsafe {
            self.stream
                .launch_builder(&self.grad_norm_phase2_kernel)
                .arg(&self.grad_norm_block_sums)
                .arg(&mut self.grad_norm_buf)
                .arg(&norm_blocks_i32)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (1, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("tlob grad_norm phase2: {e}")))?;
        }

        self.adam_step += 1;
        unsafe { *self.t_pinned = self.adam_step; }
        let adam_blocks = (tp + 255) / 256;
        // SP5 Layer B (Pearl 4): per-group Adam β1/β2/ε threaded from
        // FusedTrainingCtx via ISV[ADAM_BETA{1,2}_BASE+6] / ISV[ADAM_EPS_BASE+6].
        // TLOB shares ParamGroup::Attn=6 with GpuAttention.
        let beta1 = attn_beta1;
        let beta2 = attn_beta2;
        let eps   = attn_epsilon;
        // SP4 Layer B (Mech 9): AdamW weight_decay sourced from
        // ISV[WD_RATE[Attn=6]] by the caller and threaded through
        // `adam_step`. `.max(0.0)` upstream enforces the AdamW
        // non-negativity contract.
        let wd: f32 = weight_decay_value;
        let t_ptr = self.t_dev_ptr;
        // SP4 Task A14: Pearl C engagement-counter args. TLOB shares
        // the `attn_adam_kernel` cubin with `GpuAttention`. For Layer A
        // we silently disable Pearl C bookkeeping on this trainer:
        // null pointer + SP4_ENGAGE_OFFSET_DISABLED. Layer B can split
        // the kernel ownership and add per-trainer engagement slots.
        let nan_flags_ptr_null: u64 = 0;
        let diag_slot_disabled: i32 = -1;
        let engage_buf_ptr_null: u64 = 0;
        let engage_off_disabled: i32 =
            crate::cuda_pipeline::SP4_ENGAGE_OFFSET_DISABLED;

        unsafe {
            self.stream
                .launch_builder(&self.adam_kernel)
                .arg(&mut self.params)
                .arg(&mut self.d_params)
                .arg(&mut self.adam_m)
                .arg(&mut self.adam_v)
                .arg(&self.grad_norm_buf)
                .arg(&lr)
                .arg(&beta1)
                .arg(&beta2)
                .arg(&eps)
                .arg(&wd)
                .arg(&max_grad_norm)
                .arg(&t_ptr)
                .arg(&tp_i32)
                .arg(&weight_clamp_max_abs) // SP3 Mech 9 (close-out v2): post-Adam |p| clamp
                // SP4 A14: TLOB shares attn_adam_kernel with GpuAttention;
                // pass null/disabled to skip per-block engagement writeback
                // (Layer A — Layer B can split the kernel if needed).
                .arg(&nan_flags_ptr_null)
                .arg(&diag_slot_disabled)
                .arg(&engage_buf_ptr_null)
                .arg(&engage_off_disabled)
                .launch(LaunchConfig {
                    grid_dim: (adam_blocks as u32, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("tlob adam: {e}")))?;
        }
        Ok(())
    }

    /// Reset Adam optimizer state (m, v, step) at fold boundary.
    ///
    /// TLOB carries its own Adam state separate from the main DQN. Without
    /// this reset the TLOB optimizer enters fold N+1 with fold N's
    /// momentum, producing oversized Adam steps in the first epochs whose
    /// effect compounds through the QKV/output projections into the trunk
    /// gradient. Closes the partial-refactor gap diagnosed in the
    /// fold-boundary audit (companion to the IQN target sync + main-DQN
    /// aux Adam resets). Pinned host counter is written through the
    /// device-mapped page so the next `adam_step` launch reads `t = 1`,
    /// matching the in-place pattern used by the main DQN trainer.
    pub(crate) fn reset_adam_state(&mut self) -> Result<(), MLError> {
        self.stream.memset_zeros(&mut self.adam_m)
            .map_err(|e| MLError::ModelError(format!("reset tlob adam_m: {e}")))?;
        self.stream.memset_zeros(&mut self.adam_v)
            .map_err(|e| MLError::ModelError(format!("reset tlob adam_v: {e}")))?;
        self.adam_step = 0;
        unsafe { *self.t_pinned = 0; }
        Ok(())
    }

    /// Copy trained weight parameters from `src` into this instance via DtoD async copy.
    ///
    /// Used to synchronise the val-path TLOB (owned by `GpuBacktestEvaluator`) with
    /// the online weights trained by `FusedTrainingCtx`.  The copy is enqueued on
    /// `self.stream`; callers should ensure `src`'s stream has finished writing params
    /// before this is called (a `stream.synchronize()` on the training stream is
    /// sufficient, which already happens at the epoch boundary).
    pub(crate) fn copy_params_from(&mut self, src: &GpuTlob) -> Result<(), MLError> {
        let byte_len = (TLOB_TOTAL_PARAMS * std::mem::size_of::<f32>()) as u64;
        let (src_ptr, _src_guard) = src.params.device_ptr(&self.stream);
        let (dst_ptr, _dst_guard) = self.params.device_ptr(&self.stream);
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr,
                src_ptr,
                byte_len as usize,
                self.stream.cu_stream(),
            )
            .map_err(|e| MLError::ModelError(format!("tlob copy_params_from: {e}")))?;
        }
        Ok(())
    }

    /// Compute the mean-max attention weight across the batch for ISV diagnostic.
    /// Returns the mean of max(attn_scores[d, b] over d) across b.
    /// Runs a minimal CPU readback — only called once per epoch at the CPU monitor site.
    pub(crate) fn mean_max_attention_weight(&self, batch_size: usize) -> f32 {
        let b = batch_size;
        let d = TLOB_OUT;
        let total = d * b;
        let mut host = vec![0.0_f32; total];
        // Synchronous DtoH for diagnostic (called once/epoch outside hot path).
        // memcpy_dtoh queues async; stream.synchronize() ensures transfer completes.
        if self.stream.memcpy_dtoh(&self.attn_scores_buf, &mut host).is_ok() {
            let _ = self.stream.synchronize();
        } else {
            return 0.0;
        }
        // attn_scores col-major [D, B]: per sample b, find max over d
        let mut sum_max = 0.0_f32;
        for bi in 0..b {
            let mut mx = 0.0_f32;
            for di in 0..d {
                let v = host[di * b + bi];
                if v > mx { mx = v; }
            }
            sum_max += mx;
        }
        if b > 0 { sum_max / b as f32 } else { 0.0 }
    }
}

impl Drop for GpuTlob {
    fn drop(&mut self) {
        if !self.t_pinned.is_null() {
            unsafe { cudarc::driver::sys::cuMemFreeHost(self.t_pinned.cast()); }
        }
    }
}

unsafe impl Send for GpuTlob {}
unsafe impl Sync for GpuTlob {}

// ── Classic cuBLAS sgemm helper ─────────────────────────────────────────────
//
// Free function (not a method) so call sites can pass the already-bound
// classic handle without taking an extra `&self` borrow that would conflict
// with the `&mut self` mutable buffer borrows on the call sites in
// `forward()` / `backward()`.

#[allow(clippy::too_many_arguments)]
fn sgemm_v2(
    handle: cublas_sys::cublasHandle_t,
    transa: cublas_sys::cublasOperation_t,
    transb: cublas_sys::cublasOperation_t,
    m: i32, n: i32, k: i32,
    alpha: f32,
    a_ptr: u64, lda: i32,
    b_ptr: u64, ldb: i32,
    beta: f32,
    c_ptr: u64, ldc: i32,
    label: &str,
) -> Result<(), MLError> {
    unsafe {
        let status = cublas_sys::cublasSgemm_v2(
            handle,
            transa,
            transb,
            m, n, k,
            &alpha as *const f32,
            a_ptr as *const f32, lda,
            b_ptr as *const f32, ldb,
            &beta as *const f32,
            c_ptr as *mut f32, ldc,
        );
        if status != cublas_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
            return Err(MLError::ModelError(format!(
                "cublasSgemm_v2 tlob {label} (m={m},n={n},k={k}): {status:?}"
            )));
        }
    }
    Ok(())
}

/// Strided-batched classic SGEMM dispatch helper.  Used to fuse the three
/// Q/K/V projections (forward) and the three dW_Q/K/V backward GEMMs into
/// a single launch each (`batch_count=3` for both).  Same numerical
/// semantics as three separate `cublasSgemm_v2` calls; numerical
/// equivalence is verified in `crates/ml/tests/tlob_qkv_fusion_equivalence.rs`.
#[allow(clippy::too_many_arguments)]
fn sgemm_strided_batched(
    handle: cublas_sys::cublasHandle_t,
    transa: cublas_sys::cublasOperation_t,
    transb: cublas_sys::cublasOperation_t,
    m: i32, n: i32, k: i32,
    alpha: f32,
    a_ptr: u64, lda: i32, stride_a: i64,
    b_ptr: u64, ldb: i32, stride_b: i64,
    beta: f32,
    c_ptr: u64, ldc: i32, stride_c: i64,
    batch_count: i32,
    label: &str,
) -> Result<(), MLError> {
    unsafe {
        let status = cublas_sys::cublasSgemmStridedBatched(
            handle,
            transa,
            transb,
            m, n, k,
            &alpha as *const f32,
            a_ptr as *const f32, lda, stride_a,
            b_ptr as *const f32, ldb, stride_b,
            &beta as *const f32,
            c_ptr as *mut f32, ldc, stride_c,
            batch_count,
        );
        if status != cublas_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
            return Err(MLError::ModelError(format!(
                "cublasSgemmStridedBatched tlob {label} \
                 (m={m},n={n},k={k},batch={batch_count}): {status:?}"
            )));
        }
    }
    Ok(())
}

// ── Test-only accessors for the cuBLAS-Lt → classic-sgemm migration test ────
//
// These expose internal buffers + state so the inline `tests` module below
// can verify TLOB's forward + backward GEMM math against a CPU reference.
// Gated on `cfg(test)` so production code cannot reach inside.
#[cfg(test)]
impl GpuTlob {
    fn dump_params(&self) -> Vec<f32> {
        let mut h = vec![0.0_f32; TLOB_TOTAL_PARAMS];
        let _ = self.stream.synchronize();
        let _ = self.stream.memcpy_dtoh(&self.params, &mut h);
        let _ = self.stream.synchronize();
        h
    }
    /// Dump just the Q chunk of the fused `proj_qkv_buf` (first M·B
    /// floats). Used by the inline parity test which only checks Q.
    fn dump_proj_q(&self) -> Vec<f32> {
        let chunk = TLOB_OUT * self.batch_size;
        let mut full = vec![0.0_f32; QKV_BATCH * chunk];
        let _ = self.stream.synchronize();
        let _ = self.stream.memcpy_dtoh(&self.proj_qkv_buf, &mut full);
        let _ = self.stream.synchronize();
        full[..chunk].to_vec()
    }
    /// Dump the full fused Q/K/V buffer (3 contiguous [M, B] chunks).
    /// Used by the cross-path numerical-equivalence test.
    fn dump_proj_qkv(&self) -> Vec<f32> {
        let chunk = TLOB_OUT * self.batch_size;
        let mut h = vec![0.0_f32; QKV_BATCH * chunk];
        let _ = self.stream.synchronize();
        let _ = self.stream.memcpy_dtoh(&self.proj_qkv_buf, &mut h);
        let _ = self.stream.synchronize();
        h
    }
    fn dump_attn_out(&self) -> Vec<f32> {
        let mut h = vec![0.0_f32; TLOB_OUT * self.batch_size];
        let _ = self.stream.synchronize();
        let _ = self.stream.memcpy_dtoh(&self.attn_out_buf, &mut h);
        let _ = self.stream.synchronize();
        h
    }
    fn dump_output(&self) -> Vec<f32> {
        let mut h = vec![0.0_f32; TLOB_OUT * self.batch_size];
        let _ = self.stream.synchronize();
        let _ = self.stream.memcpy_dtoh(&self.output_buf, &mut h);
        let _ = self.stream.synchronize();
        h
    }
    /// Dump the full fused d_Q/K/V backward gradient buffer (3
    /// contiguous [M, B] chunks at offsets 0, M·B, 2·M·B). Used by
    /// the layout-alignment regression test to reconstruct the
    /// SDP-backward output for a [K, M] CPU dW reference.
    fn dump_d_proj_qkv(&self) -> Vec<f32> {
        let chunk = TLOB_OUT * self.batch_size;
        let mut h = vec![0.0_f32; QKV_BATCH * chunk];
        let _ = self.stream.synchronize();
        let _ = self.stream.memcpy_dtoh(&self.d_proj_qkv_buf, &mut h);
        let _ = self.stream.synchronize();
        h
    }
    fn dump_d_params(&self) -> Vec<f32> {
        let mut h = vec![0.0_f32; TLOB_TOTAL_PARAMS];
        let _ = self.stream.synchronize();
        let _ = self.stream.memcpy_dtoh(&self.d_params, &mut h);
        let _ = self.stream.synchronize();
        h
    }
}

#[cfg(test)]
mod tests {
    //! TLOB forward+backward parity test against a CPU reference.
    //!
    //! Verifies that the cuBLAS-Lt → classic `cublasSgemm_v2` migration is
    //! mathematically equivalent to the previous Lt path. Builds a CPU
    //! reference of every GEMM that TLOB does (Q/K/V proj, SDP, output proj,
    //! and the four backward GEMMs) and compares to the GPU outputs within
    //! float tolerance.
    //!
    //! Run: `SQLX_OFFLINE=true cargo test -p ml --lib gpu_tlob --release`.
    //! Skipped automatically on machines without a CUDA device — the test
    //! returns early when `CudaContext::new(0)` fails.
    use super::*;
    use crate::cuda_pipeline::shared_cublas_handle::PerStreamCublasHandles;
    use cudarc::driver::CudaContext;
    use ml_core::state_layout::{OFI_START, STATE_DIM_PADDED};
    use std::sync::Arc;

    /// Tolerance for cuBLAS TF32 GEMMs against an f32 CPU reference.
    /// 1e-3 absolute is comfortably above the documented TF32 mantissa
    /// rounding (~1e-4 per element); reductions over K=64 amplify it
    /// modestly. Far below any value with semantic meaning.
    const TOL_GEMM: f32 = 2.0e-3;

    /// CPU sgemm in cuBLAS column-major convention:
    ///   C[m,n] = alpha * op(A) @ op(B) + beta * C
    /// All ldX are leading dimensions in column-major.
    #[allow(clippy::too_many_arguments)]
    fn cpu_sgemm(
        transa: cublas_sys::cublasOperation_t,
        transb: cublas_sys::cublasOperation_t,
        m: usize, n: usize, k: usize,
        alpha: f32,
        a: &[f32], lda: usize,
        b: &[f32], ldb: usize,
        beta: f32,
        c: &mut [f32], ldc: usize,
    ) {
        let trans_a = transa == cublas_sys::cublasOperation_t::CUBLAS_OP_T;
        let trans_b = transb == cublas_sys::cublasOperation_t::CUBLAS_OP_T;
        // op(A) is m×k, op(B) is k×n, C is m×n (col-major).
        for col in 0..n {
            for row in 0..m {
                let mut acc = 0.0_f32;
                for kk in 0..k {
                    let a_val = if trans_a {
                        // A stored col-major [k, m] with lda>=k, op = transpose → [m, k]
                        a[row * lda + kk]
                    } else {
                        // A stored col-major [m, k] with lda>=m, op = identity → [m, k]
                        a[kk * lda + row]
                    };
                    let b_val = if trans_b {
                        // B stored col-major [n, k] with ldb>=n, op = transpose → [k, n]
                        b[kk * ldb + col]
                    } else {
                        // B stored col-major [k, n] with ldb>=k, op = identity → [k, n]
                        b[col * ldb + kk]
                    };
                    acc += a_val * b_val;
                }
                let idx = col * ldc + row;
                c[idx] = alpha * acc + beta * c[idx];
            }
        }
    }

    fn assert_close(actual: &[f32], expected: &[f32], tol: f32, label: &str) {
        assert_eq!(actual.len(), expected.len(), "{label}: length mismatch");
        let mut max_err = 0.0_f32;
        let mut max_idx = 0;
        for i in 0..actual.len() {
            let e = (actual[i] - expected[i]).abs();
            if e > max_err {
                max_err = e;
                max_idx = i;
            }
        }
        if max_err > tol {
            panic!(
                "{label}: max abs diff {max_err:.6e} exceeds tol {tol:.6e} at idx {max_idx}: gpu={} cpu={}",
                actual[max_idx], expected[max_idx]
            );
        }
    }

    #[test]
    fn tlob_sgemm_parity_with_cpu_reference() {
        // Skip cleanly if no GPU is available (CPU-only CI nodes).
        let ctx = match CudaContext::new(0) {
            Ok(c) => c,
            Err(e) => {
                eprintln!("Skipping tlob_sgemm_parity (no CUDA device): {e}");
                return;
            }
        };
        let stream = ctx.default_stream();
        let shared = Arc::new(
            PerStreamCublasHandles::new(&stream)
                .expect("PerStreamCublasHandles::new"),
        );

        let batch_size = 64_usize;
        let mut tlob = GpuTlob::new(Arc::clone(&shared), batch_size)
            .expect("GpuTlob::new must succeed after sgemm migration — \
                     graceful-degrade is gone, init must work");

        // ── Build synthetic states_buf [B, STATE_DIM_PADDED] row-major ──
        // Only the OFI slice [SL_OFI_START..+OFI_DIM] needs to be populated;
        // everything else is zero. Use a deterministic PRNG so the test is
        // reproducible.
        let mut rng_state: u64 = 0xa3a3a3a3a3a3a3a3;
        let mut prng = || -> f32 {
            rng_state = rng_state
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            ((rng_state >> 33) as f32 / (1u64 << 31) as f32) - 0.5
        };

        let mut host_states = vec![0.0_f32; batch_size * STATE_DIM_PADDED];
        // Fill only the OFI slice; cuBLAS reads `[TLOB_IN]` from `[OFI_START..]`.
        for b in 0..batch_size {
            for k in 0..TLOB_IN {
                host_states[b * STATE_DIM_PADDED + OFI_START + k] = prng();
            }
        }

        // Upload states_buf via mapped-pinned staging (canonical pattern;
        // matches the production param-upload path above).
        let mut states_buf = stream
            .alloc_zeros::<f32>(batch_size * STATE_DIM_PADDED)
            .expect("states_buf alloc");
        {
            let n = host_states.len();
            if n > 0 {
                let staging = unsafe { super::super::mapped_pinned::MappedF32Buffer::new(n) }
                    .expect("states_buf staging alloc");
                staging.write_from_slice(&host_states);
                let nbytes = n * std::mem::size_of::<f32>();
                unsafe {
                    let (dst_ptr, _g) = states_buf.device_ptr_mut(&stream);
                    cudarc::driver::result::memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream())
                        .expect("states_buf DtoD");
                }
                stream.synchronize().expect("states_buf sync");
            }
        }

        // Snapshot weights for the CPU reference (Xavier init, randomly seeded).
        let host_params = tlob.dump_params();
        let w_q = &host_params[W_Q_OFF..W_Q_OFF + TLOB_OUT * TLOB_IN];
        let w_k = &host_params[W_K_OFF..W_K_OFF + TLOB_OUT * TLOB_IN];
        let w_v = &host_params[W_V_OFF..W_V_OFF + TLOB_OUT * TLOB_IN];
        let w_o = &host_params[W_O_OFF..W_O_OFF + TLOB_OUT * TLOB_OUT];

        // ── Forward on GPU ──
        tlob.forward(&mut states_buf, batch_size).expect("forward");
        stream.synchronize().expect("sync");

        // ── CPU reference: Q/K/V proj ──
        // OFI col-major [TLOB_IN, B]: same memory layout as host_states slice
        // [OFI_START..+TLOB_IN] across the B rows but indexed differently.
        // ofi_col[b, k] = host_states[b][OFI_START + k] = ofi_row_major[b * TLOB_IN + k]
        // In col-major (lda=TLOB_IN): ofi_col[col=b, row=k] = mem[col*lda + row]=mem[b*TLOB_IN + k]. ✓
        let mut ofi_cm = vec![0.0_f32; TLOB_IN * batch_size];
        for b in 0..batch_size {
            for k in 0..TLOB_IN {
                ofi_cm[b * TLOB_IN + k] = host_states[b * STATE_DIM_PADDED + OFI_START + k];
            }
        }

        let mut ref_q = vec![0.0_f32; TLOB_OUT * batch_size];
        let mut ref_k = vec![0.0_f32; TLOB_OUT * batch_size];
        let mut ref_v = vec![0.0_f32; TLOB_OUT * batch_size];
        for (w_, out_, label) in [
            (w_q, &mut ref_q, "fwd_q"),
            (w_k, &mut ref_k, "fwd_k"),
            (w_v, &mut ref_v, "fwd_v"),
        ] {
            cpu_sgemm(
                cublas_sys::cublasOperation_t::CUBLAS_OP_T,
                cublas_sys::cublasOperation_t::CUBLAS_OP_N,
                TLOB_OUT, batch_size, TLOB_IN,
                1.0,
                w_, TLOB_IN,
                &ofi_cm, TLOB_IN,
                0.0,
                out_, TLOB_OUT,
            );
            let _ = label;
        }
        let gpu_q = tlob.dump_proj_q();
        assert_close(&gpu_q, &ref_q, TOL_GEMM, "Q projection forward");

        // ── CPU reference: SDP attention (per sample) ──
        // attn_out_cm [TLOB_OUT, B], attn_scores_cm [TLOB_OUT, B]
        let mut ref_attn_out = vec![0.0_f32; TLOB_OUT * batch_size];
        let mut ref_attn_scores = vec![0.0_f32; TLOB_OUT * batch_size];
        let scale = 1.0_f32 / (TLOB_OUT as f32).sqrt();
        for b in 0..batch_size {
            // Read q[d], k[d], v[d] from col-major buffers: qkv[d * B + b].
            let mut q = [0.0_f32; TLOB_OUT];
            let mut k_ = [0.0_f32; TLOB_OUT];
            let mut v = [0.0_f32; TLOB_OUT];
            for d in 0..TLOB_OUT {
                q[d] = ref_q[d * batch_size + b];
                k_[d] = ref_k[d * batch_size + b];
                v[d] = ref_v[d * batch_size + b];
            }
            let mut scores = [0.0_f32; TLOB_OUT];
            let mut max_score = f32::NEG_INFINITY;
            for d in 0..TLOB_OUT {
                scores[d] = q[d] * k_[d] * scale;
                if scores[d] > max_score {
                    max_score = scores[d];
                }
            }
            let mut sum_exp = 0.0_f32;
            for d in 0..TLOB_OUT {
                scores[d] = (scores[d] - max_score).exp();
                sum_exp += scores[d];
            }
            let inv_sum = if sum_exp > 1e-12 { 1.0 / sum_exp } else { 0.0 };
            for d in 0..TLOB_OUT {
                scores[d] *= inv_sum;
                ref_attn_out[d * batch_size + b] = scores[d] * v[d];
                ref_attn_scores[d * batch_size + b] = scores[d];
            }
        }
        let gpu_attn_out = tlob.dump_attn_out();
        assert_close(&gpu_attn_out, &ref_attn_out, TOL_GEMM, "SDP attn_out");

        // ── CPU reference: output projection W_O^T @ attn_out ──
        let mut ref_output = vec![0.0_f32; TLOB_OUT * batch_size];
        cpu_sgemm(
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            TLOB_OUT, batch_size, TLOB_OUT,
            1.0,
            w_o, TLOB_OUT,
            &ref_attn_out, TLOB_OUT,
            0.0,
            &mut ref_output, TLOB_OUT,
        );
        let gpu_output = tlob.dump_output();
        assert_close(&gpu_output, &ref_output, TOL_GEMM, "Output projection forward");

        // ── Backward: invent a synthetic d_concat with TLOB at offset 0 ──
        // The simplest layout: concat_dim = TLOB_DIM, tlob_concat_off = 0.
        // The kernel reads d_concat[B, concat_dim] row-major and writes
        // d_output_buf[TLOB_OUT, B] col-major. Use random gradients to
        // exercise every backward GEMM.
        let concat_dim = TLOB_OUT;
        let tlob_concat_off = 0_usize;
        let mut host_d_concat = vec![0.0_f32; batch_size * concat_dim];
        for v in host_d_concat.iter_mut() {
            *v = prng() * 0.1;
        }
        let mut d_concat = stream
            .alloc_zeros::<f32>(batch_size * concat_dim)
            .expect("d_concat alloc");
        {
            let n = host_d_concat.len();
            if n > 0 {
                let staging = unsafe { super::super::mapped_pinned::MappedF32Buffer::new(n) }
                    .expect("d_concat staging alloc");
                staging.write_from_slice(&host_d_concat);
                let nbytes = n * std::mem::size_of::<f32>();
                unsafe {
                    let (dst_ptr, _g) = d_concat.device_ptr_mut(&stream);
                    cudarc::driver::result::memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream())
                        .expect("d_concat DtoD");
                }
                stream.synchronize().expect("d_concat sync");
            }
        }

        tlob
            .backward(&d_concat, batch_size, concat_dim, tlob_concat_off)
            .expect("backward");
        stream.synchronize().expect("sync");

        // ── CPU reference of backward ──
        // d_output_cm [TLOB_OUT, B]: copy from row-major d_concat at offset 0.
        let mut ref_d_output = vec![0.0_f32; TLOB_OUT * batch_size];
        for b in 0..batch_size {
            for d in 0..TLOB_OUT {
                ref_d_output[d * batch_size + b] =
                    host_d_concat[b * concat_dim + tlob_concat_off + d];
            }
        }

        // d_attn_out = W_O @ d_output (transA=N, transB=N, M=N=TLOB_OUT, N=B, K=TLOB_OUT)
        let mut ref_d_attn_out = vec![0.0_f32; TLOB_OUT * batch_size];
        cpu_sgemm(
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            TLOB_OUT, batch_size, TLOB_OUT,
            1.0,
            w_o, TLOB_OUT,
            &ref_d_output, TLOB_OUT,
            0.0,
            &mut ref_d_attn_out, TLOB_OUT,
        );

        // dW_O = (1/B) * d_output @ attn_out^T  col-major [TLOB_OUT, TLOB_OUT]
        let inv_batch = 1.0_f32 / batch_size as f32;
        let mut ref_dw_o = vec![0.0_f32; TLOB_OUT * TLOB_OUT];
        cpu_sgemm(
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            TLOB_OUT, TLOB_OUT, batch_size,
            inv_batch,
            &ref_d_output, TLOB_OUT,
            &ref_attn_out, TLOB_OUT,
            0.0,
            &mut ref_dw_o, TLOB_OUT,
        );

        // SDP backward (per sample, with inv_batch scaling — see tlob_kernel.cu)
        // d_proj_v[d,b] = d_attn_out[d,b] * scores[d,b] * inv_batch
        // d_pre[d] = scores[d] * (d_attn_out[d] * v[d] - dot(scores, d_attn_out * v))
        // d_proj_q[d,b] = d_pre[d] * k[d] * scale * inv_batch
        // d_proj_k[d,b] = d_pre[d] * q[d] * scale * inv_batch
        let mut ref_d_q = vec![0.0_f32; TLOB_OUT * batch_size];
        let mut ref_d_k = vec![0.0_f32; TLOB_OUT * batch_size];
        let mut ref_d_v = vec![0.0_f32; TLOB_OUT * batch_size];
        for b in 0..batch_size {
            let mut s = [0.0_f32; TLOB_OUT];
            let mut v = [0.0_f32; TLOB_OUT];
            let mut q = [0.0_f32; TLOB_OUT];
            let mut k_ = [0.0_f32; TLOB_OUT];
            let mut dout = [0.0_f32; TLOB_OUT];
            for d in 0..TLOB_OUT {
                s[d] = ref_attn_scores[d * batch_size + b];
                v[d] = ref_v[d * batch_size + b];
                q[d] = ref_q[d * batch_size + b];
                k_[d] = ref_k[d * batch_size + b];
                dout[d] = ref_d_attn_out[d * batch_size + b];
            }
            for d in 0..TLOB_OUT {
                ref_d_v[d * batch_size + b] = dout[d] * s[d] * inv_batch;
            }
            let mut ds = [0.0_f32; TLOB_OUT];
            let mut dot_s_ds = 0.0_f32;
            for d in 0..TLOB_OUT {
                ds[d] = dout[d] * v[d];
                dot_s_ds += s[d] * ds[d];
            }
            let mut d_pre = [0.0_f32; TLOB_OUT];
            for d in 0..TLOB_OUT {
                d_pre[d] = s[d] * (ds[d] - dot_s_ds);
            }
            for d in 0..TLOB_OUT {
                ref_d_q[d * batch_size + b] = d_pre[d] * k_[d] * scale * inv_batch;
                ref_d_k[d * batch_size + b] = d_pre[d] * q[d] * scale * inv_batch;
            }
        }

        // dW_Q^T = ofi @ d_proj_q^T  (alpha=1, sdp_bwd already applied inv_batch).
        // Fix 20 (2026-05-03): GPU writes dW_Q col-major [K=TLOB_IN, M=TLOB_OUT]
        // matching the forward W_Q layout (so Adam's element-wise consumption
        // sees aligned (k, m) coordinates at every flat index). The CPU
        // reference computes dW^T directly in the same layout: M_sgemm=K,
        // N_sgemm=M, op_A=N (ofi col-major [K, B]), op_B=T (d_proj col-major
        // [M, B] viewed as [B, M]), ldc=K=TLOB_IN.
        let mut ref_dw_q = vec![0.0_f32; TLOB_OUT * TLOB_IN];
        let mut ref_dw_k = vec![0.0_f32; TLOB_OUT * TLOB_IN];
        let mut ref_dw_v = vec![0.0_f32; TLOB_OUT * TLOB_IN];
        for (d_proj_, dw_, _label) in [
            (&ref_d_q, &mut ref_dw_q, "dW_Q"),
            (&ref_d_k, &mut ref_dw_k, "dW_K"),
            (&ref_d_v, &mut ref_dw_v, "dW_V"),
        ] {
            cpu_sgemm(
                cublas_sys::cublasOperation_t::CUBLAS_OP_N,
                cublas_sys::cublasOperation_t::CUBLAS_OP_T,
                TLOB_IN, TLOB_OUT, batch_size,
                1.0,
                &ofi_cm, TLOB_IN,
                d_proj_, TLOB_OUT,
                0.0,
                dw_, TLOB_IN,
            );
        }

        let gpu_d_params = tlob.dump_d_params();
        let gpu_dw_q = &gpu_d_params[W_Q_OFF..W_Q_OFF + TLOB_OUT * TLOB_IN];
        let gpu_dw_k = &gpu_d_params[W_K_OFF..W_K_OFF + TLOB_OUT * TLOB_IN];
        let gpu_dw_v = &gpu_d_params[W_V_OFF..W_V_OFF + TLOB_OUT * TLOB_IN];
        let gpu_dw_o = &gpu_d_params[W_O_OFF..W_O_OFF + TLOB_OUT * TLOB_OUT];

        assert_close(gpu_dw_q, &ref_dw_q, TOL_GEMM, "dW_Q backward");
        assert_close(gpu_dw_k, &ref_dw_k, TOL_GEMM, "dW_K backward");
        assert_close(gpu_dw_v, &ref_dw_v, TOL_GEMM, "dW_V backward");
        assert_close(gpu_dw_o, &ref_dw_o, TOL_GEMM, "dW_O backward");
    }

    /// Tolerance for numerical equivalence between the new fused
    /// `cublasSgemmStridedBatched(batch=3)` path and the prior 3-SGEMM
    /// reference path.
    ///
    /// **Why not 1e-5?** The shared classic-cuBLAS handle is configured
    /// with `CUBLAS_TF32_TENSOR_OP_MATH` (see `shared_cublas_handle::
    /// create_handles_and_workspace`). TF32 truncates the multiplicand
    /// mantissa from 23 → 10 bits, so each multiply has ~1e-3 relative
    /// error and a K=32 reduction can amplify to ~5e-3 absolute. The
    /// strided-batched call and the per-call `cublasSgemm_v2` path can
    /// pick different internal algos (different SM tile sizes, different
    /// TF32 reduction order), which is the dominant source of bit-level
    /// disagreement. Both paths produce mathematically equivalent
    /// results *within TF32 precision* — the inline parity test
    /// (`tlob_sgemm_parity_with_cpu_reference`) uses the same `2.0e-3`
    /// against an f32 CPU reference for the same reason.
    ///
    /// What this catches: silent layout/stride/offset bugs (which would
    /// produce O(1) deltas, not O(1e-3)). What it does NOT catch: a
    /// re-introduction of TF32 in a different reduction order — that's
    /// fine, both fused and reference paths are TF32 and the fusion is
    /// not supposed to change Q/K/V values to better-than-TF32 accuracy.
    const TOL_FUSION: f32 = 2.0e-3;

    /// Reference implementation of the OLD 3-SGEMM Q/K/V projection path,
    /// kept private to this test module so the equivalence test can compare
    /// the new fused path against it.  Mirrors what `forward()` did before
    /// the strided-batched fusion: three separate `cublasSgemm_v2` calls
    /// writing into three separate `[M, B]` col-major buffers.  Returns the
    /// concatenated `[Q | K | V]` flat layout (3·M·B floats), matching the
    /// fused `proj_qkv_buf` layout for direct elementwise comparison.
    ///
    /// The returned `Duration` covers only the three `sgemm_v2` calls
    /// (warmup excluded, stream sync included for completion), which is
    /// the apples-to-apples comparator for the fused-only timing in
    /// `time_fused_qkv_only` below.
    fn run_three_sgemm_qkv_reference(
        stream: &Arc<cudarc::driver::CudaStream>,
        shared: &Arc<PerStreamCublasHandles>,
        params_ptr: u64,
        ofi_ptr: u64,
        batch_size: usize,
        n_iters: usize,
    ) -> (Vec<f32>, std::time::Duration) {
        let f32_sz = std::mem::size_of::<f32>();
        let m_qkv = TLOB_OUT as i32;
        let n_qkv = batch_size as i32;
        let k_qkv = TLOB_IN as i32;
        let chunk = TLOB_OUT * batch_size;

        let proj_q = stream.alloc_zeros::<f32>(chunk).expect("proj_q alloc");
        let proj_k = stream.alloc_zeros::<f32>(chunk).expect("proj_k alloc");
        let proj_v = stream.alloc_zeros::<f32>(chunk).expect("proj_v alloc");

        let classic = shared.classic_handle_for(stream).expect("classic handle");

        let do_3_sgemms = |q_dst: u64, k_dst: u64, v_dst: u64| {
            for (w_off_floats, dst, label) in [
                (W_Q_OFF, q_dst, "ref_q"),
                (W_K_OFF, k_dst, "ref_k"),
                (W_V_OFF, v_dst, "ref_v"),
            ] {
                super::sgemm_v2(
                    classic,
                    cublas_sys::cublasOperation_t::CUBLAS_OP_T,
                    cublas_sys::cublasOperation_t::CUBLAS_OP_N,
                    m_qkv, n_qkv, k_qkv,
                    1.0,
                    params_ptr + (w_off_floats * f32_sz) as u64, k_qkv,
                    ofi_ptr, k_qkv,
                    0.0,
                    dst, m_qkv,
                    label,
                ).expect("ref sgemm");
            }
        };

        // Warm-up: cuBLAS lazy heuristic / first-call cost.
        do_3_sgemms(proj_q.raw_ptr(), proj_k.raw_ptr(), proj_v.raw_ptr());
        stream.synchronize().expect("sync warmup");

        let start = std::time::Instant::now();
        for _ in 0..n_iters {
            do_3_sgemms(proj_q.raw_ptr(), proj_k.raw_ptr(), proj_v.raw_ptr());
        }
        stream.synchronize().expect("sync ref");
        let elapsed = start.elapsed();

        let mut out = vec![0.0_f32; QKV_BATCH * chunk];
        stream.memcpy_dtoh(&proj_q, &mut out[..chunk]).expect("dtoh q");
        stream.memcpy_dtoh(&proj_k, &mut out[chunk..2*chunk]).expect("dtoh k");
        stream.memcpy_dtoh(&proj_v, &mut out[2*chunk..3*chunk]).expect("dtoh v");
        stream.synchronize().expect("sync dtoh");
        (out, elapsed)
    }

    /// Time the fused `cublasSgemmStridedBatched(batch=3)` path in isolation
    /// (no OFI extract / SDP / output projection / write-states).  This
    /// isolates the QKV-SGEMM cost so the fused-vs-3-call comparison is
    /// apples-to-apples with `run_three_sgemm_qkv_reference`.
    fn time_fused_qkv_only(
        stream: &Arc<cudarc::driver::CudaStream>,
        shared: &Arc<PerStreamCublasHandles>,
        params_ptr: u64,
        ofi_ptr: u64,
        batch_size: usize,
        n_iters: usize,
    ) -> std::time::Duration {
        let f32_sz = std::mem::size_of::<f32>();
        let m_qkv = TLOB_OUT as i32;
        let n_qkv = batch_size as i32;
        let k_qkv = TLOB_IN as i32;
        let chunk = TLOB_OUT * batch_size;
        let proj_qkv = stream.alloc_zeros::<f32>(QKV_BATCH * chunk).expect("proj_qkv alloc");
        let classic = shared.classic_handle_for(stream).expect("classic handle");
        let stride_a_qkv = W_QKV_STRIDE_FLOATS as i64;
        let stride_c_qkv = (TLOB_OUT * batch_size) as i64;

        let do_fused = || {
            super::sgemm_strided_batched(
                classic,
                cublas_sys::cublasOperation_t::CUBLAS_OP_T,
                cublas_sys::cublasOperation_t::CUBLAS_OP_N,
                m_qkv, n_qkv, k_qkv,
                1.0,
                params_ptr + (W_Q_OFF * f32_sz) as u64, k_qkv, stride_a_qkv,
                ofi_ptr, k_qkv, 0_i64,
                0.0,
                proj_qkv.raw_ptr(), m_qkv, stride_c_qkv,
                QKV_BATCH as i32,
                "fused_timing",
            ).expect("fused sgemm");
        };

        // Warm-up.
        do_fused();
        stream.synchronize().expect("sync warmup");

        let start = std::time::Instant::now();
        for _ in 0..n_iters {
            do_fused();
        }
        stream.synchronize().expect("sync fused");
        start.elapsed()
    }

    /// Equivalence test for the Q/K/V SGEMM fusion.
    ///
    /// Verifies that the new `cublasSgemmStridedBatched(batch=3)` path
    /// produces bit-equivalent Q/K/V outputs to three separate
    /// `cublasSgemm_v2` calls, on the same weights and the same OFI
    /// input.  Also reports a microbenchmark of fused-vs-3-call latency
    /// so the operator can confirm the fusion actually helps in their
    /// build (per task design: revert if it doesn't).
    ///
    /// Skipped on CPU-only nodes (returns early when no CUDA device).
    /// Marked `#[ignore]` so it's only run on machines with a GPU; pass
    /// `--ignored` to `cargo test` to execute.
    #[test]
    #[ignore = "requires GPU"]
    fn tlob_qkv_fusion_equivalence() {
        let ctx = match CudaContext::new(0) {
            Ok(c) => c,
            Err(e) => {
                eprintln!("Skipping tlob_qkv_fusion_equivalence (no CUDA device): {e}");
                return;
            }
        };
        let stream = ctx.default_stream();
        let shared = Arc::new(
            PerStreamCublasHandles::new(&stream)
                .expect("PerStreamCublasHandles::new"),
        );

        // Use a representative DQN training batch size — large enough that
        // any per-call launch overhead difference is measurable.
        let batch_size = 256_usize;
        let mut tlob = GpuTlob::new(Arc::clone(&shared), batch_size)
            .expect("GpuTlob::new");

        // Same synthetic states_buf construction as the inline parity test.
        let mut rng_state: u64 = 0xa3a3a3a3a3a3a3a3;
        let mut prng = || -> f32 {
            rng_state = rng_state
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            ((rng_state >> 33) as f32 / (1u64 << 31) as f32) - 0.5
        };
        let mut host_states = vec![0.0_f32; batch_size * STATE_DIM_PADDED];
        for b in 0..batch_size {
            for k in 0..TLOB_IN {
                host_states[b * STATE_DIM_PADDED + OFI_START + k] = prng();
            }
        }
        let mut states_buf = stream
            .alloc_zeros::<f32>(batch_size * STATE_DIM_PADDED)
            .expect("states_buf alloc");
        {
            let n = host_states.len();
            if n > 0 {
                let staging = unsafe { super::super::mapped_pinned::MappedF32Buffer::new(n) }
                    .expect("states_buf staging alloc");
                staging.write_from_slice(&host_states);
                let nbytes = n * std::mem::size_of::<f32>();
                unsafe {
                    let (dst_ptr, _g) = states_buf.device_ptr_mut(&stream);
                    cudarc::driver::result::memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream())
                        .expect("states_buf DtoD");
                }
                stream.synchronize().expect("states_buf sync");
            }
        }

        // Populate ofi_col_buf via one fused forward (the OFI-extract
        // kernel is shared with the production path).  After this both
        // `tlob.ofi_col_buf` and `tlob.proj_qkv_buf` hold the fused-path
        // output for the equivalence check.
        tlob.forward(&mut states_buf, batch_size).expect("forward fused");
        stream.synchronize().expect("sync forward");
        let fused_qkv = tlob.dump_proj_qkv();

        // ── Reference 3-SGEMM path: read ofi_col_buf (already populated)
        //    and run 3 separate SGEMMs into 3 separate buffers, dumped
        //    into the [Q | K | V] flat layout for direct comparison. ──
        let params_ptr = tlob.params.raw_ptr();
        let ofi_ptr = tlob.ofi_col_buf.raw_ptr();

        // Numerical equivalence run: 1 iteration of each path, value
        // comparison.  The latency microbenchmark below uses many
        // iterations to amortise per-call timing noise.
        let (ref_qkv, _) = run_three_sgemm_qkv_reference(
            &stream, &shared, params_ptr, ofi_ptr, batch_size, 1,
        );

        let chunk = TLOB_OUT * batch_size;
        let max_abs_diff = |a: &[f32], b: &[f32]| -> f32 {
            a.iter().zip(b.iter())
                .map(|(x, y)| (x - y).abs())
                .fold(0.0_f32, f32::max)
        };
        let q_diff = max_abs_diff(&fused_qkv[..chunk], &ref_qkv[..chunk]);
        let k_diff = max_abs_diff(&fused_qkv[chunk..2*chunk], &ref_qkv[chunk..2*chunk]);
        let v_diff = max_abs_diff(&fused_qkv[2*chunk..3*chunk], &ref_qkv[2*chunk..3*chunk]);
        let max_diff = q_diff.max(k_diff).max(v_diff);

        // ── Latency microbenchmark: time JUST the QKV SGEMM portion in
        //    both paths over many iterations so the per-call overhead is
        //    measurable above timing noise.  Both paths run the same
        //    cuBLAS handle on the same stream after a warm-up. ──
        const N_ITERS: usize = 200;
        let fused_elapsed = time_fused_qkv_only(
            &stream, &shared, params_ptr, ofi_ptr, batch_size, N_ITERS,
        );
        let (_, ref_elapsed) = run_three_sgemm_qkv_reference(
            &stream, &shared, params_ptr, ofi_ptr, batch_size, N_ITERS,
        );
        let fused_per = fused_elapsed / N_ITERS as u32;
        let ref_per = ref_elapsed / N_ITERS as u32;

        eprintln!(
            "[tlob_qkv_fusion_equivalence] batch={batch_size}\n  \
             max_abs_diff: Q={q_diff:.3e} K={k_diff:.3e} V={v_diff:.3e} (tol={TOL_FUSION:.0e})\n  \
             QKV-SGEMM latency over {N_ITERS} iters:\n    \
             fused (1× cublasSgemmStridedBatched batch=3): {fused_elapsed:?} total \
             → {fused_per:?}/call\n    \
             ref   (3× cublasSgemm_v2 back-to-back):       {ref_elapsed:?} total \
             → {ref_per:?}/call\n  \
             speedup: {:.2}x",
            ref_per.as_nanos() as f64 / fused_per.as_nanos().max(1) as f64,
        );

        assert!(
            max_diff <= TOL_FUSION,
            "QKV fusion not numerically equivalent (within TF32 tolerance): \
             max abs diff {max_diff:.6e} exceeds tolerance {TOL_FUSION:.0e} \
             (Q={q_diff:.3e}, K={k_diff:.3e}, V={v_diff:.3e})",
        );
    }

    // ─────────────────────────────────────────────────────────────────
    // Fix 20 — dW_Q/K/V layout-alignment tests (2026-05-03)
    //
    // These tests exercise the cross-direction (forward write +
    // backward write + Adam consume) flat-index alignment that the
    // existing parity test cannot, because the parity test compares
    // each direction against its own col-major reference and so
    // never observes whether the two directions agree on the
    // (k, m) -> flat-index mapping.
    //
    // Pre-Fix-20: forward wrote W_Q col-major [K, M] (lda=K=32),
    // backward wrote dW_Q col-major [M, K] (ldc=M=16).  At any
    // flat index `i ≠ 0` the two directions referenced different
    // (k, m) coordinates, and Adam's element-wise consumption
    // silently updated the wrong weights.
    //
    // Post-fix: backward computes `dW^T = ofi @ d_proj^T` so C
    // lands col-major [K, M] (ldc=K=32) — bit-identical layout
    // to the forward W_Q storage.
    // ─────────────────────────────────────────────────────────────────

    /// Hand-crafted SGEMM-only reproduction.  Uses `sgemm_strided_batched`
    /// directly (the same helper as production) with controlled inputs
    /// (sentinel values 0.0 and 1.0) so the layout claim can be checked
    /// numerically without going through SDP.  Pre-fix would write `1.0`
    /// at flat dW_Q index 16 (the [M, K] convention's slot for position
    /// (m=0, k=1)); post-fix writes `1.0` at flat index 32 — the [K, M]
    /// convention's slot for the SAME logical position (m=0, k=1) — which
    /// matches the forward layout's flat slot for params[m=0, k=1] (also
    /// at flat index 32 = m*K + k = 0*32 + 1... wait, that's index 1).
    ///
    /// Working out the slot mapping carefully (M=16, K=32):
    ///   Forward W_Q col-major [K, M] with lda=K=32: position (k, m)
    ///     lives at flat index `m * K + k`. So (m=0, k=1) → flat 1.
    ///     Wait — that's the OPPOSITE of "K is fast axis" if we read
    ///     col-major [K, M] as K-rows × M-cols. Re-checking:
    ///
    ///     Col-major [rows=R, cols=C] with lda>=R: element (r, c) at
    ///     flat `c * lda + r`. So col-major [K=32, M=16] with lda=32:
    ///     position (r=k, c=m) → flat `m * 32 + k`. ✓
    ///     ⇒ For (m=0, k=1): flat = 0*32 + 1 = 1.
    ///     ⇒ For (m=0, k=16): flat = 0*32 + 16 = 16.
    ///
    ///   Backward (broken) dW_Q col-major [M=16, K=32] with ldc=16:
    ///     element (r=m, c=k) → flat `k * 16 + m`.
    ///     ⇒ For (m=0, k=1): flat = 1*16 + 0 = 16.
    ///     ⇒ For (m=0, k=16): flat = 16*16 + 0 = 256 (out of bounds —
    ///        only K=32 cols ⇒ valid range [0, M*K=512)).
    ///
    /// So the same logical position (m=0, k=1) lives at flat 1 in
    /// W_Q (forward) but at flat 16 in dW_Q (broken backward) —
    /// guaranteeing Adam updates W_Q[m=0, k=16] (flat 16, which has
    /// gradient 0 in our setup) while the position with the actual
    /// gradient (W_Q[m=0, k=1], flat 1) is left untouched.
    ///
    /// Post-fix: backward writes col-major [K=32, M=16] with ldc=32:
    /// element (r=k, c=m) → flat `m * 32 + k`. Position (m=0, k=1) ⇒
    /// flat 1 — matches forward layout. Adam updates the correct slot.
    fn run_dw_q_layout_probe(
        ctx: &Arc<cudarc::driver::CudaContext>,
        stream: &Arc<cudarc::driver::CudaStream>,
        shared: &Arc<PerStreamCublasHandles>,
        batch_size: usize,
    ) -> (Vec<f32>, Vec<f32>) {
        // Allocate the three operand buffers we need: ofi[K, B] col-major,
        // d_proj[M, B] col-major, dW[M*K] flat.
        let _ = ctx; // Reserved for future per-context buffer alloc patterns.
        let m = TLOB_OUT;
        let k = TLOB_IN;

        let mut host_ofi = vec![0.0_f32; k * batch_size];
        // Set ofi[k=1, b=0] = 1.0; everything else 0.
        host_ofi[0 * k + 1] = 1.0; // col=b=0, row=k=1 → flat = 0*K + 1 = 1
        // Set ofi[k=16, b=0] = 0 explicitly (ensures the broken-layout
        // "slot" at flat dW=16 has zero gradient in the broken path —
        // making the reproduction's "wrong slot updated" claim sharper).
        // (Already 0 from vec![0.0_f32; ...], left as documentation.)

        let mut host_d_proj = vec![0.0_f32; m * batch_size];
        // Set d_proj[m=0, b=0] = 1.0; everything else 0.
        host_d_proj[0 * m + 0] = 1.0; // col=b=0, row=m=0 → flat = 0*M + 0 = 0

        // Upload via mapped-pinned (the canonical pattern; matches all
        // other tests in this file).  For brevity, reuse the same staging
        // helper as the parity test.
        let upload = |stream: &Arc<cudarc::driver::CudaStream>,
                      host: &[f32]|
         -> CudaSlice<f32> {
            let mut dev = stream.alloc_zeros::<f32>(host.len()).expect("alloc");
            let staging = unsafe { super::super::mapped_pinned::MappedF32Buffer::new(host.len()) }
                .expect("staging alloc");
            staging.write_from_slice(host);
            let nbytes = host.len() * std::mem::size_of::<f32>();
            unsafe {
                let (dst_ptr, _g) = dev.device_ptr_mut(stream);
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(),
                )
                .expect("DtoD");
            }
            stream.synchronize().expect("sync");
            dev
        };
        let ofi_dev = upload(stream, &host_ofi);
        let d_proj_dev = upload(stream, &host_d_proj);

        // Allocate two dW slots: one for the "broken" pre-fix layout
        // ([M, K] ldc=M), one for the "fixed" post-fix layout
        // ([K, M] ldc=K).  Both are M*K floats; only the layout differs.
        let dw_broken = stream.alloc_zeros::<f32>(m * k).expect("dw_broken alloc");
        let dw_fixed = stream.alloc_zeros::<f32>(m * k).expect("dw_fixed alloc");

        let classic = shared.classic_handle_for(stream).expect("classic handle");

        // Broken layout (pre-fix): C = d_proj @ ofi^T → [M, K], ldc=M.
        super::sgemm_strided_batched(
            classic,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            m as i32, k as i32, batch_size as i32,
            1.0,
            d_proj_dev.raw_ptr(), m as i32, 0_i64,
            ofi_dev.raw_ptr(), k as i32, 0_i64,
            0.0,
            dw_broken.raw_ptr(), m as i32, 0_i64,
            1, // single batch — exercising layout, not fusion.
            "probe_broken",
        )
        .expect("broken sgemm");

        // Fixed layout (post-fix): C = ofi @ d_proj^T → [K, M], ldc=K.
        super::sgemm_strided_batched(
            classic,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            k as i32, m as i32, batch_size as i32,
            1.0,
            ofi_dev.raw_ptr(), k as i32, 0_i64,
            d_proj_dev.raw_ptr(), m as i32, 0_i64,
            0.0,
            dw_fixed.raw_ptr(), k as i32, 0_i64,
            1,
            "probe_fixed",
        )
        .expect("fixed sgemm");

        let mut h_broken = vec![0.0_f32; m * k];
        let mut h_fixed = vec![0.0_f32; m * k];
        stream.memcpy_dtoh(&dw_broken, &mut h_broken).expect("dtoh broken");
        stream.memcpy_dtoh(&dw_fixed, &mut h_fixed).expect("dtoh fixed");
        stream.synchronize().expect("sync probe");
        (h_broken, h_fixed)
    }

    /// Phase-1 reproduction: prove the [M, K] vs [K, M] layout
    /// mismatch is a real bug, not a self-consistent transposed
    /// view.  Builds the broken and fixed paths side-by-side
    /// using the SAME inputs and asserts that the gradient of
    /// position (m=0, k=1) lands at DIFFERENT flat indices in the
    /// two layouts — and that the broken layout's flat index
    /// would be consumed by Adam against `params[16]` which
    /// addresses position (m=0, k=16) under the forward layout
    /// (different (k, m) ⇒ wrong-weight update).
    #[test]
    #[ignore = "requires GPU"]
    fn tlob_dw_layout_alignment_repro() {
        let ctx = match CudaContext::new(0) {
            Ok(c) => c,
            Err(e) => {
                eprintln!("Skipping tlob_dw_layout_alignment_repro (no CUDA device): {e}");
                return;
            }
        };
        let stream = ctx.default_stream();
        let shared = Arc::new(
            PerStreamCublasHandles::new(&stream)
                .expect("PerStreamCublasHandles::new"),
        );
        let batch_size = 4_usize; // tiny; we only set b=0 anyway.
        let (h_broken, h_fixed) = run_dw_q_layout_probe(&ctx, &stream, &shared, batch_size);

        let m = TLOB_OUT; // 16
        let k = TLOB_IN;  // 32

        // ── Logical position (m=0, k=1) — the only place the gradient
        //    `d_proj[m=0, b=0] * ofi[k=1, b=0] = 1 * 1 = 1` should land. ──
        // Broken layout (col-major [M, K], ldc=M): flat = k * M + m = 1*16 + 0 = 16.
        let broken_flat_for_m0k1 = 1 * m + 0; // 16
        // Fixed layout (col-major [K, M], ldc=K): flat = m * K + k = 0*32 + 1 = 1.
        let fixed_flat_for_m0k1 = 0 * k + 1; // 1

        let broken_at_16 = h_broken[broken_flat_for_m0k1];
        let broken_at_1 = h_broken[1];
        let fixed_at_1 = h_fixed[fixed_flat_for_m0k1];
        let fixed_at_16 = h_fixed[16];

        // Sanity: the gradient lands somewhere with magnitude 1.0 ± TF32 noise.
        assert!(
            (broken_at_16 - 1.0).abs() < 1e-4,
            "broken layout: expected dW[flat=16] ≈ 1.0 (position (m=0, k=1) under [M, K] convention), got {broken_at_16}"
        );
        assert!(
            broken_at_1.abs() < 1e-6,
            "broken layout: expected dW[flat=1] ≈ 0 (no gradient at the [M, K]'s (m=1, k=0) slot), got {broken_at_1}"
        );
        assert!(
            (fixed_at_1 - 1.0).abs() < 1e-4,
            "fixed layout: expected dW[flat=1] ≈ 1.0 (position (m=0, k=1) under [K, M] convention), got {fixed_at_1}"
        );
        assert!(
            fixed_at_16.abs() < 1e-6,
            "fixed layout: expected dW[flat=16] ≈ 0 (no gradient at the [K, M]'s (m=0, k=16) slot), got {fixed_at_16}"
        );

        // ── Cross-layout disagreement: at flat index 16, the two layouts
        //    disagree by O(1) — exactly the "layout/stride/offset bug" delta
        //    the existing tests' TF32 tolerances would never permit if both
        //    paths were self-consistent. ──
        let cross_delta_at_16 = (broken_at_16 - fixed_at_16).abs();
        assert!(
            cross_delta_at_16 > 0.5,
            "layout-bug delta at flat index 16 should be O(1): broken={broken_at_16}, fixed={fixed_at_16}"
        );

        // ── Bug confirmation: at flat 16, params[i] under forward [K, M]
        //    addresses logical (k=16, m=0) — forward layout puts position
        //    (k, m) at flat `m*K + k`, so flat 16 ⇒ (k=16, m=0).  The broken
        //    backward writes dW for position (m=0, k=1) into the same flat
        //    slot.  Adam would update W_Q[m=0, k=16] using the gradient of
        //    W_Q[m=0, k=1] — exactly the silent corruption. ──
        let _adam_consumes_position_m0_k16_with_grad_for_m0_k1 = {
            // No assertion needed: the check above (broken_at_16 ≠ 0,
            // expected position-(m=0, k=1) gradient) IS the proof.  This
            // binding documents the corruption pathway.
            (broken_at_16, broken_flat_for_m0k1)
        };

        eprintln!(
            "[tlob_dw_layout_alignment_repro] M={m}, K={k}, batch={batch_size}\n  \
             broken layout (pre-fix [M, K]): dW[flat=1]={broken_at_1:.3e}, dW[flat=16]={broken_at_16:.3e}\n  \
             fixed  layout (post-fix [K, M]): dW[flat=1]={fixed_at_1:.3e}, dW[flat=16]={fixed_at_16:.3e}\n  \
             O(1) cross-layout delta at flat 16: {cross_delta_at_16:.3e}\n  \
             ⇒ Pre-fix Adam would update params[16] (forward (k=16, m=0)) \
             using d_params[16] (broken-bwd (m=0, k=1)) — silent corruption."
        );
    }

    /// Phase-3 regression test: a multi-step training-like loop
    /// (forward → backward → Adam → forward) that locks in the
    /// dW_Q layout alignment AND the cross-step consistency with
    /// the forward-layout convention.
    ///
    /// Two invariants checked:
    ///   1. **Cross-direction layout alignment** — GPU's `dW_Q` in
    ///      `d_params` matches a CPU reference computed in the
    ///      post-fix [K, M] layout (`ofi @ d_proj^T`) within TF32
    ///      tolerance. Pre-fix the GPU writes `d_proj @ ofi^T`
    ///      (col-major [M, K]) and the CPU reference (post-fix
    ///      [K, M]) disagrees element-wise at every flat index
    ///      where the two layouts diverge — which is ~all of them
    ///      when M ≠ K (TLOB: M=16, K=32). The parity test catches
    ///      the same thing at the single-pass level; this test
    ///      proves the alignment holds across an Adam step too
    ///      (no in-place permutation by the Adam kernel — there
    ///      shouldn't be one, but locked in for future refactor).
    ///
    ///   2. **Multi-step training stability** — after Adam updates
    ///      `params` in-place, the second forward Q output equals
    ///      the analytical Q computed from `params_after` under
    ///      the [K, M] layout, AND differs measurably from the
    ///      first forward Q (proving Adam actually moved weights).
    ///      This is the "training-like loop" assertion that
    ///      permanently locks in: future refactors cannot break
    ///      the forward/backward layout agreement without this
    ///      test failing.
    #[test]
    #[ignore = "requires GPU"]
    fn tlob_dw_layout_alignment_regression_full_chain() {
        let ctx = match CudaContext::new(0) {
            Ok(c) => c,
            Err(e) => {
                eprintln!("Skipping tlob_dw_layout_alignment_regression_full_chain (no CUDA device): {e}");
                return;
            }
        };
        let stream = ctx.default_stream();
        let shared = Arc::new(
            PerStreamCublasHandles::new(&stream)
                .expect("PerStreamCublasHandles::new"),
        );

        let batch_size = 32_usize;
        let mut tlob = GpuTlob::new(Arc::clone(&shared), batch_size).expect("GpuTlob::new");

        // ── Inject non-zero W_O so the gradient chain doesn't collapse ──
        // GpuTlob::new initialises W_Q/K/V via Xavier but leaves W_O at
        // zero (intentionally — production trains W_O from scratch).
        // For the regression test we need d_attn_out = W_O @ d_output to
        // be nonzero, so seed W_O with the same Xavier-like spread used
        // for the QKV block. Read-modify-write the params buffer.
        let mut rng_state: u64 = 0xfeedfacecafef00d;
        let mut prng = || -> f32 {
            rng_state = rng_state
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            ((rng_state >> 33) as f32 / (1u64 << 31) as f32) - 0.5
        };
        {
            let mut host_params = tlob.dump_params();
            let xavier_std_o = (2.0_f32 / (TLOB_OUT as f32 + TLOB_OUT as f32)).sqrt();
            for p in host_params[W_O_OFF..W_O_OFF + TLOB_OUT * TLOB_OUT].iter_mut() {
                *p = prng() * 3.46 * xavier_std_o;
            }
            let staging = unsafe { super::super::mapped_pinned::MappedF32Buffer::new(host_params.len()) }
                .expect("params staging");
            staging.write_from_slice(&host_params);
            let nbytes = host_params.len() * std::mem::size_of::<f32>();
            unsafe {
                let (dst_ptr, _g) = tlob.params.device_ptr_mut(&stream);
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(),
                )
                .expect("W_O seed DtoD");
            }
            stream.synchronize().expect("W_O seed sync");
        }
        let mut host_states = vec![0.0_f32; batch_size * STATE_DIM_PADDED];
        for b in 0..batch_size {
            for k in 0..TLOB_IN {
                host_states[b * STATE_DIM_PADDED + OFI_START + k] = prng();
            }
        }
        let mut states_buf = stream
            .alloc_zeros::<f32>(batch_size * STATE_DIM_PADDED)
            .expect("states_buf alloc");
        {
            let staging = unsafe { super::super::mapped_pinned::MappedF32Buffer::new(host_states.len()) }
                .expect("states staging");
            staging.write_from_slice(&host_states);
            let nbytes = host_states.len() * std::mem::size_of::<f32>();
            unsafe {
                let (dst_ptr, _g) = states_buf.device_ptr_mut(&stream);
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(),
                )
                .expect("states_buf DtoD");
            }
            stream.synchronize().expect("states_buf sync");
        }

        // ── Snapshot W_Q before any pass (for the second-forward
        //    analytical check below). ──
        let params_before = tlob.dump_params();
        let w_q_before = params_before[W_Q_OFF..W_Q_OFF + TLOB_OUT * TLOB_IN].to_vec();

        // Build ofi col-major [TLOB_IN, B] (same construction as parity test).
        let mut ofi_cm = vec![0.0_f32; TLOB_IN * batch_size];
        for b in 0..batch_size {
            for k in 0..TLOB_IN {
                ofi_cm[b * TLOB_IN + k] = host_states[b * STATE_DIM_PADDED + OFI_START + k];
            }
        }

        // ── Forward (loads `proj_qkv_buf` and `attn_out_buf` for backward). ──
        tlob.forward(&mut states_buf, batch_size).expect("forward");
        stream.synchronize().expect("sync forward");

        // Compute analytical Q-before from the snapshotted W_Q under [K, M]
        // layout: cpu_sgemm transA=T, lda=TLOB_IN=K. Same convention as
        // the parity test.
        let mut q_before_analytical = vec![0.0_f32; TLOB_OUT * batch_size];
        cpu_sgemm(
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            TLOB_OUT, batch_size, TLOB_IN,
            1.0,
            &w_q_before, TLOB_IN,
            &ofi_cm, TLOB_IN,
            0.0,
            &mut q_before_analytical, TLOB_OUT,
        );
        let q_before_gpu = tlob.dump_proj_q();
        // Sanity: pre-Adam Q matches the [K, M] interpretation of W_Q.
        // (Already checked by `tlob_sgemm_parity_with_cpu_reference` —
        // included here to lock the invariant for *this* test in case
        // someone refactors the parity test in isolation.)
        assert_close(&q_before_gpu, &q_before_analytical, TOL_GEMM,
            "regression: forward-1 Q matches [K, M] analytical");

        // ── Build random d_concat and run backward. ──
        let concat_dim = TLOB_OUT;
        let tlob_concat_off = 0_usize;
        let mut host_d_concat = vec![0.0_f32; batch_size * concat_dim];
        for v in host_d_concat.iter_mut() {
            *v = prng() * 0.1;
        }
        let mut d_concat = stream
            .alloc_zeros::<f32>(batch_size * concat_dim)
            .expect("d_concat alloc");
        {
            let staging = unsafe { super::super::mapped_pinned::MappedF32Buffer::new(host_d_concat.len()) }
                .expect("d_concat staging");
            staging.write_from_slice(&host_d_concat);
            let nbytes = host_d_concat.len() * std::mem::size_of::<f32>();
            unsafe {
                let (dst_ptr, _g) = d_concat.device_ptr_mut(&stream);
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(),
                )
                .expect("d_concat DtoD");
            }
            stream.synchronize().expect("d_concat sync");
        }
        tlob.backward(&d_concat, batch_size, concat_dim, tlob_concat_off)
            .expect("backward");
        stream.synchronize().expect("sync bwd");

        // ── Invariant 1 (cross-direction layout alignment) ────────
        // GPU dW_Q must match a CPU reference computed in the
        // post-fix [K, M] layout (`ofi @ d_proj^T`, M_sgemm=K,
        // N_sgemm=M, ldc=K) within TF32 tolerance. Pre-fix the
        // GPU's [M, K] layout disagrees with this reference at
        // every flat index where the two layouts diverge.
        //
        // We reconstruct `d_proj_q` from the SDP-backward kernel's
        // output (read back via dump_d_proj_qkv), then compute the
        // CPU dW_Q reference in the [K, M] convention.
        let d_proj_qkv = tlob.dump_d_proj_qkv();
        let chunk = TLOB_OUT * batch_size;
        let d_proj_q = &d_proj_qkv[..chunk];
        let mut ref_dw_q_km = vec![0.0_f32; TLOB_IN * TLOB_OUT];
        cpu_sgemm(
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            TLOB_IN, TLOB_OUT, batch_size,
            1.0,
            &ofi_cm, TLOB_IN,
            d_proj_q, TLOB_OUT,
            0.0,
            &mut ref_dw_q_km, TLOB_IN,
        );
        let d_params_post_bwd = tlob.dump_d_params();
        let gpu_dw_q = &d_params_post_bwd[W_Q_OFF..W_Q_OFF + TLOB_OUT * TLOB_IN];
        assert_close(gpu_dw_q, &ref_dw_q_km, TOL_GEMM,
            "regression: GPU dW_Q matches [K, M] CPU reference (post-fix layout)");

        // ── Adam step: lr=1e-3, no clipping, no weight-decay so the
        //    only motion in `params` is the per-element Adam update. ──
        let lr = 0.001_f32;
        let max_grad_norm = 1e9_f32;
        let weight_clamp = 0.0_f32;
        let weight_decay = 0.0_f32;
        // SP5 Layer B: adam_step signature took on Adam β1/β2/ε params from ISV.
        // Standard Adam defaults — exact β/ε don't affect this test's W_Q/dW_Q
        // layout assertion; only exercising that Adam runs at all.
        let attn_beta1 = 0.9_f32;
        let attn_beta2 = 0.999_f32;
        let attn_epsilon = 1e-8_f32;
        tlob.adam_step(lr, max_grad_norm, weight_clamp, weight_decay,
                       attn_beta1, attn_beta2, attn_epsilon)
            .expect("adam");
        stream.synchronize().expect("sync adam");

        let params_after = tlob.dump_params();
        let w_q_after = params_after[W_Q_OFF..W_Q_OFF + TLOB_OUT * TLOB_IN].to_vec();

        // ── Invariant 2 (multi-step training stability) ──────────
        // Run forward-2 and verify Q matches the [K, M]-analytical
        // interpretation of the post-Adam `w_q_after`. This locks
        // in that the forward layout convention is preserved across
        // an Adam step (Adam should be element-wise, no permutation;
        // the test would fail if some future refactor accidentally
        // re-laid-out `params` between Adam and the next forward).
        tlob.forward(&mut states_buf, batch_size).expect("forward 2");
        stream.synchronize().expect("sync forward 2");
        let q_after_gpu = tlob.dump_proj_q();
        let mut q_after_analytical = vec![0.0_f32; TLOB_OUT * batch_size];
        cpu_sgemm(
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            TLOB_OUT, batch_size, TLOB_IN,
            1.0,
            &w_q_after, TLOB_IN,
            &ofi_cm, TLOB_IN,
            0.0,
            &mut q_after_analytical, TLOB_OUT,
        );
        assert_close(&q_after_gpu, &q_after_analytical, TOL_GEMM,
            "regression: forward-2 Q matches [K, M] analytical of post-Adam W_Q");

        // ── Diff lock: forward-2 ≠ forward-1 (Adam actually moved
        //    the weights). ──
        let mut max_q_diff = 0.0_f32;
        for i in 0..q_before_gpu.len() {
            let d = (q_after_gpu[i] - q_before_gpu[i]).abs();
            if d > max_q_diff { max_q_diff = d; }
        }
        // Sanity dump on failure: print |dW_Q| stats so the operator can
        // see whether the gradient was too small to produce a Q diff.
        let dwq_max = gpu_dw_q.iter().map(|x| x.abs()).fold(0.0_f32, f32::max);
        let dwq_mean = gpu_dw_q.iter().map(|x| x.abs()).sum::<f32>() / gpu_dw_q.len() as f32;
        let mut max_param_diff_dbg = 0.0_f32;
        for i in 0..(TLOB_OUT * TLOB_IN) {
            let d = (w_q_after[i] - w_q_before[i]).abs();
            if d > max_param_diff_dbg { max_param_diff_dbg = d; }
        }
        assert!(
            max_q_diff > 1e-5,
            "regression: forward-2 Q should differ from forward-1 (Adam updated \
             W_Q); max |Δ Q| = {max_q_diff:.3e}, dW_Q max={dwq_max:.3e}, \
             dW_Q mean={dwq_mean:.3e}, max |Δ W_Q| = {max_param_diff_dbg:.3e}"
        );

        // ── Diff lock: w_q_after ≠ w_q_before (Adam wrote to params). ──
        let mut max_param_diff = 0.0_f32;
        let mut updated_count = 0usize;
        for i in 0..(TLOB_OUT * TLOB_IN) {
            let d = (w_q_after[i] - w_q_before[i]).abs();
            if d > max_param_diff { max_param_diff = d; }
            if d > 1e-6 { updated_count += 1; }
        }
        assert!(
            max_param_diff > 1e-6,
            "regression: post-Adam W_Q should differ from pre-Adam W_Q; \
             max |Δ W_Q| = {max_param_diff:.3e}"
        );

        eprintln!(
            "[tlob_dw_layout_alignment_regression_full_chain] M={}, K={}, batch={batch_size}\n  \
             GPU dW_Q matches [K, M] CPU reference (Invariant 1: cross-direction alignment)\n  \
             post-Adam W_Q updated at {updated_count}/{} slots; max |Δ W_Q| = {max_param_diff:.3e}\n  \
             max |Δ Q| forward-2 vs forward-1: {max_q_diff:.3e}\n  \
             forward-2 Q matches [K, M] analytical interpretation of post-Adam W_Q (Invariant 2: multi-step stability)",
            TLOB_OUT, TLOB_IN, TLOB_OUT * TLOB_IN,
        );
    }
}
