Skip to content

Repository files navigation

Ulysses GEMM + All-to-All Fusion

v16.0 — 大尺寸长序列 benchmark

仅更新 SM103 BF16 Graph benchmark、测量入口及结果,算子实现保持 v15.0。 64K/128K/256K/512K × CP4/8:自研融合56/80项、336个有限候选;其余因 未适配或显存限制留空。主对照为同尺寸纯 cuBLASLt,保留比例不是硬件峰值MFU。 已有 TEUB/NCCL 数据保留;KDA/MLA 特殊路由暂缓,不发布中间调优文件。

完整表、配置、样本与复现说明

v15.0 — 可选的 SM103 QKV rank 错峰

新增 FUSE_SM103_QKV_RANK_SWIZZLE 构建选项,默认 OFF。开启后,GEMM 生产顺序与 TMA 消费顺序共同按源 rank 轮转 N 分组;不改变数据布局、head 归属或 ready 同步语义。不是自动调优,也不按模型名自动开启。

同节点、同配置、无 profiling 的 Graph A/B 完成 95/96 个 QKV 点:全量 几何平均 +0.61%,长序列 47/48 点 +1.29%;有增有减,因此不统一开启。 Llama405B CP4 512K 保留显存缺项。SM90 算法与公开参数布局不变。

# 添加到原有 SM103 CMake 配置;关闭或省略该选项即使用原遍历顺序。
cmake -S . -B build/sm103-rank-swizzle -DFUSE_ARCH=sm103 \
  -DCUTLASS_ROOT=/path/to/cutlass -DFUSE_SM103_QKV_RANK_SWIZZLE=ON \
  -DFUSE_ENABLE_PROFILING=OFF
cmake --build build/sm103-rank-swizzle -j 4

目前不支持 rank 错峰与 profiling 同时开启;配置阶段明确拒绝此组合。 默认路径的 Perfetto 导出新增搬运源/目的 GPU 与字节数元数据。 完整 A/B 表、配置、样本与复现说明

v14.0 — Blackwell SM103 BF16 基线

新增 B300 / SM103a 的自研 QKVProj→A2A、A2A→OProj BF16 前向融合算子: CUTLASS Blackwell GEMM、持久化 kernel、计算/通信 CTA 功能特化、GMEM ready 同步及宏控制 profiling。源码与 benchmark 按 sm90/sm103/ 分桶;下方 历史优化和 FP8/backward 性能承诺仍属于 SM90,不扩展为 Blackwell 已实现功能。

本版定位为基线,不宣称优化收口或全局最优。Graph 全量 190/192 项通过, Llama405B CP4 512K 两方向因现有显存余量不足留空。长序列 94/96 项: QKV 几何平均 1.198 PFLOPS/卡、相对 TEUB 1.278×(47 项);OProj 为 1.114 PFLOPS/卡、相对 TEUB 1.058×(仅 23 个同口径匹配项)。 缺少基线的项不填、不跨节点配对;固定实测配置不是运行时全局最优策略。

历史 SM90 实现

单机 Ulysses Context Parallel 的 GEMM/All-to-All 融合算子。A2A+O-projection 已完成优化;QKV Projection+A2A 在 v8.0 联合选择通信 CTA 与 GEMM tile,并复用已经算好的自动配置。v9.1 提供两个面向已知锁频差异的 BF16 加权序列算子。v10.0 新增 QKV 和 OProj 的两条 BF16 反向融合路径,同时支持普通同流 B→W 与 ZeroBubble 分离 B/W;v11.0 补齐 PyTorch autograd 正确性、TE 强基线和前后向 trace;v11.2 验证四条融合边界在完整 CUDA Graph 训练中的 E2E 收益;v12.0 为这四条边界增加纯 E4M3 FP8 版本;v13.0 重构 Ulysses 源码和公开头文件层级;v13.1 将数据流原语与 Ulysses 投影语义收成两级编译期契约,设备代码与 ABI 保持不变。

Attention 输出按 head 分片:

[B, H/CP, S, D]
        │ inverse All-to-All
        ▼
[B, S/CP, H, D]
        │ O-projection GEMM
        ▼
[B, S/CP, hidden]

对应 GEMM:

M = B × S / CP
K = H × D
N = hidden

