[AMD] Small-M FP8 block-scale GEMM with fused activation quant for gfx950 (Qwen3.5)

[chuyeh]

Motivation

On MI355X, Qwen3.5-397B-A17B-FP8 decode at concurrency 1 spends 1.6 to 1.7 ms per cycle in dense FP8 projections. EAGLE with four draft tokens runs a projection at 4 to 32 tokens per rank for concurrency 1 to 8. Each projection currently launches two latency-bound kernels, taking about 4 µs and 9 µs.

This PR fuses them into one HIP kernel for M ≤ 32; larger inputs retain the AITER path.

Modifications

• New python/sglang/kernels/ops/gemm/smallm_fp8_blockscale_gfx950/ (__init__.py,
  smallm_fp8_blockscale_gemm.hip): built with hipcc on first use and launched through ctypes with the helpers of
  smallm_moe_gfx950 ([AMD] Small-M MXFP4 fused-MoE kernel for gfx950 (Qwen) #40204), no build-time or wheel change.

• Each wave loads bf16 activations and quantizes them to fp8 in registers (1x128 groups, row amax via lane shuffles),
  reproducing AITER's per-group scale and fp8 bytes exactly. It streams the AITER (16, 16)-preshuffled FP8 weight
  once with non-temporal loads, runs one v_mfma_scale_f32_16x16x128_f8f6f4 per 128-wide K block, applies
  x_scale * w_scale in fp32, and reduces across waves through LDS to a bf16 output.

• Uses the existing weight layout and weight_scale_inv ([N/128, K/128] fp32), so loading and preprocessing are
  untouched.

• Dispatch table (Qwen3.5-397B-A17B TP4 shapes), first matching M bucket wins:

  (N, K)	Layer	M buckets
  (5120, 4096)	GDN in_proj_qkvz	≤ 32
  (4608, 4096)	attention qkv_proj	≤ 32
  (4096, 2048)	GDN out_proj, attention o_proj	≤ 8, ≤ 16, ≤ 32

• fp8_utils.py: one guarded call in aiter_w8a8_block_fp8_linear, taken on the gfx950 bpreshuffle path (HIP ≥ 7.2)
  with no input_scale (input not pre-quantized) and no bias, when smallm_fp8_bs_supported accepts a bf16 input
  with 16-byte aligned rows, a weight marked is_shuffled, a contiguous [N/128, K/128] weight scale, a shape/M in
  the table, and not under torch.compile. Everything else is unchanged. Other TP sizes have different shapes and
  keep the current path.

• On by default on gfx950; SGLANG_ROCM_SMALLM_FP8_BS=0 turns it off.

• Registered AMD unit test test/registered/amd/test_smallm_fp8_blockscale_gfx950.py: matches AITER quant + GEMM
  (rel 2e-2) at M = 1 / 3 / 4 / 8 / 9 / 16 / 17 / 24 / 32 on all three shapes, quant bit-exact against
  aiter_per1x128_quant, fallback for M = 33 / fp16 / an unaligned row stride / an unshuffled weight / the off
  switch, graph replay == eager. Passes on v0.5.20-rocm720-mi35x-20260927.

Accuracy Tests

All runs use image rocm/sgl-dev:v0.5.20-rocm720-mi35x-20260927, InferenceX recipe flags with real draft acceptance,
sglang.test.few_shot_gsm8k, 1319 questions. --max-running-requests 8 keeps every verify pass at 32 tokens or
fewer, so all dense FP8 projections take the new kernel. "main" is the image's tree (main at a0ba1968b1); baseline
is the same tree with SGLANG_ROCM_SMALLM_FP8_BS=0.

Model	Tree	Baseline	This PR
Qwen/Qwen3.5-397B-A17B-FP8 TP4	main, 2 runs	0.958 / 0.955	0.952 / 0.951
	Mean of 2 runs	0.957	0.952

One run's standard error at 1319 questions is about 0.006.

Benchmarking and Profiling

Kernel microbenchmarks

MI355X, one GPU, CUDA graph replay, weights rotated over 24 copies so they are not cache resident. AITER
aiter_per1x128_quant + gemm_a8w8_blockscale_bpreshuffle (tuned Qwen3.5 config) vs this kernel, in µs:

Tokens	in_proj_qkvz 5120 x 4096	qkv_proj 4608 x 4096	out_proj / o_proj 4096 x 2048
1	13.7 → 6.8 (−51%)	13.7 → 6.6 (−51%)	6.8 → 4.0 (−42%)
4	15.8 → 7.3 (−54%)	15.8 → 7.0 (−56%)	8.8 → 4.1 (−54%)
8	16.0 → 7.7 (−52%)	16.0 → 7.8 (−51%)	8.8 → 4.2 (−53%)
16	16.5 → 9.0 (−46%)	16.6 → 8.3 (−50%)	9.0 → 4.7 (−48%)
32 (cap)	16.8 → 12.8 (−24%)	16.7 → 12.6 (−24%)	10.0 → 7.7 (−23%)

Max relative difference to AITER over all points is 3.1e-3; the quantized activations are bit-exact.

Decode step time

Time for one EAGLE draft + verify cycle at a fixed batch size. MI355X TP4, image
rocm/sgl-dev:v0.5.20-rocm720-mi35x-20260927, InferenceX recipe flags (FP8 KV cache, EAGLE 3 / 1 / 4 with simulated
acceptance 3.39). Mean of two runs with the GPU sets swapped, in ms:

Batch size	Verify tokens	Baseline	This PR	Change
1	4	11.93	11.10	−7%
4	16	16.53	15.86	−4%
8	32	16.29	16.22	~0 (at the cap)

AgentX benchmark

• Hardware: MI355X, TP4.
• Model and recipe: Qwen/Qwen3.5-397B-A17B-FP8 with the InferenceX recipe qwen3.5-fp8-mi355x-sglang-agentic,
  image rocm/sgl-dev:v0.5.20-rocm720-mi35x-20260927.
• Server flags: as in the recipe. FP8 KV cache, EAGLE 3 / 1 / 4 with simulated acceptance 3.39, admission and decode
  CUDA graph batch at 2x concurrency.
• Client: the CI AgentX command, aiperf profile --scenario inferencex-agentx-mvp on
  semianalysis_cc_traces_weka_062126_256k, 3600 s per run.
• Order: one run at a time on the same GPUs, other four GPUs idle.

Concurrency 1:

Metric	Baseline	This PR	Change
ITL p50 / p90 (ms)	3.67 / 4.14	3.45 / 3.91	−6% / −6%
TTFT p50 / p90 (ms)	397 / 947	425 / 943	+7% / ~0
Requests completed in 3600 s	559	575	+3%
Output throughput (tok/s)	159	162	+1%
Interactivity p50 / p90 (tok/s/user)	273 / 286	290 / 304	+6% / +6%

Concurrency 4:

Metric	Baseline	This PR	Change
ITL p50 / p90 (ms)	4.10 / 5.36	3.85 / 5.06	−6% / −6%
TTFT p50 / p90 (ms)	569 / 1111	535 / 1045	−6% / −6%
Requests completed in 3600 s	907	951	+5%
Output throughput (tok/s)	241	246	+2%
Interactivity p50 / p90 (tok/s/user)	244 / 280	260 / 293	+7% / +5%

Concurrency 8:

Metric	Baseline	This PR	Change
ITL p50 / p90 (ms)	5.01 / 5.68	4.74 / 5.50	−5% / −3%
TTFT p50 / p90 (ms)	756 / 1292	722 / 1260	−4% / −2%
Requests completed in 3600 s	1446	1454	+1%
Output throughput (tok/s)	371	371	~0
Interactivity p50 / p90 (tok/s/user)	200 / 273	211 / 288	+6% / +5%

Each EAGLE verify pass sends 4 tokens per running request through every projection, so the kernel covers running
batches of 8 or fewer. The running batch never exceeded 8 in these runs, so every decode step took the kernel. The
kernel does not run in prefill.

Checklist

• Format your code according to the Format code with pre-commit guide.
• Add unit tests.
• Update documentation.
• Provide accuracy and speed benchmark results.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ✅ Run #37760502231
Latest PR Test (Extra): ❌ Run #37760501789
Latest PR Test (AMD ROCm 10): ❌ Run #37760502073

[yichiche]

@chuyeh
Failed checks:

1. AMD Guard
   The diff imports AITER, but the code-path scope is gfx950 bpreshuffle path (_use_aiter_bpreshuffle_gfx95); the kernel is compiled with --offload-arch=gfx950.
   Required: Gate the AITER import with _use_aiter. is_hip alone crashes on an AMD machine that does not have AITER installed.
2. PR body
   Score 6/10. The first sentence gives token counts. MI355X and the 1.6 ms decode cycle come later, after the kernel names, so the machine and the slowdown do not connect on one read.
   Required: Improve PR body (motivation mainly, added a one line action you did in this PR as conclusion)
   Example: On MI355X, Qwen3.5-397B-A17B-FP8 decode at concurrency 1 spends 1.6 to 1.7 ms per cycle in the dense FP8 projections. EAGLE with 4 draft tokens runs each projection at 4 to 32 tokens per rank, from concurrency 1 to 8. Each projection is two latency-bound kernels, about 4 µs then 9 µs.