feat(moe): reuse prev-layer output as symm_output for FP4 routed MoE (#25379)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

python/sglang/srt/layers/communicator.py

@@ -305,6 +305,9 @@ class AttnTpContext:
         assert self.attn_inputs_ is not None
         return self.attn_inputs_.fetch_hidden_states()
 
+    def clear_attn_inputs(self) -> None:
+        self.attn_inputs_ = None
+
     @contextmanager
     def maybe_input_scattered(self, forward_batch: ForwardBatch):
         flag = self.use_input_scattered(forward_batch)

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -308,6 +308,15 @@ class FusedMoE(torch.nn.Module):
         self.quant_method.create_moe_runner(self, self.moe_runner_config)
         self.dispatcher = create_moe_dispatcher(self.moe_runner_config)
 
+        if (
+            get_moe_runner_backend().is_flashinfer_trtllm_routed()
+            or get_moe_runner_backend().is_flashinfer_trtllm()
+        ):
+            logging.warning(
+                "Setting inplace to False for FlashInfer TRTLLM MoE backend."
+            )
+            self.moe_runner_config.inplace = False
+
         self.should_fuse_routed_scaling_factor_in_topk = (
             isinstance(self.quant_method, ModelOptNvFp4FusedMoEMethod)
             or (

python/sglang/srt/layers/moe/moe_runner/base.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+import contextvars
 from abc import ABC, abstractmethod
+from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple, TypeGuard
+from typing import TYPE_CHECKING, Any, Callable, Generator, Optional, Tuple, TypeGuard
 
 import torch
 
@@ -26,6 +28,20 @@ if TYPE_CHECKING:
     )
 
 
+_moe_output_buf: contextvars.ContextVar[Optional[torch.Tensor]] = (
+    contextvars.ContextVar("moe_output_buf", default=None)
+)
+
+
+@contextmanager
+def moe_output_buffer_ctx(buf: torch.Tensor) -> Generator[None, None, None]:
+    token = _moe_output_buf.set(buf)
+    try:
+        yield
+    finally:
+        _moe_output_buf.reset(token)
+
+
 @dataclass
 class MoeRunnerConfig:
     # MoE parameters

python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py

@@ -10,6 +10,8 @@ from torch.nn.parameter import Parameter
 # Import to register custom ops for torch.compile compatibility
 from sglang.srt.distributed import get_tp_group
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    is_symmetric_memory_enabled,
+    is_tensor_in_symmetric_mempool,
     use_symmetric_memory,
 )
 from sglang.srt.layers.dp_attention import is_allocation_symmetric
@@ -21,6 +23,7 @@ from sglang.srt.layers.moe.flashinfer_trtllm_moe import (
 from sglang.srt.layers.moe.moe_runner.base import (
     MoeQuantInfo,
     MoeRunnerConfig,
+    _moe_output_buf,
     register_fused_func,
 )
 from sglang.srt.layers.quantization.fp8_kernel import (
@@ -877,14 +880,29 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
     )
     activation_type = get_activation_type(runner_config.activation)
 
-    with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
+    num_tokens = hs_fp4.shape[0]
-        num_tokens = hs_fp4.shape[0]
+    hidden_size = (
-        hidden_size = (
+        hs_fp4.shape[-1] * 2 if hs_fp4.dtype == torch.uint8 else hs_fp4.shape[-1]
-            hs_fp4.shape[-1] * 2 if hs_fp4.dtype == torch.uint8 else hs_fp4.shape[-1]
+    )
-        )
+    _provided = _moe_output_buf.get()
-        symm_output = torch.empty(
+    _symm_required = is_allocation_symmetric()
-            num_tokens, hidden_size, dtype=hidden_states.dtype, device=hs_fp4.device
+    if (
+        _provided is not None
+        and _provided.shape == (num_tokens, hidden_size)
+        and _provided.dtype == hidden_states.dtype
+        and _provided.device == hs_fp4.device
+        and (
+            not _symm_required
+            or not is_symmetric_memory_enabled()
+            or is_tensor_in_symmetric_mempool(_provided)
         )
+    ):
+        symm_output = _provided
+    else:
+        with use_symmetric_memory(get_tp_group(), disabled=not _symm_required):
+            symm_output = torch.empty(
+                num_tokens, hidden_size, dtype=hidden_states.dtype, device=hs_fp4.device
+            )
 
     if use_routed_topk:
         assert TopKOutputChecker.format_is_standard(topk_output)

python/sglang/srt/models/deepseek_v2.py

@@ -1930,6 +1930,7 @@ class DeepseekV2DecoderLayer(nn.Module):
         llama_4_scaling: Optional[torch.Tensor] = None,
         prev_topk_indices: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
+        hidden_states_orig = hidden_states
         hidden_states, residual = self.layer_communicator.prepare_attn(
             hidden_states,
             residual,
@@ -1950,6 +1951,7 @@ class DeepseekV2DecoderLayer(nn.Module):
             hidden_states, topk_indices = hidden_states
         else:
             topk_indices = None
+        get_attn_tp_context().clear_attn_inputs()
 
         hidden_states, residual = self.layer_communicator.prepare_mlp(
             hidden_states, residual, forward_batch
@@ -1969,13 +1971,24 @@ class DeepseekV2DecoderLayer(nn.Module):
         if isinstance(self.mlp, DeepseekV2MLP):
             gemm_output_zero_allocator = None
 
-        hidden_states = self.mlp(
+        if (
-            hidden_states,
+            isinstance(self.mlp, DeepseekV2MoE)
-            forward_batch,
+            and not self.mlp.experts.moe_runner_config.inplace
-            should_allreduce_fusion,
+        ):
-            use_reduce_scatter,
+            from sglang.srt.layers.moe.moe_runner.base import moe_output_buffer_ctx
-            gemm_output_zero_allocator,
+
-        )
+            _mlp_ctx = moe_output_buffer_ctx(hidden_states_orig)
+        else:
+            _mlp_ctx = nullcontext()
+
+        with _mlp_ctx:
+            hidden_states = self.mlp(
+                hidden_states,
+                forward_batch,
+                should_allreduce_fusion,
+                use_reduce_scatter,
+                gemm_output_zero_allocator,
+            )
 
         if not self.nsa_enable_prefill_cp and should_allreduce_fusion:
             hidden_states._sglang_needs_allreduce_fusion = True