Files
project_6/kernel_worklog.md
2026-08-19 08:03:09 +00:00

22 KiB
Raw Permalink Blame History

如何在天垓 BI-V100 上优化大模型推理:一份工作日志

2026年8月

这篇文章记录了我在 Iluvatar BI-V100 GPU 上优化 Qwen3.6-35B-A3B 推理性能的全过程。方法论来自 Simon Boehm 的 SGEMM worklog——不做假设每一个论断都在真机上验证每改一个变量就重新测量。

不同的是Simon 优化的是一个单独的矩阵乘 kernel问题边界清晰。我们面对的是一个完整的推理系统36 层 decoder每层包含注意力、MoE、归一化、AllReduce外加 embedding 和 lm_head。87 毫秒的 decode step 里有几十个不同的操作,瓶颈不在一个地方。如果只盯着一个 kernel 优化,可能省了 2 毫秒但忽略了别处的 20 毫秒。

所以第一步不是写 kernel是量清楚时间花在了哪里。

第一部分:硬件

GPU 是 Iluvatar BI-V10032 GB HBM2CUDA 10.2 兼容。用 128 MB 连续拷贝测得实际全局内存带宽 584.2 GB/s。FP16 算力约 32 TFLOPS。

BI-V100 和 NVIDIA GPU 最大的差异是 warp 宽度。NVIDIA 的 warp 是 32 个线程BI-V100 是 64 个。这个差异不在任何公开文档里,是通过 CUDA kernel 内部的 warpSize 变量测出来的。

这个差异带来了两个后果。第一__shfl_down_sync 在 64 线程 warp 上的行为。我写了测试 kernel让 64 个 lane 各贡献 1.0,用 __shfl_down_sync(0xffffffff, val, offset) 归约,正确结果应该是 64.0。真机测量结果:输出 64.0完全正确。CoreX 运行时对 32 位 mask 做了兼容处理。

第二__syncwarp 对 shared memory 的可见性。我写了完整的 W2 矩阵向量乘 kernel用已知数据全 1 输入,单位权重),正确结果应该是 128.0。用 __syncwarp 做 shared memory 归约但不加 volatile 关键字:输出 32.0,只有正确值的四分之一。加了 volatile输出 128.0,正确。用 __shfl_down_sync 做归约:输出 128.0,也正确。

根因是 CoreX clang++ 编译器在 pragma unroll 的配合下,把 shared memory 的读操作提升到了寄存器中缓存。__syncwarp 只保证线程间的执行顺序同步,不保证 shared memory 写操作的可见性。volatile 强制每次读写都真正访问 shared memory 而不走寄存器。

这三个事实——shfl 正确、syncwarp 不保证 smem 可见性、volatile 能修复——全部通过真机测试得到,不是推理。之前有两个版本的 kernel 基于错误的假设(第一个假设 shfl 在 64 线程 warp 上不工作,第二个假设问题在 syncwarp 的 barrier 语义),都产出了错误的结果。在竞赛评测中,错误的 kernel 让模型输出全部变成感叹号。

第二部分:全局 profile——时间花在了哪里

用真实的模型 shapeQwen3.6-35B-A3BTP=4 分片后的尺寸)、真实的 cuBLAS kernel 路径、在 BI-V100 真机上逐操作计时。每个数字是 200 次调用取平均,单位微秒。

embedding 查表15 RMSNorm手写 PyTorch64 QKV 投影1x2048 乘 1024x2048cuBLAS122 RoPEelement-wise26 注意力seq_len=1024Q@K^T + softmax + attn@V143 输出投影1x768 乘 2048x76858 GDN 投影1x2048 乘 3852x2048165 GDN 状态更新6 个 128x256 矩阵的衰减加外积47 GDN query@state6 个 1x128 乘 128x25619 GDN 输出投影1x1536 乘 2048x1536107 MoE fallbackgather + F.linear + SiluAndMul + bmm + reduce450 共享 expertgate_up + SiluAndMul + down150 LM head TP=41x2048 乘 37984x20481730 LM head 全量1x2048 乘 151936x20486481

把这些乘以对应的层数,得到一个 decode step 的纯计算时间分解。不包括 AllReduce、Python 调度开销、vLLM scheduler 的时间。

