flashinfer.comm

此模块提供分布式计算的通信原语和实用工具，包括 CUDA IPC、AllReduce 操作和内存管理实用工具。

CUDA IPC 实用工具

CudaRTLibrary([so_file])

create_shared_buffer(size_in_bytes[, group])

创建一个共享缓冲区，并返回一个指针列表，该列表表示组中所有进程的缓冲区。

free_shared_buffer(pointers[, group])

释放一个共享缓冲区。

DLPack 实用工具

pack_strided_memory(ptr, segment_size, ...)

将 GPU 内存打包到具有指定步长的 PyTorch 张量中。

映射实用工具

Mapping([world_size, rank, gpus_per_node, ...])

一个具有 8 个 GPU 的节点，tp_size = 4，cp_size = 1，pp_size = 2

TensorRT-LLM AllReduce

类型和枚举

AllReduceFusionOp()

AllReduceFusionPattern()

AllReduceStrategyConfig()

AllReduceStrategyType()

核心操作

trtllm_allreduce_fusion(allreduce_in, ...[, ...])

参数：- allreduce_in：输入张量。[token_num, hidden_dim] - world_size：进程组的大小。 - world_rank：当前进程的等级。 - token_num：序列中的 token 数量。 - hidden_dim：隐藏状态的维度。 - workspace_ptrs：工作区指针。 - launch_with_pdl：是否使用 pdl 启动。 - use_oneshot：是否使用 oneshot。如果为 None，将使用内部启发式方法。 - trigger_completion_at_end：是否在末尾触发完成。 - fp32_acc：是否使用 fp32 累积。 - pattern_code：模式代码。 - allreduce_out：输出张量。[token_num, hidden_dim] - residual_in：残差输入张量。[token_num, hidden_dim] - residual_out：残差输出张量。[token_num, hidden_dim] - norm_out：归一化输出张量。[token_num, hidden_dim] - quant_out：量化输出张量。[token_num, hidden_dim] - scale_out：缩放输出张量。 初始化参考：tests/comm/test_trtllm_allreduce_fusion.py - rms_gamma：rms gamma 张量。[hidden_dim] - rms_eps：rms epsilon 值。 - scale_factor：缩放因子。为了 cudaGraphs 的安全性，它应该是一个张量。 - layout_code：布局代码。 - metadata：来自 create_ipc_workspace_for_all_reduce_fusion 的可选工作区元数据。如果提供，则验证 token_num <= max_token_num、world_size == tp_size 且 hidden_dim == workspace hidden_dim。如果验证失败，则引发 ValueError。

trtllm_custom_all_reduce(inp, out, tp_size, ...)

参数：- inp：输入张量。

trtllm_moe_allreduce_fusion(world_size, ...)

参数：- world_size：进程组的大小。

trtllm_moe_finalize_allreduce_fusion(...)

参数：- allreduce_in：输入张量。

工作区管理

trtllm_create_ipc_workspace_for_all_reduce(...)

参数：- rank：当前进程的等级。

trtllm_create_ipc_workspace_for_all_reduce_fusion(...)

参数：- tp_rank：当前进程的等级。

trtllm_destroy_ipc_workspace_for_all_reduce(...)

注意：此函数用于销毁 all reduce 的工作区。

trtllm_destroy_ipc_workspace_for_all_reduce_fusion(...)

参数：- workspace：要销毁的工作区。

初始化和实用工具

trtllm_lamport_initialize(buffer_ptr, size, ...)

trtllm_lamport_initialize_all(buffer_0_ptr, ...)

使用负零初始化 3 个 lamport 缓冲区。

compute_fp4_swizzled_layout_sf_size(...)

计算 fp4 交错布局的填充大小的辅助函数。

vLLM AllReduce

vllm_all_reduce(fa, inp, out, reg_buffer, ...)

执行非就地 all reduce。

vllm_dispose(fa)

vllm_init_custom_ar(ipc_tensors, rank_data, ...)

vllm_register_buffer(fa, fake_ipc_ptrs)

vllm_register_graph_buffers(fa, handles, offsets)

vllm_get_graph_buffer_ipc_meta(fa)

vllm_meta_size()

TensorRT-LLM MNNVL AllReduce

trtllm_mnnvl_all_reduce(inp, ...[, out])

在多个 GPU 上执行多节点 NVLink all-reduce 操作。

trtllm_mnnvl_fused_allreduce_rmsnorm(...)

执行 MNNVL TwoShot AllReduce + RMSNorm。

mpi_barrier()

MNNVL A2A（吞吐量后端）

MoeAlltoAll(mapping, max_num_tokens, top_k, ...)

使用适当的工作区分配和同步管理 MoE All-to-All 操作。

moe_a2a_initialize(workspace, ep_rank, ...)

moe_a2a_dispatch(token_selected_experts, ...)

将 token 和 payload 分派到专家等级。

moe_a2a_combine(payload, local_num_tokens, ...)

moe_a2a_sanitize_expert_ids(expert_ids, ...)

moe_a2a_get_workspace_size_per_rank(ep_size, ...)

获取 MoE AlltoAll 操作的每个等级的工作区大小。

moe_a2a_wrap_payload_tensor_in_workspace(...)

将工作区中的偏移量包装到张量中。