Transform special struct method calls. This method marks the call node for special handling in codegen. For methods

Try to find class in ctx that has this method.

Try to find class in ctx that has this method.

Try to resolve method return type from special struct registry.

Try to resolve method return type from special struct registry.

Unparse a Module node, handling IR nodes separately.

Validate context contains `little_kernel.language` or LLType objects.

return List[(block_idx, group_idx, task_type, start_time, end_time)] sorted by start_time

Visit an expression node using mixin's visit method.

Helper method to process if statement body and else block. Args: node: ast.If node to process is_else_if: If

Write JSON to disk with proper structure

Add license header to a file that is missing it.

allgather gemm for inter-node Allgather global matrix A and do matmul with local matrix B, produces local matrix C Args: a (torch.Te

allgather gemm for inter-node Allgather global matrix A and do matmul with local matrix B, produces local matrix C Args: a (torch.Te

no-tma allgather gemm for intra-node Allgather global matrix A and do matmul with local matrix B, produces local matrix C Args: a (t

Args: input: Input tensor of shape (M, K) input_scale: Optional input scales (M,) comm_data_buffers: List of communicati

consume gemm

with the local_ptr

Write single result immediately to CSV with forced flushing

semantic should be one of ["monotonic", "release", "acquire", "acq_rel] scope should be one of ["workgroup", "agent", "system]

custom atomic_add implementation using extern_elementwise :param scope: one of "gpu", "sys". default to "gpu" :param semantic: one of "releas

Atomic non-fetch uint64 operation (thread scope). Args: dest: Symmetric address on target PE val: Value to be used in atomic

NOTE: this function should only be called with atomic support. memory over PCI-e does not support atomic r/w. DON'T use this function on such platfor

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

Benchmark performance.

Benchmark the implementation

index = bisect(sorted_ptr, values, N) off = sorted_ptr[index - 1] : let suppose values[-1] = 0 remainder = values - off It's expecte

Block index in the x dimension.

Build the compute offset kernel using LK's build() API. shared_mem_bytes = NUM_WARPS * (num_experts + 1) * sizeof(int32) The kernel uses

Build the dispatch layout kernel.

Build a launchable kernel using LittleKernel's Python binding. Parameters ---------- grid : tuple Grid dimensions (num_blocks, 1,

Build the barrier kernel.