已合并
新增CompressorGrad算子 #9763
莫允扬创建于 12 天前
新增CompressorGrad算子 #9763
已合并
莫允扬成员
12 天前
12 天前
## 描述
为 **Compressor**(KV cache 压缩算子,SMLA/QLI 前处理)新增反向算子 **CompressorGrad**,打通训练/微调场景下压缩分支的梯度计算链路。
**1. 新增反向算子 `attention/compressor_grad/`**
- **op_host**:算子定义(`compressor_grad_def.cpp`,9 输入 4 输出:d_x/d_wkv/d_wgate/d_ape)、infershape、tiling(30 字段 TilingData + workspace 分区计算 + 全量入参校验)
- **op_kernel**:PyPTO 实现的 `compressor_grad.py`,三阶段流水线(Vec scatter 反向 + Cube 双 matmul + Vec 跨核归约),支持 coff=1/2、BSH/TH 双布局、FP16/BF16,D∈{128,512}、cmp_ratio∈[2,128]
**2. 正向算子 `compressor` 增强**
- 新增 `softmax_score`/`kv` 两个 FP32 中间输出与 `grad_enabled` 属性(grad_enabled=true 时导出中间量,作为 CompressorGrad 输入)
**3. torch_extension 集成**
- `compressor.py` 注册 autograd:`backward` 通过 `aclnnCompressorGrad` 计算 d_x/d_wkv/d_wgate/d_ape,前向 `grad_enabled` 自动按 `x.requires_grad` 置位
## 关联的Issue
## 测试
## 文档更新
## 类型标签
- [ ] 🐛 Bug修复
- [x] ✨ 新特性
- [ ] ⚡ 性能优化
- [ ] ♻️ 重构
- [ ] 🧪 测试
- [ ] 📦 构建/CI
- [ ] 🔧 配置变更
- [ ] 📝 文档更新
- [ ] ⬆️ 依赖升级
- [ ] 🔒 安全修复
- [ ] 🧹 代码清理
- [ ] ❓ 其他,请描述:


Pull Request已成功合入, 合并人@CANN-robot
(感谢 莫允扬 的贡献)12 天前 update merge request[project id: 7673863, iid: 9763, commit_id: 3aa3efa84edffb581bc521a4a68ee66c6f5854e6] virtual merging success
12 天前 update merge request[project id: 7673863, iid: 9763, commit_id: 3aa3efa84edffb581bc521a4a68ee66c6f5854e6] virtual merging success
12 天前 推送了提交
12 天前 update merge request[project id: 7673863, iid: 9763, commit_id: 9742f1a9da6c9e35883872cc54dbf1d223a7295d] virtual merging success
12 天前 update merge request[project id: 7673863, iid: 9763, commit_id: 9742f1a9da6c9e35883872cc54dbf1d223a7295d] virtual merging success
12 天前 添加了label:stat/needs-squash
此处折叠了224条消息 查看更多
8 天前 添加了label:lgtm
8 天前 关闭了关联的issue
8 天前 合入了pull request