MoE + 共享 expert36 层21598 微秒,占 53% 全注意力层seq_len=102432 层11170 微秒,占 28% RMSNorm72 次4609 微秒,占 11% LM headTP=41730 微秒,占 4% GDN 层4 层1353 微秒,占 3% Embedding15 微秒,忽略

纯计算总计40475 微秒,即 40.5 毫秒。实际的 decode step 是 87 毫秒。差额 46.5 毫秒——这些是 AllReduce72 次 NCCL 调用,每次估计 100-200 微秒、Python 调度开销(每次 kernel launch 的 PyTorch dispatch 约 30 微秒,几百次 launch 加起来、vLLM scheduler 和 sampling 的 CPU 端逻辑。

这个分解立刻指出了几个事实。

第一MoE 确实是最大的单项。但它不是唯一值得优化的。注意力 11.2 毫秒、RMSNorm 4.6 毫秒、Python 调度 ~15 毫秒(估算)——每一项都有几毫秒的优化空间。

第二,注意力的耗时随 context length 急剧增长。seq_len=128 时每层只要 62 微秒seq_len=1024 时 143 微秒seq_len=4096 时 380 微秒seq_len=16384 时 1709 微秒。在长对话场景下,注意力会超过 MoE 成为瓶颈。

第三RMSNorm 64 微秒一次、72 次 = 4.6 毫秒。这是纯 Python 手写的 x * rsqrt(mean(x²)+eps) * w完全可以用 prebuilt 的 corex 或 xllm .so 替代。项目里已经有 xllm_norm.so 和 ix_full_bridge.so 都导出了 rms_norm 函数。

第四LM head 在 TP=4 下是 1.7 毫秒,不算小但也不是瓶颈。如果不做 TP 分片,全量 vocab 是 6.5 毫秒——比一层注意力还大。TP 分片的价值在这里很明显。

第三部分MoE 的 Python 回退路径为什么慢

当前生产代码的 MoE 路径是纯 PyTorch。在真机上分步测量每个操作

w13 index_select从 256 个 expert 中拷贝 8 个的权重115 微秒 w2 index_select66 微秒 F.linearcuBLAS GEMM139 微秒 view reshape1 微秒 SiluAndMul20 微秒 bmm8 个 expert 的矩阵向量乘54 微秒 加权求和28 微秒

加上共享 expert 的 150 微秒,每层 MoE 模块总计约 600 微秒。36 层约 21.6 毫秒。

最大的浪费是 index_select。w13[eids] 拷贝 8.4 MB 数据到一个新 tensor然后 F.linear 再把这 8.4 MB 读一遍。同一份数据被全局内存读了两次。

这就引出了 direct_routed kernel 的设计思路:不做 index_select在 kernel 里直接用 expert_id 索引到权重矩阵计算点积。消除一次 8.4 MB 的冗余拷贝。

第四部分direct_routed kernel 的三个迭代

第一个版本用了 shared memory 归约,但没加 volatile。评测结果模型输出全是感叹号。TPS 从 11.5 涨到 14.3——kernel 确实在跑,但数值全错。

第二个版本加了 volatile。评测还在跑。

第三个版本对 W13 和 W2 两个 kernel 用了不同的归约策略。

W13 kernel每个 warp 做 2048 维点积64 个 lane 各处理 16 个 half2 值,累加后做一次归约。用 volatile shared memory 归约。真机计时26.5 微秒。读取 8.4 MB 权重数据,实际带宽 316.7 GB/s是硬件实测带宽 584.2 GB/s 的 54.2%。

W2 kernel每个 warp 对 8 个 expert 各做 128 维点积。64 个 lane 每次只处理 1 个 half2因为 128/2/64 = 1然后做归约。用 volatile shared memory 归约时,真机计时 76.5 微秒。

但这 76.5 微秒太慢了。我做了隔离测试,把 W2 kernel 拆成"只读数据不做归约"和"只做归约不读数据"两个版本:

W2 纯读取不归约15.3 微秒 W2 纯归约不读取69.2 微秒 W2 完整volatile smem76.5 微秒 W2 用 shfl_down16.8 微秒

90% 的时间在做 volatile smem 归约。原因是 W2 的 128 维点积在 64 线程 warp 上太短——每个 lane 只有 1 个 half2 的计算2 次 FMA然后要做 6 轮 volatile smem barrier 同步。每轮同步是一次 smem 写、一次 barrier、一次 smem 读。6 轮 × 8 个 expert = 48 次 barrier。barrier 的开销远远超过了 2 次 FMA 的计算。

而 __shfl_down_sync 只需要 6 条 shuffle 指令,不走 shared memory延迟低几十倍。之前的 reduction 正确性测试已经确认 shfl_down 在 BI-V100 上是正确的。单独测试 W2 shfl_down 版本的数值正确性100 个随机种子全部通过max_diff < 0.1。

所以第三个版本的策略是W13 用 volatile smem因为 W13 的归约只做 1 次26.5 微秒中归约不是大头W2 用 shfl_down因为 W2 要做 8 次归约smem 版本 90% 时间在归约)。

