[AMD] Fix wrong MoE output at small batch with aiter fp8 block-scale (split-K with one K tile per split)

[ntgiang71096]

Motivation

What users see. On the ROCm images lmsysorg/sglang-rocm v0.5.20 and v0.5.21, Qwen3.6-35B-A3B-FP8 at TP1 produces garbage when only one request is running: a few correct tokens, then !!!!. The same bug also causes two other problems:

• a re-sent cached prompt comes back wrong when exactly one token has to be recomputed (radix cache, extra_buffer);
• the MTP draft is corrupted under NEXTN, which slows speculative decoding.

Reproduce (MI350X, lmsysorg/sglang-rocm:v0.5.21-rocm724-mi35x-20261002; SGLANG_USE_AITER=1 is the image default):

python -m sglang.launch_server --model-path Qwen/Qwen3.6-35B-A3B-FP8 --tp-size 1
curl -s localhost:30000/v1/chat/completions -H 'Content-Type: application/json' -d '{
  "messages": [{"role": "user", "content": "Explain in detail how a hash table handles collisions, with examples."}],
  "max_tokens": 64, "temperature": 0, "chat_template_kwargs": {"enable_thinking": false}}'
# Before: "# Hash Table Collision Handling\n\n## What Is!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!..."
# After:  "# Hash Table Collision Resolution: A Detailed Explanation\n\n## 1. What Is a Collision?\n\nA **hash table** maps keys to ..."
# Sending the same request again (radix-cache hit): before "#!!!!!!!!!!!!!!!!!!...", after a coherent answer.
# (Verified with an extra --mem-fraction-static 0.4 only because the GPU was shared.)

Cause. For very small batches, aiter splits the stage1 MoE GEMM along K and adds the partial results into one buffer. For Qwen3.6-35B at M=1, K=2048 is split 8 ways (256 each), so every split owns exactly one K tile. In that configuration CK skips zeroing the buffer, but the splits still accumulate into it. The buffer comes from torch.empty, so the result lands on stale memory.

aiter guarded this case in ROCm/aiter#3997, but ROCm/aiter#4394 dropped the guard. It is tracked in ROCm/aiter#4032. The v0.5.16 image (aiter d9e5ef7) still has the guard.

Modifications

At AiterRunnerCore init, sglang wraps aiter.fused_moe.ck_moe_stage1 once:

• For fp8 block-scale (per_1x128) calls whose split leaves a single K tile per split, the wrapper lowers the split to the largest count that keeps ≥ 2 K tiles per split. For example, K=2048 goes from 8 splits to 4. If no such count exists, the call runs unsplit.
• All other calls are passed through unchanged.

Why in sglang, and why it is safe:

