[Fix]Fix one-sided MoE padding sentinel for local expert maps

[Kevin-XiongC]

Summary

Fix FlashInfer NVLink one-sided MoE dispatch padding when expert_map is present.

The one-sided all2all path pads each source rank up to runtime_max_tokens_per_rank. It previously used -1 as the invalid expert id. Some downstream kernels index expert_map before filtering padded tokens, so -1 can index the last expert-map entry and corrupt local-expert routing.

Without this fix, the Marlin MoE path can see padded -1 expert ids after dispatch and trigger an IMA when those ids are used through the local expert map. When an expert map is present, this PR uses a valid global expert id that is non-local to the current EP rank instead. Paths without expert_map keep using the out-of-range num_experts sentinel.

Source-level rationale

The original padding value is unreasonable for the expert_map path because it violates the downstream kernel contract:

• FlashInferNVLinkOneSidedPrepareAndFinalize.prepare() sends topk_ids as one of the all2all payloads and passes invalid_token_expert_id to moe_alltoall.dispatch(). This means padded tokens can be materialized into the received topk_ids tensor.
• EP expert_map is a global-expert-id to local-expert-id table with shape [global_num_experts]; non-local experts are represented by values of -1 inside the table. It is not designed to be indexed with a negative global expert id.
• The CUDA MoE permute preprocess path used by expert-map-aware MoE kernels indexes the map before skipped/non-local tokens are filtered: preprocessTopkIdLauncher(..., expert_map_ptr, n_expert, ...) runs when expert_map is present, and preprocessTopkIdKernel does local_expert_idx = smem_expert_map[topk_id] for each topk_id.
• Therefore, a padded topk_id == -1 is already invalid before later filtering can happen. It becomes expert_map[-1] in the kernel rather than "a skipped token", which can read the wrong entry or hit an illegal memory access. This is the path where Marlin MoE can IMA.

The replacement sentinel is a valid global expert id owned by another EP rank. It is safe to index into expert_map, maps to -1 for the local rank, and is then treated as non-local/skipped by the existing expert-map logic.

AI assistance was used to prepare this PR.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[gemini-code-assist]

Code Review

This pull request updates the FlashInferNVLinkOneSidedPrepareAndFinalize.prepare method to calculate a non-negative invalid_token_expert_id instead of using a hardcoded -1. This change ensures that padding tokens use an expert ID that is invalid for the local rank but non-negative, avoiding potential indexing issues with expert_map. Additionally, new unit tests were added to verify the correct calculation of this ID in different scenarios. I have no feedback to provide.