For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Mojo function
zero_output_epilogue
def zero_output_epilogue[c_type: DType, /, *, c_smem_layout: Layout, block_tile_shape: IndexList[Int(3)], mma_shape: IndexList[Int(3)], c_static_N: Int, c_swizzle: TensorMapSwizzle, cta_group: Int, num_output_warps: Int, elementwise_lambda_fn: def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None, transpose_c: Bool](work_tile_coord: Tuple[Int, Int], group_end_idx: UInt32)
Zero an inactive group's output THROUGH the elementwise epilogue.
Mirror of the epilogue (unaligned/elementwise_lambda_fn) branch of
multi_stage_store_C, but with a zero accumulator: for an inactive group
(expert_id < 0) there is no MMA, so this walks the same output-tile
coordinate mapping and calls elementwise_lambda_fn with a zero vector at
each (n, m). It deliberately does NOT wait on accum_full_mbar or read
TMEM: the main-load and MMA warps skip inactive groups entirely (they never
arrive on that barrier), so waiting would deadlock. Matches the naive
grouped matmul, which likewise runs the epilogue with accum == 0 for
expert == -1 instead of storing zeros directly to C -- so callers whose
real output lives behind the epilogue (e.g. the LoRA-B QKV expand's
route_qkv) get their -1 tokens zeroed without a separate memset.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!