[Triton/Gluon] Fused Gluon KDA decode path (950/1250)

[omuhamma]

Motivation

Decode KDA path with high perf

Test Plan

Unit tests, benchmark, e2e runs (accuracy & perf)

Test Result

All unit tests pass

benchmarks:

When not updating the state pool with every state we save a lot of runtime when the tokens and batch size grows:

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4584 --add-label <label>

[vgokhale]

Can you add perf numbers for a few shapes? Otherwise LGTM

[zufayu]

Advisory review (static + hand-run; not a merge gate). Validation/Perf ran no GPU stage — reasons are on their lines below. Findings tagged [verified] are traced to code/repro.

Adds a gfx1250-only Gluon decode kernel for Kimi Delta Attention (per token S ← exp(g_t)·S, S ← S + β_t·k_t⊗(v_t − Sᵀk_t), o_t = Sᵀq_t) with paged per-token state slots, varlen, speculative decoding, cached (a,k,err) state updates, in-kernel gate/l2norm/sigmoid fusion, and opt-in TDM async-copy paths, behind a new standalone wrapper that nothing in aiter imports yet.

Review (advisory): 🔴 HIGH RISK
Validation (deterministic): NOT RUN — auto-validation disabled (REVIEW_AUTO_VALIDATE=0); the single test target op_tests/triton_tests/attention/test_kda.py is gfx1250-gated and skips wholesale on this gfx942 box (224 skipped, run here), so no deterministic run was possible
Perf (advisory): NOT RUN — the target exposes no benchmark entry point and the PR's own op_tests/op_benchmarks/triton/bench_kda.py exits "gfx1250 only" on this box; the PR description itself carries no numeric perf data (two trace screenshots only)

🔴 [verified] aiter/ops/triton/attention/kda.py sets stride_indices_seq=1 for 1-D ssm_state_indices while the kernel indexes state_indices_ptr + i_nstride_indices_seq + t, which is only correct for N==1 — the docstring's "[N, T] (or flat)" invites flat tables, and the wrapper accepts a flat [NT] int32 table with N=2, T=2 (verified here by launching it), so sequence 1 reads its seed state from indices[1] — sequence 0's token-1 slot — instead of indices[2] and writes its snapshots into other sequences' slots, silently corrupting the pool with no error. Author must assert ndim==2 (or N==1 for 1-D, or derive the row stride from the sequence length) so the flat form either works or is rejected.
🔴 [verified] aiter/ops/triton/attention/kda.py validates dt_bias (numel==HVK) and A_log contiguity and passes USE_LOWER_BOUND=True to the launch even when use_gate_in_kernel is left at its default False (verified by launching the wrapper with A_log + dt_bias + lower_bound=-5.0), but the kernel then computes a=exp(g) with no gate chain — so a caller that passes the K3 gate inputs without the flag gets silently ungated decay and garbage output, no error, while the wrapper's own shape validation signals the parameters are consumed. Author must assert A_log/dt_bias/lower_bound require use_gate_in_kernel=True (or forward them into the ungated path's error message).
⚠️ [verified] The whole 224-item suite in op_tests/triton_tests/attention/test_kda.py skips on every runner this PR's CI triggered (MI300X gfx942 / MI35X gfx950 — run here: 224 skipped on gfx942), and the only gfx1250 workflow (.github/workflows/ffm-bringup.yaml, label ci:gfx1250-ffm-triton) was never added, so "All unit tests pass" in the description rests entirely on the author's local run and no CI execution backs a 734-line kernel. Author must add the ci:gfx1250-ffm-triton label (or post a local gfx1250 run log) before merge.
⚠️ [verified] The PR's only perf evidence is two trace screenshots with no numeric values — its own benchmark op_tests/op_benchmarks/triton/bench_kda.py is gfx1250-only and published no table, maintainer vgokhale asked for numbers on Aug 6 and none were added — and no bucket in aiter/ops/triton/configs/gfx1250-KDA_DECODE-DEFAULT.json enables any advertised TDM path (the JSON sets only BV/SK/num_warps/num_buffers, and use_tdm_fused_load additionally calls tdm.async_load_fused, which stock triton 3.7.0 — checked here — does not have), so a reviewer can neither read the claimed speedup nor tell which configuration the screenshots ran. Author must publish numeric gluon-vs-fla tables for the K3 shapes (H=24/12, B=4–256, T=1–8) and name the config each number was run with.
⚠️ [inferred] aiter/ops/triton/attention/kda.py accepts num_accepted_tokens containing 0 (verified by launch) and the kernel then computes seed = load(num_accepted_ptr + i_n) − 1 = −1 and reads state_indices_ptr + i_nstride_indices_seq − 1 — four bytes before the slot row, before the tensor for sequence 0 — then uses the garbage as a state slot index; worth checking whether any speculative-decode caller can report zero accepted tokens, since vLLM/SGLang verifiers normally guarantee ≥1. Author must assert (num_accepted_tokens >= 1).all() or clamp the seed to 0.