[Perf] Parallelize AttnRes long-sequence reduction

[lbx154]

Summary

Improve the default Triton AttnRes path as one cumulative source change:

• specialize forward/backward kernels on the exact residual count instead of a padded pointer tuple;
• pass already-contiguous residual tensors directly through the autograd wrapper;
• save the forward output-RMSNorm reciprocal standard deviation for backward reuse;
• parallelize long-sequence dq/dw and optional dow reductions over the flattened token dimension.

For N >= 16384, four programs accumulate FP32 partials and a separate kernel finalizes dq, dw, and dow; shorter inputs retain the single-program reduction. Public API, checkpoint semantics, backend dispatch, and accumulation precision are unchanged.

Test plan

• pytest tests/ops/test_attnres.py -q -k T16384: 2 passed, 16 deselected
• pytest tests/ops/test_attnres.py -q: 17 passed, 1 skipped
• Hardware: 1x NVIDIA B200 (178.4 GiB)
• Software: CUDA 13.1, PyTorch 2.11.0a0+eb65b36914.nv26.02, Triton 3.6.0
• Benchmark: fused_attnres, bfloat16, fwdbwd, median latency; clean main baseline and candidate measured serially with identical shapes and benchmark settings

L	B	T	H	D	main (ms)	candidate (ms)	speedup
10	1	32768	1	4096	3.4738880395889282	3.268607974052429	1.0628035136566079x
10	1	8192	1	4096	0.9319039881229401	0.902176022529602	1.032951402886971x
64	1	8192	1	8192	12.726335525512695	11.95311975479126	1.0646873608382867x
8	1	32768	1	2048	1.45305597782135	1.295311987400055	1.1217806921851454x
8	1	8192	1	2048	0.5739999711513519	0.4638400077819824	1.237495605211243x

Geometric-mean speedup: 1.1016441159813262x. Worst-row speedup: 1.032951402886971x.

Scope and limitations

The implementation change is limited to fla/ops/attnres/fused.py; the regression coverage adds two bf16 threshold cases in tests/ops/test_attnres.py. Forward math is unchanged, but the Triton forward path now uses exact-L specialization and stores o_rstd when output RMSNorm is enabled. The Gluon backend and short reduction topology are unchanged. The long path uses temporary FP32 accumulation buffers. Performance was measured only for the five bfloat16 fwdbwd shapes above on B200; no broader hardware or end-to-end model speedup is claimed. Measurements are warmed steady-state timings; compile latency was not measured, and exact-L specialization can create more JIT/autotune variants when a model uses many distinct residual counts.

Breaking changes

None.

Checklist

• I have read CONTRIBUTING.md and follow its conventions (code style, docstrings, commit prefixes).
• I have read AGENTS.md and, where my change matches its scope, the relevant skill under .agents/skills.
• This is not a minor/cosmetic-only PR (typo, formatting, style-only tweaks).
• Dependent tests pass locally or in CI; new behavior is covered by tests where applicable.
• Kernel changes include same-hardware before/after benchmark numbers (dense + varlen where applicable).

[Copilot]

Pull request overview

This PR improves the CUDA Triton AttnRes implementation by parallelizing the long-sequence backward reduction over the flattened token dimension, reducing a serial bottleneck for large N while keeping the public API and dispatch behavior unchanged.

Changes:

• Added a split-N reduction path for backward dq/dw/(dow) where N >= 16384, using FP32 partial accumulation and a finalize kernel.
• Stored per-token output RMSNorm inverse-std (o_rstd) in forward and reused it in backward to avoid recomputation.
• Removed the padded residual pointer-table (L2) approach and specialized kernels directly on L.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[lbx154]

Thanks for checking this. The pre-existing implementation already saved optional output_rms_weight and o_pre values through save_for_backward; this PR adds o_rstd using the same supported None behavior. The output_rms_weight=None / checkpoint_level=1 paths are covered locally, including the new bf16 N=16384 case, and the updated H100 PyTorch 2.12 test-ops job passes.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

Suppressed comments (2)

fla/ops/attnres/fused.py:154

• Same concern as in the fwd kernel: making L a tl.constexpr here can cause the range(tl.cdiv(L, BL)) loops to be compile-time sized and potentially fully unrolled per-L, increasing compile time and kernel code size. If the intent is only to enable static res[i] indexing, consider separating tuple length from the runtime loop bound to avoid unrolling the full residual-source dimension.

    N,
    L: tl.constexpr,
    D: tl.constexpr,
    eps: tl.constexpr,
    scale: tl.constexpr,

fla/ops/attnres/fused.py:53

• L is now a tl.constexpr parameter. This makes loops like for i_l in range(tl.cdiv(L, BL)) compile-time sized, which can fully unroll the residual-source loop for each distinct L and BL config, significantly increasing Triton compile time / code size and multiplying cached variants (autotune key is also per-L). If L can vary across layers/models, consider keeping the outer loop runtime-sized (and bucketing the pointer tuple length separately) to reduce JIT overhead while still enabling static indexing into res[i].

This issue also appears on line 150 of the same file.

    N,
    L: tl.constexpr,
    D: tl.constexpr,
    eps: tl.constexpr,
    scale: tl.constexpr,

[lbx154]

Regarding the exact-L specialization noted in review: this is intentional because Triton tuple element selection is static, and using the real residual count removes the padded pointer tuple and enables the measured steady-state gains. The tradeoff is more JIT/autotune variants when models use many distinct L values. The benchmark table reports warmed steady-state latency and does not claim compile-time improvement; I have called that scope out in the PR. The local parity matrix covers multiple residual counts, including L=1,3,7,10,15,29, while the performance matrix includes L=8,10,64. If maintainers prefer a bucketing tradeoff to limit specialization variants, I can adjust that part separately.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

[lbx154]

Hi @zhiyuan1i, could you please review this PR when you have a chance? Thanks!