实现

  • 单个 cooperative persistent kernel,通信 CTA 与 CUTLASS GEMM CTA 同时驻留。
  • 通信 CTA 将远端 head shard 直接写入 GEMM 的最终 A layout,并按 peer/K shard 发布 ready epoch。
  • 窄 peer shard 自动将多行写回合并为约8KiB的3D TMA store,宽 shard沿用逐行路径。
  • GEMM CTA 使用 system-scope acquire 消费已到达的数据,省去独立 permutation kernel 和 kernel 间同步。
  • auto 策略在五个成熟 policy 中选择:M64N128M128N128M128N160M128N256 cluster-M2M128N320 cluster-M2;wave 以 cluster 为调度单元,并优先保证 N frontier 不被 wave 边界切开。

QKV Projection+A2A 提供 M128N64/N128/N160/N192/N256/N320 六种 policy。H200 默认 用GEMM wave成本、QKV通信量、NVLink下界和cluster前进代价联合选择通信CTA与tile; 模型名称和逐case winner不参与选择;不支持的shape或布局自动回退成熟策略。 132-SM设备复用H200整波标定;H800的通信估算使用400 GB/s双向NVLink, 其他设备按900 GB/s。完整口径见 benchmarks/QKVproj+a2a/BENCHMARK.md

comm_ctas=0 只在首次遇到某种设备和 shape 时计算通信 CTA 与 tile;后续 Eager 调用直接复用结果,但仍使用本次调用自己的地址和 epoch。这样避免 H800 训练中每层 重复读取设备信息和搜索配置。QKV 与 A2A+OProj 都会等目标显存写完后再发布 ready。

默认通信策略对短中序列使用 comm4,长序列根据 CP、N 和设备 NVLink 带宽选择 comm6/8。实验性全量成本模型可通过 FUSE_A2A_LHS_COMM_POLICY=experimental_model 启用;它默认关闭,尚不属于 Golden。模型口径和边界见 VERSION_HISTORY.md

锁频异构 CP(v9.1)

v9.0 新增独立的 weighted QKV Projection+A2A 与 weighted A2A+OProj;v9.1 保留这套接口和调度模型,并补上跨 shape 的功耗安全边界。调用方在启动前为每个 rank 提供相对 SM、HBM 和 NVLink 能力;规划器以 256-row 对齐的连续 token 区间为单位,联合选择每张卡的行数、通信 CTA 和既有 GEMM tile,使预测最慢 rank 的完成时间最短。它不读取实时频率、模型名、逐 case winner 或外部基线,也不需要调用方提供 alpha。

这不是跨 GPU 动态偷任务:每个 rank 仍只执行一个连续区间和一个 persistent kernel。全局数学结果不变,但每张卡的 local sequence length 可以不同,因此框架必须让同一分区贯穿依赖该序列布局的后续计算。QKV 默认只在全局 S≤16K 时允许重分;更长 QKV 回退原均匀算子。OProj 不使用固定的序列长度限制。

v9.1 的 shape 复核确认:把 MNK 按 H200 的 989 TFLOPS 理论峰值换算后,当单 rank 纯 GEMM 最低时间接近 1 ms 时,1980 MHz 参考卡本身也会撞 700 W 功耗墙并降到约 1500 MHz,此时标称频率比已经失效。v9.1 默认只在实测的 0.75 ms 功耗安全域内重分配,超出后 QKV 和 OProj 都会回退 uniform;若调用方提供同一负载下的有效吞吐比,可显式放开。完整约束、18 点补充结果和复现命令见 benchmarks/heterogeneous_cp/BENCHMARK.md

反向融合(v10.0)

v10 把两个前向融合边界按反向依赖倒过来实现:

QKV B:   Head→Sequence A2A(dQ,dK,dV) -> dX GEMM
QKV W:   dWqkv = dQKVᵀ × X

OProj B: dA GEMM -> Sequence→Head A2A
OProj W: dWo = dYᵀ × saved_A

通信直接读写最终布局,接口不要求框架先做 catindex_selectpermutecontiguous。普通模式在同一 stream 中严格执行 B→W,并用 beta=0 写权重梯度; ZeroBubble 模式把 B 和 W 拆成两个入口,W 用 beta=1 累加 BF16 main_grad。分离 期间调用方必须保留 W 真正需要的两个输入;能延长原 buffer 寿命时无需额外复制, 不能保留时才需要显式 stash。

两条反向算子的参数结构、自动策略和生产入口互相独立,只共用采样脚手架。完整 MNK、生命周期约束、96 行逐点结果和复现命令见 QKV backwardOProj backward

