[XPU] Use torch scaled_mm for XPU block FP8 linear (#35605)

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

@@ -66,6 +66,7 @@ from sglang.srt.layers.quantization.fp8_utils import (
     normalize_e4m3fn_to_e4m3fnuz,
     requant_block_scale_ue8m0_for_deepgemm,
     resolve_mxfp8_dense_gemm_backend,
+    torch_w8a8_block_fp8_linear,
     unshuffle_aiter_fp8_weight,
     use_aiter_bpreshuffle_gemm,
 )
@@ -799,6 +800,29 @@ class Fp8LinearMethod(LinearMethodBase):
                 layer.aiter_bpreshuffled = True
                 layer.weight.is_shuffled = True
 
+        if (
+            is_xpu()
+            and self.w8a8_block_fp8_linear is torch_w8a8_block_fp8_linear
+            and self.weight_block_size in ([1, 128], [128, 128])
+            and layer.weight_scale_inv.ndim == 2
+        ):
+            # Keep the checkpoint's logical [N-blocks, K-blocks] shape, but use
+            # transpose-contiguous storage. For [1, 128], scaled_mm transposes
+            # scale_b internally; for [128, 128], the wrapper passes scale_b.t().
+            # This avoids a per-forward contiguous/copy in either path.
+            scale = layer.weight_scale_inv.data
+            scale_b_is_contiguous = scale.t().is_contiguous()
+            if not scale_b_is_contiguous:
+                scale_reordered = torch.empty_strided(
+                    scale.shape,
+                    (1, scale.shape[0]),
+                    dtype=scale.dtype,
+                    device=scale.device,
+                )
+                scale_reordered.copy_(scale)
+                with torch.no_grad():
+                    layer.weight_scale_inv.set_(scale_reordered)
+
     def _process_mxfp8_linear_weight_scale(self, layer: Module) -> None:
         if not self.use_mxfp8:
             return

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

@@ -592,6 +592,58 @@ def dispatch_w8a8_block_fp8_linear() -> Callable:
     return _dispatch_auto_backend()
 
 
+def torch_w8a8_block_fp8_linear(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    block_size: List[int],
+    weight_scale: torch.Tensor,
+    input_scale: Optional[torch.Tensor] = None,
+    bias: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Run block-FP8 linear with Torch's scaled_mm implementation."""
+    if not isinstance(block_size, (list, tuple)) or len(block_size) != 2:
+        raise ValueError(
+            f"XPU block-FP8 scaled_mm expects a two-dimensional weight_block_size, "
+            f"but got {block_size}"
+        )
+    block_n, block_k = block_size
+    if block_k != 128 or block_n not in (1, 128):
+        raise ValueError(
+            "XPU block-FP8 scaled_mm supports weight_block_size [1, 128] or "
+            f"[128, 128], but got {block_size}"
+        )
+
+    scale_b_recipe = (
+        torch.nn.functional.ScalingType.BlockWise1x128
+        if block_n == 1
+        else torch.nn.functional.ScalingType.BlockWise128x128
+    )
+    input_2d = input.reshape(-1, input.shape[-1])
+    if input_scale is None:
+        q_input, activation_scale = per_token_group_quant_fp8(input_2d, block_k)
+    else:
+        q_input = input_2d
+        activation_scale = input_scale.reshape(-1, input_scale.shape[-1])
+
+    if q_input.stride(-1) != 1:
+        q_input = q_input.contiguous()
+    if weight.stride(-1) != 1:
+        weight = weight.contiguous()
+    weight_t = weight.t()
+    scale_b = weight_scale if block_n == 1 else weight_scale.t()
+    output = torch.nn.functional.scaled_mm(
+        q_input,
+        weight_t,
+        activation_scale,
+        torch.nn.functional.ScalingType.BlockWise1x128,
+        scale_b,
+        scale_b_recipe,
+        bias=bias,
+        output_dtype=torch.bfloat16 if input_scale is not None else input.dtype,
+    )
+    return output.view(*input.shape[:-1], weight.shape[0])
+
+
 def resolve_mxfp8_dense_gemm_backend() -> Mxfp8DenseGemmBackend:
     """Pick the MXFP8 dense linear backend, honoring `--fp8-gemm-backend` only when it
     names a backend that owns an MXFP8 dense kernel."""
@@ -804,7 +856,8 @@ def _dispatch_auto_backend() -> Callable:
     # 3. CUTLASS (if SM120 GPU and CUDA 12.8+)
     # 4. AITER (if AMD GPU with AITER enabled)
     # 5. NPU (Ascend)
-    # 6. Triton (fallback)
+    # 6. XPU (Intel GPU, PyTorch torch._scaled_mm)
+    # 7. Triton (fallback)
 
     if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
         return deepgemm_w8a8_block_fp8_linear_with_fallback
@@ -820,6 +873,8 @@ def _dispatch_auto_backend() -> Callable:
         )
 
         return npu_w8a8_mxfp8_linear
+    elif _is_xpu:
+        return torch_w8a8_block_fp8_linear
     else:
         return triton_w8a8_block_fp8_linear