add block cache

modeling.py

@@ -554,6 +554,8 @@ class Fast_dLLM_QwenForCausalLM(Fast_dLLM_QwenPreTrainedModel, GenerationMixin):
         stopping_criteria=None,
         top_p=0.95,
         temperature=0,
+        use_block_cache=False,
+        block_cache_refresh_interval=16,
         **kwargs
     ):
         num_blocks = max_new_tokens // block_size
@@ -574,18 +576,20 @@ class Fast_dLLM_QwenForCausalLM(Fast_dLLM_QwenPreTrainedModel, GenerationMixin):
             if stop_token in input_ids[:, original_input_length:]:
                 break
             prompt_length = input_ids.shape[1]
-            # 初始化x_init为mask_id
+            # Initialize x_init with mask_id
             x_init = mask_id * torch.ones((input_ids.shape[0], block_size-prompt_length%block_size), device=self.device, dtype=torch.long)
             x_init = torch.cat([input_ids, x_init], dim=1)
 
             x_t = x_init.clone()
+            step = 0
+            block_past_key_values = None
             while True:
                 if stop_token in x_t[:, prompt_length:]:
                     stop_token_idx = (x_t[:, prompt_length:] == stop_token).nonzero()[0][1]
                     if (x_t[:, prompt_length:prompt_length+stop_token_idx] == mask_id).sum() == 0:
                         break
                 mask_idx = (x_t[:, -block_size:] == mask_id)
-                # 解码完整的一个block，更新cache，并且生成下一个token
+                # Decode a complete block, update cache, and generate the next token
                 if mask_idx.sum() == 0:
                     output = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=True, block_size=block_size)
                     logits, past_key_values = output.logits, output.past_key_values
@@ -595,43 +599,59 @@ class Fast_dLLM_QwenForCausalLM(Fast_dLLM_QwenPreTrainedModel, GenerationMixin):
                 for small_block_idx in range(num_small_blocks):
                     small_block_start_idx = small_block_idx * small_block_size
                     small_block_end_idx = small_block_start_idx + small_block_size
+
+                    start = -block_size + small_block_start_idx
+                    end = None if block_size == small_block_end_idx else -block_size + small_block_end_idx
                     while True:
                         mask_idx = (x_t[:, -block_size:] == mask_id)
-                        if mask_idx[:, small_block_start_idx:small_block_end_idx].sum() == 0:
+                        if mask_idx[:, start:end].sum() == 0:
                             break
                         if stop_token in x_t[:, prompt_length:]:
                             stop_token_idx = (x_t[:, prompt_length:] == stop_token).nonzero()[0][1]
                             if (x_t[:, prompt_length:prompt_length+stop_token_idx] == mask_id).sum() == 0:
                                 break
-                        logits = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, block_size=block_size).logits
+
+                        if use_block_cache:
+                            if step % block_cache_refresh_interval == 0 or (x_t[:, -block_size+small_block_start_idx] == mask_id).any():
+                                output = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True)
+                                logits, block_past_key_values = output.logits, output.block_past_key_values
+                                logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
+                                logits = logits[:, start:end]
+                            else:
+                                logits = self.forward(input_ids=x_t[:, -block_size+small_block_start_idx:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True, block_past_key_values=block_past_key_values, replace_position=small_block_start_idx).logits
+                                logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
+                        else:
+                            logits = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False).logits
+                            logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
+                            logits = logits[:, start:end]
+
+
                         x_1, p_1t = self.sample_with_top_p(logits, top_p=top_p, temperature=temperature)
-                        # 选出p_1t中概率大于threshold的token
+                        # Select tokens with probability greater than threshold from p_1t
                         x1_p = torch.squeeze(torch.gather(p_1t, dim=-1, index=torch.unsqueeze(x_1, -1)), -1)
                         x1_p = torch.where(mask_idx, x1_p, -torch.inf)
-                        x1_p[:, small_block_end_idx:] = -torch.inf
+
                         unmask_idx = (x1_p > threshold)
+                        max_prob_idx = x1_p.argmax(dim=-1)
+                        unmask_idx[torch.arange(x_1.shape[0]), max_prob_idx] = True
+                        unmask_idx = unmask_idx & mask_idx[:, start:end]
 
-                        if unmask_idx.sum() > 0:
-                            x_t[:, -block_size:][unmask_idx] = x_1[unmask_idx]
-                        else:
-                            # 选出p_1t中概率最大的那一个token
-                            token_position = x1_p.argmax()
-                            x_t[:, -block_size:][0, token_position] = x_1[0, token_position]
+                        x_t[:, start:end][unmask_idx] = x_1[unmask_idx]
 
+                        step += 1
             input_ids = x_t
-        # 截断stop_token
+        # Truncate stop_token
         if stop_token in input_ids[:, original_input_length:]:
             stop_token_idx = (input_ids[:, original_input_length:] == stop_token).nonzero()[0][1]
             input_ids = input_ids[:, :stop_token_idx+original_input_length+1]
         return input_ids
 
     def sample_with_top_p(self, logits, top_p=0.95, temperature=1.0):
-        # 计算概率
+        # Calculate probabilities
         if temperature > 0:
             scaled_logits = logits / temperature
         else:
             p_1t = torch.softmax(logits, dim=-1)
-            p_1t = torch.cat([p_1t[:, :1, :], p_1t[:, :-1, :]], dim=1)
             x_1 = p_1t.argmax(dim=-1)
             return x_1, p_1t
                             
@@ -650,13 +670,12 @@ class Fast_dLLM_QwenForCausalLM(Fast_dLLM_QwenPreTrainedModel, GenerationMixin):
         
         probs[indices_to_remove] = 0
 
-        # 3. 重新归一化并采样
-        # 重新归一化，使得剩余 token 的概率和为 1
-        # 添加一个极小值 eps 防止除以零
+        # Renormalize so that the probabilities of remaining tokens sum to 1
+        # Add a small epsilon value to prevent division by zero
         probs_sum = torch.sum(probs, dim=-1, keepdim=True)
         normalized_probs = probs / probs_sum
 
-        p_1t = torch.cat([normalized_probs[:, :1, :], normalized_probs[:, :-1, :]], dim=1)
+        p_1t = normalized_probs
         x_1 = torch.multinomial(p_1t[0], num_samples=1).unsqueeze(0).squeeze(-1)
 
         return x_1, p_1t