Add deepseek 3.2 exp #41251

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

vasqu merged 31 commits into main from add-deepseek-exp

Jun 10, 2026

docs/source/en/_toctree.yml

-Original file line number
+Diff line change
@@ Expand Up / @@ -569,6 +569,8 @@ @@
             title: DeepSeek-V2
           - local: model_doc/deepseek_v3
             title: DeepSeek-V3
+          - local: model_doc/deepseek_v32
+            title: DeepSeek-V3.2
           - local: model_doc/deepseek_v4
             title: DeepSeek-V4
           - local: model_doc/dialogpt
@@ Expand Down @@

docs/source/en/model_doc/deepseek_v32.md

-Original file line number
+Diff line change
@@ -0,0 +1,100 @@
+    <!--Copyright 2025 the HuggingFace Team. All rights reserved.
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+        http://www.apache.org/licenses/LICENSE-2.0
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+    ⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer.
+    -->
+    *This model was published in HF papers on 2025-12-02 and contributed to Hugging Face Transformers on 2026-06-10.*
+    <div style="float: right;">
+        <div class="flex flex-wrap space-x-1">
+            <img alt="PyTorch" src="https://img.shields.io/badge/PyTorch-DE3412?style=flat&logo=pytorch&logoColor=white">
+            <img alt="FlashAttention" src="https://img.shields.io/badge/%E2%9A%A1%EF%B8%8E%20FlashAttention-eae0c8?style=flat">
+            <img alt="SDPA" src="https://img.shields.io/badge/SDPA-DE3412?style=flat&logo=pytorch&logoColor=white">
+            <img alt="FP8" src="https://img.shields.io/badge/FP8-4d8a4d?style=flat">
+        </div>
+    </div>
+    # DeepSeek-V3.2
+    ## Overview
+    [DeepSeek-V3.2-Exp](https://huggingface.co/deepseek-ai/DeepSeek-V3.2-Exp) is an experimental release from DeepSeek-AI that introduces **DeepSeek Sparse Attention (DSA)**, a trainable, fine-grained sparse attention mechanism designed to improve training and inference efficiency in long-context scenarios. It is built directly on top of [DeepSeek-V3.1-Terminus](https://huggingface.co/deepseek-ai/DeepSeek-V3.1-Terminus): the model keeps the same 685B-parameter Mixture-of-Experts (MoE) backbone and Multi-head Latent Attention (MLA), and is obtained through continued training that adds the sparse-attention indexer while deliberately aligning the training distribution with V3.1-Terminus so the two models can be compared head-to-head.
+    The work was later extended in the [DeepSeek-V3.2 technical report](https://huggingface.co/papers/2512.02556), *DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models*, which pairs DSA with a scalable reinforcement-learning framework and reports gold-medal level results on competition math (IMO) and competitive programming (IOI) benchmarks.
+    The abstract from the DeepSeek-V3.2-Exp release is the following:
+    *We introduce DeepSeek-V3.2-Exp, an experimental version of our model that incorporates DeepSeek Sparse Attention (DSA) to explore and validate optimizations for training and inference efficiency in long-context scenarios. DeepSeek Sparse Attention achieves fine-grained sparse attention for the first time with minimal impact on model output quality. Built upon DeepSeek-V3.1-Terminus, DeepSeek-V3.2-Exp delivers substantially improved efficiency in both training and inference, especially in long-context settings, while maintaining virtually identical benchmark performance.*
+    ### DeepSeek Sparse Attention (DSA)
+    DSA reduces the quadratic cost of attention over long sequences by attending only to a selected subset of past tokens. It has two components:
+. **Lightning indexer.** A lightweight, low-head-count scoring module computes an *index score* between each query and every preceding key. In the reference implementation it runs in FP8 with a Hadamard (`rotate_activation`) transform; because the transform is orthogonal (`Hq·Hk = q·k`) and FP8 is only a precision optimization, the transformers port computes the same scores directly in bf16/fp32, keeping the indexer cheap relative to the main attention.
+. **Fine-grained token selection.** For each query the indexer keeps the top-`index_topk` (2048 by default) tokens, and main MLA attention is then computed only over those tokens via an additive mask. This turns the per-query attention cost from `O(L)` to `O(index_topk)` for long sequences when using `flash_mla`, which is not supported yet 😉.
+    The indexer keeps its own small per-token key cache (single-head, `index_head_dim`) alongside the main K/V cache. In transformers this lives on a dedicated cache layer — [`DynamicIndexedLayer`] for growing caches and [`StaticIndexedLayer`] for static / `torch.compile` caches — and is updated through `past_key_values.update_indexer()`.
+    In DeepSeek-V3.2 **every layer runs its own indexer** — there is no cross-layer top-k sharing.
+    > [!NOTE]
+    > **The MLA query LoRA path (`q_lora_rank`) is required.** The indexer scores queries from the low-rank query latent `q_a_layernorm(q_a_proj(x))` (its `wq_b` projection is sized by `q_lora_rank`), so the model always uses the LoRA query path and `q_lora_rank` must be set — the released checkpoint uses `1536`. The optional non-LoRA `q_proj` path that [DeepSeek-V3](./deepseek_v3) exposes for `q_lora_rank=None` is **not supported** here: without the query latent there is nothing for the indexer to consume.
+    ## Usage examples
+    DeepSeek-V3.2-Exp is distributed as an FP8 checkpoint. The indexer projections are kept out of FP8 quantization, since the checkpoint stores them in bf16/fp32:
+    ```python
+    from transformers import FineGrainedFP8Config, AutoModelForCausalLM, AutoTokenizer
+    import torch
+    model_name = "deepseek-ai/DeepSeek-V3.2-Exp"
+    quantization_config = FineGrainedFP8Config(
+        modules_to_not_convert=["model.layers.*.mlp.gate.*", "*.self_attn.indexer.weights_proj.*"],
+        weight_block_size=(128, 128),
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype="auto",
+        device_map="auto",
+        quantization_config=quantization_config,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    inputs = tokenizer("What are we having for dinner?", return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=20)
+    print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+    ```
+    The original code can be found [here](https://github.com/deepseek-ai/DeepSeek-V3.2-Exp).
+    ## DeepseekV32Config
+    [[autodoc]] DeepseekV32Config
+    ## DeepseekV32PreTrainedModel
+    [[autodoc]] DeepseekV32PreTrainedModel
+        - forward
+    ## DeepseekV32Model
+    [[autodoc]] DeepseekV32Model
+        - forward
+    ## DeepseekV32ForCausalLM
+    [[autodoc]] DeepseekV32ForCausalLM

docs/source/en/model_doc/glm_moe_dsa.md

-Original file line number
+Diff line change
@@ Expand Up / @@ -66,6 +66,9 @@ print(tokenizer.decode(output[0], skip_special_tokens=True)) @@
     </hfoption>
     </hfoptions>
+    > [!NOTE]
+    > **The MLA query LoRA path (`q_lora_rank`) is required.** Like DeepSeek-V3.2, the DSA indexer scores queries from the low-rank query latent `q_a_layernorm(q_a_proj(x))` (its `wq_b` projection is sized by `q_lora_rank`), so the model always uses the LoRA query path and `q_lora_rank` must be set — the released checkpoint uses `2048`. The optional non-LoRA `q_proj` path that [DeepSeek-V3](./deepseek_v3) exposes for `q_lora_rank=None` is **not supported** here: without the query latent there is nothing for the indexer to consume.
     ## GlmMoeDsaConfig
     [[autodoc]] GlmMoeDsaConfig
@@ Expand Down @@

src/transformers/__init__.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -367,7 +367,9 @@ @@
         _import_structure["cache_utils"] = [
             "CacheLayerMixin",
             "DynamicLayer",
+            "DynamicIndexedLayer",
             "StaticLayer",
+            "StaticIndexedLayer",
             "StaticSlidingWindowLayer",
             "QuantoQuantizedLayer",
             "HQQQuantizedLayer",
@@ Expand Down Expand Up / @@ -487,12 +489,14 @@ @@
         from .backbone_utils import BackboneConfigMixin, BackboneMixin
         from .cache_utils import Cache as Cache
         from .cache_utils import DynamicCache as DynamicCache
+        from .cache_utils import DynamicIndexedLayer as DynamicIndexedLayer
         from .cache_utils import DynamicLayer as DynamicLayer
         from .cache_utils import EncoderDecoderCache as EncoderDecoderCache
         from .cache_utils import HQQQuantizedLayer as HQQQuantizedLayer
         from .cache_utils import QuantizedCache as QuantizedCache
         from .cache_utils import QuantoQuantizedLayer as QuantoQuantizedLayer
         from .cache_utils import StaticCache as StaticCache
+        from .cache_utils import StaticIndexedLayer as StaticIndexedLayer
         from .cache_utils import StaticLayer as StaticLayer
         from .cache_utils import StaticSlidingWindowLayer as StaticSlidingWindowLayer
         from .configuration_utils import PreTrainedConfig as PreTrainedConfig
@@ Expand Down @@

src/transformers/cache_utils.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -274,6 +274,82 @@ def crop(self, max_length: int) -> None: @@
             self.cumulative_length = self.keys.shape[-2]
+    class DynamicIndexedLayer(DynamicLayer):
+        """
+        A cache layer that extends `DynamicLayer` with an extra indexer key cache for Dynamic Sparse Attention (DSA)
+        models (e.g. GLM MoE DSA, DeepSeek V32).
+        The main K/V cache stores tensors of shape `[batch_size, num_heads, seq_len, head_dim]` (inherited).
+        The indexer key cache stores a tensor of shape `[batch_size, seq_len, index_head_dim]` (3D, single-head).
+        """
+        # Auto-registers in ``LAYER_TYPE_CACHE_MAPPING`` so ``DynamicCache`` dispatches DSA layers here.
+        layer_type = "deepseek_sparse_attention"
+        def __init__(self, config: PreTrainedConfig | None = None):
+            super().__init__(config)
+            self.indexer_keys: torch.Tensor | None = None
+            self.is_indexer_initialized: bool = False
+        def lazy_initialization_indexer(self, indexer_key_states: torch.Tensor) -> None:
+            self.indexer_dtype, self.indexer_device = indexer_key_states.dtype, indexer_key_states.device
+            self.indexer_keys = torch.tensor([], dtype=self.indexer_dtype, device=self.indexer_device)
+            self.is_indexer_initialized = True
+        def update_indexer(self, indexer_key_states: torch.Tensor) -> torch.Tensor:
+            """
+            Update the indexer key cache by concatenation, and return the full indexer keys.
+            Args:
+                indexer_key_states (`torch.Tensor`): New indexer keys, shape `[batch_size, seq_len, index_head_dim]`.
+            Returns:
+                `torch.Tensor`: The full cached indexer keys, shape `[batch_size, total_len, index_head_dim]`.
+            """
+            if not self.is_indexer_initialized:
+                self.lazy_initialization_indexer(indexer_key_states)
+            self.indexer_keys = torch.cat([self.indexer_keys, indexer_key_states], dim=1)
+            return self.indexer_keys
+        def offload(self):
+            super().offload()
+            if self.is_indexer_initialized:
+                self.indexer_keys = self.indexer_keys.to("cpu", non_blocking=True)
+        def prefetch(self):
+            super().prefetch()
+            if self.is_indexer_initialized and self.indexer_keys.device != self.device:
+                self.indexer_keys = self.indexer_keys.to(self.device, non_blocking=True)
+        def reset(self) -> None:
+            super().reset()
+            if self.is_indexer_initialized:
+                self.indexer_keys.zero_()
+        def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
+            super().reorder_cache(beam_idx)
+            if self.is_indexer_initialized and self.indexer_keys.numel() > 0:
+                self.indexer_keys = self.indexer_keys.index_select(0, beam_idx.to(self.indexer_keys.device))
+        def crop(self, max_length: int) -> None:
+            super().crop(max_length)
+            if not self.is_indexer_initialized or self.indexer_keys.numel() == 0:
+                return
+            effective = max_length if max_length >= 0 else self.indexer_keys.shape[1] - abs(max_length)
+            if self.indexer_keys.shape[1] > effective:
+                self.indexer_keys = self.indexer_keys[:, :effective, :]
+        def batch_repeat_interleave(self, repeats: int) -> None:
+            super().batch_repeat_interleave(repeats)
+            if self.is_indexer_initialized and self.indexer_keys.numel() > 0:
+                self.indexer_keys = self.indexer_keys.repeat_interleave(repeats, dim=0)
+        def batch_select_indices(self, indices: torch.Tensor) -> None:
+            super().batch_select_indices(indices)
+            if self.is_indexer_initialized and self.indexer_keys.numel() > 0:
+                self.indexer_keys = self.indexer_keys[indices, ...]
     class StaticLayer(CacheLayerMixin):
         """
         A static cache layer that stores the key and value states as static tensors of shape `[batch_size, num_heads, max_cache_len), head_dim]`.
@@ Expand Down Expand Up / @@ -511,6 +587,73 @@ def reset(self): @@
             self.cumulative_length_int = 0
+    class StaticIndexedLayer(StaticLayer):
+        """
+        A `StaticLayer` with an additional statically-allocated indexer key cache for Dynamic Sparse
+        Attention (DSA) models (e.g. GLM MoE DSA, DeepSeek V32). This is the static, `torch.compile`-friendly
+        counterpart of `DynamicIndexedLayer`: the indexer key buffer is preallocated once and mutated in-place.
+        The main K/V cache is inherited from `StaticLayer` (`[batch_size, num_heads, max_cache_len, head_dim]`).
+        The indexer key cache stores a tensor of shape `[batch_size, max_cache_len, index_head_dim]` (3D, single-head).
+        """
+        def __init__(self, max_cache_len: int):
+            super().__init__(max_cache_len=max_cache_len)
+            self.indexer_keys: torch.Tensor | None = None
+            self.is_indexer_initialized: bool = False
+            # The indexer update runs independently of (and after) the main K/V `update` in the attention
+            # forward, so it tracks its own cumulative length rather than reusing `self.cumulative_length`.
+            self.indexer_cumulative_length = torch.tensor([0], dtype=int)
+        def lazy_initialization_indexer(self, indexer_key_states: torch.Tensor) -> None:
+            self.indexer_dtype, self.indexer_device = indexer_key_states.dtype, indexer_key_states.device
+            max_batch_size, _, index_head_dim = indexer_key_states.shape
+            self.indexer_keys = torch.zeros(
+                (max_batch_size, self.max_cache_len, index_head_dim),
+                dtype=self.indexer_dtype,
+                device=self.indexer_device,
+            )
+            self.indexer_cumulative_length = self.indexer_cumulative_length.to(self.indexer_device)
+            # Tag as static addresses for cudagraphs / compile, mirroring the main K/V buffers.
+            if not is_torchdynamo_compiling():
+                torch._dynamo.mark_static_address(self.indexer_keys)
+                torch._dynamo.mark_static_address(self.indexer_cumulative_length)
+            self.is_indexer_initialized = True
+        def update_indexer(self, indexer_key_states: torch.Tensor) -> torch.Tensor:
+            """
+            Update the indexer key cache in-place at the current positions, and return the full static buffer.
+            Args:
+                indexer_key_states (`torch.Tensor`): New indexer keys, shape `[batch_size, seq_len, index_head_dim]`.
+            Returns:
+                `torch.Tensor`: The full static indexer key cache, shape `[batch_size, max_cache_len, index_head_dim]`.
+                    Unfilled positions are masked out downstream by the indexer's attention mask, exactly as the
+                    main `StaticLayer` returns its full preallocated K/V.
+            """
+            if not self.is_indexer_initialized:
+                self.lazy_initialization_indexer(indexer_key_states)
+            seq_len = indexer_key_states.shape[1]
+            cache_position = torch.arange(seq_len, device=self.indexer_device) + self.indexer_cumulative_length
+            # In-place to preserve the static data pointer (required for cudagraphs).
+            self.indexer_cumulative_length.add_(seq_len)
+            try:
+                self.indexer_keys.index_copy_(1, cache_position, indexer_key_states)
+            except NotImplementedError:
+                # Fallback for devices like MPS where index_copy_ might not be supported.
+                self.indexer_keys[:, cache_position] = indexer_key_states
+            return self.indexer_keys
+        def reset(self) -> None:
+            super().reset()
+            if self.is_indexer_initialized:
+                self.indexer_keys.zero_()
+                self.indexer_cumulative_length.zero_()
     class QuantizedLayer(DynamicLayer):
         """
         A quantized layer similar to what is described in the [KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache paper](https://huggingface.co/papers/2402.02750).
@@ Expand Down Expand Up @@
             recurrent_states = self.layers[layer_idx].update_recurrent_state(recurrent_states, **kwargs)
             return recurrent_states
+        def update_indexer(self, indexer_key_states: torch.Tensor, layer_idx: int) -> torch.Tensor:
+            """
+            Updates the indexer key cache for layer `layer_idx`.
+            Parameters:
+                indexer_key_states (`torch.Tensor`):
+                    The new indexer key states to cache, shape `[batch_size, seq_len, index_head_dim]`.
+                layer_idx (`int`):
+                    The index of the layer to cache the states for.
+            Return:
+                `torch.Tensor`: The updated indexer key states (full cache).
+            """
+            if not hasattr(self.layers[layer_idx], "update_indexer"):
+                raise ValueError(
+                    f"Cannot call `update_indexer` on layer {layer_idx} which is a "
+                    f"{type(self.layers[layer_idx]).__name__}; it has no indexer key cache "
+                    f"(expected a `DynamicIndexedLayer` or `StaticIndexedLayer`)."
+                )
+            return self.layers[layer_idx].update_indexer(indexer_key_states)
         def early_initialization(
             self,
             batch_size: int,
@@ Expand Down Expand Up / @@ -1414,6 +1578,9 @@ def __init__( @@
                 # LinearAttention layers are static by essence - using `"moe"` as well is a trick, see the comment about it on DynamicCache
                 elif layer_type in ("mamba", "conv", "linear_attention", "moe"):
                     layer = LinearAttentionLayer()
+                elif layer_type == "deepseek_sparse_attention":
+                    # Static / compile-friendly indexed layer (preallocated indexer key cache).
+                    layer = StaticIndexedLayer(max_cache_len=max_cache_len)
                 else:
                     layer = StaticLayer(max_cache_len=max_cache_len)
                 layers.append(layer)
@@ Expand Down @@

src/transformers/configuration_utils.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -73,6 +73,7 @@ @@
         "dense",
         "hybrid",  # for layers that have both mamba and attention in zamba and zamba2
         "moe",  # for nemotron_h, which uses either attention, mamba or moe
+        "deepseek_sparse_attention",  # for models with DSA indexer (GLM MoE DSA, DeepSeek V32)
     )
@@ Expand Down @@

src/transformers/conversion_mapping.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -44,6 +44,7 @@ @@
         "afmoe": "qwen2_moe",
         "deepseek_v2": "qwen2_moe",
         "deepseek_v3": "qwen2_moe",
+        "deepseek_v32": "qwen2_moe",
         "dots1": "qwen2_moe",
         "ernie4_5_moe": "qwen2_moe",
         "glm4_moe": "qwen2_moe",
@@ Expand Down @@

src/transformers/masking_utils.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -1458,6 +1458,7 @@ def create_chunked_causal_mask( @@
         "chunked_attention": create_chunked_causal_mask,
         "compressed_sparse_attention": create_sliding_window_causal_mask,
         "heavily_compressed_attention": create_sliding_window_causal_mask,
+        "deepseek_sparse_attention": create_causal_mask,
     }
@@ Expand Down Expand Up / @@ -1514,10 +1515,14 @@ def create_masks_for_generate( @@
             "block_sequence_ids": block_sequence_ids,
         }
-        # If the attribute exist, we need several masks
+        # If the attribute exist, we need several masks - unless every layer shares the same type, in which
+        # case we return a single mask.
         if hasattr(effective_config, "layer_types"):
+            layer_patterns = set(effective_config.layer_types)
+            if len(layer_patterns) == 1:
+                return LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING[next(iter(layer_patterns))](**mask_kwargs)
             causal_masks = {}
-            for layer_pattern in set(effective_config.layer_types):
+            for layer_pattern in layer_patterns:
                 causal_masks[layer_pattern] = LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING[layer_pattern](**mask_kwargs)
             return causal_masks
         # In this case, all layers are sliding
@@ Expand Down @@

src/transformers/models/__init__.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -97,6 +97,7 @@ @@
         from .deepseek_v2 import *
         from .deepseek_v3 import *
         from .deepseek_v4 import *
+        from .deepseek_v32 import *
         from .deepseek_vl import *
         from .deepseek_vl_hybrid import *
         from .deformable_detr import *
@@ Expand Down @@

Navigation Menu

Uh oh!

Add deepseek 3.2 exp #41251

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!