flashinfer.comm.trtllm_custom_all_reduce

flashinfer.comm.trtllm_custom_all_reduce(inp: Tensor, out: Tensor, tp_size: int, tp_rank: int, token_num: int, fusion_op_code: AllReduceFusionOp, strategy_code: AllReduceStrategyType, config_code: AllReduceStrategyConfig, launch_with_pdl: bool, flag_value: int, peer_comm_buffer_ptrs: Tensor, peer_barrier_ptrs_in: Tensor, peer_barrier_ptrs_out: Tensor, bias: Tensor | None, residual: Tensor | None, weight: Tensor | None, weight_pre_residual_norm: Tensor | None, eps: float | None, intermediate_buffer: Tensor | None, lamport_peer_comm_buffer_ptrs_0: Tensor | None, lamport_peer_comm_buffer_ptrs_1: Tensor | None, lamport_peer_comm_buffer_ptrs_2: Tensor | None) → None

Parameters: - inp: the input tensor. [token_num, hidden_dim] - out: the output tensor. [token_num, hidden_dim] - tp_size: the size of the process group. - tp_rank: the rank of the current process. - token_num: the number of tokens in the sequence. - fusion_op_code: the fusion operation code. - strategy_code: the strategy code. - config_code: the config code. - launch_with_pdl: whether to launch with pdl. - flag_value: monotonic counter, same value on all ranks for a given call.

Calls that dispatch to the lamport kernel (see is_lamport_supported in the kernel header) need consecutive flag values across those calls: keep a separate counter for them, or re-run trtllm_lamport_initialize_all before each one. Sharing one counter with other calls silently corrupts the lamport output.

  • peer_comm_buffer_ptrs: the peer communication buffer pointers.

  • peer_barrier_ptrs_in: the peer barrier pointers in.

  • peer_barrier_ptrs_out: the peer barrier pointers out.

  • bias: the bias tensor. [hidden_dim]

  • residual: the residual tensor. [token_num, hidden_dim]

  • weight: the weight tensor. [hidden_dim]

  • weight_pre_residual_norm: the weight pre residual norm tensor. [hidden_dim]

  • eps: the epsilon value.

  • intermediate_buffer: the intermediate buffer tensor.

  • lamport_peer_comm_buffer_ptrs_0: the lamport peer communication buffer pointers 0.

  • lamport_peer_comm_buffer_ptrs_1: the lamport peer communication buffer pointers 1.

  • lamport_peer_comm_buffer_ptrs_2: the lamport peer communication buffer pointers 2.