gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(tlob): align backward dW_Q/K/V layout with forward W_Q/K/V (Fix 20)

Browse Source 
Forward W_Q SGEMM stored col-major [K, M] (lda=K) while backward dW_Q
wrote col-major [M, K] (ldc=M). When M ≠ K (TLOB: M=16, K=32), Adam's
element-wise update applied gradients computed at position (m, k) to
weights stored at position (k, m) — silent learning corruption at
every flat index ≠ 0 (511 of 512 W_Q slots updated using wrong-position
gradients, matched in W_K/V; W_O is square so unaffected).

Standardised backward dW_Q/K/V SGEMM output to col-major [K, M] (ldc=K)
matching the forward layout (Strategy A from the audit brainstorm — the
forward layout is the definitive weight storage; Adam's flat layout
follows forward's allocation). The fix flips the cuBLAS strided-batched
operands: backward now computes `dW^T = ofi @ d_proj^T` instead of
`dW = d_proj @ ofi^T`. Same gradient values, just re-laid-out so flat
indexing matches `params`. 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).
The QKV-fusion `cublasSgemmStridedBatched(batch=3)` semantics are
preserved: ofi is the new shared operand (strideA=0), d_proj is the
per-batch operand (strideB=M·B), strideC=M·K=512 unchanged.

Phase-1 reproduction (`tlob_dw_layout_alignment_repro`,
#[ignore = "requires GPU"]) ran the broken and fixed cuBLAS dispatches
side-by-side on identical sentinel inputs (`d_proj[m=0,b=0]=1`,
`ofi[k=1,b=0]=1`, all else 0); broken `[M, K]` placed the `1.0`
gradient at flat 16, fixed `[K, M]` placed it at flat 1 — O(1)
cross-layout delta exactly matching the audit prediction. Pre-fix Adam
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]` — the silent
corruption.

Phase-3 regression (`tlob_dw_layout_alignment_regression_full_chain`,
#[ignore = "requires GPU"]) exercises the full forward → backward →
Adam → forward chain with random Xavier-init weights (W_O seeded to
break the production-zero-init that would collapse the gradient chain
to all-zero in a synthetic test). Asserts (1) GPU dW_Q matches a CPU
reference computed in the post-fix [K, M] layout within TF32 tolerance,
and (2) the second forward Q matches the analytical [K, M]
interpretation of the post-Adam W_Q — locks in cross-step layout
agreement and would fail if any future refactor accidentally
re-permutes `params` between Adam and the next forward.

Existing inline `tlob_sgemm_parity_with_cpu_reference` still passes
(its CPU dW_Q/K/V reference was updated in lockstep to the [K, M]
layout per `feedback_no_partial_refactor`; pre-fix the GPU produced
[M, K] and the new CPU reference would diverge element-wise — a clean
no-skip parity check that locks the layout convention end-to-end).
`tlob_qkv_fusion_equivalence` unchanged (the fix only touches the
backward call, forward QKV fusion is bit-identical pre/post).

Local verification (RTX 3050 Ti, batch=256 for fusion test):
  tlob_dw_layout_alignment_repro:                      PASS
  tlob_dw_layout_alignment_regression_full_chain:      PASS
  tlob_qkv_fusion_equivalence:    PASS (3.79× speedup retained)
  tlob_sgemm_parity_with_cpu_reference:                PASS

Fix 20 in docs/dqn-gpu-hot-path-audit.md updated FIXED with verdict
+ strategy + test list. Forward SGEMM call site got an inline comment
block documenting the [K, M] convention and pointing at the
`tlob_dw_layout_alignment_*` regression coverage.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-03 12:15:34 +02:00
parent dad76e1c11
commit e09757419c
2 changed files with 632 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

629
crates/ml/src/cuda_pipeline/gpu_tlob.rs
View File

@@ -448,6 +448,21 @@ impl GpuTlob {
// 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
@@ -658,31 +673,57 @@ impl GpuTlob {
}
// ── Step 4: Fused dW_Q/K/V backward ───────────────────────────
// Math (per-projection): dW_X = d_proj_X @ ofi^T (col-major
// [TLOB_OUT, TLOB_IN]):
// transA=N (d_proj_X 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.
// Layout convention (Fix 20, 2026-05-03):
//
// Fused dispatch: `cublasSgemmStridedBatched(batch=3)` reads the
// three contiguous d_proj_X chunks (`strideA = M·B`), the same
// ofi for all three (`strideB = 0`), and writes into the
// contiguous `[dW_Q | dW_K | dW_V]` slot of `d_params`
// (`strideC = M·K`). One launch in place of three.
// 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;
let stride_a_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
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_qkv_buf.raw_ptr(), m_qkv, stride_a_qkv,
ofi_ptr, n_qkv, 0_i64,
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,
@@ -1011,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<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();
@@ -1380,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];
@@ -1392,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,
);
}
@@ -1691,4 +1750,534 @@ mod tests {
(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;
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::<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,
);
}
}

23
docs/dqn-gpu-hot-path-audit.md
View File

@@ -270,3 +270,26 @@ Per `feedback_no_partial_refactor`: kernel sigs preserved (offset device pointer
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.