Environment Variables and Compiler Options

Triton-distributed-ascend environment variables and compiler options are inherited from Triton Ascend and shmem (aclshmem). They can be used directly when developing Triton operators. For detailed lists, see the corresponding official documentation:

Environment Variables

Triton-distributed-ascend environment variables are divided into two categories:

  • Runtime environment variables: Set during program execution.

  • Build-time environment variables: Build configuration that takes effect when compiling and installing the Python package (pip install -e ./python), used by this project’s build process; for detailed build steps and prerequisites (CANN, torch_npu, LLVM), see README.md and build.md.

Environment Variable Usage Example

Environment variables need to be set before running the Python program, for example:

export SHMEM_UID_SESSION_ID=192.168.1.100:1234
python run_distributed_kernel.py

Compiler Options

Compiler options are used to control the compilation strategy of individual Triton kernels. They can be passed via meta-parameters during kernel launch, triton.Config, or Autotune parameters.

Compiler Option Usage Example

This project controls communication behavior within kernels through distributed primitives in triton_dist.language (dl) on the Ascend side, for example:

import triton
import triton.language as tl
import triton_dist.language as dl
from triton.language.extra.cann.extension import sub_vec_id

@triton.jit
def notify_kernel(signal_ptr, rank):
    if sub_vec_id() == 0:
        # Send signal notification to rank 0: set 1
        dl.notify(
            signal_ptr, 0,
            signal=1, sig_op="set", comm_scope="intra_node",
        )

@triton.jit
def wait_kernel(signal_ptr, world_size):
    if sub_vec_id() == 0:
        # Wait for notifications from all other ranks (barrier scans in 64-byte stride)
        token = dl.wait(
            signal_ptr, world_size - 1,
            scope="gpu", _semantic="acquire", waitValue=1,
        )
        data_ptr = dl.consume_token(signal_ptr, token)

Compiler Option Description

NPU compiler options for Triton distributed Ascend kernels (passed via triton.Config / launch meta-parameter) are defined by Triton Ascend and can be used directly when developing Triton operators. For detailed lists, see Triton Ascend Environment Variables and Compiler Options Reference.