build-small-hackathon/MiniCPM4-8B-PaperProf
011
1# coding=utf-82# Copyright 2025 The OpenBMB Team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15""" PyTorch MiniCPM model."""16import math17import re18import warnings19from typing import Any, Dict, List, Optional, Tuple, Union20 21import torch22import torch.nn.functional as F23import torch.utils.checkpoint24from torch import nn25from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss26from transformers.activations import ACT2FN27from transformers.cache_utils import Cache, DynamicCache, CacheLayerMixin, DynamicLayer28from transformers.modeling_attn_mask_utils import (29 AttentionMaskConverter,30 _prepare_4d_attention_mask,31 _prepare_4d_causal_attention_mask,32 _prepare_4d_causal_attention_mask_for_sdpa,33)34from transformers.modeling_outputs import (35 BaseModelOutputWithPast,36 CausalLMOutputWithPast,37 SequenceClassifierOutputWithPast,38)39from transformers.modeling_utils import PreTrainedModel40from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_1341from transformers.utils import (42 add_start_docstrings,43 add_start_docstrings_to_model_forward,44 is_flash_attn_greater_or_equal_2_10,45 logging,46 replace_return_docstrings,47)48from transformers.utils.import_utils import is_torch_fx_available49 50 51 52from .configuration_minicpm import MiniCPMConfig #!一定要改53 54try:55 from flash_attn import flash_attn_func, flash_attn_varlen_func56 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa57 from infllm_v2 import (58 infllmv2_attn_stage1,59 infllmv2_attn_varlen_func,60 infllmv2_attn_with_kvcache,61 max_pooling_1d,62 max_pooling_1d_varlen63 )64except:65 pass66 67from functools import lru_cache68 69 70def compressed_attention(71 q: torch.Tensor,72 k: torch.Tensor,73 k2: torch.Tensor,74 kernel_size: int,75 kernel_stride: int,76 block_size: int,77 topk: int,78 cu_seqlens_q: torch.Tensor,79 cu_seqlens_k: torch.Tensor,80 cu_seqlens_k2: torch.Tensor,81 max_seqlen_q: int,82 max_seqlen_k: int,83 sm_scale: float = None,84 init_blocks: int = 1,85 local_blocks: int = 2,86 cache_lens=None,87) -> Tuple[torch.Tensor, torch.Tensor]:88 with torch.no_grad():89 batch_size = cu_seqlens_q.shape[0] - 190 91 # Check if it's prefilling stage92 is_prefilling = cache_lens is None or (cache_lens == 0).all().item()93 94 if is_prefilling: # prefilling stage95 # Calculate q_idx for each query position in each batch96 cache_lens = torch.zeros(batch_size, dtype=torch.int32, device=q.device) 97 q_idx = torch.cat([98 (torch.arange(cu_seqlens_q[i + 1] - cu_seqlens_q[i], device=q.device) + 99 max_seqlen_q - (cu_seqlens_q[i + 1] - cu_seqlens_q[i])) // block_size100 for i in range(batch_size)101 ], dim=0) # shape: [total_q_len]102 else: # decoding stage103 # Each batch has only one query (last position)104 q_idx = cache_lens // block_size # shape: [batch_size] = [total_q_len] in decoding105 106 # 计算attention score107 score = infllmv2_attn_stage1(108 q.contiguous(),109 k.contiguous(),110 k2.contiguous(),111 cu_seqlens_q=cu_seqlens_q,112 cu_seqlens_k=cu_seqlens_k,113 cu_seqlens_v=cu_seqlens_k2,114 max_seqlen_q=max_seqlen_q,115 max_seqlen_k=max_seqlen_k,116 causal=is_prefilling117 )118 score = score[:, :q_idx.shape[0], :] # [num_heads, total_q_len, num_blocks]119 120 block_score = max_pooling_1d_varlen(121 score.contiguous(),122 cu_seqlens_q,123 cu_seqlens_k,124 cache_lens,125 max_seqlen_q,126 max_seqlen_k,127 local_blocks=local_blocks,128 init_blocks=init_blocks,129 block_size=block_size,130 stride=kernel_stride131 ) # shape: [num_heads, total_q_len, num_blocks]132 133 134 # get topk135 topk = min(topk, block_score.shape[-1])136 topk_idx = block_score.topk(topk, dim=-1).indices.sort(-1).values137 topk_idx[topk_idx > q_idx[None, :, None]] = -1138 topk_idx = topk_idx.to(torch.int32)139 140 return topk_idx141 142 143@lru_cache(maxsize=16)144def calc_chunks_with_stride(cu_seqlen, chunk_size, kernel_stride):145 """146 Compute the chunks that require Sparse attention, with stride support.147 148 Args:149 cu_seqlen (torch.Tensor): Cumulative sequence lengths for each sample.150 chunk_size (int): Chunk size used for Sparse attention.151 kernel_stride (int): Stride size when sliding over the sequence.152 153 Returns:154 filtered_indices (torch.Tensor): Indices used to directly index into the key/value tensors.155 cu_seqlens_compressed (torch.Tensor): Cumulative sequence lengths after compression.156 """157 # 1. Compute the length of each sequence158 batch_sizes = cu_seqlen[1:] - cu_seqlen[:-1]159 160 # 2. Compute the start positions of chunks for each sequence (with stride)161 max_seq_len = torch.max(batch_sizes)162 max_num_chunks_per_seq = (max_seq_len - chunk_size) // kernel_stride + 1163 chunk_start_offsets = torch.arange(0, max_num_chunks_per_seq * kernel_stride, kernel_stride, device=cu_seqlen.device)164 seq_starts = cu_seqlen[:-1]165 chunk_start_in_seq = seq_starts[:, None] + chunk_start_offsets[None, :] # [batch_size, max_num_chunks_per_seq]166 167 # 3. Filter out chunks that exceed sequence length or are smaller than the full chunk size168 chunk_end_in_seq = chunk_start_in_seq + chunk_size169 valid_chunk_mask = (chunk_end_in_seq <= (seq_starts[:, None] + batch_sizes[:, None]))170 171 # 4. Filter valid chunk start positions using the valid_chunk_mask172 valid_chunk_starts = chunk_start_in_seq[valid_chunk_mask] # [num_valid_chunks]173 del chunk_start_in_seq174 # 5. Generate filtered_indices175 chunk_indices = torch.arange(176 0, chunk_size, device=cu_seqlen.device177 )[None, :] # [1, chunk_size]178 filtered_indices = valid_chunk_starts[:, None] + chunk_indices # [num_valid_chunks, chunk_size]179 filtered_indices = filtered_indices.view(-1) # Flatten to 1D indices180 181 # 6. Compute compressed cumulative sequence lengths182 num_filtered_chunks_per_batch = valid_chunk_mask.sum(dim=1) # Number of valid chunks per batch183 cu_seqlens_compressed = torch.zeros(184 len(cu_seqlen), dtype=torch.int32, device=cu_seqlen.device185 )186 cu_seqlens_compressed[1:] = num_filtered_chunks_per_batch.cumsum(dim=0)187 del num_filtered_chunks_per_batch, chunk_start_offsets, seq_starts, chunk_end_in_seq, valid_chunk_mask, chunk_indices188 return filtered_indices, cu_seqlens_compressed189 190 191class CompressK(torch.nn.Module):192 def __init__(self, head_num_k, head_dim, kernel_size, kernel_stride=16):193 """194 Module for compressing key (K) representations.195 196 Args:197 head_num_k (int): Number of key attention heads.198 head_dim (int): Dimension of each attention head.199 kernel_size (int): Size of each chunk used for compression.200 kernel_stride (int, optional): Stride used when dividing input into chunks. Default is 16.201 """202 super().__init__()203 self.kernel_size = kernel_size204 self.head_num_k = head_num_k205 self.head_dim = head_dim206 self.kernel_stride = kernel_stride207 208 def forward(self, k: torch.Tensor, cu_seqlens):209 """210 Forward pass for compressing the key (K) tensor.211 212 Args:213 k (torch.Tensor): Input key tensor of shape (total_seq_len, num_heads, head_dim).214 cu_seqlens (torch.Tensor): Cumulative sequence lengths for each sample in the batch, typically used for handling variable-length sequences.215 216 Returns:217 compress_k (torch.Tensor): Compressed key tensor.218 cu_seqlens_compressed (torch.Tensor): Updated cumulative sequence lengths after compression.219 220 """221 # Compute chunk-related metadata, with stride support222 filtered_k_indices, cu_seqlens_compressed = calc_chunks_with_stride(223 cu_seqlens, self.kernel_size, self.kernel_stride224 )225 226 # Extract filtered key vectors227 filtered_k = k.index_select(0, filtered_k_indices.view(-1))228 229 # split230 filtered_k = filtered_k.view(filtered_k.shape[0] // self.kernel_size, self.kernel_size, self.head_num_k, self.head_dim) # [l, block_size,h,d]231 232 compressed_k = filtered_k.mean(dim=1)233 return compressed_k, cu_seqlens_compressed234 235 236 237class InfLLMv2CacheLayer(DynamicLayer):238 def __init__(self):239 super().__init__()240 # Initialize any additional attributes specific to InfLLMv2CacheLayer241 self.no_rope_keys = torch.tensor([], dtype=torch.float32)242 self.compress_k_cache = []243 self.no_compress_k_cache = []244 self.cached_compressed_cu_seqlens = torch.tensor([], dtype=torch.int32)245 self.compress_k_cache_varlen = torch.tensor([], dtype=torch.float32)246 # Add support for compress_k2247 self.compress_k2_cache = []248 self.cached_compressed_cu_seqlens2 = torch.tensor([], dtype=torch.int32)249 self.compress_k2_cache_varlen = torch.tensor([], dtype=torch.float32)250 self.no_compress_k2_cache = []251 252 def update_no_rope_key(self, key_states):253 if self.no_rope_keys.numel() == 0:254 self.no_rope_keys = key_states255 else:256 self.no_rope_keys = torch.cat([self.no_rope_keys, key_states], dim=1)257 return self.no_rope_keys258 259 def update_compress_k(self, key_states, cu_seqlens=None):260 if len(self.compress_k_cache) == 0:261 if cu_seqlens is not None:262 self.cached_compressed_cu_seqlens = cu_seqlens.clone()263 self.compress_k_cache_varlen = key_states264 split_sizes = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()265 self.compress_k_cache = list(torch.split(key_states, split_sizes))266 else:267 for index, k in enumerate(key_states):268 if k is not None:269 self.compress_k_cache[index] = torch.cat([self.compress_k_cache[index], k], dim=0)270 new_seq_lens = torch.tensor([tensor.shape[0] for tensor in self.compress_k_cache], dtype=torch.int32)271 new_cumsum = torch.cumsum(new_seq_lens, dim=0, dtype=torch.int32)272 273 self.compress_k_cache_varlen = torch.cat(self.compress_k_cache, dim=0)274 self.cached_compressed_cu_seqlens = torch.cat([torch.tensor([0], dtype=torch.int32), new_cumsum]).to(self.compress_k_cache_varlen.device)275 return self.compress_k_cache_varlen, self.cached_compressed_cu_seqlens276 277 def update_no_compress_k(self, key_states, kernel_size=32, kernel_stride=16):278 k_chunk_list = []279 for index, k in enumerate(key_states):280 if len(self.no_compress_k_cache) <= index:281 self.no_compress_k_cache.append(k)282 else:283 self.no_compress_k_cache[index] = torch.cat([self.no_compress_k_cache[index], k], dim=0)284 current_len = self.no_compress_k_cache[index].shape[0]285 if current_len >= kernel_size:286 k_chunk_list.append(self.no_compress_k_cache[index][:kernel_size])287 self.no_compress_k_cache[index] = self.no_compress_k_cache[index][kernel_stride:]288 else:289 k_chunk_list.append(None)290 return k_chunk_list291 292 def update_compress_k2(self, key_states, cu_seqlens=None):293 if len(self.compress_k2_cache) == 0:294 if cu_seqlens is not None:295 self.cached_compressed_cu_seqlens2 = cu_seqlens.clone()296 self.compress_k2_cache_varlen = key_states297 split_sizes = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()298 self.compress_k2_cache = list(torch.split(key_states, split_sizes))299 else:300 for index, k in enumerate(key_states):301 if k is not None:302 self.compress_k2_cache[index] = torch.cat([self.compress_k2_cache[index], k], dim=0)303 new_seq_lens = torch.tensor([tensor.shape[0] for tensor in self.compress_k2_cache], dtype=torch.int32)304 new_cumsum = torch.cumsum(new_seq_lens, dim=0, dtype=torch.int32)305 306 self.compress_k2_cache_varlen = torch.cat(self.compress_k2_cache, dim=0)307 self.cached_compressed_cu_seqlens2 = torch.cat([torch.tensor([0], dtype=torch.int32), new_cumsum]).to(self.compress_k2_cache_varlen.device)308 return self.compress_k2_cache_varlen, self.cached_compressed_cu_seqlens2309 310 def update_no_compress_k2(self, key_states, kernel_size=128, kernel_stride=64):311 k_chunk_list = []312 for index, k in enumerate(key_states):313 if len(self.no_compress_k2_cache) <= index:314 self.no_compress_k2_cache.append(k)315 else:316 self.no_compress_k2_cache[index] = torch.cat([self.no_compress_k2_cache[index], k], dim=0)317 current_len = self.no_compress_k2_cache[index].shape[0]318 if current_len >= kernel_size:319 k_chunk_list.append(self.no_compress_k2_cache[index][:kernel_size])320 self.no_compress_k2_cache[index] = self.no_compress_k2_cache[index][kernel_stride:]321 else:322 k_chunk_list.append(None)323 return k_chunk_list324 325class InfLLMv2Cache(DynamicCache):326 def __init__(self, config,num_hidden_layers: Optional[int] = None) -> None:327 super().__init__(config=config)328 self.layers = [InfLLMv2CacheLayer() for _ in range(num_hidden_layers)] if num_hidden_layers else []329 self._seen_tokens = 0330 331 332 def update(self, key_states, value_states, layer_idx, cache_kwargs=None):333 if layer_idx == 0:334 self._seen_tokens += key_states.shape[-2]335 return self.layers[layer_idx].update(key_states, value_states, cache_kwargs)336 337 def update_no_rope_key(self, key_states, layer_idx, cache_kwargs=None):338 return self.layers[layer_idx].update_no_rope_key(key_states)339 340 def update_compress_k(self, key_states, layer_idx, cu_seqlens=None, cache_kwargs=None):341 return self.layers[layer_idx].update_compress_k(key_states, cu_seqlens)342 343 def update_no_compress_k(self, key_states, layer_idx, kernel_size=32, kernel_stride=16, cache_kwargs=None):344 return self.layers[layer_idx].update_no_compress_k(key_states, kernel_size, kernel_stride)345 346 def update_compress_k2(self, key_states, layer_idx, cu_seqlens=None, cache_kwargs=None):347 return self.layers[layer_idx].update_compress_k2(key_states, cu_seqlens)348 349 def update_no_compress_k2(self, key_states, layer_idx, kernel_size=128, kernel_stride=64, cache_kwargs=None):350 return self.layers[layer_idx].update_no_compress_k2(key_states, kernel_size, kernel_stride)351 352 def crop(self, max_length):353 for layer in self.layers:354 layer.crop(max_length)355 356 def batch_repeat_interleave(self, repeats):357 for layer in self.layers:358 layer.batch_repeat_interleave(repeats)359 360 def batch_select_indices(self, indices):361 for layer in self.layers:362 layer.batch_select_indices(indices)363 364 365# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.366# It means that the function will not be traced through and simply appear as a node in the graph.367if is_torch_fx_available():368 if not is_torch_greater_or_equal_than_1_13:369 import torch.fx370 371 _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)372 373 374logger = logging.get_logger(__name__)375 376_CONFIG_FOR_DOC = 'MiniCPMConfig'377 378 379def _get_unpad_data(attention_mask):380 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)381 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()382 max_seqlen_in_batch = seqlens_in_batch.max().item()383 cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))384 return (385 indices,386 cu_seqlens,387 max_seqlen_in_batch,388 )389 390 391 392 393# @torch.jit.script # type: ignore394def rms_layernorm(hidden: torch.Tensor, weight: torch.Tensor, eps: float):395 old_dtype = hidden.dtype396 variance = hidden.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)397 hidden = (hidden * torch.rsqrt(variance + eps)).to(old_dtype)398 return hidden * weight399 400 401class MiniCPMRMSNorm(nn.Module):402 def __init__(self, hidden_size, eps=1e-6):403 """404 MiniCPMRMSNorm is equivalent to T5LayerNorm405 """406 super().__init__()407 self.weight = nn.Parameter(torch.ones(hidden_size))408 self.variance_epsilon = eps409 410 def forward(self, hidden_states):411 return rms_layernorm(hidden_states, self.weight, self.variance_epsilon)412 413 414ALL_LAYERNORM_LAYERS.append(MiniCPMRMSNorm)415 416 417class MiniCPMRotaryEmbedding(nn.Module):418 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):419 super().__init__()420 421 self.dim = dim422 self.max_position_embeddings = max_position_embeddings423 self.base = base424 inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))425 self.register_buffer('inv_freq', inv_freq, persistent=False)426 427 # Build here to make `torch.jit.trace` work.428 self._set_cos_sin_cache(429 # seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()430 seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.float32431 )432 433 def _set_cos_sin_cache(self, seq_len, device, dtype):434 self.max_seq_len_cached = seq_len435 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)436 freqs = torch.outer(t, self.inv_freq)437 # Different from paper, but it uses a different permutation in order to obtain the same calculation438 emb = torch.cat((freqs, freqs), dim=-1)439 440 self.register_buffer('cos_cached', emb.cos().to(dtype), persistent=False)441 self.register_buffer('sin_cached', emb.sin().to(dtype), persistent=False)442 443 def forward(self, x, seq_len=None):444 # x: [bs, num_attention_heads, seq_len, head_size]445 if seq_len > self.max_seq_len_cached:446 self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)447 448 return (449 self.cos_cached[:seq_len].to(dtype=x.dtype),450 self.sin_cached[:seq_len].to(dtype=x.dtype),451 )452 453 454class MiniCPMLongRoPE(MiniCPMRotaryEmbedding):455 """MiniCPMRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""456 457 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, short_factor=None, long_factor=None, original_max_position_embeddings=None):458 self.short_factor = short_factor459 self.long_factor = long_factor460 self.original_max_position_embeddings = original_max_position_embeddings461 scale = (max_position_embeddings / self.original_max_position_embeddings)462 self.scaling_factor = math.sqrt(1 + math.log(scale) / math.log(self.original_max_position_embeddings))463 super().__init__(dim, max_position_embeddings, base, device)464 465 def _set_cos_sin_cache(self, seq_len, device, dtype):466 self.max_seq_len_cached = seq_len467 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)468 if seq_len > self.original_max_position_embeddings:469 ext_factors = torch.tensor(self.long_factor, dtype=torch.float32, device=device)470 else:471 ext_factors = torch.tensor(self.short_factor, dtype=torch.float32, device=device)472 473 freqs = torch.mul(474 torch.outer(t, 1.0 / ext_factors).to(device=device),475 self.inv_freq.to(device=device).to(dtype)476 )477 # Different from paper, but it uses a different permutation in order to obtain the same calculation478 emb = torch.cat((freqs, freqs), dim=-1)479 self.register_buffer('cos_cached', emb.cos().to(dtype) * self.scaling_factor, persistent=False)480 self.register_buffer('sin_cached', emb.sin().to(dtype) * self.scaling_factor, persistent=False)481 482 483class MiniCPMLinearScalingRotaryEmbedding(MiniCPMRotaryEmbedding):484 """MiniCPMRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""485 486 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):487 self.scaling_factor = scaling_factor488 super().__init__(dim, max_position_embeddings, base, device)489 490 def _set_cos_sin_cache(self, seq_len, device, dtype):491 self.max_seq_len_cached = seq_len492 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)493 t = t / self.scaling_factor494 495 freqs = torch.outer(t, self.inv_freq)496 # Different from paper, but it uses a different permutation in order to obtain the same calculation497 emb = torch.cat((freqs, freqs), dim=-1)498 self.register_buffer('cos_cached', emb.cos().to(dtype), persistent=False)499 self.register_buffer('sin_cached', emb.sin().to(dtype), persistent=False)500 501 502class MiniCPMDynamicNTKScalingRotaryEmbedding(MiniCPMRotaryEmbedding):503 """MiniCPMRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""504 505 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):506 self.scaling_factor = scaling_factor507 super().__init__(dim, max_position_embeddings, base, device)508 509 def _set_cos_sin_cache(self, seq_len, device, dtype):510 self.max_seq_len_cached = seq_len511 512 if seq_len > self.max_position_embeddings:513 base = self.base * (514 (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)515 ) ** (self.dim / (self.dim - 2))516 inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))517 self.register_buffer('inv_freq', inv_freq, persistent=False)518 519 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)520 521 freqs = torch.outer(t, self.inv_freq)522 # Different from paper, but it uses a different permutation in order to obtain the same calculation523 emb = torch.cat((freqs, freqs), dim=-1)524 525 self.register_buffer('cos_cached', emb.cos().to(dtype), persistent=False)526 self.register_buffer('sin_cached', emb.sin().to(dtype), persistent=False)527 528 529def rotate_half(x):530 """Rotates half the hidden dims of the input."""531 x1 = x[..., : x.shape[-1] // 2]532 x2 = x[..., x.shape[-1] // 2:]533 return torch.cat((-x2, x1), dim=-1)534 535 536def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):537 """Applies Rotary Position Embedding to the query and key tensors.538 539 Args:540 q (`torch.Tensor`): The query tensor.541 k (`torch.Tensor`): The key tensor.542 cos (`torch.Tensor`): The cosine part of the rotary embedding.543 sin (`torch.Tensor`): The sine part of the rotary embedding.544 position_ids (`torch.Tensor`):545 The position indices of the tokens corresponding to the query and key tensors. For example, this can be546 used to pass offsetted position ids when working with a KV-cache.547 unsqueeze_dim (`int`, *optional*, defaults to 1):548 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and549 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note550 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and551 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes552 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have553 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.554 Returns:555 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.556 """557 # cos = cos[position_ids].unsqueeze(unsqueeze_dim)558 # sin = sin[position_ids].unsqueeze(unsqueeze_dim)559 # q_embed = (q * cos) + (rotate_half(q) * sin)560 # k_embed = (k * cos) + (rotate_half(k) * sin)561 orig_dtype = k.dtype562 cos = cos[position_ids].unsqueeze(unsqueeze_dim) # [bs, 1, seq_len, dim]563 sin = sin[position_ids].unsqueeze(unsqueeze_dim) # [bs, 1, seq_len, dim]564 q_fp32 = q.to(dtype=torch.float32, device=q.device)565 k_fp32 = k.to(dtype=torch.float32, device=k.device)566 q_embed = (q_fp32 * cos) + (rotate_half(q_fp32) * sin)567 k_embed = (k_fp32 * cos) + (rotate_half(k_fp32) * sin)568 return q_embed.to(dtype=orig_dtype), k_embed.to(dtype=orig_dtype)569 570 571class MiniCPMMLP(nn.Module):572 def __init__(self, config):573 super().__init__()574 self.config = config575 self.hidden_size = config.hidden_size576 self.intermediate_size = config.intermediate_size577 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)578 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)579 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)580 self.act_fn = ACT2FN[config.hidden_act]581 582 def forward(self, x):583 if self.config.pretraining_tp > 1:584 slice = self.intermediate_size // self.config.pretraining_tp585 gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)586 up_proj_slices = self.up_proj.weight.split(slice, dim=0)587 down_proj_slices = self.down_proj.weight.split(slice, dim=1)588 589 gate_proj = torch.cat(590 [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1591 )592 up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)593 594 intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)595 down_proj = [596 F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp)597 ]598 down_proj = sum(down_proj)599 else:600 down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))601 602 return down_proj603 604def _unpad_one_tensor(hidden_states, attention_mask):605 # Unpad the hidden states using the indices606 indices, cu_seqlens, max_seqlen_in_batch = _get_unpad_data(attention_mask)607 batch_size, seq_len = hidden_states.shape[:2]608 609 # Get the remaining dimensions610 remaining_dims = hidden_states.shape[2:]611 612 # Reshape to (batch_size * seq_len, *remaining_dims)613 reshaped_states = hidden_states.reshape(batch_size * seq_len, *remaining_dims)614 615 # Apply unpadding using indices616 unpadded_states = index_first_axis(reshaped_states, indices)617 618 return unpadded_states, indices, cu_seqlens, max_seqlen_in_batch619def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:620 """621 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,622 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)623 """624 batch, num_key_value_heads, slen, head_dim = hidden_states.shape625 if n_rep == 1:626 return hidden_states627 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)628 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)629 630 631class MiniCPMAttention(nn.Module):632 """Multi-headed attention from 'Attention Is All You Need' paper"""633 634 def __init__(self, config: MiniCPMConfig, layer_idx: Optional[int] = None):635 super().__init__()636 self.config = config637 self.layer_idx = layer_idx638 if layer_idx is None:639 logger.warning_once(640 f'Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will '641 'to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` '642 'when creating this class.'643 )644 645 self.attention_dropout = config.attention_dropout646 self.hidden_size = config.hidden_size647 self.num_heads = config.num_attention_heads648 self.head_dim = self.hidden_size // self.num_heads649 self.num_key_value_heads = config.num_key_value_heads650 self.num_key_value_groups = self.num_heads // self.num_key_value_heads651 self.max_position_embeddings = config.max_position_embeddings652 self.rope_theta = config.rope_theta653 self.is_causal = True654 655 if (self.head_dim * self.num_heads) != self.hidden_size:656 raise ValueError(657 f'hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}'658 f' and `num_heads`: {self.num_heads}).'659 )660 661 self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)662 self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)663 self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)664 self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)665 self._init_rope()666 667 def _init_rope(self):668 if self.config.rope_scaling is None:669 self.rotary_emb = MiniCPMRotaryEmbedding(670 self.head_dim,671 max_position_embeddings=self.max_position_embeddings,672 base=self.rope_theta,673 )674 else:675 scaling_type = self.config.rope_scaling['rope_type']676 scaling_factor = self.config.rope_scaling.get('factor', None)677 if scaling_type == 'linear':678 self.rotary_emb = MiniCPMLinearScalingRotaryEmbedding(679 self.head_dim,680 max_position_embeddings=self.max_position_embeddings,681 scaling_factor=scaling_factor,682 base=self.rope_theta,683 )684 elif scaling_type == 'dynamic':685 self.rotary_emb = MiniCPMDynamicNTKScalingRotaryEmbedding(686 self.head_dim,687 max_position_embeddings=self.max_position_embeddings,688 scaling_factor=scaling_factor,689 base=self.rope_theta,690 )691 elif scaling_type == 'longrope':692 self.rotary_emb = MiniCPMLongRoPE(693 self.head_dim,694 max_position_embeddings=self.max_position_embeddings,695 short_factor=self.config.rope_scaling['short_factor'],696 long_factor=self.config.rope_scaling['long_factor'],697 base=self.rope_theta,698 original_max_position_embeddings=self.config.rope_scaling['original_max_position_embeddings']699 )700 else:701 raise ValueError(f'Unknown RoPE scaling type {scaling_type}')702 703 def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):704 return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()705 706 def forward(707 self,708 hidden_states: torch.Tensor,709 attention_mask: Optional[torch.Tensor] = None,710 position_ids: Optional[torch.LongTensor] = None,711 past_key_value: Optional[Cache] = None,712 output_attentions: bool = False,713 use_cache: bool = False,714 **kwargs,715 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:716 if 'padding_mask' in kwargs:717 warnings.warn(718 'Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`'719 )720 721 bsz, q_len, _ = hidden_states.size()722 723 if self.config.pretraining_tp > 1:724 key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp725 query_slices = self.q_proj.weight.split(726 (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0727 )728 key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)729 value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)730 731 query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]732 query_states = torch.cat(query_states, dim=-1)733 734 key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]735 key_states = torch.cat(key_states, dim=-1)736 737 value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]738 value_states = torch.cat(value_states, dim=-1)739 740 else:741 query_states = self.q_proj(hidden_states)742 key_states = self.k_proj(hidden_states)743 value_states = self.v_proj(hidden_states)744 745 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)746 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)747 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)748 749 kv_seq_len = position_ids.max().item() + 1750 cos, sin = self.rotary_emb(value_states.to(torch.float32), seq_len=kv_seq_len)751 752 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)753 754 if past_key_value is not None:755 cache_kwargs = {'sin': sin, 'cos': cos} # Specific to RoPE models756 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)757 758 key_states = repeat_kv(key_states, self.num_key_value_groups)759 value_states = repeat_kv(value_states, self.num_key_value_groups)760 761 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)762 if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):763 raise ValueError(764 f'Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is'765 f' {attn_weights.size()}'766 )767 768 if attention_mask is not None:769 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):770 raise ValueError(771 f'Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}'772 )773 attn_weights = attn_weights + attention_mask774 775 # upcast attention to fp32776 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)777 attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)778 attn_output = torch.matmul(attn_weights, value_states)779 780 if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):781 raise ValueError(782 f'`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is'783 f' {attn_output.size()}'784 )785 786 attn_output = attn_output.transpose(1, 2).contiguous()787 788 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)789 790 if self.config.pretraining_tp > 1:791 attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)792 o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)793 attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])794 else:795 attn_output = self.o_proj(attn_output)796 797 if not output_attentions:798 attn_weights = None799 800 return attn_output, attn_weights, past_key_value801 802 803class MiniCPMFlashAttention2(MiniCPMAttention):804 """805 MiniCPM flash attention module. This module inherits from `MiniCPMAttention` as the weights of the module stays806 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of807 flash attention and deal with padding tokens in case the input contains any of them.808 """809 810 def __init__(self, *args, **kwargs):811 super().__init__(*args, **kwargs)812 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.813 # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignment, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.814 # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).815 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()816 817 def forward(818 self,819 hidden_states: torch.Tensor,820 attention_mask: Optional[torch.LongTensor] = None,821 position_ids: Optional[torch.LongTensor] = None,822 past_key_value: Optional[Cache] = None,823 output_attentions: bool = False,824 use_cache: bool = False,825 **kwargs,826 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:827 # MiniCPMFlashAttention2 attention does not support output_attentions828 if 'padding_mask' in kwargs:829 warnings.warn(830 'Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`'831 )832 833 # overwrite attention_mask with padding_mask834 attention_mask = kwargs.pop('padding_mask')835 836 output_attentions = False837 838 bsz, q_len, _ = hidden_states.size()839 840 query_states = self.q_proj(hidden_states)841 key_states = self.k_proj(hidden_states)842 value_states = self.v_proj(hidden_states)843 844 # Flash attention requires the input to have the shape845 # batch_size x seq_length x head_dim x hidden_dim846 # therefore we just need to keep the original shape847 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)848 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)849 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)850 851 kv_seq_len = position_ids.max().item() + 1852 cos, sin = self.rotary_emb(value_states.to(torch.float32), seq_len=kv_seq_len)853 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)854 855 if past_key_value is not None:856 cache_kwargs = {'sin': sin, 'cos': cos} # Specific to RoPE models857 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)858 859 # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache860 # to be able to avoid many of these transpose/reshape/view.861 query_states = query_states.transpose(1, 2)862 key_states = key_states.transpose(1, 2)863 value_states = value_states.transpose(1, 2)864 865 dropout_rate = self.attention_dropout if self.training else 0.0866 867 # In PEFT, usually we cast the layer norms in float32 for training stability reasons868 # therefore the input hidden states gets silently casted in float32. Hence, we need869 # cast them back in the correct dtype just to be sure everything works as expected.870 # This might slowdown training & inference so it is recommended to not cast the LayerNorms871 # in fp32. (MiniCPMRMSNorm handles it correctly)872 873 input_dtype = query_states.dtype874 if input_dtype == torch.float32:875 # Handle the case where the model is quantized876 if hasattr(self.config, '_pre_quantization_dtype'):877 target_dtype = self.config._pre_quantization_dtype878 else:879 target_dtype = self.q_proj.weight.dtype880 881 logger.warning_once(882 f'The input hidden states seems to be silently casted in float32, this might be related to'883 f' the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in'884 f' {target_dtype}.'885 )886 887 query_states = query_states.to(target_dtype)888 key_states = key_states.to(target_dtype)889 value_states = value_states.to(target_dtype)890 891 attn_output = self._flash_attention_forward(892 query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate893 )894 895 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()896 attn_output = self.o_proj(attn_output)897 898 if not output_attentions:899 attn_weights = None900 901 return attn_output, attn_weights, past_key_value902 903 def _flash_attention_forward(904 self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None905 ):906 """907 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token908 first unpad the input, then computes the attention scores and pad the final attention scores.909 910 Args:911 query_states (`torch.Tensor`):912 Input query states to be passed to Flash Attention API913 key_states (`torch.Tensor`):914 Input key states to be passed to Flash Attention API915 value_states (`torch.Tensor`):916 Input value states to be passed to Flash Attention API917 attention_mask (`torch.Tensor`):918 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the919 position of padding tokens and 1 for the position of non-padding tokens.920 dropout (`int`, *optional*):921 Attention dropout922 softmax_scale (`float`, *optional*):923 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)924 """925 if not self._flash_attn_uses_top_left_mask:926 causal = self.is_causal927 else:928 # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in MiniCPMFlashAttention2 __init__.929 causal = self.is_causal and query_length != 1930 # Contains at least one padding token in the sequence931 if attention_mask is not None:932 batch_size = query_states.shape[0]933 query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(934 query_states, key_states, value_states, attention_mask, query_length935 )936 937 cu_seqlens_q, cu_seqlens_k = cu_seq_lens938 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens939 attn_output_unpad = flash_attn_varlen_func(940 query_states,941 key_states,942 value_states,943 cu_seqlens_q=cu_seqlens_q,944 cu_seqlens_k=cu_seqlens_k,945 max_seqlen_q=max_seqlen_in_batch_q,946 max_seqlen_k=max_seqlen_in_batch_k,947 dropout_p=dropout,948 softmax_scale=softmax_scale,949 causal=causal,950 )951 952 attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)953 else:954 attn_output = flash_attn_func(955 query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal956 )957 958 return attn_output959 960 def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):961 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)962 batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape963 964 key_layer = index_first_axis(965 key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k966 )967 value_layer = index_first_axis(968 value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k969 )970 if query_length == kv_seq_len:971 query_layer = index_first_axis(972 query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k973 )974 cu_seqlens_q = cu_seqlens_k975 max_seqlen_in_batch_q = max_seqlen_in_batch_k976 indices_q = indices_k977 elif query_length == 1:978 max_seqlen_in_batch_q = 1979 cu_seqlens_q = torch.arange(980 batch_size + 1, dtype=torch.int32, device=query_layer.device981 ) # There is a memcpy here, that is very bad.982 indices_q = cu_seqlens_q[:-1]983 query_layer = query_layer.squeeze(1)984 else:985 # The -q_len: slice assumes left padding.986 attention_mask = attention_mask[:, -query_length:]987 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)988 989 return (990 query_layer,991 key_layer,992 value_layer,993 indices_q,994 (cu_seqlens_q, cu_seqlens_k),995 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),996 )997 998 999class MiniCPMInfLLMv2Attention(MiniCPMAttention):1000 """1001 MiniCPM flash attention module. This module inherits from `MiniCPMAttention` as the weights of the module stays1002 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of1003 flash attention and deal with padding tokens in case the input contains any of them.1004 """1005 1006 def __init__(self, *args, **kwargs):1007 super().__init__(*args, **kwargs)1008 assert self.config._attn_implementation == 'flash_attention_2', 'Only flash_attention_2 is supported for sparse attention'1009 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.1010 # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignment, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.1011 # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).1012 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()1013 1014 # -------sparse-------1015 self.kernel_size = self.config.sparse_config.get('kernel_size', 32)1016 self.kernel_stride = self.config.sparse_config.get('kernel_stride', 16)1017 self.init_blocks = self.config.sparse_config.get('init_blocks', 1)1018 self.block_size = self.config.sparse_config.get('block_size', 64)1019 self.window_size = self.config.sparse_config.get('window_size', 2048)1020 self.dense_len = self.config.sparse_config.get('dense_len', 8192)1021 1022 self.local_blocks = self.window_size // self.block_size # local_blocks1023 self.topk = self.config.sparse_config.get('topk', 64) + (self.window_size//self.block_size)1024 self.use_nope = self.config.sparse_config.get('use_nope', False)1025 1026 self.compress_k = CompressK(self.num_key_value_heads, self.head_dim, kernel_size=self.kernel_size, kernel_stride=self.kernel_stride)1027 self.compress_k2 = CompressK(self.num_key_value_heads, self.head_dim, kernel_size=self.kernel_size*4, kernel_stride=self.kernel_stride*4)1028 1029 def forward(1030 self,1031 hidden_states: torch.Tensor,1032 attention_mask: Optional[torch.LongTensor] = None,1033 position_ids: Optional[torch.LongTensor] = None,1034 past_key_value: Optional[Cache] = None,1035 output_attentions: bool = False,1036 use_cache: bool = False,1037 **kwargs,1038 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:1039 # MiniCPMFlashAttention2 attention does not support output_attentions1040 if 'padding_mask' in kwargs:1041 warnings.warn(1042 'Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`'1043 )1044 1045 # overwrite attention_mask with padding_mask1046 attention_mask = kwargs.pop('padding_mask')1047 1048 output_attentions = False1049 1050 bsz, q_len, _ = hidden_states.size()1051 1052 1053 query_states = self.q_proj(hidden_states)1054 key_states = self.k_proj(hidden_states)1055 value_states = self.v_proj(hidden_states)1056 1057 # !save no rope1058 if self.use_nope:1059 query_states_no_rope = query_states.view(bsz, q_len, self.num_heads, self.head_dim)1060 key_states_no_rope = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim)1061 1062 # Flash attention requires the input to have the shape1063 # batch_size x seq_length x head_dim x hidden_dim1064 # therefore we just need to keep the original shape1065 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)1066 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)1067 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)1068 1069 kv_seq_len = position_ids.max().item() + 11070 cos, sin = self.rotary_emb(value_states.to(torch.float32), seq_len=kv_seq_len)1071 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)1072 1073 if past_key_value is not None:1074 cache_kwargs = {'sin': sin, 'cos': cos} # Specific to RoPE models1075 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)1076 1077 # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache1078 # to be able to avoid many of these transpose/reshape/view.1079 query_states = query_states.transpose(1, 2)1080 key_states = key_states.transpose(1, 2)1081 value_states = value_states.transpose(1, 2)1082 if self.use_nope:1083 key_states_no_rope =past_key_value.update_no_rope_key(key_states_no_rope, self.layer_idx)1084 no_rope_param = {1085 'key_states_no_rope': key_states_no_rope,1086 'query_states_no_rope': query_states_no_rope,1087 }1088 1089 else:1090 no_rope_param = None1091 1092 dropout_rate = self.attention_dropout if self.training else 0.01093 1094 # In PEFT, usually we cast the layer norms in float32 for training stability reasons1095 # therefore the input hidden states gets silently casted in float32. Hence, we need1096 # cast them back in the correct dtype just to be sure everything works as expected.1097 # This might slowdown training & inference so it is recommended to not cast the LayerNorms1098 # in fp32. (MiniCPMRMSNorm handles it correctly)1099 1100 input_dtype = query_states.dtype1101 if input_dtype == torch.float32:1102 # Handle the case where the model is quantized1103 if hasattr(self.config, '_pre_quantization_dtype'):1104 target_dtype = self.config._pre_quantization_dtype1105 else:1106 target_dtype = self.q_proj.weight.dtype1107 1108 logger.warning_once(1109 f'The input hidden states seems to be silently casted in float32, this might be related to'1110 f' the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in'1111 f' {target_dtype}.'1112 )1113 1114 query_states = query_states.to(target_dtype)1115 key_states = key_states.to(target_dtype)1116 value_states = value_states.to(target_dtype)1117 if kv_seq_len < self.dense_len:1118 attn_output = self._flash_attention_forward_dense(1119 query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate)1120 else:1121 attn_output = self._sparse_attention_forward(1122 query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate,1123 no_rope_param=no_rope_param, # if past_key_value is not None else None,1124 past_key_value=past_key_value)1125 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()1126 attn_output = self.o_proj(attn_output)1127 1128 if not output_attentions:1129 attn_weights = None1130 1131 return attn_output, attn_weights, past_key_value1132 1133 def _sparse_attention_forward(1134 self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None, no_rope_param=None, past_key_value=None1135 ):1136 """1137 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token1138 first unpad the input, then computes the attention scores and pad the final attention scores.1139 1140 Args:1141 query_states (`torch.Tensor`):1142 Input query states to be passed to Flash Attention API1143 key_states (`torch.Tensor`):1144 Input key states to be passed to Flash Attention API1145 value_states (`torch.Tensor`):1146 Input value states to be passed to Flash Attention API1147 attention_mask (`torch.Tensor`):1148 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the1149 position of padding tokens and 1 for the position of non-padding tokens.1150 dropout (`int`, *optional*):1151 Attention dropout1152 softmax_scale (`float`, *optional*):1153 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)1154 """1155 if not self._flash_attn_uses_top_left_mask:1156 causal = self.is_causal1157 else:1158 # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in MiniCPMFlashAttention2 __init__.1159 causal = self.is_causal and query_length != 11160 # Contains at least one padding token in the sequence1161 if attention_mask is not None:1162 batch_size = query_states.shape[0]1163 # assert batch_size == 1, 'Only batch_size=1 is supported at the moment.'1164 if past_key_value!=None:1165 compressed_k, compressed_cu_seqlens, compressed_k2, compressed_cu_seqlens2 = self.get_compress_k(1166 key_states=key_states if self.use_nope ==False else no_rope_param['key_states_no_rope'], # This can be optimized a bit;1167 attention_mask=attention_mask,1168 past_key_value=past_key_value,1169 1170 )1171 1172 query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(1173 query_states, key_states, value_states, attention_mask, query_length1174 )1175 1176 cu_seqlens_q, cu_seqlens_k = cu_seq_lens1177 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens1178 if no_rope_param != None:1179 if max_seqlen_in_batch_q == 1:1180 no_rope_param['query_states_no_rope'] = no_rope_param['query_states_no_rope'].squeeze(1)1181 else:1182 no_rope_param['query_states_no_rope'],_, _, _ = _unpad_one_tensor(no_rope_param['query_states_no_rope'],attention_mask=attention_mask)1183 if past_key_value==None:1184 # compress_k use varlen form1185 compressed_k, compressed_cu_seqlens = self.compress_k(key_states,cu_seqlens_k)1186 compressed_k2, compressed_cu_seqlens2 = self.compress_k2(key_states,cu_seqlens_k)1187 else:1188 # compressed_k and compressed_k2 already retrieved from get_compress_k above1189 pass1190 1191 1192 attn_output_unpad = self.sparse_forward(1193 query_states,1194 key_states,1195 value_states,1196 cu_seqlens_q,1197 cu_seqlens_k,1198 max_seqlen_in_batch_q,1199 max_seqlen_in_batch_k,1200 no_rope_param=no_rope_param,