Repository navigation
[Triton/Gluon] Fused Qwen3-Next GDN decode op for gfx950 - #5520
Conversation
…decode Defines the op contract for a fused Qwen3-Next Gated DeltaNet decode kernel before any kernel lands: packed qkvz/ba split, depthwise causal conv1d with bias + SiLU, GDN gating, the fp32 recurrent state update, gated RMSNorm with a SiLU gate, and an optional per-head group FP8 quantization epilogue. This is a sibling of fused_kda_decode, not a variant of it. KDA uses a per-channel lower-bounded sigmoid gate, a sigmoid output gate, uniform head counts and no conv bias; Qwen3-Next GDN uses a per-head exp(-exp(A_log) * softplus(a + dt_bias)) decay, a SiLU output gate, num_v_heads == 2 * num_k_heads and a biased conv. The head topology alone (lp = H*K offsets) cannot express VH == 2*KH, so a mode flag on the existing kernel is not viable. The reference has been cross-checked on MI355X against all seven Artemis Gluon variants across all nine specialized batch sizes (M = 1,2,4,8,16,32,64,128,256). conv_state matches bit-exactly on every band; the other outputs agree within the ATOL/RTOL 0.05 bar this directory already uses for fused_kda_decode, with the residual being a few bf16 ULP from differing fp32 evaluation orders. The kernel-side tests are written against the pad_slot_id convention that aiter's other GDN ops already honour (fused_conv_recurrent_norm skips state_idx < 0; causal_conv1d_update_split_qkv takes an explicit pad_slot_id), rather than the reserved-slot-0 redirection the out-of-tree Artemis adapter maintains via a monkeypatched _replay_metadata. They skip until the op lands. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
One Gluon launch replaces the four-kernel Qwen3-Next decode chain (fused_qkvzba_split_reshape_cat -> causal_conv1d_update -> fused_recurrent_gated_delta_rule_packed_decode -> layer_norm_gated) and can optionally emit the per-head group-128 FP8 activations that a block-FP8 out_proj consumes directly, removing a separate quantization launch. A sibling of fused_kda_decode, not a variant of it: KDA uses a per-channel lower-bounded sigmoid gate, a sigmoid output gate, uniform head counts and no conv bias, and its lp = H*K offsets structurally cannot express VH == 2*KH. Verified bit-identical to the upstream Artemis kernel on every output (out, quantized, scales, ssm_state, conv_state) at M = 1, 32 and 256 on MI355X, and green against the fp32 reference across M = 1..128. Two changes on top of the upstream source: * pad_slot_id: rows whose state index is the sentinel are skipped without reading or writing state, matching the contract SGLang's replay_state_indices_validator documents and the one aiter's fused_conv_recurrent_norm already implements. This removes the upstream kernel's dependence on a caller-maintained reserved slot 0. * HAS_FP8: the quantization epilogue is optional, because an unquantized out_proj cannot accept an (fp8, scales) tuple. Support is reported by fused_gdn_decode_qkvz_supported() as (bool, reason) rather than asserted, so callers fall back and can log why. Coverage is deliberately narrow and stated: gfx950, VH == 2*KH, head dims 128, conv width 4. The Gluon body is generated by build_aiter_kernel.py, which pins the source sha256 and asserts every edit it applies; the source is ~600 lines of dense Gluon with explicit register layouts, where a transcription slip would surface only as a numerical mismatch. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
build_aiter_kernel.py reads the Artemis kernel pack, which is not part of this repository, so shipping it here would put a script in the tree that cannot run. Provenance stays where a reviewer needs it -- the generated kernel's header records the source file and its sha256 -- and the generator itself lives with the pack it reads. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
The maintainer guide asks for kernel + wrapper + test + benchmark; this is the missing fourth. Reports achieved HBM bandwidth against arch peak. The kernel performs no MFMA and no gl.dot -- it is bound by the FP32 recurrent state, read and written once per token per value head, which dominates every other tensor by an order of magnitude. So bytes moved is the right roofline denominator. Two measurement choices the script documents inline, because getting either wrong silently invalidates the result: it times under CUDA graph replay, since timing the Python call measures host overhead that at small batches is ~2.5x the kernel itself; and it sizes the state pool independently of the batch, since sizing it to the batch puts small batches entirely in cache. Measured on MI355X: the kernel is latency-bound at small batch (1.4% of peak at batch 1) and becomes bandwidth-bound from batch 64, saturating at 5.0-5.2 TB/s, 62-65% of the 8 TB/s peak. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags & labels: |
The header said 'us'. The timed region is a CUDA graph enclosing the whole wrapper, so the number is per call. This op always issues a single launch, so the two coincide here -- but the label should not depend on that staying true, and the Artemis sources it derives from do use multi-launch paths at small batch where per-call and per-launch differ. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
Above batch 128 the op now uses a second decomposition that keeps the whole [V, K] state tile in registers across both matrix-vector products, instead of the latency-hiding path that serves smaller batches. This is not a tuning of the existing kernel and could not be reached by one. The two move identical bytes -- one read and one write of the fp32 state tile, confirmed by source-level count -- and three candidate explanations for the gap were tested and rejected: launch-order locality (rewriting the grid to the other kernel's order gained 0.75%), store cache policy (.cs to .wt was slightly worse), and routing large batches through the kernel that does carry a batch ladder (about 5% of the gap). What remains is the register layout and reduction structure, which is a different algorithm. At batch 256 this takes the op from 53.5 to 46.5 us, and from 65% to 74.5% of peak HBM bandwidth with the same byte count. At 128 the two are within 2%; below 128 the existing path is better by up to 15%, so the threshold sits at the measured crossover. Selection is a range check rather than a shape table, so every batch size above the threshold is covered -- there are no gaps to fall into. Correctness is verified against the fp32 reference at 128, 192 and 256, and at the non-power-of-two sizes 136, 160, 200 and 248. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
There was a problem hiding this comment.
🟡 Changes recommended
Unresolved critical operand-layout and state-pool validation issues remain, along with dispatch, predicate, and test-coverage gaps.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Adds a gfx950-only fused Qwen3-Next GDN decode operation combining convolution, recurrence, normalization, and optional FP8 quantization.
Changes:
- Adds Gluon kernels, public wrapper, dispatch, and support validation.
- Exports the operation and adds kernel package initialization.
- Adds correctness tests and CUDA-graph performance benchmarks.
File summaries
| File | Summary |
|---|---|
op_tests/triton_tests/test_fused_gdn_decode_qkvz.py |
Reference and kernel correctness tests |
op_tests/op_benchmarks/triton/bench_fused_gdn_decode_qkvz.py |
Performance benchmark |
aiter/ops/triton/gated_delta_net/fused_gdn_decode_qkvz.py |
Public API, validation, and dispatch |
aiter/ops/triton/gated_delta_net/__init__.py |
Public operation exports |
aiter/ops/triton/_gluon_kernels/gfx950/gated_delta_net/fused_gdn_decode_qkvz.py |
Small- and large-batch Gluon kernels |
aiter/ops/triton/_gluon_kernels/gfx950/gated_delta_net/__init__.py |
Kernel package initialization |
Review details
Suppressed comments (11)
aiter/ops/triton/_gluon_kernels/gfx950/gated_delta_net/fused_gdn_decode_qkvz.py:479
- The wrapper launches
_fused_decode_tiled[...]directly, making this a launchable entry kernel, but it has no config-awarerepr=. Add amake_kernel_reprcallback with its meaningful constexpr/tuning keys, as required by the existing pattern inaiter/ops/triton/README.md:269-296, so compiled specializations remain identifiable in traces.
@gluon.jit
def _fused_decode_tiled(X, BA, History, State, Indices, Weight, Bias, ALog, DTBias,
aiter/ops/triton/_gluon_kernels/gfx950/gated_delta_net/fused_gdn_decode_qkvz.py:625
- The wrapper launches
_decode_group[...]directly, making this a launchable entry kernel, but it has no config-awarerepr=. Add amake_kernel_reprcallback with its meaningful constexpr/tuning keys, as required by the existing pattern inaiter/ops/triton/README.md:269-296, so compiled specializations remain identifiable in traces.
@gluon.jit
def _decode_group(
aiter/ops/triton/gated_delta_net/fused_gdn_decode_qkvz.py:45
- This 128-token crossover is a measured performance choice rather than a functional constraint, but it is hardcoded in the wrapper. Put the dispatch threshold in the op's tuning configuration so retuning the crossover does not require changing runtime source code.
_LARGE_BATCH = 128
aiter/ops/triton/gated_delta_net/fused_gdn_decode_qkvz.py:276
- The tiled launch also hardcodes
num_warps=8, so this decomposition bypasses config-driven tuning. Route its warp count through the same tuning configuration as the other launch paths.
num_warps=8,
aiter/ops/triton/gated_delta_net/fused_gdn_decode_qkvz.py:286
- The small-batch launch repeats a hardcoded
num_warps=8, leaving this specialization untunable through configuration as well. Use the configured launch parameters here instead of another literal.
num_warps=8,
aiter/ops/triton/gated_delta_net/fused_gdn_decode_qkvz.py:40
- The public wrapper advertises both
torch.float8_e4m3fnandtorch.float8_e4m3fnuz, but every kernel test uses onlyFP8_DTYPE = torch.float8_e4m3fn. The alternate format uses a different maximum (240 versus 448), so its scale and quantization path remain unverified.
_FP8_QUANT_MAX = {torch.float8_e4m3fn: 448.0, torch.float8_e4m3fnuz: 240.0}
aiter/ops/triton/gated_delta_net/fused_gdn_decode_qkvz.py:75
- The support predicate only checks these six tensors for CUDA, while
conv_bias,A_log,dt_bias, andnorm_weightare also dereferenced by the kernel without any device, shape, or dtype validation; the flat-indexed projections, weights, bias, and indices are likewise not required to be contiguous. A CPU or strided operand can therefore makesupported()return true and still be passed under the wrong ABI. Include every kernel operand in the support contract or reject unsupported layouts/devices before launching.
tensors = (
projected_qkvz,
projected_ba,
conv_state,
ssm_state,
ssm_state_indices,
conv_weight,
)
if not all(isinstance(t, torch.Tensor) and t.is_cuda for t in tensors):
aiter/ops/triton/gated_delta_net/fused_gdn_decode_qkvz.py:98
- For an empty/zero-head state pool,
k_headsbecomes zero; because the equality check then succeeds forv_heads == 0, the next line divides byk_headswhile constructinggroup_width. The advertised(bool, reason)predicate should reject non-positive head counts before this division instead of raisingZeroDivisionError.
k_heads = (channels - v_heads * head_v_dim) // (2 * head_k_dim)
if v_heads != 2 * k_heads:
return (
False,
f"requires num_v_heads == 2 * num_k_heads, got {v_heads}/{k_heads}",
)
group_width = 2 * head_k_dim + 2 * (v_heads // k_heads) * head_v_dim
op_tests/triton_tests/test_fused_gdn_decode_qkvz.py:278
- The reference pad-slot test checks slots 3 and above, but slot 0 is the reserved slot that would be touched by an implementation that incorrectly converted the
-1sentinel into a valid index. Add an assertion for slot 0 so the test actually detects that failure mode.
assert torch.equal(inp["conv_state"][3:], conv_before[3:])
assert torch.equal(inp["ssm_state"][3:], ssm_before[3:])
op_tests/triton_tests/test_fused_gdn_decode_qkvz.py:104
- The reference keeps the SiLU convolution result in fp32, but every kernel path returns this value as BF16 and converts it back to fp32 before Q/K normalization and the recurrence (
_convolve/_conv_channelsdo.to(bfloat16).to(float32)). This oracle therefore compares different computations; round the activation at the same boundary before using it downstream.
conv_out = acc * torch.sigmoid(acc)
op_tests/triton_tests/test_fused_gdn_decode_qkvz.py:315
- The correctness parameter list jumps from 64 to 128, but the wrapper routes every batch in 65–127 to
_fused_decode_tiledwithB32=False; that entry path is never exercised by this test file. Add a representative batch in that interval (and keep the boundary cases) so the second decomposition is actually validated.
@pytest.mark.parametrize("batch", [1, 2, 4, 8, 16, 32, 64, 128, 192, 256])
def test_fused_gdn_decode_correctness(batch):
- Files reviewed: 5/6 changed files
- Comments generated: 12
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
…ions Review findings from ROCm#5520, verified against the code before fixing. The most serious was a hole in the tests: the correctness case computed ref_quant and never compared it, so the FP8 output was checked only through its scales. A broken epilogue could pass every case as long as the bf16 output, scales and state updates were right. Now asserted, in value space rather than code space -- a 1-ULP e4m3 code difference near the top of the range is enormous as a raw integer and meaningless as a value, so both sides are dequantized with their own scale. The support predicate claimed more than it checked: - conv_state and ssm_state are indexed with the same slot, but their slot counts were never compared, so a shorter conv pool passed and was then read out of bounds - only the two pools were required contiguous, while projected_qkvz, projected_ba, conv_weight, conv_bias, A_log, dt_bias and norm_weight are all addressed with flat pointer arithmetic -- a same-shaped strided view passed and was read from the wrong locations - conv_bias was missing from the device check, and A_log, dt_bias and norm_weight were not parameters at all, so the predicate could return (True, "") for a launch whose operands were on the CPU or wrongly shaped Also: the bf16-only path is now checked numerically rather than only for dtype and None, since HAS_FP8=False is a separate compiled specialization; a broken op module now fails instead of silently skipping every kernel test, with the skip narrowed to the optional Gluon dependency; the kernel module no longer imports torch, so it stays importable without PyTorch as the convention requires; and the three launchable kernels carry config-aware repr= callbacks. Dropped float8_e4m3fnuz rather than testing it: that is the gfx942 format and this op is gfx950-only, so advertising it promised a path that can never run. PAD_SLOT_ID is deliberately excluded from the repr keys. It is a protocol sentinel rather than a tuning key, and make_kernel_repr does not sanitize the sign of a negative value -- including it emits an invalid function identifier and fails at compile. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
ROCm/aiter#5520 review widened fused_gdn_decode_qkvz_supported() to take A_log, dt_bias and norm_weight, which it previously never saw -- it could return covered for a launch whose parameter tensors were on the CPU or wrongly shaped, after which the kernel dereferenced them at fixed offsets. Threads those through covered() and the backend call site so the delegation stays complete, which is the point of delegating: the two contracts cannot drift if SGLang forwards everything AITER checks. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
…ions Review findings from ROCm#5520, verified against the code before fixing. The most serious was a hole in the tests: the correctness case computed ref_quant and never compared it, so the FP8 output was checked only through its scales. A broken epilogue could pass every case as long as the bf16 output, scales and state updates were right. Now asserted, in value space rather than code space -- a 1-ULP e4m3 code difference near the top of the range is enormous as a raw integer and meaningless as a value, so both sides are dequantized with their own scale. The support predicate claimed more than it checked: - conv_state and ssm_state are indexed with the same slot, but their slot counts were never compared, so a shorter conv pool passed and was then read out of bounds - only the two pools were required contiguous, while projected_qkvz, projected_ba, conv_weight, conv_bias, A_log, dt_bias and norm_weight are all addressed with flat pointer arithmetic -- a same-shaped strided view passed and was read from the wrong locations - conv_bias was missing from the device check, and A_log, dt_bias and norm_weight were not parameters at all, so the predicate could return (True, "") for a launch whose operands were on the CPU or wrongly shaped Also: the bf16-only path is now checked numerically rather than only for dtype and None, since HAS_FP8=False is a separate compiled specialization; a broken op module now fails instead of silently skipping every kernel test, with the skip narrowed to the optional Gluon dependency; the kernel module no longer imports torch, so it stays importable without PyTorch as the convention requires; and the three launchable kernels carry config-aware repr= callbacks. Dropped float8_e4m3fnuz rather than testing it: that is the gfx942 format and this op is gfx950-only, so advertising it promised a path that can never run. PAD_SLOT_ID is deliberately excluded from the repr keys. It is a protocol sentinel rather than a tuning key, and make_kernel_repr does not sanitize the sign of a negative value -- including it emits an invalid function identifier and fails at compile. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
|
Addressing copilot raised comments. Pushed 8c1e530. Fixed Test coverage
Support predicate — it claimed more than it checked
Conventions
float8_e4m3fnuz: removed rather than tested. It's the gfx942 format and this op is gfx950-only, so advertising it promised a path that can never run. Happy to reinstate with tests if you'd rather keep the dtype surface. Two open questions for @ROCm/triton_kernels_team :
|
…ations The style checks were red on files I had not linted the same way CI does. Three distinct causes: The generated Gluon module was never formatted. It is copied from sources that are not black-formatted, so black rejected it. The generator now formats its own output rather than leaving it to be hand-patched. My local ruff runs used --select E4,E7,E9,F, which is narrower than the rule set the repository's pyproject actually enforces, so C408, ISC004, RUF059, UP006, UP035, UP045 and I001 all went unseen. Fixed at source: dict() calls rewritten as literals, implicit string concatenation inside return tuples parenthesised, typing.Optional/Tuple replaced with PEP 604/585 spellings, imports sorted, and the unread second element of the inline-asm pairs renamed to _scratch. Also removes build_aiter_kernel.py, which had been copied to the repository root by mistake and was never meant to ship here. Verified with the versions CI pins -- black 26.5.1 and ruff 0.16.0 -- run from a tree carrying this repository's pyproject, over every file the branch touches. 26 op tests still pass. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
8c1e530 to
48ec7cc
Compare
ROCm/aiter#5520 review widened fused_gdn_decode_qkvz_supported() to take A_log, dt_bias and norm_weight, which it previously never saw -- it could return covered for a launch whose parameter tensors were on the CPU or wrongly shaped, after which the kernel dereferenced them at fixed offsets. Threads those through covered() and the backend call site so the delegation stays complete, which is the point of delegating: the two contracts cannot drift if SGLang forwards everything AITER checks. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
Ruff's I001 fired on three files. My earlier local run missed it because I linted copies flattened into a scratch directory, where the aiter package path no longer matched, so isort classified aiter imports as third-party and sorted them into a different, self-consistent order. Running ruff 0.16.0 against the real paths with the repository's own pyproject reproduces CI exactly. Worth noting for anyone reproducing: the ruff step uses reviewdog with filter-mode=diff_context and -fail-on-error=true, so any finding on a changed line fails the check -- whole-repo cleanliness is not the bar, the diff is. Verified with the pinned versions, black 26.5.1 and ruff 0.16.0, on the real tree. 26 op tests still pass. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
|
@rbrugaro-amd two more changes
So the file now holds two implementations. Is this expected? If not, please make
|
Mirrors the public wrapper's category, as requested in review. Adds __init__.py to match the sibling test directories. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
Addresses the review asks to stop duplicating operations inside the kernel file and to use the repository's activation helpers. _decode_group now calls _head_dynamics and _quantization_scale instead of re-implementing the gate, softplus, exp, sigmoid and FP8 scale inline, and _sigmoid delegates to _triton_kernels/activation._sigmoid_exp2. Outputs are bit-identical to the previous kernel on all five results -- bf16 output, quantized, scales, and both mutated state pools -- at batch 1, 2, 8, 32, 64, 128, 192 and 256. Kernel time is unchanged: 10.47 -> 10.53 us at batch 32, 27.03 -> 26.76 at 128, 46.58 -> 46.45 at 256. _normalize_qk is left in place. Unlike the others it is not a call-site swap: it takes a shared-memory handle and reads through .gather() with a hardcoded 128 offset, where the large band uses .slice().load() with an explicit gl.SliceLayout tied to its register layout. The arithmetic is the same either way, so routing through it would change layouts in a kernel that is bandwidth-bound above batch 128 without changing results. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
Thanks @azaidy — addressed in Four of the five operations in your table now route through the shared Summary against your table:
1. Large band now calls the shared device helpers
2. Using
|
| batch | out (bf16) | quantized | scales | ssm_state | conv_state |
|---|---|---|---|---|---|
| 1 | ✅ | ✅ | ✅ | ✅ | ✅ |
| 2 | ✅ | ✅ | ✅ | ✅ | ✅ |
| 8 | ✅ | ✅ | ✅ | ✅ | ✅ |
| 32 | ✅ | ✅ | ✅ | ✅ | ✅ |
| 64 | ✅ | ✅ | ✅ | ✅ | ✅ |
| 128 | ✅ | ✅ | ✅ | ✅ | ✅ |
| 192 | ✅ | ✅ | ✅ | ✅ | ✅ |
| 256 | ✅ | ✅ | ✅ | ✅ | ✅ |
v_rcp_f32 and IEEE divide differ by ≤1 ULP in fp32, but every use site rounds
to bf16 immediately, so nothing propagates. Since no output bit moves, the
accuracy numbers already reported still hold unchanged. 26 op tests pass.
Kernel time is unchanged
Same-session A/B, kernel swapped in place, 3 repeats each (MI355X, 512-slot
state pool, CUDA-graph replay):
| batch | before | after | delta |
|---|---|---|---|
| 32 | 10.47 µs | 10.53 µs | +0.6% |
| 128 | 27.03 µs | 26.76 µs | −1.0% |
| 256 | 46.58 µs | 46.45 µs | −0.3% |
Full sweep after the change, median of 3:
| batch | 1 | 2 | 4 | 8 | 16 | 32 | 64 | 128 | 192 | 256 |
|---|---|---|---|---|---|---|---|---|---|---|
| µs/call | 9.41 | 9.40 | 9.48 | 9.67 | 9.66 | 10.85 | 14.44 | 26.72 | 36.57 | 46.42 |
| TB/s | 0.12 | 0.23 | 0.46 | 0.89 | 1.79 | 3.21 | 4.79 | 5.19 | 5.68 | 5.96 |
| % of peak | 1.4 | 2.9 | 5.7 | 11.2 | 22.4 | 40.1 | 59.9 | 64.8 | 71.0 | 74.5 |
_normalize_qk — the exception flagged above
This is the one row I did not unify, and it differs in kind from the others.
It is not a call-site swap. _normalize_qk takes a shared-memory handle and
reads via shared.gather(col, 0) / shared.gather(128 + col, 0). The large band
reads the same values through a different primitive:
q = conv_shared.slice(0, K).load(gl.SliceLayout(0, layout))
key = conv_shared.slice(K, K).load(gl.SliceLayout(0, layout))Three differences: .gather() vs .slice().load(); a hardcoded 128 where the
band uses constexpr K; and the band's loads carry an explicit gl.SliceLayout
tied to its register layout. Routing through the helper would change that layout
in a kernel that is bandwidth-bound above batch 128 — 74.5% of peak at 256 —
where register pressure drives occupancy.
The arithmetic is identical either way (rsqrt(sum(x*x) + 1e-6), with q
scaled), so this is a layout and occupancy question rather than a numerics one.
Happy to revisit if you'd prefer it unified regardless.
Resolves the one conflict, in aiter/ops/triton/gated_delta_net/__init__.py: upstream added fused_gdn_prefill_qkvz at the same location this branch adds fused_gdn_decode_qkvz. Both are kept, imports and __all__ alphabetically sorted, matching the repository's isort configuration. Upstream did not touch this branch's dependencies -- activation.py, kernel_repr.py and arch_info.py are unchanged -- and the decode and prefill Gluon kernels share no helper names. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
vgokhale
left a comment
vgokhale
left a comment
There was a problem hiding this comment.
Review: PR #5520 — Fused Qwen3-Next GDN decode (gfx950 Gluon, fp8 group-quant)
This is the decode counterpart to the prefill op in #5606, and structurally it's in good shape: the canonical four files (kernel in _gluon_kernels/gfx950/gated_delta_net/, wrapper in ops/triton/gated_delta_net/, test in the gated_delta_net/ test category, benchmark), no separate _launch file, repr on every launched kernel, a _supported() probe + ValueError-on-misuse wrapper, and an fp32 reference that's itself pinned. Dtypes are clean (bf16 in, fp8 e4m3fn out, fp32 reference only).
Two blockers, both already solved correctly by the prefill sibling ten files away — the fix for each is "do what fused_gdn_prefill_qkvz.py does."
1. Blocker — top-level Gluon import breaks import-safety and defeats the _supported() probe
The wrapper imports the Gluon tile kernels at module top level:
# aiter/ops/triton/gated_delta_net/fused_gdn_decode_qkvz.py:28
from aiter.ops.triton._gluon_kernels.gfx950.gated_delta_net.fused_gdn_decode_qkvz import (
_decode_group, _fused_decode, _fused_decode_tiled,
)and that kernel module does from triton.experimental import gluon at its own top level (kernel L28–29). The whole chain is unconditional — no try/function guard anywhere — and gated_delta_net/__init__.py:22 re-exports the wrapper. So import aiter.ops.triton.gated_delta_net eagerly pulls the Gluon dialect and raises ImportError on any build without triton.experimental.gluon (Triton < 3.8, the ROCm 3.7 backport, CPU / non-gluon platforms).
This is self-defeating: the wrapper docstring promises callers "fall back instead of crashing" via fused_gdn_decode_qkvz_supported() — but if importing the module already crashed, the probe is unreachable on exactly the platforms it exists for. Because the shared __init__ imports both siblings, it also poisons the carefully-deferred prefill path next to it.
The prefill wrapper is the precedent: it has no top-level Gluon import; all five of its gluon-kernel imports are lazy, inside each host launcher (e.g. fused_gdn_prefill_qkvz.py:209/381/555/860/1128), which it documents as "mirroring gemm/basic/gemm_a16w16.py."
Fix: delete the module-level import (wrapper L28–32) and import _decode_group / _fused_decode / _fused_decode_tiled lazily at the launch sites (L289 / L319 / L329), behind the gfx950 gate.
2. Blocker — no Triton ≥ 3.8 version gate
fused_gdn_decode_qkvz_supported() gates only on get_arch() != "gfx950" (L67) plus shape/dtype/contiguity. There is no version check anywhere in the wrapper or kernel.
#5606 added an explicit >= 3.8 gate (_gluon_supported()) for a specific reason: ROCm backported an earlier, incompatible Gluon into some 3.7 builds where import triton.experimental.gluon succeeds but the tile layout ops miscompile — so an import probe isn't enough, you need a numeric version floor.
This is not hypothetical for decode: the decode kernel is built from the same 3.7-fragile primitives — gl.BlockedLayout (×11, same 4-tuple form), gl.convert_layout (×14), gl.amd.slice (×4), plus gl.amd.cdna3.buffer_load/store. On a 3.7-backport environment the import sails past (gluon present), then execution reaches these ops and miscompiles. Combined with #1, that's a crash with no clean fallback.
Fix: add the >= 3.8 check to _supported(), mirroring the sibling's _gluon_supported().
3. Minor — batch sweep is a scaled-duplicate ladder
@pytest.mark.parametrize("batch", [1, 2, 4, 8, 16, 32, 64, 128, 192, 256]) appears twice (test L245, L317) → 20 cases on batch alone. The kernel has two regimes split at _LARGE_BATCH = 128, so 1/2/4/8/16/32/64 are the same small-path regime at different sizes — scaled duplicates. Keep the distinct regimes plus the crossover boundary, e.g. 1, 64, 127, 128, 256, which pins both paths and the boundary with ~half the cases.
Summary
Structure, repr, test anchoring, and dtype discipline all match where #5606 landed. The two blockers are a regression of import-safety work the prefill sibling already got right: move the Gluon imports off the module top level and into the launch sites, and add the Triton ≥ 3.8 gate to _supported(). Both are confirmed (not stylistic): the import chain is provably eager and ImportErrors without gluon, and the decode tiles actually use the layout ops the version gate exists to protect. Trim the batch ladder and it's ready.
The wrapper imported the Gluon tile kernels at module scope, so the module raised ImportError on any build without triton.experimental.gluon -- before fused_gdn_decode_qkvz_supported() could report the reason and let the caller fall back, which is the one case the probe exists for. Import the three kernels at their launch sites instead, and defer arch_info for the same reason: detection can touch the driver at import time. Mirrors fused_gdn_prefill_qkvz.py. Add _gluon_supported(), an explicit Triton >= 3.8 floor checked before the import probe. Some ROCm 3.7 builds carry an earlier, incompatible Gluon where the import succeeds but the layout ops these kernels use miscompile, so probing the import alone is not enough. Trim the batch ladder from ten sizes to one per dispatch branch (1, 32, 64, 127, 128, 256), adding 127 to cover the previously untested 65-127 path. The branch mapping is recorded above the list. The Gluon kernel is unchanged and outputs stay bit-identical across batches 1-256, so the reported accuracy and end-to-end numbers are unaffected. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
|
@vgokhale Thanks for the review, all 3 feedbacks fixed in 1. Gluon imports are lazyModule-level import removed; the three kernels are imported at their launch Verified no top-level survivors:
2. Triton ≥ 3.8 gate
Behaves as intended with Gluon blocked: and on the target machine: 3. Ladder trimmed — but 32 kept
I kept 32 against your list, because it isn't a scaled duplicate — it's its
Validation after changesRebuilt the test environment from a pristine checkout and overlaid the merged
Kernel timing unchanged — median of 3, shipped benchmark, 512-slot pool:
The Gluon kernel file is byte-unchanged — the diff is two files, host-side only |
[Triton/Gluon] Fused Qwen3-Next GDN decode op for gfx950
Motivation
Qwen3-Next's Gated DeltaNet decode step runs as four separate kernels today:
aiter has all the pieces for GDN —
chunk_gated_delta_rule,fused_recurrent_gated_delta_rule,causal_conv1d_update_split_qkv,gated_rmsnorm_fp8_group_quant— but not the single-launch fusion, and notone that consumes the packed
qkvz/baprojections directly.This adds that op. On 4x MI355X / TP4 with Qwen3-Next-80B-A3B-Instruct-FP8 it is
worth +6.0% to +9.1% end-to-end output throughput when wired into SGLang
(companion PR), with no accuracy regression.
Why a new op rather than a mode on
fused_kda_decodefused_kda_decodeis the same shape of fusion — conv, recurrence, gated norm,one launch — but it implements KDA, and the differences are structural, not
parametric:
fused_kda_decodeH; q/k/v allHheadsVH == 2*KHlower_bound * sigmoid(exp(A)*(a+dt)), vector over Kexp(-exp(A_log)*softplus(a+dt)), scalar per headsigmoid(z)z * sigmoid(z)(SiLU)mixed_qkv/gate/betaqkvz/ba, split internallyThe head topology alone settles it: KDA's
lp = H*K,q_off = i_h*K,v_off = 2*lp + i_h*Voffsets cannot expressVH = 2*KH. A mode flag woulddouble an existing kernel's branch factor to host a different algorithm, so this
ships as a sibling instead.
fused_kda_decodeis not exported fromgated_delta_net/__init__.py, so no existing interface is perturbed.Modifications
New:
aiter/ops/triton/_gluon_kernels/gfx950/gated_delta_net/fused_gdn_decode_qkvz.py— the Gluon kernels. Two decompositions behind one entry point:
_fused_decode/_fused_decode_tiled. One eight-waveworkgroup owns a token/key-head group's convolution history and both FP32
value-head states, with prefetch and tiling to hide memory latency while
the machine is under-occupied.
_decode_group, which keeps the whole[V, K]statetile in registers across both matrix-vector products. Above that threshold
the kernel is bandwidth-bound, where the latency-hiding machinery stops
paying and the simpler traffic pattern reaches a better fraction of peak.
Selection is a range check, not a shape table, so every batch size is
covered.
aiter/ops/triton/gated_delta_net/fused_gdn_decode_qkvz.py— public wrapper,shape/dtype validation and the support predicate.
op_tests/triton_tests/test_fused_gdn_decode_qkvz.py— a pure-torch fp32reference plus correctness, pad-slot, bf16-only and determinism tests.
op_tests/op_benchmarks/triton/bench_fused_gdn_decode_qkvz.py— bandwidth androofline sweep.
Modified:
aiter/ops/triton/gated_delta_net/__init__.py— export the op and itspredicate.
What it fuses
projected_qkvz-> Q/K/V/Z andprojected_ba-> b/adecay = exp(-exp(A_log) * softplus(a + dt_bias))ssm_state[indices]Stage 6 is opt-in via
quant_dtype=. WithNonethe op returns bf16 only,because an unquantized
Linearconsumer cannot accept an(fp8, scales)tuple.pad_slot_idRows whose state index equals
pad_slot_id(default-1) are skipped withoutreading or writing state. This matches what the other indexed-state ops in this
directory already do —
fused_conv_recurrent_normreturns onstate_idx < 0,causal_conv1d_update_split_qkvtakes an explicitpad_slot_id— and it is thecontract SGLang's
replay_state_indices_validatordocuments. Measured cost:0.2-0.3%, i.e. within noise.
Scope and supported shapes
Deliberately narrow, and reported rather than asserted:
fused_gdn_decode_qkvz_supported()returns(bool, reason)so callers fallback and can log why.
v_exp_f32/v_rcp_f32inline asm,explicit register layouts
num_v_heads == 2 * num_k_heads;head_k_dim == head_v_dim == 128dt_bias/norm_weight;fp32
A_logandssm_state; int32 state indicese4m3fnore4m3fnuzNo CUDA path is affected. The op is reachable only through its own module,
and every caller is expected to gate on the predicate.
Accuracy Tests
pytest op_tests/triton_tests/test_fused_gdn_decode_qkvz.py— 22 passed onMI355X / ROCm 7.0 / Triton 3.8.
test_fused_gdn_decode_correctnessat batch 1, 2, 4, 8, 16, 32, 64, 128 —output, scales, and both mutated state pools against the fp32 reference at
ATOL/RTOL 0.05, the bar this directory already uses for
fused_kda_decode.test_fused_gdn_decode_pad_slot—PAD_SLOT_IDrows leavessm_stateandconv_statebit-unchanged.test_fused_gdn_decode_bf16_only—quant_dtype=Nonereturns(out, None, None).test_fused_gdn_decode_determinism— same inputs twice, bit-identical.round-trip independently of the kernel.
The reference itself was validated before being trusted: run against seven
independent implementations of this kernel across nine batch sizes, it agrees
on every one, and
conv_statematches bit-exactly on all 63 combinations.Speed Tests and Profiling
Hardware: 4x AMD Instinct MI355X (gfx950), TP4 / EP1.
Image:
lmsysorg/sglang-rocm:v0.5.19-rocm700-mi35x-20260910(ROCm 7.0,Triton 3.8.0, torch 2.9). Model:
Qwen/Qwen3-Next-80B-A3B-Instruct-FP8.Kernel level
Per-launch time under CUDA graph replay with a fixed 512-slot state pool.
Both matter: timing the Python call instead of the kernel produces a flat ~25 us
floor that is host overhead, and sizing the pool to the batch puts small batches
entirely in cache.
Replaced chain, same conditions, per decode step over 36 GDN layers:
830.8 us -> 279.2 us (2.98x) for the four kernels listed in Motivation.
Roofline
python3 op_tests/op_benchmarks/triton/bench_fused_gdn_decode_qkvz.py, MI355X(8.0 TB/s peak HBM). The op performs no MFMA and no
gl.dot; it is bound by theFP32 recurrent state, read and written once per token per value head, which
dominates every other tensor by an order of magnitude.
Two regimes, and they are why the op carries two decompositions. Up to batch
~32 it is latency-bound — there is not enough state traffic to fill the
machine, and time is nearly flat at ~9.4-10.4 us regardless of batch, so
prefetch and tiling pay. From batch 64 it becomes bandwidth-bound, where
that machinery stops paying and the register-resident tile reaches a better
fraction of peak: 74.5% at batch 256.
A single decomposition plateaued at 5.20 TB/s (65%) there. Three candidate fixes
were tested and rejected before concluding the gap was structural: extra HBM
passes (identical byte counts, by source-level count of the state load/stores),
launch-order locality (rewriting the grid to the other kernel's order gained
0.75%), and store cache policy (
.csto.wt, slightly worse). What remains isthe register layout and reduction structure — a different algorithm, not a
tuning constant.
End to end (companion SGLang PR)
Un-profiled, random dataset ISL/OSL 1024/1024,
num-prompts == max-concurrency.GSM8K 1319q: 0.941 baseline vs 0.941 with the op,
Invalid: 0.000on both.End to end at a raised graph cap
The large-batch band only fires above batch 128, which requires
--cuda-graph-max-bs-decodeabove its default of 64. At a cap of 256, two armson the same base, un-profiled:
Reproduction
docker run --rm -it --device /dev/kfd --device /dev/dri --network host \ --group-add video --ipc host --shm-size 32g \ -v $HF_HOME:/hf lmsysorg/sglang-rocm:v0.5.19-rocm700-mi35x-20260910 pytest op_tests/triton_tests/test_fused_gdn_decode_qkvz.py -qChecklist
op_tests/triton_tests/)black/ruffcleanop_tests/op_benchmarks/triton/bench_fused_gdn_decode_qkvz.py)