[Feature]: Formalize allocator-agnostic symmetric tensor address translation

Suggestion Description

Iris device-side RMA and collective kernels are technically independent of the host allocator: given a local tensor pointer and a peer-base table, the device code can translate that pointer into the corresponding address on another rank.

However, the current host-side Iris API assumes tensors come from an Iris-owned symmetric heap and exposes a context-wide heap_bases table. We should formalize an allocator-agnostic interface so the same Iris device kernels can operate on tensors allocated through:

Current Iris TorchAllocator and VMemAllocator
rocSHMEM, including the VMM-fabric allocator
Torch Symmetric Memory
External locality-aware allocation abstractions
Imported external tensors

Iris should remain the single source of truth for device-side RMA, collective, and fused kernels. Allocation and tensor lifetime should remain owned by the respective host-side provider.

Motivation

Iris and rocSHMEM currently use the same fundamental symmetric-address translation:

  offset = local_pointer - local_allocation_base
  remote_pointer = peer_allocation_base[target_rank] + offset

The device-side implementation does not need to know:

Which allocator created the tensor
Whether the backing allocation is a process-wide heap or an individual allocation
Whether a peer is on the same node or reachable through VMM fabric
Whether the tensor was allocated directly or imported

What it does need is normalized address metadata for the tensor’s backing allocation.

Formalizing this boundary would allow all supported allocators to use the same Iris device-side library without duplicating collective kernels in allocator-specific repositories.

Current limitations

The current interface is centered on a context-wide heap:

  heap_bases = iris_ctx.get_heap_bases()

This assumes that all symmetric tensors use the same heap-base table. That is not sufficient for providers where:

Each allocation can have a different peer mapping
Imported tensors have a different allocation root
Tensor placement differs by locality
Multiple allocation providers are used in the same process or kernel
A tensor is a view into a larger registered allocation

Device context creation also embeds the Iris heap bases, coupling the context to the allocator selected when Iris is initialized.

Proposal

Introduce a per-allocation address descriptor, tentatively:

  SymmetricAddressMap:
      peer_bases: torch.Tensor      # uint64[world_size], device-resident
      local_rank: int
      allocation_base: int         # backing allocation base on this rank
      allocation_bytes: int
      capabilities: int            # optional capability mask

The important invariant is:

  peer_bases[local_rank] == local backing-allocation base

Device-side translation then becomes:

  @triton.jit
  def translate(pointer, target_rank, address_map):
      local_base = tl.load(address_map.peer_bases + address_map.local_rank)
      remote_base = tl.load(address_map.peer_bases + target_rank)
      offset = tl.cast(pointer, tl.uint64) - local_base
      return remote_base + offset

The exact Python representation is open for discussion. The stable contract is more important than the class name.

Host-provider contract

Each allocation provider should retain ownership of:

Tensor allocation and deallocation
Memory registration/import/export
Peer-handle exchange
VMM-fabric initialization
Tensor lifetime
Producing the peer-base table
Reporting supported remote-memory capabilities

Iris should consume only the normalized address descriptor.

Possible provider-facing API:

  address_map = provider.get_symmetric_address_map(tensor)

or:

  tensor, address_map = provider.allocate_symmetric(...)

Existing Iris allocation APIs can create and attach the same descriptor internally.

Per-allocation rather than context-wide metadata

The address descriptor should be associated with the tensor’s backing allocation rather than only with the Iris context.

This is needed because a single kernel may eventually consume:

One tensor allocated by rocSHMEM
Another tensor imported through an Iris allocator
Workspace allocated through Torch Symmetric Memory
A locality-aware tensor owned by another abstraction

Each tensor may require a different base table while using identical Iris device code.

Views must preserve the descriptor of their backing allocation so address translation uses the allocation root rather than the view pointer as its base.

Locality is a separate concern

Address translation and locality should remain separate interfaces.

The address map answers:

What address reaches this allocation on the target rank?

