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``: .. code-block:: python import triton_dist.language as tdl For more background on the original Triton-distributed project, refer to the `triton-distributed documentation `_. Distributed Operations ---------------------- wait ^^^^ .. py:function:: wait(barrierPtrs, numBarriers, scope: str, semantic: str, waitValue: int = 1) -> tensor Wait on barrier pointers until the specified condition is met. :param barrierPtrs: Pointer or tensor of pointers to barrier locations in memory :param numBarriers: Number of barriers to wait on :param scope: Memory scope (note: triton-distributed-ascend does not currently support scope variations) :param semantic: Memory semantic/ordering (e.g., ``"acquire"``, ``"relaxed"``) :param waitValue: Value to wait for at the barrier location (default: 1) :returns: int32 tensor indicating completion status Example: .. code-block:: python # 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 ^^^^^^^^^^^^^ .. py:function:: consume_token(value, token) -> tensor or tensor_descriptor Consume a token and associate it with a value, establishing an ordering dependency. :param value: The tensor or tensor descriptor to associate with the token :param 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 ^^^^ .. py:function:: rank(axis: int = -1) -> tensor Get the rank (ID) of the current processing element. :param axis: The axis along which to query the rank (default: -1 for global rank) :returns: int32 tensor containing the current rank Example: .. code-block:: python 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 ^^^^^^^^^ .. py:function:: num_ranks(axis: int = -1) -> tensor Get the total number of ranks (processing elements) in the system. :param 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: .. code-block:: python 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 ^^^^^^^ .. py:function:: symm_at(ptr, rank) -> tensor Get the symmetric address of a pointer at the specified remote rank. :param ptr: Scalar pointer to symmetric memory (must be a scalar pointer, not a block) :param 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: .. code-block:: python # 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 ^^^^^^ .. py:function:: 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. :param ptr: Scalar pointer to the signal location (must be scalar pointer) :param rank: Target rank to notify :param signal: Signal value to send (default: 1) :param sig_op: Signal operation type (default: ``"set"``) :param 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: .. code-block:: python # 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 ^^^^^^^^^^^ .. py:function:: 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. :param lib_name: Name of the external library :param lib_path: File system path to the library :param args: List of arguments to pass to the function :param arg_type_symbol_dict: Dictionary mapping argument type tuples to (function_name, return_type) pairs :param 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: .. code-block:: python 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 :doc:`shmem_host`). Platform Support ^^^^^^^^^^^^^^^^ These APIs are designed for Ascend NPUs. Some operations have Ascend-specific behavior or constraints as noted in the individual API descriptions.