[AMD-yanfeiwang]

Motivation

The gfx950 FlyDSL HGEMM (gemm_a16w16_gfx950) rounds the fp32 accumulator to bf16/fp16 before the LDS C shuffle. With out_dtype=fp32, the output (and every split-K partial) therefore has bf16-level error, about 2e-3 relative. That makes FlyDSL unusable for fp32-output GEMMs such as the DeepSeek-V4 compressor wkv_gate (bf16 x bf16 -> fp32, K=7168). hipBLASLt takes 11-16 us per call there, and a decode step runs 61 of these calls.

Technical Details

• Kernel (gemm_a16w16_gfx950.py, ft and hti): the shared C tile and the r2g vectors use the output dtype. fp32 output is no longer rounded to 16 bits, and the LDS size check uses the output bytes. bf16/fp16 output is unchanged.
• Cache key: the fp32-output kernel name becomes hgemm_bf16_fp32c_.... FlyDSL's disk-cache key hashes the @flyc.jit launcher and its helpers but not the @flyc.kernel bodies (checked on flydsl 0.3.2 and 0.3.4.1). Without the rename, a cache filled by the old kernel, such as a precompiled image cache, would keep serving the lossy binary.
• Tuned rows (dsv4_bf16_tuned_gemm.csv): 12 gfx950 fp32-output FlyDSL rows for DSv4 wkv_gate, N=1024 (HCA) and N=2560 (fused CSA, [ROCm][DSV4] Add opt-in C4 compressor GEMM fusion sgl-project/sglang#38920), K=7168, M in {1,2,4,8,16,32}. For M>32, FlyDSL loses to or ties hipBLASLt, so those shapes stay on torch.
• Test: op_tests/test_flydsl_hgemm.py -o fp32 runs the three DSv4 shapes. Its tolerance is 1e-2 atol / 1e-4 rtol, which any 16-bit rounding fails.
• This PR complements [DSV4] Add fp32-output untuned GEMM shapes for indexer kv_score #3836, which lists untuned fp32-output kv_score shapes. Tuning those shapes needs this fix first: the old lossy kernel still passes the tuner's 1e-2 check.
• The consumer is a follow-up sglang PR. It routes linear_bf16_fp32 to aiter tgemm.mm(otype=fp32) when the tuned row is FlyDSL.

Test Plan

• python op_tests/test_flydsl_hgemm.py -o fp32 and python op_tests/test_flydsl_hgemm.py (default bf16 set). Run on this commit applied to lmsysorg/sglang-rocm:v0.5.20-rocm724-mi35x-20260927 (flydsl 0.3.2, image cache present), 1x MI355X.

Test Result

• fp32 output: 3/3 passed, err 0. bf16 output: 34/34 passed, err 0.
• Max relative error of fp32 output vs torch.mm: 2e-3 before, 1e-6..2e-6 after (the same as torch.mm(out_dtype=fp32)).

Kernel. wkv_gate bf16 x bf16 -> fp32, K=7168, 1x MI355X. Time per call in us: CUDA-graph replay, 8 rotating weights. The tuned FlyDSL config is compared with torch.mm (hipBLASLt):

M	N=1024 torch	N=1024 FlyDSL	N=2560 torch	N=2560 FlyDSL
1	14.28	5.80	16.41	8.77
2	14.01	5.81	16.30	8.96
4	14.09	6.00	16.27	9.02
8	14.07	6.27	11.43	9.41
16	11.22	6.63	12.28	10.49
32	10.97	7.77	16.30	13.13

E2E. DSv4-Pro, TP8 on MI355X, 8k input / 1k output, T_step in ms (lower is better). Paired runs, 2 boots x 3 reps.

• Baseline: the 20260927 image plus our other pending DSv4 changes ([AMD][DSV4] fp8 unified_kv decode/verify via aiter FlyDSL v4 nm decode sgl-project/sglang#41383, [ROCm][DSV4] Support row-strided inputs in C4 v2 compression sgl-project/sglang#38919 and [ROCm][DSV4] Add opt-in C4 compressor GEMM fusion sgl-project/sglang#38920 with the env on, and others).
• This PR: the baseline plus this PR and the sglang route that uses these rows.

cell	baseline	this PR	delta
plain, c8	12.22	11.88	-2.7%
plain, c32	15.79	15.62	-1.1%
DSpark (6 drafts), c2	15.01	14.83	-1.1%
DSpark (6 drafts), c4	17.60	17.39	-1.2%

All 24 pairs are faster. gsm8k (1319 questions) accuracy per boot: 0.942 / 0.943 before and 0.944 / 0.945 after. DSpark c16/c64 (M>32) are unchanged by construction.

[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 (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
multigpu	Aiter multi-GPU tests on the 8-GPU runner
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 6005 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.