moving past keyvalue dynamically

modeling_molmo2.py

@@ -1054,31 +1054,22 @@ class Molmo2TextModel(Molmo2PreTrainedModel):
                 position_embeddings_i = position_embeddings
 
             block_device = next(decoder_block.parameters()).device
-            if hidden_states.device != block_device:
-                hidden_states = hidden_states.to(block_device)
 
-            if causal_mask_mapping is not None and causal_mask_mapping.device != block_device:
-                causal_mask_mapping = causal_mask_mapping.to(block_device)
-
-            if position_ids is not None and position_ids.device != block_device:
-                position_ids = position_ids.to(block_device)
-
-            if position_embeddings_i is not None and position_embeddings_i.device != block_device:
-                position_embeddings_i = position_embeddings_i.to(block_device)
-
-            if cache_position is not None and cache_position.device != block_device:
-                cache_position = cache_position.to(block_device)
-
-            def move_past_key_values(past_key_values, device):
-                if past_key_values is None:
-                    return None
-                return tuple(
-                    tuple(p.to(device) for p in layer)
-                    for layer in past_key_values
-                )
-
-            past_key_values = move_past_key_values(past_key_values, block_device)
-            position_embeddings_i = position_embeddings_i.to(block_device)
+            def move_if_needed(x, device):
+                if x is None:
+                    return x
+                if isinstance(x, tuple):
+                    return tuple(move_if_needed(xx, device) for xx in x)
+                if hasattr(x, "device") and x.device != device:
+                    return x.to(device)
+                return x
+
+            hidden_states = move_if_needed(hidden_states, block_device)
+            causal_mask_mapping = move_if_needed(causal_mask_mapping, block_device)
+            position_ids = move_if_needed(position_ids, block_device)
+            position_embeddings_i = move_if_needed(position_embeddings_i, block_device)
+            cache_position = move_if_needed(cache_position, block_device)
+            past_key_values = move_if_needed(past_key_values, block_device)
 
             layer_outputs = decoder_block(
                 hidden_states,