diff --git a/crates/ml/src/cuda_pipeline/gpu_tlob.rs b/crates/ml/src/cuda_pipeline/gpu_tlob.rs index 73a4e639d..75fc3f7bd 100644 --- a/crates/ml/src/cuda_pipeline/gpu_tlob.rs +++ b/crates/ml/src/cuda_pipeline/gpu_tlob.rs @@ -6,9 +6,10 @@ //! scaled-dot-product attention over the feature dimension. //! //! Architecture: -//! Q = W_Q^T @ ofi_input[32,B] → [16,B] (classic cublasSgemm_v2) -//! K = W_K^T @ ofi_input[32,B] → [16,B] (classic cublasSgemm_v2) -//! V = W_V^T @ ofi_input[32,B] → [16,B] (classic cublasSgemm_v2) +//! [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) @@ -18,9 +19,8 @@ //! 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[16,32] += ofi^T @ d_q / B (classic cublasSgemm_v2) -//! dW_K[16,32] += ofi^T @ d_k / B (classic cublasSgemm_v2) -//! dW_V[16,32] += ofi^T @ d_v / B (classic cublasSgemm_v2) +//! 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. @@ -30,17 +30,42 @@ //! 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, not cuBLAS-Lt +//! ## 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` 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. +//! `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; @@ -79,6 +104,16 @@ 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]. @@ -112,12 +147,13 @@ pub(crate) struct GpuTlob { // ── Intermediate buffers (forward, saved for backward) ──────────── /// OFI input in col-major [TLOB_IN=32, B]: extracted from states_buf. ofi_col_buf: CudaSlice, - /// Q projection [TLOB_OUT=16, B]. - proj_q_buf: CudaSlice, - /// K projection [TLOB_OUT=16, B]. - proj_k_buf: CudaSlice, - /// V projection [TLOB_OUT=16, B]. - proj_v_buf: CudaSlice, + /// 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, /// SDP output [TLOB_OUT=16, B] (pre-output-projection), saved for dW_O. attn_out_buf: CudaSlice, /// SDP softmax weights [TLOB_OUT=16, B], saved for backward. @@ -130,12 +166,12 @@ pub(crate) struct GpuTlob { d_output_buf: CudaSlice, /// Gradient w.r.t. attn_out [TLOB_OUT=16, B]. d_attn_out_buf: CudaSlice, - /// Gradient w.r.t. Q [TLOB_OUT=16, B]. - d_proj_q_buf: CudaSlice, - /// Gradient w.r.t. K [TLOB_OUT=16, B]. - d_proj_k_buf: CudaSlice, - /// Gradient w.r.t. V [TLOB_OUT=16, B]. - d_proj_v_buf: CudaSlice, + /// 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, /// Parameter gradients [TLOB_TOTAL_PARAMS]. d_params: CudaSlice, @@ -232,6 +268,11 @@ impl GpuTlob { let mut params = stream.alloc_zeros::(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 { @@ -258,19 +299,18 @@ impl GpuTlob { .map_err(|e| MLError::ModelError(format!("tlob {label} alloc: {e}"))) }; - let ofi_col_buf = alloc(db_in, "ofi_col")?; - let proj_q_buf = alloc(db_out, "proj_q")?; - let proj_k_buf = alloc(db_out, "proj_k")?; - let proj_v_buf = alloc(db_out, "proj_v")?; - 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")?; - let d_proj_q_buf = alloc(db_out, "d_proj_q")?; - let d_proj_k_buf = alloc(db_out, "d_proj_k")?; - let d_proj_v_buf = alloc(db_out, "d_proj_v")?; - let d_params = alloc(TLOB_TOTAL_PARAMS, "d_params")?; + 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")?; @@ -328,17 +368,13 @@ impl GpuTlob { read_grad_kernel, params, ofi_col_buf, - proj_q_buf, - proj_k_buf, - proj_v_buf, + proj_qkv_buf, attn_out_buf, attn_scores_buf, output_buf, d_output_buf, d_attn_out_buf, - d_proj_q_buf, - d_proj_k_buf, - d_proj_v_buf, + d_proj_qkv_buf, d_params, adam_m, adam_v, @@ -361,8 +397,10 @@ impl GpuTlob { /// TLOB forward pass. /// /// 1. Extract OFI slice from states_buf → ofi_col_buf [TLOB_IN, B] col-major. - /// 2. Q/K/V projections via cuBLAS. - /// 3. SDP (tlob_sdp_forward kernel). + /// 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. /// @@ -403,59 +441,70 @@ impl GpuTlob { let params_base = self.params.raw_ptr(); let ofi_ptr = self.ofi_col_buf.raw_ptr(); - // ── Step 2: Q/K/V projections ───────────────────────────────── - // Q = W_Q^T @ ofi: 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 we set transA = T to read it as [TLOB_OUT, TLOB_IN] col-major. - // ofi is already col-major [TLOB_IN, B], transB = N. - // M = TLOB_OUT, N = B, K = TLOB_IN, lda = ldb = TLOB_IN, ldc = TLOB_OUT. + // ── 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; - sgemm_v2( + 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, - ofi_ptr, k_qkv, + 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_q_buf.raw_ptr(), m_qkv, - "fwd_q", - )?; - 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_f32, - params_base + (W_K_OFF * f32_sz) as u64, k_qkv, - ofi_ptr, k_qkv, - 0.0_f32, - self.proj_k_buf.raw_ptr(), m_qkv, - "fwd_k", - )?; - 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_f32, - params_base + (W_V_OFF * f32_sz) as u64, k_qkv, - ofi_ptr, k_qkv, - 0.0_f32, - self.proj_v_buf.raw_ptr(), m_qkv, - "fwd_v", + 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(&self.proj_q_buf) - .arg(&self.proj_k_buf) - .arg(&self.proj_v_buf) + .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) @@ -588,17 +637,31 @@ impl GpuTlob { )?; // ── 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(&self.proj_v_buf) - .arg(&self.proj_q_buf) - .arg(&self.proj_k_buf) - .arg(&mut self.d_proj_q_buf) - .arg(&mut self.d_proj_k_buf) - .arg(&mut self.d_proj_v_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 { @@ -609,51 +672,62 @@ impl GpuTlob { .map_err(|e| MLError::ModelError(format!("tlob sdp_bwd: {e}")))?; } - // ── Step 4: dW_Q/K/V backward ──────────────────────────────── - // dW_Q = d_proj_q @ ofi^T col-major: - // [TLOB_OUT, TLOB_IN] = [TLOB_OUT, B] @ [B, TLOB_IN] - // transA=N (d_proj_q as-stored), transB=T (ofi: col [TLOB_IN, B] -> [B, TLOB_IN]) - // M=TLOB_OUT, N=TLOB_IN, K=B - // alpha=1.0 (sdp_bwd already divides by B) + // ── 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_OUT as i32; - let n_qkv = TLOB_IN as i32; + 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; - sgemm_v2( + 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, - self.d_proj_q_buf.raw_ptr(), m_qkv, - ofi_ptr, n_qkv, + 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, - "bwd_dw_q", - )?; - sgemm_v2( - classic, - cublas_sys::cublasOperation_t::CUBLAS_OP_N, - cublas_sys::cublasOperation_t::CUBLAS_OP_T, - m_qkv, n_qkv, k_qkv, - 1.0_f32, - self.d_proj_k_buf.raw_ptr(), m_qkv, - ofi_ptr, n_qkv, - 0.0_f32, - grad_base + (W_K_OFF * f32_sz) as u64, m_qkv, - "bwd_dw_k", - )?; - sgemm_v2( - classic, - cublas_sys::cublasOperation_t::CUBLAS_OP_N, - cublas_sys::cublasOperation_t::CUBLAS_OP_T, - m_qkv, n_qkv, k_qkv, - 1.0_f32, - self.d_proj_v_buf.raw_ptr(), m_qkv, - ofi_ptr, n_qkv, - 0.0_f32, - grad_base + (W_V_OFF * f32_sz) as u64, m_qkv, - "bwd_dw_v", + grad_base + (W_Q_OFF * f32_sz) as u64, m_qkv, stride_c_qkv, + QKV_BATCH as i32, + "bwd_dw_qkv", )?; Ok(()) @@ -888,6 +962,48 @@ fn sgemm_v2( 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 @@ -902,10 +1018,23 @@ impl GpuTlob { 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 { - let mut h = vec![0.0_f32; TLOB_OUT * self.batch_size]; + 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_q_buf, &mut h); + 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 { + 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 } @@ -923,6 +1052,18 @@ impl GpuTlob { 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 { + 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 { let mut h = vec![0.0_f32; TLOB_TOTAL_PARAMS]; let _ = self.stream.synchronize(); @@ -1060,7 +1201,8 @@ mod tests { } } - // Upload states_buf. + // Upload states_buf via mapped-pinned staging (canonical pattern; + // matches the production param-upload path above). let mut states_buf = stream .alloc_zeros::(batch_size * STATE_DIM_PADDED) .expect("states_buf alloc"); @@ -1291,7 +1433,13 @@ mod tests { } } - // dW_Q = d_proj_q @ ofi^T (alpha=1, sdp_bwd already applied 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]; @@ -1303,12 +1451,12 @@ mod tests { cpu_sgemm( cublas_sys::cublasOperation_t::CUBLAS_OP_N, cublas_sys::cublasOperation_t::CUBLAS_OP_T, - TLOB_OUT, TLOB_IN, batch_size, + TLOB_IN, TLOB_OUT, batch_size, 1.0, - d_proj_, TLOB_OUT, &ofi_cm, TLOB_IN, + d_proj_, TLOB_OUT, 0.0, - dw_, TLOB_OUT, + dw_, TLOB_IN, ); } @@ -1323,4 +1471,813 @@ mod tests { 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, + shared: &Arc, + params_ptr: u64, + ofi_ptr: u64, + batch_size: usize, + n_iters: usize, + ) -> (Vec, std::time::Duration) { + let f32_sz = std::mem::size_of::(); + 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::(chunk).expect("proj_q alloc"); + let proj_k = stream.alloc_zeros::(chunk).expect("proj_k alloc"); + let proj_v = stream.alloc_zeros::(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, + shared: &Arc, + params_ptr: u64, + ofi_ptr: u64, + batch_size: usize, + n_iters: usize, + ) -> std::time::Duration { + let f32_sz = std::mem::size_of::(); + 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::(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::(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::(); + 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, + stream: &Arc, + shared: &Arc, + batch_size: usize, + ) -> (Vec, Vec) { + // 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, + host: &[f32]| + -> CudaSlice { + let mut dev = stream.alloc_zeros::(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::(); + 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::(m * k).expect("dw_broken alloc"); + let dw_fixed = stream.alloc_zeros::(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::(); + 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::(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::(); + 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::(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::(); + 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; + tlob.adam_step(lr, max_grad_norm, weight_clamp, weight_decay) + .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::() / 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, + ); + } } diff --git a/docs/dqn-gpu-hot-path-audit.md b/docs/dqn-gpu-hot-path-audit.md index 558855b2b..aa4b7e1d0 100644 --- a/docs/dqn-gpu-hot-path-audit.md +++ b/docs/dqn-gpu-hot-path-audit.md @@ -247,3 +247,49 @@ Atom-stat block-sum reduction unchanged (warp shfl + shared-mem bank, determinis Expected wall-clock saving on this kernel: ~33% of 207ms = ~70ms across the smoke run; on production batch=8192 × 1382 calls per fold, roughly proportional. Lower-bound estimate; actual saving depends on memory bandwidth vs compute mix — at low N (smoke) the kernel may be launch-bound rather than load-bound. nsys re-profile after deployment will quantify. ABI unchanged; no Rust caller changes required. + +### Fix 20 (perf) — Fuse TLOB Q/K/V SGEMMs into `cublasSgemmStridedBatched` (2026-05-03) + +`crates/ml/src/cuda_pipeline/gpu_tlob.rs::forward()` previously dispatched three back-to-back `cublasSgemm_v2` calls (M=16, K=32, N=B at TF32) for the Q, K, V projections, plus three more in `backward()` for the dW_Q/K/V parameter gradients. Per task #218 nsys close-out, these dominated one of the kernel-2 hotspots: `cublasLtMatmul` heuristic selection was failing for M=16 K=32 (per-head dimensions) and the classic-sgemm fallback was paying launch + heuristic-lookup overhead 3× per forward and 3× per backward. + +Replaced with a single `cublasSgemmStridedBatched(batch=3)` per direction. Strategy 2 from the brainstorm (strided batched, NOT concatenated weights): the existing flat param block `[W_Q | W_K | W_V]` already gives strideA=M·K=512 floats per weight matrix, and a fused `proj_qkv_buf` of length 3·M·B holds three contiguous `[M, B]` col-major chunks at offsets 0, M·B, 2·M·B (strideC=M·B). The SDP forward + backward kernels are unchanged: each still receives three `const float*` Q/K/V pointers — now they're computed as offsets into the single fused buffer rather than three separate allocations. The concat-W path (Strategy 1, single SGEMM with M=3·M=48, ldc=3M) was rejected because it would have forced an SDP kernel signature change (stride-3M reads), breaking bit-equivalence with the prior 3-SGEMM path. + +Param flat layout `[W_Q | W_K | W_V | W_O]` is unchanged; no checkpoint format change (TLOB has no on-disk checkpoint anyway — Xavier random init). The strided-batched stride constants (`QKV_BATCH=3`, `W_QKV_STRIDE_FLOATS=M·K`) are derived from the existing per-projection shapes, no new tunable hyperparameters introduced. + +Numerical equivalence (RTX 3050 Ti, batch=256, TF32 mode): +- max abs Q/K/V diff vs reference 3-SGEMM path: 4.29e-4 (well under the 2e-3 TF32-aware tolerance also used by the inline `tlob_sgemm_parity_with_cpu_reference` test). +- Why not 1e-5: classic cuBLAS handle is bound to `CUBLAS_TF32_TENSOR_OP_MATH` (`shared_cublas_handle::create_handles_and_workspace`); strided-batched dispatch picks a different internal algo than back-to-back single calls and the K=32 reduction amplifies TF32 rounding to a few × 1e-4. Both paths are mathematically equivalent within TF32 precision; layout/stride/offset bugs would show up as O(1) deltas, not O(1e-3). + +Microbenchmark (RTX 3050 Ti, batch=256, 200-iter loop, QKV-SGEMM portion only — apples-to-apples): +- fused 1× `cublasSgemmStridedBatched(batch=3)`: 5.18 µs/call +- ref 3× `cublasSgemm_v2` back-to-back: 19.56 µs/call +- → **3.78× speedup** on the QKV-projection portion alone (forward; backward dW_Q/K/V fusion has the same shape, same expected gain). + +Per `feedback_no_partial_refactor`: kernel sigs preserved (offset device pointers); `proj_q_buf` / `proj_k_buf` / `proj_v_buf` (and their d_proj counterparts) replaced with `proj_qkv_buf` / `d_proj_qkv_buf` everywhere — no stale call sites. + +Verification: `tlob_sgemm_parity_with_cpu_reference` (existing inline parity test against an f32 CPU reference) still passes; `tlob_qkv_fusion_equivalence` (new, `#[ignore = "requires GPU"]`) passes locally with the latency numbers above. Workspace `cargo check` is clean (15 pre-existing warnings, 0 new). + +Pre-existing observation surfaced during analysis (NOT acted on in this commit — flagged for a separate audit pass): the cuBLAS forward + backward dW SGEMM ldas/ldcs in `gpu_tlob.rs` deserve an end-to-end Adam-update review. Forward stores `W_Q` col-major `[K, M]` (lda=K) while the backward dW_Q SGEMM emits col-major `[M, K]` (ldc=M); flat-index `i` of `params` and `d_params` therefore name (k=i%K, m=i/K) and (m=i%M, k=i/M) respectively, which agree at i=0 but diverge from i=1 when M≠K (here M=16, K=32). The inline parity test compares each direction against its own col-major reference and so does not exercise the Adam element-wise consumption that would conflate them. This is orthogonal to the QKV fusion (the fusion preserves the existing per-projection lda/ldc byte-for-byte) and likely also orthogonal to the kernel-internal `proj_X[d*B + b]` access pattern (kernel-internal, consistent across forward/backward kernels — a self-consistent transposed feature view). Worth either confirming the layouts are intentional (with a comment block in the file) or fixing in a dedicated ticket; out of scope for this perf change. + +### Fix 20-followup (correctness) — Align TLOB backward dW_Q/K/V layout with forward W_Q/K/V (2026-05-03) + +**Status**: FIXED. Verdict: **bug was real**, not a self-consistent transposed feature view. Phase-1 reproduction (`tlob_dw_layout_alignment_repro` in the inline `gpu_tlob.rs` tests module) ran the broken and fixed cuBLAS dispatches side-by-side on identical inputs (`d_proj[m=0,b=0]=1`, `ofi[k=1,b=0]=1`, all else 0); the broken `[M, K]` layout placed the `1.0` gradient at flat index 16, while the fixed `[K, M]` layout placed it at flat index 1 — an O(1) cross-layout delta exactly matching the audit's prediction. Pre-fix, Adam (which is element-wise: `params[tid] -= lr * d_params[tid] / ...` in `attention_backward_kernel.cu::attn_adam_kernel`) would have updated `W_Q[m=0, k=16]` (the forward layout's flat-16 slot) using the gradient computed for `W_Q[m=0, k=1]` — silent learning corruption, sustained for the full `dW_Q[m, k]` matrix at every i ≠ 0 when M ≠ K (TLOB: M=16, K=32; corrupted slots = 511 of 512). + +**Strategy**: A from the brainstorm — standardize on the forward `[K, M]` layout (the "definitive" weight storage; Adam's flat layout follows the forward allocation). Backward now computes `dW_Q^T = ofi @ d_proj^T` instead of `dW_Q = d_proj @ ofi^T`: same gradient values, just re-laid-out so the flat indexing matches `params`. Implementation is a single operand-swap in the `cublasSgemmStridedBatched` call: + +```diff +- A = d_proj (op=N, lda=M=16, strideA=M·B), B = ofi (op=T, ldb=N=32, strideB=0) +- M_sgemm = TLOB_OUT = 16, N_sgemm = TLOB_IN = 32, ldc = M_sgemm = 16 ++ A = ofi (op=N, lda=K=32, strideA=0), B = d_proj (op=T, ldb=M=16, strideB=M·B) ++ M_sgemm = TLOB_IN = 32, N_sgemm = TLOB_OUT = 16, ldc = M_sgemm = 32 +``` + +No new kernel; no kernel-internal layout change (the SDP forward/backward kernels still read `proj_qkv_buf` / `d_proj_qkv_buf` as `[M, B]` col-major — those buffers are untouched by the fix, only the dW destination layout in `d_params` changes). The QKV-fusion `cublasSgemmStridedBatched(batch=3)` semantics are preserved — strideA=0 (ofi is the shared operand now), strideB=M·B (d_proj is the per-batch operand), strideC=M·K=512 (unchanged: each `dW_X` slot is the same M·K floats wide regardless of which axis is the leading dim). + +**Tests**: +- `tlob_dw_layout_alignment_repro` (Phase 1, GPU-only) — proves the layout claim numerically with sentinel values; asserts O(1) cross-layout delta at flat 16 and bug-confirmation that broken layout would feed Adam a wrong-position gradient. +- `tlob_dw_layout_alignment_regression_full_chain` (Phase 3, GPU-only) — exercises the full forward → backward → Adam → params-inspection chain with a controlled `W_Q[m=0, k=1] = 0.5` injection, asserts the dominant `dW_Q` slot lands at flat 1 (post-fix [K, M]) and that Adam updates `params[1]` while leaving `params[16]` untouched. Pre-fix, this regression test fails because the dominant `dW_Q` would be at flat 16 and `params[16]` would receive the spurious update. +- Existing `tlob_sgemm_parity_with_cpu_reference` (no `#[ignore]`, runs on CPU-only nodes by skipping cleanly when no GPU) updated in lockstep: the CPU dW reference now computes `ofi @ d_proj^T` instead of `d_proj @ ofi^T` so the col-major `[K, M]` GPU output matches the col-major `[K, M]` CPU reference element-wise. The previous CPU reference was masking the bug by replicating the same wrong layout convention; updating both halves of the contract simultaneously per `feedback_no_partial_refactor`. +- `tlob_qkv_fusion_equivalence` unchanged — still asserts the post-fusion forward Q/K/V outputs match the 3-SGEMM reference forward path within TF32 tolerance (the fix only touches the backward call). + +The forward SGEMM call site in `gpu_tlob.rs::forward()` got an inline comment block documenting the [K, M] layout convention and pointing at the `tlob_dw_layout_alignment_*` tests as the canonical regression coverage for the cross-direction alignment that the per-direction parity test cannot catch.