triton-dist.language

This page documents all Python APIs in Triton-distributed.

triton_dist.language

Core distributed language extensions for Triton.

Distributed Operations

wait(barrierPtrs, numBarriers, scope, semantic, waitValue=1)

Wait for distributed barriers.

Parameters:

barrierPtrs – Pointer to barrier signals
numBarriers – Number of barriers to wait for
scope – Memory scope - "gpu" or "sys"
semantic – Memory semantic - "acquire", "release", etc.
waitValue – Value to wait for (default: 1)

Returns:

Token for consume_token

consume_token(value, token)

Consume a token to ensure proper ordering with wait operations.

Parameters:

value – Value (pointer or tensor) to consume
token – Token from wait operation

Returns:

The consumed value

rank(axis=-1)

Get the current rank ID.

Parameters:: axis – Axis for multi-dimensional rank (default: -1 for global rank)
Returns:: Rank ID as int32

num_ranks(axis=-1)

Get the total number of ranks.

Parameters:: axis – Axis for multi-dimensional rank count
Returns:: Number of ranks as int32

symm_at(ptr, rank)

Map a symmetric pointer to a remote rank.

Parameters:

ptr – Local symmetric pointer
rank – Target rank ID

Returns:

Pointer on the target rank

notify(ptr, rank, signal=1, sig_op='set', comm_scope='inter_node')

Send a signal to a remote rank.

Parameters:

ptr – Signal pointer
rank – Target rank ID
signal – Signal value (default: 1)
sig_op – Signal operation - "set" or "add"
comm_scope – Communication scope - "intra_node" or "inter_node"

SIMT Operations

simt_exec_region()

Context manager for SIMT execution region.

Usage:

with simt_exec_region():
    # SIMT code here
    pass

class vector

Vector type for SIMT operations.

__init__(args): Create a vector from a sequence of tensors.

__getitem__(idx): Get element at index.

__setitem__(idx, value): Set element at index.

__add__(other): Element-wise addition.

__sub__(other): Element-wise subtraction.

__mul__(other): Element-wise multiplication.

recast(new_elem_dtype): Bitcast to a different element type.

to(dtype, fp_downcast_rounding=None, bitcast=False): Convert to a different dtype.

make_vector(args)

Create a vector from arguments.

Parameters:: args – Sequence of tensor values
Returns:: vector object

zeros_vector(vec_size, dtype)

Create a zero-initialized vector.

Parameters:

vec_size – Vector size (constexpr)
dtype – Element data type

Returns:

vector of zeros

extract(input, indices)

Extract a scalar from a tile at given indices.

Parameters:

input – Input tensor
indices – List of indices

Returns:

Scalar value

insert(input, scalar, indices)

Insert a scalar into a tile at given indices.

Parameters:

input – Input tensor
scalar – Scalar value to insert
indices – List of indices

Returns:

Modified tensor

Core Operations

extern_call(lib_name, lib_path, args, arg_type_symbol_dict, is_pure)

Call an external library function.

Parameters:

lib_name – Name of the library
lib_path – Path to the library
args – Function arguments
arg_type_symbol_dict – Type-to-symbol mapping
is_pure – Whether the function is pure (no side effects)

triton_dist.language.extra

Extra language utilities.

language_extra Module

tid(axis)

Get the thread index within a CTA.

Parameters:: axis – Axis (0 for threadIdx.x)
Returns:: Thread index

atomic_cas(ptr, cmp_val, target_val, scope, semantic)

Atomic compare-and-swap.

Parameters:

ptr – Target pointer
cmp_val – Compare value
target_val – Target value if comparison succeeds
scope – Memory scope
semantic – Memory semantic

atomic_add(ptr, val, scope, semantic)

Atomic add operation.

Parameters:

ptr – Target pointer
val – Value to add
scope – Memory scope
semantic – Memory semantic

__syncthreads(): Block-level synchronization barrier.

ld(ptr, scope, semantic)

Load with memory ordering.

Parameters:

ptr – Pointer to load from
scope – Memory scope
semantic – Memory semantic

st(ptr, val, scope, semantic)

Store with memory ordering.

Parameters:

ptr – Pointer to store to
val – Value to store
scope – Memory scope
semantic – Memory semantic

ld_vector(ptr, vec_size, scope, semantic)

