Transformers Community Investigates Mamba2/Zamba2 Cached Decoding Discrepancy; Time-Step Limit May Not Be Applied

A report about a state space model in Transformers has put consistency between prefill and token-by-token decoding under scrutiny. The issue was opened on October 9. The following day, another community member said they reproduced a logits discrepancy in Mamba2 and Zamba2 and noted that other models share the relevant cache path. At the time of review, the issue remained open, and the page listed no related fix PR. Community report

In the code examined by the reporter, a full-sequence scan first adds a bias to the time step dt, applies softplus, and then clamps it to time_step_limit. A cached single-token update delegates bias handling and softplus to an update function, but lacks the same range limit. As a result, the same input could use different state-update equations when processed all at once versus continued from a prefix using the cache. The original validation covered standalone PyTorch functions on CPU and CUDA; it did not measure generation quality in a pretrained model. Reproduction and proposed fix

The official documentation defines Mamba2’s time_step_limit as a time-step clipping range, defaulting from zero to infinity. Whether an effective bound is configured is therefore an important condition to check. The cache stores state-space and convolution states for reuse by subsequent tokens. The documentation also cautions that different kernels, operation ordering, and low precision can cause numerical differences, so tests need to distinguish ordinary error from this missing limit. Official Mamba2 documentation

For engineering teams, the technical risk is that offline full-sequence evaluation and online cached inference may disagree. This risk is inferred from the reported mechanism; it does not establish a general decline in answer quality. A recommended check is to hold weights, inputs, precision, and padding method constant while comparing logits and final states from a full scan and cached continuation. The proposal is to apply the bias, softplus, and clipping consistently before the single-token kernel, while avoiding duplicate processing inside the kernel. The full model and accelerated backends still need validation. Released versions also need to be checked individually; a case on the development branch does not establish that they are affected. Issue validation scope