Merge: fix(tlob) align dW layout + fuse Q/K/V SGEMMs

# Conflicts:
#	crates/ml/src/cuda_pipeline/gpu_tlob.rs
This commit is contained in:
jgrusewski
2026-05-03 15:59:57 +02:00
2 changed files with 1135 additions and 132 deletions

File diff suppressed because it is too large Load Diff

View File

@@ -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.