News

One softmax, many pieces: the silent bug we kept re-making in long-context attention

A class of long-context attention bug that returns plausible numbers, survives grep and code review, and passes short-sequence tests. Here's why it hides — and why the fix is one place that combines partial softmaxes.

← All news
By Michael Harris
Engineering

A class of attention bug that returns plausible numbers, survives grep and code review, and passes short-sequence tests — and why partial-softmax combination belongs in one canonical implementation and specification, not reinvented inside every attention variant.

How you end up with twenty attention implementations

We maintain a long-context attention library. It started as one thing — dilated attention from the LongNet paper — and grew the way these libraries do: a ring-attention variant for O(n/k) memory, a block-sparse variant, a Hilbert-curve-ordered variant, an SDPA-backed one, a distributed one, Triton kernels for the hot paths. Around twenty implementations, all variations on a theme.

They share a conceptual core, but each was written separately, at a different time, often in a hurry. That detail turns out to be the whole story.

The audit, and the one bug that kept showing up

We ran a systematic correctness audit of the core library — every long-sequence implementation read for numerical correctness, with executable comparisons against a dense reference where we could run them. Roughly thirty findings survived adversarial review. But they weren't thirty different bugs. The single most common class, recurring in at least six separate implementations — the SDPA ring path, the ring autograd function, the Triton/Hilbert kernels, and all four block-sparse variants — was the same conceptual mistake:

Combining attention computed over pieces of the key space without a shared softmax denominator.

Some summed independently-normalized block outputs. Some divided by the running log-sum-exp instead of its exponential. Some averaged the per-step outputs across the ring. Different lines of code, same error — and every one of them produced output that was silently and materially wrong. Measured against a dense reference — relative error meaning the deviation as a fraction of the reference's own magnitude — the SDPA ring path ran 36–50% off, and the default block-sparse path differed from the correct joint softmax by a max-absolute difference of up to ~11.

The thing we should have seen: it's all one operation

Here's the insight that only becomes obvious once you've made the mistake six times. Every one of these “variants” is doing the same thing — attention over a subset of keys, then combining subsets:

  • Block-sparse — the pieces are selected key blocks.
  • Dilated (LongNet) — the pieces are strided key subsets.
  • Flash attention — the pieces are key tiles staged through on-chip memory; online softmax is what combines those tiles exactly.
  • Ring / sequence-parallel — the pieces are key shards on different devices.

The correct way to combine two partial results — each a normalized output out = softmax(scores) @ v over that piece's own keys, plus its log-denominator lse = logsumexp(scores) — is one associative, commutative merge (FlashAttention's running-max / running-denominator trick). Concretely, merging partials (out_a, lse_a) and (out_b, lse_b):

m    = max(lse_a, lse_b)                  # subtract a shared max, for stability
a, b = exp(lse_a - m), exp(lse_b - m)     # each piece's un-normalized weight
out  = (a * out_a + b * out_b) / (a + b)  # ONE shared denominator: a + b
lse  = m + log(a + b)                      # the combined log-denominator

The whole game is that shared denominator a + b. Each block has already divided by its own denominator, so you can't just add the outputs back up — each one has to be re-weighted by its share of the joint softmax mass, a / (a + b) above. That's what the lse values are for: a and b are the blocks' partition functions (rescaled by the shared max m), and dropping them counts every block as if it owned the entire softmax. The merge is the same operation whether the pieces are blocks in a loop or shards on different GPUs, and it's associative and commutative in exact arithmetic, so order and grouping don't change the result.

Which means there should be exactly one canonical merge — one implementation wherever the execution layer allows, one specification everywhere else — and every “variant” should be nothing more than a description of which pieces a query attends to. Instead we had the combination logic re-hand-rolled in each implementation — which is to say, six independent opportunities to get it wrong. We took all six.

Why this bug is invisible

