TileLang GEMM optimization

TileLang GEMM on C500 优化复盘

这篇文章整理的是这一轮 TileLang GEMM on MetaX C500 的完整优化过程：我们到底在测什么、baseline 是什么、TileLang kernel 做了哪些优化、profiler 告诉我们什么、为什么现在和 mcBLAS 还有明显差距，以及下一步应该怎么继续做。

这篇文章将讲清楚：为什么要这么改、每个优化解决了什么问题、哪些方向被证明不值得继续投入、以及后续应该从高层 schedule 转向更底层 lowering / layout 优化。

1. 任务内容

优化的目标是一个 fp16 GEMM：

C[M, N] = A[M, K] @ B[K, N]

M = 131072
N = 2048
K = 7168
dtype = fp16
accum = fp32

这个 shape 来自 MoE 配置。测评脚本会把 MoE 里的 batch、seq_len、top_k、hidden_size、intermediate_size 转换成一个大的 dense GEMM：

original_tokens = batch_size * seq_len
routed_tokens = original_tokens * top_k
M = routed_tokens
K = hidden_size
N = intermediate_size

当前默认配置是：

batch_size        = 4
seq_len           = 8192
top_k             = 4
hidden_size       = 7168
intermediate_size = 2048

所以：

M = 4 * 8192 * 4 = 131072
K = 7168
N = 2048

输入矩阵为：

A: [131072, 7168], fp16, contiguous
B: [7168, 2048],  fp16, contiguous
C: [131072, 2048], fp16

测评脚本是 benchmarks/collect_gemm_table.py。它会输出三类结果：GEMM baseline、TileLang GEMM、Fused MoE GEMM。其中 GEMM baseline 是通过 torch.mm(A, B, out=C_base) 测普通 GEMM，在 MetaX PyTorch 后端上通常会调用底层 mcBLAS / 厂商 GEMM 实现；TileLang GEMM 则是我们当前优化的 TileLang kernel。

优化的正式 evaluator 是：

cd /data/tilelang-metax
/opt/conda/bin/python3 benchmarks/collect_gemm_table.py \
  --warmup 10 \
  --repeat 50 \
  --output-dir <result_dir>

接受一个优化必须同时满足三个条件：

1. TileLang GEMM latency_ms_median 下降；//latency
2. max_abs_error = 0.0；//correctness
3. 生成代码仍然包含 C500 MMA：
   __builtin_mxc_mma_16x16x16f16
   MMA_COUNT=2//确认走了 C500 的高性能 MMA 路径,没有退化为普通乘加循环

2. 当前最终结果

当前最佳 TileLang GEMM 版本是：

A copy coalesced_width=4
//在A_global -> A_shared中,
//让每次 copy 尽量按 4 个元素一组进行合并/向量化访问
B prepack/interleave4
//B_global -> B_shared -> MMA,把 K 方向每 4 个 fp16 打包在一起,
//在物理存储上改为[K/4, N, 4]，
//让 B 沿 K 方向每 4 个 fp16 连续放在一起
//让 kernel 内一次能取 4 个连续值
//降低 B shared memory 访问碎片化
//让 B_shared 更容易高效喂给 MMA
panel2 + 512 threads
//让相邻 threadblock 的访问顺序更适合 B/L2 复用
//改善 block raster/order
//让更多线程一起搬数据、隐藏等待，提高计算单元利用率
C_shared swizzle256
C store coalesced_width=4
//这两个优化是改善C_local -> C_shared -> C_global的拜访访问写回

A_global
   |
   |  A copy coalesced_width=4
   v
A_shared
   |
   |  feed MMA
   v
MMA  <---------------- B_shared
                           ^
                           |
              B prepack/interleave4
              B vectorized copy
                           |
                        B_global


MMA output
   |
   v
C_local
   |
   |  C_local -> C_shared
   v
C_shared
   |
   |  C_shared swizzle256
   |  C store coalesced_width=4
   v
C_global


Block scheduling:
   panel2 + 512 threads
   affects block order, parallelism, L2/MTE behavior, MMA duty

正式结果如下：

版本 TileLang GEMM median latency TFLOPS max_abs_error
C_shared/swizzle256/panel4 36.166145 ms 106.405885 0.0
C_shared/swizzle128/panel4 36.160896 ms 106.421331 0.0
B prepack/interleave4 35.812864 ms 107.455541 0.0
panel2 + 512 threads 33.586561 ms 114.578288 0.0
当前最终 A copy coalesced_width=4 32.820992 ms 117.250897 0.0

