[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>
Divakar Verma committed
Sep 22, 2026 at 15:39 UTC
d1f63f2af98540082473cfa4aed3e55eb21dce8b
2 files changed
+57
-6
tests/kernels/moe/modular_kernel_tools/common.py
+15
-2
@@ -84,6 +84,8 @@ class Config:
84
85
world_size: int
86
87
+ activation: MoEActivation = MoEActivation.SILU
88
+
89
torch_trace_dir_path: str | None = None
90
91
def __post_init__(self):
@@ -311,6 +313,16 @@ class Config:
313
f"block={self.quant_block_shape})"
314
)
315
316
+ # Check activation support; NotImplementedError means no opinion.
317
+ try:
318
+ if not self.fused_experts_type._supports_activation(self.activation):
319
+ return False, (
320
+ f"FE {self.fused_experts_type.__name__} does not support "
321
+ f"activation {self.activation}"
322
+ )
323
+ except NotImplementedError:
324
+ pass
325
+
326
# Check block quantization support
327
is_block_quantized = self.quant_block_shape is not None
328
if is_block_quantized and self.quant_dtype is None:
@@ -605,6 +617,7 @@ def reference_moe_impl(
617
topk_ids=rank_tensors.topk_ids,
618
global_num_experts=config.E,
619
expert_map=None,
620
+ activation=config.activation,
621
w1_scale=w1_scale,
622
w2_scale=w2_scale,
623
a1_scale=a_scale,
@@ -648,7 +661,7 @@ def make_modular_kernel(
661
moe_parallel_config=moe_parallel_config,
662
in_dtype=config.dtype,
663
max_num_tokens=next_power_of_2(config.M),
651
- activation=MoEActivation.SILU,
664
+ activation=config.activation,
665
device=vllm_config.device_config.device,
666
routing_method=RoutingMethodType.DeepSeekV3,
667
)
@@ -774,7 +787,7 @@ def run_modular_kernel(
787
"w2": rank_weights.w2,
788
"topk_weights": rank_tensors.topk_weights,
789
"topk_ids": topk_ids,
777
- "activation": MoEActivation.SILU,
790
+ "activation": config.activation,
791
"expert_map": rank_tensors.expert_map,
792
"global_num_experts": config.E,
793
"apply_router_weight_on_input": config.topk == 1
tests/kernels/moe/test_modular_kernel_combinations.py
+42
-4
@@ -12,6 +12,7 @@ import torch
12
13
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
14
from vllm.config import VllmConfig, set_current_vllm_config
15
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
16
from vllm.platforms import current_platform
17
from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
18
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:
88
print(f"[AITER case] weight_key={w_key}, activation_key={a_key}")
89
90
91
+def assert_aiter_activation_case(config: Config) -> None:
92
+ """Make the AITER activation this config exercises explicit, instead of
93
+ AiterExperts being reached only indirectly through the general
94
+ activation sweep. See https://github.com/vllm-project/vllm/issues/54966."""
95
+ fe_cls = config.fused_experts_type
96
+ if getattr(fe_cls, "__name__", "") != "AiterExperts":
97
+ return
98
+
99
+ assert fe_cls._supports_activation(config.activation), (
100
+ f"AITER case activation={config.activation} reached the "
101
+ "modular-kernel harness, but AiterExperts._supports_activation "
102
+ "does not declare it supported."
103
+ )
104
+ print(f"[AITER case] activation={config.activation}")
105
+
106
+
107
def rank_worker(
108
pgi: ProcessGroupInfo,
109
vllm_config: VllmConfig,
@@ -160,6 +177,7 @@ def rank_worker(
177
continue
178
179
assert_aiter_quant_scheme_case(config)
180
+ assert_aiter_activation_case(config)
181
182
# modular kernel out
183
mk_out = run_modular_kernel(pgi, vllm_config, config, weights, rank_tensors)
@@ -225,6 +243,10 @@ Ns = [1024]
243
TOPKs = [4, 1]
244
Es = [32]
245
DTYPEs = [torch.bfloat16]
246
+MK_ACTIVATIONS = [
247
+ MoEActivation.SILU,
248
+ MoEActivation.GELU,
249
+]
250
251
252
def is_nyi_config(config: Config) -> bool:
@@ -236,7 +258,14 @@ def is_nyi_config(config: Config) -> bool:
258
unsupported_quant_config = (
259
config.is_per_act_token_quant + config.is_per_out_ch_quant
260
) == 1
239
- return unsupported_quant_config
261
+ if unsupported_quant_config:
262
+ return True
263
+
264
+ if config.activation != MoEActivation.SILU:
265
+ if getattr(config.fused_experts_type, "__name__", "") != "AiterExperts":
266
+ return True # AITER-only for this axis, for now
267
+ if config.quant_dtype is not None:
268
+ return True # unquantized-only for this axis, for now
269
270
return False
271
@@ -247,12 +276,13 @@ def generate_valid_test_cases(
276
cases = []
277
total = 0
278
250
- for k, n, e, dtype, quant_config, combination in product(
279
+ for k, n, e, dtype, quant_config, activation, combination in product(
280
Ks,
281
Ns,
282
Es,
283
DTYPEs,
284
MK_QUANT_CONFIGS,
285
+ MK_ACTIVATIONS,
286
product(prepare_finalize_types, MK_FUSED_EXPERT_TYPES),
287
):
288
total = total + 1
@@ -265,6 +295,7 @@ def generate_valid_test_cases(
295
topks=TOPKs,
296
dtype=dtype,
297
quant_config=quant_config,
298
+ activation=activation,
299
prepare_finalize_type=combination[0],
300
fused_experts_type=combination[1],
301
world_size=world_size,
@@ -292,6 +323,7 @@ def generate_valid_test_cases(
323
e,
324
dtype,
325
quant_config,
326
+ activation,
327
combination[0],
328
combination[1],
329
world_size,
@@ -304,7 +336,8 @@ def generate_valid_test_cases(
336
337
338
@pytest.mark.parametrize(
307
- "k,n,e,dtype,quant_config,prepare_finalize_type,fused_experts_type,world_size",
339
+ "k,n,e,dtype,quant_config,activation,"
340
+ "prepare_finalize_type,fused_experts_type,world_size",
341
generate_valid_test_cases(
342
world_size=2, prepare_finalize_types=MK_MULTI_GPU_PREPARE_FINALIZE_TYPES
343
),
@@ -316,6 +349,7 @@ def test_modular_kernel_combinations_multigpu(
349
e: int,
350
dtype: torch.dtype,
351
quant_config: TestMoEQuantConfig | None,
352
+ activation: MoEActivation,
353
prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,
354
fused_experts_type: mk.FusedMoEExperts,
355
world_size: int,
@@ -336,6 +370,7 @@ def test_modular_kernel_combinations_multigpu(
370
topks=TOPKs,
371
dtype=dtype,
372
quant_config=quant_config,
373
+ activation=activation,
374
prepare_finalize_type=prepare_finalize_type,
375
fused_experts_type=fused_experts_type,
376
world_size=world_size,
@@ -345,7 +380,8 @@ def test_modular_kernel_combinations_multigpu(
380
381
382
@pytest.mark.parametrize(
348
- "k,n,e,dtype,quant_config,prepare_finalize_type,fused_experts_type,world_size",
383
+ "k,n,e,dtype,quant_config,activation,"
384
+ "prepare_finalize_type,fused_experts_type,world_size",
385
generate_valid_test_cases(
386
world_size=1, prepare_finalize_types=MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES
387
),
@@ -356,6 +392,7 @@ def test_modular_kernel_combinations_singlegpu(
392
e: int,
393
dtype: torch.dtype,
394
quant_config: TestMoEQuantConfig | None,
395
+ activation: MoEActivation,
396
prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,
397
fused_experts_type: mk.FusedMoEExperts,
398
world_size: int,
@@ -372,6 +409,7 @@ def test_modular_kernel_combinations_singlegpu(
409
topks=TOPKs,
410
dtype=dtype,
411
quant_config=quant_config,
412
+ activation=activation,
413
prepare_finalize_type=prepare_finalize_type,
414
fused_experts_type=fused_experts_type,
415
world_size=world_size,