🚀 Triton Operator Generator

GPU/加速器算子智能生成助手

基于大模型的 Triton GPU Kernel 智能生成、优化、审查、解释与迁移工具

支持 NVIDIA GPU + 沐曦 Metax + 昇腾 Ascend + 寒武纪 + 海光 DCU + 燧原 GCU

📄 模型服务集成与真实调用验证文档 — 包含沐曦 Metax GPU 算力环境配置、Qwen3.5-27B 模型 vLLM 部署、API 接口封装、基于Claude Code的真实调用日志（vllm.log）和性能数据（Prompt 290.5 tokens/s、生成 10.3 tokens/s）。

📖 项目简介

Triton Operator Generator 是一个面向 GPU/加速器编程的 AI Skill，通过接入大语言模型（Gitee.AI 平台沐曦 GPU 算力），为开发者提供 Triton GPU Kernel 的全流程智能辅助，并原生支持多种国产 AI 加速器的编程适配和知识检索。

🌟 核心亮点

亮点	说明
🎯 6 大功能模块	生成、优化、审查、解释、迁移、硬件检索
🔧 国产硬件全支持	沐曦、昇腾、寒武纪、海光、燧原一站式适配
📚 10+ Few-shot 示例	覆盖 LLM 推理/训练全场景算子
🤖 硬件自动检测	输入”昇腾”、”沐曦”等关键词自动切换适配模式
📊 完整性能追踪	延迟、吞吐、Token 统计一个不落

🖥 支持的硬件平台

硬件	设备类型	Triton 兼容	编程接口	关键特性
NVIDIA GPU	`cuda`	✅ 原生	CUDA/Triton	A100/H100, 312+ TFLOPS
沐曦 Metax	`cuda`	✅ 兼容层	CUDA/Triton	C500, 200+ TFLOPS
昇腾 Ascend	`npu`	⚠️ 需适配	Ascend C/MindSpore	910B, 310 TFLOPS
寒武纪	`mlu`	⚠️ 需适配	BANG C	MLU370, FP16/INT8
海光 DCU	`cuda`	✅ ROCm	HIP/Triton	K100, FP64/FP16
燧原 GCU	`gcu`	⚠️ 需适配	TopsRider	EnI, FP32/FP16

硬件适配策略

                    用户指定硬件
                        │
            ┌───────────┼───────────┐
            │           │           │
        CUDA 兼容     需适配      查询模式
        (沐曦/海光)  (昇腾/寒武纪)  (硬件信息)
            │           │           │
            ▼           ▼           ▼
     Triton 直接运行  提供：      返回：
     + 硬件差异注释  - Triton参考  - 编程接口
     + 性能调优建议  - 原生等价实现 - 内存层次
                    - API映射表   - 兼容性分析
                    - 差异对照表   - 代码示例

🛠 Skill 功能说明

1. 🔨 算子生成 (Generate)

输入：自然语言描述（可选指定硬件）输出：完整可运行的 Kernel 代码 + launcher + 验证代码 + 硬件适配说明

# 基础生成
python -m src.skill_engine generate "生成一个向量加法 Triton 算子"

# 指定硬件生成
python -m src.skill_engine generate "生成一个 Softmax 算子" --hw ascend
python -m src.skill_engine generate "沐曦 GPU 的矩阵乘法" --hw metax

支持的算子类型：

类别	算子	说明
基础	vector_add, softmax, relu, layernorm, matmul	通用算术/归一化
LLM 推理	flash_attention, fused_add_rmsnorm, rope, swiglu	LLaMA/GPT 推理核心
训练	fused_adam, fused_sgd	优化器融合
规约	reduce_sum, reduce_max	行级/全局规约
硬件适配	softmax_ascend, vector_add_metax	跨硬件适配模板

2. ⚡ 算子优化 (Optimize)

输入：已有 Triton 代码 + 优化目标 + 目标硬件输出：优化后代码 + 逐项优化说明 + 硬件特有优化建议

