| engram support training Co-authored-by: luozhonglin222<luozhonglin1@huawei.com> # message auto-generated for no-merge-commit merge: !8914 merge master into master engram support training Created-by: luozhonglin222 Commit-by: luozhonglin222 Merged-by: cann-robot Description: ## 描述 engram算子支持训练HOST代码，前向复用现有接口，增加相应参数，反向增加新接口 ## 关联的Issue https://gitcode.com/cann/ops-transformer/issues/3838 ## 测试 本地验证 ## 文档更新 不涉及 ## 类型标签 <!-- [x] 表示选中 --> - [ ] 🐛 Bug修复 - [x] ✨ 新特性 - [ ] ⚡ 性能优化 - [ ] ♻️ 重构 - [ ] 🧪 测试 - [ ] 📦 构建/CI - [ ] 🔧 配置变更 - [ ] 📝 文档更新 - [ ] ⬆️ 依赖升级 - [ ] 🔒 安全修复 - [ ] 🧹 代码清理 - [ ] ❓ 其他，请描述： # PR #8914 代码检视报告 ## 检视概览 - **PR URL**：https://gitcode.com/cann/ops-transformer/pull/8914 - **算子名**：engram_fetch_grad / engram_fetch_train / elastic_buffer - **代码侧别**：Kernel + Tiling + 混合（torch_extension） - **变更文件数**：21 - **文件组数**：6（K1, K2, H1, H2, S, E） - **总条例数**：346 - **检视时间**：2026-07-21 - **检视模式**：大型 PR 检视（21 文件 > 10 阈值） ## 检视统计 | 状态 | 条例数 | 占比 | |------|--------|------| | PASS | 341 | 98.6% | | FAIL（发现问题） | 5 | 1.4% | | SUSPICIOUS（需关注） | 0 | 0% | ### 逐文件组统计 | 文件组 | 侧别 | 文件数 | PASS | FAIL | |--------|------|--------|------|------| | engram_fetch_grad_kernel_G1 (K1) | Kernel | 3 | 27 | 0 | | engram_fetch_train_kernel_G1 (K2) | Kernel | 3 | 37 | 0 | | engram_fetch_grad_host_G1 (H1) | Tiling | 3 | 66 | 2 | | engram_fetch_train_host_G1 (H2) | Tiling | 3 | 70 | 2 | | shared (S) | 混合 | 6 | 61 | 1 | | elastic_buffer (E) | 混合 | 3 | 80 | 0 | | **合计** | — | **21** | **341** | **5** | ### 系统性风险 **SEC-4.1 / TOPK-7 在 H1 和 H2 中均失败**：两个 Tiling 侧文件组（engram_fetch_grad 和 engram_fetch_train）均存在 sendCounts/recvLocalEntry（或 sendCountsOut/recvLocalEntryOut）的 dim(0) 值未做非负/范围校验的问题。两个算子的 tiling 代码结构高度相似，校验缺失模式相同，建议全局排查并统一修复。 --- ## 发现问题（HIGH 置信度） ### 1. [engram_fetch_grad_host_G1] SEC-4.1 外部输入数据需要做合法性校验 - **问题描述**：sendCounts 的 dim(0) 和 recvLocalEntry 的 dim(0) 作为外部输入（来自 graph shape）仅校验了维度数（1D），未对 dim(0) 的值做合法性校验，违反 SEC-4.1「外部输入数据需要做合法性校验且确保校验范围正确」。同一函数中 commContext（行 206 GetDim(0) <= 0）和 gradFetched（行 221 numTokens < 0、行 226 hiddenDim <= 0）均对 dim 值做了显式校验，sendCounts/recvLocalEntry 缺失同类校验，形成不一致的安全缺口。 **风险链路 1（sendCounts）**：dim(0) 若为负值或超过 UINT32_MAX，行 260 static_cast<uint32_t> 会发生窄化截断/回绕，rankSize 变为巨大错误值 → SetBlockDim 和 workspace 计算均受影响。 **风险链路 2（recvLocalEntry）**：dim(0) 赋值给 int64_t totalRecv（行 281），若为负值则保留负数 → wsRecvGrad = totalRecv * hiddenBytes 为负 → wsTotal 求和可能为负 → static_cast<size_t>(wsTotal) 将负数回绕为极大 size_t 值，导致 workspace 申请异常。 - **代码片段**（mc2/engram_fetch_grad/op_host/op_tiling/engram_fetch_grad_tiling.cpp 行 251-281）： cpp // sendCounts: 1D (W,) const gert::StorageShape *sendCountsShape = context->GetInputShape(IN_SEND_COUNTS); OP_TILING_CHECK(sendCountsShape == nullptr, OP_LOGE_WITH_INVALID_INPUT(nodeName, "sendCounts"), return ge::GRAPH_FAILED); OP_TILING_CHECK(sendCountsShape->GetStorageShape().GetDimNum() != DIM_ONE, OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(nodeName, "sendCounts", (std::to_string(sendCountsShape->GetStorageShape().GetDimNum()) + "D").c_str(), "The shape dim of sendCounts must be 1D."), return ge::GRAPH_FAILED); rankSize = static_cast<uint32_t>(sendCountsShape->GetStorageShape().GetDim(0)); // recvCounts: 1D (W,) const gert::StorageShape *recvCountsShape = context->GetInputShape(IN_RECV_COUNTS); OP_TILING_CHECK(recvCountsShape == nullptr, OP_LOGE_WITH_INVALID_INPUT(nodeName, "recvCounts"), return ge::GRAPH_FAILED); OP_TILING_CHECK(recvCountsShape->GetStorageShape().GetDimNum() != DIM_ONE, OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(nodeName, "recvCounts", (std::to_string(recvCountsShape->GetStorageShape().GetDimNum()) + "D").c_str(), "The shape dim of recvCounts must be 1D."), return ge::GRAPH_FAILED); // recvLocalEntry: 1D (R,) const gert::StorageShape *recvLocalEntryShape = context->GetInputShape(IN_RECV_LOCAL_ENTRY); OP_TILING_CHECK(recvLocalEntryShape == nullptr, OP_LOGE_WITH_INVALID_INPUT(nodeName, "recvLocalEntry"), return ge::GRAPH_FAILED); OP_TILING_CHECK(recvLocalEntryShape->GetStorageShape().GetDimNum() != DIM_ONE, OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(nodeName, "recvLocalEntry", (std::to_string(recvLocalEntryShape->GetStorageShape().GetDimNum()) + "D").c_str(), "The shape dim of recvLocalEntry must be 1D."), return ge::GRAPH_FAILED); totalRecv = recvLocalEntryShape->GetStorageShape().GetDim(0); - **假设检验证据**： - (+40%) 规范违反：SEC-4.1 要求外部输入数据需做合法性校验，sendCounts dim(0) 直接 cast uint32_t、recvLocalEntry dim(0) 直接赋值 int64_t，均无值范围校验 - (+30%) 上下文防御缺失：同函数内 commContext/gradFetched 均有非负校验，唯独 sendCounts/recvLocalEntry 缺失 - (+20%) PR 归属：新增文件，全部行在 diff 变更范围内 - (+15%) 数据流风险：rankSize 负值回绕 / totalRecv 负值导致 workspace 回绕为巨大值 - 自信值 = 105% ≥ 70% → 判定违规 - **修复建议**： 1. 在行 260 之前增加 sendCounts dim(0) 值校验（> 0） 2. 在行 281 之前增加 recvLocalEntry dim(0) 值校验（>= 0） 3. 参照同函数内 commContext/gradFetched 的校验模式 --- ### 2. [engram_fetch_grad_host_G1] TOPK-7 融合规则/InferShape/Tiling 外部输入校验 - **问题描述**：TOPK-7 要求「Tiling 外部输入使用时必须进行合法性校验」，校验内容包括「Shape 维度是否合法（不为负数、不超过限制）」。sendCounts dim(0) 和 recvLocalEntry dim(0) 作为外部输入被提取后直接使用，未做非负/范围校验。同函数内 commContext dim0（<=0 校验）、gradFetched dim0（<0 校验）、hiddenDim（<=0 校验）均有非负防御，唯独 sendCounts 和 recvLocalEntry 缺失。此问题与 SEC-4.1 同构。 - **代码片段**（mc2/engram_fetch_grad/op_host/op_tiling/engram_fetch_grad_tiling.cpp 行 251-281）： cpp // sendCounts: 1D (W,) const gert::StorageShape *sendCountsShape = context->GetInputShape(IN_SEND_COUNTS); OP_TILING_CHECK(sendCountsShape == nullptr, OP_LOGE_WITH_INVALID_INPUT(nodeName, "sendCounts"), return ge::GRAPH_FAILED); OP_TILING_CHECK(sendCountsShape->GetStorageShape().GetDimNum() != DIM_ONE, OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(nodeName, "sendCounts", (std::to_string(sendCountsShape->GetStorageShape().GetDimNum()) + "D").c_str(), "The shape dim of sendCounts must be 1D."), return ge::GRAPH_FAILED); rankSize = static_cast<uint32_t>(sendCountsShape->GetStorageShape().GetDim(0)); // ...（recvLocalEntry 同理） totalRecv = recvLocalEntryShape->GetStorageShape().GetDim(0); - **假设检验证据**： - (+40%) 规范违反：TOPK-7 要求 Shape 维度不为负数 - (+30%) 上下文防御缺失：同函数内其他 tensor 均有非负校验 - (+20%) PR 归属：新增文件 - (+15%) 数据流风险：负值 dim0 经 static_cast 回绕或以负 int64_t 传播至 workspace 计算 - 自信值 = 105% ≥ 70% → 判定违规 - **修复建议**：与 SEC-4.1 同一问题，修复方案一致。 --- ### 3. [engram_fetch_train_host_G1] SEC-4.1 外部输入数据需要做合法性校验 - **问题描述**：CheckTensorDim 中 sendCountsOut 和 recvLocalEntryOut 的 dim0 值作为外部输入未做范围校验。rankSize（uint32_t）通过 static_cast 从 int64_t 缩窄转换，若 dim0 为负或超过 UINT32_MAX 将静默截断/回绕为巨大值；totalRecv（int64_t）未校验非负，若为负值将导致 wsLocalData = totalRecv * hiddenBytes 为负，进而 wsTotal 为负，再经 static_cast<size_t>(wsTotal) 回绕为约 1.8e19 的巨大 workspace 申请。同函数内 commContextShape dim0 > 0（行 188）、numTokens >= 0（行 202）、numRecvOutShape dim0 == 1（行 272）均已校验，此处校验明显不一致。 - **代码片段**（mc2/engram_fetch_train/op_host/op_tiling/engram_fetch_train_tiling.cpp 行 235-262）： cpp const gert::StorageShape *sendCountsOutShape = context->GetOutputShape(OUT_SEND_COUNTS); OP_TILING_CHECK(sendCountsOutShape == nullptr, OP_LOGE_WITH_INVALID_INPUT(nodeName, "sendCountsOut"), return ge::GRAPH_FAILED); OP_TILING_CHECK(sendCountsOutShape->GetStorageShape().GetDimNum() != DIM_ONE, OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(nodeName, "sendCountsOut", (std::to_string(sendCountsOutShape->GetStorageShape().GetDimNum()) + "D").c_str(), "The shape dim of sendCountsOut must be 1D."), return ge::GRAPH_FAILED); rankSize = static_cast<uint32_t>(sendCountsOutShape->GetStorageShape().GetDim(0)); const gert::StorageShape *recvCountsOutShape = context->GetOutputShape(OUT_RECV_COUNTS); OP_TILING_CHECK(recvCountsOutShape == nullptr, OP_LOGE_WITH_INVALID_INPUT(nodeName, "recvCountsOut"), return ge::GRAPH_FAILED); OP_TILING_CHECK(recvCountsOutShape->GetStorageShape().GetDimNum() != DIM_ONE, OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(nodeName, "recvCountsOut", (std::to_string(recvCountsOutShape->GetStorageShape().GetDimNum()) + "D").c_str(), "The shape dim of recvCountsOut must be 1D."), return ge::GRAPH_FAILED); const gert::StorageShape *recvLocalEntryOutShape = context->GetOutputShape(OUT_RECV_LOCAL_ENTRY); OP_TILING_CHECK(recvLocalEntryOutShape == nullptr, OP_LOGE_WITH_INVALID_INPUT(nodeName, "recvLocalEntryOut"), return ge::GRAPH_FAILED); OP_TILING_CHECK(recvLocalEntryOutShape->GetStorageShape().GetDimNum() != DIM_ONE, OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(nodeName, "recvLocalEntryOut", (std::to_string(recvLocalEntryOutShape->GetStorageShape().GetDimNum()) + "D").c_str(), "The shape dim of recvLocalEntryOut must be 1D."), return ge::GRAPH_FAILED); totalRecv = recvLocalEntryOutShape->GetStorageShape().GetDim(0); - **假设检验证据**： - (+40%) 规范违反：外部 shape dim 值未做范围/非负校验 - (+30%) 上下文防御缺失：作用域内仅校验 dim 数，无 dim 值范围校验 - (+20%) PR 归属：新文件，全量代码均在 diff 范围 - (+15%) 数据流风险：rankSize 缩窄截断 / totalRecv 负值 → workspace 回绕 - 自信值 = 105% ≥ 70% → 判定违规 - **修复建议**： 1. 在行 243 前增加 sendCountsOut dim(0) 非负及 <= UINT32_MAX 校验 2. 在行 262 前增加 recvLocalEntryOut dim(0) >= 0 校验 3. 参照同函数内 numTokens（行 202-205）的校验模式 --- ### 4. [engram_fetch_train_host_G1] TOPK-7 融合规则/InferShape/Tiling 外部输入校验 - **问题描述**：与 SEC-4.1 同一问题。TOPK-7 要求 Tiling 外部输入的 Shape 维度不为负数、不超过限制。sendCountsOut dim0（赋值给 rankSize，缩窄为 uint32_t）和 recvLocalEntryOut dim0（赋值给 totalRecv）均未做非负/范围校验。同函数内 commContextShape dim0 > 0（行 188）、indicesShape numTokens >= 0（行 202）均已校验，存在校验不一致。 - **代码片段**（mc2/engram_fetch_train/op_host/op_tiling/engram_fetch_train_tiling.cpp 行 235-262）：（同 SEC-4.1） - **假设检验证据**： - (+40%) 规范违反：Tiling 外部 shape dim 值未做非负/范围校验 - (+30%) 上下文防御缺失：仅校验 dim 数无 dim 值校验 - (+20%) PR 归属：新文件全量在 diff 范围 - (+15%) 数据流风险：rankSize 缩窄截断 / totalRecv 负值 → work See merge request: cann/ops-transformer!8914 | 2 个月前 |
