从“跑不通”到被合并：一次 TileLang 算子优化的社区贡献实录

在 AMD ROCm 生态里折腾大模型，最让人上头的瞬间往往不是环境终于配通的那一刻，而是你发现某个算子在特定架构下效率低下，动手优化后不仅自己受益，还顺手帮整个社区填了一个坑。

最近我在尝试用 SGLang 部署 Qwen2.5 时，发现在 MI300X（gfx942）架构上，长序列生成的吞吐量始终卡在瓶颈。经过一番 profiling，问题指向了 Attention 机制中的矩阵分块策略——默认的 Tile 大小并没有完美匹配 AMD GPU 的 Wavefront 尺寸。与其在本地反复调试，不如直接把解决方案贡献给上游。于是，我开启了一次围绕 TileLang 的社区贡献旅程。

锁定问题：从性能异常到复现脚本

发现问题只是第一步，要让维护者快速理解并接纳你的修改，关键在于提供一份“无可辩驳”的复现脚本。

起初，我只是在论坛里发了一句“长序列推理慢”，这种模糊的描述很难得到实质性帮助。后来我调整策略，使用 rocprof 抓取了详细的热点数据，发现 L2 Cache 命中率极低，明显是内存访问模式出了问题。

为了证明这不是我环境配置的问题，我编写了一个最小化的复现 Demo。这个脚本不依赖庞大的模型权重，而是直接调用 TileLang 编写的 Kernel，模拟不同序列长度下的矩阵乘法场景：

# reproduce_issue.py
import torch
from tilelang import TLProgram

def benchmark_tile_performance():
    # 模拟 gfx942 架构下的典型负载
    seq_len = 32768
    head_dim = 128
    
    # 原始实现：固定 Tile 大小，未适配 Wavefront
    print("Running baseline kernel...")
    # ... 调用原始 Kernel 代码 ...
    
    # 优化预期：调整分块以匹配硬件
    print("Expected improvement with tuned tiling...")

if __name__ == "__main__":
    benchmark_tile_performance()

我把这个脚本、rocprof 的火焰图截图以及具体的硬件环境信息（ROCm 版本、驱动版本、GPU 型号）整理成了一个清晰的 GitHub Issue。在描述中，我没有抱怨“为什么这么慢”，而是客观陈述：“在 gfx942 上，当 SeqLen > 16k 时，当前分块策略导致 LDS 利用率不足 60%。”

社区协作：在讨论中打磨方案

Issue 发出后不到半天，TileLang 的一位 Maintainer 就参与了讨论。他并没有直接给我代码，而是提出了一个关键思路：是否考虑过动态调整 Block Size 以适应不同的 Sequence Length？

我们在 Issue 评论区来回交流了几轮。他建议我参考 AMD CDNA 架构的编程指南，特别是关于 Vector Register 压力的部分。这让我意识到，单纯改 Tile 大小可能不够，还需要调整寄存器分配策略。

这种交流非常高效，没有客套话，全是技术干货。我们甚至约定了一起在线 Pair Programming 的时间。在屏幕共享中，他向我展示了如何使用 TileLang 的 DSL 语法更优雅地描述数据流动，而不是硬编码魔法数字。

“试着把 BLOCK_M 设置为 128 的倍数，同时确保 BLOCK_N 能整除 Wavefront 大小（64）。”他在聊天框里敲下这行建议。

动手优化：用 TileLang 重写关键算子

有了明确的方向，我开始着手修改代码。TileLang 的魅力在于它允许你用类似 Python 的高层语法来定义底层 GPU 行为。我不需要去写晦涩的 HIP C++，只需关注计算逻辑本身。

核心的改动在于重新设计了分块策略。针对 gfx942 架构，我将原本固定的分块改为动态感知：

@tl.program
def optimized_attention(Q, K, V, O):
    # 动态获取硬件属性
    wavefront_size = tl.get_wavefront_size()
    
    # 根据序列长度自适应调整 Tile 大小
    if tl.sequence_length > 16384:
        block_m = 128
        block_n = 64  # 匹配 Wavefront
    else:
        block_m = 64
        block_n = 32
        
    # 优化 LDS 布局，减少 Bank Conflict
    q_tile = tl.load(Q, block_shape=[block_m, block_n])
    # ... 后续计算逻辑 ...
    
    tl.store(O, result)

这段代码看起来简洁，但背后是对硬件特性的深度适配。修改完成后，我在本地跑了之前的复现脚本。结果显示，在长序列场景下，Kernel 的执行时间减少了约 25%，LDS 利用率提升到了 85% 以上。

提交 PR：细节决定成败

代码跑通了，接下来就是正式的 Pull Request。在开源社区，PR 的质量往往决定了它被合并的速度。

我没有只丢一段代码上去，而是精心准备了 PR 描述：

1. 问题背景：链接到之前的 Issue，简述性能瓶颈。
2. 解决方案：解释为什么选择这种分块策略，引用了哪些架构文档。
3. 测试数据：放上了优化前后的对比表格，包含不同 SeqLen 下的 Latency 和 Throughput。
4. 兼容性说明：确认该修改不影响其他架构（如 gfx90a）的性能，甚至做了回退处理。

## Performance Improvement for gfx942 Long Context

### Problem
Fixed tiling strategy caused low LDS utilization on MI300X for seq_len > 16k.

### Solution
- Implemented dynamic tiling based on sequence length.
- Aligned block sizes with Wavefront dimensions (64).

### Benchmarks (MI300X, ROCm 6.2)
| Seq Len | Baseline (ms) | Optimized (ms) | Improvement |
|---------|---------------|----------------|-------------|
| 16384   | 45.2          | 38.1           | +15.7%      |
| 32768   | 98.5          | 72.3           | +26.6%      |

提交后的 Code Review 环节依然严谨。Maintainer 指出我的注释中有一处关于寄存器数量的描述不够准确，建议修正以避免误导后续开发者。我立刻进行了修改，并补充了相关的理论依据。

合并与反思：生态是共建出来的

两天后，PR 被正式合并。看着自己的 Commit 出现在 TileLang 的主分支上，那种成就感远比单纯跑通一个 Demo 要强烈得多。

这次经历让我深刻体会到，ROCm 生态的完善不是靠官方文档单打独斗，而是靠每一个遇到问题的开发者伸手推一把。从发现性能异常，到编写复现脚本，再到社区讨论、代码实现、最终合并，这是一条完整的闭环。

如果你也在使用 SGLang、TileLang 或 LLaMA-Factory 时遇到了奇怪的报错或性能瓶颈，别急着绕道走。花点时间定位问题，提个 Issue，甚至尝试修一下。你的每一次贡献，都在让这块硬件更好用一点。开源的魅力，就在于此。

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