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

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

Wait on barrier pointers until the specified condition is met.

Parameters:

  • barrierPtrs – Pointer or tensor of pointers to barrier locations in memory

  • numBarriers – Number of barriers to wait on

  • scope – Memory scope (note: triton-distributed-ascend does not currently support scope variations)

  • semantic – Memory semantic/ordering (e.g., "acquire", "relaxed")

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

Returns:

int32 tensor indicating completion status

Example:

# Wait on a barrier with acquire semantics
status = tdl.wait(barrier_ptr, num_barriers=1, scope="", semantic="acquire", waitValue=1)

Source: python/triton_dist/language/distributed_ops.py:58-85

consume_token

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 token value (must be of integer type)

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.

Source: python/triton_dist/language/distributed_ops.py:89-97

rank

rank(axis: int = -1) → tensor

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

Parameters:

axis – The axis along which to query the rank (default: -1 for global rank)

Returns:

int32 tensor containing the current rank

Example:

my_rank = tdl.rank()  # Get global rank
rank_on_axis0 = tdl.rank(axis=0)  # Get rank along specific axis

Source: python/triton_dist/language/distributed_ops.py:101-103

num_ranks

num_ranks(axis: int = -1) → tensor

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

Parameters:

axis – The axis along which to query the number of ranks (default: -1 for global)

Returns:

int32 tensor containing the total number of ranks

Example:

total_ranks = tdl.num_ranks()  # Get total number of ranks
ranks_on_axis0 = tdl.num_ranks(axis=0)  # Get number of ranks along specific axis

Source: python/triton_dist/language/distributed_ops.py:107-109

symm_at

symm_at(ptr, rank) → tensor

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

Parameters:

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

  • 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)

Source: python/triton_dist/language/distributed_ops.py:113-116

notify

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 (default: "inter_node")

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")

Source: python/triton_dist/language/distributed_ops.py:120-146

External Library Calls

extern_call

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
)

Source: python/triton_dist/language/core.py:100-146

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.