From 757b9eaf1c4990b3ed34aa69e7688fd804f4a9ad Mon Sep 17 00:00:00 2001 From: ntgiang71096 Date: Fri, 9 Oct 2026 10:10:10 +0000 Subject: [PATCH] [AMD] Keep aiter block-fp8 MoE split-K stage1 at >= 2 K tiles per split aiter's ck_moe_stage1 accumulates split-K partials into a torch.empty fp32 buffer that CK only zeroes when every split owns at least two K tiles (KBatch = K // (splitk * 256) > 1). For small M the default heuristic picks a split with KBatch == 1 (e.g. Qwen3.6-35B-A3B-FP8 at TP1, M=1), so the MoE output is garbage: single-request decode emits "!!!!", 1-token radix extends are wrong, and the MTP draft is corrupted. aiter guarded this in #3997 and lost the guard in #4394 (ROCm/aiter#4032). Wrap ck_moe_stage1 once from AiterRunnerCore and lower such splits to the largest count that keeps >= 2 K tiles per split (0 if none). All other calls are passed through unchanged. --- .../sglang/srt/layers/moe/moe_runner/aiter.py | 46 +++++++ .../amd/test_aiter_moe_splitk_kbatch_guard.py | 118 ++++++++++++++++++ 2 files changed, 164 insertions(+) create mode 100644 test/registered/amd/test_aiter_moe_splitk_kbatch_guard.py diff --git a/python/sglang/srt/layers/moe/moe_runner/aiter.py b/python/sglang/srt/layers/moe/moe_runner/aiter.py index 18dc7c773452..f4d1979a89b4 100644 --- a/python/sglang/srt/layers/moe/moe_runner/aiter.py +++ b/python/sglang/srt/layers/moe/moe_runner/aiter.py @@ -46,6 +46,51 @@ _MOE_SORTING_DISPATCH_POLICY = envs.SGLANG_AITER_MOE_SORTING_DISPATCH_POLICY.get() +# aiter's block-fp8 split-K stage1 (ck_moe_stage1) accumulates the splits into +# a torch.empty fp32 buffer, which CK zeroes only when every split owns at +# least two K tiles (KBatch = K // (splitk * 256) > 1). For tiny M, aiter's +# default heuristic can still choose a split with KBatch == 1; e.g. M=1 for +# Qwen3.5/3.6-FP8 at TP1, M<=4 at TP4. The stage1 output is then garbage. +# aiter #3997 guarded this by running such calls unsplit and #4394 dropped the +# guard. Lowering the split to the largest count that keeps two K tiles per +# split is also correct, and faster than running unsplit. +# TODO(ntgiang71096): remove once the minimum supported aiter carries the +# guard again (ROCm/aiter#4032). +_AITER_SPLITK_K_PER_BLOCK = 256 + + +def _aiter_safe_splitk(k: int, splitk: int) -> int: + """Largest split <= splitk whose splits each own >= 2 K tiles; 0 runs unsplit.""" + for split in range(splitk, 1, -1): + tiles, rem = divmod(k, split * _AITER_SPLITK_K_PER_BLOCK) + if rem == 0 and tiles >= 2: + return split + return 0 + + +@functools.cache +def apply_aiter_splitk_kbatch_guard() -> None: + """Keep aiter's block-fp8 split-K stage1 at >= 2 K tiles per split.""" + try: + import aiter.fused_moe as fm + from aiter import QuantType + except ImportError: + return + orig_stage1 = fm.ck_moe_stage1 + + @functools.wraps(orig_stage1) + def ck_moe_stage1_guarded(*args, **kwargs): + splitk = kwargs.get("splitk", 1) + if splitk > 1 and kwargs.get("quant_type") == QuantType.per_1x128: + hidden_states = args[0] if args else kwargs["hidden_states"] + k = hidden_states.shape[-1] + if k // (splitk * _AITER_SPLITK_K_PER_BLOCK) < 2: + kwargs["splitk"] = _aiter_safe_splitk(k, splitk) + return orig_stage1(*args, **kwargs) + + fm.ck_moe_stage1 = ck_moe_stage1_guarded + logger.info("aiter split-K KBatch guard applied") + class AiterQuantType(str, Enum): NONE = "No" @@ -255,6 +300,7 @@ def __init__(self, *args, **kwargs): ) apply_aiter_small_moe_sort_patch() + apply_aiter_splitk_kbatch_guard() def run( self, diff --git a/test/registered/amd/test_aiter_moe_splitk_kbatch_guard.py b/test/registered/amd/test_aiter_moe_splitk_kbatch_guard.py new file mode 100644 index 000000000000..5b63e55f1d5b --- /dev/null +++ b/test/registered/amd/test_aiter_moe_splitk_kbatch_guard.py @@ -0,0 +1,118 @@ +"""aiter block-fp8 fused MoE at tiny M vs a float32 torch reference. + +For K=2048, topk=8, inter=512 (Qwen3.5/3.6-FP8 MoE at TP1), aiter's default +heuristic splits stage1 eight ways at M=1, so each split owns a single 256-wide +K tile. The split-K accumulator is then never zeroed and the output is +garbage. The sglang aiter runner installs a guard that runs such calls unsplit +(ROCm/aiter#4032). +""" + +import unittest + +import torch +import torch.nn.functional as F + +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.test_utils import CustomTestCase + +register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd") + + +def _rocm_with_aiter(): + if not (torch.version.hip and torch.cuda.is_available()): + return False + try: + import aiter # noqa: F401 + except ImportError: + return False + return True + + +def _block_quant_128(w): + # [E, N, K] -> fp8 [E, N, K] with one fp32 scale per 128x128 block + from aiter import dtypes, pertoken_quant + + e, n, k = w.shape + blocks = w.view(e, n // 128, 128, k // 128, 128).permute(0, 1, 3, 2, 4) + q, s = pertoken_quant( + blocks.reshape(e, -1, 128 * 128).contiguous(), quant_dtype=dtypes.fp8 + ) + q = q.view(e, n // 128, k // 128, 128, 128).permute(0, 1, 3, 2, 4) + return q.reshape(e, n, k).contiguous(), s.view(e, n // 128, k // 128) + + +def _dequant_128(q, s): + return q.float() * s.repeat_interleave(128, 1).repeat_interleave(128, 2) + + +@unittest.skipUnless(_rocm_with_aiter(), "ROCm with aiter only") +class TestAiterMoeSplitKKBatchGuard(CustomTestCase): + E, INTER, DIM, TOPK = 16, 512, 2048, 8 + + @classmethod + def setUpClass(cls): + from aiter.ops.shuffle import shuffle_weight + + from sglang.srt.layers.moe.moe_runner.aiter import ( + apply_aiter_splitk_kbatch_guard, + ) + + apply_aiter_splitk_kbatch_guard() + torch.manual_seed(0) + dev = torch.device("cuda", 0) + w1 = torch.randn(cls.E, 2 * cls.INTER, cls.DIM, device=dev) * 0.05 + w2 = torch.randn(cls.E, cls.DIM, cls.INTER, device=dev) * 0.05 + w1_q, cls.w1_scale = _block_quant_128(w1.to(torch.bfloat16)) + w2_q, cls.w2_scale = _block_quant_128(w2.to(torch.bfloat16)) + cls.w1_ref = _dequant_128(w1_q, cls.w1_scale) + cls.w2_ref = _dequant_128(w2_q, cls.w2_scale) + cls.w1 = shuffle_weight(w1_q, (16, 16)) + cls.w2 = shuffle_weight(w2_q, (16, 16)) + # same layout flag the fp8 MoE method sets on its shuffled weights + cls.w1.is_shuffled = True + cls.w2.is_shuffled = True + cls.dev = dev + + def _reference(self, x, topk_weights, topk_ids): + from aiter import dtypes, pertoken_quant + + # aiter quantizes activations per 1x128 group before stage1 + x_q, x_s = pertoken_quant(x.view(x.shape[0], -1, 128), quant_dtype=dtypes.fp8) + x_deq = (x_q.float() * x_s).view(x.shape[0], -1) + out = torch.zeros(x.shape[0], self.DIM, device=self.dev) + for t in range(x.shape[0]): + for slot in range(self.TOPK): + e = int(topk_ids[t, slot]) + h = x_deq[t] @ self.w1_ref[e].t() + act = F.silu(h[: self.INTER]) * h[self.INTER :] + out[t] += float(topk_weights[t, slot]) * (act @ self.w2_ref[e].t()) + return out + + def test_small_m_matches_reference(self): + from aiter import ActivationType, QuantType + from aiter.fused_moe import fused_moe, fused_topk + + for m in (1, 2, 4, 8): + with self.subTest(m=m): + x = torch.randn(m, self.DIM, device=self.dev).to(torch.bfloat16) + logits = torch.randn(m, self.E, device=self.dev).to(torch.bfloat16) + topk_weights, topk_ids = fused_topk(x, logits, self.TOPK, True) + out = fused_moe( + x, + self.w1, + self.w2, + topk_weights, + topk_ids, + quant_type=QuantType.per_128x128, + activation=ActivationType.Silu, + w1_scale=self.w1_scale, + w2_scale=self.w2_scale, + ).float() + ref = self._reference(x, topk_weights, topk_ids) + self.assertTrue(torch.isfinite(out).all()) + cos = F.cosine_similarity(out.flatten(), ref.flatten(), dim=0) + self.assertGreater(cos.item(), 0.99) + + +if __name__ == "__main__": + unittest.main()