Qubitium commited on
Commit
728d0de
·
verified ·
1 Parent(s): 75fab16

Disambiguate exact-capacity static masks

Browse files

Follow-up: the exact-capacity static-mask case is now covered. When the mask width equals `past_length + input_length`, its final slot distinguishes a full prompt with unused cache capacity from a new-only token block. A mask with unused capacity at the end keeps the full-prompt path.

[GPT-QModel #3574](https://github.com/ModelCloud/GPTQModel/pull/3574) contains regression tests for both exact-capacity and longer padded masks, plus existing new-only block coverage. It builds on merged [GPT-QModel #3571](https://github.com/ModelCloud/GPTQModel/pull/3571). The Ripple suite passed **35 tests** against this upstream source (`tests/test_ricursive_support.py`).

Files changed (1) hide show
  1. modeling_ricursive.py +11 -9
modeling_ricursive.py CHANGED
@@ -911,17 +911,19 @@ class RicursiveForCausalLM(RicursivePreTrainedModel, GenerationMixin):
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 with exactly the cached prefix plus new
914
- # tokens identifies new-only input. A longer mask may be padded to a
915
- # static cache's capacity, so it cannot identify a new-only block.
 
916
  input_length = input_ids.shape[1]
917
- full_prompt = (
918
- past_length > 0
919
- and input_length > past_length
920
- and (
921
- attention_mask is None
922
- or attention_mask.shape[-1] != past_length + input_length
923
- )
924
  )
 
925
  new_len = input_length - past_length if full_prompt else input_length
926
  cache_position = torch.arange(
927
  past_length, past_length + new_len, device=input_ids.device
 
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 with exactly the cached prefix plus new
914
+ # tokens identifies new-only input when its final slot is populated.
915
+ # A longer mask, or one with unused capacity at the end, may be padded
916
+ # to a static cache's capacity and describes a full prompt instead.
917
  input_length = input_ids.shape[1]
918
+ candidate_full_prompt = past_length > 0 and input_length > past_length
919
+ new_block_mask = (
920
+ candidate_full_prompt
921
+ and attention_mask is not None
922
+ and attention_mask.dim() == 2
923
+ and attention_mask.shape[-1] == past_length + input_length
924
+ and bool(attention_mask[..., -1].any())
925
  )
926
+ full_prompt = candidate_full_prompt and not new_block_mask
927
  new_len = input_length - past_length if full_prompt else input_length
928
  cache_position = torch.arange(
929
  past_length, past_length + new_len, device=input_ids.device