flashinfer.moe_ep.CakeAlltoAll

class flashinfer.moe_ep.CakeAlltoAll(bootstrap: BootstrapConfig, params: MoEEpCommParams, config: CakeAlltoAllConfig | None = None)

Bases: MoEEpCommunication

Rank-major dispatch/combine with the generated Cake kernels.

hidden_states, its optional scale factors, topk_ids and topk_weights travel as payloads of a single dispatch. Received rows beyond each source rank’s token count get invalid_expert_id routing.

classmethod is_platform_supported() → bool

Whether this backend can run on the current machine.

Callable before construction; performs no collective work.

dispatch(hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor | None = None, *, hidden_states_scale: torch.Tensor | None = None, runtime_max_tokens_per_rank: int | None = None, eplb_local_stats: torch.Tensor | None = None, active_rank_mask: torch.Tensor | None = None) → MoEEpDispatchResult

See MoEEpCommunication.dispatch().

active_rank_mask is a CPU uint64 mask from flashinfer.moe_ep.backends.split.comm.cake.kernels.active_rank_mask(); tokens routed to a masked-off rank are dropped. Requires enable_rank_mask=True.

get_combine_input_buffer(dtype: torch.dtype) → torch.Tensor

[ep_size * tokens_per_rank, hidden_size] view of the combine payload region of this rank’s workspace.

combine(expert_output: torch.Tensor, *, output: torch.Tensor | None = None, active_rank_mask: torch.Tensor | None = None) → torch.Tensor

See MoEEpCommunication.combine().

active_rank_mask must match the mask passed to dispatch().

checkpoint_prepare() → None

Release the workspace’s physical memory before a process checkpoint.

A CUDA process checkpoint cannot save the workspace’s MNNVL physical memory, which peer processes import. This method releases it but keeps the workspace’s virtual addresses reserved, so CUDA graphs captured on the workspace replay correctly once checkpoint_restore() has backed it with new memory:

communication.checkpoint_prepare()
# checkpoint and restore the process
communication.checkpoint_restore(comm_backend)

Collective over the EP group: every rank calls it in the same order, after all work on the workspace has completed and only between rounds. Instances sharing the workspace may call it again, which does nothing. If it raises once unmapping has begun, the workspace cannot be reused: restart the rank rather than retry.

checkpoint_restore(comm_backend: CommBackend) → None

Back the workspace with new memory after checkpoint_prepare() and re-initialize it.

Collective over comm_backend, a fresh backend that must span the EP group with the original ranks and world size. Every rank calls it in the same order. Does nothing when the workspace is mapped. If it raises once remapping has begun, the workspace cannot be reused: restart the rank rather than retry.

destroy() → None

Release this instance’s share of the workspace.

The last instance on a workspace frees it; every rank must have finished its last combine on it.