RBLN PyTorch APIs and Supported Operations¶

Check if any RBLN devices are available.

Returns:

Get the number of available RBLN devices.

Returns:

Get the number of physical RBLN devices in the system.

This function returns the actual number of physical devices, regardless of whether logical device aggregation is enabled.

Returns:

Get the index of the currently selected RBLN device.

Returns:

Set the current device.

Parameters:

Wait for all work on the given RBLN device: its pending asynchronous transfers, then every stream on it (torch.cuda.synchronize parity).

If no device is specified, the current device is used.

Parameters:

1
2

cpu_tensor = rbln_tensor.to("cpu", non_blocking=True)
torch.rbln.synchronize()  # wait for the transfer to complete

True once set_device() has run, or when there are no devices at all.

The no-device case is deliberate: it makes torch.distributed's init_device_mesh skip its get_rank() % device_count() auto-select, which would otherwise fail on a host with no NPU.

Note: this calls device_count(), which plans the RBLN device mapping from the current RBLN_* environment. It does not claim an NPU or freeze the mapping; both happen on first device use, so a launcher may still assign RBLN_DEVICES afterwards.

Whether the process runs on a host-backed dummy device rather than an NPU.

is_available() is True on a dummy device too, so use this to tell a dummy device apart from real NPU availability.

Returns:

Print a summary of the RBLN device topology showing the mapping between logical devices and physical NPU IDs.

The table shows:

• Logical device indices (rbln:0, rbln:1, ...)
• Physical NPU IDs mapped to each logical device
• Status of each device (Active, Active (Aggregated), or Unused)
• A warning about unused physical devices, if any

The device mapping is determined by the following environment variables (in order of priority):

1. RBLN_DEVICE_MAP: explicit groups, for example "[0,1],[2,3,4,5]"
2. RBLN_NPUS_PER_DEVICE: uniform groups of the given size, for example "2"

1
2
3

import torch

torch.rbln.device_summary()

1
2
3
4
5
6
7

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

Describe the NPUs behind an RBLN logical device (torch.cuda parity).

A logical device may span several physical NPUs (RBLN_DEVICE_MAP / RBLN_NPUS_PER_DEVICE), so total_memory sums them and npu_count says how many. num_chiplet and memory_per_chiplet stay per NPU: a device runs out on its heaviest chiplet, which total_memory hides, so size a pool by memory_per_chiplet.

The query does not allocate memory or claim the device, so it is safe to call before the runtime initializes. Properties are not cached: each call queries every NPU behind the logical device.

Parameters:

Returns:

Raises:

1
2

props = torch.rbln.get_device_properties(0)
props.total_memory  # 150323855360

Return the NPU name behind an RBLN logical device, for example "RBLN-CA25" (torch.cuda parity).

Same as get_device_properties(device).name.

Parameters:

Returns:

Get a list of data types supported by automatic mixed precision (AMP) on RBLN devices.

Returns:

Return the random number generator state of device as a ByteTensor (torch.cuda.get_rng_state parity).

Parameters:

Returns:

Return the random number generator state of every device, in device-index order.

Returns:

Set the random number generator state of device.

Parameters:

Set the random number generator state of every device.

Parameters:

Return the currently selected stream for a given device.

Parameters:

Returns:

Return the default stream for a given device.

Parameters:

Returns:

Return a context manager that selects the given stream.

Parameters:

Returns:

1
2
3

s = torch.rbln.Stream()
with torch.rbln.stream(s):
    y = x * 2

Select a stream for all subsequent work, and select its device as well.

This sets the stream for the remainder of the calling thread. Prefer the stream context manager, which restores the previous stream on exit.

Parameters:

Return the current device memory occupied by tensors in bytes for a given device.

Parameters:

Returns:

Return the maximum device memory occupied by tensors in bytes for a given device.

Parameters:

Returns:

Return the current device memory managed by the caching allocator in bytes for a given device.

Parameters:

Returns:

Return the maximum device memory managed by the caching allocator in bytes for a given device.

Parameters:

Returns:

Return a dictionary of device memory allocator statistics for a given device.

The returned dictionary contains various memory statistics including:

• allocated.current: Current memory occupied by tensors
• allocated.peak: Peak memory occupied by tensors
• allocated.total_allocated: Total memory allocated to tensors (cumulative)
• allocated.total_freed: Total memory freed from tensors (cumulative)
• reserved.current: Current memory managed by the caching allocator
• reserved.peak: Peak memory managed by the caching allocator
• reserved.total_allocated: Total memory allocated by the caching allocator (cumulative)
• reserved.total_freed: Total memory freed by the caching allocator (cumulative)
• active.current: Current size of blocks in use (may differ from allocated due to block granularity)
• active.peak: Peak size of blocks in use
• cached.current: Current size of cached blocks available for reuse
• cached.peak: Peak size of cached blocks
• num_alloc_retries: Number of allocation retries after cache flush
• num_ooms: Number of out-of-memory errors
• num_device_alloc: Number of device memory acquisitions
• num_device_free: Number of device memory releases