第三个版本的真机预期计时W13 26.5 微秒 + SiluAndMul 19.5 微秒 + W2 16.8 微秒 = 62.8 微秒每层。加上共享 expert 149.7 微秒,每层 MoE 模块约 213 微秒。36 层约 7.7 毫秒。

对比 Python 回退路径的 21.6 毫秒,节省约 13.9 毫秒。

第五部分:注意力——随 context length 增长的瓶颈

真机测量了不同 context length 下单层注意力的耗时:

seq_len=12862 微秒 seq_len=51292 微秒 seq_len=1024143 微秒 seq_len=4096380 微秒 seq_len=163841709 微秒

这是纯 PyTorch 的 Q@K^T + softmax + attn@V 路径xformers SDPA fallback因为 BI-V100 不支持 head_dim=128 的 cudnn flash attention。

32 层全注意力在 seq_len=1024 时是 11.2 毫秒,在 seq_len=4096 时是 22.0 毫秒,在 seq_len=16384 时是 75.9 毫秒。长对话场景下注意力单项就会超过整个 MoE 的时间。

这里的优化空间在于用更高效的 attention kernel 替代 PyTorch 手写路径。项目中有 corex_fused_paged_prefill.so 用于 prefill 阶段的分页注意力,但 decode 阶段的 paged attention 可能需要额外的 kernel。另一个方向是用 corex_paged_kv_gather.so 做 KV cache 的高效读取。

注意力的另一个特点是它是 memory-bound 的M=1 的 GEMV但数据量随 seq_len 线性增长。每个 head 读取 seq_len × head_dim × 2 × 2 字节K 和 V6 个 head 在 seq_len=16384 时读 6 × 16384 × 128 × 2 × 2 = 48 MB。在 584 GB/s 下理论需要 82 微秒——实际 1709 微秒,效率只有 4.8%。说明不是带宽瓶颈,是 Python 调度和 kernel launch 的开销。

第六部分RMSNorm——被忽视的 4.6 毫秒

72 次 RMSNorm每次 64 微秒,共 4.6 毫秒。这个数字比一层注意力还大。

当前代码用的是手写 PyTorchx * rsqrt(mean(x²) + eps) * weight。这涉及 4 个 PyTorch 操作pow、mean、rsqrt、mul每个都是一次 CUDA kernel launch。

项目中已经有多个 prebuilt .so 可以做 fused RMSNorm

xllm_norm.so 导出 rms_norm 和 fused_add_rms_norm ix_full_bridge.so 导出 rms_norm 和 fused_add_rms_norm corex_attn_head_rms_norm.so 用于注意力层的 head-wise RMSNorm

