[Bugfix] bugfix for moe_mlp

[Clorist33]

What this PR does / why we need it?

This PR fixes a bug in the moe_mlp module by correcting the arguments passed to the torch_npu.npu_dequant_swiglu_quant function.It properly converts group_list from a cumulative sum to counts for the group_index parameter.

Does this PR introduce any user-facing change?

No

• vLLM version: v0.12.0
• vLLM main: vllm-project/vllm@ad32e3e

[gemini-code-assist]

Code Review

This pull request addresses a bug in the moe_mlp module by correctly converting the group_list argument from a cumulative sum to counts before passing it to torch_npu.npu_dequant_swiglu_quant. The logic for this conversion is sound. However, I've identified a potential edge case in the new code: it doesn't handle an empty group_list, which could lead to a runtime error. I have provided a suggestion to make the implementation more robust against this scenario.

[gemini-code-assist]

Using group_list[0].unsqueeze(0) will raise an IndexError if group_list is an empty tensor. This can occur if no tokens are routed to experts on the current device. Using slicing group_list[:1] is more robust as it returns an empty tensor for an empty input, preventing a crash.

            new_group = torch.cat([group_list[:1], group_diff], dim=0)

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

[zhoux77899]

It’s better to first check if there are cases where group_list_type != 1. You can extract a function to convert other types of group_list into count format, taking group_list and group_list_type as parameters, and determine how to perform the conversion within the function (refer to the cumsum_group_list function).

def count_group_list(group_list: torch.Tensor,
                     group_list_type: int) -> torch.Tensor
    if group_list_type not in [0, 1, 2]:
        raise ValueError(
            f"group_list_type should be in [0, 1, 2], but received {group_list_type}"
        )

    if group_list_type == 0:
        return torch.cat((group_list[:1], torch.diff(group_list)))
    if group_list_type == 1:
        return group_list

    # group_list_type == 2
    ...

[Clorist33]

Thanks for your suggestions. We have incorporated your proposal into the cumsum_group_list function. Additionally, could you please clarify what scenario corresponds to group_list_type == 2?

[Clorist33]

The changes are here：

def cumsum_group_list(group_list: torch.Tensor,
                      src_list_type: int,
                      dst_list_type: int,
                      active_num: int = 0,
                      expert_num: int = 0) -> torch.Tensor:
    if src_list_type not in [0, 1, 2]:
        raise ValueError(
            f"group_list_type should be in [0, 1, 2], but received {src_list_type}"
        )
    if src_list_type == dst_list_type:
        return group_list
    if src_list_type == 1 and dst_list_type == 0:
        return group_list.cumsum(dim=0)
    if src_list_type == 0 and dst_list_type == 1:
        group_diff = torch.diff(group_list)
        new_group = torch.cat([group_diff[0].unsqueeze(0), group_diff], dim=0)
        return new_group

    experts = pad(group_list[:, 0], (1, 0))
    tokens = pad(group_list[:, 1].cumsum(dim=0), (1, 0))
    cumsum_group_list = torch.full(size=(expert_num, ),
                                   fill_value=active_num,
                                   dtype=group_list.dtype,
                                   device=group_list.device)

    for i, (start, end) in enumerate(zip(experts[:-1], experts[1:])):
        if end > start:
            cumsum_group_list[start:end] = tokens[i]

    return cumsum_group_list

[zhoux77899]

And use the function here.

Suggested change

	group_index=new_group,
	group_index=count_group_list(group_list, group_list_type)

[zzzzwwjj]

what meanings src_list_type==2? How to handle src_list_type=2 and dst_list_type=0?

[Clorist33]

We also have the same confusion regarding the scenario where src_list_type == 2 in the file ops/fused_moe/moe_mlp.py on the main branch of the vllm-ascend repository. @zhoux77899 Would you please clarify this point ?

[zhoux77899]

Some ops like moe_init_routing_v2 can output the type 2 (key_value) group_list, but I’ve never seen this type of group_list actually used anywhere.

If there is a type 1 group_list like [0, 2, 1, 0]:

• group_list_type = 0 means cumsum group_list, it will be [0, 2, 3, 3];
• group_list_type = 1 means count group_list, it will be [0, 2, 1, 0];
• group_list_type = 2 means key_value group_list, it will be [[1, 2], [2, 1], [0, 0], [0, 0]];

[Clorist33]

group_list_type == 2 means key_value group_list, it will be [[0, 0], [1, 2], [2, 1], [3, 0]]？

[zhoux77899]

Apologies, when group_list_type = 2, group_list will be [[1, 2], [2, 1], [0, 0], [0, 0]]. It only contains tokens of active_num but is padded to expert_num. Maybe you can compare their differences using the script below.

import torch
import torch_npu

from vllm_ascend.ops.fused_moe.moe_mlp import cumsum_group_list


class GroupListTypeTester:
    def __init__(
        self,
        batch_size: int = 1,
        hidden_size: int = 768,
        active_experts: int = 2,
        num_experts: int = 4,
    ) -> None:
        self.batch_size = batch_size
        self.hidden_size = hidden_size
        self.active_experts = active_experts
        self.num_experts = num_experts

        self.x = torch.randn(size=(self.batch_size, self.hidden_size), dtype=torch.bfloat16).npu()
        self.expert_idx = torch.randint(low=0, high=self.num_experts, size=(self.batch_size, self.active_experts), dtype=torch.int32).npu()
        self.scale = torch.randn(size=(self.batch_size, ), dtype=torch.float32).npu()
        self.offset = None

        self.init_routing_kwargs = {
            "x": self.x,
            "expert_idx": self.expert_idx,
            "scale": self.scale,
            "offset": self.offset,
            "active_num": self.active_experts,
            "expert_num": self.num_experts,
            "expert_tokens_num_flag": True,
            "quant_mode": -1,
            "active_expert_range": [0, self.num_experts],
            "row_idx_type": 0,
        }

    def __call__(self) -> None:
        count_group_list = self.output_count_group_list()
        kv_group_list = self.output_kv_group_list()
        print(f"{count_group_list=}, cumsum_group_list_from_count={cumsum_group_list(count_group_list, 1)}")
        print(f"{kv_group_list=}, cumsum_group_list_from_kv={cumsum_group_list(kv_group_list, 2, self.active_experts, self.num_experts)}")

    def output_count_group_list(self) -> torch.Tensor:
        _, _, group_list, _ = torch_npu.npu_moe_init_routing_v2(
            **self.init_routing_kwargs,
            expert_tokens_num_type=1,
        )
        return group_list

    def output_kv_group_list(self) -> torch.Tensor:
        _, _, group_list, _ = torch_npu.npu_moe_init_routing_v2(
            **self.init_routing_kwargs,
            expert_tokens_num_type=2,
        )
        return group_list


if __name__ == "__main__":
    tester = GroupListTypeTester()
    tester()

I think consolidating all types of group_list computations into a single function might be overly complex, as it requires handling 6 different cases, and many call sites would also need modifications. Splitting them into separate functions would be more reasonable.

You may also consider whether to include the group_list_type = 2 scenario. The reason cumsum_group_list includes it is for safety redundancy, but I’m unsure about the specific use case it was designed for and I’ve never seen actually uses it at anywhere.

[github-actions]

This pull request has conflicts, please resolve those before we can evaluate the pull request.

[weijinqian0]

Need to check the range of group_list supported by the corresponding operator.

[Clorist33]

Confirmed. Code updated.

[wangxiyuan]

This need be merged to dev branch. So I merged this now.