昇腾triton算子TritonBench

# TritonBench-Ascend Ascend NPU 上的 Triton 算子代码生成基准数据集，用于评测和微调大模型生成 Triton kernel 的能力。 ## 数据集概览 | 项目 | 说明 | |------|------| | 样本数 | 25 个已验证算子 | | 任务 | Triton kernel + Python wrapper 代码生成 | | 平台 | Ascend NPU（需 `torch_npu`，`device='npu'`） | | 格式 | Alpaca SFT + 元数据 JSON + 参考源码 | ## 目录结构 ``` data/ ├── TritonBench_A_v1.json # 元数据索引（算子名、kernel、难度等） ├── TritonBench_A_comp_alpac_v1.json # Alpaca 格式，完整 instruction ├── TritonBench_A_simp_alpac_v1.json # Alpaca 格式，简化 instruction └── TritonBench_A_v1/ # 25 个参考 kernel 源码 (.py) ``` ## 数据格式 ### 元数据 (`TritonBench_A_v1.json`) 每条样本包含： - `file` / `op`: 算子文件名与名称 - `kernels` / `wrappers`: kernel 与 wrapper 函数名 - `apis`: 使用的 Triton 语言原语 - `has_dot` / `has_autotune`: 是否含 dot / autotune - `difficulty`: 难度等级 (1–3) - `npu_validated`: 是否在 Ascend NPU 上验证通过 ### Alpaca SFT 格式 ```json { "instruction": "You are an expert in writing Triton operators for the Ascend NPU...", "input": "", "output": "import torch\nimport torch_npu\n..." } ``` - **comp**: 完整 instruction，含 kernel 签名、Triton 原语、Ascend 注意事项 - **simp**: 简化 instruction，仅保留核心任务描述 ## 算子列表 涵盖 vector add、matmul、fused softmax、layer norm、fused attention、grouped GEMM、 persistent matmul 等基础算子，以及 vLLM 推理相关 kernel（RoPE、chunk attention 等）。 ## 加载示例 ```python import json from modelscope.msdatasets import MsDataset # 从 ModelScope 加载（需替换 namespace） # ds = MsDataset.load('your_namespace/TritonBench-Ascend', split='train') # 或直接读取本地 JSON with open("data/TritonBench_A_comp_alpac_v1.json") as f: data = json.load(f) print(len(data), data[0]["instruction"][:80]) ``` ## 许可证 Apache 2.0 ## 引用 如在研究中使用本数据集，请引用 TritonBench-Ascend 及 Triton-Ascend 项目。