[NVIDIA] Support flashinfer Mega Moe (#31470)

Co-authored-by: djns99 <40156487+djns99@users.noreply.github.com>
Co-authored-by: 云挚 <ningyunxiao.nyx@antgroup.com>
Co-authored-by: Yangmin Li <yangminl@nvidia.com>
Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com>

test/manual/ep/test_flashinfer_dispatcher.py

@@ -4,13 +4,15 @@ import torch
 
 from sglang.srt.distributed import init_distributed_environment
 from sglang.srt.distributed.parallel_state import (
+    destroy_distributed_environment,
+    destroy_model_parallel,
     get_tp_group,
     initialize_model_parallel,
 )
 from sglang.srt.layers.dp_attention import set_dp_buffer_len
 from sglang.srt.layers.moe.token_dispatcher.flashinfer import FlashinferDispatcher
 from sglang.srt.layers.moe.utils import initialize_moe_config
-from sglang.srt.runtime_context import publish
+from sglang.srt.runtime_context import get_context, publish
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.test.test_utils import CustomTestCase
 
@@ -21,6 +23,7 @@ class TestFlashinferDispatcher(CustomTestCase):
         server_args = ServerArgs(model_path="dummy")
         server_args.moe_runner_backend = "flashinfer_cutlass"
         server_args.moe_a2a_backend = "flashinfer"
+        cls.server_args = server_args
         set_global_server_args_for_scheduler(server_args)
         publish(server_args, role="scheduler")
         initialize_moe_config()
@@ -41,9 +44,18 @@ class TestFlashinferDispatcher(CustomTestCase):
 
     @classmethod
     def tearDownClass(cls):
-        # Clean up distributed environment
-        if torch.distributed.is_initialized():
-            torch.distributed.destroy_process_group()
+        try:
+            from flashinfer.comm.trtllm_moe_alltoall import MoeAlltoAll
+
+            for workspace in MoeAlltoAll._WORKSPACE_CACHE.values():
+                mnnvl_mem = workspace.get("mnnvl_mem")
+                if mnnvl_mem is not None and "ptr" in vars(mnnvl_mem):
+                    del mnnvl_mem.ptr
+            MoeAlltoAll._WORKSPACE_CACHE.clear()
+        except ImportError:
+            pass
+        destroy_model_parallel()
+        destroy_distributed_environment()
 
     def create_dispatcher(
         self, router_topk=2, num_experts=8, num_local_experts=4, hidden_size=128
@@ -58,8 +70,40 @@ class TestFlashinferDispatcher(CustomTestCase):
             params_dtype=torch.bfloat16,
         )
 
+    def set_dispatch_type(self, dispatch_type):
+        get_context().override(
+            "test_flashinfer_dispatcher",
+            flashinfer_a2a_dispatch_type=dispatch_type,
+        )
+
+    def _zero_moe_a2a_dispatch_payloads(self):
+        # Shared MoeAlltoAll workspaces keep stale recv payloads across tests.
+        # Zero only the payload region so unused-source == 0 asserts stay valid.
+        try:
+            from flashinfer.comm.trtllm_moe_alltoall import (
+                MoeAlltoAll,
+                get_moe_alltoall_module,
+            )
+        except ImportError:
+            return
+
+        module = get_moe_alltoall_module()
+        for ws in MoeAlltoAll._WORKSPACE_CACHE.values():
+            workspace = ws["workspace"]
+            aux = int(
+                module.moe_a2a_get_aux_data_size(
+                    ws["ep_size"],
+                    ws["max_num_tokens"],
+                    ws["eplb_stats_num_experts"],
+                )
+            )
+            aux = ((aux + 127) // 128) * 128
+            if aux < workspace.shape[1]:
+                workspace[:, aux:].zero_()
+
     def test_dispatch_basic(self):
         """Test basic dispatch functionality"""
+        self.set_dispatch_type("bf16")
         num_tokens = 16
         hidden_size = 128
         router_topk = 1  # Single expert per token for simplicity
@@ -143,9 +187,10 @@ class TestFlashinferDispatcher(CustomTestCase):
 
     def test_dispatch_with_empty_tokens(self):
         """Test dispatch when there are no tokens (edge case)"""
+        self.set_dispatch_type("bf16")
         # This tests the dummy token handling
         num_tokens = 16
-        hidden_size = 1
+        hidden_size = 128
         router_topk = 1  # Single expert per token for simplicity
         world_size = torch.distributed.get_world_size()
         rank = torch.distributed.get_rank()
@@ -195,6 +240,9 @@ class TestFlashinferDispatcher(CustomTestCase):
             topk_weights=topk_weights, topk_ids=topk_ids, router_logits=None
         )
 
+        self._zero_moe_a2a_dispatch_payloads()
+        torch.distributed.barrier()
+
         dispatcher = self.create_dispatcher(
             router_topk=router_topk,
             num_experts=num_experts,
@@ -250,6 +298,7 @@ class TestFlashinferDispatcher(CustomTestCase):
 
     def test_dispatch_with_fp4_quantization(self):
         """Test dispatch with FP4 quantization enabled"""
+        self.set_dispatch_type("nvfp4")
         num_tokens = 128
         hidden_size = 128
         router_topk = 1  # Single expert per token for simplicity
@@ -312,6 +361,133 @@ class TestFlashinferDispatcher(CustomTestCase):
         )
         self.assertEqual(dispatch_output.hidden_states_scale.dtype, torch.uint8)
 
+    def test_dispatch_with_mxfp8_quantization(self):
+        """Test dispatch with MXFP8 quantization enabled"""
+        self.set_dispatch_type("mxfp8")
+        num_tokens = 128
+        hidden_size = 128
+        router_topk = 1
+        world_size = torch.distributed.get_world_size()
+        rank = torch.distributed.get_rank()
+        num_experts = world_size
+        num_local_experts = 1
+
+        set_dp_buffer_len(
+            global_dp_buffer_len=num_tokens * world_size,
+            local_dp_buffer_len=num_tokens,
+            dp_max_padding=True,
+            global_num_tokens=None,
+        )
+
+        hidden_states = torch.randn(
+            (num_tokens, hidden_size), dtype=torch.bfloat16, device="cuda"
+        )
+
+        target_rank = (rank + 1) % world_size
+        target_expert = target_rank
+        topk_ids = torch.full(
+            (num_tokens, router_topk), target_expert, dtype=torch.int32, device="cuda"
+        )
+        topk_weights = torch.ones(
+            (num_tokens, router_topk), dtype=torch.float32, device="cuda"
+        )
+
+        from sglang.srt.layers.moe.topk import StandardTopKOutput
+
+        topk_output = StandardTopKOutput(
+            topk_weights=topk_weights, topk_ids=topk_ids, router_logits=None
+        )
+
+        dispatcher = self.create_dispatcher(
+            router_topk=router_topk,
+            num_experts=num_experts,
+            num_local_experts=num_local_experts,
+            hidden_size=hidden_size,
+        )
+        dispatcher.set_quant_config({"input_global_scale": None, "use_mxfp8": True})
+
+        dispatch_output = dispatcher.dispatch(hidden_states, topk_output)
+
+        self.assertEqual(
+            dispatch_output.hidden_states.shape,
+            (num_tokens * world_size, hidden_size),
+        )
+        self.assertEqual(dispatch_output.hidden_states.dtype, torch.float8_e4m3fn)
+        self.assertEqual(dispatch_output.output_dtype, torch.bfloat16)
+
+        self.assertIsNotNone(dispatch_output.hidden_states_scale)
+        self.assertEqual(
+            dispatch_output.hidden_states_scale.shape,
+            (num_tokens * world_size, hidden_size // 32),
+        )
+        self.assertEqual(dispatch_output.hidden_states_scale.dtype, torch.uint8)
+        self.assertEqual(
+            dispatch_output.topk_output.topk_ids.shape,
+            (num_tokens * world_size, router_topk),
+        )
+        self.assertEqual(dispatch_output.topk_output.topk_ids.dtype, torch.int32)
+
+    def test_dispatch_with_mxfp8_quantization_and_empty_rank(self):
+        """All ranks must contribute the same payload dtypes, including empty ranks."""
+        self.set_dispatch_type("mxfp8")
+        num_tokens = 16
+        hidden_size = 128
+        router_topk = 1
+        world_size = torch.distributed.get_world_size()
+        rank = torch.distributed.get_rank()
+        empty_rank = 1
+
+        global_num_tokens = [num_tokens] * world_size
+        global_num_tokens[empty_rank] = 0
+        set_dp_buffer_len(
+            global_dp_buffer_len=num_tokens * world_size,
+            local_dp_buffer_len=num_tokens,
+            dp_max_padding=False,
+            global_num_tokens=global_num_tokens,
+        )
+
+        local_tokens = 0 if rank == empty_rank else num_tokens
+        hidden_states = torch.randn(
+            (local_tokens, hidden_size), dtype=torch.bfloat16, device="cuda"
+        )
+        target_expert = (rank + 1) % world_size
+        topk_ids = torch.full(
+            (local_tokens, router_topk),
+            target_expert,
+            dtype=torch.int32,
+            device="cuda",
+        )
+        topk_weights = torch.ones(
+            (local_tokens, router_topk), dtype=torch.float32, device="cuda"
+        )
+
+        from sglang.srt.layers.moe.topk import StandardTopKOutput
+
+        dispatcher = self.create_dispatcher(
+            router_topk=router_topk,
+            num_experts=world_size,
+            num_local_experts=1,
+            hidden_size=hidden_size,
+        )
+        dispatcher.set_quant_config({"input_global_scale": None, "use_mxfp8": True})
+        self._zero_moe_a2a_dispatch_payloads()
+        torch.distributed.barrier()
+        dispatch_output = dispatcher.dispatch(
+            hidden_states,
+            StandardTopKOutput(
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                router_logits=None,
+            ),
+        )
+
+        self.assertEqual(dispatch_output.hidden_states.dtype, torch.float8_e4m3fn)
+        self.assertEqual(dispatch_output.hidden_states_scale.dtype, torch.uint8)
+        self.assertEqual(
+            dispatch_output.hidden_states.shape,
+            (num_tokens * world_size, hidden_size),
+        )
+
 
 if __name__ == "__main__":
     """