python -m src.skill_engine optimize kernel.py --hw ascend

3. 🔍 算子审查 (Review)

输入：Triton 代码输出：正确性 ✅⚠️ + 性能 🟢🟡🔴 + 硬件兼容性 🔧✅ + 修复代码

python -m src.skill_engine review kernel.py --hw metax

4. 📚 算子解释 (Explain)

输入：Triton 代码输出：逐行中文解释 + 数据流图 + PyTorch 等价实现 + 硬件映射表

python -m src.skill_engine explain kernel.py

5. 🔄 PyTorch → Triton 迁移 (Migrate)

输入：PyTorch 代码 + 目标硬件输出：Triton Kernel + 迁移说明 + 硬件适配方案

python -m src.skill_engine migrate pytorch_op.py --hw ascend

6. 🔎 硬件知识检索 (Search Hardware)

输入：硬件名称 + 查询内容输出：编程接口、性能特征、Triton 兼容性、API 映射对照表

# 查询昇腾的编程接口
python -m src.skill_engine search_hardware "昇腾"

# 查询沐曦的内存层次
python -m src.skill_engine search_hardware "沐曦"  # 会自动查询

🧠 使用的模型与算力环境

模型服务

模型	说明	适用场景
DeepSeek-V3	主要模型，代码生成能力极强	算子生成、迁移
DeepSeek-R1	推理型模型	复杂算子设计、审查
Qwen2.5-72B	中文能力强	硬件知识检索、解释
GLM-4	国产模型	中文场景备选

算力环境

平台：Gitee.AI
GPU：沐曦 GPU 算力（C500 系列）
接入方式：OpenAI 兼容 API 接口

API 接口封装

from src.api_client import create_client

# 创建客户端
client = create_client(api_key="your-api-key")

# 调用模型
result = client.chat(
    messages=[
        {"role": "system", "content": "你是 Triton 算子生成专家..."},
        {"role": "user", "content": "生成一个向量加法 Triton 算子"},
    ],
    temperature=0.3,
    max_tokens=4096,
)

print(result["content"])
print(f"延迟: {result['latency_seconds']:.2f}s")
print(f"吞吐: {result['tokens_per_second']:.1f} tokens/s")

🏗 技术架构

┌──────────────────────────────────────────────────────────┐
│                      用户输入层                           │
│  "生成一个 Softmax，目标硬件昇腾"                          │
│  / Triton 代码 / PyTorch 代码 / 硬件查询                  │
└────────────────────────┬─────────────────────────────────┘
                         │
                         ▼
┌──────────────────────────────────────────────────────────┐
│                    意图识别 + 硬件检测                      │
│  ┌──────────┐  ┌──────────────────────────────────────┐  │
│  │ 功能识别  │  │ 硬件自动检测（昇腾/沐曦/寒武纪/海光）  │  │
│  └──────────┘  └──────────────────────────────────────┘  │
└────────────────────────┬─────────────────────────────────┘
                         │
                         ▼
┌──────────────────────────────────────────────────────────┐
│                  Prompt 工程层                            │
│  ┌────────────────────────────────────────────────────┐  │
│  │ System Prompt（GPU专家+国产硬件知识库）              │  │
│  ├────────────────────────────────────────────────────┤  │
│  │ Few-shot 示例（10+ 个，按任务+硬件智能选择）        │  │
│  ├────────────────────────────────────────────────────┤  │
│  │ 硬件上下文注入（目标硬件特有约束+API映射）           │  │
│  └────────────────────────────────────────────────────┘  │
└────────────────────────┬─────────────────────────────────┘
                         │
                         ▼
┌──────────────────────────────────────────────────────────┐
│                  API 客户端层                             │
│  Gitee.AI 封装 | 自动重试 | 性能追踪 | 日志记录          │
└────────────────────────┬─────────────────────────────────┘
                         │
                         ▼
