diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 3ec522c8ec..1944aaac57 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/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 diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index fc1d8ac07b..9bfd3b5a2e 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/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,