01.3-task3

2.3 任务三：TileLang Softmax 与 NineToothed GEMM

目标

1. TileLang 任务：Softmax：实现二维矩阵按行 Softmax
2. 九齿任务：GEMM：实现 C = A @ B

2.3.1 TileLang 任务：Softmax

实现二维矩阵按行 Softmax。输入 A:[N,M] float32，输出同 shape/dtype，沿最后一维计算稳定公式 exp(x-max)/sum(exp(x-max))。

要求：在 solution.py 实现 tl_softmax(A,BLOCK_N,BLOCK_M)；完成最大值归约、指数、求和归约和写回；综合 benchmark 使用 (1,1)、(3,7)、(16,256)、(17,513)、(64,4096) 五组配置。

运行：

cd /data/gollamago
source ./setup_env.sh
cd assignment/task3
python -m pytest -q test_softmax.py
python benchmark_softmax.py

步骤

步骤1：补全 tl_softmax 函数

打开 gollamago/assignment/task3/solution.py ，补全 tl_softmax 函数

补全参考如下：

@tilelang.jit
def tl_softmax(A, BLOCK_N: int, BLOCK_M: int):
    N, M = T.const("N, M")
    A: T.Tensor((N, M), T.float32)
    B = T.empty((N, M), T.float32)
 
    log2_e = 1.4426950408889634
 
    # Step 1: ceildiv 覆盖任意 N；
    num_row_blocks = T.ceildiv(N, BLOCK_N)
 
    # Step 2: 分配 tile fragment。
    with T.Kernel(num_row_blocks, threads=256) as pid:
        a_tile = T.alloc_fragment((BLOCK_N, BLOCK_M), T.float32)
        exp_tile = T.alloc_fragment((BLOCK_N, BLOCK_M), T.float32)
        tile_max = T.alloc_fragment((BLOCK_N,), T.float32)
        tile_sum = T.alloc_fragment((BLOCK_N,), T.float32)
        lse = T.alloc_fragment((BLOCK_N,), T.float32)
 
        T.fill(lse, -T.infinity(T.float32))
 
        # Step 3: ceildiv 扫描任意 M；
        for m_blk_id in T.Serial(T.ceildiv(M, BLOCK_M)):
            # Step 4: 显式填充行尾和列尾，避免越界位置的 0 污染 max/sum。
            for i, j in T.Parallel(BLOCK_N, BLOCK_M):
                row = pid * BLOCK_N + i
                col = m_blk_id * BLOCK_M + j
                if row < N:
                    if col < M:
                        a_tile[i, j] = A[row, col]
                    else:
                        a_tile[i, j] = -T.infinity(T.float32)
                else:
                    a_tile[i, j] = -T.infinity(T.float32)
 
            # Step 5: reduce_max；
            T.reduce_max(a_tile, tile_max, dim=1, clear=True)
 
            # Step 6: exp2 稳定指数；
            for i, j in T.Parallel(BLOCK_N, BLOCK_M):
                exp_tile[i, j] = T.exp2(
                    a_tile[i, j] * log2_e - tile_max[i] * log2_e
                )
 
            # Step 7: reduce_sum。
            T.reduce_sum(exp_tile, tile_sum, dim=1, clear=True)
 
            # Step 8: online 更新 running max/sum（或 lse）。
            for i in T.Parallel(BLOCK_N):
                lse[i] = tile_max[i] * log2_e + T.log2(
                    T.exp2(lse[i] - tile_max[i] * log2_e) + tile_sum[i]
                )
 
        # Step 9: 第二遍扫描并归一化；
        for m_blk_id in T.Serial(T.ceildiv(M, BLOCK_M)):
            for i, j in T.Parallel(BLOCK_N, BLOCK_M):
                row = pid * BLOCK_N + i
                col = m_blk_id * BLOCK_M + j
                if row < N:
                    if col < M:
                        B[row, col] = T.exp2(A[row, col] * log2_e - lse[i])
 
    # Step 10: 只写回有效行列。
    return B

测试Add是否补全正确，进入终端执行如下命令

cd assignment/task3
python -m pytest -q test_softmax.py

返回 6 passed 验证通过 3. 继续在task3目录执行如下命令，运行测试用例

python benchmark_softmax.py

图片教程

在 Jupyterlab 打开 gollamago/assignment/task3/solution.py，按要求补全函数

测试通过的结果如下，显示绿色的passed表示通过测试

运行 python benchmark_softmax.py 的结果如下，5个pass表示5个测试用例运行正确

提交

运行python benchmark_softmax.py后的终端结果截图

2.3.2 九齿任务：GEMM

在 ninetoothed_gemm.py 中完成 arrangement() 和 application()，实现 C = A @ B：输入和输出使用 float16，中间结果使用 float32 累加。默认 block size 为 64 x 64 x 64。

arrangement() 负责按照 M、N、K 三个维度组织矩阵 tile，application() 负责遍历 K 方向的 tile、使用 ntl.dot 累加，并将结果写回输出矩阵。输入规模不是 block size 的整数倍时也必须正确处理边界。

参考：NineToothed Matrix Multiplication。

python -m pip install ninetoothed
python -m pytest -q test_ninetoothed_gemm.py
python benchmark_ninetoothed_gemm.py

需要尝试自动调优时运行：

NINETOOTHED_AUTOTUNE=1 python benchmark_ninetoothed_gemm.py

测试使用 M=N=K=512；benchmark 对 M=N=K=2^3 到 2^12 的规模比较九齿与 torch.mm 的耗时。提交完成的 ninetoothed_gemm.py的测试输出和基线性能数据截图。

studypath

目录