[CPU] Fix TP padding issue on Phi-4 (#8289)

2025-08-18 07:25:26 +08:00
parent b341b7dbce
commit ebbb75e917
5 changed files with 93 additions and 27 deletions
--- a/python/sglang/srt/layers/quantization/unquant.py
+++ b/python/sglang/srt/layers/quantization/unquant.py
@@ -116,9 +116,15 @@ class UnquantizedLinearMethod(LinearMethodBase):
    ) -> torch.Tensor:

        if use_intel_amx_backend(layer):
-            return torch.ops.sgl_kernel.weight_packed_linear(
+            x_shapes = x.shape
+            if len(x_shapes) == 3:
+                x = x.view(-1, x.shape[-1])
+            output = torch.ops.sgl_kernel.weight_packed_linear(
                x, layer.weight, bias, True  # is_vnni
            )
+            if len(x_shapes) == 3:
+                output = output.view(x_shapes[0], x_shapes[1], -1)
+            return output

        return F.linear(x, layer.weight, bias)