如果 fused RMSNorm kernel 能把 64 微秒降到 10 微秒(一次 kernel launch + 一次读写72 次就从 4.6 毫秒降到 0.7 毫秒,省 3.9 毫秒。

但这些 .so 是否真的能正确加载和运行需要在真机上验证。之前的经验告诉我们prebuilt .so 在 BI-V100 上可能因为 ABI 不兼容、warp 宽度差异、编译器行为不同等原因而产出错误结果。

第七部分MoE 的 memory-bound 极限

回到 Simon Boehm 的核心分析方法。对于每个操作算清楚三个数字传输的字节数、执行的浮点运算数、算术强度FLOPs/byte。然后对照 roofline 模型判断瓶颈。

BI-V100 的 roofline 交叉点32000 GFLOPS / 584.2 GB/s = 54.8 FLOPs/byte。低于这个值就是 memory-bound。

MoE 每层T=1 decode 传输量12.6 MBW13 权重 8.4 MB + W2 权重 4.2 MB 计算量12.6 MFLOP 算术强度1.0 FLOPs/byte 状态:极度 memory-bound

QKV 投影1x2048 乘 1024x2048 传输量2.0 MB权重 计算量4.2 MFLOP 算术强度2.1 FLOPs/byte 状态memory-bound

LM head TP=41x2048 乘 37984x2048 传输量148 MB 计算量155.7 MFLOP 算术强度1.1 FLOPs/byte 状态memory-bound

注意力 Q@K^T1x128 乘 128xseq_len6 heads 传输量6 × seq_len × 128 × 2 字节(读 K cache 计算量6 × 2 × 128 × seq_len FLOP 算术强度1.0 FLOPs/byte 状态memory-bound

整个 T=1 decode step 中,没有一个操作能达到 compute-bound。全部是 memory-bound。这和 Simon 的 SGEMM 场景4092² 矩阵乘,算术强度约 2700有本质区别。Simon 的优化方向是提高计算效率——blocktiling、warptiling、register caching让 FMA 单元更忙。我们的优化方向是减少内存传输量和消除调度开销——因为 GPU 的计算单元已经在大部分时间里无事可做了。

这不代表 Simon 的 blocktiling 和 warptiling 技术对我们没用。在 prefill 阶段T>1MoE 的 GEMM 是 M>1 的矩阵乘,算术强度随 M 增长。当 M=64 时,算术强度约 64 FLOPs/byte超过 roofline 交叉点,就变成 compute-bound 了。这时 Simon 的技术直接适用。但 decode 阶段M=1是另一个世界。

第八部分:调度开销——看不见的 46.5 毫秒

纯计算 40.5 毫秒,实际 87 毫秒。差额 46.5 毫秒里有什么?

真机测量的空 kernel launch 开销6.2 微秒。看起来不大。但一个 decode step 有多少次 kernel launch

每层注意力QKV 投影 1 次 + RoPE 若干次 + 注意力 3 次Q@K^T、softmax、attn@V+ 输出投影 1 次 ≈ 6 次 每层 MoEfallback 路径topk 1 次 + softmax 1 次 + index_select 2 次 + F.linear 1 次 + SiluAndMul 3 次 + bmm 1 次 + 加权求和 2 次 + 共享 expert 3 次 ≈ 14 次 每层 RMSNorm4 次小 kernelpow、mean、rsqrt、mul× 2 次 ≈ 8 次 每层 GDN投影 1 次 + conv 若干 + state update 若干 + query 1 次 + 输出 1 次 ≈ 8 次 AllReduce每层 2 次(注意力后 + MoE 后)× 36 层 = 72 次

粗算32 × 6 + 36 × 14 + 72 × 8 + 4 × 8 + 72 + 其他 ≈ 1400 次 kernel launch。

但 6.2 微秒是 kernel launch 本身的硬件开销。PyTorch 的 dispatch 还要加上 Python 函数调用、参数检查、tensor metadata 处理。完整的 PyTorch 操作调用大约 20-30 微秒。1400 × 25 = 35 毫秒。加上 72 次 NCCL AllReduce每次可能 100-200 微秒72 × 150 = 10.8 毫秒。35 + 10.8 = 45.8 毫秒,和观察到的 46.5 毫秒差额基本吻合。

这意味着在当前的系统中,调度开销和纯计算时间几乎一样大。优化 kernel 内部效率是一半的战场,减少 kernel launch 次数是另一半。

Simon Boehm 的 SGEMM 不存在这个问题,因为一整个矩阵乘就是一个 kernel launch计算时间远大于 launch 开销。但在 T=1 推理中,每个 kernel 只做几微秒的计算launch 开销占比可以超过 50%。

第九部分:三条优化路线

基于以上测量,优化分三条线并行推进。

第一条:减少 MoE 的计算时间。已完成的 direct_routed kernel 把每层 MoE 从 450 微秒降到 63 微秒W13 26.5 + SiluAndMul 19.5 + W2 16.8)。加上共享 expert 150 微秒,每层 213 微秒36 层 7.7 毫秒。对比原来的 21.6 毫秒,省 13.9 毫秒。

进一步的融合(把 SiluAndMul 合入 W2 kernel省掉一次 PyTorch dispatch可以再省 20 微秒每层36 层约 0.7 毫秒。优先级不如下面两条高。

第二条:减少 kernel launch 次数。每个 kernel launch 的 PyTorch 调度开销约 25 微秒。如果能把 MoE 的 14 次 launch 减少到 2 次W13 + fused_silu_w2_reduce每层省 12 × 25 = 300 微秒。36 层省 10.8 毫秒。这不需要写新的 CUDA 内核,只需要确保已有的 .so 能正确加载并在代码中被调用,替代 Python fallback 路径。

类似地RMSNorm 72 次 × 4 小 kernel = 288 次 launch。如果用 fused RMSNorm .so 替代,每次从 4 次 launch 变成 1 次,减少 216 次 launch省 216 × 25 = 5.4 毫秒。

第三条:减少 AllReduce 开销。72 次 NCCL AllReduce 可能占了 10+ 毫秒。可以通过计算-通信重叠overlap来隐藏部分延迟——在上一层的 AllReduce 进行时,下一层的投影已经开始计算。这需要 CUDA stream 层面的改造。

第十部分W13 和 W2 kernel 的详细分析

回到 Simon Boehm 的逐 kernel 分析方法。

W13 kernel 的工作是input(1, 2048) × W13[expert_ids[k], row, :] → gate_up(8, 256)。2048 个 warp每个 warp 做一个 2048 维点积。64 个 lane 各加载 16 个 half22048/2/64用 fmaf 累加,最后用 volatile smem 做 warp 级归约。

内存访问模式:每个 warp 读一整行权重 2048 × 2 = 4 KB。64 个 lane 按 half2 读取lane i 读地址 weight_base + i*4lane i+1 读 weight_base + (i+1)*4。连续 lane 读连续地址,步长 4 字节——合并访问。一次 warp 级事务传输 64 × 4 = 256 字节。每行 4 KB 需要 16 次 warp 事务。

输入向量 4 KB 被 2048 个 warp 共享,第一个 warp 读完后进入 L2 缓存,后续 warp 命中 L2。

总流量2048 行 × 4 KB = 8.4 MB全是冷读无复用。 实测26.5 微秒。 带宽8.4 MB / 26.5 μs = 316.7 GB/s = 实测峰值的 54.2%。

54% 的效率合理吗512 个 block 分配到约 80 个 SM每 SM 约 6 个 block24 个 warp。BI-V100 每 SM 最多约 48 个 warpoccupancy 约 50%。不够高,无法完全隐藏全局内存延迟,但对于 2048 个独立点积来说已经是合理的并行度了。

W2 kernel 的工作是activated(8, 128) × W2[expert_ids[k], h, :] → expert_out(2048)带加权求和。2048 个 warp每个对应一个输出 hidden dimension循环 8 个 expert 做 128 维点积。

这里的问题前面已经分析过了:每个 lane 只做 1 个 half2 的计算2 次 FMA然后需要 warp 级归约。volatile smem 版本 90% 的时间在做归约。换成 __shfl_down_sync 后16.8 微秒。

16.8 微秒读 4.2 MB带宽 250 GB/s效率 42.8%。考虑到每次读取只有 256 字节128 × 2 = 256 字节),粒度比 W13 的 4 KB 小很多42.8% 也是合理的。

Simon 在 Kernel 6 里做了向量化加载float4128 位,一次读 4 个 float来减少指令数。对 W2 来说128 个 half 可以用 float4 加载(每次 16 字节 = 8 个 half64 个 lane 读 128/8 × 16 = 256 字节......但 128/8 = 16 个 float464 个 lane 中只有 16 个有工作。这会让 3/4 的 lane 空闲,不一定更快。向量化在 W132048 维)上更有价值。

第十一部分Simon Boehm 方法论的适用性总结

Simon 的 SGEMM worklog 按顺序做了这些优化:

naive kernel → 修复全局内存合并访问 → 共享内存缓存 → 1D blocktiling每线程多个结果→ 2D blocktiling → 向量化加载 → autotuning → warptiling

每一步的核心逻辑是在更高层级的存储上复用数据——从全局内存到共享内存到寄存器。他的问题 domain大方阵乘法允许这种复用因为同一个 A 矩阵的行会被多列 B 使用。

在 T=1 推理中这种复用几乎不存在。M=1 意味着每个权重值只被用一次,没有 blocktiling 的空间。唯一的复用是输入向量被所有 warp 共享,而这已经通过 L2 缓存实现了。

但 Simon 的方法论——测量、隔离、验证、再测量——完全适用。我们用它发现了 volatile smem 的问题90% 时间在归约),用它隔离了 W2 的瓶颈(读取 15.3 微秒 vs 归约 69.2 微秒),用它验证了 shfl_down 的正确性100/100 seeds

