Update: modeling_ricursive.py in this PR now also carries the DFlash2 verification fixes (multi-token warm-cache mask, logits_to_keep, chunked MLP). See the PR description for the full bug/fix table and validation.

Tracking PR in GPTQModel: https://github.com/ModelCloud/GPTQModel/pull/3545

This upstream fix is the same code that ships in GPTQModel's ricursive model
definition. That PR adds Ripple support plus the load-time compatibility patch and
carries the full test suite (tests/test_ricursive_support.py) and benchmark numbers.

Fix generate() and attention masking under Transformers 5.x

Symptom

Under Transformers 5.x, generate() fails on the very first prefill:

RuntimeError: not enough values to unpack (expected 3, got 2)

prepare_inputs_for_generation does input_ids[:, cache_position] whenever
input_ids.shape[1] != 1. Transformers 5 pre-creates an empty DynamicCache and passes
cache_position=None at prefill, so that evaluates input_ids[:, None] and adds a
spurious axis: (batch, seq) becomes (batch, 1, seq). The embedding lookup then yields
4-D hidden states and the first decoder layer dies on
bsz, q_len, _ = hidden_states.shape.

That line is only the first of six defects. Padded batches decode to garbage and
generations longer than sliding_window crash.

Bug → Fix

Notes on the two that are easy to miss:

• #2 only shows up with an explicit padding mask — an all-ones mask returns early. So
  single-sequence generation looked healthy while batched evaluation silently scored 0%.
  The old code also wrote 1.0 for visible keys, adding a constant bias to every score.
• #4 is a silent correctness bug, not a crash. The window mask progressively excluded
  keys that were still valid, degrading long-context output without any error.

Also fixed: position_ids derived by transformers from caller tensors arrive on CPU while
the model is on GPU, so RoPE fails with mat2 is on cpu. input_ids, attention_mask,
inputs_embeds and position_ids are now all moved to the model device.

A module-level RICURSIVE_TF5_COMPAT marker is set on the class so downstream loaders that
monkeypatch older checkpoints stand down instead of stacking patches.

Validation

Transformers 5.22.1. "Bit-identical" means exact tensor equality of generated token
ids / logits — not a tolerance comparison.

The native implementation does not merely approximate the reference patch — it produces
exactly the same tokens on every path exercised.

Sync

The identical modeling_ricursive.py (sha256 b2ca81fa07dcdeab…, 30790 bytes) is
published in all three places:

