Fix warm-cache new token blocks in generation preparation

modeling_ricursive.py

@@ -908,17 +908,26 @@ class RicursiveForCausalLM(RicursivePreTrainedModel, GenerationMixin):
         past_length = _cache_seq_length(past_key_values)
 
         if cache_position is None:
-            # With a non-empty cache and the full sequence supplied, only the trailing
-            # (uncached) tokens are new. Otherwise every token in `input_ids` is new.
-            if past_length > 0 and input_ids.shape[1] > 1:
-                new_len = input_ids.shape[1] - past_length
-                cache_position = torch.arange(
-                    past_length, past_length + new_len, device=input_ids.device
-                )
-            else:
-                cache_position = torch.arange(
-                    past_length, past_length + input_ids.shape[1], device=input_ids.device
-                )
+            # A full prompt can include cached tokens, while verification may pass
+            # only a new token block. A block no longer than the cache must be new;
+            # for a longer block, a mask with exactly the cached prefix plus new
+            # tokens identifies new-only input when its final slot is populated.
+            # A longer mask, or one with unused capacity at the end, may be padded
+            # to a static cache's capacity and describes a full prompt instead.
+            input_length = input_ids.shape[1]
+            candidate_full_prompt = past_length > 0 and input_length > past_length
+            new_block_mask = (
+                candidate_full_prompt
+                and attention_mask is not None
+                and attention_mask.dim() == 2
+                and attention_mask.shape[-1] == past_length + input_length
+                and bool(attention_mask[..., -1].any())
+            )
+            full_prompt = candidate_full_prompt and not new_block_mask
+            new_len = input_length - past_length if full_prompt else input_length
+            cache_position = torch.arange(
+                past_length, past_length + new_len, device=input_ids.device
+            )
 
         # Only slice when the cache is warm and `input_ids` still carries cached
         # tokens. Slicing unconditionally turned a `None` cache position into