[Feat] Add GQA Sink Backward Varlen - #1640
Open
wch0810 wants to merge 4 commits into
Open
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
…of op full-pass
- Add example_gqa_sink_bwd_varlen.py: 12-kernel backward (k0 preprocess + k1,k2,k3c,k3,k3b,k4,k5c,k5,k5b + k6 postprocess) + forward kernel, Developer mode (pass_configs 4 True, 0 T.Scope/flag/barrier_all), Compensated GEMM (k3c/k3b for dV correction, k5c/k5b for dK correction) recovering fp16 cast precision to near-fp32, host pipeline with 11-kernel split + GM workspace
- Add test_gqa_sink_bwd_varlen.py: L0(8) + 3x determinism + L1(13) + L2(10) + Boundary(11) + do_bench + msprof op (sample + full-pass), 169-line precision standard (fp16 atol=6.10e-5, rtol=1.95e-3, max_abs_limit=0.1, required_ratio=0.99), check_precision dual-gate
- msprof op design: bwd-only target (host-side PyTorch ref_fwd_varlen avoids forward kernel launch contamination), --capture-mode {sample, full-pass} for per-kernel analysis and kernel-only total
- Verified: do_bench backward 266ms (full causal) / 221ms (SWA), msprof kernel-only total 221.26ms vs GPU 28574us (7.74x), L0 8/8 + L1 13/13 PRECISION_PASS, 7/7 bit-exact determinism
- CI compliance: ruff check/format passed, all --level modes output 'Test Passed!', AST normalization confirms kernel semantics unchanged
wch0810
force-pushed
the
feat/gqa_sink_bwd_varlen_pr
branch
from
August 20, 2026 03:54
4a67913 to
295f741
Compare
…er AUTO_CV_COMBINE Root cause: k3/k5a loaded P/dS via GM(fp32)→UB→cast→L1 (Vector→Cube virtual-channel transfer). Under AUTO_CV_COMBINE=True, this corrupted the L0A transpose path in transpose_A=True GEMM, causing only the first 16 rows (1 MMAD M-tile) to be computed. dQ ratio dropped from 0.9951 to 0.0010. Fix: k3c/k5c now write fp16 P/dS to new GM workspace (ws_p_fp16/ws_ds_fp16) during their existing cast step. k3/k5a load directly GM(fp16)→L1 (pure Cube path, matching k3b/k5b correction kernels and example_gqa_sink_bwd_bhsd pattern). Changes (5 spots, example_gqa_sink_bwd_varlen.py only): - k3c: add ws_p_fp16 output parameter + T.copy(p_half_ub, ws_p_fp16) - k3: replace UB cast with T.copy(ws_p_fp16, p_l1) direct GM→L1 - k5c: add ws_ds_fp16 output parameter + T.copy(ds_half_ub, ws_ds_fp16) - k5: replace UB cast with T.copy(ws_ds_fp16, ds_l1) direct GM→L1 - run_bwd_pipeline: add ws_p_fp16/ws_ds_fp16 workspace + update k3c/k3/k5c/k5 calls Verified: dQ=0.9951 dK=0.9987 dV=1.0000, L0 8/8 + L1 13/13 PASS, 7/7 bit-exact, do_bench bwd 233ms (was 266ms, -12.4%), msprof kernel-only 201ms (was 221ms, -9%)
wch0810
changed the title
Integrate all backward kernels into one kernel launch (flashattn_bwd_single), addressing reviewer feedback on PR tile-ai#1640: - Single kernel: 2 total (fwd + bwd), 1 backward launch (was 5 kernels / 34 launches) - On-chip: all intermediates (S/P/dP/dS/p_delta/ds_delta) stay in L1/UB/L0C - No GM workspace: p_delta/ds_delta via fp16 UB→L1 direct (8KB each) - No host sync: K-loop T.serial inside kernel, dSinks in-kernel Phase 6 - Compensated GEMM via L0C init=False accumulation (no separate correction kernel) - dSinks: kernel-internal T.tile.exp, golden uses kernel lse/Delta (ratio=1.0) - P_fp32 retention in s_ub across phases (eliminates GM roundtrip) Techniques adapted from example_gqa_sink_bwd_bhsd rev3 pattern. Performance: backward 114ms (was 274ms), 4.0x vs GPU (was 9.6x). Precision: L0 8/8 + L1 13/13 + determinism 7/7 bit-exact, all 0.99 blocking.
Contributor
Author
|
/re-test |
|
🔄 Re-running failed jobs Original workflow run: View details Only the failed jobs will be re-executed. |
hedi515
reviewed
Aug 25, 2026
| for b in range(B): | ||
| ks = cu_seqlens_k[b] | ||
| ke = cu_seqlens_k[b + 1] | ||
| # dK/dV 从 host .half() 来 (fp32 atomic_add + host cast) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Refactor GQA + Attention Sink Flash Attention Backward (Varlen) into a single-kernel backward. All backward computation (Delta + 5-GEMM + 2-softmax-bwd + Compensated GEMM + dSinks) is now fused into one
flashattn_bwd_singlekernel launch.Strategy: single kernel with
T.serialK-loop, on-chip direct CV handoff (L0C↔UB, independent buffer per handoff), Compensated GEMM via L0Cinit=Falseaccumulation (no separate correction kernel), dSinks computed in-kernel (Phase 6). UsescombineCV+auto_syncpasses (no manualT.Scope/T.barrier_all/cross_flag).Tested on Ascend NPU, fp16, B=8, H=64, G=16, q=k=2048, D=128, causal. Precision verified against PyTorch golden using 169-line mixed-tolerance standard (atol=6.10e-5, rtol=1.95e-3, ratio=0.99).
Performance (B=8, H=64, G=16, q=k=2048, D=128, fp16, causal)
Precision (169-line standard, all 4 outputs blocking)
Notes
T.serialinside kernel, no host for-loop, notorch.npu.synchronize()between stages.s_ubretains P_fp32 across phases (Phase 2 → Phase 4), eliminating GM roundtrip for dS compute.T.tile.exp, fp32 output. Golden recomputes using kernel's lse/Delta to isolateT.tile.expprecision (ratio=1.0)..half(): GQA groups=16 requires fp32 cross-head accumulation (fp16 atomic_add precision insufficient).example_gqa_sink_bwd_bhsd(rev3 single-kernel pattern).threads=1(AscendAUTO_CV_SYNCdoes not supportthreads=2for hybrid kernels withatomic_add— workspace race). This is the primary source of the 4.0x gap vs GPU (128 threads). The single-kernel on-chip architecture is correct; performance will improve once the framework supports multi-thread CV fusion.