┌──────────────────────────────────────────────────────────┐
│              沐曦 GPU 算力 + 大模型推理                    │
│  DeepSeek-V3 | Qwen2.5-72B | GLM-4 | Llama-3.1         │
└──────────────────────────────────────────────────────────┘

🚀 运行与部署方法

环境要求

Python 3.8+
CUDA GPU（运行 Triton Kernel 需要）或昇腾 NPU/沐曦 GPU 等
Gitee.AI API Key

安装

# 克隆仓库
git clone https://gitlink.org.cn/your-username/triton-operator-generator.git
cd triton-operator-generator

# 安装依赖
pip install -r requirements.txt

# 配置 API Key
cp .env.example .env
# 编辑 .env 填入你的 GITEE_AI_API_KEY

命令行使用

# 设置 API Key
export GITEE_AI_API_KEY="your-api-key"

# 基础功能
python -m src.skill_engine generate "生成一个向量加法 Triton 算子"
python -m src.skill_engine optimize kernel.py
python -m src.skill_engine review kernel.py
python -m src.skill_engine explain kernel.py
python -m src.skill_engine migrate pytorch_op.py

# 指定目标硬件
python -m src.skill_engine generate "生成 Softmax" --hw ascend
python -m src.skill_engine generate "矩阵乘法" --hw metax
python -m src.skill_engine optimize kernel.py --hw dcu

# 硬件知识检索
python -m src.skill_engine search_hardware "昇腾"
python -m src.skill_engine search_hardware "沐曦 Metax 的编程接口"

# 带性能统计
python -m src.skill_engine generate "ReLU 算子" --stats

Python API 使用

from src.skill_engine import TritonOperatorGenerator, detect_hardware

engine = TritonOperatorGenerator(api_key="your-api-key")

# 自动检测硬件
hw = detect_hardware("生成一个 Softmax，在昇腾上跑")
# → "昇腾 Ascend NPU"

# 生成算子（自动适配硬件）
result = engine.generate("生成一个 Softmax，在昇腾上跑")
print(result["code"])
print(result["hardware"])  # "昇腾 Ascend NPU"

# 指定硬件生成
result = engine.generate("矩阵乘法", hardware="沐曦 Metax GPU")

# 硬件知识检索
info = engine.search_hardware_info("昇腾", query="内存层次和编程接口")
print(info["info"])

# 性能统计
engine.client.print_stats()

📋 示例输入输出

示例1：昇腾 NPU 适配的 Softmax

输入：

生成一个 Softmax 算子，需要支持在昇腾 NPU 上运行

输出（关键部分）：

def softmax_ascend(x: torch.Tensor) -> torch.Tensor:
    """昇腾 NPU 适配的 Softmax"""
    device_type = str(x.device).split(':')[0]

    if device_type == 'npu':
        import torch_npu
        return torch.nn.functional.softmax(x, dim=-1)
    elif device_type == 'cuda':
        return softmax_triton(x)  # Triton kernel
    else:
        return torch.nn.functional.softmax(x, dim=-1)

硬件差异对照表：

Triton 操作	昇腾 Ascend C 等价	说明
tl.load	DataCopy(GM→UB)	全局内存到统一缓冲区
tl.store	DataCopy(UB→GM)	统一缓冲区到全局内存
tl.max(axis=0)	ReduceMax	行内归约
tl.exp	Exp	逐元素指数
tl.sum(axis=0)	ReduceSum	行内求和

示例2：Flash Attention 生成

输入：

生成一个 Flash Attention 前向传播的 Triton 算子，支持 causal mask

输出：完整的 Flash Attention kernel 实现，包含 online softmax、causal mask、分块计算（详见 Few-shot 示例）。

示例3：硬件知识检索

输入：

查询沐曦 Metax GPU 的编程接口

输出：

沐曦 Metax GPU 知识：

