[AMD][Kimi-K3] Add FlyDSL KDA fused decode kernel for gfx950

[nehaprakriya]

Motivation

The Kimi-K3 KDA (Kimi Delta Attention) linear attention decode path launches 3 separate GPU kernels per layer across 69 layers per forward pass: causal_conv1d_update → kda_packed_decode → rms_norm_gated. This adds 138 kernel launches and intermediate HBM round-trips per decode step.

A FlyDSL kernel for AMD gfx950 (MI355X) fuses all three operations into a single dispatch.

Modifications

python/sglang/srt/layers/attention/linear/kda_backend.py: Add a gfx950 fast path in forward_decode that checks for _k3_flydsl_decode_args (stashed by the model's prepare_weights for Kimi-K3 on gfx950) and calls flydsl_kimi_k3_kda_decode when conditions are met. Falls through to the existing causal_conv1d_update chain if unavailable. CUDA-graph compatible.

The conv_state pool layout [slots, 3, 4608] is passed via .transpose(-1, -2); the FlyDSL kernel reads layout via explicit stride arguments — no tensor copy needed.

Accuracy Tests

Correctness verified on MI355X (gfx950): output matches the triton fallback chain to within bf16 numerical precision across all tested shapes at B=1 decode.

Speed Tests and Profiling

Eliminates 138 kernel launches per decode step on Kimi-K3. Effect is neutral at high concurrency (MoE compute dominates); latency reduction visible at CONC=1 where kernel launch overhead is a larger fraction of step time.

Repro Steps

Hardware: AMD MI355X (gfx950), ROCm 7.x

Unit test (correctness):

python test/manual/test_amd_flydsl_kda_fused_decode.py

End-to-end with FlyDSL path enabled (Kimi-K3 on gfx950):

# FlyDSL path activates automatically on gfx950 when Kimi-K3 model is loaded
python -m sglang.launch_server \
  --model <kimi-k3-path> --tp 8 \
  --trust-remote-code

# Benchmark at low concurrency where launch savings are visible:
python -m sglang.bench_serving \
  --model <kimi-k3-path> --tp 8 \
  --num-prompts 200 --input-len 1024 --output-len 1024 \
  --concurrency 1

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.

---

CI States

Latest PR Test (Base): ❌ Run #30775092198
Latest PR Test (Extra): ❌ Run #30775092144

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[RolaoDenthu]

Hi, Thanks for the contribution. I think this pr is built on ROCm/aiter#4495, which is still a draft. We'd need that merged into main before this can be merged.
Could you also add the dependency to the PR description?

[RolaoDenthu]

Can't find this kernel in aiter main.