Lazy Tensor Memory Allocation: All memory-related functions in this module (including memory_allocated(), memory_reserved(), max_memory_allocated(), max_memory_reserved(), and this function) reflect device memory only, not CPU memory.

RBLN tensors use lazy memory allocation for device memory. When you create a tensor on an RBLN device:

• The tensor is initially allocated in CPU memory immediately upon creation
• Device memory allocation is deferred until the tensor is actually needed for device operations
• When a device operation is required, the tensor data is lazily transferred from CPU to device memory

This lazy allocation strategy means that memory statistics may be lower than expected immediately after tensor creation until the tensors are used in device computations. Device memory statistics will increase when tensors are materialized on the device during actual computation.

The statistics also include lazy tensor related metrics, which provide insights into the memory management for tensors that have not yet been materialized on the device. The specific lazy tensor statistics fields may vary depending on the implementation version.

Parameters:

Returns:

Return memory allocator statistics broken down per chiplet.

Same keys as memory_stats(), each prefixed with npu.<n>.chiplet.<c>. -- e.g. npu.0.chiplet.0.allocated.current. A device runs out on its heaviest chiplet, which the aggregate memory_stats() hides.

npu.<n> is the n-th physical NPU of this logical device (RBLN_NPUS_PER_DEVICE / RBLN_DEVICE_MAP), not a physical NPU id; a 1:1 mapping yields npu.0 only. Same scope as memory_stats().

Parameters:

Returns:

Return a human-readable printout of the current memory allocator statistics.

Rows are per (NPU, chiplet), so an imbalance is visible at a glance. The scope line under the title names what the numbers cover: they are this process's allocator, not the NPU's occupancy (see memory_stats()). The npu column is the physical NPU id as torch.rbln.device_summary() prints it: relative to RBLN_VISIBLE_DEVICES, so it differs from rbln-smi when that variable is set. A logical device grouping several NPUs (RBLN_NPUS_PER_DEVICE / RBLN_DEVICE_MAP) gets one row per NPU.

Peaks are tracked per chiplet, so the peak columns of the total row bound the joint peak from above rather than reporting it. memory_stats() carries the exact joint peak.

Parameters:

Returns:

Return the free and total device DRAM of device in bytes, as (free, total).

Same contract as torch.cuda.mem_get_info, and what torch.accelerator.get_memory_info reports for rbln: the NPU's figures as the kernel driver sees them, across every process, not this process's caching allocator, which memory_stats() covers. total is the pool the driver can hand out, so it sits below the part's nominal DRAM. A logical device spanning several physical NPUs (RBLN_NPUS_PER_DEVICE / RBLN_DEVICE_MAP) reports their sum; one tensor still lives on one NPU. The result is a reading, not a reservation: another process may allocate right after.

Parameters:

Returns:

Raises:

Return the driver's device DRAM usage of device broken down per chiplet, in bytes.

Same figures as mem_get_info() before summation, keyed like memory_stats_per_chiplet(): npu.<n>.chiplet.<c>.{total,used,free,largest_free, largest_free_huge} for each chiplet, and npu.<n>.{total,used,free,granularity, huge_granularity} for each physical NPU of the logical device. npu.<n> is the NPU's position within the logical device, not a physical NPU id.

Each chiplet allocates from its own pool, so a physically contiguous buffer is bounded by one chiplet's largest_free (largest_free_huge for the huge granule; 0 where the part has none), which the device-wide free hides. largest_free may exceed free in the same reply: both are estimates from a lock-free walk.

Parameters:

Returns:

Raises:

Reset the "peak" stats tracked by the caching allocator for a given device.

This function resets the peak values to their current values for the following stats:

• allocated.peak: Reset to allocated.current
• reserved.peak: Reset to reserved.current
• active.peak: Reset to active.current
• cached.peak: Reset to cached.current

Parameters:

Reset the "accumulated" (historical) stats tracked by the caching allocator for a given device.

This function resets the following accumulated stats to zero:

• allocated.total_allocated
• allocated.total_freed
• reserved.total_allocated
• reserved.total_freed
• num_alloc_retries
• num_ooms
• num_device_alloc
• num_device_free

Parameters:

Release all unoccupied cached memory currently held by the caching allocator so that those can be used in other application.

Note that only unfragmented (non-split) blocks can be released; fragmented blocks that have been split will remain in the cache until they can be coalesced.

Parameters:

Materialize tensor's device allocation instead of leaving it lazy.

A device allocation reserves a virtual address and materializes the physical memory behind it on first use through a torch op. A consumer that reads those physical buffers out of band, such as a collective library or direct storage (NVMe) DMA, never runs such an op, so it would find nothing there. Call this after allocating a buffer you are handing to one.

