Generated codegen function. args_cpp are the C++ code strings for function arguments.

Default handler for unimplemented AST nodes.

Check if this is the first PP stage

Reference PyTorch forward pass using Tensor Parallelism. Final output is AllReduced. x: input tensor, shape [batch_size * seq

Initializes and shards MLP parameters for Tensor Parallelism. mlp: A standard nn.Module MLP (e.g., from HuggingFace Transformers).

Checks if 'mx-smi' is available on the system's PATH.

split src/dest outside of _recv_ll. this function is designed for a threadblock num_ints: of the pre-LL-packed num_ints.

symm_flags is expected to: 1. of int32 dtype 2. has at least num_ranks * 2 + 1 elements 3. of symmetric pointer

Ceiling division: (a + b - 1) // b Returns the same type as the first argument (or second if first is not LLType). If both arguments are const

triton_dist forward pass for TP. This version uses ag_gemm and gemm_rs. x: input tensor, shape [batch_size, seq_len, hidden_s

Format an error message with location information and code snippet. Args: message: Error message node: AST node where th

Run fused dispatch + grouped GEMM. Returns ``(gemm_output[M, N], meta)``. ``meta`` carries the layout (recv offsets, scatter idx, per-rank token

Get the total nvshmem memory size in MB.

Create int64 tensor of data_ptr() values on target_device.

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

Load data from global memory with PTX instructions multimem.ld_reduce Args: ptr: Pointer to the global memory acc_prec: Accum

Write a line to header buffer (for mixin compatibility).

Device-kernel version of ``_build_dispatch_metadata`` (same meta keys). Replaces the host ``bincount`` + ``all_gather_into_tensor`` + recv-offset

Return SM architecture integer (e.g. 90, 100), or None.

split src/dest outside of _recv_ll. this function is designed for a threadblock nbytes: of the pre-LL-packed bytes. BLOCK_SIZE: count by ints

Performs an all-reduce operation on the input tensor `x`. This is the main entry point for the all-reduce collective. It selects the specifi

GPU bincount for AMD (no AOT). input_tensor: 1D int32 on device; output: length elements.

create context for allgather gemm intra-node Args: rank (int): current rank num_ranks (int): total number of ranks max_M:

Enter a new scope (e.g. function body, loop body). Args: initial_types: Optional initial type mappings for this

low-latency all-to-all communication

gemm swizzle Nz

Get all registered special structs.