Load a vector.

Parameters:

ptr – Pointer to load from
vec_size – Vector size (constexpr)
scope – Memory scope
semantic – Memory semantic

Returns:

vector object

st_vector(ptr, vec, scope, semantic)

Store a vector.

Parameters:

ptr – Pointer to store to
vec – vector object to store
scope – Memory scope
semantic – Memory semantic

pack(src, dst_type): Pack values into a larger type.

unpack(src, dst_type): Unpack a value into smaller types.

threads_per_warp(): Get threads per warp (32 for NVIDIA, 64 for AMD).

num_threads(): Get total number of threads (num_warps * threads_per_warp).

num_warps(): Get number of warps.

libshmem_device Module

NVSHMEM/ROCSHMEM device API bindings. The API automatically dispatches to the appropriate backend.

Initialization:

my_pe(): Get current PE (processing element) ID.

n_pes(): Get total number of PEs.

team_my_pe(team): Get PE ID within a team.

team_n_pes(team): Get number of PEs in a team.

Memory Operations:

remote_ptr(local_ptr, pe): Get remote pointer on another PE.

remote_mc_ptr(team, ptr): Get multicast remote pointer.

Synchronization:

barrier_all(): Global barrier across all PEs.

barrier_all_block(): Block-level global barrier.

barrier_all_warp(): Warp-level global barrier.

barrier(team): Barrier within a team.

sync_all(): Global sync across all PEs.

quiet(): Ensure all outstanding operations complete.

fence(): Memory fence.

Get Operations:

getmem_nbi_block(dest, source, bytes, pe): Non-blocking block-level get.

getmem_block(dest, source, bytes, pe): Blocking block-level get.

getmem_nbi_warp(dest, source, bytes, pe): Non-blocking warp-level get.

getmem_warp(dest, source, bytes, pe): Blocking warp-level get.

getmem_nbi(dest, source, bytes, pe): Non-blocking get.

getmem(dest, source, bytes, pe): Blocking get.

Put Operations:

putmem_block(dest, source, bytes, pe): Block-level put.

putmem_nbi_block(dest, source, bytes, pe): Non-blocking block-level put.

putmem_warp(dest, source, bytes, pe): Warp-level put.

putmem_nbi_warp(dest, source, bytes, pe): Non-blocking warp-level put.

putmem(dest, source, bytes, pe): Put operation.

putmem_nbi(dest, source, bytes, pe): Non-blocking put.

Signal Operations:

putmem_signal_nbi_block(dest, source, bytes, sig_addr, signal, sig_op, pe): Non-blocking block-level put with signal.

putmem_signal_block(dest, source, bytes, sig_addr, signal, sig_op, pe): Block-level put with signal.

signal_op(sig_addr, signal, sig_op, pe): Signal operation.

signal_wait_until(sig_addr, cmp_, cmp_val): Wait until signal meets condition.

Collective Operations:

broadcastmem_block(team, dest, source, nelems, pe_root): Block-level broadcast.

fcollectmem_block(team, dest, source, nelems): Block-level fixed collect.

Constants:

NVSHMEM_CMP_EQ: Compare equal

NVSHMEM_CMP_NE: Compare not equal

NVSHMEM_CMP_GT: Compare greater than

NVSHMEM_CMP_LT: Compare less than

NVSHMEM_CMP_GE: Compare greater or equal

NVSHMEM_CMP_LE: Compare less or equal

NVSHMEM_SIGNAL_SET: Signal set operation

NVSHMEM_SIGNAL_ADD: Signal add operation

NVSHMEM_TEAM_WORLD: World team

NVSHMEMX_TEAM_NODE: Node team

triton_dist.utils

Utility functions for distributed programming.

Initialization

initialize_distributed(seed=None, initialize_shmem=True)

Initialize distributed runtime.

Parameters:

seed – Random seed (default: rank)
initialize_shmem – Whether to initialize NVSHMEM/ROCSHMEM

Returns:

Tensor parallel process group

init_nvshmem_by_torch_process_group(pg): Initialize NVSHMEM using a PyTorch process group.

init_rocshmem_by_torch_process_group(pg): Initialize ROCSHMEM using a PyTorch process group.

finalize_distributed(): Finalize distributed runtime.

is_shmem_initialized(): Check if SHMEM is initialized.

