Disambiguate exact-capacity static masks
Browse filesFollow-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`).
- 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
|
| 915 |
-
#
|
|
|
|
| 916 |
input_length = input_ids.shape[1]
|
| 917 |
-
|
| 918 |
-
|
| 919 |
-
|
| 920 |
-
and
|
| 921 |
-
|
| 922 |
-
|
| 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
|