[ROCm][CI] Add GELU activation for AiterExperts in the modular-kernel coverage (#58030)

Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>

tests/kernels/moe/modular_kernel_tools/common.py

@@ -84,6 +84,8 @@ class Config:
 
     world_size: int
 
+    activation: MoEActivation = MoEActivation.SILU
+
     torch_trace_dir_path: str | None = None
 
     def __post_init__(self):
@@ -311,6 +313,16 @@ class Config:
                 f"block={self.quant_block_shape})"
             )
 
+        # Check activation support; NotImplementedError means no opinion.
+        try:
+            if not self.fused_experts_type._supports_activation(self.activation):
+                return False, (
+                    f"FE {self.fused_experts_type.__name__} does not support "
+                    f"activation {self.activation}"
+                )
+        except NotImplementedError:
+            pass
+
         # Check block quantization support
         is_block_quantized = self.quant_block_shape is not None
         if is_block_quantized and self.quant_dtype is None:
@@ -605,6 +617,7 @@ def reference_moe_impl(
         topk_ids=rank_tensors.topk_ids,
         global_num_experts=config.E,
         expert_map=None,
+        activation=config.activation,
         w1_scale=w1_scale,
         w2_scale=w2_scale,
         a1_scale=a_scale,
@@ -648,7 +661,7 @@ def make_modular_kernel(
         moe_parallel_config=moe_parallel_config,
         in_dtype=config.dtype,
         max_num_tokens=next_power_of_2(config.M),
-        activation=MoEActivation.SILU,
+        activation=config.activation,
         device=vllm_config.device_config.device,
         routing_method=RoutingMethodType.DeepSeekV3,
     )
@@ -774,7 +787,7 @@ def run_modular_kernel(
         "w2": rank_weights.w2,
         "topk_weights": rank_tensors.topk_weights,
         "topk_ids": topk_ids,
-        "activation": MoEActivation.SILU,
+        "activation": config.activation,
         "expert_map": rank_tensors.expert_map,
         "global_num_experts": config.E,
         "apply_router_weight_on_input": config.topk == 1

tests/kernels/moe/test_modular_kernel_combinations.py

@@ -12,6 +12,7 @@ import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.platforms import current_platform
 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
 from vllm.utils.import_utils import has_deep_ep, has_deep_gemm
@@ -87,6 +88,22 @@ def assert_aiter_quant_scheme_case(config: Config) -> None:
     print(f"[AITER case] weight_key={w_key}, activation_key={a_key}")
 
 
+def assert_aiter_activation_case(config: Config) -> None:
+    """Make the AITER activation this config exercises explicit, instead of
+    AiterExperts being reached only indirectly through the general
+    activation sweep. See https://github.com/vllm-project/vllm/issues/54966."""
+    fe_cls = config.fused_experts_type
+    if getattr(fe_cls, "__name__", "") != "AiterExperts":
+        return
+
+    assert fe_cls._supports_activation(config.activation), (
+        f"AITER case activation={config.activation} reached the "
+        "modular-kernel harness, but AiterExperts._supports_activation "
+        "does not declare it supported."
+    )
+    print(f"[AITER case] activation={config.activation}")
+
+
 def rank_worker(
     pgi: ProcessGroupInfo,
     vllm_config: VllmConfig,
@@ -160,6 +177,7 @@ def rank_worker(
                 continue
 
             assert_aiter_quant_scheme_case(config)
+            assert_aiter_activation_case(config)
 
             # modular kernel out
             mk_out = run_modular_kernel(pgi, vllm_config, config, weights, rank_tensors)
@@ -225,6 +243,10 @@ Ns = [1024]
 TOPKs = [4, 1]
 Es = [32]
 DTYPEs = [torch.bfloat16]
+MK_ACTIVATIONS = [
+    MoEActivation.SILU,
+    MoEActivation.GELU,
+]
 
 
 def is_nyi_config(config: Config) -> bool:
@@ -236,7 +258,14 @@ def is_nyi_config(config: Config) -> bool:
         unsupported_quant_config = (
             config.is_per_act_token_quant + config.is_per_out_ch_quant
         ) == 1
-        return unsupported_quant_config
+        if unsupported_quant_config:
+            return True
+
+    if config.activation != MoEActivation.SILU:
+        if getattr(config.fused_experts_type, "__name__", "") != "AiterExperts":
+            return True  # AITER-only for this axis, for now
+        if config.quant_dtype is not None:
+            return True  # unquantized-only for this axis, for now
 
     return False
 
@@ -247,12 +276,13 @@ def generate_valid_test_cases(
     cases = []
     total = 0
 
-    for k, n, e, dtype, quant_config, combination in product(
+    for k, n, e, dtype, quant_config, activation, combination in product(
         Ks,
         Ns,
         Es,
         DTYPEs,
         MK_QUANT_CONFIGS,
+        MK_ACTIVATIONS,
         product(prepare_finalize_types, MK_FUSED_EXPERT_TYPES),
     ):
         total = total + 1
@@ -265,6 +295,7 @@ def generate_valid_test_cases(
             topks=TOPKs,
             dtype=dtype,
             quant_config=quant_config,
+            activation=activation,
             prepare_finalize_type=combination[0],
             fused_experts_type=combination[1],
             world_size=world_size,
@@ -292,6 +323,7 @@ def generate_valid_test_cases(
                 e,
                 dtype,
                 quant_config,
+                activation,
                 combination[0],
                 combination[1],
                 world_size,
@@ -304,7 +336,8 @@ def generate_valid_test_cases(
 
 
 @pytest.mark.parametrize(
-    "k,n,e,dtype,quant_config,prepare_finalize_type,fused_experts_type,world_size",
+    "k,n,e,dtype,quant_config,activation,"
+    "prepare_finalize_type,fused_experts_type,world_size",
     generate_valid_test_cases(
         world_size=2, prepare_finalize_types=MK_MULTI_GPU_PREPARE_FINALIZE_TYPES
     ),
@@ -316,6 +349,7 @@ def test_modular_kernel_combinations_multigpu(
     e: int,
     dtype: torch.dtype,
     quant_config: TestMoEQuantConfig | None,
+    activation: MoEActivation,
     prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,
     fused_experts_type: mk.FusedMoEExperts,
     world_size: int,
@@ -336,6 +370,7 @@ def test_modular_kernel_combinations_multigpu(
         topks=TOPKs,
         dtype=dtype,
         quant_config=quant_config,
+        activation=activation,
         prepare_finalize_type=prepare_finalize_type,
         fused_experts_type=fused_experts_type,
         world_size=world_size,
@@ -345,7 +380,8 @@ def test_modular_kernel_combinations_multigpu(
 
 
 @pytest.mark.parametrize(
-    "k,n,e,dtype,quant_config,prepare_finalize_type,fused_experts_type,world_size",
+    "k,n,e,dtype,quant_config,activation,"
+    "prepare_finalize_type,fused_experts_type,world_size",
     generate_valid_test_cases(
         world_size=1, prepare_finalize_types=MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES
     ),
@@ -356,6 +392,7 @@ def test_modular_kernel_combinations_singlegpu(
     e: int,
     dtype: torch.dtype,
     quant_config: TestMoEQuantConfig | None,
+    activation: MoEActivation,
     prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,
     fused_experts_type: mk.FusedMoEExperts,
     world_size: int,
@@ -372,6 +409,7 @@ def test_modular_kernel_combinations_singlegpu(
         topks=TOPKs,
         dtype=dtype,
         quant_config=quant_config,
+        activation=activation,
         prepare_finalize_type=prepare_finalize_type,
         fused_experts_type=fused_experts_type,
         world_size=world_size,