RBLN PyTorch API¶

PyTorch Native APIs¶

Most of PyTorch native functions can be used on RBLN NPU.

RBLN Specific APIs¶

The following APIs are exposed to users through the torch.rbln module.

Classes¶

`device(device)` ¶

Context-manager that changes the selected device.

Parameters:

Name	Type	Description	Default
`device`	`torch.device, int, str, or None`	device index to select. It is a no-op if this argument is a negative integer or `None`.	required

Example

with torch.rbln.device(0):
    x = torch.randn(2, 2, device="rbln:0")

`device_of(obj)` ¶

Context-manager that changes the current device to that of given object.

You can use both tensors and storages as arguments. If a given object is not allocated on an RBLN device, this is a no-op.

Parameters:

Name	Type	Description	Default
`obj`	`Tensor or Storage`	object allocated on the selected device.	required

Example

tensor = torch.tensor([2, 64], device="rbln:1")
with torch.rbln.device_of(tensor):
    # operations here will use the same device as tensor (rbln:1)
    new_tensor = torch.zeros(3, device="rbln")

`RBLNExplain` ¶

A hidden-overhead explain region. Create one with explain() and use it as a context manager.

Regions must not overlap: the underlying counters and timers are process-global, so opening a nested or concurrent region raises a RuntimeError. The counters are also process-wide, meaning the values include activity from every thread that ran during the region, not only the calling thread.

Methods:¶

`report()` ¶

Render the human-readable report: a [clean] / [overhead: N signals] marker, a torch.profiler-style signal table with a Note column carrying the remedy for each actionable row, and per-signal detail blocks.

Returns:

Name	Type	Description
`str`	`str`	The report text. Print this.

`verdict()` ¶

Summarize the region as a clean flag plus the factual reasons behind it.

The result is intended for CI gating. clean states whether any hidden signal fired; it is not a severity grade.

Returns:

Type	Description
`Dict[str, Any]`	Dict[str, Any]: A dictionary containing `clean` and `reasons`, among others.

`dump()` ¶

Return every signal as raw numbers, for programmatic checks.

Returns:

Type	Description
`Dict[str, Any]`	Dict[str, Any]: Host-bounce counts and bytes by site, dispatch counters
`Dict[str, Any]`	(CPU fallbacks, recompilations, warm hits), per-operator breakdowns,
`Dict[str, Any]`	runtime residency, and the region wall time.

Functions:¶

`get_amp_supported_dtype()` ¶

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

Returns:

Type	Description
`List[dtype]`	List[torch.dtype]: A list of data types supported by AMP.

Note

Automatic mixed precision is not implemented for RBLN devices yet, so this currently returns an empty list. As a result, entering torch.autocast("rbln") emits a warning and leaves autocast disabled instead of raising an error.

`is_available()` ¶

Check if any RBLN devices are available.

Returns:

Name	Type	Description
`bool`	`bool`	True if at least one RBLN device is available, False otherwise.

`current_device()` ¶

Get the index of the currently selected RBLN device.

Returns:

Name	Type	Description
`int`	`int`	The index of the currently selected RBLN device.

`device_count()` ¶

Get the number of available RBLN devices.

Returns:

Name	Type	Description
`int`	`int`	The number of available RBLN devices.

`physical_device_count()` ¶

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:

Name	Type	Description
`int`	`int`	The number of physical RBLN devices.

`set_device(device)` ¶

Set the current device.

Parameters:

Name	Type	Description	Default
`device`	`device or int or str`	selected device.	required

`synchronize(device=None)` ¶

Wait for all pending asynchronous transfers on the given RBLN device to complete.

Parameters:

Name	Type	Description	Default
`device`	`device or int or str`	The device to synchronize. If None, uses the current device. Defaults to None.	`None`

Example

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

`empty_cache(device=None)` ¶

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:

Name	Type	Description	Default
`device`	`int, str, or torch.device`	The device to empty cache for. If None, uses the current device. Defaults to None.	`None`

`memory_allocated(device=None)` ¶

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

Parameters:

Name	Type	Description	Default
`device`	`int, str, or torch.device`	The device to query. If None, uses the current device. Defaults to None.	`None`

Returns:

Name	Type	Description
`int`	`int`	The current memory occupied by tensors in bytes.

Note

This function reflects device memory only. For information about lazy memory allocation, see :func:memory_stats.

`memory_reserved(device=None)` ¶

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

Parameters:

Name	Type	Description	Default
`device`	`int, str, or torch.device`	The device to query. If None, uses the current device. Defaults to None.	`None`

