diff --git a/src/transformers/masking_utils.py b/src/transformers/masking_utils.py index 130b1a1dec66..321ef27de913 100644 --- a/src/transformers/masking_utils.py +++ b/src/transformers/masking_utils.py @@ -229,8 +229,10 @@ def maybe_pad_block_sequence_ids( def _can_skip_causal_mask_xpu( padding_mask: torch.Tensor | None, - query_length: int, + q_length: int, kv_length: int, + q_offset: int, + kv_offset: int, local_attention_size: int | None, ) -> bool: """ @@ -251,22 +253,23 @@ def _can_skip_causal_mask_xpu( if padding_mask is None: # Without padding mask, can skip if single query token or full causal attention - return query_length == 1 or kv_length == query_length + return q_length == 1 or kv_length == q_length # XPU allows skipping under additional conditions when padding_mask is provided - if query_length == 1: + if q_length == 1: # Single query token: skip only if no padding tokens present return padding_mask.all() - # XPU-specific: check if query window is all True and rest is all False - # This allows XPU to optimize the 1st token in static cache - return padding_mask[:, :query_length].all() and not padding_mask[:, query_length:].any() + # XPU-specific: check if query window is all True and rest is all False. + # This allows XPU to optimize the 1st token in static cache when the cache is empty. + return q_offset == 0 and padding_mask[:, :q_length].all() and not padding_mask[:, q_length:].any() def _ignore_causal_mask_sdpa( padding_mask: torch.Tensor | None, - query_length: int, + q_length: int, kv_length: int, + q_offset: int, kv_offset: int, local_attention_size: int | None = None, ) -> bool: @@ -288,7 +291,7 @@ def _ignore_causal_mask_sdpa( # - Single query tokens use the same logic as CUDA # - Multi-query tokens can skip if padding_mask is provided and correctly structured # (all True in query window, all False after) - return _can_skip_causal_mask_xpu(padding_mask, query_length, kv_length, local_attention_size) + return _can_skip_causal_mask_xpu(padding_mask, q_length, kv_length, q_offset, kv_offset, local_attention_size) # When using `torch.export` or `torch.onnx.dynamo_export`, we must pass an example input, and `is_causal` behavior is # hard-coded to the forward. If a user exports a model with query_length > 1, the exported model will hard-code `is_causal=True` # which is in general wrong (see https://github.com/pytorch/pytorch/issues/108108). Thus, we only set @@ -296,7 +299,7 @@ def _ignore_causal_mask_sdpa( if ( not is_tracing(padding_mask) # only cases when lower and upper diags are the same, see https://github.com/pytorch/pytorch/issues/108108 - and (query_length == 1 or kv_length == query_length) + and (q_length == 1 or kv_length == q_length) # in this case we need to add special patterns to the mask so cannot be skipped otherwise and (local_attention_size is None or kv_length < local_attention_size) # In this case, we need to add padding to the mask, so cannot be skipped otherwise @@ -522,7 +525,9 @@ def sdpa_mask( # Under specific conditions, we can avoid materializing the mask # 1. Causal masks can rely on the `is_causal` argument # 2. Bidirectional do not need any further processing (no bias) - if allow_is_causal_skip and _ignore_causal_mask_sdpa(padding_mask, q_length, kv_length, kv_offset, local_size): + if allow_is_causal_skip and _ignore_causal_mask_sdpa( + padding_mask, q_length, kv_length, q_offset, kv_offset, local_size + ): return None if allow_is_bidirectional_skip and _ignore_bidirectional_mask_sdpa(padding_mask, kv_length, local_size): return None @@ -1552,15 +1557,12 @@ def create_masks_for_generate( "block_sequence_ids": block_sequence_ids, } - # If the attribute exists, we need several masks - unless every layer shares the same type, in which - # case we return a single mask. + # If the attribute exists, we need several masks keyed by layer type. if hasattr(effective_config, "layer_types"): layer_patterns = set(effective_config.layer_types) # Without a registered attention-mask function, defer to the model by returning the raw attention mask if any(layer_type not in LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING for layer_type in layer_patterns): return attention_mask - if len(layer_patterns) == 1: - return LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING[next(iter(layer_patterns))](**mask_kwargs) causal_masks = {} for layer_pattern in layer_patterns: causal_masks[layer_pattern] = LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING[layer_pattern](**mask_kwargs) diff --git a/src/transformers/models/deepseek_v32/modeling_deepseek_v32.py b/src/transformers/models/deepseek_v32/modeling_deepseek_v32.py index 037e8064a77f..782ed117c641 100644 --- a/src/transformers/models/deepseek_v32/modeling_deepseek_v32.py +++ b/src/transformers/models/deepseek_v32/modeling_deepseek_v32.py @@ -719,21 +719,24 @@ def forward( position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens position_ids = position_ids.unsqueeze(0) - causal_mask = create_causal_mask( - config=self.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - past_key_values=past_key_values, - position_ids=position_ids, - ) + # It may already have been prepared by e.g. `generate` + if not isinstance(causal_mask_mapping := attention_mask, dict): + mask_kwargs = { + "config": self.config, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + causal_mask_mapping = {"deepseek_sparse_attention": create_causal_mask(**mask_kwargs)} hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) - for decoder_layer in self.layers[: self.config.num_hidden_layers]: + for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): hidden_states = decoder_layer( hidden_states, - attention_mask=causal_mask, + attention_mask=causal_mask_mapping[self.config.layer_types[i]], position_embeddings=position_embeddings, position_ids=position_ids, past_key_values=past_key_values, diff --git a/src/transformers/models/deepseek_v32/modular_deepseek_v32.py b/src/transformers/models/deepseek_v32/modular_deepseek_v32.py index 1fae48192b35..4410069576dd 100644 --- a/src/transformers/models/deepseek_v32/modular_deepseek_v32.py +++ b/src/transformers/models/deepseek_v32/modular_deepseek_v32.py @@ -29,12 +29,14 @@ import torch.nn.functional as F from huggingface_hub.dataclasses import strict -from ...cache_utils import Cache +from ...cache_utils import Cache, DynamicCache +from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_outputs import BaseModelOutputWithPast from ...modeling_rope_utils import RotaryEmbeddingConfigMixin from ...modeling_utils import ALL_ATTENTION_FUNCTIONS from ...processing_utils import Unpack -from ...utils import auto_docstring, logging +from ...utils import TransformersKwargs, auto_docstring, logging from ..deepseek_v3.modeling_deepseek_v3 import ( DeepseekV3Attention, DeepseekV3ForCausalLM, @@ -368,7 +370,60 @@ class DeepseekV32PreTrainedModel(DeepseekV3PreTrainedModel): class DeepseekV32Model(DeepseekV3Model): - pass + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + # It may already have been prepared by e.g. `generate` + if not isinstance(causal_mask_mapping := attention_mask, dict): + mask_kwargs = { + "config": self.config, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + causal_mask_mapping = {"deepseek_sparse_attention": create_causal_mask(**mask_kwargs)} + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + + for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask_mapping[self.config.layer_types[i]], + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) class DeepseekV32ForCausalLM(DeepseekV3ForCausalLM): diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index 52f8b6cd3965..941522423c74 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -698,22 +698,25 @@ def forward( position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens position_ids = position_ids.unsqueeze(0) - causal_mask = create_causal_mask( - config=self.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - past_key_values=past_key_values, - position_ids=position_ids, - ) + # It may already have been prepared by e.g. `generate` + if not isinstance(causal_mask_mapping := attention_mask, dict): + mask_kwargs = { + "config": self.config, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + causal_mask_mapping = {"deepseek_sparse_attention": create_causal_mask(**mask_kwargs)} hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) topk_indices = None # MAIN DIFF with DSV3.2 - for decoder_layer in self.layers[: self.config.num_hidden_layers]: + for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): hidden_states, topk_indices = decoder_layer( hidden_states, - attention_mask=causal_mask, + attention_mask=causal_mask_mapping[self.config.layer_types[i]], position_embeddings=position_embeddings, position_ids=position_ids, past_key_values=past_key_values, diff --git a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py index e9b73c57407f..c2c347fcfea9 100644 --- a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py @@ -358,22 +358,25 @@ def forward( position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens position_ids = position_ids.unsqueeze(0) - causal_mask = create_causal_mask( - config=self.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - past_key_values=past_key_values, - position_ids=position_ids, - ) + # It may already have been prepared by e.g. `generate` + if not isinstance(causal_mask_mapping := attention_mask, dict): + mask_kwargs = { + "config": self.config, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + causal_mask_mapping = {"deepseek_sparse_attention": create_causal_mask(**mask_kwargs)} hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) topk_indices = None # MAIN DIFF with DSV3.2 - for decoder_layer in self.layers[: self.config.num_hidden_layers]: + for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): hidden_states, topk_indices = decoder_layer( hidden_states, - attention_mask=causal_mask, + attention_mask=causal_mask_mapping[self.config.layer_types[i]], position_embeddings=position_embeddings, position_ids=position_ids, past_key_values=past_key_values,