diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 8faaa984b..593c65372 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/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): diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index fe78b4fcc..7eaa1e7f2 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/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) diff --git a/test/registered/xpu/test_moe_ld_padding.py b/test/registered/xpu/test_moe_ld_padding.py new file mode 100644 index 000000000..ca9cf0d63 --- /dev/null +++ b/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()