update modeling and readme

README.md

@@ -1,3 +1,276 @@
----
-license: mit
----
+---
+license: mit
+base_model:
+- inclusionAI/Ling-mini-base-2.0
+---
+
+
+
+<p align="center">
+    <img src="https://mdn.alipayobjects.com/huamei_qa8qxu/afts/img/A*4QxcQrBlTiAAAAAAQXAAAAgAemJ7AQ/original" width="100"/>
+<p>
+
+<p align="center">🤗 <a href="https://huggingface.co/inclusionAI">Hugging Face</a>&nbsp&nbsp | &nbsp&nbsp🤖 <a href="https://modelscope.cn/organization/inclusionAI">ModelScope</a></p>
+
+
+## Introduction
+
+Today, we are excited to announce the open-sourcing of __Ling 2.0__ — a family of MoE-based large language models that combine __SOTA performance__ with __high efficiency__.
+The first released version, Ling-mini-2.0, is compact yet powerful. It has __16B total parameters__, but only __1.4B__ are activated per input token (non-embedding 789M). Trained on more than __20T tokens__ of high-quality data and enhanced through multi-stage supervised fine-tuning and reinforcement learning, Ling-mini-2.0 achieves remarkable improvements in complex reasoning and instruction following. With just 1.4B activated parameters, it still reaches the top-tier level of sub-10B dense LLMs and even matches or surpasses much larger MoE models.
+
+<p align="center"><img src="https://mdn.alipayobjects.com/huamei_fi95qp/afts/img/2NKZS5LVXzcAAAAASBAAAAgADkZ7AQFr/fmt.webp" /></p>
+
+### Strong General and Professional Reasoning
+
+We evaluated Ling-mini-2.0 on challenging general reasoning tasks in coding (LiveCodeBench, CodeForces) and mathematics (AIME 2025, HMMT 2025), as well as knowledge-intensive reasoning tasks across multiple domains (MMLU-Pro, Humanity's Last Exam). Compared with sub-10B dense models (e.g., Qwen3-4B-instruct-2507, Qwen3-8B-nothinking) and larger-scale MoE models (Ernie-4.5-21B-A3B-PT, GPT-OSS-20B/low), Ling-mini-2.0 demonstrated outstanding overall reasoning capabilities.
+
+### 7× Equivalent Dense Performance Leverage
+
+Guided by [Ling Scaling Laws](https://arxiv.org/abs/2507.17702), Ling 2.0 adopts a __1/32 activation ratio__ MoE architecture, with empirically optimized design choices in expert granularity, shared expert ratio, attention ratio, aux-loss free + sigmoid routing strategy, MTP loss, QK-Norm, half RoPE, and more. This enables small-activation MoE models to achieve over __7× equivalent dense performance__. In other words, __Ling-mini-2.0 with only 1.4B activated parameters (non-embedding 789M) can deliver performance equivalent to a 7–8B dense model__.
+
+### High-speed Generation at 300+ token/s
+
+<p align="center"><img src="https://mdn.alipayobjects.com/huamei_fi95qp/afts/img/bnxIRaK9tzcAAAAAgSAAAAgADkZ7AQFr/original" /></p>
+
+The highly sparse small-activation MoE architecture also delivers significant training and inference efficiency. In simple QA scenarios (within 2000 tokens), __Ling-mini-2.0 generates at 300+ token/s (on H20 deployment)__ — more than __2× faster__ than an 8B dense model. Ling-mini-2.0 is able to handle __128K context length__ with YaRN, as sequence length increases, the relative speedup can reach __over 7×__.
+
+<p align="center"><img src="https://raw.githubusercontent.com/inclusionAI/Ling-V2/refs/heads/main/figures/needle_in_a_haystack.webp" /></p>
+
+### Open-sourced FP8 Efficient Training Solution
+
+Ling 2.0 employs __FP8 mixed-precision training__ throughout. Compared with BF16, experiments with over 1T training tokens show nearly identical loss curves and downstream benchmark performance. To support the community in efficient continued pretraining and fine-tuning under limited compute, we are also open-sourcing our __FP8 training solution__. Based on tile/blockwise FP8 scaling, it further introduces FP8 optimizer, FP8 on-demand transpose weight, and FP8 padding routing map for extreme memory optimization. On 8/16/32 80G GPUs, compared with LLaMA 3.1 8B and Qwen3 8B, __Ling-mini-2.0 achieved 30–60% throughput gains with MTP enabled, and 90–120% throughput gains with MTP disabled__.
+
+### A More Open Opensource Strategy
+
+We believe Ling-mini-2.0 is an ideal starting point for MoE research. For the first time at this scale, it integrates 1/32 sparsity, MTP layers, and FP8 training — achieving both strong effectiveness and efficient training/inference performance, making it a prime candidate for the small-size LLM segment.
+To further foster community research, in addition to releasing the post-trained version, we are also open-sourcing __five pretraining checkpoints__: the pre-finetuning Ling-mini-2.0-base, along with four base models trained on 5T, 10T, 15T, and 20T tokens, enabling deeper research and broader applications.
+
+
+## Model Downloads
+
+You can download the following table to see the various stage of Ling-mini-2.0 models(1.43B activated of 16.26B total params). If you are located in mainland China, we also provide the model on ModelScope.cn to speed up the download process.
+
+<center>
+
+|       **Model**        | **Context Length** |                                                                             **Download**                                                                             |
+|:----------------------:| :----------------: |:--------------------------------------------------------------------------------------------------------------------------------------------------------------------:|
+|   Ling-mini-base-2.0   |        32K -> 128K (YaRN)         |     [🤗 HuggingFace](https://huggingface.co/inclusionAI/Ling-mini-base-2.0) <br>[🤖 ModelScope](https://www.modelscope.cn/models/inclusionAI/Ling-mini-base-2.0)     |
+| Ling-mini-base-2.0-5T  |        4K         |  [🤗 HuggingFace](https://huggingface.co/inclusionAI/Ling-mini-base-2.0-5T) <br>[🤖 ModelScope](https://www.modelscope.cn/models/inclusionAI/Ling-mini-base-2.0-5T)  |
+| Ling-mini-base-2.0-10T |        4K         | [🤗 HuggingFace](https://huggingface.co/inclusionAI/Ling-mini-base-2.0-10T) <br>[🤖 ModelScope](https://www.modelscope.cn/models/inclusionAI/Ling-mini-base-2.0-10T) |
+| Ling-mini-base-2.0-15T |        4K         | [🤗 HuggingFace](https://huggingface.co/inclusionAI/Ling-mini-base-2.0-15T) <br>[🤖 ModelScope](https://www.modelscope.cn/models/inclusionAI/Ling-mini-base-2.0-15T) |
+| Ling-mini-base-2.0-20T |        4K         | [🤗 HuggingFace](https://huggingface.co/inclusionAI/Ling-mini-base-2.0-20T) <br>[🤖 ModelScope](https://www.modelscope.cn/models/inclusionAI/Ling-mini-base-2.0-20T) |
+|     Ling-mini-2.0      |        32K -> 128K (YaRN)        |          [🤗 HuggingFace](https://huggingface.co/inclusionAI/Ling-mini-2.0) <br>[🤖 ModelScope](https://www.modelscope.cn/models/inclusionAI/Ling-mini-2.0)          |
+
+</center>
+
+Note: If you are interested in previous version, please visit the past model collections in [Huggingface](https://huggingface.co/inclusionAI) or [ModelScope](https://modelscope.cn/organization/inclusionAI).
+
+
+## Quickstart
+
+### Convert to safetensors
+
+Models with safetensors format can be downloaded from [HuggingFace](https://huggingface.co/inclusionAI) or [ModelScope](https://modelscope.cn/organization/inclusionAI).
+If you want to train your model and eval it, you can convert from dcp produced by training.
+```shell
+python tools/convert_dcp_to_safe_tensors.py --checkpoint-path ${DCP_PATH} --target-path ${SAFETENSORS_PATH}
+```
+
+Currently, BF16 and FP8 formats are supported, you can use convert parameter to handle it:
+- `--force-bf16` for BF16 format.
+- `--force-fp8` for FP8 format.
+
+### 🤗 Hugging Face Transformers
+
+Here is a code snippet to show you how to use the chat model with `transformers`:
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_name = "inclusionAI/Ling-mini-2.0"
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    dtype="auto",
+    device_map="auto",
+    trust_remote_code=True,
+)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+prompt = "Give me a short introduction to large language models."
+messages = [
+    {"role": "system", "content": "You are Ling, an assistant created by inclusionAI"},
+    {"role": "user", "content": prompt}
+]
+text = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True
+)
+model_inputs = tokenizer([text], return_tensors="pt", return_token_type_ids=False).to(model.device)
+
+generated_ids = model.generate(
+    **model_inputs,
+    max_new_tokens=512
+)
+generated_ids = [
+    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+]
+
+response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+```
+
+### 🤖 ModelScope
+
+If you're in mainland China, we strongly recommend you to use our model from 🤖 <a href="https://modelscope.cn/organization/inclusionAI">ModelScope</a>.
+
+## Deployment
+
+### vLLM
+
+vLLM supports offline batched inference or launching an OpenAI-Compatible API Service for online inference.
+
+#### Environment Preparation
+
+Since the Pull Request (PR) has not been submitted to the vLLM community at this stage, please prepare the environment by following the steps below:
+
+```bash
+git clone -b v0.10.0 https://github.com/vllm-project/vllm.git
+cd vllm
+git apply Ling-V2/inference/vllm/bailing_moe_v2.patch
+pip install -e .
+```
+
+#### Offline Inference:
+
+```bash
+from transformers import AutoTokenizer
+from vllm import LLM, SamplingParams
+
+tokenizer = AutoTokenizer.from_pretrained("inclusionAI/Ling-mini-2.0")
+
+sampling_params = SamplingParams(temperature=0.7, top_p=0.8, repetition_penalty=1.05, max_tokens=16384)
+
+llm = LLM(model="inclusionAI/Ling-mini-2.0", dtype='bfloat16')
+prompt = "Give me a short introduction to large language models."
+messages = [
+    {"role": "system", "content": "You are Ling, an assistant created by inclusionAI"},
+    {"role": "user", "content": prompt}
+]
+
+text = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True
+)
+outputs = llm.generate([text], sampling_params)
+
+```
+
+#### Online Inference:
+
+```bash
+vllm serve inclusionAI/Ling-mini-2.0 \
+              --tensor-parallel-size 2 \
+              --pipeline-parallel-size 1 \
+              --use-v2-block-manager \
+              --gpu-memory-utilization 0.90
+```
+
+To handle long context in vLLM using YaRN, we need to follow these two steps:
+1. Add a `rope_scaling` field to the model's `config.json` file, for example:
+```json
+{
+  ...,
+  "rope_scaling": {
+    "factor": 4.0,
+    "original_max_position_embeddings": 32768,
+    "type": "yarn"
+  }
+}
+```
+2. Use an additional parameter `--max-model-len` to specify the desired maximum context length when starting the vLLM service.
+
+For detailed guidance, please refer to the vLLM [`instructions`](https://docs.vllm.ai/en/latest/).
+
+
+### SGLang
+
+#### Environment Preparation
+
+We will later submit our model to SGLang official release, now we can prepare the environment following steps:
+```shell
+pip3 install sglang==0.5.2rc0 sgl-kernel==0.3.7.post1
+```
+You can use docker image as well:
+```shell
+docker pull lmsysorg/sglang:v0.5.2rc0-cu126
+```
+Then you should apply patch to sglang installation:
+```shell
+# patch command is needed, run `yum install -y patch` if needed
+patch -d `python -c 'import sglang;import os; print(os.path.dirname(sglang.__file__))'` -p3 < inference/sglang/bailing_moe_v2.patch
+```
+
+#### Run Inference
+
+BF16 and FP8 models are supported by SGLang now, it depends on the dtype of the model in ${MODEL_PATH}. They both share the same command in the following:  
+
+- Start server:
+```shell
+python -m sglang.launch_server \
+    --model-path $MODLE_PATH \
+    --host 0.0.0.0 --port $PORT \
+    --trust-remote-code \
+    --attention-backend fa3
+```
+MTP is supported for base model, and not yet for chat model. You can add parameter `--speculative-algorithm NEXTN`
+to start command.
+
+- Client:
+```shell
+curl -s http://localhost:${PORT}/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{"model": "auto", "messages": [{"role": "user", "content": "What is the capital of France?"}]}'
+"""
+```
+More usage can be found [here](https://docs.sglang.ai/basic_usage/send_request.html)
+
+## Training
+
+We also provide a complete and efficient training framework that covers both pre-training and finetune. Based on this framework, continue training can be performed on the Ling-mini-2.0 checkpoint. With our training framework, the training throughput of the Ling-mini-2.0 model is significantly better than that of the existing Dense 8B model (Qwen3-8B, Llama3-8B).
+
+### Pre-training
+
+[Pretraining demo](https://github.com/inclusionAI/Ling-V2/blob/main/docs/gpu_based_training.md) to Continue pretraining Ling models.
+
+#### Performance Benchmark
+
+The table below shows the pre-training performance of several models, measured in **tokens per second** on 8, 16, and 32 80G GPUs. Ling-mini-2.0 achieves significantly higher training efficiency compared to the baseline, making it easier and more cost-effective to continue pre-training with our [demo scripts](https://github.com/inclusionAI/Ling-V2/blob/main/docs/gpu_based_training.md).
+
+<center>
+
+|        **Model**        | **8 x 80G GPUs (GBS=128)** | **16 x 80G GPUs (GBS=256)** | **32 x 80G GPUs (GBS=512)** |
+|:-----------------------:| :--------------------: | :---------------------: | :---------------------: |
+| LLaMA 3.1 8B (baseline) |         81222          |         161319          |         321403          |
+|        Qwen3 8B         |    55775 (-31.33%)     |    109799 (-31.94%)     |    219943 (-31.57%)     |
+|      Ling-mini-2.0      |    109532 (+34.86%)    |    221585 (+37.36%)     |    448726 (+39.61%)     |
+|  Ling-mini-2.0 w/o MTP  |    128298 (+57.96%)    |    307264 (+90.47%)     |    611466 (+90.25%)     |
+
+</center>
+
+### Finetuning
+
+We recommend you to use [Llama-Factory](https://github.com/hiyouga/LLaMA-Factory) to [finetune Ling](https://github.com/inclusionAI/Ling-V2/blob/main/docs/llamafactory_finetuning.md). In addition to that, you can also use [Megatron for finetuning](https://github.com/inclusionAI/Ling-V2/blob/main/docs/megatron_sft_training.md).
+
+## License
+
+This code repository is licensed under [the MIT License](https://github.com/inclusionAI/Ling-V2/blob/master/LICENCE).
+
+## Citation
+
+If you find our work helpful, feel free to give us a cite.
+
+```
+
+```

config.json

@@ -47,5 +47,5 @@
     "use_qk_norm": true,
     "score_function": "sigmoid",
     "moe_shared_expert_intermediate_size": 512,
-    "num_nextn_predict_layers": 1
+    "num_nextn_predict_layers": 0
 }

configuration_bailing_moe_v2.py

@@ -5,8 +5,6 @@ from transformers.configuration_utils import PretrainedConfig
 
 class BailingMoeV2Config(PretrainedConfig):
 
-    model_type = "bailing_moe"
-
     def __init__(
         self,
         vocab_size=157184,
@@ -41,7 +39,7 @@ class BailingMoeV2Config(PretrainedConfig):
         head_dim=128,
         output_router_logits=False,
         use_qk_norm=True,
-        num_mtp_layers=0,
+        num_nextn_predict_layers=0,
         mtp_loss_scaling_factor=0,
         moe_router_enable_expert_bias=True,
         routed_scaling_factor=1.0,
@@ -60,7 +58,7 @@ class BailingMoeV2Config(PretrainedConfig):
         self.embedding_dropout = embedding_dropout
         self.attention_dropout = attention_dropout
         self.output_dropout = output_dropout
-        self.num_mtp_layers = num_mtp_layers
+        self.num_nextn_predict_layers = num_nextn_predict_layers
         self.mtp_loss_scaling_factor = mtp_loss_scaling_factor
         self.initializer_range = initializer_range
         self.max_position_embeddings = max_position_embeddings
@@ -84,4 +82,3 @@ class BailingMoeV2Config(PretrainedConfig):
         self.output_router_logits = output_router_logits
 
         super().__init__(pad_token_id=pad_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs)
-

modeling_bailing_moe_v2.py

@@ -25,9 +25,7 @@ from typing import List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
-import torch.utils.checkpoint
 from torch import nn
-from torch.nn import CrossEntropyLoss
 
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
@@ -234,7 +232,7 @@ def rotate_half(x):
 
 
 # Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
     """Applies Rotary Position Embedding to the query and key tensors.
 
     Args:
@@ -242,9 +240,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
         k (`torch.Tensor`): The key tensor.
         cos (`torch.Tensor`): The cosine part of the rotary embedding.
         sin (`torch.Tensor`): The sine part of the rotary embedding.
-        position_ids (`torch.Tensor`):
-            The position indices of the tokens corresponding to the query and key tensors. For example, this can be
-            used to pass offsetted position ids when working with a KV-cache.
         unsqueeze_dim (`int`, *optional*, defaults to 1):
             The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
             sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
@@ -403,7 +398,6 @@ class BailingMoeV2SparseMoeBlock(nn.Module):
         tokens_per_expert = cnts.sum(dim=0)
         idxs = topk_ids.view(-1).argsort()
         sorted_tokens = x[idxs // topk_ids.shape[1]]
-        sorted_tokens_shape = sorted_tokens.shape
         tokens_per_expert = tokens_per_expert.cpu().numpy()
         outputs = []
         start_idx = 0
@@ -495,10 +489,6 @@ class BailingMoeV2Attention(nn.Module):
         position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
         **kwargs,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-        if "padding_mask" in kwargs:
-            warnings.warn(
-                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
-            )
 
         bsz, q_len, _ = hidden_states.size()
 
@@ -516,7 +506,9 @@ class BailingMoeV2Attention(nn.Module):
             query_states = self.query_layernorm(query_states)
             key_states = self.key_layernorm(key_states)
 
-        kv_seq_len = key_states.shape[-2]
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
         if past_key_value is not None:
             if self.layer_idx is None:
                 raise ValueError(
@@ -524,12 +516,7 @@ class BailingMoeV2Attention(nn.Module):
                     "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
                     "with a layer index."
                 )
-            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
-        cos, sin = position_embeddings
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
-
-        if past_key_value is not None:
-            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
+            cache_kwargs = {"sin": sin, "cos": cos}
             key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
 
         key_states = repeat_kv(key_states, self.num_key_value_groups)
@@ -537,6 +524,7 @@ class BailingMoeV2Attention(nn.Module):
 
         attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
 
+        kv_seq_len = key_states.shape[-2]
         if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
             raise ValueError(
                 f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
@@ -601,14 +589,6 @@ class BailingMoeV2FlashAttention2(BailingMoeV2Attention):
         **kwargs,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
         # BailingMoeV2FlashAttention2 attention does not support output_attentions
-        if "padding_mask" in kwargs:
-            warnings.warn(
-                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
-            )
-
-            # overwrite attention_mask with padding_mask
-            attention_mask = kwargs.pop("padding_mask")
-
         output_attentions = False
 
         bsz, q_len, _ = hidden_states.size()
@@ -631,14 +611,11 @@ class BailingMoeV2FlashAttention2(BailingMoeV2Attention):
             query_states = self.query_layernorm(query_states)
             key_states = self.key_layernorm(key_states)
 
-        kv_seq_len = key_states.shape[-2]
-        if past_key_value is not None:
-            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
         cos, sin = position_embeddings
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
         if past_key_value is not None:
-            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
+            cache_kwargs = {"sin": sin, "cos": cos}
             key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
 
         # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
@@ -840,21 +817,18 @@ class BailingMoeV2SdpaAttention(BailingMoeV2Attention):
             query_states = self.query_layernorm(query_states)
             key_states = self.key_layernorm(key_states)
 
-        kv_seq_len = key_states.shape[-2]
-        if past_key_value is not None:
-            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
         cos, sin = position_embeddings
-
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
         if past_key_value is not None:
-            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
+            cache_kwargs = {"sin": sin, "cos": cos}
             key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
 
         key_states = repeat_kv(key_states, self.num_key_value_groups)
         value_states = repeat_kv(value_states, self.num_key_value_groups)
 
         if attention_mask is not None:
+            kv_seq_len = key_states.shape[-2]
             if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
                 raise ValueError(
                     f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
@@ -1012,10 +986,6 @@ class BailingMoeV2DecoderLayer(nn.Module):
                 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
                 (see `past_key_values`).
         """
-        if "padding_mask" in kwargs:
-            warnings.warn(
-                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
-            )
         residual = hidden_states
 
         hidden_states = self.input_layernorm(hidden_states)
@@ -1185,11 +1155,11 @@ class BailingMoeV2Model(BailingMoeV2PreTrainedModel):
         super().__init__(config)
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
-        self.num_mtp_layers = config.num_mtp_layers
+        self.num_nextn_predict_layers = config.num_nextn_predict_layers
 
         self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
         self.layers = []
-        for layer_idx in range(config.num_hidden_layers + config.num_mtp_layers):
+        for layer_idx in range(config.num_hidden_layers + config.num_nextn_predict_layers):
             layer_cls = BailingMoeV2DecoderLayer if layer_idx < config.num_hidden_layers else BailingMoeV2MTPLayer
             self.layers.append(layer_cls(config, layer_idx))
 
@@ -1252,23 +1222,20 @@ class BailingMoeV2Model(BailingMoeV2PreTrainedModel):
                 )
                 use_cache = False
 
-        past_key_values_length = 0
-        if use_cache:
-            use_legacy_cache = not isinstance(past_key_values, Cache)
-            if use_legacy_cache:
-                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
-            past_key_values_length = past_key_values.get_usable_length(seq_length)
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+
+        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
 
         if position_ids is None:
-            device = input_ids.device if input_ids is not None else inputs_embeds.device
             position_ids = torch.arange(
-                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
             )
             position_ids = position_ids.unsqueeze(0)
 
-        if inputs_embeds is None:
-            inputs_embeds = self.word_embeddings(input_ids)
-
         if self._use_flash_attention_2:
             # 2d mask is passed through the layers
             attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
@@ -1279,12 +1246,12 @@ class BailingMoeV2Model(BailingMoeV2PreTrainedModel):
                 attention_mask,
                 (batch_size, seq_length),
                 inputs_embeds,
-                past_key_values_length,
+                past_seen_tokens,
             )
         else:
             # 4d mask is passed through the layers
             attention_mask = _prepare_4d_causal_attention_mask(
-                attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
+                attention_mask, (batch_size, seq_length), inputs_embeds, past_seen_tokens
             )
 
         # embed positions
@@ -1298,8 +1265,8 @@ class BailingMoeV2Model(BailingMoeV2PreTrainedModel):
         all_self_attns = () if output_attentions else None
         all_router_logits = () if output_router_logits else None
         next_decoder_cache = None
-        layers = self.layers[: -self.num_mtp_layers] if self.num_mtp_layers > 0 else self.layers
-        mtp_layers = self.layers[-self.num_mtp_layers :] if self.num_mtp_layers > 0 else None
+        layers = self.layers[: -self.num_nextn_predict_layers] if self.num_nextn_predict_layers > 0 else self.layers
+        mtp_layers = self.layers[-self.num_nextn_predict_layers :] if self.num_nextn_predict_layers > 0 else None
 
         for decoder_layer in layers:
             if output_hidden_states:
@@ -1397,7 +1364,7 @@ class BailingMoeV2Model(BailingMoeV2PreTrainedModel):
 
         next_cache = None
         if use_cache:
-            next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
+            next_cache = next_decoder_cache
         if not return_dict:
             return tuple(
                 v
@@ -1422,7 +1389,7 @@ class BailingMoeV2ForCausalLM(BailingMoeV2PreTrainedModel, GenerationMixin):
         self.model = BailingMoeV2Model(config)
         self.vocab_size = config.vocab_size
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        self.num_mtp_layers = config.num_mtp_layers
+        self.num_nextn_predict_layers = config.num_nextn_predict_layers
         self.mtp_loss_scaling_factor = config.mtp_loss_scaling_factor
 
         # Initialize weights and apply final processing
@@ -1519,29 +1486,24 @@ class BailingMoeV2ForCausalLM(BailingMoeV2PreTrainedModel, GenerationMixin):
         logits = logits.float()
 
         if labels is not None:
-            # Shift so that tokens < n predict n
-            # Flatten the tokens
-            loss_fct = CrossEntropyLoss()
-            logits = logits.view(-1, self.config.vocab_size)
-            # Enable model parallelism
-            loss = loss_fct(logits, labels.to(logits.device).view(-1))
+            loss = self.loss_function(logits, labels, self.config.vocab_size, **kwargs)
 
         all_mtp_logits = None
-        if self.num_mtp_layers > 0:
+        if self.num_nextn_predict_layers > 0:
             mtp_hidden_states = outputs.mtp_hidden_states
-            for i in range(self.num_mtp_layers):
+            shift_labels_mtp = None
+            for i in range(self.num_nextn_predict_layers):
                 mtp_hidden_states = mtp_hidden_states[i]
                 mtp_logits = self.lm_head(mtp_hidden_states).float()
                 if all_mtp_logits is None:
                     all_mtp_logits = []
                 all_mtp_logits.append(mtp_logits)
                 if labels is not None:
-                    labels, _ = roll_tensor(labels, shifts=-1, dims=-1, fill_value=-100)
-                    # Flatten the tokens
-                    loss_fct = CrossEntropyLoss()
+                    if shift_labels_mtp is None:
+                        shift_labels_mtp = labels.clone()
+                    shift_labels_mtp, _ = roll_tensor(shift_labels_mtp, shifts=-1, dims=-1, fill_value=-100)
                     mtp_logits_ = mtp_logits.view(-1, self.config.vocab_size)
-                    # Enable model parallelism
-                    mtp_loss = loss_fct(mtp_logits_, labels.to(mtp_logits_.device).view(-1))
+                    mtp_loss = self.loss_function(mtp_logits_, shift_labels_mtp.to(mtp_logits_.device).view(-1), self.config.vocab_size, **kwargs)
                     if loss is not None:
                         loss += self.mtp_loss_scaling_factor * mtp_loss
                     else:
@@ -1569,70 +1531,3 @@ class BailingMoeV2ForCausalLM(BailingMoeV2PreTrainedModel, GenerationMixin):
             router_logits=outputs.router_logits,
         )
 
-    def prepare_inputs_for_generation(
-        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, token_type_ids=None, **kwargs
-    ):
-        if past_key_values is not None:
-            if isinstance(past_key_values, Cache):
-                cache_length = past_key_values.get_seq_length()
-                past_length = past_key_values.seen_tokens
-                max_cache_length = (
-                    past_key_values.get_max_length()
-                    if hasattr(past_key_values, "get_max_length")
-                    else past_key_values.get_max_cache_shape()
-                )
-            else:
-                cache_length = past_length = past_key_values[0][0].shape[2]
-                max_cache_length = None
-
-            # Keep only the unprocessed tokens:
-            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
-            # some of the inputs are exclusivelly passed as part of the cache (e.g. when passing input_embeds as input)
-            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
-                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
-            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
-            # input_ids based on the past_length.
-            elif past_length < input_ids.shape[1]:
-                input_ids = input_ids[:, past_length:]
-            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
-
-            # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
-            if (
-                max_cache_length is not None
-                and attention_mask is not None
-                and cache_length + input_ids.shape[1] > max_cache_length
-            ):
-                attention_mask = attention_mask[:, -max_cache_length:]
-
-        position_ids = kwargs.get("position_ids", None)
-        if attention_mask is not None and position_ids is None:
-            # create position_ids on the fly for batch generation
-            position_ids = attention_mask.long().cumsum(-1) - 1
-            position_ids.masked_fill_(attention_mask == 0, 1)
-            if past_key_values:
-                position_ids = position_ids[:, -input_ids.shape[1] :]
-
-        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
-        if inputs_embeds is not None and past_key_values is None:
-            model_inputs = {"inputs_embeds": inputs_embeds}
-        else:
-            model_inputs = {"input_ids": input_ids}
-
-        model_inputs.update(
-            {
-                "position_ids": position_ids,
-                "past_key_values": past_key_values,
-                "use_cache": kwargs.get("use_cache"),
-                "attention_mask": attention_mask,
-            }
-        )
-        return model_inputs
-
-    @staticmethod
-    def _reorder_cache(past_key_values, beam_idx):
-        reordered_past = ()
-        for layer_past in past_key_values:
-            reordered_past += (
-                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
-            )
-        return reordered_past