Mojo struct
EpilogueKContext
struct EpilogueKContext[origin: MutOrigin, input_origin: MutOrigin, opc: OutputPipelineConfig, num_input_stages: Int]
Per-K context manager for epilogue warp in blockwise FP8.
Bundles output pipeline (MMA→Epilogue sync) and input pipeline (A-scales) into a single context manager for clean per-K iteration handling.
Example usage: for k_iter in range(num_iters): with epi_ctx.per_k_stage(input_pipeline) as epi_stage: accum.promote(epi_stage, ...) # exit signals BOTH pipelines
enter: Waits for MMA to complete this K iteration, returns EpilogueKStage exit: Signals both output consumer barrier AND input consumer_step
Fields
- output_pipeline_ptr (
Pointer[OutputTilePipeline[opc], origin]): - input_pipeline_ptr (
Pointer[ProducerConsumerPipeline[num_input_stages], input_origin]): - output_stage (
EpilogueKContext[origin, input_origin, opc, num_input_stages].OutputStageType): - input_stage_index (
UInt32):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
CombinedStageType
comptime CombinedStageType = EpilogueKStage[opc, num_input_stages]
cta_group
comptime cta_group = opc.cta_group
InputPipelineType
comptime InputPipelineType = ProducerConsumerPipeline[num_input_stages]
num_output_stages
comptime num_output_stages = opc.num_stages
OutputPipelineType
comptime OutputPipelineType = OutputTilePipeline[opc]
OutputStageType
comptime OutputStageType = OutputStage[opc]
Methods
__init__
__init__(output_pipeline_ptr: Pointer[OutputTilePipeline[opc], origin], input_pipeline_ptr: Pointer[ProducerConsumerPipeline[num_input_stages], input_origin]) -> Self
__enter__
__enter__(mut self) -> EpilogueKContext[origin, input_origin, opc, num_input_stages].CombinedStageType
Returns:
EpilogueKContext[origin, input_origin, opc, num_input_stages].CombinedStageType
__exit__
__exit__(mut self)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!