用数据说话：rocprof 定位内核热点的实战心法

在 AMD GPU 上进行大模型算子优化时，最忌讳的就是“凭感觉”写代码。很多时候，我们觉得某个 Kernel 慢是因为计算太复杂，但实际跑起来才发现是内存带宽被吃光了，或者是线程束（Wavefront）发生了严重的发散。TileLang 作为一个强大的 DSL 工具，能让我们灵活地定义分块和流水线，但如果缺乏精准的 Profiling 数据支撑，这些优化往往是在盲人摸象。

我习惯在动手修改 TileLang 代码前，先让 rocprof 跑一轮。这个工具就像是 GPU 的“听诊器”，能清晰地告诉我们时间都花哪儿了。基本的用法很简单，直接在启动命令前加上 rocprof --rocm-trace-input 或者使用 rocprofv2 生成 HTML 报告即可。但我更关注的是其中的 PMCs（性能计数器）数据，特别是 L2CacheHit 和 MemUnitStalled 这两个指标。

有一次我在优化一个自定义的 Attention 算子时，直觉告诉我问题出在矩阵乘法的指令调度上，于是花了一下午调整指令顺序。结果一测，性能提升不到 2%。后来静下心来仔细看 rocprof 的报告，发现该 Kernel 的 MemUnitStalled 占比高达 70%，而计算单元利用率极低。这说明瓶颈根本不在计算，而在全球内存访问上。那一刻我才意识到，没有数据的优化就是猜谜。只有先通过工具锁定了真正的热点——是 L1 Cache 命中率低？还是 Shared Memory bank conflict？——后续的 TileLang 代码调整才能有的放矢。

案例一：打破内存墙，动态调整分块策略

找到瓶颈后，第一个要动刀的地方通常是分块（Tiling）策略。在 TileLang 中，我们可以非常直观地修改 block_size 和 tile_size，但这其中的学问在于如何匹配硬件的物理特性。

记得在处理一个大规模矩阵乘法场景时，初始版本我沿用了 CUDA 社区常见的 128x128 分块。在 NVIDIA 卡上这或许是个经验值，但在 AMD MI300X 上，性能却迟迟上不去。通过 rocprof 分析，我发现全局内存的读取请求过于细碎，导致带宽利用率只有理论值的 40% 左右。AMD 的架构对合并访问（Coalesced Access）有着严格的要求，如果线程束内的线程访问地址不连续，就会触发多次事务。

解决思路很直接：在 TileLang 中增大 L2 层的分块尺寸，并显式指定数据在 LDS（本地共享内存）中的布局。我将代码中的分块参数从默认的静态值改为根据 wave_size 动态计算的模式。具体来说，利用 TileLang 的语法特性，让每个 Block 一次性加载更多连续的数据到 LDS 中，减少全局内存的访问频次。

# 伪代码示例：调整 TileLang 中的分块逻辑
@tilelang.kernel
def matmul_optimized(A, B, C):
    # 根据硬件 Wavefront 大小动态调整分块，而非硬编码
    block_size = (128, 256) 
    tile_size = (32, 32)
    
    # 显式声明 LDS 缓存策略，减少 Bank Conflict
    A_local = allocate_shared(block_size[0], tile_size[0])
    B_local = allocate_shared(tile_size[1], block_size[1])
    
    # 加载数据时确保合并访问
    load_global_to_shared(A, A_local, coalesce=True)
    load_global_to_shared(B, B_local, coalesce=True)
    
    # 执行计算...

修改后的测试结果显示，全局内存访问量下降了近一半，整体 Kernel 的执行时间缩短了 35%。这个案例让我深刻体会到，所谓的“通用最佳实践”在不同架构下可能完全失效，必须依据 Profiling 数据量身定制分块大小。

案例二：消除线程束发散，精细化指令调度

解决了内存问题，下一个拦路虎往往是控制流带来的开销。在早期的一个 Softmax 实现中，我发现虽然内存带宽跑满了，但 SM 占用率却一直上不去。查看 rocprof 的分支效率指标，发现存在大量的线程束发散（Branch Divergence）。

原因出在一个边界条件的判断上。原本为了代码简洁，我在 Kernel 内部写了一个 if (idx < bound) 的判断来处理矩阵边缘的非对齐部分。在数据量巨大时，这个判断导致同一个 Wavefront 内的线程走了不同的分支路径，一半线程在计算，另一半在等待。

利用 TileLang，我可以更精细地控制循环展开和边界处理。我没有选择在运行时做判断，而是通过模板参数在编译期生成专门处理“非对齐尾块”的特化 Kernel。同时，利用 TileLang 提供的 unroll 指令，将核心循环完全展开，避免了指针跳转带来的开销。

这种“空间换时间”的策略在 AMD 架构上效果显著。因为消除了分支判断，指令流水线不再停顿，计算单元的活跃度大幅提升。虽然代码体积稍微增加了一点，但对于性能敏感的核心算子来说，这点代价完全值得。这也提醒我们，在高性能计算领域，有时候“笨拙”的展开比“聪明”的条件判断更高效。

案例三：流水线重叠，隐藏数据加载延迟

最后一个技巧是关于流水线的。在很多复杂的算子中，计算和访存是串行的：等数据加载完再算，算完再存。这中间的空窗期就是性能的浪费。

在一次优化 LayerNorm 算子时，我尝试引入双缓冲（Double Buffering）机制。思路很简单：当 GPU 正在计算当前这块数据时，预取下一块数据到共享内存。TileLang 对此提供了很好的支持，可以通过 async_copy 类似的语义来描述这种异步行为。

起初 implementation 并不顺利，出现了数据竞争的问题。通过反复迭代测试，并在 rocprof 中观察 GDWS（Global Data Write Stall）和 GLDS（Global Load Store）的时间重叠情况，我逐渐调整了预取的时机和缓冲区的大小。最终，成功实现了计算与访存的完美重叠。从时序图上看，原本明显的“锯齿状”执行波形变得平滑连续，延迟被有效隐藏了起来。

这个过程并非一蹴而就。我经历了至少五六个版本的迭代，每次只改动一个小参数，然后立刻跑测试对比。过早优化是大忌，有时候强行上流水线反而会增加寄存器压力，导致 occupancy 下降。只有在前两步（内存和分支）都优化到位后，流水线重叠才能发挥出最大威力。

写在最后：迭代才是王道

回顾这三次调试经历，核心方法论其实就一条：不要猜，去测。TileLang 给了我们极大的自由度去尝试各种奇思妙想，但如果没有 rocprof 这样的工具作为导航，很容易迷失在细节里。性能优化是一个螺旋上升的过程，从定位热点到调整分块，再到指令微调和流水线设计，每一步都需要数据的验证。

很多开发者容易陷入“一次成型”的误区，指望写出完美的第一版代码。但在真实的工程实践中，尤其是面对 AMD 这样不断进化的架构，快速迭代、小步快跑才是正道。哪怕只是提升 5% 的性能，积少成多，对于大规模训练推理集群来说，节省下来的算力和时间都是巨大的。

如果你也想亲手试试这些优化技巧，或者正在为找不到合适的实验环境发愁，现在有个不错的机会。200 小时 GPU 算力已就位，快来领取：https://marketing.csdn.net/questions/Q2604140858304426315?utm_source=AIpaper 。有了充足的算力资源，配合科学的调试方法，相信你也能在 ROCm 生态中挖掘出令人惊喜的性能表现。