From f614c2548dc7cadd30311fdb47cc5f79c3f4edc1 Mon Sep 17 00:00:00 2001 From: kaixuanliu Date: Thu, 2 Jul 2026 08:30:38 +0000 Subject: [PATCH 1/8] fix mask return-type contract regression and add correctness guard for xpu Signed-off-by: kaixuanliu --- src/transformers/masking_utils.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/transformers/masking_utils.py b/src/transformers/masking_utils.py index 130b1a1dec66..483c23b2b40d 100644 --- a/src/transformers/masking_utils.py +++ b/src/transformers/masking_utils.py @@ -231,6 +231,7 @@ def _can_skip_causal_mask_xpu( padding_mask: torch.Tensor | None, query_length: int, kv_length: int, + kv_offset: int, local_attention_size: int | None, ) -> bool: """ @@ -258,8 +259,10 @@ def _can_skip_causal_mask_xpu( # 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 + # 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. + if kv_offset != 0: + return False return padding_mask[:, :query_length].all() and not padding_mask[:, query_length:].any() @@ -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, query_length, kv_length, 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 @@ -1552,15 +1555,13 @@ 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. Some models use this mapping as the + # signal that the mask was already prepared by generate, even when all layers share the same 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) From 8f3d6dc3bedd6bab4c5d230bd084a1defb3f9bd6 Mon Sep 17 00:00:00 2001 From: kaixuanliu Date: Thu, 2 Jul 2026 08:43:56 +0000 Subject: [PATCH 2/8] update comment Signed-off-by: kaixuanliu --- src/transformers/masking_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/transformers/masking_utils.py b/src/transformers/masking_utils.py index 483c23b2b40d..9d87a25c6ea2 100644 --- a/src/transformers/masking_utils.py +++ b/src/transformers/masking_utils.py @@ -1555,8 +1555,7 @@ def create_masks_for_generate( "block_sequence_ids": block_sequence_ids, } - # If the attribute exists, we need several masks keyed by layer type. Some models use this mapping as the - # signal that the mask was already prepared by generate, even when all layers share the same type. + # 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 From 61b84009cbc11b84ef064003df33862c4869299f Mon Sep 17 00:00:00 2001 From: kaixuanliu Date: Thu, 2 Jul 2026 09:26:01 +0000 Subject: [PATCH 3/8] fix failed CI cases for deepseek_v32 Signed-off-by: kaixuanliu --- .../deepseek_v32/modeling_deepseek_v32.py | 19 +++--- .../deepseek_v32/modular_deepseek_v32.py | 63 ++++++++++++++++++- .../glm_moe_dsa/modeling_glm_moe_dsa.py | 19 +++--- .../models/glm_moe_dsa/modular_glm_moe_dsa.py | 19 +++--- 4 files changed, 96 insertions(+), 24 deletions(-) diff --git a/src/transformers/models/deepseek_v32/modeling_deepseek_v32.py b/src/transformers/models/deepseek_v32/modeling_deepseek_v32.py index 037e8064a77f..d4a7c6dd67f2 100644 --- a/src/transformers/models/deepseek_v32/modeling_deepseek_v32.py +++ b/src/transformers/models/deepseek_v32/modeling_deepseek_v32.py @@ -719,13 +719,18 @@ 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, - ) + # `generate` may pass an already prepared mask keyed by layer type (see `create_masks_for_generate`). + # As all layers share the same attention type, we simply unwrap the single mask it contains. + if isinstance(attention_mask, dict): + causal_mask = attention_mask[self.config.layer_types[0]] + else: + 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, + ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) diff --git a/src/transformers/models/deepseek_v32/modular_deepseek_v32.py b/src/transformers/models/deepseek_v32/modular_deepseek_v32.py index 1fae48192b35..1704455d7be0 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,62 @@ 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) + + # `generate` may pass an already prepared mask keyed by layer type (see `create_masks_for_generate`). + # As all layers share the same attention type, we simply unwrap the single mask it contains. + if isinstance(attention_mask, dict): + causal_mask = attention_mask[self.config.layer_types[0]] + else: + 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, + ) + + 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]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + 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 fcf03648a1aa..d04bbbf4b1d4 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 @@ -736,13 +736,18 @@ 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, - ) + # `generate` may pass an already prepared mask keyed by layer type (see `create_masks_for_generate`). + # As all layers share the same attention type, we simply unwrap the single mask it contains. + if isinstance(attention_mask, dict): + causal_mask = attention_mask[self.config.layer_types[0]] + else: + 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, + ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) 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 72b2095376eb..03abdb2c76a9 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 @@ -297,13 +297,18 @@ 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, - ) + # `generate` may pass an already prepared mask keyed by layer type (see `create_masks_for_generate`). + # As all layers share the same attention type, we simply unwrap the single mask it contains. + if isinstance(attention_mask, dict): + causal_mask = attention_mask[self.config.layer_types[0]] + else: + 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, + ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) From c1beedc466889f2587875047283dae9152b9843d Mon Sep 17 00:00:00 2001 From: kaixuanliu Date: Thu, 2 Jul 2026 22:43:46 +0800 Subject: [PATCH 4/8] Update src/transformers/masking_utils.py Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> --- src/transformers/masking_utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/transformers/masking_utils.py b/src/transformers/masking_utils.py index 9d87a25c6ea2..cf99e471ab42 100644 --- a/src/transformers/masking_utils.py +++ b/src/transformers/masking_utils.py @@ -261,9 +261,7 @@ def _can_skip_causal_mask_xpu( # 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. - if kv_offset != 0: - return False - return padding_mask[:, :query_length].all() and not padding_mask[:, query_length:].any() +return kv_offset == 0 and padding_mask[:, :query_length].all() and not padding_mask[:, query_length:].any() def _ignore_causal_mask_sdpa( From 77ce3778fc3e655abcc65caf199f7531b375acd5 Mon Sep 17 00:00:00 2001 From: kaixuanliu Date: Thu, 2 Jul 2026 14:48:09 +0000 Subject: [PATCH 5/8] update code Signed-off-by: kaixuanliu --- .../deepseek_v32/modeling_deepseek_v32.py | 26 +++++++++---------- .../deepseek_v32/modular_deepseek_v32.py | 26 +++++++++---------- .../glm_moe_dsa/modeling_glm_moe_dsa.py | 26 +++++++++---------- .../models/glm_moe_dsa/modular_glm_moe_dsa.py | 26 +++++++++---------- 4 files changed, 48 insertions(+), 56 deletions(-) diff --git a/src/transformers/models/deepseek_v32/modeling_deepseek_v32.py b/src/transformers/models/deepseek_v32/modeling_deepseek_v32.py index d4a7c6dd67f2..c70b2397ff13 100644 --- a/src/transformers/models/deepseek_v32/modeling_deepseek_v32.py +++ b/src/transformers/models/deepseek_v32/modeling_deepseek_v32.py @@ -719,26 +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) - # `generate` may pass an already prepared mask keyed by layer type (see `create_masks_for_generate`). - # As all layers share the same attention type, we simply unwrap the single mask it contains. - if isinstance(attention_mask, dict): - causal_mask = attention_mask[self.config.layer_types[0]] - else: - 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 = {self.config.layer_types[0]: 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 1704455d7be0..8d2d7fe486fe 100644 --- a/src/transformers/models/deepseek_v32/modular_deepseek_v32.py +++ b/src/transformers/models/deepseek_v32/modular_deepseek_v32.py @@ -394,26 +394,24 @@ def forward( position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens position_ids = position_ids.unsqueeze(0) - # `generate` may pass an already prepared mask keyed by layer type (see `create_masks_for_generate`). - # As all layers share the same attention type, we simply unwrap the single mask it contains. - if isinstance(attention_mask, dict): - causal_mask = attention_mask[self.config.layer_types[0]] - else: - 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 = {self.config.layer_types[0]: 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/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index 51cebb24c579..e72aff114b80 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,27 +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) - # `generate` may pass an already prepared mask keyed by layer type (see `create_masks_for_generate`). - # As all layers share the same attention type, we simply unwrap the single mask it contains. - if isinstance(attention_mask, dict): - causal_mask = attention_mask[self.config.layer_types[0]] - else: - 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 = {self.config.layer_types[0]: 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 85943078bc92..7f92598146bd 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,27 +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) - # `generate` may pass an already prepared mask keyed by layer type (see `create_masks_for_generate`). - # As all layers share the same attention type, we simply unwrap the single mask it contains. - if isinstance(attention_mask, dict): - causal_mask = attention_mask[self.config.layer_types[0]] - else: - 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 = {self.config.layer_types[0]: 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, From 16ec359e6cd1747b4cbbd4f94181d8c3477a6b59 Mon Sep 17 00:00:00 2001 From: kaixuanliu Date: Thu, 2 Jul 2026 14:59:28 +0000 Subject: [PATCH 6/8] fix LINT issue Signed-off-by: kaixuanliu --- src/transformers/masking_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/masking_utils.py b/src/transformers/masking_utils.py index cf99e471ab42..a8d27f33ab91 100644 --- a/src/transformers/masking_utils.py +++ b/src/transformers/masking_utils.py @@ -261,7 +261,7 @@ def _can_skip_causal_mask_xpu( # 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 kv_offset == 0 and padding_mask[:, :query_length].all() and not padding_mask[:, query_length:].any() + return kv_offset == 0 and padding_mask[:, :query_length].all() and not padding_mask[:, query_length:].any() def _ignore_causal_mask_sdpa( From 132c6c48a2dde42035645abcd918fe0ef54bf661 Mon Sep 17 00:00:00 2001 From: kaixuanliu Date: Fri, 3 Jul 2026 01:18:58 +0000 Subject: [PATCH 7/8] update code Signed-off-by: kaixuanliu --- src/transformers/models/deepseek_v32/modeling_deepseek_v32.py | 2 +- src/transformers/models/deepseek_v32/modular_deepseek_v32.py | 2 +- src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py | 2 +- src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/transformers/models/deepseek_v32/modeling_deepseek_v32.py b/src/transformers/models/deepseek_v32/modeling_deepseek_v32.py index c70b2397ff13..782ed117c641 100644 --- a/src/transformers/models/deepseek_v32/modeling_deepseek_v32.py +++ b/src/transformers/models/deepseek_v32/modeling_deepseek_v32.py @@ -728,7 +728,7 @@ def forward( "past_key_values": past_key_values, "position_ids": position_ids, } - causal_mask_mapping = {self.config.layer_types[0]: create_causal_mask(**mask_kwargs)} + 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) diff --git a/src/transformers/models/deepseek_v32/modular_deepseek_v32.py b/src/transformers/models/deepseek_v32/modular_deepseek_v32.py index 8d2d7fe486fe..4410069576dd 100644 --- a/src/transformers/models/deepseek_v32/modular_deepseek_v32.py +++ b/src/transformers/models/deepseek_v32/modular_deepseek_v32.py @@ -403,7 +403,7 @@ def forward( "past_key_values": past_key_values, "position_ids": position_ids, } - causal_mask_mapping = {self.config.layer_types[0]: create_causal_mask(**mask_kwargs)} + 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) 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 e72aff114b80..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 @@ -707,7 +707,7 @@ def forward( "past_key_values": past_key_values, "position_ids": position_ids, } - causal_mask_mapping = {self.config.layer_types[0]: create_causal_mask(**mask_kwargs)} + 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) 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 7f92598146bd..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 @@ -367,7 +367,7 @@ def forward( "past_key_values": past_key_values, "position_ids": position_ids, } - causal_mask_mapping = {self.config.layer_types[0]: create_causal_mask(**mask_kwargs)} + 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) From 55f628e66ecfae6fe2e738abe11944fde34cc557 Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Mon, 6 Jul 2026 10:41:21 +0900 Subject: [PATCH 8/8] it needs to be skipped only if cache is empty --- src/transformers/masking_utils.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/src/transformers/masking_utils.py b/src/transformers/masking_utils.py index a8d27f33ab91..321ef27de913 100644 --- a/src/transformers/masking_utils.py +++ b/src/transformers/masking_utils.py @@ -229,8 +229,9 @@ 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: @@ -252,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 when the cache is empty. - return kv_offset == 0 and padding_mask[:, :query_length].all() and not padding_mask[:, query_length:].any() + 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: @@ -289,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, kv_offset, 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 @@ -297,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 @@ -523,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