[Bug]: EPLB load statistics problem

[david6666666]

Your current environment

The output of python collect_env.py

Your output of `python collect_env.py` here

A bug in load statistics: load statistics are collected during expert selection. However, only the top k experts for the current token that are routed to the local node are counted. The load of experts routed to other nodes is not counted.

🐛 Describe the bug

vllm/vllm/model_executor/layers/fused_moe/layer.py

Line 1336 in 555e722

if expert_map is not None:

            if expert_map is not None:
                topk_ids_local = expert_map[topk_ids]
                topk_ids_flatten = topk_ids_local.flatten()
            else:
                topk_ids_flatten = topk_ids.flatten()

            # Should be equivalent to:
            # ```
            # topk_ids_masked = topk_ids_local[topk_ids_local >= 0]
            # expert_load_view += topk_ids_masked.bincount(
            #     minlength=expert_load_view.shape[0])
            # ```
            # We use `scatter_add_` since `bincount` cannot be compiled

            # Performance optimization:
            # `masked_fill` is significantly faster than `masked_select`
            invalid_mask = topk_ids_flatten < 0
            # Replace invalid expert ids with 0 (just a dummy position)
            # to avoid out-of-bounds errors in scatter_add_
            index = topk_ids_flatten.masked_fill_(invalid_mask, 0)
            # `src` is the valid mask, which is 1 for valid and 0 for invalid
            src = ~invalid_mask

            expert_load_view.scatter_add_(dim=0,
                                          index=index.long(),
                                          src=src.to(expert_load_view))

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[david6666666]

@ycyaw66 @hsliuustc0106 will fix this bug later

@abmfy @Gongzq5

[abmfy]

Previously, I was testing with DP and naive dispatch, which doesn’t cause any issue here since the calculation is based on all-gathered tokens (see this line, where in NaiveAll2AllManager it's just sending all tokens to all ranks).

However, with specialized communication kernels like pplx-kernels or DeepEP, which follow a different code path, the dispatch might not have been executed yet at that point. In that case, this issue could indeed occur. We might need to revise the logic accordingly — for example, by checking whether dispatch has already happened and masking out non-local experts only if needed.

Could you please confirm if my understanding aligns with yours?

[david6666666]

Previously, I was testing with DP and naive dispatch, which doesn’t cause any issue here since the calculation is based on all-gathered tokens (see this line, where in NaiveAll2AllManager it's just sending all tokens to all ranks).

However, with specialized communication kernels like pplx-kernels or DeepEP, which follow a different code path, the dispatch might not have been executed yet at that point. In that case, this issue could indeed occur. We might need to revise the logic accordingly — for example, by checking whether dispatch has already happened and masking out non-local experts only if needed.

Could you please confirm if my understanding aligns with yours?

Yes, we have the same understanding on this issue