Fix warm-cache new token blocks in generation preparation
Browse files## Problem
`RicursiveForCausalLM.prepare_inputs_for_generation` assumes every multi-token
`input_ids` input includes the cached prefix when `cache_position` is absent.
If the cache already holds 10 tokens and a caller supplies only two new tokens,
the calculated new length is -8 and `torch.arange` raises
`RuntimeError: upper bound and lower bound inconsistent with step sign`.
## Fix
Distinguish new-only token blocks from full prompts containing cached tokens.
An input block no longer than the cache is new; a longer input with an attention
mask spanning the cache is also new. Keep existing full-prompt slicing when
neither clue is present.
## Validation
- The three regression cases pass against this patched upstream function:
short new block, longer new block with a spanning mask, and full prompt.
- The tiny-model Ripple suite passed against this source:
`GPTQMODEL_RIPPLE_REF=<patched Ripple checkout> python -m pytest -q -o log_cli=false tests/test_ricursive_support.py`
(**33 passed**).
- `python -m py_compile modeling_ricursive.py` passed.
- Unit test coverage and the corresponding GPT-QModel compatibility fix are in
merged [ModelCloud/GPTQModel#3571](https://github.com/ModelCloud/GPTQModel/pull/3571),
especially `test_generate_patch_distinguishes_new_blocks_from_full_prompts`.
The current BF16 checkpoint marks its Transformers 5 support as native, so
GPT-QModel intentionally skips that compatibility patch; this repository needs
the fix directly.
- modeling_ricursive.py +14 -11
|
@@ -908,17 +908,20 @@ class RicursiveForCausalLM(RicursivePreTrainedModel, GenerationMixin):
|
|
| 908 |
past_length = _cache_seq_length(past_key_values)
|
| 909 |
|
| 910 |
if cache_position is None:
|
| 911 |
-
#
|
| 912 |
-
#
|
| 913 |
-
|
| 914 |
-
|
| 915 |
-
|
| 916 |
-
|
| 917 |
-
|
| 918 |
-
|
| 919 |
-
|
| 920 |
-
|
| 921 |
-
|
|
|
|
|
|
|
|
|
|
| 922 |
|
| 923 |
# Only slice when the cache is warm and `input_ids` still carries cached
|
| 924 |
# tokens. Slicing unconditionally turned a `None` cache position into
|
|
|
|
| 908 |
past_length = _cache_seq_length(past_key_values)
|
| 909 |
|
| 910 |
if cache_position is None:
|
| 911 |
+
# A full prompt can include cached tokens, while verification may pass
|
| 912 |
+
# only a new token block. A block no longer than the cache must be new;
|
| 913 |
+
# for a longer block, a mask spanning the cache also identifies new-only
|
| 914 |
+
# input. Keep the full-prompt behavior when neither clue is present.
|
| 915 |
+
input_length = input_ids.shape[1]
|
| 916 |
+
full_prompt = (
|
| 917 |
+
past_length > 0
|
| 918 |
+
and input_length > past_length
|
| 919 |
+
and (attention_mask is None or attention_mask.shape[-1] <= input_length)
|
| 920 |
+
)
|
| 921 |
+
new_len = input_length - past_length if full_prompt else input_length
|
| 922 |
+
cache_position = torch.arange(
|
| 923 |
+
past_length, past_length + new_len, device=input_ids.device
|
| 924 |
+
)
|
| 925 |
|
| 926 |
# Only slice when the cache is warm and `input_ids` still carries cached
|
| 927 |
# tokens. Slicing unconditionally turned a `None` cache position into
|