Mojo module

tile_pipeline

Tile pipeline for SM100 producer-consumer synchronization.

Provides staged tile storage with producer-consumer barrier synchronization for TMA-MMA pipeline coordination. All barrier operations are encapsulated in context managers for safety and clarity.

Key Abstractions

InputTilePipeline[Payload]: Generic pipeline with payload abstraction
TilePipeline: Standard pipeline with explicit A/B tile types
OutputTilePipeline: TMEM accumulator stages for MMA→Epilogue pipeline

Naming Conventions

*Pipeline: Multi-stage buffer (InputTilePipeline, OutputTilePipeline)
*Producer/*Consumer: Role handles (InputProducer, OutputConsumer)
acquire(): Context manager to get one pipeline stage

Context Manager Semantics

Each with block handles barrier synchronization automatically:

with producer.acquire() as tiles:   # BLOCKS until consumer releases stage
    load_tiles(tiles)                # safe to write
                                     # EXIT: signals producer barrier, advances

with consumer.acquire() as tiles:   # BLOCKS until producer fills stage
    use_tiles(tiles)                 # safe to read
                                     # EXIT: signals consumer barrier, advances

Example: TMA Load Warp (Producer)

with input_pipeline.producer() as producer:  # producer role for this warp
    while work_iter.has_work():
        with work_iter.next() as current:
            for i in range(num_iters):
                with producer.acquire() as tiles:  # waits for consumer
                    tma_load(tiles.a_tile(), tiles.b_tile())
    producer.drain()  # wait for all stages consumed before CTA exits

Example: MMA Warp (Consumer + Output Producer)

with mma_ctx:  # TMEM lifecycle
    while work_iter.has_work():
        with work_iter.wait_and_advance():  # blocks on CLC response
            with output_pipeline.producer() as output_stage:  # waits for epilogue
                with input_pipeline.consumer() as consumer:
                    for i in range(num_iters):
                        with consumer.acquire() as input_tiles:  # waits for TMA
                            mma(output_stage.tmem, input_tiles)

Example: Epilogue Warp (Output Consumer)

with epi_ctx:  # signals TMEM dealloc on exit
    while work_iter.has_work():
        with work_iter.next() as current:
            with output_pipeline.consumer() as output_stage:  # waits for MMA
                write_output(output_stage)

`comptime` values

`MbarPtr`

comptime MbarPtr = LegacyUnsafePointer[SharedMemBarrier, address_space=AddressSpace.SHARED]

Structs

BlockScaledTilePayload: Tile payload for block-scaled matmul (A, B, SFA, SFB tiles).
BlockwiseFP8TilePayload: Tile payload for blockwise FP8 matmul (A, B, A-scales tiles).
EpilogueKContext: Per-K context manager for epilogue warp in blockwise FP8.
EpilogueKStage: Per-K stage for epilogue warp in blockwise FP8.
InputConsumer: Consumer view for MMA warp. Use acquire() to get stages.
InputConsumerStage: Consumer stage context manager. Released on scope exit.
InputProducer: Producer view for TMA Load warp. Use acquire() to get stages.
InputProducerStage: Producer stage context manager. Released on scope exit.
InputTilePipeline: Tile pipeline with configurable payload type.
MmaKStage: Per-K stage context for MMA warp in blockwise FP8.
OutputConsumer: Consumer view for epilogue warp (output pipeline).
OutputKPipeline: Per-K-iteration view of OutputTilePipeline.
OutputProducer: Producer view for MMA warp (output pipeline).
OutputStage: Acquired output stage with TMEM handle and pipeline reference.
OutputTilePipeline: Pipeline for MMA→Epilogue TMEM stage synchronization.
PerKConsumerStage: Context manager for per-K epilogue consumption.
StandardConsumerStage: Context manager for consumer tile access with encapsulated stage indexing.
StandardProducerStage: Context manager for producer tile access with encapsulated stage indexing.
StandardTileConsumer: Consumer view for MMA warp (standard tile pipeline).
StandardTilePayload: Tile payload for standard matmul (A and B tiles).
StandardTileProducer: Producer view for TMA Load warp (standard tile pipeline).
TilePipeline: Staged tile storage with producer-consumer synchronization for SM100.

Traits

TilePayload: Trait for tile payload types. Must be @register_passable("trivial").

Key Abstractions​

Naming Conventions​

Context Manager Semantics​

Example: TMA Load Warp (Producer)​

Example: MMA Warp (Consumer + Output Producer)​

Example: Epilogue Warp (Output Consumer)​

comptime values​

MbarPtr​

Structs​

Traits​