forked from metax-maca/op_optimization
ADD file via upload
This commit is contained in:
yyyymmm
parent
7f89db4302
commit
b3cf281f26
|
|
@ -0,0 +1,267 @@
|
|||
# TileLang-MetaX Race 分支 `race_tests` 三大算子测试跑通指南
|
||||
|
||||
> 本文档记录如何在**模力方舟**（MetaX GPU）环境下，完成 `tilelang-metax` **race** 分支中 `race_tests` 目录下三个核心算子（MLA、MoE、NSA）的功能与性能测试。
|
||||
|
||||
---
|
||||
|
||||
## 1. 前置环境与硬件信息
|
||||
|
||||
### 1.1 硬件与驱动
|
||||
|
||||
```text
|
||||
MX-SMI 2.2.12
|
||||
Kernel Mode Driver Version: 3.0.11
|
||||
MACA Version: 3.5.3.20
|
||||
GPU: MetaX C500
|
||||
VRAM: 65536 MiB
|
||||
Sliced GPU: 50% Compute, 32000 MiB Vram Quota
|
||||
```
|
||||
|
||||
> 本次测试在 **模力方舟** 平台提供的 MetaX C500 容器镜像中完成，镜像已预装 MACA 工具链与 PyTorch（`torch 2.8.0+metax3.5.3.9`）。
|
||||
|
||||
### 1.2 软件版本
|
||||
|
||||
| 组件 | 版本 |
|
||||
|------|------|
|
||||
| Python | 3.12.11 |
|
||||
| PyTorch | 2.8.0+metax3.5.3.9 |
|
||||
| TileLang | 0.1.9+maca.gitf1ca0fb9 |
|
||||
| CMake | 4.3.2 |
|
||||
| GCC | 11.4.0 |
|
||||
|
||||
---
|
||||
|
||||
## 2. 仓库准备与编译安装
|
||||
|
||||
### 2.1 拉取代码并切换到 race 分支
|
||||
|
||||
```bash
|
||||
cd /data
|
||||
git clone https://github.com/tile-ai/tilelang-metax.git
|
||||
cd tilelang-metax
|
||||
git checkout race
|
||||
```
|
||||
|
||||
### 2.2 编译安装 TileLang（MACA 后端）
|
||||
|
||||
在容器内执行以下命令，启用 MACA 后端并重新编译：
|
||||
|
||||
```bash
|
||||
cd /app/tilelang-metax # 或 /data/tilelang-metax
|
||||
rm -rf build
|
||||
export USE_MACA=ON
|
||||
pip install -e . -v
|
||||
```
|
||||
|
||||
编译完成后，验证安装：
|
||||
|
||||
```bash
|
||||
python -c "import tilelang; print(tilelang.__version__)"
|
||||
# 输出：0.1.9+maca.gitf1ca0fb9
|
||||
```
|
||||
|
||||
> 编译日志显示成功构建了 692 个目标，包括 `libtilelang.so`、`libtvm.so` 及 Cython wrapper，最终生成 wheel 并安装到当前 Python 环境。
|
||||
|
||||
---
|
||||
|
||||
## 3. 算子测试总览
|
||||
|
||||
race 分支在 `race_tests/` 目录下提供了三个竞赛算子测试：
|
||||
|
||||
| 算子 | 目录 | 说明 | 测试脚本 |
|
||||
|------|------|------|----------|
|
||||
| **MLA** | `race_tests/mla/` | Multi-Head Latent Attention（多头潜在注意力）FlashAttention 风格分块实现 | `test_tilelang_mla.py` |
|
||||
| **MoE** | `race_tests/moe/` | Mixture of Experts（混合专家模型）DeepSeek 风格 Routed Grouped GEMM | `fusedmoe_benchmark.py` |
|
||||
| **NSA** | `race_tests/nsa/` | Native Sparse Attention（原生稀疏注意力）前向推理 | `test_tilelang_nsa_fwd.py` |
|
||||
|
||||
---
|
||||
|
||||
## 4. MLA（Multi-Head Latent Attention）测试
|
||||
|
||||
### 4.1 算子简介
|
||||
|
||||
MLA 测试基于 TileLang 实现了类 FlashAttention 的 split / no_split 双路径内核，支持：
|
||||
- **Q/KV 分离**：`Q` + `Q_pe` 与 `KV` + `K_pe`
|
||||
- **分块计算**：通过 `block_N`、`block_H` 控制 shared memory 分块
|
||||
- **精度校验**：与 PyTorch 参考实现 (`ref_program`) 对比，`rtol=2e-4, atol=1e-4`
|
||||
|
||||
### 4.2 运行方式
|
||||
|
||||
#### 单 case 快速测试
|
||||
|
||||
```bash
|
||||
cd race_tests/mla
|
||||
python test_tilelang_mla.py \
|
||||
--no-json \
|
||||
--batch 1 --heads 16 --kv_heads 1 \
|
||||
--kv_ctx 2048 --dim 512 --pe_dim 64
|
||||
```
|
||||
|
||||
#### 批量回归测试（JSON 用例）
|
||||
|
||||
```bash
|
||||
cd /data/tilelang-metax
|
||||
python race_tests/mla/test_tilelang_mla.py
|
||||
```
|
||||
|
||||
用例文件：`race_tests/mla/test_cases_mla_batch_ctx.json`，共 **31 组** 不同 batch 与 context length 的组合。
|
||||
|
||||
### 4.3 测试结果
|
||||
|
||||
**全部 31 组用例通过（31/31 PASS）**，结果摘要如下：
|
||||
|
||||
```text
|
||||
=== Summary: 31/31 passed, 0/31 failed ===
|
||||
CSV saved to: race_tests/mla/test_cases_mla_batch_ctx_results.csv
|
||||
```
|
||||
|
||||
#### 关键数据节选
|
||||
|
||||
| case_id | batch | kv_ctx | latency_ms | tflops | status |
|
||||
|---------|-------|--------|------------|--------|--------|
|
||||
| 1 | 1 | 2048 | 1.1888 | 0.0600 | PASS |
|
||||
| 3 | 1 | 8192 | 4.5795 | 0.0623 | PASS |
|
||||
| 6 | 1 | 65536 | 36.3962 | 0.0627 | PASS |
|
||||
| 12 | 4 | 2048 | 1.3289 | 0.2146 | PASS |
|
||||
| 16 | 4 | 65536 | 41.3467 | 0.2207 | PASS |
|
||||
| 27 | 32 | 2048 | 1.3790 | 1.6546 | PASS |
|
||||
| 31 | 32 | 65536 | 43.4431 | 1.6807 | PASS |
|
||||
|
||||
**观察**：
|
||||
- 随着 `batch` 增大，吞吐量（TFlops）线性提升，最高达到 **~1.74 TFlops**（batch=32）。
|
||||
- Latency 随 `kv_ctx` 增长而增加，但保持在合理范围内。
|
||||
|
||||
---
|
||||
|
||||
## 5. MoE（Mixture of Experts）测试
|
||||
|
||||
### 5.1 算子简介
|
||||
|
||||
MoE 测试实现了 DeepSeek 风格的 **Routed Grouped GEMM**：
|
||||
- **Gating Network**：Top-K 路由选择专家
|
||||
- **Grouped GEMM**：通过 `RoutedMoEKernel` 完成各专家的 gate/up/down 投影
|
||||
- **Scatter-Reduce**：将各专家输出按索引汇总回原始 token 位置
|
||||
|
||||
测试包含两类：
|
||||
- **Functional**：与 `ref_fusedmoe.py` 的参考实现逐元素对比，`atol=1e-2, rtol=1e-2`
|
||||
- **Performance**：CUDA Event 计时，warmup 10 轮，迭代 100 轮取平均
|
||||
|
||||
### 5.2 运行方式
|
||||
|
||||
```bash
|
||||
cd race_tests/moe
|
||||
python fusedmoe_benchmark.py
|
||||
# 或
|
||||
bash run.sh
|
||||
```
|
||||
|
||||
用例配置：`moe_test_configs.json`，包含 2 组功能测试 + 2 组性能测试。
|
||||
|
||||
### 5.3 测试结果
|
||||
|
||||
```text
|
||||
=== Running functional tests ===
|
||||
✅ Functional test passed for config: {'dhidden': 7168, 'dexpert': 2048, 'nroutedexperts': 8, 'nexpertspertoken': 4, 'bs': 1, 'seqlen': 8192, 'seed': 81394}
|
||||
✅ Functional test passed for config: {'dhidden': 3584, 'dexpert': 1024, 'nroutedexperts': 4, 'nexpertspertoken': 2, 'bs': 2, 'seqlen': 4096, 'seed': 81394}
|
||||
|
||||
=== Running performance tests ===
|
||||
⏱ Performance test: 460.92968750ms for config: {'dhidden': 7168, 'dexpert': 2048, 'nroutedexperts': 8, 'nexpertspertoken': 4, 'bs': 4, 'seqlen': 8192, 'seed': 81394}
|
||||
⏱ Performance test: 66.26872559ms for config: {'dhidden': 3584, 'dexpert': 1024, 'nroutedexperts': 4, 'nexpertspertoken': 2, 'bs': 8, 'seqlen': 4096, 'seed': 81394}
|
||||
```
|
||||
|
||||
**结论**：
|
||||
- **2/2 功能测试全部通过**，TileLang 实现与参考实现数值一致。
|
||||
- **2/2 性能测试完成**，在大 batch（bs=4, seqlen=8192, 8 专家）场景下单次前向约 **460.9 ms**；较小规模（bs=8, seqlen=4096, 4 专家）约 **66.3 ms**。
|
||||
|
||||
---
|
||||
|
||||
## 6. NSA（Native Sparse Attention）测试
|
||||
|
||||
### 6.1 算子简介
|
||||
|
||||
NSA 测试实现了**原生稀疏注意力前向推理**：
|
||||
- **Block-wise Sparse**：通过 `block_indices` 仅加载选定的 KV block，降低显存与计算量
|
||||
- **Causal Mask 支持**：`is_causal=True` 时自动屏蔽未来位置
|
||||
- **Shared Memory 预估**：测试脚本内置了 shared memory 占用估算，超出 65536 Bytes 的 case 自动 skip
|
||||
|
||||
### 6.2 运行方式
|
||||
|
||||
```bash
|
||||
cd race_tests/nsa
|
||||
python test_tilelang_nsa_fwd.py
|
||||
```
|
||||
|
||||
用例文件：`test_cases_nsa_fwd.json`，共 **109 组** 参数组合，覆盖不同 batch、seq_len、head dim、block_size、selected_blocks 等。
|
||||
|
||||
### 6.3 测试结果
|
||||
|
||||
测试持续运行中，前 34 组 case 全部通过，典型输出如下：
|
||||
|
||||
```text
|
||||
[1/109] B=1 SEQ_LEN=64 H=1 HQ=16 D=32 S=1 block_size=16 is_causal=True
|
||||
GPU latency: 0.0309 ms
|
||||
[2/109] B=1 SEQ_LEN=64 H=1 HQ=16 D=64 S=1 block_size=16 is_causal=True
|
||||
GPU latency: 0.0236 ms
|
||||
...
|
||||
[29/109] B=2 SEQ_LEN=1024 H=1 HQ=16 D=128 S=1 block_size=16 is_causal=True
|
||||
GPU latency: 0.0891 ms
|
||||
[30/109] B=4 SEQ_LEN=64 H=1 HQ=16 D=32 S=1 block_size=16 is_causal=True
|
||||
GPU latency: 0.0227 ms
|
||||
```
|
||||
|
||||
结果自动追加写入：`benchmark_results_nsa_fwd.csv`
|
||||
|
||||
#### CSV 结果节选
|
||||
|
||||
| idx | B | SEQ_LEN | H | HQ | D | S | block_size | latency_ms | status |
|
||||
|-----|---|---------|---|----|---|---|------------|------------|--------|
|
||||
| 1 | 1 | 64 | 1 | 16 | 32 | 1 | 16 | 0.030863 | PASS |
|
||||
| 5 | 1 | 128 | 1 | 16 | 64 | 1 | 16 | 0.029292 | PASS |
|
||||
| 14 | 1 | 1024 | 1 | 16 | 64 | 1 | 16 | 0.025723 | PASS |
|
||||
| 29 | 2 | 1024 | 1 | 16 | 128 | 1 | 16 | 0.089119 | PASS |
|
||||
|
||||
**观察**：
|
||||
- 小规模（B=1, SEQ_LEN=64~128）latency 约 **0.02~0.03 ms**
|
||||
- 中等规模（B=2, SEQ_LEN=1024, D=128）latency 约 **0.09 ms**
|
||||
- 所有已运行 case 状态均为 **PASS**，无 shared memory 超限导致的 skip
|
||||
|
||||
---
|
||||
|
||||
## 7. 总结
|
||||
|
||||
| 算子 | 测试类型 | 用例数 | 通过数 | 关键结论 |
|
||||
|------|----------|--------|--------|----------|
|
||||
| **MLA** | 功能+性能 | 31 | 31/31 | FlashAttention 风格实现正确，batch=32 时达 ~1.74 TFlops |
|
||||
| **MoE** | 功能+性能 | 4 | 4/4 | Routed Grouped GEMM 数值正确，大模型配置单次前向约 460 ms |
|
||||
| **NSA** | 功能+性能 | 109 | 持续通过中 | 稀疏注意力前向推理稳定，小尺度 latency < 0.1 ms |
|
||||
|
||||
**环境确认**：
|
||||
- ✅ MetaX C500 + MACA 3.5.3.20 驱动就绪
|
||||
- ✅ TileLang-MetaX race 分支编译安装成功
|
||||
- ✅ `race_tests` 三大算子均可在当前环境正常编译、运行并通过精度校验
|
||||
|
||||
---
|
||||
|
||||
## 8. 附录：常用命令速查
|
||||
|
||||
```bash
|
||||
# 查看 GPU 状态
|
||||
mx-smi
|
||||
|
||||
# 验证 TileLang 版本
|
||||
python -c "import tilelang; print(tilelang.__version__)"
|
||||
|
||||
# MLA 测试
|
||||
python race_tests/mla/test_tilelang_mla.py
|
||||
|
||||
# MoE 测试
|
||||
cd race_tests/moe && python fusedmoe_benchmark.py
|
||||
|
||||
# NSA 测试
|
||||
cd race_tests/nsa && python test_tilelang_nsa_fwd.py
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
*文档生成时间：2026-05-25*
|
||||
*测试平台：模力方舟 MetaX C500*
|
||||
Loading…
Reference in New Issue