diff --git a/python/sglang/srt/layers/flashinfer_comm_fusion.py b/python/sglang/srt/layers/flashinfer_comm_fusion.py index 81280db0a..19efc8d44 100644 --- a/python/sglang/srt/layers/flashinfer_comm_fusion.py +++ b/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, diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 03e3cdf34..bc0e1061e 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/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), diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 23dbda5e3..bbb2daa7d 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/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