Here's the part that should worry you most: this wasn't code we'd just written. The variants went in across 2025, and we'd been shipping the violated invariant for months before the audit surfaced it in May 2026. It lasted that long because it defeats every cheap way you'd normally catch a bug:

  • It doesn't crash. No exception, usually no NaN. You get a tensor of the right shape, full of reasonable-looking numbers. Nothing draws your eye.
  • grep and code search can't see it. The bug is a missing rescale — the absence of an operation, and there's no keyword for a line that isn't there. Text search finds the accumulation sites easily; what it can't do is tell a correct online-softmax merge from one missing its denominator, because the difference is semantic. (It doesn't help that the tokens that would localize it, += and @, are the ones code search tends to discard.)
  • It passes the obvious test. The math is exact when there's only one piece — one block, one shard — so any test that exercises a single piece passes clean. In our setup those were the short-sequence tests. The error only appears once you combine two or more pieces, which is precisely the long-context regime the library exists for. Our short-sequence tests were green.
  • It fools a quick read — in both directions. While looking at other people's code we found one “flash attention” that computes a per-block softmax and sums the outputs with no rescale; it looks fine and is broken — exact for one block, then 106%, 293%, 620% off in relative error at 2, 4, 8 blocks. We found another with the rescaling line commented out right before the accumulation; it looks broken and is correct, because it uses the other valid formulation two lines up. You often can't tell from the accumulation line alone — you need the full derivation across the surrounding code, or, more reliably, an executable comparison against a dense reference.

That last point is the practical takeaway: for this class of bug, an executable check against ground truth isn't a nice-to-have — it's the most reliable arbiter, the check I'd no longer ship this code without.

The fix is boring, which is the point

We wrote one primitive — merge_attention_partials(acc_out, acc_lse, blk_out, blk_lse) — the single canonical place softmax combination is allowed to live. It's ~15 lines, associative and commutative in exact arithmetic (and invariant within floating-point tolerance under reordering), so the same merge serves a single-GPU block loop and a multi-node ring reduction. The point is architectural: one canonical merge with one specification, so that at the semantic level every variant is just a descriptor of which key pieces a query sees, and backend code implements that same tested contract. Routing the existing variants onto it is in progress — the block-sparse path already calls it; the ring and Hilbert accumulators are being migrated off their bespoke copies. (The hand-written Triton kernels can't call a Python helper, so they stay a separate implementation held to the same specification rather than calling this one.) We tested the primitive as a specification, not a smoke test: equivalence to a dense joint softmax, order- and grouping-invariance within tolerance, forward-and-backward gradcheck, and — importantly — a deliberately broken control asserting that the naive “sum of normalized blocks” fails, so the test can't pass vacuously.

Is this everyone's problem? No — and that's the real lesson

We wondered whether we'd stumbled onto a widespread issue, so we checked. We ran an executable correctness oracle over ~50 public attention implementations — ring, block-sparse, sliding-window, from-scratch flash, LongNet reimplementations. The result was humbling in a useful way: it's rare. Focused, single-purpose implementations overwhelmingly get this right. The canonical ring-attention libraries are correct. Even most “flash attention from scratch” tutorials are correct — unsurprisingly, since online softmax is the lesson they're teaching. Across the whole external sample we found the mistake in the wild essentially once. Treat that as a convenience sample, not a prevalence study — repositories we could find and run, on the configurations the oracle covered — so it means “rare in the code we checked,” not a population estimate.

So the honest conclusion isn't “the field is full of broken attention.” It's narrower and more actionable:

The risk isn't writing attention. It's writing attention twenty times. Every hand-rolled re-combination of partial softmaxes is a fresh chance to silently violate an invariant that has exactly one correct form. Breadth of variants concentrates this risk — and the mitigation is to make the invariant un-re-implementable: one merge, one test, everything else a descriptor.

We made the mistake because we had twenty front doors and one lock we kept re-cutting by hand. The fix wasn't cleverness. It was refusing to cut the lock more than once.

Takeaways

  • If your codebase combines attention over pieces (blocks, shards, tiles, strides), route all of it through one online-softmax merge. Don't let variants re-implement combination.
  • Test that merge as a specification — dense-equivalence, order-invariance, gradcheck, and a deliberately broken control that the naive version fails.
  • Distrust short-sequence smoke tests for long-context code: the combination bug is exact at one piece and only appears at two or more.
  • For silent numerical bugs, a careful read isn't enough on its own — run it against a dense reference.