PUMADOMINO昇腾 Atlas 910 在线推理实战基于 PyTorch 的 policy server 部署与优化解析【免费下载链接】cann-recipes-embodied-ai本项目针对具身智能业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-ai本文是基于 manipulation/puma/infer_with_torch/README.md 编写的 PUMA 昇腾在线推理技术指南。PUMA 是 DOMINOECCV 2026提出的动态操作 VLAVision-Language-Action模型以 Qwen3-VL-4B 为骨干结合场景级历史光流与世界查询world queries对物体未来状态做短时预测从而在动态环境中更及时地做出反应DOMINO 同时提供 35 个层次化动态操作任务、超过 110K 条专家轨迹的数据集以及基于 RoboTwin 仿真的多维评测套件参见 manipulation/puma/README.md。本文将完整介绍如何把已有 PUMA checkpoint 原样加载到昇腾 Atlas 910 上、以 policy server 形式对 RoboTwin 仿真评测端提供动作推理服务并深入解析昇腾侧适配的五个优化点FlashAttention→SDPA 覆盖、Conv3d patch embed 线性化、Qwen3 推理请求计划、首请求预热与依赖治理帮助读者独立完成环境搭建、服务启动、RoboTwin 联调与稳态时延统计的全流程。1. 样例定位与总体设计原则本目录提供的在线推理样例遵循四个核心设计原则理解它们有助于后续排查问题权重不做转换上游发布的 PUMA checkpoint 在昇腾 NPU 上原样加载不修改任何 checkpoint key所有结构性替换均在加载时完成且数值等价详见 docs/manipulation/puma/infer_with_torch/README.md。后端中立通用代码VLM 封装、部署入口不出现if npu分支昇腾专用逻辑集中在PUMA/model/modules/vlm/ascend/包内通过--device npu一个开关注入。启动即校验环境不满足要求缺 torch_npu、NPU 不可用时服务在启动阶段就会报错退出不会带着错误配置继续服务。上游行为不变所有适配仅在 NPU 上激活上游原有推理路径保持原状。由于昇腾适配代码已合入 DOMINO 上游主干本样例只保存脚本与文档不携带任何补丁。服务端只需加--device npu仿真评测端完全不感知设备差异。2. 适用场景与已验证软件栈项目要求硬件昇腾 Atlas 910 系列CANN8.5.2任务PUMA policy server 在线推理 RoboTwin/DOMINO 仿真评测外部代码仓H-EmbodVis/DOMINO固定 commit见下文模型权重H-EmbodVis/PUMA 预训练 checkpoint已验证软件栈组件版本NPUAtlas 910 系列CANN8.5.2torch / torch-npu2.5.1 / 2.5.1.post1transformers4.57.0Python3.10相邻的 CANN / torch-npu 版本组合大概率可用但只有上述软件栈经过端到端验证。请务必确保 CANN 版本与torch-npu构建的要求匹配参见 torch-npu 版本配套表。3. 目录结构与外部依赖固定本样例不内嵌 DOMINO 源码而是将其作为外部依赖单独 clone 并固定到已验证 commit9c94f2d3a700fff3b65f041df4038d497139ed1f该 commit 已包含全部昇腾训练与推理适配。样例目录结构如下manipulation/puma/infer_with_torch/ ├── README.md └── src/ └── scripts/ ├── setup.sh # 复用训练 setup 并追加 RoboTwin 评测通信依赖 └── run_server.sh # policy server 启动封装--device npu推理侧昇腾适配与优化文档位于 docs/manipulation/puma/infer_with_torch/README.md训练样例见 manipulation/puma/train/README.md。3.1 setup.sh 的实现细节从 setup.sh 源码可以看到推理 setup 本身是一个极薄封装它定位到manipulation/puma/train/src/scripts/setup.sh校验其存在后以--with-eval-deps参数转发执行exec bash $TRAIN_SETUP --with-eval-deps $因此推理与训练共享同一个固定的 DOMINO commit 和昇腾运行时仅在其上追加 RoboTwin 评测通信协议依赖examples/Robotwin/eval_files/requirements.txt与上游一致。训练侧 setupmanipulation/puma/train/src/scripts/setup.sh会完成在cann-recipes-embodied-ai同级目录workspace/准备DOMINO/代码仓可通过DOMINO_REPO_URL指定镜像地址或提前手动放置已 checkout 的DOMINO/目录供其复用checkout 到固定 commit9c94f2d3a700fff3b65f041df4038d497139ed1f先安装requirements-ascend.txt固定 torch2.5.1 / torch-npu2.5.1.post1避免后续依赖解析拉取 GPU 版 torch以pip install --no-build-isolation -e .可编辑方式安装 PUMA最终校验import torch, torch_npu, transformers, deepspeed是否成功失败则给出排障提示并退出。其支持的参数包括--create-conda、--env-name默认puma-ascend、--python-version默认3.10、--with-eval-deps、--skip-torch-check。4. 环境准备4.1 clone 代码git clone https://gitcode.com/cann/cann-recipes-embodied-ai.git cd cann-recipes-embodied-ai4.2 准备 DOMINO 与运行时# 先加载 CANN 工具链环境路径按实际安装位置调整 source /usr/local/Ascend/ascend-toolkit/set_env.sh chmod x manipulation/puma/infer_with_torch/src/scripts/setup.sh ./manipulation/puma/infer_with_torch/src/scripts/setup.sh --create-conda推荐工作区布局workspace/ ├── cann-recipes-embodied-ai/ └── DOMINO/ └── policy/PUMA/ └── playground/Pretrained_models/ # 模型权重放置目录4.3 注意事项依赖治理红线不要安装flash-attn、decord、eva-decord它们会拉取 GPU 专用依赖并破坏环境requirements-ascend.txt已刻意排除昇腾侧使用sdpa注意力与torchvision_av视频解码后端。保持numpy1.26.4后续安装supervision、opencv-python等包可能将 NumPy 静默升级到 2.x导致 torch-npu 运行时崩溃如发生请重新固定。权重准备方式与上游原有流程完全一致按 DOMINO 上游权重下载说明下载 PUMA checkpoint 后放置或软链到DOMINO/policy/PUMA/playground/Pretrained_models下。5. 启动 policy server5.1 使用封装脚本./manipulation/puma/infer_with_torch/src/scripts/run_server.sh \ --ckpt /absolute/path/to/checkpoints/steps_100000_pytorch_model.pt \ --port 9001 \ --npu 0从 run_server.sh 源码可知其支持的参数与环境变量参数 / 变量默认值含义--ckpt PATH必填PUMA checkpoint 绝对路径如steps_100000_pytorch_model.pt文件不存在会直接报错退出--port PORT9001policy server 监听端口--npu ID0指定使用的 NPU 卡 id内部转换为ASCEND_RT_VISIBLE_DEVICES-- ...—--之后的参数原样透传给server_policy.pyDOMINO_ROOTworkspace/DOMINO覆盖 DOMINO 仓库根目录适用于已有 checkout 的场景ASCEND_SET_ENV/usr/local/Ascend/ascend-toolkit/set_env.shCANNset_env.sh路径文件不存在时仅告警并假定环境已加载脚本执行流程为解析参数 → 校验 checkpoint 与PUMA_ROOT存在 → 按需 source CANN 环境 → 设置PYTHONPATH与ASCEND_RT_VISIBLE_DEVICES→ 以exec方式调用上游部署入口exec python deployment/model_server/server_policy.py \ --ckpt_path $CKPT_PATH \ --port $PORT \ --device npu \ --use_bf16 \ ${EXTRA_ARGS[]}5.2 等价的直接调用方式与上游文档一致cd ../DOMINO/policy/PUMA ASCEND_RT_VISIBLE_DEVICES0 python deployment/model_server/server_policy.py \ --ckpt_path /absolute/path/to/checkpoints/steps_100000_pytorch_model.pt \ --port 9001 \ --device npu \ --use_bf16说明ASCEND_RT_VISIBLE_DEVICES用于指定服务使用的 NPU 卡首个请求会明显慢于稳态请求设备与内存池初始化、权重首次搬运、torch_npu 首次算子下发等一次性开销建议压测或评测前先发一个预热请求再统计稳态时延服务端加载权重时自动应用昇腾配置覆盖FlashAttention→SDPA、Conv3d patch embed 线性化等细节见 docs/manipulation/puma/infer_with_torch/README.md。6. 昇腾推理适配与优化要点以下内容来自 docs/manipulation/puma/infer_with_torch/README.md是该样例的技术内核。全部适配代码已合入 DOMINO 上游仓库文中路径均相对DOMINO/policy/PUMA/。6.1 加载期配置覆盖FlashAttention → SDPA、bf16动机PUMA 上游配方默认使用 FlashAttention该库为 GPU 专用昇腾上不可用也不应安装。实现服务端检测到--device npu后通过ascend_inference_config_overrides()对 checkpoint 携带的配置做内存内覆盖attn_implementationsdpa、model_dtypebfloat16、linearize_vision_patch_embedtrue、enable_ascend_inference_adaptertrue。checkpoint 文件本身不做任何修改。收益同一份 checkpoint 无需转换即可在昇腾上服务SDPA 走 torch_npu 的融合注意力实现bf16 与训练精度口径一致。6.2 Conv3d 视觉 patch embed 线性化动机Qwen3-VL 视觉塔的 patch embed 是kernel_size stride的 Conv3d。该形态的 Conv3d 在 NPU 上不支持/性能差且是视觉编码的入口热点。实现ascend/patch_embed.py在加载时将 Conv3d 投影替换为数学等价的F.linearLinearizedConv3dPatchEmbed当 kernel 与 stride 相等时Conv3d 严格等价于对展平 patch 的一次线性投影权重直接 reshape 复用无需重训或转换。替换带守卫仅当proj确为 Conv3d 且 kernelstride失败时报错而非静默降级。收益绕开 NPU 上的 Conv3d 劣化路径改走高度优化的矩阵乘数值与原 Conv3d 完全等价。6.3 Qwen3 推理请求计划CPU 预计算控制流动机Qwen3-VL 的视觉前处理包含大量小规模、数据依赖的控制流计算position embedding 双线性插值索引、cu_seqlens、rotary 位置 id、按 grid 切分长度等。这些计算在 NPU 上会带来形态多变的小算子下发与频繁的 host-device 同步是时延与稳定性的主要来源。实现ascend/qwen3_inference.py引入“推理请求计划”Qwen3AscendInferencePlan每个请求先在 CPU 上基于grid_thw元数据构建完整计划强制校验grid_thw必须位于 CPU一次性生成全部索引/权重/切分信息再通过ContextVar注入到运行时被覆盖的模型类Qwen3VLModel/Qwen3VLVisionModel等中执行在可安全省略处直接跳过 attention mask 构建。类覆盖不改动 checkpoint key也不影响上游原有路径。收益NPU 侧只执行静态、规整的张量运算消除形态波动与逐步同步对相同分辨率输入vision 计划签名可复用。6.4 首请求预热动机稳态请求时延正常但服务启动后的首个请求明显更慢。CANN 的二进制算子库已覆盖绝大多数常见算子形态通常不涉及即时编译首请求的额外耗时主要来自一次性开销ACL/设备初始化、显存池首次扩张、权重首次搬运到设备以及 torch_npu 侧首次算子下发。做法不改变默认行为仅在启动脚本与文档中提示该现象run_server.sh也会在启动时打印对应提示同时通过 6.3 的请求计划把视觉前处理的动态 shape 收敛为静态规整形态减少形态波动带来的首包抖动。收益避免把首请求耗时误判为性能问题压测与评测按“先预热一次、再统计稳态时延”的口径进行。6.5 依赖治理requirements-ascend.txt动机flash-attn、decord、eva-decord均会拉取 GPU 专用 wheel 或在昇腾 aarch64 上不可用依赖解析顺序不当会导致 pip 先装 GPU 版 torch。实现独立的requirements-ascend.txt置顶固定 torch 三件套torch2.5.1 / torchvision0.20.1 / torch-npu2.5.1.post1刻意排除上述 GPU 专用包视频解码改用torchvision_av后端numpy固定 1.26.4 以匹配 torch-npu 运行时。收益一条命令得到可复现的昇腾环境避免 GPU 版 wheel 混入。7. RoboTwin 仿真评测仿真评测端流程与上游原有流程完全一致按 DOMINO 上游评测章节配置 RoboTwin 环境将deploy_policy.yml指向本 policy server 的 host 与端口即可。架构上仿真在主机侧执行policy 在 NPU 侧执行二者通过 WebSocket 通信。由于 policy server 对外呈现的接口与上游完全一致仿真评测端完全不感知后端设备差异——这正是“后端中立”设计原则的直接收益。8. 已验证结果与后续工作8.1 已验证结果摘要已在 Atlas 910CANN 8.5.2上完成 policy server 端到端验证PUMA checkpoint 直接加载bf16RoboTwin 评测链路与上游原有流程一致环境不满足要求时无 torch_npu / NPU 不可用服务在启动阶段直接报错退出不会带着错误配置继续服务训练所得 checkpoint 可通过本在线推理样例在昇腾上直接服务训练样例见 manipulation/puma/train/README.md。8.2 后续工作来自优化文档探索 torchair 图模式进一步降低稳态时延上游 transformers 版本升级后回归验证 Qwen3 类覆盖的兼容性当前针对 4.57.0 验证。9. 相关说明与排障提示昇腾适配的实现要点、优化动机与收益docs/manipulation/puma/infer_with_torch/README.md训练样例manipulation/puma/train/README.md含 8 卡 DeepSpeed ZeRO-2 训练、世界模型监督、常见问题排查PUMA 样例总览与已验证环境manipulation/puma/README.md训练侧适配与优化文档docs/manipulation/puma/train/README.md。常见排障口径setup 末尾校验失败确认 CANNset_env.sh已 source、CANN 版本与torch-npu2.5.1.post1的配套要求匹配、numpy仍为 1.26.4首个请求很慢属于预期的一次性开销ACL/设备初始化、显存池扩张、权重搬运先预热一次再统计稳态时延启动即报错退出多半是 torch_npu 缺失或 NPU 不可用被“启动即校验”逻辑拦下属保护性行为而非故障。【免费下载链接】cann-recipes-embodied-ai本项目针对具身智能业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-ai创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考