From 6330d6641ba30c50952531f4148438cc0aacf507 Mon Sep 17 00:00:00 2001
From: alisonshao <54658187+alisonshao@users.noreply.github.com>
Date: Wed, 26 Nov 2025 17:09:02 -0800
Subject: [PATCH] Fix flashinfer cutlass MoE output shape for non-FP4-packed
 inputs (#14028)

---
 python/sglang/srt/layers/quantization/modelopt_quant.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py
index 2809c7a19..8d44584e7 100755
--- a/python/sglang/srt/layers/quantization/modelopt_quant.py
+++ b/python/sglang/srt/layers/quantization/modelopt_quant.py
@@ -1654,11 +1654,15 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
 
             output_dtype = torch.bfloat16
 
+            # If x_sf is not None, x is FP4 packed (half size), so we need * 2
+            # If x_sf is None, x is not packed, so output_col = x.shape[1]
+            output_col = x.shape[1] * 2 if x_sf is not None else x.shape[1]
+
             with use_symmetric_memory(
                 get_tp_group(), disabled=not is_allocation_symmetric()
             ):
                 symm_output = torch.empty(
-                    x.shape[0], x.shape[1] * 2, dtype=output_dtype, device=x.device
+                    x.shape[0], output_col, dtype=output_dtype, device=x.device
                 )
 
             output = flashinfer_cutlass_fused_moe(