torch-rbln 0.1.6

Rebellions Extension for PyTorch

PyTorch for Rebellions' NPU

This package provides PyTorch integration for Rebellions' NPU.

Getting Started (torch: Python package, rebel-compiler: Python package)

Prerequisites

• Python 3.9 or later
• Git
• CMake 3.18 or later
• Ninja build system
• LDAP credentials for Rebellions' package repository

Update Git submodules

Clone submodules recursively. This will download required third-party libraries such as Rebel Compiler headers from third_party/rebel_compiler.

git submodule update --init ./

Create Python Virtual Environment

Create Python virtual environment. This will create a directory named .venv in the current directory.

python3 -m venv ./.venv && source ./.venv/bin/activate

Install Dependencies

Install Python package manager poetry. This will manage dependencies, building, packaging and installing.

pip3 install poetry==2.0.1

Save credentials for https://gate-keeper.rebellions.in. Authorization for https://pypi.rbln.in is required. But, this is not safe way to save credentials. See ~/.config/pypoetry/auth.toml.

export LDAP_USERNAME=daekyeong.kim     # Put your username
export LDAP_PASSWORD=mysecretpassword  # Put your password
poetry config keyring.enabled false    # Optional, if building freezes while auth
poetry config http-basic.rbln-internal $LDAP_USERNAME $LDAP_PASSWORD

NOTE During development we have to use rbln-internal instead of rbln. If you want to download rebel compiler from rbln (external pypi server of rebellions), do the following.

poetry config http-basic.rbln <rbln username> <rbln password>

Install dependencies written in poetry.lock using poetry, except the root package torch-rbln. Be careful, below command uninstall packages which is not on poetry.lock.

poetry sync --no-root

Choose Build Type (Optional)

Choose build type like below. Default is Release.

export RBLN_BUILD_TYPE=Debug

Install Editable Package

Build C++ project and install editable torch-rbln package.

poetry install --only-root

Control RBLN Log Verbosity

To control the verbosity of logs emitted by torch_rbln, including detailed debug information such as CPU fallback operations, you can set the following environment variable:

export TORCH_RBLN_LOG=DEBUG

Available log levels (from most to least verbose):

• VERBOSE – Show very detailed verbose messages including function entry/exit and parameter values. Note: VERBOSE logs are only available in Debug builds (when RBLN_BUILD_TYPE=Debug).
• DEBUG – Show debug messages including CPU fallbacks and trace markers.
• INFO – General runtime information without detailed debug context.
• WARNING – Only show important warnings.
• ERROR – Only show errors and critical failures.

The default level is WARNING. Setting VERBOSE provides the most detailed logging for deep debugging, but requires a Debug build. DEBUG is useful for general development or troubleshooting execution issues on RBLN devices and works in both Debug and Release builds.

Performance Optimization Flag (Optional)

To reduce runtime overhead (e.g., skipping unnecessary NaN/Inf checks), set the following environment variable:

export TORCH_RBLN_DEPLOY=ON

This enables lightweight execution for deployment scenarios.

Device Mapping Configuration

By default, each physical NPU device is mapped to a logical device with a 1:1 relationship (equivalent to RBLN_NPUS_PER_DEVICE=1). This is called Direct Mapping and provides the standard PyTorch device usage experience.

You can configure device mapping using the following environment variables to enable Aggregated Mapping, which groups multiple physical NPUs into a single logical device for RSD (Rebellions Scalable Design) functionality.

RBLN_NPUS_PER_DEVICE

Groups physical NPUs together to create logical devices. Each logical device will contain the specified number of physical NPUs. This is designed for Normal Users who want simple configuration.

Constraint: Must be one of the supported sizes: 1, 2, 4, 8, 16, or 32. These values match the base_sizes defined in rebel/core/compilation/_impl.py for production environments.

export RBLN_NPUS_PER_DEVICE=2

Examples:

With 4 physical devices (RBLN_DEVICES=0,1,2,3 or default):

• RBLN_NPUS_PER_DEVICE=2 → rbln:0 maps to NPUs [0, 1], rbln:1 maps to NPUs [2, 3]
• RBLN_NPUS_PER_DEVICE=4 → rbln:0 maps to NPUs [0, 1, 2, 3] (full aggregation)

