[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)