Centralized type inferencer for LLType system. Uses modular components for maintainability and extensibility. Uses the unified Scope

Literal expression (integer, float, boolean constants)

C++ code emitter with context-aware LLType resolution and scope tracking. - Avoids redundant type declarations for in-scope variables (parame

Raised when types are incompatible (e.g., int + float)

Binary operation expression (e.g., a + b, a && b)

Unified CSV result writer with metadata header.

Raised when an AST node type is not supported for inference

Variable expression (access to local/global variables or parameters)

Base exception for type inference failures (with AST node context)

Unary operation expression (e.g., !a, *p)

Tuple type parameterized by a list of element types (e.g., Tuple[bool, int32, int32]).

A class to filter integers based on a defined rule. The rule can be an integer, a [min, max] list for a range, or None for a wildcard.

You may use this to barrier all ranks in global, or just in intra-node team. NOTE: nvshmem_barrier_all is slower for intra-node only.

Convert Python class AST to C++ struct definition. This visitor traverses a Python class definition and generates the corresponding

A block of sequential statements (e.g., body of an if/loop). Represents a scoped sequence of statements, common in structured control flow.

Raised when referencing a variable not in scope/context

Memory allocation statement (dynamic memory management). Represents allocation of memory in CUDA's memory spaces (e.g., `int* x = new int[N];

Expression for allocating a tensor (returns a tensor object). Mimics PyTorch's `torch.empty`, `torch.zeros`, etc. Supports multi-dimensional

Assignment statement (lhs = rhs). Assigns the value of an expression (rhs) to a mutable target (lhs).

Raw cudaMalloc buffer that is P2P accessible (unlike PyTorch tensors).

Fused EP All-to-All + grouped GEMM MoE op for AMD GPUs.

Raw cudaMalloc buffer for P2P-accessible cross-GPU memory.

Base class for allocation expressions (returns a pointer/handle to allocated memory). Serves as a common parent for raw memory allocation (e.

Function call expression with a return value. Unlike CallStmt (which is a statement with no return value), CallExpr is an expression that

For-loop statement (structured iteration). Represents a counted loop with initialization, condition, and increment. Common in CUDA for th

Unified symbol scope manager for all passes. Manages symbol scopes for passes, ensuring that local definitions take priority over g

Type for special structs like Scheduler that need custom code generation.

Tensor Parallel Attention. QKV Projection: Column Parallelism on weights (sharded over head dimension). Output Projection: Row Parallelis

Tensor Parallel MLP. This MLP uses a common TP strategy: 1. First linear layer (gate/up) weights are column-parallel. 2. Second linea

Function call statement (for void functions or side effects). Differs from CallExpr: CallStmt is a standalone statement (no return value),

Unified CUDA code emitter. Generates standalone CUDA code using raw PTX intrinsics. No CuTe/CUTLASS dependency. Uses CUtensorMap directly for

Variable declaration statement (with optional initialization). Declares a new variable with a type, name, and optional initial value. Sup

Dense model implementation for tensor parallel training. This model initializes the parameters, sets the forward pass method, and provides an

IR node representing an enum definition. This node is inserted into the Module body during materialize pass and processed during cod

Conditional statement (if-else). Executes one of two blocks based on a boolean condition.

Converts LLType to C++ type strings using registry.

Statement for allocating a tensor and assigning it to a variable. Combines tensor allocation (via `AllocateTensorExpr`) with variable assignm

Handles resolution of type annotations.

Handles type inference for function/method calls.

Compiled CUDA kernel that can be launched from Python.

Handles type inference for constant values.

EP MoE

A class to filter integers based on a defined rule. The rule can be an integer, a [min, max] list for a range, or None for a wildcard.

Root base class for all types in the system. Ensures consistent interface for equality, hashing, and string representation. All types are

Represents a loop modifier (e.g., ll.unroll, ll.parallel, ll.serial).

Centralized method return type resolver for type inference.

Information about a method in a struct stub.

Annotate for pointer types (e.g., ptr[uint32]).

Return statement (exits a function). May optionally return a value (for non-void functions).

Tensor Parallel MoE. This MoE uses a common TP strategy: 1. First linear layer (gate/up) weights are column-parallel. 2. Second linea

Annotate for template parameter types (e.g., template[uint32]).

Tensor type parameterized by its element scalar type (e.g., tensor[int32]).

TMA (Tensor Memory Accelerator) descriptor type (singleton).

this tracker contains multiple tracks to support overlaped tracks

While-loop statement (condition-based iteration). Executes a block repeatedly while a condition remains true.

Handles type inference for various AST node types.