Online Softmax Merge

Problem

Merge two partial online-softmax summaries, each given as a running maximum m and running denominator l: m_new = max(m1, m2) and l_new = l1 * exp(m1 - m_new) + l2 * exp(m2 - m_new).

The signature is online_softmax_merge(m1, l1, m2, l2). Return [m_new, l_new]. This is how flash decoding combines partial attention tiles.

Hint

Only the summary with the smaller maximum needs rescaling, because the larger one already has the winning max.

Test cases

case 1

in: [0,1,0,1]

expected: [0,2]

case 2

in: [1,1,0,1]

expected: [1,1.3678794412]

case 3

in: [5,0.5,1,0.5]

expected: [5,0.5091578194]

case 4

in: [-1000000000,1,0,0]

expected: [0,0]

Starter codePython

import math
def online_softmax_merge(m1, l1, m2, l2):
    # Your code here
    pass

4 test cases · run your code against them in the editor