Implement the split-KV merge: combine partial attention outputs with their log-sum-exp

mediumFree drillimplement merge_partials

Before you code

Two attention partitions have different log-sum-exp values. Would averaging their normalized outputs preserve the result? Predict which partition should carry more weight before writing the merge.

Review the prerequisite concept

The concept page shows its own access level. The statement and visible tests below are free.

A decode attention kernel that splits one long context across thread blocks produces, per split $s$ and query head $h$ , a partial output $o_{s,h}$ already normalized within that split, and the split's log-sum-exp $L_{s,h} = m_{s,h} + \ln \ell_{s,h}$ of its scaled scores. Write the reduction that turns those partials into the exact output over the whole context.

merge_partials(outs, lses) -> numpy array [H, d]
  outs  [S, H, d]  partial outputs, each normalized within its split
  lses  [S, H]     log-sum-exp of each split's scores; -inf for an empty split
                   (an empty split's outs row is all zeros)

At least one split per head is non-empty. Scores can be large: a log-sum-exp of 1,000 is valid input, and the merge must not overflow in float64. The result must not depend on the order of the splits.

The visible tests compare against attention computed in one pass. The full suite adds empty splits, extreme log-sum-exps, a single split and reordered splits.

YOUR IMPLEMENTATION · PYTHON 3.14 + NUMPY

Free drill: the visible tests run for everyone, the edge cases and the solution for any signed-in reader.

The visible tests

import numpy as np


def _partials(scores, values, bounds):
    """Per-split normalized outputs and log-sum-exps for one head.
    scores: [T], values: [T, d]."""
    outs, lses = [], []
    for lo, hi in zip(bounds, bounds[1:]):
        s = scores[lo:hi]
        if s.size == 0:
            outs.append(np.zeros(values.shape[1]))
            lses.append(-np.inf)
            continue
        m = s.max()
        p = np.exp(s - m)
        outs.append(p @ values[lo:hi] / p.sum())
        lses.append(m + np.log(p.sum()))
    return np.array(outs), np.array(lses)


def _case(seed, T, H, d, bounds, scale=3.0):
    rng = np.random.default_rng(seed)
    scores = rng.normal(scale=scale, size=(H, T))
    values = rng.normal(size=(H, T, d))
    outs = np.zeros((len(bounds) - 1, H, d))
    lses = np.zeros((len(bounds) - 1, H))
    full = np.zeros((H, d))
    for h in range(H):
        o, l = _partials(scores[h], values[h], bounds)
        outs[:, h], lses[:, h] = o, l
        p = np.exp(scores[h] - scores[h].max())
        full[h] = p @ values[h] / p.sum()
    return outs, lses, full


def _check(ns, outs, lses, full):
    got = np.asarray(ns["merge_partials"](outs, lses))
    assert got.shape == full.shape, f"expected shape {full.shape}, got {got.shape}"
    err = np.abs(got - full).max()
    assert err < 1e-10, f"largest error {err:.3g} against one-pass attention"


def test_two_equal_splits_match_full_attention(ns):
    _check(ns, *_case(1, 64, 2, 8, [0, 32, 64]))


def test_uneven_splits_match_full_attention(ns):
    _check(ns, *_case(2, 1000, 4, 16, [0, 7, 300, 301, 1000]))

The reasoning behind this component: the companion page.