[Bug]: [CuteDSL MoE] `generate_token_selected_experts`: ghost-token assignment to zero-allocation experts + dead-on-multi-rank prioritization

System Info

Description

Two related issues in generate_token_selected_experts (tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py), both in the loop that builds token_selected_experts for autotune profile inputs.

Bug 1: Ghost-token assignment when num_tokens_per_expert[j] == 0

cute_dsl_custom_ops.py#L143-L159

for j, num_tokens_j in enumerate(num_tokens_per_expert):
    ...
    for i in selection_order_j:
        if num_selected_experts[i] < self.top_k:
            token_selected_experts[
                i, num_selected_experts[i]] = j + self.local_expert_offset
            num_selected_experts[i] += 1
            num_tokens_j -= 1
            if num_tokens_j <= 0:
                break

When num_tokens_per_expert[j] == 0, the outer loop still enters the inner loop with num_tokens_j = 0. The first eligible token gets expert j assigned (num_selected_experts[i] → 1, num_tokens_j → -1), then break exits. So an expert that the per-expert allocation said should receive 0 tokens still receives 1.

Concrete example

DeepSeek-V3 with ep_size=16, top_k=8, num_tokens=1:

• average_num_tokens_per_rank = 1 × 8 / 16 = 0.5
• extra ≈ 0.83 (from the approx-max-load term)
• num_tokens_on_curr_rank = ceil(0.5 + 0.83) = 2
• divmod(2, 16) = (0, 2) → num_tokens_per_expert = [1, 1, 0, 0, ..., 0]

Intended: experts 0 and 1 each receive 1 token; experts 2–15 each receive 0.
Actual: experts 0 through 7 each receive 1 token (the inner loop ghosts experts 2–7 onto the one available token, then stops once that token has all top_k=8 slots filled).

Suggested fix

Skip the outer-loop body when num_tokens_j <= 0:

for j, num_tokens_j in enumerate(num_tokens_per_expert):
      if num_tokens_j <= 0:
          continue
      ...

Bug 2: Prioritization uses local expert index instead of global

cute_dsl_custom_ops.py#L145-L146

  prioritized = torch.nonzero(num_selected_experts <= (
      self.top_k - (self.num_experts - j))).squeeze(-1).tolist()

The limit expression uses the loop variable j, but the function iterates over local experts (enumerate(num_tokens_per_expert) runs j from 0 to num_local_experts - 1). The condition appears to reason about how many experts remain in the global iteration, which would need j + self.local_expert_offset.

Concrete impact

For ep_size=16, top_k=8, num_experts=256: at the last local iteration j = 15, limit = 8 - (256 - 15) = -233. The condition num_selected_experts[i] <= -233 is never true, so prioritized is always empty on any rank with local_expert_offset > 0. The "at-risk token" prioritization is effectively dead code on multi-rank EP.

Suggested fix

Use the global expert index:

global_j = j + self.local_expert_offset
prioritized = torch.nonzero(num_selected_experts <= (
    self.top_k - (self.num_experts - global_j))).squeeze(-1).tolist()

(If the multi-rank no-op behavior is intentional, then maybe document it explicitly — the current code reads as if prioritization is meant to apply across ranks.)

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Invoke TRT-LLM with DeepSeek R1 and the CuteDSL MoE backend. The code paths:

#L143-L159

cute_dsl_custom_ops.py#L145-L146

should be hit.

Expected behavior

See description for more detail, but in summary:

1. Skip the outer-loop body when num_tokens_j <= 0
2. Use the global expert index

actual behavior

See description for more detail, but in summary:

1. per-expert allocation said should receive 0 tokens still receives 1
2. condition appears to reason about how many experts remain in the global iteration, but uses local

additional notes

Note these were flagged by Gemini in this PR in FlashInfer:

https://github.com/flashinfer-ai/flashinfer/pull/3286

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

Start in tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py, function generate_token_selected_experts, at the referenced loop. Reproduce with DeepSeek R1 and the CuteDSL MoE backend, checking the zero-allocation and multi-rank cases described. Done means zero-token experts receive no assignments and prioritization uses the global expert index as expected.