The region is laid out flat and 1:1 with no dtype transform, on the device's main node, which is what a consumer treating it as bytes expects. Use set_device_layout_like() instead when the layout has to match another tensor's. Binding an already-bound region is allowed.

Parameters:

Raises:

1
2

staging = torch.empty(nbytes, dtype=torch.uint8, device="rbln:0")
torch.rbln.bind_device_memory(staging)

Configure target's device-allocation layout to match ref (no copy).

Both must be RBLN tensors with the same dtype, on the same device, and each covering its whole storage. ref must be device-resident. target adopts ref's layout and dtype while keeping its own size; no data is transferred. A subsequent device-to-device copy between target and ref then stays on the fast path.

Typical use: make a host→device staging buffer match a KV cache's layout so the bulk upload and the per-slot device-to-device scatter are both fast.

Parameters:

Raises:

Allocate a host buffer the device can DMA into without a staging copy.

An ordinary CPU tensor is 64 B aligned; the runtime's copy path stages any host address that is not page aligned through a bounce buffer, and even an aligned one pays a page fault per 4 KiB the first time the runtime resolves its host addresses, a cost that lands in the transfer, not in setup. This returns 2 MiB-aligned memory instead, prefaulted, so neither happens. It also asks for transparent huge pages, but that part is best effort: with THP disabled the alignment still holds and nothing is huge-page backed.

The buffer is released when the last reference to the returned tensor (or to a view of it) goes away.

Parameters:

Returns:

Raises:

1
2

slab = torch.rbln.huge_host_empty(1 << 30)
slab.view(...)  # hand to whatever consumes the host side

Context manager that enables RBLN file offloading for its scope.

Inside the with block, the process-wide file offloading switch is on, so host-side regions backing RBLN tensors allocated within the block may be paged out to disk. Use this around code paths that allocate large host-resident tensors (for example, KV-cache initialization) where host RAM pressure matters.

Nested offload blocks are tracked via a thread-safe depth counter; the switch is flipped back off only when the outermost context exits.

1
2

with torch.rbln.offload():
    tensor = torch.zeros(1 << 30, device="rbln:0")  # offloaded

Remove this process's file offloading temp files and directories.

offload() writes into a per-process directory under RBLN_OFFLOAD_DIR (default $HOME/.cache/rbln_cache/offload) that the runtime removes on teardown. Call this on a shutdown path that may be killed first. Offloaded tensors must not be used afterwards.

Returns:

Return a hidden-overhead explain region, usable as a context manager.

A normal PyTorch operation can silently round-trip the host (NPU to CPU to NPU), fall back to a CPU kernel, or trigger a recompilation. None of that is visible in your Python code. explain() counts those hidden events, attributes a cause and a remedy to each, and renders a torch.profiler-style report.

This is a hidden-overhead explainer, not a timing profiler: it does not report how long a forward pass took or which layer is slow. A region with no hidden overhead can still be slow because it is device-compute bound. For wall-clock timing, use torch.profiler.

Parameters:

Returns:

1
2
3
4
5
6

with torch.rbln.explain() as p:
    model(x)

print(p.report())  # [clean] / [overhead] marker plus a signal table
p.verdict()        # {'clean': bool, 'reasons': [...]} for CI gating
p.dump()           # every signal as raw numbers

Place two explain regions around fn: one around the first call it makes and one around a later call after warmup more calls.

The helper only automates region placement. Its cold and warm labels mean the first call it made and a later call it made, nothing more. Reading the result as one-time versus recurring cost rests on two assumptions that the caller has to guarantee, because the helper cannot verify them:

1. fn was not run or compiled before this call. Otherwise the one-time cost happened before the cold sample and is not captured.
2. Every fn() call does the same work. Otherwise the recompiles in the warm sample come from different shapes, not from recurring overhead.

When both hold, cold.diff(warm) shows what changed: signals gone in the warm region did not recur across the two calls, and signals that persist are the recurring overhead. The helper is plain Python orchestration and adds nothing to the hot path.

Parameters:

Returns:

1
2

d = torch.rbln.explain_steady(lambda: model(x), warmup=2, as_diff=True)
print(d.report())  # signals marked as persisting recur across the two calls

Context manager that records the programs the rbln backend builds inside it.

torch.compile keeps the callable a backend returns inside Dynamo's code cache, so the caller never receives a handle to what was compiled. This scope yields a list that fills, in build order, as torch.compile invokes the backend, the way torch.profiler.profile() collects the events inside it.

Every open scope receives each program, so nested scopes see the same programs; an inner scope simply sees fewer. The scope is thread-local: it records the backend calls made from the thread that opened it.

Returns:

1
2
3
4
5
6
7
8
9

with torch.rbln.capture_programs() as programs:
    compiled = torch.compile(model, backend="rbln")
    compiled(*warmup_inputs)  # Dynamo builds the graphs here

for program in programs:
    print(program.name, program.device)
    for spec in program.input_specs:
        if spec.physical_placement is not None:
            print(spec.name, spec.physical_placement.shards)