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 module

ep_comm

Expert-parallelism (EP) communication kernels for MoE token dispatch and combine.

Implements the token dispatch (scatter) and combine (gather) collectives used in Mixture-of-Experts inference across multiple GPUs, including FP8 and FP4 quantization paths for bandwidth-limited transfers.

comptime values

BLOCK_SCOPE

comptime BLOCK_SCOPE = _BLOCK_SCOPE()

DEVICE_SCOPE

comptime DEVICE_SCOPE = _DEVICE_SCOPE()

DISPATCH_COPY_BEGIN

comptime DISPATCH_COPY_BEGIN = 0

Comm warp index starts copying its first owned token.

DISPATCH_CTA_INFO

comptime DISPATCH_CTA_INFO = 12

The CTA's send share as tokens | first_token << 16 | n_signal_sms << 40 | n_send_sms << 48 (index 0).

DISPATCH_PACK_DONE

comptime DISPATCH_PACK_DONE = 1

Comm warp index finishes packing an owned token; reported once per token.

DISPATCH_PAIR_COPY_DONE

comptime DISPATCH_PAIR_COPY_DONE = 9

Publisher index completes its first pair's copy and scale placement.

DISPATCH_PAIR_COPY_ISSUE

comptime DISPATCH_PAIR_COPY_ISSUE = 8

Publisher index issues its first pair's peer copy.

DISPATCH_PAIR_FINISHED

comptime DISPATCH_PAIR_FINISHED = 10

Publisher index issues its first pair's finished-count increment.

DISPATCH_PAIR_INFO

comptime DISPATCH_PAIR_INFO = 11

Publisher index's first pair as token | dst_rank << 16 | expert << 24 | slot << 32, with its pair count in bits 48 and up.

DISPATCH_PAIR_KEY

comptime DISPATCH_PAIR_KEY = 7

Publisher index's first pair as final_row | counter_offset << 32.

DISPATCH_PAIR_RESERVE_ISSUE

comptime DISPATCH_PAIR_RESERVE_ISSUE = 4

Publisher index issues its first pair's slot reservation.

DISPATCH_PAIR_RESERVED

comptime DISPATCH_PAIR_RESERVED = 5

Publisher index's first slot reservation returns.

DISPATCH_PAIR_ROW

comptime DISPATCH_PAIR_ROW = 6

Publisher index resolves its first pair's final row.

DISPATCH_SEND_JOIN_ARRIVE

comptime DISPATCH_SEND_JOIN_ARRIVE = 2

Comm warp index arrives at the send join.

DISPATCH_SEND_JOIN_PASS

comptime DISPATCH_SEND_JOIN_PASS = 3

Comm warp index leaves the send join.

elementwise_epilogue_type

comptime elementwise_epilogue_type = def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None

EP_COPY_SRC_WIDTH

comptime EP_COPY_SRC_WIDTH = 8

Source elements a thread handles per copy item in the quantizing token formats. Hardcoded there rather than derived from simd_width_of, so the role split has to agree with it explicitly.

EP_DATA_READY_FLAG

comptime EP_DATA_READY_FLAG = 1024

EP_LOCAL_SYNC_RESERVED_WORDS

comptime EP_LOCAL_SYNC_RESERVED_WORDS = 262144

Int32 words (1 MiB) reserved past the regions of an EPLocalSyncCounters buffer for one fused EP kernel's own counters: the fused MegaMoE path keeps a typed view there. Only a buffer of allocation_size() has them, and ep.init zeroes them with the rest of the buffer.

input_scales_wrapper_type

comptime input_scales_wrapper_type = def[dtype: DType](Int) capturing thin -> Scalar[dtype]

MAX_GPUS_PER_NODE

comptime MAX_GPUS_PER_NODE = 8

MAX_SMS_PER_DEVICE

comptime MAX_SMS_PER_DEVICE = 304

router_weights_wrapper_type

comptime router_weights_wrapper_type = def[width: Int](token_idx: Int, topk_id: Int) capturing thin -> SIMD[.float32, width]

SF_ATOM_ELEMS_LOCAL

comptime SF_ATOM_ELEMS_LOCAL = ((load_from_mem SF_ATOM_M.__getitem_param__[Int(0)]()) * SF_ROW_STRIDE_LOCAL)