With 6 physical devices and RBLN_NPUS_PER_DEVICE=4:

• rbln:0 maps to NPUs [0, 1, 2, 3]
• NPUs [4, 5] remain unused (warning will be displayed)

RBLN_DEVICE_MAP

Provides explicit mapping between logical devices and physical NPU IDs. This is designed for Advanced Users who need fine-grained control over device topology.

Constraint: Each device group must contain one of the supported sizes: 1, 2, 4, 8, 16, or 32 devices.

export RBLN_DEVICE_MAP="[0,1],[2,3,4,5]"

Format: Comma-separated groups of NPU IDs, each group enclosed in square brackets.

Example: With 6 physical devices:

• RBLN_DEVICE_MAP="[0,1],[2,3,4,5]" → rbln:0 maps to NPUs [0, 1], rbln:1 maps to NPUs [2, 3, 4, 5]

Configuration Priority and Conflict Resolution

Priority order: RBLN_DEVICE_MAP > RBLN_NPUS_PER_DEVICE > default (1:1 mapping)

Viewing Device Topology

You can view the current device topology using torch.rbln.device_summary():

import torch_rbln
torch.rbln.device_summary()

Example output:

[RBLN] Device Topology Initialized:
+-------------------+-------------------+----------------------+
| Logical Device    | Physical NPU IDs  | Status               |
+-------------------+-------------------+----------------------+
| rbln:0            | [ 0, 1 ]          | Active (Aggregated)  |
| rbln:1            | [ 2, 3 ]          | Active (Aggregated)  |
+-------------------+-------------------+----------------------+

Install Wheel Package (Optional)

If you want to make *.whl and install that, run below command.

poetry build
pip install ./dist/torch_rbln*.whl

When you change C++ or Python source code, you just run Install Editable Package or Install Wheel Package again.

Apply Custom rebel-compiler

You have 2 choices:

• Use built-in one
• Use external one

Use torch-rbln built-in rebel-compiler (torch: Python package, rebel-compiler: third_party/rebel_compiler)

This way is strongly recommended. Those are same with Getting Started.

git submodule update --init ./
python3 -m venv ./.venv && source ./.venv/bin/activate
pip3 install poetry==2.0.1
export LDAP_USERNAME=daekyeong.kim     # Put your username
export LDAP_PASSWORD=mysecretpassword  # Put your password
poetry config http-basic.rbln $LDAP_USERNAME $LDAP_PASSWORD

Without poetry sync, checkout rebel-compiler where ./third_party/rebel_compiler to your custom branch.

pushd ./third_party/rebel_compiler
  git checkout my_custom_branch
popd

It will make a package and install into your environment with syncing.

./tools/apply-custom-rebel.sh

Above script edits pyproject.toml and poetry.lock files. If you want to apply custom rebel-compiler temporarily, keep your eyes on those files.

(Optional) You can choose build type like below.

RBLN_BUILD_TYPE=Debug ./tools/apply-custom-rebel.sh

Then, you can build or install torch-rbln package on the custom rebel-compiler package.

poetry install --only-root

Use external rebel-compiler (for rebel-compiler developers)

Prereqs

• You’ve already built rebel-compiler.
• ${REBEL_HOME} points to the rebel-compiler repo root.

1) Create and activate a virtualenv

python3 -m venv .venv
source .venv/bin/activate

2) Add your local rebel-compiler in editable mode

poetry add --editable "${REBEL_HOME}/python"

3) Install this project, using the external compiler

RBLN_USE_EXTERNAL_REBEL_COMPILER=1 poetry install --only-root

Create Git Commit

Git pre-commit hook is working. So, when you create Git commit, linting would be triggered. For prepare linting, you MUST initialize lintrunner.

source ./.venv/bin/activate

lintrunner init

Once lintrunner initalized, no need to initialize again. You can commit now.

git commit

Some failures can be fixed automatically. Run below command for auto fixing.

lintrunner -m main -a

Run Tests

Assume that you are in Python virtual environment, and install torch-rbln package successfully.

C++ Tests

Making package runs in new isolated environment. Although you build your C++ project using poetry install --only-root, can't find that directory. So, for CTest you MUST build C++ project manually.

./tools/build-libtorch-rbln.sh

ctest --test-dir ./build

Python Tests

pytest ./test

---