[TRITON][GLUON] Add DSv4 sparse MLA prefill Triton and unified with MLA Gluon kernel on GFX950 - #3833
Merged
Dewei-Wang-sh merged 12 commits intoJul 5, 2026
Conversation
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
Contributor
There was a problem hiding this comment.
Pull request overview
Adds DeepSeek V4 sparse MLA prefill support to AITer, including a Triton implementation (with autotune) and an optimized Gluon (CDNA4/gfx950) kernel, plus wrappers/tests/benchmarks to validate and measure performance.
Changes:
- Introduces Triton sparse-attention prefill kernel + ragged-index preparation kernels and autotune configs.
- Adds a gfx950 Gluon sparse-attention prefill kernel and a wrapper that selects Gluon when available.
- Adds an accuracy/perf test and a benchmark driver, and documents the new Gluon kernel.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
op_tests/test_sparse_attention_dsv4.py |
Adds a torch reference + correctness/perf test harness for sparse prefill. |
op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4.py |
Adds a Triton vs Gluon benchmark runner and reporting. |
aiter/ops/triton/gluon/sparse_attention_dsv4.py |
Implements the gfx950 Gluon sparse MLA prefill kernel + host launcher. |
aiter/ops/triton/gluon/README.md |
Documents the new Gluon sparse-attention kernel and how to run tests/benchmarks. |
aiter/ops/triton/attention/sparse_attention_dsv4.py |
Adds the public wrapper API, dense→CSR packing helpers, and Gluon gating logic. |
aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4.py |
Adds Triton kernels for index packing and sparse prefill (autotuned). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
leonling-ll
force-pushed
the
liyang/dsv4_sparse_attn_prefill
branch
from
June 22, 2026 05:31
d605eb9 to
3e79696
Compare
Dewei-Wang-sh
requested changes
Jun 23, 2026
Dewei-Wang-sh
left a comment
Dewei-Wang-sh
left a comment
Contributor
There was a problem hiding this comment.
reuse the mla_decode_gluon to relieve maintain burden.
leonling-ll
force-pushed
the
liyang/dsv4_sparse_attn_prefill
branch
3 times, most recently
from
June 24, 2026 05:58
356fd58 to
7813e7a
Compare
Prefill-only slice of the DSV4 compressed sparse attention work: bf16 KV flat-buffer prefill kernel (Triton + persistent Gluon), the ragged-index builders + host launcher, and the prefill test/benchmark. Decode (fp8_ds_mla paged cache) is intentionally excluded from this branch. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Resolve the code style check failure Address Comments
Fix format Unify dsv4 sparse attn and mla_decode
leonling-ll
force-pushed
the
liyang/dsv4_sparse_attn_prefill
branch
from
June 24, 2026 07:48
7813e7a to
e31de73
Compare
leonling-ll
changed the title
leonling-ll
changed the title
Contributor
|
overall, we should keep it additive and add comments about env/change, keep it concise and clear. |
leonling-ll
force-pushed
the
liyang/dsv4_sparse_attn_prefill
branch
from
June 30, 2026 16:21
a4a6e7a to
a89a80b
Compare
leonling-ll
commented
Jul 1, 2026
leonling-ll
commented
Jul 1, 2026
leonling-ll
commented
Jul 1, 2026
leonling-ll
commented
Jul 1, 2026
leonling-ll
force-pushed
the
liyang/dsv4_sparse_attn_prefill
branch
3 times, most recently
from
July 2, 2026 08:53
5e7ba56 to
708895d
Compare
Address review comments Minimize changes The dev branch selected DeepSeek V4 sparse prefill via a `mode="decode" | "dsv4-prefill"` string that drove a separate wrapper branch duplicating the whole bh64 decode setup. Replace it with a `has_pe=True` flag (HAS_PE=False => prefill) and collapse the duplicated branch back onto main's single-path decode wrapper, adding prefill support as minimal inline guards: - q_pe may be None (fixed head_dim_kpe=64), small `if not has_pe` block aliases the unused q_pe/k_pe placeholders and forces kv_pe_offset=0 / use_2d_view=False - bh64 regime: has_pe keeps the split auto-pick + asserts; else NUM_KV_SPLITS=1 with a batch % 64 guard - unified bh64 grid (cdiv(batch, NUM_XCDS) * NUM_KV_SPLITS) covers both paths Callers (sparse_attention_dsv4, bench_sparse_attention_dsv4) pass has_pe=False. Verified on gfx950: test_mla decode (bh64/bh16bn64/bh16bn128, +lse/varlen) and test_sparse_attention_dsv4 prefill (+sink) all pass; benchmark shows prefill is performance-neutral vs the pre-refactor version (within +-1%).
leonling-ll
force-pushed
the
liyang/dsv4_sparse_attn_prefill
branch
from
July 2, 2026 09:00
708895d to
0f4d2e2
Compare
leonling-ll
commented
Jul 2, 2026
leonling-ll
commented
Jul 2, 2026
leonling-ll
commented
Jul 2, 2026
leonling-ll
force-pushed
the
liyang/dsv4_sparse_attn_prefill
branch
2 times, most recently
from
July 3, 2026 04:32
cd2912a to
bdbd317
Compare
leonling-ll
force-pushed
the
liyang/dsv4_sparse_attn_prefill
branch
from
July 3, 2026 04:55
bdbd317 to
87e062c
Compare
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
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.
This PR aims to add sparse attention prefill kernel (triton + gluon gfx950).
The triton kernels are adapted from vLLM mla_sparse_dsv4 with enabling triton autotune.
It reuses
mla_decode_gluonand unifies asmla_gluonkernel for both mla decode and sparse prefill.The prefill performance status:
This table can be reproduce with:
This PR is split from #3456 as the first part.