From 11153736681561f79f9a3624f43eab1041371400 Mon Sep 17 00:00:00 2001 From: Zhiyao Jiang Date: Sat, 20 Jun 2026 04:25:01 -0400 Subject: [PATCH] [AMD] Fix garbled unquantized Qwen3-30B-A3B output on ROCm/aiter where the aiter CK fused-MoE falls back to Triton with pre-shuffled weights (#28244) Co-authored-by: Xinyu Jiang --- .../sglang/srt/layers/quantization/unquant.py | 53 ++-- .../mi35x/test_qwen3_moe_eval_mi35x.py | 230 ++++++++++++++++++ 2 files changed, 265 insertions(+), 18 deletions(-) create mode 100644 test/registered/amd/accuracy/mi35x/test_qwen3_moe_eval_mi35x.py diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index a30872231..2d1628d1d 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -238,8 +238,13 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): set_weight_attrs(w2_weight_bias, extra_weight_attrs) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - _should_use_aiter_moe = _use_aiter and ( - get_moe_runner_backend().is_auto() or get_moe_runner_backend().is_aiter() + _should_use_aiter_moe = ( + _use_aiter + and ( + get_moe_runner_backend().is_auto() + or get_moe_runner_backend().is_aiter() + ) + and self._aiter_ck_moe_supported(layer) ) if _should_use_aiter_moe: copy_or_rebind_param( @@ -392,6 +397,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): param.data = param.data.reshape(expected_shape) + def _aiter_ck_moe_supported(self, layer) -> bool: + # aiter CK fused-MoE requires intermediate_size_per_partition to be 128-aligned + # (GemmSpec=Default; otherwise CK raises "not support this GEMM problem"). + return layer.intermediate_size_per_partition % 128 == 0 + def create_moe_runner( self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): @@ -410,7 +420,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): backend = MoeRunnerBackend.TRITON self.runner = MoeRunner(backend, moe_runner_config) - # Separate runner so CK-shape errors fall back to self.runner on every call. + # aiter CK fused-MoE only supports 128-aligned shapes; otherwise use triton. self._aiter_runner: Optional[MoeRunner] = None if ( _use_aiter @@ -420,7 +430,22 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): ) and get_moe_a2a_backend().supports_aiter() ): - self._aiter_runner = MoeRunner(MoeRunnerBackend.AITER, moe_runner_config) + if self._aiter_ck_moe_supported(layer): + self._aiter_runner = MoeRunner( + MoeRunnerBackend.AITER, moe_runner_config + ) + elif get_moe_runner_backend().is_aiter(): + raise ValueError( + "moe_runner_backend=aiter is not supported for " + f"intermediate_size_per_partition={layer.intermediate_size_per_partition}; " + "use --moe-runner-backend triton." + ) + else: + logger.warning_once( + "aiter CK fused-MoE does not support " + f"intermediate_size_per_partition={layer.intermediate_size_per_partition}; " + "using triton MoE runner." + ) @property def load_up_proj_weight_first(self) -> bool: @@ -522,20 +547,12 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): AiterMoeQuantInfo, ) - try: - quant_info = AiterMoeQuantInfo( - w13_weight=layer.w13_weight, - w2_weight=layer.w2_weight, - expert_mask=layer.dispatcher.expert_mask_gpu, - ) - return self._aiter_runner.run(dispatch_output, quant_info) - except RuntimeError as e: - # AITER CK fused_moe may not support all GEMM dimensions - # (e.g. Gemma4 MoE with 128 experts x 704 intermediate size) - logger.warning_once( - f"AITER CK fused_moe failed ({e}), " - "falling back to Triton MoE runner." - ) + quant_info = AiterMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + expert_mask=layer.dispatcher.expert_mask_gpu, + ) + return self._aiter_runner.run(dispatch_output, quant_info) quant_info = TritonMoeQuantInfo( w13_weight=layer.w13_weight, diff --git a/test/registered/amd/accuracy/mi35x/test_qwen3_moe_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_qwen3_moe_eval_mi35x.py new file mode 100644 index 000000000..a3195eb9b --- /dev/null +++ b/test/registered/amd/accuracy/mi35x/test_qwen3_moe_eval_mi35x.py @@ -0,0 +1,230 @@ +"""MI35x Qwen3 MoE (unquantized) GSM8K Completion Evaluation Test (8-GPU) + +Tests unquantized (bf16) Qwen3 MoE (Qwen/Qwen3-30B-A3B) using a few-shot GSM8K +completion benchmark on MI35x with aiter enabled (SGLANG_USE_AITER=1). + +Registry: nightly-amd-8-gpu-mi35x suite +""" + +import ast +import os +import re +import time +import unittest +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import numpy as np + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + is_in_ci, + popen_launch_server, + write_github_step_summary, +) +from sglang.utils import download_and_cache_file, read_jsonl + +# Register for AMD CI - MI35x Qwen3 MoE accuracy test (~15 min) +register_amd_ci(est_time=1200, suite="nightly-amd-8-gpu-mi35x", nightly=True) + +INVALID = -9999999 + + +@dataclass +class ModelConfig: + """Configuration for a model to test.""" + + model_path: str + tp_size: int = 8 + accuracy_threshold: float = 0.50 + other_args: Optional[List[str]] = None + env_vars: Optional[dict] = None + timeout: Optional[int] = None + + def __post_init__(self): + if self.other_args is None: + self.other_args = [] + if self.env_vars is None: + self.env_vars = {} + + +# Unquantized (bf16) Qwen3 MoE on MI35x with aiter enabled. +MI35X_QWEN3_MOE_MODELS = [ + ModelConfig( + model_path="Qwen/Qwen3-30B-A3B", + tp_size=8, + accuracy_threshold=0.75, + other_args=[ + "--max-running-requests", + "128", + "--mem-fraction-static", + "0.8", + "--trust-remote-code", + ], + # ServerArgs routes the unquantized Qwen3 MoE runner to triton on ROCm with + # aiter (aiter CK can't run the TP-sharded intermediate, e.g. 768 // 8). + env_vars={"SGLANG_USE_AITER": "1"}, + ), +] + + +def get_one_example(lines, i, include_answer): + """Format a single GSM8K example.""" + ret = "Question: " + lines[i]["question"] + "\nAnswer:" + if include_answer: + ret += " " + lines[i]["answer"] + return ret + + +def get_few_shot_examples(lines, k): + """Get k few-shot examples for prompting.""" + ret = "" + for i in range(k): + ret += get_one_example(lines, i, True) + "\n\n" + return ret + + +def get_answer_value(answer_str): + """Extract numerical answer from response.""" + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +def run_gsm8k_benchmark( + base_url: str, + num_questions: int = 200, + num_shots: int = 5, + parallel: int = 64, +) -> Tuple[float, float, float]: + """Run GSM8K few-shot completion benchmark.""" + import sglang as sgl + from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint + + url = "https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl" + data_path = download_and_cache_file(url) + lines = list(read_jsonl(data_path)) + + few_shot_examples = get_few_shot_examples(lines, num_shots) + + questions = [] + labels = [] + for i in range(len(lines[:num_questions])): + questions.append(get_one_example(lines, i, False)) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + arguments = [{"question": q} for q in questions] + + @sgl.function + def few_shot_gsm8k(s, question): + s += few_shot_examples + question + s += sgl.gen( + "answer", max_tokens=512, stop=["Question", "Assistant:", "<|separator|>"] + ) + + backend = RuntimeEndpoint(base_url) + sgl.set_default_backend(backend) + + tic = time.perf_counter() + states = few_shot_gsm8k.run_batch( + arguments, temperature=0, num_threads=parallel, progress_bar=True + ) + latency = time.perf_counter() - tic + + preds = [get_answer_value(states[i]["answer"]) for i in range(len(states))] + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + + return float(acc), float(invalid), float(latency) + + +class TestQwen3MoeEvalMI35x(unittest.TestCase): + """Unquantized Qwen3 MoE GSM8K Completion Evaluation Test for AMD MI35x.""" + + @classmethod + def setUpClass(cls): + cls.models = MI35X_QWEN3_MOE_MODELS + cls.base_url = DEFAULT_URL_FOR_TEST + cls.num_questions = int(os.environ.get("GSM8K_NUM_QUESTIONS", "200")) + + def test_qwen3_moe_accuracy(self): + """Test unquantized Qwen3 MoE with GSM8K completion benchmark.""" + all_results = [] + summary = "### Qwen3 MoE Models (MI35x)\n\n" + summary += "| Model | TP | Accuracy | Threshold | Status |\n" + summary += "| ----- | -- | -------- | --------- | ------ |\n" + + for config in self.models: + with self.subTest(model=config.model_path): + print(f"\n{'='*60}") + print(f"Testing: {config.model_path}") + print(f"{'='*60}") + + env = os.environ.copy() + for key, value in config.env_vars.items(): + env[key] = value + + other_args = list(config.other_args) + other_args.extend(["--tp", str(config.tp_size)]) + timeout = config.timeout or DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH + + try: + process = popen_launch_server( + model=config.model_path, + base_url=self.base_url, + timeout=timeout, + other_args=other_args, + env=env, + ) + + try: + acc, invalid, latency = run_gsm8k_benchmark( + self.base_url, num_questions=self.num_questions + ) + passed = acc >= config.accuracy_threshold + status = "✅ PASS" if passed else "❌ FAIL" + print( + f" accuracy={acc:.3f} threshold={config.accuracy_threshold} {status}" + ) + + all_results.append( + { + "model": config.model_path, + "accuracy": acc, + "passed": passed, + } + ) + summary += f"| {config.model_path} | {config.tp_size} | {acc:.3f} | {config.accuracy_threshold} | {status} |\n" + + finally: + kill_process_tree(process.pid) + + except Exception as e: + summary += f"| {config.model_path} | {config.tp_size} | N/A | {config.accuracy_threshold} | ❌ ERROR |\n" + all_results.append( + { + "model": config.model_path, + "accuracy": None, + "passed": False, + "error": str(e), + } + ) + + if is_in_ci(): + write_github_step_summary(summary) + + failed = [r for r in all_results if not r["passed"]] + if failed: + raise AssertionError(f"Failed models: {[r['model'] for r in failed]}") + + +if __name__ == "__main__": + unittest.main()