Pass through more inputs for cache support

Author: tarun-menta

surya/common/surya/__init__.py

@@ -296,6 +296,8 @@ def forward(
         logits_to_keep=None,
         encoder_chunk_size=None,
         cache_idxs=None,
+        valid_tokens=None,
+        prefill=False,
         **kwargs: KwargsForCausalLM,
     ):
         # Process the mixed batch if provided
@@ -351,6 +353,8 @@ def forward(
             return_dict=True,
             use_cache=use_cache,
             cache_idxs=cache_idxs,
+            valid_tokens=valid_tokens,
+            prefill=prefill
             **kwargs,
         )
 

surya/common/surya/decoder/__init__.py

@@ -161,6 +161,8 @@ def forward(
         past_key_value: Optional[Cache] = None,
         cache_position: Optional[torch.LongTensor] = None,
         cache_idxs: Optional[List[int]] = None,
+        valid_tokens: Optional[List[int]] = None,
+        prefill: bool = False,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
         input_shape = hidden_states.shape[:-1]
@@ -185,7 +187,15 @@ def forward(
 
         if past_key_value is not None:
             # sin and cos are specific to RoPE models; cache_position needed for the static cache
-            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position, "cache_idxs": cache_idxs}
+            # cache_idxs, valid_tokens, and prefill add support for our new caching mechanism 
+            cache_kwargs = {
+                "sin": sin,
+                "cos": cos,
+                "cache_position": cache_position,
+                "cache_idxs": cache_idxs,
+                "valid_tokens": valid_tokens,
+                "prefill": prefill
+            }
             key_states, value_states = past_key_value.update(
                 key_states, value_states, self.layer_idx, cache_kwargs
             )
@@ -279,6 +289,8 @@ def forward(
         use_cache: Optional[bool] = False,
         cache_position: Optional[torch.LongTensor] = None,
         cache_idxs: Optional[List[int]] = None,
+        valid_tokens: Optional[List[int]] = None,
+        prefill: bool = False,
         position_embeddings: Optional[
             Tuple[torch.Tensor, torch.Tensor]
         ] = None,  # necessary, but kept here for BC
@@ -300,7 +312,9 @@ def forward(
             use_cache=use_cache,
             cache_position=cache_position,
             position_embeddings=position_embeddings,
-            cache_idxs=cache_idxs
+            cache_idxs=cache_idxs,
+            valid_tokens=valid_tokens,
+            prefill=prefill,
             **kwargs,
         )
         hidden_states = residual + hidden_states
@@ -461,6 +475,8 @@ def forward(
         return_dict: Optional[bool] = None,
         cache_position: Optional[torch.LongTensor] = None,
         cache_idxs: Optional[List[int]] = None,
+        valid_tokens: Optional[List[int]] = None,
+        prefill: bool = False,
         **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
     ) -> Union[Tuple, BaseModelOutputWithPast]:
         use_cache = use_cache if use_cache is not None else self.config.use_cache
@@ -501,6 +517,8 @@ def forward(
                 cache_position=cache_position,
                 position_embeddings=position_embeddings,
                 cache_idxs=cache_idxs,
+                valid_tokens=valid_tokens,
+                prefill=prefill,
                 **flash_attn_kwargs,
             )