Fix FlashInfer A2A top-k ID dtype (#29929)

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)

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)."""