triton_dist.language API

This module provides distributed computing primitives for Triton kernels on Ascend NPUs, enabling communication and synchronization across multiple processing elements in distributed systems.

All APIs are available through triton_dist.language or the shorthand tdl:

import triton_dist.language as tdl

For more background on the original Triton-distributed project, refer to the triton-distributed documentation.

Distributed Operations

wait(barrierPtrs, numBarriers, scope: str, semantic: str, waitValue: int = 1) → tensor

Wait on barrier pointers until the specified condition is met.

Parameters:

  • barrierPtrs – Scalar pointer or block pointer to barrier locations in memory

  • numBarriers – Number of barriers to wait on

  • scope – Memory scope. Ascend currently only supports scope "sys"; other parameter values are auto-promoted to "sys".

  • semantic – Memory semantic/ordering: "acquire", "release", "acq_rel", or "relaxed"

  • waitValue – Value to wait for at the barrier location (default: 1)

Returns:

int32 synchronization token. It can be passed as the second argument to consume_token to establish a dependency between the wait and subsequent data accesses

Example:

# Wait on a barrier with acquire semantics
token = tdl.wait(barrier_ptr, numBarriers=1, scope="sys", semantic="acquire", waitValue=1)
data_ptr = tdl.consume_token(data_ptr, token)
data = tl.load(data_ptr)
consume_token(value, token) → tensor or tensor_descriptor

Consume a token and associate it with a value, establishing an ordering dependency.

Parameters:

  • value – The tensor or tensor descriptor to associate with the token

  • token – Integer synchronization token returned by wait

Returns:

The input value with the token dependency attached (same type as input)

This operation is used to establish ordering constraints in distributed pipelines. The returned value maintains the same shape and type as the input.

rank(axis: int = -1) → tensor

Get the rank (ID) of the current processing element.

Parameters:

axis – Reserved for device mesh support. This parameter is ignored and the global rank is returned.

Returns:

int32 tensor containing the current rank

Example:

my_rank = tdl.rank()  # Get global rank
num_ranks(axis: int = -1) → tensor

Get the total number of ranks (processing elements) in the system.

Parameters:

axis – Reserved for device mesh support. This parameter is ignored and the total number of ranks is returned.

Returns:

int32 tensor containing the total number of ranks

Example:

total_ranks = tdl.num_ranks()  # Get total number of ranks
symm_at(ptr, rank) → tensor

Get the symmetric address of a pointer at the specified remote rank.

Parameters:

  • ptr – Scalar pointer to single symmetric memory address (must be a scalar pointer, not a block pointer)

  • rank – Target rank whose symmetric address to retrieve

Returns:

Tensor containing the symmetric pointer at the specified rank (same type as input pointer)

Only supports scalar pointers. Requires that the memory was allocated as symmetric across all ranks.

Example:

# Get symmetric address on rank 1
remote_ptr = tdl.symm_at(local_ptr, rank=1)
notify(ptr, rank, signal: int = 1, sig_op: str = 'set', comm_scope: str = 'inter_node') → tensor

Send a notification signal to a remote rank by updating a remote memory location.

Parameters:

  • ptr – Scalar pointer to the signal location (must be scalar pointer)

  • rank – Target rank to notify

  • signal – Signal value to send (default: 1)

  • sig_op – Signal operation type (default: "set")

  • comm_scope – Communication scope. Ascend currently ignores this parameter; the actual path is automatically selected by ACLSHMEM based on the target rank and runtime topology.

Returns:

Void tensor (operation has side effects but no return value)

Signal operations:

  • "set" — Set the value at the target location

  • "add" — Add to the value at the target location

Communication scopes:

  • "intra_node" — Within a single node

  • "inter_node" — Across nodes

Platform-specific Requirements (Ascend):

  • Signal pointer must be int32, int64, or uint64 type

  • For int64/uint64 types, only sig_op="set" is supported

  • For int32 type, both "set" and "add" operations are supported

Example:

# Notify rank 1 by setting signal to 1
tdl.notify(signal_ptr, rank=1, signal=1, sig_op="set", comm_scope="inter_node")

# Notify rank 2 by incrementing signal
tdl.notify(signal_ptr, rank=2, signal=1, sig_op="add", comm_scope="intra_node")

External Library Calls

extern_call(lib_name: str, lib_path: str, args: list, arg_type_symbol_dict: dict, is_pure: bool) → tensor

Invoke an external library function from within a Triton kernel.

Parameters:

  • lib_name – Name of the external library

  • lib_path – File system path to the library

  • args – List of arguments to pass to the function

  • arg_type_symbol_dict – Dictionary mapping argument type tuples to (function_name, return_type) pairs

  • is_pure – Whether the function is pure (no side effects)

Returns:

Tensor with the return value from the external function

All arguments must be scalar types (no block/tensor arguments). The function dispatches to the appropriate external symbol based on argument types. This is used internally by device-side SHMEM operations.

Example:

result = tdl.extern_call(
    "libshmem_device",
    "",
    [ptr, value, pe],
    {
        (tl.pointer_type(tl.int32), tl.int32, tl.int32): ("aclshmem_int32_p", ()),
    },
    is_pure=False
)

Notes

Memory Ordering and Synchronization

The distributed operations in this module interact with hardware memory hierarchies and interconnects. When using barrier operations, wait primitives, and notifications:

  • Choose appropriate semantic parameters based on your memory consistency requirements

  • Use acquire semantics when you need to observe writes from other threads

  • Use release semantics when publishing data to other threads

Symmetric Memory

Several operations (symm_at, notify) require symmetric memory — memory allocated at the same offset in each rank’s address space. Ensure memory is allocated using symmetric allocation primitives (see SHMEM Host API).

Platform Support

These APIs are designed for Ascend NPUs. Some operations have Ascend-specific behavior or constraints as noted in the individual API descriptions.