M2: the scaled accumulator as a monoid action, and fold pipelines as fusion (ideas from Haskell) #50

🪭 Claude (einfold)

Two patterns from Haskell that fit milestone M2, demo 1's FlashAttention from SQL (design §9.6 and §16). They came out of the maintainer's question on #43 about what Haskell could teach this IR. Neither changes the design. Both make its correctness easier to state and test.

1. The scaled exponential accumulator is a monoid action

Design §8.6 item 6 holds a sum of exponentials as a pair (m, Σ e^(s − m)·y): a running maximum, and a sum scaled by it. Merging two pairs raises m to the larger maximum and rescales the other sum by e^(m_old − m_new). This is the online softmax.

Brent Yorgey's functional pearl "Monoids: Theme and Variations" names the structure. The maxima form a monoid (max, with identity −∞) that acts on the sums by rescaling, and the pair is the semidirect product of the two. Writing it that way has two uses:

• Correctness by construction. The pair's merge is associative and has an identity, which the kernel needs in order to merge tiles in any grouping. That follows from two smaller laws, each easy to property-test on its own, special floats included:

  • the action respects the monoid: rescaling by a then by b is rescaling by max(a, b);
  • rescaling distributes over +.

  Today we would test the pair's associativity directly, as one law over four-tuples.

• Reuse. Any fold "normalized by a running statistic" has the same shape: log-sum-exp losses, the Nadaraya–Watson control in demos.md, and running MAX with values attached. In Describe a fold by its aggregate, outputs and domain #43's IR it's one more Monoid part.

2. Fold pipelines are fusion, in the sense of deforestation

Fusing a producer fold into its consumer, so that attention's N² scores never exist, is what functional compilers call deforestation (Wadler, 1988). GHC does it with short-cut fusion and stream fusion (Coutts, Leshchinskiy and Stewart, "Stream Fusion: From Lists to Streams to Nothing at All", ICFP 2007). Their lesson for einfold: fusion is reliable when producers and consumers are written in a small set of combinators whose fusion rules are equations. GHC's RULES are rewrite rules, applied the way egglog would apply ours.

What that means for M2:

• Keep pipelines a visible shape in the IR. Make "a fold whose operand is a map over another fold's result" a node of its own, rather than recovering it from plans each time.
• State the fusion rule as an equation. Write §9.6's fusion condition (the producer's groups are the consumer's groups paired with values of a summed dimension) as a rewrite with its side condition, so the same rule fires on the Nadaraya–Watson control without anything attention-specific.

Also worth reading

• Stephen Dolan, "Fun with semirings" (ICFP 2013): one closed-semiring type class gives shortest paths, transitive closure and more from the same matrix code. It's the model for M5's semirings beyond sum-product, with laws attached to each.

References:

• Yorgey, Monoids: Theme and Variations (Functional Pearl)
• Coutts, Leshchinskiy and Stewart, Stream Fusion: From Lists to Streams to Nothing at All
• Dolan, Fun with semirings