[FlyDSL] Add gfx950 Kimi Delta Attention prefill kernel

[amd-wsung102]

Summary

This PR adds a FlyDSL chunkwise Kimi Delta Attention (KDA) prefill kernel to aiter, exposed as flydsl_chunk_kda in aiter.ops.flydsl. It is a drop-in for the KDA prefill recurrence that Kimi-K3 runs (same signature as fla.ops.kda.chunk_kda / aiter's chunk_kimi_delta_attn): packed varlen [1, total_tokens, H, D] inputs, raw pre-activation gate/beta, and a V-first recurrent state.

A single FlyDSL path replaces the multi-kernel Triton stack (q/k L2-norm, gate cumsum, beta sigmoid, the two intra-chunk kernels, the w/u recompute, the state scan, and the output matmul) with:

• A fused Triton pack pass (varlen → padded dense, folding in L2-norm + gate activation + beta sigmoid),
• Two FlyDSL kernels (prep, parallel over chunks; scan, the serial state walk),
• A Triton unpack back to packed layout.

Performance improvement: 1.35x speedup on the KDA block.

Relevant Files

• aiter/ops/flydsl/kernels/kda_kernel.py - fused single-kernel variant (included; not used by the wrapper, see below)
• aiter/ops/flydsl/kernels/kda_split.py - sequence-parallel split kernels (prep + scan), the path in use
• aiter/ops/flydsl/kda_kernels.py - host wrapper / dispatch (flydsl_chunk_kda, kda_chunk_fwd, flydsl_kda_supported)
• aiter/ops/flydsl/kda_varlen.py - fused Triton pack / unpack
• aiter/ops/flydsl/__init__.py - exports
• op_tests/flydsl_tests/test_flydsl_kda.py - 32 correctness tests

Scope: gfx950, bf16, 128-wide K/V heads. Validated on MI355X, TP=8, serving Kimi-K3.

Optimizations Applied

• Sequence-parallel split path (not the fused kernel). The fused single-kernel variant factors its C×C tiles against one per-chunk reference row; at Kimi's gate_lower_bound = -5 a 32-token chunk accumulates up to 160 nats of decay, which overflows fp32 and silently produces NaN (the final state stays correct, so it passes a whole prefill before surfacing). The split path builds the same tiles with a bounded construction and stays exact — verified NaN-free out to 640 nats of decay, vs the fused path failing at ~90–115 nats. The split path is also the faster of the two at prefill chunk counts. The fused kernel is included for reference but the wrapper uses split exclusively.
• C=32 chunk with a 4-way value-channel split. Keeps LDS at ~71 KB so two workgroups fit per CU (160 KB); with the O(C²) shrink of the sequential triangular solve this is ~2.9× faster than C=64 despite the extra chunks.
• Occupancy-aware dv_split, chosen from the launch shape rather than a constant: the finest split that still fills idle CUs and keeps two workgroups resident per CU.
• Fused pack kernel. Scatters packed varlen → padded [N, H, Tpad, D] and folds in the q/k L2-norm, the lower-bounded-sigmoid gate activation, and the beta sigmoid — replacing FLA's separate l2norm_fwd, kda_gate_chunk_cumsum, and fused_beta_sigmoid passes. Padding slots are written as zeros, which is an exact no-op for the recurrence (beta = 0, g = 0, state carried unchanged).
• Single workspace pool — the six per-chunk tiles are carved from one allocation instead of six torch.empty calls (the per-call host cost is visible at short sequences).
• Module cache (lru_cache, 256 entries) amortizes the ~35 ms FlyDSL build across the 69 KDA layers that share a shape within a step.
• Perf-aware dispatch guard. Returns None (fall back) for: non-gfx950 / non-bf16 / non-128 heads; lopsided batches where right-padding to the max length would dominate; memory pressure (the guard counts the caching allocator's free pool, not just driver-free, so a KV-cache-heavy server still qualifies); and sequences past the perf crossover — ≤ 2048 tokens/seq for batches ≥ 4 sequences, ≤ 1024 for narrower ones — because the serial scan loses to the Triton kernel beyond ~2K tokens/seq. Thresholds are env-tunable (AITER_KDA_FLYDSL_MAX_SEQLEN, AITER_KDA_FLYDSL_MAX_SEQLEN_NARROW, AITER_KDA_FLYDSL_MAX_PAD_RATIO).

Performance

All measured on ATOM MI355X (gfx950), TP=8, Kimi-K3, bf16, head_dim = 128.
ATOM image used: rocm/atom-dev:latest.

End-to-end serving profile — 1k/1k, 128 prompts

Per-rank trace, identical 1035 KDA layer-calls each. The four kernels from the original profile plus the surrounding preprocessing / output kernels are replaced:

FLA path (µs / layer-call)	µs	FlyDSL path (µs / layer-call)	µs
chunk_kda_fwd_kernel_intra_sub_chunk	86.6	kda_scan	116.5
chunk_gated_delta_rule_fwd_kernel_h	73.8	kda_prep	101.2
recompute_w_u_fwd_kda_kernel	54.1	kda_pack (Triton)	40.6
chunk_kda_fwd_kernel_inter_solve	48.1	kda_unpack (Triton)	10.9
chunk_gla_fwd_kernel_o	46.9
kda_gate_chunk_cumsum	30.1
l2norm_fwd (×2)	20.2
fused_beta_sigmoid	4.6
Total	364.5	Total	269.1

1.35× on the KDA block (≈ 1.44× counting the elementwise / fill scaffolding the pack kernel also removes). Per prefill iteration (median over 12 bs=10, tok=10240 iterations): KDA time 27.6 ms → 21.1 ms; whole prefill iteration 662.3 → 653.3 ms; decode iterations unchanged (41.4 ms — decode uses the fused update kernel, not this path).

Accuracy

• Unit tests (test_flydsl_kda.py, 32 passing). flydsl_chunk_kda vs fla.ops.kda.chunk_kda at bf16 precision (relative error < 2e-2) across: multi-sequence packed prefills, ragged and non-chunk-aligned lengths, with / without initial state, saturated (worst-case) gates, raw-bf16 vs fp32 beta, NC=1 / single-token edges, out= aliasing, and the dispatch-guard fall-back cases.

• Long-context retrieval (needle-in-haystack), served through ATOM, FlyDSL forced on at all lengths, vs Triton baseline. Accuracy = planted 5-digit codes recalled correctly:

  Config	FlyDSL	Baseline
  1k in / 1k out	32/32 = 100%	100%
  8k in / 1k out	24/24 = 100%	100%
  1k in / 8k out	16/16 = 100%	100%

  (8k forced FlyDSL on so the kernel — not the fallback — was exercised; confirmed via trace that kda_prep / kda_scan ran for all KDA layers.)

• GSM8K (lm-eval-harness, 5-shot, full 1319-question test set), served through ATOM:

  Filter	FlyDSL	Baseline
  strict-match	95.98%	95.68%
  flexible-extract	95.98%	95.75%

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:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
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 4741 --add-label <label>