studypath

01.2-task2

6分钟阅读

2.2 任务二：TileLang Add 与 NineToothed Vector Add

目标
1. TileLang 任务：Vector Add ：实现 1-D float16 向量加法，并通过测试。
2. 九齿任务：Vector Add ：实现一维 float16 向量加法并处理尾块，并通过测试。
2.2.1 TileLang 任务：Vector Add

实现 1-D float16 向量加法：C[i]=A[i]+B[i]。1<=N<=1M，必须处理尾块（N 不整除 BLOCK_N）。

要求：在 solution.py 实现 tl_add_1d(A,B,BLOCK_N)；使用 T.Kernel、T.Parallel； 测试 N=1、127、1024、100003、1048576；与 PyTorch 逐元素校验；比较至少两种 BLOCK_N。

运行：

  1. cd /data/gollamago
  2. source ./setup_env.sh
  3. cd assignment/task2
  4. python -m pytest -q test_add.py
  5. python benchmark_add.py

benchmark_add.py 会对 N=1、127、1024、100003、1048576 逐一做正确性校验，并同时输出 TileLang 与 PyTorch 的平均耗时和最大误差。(不必要求TileLang比PyTorch更快)

步骤
步骤1：clone 项目到服务器
  1. 进入 Jupyterlab 的终端，执行clone命令
git clone https://gitlink.org.cn/ccf-ai-infra/gollamago.git
图片教程
(base) root@4a02ff7ad599:/data# ls
huggingface_home  lost+found
(base) root@4a02ff7ad599:/data# cd huggingface_home/hub/
(base) root@4a02ff7ad599:/data/huggingface_home/hub# git clone https://gitlink.org.cn/ccf-ai-infra/gollamago.git #进入工作区clone项目
Cloning into 'gollamago'...
remote: Enumerating objects: 337, done.
remote: Counting objects: 100% (337/337), done.
remote: Compressing objects: 100% (249/249), done.
remote: Total 337 (delta 158), reused 176 (delta 80), pack-reused 0
Receiving objects: 100% (337/337), 152.89 KiB | 2.64 MiB/s, done.
Resolving deltas: 100% (158/158), done. #这样的结果表示clone成功
(base) root@4a02ff7ad599:/data/huggingface_home/hub# ls
gollamago
(base) root@4a02ff7ad599:/data/huggingface_home/hub# cd gollamago/
(base) root@4a02ff7ad599:/data/huggingface_home/hub/gollamago# ls #查看项目
LICENSE  README.md  assignment  backends.py  benchmarks  image.png  infer.py  llama.py  operators  pytest.ini  requirements.txt  setup_env.sh  tests
(base) root@4a02ff7ad599:/data/huggingface_home/hub/gollamago#

步骤2：配置 TileLang 环境
  1. 在终端继续执行以下命令，配置 TileLang + 沐曦 MXMACA 运行环境
source ./setup_env.sh
图片教程
(base) root@4a02ff7ad599:/data/huggingface_home/hub/gollamago# source ./setup_env.sh
TileLang environment configured:
  source: /app/tilelang-metax
  native: /app/tilelang-metax/build/lib/libtilelang.so
  python: /opt/conda/bin/python
步骤3：补全Add函数，进行验证测试
  1. 打开 gollamago/assignment/task2/solution.py ，补全tl_add_1d函数

补全参考如下：

@tilelang.jit
def tl_add_1d(A, B, BLOCK_N: int):
    N = T.const("N")
    A: T.Tensor((N,), T.float16)
    B: T.Tensor((N,), T.float16)
    C = T.empty((N,), T.float16)
    # Step 1: 用 T.ceildiv 计算 block 数。
    num_blocks = T.ceildiv(N, BLOCK_N) # 需要多少个 block（向上取整，处理尾块）
    
    # Step 2: 在 T.Kernel 中取得 pid 并计算 base_idx。
    with T.Kernel(num_blocks, threads=BLOCK_N) as pid: # 启动 kernel；pid = 当前 block 编号
        base_idx = pid * BLOCK_N # 这个 block 负责起始下标
    
        # Step 3: 用 T.Parallel 遍历 tile 元素。
        for i in T.Parallel(BLOCK_N):
            index = base_idx + i
            
            # Step 4: 判断 index < N，保护尾块越界。
            if index < N:
            
                # Step 5: 写回 C[index] = A[index] + B[index]。
                C[index] = A[index] + B[index]
 
    return C
  1. 测试Add是否补全正确，进入终端执行如下命令
cd assignment/task2
python -m pytest -q test_add.py

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

python benchmark_add.py
图片教程

进入 gollamago/assignment/task2/solution.py 按要求补全任务

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

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

报错处理
  1. 缺gpu报错：如果你执行python -m pytest -q test_add.py看到了这样的报错，是因为你在无GPU开机模式，解决办法是关机重启，选择GPU开机。

  1. 未实现add报错：gollamago/assignment/task2/solution.py 补全任务未完成，按步骤实现 Add函数即可

  1. TypeError: ‘module’ object is not callable 报错：检测Kernel/Parallel首字母是否写成了小写

提交

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

2.2.2 九齿任务：Vector Add

在 ninetoothed_add.py 中完成 arrangement() 和 application()，实现一维 float16 向量加法并处理尾块。默认使用 BLOCK_SIZE = 1024。

arrangement() 负责按 block 对齐输入和输出，application() 负责完成 tile 内的逐元素 相加。NINETOOTHED_AUTOTUNE=1 时，block size 会在 256 到 1024 之间自动搜索。

参考：NineToothed Vector Addition。

  1. python -m pip install ninetoothed
  2. python -m pytest -q test_ninetoothed_add.py
  3. python benchmark_ninetoothed_add.py

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

  1. NINETOOTHED_AUTOTUNE=1 python benchmark_ninetoothed_add.py

测试使用 size=98432；benchmark 对 2^18 到 2^27 的向量长度比较九齿与 PyTorch 的耗时。

提交完成ninetoothed_add.py的测试与 benchmark 输出截图，并用 3 到 5 句话记录九齿的 分块方式和开发体验。

步骤
步骤1：安装 ninetoothed 环境
  1. 进入 jupyterlab 的终端执行以下命令
python -m pip install ninetoothed
图片教程

步骤2：补全 切分与加法 函数
  1. 打开 gollamago/assignment/task2/ninetoothed_add.py，补全arrangement() 和 application()函数

补全参考如下：

def arrangement(lhs, rhs, output):
    return (
        lhs.tile((BLOCK_SIZE,)),
        rhs.tile((BLOCK_SIZE,)),
        output.tile((BLOCK_SIZE,)),
    )
 
 
def application(lhs, rhs, output):
    output = lhs + rhs
  1. 测试arrangement() 和 application()是否补全正确，进入终端执行如下命令
python -m pytest -q test_ninetoothed_add.py

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

python benchmark_ninetoothed_add.py
图片教程

进入 gollamago/assignment/task2/ninetoothed_add.py，按要求补全arrangement() 和 application()函数

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

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

关系图谱

反向链接