【合入来源】

如有社区issue，请关联issue链接
请勿携带内部流程信息（需求链接、问题单、内部issue等）

• 需求
• 问题单
• issue/工单
• 重构优化
• 资料更新

【修改方案】

将 NPU Caching Allocator 对齐到上游 CUDA 实现，补齐 Allocator Trace Tracker 回调机制，并将 mempool 关联方式从间接反查改为直接存储。

1. 新增 AllocatorTraceTracker 回调机制（NPUCachingAllocator.h + .cpp）

• 新增 AllocatorTraceTracker 类型、虚方法、trace_trackers_ 成员、公共接口层转发及便捷函数，与 CUDA 接口对齐

2. record_trace 支持 tracker 回调分发（NPUCachingAllocator.cpp）

• 入口条件改为 !record_history && trace_trackers_.empty()，构造 TraceEntry 后遍历 trace_trackers_ 调用回调

3. TraceEntry 新增 mempool_ 字段（NPUCachingAllocator.h）

• 构造函数新增 MempoolId_t mempool = {0, 0} 参数，新增 MempoolId_t mempool_ 成员

4. PrivatePool 改为直接存储 mempool ID（NPUCachingAllocator.cpp）

• 新增 id 成员，构造函数接受 MempoolId_t id
• 新增 BlockPool::owner_MempoolId() 方法
• 新增 create_or_incref_pool / get_private_pool，重构 beginAllocateToPool / endAllocateToPool

5. 所有 record_trace 调用点补齐 mempool_id 参数（NPUCachingAllocator.cpp）

• 8 处调用点传入 block->pool->owner_MempoolId()

6. NPUWorkspaceAllocator 补齐 TraceEntry 构造参数（NPUWorkspaceAllocator.cpp）

• 3 处构造补齐 MempoolId_t{0, 0}

7. NPUPluggableAllocator 补齐接口（NPUPluggableAllocator.h + .cpp）

• 新增 attachAllocatorTraceTracker override，TORCH_CHECK(false) 拒绝调用

8. 新增测试

• C++ 扩展注册回调统计 SEGMENT_ALLOC/FREE 事件，Python 测试验证两种场景下回调正确触发

【资料变更】

不涉及

【接口变更】

不涉及

【功能验证】

test/npu/test_allocator_trace_tracker.py

【CheckList】

PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x]

• 代码注释完备，正确记录错误日志
• 代码实现进行了返回值、空指针等校验
• PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等
• PR持续集成流水线（CI）执行通过，代码检查无异常