Skip to content

RBLN PyTorch: A PyTorch extension for RBLN NPUs

RBLN PyTorch beta status

RBLN PyTorch is currently in beta and is actively developed. APIs may change or be removed between releases, and backward compatibility is not guaranteed. Supported operations are currently limited. Try it and share feedback to help stabilize it for general availability.

Overview

RBLN PyTorch (torch-rbln) is a PyTorch extension that enables RBLN NPUs to be used directly from PyTorch. By supporting eager (define-by-run) execution, it preserves the familiar PyTorch workflow across model development, debugging, and serving.

torch-rbln is designed so that RBLN NPUs can be used naturally from standard PyTorch code. To that end it implements the rbln device string, the torch.rbln module, and torch.compile with the rbln backend, among other features. It runs on both NPU families, Rebel100™ (RBLN-CR cards) and ATOM™ (RBLN-CA cards); the differences that matter to PyTorch code are collected in Execution Semantics › Rebel100™ and ATOM™. This page explains what RBLN PyTorch provides, how it works, the software layers that make it possible, and where to go next.

What RBLN PyTorch provides

Surface What it gives you Where to read more
device="rbln" Tensor creation, arithmetic, indexing, and copies on RBLN NPUs, with the same code you write for cuda or cpu. Eager operations are compiled per operator and cached. Execution Semantics, Supported operations (eager)
torch.rbln Device management and properties (device_count, set_device, synchronize, device_summary, get_device_properties), streams and events, memory statistics and allocator controls, the driver's memory query (mem_get_info), the explain() hidden-overhead explainer, and capture_programs() for inspecting the programs torch.compile builds. APIs, Compiled programs
torch.compile(backend="rbln") Graph compilation through the RBLN Compiler. The backend registers itself on the first torch.compile call. Execution Semantics
torch.distributed backend rbln-ccl A c10d process group for collectives across RBLN devices. Requires real hardware. Tutorial: Collectives with rbln-ccl, Execution Semantics › Distributed
torch.profiler bridge Setting RBLN_PROFILER=1 before the runtime starts includes RBLN NPU activity in torch.profiler traces through the Kineto bridge. Rebel100™ NPUs only; on ATOM™ NPUs the trace shows host activity. Tutorial: Profiling on Rebel100™
python -m torch_rbln.diagnose Prints the library search paths, package locations, the librbln.so that loaded, and the ABI (Application Binary Interface) verdict. Installation, Troubleshoot

Execution flow

When a PyTorch operation runs, the dispatcher selects the backend based on tensor device metadata. For tensors on the rbln device, the dispatcher routes the operation to the RBLN path through the PrivateUse1 key and then on through torch.compile to the RBLN Compiler.

RBLN runtime path in PyTorch: the main flow goes from ATen op to the Dispatcher, which routes RBLN tensors through PrivateUse1, then to torch.compile, the RBLN Compiler, and the RBLN Driver; an inset shows the dispatcher matrix and highlights the PrivateUse1 column. RBLN runtime path in PyTorch: the main flow goes from ATen op to the Dispatcher, which routes RBLN tensors through PrivateUse1, then to torch.compile, the RBLN Compiler, and the RBLN Driver; an inset shows the dispatcher matrix and highlights the PrivateUse1 column.

  1. User-facing operations such as torch.add resolve to ATen operators such as aten::add.
  2. The PyTorch dispatcher routes rbln tensors through PrivateUse1, the dispatch key PyTorch reserves for out-of-tree backends such as RBLN.
  3. torch.compile serves as the entry point to the RBLN path. In eager mode each operator is compiled on its own; a graph captured explicitly with torch.compile(backend="rbln") is compiled as a whole.
  4. The RBLN Compiler handles compilation and execution for the RBLN path.
  5. The RBLN Driver provides the low-level device interface.

Unlike operators that compute, operators that only move or index data (copy_, view, cat, index_select, and similar) are implemented as native C++ implementations in RBLN PyTorch and do not go through the compiler. The Supported operations tables mark which path each operator takes in eager execution; inside a torch.compile graph the RBLN Compiler decides.

Software stack

The runtime flow above shows execution order. The stack below shows the same system as software layers: PyTorch, torch-rbln, the RBLN Compiler, and the RBLN Driver. The compiler layer includes runtime components for compilation and execution, while the driver provides the low-level hardware interface to RBLN NPUs.

RBLN PyTorch software stack: four numbered layers (1–4), from PyTorch to the RBLN Driver. The compiler layer includes runtime components for compilation and execution, and the driver is the low-level interface to RBLN NPUs. RBLN PyTorch software stack: four numbered layers (1–4), from PyTorch to the RBLN Driver. The compiler layer includes runtime components for compilation and execution, and the driver is the low-level interface to RBLN NPUs.

Layer Library / file Role
PyTorch (torch) libtorch_python.so Interface between Python and the PyTorch C++ backend
libtorch.so Main PyTorch library; provides PyTorch C++ APIs
libtorch_cpu.so LibTorch CPU backend; also executes CPU fallback operations
libc10.so Low-level PyTorch utilities and core structures (tensor management, device abstraction, memory allocation, etc.)
RBLN PyTorch (torch-rbln) libtorch_rbln.so ATen operator C++ library: copies, views, indexing, and other data-movement kernels; the c10d process group; the Kineto bridge
register_ops.py Python operator bindings for compute operators, generated at build time from native_functions.yaml; each binding compiles its operator through torch.compile
libc10_rbln.so RBLN C10 extension: device type, allocator, streams, logging, and device-mapping hooks for libc10
torch.rbln Python module registered as the rbln device module: device management, streams and events, memory, explain()
RBLN Compiler (rebel-compiler) librbln.so Compiler library with runtime components for RBLN NPUs; exports the ABI entry points torch-rbln checks at import
RBLN Driver librbln-thunk.so Low-level device / hardware interface for RBLN NPUs

torch-rbln and rebel-compiler are released together. The release notes list the versions that ship in each SDK release, and the import-time ABI check enforces that the two you installed are compatible.

Next steps