A separate topology/locality descriptor should answer:

What is the relationship and expected cost between these ranks?

For example:

  LocalityMap:
      group_ranks
      rank_stride
      peer_class       # same GPU/XCD/node/fabric pod/etc.
      direct_access
      system_atomics

A locality-aware allocator may provide both descriptors, but device address translation should not depend on the locality implementation.

This separation is important on MI450 VMM fabric: both intra-node and inter-node peers may have valid direct addresses, while locality remains useful for hierarchical collective selection and traffic shaping.

Device-side API direction

  Device-side RMA operations should accept an explicit address map or symmetric tensor view:

  iris.load(pointer, from_rank, address_map)
  iris.store(pointer, value, to_rank, address_map)
  iris.atomic_add(pointer, value, to_rank, address_map)

An object-oriented form could be:

  symmetric_view = IrisSymmetricView(pointer, address_map)
  value = symmetric_view.load(from_rank)

The existing heap_bases arguments can remain temporarily as a compatibility path.

Torch Symmetric Memory

Torch Symmetric Memory uses a different representation from the current Iris and rocSHMEM heaps.

Preferred integration: TBD, need exploration.

Backward compatibility

Suggested migration:

Add the per-allocation descriptor without removing Iris.get_heap_bases().
Have current Iris allocators produce descriptors.
Update DeviceContext and RMA functions to accept explicit descriptors.
Keep context-wide heap bases as a deprecated convenience for tensors from the default Iris heap.
Add rocSHMEM and Torch Symmetric integration tests.
Migrate collectives and fused operations to the allocator-agnostic interface.

Acceptance criteria

Non-goals

This issue does not propose:

Replacing rocSHMEM, Torch Symmetric, or external host allocators
Moving their memory-management implementations into Iris
Making Iris own externally allocated tensor lifetimes
Duplicating collective implementations per allocator
Defining the final locality-aware collective-selection policy

The goal is to establish a stable boundary where multiple host-side allocation providers can use one Iris device-side kernel library.

Operating System

No response

GPU

No response

ROCm Component

No response

Thanks for putting this together. I agree with the main direction. I want to propose one refinement to the client-facing API, though. Passing a pointer and address map separately is flexible but it allows accidental combinations such as a rocSHMEM pointer with an Iris address map and seem to introduce some cognitive load on developers. Instead, the provider could bind them once into a Triton-compatible structured argument like this:

  src = rocshmem_adapter.bind(roc_tensor)
  dst = iris_ctx.bind(iris_tensor)

  kernel[grid](
      iris_ctx.get_device_context(),
      src,
      dst,
  )

The kernel would then use:

  values = ctx.load(
      src,
      offsets=offsets,
      from_rank=source_rank,
      mask=mask,
  )

  ctx.store(
      dst,
      offsets=offsets,
      value=values,
      to_rank=target_rank,
      mask=mask,
  )

Note that bind function returns a buffer of named tuple SymmetricBuffer, instead of raw pointer, containing only the device-relevant fields:

SymmetricBuffer:
    ptr
    peer_bases
    capabilities

Another note: Triton can represent the bound argument as a structured tuple and flatten its pointer fields into normal kernel arguments:

https://github.com/triton-lang/triton/blob/b1233aa326fa485b08de8593da2d08cb853c346b/python/src/specialize.cc#L439
https://github.com/triton-lang/triton/blob/b1233aa326fa485b08de8593da2d08cb853c346b/python/test/unit/language/test_tuple.py#L146-L203

The complete SymmetricAddressMap from the proposal would still be the canonical provider contract and would contain allocation base, size, local rank, capabilities, and other validation metadata. bind() would validate the tensor against that map and produce the compact kernel argument.

Backward compatibility would be maintained:

# Existing API: continues using the default Iris heap bases.
ctx.load(pointer + offsets, from_rank=peer, mask=mask)

