Mojo struct
EpilogueKContext
@register_passable(trivial)
struct EpilogueKContext[origin: MutOrigin, input_origin: MutOrigin, num_output_stages: Int, stage_stride_cols: Int, cta_group: Int, 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[EpilogueKContext[origin, input_origin, num_output_stages, stage_stride_cols, cta_group, num_input_stages].OutputPipelineType, origin]): - input_pipeline_ptr (
Pointer[EpilogueKContext[origin, input_origin, num_output_stages, stage_stride_cols, cta_group, num_input_stages].InputPipelineType, input_origin]): - output_stage (
EpilogueKContext[origin, input_origin, num_output_stages, stage_stride_cols, cta_group, num_input_stages].OutputStageType): - input_stage_index (
UInt32):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable
comptime members
__copyinit__is_trivial
comptime __copyinit__is_trivial = True
__del__is_trivial
comptime __del__is_trivial = True
__moveinit__is_trivial
comptime __moveinit__is_trivial = True
CombinedStageType
comptime CombinedStageType = EpilogueKStage[num_output_stages, stage_stride_cols, cta_group, num_input_stages]
InputPipelineType
comptime InputPipelineType = ProducerConsumerPipeline[num_input_stages]
OutputPipelineType
comptime OutputPipelineType = OutputTilePipeline[num_output_stages, stage_stride_cols, cta_group]
OutputStageType
comptime OutputStageType = OutputStage[num_output_stages, stage_stride_cols, cta_group]
Methods
__init__
__init__(output_pipeline_ptr: Pointer[EpilogueKContext[origin, input_origin, num_output_stages, stage_stride_cols, cta_group, num_input_stages].OutputPipelineType, origin], input_pipeline_ptr: Pointer[EpilogueKContext[origin, input_origin, num_output_stages, stage_stride_cols, cta_group, num_input_stages].InputPipelineType, input_origin]) -> Self
__enter__
__enter__(mut self) -> EpilogueKContext[origin, input_origin, num_output_stages, stage_stride_cols, cta_group, num_input_stages].CombinedStageType
Returns:
EpilogueKContext
__exit__
__exit__(mut self)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!