[karverma-amd]

[Qwen3-Next] Consume the AITER fused GDN prefill kernel

Motivation

The Qwen3-Next GDN (gated delta-net) prefill runs as a four-kernel chain
today (causal_conv1d_split_qkv → fused gating → the fp32 chunked delta-rule
scan → gated RMSNorm), so the large intermediate tensors round-trip through HBM
between each launch. AITER now ships a fused single-kernel prefill
(fused_gdn_prefill_qkvz, ROCm/aiter#5606) that folds that chain into a tight set
of Gluon launches and emits the per-head group-128 FP8 activations a block-FP8
out_proj consumes directly.

This PR wires sglang to use it. It is the prefill sibling of #39554 (the decode
consumer of ROCm/aiter#5520) and rides the standard AMD switch, so one flag enables
both phases:

--enable-aiter          # env: SGLANG_USE_AITER

On by default wherever AITER is enabled on a gfx950 host — there is no bespoke
per-kernel flag. Everything is probed, nothing assumed: unsupported platform,
older Triton, or an uncovered call runs the existing unfused chain unchanged.

What it does

• Adapter kernels/ops/attention/gdn_fused_prefill_aiter.py — fail-closed
  available() / covered() / run(): platform + SGLANG_USE_AITER, a Triton ≥ 3.8
  Gluon probe, and a per-call shape/dtype contract. covered() delegates the tensor
  contract to AITER's own predicate so the two can't drift, then additionally checks
  the parameters that predicate does not see — the SiLU output gate and the
  A_log/dt_bias/norm_weight element counts/dtypes — so a mismatch falls back
  instead of raising deep inside the kernel. Two small pure helpers
  (out_proj_accepts_group128_fp8, prefill_rows_match_real_tokens) hold the
  out_proj-compat and DP-padding decisions and are unit-tested directly.
• GDNAttnBackend.try_fused_gdn_prefill — builds the ragged-batch args
  (cu_seqlens from query_start_loc, has_initial_state from extend_prefix_lens,
  cache_indices from the mamba pool) and the conv/state pools, calls
  covered() → run(), and stashes the post-norm output + group-128 FP8 activations.
  Declines MIS / target-verify / tracked-conv batches, and declines DP-padded
  batches — the fused kernel has no padding trim, so it must not run when the
  projected rows exceed the real varlen tokens (global_num_token_non_padded_cpu
  below the row count, or cu_seqlens[-1] != rows).
• RadixLinearAttention.try_fused_gdn_prefill — unwraps HybridLinearAttnBackend
  to reach the GDN backend; guards extend / not-verify / not-piecewise.
• qwen3_next.py — the fused path is AMD/AITER-only and stays off the NVIDIA
  common path (_is_hip fence + lazy import). Static readiness gate (available() +
  num_v_heads == 2·num_k_heads + head_k == head_v == 128 + conv == 4), publishes
  the bf16 norm weight after load, attempts the fused path first in forward. It feeds
  the kernel's group-128 FP8 activations straight into out_proj only when that
  GEMM is a group-128 block-FP8 (block K == 128, not MXFP8); every other method takes
  the bf16 normalized output and re-quantizes as usual. The FP8 stash is cleared right
  after it is consumed.

Files changed

File	+	What
python/sglang/kernels/ops/attention/gdn_fused_prefill_aiter.py	232 (new)	Fail-closed adapter (available/covered/run) + out_proj-compat and DP-padding helpers
python/sglang/srt/models/qwen3_next.py	90	readiness gate, NVIDIA fence, norm-weight publish, attempt-first + gated block-FP8 out_proj feed
python/sglang/srt/layers/attention/linear/gdn_backend.py	76	try_fused_gdn_prefill dispatch + DP-padding decline + stash
python/sglang/srt/layers/radix_linear_attention.py	46	backend unwrap + guards
test/registered/attention/test_gdn_fused_prefill_guards.py	156 (new)	CPU unit tests for the DP-padding, out_proj-compat, and covered() decline branches

Dependency

Requires the AITER fused prefill kernel ROCm/aiter#5606. Fail-closed, so this
merges safely before that lands — the fused path simply activates once the kernel
is importable.

Accuracy — run command

Same model/args with the AITER fused kernel present vs absent; GSM8K must stay at
parity. Needs a gfx950 (MI350-series) host with Triton ≥ 3.8.

# Qwen3-Next-80B FP8, TP4. Fused prefill is on whenever --enable-aiter is set and
# the AITER kernel is importable; it fails closed to the unfused chain otherwise.
python3 -m sglang.launch_server --model-path <Qwen3-Next-80B-A3B-Instruct-FP8> \
  --tp-size 4 --quantization fp8 --kv-cache-dtype fp8_e4m3 \
  --attention-backend aiter --enable-aiter --context-length 8192 &

python3 -m sglang.test.few_shot_gsm8k --num-questions 1319 --parallel 32 --port 30000

Result (MI355X): GSM8K parity — fused 0.953 vs stock 0.950.

Perf

End-to-end A/B, fused kernel present vs absent, same Triton-3.8 image (isolates the
kernel), TP4, Qwen3-Next-80B FP8, MI355X.

Prefill-focused (ISL 6144 / OSL 8):

metric	stock	fused	Δ
Median TTFT (ms)	647.9	471.6	−27%
Input throughput (tok/s)	92.8k	125.4k	+35%

Balanced 1k/1k, concurrency sweep (Median TTFT, ms):

conc	stock	fused	Δ
1	69.8	61.1	−12%
4	137.2	115.4	−16%
8	141.4	138.9	−2%
16	205.9	193.3	−6%
32	356.3	339.5	−5%

TPOT is identical on both arms (decode path untouched); the win is prefill-side,
largest where prefill is a bigger fraction of the request (low concurrency / long
prompts) and diluting as the decode-heavy steady state dominates.

---

CI States

Latest PR Test (Base): 🚫 Run #41252100316
Latest PR Test (Extra): ❌ Run #41252100257
Latest PR Test (AMD ROCm 10): ❌ Run #41252100266

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

[yichiche]

Please add a guard with if_hip or if_rocm to make sure your change doesn't affect the NVIDIA common path. Use if_hip if this kernel can be used by all AMD GPUs; use use_aiter if this kernel imports from the AITER library.

[karverma-amd]

Thanks @yichiche — good call, addressed in ae2083a.

Guard the NVIDIA common path. In qwen3_next.py, the fused-prefill defaults (_gdn_fused_prefill_ready=False, etc.) are now
set on every platform, and the AITER adapter import + setup is fenced behind if _is_hip: (lazy import). So the NVIDIA
path never imports or touches the AMD-only code. The runtime hand-off and the post-load publish already gate on
_gdn_fused_prefill_ready, which stays False off-HIP.

No bespoke flag / on-by-default via use_aiter. Dropped the SGLANG_QWEN3_NEXT_GDN_FUSED_BACKEND opt-in env var. The adapter now gates on the standard switch:
use_aiter = bool(envs.SGLANG_USE_AITER.get()) and is_hip()
so it's enabled by default wherever AITER is (--enable-aiter), no manual per-kernel flag. Since the kernel imports from
AITER, use_aiter is the predicate (the if_hip guard fences the common-code touchpoint; the gfx950 / Triton-3.8 specifics
gate inside available()). This also removed a raw os.getenv("SGLANG...") call site in favor of the envs descriptor.

[yichiche]

Please also added unit test for this modification.

[yichiche]

_aiter_gdn_prefill_reject_logged: we should remove this log from production code.

[yichiche]

The fused prefill path skips the DP-padding trim in forward. This docstring says DP-padding envelopes are declined, but the only extra guard is the piecewise-graph check.

forward narrows extend inputs to global_num_token_non_padded_cpu and zeros the pad tail, because those rows go into the residual, MoE, and collectives. try_fused_gdn_prefill runs on the full projected_qkvz. AITER allocates outputs with torch.empty((M, ...)) and does not require cu_seqlens[-1] == M, so a padded prefill either folds the pad rows into conv/SSM state or leaves an uninitialized tail.

Please return None (unfused chain) when global_num_token_non_padded_cpu is set and smaller than projected_qkvz.shape[0], or when cu_seqlens[-1] != projected_qkvz.shape[0].

Two related holes on the same path:

1. covered() does not pass A_log, dt_bias, or norm_weight. AITER re-checks those inside run() and raises ValueError (dt_bias must be bf16 [8], norm_weight bf16 [128]). Pass them into covered() so a mismatch falls back instead of failing the request.
2. _gdn_out_proj_fp8 is any Fp8LinearMethod.block_quant, including MXFP8 and block K other than 128. The kernel scales are fp32 group-128 [M, 8]. Gate the direct FP8 feed on block K == 128 and not MXFP8, and clear attn._gdn_fp8_out after out_proj.

Also, the PR body still documents SGLANG_QWEN3_NEXT_GDN_FUSED_BACKEND=aiter, which this commit removed. The repro should use --enable-aiter.

[karverma-amd]

Thanks @yichiche for the careful review — all points addressed across two commits
(93f49f1834 for the guards/logging/tests, c6cf779a34 for the dtype contract).

---

1. Remove the _aiter_gdn_prefill_reject_logged production log

Done — dropped the class flag and the rank0_log call; covered() now just falls
back on a miss (ok, _ = covered(...); if not ok: return None).

2. Fused path skips the DP-padding trim (radix_linear_attention.py:107)

Fixed. The backend now declines any batch whose projected rows exceed the real
varlen tokens, so it never folds pad rows into the conv/SSM state or leaves the
output tail uninitialized. New helper prefill_rows_match_real_tokens() returns
False when global_num_token_non_padded_cpu < rows or cu_seqlens[-1] != rows
— computed host-side (cu_seqlens[-1] == sum(extend_seq_lens_cpu)) to avoid a
per-layer device sync. The RadixLinearAttention.try_fused_gdn_prefill docstring is
corrected to match.

2.1 covered() didn't pass A_log / dt_bias / norm_weight

Fixed. covered() now takes and validates all three after AITER's own predicate
passes. Rather than a generic float check I pinned the exact contract the kernel
indexes off the pointer, cross-checked against AITER's reference harness:

• A_log → fp32 [v_heads] (read at full precision for the exp(-exp(A_log)) decay)
• dt_bias → bf16 [v_heads]
• norm_weight → bf16 [head_v_dim]

To keep the fused path enabled (the model stores dt_bias in fp32), the model now
casts dt_bias to bf16 at load — same as norm_weight — and publishes it through
the stash; the backend feeds that bf16 tensor. Verified on MI355X (gfx950):
covered() returns True for the bf16 contract and declines a fp32 dt_bias with
the expected message, and the fused output is bit-identical between fp32 and bf16
dt_bias (the tile upcasts to fp32, so bf16 rounding of the small bias vanishes in
the fp8 e4m3 activations — relMAE 0.00000 on activations, normalized, and delta
state).

2.2 _gdn_out_proj_fp8 matched any block_quant (incl. MXFP8 / block K ≠ 128)

Fixed via out_proj_accepts_group128_fp8(): the direct FP8 feed is now gated on a
group-128 block-FP8 GEMM (block_quant, not MXFP8, weight_block_size[-1] == 128);
every other method takes the bf16 re-quant path. The (fp8, scales) stash is also
cleared right after out_proj consumes it, so no later forward can read a stale pair.

3. Add a unit test

Added test/registered/attention/test_gdn_fused_prefill_guards.py (CPU CI,
base-a-test-cpu) — 16 cases covering the DP-padding decline, the out_proj-compat
gate (MXFP8 / K≠128 / per-tensor all decline), the covered() early declines
(non-SiLU, missing conv bias), and the per-vector dtype contract (fp32 dt_bias,
bf16 A_log, fp32 norm_weight, and wrong element count all decline).

4. PR body still documented the removed env var

Updated — the flag is gone; the repro now uses --enable-aiter (the fused path is
on by default wherever AITER is enabled on gfx950, fail-closed otherwise).

[yichiche]

#39873 Consume the AITER fused GDN prefill kernel — +703/−0 across 5 files. It reaches shared code only additively (gdn_backend.py, radix_linear_attention.py), so nothing that already ships changes behaviour. Guarded by use_aiter (SGLANG_USE_AITER and is_hip()), with the NVIDIA path fenced by _is_hip and gfx950 / Triton 3.8 inside available(). No new env var, no new global, no shared signature rewritten. /sglang-pr-review found no CRITICAL (0 high, 0 medium, 1 low — none blocking). The PR body carries accuracy and perf numbers, and a test lands in the diff. LGTM.