开箱运行

依赖 CUDA 12.8、CMake、Ninja 和 CUTLASS。Golden 使用以下源码版本:

TransformerEngine  a7aec214eb5c3969984a40c3accb6d66987d8f25
git clone https://github.com/CyberSecurityErial/fuse.git
cd fuse

CUTLASS_ROOT=/path/to/TransformerEngine/3rdparty/cutlass bash scripts/build.sh

./build/fuse_smoke --quick

默认入口与 SOTA 复现

下面是 CP8、中型宽度、全局 S=4K 的 Golden 命令。直接照抄,不要改 tile、通信 CTA、raster 或 swizzle:

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 ./build/fuse_bench \
  --mode a2a_gemm_lhs \
  --m 512 --n 5120 --k 5120 \
  --batch 1 --q-heads 40 --head-dim 128 \
  --comm-ctas 0 --lhs-policy auto \
  --raster n --swizzle 1 \
  --warmup 10 --iterations 50 \
  --json-out oproj_cp8.json

正确输出应包含:

comm_ctas=4 policy=m128n160 tile=128x160
fused p50 ≈ 82.9 μs

这是不需要逐 shape 调参的生产入口。v4.0 的极限表保持 --lhs-policy auto --raster n --swizzle 1,只扫描通信 CTA 数;核心 kernel 和 tile 选择逻辑没有修改。具体映射见 benchmarks/a2a+Oproj/BENCHMARK.md

相同软件和硬件下,短程复跑落在 Golden 的 ±5% 可视为正常。明显偏慢时按顺序检查:

  1. bash scripts/require_idle_gpus.sh 必须通过;
  2. nvidia-smi topo -m 中参与设备必须走 NVLink/P2P;
  3. 各卡 SM clock 保持稳定,结果由最慢 rank 决定;
  4. 依赖 commit 与上文一致,使用 Release/SM90a 构建;
  5. 默认模式保留 --comm-ctas 0 --lhs-policy auto --raster n --swizzle 1;复现 v4.0 极限表时只替换 --comm-ctas

完整 shape matrix:

# OProj:正式表分别运行 eager 与 CUDA Graph。
python3 'benchmarks/a2a+Oproj/oproj_shape_bench.py' \
  --phase fuse-launch-formal --phase fuse-launch-aggregate \
  --phase shape-table \
  --models representative_small,representative_medium,representative_large \
  --seqs 1024,4096,16384 --cps 8 \
  --results results/reproduce_cp8

# QKV:正式表必须使用 MPI 一进程一卡;eager 与 Graph 分列。
unset FUSE_QKV_COMM_POLICY FUSE_QKV_GEMM_POLICY
python3 'benchmarks/QKVproj+a2a/qkv_shape_bench.py' \
  --phase fuse-mpi-formal --phase fuse-mpi-aggregate \
  --phase comparison-table

性能

实验设置:单机 8×H200、NVLink、每卡 132 SM、BF16、CUDA 12.8;10 次 warmup + 50 次采样,表内延迟为跨 rank 最大值的 p50。最优分离实现取调优后的 TE+NCCL 与 cuBLASLt+NCCL 中较快者;纯 GEMM 百分比固定对比经典 cuBLAS。吞吐只计算 GEMM FLOPs,延迟包含通信。

当前版本:v13.1(Ulysses 数据流原语与投影语义两级契约)

v13.1 不更新性能表。与 v12.0/v13.0 的干净 Release 构建逐函数比较,69 个 device function 的 SASS 全部一致,58 个公开符号也完全一致,因此下表继续引用对应算子的 已发布 benchmark。新代码从 operators/primitives 的两种数据流开始,再由 operators/semantics/ulysses/projection.h 的四份语义契约组合 QKV/OProj 与 Forward/Backward;旧的 operators/ulysses/projection_dataflow.h 仅作兼容入口。

启动口径 CP4 对最强外部 CP8 对最强外部 总胜场 纯 GEMM 中位数(CP4 / CP8)
Eager 47/48,1.110× 中位 48/48,1.178× 中位 95/96 86.6% / 84.9%
CUDA Graph 46/48,1.130× 中位 48/48,1.199× 中位 94/96 86.8% / 86.7%

“最强外部”逐 setting 取 min(TE Userbuffers, 最强 TE/cuBLASLt+NCCL 分离方案)。 Eager 与 Graph 各自做10+50正式采样,不拿两列之间的差值当算子收益。Graph 在采样前 完成 capture、instantiate 和显式 upload。上述极限表包含 per-shape comm_ctas 单变量标定;默认自动入口不承诺零调参复现全部极限点。

