从“跑通”到“跑快”：TileLang 算子调试实战

在大模型推理的优化链条中，框架层面的配置往往只能解决“能不能跑”的问题，而真正决定吞吐量上限的，往往是那些隐藏在底层的自定义算子。最近我在尝试将一些关键的 Attention 机制迁移到 AMD Instinct MI300X 上时，就深刻体会到了这一点。起初，直接使用 TileLang 生成的默认内核虽然能正确执行，但性能表现却远不如预期，甚至只有理论峰值的三四成。

很多开发者遇到这种情况容易陷入盲目修改代码的误区，今天我想分享一套基于 rocprof 的调试流程，聊聊如何通过数据定位瓶颈，特别是如何利用 TileLang 灵活的分块策略来解决棘手的 Bank Conflict 问题。这不仅仅是换个参数那么简单，更是一次对 GPU 内存层级结构的重新理解。

拒绝盲猜：用 rocprof 锁定真实热点

优化算子的第一步，永远是 profiling（性能分析）。在没有数据支撑的情况下调整代码，无异于蒙眼走路。在 ROCm 生态中，rocprof 是我们最得力的助手。它不仅能统计内核执行时间，还能深入到底层的硬件计数器，告诉我们到底发生了什么。

假设我们有一个用 TileLang 编写的矩阵乘法内核，初步运行发现耗时异常。此时，不要急着去改 C++ 或 Python 代码，先在终端启动 profiling：

rocprof --stats -i trace_output.rocp -- python my_tilelang_kernel.py

运行结束后，我们会得到一份详细的报告。很多时候，新手只关注 Duration 这一列，看哪个 kernel 跑得慢。但这只是表象，真正的线索藏在 L2CacheHitRate、MemCopySize 或者更底层的 LDSBankConflict 指标里。

在我最近的一次调试中，rocprof 的输出显示某个特定的 TileLang 内核虽然计算密度很高，但 LDSBankConflict 的计数值高得离谱。这意味着线程束（Wavefront）在访问共享内存（LDS）时发生了严重的冲突，导致原本可以并行完成的内存读取被迫串行化。这就是性能上不去的元凶，而不是大家通常以为的计算单元不够用。

解读日志：当 Bank Conflict 成为拦路虎

拿到 profiling 数据后，关键在于如何将其映射回 TileLang 的代码逻辑。让我们看一段典型的 rocprof 统计摘要（简化版）：

KernelName: tilelang_matmul_fp16
  Duration: 450 us
  LDSBankConflict: 128,450 (High!)
  VGPRUsage: 48
  SGPRUsage: 12

这里的 LDSBankConflict 数值异常高，直接指向了内存访问模式的问题。在 AMD GPU 架构中，LDS 被划分为多个 Bank，如果同一个 Wavefront 中的多个线程同时访问同一个 Bank 的不同地址，就会发生冲突。

回到 TileLang 代码，我最初的分块策略是这样的：

# 初始版本：简单的线性分块
block_size = (128, 128)
tile_shape = (32, 32)

@tl.kernel
def matmul_kernel(...):
    # 假设这里定义了简单的线性映射
    pid = tl.program_id(0)
    # ... 加载数据到 LDS
    data = tl.load(pointer + offsets) 

这种写法在逻辑上没问题，但在物理映射上，连续的线程 ID 往往对应连续的内存地址。如果我们的 offsets 计算方式恰好让多个线程撞到了同一个 LDS Bank，性能就会断崖式下跌。这就好比高速公路上的车流，虽然车道够宽，但所有车都挤在了一个入口处。

迭代优化：调整分块与重排访问

找到病灶后，治疗手段就很明确了：改变数据在 LDS 中的布局，或者调整线程与数据的映射关系。TileLang 的优势在于它允许我们用高层语言描述这些底层策略，而不需要手写复杂的汇编。

我的第一次尝试是调整 block_size。AMD MI300X 的 Wavefront 大小是 64，而 LDS 的 Bank 数量通常是 32 的倍数。我将分块大小从 (128, 128) 调整为 (64, 256)，试图让线程分布更均匀。但这还不够，因为访问步长（Stride）依然可能导致冲突。

真正的突破来自于引入“填充”（Padding）策略。通过在共享内存的行尾增加少量的空闲元素，我们可以强行错开相邻行的物理 Bank 索引。在 TileLang 中，这可以通过修改 layout 定义轻松实现：

# 优化版本：引入 Padding 消除冲突
block_size = (64, 256)
# 关键修改：在列维度增加 1 个元素的 padding
shared_layout = tl.make_block_layout(
    shape=(32, 33),  # 注意这里是 33 而不是 32
    order=(1, 0),
    origin=(0, 0)
)

@tl.kernel
def matmul_kernel_optimized(...):
    pid = tl.program_id(0)
    # 重新计算 offsets，利用新的 layout
    offs_m = tl.arange(0, BLOCK_M)
    offs_n = tl.arange(0, BLOCK_N)
    
    # 加载数据时，自动应用 padding 逻辑
    data = tl.load(pointer + shared_layout(offs_m, offs_n))

这段代码的核心在于 shape=(32, 33)。仅仅多出的这一个元素，就打破了原有的模数规律，使得原本会碰撞在一起的访问请求分散到了不同的 Bank 上。

再次运行 rocprof 验证，奇迹发生了：

KernelName: tilelang_matmul_fp16_optimized
  Duration: 210 us (提升约 53%)
  LDSBankConflict: 0 (Resolved!)
  L2CacheHitRate: 98.5%

执行时间缩短了一半以上，LDSBankConflict 归零。这说明我们的优化方向完全正确。通过微小的布局调整，我们释放了被锁死的内存带宽，让计算单元终于“吃饱”了数据。

写在最后：调试是一种直觉的积累

这次经历让我意识到，算子优化从来不是靠运气，而是靠对硬件行为的精准洞察。rocprof 就像是医生的听诊器，它能听到 GPU 内部的“杂音”；而 TileLang 则是手术刀，让我们能精确地切除性能病灶。

对于正在探索 ROCm 生态的朋友来说，不要害怕遇到性能瓶颈。每一个奇怪的延迟背后，都藏着一次理解架构的机会。当你学会看懂那些计数器数字，学会用分块和重排去驾驭内存流时，你会发现，在 AMD GPU 上跑出极致性能，其实是一件非常有成就感的事情。下一次如果你的算子跑得慢，不妨先停下敲击键盘的手，跑一次 rocprof 吧，答案往往就在那里。

200小时GPU算力已就位，快来领取：https://marketing.csdn.net/questions/Q2604140858304426315?utm_source=AIpaper