Fix background filter eval when draft model used

Author: turboderp

exllamav2/generator/dynamic.py

@@ -233,9 +233,9 @@ class ExLlamaV2DynamicGenerator:
 
     max_sampling_threads: int
     min_sampling_threads: int
-    sampling_pool: ThreadPoolExecutor
-    filter_pool: ThreadPoolExecutor
-    filter_queue: list
+    sampling_pool: ThreadPoolExecutor | None
+    filter_pool: ThreadPoolExecutor | None
+    filter_queue: list | None
 
 
     def __init__(
@@ -255,6 +255,7 @@ def __init__(
         max_sampling_threads: int = 16,
         min_sampling_threads: int = 3,
         paged: bool = True,
+        filter_background_eval: bool = True,
         **kwargs
     ):
         """
@@ -316,6 +317,10 @@ def __init__(
             does not require paged attention support, but in which the max supported batch size is 1. CFG also will
             not work in this mode.
 
+        :param filter_background_eval:
+            Try to overlap filter evaluation with model forward pass. This should generally have no downside since
+            filters are evaluated by the CPU which will otherwise be busywaiting after CUDA workload is scheduled.
+
         :param kwargs:
         """
 
@@ -449,8 +454,12 @@ def __init__(
 
         # Filter threads
 
-        self.filter_pool = ThreadPoolExecutor(max_workers = 16)
-        self.filter_queue = []
+        if filter_background_eval:
+            self.filter_pool = ThreadPoolExecutor(max_workers = 16)
+            self.filter_queue = []
+        else:
+            self.filter_pool = None
+            self.filter_queue = None
 
         # Temp buffers for defrag
 
@@ -1243,7 +1252,8 @@ def iterate_gen(self, results: list, draft_tokens: torch.Tensor | None = None):
                     next_k_probs,
                     next_prob,
                     filter_eos,
-                    results
+                    results,
+                    i == 0
                 )
 
                 if eos:
@@ -1867,7 +1877,8 @@ def receive_sample(
             next_k_probs: torch.Tensor | None,
             next_prob: torch.Tensor | None,
             filter_eos: bool | None,
-            results: list
+            results: list,
+            first_sample_in_sd_batch: bool = True
     ):
         page_size = self.generator.page_size
 
@@ -1879,15 +1890,16 @@ def receive_sample(
                 f.feed(next_token)
                 if not f.can_mask_logits() or not f.use_background_worker():
                     all_mask = False
-            if all_mask:
-                # Using logit mask(s)
-                for f in self.filters:
-                    self.generator.filter_queue.append((f, True))
-            else:
-                # Using allowed token list(s)
-                for f in self.filters:
-                    if f.use_background_worker():
-                        self.generator.filter_queue.append((f, False))
+            if first_sample_in_sd_batch and self.generator.filter_queue is not None:
+                if all_mask:
+                    # Using logit mask(s)
+                    for f in self.filters:
+                        self.generator.filter_queue.append((f, True))
+                else:
+                    # Using allowed token list(s)
+                    for f in self.filters:
+                        if f.use_background_worker():
+                            self.generator.filter_queue.append((f, False))
 
         # Accept token
 

exllamav2/generator/sampler.py

@@ -442,6 +442,9 @@ def prep_logit_filter(lf):
                     "Attempting to use precomputed logit mask, but filter is not precomputing mask"
                 flat_logits = logits[0][0]
                 logits = f.mask_logits(flat_logits).view(1, 1, -1)
+                # not_inf_indices = torch.nonzero(logits != -float('inf'), as_tuple = True)
+                # txt = [tokenizer.get_id_to_piece_list()[i] for i in not_inf_indices[2].tolist()]
+                # print(txt)
             end_tokens = None
 
         elif len(filters) > 0: