The Bug Was One Axis — And It Only Fires When the Fast Kernels Aren't Installed
Two released hybrid models leak future information across chunk boundaries. The root cause is a single reduction over the wrong axis. And the execution path it lives on is the one your CI is almost certainly using.
Paper: arXiv:2608.22876 (v2)
Start with the line that should worry you
A model can pass every fused-kernel test and still produce leaking logits the moment it is run without the kernels — e.g. on CPU or in CI.
That sentence is the practical core of this work, and it took a source-level census to earn it. Here is how we got there.
The check that stopped working
An autoregressive model is only meaningful under one constraint: the representation at position t depends on positions ≤ t and nothing else. We call it prefix invariance.
For a decoder-only Transformer, the mask enforces it, so people inspected the mask and called that an audit. That worked while attention was the only temporal mixer.
It does not work now. Modern stacks interleave attention with linear recurrences, state-space scans, sliding windows, and convolutions. A scan has no mask. Causality in these models is a property of the whole computation graph, not of one attribute on one module — and leaks can arrive through scans, aggregations, or normalization while every mask in the model is perfectly correct.
Our audit tests the property directly instead of its proxy:
1 x₁ ← random token sequence of length T
2 x₂ ← copy of x₁ with position T−1 replaced by a different token
3 attach a forward hook to every layer
4 h₁ ← f(x₁) caching disabled
5 h₂ ← f(x₂) caching disabled
6 Δℓ ← max |h₁ℓ[0:T−1] − h₂ℓ[0:T−1]| for each layer ℓ
7 ℓ* ← min{ℓ : Δℓ > τ}
8 return (LEAK, ℓ*) if ℓ* exists else CLEAN
Two forward passes. No gradients, no labels, no accelerator. Against injected faults it is decisive: 192 injected faults across 8 checkpoints — mask inspection detected 0, the per-layer audit localized 192/192 to the exact layer. A logits-only variant detects all 192 and localizes none, which is why the hooks are not optional.
Reading the source before running anything
Auditing checkpoints one by one is a search. Reading the implementation is a prediction.
The inter-chunk recurrence of a Mamba-2-style chunked scan is a small, recognizable computation, and the reference implementation in transformers fixes a specific orientation of the input- and output-chunk axes. So we ran a static census over every chunked-scan implementation shipped in transformers 5.7.0:
| Implementation | Inter-chunk axis handling | |
|---|---|---|
modeling_mamba2.py |
transpose(1,3) → reduce over input chunk axis |
reference |
modeling_bamba.py |
matches reference | conformant |
modeling_falcon_h1.py |
matches reference | conformant |
modeling_granitemoehybrid.py |
matches reference | conformant |
modeling_zamba2.py |
permute → reduce over output chunk axis |
departs |
modeling_nemotron_h.py |
permute → reduce over output chunk axis |
departs |
That is the whole defect. One reduction, one axis.
It is worth sitting with how ordinary this is. Nobody wrote something exotic. Two implementations transposed a tensor a slightly different way than the reference did, and the resulting recurrence carries information backward across a chunk boundary. There is no exception, no NaN, no shape error. The shapes all line up — that is precisely why it survives.
The prediction, tested
A static census is a hypothesis. We then audited the models dynamically, at sequence lengths chosen to exceed each model's declared chunk or window parameter.
| Checkpoint | Declared parameter | Verdict |
|---|---|---|
| Zamba2-1.2B | chunk size 256 | leak, onset = 256 |
| Nemotron-H-8B | chunk size 128 | leak, onset = 128 |
| Bamba-9B | mamba chunk size 256 | clean |
| Falcon-Mamba-7B | conv kernel 4 | clean |
| Falcon-H1-1.5B | mamba chunk size 128 | clean |
| Granite-4.0-H-Tiny | mamba chunk size 256 | clean |
| Gemma-3-1B | sliding window 512 | clean at T = 1536 |
| RecurrentGemma-2B | sliding window 2048 | clean at T = 3072 |
| Gemma-2-2B | sliding window 4096 | clean |
| Mamba2-130M / 370M | chunk size 256 | clean |
| RWKV-6-1.6B, Qwen3-0.6B, LFM2-1.2B, ZR1-1.5B | — | clean |
The leak onset equals the declared chunk size, exactly, in both cases. Not approximately. That is not a number you argue about — it is a fingerprint pointing at a specific line.
And the negative result matters as much: every implementation the census called conformant came back clean. The two the census called out are the two that leaked.
🔴 Now the part that changes what you should do on Monday
The defect lives in the chunked scan written in PyTorch (torch_forward / segment_sum).
That path executes whenever the optional fused kernels — mamba_ssm, causal_conv1d — are absent. Which is:
- all CPU execution
- any GPU environment without those packages installed
- a stock
transformersinstall that has not added them
Those are separate dependencies that are famously difficult to build. In our environment they refuse to compile against the current torch. This is the path a large fraction of evaluation, research, reproduction, and CI workflows actually take.
We scope the claim precisely, and we want to be equally precise about what we did not test: an install with the fused CUDA kernels present dispatches to a different implementation, which we could not audit because those kernels do not build here. Whether the fast path shares the defect is untested and open.
But notice that the caveat cuts both ways, and the second direction is the uncomfortable one:
A model can pass every fused-kernel test and still produce leaking logits the moment it is run without the kernels.
If your evaluation harness runs on CPU — and many do — or your CI container skips the optional CUDA extensions — and most do — you may be measuring a model that is not causal, and every number that comes out of it is contaminated in the direction that flatters it.
Because that is the other thing about causal leakage: it does not crash. It lowers training loss and perplexity. Future information makes next-token prediction artificially easier. Your dashboard says the change worked. The failure only surfaces during free-running generation or deployment, by which time the architecture comparisons that justified your design are already poisoned.
Two auditing norms we learned the hard way
1. A CLEAN verdict means nothing without a positive control on the same loaded checkpoint.
We hit checkpoints that produced bit-identical outputs for different inputs: full logit difference 0.0, difference at the perturbed position 0.0, while still emitting plausibly-scaled, position-varying logits with no NaNs. Under that load, Δ = 0 everywhere is guaranteed regardless of whether the architecture is causal. The audit reports CLEAN and the report is worthless.
So we distinguish gated clean verdicts — where a known fault was injected into that exact loaded checkpoint and the detector caught it — from provisional ones, where the harness could not inject (some architectures' layer return signatures do not accept it). We label them differently in the paper, and we applied the same standard to our own released model rather than only to other people's.
2. Your audit length must exceed the architecture's chunk, kernel, or window parameter.
Our default was T = 48. At T = 48, Zamba2 is clean. Its chunk is 256, so the inter-chunk recurrence — the thing that is broken — is never exercised at all. A chunked scan over a sequence shorter than one chunk degenerates to a single chunk with no inter-chunk step; a sliding window wider than the sequence degenerates to full attention. An audit below those thresholds cannot observe a whole class of defects and will confidently report the model as fine.
To check that the positive findings weren't just an artifact of longer sequences, we re-audited two clean models at three times their own window width — Gemma-3-1B at T=1536 and RecurrentGemma-2B at T=3072 — and both still returned exact-zero deltas.
This generalizes past causality: your test has to exercise the path the bug lives on, and architectural constants tell you where those paths begin.
Why there is no repository
We publish the method and release the measurements, not a package. That is deliberate, and Appendix A.1 says why.
The method is five lines of arithmetic on top of standard forward hooks. A competent reader can reimplement it in about an hour against any model whose layers are enumerable — and an independent reimplementation is a stronger reproduction than running our binary, because it is independent. Running someone's code and getting their answer is weak evidence. Reading a spec, building your own, and getting the same answer is strong evidence.
So what we release is the part that is harder to wave away: complete audit logs, exact checkpoint identifiers, per-layer delta arrays for every clean scan and every injected trial, the injected-fault specifications, and the full environment description. Everything needed to reproduce — or contest — any number in the paper.
What we are asking for
Every model card states how many parameters, how long a context, how many benchmark points. None of them state whether the released implementation is causally correct — and after this census we no longer think that is a reasonable omission. The certificate belongs on the card.
Not because these teams were careless. The two implementations we flagged are serious work from serious groups, and the defect is one axis in one reduction inside a computation that is genuinely fiddly. That is exactly the argument: if a bug class is invisible to the tool everyone uses, the fix is a different tool, cheap enough that everyone runs it.
Two forward passes. No gradients. Runs on a CPU in seconds.
If you maintain a hybrid or state-space model, the concrete ask is small: run your own audit at a sequence length above your chunk size, with a positive control, on the PyTorch path. That is where we found it.
Paper: arXiv:2608.22876 — The Mask Is Not the Model: Auditing Prefix Invariance in Attention, State-Space, and Hybrid Sequence Models
Authors: Taebong Kim, Youngsik Hong, Minsik Kim, Sunyoung Choi, Jaewon Jang, Minseo Kim — VIDRAFT AI Research
AX-RAY model diagnostics: huggingface.co/spaces/FINAL-Bench/AX-RAY
Top comments (0)