For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo function

launch

def launch[Body: RowBody, //, axis: Int, simd_width: Int, supports_splitk: Bool = True, computationally_expensive: Bool = False, BLOCK_SIZE: Int = Int(256), TILED_BLOCK_SIZE: Int = Int(32), COOPERATIVE_BLOCK_SIZE: Int = Int(128), WARP_BLOCK_WARPS: Int = Int(4), cache_dtype: Optional[DType] = None, cache_count: Int = Int(0), num_phases: Int = Int(0), dtype_size: Int = Int(0)](body: Body, shape: Coord, ctx: DeviceContext)

Top-level scaffolder. Picks the tier from shape + axis, instantiates the matching kernel, and launches.

Tier choice:

• Inner axis (axis == rank - 1) — contiguous reduce axis. Warp-per-row when row_size <= WARP_SIZE, else block-per-row with a 2D SIMD-width heuristic over (num_rows, row_size, sm_count).
• Non-inner axis, output count saturates the device — tiled kernel (one thread per row tile, coalesced SIMD load + store on the innermost non-axis dim).
• Non-inner axis otherwise — block-per-output cooperative (one block per output row, threads collaborate on the strided reduce axis with simd_width = 1).

Parameters:

• ​axis (Int): Axis being reduced.
• ​simd_width (Int): The SIMD width the scaffolder uses for tile dispatch. Bodies compute this via rowwise.pick_simd_width[...] and pass it directly.
• ​supports_splitk (Bool): Whether the body can run in the split-K tier. Default True — single-state reductions opt in implicitly. Bodies that maintain multiple monoid states in one rowwise.reduce call (fused dual-reduce, fused mean+M2, ...) must set this False until the split-K partials buffer can hold N states per slot.
• ​computationally_expensive (Bool): Author hint that per-element work is heavy (exp, tanh, sqrt, a normalize step), so per-load address math is in the noise. When True, the inner-axis block tier picks SIMD-full as soon as alignment permits, bypassing the iters_full threshold that amortizes address math for cheap reductions. Softmax / log-softmax / layernorm set this; plain sum / max do not.
• ​BLOCK_SIZE (Int): Threads per block for the inner-axis block tier.
• ​TILED_BLOCK_SIZE (Int): Threads per block for the tiled tier (small by design — more blocks per SM under scattered-write pressure).
• ​COOPERATIVE_BLOCK_SIZE (Int): Threads per block for the non-inner cooperative tier.
• ​WARP_BLOCK_WARPS (Int): Warps per block for the warp tier.
• ​cache_dtype (Optional[DType]): Optional persistent-cache dtype for the warp-tier multi-row scratch (None for no cache).
• ​cache_count (Int): Number of cache entries per row.
• ​num_phases (Int): Per-element-output split-K opt-in for normalize-shaped bodies. <= 1 disables it (comptime-dead, byte-identical codegen). N > 1 = the body has N - 1 dependent row.reduce phases before its final per-element write; on an under-occupied inner-axis shape (num_rows < sm_count, num_rows <= _SPLITK_MAX_ROWS_FOR_SPLIT, and row bytes >= _SPLITK_MIN_ROW_BYTES) the row is split across many blocks via a phase-aware K+1-launch schedule (softmax / log-softmax: N = 3). Otherwise falls through to the warp / block tiers unchanged.
• ​dtype_size (Int): Byte size of the body's primary dtype, used only by the per-element-output split-K byte gate above (0 when the tier is off). Bodies pass size_of[dtype]().

Args:

• ​body (Body): The per-row computation. Receives a Context and the row's Coord; uses reduce / map / once / simd to compose the algorithm.
• ​shape (Coord): Tensor shape.
• ​ctx (DeviceContext): Device context.

Raises:

If the underlying GPU kernel launch fails.