跳到主要内容
返回时间线
NVIDIA Technical Blog来源发表:

Jetson上的DLA推理改由Python直接调度:cuda-bindings 13.3.0开放cuDLA接口,GPU与DLA可同流排队

核心概要

NVIDIA在cuda-bindings 13.3.0中发布cuda.bindings.cudla模块,把cuDLA的C++运行时接口以Cython生成的薄层暴露给Python,使Jetson AGX Orin上的DLA推理无需编写C++即可完成设备创建、网络加载、缓冲区注册与任务提交。该模块面向混合模式,DLA任务通过CUDA流与GPU前后处理按序执行;同一套API在RTX Spark上保持一致。发布说明同时指出,当前版本不支持需要NvSci缓冲区与同步原语的独立模式。

Source-provided article image: Run Native DLA Inference in Python on NVIDIA Jetson AGX Orin Using cuDLA
Figure 1 · NVIDIA 开发者技术博客

深度剖析

边缘AI应用常要在有限算力与功耗下同时跑多个神经网络,Jetson AGX Orin上的DLA是专门执行深度学习推理的硬件,可把推理从GPU卸载出去,从而改善能效并给GPU留出容量。 此前管理DLA执行需要走cuDLA的C++运行时接口,应用要自己用C++处理设备、网络与缓冲区。 cuda.bindings.cudla作为cuDLA C++ API之上的薄Python层随cuda-bindings 13.3.0发布,暴露十三个函数、九种数据类型和八个枚举,推理仍在DLA硬件上执行。

工作流从创建DLA设备开始,加载编译好的网络,按模块自报的张量描述符分配并注册缓冲区,再把推理任务提交到CUDA流上。 cudla.Mode.CUDA_DLA选择混合模式,预处理内核、DLA推理和后处理内核可以排在同一条流上按序执行。 示例以ResNet-50骨干编译成的DLA loadable运行,输出显示一个输入张量resnetv17_pool0_fwd和一个输出张量resnetv17_stage4_activation2,形状分别为(1,64,56,56)与(1,2048,7,7)。

缓冲区大小不写死在代码里,而是从模块查询到的张量描述符中读取,因此同一套代码路径可处理任意loadable,只有预处理与后处理随模型变化。 loadable以Python字节对象直接传入module_load_from_memory,不需要文件路径或文件系统句柄。 编译侧用trtexec的--buildDLAStandalone生成推理时不需要TensorRT的裸DLA loadable,示例针对Jetson AGX Orin上JetPack 7.2.2或更高版本的TensorRT 11.4.1与CUDA 13.4 Update 1。

启示与展望

演示范围是Jetson AGX Orin上JetPack环境中的图像分类示例,从编译loadable到读取输出的完整生命周期;同一API在RTX Spark上保持一致,因此文中cuDLA接口部分可直接适用。

当前版本不支持独立模式,请求该模式会抛出ErrorUnsupportedOperation;并非所有网络层都与DLA兼容,常需预处理ONNX图以裁掉不兼容的首尾层;TensorRT 11从模型张量类型推导精度,--fp16与--int8不再受支持。