[Misc][ROCm] Restrict Aiter moe to specific models.

[charlifu]

This PR adds model restriction to aiter moe kernels. On rocm, only Mixtral and DeepSeek models will use aiter moe kernels.

• Add cache decorator to is_rocm_aiter_moe_enabled()
• Set vllm_config global value for _process_weights_after_loading function.
• Use get_current_vllm_config to obtain the model name in the first time is_rocm_aiter_moe_enabled invoked.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[SageMoore]

Looks good. Just one question.

[SageMoore]

I don't fully appreciate the implications of setting the config around _process_weights_after_loading. Could you explain a bit why it's necessary?

[charlifu]

The current_vllm_config value is not set when _process_weights_after_loading is being callled. Currently, this value is set only when creating the model class. So we do not have the information when determining whether to use aiter moe.

We definitely have other options.

• Add model config parameter to _process_weights_after_loading and is_rocm_aiter_moe_enabled functions. We might have to add more cases to here.
• Add a private member of model config to fused_moe layer class and set the value when creating the layer class.

Both require changing the interface. If you are concerned that exposing current vllm config during the execution of _process_weigths_after_loading could cause any potential issues. I prefer second option, since we call is_rocm_aiter_moe_enabled function during model execution as well.

[tlrmchlsmth]

Why restrict the AITER kernels to specific models? Is this a temporary performance hack or do you plan to maintain this list of models long term?

[tlrmchlsmth]

It looks like is_rocm_aiter_moe_enabled is called during the actual forward pass, when (IIUC) it's not valid to call get_current_vllm_config(), as it's not set (e.g. in dispatch_fused_experts_func). That should be resolved before landing otherwise users will get spammed with warnings

[charlifu]

That's why I add a cache decorator, it will use the cached value during the actually forward pass.

[tlrmchlsmth]

I see. But then is_rocm_aiter_moe_enabled would not be valid when instantiating a second LLM

[charlifu]

Oh, yeah. You are right.

[tlrmchlsmth]

Maybe we should just set the current vllm config during the forward pass in the model runner. @youkaichao WDYT?

[charlifu]

Why restrict the AITER kernels to specific models? Is this a temporary performance hack or do you plan to maintain this list of models long term?

We plan to maintain this list long term: AITER assembly MoE is manually tuned per model characteristics. We would like to enable it only for models that have been fully vetted to avoid performance pitfalls

[charlifu]

@SageMoore Thanks to @tlrmchlsmth, current way to make this change is not valid when creating a second LLM.

[ProExpertProg]

What if we did something similar to FpLinearOp where we initialize an object that dispatched at init time instead of forward time? I think we're trying to move all kernels in that direction anyway.

[tjtanaa]

