@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 65c5e24689

[ROCm][CI][AITER Coverage] Harden MoE sorting-backend/dispatch env-var test matrix (#58393)

Signed-off-by: Divakar Verma <divakar.verma@amd.com>

Divakar Verma committed Sep 24, 2026 at 22:14 UTC 65c5e2468949410a0328f61fc941bbb8ee80fd04
4 files changed +355 -12
tests/kernels/moe/test_modular_kernel_combinations.py
+273 -7
@@ -2,20 +2,24 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import copy
5 +import inspect
6 import textwrap
7 import traceback
8 from itertools import product
9 from typing import Any
10 +from unittest import mock
11
12 import pytest
13 import torch
14
15 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
16 +from vllm._aiter_ops import rocm_aiter_ops
17 from vllm.config import VllmConfig, set_current_vllm_config
18 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
19 +from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import AiterExperts
20 from vllm.platforms import current_platform
21 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
18 -from vllm.utils.import_utils import has_deep_ep, has_deep_gemm
22 +from vllm.utils.import_utils import has_aiter, has_deep_ep, has_deep_gemm
23 from vllm.utils.torch_utils import set_random_seed
24 from vllm.v1.worker.workspace import init_workspace_manager
25
@@ -72,7 +76,7 @@ def assert_aiter_quant_scheme_case(config: Config) -> None:
76 indirectly through the general quant-config sweep.
77 See https://github.com/vllm-project/vllm/issues/54966."""
78 fe_cls = config.fused_experts_type
75 - if getattr(fe_cls, "__name__", "") != "AiterExperts":
79 + if fe_cls is not AiterExperts:
80 return
81
82 if config.quant_config is None:
@@ -93,7 +97,7 @@ def assert_aiter_activation_case(config: Config) -> None:
97 AiterExperts being reached only indirectly through the general
98 activation sweep. See https://github.com/vllm-project/vllm/issues/54966."""
99 fe_cls = config.fused_experts_type
96 - if getattr(fe_cls, "__name__", "") != "AiterExperts":
100 + if fe_cls is not AiterExperts:
101 return
102
103 assert fe_cls._supports_activation(config.activation), (
@@ -148,7 +152,7 @@ def rank_worker(
152 if (
153 topk == 1
154 and config.supports_apply_weight_on_input()
151 - and getattr(config.fused_experts_type, "__name__", "") == "AiterExperts"
155 + and config.fused_experts_type is AiterExperts
156 and config.quant_block_shape is not None
157 ):
158 print(
@@ -164,7 +168,7 @@ def rank_worker(
168 # https://github.com/vllm-project/vllm/issues/57029
169 if (
170 config.world_size > 1
167 - and getattr(config.fused_experts_type, "__name__", "") == "AiterExperts"
171 + and config.fused_experts_type is AiterExperts
172 and getattr(config.prepare_finalize_type, "__name__", "")
173 in ("DeepEPHTPrepareAndFinalize", "MoriPrepareAndFinalize")
174 ):
@@ -202,7 +206,7 @@ def rank_worker(
206
207 is_aiter_fp8 = (
208 _cp.is_rocm()
205 - and getattr(config.fused_experts_type, "__name__", "") == "AiterExperts"
209 + and config.fused_experts_type is AiterExperts
210 and config.quant_config is not None
211 )
212 if is_aiter_fp8:
@@ -262,7 +266,7 @@ def is_nyi_config(config: Config) -> bool:
266 return True
267
268 if config.activation != MoEActivation.SILU:
265 - if getattr(config.fused_experts_type, "__name__", "") != "AiterExperts":
269 + if config.fused_experts_type is not AiterExperts:
270 return True # AITER-only for this axis, for now
271 if config.quant_dtype is not None:
272 return True # unquantized-only for this axis, for now
@@ -425,6 +429,268 @@ def test_modular_kernel_combinations_singlegpu(
429 run(config, verbosity > 0)
430
431
432 +# AITER sorting-backend dispatch env-var matrix (issue #54966 step 3) -------
433 +#
434 +# AITER_USE_CK_MOE_SORTING / AITER_USE_FLYDSL_MOE_SORTING are aiter globals
435 +# read once at import time, so each combo needs a fresh child process (see
436 +# parallel_launch_with_config) rather than in-process monkeypatching.
437 +#
438 +# AITER_MOE_SORT_BACKEND is not covered: it only matters when output_aux=True,
439 +# which vLLM only sets on the MXFP4 path -- not yet wired into AiterExperts.
440 +
441 +
442 +def _aiter_sort_backend_hooks_available() -> bool:
443 + """Whether the installed aiter build exposes the sort-backend hooks that
444 + AITER_USE_CK_MOE_SORTING / AITER_USE_FLYDSL_MOE_SORTING select between.
445 + """
446 + try:
447 + import aiter.fused_moe as aiter_fused_moe
448 + except ImportError:
449 + return False
450 +
451 + required = (
452 + "_moe_sorting_impl",
453 + "_flydsl_moe_sorting",
454 + "_USE_CK_MOE_SORTING",
455 + "_USE_FLYDSL_MOE_SORTING",
456 + )
457 + return all(hasattr(aiter_fused_moe, name) for name in required)
458 +
459 +
460 +# Deliberately does NOT call _aiter_sort_backend_hooks_available() here: this
461 +# marker is evaluated at module-import time, which also happens inside each
462 +# spawned child (to unpickle its worker) -- before that case's env vars are
463 +# applied, permanently freezing aiter's sort-backend globals. The hooks check
464 +# instead runs inside the test body below (parent process only).
465 +require_aiter_moe = pytest.mark.skipif(
466 + not (
467 + current_platform.is_rocm()
468 + and has_aiter()
469 + and rocm_aiter_ops.is_fused_moe_enabled()
470 + ),
471 + reason=(
472 + "AITER MoE sorting-backend dispatch needs ROCm + AITER, with "
473 + "VLLM_ROCM_USE_AITER=1 and VLLM_ROCM_USE_AITER_MOE=1 set before "
474 + "the test process starts."
475 + ),
476 +)
477 +
478 +# (AITER_USE_CK_MOE_SORTING, AITER_USE_FLYDSL_MOE_SORTING, expected_backend)
479 +AITER_SORTING_BACKEND_ENV_CASES = [
480 + (1, 0, "ck"),
481 + (1, 1, "ck"), # CK wins over FlyDSL even when both are requested.
482 + (0, 1, "flydsl"),
483 + (0, 0, "opus"), # default when neither flag is set.
484 +]
485 +
486 +
487 +def _aiter_sorting_backend_worker(
488 + pgi: ProcessGroupInfo,
489 + vllm_config: VllmConfig,
490 + cpu_group,
491 + config: Config,
492 + weights: WeightTensors,
493 + verbose: bool,
494 + ck: int,
495 + flydsl: int,
496 + expected_backend: str,
497 +):
498 + # Imported lazily in the spawned child, after env vars are applied, so
499 + # aiter's import-time globals pick up this case's env.
500 + import aiter.fused_moe as aiter_fused_moe
501 +
502 + # Assert the env vars actually stuck in aiter's import-time globals,
503 + # rather than inferring it indirectly from which sorting function fires.
504 + assert bool(ck) == aiter_fused_moe._USE_CK_MOE_SORTING, (
505 + f"aiter.fused_moe._USE_CK_MOE_SORTING={aiter_fused_moe._USE_CK_MOE_SORTING} "
506 + f"but AITER_USE_CK_MOE_SORTING={ck} was set before this process started -- "
507 + "the env var never reached aiter's import-time globals."
508 + )
509 + assert bool(flydsl) == aiter_fused_moe._USE_FLYDSL_MOE_SORTING, (
510 + "aiter.fused_moe._USE_FLYDSL_MOE_SORTING="
511 + f"{aiter_fused_moe._USE_FLYDSL_MOE_SORTING} but AITER_USE_FLYDSL_MOE_SORTING="
512 + f"{flydsl} was set before this process started -- the env var never "
513 + "reached aiter's import-time globals."
514 + )
515 +
516 + # Captured before patching, since the patched attribute would resolve to
517 + # the mock's own (*args, **kwargs) signature instead of the real one.
518 + original_moe_sorting_impl = aiter_fused_moe._moe_sorting_impl
519 +
520 + with (
521 + mock.patch.object(
522 + aiter_fused_moe,
523 + "_moe_sorting_impl",
524 + wraps=original_moe_sorting_impl,
525 + ) as sorting_impl_mock,
526 + mock.patch.object(
527 + aiter_fused_moe,
528 + "_flydsl_moe_sorting",
529 + wraps=aiter_fused_moe._flydsl_moe_sorting,
530 + ) as flydsl_mock,
531 + ):
532 + # Reuses the real correctness/accuracy checking rank_worker already
533 + # does for every other AiterExperts combo, instead of duplicating it.
534 + rank_worker(pgi, vllm_config, cpu_group, config, weights, verbose)
535 +
536 + if expected_backend == "flydsl":
537 + assert flydsl_mock.call_count > 0, "Expected FlyDSL sorting to fire."
538 + assert sorting_impl_mock.call_count == 0, (
539 + "FlyDSL was requested and eligible, but the opus/CK sorting "
540 + "path fired instead -- a silent fallback."
541 + )
542 + else:
543 + assert flydsl_mock.call_count == 0, (
544 + f"Expected the '{expected_backend}' sorting path, but FlyDSL "
545 + "fired instead."
546 + )
547 + assert sorting_impl_mock.call_count > 0, (
548 + f"Expected the '{expected_backend}' sorting path to fire."
549 + )
550 + # Bind by signature so this works regardless of whether use_opus
551 + # is passed positionally or by keyword.
552 + call = sorting_impl_mock.call_args
553 + bound = inspect.signature(original_moe_sorting_impl).bind(
554 + *call.args, **call.kwargs
555 + )
556 + use_opus = bound.arguments["use_opus"]
557 + assert use_opus == (expected_backend == "opus"), (
558 + f"Expected use_opus={expected_backend == 'opus'} for the "
559 + f"'{expected_backend}' sorting path, got use_opus={use_opus}."
560 + )
561 +
562 +
563 +@require_aiter_moe
564 +@pytest.mark.parametrize("ck,flydsl,expected_backend", AITER_SORTING_BACKEND_ENV_CASES)
565 +def test_aiter_moe_sorting_backend_dispatch_env_matrix(
566 + ck: int, flydsl: int, expected_backend: str
567 +):
568 + """See https://github.com/vllm-project/vllm/issues/54966 ("Test AITER
569 + backend and dispatch settings")."""
570 + if not _aiter_sort_backend_hooks_available():
571 + pytest.skip(
572 + "Installed aiter build predates the sort-backend hooks "
573 + "(_moe_sorting_impl, _flydsl_moe_sorting, _USE_CK_MOE_SORTING, "
574 + "_USE_FLYDSL_MOE_SORTING)."
575 + )
576 +
577 + from vllm.model_executor.layers.fused_moe.prepare_finalize import (
578 + MoEPrepareAndFinalizeNoDPEPModular,
579 + )
580 +
581 + config = Config(
582 + Ms=[32],
583 + K=2048,
584 + N=1024,
585 + E=32,
586 + topks=[4],
587 + dtype=torch.bfloat16,
588 + quant_config=None,
589 + prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular,
590 + fused_experts_type=AiterExperts,
591 + world_size=1,
592 + )
593 + assert config.is_valid()[0]
594 +
595 + weights = WeightTensors.make(config)
596 + vllm_config, env_dict = config.make_env_data()
597 + env_dict = {
598 + **env_dict,
599 + "AITER_USE_CK_MOE_SORTING": str(ck),
600 + "AITER_USE_FLYDSL_MOE_SORTING": str(flydsl),
601 + }
602 +
603 + parallel_launch_with_config(
604 + config.world_size,
605 + _aiter_sorting_backend_worker,
606 + vllm_config,
607 + env_dict,
608 + config,
609 + weights,
610 + False,
611 + ck,
612 + flydsl,
613 + expected_backend,
614 + )
615 +
616 +
617 +def _aiter_dispatch_policy_worker(
618 + pgi: ProcessGroupInfo,
619 + vllm_config: VllmConfig,
620 + cpu_group,
621 + config: Config,
622 + weights: WeightTensors,
623 + verbose: bool,
624 + dispatch_policy: int,
625 +):
626 + # Drives the real AiterExperts.apply() -> rocm_aiter_fused_experts() call
627 + # through the modular kernel, so it also catches a regression in apply()'s
628 + # own forwarding (unlike test_rocm_aiter_moe.py's version, which supplies
629 + # the policy value directly to rocm_aiter_fused_experts()).
630 + with mock.patch.object(
631 + rocm_aiter_ops, "fused_moe", wraps=rocm_aiter_ops.fused_moe
632 + ) as fused_moe_mock:
633 + rank_worker(pgi, vllm_config, cpu_group, config, weights, verbose)
634 +
635 + assert fused_moe_mock.call_count > 0, (
636 + "Expected AiterExperts.apply() to call rocm_aiter_ops.fused_moe."
637 + )
638 + for call in fused_moe_mock.call_args_list:
639 + forwarded = call.kwargs["moe_sorting_dispatch_policy"]
640 + assert forwarded == dispatch_policy, (
641 + f"AiterExperts.apply() forwarded moe_sorting_dispatch_policy="
642 + f"{forwarded}, but VLLM_ROCM_AITER_MOE_DISPATCH_POLICY="
643 + f"{dispatch_policy} was set before this process started."
644 + )
645 +
646 +
647 +@require_aiter_moe
648 +@pytest.mark.parametrize("dispatch_policy", [0, 1, 2])
649 +def test_aiter_moe_dispatch_policy_forwarded_through_apply(dispatch_policy: int):
650 + """See https://github.com/vllm-project/vllm/issues/54966 ("Test AITER
651 + backend and dispatch settings").
652 +
653 + Complements test_aiter_moe_dispatch_policy_forwarded_to_fused_moe in
654 + test_rocm_aiter_moe.py by driving the real AiterExperts.apply() call
655 + instead of supplying moe_sorting_dispatch_policy directly.
656 + """
657 + from vllm.model_executor.layers.fused_moe.prepare_finalize import (
658 + MoEPrepareAndFinalizeNoDPEPModular,
659 + )
660 +
661 + config = Config(
662 + Ms=[32],
663 + K=2048,
664 + N=1024,
665 + E=32,
666 + topks=[4],
667 + dtype=torch.bfloat16,
668 + quant_config=None,
669 + prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular,
670 + fused_experts_type=AiterExperts,
671 + world_size=1,
672 + )
673 + assert config.is_valid()[0]
674 +
675 + weights = WeightTensors.make(config)
676 + vllm_config, env_dict = config.make_env_data()
677 + env_dict = {
678 + **env_dict,
679 + "VLLM_ROCM_AITER_MOE_DISPATCH_POLICY": str(dispatch_policy),
680 + }
681 +
682 + parallel_launch_with_config(
683 + config.world_size,
684 + _aiter_dispatch_policy_worker,
685 + vllm_config,
686 + env_dict,
687 + config,
688 + weights,
689 + False,
690 + dispatch_policy,
691 + )
692 +
693 +
694 if __name__ == "__main__":
695 # Ability to test individual PrepareAndFinalize and FusedExperts combination
696 from .modular_kernel_tools.cli_args import make_config, make_config_arg_parser
tests/kernels/moe/test_rocm_aiter_moe.py
+71
@@ -22,6 +22,7 @@ import importlib
22 import math
23 import warnings
24 from typing import Any, NamedTuple
25 +from unittest import mock
26
27 import pytest
28 import torch
@@ -645,6 +646,76 @@ def test_aiter_moe_padding_env_var(
646 assert envs.VLLM_ROCM_MOE_PADDING is moe_padding
647
648
649 +# Dispatch-policy forwarding test ------------------------------------------
650 +
651 +
652 +@pytest.mark.parametrize("dispatch_policy", [0, 1, 2])
653 +def test_aiter_moe_dispatch_policy_forwarded_to_fused_moe(
654 + dispatch_policy: int,
655 + monkeypatch: pytest.MonkeyPatch,
656 +):
657 + """VLLM_ROCM_AITER_MOE_DISPATCH_POLICY should reach rocm_aiter_ops.fused_moe
658 + unchanged, the same forwarding AiterExperts.apply does via
659 + rocm_aiter_fused_experts, for every documented policy value (0=auto,
660 + 1=always single-pass, 2=always multi-pass). See vllm-project/vllm#54966.
661 + """
662 + from tests.kernels.moe.utils import make_dummy_moe_config
663 + from vllm._aiter_ops import rocm_aiter_ops
664 + from vllm.model_executor.layers.fused_moe.activation import MoEActivation
665 + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
666 + rocm_aiter_fused_experts,
667 + )
668 +
669 + _assert_aiter_supported()
670 + case = _make_moe_case(
671 + num_tokens=8,
672 + hidden_dim=256,
673 + intermediate_dim=512,
674 + num_experts=4,
675 + topk=2,
676 + seed=5,
677 + )
678 + w1_shuffled, w2_shuffled = _shuffle_moe_weights(case["w1"], case["w2"])
679 + moe_config = make_dummy_moe_config(
680 + num_experts=4,
681 + experts_per_token=2,
682 + hidden_dim=256,
683 + intermediate_size=512,
684 + in_dtype=torch.bfloat16,
685 + )
686 +
687 + with monkeypatch.context() as mp:
688 + mp.setenv("VLLM_ROCM_AITER_MOE_DISPATCH_POLICY", str(dispatch_policy))
689 + _reload_envs()
690 + rocm_aiter_ops.refresh_env_variables()
691 +
692 + assert rocm_aiter_ops.get_moe_dispatch_policy() == dispatch_policy, (
693 + "rocm_aiter_ops cached a stale dispatch policy after the env var changed."
694 + )
695 +
696 + with mock.patch.object(
697 + rocm_aiter_ops, "fused_moe", wraps=rocm_aiter_ops.fused_moe
698 + ) as fused_moe_mock:
699 + rocm_aiter_fused_experts(
700 + hidden_states=case["hidden_states"],
701 + w1=w1_shuffled,
702 + w2=w2_shuffled,
703 + topk_weights=case["topk_weights"],
704 + topk_ids=case["topk_ids"],
705 + moe_config=moe_config,
706 + activation=MoEActivation.SILU,
707 + moe_sorting_dispatch_policy=rocm_aiter_ops.get_moe_dispatch_policy(),
708 + )
709 +
710 + assert (
711 + fused_moe_mock.call_args.kwargs["moe_sorting_dispatch_policy"]
712 + == dispatch_policy
713 + ), (
714 + "VLLM_ROCM_AITER_MOE_DISPATCH_POLICY was not forwarded to "
715 + "rocm_aiter_ops.fused_moe."
716 + )
717 +
718 +
719 # Enum tests --------------------------------------------------------------
720
721
vllm/_aiter_ops.py
+1
@@ -2019,6 +2019,7 @@ class rocm_aiter_ops:
2019 cls._MOE_SITUV2 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2
2020 _sync_aiter_situv2_moe_env()
2021 cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
2022 + cls._MOE_DISPATCH_POLICY = envs.VLLM_ROCM_AITER_MOE_DISPATCH_POLICY
2023
2024 @staticmethod
2025 def get_aiter_activation_type(activation_str: str) -> "ActivationType | None":
vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py
+10 -5
@@ -7,11 +7,6 @@ import json
7 import os
8
9 import torch
10 -from aiter.fused_moe import moe_sorting as aiter_moe_sorting
11 -from aiter.ops.flydsl.kernels.moe_2stage_a16wmix import (
12 - flydsl_a16w4_gemm1,
13 - flydsl_a16w4_gemm2,
14 -)
10
11 from vllm.logger import init_logger
12 from vllm.utils.platform_utils import get_device_name_as_file_name
@@ -81,6 +76,10 @@ def moe_sorting(
76 model_dim: int,
77 block_m: int,
78 ):
79 + # Imported lazily so aiter's import-time sort-backend env vars are read
80 + # when this is called, not whenever this module happens to be imported.
81 + from aiter.fused_moe import moe_sorting as aiter_moe_sorting
82 +
83 topk_ids_i32 = topk_ids.to(torch.int32)
84 topk_w_f32 = topk_weights.to(torch.float32)
85 sorted_ids, sorted_w, sorted_expert_ids, num_valid_ids, _moe_buf = (
@@ -185,6 +184,12 @@ def fused_flydsl_moe_impl(
184 tile_n2: int | None = None,
185 tile_k2: int | None = None,
186 ) -> torch.Tensor:
187 + # Imported lazily -- see the comment in moe_sorting() above.
188 + from aiter.ops.flydsl.kernels.moe_2stage_a16wmix import (
189 + flydsl_a16w4_gemm1,
190 + flydsl_a16w4_gemm2,
191 + )
192 +
193 device = hidden_states.device
194 tokens = hidden_states.shape[0]
195 model_dim = hidden_states.shape[1]