# 海光 DCU 多卡推理预筛实录：TP2/TP4 为什么不能照搬 TP1 的 kernel 白名单 我开始准备多卡推理时，最先冒出来的想法很朴素：TP1 上已经测快的 kernel，到了 TP2、TP4，把同一份白名单扩过去不就行了？ 真正把每个 rank 的矩阵形状列出来后，我很快发现这件事没有这么省心。 Tensor Parallel 不只是把一次计算分给多张卡，它会直接改变单个 rank 看到的 N、K、Q heads、KV heads，甚至改变一个 kernel 的线程块是否合法。 在海光 DCU gfx936 上，同一套实现面对 TP2 和 TP4，可能出现三种完全不同 的结果：继续明显提速、收益缩到几乎没有，或者连 launch geometry 都不再 成立。 这篇文章记录的是一次单卡 rank-local 预筛。测试环境为 gfx936， PyTorch 2.10.0，HIP runtime 6.3.26093。我在一张 DCU 上构造 TP2、TP4 各 rank 实际会遇到的局部 shape，用 CUDAGraph 测 kernel 本身，再检查 输出、实际 dispatch 和静态硬件约束。 这里必须先写清边界：这些数据不包含 RCCL 通信，也不包含多进程图捕获和 完整服务调度，所以它们不是多卡吞吐预测。它们的价值是提前排除肯定不能 用的候选，把昂贵的真实 TP 端点时间留给更有希望的组合。 ## 一、TP 改变的首先是每个 rank 的 shape 以几类 BF16 线性层为例，TP2 和 TP4 下的局部输出宽度并不是简单的标签 变化，而是直接进入 kernel dispatch 的实参。 我先测了当前用于 M=1 decode 的专用 LLMM1 kernel。下面的加速比均为 CUDAGraph 中位数，相对 torch.nn.functional.linear。 | 局部算子 | TP2 shape | TP2 结果 | TP4 shape | TP4 结果 | |---|---|---:|---|---:| | Attention QKV | 7168×5120 | 1.110x，保留 | 3584×5120 | 1.080x，边界保留 | | GDN qkvz | 8192×5120 | 1.100x，保留 | 4096×5120 | 1.019x，拒绝 | | MLP gate/up | 17408×5120 | 1.914x，强保留 | 8704×5120 | 1.993x，强保留 | | GDN BA | 48×5120 | 1.020x，拒绝 | 24×5120 | 1.536x，保留 | | MLP down | 5120×8704 | launch 非法 | 5120×4352 | 1.258x，保留 | | out projection | 5120×3072 | 1.339x，保留 | 5120×1536 | 1.014x，拒绝 | 这张表最有用的地方，不是某个 1.9x，而是同一家族在不同 TP 下发生的 翻转。TP2 的 GDN BA 太窄，只有 1.020x；到了 TP4，N 继续减半，却因为 调度形态变化获得 1.536x。out projection 正好相反：TP2 值得保留， TP4 已经只剩 1.014x。 这说明白名单的键至少要包含真实的 N、K 和 kernel mode，不能只写 “这是 GDN”或“这是 out projection”。 ## 二、最快的错误结果，往往来自静态约束没先检查 TP2 MLP-down 是这轮预筛里最值得记住的反例。当前 kernel 的线程数公式是： ~~~text NUM_THREADS = round_up(K / 8, 64) ~~~ 当 K=8704 时： ~~~text 8704 / 8 = 1088 ~~~ gfx936 的单个线程块上限是 1024。也就是说，这个 shape 请求 1088 个线程， 从一开始就不合法。此时即便计时脚本吐出一个漂亮数字，也不能把它当成性能 结果，因为 launch geometry 已经越过硬件边界。 同一轮检查还发现，调度参数 rows_per_block=1 并不是一个真实存在的 “RPB1 模式”。C++ dispatcher 只显式支持 2、4、8、16，传入 1 会落入 默认的 kernel<4> 分支。若只看外部参数名，会以为自己比较了一个新模式； 只有顺着 dispatcher 看实际实例化，才知道测到的仍是 RPB4。 从这以后，我把静态检查放到了计时之前： 1. 线程数是否超过设备上限； 2. shared memory 是否超限； 3. dispatcher 是否真的存在这个 mode； 4. shape guard 是否会命中预期分支； 5. 结果文件中的 kernel 名是否变化。 这样做很便宜，却能挡住一类“数字很好、实验无效”的问题。 ## 三、Tensile 调优结果也不能直接当作生效证明 Prefill 的 q=4096 大 GEMM，我又按 TP2、TP4 的局部 shape 做了 Tensile override 验证。每个候选都跑两组 baseline/override 交错对照， 同时检查完整输出 hash 和实际 kernel 名。 TP2 中值得带入真实多卡验证的四项是： | TP2 q=4096 局部算子 | 平均加速 | |---|---:| | gate/up | 1.0932x | | out projection | 1.0582x | | GDN qkvz | 1.1103x | | Attention QKV | 1.0785x | TP2 MLP-down 虽然确实切换了 kernel，但只有 1.0184x，低于我设置的 1.05x 预筛门槛，因此拒绝。 TP4 的情况更能说明问题： | TP4 q=4096 局部算子 | 平均加速 | 结论 | |---|---:|---| | gate/up | 1.0896x | 保留 | | GDN qkvz | 1.0772x | 保留 | | Attention QKV | 1.0680x | 保留 | | out projection | 1.0337x | 拒绝 | | MLP-down | 0.9887x | 拒绝 | TP4 MLP-down 的调优器给出了一个 solution 编号，但 baseline 与 override 最终使用的 kernel 名没有变化，而且 override 还慢了一点。这里的教训很 直接：调优器“找到了一行”只说明数据库有候选，不等于运行时完成了切换。 若没有 kernel-name 证据，调优结果只能算线索。 ## 四、Attention 的 head 布局也要按 TP 重做 guard Attention 的局部布局变化更加直观： ~~~text TP1: 24Q / 4KV TP2: 12Q / 2KV TP4: 6Q / 1KV ~~~ 我用这两组新的局部 head tuple 做了独立验证。vendor paged-prefill 相对当前 Triton 路径，在六组 q/k 长度组合上的结果为： | 布局 | 最小加速 | 中位加速 | 数值表现 | |---|---:|---:|---| | TP2，12Q/2KV | 1.4059x | 2.8168x | 全部确定、通过 | | TP4，6Q/1KV | 1.4069x | 2.5267x | 全部确定、通过 | 这些 prefill case 的相对 L2 约为 0.0058 到 0.00643，属于当前 vendor 数值路径的既有口径，并非逐位一致。 q=1 decode 的 page-contiguous kernel 则是另一种表现： | 布局 | 8K | 16K | 32K | 几何平均 | |---|---:|---:|---:|---:| | TP2，12Q/2KV | 1.4037x | 1.4096x | 1.4469x | 1.4199x | | TP4，6Q/1KV | 1.3922x | 1.4022x | 1.4430x | 1.4123x | 两组 decode 输出均逐位一致。这证明底层 kernel 能处理新的局部 head 布局，但仍不代表现有运行时已经会 dispatch 到它。真正的 TP 服务里， 还要核对 block size 推导、guard、图捕获和 RCCL 交互。 ## 五、我现在采用的多卡 kernel 迁移顺序 这轮实验以后，我不再把 TP1 白名单复制到 TP2/TP4，而是按下面的顺序做： 1. 从模型配置计算每个 rank 的真实 N、K、Q heads 和 KV heads； 2. 先做线程数、共享内存和 dispatcher 分支的静态检查； 3. 在单卡上构造 rank-local shape，做 CUDAGraph 与正确性预筛； 4. 对库 kernel 同时验证输出 hash、实际 kernel 名和重复 A/B； 5. 只把过门槛的 shape 写入 TP-aware 白名单； 6. 最后才在真实 TP 端点检查 RCCL、图捕获、TTFT、TPOT 和输出。 单卡 rank-local 预筛不能回答多卡最终有多快，但它很擅长回答另一件事： 哪些候选根本不值得占用多卡调试时间。 对我来说，这就是它最大的价值。多卡优化的第一步不是“把卡加上”，而是 承认每个 rank 已经面对了一套新的 kernel 问题。