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, *, backend: Literal['trtllm', 'cake'] = 'trtllm')

Sanitize invalid slots that contain no token routed to this rank by setting their expert IDs to invalid_expert_id.

backend="trtllm" is the default. For a workspace initialized with backend="cake", explicitly pass backend="cake" here as well.

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.

  • backend (MoeAlltoAllBackend) – Communication backend. Defaults to "trtllm"; use "cake" only with a workspace initialized by the Cake backend.