人工智能大模型算子库AI 技能/插件【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址https://gitcode.com/cann/pypto-gym点击查看免费下载导读本文围绕 CANN PyPTO-Gym 仓库中 PyPTO 算子设计模式体系里的线性投影骨架 SK-03Linear ProjectionNorm→MatMul展开系统讲解以「RMSNorm 前置归一化 MatMul 线性投影」为核心的计算骨架的代码结构、V→C 排布原则、TileShape 切换要点、开箱性能优化配置以及面向「大归约维 × 极小输出宽」形态mhc_pre 类的专项变体设计。读完本文你将能够识别一个算子是否命中 SK-03 骨架按骨架模板写出可运行的 linear_projection 内核正确配置cube_l1_reuse_setting、vec_nbuffer_setting、set_cache_policy(NONE_CACHEABLE)等关键性能参数并在遇到输出宽 ≤ 64、归约维数万的极端形状时直接套用loop_unroll变体结构而非盲目扫描 tile 尺寸。SK-03 骨架定位与适用场景SK-03 是 PyPTO 计算骨架索引 中定义的 16 个计算骨架之一索引中将其标注为projection类骨架flow_pattern 为V1, C1。适用场景单阶段或多阶段线性投影通常带前置 RMSNorm。典型形态是「输入 → RMSNorm → MatMul →可选后处理→ 输出」即一次 V 阶段归一化加一次 C 阶段矩阵乘即可完成的算子。CV 排布V(RMSNorm) → C(MatMul) [→ V(后处理)]其中 V 表示 Vector 单元向量计算C 表示 Cube 单元矩阵乘计算。该排布也是 AT-11RMSNorm Linear Quant Fused原子模式的排布基础参见 AT-11-norm-quant-linear.md。展开因子候选为 128、64、32、16、8、1初始设计每次只选一个值并验证不能整除时的处理其余值分别作为调优候选。这一候选列表在仓库源码中直接体现IndexerPrologQuantConfig(unroll_list[128, 64, 32, 16, 8, 1])见 lightning_indexer_prolog_quant_v4_impl.py以及MlaPrologV4Configs(unroll_list[128, 64, 32, 16, 1], ...)见 mla_prolog_v4_impl.py。典型算子MLAProlog部分、Qwen3PreAttn。仓库中与 SK-03 直接对应的实现包括deepseek_v4目录下的 mla_prolog_v4_impl.pywqa-linear → q-rmsnorm 的 Norm→MatMul 段与 lightning_indexer_prolog_quant_v4_impl.pyQuery-Linear 段对应测试见 test_mla_prolog_v4.py 与 test_lightning_indexer_prolog_quant_v4.py。骨架结构标准模板代码SK-03 骨架的核心代码结构如下def linear_projection_kernel(input, weight, gamma, output, ...): gamma_fp32 pypto.cast(gamma, DT_FP32) weight_bf16 ... for idx in pypto.loop(tile_count): # Loop: Tile (or loop_unroll) x_tile pypto.view(input, [TILE, K], [offset, 0], valid_shape[tile_len, K]) # V: Pre-Norm (RMSNorm) pypto.set_vec_tile_shapes(v_tiles) normed rms_norm(x_tile, gamma_fp32, eps) normed_bf16 pypto.cast(normed, DT_BF16) # V: Quant (optional) # ... # C: Linear Projection pypto.set_cube_tile_shapes(c_tiles) projected pypto.matmul(normed_bf16, weight, dtype, b_transTrue) # V: Dequant / Post-Proc (optional) # ... pypto.assemble(cast(projected, output_dtype), [offset, 0], output)对照仓库中 mla_prolog_v4_impl.py 的实际写法可以看到 wqa-linear 段正是这一骨架的直接实现x_tile pypto.view(x, [t_tile, h], [tIdx, 0], valid_shape[t_tile, h]) pypto.set_semantic_label(wqa-linear) pypto.set_cube_tile_shapes([32, 32], [512, 512], [64, 64]) q pypto.matmul(x_tile, wq_a, pypto.DataType.DT_BF16) pypto.set_semantic_label(q-rmsnorm with weight) pypto.set_vec_tile_shapes(8, q_lora_rank) qr rms_norm(q, attrs.eps) qr pypto.mul(qr, gamma_cq_2d_fp32) qr pypto.cast(qr, pypto.DataType.DT_BF16) pypto.assemble(qr, [tIdx, 0], qr_out)注意这里实际是「先 C 后 V」MatMul 在前、RMSNorm 在后这体现了 SK-03 的灵活性骨架的 V→C 排布描述的是典型主形态当投影与归一化的顺序调换时仍可复用同一套结构原则单层 Loop、TileShape 切换、semantic label 分阶段。关键编码特征特征规则V→C 排布先 VecNorm再 CubeMatMul可选后接 Vec单层 Loop仅 batch/token 一维循环loop 内完成 VCTileShape 切换V 阶段和 C 阶段前各自set_vec/cube_tile_shapes后处理 V 阶段可选split、reshape、per-head norm、RoPE非必须量化变体V→C 可扩展为 V→Quant→C→DequantAT-11典型算子MLAProlog部分、Qwen3PreAttn各特征的仓库佐证单层 LoopSK-03 只在 batch/token 维度上循环。mla_prolog_v4_compute使用pypto.loop_unroll(0, t, 1, nameMLA_BS_LOOP, ...)单层循环完成全部 wqa/wqb/wkv 三个投影见 mla_prolog_v4_impl.pyquant_lightning_indexer_prolog_kernel同样以pypto.loop_unroll(0, t, 1, nameIndexerPrologLoop, ...)单层循环完成 Query 与 Weight 两路计算见 lightning_indexer_prolog_quant_v4_impl.py。TileShape 切换在每次 V/C 阶段切换前分别调用pypto.set_vec_tile_shapes(...)与pypto.set_cube_tile_shapes(...)。例如 lightning_indexer_prolog_quant_v4_impl.py 中 matmul 前设置 cube tile[128, 128], [256, 1024], [256, 256]随后 cast/dequant 前切换为 vec tile(1, idx_nq * head_dim)。后处理 V 阶段mla_prolog_v4在 wqb-linear 的 matmul 之后继续做 RMSNorm、RoPE、concat 等后处理见 mla_prolog_v4_impl.py印证后处理 V 阶段可选的扩展方式。量化变体lightning_indexer_prolog_quant_v4的 Query 路径即 V→C→V 的量化投影链qrINT8经 cube matmul 得 INT32cast 到 FP32 后乘 scale 完成 dequant见 lightning_indexer_prolog_quant_v4_impl.py与 AT-11 描述的 Norm→Quant→MatMul→Dequant 流水一致。开箱性能优化提示实证来源lightning_indexer_prolog_quant_v4_impl.py、mla_prolog_v4_impl.py维度推荐配置取值经验作用pass_options.cube_l1_reuse_setting必配{-1: 2, 1: 1}分轴权重轴用 1不复用激活轴用 2双缓冲匹配权重静态、激活动态特性pass_options.vec_nbuffer_setting推荐{0: 2}最 Norm 阶段RMSNorm 在 V 阶段nbuffer2 即可pypto.set_cache_policy(NONE_CACHEABLE, True)条件配静态权重张量仅当权重在 loop 内被单次消费权重只读一次时不需占用 L2避免与激活竞争 cache⚠️ 若权重跨 loop 迭代复用如 mhc_pre 变体的 phi 被 8 次 unroll 复用且可驻留 L2标记 NONE_CACHEABLE 会强制每迭代回 HBM 重读反而劣化——勿用token 循环展开按需每次选一个因子128 仅为候选比较编译成本和运行耗时验证余数处理pypto.set_semantic_label(...)推荐每阶段一个标签帮助编译器识别阶段边界便于 Pass 调度pypto.reshape(..., inplaceTrue)推荐tile 内 reshape避免临时张量分配TileShape推荐V 阶段set_vec_tile_shapes(4, hidden)C 阶段set_cube_tile_shapes([M,K],[K,N],[M,N])V/C 切换前各设一次配置项源码级展开cube_l1_reuse_setting。以lightning_indexer_prolog_quant_v4为例JIT 装饰器中的pass_options直接写入{cube_l1_reuse_setting: {-1: 2, 1: 1}, vec_nbuffer_setting: {0: 2}}见 lightning_indexer_prolog_quant_v4_impl.py。其中 key-1对应激活轴双缓冲、复用 L1key1对应权重轴不复用与权重静态、激活动态的访存特征匹配——激活每迭代都变化值得双缓冲隐藏 L1 加载延迟权重每迭代内容相同复用得越多反而挤占 L1 容量。set_cache_policy(NONE_CACHEABLE)。同一实现中idx_wq_b、weights_proj、hadamard三个静态权重张量均在进入 loop 前调用set_cache_policy(pypto.CachePolicy.NONE_CACHEABLE, True)见 lightning_indexer_prolog_quant_v4_impl.py。这些权重在 loop 内单次消费INT8 量化场景matmul 输入是动态量化后的激活标记 NONE_CACHEABLE 后不再占用 L2把 cache 让给激活数据流从而恢复激活吞吐。这是 SK-03 骨架权重 cache 抢占激活这一主要性能瓶颈的对应解法。vec_nbuffer_setting。{0: 2}作用于最靠近 Norm 的 V 阶段RMSNorm 涉及平方、求和、开方等多步链式操作nbuffer2 提供双缓冲流水即可不需要更高的多缓冲。semantic label 与 reshape inplace。仓库中每个计算阶段前都设置标签如Query-Linear、Query-Dequant、Hadamard-Compute、Weight-Compute见 lightning_indexer_prolog_quant_v4_impl.py帮助编译器识别阶段边界。w_qb_scale pypto.reshape(idx_wq_b_scale, [1, idx_nq * head_dim], inplaceTrue)以 inplace 方式 reshape源码注释明确指出Reshape inplace will not generate data move见 lightning_indexer_prolog_quant_v4_impl.py即 inplace reshape 不产生数据搬移避免临时张量分配。该骨架特有的性能方向权重 NONE_CACHEABLE loop_unroll 自适应。瓶颈通常在权重 cache 抢占激活——用set_cache_policy(NONE_CACHEABLE)标记权重就能恢复激活吞吐。量化变体AT-11在源码中的体现SK-03 的量化变体 V→Quant→C→Dequant 在 common.py 中具备完整的原子函数支撑quant对称/非对称 INT8 量化。对称模式按scale max(|x|)/127.0计算经CAST_RINT→CAST_ROUND→CAST_TRUNC SaturationMode.ON多级 cast 落到 INT8并返回 dequant 用 scale见 common.pyquant_tensorper-token 动态量化同样输出(int8, dequant_scale)见 common.py。这与 AT-11 定义的Norm → Quant → MatMul → Dequant完整流水一一对应可作为骨架扩展 Quant 阶段的直接参考实现。⚠️ 形态变体大归约维 ND × 极小输出宽 Kmhc_pre 类当算子为「单 matmul norm」但形状落在输出宽 ≤ 64、归约维 N·D ≥ 数万、FP32 计算典型mhc_pre的matmul [B,28672]×[28672,24]时用下列变体结构替代上方开箱提示维度变体配置原理循环loop_unroll(0, BS, 1, unroll_list[16])Munroll_length平铺 BT-loop 顾此失彼BT 大则 cube 只有一个任务无并行BT 小则 vec 归约被切碎。loop_unroll 让 vec 连续处理 16 行整 Dcube 按 M16 出多个任务——两侧并行同时成立vec tileD 轴大 tile(8,2048)/(1,2048)/(1,N,2048)计算集中在尾轴单任务连续扫得越长DMA/计算比越好小行 tile 会把长归约切碎权重布局wrapper 预转置phi.T.contiguous()matmul 不加b_trans转置 host 侧只做一次b_trans是每个 cube 任务重复付跨步寻址代价cube[16,16],[512,1024],[128,128]enable_split_kTrue输出极小cube 并行只能切归约维分核算部分和再归并M16 与 unroll_length 对齐权重 cache不设 NONE_CACHEABLEphi 被各 unroll 迭代反复复用且可驻留 L2禁缓存等于强制每轮回 HBM 重读装饰器stitch_function_max_num128、device_sched_mode2、cube_nbuffer{-1:4}、vec_nbuffer{DEFAULT:4,...}、sg_set_scope分段、combine_axisTruestitch 把多次 unroll 迭代缝成大图做流水调度动态调度只在图内任务足够多时优于静态分配UB 约束勿提前 cast 大中间量FP32 全 tile UB 会 spillx 的 FP32 形态是 BF16 两倍跨 matmul 驻留 UB 装不下用时重读 BF16 即时 cast 反而更快识别条件matmul_count1 AND 输出宽 ≤ 64 AND 归约维 N·D ≥ 2^14 AND FP32 计算命中即直接用变体跳过 BT sweep。警示变体各配置的有效性都依赖 loop_unroll 这个结构前提拆开单项套到平铺 BT-loop 上只会劣化——要么整体用要么不用。变体在源码中的对应实现仓库 hc_pre_impl.py 正是这一变体的完整落点对应测试 test_hc_pre.pyloop_unroll 结构for t_idx, unrollLength in pypto.loop_unroll(0, t, 1, namet_loop, idx_namet_idx, unroll_listunroll_list)decode 形态 unroll_list 为[256, 64, 16, 4, 1]见 hc_pre_impl.py。cube 与 split_k 的形态自适应代码按tile_t分三档设置 cube tile——tile_t ≤ 32时[16,16],[512,1024],[128,128]且enable_split_kFalsetile_t ≤ 64时启用 split_k更大时[16,16],[512,2*1024],[128,128]且enable_split_kTrue见 hc_pre_impl.py。这正是输出极小24 列、归约维 28672 巨大场景下 split_k 按需开启的工程实现。权重预转置hc_fn权重以(mix_hc24, hc*d16384)的布局参与pypto.matmul(x_fp32, hc_fn, pypto.DT_FP32, b_transTrue)通过 host 侧一次性预置为转置友好的布局避免每 cube 任务重复跨步寻址。sg_set_scope 分段与 FP32 中间量控制代码用pypto.set_pass_options(sg_set_scope1)/sg_set_scope-1将 x 的 FP32 cast 与 RMSNorm 归入同一 scope避免大 FP32 中间量跨阶段驻留 UB见 hc_pre_impl.py与变体表的UB 约束条目直接对应。stitch 与调度pypto.frontend.jit(runtime_options{stitch_function_max_num: 128, device_sched_mode: 0})见 hc_pre_impl.pystitch 把多次 unroll 迭代缝成一个大图做整体流水调度。何时选择 SK-03 以及如何验证在实际算子设计流程中可通过以下步骤确认是否命中 SK-03 骨架检查计算流算子主体是否为「可选前置 Norm 一次核心 MatMul 可选后处理」的单阶段线性投影若是多 matmul 串行应优先参考 SK-04-multi-stage-prolog.md 或 SK-05-fused-pre-attention.md。检查形状极端性若命中matmul_count1 AND 输出宽 ≤ 64 AND 归约维 ≥ 2^14 AND FP32直接采用 mhc_pre 变体结构跳过 BT sweep。按展开因子候选调优从[128, 64, 32, 16, 8, 1]中每次选一个值编译运行验证不可整除时的余数处理逻辑再比较编译成本与运行耗时确定最优因子。回归验证参照仓库中的测试用例如 test_mla_prolog_v4.py、test_lightning_indexer_prolog_quant_v4.py、test_hc_pre.py做精度与性能双向验证。小结SK-03 线性投影骨架以「单层 token/batch 循环 V/C 阶段切换 按需展开」为核心覆盖了从标准 Norm→MatMul 到量化投影AT-11再到极端形状mhc_pre 变体的完整形态谱系。其性能调优的关键在于两点一是让权重不抢占激活的 cacheNONE_CACHEABLE 与 cube_l1_reuse_setting 分轴配置二是用 loop_unroll 适配大归约维 × 小输出的并行困境。仓库中deepseek_v4系列实现与对应测试为骨架的每个条目提供了可直接对照的工程样例是学习与实践 SK-03 的最佳参照。赞分享人工智能大模型算子库AI 技能/插件【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址https://gitcode.com/cann/pypto-gym点击查看免费下载相关推荐PyPTO-Gym 算子设计骨架 SK-06FFN/SwiGLU 双 MatMul 激活融合的编写与调优指南PyPTO Gym 算子设计骨架 SK 06FFN/SwiGLU 双 MatMul 激活融合的编写与调优指南 PyPTO Gym 的算子设计模式库将前馈网络人工智能大模型算子库AI 技能/插件pypto-gym 中的 SK-14 通用多矩阵乘骨架PyPTO 多 MatMul 算子设计的兜底范式pypto gym 中的 SK 14 通用多矩阵乘骨架PyPTO 多 MatMul 算子设计的兜底范式 SK 14General Multi MatMul人工智能大模型算子库AI 技能/插件PyPTO-Gym SK-13 通用纯向量兜底骨架无 MatMul 算子的设计范式、变体选型与性能配置PyPTO Gym SK 13 通用纯向量兜底骨架无 MatMul 算子的设计范式、变体选型与性能配置 本篇技术指南围绕 PyPTO Gym 算子设计模式库中人工智能大模型算子库AI 技能/插件创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
