tilelang-cuda-memory

TileLang CUDA 内存访问优化

内存访问是 GPU 性能的关键瓶颈。本文档提供 TileLang CUDA 的内存访问优化策略。

---

1. GPU 内存层次

内存带宽和延迟

内存类型	带宽 (A100)	延迟	容量	TileLang API
寄存器	~19 TB/s	1 cycle	256 KB/SM	T.alloc_fragment / T.alloc_local
共享内存	~19 TB/s	~20 cycles	164 KB/SM	T.alloc_shared
L2 缓存	~5 TB/s	~100 cycles	40 MB	T.use_swizzle 优化
全局内存 (HBM)	~2 TB/s	~400 cycles	40/80 GB	T.Tensor

优化原则

• 减少全局内存访问: 利用共享内存和寄存器缓存数据
• 合并访问 (Coalesced Access): 使用 T.copy 进行高效数据传输
• 提高 L2 缓存命中率: 使用 T.use_swizzle 优化数据局部性

---

2. 内存分配最佳实践

共享内存（频繁访问的数据）

# 共享内存用于缓存从全局内存加载的数据块
A_shared = T.alloc_shared((block_M, block_K), "float16")
B_shared = T.alloc_shared((block_K, block_N), "float16")

# 高效数据加载
T.copy(A[by * block_M, ko * block_K], A_shared)

适用场景：

• 矩阵乘法中的输入块
• 多次访问的中间数据
• 需要线程间共享的数据

寄存器片段（累加器和临时存储）

# 寄存器用于累加和局部计算
C_local = T.alloc_fragment((block_M, block_N), "float")
T.clear(C_local)

# 累加操作
T.gemm(A_shared, B_shared, C_local)

适用场景：

• 矩阵乘法累加器
• 归约的临时结果
• 局部计算结果

本地内存（线程私有存储）

# 本地内存用于线程私有变量
C_reg = T.alloc_local((1,), "float")
T.clear(C_reg)

适用场景：

• 单个线程的累加值
• 线程局部的临时变量

---

3. 数据传输优化

使用 T.copy 进行合并访问

# ✅ 推荐：使用 T.copy 自动合并访问
T.copy(A[by * block_M, ko * block_K], A_shared)
T.copy(B[ko * block_K, bx * block_N], B_shared)

# ✅ 结果写回
T.copy(C_local, C[by * block_M, bx * block_N])

使用 T.Parallel 进行并行数据复制

# 并行数据复制
for k, j in T.Parallel(block_K, block_N):
    B_shared[k, j] = B[ko * block_K + k, bx * block_N + j]

使用向量化加载

# 向量化加载以提高带宽利用率
for k in T.vectorized(TILE_K):
    A_local[k] = A[bk * BLOCK_K + tk * TILE_K + k]
    B_local[k] = B[bn * BLOCK_N + tn, bk * BLOCK_K + tk * TILE_K + k]

---

4. 软件流水线

基本用法

# 使用 T.Pipelined 实现软件流水线
for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
    # 数据加载和计算自动重叠
    T.copy(A[by * block_M, ko * block_K], A_shared)
    T.copy(B[ko * block_K, bx * block_N], B_shared)
    T.gemm(A_shared, B_shared, C_local)