add sampling

Author: Viol2000

README.md

@@ -4,6 +4,11 @@
 | <a href="https://arxiv.org/abs/2402.02057"><b>Paper</b></a> | <a href="https://lmsys.org/blog/2023-11-21-lookahead-decoding/"><b>Blog</b></a> | <a href="https://github.com/hao-ai-lab/LookaheadDecoding/issues/13"><b>Roadmap</b></a> | 
 </p>
 
+---
+*News* 🔥
+- [2024/2] Lookahead Decoding Paper now available on [arXiv](https://arxiv.org/abs/2402.02057). Sampling and FlashAttention are supported. Advanced features for better token prediction are updated.
+
+---
 ## Introduction 
 We introduce lookahead decoding:
 - A parallel decoding algorithm to accelerate LLM inference.
@@ -148,14 +153,22 @@ lade.config_lade(LEVEL=5, WINDOW_SIZE=7, GUESS_SET_SIZE=7, DEBUG=0)
 #You can obtain a better performance by tuning LEVEL/WINDOW_SIZE/GUESS_SET_SIZE on your own device.
 ```
 
-Then you can speedup the decoding process.
+Then you can speedup the decoding process. Here is an example using greedy search:
 ```
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map=torch_device)
 model_inputs = tokenizer(input_text, return_tensors='pt').to(torch_device)
 greedy_output = model.generate(**model_inputs, max_new_tokens=1024) #speedup obtained
 ```
 
+Then you can speedup the decoding process. Here is an example using sampling:
+```
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map=torch_device)
+model_inputs = tokenizer(input_text, return_tensors='pt').to(torch_device)
+sample_output = model.generate(**model_inputs, max_new_tokens=1024, temperature=0.7) #speedup obtained
+```
+
 ## Citation
 ```bibtex
 @misc{fu2024break,

lade/decoding.py

@@ -374,7 +374,7 @@ def copy_from_last():
 
     for warper in logits_warper:
         #assert type(warper) == TemperatureLogitsWarper or type(warper) == TopPLogitsWarper or type(warper) == TopKLogitsWarper,  f"please set top_k=0 {warper}"
-        assert type(warper) == TemperatureLogitsWarper,  f"please set top_k=0.0 and top_p=1.0 {warper}"
+        assert type(warper) == TemperatureLogitsWarper or type(warper) == TopKLogitsWarper or type(warper) == TopPLogitsWarper,  f"please set top_k=0.0 and top_p=1.0 {warper}"
 
     # auto-regressive generation
     while True:
@@ -485,9 +485,8 @@ def copy_from_last():
                 probs_next = torch.nn.functional.softmax(next_token_scores, dim=-1)[0]
                 hits = []
                 #= original model output
-                #print("size: ", input_ids.size(), outputs.guess_logits.size())
-                guess_logits = logits_warper(input_ids, outputs.guess_logits)
-                guess_probs = torch.nn.functional.softmax(guess_logits, dim=-1)[0] #
+                guess_logits = logits_warper(input_ids, outputs.guess_logits[0])
+                guess_probs = torch.nn.functional.softmax(guess_logits, dim=-1) #
                 #guess_results = torch.argmax(outputs.guess_logits, dim=-1)[0].tolist()
                 guess_indices = list(range(outputs.guess_logits.size(1) // GUESS_SIZE))
                 #algorithm modified from specinfer
@@ -887,7 +886,7 @@ def random_set():
 
     def copy_from():
         return random.choice(all_old_tokens)
-    
+
     def order_copy_from():
         if order_copy_from_idx[0] >= len(all_old_tokens):
             order_copy_from_idx[0] = 0
@@ -915,12 +914,12 @@ def copy_from_last():
 
     if POOL_FROM_PROMPT:
         fill_pool_with_prompt(all_old_tokens, token_map, LEVEL, GUESS_SET_SIZE)
-
+        
     if chat:
         init = self.tokenizer.decode(all_old_tokens, skip_special_tokens=True, \
                                    spaces_between_special_tokens=False, clean_up_tokenization_spaces=True,)
         prev = len(init)
-    
+
     while True:
         if synced_gpus:
             # Under synced_gpus the `forward` call must continue until all gpus complete their sequence.
@@ -964,7 +963,7 @@ def copy_from_last():
                     guess_tokens = None
         else:
             guess_tokens = None
-        
+
         assert return_dict_in_generate == False
         assert len(logits_processor) == 0
         # forward pass to get next token        
@@ -985,7 +984,7 @@ def copy_from_last():
                 past_tokens_inp.append(tokens[window_start: window_end] if tokens is not None else None)
         else:
             past_tokens_inp = past_tokens
-
+            
         outputs = self.jforward_multilevel(
             **model_inputs,
             past_tokens=past_tokens_inp,
@@ -1040,7 +1039,7 @@ def copy_from_last():
             assert fill_level == 0
             past_tokens[0] = past_tokens[0][1:] 
             past_tokens[1] = torch.argmax(outputs.inp_logits, dim=-1)[0].tolist()
-
+            
             if DIST_WORKERS > 1:
                 nn_past_tokens = [copy.deepcopy(past_tokens[1])]
                 torch.distributed.broadcast_object_list(nn_past_tokens, src=DIST_WORKERS - 1)
@@ -1051,7 +1050,7 @@ def copy_from_last():
             for level in range(fill_level + 1):
                 past_tokens[level] = past_tokens[level][1:] 
             current_past_tokens = torch.argmax(outputs.inp_logits, dim=-1)[0].tolist()
-
+            
 
             if DIST_WORKERS > 1:
                 nn_past_tokens = [None] * DIST_WORKERS
@@ -1063,9 +1062,9 @@ def copy_from_last():
             past_tokens[fill_level + 1] = current_past_tokens[1:]
             #print("new past: ", (LOCAL_RANK, past_tokens))
 
-            
+
             fill_level += 1
-        else:
+        else: 
             #time.sleep(10000)
             #multi-level window is filled
             #match guess tokens 
@@ -1101,7 +1100,7 @@ def copy_from_last():
             #    print("rank: ",hits, max_hit)
             #sync new_results
             new_results = torch.argmax(outputs.inp_logits, dim=-1)[0].tolist()
-            
+
             if DIST_WORKERS > 1:
                 nn_past_tokens = [None] * DIST_WORKERS
                 torch.distributed.all_gather_object(nn_past_tokens, new_results)
@@ -1149,7 +1148,7 @@ def copy_from_last():
         if DIST_WORKERS > 1 and max_hit > 0:
 
             guess_skip_dist = max_hit
-        for idx, kv in enumerate(outputs.past_key_values):
+            for idx, kv in enumerate(outputs.past_key_values):
                 past_key_values.append( (kv[0][:,:,:outputs.kvcache_len,:], kv[1][:,:,:outputs.kvcache_len,:]) )
             outputs.past_key_values = past_key_values
         else:
@@ -1160,8 +1159,8 @@ def copy_from_last():
                 if max_hit > 0:
                     kv[0][:,:,outputs.kvcache_len:outputs.kvcache_len+max_hit,:] = kv[0][:,:,offset_kv_cache:offset_kv_cache+max_hit,:]
                     kv[1][:,:,outputs.kvcache_len:outputs.kvcache_len+max_hit,:] = kv[1][:,:,offset_kv_cache:offset_kv_cache+max_hit,:]
-            past_key_values.append( (kv[0][:,:,:outputs.kvcache_len + max_hit,:], kv[1][:,:,:outputs.kvcache_len + max_hit,:]) )
-        outputs.past_key_values = past_key_values
+                past_key_values.append( (kv[0][:,:,:outputs.kvcache_len + max_hit,:], kv[1][:,:,:outputs.kvcache_len + max_hit,:]) )
+            outputs.past_key_values = past_key_values
 
         lst_token = hits[max_hit]
 
@@ -1176,7 +1175,7 @@ def copy_from_last():
                 all_old_tokens.append(hits[max_hit])
                 if POOL_FROM_PROMPT:
                     append_new_generated_pool(all_old_tokens[-LEVEL:], token_map, LEVEL, GUESS_SET_SIZE)
-        
+
 
         if chat and LOCAL_RANK == 0:
             all_str = self.tokenizer.decode(all_old_tokens, skip_special_tokens=True, \
@@ -1188,7 +1187,7 @@ def copy_from_last():
                                     spaces_between_special_tokens=False, clean_up_tokenization_spaces=True,) 
                     pt = colored(not_hit[prev:],"blue") +  colored(all_str[len(not_hit):], "blue")
                 else:
-                    pt = all_str[prev:]                    
+                    pt = all_str[prev:]                
                 print(pt,  flush=True, end="")
             else:
                 print(all_str[prev:],  flush=True, end="")
@@ -1440,7 +1439,7 @@ def greedy_search_chat(
 
         # prepare model inputs
         model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
-
+        
         # forward pass to get next token
         outputs = self(
             **model_inputs,