[Perf] TileLang backends for linear-attention chunk backward (simple_gla / gla / delta_rule / gdn) - #1046
Conversation
Shared hand-written TileLang backend for FLA linear-attention chunk backward (dqkwg staging + chunk_h), wired into the backend dispatch. Includes a CUDA 12.6/12.7 __nv_fp8_e8m0 stub so TileLang's FP8 header compiles on this toolchain. Backend is gated by FLA_TILELANG and falls back to Triton when TileLang has no usable nvcc.
V-first direct-state dqkwg with K-inner staging and on-chip shadow-state reuse, fusing the hdh epilogue to cut HBM round trips. Frozen N=3 scorer, pure-Triton baseline (FLA_TILELANG=0 vs 1) on H100 NVL: fwdbwd geomean +16.1%, all-row +8.0%, peak 1.22x; fwd+bwd NaN-poison gate green.
Op-local fused K-tile backward consumer for dense D128/D256, keeping the existing dh/dv/dA producers. Frozen N=3 scorer, pure-Triton baseline on H100 NVL: fwdbwd geomean +20.0%, all-row +9.2%, peak 1.32x; gate green.
D256-focused fused producer-consumer WY-backward route, collapsing the WY representation backward into one on-chip pipeline to remove dhu / WY-backward HBM round trips. Frozen N=3 scorer (verify --base main): fwdbwd geomean +5.6%, all-row +4.4%, D256 peak 1.155x; gate green.
Split rewrite of the gated delta net (GDN) prepare_wy backward, compiled via the CUDA 12.8 toolkit (Hopper FP8 e8m0 + TMA). Frozen N=3 scorer: fwdbwd geomean +3.7%, all-row +2.1%, three repeats all GATE PASSED; correctness green.
|
Thanks to the all-powerful agents. This MR is simply enormous… |
|
Please restore the It was added deliberately in #944: without it, this kernel misbehaves on the triton-ascend backend (a timing/synchronization issue in that Triton variant — the barrier orders the preceding If you have evidence the underlying backend issue is fixed (a specific Triton/backend changelog), please gate the removal on a minimum-version check instead of deleting the barrier outright. |
The barrier was added deliberately in fla-org#944: without it, triton-ascend can misorder the b_dA update before the dk load-modify-store in the K-loop, failing silently on that backend. Restore it with a comment so it is not removed again as an NVIDIA-only-looking no-op. Patch from @lbx154, applied via maintainer edit.
|
Follow-up on the merge conflict in #1046: I resolved the current conflict locally and pushed the prepared branch here: chen-xiao-yu-1:pr-1046-conflict-resolved 4906a2e — merge latest origin/main into the PR branch and resolve the conflict |
|
The current merge conflict is resolved in markovchain-builder#1. |
Resolve chunk_h.py: main migrated off tl.make_block_ptr to raw pointer arithmetic; keep the STORE_MMA_STATE dual-write and FUSE_HDH_LAST paths, rewritten with the same raw-pointer + mask style.
zhiyuan1i
left a comment
zhiyuan1i
left a comment
There was a problem hiding this comment.
Thanks for your patience on this one. I've gone ahead and pushed two things to the branch to help move it along:
- Resolved the textual merge conflict in
fla/ops/common/chunk_h.py(main's raw-pointer migration × yourSTORE_MMA_STATEdual-write andFUSE_HDH_LASTpaths, rewritten in the raw-pointer + mask style; cross-checked against @chen-xiao-yu-1's resolution — semantically equivalent). - Restored the
tl.debug_barrier()inprepare_wy_repr_bwd_kernelwith @lbx154's patch from the comments (load-bearing on triton-ascend, now with a why-comment).
However, CI on the merged head shows this PR now needs a rebase-level adaptation to main, which is beyond a conflict fix and really needs the author:
-
test-ops fails: the NPU call sequence in
fla/ops/gated_delta_rule/wy_fast.py(the per-kernelprepare_wy_repr_bwd_k_npu/_v_npu/_da_mask_npu/ ... launches) no longer matches main's triton_ascend backend, which was restructured into a fusedprepare_wy_repr_bwd_kv_npuwith different parameters and grid semantics (task_num/num_corepersistence,G_EXP_PRECOMP, etc.). The test failure isAttributeError: ... no attribute 'prepare_wy_repr_bwd_k_npu'. The dispatch sequence needs to be rewritten against the new fused backend (or dropped in favor of the existing backend dispatch). -
lint fails: main added a
no-make-block-ptrpre-commit hook (kernels migrated offtl.make_block_ptr/tl.advancefor triton-main compatibility).fla/ops/delta_rule/backends/triton/chunk_bwd.pyon this branch still has 19 block_ptr uses that need migrating to the raw-pointer + mask style. -
check-pr-title fails: main now requires the PR body checklist to be fully ticked. Please refresh the body against the current template.
Could you rebase the branch onto current main and address these three? Happy to re-review quickly once CI is green.
|
@zhiyuan1i I prepared the requested fixes on Could you merge current
The prepared branch preserves the reviewed The PR body checklist also needs to be refreshed against the current template. |
4 participants
Summary
Hand-written TileLang backends for the backward path of four core linear-attention chunk kernels, gated by
FLA_TILELANGwith automatic Triton fallback. All four pass the frozen fwd+bwd gate (NaN-poison) and show measured speedups over the existing Triton kernels on H100 NVL.Results (N=3 geomean, fwd+bwd)
chunk_glaFLA_TILELANG=0vs=1)chunk_simple_glachunk_delta_ruleverify --base mainchunk_gdnverify --base mainHow
All speedups come from the backward path — fusing the separate Triton backward kernels (dq/dk/dw/dg + WY-representation backward) into single producer–consumer TileLang pipelines that keep intermediate state on-chip and cut HBM round trips:
chunk_h) + a CUDA 12.6/12.7__nv_fp8_e8m0stub so TileLang's FP8 header compiles on this toolchain.prepare_wybackward (compiled via CUDA 12.8 for Hopper FP8 e8m0 + TMA).Methodology
tests/ops/test_*.py(fwd+bwd, NaN memory poisoning), unchanged; naive references unchanged.benchmarks/ops/run.py/verify.py, N=3, geomean, identical shape/precision/numeric flags on both sides.FLA_TILELANG=1and falls back to Triton when no usable nvcc is present.