[Misc] Enable test_fused_moe_wn16 on XPU (#51672)
Signed-off-by: pmanczak <pawel.manczak@intel.com>
pmanczak committed
Aug 10, 2026 at 15:57 UTC
ec21f61d82cfe74723bf51e94beb1a6665fb736e
1 file changed
+23
-13
tests/kernels/moe/test_moe.py
+23
-13
@@ -575,6 +575,10 @@ def test_naive_block_assignment_moe(
575
@pytest.mark.parametrize("group_size", [64, 128])
576
@pytest.mark.parametrize("has_zp", [True, False])
577
@pytest.mark.parametrize("weight_bits", [4, 8])
578
+@pytest.mark.skipif(
579
+ not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
580
+ reason="MoE WNA16 Triton kernels require CUDA/ROCm or XPU.",
581
+)
582
def test_fused_moe_wn16(
583
m: int,
584
n: int,
@@ -587,10 +591,10 @@ def test_fused_moe_wn16(
591
has_zp: bool,
592
weight_bits: int,
593
):
590
- a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
591
- w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
592
- w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
593
- score = torch.randn((m, e), device="cuda", dtype=dtype)
594
+ a = torch.randn((m, k), device=DEVICE_TYPE, dtype=dtype) / 10
595
+ w1 = torch.randn((e, 2 * n, k), device=DEVICE_TYPE, dtype=dtype) / 10
596
+ w2 = torch.randn((e, k, n), device=DEVICE_TYPE, dtype=dtype) / 10
597
+ score = torch.randn((m, e), device=DEVICE_TYPE, dtype=dtype)
598
599
if weight_bits == 4:
600
pack_factor = 2
@@ -602,16 +606,22 @@ def test_fused_moe_wn16(
606
w1_ref = w1.clone()
607
w2_ref = w2.clone()
608
w1_qweight = torch.empty(
605
- (e, 2 * n, k // pack_factor), device="cuda", dtype=torch.uint8
609
+ (e, 2 * n, k // pack_factor), device=DEVICE_TYPE, dtype=torch.uint8
610
+ )
611
+ w2_qweight = torch.empty(
612
+ (e, k, n // pack_factor), device=DEVICE_TYPE, dtype=torch.uint8
613
)
607
- w2_qweight = torch.empty((e, k, n // pack_factor), device="cuda", dtype=torch.uint8)
608
- w1_scales = torch.empty((e, 2 * n, k // group_size), device="cuda", dtype=dtype)
609
- w2_scales = torch.empty((e, k, n // group_size), device="cuda", dtype=dtype)
614
+ w1_scales = torch.empty(
615
+ (e, 2 * n, k // group_size), device=DEVICE_TYPE, dtype=dtype
616
+ )
617
+ w2_scales = torch.empty((e, k, n // group_size), device=DEVICE_TYPE, dtype=dtype)
618
w1_qzeros = torch.empty(
611
- (e, 2 * n // pack_factor, k // group_size), device="cuda", dtype=torch.uint8
619
+ (e, 2 * n // pack_factor, k // group_size),
620
+ device=DEVICE_TYPE,
621
+ dtype=torch.uint8,
622
)
623
w2_qzeros = torch.empty(
614
- (e, k // pack_factor, n // group_size), device="cuda", dtype=torch.uint8
624
+ (e, k // pack_factor, n // group_size), device=DEVICE_TYPE, dtype=torch.uint8
625
)
626
627
for i in range(e * 2):
@@ -653,9 +663,9 @@ def test_fused_moe_wn16(
663
664
if ep_size > 1:
665
local_e = e // ep_size
656
- e_ids = torch.randint(0, e, (local_e,), device="cuda", dtype=torch.int32)
657
- e_map = torch.full((e,), -1, device="cuda", dtype=torch.int32)
658
- e_map[e_ids] = torch.arange(local_e, device="cuda", dtype=torch.int32)
666
+ e_ids = torch.randint(0, e, (local_e,), device=DEVICE_TYPE, dtype=torch.int32)
667
+ e_map = torch.full((e,), -1, device=DEVICE_TYPE, dtype=torch.int32)
668
+ e_map[e_ids] = torch.arange(local_e, device=DEVICE_TYPE, dtype=torch.int32)
669
w1_ref = w1_ref[e_ids]
670
w2_ref = w2_ref[e_ids]
671
w1_qweight = w1_qweight[e_ids]