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 struct
DistributedEPDispatchNVFP4
struct DistributedEPDispatchNVFP4
Implemented traits
AnyType,
ImplicitlyDestructible
Methods
execute
static def execute[input_dtype: DType, dispatch_dtype: DType, dispatch_scale_dtype: DType, hidden_size: Int, top_k: Int, n_experts: Int, max_token_per_rank: Int, n_gpus_per_node: Int, n_nodes: Int, fused_shared_expert: Bool, //, target: StringSlice[StaticConstantOrigin], _trace_name: StringSlice[StaticConstantOrigin]](output_tokens: VariadicTensors[Output, static_specs=output_tokens.static_specs], output_scales: VariadicTensors[Output, static_specs=output_scales.static_specs], row_offsets: VariadicTensors[Output, static_specs=row_offsets.static_specs], scales_offsets: VariadicTensors[Output, static_specs=scales_offsets.static_specs], expert_ids: VariadicTensors[Output, static_specs=expert_ids.static_specs], src_info: VariadicTensors[Output, static_specs=src_info.static_specs], input_tokens: VariadicTensors[Input, static_specs=input_tokens.static_specs], topk_ids: VariadicTensors[Input, static_specs=topk_ids.static_specs], send_ptrs: VariadicTensors[Input, static_specs=send_ptrs.static_specs], recv_ptrs: VariadicTensors[Input, static_specs=recv_ptrs.static_specs], recv_count_ptrs: VariadicTensors[Input, static_specs=recv_count_ptrs.static_specs], input_scales: VariadicTensors[Input, static_specs=input_scales.static_specs], atomic_counters: VariadicTensors[MutableInput, static_specs=atomic_counters.static_specs], dev_ctxs: DeviceContextList)
Multi-device fused Expert Parallelism NVFP4 dispatch.
Launches the EP dispatch kernel on all devices simultaneously via _ep_launch_device_collective. Each device routes its tokens to experts based on top-k IDs, quantizes them to NVFP4 format, and sends them to the appropriate peer devices.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!