[CPU] Fix model failures on Xeon (#29497)

python/pyproject_cpu.toml

@@ -49,6 +49,7 @@ dependencies = [
   "pybase64",
   "pydantic",
   "python-multipart",
+  "pytest",
   "pyzmq>=25.1.2",
   "requests",
   "scipy",
@@ -69,6 +70,7 @@ dependencies = [
   "uvicorn",
   "uvloop",
   "xgrammar==0.2.1",
+  "zstandard",
 ]
 
 [project.optional-dependencies]
@@ -105,7 +107,6 @@ test = [
   "matplotlib",
   "pandas",
   "peft>=0.18.0",
-  "pytest",
   "sentence_transformers",
 ]
 all = []

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

@@ -857,6 +857,29 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                     layer.w2_weight_bias = Parameter(
                         layer.w2_weight_bias.float(), requires_grad=False
                     )
+                return
+            # Fallback if the TP-sharded layer cannot be AMX-packed
+            from sglang.srt.layers.quantization.mxfp4_tensor import MXFP4QuantizeUtil
+
+            w13_weight = MXFP4QuantizeUtil.dequantize(
+                quantized_data=layer.w13_weight,
+                dtype=torch.bfloat16,
+                scale=layer.w13_weight_scale,
+                block_sizes=[32],
+            )
+            w2_weight = MXFP4QuantizeUtil.dequantize(
+                quantized_data=layer.w2_weight,
+                dtype=torch.bfloat16,
+                scale=layer.w2_weight_scale,
+                block_sizes=[32],
+            )
+            del layer.w13_weight
+            del layer.w2_weight
+            del layer.w13_weight_scale
+            del layer.w2_weight_scale
+            layer.w13_weight = Parameter(w13_weight, requires_grad=False)
+            layer.w2_weight = Parameter(w2_weight, requires_grad=False)
+
             return
         else:
             from triton_kernels.numerics_details.mxfp import upcast_from_mxfp
@@ -1085,32 +1108,42 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
 
         x = dispatch_output.hidden_states
         topk_output = dispatch_output.topk_output
-        if use_intel_amx_backend(layer):
-            from sglang.srt.layers.moe.topk import apply_topk_weights_cpu
+        if _is_cpu:
+            if use_intel_amx_backend(layer):
+                from sglang.srt.layers.moe.topk import apply_topk_weights_cpu
 
-            topk_weights, topk_ids, _ = dispatch_output.topk_output
-            x, topk_weights = apply_topk_weights_cpu(
-                self.moe_runner_config.apply_router_weight_on_input, topk_weights, x
-            )
-            output = torch.ops.sgl_kernel.fused_experts_cpu(
-                x,
-                layer.w13_weight,
-                layer.w2_weight,
-                topk_weights,
-                topk_ids,
-                False,  # inplace See [Note] inplace should be False in fused_experts.
-                CPUQuantMethod.MXFP4,
-                layer.w13_weight_scale,  # w1_scale
-                layer.w2_weight_scale,  # w2_scale
-                None,  # w1_zp
-                None,  # w2_zp
-                None,  # block_size
-                getattr(layer, "w13_weight_bias", None),
-                getattr(layer, "w2_weight_bias", None),
-                layer.moe_runner_config.gemm1_alpha,
-                layer.moe_runner_config.gemm1_clamp_limit,
-                True,  # is_vnni
-            )
+                topk_weights, topk_ids, _ = dispatch_output.topk_output
+                x, topk_weights = apply_topk_weights_cpu(
+                    self.moe_runner_config.apply_router_weight_on_input, topk_weights, x
+                )
+                output = torch.ops.sgl_kernel.fused_experts_cpu(
+                    x,
+                    layer.w13_weight,
+                    layer.w2_weight,
+                    topk_weights,
+                    topk_ids,
+                    False,  # inplace See [Note] inplace should be False in fused_experts.
+                    CPUQuantMethod.MXFP4,
+                    layer.w13_weight_scale,  # w1_scale
+                    layer.w2_weight_scale,  # w2_scale
+                    None,  # w1_zp
+                    None,  # w2_zp
+                    None,  # block_size
+                    getattr(layer, "w13_weight_bias", None),
+                    getattr(layer, "w2_weight_bias", None),
+                    layer.moe_runner_config.gemm1_alpha,
+                    layer.moe_runner_config.gemm1_clamp_limit,
+                    True,  # is_vnni
+                )
+            else:
+                from sglang.srt.layers.moe.fused_moe_native import moe_forward_native
+
+                output = moe_forward_native(
+                    layer,
+                    x,
+                    topk_output,
+                    self.moe_runner_config,
+                )
             return StandardCombineInput(hidden_states=output)
 
         if self.use_marlin:

python/sglang/srt/model_executor/cuda_graph_buffer_registry.py

@@ -514,6 +514,7 @@ def build_decode_registry(
     enable_mamba_track: bool = False,
     is_encoder_decoder: bool = False,
     encoder_len_fill_value: int = 0,
+    encoder_lens_dtype: torch.dtype = torch.int32,
     enable_num_token_non_padded: bool = False,
     require_gathered_buffer: bool = False,
     enable_prefill_cp: bool = False,
@@ -632,7 +633,7 @@ def build_decode_registry(
             GraphSlot(
                 "encoder_lens",
                 _bs,
-                torch.int32,
+                encoder_lens_dtype,
                 axis="bs",
                 padding_policy=PaddingPolicy.FILL_ONCE,
                 pad_value=encoder_len_fill_value,
@@ -898,6 +899,7 @@ def build_eager_registry(
     enable_mamba_track: bool = False,
     is_encoder_decoder: bool = False,
     encoder_len_fill_value: int = 0,
+    encoder_lens_dtype: torch.dtype = torch.int32,
     dp_size: int = 1,
 ) -> CudaGraphBufferRegistry:
     """One fixed-max input registry for the ``EagerRunner``, serving BOTH eager
@@ -924,6 +926,7 @@ def build_eager_registry(
         enable_mamba_track=enable_mamba_track,
         is_encoder_decoder=is_encoder_decoder,
         encoder_len_fill_value=encoder_len_fill_value,
+        encoder_lens_dtype=encoder_lens_dtype,
         enable_num_token_non_padded=False,
         register_global_num_tokens=False,
         require_gathered_buffer=False,

python/sglang/srt/model_executor/runner/eager_runner.py

@@ -130,6 +130,9 @@ class EagerRunner(BaseRunner):
                 if is_encoder_decoder
                 else 0
             ),
+            encoder_lens_dtype=(
+                torch.int64 if torch.device(mr.device).type == "cpu" else torch.int32
+            ),
             dp_size=sa.dp_size,
         )
         # Eager has no capture step, so warm up here (run-once via mr._kernel_warmed_up).