An Algebraic Interface for Backpropagation through Fused Map-Folds

Memory-efficient neural kernels often have the structure of a fused map-fold: a large logical interaction space is mapped to monoidal factors and immediately folded, avoiding materialization of tensors such as attention logits. While this explains the forward pass of kernels such as FlashAttention and fused cross entropy, their backward passes are usually derived and implemented separately. We identify a sufficient algebraic condition, the Local Gradient Property, under which the reverse pass of a map-fold can be computed with the same fused structure: local factors are recomputed tile by tile, and their cotangents are recovered from compact saved fold state. We extend this condition with a finalize-embed factorization that exposes a compact algebraic backward interface, and cover both commutative folds and ordered non-commutative folds with prefix state. We present CuTileReduce, an implementation of the interface built on CuTile, and instantiate it for cross entropy, attention, and a non-commutative attention variant, demonstrating correct memory-efficient forward and backward kernels generated from the same map-fold abstraction.