Add AITER GDN prefill and decode backends

[IzacharyI]

Integrate ROCm HIP and FlyDSL kernels through the canonical VK dispatcher while preserving Triton fallbacks for graph capture, padding, and unsupported configurations.

Related AITER PR: ROCm/aiter#4540

Motivation

The existing GDN backend relies on Triton kernels on ROCm even when optimized AITER HIP and FlyDSL implementations are available.

This PR integrates those kernels without changing the canonical GDN state layout or scheduler interfaces. AITER is selected only for validated workloads, while Triton remains the correctness fallback for unsupported shapes, dependencies, padding, graph configurations, and execution modes.

Modifications

Modifications

• Add aiter as an independent linear-attention decode and prefill backend.
• Implement AiterGDNKernel through the existing LinearAttnKernelBase and GDNKernelDispatcher interfaces.
• Enable AITER for both decode and prefill with:
  • --linear-attn-backend aiter
• Retain optional per-mode overrides for A/B testing:
  • --linear-attn-decode-backend aiter
  • --linear-attn-prefill-backend aiter
• Integrate the following decode kernels:
  • PR [Fix] Not skip NVML Check on AMD Platform sgl-project/sglang#3135 FlyDSL VK flydsl_gdr_decode.
  • ROCm HIP fused sigmoid-gating delta-rule update.
  • Fused causal-convolution update and QKV split when available.
• Route decode kernels as follows:
  • Qwen3.5-397B TP8 (HK, HV)=(2, 8), batch >=24: HIP.
  • All other valid head layouts and smaller 397B batches: FlyDSL.
  • Missing or incompatible AITER kernels: Triton.
• Preserve the canonical [slot, HV, V, K] recurrent-state layout without introducing state transposes or legacy bitmap management.
• Support CUDA Graph capture and replay with graph-safe state indexing and padded-batch handling.
• Disable the HIP host-side sorted-index cache internally to ensure graph replay safety.
• Integrate indexed AITER VK prefill:
  • High-level in-place VK path for supported single-sequence prefill.
  • Fused Q/K L2Norm before the high-level VK path.
  • Low-level K1–K6 path when intermediate-state tracking is required.
• Preserve Triton fallbacks for:
  • Missing or incompatible AITER APIs.
  • Unsupported architectures, head layouts, dtypes, and state layouts.
  • Padded or multi-sequence prefill batches.
  • Invalid state slots.
  • Target verification and unsupported graph configurations.
• Keep ReplaySSM on the Triton linear-attention backend.

Accuracy Tests

5-shot GSM8K evaluation using 100 examples, temperature 0, top-p 1, and maximum output length 4096:

Model	Parallelism	Triton baseline	AITER backend	Result
Qwen3.5-27B-FP8	TP1	0.89	0.88	One-question variation
Qwen3.5-397B-A17B-PTPC	TP8	0.75	0.75	No regression

Focused correctness tests:

• 39 tests passed across AITER dispatch, model registration, and non-contiguous-stride coverage.
• Existing Triton suite: 8 passed, 1 skipped, 24 subtests passed.
• Real ROCm output and recurrent-state parity validated against Triton.
• CUDA Graph replay validated with changed state-slot indices.
• Eager and graph padding, prefix tracking, fallback routing, and intermediate-state tracking validated.

Speed Tests and Profiling

Environment:

• AMD MI308X
• Qwen3.5-397B-A17B-PTPC
• TP8
• FP8 KV cache
• AITER commit 702aacd62
• FlyDSL 0.2.4
• ISL/OSL: 8000/500
• Radix cache disabled
• Baseline commit: b008d42d8521d3ec5a4052e13320d623356e187f

Kernel profiling

Model shape	Operation	Workload	Speedup vs. Triton
27B TP1	Fused conv-split + FlyDSL decode	B=2/4/8/16	1.18× / 1.18× / 1.19× / 1.13×
397B TP8	Fused conv-split + FlyDSL decode	B=64	1.21×
27B TP1	AITER VK prefill	T=8192	1.26×
397B TP8	AITER VK prefill	T=4096/8192	1.22× / 1.24×

End-to-end concurrency benchmark

Concurrency	Baseline QPS	AITER QPS	QPS change	Baseline step latency	AITER step latency
4	0.3868	0.3960	+2.37%	19.589 ms	19.248 ms
8	0.5994	0.6229	+3.92%	25.583 ms	24.721 ms
16	0.7829	0.8854	+13.09%	39.576 ms	35.282 ms
20	0.8907	0.9075	+1.88%	43.828 ms	43.182 ms
24	0.9606	1.0309	+7.32%	48.894 ms	45.633 ms

