已合并
新增CompressorGrad算子 #9763
新增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 - [ ] 🔧 配置变更 - [ ] 📝 文档更新 - [ ] ⬆️ 依赖升级 - [ ] 🔒 安全修复 - [ ] 🧹 代码清理 - [ ] ❓ 其他,请描述:
likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 莫允扬 的贡献)
莫允扬莫允扬成员
12 天前 创建了 pull request,commit e560ddd0
莫允扬莫允扬成员
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
CANN-robotCANN-robot成员
12 天前 添加了label:stat/needs-squash
此处折叠了224条消息 查看更多
shasha_an成员
8 天前 评论:

/lgtm

likedislike
mabing成员
8 天前 评论:

/lgtm
/approve

likedislike
CANN-robotCANN-robot成员
8 天前 添加了label:lgtm
CANN-robotCANN-robot成员
8 天前 关闭了关联的issue
CANN-robotCANN-robot成员
8 天前 合入了pull request