flashinfer.comm.moe_a2a_sanitize_expert_ids

flashinfer.comm.moe_a2a_sanitize_expert_ids(expert_ids: Tensor, workspace: Tensor, metainfo: Tensor, ep_rank: int, invalid_expert_id: int)

Replace expert IDs not owned by this rank with invalid_expert_id.

Parameters:

  • expert_ids (torch.Tensor) – [local_num_tokens, top_k] int32 tensor of expert assignments (mutated in place).

  • workspace (torch.Tensor) – Shared workspace tensor.

  • metainfo (torch.Tensor) – Metainfo tensor returned by moe_a2a_initialize().

  • ep_rank (int) – Current expert-parallel rank.

  • invalid_expert_id (int) – Value to write where the original expert lies outside this rank’s local range.