Repository navigation
[8/8] perf(models): per-arch short-conv and KDA backward kernels for linear attention - #3634
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude Code Review
This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.
Tip: disable this comment in your organization's Code Review settings.
33bbef9 to
e1f583e
Compare
e1f583e to
0cda557
Compare
0cda557 to
388d7a9
Compare
388d7a9 to
61ea268
Compare
0de9f79 to
ff6193f
Compare
Zhichenzzz
added
run-ci-megatron
run-ci-precision
run-ci-low-precision
run-ci-model-scripts
5173f58 to
507bd28
Compare
507bd28 to
28ba0a8
Compare
28ba0a8 to
cbedd6d
Compare
cbedd6d to
17f641e
Compare
Zhichenzzz
changed the title
17f641e to
7a81d7f
Compare
|
Running the three files whose home suite is |
|
/rerun-test tests/fast-gpu/test_nvfp4_quantizer.py |
|
/rerun-test tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py |
|
/rerun-test tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4_w4a16.py |
|
⚪ The run was cancelled. |
|
⚪ The run was cancelled. |
|
⚪ The run was cancelled. |
|
/bypass-fastfail |
7a81d7f to
4dd7a68
Compare
radixark-miles-ci
Bot
added
the
bypass-fastfail
4dd7a68 to
2e5d755
Compare
|
|
…d backward The head-sharded short conv used fla's Triton conv for forward and backward. When causal-conv1d's CUDA backward imports (it does in the miles image), it now uses fla's "mix" backend: Triton forward, causal-conv1d CUDA backward. FLA_CONV_BACKEND overrides the choice, as it does for fla's own ShortConvolution, and short_conv_backend logs it once. The CP path is unaffected; fla routes it to its own CP conv whatever the backend. GB300, one conv at Kimi-K3 shape (12288 channels), forward + backward: 64k tokens 7.45 -> 5.27 ms (one sequence) and 7.43 -> 4.31 ms (three packed sequences); 16k 2.12 -> 1.17 ms. Against a float64 reference the mix gradients are slightly closer than Triton's (dx 1.66e-3 vs 1.76e-3; outputs are identical). The conv test gains a float64 torch reference over packed sequences, run on Triton and on the selected backend.
fla dispatches the fused KDA backward (chunk_kda_bwd_wy_dqkg_fused) to its TileLang kernel whenever TileLang is installed. On GB300 that kernel is slower than fla's own Triton implementation at every tiling fla offers; the best TileLang config (BK=BV=128, 8 warps) still trails Triton. kda_kernel now defaults FLA_TILELANG=0 on compute capability >= 10; an explicit FLA_TILELANG still wins, and the value is logged once. Hopper keeps fla's default until it is measured there. GDN has no TileLang backend. GB300, chunk_kda forward + backward at Kimi-K3 shape (96 heads x 128): 64k tokens 77.2 -> 66.9 ms, 16k 21.5 -> 17.0 ms; gradients match the TileLang path to 6e-5 (g) and bitwise-or-1e-9 (q, k, v). Whole head-sharded layer, 64k tokens: Kimi-K3 KDA TP=1 194.8 -> 180.3 ms, TP=4 58.3 -> 55.7 ms, CP=4 60.3 -> 57.7 ms; GLM-5.3-Flash KDA TP=1 99.9 -> 92.4 ms. tests/fast/models/test_linear_attn_kernel_selection.py covers the Blackwell default, Hopper, the explicit override and the FLA_CONV_BACKEND override on CPU.
|
Pushed two fixes at the bottom of the stack (rebased on main
|
2 participants
What
Two backend choices for the linear-attention layer, measured on GB300. Both stay within fla's own kernels.
mixbackend (Triton forward, causal-conv1d CUDA backward) when causal-conv1d is installed;FLA_CONV_BACKENDoverrides.FLA_TILELANG=0by default on compute capability ≥ 10, because fla's TileLang dqkg kernel is slower than its Triton one there. An explicit setting wins; Hopper keeps fla's default.vs current main (GB300, 2026-10-08)
K3 KDA layer (hidden 7168, 96 x 128 heads), main's
KimiK3Attentionvs this stack's tip, forward + backward, SP when TP > 1:This stack's tip includes this PR; see #3632 for the split.
Kernel measurements (GB300)
One short conv at Kimi-K3 shape (12288 channels), forward + backward:
Outputs are identical. The float64 comparison was run on two packed sequences, 16k tokens in total.
chunk_kdaforward + backward at Kimi-K3 shape (96 heads × 128):The Triton gradients match TileLang's to 6e-5 for g, and bitwise or within 1e-9 for q, k and v.
Layer measurements (GB300)
tests/manual/bench_linear_attn.py, head-sharded layer, forward + backward. The #3632 tip (99d3445db) and this branch (33bbef952) ran back to back on an otherwise idle node.* In the base run, two cells were outliers: GLM TP=4 64k at 155 ms and GDN CP=4 16k at 14.7 ms. For GLM TP=4 64k, the base value shown is from the preceding run with TileLang on. The GDN CP=4 16k cell has no clean base value and is left out.
Against Kimi-K3's KDA layer before #3632, with identical weights and inputs (
bench_k3_xcommit.py, forward + backward, parameter gradients compared afterfinalize_model_grads' TP sum):Hopper (H200)
FLA_TILELANGis left at fla's default on Hopper because there, fla's TileLang dqkg is slightly faster than Triton:chunk_kdaforward + backward, Kimi-K3 shapeThe last row is 12% faster than fla's default tiling on H200, with bitwise-equal gradients. That tiling is hard-coded inside fla, so it belongs upstream in fla rather than in miles.
The short conv on H200 at Kimi-K3 shape (12288 channels) goes from 9.53 ms (Triton) to 7.01 ms (mix, one sequence) and 5.86 ms (mix, three packed sequences) at 64k tokens, and from 2.43 ms to 1.55 ms at 16k.
Head-sharded layer, TP=1, forward + backward, #3632 tip → this branch:
The Hopper numbers were measured with Megatron
miles-mainat96fe16ecb.Precision and tests (GB300,
33bbef952)test_linear_attn_short_conv.py(new float64 reference on Triton and mix, plus chunked vs single call) +tests/fast/modelstest_linear_attn_head_sharded.py: Qwen3.5 / Qwen3-Next / KDA vs replicated, TP=2 and TP=4test_qwen3_5_cp_correctness.py, CP=2: GDN / KDAtest_qwen3_5_cp_correctness.py, CP=4: GDN / KDAThe CP tests compare the layer with CP against the same layer without it. The dx distance grew from about 1e-3 (GDN) and 3.7e-3 (KDA) because the two sides now run different conv backwards: mix without CP, fla's Triton CP conv with CP. Each side is the same distance from float64 as before.
Stack
miles/kernelsmiles/kernels/attentionmiles_plugins/models/linear_attn.py)Rules for
miles/kernelsare in its README: only kernels we write, one directory per op, names say what is computed (never the backend), callers import the entry, a torch-reference test per op.