Reenable MNNVL backend for FlashInfer allreduce fusion (#23402)

test/registered/lora/test_lora_qwen3_30b_a3b_instruct_2507_logprob_diff.py

@@ -82,6 +82,7 @@ class TestLoRAQwen3_30B_A3B_Instruct_2507_LogprobDiff(CustomTestCase):
             lora_paths={"my_lora": adapter_path},
             lora_backend=LORA_BACKEND,
             attention_backend="flashinfer",
+            flashinfer_allreduce_fusion_backend="trtllm",
             moe_runner_backend=MOE_RUNNER_BACKEND,
             experts_shared_outer_loras=EXPERTS_SHARED_OUTER_LORAS,
             prefill_attention_backend=PREFILL_ATTENTION_BACKEND,

test/registered/models_e2e/test_deepseek_v32_fp4_mtp_tp.py

@@ -46,7 +46,15 @@ class TestDeepseekV32FP4TPSpec(GSM8KMixin, DefaultServerBase):
     gsm8k_accept_length_thres = 2.7
 
     def test_z_bs_1_speed(self):
-        args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
+        args = BenchArgs(
+            port=int(self.base_url.split(":")[-1]),
+            max_new_tokens=2048,
+            prompt=(
+                "Human: Think carefully before answering. Build a fully functional FastAPI todo server. "
+                "Start with a short design plan, then output the complete Python code, then show how to run it "
+                "and test three endpoints.\n\nAssistant:"
+            ),
+        )
         acc_length, speed = send_one_prompt(args)
 
         print(f"{acc_length=:.2f} {speed=:.2f}")

test/registered/unit/layers/test_flashinfer_comm_fusion.py

@@ -0,0 +1,181 @@
+import types
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.layers import flashinfer_comm_fusion as fusion
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-h100")
+register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-b200")
+register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-gb300")
+
+
+class _FakeWorkspace:
+    def __init__(self, backend, world_size):
+        self.backend = backend
+        self.world_size = world_size
+
+    def is_buffer_size_sufficient(self, **_kwargs):
+        return True
+
+
+class _FakeFlashInferComm:
+    class AllReduceFusionPattern:
+        kARResidualRMSNorm = object()
+
+    def __init__(self):
+        self.calls = []
+
+    def create_allreduce_fusion_workspace(self, **kwargs):
+        self.calls.append(kwargs)
+        return _FakeWorkspace(kwargs["backend"], kwargs["world_size"])
+
+    def allreduce_fusion(
+        self,
+        *,
+        input,
+        workspace,
+        residual_out,
+        norm_out,
+        residual_in,
+        rms_gamma,
+        rms_eps,
+        **_kwargs,
+    ):
+        allreduced = input * workspace.world_size
+        expected_residual = allreduced + residual_in
+        variance = expected_residual.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)
+        expected_norm = (
+            expected_residual.to(torch.float32)
+            * torch.rsqrt(variance + rms_eps)
+            * rms_gamma.to(torch.float32)
+        ).to(input.dtype)
+        residual_out.copy_(expected_residual)
+        norm_out.copy_(expected_norm)
+
+
+def _torch_allreduce_residual_rmsnorm_baseline(
+    input_tensor, residual, weight, world_size, eps
+):
+    allreduced = input_tensor * world_size
+    residual_out = allreduced + residual
+    variance = residual_out.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)
+    norm_out = (
+        residual_out.to(torch.float32)
+        * torch.rsqrt(variance + eps)
+        * weight.to(torch.float32)
+    ).to(input_tensor.dtype)
+    return norm_out, residual_out
+
+
+class TestFlashInferCommFusion(unittest.TestCase):
+    def test_auto_backend_resolves_by_arch(self):
+        single_node = types.SimpleNamespace(
+            flashinfer_allreduce_fusion_backend="auto", nnodes=1
+        )
+        multi_node = types.SimpleNamespace(
+            flashinfer_allreduce_fusion_backend="auto", nnodes=2
+        )
+
+        # Blackwell: mnnvl regardless of node count.
+        with patch.object(fusion, "is_sm100_supported", return_value=True):
+            self.assertEqual(
+                fusion.resolve_flashinfer_allreduce_fusion_backend(single_node), "mnnvl"
+            )
+            self.assertEqual(
+                fusion.resolve_flashinfer_allreduce_fusion_backend(multi_node), "mnnvl"
+            )
+
+        # SM90: mnnvl on single-node, trtllm fallback on multi-node.
+        with (
+            patch.object(fusion, "is_sm100_supported", return_value=False),
+            patch.object(fusion, "is_sm90_supported", return_value=True),
+        ):
+            self.assertEqual(
+                fusion.resolve_flashinfer_allreduce_fusion_backend(single_node), "mnnvl"
+            )
+            self.assertEqual(
+                fusion.resolve_flashinfer_allreduce_fusion_backend(multi_node), "trtllm"
+            )
+
+        # Pre-SM90: trtllm everywhere.
+        with (
+            patch.object(fusion, "is_sm100_supported", return_value=False),
+            patch.object(fusion, "is_sm90_supported", return_value=False),
+        ):
+            self.assertEqual(
+                fusion.resolve_flashinfer_allreduce_fusion_backend(single_node),
+                "trtllm",
+            )
+
+    def test_allreduce_fusion_backends_match_torch_baseline(self):
+        fake_comm = _FakeFlashInferComm()
+        original_comm = fusion._flashinfer_comm
+        original_create = fusion._create_allreduce_fusion_workspace
+        original_manager = fusion._attn_tp_workspace_manager
+        original_unavailable = fusion._flashinfer_allreduce_unavailable
+        try:
+            fusion._flashinfer_comm = fake_comm
+            fusion._create_allreduce_fusion_workspace = (
+                fake_comm.create_allreduce_fusion_workspace
+            )
+            fusion._flashinfer_allreduce_unavailable = False
+
+            for backend in ("trtllm", "mnnvl"):
+                with self.subTest(backend=backend):
+                    world_size = 4
+                    manager = fusion.FlashInferWorkspaceManager()
+                    manager.workspace = _FakeWorkspace(backend, world_size)
+                    manager.initialized = True
+                    fusion._attn_tp_workspace_manager = manager
+                    if not torch.cuda.is_available():
+                        self.skipTest("FlashInfer allreduce custom op is CUDA-only")
+                    device = torch.device("cuda")
+                    torch.manual_seed(0)
+                    input_tensor = torch.randn(4, 8, dtype=torch.float32, device=device)
+                    residual = torch.randn(4, 8, dtype=torch.float32, device=device)
+                    weight = torch.randn(8, dtype=torch.float32, device=device)
+                    eps = 1e-6
+
+                    expected_norm, expected_residual = (
+                        _torch_allreduce_residual_rmsnorm_baseline(
+                            input_tensor, residual, weight, world_size, eps
+                        )
+                    )
+
+                    with (
+                        patch.object(
+                            fusion, "is_flashinfer_available", return_value=True
+                        ),
+                        patch.object(
+                            fusion,
+                            "get_attn_tensor_model_parallel_world_size",
+                            return_value=world_size,
+                        ),
+                        patch.object(
+                            fusion, "ensure_workspace_initialized", return_value=True
+                        ),
+                    ):
+                        norm_out, residual_out = (
+                            fusion.flashinfer_allreduce_residual_rmsnorm(
+                                input_tensor=input_tensor,
+                                residual=residual,
+                                weight=weight,
+                                eps=eps,
+                                max_token_num=8,
+                            )
+                        )
+
+                    torch.testing.assert_close(norm_out, expected_norm)
+                    torch.testing.assert_close(residual_out, expected_residual)
+        finally:
+            fusion._flashinfer_comm = original_comm
+            fusion._create_allreduce_fusion_workspace = original_create
+            fusion._attn_tp_workspace_manager = original_manager
+            fusion._flashinfer_allreduce_unavailable = original_unavailable
+
+
+if __name__ == "__main__":
+    unittest.main()