Enable defrag for paged TP cache

Author: turboderp

exllamav2/attn.py

@@ -757,16 +757,16 @@ def forward_paged_tp_old(
         else:
             cache_seqlens_a = attn_params.cache_seqlens_tp
 
-        if cache.q_block == 1:
-            cache.get_kv_state(
-                self.layer_idx,
-                batch_size,
-                0,
-                attn_params.max_cache_seqlen,
-                page_size,
-                attn_params.cache_seqlens_tp,
-                attn_params.block_index_tp
-            )
+        # if cache.q_block == 1:
+        #     cache.get_kv_state(
+        #         self.layer_idx,
+        #         batch_size,
+        #         0,
+        #         attn_params.max_cache_seqlen,
+        #         page_size,
+        #         attn_params.cache_seqlens_tp,
+        #         attn_params.block_index_tp
+        #     )
 
         flash_kwargs = {}
         if self.sliding_window:

exllamav2/cache.py

@@ -845,7 +845,6 @@ def copy_states(
         to_row: int,
         to_rows: int
     ):
-        # TODO: Parallel implementation
         for cache, tcache in zip(self.caches, target.caches):
             cache.copy_states(
                 tcache,
@@ -865,8 +864,10 @@ def touch_device(self, device):
 
 
     def all_tensors(self):
-        # TODO: Support defrag with TP cache
-        return []
+        tensors = []
+        for cache in self.caches:
+            tensors += cache.all_tensors()
+        return tensors
 
 
     def reset(self):

exllamav2/generator/dynamic.py

@@ -2222,6 +2222,7 @@ def prefill(self, results: list):
                             best_match_page = page
                     if best_match_page and best_match > 1:
                         page = seq.allocated_pages[p0]
+                        # print([sap.page_index for sap in seq.allocated_pages])
                         for c in [self.generator.cache] if not self.generator.draft_model else \
                             [self.generator.cache, self.generator.draft_cache]:
                             c.copy_states(