• 高性能计算
• 通信
• 深度学习
• 分布式训练

UB-X（UltraBandwidth — eXperimental）是 NCCL 源码树contrib/nccl_ubx/目录下的一个实验性 GPU 通信库原型，目标是实现最低的集合通信延迟与最高的持续 NVLink 带宽。它的核心思路有两个：其一是计算融合——把集合通信前后的计算工作（残差相加、RMSNorm、mxfp8 量化）直接折叠进通信 kernel；其二是与计算 kernel 的高效重叠，让集合通信能与模型计算在同一块 GPU 上并发运行。读完本文，你将掌握 UB-X 的安装与构建配置、对称内存分配器的使用方式、各集合通信变体的选择策略、MoE token dispatch/combine 的完整调用链，以及基准测试与测试套件的运行方法。

API 稳定性提示（原文档声明）：UB-X 是新的代码，公开 API 尚处于演进期，在未来一段时间内不保证稳定。

1. 设计核心：对称内存 + NCCL Device API

UB-X 区别于传统 host 驱动 ring 算法的关键，在于它使用NCCL 对称内存设备 API让集合通信以单 kernel、device-side 原语（multimem 多播、Lamport 轮询、UC peer 写）的形式执行，而非由 CPU 侧代理线程驱动多轮通信。

从源码结构看，这一设计贯穿整个调用链：

• 池分配层 NcclSymPool 基于 nccl4py 封装了ncclMemAlloc+ncclCommWindowRegister+ncclDevCommCreate的完整生命周期；
• CUDA kernel 层 ubx.cu 直接包含nccl_device.h，在 kernel 内部调用ncclGetLsaPointer解析 peer 指针、调用ncclGetLsaMultimemDevicePointer解析多播指针，并用内联 PTX 发射multimem.ld_reduce/multimem.st/multimem.red.add指令（见 ubx.cu 的 PTX 宏定义）；
• 上层 SymmAllocator 把整个池暴露为 NCCL 注册窗口，分配出的SymmTensor都是普通torch.Tensor子类，可无缝接入现有 PyTorch 代码，实现零拷贝的集合通信输入/输出。

池的内部布局为[ REG0 | graph_pool | non_graph_pool ]：REG0 持有 NCCL barrier 标志区与跨 rank 同步 flag 槽位，graph_pool 为 CUDA graph 重放提供地址稳定的分配区，non_graph_pool 服务 eager 分配（见 allocator.py 的池布局注释）。这一点在实操中非常关键——任何与 REG0 重叠的分配都会静默破坏跨 rank 同步协议。

2. 硬件与软件前提

硬件要求（原文档明确列出）：

• GPU：SM 9.0+——Hopper（H100、H200）、Blackwell（B200、GB200、GB300）。MC（multicast）kernel 路径要求 NVLink multicast 硬件支持。
• SM 8.0（A100）不受支持——Ampere 没有 NVLink multicast 硬件，内联 PTX（multimem.*）无法为 arch8.0汇编。

软件要求：

setup.py 显式链接libnccl（libraries=["nccl"]），并注释说明了为何不手动链接-lcuda/-lnvidia-ml（Docker 环境中libcuda.so常以 stub 形式存在于非常规库路径，会导致链接失败）。

3. 安装与编译选项

在 contrib/nccl_ubx/ 目录下执行：

pip install -e . # 默认: SM 9.0a;10.0a pip install -e ".[dev]" # 额外安装 pytest, ruff TORCH_CUDA_ARCH_LIST="9.0a;10.0a" pip install -e . # 显式指定架构列表

同一条命令会同时安装基准测试套件ubx_bench（来自树内 bench/ 目录）。pyproject.toml 中通过tool.setuptools.packages.find将ubx（根目录）与ubx_bench（bench/下）作为平级包一起打包，二者 import 路径相互独立。

构建时环境变量（由 setup.py 消费）：

从源码可进一步确认两个文档表中未列出的构建开关（见 setup.py）：

• UBX_BUILD_TIMEOUT=1会向 cxx/nvcc 同时注入-DUB_TIMEOUT_ENABLED宏；ubx.cu 中定义了__device__ __constant__超时计数器，host 端通过ubx_set_timeout()从UBX_TIMEOUT_SEC覆写其值；
• UBX_BUILD_SKIP_GRID_DEP_SYNC=1可编译掉 kernel 中的cudaGridDependencySynchronize()（PDL 上游依赖等待）——某些 driver/runtime 组合下即使 PDL 未启用也可能在该 intrinsic 上卡住。