# New API: uses the peer bases attached to this allocation.
ctx.load(buffer, offsets=offsets, from_rank=peer, mask=mask)

Note that the additional offsets parameter is needed because the bound buffer represents the allocation rather than an already-offset raw pointer. However, we'll handle such details, for example, in the load function by checking if offsets is None or not, something like this:

      if offsets is None:
          # Legacy interface:
          #
          # ctx.load(pointer + offsets, from_rank=peer)
          pointer = operand
          peer_bases = self.heap_bases
      else:
          # Allocator-agnostic interface:
          #
          # ctx.load(buffer, offsets=offsets, from_rank=peer)
          pointer = operand.ptr + offsets
          peer_bases = operand.peer_bases

I also agree that locality should remain separate from address translation. A provider may expose both, but the core RMA translation should depend only on the bound address metadata.

Thanks! For P0 we agreed on the following:

Iris host-side will provide a new API:

tensor, address_map = provider.allocate_symmetric(...)

The tensor and its address map are passed as two separate kernel arguments. We are deliberately not wrapping them in a single struct for now: the named-tuple flattening that would make that ergonomic requires a recent Triton, and several environments we care about are pinned to older ROCm/Triton stacks.

We will implement TBD kernels (provided by Dimple) that directly use the pointer and address maps, manually doing the translation inside the kernel. Later on, we will revisit the device-side APIs for translation and other RMA operations.

Example:

# Host
src, src_map = provider.allocate_symmetric((N,), dtype=torch.float32)
dst, dst_map = provider.allocate_symmetric((N,), dtype=torch.float32)

put_kernel[grid](
    src,
    dst, dst_map.peer_bases,   # uint64[world_size], device-resident
    N,
    peer,
    CUR_RANK=dst_map.local_rank,
    BLOCK=1024,
)

@triton.jit
def put_kernel(src,
               dst, dst_peer_bases,
               n_elements,
               target_rank,
               CUR_RANK: tl.constexpr,
               BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n_elements

    # Hoisted once, outside the access path: two loads regardless of how
    # many translations the kernel performs.
    local_base = tl.load(dst_peer_bases + CUR_RANK)
    remote_base = tl.load(dst_peer_bases + target_rank)
    ptr_delta = remote_base - local_base

    # Translate the allocation base once, then index into it. The offset
    # is relative to the backing allocation, so views translate against
    # the allocation root rather than the view pointer.
    remote_dst = (dst.to(tl.uint64) + ptr_delta).to(dst.dtype)

    values = tl.load(src + offs, mask=mask)
    tl.store(remote_dst + offs, values, mask=mask)

Existing get_heap_bases() and the current RMA APIs are unchanged — this is additive.

Drop the descriptor struct and return the peer base table directly

Reduce allocate_symmetric to one method

Tighten the tests and drop an unsound hint

Name the heap offset for what it is

Build peer views through uint8 to dodge the bf16 typestr bug

Correct the per-allocation claim: one context-wide table suffices

Add a CI-shaped test, and record why this file lives in Iris

Navigation Menu

Repository navigation

[Feature]: Formalize allocator-agnostic symmetric tensor address translation #546

Suggestion Description

Motivation

Current limitations

Proposal

Host-provider contract

Per-allocation rather than context-wide metadata

Locality is a separate concern

Device-side API direction

Torch Symmetric Memory

Backward compatibility

Acceptance criteria

Non-goals

Operating System

GPU

ROCm Component

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[Feature]: Formalize allocator-agnostic symmetric tensor address translation #546

Description

Suggestion Description

Motivation

Current limitations

Proposal

Host-provider contract

Per-allocation rather than context-wide metadata

Locality is a separate concern

Device-side API direction

Torch Symmetric Memory

Backward compatibility

Acceptance criteria

Non-goals

Operating System

GPU

ROCm Component

Activity

artulab commented on Aug 28, 2026

mawad-amd commented on Aug 28, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions