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 trait

OutputWriter

Injected output-writer policy for BlackwellMatmulSM100Kernel.

Carries the regime scalars and a static factory-and-write that resolves the concrete writer from the kernel's config-derived parameters. The kernel takes this as the comptime output_writer_type parameter; the default is StandardOutputWriter (local TMA store, no sync, one peer).

Implemented traits

AnyType

comptime members

needs_sync

comptime needs_sync

Whether the produced writer needs cross-GPU barrier/signal sync around the epilogue (reduce-scatter), False for a local TMA store.

num_peers

comptime num_peers

Number of C TMA descriptors / peer buffers the kernel must supply in c_tma_ops (1 for the standard local store).

Required methods

write_batched

static def write_batched[tma_origin: ImmOrigin, c_type: DType, c_rank: Int, c_tile_shape: IndexList[c_rank], c_desc_shape: IndexList[c_rank], a_type: DType, accum_type: DType, block_tile_shape: IndexList[Int(3)], mma_shape: IndexList[Int(3)], opc: OutputPipelineConfig, c_swizzle: TensorMapSwizzle, transpose_c: Bool, c_smem_dim0: Int, c_smem_dim1: Int, num_output_stages: Int, num_output_warps: Int, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None], elementwise_compute_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]], register_based_epilogue: Bool](c_tma_ops: Pointer[Array[TMATensorTile[c_type, c_rank, c_tile_shape, c_desc_shape], _Self.num_peers], tma_origin], c_tiles: SMemTileArray2DRowMajor[c_type, c_smem_dim0, c_smem_dim1, num_output_stages], stage: OutputStage[opc], tile_coord: Tuple[UInt32, UInt32, UInt32], shape: Tuple[UInt32, UInt32], alpha: Float32 = 1)

Construct the concrete writer from the c_tma_ops array and write one batched (3D-coord) output tile.

The standard policy stores descriptor [0]; the reduce-scatter policy retains all num_peers descriptors. The c_tiles / stage argument types are the shared SMemTileArray2DRowMajor[...] / OutputStage[opc] both writers consume, so the kernel passes its smem.c_tiles() and pipeline stage directly, with no rebind needed.

Parameters:

• ​tma_origin (ImmOrigin): Origin of the C TMA descriptor memory.
• ​c_type (DType): Element DType of the output C tensor.
• ​c_rank (Int): Rank (number of dimensions) of the C tensor.
• ​c_tile_shape (IndexList[c_rank]): Per-dimension tile shape of the C TMA descriptor.
• ​c_desc_shape (IndexList[c_rank]): Per-dimension shape of the C output tensor.
• ​a_type (DType): Element DType of the A operand.
• ​accum_type (DType): Element DType of the MMA accumulator.
• ​block_tile_shape (IndexList[Int(3)]): (BM, BN, BK) block tile shape in elements.
• ​mma_shape (IndexList[Int(3)]): (MMA_M, MMA_N, MMA_K) shape of one MMA instruction.
• ​opc (OutputPipelineConfig): Output pipeline config for TMEM accumulator staging.
• ​c_swizzle (TensorMapSwizzle): Swizzle mode of the C TMA descriptor.
• ​transpose_c (Bool): Whether the C output is stored transposed (A/B swapped).
• ​c_smem_dim0 (Int): First dimension of the C SMEM tile array.
• ​c_smem_dim1 (Int): Second dimension of the C SMEM tile array.
• ​num_output_stages (Int): Number of stages in the C SMEM tile array.
• ​num_output_warps (Int): Number of warps participating in the epilogue.
• ​elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional fused elementwise epilogue that writes the output tile to global memory, or None.
• ​elementwise_compute_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]]): Optional fused elementwise compute applied to the accumulator before store, or None.
• ​register_based_epilogue (Bool): Whether the epilogue runs from register fragments instead of staging through SMEM.

Args:

• ​c_tma_ops (Pointer[Array[TMATensorTile[c_type, c_rank, c_tile_shape, c_desc_shape], _Self.num_peers], tma_origin]): Pointer to the array of C TMA descriptors, one per peer.
• ​c_tiles (SMemTileArray2DRowMajor[c_type, c_smem_dim0, c_smem_dim1, num_output_stages]): Shared-memory C tile array staging the output tile.
• ​stage (OutputStage[opc]): Acquired output pipeline stage holding the TMEM accumulator to read.
• ​tile_coord (Tuple[UInt32, UInt32, UInt32]): (m, n, k_start) block coordinates of the output tile.
• ​shape (Tuple[UInt32, UInt32]): (M, N) problem dimensions for bounds checking.
• ​alpha (Float32): Scaling factor applied to the accumulator (defaults to 1).