From 8ae962021dbb5eb8ee408c578e772c13e2c934fe Mon Sep 17 00:00:00 2001 From: Carrie Chen <139071206+xikronz@users.noreply.github.com> Date: Mon, 7 Sep 2026 04:57:07 -0400 Subject: [PATCH] Use fp32 in TRTLLM all reduce buffers (#36143) Co-authored-by: Brayden Zhong --- .../sglang/srt/layers/flashinfer_comm_fusion.py | 13 ++++++++----- .../unit/layers/test_flashinfer_comm_fusion.py | 16 +++++++++++++--- 2 files changed, 21 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/layers/flashinfer_comm_fusion.py b/python/sglang/srt/layers/flashinfer_comm_fusion.py index 00dab3996..3d6e8770f 100644 --- a/python/sglang/srt/layers/flashinfer_comm_fusion.py +++ b/python/sglang/srt/layers/flashinfer_comm_fusion.py @@ -765,7 +765,7 @@ def fake_flashinfer_allreduce_residual_rmsnorm( max_token_num: int = 16384, use_oneshot: Optional[bool] = None, trigger_completion_at_end: bool = False, - fp32_acc: bool = False, + fp32_acc: bool = True, use_attn_tp_group: bool = True, ) -> Tuple[torch.Tensor, torch.Tensor]: residual_out = torch.empty_like(residual) @@ -785,7 +785,7 @@ def flashinfer_allreduce_residual_rmsnorm( max_token_num: int = 2048, use_oneshot: Optional[bool] = None, trigger_completion_at_end: bool = False, - fp32_acc: bool = False, + fp32_acc: bool = True, use_attn_tp_group: bool = True, ) -> Tuple[torch.Tensor, torch.Tensor]: """ @@ -801,7 +801,8 @@ def flashinfer_allreduce_residual_rmsnorm( max_token_num: Maximum token number use_oneshot: Whether to use oneshot mode trigger_completion_at_end: Whether to trigger completion at end - fp32_acc: Whether to use fp32 precision + fp32_acc: Accumulate the allreduce in fp32 (trtllm backend only; the + mnnvl backends always accumulate in fp32) use_attn_tp_group: If True, use attention TP group; otherwise use MoE TP group Returns: @@ -865,8 +866,9 @@ def flashinfer_allreduce_residual_rmsnorm( rms_gamma=weight, rms_eps=eps, use_oneshot=use_oneshot, - fp32_acc=fp32_acc, ) + if workspace_manager.backend == "trtllm": + kwargs["fp32_acc"] = fp32_acc if _flashinfer_allreduce_supports_trigger_completion: kwargs["trigger_completion_at_end"] = trigger_completion_at_end _flashinfer_comm.allreduce_fusion(**kwargs) @@ -976,9 +978,10 @@ def flashinfer_allreduce( workspace=workspace_manager.workspace, pattern=_flashinfer_comm.AllReduceFusionPattern.kAllReduce, launch_with_pdl=True, - fp32_acc=False, output=output, ) + if workspace_manager.backend == "trtllm": + kwargs["fp32_acc"] = True if _flashinfer_allreduce_supports_trigger_completion: kwargs["trigger_completion_at_end"] = False _flashinfer_comm.allreduce_fusion(**kwargs) diff --git a/test/registered/unit/layers/test_flashinfer_comm_fusion.py b/test/registered/unit/layers/test_flashinfer_comm_fusion.py index a82ad8107..ac9189a09 100644 --- a/test/registered/unit/layers/test_flashinfer_comm_fusion.py +++ b/test/registered/unit/layers/test_flashinfer_comm_fusion.py @@ -31,6 +31,7 @@ class _FakeFlashInferComm: def __init__(self): self.calls = [] + self.fusion_kwargs = None def create_allreduce_fusion_workspace(self, **kwargs): self.calls.append(kwargs) @@ -52,6 +53,7 @@ class _FakeFlashInferComm: rms_eps=None, **_kwargs, ): + self.fusion_kwargs = _kwargs if pattern is self.AllReduceFusionPattern.kAllReduce: allreduced = input * workspace.world_size if output is None: @@ -202,6 +204,7 @@ class TestFlashInferCommFusion(CustomTestCase): world_size = 4 manager = fusion.FlashInferWorkspaceManager() manager.workspace = _FakeWorkspace(backend, world_size) + manager.backend = backend manager.initialized = True buffers[manager_key] = manager if not torch.cuda.is_available(): @@ -240,6 +243,10 @@ class TestFlashInferCommFusion(CustomTestCase): torch.testing.assert_close(norm_out, expected_norm) torch.testing.assert_close(residual_out, expected_residual) + self.assertEqual( + fake_comm.fusion_kwargs.get("fp32_acc", False), + backend == "trtllm", + ) finally: fusion._flashinfer_comm = original_comm fusion._create_allreduce_fusion_workspace = original_create @@ -279,10 +286,11 @@ class TestFlashInferAllReduceOnly(CustomTestCase): original_unavailable = fusion._flashinfer_allreduce_unavailable buffers[manager_key] = manager - fusion._flashinfer_comm = _FakeFlashInferComm() + fake_comm = _FakeFlashInferComm() + fusion._flashinfer_comm = fake_comm fusion._flashinfer_allreduce_unavailable = False try: - yield + yield fake_comm finally: fusion._flashinfer_comm = original_comm fusion._flashinfer_allreduce_unavailable = original_unavailable @@ -306,7 +314,7 @@ class TestFlashInferAllReduceOnly(CustomTestCase): manager = self._make_manager(world_size) manager.dtype = torch.bfloat16 manager.use_fp32_lamport = False - with self._patched_attn_workspace(manager): + with self._patched_attn_workspace(manager) as fake_comm: input_ = torch.randn(8, 16, dtype=torch.bfloat16, device="cuda") expected = input_ * world_size @@ -315,6 +323,8 @@ class TestFlashInferAllReduceOnly(CustomTestCase): result = fusion.flashinfer_allreduce(input_, use_attn_tp_group=True) torch.testing.assert_close(result, expected) + # trtllm rounds to the input dtype on every rank without this + self.assertTrue(fake_comm.fusion_kwargs["fp32_acc"]) def test_shape_guard_rejects_non_2d(self): with self._patched_attn_workspace(self._make_manager(4)):