minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[NVIDIA] Add flashinfer MNNVL backend for allreduce only (#30700)

Browse Source
This commit is contained in:
Shu Wang
2026-08-11 16:46:46 -07:00
committed by GitHub
parent 983dfd6a9a
commit c7c03ec53b
7 changed files with 476 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/layers/test_flashinfer_comm_fusion.py
+206 -6
View File
@@ -1,3 +1,4 @@
import contextlib
import types
import unittest
from unittest.mock import patch
@@ -7,6 +8,7 @@ import torch
from sglang.srt.layers import flashinfer_comm_fusion as fusion
from sglang.srt.runtime_context import get_parallel
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
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")
@@ -24,6 +26,7 @@ class _FakeWorkspace:
class _FakeFlashInferComm:
class AllReduceFusionPattern:
kAllReduce = object()
kARResidualRMSNorm = object()
def __init__(self):
@@ -38,13 +41,25 @@ class _FakeFlashInferComm:
*,
input,
workspace,
residual_out,
norm_out,
residual_in,
rms_gamma,
rms_eps,
pattern,
output=None,
residual_out=None,
norm_out=None,
residual_in=None,
rms_gamma=None,
rms_eps=None,
**_kwargs,
):
if pattern is self.AllReduceFusionPattern.kAllReduce:
allreduced = input * workspace.world_size
if output is None:
return allreduced
output.copy_(allreduced)
return output
if pattern is not self.AllReduceFusionPattern.kARResidualRMSNorm:
raise ValueError(f"Unexpected pattern: {pattern}")
allreduced = input * workspace.world_size
expected_residual = allreduced + residual_in
variance = expected_residual.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)
@@ -71,7 +86,7 @@ def _torch_allreduce_residual_rmsnorm_baseline(
return norm_out, residual_out
class TestFlashInferCommFusion(unittest.TestCase):
class TestFlashInferCommFusion(CustomTestCase):
def test_auto_backend_resolves_by_arch(self):
single_node = types.SimpleNamespace(
flashinfer_allreduce_fusion_backend="auto", nnodes=1
@@ -240,5 +255,190 @@ class TestFlashInferCommFusion(unittest.TestCase):
fusion._flashinfer_allreduce_unavailable = original_unavailable
_GROUP_KEY = ("device_group", "cpu_group")
_OTHER_GROUP_KEY = ("other_device_group", "other_cpu_group")
class TestFlashInferAllReduceOnly(CustomTestCase):
def _make_manager(self, world_size, group_key=_GROUP_KEY):
manager = fusion.FlashInferWorkspaceManager()
manager.workspace = _FakeWorkspace(None, world_size)
manager.initialized = True
manager.world_size = world_size
manager.group = group_key
manager.max_token_num = 2048
manager.hidden_dim = 4096
manager.dtype = torch.float32
return manager
@contextlib.contextmanager
def _patched_attn_workspace(self, manager):
from sglang.srt.runtime_context import get_resources
buffers = get_resources().buffers
manager_key = "flashinfer_fusion_attn_tp_workspace"
original_manager = buffers.get(manager_key)
original_comm = fusion._flashinfer_comm
original_unavailable = fusion._flashinfer_allreduce_unavailable
buffers[manager_key] = manager
fusion._flashinfer_comm = _FakeFlashInferComm()
fusion._flashinfer_allreduce_unavailable = False
try:
yield
finally:
fusion._flashinfer_comm = original_comm
fusion._flashinfer_allreduce_unavailable = original_unavailable
if original_manager is None:
buffers.pop(manager_key, None)
else:
buffers[manager_key] = original_manager
def _can_use(self, input_, world_size=4, group_key=_GROUP_KEY):
return fusion.can_use_flashinfer_allreduce(
input_,
use_attn_tp_group=True,
expected_world_size=world_size,
expected_group_key=group_key,
)
def test_allreduce_output_equals_input_times_world_size(self):
if not torch.cuda.is_available():
self.skipTest("CUDA required for flashinfer custom op")
world_size = 4
with self._patched_attn_workspace(self._make_manager(world_size)):
input_ = torch.randn(8, 16, dtype=torch.bfloat16, device="cuda")
expected = input_ * world_size
with get_parallel().override(attn_tp_size=world_size):
self.assertTrue(self._can_use(input_, world_size=world_size))
result = fusion.flashinfer_allreduce(input_, use_attn_tp_group=True)
torch.testing.assert_close(result, expected)
def test_shape_guard_rejects_non_2d(self):
with self._patched_attn_workspace(self._make_manager(4)):
self.assertFalse(self._can_use(torch.randn(16)))
self.assertFalse(self._can_use(torch.randn(2, 8, 16)))
def test_shape_guard_rejects_non_contiguous(self):
with self._patched_attn_workspace(self._make_manager(4)):
non_contiguous = torch.randn(16, 8).t()
self.assertFalse(non_contiguous.is_contiguous())
self.assertFalse(self._can_use(non_contiguous))
def test_rejects_when_unavailable(self):
original_unavailable = fusion._flashinfer_allreduce_unavailable
try:
fusion._flashinfer_allreduce_unavailable = True
self.assertFalse(self._can_use(torch.randn(8, 16)))
finally:
fusion._flashinfer_allreduce_unavailable = original_unavailable
def test_rejects_when_workspace_uninitialized(self):
with self._patched_attn_workspace(fusion.FlashInferWorkspaceManager()):
with get_parallel().override(attn_tp_size=4):
self.assertFalse(self._can_use(torch.randn(8, 16)))
def test_rejects_when_workspace_group_differs(self):
"""A workspace rendezvoused on other peers must not be reused.
Under hybrid EP+TP (e.g. tp=4, ep=2) the MoE-TP and MoE-EP groups have
the same world size but pair different ranks, so a workspace built for
one silently reduces across the wrong peers when used by the other --
wrong output rather than a crash.
"""
with self._patched_attn_workspace(self._make_manager(2)):
self.assertFalse(
self._can_use(
torch.randn(8, 16), world_size=2, group_key=_OTHER_GROUP_KEY
)
)
def test_rejects_when_workspace_world_size_differs(self):
with self._patched_attn_workspace(self._make_manager(4)):
self.assertFalse(self._can_use(torch.randn(8, 16), world_size=2))
def test_rejects_when_token_num_exceeds_workspace_capacity(self):
"""Under Dynamo the capacity check replaces is_buffer_size_sufficient().
_FakeWorkspace.is_buffer_size_sufficient() always says yes, so this only
passes if the compiling branch consults the manager's own allocation.
"""
manager = self._make_manager(4)
manager.max_token_num = 8
with self._patched_attn_workspace(manager):
with patch.object(torch.compiler, "is_compiling", return_value=True):
self.assertTrue(self._can_use(torch.randn(8, 16)))
self.assertFalse(self._can_use(torch.randn(9, 16)))
def test_rejects_when_hidden_dim_exceeds_workspace_capacity(self):
manager = self._make_manager(4)
manager.hidden_dim = 16
with self._patched_attn_workspace(manager):
with patch.object(torch.compiler, "is_compiling", return_value=True):
self.assertTrue(self._can_use(torch.randn(8, 16)))
self.assertFalse(self._can_use(torch.randn(8, 17)))
def test_rejects_when_dtype_mismatches_workspace(self):
manager = self._make_manager(4)
manager.dtype = torch.bfloat16
with self._patched_attn_workspace(manager):
with patch.object(torch.compiler, "is_compiling", return_value=True):
self.assertTrue(self._can_use(torch.randn(8, 16, dtype=torch.bfloat16)))
self.assertFalse(self._can_use(torch.randn(8, 16, dtype=torch.float32)))
class _FakeGroupCoordinator:
def __init__(self, world_size):
self.world_size = world_size
self._fi_workspace_hint = None
class TestTagGroupsForFlashInferAllReduceOnly(CustomTestCase):
"""The MoE workspace rendezvouses on the EP group when moe_ep_size > 1 and
on the MoE-TP group otherwise, so only that one group may be tagged."""
def _tag(self, *, attn_tp, moe_ep, moe_tp):
from sglang.srt.distributed import parallel_state as ps
with patch.object(ps, "_ENABLE_FLASHINFER_ALLREDUCE_ONLY", True), patch.object(
ps, "_ATTN_TP", attn_tp
), patch.object(ps, "_MOE_EP", moe_ep), patch.object(ps, "_MOE_TP", moe_tp):
ps._tag_groups_for_flashinfer_allreduce_only()
def test_hybrid_ep_tp_tags_only_the_ep_group(self):
attn_tp = _FakeGroupCoordinator(4)
moe_ep = _FakeGroupCoordinator(2)
moe_tp = _FakeGroupCoordinator(2)
self._tag(attn_tp=attn_tp, moe_ep=moe_ep, moe_tp=moe_tp)
self.assertEqual(attn_tp._fi_workspace_hint, "attn_tp")
self.assertEqual(moe_ep._fi_workspace_hint, "moe")
self.assertIsNone(moe_tp._fi_workspace_hint)
def test_pure_moe_tp_tags_only_the_moe_tp_group(self):
attn_tp = _FakeGroupCoordinator(4)
moe_ep = _FakeGroupCoordinator(1)
moe_tp = _FakeGroupCoordinator(4)
self._tag(attn_tp=attn_tp, moe_ep=moe_ep, moe_tp=moe_tp)
self.assertEqual(moe_tp._fi_workspace_hint, "moe")
self.assertIsNone(moe_ep._fi_workspace_hint)
def test_shared_coordinator_prefers_attn_tp(self):
# tp=4, ep=4: _ATTN_TP is _MOE_EP is _TP. Either workspace spans the
# same peers, but the choice must be deterministic.
shared = _FakeGroupCoordinator(4)
moe_tp = _FakeGroupCoordinator(1)
self._tag(attn_tp=shared, moe_ep=shared, moe_tp=moe_tp)
self.assertEqual(shared._fi_workspace_hint, "attn_tp")
self.assertIsNone(moe_tp._fi_workspace_hint)
if __name__ == "__main__":
unittest.main()
test/registered/unit/layers/test_layer_communicator_fusion_gate.py
+65
View File
@@ -0,0 +1,65 @@
import types
import unittest
from unittest.mock import patch
from sglang.srt.layers import communicator as comm
from sglang.srt.layers.communicator import LayerCommunicator, ScatterMode
from sglang.srt.runtime_context import get_parallel
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
def _fake_communicator():
return types.SimpleNamespace(
_speculative_algo=None,
layer_scatter_modes=types.SimpleNamespace(mlp_mode=ScatterMode.TP_ATTN_FULL),
is_last_layer=False,
_context=types.SimpleNamespace(tp_size=4),
)
class TestFuseMlpAllReduceGate(CustomTestCase):
"""Hybrid EP+TP must not fuse the post-experts all-reduce away.
The fused residual+LN reduces over a single group, but with moe_ep_size > 1
and moe_tp_size > 1 the post-experts reduction spans two disjoint groups
(_MOE_EP then _MOE_TP) and should_skip_post_experts_all_reduce() drops both
once fusion is published. The result is activations reduced over only half
the peers -- wrong output, no crash. Observed as garbage completions on
Qwen3-30B-A3B with --tp-size 4 --ep-size 2.
"""
def _should_fuse(self, *, moe_ep_size, moe_tp_size):
forward_batch = types.SimpleNamespace(
input_ids=types.SimpleNamespace(shape=(8,))
)
with (
patch.object(comm, "is_enable_moe_cp_allgather", return_value=False),
patch.object(comm, "apply_flashinfer_allreduce_fusion", return_value=True),
patch.object(
comm,
"get_attn_tp_context",
return_value=types.SimpleNamespace(input_scattered=False),
),
get_parallel().override(
moe_ep_size=moe_ep_size, moe_tp_size=moe_tp_size, tp_size=4
),
):
return LayerCommunicator.should_fuse_mlp_allreduce_with_next_layer(
_fake_communicator(), forward_batch
)
def test_hybrid_ep_tp_does_not_fuse(self):
self.assertFalse(self._should_fuse(moe_ep_size=2, moe_tp_size=2))
def test_pure_tp_still_fuses(self):
self.assertTrue(self._should_fuse(moe_ep_size=1, moe_tp_size=4))
def test_pure_ep_still_fuses(self):
self.assertTrue(self._should_fuse(moe_ep_size=4, moe_tp_size=1))
if __name__ == "__main__":
unittest.main()