1. 架构：兼容 CUDA 生态，支持 Triton
2. 编程模型：与 NVIDIA GPU 高度兼容，使用 CUDA/Triton 编程
3. 内存层次：Global Memory → L2 Cache → Shared Memory → Registers
4. Triton 兼容性：✅ 直接支持，无需额外适配
5. 设备类型：'cuda'（通过 CUDA 兼容层运行）
6. 沐曦 C500 系列：FP32/FP16/BF16/INT8，约 200+ TFLOPS

示例代码：
  import torch
  x = torch.randn(1000, device='cuda')  # 沐曦 GPU 使用 cuda 设备
  ...

示例4：性能数据

📊 API 调用性能统计
==================================================
总调用次数:   10
成功/失败:    10/0
平均延迟:     2.356s
P50 延迟:     2.123s
P99 延迟:     4.567s
平均吞吐量:   85.3 tokens/s
总 Token 消耗: 25,678
==================================================

📁 项目结构

triton-operator-generator/
├── .trae/
│   └── skills/
│       └── triton-operator-generator/
│           └── SKILL.md                    # Skill 定义文件（核心交付物）
├── src/
│   ├── __init__.py                         # 包初始化
│   ├── api_client.py                       # Gitee.AI API 客户端封装
│   ├── prompts.py                          # Prompt 模板 + 国产硬件知识库 + 13 个 Few-shot
│   └── skill_engine.py                     # 核心引擎 + 硬件检测 + 6 大功能
├── examples/
│   ├── demo_generate.py                    # 功能演示脚本
│   └── verify_call.py                      # 真实调用验证脚本
├── requirements.txt                        # 依赖
├── .env.example                            # 环境变量模板
├── .gitignore                              # Git 忽略配置
├── LICENSE                                 # MIT 开源协议
└── README.md                               # 本文档

🔧 国产硬件知识库

本 Skill 内置了完整的国产 AI 加速器知识库，包含：

硬件	架构	编程接口	内存层次	Triton 兼容
沐曦 Metax C500	CUDA 兼容	CUDA/Triton	GM→L2→SM→Reg	✅ 直接运行
昇腾 Ascend 910B	达芬奇	Ascend C	GM→L1→L0→UB	⚠️ 需适配
寒武纪 MLU370	MLUarch	BANG C	GM→Cluster→NRAM→WRAM	⚠️ 需适配
海光 DCU K100	AMD ROCm	HIP/Triton	GM→L2→LDS→Reg	✅ ROCm 后端
燧原 EnI	自研	TopsRider	自定义	⚠️ 需适配

API 映射对照表（Triton → 国产硬件原生 API）

Triton 操作	昇腾 Ascend C	寒武纪 BANG C	海光 HIP
`tl.load`	`DataCopy(GM→UB)`	`__memcpy`	相同
`tl.store`	`DataCopy(UB→GM)`	`__memcpy`	相同
`tl.dot`	`CubeCompute`	`__bang_matmul`	相同
`tl.program_id`	`GetBlockIdx`	`taskId`	相同
`BLOCK_SIZE`	`tiling 因子`	`拆分因子`	相同
`tl.arange`	`计算偏移`	`计算偏移`	相同
`tl.where`	`Select`	`__bang_select`	相同

📊 真实调用验证

验证方法

export GITEE_AI_API_KEY="your-api-key"
cd examples
python verify_call.py

验证内容

算子生成：5 个不同类型的算子生成
算子优化：Softmax 优化
算子审查：含 bug 代码审查
算子迁移：PyTorch → Triton 迁移
硬件检索：国产硬件知识查询

性能数据记录

验证脚本自动记录：

每次 API 调用的延迟（秒）
Token 吞吐量（tokens/s）
输入/输出 Token 数
P50/P99 延迟
成功率

📄 License

MIT License

🙏 致谢

Triton - OpenAI 开发的 GPU 编程语言
Gitee.AI - 模型服务平台
沐曦 GPU - 算力支持
昇腾 Ascend - NPU 参考文档
寒武纪 Cambricon - MLU 参考文档
海光 Hygon - DCU 参考文档