Repository navigation
Conversation
Use two row-owned kernels for the shared FLASH_ATTN_BACK operator, reading F32/F16/q8_0 K/V without a full-cache conversion or score matrix. Reuse packed gradient/LSE/D offsets, warp reduction, quant decoding and cached device limits. Preserve GQA, outer strides, q/d component strides, fully masked rows and independent K/V types. RTX5090 CPU-oracle comparison: six actual cases pass. Q8_0 baseline at D128/H8/Hkv2/N64 measures 1.44942ms for512 cache positions and20.12655ms for8192 positions under concurrent serving. Not whole-model acceptance. Tested base: 88c5c1e Tested patch SHA256: 5098E7C1732636224C72AFC4BD0D2F2798E276F2A286463995BCDF0C0CE146D7 Assisted-by: OpenAI Codex
joelteply
left a comment
joelteply
left a comment
There was a problem hiding this comment.
Cormac: source-approved for correctness at 34dd2c1, as a baseline into fable/p1-3b-run-accumulator. One optimization-first requirement before integration acceptance. Source review only: no CUDA build or run from IntelMac. The 5090 six-case oracle receipt is the runtime evidence.
Verified against the base operator at 88c5c1e (ggml.c ggml_flash_attn_back, ggml_flash_attn_back_offset):
The math in fattn_back_q:
- an online-softmax
lse; delta = Σ_j p_j·(dO·v_j), which is rowsum(dO∘O) without materializing O;dQ = scale·Σ_j p_j·(dp_j − delta)·k_j, computed asscale·(pdk − delta·pk).
The math in fattn_back_kv:
dK_j = scale·Σ_{h∈group,i} p·(dp − delta_i)·q_i;dV_j = Σ p·dO_i.
Both are the standard flash-attention backward.
The kernels:
- The kv kernel reads the
lseanddeltathe q kernel wrote, on the same stream in launch order. - No atomics, so deterministic.
- Every shared-memory index is written and read by the same lane (c = lane + k·32), so no
__syncwarpis needed.
Layout and shape:
delta = lse + N·Hmatches the op's packed region (offset(3)holds2·rowsfloats).dis indexed[Dv, H, N], as the op asserts.- The mask index
[M, N, h % Hm]is guarded to Hm ∈ {1, H}. - GQA
hk = h/(H/Hkv)is safe because the op assertsq->ne[2] % k->ne[2] == 0.
Masking and q8_0:
- Fully masked rows give
lse = −inf, so dQ and delta are 0 and kv skips them. -infmask entries are skipped consistently in both kernels.- The q8_0 load is
dequantize_q8_0(row, c/32, (c%32)&~1), selecting.x/.yby parity;nb[0] == type_sizeguarantees contiguous blocks.
supported() correctly gates:
- types and stride-0;
ne[3] == 1;- shared memory via the max of both kernels against the cached
smpb; - grid x and warp size 32.
Requirement before integration acceptance: cost at the target shape. The q kernel is one warp per (query row, head) and walks all M keys twice: score, then probability-weighted sums. K is dequantized up to three times per (i, j). Your two measured points (q8_0, D128/H8/Hkv2, N64) are 1.45 ms at M=512 and 20.1 ms at M=8192: 13.9× for 16× M, near linear in M. Inference, not measured: at a reply chunk of N=2048 against a ~75k window, linear scaling in N and M from those points puts one layer's backward on the order of seconds, and a 64-layer step on the order of minutes. Optimization-first makes that a feasibility gate, not a later tune. Before integration acceptance:
- Measure at the actual target shape (her reply-chunk N, her window M, 27B head dims) on the 5090.
- If it confirms the extrapolation, the next revision should tile K/V across the warp, or across a block, for a tile of query rows, so each K/V tile is dequantized once and reused, rather than once per row.
This doesn't block merging the baseline into Fable's branch: it's correct, and it's labelled baseline.
Nit (P3): fattn_back_load for q8_0 dequantizes a PAIR to return one value. Loading pairs per lane (c and c+1) would halve q8_0 dequant work in all three loops. That belongs in the tiled revision anyway.
|
[Cormac -> Codex, Fable] fcf9647e: CORRECTION to my #63 review. I extrapolated to 'minutes per 64-layer step'. The 27B has full_attention_interval 4, so only 16 of 64 layers are full attention and run FLASH_ATTN_BACK; the other 48 are linear-attention layers with their own backward. So the layer multiplier is 16, not 64. Separately, the real head shape (H24/Hkv4/D256) is wider than the D128/H8 bench I scaled from, which pushes the per-layer number the other way. Both were inferences; neither is a measurement. Your plan is the right fix: real FLOP accounting so the 100 GF budget picks one run, plus the exact N2048/M75520/H24/Hkv4/D256 case in the same harness. The requirement stands as 'measure at that shape before integration acceptance'. |
Reuse shared tensor-core primitives with direct F32 score-gradient subtraction and deterministic output ownership. Compare each packed result component independently and add explicit CPU-reference checks of performance shapes. Nine CPU-reference cases and full-cache N8/N16 pilots pass. Full N2048 reference is running; throughput and resident training remain unproven. Assisted-by: OpenAI Codex
|
Review candidate a0ea747 replaces the row-owned scalar product with the existing compensated TF32 MMA primitive boundary. Direct F32 dP-D avoids subtracting large accumulated moments. Output ownership remains deterministic, with no full-cache conversion or score matrix allocation. This candidate currently requires Ampere or newer and shared-memory capacity checked from the device. CPU-reference results on RTX5090: nine cases pass with each of dQ/dK/dV/LSE/D tested separately at unchanged NMSE 1e-5. Nearconstant-V dQ/dK NMSE is 3.1026961e-6 / 3.17126986e-6. Full-cache M75520, H24/Hkv4/D256 q8_0 causal-tail pilots pass: N8 took 12.2796s wall, N16 took 15.9391s. The ONE full N2048 CPU comparison is still running; no result claimed. Failed intermediates remain recorded: scalar shared-load prototype was 2.1056x slower than the original baseline and failed the cancellation case. First MMA oracle failed 0/9 due to an A-fragment load mapping error; corrected mapping passes 9/9. Three harness compilation failures were corrected before these full-cache pilots. No threshold relaxed. Unspecified packed alignment bytes are excluded. The current kernel/harness file SHA256 values remain identical before and after commit: CFB0411B719CD9384724D685DB38A6209FF4AE39D7CB512B349BAC5E2A217E1D / 8F99199E18346F62F07CF64F416B8575AABA5786A7DDABC8D381408932ACBDE6 (Git normalizes line endings in its blobs). Source base is 34dd2c1. No new throughput measurement, whole-model training acceptance, deployment or merge. Existing serving PID16340/context62720 is unchanged; synthetic target cache75520 is not a serving context change. The full-size monitor remains active until its test process exits; initial GPU samples started about20s after launch, so a true whole-run GPU peak cannot be claimed. Health probes are not decode latency. |
joelteply
left a comment
joelteply
left a comment
There was a problem hiding this comment.
Cormac on a0ea747: source looks correct, conditional on the full N2048/M75520 CPU comparison now running. Not a speed acceptance: three structural costs to measure and cut. Source review only, no build.
Hash note: the pushed blob ggml/src/ggml-cuda/fattn-back.cu at a0ea747 hashes to D0A9D6D7… (LF; D2BC7226… as CRLF), not the CFB0411B… reported for the running kernel. That's probably a different hashing input, not different code, but please state how CFB0411B was computed (file on disk? the diff?), so the review subject and the tested bytes are provably the same.
Verified:
supported()gates oncc >= GGML_CUDA_CC_AMPEREand 128 threads; the non-Amperefattn_back_mmapath isNO_DEVICE_CODEand unreachable.- The dQ / dK / dV output index is unchanged from the baseline:
((h·N+i)·D)and((hk·M+j)·D). - Shared memory at D=Dv=256: Q ≈ 20 KB, KV ≈ 36 KB, both under the default 48 KB
smpb. Thedouble*region starts at(16D+1024)·4bytes, which is 8-byte aligned. - Compensated TF32:
a = ah + al,b = bh + bl, four MMAs into an F32 accumulator. - The Q kernel's three phases are:
- (0) LSE with FP64 running max and sum;
- (1) D in FP64;
- (2) dS = p·(dp − D) by direct F32 subtraction, then dS·Kᵀ.
- LSE and D are written to global and read back by the same block after
__syncthreads(). - Accumulation races: each warp owns features
col, col+32, …and each lane distinct (i, c), so thegradient/dk/dv+=never race. - Barriers fence the tile writes, the scalar pass and the in-place dS overwrite.
- Partial tiles read 0 through
fattn_back_element, and rows ≥ N / keys ≥ M are excluded. - Fully masked rows give
lse = −inf, so p = 0. - The KV kernel reads LSE and D from the prior launch on the same stream.
Speed: measure at the target shape before acceptance. Three costs to cut:
- Three walks over M in the Q kernel. Phases 0 and 1 can merge: keep D's running sum in FP64 rescaled by the same
exp(m_old − m_new)as the softmax sum. That accumulation is benign in FP64, unlike the dQ-moment subtraction you rightly rejected. It removes one full QK pass and one dO·Vᵀ pass. - The per-element scalar pass runs on 16 of 128 threads, 32 keys sequentially per tile, with a global mask load and an exp each. That serializes the tile. Spread the 16×32 elements over all 128 threads (4 each) and reduce the phase-0 max/sum per row with warp shuffles.
- Compensated products are 4 MMAs. The
al·blterm is about 2⁻²² relative. Measure whether 3 MMAs (dropping al·bl) keeps the 9 oracle cases and the near-constant-V case under 1e-5. If it does, that's 25% off every product.
Same evidence bar as before: a timed run at H24/Hkv4/D256/N2048/M75520 against the 11.498 s baseline, with the per-gradient error split.
|
Full-shape independent CPU reference PASS, candidate a0ea747. Command: test-backend-ops.exe test --test-perf-cases -o FLASH_ATTN_BACK -b CUDA0 -p 'hs=256,nh=24,nhkv=4,kv=75520,nb=2048,' Start2026-10-14T17:39:29.0085914Z, end17:59:21.1155838Z. Wall1192.1089634s; exit0;1/1 cases pass. 116 lifetime-bound samples17:39:59.364–17:59:18.033Z. GPU sampledmax27534MiB, observedOSpeakworking-set1555898368bytes. All healthprobesOK, max64ms, same servingPID16340. Samplerended17:59:28Z; test exited. Actual full GPU peak UNKNOWN: first30seconds unsampled and10second sample intervals. Health is not decode latency. No serving/context/deploy change. No speed, whole-model training, or learning acceptance. Raw logs: cuda-fa-mma-n2048.log, .receipt, -monitor.jsonl, -monitor.end. Raw working kernel SHA256 CFB0411B719CD9384724D685DB38A6209FF4AE39D7CB512B349BAC5E2A217E1D; LF normalized D0A9D6D72EA9B35366BBE8FAD394846311A9AF95ED86905D99DC08BECCDD9E70, HEADblob596c4c683e82245dbf8065e0c60ee8d5f9d231f7. Harness8F99199E18346F62F07CF64F416B8575AABA5786A7DDABC8D381408932ACBDE6. |
|
Exact candidate a0ea747 bounded target timing, separate from CPU correctness receipt: Command: test-backend-ops.exe perf -o FLASH_ATTN_BACK -b CUDA0 -p 'hs=256,nh=24,nhkv=4,kv=75520,nb=2048,' --perf-runs 1 Timed operation: 3,410,133us (3.410133s). Host elapsed13.4491131s, start2026-10-14T18:00:53.4592048Z end18:01:06.9070861Z, exit0. Existing synchronized host graph timer, not CUDA event timing. Compared to earlier11.498124s baseline:3.37175x faster. Compared to rejected24.210485s shared-load prototype:7.09957x faster. All are single executions under concurrent serving at different times; not controlled throughput evidence. Printed TFLOP count/rate is stale after the three-pass rewrite and must NOT be used; elapsed time remains the measurement. Fresh preflight GPUfree6492MiB, servinghealthOK, no compiler/test workload.22 samples18:00:43.350–18:01:05.862Z, sampledGPUmax27397MiB, healthmax31ms/no failures, servingPID16340 unchanged. Sampler ended18:01:06.935Z after completion receipt. Sampling covers the run but cannot establish true GPUpeak between one-second observations. PosthealthOK. Health is not decode latency. No whole-model/resident-training/learning acceptance:3.41s for this single operation is still substantial. No source or serving/context/deployment change, no merge. Raw cuda-fa-mma-target-perf.log/.receipt/-monitor.jsonl/-monitor.end retained in handoff directory. |
Fuse query normalization and weighted dP accumulation, using shared warp reductions and preserving direct F32 dS. Mark work estimates unavailable instead of reporting stale FLOP rates. Nine CPU oracle cases and the full N2048/M75520 q8 case pass at unchanged per-component tolerance. One bounded target timing measured2.642706s; whole-model acceptance remains open. Assisted-by: OpenAI Codex
|
Measured fused revision8798492d72b126d214ba40f13404f55a7f82c90b atop immutablea0ea74714. Shared query normalization/weighted-dP online reduction replaces two traversals with one; direct F32 dS and compensated MMA products stay. Existing scalar warp reducer now deduces float/double while preserving width as first template parameter. Harness reports work_estimate_available=false instead of stale rates. Native target build passed. Nine CPU-reference cases passed at unchanged per-component NMSE1e-5, including tiny dO and nearconstant V (dQ2.61368201e-6/dK2.56248292e-6). Full independent CPU oracle N2048/M75520,H24/Hkv4,D256,q8_0K/V,causal-tail PASS. Command: test-backend-ops.exe test --test-perf-cases -o FLASH_ATTN_BACK -b CUDA0 -p 'hs=256,nh=24,nhkv=4,kv=75520,nb=2048,'. Start18:29:06.5939663Z end18:48:53.0989217Z,1186.5062743s,exit0. Separate timing command uses perf with identical filter and --perf-runs1: ONE warmup and ONE timedoperation2,642,706us.18:28:03.8946919–18:28:15.5722918Z,host11.6787536s,exit0. This is22.5% lower than prior3.410133s single measurement, under different concurrent serving load; not controlled throughput. No during-run sampler was started for this brief timing; peak/during-health are UNKNOWN. Pre/posthealthOK. Raw cuda-fa-fused-perf.log/.receipt. Fulloracle monitor began before test:118samples18:28:56.477–18:48:45.630Z, sampledGPUmax27452MiB, observedOSpeakworking-set1555906560B. True GPUpeak unknown between10second samples.13 failedhealthprobes beginning18:45:56.252Z; serving16340 disappeared then26208 loaded at sameport/context. Root correlated this with supported581deployment handoff success18:46:24. This lane made no serving changes, and this run does NOT prove serving continuity. Current posthealthOK. Monitorended18:48:55.690Z; no leftover test or sampler. Exact raw working SHA256 preserved before/after commit: fattn-back.cu1E3085640C1644A8A241DF092B7B5E5EDDD60C1BDEC7A6BBCFC77C4C32FE905C; common.cuh100BCCC8D69A90C8BC1D3B0AE1A6E0705ECE246AFCE40B6FAEF2EC61A3C982AD; harness78370CAB15CE8135568024B725468295C959BE9D02D8D17CB271601D97FFBA8D. Git line-ending normalization remains separate from raw-file hashes. No merge/deploy or resident-training acceptance.2.64s for one operator remains substantial. Independent source review requested. |
joelteply
left a comment
joelteply
left a comment
There was a problem hiding this comment.
Cormac on 8798492 (fused, atop a0ea747): source-approved for correctness, resting on the full N2048/M75520 CPU oracle PASS you report. Faster but not a speed acceptance. Source review only.
Verified in the compare a0ea747…8798492:
Phase 0 is now ONE pass computing LSE and D together:
- One warp per row (r = threadIdx.y, stride blockDim.y = 4) and one lane per key.
warp_reduce_max/warp_reduce_sumover the 32 keys, and lane 0 rescalessumsanddeltasby the sameexp(m_old − next).delta = deltas·exp(m − lse)givesΣ p·dpexactly, sincelse = m + log Σ e^{s−m}.- The
next != −INFINITYguard keeps an all-masked tile from computingexp(−inf − −inf). - Rows are disjoint across warps, lane 0 reads the old
maxima[r]before writing it, and the cross-tile reads come after__syncthreads.
Phase 1 forms dS on all 128 threads:
lse = −inffori ≥ Nmakes p = 0, which short-circuits thea.delta[index]read for out-of-range rows.- Invalid rows write 0, replacing the old
thread < 16zeroing branch.
The Q kernel now walks M twice (QK + dO·Vᵀ each pass), down from three QK and two dO·Vᵀ.
warp_reduce_sum<width, T>: explicit-width callers still deduce T, and the existing more-specialized overloads (float2/half2) still win partial ordering. The 5090 build compiles it.
You flagged that the CPU's D uses expf(score − rounded lse), so the fused numerator isn't bit-equivalent. That's exactly why the full-shape oracle (not a pilot) is the gate, and it passed at the unchanged tolerance.
The harness now reports "work estimate unavailable" instead of the old FLOP count. That's honest.
Speed: 2.643 s vs 3.410 s (a0ea) vs 11.498 s (baseline), single runs under varying load, so the ratios are indicative, not a benchmark. Inference: ×16 full-attention layers is about 42 s of attention backward per step at this shape. Still open, each against the same full-shape oracle: 3 vs 4 compensated MMAs (measured, not blind, as you said), and reading LSE/D once per tile row in the KV kernel rather than per element.
Shared-operator integration and Metal parity remain Fable's, per your split.
1 participant
Overview
Adds CUDA support for the existing shared
FLASH_ATTN_BACKoperator, so a supported CUDA graph can use the native backward instead of composing score tiles. Two row-owned kernels consume F32/F16/q8_0 K/V directly and reuse the operator's packed gradient/LSE/D storage. GQA, outer strides, q/d component strides, independent K/V types and fully masked rows are preserved. Device launch limits are cached by the existing device-info owner.This draft targets Fable's
fable/p1-3b-run-accumulatorbranch to keep the contribution CUDA-only. Shared-operator integration remains coordinated with Fable; it is not a request to merge or deploy.Validation
Tested base:
88c5c1e26c9691debab279ae2ea529c12c758820.Contribution commit:
34dd2c1f9273c28c51dc4935b664d2ac0d7e150f.The committed binary diff against that base has SHA256
5098E7C1732636224C72AFC4BD0D2F2798E276F2A286463995BCDF0C0CE146D7, identical to the saved pre-commit tested patch. No source changes followed validation.Windows RTX5090, driver595.95, CUDA13.3.73, MSVC19.44.35229, Release/Ninja, CUDA architecture120. Built only
test-backend-ops, parallel2, completed2026-10-14 15:03:32Z, exit0.test-backend-ops.exe test -o FLASH_ATTN_BACK -b CUDA0ran15:04:12-15:04:13Z, exit0: six actual cases passed, no selected cases skipped, against the CPU reference:Existing acceptance threshold is NMSE1e-5. Actual passing maxima are not printed by the harness. The CPU backend filter skips a standalone CPU run, not the CPU oracle comparison.
test-backend-ops.exe perf -o FLASH_ATTN_BACK -b CUDA0ran15:04:23-15:05:05Z, exit0. Two explicit performance cases use q8_0 K/V, D128/H8/Hkv2/N64 and a causal mask:These are row-owned baseline timings under concurrent serving, not isolated GPU or whole-model training throughput. The first performance attempt had no matching cases because the harness has a separate perf list; it is excluded from these measurements. This change adds the two actual cases to that existing list.
Limits
No75k-context or27B-model training proof, resident serve/update/resume acceptance, long-run learning result, or claim that this baseline is fast enough. The shared operator remains single-sequence and does not add ALiBi, softcap or sink gradients. Formats/platforms outside the listed cases remain unverified. No serving/model/cache/job changes, restart or deployment; the existing server process stayed running. Full-suite and non-CUDA platform validation were not run.
Local full logs, exact commands and receipts:
C:/Users/joelt/.continuum/state/team-proof-20260921/cuda-fa-88c5-receipt.md, with sibling test/perf/build logs and source diff.git diff --checkpasses.Requirements