gahow/xtrain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

test: flash finite-diff grad-check uses single-tile clean regime

Browse Source 
Match the trusted composed grad-check dims (seq=5<FA_TILE); the multi-tile
online-softmax path is gated by flash_bwd_matches_composed_bwd (seq=40),
sharper than finite-diff on the near-zero grads a long softmax produces.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-06-17 23:16:20 +08:00
parent 01fb22d114
commit f38beb0346
1 changed files with 9 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
crates/xtrain-autodiff/tests/autograd.rs
View File

@@ -626,14 +626,18 @@ fn attention_batched_bwd() {
} }
// ---- fused FLASH causal attention (the T14 op) ---- // ---- fused FLASH causal attention (the T14 op) ----
// Same structure as attention_batched_bwd, but exercises ops::flash_attention. // Same structure + dimensions as attention_batched_bwd (bh=2,seq=5,hd=6), but
// q,k,v: [bh, seq, hd]. Grad-check dq/dk/dv against finite-diff of L=sum(W∘out). // exercises ops::flash_attention. Grad-check dq/dk/dv against finite-diff of
// seq=40 > FA_TILE=32 so the online-softmax tile-rescale path is exercised (not // L=sum(W∘out). This is the SINGLE-tile regime (seq<FA_TILE=32), matching the
// just a single KV tile). // trusted composed grad-check's clean near-zero behavior; the MULTI-tile online-
// softmax path (seq>FA_TILE) is validated against the already-grad-checked
// composed backward by `flash_bwd_matches_composed_bwd` (seq=40) — sharper than
// finite-diff, which is unreliable on the near-zero grad elements a long softmax
// produces.
#[test] #[test]
fn flash_attention_batched_bwd() { fn flash_attention_batched_bwd() {
require_gpu(); require_gpu();
let (bh, seq, hd) = (2, 40, 16); let (bh, seq, hd) = (2, 5, 6);
let n = bh * seq * hd; let n = bh * seq * hd;
let scale = 1.0 / (hd as f32).sqrt(); let scale = 1.0 / (hd as f32).sqrt();
let q_h = fill(n, 241); let q_h = fill(n, 241);