Tensor Allocation

nvshmem_create_tensor(shape, dtype)

Create a symmetric tensor with NVSHMEM.

Parameters:

shape – Tensor shape
dtype – Data type

Returns:

Symmetric torch.Tensor

nvshmem_create_tensors(shape, dtype, rank, local_world_size)

Create a list of peer tensors.

Returns:: List of tensors for each local rank

nvshmem_free_tensor_sync(tensor): Free a symmetric tensor.

Synchronization

nvshmem_barrier_all_on_stream(stream=None): Barrier all on a CUDA stream.

rocshmem_barrier_all_on_stream(stream=None): Barrier all on a HIP stream.

Platform Detection

is_cuda(): Check if running on CUDA platform.

is_hip(): Check if running on HIP platform.

has_tma(): Check if TMA is supported (SM90+).

is_nvshmem_multimem_supported(): Check if NVSHMEM multicast memory is supported.

Device Info

get_numa_node(device_id): Get NUMA node for a device.

get_max_gpu_clock_rate_in_khz(device_id): Get maximum GPU clock rate.

get_current_gpu_clock_rate_in_khz(device_id): Get current GPU clock rate.

get_device_max_shared_memory_size(device_id): Get maximum shared memory size.

has_fullmesh_nvlink(): Check if NVLink full mesh is available.

supports_p2p_native_atomic(): Check if P2P native atomic is supported.

Utility Functions

dist_print(*args, allowed_ranks=[0], prefix=False, need_sync=False, **kwargs)

Print from specific ranks.

Parameters:

allowed_ranks – List of ranks to print from, or “all”
prefix – Whether to prefix with rank
need_sync – Whether to synchronize before printing

rand_tensor(shape, dtype, device='cuda')

Generate random tensor.

Parameters:

shape – Tensor shape
dtype – Data type
device – Device

Returns:

Random tensor with values in [-1, 1] for floats

init_seed(seed=0): Initialize random seeds for reproducibility.

get_bool_env(env, default_value): Get boolean environment variable.

get_int_env(env, default_value): Get integer environment variable.

NVSHMEM_SIGNAL_DTYPE: Default signal dtype (torch.int64)

triton_dist.tools

Development and debugging tools.

Compile Tools

make_ast_source(...): Generate AST source for AOT compilation.

kernel_name_suffix(...): Generate kernel name suffix.

materialize_c_params(...): Materialize C parameters.

dump_c_code(...): Dump generated C code.

aot_compile_spaces(...): AOT compile with configuration spaces.

apply_triton340_inductor_patch(): Apply compatibility patch for torch.compile with Triton 3.4.0.

Profiler

class Profiler: Intra-kernel profiler for performance analysis.

alloc_profiler_buffer(...): Allocate profiler buffer.

reset_profiler_buffer(...): Reset profiler buffer.

class ProfilerBuffer: Profiler buffer class.

export_to_perfetto_trace(...): Export profile to Perfetto trace format.

parse_to_tracks(...): Parse profile data to tracks.

triton_dist.mega_triton_kernel

MegaTritonKernel for fused multi-layer kernels.

Core

class ModelBuilder: Build fused models with MegaTritonKernel.

Tasks

Available tasks for MegaTritonKernel:

Activation tasks
AllReduce tasks
Barrier tasks
Elementwise tasks
Flash Attention tasks
Flash Decode tasks
Linear tasks
Normalization tasks
Prefetch tasks

See python/triton_dist/mega_triton_kernel/tasks/ for implementation details.

Example Usage

import triton_dist
import triton.language as tl
import triton_dist.language as dl
from triton_dist.language.extra import libshmem_device
from triton_dist.language.extra.language_extra import tid, __syncthreads

@triton_dist.jit
def my_distributed_kernel(input_ptr, output_ptr, ...):
    rank = dl.rank()
    num_ranks = dl.num_ranks()

    # Use symmetric pointers
    remote_ptr = dl.symm_at(input_ptr, peer_rank)

    # Wait for data
    token = dl.wait(signal_ptr, 1, "sys", "acquire")
    input_ptr = dl.consume_token(input_ptr, token)

    # Notify completion
    dl.notify(signal_ptr, peer_rank, signal=1, sig_op="set", comm_scope="intra_node")