[XPU] Pad MoE expert weight row stride to avoid L3 aliasing (#33905)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Alex Nails <alex.nails@radixark.ai>

python/sglang/srt/layers/moe/utils.py

@@ -685,6 +685,41 @@ class RoutingMethodType(IntEnum):
 AITER_PADDING_SIZE = 128
 TRITON_PADDING_SIZE = 128
 
+# Row-stride padding, in bytes, applied to XPU MoE expert weights whose K dim
+# lands on an L3 aliasing stride (see xpu_moe_ld_padding_elems). 64B matches
+# the 32 bf16 elements used by the sgl-kernel-xpu MoE benchmark. Expressed in
+# bytes because the aliasing is a property of the row's byte size, so this
+# stays correct if the path ever carries a non-bf16 weight dtype.
+#
+# Measured on BMG: halving this to 32B still clears the aliasing but runs ~6%
+# slower than not padding at all on hidden=7168 shapes (0.94x), presumably by
+# misaligning the grouped GEMM's row loads. Doubling it to 128B gains nothing
+# over 64B. Re-measure before changing.
+XPU_MOE_LD_PADDING_BYTES = 64
+
+
+def xpu_moe_ld_padding_elems(k_dim: int, itemsize: int) -> int:
+    """Extra elements to add to an XPU MoE weight's row stride (leading dim).
+
+    The Xe20 grouped GEMM walks B row-by-row over the K dim, so the row stride
+    in bytes decides which L3 set each row lands in. The L3 set index is
+    derived by XOR-folding address bits; when the row byte size is a multiple
+    of 2048 with an odd cofactor >= 3 (K = 3072, 7168, ... in bf16) successive
+    rows collapse onto a small number of sets and thrash. Padding the stride
+    (without changing the logical shape) breaks the aliasing.
+
+    Returns 0 when the shape is already well distributed, so callers can use
+    this to decide whether to allocate a padded buffer at all.
+    """
+    row_bytes = k_dim * itemsize
+    if row_bytes <= 0 or XPU_MOE_LD_PADDING_BYTES % itemsize != 0:
+        return 0
+    trailing_zeros = (row_bytes & -row_bytes).bit_length() - 1
+    odd_cofactor = row_bytes >> trailing_zeros
+    if trailing_zeros >= 11 and odd_cofactor >= 3:
+        return XPU_MOE_LD_PADDING_BYTES // itemsize
+    return 0
+
 
 # Unit of padding - context dependent
 def get_moe_padding_size(is_aiter_moe):

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

@@ -24,6 +24,7 @@ from sglang.srt.layers.moe import (
     get_moe_runner_backend,
 )
 from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
+from sglang.srt.layers.moe.utils import xpu_moe_ld_padding_elems
 from sglang.srt.layers.quantization.base_config import (
     FusedMoEMethodBase,
     LinearMethodBase,
@@ -296,6 +297,52 @@ class UnquantizedLinearMethod(LinearMethodBase):
         return output
 
 
+def _use_xpu_moe_ld_padding(use_triton_kernels: bool) -> bool:
+    """Whether MoE expert weights should get a padded row stride for XPU.
+
+    use_intel_xpu_backend() only tells us an XPU exists on this machine, not
+    that the weights being created land on it -- the env var can be set while
+    serving on CPU/CUDA. create_weights takes no device argument and allocates
+    under the model loader's ambient device context, so check that context too:
+    padding a non-XPU weight would make it non-contiguous for no benefit, and
+    other backends' MoE kernels expect contiguous expert tensors.
+
+    The Triton path stores B transposed and does not read a row stride, so it
+    is excluded even on XPU.
+    """
+    return (
+        use_intel_xpu_backend()
+        and torch.get_default_device().type == "xpu"
+        and not use_triton_kernels
+    )
+
+
+def _empty_xpu_moe_expert_weight(
+    num_experts: int,
+    n_dim: int,
+    k_dim: int,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Allocate an [E, N, K] XPU expert weight, over-allocating K when padding
+    its row stride would avoid L3 set aliasing.
+
+    Some K dims (3072, 7168 in bf16) put every weight row in the same handful
+    of L3 sets, which throttles the grouped GEMM's B loads. Over-allocating K
+    and returning a narrowed view keeps the logical [E, N, K] shape (so the
+    weight loader is unchanged) while giving the rows a non-aliasing stride.
+    The Xe20 grouped GEMM reads B's row stride from the tensor, so the padding
+    is transparent to it.
+
+    Callers must have checked _use_xpu_moe_ld_padding() first. K dims that are
+    already well distributed get no padding and allocate normally.
+    """
+    pad = xpu_moe_ld_padding_elems(k_dim, dtype.itemsize)
+    if pad == 0:
+        return torch.empty(num_experts, n_dim, k_dim, dtype=dtype)
+    # The view is non-contiguous; only the K slice is ever read or written.
+    return torch.empty(num_experts, n_dim, k_dim + pad, dtype=dtype)[:, :, :k_dim]
+
+
 class UnquantizedFusedMoEMethod(FusedMoEMethodBase, BaseFusedOp):
     """MoE method without quantization."""
 
@@ -325,6 +372,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, BaseFusedOp):
     ):
         self.with_bias = with_bias
 
+        # XPU only: the sgl-kernel-xpu grouped GEMM honours the weights' row
+        # stride, so it can be padded to dodge L3 set aliasing on unlucky K
+        # dims. Every other device allocates plainly, exactly as before.
+        pad_ld_for_xpu = _use_xpu_moe_ld_padding(self.use_triton_kernels)
+
         # Fused gate_up_proj (column parallel)
         w13_up_dim = (
             2 * intermediate_size_per_partition
@@ -334,10 +386,15 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, BaseFusedOp):
         w13_weight_n, w13_weight_k = (w13_up_dim, hidden_size)
         if self.use_triton_kernels:
             w13_weight_n, w13_weight_k = w13_weight_k, w13_weight_n
-        w13_weight = torch.nn.Parameter(
-            torch.empty(num_experts, w13_weight_n, w13_weight_k, dtype=params_dtype),
-            requires_grad=False,
-        )
+        if pad_ld_for_xpu:
+            w13_weight_data = _empty_xpu_moe_expert_weight(
+                num_experts, w13_weight_n, w13_weight_k, params_dtype
+            )
+        else:
+            w13_weight_data = torch.empty(
+                num_experts, w13_weight_n, w13_weight_k, dtype=params_dtype
+            )
+        w13_weight = torch.nn.Parameter(w13_weight_data, requires_grad=False)
         layer.register_parameter("w13_weight", w13_weight)
         set_weight_attrs(w13_weight, extra_weight_attrs)
 
@@ -356,10 +413,15 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, BaseFusedOp):
         )
         if self.use_triton_kernels:
             w2_weight_n, w2_weight_k = w2_weight_k, w2_weight_n
-        w2_weight = torch.nn.Parameter(
-            torch.empty(num_experts, w2_weight_n, w2_weight_k, dtype=params_dtype),
-            requires_grad=False,
-        )
+        if pad_ld_for_xpu:
+            w2_weight_data = _empty_xpu_moe_expert_weight(
+                num_experts, w2_weight_n, w2_weight_k, params_dtype
+            )
+        else:
+            w2_weight_data = torch.empty(
+                num_experts, w2_weight_n, w2_weight_k, dtype=params_dtype
+            )
+        w2_weight = torch.nn.Parameter(w2_weight_data, requires_grad=False)
         layer.register_parameter("w2_weight", w2_weight)
         set_weight_attrs(w2_weight, extra_weight_attrs)
 

test/registered/xpu/test_moe_ld_padding.py

@@ -0,0 +1,180 @@
+"""
+python3 -m unittest test_moe_ld_padding.py
+"""
+
+import unittest
+import unittest.mock
+
+import torch
+
+from sglang.srt.layers.moe.utils import (
+    XPU_MOE_LD_PADDING_BYTES,
+    xpu_moe_ld_padding_elems,
+)
+from sglang.srt.layers.quantization.unquant import _empty_xpu_moe_expert_weight
+from sglang.test.ci.ci_register import register_xpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_xpu_ci(est_time=30, suite="stage-b-test-1-gpu-xpu")
+
+
+class TestXpuMoeLdPadding(CustomTestCase):
+    def test_padding_selects_aliasing_shapes(self):
+        # bf16: row bytes = 2 * K. Aliasing when row bytes is a multiple of
+        # 2048 with an odd cofactor >= 3.
+        pad = XPU_MOE_LD_PADDING_BYTES // 2
+        for k in (3072, 5120, 6144, 7168, 14336):
+            self.assertEqual(xpu_moe_ld_padding_elems(k, 2), pad, f"K={k}")
+        # Pure powers of two are already well distributed, so are shapes whose
+        # row size is not a multiple of 2048.
+        for k in (1024, 2048, 4096, 8192, 1536, 2880):
+            self.assertEqual(xpu_moe_ld_padding_elems(k, 2), 0, f"K={k}")
+
+    def test_padding_scales_with_itemsize(self):
+        # The pad is a fixed byte count, so the element count scales inversely
+        # with itemsize, and the aliasing test is on bytes not elements.
+        self.assertEqual(
+            xpu_moe_ld_padding_elems(3072, 4), XPU_MOE_LD_PADDING_BYTES // 4
+        )
+        self.assertEqual(xpu_moe_ld_padding_elems(6144, 1), XPU_MOE_LD_PADDING_BYTES)
+        # 3072 bytes is not a multiple of 2048, so fp8/int8 K=3072 is fine.
+        self.assertEqual(xpu_moe_ld_padding_elems(3072, 1), 0)
+
+    def test_allocation_keeps_shape_and_pads_stride(self):
+        E, N, K = 4, 64, 3072
+        pad = xpu_moe_ld_padding_elems(K, 2)
+        self.assertGreater(pad, 0)
+
+        padded = _empty_xpu_moe_expert_weight(E, N, K, torch.bfloat16)
+        plain = torch.empty(E, N, K, dtype=torch.bfloat16)
+
+        # Logical shape is identical -- this is what keeps the weight loader,
+        # which indexes purely by shape, working unchanged.
+        self.assertEqual(padded.shape, plain.shape)
+        self.assertEqual(padded.stride(1), K + pad)
+        self.assertFalse(padded.is_contiguous())
+        self.assertTrue(plain.is_contiguous())
+
+        # A non-aliasing K allocates normally even on the XPU path.
+        unpadded = _empty_xpu_moe_expert_weight(E, N, 1024, torch.bfloat16)
+        self.assertTrue(unpadded.is_contiguous())
+
+    def test_only_pads_weights_that_land_on_xpu(self):
+        # SGLANG_USE_SGL_XPU only says an XPU exists on the machine; the weights
+        # can still be built for CPU/CUDA. create_weights takes no device
+        # argument, so the gate reads the ambient device context. Padding a
+        # non-XPU weight would make it non-contiguous for no benefit.
+        from sglang.srt.layers.moe import MoeRunnerConfig
+        from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod
+
+        class _Layer(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.moe_runner_config = MoeRunnerConfig(activation="silu")
+                self.moe_runner_config.is_gated = True
+
+        def build(device, use_triton_kernels=False):
+            method = UnquantizedFusedMoEMethod(use_triton_kernels=use_triton_kernels)
+            layer = _Layer()
+            with torch.device(device):
+                method.create_weights(
+                    layer=layer,
+                    num_experts=8,
+                    hidden_size=3072,
+                    intermediate_size_per_partition=3072,
+                    params_dtype=torch.bfloat16,
+                    with_bias=False,
+                )
+            return layer.w13_weight, layer.w2_weight
+
+        with unittest.mock.patch(
+            "sglang.srt.layers.quantization.unquant.use_intel_xpu_backend",
+            return_value=True,
+        ):
+            # Env var on but building for CPU -> must stay contiguous.
+            w13_cpu, w2_cpu = build("cpu")
+            self.assertTrue(w13_cpu.is_contiguous())
+            self.assertTrue(w2_cpu.is_contiguous())
+
+            if torch.xpu.is_available():
+                w13_xpu, _ = build("xpu")
+                self.assertFalse(w13_xpu.is_contiguous())
+                # The Triton path stores B transposed and ignores row stride.
+                w13_triton, _ = build("xpu", use_triton_kernels=True)
+                self.assertTrue(w13_triton.is_contiguous())
+
+        # Backend off entirely -> never padded, even on XPU.
+        with unittest.mock.patch(
+            "sglang.srt.layers.quantization.unquant.use_intel_xpu_backend",
+            return_value=False,
+        ):
+            device = "xpu" if torch.xpu.is_available() else "cpu"
+            w13, w2 = build(device)
+            self.assertTrue(w13.is_contiguous())
+            self.assertTrue(w2.is_contiguous())
+
+    def test_loader_style_copy_into_padded_view(self):
+        # Mirrors _load_w13 / _load_w2: narrow the destination along a dim and
+        # copy_ the checkpoint slice in. Must be exact despite the row gaps.
+        E, N, K = 4, 64, 3072
+        dst = _empty_xpu_moe_expert_weight(E, N, K, torch.bfloat16)
+        dst.zero_()
+        ref = torch.empty(E, N, K, dtype=torch.bfloat16).normal_()
+        half = N // 2
+        for e in range(E):
+            dst[e].narrow(0, 0, half).copy_(ref[e].narrow(0, 0, half))
+            dst[e].narrow(0, half, half).copy_(ref[e].narrow(0, half, half))
+        self.assertTrue(torch.equal(dst, ref))
+        # Still a padded view after the copies.
+        self.assertEqual(dst.stride(1), K + xpu_moe_ld_padding_elems(K, 2))
+
+
+@unittest.skipUnless(
+    torch.xpu.is_available(), "sgl-kernel-xpu grouped GEMM requires an XPU"
+)
+class TestXpuMoePaddedWeightsNumerics(CustomTestCase):
+    """The Xe20 grouped GEMM reads B's row stride from the tensor, so padded
+    weights must give bit-identical results to contiguous ones."""
+
+    def _run(self, hidden, inter, num_tokens, num_experts=8, topk=2):
+        from sgl_kernel import fused_experts
+
+        dtype, dev = torch.bfloat16, "xpu"
+        torch.manual_seed(0)
+        x = torch.empty(num_tokens, hidden, dtype=dtype, device=dev).normal_(0, 0.02)
+        gate = torch.randn(num_tokens, num_experts, device=dev, dtype=torch.float32)
+        topk_weights, topk_ids = torch.topk(torch.softmax(gate, -1), topk, -1)
+        topk_weights = topk_weights.to(dtype)
+
+        def alloc(n_dim, k_dim, pad):
+            with torch.device(dev):
+                if pad:
+                    return _empty_xpu_moe_expert_weight(
+                        num_experts, n_dim, k_dim, dtype
+                    )
+                return torch.empty(num_experts, n_dim, k_dim, dtype=dtype)
+
+        w13 = alloc(2 * inter, hidden, False).normal_(0, 0.02)
+        w2 = alloc(hidden, inter, False).normal_(0, 0.02)
+        w13_pad = alloc(2 * inter, hidden, True)
+        w2_pad = alloc(hidden, inter, True)
+        w13_pad.copy_(w13)
+        w2_pad.copy_(w2)
+
+        out = fused_experts(x, w13, w2, topk_weights, topk_ids)
+        out_pad = fused_experts(x, w13_pad, w2_pad, topk_weights, topk_ids)
+        torch.xpu.synchronize()
+        self.assertTrue(
+            torch.equal(out, out_pad),
+            f"padded weights changed the result for hidden={hidden} inter={inter}",
+        )
+
+    def test_padded_weights_bitwise_identical(self):
+        for hidden, inter in ((3072, 3072), (7168, 1024), (2880, 2880)):
+            for num_tokens in (64, 256):
+                with self.subTest(hidden=hidden, inter=inter, num_tokens=num_tokens):
+                    self._run(hidden, inter, num_tokens)
+
+
+if __name__ == "__main__":
+    unittest.main()