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, enable_pdl: bool | None = None)

Sanitize invalid slots that contain no token routed to this rank by setting their expert IDs to 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 into slots that received no token (per recv_counters, e.g. padding beyond a source rank’s valid count).

  • enable_pdl (Optional[bool]) – Whether to use programmatic dependent launch. None auto-detects from the device.