flashinfer.comm.trtllm_moe_allreduce_fusion

flashinfer.comm.trtllm_moe_allreduce_fusion(world_size: int, world_rank: int, token_num: int, hidden_dim: int, workspace_ptrs: Tensor, launch_with_pdl: bool, residual_in: Tensor, rms_gamma: Tensor, rms_eps: float, scale_factor: float, moe_reduction_device_num_experts: int, moe_reduction_scale_input: Tensor, moe_reduction_active_experts_token_input: Tensor, moe_reduction_token_input: Tensor, layout_code: QuantizationSFLayout | None, moe_allreduce_out: Tensor | None, residual_out: Tensor | None, norm_out: Tensor | None, quant_out: Tensor | None, scale_out: Tensor | None, weight_bias: float | None = None, *, backend: Literal['trtllm', 'cake'] = 'trtllm') → None

Parameters: - world_size: the size of the process group. - world_rank: the rank of the current process. - token_num: the number of tokens in the sequence. - hidden_dim: the dimension of the hidden states. - workspace_ptrs: the workspace pointers. - launch_with_pdl: whether to launch with pdl. - residual_in: the residual input tensor. [token_num, hidden_dim] - rms_gamma: the rms gamma tensor. [hidden_dim] - rms_eps: the rms epsilon value. - scale_factor: the scale factor. - moe_reduction_device_num_experts: the number of experts. - moe_reduction_scale_input: the scale input tensor. [token_num, hidden_dim] - moe_reduction_active_experts_token_input: the active experts token input tensor. [token_num, hidden_dim] - moe_reduction_token_input: the token input tensor. [token_num, hidden_dim] - layout_code: the layout code. - moe_allreduce_out: the moe allreduce output tensor. [token_num, hidden_dim] - residual_out: the residual output tensor. [token_num, hidden_dim] - norm_out: the norm output tensor. [token_num, hidden_dim] - quant_out: the quant output tensor. [token_num // 4, hidden_dim], fp16/bf16 -> fp4 - scale_out: the scale output tensor. Initialization referece: tests/comm/test_trtllm_moe_allreduce_fusion.py - weight_bias: bias added to rms_gamma before scaling.

None or 0.0 -> standard RMSNorm (out = gamma * x * rsqrt(…)). 1.0 -> Gemma / Qwen3.5 RMSNorm (out = (1 + gamma) * x * rsqrt(…)).

• backend: "trtllm" (default) or the constrained "cake" SM100/SM103 backend. The public Python API has 22 parameters; the isolated source module’s run_reduction entry has an exact 18-argument FFI ABI. The optional backend supports contiguous FP16/BF16 tensors, world sizes 2, 4, and 8, hidden_dim=7168, token payloads within the existing Lamport MAX_COMM_SIZE byte limit, and residual plus norm outputs. It does not support quantization. weight_bias remains a runtime value; None is passed to the kernel as 0.0. On SM100 and SM103 with world_size 2, 4 or 8 and an moe_allreduce_out tensor, "cake" runs the verified source export of the Cake all-reduce union (flashinfer.jit.cake_trtllm_moe_allreduce_union); the route binds the workspace pointer table the way trtllm_create_ipc_workspace_for_all_reduce_fusion registers it and needs no device readback. Every other Cake configuration keeps the isolated source bundle.