相对上一版 panel2 + 512 threads，最终版本又提升了：

latency improvement = 2.3326%
speedup             = 1.0233x

从最早的 36.16ms 左右到现在的 32.82ms，整体 TileLang GEMM 性能从大约 106 TFLOPS 提升到 117 TFLOPS。

3. 与 mcBLAS baseline 的关系

当前 benchmark 里，mcBLAS baseline 对应脚本输出里的：

GEMM baseline

它使用的是：

torch.mm(A, B, out=C_base)

输入是原始 contiguous A 和原始 contiguous B：

A: [131072, 7168], fp16, contiguous
B: [7168, 2048],  fp16, contiguous
C: [131072, 2048], fp16, contiguous output

注意这里没有显式调用 mcBLAS C API，而是通过 MetaX PyTorch 的 torch.mm 走底层 GEMM 库路径，通常就是 mcBLAS / 厂商 BLAS。

当前已知正式对比结果是：

GEMM baseline: 16.474111557006836 ms, 233.59625093586484 TFLOPS
TileLang GEMM: 32.82099151611328 ms, 117.25089704638276 TFLOPS
TileLang max_abs_error: 0.0

也就是说，TileLang 当前大约是 mcBLAS dense GEMM baseline 的一半速度。文档里也明确指出，当前 TileLang kernel 仍慢于 mcBLAS baseline，但 TileLang 的定位是一个可控 schedule、可插入 MoE 或自定义算子逻辑的实验路径；优化时主要比较 TileLang 自己不同版本之间的 latency 和 TFLOPS，而不是要求短期超过 mcBLAS。

mcBLAS 是厂商高度优化的 GEMM 库，规则大 GEMM 是它擅长的场景。TileLang 的价值不应该只用裸 GEMM 能不能打赢 mcBLAS 来衡量，而应该看它是否能在 MoE、fused kernel、特殊 layout、grouped GEMM、减少中间写回等场景里发挥可控 schedule 的优势。

4. TileLang GEMM 和 mcBLAS GEMM 的输入差异

数学上，TileLang GEMM 和 mcBLAS GEMM 算的是同一个东西：

A[M, K] @ B[K, N] = C[M, N]

但是物理输入 layout 不完全一样。

mcBLAS baseline 使用的是原始 B：

B: [7168, 2048], row-major contiguous

TileLang GEMM 为了优化 C500 上 shared-memory / MMA feed 路径，不直接把原始 B 传给 kernel，而是先做了 B 预打包：

