Duplicate from meituan-longcat/LongCat-Flash-Lite

Co-authored-by: LongCat <LongCat0830@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+tech_report.pdf filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Meituan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in 
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,260 @@
+---
+license: mit
+library_name: LongCat-Flash-Lite
+pipeline_tag: text-generation
+tags:
+- transformers
+---
+
+# LongCat-Flash-Lite
+
+<div align="center">
+  <img src="https://raw.githubusercontent.com/meituan-longcat/LongCat-Flash-Chat/main/figures/longcat_logo.svg" 
+       width="300" 
+       alt="LongCat Logo"/>
+</div>
+
+<hr>
+
+
+<div align="center" style="line-height: 1;">
+
+  <a href="https://huggingface.co/meituan-longcat" target="_blank" style="margin: 2px;">
+    <img alt="Hugging Face" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-LongCat-ffc107?color=ffc107&logoColor=white" style="display: inline-block; vertical-align: middle;"/>
+  </a>
+</div>
+
+<div align="center" style="line-height: 1;">
+  <a href="https://github.com/meituan-longcat/LongCat-Flash-Chat/blob/main/figures/wechat_official_accounts.png" target="_blank" style="margin: 2px;">
+    <img alt="Wechat" src="https://img.shields.io/badge/WeChat-LongCat-brightgreen?logo=wechat&logoColor=white" style="display: inline-block; vertical-align: middle;"/>
+  </a>
+  <a href="https://x.com/Meituan_LongCat" target="_blank" style="margin: 2px;">
+    <img alt="Twitter Follow" src="https://img.shields.io/badge/Twitter-LongCat-white?logo=x&logoColor=white" style="display: inline-block; vertical-align: middle;"/>
+  </a>
+</div>
+
+<div align="center" style="line-height: 1;">
+  <a href="https://huggingface.co/meituan-longcat/LongCat-Flash-Chat/blob/main/LICENSE" style="margin: 2px;">
+    <img alt="License" src="https://img.shields.io/badge/License-MIT-f5de53?&color=f5de53" style="display: inline-block; vertical-align: middle;"/>
+  </a>
+</div>
+
+<p align="center">
+  <a href="https://arxiv.org/abs/2601.21204"><b>Tech Report</b>&nbsp;📄</a>
+</p>
+
+## Model Introduction
+We introduce LongCat-Flash-Lite, a non-thinking 68.5B parameter Mixture-of-Experts (MoE) model with approximately 3B activated parameters, supporting a 256k context length through the YaRN method. Building upon the LongCat-Flash architecture, LongCat-Flash-Lite distinguishes itself through the integration of an **N-gram embedding** table designed to enhance both model performance and inference speed. Despite allocating over 30B parameters to embeddings, LongCat-Flash-Lite not only outperforms parameter-equivalent MoE baselines but also demonstrates exceptional competitiveness against existing models of comparable scale, particularly in the agentic and coding domains.
+
+
+
+### Key Features
+
+#### 🌟  Superior Scaling Efficiency: A Better Alternative to MoE
+Through comprehensive scaling experiments across diverse scenarios, we identify specific regimes where embedding scaling achieves a superior Pareto frontier compared to increasing the number of experts, thereby offering a highly efficient alternative for model scaling. We further delineate a comprehensive set of architectural factors that determine embedding scaling efficacy, encompassing integration timing, parameter budgeting, hash collision mitigation, hyperparameter configuration, and embedding initialization, alongside the impacts of model width and depth. 
+
+#### 🌟 Superior Inference Efficiency with Specialized System Optimization
+In contrast to FFN-based experts, the N-gram embedding table inherently mitigates I/O bottlenecks within MoE layers, yielding substantial improvements in inference latency. Furthermore, we introduce a specialized N-gram Cache and develop synchronized kernels, which collectively and significantly boost inference efficiency.
+
+#### 🌟 Strong Agentic and Coding Performance
+LongCat-Flash-Lite demonstrates robust capabilities in agentic tool use and coding proficiency that are highly competitive relative to its model scale. 
+
+Please refer to our [technical report](https://arxiv.org/abs/2601.21204) for details!
+
+## Evaluation Results
+
+| Benchmark | Kimi-Linear-48B-A3B | Qwen3-Next-80B-A3B-Instruct | Gemini 2.5 Flash-Lite | LongCat-Flash-Lite |
+|----------|---------------------|----------------------------|----------------------|---------|
+| **Architecture** | MoE | MoE | - | MoE + NE |
+| **# Total Params** | 48B | 80B | - | 68.5B |
+| **# Activated Params** | 3B | 3B | - | 2.9B~4.5B |
+| **Agentic Tool Use** | | | | |
+| Tau2-Airline(avg@8) | 44.00 | 45.5* | 35.00 | 58.00 |
+| Tau2-Retail(avg@8) | 18.86 | 57.3* | 37.50 | 73.10 |
+| Tau2-Telecom(avg@8) | 15.68 | 13.2* | 21.93 | 72.80 |
+| **Agentic Coding** | | | | |
+| SWE-Bench(acc) | 32.80 | 37.60 | 41.3* | 54.40 |
+| TerminalBench(acc) | 20.00 | 15.19 | 20.00 | 33.75 |
+| SWE-Bench Multiligual | 37.20 | 31.30 | - | 38.10 |
+|PRDBench | - | 15.36 | - |  39.63 |
+| **General Domains** | | | | |
+| GPQA-Diamond(avg@16) | 69.89 | 74.33 | 70.20* | 66.78 |
+| MMLU(acc) | 79.91 | 89.28 | 84.68 | 85.52 |
+| MMLU-Pro(acc) | 67.22 | 82.93 | 78.95 | 78.29 |
+| CEval(acc) | 78.48 | 90.91 | 75.16 | 86.55 |
+| CMMLU(acc) | 76.26 | 86.50 | 72.06 | 82.48 |
+| **Mathematical Reasoning** | | | | |
+| MATH500(acc) | 94.20 | 98.00 | 95.20 | 96.80 |
+| AIME24(avg@32) | 70.52 | 81.35 | 63.33 | 72.19 |
+| AIME25(avg@32) | 59.58 | 68.44 | 50.1* | 63.23 |
+
+> **Note:** Values marked with * are sourced from public reports. NE is an abbreviation of N-gram Embedding.
+
+## Quick Start
+To use LongCat-Flash-Lite with transformers, we need at least 2 GPUs (80GB VRAM each, e.g., H100/A100 80GB), and we recommend the following environment:
+
+* `python` >= 3.10
+* `torch` >= 2.6
+* `transformers` >= 4.57.6
+* `accelerate` >= 1.10.0
+
+```shell
+pip install -U transformers==4.57.6 accelerate==1.10.0
+```
+
+Basic Usage Example:
+```py
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_name = "meituan-longcat/LongCat-Flash-Lite"
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype="auto",
+    device_map="auto",
+    trust_remote_code=True
+)
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+
+messages = [
+    {"role": "system", "content": "You are a helpful assistant."},
+    {"role": "user", "content": "Give me a brief introduction to large language models."}
+]
+input_ids = tokenizer.apply_chat_template(
+    messages, 
+    add_generation_prompt=True, 
+    return_tensors="pt"
+).to(model.device)
+generated_ids = model.generate(inputs=input_ids, max_new_tokens=256)
+output_ids = generated_ids[0][len(input_ids[0]):].tolist()
+response = tokenizer.decode(output_ids, skip_special_tokens=True).strip("\n")
+print(response)
+```
+
+Tool Calling Example:
+```py
+tools = [
+    {
+        "type": "function",
+        "function": {
+            "name": "func_add",
+            "description": "Calculate the sum of two numbers",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "x1": {"type": "number", "description": "The first addend"},
+                    "x2": {"type": "number", "description": "The second addend"}
+                },
+                "required": ["x1", "x2"]
+            }
+        }
+    }
+]
+messages = [
+    {"role": "system", "content": "You are a helpful assistant."},
+    {"role": "user", "content": "Please tell me what is $$125679 + 234519$$?"},
+    {
+        "role": "assistant", 
+        "content": "I'll calculate the sum of 125679 and 234519 for you.", 
+        "tool_calls": [{"type": "function", "function": {"name": "func_add", "arguments": {"x1": 125679, "x2": 234519}}}]
+    },
+    {"role": "tool", "name": "func_add", "content": '{"ans": 360198}'}
+]
+
+input_ids = tokenizer.apply_chat_template(
+    messages, 
+    tools=tools,
+    add_generation_prompt=True, 
+    return_tensors="pt"
+).to(model.device)
+generated_ids = model.generate(inputs=input_ids, max_new_tokens=256)
+output_ids = generated_ids[0][len(input_ids[0]):].tolist()
+response = tokenizer.decode(output_ids, skip_special_tokens=True).strip("\n")
+print(response)
+```
+
+Response Parsing:
+
+```python
+from parse_model_response import parse_model_response
+
+response = tokenizer.decode(output_ids, skip_special_tokens=True).strip("\n")
+parsed_message = parse_model_response(response, tools)
+```
+See [`parse_model_response.py`](./parse_model_response.py) for detailed implementation and examples.
+
+<br>
+
+Recommended Sampling Setting:
+```shell
+{ "repetition_penalty": 1.06, "temperature": 0.7, "top_p": 0.95, "top_k": 4 }
+```
+
+## Deployment
+
+We have implemented basic adaptations in SGLang ([PR](https://github.com/sgl-project/sglang/pull/17838)) to support the deployment of LongCat-Flash-Lite.
+
+LongCat-Flash-Lite can be served on a single node (e.g., 8xH20-141G) using a combination of Tensor Parallelism and Expert Parallelism.
+
+Compile and update sgl-kernel first.
+
+```shell
+cd sgl-kernel
+python3 -m uv build --wheel --color=always --no-build-isolation \
+        -Ccmake.define.SGL_KERNEL_ENABLE_SM90A=1 \
+        -Ccmake.define.CMAKE_POLICY_VERSION_MINIMUM=3.5 \
+        -Cbuild-dir=build .
+pip3 install dist/sgl_kernel-0.3.21-cp310-abi3-linux_x86_64.whl --force-reinstall
+```
+Then launch the server.
+```py
+python3 -m sglang.launch_server \
+    --model meituan-longcat/LongCat-Flash-Lite \
+    --port 8080 \
+    --host 0.0.0.0 \
+    --mem-fraction-static 0.9 \
+    --max-running-requests 64 \
+    --trust-remote-code \
+    --skip-server-warmup \
+    --attention-backend flashinfer \
+    --ep 8 \
+    --tp 8 \
+    --disable-cuda-graph
+```
+
+## License Agreement
+
+This repository, including both the model weights and the source code, is released under the **MIT License**. 
+
+Any contributions to this repository are licensed under the MIT License, unless otherwise stated. This license does not grant any rights to use Meituan trademarks or patents. 
+
+For details, see the [LICENSE](./LICENSE) file. 
+
+## Usage Considerations 
+This model has not been specifically designed or comprehensively evaluated for every possible downstream application. 
+
+Developers should take into account the known limitations of large language models, including performance variations across different languages, and carefully assess accuracy, safety, and fairness before deploying the model in sensitive or high-risk scenarios. 
+It is the responsibility of developers and downstream users to understand and comply with all applicable laws and regulations relevant to their use case, including but not limited to data protection, privacy, and content safety requirements. 
+
+Nothing in this Model Card should be interpreted as altering or restricting the terms of the MIT License under which the model is released. 
+
+
+## Citation
+
+We kindly encourage citation of our work if you find it useful.
+
+```
+@misc{liu2026scalingembeddingsoutperformsscaling,
+      title={Scaling Embeddings Outperforms Scaling Experts in Language Models}, 
+      author={Hong Liu and Jiaqi Zhang and Chao Wang and Xing Hu and Linkun Lyu and Jiaqi Sun and Xurui Yang and Bo Wang and Fengcun Li and Yulei Qian and Lingtong Si and Yerui Sun and Rumei Li and Peng Pei and Yuchen Xie and Xunliang Cai},
+      year={2026},
+      eprint={2601.21204},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2601.21204}, 
+}
+```
+
+
+## Contact
+Please contact us at <a href="mailto:longcat-team@meituan.com">longcat-team@meituan.com</a> or open an issue if you have any questions.
+

config.json

@@ -0,0 +1,51 @@
+{
+  "architectures": [
+    "LongcatFlashNgramForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_longcat_ngram.LongcatFlashNgramConfig",
+    "AutoModel": "modeling_longcat_ngram.LongcatFlashNgramModel",
+    "AutoModelForCausalLM": "modeling_longcat_ngram.LongcatFlashNgramForCausalLM"
+  },
+  "vocab_size": 131072,
+  "hidden_size": 3072,
+  "ffn_hidden_size": 6144,
+  "expert_ffn_hidden_size": 1024,
+  "num_layers": 14,
+  "num_attention_heads": 32,
+  "kv_lora_rank": 512,
+  "q_lora_rank": 1536,
+  "qk_rope_head_dim": 64,
+  "v_head_dim": 128,
+  "qk_nope_head_dim": 128,
+  "mla_scale_q_lora": true,
+  "mla_scale_kv_lora": true,
+  "routed_scaling_factor": 6.0,
+  "n_routed_experts": 256,
+  "rms_norm_eps": 1e-5,
+  "use_cache": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "rope_theta": 5000000.0,
+  "max_position_embeddings": 327680,
+  "rope_scaling": {
+      "original_max_position_embeddings": 32768,
+      "rope_type": "yarn",
+      "factor": 10,
+      "beta_fast": 32,
+      "beta_slow": 1,
+      "mscale": 1,
+      "mscale_all_dim": 1
+  },
+  "zero_expert_num": 128,
+  "zero_expert_type": "identity",
+  "moe_topk": 12,
+  "ngram_vocab_size_ratio": 78,
+  "emb_neighbor_num": 4,
+  "emb_split_num": 4,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.57.6"
+}
+

configuration_longcat_ngram.py

@@ -0,0 +1,216 @@
+from transformers.models.longcat_flash import LongcatFlashConfig
+
+
+class LongcatFlashNgramConfig(LongcatFlashConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`LongcatFlashNgramModel`]. It is used to instantiate
+    a LongCat Flash model with N-gram enhanced embeddings according to the specified arguments, defining the model architecture. 
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 131072):
+            Vocabulary size of the LongCat Flash model. Defines the number of different tokens that can be represented by the
+            `input_ids` passed when calling [`LongcatFlashNgramModel`]
+        hidden_size (`int`, *optional*, defaults to 6144):
+            Dimension of the hidden representations.
+        num_hidden_layers (`int`, *optional*, defaults to 56):
+            Number of hidden layers in the Transformer decoder.
+        num_layers (`int`, *optional*, defaults to 28):
+            Number of layers, each with 2 sublayers.
+        num_attention_heads (`int`, *optional*, defaults to 64):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        num_key_value_heads (`int`, *optional*):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting from a multi-head checkpoint to a GQA checkpoint, each group key and value head should be
+            constructed by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
+            `num_attention_heads`.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 131072):
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-05):
+            The epsilon value used by the RMS normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        pad_token_id (`int`, *optional*):
+            Padding token id.
+        bos_token_id (`int`, *optional*, defaults to 1):
+            Beginning of stream token id.
+        eos_token_id (`int`, *optional*, defaults to 2):
+            End of stream token id.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether to tie input and output embeddings.
+        rope_theta (`float`, *optional*, defaults to 10000000.0):
+            The base period of the RoPE embeddings.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
+            strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
+            `{"type": strategy name, "factor": scaling factor}`.
+        attention_bias (`bool`, *optional*, defaults to `False`):
+            Whether to use a bias in the query, key, value and output projection layers during self-attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        ffn_hidden_size (`int`, *optional*, defaults to 12288):
+            Dimension of the MLP representations.
+        q_lora_rank (`int`, *optional*, defaults to 1536):
+            The rank of the query LoRA projection in MLA (Multi-head Latent Attention).
+        kv_lora_rank (`int`, *optional*, defaults to 512):
+            The rank of the key-value LoRA projection in MLA.
+        qk_nope_head_dim (`int`, *optional*, defaults to 128):
+            The dimension of the non-position encoding part of query/key heads.
+        qk_rope_head_dim (`int`, *optional*, defaults to 64):
+            The dimension of the RoPE part of query/key heads.
+        head_dim (`int`, *optional*, defaults to 64):
+            Standard dimension of qk heads, unused except for CI.
+        v_head_dim (`int`, *optional*, defaults to 128):
+            The dimension of value heads.
+        qk_head_dim (`int`, *optional*):
+            The total dimension of query/key heads. If not specified, set to `qk_nope_head_dim + qk_rope_head_dim`.
+        moe_topk (`int`, *optional*, defaults to 12):
+            Number of experts to route to for each token in the MoE layer.
+        n_routed_experts (`int`, *optional*, defaults to 512):
+            Number of routed experts in the MoE layer.
+        zero_expert_num (`int`, *optional*, defaults to 256):
+            Number of zero experts (identity function) to add to the expert pool.
+        expert_ffn_hidden_size (`int`, *optional*, defaults to 2048):
+            Hidden size of individual expert FFN layers.
+        routed_scaling_factor (`float`, *optional*, defaults to 6.0):
+            Scaling factor applied to the routing weights.
+        emb_neighbor_num (`int`, *optional*):
+            Maximum N-gram length for N-gram embeddings. This parameter determines the context window size for N-gram computation. Higher values capture 
+            longer-range lexical patterns but increase memory usage.
+        emb_split_num (`int`, *optional*):
+            Number of hash functions (or splits) to use for N-gram embeddings. Multiple hash functions help improve the quality of N-gram representations. 
+        ngram_vocab_size_ratio (`float`, *optional*):
+            Ratio multiplier for N-gram vocabulary size relative to the base vocabulary size. The N-gram vocabulary 
+            size is calculated as `vocab_size * ngram_vocab_size_ratio`. 
+
+    Example:
+    ```python
+    >>> from transformers import LongcatFlashNgramModel, LongcatFlashNgramConfig
+
+    >>> # Initializing a LongCat Flash N-gram style configuration
+    >>> configuration = LongcatFlashNgramConfig(
+    ...     emb_neighbor_num=3,
+    ...     emb_split_num=4,
+    ...     ngram_vocab_size_ratio=1.5
+    ... )
+
+    >>> # Initializing a model from the configuration
+    >>> model = LongcatFlashNgramModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "longcat_flash_ngram"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    base_model_tp_plan = {
+        "layers.*.self_attn.*.q_b_proj": "colwise",
+        "layers.*.self_attn.*.kv_b_proj": "colwise",
+        "layers.*.self_attn.*.o_proj": "rowwise",
+        "layers.*.mlps.*.gate_proj": "colwise",
+        "layers.*.mlps.*.up_proj": "colwise",
+        "layers.*.mlps.*.down_proj": "rowwise",
+        "layers.*.mlp.experts.*.gate_proj": "colwise",
+        "layers.*.mlp.experts.*.up_proj": "colwise",
+        "layers.*.mlp.experts.*.down_proj": "rowwise",
+    }
+
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size=131072,
+        hidden_size=6144,
+        num_hidden_layers=56,
+        num_layers=28,
+        num_attention_heads=64,
+        num_key_value_heads=None,
+        hidden_act="silu",
+        max_position_embeddings=131072,
+        initializer_range=0.02,
+        rms_norm_eps=1e-5,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=False,
+        rope_theta=10000000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        ffn_hidden_size=12288,
+        q_lora_rank=1536,
+        kv_lora_rank=512,
+        qk_nope_head_dim=128,
+        qk_rope_head_dim=64,
+        head_dim=64,
+        v_head_dim=128,
+        qk_head_dim=None,
+        moe_topk=12,
+        n_routed_experts=512,
+        zero_expert_num=256,
+        expert_ffn_hidden_size=2048,
+        routed_scaling_factor=6.0,
+        emb_neighbor_num=None,
+        emb_split_num=None,
+        ngram_vocab_size_ratio=None,
+        **kwargs,
+    ):
+        # N-gram embedding specific parameters
+        self.emb_neighbor_num = emb_neighbor_num
+        self.emb_split_num = emb_split_num
+        self.ngram_vocab_size_ratio = ngram_vocab_size_ratio
+
+        super().__init__(
+            vocab_size=vocab_size,
+            hidden_size=hidden_size,
+            num_hidden_layers=num_hidden_layers,
+            num_layers=num_layers,
+            num_attention_heads=num_attention_heads,
+            num_key_value_heads=num_key_value_heads,
+            hidden_act=hidden_act,
+            max_position_embeddings=max_position_embeddings,
+            initializer_range=initializer_range,
+            rms_norm_eps=rms_norm_eps,
+            use_cache=use_cache,
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            rope_theta=rope_theta,
+            rope_scaling=rope_scaling,
+            attention_bias=attention_bias,
+            attention_dropout=attention_dropout,
+            ffn_hidden_size=ffn_hidden_size,
+            q_lora_rank=q_lora_rank,
+            kv_lora_rank=kv_lora_rank,
+            qk_nope_head_dim=qk_nope_head_dim,
+            qk_rope_head_dim=qk_rope_head_dim,
+            head_dim=head_dim,
+            v_head_dim=v_head_dim,
+            qk_head_dim=qk_head_dim,
+            moe_topk=moe_topk,
+            n_routed_experts=n_routed_experts,
+            zero_expert_num=zero_expert_num,
+            expert_ffn_hidden_size=expert_ffn_hidden_size,
+            routed_scaling_factor=routed_scaling_factor,
+            **kwargs,
+        )
+
+
+__all__ = ["LongcatFlashNgramConfig"]

generation_config.json

@@ -0,0 +1,11 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 3,
+  "transformers_version": "4.57.6",
+  "repetition_penalty": 1.06,
+  "temperature": 0.7,
+  "top_p": 0.95,
+  "top_k": 4 
+}

model-00000-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:19412c9358bc484c8c9af94bb2e135a6939a39441488c6c042b16261bfe746ad
+size 5313604160

model-00001-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bff23594192d36756bdac96c7ee67dd641b8969f7d37c12d2a68a5e3ed9b8d87
+size 5313604768

model-00002-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3dcfc95fd273d4a8a57674916db168d977c10cddb87e0cc5b1fd5fc682ad121f
+size 5313604016

model-00003-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b465510bc231ab86c9f7403bd82e03a12b72b323cbe319547453b786577791d1
+size 5313604136

model-00004-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd0ca68e274a10a89944da1984cf583adb7b51a9f54782f4e72ba6e6305aa7c8
+size 5315503392

model-00005-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:edfcfed05027a48b2cab645e44b7dc8ec0544c1a64a0277bb0b76de703e85913
+size 5313604136

model-00006-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:29959ee6a47ed392e05966031f46a589d9c41cfce369ee0be0f32a83246bc3ba
+size 5313604136

model-00007-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba85ee575bdafb34270bbab3f7256c87178c3a2c05a5e637bb9ee83e565fdf2b
+size 5313604760

model-00008-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0533404a0b7ee44b0ade17bf9b9c00969e37cf864ff5d5f7818b030c702abb06
+size 5313604216

model-00009-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b65ca6876f65b690af110232c4d460c912beab21c195c4e0170e9f01ca1d8bb2
+size 5315105048

model-00010-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8d1032fedf9f1029544ac164405278774fd2f9e3b8399d8466b797d08f756fc6
+size 5315104024

model-00011-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e13ef0f626c36a424cbed2e4181b853384de4cb8ad0660b65084c7408cfc729
+size 5315103008

model-00012-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d1c91f60ac661af3c195337d4ca771e274947ef8e7f258607b79e0e19a02c30d
+size 5315190600

model-00013-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2b6994f58b462df2107a899881309116d8b668e78588786e33fdcda5deee7308
+size 5315190584

model-00014-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a646727dce5f710bb23f38b01faae5a97e22966c0152a7f538ba711ec730ea9e
+size 5315190384

model-00015-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2fe62f582173fe86bac2a206b5b0d358a745355136c968401203305456b3e220
+size 5315199112

model-00016-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:97d63ab634d0d6ee498b3e97db332c283b7d4c2c4aa65f2d5d949d7451a6c75c
+size 5315199120

model-00017-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d3cc16384a3b881669a1616091877297f08c6e6fd1a212b79a991f1adc7d3192
+size 5315199120

model-00018-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6130194e56cd991ffc0bcaab6f2ca2198e2a38db30b9cf3eef1833da86b97c80
+size 5315199168

model-00019-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:19b89d889c49b935ba4c6579c2b5f9fd3a5005b97cd719059ac9ee3237f99e7c
+size 5315199192

model-00020-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2d274ec0a1f29bd9fa060c39de83e8dafd62e5b3882d6f907bf6f7fc40535a4b
+size 5314674968

model-00021-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:701ed226f316cd6d608234d2f77714b54fc071445beba1fd9da9375d2115b831
+size 5314674952

model-00022-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6c194582d2c6af1bcf698927aeff47128401c0302047c8465b5ab5882e7e697d
+size 5315068040

model-00023-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a084318a109e2776edd5a0fe71f87db3a513843c1eda6c374eed9469ad2c965
+size 5315068048

model-00024-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b2b2af8fa51ccc188f5c55afac3e3e0b59fb4c713dda44f9ec7820749cd851d6
+size 5315199128

model-00025-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c7643ca1f768712f319b01e552071b67f8550fb42c633bcee98b5fc40ab5501a
+size 5315199232

model_hashes.txt

@@ -0,0 +1,38 @@
+022ef6dfbf2f526e0ef9a66d636ab246675815f9877f6654f45d6cda522ec9e9  ./model-00035-of-00038.safetensors
+0c9d09758a0617c2b3ecc401f55141537114cf484b549a161628751ee9b699fa  ./model-00037-of-00038.safetensors
+0cb681b76b8cf18fb5a8ae136923bcb1ff79b2d706b8e73c8677decad7ff5648  ./model-00029-of-00038.safetensors
+139cc01fc2445267a046da2939d589367d17cc146be8259f69bfabc9e23eca63  ./model-00025-of-00038.safetensors
+15b59cf9ecdb92b277d5b7a10f25f35db7373daa5050d55a65f8be2a1f283624  ./model-00011-of-00038.safetensors
+1a66b7d59d3bb94d1874dfa7fc7b91afa2632af4d2dd19b14ee503596b38f4e2  ./model-00007-of-00038.safetensors
+1ccc5afa878158ddab54d0e8381e4e91ecc524ebd2f99aa28cdb8e71fc6ed155  ./model-00034-of-00038.safetensors
+2151fa0ae1c16356c7606f7fdd0778aa49f2727f5341b7daa19ac3b8afbf6e0e  ./model-00031-of-00038.safetensors
+239b6c3f708e901b273711c71f7b85619d462dcca9130e3506d6151f75b5387c  ./model-00026-of-00038.safetensors
+244ba1cc3f7da33f1b3bfce6eec2a38c1b1410eb1aea22c7796bd0e2bec53dd5  ./model-00002-of-00038.safetensors
+3647fdc2158e2201da25a6e6b93ebc77d52e085f28995920c94346d79f5da845  ./model-00016-of-00038.safetensors
+367a91bc240583a8a701b41ded917297cebd04113cdd13c3b1703be396a1fdab  ./model-00006-of-00038.safetensors
+3c91728af1195b4b7172a2190606fd3cfb53604aef9f0a27f951af509824e740  ./model-00028-of-00038.safetensors
+3ecba8f4358c5fa7a84c4e2a38b32e1c579e0601d661822288cd4ac629b58f71  ./model-00014-of-00038.safetensors
+411ac9baf753252971f9b7072ac08c950ff66c2acd410207a5c48480d898f532  ./model-00005-of-00038.safetensors
+4862c4cc80ac36701781ed02a84822a152e7a285b57cf1260f8fd23209d99cce  ./model-00001-of-00038.safetensors
+4c5828059c2f5e3fb067e07232856190a018863fae38a1e1858fc2dd1fb86ca1  ./model-00019-of-00038.safetensors
+58fa5258626248582925148ffaec79f752045d76b3073eff9035f78244174e27  ./model-00000-of-00038.safetensors
+590f6a7b7577a82bd7edf6434462316190aea55ccf2529053bd3465d0c232e52  ./model-00032-of-00038.safetensors
+5f964d48cf6d02f34fc4ba242bbe1a11d83c4df5fa660695f9fee53b53b84eab  ./model-00010-of-00038.safetensors
+70511491f956b62c2418781e3d9db56df42510e807735ec6d8f395ccb3340e82  ./model-00012-of-00038.safetensors
+7298c81efffeee772390b48e182e7d033be69c1d5df86c24a54d4f68f93b726a  ./model-00030-of-00038.safetensors
+73aebcb63519a519befe68cb1c26b05bcc310ee14dfac44ca49b1f18d82c69bf  ./model-00003-of-00038.safetensors
+7d14972288a2e73168f1fde0536b53bd9cd4e83daa3902cd8679fce38b8f677e  ./model-00036-of-00038.safetensors
+9372820581b918bfa6866b6628013f331d91946846c65062374750cd5ca16fe9  ./model-00018-of-00038.safetensors
+94f4f1b53f863bf2e4fb7f817bfd722c4e3f712180a47b4c173efe86f9f708ff  ./model-00020-of-00038.safetensors
+978483a120263813f2f0d1f0a4f741390202d0364e5c449b79e8f0fc05819a1c  ./model-00017-of-00038.safetensors
+982852a74d080f1567e6a02aee0b60ef56bf37f5073cf2ae904bd36561f5f395  ./model-00027-of-00038.safetensors
+99d1e0f997b7f6e4cc7c2245649f0facbcef377bf88f883a31f2faa557cd8ac9  ./model-00009-of-00038.safetensors
+a8fee6ecd8a561d353cd135a681252f923dc247d27c624d083b45c21354d4080  ./model-00023-of-00038.safetensors
+ad229f37d38cd12675367219d16f803f2a07d5fb71bc84c9d9c62598692a4100  ./model-00024-of-00038.safetensors
+ae5f05ffeeeceaa0d09a4703340fd7e7834d87d23d8a7b5f928a3a767a4d29cf  ./model-00021-of-00038.safetensors
+c549541ab83b5013190c191de4db8ff7dcabc5fa78cb68796970dabda4a14636  ./model-00015-of-00038.safetensors
+e10c54ec49ab87bb7abf1b9df07b2767a37a26ddf1c62014d464144c01ff8ea9  ./model-00008-of-00038.safetensors
+e614af5d55350bda182d21e2e39a69c4436f416f301bc536a8280d6e5fbbc465  ./model-00013-of-00038.safetensors
+ebae9b3f3ff94b0b36b352ec72a0cc228fbe611dc59b3026a44d932961654791  ./model-00022-of-00038.safetensors
+ed6f0578b36b47680355b736a48ab732380b77d29d29f8c0e18386397523ca59  ./model-00004-of-00038.safetensors
+f95d2a591d0ca4cb93db00a13c691f006ac8aaf0a4aeb09fc84447f0176aab10  ./model-00033-of-00038.safetensors

modeling_longcat_ngram.py

@@ -0,0 +1,338 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2025 Meituan
+# This code is licensed under the MIT License, for details, see the ./LICENSE file. 
+
+from typing import Optional, Tuple, Dict, List
+
+import torch
+from torch import nn
+
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.processing_utils import Unpack
+from transformers.utils import auto_docstring, logging
+from transformers.models.longcat_flash.modeling_longcat_flash import (
+    LongcatFlashForCausalLM,
+    LongcatFlashModel,
+    LongcatFlashRMSNorm,
+    LongcatFlashRotaryEmbedding,
+    LongcatFlashDecoderLayer,
+    LongcatFlashPreTrainedModel,
+)
+from .configuration_longcat_ngram import LongcatFlashNgramConfig
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring
+class LongcatFlashNgramPreTrainedModel(LongcatFlashPreTrainedModel):
+    pass
+
+
+class NgramCache(DynamicCache):
+    """
+    Extended DynamicCache for storing N-gram context alongside KV cache.
+    """
+    def __init__(self, config=None):
+        super().__init__()
+        self.ngram_context = None
+        # Keep only n-1 tokens (minimum needed for N-gram computation)
+        self.max_context_len = config.emb_neighbor_num - 1
+    
+    def update_ngram_context(self, new_tokens: torch.Tensor) -> None:
+        """
+        Update N-gram context with window management.
+        
+        Args:
+            new_tokens: New tokens to append, shape (batch_size, seq_len)
+        """
+        if self.ngram_context is None:
+            self.ngram_context = new_tokens.clone()
+        else:
+            self.ngram_context = torch.cat([self.ngram_context, new_tokens], dim=-1)
+        
+        # Truncate to maintain constant memory footprint
+        if self.ngram_context.size(-1) > self.max_context_len:
+            self.ngram_context = self.ngram_context[..., -self.max_context_len:]
+    
+    def reorder_cache(self, beam_idx: torch.LongTensor) -> "Cache":
+        """Reorder cache for beam search."""
+        # Reorder parent's KV cache
+        super().reorder_cache(beam_idx)
+        
+        # Reorder N-gram context
+        if self.ngram_context is not None:
+            self.ngram_context = self.ngram_context.index_select(0, beam_idx.to(self.ngram_context.device))
+        
+        return self
+
+
+class NgramEmbedding(nn.Module):
+    """
+    Computes embeddings enriched with N-gram features without maintaining internal state.
+    """
+    def __init__(self, config, base_embeddings):
+        super().__init__()
+        self.config = config
+        self.word_embeddings = base_embeddings
+        
+        self.m = config.ngram_vocab_size_ratio * config.vocab_size
+        self.k = config.emb_split_num
+        self.n = config.emb_neighbor_num
+        
+        self._init_ngram_embeddings()
+        self._vocab_mods_cache = None
+        
+    def _init_ngram_embeddings(self) -> None:
+        """Initialize N-gram embedding and projection layers."""
+        num_embedders = self.k * (self.n - 1)
+        emb_dim = self.config.hidden_size // num_embedders
+        
+        embedders = []
+        post_projs = []
+        
+        for i in range(num_embedders):
+            vocab_size = int(self.m + i * 2 + 1)
+            emb = nn.Embedding(vocab_size, emb_dim, padding_idx=self.config.pad_token_id)
+            proj = nn.Linear(emb_dim, self.config.hidden_size, bias=False)
+            embedders.append(emb)
+            post_projs.append(proj)
+        
+        self.embedders = nn.ModuleList(embedders)
+        self.post_projs = nn.ModuleList(post_projs)
+    
+    def _shift_right_ignore_eos(self, tensor: torch.Tensor, n: int, eos_token_id: int = 2) -> torch.Tensor:
+        """Shift tensor right by n positions, resetting at EOS tokens."""
+        batch_size, seq_len = tensor.shape
+        result = torch.zeros_like(tensor)
+        eos_mask = (tensor == eos_token_id)
+        
+        for i in range(batch_size):
+            eos_positions = eos_mask[i].nonzero(as_tuple=True)[0]
+            prev_idx = 0
+            
+            for eos_idx in eos_positions:
+                end_idx = eos_idx.item() + 1
+                if end_idx - prev_idx > n:
+                    result[i, prev_idx+n:end_idx] = tensor[i, prev_idx:end_idx-n]
+                prev_idx = end_idx
+            
+            if prev_idx < seq_len and seq_len - prev_idx > n:
+                result[i, prev_idx+n:seq_len] = tensor[i, prev_idx:seq_len-n]
+        
+        return result
+    
+    def _precompute_vocab_mods(self) -> Dict[Tuple[int, int], List[int]]:
+        """Precompute modular arithmetic values for vocabulary."""
+        if self._vocab_mods_cache is not None:
+            return self._vocab_mods_cache
+        
+        vocab_mods = {}
+        vocab_size = self.config.vocab_size
+        
+        for i in range(2, self.n + 1):
+            for j in range(self.k):
+                index = (i - 2) * self.k + j
+                emb_vocab_dim = int(self.m + index * 2 + 1)
+                
+                mods = []
+                power_mod = 1
+                for _ in range(i - 1):
+                    power_mod = (power_mod * vocab_size) % emb_vocab_dim
+                    mods.append(power_mod)
+                
+                vocab_mods[(i, j)] = mods
+        
+        self._vocab_mods_cache = vocab_mods
+        return vocab_mods
+    
+    def _get_ngram_ids(
+        self, 
+        input_ids: torch.Tensor, 
+        shifted_ids: Dict[int, torch.Tensor], 
+        vocab_mods: List[int], 
+        ngram: int
+    ) -> torch.Tensor:
+        """Compute N-gram hash IDs using polynomial rolling hash."""
+        ngram_ids = input_ids.clone()
+        for k in range(2, ngram + 1):
+            ngram_ids = ngram_ids + shifted_ids[k] * vocab_mods[k - 2]
+        return ngram_ids
+    
+    def forward(
+        self, 
+        input_ids: torch.Tensor, 
+        ngram_context: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Stateless forward pass.
+        
+        Args:
+            input_ids: Current input token IDs of shape (batch_size, seq_len)
+            ngram_context: Optional historical context of shape (batch_size, context_len)
+            
+        Returns:
+            Embedding tensor of shape (batch_size, seq_len, hidden_size)
+        """
+        seq_len = input_ids.size(-1)
+        
+        # Determine complete context
+        if ngram_context is not None:
+            context = torch.cat([ngram_context[..., -(self.n-1):], input_ids], dim=-1)
+        else:
+            context = input_ids
+        
+        # Base word embeddings
+        device = self.word_embeddings.weight.device
+        x = self.word_embeddings(input_ids.to(device)).clone()
+        
+        # Precompute modular values
+        vocab_mods = self._precompute_vocab_mods()
+        
+        # Compute shifted IDs
+        shifted_ids = {}
+        for i in range(2, self.n + 1):
+            shifted_ids[i] = self._shift_right_ignore_eos(
+                context, i - 1, eos_token_id=self.config.eos_token_id
+            )
+        
+        # Add N-gram embeddings
+        for i in range(2, self.n + 1):
+            for j in range(self.k):
+                index = (i - 2) * self.k + j
+                emb_vocab_dim = int(self.m + index * 2 + 1)
+                
+                ngram_ids = self._get_ngram_ids(context, shifted_ids, vocab_mods[(i, j)], ngram=i)
+                new_ids = (ngram_ids % emb_vocab_dim)[..., -seq_len:]
+                
+                embedder_device = self.embedders[index].weight.device
+                x_ngram = self.embedders[index](new_ids.to(embedder_device))
+                
+                proj_device = self.post_projs[index].weight.device
+                x_proj = self.post_projs[index](x_ngram.to(proj_device))
+                x = x + x_proj.to(x.device)
+        
+        # Normalize
+        x = x / (1 + self.k * (self.n - 1))
+        
+        return x
+
+
+class LongcatFlashNgramModel(LongcatFlashModel):
+    """LongcatFlash model with N-gram enhanced embeddings."""
+    _keys_to_ignore_on_load_unexpected = [r"model\.mtp.*"]
+    config_class = LongcatFlashNgramConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.ngram_embeddings = NgramEmbedding(config, self.embed_tokens)
+
+        self.layers = nn.ModuleList(
+            [LongcatFlashDecoderLayer(config, layer_idx) for layer_idx in range(config.num_layers)]
+        )
+        
+        self.head_dim = config.head_dim  
+        self.config.num_hidden_layers = 2 * config.num_layers
+        self.norm = LongcatFlashRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = LongcatFlashRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        **kwargs
+    ) -> BaseModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        # Extract N-gram context if available
+        ngram_context = None
+        if isinstance(past_key_values, NgramCache) and past_key_values.ngram_context is not None:
+            ngram_context = past_key_values.ngram_context
+
+        if inputs_embeds is None:
+            inputs_embeds = self.ngram_embeddings(input_ids, ngram_context=ngram_context)
+
+        # Initialize NgramCache if needed
+        if use_cache and past_key_values is None:
+            past_key_values = NgramCache(config=self.config)
+        
+        # Update N-gram context
+        if use_cache and isinstance(past_key_values, NgramCache):
+            past_key_values.update_ngram_context(input_ids)
+
+        # Prepare cache position
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                inputs_embeds.shape[1], device=inputs_embeds.device
+            ) + past_seen_tokens
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        # Create causal mask
+        causal_mask = create_causal_mask(
+            config=self.config,
+            input_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+
+        # Forward through decoder layers
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.layers[: self.config.num_layers]:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+            hidden_states=None,
+            attentions=None,
+        )
+
+
+class LongcatFlashNgramForCausalLM(LongcatFlashForCausalLM):
+    """LongcatFlash model for causal language modeling with N-gram embeddings."""
+    _keys_to_ignore_on_load_unexpected = [r"model\.mtp.*"]
+    config_class = LongcatFlashNgramConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = LongcatFlashNgramModel(config)
+
+    @torch.no_grad()
+    def generate(self, inputs=None, generation_config=None, **kwargs):
+        """Override to ensure NgramCache is used."""
+
+        if "past_key_values" not in kwargs or kwargs["past_key_values"] is None:
+            kwargs["past_key_values"] = NgramCache(config=self.config)
+        
+        return super().generate(inputs=inputs, generation_config=generation_config, **kwargs)
+
+__all__ = ["LongcatFlashNgramPreTrainedModel", "LongcatFlashNgramModel", "LongcatFlashNgramForCausalLM"]

parse_model_response.py

@@ -0,0 +1,236 @@
+import re
+import json
+import uuid
+
+def parse_arguments(json_value):
+    """
+    Attempt to parse a string as JSON
+    
+    Args:
+        json_value: String to parse
+        
+    Returns:
+        tuple: (parsed_value, is_valid_json)
+    """
+    try:
+        parsed_value = json.loads(json_value)
+        return parsed_value, True
+    except:
+        return json_value, False
+
+def get_argument_type(func_name: str, arg_key: str, defined_tools: list):
+    """
+    Get the type definition of a tool parameter
+    
+    Args:
+        func_name: Name of the function/tool
+        arg_key: Parameter key name
+        defined_tools: List of tool definitions
+        
+    Returns:
+        str or None: Type of the parameter ('string', 'object', 'array', 'integer', 'number', 'boolean')
+    """
+    name2tool = {tool["name"]: tool for tool in defined_tools}
+    if func_name not in name2tool:
+        return None
+    tool = name2tool[func_name]
+    if "parameters" not in tool or "properties" not in tool["parameters"]:
+        return None
+    if arg_key not in tool["parameters"]["properties"]:
+        return None
+    return tool["parameters"]["properties"][arg_key].get("type")
+
+def parse_model_response(response: str, defined_tools: list=[]):
+    """
+    Parse model response to extract reasoning_content, content, and tool_calls
+    
+    Args:
+        response: Raw response text from the model
+        defined_tools: List of tool definitions
+
+    Returns:
+        dict: Message containing role, reasoning_content (optional), content (optional), 
+              and tool_calls (optional)
+    """
+    text = response
+    reasoning_content = None
+    content = None
+    tool_calls = []
+    
+    formatted_tools = []
+    for tool in defined_tools:
+        if "function" in tool:
+            formatted_tools.append(tool['function'])
+        else:
+            formatted_tools.append(tool)
+                
+    if '</longcat_think>' in text:
+        text = text.replace('<longcat_think>', '')
+        thinking_end = text.find('</longcat_think>')
+        reasoning_content = text[: thinking_end].strip()
+        text = text[thinking_end + len('</longcat_think>'):].lstrip()
+    
+    assert '<longcat_think>' not in text, "Unclosed <longcat_think> tag found in remaining text"
+    assert '</longcat_think>' not in text, "Unexpected </longcat_think> tag found without opening tag"
+    
+    if '<longcat_tool_call>' in text:
+        index = text.find('<longcat_tool_call>')
+        content = text[:index]
+        text = text[index:].strip()
+    else:
+        content = text
+        text = ""
+    
+    open_tags = text.count('<longcat_tool_call>')
+    close_tags = text.count('</longcat_tool_call>')
+    assert open_tags == close_tags, \
+        f"Mismatched tool_call tags: {open_tags} opening tags, {close_tags} closing tags"
+    
+    tool_call_strs = re.findall(
+        r'<longcat_tool_call>(.*?)</longcat_tool_call>', 
+        text, 
+        re.DOTALL
+    )
+    
+    for call in tool_call_strs:
+        func_name_match = re.match(r'([^\n<]+)', call.strip())
+        assert func_name_match, f"Missing function name in tool call: {call[:100]}"
+        
+        func_name = func_name_match.group(1).strip()
+        assert func_name, "Empty function name in tool call"
+        
+        # Verify argument tags are properly paired
+        arg_key_count = call.count('<longcat_arg_key>')
+        arg_key_close_count = call.count('</longcat_arg_key>')
+        arg_value_count = call.count('<longcat_arg_value>')
+        arg_value_close_count = call.count('</longcat_arg_value>')
+        
+        assert arg_key_count == arg_key_close_count, \
+            f"Mismatched arg_key tags in function {func_name}: {arg_key_count} opening, {arg_key_close_count} closing"
+        assert arg_value_count == arg_value_close_count, \
+            f"Mismatched arg_value tags in function {func_name}: {arg_value_count} opening, {arg_value_close_count} closing"
+        assert arg_key_count == arg_value_count, \
+            f"Mismatched arg_key and arg_value count in function {func_name}: {arg_key_count} keys, {arg_value_count} values"
+        
+        pairs = re.findall(
+            r'<longcat_arg_key>(.*?)</longcat_arg_key>\s*<longcat_arg_value>(.*?)</longcat_arg_value>', 
+            call, 
+            re.DOTALL
+        )
+        
+        assert len(pairs) == arg_key_count, \
+            f"Failed to parse all arguments in function {func_name}: expected {arg_key_count}, got {len(pairs)}"
+        
+        arguments = {}
+        for arg_key, arg_value in pairs:
+            arg_key = arg_key.strip()
+            arg_value = arg_value.strip()
+            
+            assert arg_key, f"Empty argument key in function {func_name}"
+            assert arg_key not in arguments, \
+                f"Duplicate argument key '{arg_key}' in function {func_name}"
+            
+            arg_type = get_argument_type(func_name, arg_key, formatted_tools)
+            
+            if arg_type and arg_type != 'string':
+                parsed_value, is_good_json = parse_arguments(arg_value)
+                arg_value = parsed_value
+            
+            arguments[arg_key] = arg_value
+        
+        tool_calls.append({
+            'id': "tool-call-" + str(uuid.uuid4()),
+            'type': "function",
+            'function': {
+                'name': func_name,
+                'arguments': arguments
+            }
+        })
+    
+    message = {'role': 'assistant'}
+    
+    if reasoning_content:
+        message['reasoning_content'] = reasoning_content
+    message['content'] = content
+    if tool_calls:
+        message['tool_calls'] = tool_calls
+    
+    return message
+
+if __name__=="__main__":
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from parse_model_response import parse_model_response
+
+    model_name = "meituan-longcat/LongCat-Flash-Lite"
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype="auto",
+        device_map="auto",
+        trust_remote_code=True
+    )
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "Give me a brief introduction to large language models."}
+    ]
+    input_ids = tokenizer.apply_chat_template(
+        messages, 
+        add_generation_prompt=True, 
+        return_tensors="pt"
+    ).to(model.device)
+    generated_ids = model.generate(inputs=input_ids, max_new_tokens=256)
+    output_ids = generated_ids[0][len(input_ids[0]):].tolist()
+    response = tokenizer.decode(output_ids, skip_special_tokens=True).strip("\n")
+    print("Example 1: sample response.")
+    print("\nRaw response:")
+    print(response)
+    print("\nParsed result:")
+
+    response = tokenizer.decode(output_ids, skip_special_tokens=True).strip("\n")
+    parsed_message = parse_model_response(response)
+    print(json.dumps(parsed_message, indent=2, ensure_ascii=False))
+
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "func_add",
+                "description": "Calculate the sum of two numbers",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "x1": {"type": "number", "description": "The first addend"},
+                        "x2": {"type": "number", "description": "The second addend"}
+                    },
+                    "required": ["x1", "x2"]
+                }
+            }
+        }
+    ]
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "Please tell me what is $$125679 + 234519$$?"},
+        # {
+        #     "role": "assistant", 
+        #     "content": "I'll calculate the sum of 125679 and 234519 for you.", 
+        #     "tool_calls": [{"type": "function", "function": {"name": "func_add", "arguments": {"x1": 125679, "x2": 234519}}}]
+        # },
+        # {"role": "tool", "name": "func_add", "content": '{"ans": 360198}'}
+    ]
+
+    input_ids = tokenizer.apply_chat_template(
+        messages, 
+        tools=tools,
+        add_generation_prompt=True, 
+        return_tensors="pt"
+    ).to(model.device)
+    generated_ids = model.generate(inputs=input_ids, max_new_tokens=256)
+    output_ids = generated_ids[0][len(input_ids[0]):].tolist()
+    response = tokenizer.decode(output_ids, skip_special_tokens=True).strip("\n")
+    print("Example 2: tool call response.")
+    print("\nRaw response:")
+    print(response)
+    print("\nParsed result:")
+    parsed_message = parse_model_response(response, tools)
+    print(json.dumps(parsed_message, indent=2, ensure_ascii=False))

special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<longcat_s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</longcat_s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<longcat_pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<longcat_unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tech_report.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cae927e0ba049046b7872dba03467158db9579914f534143db56a17c5f1bcd02
+size 736397

tokenizer_config.json

@@ -0,0 +1,42 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": true,
+  "add_prefix_space": false,
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "<longcat_s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "clean_up_tokenization_spaces": false,
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "</longcat_s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "model_max_length": 131072,
+  "pad_token": {
+    "__type": "AddedToken",
+    "content": "<longcat_pad>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sp_model_kwargs": {},
+  "tokenizer_class": "BloomTokenizer",
+  "unk_token": {
+    "__type": "AddedToken",
+    "content": "<longcat_unk>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "chat_template": "{%- set tool_choice = tool_choice | default('auto') %}\n{%- set ns = namespace(tool_types = [], last_query_index = -1) %}\n\n{%- if tools and tool_choice != 'none' %}\n    {{- \"<longcat_tool_declare>\\n\"-}}\n    {{- \"# Tools\\n\" }}\n    {{- \"You have access to the following tools:\\n\\n\" }}\n    {%- for tool in tools %}\n        {%- if tool.type not in ns.tool_types %}\n            {%- set ns.tool_types = ns.tool_types + [tool.type] %}\n            {{- \"## Tool namespace: \" ~ tool.type ~ \"\\n\\n\" }}\n        {%- endif %}\n        {%- if tool.type == 'code_interpreter' %}\n            {%- set tool = {\"type\":\"code_interpreter\",\"function\":{\"name\":\"code_interpreter_preview\",\"description\":\"The code will be executed in a stateful Jupyter notebook sandbox environment, only supports local computation, data processing, and file operations.\\nCode sandbox environment (network isolated) Any external network requests or online API calls are prohibited.\\nIf online functionality is needed, please use other permitted tools.\\nCode will respond with the output of the execution or time out after 60.0 seconds. \",\"parameters\":{\"type\":\"object\",\"properties\":{\"language\":{\"type\":\"string\",\"description\":\"The programming language of the code to be executed. Available values: python (Default), java, go, js, ts, c, c++.\"},\"code\":{\"type\":\"string\",\"description\":\"Python code to be executed must not include the following:\\n- Importing network libraries such as requests, httplib, etc.\\n- Any form of HTTP requests.\\n- External API calls.\\n- Network port operations. Example: ```python\\nimport pandas as pd\\npd.DataFrame({'A':[1,2]})\\n```\"},\"timeout\":{\"type\":\"number\",\"description\":\"The maximum execution time of the code, in seconds. Default is 60.0.\"}}},\"required\":[\"code\"]}} %}\n        {%- endif %}\n        {{- \"### Tool name: \" + tool.function.name + \"\\n\" }}\n        {{- \"Description: \" + tool.function.description + \"\\n\\n\" }}\n        {{- \"InputSchema: \" + tool.function.parameters | tojson(ensure_ascii=False) + \"\\n\\n\" }}\n    {%- endfor %}\n    {{- '**Note**: For each function call, output the function name and arguments within the following XML format:\\n<longcat_tool_call>{function-name}\\n<longcat_arg_key>{arg-key-1}</longcat_arg_key>\\n<longcat_arg_value>{arg-value-1}</longcat_arg_value>\\n<longcat_arg_key>{arg-key-2}</longcat_arg_key>\\n<longcat_arg_value>{arg-value-2}</longcat_arg_value>\\n...\\n</longcat_tool_call>\\n' }}\n    {{- \"</longcat_tool_declare>\"-}}\n    {%- for idx in range(messages|length - 1) %}\n        {%- set msg = messages[idx] %}\n        {%- if msg.role == 'assistant' and not msg.tool_calls %}\n            {%- set ns.last_query_index = idx %}\n        {%- endif %}\n    {%- endfor%}\n{%- endif %}\n\n{%- for msg in messages %}\n    {%- if msg.role == \"system\" %}\n        {{- \"<longcat_system>\" + msg.content }}\n    {%- elif msg.role == \"user\" %}\n        {{- \"<longcat_user>\" }}\n        {%- if msg[\"files\"] %}\n            {{- '<longcat_files>\\n' ~ msg.files | tojson(indent=2) ~ '\\n</longcat_files>' }}\n        {%- endif %}\n        {{- msg.content }}\n    {%- elif msg.role == \"assistant\" %}\n        {{- \"<longcat_assistant>\" }}\n        {%- if enable_thinking == true and msg.reasoning_content and ns.tool_types != [] and loop.index0 > ns.last_query_index %}\n            {{- \"\\n<longcat_think>\\n\" ~ msg.reasoning_content ~ \"\\n</longcat_think>\\n\" }}\n        {%- endif %}\n        {%- if msg.content%}\n            {{- msg.content }}\n        {%- endif %}\n        {%- if msg.tool_calls %}\n            {%- for tool_call in msg.tool_calls -%}\n                {{- \"<longcat_tool_call>\" ~ tool_call.function.name ~ \"\\n\" -}}\n                {% set _args = tool_call.function.arguments %}\n                {% for k, v in _args.items() %}\n                    {{- \"<longcat_arg_key>\" ~ k ~ \"</longcat_arg_key>\\n\" -}}\n                    {{- \"<longcat_arg_value>\" ~ (v if v is string else v | tojson(ensure_ascii=False)) ~ \"</longcat_arg_value>\\n\" -}}\n                {% endfor %}\n                {{- \"</longcat_tool_call>\\n\" }}\n            {%- endfor %}\n        {%- endif %}\n        {{- \"</longcat_s>\" -}}\n    {%- elif msg.role == \"tool\" %}\n        {%- if messages[loop.index0 - 1].role != \"tool\"%}\n            {{- \"<longcat_user>\" -}}\n        {%- endif %}\n        {{- \"<longcat_tool_response>\" ~ msg.content ~ \"</longcat_tool_response>\"-}}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {%- if enable_thinking == true %}\n        {{- \" /think_on\" }}\n        {%- if thinking_budget %}\n            {%- if thinking_budget < 1024 %}\n                {%- set thinking_budget = 1024 %}\n            {%- endif%}\n            {{- \"\\nthinking_budget: < \" ~ thinking_budget ~ \".\"}}\n        {%- endif %}\n        {{- \" <longcat_assistant><longcat_think>\\n\"}}\n    {%- elif enable_thinking == false %}\n        {{- \" /think_off <longcat_assistant><longcat_think>\\n\\n</longcat_think>\\n\" }}\n    {%- else %}\n        {{- \"<longcat_assistant>\" }}\n    {%- endif %}\n{%- endif %}"
+}