Get source lines from a function or source code string. Args: func_or_source: Either a function object or a string containing so

Share a RawCudaBuffer across processes via CUDA IPC. Uses int-array serialization to avoid null-byte truncation. Returns list[world_size] of d

Checks if a given integer value matches the filter's rule. Args: val: The integer value to check. Returns:

A helper function to encapsulate the benchmarking logic.

Shuffle down: each lane reads from (laneid + delta), clamped to 63.

Build a kernel ensuring it's loaded on the specified GPU.

no check dtype. no check device/host. no check tensor size. no check contiguous.

Low-level ds_bpermute_b32: read *value* from the lane at *byte_offset/4*.

Lazy initialization of Triton-dist backend

Convert an LLType instance to a valid C++ type string.

no check dtype. no check device/host. no check tensor size. no check contiguous.

Precomputes cosine and sine cache for rotary position embeddings.

Convert assignment target (Store ctx) to expression for use as Call argument (Load ctx). Required for A[0] etc.: ast.Name(id='A[0]') is invalid; m

Transform special struct constructor call into struct declaration. Example: scheduler = Scheduler(...) b

allgather gemm for intra-node return C = all_gather(A) @ B.T Args: A (torch.Tensor<float>): local matmul A matrix. shape: [M_per_ran

Applies Rotary Position Embedding inplace.

barrier_all_on_stream does not support CUDAGraph

Check if two tensors are close within a tolerance.

create context for allgather gemm inter-node Args: tensor_A (torch.Tensor<float>): local matmul A matrix. shape: [M_per_rank, K]

create context for allgather gemm intra-node Args: max_M: max number of M shape N(int): N K(int): K input_dtype(t

symmetric buffer requirement for input tensor x with x.nbytes = N. method | symmetric buffer size double_tree

create context for gemm reduce-scatter intra-node Args: max_M (int): max M N(int): N output_dtype(torch.dtype): dtype of

Usage: def test_xxx(dist_test): dist_test(my_worker_fn, world_size=2)

Exit the current scope.

Functional (inference) entry: full fused MoE forward; returns ``[num_tokens, hidden]``.

GEMM Reduce-Scatter for Multi-Node computes local GEMM (A x B) to generate partial results, followed by `reduce_scatter` to produce c Args:

GEMM Reduce-Scatter for Intra-Node return C = reduce_scatter(A @ B.T) Args: A (torch.Tensor<bfloat16/float16>): local matmul A matri

Get all registered Enum classes.

Get C++ operator string for binary operator.

Get the total nvshmem memory size in GB.

Get the global operator codegen registry.