若希望 NCCL、libnccl 与 UB-X 来自同一份源码树（保证 device API 特性不缺失），仓库提供了 Dockerfile：它从单一 NCCL 检出同时构建libnccl.so和 UB-X，并支持linux/arm64与linux/amd64双平台。

4. 两种使用方式：便捷 API 与直接分配器 API

4.1 便捷 API（ubx.ops）

全局注册表把每个torch.distributedprocess group 映射到一个惰性创建的SymmAllocator（实现见 ops.py）。你先声明最大 tensor 形状，之后按需取对称 tensor、调用自由函数跑集合通信，适合嵌入现有模型的 drop-in 用法：

import torch import torch.distributed as dist from ubx.ops import request_allocator, get_sym_tensor, allreduce # 模型初始化阶段：声明分配器所属 process group + 最大形状 dist.init_process_group(backend="nccl") group = dist.group.WORLD request_allocator(group, shape=(1024, 2048), dtype=torch.bfloat16) # 前向传播阶段：分配对称 tensor 并执行集合通信 tensor = get_sym_tensor((1024, 2048), torch.bfloat16, group) tensor.fill_(1.0) output = allreduce(tensor) # 按 size 自动选择 Lamport / MC

实现细节值得注意：request_allocator是幂等的，多次调用会累积最大 tensor 尺寸；get_sym_tensor首次调用时才真正构造SymmAllocator，此时池大小取环境变量UBX_SYMM_POOL_SIZE，或默认启发式6 × max_tensor_size（见 ops.py）。此外该路径目前仅支持 bf16，其他 dtype 返回None以便回退 NCCL。

4.2 直接分配器 API（ubx.SymmAllocator）

自行构造 allocator、显式控制池大小、生命周期或多 group 混用时使用：

from ubx import SymmAllocator allocator = SymmAllocator(pool_bytes, device, dist.group.WORLD) sym = allocator.create_tensor(shape, torch.bfloat16) out = allocator.allreduce(sym)

构造是集合通信操作：dist_group内每个 rank 必须用相同的size_bytes调用。从 allocator.py 可以看到构造流程：

1. 池大小向上取整到 2 MB 对齐（NCCL 以 2 MB 页分配）；
2. 优先复用torch.distributedprocess group 底层的ncclComm_t（get_or_create_nccl_comm 会探测多个属性名），失败则用 rank 0 广播 NCCL unique ID 的方式 bootstrap 一个新 communicator；
3. NcclSymPool完成ncclMemAlloc+ 窗口注册 + devcomm 创建；
4. 默认执行 NCCL prewarm（UBX_PREWARM_NCCL=1）：提前跑一次all_gather_into_tensor与all_to_all_single，让各 op 的懒加载 channel 初始化发生在构造期而非首次运行时调用——文档注释指出，不做这一步在 NVLS 禁用模式下可能触发 10 分钟 watchdog 超时。

常用属性：

4.3SymmTensor：零拷贝的 torch.Tensor 子类

SymmTensor 是对池内 uint8 字节段做view(dtype)+ reshape 得到的torch.Tensor子类，带引用计数：最后一个 view 被丢弃时内存归还 free list（first-fit 分配，相邻空闲段在释放时合并）。特殊语义：

• blocked='mxfp8'时，池内预留对齐的数据区后跟 E8M0 scale 字节（每 32 个元素 1 字节）；tensor 的 shape/dtype 只描述数据，scale 通过tensor.metadata_ptr访问（实现见 tensor.py）；
• view()返回同内存的SymmTensor视图；clone()回退为普通torch.Tensor（clone 不再对称）；blocked tensor 不支持view()；
• ops.restore() 可把data_ptr落在池内的普通 tensor 重新包装为SymmTensor。

5. 支持的集合通信与自动选择策略

0.25 MB 的切换阈值在 allocator.py 中定义为AUTO_SWITCH_BYTES = 256 * 1024，SymmAllocator.allreduce()/alltoall_auto()均按此阈值在稳态最快变体间切换；直接调用显式方法（allreduce_mc、alltoall_lamport等）则绕过自动选择。各变体的适用场景（结合 docs/api.md）：

