Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2024 IBM and the HuggingFace Inc. team. All rights reserved.3#4# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX5# and OPT implementations in this library. It has been modified from its6# original forms to accommodate minor architectural differences compared7# to GPT-NeoX and OPT used by the Meta AI team that trained the model.8#9# Licensed under the Apache License, Version 2.0 (the "License");10# you may not use this file except in compliance with the License.11# You may obtain a copy of the License at12#13# http://www.apache.org/licenses/LICENSE-2.014#15# Unless required by applicable law or agreed to in writing, software16# distributed under the License is distributed on an "AS IS" BASIS,17# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.18# See the License for the specific language governing permissions and19# limitations under the License.20"""PyTorch Bamba model."""21 22from typing import Optional, TypedDict, Union23 24import torch25from torch import nn26 27from transformers.activations import ACT2FN28from transformers.models.jamba.modeling_jamba import HybridMambaAttentionDynamicCache, JambaAttentionDecoderLayer29from transformers.models.llama.modeling_llama import (30 LlamaAttention,31 LlamaForCausalLM,32 LlamaMLP,33 LlamaRMSNorm,34 LlamaRotaryEmbedding,35 rotate_half,36)37from transformers.models.mamba2.modeling_mamba2 import (38 MambaRMSNormGated,39 pad_tensor_by_size,40 reshape_into_chunks,41 segment_sum,42)43 44from ...modeling_attn_mask_utils import AttentionMaskConverter45from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast46from ...modeling_utils import PreTrainedModel47from ...processing_utils import Unpack48from ...utils import (49 auto_docstring,50 can_return_tuple,51 logging,52)53from ...utils.deprecation import deprecate_kwarg54from ...utils.import_utils import is_causal_conv1d_available, is_mamba_2_ssm_available55from .configuration_bamba import BambaConfig56 57 58if is_mamba_2_ssm_available():59 from mamba_ssm.ops.triton.selective_state_update import selective_state_update60 from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined61else:62 selective_state_update = None63 64if is_causal_conv1d_available():65 from causal_conv1d import causal_conv1d_fn, causal_conv1d_update66else:67 causal_conv1d_update, causal_conv1d_fn = None, None68 69is_fast_path_available = all((selective_state_update, causal_conv1d_fn, causal_conv1d_update))70 71 72logger = logging.get_logger(__name__)73 74 75class BambaFlashAttentionKwargs(TypedDict, total=False):76 """77 Keyword arguments for advanced Flash Attention, causal-conv1d, and mamba_ssm kernel usage.78 Use cases include padding-free training and fewer `torch.compile` graph breaks.79 80 Attributes:81 cu_seq_lens_q (`torch.LongTensor`)82 Gets cumulative sequence length for query state.83 cu_seq_lens_k (`torch.LongTensor`)84 Gets cumulative sequence length for key state.85 max_length_q (`int`):86 Maximum sequence length for query state.87 max_length_k (`int`):88 Maximum sequence length for key state.89 seq_idx (`torch.IntTensor):90 Index of each packed sequence.91 """92 93 cu_seq_lens_q: torch.LongTensor94 cu_seq_lens_k: torch.LongTensor95 max_length_q: int96 max_length_k: int97 seq_idx: torch.IntTensor98 99 100# Adapted from transformers.models.jamba.modeling_jamba.HybridMambaAttentionDynamicCache for the v2 mixer101class HybridMambaAttentionDynamicCache(HybridMambaAttentionDynamicCache):102 """103 A dynamic cache that can handle both the attention cache (which has a seq_len dimension) and the mamba cache104 (which has a constant shape regardless of seq_len).105 106 This cache has two sets of lists of tensors: `key_cache` and `value_cache` for attention cache and `conv_states`107 and `ssm_states` for mamba cache. Each of these lists has `num_layers` tensors. The expected shape for each tensor108 For attention layers, `key_cache` and `value_cache` have a shape of `(batch_size, num_heads, seq_len, head_dim)`,109 while `conv_states` and `ssm_states` have a shape of `(batch_size, 0)` (empty tensors).110 For mamba layers, `key_cache` and `value_cache` have a shape of `(batch_size, 0)` (empty tensors),111 while `conv_states` represents the convolution state and has a shape of `(batch_size, d_inner, d_conv)`,112 and `ssm_states` represents the ssm state and has a shape of `(batch_size, d_inner, d_state)`.113 """114 115 def __init__(self, config: BambaConfig, batch_size, dtype=torch.float16, device=None):116 self.layers_block_type = config.layers_block_type117 self.has_previous_state = False # only used by mamba118 conv_kernel_size = config.mamba_d_conv119 ssm_state_size = config.mamba_d_state120 121 self.conv_states = []122 self.ssm_states = []123 self.transformer_layers = []124 for i in range(config.num_hidden_layers):125 if self.layers_block_type[i] == "mamba":126 self.conv_states += [127 torch.zeros(128 batch_size,129 (config.mamba_expand * config.hidden_size + 2 * config.mamba_n_groups * ssm_state_size),130 conv_kernel_size,131 device=device,132 dtype=dtype,133 )134 ]135 self.ssm_states += [136 torch.zeros(137 batch_size,138 config.mamba_n_heads,139 config.mamba_d_head,140 ssm_state_size,141 device=device,142 dtype=dtype,143 )144 ]145 else:146 self.conv_states += [torch.tensor([[]] * batch_size, device=device)]147 self.ssm_states += [torch.tensor([[]] * batch_size, device=device)]148 self.transformer_layers.append(i)149 150 self.key_cache = [torch.tensor([[]] * batch_size, device=device) for _ in range(config.num_hidden_layers)]151 self.value_cache = [torch.tensor([[]] * batch_size, device=device) for _ in range(config.num_hidden_layers)]152 153 154class BambaRotaryEmbedding(LlamaRotaryEmbedding):155 pass156 157 158# Adapted from transformers.models.glm.modular_glm.apply_rotary_pos_emb159def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):160 """Applies Rotary Position Embedding to the query and key tensors.161 162 Removes the interleaving of cos and sin from GLM163 164 Args:165 q (`torch.Tensor`): The query tensor.166 k (`torch.Tensor`): The key tensor.167 cos (`torch.Tensor`): The cosine part of the rotary embedding.168 sin (`torch.Tensor`): The sine part of the rotary embedding.169 position_ids (`torch.Tensor`, *optional*):170 Deprecated and unused.171 unsqueeze_dim (`int`, *optional*, defaults to 1):172 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and173 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note174 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and175 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes176 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have177 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.178 Returns:179 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.180 """181 cos = cos.unsqueeze(unsqueeze_dim)182 sin = sin.unsqueeze(unsqueeze_dim)183 184 # Keep half or full tensor for later concatenation185 rotary_dim = cos.shape[-1]186 q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]187 k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]188 189 # Apply rotary embeddings on the first half or full tensor190 q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)191 k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)192 193 # Concatenate back to full shape194 q_embed = torch.cat([q_embed, q_pass], dim=-1)195 k_embed = torch.cat([k_embed, k_pass], dim=-1)196 return q_embed, k_embed197 198 199class BambaAttention(LlamaAttention):200 pass201 202 203class BambaRMSNormGated(MambaRMSNormGated):204 pass205 206 207def apply_mask_to_padding_states(hidden_states, attention_mask):208 """209 Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66210 """211 if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:212 dtype = hidden_states.dtype213 hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)214 215 return hidden_states216 217 218# Adapted from transformers.models.mamba2.modeling_mamba2.Mamba2Mixer219class BambaMixer(nn.Module):220 """221 Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.222 A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective)223 ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4,224 and is why Mamba is called **selective** state spaces)225 226 The are a few differences between this and Mamba2Mixer:227 - The variable use_precomputed_states is slightly different due to the hybrid cache structure228 - There's a few non-obvious bugs fixed with batching in the slow path that exist in main229 - Some extra variables that our layer doesn't need have been removed230 - We ported most of the refactors in https://github.com/huggingface/transformers/pull/35154, which is (as of Dec 18, 2024) unmerged231 """232 233 def __init__(self, config: BambaConfig, layer_idx: int):234 super().__init__()235 self.num_heads = config.mamba_n_heads236 self.hidden_size = config.hidden_size237 self.ssm_state_size = config.mamba_d_state238 self.conv_kernel_size = config.mamba_d_conv239 self.intermediate_size = int(config.mamba_expand * self.hidden_size)240 self.layer_idx = layer_idx241 self.use_conv_bias = config.mamba_conv_bias242 self.activation = config.hidden_act243 self.act = ACT2FN[config.hidden_act]244 self.use_bias = config.mamba_proj_bias245 246 self.layer_norm_epsilon = config.rms_norm_eps247 248 self.n_groups = config.mamba_n_groups249 self.head_dim = config.mamba_d_head250 self.chunk_size = config.mamba_chunk_size251 252 # FIXME:253 self.time_step_limit = (0.0, float("inf"))254 self.time_step_min = 0.001255 self.time_step_max = 0.1256 257 self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size258 self.conv1d = nn.Conv1d(259 in_channels=self.conv_dim,260 out_channels=self.conv_dim,261 bias=config.mamba_conv_bias,262 kernel_size=self.conv_kernel_size,263 groups=self.conv_dim,264 padding=self.conv_kernel_size - 1,265 )266 267 # projection of the input hidden states268 projection_size = self.intermediate_size + self.conv_dim + self.num_heads269 self.in_proj = nn.Linear(270 self.hidden_size,271 projection_size,272 bias=self.use_bias,273 )274 # selective projection used to make dt, B and C input dependent275 276 # time step projection (discretization)277 # instantiate once and copy inv_dt in init_weights of PretrainedModel278 self.dt_bias = nn.Parameter(torch.ones(self.num_heads))279 280 # S4D real initialization. These are not discretized!281 # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded282 A = torch.arange(1, self.num_heads + 1)283 self.A_log = nn.Parameter(torch.log(A))284 self.norm = BambaRMSNormGated(self.intermediate_size, eps=self.layer_norm_epsilon)285 self.D = nn.Parameter(torch.ones(self.num_heads))286 287 self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=self.use_bias)288 289 if not is_fast_path_available:290 logger.warning_once(291 "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"292 " is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and"293 " https://github.com/Dao-AILab/causal-conv1d"294 )295 else:296 logger.warning_once("The fast path for Bamba will be used when running the model on a GPU")297 298 def cuda_kernels_forward(299 self,300 hidden_states: torch.Tensor,301 cache_params: Optional[HybridMambaAttentionDynamicCache] = None,302 cache_position: Optional[torch.LongTensor] = None,303 attention_mask: Optional[torch.Tensor] = None,304 seq_idx: Optional[torch.IntTensor] = None,305 ):306 # 1. Gated MLP's linear projection307 hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)308 projected_states = self.in_proj(hidden_states)309 310 # Set up dimensions for reshapes later311 batch_size, seq_len, _ = hidden_states.shape312 groups_time_state_size = self.n_groups * self.ssm_state_size313 314 use_precomputed_states = (315 cache_params is not None316 and cache_params.has_previous_state317 and seq_len == 1318 and cache_params.conv_states[self.layer_idx].shape[0]319 == cache_params.ssm_states[self.layer_idx].shape[0]320 == batch_size321 and cache_position is not None322 and cache_position[0] > 0323 )324 325 # getting projected states from cache if it exists326 if use_precomputed_states:327 gate, hidden_states_B_C, dt = projected_states.squeeze(1).split(328 [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1329 )330 331 # 2. Convolution sequence transformation332 hidden_states_B_C = causal_conv1d_update(333 hidden_states_B_C,334 cache_params.conv_states[self.layer_idx],335 self.conv1d.weight.squeeze(1),336 self.conv1d.bias,337 self.activation,338 )339 340 hidden_states, B, C = torch.split(341 hidden_states_B_C,342 [self.intermediate_size, groups_time_state_size, groups_time_state_size],343 dim=-1,344 )345 346 # 3. SSM transformation347 A = -torch.exp(self.A_log.float()) # (nheads,)348 A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)349 dt = dt[:, :, None].expand(-1, -1, self.head_dim)350 dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim)351 D = self.D[:, None, ...].expand(-1, self.head_dim)352 B = B.view(batch_size, self.n_groups, B.shape[1] // self.n_groups)353 C = C.view(batch_size, self.n_groups, C.shape[1] // self.n_groups)354 hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim)355 hidden_states = selective_state_update(356 cache_params.ssm_states[self.layer_idx],357 hidden_states_reshaped,358 dt,359 A,360 B,361 C,362 D,363 z=None,364 dt_bias=dt_bias,365 dt_softplus=True,366 )367 hidden_states = hidden_states.view(batch_size, self.num_heads * self.head_dim)368 hidden_states = self.norm(hidden_states, gate)369 370 # 4. Final linear projection371 out = self.out_proj(hidden_states)[:, None, ...]372 # Fused calculations or step by step if no initialized cache is found373 else:374 A = -torch.exp(self.A_log.float()) # (num_heads) or (intermediate_size, state_size)375 dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit}376 377 # 2-4. Fused kernel for conv1d, SSM, and the final projection378 if self.training and cache_params is None:379 out = mamba_split_conv1d_scan_combined(380 projected_states,381 self.conv1d.weight.squeeze(1),382 self.conv1d.bias,383 self.dt_bias,384 A,385 D=self.D,386 chunk_size=self.chunk_size,387 seq_idx=seq_idx,388 activation=self.activation,389 rmsnorm_weight=self.norm.weight,390 rmsnorm_eps=self.norm.variance_epsilon,391 outproj_weight=self.out_proj.weight,392 outproj_bias=self.out_proj.bias,393 headdim=self.head_dim,394 ngroups=self.n_groups,395 norm_before_gate=False,396 return_final_states=False,397 **dt_limit_kwargs,398 )399 400 else:401 gate, hidden_states_B_C, dt = projected_states.split(402 [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1403 )404 405 # 2. Convolution sequence transformation406 # Init cache407 if cache_params is not None:408 # storing the states409 # If we just take xBC[:, :, -self.d_conv :], it will error if seqlen < self.d_conv410 # Instead F.pad will pad with zeros if seqlen < self.d_conv, and truncate otherwise.411 hidden_states_B_C_transposed = hidden_states_B_C.transpose(1, 2)412 conv_states = nn.functional.pad(413 hidden_states_B_C_transposed,414 (self.conv_kernel_size - hidden_states_B_C_transposed.shape[-1], 0),415 )416 cache_params.conv_states[self.layer_idx].copy_(conv_states)417 418 if self.activation not in ["silu", "swish"]:419 hidden_states_B_C = self.act(420 self.conv1d(hidden_states_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2)421 )422 else:423 hidden_states_B_C = causal_conv1d_fn(424 x=hidden_states_B_C.transpose(1, 2),425 weight=self.conv1d.weight.squeeze(1),426 bias=self.conv1d.bias,427 activation=self.activation,428 seq_idx=seq_idx,429 ).transpose(1, 2)430 431 hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)432 hidden_states, B, C = torch.split(433 hidden_states_B_C,434 [self.intermediate_size, groups_time_state_size, groups_time_state_size],435 dim=-1,436 )437 438 # 3. SSM transformation439 scan_output, ssm_state = mamba_chunk_scan_combined(440 hidden_states.view(batch_size, seq_len, -1, self.head_dim),441 dt,442 A,443 B.view(batch_size, seq_len, self.n_groups, -1),444 C.view(batch_size, seq_len, self.n_groups, -1),445 chunk_size=self.chunk_size,446 D=self.D,447 z=None,448 seq_idx=seq_idx,449 return_final_states=True,450 dt_bias=self.dt_bias,451 dt_softplus=True,452 **dt_limit_kwargs,453 )454 455 # Init cache456 if ssm_state is not None and cache_params is not None:457 cache_params.ssm_states[self.layer_idx].copy_(ssm_state)458 459 scan_output = scan_output.view(batch_size, seq_len, -1)460 # Multiply "gate" branch and apply extra normalization layer461 scan_output = self.norm(scan_output, gate)462 463 # 4. Final linear projection464 out = self.out_proj(scan_output)465 return out466 467 # fmt: off468 def torch_forward(469 self,470 input_states,471 cache_params: Optional[HybridMambaAttentionDynamicCache] = None,472 cache_position: Optional[torch.LongTensor] = None,473 attention_mask: Optional[torch.Tensor] = None,474 ):475 batch_size, seq_len, _ = input_states.shape476 dtype = input_states.dtype477 478 # 1. Gated MLP's linear projection479 input_states = apply_mask_to_padding_states(input_states, attention_mask)480 projected_states = self.in_proj(input_states)481 gate, hidden_states_B_C, dt = projected_states.split(482 [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1483 )484 485 use_precomputed_states = (486 cache_params is not None487 and cache_params.has_previous_state488 and seq_len == 1489 and cache_params.conv_states[self.layer_idx].shape[0]490 == cache_params.ssm_states[self.layer_idx].shape[0]491 == batch_size492 and cache_position is not None493 and cache_position[0] > 0494 )495 496 # 2. Convolution sequence transformation497 if use_precomputed_states:498 cache_params.conv_states[self.layer_idx] = cache_params.conv_states[self.layer_idx].roll(shifts=-1, dims=-1)499 cache_params.conv_states[self.layer_idx][:, :, -1] = hidden_states_B_C[:, 0, :].to(cache_params.conv_states[self.layer_idx].device)500 501 # We need to guarantee that anything regarding the cache is on the same device502 conv_states = cache_params.conv_states[self.layer_idx].to(device=self.conv1d.weight.device)503 504 hidden_states_B_C = torch.sum(505 conv_states * self.conv1d.weight.squeeze(1), dim=-1506 )507 if self.use_conv_bias:508 hidden_states_B_C = hidden_states_B_C + self.conv1d.bias509 hidden_states_B_C = self.act(hidden_states_B_C)510 else:511 # Init cache512 if cache_params is not None:513 hidden_states_B_C_transposed = hidden_states_B_C.transpose(1, 2)514 conv_states = nn.functional.pad(515 hidden_states_B_C_transposed, (self.conv_kernel_size - hidden_states_B_C_transposed.shape[-1], 0)516 )517 cache_params.conv_states[self.layer_idx].copy_(conv_states)518 519 hidden_states_B_C = self.act(self.conv1d(hidden_states_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2))520 521 hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)522 hidden_states, B, C = torch.split(523 hidden_states_B_C,524 [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size],525 dim=-1526 )527 528 # 3. SSM transformation529 A = -torch.exp(self.A_log.float()) # [num_heads]530 if use_precomputed_states:531 # We need to guarantee that anything regarding the cache is on the same device532 cache_device = cache_params.ssm_states[self.layer_idx].device533 534 # Note: there is no need to pad parameter matrices here, as there is just one new token535 # for batched generation536 dt = dt[:, 0, :][:, None, ...]537 dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim)538 # [num_heads] -> [num_heads, head_dim]539 dt_bias = self.dt_bias[..., None].expand(self.dt_bias.shape[0], self.head_dim)540 541 dt = torch.nn.functional.softplus(dt + dt_bias.to(dt.dtype))542 dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1])543 A = A[..., None, None].expand(self.num_heads, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)544 # [bsz, num_heads, head_dim, state_size]545 dA = (torch.exp(dt[..., None] * A)).to(device=cache_device)546 547 # Discretize B548 # [bsz, n_groups * state_size] -> [bsz, n_groups, 1, state_size] ->549 # -> [bsz, n_groups, group to head repetition factor, state_size] -> [bsz, num_heads, state_size]550 B = B.reshape(batch_size, self.n_groups, -1)[..., None, :]551 B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous()552 B = B.reshape(batch_size, -1, B.shape[-1])553 # [bsz, num_heads, head_dim, state_size]554 dB = dt[..., None] * B[..., None, :]555 556 # Discretize x into dB557 # [bsz, intermediate_size] -> [bsz, num_heads, head_dim]558 hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim)559 dBx = (dB * hidden_states[..., None]).to(device=cache_device)560 561 # State calculation562 cache_params.ssm_states[self.layer_idx].copy_(563 cache_params.ssm_states[self.layer_idx] * dA + dBx564 )565 566 # Subsequent output567 # [bsz, n_groups * state_size] -> [bsz, num_heads, state_size]568 C = C.reshape(batch_size, self.n_groups, -1)[..., None, :]569 C = C.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, C.shape[-1]).contiguous()570 C = C.reshape(batch_size, -1, C.shape[-1])571 # [bsz, num_heads, head_dim]572 573 ssm_states = cache_params.ssm_states[self.layer_idx].to(device=C.device, dtype=C.dtype) # Shape: [b, h, d, n]574 # Reshape ssm_states to merge the first two dimensions575 ssm_states_reshaped = ssm_states.view(batch_size * self.num_heads, self.head_dim, self.ssm_state_size) # Shape: [b*h, d, n]576 C_reshaped = C.view(batch_size * self.num_heads, self.ssm_state_size, 1) # Shape: [b*h, n, 1]577 y = torch.bmm(ssm_states_reshaped, C_reshaped)578 y = y.view(batch_size, self.num_heads, self.head_dim)579 580 # D skip connection581 # [num_heads] -> [num_heads, head_dim]582 D = self.D[..., None].expand(self.D.shape[0], self.head_dim)583 y = (y + hidden_states * D).to(y.dtype)584 585 # [bsz, num_heads, head_dim] -> [bsz, 1, intermediate_size]586 y = y.reshape(batch_size, -1)[:, None, ...]587 else:588 # begin ssd naive implementation without einsums589 dt = nn.functional.softplus(dt + self.dt_bias)590 dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1])591 hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float()592 B = B.reshape(batch_size, seq_len, -1, self.ssm_state_size).float()593 C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float()594 B = B.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)595 C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)596 pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size597 598 D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size)599 600 # Discretize x and A601 hidden_states = hidden_states * dt[..., None]602 A = A.to(hidden_states.dtype) * dt603 604 # Rearrange into blocks/chunks605 hidden_states, A, B, C = [reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states, A, B, C)]606 607 # [bsz, -1, chunk_size, num_heads] -> [bsz, num_heads, -1, chunk_size]608 A = A.permute(0, 3, 1, 2)609 A_cumsum = torch.cumsum(A, dim=-1)610 611 # 1. Compute the output for each intra-chunk (diagonal blocks)612 # This is the analog of a causal mask613 L = torch.exp(segment_sum(A))614 615 # Contraction of C and B to get G (attention-weights like)616 G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :] # shape: (b, c, l, s, h, n)617 G = G_intermediate.sum(dim=-1) # shape: (b, c, l, s, h)618 619 # Compute M, equivalent to applying attention mask to weights620 M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]621 M = M_intermediate.sum(dim=-1)622 623 # Compute Y_diag (apply to values)624 Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3)625 626 # 2. Compute the state for each intra-chunk627 # (right term of low-rank factorization of off-diagonal blocks; B terms)628 decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum)629 B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None]630 states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2)631 632 # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries633 # (middle term of factorization of off-diag blocks; A terms)634 if use_precomputed_states:635 previous_states = cache_params.ssm_states[self.layer_idx][:, None, ...].to(device=states.device)636 else:637 previous_states = torch.zeros_like(states[:, :1])638 states = torch.cat([previous_states, states], dim=1)639 decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0))))640 decay_chunk = decay_chunk.transpose(1, 3)641 new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)642 states, ssm_state = new_states[:, :-1], new_states[:, -1]643 644 # 4. Compute state -> output conversion per chunk645 # (left term of low-rank factorization of off-diagonal blocks; C terms)646 state_decay_out = torch.exp(A_cumsum)647 C_times_states = (C[..., None, :] * states[:, :, None, ...])648 state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1)649 Y_off = (C_times_states.sum(-1) * state_decay_out_permuted[..., None])650 651 # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)652 y = Y_diag + Y_off653 # [bsz, -1, self.chunk_size, num_heads, head_dim] -> [bsz, (padded) seq_len, num_heads, head_dim]654 y = y.reshape(batch_size, -1, self.num_heads, self.head_dim)655 656 y = y + D_residual657 # Cutting off padded chunks658 if pad_size > 0:659 y = y[:, :seq_len, :, :]660 y = y.reshape(batch_size, seq_len, -1)661 662 # Init cache663 if ssm_state is not None and cache_params is not None:664 cache_params.ssm_states[self.layer_idx].copy_(ssm_state)665 666 scan_output = self.norm(y, gate)667 668 # end ssd naive669 670 # 4. Final linear projection671 contextualized_states = self.out_proj(scan_output.to(dtype)) # [batch, seq_len, hidden_size]672 return contextualized_states673 # fmt: on674 675 def forward(676 self,677 hidden_states,678 cache_params: Optional[HybridMambaAttentionDynamicCache] = None,679 cache_position: Optional[torch.LongTensor] = None,680 attention_mask: Optional[torch.Tensor] = None,681 seq_idx: Optional[torch.IntTensor] = None,682 **kwargs,683 ):684 if is_fast_path_available and "cuda" in self.in_proj.weight.device.type:685 return self.cuda_kernels_forward(hidden_states, cache_params, cache_position, attention_mask, seq_idx)686 if seq_idx is not None:687 raise NotImplementedError(688 "`seq_idx` support requires fast path support. Please install `mamba_ssm` and `causal_conv1d`"689 )690 dtype = hidden_states.dtype691 if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:692 # tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66693 hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)694 695 return self.torch_forward(hidden_states, cache_params, cache_position, attention_mask)696 697 698class BambaMLP(LlamaMLP):699 pass700 701 702class BambaRMSNorm(LlamaRMSNorm):703 pass704 705 706class BambaDecoderLayer(JambaAttentionDecoderLayer):707 def __init__(self, config: BambaConfig, layer_idx: int, layer_type: str = "mamba"):708 super().__init__(config, layer_idx)709 710 del self.self_attn711 712 num_experts = 1713 ffn_layer_class = BambaMLP if num_experts == 1 else None714 self.feed_forward = ffn_layer_class(config)715 716 self.layer_type = layer_type717 if layer_type == "mamba":718 self.mamba = BambaMixer(config=config, layer_idx=layer_idx)719 elif layer_type == "attention":720 self.self_attn = BambaAttention(config, layer_idx)721 else:722 raise ValueError("Invalid layer_type")723 724 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")725 def forward(726 self,727 hidden_states: torch.Tensor,728 attention_mask: Optional[torch.Tensor] = None,729 position_ids: Optional[torch.LongTensor] = None,730 past_key_values: Optional[HybridMambaAttentionDynamicCache] = None,731 output_attentions: Optional[bool] = False,732 use_cache: Optional[bool] = False,733 cache_position: Optional[torch.LongTensor] = None,734 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC735 **kwargs: Unpack[BambaFlashAttentionKwargs],736 ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:737 """738 Args:739 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`740 attention_mask (`torch.FloatTensor`, *optional*): attention mask of size741 `(batch, sequence_length)` where padding elements are indicated by 0.742 past_key_values (`HybridMambaAttentionDynamicCache`, *optional*): cached past key and value projection states743 output_attentions (`bool`, *optional*):744 Whether or not to return the attentions tensors of all attention layers. See `attentions` under745 returned tensors for more detail.746 use_cache (`bool`, *optional*):747 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding748 (see `past_key_values`).749 cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):750 Indices depicting the position of the input sequence tokens in the sequence.751 position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):752 Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,753 with `head_dim` being the embedding dimension of each attention head.754 kwargs (`dict`, *optional*):755 Arbitrary kwargs. Can be used to provide `BambaFlashAttentionKwargs` for756 padding-free training and/or improve torch.compile performance.757 """758 759 residual = hidden_states760 761 hidden_states = self.input_layernorm(hidden_states)762 763 # this is a hybrid decoder layer764 if self.layer_type == "mamba":765 hidden_states = self.mamba(766 hidden_states=hidden_states,767 cache_params=past_key_values,768 cache_position=cache_position,769 attention_mask=attention_mask,770 **kwargs,771 )772 self_attn_weights = None773 elif self.layer_type == "attention":774 hidden_states, self_attn_weights = self.self_attn(775 hidden_states=hidden_states,776 attention_mask=attention_mask,777 position_ids=position_ids,778 past_key_values=past_key_values,779 output_attentions=output_attentions,780 use_cache=use_cache,781 cache_position=cache_position,782 position_embeddings=position_embeddings,783 **kwargs,784 )785 786 # residual connection after attention787 hidden_states = residual + hidden_states788 789 # feed-forward790 residual = hidden_states791 hidden_states = self.pre_ff_layernorm(hidden_states)792 hidden_states = self.feed_forward(hidden_states)793 hidden_states = residual + hidden_states794 795 outputs = (hidden_states,)796 797 if output_attentions:798 outputs += (self_attn_weights,)799 800 return outputs801 802 803@auto_docstring804class BambaPreTrainedModel(PreTrainedModel):805 config: BambaConfig806 base_model_prefix = "model"807 supports_gradient_checkpointing = True808 _no_split_modules = ["BambaDecoderLayer"]809 _skip_keys_device_placement = "past_key_values"810 _supports_flash_attn = True811 _supports_sdpa = True812 # Note: only supports HybridMambaAttentionDynamicCache813 _is_stateful = True814 815 def _init_weights(self, module):816 super()._init_weights(module)817 if isinstance(module, BambaMixer):818 module.dt_bias.data.fill_(1.0)819 module.A_log.data = torch.log(torch.arange(1, module.num_heads + 1))820 module.D.data.fill_(1.0)821 822 823@auto_docstring824class BambaModel(BambaPreTrainedModel):825 def __init__(self, config: BambaConfig):826 super().__init__(config)827 self.padding_idx = config.pad_token_id828 self.vocab_size = config.vocab_size829 830 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)831 decoder_layers = []832 for i in range(config.num_hidden_layers):833 decoder_layers.append(BambaDecoderLayer(config, layer_idx=i, layer_type=config.layers_block_type[i]))834 self.layers = nn.ModuleList(decoder_layers)835 836 self._attn_implementation = config._attn_implementation837 self.final_layernorm = BambaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)838 self.rotary_emb = BambaRotaryEmbedding(config=config)839 840 self.gradient_checkpointing = False841 # Initialize weights and apply final processing842 self.post_init()843 844 @can_return_tuple845 @auto_docstring846 def forward(847 self,848 input_ids: Optional[torch.LongTensor] = None,849 attention_mask: Optional[torch.Tensor] = None,850 position_ids: Optional[torch.LongTensor] = None,851 past_key_values: Optional[HybridMambaAttentionDynamicCache] = None,852 inputs_embeds: Optional[torch.FloatTensor] = None,853 use_cache: Optional[bool] = None,854 output_attentions: Optional[bool] = None,855 output_hidden_states: Optional[bool] = None,856 cache_position: Optional[torch.LongTensor] = None,857 **kwargs: Unpack[BambaFlashAttentionKwargs],858 ) -> BaseModelOutputWithPast:859 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions860 output_hidden_states = (861 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states862 )863 use_cache = use_cache if use_cache is not None else self.config.use_cache864 865 if (input_ids is None) ^ (inputs_embeds is not None):866 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")867 868 if self.gradient_checkpointing and self.training and use_cache:869 logger.warning_once(870 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."871 )872 use_cache = False873 874 if inputs_embeds is None:875 inputs_embeds = self.embed_tokens(input_ids)876 hidden_states = inputs_embeds877 878 if use_cache and past_key_values is None:879 logger.warning_once(880 "Bamba requires an initialized `HybridMambaAttentionDynamicCache` to return a cache. None was "881 "provided, so no cache will be returned."882 )883 884 if cache_position is None:885 cache_position = torch.arange(hidden_states.shape[1], device=hidden_states.device)886 if position_ids is None:887 position_ids = cache_position.unsqueeze(0)888 889 causal_mask = self._update_causal_mask(890 attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions891 )892 mamba_mask = self._update_mamba_mask(attention_mask, cache_position)893 894 # create position embeddings to be shared across the decoder layers895 position_embeddings = self.rotary_emb(hidden_states, position_ids)896 897 all_hidden_states = () if output_hidden_states else None898 all_self_attns = () if output_attentions else None899 900 for decoder_layer in self.layers:901 # Depending on the layer type we opt for 2D base attention mask (Mamba) or 4D causal mask (Attention)902 layer_mask = mamba_mask if decoder_layer.layer_type == "mamba" else causal_mask903 904 if output_hidden_states:905 all_hidden_states += (hidden_states,)906 907 layer_outputs = decoder_layer(908 hidden_states,909 attention_mask=layer_mask,910 position_ids=position_ids,911 past_key_values=past_key_values,912 output_attentions=output_attentions,913 use_cache=use_cache,914 cache_position=cache_position,915 position_embeddings=position_embeddings,916 **kwargs,917 )918 919 hidden_states = layer_outputs[0]920 921 if output_attentions:922 if layer_outputs[1] is not None:923 # append attentions only of attention layers. Mamba layers return `None` as the attention weights924 all_self_attns += (layer_outputs[1],)925 926 hidden_states = self.final_layernorm(hidden_states)927 928 # add hidden states from the last decoder layer929 if output_hidden_states:930 all_hidden_states += (hidden_states,)931 932 if past_key_values and not past_key_values.has_previous_state:933 past_key_values.has_previous_state = True934 935 next_cache = None if not use_cache else past_key_values936 937 return BaseModelOutputWithPast(938 last_hidden_state=hidden_states,939 past_key_values=next_cache,940 hidden_states=all_hidden_states,941 attentions=all_self_attns,942 )943 944 def _update_causal_mask(945 self,946 attention_mask: torch.Tensor,947 input_tensor: torch.Tensor,948 cache_position: torch.Tensor,949 past_key_values: HybridMambaAttentionDynamicCache,950 output_attentions: bool,951 ):952 if self.config._attn_implementation == "flash_attention_2":953 if attention_mask is not None and 0.0 in attention_mask:954 return attention_mask955 return None956 957 # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in958 # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail959 # to infer the attention mask.960 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0961 962 # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward963 if self.config._attn_implementation == "sdpa" and not output_attentions:964 if AttentionMaskConverter._ignore_causal_mask_sdpa(965 attention_mask,966 inputs_embeds=input_tensor,967 past_key_values_length=past_seen_tokens,968 is_training=self.training,969 ):970 return None971 972 dtype = input_tensor.dtype973 sequence_length = input_tensor.shape[1]974 target_length = (975 attention_mask.shape[-1]976 if isinstance(attention_mask, torch.Tensor)977 else past_seen_tokens + sequence_length + 1978 )979 980 # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).981 causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(982 attention_mask,983 sequence_length=sequence_length,984 target_length=target_length,985 dtype=dtype,986 cache_position=cache_position,987 batch_size=input_tensor.shape[0],988 )989 990 if (991 self.config._attn_implementation == "sdpa"992 and attention_mask is not None993 and attention_mask.device.type in ["cuda", "xpu", "npu"]994 and not output_attentions995 ):996 # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when997 # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.998 # Details: https://github.com/pytorch/pytorch/issues/110213999 min_dtype = torch.finfo(dtype).min1000 causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)1001 1002 return causal_mask1003 1004 @staticmethod1005 def _prepare_4d_causal_attention_mask_with_cache_position(1006 attention_mask: torch.Tensor,1007 sequence_length: int,1008 target_length: int,1009 dtype: torch.dtype,1010 cache_position: torch.Tensor,1011 batch_size: int,1012 **kwargs,1013 ):1014 """1015 Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape1016 `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.1017 1018 Args:1019 attention_mask (`torch.Tensor`):1020 A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape1021 `(batch_size, 1, query_length, key_value_length)`.1022 sequence_length (`int`):1023 The sequence length being processed.1024 target_length (`int`):1025 The target length: when generating with static cache, the mask should be as long as the static cache,1026 to account for the 0 padding, the part of the cache that is not filled yet.1027 dtype (`torch.dtype`):1028 The dtype to use for the 4D attention mask.1029 cache_position (`torch.Tensor`):1030 Indices depicting the position of the input sequence tokens in the sequence.1031 batch_size (`torch.Tensor`):1032 Batch size.1033 """1034 if attention_mask is not None and attention_mask.dim() == 4:1035 # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.1036 causal_mask = attention_mask1037 else:1038 min_dtype = torch.finfo(dtype).min1039 causal_mask = torch.full(1040 (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device1041 )1042 if sequence_length != 1:1043 causal_mask = torch.triu(causal_mask, diagonal=1)1044 causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1)1045 causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)1046 if attention_mask is not None:1047 causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit1048 mask_length = attention_mask.shape[-1]1049 padding_attention_mask = (attention_mask[:, None, None, :] == attention_mask[:, None, :, None])[1050 :, :, -sequence_length:, :1051 ].to(dtype)1052 padding_mask = causal_mask[:, :, :, :mask_length] + padding_attention_mask1053 padding_mask = padding_mask == 01054 causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(1055 padding_mask, min_dtype1056 )1057 1058 return causal_mask1059 1060 def _update_mamba_mask(self, attention_mask, cache_position):1061 """1062 No need for zeroing states when1063 1. Cached forward1064 2. Attending to all inputs1065 """1066 mamba_mask = attention_mask1067 if cache_position[0] > 0 or (attention_mask is not None and torch.all(attention_mask == 1)):1068 mamba_mask = None1069 return mamba_mask1070 1071 1072class BambaForCausalLM(LlamaForCausalLM):1073 def __init__(self, config):1074 super().__init__(config)1075 self.z_loss_coefficient = config.z_loss_coefficient1076 1077 # Initialize weights and apply final processing1078 self.post_init()1079 1080 def forward(1081 self,1082 input_ids: Optional[torch.LongTensor] = None,1083 attention_mask: Optional[torch.Tensor] = None,1084 position_ids: Optional[torch.LongTensor] = None,1085 past_key_values: Optional[HybridMambaAttentionDynamicCache] = None,1086 inputs_embeds: Optional[torch.FloatTensor] = None,1087 labels: Optional[torch.LongTensor] = None,1088 use_cache: Optional[bool] = None,1089 output_attentions: Optional[bool] = None,1090 output_hidden_states: Optional[bool] = None,1091 cache_position: Optional[torch.LongTensor] = None,1092 logits_to_keep: Union[int, torch.Tensor] = 0,1093 **kwargs,1094 ) -> CausalLMOutputWithPast:1095 r"""1096 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):1097 Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,1098 config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored1099 (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.1100 1101 Example:1102 1103 ```python1104 >>> from transformers import AutoTokenizer, BambaForCausalLM1105 1106 >>> model = BambaForCausalLM.from_pretrained("...")1107 >>> tokenizer = AutoTokenizer.from_pretrained("...")1108 1109 >>> prompt = "Hey, are you conscious? Can you talk to me?"1110 >>> inputs = tokenizer(prompt, return_tensors="pt")1111 1112 >>> # Generate1113 >>> generate_ids = model.generate(inputs.input_ids, max_length=30)1114 >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]1115 "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."1116 ```"""1117 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1118 output_hidden_states = (1119 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1120 )1121 1122 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)1123 outputs: BaseModelOutputWithPast = self.model(1124 input_ids=input_ids,1125 attention_mask=attention_mask,1126 position_ids=position_ids,1127 past_key_values=past_key_values,1128 inputs_embeds=inputs_embeds,1129 use_cache=use_cache,1130 output_attentions=output_attentions,1131 output_hidden_states=output_hidden_states,1132 cache_position=cache_position,1133 **kwargs,1134 )1135 1136 hidden_states = outputs.last_hidden_state1137 # Only compute necessary logits, and do not upcast them to float if we are not computing the loss1138 slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep1139 logits = self.lm_head(hidden_states[:, slice_indices, :])1140 1141 loss = None1142 if labels is not None:1143 loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)1144 if self.z_loss_coefficient > 0:1145 # Type-match loss, but avoid upcasting large logits tensor until after it's been reduced on dim -11146 z_loss = logits.logsumexp(dim=-1).to(dtype=loss.dtype).pow(2).mean()1147 loss = loss + self.z_loss_coefficient * z_loss1148 1149 return CausalLMOutputWithPast(1150 loss=loss,1151 logits=logits,1152 past_key_values=outputs.past_key_values,1153 hidden_states=outputs.hidden_states,1154 attentions=outputs.attentions,1155 )1156 1157 def prepare_inputs_for_generation(1158 self,1159 input_ids,1160 past_key_values=None,1161 attention_mask=None,1162 inputs_embeds=None,1163 cache_position=None,1164 position_ids=None,1165 use_cache=True,1166 **kwargs,1167 ):1168 # Overwritten -- has a unique cache type, `HybridMambaAttentionDynamicCache`1169 1170 empty_past_kv = past_key_values is None1171 1172 # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens1173 # Exception 1: when passing input_embeds, input_ids may be missing entries1174 # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here1175 # Exception 3: with synced GPUs cache_position may go out of bounds, but we only want dummy token in that case.1176 # (we can't check exception 3 while compiling)1177 if not empty_past_kv:1178 if (1179 inputs_embeds is not None # Exception 11180 or cache_position[-1] >= input_ids.shape[1] # Exception 31181 ):1182 input_ids = input_ids[:, -cache_position.shape[0] :]1183 elif input_ids.shape[1] != cache_position.shape[0]: # Default case (the "else", a no op, is Exception 2)1184 input_ids = input_ids[:, cache_position]1185 else:1186 past_key_values = HybridMambaAttentionDynamicCache(1187 self.config, input_ids.shape[0], self.dtype, device=self.device1188 )1189 1190 if attention_mask is not None and position_ids is None:1191 # create position_ids on the fly for batch generation1192 position_ids = attention_mask.long().cumsum(-1) - 11193 position_ids.masked_fill_(attention_mask == 0, 1)1194 if not empty_past_kv:1195 position_ids = position_ids[:, -input_ids.shape[1] :]1196 1197 # if `inputs_embeds` are passed, we only want to use them in the 1st generation step1198 if inputs_embeds is not None and empty_past_kv:1199 model_inputs = {"inputs_embeds": inputs_embeds}1200 else: