dsv4.1-amd: gfx950 sparse decode attention and sorted top-k by kevin-mii · Pull Request #41020 · sgl-project/sglang

kevin-mii · 2026-09-24T03:03:56Z

This PR was "stack 3/4" of the DeepSeek-V4.1 AMD series (#41018 to #41021). It now follows the CUDA dsv4.1 layout (#39646, #39652, #39653, #39656, #39664, then #38798): kernel PRs by domain, each on main with no callers, then one integration PR. This is the third kernel PR; it does not depend on #41018 or #41019.

Summary

ops/attention/sparse_decode_reduce_hip.py: gfx_sparse_split_reduce combines the split-KV partials of aiter's gfx950 pa_decode_sparse(..., skip_reduce=True), bitwise equal to aiter's own reduce, with the attention sink and an optional inverse RoPE of each head's tail.
ops/attention/dsv4/compact_attention_hip.py: sparse decode attention that reads the compact V4.1 pages (528 B fp8 + ue8m0, 288 B e2m1 + e4m3) directly through aiter's gfx950 _qkpv_fp8, with main and extra slot ranges and 16-head groups.
ops/attention/dsv4/swapab_gluon_hip.py: the 16-head swap-AB Gluon kernel for HIP, argument-compatible with the SM100 kernel (CDNA4 MFMA, bf16 hi + lo compensated PV). It reuses the SM100 module's _load_v4, which is plain Gluon; swapab_attention gains inv_rope=None.
ops/attention/dsv4/attn_glue_hip.py: single-launch Triton replacements for the torch glue around the V4.1 sparse decode metadata, each bitwise equal to the torch expression in its docstring.
ROCm AOT top-k (deepseek_v4_topk.cu, sgl_kernel/top_k.py): sort_output sorts each row's picks ascending with a block-wide bitonic sort, -1 padding last; the selected set is unchanged; the sort is ROCm-only (under USE_ROCM). Strides below the wave width use lane shuffles and the rest LDS, so it is also right on 32-lane waves (gfx1250). benchmark/bench_dsv4_topk_transform.py times it: on MI350X the sort adds ~3.4 us per call, where a separate torch.sort costs ~17-20 us.
q_rope_store.py on HIP rounds like the batched HIP RoPE kernel.
ops/attention/dsv4/candidate_blocks_hip.py, beside main's candidate_blocks.py: the V4.1 two-level indexer's level-one block scores and top blocks, the top-k within the chosen blocks and the paged top-k sorted in the AOT epilogue, graph-safe.
Model call sites are in the integration PR of this stack.

Changes to existing kernels

decode_attention_sm100.py: swapab_attention takes an optional partial_gluon (swapab_gluon_hip.swapab_attention binds the gfx950 kernel), with no HIP branch in the SM100 file; _combine compiles the inverse RoPE out when inv_rope=None, so SM100 output is unchanged. decode_attention_sm100_gluon.py is untouched.
q_rope_store.py: the new branch is a HIP-only constexpr; CUDA arithmetic is unchanged.
The top-k epilogue, op schema and declaration are ROCm-only (setup_rocm.py); with sort_output=False the kernel is the same, plus a 4 KB static LDS array. The wrapper always passes sort_output, so sgl-kernel must be rebuilt from this PR.
Rollout: this PR is kernel-only, and no caller passes sort_output=True yet (the one existing caller, ops/attention/dsv4/topk.py, still uses the 6-argument form, which the default keeps working). Callers will only pass sort_output=True after sgl-kernel is released and the sglang-kernel pin is bumped to 0.4.8 or later, following How to update kernels in SGLang.

Verification

New tests: test_swapab_attention.py (SM100 and gfx950: fp64 reference over padded and poisoned pages, -1 / out-of-range slots, sinks, graph replay, and the combine's inverse RoPE bitwise to fused_rope_inplace); test_compact_attention_hip.py (both packed layouts, a second cache, -1 slots and a key-less row against a torch softmax with sink); test_topk.py::test_deepseek_v4_topk_transform_sorted_output; test_dsv41_decode_glue_hip.py (candidate blocks against main's select_candidate_block_ids with garbage past each row's reach, the sorted mapping, and the attn_glue_hip builders bitwise to their torch forms). test_dsv4_q_rope_store.py is now also registered for AMD.
MI350X (ROCm 10.0.0, torch 2.11, triton 3.8, aiter at the Dockerfile pin): these tests pass on this branch. The top-k epilogue built from this source matches the installed op over 120 shapes (sorted rows bitwise, unsorted rows the same set) and compiles for gfx1250.
With the integration PR's callers, DeepSeek-V4.1-Flash at TP4 / EP4 with DSpark serves and scores 0.900-0.905 on GSM8K (200 questions, two runs).
Pending on an SM100 runner: test_swapab_attention.py, test_dsv4_q_rope_store.py.

Stack

Base: main. Independent of dsv4.1-amd: gfx950 MXFP8 matmul kernels and fp8-grid producers #41018 (gfx950 MXFP8 matmul kernels) and dsv4.1-amd: KV cache layouts, FP4 indexer, compressor and router kernels #41019 (KV cache layouts, FP4 indexer, compressor and router kernels).

🤖 Generated with Claude Code

CI States

Latest PR Test (Base): ❌ Run #36477152106
Latest PR Test (Extra): ✅ Run #36477151594
Latest PR Test (AMD ROCm 10): ⏳ Run #36477151949

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

candidate_blocks_hip.py sits beside main's candidate_blocks.py: the level-one block scores and top blocks, the top-k within the chosen blocks and the paged top-k sorted in the AOT epilogue, all graph-safe. test_dsv41_decode_glue_hip.py pins them and this PR's attn_glue_hip builders against the torch references. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

HaiShaw

Just a few comments to address.

HaiShaw · 2026-09-27T05:31:45Z

 import triton.language as tl

+from sglang.srt.utils import is_hip
+


Consider to move HIP out of _sm100.py, or rename file is possible.

Done in 7055e0c: no HIP code in _sm100.py now; the gfx950 entry is swapab_gluon_hip.swapab_attention.

HaiShaw · 2026-09-27T05:36:53Z

+
+
+@triton.jit
+def _aiter_sparse_decode_reduce_kernel(


In-tree kernel no need to start with aiter when no origin relevant (in that case use _gfx_sparse_decode_recue_kernel), othereise we do import from it.

Done: _gfx_sparse_decode_reduce_kernel; file renamed to sparse_decode_reduce_hip.py to match the other _hip.py kernels.

HaiShaw · 2026-09-27T05:37:16Z

+    tl.store(out_ptr + t * out_stride_t + h * out_stride_h + d, out)
+
+
+def aiter_sparse_split_reduce(


Same for method name

Done: gfx_sparse_split_reduce.

… in-tree split reduce - decode_attention_sm100.py no longer branches on is_hip: swapab_attention takes an optional partial_gluon, and swapab_gluon_hip.py exports the gfx950 swapab_attention bound to its kernel. - aiter_sparse_decode_reduce.py -> sparse_decode_reduce_hip.py (_gfx_sparse_decode_reduce_kernel, gfx_sparse_split_reduce): the kernel is in-tree, only its input layout comes from aiter. - deepseek_v4_topk.cu: the wave-size, DPP/permlane and bitonic helpers and the sort_output epilogue are under USE_ROCM; a CUDA build rejects sort_output. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

kevin-mii · 2026-09-27T05:55:15Z

Note on the AMD CI failure in run 36227768892 / stage-c-test-4-gpu-amd (rocm10, linux-mi300-4gpu-sglang, 0) (at 4421656): it is pre-existing on main, not from this PR.

Failing test: test/registered/lora/test_lora_gpt_oss_20b_logprob_diff.py, KL(sglang, trainer) 1.13-1.51 vs threshold 0.005 on every retry.
The same test fails the same way in main's scheduled PR Test (AMD) runs on the same job: 36241615319 (rocm10, KL up to 23.9), 36205773263 (rocm10, KL up to 15.4), 36168295472 (rocm720, KL up to 9.5).
This PR only touches DeepSeek-V4 kernels (ops/attention/dsv4/*, the split-KV reduce) and the ROCm DeepSeek-V4 top-k; gpt-oss LoRA does not reach any of them.

Review comments are addressed in 7055e0c.

HaiShaw · 2026-09-27T07:06:28Z

Just in case - to double check
https://lmsysorg.mintlify.app/docs/developer_guide/contribution_guide#how-to-update-kernels-in-sglang

bench_dsv4_topk_transform.py times deepseek_v4_topk_transform_512 unsorted, with sort_output, and unsorted followed by torch.sort (MI350X, decode batch 1-128, 1k-32k keys: the epilogue sort adds ~3.4 us, a separate torch.sort ~17-20 us). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

kevin-mii · 2026-09-28T20:08:54Z

https://github.com/sgl-project/sglang/actions/runs/36334848845 all PR base passed except for main failure for JIT test

kevin-mii and others added 2 commits September 26, 2026 05:45

Merge main into dsv41-amd-3-hip-attn-backend

e5242cc

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

bingxche mentioned this pull request Sep 26, 2026

[CI Monitor] Daily Report - 2026-09-26 bingxche/sglang-ci-bot#207

Open

HaiShaw reviewed Sep 27, 2026

View reviewed changes

michaelzhang-ai mentioned this pull request Sep 27, 2026

[AMD] Add DeepSeek-V4.1-Flash MI35x nightly accuracy test #41476

Merged

1 task done

HaiShaw approved these changes Sep 28, 2026

View reviewed changes

Merge branch 'main' into dsv41-amd-3-hip-attn-backend

05564d1

hnyls2002 approved these changes Sep 28, 2026

View reviewed changes

jiaryang mentioned this pull request Sep 29, 2026

[ROCm] Select DSA indexer top-k wave size by arch (wave32 on gfx1250) #41687

Merged

4 of 5 tasks

This was referenced Oct 7, 2026

[CI Monitor] Daily Report - 2026-10-07 bingxche/sglang-ci-bot#218

Open

[CI Monitor] Daily Report - 2026-10-08 bingxche/sglang-ci-bot#219

Open

michaelzhang-ai mentioned this pull request Oct 9, 2026

[AMD] DSV4 compact_attention_hip imports an aiter module that was removed #43257

Closed

michaelzhang-ai mentioned this pull request Oct 12, 2026

[Bug][AMD] DSV4 compact_attention_hip imports aiter's pa_decode_sparse, which aiter removed on 2026-09-18 #43765

Closed

1am9trash mentioned this pull request Oct 12, 2026

[AMD] Fix test_compact_attention_hip.py CI #43836

Merged

5 tasks

Navigation Menu

Repository navigation

dsv4.1-amd: gfx950 sparse decode attention and sorted top-k - #41020

dsv4.1-amd: gfx950 sparse decode attention and sorted top-k#41020
hnyls2002 merged 6 commits into
sgl-project:mainfrom
kevin-mii:dsv41-amd-3-hip-attn-backend

kevin-mii commented Sep 24, 2026 •

edited by github-actions Bot

Loading

Uh oh!

HaiShaw left a comment

Uh oh!

HaiShaw Sep 27, 2026

Uh oh!

kevin-mii Sep 27, 2026

Uh oh!

HaiShaw Sep 27, 2026

Uh oh!

kevin-mii Sep 27, 2026

Uh oh!

HaiShaw Sep 27, 2026

Uh oh!

kevin-mii Sep 27, 2026

Uh oh!

Uh oh!

kevin-mii commented Sep 27, 2026

Uh oh!

HaiShaw commented Sep 27, 2026

Uh oh!

kevin-mii commented Sep 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

		import triton.language as tl

		from sglang.srt.utils import is_hip

		tl.store(out_ptr + t * out_stride_t + h * out_stride_h + d, out)


		def aiter_sparse_split_reduce(

Conversation

kevin-mii commented Sep 24, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes to existing kernels

Verification

Stack

CI States

Uh oh!

HaiShaw left a comment

Choose a reason for hiding this comment

Uh oh!

HaiShaw Sep 27, 2026

Choose a reason for hiding this comment

Uh oh!

kevin-mii Sep 27, 2026

Choose a reason for hiding this comment

Uh oh!

HaiShaw Sep 27, 2026

Choose a reason for hiding this comment

Uh oh!

kevin-mii Sep 27, 2026

Choose a reason for hiding this comment

Uh oh!

HaiShaw Sep 27, 2026

Choose a reason for hiding this comment

Uh oh!

kevin-mii Sep 27, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

kevin-mii commented Sep 27, 2026

Uh oh!

HaiShaw commented Sep 27, 2026

Uh oh!

kevin-mii commented Sep 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

kevin-mii commented Sep 24, 2026 •

edited by github-actions Bot

Loading