For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

ep_combine_wait_kernel_api

def ep_combine_wait_kernel_api[combine_dtype: DType, //, hidden_size: Int, top_k: Int, n_experts: Int, max_token_per_rank: Int, n_gpus_per_node: Int, n_nodes: Int, target: StringSlice[ImmStaticOrigin], router_weights_wrapper: Optional[def[width: Int](token_idx: Int, topk_id: Int) capturing thin -> SIMD[DType.float32, width]] = None, epilogue_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None](output_tokens: TileTensor[combine_dtype, address_space=output_tokens.address_space, linear_idx_type=output_tokens.linear_idx_type], atomic_counters: TileTensor[DType.int32, address_space=atomic_counters.address_space, linear_idx_type=atomic_counters.linear_idx_type], recv_ptrs: TileTensor[DType.uint64, address_space=recv_ptrs.address_space, linear_idx_type=recv_ptrs.linear_idx_type], recv_count_ptrs: TileTensor[DType.uint64, address_space=recv_count_ptrs.address_space, linear_idx_type=recv_count_ptrs.linear_idx_type], context: DeviceContext, num_input_tokens: Int = Int(-1))

Execute the Expert Parallelism combine completion kernel.

This function launches the combine_wait_kernel from ep_comm.mojo to complete the token combine phase. It waits for all inter-device communication to complete, then computes the weighted sum of routed expert outputs for each token.

Arguments: output_tokens: Final output tensor with expert results. atomic_counters: EP kernel synchronization counters. Used to coordinate between different thread blocks. recv_ptrs: Receive buffer pointers for each local GPU. recv_count_ptrs: Receive count buffer pointers for each local GPU. context: Device context pointer. num_input_tokens: Per-rank input token count for this layer. When >= 0 enables the decode-fast-path grid sizing. Default -1 keeps the full-sm_count grid for backwards compatibility.

Parameters:

combine_dtype (DType): Data type for tokens during combine phase.
hidden_size (Int): Model hidden dimension size.
top_k (Int): Number of experts each token was routed to.
n_experts (Int): Total experts across all devices.
max_token_per_rank (Int): Maximum tokens any device can receive.
n_gpus_per_node (Int): GPUs per physical node.
n_nodes (Int): Number of physical nodes.
target (StringSlice[ImmStaticOrigin]): Target.
router_weights_wrapper (Optional[def[width: Int](token_idx: Int, topk_id: Int) capturing thin -> SIMD[DType.float32, width]]): Wrapper for the optional router weights.
epilogue_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional elementwise epilogue function applied after computing combined output.