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 trait

EpiloguePeerSink

A second destination for the epilogue's SMEM output tile.

The epilogue normally stores its tile to the local C tensor. A sink can take that tile somewhere else instead -- for expert parallelism, straight to the peers that own the rows -- in which case the local store is dead and the caller need not allocate C at all.

A sink is a STATELESS comptime policy: every method is a @staticmethod and the destination's runtime state arrives as p3_cfg, exactly as it did when this epilogue owned the send. So a sink is never a value, contributes no kernel-argument bytes whether enabled or not, and needs no construction threading through the writer.

Implementations:

• NullPeerSink: Enabled is False; every body is pass and every call site sits inside a comptime if, so the no-sink build emits nothing for it.
• EpPeerSink (closed tree): the EP-combine peer send.

Implemented traits

AnyType

comptime members

Enabled

comptime Enabled

Whether this sink owns the output. When True the local C store is dead, and the epilogue skips it along with its TMA descriptor.

SendWarps

comptime SendWarps

Warps the sink needs above the scheduler, or 0 to run inside the output warps. The warp-role geometry reserves exactly this many.

Required methods

send_tile

static def send_tile[stage: Int, num_threads: Int, TileT: AnyType](tile: TileT, m_abs: UInt32, n_abs: UInt32, m_end: UInt32, expert_id: Int32, p3_cfg: P3PeerSendConfig, p3_control: Int, tid: Int)

Takes one finished output tile in place of the local store.

Parameters:

• ​stage (Int): Output pipeline stage the tile belongs to.
• ​num_threads (Int): Threads cooperating on the tile.
• ​TileT (AnyType): The epilogue's SMEM tile view type.

Args:

• ​tile (TileT): The finished, barrier-visible output tile in SMEM.
• ​m_abs (UInt32): Absolute row of the tile's first row.
• ​n_abs (UInt32): Absolute column of the tile's first column.
• ​m_end (UInt32): Exclusive row bound for the tile's group.
• ​expert_id (Int32): Logical per-rank expert slot of this tile.
• ​p3_cfg (P3PeerSendConfig): The sink's runtime state.
• ​p3_control (Int): Runtime gate; negative disables the sink.
• ​tid (Int): Calling thread's index within num_threads.

service

static def service[TilesT: AnyType](tiles: TilesT, send_tid: Int, p3_control: Int, p3_cfg: P3PeerSendConfig)

Runs the sink's own warps until the epilogue terminates them.

Parameters:

• ​TilesT (AnyType): The epilogue's SMEM tile array type.

Args:

• ​tiles (TilesT): The epilogue's SMEM output tile array.
• ​send_tid (Int): Calling thread's index within SendWarps warps.
• ​p3_control (Int): Runtime gate; negative leaves these warps idle.
• ​p3_cfg (P3PeerSendConfig): The sink's runtime state.

drain

static def drain(p3_control: Int, p3_cfg: P3PeerSendConfig)

Collects work deferred from the previous tile.

Args:

• ​p3_control (Int): Runtime gate; negative means the warps never ran.
• ​p3_cfg (P3PeerSendConfig): The sink's runtime state.

shutdown

static def shutdown(p3_control: Int, p3_cfg: P3PeerSendConfig, drain_pending: Bool)

Retires the sink's warps with arrivals matched.

Args:

• ​p3_control (Int): Runtime gate; negative means the warps never ran.
• ​p3_cfg (P3PeerSendConfig): The sink's runtime state.
• ​drain_pending (Bool): Whether a deferred drain is outstanding.