环境变量与编译选项

Triton-distributed-ascend 环境变量与编译选项均来自 Triton Ascendshmem(aclshmem),开发 Triton 算子时可直接使用,详细清单见对应官方文档:

环境变量

Triton-distributed-ascend 的环境变量分为以下两类:

  • 运行时环境变量:程序运行期间设置。

  • 构建期环境变量:编译安装 Python 包(pip install -e ./python)时生效的构建配置,由本项目构建流程使用;详细的构建步骤与前置依赖(CANN、torch_npu、LLVM)见 README.mdbuild.md

环境变量用法示例

环境变量需在运行 Python 程序前设置,例如:

export SHMEM_UID_SESSION_ID=192.168.1.100:1234
python run_distributed_kernel.py

编译选项

编译选项用于控制单个 Triton kernel 的编译策略,可通过 kernel launch 时传入的 meta-parameter、triton.Config 或 Autotune 参数传入。

编译选项用法示例

本项目 Ascend 侧通过 triton_dist.languagedl)的分布式原语在 kernel 内控制通信行为,例如:

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:
        # 向 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:
        # 等待所有其他 rank 的通知(barrier 按 64 字节步长扫描)
        token = dl.wait(
            signal_ptr, world_size - 1,
            scope="gpu", _semantic="acquire", waitValue=1,
        )
        data_ptr = dl.consume_token(signal_ptr, token)

编译选项说明

Triton distributed Ascend kernel 的 NPU 编译选项(通过 triton.Config / launch meta-parameter 传入)由 Triton Ascend 定义,开发 Triton 算子时可直接使用,详细清单见 Triton Ascend 环境变量与编译选项参考