• QPS improved at all five tested concurrency levels.
• Arithmetic mean of the five observed QPS improvements: +5.71%.
• Error request rate: 0% for both configurations.
• Generated dataset MD5 values matched between baseline and AITER runs.
• Results are single-run observations and are not presented as confidence bounds.

The 397B TP8 prefill path remains on Triton by default because controlled 8K end-to-end measurements showed a 0.81% mean TTFT regression. The AITER prefill path remains available through an explicit threshold override.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

[apinge]

Decode 8K in C1

Before

After

Prefill 8K in C1

Before

After

[sammysun0711]

@IzacharyI for APC function mamba cache support with linear attention.
Please update following patch:

diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py
index 9439b6da4..dc5443bfa 100644
--- a/python/sglang/srt/arg_groups/overrides.py
+++ b/python/sglang/srt/arg_groups/overrides.py
@@ -999,7 +999,9 @@ def supports_mamba_cache_extra_buffer(view: Any, model_arch: str) -> bool:
     """Whether ``model_arch`` supports the extra_buffer strategy on the
     configured linear-attention backend (pure read)."""
     if model_arch in _MAMBA_EXTRA_BUFFER_ARCHS:
-        return view.linear_attn_backend == "triton"
+        return view.linear_attn_backend == "triton" or (
+            is_hip() and view.linear_attn_backend == "aiter"
+        )
     return False

[Copilot]

Pull request overview

This PR adds an AITER-based GDN linear-attention backend on ROCm, integrating AITER HIP + FlyDSL kernels into the existing GDN dispatcher while keeping Triton as the fallback for unsupported shapes/modes and for ReplaySSM.

Changes:

• Add aiter as a selectable LinearAttnKernelBackend and CLI choice, and wire it into GDNKernelDispatcher for decode/prefill.
• Introduce AiterGDNKernel with runtime guards and fallback routing (HIP vs FlyDSL vs Triton), plus graph-padding aware decode handling via active_batch_size.
• Improve Mamba state tracking/scatter paths with direct-copy helpers and add/expand AMD-focused tests for AITER GDN and state-copy correctness.

Reviewed changes

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

Show a summary per file

File	Description
test/registered/unit/test_model_overrides.py	Extends override tests to cover aiter + ROCm gating for mamba extra buffer support.
test/registered/unit/layers/test_mamba_state_scatter_triton.py	Adds unit tests for new extend-state copy flag derivation and copy behavior (incl. overlap).
test/registered/attention/test_gdn_aiter_backend.py	New AMD CI suite validating AITER backend registration, routing, graph padding behavior, and parity vs Triton.
python/sglang/srt/server_args.py	Adds aiter to linear-attention backend CLI choices.
python/sglang/srt/layers/attention/mamba/mamba2_metadata.py	Adds trusted/disjoint flags to forward metadata and threads them through prepare helpers.
python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py	Adds trusted/disjoint derivation plus a Triton row-copy kernel and extend-row copy helper.
python/sglang/srt/layers/attention/linear/utils.py	Registers LinearAttnKernelBackend.AITER and helper predicate.
python/sglang/srt/layers/attention/linear/kernels/gdn_aiter.py	New AITER GDN kernel implementation with HIP/FlyDSL decode and VK/low-level prefill paths plus Triton fallback.
python/sglang/srt/layers/attention/linear/gdn_backend.py	Integrates AITER into kernel dispatch and adds decode conv-split + active-batch handling hooks.
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py	Uses new trusted/disjoint flags and direct-copy helper for extend-state tracking; minor formatting changes.
python/sglang/srt/layers/attention/fla/l2norm.py	Adds fused Q/K L2Norm Triton kernels and fused_l2norm_qk helper.
python/sglang/srt/arg_groups/overrides.py	Allows mamba extra buffer on ROCm when linear_attn_backend == "aiter".

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[sammysun0711]

@IzacharyI do you think it has out of bound read risk?

[IzacharyI]

I checked the actual AITER Gluon conv-split kernel. It explicitly returns when conv_state_indices is PAD_SLOT_ID (-1), before calculating or accessing the state address, so padded rows cannot modify conv_state[-1].

The proposed Python guard would not protect CUDA Graph replay because it is evaluated only during graph capture, not again when the active batch changes during replay.

I added a regression test that captures with valid indices, replays after changing one index to -1, and verifies that valid Q/K/V match the reference and the last cache row is unchanged. So, may the fused fast path is kept enabled.