def pack_b_interleave4(B: torch.Tensor) -> torch.Tensor:
    K, N = B.shape
    return B.reshape(K // 4, 4, N).permute(0, 2, 1).contiguous().view(K, N)

也就是逻辑上还是：

B[K, N]

但物理上变成类似：

B_packed[K/4, N, 4]

当前 shape 下就是：

logical B:  [7168, 2048]
physical:   [1792, 2048, 4]

TileLang 版本的输入差异可以总结成：

项目 mcBLAS / GEMM baseline TileLang GEMM
A 输入 原始 A [M, K] 同一个原始 A [M, K]
B 输入 原始 B [K, N] pack_b_interleave4(B) 后的 B_tile
B 物理 layout row-major contiguous [K/4, N, 4] interleave4
输出 C_base [M, N] C_tile [M, N]
正确性对比 baseline 自身作为参考 C_tile 对比 C_base
预打包耗时 无 不计入 GEMM timed region

这意味着当前对比不是完全相同物理 layout 的 apples-to-apples 对比。它的真实含义是：

mcBLAS: 原始 row-major B 的通用 GEMM
TileLang: 预打包权重布局下的专用 GEMM kernel

但数学结果仍然严格对比 torch.mm(A, original_B)，并要求 max_abs_error=0.0。

这个设置是否合理，取决于我们的目标场景。如果 B 是动态输入，那么 B pack 成本必须计入；但如果是 MoE 推理里的静态专家权重，B 预打包可以在模型加载或初始化阶段完成，不计入 timed region 是合理的。报告里写清楚：当前 TileLang GEMM 是静态权重 / 预打包权重场景下的专用 kernel。

5. 早期瓶颈

最早 profiler 看到的现象是：

指标 现象
BSM / Workgroup Memory instructions 很高，甚至超过 MMA 指令数量
AP MMA Duty 约 43%～44%，不够高
AP MTE Duty 约 10%～11%
L2C Duty 很低
L2 hit rate 较高

这个组合非常关键。

BSM / Workgroup Memory instructions 很高，甚至超过 MMA 指令数量，说明MMA这种真正用来计算的指令占比小，共享存储访问命令占比高，是需要注意的。
如果是纯 HBM 带宽瓶颈，我们通常会看到 global memory / L2 相关指标特别紧张，计算单元在等外部内存。但这里 L2 hit rate 命中率较高，L2C Duty 负载很低，说明问题不是 HBM 或 L2 带宽不够。真正明显的是 BSM / Workgroup Memory instructions 很高，AP MMA Duty 不够高。memory / 片上数据搬运问题严重

因此当时的判断是：

不是 HBM/L2 带宽瓶颈。
主要瓶颈在 shared/workgroup memory 路径，
也就是 BSM 和 A/B/C shared tile 访问。

所以后续优化重点不是盲目扩大 block size，也不是盲目增加 pipeline stage，而是围绕下面三条路径做：

global -> shared
shared -> MMA
C_shared -> global

这个判断决定了后面所有优化的方向。

6. GEMM kernel 的数据流怎么理解？

可以把当前 TileLang GEMM 分成四个阶段：

1. 从 global memory 读取 A tile 和 B tile
2. 把 A/B tile 搬到 shared memory
3. 从 shared memory 喂给 C500 MMA 做矩阵乘
4. 把 C_local 的结果写回 global memory

更具体一点：

A_global -> A_shared -> MMA
B_global -> B_shared -> MMA
MMA accumulates into C_local
C_local -> C_shared -> C_global

这里最理想的状态是 MMA 一直有数据可计算，global load、shared load、store 都尽量规整、合并、低冲突。但当前 profiler 显示，MMA 并没有被完全喂饱，shared/workgroup memory 路径压力很大。

所以这轮优化的核心目标不是“让数学计算变少”，而是：

让数据更顺地从 global 进入 shared，
再从 shared 喂给 MMA，
最后更顺地写回 global。

7. 优化一：C_shared staging，先让 C store 路径稳定下来

7.1 改了什么？

C 结果的写回没有直接走：

C_local -> C_global

而是走：

C_local -> C_shared -> C_global

对应代码是：

T.copy(C_local, C_shared)
T.copy(C_shared, C[by * block_M, bx * block_N], coalesced_width=4)

7.2 为什么多走一次 shared 反而更快？

直觉上，多经过一次 shared memory 应该更慢。但在 GPU kernel 里，真正重要的不是“搬了几次”，而是最后生成的访存指令是否规整、是否合并、是否符合后端 lowering 的最佳路径。

当时测试发现，直接 C_local -> C_global 在 C500 上生成的 store 路径较差，latency 反而更慢；通过 C_shared staging 后，global store 更稳定，并且最终 store 可以使用 coalesced_width=4。

所以这个优化的作用是：

用一次 shared staging 的额外成本，
换取更稳定、更规整的 global store path。
//这里的global stroe path 是 kernel 算完 C 之后，把结果写回全局显存 C 矩阵的那条路径。

7.3 But at what cost 但优化的代价是什么

它不是完美方案。C_shared staging 会引入额外 shared memory 指令，因此当前剩余瓶颈里仍然包括：

C_shared staging 的额外 BSM/store 指令

从长期看，更理想的方向是让 direct global store 也能生成高效 store path。但在当前 TileLang/MXMACA lowering 质量下，C_shared staging 是一个有效 workaround。

8. 优化二：C_shared swizzle，减少 shared memory 冲突

8.1 改了什么？

swizzle是有规律地“打乱 / 重排”数据的存储位置，让并行访问更均匀、更少冲突。
C_shared layout 加了 swizzle：

C_shared: T.Layout(
    C_shared.shape,
    lambda i, j: get_swizzle_layout(
        i,
        j,
        C_shared.shape[-1],
        C_shared.dtype,
        swizzle_bytes=256,
    ),
)

最终当前版本用的是：

swizzle_bytes = 256

8.2 为什么要 swizzle？

shared memory 可以理解成很多 bank。如果多个线程同时访问的地址落到冲突的 bank 上，就会产生 bank conflict，导致 shared memory access efficiency 下降、conflict cycles 增加。

C_shared swizzle 的作用是改变数据在 shared memory 中的物理排列方式，让线程访问 C_shared 时更分散、更规整，从而减少冲突。

这一轮测试过：

swizzle_bytes = 32 / 64 / 128 / 256 / 512 / 1024

最后保留 256 的原因是：512/1024 可能超 shared memory 限制或不稳定；128/256 在不同阶段表现较好；当前 panel2 + 512 threads 配置下 256 最好。

9. 优化三：B prepack / interleave4，解决 B feed MMA 的低效访问

9.1 改了什么？

这是这一轮非常关键的优化。
gemm运算中A[M, K] @ B[K, N] = C[M, N]

原始 B 是：

B[K, N]

TileLang 版本先把 B 预打包成 interleave4 layout：

def pack_b_interleave4(B: torch.Tensor) -> torch.Tensor:
    """Pack logical B[K, N] as physical [K/4, N, 4] outside timed regions."""
    if B.dim() != 2 or B.shape[0] % 4 != 0:
        raise ValueError("pack_b_interleave4 expects B with shape (K, N) and K divisible by 4")
    K, N = B.shape
    return B.reshape(K // 4, 4, N).permute(0, 2, 1).contiguous().view(K, N)

逻辑上仍然是：

B[K, N]

但物理存储变成：

[K/4, N, 4]

这样 kernel 内沿 K 方向可以一次访问 4 个连续 fp16，减少 B shared-memory 访问里的低效 scalar/BSM pattern。

kernel 内还配合了 layout annotation：

B: T.Layout(lambda i, j: (i // 4, j, i % 4))
B_shared: T.Layout(lambda i, j: (i // 4, j, i % 4))

以及手写 vectorized copy：

for kg, jj in T.Parallel(block_K // b_interleave, block_N):
    for rr in T.vectorized(b_interleave):
        B_shared[kg * b_interleave + rr, jj] = \
            B[k * block_K + kg * b_interleave + rr, bx * block_N + jj]

9.2 优化的作用？

MMA 本身吃的是 fragment。B 从 global memory 进 shared memory，再从 shared memory 喂给 MMA。如果 B 的物理布局不适合 MMA 消费方式，就会出现碎片化 shared load、更多 BSM 指令、更多低效访问。

B interleave4 的作用就是让 B 的物理 layout 更贴近 C500 MMA feed 的访问模式。

通俗讲，就是原来 B 的数据摆放方式不够顺，线程取 B 时有点“东一块西一块”；现在把 B 按 K 方向每 4 个打包在一起，让每次取数更连续。

9.3 profiler 证据

B prepack/interleave4 后，对比 swizzle128 版本，profiler 指标变化非常明显：

指标 swizzle128 B prepack/interleave4 变化
Total Instructions 24.603B 18.658B -24.16%
Memory Instructions 9.208B 4.820B -47.65%
BSM instructions 8.424B 3.846B -54.35%
MTE instructions 5.924B 4.087B -31.00%
MMA instructions 6.107B 6.109B 基本不变

这里最关键的是：

MMA instructions 基本不变
BSM / Memory instructions 大幅下降

这说明数学计算路径没有变，MMA 计算量没有减少，只是数据搬运变得更好了。最终 latency 从约 36.16ms 推进到 35.81ms。

10. 优化四：panel2 + 512 threads，改善 block 调度和 AP 利用率

10.1 改了什么？

这一版主要改了两个点：

T.use_swizzle(panel_size=2, order="column", enable=True)
threads = 512

也就是：

threadblock swizzle: column, panel_size=2
block threads: 512

10.2 为什么这个会快？

B prepack 之后，shared 路径的一部分压力被降低了，global/L2/MTE 路径的重要性上升。panel_size=2 改善 threadblock/raster 顺序，让相邻 block 对 B/L2 的访问更友好；512 threads 则改善 wave/occupancy 和 AP 利用率。

这一版结果从：

35.812864 ms -> 33.586561 ms

是这一轮中比较明显的一次提升。

10.3 profiler 证据

对比上一版 B prepack/panel4/256t，panel2 + 512 threads 的 profiler 变化是：

Metric previous bpack/panel4/256t panel2/512t Change
Total Instructions 18.658B 21.254B +13.92%
Memory Instructions 4.820B 6.159B +27.77%
Global Read Instructions 952.760M 763.565M -19.86%
L2C Read Instructions 12.188B 9.161B -24.83%
Workgroup total instructions 3.846B 5.371B +39.64%
Average WG conflict cycles / instruction 2.40 1.72 -28.33%
Shared memory access efficiency 66.67% 72.70% +9.04%
AP MTE Duty ratio 7.63% 5.18% -32.11%
AP MMA Duty ratio 44.51% 47.10% +5.82%

这个结果很有意思：它不是所有指标都变好了。

变好的地方是：

Global Read Instructions 下降
L2C Read Instructions 下降
Average WG conflict cycles 下降
Shared memory access efficiency 提升
AP MTE Duty ratio 下降
AP MMA Duty ratio 提升

变差的地方是：

Total Instructions 上升
Memory Instructions 上升
Workgroup total instructions 上升

所以这一版的本质是：

用更多 workgroup/shared 指令，
换来了更好的 global/L2/MTE 行为和更高 MMA 利用率，
最终 latency 下降。

profiler 的判断也很明确：panel2 + 512 threads 改善了 global/L2/MTE 路径，AP MMA Duty 提高，说明计算单元利用率更好；但 Workgroup instructions 上升，说明 shared-memory 路径仍然有压力。因此下一步不应该继续盲目加 threads 或改大 block，而应该针对 copy 路径做小优化。

11. 优化五：A copy coalesced_width=4，改善 A global-to-shared copy

11.1 改了什么？

最后一轮 accepted 优化是给 A global-to-shared copy 加上：

T.copy(A[by * block_M, k * block_K], A_shared, coalesced_width=4)

也就是：

A copy coalesced_width = 4

11.2 为什么优化 A copy？

每个 K iteration 都需要把 A tile 从 global memory 搬到 A_shared：

A_global -> A_shared

前面我们已经优化了 B 的 layout 和 B_shared feed；panel2 + 512 也改善了 global/L2/MTE 路径。但 Workgroup instructions 仍然高，copy path 仍然有优化空间。

coalesced_width=4 的作用是让 fp16 路径生成更合理的 vectorized global load。source verify 里也看到 A load 使用了 uint2。这个优化不改变数学结果，也不破坏 C500 MMA lowering。

最终结果是：

33.586561 ms -> 32.820992 ms

并且生成代码确认仍然包含：

__launch_bounds__(512, 1)
MMA_COUNT=2
__builtin_mxc_mma_16x16x16f16
A load 使用 uint2

这说明最后一轮优化是一个相对干净的 copy path 优化：数学不变，MMA 不变，只是 A load 更规整。

12. 当前最终 kernel 可以怎么概括？

当前最终版本可以概括成下面这几层：

B 侧：
    B 预打包成 interleave4
    B 和 B_shared 使用 interleave4 layout
    B copy 使用 vectorized(4)

A 侧：
    A global-to-shared copy 使用 coalesced_width=4

C 侧：
    C_local -> C_shared -> C_global
    C_shared 使用 swizzle_bytes=256
    C_shared -> C_global 使用 coalesced_width=4

调度侧：
    threadblock swizzle 使用 column panel_size=2
    threads = 512

lowering 配置：
    启用 LDG/STG lowering
    保留当前 C500 上表现稳定的 TileLang pass config
    保留 C500 MMA lowering

换一种更直观的说法：

B 重新打包，让 B 更好喂给 MMA；
A copy 做合并访问，让 A 进 shared 更顺；
C 先在 shared 整理，再更稳定地写回 global；
block 调度改成对 L2/MTE 更友好；
线程数提升到 512，提高 wave/occupancy 和 AP 利用率。

13. 被测试但没有接受的方向

这轮优化不是只保留成功的结果，也测试了很多失败方向。

13.1 更大 block shape

测试过：

256x128x64
128x256x64
256x256x64

问题是：

shared memory 超 64KB
或者 latency 变慢

这说明这个 kernel 不能简单靠扩大 tile 来提升性能。更大的 tile 会带来 shared memory 压力、occupancy 下降、寄存器压力增加等问题。

13.2 更多 pipeline stages

测试过：

num_stages = 2 / 3 / 4

结果是：

stage=2 明显变慢
stage=3/4 通常 shared memory 超限

这说明当前瓶颈不是简单靠更多 pipeline stage 就能隐藏掉。更多 stage 带来的资源压力超过了它隐藏访存的收益。

13.3 direct C store

尝试去掉 C_shared：

C_local -> C_global

但结果变慢。原因是虽然少了一次 shared staging，但生成的 global store 路径更差。

13.4 B interleave2/8/16

测试过：

interleave2
interleave8
interleave16

结果是：

interleave2 慢
interleave8/16 更慢
interleave4 最合适

说明 interleave4 是当前 C500 fp16 MMA feed 路径下比较合适的折中。interleave 太小，连续访问不够；interleave 太大，可能破坏 layout、增加冲突或地址计算压力。

13.5 C 分片 staging

测试过：

C stage N=64/32/16
C stage M=64/32

问题是有些看似更快，但 max_abs_error 很大，直接拒绝；有些 layout inference 失败；正确版本又不如当前方案。

13.6 GEMM warp policy

测试过：

Square
FullRow
FullCol

结果是：

FullRow / FullCol 明显变慢
Square 基本等价，无收益

这些失败方向也很重要，因为它们说明当前优化已经不是简单扫高层参数可以继续大幅推进的阶段。

14. 为什么 TileLang GEMM 和 mcBLAS 速度差距这么大？

当前结果是：

mcBLAS baseline: 16.474 ms, 233.60 TFLOPS
TileLang GEMM:   32.821 ms, 117.25 TFLOPS

也就是 TileLang 约为 mcBLAS 的 50%。

我认为主要原因有五个。

14.1 mcBLAS 是厂商库，规则大 GEMM 是它的主场

当前 shape 是：

M = 131072
N = 2048
K = 7168

这是一个非常规则的大 GEMM：

A [131072, 7168] @ B [7168, 2048]

没有 ragged tile，没有 grouped expert，没有 routing，没有 scatter/gather，也没有 activation fusion。这种大而规则的 dense GEMM 正是 mcBLAS 这类厂商库最擅长的场景。

mcBLAS 大概率已经做了大量 C500-specific 优化：

更成熟的 block/wave 映射
更好的 shared memory layout
更低冲突的 fragment feed
更好的 global load / store 指令选择
更强的 pipeline
更好的 epilogue
更精细的 instruction scheduling

TileLang 当前虽然已经用上了 C500 MMA，但 shared feed、store path、lowering 质量还没有达到厂商库水平。

14.2 TileLang 当前 AP MMA Duty 还不够高

早期 profiler 里 AP MMA Duty 大约 43%～44%；panel2 + 512 threads 后提升到约 47.10%。这说明 MMA 单元利用率有提升，但还远没有达到理想状态。

换句话说：

MMA 单元没有被持续喂饱。

造成这个问题的可能原因包括：

A/B shared tile feed 不够高效
BSM load/store mapping 不够好
fragment load mapping 不够细
copy 与 compute overlap 不够强
pipeline stage 受 shared/register 限制
C_shared staging 引入额外 shared 指令

14.3 TileLang 有 C_shared staging 的额外成本

当前 TileLang 为了让 C store 路径稳定，使用：

C_local -> C_shared -> C_global

这在当前后端下是有效优化，但它本质上是 workaround，会带来额外 shared memory 指令。mcBLAS 很可能有更成熟的 epilogue store path，不需要这么绕一圈。

所以 TileLang 的 C store 当前是：

为了得到更好的 global store，
付出了一次 shared staging 的成本。

14.4 TileLang 和 mcBLAS 的物理输入 layout 不完全相同

TileLang 用的是预打包后的 B，而 mcBLAS 用的是原始 row-major B。当前测评文档也明确说明，这不是完全同一物理输入 layout 的 apples-to-apples 对比；TileLang 是预打包权重布局下的专用 GEMM kernel，mcBLAS 是原始 row-major B 的通用 GEMM。

不过这点并不是 TileLang 慢的主要借口。相反，TileLang 已经拿到了预打包 B 的优势，但仍然只有 mcBLAS 一半速度，说明瓶颈更深层，主要在 shared feed、lowering、pipeline 和 epilogue。

14.5 当前高层 schedule knob 的收益已经变小

已经测试过 block shape、num_stages、warp policy、direct C store、B interleave 多种参数。很多方向不是变慢，就是 shared memory 超限，或者 correctness 失败。

当前文档里的判断是：后续单纯高层 schedule knob 的收益会变小，更可能需要更底层的 TileLang/MXMACA lowering 修改、C500-specific shared layout、让 direct global store 生成高效 store path、以及更精细控制 A/B fragment load mapping。

15. 这个任务设置合理吗？

我觉得要分两层看。

如果目标是：

用 TileLang 手写裸 GEMM 短期超过 mcBLAS

那这个任务设置不太合理。因为当前 shape 是规则大 dense GEMM，mcBLAS 正好是厂商库的最强场景。

但如果目标是：

通过一个标准大 GEMM 练习 TileLang/C500 优化方法，
验证 MMA lowering、shared layout、copy coalescing、threadblock swizzle、profiler 分析方法，
并为后续 MoE fused kernel 打基础

那这个任务非常合理。

所以汇报时建议把 mcBLAS 定位成：

C500 dense GEMM 性能上限 / roofline reference

而不是把它当成 TileLang 短期必须超过的目标。

TileLang 当前应该重点和自己的历史版本比：

36.16 ms -> 35.81 ms -> 33.59 ms -> 32.82 ms

这条曲线说明优化方向是有效的。

16. 当前剩余瓶颈判断

当前最终版本已经做了：

B prepack/interleave4
A copy coalesced_width=4
panel2 + 512 threads
C_shared swizzle256
C store coalesced_width=4

目前判断：

1. 已经不是纯 HBM 带宽瓶颈；
2. B feed 的 BSM/shared 压力已经被 interleave4 明显降低；
3. panel2 + 512 降低了 global/L2/MTE 压力，并提高了 AP MMA Duty；
4. 剩余瓶颈大概率仍在：
   - A/B shared tile feed MMA 的 shared-memory 指令量；
   - C_shared staging 的额外 BSM/store 指令；
   - TileLang/MXMACA lowering 生成的 BSM load/store mapping。

也就是说，当前 kernel 的主要问题不在“有没有用 MMA”，而在：

MMA 前后的数据流是否足够顺。

更具体就是：

A/B 从 global 到 shared 是否足够 coalesced；
A/B 从 shared 喂给 MMA 是否足够低冲突；
C 从 local/shared 写回 global 是否足够高效；
TileLang lowering 是否能生成接近 mcBLAS 的底层指令模式。

17. 接下来应该做什么？

我建议下一步分成四条线。

17.1 重新采集当前最终版本 mcProfiler

当前最终版本是：

A_copy_coalesced_width=4
panel2
threads=512
B interleave4
C_shared swizzle256

下一步应该重新采集这个版本的 mcProfiler，和上一版 output20260612190053 对比。重点看：

Global Read Instructions
L2C Read Instructions
MTE instruction count
AP MTE Duty ratio
Workgroup load/store/total instructions
Average conflict cycles
Shared memory access efficiency
AP MMA Duty ratio

如果这些指标确认 A copy coalescing 的收益，就继续围绕 A/B shared feed 做低层 layout 优化；如果指标变化不明显但 latency 下降，那可能要重点看 source diff、instruction scheduling 或 lowering 差异；如果指标没有证明收益，下一步就不应该继续盲扫高层参数，而应该转向 lowering 层。

17.2 不要继续大规模盲扫高层参数

已经测试过的方向说明：

更大 block shape 容易超 shared memory 或变慢
更多 pipeline stages 不一定能隐藏访存，反而增加资源压力
direct C store 当前 lowering 路径较差
B interleave4 是当前最优折中
warp policy 没有明显收益

所以继续大规模扫：

block_M / block_N / block_K
num_stages
threads
warp policy

收益可能很低。

后续更值得做的是：

C500-specific shared layout
A/B fragment load mapping
BSM 指令生成模式
direct global store lowering
TileLang/MXMACA pass 优化

17.3 回到 MoE 真实场景，不要只看裸 GEMM

裸 GEMM 是 mcBLAS 的主场。TileLang 的优势应该体现在 MoE 场景里。

MoE 的真实路径不只是一个 GEMM，它还包括：

routing / topk
token dispatch
expert grouping
grouped GEMM
activation
multiply
down projection
scatter / combine
中间 tensor 写回
kernel launch / sync / allocation

mcBLAS 在单个大 GEMM 上非常强，但 MoE 里可能有很多不规则、小 batch、grouped、fused 的计算。TileLang 真正的价值是：

把多个操作融合
减少中间 tensor
减少 global memory 写回
支持特殊 layout
针对 expert token distribution 做专用 schedule

所以后续 benchmark 应该分层：

1. Dense GEMM：作为 roofline reference
2. Grouped expert GEMM：模拟每个 expert token 数不同
3. gate/up fused GEMM：融合两个 projection
4. gate/up + activation + multiply
5. full MoE body：dispatch + expert GEMM + combine

如果只拿裸 GEMM 和 mcBLAS 比，很容易低估 TileLang 的价值。

17.4 规范每次优化的记录方式

后续每次改 kernel，都建议记录：

改动点
命令
输入 shape
TileLang latency / TFLOPS / max_abs_error
GEMM baseline latency / TFLOPS
是否仍有 C500 MMA
source verify 证据
mcProfiler 关键指标
是否接受
拒绝原因

测评流程建议保持：

Step 1: 跑正式 evaluator
Step 2: 确认 max_abs_error = 0.0
Step 3: grep source 确认 C500 MMA 仍存在
Step 4: 必要时跑 mcProfiler
Step 5: 记录指标和结论

文档里也建议重点关注 Total Instructions、AP MMA Duty ratio、AP MTE Duty ratio、MMA/MTE/BSM instruction count、Global Read Instructions、L2C Read Instructions / Hit Rate、Workgroup Memory load/store/total instructions、average conflict cycles、shared memory access efficiency、Achieved/Dispatched waves 等指标。

18. 可以直接用于汇报的总结

这轮工作做的不是“随便调几个参数”，而是一次比较完整的 profiler-driven GEMM 优化。

我们首先确定了目标 shape：

M=131072, N=2048, K=7168, fp16 accumulation fp32

然后建立了严格的接受标准：

latency 下降
max_abs_error = 0.0
C500 MMA lowering 仍然存在

接着通过 mcProfiler 判断瓶颈并不是单纯 HBM/L2 带宽，而是 shared/workgroup memory 路径，也就是 A/B/C shared tile 访问和 BSM 指令压力。基于这个判断，优化方向从“盲目扩大 block 或增加 stage”转成了“优化 global->shared、shared->MMA、C_shared->global”。

具体做了五类优化：

1. C_shared staging：
   用 C_local -> C_shared -> C_global 替代 direct C store，
   让 store path 更稳定。

2. C_shared swizzle256：
   降低 C staging 里的 shared memory conflict。

3. B prepack/interleave4：
   把 B 的物理 layout 从 [K, N] 改成 [K/4, N, 4]，
   让 B shared feed 更贴近 MMA 消费方式，
   大幅降低 BSM/workgroup memory 指令。

4. panel2 + 512 threads：
   改善 threadblock 调度、L2/MTE 行为和 AP MMA Duty。

5. A copy coalesced_width=4：
   让 A global-to-shared copy 生成更合理的 vectorized load，
   source 中看到 A load 使用 uint2。

最终 TileLang GEMM 从大约 36.16ms 优化到 32.82ms，性能从约 106 TFLOPS 提升到 117 TFLOPS，且 max_abs_error=0.0，仍然保留 C500 MMA。

当前和 mcBLAS baseline 仍有明显差距：mcBLAS 约 16.47ms / 233.6 TFLOPS，TileLang 约 32.82ms / 117.25 TFLOPS。这个差距主要因为 mcBLAS 是厂商高度优化的 dense GEMM 库，而当前 shape 又是非常规则的大 GEMM，正好是 mcBLAS 的主场。TileLang 当前的价值不是短期在裸 dense GEMM 上打赢 mcBLAS，而是通过可控 schedule 为 MoE fused/grouped/irregular kernel 打基础。

下一步不建议继续盲扫高层 schedule knob，而应该重新采集当前最终版本的 mcProfiler，确认 A copy coalescing 的硬件指标变化；之后重点转向 A/B shared feed、C500-specific shared layout、direct global store lowering、A/B fragment load mapping，以及 TileLang/MXMACA lowering 层优化。同时，评测应该逐步回到真实 MoE 场景，用 fused MoE body、grouped expert GEMM、activation fusion、减少中间写回等指标来体现 TileLang 的价值。

19. 一句话结论

这轮优化的核心成果是：我们把 TileLang GEMM 从一个能跑 C500 MMA 的基础版本，推进到了一个有明确 B layout、A copy、C store、threadblock swizzle 和 profiler 证据支撑的版本；性能从 36ms 推到 32.8ms，但剩余差距已经主要不是高层参数问题，而是 C500-specific shared feed 和 TileLang/MXMACA lowering 质量问题。

posted @ 2026-06-15 09:40  White_Swan  阅读(7)  评论(0)    收藏  举报