diff --git a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py index 2f5ac95f2..53e1b086a 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py @@ -21,7 +21,11 @@ from sglang.srt.layers.moe.token_dispatcher import ( from sglang.srt.layers.moe.token_dispatcher.flashinfer_utils import ( TorchDistributedCommBackend, ) -from sglang.srt.layers.moe.topk import StandardTopKOutput, TopKOutput +from sglang.srt.layers.moe.topk import ( + StandardTopKOutput, + TopKOutput, + TopKOutputChecker, +) from sglang.srt.layers.moe.utils import get_moe_runner_backend from sglang.srt.runtime_context import get_server_args from sglang.srt.speculative.spec_info import SpeculativeAlgorithm @@ -176,7 +180,12 @@ class FlashinferDispatcher(BaseDispatcher): output_dtype = hidden_states.dtype x = hidden_states x_sf = None - topk_ids = topk_output.topk_ids + # FlashInfer dispatch requires materialized top-k IDs and weights. + if TopKOutputChecker.format_is_bypassed(topk_output): + topk_output = topk_output.to_standard() + # FlashInfer MoeAlltoAll's expert-ID ABI is int32. This dispatcher is + # only selected for moe_a2a_backend="flashinfer". + topk_ids = topk_output.topk_ids.to(torch.int32) topk_weights = topk_output.topk_weights global_scale = self.quant_config.get("input_global_scale", None) diff --git a/test/registered/ep/test_flashinfer_a2a.py b/test/registered/ep/test_flashinfer_a2a.py index 7a2a078aa..3ffcdc295 100644 --- a/test/registered/ep/test_flashinfer_a2a.py +++ b/test/registered/ep/test_flashinfer_a2a.py @@ -1,6 +1,8 @@ import unittest from types import SimpleNamespace +import requests + from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.run_eval import run_eval @@ -13,8 +15,12 @@ from sglang.test.test_utils import ( register_cuda_ci(est_time=500, stage="base-c", runner_config="4-gpu-gb300") DEEPSEEK_V3_FP4_MODEL = "nvidia/DeepSeek-V3-0324-FP4" +GLM52_NVFP4_MODEL = "nvidia/GLM-5.2-NVFP4" QWEN3_FP8_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8" SERVER_LAUNCH_TIMEOUT = 1000 +FLASHINFER_A2A_ENV = { + "SGLANG_FLASHINFER_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "4096", +} class TestFlashinferA2ATrtllmRoutedFP4(CustomTestCase): @@ -68,6 +74,65 @@ class TestFlashinferA2ATrtllmRoutedFP4(CustomTestCase): self.assertGreater(metrics["score"], 0.90) +class TestFlashinferA2ACutedslStaticFP4(CustomTestCase): + """flashinfer A2A + static EP + flashinfer_cutedsl with GLM-5.2 NVFP4.""" + + @classmethod + def setUpClass(cls): + cls.model = GLM52_NVFP4_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=[ + "--tp", + "4", + "--ep", + "4", + "--dp", + "4", + "--enable-dp-attention", + "--moe-a2a-backend", + "flashinfer", + "--moe-runner-backend", + "flashinfer_cutedsl", + "--ep-dispatch-algorithm", + "static", + "--quantization", + "modelopt_fp4", + "--trust-remote-code", + "--chunked-prefill-size", + "4096", + "--mem-fraction-static", + "0.78", + "--cuda-graph-max-bs-decode", + "16", + "--disable-flashinfer-autotune", + "--model-loader-extra-config", + '{"enable_multithread_load": true}', + ], + env=FLASHINFER_A2A_ENV, + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + def test_generate(self): + response = requests.post( + self.base_url + "/generate", + json={ + "text": "What is 2 + 2?", + "sampling_params": {"temperature": 0, "max_new_tokens": 8}, + }, + timeout=120, + ) + self.assertEqual(response.status_code, 200, response.text) + self.assertTrue(response.json()["text"]) + + class TestFlashinferA2ATrtllmRoutedFP8(CustomTestCase): """flashinfer A2A + flashinfer_trtllm_routed with fp8 (Qwen3-Next)."""