IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo function

ep_signal_completion

def ep_signal_completion[p2p_world_size: Int, //, use_shmem: Bool, n_experts_per_device: Int = Int(0), skip_a2a: Bool = False, has_rank_flag: Bool = False, ep_ord_r: Bool = False](my_rank: Int32, dst_rank: Int32, recv_count_ptrs: Array[Pointer[UInt64, MutUntrackedOrigin], p2p_world_size], signal_offset: Int32, signal: UInt64, rank_completion_counter: Pointer[Int32, MutUntrackedOrigin], rank_flag_offset: Int32 = Int32(0))

Signals the completion of the communication by writing to the receive count buffer. Will use direct memory access if the target device is on the same node, and use the SHMEM API if the target device is on a different node.

For same-node signaling, uses normal stores and only issues a release store when the last expert for a destination rank is completed. This reduces the number of release stores from n_experts to p2p_world_size, and the destination checks each per-expert signal individually.

With ep_ord_r the routine implements rank-level completion instead, which a fused consumer needs when it acquires ONCE per source rank rather than per expert: the per-expert count word becomes a plain progress store issued before an ACQUIRE_RELEASE election, and the elected last producer release-stores a dedicated rank-completion flag.

That election ordering is load-bearing correctness, not an optimization. Each producer's RMW releases its own payload and count stores into the counter's modification order and acquires its predecessor's, so by induction the elected producer happens-after every other producer; its single system-scope release then carries the whole rank's state to a destination that acquires only the rank flag. A relaxed election would publish nothing but the elected producer's own stores.

Parameters:

  • ​p2p_world_size (Int): Number of ranks reachable by direct peer access.
  • ​use_shmem (Bool): Whether to signal remote ranks through the SHMEM API.
  • ​n_experts_per_device (Int): Local expert count, used for the election bound.
  • ​skip_a2a (Bool): Whether all-to-all traffic is skipped (device-scope signals).
  • ​has_rank_flag (Bool): Whether the caller reserved a dedicated rank-completion word per source rank in the receive-count buffer tail.
  • ​ep_ord_r (Bool): Whether to use rank-level completion (see above).

Args:

Was this page helpful?