# lmcustomop

## 简介

`lmcustomop` 是一个面向 DCU/ROCm 环境的轻量融合算子包，当前聚焦于 **RMSNorm + RoPE 融合前向**。

当前版本提供的核心能力：

- `rms_rotary_embedding_fuse`
- 将 `RMSNorm` 与 `Rotary Embedding` 融合在同一个自定义算子中执行
- 对 `query` / `key` 做原地更新，减少中间访存

---

## 安装

### 环境依赖

- Python 3.10+
- PyTorch（带 ROCm/CUDA Extension 编译能力）
- 对应 DCU 驱动与编译工具链

> 说明：本仓库通过 `torch.utils.cpp_extension.CUDAExtension` 构建。

### 源码安装

在仓库目录执行：

```bash
python setup_lmcustomop.py install
```

如果需要指定架构（示例）：

```bash
PYTORCH_ROCM_ARCH='gfx906;gfx926' python setup_lmcustomop.py install
```

### 构建 wheel

```bash
python setup_lmcustomop.py bdist_wheel
```

构建完成后，wheel 位于 `dist/` 目录。

---

## 算子介绍

### 核心算子

| 算子 | 说明 |
| --- | --- |
| `rms_rotary_embedding_fuse` | 对 `query/key` 执行 RMSNorm 与 RoPE 融合计算（in-place） |

---

## 安装验证

```bash
python -c "import lmcustomop; print(lmcustomop.rms_rotary_embedding_fuse)"
```

若能正常打印函数对象，说明安装成功。