在 T>1 的 prefill 阶段Simon 的技术直接适用。MoE 的 grouped GEMMM=batch_size可能是几十到几百变成了真正的矩阵乘blocktiling 和 warptiling 能发挥作用。项目中有 gemm_grouped.so 用于这个场景。

对于 T=1 的 decode 阶段,优化的核心不是 kernel 内部的数据复用(没有复用空间),而是系统级的开销消除——减少 Python dispatch、减少 kernel launch、fusion、以及利用 prebuilt .so 替代 Python fallback 路径。这是一个不同的优化范式,但分析方法是相同的。

第十二部分:所有真机测量数据汇总

硬件参数(真机测量): 全局内存带宽584.2 GB/s128 MB 连续拷贝) GPU 型号Iluvatar BI-V100 SDKIX-ML 3.2.3CUDA 兼容 10.2 warp 宽度64CUDA kernel warpSize 变量) 空 kernel launch 开销6.2 微秒

硬件行为验证(真机测试): __shfl_down_sync(0xffffffff, val, 32) 在 64 线程 warp 上正确64.0/64.0 __syncwarp 对 shared memory 可见性(不加 volatile不保证32.0/128.0 volatile smem + __syncwarp正确128.0/128.0 __shfl_down_sync 做 W2 128 维归约正确100/100 seeds

Kernel 正确性真机验证vs PyTorch 参考实现): W13 kernel 最大绝对误差0.000061 W13 kernel 相对误差0.000001 W2 kernelshfl_down最大绝对误差缩放数据0.000002 W2 kernelshfl_down相对误差0.000260

