[ROCm][CI] Enable modular OAI Triton MoE tests (#46434)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Andreas Karatzas committed
Aug 19, 2026 at 13:46 UTC
480d4f0d154a00b719abf876dd9c1a98bc5acbf7
1 file changed
+9
-8
tests/kernels/moe/test_modular_oai_triton_moe.py
+9
-8
@@ -257,7 +257,8 @@ def oai_triton_moe_impl(
257
258
259
@pytest.mark.skipif(
260
- not current_platform.is_cuda_alike(), reason="Requires CUDA-alike platform."
260
+ not OAITritonExperts._supports_current_device(),
261
+ reason="OAI Triton MoE is not supported on this device.",
262
)
263
@pytest.mark.parametrize("dtype", [torch.bfloat16])
264
@pytest.mark.parametrize("m,n,k", MNK)
@@ -319,7 +320,8 @@ def test_oai_triton_moe(
320
321
322
@pytest.mark.skipif(
322
- not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
323
+ not UnfusedOAITritonExperts._supports_current_device(),
324
+ reason="Unfused OAI Triton MoE is not supported on this device.",
325
)
326
def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_init):
327
"""Exercise ``UnfusedOAITritonExperts.apply`` with explicit workspaces.
@@ -345,10 +347,11 @@ def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_
347
w2_bias_tri,
348
w1_precision_config,
349
w2_precision_config,
348
- _x_pad,
350
+ x_pad,
351
) = make_weights(dtype, k, n, num_experts)
352
353
x = torch.randn((m, k), dtype=dtype, device="cuda")
354
+ x_tri = F.pad(x, (0, x_pad, 0, 0))
355
router_logits = torch.randn(m, num_experts, device="cuda", dtype=dtype)
356
topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1, sorted=True)
357
topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1)
@@ -367,10 +370,7 @@ def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_
370
)
371
experts = UnfusedOAITritonExperts(moe_config, quant_config)
372
370
- if not UnfusedOAITritonExperts._supports_current_device():
371
- pytest.skip("UnfusedOAITritonExperts does not support this device")
372
-
373
- _, _, N, K, top_k = experts.moe_problem_size(x, w1_tri, w2_tri, topk_ids)
373
+ _, _, N, K, top_k = experts.moe_problem_size(x_tri, w1_tri, w2_tri, topk_ids)
374
assert top_k == topk
375
ws13_shape, ws2_shape, out_shape = experts.workspace_shapes(
376
m,
@@ -390,7 +390,7 @@ def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_
390
out_ref = torch_moe_impl(x, w1, w2, w1_bias, w2_bias, topk_weights, topk_ids)
391
experts.apply(
392
output=output,
393
- hidden_states=x,
393
+ hidden_states=x_tri,
394
w1=w1_tri,
395
w2=w2_tri,
396
topk_weights=topk_weights,
@@ -405,5 +405,6 @@ def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_
405
expert_tokens_meta=None,
406
apply_router_weight_on_input=False,
407
)
408
+ output = output[..., :k]
409
410
assert_close(ref=out_ref, tri=output, maxtol=0.025, rmstol=0.005)