chenchenghao/DeepGEMM

Fork 0

Files

Xinyi Liu 1dafd0f173 feat: add sm90 megamoe phase1 interfaces

2026-06-17 23:54:49 +08:00

8.6 KiB

Raw Blame History

DeepGEMM 算子开发工作流 (AGENTS.md)

项目概述

DeepGEMM 是一个高性能 CUDA kernel 库（C++20/CUDA，JIT 编译），支持 SM90 (H100/H200) 与 SM100 (B200) 架构。

核心开发文件：

Kernel 实现: deep_gemm/include/deep_gemm/impls/*.cuh
JIT / host 端代码: csrc/**/*.hpp, csrc/**/*.cpp, csrc/**/*.cu
Python API / 测试: deep_gemm/**/*.py, tests/*.py

SSH 远程开发工作流（核心）

原则：单向同步 Local → Remote，禁止反向

本地编辑 → scp 传输 → ssh 远程测试 → 循环
   ↑                                        |
   └──────────── 禁止 scp remote→local ─────┘

标准流程

本地编辑 — 在本地修改 kernel/host 代码

scp 传输 — 将变更文件推送到远程服务器

# 示例：传输单个文件
scp deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh <REMOTE_HOST>:<REMOTE_PATH>/deep_gemm/include/deep_gemm/impls/

# 示例：传输整个目录
scp -r deep_gemm/include/deep_gemm/impls/ <REMOTE_HOST>:<REMOTE_PATH>/deep_gemm/include/deep_gemm/impls/

ssh 远程测试 — 在远程服务器上编译并运行测试，长耗时或可能 hang 的命令必须设置合理 timeout
```
ssh <REMOTE_HOST> "cd <REMOTE_PATH> && timeout 60s python tests/test_fp8_fp4.py"
```
迭代 — 根据远程测试结果，在本地继续修改，重复步骤 2-3

连接配置

远程主机: <REMOTE_HOST>（从 ~/.ssh/config 或环境变量读取）
远程路径: <REMOTE_PATH>（项目在远程服务器的根目录）
传输前确认远程路径存在: ssh <REMOTE_HOST> "ls <REMOTE_PATH>"

Timeout 规则

运行远程 build / pytest / torchrun / benchmark 命令时，必须根据预期耗时设置 timeout，避免 kernel hang、NCCL hang 或 barrier hang 长时间占用会话。
远端 Linux/container 命令优先使用 shell timeout 包裹实际测试命令，例如 timeout 60s python ... 或 timeout 60s torchrun ...。
如果调用工具本身也支持 timeout 参数，也必须设置外层 timeout；外层 timeout 应略大于远端 timeout，便于收集远端退出信息。
经验值：常规 smoke / build / torchrun 测试优先设置 60 秒左右 timeout；性能 benchmark 按 case 数量单独估算。

Debug 代码规则（关键）

必须加 `DEBUG` 注释

每行 debug 代码末尾必须加 // DEBUG 注释（C++/CUDA）或 # DEBUG（Python）。

// 正确示例
printf("m=%d n=%d k=%d\n", m, n, k);  // DEBUG
float* debug_buf = new float[M * N];   // DEBUG

// 错误示例（禁止）
printf("m=%d n=%d k=%d\n", m, n, k);  // 缺少 DEBUG 标记

# 正确示例
print(f"shape: {x.shape}, dtype: {x.dtype}")  # DEBUG
torch.save(result, "/tmp/debug_result.pt")     # DEBUG

Debug 代码是临时的，禁止提交

Debug 代码仅用于定位问题，严禁进入 git commit。

Debug → 正式修复的标准流程

1. 用 debug 代码定位到问题根因
       ↓
2. git stash -m "DEBUG: <描述本次调试目的>"
   （debug 代码被暂存，工作区恢复干净）
       ↓
3. git stash/git restore
       ↓
4. 编写正式修复（干净代码，不含 DEBUG 注释）
       ↓
5. git add <修复的文件>
   git commit -m "<修复描述>"

stash 消息必须有意义，清晰描述调试目标：

# 好
git stash -m "DEBUG: 排查 sm90_fp8_gemm 的 block scheduling index 越界问题"

# 差
git stash -m "test"
git stash -m "debug"

测试脚本规则

存放位置：`./megamoe_dev_test_scripts/phase<phase id>/`

MegaMoE / SM90 开发测试脚本必须按 phase 存放到仓库内，并随对应工作提交到 git：

mkdir -p megamoe_dev_test_scripts/phase2

示例路径：

megamoe_dev_test_scripts/phase0/megamoe_phase0_smoke.py
megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py

生命周期：随代码提交，禁止 stash-only

# 1. 编写或更新测试脚本
#    megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py

# 2. 与对应实现一起暂存
git add deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh
git add megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py

# 3. 随 clean 工作提交进入 git history
git commit -m "feat: 实现 sm90 megamoe dispatch-only 路径"

测试脚本是开发过程的一部分，必须可追溯、可复跑。不要再把 MegaMoE 开发测试脚本放到 .tmp/test_scripts/，也不要通过 git stash 作为唯一保存方式。

禁止操作清单

操作	说明
`sed` 编辑代码	严格禁止。使用 Edit 工具修改文件
`scp remote→local`	禁止从远程拉回变更，所有修改在本地进行
Debug 代码 commit	禁止将含 `// DEBUG` / `# DEBUG` 的代码提交到 git
MegaMoE 测试脚本 stash-only	测试脚本必须放入 `megamoe_dev_test_scripts/phase<phase id>/` 并随 commit 提交
MegaMoE 测试脚本放 `.tmp/test_scripts/`	测试脚本不再放 `.tmp`，`.tmp` 只用于非提交的临时文件
无 timeout 的远程长命令	build / pytest / torchrun / benchmark 必须设置合理 timeout，避免 hang
git stash pop	避免git conflict
ssh xxxx bash -c "大段测试脚本"	避免直接运行测试，必须持久化成临时脚本

Git 管理规范

每次 clean 工作提交后的开发日志

每次完成用户请求并创建一次 clean 的工作 commit 后，必须把本次工作内容和详细流程追加到 MEGAMOE_SM90_DEV.md，并把该开发日志文档提交到 git。

由于 git commit hash 只有在 commit 创建后才确定，标准流程是两步提交：

# 1. 提交干净的代码/测试/文档变更
git add <本次工作文件>
git commit -m "<本次工作描述>"

# 2. 获取刚完成的工作 commit hash
git rev-parse HEAD

# 3. 追加开发日志到 MEGAMOE_SM90_DEV.md
# 日志必须记录上一步输出的 HEAD hash

# 4. 单独提交开发日志
git add MEGAMOE_SM90_DEV.md
git commit -m "docs: 记录 <本次工作描述> 开发日志"

MEGAMOE_SM90_DEV.md 每条日志至少包含：

日期和时间
对应 clean 工作 commit 的 git HEAD hash
用户请求摘要
本次提交的核心改动
关键文件列表
详细开发流程（本地修改、远程同步、编译/测试命令）
测试结果和已知问题
后续待办

日志 commit 记录的是刚完成的 clean 工作 commit hash，不要求记录日志 commit 自身 hash。

Debug 周期完整示例

# === 阶段 1: 添加 debug 代码 ===
# 编辑 deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh
# 加入 printf + DEBUG 标记

# === 阶段 2: scp 传输到远程 ===
scp deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh gpu01:/workspace/DeepGEMM/deep_gemm/include/deep_gemm/impls/

# === 阶段 3: ssh 远程测试 ===
ssh gpu01 "cd /workspace/DeepGEMM && python tests/test_fp8_fp4.py"

# === 阶段 4: 定位到问题后，stash debug 代码 ===
git stash -m "DEBUG: 排查 fp8_gemm_1d1d 的 warp scheduling 偏移错误"

# === 阶段 5: 编写正式修复 ===
# 编辑文件，写入干净的修复代码（无 DEBUG 注释）

# === 阶段 6: 提交修复 ===
git add deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh
git commit -m "fix: 修正 sm90_fp8_gemm 1d1d warp scheduling 偏移计算"

测试脚本管理示例

# 编写 Phase 2 开发测试
mkdir -p megamoe_dev_test_scripts/phase2
# megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py

# 与对应实现一起提交
git add megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py
git add deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh
git commit -m "feat: 实现 sm90 megamoe phase2 dispatch-only"

远程常用命令参考

# 传输单个 kernel 文件
scp deep_gemm/include/deep_gemm/impls/<kernel>.cuh <HOST>:<PATH>/deep_gemm/include/deep_gemm/impls/

# 传输整个 impls 目录
scp -r deep_gemm/include/deep_gemm/impls/ <HOST>:<PATH>/deep_gemm/include/deep_gemm/impls/

# 传输 host 端代码
scp csrc/**/*.hpp <HOST>:<PATH>/csrc/

# 远程运行测试
ssh <HOST> "cd <PATH> && timeout 60s python tests/test_fp8_fp4.py"

# 远程运行单个测试函数
ssh <HOST> "cd <PATH> && timeout 60s python -c 'from tests.test_fp8_fp4 import test_fp8_gemm_nt; test_fp8_gemm_nt()'"

# 远程 build（如需重新编译 _C.so）
ssh <HOST> "cd <PATH> && timeout 60s bash develop.sh"

# 远程多 rank 测试
ssh <HOST> "cd <PATH> && timeout 60s torchrun --standalone --nproc_per_node=2 megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py"

8.6 KiB Raw Blame History Unescape Escape