flydsl: gfx942 FP8 MQA logits indexer kernel (+ Triton FN/FNUZ fix)

[akii96]

Motivation

DeepSeek-V4 Flash's sparse attention uses a lightning indexer (fp8 MQA logits). This PR adds a FlyDSL gfx942 drop-in for the Triton fp8_mqa_logits that is faster on the DSv4 prefill shapes for both the mono-dtype config (q/k both e4m3fnuz) as well as the currently vllm upstream's mixed config.

The uplift does not depend on operand dtype: even with both operands fnuz (no conversion on either side), FlyDSL is 1.5–2.5× faster at the kernel level and cuts end-to-end prefill TTFT by up to 17% at 128K and increases with increasing ctx size.

Technical Details

1. New FlyDSL gfx942 kernel which is a drop-in for the Triton one (same args, return, and clean_logits semantics).
2. (Optional safety net) Triton fp8_mqa_logits FN/FNUZ fix by @vpietila-amd; a no-op for fnuz/fnuz, so existing callers are unchanged. This only matters while callers still send mixed q=fn/k=fnuz (see note below); it can be dropped if maintainers prefer the source-side fix only.

FlyDSL kernel initially translated from GEAK by @peyron-amd, with perf iterated by myself and @vpietila-amd.

Test Plan

pytest op_tests/flydsl_tests/test_flydsl_fp8_mqa_logits.py — validated against the torch reference and cross-checked vs the Triton kernel (calc_diff < 1e-3): both fp8 dtype combos

Test Result

335/335 pass.

Kernel speedup (gfx942, H=64, D=128). FlyDSL wins in both operand configs:

Mixed q=e4m3fn / k=e4m3fnuz (current vLLM behavior):

S_q × S_kv	Triton	FlyDSL	speedup
4096 × 4096	0.47 ms	0.21 ms	2.2×
8192 × 8192	1.48 ms	0.70 ms	2.1×
128 × 32768	0.30 ms	0.12 ms	2.4×
671 × 131072	3.14 ms	1.94 ms	1.6×

Mono fnuz/fnuz (the production config; gfx942, H=64, D=128):

S_q × S_kv	Triton	FlyDSL	speedup
4096 × 4096	0.34 ms	0.22 ms	1.5×
8192 × 8192	1.15 ms	0.74 ms	1.6×
128 × 32768	0.29 ms	0.12 ms	2.5×
671 × 131072	3.14 ms	1.78 ms	1.8×

End-to-end DeepSeek-V4 Flash TTFT (TP4, MI300-class gfx942, OSL=27, concurrency=1; Triton gfx942 indexer vs FlyDSL, leveraging PR 46730 on vLLM which fixes the mixed dtype):

ISL	Triton	FlyDSL	speedup	TTFT reduction
32K	2099 ms	1981 ms	1.06×	−5.7%
128K	11337 ms	9401 ms	1.21×	−17.0%

The E2E gain grows with context length (the indexer logits cost scales with N), which is where prefill latency hurts most.

Note on the mixed FN/FNUZ path (re: @valarLip)

You were right to push on the mixed dtype usage. The mixed q=fn/k=fnuz inputs come from vLLM, not aiter. Currently DSv4-Flash's fused_indexer_q.py quantizes Q to e4m3fn while the gfx942 KV cache is fnuz. The proper fix is at the source: vllm-project/vllm#46730 makes Q fnuz so both operands match. The Triton FN/FNUZ fix here is an optional bridge for the current vLLM behavior and can be dropped if you'd rather rely on the source fix.

[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 3913 --add-label <label>

[valarLip]

why there is a mixed FN/FNUZ path?

[akii96]

@valarLip

The mixed types come from vllm. On gfx942 the KV cache is fnuz, but DSv4 Flash quantizes Q to fn (fused_indexer_q.py only emits fn), so the kernel gets q=fn / k=fnuz.

We can fix that at the source and give the Q-quant an fnuz path like the KV cache already has (I can look at it later today!)

Even with both fnuz (no conversion on either side), flydsl is faster (compared below) so my thinking is that we can still have this PR

fnuz/fnuz	triton	flydsl
4096×4096	0.34 ms	0.22 ms (1.5×)
8192×8192	1.15 ms	0.74 ms (1.6×)
128×32768	0.29 ms	0.12 ms (2.5×)
671×131072	3.14 ms	1.78 ms (1.8×)

I am still polishing this work up but do you think the flydsl kernel improvements for the mono dtype are upstream worthy or do you have other concerns ? 😅

[akii96]

@valarLip I've updated the PR description significantly to give a fuller picture of this work, could you have a second look and advise how you'd like to scope this PR?

• flydsl kernel only: drop the Triton FN/FNUZ fix and rely on the vLLM source-side fix (Q→fnuz) to eliminate the mixed dtype, or
• flydsl kernel + keep the mixed-dtype bridge: the earlier source-side attempts upstream never merged and went stale, and the current one ([ROCm][Perf][Bugfix] DSv4 indexer: use platform FP8 dtype (fnuz) for Q-quant on gfx942 vllm-project/vllm#46730) isn't guaranteed to land, so carrying mixed-dtype support in aiter short-term may be a necessary evil

Thankful for any insight you can provide here!

[vpietila-amd]

We should add

if arch == "gfx942":

such that we don't accidentally convert FP8-FN on other architectures that support FP8.

[akii96]

fixed now in 99bd4d3! Thanks

[vpietila-amd]

We need to check

if arch == "gfx942":

to ensure that we don't convert to FP8-FNUZ on other architectures.

This kernel is optimized for gfx942 and while it should run also on other FP8 supported architectures, the performance will probably be suboptimal as the it is not using e.g. async loads that are available in new architectures. At minimum, we should test that the current implementation works on gfx950.

[akii96]

fixed now in 99bd4d3! Thanks. The decision to claify for others is to gate this kernel for now on gfx942

[vpietila-amd]

We should add here named constants

MFMA_M = 16
MFMA_N = 16
MFMA_K = 32

And use those constants in the rest of the code. Now there too many magic number 16 and 32 floating around.

[akii96]

Yeah made sense! and now addressed in d0d58f0. Added MFMA_M/MFMA_N/MFMA_K and used them for the tile dims everywhere, and renamed the lane_div_16/lane_mod_16 vars to lane_div_N/lane_mod_N to match

[vpietila-amd]

Reviewed the implementation and the new FlyDSL kernel looks good to me.

[akii96]

@valarLip we gathered more results with the new FlyDSL kernel introduced in this PR

Outside of DSv4 Flash we also tried replacing the Triton kernel with the FlyDSL counterpart for GLM-5.2-FP8:

Max Conc	TTFT (ms)
2	20389.85 → 15427.41 (−24%)
4	32335.73 → 25395.11 (−21%)
6	33781.63 → 27584.69 (−18%)
8	39559.43 → 30936.41 (−22%)
GeoMean	30637.75 → 24046.14 (−21%)

TTFT reduction is consistent (18–24%) across concurrency levels

Hoping to get your review on this and get this merged as we noticed similar improvements across DS3.2/GLM5.1 as well (happy to provide more results on request!)

[valarLip]

please refact ut follow this one https://github.com/ROCm/aiter/tree/main/.claude/skills/aiter-op-test

[akii96]

Refactored in 9e8797b!