Invisible to grep: a portable check for the silent-softmax bug

The third of three. The first post was the bug — attention combined over separate pieces of the key space without a shared softmax denominator, wrong but plausible, no crash. The second was how we found it in our own code. This one is a check you can run on code you didn't write: why the bug is invisible to search, what the check actually is, and what turned up when we ran it across fifty-two public repositories.

The bug you can't search for

A quick recap, since this post leans on the last two. The failure is combining attention computed over separate pieces of the key space — blocks, ring steps, KV shards — by mixing their independently-normalized outputs instead of carrying one shared softmax denominator across them. Nothing throws; the output looks like an attention output. It's just quietly wrong, and more wrong the more pieces you split into.

What makes it worse than an ordinary bug is that you can't text-search for it. The defect isn't a token you can match — it's a missing rescale, and an absence has no keyword. A correct combiner and a broken one can differ only in whether a running denominator is threaded through the loop; no grep pattern separates them, because the distinguishing feature is a computation that isn't there. It doesn't reliably surface in review either — the broken version looks like exactly the loop you'd expect to see. If you can't find a bug by text and can't reliably eyeball it, what's left is to test for it — and here a test works cleanly.

An invariant you can turn into a test

One property decides whether a combiner is correct, no matter how the code is written: combining the pieces has to reproduce the result of a single softmax taken over the concatenation of the same key/value pieces. That's the whole specification. It doesn't care whether a repo uses a running maximum, a log-sum-exp merge, a Triton kernel, or plain concatenation — only whether the combined answer matches that joint result. One qualification saves a lot of false alarms: the reference has to use the same keys under the same attention semantics — the same scaling, masks, bias, and positional and head-mapping rules, with dropout off. For sparse attention especially, “dense” doesn't mean over every key; it means computing those semantics directly over exactly the pieces the query is allowed to see, rather than through the blocked implementation. In the interface we use, a partial is the pair (out, lse) — the block's own normalized output and its log-denominator (a correct implementation might instead carry a running max and denominator internally; the invariant doesn't care which). The check compares a combiner against that dense joint:

# a "partial" is (out, lse) computed over ONE piece of the key space:
#   out = softmax(scores) @ v         # normalized over its own keys only
#   lse = logsumexp(scores, dim=-1)   # that piece's log-denominator

def joint(q, ks, vs, scale):          # ground truth: ONE softmax over the pieces
    k, v = cat(ks, dim=-2), cat(vs, dim=-2)
    scores = (q @ k.transpose(-2, -1)) * scale
    return softmax(scores, dim=-1) @ v

# D - dense-equivalence: the combiner must reproduce the joint softmax
assert allclose(combine(partials), joint(q, ks, vs, scale))

# the deliberately-broken control: naively summing normalized blocks MUST
# fail D - else the test is vacuous, passing a combiner that ignores lse
assert not allclose(sum(out for out, _ in partials), joint(q, ks, vs, scale))

That second assertion is the one people forget, and it's load-bearing. Without it, a combiner that ignores the denominators entirely could still pass on a lucky test case, and you'd have written a check that certifies the very bug it's meant to catch. The full battery runs four properties — dense-equivalence, commutativity, associativity, and the deliberately-broken naive-sum control — but dense-equivalence guarded by that naive-sum control is what actually does the work.

Two layers: match the shape, then run it

An invariant only helps at scale if you can apply it without hand-wrapping every repository's calling convention — and every repo spells attention differently. So the check has two layers.

The first is static and cheap. It reads the code as a syntax tree, not as text, and looks for the shape of the bug: an output accumulator that gets +='d inside a loop, with a matmul or softmax on the right-hand side, and no sign of a rescale anywhere in the function — no running maximum, no log-sum-exp, no correction-multiply, no call to a known merge routine. This is the part ordinary text search can't do reliably. It checks that the thing being accumulated is an output buffer, that its right-hand side is real attention math, that it sits in a loop, and that a rescale is absent — a conjunction of facts, none of them a searchable string. The two shapes look like this:

# the shape the static pass flags: accumulate in a loop, softmax on the
# right, and NO rescale of the running output anywhere in the function
out = 0
for k_blk, v_blk in blocks:
    scores = (q @ k_blk.transpose(-2, -1)) * scale
    out += softmax(scores, dim=-1) @ v_blk   # sums independently-normalized blocks
return out                                    # exact for 1 block, wrong for more

# a correct combiner keeps a running max / lse and rescales as it goes -
# or concatenates the blocks first and takes ONE softmax over all of them

The second layer only fires on what the first flags. For each suspect function it does the obvious thing the static pass can't: imports the function, feeds it random q, k, v with the block count forced above one, and compares the output against a plain dense-softmax reference evaluated in double precision, with the tolerance sized to the precision of the code under test — a fair bar in float64 is a knife-edge in bfloat16. This bug makes that easy: its error isn't a rounding sliver, it's the output off by multiples of its own magnitude, so a generous tolerance separates it cleanly and a second random input guards against a fluke of the first. The layering keeps the expensive, fragile part — executing someone else's code — pointed only at the handful of functions worth executing.

One case is worth keeping because it shows why the static layer earns its place. A repo had a correction-multiply — the rescale that makes the loop correct — sitting right above the accumulation, but commented out. To grep, and to a quick read, it looks broken. The shape-matcher saw that the rest of the online-softmax machinery was still live and left it alone; reading the function ourselves afterward settled it — a different, still-live rescale was doing the work, so the commented-out line was merely redundant. A keyword search would have raised it as a false alarm; the structural check, backed by our own reading, did not.

What it found

We pointed the pipeline at fifty-two public repositories — canonical ring-attention libraries, block-sparse implementations, sequence-parallel experiments, and a pile of “flash attention from scratch” tutorials, scanned at a fixed point in time. It surfaced one instance, which we then dynamically confirmed.

That one was a grab-bag repository — a personal collection of optimizations — whose flash-attention routine sums per-block softmax outputs with no shared denominator. It is exact when you hand it a single KV block and wrong the moment there are two; in the case we measured, the error climbed to several times the magnitude of the correct output by eight blocks. The dependable part of that signature is the one-block property: a naive sum over a single piece simply is the joint softmax, so the bug is provably invisible there and can only appear once several independently-normalized pieces have to be combined. How fast the error then grows is input-dependent — it needn't climb monotonically — but exactness at one piece, with failure appearing only once several independently-normalized pieces are combined, is a strong signature for this bug class (not a unique one — other block-boundary bugs can also stay quiet at a single piece — but a sharp place to start). The pipeline flagged none of the other fifty-one: they either showed the running-softmax machinery the invariant wants or never split the key space to begin with, tutorials among them — which figures, because in a from-scratch flash tutorial the running softmax is the lesson.

What the count does and doesn't say

This is the section that keeps the rest honest — the same reflex as the last post.

• “One in fifty-two” is not a prevalence rate, and it's worth being exact about why. The dynamic check only ever ran on functions the static pass flagged, and we have never measured that pass's recall — so what we observed is one instance surfaced and confirmed among fifty-two repositories scanned, not one repository in fifty-two contained the bug. There may be missed instances among the other fifty-one. The honest reading is that the bug looks uncommon in the code we scanned; calling it rare in the wild would need either ungated dynamic coverage or a measured recall we don't have. The fix is known — run every attention function against the reference regardless of the static verdict, so coverage stops depending on the triage — we just haven't built it.
• Keep the two halves separate, because they carry very different weight. The dynamic invariant is general: any combiner that gets the joint softmax wrong fails it. The static triage that decides which functions are worth executing is a deliberately narrow heuristic — better called a targeted scanner than a detector. It was calibrated on a dozen hand-checked repositories, only one of which was a real bug. One true positive tells you almost nothing about the scanner's sensitivity or its false-positive rate on genuinely diverse code — so trust the invariant, and treat the triage as a way to spend execution where it's likeliest to pay, not as proof about everything it leaves alone.
• And the triage is shallow on purpose. It keys off ordinary variable names, reads only Python, and looks at one function at a time — so a rescale hidden behind a helper call, an alias, or a custom kernel reads as absent when it isn't, and an accumulator with an unusual name or a kernel in another language slips past entirely. That intraprocedural, name-based heuristic is exactly why its recall is unknown.
• Absence of a flag is not a certificate — the same caveat as the audit. The check tells you where it found the bug; it never tells you the bug is absent.
• The one genuinely interesting pattern is a contrast, and it's under-powered. The scanner surfaced the bug only once across fifty-odd focused, single-purpose repos, while our own exhaustive audit turned it up again and again in exactly one place — a fast-moving library with twenty-odd attention variants written quickly. That library is ours, so the “concentrates in multi-variant libraries” side of the contrast is a sample of one — and note the two sides aren't measured the same way: a thorough audit on one side, a recall-unknown scan on the other. Take it as a hypothesis about where to look — grab-bag collections and bolt-on attention inside large codebases, not dedicated single-purpose repos — not a law we've established.

If you want to run it

• The portable part is the invariant, not our code. Wrap whatever combine routine you're checking and compare its output to a single softmax over the concatenated pieces — reference and combiner sharing the same keys, scaling, masks, bias, and positional and head-mapping rules — and keep the naive-sum control alongside, so the test can't pass by ignoring the denominators.
• Force more than one block. The bug is exact at a single piece; test with at least two and watch whether the error grows as you add more. If it does, this is what you have.
• Read to triage, execute to adjudicate. Structural reading is good for narrowing down where the bug could be; the dense reference is what actually decides. Reading for correctness is the part that fails — a small differential test against dense attention settles it directly, once you exercise the affected multi-piece path, which staring at a hundred lines of loop will not.
• Then encode it as a test and it stops recurring — the same move that closed it in our own code. The check finds the open seam once; the test keeps it shut.

A bug you can't text-search for isn't a bug you're stuck with. Reading tells you where to look; running the code against an answer you already trust is what tells you for sure.