QKV的96点v8结果中,Eager/Graph对TE Userbuffers均为96/96胜场,几何平均 分别领先1.282×1.337×;对最强外部基线也均为96/96,几何平均分别领先 1.215×1.267×。相对v7的全量p50几何平均提升为1.0116×1.0135×。 融合吞吐达到经典cuBLAS的全量中位数为Eager 91.0%、Graph 92.2%;CP4分别为 92.6%/93.4%,CP8为88.8%/90.0%。所有case使用同一自动入口,运行时不读取 逐shape winner或TE结果。

v9 锁频异构矩阵使用三张 1500 MHz 卡与 1980 MHz 参考卡,HBM 均为 3201 MHz。CP2/4/6/8 共 22 个 setting 使用 5 次 warmup + 30 次正式采样并逐样本取 max-rank p50。短 QKV 实际启用的 7 点全部提升 1.0619×~1.0901×;长 QKV 全部回退为 1.0000×。OProj 启用的 13 点全部提升 1.1593×~1.4216×。所有启用点均通过 BF16 逐元素完全一致检查。

v10 反向矩阵每个算子覆盖 96 个 setting,并分别测 Eager/Graph 与普通/ZeroBubble, 共 768 份融合结果。下表的“前向占比”是 B+W 总 FLOPs 吞吐除以同 shape 已发布 前向融合吞吐;“cuBLAS 中位”使用本轮同卡组、同 MNK、匹配 beta 的两次经典 cuBLAS 纯 GEMM。

反向算子 调度 Eager 前向占比几何平均 Graph 前向占比几何平均 Eager / Graph cuBLAS 中位 Eager / Graph 989T MFU中位
QKV ZeroBubble B/W分离,beta=1 100.5% 98.0% 91.0% / 90.4% 63.7% / 64.0%
QKV 普通同流 B→W,beta=0 104.3% 102.3% 88.8% / 89.6% 64.7% / 64.6%
OProj ZeroBubble B/W分离,beta=1 99.6% 100.1% 89.7% / 90.3% 62.4% / 62.4%
OProj 普通同流 B→W,beta=0 104.0% 104.6% 89.5% / 90.1% 63.1% / 63.4%

所有四组几何平均都达到同 shape 前向吞吐的 98.0% 以上。小矩阵 exact smoke 覆盖 CP4/CP8、rank-major/causal、batch=2、同流/分离和 beta=1;正式 S=1K 再检查跨 rank route、epoch 与完整写入。v9 的 43 个旧前向 device-kernel 指令体 在 v10 Release 构建中保持一致。

适配版 TE Userbuffers 的 96 点正式强基线也已完成。QKV 的 Eager 普通/ ZeroBubble 几何平均加速为 1.463×/1.475×,Graph 为 1.222×/1.216×;OProj 对应为 1.428×/1.430×1.146×/1.146×。PyTorch 原生 forward→autograd backward 对照覆盖两种布局、CP4/CP8、batch=2、宽 GQA、普通与 ZeroBubble 共 16 组,全部通过;ZeroBubble 连续两次 beta=1 累加的最大绝对误差为 0.0009765625

v11.2 的完整训练对照同时替换 QKV/OProj 前向和反向,原生侧与融合侧均开启 full-iteration CUDA Graph。Nanbeige、Llama-3 8B geometry、Qwen2.5 7B geometry 的 1K–128K 共 15 个 setting 全部提速,完整 step 吞吐几何平均提升 2.09%,最大提升为 Nanbeige 16K 的 4.26%。该时间包含前向、激活重计算、 反向与优化器,不是把四个 microbenchmark 的收益简单相加。

v12.0 的 FP8 benchmark 使用纯 E4M3 输入、权重、通信数据和输出,FP32 累加; 量化、amax 和 scale 由调用方负责。最终 QKV CUDA Graph 96 点相对同 shape BF16 Graph 的 p50 几何平均为 1.763×85/96 达到 1.5×。没有达到 1.5× 的 短矩阵继续保留,因为固定通信与同步成本不会随 Tensor Core FLOPS 同比例下降。 最终机器可读结果与四条路径的确定性正确性入口见 四算子 FP8 benchmark

完整数据与复现流程:

About

No description, website, or topics provided.

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages