Mojo struct
OutputKPipeline
struct OutputKPipeline[origin: MutOrigin, opc: OutputPipelineConfig]
Per-K-iteration view of OutputTilePipeline.
Unlike standard producer()/consumer() which signal once per tile (after all K iterations), this view signals after each K iteration. Use for kernels with per-K accumulation patterns (e.g., blockwise FP8).
Example (MMA warp): for i in range(num_iters): with mma_ctx.output_pipeline.per_k().produce() as stage: mma(stage.tmem, ...) # exit signals mma_arrive for this K iteration
Example (Epilogue warp): for k_iter in range(num_iters): with epi_ctx.output_pipeline.per_k().consume() as stage: promote(stage.tmem, ...) # exit signals consumer_step for this K iteration
Fields
- pipeline_ptr (
Pointer[OutputTilePipeline[opc], origin]):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
TilePipelineType
comptime TilePipelineType = OutputTilePipeline[opc]
Methods
__init__
__init__(pipeline_ptr: Pointer[OutputTilePipeline[opc], origin]) -> Self
produce
produce(self) -> MmaKStage[origin, opc]
Get MMA stage context manager for one K iteration.
Returns:
MmaKStage[origin, opc]: Context manager that acquires stage on enter and signals
mma_arrive on exit.
consume
consume(self) -> PerKConsumerStage[origin, opc]
Get consumer context manager for one K iteration.
Returns:
PerKConsumerStage[origin, opc]: Context manager that waits for MMA on enter and signals
consumer_step on exit.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!