Returns:

Name	Type	Description
`int`	`int`	The current memory managed by the caching allocator in bytes.

Note

This function reflects device memory only. For information about lazy memory allocation, see :func:memory_stats.

`max_memory_allocated(device=None)` ¶

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

Parameters:

Name	Type	Description	Default
`device`	`int, str, or torch.device`	The device to query. If None, uses the current device. Defaults to None.	`None`

Returns:

Name	Type	Description
`int`	`int`	The maximum memory occupied by tensors in bytes.

Note

This function reflects device memory only. For information about lazy memory allocation, see :func:memory_stats.

`max_memory_reserved(device=None)` ¶

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

Parameters:

Name	Type	Description	Default
`device`	`int, str, or torch.device`	The device to query. If None, uses the current device. Defaults to None.	`None`

Returns:

Name	Type	Description
`int`	`int`	The maximum memory managed by the caching allocator in bytes.

Note

This function reflects device memory only. For information about lazy memory allocation, see :func:memory_stats.

`memory_stats(device=None)` ¶

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 :func:memory_allocated, :func:memory_reserved, :func:max_memory_allocated, :func: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:

Name	Type	Description	Default
`device`	`int, str, or torch.device`	The device to query. If None, uses the current device. Defaults to None.	`None`

Returns:

Type	Description
`Dict[str, int]`	Dict[str, int]: A dictionary containing device memory statistics. Note that these statistics
`Dict[str, int]`	reflect device memory only (not CPU memory) and may not include memory for tensors that have
`Dict[str, int]`	not yet been transferred to the device.

Note

To see accurate device memory usage, check statistics after performing operations that require the tensors to be materialized on the device, as device memory is allocated lazily when needed. This applies to all memory-related functions in this module.

`reset_peak_memory_stats(device=None)` ¶

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:

Name	Type	Description	Default
`device`	`int, str, or torch.device`	The device to reset stats for. If None, uses the current device. Defaults to None.	`None`

`reset_accumulated_memory_stats(device=None)` ¶

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:

Name	Type	Description	Default
`device`	`int, str, or torch.device`	The device to reset stats for. If None, uses the current device. Defaults to None.	`None`

`offload()` ¶

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.

Example

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

`explain(with_stack=False)` ¶

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:

Name	Type	Description	Default
`with_stack`	`bool`	If True, capture the Python call site of each fallback, recompilation, and host bounce (deduplicated per operator) and show it in the report. Off by default, so a plain region adds nothing to any code path. Defaults to False.	`False`

Returns:

Name	Type	Description
`RBLNExplain`	`RBLNExplain`	A region object usable as a context manager.

Example

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

Note

torch.rbln.profile and torch.rbln.RBLNProfile are backward-compatible aliases of explain and RBLNExplain. The trace= keyword is a deprecated alias for with_stack=.

RBLN PyTorch API¶

PyTorch Native APIs¶

RBLN Specific APIs¶

Classes¶

device(device) ¶

device_of(obj) ¶

RBLNExplain ¶

Methods:¶

report() ¶

verdict() ¶

dump() ¶

Functions:¶

get_amp_supported_dtype() ¶

is_available() ¶

current_device() ¶

device_count() ¶

physical_device_count() ¶

set_device(device) ¶

synchronize(device=None) ¶

empty_cache(device=None) ¶

memory_allocated(device=None) ¶

memory_reserved(device=None) ¶

max_memory_allocated(device=None) ¶

max_memory_reserved(device=None) ¶

memory_stats(device=None) ¶

reset_peak_memory_stats(device=None) ¶

reset_accumulated_memory_stats(device=None) ¶

offload() ¶

explain(with_stack=False) ¶

`device(device)` ¶

`device_of(obj)` ¶

`RBLNExplain` ¶

`report()` ¶

`verdict()` ¶

`dump()` ¶

`get_amp_supported_dtype()` ¶

`is_available()` ¶

`current_device()` ¶

`device_count()` ¶

`physical_device_count()` ¶

`set_device(device)` ¶

`synchronize(device=None)` ¶

`empty_cache(device=None)` ¶

`memory_allocated(device=None)` ¶

`memory_reserved(device=None)` ¶

`max_memory_allocated(device=None)` ¶

`max_memory_reserved(device=None)` ¶

`memory_stats(device=None)` ¶

`reset_peak_memory_stats(device=None)` ¶

`reset_accumulated_memory_stats(device=None)` ¶

`offload()` ¶

`explain(with_stack=False)` ¶