@charlifu @ProExpertProg @tlrmchlsmth I think the issue addressed by this PR is being resolved in this PR (#16727).

The AITER MoE kernel selection logic has been cleaned up and now it is determined based on the arguments passed into therocm_aiter_fused_experts. There is no need to maintain a list of support models

https://github.com/vllm-project/vllm/pull/16727/files#diff-033e75fc2b8c4fd797f15d1ef0d8b6079996e8f0f9b1bd50cdfe53bcc3c593bc

def rocm_aiter_fused_experts(
    hidden_states: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    inplace: bool = False,
    activation: str = "silu",
+    apply_router_weight_on_input: bool = False,
+    use_fp8_w8a8: bool = False,
    use_int8_w8a8: bool = False,
    use_int8_w8a16: bool = False,
    use_int4_w4a16: bool = False,
+    per_channel_quant: bool = False,
    global_num_experts: int = -1,
    expert_map: Optional[torch.Tensor] = None,
    w1_scale: Optional[torch.Tensor] = None,
    w2_scale: Optional[torch.Tensor] = None,
    w1_zp: Optional[torch.Tensor] = None,
    w2_zp: Optional[torch.Tensor] = None,
    a1_scale: Optional[torch.Tensor] = None,
    a2_scale: Optional[torch.Tensor] = None,
+    block_shape: Optional[List[int]] = None,
    allow_deep_gemm: bool = False,
) -> torch.Tensor:

    import aiter as rocm_aiter
    import aiter.fused_moe_bf16_asm as rocm_aiter_asm_fmoe

    from vllm.model_executor.layers.quantization.utils.fp8_utils import (
        per_token_group_quant_fp8)

    # All AITER Fused MoE kernels are expecting the following datatypes
    topk_weights = topk_weights.to(torch.float32)
    topk_ids = topk_ids.to(torch.int32)

+    if (block_shape is not None) and use_fp8_w8a8:
        assert not apply_router_weight_on_input, (
            "apply_router_weight_on_input is not supported for block scaled moe"
        )

        assert w1_scale is not None
        assert w2_scale is not None

        local_E = E = w1.shape[0]
        if expert_map is not None:
            E = expert_map.numel()

        topk = topk_ids.shape[1]
        model_dim = w1.shape[-1]
        dtype = hidden_states.dtype
        # The default block sizes are 128 in AITER.
        if block_shape is None:
            block_shape = [128, 128]

        scale_blk_k = block_shape[1]

        (
            sorted_token_ids,
            sorted_weight_buf,
            sorted_expert_ids,
            num_valid_ids,
            out_asm,
        ) = rocm_aiter_asm_fmoe.moe_sorting_ck(topk_ids,
                                               topk_weights,
                                               E,
                                               model_dim,
                                               dtype,
                                               expert_mask=expert_map)

        a1, a1_scale = per_token_group_quant_fp8(hidden_states, scale_blk_k)
        rocm_aiter.fmoe_fp8_blockscale_g1u1(
            out_asm,
            a1,
            w1,
            w2,
            sorted_token_ids,
            sorted_weight_buf,
            sorted_expert_ids,
            num_valid_ids,
            topk,
            w1_scale.view(local_E, -1),
            w2_scale.view(local_E, -1),
            a1_scale.t().contiguous(),
            block_shape[0],
            block_shape[1],
            None,
        )
        return out_asm

+    elif per_channel_quant and apply_router_weight_on_input and use_fp8_w8a8:
        # AITER tkw1 kernel for FP8 models with `apply_router_weight_on_input`
        # This applies topk_weights on the GEMM output of the first FC layer
        #  rather than the second FC.
        assert (topk_weights.dim() == 2
                ), "`topk_weights` should be in shape (num_tokens, topk)"
        assert topk_weights.shape[-1] == 1, (
            "Only support topk=1 when"
            " `apply_router_weight_on_input` is True")

        return rocm_aiter_asm_moe_tkw1(hidden_states,
                                       w1,
                                       w2,
                                       topk_weights,
                                       topk_ids,
                                       fc1_scale=w1_scale,
                                       fc2_scale=w2_scale,
                                       fc1_smooth_scale=None,
                                       fc2_smooth_scale=None,
                                       a16=False,
                                       per_tensor_quant_scale=None,
                                       expert_mask=expert_map,
                                       activation_str=activation)

    elif use_fp8_w8a8:
        assert not apply_router_weight_on_input, (
            "apply_router_weight_on_input is not supported for fp8_w8a8")
        return rocm_aiter_asm_fmoe.asm_moe(hidden_states=hidden_states,
                                           w1=w1,
                                           w2=w2,
                                           topk_weight=topk_weights,
                                           topk_ids=topk_ids,
                                           fc1_scale=w1_scale,
                                           fc2_scale=w2_scale,
                                           fc1_smooth_scale=None,
                                           fc2_smooth_scale=None,
                                           a16=False)

    if apply_router_weight_on_input:
        assert (topk_weights.dim() == 2
                ), "`topk_weights` should be in shape (num_tokens, topk)"
        _, topk = topk_weights.shape
        assert (
            topk == 1
        ), "Only support topk=1 when `apply_router_weight_on_input` is True"

        hidden_states = hidden_states * topk_weights.to(hidden_states.dtype)
        topk_ids = topk_ids.to(torch.int32)
        topk_weights = torch.ones_like(topk_weights, dtype=torch.float32)

    return rocm_aiter.ck_moe(hidden_states=hidden_states,
                             w1=w1,
                             w2=w2,
                             topk_weights=topk_weights,
                             topk_ids=topk_ids)

[charlifu]

@tjtanaa Thank you for the information. Will close this.