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) → None

参数: - world_size: 进程组的大小。 - world_rank: 当前进程的等级。 - token_num: 序列中的 token 数量。 - hidden_dim: 隐藏状态的维度。 - workspace_ptrs: 工作区指针。 - launch_with_pdl: 是否使用 pdl 启动。 - residual_in: 残差输入张量。 [token_num, hidden_dim] - rms_gamma: rms gamma 张量。 [hidden_dim] - rms_eps: rms epsilon 值。 - scale_factor: 缩放因子。 - moe_reduction_device_num_experts: 专家数量。 - moe_reduction_scale_input: 缩放输入张量。 [token_num, hidden_dim] - moe_reduction_active_experts_token_input: 活动专家 token 输入张量。 [token_num, hidden_dim] - moe_reduction_token_input: token 输入张量。 [token_num, hidden_dim] - layout_code: 布局代码。 - moe_allreduce_out: moe allreduce 输出张量。 [token_num, hidden_dim] - residual_out: 残差输出张量。 [token_num, hidden_dim] - norm_out: norm 输出张量。 [token_num, hidden_dim] - quant_out: quant 输出张量。 [token_num // 4, hidden_dim], fp16/bf16 -> fp4 - scale_out: 缩放输出张量。 初始化参考: tests/comm/test_trtllm_moe_allreduce_fusion.py