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 struct

StandardOutputWriter

struct StandardOutputWriter

Default OutputWriter policy: local TMA store via TileWriter.

One peer, no cross-GPU synchronization. This is the writer policy BlackwellMatmulSM100Kernel uses unless a reduce-scatter policy is injected. Target hardware: SM100 (B200).

Implemented traits

AnyType, Deinitable, Movable, OutputWriter

comptime members

needs_sync

comptime needs_sync = False

num_peers

comptime num_peers = 1

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], Int(1)], 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)

Local TMA store of one batched output tile (uses descriptor [0]).

Parameters:

• ​tma_origin (ImmOrigin): Memory origin of the TMA descriptor pointer (inferred).
• ​c_type (DType): Element dtype of the C output tensor (inferred).
• ​c_rank (Int): Rank of the C output tensor (inferred).
• ​c_tile_shape (IndexList[c_rank]): Per-tile shape of the C output (inferred).
• ​c_desc_shape (IndexList[c_rank]): TMA descriptor shape for C (inferred).
• ​a_type (DType): Element dtype of the A input matrix.
• ​accum_type (DType): Accumulator dtype stored in TMEM.
• ​block_tile_shape (IndexList[Int(3)]): Block tile shape as (BM, BN, BK).
• ​mma_shape (IndexList[Int(3)]): MMA instruction shape as (MMA_M, MMA_N, MMA_K).
• ​opc (OutputPipelineConfig): Output pipeline config bundling accumulator stages, stage stride, and CTA group.
• ​c_swizzle (TensorMapSwizzle): TMA swizzle pattern for the C SMEM layout.
• ​transpose_c (Bool): Whether C is stored transposed.
• ​c_smem_dim0 (Int): Row dimension of the C SMEM tile.
• ​c_smem_dim1 (Int): Column dimension of the C SMEM tile.
• ​num_output_stages (Int): Number of C SMEM pipeline stages.
• ​num_output_warps (Int): Number of warps driving the output pipeline.
• ​elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional elementwise epilogue applied to fragments before the store.
• ​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 compute epilogue fused into the register path.
• ​register_based_epilogue (Bool): Whether the compute epilogue runs in registers (true) or SMEM (false).

Args:

• ​c_tma_ops (Pointer[Array[TMATensorTile[c_type, c_rank, c_tile_shape, c_desc_shape], Int(1)], tma_origin]): Pointer to the array of TMA store descriptors for C.
• ​c_tiles (SMemTileArray2DRowMajor[c_type, c_smem_dim0, c_smem_dim1, num_output_stages]): SMEM tile array for the C output.
• ​stage (OutputStage[opc]): OutputStage with pipeline, index, and TMEM handle.
• ​tile_coord (Tuple[UInt32, UInt32, UInt32]): (m_tile, n_tile, batch) tile coordinates.
• ​shape (Tuple[UInt32, UInt32]): (M, N) problem dimensions.
• ​alpha (Float32): Scalar applied to fragments before the store (defaults to 1.0).