SF_ROW_STRIDE_LOCAL

comptime SF_ROW_STRIDE_LOCAL = ((load_from_mem SF_ATOM_M.__getitem_param__[Int(1)]()) * Int(4))

Structs

• ​BF16TokenFormat: Token format that transmits the full hidden state in BFloat16.
• ​BlockwiseFP8TokenFormat: Token format that quantizes the hidden state to FP8 with block-wise scales.
• ​EPCombineKernel: Implements combine_async and combine_wait kernel logic for Expert Parallelism.
• ​EPDispatchKernel: Implements dispatch_async and dispatch_wait kernel logic for Expert Parallelism.
• ​EPLocalSyncCounters: Manages atomic counters for EP kernel synchronization within a device.
• ​EPRoleSplit: Splits one dispatch CTA into a copy role and a publisher role.
• ​MXTokenFormat: Token format for MX (microscaling) FP4 quantization.
• ​NullDispatchTrace: Discards every dispatch-phase event.
• ​NVBlockScaledTokenFormat: Token format for NVIDIA block-scaled FP4/FP8 quantization.

Traits

• ​DispatchTraceSink: Receives the dispatch-phase events (DISPATCH_*).
• ​TokenFormat: Specifies the wire format for a single MoE token in EP dispatch/combine.

Functions

• ​block_memcpy: Copies a memory area from source to destination. This function will use the vectorized store and load instructions to copy the memory area. User should make sure pointers are aligned to the simd width.
• ​block_prefix_sum: Performs a prefix sum (scan) operation across all threads in a block.
• ​combine_async_kernel: Send tokens to the original rank based on the src_info tensor. This kernel utilizes the non-blocking SHMEM API, and would return immediately after initiating the communication. The communication is considered complete after calling the combine_wait_kernel.
• ​combine_kernel: Fused combine kernel that combines combine_async and combine_wait functionality in a single kernel launch.
• ​combine_wait_kernel: This kernel is called after the combine_kernel to complete the communication. It will keep polling the receive count buffer, and once the count is no longer MAX_FINITE, it can confirm that the communication is complete from a remote rank.
• ​dispatch_async_kernel: Dispatch tokens to experts on remote ranks based on the top-k expert IDs. This kernel utilizes the non-blocking SHMEM API if use_shmem is True, and would return immediately after initiating the communication. The communication is considered complete after calling the dispatch_wait_kernel.
• ​dispatch_kernel: Fused dispatch kernel that combines dispatch_async and dispatch_wait functionality in a single kernel launch.
• ​dispatch_wait_kernel: This kernel is called after the dispatch_kernel to complete the communication. It will keep polling the receive count buffer, and once the count is no longer MAX_FINITE, it can confirm that the communication is complete from a remote rank.
• ​ep_signal_completion: 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.
• ​fused_silu_kernel: This kernel performs the SILU operation for all the MLPs in the EP MoE module. We need to manually implement the kernel here is because after the EP dispatch phase, the actual number of received tokens is not known to the host. This kernel will read the row offsets to determine the actual number of received tokens in the input tensor.
• ​fused_silu_mx_kernel: This kernel performs the SILU operation for all the MLPs in the EP MoE module. We need to manually implement the kernel here is because after the EP dispatch phase, the actual number of received tokens is not known to the host. This kernel will read the row offsets to determine the actual number of received tokens in the input tensor.
• ​fused_silu_mxfp6_kernel: SwiGLU + MXFP6 quantization for the EP MoE down-projection input.
• ​fused_silu_mxfp8_interleaved_kernel: SwiGLU + MXFP8 quantization for interleaved gate/up layout.
• ​fused_silu_nvfp4_interleaved_kernel: SwiGLU + NVFP4 quantization for interleaved gate/up layout.
• ​fused_silu_nvfp4_kernel: This kernel performs the SILU operation for all the MLPs in the EP MoE module. We need to manually implement the kernel here is because after the EP dispatch phase, the actual number of received tokens is not known to the host. This kernel will read the row offsets to determine the actual number of received tokens in the input tensor.
• ​get_device_alignment: Returns the natural SIMD alignment in bytes for the current GPU target.
• ​sigmoid: