Skip to content

[4/8] perf(kernels): per-arch DSA sparse attention dispatch with a FlashMLA forward - #3608

Merged
Zhichenzzz merged 1 commit into
zhichen/dsa-merge-kernelsfrom
zhichen/dsa-faster-kernels
Oct 8, 2026
Merged

[4/8] perf(kernels): per-arch DSA sparse attention dispatch with a FlashMLA forward#3608
Zhichenzzz merged 1 commit into
zhichen/dsa-merge-kernelsfrom
zhichen/dsa-faster-kernels

Conversation

@Zhichenzzz

@Zhichenzzz Zhichenzzz commented Sep 22, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

What

  • Forward: FlashMLA's sparse prefill kernel or the TileLang kernel, chosen per shape. FlashMLA (heads padded to 64) wins every GLM-5 shape and everything on Blackwell; TileLang wins DeepSeek-V4.1 at 8–16 heads on H200.
  • Padding: the host passes kv_indices = indices.clamp(min=0); both kernels gather through it and mask from the original indices. No out-of-bounds access, no per-element clamp.
  • The TileLang forward runs unpipelined. With the unconditional gather, TileLang's pipeline read the wrong stage when the block count was stages + 1 (16 heads, topk 128 / 192): rows came back as 1e30. main's DeepSeek-V4.1 forward has the same pattern and is wrong at 16 heads, topk 128; this PR replaces it.
  • Dynamic shapes: the backward no longer recompiles per packed length (main's V4.1 copy spends 6–8 s per new length on H200).
  • SparseAttentionConfig: forward backend, forward threads, backward stages/split_store, from a measured table keyed by (arch, RoPE tail, TP-local heads). B300 and GB300 pick the same winners; H200 differs. Callers may pass their own.
  • tests/manual/bench_dsa.py benchmarks GLM-5.3-Flash, DeepSeek-V4.1-Flash and GLM-5 at 4k–128k tokens.

vs current main (2026-10-08, this push)

Forward + backward per layer; "main" is the kernel each model runs on main (b7b2870daf). 16k / 64k tokens, idle GPUs.

Model, TP-local heads H200 B300 GB300
GLM-5 / 5.3-Flash, 8 (TP=8) 3.40x / 3.28x 3.02x / 3.15x 3.01x / 3.15x
GLM-5 / 5.3-Flash, 16 (TP=4) 3.24x / 3.02x 3.02x / 3.15x 3.01x / 3.15x
GLM-5 / 5.3-Flash, 64 (TP=1) 9.41x / 10.06x 5.43x / 6.37x 5.42x / 6.38x
DeepSeek-V4.1 vs V4 kernel, 8 (TP=8) 1.12x / 1.18x 1.89x / 2.01x 1.88x / 2.01x
DeepSeek-V4.1 vs V4 kernel, 32 (TP=2) 1.32x / 1.51x 1.88x / 2.03x 1.88x / 2.03x
DeepSeek-V4.1 vs V4 kernel, 64 (TP=1) 3.27x / 2.96x 3.27x / 3.42x 3.27x / 3.42x
DeepSeek-V4.1 vs its own copy, 8 (TP=8) 0.83x / 0.88x 1.61x / 1.66x 1.60x / 1.67x
DeepSeek-V4.1 vs its own copy, 32 (TP=2) 1.03x / 1.03x 1.70x / 1.74x 1.70x / 1.74x
DeepSeek-V4.1 vs its own copy, 64 (TP=1) 3.13x / 2.85x 3.57x / 3.71x 3.57x / 3.71x

The H200 8-head V4.1 row loses to the copy's static-shape backward at a fixed length; at a new length the copy recompiles for 6–8 s.

Correctness

  • test_dsa_sparse_attention.py 89 passed, test_dsa_indexer.py 13 passed (H200, B300, GB300). New cases: odd lengths, 32 heads, packed sequences end to end (causal_ranges → indexer → top-k → attention, 1-token segments).
  • Wider matrix, fwd + bwd vs fp32, 3 repeats: 5 head counts × 6 topk × {packed, 1-token segments, causal, batch 2} × {tail, shuffled} padding × 8 configs. 0 failures on H200, B300, GB300; the only skips are stage-2 backward configs that exceed shared memory at 64 / 128 heads, which the table never picks.
  • Odd lengths (4097–30011) and packed sequences up to 16k vs main's kernels: out / dq bitwise or within 1e-6 on the TileLang path, 3e-3 on FlashMLA; dkv within 1e-4; all finite.

Summary of an earlier run (H200, before the rebase)

Forward + backward per layer, before this stack versus after:

Config Before After Speedup
GLM-5.3-Flash TP=8, 4k 82.6 ms 37.4 ms 2.2x
GLM-5.3-Flash TP=8, 128k 1492 ms 940 ms 1.6x
GLM-5.3-Flash TP=1, 4k 208 ms 37.9 ms 5.5x
GLM-5.3-Flash TP=1, 128k 8787 ms 1073 ms 8.2x
DeepSeek-V4.1-Flash TP=8, 4k 13.8 ms 5.2 ms 2.7x
DeepSeek-V4.1-Flash TP=8, 128k 405 ms 219 ms 1.85x
DeepSeek-V4.1-Flash TP=1, 4k 53.1 ms 7.8 ms 6.8x
DeepSeek-V4.1-Flash TP=1, 128k 1948 ms 301 ms 6.5x

Precision:

  • Both forward backends are within 2e-6 relative of a dense fp32 reference, the same as the old kernels.
  • The two backends agree to 5e-7 on out, dq and dkv.
  • The TileLang forward's output and lse are bitwise unchanged, and dq is bitwise unchanged.
  • dkv differs at the 4e-7 level, which is fp32 atomic ordering.
  • Every test runs on both forward backends.

This branch (91c99cd02)

Test H200 GB300 (sm_103, FlashMLA built for sm_100f; run at 91c99cd02)
test_dsa_sparse_attention.py, both forward backends 41 passed pass
test_dsa_indexer.py 13 passed pass
tests/fast/test_dsv4_thd.py 88 passed –
test_dsv4_thd_cp_correctness.py, CP=2 pass –

H200 results

8×H200 on the current radixark/miles:dev image (pulled 2026-09-23): torch 2.13 + CUDA 13.0, sglang 0.5.21.dev55, tilelang 0.1.14, fla 0.5.2, FlashQLA 0.1.2, FlashMLA; Megatron-LM 7a1f7fe3f (the image's); this stack at dd2cc6f6b.

Forward + backward per layer, same columns as the GB300 table below.

Config TP Tokens Before fwd+bwd After fwd+bwd Speedup vs fp32 / backends agree
glm5.3-flash (as in #2786: 64-wide zero tail) 8 4k 82.6 32.37 2.6x tilelang 2e-06 flash_mla 2e-06 / out 2e-07 dq 2e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 8 8k 124.87 47.18 2.6x tilelang 2e-06 flash_mla 2e-06 / out 3e-07 dq 2e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 8 16k 208.45 77.81 2.7x / out 2e-07 dq 2e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 8 32k 380.65 158.5 2.4x / out 1e-07 dq 4e-07 dkv 1e-07
glm5.3-flash (as in #2786: 64-wide zero tail) 8 64k 744.15 351.9 2.1x / out 3e-07 dq 3e-07 dkv -1e-07
glm5.3-flash (as in #2786: 64-wide zero tail) 8 128k 1484.67 764.51 1.9x / out 1e-07 dq 3e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 4 4k 82.8 32.58 2.5x tilelang 2e-06 flash_mla 2e-06 / out 2e-07 dq 4e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 4 8k 124.7 47.8 2.6x tilelang 2e-06 flash_mla 2e-06 / out 2e-07 dq 4e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 4 16k 208.96 78.83 2.7x / out 2e-07 dq 4e-07 dkv -1e-07
glm5.3-flash (as in #2786: 64-wide zero tail) 4 32k 381.61 160.7 2.4x / out 2e-07 dq 3e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 4 64k 744.17 354.92 2.1x / out 3e-07 dq 3e-07 dkv -1e-07
glm5.3-flash (as in #2786: 64-wide zero tail) 4 128k 1487.62 775.63 1.9x / out 3e-07 dq 3e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 1 4k 198.97 38.24 5.2x tilelang 2e-06 flash_mla 2e-06 / out 3e-07 dq 3e-07 dkv 1e-07
glm5.3-flash (as in #2786: 64-wide zero tail) 1 8k 427.32 62.18 6.9x tilelang 2e-06 flash_mla 2e-06 / out 2e-07 dq 3e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 1 16k 927.46 114.24 8.1x / out 2e-07 dq 2e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 1 32k 1966.35 239.5 8.2x / out 2e-07 dq 4e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 1 64k 4108.12 515.17 8.0x / out 1e-07 dq 4e-07 dkv 1e-07
glm5.3-flash (as in #2786: 64-wide zero tail) 1 128k 8438.84 1094.33 7.7x / out 2e-07 dq 1e-07 dkv 0e+00
glm5.3-flash (no zero tail) 8 4k 80.7 31.31 2.6x tilelang 2e-06 flash_mla 2e-06 / out 3e-07 dq 4e-07 dkv 0e+00
glm5.3-flash (no zero tail) 8 8k 119.17 44.66 2.7x tilelang 2e-06 flash_mla 2e-06 / out 2e-07 dq 2e-07 dkv 0e+00
glm5.3-flash (no zero tail) 8 16k 195.31 71.69 2.7x / out 4e-07 dq 3e-07 dkv 1e-07
glm5.3-flash (no zero tail) 8 32k 348.85 138.68 2.5x / out 2e-07 dq 2e-07 dkv 0e+00
glm5.3-flash (no zero tail) 8 64k 673.99 303.2 2.2x / out 2e-07 dq 4e-07 dkv 2e-07
glm5.3-flash (no zero tail) 8 128k 1342.77 661.1 2.0x / out 3e-07 dq 4e-07 dkv 0e+00
glm5.3-flash (no zero tail) 4 4k 80.2 31.56 2.5x tilelang 2e-06 flash_mla 2e-06 / out 3e-07 dq 3e-07 dkv 0e+00
glm5.3-flash (no zero tail) 4 8k 119.01 45.08 2.6x tilelang 2e-06 flash_mla 2e-06 / out 3e-07 dq 4e-07 dkv 0e+00
glm5.3-flash (no zero tail) 4 16k 195.23 72.51 2.7x / out 2e-07 dq 3e-07 dkv 0e+00
glm5.3-flash (no zero tail) 4 32k 348.55 140.07 2.5x / out 2e-07 dq 2e-07 dkv 0e+00
glm5.3-flash (no zero tail) 4 64k 671.33 305.95 2.2x / out 3e-07 dq 3e-07 dkv 1e-07
glm5.3-flash (no zero tail) 4 128k 1338.78 669.58 2.0x / out 4e-07 dq 4e-07 dkv 0e+00
glm5.3-flash (no zero tail) 1 4k 165.56 35.62 4.6x tilelang 2e-06 flash_mla 2e-06 / out 2e-07 dq 4e-07 dkv -1e-07
glm5.3-flash (no zero tail) 1 8k 336.65 55.53 6.1x tilelang 2e-06 flash_mla 2e-06 / out 2e-07 dq 2e-07 dkv 0e+00
glm5.3-flash (no zero tail) 1 16k 707.03 96.57 7.3x / out 2e-07 dq 3e-07 dkv 0e+00
glm5.3-flash (no zero tail) 1 32k 1468.96 194.31 7.6x / out 2e-07 dq 2e-07 dkv 1e-07
glm5.3-flash (no zero tail) 1 64k 3037.81 418.77 7.3x / out 1e-07 dq 4e-07 dkv -1e-07
glm5.3-flash (no zero tail) 1 128k 6224.98 893.5 7.0x / out 3e-07 dq 4e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 8 4k 13.84 5.23 2.6x tilelang 2e-06 flash_mla 2e-06 / out 4e-07 dq 4e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 8 8k 25.59 9.54 2.7x tilelang 2e-06 flash_mla 2e-06 / out 4e-07 dq 4e-07 dkv 6e-08
dsv4.1-flash ratio-1 layers 8 16k 49.22 20.02 2.5x / out 4e-07 dq 4e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 8 32k 98.94 45.78 2.2x / out 4e-07 dq 5e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 8 64k 201.15 102.51 2.0x / out 4e-07 dq 5e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 8 128k 408.0 220.81 1.8x / out 4e-07 dq 5e-07 dkv 1e-07
dsv4.1-flash ratio-1 layers 2 4k 21.37 8.81 2.4x tilelang 2e-06 flash_mla 2e-06 / out 4e-07 dq 4e-07 dkv 6e-08
dsv4.1-flash ratio-1 layers 2 8k 41.53 17.38 2.4x tilelang 2e-06 flash_mla 2e-06 / out 4e-07 dq 4e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 2 16k 82.33 37.09 2.2x / out 4e-07 dq 5e-07 dkv -1e-07
dsv4.1-flash ratio-1 layers 2 32k 165.38 80.9 2.0x / out 4e-07 dq 5e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 2 64k 333.35 173.66 1.9x / out 5e-07 dq 5e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 2 128k 669.49 364.61 1.8x / out 5e-07 dq 4e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 1 4k 54.17 7.72 7.0x tilelang 2e-06 flash_mla 2e-06 / out 4e-07 dq 4e-07 dkv -1e-07
dsv4.1-flash ratio-1 layers 1 8k 111.0 14.95 7.4x tilelang 2e-06 flash_mla 2e-06 / out 4e-07 dq 4e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 1 16k 230.95 31.39 7.4x / out 4e-07 dq 5e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 1 32k 475.6 67.53 7.0x / out 4e-07 dq 5e-07 dkv -1e-07
dsv4.1-flash ratio-1 layers 1 64k 975.06 142.99 6.8x / out 5e-07 dq 4e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 1 128k 1971.65 302.39 6.5x / out 3e-07 dq 5e-07 dkv 0e+00
dsv4.1-flash ratio-0 layers 8 4k 15.82 5.98 2.6x tilelang 2e-06 flash_mla 2e-06 / out 3e-07 dq 5e-07 dkv 0e+00
dsv4.1-flash ratio-0 layers 8 8k 27.31 10.25 2.7x tilelang 2e-06 flash_mla 2e-06 / out 4e-07 dq 4e-07 dkv 0e+00
dsv4.1-flash ratio-0 layers 8 16k 50.81 18.9 2.7x / out 4e-07 dq 4e-07 dkv 1e-07
dsv4.1-flash ratio-0 layers 8 32k 98.14 40.08 2.4x / out 4e-07 dq 4e-07 dkv 0e+00
dsv4.1-flash ratio-0 layers 8 64k 197.93 92.17 2.1x / out 4e-07 dq 4e-07 dkv -1e-07
dsv4.1-flash ratio-0 layers 8 128k 402.15 204.45 2.0x / out 5e-07 dq 5e-07 dkv 0e+00
glm5 8 4k 79.05 29.83 2.7x tilelang 2e-06 flash_mla 2e-06 / out 3e-07 dq 2e-07 dkv 0e+00
glm5 8 8k 119.75 43.63 2.7x tilelang 2e-06 flash_mla 2e-06 / out 3e-07 dq 2e-07 dkv 1e-07
glm5 8 16k 201.49 72.87 2.8x / out 2e-07 dq 3e-07 dkv -1e-07
glm5 8 32k 368.63 150.22 2.5x / out 2e-07 dq 4e-07 dkv 0e+00
glm5 8 64k 721.18 334.95 2.2x / out 3e-07 dq 3e-07 dkv 0e+00
glm5 8 128k 1441.32 736.58 2.0x / out 3e-07 dq 4e-07 dkv 0e+00
glm5 4 4k 79.14 29.77 2.7x tilelang 2e-06 flash_mla 2e-06 / out 3e-07 dq 4e-07 dkv 0e+00
glm5 4 8k 120.2 44.04 2.7x tilelang 2e-06 flash_mla 2e-06 / out 4e-07 dq 3e-07 dkv -1e-07
glm5 4 16k 202.13 73.76 2.7x / out 2e-07 dq 4e-07 dkv 1e-07
glm5 4 32k 369.45 151.98 2.4x / out 2e-07 dq 3e-07 dkv 1e-07
glm5 4 64k 721.26 339.9 2.1x / out 2e-07 dq 3e-07 dkv 0e+00
glm5 4 128k 1446.7 744.71 1.9x / out 2e-07 dq 2e-07 dkv 0e+00
glm5 1 4k 191.87 35.3 5.4x tilelang 2e-06 flash_mla 2e-06 / out 2e-07 dq 2e-07 dkv 0e+00
glm5 1 8k 416.54 58.08 7.2x tilelang 2e-06 flash_mla 2e-06 / out 3e-07 dq 4e-07 dkv 0e+00
glm5 1 16k 900.45 107.2 8.4x / out 2e-07 dq 3e-07 dkv 1e-07
glm5 1 32k 1904.32 225.57 8.4x / out 2e-07 dq 2e-07 dkv 0e+00
glm5 1 64k 3978.33 483.78 8.2x / out 3e-07 dq 3e-07 dkv -1e-07
glm5 1 128k 8163.86 1022.31 8.0x / out 2e-07 dq 4e-07 dkv 0e+00

GB300 results

8×GB300 (sm_103) on two 4-GPU nodes, GPUs otherwise idle; torch 2.13 + CUDA 13.0, tilelang 0.1.12, fla 0.5.2, FlashMLA b7643bd built for sm_100f; commits as pushed to this stack.

Forward + backward per layer. "Before" is the merged kernel before this PR (#3607, bitwise equal to main's kernels); "after" is this PR with the default forward backend (FlashMLA where it applies). The last column is the relative error against a dense fp32 reference (up to 8k tokens) and between the two forward backends.

Config TP Tokens Before fwd+bwd After fwd+bwd Speedup vs fp32 / backends agree
glm5.3-flash (as in #2786: 64-wide zero tail) 8 4k 67.05 31.48 2.1x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 8e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 8 8k 98.93 49.57 2.0x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 8e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 8 16k 162.63 86.44 1.9x / out 2e-06 dq 7e-07 dkv 1e-07
glm5.3-flash (as in #2786: 64-wide zero tail) 8 32k 290.69 165.2 1.8x / out 2e-06 dq 7e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 8 64k 552.94 337.66 1.6x / out 3e-06 dq 7e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 8 128k 1088.68 699.69 1.6x / out 3e-06 dq 8e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 4 4k 67.82 31.36 2.2x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 7e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 4 8k 100.03 49.97 2.0x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 7e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 4 16k 165.98 86.86 1.9x / out 2e-06 dq 7e-07 dkv 1e-07
glm5.3-flash (as in #2786: 64-wide zero tail) 4 32k 297.86 166.21 1.8x / out 2e-06 dq 7e-07 dkv -1e-07
glm5.3-flash (as in #2786: 64-wide zero tail) 4 64k 565.32 340.34 1.7x / out 2e-06 dq 7e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 4 128k 1112.54 706.56 1.6x / out 3e-06 dq 6e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 1 4k 113.98 33.03 3.5x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 7e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 1 8k 213.54 55.49 3.8x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 5e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 1 16k 421.18 100.53 4.2x / out 2e-06 dq 7e-07 dkv -1e-07
glm5.3-flash (as in #2786: 64-wide zero tail) 1 32k 863.49 195.58 4.4x / out 2e-06 dq 7e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 1 64k 1780.1 410.31 4.3x / out 2e-06 dq 7e-07 dkv 0e+00
glm5.3-flash (as in #2786: 64-wide zero tail) 1 128k 3748.82 865.48 4.3x / out 3e-06 dq 7e-07 dkv 0e+00
glm5.3-flash (no zero tail) 8 4k 65.29 25.35 2.6x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 7e-07 dkv 0e+00
glm5.3-flash (no zero tail) 8 8k 94.01 35.57 2.6x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 8e-07 dkv 0e+00
glm5.3-flash (no zero tail) 8 16k 152.29 56.17 2.7x / out 2e-06 dq 7e-07 dkv 1e-07
glm5.3-flash (no zero tail) 8 32k 268.89 97.53 2.8x / out 2e-06 dq 7e-07 dkv -1e-07
glm5.3-flash (no zero tail) 8 64k 504.56 191.65 2.6x / out 2e-06 dq 6e-07 dkv 0e+00
glm5.3-flash (no zero tail) 8 128k 982.5 406.99 2.4x / out 3e-06 dq 7e-07 dkv 0e+00
glm5.3-flash (no zero tail) 4 4k 65.12 25.59 2.5x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 7e-07 dkv 0e+00
glm5.3-flash (no zero tail) 4 8k 94.33 35.97 2.6x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 7e-07 dkv 0e+00
glm5.3-flash (no zero tail) 4 16k 152.92 56.8 2.7x / out 2e-06 dq 7e-07 dkv 0e+00
glm5.3-flash (no zero tail) 4 32k 270.1 98.81 2.7x / out 2e-06 dq 7e-07 dkv 0e+00
glm5.3-flash (no zero tail) 4 64k 506.85 194.81 2.6x / out 2e-06 dq 6e-07 dkv 0e+00
glm5.3-flash (no zero tail) 4 128k 988.3 413.32 2.4x / out 3e-06 dq 7e-07 dkv 0e+00
glm5.3-flash (no zero tail) 1 4k 99.66 30.24 3.3x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 7e-07 dkv 0e+00
glm5.3-flash (no zero tail) 1 8k 179.49 49.09 3.7x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 6e-07 dkv 1e-07
glm5.3-flash (no zero tail) 1 16k 342.14 86.67 3.9x / out 2e-06 dq 6e-07 dkv 0e+00
glm5.3-flash (no zero tail) 1 32k 691.57 162.97 4.2x / out 2e-06 dq 7e-07 dkv 1e-07
glm5.3-flash (no zero tail) 1 64k 1430.83 332.75 4.3x / out 2e-06 dq 6e-07 dkv 1e-07
glm5.3-flash (no zero tail) 1 128k 3002.8 703.03 4.3x / out 3e-06 dq 7e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 8 4k 10.72 4.11 2.6x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 9e-07 dkv -1e-07
dsv4.1-flash ratio-1 layers 8 8k 19.7 7.4 2.7x tilelang 2e-06 flash_mla 2e-06 / out 3e-06 dq 1e-06 dkv 0e+00
dsv4.1-flash ratio-1 layers 8 16k 37.61 14.08 2.7x / out 3e-06 dq 9e-07 dkv 1e-07
dsv4.1-flash ratio-1 layers 8 32k 73.92 29.24 2.5x / out 3e-06 dq 1e-06 dkv 0e+00
dsv4.1-flash ratio-1 layers 8 64k 147.65 63.77 2.3x / out 3e-06 dq 9e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 8 128k 296.4 137.84 2.2x / out 3e-06 dq 9e-07 dkv 1e-07
dsv4.1-flash ratio-1 layers 2 4k 15.62 6.91 2.3x tilelang 2e-06 flash_mla 2e-06 / out 3e-06 dq 1e-06 dkv 0e+00
dsv4.1-flash ratio-1 layers 2 8k 29.66 13.07 2.3x tilelang 2e-06 flash_mla 2e-06 / out 3e-06 dq 1e-06 dkv 1e-07
dsv4.1-flash ratio-1 layers 2 16k 58.37 25.98 2.2x / out 3e-06 dq 9e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 2 32k 118.32 56.1 2.1x / out 3e-06 dq 9e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 2 64k 242.55 123.4 2.0x / out 3e-06 dq 8e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 2 128k 494.23 265.58 1.9x / out 3e-06 dq 1e-06 dkv 0e+00
dsv4.1-flash ratio-1 layers 1 4k 26.7 6.98 3.8x tilelang 2e-06 flash_mla 2e-06 / out 3e-06 dq 8e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 1 8k 53.32 13.31 4.0x tilelang 2e-06 flash_mla 2e-06 / out 3e-06 dq 8e-07 dkv 6e-08
dsv4.1-flash ratio-1 layers 1 16k 108.48 26.42 4.1x / out 3e-06 dq 9e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 1 32k 226.62 54.93 4.1x / out 3e-06 dq 8e-07 dkv 0e+00
dsv4.1-flash ratio-1 layers 1 64k 474.96 115.73 4.1x / out 3e-06 dq 9e-07 dkv -1e-07
dsv4.1-flash ratio-1 layers 1 128k 980.66 241.98 4.1x / out 3e-06 dq 9e-07 dkv 0e+00
dsv4.1-flash ratio-0 layers 8 4k 12.55 4.78 2.6x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 8e-07 dkv 0e+00
dsv4.1-flash ratio-0 layers 8 8k 21.4 8.06 2.7x tilelang 2e-06 flash_mla 2e-06 / out 3e-06 dq 9e-07 dkv 6e-08
dsv4.1-flash ratio-0 layers 8 16k 39.25 14.66 2.7x / out 3e-06 dq 1e-06 dkv 0e+00
dsv4.1-flash ratio-0 layers 8 32k 75.21 27.97 2.7x / out 3e-06 dq 9e-07 dkv 0e+00
dsv4.1-flash ratio-0 layers 8 64k 147.84 58.26 2.5x / out 3e-06 dq 9e-07 dkv 1e-07
dsv4.1-flash ratio-0 layers 8 128k 294.75 127.17 2.3x / out 3e-06 dq 9e-07 dkv 0e+00
glm5 8 4k 64.14 28.68 2.2x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 8e-07 dkv 0e+00
glm5 8 8k 94.95 46.06 2.1x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 7e-07 dkv 0e+00
glm5 8 16k 156.96 80.8 1.9x / out 2e-06 dq 8e-07 dkv 0e+00
glm5 8 32k 281.36 155.17 1.8x / out 2e-06 dq 7e-07 dkv 0e+00
glm5 8 64k 535.15 319.29 1.7x / out 3e-06 dq 7e-07 dkv 0e+00
glm5 8 128k 1055.24 663.23 1.6x / out 3e-06 dq 7e-07 dkv 1e-07
glm5 4 4k 64.71 28.95 2.2x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 7e-07 dkv 0e+00
glm5 4 8k 96.68 46.41 2.1x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 7e-07 dkv 0e+00
glm5 4 16k 159.95 81.62 2.0x / out 2e-06 dq 7e-07 dkv 0e+00
glm5 4 32k 287.84 156.51 1.8x / out 2e-06 dq 7e-07 dkv 0e+00
glm5 4 64k 547.48 322.21 1.7x / out 2e-06 dq 7e-07 dkv 0e+00
glm5 4 128k 1077.99 670.54 1.6x / out 3e-06 dq 6e-07 dkv 0e+00
glm5 1 4k 110.34 30.58 3.6x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 6e-07 dkv 0e+00
glm5 1 8k 207.58 51.85 4.0x tilelang 2e-06 flash_mla 2e-06 / out 2e-06 dq 7e-07 dkv 0e+00
glm5 1 16k 409.05 119.56 3.4x

Stack

PR Change
1 #3605 non-attention kernels → miles/kernels
2 #3606 attention kernels → miles/kernels/attention
3 #3607 one DSA kernel pair (GLM-5, DeepSeek-V4, V4.1)
4 #3608 per-arch DSA sparse attention dispatch
5 #3609 shared head-sharded linear-attention layer (miles_plugins/models/linear_attn.py)
6 #3610 Qwen3.5 / Qwen3-Next GDN on it
7 #3632 Kimi-K3 KDA on it
8 #3634 short-conv and KDA backward backends

Rules for miles/kernels are 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.

claude[bot]
claude Bot reviewed Sep 22, 2026
View reviewed changes

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@Zhichenzzz Zhichenzzz changed the title perf(kernels): 3-9x faster DSA sparse attention perf(kernels): DSA sparse attention kernels Sep 22, 2026
@Zhichenzzz
Zhichenzzz added this pull request to stack #3611 September 22, 2026 23:16
This was referenced Sep 22, 2026
Merged
Merged
Merged
Merged
Merged
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/dsa-faster-kernels branch 2 times, most recently from 91c99cd to 1aa9c67 Compare September 23, 2026 00:51
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/dsa-faster-kernels branch from 1aa9c67 to 47d0bf1 Compare September 23, 2026 01:02
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/dsa-faster-kernels branch from 47d0bf1 to c1cfca0 Compare September 23, 2026 06:19
This was referenced Sep 23, 2026
Merged
Merged
Open
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/dsa-faster-kernels branch from c1cfca0 to 671d37b Compare September 23, 2026 22:08
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/dsa-faster-kernels branch 2 times, most recently from 7a89cd7 to 9e38836 Compare September 24, 2026 01:05
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/dsa-faster-kernels branch from 9e38836 to 176c8e5 Compare September 24, 2026 04:58
This was referenced Sep 30, 2026
Merged
Merged
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/dsa-faster-kernels branch from 5adc002 to ae5dd36 Compare October 2, 2026 22:31
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/dsa-faster-kernels branch from ae5dd36 to 3a5ab7f Compare October 3, 2026 07:43
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/dsa-faster-kernels branch 2 times, most recently from 358ad9f to f388879 Compare October 7, 2026 23:45
@Zhichenzzz Zhichenzzz changed the title perf(kernels): DSA sparse attention kernels [4/8] perf(kernels): 3-9x faster DSA sparse attention Oct 7, 2026
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/dsa-faster-kernels branch from f388879 to c7192e9 Compare October 8, 2026 00:03
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/dsa-faster-kernels branch from c7192e9 to 9108623 Compare October 8, 2026 00:21
@Zhichenzzz Zhichenzzz changed the title [4/8] perf(kernels): 3-9x faster DSA sparse attention [4/8] perf(kernels): DSA sparse attention kernel optimization Oct 8, 2026
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/dsa-faster-kernels branch from 9108623 to 1370eef Compare October 8, 2026 02:17
@Zhichenzzz Zhichenzzz changed the title [4/8] perf(kernels): DSA sparse attention kernel optimization [4/8] perf(kernels): 3-10x faster DSA sparse attention, tuned per arch Oct 8, 2026
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/dsa-faster-kernels branch from 1370eef to 2dbb126 Compare October 8, 2026 03:25
@Zhichenzzz Zhichenzzz changed the title [4/8] perf(kernels): 3-10x faster DSA sparse attention, tuned per arch [4/8] perf(kernels): per-arch DSA sparse attention dispatch with a FlashMLA forward Oct 8, 2026
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/dsa-faster-kernels branch from 2dbb126 to d31c21d Compare October 8, 2026 07:51
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/dsa-faster-kernels branch 2 times, most recently from 6e2f33e to e48e5e2 Compare October 8, 2026 16:47
@Zhichenzzz
perf(kernels): 3-10x faster DSA sparse attention
d12726f 
One kernel pair for GLM-5 and DeepSeek-V4 / V4.1, tuned per arch and
TP-local head count.

- Forward: FlashMLA's sparse prefill kernel or the TileLang kernel, chosen
  per shape. FlashMLA pads query heads to 64; it wins every GLM-5 shape and
  every shape on Blackwell; the TileLang forward wins DeepSeek-V4.1 at 8-16
  heads on H200. FlashMLA's natural-log lse is converted to the log2 lse with
  the sink folded in that the backward takes.
- Padded -1 indices: the host passes kv_indices = indices.clamp(min=0);
  the forward gathers KV through it unconditionally and builds the -inf mask
  from the original indices, the backward does the same and atomic-adds dKV
  through kv_indices. No out-of-bounds access, no per-element clamp (which
  cost up to 2x at 8-16 heads on H200).
- The TileLang forward runs unpipelined. With the unconditional gather,
  TileLang's software pipeline turned it into async copies that read the
  wrong stage when the block count was stages + 1 (16 heads, topk 128 at one
  stage / 192 at two): rows came back as 1e30. main's DeepSeek-V4.1 forward
  has the same pattern and is wrong at 16 heads, topk 128 (rel err 5e-4).
- Backward takes batch and sequence lengths as dynamic shapes: a new
  packed length no longer compiles a new kernel (main's V4.1 copy spends
  6-8 s per new length on H200).
- SparseAttentionConfig: forward backend, forward threads, backward
  stages/split_store. Measured rows keyed by (arch major, RoPE
  tail, TP-local heads); B300 and GB300 pick identical winners, H200
  differs (e.g. GLM-5 8 heads backward: stages 1 / split 2 on H200,
  stages 2 / split 1 on Blackwell, 1.5x apart). Callers may pass a config.
- Indexer forward tiles: block_N 128, 2 stages, block_Q 256 / heads,
  2x faster on GB300.
- Backward threads 256 when a block holds >= 32 heads on CUDA; ROCm keeps
  32-head blocks above 64 heads and 1 stage.

fwd + bwd per layer vs main, 16k / 64k tokens:

                             H200            GB300
GLM-5 8 heads (TP=8)         3.40x / 3.28x   3.01x / 3.15x
GLM-5 16 heads (TP=4)        3.24x / 3.02x   3.01x / 3.15x
GLM-5 64 heads (TP=1)        9.41x / 10.06x  5.42x / 6.38x
V4.1 vs V4 kernel, 8 heads   1.12x / 1.18x   1.88x / 2.01x
V4.1 vs V4 kernel, 32 heads  1.32x / 1.51x   1.88x / 2.03x
V4.1 vs V4 kernel, 64 heads  3.27x / 2.96x   3.27x / 3.42x
V4.1 vs its copy, 8 heads    0.83x / 0.88x   1.60x / 1.67x
V4.1 vs its copy, 32 heads   1.03x / 1.03x   1.70x / 1.74x
V4.1 vs its copy, 64 heads   3.13x / 2.85x   3.57x / 3.71x

The H200 8-head V4.1 row is the copy's static-shape backward; at a fixed
length it is ~2 ms faster, at a new length it recompiles for 6-8 s.
Outputs within 1e-5 (TileLang) / 3e-3 (FlashMLA) of main, dkv within 1e-4
(fp32 atomic order). Tests cover every case on both forwards and on the
tuned default config, -1 padding, all-padded rows, sinks, odd lengths and
packed sequences (causal_ranges -> indexer -> top-k -> attention, 1-token
segments). A wider matrix (5 head counts x 6 topk x packed / causal / batch-2
x tail / shuffled padding x 8 configs, 3 repeats, fwd + bwd vs fp32) passes
on H200, B300 and GB300.
@Zhichenzzz
Zhichenzzz force-pushed the zhichen/dsa-faster-kernels branch from e48e5e2 to d12726f Compare October 8, 2026 19:16
@Zhichenzzz
Zhichenzzz merged commit 7779497 into main Oct 8, 2026
38 of 41 checks passed
@Zhichenzzz
Zhichenzzz deleted the zhichen/dsa-faster-kernels branch October 8, 2026 19:57
@zhouhelena1 zhouhelena1 mentioned this pull request Oct 9, 2026
Closed
4 of 7 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@claude claude[bot] claude[bot] left review comments
@yueming-yuan yueming-yuan yueming-yuan approved these changes
@fzyzcjy fzyzcjy Awaiting requested review from fzyzcjy fzyzcjy is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@Zhichenzzz @yueming-yuan