Qubitium commited on
Commit
d160eae
·
verified ·
1 Parent(s): da94cc0

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.

Files changed (1) hide show
  1. modeling_ricursive.py +14 -11
modeling_ricursive.py CHANGED
@@ -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
- # With a non-empty cache and the full sequence supplied, only the trailing
912
- # (uncached) tokens are new. Otherwise every token in `input_ids` is new.
913
- if past_length > 0 and input_ids.shape[1] > 1:
914
- new_len = input_ids.shape[1] - past_length
915
- cache_position = torch.arange(
916
- past_length, past_length + new_len, device=input_ids.device
917
- )
918
- else:
919
- cache_position = torch.arange(
920
- past_length, past_length + input_ids.shape[1], device=input_ids.device
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