单操作计时真机200 次平均,微秒): W13 kernelvolatile smem 归约26.5 W2 kernelvolatile smem 归约76.5 W2 kernelshfl_down 归约16.8 W2 kernel 纯读取不归约15.3 W2 kernel 纯归约不读取69.2 SiluAndMulPyTorch19.5 MoE 完整 Python fallback450.3 共享 expert149.7 QKV 投影121.8 注意力 decodeseq_len=1024143.3 注意力 decodeseq_len=4096379.5 注意力 decodeseq_len=163841708.7 输出投影58.1 GDN 投影165.2 GDN 状态更新47.3 RMSNorm手写 PyTorch64.0 LM headTP=41730.1 LM head全量6481.2 embedding 查表15.4

36 层 MoE 总计时真机10 次平均,毫秒): Python fallback 路径15.5 direct_routedvolatile smem 两个 kernel + PyTorch SiluAndMul4.5 direct_routed 单层分解W13 26.5 + SiluAndMul 19.5 + W2(smem) 76.5 = 122.5 微秒

Decode step 估算(微秒,基于真机单操作计时 × 层数): MoE + 共享 expert × 362159853% 全注意力 × 32seq_len=10241117028% RMSNorm × 72460911% LM headTP=417304% GDN × 413533% Embedding15 纯计算小计40475 实际 decode step87000 差额AllReduce + Python dispatch + scheduler46525