PyPTO 定制算子 gather_in_l1面向 PagedAttention KV Cache 的 GM 到 L1 离散行搬运接口解析【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto本文以 PyPTO 的定制接口pypto.experimental.gather_in_l1为核心完整解析其函数原型、参数语义与调用方式并结合 C 接口层、代码生成层与设备侧实现的源码说明该算子如何按页表block_table把 GM 上 KV Cache 的离散 token 行高效搬运到 L1以及它与pypto.matmul配合完成 PagedAttention 类稀疏矩阵乘的完整用法。产品支持情况根据接口文档 pypto-experimental-gather_in_l1.md该接口的产品支持情况如下Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持需要特别注意文档明确标注该接口为定制接口不保证稳定性其语义与底层实现均围绕特定场景页式 KV Cache 的离散行 gather定制接口签名与行为可能随版本演进变化生产使用前应锁定版本并关注接口变更。功能定位为什么需要 gather_in_l1在 PagedAttention 这类推理加速机制中KV Cache 被切分为固定大小的“页”block存放在 GMGlobal Memory中通过一张页表block_table记录逻辑块到物理块的映射。计算 attention 时只有 TopK 选中的少量 token 会被使用这些 token 在 GM 上的物理位置是离散的——逐个用普通向量式 gather 从 GM 读回再参与 Cube 矩阵乘会带来大量碎搬运开销。gather_in_l1正是为此场景定制的算子从 GM 上的 Tensor 离散搬运指定行的数据同时每行搬运前 size 个数据至 L1。从源码文件头部的注释可以印证这一定位gather_in_l1_impl.h 中写明该实现 “only used in deepseek model (Atlas A3, Ascend 950PR/Ascend 950DT)”即它是面向页式 KV Cache 稀疏读取的定制搬运算子。它直接生成 GM 到 L1 的TLOAD搬运指令使搬运行号通过页表在硬件侧完成“逻辑 token 号 → 物理行号”的换算从而把离散的 KV 行以 Tile 形式落入 L1供后续matmul直接消费。函数原型gather_in_l1(src: Tensor, indices: Tensor, block_table: Tensor, block_size: int, size: int, is_b_matrix: bool, is_trans: bool) - TensorPython 侧的定义位于 gather.py它是一个薄封装op_wrapper def gather_in_l1( src: Tensor, indices: Tensor, block_table: Tensor, block_size: int, size: int, is_b_matrix: bool, is_trans: bool ) - Tensor: gather_in_l1. return pypto_impl.gather_in_l1(src, indices, block_table, block_size, size, is_b_matrix, is_trans)并通过 experimental/init.py 导出因此调用方使用pypto.experimental.gather_in_l1访问。同文件还定义了语义相近的姊妹算子gather_in_ub结果写入 UB 而非 L1二者共享“KV Cache 页表 逻辑索引”的参数范式可对照阅读。参数说明参数名输入/输出说明src输入源操作数。支持的数据类型为DT_FP32DT_FP16DT_BF16DT_INT8。不支持空 Tensor支持两维。在页式场景中即 KV Cache 的物理 buffer形状为[物理token总数, hidden_dim]。indices输入源操作数的行偏移。支持的数据类型为DT_INT32DT_INT64。不支持空 Tensor支持两维。Shape 形状为[1,n]。语义上是逻辑 token 索引如 TopK 结果n 即选出的 token 个数。block_table输入页表。支持的数据类型为 DT_INT32。不支持空 Tensor支持两维。在实际使用中表示为 Page Attention 中的页表形状为[1,block_table_size]其中 block_table_size 表示页表的长度内容为“逻辑块 ID → 物理块 ID”的映射。block_size输入int 类型。表示 Page Attention 中一个块可以放多少个 token即每页容纳的 token 数。size输入每行搬运的数据数。数据数要小于源操作数的列数接口层实现中的校验为size src.shape[1]。is_b_matrix输入搬运后的结果即输出 Tensor 是否作为 matmul 的 B 矩阵。is_trans输入搬运后的结果即输出 Tensor 是否转置。参数语义的源码印证C 接口层的入口是 operation_impl.cpp 中的experimental::GatherInL1模板函数其模板参数isB、isTrans分别对应 Python 层的is_b_matrix与is_trans四个组合false,false / false,true / true,false / true,true均被显式实例化。该函数中包含以下关键校验与输出形状推导constexpr int32_t NUM_SIZE 2; CHECK_OP(src.GetShape().size() NUM_SIZE); CHECK_OP(offsets.GetShape().size() NUM_SIZE); // offsets必须是两维是因为不支持1维的Tensor CHECK_OP(offsets.GetShape()[0] 1); CHECK_OP(size src.GetShape()[1]); Tensor dst(src.GetStorage()-Datatype(), {offsets.GetShape()[1], size});可以从中确认三条文档之外的硬性约束src、indices必须是两维Tensor注释说明了indices必须两维的原因当前实现不支持一维 Tensorindices的第一维必须为 1即形状严格为[1, n]size不得超过src的列数。输出 Tensor 的数据类型与src相同形状为[indices.shape[1], size]——即选出的 n 个逻辑 token每个取出前size个元素。动态有效形状也会同步设置UpdateDynValidShape保证动态 shape 场景下的正确性。返回值说明返回输出 Tensor数据类型与src一致形状为[n, size]n为indices第二维。在 tile 图层面当算子被按 Cube Tile 切分时每个 tile 的有效形状由 op_infer_shape_impl.cpp 中的InferFuncGatherInL1推导为[offsets.shape[1], min(src.shape[1] - startOffset, output.shape[1])]其中startOffset是该 tile 在源张量列方向上的起始偏移用于处理size跨 tile 切列的边界情况。约束说明文档标注“约束说明无”即官方未声明额外的使用约束。但结合接口层与代码生成层的校验实际使用时需满足三个 Tensor 参数均不支持空 Tensor且均为两维indices形状为[1, n]block_table形状为[1, block_table_size]size src.shape[1]src数据类型限 DT_FP32 / DT_FP16 / DT_BF16 / DT_INT8indices限 DT_INT32 / DT_INT64block_table限 DT_INT32该接口为定制接口不保证稳定性接口行为可能随版本变化。调用示例接口文档给出的最小示例如下src pypto.tensor([16, 32], pypto.DT_FP32, tensor_src) offset pypto.tensor([1, 32], pypto.DT_INT32, tensor_offset) block_table pypto.tensor([1, 4], pypto.DT_INT32, block_table) block_size 2 size 16 is_b_matrix False is_trans False out pypto.experimental.gather_in_l1(src, offset, block_table, block_size, size, is_b_matrix, is_trans)示例中src是[16, 32]的物理 buffer16 个 token每行 32 列offset含 32 个逻辑行号block_table含 4 个页表项block_size2时页表覆盖4 * 2 8个逻辑 token每行只搬运前 16 个数据输出形状为[32, 16]。页表到物理行号的映射规则该算子的核心语义是“逻辑 token 号 → 物理行号”的换算。从 C 系统测试 test_gather_in_l1.cpp 中的 golden 计算函数可以清晰看到这一规则IndexType logical_block_id logical_index / static_castIndexType(cfg.block_size); IndexType physical_block_id page_table[logical_block_id]; IndexType block_offset logical_index % static_castIndexType(cfg.block_size); IndexType physical_index physical_block_id * static_castIndexType(cfg.block_size) block_offset;即物理行号 page_table[逻辑token // block_size] * block_size 逻辑token % block_size。Python 侧系统测试 test_gather_in_l1.py 用同样的公式生成 goldenextracted ( page_table.view(-1)[indices.view(-1) // config.block_size] * config.block_size indices.view(-1) % config.block_size ) extracted_tokens all_buffer[extracted]设备侧实现 gather_in_l1_impl.h 中这一换算由CalaOffset2PageAttention在搬运循环内逐行完成随后用TASSIGN/TLOAD把该行前size列的数据落入 L1for (int64_t i 0; i loop; i) { uint64_t gatherOffset offset.GetAddr()[i offsetsStartOffset]; gatherOffset CalaOffset2PageAttentionuint64_t, typename BlockT::Type, blockSize( block.GetAddr() GMBlockTableOffset, gatherOffset); globalData src0Global( (__gm__ typename GlobalData::Type*)(src.GetAddr() gatherOffset * srcCol srcColumnStartOffset), pto::Shape1, 1, 1, -1, -1(srcShape0, srcShape1), pto::Stride1, 1, 1, -1, -1(srcStride0, srcStride1)); tileData dstL1(dstShape0, dstShape1); pto::TASSIGN(dstL1, (uint64_t)((__cbuf__ typename TileData::Type*)dst.GetAddr() i * c0Size)); pto::TLOAD(dstL1, src0Global); }该实现按 dst 列数是否整除c0SizeBLOCK_ALIGN_BYTE / 元素字节数分为两条路径不整除时按普通 ND 布局逐行搬运整除时切换为性能更高的 ND2ND 搬运方式源码注释说明GM 上把(1, dstShape1)视为(dstShape1 / c0Size, c0Size)L1 上按 NZ 分块展平以减少 Cube 消费前的格式转换开销。与 pypto.matmul 配合的完整 Kernel 示例is_b_matrix与is_trans的设计意图是让 gather 结果直接作为 matmul 的 A 或 B 矩阵避免落 GM 后再读回。仓库中 Python 系统测试 test_gather_in_l1.py 给出了一个完整可运行的 kernelpypto.frontend.jit(debug_options{runtime_debug_mode: 0, compile_debug_mode: 0}) def gather_matmul_pto_kernel( src: pypto.Tensor(), indices: pypto.Tensor(), page_table: pypto.Tensor(), mat2_tensor: pypto.Tensor(), out_tensor: pypto.Tensor(), config: GatherInL1Config, ): pypto.set_cube_tile_shapes([128, 128], [128, 128], [128, 128]) dyn_src pypto.view(src, src.shape, [0, 0], valid_shapesrc.shape) dyn_offsets pypto.view(indices, indices.shape, [0, 0], valid_shapeindices.shape) dyn_page_table pypto.view(page_table, page_table.shape, [0, 0], valid_shapepage_table.shape) input_tensor pypto.experimental.gather_in_l1( dyn_src, dyn_offsets, dyn_page_table, config.block_size, config.token_dim, is_b_matrixconfig.is_gather_b_matrix, is_transconfig.is_gather_trans, ) if not config.is_gather_b_matrix: out_tensor[:] pypto.matmul( input_tensor, mat2_tensor, config.out_dtype, a_transconfig.is_gather_trans, b_transFalse ) else: out_tensor[:] pypto.matmul( mat2_tensor, input_tensor, config.out_dtype, a_transFalse, b_transconfig.is_gather_trans )要点解析动态 view三个输入先经pypto.view(..., valid_shape...)包装成动态 shape 张量这与接口层对DynValidShape的依赖一致CHECK_OP(!offsets.GetStorage()-GetDynValidShape().empty())要求 offsets 具备动态有效形状is_b_matrix与is_trans的四种组合gather 结果作为 A 矩阵时is_b_matrixFalsea_trans取is_trans作为 B 矩阵时is_b_matrixTrueb_trans取is_trans。四种组合与GatherInL1isB, isTrans的四个模板实例一一对应测试数据构造测试通过torch.randperm随机打乱页表模拟物理块复用torch.randint生成逻辑索引再按页表公式生成 golden 结果最后用torch.allclose校验覆盖[topk, hidden]的 gather 语义。Tile 切分与代码生成链路当 gather 结果参与大矩阵乘时算子会被按 Cube Tile 切分。operation_impl.cpp 的TiledGatherInL1展示了切分逻辑按 dst 的行维A 矩阵取cubeTile.m[1]、B 矩阵取cubeTile.n[1]转置时两维交换与列维循环为每个 tile 截取对应的offsetsTile行维切片[1, shape0]从{0, i}起和整张blockTableTile并写入两个关键属性startOffset该 tile 在源张量列方向的起始偏移供形状推导与设备端srcColumnStartOffset使用op_attr_blocksize页大小透传到设备端模板参数。代码生成侧codegen_mte_gather.cpp 的GenGatherInL1将 tile 算子翻译为 CCE 调用并在此处再次做防御性校验三个操作数均只支持两维、offsets 类型必须是int32_t或int64_t、dst 与 src 类型必须一致最终生成的调用形如TGatherInL1dstType, offsetsType, blockTableType, dstShape0, offsetsShape1, srcColumnStartOffset, blockSize(dst, dstValidShape, src, srcShape1, offsets, blockTable, offsetsStartOffset, blockTableGMStride, blockTableStartOffset)即把逻辑→物理换算与 GM→L1 搬运全部下沉到 tile 内完成。tile 算子TGatherInL1的声明位于 cube_pto.h入口处通过CheckShapeValid做动态形状合法性检查。测试与验证路径Python 系统测试test_gather_in_l1.py参数化用例配置在 gather_in_l1_test_case.py包含小容量num_buffer_tokens33, token_dim30, block_size10, DT_FP16与大容量num_buffer_tokens82816, topk2048, block_size128, DT_INT8 → DT_INT32两类用例并在 NPU / SIM 两种pypto.RunMode下运行C 系统测试test_gather_in_l1.cpp覆盖isB / isTrans组合下的 matmul 联动验证并用页表 golden 函数逐元素比对代码生成单测test_codegen_dyn_copy.cpp 中包含GATHER_IN_L1的动态代码生成检查。小结pypto.experimental.gather_in_l1是 PyPTO 面向 PagedAttention KV Cache 场景定制的 GM→L1 离散行搬运算子以[1, n]的逻辑行偏移、[1, block_table_size]的页表和block_size页大小完成“逻辑 token → 物理行”换算把每行前size列数据直接以 L1 Tile 形式落地并通过is_b_matrix/is_trans让结果无缝对接pypto.matmul的 A/B 矩阵含转置消费。使用时应牢记它是两维 Tensor 定制接口、不保证稳定性参数形状与类型约束建议对照接口层校验见上文“约束说明”逐项核对。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 PyPTO 定制算子 gather_in_l1：面向 PagedAttention KV Cache 的 GM 到 L1 离散行搬运接口解析
PyPTO 定制算子 gather_in_l1面向 PagedAttention KV Cache 的 GM 到 L1 离散行搬运接口解析【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto本文以 PyPTO 的定制接口pypto.experimental.gather_in_l1为核心完整解析其函数原型、参数语义与调用方式并结合 C 接口层、代码生成层与设备侧实现的源码说明该算子如何按页表block_table把 GM 上 KV Cache 的离散 token 行高效搬运到 L1以及它与pypto.matmul配合完成 PagedAttention 类稀疏矩阵乘的完整用法。产品支持情况根据接口文档 pypto-experimental-gather_in_l1.md该接口的产品支持情况如下Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持需要特别注意文档明确标注该接口为定制接口不保证稳定性其语义与底层实现均围绕特定场景页式 KV Cache 的离散行 gather定制接口签名与行为可能随版本演进变化生产使用前应锁定版本并关注接口变更。功能定位为什么需要 gather_in_l1在 PagedAttention 这类推理加速机制中KV Cache 被切分为固定大小的“页”block存放在 GMGlobal Memory中通过一张页表block_table记录逻辑块到物理块的映射。计算 attention 时只有 TopK 选中的少量 token 会被使用这些 token 在 GM 上的物理位置是离散的——逐个用普通向量式 gather 从 GM 读回再参与 Cube 矩阵乘会带来大量碎搬运开销。gather_in_l1正是为此场景定制的算子从 GM 上的 Tensor 离散搬运指定行的数据同时每行搬运前 size 个数据至 L1。从源码文件头部的注释可以印证这一定位gather_in_l1_impl.h 中写明该实现 “only used in deepseek model (Atlas A3, Ascend 950PR/Ascend 950DT)”即它是面向页式 KV Cache 稀疏读取的定制搬运算子。它直接生成 GM 到 L1 的TLOAD搬运指令使搬运行号通过页表在硬件侧完成“逻辑 token 号 → 物理行号”的换算从而把离散的 KV 行以 Tile 形式落入 L1供后续matmul直接消费。函数原型gather_in_l1(src: Tensor, indices: Tensor, block_table: Tensor, block_size: int, size: int, is_b_matrix: bool, is_trans: bool) - TensorPython 侧的定义位于 gather.py它是一个薄封装op_wrapper def gather_in_l1( src: Tensor, indices: Tensor, block_table: Tensor, block_size: int, size: int, is_b_matrix: bool, is_trans: bool ) - Tensor: gather_in_l1. return pypto_impl.gather_in_l1(src, indices, block_table, block_size, size, is_b_matrix, is_trans)并通过 experimental/init.py 导出因此调用方使用pypto.experimental.gather_in_l1访问。同文件还定义了语义相近的姊妹算子gather_in_ub结果写入 UB 而非 L1二者共享“KV Cache 页表 逻辑索引”的参数范式可对照阅读。参数说明参数名输入/输出说明src输入源操作数。支持的数据类型为DT_FP32DT_FP16DT_BF16DT_INT8。不支持空 Tensor支持两维。在页式场景中即 KV Cache 的物理 buffer形状为[物理token总数, hidden_dim]。indices输入源操作数的行偏移。支持的数据类型为DT_INT32DT_INT64。不支持空 Tensor支持两维。Shape 形状为[1,n]。语义上是逻辑 token 索引如 TopK 结果n 即选出的 token 个数。block_table输入页表。支持的数据类型为 DT_INT32。不支持空 Tensor支持两维。在实际使用中表示为 Page Attention 中的页表形状为[1,block_table_size]其中 block_table_size 表示页表的长度内容为“逻辑块 ID → 物理块 ID”的映射。block_size输入int 类型。表示 Page Attention 中一个块可以放多少个 token即每页容纳的 token 数。size输入每行搬运的数据数。数据数要小于源操作数的列数接口层实现中的校验为size src.shape[1]。is_b_matrix输入搬运后的结果即输出 Tensor 是否作为 matmul 的 B 矩阵。is_trans输入搬运后的结果即输出 Tensor 是否转置。参数语义的源码印证C 接口层的入口是 operation_impl.cpp 中的experimental::GatherInL1模板函数其模板参数isB、isTrans分别对应 Python 层的is_b_matrix与is_trans四个组合false,false / false,true / true,false / true,true均被显式实例化。该函数中包含以下关键校验与输出形状推导constexpr int32_t NUM_SIZE 2; CHECK_OP(src.GetShape().size() NUM_SIZE); CHECK_OP(offsets.GetShape().size() NUM_SIZE); // offsets必须是两维是因为不支持1维的Tensor CHECK_OP(offsets.GetShape()[0] 1); CHECK_OP(size src.GetShape()[1]); Tensor dst(src.GetStorage()-Datatype(), {offsets.GetShape()[1], size});可以从中确认三条文档之外的硬性约束src、indices必须是两维Tensor注释说明了indices必须两维的原因当前实现不支持一维 Tensorindices的第一维必须为 1即形状严格为[1, n]size不得超过src的列数。输出 Tensor 的数据类型与src相同形状为[indices.shape[1], size]——即选出的 n 个逻辑 token每个取出前size个元素。动态有效形状也会同步设置UpdateDynValidShape保证动态 shape 场景下的正确性。返回值说明返回输出 Tensor数据类型与src一致形状为[n, size]n为indices第二维。在 tile 图层面当算子被按 Cube Tile 切分时每个 tile 的有效形状由 op_infer_shape_impl.cpp 中的InferFuncGatherInL1推导为[offsets.shape[1], min(src.shape[1] - startOffset, output.shape[1])]其中startOffset是该 tile 在源张量列方向上的起始偏移用于处理size跨 tile 切列的边界情况。约束说明文档标注“约束说明无”即官方未声明额外的使用约束。但结合接口层与代码生成层的校验实际使用时需满足三个 Tensor 参数均不支持空 Tensor且均为两维indices形状为[1, n]block_table形状为[1, block_table_size]size src.shape[1]src数据类型限 DT_FP32 / DT_FP16 / DT_BF16 / DT_INT8indices限 DT_INT32 / DT_INT64block_table限 DT_INT32该接口为定制接口不保证稳定性接口行为可能随版本变化。调用示例接口文档给出的最小示例如下src pypto.tensor([16, 32], pypto.DT_FP32, tensor_src) offset pypto.tensor([1, 32], pypto.DT_INT32, tensor_offset) block_table pypto.tensor([1, 4], pypto.DT_INT32, block_table) block_size 2 size 16 is_b_matrix False is_trans False out pypto.experimental.gather_in_l1(src, offset, block_table, block_size, size, is_b_matrix, is_trans)示例中src是[16, 32]的物理 buffer16 个 token每行 32 列offset含 32 个逻辑行号block_table含 4 个页表项block_size2时页表覆盖4 * 2 8个逻辑 token每行只搬运前 16 个数据输出形状为[32, 16]。页表到物理行号的映射规则该算子的核心语义是“逻辑 token 号 → 物理行号”的换算。从 C 系统测试 test_gather_in_l1.cpp 中的 golden 计算函数可以清晰看到这一规则IndexType logical_block_id logical_index / static_castIndexType(cfg.block_size); IndexType physical_block_id page_table[logical_block_id]; IndexType block_offset logical_index % static_castIndexType(cfg.block_size); IndexType physical_index physical_block_id * static_castIndexType(cfg.block_size) block_offset;即物理行号 page_table[逻辑token // block_size] * block_size 逻辑token % block_size。Python 侧系统测试 test_gather_in_l1.py 用同样的公式生成 goldenextracted ( page_table.view(-1)[indices.view(-1) // config.block_size] * config.block_size indices.view(-1) % config.block_size ) extracted_tokens all_buffer[extracted]设备侧实现 gather_in_l1_impl.h 中这一换算由CalaOffset2PageAttention在搬运循环内逐行完成随后用TASSIGN/TLOAD把该行前size列的数据落入 L1for (int64_t i 0; i loop; i) { uint64_t gatherOffset offset.GetAddr()[i offsetsStartOffset]; gatherOffset CalaOffset2PageAttentionuint64_t, typename BlockT::Type, blockSize( block.GetAddr() GMBlockTableOffset, gatherOffset); globalData src0Global( (__gm__ typename GlobalData::Type*)(src.GetAddr() gatherOffset * srcCol srcColumnStartOffset), pto::Shape1, 1, 1, -1, -1(srcShape0, srcShape1), pto::Stride1, 1, 1, -1, -1(srcStride0, srcStride1)); tileData dstL1(dstShape0, dstShape1); pto::TASSIGN(dstL1, (uint64_t)((__cbuf__ typename TileData::Type*)dst.GetAddr() i * c0Size)); pto::TLOAD(dstL1, src0Global); }该实现按 dst 列数是否整除c0SizeBLOCK_ALIGN_BYTE / 元素字节数分为两条路径不整除时按普通 ND 布局逐行搬运整除时切换为性能更高的 ND2ND 搬运方式源码注释说明GM 上把(1, dstShape1)视为(dstShape1 / c0Size, c0Size)L1 上按 NZ 分块展平以减少 Cube 消费前的格式转换开销。与 pypto.matmul 配合的完整 Kernel 示例is_b_matrix与is_trans的设计意图是让 gather 结果直接作为 matmul 的 A 或 B 矩阵避免落 GM 后再读回。仓库中 Python 系统测试 test_gather_in_l1.py 给出了一个完整可运行的 kernelpypto.frontend.jit(debug_options{runtime_debug_mode: 0, compile_debug_mode: 0}) def gather_matmul_pto_kernel( src: pypto.Tensor(), indices: pypto.Tensor(), page_table: pypto.Tensor(), mat2_tensor: pypto.Tensor(), out_tensor: pypto.Tensor(), config: GatherInL1Config, ): pypto.set_cube_tile_shapes([128, 128], [128, 128], [128, 128]) dyn_src pypto.view(src, src.shape, [0, 0], valid_shapesrc.shape) dyn_offsets pypto.view(indices, indices.shape, [0, 0], valid_shapeindices.shape) dyn_page_table pypto.view(page_table, page_table.shape, [0, 0], valid_shapepage_table.shape) input_tensor pypto.experimental.gather_in_l1( dyn_src, dyn_offsets, dyn_page_table, config.block_size, config.token_dim, is_b_matrixconfig.is_gather_b_matrix, is_transconfig.is_gather_trans, ) if not config.is_gather_b_matrix: out_tensor[:] pypto.matmul( input_tensor, mat2_tensor, config.out_dtype, a_transconfig.is_gather_trans, b_transFalse ) else: out_tensor[:] pypto.matmul( mat2_tensor, input_tensor, config.out_dtype, a_transFalse, b_transconfig.is_gather_trans )要点解析动态 view三个输入先经pypto.view(..., valid_shape...)包装成动态 shape 张量这与接口层对DynValidShape的依赖一致CHECK_OP(!offsets.GetStorage()-GetDynValidShape().empty())要求 offsets 具备动态有效形状is_b_matrix与is_trans的四种组合gather 结果作为 A 矩阵时is_b_matrixFalsea_trans取is_trans作为 B 矩阵时is_b_matrixTrueb_trans取is_trans。四种组合与GatherInL1isB, isTrans的四个模板实例一一对应测试数据构造测试通过torch.randperm随机打乱页表模拟物理块复用torch.randint生成逻辑索引再按页表公式生成 golden 结果最后用torch.allclose校验覆盖[topk, hidden]的 gather 语义。Tile 切分与代码生成链路当 gather 结果参与大矩阵乘时算子会被按 Cube Tile 切分。operation_impl.cpp 的TiledGatherInL1展示了切分逻辑按 dst 的行维A 矩阵取cubeTile.m[1]、B 矩阵取cubeTile.n[1]转置时两维交换与列维循环为每个 tile 截取对应的offsetsTile行维切片[1, shape0]从{0, i}起和整张blockTableTile并写入两个关键属性startOffset该 tile 在源张量列方向的起始偏移供形状推导与设备端srcColumnStartOffset使用op_attr_blocksize页大小透传到设备端模板参数。代码生成侧codegen_mte_gather.cpp 的GenGatherInL1将 tile 算子翻译为 CCE 调用并在此处再次做防御性校验三个操作数均只支持两维、offsets 类型必须是int32_t或int64_t、dst 与 src 类型必须一致最终生成的调用形如TGatherInL1dstType, offsetsType, blockTableType, dstShape0, offsetsShape1, srcColumnStartOffset, blockSize(dst, dstValidShape, src, srcShape1, offsets, blockTable, offsetsStartOffset, blockTableGMStride, blockTableStartOffset)即把逻辑→物理换算与 GM→L1 搬运全部下沉到 tile 内完成。tile 算子TGatherInL1的声明位于 cube_pto.h入口处通过CheckShapeValid做动态形状合法性检查。测试与验证路径Python 系统测试test_gather_in_l1.py参数化用例配置在 gather_in_l1_test_case.py包含小容量num_buffer_tokens33, token_dim30, block_size10, DT_FP16与大容量num_buffer_tokens82816, topk2048, block_size128, DT_INT8 → DT_INT32两类用例并在 NPU / SIM 两种pypto.RunMode下运行C 系统测试test_gather_in_l1.cpp覆盖isB / isTrans组合下的 matmul 联动验证并用页表 golden 函数逐元素比对代码生成单测test_codegen_dyn_copy.cpp 中包含GATHER_IN_L1的动态代码生成检查。小结pypto.experimental.gather_in_l1是 PyPTO 面向 PagedAttention KV Cache 场景定制的 GM→L1 离散行搬运算子以[1, n]的逻辑行偏移、[1, block_table_size]的页表和block_size页大小完成“逻辑 token → 物理行”换算把每行前size列数据直接以 L1 Tile 形式落地并通过is_b_matrix/is_trans让结果无缝对接pypto.matmul的 A/B 矩阵含转置消费。使用时应牢记它是两维 Tensor 定制接口、不保证稳定性参数形状与类型约束建议对照接口层校验见上文“约束说明”逐项核对。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考