support python api

.gitignore

@@ -0,0 +1,2 @@
+kvcache/
+__pycache__/

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8/.gitattributes

@@ -0,0 +1,35 @@
+*.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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8/LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2024 Alibaba Cloud
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8/README.md

@@ -0,0 +1,108 @@
+---
+base_model: Qwen/Qwen2.5-1.5B-Instruct
+language:
+- en
+library_name: transformers
+license: apache-2.0
+license_link: https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct-GPTQ-Int8/blob/main/LICENSE
+pipeline_tag: text-generation
+tags:
+- chat
+---
+# Qwen2.5-1.5B-Instruct-GPTQ-Int8
+
+## Introduction
+
+Qwen2.5 is the latest series of Qwen large language models. For Qwen2.5, we release a number of base language models and instruction-tuned language models ranging from 0.5 to 72 billion parameters. Qwen2.5 brings the following improvements upon Qwen2:
+
+- Significantly **more knowledge** and has greatly improved capabilities in **coding** and **mathematics**, thanks to our specialized expert models in these domains.
+- Significant improvements in **instruction following**, **generating long texts** (over 8K tokens), **understanding structured data** (e.g, tables), and **generating structured outputs** especially JSON. **More resilient to the diversity of system prompts**, enhancing role-play implementation and condition-setting for chatbots.
+- **Long-context Support** up to 128K tokens and can generate up to 8K tokens.
+- **Multilingual support** for over 29 languages, including Chinese, English, French, Spanish, Portuguese, German, Italian, Russian, Japanese, Korean, Vietnamese, Thai, Arabic, and more. 
+
+**This repo contains the GPTQ-quantized 8-bit instruction-tuned 1.5B Qwen2.5 model**, which has the following features:
+- Type: Causal Language Models
+- Training Stage: Pretraining & Post-training
+- Architecture: transformers with RoPE, SwiGLU, RMSNorm, Attention QKV bias and tied word embeddings
+- Number of Parameters: 1.54B
+- Number of Paramaters (Non-Embedding): 1.31B
+- Number of Layers: 28
+- Number of Attention Heads (GQA): 12 for Q and 2 for KV
+- Context Length: Full 32,768 tokens and generation 8192 tokens
+- Quantization: GPTQ 8-bit
+
+For more details, please refer to our [blog](https://qwenlm.github.io/blog/qwen2.5/), [GitHub](https://github.com/QwenLM/Qwen2.5), and [Documentation](https://qwen.readthedocs.io/en/latest/).
+
+## Requirements
+
+The code of Qwen2.5 has been in the latest Hugging face `transformers` and we advise you to use the latest version of `transformers`.
+
+With `transformers<4.37.0`, you will encounter the following error:
+```
+KeyError: 'qwen2'
+```
+
+Also check out our [GPTQ documentation](https://qwen.readthedocs.io/en/latest/quantization/gptq.html) for more usage guide.
+
+## Quickstart
+
+Here provides a code snippet with `apply_chat_template` to show you how to load the tokenizer and model and how to generate contents.
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+model_name = "Qwen/Qwen2.5-1.5B-Instruct-GPTQ-Int8"
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype="auto",
+    device_map="auto"
+)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+prompt = "Give me a short introduction to large language model."
+messages = [
+    {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},
+    {"role": "user", "content": prompt}
+]
+text = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True
+)
+model_inputs = tokenizer([text], return_tensors="pt").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]
+```
+
+
+## Evaluation & Performance
+
+Detailed evaluation results are reported in this [📑 blog](https://qwenlm.github.io/blog/qwen2.5/).
+
+For quantized models, the benchmark results against the original bfloat16 models can be found [here](https://qwen.readthedocs.io/en/latest/benchmark/quantization_benchmark.html)
+
+For requirements on GPU memory and the respective throughput, see results [here](https://qwen.readthedocs.io/en/latest/benchmark/speed_benchmark.html).
+
+## Citation
+
+If you find our work helpful, feel free to give us a cite.
+
+```
+@misc{qwen2.5,
+    title = {Qwen2.5: A Party of Foundation Models},
+    url = {https://qwenlm.github.io/blog/qwen2.5/},
+    author = {Qwen Team},
+    month = {September},
+    year = {2024}
+}
+@article{qwen2,
+      title={Qwen2 Technical Report}, 
+      author={An Yang and Baosong Yang and Binyuan Hui and Bo Zheng and Bowen Yu and Chang Zhou and Chengpeng Li and Chengyuan Li and Dayiheng Liu and Fei Huang and Guanting Dong and Haoran Wei and Huan Lin and Jialong Tang and Jialin Wang and Jian Yang and Jianhong Tu and Jianwei Zhang and Jianxin Ma and Jin Xu and Jingren Zhou and Jinze Bai and Jinzheng He and Junyang Lin and Kai Dang and Keming Lu and Keqin Chen and Kexin Yang and Mei Li and Mingfeng Xue and Na Ni and Pei Zhang and Peng Wang and Ru Peng and Rui Men and Ruize Gao and Runji Lin and Shijie Wang and Shuai Bai and Sinan Tan and Tianhang Zhu and Tianhao Li and Tianyu Liu and Wenbin Ge and Xiaodong Deng and Xiaohuan Zhou and Xingzhang Ren and Xinyu Zhang and Xipin Wei and Xuancheng Ren and Yang Fan and Yang Yao and Yichang Zhang and Yu Wan and Yunfei Chu and Yuqiong Liu and Zeyu Cui and Zhenru Zhang and Zhihao Fan},
+      journal={arXiv preprint arXiv:2407.10671},
+      year={2024}
+}
+```

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8/config.json

@@ -0,0 +1,51 @@
+{
+  "architectures": [
+    "Qwen2ForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "eos_token_id": 151645,
+  "hidden_act": "silu",
+  "hidden_size": 1536,
+  "initializer_range": 0.02,
+  "intermediate_size": 8960,
+  "max_position_embeddings": 32768,
+  "max_window_layers": 21,
+  "model_type": "qwen2",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 2,
+  "quantization_config": {
+    "batch_size": 1,
+    "bits": 8,
+    "block_name_to_quantize": null,
+    "cache_block_outputs": true,
+    "damp_percent": 0.01,
+    "dataset": null,
+    "desc_act": false,
+    "exllama_config": {
+      "version": 1
+    },
+    "group_size": 128,
+    "max_input_length": null,
+    "model_seqlen": null,
+    "module_name_preceding_first_block": null,
+    "modules_in_block_to_quantize": null,
+    "pad_token_id": null,
+    "quant_method": "gptq",
+    "sym": true,
+    "tokenizer": null,
+    "true_sequential": true,
+    "use_cuda_fp16": false,
+    "use_exllama": true
+  },
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 1000000.0,
+  "sliding_window": 32768,
+  "tie_word_embeddings": true,
+  "torch_dtype": "float16",
+  "transformers_version": "4.39.3",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8/generation_config.json

@@ -0,0 +1,14 @@
+{
+  "bos_token_id": 151643,
+  "do_sample": true,
+  "eos_token_id": [
+    151645,
+    151643
+  ],
+  "pad_token_id": 151643,
+  "repetition_penalty": 1.1,
+  "temperature": 0.7,
+  "top_k": 20,
+  "top_p": 0.8,
+  "transformers_version": "4.39.3"
+}

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8/tokenizer_config.json

@@ -0,0 +1,207 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/model.embed_tokens.weight.npy

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l0_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l10_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l11_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l12_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l13_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l14_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l15_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l16_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l17_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l18_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l19_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l1_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l20_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l21_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l22_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l23_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l24_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l25_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l26_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l27_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l2_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l3_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l4_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l5_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l6_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l7_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l8_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_p128_l9_together.axmodel

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

python/Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/qwen2_post.axmodel

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

python/README.md

@@ -0,0 +1,66 @@
+在 `AXERA 650N` 开发板上使用 `python api` 进行模型推理.
+
+在当前目录执行以下命令:
+
+```sh
+$ python3 chat.py
+```
+
+当出现 `prompt (输入 q 退出对话) >>` 提示时输入文字, 等待模型输出, 具体示例如下:
+
+```bash
+$ python3 chat.py
+...
+The models have been loaded!
+2025-07-21 14:23:46.137 | DEBUG    | __main__:<module>:143 - >>> 创建 LlamaChatSession >>>
+>>> 系统提示: 你的名字叫小智(allen), 你是一个人畜无害的 AI 助手. 深圳市今天(4月1日)阴天, 愚人节, 气温在 14°C 至 19°C
+之间, 微风.
+2025-07-21 14:23:46.137 | INFO     | __main__:chat_loop:69 - Type 'q' to exit, Ctrl+c to stop current generation
+
+prompt (输入 q 退出对话) >> 定义函数y=3x^3+2x+1,求解它的导数.
+answer: >> 要找到函数 \( y = 3x^3 + 2x + 1 \) 的导数，我们需要对每个项分别求导，然后将它们相加起来。
+
+1. 对 \( 3x^3 \) ���导，结果是 \( 3 \cdot 3x^{3-1} = 9x^2 \)。
+2. 对 \( 2x \) ���导，结果是 \( 2 \cdot 1x^{1-1} = 2 \)。
+3. 对常数项 \( 1 \) ���导，结果是 \( 0 \)。
+
+将这些结果相加，我们得到：
+
+\[ y' = 9x^2 + 2 \]
+
+所以，函数 \( y = 3x^3 + 2x + 1 \) 的导数是 \( y' = 9x^2 + 2 \)。
+
+prompt (输入 q 退出对话) >> 这个函数中自变量和因变量分别是什么?
+answer: >> 在数学中，函数通常由两个变量组成：自变量（也称为输入变量）和因变量（也称为输出变量）。自变量是函数中的一个
+  量，它的值决定了因变量的值。
+
+在你提供的函数 \( y = 3x^3 + 2x + 1 \) 中：
+
+- \( x \) 是自变量。
+- \( y \) 是因变量。
+
+自变量 \( x \) 的值决定了因变量 \( y \) 的值。例如，如果你给 \( x \) ���值为 2，那么 \( y \) ��等于 \( 3(2)^3 + 2(2) +
+ 1 = 24 + 4 + 1 = 29 \)。
+
+因此，这个函数描述了一个关于 \( x \) 和 \( y \) 的关系，其中 \( x \) 是自变量，而 \( y \) 是因变量。通过改变 \( x \)
+  值，你可以计算出相应的 \( y \) ���。
+
+prompt (输入 q 退出对话) >> 这个函数中最高幂次和最低幂次分别是多少?
+answer: >> 在函数 \( y = 3x^3 + 2x + 1 \) 中，最高次幂（最高幂次）是 \( x^3 \)，因此最高幂次是 3。
+
+最低次幂（最低幂次）是 \( x^0 \)，因为 \( x^0 = 1 \) 对于任何 \( x \) ���成立，所以最低幂次是 0。
+
+因此，这个函数的最高幂次是 3，最低幂次是 0。最高幂次和最低幂次的差值是 \( 3 - 0 = 3 \)。这意味着函数的图形是一个三次多
+  式，它有一个顶点（如果最高幂次是偶数）或一个拐点（如果最高幂次是奇数）。在这个例子中，由于最高幂次是奇数，函数的图形
+  有一个拐点。
+
+```
+
+当上下文窗口达到上限, 可以输入 `reset` 命令重置, 例如:
+
+```sh
+prompt (输入 q 退出对话) >> reset
+上下文已重置
+prompt (输入 q 退出对话) >> 你是谁?今天天气如何?
+answer: >> 我是小智,一名人工智能助手。今天是阴天,愚人节,气温在14°C至19°C之间,微风。
+```

python/chat.py

@@ -0,0 +1,148 @@
+import sys
+import numpy as np
+import os
+import json
+from pathlib import Path
+from typing import List, Tuple, Dict, Optional
+from ml_dtypes import bfloat16
+import dataclasses
+from transformers import AutoTokenizer, AutoConfig
+import torch
+from torchvision.transforms.functional import InterpolationMode
+from axengine import InferenceSession
+from tqdm import tqdm
+import torchvision.transforms as T
+from PIL import Image
+import argparse
+from loguru import logger
+from copy import deepcopy
+from utils.infer_func import InferManager, KVCacheTools
+
+
+class LlamaChatSession:
+    def __init__(self, builder_instance):
+        self.system_prompt = builder_instance.system_prompt
+        self.builder_instance = builder_instance
+        self.last_reply = ""
+
+    def encode(self, prompt: str) -> Tuple[List[int], List[int]]:
+        """
+        keys: "message", "model_inputs", "input_ids", "input_embeds", "input_ids_len"
+        """
+        return self.builder_instance.encoder_prompt(prompt)
+
+    def get_kvcache(self) -> Tuple[List[np.ndarray], List[np.ndarray], int]:
+        return self.builder_instance.k_caches, self.builder_instance.v_caches
+
+    def generate(self,
+        model_inputs
+    ):
+        token_ids = model_inputs["input_ids"]
+        self.builder_instance.decode(token_ids)
+        return None
+
+    def run(self, model_inputs) -> str:
+        response = self.generate(
+            model_inputs
+        )
+        return response
+
+    def reset_context(self, system_prompt: str = None):
+        """
+        reset 只需要把 kv cache 清空即可 (甚至可以直接复写)
+        但如果 system_prompt 变化了，则需要重新计算 kv cache
+        """
+        if system_prompt is not None:
+            self.system_prompt = system_prompt
+
+        self.builder_instance.precompute_len = self.builder_instance.system_input_ids_len
+
+        for i in range(len(self.builder_instance.k_caches)):
+            self.builder_instance.k_caches[i][:, self.builder_instance.precompute_len:, :].fill(0)
+            self.builder_instance.v_caches[i][:, self.builder_instance.precompute_len:, :].fill(0)
+
+    def chat_loop(self, live_print: bool = False):
+
+        if self.system_prompt:
+            print(f">>> 系统提示: {self.system_prompt}")
+
+        logger.info("Type 'q' to exit, Ctrl+c to stop current generation\n")
+
+        while True:
+            try:
+                prompt = input("prompt (输入 q 退出对话) >> ")
+
+                if prompt.lower() == "q" or prompt.lower() == "exit":
+                    print("\nOK, 已退出对话.")
+                    return
+
+                if prompt.lower() == "debug":
+                    print(f"\n>>> DEBUG INFO >>>\n precompute_len is {self.builder_instance.precompute_len}\n<<< DEBUG INFO <<<\n")
+                    continue
+
+                if not prompt.strip():
+                    print(f"\n{self.system_prompt}")
+                    continue
+
+                if prompt.strip() == "reset":
+                    self.reset_context()
+                    print("上下文已重置")
+                    continue
+
+                model_inputs = self.encode(prompt)
+
+                if self.builder_instance.precompute_len + 128 >= 2559:
+                    logger.info("ERROR: 上下文窗口已满! 请使用 `reset` 命令重置上下文")
+                    continue
+
+                response = self.run(model_inputs)
+
+            except KeyboardInterrupt:
+                # 处理用户按 Ctrl+C 中断生成
+                print("\n好的, 已成功退出对话.")
+                exit()
+                
+            except Exception as e:
+                print(f"ERROR: {str(e)}")
+
+
+if __name__ == "__main__":
+
+    hf_model_path = './Qwen2.5-1.5B-Instruct-GPTQ-Int8/'
+    axmodel_model_path = './Qwen2.5-1.5B-Instruct-GPTQ-Int8_axmodel/'
+
+    builder = InferManager(hf_model_path, axmodel_model_path) # init tokenizer & hf_config & system prompt
+    builder.build_system_prompt()
+    builder.build_kvcache()
+    builder.build_infer_model()
+
+    cache_manager = KVCacheTools(axmodel_num=28, dtype=bfloat16)
+
+    if not os.path.exists("./kvcache"):
+        # 预计算 system prompt k,v 缓存
+        update_kv_cache = builder.prefill(
+            builder.model_inputs,
+            slice_len=128,
+        )
+        if cache_manager.save_kvcache(
+            target_dir="./kvcache",
+            system_prompt=builder.system_prompt,
+            precompute_len=builder.system_input_ids_len,
+            k_caches=update_kv_cache[0],
+            v_caches=update_kv_cache[1],
+            metadata={"model_version": "v0.1"}
+        ):
+            logger.info(">>> 预计算 system prompt kvcache 保存到 ./kvcache 目录, 下次启动可直接加载缓存 <<<")
+        else:
+            logger.error(">>> kvcache 缓存保存失败, 程序退出! <<<")
+            exit()
+    else:
+        update_kv_cache, prompt, plen, meta = cache_manager.load_kvcache("./kvcache")
+        builder.update_kvcache(update_kv_cache)
+
+    logger.debug(">>> 创建 LlamaChatSession >>>")
+
+    session = LlamaChatSession(
+        builder_instance=builder
+    )
+    session.chat_loop(live_print=False)

python/infer.py

@@ -0,0 +1,293 @@
+# import llm_utils
+import dataclasses
+import json
+from transformers import AutoTokenizer, AutoConfig
+import torch
+from torchvision.transforms.functional import InterpolationMode
+import numpy as np
+from ml_dtypes import bfloat16
+from axengine import InferenceSession
+from tqdm import tqdm
+import torchvision.transforms as T
+from PIL import Image
+import argparse
+
+
+def post_process(data, topk=1, topp=0.9, temperature=0.6):
+    def top_p(l: np.ndarray, p: float) -> np.ndarray:
+        index = np.argsort(l)
+        res = l.copy()
+        sum_p = 0
+        for i in index[::-1]:
+            if sum_p >= p:
+                res[i] = 0
+            sum_p += res[i]
+        return res / sum_p
+
+    def softmax(l: np.ndarray) -> np.ndarray:
+        l_max = l - l.max()
+        l_exp = np.exp(l_max)
+        res = l_exp / np.sum(l_exp)
+        return res.astype(np.float64)
+
+    r = data.astype(np.float32)
+    r = r.flatten()
+    # topk
+    candidate_index = np.argpartition(r, -topk)[-topk:]
+    candidate_value = r[candidate_index]
+    # temperature
+    candidate_value /= temperature
+    # softmax
+    candidate_soft = softmax(candidate_value)
+    # topp
+    candidate_soft = top_p(candidate_soft, topp)
+    candidate_soft = candidate_soft.astype(np.float64) / candidate_soft.sum()
+    pos = np.random.multinomial(1, candidate_soft).argmax()
+    next_token = candidate_index[pos]
+    return next_token, candidate_index, candidate_soft
+
+
+def generate_slice_indices(token_len, prefill=128, expand=512):
+    remaining = max(0, token_len - prefill)
+    extra_blocks = (remaining + expand - 1) // expand
+    return list(range(extra_blocks + 1))
+
+
+if __name__ == "__main__":
+
+    prompt = None
+    parser = argparse.ArgumentParser(description="Model configuration parameters")
+    parser.add_argument("--hf_model", type=str, default="../qwen2.5_tokenizer",
+                        help="Path to HuggingFace model")
+    parser.add_argument("--axmodel_path", type=str, default="../qwen2.5-1.5b-ctx-ax650",
+                        help="Path to save compiled axmodel of llama model")
+    parser.add_argument("-q", "--question", type=str, default="Please calculate the derivative of the function y=2x^2.",
+                        help="Your question that you want to ask the model.")
+    args = parser.parse_args()
+
+    device = "cpu"
+    hf_model_path = args.hf_model
+    axmodel_path = args.axmodel_path
+
+    cfg = AutoConfig.from_pretrained(hf_model_path, trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained(hf_model_path, trust_remote_code=True, use_fast=False)
+
+    prompt = args.question
+
+    print("注意: 这里预先设定了长 prompt 测试. ")
+    # input_ids: 438 > 128
+    prompt = '''你能将英文翻译成中文吗? 比如这句话:
+    Once when I was six years old I saw a magnificent picture in a book, called True Stories from Nature, about the primeval forest. It was a picture of a boa constrictor in the act of swallowing an animal. Here is a copy of the drawing. 
+    In the book it said: "Boa constrictors swallow their prey whole, without chewing it. After that they are not able to move, and they sleep through the six months that they need for digestion." 
+    I pondered deeply, then, over the adventures of the jungle. And after some work with a colored pencil I succeeded in making my first drawing. My Drawing Number One. It looked like this: 
+    I showed my masterpiece to the grown-ups, and asked them whether the drawing frightened them.
+    But they answered: "Frighten? Why should any one be frightened by a hat?"
+    My drawing was not a picture of a hat. It was a picture of a boa constrictor digesting an elephant. But since the grown-ups were not able to understand it, I made another drawing: I drew the inside of the boa constrictor, so that the grown-ups could see it clearly. They always need to have things explained. My Drawing Number Two looked like this: 
+    The grown-ups‘ response, this time, was to advise me to lay aside my drawings of boa constrictors, whether from the inside or the outside, and devote myself instead to geography, history, arithmetic and grammar. That is why, at the age of six, I gave up what might have been a magnificent career as a painter. I had been disheartened by the failure of my Drawing Number One and my Drawing Number Two. Grown-ups never understand anything by themselves, and it is tiresome for children to be always and forever explaining things to them.
+    So then I chose another profession, and learned to pilot airplanes.
+    '''
+
+    # prompt = "介绍一下你自己"
+    # prompt = "今天是几号，天气怎么样"
+    prompt = "你知道 `床前明月光,疑是地上霜`是谁写的吗?"
+
+    messages = [
+        {"role": "system", "content": "你的名字叫小智(allen), 你是一个人畜无害的 AI 助手. 深圳市今天(4月1日)阴天, 愚人节, 气温在 14°C 至 19°C 之间, 微风."},
+        {"role": "user", "content": prompt}
+    ]
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+    model_inputs = tokenizer([text], return_tensors="pt").to(device)
+    token_ids = model_inputs.input_ids[0].cpu().numpy().tolist()
+
+    embeds = np.load(f"{axmodel_path}/model.embed_tokens.weight.npy")
+    prefill_data = np.take(embeds, token_ids, axis=0)
+    token_len = len(token_ids)
+    import pdb; pdb.set_trace()
+
+    ##################
+    lastN = 2559
+
+    kv_dim = cfg.hidden_size // cfg.num_attention_heads * cfg.num_key_value_heads
+    k_caches = [
+        np.zeros((1, lastN, kv_dim), dtype=bfloat16)
+        for _ in range(cfg.num_hidden_layers)
+    ]
+    v_caches = [
+        np.zeros((1, lastN, kv_dim), dtype=bfloat16)
+        for _ in range(cfg.num_hidden_layers)
+    ]
+
+    prefill_decoder_sessins = []
+
+    for i in tqdm(range(cfg.num_hidden_layers), desc="Init InferenceSession"):
+        session = InferenceSession(
+            f"{axmodel_path}/qwen2_p128_l{i}_together.axmodel"
+        )
+        prefill_decoder_sessins.append(session)
+
+    post_process_session = InferenceSession(
+        f"{axmodel_path}/qwen2_post.axmodel"
+    )
+    print("model load done!")
+    print("prefill token_len: ", token_len)
+
+    """
+    Model input shape:
+        - kv_cache: g1:[1, 1, hidden_size] -> g2:[1, kv_mask_expand_lens, kv_dim] -> g3:[1, kv_mask_expand_lens * 2, kv_dim]
+        - mask: g1:[1, input_prefill_len, input_prefill_len] -> g2:[1, input_prefill_len, input_prefill_len+kv_mask_expand_lens] -> g3:[1, input_prefill_len, input_prefill_len+kv_mask_expand_lens*2]
+        - indices: g1:[1, input_prefill_len] -> g2:[1, input_prefill_len] -> g3:[1, input_prefill_len]
+        - input: g1:[1, input_prefill_len, hidden_size] -> g2:[1, input_prefill_len, hidden_size] -> g3:[1, input_prefill_len, hidden_size]
+    """
+
+    input_prefill_len = 128
+    kv_mask_expand_len = 128 # 512
+
+    """
+    Model output shape:
+        - kv_cache: g1:[1, input_prefill_len, kv_dim] -> g2:[1, input_prefill_len, kv_dim] -> g3:[1, input_prefill_len, kv_dim]
+        - output: g1:[1, input_prefill_len, hidden_size] -> g2:[1, input_prefill_len, hidden_size] -> g3:[1, input_prefill_len, hidden_size]
+    """
+    slice_indexs = generate_slice_indices(token_len, input_prefill_len, input_prefill_len)
+    print(f"slice_indexs is {slice_indexs}")
+
+    """
+        prefill
+    """
+    if input_prefill_len > 0:
+        for slice_index in slice_indexs:
+            if slice_index == 0:
+                current_slice_len = input_prefill_len
+            else:
+                current_slice_len = kv_mask_expand_len
+
+            indices = np.array(
+                list(
+                    range(
+                        slice_index * input_prefill_len,
+                        (slice_index + 1) * input_prefill_len,
+                    )
+                ),
+                np.uint32,
+            ).reshape((1, input_prefill_len))
+
+            mask = (
+                np.zeros((1, input_prefill_len, current_slice_len * slice_index + input_prefill_len))
+                - 65536
+            )
+            data = np.zeros((1, input_prefill_len, cfg.hidden_size)).astype(bfloat16)
+            for i, t in enumerate(
+                range(
+                    slice_index * input_prefill_len,
+                    (slice_index + 1) * input_prefill_len,
+                )
+            ):
+                if t < len(token_ids):
+                    mask[:, i, : slice_index * input_prefill_len + i + 1] = 0
+                    data[:, i : i + 1, :] = (
+                        prefill_data[t]
+                        .reshape((1, 1, cfg.hidden_size))
+                        .astype(bfloat16)
+                    )
+
+            if slice_index == slice_indexs[-1]:
+                curlen_procd = token_len - slice_index * input_prefill_len # curlen_procd 是当前处理数据的长度
+            else:
+                curlen_procd = input_prefill_len
+
+            mask = mask.astype(bfloat16)
+            for i in range(cfg.num_hidden_layers):
+                input_feed = {
+                    "K_cache": (
+                        k_caches[i][:, 0: current_slice_len * slice_index, :]
+                        if slice_index
+                        else np.zeros((1, 1, cfg.hidden_size), dtype=bfloat16)
+                    ),
+                    "V_cache": (
+                        v_caches[i][:, 0: current_slice_len * slice_index, :]
+                        if slice_index
+                        else np.zeros((1, 1, cfg.hidden_size), dtype=bfloat16)
+                    ),
+                    "indices": indices,
+                    "input": data,
+                    "mask": mask,
+                }
+                outputs = prefill_decoder_sessins[i].run(None, input_feed, shape_group=slice_index + 1)
+
+                k_caches[i][
+                    :,
+                    slice_index
+                    * input_prefill_len : slice_index
+                    * input_prefill_len + curlen_procd, # current_slice_len
+                    :,
+                ] = outputs[0][:, :curlen_procd, :]
+
+                v_caches[i][
+                    :,
+                    slice_index
+                    * input_prefill_len : slice_index
+                    * input_prefill_len + curlen_procd, # current_slice_len
+                    :,
+                ] = outputs[1][:, :curlen_procd, :]
+
+                data = outputs[2]
+
+            print("slice prefill done", slice_index)
+
+        post_out = post_process_session.run(
+            None,
+            {
+                "input": data[
+                    :, token_len - (len(slice_indexs) - 1) * input_prefill_len - 1, None, :
+                ]
+            }
+        )[0]
+        next_token, posssible_tokens, possible_soft = post_process(post_out)
+        posibles = [tokenizer.decode([t]) for t in posssible_tokens]
+        posible_soft = [str((t, s)) for t, s in zip(posibles, possible_soft)]
+        token_ids.append(next_token)
+
+    print("answer >>", tokenizer.decode(token_ids[token_len], skip_special_tokens=True), end='', flush=True)
+    # print("answer >>", end='', flush=True)
+
+    # set to decoder
+    kv_cache_len = lastN
+    mask = np.zeros((1, 1, kv_cache_len + 1), dtype=np.float32).astype(bfloat16)
+    mask[:, :, :kv_cache_len] -= 65536
+    if input_prefill_len > 0:
+        mask[:, :, :token_len] = 0
+
+    # for start_indice in tqdm(range(kv_cache_len), desc="Decode"):
+    for start_indice in range(kv_cache_len):
+        if input_prefill_len > 0 and start_indice < token_len:
+            continue
+
+        next_token = token_ids[start_indice]
+        indices = np.array([start_indice], np.uint32).reshape((1, 1))
+        data = embeds[next_token, :].reshape((1, 1, cfg.hidden_size)).astype(bfloat16)
+        for i in range(cfg.num_hidden_layers):
+            input_feed = {
+                "K_cache": k_caches[i],
+                "V_cache": v_caches[i],
+                "indices": indices,
+                "input": data,
+                "mask": mask,
+            }
+            outputs = prefill_decoder_sessins[i].run(None, input_feed, shape_group=0)
+            k_caches[i][:, start_indice, :] = outputs[0][:, :, :]
+            v_caches[i][:, start_indice, :] = outputs[1][:, :, :]
+            data = outputs[2]
+        mask[..., start_indice] = 0
+        if start_indice < token_len - 1:
+            pass
+        else:
+            post_out = post_process_session.run(None, {"input": data})[0]
+            next_token, posssible_tokens, possible_soft = post_process(post_out)
+            token_ids.append(next_token)
+            if next_token == tokenizer.eos_token_id and next_token > token_len:
+                break
+
+        print(tokenizer.decode(next_token, skip_special_tokens=True), end='', flush=True)

python/infer_torch.py

@@ -0,0 +1,33 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from loguru import logger
+
+
+model_name = "./Qwen2.5-1.5B-Instruct-GPTQ-Int8"
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype="auto",
+    device_map="auto"
+)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+prompt = "请你对 large language model 做一个简短的介绍."
+messages = [
+    {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},
+    {"role": "user", "content": prompt}
+]
+text = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True
+)
+model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+import pdb; pdb.set_trace()
+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]
+logger.debug(response)

python/utils/infer_func.py

@@ -0,0 +1,383 @@
+import torch
+import numpy as np
+from typing import List, Tuple, Optional, Dict
+from pathlib import Path
+from tqdm import tqdm
+from axengine import InferenceSession
+from ml_dtypes import bfloat16
+from transformers import AutoTokenizer, AutoConfig
+import json
+from loguru import logger
+
+
+class KVCacheTools:
+    """
+    k, v cache 的本地保存和加载
+    """
+    def __init__(self, axmodel_num: int, dtype=np.float32):
+        self.axmodel_num = axmodel_num
+        self.dtype = dtype
+
+    def save_kvcache(
+        self,
+        target_dir: str,
+        system_prompt: str,
+        precompute_len: int,
+        k_caches: List[np.ndarray],
+        v_caches: List[np.ndarray],
+        metadata: Optional[Dict] = None
+    ) -> bool:
+        try:
+            target_path = Path(target_dir)
+            target_path.mkdir(parents=True, exist_ok=True)
+
+            for i, (k, v) in enumerate(zip(k_caches, v_caches)):
+                k.astype(self.dtype).tofile(target_path / f"k_cache_{i}.bin")
+                v.astype(self.dtype).tofile(target_path / f"v_cache_{i}.bin")
+
+            config = {
+                "precompute_len": precompute_len,
+                "system_prompt": system_prompt,
+                "axmodel_num": self.axmodel_num,
+                "dtype": str(self.dtype),
+                "metadata": metadata or {},
+            }
+            with open(target_path / "config.json", "w", encoding="utf8") as f:
+                json.dump(config, f, indent=2, ensure_ascii=False)
+
+            return True
+        except Exception as e:
+            print(f"Save failed: {str(e)}")
+            return False
+
+    def load_kvcache(
+        self,
+        cache_dir: str
+    ) -> Tuple[
+        List[np.ndarray], 
+        List[np.ndarray], 
+        str, 
+        int,
+        Dict
+    ]:
+        try:
+            cache_path = Path(cache_dir)
+            k_caches, v_caches = [], []
+
+            with open(cache_path / "config.json") as f:
+                config = json.load(f)
+
+            if config["axmodel_num"] != self.axmodel_num:
+                raise ValueError(
+                    f"Model layer mismatch: "
+                    f"Expected {self.axmodel_num}, got {config['axmodel_num']}"
+                )
+
+            for i in range(self.axmodel_num):
+                k_data = np.fromfile(cache_path / f"k_cache_{i}.bin", dtype=self.dtype).reshape(1, -1, 256)
+                v_data = np.fromfile(cache_path / f"v_cache_{i}.bin", dtype=self.dtype).reshape(1, -1, 256)
+                k_caches.append(k_data)
+                v_caches.append(v_data)
+
+            return (
+                (k_caches, v_caches),
+                config["system_prompt"],
+                config["precompute_len"],
+                config.get("metadata", {})
+            )
+        except Exception as e:
+            print(f"Load failed: {str(e)}")
+            exit()
+
+
+class InferManager:    
+    def __init__(self, hf_model_path: str, axmodel_path: str):
+        self.device = "cpu"
+        self.hf_model_path = hf_model_path
+        self.axmodel_path = axmodel_path
+
+        self.hf_config = AutoConfig.from_pretrained(self.hf_model_path, trust_remote_code=True)
+        self.tokenizer = AutoTokenizer.from_pretrained(self.hf_model_path, trust_remote_code=True, use_fast=False)
+        self.system_prompt = "你的名字叫小智(allen), 你是一个人畜无害的 AI 助手. 深圳市今天(4月1日)阴天, 愚人节, 气温在 14°C 至 19°C 之间, 微风."
+        self.embeds = np.load(f"{self.axmodel_path}/model.embed_tokens.weight.npy")
+
+    def build_system_prompt(self):
+
+        messages = [
+            {"role": "system", "content": self.system_prompt},
+            # {"role": "user", "content": prompt}
+        ]
+        text = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False
+        )
+        self.system_inputs = self.tokenizer([text], return_tensors="pt").to(self.device)
+        self.system_input_ids = self.system_inputs.input_ids[0].cpu().numpy().tolist()
+        self.system_input_embeds = np.take(self.embeds, self.system_input_ids, axis=0)
+        self.system_input_ids_len = len(self.system_input_ids)
+        self.model_inputs = {
+            "input_ids": self.system_input_ids,
+            "input_embeds": self.system_input_embeds,
+            "input_ids_len": self.system_input_ids_len
+        }
+        self.precompute_len = self.system_input_ids_len
+        # logger.info(f"system prompt prompt ids len: {self.system_input_ids_len}")
+
+    def encoder_prompt(self, prompt):
+
+        text = f'<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n'
+        model_inputs = self.tokenizer([text], return_tensors="pt").to(self.device)
+        input_ids = model_inputs.input_ids[0].cpu().numpy().tolist()
+        input_embeds = np.take(self.embeds, input_ids, axis=0)
+        input_ids_len = len(input_ids)
+        # logger.info(f"user prompt token_len: {input_ids_len}")
+
+        model_inputs = {
+            "message": text,
+            "model_inputs": model_inputs,
+            "input_ids": input_ids,
+            "input_embeds": input_embeds,
+            "input_ids_len": input_ids_len
+        }
+        return model_inputs
+
+    def build_kvcache(self, kv_cache_len: int = 2559):
+
+        kv_dim = self.hf_config.hidden_size // self.hf_config.num_attention_heads * self.hf_config.num_key_value_heads
+        self.k_caches = [
+            np.zeros((1, kv_cache_len, kv_dim), dtype=bfloat16)
+            for _ in range(self.hf_config.num_hidden_layers)
+        ]
+        self.v_caches = [
+            np.zeros((1, kv_cache_len, kv_dim), dtype=bfloat16)
+            for _ in range(self.hf_config.num_hidden_layers)
+        ]
+
+    def get_kvcache(self):
+        return [self.k_caches, self.v_caches]
+
+    def update_kvcache(self, update_kv_cache):
+        self.k_caches = update_kv_cache[0]
+        self.v_caches = update_kv_cache[1]
+
+    def get_tokenizer(self):
+        return self.tokenizer
+
+    def get_system_prompt(self):
+        return self.system_prompt
+
+    def set_system_prompt(self, prompt):
+        self.system_prompt = prompt
+
+    def build_infer_model(self, ):
+        self.prefill_decoder_sessins = []
+
+        for i in tqdm(range(self.hf_config.num_hidden_layers), desc="Init InferenceSession"):
+            session = InferenceSession(
+                f"{self.axmodel_path}/qwen2_p128_l{i}_together.axmodel"
+            )
+            self.prefill_decoder_sessins.append(session)
+
+        self.post_process_session = InferenceSession(
+            f"{self.axmodel_path}/qwen2_post.axmodel"
+        )
+        print("The models have been loaded!")
+
+    def get_infer_session(self):
+        return [self.prefill_decoder_sessins, self.post_process_session]
+
+    @staticmethod
+    def _top_p(probs: np.ndarray, p: float) -> np.ndarray:
+        sorted_indices = np.argsort(probs)
+        filtered = probs.copy()
+        cumulative = 0
+        for idx in sorted_indices[::-1]:
+            if cumulative >= p:
+                filtered[idx] = 0
+            cumulative += filtered[idx]
+        return filtered / cumulative
+
+    @staticmethod
+    def _softmax(logits: np.ndarray) -> np.ndarray:
+        logits = logits - logits.max()
+        exp_logits = np.exp(logits)
+        return (exp_logits / np.sum(exp_logits)).astype(np.float64)
+
+    def post_process(self, logits, top_k=1, top_p=0.9, temperature=0.6):
+        logits = logits.astype(np.float32).flatten()
+        candidate_indices = np.argpartition(logits, -top_k)[-top_k:]
+        candidate_logits = logits[candidate_indices] / temperature
+        candidate_probs = self._softmax(candidate_logits)
+        candidate_probs = self._top_p(candidate_probs, top_p)
+        candidate_probs = candidate_probs.astype(np.float64) / candidate_probs.sum()
+        chosen_idx = np.random.multinomial(1, candidate_probs).argmax()
+        next_token = candidate_indices[chosen_idx]
+        return next_token, candidate_indices, candidate_probs
+
+    def gen_slice_indices(self, token_len, prefill=128, expand=128):
+        remaining = max(0, token_len - prefill)
+        extra_blocks = (remaining + expand - 1) // expand
+        return list(range(extra_blocks + 1))
+
+    def prefill(
+        self,
+        model_inputs,
+        slice_len=128,
+        precompute_len=0, # system prompt prefill 的时候, 只能设置为 0
+    ):
+        """
+        Prefill step for chunked inference.
+        """
+        token_ids = model_inputs["input_ids"]
+        token_embeds = model_inputs["input_embeds"]
+        token_len = model_inputs["input_ids_len"]
+
+        seq_len = len(token_ids)
+        slice_indices = [i for i in range(seq_len // slice_len + 1)]
+        print(f"slice_indices: {slice_indices}")
+        # total_prefill_len = (
+        #     slice_len * slice_indices[-1]
+        #     if slice_indices[-1] != 0
+        #     else slice_len
+        # )
+        # slice_indices = self.gen_slice_indices(seq_len)
+        total_prefill_len = slice_len * (slice_indices[-1] + 1)
+        kv_mask_expand_len = 128
+
+        if total_prefill_len > 0:
+            for slice_index in slice_indices:
+                if slice_index == 0:
+                    current_slice_len = slice_len
+                else:
+                    current_slice_len = kv_mask_expand_len
+
+                indices = np.array(
+                    list(
+                        range(
+                            precompute_len + slice_index * slice_len,
+                            precompute_len + (slice_index + 1) * slice_len,
+                        )
+                    ),
+                    np.uint32,
+                ).reshape((1, slice_len))
+                indices[:, min(token_len, slice_len):] = 0
+
+                mask = (
+                    np.zeros((1, slice_len, current_slice_len * slice_index + slice_len))
+                    - 65536
+                )
+                data = np.zeros((1, slice_len, self.hf_config.hidden_size)).astype(bfloat16)
+
+                for i, t in enumerate(
+                    range(
+                        slice_index * slice_len,
+                        (slice_index + 1) * slice_len,
+                    )
+                ):
+                    if t < len(token_ids):
+                        # mask[:, i, 0: slice_index * slice_len + i + 1] = 0
+                        data[:, i : i + 1, :] = (
+                            token_embeds[t]
+                            .reshape((1, 1, self.hf_config.hidden_size))
+                            .astype(bfloat16)
+                        )
+                    if t < len(token_ids) + precompute_len:
+                        mask[:, i, 0: slice_index * slice_len + i + 1] = 0
+
+                if slice_index == slice_indices[-1]:
+                    curlen_procd = token_len - slice_index * slice_len # curlen_procd 是当前处理数据的长度
+                else:
+                    curlen_procd = slice_len
+
+                mask = mask.astype(bfloat16)
+                for i in range(self.hf_config.num_hidden_layers):
+                    input_feed = {
+                        "K_cache": (
+                            self.k_caches[i][:, 0: current_slice_len * slice_index, :]
+                            if slice_index
+                            else np.zeros((1, 1, self.hf_config.hidden_size), dtype=bfloat16)
+                        ),
+                        "V_cache": (
+                            self.v_caches[i][:, 0: current_slice_len * slice_index, :]
+                            if slice_index
+                            else np.zeros((1, 1, self.hf_config.hidden_size), dtype=bfloat16)
+                        ),
+                        "indices": indices,
+                        "input": data,
+                        "mask": mask,
+                    }
+                    outputs = self.prefill_decoder_sessins[i].run(None, input_feed, shape_group=slice_index + 1)
+                    self.k_caches[i][
+                        :,
+                        slice_index
+                        * slice_len + precompute_len : slice_index
+                        * slice_len + curlen_procd + precompute_len,
+                        :,
+                    ] = outputs[0][:, :curlen_procd, :]
+
+                    self.v_caches[i][
+                        :,
+                        slice_index
+                        * slice_len + precompute_len: slice_index
+                        * slice_len + curlen_procd + precompute_len,
+                        :,
+                    ] = outputs[1][:, :curlen_procd, :]
+
+                    data = outputs[2]
+
+                print("slice prefill done", slice_index)
+        else:
+            print("No prefill needed.")
+        # return "Calculated the kv cache of the system prompt."
+        return (self.k_caches, self.v_caches)
+
+    def decode(
+        self,
+        token_ids,
+        prefill_len=128,
+        slice_len=128
+    ):
+        token_len = len(token_ids)
+        # set to decoder
+        print("answer: >> ", end='', flush=True)
+        kv_cache_len = 2559
+        mask = np.zeros((1, 1, kv_cache_len + 1), dtype=np.float32).astype(bfloat16)
+        mask[:, :, :kv_cache_len] -= 65536
+        if prefill_len > 0:
+            mask[:, :, :token_len + self.precompute_len] = 0
+
+        for start_indice in range(kv_cache_len):
+            if self.precompute_len > 0 and start_indice < self.precompute_len:
+                continue
+            next_token = token_ids[start_indice - self.precompute_len]
+            indices = np.array([start_indice], np.uint32).reshape((1, 1))
+            data = self.embeds[next_token, :].reshape((1, 1, self.hf_config.hidden_size)).astype(bfloat16)
+            for i in range(self.hf_config.num_hidden_layers):
+                input_feed = {
+                    "K_cache": self.k_caches[i],
+                    "V_cache": self.v_caches[i],
+                    "indices": indices,
+                    "input": data,
+                    "mask": mask,
+                }
+                outputs = self.prefill_decoder_sessins[i].run(None, input_feed, shape_group=0)
+                self.k_caches[i][:, start_indice, :] = outputs[0][:, :, :]
+                self.v_caches[i][:, start_indice, :] = outputs[1][:, :, :]
+                data = outputs[2]
+            mask[..., start_indice] = 0
+            if start_indice < token_len + self.precompute_len - 1:
+                pass
+            else:
+                post_out = self.post_process_session.run(None, {"input": data})[0]
+                next_token, posssible_tokens, possible_soft = self.post_process(post_out)
+                token_ids.append(next_token)
+                print(self.tokenizer.decode(next_token, skip_special_tokens=True), end='', flush=True)
+
+                if next_token == self.tokenizer.eos_token_id and start_indice > token_len + self.precompute_len:
+                    # print("\n>> HINT: The next_token encountered EOS token, generation completed.")
+                    break
+        print("\n")
+        self.precompute_len = len(token_ids) + self.precompute_len - 1
+        return self.tokenizer.decode(token_ids[self.precompute_len - 1:], skip_special_tokens=True)
+