The GPTQModel PR (#3535) carries the
equivalent fix as a load-time runtime patch (gptqmodel/utils/ricursive_patch.py) for
checkpoints that still hold the original file. The two implementations were compared
directly on the same checkpoint:

Compatibility

Backward compatible with Transformers 4.x: GenerationMixin is importable from
transformers.generation.utils, and listing it after PreTrainedModel is valid when
PreTrainedModel already inherits it.

Files

modeling_ricursive.py   patched drop-in replacement (only file changed)

Mirror of the patched file and a unified diff:
https://huggingface.co/ModelCloud/Ripple-65B-Preview-tf5-patch

Thank you for the PR, couple of comments before merge:

1. Preallocated caches can break causal attention. Line 298 treats cache capacity as occupied context. A six-token prompt in a 16-slot StaticCache gets query positions 10–15. Changing token six then changes token one’s output in full-attention layers, which must be impossible. The previous commit passes this check. Build the mask from actual cache positions.
2. labels with logits_to_keep=1 returns NaN loss. Line 705 trims logits and labels to one position, then the loss calculation shifts both to empty tensors. The same call previously returned a finite loss. Preserve the labeled loss calculation or explicitly reject this combination.

@ebrahim-ricursive

https://github.com/ModelCloud/GPTQModel/pull/3545

both issues are fixed in modeling_ricursive.py (commit 0d2fa09).

1. Preallocated caches no longer define attention context.
StaticCache returns its full preallocated capacity from update(), so the mask was built from capacity instead of occupied positions: a 6-token prompt in a 16-slot cache was treated as a 10-token prefix. The mask is now built from the layer's actual get_seq_length() and cache_position; the unwritten tail of k/v is trimmed, and an incoming mask longer than the keys actually present is aligned to them.

Measured on the full model: changing the last token of a 6-token prompt in a 16-slot StaticCache now gives 0.0 difference on tokens 1–5 (previously 0.4830).

2. labels + logits_to_keep=1 no longer returns NaN.
The full sequence is kept whenever labels are supplied, so the shifted causal-LM loss stays defined; only the returned logits are trimmed. logits_to_keep=1 returns (batch, 1, vocab) with the same loss as the untrimmed forward. Unlabeled calls still slice the LM-head input, so generation keeps the memory win.

Full-model validation was run on the quantized checkpoint ModelCloud/Ripple-65B-Preview-GPTQ-W4A16-v2 (Ripple-65B-Preview, bfloat16), with the load-time patches enabled and, separately, disabled so only this file is active — same numbers in both runs:

All numbers above were measured on the quantized checkpoint linked above. The verify-block error (0.25) is within 1.4× of the no-mask control (0.156) and about 130× smaller than the previous 32.31, so the residual is bfloat16 kernel shape noise, not masking. The DFlash2 token difference is a genuine tie at bfloat16 precision.

The same fixes are mirrored in GPTQModel (ModelCloud/GPTQModel#3551; earlier work in #3545), which applies them at load time for quantized checkpoints.

Thanks for the fix. Almost there, there are two cache issues need fixing:

1. Wrong mask after sliding-window eviction. With a mask sized to the allocated cache capacity, line 315 selects the unused tail instead of the positions occupied by retained keys. My test produces a maximum logit error of 0.345, versus 0.00000016 with the correctly sized mask. Slice [total_kv - kv_len : total_kv]; otherwise this path silently returns incorrect outputs.

2. Transformers 4 StaticCache decoding overwrites cached tokens. Line 289 omits cache_position from the cache update. Under 4.57.6, the first decode step overwrites the beginning of the cache, producing a 0.145 maximum logit error. Forward the cache positions and add a decoding regression test. This is a remaining compatibility issue, not introduced by the latest commits.

Companion implementation: https://github.com/ModelCloud/GPTQModel/pull/3545

Both issues from the last review are fixed in modeling_ricursive.py (commit ee974ea). Summary of the two new fixes, then the full bug/fix table and the validation numbers.

Review issues (this round)

All bugs fixed in this PR

Plus two inference fixes: an end-anchored causal mask for multi-token verify blocks against a warm cache, and logits_to_keep slicing the LM head input (labeled calls keep the full sequence so the shifted loss stays defined).

Validation

[PASS] issue1 capacity-sized mask aligns to [total_kv-kv_len, total_kv) (max abs err = 0.000e+00)
[PASS] issue1 test is sensitive: old tail slice is wrong        (max abs err = 4.260e-01)
[PASS] issue2 StaticCache decode matches a full forward          (max abs err = 3.576e-07)
[PASS] issue2 cache_position is forwarded into update()          (received {'cache_position': tensor([9])})
[PASS] labeled loss stays finite with logits_to_keep=1           (loss=4.8866, logits=(1,1,128))
ALL_REVIEW_FIX_TESTS_PASSED

• Tested on transformers 5.22.1; the cache_position path is the 4.x compatibility contract (5.x caches ignore the kwarg, 4.x StaticCache needs it).
• generate() parity: cached and uncached outputs are identical.
• Padded batches: bit-identical to unpadded.
• The 65B checkpoint loads, generates, and produces deterministic output with this file.

The same file is published byte-identically in ModelCloud/Ripple-65B-Preview-GPTQ-W4A16-v2 (sha256 da08c5b4...), so downstream users get the identical code.

Great, Thank you. Justed tested on my side and it works.