Update modeling_minicpm.py

Modify modeling_minicpm.py to use LSE compression.

modeling_minicpm.py

@@ -21,7 +21,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import torch
 import torch.nn.functional as F
 import torch.utils.checkpoint
-from torch import nn
+from torch import  nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache, CacheLayerMixin, DynamicLayer
@@ -47,7 +47,9 @@ from transformers.utils import (
 )
 from transformers.utils.import_utils import is_torch_fx_available
 
-from .configuration_minicpm import MiniCPMConfig
+
+
+from .configuration_minicpm import MiniCPMConfig    #!一定要改
 
 try:
     from flash_attn import flash_attn_func, flash_attn_varlen_func
@@ -68,50 +70,28 @@ from functools import lru_cache
 def compressed_attention(
     q: torch.Tensor,
     k: torch.Tensor,
-    v: torch.Tensor,
+    k2: torch.Tensor,
     kernel_size: int,
     kernel_stride: int,
     block_size: int,
     topk: int,
     cu_seqlens_q: torch.Tensor,
     cu_seqlens_k: torch.Tensor,
+    cu_seqlens_k2: torch.Tensor,
     max_seqlen_q: int,
     max_seqlen_k: int,
     sm_scale: float = None,
     init_blocks: int = 1,
     local_blocks: int = 2,
-    cache_lens: torch.Tensor = None,
+    cache_lens=None,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Attention between query and compressed key and value. Compute attention output and topk block idx used in topk_sparse_attention.
-
-    Args:
-        q (torch.Tensor): shape [total_q_len, num_q_heads, head_dim]
-        k (torch.Tensor): shape [total_kv_len, num_kv_heads, head_dim]
-        v (torch.Tensor): shape [total_kv_len, num_kv_heads, head_dim]
-        kernel_size (int): kernel size in compress_key_value
-        kernel_stride (int): stride of compress_key_value
-        block_size (int): key value block size for topk sparse attention.
-        topk (int): number of blocks for each query.
-        cu_seqlens_q (torch.Tensor): shape [batch_size + 1], similar to cu_seqlens_q in flash_attn_func_varlen.
-        cu_seqlens_k (torch.Tensor): shape [batch_size + 1], similar to cu_seqlens_k in flash_attn_func_varlen.
-        max_seqlen_q (int): max q len of the batch.
-        max_seqlen_k (int): max k len of the batch.
-        sm_scale (float, optional): softmax scale. Defaults to None, means 1/sqrt(head_dim).
-        init_blocks (int, optional): Number of init blocks for each query. Defaults to 1.
-        local_blocks (int, optional): Number of local blocks for each query. Defaults to 2.
-        cache_lens (torch.Tensor, optional): shape [batch_size], used to record the cache length of each query. Defaults to None.
-
-    Returns:
-        Tuple[torch.Tensor, torch.Tensor]: attention output and topk_idx used in topk_sparse_attention
-    """
     with torch.no_grad():
         batch_size = cu_seqlens_q.shape[0] - 1
         
         # Check if it's prefilling stage
         is_prefilling = cache_lens is None or (cache_lens == 0).all().item()
-
-        # prefilling stage
-        if is_prefilling:
+        
+        if is_prefilling:  # prefilling stage
             # Calculate q_idx for each query position in each batch
             cache_lens = torch.zeros(batch_size, dtype=torch.int32, device=q.device) 
             q_idx = torch.cat([
@@ -119,25 +99,24 @@ def compressed_attention(
                  max_seqlen_q - (cu_seqlens_q[i + 1] - cu_seqlens_q[i])) // block_size
                 for i in range(batch_size)
             ], dim=0)  # shape: [total_q_len]
-        # decoding stage
-        else:
-            # Each batch has only one query (last position). Shape: [batch_size] = [total_q_len] in decoding
-            q_idx = cache_lens // block_size
+        else:  # decoding stage
+            # Each batch has only one query (last position)
+            q_idx = cache_lens // block_size  # shape: [batch_size] = [total_q_len] in decoding
 
-        # compute attention score
+        # 计算attention score
         score = infllmv2_attn_stage1(
             q.contiguous(),
             k.contiguous(),
-            v.contiguous(),
+            k2.contiguous(),
             cu_seqlens_q=cu_seqlens_q,
             cu_seqlens_k=cu_seqlens_k,
+            cu_seqlens_v=cu_seqlens_k2,
             max_seqlen_q=max_seqlen_q,
             max_seqlen_k=max_seqlen_k,
-            causal=is_prefilling)
-        # Shape: [num_heads, total_q_len, num_blocks]
-        score = score[:, :q_idx.shape[0], :]
-
-        # Shape: [num_heads, total_q_len, num_blocks]
+            causal=is_prefilling
+        )
+        score = score[:, :q_idx.shape[0], :]  # [num_heads, total_q_len, num_blocks]
+        
         block_score = max_pooling_1d_varlen(
             score.contiguous(),
             cu_seqlens_q,
@@ -148,7 +127,9 @@ def compressed_attention(
             local_blocks=local_blocks,
             init_blocks=init_blocks,
             block_size=block_size,
-            stride=kernel_stride)
+            stride=kernel_stride
+        )  # shape: [num_heads, total_q_len, num_blocks]
+        
 
         # get topk
         topk = min(topk, block_score.shape[-1])
@@ -262,6 +243,11 @@ class InfLLMv2CacheLayer(DynamicLayer):
         self.no_compress_k_cache = []
         self.cached_compressed_cu_seqlens = torch.tensor([], dtype=torch.int32)
         self.compress_k_cache_varlen = torch.tensor([], dtype=torch.float32)
+        # Add support for compress_k2
+        self.compress_k2_cache = []
+        self.cached_compressed_cu_seqlens2 = torch.tensor([], dtype=torch.int32)
+        self.compress_k2_cache_varlen = torch.tensor([], dtype=torch.float32)
+        self.no_compress_k2_cache = []
 
     def update_no_rope_key(self, key_states):
         if self.no_rope_keys.numel() == 0:
@@ -303,12 +289,45 @@ class InfLLMv2CacheLayer(DynamicLayer):
                     k_chunk_list.append(None)
         return k_chunk_list
 
+    def update_compress_k2(self, key_states, cu_seqlens=None):
+        if len(self.compress_k2_cache) == 0:
+            if cu_seqlens is not None:
+                self.cached_compressed_cu_seqlens2 = cu_seqlens.clone()
+            self.compress_k2_cache_varlen = key_states
+            split_sizes = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
+            self.compress_k2_cache = list(torch.split(key_states, split_sizes))
+        else:
+            for index, k in enumerate(key_states):
+                if k is not None:
+                    self.compress_k2_cache[index] = torch.cat([self.compress_k2_cache[index], k], dim=0)
+            new_seq_lens = torch.tensor([tensor.shape[0] for tensor in self.compress_k2_cache], dtype=torch.int32)
+            new_cumsum = torch.cumsum(new_seq_lens, dim=0, dtype=torch.int32)
+            
+            self.compress_k2_cache_varlen = torch.cat(self.compress_k2_cache, dim=0)
+            self.cached_compressed_cu_seqlens2 = torch.cat([torch.tensor([0], dtype=torch.int32), new_cumsum]).to(self.compress_k2_cache_varlen.device)
+        return self.compress_k2_cache_varlen, self.cached_compressed_cu_seqlens2
+
+    def update_no_compress_k2(self, key_states, kernel_size=128, kernel_stride=64):
+        k_chunk_list = []
+        for index, k in enumerate(key_states):
+            if len(self.no_compress_k2_cache) <= index:
+                self.no_compress_k2_cache.append(k)
+            else:
+                self.no_compress_k2_cache[index] = torch.cat([self.no_compress_k2_cache[index], k], dim=0)
+                current_len = self.no_compress_k2_cache[index].shape[0]
+                if current_len >= kernel_size:
+                    k_chunk_list.append(self.no_compress_k2_cache[index][:kernel_size])
+                    self.no_compress_k2_cache[index] = self.no_compress_k2_cache[index][kernel_stride:]
+                else:
+                    k_chunk_list.append(None)
+        return k_chunk_list
+
 class InfLLMv2Cache(DynamicCache):
-    def __init__(self,
-                 config,num_hidden_layers: Optional[int] = None) -> None:
+    def __init__(self, config,num_hidden_layers: Optional[int] = None) -> None:
         super().__init__(config=config)
         self.layers = [InfLLMv2CacheLayer() for _ in range(num_hidden_layers)] if num_hidden_layers else []
         self._seen_tokens = 0
+        
 
     def update(self, key_states, value_states, layer_idx, cache_kwargs=None):
         if layer_idx == 0:
@@ -324,6 +343,12 @@ class InfLLMv2Cache(DynamicCache):
     def update_no_compress_k(self, key_states, layer_idx, kernel_size=32, kernel_stride=16, cache_kwargs=None):
         return self.layers[layer_idx].update_no_compress_k(key_states, kernel_size, kernel_stride)
 
+    def update_compress_k2(self, key_states, layer_idx, cu_seqlens=None, cache_kwargs=None):
+        return self.layers[layer_idx].update_compress_k2(key_states, cu_seqlens)
+
+    def update_no_compress_k2(self, key_states, layer_idx, kernel_size=128, kernel_stride=64, cache_kwargs=None):
+        return self.layers[layer_idx].update_no_compress_k2(key_states, kernel_size, kernel_stride)
+
     def crop(self, max_length):
         for layer in self.layers:
             layer.crop(max_length)
@@ -591,7 +616,6 @@ def _unpad_one_tensor(hidden_states, attention_mask):
     unpadded_states = index_first_axis(reshaped_states, indices)
     
     return unpadded_states, indices, cu_seqlens, max_seqlen_in_batch
-
 def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
     """
     This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
@@ -998,7 +1022,9 @@ class MiniCPMInfLLMv2Attention(MiniCPMAttention):
         self.local_blocks = self.window_size // self.block_size  # local_blocks
         self.topk = self.config.sparse_config.get('topk', 64) + (self.window_size//self.block_size)
         self.use_nope = self.config.sparse_config.get('use_nope', False)
+        
         self.compress_k = CompressK(self.num_key_value_heads, self.head_dim, kernel_size=self.kernel_size, kernel_stride=self.kernel_stride)
+        self.compress_k2 = CompressK(self.num_key_value_heads, self.head_dim, kernel_size=self.kernel_size*4, kernel_stride=self.kernel_stride*4)
 
     def forward(
         self,
@@ -1023,6 +1049,7 @@ class MiniCPMInfLLMv2Attention(MiniCPMAttention):
 
         bsz, q_len, _ = hidden_states.size()
 
+
         query_states = self.q_proj(hidden_states)
         key_states = self.k_proj(hidden_states)
         value_states = self.v_proj(hidden_states)
@@ -1053,11 +1080,12 @@ class MiniCPMInfLLMv2Attention(MiniCPMAttention):
         key_states = key_states.transpose(1, 2)
         value_states = value_states.transpose(1, 2)
         if self.use_nope:
-            key_states_no_rope = past_key_value.update_no_rope_key(key_states_no_rope, self.layer_idx)
+            key_states_no_rope =past_key_value.update_no_rope_key(key_states_no_rope, self.layer_idx)
             no_rope_param = {
                 'key_states_no_rope': key_states_no_rope,
                 'query_states_no_rope': query_states_no_rope,
             }
+            
         else:
             no_rope_param = None
 
@@ -1103,16 +1131,8 @@ class MiniCPMInfLLMv2Attention(MiniCPMAttention):
         return attn_output, attn_weights, past_key_value
 
     def _sparse_attention_forward(
-            self,
-            query_states,
-            key_states,
-            value_states,
-            attention_mask,
-            query_length,
-            dropout=0.0,
-            softmax_scale=None,
-            no_rope_param=None,
-            past_key_value=None):
+            self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None, no_rope_param=None, past_key_value=None
+        ):
             """
             Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
             first unpad the input, then computes the attention scores and pad the final attention scores.
@@ -1142,15 +1162,17 @@ class MiniCPMInfLLMv2Attention(MiniCPMAttention):
                 batch_size = query_states.shape[0]
                 # assert batch_size == 1, 'Only batch_size=1 is supported at the moment.'
                 if past_key_value!=None:
-                    compressed_k, compressed_cu_seqlens = self.get_compress_k(
+                    compressed_k, compressed_cu_seqlens, compressed_k2, compressed_cu_seqlens2 = self.get_compress_k(
                         key_states=key_states if self.use_nope ==False else no_rope_param['key_states_no_rope'],  # This can be optimized a bit;
                         attention_mask=attention_mask,
-                        past_key_value=past_key_value)
+                        past_key_value=past_key_value,
+                        
+                    )
 
                 query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
                     query_states, key_states, value_states, attention_mask, query_length
                 )
-
+                
                 cu_seqlens_q, cu_seqlens_k = cu_seq_lens
                 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
                 if no_rope_param != None:
@@ -1161,7 +1183,12 @@ class MiniCPMInfLLMv2Attention(MiniCPMAttention):
                 if past_key_value==None:
                     # compress_k use varlen form
                     compressed_k, compressed_cu_seqlens = self.compress_k(key_states,cu_seqlens_k)
+                    compressed_k2, compressed_cu_seqlens2 = self.compress_k2(key_states,cu_seqlens_k)
+                else:
+                    # compressed_k and compressed_k2 already retrieved from get_compress_k above
+                    pass
 
+                
                 attn_output_unpad = self.sparse_forward(
                     query_states,
                     key_states,
@@ -1171,15 +1198,16 @@ class MiniCPMInfLLMv2Attention(MiniCPMAttention):
                     max_seqlen_in_batch_q,
                     max_seqlen_in_batch_k,
                     no_rope_param=no_rope_param,
-                    compressed_k=compressed_k,
-                    compressed_cu_seqlens=compressed_cu_seqlens)
+                    compressed_k=compressed_k, compressed_cu_seqlens=compressed_cu_seqlens,
+                    compressed_k2=compressed_k2, compressed_cu_seqlens2=compressed_cu_seqlens2
+                )
 
                 attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
+                
             else:
                 raise ValueError('Need attention mask')
 
             return attn_output
-
     def get_compress_k(self, key_states, attention_mask, past_key_value):
         """
         Get compressed key states and corresponding cumulative sequence lengths.
@@ -1191,34 +1219,51 @@ class MiniCPMInfLLMv2Attention(MiniCPMAttention):
             no_rope_param: Optional parameter containing key states without rope
             
         Returns:
-            Tuple of (compressed_k, compressed_cu_seqlens)
+            Tuple of (compressed_k, compressed_cu_seqlens, compressed_k2, compressed_cu_seqlens2)
         """
+        
         # Check if this is prefilling or initial compression condition
+
         is_prefilling = (
             key_states.shape[1] >= self.dense_len and
             (
                 not past_key_value.layers[self.layer_idx].compress_k_cache
             )
         )
-
+        
         if is_prefilling:
             unpadded_key_states, indices, cu_seqlens, max_seqlen_in_batch = _unpad_one_tensor(key_states,attention_mask=attention_mask)
             # Compress the keys
             compressed_k, compressed_cu_seqlens = self.compress_k(unpadded_key_states, cu_seqlens)
-
+            compressed_k2, compressed_cu_seqlens2 = self.compress_k2(unpadded_key_states, cu_seqlens)
+            
             past_key_value.update_compress_k(
                 compressed_k, self.layer_idx, compressed_cu_seqlens)
-
+            past_key_value.update_compress_k2(
+                compressed_k2, self.layer_idx, compressed_cu_seqlens2)
+            
             no_compress_k_list = []
             # Compute and update no_compress_k
             for i in range(len(compressed_cu_seqlens)-1):
                 no_compress_k_start = (compressed_cu_seqlens[i+1]- compressed_cu_seqlens[i]) * self.kernel_stride
-
+                
                 no_compress_k_list.append(unpadded_key_states[cu_seqlens[i]+no_compress_k_start:cu_seqlens[i+1]].clone())
 
             past_key_value.update_no_compress_k(
                 no_compress_k_list, self.layer_idx,kernel_stride=self.kernel_stride, 
                 kernel_size=self.kernel_size)
+            
+            # Also update no_compress_k2
+            no_compress_k2_list = []
+            for i in range(len(compressed_cu_seqlens2)-1):
+                no_compress_k2_start = (compressed_cu_seqlens2[i+1]- compressed_cu_seqlens2[i]) * self.kernel_stride * 4
+                
+                no_compress_k2_list.append(unpadded_key_states[cu_seqlens[i]+no_compress_k2_start:cu_seqlens[i+1]].clone())
+
+            past_key_value.update_no_compress_k2(
+                no_compress_k2_list, self.layer_idx,kernel_stride=self.kernel_stride*4, 
+                kernel_size=self.kernel_size*4)
+                
         else:
             # Decode case: incremental update
             batch_size = key_states.shape[0] # key_states.shape = [batch_size, seq, k_head_num, head_dim]
@@ -1233,16 +1278,32 @@ class MiniCPMInfLLMv2Attention(MiniCPMAttention):
                 kernel_size=self.kernel_size)
             new_compressed_k_list = []
             for no_compress_k in no_compress_k_list:
+
                 if no_compress_k is not None:
                     # We have enough tokens to compress
                     new_compressed_k = no_compress_k.mean(dim=0, keepdim=True)  # [1, n_heads_k, head_dim]
+                    
                     new_compressed_k_list.append(new_compressed_k)
                 else:
                     new_compressed_k_list.append(None)
             compressed_k, compressed_cu_seqlens = past_key_value.update_compress_k(new_compressed_k_list, self.layer_idx,)
-
-        return compressed_k, compressed_cu_seqlens
-
+            
+            # For compress_k2, update no_compress_k2 buffer and compress when ready
+            no_compress_k2_list = past_key_value.update_no_compress_k2(
+                key_states_split, self.layer_idx, 
+                kernel_stride=self.kernel_stride*4, 
+                kernel_size=self.kernel_size*4)
+            new_compressed_k2_list = []
+            for no_compress_k2 in no_compress_k2_list:
+                if no_compress_k2 is not None:
+                    # We have enough tokens to compress for k2
+                    new_compressed_k2 = no_compress_k2.mean(dim=0, keepdim=True)  # [1, n_heads_k, head_dim]
+                    new_compressed_k2_list.append(new_compressed_k2)
+                else:
+                    new_compressed_k2_list.append(None)
+            compressed_k2, compressed_cu_seqlens2 = past_key_value.update_compress_k2(new_compressed_k2_list, self.layer_idx,)
+        
+        return compressed_k, compressed_cu_seqlens, compressed_k2, compressed_cu_seqlens2
     def sparse_forward(self,
                        query_layer,
                        key_layer,
@@ -1252,8 +1313,8 @@ class MiniCPMInfLLMv2Attention(MiniCPMAttention):
                        max_seqlen_in_batch_q,
                        max_seqlen_in_batch_k,
                        no_rope_param=None,
-                       compressed_k=None,
-                       compressed_cu_seqlens=None):
+                       compressed_k=None, compressed_cu_seqlens=None,
+                       compressed_k2=None, compressed_cu_seqlens2=None):
         compressed_seqlens = compressed_cu_seqlens[1:] - compressed_cu_seqlens[:-1]
         cache_lens = None
         if max_seqlen_in_batch_q==1 and max_seqlen_in_batch_k>1: #decoding
@@ -1263,13 +1324,14 @@ class MiniCPMInfLLMv2Attention(MiniCPMAttention):
         topk_idx = compressed_attention(
             query_layer if no_rope_param is None else no_rope_param['query_states_no_rope'],
             compressed_k,
-            compressed_k.clone(),
+            compressed_k2,
             self.kernel_size,
             self.kernel_stride,
             self.block_size,
             self.topk,
             cu_seqlens_q,
             compressed_cu_seqlens,
+            compressed_cu_seqlens2,
             max_seqlen_in_batch_q,
             compressed_seqlens.max().item(),
             None,