Qwen3-1.7B 昇腾 NPU 迁移实战：基于 PyPTO 的 Q/K RMSNorm+RoPE 算子融合指南

Qwen3-1.7B 昇腾 NPU 迁移实战基于 PyPTO 的 Q/K RMSNormRoPE 算子融合指南【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym导读本文基于 CANN / pypto-gym 仓库中 Qwen3-1.7B 迁移说明完整还原 Qwen3-1.7B 从 HuggingFace 权重下载、transformers 4.51.0 原生代码适配、PyPTO 算子入网到单次问答推理的全流程。文章不仅保留了原始文档中的环境版本、性能数据与运行命令还深入剖析了Qwen3Attention.forward()中 6 行 PyPTO dispatch 的注入方式、融合内核rrms_norm_rope_impl.py的 tile 计算流程与 golden 测试验证方法。读完本文你将掌握如何把一款开源 LLM 以最小改动 部分算子融合的方式迁移到 Ascend NPU并在 prefill 与 decode 两个阶段分别选择 PyPTO 内核与原生路径。迁移目标与最小化原则Qwen3-1.7B 是 Qwen 系列中面向轻量级部署的因果语言模型。本次迁移的核心思路是最小化改动字段说明HuggingFaceQwen/Qwen3-1.7B权重目录{model_weight_dir}如/path/to/models/Qwen3-1.7B代码来源transformers 4.51.0 内置/models/qwen3/import 已修复为from transformers.xxxtransformers 版本4.51.0代码位置modeling_qwen3.py, configuration_qwen3.py根目录auto_map 指向根目录修改内容最小化仅Qwen3Attention.forward()增加 6 行 PyPTO dispatchprefill 用 PTOdecode 退原生运行命令python3 scripts/ask_Qwen3-1.7B.py --device 0 --prompt 你好仓库内已归档的适配产物位于 src/pypto_gym/transformers/qwen3_1_7b/其中 modeling_qwen3.py 与 configuration_qwen3.py 放在权重目录根级并在 config.json 中通过auto_map指向根目录从而支持trust_remote_codeTrue直接加载。环境信息与前置条件原始文档给出了经过验证的软件栈版本迁移到其他环境时需要保证核心组件对齐组件版本torch2.9.0torch_npu2.9.0.post2torchvision0.24.0transformers4.51.0CANN9.0.0NPUAscend 910B2accelerate1.10.1safetensors0.6.2tokenizers0.21.4numpy2.2.6其中torch_npu提供torch.npu.set_device等 NPU 设备接口CANN 9.0.0提供算子运行所需的底层昇腾运行时。推理脚本通过device_map{: fnpu:{args.device}}将模型整体放置到指定 NPU 卡上。模型下载使用仓库内置的 download_hf_model.py 从 HuggingFace 拉取权重python3 cannbot-skills/model/pypto-fused-op-integration/scripts/download_hf_model.py \ --model-id Qwen/Qwen3-1.7B \ --output-dir {model_weight_dir}下载完成后权重目录应包含model-*.safetensors分片权重与tokenizer相关文件分词器配置、词表等。PyPTO 入网适配补丁还原脚本下载的原始权重目录里是 transformers 官方实现尚未包含 PyPTO 融合算子。需要执行 restore_model_patch.sh 完成入网适配将 HF 原始模型替换为华为修改版bash cannbot-skills/model/pypto-fused-op-integration/scripts/restore_model_patch.sh \ {model_weight_dir} qwen3该脚本以model_weight_dir与model_name此处为qwen3两个参数运行执行四个步骤替换 modeling/configuration 代码将 src/pypto_gym/transformers/qwen3_1_7b/ 下的modeling_qwen3.py、configuration_qwen3.py拷贝到权重目录根级若存在原文件会先备份为*.hf_orig保证可回退写入 pto_kernels 融合算子模块将 src/pypto_gym/ops/pypto_tensor/qwen3_1_7b/ 完整拷贝为权重目录下的qwen3_pto_kernels/确保 config.json 含auto_map写入AutoConfig/AutoModelForCausalLM指向根目录的configuration_qwen3.Qwen3Config与modeling_qwen3.Qwen3ForCausalLM参见 config.json 中auto_map字段这是trust_remote_code加载的前提清除 HF 模块缓存清理~/.cache/huggingface/modules/transformers_modules下与模型名匹配的缓存目录避免加载旧版远端代码。说明脚本在第 2 步使用rm -rf重建目标pto_kernels目录这是脚本自身的正常清理逻辑运行前请确保{model_weight_dir}路径正确。运行命令基线推理与 PyPTO 融合推理推理入口是 ask_Qwen3-1.7B.py支持--device、--prompt、--model-path、--use-pto四个参数# 基线推理原生 PyTorch python3 scripts/ask_Qwen3-1.7B.py --device 0 --prompt 你好 # PyPTO 融合推理 python3 scripts/ask_Qwen3-1.7B.py --device 0 --prompt 你好 --use-pto两条命令仅差一个--use-pto开关。脚本内部的关键逻辑见 ask_Qwen3-1.7B.py# PyPTO injection (must happen before transformers import) if args.use_pto: sys.path.insert(0, args.model_path) import qwen3_pto_kernels sys.modules[qwen3_pto_kernels] qwen3_pto_kernels qwen3_pto_kernels.USE_PTO_ROPE True logger.info(PyPTO RoPE mode enabled)需要注意两个设计细节注入时机必须在import transformers之前。modeling_qwen3.py在运行时通过sys.modules.get(qwen3_pto_kernels)查找模块因此必须先把权重目录加入sys.path并注册到sys.modules模型代码才能看得到融合算子包开关变量USE_PTO_ROPE在推理脚本中显式置为True而算子包init.py 中的默认值为False两者配合实现不开--use-pto则完全等价于原生 PyTorch的可回退语义。模型加载与生成部分使用 float16 精度model AutoModelForCausalLM.from_pretrained( args.model_path, torch_dtypetorch.float16, device_map{: fnpu:{args.device}}, local_files_onlyTrue, trust_remote_codeTrue ) outputs model.generate(**inputs, max_new_tokens512, temperature0.7, do_sampleTrue)性能对比与开销归因原始文档在 Ascend 910B2 上给出了单条 prompt 的实测数据模式命令模型加载推理耗时吞吐峰值显存baselinepython3 scripts/ask_Qwen3-1.7B.py --device 0 --prompt 你好2.5s7.2s13.9 tok/s4470 MBptopython3 scripts/ask_Qwen3-1.7B.py --device 0 --prompt 你好 --use-pto2.4s8.3s12.0 tok/s4470 MB单算子替换时 PTO 比基线慢 ~15%JIT 首编 kernel launch 开销收益来自多算子融合。这组数据揭示了一个重要的工程认知单个算子替换本身不会带来收益。PyPTO 的 JIT 首编首次运行时编译生成内核与 kernel launch内核启动开销会摊薄甚至超过算子本身的节省--use-pto模式下多出的约 1.1s 正是这两部分开销。PyPTO 的价值体现在把多个小算子融合进一个内核——例如本文的 Q/K RMSNorm RoPE 融合——从而减少多次 kernel launch 与中间张量落盘。峰值显存两侧一致4470 MB说明融合并未引入额外显存占用。目录结构权重目录Qwen3-1.7B/经过补丁还原后形成如下结构Qwen3-1.7B/ ├── modeling_qwen3.py # 模型实现4.51.0内置 6行PTO dispatch ├── configuration_qwen3.py # 配置实现 ├── config.json # 模型配置auto_map指向根目录 ├── scripts/ │ ├── ask_Qwen3-1.7B.py # 推理脚本支持--use-pto │ └── README.md # 本文档 ├── qwen3_pto_kernels/ # PyPTO融合算子模块 │ ├── __init__.py # USE_PTO_ROPE开关 qk_rope_wrapper │ └── rope/ │ ├── __init__.py │ ├── README.md │ └── rrms_norm_rope_impl.py # Q/K RMSNormRoPE融合kernel ├── model-*.safetensors # 模型权重 └── tokenizer相关文件 # 分词器PyPTO 融合范围Qwen3-1.7B 的整体算子分布中本次迁移只对注意力前端的两个高频小算子做了融合算子融合?说明Q/K RMSNorm RoPE✅prefill (S1) 走 PyPTO kerneldecode (S1) 退原生Q/K/V projection❌PyTorch LinearAttention❌eager/sdpa 原生MLP❌PyTorch Linear整体 RMSNorm❌PyTorch Qwen3RMSNorm选择 Q/K RMSNorm RoPE 作为首个融合目标并非偶然该位置每层都要执行两次 RMSNormq_norm/k_norm与两次 RoPE 旋转涉及[B, S, num_heads, head_dim]级别的张量搬运是 prefill 阶段可观的访存与启动开销来源。而 Attention 本身在 910B2 上已有成熟的 eager/sdpa 原生实现Q/K/V projection 与 MLP 属于大 GEMM单算子替换收益有限因此均保持原生。融合内核解析rrms_norm_rope_impl.py融合内核位于 rrms_norm_rope_impl.py由_make_qk_rope_kernel(num_heads)工厂函数分别实例化 Q/K 两个专用内核qwen3_qk_rope_q _make_qk_rope_kernel(16) # Q: N_q16 heads qwen3_qk_rope_k _make_qk_rope_kernel(8) # K: N_kv8 heads这与 config.json 中的多头配置严格对应num_attention_heads16、num_key_value_heads8、head_dim128。内核按固定D128、EPS1e-6、BS_TILE8编译通过pypto.frontend.jit声明动态序列维度pypto.DYNAMIC并指定运行时选项pypto.frontend.jit( runtime_options{stitch_function_max_num: 128, device_sched_mode: 1}, ) def kernel( x: pypto.Tensor([pypto.DYNAMIC, num_heads, D], pypto.DT_BF16), cos: pypto.Tensor([pypto.DYNAMIC, D], pypto.DT_BF16), sin: pypto.Tensor([pypto.DYNAMIC, D], pypto.DT_BF16), w_norm: pypto.Tensor([D], pypto.DT_BF16), out: pypto.Tensor([pypto.DYNAMIC, num_heads, D], pypto.DT_BF16), ):内核主体是一个按BS_TILE8步长遍历 sequence 维度的 tile 循环LOOP_BS_QKROPE每个 tile 内部依次完成五步取 tile 输入pypto.view(x, [BS_TILE, num_heads, D], [bs_idx*BS_TILE, 0, 0], valid_shape[cur_bs, num_heads, D])用valid_shape处理末尾不足 8 的尾块RMSNorm先pypto.cast到 FP32执行mul → sum(keepdim) → rsqrt(add eps) → mul最后乘上广播后的w_norm权重_rms_norm_per_d函数准备 cos/sin同样按 tile 切出cos_tile/sin_tile再取前半维HALF_D64因为 RoPE 旋转只需要一半维度的三角函数值RoPE 计算把归一化结果拆成左半x_left与右半x_right执行o1 x_left*cos - x_right*sin、o2 x_right*cos x_left*sin再pypto.concat([o1, o2], 2)拼回完整维度写回pypto.assemble(roped_bf, [bs_idx*BS_TILE, 0, 0], out)。全程 BF16 输入、FP32 中间计算、BF16 输出的精度策略配合pypto.set_vec_tile_shapes逐阶段切换向量 tile 形状保证计算符合昇腾 UBUnified Buffer容量约束。Dispatch 机制modeling_qwen3.py 中的 6 行注入融合内核如何被模型调用答案在 modeling_qwen3.py 的Qwen3Attention.forward()中这是整个迁移唯一的模型代码改动点# PyPTO fused Q/K RMSNorm RoPE (prefill only; decode falls back to eager) import sys pto_kernels sys.modules.get(qwen3_pto_kernels) if pto_kernels is not None and pto_kernels.USE_PTO_ROPE and hidden_states.shape[1] 1: query_states, key_states pto_kernels.qk_rope_wrapper( self.q_proj(hidden_states), self.k_proj(hidden_states), cos, sin, self.q_norm.weight, self.k_norm.weight, self.config.num_attention_heads, self.config.num_key_value_heads, self.head_dim ) value_states self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) else: query_states self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) key_states self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) value_states self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) query_states, key_states apply_rotary_pos_emb(query_states, key_states, cos, sin)dispatch 条件包含三重保护缺一不可pto_kernels is not None未执行--use-pto时模块不存在走原生路径pto_kernels.USE_PTO_ROPE开关变量显式开启hidden_states.shape[1] 1序列长度大于 1 才算 prefillS1 的 decode 阶段自动退回原生。原因在于 decode 每次只处理一个 token融合内核的 JIT 与启动开销占比过高且原生路径对单 token 场景已有足够好的表现。qk_rope_wrapper的封装逻辑见 qwen3_1_7b/init.py它把[B, S, q_num_heads*head_dim]的投影输出展平为[-1, num_heads, head_dim]三维张量等价于把 B×S 拼成内核要求的动态序列维cos/sin 与 norm 权重统一转成 BF16分别调用 Q/K 内核后转回原 dtype再transpose(1, 2).contiguous()恢复成注意力期望的[B, num_heads, S, head_dim]布局。开关变量qwen3_pto_kernels.USE_PTO_ROPEbool默认 False值行为False默认走原生 PyTorchq_norm/k_normapply_rotary_pos_embTrueQ/K RMSNorm RoPE 经 PyPTO 内核融合执行prefill 阶段该开关同时承担回归对比职责不开开关即可获得与基线完全一致的原生行为方便在相同 prompt 下对拍数值与性能。正确性验证golden 参考实现与测试仓库在 tests/ops/qwen3_1_7b/ 提供了融合算子的 golden 验证链路rms_norm_rope_golden.py纯 PyTorch 参考实现。rms_norm_torch用x * torch.rsqrt(variance eps) * weight复现 RMSNormrope_torch复现左半/右半的旋转拼接rms_norm_rope_golden将两者串联且 cos/sin 同样只取前半维——与 PyPTO 内核的数学定义逐一对齐test_rms_norm_rope.py读取 test_cases.json 中定义的多个 shape/dtype 用例float16/float32/bfloat16断言 golden 输出在形状、dtype、无 NaN/Inf、数值量级max_abs 100四个维度全部合法可直接运行python3 tests/ops/qwen3_1_7b/test_rms_norm_rope.py做自检。该 golden 实现与内核共享同一套[S, N, D]约定与eps1e-6可作为后续把内核输出与 PyTorch 结果做逐元素对比的数值基准。仓库还提供了modeling/transformers/bench_qwen3_1_7b.sh等批量基准脚本用于在更长 prompt、多 batch 场景下评估融合收益。归档映射从权重目录到仓库为便于复现与维护原始文档给出了权重目录与 pypto-gym 仓库之间的归档对应关系来源 ({model_weight_dir}/)目标 (pypto-gym/)scripts/modeling/transformers/qwen3_1_7b/config.json,modeling_qwen3.py,configuration_qwen3.pysrc/pypto_gym/transformers/qwen3_1_7b/qwen3_pto_kernels/src/pypto_gym/ops/pypto_tensor/qwen3_1_7b/这条映射同时解释了 restore_model_patch.sh 的反向行为它正是把仓库内这三个位置的内容重新铺回权重目录形成可加载的模型包。迁移总结与注意事项改动面可控模型代码仅Qwen3Attention.forward()一处、约 6 行 dispatch其余全部继承 transformers 4.51.0 官方实现方便随上游升级阶段感知调度prefillS1与 decodeS1分别走 PyPTO 与原生路径避免单 token 场景的 JIT/启动开销可回退设计USE_PTO_ROPE默认 False *.hf_orig备份保证基线对拍与快速回滚预期管理单算子融合阶段 PTO 比基线慢约 15%真实收益依赖多算子融合把多次 launch 合并为一次峰值显存不变4470 MB说明融合不增加显存开销运行前提需 CANN 9.0.0 torch_npu 2.9.0.post2 Ascend 910B2 环境模型加载使用local_files_onlyTrue因此务必先完成权重下载与补丁还原。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考