# mp-gemm

**Repository Path**: riverclouds/mp-gemm

## Basic Information

- **Project Name**: mp-gemm
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: main
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2025-07-25
- **Last Updated**: 2025-07-25

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# 高性能GEMM代码生成算法

基于LLM的自动化Triton GEMM算子优化系统。

## 项目概述

这个项目实现了一个使用大语言模型(LLM)自动生成和优化高性能GEMM(通用矩阵乘法)算子的算法。该算法采用搜索树结构，通过迭代应用优化原语，最终生成高性能的Triton代码。

## 核心算法

### 输入参数
- `pA`: 一组元提示(meta-prompts)，包含平台无关的优化目标和Triton特定的实现提示
- `k`: 搜索宽度(beam width)，控制每轮迭代保留的候选数量
- `T`: 最大迭代次数

### 算法流程

1. **初始化**
   - 创建朴素GEMM实现作为搜索树根节点
   - 测试基线性能

2. **迭代优化** (共T轮)
   - 对每个候选代码，让LLM建议优化原语
   - 应用每个优化原语生成新代码
   - 性能测试所有新代码
   - 剪枝：保留性能最好的前k个候选

3. **返回结果**
   - 输出性能最优的代码

## 文件结构

```
.
├── meta-prompt.json          # 元提示配置
├── config.json               # 算法配置参数
├── gemm_generator.py         # 主算法实现
├── test_gemm_generator.py    # 测试脚本(使用模拟组件)
├── demo.py                   # 算法演示
└── README.md                # 项目说明
```

## 关键组件

### 1. LLMClient
- 负责与DeepSeek API通信
- 模型：`deepseek-reasoner`
- 处理优化原语建议和代码生成

### 2. TritonTester  
- 自动测试生成代码的性能
- 创建临时测试脚本
- 返回平均执行时间(ms)

### 3. GEMMOptimizer
- 核心优化算法实现
- 管理搜索树和候选剪枝
- 集成LLM和性能测试

### 4. 元提示系统
- **平台无关描述**：数据访问优化、计算/内存重叠、算子融合、粒度控制
- **Triton特定提示**：块指针、autotuning、高效融合实现

## 使用方法

### 基本运行
```bash
python gemm_generator.py
```

### 演示算法流程
```bash
python demo.py
```

### 测试(使用模拟组件)
```bash  
python test_gemm_generator.py
```

## 配置参数

项目使用 `config.json` 文件进行配置：

```json
{
    "llm_config": {
        "base_url": "https://api.deepseek.com",
        "model": "deepseek-reasoner", 
        "api_key": "",
        "temperature": 0.3,
        "max_tokens": 2000
    },
    "optimization_config": {
        "beam_width": 3,
        "max_iterations": 5,
        "matrix_size": 1024,
        "max_fix_attempts": 3
    },
    "output_config": {
        "verbose": true,
        "save_intermediate_results": true
    }
}
```

### 主要参数说明
- `beam_width`: 搜索宽度，控制每轮保留的候选数量
- `max_iterations`: 最大迭代次数
- `matrix_size`: 性能测试时使用的矩阵大小
- `max_fix_attempts`: 代码错误修复的最大尝试次数
- `model`: LLM模型名称（deepseek-reasoner）

## 优化策略示例

算法会自动发现并应用以下优化模式：

1. **Autotuning装饰器**
   ```python
   @triton.autotune(configs=[...], key=['M', 'N', 'K'])
   ```

2. **内存访问优化**
   ```python
   # 连续块访问，优化缓存命中率
   a_mask = (offs_am[:, None] < M) & (offs_k[None, :] < K)
   ```

3. **算子融合**
   ```python
   # 在同一kernel中完成多个操作
   accumulator = tl.dot(a, b) + bias
   result = tl.exp(accumulator)  # 融合激活函数
   ```

4. **计算粒度控制**
   ```python
   # 优化块大小平衡并行度和资源利用率
   BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 128, 256, 64
   ```

## 性能预期

基于典型优化场景：
- 基线性能: ~15ms (朴素实现)
- 第1轮优化: ~8ms (autotuning)
- 第2轮优化: ~7ms (内存优化)
- 最终优化: ~6ms (组合策略)
- **总体提升: 55%+**

## 算法特点

- **自动化**: 完全基于LLM的代码生成和优化
- **平台感知**: 结合Triton特定的优化提示
- **性能导向**: 以执行时间为唯一优化目标
- **错误修复**: 自动检测和修复生成代码中的错误
- **可扩展**: 容易添加新的优化原语和平台支持

## 依赖要求

- Python 3.8+
- torch (用于性能测试)
- triton (目标平台)
- requests (LLM API调用)

## 注意事项

1. 需要有效的DeepSeek API访问
2. 需要CUDA环境用于Triton代码测试
3. 性能测试依赖于具体硬件配置
4. LLM生成的代码可能需要语法检查

## 扩展方向

- 支持更多矩阵运算类型(Conv, Attention等)
- 集成更多优化原语
- 支持多GPU并行优化
- 添加代码正确性验证
- 优化LLM prompt设计

---

*该项目展示了LLM在高性能计算代码生成领域的应用潜力，为自动化性能优化提供了新的思路。*