[ROCm][CI TG] refactor and fix deepep_moe test group (#46758)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
Jun 26, 2026 at 22:45 UTC
00e045b7c7b82599f626779e111233abd4d0a64e
3 files changed
+84
-58
tests/kernels/moe/test_deepep_moe.py
+60
-35
@@ -11,7 +11,7 @@ import torch.distributed
11
from torch.distributed import ProcessGroup
12
13
import vllm.envs as envs
14
-from tests.kernels.moe.utils import make_dummy_moe_config
14
+from tests.kernels.moe.utils import check_accuracy, make_dummy_moe_config
15
from vllm import _custom_ops as ops
16
from vllm.config import VllmConfig, set_current_vllm_config
17
from vllm.model_executor.layers.activation import SiluAndMul
@@ -227,39 +227,44 @@ def deep_ep_moe_impl(
227
out_hidden_states = torch.empty_like(test_tensors.rank_tokens)
228
total_num_tokens = test_tensors.rank_tokens.size(0)
229
230
+ quant_config = FusedMoEQuantConfig.make(
231
+ q_dtype,
232
+ w1_scale=w1_scale,
233
+ w2_scale=w2_scale,
234
+ per_act_token_quant=per_act_token_quant,
235
+ a1_scale=test_tensors.rank_token_scales,
236
+ )
237
+
238
+ # Build the kernel (and its DeepEP buffer) once and reuse it across chunks.
239
+ # Re-creating it per chunk re-inits rocSHMEM, which only allows one
240
+ # allocation per process on ROCm. The buffer is sized by max_tokens_per_rank
241
+ # so it is valid for every chunk (mirrors production's cached all2all handle).
242
+ mk: FusedMoEKernel = make_modular_kernel(
243
+ pg,
244
+ pgi,
245
+ low_latency_mode,
246
+ hidden_size,
247
+ dp_size,
248
+ num_experts,
249
+ num_local_experts,
250
+ q_dtype,
251
+ use_fp8_dispatch,
252
+ quant_config,
253
+ )
254
+
255
def process_chunk(chunk_start, chunk_end, skip_result_store=False):
256
rank_tokens_chunk = test_tensors.rank_tokens[chunk_start:chunk_end]
257
topk_weights_chunk = test_tensors.topk_weights[chunk_start:chunk_end]
258
topk_chunk = test_tensors.topk[chunk_start:chunk_end]
234
- rank_token_scales_chunk = test_tensors.rank_token_scales
235
- if (
236
- rank_token_scales_chunk is not None
237
- and rank_token_scales_chunk.size(0) == total_num_tokens
238
- ):
239
- # per act token
240
- rank_token_scales_chunk = rank_token_scales_chunk[chunk_start:chunk_end]
241
-
242
- quant_config = FusedMoEQuantConfig.make(
243
- q_dtype,
244
- w1_scale=w1_scale,
245
- w2_scale=w2_scale,
246
- per_act_token_quant=per_act_token_quant,
247
- a1_scale=rank_token_scales_chunk,
248
- )
259
250
- # Make modular kernel
251
- mk: FusedMoEKernel = make_modular_kernel(
252
- pg,
253
- pgi,
254
- low_latency_mode,
255
- hidden_size,
256
- dp_size,
257
- num_experts,
258
- num_local_experts,
259
- q_dtype,
260
- use_fp8_dispatch,
261
- quant_config,
262
- )
260
+ if low_latency_mode:
261
+ # Reusing one buffer leaves it dirty; the low-latency kernels need
262
+ # the zero-initialized regions reset before each dispatch.
263
+ mk.prepare_finalize.buffer.clean_low_latency_buffer(
264
+ MAX_TOKENS_PER_RANK,
265
+ hidden_size,
266
+ num_experts,
267
+ )
268
269
out = mk.apply(
270
hidden_states=rank_tokens_chunk,
@@ -350,6 +355,28 @@ def torch_moe_impl(
355
return out
356
357
358
+def assert_deepep_close(
359
+ expected: torch.Tensor,
360
+ actual: torch.Tensor,
361
+ k: int,
362
+ use_fp8_dispatch: bool,
363
+) -> None:
364
+ if use_fp8_dispatch and current_platform.is_fp8_fnuz():
365
+ # ROCm e4m3fnuz rounds differently than the reference quant,
366
+ # so DeepEP's fp8 dispatch can yield a few outliers even with
367
+ # a correct kernel; allow a small fraction of mismatches here.
368
+ atol = rtol = 1.5e-1
369
+ check_accuracy(expected, actual, atol=atol, rtol=rtol, percent=0.95)
370
+ return
371
+
372
+ torch.testing.assert_close(
373
+ expected,
374
+ actual,
375
+ atol=6e-2,
376
+ rtol=6e-2,
377
+ )
378
+
379
+
380
def _deep_ep_moe(
381
pgi: ProcessGroupInfo,
382
low_latency_mode: bool,
@@ -362,6 +389,9 @@ def _deep_ep_moe(
389
use_fp8_dispatch: bool,
390
per_act_token_quant: bool,
391
):
392
+ # Set seed in worker process for deterministic tensor generation.
393
+ set_random_seed(7)
394
+
395
device = torch.device(f"cuda:{pgi.local_rank}")
396
init_workspace_manager(device)
397
@@ -426,12 +456,7 @@ def _deep_ep_moe(
456
per_act_token_quant,
457
)
458
429
- torch.testing.assert_close(
430
- torch_combined,
431
- deepep_combined,
432
- atol=6e-2,
433
- rtol=6e-2,
434
- )
459
+ assert_deepep_close(torch_combined, deepep_combined, config.k, use_fp8_dispatch)
460
461
462
MNKs = [
tests/kernels/moe/test_ocp_mx_moe.py
+1
-23
@@ -9,6 +9,7 @@ import pytest
9
import torch
10
from packaging import version
11
12
+from tests.kernels.moe.utils import check_accuracy
13
from vllm._aiter_ops import is_aiter_found
14
from vllm.platforms import current_platform
15
from vllm.utils.flashinfer import has_flashinfer
@@ -515,29 +516,6 @@ def tg_mxfp4_moe(
516
return tg_result
517
518
518
-def check_accuracy(a, b, atol, rtol, percent):
519
- """Allow a mismatch percentage of 1 - percent."""
520
- if torch.any(torch.isnan(a)):
521
- raise Exception("NaN in reference output")
522
- if torch.any(torch.isnan(b)):
523
- raise Exception("NaN in actual output")
524
- if torch.any(torch.isinf(a)):
525
- raise Exception("Inf in reference output")
526
- if torch.any(torch.isinf(b)):
527
- raise Exception("Inf in actual output")
528
- assert a.shape == b.shape, f"Shape mismatch: {a.shape} vs {b.shape}"
529
-
530
- left = torch.abs(a - b)
531
- right = atol + rtol * torch.abs(b)
532
- count = torch.sum(left > right)
533
- mismatch_percent = count / a.numel()
534
- if mismatch_percent > 1 - percent:
535
- raise Exception(
536
- f"Mismatch percentage is {mismatch_percent:.4f} for rtol {rtol} "
537
- f"(threshold: {1 - percent:.4f})"
538
- )
539
-
540
-
519
@pytest.mark.parametrize("topk", [1, 4])
520
@pytest.mark.parametrize("num_experts", [32, 128])
521
@pytest.mark.parametrize("num_tokens", [1, 128, 1024])
tests/kernels/moe/utils.py
+23
@@ -653,3 +653,26 @@ def make_shared_experts(
653
return make_shared_experts_with_weights(
654
N, K, in_dtype, w1, w2, w1_s=w1_s, w2_s=w2_s, quant_dtype=quant_dtype
655
)
656
+
657
+
658
+def check_accuracy(a, b, atol, rtol, percent):
659
+ """Allow a mismatch percentage of 1 - percent."""
660
+ if torch.any(torch.isnan(a)):
661
+ raise Exception("NaN in reference output")
662
+ if torch.any(torch.isnan(b)):
663
+ raise Exception("NaN in actual output")
664
+ if torch.any(torch.isinf(a)):
665
+ raise Exception("Inf in reference output")
666
+ if torch.any(torch.isinf(b)):
667
+ raise Exception("Inf in actual output")
668
+ assert a.shape == b.shape, f"Shape mismatch: {a.shape} vs {b.shape}"
669
+
670
+ left = torch.abs(a - b)
671
+ right = atol + rtol * torch.abs(b)
672
+ count = torch.sum(left > right)
673
+ mismatch_percent = count / a.numel()
674
+ if mismatch_percent > 1 - percent:
675
+ raise Exception(
676
+ f"Mismatch percentage is {mismatch_percent:.4f} for rtol {rtol} "
677
+ f"(threshold: {1 - percent:.4f})"
678
+ )