• AllReduce：全部为 in-place 操作。allreduce_mc强制多播，适合大 tensor；allreduce_lamport强制 Lamport 轮询，适合小 tensor（池耗尽时回退 MC，multicast 不可用时回退 UC）；allreduce_uc为无多播 fabric 的软件归约回退路径。
• AllToAll：alltoall(t)为 UC + barrier；alltoall_lamport(t)采用三缓冲、稳态无 barrier，小尺寸下比 UC 快 1.5–2×。注意alltoall的smlimit是精确设定 SM 数（便于上下扫描），而alltoall_lamport的smlimit是上限（与 allreduce 语义一致）。
• AllToAllV（变长）提供两步热路径 API：alltoallv_prepare(output_split_sizes, input_split_sizes, dtype)做一次all_gather计算远端接收偏移并分配输出缓冲，之后可重复alltoallv_run(input_tensor, state)；split 尺寸每次都变时用一次性封装alltoallv()。

计算融合：allreduce_mc()/allreduce_lamport()接受可选gamma、residual_in、residual_out、eps、hidden_size参数，把残差相加 + RMSNorm 折叠进同一 kernel。fused.py 提供可直接调用融合参数的包装（该模块是可拆卸的——删掉它不影响核心集合通信 API）：

from ubx.fused import allreduce_fused out = allreduce_fused( allocator, sym, hidden_size, residual_in=residual, # 可选 residual_out=new_residual, # 可选 gamma=rmsnorm_weight, # 与 eps 同时提供时启用 RMSNorm 融合 eps=1e-6, )

便捷层ubx.ops.allreduce在传入gamma/eps时也会自动启用融合，并在 allocator 上维护内部 residual 缓冲（跨 transformer FFN/attention 块复用），配free_residual(sym)释放。ops.allreduce 的注释还给出了四种子模式：独立 allreduce、首层 PROJ（融合 layernorm + 全局 residual 输入）、中间层（内部 residual 进出）、末层 FC2（无 layernorm、无 residual 输出）。

6. MoE Token Dispatch 与 mxfp8 量化

a2av_token_bf16_mxfp8用单个 GPU kernel完成两件事：把 bf16 token 路由到远端 rank，同时在传输路上把数据量化为 mxfp8（每 32 个元素一个 E8M0 scale）。配合compute_token_offsets把路由矩阵转换为 kernel 期望的 slot 分配 tensor：

from ubx import compute_token_offsets, SymmAllocator token_offsets, max_tokens_per_rank, _, _ = compute_token_offsets( routing, # [ntokens, total_experts] (uint8/bool/int) experts_per_rank, myrank=rank, nranks=world_size, ) output = allocator.create_tensor( [max_tokens_per_rank, hidden], torch.float8_e4m3fn, blocked="mxfp8", ) output = allocator.a2av_token_bf16_mxfp8( tokens_bf16, token_offsets, experts_per_rank, output, )

要点（结合 docs/api.md 与 ops.py 的 docstring）：

• hidden必须是 32 的倍数（一个 mxfp8 block）；量化误差约为max_block_scale × 2^-4；调用后 fp8 数据在output.data_ptr()、E8M0 scale 字节在output.metadata_ptr；
• routing矩阵要求所有 rank 产出完全相同的矩阵（用共享 seed 或显式 broadcast），ntokens必须能被nranks整除，total_experts = nranks × experts_per_rank；token 按块分布：rank r 拥有全局 token 行[r*local_n, (r+1)*local_n)；
• 异步调度：a2av_token_bf16_mxfp8(..., sync=False)让 kernel 在 barrier flag 置位后立即返回，读取输出前必须调用a2av_wait()；
• 流水线调度：a2av_token_bf16_mxfp8_persistent(..., nexperts_per_chunk=2)用一个 kernel 内执行nchunks = ceil(experts_per_rank / nexperts_per_chunk)个内部 chunk（nchunks上限 32），每个 chunk 自带与同步版一致的 barrier 协议，适合下游是 per-expert GEMM 的场景；
• bf16→bf16 变体：a2av_token_bf16_bf16走同样的路由/slot 方案但不量化，输出必须是普通（非 blocked）bf16SymmTensor。

Token combine（dispatch 的逆过程）：每个 rank 持有自己收到的 token 的 FFN expert 输出，combine 把结果路由回源 rank 并做 gate 加权求和，产出[local_ntokens, hidden]bf16。两大族：

• PULL-barrier（combine_bf16_bf16、combine_mxfp8_bf16）：Phase 1 各 rank 把 expert 输出暂存到临时对称缓冲 → 全体 barrier → Phase 2 各 rank 拉取自己 token 的 peer 贡献、gate 加权、fp32 累加、下转。复用 dispatch 的token_offsets矩阵；combine_mxfp8_bf16在跨 rank 交换前把 expert 输出量化为 mxfp8（NVLink 载荷减半），累加与返回仍保持 bf16。支持sync=False两阶段异步模式（combine_wait()执行 Phase 2）；
• PUSH（combine_bf16_bf16_push、combine_bf16_bf16_lamport_push）：各 rank 直接把 FFN 输出推到 peer 目的缓冲的已知(origin_token, k_idx)坐标，接收端只读自己的缓冲——与alltoall_lamport相同的无竞态模式。需要compute_combine_push_map构造的逆路由表。PUSH 变体目前不提供sync=False异步形式。

from ubx import compute_combine_push_map inverse_map, topk_idx, max_tokens_per_rank = compute_combine_push_map( routing, experts_per_rank, myrank=rank, nranks=world_size, ) out = allocator.combine_bf16_bf16_lamport_push( expert_outputs, inverse_map, topk_idx, experts_per_rank, max_tokens_per_rank, gate_weights=g, )

gate_weights（可选）是完整的[local_ntokens, total_experts]fp32 router 输出，kernel 按贡献逐条索引；传None表示无权求和。

7. 端到端示例

examples/ 目录包含使用公开 API 的最小多 GPU 程序，均可用两台以上 GPU 的单机直接启动：

torchrun --nproc-per-node=2 examples/allreduce.py torchrun --nproc-per-node=2 examples/alltoall.py torchrun --nproc-per-node=2 examples/moe_dispatch.py

以 allreduce.py 为例，它的验证方式是：用按 rank 区分的随机数据填充对称 tensor，跑 UB-X auto allreduce，再对参考值用 fp32 的 NCCLall_reduce（规避 NCCL bf16 ring 归约的累加顺序误差），最后torch.testing.assert_close(out.float(), ref, atol=0.0625, rtol=0.02)比对。alltoall 示例中，65536 个 bf16 元素（128 KB）会命中 Lamport 路径；把元素数提到约 256K（bf16 下 0.5 MB）即可观察到切换到 UC 路径。

8. 基准测试：ubx_bench

ubx_bench提供nccl-tests 兼容的 CLI，可把 UB-X 与 NCCL 的集合通信并排测量，输出列（size、count、type、redop、time、algbw、busbw）可直接粘贴进既有对比表。CUDA graph 模式是默认（-G 10000次图重放），需要把每次 op 的 launch + sync 开销也计入时传--no-cudagraph：

# AllReduce: UB-X vs NCCL torchrun --nproc-per-node=8 -m ubx_bench all_reduce \ -b 1K -e 128M -f 2 --backend ubx,nccl -n 5 -w 2 # AllToAll torchrun --nproc-per-node=8 -m ubx_bench all_to_all \ -b 1K -e 8M -f 2 --backend ubx,nccl -n 5 -w 2

完整 CLI 参考与基准环境说明见 bench/README.md。常用 flag（原文档表格继承）：

可用子命令包括all_reduce、all_to_all、all_gather、alltoallv、reduce_scatter（仅 nccl 后端，无 UB-X kernel）、a2av_mxfp8、a2av_dispatch、a2av_combine。带宽校正系数与 nccl-tests 一致：AllReduce 为2×(n-1)/n，AllGather/ReduceScatter/AllToAll/AllToAllV/A2AV 为(n-1)/n。基准环境建议：CUDA_DEVICE_MAX_CONNECTIONS=1（降低 graph 模式下的 GPU 调度开销）；UBX_GRAPH_POOL_SHARE由 bench 按模式自动设置（graph 模式0.5、eager 模式0.1）。

9. 测试套件与正确性验证

# 单 GPU 单元测试（allocator、tensor、backend 检测） pytest tests/ # 多 GPU 正确性测试（需要 2 块以上 GPU） torchrun --nproc-per-node=2 -m pytest tests/distributed/

• 单 GPU 测试位于 tests/，覆盖分配器池管理（如 test_allocator_pool.py，专门回归 REG0 与 graph 池越界破坏 BAR 计数器的问题）、API 注册表、backend 检测、路由计算；
• 分布式测试采用子进程模式（见 tests/distributed/conftest.py）：外层 pytest 调用run_distributed_test(script, num_procs, args)，优先用srun（自动设置 SLURM 变量），回退torchrun（正确处理env://rendezvous；刻意避免mpirun，因为它不设置env://所需的变量）；worker 脚本位于tests/distributed/_workers/，执行 op 后与torch.distributed（NCCL）参考输出比对；
• 容差：bf16 为atol=0.0625, rtol=0.02；FP8 为atol=0.0625, rtol=0.125。

10. CUDA Graph 支持

池的 graph 区域在重放之间保持地址稳定：在cudaStreamBeginCapture之前分配好要捕获的 tensor，捕获期间执行集合通信，之后同一个SymmTensor对象可无限次重放。Lamport allreduce 与 alltoall 路径为 graph 与 eager 模式分别维护独立的 "next poison" 与三缓冲状态，二者互不干扰（allocator 中lamport_a2a_state等状态字典按 graph/eager 分键）。

11. 环境变量总览

docs/api.md 还列出了若干进阶变量：UBX_MAXSM（0–128，覆写每个 kernel 的 SM grid，0 为编译默认）、UBX_DUMMY（1时绕过 UB-X 原样返回输入）、UBX_PREWARM_NCCL/UBX_PREWARM_NCCL_ELEMS（控制构造期 NCCL 预热）、UBX_FORCE_BOOTSTRAP（跳过 PG handle 复用、强制 uid broadcast bootstrap，用于多 PG 环境怀疑复用了错误 handle 时）。

12. 常见陷阱（来自官方 API 指南）

docs/api.md 的 "Common pitfalls" 一节值得完整保留，它们都对应源码中真实的修复历史：

• 池过小：Lamport allreduce/alltoall 稳态会同时保持两三个输入形状的缓冲存活，池大小至少取6 × max_tensor_size（默认启发式）；
• 在不支持 multicast 的 fabric 上未禁用它：mc0_ptr会是 0，但 NCCL 仍会尝试 MC rendezvous——必须在构造 allocator之前设置NCCL_NVLS_ENABLE=0（该信号在 allocator.py 中统一接入NCCLDevCommRequirements.lsa_multimem与 NCCL 自身集合通信）；
• mxfp8 dispatch 中hidden % 32 != 0：向上 padding 到 32 的倍数或在上游 reshape；
• 各 rank 的routing矩阵不一致：compute_token_offsets假定所有 rank 产出相同矩阵，务必用共享 seed 或显式 broadcast；
• sync=True之后调用a2av_wait()：无害但多余，kernel 已轮询到完成；
• kernel 仍在运行时丢弃最后一个SymmTensorview：引用计数是 host 侧的，对生命周期敏感的代码应依赖torch.cuda.synchronize()或流排序后再释放。

13. 目录结构与 API 版本治理

仓库结构（contrib/nccl_ubx/）：

ubx/ Python 包（allocator、tensor、ops、fused） csrc/ C++/CUDA kernel 与 pybind11 绑定 tests/ 单 GPU + 分布式测试套件 bench/ 树内 ubx_bench 基准测试框架 examples/ 使用公开 API 的最小可运行程序 docs/ 面向用户的 API 文档（api.md）

由于 API 处于快速演进期，UB-X 用一套显式的 API 契约机制管理漂移（见 api.version.md 与 ubx/_api_registry.py）：ubx.version是 API 契约版本（区别于包版本ubx.__version__，当前为 1.1.0），下游代码可用ubx.query_api(name, params=[...])检查自己的调用是否仍在当前签名内（返回ok/removed/renamed/params_changed及人类可读诊断），用ubx.query_supported_hw()查询本构建可用的 transport 与最低 SM 架构。注册表是一个 append-only 的事件日志，模块加载时重放得到当前公开面，因此能回答"你的参数是在哪个版本被改名的"这类问题。

小结

UB-X 展示了 NCCL device API 的一种激进用法：把池、窗口、devcomm 的生命周期交给SymmAllocator一次性建好，之后所有集合通信都变成单 kernel 的 device-side 多播/轮询/peer 写操作，并借机把残差、RMSNorm、FP8 量化等前后处理融进通信 kernel，配合 CUDA graph 池区实现可重放的低延迟路径。需要再次强调的是：它要求 SM 9.0+ 硬件（MC 路径额外要求 NVLink multicast）、NCCL ≥ 2.30 与配套构建的 nccl4py，且 API 尚不稳定——建议在生产前用ubx.query_api/ubx.query_supported_hw做漂移检测，并用 examples/ 与 tests/ 作为升级后的回归基线。

• 高性能计算
• 通信
• 深度学习
• 分布式训练