• No released image has the aiter fix, so v0.5.20/v0.5.21 users hit this now.
• Narrow: only calls that currently return wrong results change. sglang already patches aiter the same way (apply_aiter_small_moe_sort_patch, MoE: small-batch sorting path with fused mxfp8 quantisation #36559).
• Clear removal condition: a TODO drops it once the minimum supported aiter carries the guard again. If aiter fixes it first, the wrapper is harmless.
• Not AITER_KSPLIT=1: that disables split-K for every call and slows small-batch decode.
• Lowers the split instead of disabling split-K (what My SGLang doesn't print resources usages #3997 does). Running unsplit leaves most of the GPU idle (32 workgroups on 256 CUs) and is slower. Lowering the split is correct and slightly faster than the original: 27.7 µs vs 28.8 µs per MoE call.
• Supersedes [AMD][Bugfix] Mamba radix cache: avoid the degenerate single-token (M=1) recompute that corrupts the first token #29633. That PR only avoided one path to M=1 (the radix cache), so it could not fix decode, the spec draft, or TP > 1.
• Interaction with [AMD] gfx950 small-batch MoE: expert-count gate for the small sort and FP8 block-scale small-M kernel #41982: on gfx950, small-M FP8 calls with the Qwen3.5-397B-A17B shapes now go to sglang's smallm_moe_gfx950 kernel and never reach ck_moe_stage1. This change covers the remaining shapes, and gfx942.

Accuracy Tests

Qwen3.6-35B-A3B-FP8, TP1, MI350X, v0.5.21-rocm724-mi35x-20261002, --attention-backend triton --mamba-radix-cache-strategy extra_buffer. Serving logprobs are compared against logprobs recomputed by a clean prefill, following sglang's kl_test_utils.

Case	Before	After
Single-request decode, no spec	NaN in 2032/2048 tokens	KL 1.6e-3 (sglang's threshold for hybrid models: 2e-3)
Re-sent cached prompt, 1 token to recompute	NaN in 8/8	KL 4.6e-3, same range as an unaffected control (2-token recompute)
Batch ≥ 2 (not touched by this change)	KL 1.8e-3	KL 1.2e-3, within noise

Details: equivalence on unchanged paths, kernel-level check, affected models and TP

Equivalence on unchanged paths

• Byte-identical output is not a usable criterion. aiter's fused MoE is not deterministic run to run, even without this change: two runs of the unpatched server on the same inputs give identical tokens for only 6/16 sequences.
• Equivalence is therefore measured against that noise. On paths this change does not touch, KL between unpatched and patched is 6.6e-4. The unpatched-vs-unpatched noise is 5.4e-4 to 7.6e-4.
• The prefill-logprob |Δ| quantiles across servers match unpatched-vs-unpatched reruns (p50 0.007–0.010 vs 0.007–0.008).

Kernel level

• fused_moe is compared against an fp32 reference, including with the caching allocator pre-filled with NaN.
• Before: wrong exactly in the single-tile-split cases. After: correct in every case.

Affected models and TP (MI350X/MI355X, default heuristic):

• Qwen3.6-35B-A3B-FP8:
  • TP1 at M=1
  • TP2 at M≤2
  • TP4/TP8 at M≤4
• Qwen3.5-397B-A17B-FP8: TP8 at M=1. On gfx950 this is now covered by [AMD] gfx950 small-batch MoE: expert-count gate for the small sort and FP8 block-scale small-M kernel #41982.
• Qwen3.8-Flash-Next-FP8:
  • TP4 at M=1
  • TP8 at M≤2
• Dense models are not affected.
• 397B and Flash-Next were verified at kernel level only.

Speed Tests and Profiling

Qwen3.6-35B-A3B-FP8, TP1, MI350X, single request, same image and flags as above:

	Before	After
NEXTN, 10,000-token generation (3 long-form prompts)	360 tok/s	426 tok/s (+18%)
NEXTN, 256-token answers	TPOT 3.23 ms	2.13 ms
No spec	3.40 ms (produces garbage, not a valid baseline)	3.44 ms

• The NEXTN speedup comes from the MTP draft being correct again: accept length goes up ×1.13–1.46, depending on content.
• With batch ≥ 2 at TP1 this change does not act, so speed is unchanged.

GPU time per MoE call at M=1

Shape	Before (wrong)	Split-K disabled	This PR
35B TP1	28.8 µs	31.5 µs	27.7 µs
397B TP8	32.3 µs	40.5 µs	26.5 µs

Tests

Added test/registered/amd/test_aiter_moe_splitk_kbatch_guard.py:

• It runs fused_moe at M=1/2/4/8 and compares against an fp32 reference.
• It fails at M=1 without the guard and passes on MI350X.

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. (not applicable)
• 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 #37916671528
Latest PR Test (Extra): ❌ Run #37916671150
Latest PR Test (AMD ROCm 10): ❌ Run #37916671514

[ntgiang71096]

@hnyls2002 Following up on your note in #29633: the M=1 issue does still reproduce on the unified-cache tree. The root cause turned out to be in aiter's split-K path rather than the radix cache: ck_moe_stage1 with a single K tile per split, whose guard from ROCm/aiter#3997 was dropped in ROCm/aiter#4394. This PR fixes it in the aiter MoE runner. That also covers plain single-request decode and the MTP draft, not just the 1-token extend.

Could you take a look and run /tag-and-rerun-ci when you have a moment? Thanks!