[AMD-yanfeiwang]

Motivation

On ROCm, linear_bf16_fp32 (bf16 x bf16 -> fp32) runs torch.mm(out_dtype=fp32), i.e. hipBLASLt. For the DeepSeek-V4 compressor wkv_gate (K=7168, N=1024 HCA / N=2560 fused CSA) that takes 11-16 us per call at decode sizes (M<=32), and a DSv4-Pro decode step makes 61 of these calls. ROCm/aiter#6005 makes the FlyDSL HGEMM fp32-accurate for fp32 output and adds tuned FlyDSL rows for these shapes (6-14 us).

Modifications

• kernels/ops/gemm/bf16_fp32.py: on ROCm, _linear_bf16_fp32_cublas looks up aiter's tuned bf16 -> fp32 config for (M, N, K). If it is a FlyDSL row, it calls aiter.tuned_gemm.tgemm.mm(x, w, otype=fp32). Otherwise it uses torch.mm(out_dtype=fp32) as before.
• Depends on [FlyDSL] HGEMM: keep the C shuffle in fp32 for fp32 output; DSv4 wkv_gate fp32 rows ROCm/aiter#6005 for the fp32-accurate kernel and the DSv4 rows (M<=32 only; larger M stays on torch). Without it, none of the current callers' shapes has a FlyDSL fp32 row (the only ones on aiter main are N=6144, K=2048/3072), so everything stays on torch.mm. The same holds when aiter is not installed.

Accuracy Tests

• DSv4 wkv_gate shapes (N=1024/2560, K=7168, M=1..112) vs fp64, 1x MI355X, image lmsysorg/sglang-rocm:v0.5.20-rocm724-mi35x-20260927:
  • stock aiter: never routed, output bitwise equal to torch.mm;
  • aiter + [FlyDSL] HGEMM: keep the C shuffle in fp32 for fp32 output; DSv4 wkv_gate fp32 rows ROCm/aiter#6005: routed for M<=32, max rel err 2e-7..4e-7 (torch.mm 1e-7..5e-7); M>32 not routed, bitwise equal to torch.mm.
• test/registered/kernels/ops/gemm/test_router_gemv_hip.py: 2/2 pass, with and without [FlyDSL] HGEMM: keep the C shuffle in fp32 for fp32 output; DSv4 wkv_gate fp32 rows ROCm/aiter#6005.
• DSv4-Pro TP8, DSpark (6 drafts), gsm8k 1319 questions, per boot: before 0.942 / 0.943, after 0.944 / 0.945.

Speed Tests and Profiling

Kernel, wkv_gate through linear_bf16_fp32, 1x MI355X, CUDA graph of 64 calls with 8 rotating weights, us per call:

M	N=1024 torch.mm	N=1024 this PR	N=2560 torch.mm	N=2560 this PR
1	14.28	5.80	16.43	9.33
8	14.07	6.27	11.64	10.02
16	11.32	6.65	12.34	11.16
32	11.08	7.79	16.32	14.10
48	8.98	8.97	16.14	16.03

N=1024 M=1/8 come from the aiter tuning sweep (same timing method); the other cells time the sglang route.

E2E, DeepSeek-V4-Pro fp8, TP8 MI355X, image lmsysorg/sglang-rocm:v0.5.20-rocm724-mi35x-20260927, 8k in / 1k out, step time (ms). Before = image + our DSv4 decode stack: #41383 (with ROCm/aiter#5882), #38919 + #38920 (SGLANG_OPT_DSV4_C4_COMPRESSOR_GEMM_FUSION=1), the DSpark verify ring write and the MoE small-batch sort. After = before + ROCm/aiter#6005 + this PR. Paired, 2 boots x 3 reps:

cell	before	after	delta mean [min, max]
plain c8	12.22	11.88	-2.71% [-0.365, -0.304 ms]
plain c32	15.79	15.62	-1.07% [-0.195, -0.137 ms]
DSpark b6 c2	15.01	14.83	-1.14% [-0.200, -0.127 ms]
DSpark b6 c4	17.60	17.39	-1.17% [-0.258, -0.125 ms]

All 24 pairs are faster. DSpark c16/c64 (M>32) are unchanged by construction.

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 #36731607608
Latest PR Test (Extra): ❌ Run #36731607096
Latest PR Test (AMD ROCm 10): ❌ Run #36731607442

[Duyi-Wang]

Coverd by #41931