[Bugfix] Fix illegal memory access (#12758)

python/sglang/srt/layers/flashinfer_comm_fusion.py

@@ -96,7 +96,7 @@ _workspace_manager = FlashInferWorkspaceManager()
 
 
 def ensure_workspace_initialized(
-    max_token_num: int = 2048, hidden_dim: int = 4096, use_fp32_lamport: bool = False
+    max_token_num: int = 16384, hidden_dim: int = 4096, use_fp32_lamport: bool = False
 ):
     """Ensure workspace is initialized"""
     if not is_flashinfer_available() or _flashinfer_comm is None:
@@ -128,7 +128,7 @@ def flashinfer_allreduce_residual_rmsnorm(
     residual: torch.Tensor,
     weight: torch.Tensor,
     eps: float = 1e-6,
-    max_token_num: int = 2048,
+    max_token_num: int = 16384,
     use_oneshot: Optional[bool] = None,
     trigger_completion_at_end: bool = False,
     fp32_acc: bool = False,
@@ -160,6 +160,15 @@ def flashinfer_allreduce_residual_rmsnorm(
         logger.debug("Single GPU, no need for allreduce fusion")
         return None, None
 
+    if input_tensor.shape[0] > max_token_num:
+        logger.debug(
+            "Input token(%d) is greater than max_token_num(%d), "
+            "falling back to standard implementation",
+            input_tensor.shape[0],
+            max_token_num,
+        )
+        return None, None
+
     if not ensure_workspace_initialized(
         max_token_num=max_token_num,
         hidden_dim=input_tensor.shape[-1],
@@ -205,7 +214,7 @@ def fake_flashinfer_allreduce_residual_rmsnorm(
     residual: torch.Tensor,
     weight: torch.Tensor,
     eps: float = 1e-6,
-    max_token_num: int = 2048,
+    max_token_num: int = 16384,
     use_oneshot: Optional[bool] = None,
     trigger_completion_at_end: bool = False,
     fp32_acc: bool = False,

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

@@ -1052,7 +1052,15 @@ class FlashInferFP4MoE(FusedMoE):
         with use_symmetric_memory(
             get_tp_group(), disabled=not is_allocation_symmetric()
         ):
-            symm_output = torch.empty_like(hidden_states)
+            num_tokens = hs_fp4.shape[0]
+            hidden_size = (
+                hs_fp4.shape[-1] * 2
+                if hs_fp4.dtype == torch.uint8
+                else hs_fp4.shape[-1]
+            )
+            symm_output = torch.empty(
+                num_tokens, hidden_size, dtype=torch.bfloat16, device=hs_fp4.device
+            )
         result = trtllm_fp4_block_scale_moe(
             routing_logits=router_logits,
             routing_bias=topk_config.correction_bias.to(hidden_states.dtype),

python/sglang/srt/layers/quantization/mxfp4.py

@@ -644,7 +644,15 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             with use_symmetric_memory(
                 get_tp_group(), disabled=not is_allocation_symmetric()
             ):
-                symm_output = torch.empty_like(x)
+                num_tokens = x_quant.shape[0]
+                hidden_size = (
+                    x_quant.shape[-1] * 2
+                    if x_quant.dtype == torch.uint8
+                    else x_quant.shape[-1]
+                )
+                symm_output = torch.empty(
+                    num_tokens, hidden_size, dtype=torch.bfloat16, device=x_quant.device
+                )
             trtllm_gen_output = trtllm_fp4_block_scale_moe(
                 router_logits.to(torch.bfloat16),
                 None,  # routing_bias