Run Native DLA Inference in Python on NVIDIA Jetson AGX Orin Using cuDLA - NVIDIA Developer

Edge AI applications often require running multiple neural networks within strict compute and power budgets. NVIDIA's Deep Learning Accelerator (DLA) platforms offer dedicated hardware for this purpose, executing supported neural network operations on the Jetson AGX Orin. This allows applications to offload inference work from the GPU, improving power efficiency and freeing up GPU capacity for other tasks.

The company has now introduced a Python API called cuda.bindings.cudla, which exposes the cuDLA runtime directly to Python developers. This module provides thirteen functions and eight enums for hybrid-mode workflows, where CUDA streams coordinate DLA execution with GPU work. Previously, developers had to write C++ code to manage the runtime, but this new layer allows them to integrate DLA execution into their applications without writing C++.

New Python bindings for cuDLA runtime

The cuda.bindings.cudla module is a thin, Cython-generated Python layer over the cuDLA C++ API. It exposes the runtime directly to the application, while inference executes on the DLA hardware. This complements the existing C++ workflow by allowing runtime management and inference orchestration to remain in Python. The API is identical on both Jetson AGX Orin and the new RTX Spark platform, extending DLA inference to mainstream Windows AI developer laptops for the first time.

Building a DLA loadable for inference

To run inference on the DLA, a developer must first compile a trained model into a DLA loadable binary artifact. This is done using the TensorRT trtexec tool, which encodes the network in a format the DLA hardware can execute directly. The example targets TensorRT 11.4.1 with CUDA 13.4 Update 1 on Jetson AGX Orin running JetPack 7.2.2 or later.

The compilation process involves setting specific flags to target the first DLA engine, select the CHW16 data format, and produce a standalone DLA loadable. The command uses the --useDLACore=0 flag to target the first DLA engine and the --buildDLAStandalone flag to produce a raw DLA loadable that does not require TensorRT at inference time. This ensures the compiled model is optimized for the DLA's fixed-function accelerator.

Not every network layer is DLA-compatible, so preprocessing the ONNX graph to trim DLA-incompatible head and tail layers is often necessary. The resulting loadable can be a .bin or .nvdla file, which can then be loaded directly into the Python application. This workflow separates the one-time build step from the on-device runtime, allowing developers to compile a model once and run inference efficiently on the device.

Python workflow for hybrid GPU-DLA execution

The Python workflow for running inference involves several key steps, starting with initializing CUDA and creating a DLA device. The cudla.Mode.CUDA_DLA mode selects hybrid mode, where DLA tasks are submitted through a CUDA stream. This enables GPU-DLA pipelines, where preprocessing kernels, DLA inference, and postprocessing kernels can all be enqueued on the same stream and executed in order.

The developer loads the compiled loadable into memory using the module_load_from_memory function, which accepts the loadable as a Python bytes object. Before allocating memory, the developer introspects the loaded module to query the exact tensors it expects, including the number of inputs and outputs and their descriptors. This introspection step makes the workflow model-agnostic, as the same code paths handle any loadable regardless of network architecture.

After allocating GPU memory for the inputs and outputs, the developer must register the GPU buffers with the DLA device. This is a unique step in the cuDLA workflow, as the buffers must be registered with the DLA device before they can appear in a task. The registration returns a DLA-side handle that the hardware uses to access the buffer. The developer then prepares the input data and copies it to the GPU, followed by submitting the inference task to the stream.