Delete .github/workflows/distribute-v100.yml

* change_save_meta_in_zcc_crossponding_to_paddle

* delete

* fix

* change_save_meta_in_zcc_crossponding_to_paddle and fix args

* fix deadcode

---------

Co-authored-by: liufengwei02 <liufengwei02@baidu.com>

.copyright.hook

@@ -0,0 +1,134 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import argparse
+import io
+import re
+import sys
+import os
+import datetime
+
+COPYRIGHT = '''Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.'''
+
+def _generate_copyright(comment_mark):
+    copyright=COPYRIGHT.split(os.linesep)
+    header = copyright[0].rstrip()
+
+    p = re.search('(\d{4})', header).group(0)
+    now = datetime.datetime.now()
+
+    header = header.replace(p,str(now.year))
+
+    ans=[comment_mark + " " + header + os.linesep]
+    for idx, line in enumerate(copyright[1:]):
+        ans.append(comment_mark + " " + line.rstrip() + os.linesep)
+
+    return ans
+
+def _get_comment_mark(path):
+    lang_type=re.compile(r"\.(py|sh)$")
+    if lang_type.search(path) is not None:
+        return "#"
+
+    lang_type=re.compile(r"\.(h|c|hpp|cc|cpp|cu|go|cuh|proto)$")
+    if lang_type.search(path) is not None:
+        return "//"
+
+    return None
+
+
+RE_ENCODE = re.compile(r"^[ \t\v]*#.*?coding[:=]", re.IGNORECASE)
+RE_COPYRIGHT = re.compile(r".*Copyright( \(c\))* \d{4}", re.IGNORECASE)
+RE_SHEBANG = re.compile(r"^[ \t\v]*#[ \t]?\!")
+
+def _check_copyright(path):
+    head=[]
+    try:
+        with open(path, encoding="utf-8") as f:
+            head = [next(f) for x in range(4)]
+    except StopIteration:
+        pass
+
+    for idx, line in enumerate(head):
+        if RE_COPYRIGHT.search(line) is not None:
+            return True
+
+    return False
+
+def generate_copyright(path, comment_mark):
+    original_contents = io.open(path, encoding="utf-8").readlines()
+    head = original_contents[0:4]
+
+    insert_line_no=0
+    for i, line in enumerate(head):
+        if RE_ENCODE.search(line) or RE_SHEBANG.search(line):
+            insert_line_no=i+1
+
+    copyright = _generate_copyright(comment_mark)
+    if insert_line_no == 0:
+        new_contents = copyright
+        if len(original_contents) > 0 and len(original_contents[0].strip()) != 0:
+            new_contents.append(os.linesep)
+        new_contents.extend(original_contents)
+    else:
+        new_contents=original_contents[0:insert_line_no]
+        new_contents.append(os.linesep)
+        new_contents.extend(copyright)
+        if len(original_contents) > insert_line_no and len(original_contents[insert_line_no].strip()) != 0:
+            new_contents.append(os.linesep)
+        new_contents.extend(original_contents[insert_line_no:])
+    new_contents="".join(new_contents)
+
+    with io.open(path, 'w') as output_file:
+        output_file.write(new_contents)
+
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(
+        description='Checker for copyright declaration.')
+    parser.add_argument('filenames', nargs='*', help='Filenames to check')
+    args = parser.parse_args(argv)
+
+    retv = 0
+    for path in args.filenames:
+        comment_mark = _get_comment_mark(path)
+        if comment_mark is None:
+            print("warning:Unsupported file", path, file=sys.stderr)
+            continue
+
+        if _check_copyright(path):
+            continue
+
+        generate_copyright(path, comment_mark)
+
+
+if __name__ == '__main__':
+    exit(main())

.flake8

@@ -0,0 +1,7 @@
+[flake8]
+ignore = E203, E402, E501, E731, E741, W503, W605, E722
+max-line-length = 119
+
+# E402: module level import not at top of file
+per-file-ignores =
+    __init__.py:F401,F403,E402

.github/CODE_OF_CONDUCT.md

@@ -0,0 +1,89 @@
+**简体中文**🀄 | [English🌎](./CODE_OF_CONDUCT_en.md)
+
+# 贡献者公约
+
+## 我们的承诺
+
+身为社区成员、贡献者和领袖，我们承诺使社区参与者不受骚扰，无论其年龄、体型、可见或不可见的缺陷、族裔、性征、性别认同和表达、经验水平、教育程度、社会与经济地位、国籍、相貌、种族、种姓、肤色、宗教信仰、性倾向或性取向如何。
+
+我们承诺以有助于建立开放、友善、多样化、包容、健康社区的方式行事和互动。
+
+## 我们的准则
+
+有助于为我们的社区创造积极环境的行为例子包括但不限于：
+
+* 表现出对他人的同情和善意
+* 尊重不同的主张、观点和感受
+* 提出和大方接受建设性意见
+* 承担责任并向受我们错误影响的人道歉
+* 注重社区共同诉求，而非个人得失
+
+不当行为例子包括：
+
+* 使用情色化的语言或图像，及性引诱或挑逗
+* 嘲弄、侮辱或诋毁性评论，以及人身或政治攻击
+* 公开或私下的骚扰行为
+* 未经他人明确许可，公布他人的私人信息，如物理或电子邮件地址
+* 其他有理由认定为违反职业操守的不当行为
+
+## 责任和权力
+
+社区领袖有责任解释和落实我们所认可的行为准则，并妥善公正地对他们认为不当、威胁、冒犯或有害的任何行为采取纠正措施。
+
+社区领导有权力和责任删除、编辑或拒绝或拒绝与本行为准则不相符的评论（comment）、提交（commits）、代码、维基（wiki）编辑、议题（issues）或其他贡献，并在适当时机知采取措施的理由。
+
+## 适用范围
+
+本行为准则适用于所有社区场合，也适用于在公共场所代表社区时的个人。
+
+代表社区的情形包括使用官方电子邮件地址、通过官方社交媒体帐户发帖或在线上或线下活动中担任指定代表。
+
+## 监督
+
+辱骂、骚扰或其他不可接受的行为可通过 paddlenlp@baidu.com 向负责监督的社区领袖报告。
+所有投诉都将得到及时和公平的审查和调查。
+
+所有社区领袖都有义务尊重任何事件报告者的隐私和安全。
+
+## 处理方针
+
+社区领袖将遵循下列社区处理方针来明确他们所认定违反本行为准则的行为的处理方式：
+
+### 1. 纠正
+
+**社区影响**：使用不恰当的语言或其他在社区中被认定为不符合职业道德或不受欢迎的行为。
+
+**处理意见**：由社区领袖发出非公开的书面警告，明确说明违规行为的性质，并解释举止如何不妥。或将要求公开道歉。
+
+### 2. 警告
+
+**社区影响**：单个或一系列违规行为。
+
+**处理意见**：警告并对连续性行为进行处理。在指定时间内，不得与相关人员互动，包括主动与行为准则执行者互动。这包括避免在社区场所和外部渠道中的互动。违反这些条款可能会导致临时或永久封禁。
+
+### 3. 临时封禁
+
+**社区影响**: 严重违反社区准则，包括持续的不当行为。
+
+**处理意见**: 在指定时间内，暂时禁止与社区进行任何形式的互动或公开交流。在此期间，不得与相关人员进行公开或私下互动，包括主动与行为准则执行者互动。违反这些条款可能会导致永久封禁。
+
+### 4. 永久封禁
+
+**社区影响**：行为模式表现出违反社区准则，包括持续的不当行为、骚扰个人或攻击或贬低某个类别的个体。
+
+**处理意见**：永久禁止在社区内进行任何形式的公开互动。
+
+## 参见
+
+本行为准则改编自 [Contributor Covenant][homepage] 2.1 版, 参见 [https://www.contributor-covenant.org/version/2/1/code_of_conduct.html][v2.1]。
+
+社区处理方针灵感来源于 [Mozilla's code of conduct enforcement ladder][Mozilla CoC]。
+
+有关本行为准则的常见问题的答案，参见 [https://www.contributor-covenant.org/faq][FAQ]。
+其他语言翻译参见 [https://www.contributor-covenant.org/translations][translations]。
+
+[homepage]: https://www.contributor-covenant.org
+[v2.1]: https://www.contributor-covenant.org/version/2/1/code_of_conduct.html
+[Mozilla CoC]: https://github.com/mozilla/diversity
+[FAQ]: https://www.contributor-covenant.org/faq
+[translations]: https://www.contributor-covenant.org/translations

.github/CODE_OF_CONDUCT_en.md

@@ -0,0 +1,134 @@
+[简体中文🀄](./CODE_OF_CONDUCT.md) |  **English**🌎
+
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, caste, color, religion, or sexual
+identity and orientation.
+
+We pledge to act and interact in ways that contribute to an open, welcoming,
+diverse, inclusive, and healthy community.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment for our
+community include:
+
+* Demonstrating empathy and kindness toward other people
+* Being respectful of differing opinions, viewpoints, and experiences
+* Giving and gracefully accepting constructive feedback
+* Accepting responsibility and apologizing to those affected by our mistakes,
+  and learning from the experience
+* Focusing on what is best not just for us as individuals, but for the overall
+  community
+
+Examples of unacceptable behavior include:
+
+* The use of sexualized language or imagery, and sexual attention or advances of
+  any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or email address,
+  without their explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Enforcement Responsibilities
+
+Community leaders are responsible for clarifying and enforcing our standards of
+acceptable behavior and will take appropriate and fair corrective action in
+response to any behavior that they deem inappropriate, threatening, offensive,
+or harmful.
+
+Community leaders have the right and responsibility to remove, edit, or reject
+comments, commits, code, wiki edits, issues, and other contributions that are
+not aligned to this Code of Conduct, and will communicate reasons for moderation
+decisions when appropriate.
+
+## Scope
+
+This Code of Conduct applies within all community spaces, and also applies when
+an individual is officially representing the community in public spaces.
+Examples of representing our community include using an official e-mail address,
+posting via an official social media account, or acting as an appointed
+representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the community leaders responsible for enforcement at
+paddlenlp@baidu.com.
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+**Community Impact**: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+**Consequence**: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+**Community Impact**: A violation through a single incident or series of
+actions.
+
+**Consequence**: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or permanent
+ban.
+
+### 3. Temporary Ban
+
+**Community Impact**: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+**Consequence**: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+**Community Impact**: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior, harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+**Consequence**: A permanent ban from any sort of public interaction within the
+community.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 2.1, available at
+[https://www.contributor-covenant.org/version/2/1/code_of_conduct.html][v2.1].
+
+Community Impact Guidelines were inspired by
+[Mozilla's code of conduct enforcement ladder][Mozilla CoC].
+
+For answers to common questions about this code of conduct, see the FAQ at
+[https://www.contributor-covenant.org/faq][FAQ]. Translations are available at
+[https://www.contributor-covenant.org/translations][translations].
+
+[homepage]: https://www.contributor-covenant.org
+[v2.1]: https://www.contributor-covenant.org/version/2/1/code_of_conduct.html
+[Mozilla CoC]: https://github.com/mozilla/diversity
+[FAQ]: https://www.contributor-covenant.org/faq
+[translations]: https://www.contributor-covenant.org/translations

.github/CONTRIBUTING_en.md

@@ -0,0 +1,154 @@
+[简体中文🀄](../CONTRIBUTING.md) |  **English**🌎
+
+# Contributing to PaddleNLP
+
+We highly welcome and value your contributions to `PaddleNLP`. The first step to start your contribution is to sign the [PaddlePaddle Contributor License Agreement](https://cla-assistant.io/PaddlePaddle/PaddleNLP).
+
+This document explains our workflow and work style:
+
+## Finding out what to work on Workflow
+
+## Development Workflow
+
+PaddleNLP uses the [Git branching model](http://nvie.com/posts/a-successful-git-branching-model/).  The following steps guide usual contributions.
+
+#### 1. Fork
+
+   Our development community has been growing fastly; it doesn't make sense for everyone to write into the official repo.  So, please file Pull Requests from your fork.  To make a fork,  just head over to the GitHub page and click the ["Fork" button](https://help.github.com/articles/fork-a-repo/).
+
+#### 2. Clone
+
+   To make a copy of your fork to your local computers, please run
+
+   ```bash
+   git clone https://github.com/<your-github-account>/PaddleNLP
+   cd PaddleNLP
+   ```
+
+#### 3. Create the local feature branch
+
+   For daily works like adding a new feature or fixing a bug, please open your feature branch before coding:
+
+   ```bash
+   git checkout -b my-cool-feature
+   ```
+
+#### 4. Set up the development environment
+
+   Before you start coding, you need to setup the development environment. We highly recommend doing all your development in a virtual environment such as
+   [venv](https://docs.python.org/3/library/venv.html) or [conda](https://docs.conda.io/en/latest/). After you setup and activated your virtual environment,
+   run the following command:
+
+   ```bash
+   make install
+   ```
+
+   This will setup all the dependencies of `PaddleNLP` as well as the [`pre-commit`](http://pre-commit.com/) tool.
+
+   If you are working on the `examples` or `applications` module and require importing from `PaddleNLP`, make sure you install `PaddleNLP` in editable mode.
+   If `PaddleNLP` is already installed in the virtual environment, remove it with `pip uninstall paddlenlp` before reinstalling it in editable mode with
+   `pip install -e .`
+
+#### 5. Develop
+
+   As you develop your new exciting feature, keep in mind that it should be covered by unit tests. All of our unit tests can be found under the `tests` directory.
+   You can either modify existing unit test to cover the new feature, or create a new test from scratch.
+   As you finish up the your code, you should make sure the test suite passes. You can run the tests impacted by your changes like this:
+
+   ```bash
+   pytest tests/<test_to_run>.py
+   ```
+
+#### 6. Commit
+
+   We utilizes [`pre-commit`](http://pre-commit.com/) (with [black](https://black.readthedocs.io/en/stable/), [isort](https://pycqa.github.io/isort/) and
+   [flake8](https://flake8.pycqa.org/en/latest/) under the hood) to check the style of code and documentation in every commit. When you run run `git commit`, you will see
+   something like the following:
+
+   ```
+    ➜  (my-virtual-env) git commit -m "commiting my cool feature"
+    black....................................................................Passed
+    isort....................................................................Passed
+    flake8...................................................................Passed
+    check for merge conflicts................................................Passed
+    check for broken symlinks............................(no files to check)Skipped
+    detect private key.......................................................Passed
+    fix end of files.....................................(no files to check)Skipped
+    trim trailing whitespace.............................(no files to check)Skipped
+    CRLF end-lines checker...............................(no files to check)Skipped
+    CRLF end-lines remover...............................(no files to check)Skipped
+    No-tabs checker......................................(no files to check)Skipped
+    Tabs remover.........................................(no files to check)Skipped
+    copyright_checker........................................................Passed
+   ```
+
+   But most of the time things don't go so smoothly. When your code or documentation doesn't meet the standard, the `pre-commit` check will fail.
+   ```
+    ➜  (my-virtual-env) git commit -m "commiting my cool feature"
+    black....................................................................Passed
+    isort....................................................................Failed
+    - hook id: isort
+    - files were modified by this hook
+
+    Fixing examples/information_extraction/waybill_ie/run_ernie_crf.py
+
+    flake8...................................................................Passed
+    check for merge conflicts................................................Passed
+    check for broken symlinks............................(no files to check)Skipped
+    detect private key.......................................................Passed
+    fix end of files.....................................(no files to check)Skipped
+    trim trailing whitespace.............................(no files to check)Skipped
+    CRLF end-lines checker...............................(no files to check)Skipped
+    CRLF end-lines remover...............................(no files to check)Skipped
+    No-tabs checker......................................(no files to check)Skipped
+    Tabs remover.........................................(no files to check)Skipped
+    copyright_checker........................................................Passed
+   ```
+
+   But **don't panic**!
+   Our tooling will fix most of the style errors automatically. Some errors will need to be addressed manually. Fortunately, the error messages are straight forward and
+   the errors are usually simple to fix. After addressing the errors, you can run `git add <files>` and `git commit` again, which will trigger `pre-commit` again.
+   Once the `pre-commit` checks pass, you are ready to push the code.
+
+   [Google][http://google.com/] or [StackOverflow](https://stackoverflow.com/) are great tools to help you understand the code style errors.
+   Don't worry if you still can't figure it out. You can commit with `git commit -m "style error" --no-verify` and we are happy to help you once you create a Pull Request.
+
+#### 7. Keep pulling
+
+   An experienced Git user pulls from the official repo often -- daily or even hourly, so they notice conflicts with others work early, and it's easier to resolve smaller conflicts.
+
+   ```bash
+   git remote add upstream https://github.com/PaddlePaddle/PaddleNLP
+   git pull upstream develop
+   ```
+
+#### 8. Push and file a pull request
+
+   You can "push" your local work into your forked repo:
+
+   ```bash
+   git push origin my-cool-stuff
+   ```
+
+   The push allows you to create a pull request, requesting owners of this [official repo](https://github.com/PaddlePaddle/PaddleNLP) to pull your change into the official one.
+
+   To create a pull request, please follow [these steps](https://help.github.com/articles/creating-a-pull-request/).
+
+#### 9. Delete local and remote branches
+
+   To keep your local workspace and your fork clean, you might want to remove merged branches:
+
+   ```bash
+   git push origin my-cool-stuff
+   git checkout develop
+   git pull upstream develop
+   git branch -d my-cool-stuff
+   ```
+
+## Code Review
+
+-  Please feel free to ping your reviewers by @-mentioning the in the Pull Request.  Please do this after your pull request passes the CI.
+
+- Please answer reviewers' every comment.  If you are to follow the comment, please write "Done"; Otherwise, please start a discussion under the comment.
+
+- If you don't want your reviewers to get overwhelmed by email notifications, you might reply their comments by [in a batch](https://help.github.com/articles/reviewing-proposed-changes-in-a-pull-request/).

.github/ISSUE_TEMPLATE/ask-question.yml

@@ -0,0 +1,23 @@
+name: 🐛 Ask Question
+description: 请描述您使用PaddleNLP时遇到的问题
+title: "[Question]: "
+labels: 
+  - question
+body:
+- type: markdown
+  attributes:
+    value: >
+      #### 你可以在这里提出一个使用/咨询问题，提问之前请确保：
+      
+      - 1）已经百度/谷歌搜索过你的问题，但是没有找到解答；
+      
+      - 2）已经在官网查询过[API文档](https://www.paddlepaddle.org.cn/documentation/docs/zh/api/index_cn.html)与[FAQ](https://www.paddlepaddle.org.cn/documentation/docs/zh/faq/index_cn.html)，但是没有找到解答；
+      
+      - 3）已经在[历史issue](https://github.com/PaddlePaddle/Paddle/issues)中搜索过，没有找到同类issue或issue未被解答。
+      
+- type: textarea
+  id: question
+  attributes:
+    label: 请提出你的问题
+  validations:
+    required: true

.github/ISSUE_TEMPLATE/bug-report.yml

@@ -0,0 +1,45 @@
+name: 🐛 Bug Report
+description: PaddleNLP问题反馈
+title: "[Bug]: "
+labels: bug
+body: 
+  - type: textarea
+    id: environment
+    attributes:
+      label: 软件环境
+      description: |
+        请使用以下命令给出您本地Paddle相关包信息
+          ```sh
+          pip list | grep paddle
+          
+          ```
+      value: |
+          - paddlepaddle:
+          - paddlepaddle-gpu: 
+          - paddlenlp: 
+      render: Markdown
+    validations:
+      required: true
+  - type: checkboxes
+    id: dumplicated-problem
+    attributes:
+      label: 重复问题
+      description: 是否已在issues中搜索相关问题
+      options:
+      - label: I have searched the existing issues
+        required: true
+  - type: textarea
+    id: descripton
+    attributes:
+      label: 错误描述
+      description: 给出错误详细描述，以便能够更好的追踪相关问题
+      render: Markdown
+    validations:
+      required: true
+  - type: textarea
+    id: mvp-code
+    attributes:
+      label: 稳定复现步骤 & 代码
+      description: 请给出稳定复现该问题的步骤 & 代码，以便相关人员能够快速定位到具体问题。
+    validations:
+      required: true

.github/ISSUE_TEMPLATE/docs-report.yml

@@ -0,0 +1,32 @@
+name: 🐛 Docs Report
+description: PaddleNLP文档反馈
+title: "[Docs]: "
+labels: 
+  - documentation
+
+body: 
+  - type: textarea
+    id: environment
+    attributes:
+      label: 软件环境
+      description: |
+        请使用以下命令给出您本地Paddle相关包信息
+          ```sh
+          pip list | grep paddle
+          
+          ```
+      value: |
+          - paddlepaddle:
+          - paddlepaddle-gpu: 
+          - paddlenlp: 
+      render: Markdown
+    validations:
+      required: true
+  - type: textarea
+    id: description
+    attributes:
+      label: 详细描述
+      description: 请详细描述您想要反馈的具体问题
+      render: Markdown
+    validations:
+      required: true

.github/ISSUE_TEMPLATE/feature-request.yml

@@ -0,0 +1,30 @@
+name: "\U0001F680 Feature request"
+description: 请详细描述您所需功能
+labels: [ "feature" ]
+body:
+  - type: textarea
+    id: feature-request
+    validations:
+      required: true
+    attributes:
+      label: Feature request
+      description: |
+        对特性提案的清晰而简明的描述。如果论文和代码存在，请提供链接。 
+
+  - type: textarea
+    id: motivation
+    validations:
+      required: true
+    attributes:
+      label: Motivation
+      description: |
+        请概述这项建议的动机。您的特性要求与问题有关吗?
+        
+  - type: textarea
+    id: contribution
+    validations:
+      required: true
+    attributes:
+      label: Your contribution
+      description: |
+        Is there any way that you could help, e.g. by submitting a PR? Make sure to read the CONTRIBUTING.MD [readme](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md)

.github/ISSUE_TEMPLATE/new-model.yaml

@@ -0,0 +1,28 @@
+name: "\U0001F31F 添加新模型"
+description: 请为新模型提交一份说明
+labels: [ "New model" ]
+
+body:
+  - type: textarea
+    id: description-request
+    validations:
+      required: true
+    attributes:
+      label: 简要描述
+      description: |
+        请简要描述模型的类型、解决的问题等。
+
+  - type: checkboxes
+    id: information-tasks
+    attributes:
+      label: 是否已开源
+      options:
+        - label: 已开源
+        - label: 未开源
+
+  - type: textarea
+    id: additional-info
+    attributes:
+      label: 模型详细信息
+      description: |
+        请给出新模型相关信息，如论文地址、现存代码地址等。

.github/ISSUE_TEMPLATE/others.yml

@@ -0,0 +1,23 @@
+name: 🧩 其他 Others
+description: 提出其他问题。
+labels: [others]
+
+body:
+- type: markdown
+  attributes:
+    value: >
+      #### 你可以在这里提出任何前面几类模板不适用的问题，包括但不限于：优化性建议、框架使用体验反馈、版本兼容性问题、报错信息不清楚等。
+
+- type: textarea
+  id: others
+  attributes:
+    label: 问题描述
+  validations:
+    required: true
+    
+- type: markdown
+  attributes:
+    value: >
+      感谢你的贡献 🎉！
+
+

.github/PULL_REQUEST_TEMPLATE.md

@@ -1,4 +1,18 @@
 <!-- Demo: https://github.com/PaddlePaddle/PaddleNLP/pull/26 -->
+#### Before submitting
+
+- [ ] Lint code. If there are lint issues, please format the code first.
+
+```shell
+# Install and register `pre-commit` in the project folder
+pip install pre-commit && pre-commit install
+
+# Process previous code files separately
+pre-commit run --file XXXX.py
+```
+
+- [ ] Add test cases into `tests` folder. If there are codecov issues, please add tests cases first.
+
 ### PR types
 <!-- One of [ New features | Bug fixes | Function optimization | Performance optimization | Breaking changes | Others ] -->
 

.github/actions/check-bypass/action.yml

@@ -0,0 +1,41 @@
+name: "Check bypass"
+description: "A custom action to encapsulate PFCCLab/ci-bypass"
+inputs:
+  github-token:
+    description: "GitHub token"
+    required: true
+  workflow-name:
+    description: "Workflow name"
+    required: true
+outputs:
+  can-skip:
+    description: "Whether the workflow can be skipped."
+    value: ${{ steps.check-bypass.outputs.can-skip }}
+
+runs:
+  using: "composite"
+  steps:
+    - id: check-bypass
+      name: Check Bypass
+      env:
+        CI_TEAM_MEMBERS: '["tianshuo78520a", "swgu98", "risemeup1", "XieYunshen","luotao1","From00"]'
+      uses: PFCCLab/ci-bypass@v1
+      with:
+        github-token: ${{ inputs.github-token }}
+        non-pull-request-event-strategy: 'never-skipped'
+        type: 'composite'
+        composite-rule: |
+          {
+            "any": [
+              {
+                "type": "labeled",
+                "label": ["skip-ci: ${{ inputs.workflow-name }}", "skip-ci: all"],
+                "username": ${{ env.CI_TEAM_MEMBERS }}
+              },
+              {
+                "type": "commented",
+                "comment-pattern": [".*/skip-ci ${{ inputs.workflow-name }}.*", ".*/skip-ci all.*"],
+                "username": ${{ env.CI_TEAM_MEMBERS }}
+              }
+            ]
+          }

.github/actions/rerun-workflow/action.yml

@@ -0,0 +1,30 @@
+name: 'Rerun Workflow'
+description: 'Re-run GitHub Actions workflow for a given Pull Request'
+inputs:
+  GITHUB_TOKEN:
+    description: 'GitHub token with repo scope'
+    required: true
+  OWNER:
+    description: 'Repository owner'
+    required: true
+  REPO:
+    description: 'Repository name'
+    required: true
+  PR_ID:
+    description: 'Pull Request ID'
+    required: true
+  JOB_NAME:
+    description: 'Job name to rerun'
+    required: true
+
+runs:
+  using: 'composite'
+  steps:
+    - run: bash ./.github/actions/rerun-workflow/rerun.sh
+      shell: bash
+      env:
+        GITHUB_TOKEN: ${{ inputs.GITHUB_TOKEN }}
+        OWNER: ${{ inputs.OWNER }}
+        REPO: ${{ inputs.REPO }}
+        PR_ID: ${{ inputs.PR_ID }}
+        JOB_NAME: ${{ inputs.JOB_NAME }}

.github/actions/rerun-workflow/rerun.sh

@@ -0,0 +1,77 @@
+# Copyright (c) 2025 PaddleNLP Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -e
+
+COMMIT_SHA=$(curl -s -H "Authorization: token $GITHUB_TOKEN" \
+  "https://api.github.com/repos/$OWNER/$REPO/pulls/$PR_ID" | jq -r '.head.sha')
+
+echo "Commit SHA: $COMMIT_SHA"
+
+response=$(curl -s -H "Authorization: token $GITHUB_TOKEN" \
+  "https://api.github.com/repos/$OWNER/$REPO/actions/runs?head_sha=$COMMIT_SHA&per_page=100")
+
+echo "Response: $response"
+
+run_ids=$(echo "$response" | jq -r '.workflow_runs[].id')
+
+if [ -n "$run_ids" ]; then
+  echo "Found run_ids for commit $COMMIT_SHA: $run_ids"
+
+  for run_id in $run_ids; do
+    if [ "$JOB_NAME" = "all-failed" ]; then
+      echo "Rerunning all failed jobs for run_id: $run_id"
+
+      rerun_response=$(curl -X POST -s -w "%{http_code}" -o /dev/null \
+        -H "Accept: application/vnd.github.v3+json" \
+        -H "Authorization: Bearer $GITHUB_TOKEN" \
+        "https://api.github.com/repos/$OWNER/$REPO/actions/runs/$run_id/rerun-failed-jobs")
+      if [ "$rerun_response" -eq 201 ]; then
+        echo "Successfully requested rerun for all blocked jobs in run_id: $run_id"
+      else
+        echo "Failed to request rerun for run_id: $run_id with status code $rerun_response"
+      fi
+
+    else
+      jobs_response=$(curl -s -H "Authorization: token $GITHUB_TOKEN" \
+      "https://api.github.com/repos/$OWNER/$REPO/actions/runs/$run_id/jobs")
+
+      echo "Jobs Response for run_id $run_id: $jobs_response"
+
+      # if [[ "$JOB_NAME" == *"bypass"* ]]; then
+        block_jobs=$(echo "$jobs_response" | jq -r --arg job_name "$JOB_NAME" \
+        '.jobs[] | select(.name == $job_name) | .id')
+      # else
+      #   block_jobs=$(echo "$jobs_response" | jq -r --arg job_name "$JOB_NAME" \
+      #   '.jobs[] | select(.name == $job_name and .conclusion != "success") | .id')
+      # fi
+
+      if [ -n "$block_jobs" ]; then
+        echo "Found block jobs for run_id $run_id: $block_jobs"
+
+        for job_id in $block_jobs; do
+          echo "Rerunning job_id: $job_id"
+          curl -X POST -H "Accept: application/vnd.github.v3+json" \
+            -H "Authorization: token $GITHUB_TOKEN" \
+            "https://api.github.com/repos/$OWNER/$REPO/actions/jobs/$job_id/rerun"
+        done
+      else
+        echo "No block jobs found for run_id $run_id with name $JOB_NAME."
+      fi
+    fi
+  done
+else
+  echo "No matching workflow runs found for commit $COMMIT_SHA."
+  exit 1
+fi

.github/codecov.yml

@@ -0,0 +1,13 @@
+codecov:
+  notify:
+    wait_for_ci: false
+
+coverage:
+  status:
+    project:
+      default: 
+        target: 58% # overall project Coverage 
+        threshold: 1% # Allow the coverage to drop by 1%, and posting a success status.
+    patch:
+      default: 
+        target: 80% # lines adjusted  Coverage < 80%  CI will fail

.github/workflows/PaddleFormers-PR-Check.yml

@@ -0,0 +1,47 @@
+name: PaddleFormers PR Check
+
+on:
+  pull_request:
+    types: [opened, synchronize]
+
+env:
+  BRANCH: ${{ github.base_ref }}
+  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+  PR_ID: ${{ github.event.pull_request.number }}
+  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
+  token: ${{ vars.ACTION_GITHUB_TOKEN }}
+
+jobs:
+  check-approvers:
+    name: Check approval
+    runs-on: 
+      group: APPROVAL
+    steps:
+      - name: Cleanup
+        run: |
+          rm -rf * .[^.]*
+      - name: Update paddle
+        run: |
+          wget -q --no-proxy https://xly-devops.bj.bcebos.com/PaddleTest/PaddleNLP/PaddleNLP-develop.tar.gz --no-check-certificate
+          tar zxf PaddleNLP-develop.tar.gz --strip-components=1 >/dev/null
+          rm -rf PaddleNLP-develop.tar.gz >/dev/null
+          git fetch origin pull/${PR_ID}/head
+          git checkout -b origin_pr FETCH_HEAD
+          git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
+          git fetch upstream
+          git checkout -b test_pr upstream/${BRANCH}
+          git merge --no-edit origin_pr
+          git log --pretty=oneline -10
+
+      - name: Check bypass
+        id: check-bypass
+        uses: ./.github/actions/check-bypass
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          workflow-name: approval
+        
+      - name: Display Required Approvers
+        if: steps.check-bypass.outputs.can-skip != 'true'
+        run: |
+          cd scripts/ci_approval
+          bash run_ci_approval_cherry-pick.sh

.github/workflows/approval.yml

@@ -0,0 +1,47 @@
+name: Approval
+
+on:
+  pull_request:
+    types: [opened, synchronize]
+
+env:
+  BRANCH: ${{ github.base_ref }}
+  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+  PR_ID: ${{ github.event.pull_request.number }}
+  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
+  token: ${{ vars.ACTION_GITHUB_TOKEN }}
+
+jobs:
+  check-approvers:
+    name: Check approval
+    runs-on: 
+      group: APPROVAL
+    steps:
+      - name: Cleanup
+        run: |
+          rm -rf * .[^.]*
+      - name: Update paddle
+        run: |
+          wget -q --no-proxy https://xly-devops.bj.bcebos.com/PaddleTest/PaddleNLP/PaddleNLP-develop.tar.gz --no-check-certificate
+          tar zxf PaddleNLP-develop.tar.gz --strip-components=1 >/dev/null
+          rm -rf PaddleNLP-develop.tar.gz >/dev/null
+          git fetch origin pull/${PR_ID}/head
+          git checkout -b origin_pr FETCH_HEAD
+          git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
+          git fetch upstream
+          git checkout -b test_pr upstream/${BRANCH}
+          git merge --no-edit origin_pr
+          git log --pretty=oneline -10
+
+      - name: Check bypass
+        id: check-bypass
+        uses: ./.github/actions/check-bypass
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          workflow-name: approval
+        
+      - name: Display Required Approvers
+        if: steps.check-bypass.outputs.can-skip != 'true'
+        run: |
+          cd scripts/ci_approval
+          bash -x run_ci_approval.sh

.github/workflows/check-bypass.yml

@@ -0,0 +1,52 @@
+on:
+  workflow_call:
+    inputs:
+      workflow-name:
+        required: true
+        type: string
+    secrets:
+      github-token:
+        required: true
+    outputs:
+      can-skip:
+        description: "Whether the workflow can be skipped."
+        value: ${{ jobs.check-bypass.outputs.can-skip }}
+
+jobs:
+  check-bypass:
+    name: Check bypass
+    runs-on:
+      group: APPROVAL
+    permissions:
+      contents: read
+    env:
+      CI_TEAM_MEMBERS: '["tianshuo78520a", "swgu98", "risemeup1" , "XieYunshen","luotao1","From00"]'
+    outputs:
+      can-skip: ${{ steps.check-bypass.outputs.can-skip }}
+    steps:
+      - name: Cleanup
+        run: |
+          rm -rf * .[^.]*
+
+      - id: check-bypass
+        name: Check Bypass
+        uses: PFCCLab/ci-bypass@v1
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          non-pull-request-event-strategy: 'never-skipped'
+          type: 'composite'
+          composite-rule: |
+            {
+              "any": [
+                {
+                  "type": "labeled",
+                  "label": ["skip-ci: ${{ inputs.workflow-name }}", "skip-ci: all"],
+                  "username": ${{ env.CI_TEAM_MEMBERS }}
+                },
+                {
+                  "type": "commented",
+                  "comment-pattern": [".*/skip-ci ${{ inputs.workflow-name }}.*", ".*/skip-ci all.*"],
+                  "username": ${{ env.CI_TEAM_MEMBERS }}
+                }
+              ]
+            }

.github/workflows/lint.yml

@@ -0,0 +1,113 @@
+name: Codestyle Check
+
+on: [push, pull_request]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
+  cancel-in-progress: true
+
+env:
+  PR_ID: ${{ github.event.pull_request.number }}
+  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
+  BRANCH: ${{ github.event.pull_request.base.ref }}
+  TASK: PaddleNLP-CI-Lint-${{ github.event.pull_request.number }}
+
+jobs:
+  check-bypass:
+    name: Check bypass
+    uses: ./.github/workflows/check-bypass.yml
+    with:
+      workflow-name: 'lint'
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+
+  Lint:
+    name: Lint
+    needs: check-bypass
+    if: ${{ needs.check-bypass.outputs.can-skip != 'true' }}
+    runs-on: [self-hosted, ernie-cpu]
+    steps:
+      - name: Run Container
+        env:
+          work_dir: ${{ github.workspace }}
+          python_version: "3.10"
+        run: |
+          container_name=${TASK}-$(date +%Y%m%d-%H%M%S)
+          echo "container_name=${container_name}" >> ${{ github.env }}
+          docker_image="iregistry.baidu-int.com/paddlecloud/base-images:paddlecloud-ubuntu20.04-gcc12.2-cuda12.3-cudnn9.0-nccl2.20.3.1-openmpi4.1.5-latest"
+          docker run -d -t --name ${container_name} --net=host -v /dev/shm:/dev/shm --shm-size=32G \
+            -v $work_dir/../../..:$work_dir/../../.. \
+            -v $work_dir:/workspace \
+            -v /home/.cache/pip:/home/.cache/pip \
+            -e BRANCH \
+            -e PR_ID \
+            -e COMMIT_ID \
+            -e work_dir \
+            -e no_proxy \
+            -e python_version \
+            -w /workspace ${docker_image}
+
+      - name: Download Code
+        env:
+          work_dir: ${{ github.workspace }}
+        run: |
+          docker exec -t ${container_name} /bin/bash -c '
+          rm -rf * .[^.]*
+          echo "Downloading PaddleNLP.tar"
+          wget -q --no-proxy  https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddleNLP.tar --no-check-certificate
+          echo "Extracting PaddleNLP.tar"
+          tar xf PaddleNLP.tar && rm -rf PaddleNLP.tar
+          source $work_dir/../../../proxy
+          cd PaddleNLP
+          git config --global user.name "PaddleCI"
+          git config --global user.email "paddle_ci@example.com"
+          git pull
+          git submodule update --init --recursive --force
+          if [ -n "${PR_ID}" ]; then
+            git fetch origin pull/${PR_ID}/head
+            git checkout -b PR_${PR_ID} FETCH_HEAD
+            git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
+            git fetch upstream ${BRANCH}
+            git merge ${BRANCH} --no-edit
+            git diff --numstat ${BRANCH} -- | awk "{print \$NF}"
+          else
+            echo "Not in a pull_request event. Skipping PR-specific operations."
+          fi
+          git log --pretty=oneline -10
+
+          if ! git show-ref --quiet refs/heads/develop; then \
+              echo "local develop branch is missing, creating local develop branch that tracks remote develop branch"
+              git fetch origin develop
+              git branch develop --track origin/develop
+          else
+            echo "local develop branch exist, skipping"
+          fi
+          '
+          
+      - name: Setup Environment
+        run: |
+          docker exec -t $container_name /bin/bash -c '
+          unlink /usr/local/bin/python
+          ln -sf $(which python${python_version}) /usr/local/bin/python
+          set -e
+          python -m pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
+          python -m pip config set global.cache-dir "/home/.cache/pip"
+          source $work_dir/../../../proxy
+          python -m pip install --upgrade pip
+          cd /workspace/PaddleNLP && git config --global --add safe.directory $PWD
+          make install
+          '
+
+      - name: Test
+        run: |
+          docker exec -t $container_name /bin/bash -c '
+          set -e
+          source $work_dir/../../../proxy
+          cd /workspace/PaddleNLP
+          make lint
+          '
+      
+      - name: Terminate And Delete the Container
+        if: always()
+        run: |
+          docker rm -f $container_name 2>/dev/null || true

.github/workflows/llm.yml

@@ -0,0 +1,207 @@
+# Disabled LLM CI workflow - manually enabled when needed
+name: LLM CI (DISABLED)
+
+#on:
+#  pull_request:
+#    types: [opened, synchronize, reopened]
+#    branches: [develop]
+#    schedule:
+#    - cron: "2 0 * * *"
+#  workflow_call:
+#    inputs:
+#      run_downstream:
+#        required: true
+#        type: string
+#      image_name:
+#        required: true
+#        type: string
+
+#concurrency:
+#  group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
+#  cancel-in-progress: true
+
+#env:
+#  PR_ID: ${{ github.event.pull_request.number }}
+  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
+  TASK: paddlenlp-CI-${{ github.event.pull_request.number }}-llm
+  ci_scripts: /workspace/PaddleNLP/scripts/regression
+  BRANCH: ${{ github.event.pull_request.base.ref }}
+  AGILE_COMPILE_BRANCH: ${{ github.event.pull_request.base.ref }}
+  CI_name: llm-ci
+  no_proxy: "localhost,bj.bcebos.com,su.bcebos.com,bcebos.com,apiin.im.baidu.com,gitee.com,aliyun.com,.baidu.com,.tuna.tsinghua.edu.cn"
+  HF_ENDPOINT: https://hf-mirror.com
+  STUDIO_GIT_HOST: http://git.prod.idc-to-cloud.aistudio.baidu-int.com
+  PPNLP_HOME: /ssd1/paddlenlp
+  HF_DATASETS_CACHE: /ssd1/paddlenlp/huggingface/datasets
+  TRANSFORMERS_CACHE: /ssd1/paddlenlp/huggingface
+  CCACHE_DIR: /home/data/gzcfs/.ccache/gpubox
+  RUN_DOWNSTREAM: ${{ inputs.run_downstream }}
+
+#defaults:
+#  run:
+#    shell: bash
+
+#jobs:
+#  llm-ci:
+#    name: llm-ci (DISABLED)
+    runs-on: [self-hosted, ernie-8gpu]
+    steps:
+      - name: Determine Image Name
+        env:
+          IMAGE_NAME: ${{ inputs.image_name }}
+        run: |
+          if [[ -n "${IMAGE_NAME}" ]]; then
+            echo "IMAGE_NAME=${IMAGE_NAME}" >> "$GITHUB_ENV"
+          else
+            echo "IMAGE_NAME=iregistry.baidu-int.com/paddlecloud/base-images:paddlecloud-ubuntu18.04-gcc8.2-cuda11.8-cudnn8.6-nccl2.15.5-paddlenlp-latest" >> "$GITHUB_ENV"
+          fi
+
+      - name: Run Container
+        env:
+          work_dir: ${{ github.workspace }}
+          CACHE_DIR: /home/data/cfs/.cache
+          FLAGS_dynamic_static_unified_comm: "True"
+          python_version: "3.10"
+          paddle_whl: https://paddle-qa.bj.bcebos.com/paddle-pipeline/Develop-GpuSome-LinuxCentos-Gcc82-Cuda118-Cudnn86-Trt85-Py310-CINN-Compile/latest/paddlepaddle_gpu-0.0.0-cp310-cp310-linux_x86_64.whl
+        run: |
+          container_name=${TASK}-$(date +%Y%m%d-%H%M%S)
+          echo "container_name=${container_name}" >> "$GITHUB_ENV"
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping..."
+          else
+            docker run -d -t --name ${container_name} --net=host -v /dev/shm:/dev/shm --shm-size=32G \
+              -v $work_dir/../../..:$work_dir/../../.. \
+              -v $work_dir:/workspace \
+              -v /home/.cache/pip:/home/.cache/pip \
+              -v /ssd1/paddlenlp:/ssd1/paddlenlp \
+              -v /home/data/gzcfs/.ccache/gpubox:/home/data/gzcfs/.ccache/gpubox \
+              -e BRANCH \
+              -e AGILE_COMPILE_BRANCH \
+              -e PR_ID \
+              -e COMMIT_ID \
+              -e work_dir \
+              -e ci_scripts \
+              -e no_proxy \
+              -e CI_name \
+              -e paddle_whl \
+              -e HF_ENDPOINT \
+              -e STUDIO_GIT_HOST \
+              -e PPNLP_HOME \
+              -e HF_DATASETS_CACHE \
+              -e TRANSFORMERS_CACHE \
+              -e CACHE_DIR \
+              -e FLAGS_dynamic_static_unified_comm \
+              -e python_version \
+              -w /workspace --runtime=nvidia $IMAGE_NAME
+          fi
+
+      - name: Download Code
+        env:
+          work_dir: ${{ github.workspace }}
+        run: |
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping.."
+          else
+            docker exec -t $container_name /bin/bash -c '
+            rm -rf * .[^.]*
+            echo "Downloading PaddleNLP.tar.gz"
+            wget -q --no-proxy  https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddleNLP.tar --no-check-certificate
+            echo "Extracting PaddleNLP.tar.gz"
+            tar xf PaddleNLP.tar && rm -rf PaddleNLP.tar
+            source $work_dir/../../../proxy
+            cd PaddleNLP
+            git config --global user.name "PaddleCI"
+            git config --global user.email "paddle_ci@example.com"
+            git pull
+            git submodule update --init --recursive --force
+            if [ -n "${PR_ID}" ]; then
+              git fetch origin pull/${PR_ID}/head
+              git checkout -b PR_${PR_ID} FETCH_HEAD
+              git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
+              git fetch upstream ${BRANCH}
+              git merge ${BRANCH} --no-edit
+              git diff --numstat ${BRANCH} -- | awk "{print \$NF}"
+            else
+              echo "Not in a pull_request event. Skipping PR-specific operations."
+            fi
+            git log --pretty=oneline -10
+            '
+          fi
+
+      - name: Skip For Bug
+        run: |
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping..."
+          else
+            docker exec -t $container_name /bin/bash -c '
+            cd /workspace/PaddleNLP
+            git revert f2477c07272d04244cd3287d1f21c70482a4a85f --no-edit  # 套件PR#10413引入bug-待修复
+            git revert d74c950e15a35b7a100d1688c89318195cc83bca --no-edit  # 框架升级后配合修改，当前因框架bug固定commit适配
+            '
+          fi
+      
+      - name: Test
+        run: |
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping..."
+          else
+            docker exec -t $container_name /bin/bash -c '
+            ldconfig
+            unlink /usr/bin/python3
+            ln -sf $(which python${python_version}) /usr/bin/python3
+            pip config set global.cache-dir "/home/.cache/pip"
+            set -e
+            source $work_dir/../../../proxy
+            cd /workspace/PaddleNLP && git config --global --add safe.directory $PWD
+            echo "Paddle PR#73283  PR#74484 import bug, set paddle_commit=8ae742 to skip "
+            export paddle_whl=https://paddle-qa.bj.bcebos.com/paddle-pipeline/Develop-GpuSome-LinuxCentos-Gcc82-Cuda118-Cudnn86-Trt85-Py310-CINN-Compile/8ae7423e99b2ea96e410968a0ebb3f1795e37205/paddlepaddle_gpu-0.0.0-cp310-cp310-linux_x86_64.whl
+            timeout 2h bash scripts/regression/run_ci.sh python${python_version} ${paddle_whl}
+            '
+          fi
+
+      - name: Upload Allure-reports & Logs
+        if: always()
+        env:
+          home_path: ${{ github.workspace }}/../../..
+          bos_file: ${{ github.workspace }}/../../../bos/BosClient.py
+          allure_file: ${{ github.workspace }}/../../../allure-2.19.0/bin/allure
+        run: |
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping..."
+          else
+            docker exec -t $container_name /bin/bash -c '
+            unset http_proxy && unset https_proxy
+            if [ ! -f "${{ env.bos_file }}" ]; then
+              wget -q --no-proxy -O ${{ env.home_path }}/bos_new.tar.gz https://xly-devops.bj.bcebos.com/home/bos_new.tar.gz --no-check-certificate
+              mkdir ${{ env.home_path }}/bos
+              tar xf ${{ env.home_path }}/bos_new.tar.gz -C ${{ env.home_path }}/bos
+            fi
+            if [ ! -f "${{ env.allure_file }}" ]; then
+              wget -q --no-proxy -O ${{ env.home_path }}/allure-2.19.0.zip https://xly-devops.bj.bcebos.com/tools/allure-2.19.0.zip --no-check-certificate
+              unzip -q ${{ env.home_path }}/allure-2.19.0.zip -d ${{ env.home_path }}/
+            fi
+            if [[ "${{ env.RUN_DOWNSTREAM }}" == "" && -n "${PR_ID}" ]]; then
+              bos_prefix="${PR_ID}/${COMMIT_ID}"
+            elif [[ "${{ env.RUN_DOWNSTREAM }}" == "true" && -n "${PR_ID}" ]]; then
+              bos_prefix="${PR_ID}/${COMMIT_ID}/test_build"
+            else
+              bos_prefix="schedule/$(date +%Y%m%d)"
+            fi
+            cd /workspace/PaddleNLP/model_logs
+            for FILE in /workspace/PaddleNLP/model_logs/*; do
+              file=$(basename "$FILE")
+              python ${{ env.bos_file }} $file paddle-github-action/PR/PaddleNLP/llm/${bos_prefix}/logs
+              echo "$file: https://paddle-github-action.bj.bcebos.com/PR/PaddleNLP/llm/${bos_prefix}/logs/$file"
+            done
+            cd /workspace/PaddleNLP/
+            ${{ env.allure_file }} generate result -o report
+            tar -czf products.tar.gz report model_logs
+            python ${{ env.bos_file }} products.tar.gz paddle-github-action/PR/PaddleNLP/llm/${bos_prefix}/logs
+            echo "products: https://paddle-github-action.bj.bcebos.com/PR/PaddleNLP/llm/${bos_prefix}/logs/products.tar.gz"
+            '
+          fi
+
+      - name: Terminate And Delete the Container
+        if: always()
+        run: |
+          docker rm -f ${{ env.container_name }} 2>/dev/null || true

.github/workflows/pipelines.yml

@@ -0,0 +1,30 @@
+name: Pipelines-Test
+
+on:
+  push:
+    paths:
+      - 'slm/pipelines/*'
+  pull_request:
+    paths:
+      - 'slm/pipelines/*'
+
+
+jobs:
+  Pipelines-Test:
+    name: Pipelines-Test
+    runs-on: ubuntu-24.04
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: 'pip' # caching pip dependencies
+      - name: Install dependencies
+        working-directory: ./slm/pipelines
+        run: |
+          python -m pip install --upgrade pip
+          make install
+          pip install -r tests/requirements.txt
+      - name: run the command
+        working-directory: ./slm/pipelines
+        run: make test

.github/workflows/release.yml

@@ -0,0 +1,58 @@
+name: Release
+
+on: 
+  workflow_dispatch:
+    inputs:
+      package:
+        description: 'Package to be deployed'
+        required: true
+        default: 'paddlenlp'
+        type: choice
+        options:
+        - paddlenlp
+        - paddle-pipelines
+        - ppdiffusers
+
+jobs:
+  Release:
+    name: Release
+    runs-on: ubuntu-24.04
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions/setup-python@v1
+        with:
+          python-version: '3.10'
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install setuptools wheel twine
+      - name: Should Deploy
+        id: should-deploy
+        env: 
+          PACKAGE: ${{ inputs.package }}
+        run: python scripts/should_deploy.py  --name $PACKAGE >> $GITHUB_OUTPUT
+
+      - name: Check Branch
+        id: check-branch
+        run: |
+          if [[ ${{ github.ref }} =~ ^refs/tags/v[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
+              echo "match=true" >> $GITHUB_OUTPUT
+          fi  # See: https://stackoverflow.com/a/58869470/1123955
+
+      - name: Is A Publish Branch
+        if: steps.check-branch.outputs.match == 'true' && steps.should-deploy.outputs.should_deploy == 'true'
+        env:
+          TWINE_USERNAME: __token__
+          TWINE_PASSWORD: ${{ secrets.paddlenlp }}
+          PACKAGE: ${{ inputs.package }}
+          PADDLENLP_STABLE_VERSION: "true"
+        run: |
+          make deploy-$PACKAGE
+
+      - name: Should Not Deploy To Pypi Server
+        if: steps.should-deploy.outputs.should_deploy != 'true'
+        run: echo 'should not deploy pypi server'
+
+      - name: Is Not A Publish Branch
+        if: steps.check-branch.outputs.match != 'true'
+        run: echo 'Not A Publish Branch'

.github/workflows/rerun.yml

@@ -0,0 +1,73 @@
+name: Re-run
+
+on:
+  issue_comment:
+    types: [created]
+
+jobs:
+  re-run:
+    if: ${{ github.event.issue.pull_request && contains(github.event.comment.body, '/re-run')  && github.event.comment.user.login == github.event.issue.user.login }}
+    runs-on: [self-hosted, ernie-cpu]
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v3
+
+      - name: Rerun all failed jobs
+        if: ${{ contains(github.event.comment.body, 'all-failed') }}
+        uses: ./.github/actions/rerun-workflow
+        with:
+          PR_ID: ${{ github.event.issue.number }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          OWNER: ${{ github.repository_owner }}
+          REPO: ${{ github.event.repository.name }}
+          JOB_NAME: 'all-failed'
+
+      - name: Rerun LLM
+        if: ${{ contains(github.event.comment.body, 'LLM') }}
+        uses: ./.github/actions/rerun-workflow
+        with:
+          PR_ID: ${{ github.event.issue.number }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          OWNER: ${{ github.repository_owner }}
+          REPO: ${{ github.event.repository.name }}
+          JOB_NAME: 'LLM CI / llm-ci'
+
+      - name: Unittest GPU
+        if: ${{ contains(github.event.comment.body, 'Unittest GPU') }}
+        uses: ./.github/actions/rerun-workflow
+        with:
+          PR_ID: ${{ github.event.issue.number }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          OWNER: ${{ github.repository_owner }}
+          REPO: ${{ github.event.repository.name }}
+          JOB_NAME: 'Unittest GPU CI / unittest-gpu-ci'
+
+      - name: Rerun Unittest CPU
+        if: ${{ contains(github.event.comment.body, 'Unittest CPU') }}
+        uses: ./.github/actions/rerun-workflow
+        with:
+          PR_ID: ${{ github.event.issue.number }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          OWNER: ${{ github.repository_owner }}
+          REPO: ${{ github.event.repository.name }}
+          JOB_NAME: 'Unittest CPU CI / unit-cpu-ci'
+
+      - name: Rerun Distribute (V100)
+        if: ${{ contains(github.event.comment.body, 'Distribute (V100)') }}
+        uses: ./.github/actions/rerun-workflow
+        with:
+          PR_ID: ${{ github.event.issue.number }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          OWNER: ${{ github.repository_owner }}
+          REPO: ${{ github.event.repository.name }}
+          JOB_NAME: 'Distribute CI (V100)/ distribute-v100-ci'
+
+      - name: Rerun Codestyle Check
+        if: ${{ contains(github.event.comment.body, 'Codestyle Check') }}
+        uses: ./.github/actions/rerun-workflow
+        with:
+          PR_ID: ${{ github.event.issue.number }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          OWNER: ${{ github.repository_owner }}
+          REPO: ${{ github.event.repository.name }}
+          JOB_NAME: 'Codestyle Check / Lint'

.github/workflows/stale.yml

@@ -0,0 +1,35 @@
+name: Stale
+
+on:
+  # Allow manual run via GitHub web or CLI
+  workflow_dispatch:
+  schedule:
+    # Run daily at midnight UTC
+    - cron: 0 0 * * *
+
+permissions:
+  issues: write
+  pull-requests: write
+
+jobs:
+  stale:
+    runs-on: ubuntu-24.04
+
+    permissions:
+      issues: write
+      pull-requests: write
+
+    steps:
+      - uses: actions/stale@v6.0.1
+        with:
+          days-before-issue-stale: 60
+          days-before-issue-close: 14
+          stale-issue-label: "stale"
+          stale-issue-message: "This issue is stale because it has been open for 60 days with no activity. 当前issue 60天内无活动，被标记为stale。"
+          close-issue-message: "This issue was closed because it has been inactive for 14 days since being marked as stale. 当前issue 被标记为stale已有14天，即将关闭。"
+          exempt-issue-labels: 'triage,keep'
+          days-before-pr-stale: 60
+          days-before-pr-close: -1
+          stale-pr-label: "stale"
+          stale-pr-message: "This Pull Request is stale because it has been open for 60 days with no activity. 当前Pull Request 60天内无活动，被标记为stale。"
+          operations-per-run: 400

.github/workflows/unittest-cpu.yml

@@ -0,0 +1,120 @@
+# Disabled unittest CPU workflow - manually enabled when needed
+name: Unittest CPU CI (DISABLED)
+
+#on: [pull_request]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
+  cancel-in-progress: true
+
+env:
+  PR_ID: ${{ github.event.pull_request.number }}
+  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
+  BRANCH: ${{ github.event.pull_request.base.ref }}
+  TASK: PaddleNLP-CI-${{ github.event.pull_request.number }}-unittest-cpu
+  HF_ENDPOINT: https://hf-mirror.com
+  STUDIO_GIT_HOST: http://git.prod.idc-to-cloud.aistudio.baidu-int.com
+  PPNLP_HOME: /home/disk1/cache
+  HF_DATASETS_CACHE: /home/disk1/cache/huggingface/datasets
+  TRANSFORMERS_CACHE: /home/disk1/cache/huggingface
+
+#jobs:
+#  unittest-cpu-ci:
+#    name: unittest-cpu-ci (DISABLED)
+    runs-on: [self-hosted, ernie-cpu]
+    permissions:
+      pull-requests: write
+      contents: read
+      id-token: write
+    steps:
+      - name: Run Container
+        env:
+          work_dir: ${{ github.workspace }}
+          python_version: "3.10"
+        run: |
+          container_name=${TASK}-$(date +%Y%m%d-%H%M%S)
+          echo "container_name=${container_name}" >> ${{ github.env }}
+          docker_image="iregistry.baidu-int.com/paddlecloud/base-images:paddlecloud-ubuntu20.04-gcc12.2-cuda12.3-cudnn9.0-nccl2.20.3.1-openmpi4.1.5-latest"
+          docker run -d -t --name ${container_name} --net=host -v /dev/shm:/dev/shm --shm-size=32G \
+            -v $work_dir/../../..:$work_dir/../../.. \
+            -v $work_dir:/workspace \
+            -v /home/.cache/pip:/home/.cache/pip \
+            -v /home/disk1/cache:/home/disk1/cache \
+            -e BRANCH \
+            -e PR_ID \
+            -e COMMIT_ID \
+            -e work_dir \
+            -e no_proxy \
+            -e python_version \
+            -e HF_ENDPOINT \
+            -e STUDIO_GIT_HOST \
+            -e PPNLP_HOME \
+            -e HF_DATASETS_CACHE \
+            -e TRANSFORMERS_CACHE \
+            -w /workspace ${docker_image}
+
+      - name: Download Code
+        env:
+          work_dir: ${{ github.workspace }}
+        run: |
+          docker exec -t ${container_name} /bin/bash -c '
+          rm -rf * .[^.]*
+          echo "Downloading PaddleNLP.tar"
+          wget -q --no-proxy  https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddleNLP.tar --no-check-certificate
+          echo "Extracting PaddleNLP.tar"
+          tar xf PaddleNLP.tar && rm -rf PaddleNLP.tar
+          source $work_dir/../../../proxy
+          cd PaddleNLP
+          git config --global user.name "PaddleCI"
+          git config --global user.email "paddle_ci@example.com"
+          git pull
+          git submodule update --init --recursive --force
+          if [ -n "${PR_ID}" ]; then
+            git fetch origin pull/${PR_ID}/head
+            git checkout -b PR_${PR_ID} FETCH_HEAD
+            git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
+            git fetch upstream ${BRANCH}
+            git merge ${BRANCH} --no-edit
+            git diff --numstat ${BRANCH} -- | awk "{print \$NF}"
+          else
+            echo "Not in a pull_request event. Skipping PR-specific operations."
+          fi
+          git log --pretty=oneline -10
+          '
+
+      - name: Setup Environment
+        run: |
+          docker exec -t $container_name /bin/bash -c '
+          unlink /usr/bin/python3
+          ln -sf $(which python${python_version}) /usr/bin/python3
+          set -e
+          python -m pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
+          python -m pip config set global.cache-dir "/home/.cache/pip"
+          source $work_dir/../../../proxy
+          python -m pip install --upgrade pip
+          cd /workspace/PaddleNLP && git config --global --add safe.directory $PWD
+          pip install -r tests/requirements.txt
+          make install
+          '
+
+      - name: Test
+        run: |
+          docker exec -t $container_name /bin/bash -c '
+          source $work_dir/../../../proxy
+          cd /workspace/PaddleNLP
+          set -e
+          make test
+          '
+
+      # - name: Upload Coverage To Codecov
+      #   uses: codecov/codecov-action@v3
+      #   env:
+      #     CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
+      #   with:
+      #     token: ${{ secrets.CODECOV_TOKEN }}
+      #     files: ./PaddleNLP/coverage.xml
+      
+      - name: Terminate And Delete the Container
+        if: always()
+        run: |
+          docker rm -f $container_name 2>/dev/null || true

.github/workflows/unittest-gpu.yml

@@ -0,0 +1,204 @@
+# Disabled unittest GPU workflow - manually enabled when needed
+name: Unittest GPU CI (DISABLED)
+
+#on:
+#  pull_request:
+#    types: [opened, synchronize, reopened]
+#    branches: [develop]
+#  schedule:
+#    - cron: "3 0 * * *"
+#  workflow_call:
+#    inputs:
+#      run_downstream:
+#        required: true
+#        type: string
+#      image_name:
+#        required: true
+#        type: string
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
+  cancel-in-progress: true
+
+env:
+  PR_ID: ${{ github.event.pull_request.number }}
+  COMMIT_ID: ${{ github.event.pull_request.head.sha }}
+  TASK: paddlenlp-CI-${{ github.event.pull_request.number }}-unit-gpu
+  ci_scripts: /workspace/PaddleNLP/scripts/unit_test
+  BRANCH: ${{ github.event.pull_request.base.ref }}
+  AGILE_COMPILE_BRANCH: ${{ github.event.pull_request.base.ref }}
+  CI_name: unittest-gpu-ci
+  no_proxy: "localhost,bj.bcebos.com,su.bcebos.com,bcebos.com,apiin.im.baidu.com,gitee.com,aliyun.com,.baidu.com,.tuna.tsinghua.edu.cn"
+  HF_ENDPOINT: https://hf-mirror.com
+  STUDIO_GIT_HOST: http://git.prod.idc-to-cloud.aistudio.baidu-int.com
+  PPNLP_HOME: /ssd1/paddlenlp
+  HF_DATASETS_CACHE: /ssd1/paddlenlp/huggingface/datasets
+  TRANSFORMERS_CACHE: /ssd1/paddlenlp/huggingface
+  CCACHE_DIR: /home/data/gzcfs/.ccache/gpubox
+  RUN_DOWNSTREAM: ${{ inputs.run_downstream }}
+
+defaults:
+  run:
+    shell: bash
+
+#jobs:
+#  unittest-gpu-ci:
+#    name: unittest-gpu-ci (DISABLED)
+    runs-on: [self-hosted, ernie-8gpu]
+    steps:
+      - name: Determine Image Name
+        env:
+          IMAGE_NAME: ${{ inputs.image_name }}
+        run: |
+          if [[ -n "${IMAGE_NAME}" ]]; then
+            echo "IMAGE_NAME=${IMAGE_NAME}" >> "$GITHUB_ENV"
+          else
+            echo "IMAGE_NAME=iregistry.baidu-int.com/paddlecloud/base-images:paddlecloud-ubuntu18.04-gcc8.2-cuda11.8-cudnn8.6-nccl2.15.5-paddlenlp-latest" >> "$GITHUB_ENV"
+          fi
+      
+      - name: Run Container
+        env:
+          work_dir: ${{ github.workspace }}
+          CACHE_DIR: /home/data/cfs/.cache
+          FLAGS_dynamic_static_unified_comm: "True"
+          python_version: "3.10"
+          paddle_whl: https://paddle-qa.bj.bcebos.com/paddle-pipeline/Develop-GpuSome-LinuxCentos-Gcc82-Cuda118-Cudnn86-Trt85-Py310-CINN-Compile/latest/paddlepaddle_gpu-0.0.0-cp310-cp310-linux_x86_64.whl
+        run: |
+          container_name=${TASK}-$(date +%Y%m%d-%H%M%S)
+          echo "container_name=${container_name}" >> "$GITHUB_ENV"
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping..."
+          else
+            docker run -d -t --name ${container_name} --net=host -v /dev/shm:/dev/shm --shm-size=32G \
+              -v $work_dir/../../..:$work_dir/../../.. \
+              -v $work_dir:/workspace \
+              -v /home/.cache/pip:/home/.cache/pip \
+              -v /ssd1/paddlenlp:/ssd1/paddlenlp \
+              -v /home/data/gzcfs/.ccache/gpubox:/home/data/gzcfs/.ccache/gpubox \
+              -e BRANCH \
+              -e AGILE_COMPILE_BRANCH \
+              -e PR_ID \
+              -e COMMIT_ID \
+              -e work_dir \
+              -e ci_scripts \
+              -e no_proxy \
+              -e CI_name \
+              -e paddle_whl \
+              -e HF_ENDPOINT \
+              -e STUDIO_GIT_HOST \
+              -e PPNLP_HOME \
+              -e HF_DATASETS_CACHE \
+              -e TRANSFORMERS_CACHE \
+              -e CACHE_DIR \
+              -e FLAGS_dynamic_static_unified_comm \
+              -e python_version \
+              -w /workspace --runtime=nvidia $IMAGE_NAME
+          fi
+          
+      - name: Download Code
+        env:
+          work_dir: ${{ github.workspace }}
+        run: |
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping.."
+          else
+            docker exec -t $container_name /bin/bash -c '
+            rm -rf * .[^.]*
+            echo "Downloading PaddleNLP.tar.gz"
+            wget -q --no-proxy  https://paddle-qa.bj.bcebos.com/CodeSync/develop/PaddleNLP.tar --no-check-certificate
+            echo "Extracting PaddleNLP.tar.gz"
+            tar xf PaddleNLP.tar && rm -rf PaddleNLP.tar
+            source $work_dir/../../../proxy
+            cd PaddleNLP
+            git config --global user.name "PaddleCI"
+            git config --global user.email "paddle_ci@example.com"
+            git pull
+            git submodule update --init --recursive --force
+            if [ -n "${PR_ID}" ]; then
+              git fetch origin pull/${PR_ID}/head
+              git checkout -b PR_${PR_ID} FETCH_HEAD
+              git remote add upstream https://github.com/PaddlePaddle/PaddleNLP.git
+              git fetch upstream ${BRANCH}
+              git merge ${BRANCH} --no-edit
+              git diff --numstat ${BRANCH} -- | awk "{print \$NF}"
+            else
+              echo "Not in a pull_request event. Skipping PR-specific operations."
+            fi
+            git log --pretty=oneline -10
+            '
+          fi
+      
+      - name: Skip For Bug
+        run: |
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping..."
+          else
+            docker exec -t $container_name /bin/bash -c '
+            cd /workspace/PaddleNLP
+            echo "no skip for bug"
+            '
+          fi
+      
+      - name: Test
+        run: |
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping..."
+          else
+            docker exec -t $container_name /bin/bash -c '
+            ldconfig
+            unlink /usr/bin/python3
+            ln -sf $(which python${python_version}) /usr/bin/python3
+            pip config set global.cache-dir "/home/.cache/pip"
+            source $work_dir/../../../proxy
+            set -e
+            cd /workspace/PaddleNLP && git config --global --add safe.directory $PWD
+            timeout 50m bash scripts/unit_test/ci_unit.sh ${paddle_whl}
+            '
+          fi
+
+      - name: Upload Allure-reports & Logs
+        if: always()
+        env:
+          home_path: ${{ github.workspace }}/../../..
+          bos_file: ${{ github.workspace }}/../../../bos/BosClient.py
+          allure_file: ${{ github.workspace }}/../../../allure-2.19.0/bin/allure
+        run: |
+          if [[ "$RUN_DOWNSTREAM" == "false" ]]; then
+            echo "Not in a pull_request or test_build event. Skipping..."
+          else
+            docker exec -t $container_name /bin/bash -c '
+            unset http_proxy && unset https_proxy
+            if [ ! -f "${{ env.bos_file }}" ]; then
+              wget -q --no-proxy -O ${{ env.home_path }}/bos_new.tar.gz https://xly-devops.bj.bcebos.com/home/bos_new.tar.gz --no-check-certificate
+              mkdir ${{ env.home_path }}/bos
+              tar xf ${{ env.home_path }}/bos_new.tar.gz -C ${{ env.home_path }}/bos
+            fi
+            if [ ! -f "${{ env.allure_file }}" ]; then
+              wget -q --no-proxy -O ${{ env.home_path }}/allure-2.19.0.zip https://xly-devops.bj.bcebos.com/tools/allure-2.19.0.zip --no-check-certificate
+              unzip -q ${{ env.home_path }}/allure-2.19.0.zip -d ${{ env.home_path }}/
+            fi
+            if [[ "${{ env.RUN_DOWNSTREAM }}" == "" && -n "${PR_ID}" ]]; then
+              bos_prefix="${PR_ID}/${COMMIT_ID}"
+            elif [[ "${{ env.RUN_DOWNSTREAM }}" == "true" && -n "${PR_ID}" ]]; then
+              bos_prefix="${PR_ID}/${COMMIT_ID}/test_build"
+            else
+              bos_prefix="schedule/$(date +%Y%m%d)"
+            fi
+            cd /workspace/PaddleNLP/unittest_logs
+            for FILE in /workspace/PaddleNLP/unittest_logs/*; do
+              file=$(basename "$FILE")
+              python ${{ env.bos_file }} $file paddle-github-action/PR/PaddleNLP/unittest-gpu/${bos_prefix}/logs
+              echo "$file: https://paddle-github-action.bj.bcebos.com/PR/PaddleNLP/unittest-gpu/${bos_prefix}/logs/$file"
+            done
+            cd /workspace/PaddleNLP/
+            ${{ env.allure_file }} generate result -o report
+            tar -czf products.tar.gz report unittest_logs
+            python ${{ env.bos_file }} products.tar.gz paddle-github-action/PR/PaddleNLP/unittest-gpu/${bos_prefix}/logs
+            echo "report: https://paddle-github-action.bj.bcebos.com/PR/PaddleNLP/unittest-gpu/${bos_prefix}/logs/products.tar.gz"
+            '
+          fi
+
+      - name: Terminate And Delete the Container
+        if: always()
+        run: |
+          docker rm -f $container_name 2>/dev/null || true

.gitignore

@@ -9,6 +9,7 @@ __pycache__/
 # Distribution / packaging
 .Python
 build/
+build*
 develop-eggs/
 dist/
 downloads/
@@ -52,6 +53,7 @@ coverage.xml
 # Translations
 *.mo
 *.pot
+*.doctree
 
 # Django stuff:
 *.log
@@ -121,3 +123,21 @@ FETCH_HEAD
 
 # vscode
 .vscode
+./ppdiffusers/ppdiffusers/version.py
+
+# third party
+csrc/third_party/
+dataset/
+output/
+
+# gen codes
+autogen/
+
+# cutlass kernel
+!csrc/gpu/cutlass_kernels/gemm/collective/builders
+
+
+#fp8
+ops/csrc/fp8/deep_gemm/include/cutlass
+ops/csrc/fp8/deep_gemm/include/cute
+.ccls-cache

.gitmodules

@@ -0,0 +1,6 @@
+[submodule "csrc/third_party/cutlass"]
+	path = csrc/third_party/cutlass
+	url = https://github.com/NVIDIA/cutlass.git
+[submodule "csrc/third_party/nlohmann_json"]
+	path = csrc/third_party/nlohmann_json
+	url = https://github.com/nlohmann/json.git

.pre-commit-config.yaml

@@ -1,10 +1,22 @@
--   repo: https://github.com/PaddlePaddle/mirrors-yapf.git
-    sha: 0d79c0c469bab64f7229c9aca2b1186ef47f0e37
+exclude: 'slm/model_zoo/gpt-3;csrc/third_party'
+repos:
+# For Python files
+-   repo: https://github.com/psf/black.git
+    rev: 22.8.0
     hooks:
-    -   id: yapf
-        files: \.py$
+    -   id: black
+        files: \.(py|pyi)$
+        additional_dependencies: [toml]
+-   repo: https://github.com/PyCQA/isort
+    rev: 5.11.5
+    hooks:
+    -   id: isort
+-   repo: https://github.com/PyCQA/flake8
+    rev: 4.0.1
+    hooks:
+    -   id: flake8
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    sha: a11d9314b22d8f8c7556443875b731ef05965464
+    rev: v4.1.0
     hooks:
     -   id: check-merge-conflict
     -   id: check-symlinks
@@ -15,7 +27,7 @@
     -   id: trailing-whitespace
         files: \.md$
 -   repo: https://github.com/Lucas-C/pre-commit-hooks
-    sha: v1.0.1
+    rev: v1.1.14
     hooks:
     -   id: forbid-crlf
         files: \.md$
@@ -25,11 +37,37 @@
         files: \.md$
     -   id: remove-tabs
         files: \.md$
+-   repo: local
+    hooks:
+    -   id: copyright_checker
+        name: copyright_checker
+        entry: python .copyright.hook
+        language: system
+        files: \.(c|cc|cxx|cpp|cu|h|hpp|hxx|proto|xpu|kps|py|sh)$
+# For Markdown files
+-   repo: local
+    hooks:
+    -   id: add-spaces-between-chinese-and-english
+        name: Add spaces between Chinese and English characters
+        entry: python scripts/codestyle/check_spaces.py
+        language: python
+        files: \.(md|markdown)$
+        pass_filenames: true
+# For dead links
+-   repo: local
+    hooks:
+    -   id: check-dead-links
+        name: Check dead links
+        entry: python scripts/codestyle/check_dead_links.py
+        language: python
+        files: \.(md|markdown|rst)$
+        pass_filenames: true
+
 -   repo: local
     hooks:
     -   id: clang-format
         name: clang-format
-        description: Format files with ClangFormat
-        entry: bash .clang_format.hook -i
+        description: Format files with ClangFormat.
+        entry: bash ./tools/codestyle/clang_format.sh -i
         language: system
-        files: \.(c|cc|cxx|cpp|cu|h|hpp|hxx|cuh|proto)$
+        files: \.(c|cc|cxx|cpp|cu|h|hpp|hxx|xpu|kps)$

.readthedocs.yaml

@@ -4,6 +4,10 @@
 
 # Required
 version: 2
+build:
+  os: "ubuntu-20.04"
+  tools:
+    python: "3.10"
 
 submodules:
   include: all
@@ -11,7 +15,7 @@ submodules:
 
 # Build documentation in the docs/ directory with Sphinx
 sphinx:
-   configuration: docs/conf.py
+   configuration: docs/zh/conf.py
 
 # Optionally build your docs in additional formats such as PDF
 #formats:
@@ -19,10 +23,5 @@ sphinx:
 
 # Optionally set the version of Python and requirements required to build your docs
 python:
-  version: 3.7
   install:
     - requirements: docs/requirements.txt
-    - requirements: requirements.txt
-    - method: setuptools
-      path: .
-  system_packages: true

.style.yapf

@@ -1,3 +0,0 @@
-[style]
-based_on_style = pep8
-column_limit = 80

CONTRIBUTING.md

@@ -0,0 +1,158 @@
+**简体中文**🀄 | [English🌎](.github/CONTRIBUTING_en.md)
+
+# Contributing to PaddleNLP
+
+我们非常欢迎并希望您对`PaddleNLP`做出开源贡献。在您开始提交您的贡献之前，请先行签署[PaddlePaddle 贡献者许可协议](https://cla-assistant.io/PaddlePaddle/PaddleNLP)。
+本文接下来将介绍我们的开发与贡献流程：
+
+## 贡献方式
+
+我们欢迎不同的向`PaddleNLP`做出贡献的方式，例如：
+
+- 修复已知的 Issue
+- 提交新的 Issue，例如提出功能需求或者 bug 报告
+- 实现新的模型结构
+
+如果您不知道从哪里开始，请查看 Issues 板块中的`Good First Issue`标签。它为您提供一个对初学者友好的已知 Issue 列表，可以降低贡献的门槛，帮助您开始为开源做出贡献。您只需在您想处理的 Issue 中告知我们您想负责此 Issue 即可。
+
+## 开发流程
+
+PaddleNLP 使用 [Git 分支模型](http://nvie.com/posts/a-successful-git-branching-model/)。对于常见的开源贡献，我们有以下的贡献流程：
+
+### 1. Fork
+
+   因为 PaddleNLP 的开发社区一直在发展，如果每位贡献者都直接向官方 Repo 提交 commit 将会难以管理。因此，请从您的分支中提交 Pull Requests。建议您通过 GitHub 的[“Fork”按钮](https://help.github.com/articles/fork-a-repo/)来创建您的 Fork 分支。
+
+### 2. Clone
+
+   请运行一下命令将您的分支 clone 到本地
+
+   ```bash
+   git clone https://github.com/<your-github-account>/PaddleNLP
+   cd PaddleNLP
+   ```
+
+### 3. 创建本地开发分支
+
+   对于添加新功能或修复错误等日常工作，请在开发前创建您的本地开发分支：
+
+   ```bash
+   git checkout -b my-cool-feature
+   ```
+
+### 4. 配置开发环境
+
+   在开始编码之前，您需要设置开发环境。我们强烈建议您在虚拟环境中进行所有开发，例如[venv](https://docs.python.org/3/library/venv.html)或[conda](https://docs.conda.io/en/latest/)。
+   请您设置并激活虚拟环境后，运行以下命令：
+
+   ```bash
+   make install
+   ```
+
+   这将设置 `PaddleNLP` 的所有依赖以及 [`pre-commit`](http://pre-commit.com/) 工具。
+
+   如果您需要开发 `examples` 或 `applications` 模块并加载 `PaddleNLP`，请确保以可编辑模式（`-e`）安装 `PaddleNLP`。
+   如果在虚拟环境中已经安装 `PaddleNLP` ，请使用 `pip uninstall paddlenlp` 将其删除，然后以可编辑模式重新安装它
+   `pip install -e .`
+
+### 5. 开发
+
+   当您开发时，请确保您新增的代码会被单元测试所覆盖。我们所有的单元测试都可以在 `tests` 目录下找到。
+   您可以修改现有单元测试以覆盖新功能，也可以从头开始创建新测试。
+   当您完成代码时，您应该确保相关的单元测试可以通过。您可以像这样运行受更改影响的测试：
+
+   ```bash
+   pytest tests/<test_to_run>.py
+   ```
+
+### 6. Commit
+
+   我们使用 [`pre-commit`](http://pre-commit.com/)工具（包括[black](https://black.readthedocs.io/en/stable/)、[isort](https:/ /pycqa.github.io/isort/) 和
+   [flake8](https://flake8.pycqa.org/en/latest/)）来检查每次提交中的代码和文档的风格。当你运行 `git commit` 时，你会看到
+   类似于以下内容：
+
+   ```text
+    ➜  (my-virtual-env) git commit -m "committing my cool feature"
+    black....................................................................Passed
+    isort....................................................................Passed
+    flake8...................................................................Passed
+    check for merge conflicts................................................Passed
+    check for broken symlinks............................(no files to check)Skipped
+    detect private key.......................................................Passed
+    fix end of files.....................................(no files to check)Skipped
+    trim trailing whitespace.............................(no files to check)Skipped
+    CRLF end-lines checker...............................(no files to check)Skipped
+    CRLF end-lines remover...............................(no files to check)Skipped
+    No-tabs checker......................................(no files to check)Skipped
+    Tabs remover.........................................(no files to check)Skipped
+    copyright_checker........................................................Passed
+   ```
+
+   但大多数时候事情并没有那么顺利。当您的代码或文档不符合标准时，`pre-commit` 检查将失败。
+
+   ```text
+    ➜  (my-virtual-env) git commit -m "committing my cool feature"
+    black....................................................................Passed
+    isort....................................................................Failed
+    - hook id: isort
+    - files were modified by this hook
+
+    Fixing examples/information_extraction/waybill_ie/run_ernie_crf.py
+
+    flake8...................................................................Passed
+    check for merge conflicts................................................Passed
+    check for broken symlinks............................(no files to check)Skipped
+    detect private key.......................................................Passed
+    fix end of files.....................................(no files to check)Skipped
+    trim trailing whitespace.............................(no files to check)Skipped
+    CRLF end-lines checker...............................(no files to check)Skipped
+    CRLF end-lines remover...............................(no files to check)Skipped
+    No-tabs checker......................................(no files to check)Skipped
+    Tabs remover.........................................(no files to check)Skipped
+    copyright_checker........................................................Passed
+   ```
+
+   我们的工具将自动修复大部分样式错误，但是有些错误需要手动解决。幸运的是，错误信息一般通俗易懂，很容易修复。
+   解决错误后，您可以再次运行 `git add <files>` 和 `git commit`，这将再次触发 pre-commit 。
+   一旦 pre-commit 检查通过，您就可以推送代码了。
+
+   [Google](https://google.com/) 或 [StackOverflow](https://stackoverflow.com/) 是帮助您了解代码风格错误的好工具。
+   如果您仍然无法弄清楚，请不要担心。您可以使用 `git commit -m "style error" --no-verify` 提交，我们很乐意在您创建 Pull Request 后帮助您。
+
+### 7. git pull 与代码冲突
+
+   有经验的 Git 用户经常从官方 Repo 中 git pull。因为这样子他们会及早注意到与其他人的代码冲突，并且让代码冲突更容易解决
+
+   ```bash
+   git remote add upstream https://github.com/PaddlePaddle/PaddleNLP
+   git pull upstream develop
+   ```
+
+### 8. git push 与提交 Pull Request
+
+   您可以将您的本地开发分支中的工作 push 到您的 fork 的分支中：
+
+   ```bash
+   git push origin my-cool-stuff
+   ```
+
+   git push 之后，您可以提交 Pull Request，请求[官方 repo](https://github.com/PaddlePaddle/PaddleNLP) 采纳您的开发工作。请您依照[这些步骤](https://help.github.com/articles/creating-a-pull-request/)创建 Pull Request。
+
+### 9. 删除已经合入的本地和远程分支
+
+   为了保持您本地的工作区和 fork 分支的干净整洁，建议您在 Pull Request 合入之后删除本地的残余分支：
+
+   ```bash
+   git push origin my-cool-stuff
+   git checkout develop
+   git pull upstream develop
+   git branch -d my-cool-stuff
+   ```
+
+## 代码 Review
+
+- 在您的 Pull Request 能够顺利通过本地测试以及 CI 的情况下，您可以在 Pull Request 中 @ 相关的 Reviewer，提醒他们尽快对您的 Pull Request 进行 Review。
+
+- 请处理 Reviewer 的每一条评论。如果您已按照评论修改，请回复“完成”；否则，可以在评论下展开讨论。
+
+- 如果您不希望您的 Reviewer 被电子邮件通知淹没，您可以[批量回复](https://help.github.com/articles/reviewing-proposed-changes-in-a-pull-request/)。

Makefile

@@ -0,0 +1,79 @@
+# Makefile for PaddleNLP
+#
+# 	GitHb: https://github.com/PaddlePaddle/PaddleNLP
+# 	Author: Paddle Team https://github.com/PaddlePaddle
+#
+
+.PHONY: all
+all : lint test
+check_dirs := applications examples model_zoo paddlenlp pipelines ppdiffusers scripts tests 
+# # # # # # # # # # # # # # # Format Block # # # # # # # # # # # # # # # 
+
+format:
+	pre-commit run isort
+	pre-commit run black
+
+# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
+
+# # # # # # # # # # # # # # # Lint Block # # # # # # # # # # # # # # # 
+
+.PHONY: lint
+lint:
+	$(eval modified_py_files := $(shell python scripts/get_modified_files.py $(check_dirs)))
+	@if test -n "$(modified_py_files)"; then \
+		echo ${modified_py_files}; \
+		pre-commit run --files ${modified_py_files}; \
+	else \
+		echo "No library .py files were modified"; \
+	fi	
+
+# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
+
+# # # # # # # # # # # # # # # Test Block # # # # # # # # # # # # # # # 
+
+.PHONY: test
+test: unit-test
+
+unit-test:
+	PYTHONPATH=$(shell pwd) pytest -v \
+		-n auto \
+		--retries 1 --retry-delay 1 \
+		--durations 20 \
+		--cov paddlenlp \
+		--cov-report xml:coverage.xml
+
+# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
+
+.PHONY: install
+install:
+	pip install --pre paddlepaddle -i https://www.paddlepaddle.org.cn/packages/nightly/cpu/
+	pip install -r requirements-dev.txt
+	pip install -r requirements.txt
+	pip install -r paddlenlp/experimental/autonlp/requirements.txt
+	pre-commit install
+
+
+.PHONY: deploy-ppdiffusers
+deploy-ppdiffusers:
+	cd ppdiffusers && make install && make
+
+.PHONY: deploy-paddle-pipelines
+deploy-paddle-pipelines:
+	cd pipelines && make install && make
+
+.PHONY: deploy-paddlenlp
+deploy-paddlenlp:
+	# install related package
+	make install
+	# build
+	python3 setup.py sdist bdist_wheel
+	# upload
+	twine upload --skip-existing dist/*
+
+.PHONY: regression-all
+release: 
+	bash ./scripts/regression/run_release.sh 0 0,1 all
+
+.PHONY: regression-key
+key: 
+	bash ./scripts/regression/run_release.sh 0 0,1 p0

README.md

@@ -1,185 +1,298 @@
-简体中文 | [English](./README_en.md)
+**简体中文**🀄 | [English🌎](./README_en.md)
 
 <p align="center">
-  <img src="./docs/imgs/paddlenlp.png" width="720" height ="100" />
+  <img src="https://user-images.githubusercontent.com/1371212/175816733-8ec25eb0-9af3-4380-9218-27c154518258.png" align="middle"  width="500" />
 </p>
 
 ------------------------------------------------------------------------------------------
 
-[![PyPI - PaddleNLP Version](https://img.shields.io/pypi/v/paddlenlp.svg?label=pip&logo=PyPI&logoColor=white)](https://pypi.org/project/paddlenlp/)
-[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/paddlenlp)](https://pypi.org/project/paddlenlp/)
-[![PyPI Status](https://pepy.tech/badge/paddlenlp/month)](https://pepy.tech/project/paddlenlp)
-![python version](https://img.shields.io/badge/python-3.6+-orange.svg)
-![support os](https://img.shields.io/badge/os-linux%2C%20win%2C%20mac-yellow.svg)
-![GitHub](https://img.shields.io/github/license/paddlepaddle/paddlenlp)
+<p align="center">
+    <a href="https://paddlenlp.readthedocs.io/en/latest/?badge=latest"><img src="https://readthedocs.org/projects/paddlenlp/badge/?version=latest">
+    <a href="https://github.com/PaddlePaddle/PaddleNLP/releases"><img src="https://img.shields.io/github/v/release/PaddlePaddle/PaddleNLP?color=ffa"></a>
+    <a href=""><img src="https://img.shields.io/badge/python-3.7+-aff.svg"></a>
+    <a href=""><img src="https://img.shields.io/badge/os-linux%2C%20win%2C%20mac-pink.svg"></a>
+    <a href="https://github.com/PaddlePaddle/PaddleNLP/graphs/contributors"><img src="https://img.shields.io/github/contributors/PaddlePaddle/PaddleNLP?color=9ea"></a>
+    <a href="https://github.com/PaddlePaddle/PaddleNLP/commits"><img src="https://img.shields.io/github/commit-activity/m/PaddlePaddle/PaddleNLP?color=3af"></a>
+    <a href="https://pypi.org/project/paddlenlp/"><img src="https://img.shields.io/pypi/dm/paddlenlp?color=9cf"></a>
+    <a href="https://github.com/PaddlePaddle/PaddleNLP/issues"><img src="https://img.shields.io/github/issues/PaddlePaddle/PaddleNLP?color=9cc"></a>
+    <a href="https://github.com/PaddlePaddle/PaddleNLP/stargazers"><img src="https://img.shields.io/github/stars/PaddlePaddle/PaddleNLP?color=ccf"></a>
+    <a href="./LICENSE"><img src="https://img.shields.io/badge/license-Apache%202-dfd.svg"></a>
+</p>
 
-## 简介
+<h4 align="center">
+  <a href=#特性> 特性 </a> |
+  <a href=#模型支持> 模型支持 </a> |
+  <a href=#安装> 安装 </a> |
+  <a href=#快速开始> 快速开始 </a> |
+  <a href=#社区交流> 社区交流 </a>
+</h4>
 
-PaddleNLP 2.0拥有**覆盖多场景的模型库**、**简洁易用的全流程API**与**动静统一的高性能分布式训练**能力，旨在为飞桨开发者提升文本领域建模效率，并提供基于PaddlePaddle 2.0的NLP领域最佳实践。
+**PaddleNLP**是一款基于飞桨深度学习框架的大语言模型(LLM)开发套件，支持在多种硬件上进行高效的大模型训练、无损压缩以及高性能推理。PaddleNLP 具备**简单易用**和**性能极致**的特点，致力于助力开发者实现高效的大模型产业级应用。
+
+<a href="https://trendshift.io/repositories/2246" target="_blank"><img src="https://trendshift.io/api/badge/repositories/2246" alt="PaddlePaddle%2FPaddleNLP | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
+
+## News 📢
+
+* **2025.04.29 PaddleNLP 现已支持 Qwen3 系列模型**: Qwen3 系列模型支持持两种思考模式，预训练约 36 万亿个 token、119 种语言和方言。包括六个 Dense 模型, Qwen3-32B、Qwen3-14B、Qwen3-8B、Qwen3-4B、Qwen3-1.7B 和 Qwen3-0.6B。两个 MoE 模型的权重：Qwen3-235B-A22B，Qwen3-30B-A3B。
+
+* **2025.03.12 [PaddleNLP v3.0 Beta4](https://github.com/PaddlePaddle/PaddleNLP/releases/tag/v3.0.0-beta4)**：全面支持 DeepSeek V3/R1/R1-Distill, 及 QwQ-32B 等热门思考模型。**DeepSeek V3/R1完整版支持 FP8、INT8、4-bit 量化推理，MTP 投机解码**。单机 FP8推理输出超**1000 tokens/s**; 4-bit 推理输出超**2100 tokens/s**! 发布新版推理部署镜像，热门模型[一键部署](https://paddlenlp.readthedocs.io/zh/latest/llm/server/docs/general_model_inference.html)。推理部署[使用文档](https://paddlenlp.readthedocs.io/zh/latest/llm/docs/predict/index.html)全面更新，体验全面提升！自研下一代通用信息抽取模型 PP-UIE [全新发布](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/application/information_extraction)，支持8K 长度信息抽取。新增大模型 Embedding 训练，支持 INF-CL 超大 batch size 训练。新增[MergeKit](https://paddlenlp.readthedocs.io/zh/latest/llm/docs/mergekit.html)模型融合工具，缓解对齐代价。低资源训练全面优化，16G 小显存可以流畅训练。
+
+
+* **2025.02.10 PaddleNLP 现已支持 DeepSeek-R1系列模型，[在线使用](https://aistudio.baidu.com/projectdetail/8775758)**：依托全新的 PaddleNLP 3.0套件，DeepSeek-R1系列模型现已全面支持。凭借数据并行、数据分组切分并行、模型并行、流水线并行以及专家并行等一系列先进的分布式训练能力，结合 Paddle 框架独有的列稀疏注意力掩码表示技术——FlashMask 方法，DeepSeek-R1系列模型在训练过程中显著降低了显存消耗，同时取得了卓越的训练性能提升。
+
+<details><summary> <b>点击展开</b> </summary><div>
+
+* **2025.03.17 《DeepSeek-R1满血版单机部署实测》** 🔥🔥🔥 飞桨框架3.0大模型推理部署全面升级，支持多款主流大模型，DeepSeek-R1满血版实现单机部署，吞吐提升一倍！欢迎广大用户开箱体验～现已开启有奖活动：完成 DeepSeek-R1-MTP 单机部署任务、提交高质量测评 blog，即可实时赢取奖金！💰💰💰
+报名[地址](https://www.wjx.top/vm/OlzzmbG.aspx#)， 活动详情：https://github.com/PaddlePaddle/PaddleNLP/issues/10166 ， 参考文档：https://github.com/PaddlePaddle/PaddleNLP/issues/10157 。
+
+* **2025.03.06 PaddleNLP 现已支持 Qwen/QwQ-32B 模型**: 其模型参数仅有 32B，但其数学推理、编程能力和通用能力可与具备 671B 参数（其中 37B 被激活）的 DeepSeek-R1 媲美。借助 PaddleNLP 3.0套件，现可实现多种并行策略[微调训练](./llm/README.md)、[高性能推理、低比特量化](./llm/docs/predict/qwen.md)和[服务化部署](./llm/server/README.md)。
+
+* **2025.02.20 🔥🔥《PP-UIE 信息抽取智能引擎全新升级》** 强化零样本学习能力，支持极少甚至零标注数据实现高效冷启动与迁移学习，显著降低数据标注成本；具备处理长文本能力，支持 8192 个 Token 长度文档信息抽取，实现跨段落识别关键信息，形成完整理解；提供完整可定制化的训练和推理全流程，训练效率相较于 LLama-Factory 实现了1.8倍的提升。
+2月26日（周三）19：00为您深度解析全新 PP-UIE 技术方案及在部署方面的功能、优势与技巧。报名链接：https://www.wjx.top/vm/mBKC6pb.aspx?udsid=606418
+
+* **2024.12.16 [PaddleNLP v3.0 Beta3](https://github.com/PaddlePaddle/PaddleNLP/releases/tag/v3.0.0-beta3)**：大模型功能全新升级，新增了 Llama-3.2、DeepSeekV2模型，升级了 TokenizerFast，快速分词，重构了 SFTTrainer，一键开启 SFT 训练。此外，PaddleNLP 还支持了优化器状态的卸载和重载功能，实现了精细化的重新计算，训练性能提升7%。在 Unified Checkpoint 方面，进一步优化了异步保存逻辑，新增 Checkpoint 压缩功能，可节省78.5%存储空间。
+最后，在大模型推理方面，升级 Append Attention，支持了 FP8量化，支持投机解码。
+
+* **2024.12.13 📚《飞桨大模型套件 Unified Checkpoint 技术》**，加速模型存储95%，节省空间78%。支持全分布式策略调整自适应转换，提升模型训练的灵活性与可扩展性。训练-压缩-推理统一存储协议，无需手动转换提升全流程体验。Checkpoint 无损压缩结合异步保存，实现秒级存储并降低模型存储成本。适用于智能制造、指挥交通、医疗健康、金融服务等产业实际场景。12月24日（周二）19：00直播为您详细解读该技术如何优化大模型训练流程。报名链接：https://www.wjx.top/vm/huZkHn9.aspx?udsid=787976
+
+* **2024.11.28 📚《FlashRAG-Paddle | 基于 PaddleNLP 的高效开发与评测 RAG 框架》**，为文本更快更好构建准确嵌入表示、加速推理生成速度。PaddleNLP 支持超大 Batch 嵌入表示学习与多硬件高性能推理，涵盖 INT8/INT4量化技术及多种高效注意力机制优化与 TensorCore 深度优化。内置全环节算子融合技术，使得 FlashRAG 推理性能相比 transformers 动态图提升70%以上，结合检索增强知识输出结果更加准确，带来敏捷高效的使用体验。直播时间：12月3日（周二）19：00。报名链接：https://www.wjx.top/vm/eaBa1vA.aspx?udsid=682361
+
+* **2024.08.08 📚《飞桨产业级大语言模型开发利器 PaddleNLP 3.0 重磅发布》**，训压推全流程贯通，主流模型全覆盖。大模型自动并行，千亿模型训推全流程开箱即用。提供产业级高性能精调与对齐解决方案，压缩推理领先，多硬件适配。覆盖产业级智能助手、内容创作、知识问答、关键信息抽取等应用场景。直播时间：8月22日（周四）19：00。报名链接：https://www.wjx.top/vm/Y2f7FFY.aspx?udsid=143844
+
+* **2024.06.27 [PaddleNLP v3.0 Beta](https://github.com/PaddlePaddle/PaddleNLP/releases/tag/v3.0.0-beta0)**：拥抱大模型，体验全升级。统一大模型套件，实现国产计算芯片全流程接入；全面支持飞桨4D 并行配置、高效精调策略、高效对齐算法、高性能推理等大模型产业级应用流程；自研极致收敛的 RsLoRA+算法、自动扩缩容存储机制 Unified Checkpoint 和通用化支持的 FastFFN、FusedQKV 助力大模型训推；主流模型持续支持更新，提供高效解决方案。
+
+* **2024.04.24 [PaddleNLP v2.8](https://github.com/PaddlePaddle/PaddleNLP/releases/tag/v2.8.0)**：自研极致收敛的 RsLoRA+算法，大幅提升 PEFT 训练收敛速度以及训练效果；引入高性能生成加速到 RLHF PPO 算法，打破 PPO 训练中生成速度瓶颈，PPO 训练性能大幅领先。通用化支持 FastFFN、FusedQKV 等多个大模型训练性能优化方式，大模型训练更快、更稳定。
+</div></details>
 
 ## 特性
 
-- **覆盖多场景的模型库**
-  - PaddleNLP集成了RNN与Transformer等多种主流模型结构，涵盖从[词向量](./examples/word_embedding/)、[词法分析](./examples/lexical_analysis/)、[命名实体识别](./examples/information_extraction/msra_ner/)、[语义表示](./examples/language_model/)等NLP基础技术，到[文本分类](./examples/text_classification/)、[文本匹配](./examples/text_matching/)、[文本生成](./examples/text_generation/)、[文本图学习](./examples/text_graph/erniesage/)、[信息抽取](./examples/information_extraction)等NLP核心技术。同时针对[机器翻译](./examples/machine_translation/)、[通用对话](./examples/dialogue/)、[阅读理解](./examples/machine_reading_comprehension/)等系统应用提供相应核心组件与预训练模型。更多详细介绍请查看[PaddleNLP应用示例](./examples/)。
+### <a href=#多硬件训推一体> 🔧 多硬件训推一体 </a>
+
+支持英伟达 GPU、昆仑 XPU、昇腾 NPU、燧原 GCU 和海光 DCU 等多个硬件的大模型和自然语言理解模型训练和推理，套件接口支持硬件快速切换，大幅降低硬件切换研发成本。
+当前支持的自然语言理解模型：[多硬件自然语言理解模型列表](./docs/zh/model_zoo/model_list_multy_device.md)
+
+### <a href=#高效易用的预训练> 🚀 高效易用的预训练 </a>
+
+支持纯数据并行策略、分组参数切片的数据并行策略、张量模型并行策略和流水线模型并行策略的4D 高性能训练，Trainer 支持分布式策略配置化，降低复杂分布式组合带来的使用成本；
+[Unified Checkpoint 大模型存储工具](./llm/docs/unified_checkpoint.md)可以使得训练断点支持机器资源动态扩缩容恢复。此外，异步保存，模型存储可加速95%，Checkpoint 压缩，可节省78.5%存储空间。
+
+### <a href=#高效精调> 🤗 高效精调 </a>
+
+精调算法深度结合零填充数据流和 [FlashMask](./llm/docs/flashmask.md) 高性能算子，降低训练无效数据填充和计算，大幅提升精调训练吞吐。
+
+### <a href=#无损压缩和高性能推理> 🎛️ 无损压缩和高性能推理 </a>
+
+大模型套件高性能推理模块内置动态插入和全环节算子融合策略，极大加快并行推理速度。底层实现细节封装化，实现开箱即用的高性能并行推理能力。
+
+## 文档
+更多详细文档, 请访问 [PaddleNLP Documentation](https://paddlenlp.readthedocs.io/).
+
+------------------------------------------------------------------------------------------
+
+## 模型支持
+
+* 模型参数已支持 LLaMA 系列、Baichuan 系列、Bloom 系列、ChatGLM 系列、Gemma 系列、Mistral 系列、OPT 系列和 Qwen 系列，详细列表👉【LLM】模型参数支持列表如下：
+
+|                                                模型系列                                                 | 模型名称                                                                                                                                                                                                                                                                                                                                                                                      |
+|:-------------------------------------------------------------------------------------------------------:|:----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| [PP-UIE](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/application/information_extraction) | paddlenlp/PP-UIE-0.5B, paddlenlp/PP-UIE-1.5B, paddlenlp/PP-UIE-7B, paddlenlp/PP-UIE-14B                                                                                                                                                                                                                                                                                                       |
+|            [LLaMA](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/llama)             | facebook/llama-7b, facebook/llama-13b, facebook/llama-30b, facebook/llama-65b                                                                                                                                                                                                                                                                                                                 |
+|            [Llama2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/llama)            | meta-llama/Llama-2-7b, meta-llama/Llama-2-7b-chat, meta-llama/Llama-2-13b, meta-llama/Llama-2-13b-chat, meta-llama/Llama-2-70b, meta-llama/Llama-2-70b-chat                                                                                                                                                                                                                                   |
+|            [Llama3](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/llama)            | meta-llama/Meta-Llama-3-8B, meta-llama/Meta-Llama-3-8B-Instruct, meta-llama/Meta-Llama-3-70B, meta-llama/Meta-Llama-3-70B-Instruct                                                                                                                                                                                                                                                            |
+|           [Llama3.1](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/llama)           | meta-llama/Meta-Llama-3.1-8B, meta-llama/Meta-Llama-3.1-8B-Instruct, meta-llama/Meta-Llama-3.1-70B, meta-llama/Meta-Llama-3.1-70B-Instruct, meta-llama/Meta-Llama-3.1-405B, meta-llama/Meta-Llama-3.1-405B-Instruct, meta-llama/Llama-Guard-3-8B                                                                                                                                              |
+|           [Llama3.2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/llama)           | meta-llama/Llama-3.2-1B, meta-llama/Llama-3.2-1B-Instruct, meta-llama/Llama-3.2-3B, meta-llama/Llama-3.2-3B-Instruct, meta-llama/Llama-Guard-3-1B                                                                                                                                                                                                                                             |
+|           [Llama3.3](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/llama)           | meta-llama/Llama-3.3-70B-Instruct                                                                                                                                                                                                                                                                                                                                                             |
+|         [Baichuan](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/baichuan)          | baichuan-inc/Baichuan-7B, baichuan-inc/Baichuan-13B-Base, baichuan-inc/Baichuan-13B-Chat                                                                                                                                                                                                                                                                                                      |
+|         [Baichuan2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/baichuan)         | baichuan-inc/Baichuan2-7B-Base, baichuan-inc/Baichuan2-7B-Chat, baichuan-inc/Baichuan2-13B-Base, baichuan-inc/Baichuan2-13B-Chat                                                                                                                                                                                                                                                              |
+|            [Bloom](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/bloom)             | bigscience/bloom-560m, bigscience/bloom-560m-bf16, bigscience/bloom-1b1, bigscience/bloom-3b, bigscience/bloom-7b1, bigscience/bloomz-560m, bigscience/bloomz-1b1, bigscience/bloomz-3b, bigscience/bloomz-7b1-mt, bigscience/bloomz-7b1-p3, bigscience/bloomz-7b1, bellegroup/belle-7b-2m                                                                                                    |
+|          [ChatGLM](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/chatglm/)          | THUDM/chatglm-6b, THUDM/chatglm-6b-v1.1                                                                                                                                                                                                                                                                                                                                                       |
+|         [ChatGLM2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/chatglm2)          | THUDM/chatglm2-6b                                                                                                                                                                                                                                                                                                                                                                             |
+|         [ChatGLM3](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/chatglm2)          | THUDM/chatglm3-6b                                                                                                                                                                                                                                                                                                                                                                             |
+|       [DeepSeekV2](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm/config/deepseek-v2)       | deepseek-ai/DeepSeek-V2, deepseek-ai/DeepSeek-V2-Chat, deepseek-ai/DeepSeek-V2-Lite, deepseek-ai/DeepSeek-V2-Lite-Chat, deepseek-ai/DeepSeek-Coder-V2-Base, deepseek-ai/DeepSeek-Coder-V2-Instruct, deepseek-ai/DeepSeek-Coder-V2-Lite-Base, deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct                                                                                                      |
+|       [DeepSeekV3](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm/config/deepseek-v2)       | deepseek-ai/DeepSeek-V3, deepseek-ai/DeepSeek-V3-Base                                                                                                                                                                                                                                                                                                                                         |
+|      [DeepSeek-R1](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm/config/deepseek-v2)       | deepseek-ai/DeepSeek-R1, deepseek-ai/DeepSeek-R1-Zero, deepseek-ai/DeepSeek-R1-Distill-Llama-70B, deepseek-ai/DeepSeek-R1-Distill-Llama-8B, deepseek-ai/DeepSeek-R1-Distill-Qwen-14B, deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B, deepseek-ai/DeepSeek-R1-Distill-Qwen-32B, deepseek-ai/DeepSeek-R1-Distill-Qwen-7B                                                                            |
+|            [Gemma](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/gemma)             | google/gemma-7b, google/gemma-7b-it, google/gemma-2b, google/gemma-2b-it                                                                                                                                                                                                                                                                                                                      |
+|          [Mistral](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/mistral)           | mistralai/Mistral-7B-Instruct-v0.3, mistralai/Mistral-7B-v0.1                                                                                                                                                                                                                                                                                                                                 |
+|          [Mixtral](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/mixtral)           | mistralai/Mixtral-8x7B-Instruct-v0.1                                                                                                                                                                                                                                                                                                                                                          |
+|              [OPT](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/opt)               | facebook/opt-125m, facebook/opt-350m, facebook/opt-1.3b, facebook/opt-2.7b, facebook/opt-6.7b, facebook/opt-13b, facebook/opt-30b, facebook/opt-66b, facebook/opt-iml-1.3b, opt-iml-max-1.3b                                                                                                                                                                                                  |
+|             [Qwen](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)             | qwen/qwen-7b, qwen/qwen-7b-chat, qwen/qwen-14b, qwen/qwen-14b-chat, qwen/qwen-72b, qwen/qwen-72b-chat,                                                                                                                                                                                                                                                                                        |
+|           [Qwen1.5](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)            | Qwen/Qwen1.5-0.5B, Qwen/Qwen1.5-0.5B-Chat, Qwen/Qwen1.5-1.8B, Qwen/Qwen1.5-1.8B-Chat, Qwen/Qwen1.5-4B, Qwen/Qwen1.5-4B-Chat, Qwen/Qwen1.5-7B, Qwen/Qwen1.5-7B-Chat, Qwen/Qwen1.5-14B, Qwen/Qwen1.5-14B-Chat, Qwen/Qwen1.5-32B, Qwen/Qwen1.5-32B-Chat, Qwen/Qwen1.5-72B, Qwen/Qwen1.5-72B-Chat, Qwen/Qwen1.5-110B, Qwen/Qwen1.5-110B-Chat, Qwen/Qwen1.5-MoE-A2.7B, Qwen/Qwen1.5-MoE-A2.7B-Chat |
+|            [Qwen2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)             | Qwen/Qwen2-0.5B, Qwen/Qwen2-0.5B-Instruct, Qwen/Qwen2-1.5B, Qwen/Qwen2-1.5B-Instruct, Qwen/Qwen2-7B, Qwen/Qwen2-7B-Instruct, Qwen/Qwen2-72B, Qwen/Qwen2-72B-Instruct, Qwen/Qwen2-57B-A14B, Qwen/Qwen2-57B-A14B-Instruct                                                                                                                                                                       |
+|          [Qwen2-Math](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)          | Qwen/Qwen2-Math-1.5B, Qwen/Qwen2-Math-1.5B-Instruct, Qwen/Qwen2-Math-7B, Qwen/Qwen2-Math-7B-Instruct, Qwen/Qwen2-Math-72B, Qwen/Qwen2-Math-72B-Instruct, Qwen/Qwen2-Math-RM-72B                                                                                                                                                                                                               |
+|           [Qwen2.5](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)            | Qwen/Qwen2.5-0.5B, Qwen/Qwen2.5-0.5B-Instruct, Qwen/Qwen2.5-1.5B, Qwen/Qwen2.5-1.5B-Instruct, Qwen/Qwen2.5-3B, Qwen/Qwen2.5-3B-Instruct, Qwen/Qwen2.5-7B, Qwen/Qwen2.5-7B-Instruct, Qwen/Qwen2.5-7B-Instruct-1M, Qwen/Qwen2.5-14B, Qwen/Qwen2.5-14B-Instruct, Qwen/Qwen2.5-14B-Instruct-1M, Qwen/Qwen2.5-32B, Qwen/Qwen2.5-32B-Instruct, Qwen/Qwen2.5-72B, Qwen/Qwen2.5-72B-Instruct          |
+|         [Qwen2.5-Math](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)         | Qwen/Qwen2.5-Math-1.5B, Qwen/Qwen2.5-Math-1.5B-Instruct, Qwen/Qwen2.5-Math-7B, Qwen/Qwen2.5-Math-7B-Instruct, Qwen/Qwen2.5-Math-72B, Qwen/Qwen2.5-Math-72B-Instruct, Qwen/Qwen2.5-Math-RM-72B                                                                                                                                                                                                 |
+|        [Qwen2.5-Coder](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)         | Qwen/Qwen2.5-Coder-1.5B, Qwen/Qwen2.5-Coder-1.5B-Instruct, Qwen/Qwen2.5-Coder-7B, Qwen/Qwen2.5-Coder-7B-Instruct                                                                                                                                                                                                                                                                              |
+|           [Qwen3](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)              | Qwen/Qwen3-0.6B, Qwen/Qwen3-1.7B, Qwen/Qwen3-4B, Qwen/Qwen3-8B, Qwen/Qwen3-14B, Qwen/Qwen3-32B, Qwen/Qwen3-30B-A3B, Qwen/Qwen3-235B-A22B, Qwen/Qwen3-0.6B-Base, Qwen/Qwen3-1.7B-Base, Qwen/Qwen3-4B-Base, Qwen/Qwen3-8B-Base, Qwen/Qwen3-14B-Base, Qwen/Qwen3-30B-A3B-Base                                                                                                                    |
+|             [QwQ](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)              | Qwen/QwQ-32B, Qwen/QwQ-32B-Preview                                                                                                                                                                                                                                                                                                                                                            |
+|            [Yuan2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/yuan/)             | IEITYuan/Yuan2-2B, IEITYuan/Yuan2-51B, IEITYuan/Yuan2-102B                                                                                                                                                                                                                                                                                                                                    |
+
+* 4D 并行和算子优化已支持 LLaMA 系列、Baichuan 系列、Bloom 系列、ChatGLM 系列、Gemma 系列、Mistral 系列、OPT 系列和 Qwen 系列，【LLM】模型4D 并行和算子支持列表如下：
+
+| 模型名称/并行能力支持 | 数据并行 | 张量模型并行 |          | 参数分片并行 |        |        | 流水线并行 |
+|:---------------------:|:--------:|:------------:|:--------:|:------------:|:------:|:------:|:----------:|
+|                       |          |   基础能力   | 序列并行 |    stage1    | stage2 | stage3 |            |
+|         Llama         |    ✅     |      ✅       |    ✅     |      ✅       |   ✅    |   ✅    |     ✅      |
+|         Qwen          |    ✅     |      ✅       |    ✅     |      ✅       |   ✅    |   ✅    |     ✅      |
+|        Qwen1.5        |    ✅     |      ✅       |    ✅     |      ✅       |   ✅    |   ✅    |     ✅      |
+|         Qwen2         |    ✅     |      ✅       |    ✅     |      ✅       |   ✅    |   ✅    |     ✅      |
+|     Mixtral(moe)      |    ✅     |      ✅       |    ✅     |      ✅       |   ✅    |   ✅    |     🚧     |
+|        Mistral        |    ✅     |      ✅       |    🚧    |      ✅       |   ✅    |   ✅    |     🚧     |
+|       Baichuan        |    ✅     |      ✅       |    ✅     |      ✅       |   ✅    |   ✅    |     ✅      |
+|       Baichuan2       |    ✅     |      ✅       |    ✅     |      ✅       |   ✅    |   ✅    |     ✅      |
+|        ChatGLM        |    ✅     |      ✅       |    🚧    |      ✅       |   ✅    |   ✅    |     🚧     |
+|       ChatGLM2        |    ✅     |      🚧      |    🚧    |      ✅       |   ✅    |   ✅    |     🚧     |
+|       ChatGLM3        |    ✅     |      🚧      |    🚧    |      ✅       |   ✅    |   ✅    |     🚧     |
+|         Bloom         |    ✅     |      ✅       |    🚧    |      ✅       |   ✅    |   ✅    |     🚧     |
+|      GPT-2/GPT-3      |    ✅     |      ✅       |    ✅     |      ✅       |   ✅    |   ✅    |     ✅      |
+|          OPT          |    ✅     |      ✅       |    🚧    |      ✅       |   ✅    |   ✅    |     🚧     |
+|         Gemma         |    ✅     |      ✅       |    ✅     |      ✅       |   ✅    |   ✅    |     ✅      |
+|         Yuan2         |    ✅     |      ✅       |    ✅     |      ✅       |   ✅    |   ✅    |     🚧     |
+
+* 大模型预训练、精调（包含 SFT、PEFT 技术）、对齐、量化已支持 LLaMA 系列、Baichuan 系列、Bloom 系列、ChatGLM 系列、Mistral 系列、OPT 系列和 Qwen 系列，【LLM】模型预训练、精调、对齐、量化支持列表如下：
 
 
-- **简洁易用的全流程API**
-  - 深度兼容飞桨2.0的[高层API](https://www.paddlepaddle.org.cn/documentation/docs/zh/tutorial/quick_start/high_level_api/high_level_api.html)体系，内置可复用的文本建模模块([Embedding](./docs/embeddings.md), [CRF](./paddlenlp/layers/crf.py), [Seq2Vec](./paddlenlp/seq2vec/encoder.py), [Transformer](./docs/transformers.md))，可大幅度减少在数据处理、模型组网、训练与评估、推理部署环节的开发量，提升NLP任务迭代与落地的效率。
-
-- **动静统一的高性能分布式训练**
-  - 基于飞桨2.0核心框架『动静统一』的特性与领先的混合精度优化策略，结合Fleet分布式训练API，可充分利用GPU集群资源，高效完成大规模预训练模型的分布式训练。
+| Model                                      | Pretrain | SFT | LoRA | FlashMask | Prefix Tuning | DPO/SimPO/ORPO/KTO | RLHF | Mergekit | Quantization |
+|--------------------------------------------|:--------:|:---:|:----:|:---------:|:-------------:|:------------------:|:----:|:--------:|:------------:|
+| [Llama](./llm/config/llama)                |    ✅     |  ✅  |  ✅   |     ✅     |       ✅       |         ✅          |  ✅   |    ✅     |      ✅       |
+| [Qwen](./llm/config/qwen)                  |    ✅     |  ✅  |  ✅   |     ✅     |       ✅       |         ✅          |  🚧  |    ✅     |      🚧      |
+| [Mixtral](./llm/config/mixtral)            |    ✅     |  ✅  |  ✅   |    🚧     |      🚧       |         ✅          |  🚧  |    ✅     |      🚧      |
+| [Mistral](./llm/config/mistral)            |    ✅     |  ✅  |  ✅   |    🚧     |       ✅       |         ✅          |  🚧  |    ✅     |      🚧      |
+| [Baichuan/Baichuan2](./llm/config/llama)   |    ✅     |  ✅  |  ✅   |     ✅     |       ✅       |         ✅          |  🚧  |    ✅     |      ✅       |
+| [ChatGLM-6B](./llm/config/chatglm)         |    ✅     |  ✅  |  ✅   |    🚧     |       ✅       |         🚧         |  🚧  |    ✅     |      ✅       |
+| [ChatGLM2/ChatGLM3](./llm/config/chatglm2) |    ✅     |  ✅  |  ✅   |    🚧     |       ✅       |         ✅          |  🚧  |    ✅     |      ✅       |
+| [Bloom](./llm/config/bloom)                |    ✅     |  ✅  |  ✅   |    🚧     |       ✅       |         🚧         |  🚧  |    ✅     |      ✅       |
+| [GPT-3](./llm/config/gpt-3)                |    ✅     |  ✅  |  🚧  |    🚧     |      🚧       |         🚧         |  🚧  |    ✅     |      🚧      |
+| [OPT](./llm/config/opt)                    |    ✅     |  ✅  |  ✅   |    🚧     |      🚧       |         🚧         |  🚧  |    ✅     |      🚧      |
+| [Gemma](./llm/config/gemma)                |    ✅     |  ✅  |  ✅   |    🚧     |      🚧       |         ✅          |  🚧  |    ✅     |      🚧      |
+| [Yuan](./llm/config/yuan)                  |    ✅     |  ✅  |  ✅   |    🚧     |      🚧       |         ✅          |  🚧  |    ✅     |      🚧      |
+* [大模型推理](./llm/docs/predict/inference.md)已支持 LLaMA 系列、Qwen 系列、DeepSeek 系列、Mistral 系列、ChatGLM 系列、Bloom 系列和 Baichuan 系列，支持 Weight Only INT8及 INT4推理，支持 WAC（权重、激活、Cache KV）进行 INT8、FP8量化的推理，【LLM】模型推理支持列表如下：
 
+|           模型名称/量化类型支持            | FP16/BF16 | WINT8 | WINT4 | INT8-A8W8 | FP8-A8W8 | INT8-A8W8C8 |
+|:------------------------------------------:|:---------:|:-----:|:-----:|:---------:|:--------:|:-----------:|
+|    [LLaMA](./llm/docs/predict/llama.md)    |     ✅     |   ✅   |   ✅   |     ✅     |    ✅     |      ✅      |
+|     [Qwen](./llm/docs/predict/qwen.md)     |     ✅     |   ✅   |   ✅   |     ✅     |    ✅     |      ✅      |
+| [DeepSeek](./llm/docs/predict/deepseek.md) |     ✅     |   ✅   |   ✅   |    🚧     |    ✅     |     🚧      |
+|   [Qwen-Moe](./llm/docs/predict/qwen.md)   |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
+|  [Mixtral](./llm/docs/predict/mixtral.md)  |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
+|                  ChatGLM                   |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
+|                   Bloom                    |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
+|                  BaiChuan                  |     ✅     |   ✅   |   ✅   |     ✅     |    ✅     |     🚧      |
 
 ## 安装
 
 ### 环境依赖
 
-- python >= 3.6
-- paddlepaddle >= 2.0.1
+* python >= 3.8
+* paddlepaddle >= 3.0.0rc1
 
-### pip安装
+如果您尚未安装 PaddlePaddle，请参考 [飞桨官网](https://www.paddlepaddle.org.cn/) 进行安装。
 
-```
-pip install --upgrade paddlenlp -i https://pypi.org/simple
+### pip 安装
+
+```shell
+pip install --upgrade paddlenlp==3.0.0b4
 ```
 
-### 源码安装
+或者可通过以下命令安装最新 develop 分支代码：
 
-如果您想体验最新的版本，可以使用以下命令进行源码安装，支持GitHub和Gitee两种方式。
-```
-pip install --upgrade git+https://github.com/PaddlePaddle/PaddleNLP.git
-
-pip install --upgrade git+https://gitee.com/PaddlePaddle/PaddleNLP.git
+```shell
+pip install --pre --upgrade paddlenlp -f https://www.paddlepaddle.org.cn/whl/paddlenlp.html
 ```
 
+更多关于 PaddlePaddle 和 PaddleNLP 安装的详细教程请查看[Installation](./docs/zh/get_started/installation.rst)。
 
-更多关于PaddlePaddle的安装和PaddleNLP安装详细教程请查看[Installation](./docs/get_started/installation.rst)
+------------------------------------------------------------------------------------------
 
 ## 快速开始
 
-### 数据集快速加载
+### 大模型文本生成
+
+PaddleNLP 提供了方便易用的 Auto API，能够快速的加载模型和 Tokenizer。这里以使用 `Qwen/Qwen2-0.5B` 模型做文本生成为例：
 
 ```python
-from paddlenlp.datasets import load_dataset
-
-train_ds, dev_ds, test_ds = load_dataset("chnsenticorp", splits=["train", "dev", "test"])
+from paddlenlp.transformers import AutoTokenizer, AutoModelForCausalLM
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B")
+# if using CPU, please change float16 to float32
+model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B", dtype="float16")
+input_features = tokenizer("你好！请自我介绍一下。", return_tensors="pd")
+outputs = model.generate(**input_features, max_new_tokens=128)
+print(tokenizer.batch_decode(outputs[0], skip_special_tokens=True))
+# ['我是一个AI语言模型，我可以回答各种问题，包括但不限于：天气、新闻、历史、文化、科学、教育、娱乐等。请问您有什么需要了解的吗？']
 ```
 
-可参考[Dataset文档](./docs/datasets.md)查看更多数据集。
+### 大模型预训练
 
-### 一键加载预训练中文词向量
+```shell
+git clone https://github.com/PaddlePaddle/PaddleNLP.git && cd PaddleNLP # 如已clone或下载PaddleNLP可跳过
+mkdir -p llm/data && cd llm/data
+wget https://bj.bcebos.com/paddlenlp/models/transformers/llama/data/llama_openwebtext_100k.bin
+wget https://bj.bcebos.com/paddlenlp/models/transformers/llama/data/llama_openwebtext_100k.idx
+cd .. # change folder to PaddleNLP/llm
+# 如需使用use_fused_rms_norm=true，需要前往slm/model_zoo/gpt-3/external_ops安装fused_ln
+python -u run_pretrain.py ./config/qwen/pretrain_argument_0p5b.json
+```
+
+### 大模型 SFT 精调
+
+```shell
+git clone https://github.com/PaddlePaddle/PaddleNLP.git && cd PaddleNLP # 如已clone或下载PaddleNLP可跳过
+mkdir -p llm/data && cd llm/data
+wget https://bj.bcebos.com/paddlenlp/datasets/examples/AdvertiseGen.tar.gz && tar -zxvf AdvertiseGen.tar.gz
+cd .. # change folder to PaddleNLP/llm
+python -u run_finetune.py ./config/qwen/sft_argument_0p5b.json
+```
+
+更多大模型全流程步骤，请参考[飞桨大模型套件](./llm)介绍。
+另外我们还提供了快速微调方式, 无需 clone 源代码：
 
 ```python
-from paddlenlp.embeddings import TokenEmbedding
+from paddlenlp.trl import SFTConfig, SFTTrainer
+from datasets import load_dataset
 
-wordemb = TokenEmbedding("w2v.baidu_encyclopedia.target.word-word.dim300")
-print(wordemb.cosine_sim("国王", "王后"))
->>> 0.63395125
-wordemb.cosine_sim("艺术", "火车")
->>> 0.14792643
+dataset = load_dataset("ZHUI/alpaca_demo", split="train")
+
+training_args = SFTConfig(output_dir="Qwen/Qwen2.5-0.5B-SFT", device="gpu")
+trainer = SFTTrainer(
+    args=training_args,
+    model="Qwen/Qwen2.5-0.5B-Instruct",
+    train_dataset=dataset,
+)
+trainer.train()
 ```
 
-内置50+中文词向量，更多使用方法请参考[Embedding文档](./examples/word_embedding/README.md)。
+更多 PaddleNLP 内容可参考：
 
+* [精选模型库](./slm/model_zoo)，包含优质预训练模型的端到端全流程使用。
+* [多场景示例](./slm/examples)，了解如何使用 PaddleNLP 解决 NLP 多种技术问题，包含基础技术、系统应用与拓展应用。
+* [交互式教程](https://aistudio.baidu.com/aistudio/personalcenter/thirdview/574995)，在🆓免费算力平台 AI Studio 上快速学习 PaddleNLP。
 
-### 一键加载高质量中文预训练模型
+------------------------------------------------------------------------------------------
 
-```python
-from paddlenlp.transformers import ErnieModel, BertModel, RobertaModel, ElectraModel, GPTForPretraining
+## 社区交流
 
-ernie = ErnieModel.from_pretrained('ernie-1.0')
-bert = BertModel.from_pretrained('bert-wwm-chinese')
-roberta = RobertaModel.from_pretrained('roberta-wwm-ext')
-electra = ElectraModel.from_pretrained('chinese-electra-small')
-gpt = GPTForPretraining.from_pretrained('gpt-cpm-large-cn')
-```
-
-### 便捷获取文本特征
-
-```python
-import paddle
-from paddlenlp.transformers import ErnieTokenizer, ErnieModel
-
-tokenizer = ErnieTokenizer.from_pretrained('ernie-1.0')
-model = ErnieModel.from_pretrained('ernie-1.0')
-
-text = tokenizer('自然语言处理')
-pooled_output, sequence_output = model(input_ids=paddle.to_tensor([text['input_ids']]))
-```
-
-请参考[Transformer API文档](./docs/transformers.md)查看目前支持的预训练模型。
-
-## 模型库及其应用
-
-PaddleNLP模型库整体介绍请参考文档[PaddleNLP Model Zoo](./docs/model_zoo.md)。
-模型应用场景介绍请参考[PaddleNLP Examples](./examples/)。
-
-- [词向量](./examples/word_embedding/)
-- [词法分析](./examples/lexical_analysis/)
-- [命名实体识别](./examples/information_extraction/msra_ner/)
-- [语言模型](./examples/language_model/)
-- [文本分类](./examples/text_classification/)
-- [文本生成](./examples/text_generation/)
-- [语义匹配](./examples/text_matching/)
-- [文本图学习](./examples/text_graph/erniesage/)
-- [信息抽取](./examples/information_extraction/)
-- [通用对话](./examples/dialogue/)
-- [机器翻译](./examples/machine_translation/)
-- [阅读理解](./examples/machine_reading_comprehension/)
-
-## 进阶应用
-
-- [模型压缩](./examples/model_compression/)
-
-## API 使用文档
-
-- [Transformer API](./docs/transformers.md)
-  * 基于Transformer结构相关的预训练模型API，包含ERNIE, BERT, RoBERTa, Electra等主流经典结构和下游任务。
-- [Data API](./docs/data.md)
-  * 文本数据处理Pipeline的相关API说明。
-- [Dataset API](./docs/datasets.md)
-  * 数据集相关API，包含自定义数据集，数据集贡献与数据集快速加载等功能说明。
-- [Embedding API](./docs/embeddings.md)
-  * 词向量相关API，支持一键快速加载包预训练的中文词向量，VisulDL高维可视化等功能说明。
-- [Metrics API](./docs/metrics.md)
-  * 针对NLP场景的评估指标说明，与飞桨2.0框架高层API兼容。
-
-
-## 交互式Notebook教程
-
-- [使用Seq2Vec模块进行句子情感分类](https://aistudio.baidu.com/aistudio/projectdetail/1283423)
-- [如何通过预训练模型Fine-tune下游任务](https://aistudio.baidu.com/aistudio/projectdetail/1294333)
-- [使用BiGRU-CRF模型完成快递单信息抽取](https://aistudio.baidu.com/aistudio/projectdetail/1317771)
-- [使用预训练模型ERNIE优化快递单信息抽取](https://aistudio.baidu.com/aistudio/projectdetail/1329361)
-- [使用Seq2Seq模型完成自动对联](https://aistudio.baidu.com/aistudio/projectdetail/1321118)
-- [使用预训练模型ERNIE-GEN实现智能写诗](https://aistudio.baidu.com/aistudio/projectdetail/1339888)
-- [使用TCN网络完成新冠疫情病例数预测](https://aistudio.baidu.com/aistudio/projectdetail/1290873)
-
-更多教程参见[PaddleNLP on AI Studio](https://aistudio.baidu.com/aistudio/personalcenter/thirdview/574995)。
-
-
-
-## 最近更新（自2.0.0版本）
-
-### paddlenlp文本领域API
-
-- 修复Windows CPU环境下的import问题。
-
-## 社区贡献与技术交流
-
-### 特殊兴趣小组
-- 欢迎您加入PaddleNLP的SIG社区，贡献优秀的模型实现、公开数据集、教程与案例、外围小工具。
-
-### QQ
-- 现在就加入PaddleNLP的QQ技术交流群，一起交流NLP技术吧！⬇️
+* 微信扫描二维码并填写问卷，即可加入交流群与众多社区开发者以及官方团队深度交流.
 
 <div align="center">
-  <img src="./docs/imgs/qq.png" width="200" height="200" />
-</div>  
+    <img src="https://github.com/user-attachments/assets/3a58cc9f-69c7-4ccb-b6f5-73e966b8051a" width="150" height="150" />
+</div>
 
-### Slack
-- 欢迎加入[PaddleNLP Slack channel](https://paddlenlp.slack.com/)与我们的开发者进行技术交流。
+## Citation
+
+如果 PaddleNLP 对您的研究有帮助，欢迎引用
+
+```bibtex
+@misc{=paddlenlp,
+    title={PaddleNLP: An Easy-to-use and High Performance NLP Library},
+    author={PaddleNLP Contributors},
+    howpublished = {\url{https://github.com/PaddlePaddle/PaddleNLP}},
+    year={2021}
+}
+```
+
+## Acknowledge
+
+我们借鉴了 Hugging Face 的[Transformers](https://github.com/huggingface/transformers)🤗关于预训练模型使用的优秀设计，在此对 Hugging Face 作者及其开源社区表示感谢。
 
 ## License
 
-PaddleNLP遵循[Apache-2.0开源协议](./LICENSE)。
+PaddleNLP 遵循[Apache-2.0开源协议](./LICENSE)。

README_en.md

@@ -1,171 +1,168 @@
-English | [简体中文](./README.md)
+[简体中文🀄](./README.md) | **English🌎**
 
 <p align="center">
-  <img src="./docs/imgs/paddlenlp.png" width="520" height ="100" />
+  <img src="https://user-images.githubusercontent.com/1371212/175816733-8ec25eb0-9af3-4380-9218-27c154518258.png" align="middle"  width="500" />
 </p>
 
----------------------------------------------------------------------------------
+------------------------------------------------------------------------------------------
 
-[![PyPI - PaddleNLP Version](https://img.shields.io/pypi/v/paddlenlp.svg?label=pip&logo=PyPI&logoColor=white)](https://pypi.org/project/paddlenlp/)
-[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/paddlenlp)](https://pypi.org/project/paddlenlp/)
-[![PyPI Status](https://pepy.tech/badge/paddlenlp/month)](https://pepy.tech/project/paddlenlp)
-![support os](https://img.shields.io/badge/os-linux%2C%20win%2C%20mac-yellow.svg)
-![GitHub](https://img.shields.io/github/license/paddlepaddle/paddlenlp)
+<p align="center">
+    <a href="https://paddlenlp.readthedocs.io/en/latest/?badge=latest"><img src="https://readthedocs.org/projects/paddlenlp/badge/?version=latest">
+    <a href="https://github.com/PaddlePaddle/PaddleNLP/releases"><img src="https://img.shields.io/github/v/release/PaddlePaddle/PaddleNLP?color=ffa"></a>
+    <a href=""><img src="https://img.shields.io/badge/python-3.7+-aff.svg"></a>
+    <a href=""><img src="https://img.shields.io/badge/os-linux%2C%20win%2C%20mac-pink.svg"></a>
+    <a href="https://github.com/PaddlePaddle/PaddleNLP/graphs/contributors"><img src="https://img.shields.io/github/contributors/PaddlePaddle/PaddleNLP?color=9ea"></a>
+    <a href="https://github.com/PaddlePaddle/PaddleNLP/commits"><img src="https://img.shields.io/github/commit-activity/m/PaddlePaddle/PaddleNLP?color=3af"></a>
+    <a href="https://pypi.org/project/paddlenlp/"><img src="https://img.shields.io/pypi/dm/paddlenlp?color=9cf"></a>
+    <a href="https://github.com/PaddlePaddle/PaddleNLP/issues"><img src="https://img.shields.io/github/issues/PaddlePaddle/PaddleNLP?color=9cc"></a>
+    <a href="https://github.com/PaddlePaddle/PaddleNLP/stargazers"><img src="https://img.shields.io/github/stars/PaddlePaddle/PaddleNLP?color=ccf"></a>
+    <a href="./LICENSE"><img src="https://img.shields.io/badge/license-Apache%202-dfd.svg"></a>
+</p>
 
-## Introduction
+<h4 align="center">
+    <a href=#Features> Features </a> |
+    <a href=#Support-Models> Supported Models </a> |
+    <a href=#Installation> Installation </a> |
+    <a href=#Quick-start> Quick Start </a> |
+    <a href=#community> Community </a>
+</h4>
 
-PaddleNLP aims to accelerate NLP applications through powerful model zoo, easy-to-use API and high performance distributed training. It's also the NLP best practice for PaddlePaddle 2.0 API system.
+**PaddleNLP** is a Large Language Model (LLM) development suite based on the PaddlePaddle deep learning framework, supporting efficient large model training, lossless compression, and high-performance inference on various hardware devices. With its **simplicity** and **ultimate performance**, PaddleNLP is dedicated to helping developers achieve efficient industrial applications of large models.
+
+## News 📢
+
+* **2024.06.27 [PaddleNLP v3.0 Beta](https://github.com/PaddlePaddle/PaddleNLP/releases/tag/v3.0.0-beta0)**：Embrace large models and experience a complete upgrade. With a unified large model suite, we achieve full-process access to domestically produced computing chips. We fully support industrial-level application processes for large models, such as PaddlePaddle's 4D parallel configuration, efficient fine-tuning strategies, efficient alignment algorithms, and high-performance reasoning. Our developed RsLoRA+ algorithm, full checkpoint storage mechanism Unified Checkpoint, and generalized support for FastFNN and FusedQKV all contribute to the training and inference of large models. We continuously support updates to mainstream models for providing efficient solutions.
+
+* **2024.04.24 [PaddleNLP v2.8](https://github.com/PaddlePaddle/PaddleNLP/releases/tag/v2.8.0)**：Our self-developed RsLoRA+ algorithm with extreme convergence significantly improves the convergence speed and training effectiveness of PEFT training. By introducing high-performance generation acceleration into the RLHF PPO algorithm, we have broken through the generation speed bottleneck in PPO training, achieving a significant lead in PPO training performance. We generally support multiple large model training performance optimization methods such as FastFFN and FusedQKV, making large model training faster and more stable.
 
 ## Features
 
-* **Powerful Model Zoo for Rich Senario**
-  - Our Model Zoo covers mainstream NLP applications, including Lexical Analysis, Text Classification, Text Generation, Text Matching, Text Graph, Information Extraction, Machine Translation, General Dialogue and Question Answering etc.
+### <a href=#Integrated training and inference on multiple hardware platforms> 🔧 Integrated training and inference on multiple hardware platforms </a>
 
-* **Easy-to-Use and End-to-End API**
-  - The API is fully integrated with PaddlePaddle 2.0 high-level API system. It minimizes the number of user actions required for common use cases like data loading, text pre-processing, training and evaluation, which enables you to deal with text problems more productively.
+Our development suit supports large model training and inference on multiple hardware platforms, including NVIDIA GPUs, Kunlun XPUs, Ascend NPUs, Enflame GCUs, and Hygon DCUs. The toolkit's interface allows for quick hardware switching, significantly reducing research and development costs associated with hardware transitions.
 
-* **High Performance and Distributed Training**
--  We provide a highly optimized ditributed training implementation for BERT with Fleet API, and mixed precision training strategy based on PaddlePaddle 2.0, it can fully utilize GPU clusters for large-scale model pre-training.
+### <a href=Efficient and easy-to-use pre-training> 🚀 Efficient and easy-to-use pre-training </a>
 
+We support 4D high-performance training with data parallelism, sharding parallelism, tensor parallelism, and pipeline parallelism. The Trainer supports configurable distributed strategies, reducing the cost associated with complex distributed combinations. The Unified Checkpoint large model storage format supports dynamic scaling of model parameter distribution during training, thereby reducing the migration cost caused by hardware switching.
+
+### <a href=#Efficient fine-tuning> 🤗 Efficient fine-tuning </a>
+
+The fine-tuning algorithms are deeply integrated with zero-padding data streams and high-performance FlashMask operators, reducing invalid data padding and computation during training, and significantly improving the throughput of fine-tuning training.
+
+### <a href=#Lossless compression and high-performance inference> 🎛️ Lossless compression and high-performance inference </a>
+
+The high-performance inference module of the large model toolkit incorporates dynamic insertion and operator fusion strategies throughout the entire process, greatly accelerating parallel inference speed. The underlying implementation details are encapsulated, enabling out-of-the-box high-performance parallel inference capabilities.
+
+## Documentation
+For detailed documentation, visit the [PaddleNLP Documentation](https://paddlenlp.readthedocs.io/).
+
+------------------------------------------------------------------------------------------
+
+## Support Models
+
+Detailed list 👉 [Supported Model List](https://github.com/PaddlePaddle/PaddleNLP/issues/8663)
 
 ## Installation
 
 ### Prerequisites
 
-* python >= 3.6
-* paddlepaddle >= 2.0.1
+* python >= 3.8
+* paddlepaddle >= 3.0.0b0
 
-More information about PaddlePaddle installation please refer to [PaddlePaddle Install](https://www.paddlepaddle.org.cn/install/quick?docurl=/documentation/docs/zh/install/conda/linux-conda.html)
+### Pip Installation
 
-### PIP Installation
-
-```
-pip install --upgrade paddlenlp -i https://pypi.org/simple
+```shell
+pip install --upgrade paddlenlp==3.0.0b3
 ```
 
-### Install from Source
+or you can install the latest develop branch code with the following command:
 
+```shell
+pip install --pre --upgrade paddlenlp -f https://www.paddlepaddle.org.cn/whl/paddlenlp.html
 ```
-pip install --upgrade git+https://github.com/PaddlePaddle/PaddleNLP.git
 
-pip install --upgrade git+https://gitee.com/PaddlePaddle/PaddleNLP.git
-```
+More information about PaddlePaddle installation please refer to [PaddlePaddle's Website](https://www.paddlepaddle.org.cn).
+
+------------------------------------------------------------------------------------------
 
 ## Quick Start
 
-### Quick Dataset Loading
+### Text generation with large language model
+
+PaddleNLP provides a convenient and easy-to-use Auto API, which can quickly load models and Tokenizers. Here, we use the `Qwen/Qwen2-0.5B` large model as an example for text generation:
 
 ```python
-from paddlenlp.datasets import load_dataset
-
-train_ds, dev_ds, test_ds = load_dataset("chnsenticorp", splits=["train", "dev", "test"])
+>>> from paddlenlp.transformers import AutoTokenizer, AutoModelForCausalLM
+>>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B")
+>>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B", dtype="float16")
+>>> input_features = tokenizer("你好！请自我介绍一下。", return_tensors="pd")
+>>> outputs = model.generate(**input_features, max_length=128)
+>>> print(tokenizer.batch_decode(outputs[0], skip_special_tokens=True))
+['我是一个AI语言模型，我可以回答各种问题，包括但不限于：天气、新闻、历史、文化、科学、教育、娱乐等。请问您有什么需要了解的吗？']
 ```
 
-For more dataset API usage please refer to [Dataset API](./docs/datasets.md).
+### Pre-training for large language model
 
-### Pre-trained Text Embedding Loading
-
-```python
-
-from paddlenlp.embeddings import TokenEmbedding
-
-wordemb = TokenEmbedding("fasttext.wiki-news.target.word-word.dim300.en")
-wordemb.cosine_sim("king", "queen")
->>> 0.77053076
-wordemb.cosine_sim("apple", "rail")
->>> 0.29207364
+```shell
+git clone https://github.com/PaddlePaddle/PaddleNLP.git && cd PaddleNLP # if cloned or downloaded, can skip this step
+mkdir -p llm/data && cd llm/data
+wget https://bj.bcebos.com/paddlenlp/models/transformers/llama/data/llama_openwebtext_100k.bin
+wget https://bj.bcebos.com/paddlenlp/models/transformers/llama/data/llama_openwebtext_100k.idx
+cd .. # change folder to PaddleNLP/llm
+python -u -m paddle.distributed.launch --gpus "0,1,2,3,4,5,6,7" run_pretrain.py ./config/llama/pretrain_argument.json
 ```
 
-For more TokenEmbedding usage, please refer to [Embedding API](./docs/embeddings.md)
+### SFT finetuning forlarge language model
 
-### Rich Chinese Pre-trained Models
-
-```python
-from paddlenlp.transformers import ErnieModel, BertModel, RobertaModel, ElectraModel, GPTForPretraining
-
-ernie = ErnieModel.from_pretrained('ernie-1.0')
-bert = BertModel.from_pretrained('bert-wwm-chinese')
-roberta = RobertaModel.from_pretrained('roberta-wwm-ext')
-electra = ElectraModel.from_pretrained('chinese-electra-small')
-gpt = GPTForPretraining.from_pretrained('gpt-cpm-large-cn')
+```shell
+git clone https://github.com/PaddlePaddle/PaddleNLP.git && cd PaddleNLP # if cloned or downloaded, can skip this step
+mkdir -p llm/data && cd llm/data
+wget https://bj.bcebos.com/paddlenlp/datasets/examples/AdvertiseGen.tar.gz && tar -zxvf AdvertiseGen.tar.gz
+cd .. # change folder to PaddleNLP/llm
+python -u -m paddle.distributed.launch --gpus "0,1,2,3,4,5,6,7" run_finetune.py ./config/llama/sft_argument.json
 ```
 
-For more pretrained model selection, please refer to [Transformer API](./docs/transformers.md)
+For more steps in the entire large model process, please refer to the[Large Model Full-Process Suite](./llm).
 
-### Extract Feature Through Pre-trained Model
+For more PaddleNLP content, please refer to:
 
-```python
-import paddle
-from paddlenlp.transformers import ErnieTokenizer, ErnieModel
-
-tokenizer = ErnieTokenizer.from_pretrained('ernie-1.0')
-model = ErnieModel.from_pretrained('ernie-1.0')
-
-text = tokenizer('自然语言处理')
-pooled_output, sequence_output = model.forward(input_ids=paddle.to_tensor([text['input_ids']]))
-```
-
-## Model Zoo and Applications
-
-For model zoo introduction please refer to[PaddleNLP Model Zoo](./docs/model_zoo.md). As for applicaiton senario please refer to [PaddleNLP Examples](./examples/)。
-
-- [Word Embedding](./examples/word_embedding/)
-- [Lexical Analysis](./examples/lexical_analysis/)
-- [Named Entity Recognition](./examples/information_extraction/msra_ner/)
-- [Language Model](./examples/language_model/)
-- [Text Classification](./examples/text_classification/)
-- [Text Gneeration](./examples/text_generation/)
-- [Semantic Maching](./examples/text_matching/)
-- [Text Graph](./examples/text_graph/erniesage/)
-- [Information Extraction](./examples/information_extraction/)
-- [General Dialogue](./examples/dialogue/)
-- [Machine Translation](./examples/machine_translation/)
-- [Machine Readeng Comprehension](./examples/machine_reading_comprehension/)
-
-## Advanced Application
-
-- [Model Compression](./examples/model_compression/)
-
-## API Usage
-
-- [Transformer API](./docs/transformers.md)
-- [Data API](./docs/data.md)
-- [Dataset API](./docs/datasets.md)
-- [Embedding API](./docs/embeddings.md)
-- [Metrics API](./docs/metrics.md)
-
-
-## Tutorials
-
-Please refer to our official AI Studio account for more interactive tutorials: [PaddleNLP on AI Studio](https://aistudio.baidu.com/aistudio/personalcenter/thirdview/574995)
-
-* [What's Seq2Vec?](https://aistudio.baidu.com/aistudio/projectdetail/1283423) shows how to use simple API to finish LSTM model and solve sentiment analysis task.
-
-* [Sentiment Analysis with ERNIE](https://aistudio.baidu.com/aistudio/projectdetail/1294333) shows how to exploit the pretrained ERNIE to solve sentiment analysis problem.
-
-* [Waybill Information Extraction with BiGRU-CRF Model](https://aistudio.baidu.com/aistudio/projectdetail/1317771) shows how to make use of Bi-GRU plus CRF to finish information extraction task.
-
-* [Waybill Information Extraction with ERNIE](https://aistudio.baidu.com/aistudio/projectdetail/1329361) shows how to use ERNIE, the Chinese pre-trained model improve information extraction performance.
-
-* [Use TCN Model to predict COVID-19 confirmed cases](https://aistudio.baidu.com/aistudio/projectdetail/1290873)
+* [Model Library](./slm/model_zoo)，which includes end-to-end usage of high-quality pre-trained models.
+* [Multi-scenario Examples](./slm/examples)，to understand how to use PaddleNLP to solve various NLP technical problems, including basic techniques, system applications, and extended applications.
+* [Interactive Tutorial](https://aistudio.baidu.com/aistudio/personalcenter/thirdview/574995)，to quickly learn PaddleNLP on the free computing platform AI Studio.
 
+------------------------------------------------------------------------------------------
 
 ## Community
 
-### Special Interest Group(SIG)
-
-Welcome to join [PaddleNLP SIG](https://iwenjuan.baidu.com/?code=bkypg8) for contribution, eg. Dataset, Models and Toolkit.
-
 ### Slack
+
 To connect with other users and contributors, welcome to join our [Slack channel](https://paddlenlp.slack.com/).
 
-### QQ
-Join our QQ Technical Group for technical exchange right now! ⬇️
+### WeChat
+
+Scan the QR code below with your Wechat⬇️. You can access to official technical exchange group. Look forward to your participation.
 
 <div align="center">
-  <img src="./docs/imgs/qq.png" width="200" height="200" />
+    <img src="https://github.com/user-attachments/assets/3a58cc9f-69c7-4ccb-b6f5-73e966b8051a" width="150" height="150" />
 </div>
 
+## Citation
+
+If you find PaddleNLP useful in your research, please consider citing
+
+```bibtext
+@misc{=paddlenlp,
+    title={PaddleNLP: An Easy-to-use and High Performance NLP Library},
+    author={PaddleNLP Contributors},
+    howpublished = {\url{https://github.com/PaddlePaddle/PaddleNLP}},
+    year={2021}
+}
+```
+
+## Acknowledge
+
+We have borrowed from Hugging Face's [Transformers](https://github.com/huggingface/transformers)🤗 excellent design on pretrained models usage, and we would like to express our gratitude to the authors of Hugging Face and its open source community.
+
 ## License
 
 PaddleNLP is provided under the [Apache-2.0 License](./LICENSE).

csrc/README.md

@@ -0,0 +1,28 @@
+# PaddleNLP 大模型高性能自定义推理算子
+
+此文档介绍如何编译安装 PaddleNLP 大模型高性能自定义推理算子的安装教程。
+
+使用这些高性能算子，可以大幅提升大模型推理速度。
+大模型推理相关教程详见[此处](https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm/README.md#6-%E6%8E%A8%E7%90%86)。
+
+## 安装 C++ 依赖
+
+```shell
+pip install -r requirements.txt
+```
+
+## 编译 Cuda 算子
+
+```shell
+python setup_cuda.py install
+```
+
+### FP8 GEMM 自动调优
+
+确保 `cutlass` 库已经安装，然后执行以下命令进行自动调优。
+- 对于89架构的 GPU，CUDA 版本至少为12.4
+- 对于90架构的 GPU，CUDA 版本至少为12.0
+```shell
+cd tools
+sh tune_fp8_gemm.sh
+```

csrc/cpu/0001-fp16_bf16.patch

@@ -0,0 +1,619 @@
+From e479452fdedfd1f0a54cc1134966f2235df14874 Mon Sep 17 00:00:00 2001
+From: bukejiyu <395822456@qq.com>
+Date: Wed, 21 Aug 2024 22:01:52 +0800
+Subject: [PATCH] fp16_bf16
+
+---
+ cmake/mkl.cmake                  |   2 +-
+ cmake/mklml.cmake                |   2 +-
+ cmake/onednn.cmake               |   1 +
+ cmake/xdnn.cmake                 |   2 +-
+ include/dtype.h                  |   5 +
+ include/layers_decoder.h         |  21 +--
+ src/layers/attention.h           |   9 +-
+ src/layers/decoder_layer.cpp     | 256 ++++++++++++++++++++++++-------
+ src/layers/decoder_layer.h       |   4 +-
+ tests/ut/layers_decoder_test.cpp |  52 ++++---
+ 10 files changed, 255 insertions(+), 99 deletions(-)
+
+diff --git a/cmake/mkl.cmake b/cmake/mkl.cmake
+index 0ef2e66..92c6d06 100644
+--- a/cmake/mkl.cmake
++++ b/cmake/mkl.cmake
+@@ -25,7 +25,7 @@ set(MKL_3rdparty_DIR "${CMAKE_SOURCE_DIR}/3rdparty/mkl")
+ if(NOT EXISTS ${MKL_3rdparty_DIR})
+     find_package(Python COMPONENTS Interpreter Development)
+     execute_process(COMMAND ${Python_EXECUTABLE} -m pip install --force-reinstall
+-                            --prefix=${MKL_3rdparty_DIR} mkl-static==2024.0.0 mkl-include==2024.0.0
++                            --prefix=${MKL_3rdparty_DIR} mkl-static==2024.0.0 mkl-include==2024.0.0 -i https://mirrors.aliyun.com/pypi/simple
+                     RESULT_VARIABLE EXIT_CODE)
+ 
+     if(NOT ${EXIT_CODE} EQUAL 0)
+diff --git a/cmake/mklml.cmake b/cmake/mklml.cmake
+index 4baec46..d89ce46 100644
+--- a/cmake/mklml.cmake
++++ b/cmake/mklml.cmake
+@@ -28,7 +28,7 @@ include(ExternalProject)
+ ExternalProject_Add(mklml
+   URL               https://github.com/oneapi-src/oneDNN/releases/download/v0.21/mklml_lnx_2019.0.5.20190502.tgz
+   URL_HASH          MD5=dfcea335652dbf3518e1d02cab2cea97
+-  TIMEOUT           60
++  TIMEOUT           360
+   SOURCE_DIR        ${CMAKE_SOURCE_DIR}/3rdparty/mklml
+   CONFIGURE_COMMAND ""
+   BUILD_COMMAND     ""
+diff --git a/cmake/onednn.cmake b/cmake/onednn.cmake
+index 8efabd1..bfa558a 100644
+--- a/cmake/onednn.cmake
++++ b/cmake/onednn.cmake
+@@ -36,6 +36,7 @@ if(NOT EXISTS ${ONEDNN_3rdparty_DIR})
+     ExternalProject_Add(onednn
+       GIT_REPOSITORY    https://github.com/oneapi-src/oneDNN.git
+       GIT_TAG           v3.5
++      TIMEOUT           360
+       SOURCE_DIR        ${ONEDNN_3rdparty_DIR}
+       BINARY_DIR        ${ONEDNN_3rdparty_DIR}
+       CONFIGURE_COMMAND ${CMAKE_COMMAND} -E make_directory "build" && ${CMAKE_COMMAND} -E chdir "build" ${CMAKE_COMMAND} ${ONEDNN_BUILD_OPTIONS} ..
+diff --git a/cmake/xdnn.cmake b/cmake/xdnn.cmake
+index 7c0e051..721bfe3 100644
+--- a/cmake/xdnn.cmake
++++ b/cmake/xdnn.cmake
+@@ -28,7 +28,7 @@ include(ExternalProject)
+ ExternalProject_Add(xdnn_lib
+   URL               https://github.com/intel/xFasterTransformer/releases/download/IntrinsicGemm/xdnn_v1.5.2.tar.gz
+   URL_HASH          MD5=884f2e1e2c846ff19f33c889681f8dc2
+-  TIMEOUT           120
++  TIMEOUT           360
+   SOURCE_DIR        ${CMAKE_SOURCE_DIR}/3rdparty/xdnn
+   CONFIGURE_COMMAND ""
+   BUILD_COMMAND     ""
+diff --git a/include/dtype.h b/include/dtype.h
+index de72bce..9e0a448 100644
+--- a/include/dtype.h
++++ b/include/dtype.h
+@@ -31,6 +31,7 @@ enum DataType {
+     w8a8_int8,
+     w8a8_int4,
+     w8a8_nf4,
++    fp16_int8,
+     unknown,
+ };
+ 
+@@ -51,4 +52,8 @@ enum ActivationType {
+     SILU,
+ };
+ 
++enum RopeType {
++    LLAMA_ROPE = 0,
++};
++
+ } // namespace xft
+diff --git a/include/layers_decoder.h b/include/layers_decoder.h
+index 34f6aa5..a30e34d 100644
+--- a/include/layers_decoder.h
++++ b/include/layers_decoder.h
+@@ -17,13 +17,14 @@
+ #include "dtype.h"
+ 
+ namespace xft {
+-
+-void invokeLayerLLaMA(DataType dt, ActivationType at, NormType nt, int batchSize, int inputSeqLen, int attHeadDim,
+-        int attHeadNum, int kvHeadNum, int maxPositions, int maxPosEmbed, int pastSeqLen, int currentSeqLen, int step,
+-        int hiddenSize, int intermediateSize, void *output, int outputStride, const void *input, int inputStride,
+-        const float *ln1Gamma, const float *ln1Beta, const void *queryWeight, const void *keyWeight,
+-        const void *valueWeight, const void *attnOutWeight, const float *ln2Gamma, const float *ln2Beta,
+-        const void *gateWeight, const void *upWeight, const void *downWeight, const float *queryBias = nullptr,
+-        const float *keyBias = nullptr, const float *valueBias = nullptr, const float *attnOutBias = nullptr);
+-
+-} // namespace xft
+\ No newline at end of file
++void invokeLayerLLaMA(DataType dt, DataType kvcdt, RopeType rt, ActivationType at, NormType nt, int batchSize,
++        int inputSeqLen, int attHeadDim, int attHeadNum, int kvHeadNum, int maxPositions, int maxPosEmbed,
++        int pastSeqLen, int currentSeqLen, int step, int hiddenSize, int intermediateSize, void *output,
++        int outputStride, const void *input, int inputStride, const float *ln1Gamma, const float *ln1Beta,
++        const void *queryWeight, const void *keyWeight, const void *valueWeight, const void *attnOutWeight,
++        const float *ln2Gamma, const float *ln2Beta, const void *gateWeight, const void *upWeight,
++        const void *downWeight, const float *queryBias = nullptr, const float *keyBias = nullptr,
++        const float *valueBias = nullptr, const float *attnOutBias = nullptr, const void *myqkvWeight = nullptr,
++        const float *gateBias = nullptr, const float *upBias = nullptr, const float *downBias = nullptr,
++        const float *myqkvBias = nullptr);
++} // namespace xft
+diff --git a/src/layers/attention.h b/src/layers/attention.h
+index 092b3d6..b438837 100644
+--- a/src/layers/attention.h
++++ b/src/layers/attention.h
+@@ -84,7 +84,8 @@ public:
+             const float *queryBias, const OriWeiT *keyWeight, const float *keyScale, const float *keyZero,
+             const float *keyBias, const OriWeiT *valueWeight, const float *valueScale, const float *valueZero,
+             const float *valueBias, const OriWeiT *attnOutWeight, const float *attnOutScale, const float *attnOutZero,
+-            const float *attnOutBias, bool doLNorm, const float *gamma1, const float *beta1, bool trans = true) {
++            const float *attnOutBias, bool doLNorm, const float *gamma1, const float *beta1, bool trans = true,
++            const OriWeiT *myqkvWeight = nullptr) {
+         int hiddenSize = ctx->hiddenSize;
+         int headSize = ctx->attHeadSize;
+ 
+@@ -107,7 +108,10 @@ public:
+                     valueWeight + this->startKVHead * headSize * hiddenSize / sizeFactor,
+                     hiddenSize * kvResponsibleCols * sizeof(OriWeiT) / sizeFactor);
+         } else {
+-            int qkvStride = (ctx->attHeadNum + ctx->kvHeadNum + ctx->kvHeadNum) * ctx->attHeadSize;
++            if (myqkvWeight != nullptr) {
++                memcpy(concatBuf, myqkvWeight, hiddenSize * responsibleCols * sizeof(OriWeiT) / sizeFactor);
++            } else {
++                int qkvStride = (ctx->attHeadNum + ctx->kvHeadNum + ctx->kvHeadNum) * ctx->attHeadSize;
+ #pragma omp parallel for
+             for (int i = 0; i < hiddenSize; ++i) {
+                 memcpy(concatBuf + i * responsibleCols / sizeFactor,
+@@ -120,6 +124,7 @@ public:
+                                 + kvResponsibleCols / sizeFactor,
+                         valueWeight + i * qkvStride / sizeFactor + this->startKVHead * headSize / sizeFactor,
+                         kvResponsibleCols * sizeof(OriWeiT) / sizeFactor);
++                }
+             }
+         }
+         float *concatScale = nullptr;
+diff --git a/src/layers/decoder_layer.cpp b/src/layers/decoder_layer.cpp
+index 02f13cb..0f30f21 100644
+--- a/src/layers/decoder_layer.cpp
++++ b/src/layers/decoder_layer.cpp
+@@ -21,19 +21,21 @@
+ #include "layers_mlp.h"
+ #include "mlp_llama.h"
+ #include "rms_norm.h"
++#include "numa_allocator.h"
+ 
+ #include <unordered_map>
+ 
+ namespace xft {
+ 
+-template <typename DataT, typename NormT>
++template <typename DataT, typename KVCacheT, typename RopeT, typename NormT>
+ void LayerLLaMAImpl(DataType dt, ActivationType at, NormType nt, int batchSize, int inputSeqLen, int attHeadDim,
+         int attHeadNum, int kvHeadNum, int maxPositions, int maxPosEmbed, int pastSeqLen, int currentSeqLen, int step,
+         int hiddenSize, int intermediateSize, void *output, int outputStride, const void *input, int inputStride,
+         const float *ln1Gamma, const float *ln1Beta, const void *queryWeight, const void *keyWeight,
+         const void *valueWeight, const void *attnOutWeight, const float *ln2Gamma, const float *ln2Beta,
+         const void *gateWeight, const void *upWeight, const void *downWeight, const float *queryBias,
+-        const float *keyBias, const float *valueBias, const float *attnOutBias) {
++        const float *keyBias, const float *valueBias, const float *attnOutBias,
++        MMHelper *mmHelper, DecoderContext *ctx, KVCacheManager<KVCacheT> *kvCacheMgr,const void *myqkvWeight = nullptr) {
+ 
+     // TODO: will deprecate attention mask in future, so need to change this
+     auto prepareAttnMask = [&](DecoderContext *ctx, int step) {
+@@ -83,67 +85,46 @@ void LayerLLaMAImpl(DataType dt, ActivationType at, NormType nt, int batchSize,
+         return mask;
+     };
+ 
+-    using DECODER = Decoder<Attention<DataT, LlamaRotaryEmbedding, NormT>, LlamaMLP<DataT>>;
+-    static std::unordered_map<std::string, DECODER *> llama_layer_hub;
+-    static MMHelper *mmHelper;
+-    static DecoderContext *ctx;
+-    static KVCacheManager<float16_t> *kvCacheMgr;
+-
+-    std::string actType;
+-    if (at == ActivationType::SILU)
+-        actType = "silu";
+-    else if (at == ActivationType::RELU)
+-        actType = "relu";
+-    else if (at == ActivationType::GELU)
+-        actType = "gelu";
+-    else if (at == ActivationType::SWIGLU)
+-        actType = "swiglu";
+-    else
+-        printf(">> unsupported activation type\n");
+-
+-    if (ctx == nullptr
+-            || (ctx != nullptr && (ctx->hiddenSize != hiddenSize || ctx->intermediateSize != intermediateSize))) {
+-        if (ctx != nullptr) delete ctx;
+-        printf(">> create context: %d %d\n", hiddenSize, intermediateSize);
+-        mmHelper = new MMHelper(Env::getInstance().getEngineKind(), Env::getInstance().getEngineIndex());
+-        ctx = new DecoderContext(1, hiddenSize, attHeadDim, attHeadNum, kvHeadNum, intermediateSize, actType, 1e-6, 0,
+-                0, maxPositions, maxPosEmbed, -1, 0, 1, mmHelper);
+-        if (kvCacheMgr != nullptr) delete kvCacheMgr;
+-        kvCacheMgr = new KVCacheManager<float16_t>(1);
+-    }
+-
++    using DECODER = Decoder<Attention<DataT, RopeT, NormT>, LlamaMLP<DataT>>;
++    DECODER *llama_layer;
++    static xft::Matrix<float> actBuffers ;
++    //static std::unordered_map<std::string, DECODER *> llama_layer_hub;
++    static std::unordered_map<std::string, std::tuple<DECODER*>> llama_layer_hub;
+     // create hash key and value: if hidden and intermediateSize is changed , then memory pointer is also changed.
+     std::stringstream weights_addr;
+     weights_addr << queryWeight << "_" << keyWeight << "_" << valueWeight << "_" << attnOutWeight << "_" << gateWeight
+                  << "_" << upWeight << "_" << downWeight << "_" << dt << "_" << at << "_" << nt << "_" << attHeadDim
+                  << "_" << attHeadNum << "_" << kvHeadNum;
+     std::string llama_layer_key = weights_addr.str();
+-    DECODER *llama_layer;
+ 
+     auto it_created = llama_layer_hub.find(llama_layer_key);
+     if (it_created == llama_layer_hub.end()) {
++        int firstNode = getenv("FIRST_TOKEN_WEIGHT_LOCATION") ? atoi(getenv("FIRST_TOKEN_WEIGHT_LOCATION")) : -1;
++        int nextNode = getenv("NEXT_TOKEN_WEIGHT_LOCATION") ? atoi(getenv("NEXT_TOKEN_WEIGHT_LOCATION")) : -1;
++        if (step == 0)
++                xft_set_preferred_node(firstNode);
++        else
++                xft_set_preferred_node(nextNode);
+         llama_layer = new DECODER(ctx, 0);
+         llama_layer->setWeights(ctx, (const float *)queryWeight, nullptr, nullptr, queryBias, (const float *)keyWeight,
+                 nullptr, nullptr, keyBias, (const float *)valueWeight, nullptr, nullptr, valueBias,
+                 (const float *)attnOutWeight, nullptr, nullptr, attnOutBias, ln1Gamma, ln1Beta,
+                 (const float *)gateWeight, nullptr, nullptr, nullptr, (const float *)upWeight, nullptr, nullptr,
+-                nullptr, ln2Gamma, ln2Beta, (const float *)downWeight, nullptr, nullptr, false);
+-        llama_layer_hub[llama_layer_key] = llama_layer;
+-        printf(">> create llama_layer_key: %s\n", llama_layer_key.c_str());
++                nullptr, ln2Gamma, ln2Beta, (const float *)downWeight, nullptr, nullptr, false,(const float *)myqkvWeight);
++
++
++        llama_layer_hub[llama_layer_key] = std::make_tuple(llama_layer);;
++        // printf(">> create llama_layer_key: %s\n", llama_layer_key.c_str());
++        xft_set_preferred_node(-1);
+     } else {
+-        llama_layer = it_created->second;
++        llama_layer = std::get<0>(it_created->second);
+     }
+-
+-    ctx->resize(batchSize, inputSeqLen, pastSeqLen);
+-    xft::Matrix<float> actBuffers;
+     actBuffers.Resize(batchSize * inputSeqLen * 2, hiddenSize);
++    ctx->resize(batchSize, inputSeqLen, pastSeqLen);
+     float *attnMask = prepareAttnMask(ctx, step);
+ 
+-    int workers = 1;
+-    int headsPerSplit = (ctx->kvHeadNum + workers - 1) / workers;
+-    kvCacheMgr->resize(maxPositions, batchSize, headsPerSplit, attHeadDim);
+-    KVCacheTensor<float16_t> &presentKey = kvCacheMgr->getKey(0);
+-    KVCacheTensor<float16_t> &presentValue = kvCacheMgr->getValue(0);
++    KVCacheTensor<KVCacheT> &presentKey = kvCacheMgr->getKey(0);
++    KVCacheTensor<KVCacheT> &presentValue = kvCacheMgr->getValue(0);
+ 
+     float *attnOut = (float *)(ctx->tmpBuf.Data());
+ 
+@@ -159,45 +140,168 @@ void LayerLLaMAImpl(DataType dt, ActivationType at, NormType nt, int batchSize,
+     llama_layer->forwardFFN(ctx, attnOut, (float *)output, inputStride, outputStride, true);
+ }
+ 
+-void invokeLayerLLaMA(DataType dt, ActivationType at, NormType nt, int batchSize, int inputSeqLen, int attHeadDim,
++template <typename KVCacheT, typename RopeT>
++void LayerLLaMAWrapper(DataType dt, ActivationType at, NormType nt, int batchSize, int inputSeqLen, int attHeadDim,
+         int attHeadNum, int kvHeadNum, int maxPositions, int maxPosEmbed, int pastSeqLen, int currentSeqLen, int step,
+         int hiddenSize, int intermediateSize, void *output, int outputStride, const void *input, int inputStride,
+         const float *ln1Gamma, const float *ln1Beta, const void *queryWeight, const void *keyWeight,
+         const void *valueWeight, const void *attnOutWeight, const float *ln2Gamma, const float *ln2Beta,
+         const void *gateWeight, const void *upWeight, const void *downWeight, const float *queryBias,
+-        const float *keyBias, const float *valueBias, const float *attnOutBias) {
++        const float *keyBias, const float *valueBias, const float *attnOutBias,const void *myqkvWeight=nullptr) {
+     static std::mutex mutex;
+     std::lock_guard<std::mutex> lock(mutex);
+ 
++    std::string actType;
++    if (at == ActivationType::SILU)
++        actType = "silu";
++    else if (at == ActivationType::RELU)
++        actType = "relu";
++    else if (at == ActivationType::GELU)
++        actType = "gelu";
++    else if (at == ActivationType::SWIGLU)
++        actType = "swiglu";
++    else {
++        printf(">> unsupported activation type\n");
++        return;
++    }
++
++    static MMHelper *mmHelper;
++    static DecoderContext *ctx;
++    if (ctx == nullptr
++            || (ctx != nullptr && (ctx->hiddenSize != hiddenSize || ctx->intermediateSize != intermediateSize))) {
++        if (ctx != nullptr) delete ctx;
++        // printf(">> create context: %d %d\n", hiddenSize, intermediateSize);
++        mmHelper = new MMHelper(Env::getInstance().getEngineKind(), Env::getInstance().getEngineIndex());
++        ctx = new DecoderContext(1, hiddenSize, attHeadDim, attHeadNum, kvHeadNum, intermediateSize, actType, 1e-6, 0,
++                0, maxPositions, maxPosEmbed, -1, 0, 1, mmHelper);
++    }
++
++    KVCacheManager<KVCacheT> *kvCacheMgr;
++    static std::unordered_map<std::string, KVCacheManager<KVCacheT> *> kv_hub;
++
++    // create hash key and value: if hidden and intermediateSize is changed , then memory pointer is also changed.
++    std::stringstream layer_key;
++    layer_key << queryWeight << "_" << keyWeight << "_" << valueWeight << "_" << attnOutWeight << "_" << gateWeight
++                 << "_" << upWeight << "_" << downWeight << "_" << dt << "_" << at << "_" << nt << "_" << attHeadDim
++                 << "_" << attHeadNum << "_" << kvHeadNum;
++    std::string kv_hub_key = layer_key.str();
++
++    auto it_created = kv_hub.find(kv_hub_key);
++    if (it_created == kv_hub.end()) {
++        int kvcNode = getenv("KVCACHE_LOCATION") ? atoi(getenv("KVCACHE_LOCATION")) : -1;
++        xft_set_preferred_node(kvcNode);
++        kvCacheMgr = new KVCacheManager<KVCacheT>(1);
++        int workers = 1;
++        int headsPerSplit = (ctx->kvHeadNum + workers - 1) / workers;
++        kvCacheMgr->resize(maxPositions, batchSize, headsPerSplit, attHeadDim);
++        kv_hub[kv_hub_key] = kvCacheMgr;
++        // printf(">> create kv_hub_key: %s\n", kv_hub_key.c_str());
++        xft_set_preferred_node(-1);
++    } else {
++        kvCacheMgr = it_created->second;
++    }
++
+     if (dt == DataType::bf16) {
+-        if (nt == NormType::RMS)
+-            LayerLLaMAImpl<bfloat16_t, RmsNorm>(dt, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++        if (nt == NormType::RMS) {
++            LayerLLaMAImpl<bfloat16_t, KVCacheT, RopeT, RmsNorm>(dt, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
+                     maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
+                     outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
+                     attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
+-                    attnOutBias);
+-        else if (nt == NormType::LN) {
+-            LayerLLaMAImpl<bfloat16_t, LayerNorm>(dt, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                    attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++        } else if (nt == NormType::LN) {
++            LayerLLaMAImpl<bfloat16_t, KVCacheT, RopeT, LayerNorm>(dt, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
+                     maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
+                     outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
+                     attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
+-                    attnOutBias);
++                    attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
+         } else {
+             printf(">> unsupported norm type\n");
+         }
+     } else if (dt == DataType::fp16) {
+-        if (nt == NormType::RMS)
+-            LayerLLaMAImpl<float16_t, RmsNorm>(dt, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++        if (nt == NormType::RMS) {
++            LayerLLaMAImpl<float16_t, KVCacheT, RopeT, RmsNorm>(dt, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
+                     maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
+                     outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
+                     attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
+-                    attnOutBias);
+-        else if (nt == NormType::LN) {
+-            LayerLLaMAImpl<float16_t, LayerNorm>(dt, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                    attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++        } else if (nt == NormType::LN) {
++            LayerLLaMAImpl<float16_t, KVCacheT, RopeT, LayerNorm>(dt, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
+                     maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
+                     outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
+                     attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
+-                    attnOutBias);
++                    attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++        } else {
++            printf(">> unsupported norm type\n");
++        }
++    } else if (dt == DataType::bf16_int8) {
++        if (nt == NormType::RMS) {
++            auto firstTokenFunc = LayerLLaMAImpl<bfloat16_t, KVCacheT, RopeT, RmsNorm>;
++            auto nextTokenFunc = LayerLLaMAImpl<int8_t, KVCacheT, RopeT, RmsNorm>;
++            if (step == 0) {
++                    firstTokenFunc(DataType::bf16, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                        maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
++                        outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
++                        attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
++                        attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++
++            } else {
++                    nextTokenFunc(DataType::int8, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                        maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
++                        outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
++                        attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
++                        attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++            }
++        } else if (nt == NormType::LN) {
++            auto firstTokenFunc = LayerLLaMAImpl<bfloat16_t, KVCacheT, RopeT, LayerNorm>;
++            auto nextTokenFunc = LayerLLaMAImpl<int8_t, KVCacheT, RopeT, LayerNorm>;
++            if (step == 0)
++                    firstTokenFunc(DataType::bf16, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                        maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
++                        outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
++                        attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
++                        attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++            else
++                    nextTokenFunc(DataType::int8, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                        maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
++                        outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
++                        attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
++                        attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++        } else {
++            printf(">> unsupported norm type\n");
++        }
++    } else if (dt == DataType::fp16_int8) {
++        if (nt == NormType::RMS) {
++            auto firstTokenFunc = LayerLLaMAImpl<float16_t, KVCacheT, RopeT, RmsNorm>;
++            auto nextTokenFunc = LayerLLaMAImpl<int8_t, KVCacheT, RopeT, RmsNorm>;
++            if (step == 0) {
++                    firstTokenFunc(DataType::fp16, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                        maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
++                        outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
++                        attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
++                        attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++
++            } else {
++                    nextTokenFunc(DataType::int8, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                        maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
++                        outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
++                        attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
++                        attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++            }
++        } else if (nt == NormType::LN) {
++            auto firstTokenFunc = LayerLLaMAImpl<bfloat16_t, KVCacheT, RopeT, LayerNorm>;
++            auto nextTokenFunc = LayerLLaMAImpl<int8_t, KVCacheT, RopeT, LayerNorm>;
++            if (step == 0)
++                    firstTokenFunc(DataType::fp16, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                        maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
++                        outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
++                        attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
++                        attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++            else
++                    nextTokenFunc(DataType::int8, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                        maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
++                        outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
++                        attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
++                        attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
+         } else {
+             printf(">> unsupported norm type\n");
+         }
+@@ -206,4 +310,40 @@ void invokeLayerLLaMA(DataType dt, ActivationType at, NormType nt, int batchSize
+     }
+ }
+ 
++void invokeLayerLLaMA(DataType dt, DataType kvcdt, RopeType rt, ActivationType at, NormType nt, int batchSize, int inputSeqLen, int attHeadDim,
++        int attHeadNum, int kvHeadNum, int maxPositions, int maxPosEmbed, int pastSeqLen, int currentSeqLen, int step,
++        int hiddenSize, int intermediateSize, void *output, int outputStride, const void *input, int inputStride,
++        const float *ln1Gamma, const float *ln1Beta, const void *queryWeight, const void *keyWeight,
++        const void *valueWeight, const void *attnOutWeight, const float *ln2Gamma, const float *ln2Beta,
++        const void *gateWeight, const void *upWeight, const void *downWeight, const float *queryBias,
++        const float *keyBias, const float *valueBias, const float *attnOutBias, const void *myqkvWeight ,
++        const float *gateBias , const float *upBias , const float *downBias, const float *myqkvBias) {
++
++    if (kvcdt == DataType::fp16) {
++        if (rt == RopeType::LLAMA_ROPE)
++            return LayerLLaMAWrapper<float16_t, LlamaRotaryEmbedding>(dt, at, nt, batchSize, inputSeqLen, attHeadDim,
++                attHeadNum, kvHeadNum, maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step,
++                hiddenSize, intermediateSize, output, outputStride, input, inputStride,
++                ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight, attnOutWeight, ln2Gamma, ln2Beta,
++                gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias, attnOutBias,myqkvWeight) ;
++        else {
++            printf(">> unsupported Rope type: %d\n", rt);
++        }
++    } else if (kvcdt == DataType::int8) {
++        if (rt == RopeType::LLAMA_ROPE)
++            return LayerLLaMAWrapper<int8_t, LlamaRotaryEmbedding>(dt, at, nt, batchSize, inputSeqLen, attHeadDim,
++                attHeadNum, kvHeadNum, maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step,
++                hiddenSize, intermediateSize, output, outputStride, input, inputStride,
++                ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight, attnOutWeight, ln2Gamma, ln2Beta,
++                gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias, attnOutBias,myqkvWeight) ;
++        else {
++            printf(">> unsupported Rope type: %d\n", rt);
++        }
++    } else {
++        printf(">> unsupported KVcache data type: %d\n", kvcdt);
++        return;
++    }
++
++}
++
+ } // namespace xft
+diff --git a/src/layers/decoder_layer.h b/src/layers/decoder_layer.h
+index 3cb5873..570b267 100644
+--- a/src/layers/decoder_layer.h
++++ b/src/layers/decoder_layer.h
+@@ -83,10 +83,10 @@ public:
+             const float *fc1Scales, const float *fc1Zeros, const float *fc1Bias, const OriWeiT *fc2Weight,
+             const float *fc2Scales, const float *fc2Zeros, const float *fc2Bias, const float *ln2Gamma,
+             const float *ln2Beta, const OriWeiT *fc3Weight, const float *fc3Scales, const float *fc3Zeros,
+-            bool trans = true) {
++            bool trans = true,const OriWeiT *myqkvWeight = nullptr) {
+         attn.setWeights(ctx, queryWeight, queryScale, queryZero, queryBias, keyWeight, keyScale, keyZero, keyBias,
+                 valueWeight, valueScale, valueZero, valueBias, attnOutWeight, attnOutScale, attnOutZero, attnOutBias,
+-                true, ln1Gamma, ln1Beta, trans);
++                true, ln1Gamma, ln1Beta, trans,myqkvWeight);
+ 
+         mlp.setWeights(ctx, fc1Weight, fc1Scales, fc1Zeros, fc1Bias, fc2Weight, fc2Scales, fc2Zeros, fc2Bias, ln2Gamma,
+                 ln2Beta, fc3Weight, fc3Scales, fc3Zeros, trans);
+diff --git a/tests/ut/layers_decoder_test.cpp b/tests/ut/layers_decoder_test.cpp
+index be75d94..0e56b10 100644
+--- a/tests/ut/layers_decoder_test.cpp
++++ b/tests/ut/layers_decoder_test.cpp
+@@ -21,8 +21,8 @@
+ #include "layers_decoder.h"
+ #include "gtest/gtest.h"
+ 
+-template <typename T>
+-static void compareLayerLLaMA(int step, int batchSize, int inputSeqLen, int pastSeqLen, int currentSeqLen,
++static void compareLayerLLaMA(xft::DataType dt, xft::DataType kvcdt, int step,
++        int batchSize, int inputSeqLen, int pastSeqLen, int currentSeqLen,
+         int attHeadDim, int attHeadNum, int kvHeadNum, int maxPositions, int maxPosEmbed, int hiddenSize,
+         int intermediateSize, const float *ln1Gamma, const float *ln1Beta, const void *queryWeight,
+         const void *keyWeight, const void *valueWeight, const void *attnOutWeight, const float *ln2Gamma,
+@@ -36,19 +36,8 @@ static void compareLayerLLaMA(int step, int batchSize, int inputSeqLen, int past
+         input[i] = static_cast<float>(1.0f * rand() / RAND_MAX);
+     }
+ 
+-    xft::DataType dt = xft::DataType::unknown;
+-    if constexpr (std::is_same<T, bfloat16_t>::value) {
+-        dt = xft::DataType::bf16;
+-    } else if constexpr (std::is_same<T, float16_t>::value) {
+-        dt = xft::DataType::fp16;
+-    } else {
+-        printf("Unsupported data type\n");
+-        GTEST_FAIL();
+-        return;
+-    }
+-
+     auto start = std::chrono::high_resolution_clock::now();
+-    invokeLayerLLaMA(dt, xft::ActivationType::SILU, xft::NormType::RMS, batchSize, inputSeqLen, attHeadDim, attHeadNum,
++    invokeLayerLLaMA(dt, kvcdt, xft::RopeType::LLAMA_ROPE, xft::ActivationType::SILU, xft::NormType::RMS, batchSize, inputSeqLen, attHeadDim, attHeadNum,
+             kvHeadNum, maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize,
+             (void *)ourOutput, hiddenSize, input, hiddenSize, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
+             attnOutWeight, ln2Gamma, ln2Beta, gateW, upW, downW);
+@@ -60,8 +49,7 @@ static void compareLayerLLaMA(int step, int batchSize, int inputSeqLen, int past
+     free(ourOutput);
+ }
+ 
+-template <typename T>
+-void test_LayerLLaMA(void) {
++void test_LayerLLaMA(xft::DataType dt, xft::DataType kvcdt) {
+     int maxPosEmbed = 4096;
+     int maxPositions = maxPosEmbed;
+     int hiddenSize = 4096;
+@@ -111,16 +99,16 @@ void test_LayerLLaMA(void) {
+     int currentSeqLen = inputSeqLen;
+     int nextTokenNum = 1;
+ 
+-    compareLayerLLaMA<T>(step++, batchSize, inputSeqLen, pastSeqLen, currentSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++    compareLayerLLaMA(dt, kvcdt, step++, batchSize, inputSeqLen, pastSeqLen, currentSeqLen, attHeadDim, attHeadNum, kvHeadNum,
+             maxPositions, maxPosEmbed, hiddenSize, intermediateSize, ln1Gamma, ln1Beta, qkvProj, qkvProj + qSize,
+             qkvProj + kvSize, oProj, ln2Gamma, ln2Beta, gateW, upW, downW);
+     pastSeqLen += inputSeqLen;
+     currentSeqLen = nextTokenNum;
+-    compareLayerLLaMA<T>(step++, batchSize, inputSeqLen, pastSeqLen, currentSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++    compareLayerLLaMA(dt, kvcdt, step++, batchSize, inputSeqLen, pastSeqLen, currentSeqLen, attHeadDim, attHeadNum, kvHeadNum,
+             maxPositions, maxPosEmbed, hiddenSize, intermediateSize, ln1Gamma, ln1Beta, qkvProj, qkvProj + qSize,
+             qkvProj + kvSize, oProj, ln2Gamma, ln2Beta, gateW, upW, downW);
+     pastSeqLen += nextTokenNum;
+-    compareLayerLLaMA<T>(step++, batchSize, inputSeqLen, pastSeqLen, currentSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++    compareLayerLLaMA(dt, kvcdt, step++, batchSize, inputSeqLen, pastSeqLen, currentSeqLen, attHeadDim, attHeadNum, kvHeadNum,
+             maxPositions, maxPosEmbed, hiddenSize, intermediateSize, ln1Gamma, ln1Beta, qkvProj, qkvProj + qSize,
+             qkvProj + kvSize, oProj, ln2Gamma, ln2Beta, gateW, upW, downW);
+ 
+@@ -135,15 +123,31 @@ void test_LayerLLaMA(void) {
+     free(downW);
+ }
+ 
+-TEST(LayerLLaMA, bfloat16_t) {
+-    test_LayerLLaMA<bfloat16_t>();
++TEST(LayerLLaMA, w_bf16_kv_fp16) {
++    test_LayerLLaMA(xft::DataType::bf16, xft::DataType::fp16);
++}
++
++TEST(LayerLLaMA, w_bf16_kv_int8) {
++    test_LayerLLaMA(xft::DataType::bf16, xft::DataType::int8);
++}
++
++TEST(LayerLLaMA, w_fp16_kv_fp16) {
++    test_LayerLLaMA(xft::DataType::fp16, xft::DataType::fp16);
+ }
+ 
+-TEST(LayerLLaMA, float16_t) {
+-    test_LayerLLaMA<float16_t>();
++TEST(LayerLLaMA, w_fp16_kv_int8) {
++    test_LayerLLaMA(xft::DataType::fp16, xft::DataType::int8);
++}
++
++TEST(LayerLLaMA, w_bf16_int8_kv_fp16) {
++    test_LayerLLaMA(xft::DataType::bf16_int8, xft::DataType::fp16);
++}
++
++TEST(LayerLLaMA, w_bf16_int8_kv_int8) {
++    test_LayerLLaMA(xft::DataType::bf16_int8, xft::DataType::int8);
+ }
+ 
+ int main(int argc, char **argv) {
+     ::testing::InitGoogleTest(&argc, argv);
+     return RUN_ALL_TESTS();
+-}
+\ No newline at end of file
++}
+-- 
+2.25.1
+

csrc/cpu/0001-fp32.patch

@@ -0,0 +1,647 @@
+From 8bb1110823ba646f56aaa459394ffb0b3a940352 Mon Sep 17 00:00:00 2001
+From: bukejiyu <395822456@qq.com>
+Date: Wed, 21 Aug 2024 22:03:13 +0800
+Subject: [PATCH] fp32
+
+---
+ CMakeLists.txt                   |  16 +-
+ cmake/mkl.cmake                  |   2 +-
+ cmake/mklml.cmake                |   2 +-
+ cmake/onednn.cmake               |   1 +
+ cmake/xdnn.cmake                 |   2 +-
+ include/dtype.h                  |   5 +
+ include/layers_decoder.h         |  21 +--
+ src/layers/attention.h           |   9 +-
+ src/layers/decoder_layer.cpp     | 256 ++++++++++++++++++++++++-------
+ src/layers/decoder_layer.h       |   4 +-
+ tests/ut/layers_decoder_test.cpp |  52 ++++---
+ 11 files changed, 263 insertions(+), 107 deletions(-)
+
+diff --git a/CMakeLists.txt b/CMakeLists.txt
+index aa70d4b..9149154 100644
+--- a/CMakeLists.txt
++++ b/CMakeLists.txt
+@@ -178,14 +178,14 @@ if(WITH_GPU)
+     add_definitions(-DAVX512_FP32_WEIGHT_ONLY_NF4=true)
+ else()
+     # Enable AVX512_FP16 optimization
+-    # add_definitions(-DAVX512_FP32_WEIGHT_ONLY_FP16=true)
+-    add_definitions(-DAVX512_FP16_WEIGHT_ONLY_FP16=true)
+-    add_definitions(-DAVX512_BF16_WEIGHT_ONLY_BF16=true)
+-    # add_definitions(-DAVX512_FP32_WEIGHT_ONLY_INT8=true)
+-    add_definitions(-DAVX512_FP16_WEIGHT_ONLY_INT8=true)
+-    # add_definitions(-DAVX512_FP32_WEIGHT_ONLY_INT4=true)
+-    add_definitions(-DAVX512_FP16_WEIGHT_ONLY_INT4=true)
+-    add_definitions(-DAVX512_FP32_WEIGHT_ONLY_NF4=true)
++    add_definitions(-DAVX512_FP32_WEIGHT_ONLY_FP16=true)
++    # add_definitions(-DAVX512_FP16_WEIGHT_ONLY_FP16=true)
++    # add_definitions(-DAVX512_BF16_WEIGHT_ONLY_BF16=true)
++    add_definitions(-DAVX512_FP32_WEIGHT_ONLY_INT8=true)
++    # add_definitions(-DAVX512_FP16_WEIGHT_ONLY_INT8=true)
++    add_definitions(-DAVX512_FP32_WEIGHT_ONLY_INT4=true)
++    # add_definitions(-DAVX512_FP16_WEIGHT_ONLY_INT4=true)
++    # add_definitions(-DAVX512_FP32_WEIGHT_ONLY_NF4=true)
+     # add_definitions(-DAVX512_FP16_WEIGHT_ONLY_NF4=true)
+     # Enable AMX_FP16 optimization 
+     # add_definitions(-DAMX_FP16_WEIGHT_ONLY_FP16=true)
+diff --git a/cmake/mkl.cmake b/cmake/mkl.cmake
+index 0ef2e66..92c6d06 100644
+--- a/cmake/mkl.cmake
++++ b/cmake/mkl.cmake
+@@ -25,7 +25,7 @@ set(MKL_3rdparty_DIR "${CMAKE_SOURCE_DIR}/3rdparty/mkl")
+ if(NOT EXISTS ${MKL_3rdparty_DIR})
+     find_package(Python COMPONENTS Interpreter Development)
+     execute_process(COMMAND ${Python_EXECUTABLE} -m pip install --force-reinstall
+-                            --prefix=${MKL_3rdparty_DIR} mkl-static==2024.0.0 mkl-include==2024.0.0
++                            --prefix=${MKL_3rdparty_DIR} mkl-static==2024.0.0 mkl-include==2024.0.0 -i https://mirrors.aliyun.com/pypi/simple
+                     RESULT_VARIABLE EXIT_CODE)
+ 
+     if(NOT ${EXIT_CODE} EQUAL 0)
+diff --git a/cmake/mklml.cmake b/cmake/mklml.cmake
+index 4baec46..d89ce46 100644
+--- a/cmake/mklml.cmake
++++ b/cmake/mklml.cmake
+@@ -28,7 +28,7 @@ include(ExternalProject)
+ ExternalProject_Add(mklml
+   URL               https://github.com/oneapi-src/oneDNN/releases/download/v0.21/mklml_lnx_2019.0.5.20190502.tgz
+   URL_HASH          MD5=dfcea335652dbf3518e1d02cab2cea97
+-  TIMEOUT           60
++  TIMEOUT           360
+   SOURCE_DIR        ${CMAKE_SOURCE_DIR}/3rdparty/mklml
+   CONFIGURE_COMMAND ""
+   BUILD_COMMAND     ""
+diff --git a/cmake/onednn.cmake b/cmake/onednn.cmake
+index 8efabd1..bfa558a 100644
+--- a/cmake/onednn.cmake
++++ b/cmake/onednn.cmake
+@@ -36,6 +36,7 @@ if(NOT EXISTS ${ONEDNN_3rdparty_DIR})
+     ExternalProject_Add(onednn
+       GIT_REPOSITORY    https://github.com/oneapi-src/oneDNN.git
+       GIT_TAG           v3.5
++      TIMEOUT           360
+       SOURCE_DIR        ${ONEDNN_3rdparty_DIR}
+       BINARY_DIR        ${ONEDNN_3rdparty_DIR}
+       CONFIGURE_COMMAND ${CMAKE_COMMAND} -E make_directory "build" && ${CMAKE_COMMAND} -E chdir "build" ${CMAKE_COMMAND} ${ONEDNN_BUILD_OPTIONS} ..
+diff --git a/cmake/xdnn.cmake b/cmake/xdnn.cmake
+index 7c0e051..721bfe3 100644
+--- a/cmake/xdnn.cmake
++++ b/cmake/xdnn.cmake
+@@ -28,7 +28,7 @@ include(ExternalProject)
+ ExternalProject_Add(xdnn_lib
+   URL               https://github.com/intel/xFasterTransformer/releases/download/IntrinsicGemm/xdnn_v1.5.2.tar.gz
+   URL_HASH          MD5=884f2e1e2c846ff19f33c889681f8dc2
+-  TIMEOUT           120
++  TIMEOUT           360
+   SOURCE_DIR        ${CMAKE_SOURCE_DIR}/3rdparty/xdnn
+   CONFIGURE_COMMAND ""
+   BUILD_COMMAND     ""
+diff --git a/include/dtype.h b/include/dtype.h
+index de72bce..9e0a448 100644
+--- a/include/dtype.h
++++ b/include/dtype.h
+@@ -31,6 +31,7 @@ enum DataType {
+     w8a8_int8,
+     w8a8_int4,
+     w8a8_nf4,
++    fp16_int8,
+     unknown,
+ };
+ 
+@@ -51,4 +52,8 @@ enum ActivationType {
+     SILU,
+ };
+ 
++enum RopeType {
++    LLAMA_ROPE = 0,
++};
++
+ } // namespace xft
+diff --git a/include/layers_decoder.h b/include/layers_decoder.h
+index 34f6aa5..a30e34d 100644
+--- a/include/layers_decoder.h
++++ b/include/layers_decoder.h
+@@ -17,13 +17,14 @@
+ #include "dtype.h"
+ 
+ namespace xft {
+-
+-void invokeLayerLLaMA(DataType dt, ActivationType at, NormType nt, int batchSize, int inputSeqLen, int attHeadDim,
+-        int attHeadNum, int kvHeadNum, int maxPositions, int maxPosEmbed, int pastSeqLen, int currentSeqLen, int step,
+-        int hiddenSize, int intermediateSize, void *output, int outputStride, const void *input, int inputStride,
+-        const float *ln1Gamma, const float *ln1Beta, const void *queryWeight, const void *keyWeight,
+-        const void *valueWeight, const void *attnOutWeight, const float *ln2Gamma, const float *ln2Beta,
+-        const void *gateWeight, const void *upWeight, const void *downWeight, const float *queryBias = nullptr,
+-        const float *keyBias = nullptr, const float *valueBias = nullptr, const float *attnOutBias = nullptr);
+-
+-} // namespace xft
+\ No newline at end of file
++void invokeLayerLLaMA(DataType dt, DataType kvcdt, RopeType rt, ActivationType at, NormType nt, int batchSize,
++        int inputSeqLen, int attHeadDim, int attHeadNum, int kvHeadNum, int maxPositions, int maxPosEmbed,
++        int pastSeqLen, int currentSeqLen, int step, int hiddenSize, int intermediateSize, void *output,
++        int outputStride, const void *input, int inputStride, const float *ln1Gamma, const float *ln1Beta,
++        const void *queryWeight, const void *keyWeight, const void *valueWeight, const void *attnOutWeight,
++        const float *ln2Gamma, const float *ln2Beta, const void *gateWeight, const void *upWeight,
++        const void *downWeight, const float *queryBias = nullptr, const float *keyBias = nullptr,
++        const float *valueBias = nullptr, const float *attnOutBias = nullptr, const void *myqkvWeight = nullptr,
++        const float *gateBias = nullptr, const float *upBias = nullptr, const float *downBias = nullptr,
++        const float *myqkvBias = nullptr);
++} // namespace xft
+diff --git a/src/layers/attention.h b/src/layers/attention.h
+index 092b3d6..b438837 100644
+--- a/src/layers/attention.h
++++ b/src/layers/attention.h
+@@ -84,7 +84,8 @@ public:
+             const float *queryBias, const OriWeiT *keyWeight, const float *keyScale, const float *keyZero,
+             const float *keyBias, const OriWeiT *valueWeight, const float *valueScale, const float *valueZero,
+             const float *valueBias, const OriWeiT *attnOutWeight, const float *attnOutScale, const float *attnOutZero,
+-            const float *attnOutBias, bool doLNorm, const float *gamma1, const float *beta1, bool trans = true) {
++            const float *attnOutBias, bool doLNorm, const float *gamma1, const float *beta1, bool trans = true,
++            const OriWeiT *myqkvWeight = nullptr) {
+         int hiddenSize = ctx->hiddenSize;
+         int headSize = ctx->attHeadSize;
+ 
+@@ -107,7 +108,10 @@ public:
+                     valueWeight + this->startKVHead * headSize * hiddenSize / sizeFactor,
+                     hiddenSize * kvResponsibleCols * sizeof(OriWeiT) / sizeFactor);
+         } else {
+-            int qkvStride = (ctx->attHeadNum + ctx->kvHeadNum + ctx->kvHeadNum) * ctx->attHeadSize;
++            if (myqkvWeight != nullptr) {
++                memcpy(concatBuf, myqkvWeight, hiddenSize * responsibleCols * sizeof(OriWeiT) / sizeFactor);
++            } else {
++                int qkvStride = (ctx->attHeadNum + ctx->kvHeadNum + ctx->kvHeadNum) * ctx->attHeadSize;
+ #pragma omp parallel for
+             for (int i = 0; i < hiddenSize; ++i) {
+                 memcpy(concatBuf + i * responsibleCols / sizeFactor,
+@@ -120,6 +124,7 @@ public:
+                                 + kvResponsibleCols / sizeFactor,
+                         valueWeight + i * qkvStride / sizeFactor + this->startKVHead * headSize / sizeFactor,
+                         kvResponsibleCols * sizeof(OriWeiT) / sizeFactor);
++                }
+             }
+         }
+         float *concatScale = nullptr;
+diff --git a/src/layers/decoder_layer.cpp b/src/layers/decoder_layer.cpp
+index 02f13cb..0f30f21 100644
+--- a/src/layers/decoder_layer.cpp
++++ b/src/layers/decoder_layer.cpp
+@@ -21,19 +21,21 @@
+ #include "layers_mlp.h"
+ #include "mlp_llama.h"
+ #include "rms_norm.h"
++#include "numa_allocator.h"
+ 
+ #include <unordered_map>
+ 
+ namespace xft {
+ 
+-template <typename DataT, typename NormT>
++template <typename DataT, typename KVCacheT, typename RopeT, typename NormT>
+ void LayerLLaMAImpl(DataType dt, ActivationType at, NormType nt, int batchSize, int inputSeqLen, int attHeadDim,
+         int attHeadNum, int kvHeadNum, int maxPositions, int maxPosEmbed, int pastSeqLen, int currentSeqLen, int step,
+         int hiddenSize, int intermediateSize, void *output, int outputStride, const void *input, int inputStride,
+         const float *ln1Gamma, const float *ln1Beta, const void *queryWeight, const void *keyWeight,
+         const void *valueWeight, const void *attnOutWeight, const float *ln2Gamma, const float *ln2Beta,
+         const void *gateWeight, const void *upWeight, const void *downWeight, const float *queryBias,
+-        const float *keyBias, const float *valueBias, const float *attnOutBias) {
++        const float *keyBias, const float *valueBias, const float *attnOutBias,
++        MMHelper *mmHelper, DecoderContext *ctx, KVCacheManager<KVCacheT> *kvCacheMgr,const void *myqkvWeight = nullptr) {
+ 
+     // TODO: will deprecate attention mask in future, so need to change this
+     auto prepareAttnMask = [&](DecoderContext *ctx, int step) {
+@@ -83,67 +85,46 @@ void LayerLLaMAImpl(DataType dt, ActivationType at, NormType nt, int batchSize,
+         return mask;
+     };
+ 
+-    using DECODER = Decoder<Attention<DataT, LlamaRotaryEmbedding, NormT>, LlamaMLP<DataT>>;
+-    static std::unordered_map<std::string, DECODER *> llama_layer_hub;
+-    static MMHelper *mmHelper;
+-    static DecoderContext *ctx;
+-    static KVCacheManager<float16_t> *kvCacheMgr;
+-
+-    std::string actType;
+-    if (at == ActivationType::SILU)
+-        actType = "silu";
+-    else if (at == ActivationType::RELU)
+-        actType = "relu";
+-    else if (at == ActivationType::GELU)
+-        actType = "gelu";
+-    else if (at == ActivationType::SWIGLU)
+-        actType = "swiglu";
+-    else
+-        printf(">> unsupported activation type\n");
+-
+-    if (ctx == nullptr
+-            || (ctx != nullptr && (ctx->hiddenSize != hiddenSize || ctx->intermediateSize != intermediateSize))) {
+-        if (ctx != nullptr) delete ctx;
+-        printf(">> create context: %d %d\n", hiddenSize, intermediateSize);
+-        mmHelper = new MMHelper(Env::getInstance().getEngineKind(), Env::getInstance().getEngineIndex());
+-        ctx = new DecoderContext(1, hiddenSize, attHeadDim, attHeadNum, kvHeadNum, intermediateSize, actType, 1e-6, 0,
+-                0, maxPositions, maxPosEmbed, -1, 0, 1, mmHelper);
+-        if (kvCacheMgr != nullptr) delete kvCacheMgr;
+-        kvCacheMgr = new KVCacheManager<float16_t>(1);
+-    }
+-
++    using DECODER = Decoder<Attention<DataT, RopeT, NormT>, LlamaMLP<DataT>>;
++    DECODER *llama_layer;
++    static xft::Matrix<float> actBuffers ;
++    //static std::unordered_map<std::string, DECODER *> llama_layer_hub;
++    static std::unordered_map<std::string, std::tuple<DECODER*>> llama_layer_hub;
+     // create hash key and value: if hidden and intermediateSize is changed , then memory pointer is also changed.
+     std::stringstream weights_addr;
+     weights_addr << queryWeight << "_" << keyWeight << "_" << valueWeight << "_" << attnOutWeight << "_" << gateWeight
+                  << "_" << upWeight << "_" << downWeight << "_" << dt << "_" << at << "_" << nt << "_" << attHeadDim
+                  << "_" << attHeadNum << "_" << kvHeadNum;
+     std::string llama_layer_key = weights_addr.str();
+-    DECODER *llama_layer;
+ 
+     auto it_created = llama_layer_hub.find(llama_layer_key);
+     if (it_created == llama_layer_hub.end()) {
++        int firstNode = getenv("FIRST_TOKEN_WEIGHT_LOCATION") ? atoi(getenv("FIRST_TOKEN_WEIGHT_LOCATION")) : -1;
++        int nextNode = getenv("NEXT_TOKEN_WEIGHT_LOCATION") ? atoi(getenv("NEXT_TOKEN_WEIGHT_LOCATION")) : -1;
++        if (step == 0)
++                xft_set_preferred_node(firstNode);
++        else
++                xft_set_preferred_node(nextNode);
+         llama_layer = new DECODER(ctx, 0);
+         llama_layer->setWeights(ctx, (const float *)queryWeight, nullptr, nullptr, queryBias, (const float *)keyWeight,
+                 nullptr, nullptr, keyBias, (const float *)valueWeight, nullptr, nullptr, valueBias,
+                 (const float *)attnOutWeight, nullptr, nullptr, attnOutBias, ln1Gamma, ln1Beta,
+                 (const float *)gateWeight, nullptr, nullptr, nullptr, (const float *)upWeight, nullptr, nullptr,
+-                nullptr, ln2Gamma, ln2Beta, (const float *)downWeight, nullptr, nullptr, false);
+-        llama_layer_hub[llama_layer_key] = llama_layer;
+-        printf(">> create llama_layer_key: %s\n", llama_layer_key.c_str());
++                nullptr, ln2Gamma, ln2Beta, (const float *)downWeight, nullptr, nullptr, false,(const float *)myqkvWeight);
++
++
++        llama_layer_hub[llama_layer_key] = std::make_tuple(llama_layer);;
++        // printf(">> create llama_layer_key: %s\n", llama_layer_key.c_str());
++        xft_set_preferred_node(-1);
+     } else {
+-        llama_layer = it_created->second;
++        llama_layer = std::get<0>(it_created->second);
+     }
+-
+-    ctx->resize(batchSize, inputSeqLen, pastSeqLen);
+-    xft::Matrix<float> actBuffers;
+     actBuffers.Resize(batchSize * inputSeqLen * 2, hiddenSize);
++    ctx->resize(batchSize, inputSeqLen, pastSeqLen);
+     float *attnMask = prepareAttnMask(ctx, step);
+ 
+-    int workers = 1;
+-    int headsPerSplit = (ctx->kvHeadNum + workers - 1) / workers;
+-    kvCacheMgr->resize(maxPositions, batchSize, headsPerSplit, attHeadDim);
+-    KVCacheTensor<float16_t> &presentKey = kvCacheMgr->getKey(0);
+-    KVCacheTensor<float16_t> &presentValue = kvCacheMgr->getValue(0);
++    KVCacheTensor<KVCacheT> &presentKey = kvCacheMgr->getKey(0);
++    KVCacheTensor<KVCacheT> &presentValue = kvCacheMgr->getValue(0);
+ 
+     float *attnOut = (float *)(ctx->tmpBuf.Data());
+ 
+@@ -159,45 +140,168 @@ void LayerLLaMAImpl(DataType dt, ActivationType at, NormType nt, int batchSize,
+     llama_layer->forwardFFN(ctx, attnOut, (float *)output, inputStride, outputStride, true);
+ }
+ 
+-void invokeLayerLLaMA(DataType dt, ActivationType at, NormType nt, int batchSize, int inputSeqLen, int attHeadDim,
++template <typename KVCacheT, typename RopeT>
++void LayerLLaMAWrapper(DataType dt, ActivationType at, NormType nt, int batchSize, int inputSeqLen, int attHeadDim,
+         int attHeadNum, int kvHeadNum, int maxPositions, int maxPosEmbed, int pastSeqLen, int currentSeqLen, int step,
+         int hiddenSize, int intermediateSize, void *output, int outputStride, const void *input, int inputStride,
+         const float *ln1Gamma, const float *ln1Beta, const void *queryWeight, const void *keyWeight,
+         const void *valueWeight, const void *attnOutWeight, const float *ln2Gamma, const float *ln2Beta,
+         const void *gateWeight, const void *upWeight, const void *downWeight, const float *queryBias,
+-        const float *keyBias, const float *valueBias, const float *attnOutBias) {
++        const float *keyBias, const float *valueBias, const float *attnOutBias,const void *myqkvWeight=nullptr) {
+     static std::mutex mutex;
+     std::lock_guard<std::mutex> lock(mutex);
+ 
++    std::string actType;
++    if (at == ActivationType::SILU)
++        actType = "silu";
++    else if (at == ActivationType::RELU)
++        actType = "relu";
++    else if (at == ActivationType::GELU)
++        actType = "gelu";
++    else if (at == ActivationType::SWIGLU)
++        actType = "swiglu";
++    else {
++        printf(">> unsupported activation type\n");
++        return;
++    }
++
++    static MMHelper *mmHelper;
++    static DecoderContext *ctx;
++    if (ctx == nullptr
++            || (ctx != nullptr && (ctx->hiddenSize != hiddenSize || ctx->intermediateSize != intermediateSize))) {
++        if (ctx != nullptr) delete ctx;
++        // printf(">> create context: %d %d\n", hiddenSize, intermediateSize);
++        mmHelper = new MMHelper(Env::getInstance().getEngineKind(), Env::getInstance().getEngineIndex());
++        ctx = new DecoderContext(1, hiddenSize, attHeadDim, attHeadNum, kvHeadNum, intermediateSize, actType, 1e-6, 0,
++                0, maxPositions, maxPosEmbed, -1, 0, 1, mmHelper);
++    }
++
++    KVCacheManager<KVCacheT> *kvCacheMgr;
++    static std::unordered_map<std::string, KVCacheManager<KVCacheT> *> kv_hub;
++
++    // create hash key and value: if hidden and intermediateSize is changed , then memory pointer is also changed.
++    std::stringstream layer_key;
++    layer_key << queryWeight << "_" << keyWeight << "_" << valueWeight << "_" << attnOutWeight << "_" << gateWeight
++                 << "_" << upWeight << "_" << downWeight << "_" << dt << "_" << at << "_" << nt << "_" << attHeadDim
++                 << "_" << attHeadNum << "_" << kvHeadNum;
++    std::string kv_hub_key = layer_key.str();
++
++    auto it_created = kv_hub.find(kv_hub_key);
++    if (it_created == kv_hub.end()) {
++        int kvcNode = getenv("KVCACHE_LOCATION") ? atoi(getenv("KVCACHE_LOCATION")) : -1;
++        xft_set_preferred_node(kvcNode);
++        kvCacheMgr = new KVCacheManager<KVCacheT>(1);
++        int workers = 1;
++        int headsPerSplit = (ctx->kvHeadNum + workers - 1) / workers;
++        kvCacheMgr->resize(maxPositions, batchSize, headsPerSplit, attHeadDim);
++        kv_hub[kv_hub_key] = kvCacheMgr;
++        // printf(">> create kv_hub_key: %s\n", kv_hub_key.c_str());
++        xft_set_preferred_node(-1);
++    } else {
++        kvCacheMgr = it_created->second;
++    }
++
+     if (dt == DataType::bf16) {
+-        if (nt == NormType::RMS)
+-            LayerLLaMAImpl<bfloat16_t, RmsNorm>(dt, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++        if (nt == NormType::RMS) {
++            LayerLLaMAImpl<bfloat16_t, KVCacheT, RopeT, RmsNorm>(dt, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
+                     maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
+                     outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
+                     attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
+-                    attnOutBias);
+-        else if (nt == NormType::LN) {
+-            LayerLLaMAImpl<bfloat16_t, LayerNorm>(dt, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                    attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++        } else if (nt == NormType::LN) {
++            LayerLLaMAImpl<bfloat16_t, KVCacheT, RopeT, LayerNorm>(dt, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
+                     maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
+                     outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
+                     attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
+-                    attnOutBias);
++                    attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
+         } else {
+             printf(">> unsupported norm type\n");
+         }
+     } else if (dt == DataType::fp16) {
+-        if (nt == NormType::RMS)
+-            LayerLLaMAImpl<float16_t, RmsNorm>(dt, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++        if (nt == NormType::RMS) {
++            LayerLLaMAImpl<float16_t, KVCacheT, RopeT, RmsNorm>(dt, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
+                     maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
+                     outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
+                     attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
+-                    attnOutBias);
+-        else if (nt == NormType::LN) {
+-            LayerLLaMAImpl<float16_t, LayerNorm>(dt, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                    attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++        } else if (nt == NormType::LN) {
++            LayerLLaMAImpl<float16_t, KVCacheT, RopeT, LayerNorm>(dt, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
+                     maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
+                     outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
+                     attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
+-                    attnOutBias);
++                    attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++        } else {
++            printf(">> unsupported norm type\n");
++        }
++    } else if (dt == DataType::bf16_int8) {
++        if (nt == NormType::RMS) {
++            auto firstTokenFunc = LayerLLaMAImpl<bfloat16_t, KVCacheT, RopeT, RmsNorm>;
++            auto nextTokenFunc = LayerLLaMAImpl<int8_t, KVCacheT, RopeT, RmsNorm>;
++            if (step == 0) {
++                    firstTokenFunc(DataType::bf16, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                        maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
++                        outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
++                        attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
++                        attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++
++            } else {
++                    nextTokenFunc(DataType::int8, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                        maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
++                        outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
++                        attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
++                        attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++            }
++        } else if (nt == NormType::LN) {
++            auto firstTokenFunc = LayerLLaMAImpl<bfloat16_t, KVCacheT, RopeT, LayerNorm>;
++            auto nextTokenFunc = LayerLLaMAImpl<int8_t, KVCacheT, RopeT, LayerNorm>;
++            if (step == 0)
++                    firstTokenFunc(DataType::bf16, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                        maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
++                        outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
++                        attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
++                        attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++            else
++                    nextTokenFunc(DataType::int8, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                        maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
++                        outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
++                        attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
++                        attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++        } else {
++            printf(">> unsupported norm type\n");
++        }
++    } else if (dt == DataType::fp16_int8) {
++        if (nt == NormType::RMS) {
++            auto firstTokenFunc = LayerLLaMAImpl<float16_t, KVCacheT, RopeT, RmsNorm>;
++            auto nextTokenFunc = LayerLLaMAImpl<int8_t, KVCacheT, RopeT, RmsNorm>;
++            if (step == 0) {
++                    firstTokenFunc(DataType::fp16, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                        maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
++                        outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
++                        attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
++                        attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++
++            } else {
++                    nextTokenFunc(DataType::int8, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                        maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
++                        outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
++                        attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
++                        attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++            }
++        } else if (nt == NormType::LN) {
++            auto firstTokenFunc = LayerLLaMAImpl<bfloat16_t, KVCacheT, RopeT, LayerNorm>;
++            auto nextTokenFunc = LayerLLaMAImpl<int8_t, KVCacheT, RopeT, LayerNorm>;
++            if (step == 0)
++                    firstTokenFunc(DataType::fp16, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                        maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
++                        outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
++                        attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
++                        attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
++            else
++                    nextTokenFunc(DataType::int8, at, nt, batchSize, inputSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++                        maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize, output,
++                        outputStride, input, inputStride, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
++                        attnOutWeight, ln2Gamma, ln2Beta, gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias,
++                        attnOutBias, mmHelper, ctx, kvCacheMgr,myqkvWeight);
+         } else {
+             printf(">> unsupported norm type\n");
+         }
+@@ -206,4 +310,40 @@ void invokeLayerLLaMA(DataType dt, ActivationType at, NormType nt, int batchSize
+     }
+ }
+ 
++void invokeLayerLLaMA(DataType dt, DataType kvcdt, RopeType rt, ActivationType at, NormType nt, int batchSize, int inputSeqLen, int attHeadDim,
++        int attHeadNum, int kvHeadNum, int maxPositions, int maxPosEmbed, int pastSeqLen, int currentSeqLen, int step,
++        int hiddenSize, int intermediateSize, void *output, int outputStride, const void *input, int inputStride,
++        const float *ln1Gamma, const float *ln1Beta, const void *queryWeight, const void *keyWeight,
++        const void *valueWeight, const void *attnOutWeight, const float *ln2Gamma, const float *ln2Beta,
++        const void *gateWeight, const void *upWeight, const void *downWeight, const float *queryBias,
++        const float *keyBias, const float *valueBias, const float *attnOutBias, const void *myqkvWeight ,
++        const float *gateBias , const float *upBias , const float *downBias, const float *myqkvBias) {
++
++    if (kvcdt == DataType::fp16) {
++        if (rt == RopeType::LLAMA_ROPE)
++            return LayerLLaMAWrapper<float16_t, LlamaRotaryEmbedding>(dt, at, nt, batchSize, inputSeqLen, attHeadDim,
++                attHeadNum, kvHeadNum, maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step,
++                hiddenSize, intermediateSize, output, outputStride, input, inputStride,
++                ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight, attnOutWeight, ln2Gamma, ln2Beta,
++                gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias, attnOutBias,myqkvWeight) ;
++        else {
++            printf(">> unsupported Rope type: %d\n", rt);
++        }
++    } else if (kvcdt == DataType::int8) {
++        if (rt == RopeType::LLAMA_ROPE)
++            return LayerLLaMAWrapper<int8_t, LlamaRotaryEmbedding>(dt, at, nt, batchSize, inputSeqLen, attHeadDim,
++                attHeadNum, kvHeadNum, maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step,
++                hiddenSize, intermediateSize, output, outputStride, input, inputStride,
++                ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight, attnOutWeight, ln2Gamma, ln2Beta,
++                gateWeight, upWeight, downWeight, queryBias, keyBias, valueBias, attnOutBias,myqkvWeight) ;
++        else {
++            printf(">> unsupported Rope type: %d\n", rt);
++        }
++    } else {
++        printf(">> unsupported KVcache data type: %d\n", kvcdt);
++        return;
++    }
++
++}
++
+ } // namespace xft
+diff --git a/src/layers/decoder_layer.h b/src/layers/decoder_layer.h
+index 3cb5873..570b267 100644
+--- a/src/layers/decoder_layer.h
++++ b/src/layers/decoder_layer.h
+@@ -83,10 +83,10 @@ public:
+             const float *fc1Scales, const float *fc1Zeros, const float *fc1Bias, const OriWeiT *fc2Weight,
+             const float *fc2Scales, const float *fc2Zeros, const float *fc2Bias, const float *ln2Gamma,
+             const float *ln2Beta, const OriWeiT *fc3Weight, const float *fc3Scales, const float *fc3Zeros,
+-            bool trans = true) {
++            bool trans = true,const OriWeiT *myqkvWeight = nullptr) {
+         attn.setWeights(ctx, queryWeight, queryScale, queryZero, queryBias, keyWeight, keyScale, keyZero, keyBias,
+                 valueWeight, valueScale, valueZero, valueBias, attnOutWeight, attnOutScale, attnOutZero, attnOutBias,
+-                true, ln1Gamma, ln1Beta, trans);
++                true, ln1Gamma, ln1Beta, trans,myqkvWeight);
+ 
+         mlp.setWeights(ctx, fc1Weight, fc1Scales, fc1Zeros, fc1Bias, fc2Weight, fc2Scales, fc2Zeros, fc2Bias, ln2Gamma,
+                 ln2Beta, fc3Weight, fc3Scales, fc3Zeros, trans);
+diff --git a/tests/ut/layers_decoder_test.cpp b/tests/ut/layers_decoder_test.cpp
+index be75d94..0e56b10 100644
+--- a/tests/ut/layers_decoder_test.cpp
++++ b/tests/ut/layers_decoder_test.cpp
+@@ -21,8 +21,8 @@
+ #include "layers_decoder.h"
+ #include "gtest/gtest.h"
+ 
+-template <typename T>
+-static void compareLayerLLaMA(int step, int batchSize, int inputSeqLen, int pastSeqLen, int currentSeqLen,
++static void compareLayerLLaMA(xft::DataType dt, xft::DataType kvcdt, int step,
++        int batchSize, int inputSeqLen, int pastSeqLen, int currentSeqLen,
+         int attHeadDim, int attHeadNum, int kvHeadNum, int maxPositions, int maxPosEmbed, int hiddenSize,
+         int intermediateSize, const float *ln1Gamma, const float *ln1Beta, const void *queryWeight,
+         const void *keyWeight, const void *valueWeight, const void *attnOutWeight, const float *ln2Gamma,
+@@ -36,19 +36,8 @@ static void compareLayerLLaMA(int step, int batchSize, int inputSeqLen, int past
+         input[i] = static_cast<float>(1.0f * rand() / RAND_MAX);
+     }
+ 
+-    xft::DataType dt = xft::DataType::unknown;
+-    if constexpr (std::is_same<T, bfloat16_t>::value) {
+-        dt = xft::DataType::bf16;
+-    } else if constexpr (std::is_same<T, float16_t>::value) {
+-        dt = xft::DataType::fp16;
+-    } else {
+-        printf("Unsupported data type\n");
+-        GTEST_FAIL();
+-        return;
+-    }
+-
+     auto start = std::chrono::high_resolution_clock::now();
+-    invokeLayerLLaMA(dt, xft::ActivationType::SILU, xft::NormType::RMS, batchSize, inputSeqLen, attHeadDim, attHeadNum,
++    invokeLayerLLaMA(dt, kvcdt, xft::RopeType::LLAMA_ROPE, xft::ActivationType::SILU, xft::NormType::RMS, batchSize, inputSeqLen, attHeadDim, attHeadNum,
+             kvHeadNum, maxPositions, maxPosEmbed, pastSeqLen, currentSeqLen, step, hiddenSize, intermediateSize,
+             (void *)ourOutput, hiddenSize, input, hiddenSize, ln1Gamma, ln1Beta, queryWeight, keyWeight, valueWeight,
+             attnOutWeight, ln2Gamma, ln2Beta, gateW, upW, downW);
+@@ -60,8 +49,7 @@ static void compareLayerLLaMA(int step, int batchSize, int inputSeqLen, int past
+     free(ourOutput);
+ }
+ 
+-template <typename T>
+-void test_LayerLLaMA(void) {
++void test_LayerLLaMA(xft::DataType dt, xft::DataType kvcdt) {
+     int maxPosEmbed = 4096;
+     int maxPositions = maxPosEmbed;
+     int hiddenSize = 4096;
+@@ -111,16 +99,16 @@ void test_LayerLLaMA(void) {
+     int currentSeqLen = inputSeqLen;
+     int nextTokenNum = 1;
+ 
+-    compareLayerLLaMA<T>(step++, batchSize, inputSeqLen, pastSeqLen, currentSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++    compareLayerLLaMA(dt, kvcdt, step++, batchSize, inputSeqLen, pastSeqLen, currentSeqLen, attHeadDim, attHeadNum, kvHeadNum,
+             maxPositions, maxPosEmbed, hiddenSize, intermediateSize, ln1Gamma, ln1Beta, qkvProj, qkvProj + qSize,
+             qkvProj + kvSize, oProj, ln2Gamma, ln2Beta, gateW, upW, downW);
+     pastSeqLen += inputSeqLen;
+     currentSeqLen = nextTokenNum;
+-    compareLayerLLaMA<T>(step++, batchSize, inputSeqLen, pastSeqLen, currentSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++    compareLayerLLaMA(dt, kvcdt, step++, batchSize, inputSeqLen, pastSeqLen, currentSeqLen, attHeadDim, attHeadNum, kvHeadNum,
+             maxPositions, maxPosEmbed, hiddenSize, intermediateSize, ln1Gamma, ln1Beta, qkvProj, qkvProj + qSize,
+             qkvProj + kvSize, oProj, ln2Gamma, ln2Beta, gateW, upW, downW);
+     pastSeqLen += nextTokenNum;
+-    compareLayerLLaMA<T>(step++, batchSize, inputSeqLen, pastSeqLen, currentSeqLen, attHeadDim, attHeadNum, kvHeadNum,
++    compareLayerLLaMA(dt, kvcdt, step++, batchSize, inputSeqLen, pastSeqLen, currentSeqLen, attHeadDim, attHeadNum, kvHeadNum,
+             maxPositions, maxPosEmbed, hiddenSize, intermediateSize, ln1Gamma, ln1Beta, qkvProj, qkvProj + qSize,
+             qkvProj + kvSize, oProj, ln2Gamma, ln2Beta, gateW, upW, downW);
+ 
+@@ -135,15 +123,31 @@ void test_LayerLLaMA(void) {
+     free(downW);
+ }
+ 
+-TEST(LayerLLaMA, bfloat16_t) {
+-    test_LayerLLaMA<bfloat16_t>();
++TEST(LayerLLaMA, w_bf16_kv_fp16) {
++    test_LayerLLaMA(xft::DataType::bf16, xft::DataType::fp16);
++}
++
++TEST(LayerLLaMA, w_bf16_kv_int8) {
++    test_LayerLLaMA(xft::DataType::bf16, xft::DataType::int8);
++}
++
++TEST(LayerLLaMA, w_fp16_kv_fp16) {
++    test_LayerLLaMA(xft::DataType::fp16, xft::DataType::fp16);
+ }
+ 
+-TEST(LayerLLaMA, float16_t) {
+-    test_LayerLLaMA<float16_t>();
++TEST(LayerLLaMA, w_fp16_kv_int8) {
++    test_LayerLLaMA(xft::DataType::fp16, xft::DataType::int8);
++}
++
++TEST(LayerLLaMA, w_bf16_int8_kv_fp16) {
++    test_LayerLLaMA(xft::DataType::bf16_int8, xft::DataType::fp16);
++}
++
++TEST(LayerLLaMA, w_bf16_int8_kv_int8) {
++    test_LayerLLaMA(xft::DataType::bf16_int8, xft::DataType::int8);
+ }
+ 
+ int main(int argc, char **argv) {
+     ::testing::InitGoogleTest(&argc, argv);
+     return RUN_ALL_TESTS();
+-}
+\ No newline at end of file
++}
+-- 
+2.25.1
+

csrc/cpu/README.md

@@ -0,0 +1,17 @@
+# cpu-custom-ops
+
+## 快速开始
+### 1.环境准备
+```shell
+# 查询机器是否支持 avx512指令
+lscpu | grep avx512*
+```
+
+### 2.安装 cpu 自定义算子和第三方库
+```shell
+#建议在 gcc 9.4.0 下安装第三方库
+bash setup.sh
+```
+**Note:**
+
+包含 avx 指令 CPU 机器大模型推理教程 [X86 CPU](../../llm/docs/cpu_install.md)

csrc/cpu/setup.sh

@@ -0,0 +1,68 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+#0.环境准备:安装numactl
+# apt-get update
+# apt-get install numactl
+
+# 1. download XFT
+if [ ! -d xFasterTransformer ]; then
+    git clone https://github.com/intel/xFasterTransformer.git
+fi
+
+#2.cp patch
+cd xFasterTransformer
+git reset --hard 420a493f5c3c74f5fdd786f5399aacd04e021df7
+cd ..
+
+if lscpu | grep -q "avx512_bf16"; then
+    echo "apply bf16 and fp16."
+    if [ ! -f 0001-fp16_bf16.patch ]; then
+        echo "Error:  0001-fp16_bf16.patch not exist."
+        exit 1
+    fi
+    # apply patch
+    cp ./0001-fp16_bf16.patch  ./xFasterTransformer/paddle.patch
+else
+    echo "apply fp32 "
+    if [ ! -f 0001-fp32.patch ]; then
+        echo "Error:  does 0001-fp32.patch not exist."
+        exit 1
+    fi
+    cp ./0001-fp32.patch  ./xFasterTransformer/paddle.patch
+fi
+
+#3. apply patch
+cd xFasterTransformer
+git apply paddle.patch
+
+#4. build xFasterTransformer
+cd 3rdparty
+bash prepare_oneccl.sh
+source ./oneccl/build/_install/env/setvars.sh
+cd ..
+
+rm -rf build
+mkdir build && cd build
+cmake ..
+make -j
+cd ..
+
+#xft
+export XFT_HEADER_DIR=$PWD
+export XFT_LIB_DIR=$XFT_HEADER_DIR/build
+export LD_LIBRARY_PATH=$XFT_LIB_DIR:$LD_LIBRARY_PATH
+#setup cpu paddle_nlp ops
+cd ..
+python ./src/setup_cpu.py install --user

csrc/cpu/src/avx_weight_only.cc

@@ -0,0 +1,160 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "dtype.h"
+#include "matmul_helper.h"
+#include "my_types.h"
+#include "paddle/extension.h"
+template <typename T>
+void AvxCompute(const paddle::Tensor &x,
+                const paddle::Tensor &weight,
+                bool trans,
+                const std::string alog,
+                paddle::Tensor &out,
+                xft::Matrix<T> &quantizedWeight,
+                xft::Vector<float> &WeightScale,
+                xft::Vector<float> &WeightZero,
+                xft::Vector<float> &WeightSum,
+                MMHelper *mmHelper) {
+  auto out_data = out.data<float>();
+  const float *x_data = reinterpret_cast<const float *>(x.data<float>());
+  const float *bias_data = nullptr;
+  int m = 1;
+  for (int i = 0; i < x.shape().size() - 1; i++) {
+    m = m * x.shape()[i];
+  }
+  int k = x.shape()[x.shape().size() - 1];
+  int l = weight.shape()[1];
+  int n = weight.shape()[1];
+  
+  mmHelper->compute(false,
+                    m,
+                    n,
+                    k,
+                    1.0f,
+                    x_data,
+                    k,
+                    quantizedWeight.Data(),
+                    WeightScale.Data(),
+                    WeightZero.Data(),
+                    WeightSum.Data(),
+                    0.0,
+                    out_data,
+                    l);
+};
+template <typename T>
+void AvxWeightOnly(const paddle::Tensor &x,
+                   const paddle::Tensor &weight,
+                   bool trans,
+                   const std::string alog,
+                   paddle::Tensor &out) {
+  static std::unordered_map<std::string,
+                            std::tuple<xft::Matrix<T> *,
+                                       xft::Vector<float> *,
+                                       xft::Vector<float> *,
+                                       xft::Vector<float> *>>
+      weight_only_hub;
+  std::stringstream weights_addr;
+  weights_addr << weight.data<float>() << alog;
+  std::string weight_only_key = weights_addr.str();
+  auto it_created = weight_only_hub.find(weight_only_key);
+  static MMHelper *mmHelper;
+  int rows = weight.shape()[0], cols = weight.shape()[1];
+  xft::Vector<float> *WeightScale =
+      new xft::Vector<float>();  // if weight is int8
+  xft::Vector<float> *WeightZero =
+      new xft::Vector<float>();  // if weight is int8
+  xft::Vector<float> *WeightSum =
+      new xft::Vector<float>();  // if weight is int8
+  xft::Matrix<T> *quantizedWeight = new xft::Matrix<T>();
+  if (it_created == weight_only_hub.end()) {
+    auto weight_ptr = reinterpret_cast<const float *>(weight.data<float>());
+    xft::Matrix<T> convertedWeight;
+    mmHelper = new MMHelper(xft::DeviceKind::iCPU, 0);
+    mmHelper->convertWeight(trans,
+                            rows,
+                            cols,
+                            weight_ptr,
+                            nullptr,
+                            nullptr,
+                            convertedWeight,
+                            *WeightScale,
+                            *WeightZero,
+                            *WeightSum);
+    quantizedWeight->Resize(rows, cols);
+    mmHelper->packWeight(trans, convertedWeight, *quantizedWeight);
+    weight_only_hub[weight_only_key] =
+        std::make_tuple(quantizedWeight, WeightScale, WeightZero, WeightSum);
+    AvxCompute<T>(x,
+                  weight,
+                  trans,
+                  alog,
+                  out,
+                  *quantizedWeight,
+                  *WeightScale,
+                  *WeightZero,
+                  *WeightSum,
+                  mmHelper);
+  } else {
+    AvxCompute<T>(x,
+                  weight,
+                  trans,
+                  alog,
+                  out,
+                  *(std::get<0>(it_created->second)),
+                  *(std::get<1>(it_created->second)),
+                  *(std::get<2>(it_created->second)),
+                  *(std::get<3>(it_created->second)),
+                  mmHelper);
+  }
+}
+std::vector<paddle::Tensor> InvokeAvxWeightOnly(const paddle::Tensor &x,
+                                                const paddle::Tensor &weight,
+                                                const std::string &alog,
+                                                bool trans) {
+  auto out_shape = x.shape();
+  out_shape[out_shape.size() - 1] = weight.shape()[1];
+  auto out = paddle::empty(out_shape, x.dtype(), paddle::CPUPlace());
+  if (alog == "int8") {
+    AvxWeightOnly<int8_t>(x, weight, trans, alog, out);
+  } else if (alog == "fp16") {
+    AvxWeightOnly<float16_t>(x, weight, trans, alog, out);
+  } else {
+    AvxWeightOnly<float16_t>(x, weight, trans, alog, out);
+  }
+  return {out};
+}
+
+std::vector<std::vector<int64_t>> AvxWeightOnlyInferShape(
+    std::vector<int64_t> x_shape,
+    std::vector<int64_t> weight_shape) {
+  int m = 1;
+  for (int i = 0; i < x_shape.size() - 1; i++) {
+    m = m * x_shape[i];
+  }
+  return {std::vector<int64_t>{m, weight_shape[1]}};
+}
+
+std::vector<paddle::DataType> AvxWeightOnlyInferDtype(
+    paddle::DataType x_dtype,
+    paddle::DataType weight_dtype) {
+  return {x_dtype};
+}
+
+PD_BUILD_OP(avx_weight_only)
+    .Inputs({"x", "weight"})
+    .Outputs({"out"})
+    .Attrs({"alog: std::string", "trans:bool"})
+    .SetKernelFn(PD_KERNEL(InvokeAvxWeightOnly))
+    .SetInferShapeFn(PD_INFER_SHAPE(AvxWeightOnlyInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(AvxWeightOnlyInferDtype));

csrc/cpu/src/set_value_by_flags.cc

@@ -0,0 +1,56 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+// 
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// 
+//     http://www.apache.org/licenses/LICENSE-2.0
+// 
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/extension.h"
+
+void set_value_by_flag_and_id(const bool *stop_flags, int64_t *pre_ids_all, const int64_t *pre_ids, const int64_t *step_idx, int bs, int length) {
+    for (int bi=0;bi<bs;bi++){
+        if(!stop_flags[bi]){
+             int64_t *pre_ids_all_now = pre_ids_all + bi * length;
+             if (step_idx[bi] >= 0) {
+                pre_ids_all_now[step_idx[bi]] = pre_ids[bi];
+             }
+        }
+    }
+}
+
+std::vector<paddle::Tensor> SetValueByFlagsAndIdx(const paddle::Tensor& pre_ids_all, const paddle::Tensor& pre_ids_now, const paddle::Tensor& step_idx, const paddle::Tensor& stop_flags) {
+    std::vector<int64_t> pre_ids_all_shape = pre_ids_all.shape();
+    auto stop_flags_out = stop_flags.copy_to(stop_flags.place(), false); 
+
+    int bs = stop_flags.shape()[0];
+    int length = pre_ids_all_shape[1];
+
+    set_value_by_flag_and_id(stop_flags.data<bool>(), const_cast<int64_t*>(pre_ids_all.data<int64_t>()), pre_ids_now.data<int64_t>(), step_idx.data<int64_t>(), bs, length);
+    return {stop_flags_out};
+}
+
+std::vector<std::vector<int64_t>> SetValueByFlagsAndIdxInferShape(const std::vector<int64_t>& pre_ids_all_shape, const std::vector<int64_t>& pre_ids_now_shape,
+                                                                  const std::vector<int64_t>& step_idx_shape, const std::vector<int64_t>& stop_flags_shape) {
+    return {stop_flags_shape};
+}
+
+std::vector<paddle::DataType> SetValueByFlagsAndIdxInferDtype(const paddle::DataType& pre_ids_all_dtype,
+                                                              const paddle::DataType& pre_ids_now_dtype,
+                                                              const paddle::DataType& step_idx_dtype,
+                                                              const paddle::DataType& stop_flags_dtype) {
+    return {stop_flags_dtype};
+}
+
+PD_BUILD_OP(set_value_by_flags_and_idx)
+    .Inputs({"pre_ids_all", "pre_ids_now", "step_idx", "stop_flags"})
+    .Outputs({"stop_flags_out"})
+    .SetKernelFn(PD_KERNEL(SetValueByFlagsAndIdx))
+    .SetInferShapeFn(PD_INFER_SHAPE(SetValueByFlagsAndIdxInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(SetValueByFlagsAndIdxInferDtype));

csrc/cpu/src/setup_cpu.py

@@ -0,0 +1,134 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import site
+import subprocess
+
+from paddle.utils.cpp_extension import CppExtension, setup
+
+# from setuptools import Extension, setup
+from setuptools.command.build_ext import build_ext
+
+
+# refer: https://note.qidong.name/2018/03/setup-warning-strict-prototypes
+# Avoid a gcc warning below:
+# cc1plus: warning: command line option ‘-Wstrict-prototypes’ is valid
+# for C/ObjC but not for C++
+class BuildExt(build_ext):
+    def build_extensions(self):
+        if "-Wstrict-prototypes" in self.compiler.compiler_so:
+            self.compiler.compiler_so.remove("-Wstrict-prototypes")
+        super().build_extensions()
+
+
+def check_avx512_bf16__support():
+    try:
+        result = subprocess.run(
+            ["lscpu", "|", "grep", '"avx512_bf16"'],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            shell=True,
+        )
+
+        if "avx512_bf16" in result.stdout.lower():
+            return True
+        else:
+            return False
+
+    except Exception as e:
+        print(f"Error checking AVX512 support: {e}")
+        return False
+
+
+paddle_extra_compile_args = [
+    "-std=c++17",
+    "-shared",
+    "-fPIC",
+    "-Wno-parentheses",
+    "-DPADDLE_WITH_CUSTOM_KERNEL",
+    "-mavx512f",
+    "-mavx512vl",
+    "-fopenmp",
+    "-mavx512bw",
+    "-mno-mmx",
+    "-Wall",
+    "-march=skylake-avx512",
+    "-O3",
+    "-g",
+]
+
+if check_avx512_bf16__support():
+    paddle_extra_compile_args += [
+        "-DAVX512_BF16_WEIGHT_ONLY_BF16=true",
+        "-DAVX512_FP16_WEIGHT_ONLY_INT8=true",
+        "-DAVX512_FP16_WEIGHT_ONLY_FP16=true",
+    ]
+else:
+    paddle_extra_compile_args += [
+        "-DAVX512_FP32_WEIGHT_ONLY_FP16=true",
+        "-DAVX512_FP32_WEIGHT_ONLY_INT8=true",
+    ]
+# include path
+site_packages_path = site.getsitepackages()
+paddle_custom_kernel_include = [os.path.join(path, "paddle", "include") for path in site_packages_path]
+
+XFT_INCLUDE_DIR = os.environ["XFT_HEADER_DIR"]
+XFT_LIBRARY_DIR = os.environ["XFT_LIB_DIR"]
+
+# include path third_party
+paddle_custom_kernel_include += [
+    os.path.join(XFT_INCLUDE_DIR, "include"),
+    os.path.join(XFT_INCLUDE_DIR, "src/common"),
+    os.path.join(XFT_INCLUDE_DIR, "src/kernel"),
+    os.path.join(XFT_INCLUDE_DIR, "src/layers"),
+    os.path.join(XFT_INCLUDE_DIR, "src/models"),
+    os.path.join(XFT_INCLUDE_DIR, "src/utils"),
+    os.path.join(XFT_INCLUDE_DIR, "3rdparty/onednn/include"),
+    os.path.join(XFT_INCLUDE_DIR, "3rdparty/onednn/build/include"),
+    os.path.join(XFT_INCLUDE_DIR, "3rdparty/xdnn"),
+    os.path.join(XFT_INCLUDE_DIR, "3rdparty"),
+    os.path.join(XFT_INCLUDE_DIR, "3rdparty/mkl/include"),
+]
+
+# libs path
+paddle_custom_kernel_library_dir = [os.path.join(path, "paddle", "base") for path in site_packages_path]
+paddle_custom_kernel_library_dir += [XFT_LIBRARY_DIR]
+
+
+libs = [":libxfastertransformer.so", ":libxft_comm_helper.so"]
+
+custom_kernel_dot_module = CppExtension(
+    sources=[
+        "../gpu/save_with_output.cc",
+        "./src/token_penalty_multi_scores.cc",
+        "./src/stop_generation_multi_ends.cc",
+        "./src/set_value_by_flags.cc",
+        "./src/xft_transformer.cc",
+        "./src/avx_weight_only.cc",
+        "./src/xft_greedy_search.cc",
+    ],
+    include_dirs=paddle_custom_kernel_include,
+    library_dirs=paddle_custom_kernel_library_dir,
+    libraries=libs,
+    extra_compile_args=paddle_extra_compile_args,
+)
+
+setup(
+    name="paddlenlp_ops",
+    version="1.0",
+    description="custom kernel for compiling",
+    ext_modules=[custom_kernel_dot_module],
+)

csrc/cpu/src/stop_generation_multi_ends.cc

@@ -0,0 +1,87 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <stdlib.h>
+#include <string.h>
+#include <stdio.h>
+#include <vector>
+#include "paddle/extension.h"
+
+
+bool is_in_end(const int64_t id, const int64_t *end_ids, int length) {
+    bool flag = false;
+    for (int i = 0; i < length; i++) {
+        if (id == end_ids[i]) {
+            return true;
+        }
+    }
+    return flag;
+}
+
+void set_value_by_flags(const bool* stop_flags,
+                        const int64_t* end_ids,
+                        int64_t* topk_ids,
+                        bool* stop_flags_out,
+                        const int bs,
+                        int end_length) {
+  for (int bi = 0; bi < bs; bi++) {
+    topk_ids[bi] = stop_flags[bi] ? end_ids[0] : topk_ids[bi];
+    if (is_in_end(topk_ids[bi], end_ids, end_length)) {
+      stop_flags_out[bi] = true;
+    }
+  }
+}
+
+
+std::vector<paddle::Tensor> GetStopFlagsMulti(const paddle::Tensor& topk_ids,
+                                              const paddle::Tensor& stop_flags,
+                                              const paddle::Tensor& end_ids) {
+  PD_CHECK(topk_ids.dtype() == paddle::DataType::INT64);
+  PD_CHECK(stop_flags.dtype() == paddle::DataType::BOOL);
+
+  std::vector<int64_t> shape = topk_ids.shape();
+  int64_t bs_now = shape[0];
+  int64_t end_length = end_ids.shape()[0];
+  auto topk_ids_out = topk_ids.copy_to(topk_ids.place(), false);
+  auto stop_flags_out = stop_flags.copy_to(stop_flags.place(), false);
+  set_value_by_flags(stop_flags.data<bool>(),
+                     end_ids.data<int64_t>(),
+                     topk_ids_out.data<int64_t>(),
+                     stop_flags_out.data<bool>(),
+                     bs_now,
+                     end_length);
+
+  return {topk_ids_out, stop_flags_out};
+}
+
+std::vector<std::vector<int64_t>> GetStopFlagsMultiInferShape(
+    const std::vector<int64_t>& topk_ids_shape,
+    const std::vector<int64_t>& stop_flags_shape,
+    const std::vector<int64_t>& end_ids_shape) {
+  return {topk_ids_shape, stop_flags_shape};
+}
+
+std::vector<paddle::DataType> GetStopFlagsMultiInferDtype(
+    const paddle::DataType& topk_ids_dtype,
+    const paddle::DataType& stop_flags_dtype,
+    const paddle::DataType& end_ids_dtype) {
+  return {topk_ids_dtype, stop_flags_dtype};
+}
+
+PD_BUILD_OP(set_stop_value_multi_ends)
+    .Inputs({"topk_ids", "stop_flags", "end_ids"})
+    .Outputs({"topk_ids_out", "stop_flags_out"})
+    .SetKernelFn(PD_KERNEL(GetStopFlagsMulti))
+    .SetInferShapeFn(PD_INFER_SHAPE(GetStopFlagsMultiInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(GetStopFlagsMultiInferDtype));

csrc/cpu/src/token_penalty_multi_scores.cc

@@ -0,0 +1,192 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <paddle/extension.h>
+#include <vector>
+
+template <typename T>
+void min_length_logits_process(T* logits,
+                               const int64_t* cur_len,
+                               const int64_t* min_len,
+                               const int64_t* eos_token_id,
+                               const int64_t bs,
+                               const int64_t length,
+                               const int64_t end_length) {
+  for (int bi = 0; bi < bs; ++bi) {
+    if (cur_len[bi] < 0) {
+      continue;
+    }
+    if (cur_len[bi] < min_len[bi]) {
+      for (int i = 0; i < end_length; ++i) {
+        logits[bi * length + eos_token_id[i]] = -1e10;
+      }
+    }
+  }
+}
+
+
+void update_repeat_times(const int64_t* pre_ids,
+                         const int64_t* cur_len,
+                         int* repeat_times,
+                         const int64_t bs,
+                         const int64_t length,
+                         const int64_t length_id) {
+  for (int bi = 0; bi < bs; ++bi) {
+    if (cur_len[bi] < 0) {
+      continue;
+    }
+    const int64_t* pre_ids_now = pre_ids + bi * length_id;
+    int* repeat_times_now = repeat_times + bi * length;
+    for (int i = 0; i < length_id; i++) {
+      int64_t id = pre_ids_now[i];
+      if (id < 0) {
+        break;
+      }
+      repeat_times_now[id] += 1;
+    }
+  }
+}
+
+template <typename T>
+void update_value_by_repeat_times(const int* repeat_times,
+                                  const T* penalty_scores,
+                                  const T* frequency_score,
+                                  const T* presence_score,
+                                  T* logits,
+                                  const int64_t bs,
+                                  const int64_t length) {
+  for (int bi = 0; bi < bs; ++bi) {
+    T* logits_now = logits + bi * length;
+    const int* repeat_times_now = repeat_times + bi * length;
+    float alpha = static_cast<float>(penalty_scores[bi]);
+    float beta = static_cast<float>(frequency_score[bi]);
+    float gamma = static_cast<float>(presence_score[bi]);
+    for (int i = 0; i < length; ++i) {
+      int times = repeat_times_now[i];
+      if (times == 0) continue;
+      float logit_now = static_cast<float>(logits_now[i]);
+      logit_now = logit_now < 0 ? logit_now * alpha : logit_now / alpha;
+      logits_now[i] = static_cast<T>(logit_now - times * beta - gamma);
+    }
+  }
+}
+template <paddle::DataType D>
+std::vector<paddle::Tensor> token_penalty_multi_scores_kernel(
+    const paddle::Tensor& pre_ids,
+    const paddle::Tensor& logits,
+    const paddle::Tensor& penalty_scores,
+    const paddle::Tensor& frequency_score,
+    const paddle::Tensor& presence_score,
+    const paddle::Tensor& cur_len,
+    const paddle::Tensor& min_len,
+    const paddle::Tensor& eos_token_id) {
+  std::vector<int64_t> shape = logits.shape();
+  auto repeat_times =
+      paddle::full(shape, 0, paddle::DataType::INT32, pre_ids.place());
+  int64_t bs = shape[0];
+  int64_t length = shape[1];
+  int64_t length_id = pre_ids.shape()[1];
+  auto logits_out = logits.copy_to(logits.place(), false);
+  int64_t end_length = eos_token_id.shape()[0];
+
+  min_length_logits_process(const_cast<float*>(logits_out.data<float>()),
+                            cur_len.data<int64_t>(),
+                            min_len.data<int64_t>(),
+                            eos_token_id.data<int64_t>(),
+                            bs,
+                            length,
+                            end_length);
+  update_repeat_times(pre_ids.data<int64_t>(),
+                      cur_len.data<int64_t>(),
+                      repeat_times.data<int>(),
+                      bs,
+                      length,
+                      length_id);
+  update_value_by_repeat_times(
+      repeat_times.data<int>(),
+      const_cast<float*>(penalty_scores.data<float>()),
+      const_cast<float*>(frequency_score.data<float>()),
+      const_cast<float*>(presence_score.data<float>()),
+      const_cast<float*>(logits_out.data<float>()),
+      bs,
+      length);
+  return {logits_out};
+}
+
+std::vector<paddle::Tensor> TokenPenaltyMultiScores(
+    const paddle::Tensor& pre_ids,
+    const paddle::Tensor& logits,
+    const paddle::Tensor& penalty_scores,
+    const paddle::Tensor& frequency_scores,
+    const paddle::Tensor& presence_scores,
+    const paddle::Tensor& cur_len,
+    const paddle::Tensor& min_len,
+    const paddle::Tensor& eos_token_id) {
+  switch (logits.type()) {
+    case paddle::DataType::FLOAT32: {
+      return token_penalty_multi_scores_kernel<paddle::DataType::FLOAT32>(
+          pre_ids,
+          logits,
+          penalty_scores,
+          frequency_scores,
+          presence_scores,
+          cur_len,
+          min_len,
+          eos_token_id);
+    }
+    default: {
+      PD_THROW(
+          "NOT supported data type. "
+          "Only float32 are supported. ");
+      break;
+    }
+  }
+}
+
+std::vector<std::vector<int64_t>> TokenPenaltyMultiScoresInferShape(
+    const std::vector<int64_t>& pre_ids_shape,
+    const std::vector<int64_t>& logits_shape,
+    const std::vector<int64_t>& penalty_scores_shape,
+    const std::vector<int64_t>& frequency_scores_shape,
+    const std::vector<int64_t>& presence_scores_shape,
+    const std::vector<int64_t>& cur_len_shape,
+    const std::vector<int64_t>& min_len_shape,
+    const std::vector<int64_t>& eos_token_id_shape) {
+  return {logits_shape};
+}
+
+std::vector<paddle::DataType> TokenPenaltyMultiScoresInferDtype(
+    const paddle::DataType& pre_ids_dtype,
+    const paddle::DataType& logits_dtype,
+    const paddle::DataType& penalty_scores_dtype,
+    const paddle::DataType& frequency_scores_dtype,
+    const paddle::DataType& presence_scores_dtype,
+    const paddle::DataType& cur_len_dtype,
+    const paddle::DataType& min_len_dtype,
+    const paddle::DataType& eos_token_id_dtype) {
+  return {logits_dtype};
+}
+PD_BUILD_OP(get_token_penalty_multi_scores)
+    .Inputs({"pre_ids",
+             "logits",
+             "penalty_scores",
+             "frequency_scores",
+             "presence_scores",
+             "cur_len",
+             "min_len",
+             "eos_token_id"})
+    .Outputs({"logits_out"})
+    .SetKernelFn(PD_KERNEL(TokenPenaltyMultiScores))
+    .SetInferShapeFn(PD_INFER_SHAPE(TokenPenaltyMultiScoresInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(TokenPenaltyMultiScoresInferDtype));

csrc/cpu/src/xft_greedy_search.cc

@@ -0,0 +1,131 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <omp.h>
+
+#include <cstdio>
+#include <iostream>
+
+#include "paddle/extension.h"
+
+void GreedySearch(const float *probs,
+                  int64_t *next_token_ids,
+                  int bsz,
+                  int vocab_size) {
+  int numThreads = 0;
+#pragma omp parallel
+  {
+    int tid = omp_get_thread_num();
+    if (tid == 0) {
+      numThreads = omp_get_num_threads();
+    }
+  }
+  // Max ID and value for each sample
+  // std::vector<int> maxIds(batchSize);
+  float maxVals[bsz];
+
+  // Small batch size (each sample can have at least 2 threads)
+  if (numThreads / bsz >= 2) {
+    int thrPerSample = numThreads / bsz;
+    int sizePerThr = (vocab_size + thrPerSample - 1) / thrPerSample;
+    int maxIndices[bsz * thrPerSample];
+    float maxValues[bsz * thrPerSample];
+
+    // TODO: if size is small, possible to cause out of boundary
+#pragma omp parallel for collapse(2)
+    for (int b = 0; b < bsz; ++b) {
+      for (int t = 0; t < thrPerSample;
+           ++t) {  // thread index inside the sample
+        int start = t * sizePerThr;
+        int end = (start + sizePerThr) > vocab_size ? vocab_size
+                                                    : (start + sizePerThr);
+        const float *p = probs + b * vocab_size;
+
+        int maxIdx = start;
+        float maxVal = p[start];
+        for (int off = start + 1; off < end; ++off) {
+          if (p[off] > maxVal) {
+            maxVal = p[off];
+            maxIdx = off;
+          }
+        }
+
+        // False sharing happens, but since only one time, not avoided
+        maxIndices[b * thrPerSample + t] = maxIdx;
+        maxValues[b * thrPerSample + t] = maxVal;
+      }
+    }
+
+    // Local reduction
+    for (int i = 0; i < bsz; ++i) {
+      int *pIndices = maxIndices + i * thrPerSample;
+      float *pValues = maxValues + i * thrPerSample;
+      int maxIdx = pIndices[0];
+      float maxVal = pValues[0];
+      for (int j = 1; j < thrPerSample; ++j) {
+        if (pValues[j] > maxVal) {
+          maxVal = pValues[j];
+          maxIdx = pIndices[j];
+        }
+      }
+      next_token_ids[i] = maxIdx;
+      maxVals[i] = maxVal;
+    }
+  }
+
+  // Each thread handle one sample (one row)
+  else {
+#pragma omp parallel for
+    for (int i = 0; i < bsz; ++i) {
+      int maxId = 0;
+      const float *p = probs + i * vocab_size;
+      float maxVal = p[0];
+      for (int j = 1; j < vocab_size; ++j) {
+        if (p[j] > maxVal) {
+          maxVal = p[j];
+          maxId = j;
+        }
+      }
+      next_token_ids[i] = maxId;
+      maxVals[i] = maxVal;
+    }
+  }
+  return;
+}
+std::vector<paddle::Tensor> XftGreedySearch(const paddle::Tensor &probs) {
+  const int bsz = probs.shape()[0];
+  const int vocab_size = probs.shape()[1];
+  auto next_tokens =
+      paddle::empty({bsz, 1}, paddle::DataType::INT64, probs.place());
+  GreedySearch(probs.data<float>(),
+               const_cast<int64_t *>(next_tokens.data<int64_t>()),
+               bsz,
+               vocab_size);
+  return {next_tokens};
+}
+std::vector<std::vector<int64_t>> XftGreedySearchInferShape(
+    const std::vector<int64_t> &probs_shape) {
+  int64_t bsz = probs_shape[0];
+  return {{bsz, 1}};
+}
+std::vector<paddle::DataType> XftGreedySearchInferDtype(
+    const paddle::DataType &probs_dtype) {
+  return {paddle::DataType::INT64};
+}
+PD_BUILD_OP(xft_greedy_search)
+    .Inputs({"probs"})
+    .Outputs({"next_tokens_ids"})
+    .SetInferShapeFn(PD_INFER_SHAPE(XftGreedySearchInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(XftGreedySearchInferDtype))
+    .SetKernelFn(PD_KERNEL(XftGreedySearch));

csrc/cpu/src/xft_transformer.cc

@@ -0,0 +1,190 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "layers_decoder.h"
+#include "paddle/extension.h"
+
+
+std::vector<paddle::Tensor> InvokeXftTransformer(
+    const paddle::Tensor &input,
+    const std::vector<paddle::Tensor> &ln1Gamma,
+    const std::vector<paddle::Tensor> &qkvWeight,
+    const std::vector<paddle::Tensor> &attnOutWeight,
+    const std::vector<paddle::Tensor> &ln2Gamma,
+    const std::vector<paddle::Tensor> &gateWeight,
+    const std::vector<paddle::Tensor> &upWeight,
+    const std::vector<paddle::Tensor> &downWeight,
+    const paddle::Tensor &pastSeqLen,
+    const paddle::Tensor &currentSeqLen,
+    const paddle::Tensor &step,
+    int hiddensize,
+    int totalLayer,
+    const std::string &computeType,
+    const std::string &cacheType,
+    const std::string &activation,
+    const std::string &normType,
+    int attHeadDim,
+    int attHeadNum,
+    int kvHeadNum,
+    int maxPositions,
+    int maxPosEmbed,
+    int intermediateSize) {
+  auto out = paddle::empty_like(input);
+  auto batchSize = input.shape()[0];
+  auto inputSeqLen = input.shape()[1];
+  auto past_seq_len = pastSeqLen.data<int64_t>()[0];
+
+  auto cur_seq_len = currentSeqLen.data<int64_t>()[0];
+
+  auto step_id = step.data<int64_t>()[0];
+  auto output_ptr = reinterpret_cast<void *>(out.data<float>());
+  auto xft_data_type = xft::DataType::fp16;
+  if (computeType == "bf16") {
+    xft_data_type = xft::DataType::bf16;
+  } else if (computeType == "bf16_int8") {
+    xft_data_type = xft::DataType::bf16_int8;
+  } else if (computeType == "fp16_int8") {
+    xft_data_type = xft::DataType::fp16_int8;
+  }
+  auto kvc_type = xft::DataType::fp16;
+  if (cacheType == "int8") {
+    kvc_type = xft::DataType::int8;
+  }
+  auto rope_type = xft::RopeType::LLAMA_ROPE;
+  auto act_type = xft::ActivationType::SILU;
+  if (activation == "relu") {
+    act_type = xft::ActivationType::RELU;
+  } else if (activation == "gelu") {
+    act_type = xft::ActivationType::GELU;
+  } else if (activation == "swiglu") {
+    act_type = xft::ActivationType::SWIGLU;
+  }
+  auto norm_type = xft::NormType::RMS;
+  if (normType == "layernorm") {
+    norm_type = xft::NormType::LN;
+  }
+  auto input_ptr = reinterpret_cast<const void *>(input.data<float>());
+  auto gate_bias_ptr = nullptr;
+  auto up_bias_ptr = nullptr;
+  auto down_bias_ptr = nullptr;
+  auto ln2Beta_ptr = nullptr;
+  auto qkvBiasWeight_ptr = nullptr;
+  auto attnOutBias_ptr = nullptr;
+  auto ln1Beta_ptr = nullptr;
+  for (int i = 0; i < totalLayer; ++i) {
+    auto ln1Gamma_ptr =
+        reinterpret_cast<const float *>(ln1Gamma[i].data<float>());
+    auto qkvWeight_ptr =
+        reinterpret_cast<const void *>(qkvWeight[i].data<float>());
+    auto attnOutWeight_ptr =
+        reinterpret_cast<const void *>(attnOutWeight[i].data<float>());
+    auto ln2Gamma_ptr =
+        reinterpret_cast<const float *>(ln2Gamma[i].data<float>());
+    auto gate_weight_ptr =
+        reinterpret_cast<const void *>(gateWeight[i].data<float>());
+    auto up_weight_ptr =
+        reinterpret_cast<const void *>(upWeight[i].data<float>());
+    auto down_weight_ptr =
+        reinterpret_cast<const void *>(downWeight[i].data<float>());
+    invokeLayerLLaMA(
+        xft_data_type,               // dt
+        kvc_type,                    // kvcdt
+        rope_type,                   // rope
+        act_type,                    // at
+        norm_type,                   // nt
+        batchSize,                   // batchSize
+        inputSeqLen,                 // inputSeqLen
+        attHeadDim,                  // attHeadDim
+        attHeadNum,                  // attHeadNum
+        kvHeadNum,                   // kvHeadNum
+        maxPositions,                // maxPositions
+        maxPosEmbed,                 // maxPosEmbed
+        past_seq_len,                // pastSeqLen
+        cur_seq_len,                 // currentSeqLen
+        step_id,                     // step
+        hiddensize,                  // hiddenSize
+        intermediateSize,            // intermediateSize
+        (void *)output_ptr,          // output
+        hiddensize,                  // outputStride
+        input_ptr,                   // input
+        hiddensize,                  // inputStride
+        ln1Gamma_ptr,                // ln1Gamma
+        ln1Beta_ptr,                 // ln1Beta
+        qkvWeight_ptr,               // queryWeight
+        qkvWeight_ptr + hiddensize,  // keyWeight
+        qkvWeight_ptr + hiddensize + kvHeadNum * attHeadDim,  // valueWeight
+        attnOutWeight_ptr,                                    // attnOutWeight
+        ln2Gamma_ptr,                                         // ln2Gamma
+        ln2Beta_ptr,                                          // ln2Beta
+        gate_weight_ptr,
+        up_weight_ptr,
+        down_weight_ptr,
+        nullptr,          // queryBias
+        nullptr,          // keyBias
+        nullptr,          // valueBias
+        attnOutBias_ptr,  // attnOutBias
+        qkvWeight_ptr,    // myqkvWeight
+        gate_bias_ptr,
+        up_bias_ptr,
+        down_bias_ptr,
+        qkvBiasWeight_ptr);
+    if (i < totalLayer - 1) {
+      memcpy(const_cast<void *>(input_ptr),
+             output_ptr,
+             batchSize * inputSeqLen * hiddensize * sizeof(float));
+    }
+  }
+  return {out};
+}
+
+std::vector<std::vector<int64_t>> XftTransformerInferShape(
+    std::vector<int64_t> x_shape) {
+  return {x_shape};
+}
+
+std::vector<paddle::DataType> XftTransformerInferDtype(
+    paddle::DataType x_dtype) {
+  return {x_dtype};
+}
+
+PD_BUILD_OP(xft_transformer)
+    .Inputs({
+        "x",
+        paddle::Vec("ln1Gamma"),
+        paddle::Vec("qkvWeight"),
+        paddle::Vec("attnOutWeight"),
+        paddle::Vec("ln2Gamma"),
+        paddle::Vec("gateWeight"),
+        paddle::Vec("upWeight"),
+        paddle::Vec("downWeight"),
+        "pastSeqLen",
+        "currentSeqLen",
+        "step",
+    })
+    .Outputs({"out"})
+    .Attrs({"hiddensize :int",
+            "totalLayer :int",
+            "computeType : std::string",
+            "cacheType :std::string",
+            "activation :std::string",
+            "normType :std::string",
+            "attHeadDim: int",
+            "attHeadNum: int",
+            "kvHeadNum: int",
+            "maxPositions: int",
+            "maxPosEmbed: int",
+            "intermediateSize: int"})
+    .SetKernelFn(PD_KERNEL(InvokeXftTransformer))
+    .SetInferShapeFn(PD_INFER_SHAPE(XftTransformerInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(XftTransformerInferDtype));

csrc/gpu/all_reduce.cu

@@ -0,0 +1,125 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "helper.h"
+#include "all_reduce.cuh"
+
+// Fake pointer type, must match fptr_t type in ops.h.
+// We use this type alias to indicate when pointers are passed in as int64_t.
+using fptr_t = int64_t;
+static_assert(sizeof(void*) == sizeof(fptr_t));
+
+fptr_t init_custom_all_reduce(const std::vector<fptr_t>& fake_ipc_ptrs,
+                      paddle::Tensor& rank_data, int64_t rank,
+                      bool full_nvlink) {
+  int world_size = fake_ipc_ptrs.size();
+  if (world_size > 8)
+    throw std::invalid_argument("world size > 8 is not supported");
+  if (world_size % 2 != 0)
+    throw std::invalid_argument("Odd num gpus is not supported for now");
+  if (rank < 0 || rank >= world_size)
+    throw std::invalid_argument("invalid rank passed in");
+
+  paddle::Signal* ipc_ptrs[8];
+  for (int i = 0; i < world_size; i++) {
+    ipc_ptrs[i] = reinterpret_cast<paddle::Signal*>(fake_ipc_ptrs[i]);
+  }
+  return (fptr_t) new paddle::CustomAllreduce(ipc_ptrs, rank_data.data(),
+                                            rank_data.numel(), rank, world_size,
+                                            full_nvlink);
+}
+
+/**
+ * Performs an out-of-place allreduce and stores result in out.
+ *
+ * If _reg_buffer is null, assumes inp.data() is already IPC-registered.
+ * Otherwise, _reg_buffer is assumed to be IPC-registered and inp is first
+ * copied into _reg_buffer.
+ */
+void all_reduce(fptr_t _fa, paddle::Tensor& inp, paddle::Tensor& out,
+                fptr_t _reg_buffer, int64_t reg_buffer_sz_bytes) {
+  auto fa = reinterpret_cast<paddle::CustomAllreduce*>(_fa);
+  auto stream = inp.stream();
+
+  auto input_size = inp.numel() * 2;
+  auto reg_buffer = reinterpret_cast<void*>(_reg_buffer);
+  if (reg_buffer) {
+    cudaMemcpyAsync(reg_buffer, inp.data(), input_size,
+                                  cudaMemcpyDeviceToDevice, stream);
+  } else {
+    reg_buffer = inp.data();
+  }
+  switch (out.dtype()) {
+    case phi::DataType::FLOAT32: {
+      fa->allreduce<float>(stream, reinterpret_cast<float*>(reg_buffer),
+                           reinterpret_cast<float*>(out.data()),
+                           out.numel());
+      break;
+    }
+    case phi::DataType::FLOAT16: {
+      fa->allreduce<half>(stream, reinterpret_cast<half*>(reg_buffer),
+                          reinterpret_cast<half*>(out.data()), out.numel());
+      break;
+    }
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800)
+    case phi::DataType::BFLOAT16: {
+      fa->allreduce<nv_bfloat16>(
+          stream, reinterpret_cast<nv_bfloat16*>(reg_buffer),
+          reinterpret_cast<nv_bfloat16*>(out.data()), out.numel());
+      break;
+    }
+#endif
+    default:
+      throw std::runtime_error(
+          "custom allreduce only supports float32, float16 and bfloat16");
+  }
+}
+
+void dispose(fptr_t _fa) {
+  delete reinterpret_cast<paddle::CustomAllreduce*>(_fa);
+}
+
+int64_t meta_size() { return sizeof(paddle::Signal); }
+
+void register_buffer(fptr_t _fa, const std::vector<fptr_t>& fake_ipc_ptrs) {
+  auto fa = reinterpret_cast<paddle::CustomAllreduce*>(_fa);
+  void* ipc_ptrs[8];
+  for (int i = 0; i < fake_ipc_ptrs.size(); i++) {
+    ipc_ptrs[i] = reinterpret_cast<void*>(fake_ipc_ptrs[i]);
+  }
+  fa->register_buffer(ipc_ptrs);
+}
+
+// Use vector<int64_t> to represent byte data for python binding compatibility.
+std::tuple<std::vector<int64_t>, std::vector<int64_t>>
+get_graph_buffer_ipc_meta(fptr_t _fa) {
+  auto fa = reinterpret_cast<paddle::CustomAllreduce*>(_fa);
+  auto [handle, offsets] = fa->get_graph_buffer_ipc_meta();
+  std::vector<int64_t> bytes(handle.begin(), handle.end());
+  return std::make_tuple(bytes, offsets);
+}
+
+// Use vector<int64_t> to represent byte data for python binding compatibility.
+void register_graph_buffers(fptr_t _fa,
+                            const std::vector<std::vector<int64_t>>& handles,
+                            const std::vector<std::vector<int64_t>>& offsets) {
+  auto fa = reinterpret_cast<paddle::CustomAllreduce*>(_fa);
+  std::vector<std::string> bytes;
+  bytes.reserve(handles.size());
+  for (int i = 0; i < handles.size(); i++) {
+    bytes.emplace_back(handles[i].begin(), handles[i].end());
+  }
+  bytes.reserve(handles.size());
+  fa->register_graph_buffers(bytes, offsets);
+}

csrc/gpu/all_reduce.cuh

@@ -0,0 +1,528 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#pragma once
+
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include <iostream>
+#include <array>
+#include <limits>
+#include <map>
+#include <unordered_map>
+#include <vector>
+
+#define CUDACHECK(cmd)                                              \
+  do {                                                              \
+    cudaError_t e = cmd;                                            \
+    if (e != cudaSuccess) {                                         \
+      printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, \
+             cudaGetErrorString(e));                                \
+      exit(EXIT_FAILURE);                                           \
+    }                                                               \
+  } while (0)
+
+namespace paddle {
+
+constexpr int kMaxBlocks = 36;
+// Counter may overflow, but it's fine since unsigned int overflow is
+// well-defined behavior.
+using FlagType = uint32_t;
+struct Signal {
+  alignas(128) FlagType self_counter[kMaxBlocks][8];
+  // Two sets of peer counters are needed for two syncs. The reason is that
+  // it's possible for peer GPU block to arrive at the second sync point while
+  // the current GPU block haven't passed the first sync point. Thus, peer GPU
+  // may write counter+1 while current GPU is busy waiting for counter. We use
+  // alternating counter array to avoid this possibility.
+  alignas(128) FlagType peer_counter[2][kMaxBlocks][8];
+};
+
+struct __align__(16) RankData {
+  const void* __restrict__ ptrs[8];
+};
+
+struct __align__(16) RankSignals {
+  Signal* signals[8];
+};
+
+// like std::array, but aligned
+template <typename T, int sz>
+struct __align__(alignof(T) * sz) array_t {
+  T data[sz];
+  using type = T;
+  static constexpr int size = sz;
+};
+
+// use packed type to maximize memory efficiency
+// goal: generate ld.128 and st.128 instructions
+template <typename T>
+struct packed_t {
+  // the (P)acked type for load/store
+  using P = array_t<T, 16 / sizeof(T)>;
+  // the (A)ccumulator type for reduction
+  using A = array_t<float, 16 / sizeof(T)>;
+};
+
+#define DINLINE __device__ __forceinline__
+
+// scalar cast functions
+DINLINE float upcast_s(half val) { return __half2float(val); }
+
+template <typename T>
+DINLINE T downcast_s(float val);
+template <>
+DINLINE half downcast_s(float val) {
+  return __float2half(val);
+}
+
+// scalar add functions
+// for some reason when compiling with Paddle, the + operator for half and
+// bfloat is disabled so we call the intrinsics directly
+DINLINE half& assign_add(half& a, half b) {
+  a = __hadd(a, b);
+  return a;
+}
+DINLINE float& assign_add(float& a, float b) { return a += b; }
+
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800)
+DINLINE float upcast_s(nv_bfloat16 val) { return __bfloat162float(val); }
+template <>
+DINLINE nv_bfloat16 downcast_s(float val) {
+  return __float2bfloat16(val);
+}
+DINLINE nv_bfloat16& assign_add(nv_bfloat16& a, nv_bfloat16 b) {
+  a = __hadd(a, b);
+  return a;
+}
+#endif
+
+template <typename T, int N>
+DINLINE array_t<T, N>& packed_assign_add(array_t<T, N>& a, array_t<T, N> b) {
+#pragma unroll
+  for (int i = 0; i < N; i++) {
+    assign_add(a.data[i], b.data[i]);
+  }
+  return a;
+}
+
+template <typename T, int N>
+DINLINE array_t<float, N> upcast(array_t<T, N> val) {
+  if constexpr (std::is_same<T, float>::value) {
+    return val;
+  } else {
+    array_t<float, N> out;
+#pragma unroll
+    for (int i = 0; i < N; i++) {
+      out.data[i] = upcast_s(val.data[i]);
+    }
+    return out;
+  }
+}
+
+template <typename O>
+DINLINE O downcast(array_t<float, O::size> val) {
+  if constexpr (std::is_same<typename O::type, float>::value) {
+    return val;
+  } else {
+    O out;
+#pragma unroll
+    for (int i = 0; i < O::size; i++) {
+      out.data[i] = downcast_s<typename O::type>(val.data[i]);
+    }
+    return out;
+  }
+}
+
+static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
+  asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag),
+               "l"(flag_addr));
+#else
+  asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag),
+               "l"(flag_addr));
+#endif
+}
+
+static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) {
+  FlagType flag;
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
+  asm volatile("ld.acquire.sys.global.u32 %0, [%1];"
+               : "=r"(flag)
+               : "l"(flag_addr));
+#else
+  asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;"
+               : "=r"(flag)
+               : "l"(flag_addr));
+#endif
+  return flag;
+}
+
+static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) {
+  asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr));
+}
+
+static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) {
+  FlagType flag;
+  asm volatile("ld.volatile.global.u32 %0, [%1];"
+               : "=r"(flag)
+               : "l"(flag_addr));
+  return flag;
+}
+
+// is_start: whether this is the very first synchronization barrier.
+// need_fence: whether a memory fence is needed. If true, a release-acquire
+// semantic is used to enforce memory access order before and after this
+// barrier.
+template <int ngpus, bool is_start, bool need_fence = false>
+DINLINE void multi_gpu_barrier(const RankSignals& sg, Signal* self_sg,
+                               int rank) {
+  if constexpr (!is_start) __syncthreads();
+  static_assert(
+      !(is_start && need_fence));  // Start barrier shouldn't need fence.
+  if (threadIdx.x < ngpus) {
+    // Increment the counter. Technically we only need one counter, but we use
+    // multiple per block to eliminate the need to share the counter via smem.
+    auto val = self_sg->self_counter[blockIdx.x][threadIdx.x] += 1;
+    // Write the expected counter value to peer and wait for correct value from
+    // peer.
+    auto peer_counter_ptr =
+        &sg.signals[threadIdx.x]->peer_counter[val % 2][blockIdx.x][rank];
+    auto self_counter_ptr =
+        &self_sg->peer_counter[val % 2][blockIdx.x][threadIdx.x];
+    if constexpr (need_fence) {
+      st_flag_release(peer_counter_ptr, val);
+      while (ld_flag_acquire(self_counter_ptr) != val);
+    } else {
+      st_flag_volatile(peer_counter_ptr, val);
+      while (ld_flag_volatile(self_counter_ptr) != val);
+    }
+  }
+  if constexpr (is_start || need_fence) __syncthreads();
+}
+
+template <typename P, int ngpus, typename A>
+DINLINE P packed_reduce(const P* ptrs[], int idx) {
+  A tmp = upcast(ptrs[0][idx]);
+#pragma unroll
+  for (int i = 1; i < ngpus; i++) {
+    packed_assign_add(tmp, upcast(ptrs[i][idx]));
+  }
+  return downcast<P>(tmp);
+}
+
+template <typename T, int ngpus>
+__global__ void __launch_bounds__(512, 1)
+    cross_device_reduce_1stage(RankData* _dp, RankSignals sg, Signal* self_sg,
+                               T* __restrict__ result, int rank, int size) {
+  using P = typename packed_t<T>::P;
+  using A = typename packed_t<T>::A;
+  // note: we don't reorder the address so the accumulation order is the same
+  // for all ranks, ensuring bitwise identical results
+  auto dp = *_dp;
+  multi_gpu_barrier<ngpus, true>(sg, self_sg, rank);
+  // do the actual reduction
+  for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
+       idx += gridDim.x * blockDim.x) {
+    ((P*)result)[idx] = packed_reduce<P, ngpus, A>((const P**)&dp.ptrs[0], idx);
+  }
+  multi_gpu_barrier<ngpus, false>(sg, self_sg, rank);
+}
+
+template <typename P>
+DINLINE P* get_tmp_buf(Signal* sg) {
+  return (P*)(((Signal*)sg) + 1);
+}
+
+template <typename T, int ngpus>
+__global__ void __launch_bounds__(512, 1)
+    cross_device_reduce_2stage(RankData* _dp, RankSignals sg, Signal* self_sg,
+                               T* __restrict__ result, int rank, int size) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = gridDim.x * blockDim.x;
+  using P = typename packed_t<T>::P;
+  using A = typename packed_t<T>::A;
+  int part = size / ngpus;
+  int start = rank * part;
+  int end = rank == ngpus - 1 ? size : start + part;
+  int largest_part = part + size % ngpus;
+  const P* ptrs[ngpus];
+  P* tmps[ngpus];
+#pragma unroll
+  for (int i = 0; i < ngpus; i++) {
+    int target = (rank + i) % ngpus;
+    ptrs[i] = (const P*)_dp->ptrs[target];
+    tmps[i] = get_tmp_buf<P>(sg.signals[target]);
+  }
+  auto tmp_out = tmps[0];
+  multi_gpu_barrier<ngpus, true>(sg, self_sg, rank);
+  // stage 1: reduce scatter
+  for (int idx = start + tid; idx < end; idx += stride) {
+    tmp_out[idx - start] = packed_reduce<P, ngpus, A>(ptrs, idx);
+  }
+  multi_gpu_barrier<ngpus, false, true>(sg, self_sg, rank);
+
+  // stage 2: allgather. Note: it's important to match the tid between
+  // the two stages, because visibility across devices is only guaranteed
+  // between threads that have the same tid. If thread i computes the sum of
+  // start + i in the first stage, then thread i also gathers start + i from all
+  // ranks.
+  for (int idx = tid; idx < largest_part; idx += stride) {
+#pragma unroll
+    for (int i = 0; i < ngpus; i++) {
+      int gather_from_rank = ((rank + i) % ngpus);
+      if (gather_from_rank == ngpus - 1 || idx < part) {
+        int dst_idx = gather_from_rank * part + idx;
+        ((P*)result)[dst_idx] = tmps[i][idx];
+      }
+    }
+  }
+}
+
+using IPC_KEY = std::array<uint8_t, sizeof(cudaIpcMemHandle_t)>;
+static_assert(sizeof(IPC_KEY) == sizeof(cudaIpcMemHandle_t));
+static_assert(alignof(IPC_KEY) == alignof(cudaIpcMemHandle_t));
+
+class CustomAllreduce {
+ public:
+  int rank_;
+  int world_size_;
+  bool full_nvlink_;
+
+  RankSignals sg_;
+  // Stores an map from a pointer to its peer pointters from all ranks.
+  std::unordered_map<void*, RankData*> buffers_;
+  Signal* self_sg_;
+
+  // Stores rank data from all ranks. This is mainly for cuda graph purposes.
+  // For cuda graph to work, all kernel arguments must be fixed during graph
+  // capture time. However, the peer pointers are not known during graph capture
+  // time. Therefore, during capture, we increment the rank data pointer and use
+  // that as the argument to the kernel. The kernel arguments are stored in
+  // graph_unreg_buffers_. The actual peer pointers will be filled in at the
+  // memory pointed to by the pointers in graph_unreg_buffers_ when
+  // the IPC handles are exchanged between ranks.
+  //
+  // The overall process looks like this:
+  // 1. Graph capture.
+  // 2. Each rank obtains the IPC handles for each addresses used during cuda
+  // graph capture using get_graph_buffer_ipc_meta.
+  // 3. (In Python) all gather the IPC handles.
+  // 4. Obtain the peer pointers by opening the IPC handles, and store them in
+  // the rank data array at corresponding positions.
+  RankData *d_rank_data_base_, *d_rank_data_end_;
+  std::vector<void*> graph_unreg_buffers_;
+  // a map from IPC handles to opened IPC pointers
+  std::map<IPC_KEY, char*> ipc_handles_;
+
+  /**
+   * Signals are an array of ipc-enabled buffers from all ranks.
+   * For each of the buffer, the layout is as follows:
+   * | -- sizeof(Signal) -- | ------ a few MB ----- |
+   * The first section is for allreduce synchronization, and the second section
+   * is for storing the intermediate results required by some allreduce algos.
+   *
+   * Note: this class does not own any device memory. Any required buffers
+   * are passed in from the constructor.
+   */
+  CustomAllreduce(Signal** signals, void* rank_data, size_t rank_data_sz,
+                  int rank, int world_size, bool full_nvlink = true)
+      : rank_(rank),
+        world_size_(world_size),
+        full_nvlink_(full_nvlink),
+        self_sg_(signals[rank]),
+        d_rank_data_base_(reinterpret_cast<RankData*>(rank_data)),
+        d_rank_data_end_(d_rank_data_base_ + rank_data_sz / sizeof(RankData)) {
+    for (int i = 0; i < world_size_; i++) {
+      sg_.signals[i] = signals[i];
+    }
+  }
+
+  char* open_ipc_handle(const void* ipc_handle) {
+    auto [it, new_handle] =
+        ipc_handles_.insert({*((IPC_KEY*)ipc_handle), nullptr});
+    if (new_handle) {
+      char* ipc_ptr;
+      CUDACHECK(cudaIpcOpenMemHandle((void**)&ipc_ptr,
+                                     *((const cudaIpcMemHandle_t*)ipc_handle),
+                                     cudaIpcMemLazyEnablePeerAccess));
+      it->second = ipc_ptr;
+    }
+    return it->second;
+  }
+
+  std::pair<std::string, std::vector<int64_t>> get_graph_buffer_ipc_meta() {
+    auto num_buffers = graph_unreg_buffers_.size();
+    auto handle_sz = sizeof(cudaIpcMemHandle_t);
+    std::string handles(handle_sz * num_buffers, static_cast<char>(0));
+    std::vector<int64_t> offsets(num_buffers);
+    for (int i = 0; i < num_buffers; i++) {
+      auto ptr = graph_unreg_buffers_[i];
+      void* base_ptr;
+      // note: must share the base address of each allocation, or we get wrong
+      // address
+      if (cuPointerGetAttribute(&base_ptr,
+                                CU_POINTER_ATTRIBUTE_RANGE_START_ADDR,
+                                (CUdeviceptr)ptr) != CUDA_SUCCESS)
+        throw std::runtime_error("failed to get pointer attr");
+      CUDACHECK(cudaIpcGetMemHandle(
+          (cudaIpcMemHandle_t*)&handles[i * handle_sz], base_ptr));
+      offsets[i] = ((char*)ptr) - ((char*)base_ptr);
+    }
+    return std::make_pair(handles, offsets);
+  }
+
+  void check_rank_data_capacity(size_t num = 1) {
+    if (d_rank_data_base_ + num > d_rank_data_end_)
+      throw std::runtime_error(
+          "Rank data buffer is overflowed by " +
+          std::to_string(d_rank_data_base_ + num - d_rank_data_end_));
+  }
+
+  /**
+   * Register already-shared IPC pointers.
+   */
+  void register_buffer(void** ptrs) {
+    check_rank_data_capacity();
+    RankData data;
+    for (int i = 0; i < world_size_; i++) {
+      data.ptrs[i] = ptrs[i];
+    }
+    auto d_data = d_rank_data_base_++;
+    CUDACHECK(
+        cudaMemcpy(d_data, &data, sizeof(RankData), cudaMemcpyHostToDevice));
+    buffers_[ptrs[rank_]] = d_data;
+  }
+
+  // Note: when registering graph buffers, we intentionally choose to not
+  // deduplicate the addresses. That means if the allocator reuses some
+  // addresses, they will be registered again. This is to account for the remote
+  // possibility of different allocation patterns between ranks. For example,
+  // rank 1 may get the same input address for the second allreduce, but rank 2
+  // got a different address. IPC handles have internal reference counting
+  // mechanism so overhead should be small.
+  void register_graph_buffers(
+      const std::vector<std::string>& handles,
+      const std::vector<std::vector<int64_t>>& offsets) {
+    auto num_buffers = graph_unreg_buffers_.size();
+    check_rank_data_capacity(num_buffers);
+    std::vector<RankData> rank_data(num_buffers);
+    for (int i = 0; i < num_buffers; i++) {
+      auto self_ptr = graph_unreg_buffers_[i];
+      auto& rd = rank_data[i];
+      for (int j = 0; j < world_size_; j++) {
+        if (j != rank_) {
+          char* handle =
+              open_ipc_handle(&handles[j][i * sizeof(cudaIpcMemHandle_t)]);
+          handle += offsets[j][i];
+          rd.ptrs[j] = handle;
+        } else {
+          rd.ptrs[j] = self_ptr;
+        }
+      }
+    }
+    CUDACHECK(cudaMemcpy(d_rank_data_base_, rank_data.data(),
+                         sizeof(RankData) * num_buffers,
+                         cudaMemcpyHostToDevice));
+    d_rank_data_base_ += num_buffers;
+    graph_unreg_buffers_.clear();
+  }
+
+  /**
+   * Performs allreduce, assuming input has already been registered.
+   *
+   * Block and grid default configs are results after careful grid search. Using
+   * 36 blocks give the best or close to the best runtime on the devices I
+   * tried: A100, A10, A30, T4, V100. You'll notice that NCCL kernels also only
+   * take a small amount of SMs. Not quite sure the underlying reason, but my
+   * guess is that too many SMs will cause contention on NVLink bus.
+   */
+  template <typename T>
+  void allreduce(cudaStream_t stream, T* input, T* output, int size,
+                 int threads = 512, int block_limit = 36) {
+    // std::cout<<"input "<< input <<" output "<< output<<std::endl;
+    auto d = packed_t<T>::P::size;
+    if (size % d != 0)
+      throw std::runtime_error(
+          "custom allreduce currently requires input length to be multiple "
+          "of " +
+          std::to_string(d));
+    if (block_limit > kMaxBlocks)
+      throw std::runtime_error("max supported block limit is " +
+                               std::to_string(kMaxBlocks) + ". Got " +
+                               std::to_string(block_limit));
+
+    RankData* ptrs;
+    cudaStreamCaptureStatus status;
+    CUDACHECK(cudaStreamIsCapturing(stream, &status));
+    if (status == cudaStreamCaptureStatusActive) {
+      ptrs = d_rank_data_base_ + graph_unreg_buffers_.size();
+      graph_unreg_buffers_.push_back(input);
+    } else {
+      auto it = buffers_.find(input);
+      if (it == buffers_.end())
+        throw std::runtime_error(
+            "buffer address " +
+            std::to_string(reinterpret_cast<uint64_t>(input)) +
+            " is not registered!");
+      ptrs = it->second;
+    }
+
+    size /= d;
+    auto bytes = size * sizeof(typename packed_t<T>::P);
+    int blocks = std::min(block_limit, (size + threads - 1) / threads);
+#define KL(ngpus, name)                                                       \
+  name<T, ngpus><<<blocks, threads, 0, stream>>>(ptrs, sg_, self_sg_, output, \
+                                                 rank_, size);
+    // TODO(hanzhi713): Threshold is different for A100 and H100.
+    // Add per device threshold.
+#define REDUCE_CASE(ngpus)                            \
+  case ngpus: {                                       \
+    if (world_size_ == 2) {                           \
+      KL(ngpus, cross_device_reduce_1stage);          \
+    } else if (full_nvlink_) {                        \
+      if ((world_size_ <= 4 && bytes < 512 * 1024) || \
+          (world_size_ <= 8 && bytes < 256 * 1024)) { \
+        KL(ngpus, cross_device_reduce_1stage);        \
+      } else {                                        \
+        KL(ngpus, cross_device_reduce_2stage);        \
+      }                                               \
+    }                                                 \
+    break;                                            \
+  }
+
+    switch (world_size_) {
+      REDUCE_CASE(2)
+      REDUCE_CASE(4)
+      REDUCE_CASE(6)
+      REDUCE_CASE(8)
+      default:
+        throw std::runtime_error(
+            "custom allreduce only supports num gpus in (2,4,6,8). Actual num "
+            "gpus = " +
+            std::to_string(world_size_));
+    }
+#undef REDUCE_CASE
+#undef KL
+  }
+
+  ~CustomAllreduce() {
+    for (auto [_, ptr] : ipc_handles_) {
+      CUDACHECK(cudaIpcCloseMemHandle(ptr));
+    }
+  }
+};
+}  // namespace paddle

csrc/gpu/all_reduce.h

@@ -0,0 +1,30 @@
+
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "helper.h"
+
+using fptr_t = int64_t;
+fptr_t init_custom_all_reduce(const std::vector<int64_t>& fake_ipc_ptrs,
+                      paddle::Tensor& rank_data, int64_t rank, bool full_nvlink);
+void all_reduce(fptr_t _fa, paddle::Tensor& inp, paddle::Tensor& out,
+                fptr_t reg_buffer, int64_t reg_buffer_sz_bytes);
+void dispose(fptr_t _fa);
+int64_t meta_size();
+void register_buffer(fptr_t _fa, const std::vector<int64_t>& fake_ipc_ptrs);
+std::tuple<std::vector<int64_t>, std::vector<int64_t>>
+get_graph_buffer_ipc_meta(fptr_t _fa);
+void register_graph_buffers(fptr_t _fa,
+                            const std::vector<std::vector<int64_t>>& handles,
+                            const std::vector<std::vector<int64_t>>& offsets);

csrc/gpu/append_attention.cu

@@ -0,0 +1,980 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "append_attn/append_attention_kernel.h"
+#include "append_attn/decoder_write_cache_with_rope_kernel.h"
+#include "append_attn/speculate_write_cache_with_rope_kernel.h"
+#include "append_attn/encoder_write_cache_with_rope_kernel.h"
+
+template <paddle::DataType D>
+std::vector<paddle::Tensor> AppendAttentionKernel(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,
+    const paddle::Tensor& key_cache,
+    const paddle::Tensor& value_cache,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& encoder_batch_ids,
+    const paddle::Tensor& encoder_tile_ids_per_batch,
+    const paddle::Tensor& encoder_num_blocks,
+    const paddle::Tensor& kv_batch_ids,
+    const paddle::Tensor& kv_tile_ids_per_batch,
+    const paddle::Tensor& kv_num_blocks,
+    const paddle::Tensor& decoder_batch_ids,
+    const paddle::Tensor& decoder_tile_ids_per_batch,
+    const paddle::Tensor& decoder_num_blocks,
+    const paddle::Tensor& max_enc_len_this_time,
+    const paddle::Tensor& max_dec_len_this_time,
+    const paddle::Tensor& max_len_kv,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>& qkv_bias,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_quant_scales,
+    const paddle::optional<paddle::Tensor>& cache_v_quant_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_dequant_scales,
+    const paddle::optional<paddle::Tensor>& cache_v_dequant_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const paddle::optional<paddle::Tensor>& out_linear_shifts,
+    const paddle::optional<paddle::Tensor>& out_linear_smooths,
+    const paddle::optional<paddle::Tensor>& excess_blocks,
+    const std::string& cache_quant_type_str,
+    const bool use_neox_rotary_style,
+    const int max_input_length,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float out_linear_in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool speculate_decoder) {
+  typedef PDTraits<D> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
+
+  int encoder_num_blocks_data = encoder_num_blocks.data<int>()[0];
+  int kv_num_blocks_data = kv_num_blocks.data<int>()[0];
+  int decoder_num_blocks_data = decoder_num_blocks.data<int>()[0];
+  int max_enc_len_this_time_data = max_enc_len_this_time.data<int>()[0];
+  int max_dec_len_this_time_data = max_dec_len_this_time.data<int>()[0];
+  int max_len_kv_data = max_len_kv.data<int>()[0];
+  const int encoder_block_shape_q = get_encoder_block_shape_q();
+  const int decoder_block_shape_q = get_decoder_block_shape_q();
+  auto main_stream = qkv.stream();
+  static cudaEvent_t main_event;
+  static cudaEvent_t decoder_event;       
+  static cudaStream_t decoder_stream;
+  static bool init_flag = false;
+  if (max_enc_len_this_time_data > 0 && max_dec_len_this_time_data > 0 &&
+      !init_flag) {
+    cudaEventCreateWithFlags(&main_event, cudaEventDisableTiming);
+    cudaEventCreateWithFlags(&decoder_event, cudaEventDisableTiming);
+    cudaStreamCreateWithFlags(&decoder_stream, cudaStreamNonBlocking);
+    init_flag = true;
+  }
+
+  paddle::Tensor qkv_out;
+  if (qkv_out_scales) {
+    qkv_out = GetEmptyTensor(qkv.dims(), D, qkv.place());
+  } else {
+    qkv_out = qkv;
+  }
+  paddle::Tensor fmha_out;
+  if (out_linear_in_scale > 0.0) {
+    if (fabs(quant_max_bound - 127.0f) < 0.000001) {
+      fmha_out = GetEmptyTensor(
+        {meta_data.token_nums, meta_data.q_num_heads * meta_data.head_dims_v},
+        paddle::DataType::INT8,
+        qkv.place());
+    } 
+    else if (fabs(quant_max_bound - 448.0f) < 0.000001) {
+      fmha_out = GetEmptyTensor(
+        {meta_data.token_nums, meta_data.q_num_heads * meta_data.head_dims_v},
+        paddle::DataType::FLOAT8_E4M3FN,
+        qkv.place());
+    }else{
+      PD_THROW("Only supported attr of quant_max_bound in ['127.0', '448.0'].");
+    }
+  } else {
+    fmha_out = GetEmptyTensor(
+        {meta_data.token_nums, meta_data.q_num_heads * meta_data.head_dims_v},
+        D,
+        qkv.place());
+  }
+
+  if (max_enc_len_this_time_data > 0) {
+    if (max_dec_len_this_time_data > 0) {
+      cudaEventRecord(main_event, main_stream);
+    }
+    if (qkv_out_scales) {
+      EncoderWriteCacheWithRopeKernel<data_t, int>(
+          meta_data,
+          qkv,
+          seq_lens_this_time,
+          seq_lens_encoder,
+          seq_lens_decoder,
+          padding_offsets,
+          cum_offsets,
+          block_tables,
+          kv_batch_ids,
+          kv_tile_ids_per_batch,
+          rotary_embs,
+          qkv_out_scales,
+          qkv_bias,
+          cache_k_quant_scales,
+          cache_v_quant_scales,
+          cache_k_zp,
+          cache_v_zp,
+          excess_blocks,
+          cache_quant_type_str,
+          kv_num_blocks_data,
+          max_input_length,
+          use_neox_rotary_style,
+          main_stream,
+          &qkv_out,
+          const_cast<paddle::Tensor*>(&key_cache),
+          const_cast<paddle::Tensor*>(&value_cache));
+    } else {
+      EncoderWriteCacheWithRopeKernel<data_t, data_t>(
+          meta_data,
+          qkv_out,
+          seq_lens_this_time,
+          seq_lens_encoder,
+          seq_lens_decoder,
+          padding_offsets,
+          cum_offsets,
+          block_tables,
+          kv_batch_ids,
+          kv_tile_ids_per_batch,
+          rotary_embs,
+          qkv_out_scales,
+          qkv_bias,
+          cache_k_quant_scales,
+          cache_v_quant_scales,
+          cache_k_zp,
+          cache_v_zp,
+          excess_blocks,
+          cache_quant_type_str,
+          kv_num_blocks_data,
+          max_input_length,
+          use_neox_rotary_style,
+          main_stream,
+          &qkv_out,
+          const_cast<paddle::Tensor*>(&key_cache),
+          const_cast<paddle::Tensor*>(&value_cache));
+    }
+    if (out_linear_in_scale > 0.0) {
+      switch (fmha_out.dtype()) {
+        case paddle::DataType::INT8:{
+          CascadeAppendAttentionKernel<data_t, int8_t>(
+            meta_data,
+            qkv_out,
+            key_cache,
+            value_cache,
+            attn_mask,
+            cache_k_dequant_scales,
+            cache_v_dequant_scales,
+            cache_k_zp,
+            cache_v_zp,
+            out_linear_shifts,
+            out_linear_smooths,
+            seq_lens_this_time,
+            seq_lens_decoder,
+            seq_lens_encoder,
+            padding_offsets,
+            cum_offsets,
+            block_tables,
+            encoder_batch_ids,
+            encoder_tile_ids_per_batch,
+            cache_quant_type_str,
+            encoder_num_blocks_data,
+            encoder_block_shape_q,
+            max_input_length,
+            max_enc_len_this_time_data,
+            softmax_scale,
+            quant_max_bound,
+            quant_min_bound,
+            out_linear_in_scale,
+            speculate_max_draft_token_num,
+            causal,
+            false,
+            true,
+            main_stream,
+            &fmha_out);
+          break;
+        }
+        case paddle::DataType::FLOAT8_E4M3FN:{
+          CascadeAppendAttentionKernel<data_t, phi::dtype::float8_e4m3fn>(
+            meta_data,
+            qkv_out,
+            key_cache,
+            value_cache,
+            attn_mask,
+            cache_k_dequant_scales,
+            cache_v_dequant_scales,
+            cache_k_zp,
+            cache_v_zp,
+            out_linear_shifts,
+            out_linear_smooths,
+            seq_lens_this_time,
+            seq_lens_decoder,
+            seq_lens_encoder,
+            padding_offsets,
+            cum_offsets,
+            block_tables,
+            encoder_batch_ids,
+            encoder_tile_ids_per_batch,
+            cache_quant_type_str,
+            encoder_num_blocks_data,
+            encoder_block_shape_q,
+            max_input_length,
+            max_enc_len_this_time_data,
+            softmax_scale,
+            quant_max_bound,
+            quant_min_bound,
+            out_linear_in_scale,
+            speculate_max_draft_token_num,
+            causal,
+            false,
+            true,
+            main_stream,
+            &fmha_out);
+          break;
+        }
+        default:{
+          PD_THROW("Only supported output fmha_out of quant dtype in ['int8', 'fp8_e4m3'].");
+          break;
+        }
+      }
+    } else {
+      CascadeAppendAttentionKernel<data_t, data_t>(
+          meta_data,
+          qkv_out,
+          key_cache,
+          value_cache,
+          attn_mask,
+          cache_k_dequant_scales,
+          cache_v_dequant_scales,
+          cache_k_zp,
+          cache_v_zp,
+          out_linear_shifts,
+          out_linear_smooths,
+          seq_lens_this_time,
+          seq_lens_decoder,
+          seq_lens_encoder,
+          padding_offsets,
+          cum_offsets,
+          block_tables,
+          encoder_batch_ids,
+          encoder_tile_ids_per_batch,
+          cache_quant_type_str,
+          encoder_num_blocks_data,
+          encoder_block_shape_q,
+          max_input_length,
+          max_enc_len_this_time_data,
+          softmax_scale,
+          quant_max_bound,
+          quant_min_bound,
+          out_linear_in_scale,
+          speculate_max_draft_token_num,
+          causal,
+          false,
+          true,
+          main_stream,
+          &fmha_out);
+    }
+  }
+
+  if (max_dec_len_this_time_data > 0) {
+    cudaStream_t exec_stream;
+    if (max_enc_len_this_time_data > 0) {
+      cudaStreamWaitEvent(decoder_stream, main_event);
+      exec_stream = decoder_stream;
+    } else {
+      exec_stream = main_stream;
+    }
+    if (speculate_decoder) {
+      if (qkv_out_scales) {
+        SpeculateWriteCacheWithRoPEKernel<data_t, int>(
+            meta_data,
+            qkv,  // [token_num, num_heads, head_dim]
+            seq_lens_decoder,
+            seq_lens_encoder,
+            padding_offsets,
+            cum_offsets,
+            block_tables,
+            rotary_embs,
+            qkv_out_scales,
+            qkv_bias,
+            cache_k_quant_scales,
+            cache_v_quant_scales,
+            cache_k_zp,
+            cache_v_zp,
+            cache_quant_type_str,
+            use_neox_rotary_style,
+            max_input_length,
+            exec_stream,
+            &qkv_out,
+            const_cast<paddle::Tensor*>(&key_cache),
+            const_cast<paddle::Tensor*>(&value_cache));
+      } else {
+        SpeculateWriteCacheWithRoPEKernel<data_t, data_t>(
+            meta_data,
+            qkv_out,  // [token_num, num_heads, head_dim]
+            seq_lens_decoder,
+            seq_lens_encoder,
+            padding_offsets,
+            cum_offsets,
+            block_tables,
+            rotary_embs,
+            qkv_out_scales,
+            qkv_bias,
+            cache_k_quant_scales,
+            cache_v_quant_scales,
+            cache_k_zp,
+            cache_v_zp,
+            cache_quant_type_str,
+            use_neox_rotary_style,
+            max_input_length,
+            exec_stream,
+            &qkv_out,
+            const_cast<paddle::Tensor*>(&key_cache),
+            const_cast<paddle::Tensor*>(&value_cache));
+      }
+    } else {
+      if (qkv_out_scales) {
+        DecoderWriteCacheWithRoPEKernel<data_t, int>(
+            meta_data,
+            qkv,  // [token_num, num_heads, head_dim]
+            seq_lens_decoder,
+            seq_lens_encoder,
+            padding_offsets,
+            cum_offsets,
+            block_tables,
+            rotary_embs,
+            qkv_out_scales,
+            qkv_bias,
+            cache_k_quant_scales,
+            cache_v_quant_scales,
+            cache_k_zp,
+            cache_v_zp,
+            cache_quant_type_str,
+            use_neox_rotary_style,
+            max_input_length,
+            exec_stream,
+            &qkv_out,
+            const_cast<paddle::Tensor*>(&key_cache),
+            const_cast<paddle::Tensor*>(&value_cache));
+      } else {
+        DecoderWriteCacheWithRoPEKernel<data_t, data_t>(
+            meta_data,
+            qkv_out,  // [token_num, num_heads, head_dim]
+            seq_lens_decoder,
+            seq_lens_encoder,
+            padding_offsets,
+            cum_offsets,
+            block_tables,
+            rotary_embs,
+            qkv_out_scales,
+            qkv_bias,
+            cache_k_quant_scales,
+            cache_v_quant_scales,
+            cache_k_zp,
+            cache_v_zp,
+            cache_quant_type_str,
+            use_neox_rotary_style,
+            max_input_length,
+            exec_stream,
+            &qkv_out,
+            const_cast<paddle::Tensor*>(&key_cache),
+            const_cast<paddle::Tensor*>(&value_cache));
+      }
+    }
+
+    if (out_linear_in_scale > 0.0) {
+      switch (fmha_out.dtype()) {
+        case paddle::DataType::INT8:{
+          CascadeAppendAttentionKernel<data_t, int8_t>(
+            meta_data,
+            qkv_out,
+            key_cache,
+            value_cache,
+            attn_mask,
+            cache_k_dequant_scales,
+            cache_v_dequant_scales,
+            cache_k_zp,
+            cache_v_zp,
+            out_linear_shifts,
+            out_linear_smooths,
+            seq_lens_this_time,
+            seq_lens_decoder,
+            seq_lens_encoder,
+            padding_offsets,
+            cum_offsets,
+            block_tables,
+            decoder_batch_ids,
+            decoder_tile_ids_per_batch,
+            cache_quant_type_str,
+            decoder_num_blocks_data,
+            decoder_block_shape_q,
+            max_input_length,
+            max_len_kv_data,
+            softmax_scale,
+            quant_max_bound,
+            quant_min_bound,
+            out_linear_in_scale,
+            speculate_max_draft_token_num,
+            causal,
+            !speculate_decoder,
+            !speculate_decoder,
+            exec_stream,
+            &fmha_out);
+          break;
+        }
+        case paddle::DataType::FLOAT8_E4M3FN:{
+          CascadeAppendAttentionKernel<data_t, phi::dtype::float8_e4m3fn>(
+            meta_data,
+            qkv_out,
+            key_cache,
+            value_cache,
+            attn_mask,
+            cache_k_dequant_scales,
+            cache_v_dequant_scales,
+            cache_k_zp,
+            cache_v_zp,
+            out_linear_shifts,
+            out_linear_smooths,
+            seq_lens_this_time,
+            seq_lens_decoder,
+            seq_lens_encoder,
+            padding_offsets,
+            cum_offsets,
+            block_tables,
+            decoder_batch_ids,
+            decoder_tile_ids_per_batch,
+            cache_quant_type_str,
+            decoder_num_blocks_data,
+            decoder_block_shape_q,
+            max_input_length,
+            max_len_kv_data,
+            softmax_scale,
+            quant_max_bound,
+            quant_min_bound,
+            out_linear_in_scale,
+            speculate_max_draft_token_num,
+            causal,
+            !speculate_decoder,
+            !speculate_decoder,
+            exec_stream,
+            &fmha_out);
+          break;
+        }
+        default:{
+          PD_THROW("Only supported output fmha_out of quant dtype in ['int8', 'fp8_e4m3'].");
+          break;
+        }
+      }
+      
+    } else {
+      CascadeAppendAttentionKernel<data_t, data_t>(
+          meta_data,
+          qkv_out,
+          key_cache,
+          value_cache,
+          attn_mask,
+          cache_k_dequant_scales,
+          cache_v_dequant_scales,
+          cache_k_zp,
+          cache_v_zp,
+          out_linear_shifts,
+          out_linear_smooths,
+          seq_lens_this_time,
+          seq_lens_decoder,
+          seq_lens_encoder,
+          padding_offsets,
+          cum_offsets,
+          block_tables,
+          decoder_batch_ids,
+          decoder_tile_ids_per_batch,
+          cache_quant_type_str,
+          decoder_num_blocks_data,
+          decoder_block_shape_q,
+          max_input_length,
+          max_len_kv_data,
+          softmax_scale,
+          quant_max_bound,
+          quant_min_bound,
+          out_linear_in_scale,
+          speculate_max_draft_token_num,
+          causal,
+          !speculate_decoder,
+          !speculate_decoder,
+          exec_stream,
+          &fmha_out);
+    }
+    if (max_enc_len_this_time_data > 0) {
+      cudaEventRecord(decoder_event, exec_stream);
+      cudaStreamWaitEvent(main_stream, decoder_event);
+    }
+  }
+
+  return {fmha_out, qkv_out};
+}
+
+std::vector<paddle::Tensor> AppendAttention(
+    const paddle::Tensor& qkv,
+    const paddle::Tensor& key_cache,
+    const paddle::Tensor& value_cache,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& encoder_batch_ids,
+    const paddle::Tensor& encoder_tile_ids_per_batch,
+    const paddle::Tensor& encoder_num_blocks,
+    const paddle::Tensor& kv_batch_ids,
+    const paddle::Tensor& kv_tile_ids_per_batch,
+    const paddle::Tensor& kv_num_blocks,
+    const paddle::Tensor& decoder_batch_ids,
+    const paddle::Tensor& decoder_tile_ids_per_batch,
+    const paddle::Tensor& decoder_num_blocks,
+    const paddle::Tensor& max_enc_len_this_time,
+    const paddle::Tensor& max_dec_len_this_time,
+    const paddle::Tensor& max_len_kv,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>& qkv_bias,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_quant_scales,
+    const paddle::optional<paddle::Tensor>& cache_v_quant_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_dequant_scales,
+    const paddle::optional<paddle::Tensor>& cache_v_dequant_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const paddle::optional<paddle::Tensor>& out_linear_shifts,
+    const paddle::optional<paddle::Tensor>& out_linear_smooths,
+    const paddle::optional<paddle::Tensor>& excess_blocks,
+    const std::string& compute_dtype,
+    const std::string& cache_quant_type_str,
+    const bool use_neox_rotary_style,
+    const int max_input_length,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float out_linear_in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool speculate_decoder) {
+  AppendAttnMetaData meta_data;
+
+  const auto& qkv_dims = qkv.dims();
+  const auto& key_cache_dims = key_cache.dims();
+  meta_data.token_nums = qkv_dims[0];
+  meta_data.kv_num_heads = key_cache_dims[1];
+  meta_data.head_dims = key_cache_dims[3];
+  meta_data.head_dims_v = value_cache.dims()[3];
+  const int q_hidden_size =
+      qkv_dims[qkv_dims.size() - 1] - meta_data.kv_num_heads * (meta_data.head_dims + meta_data.head_dims_v);
+  meta_data.q_num_heads = q_hidden_size / meta_data.head_dims;
+
+  meta_data.max_blocks_per_seq = block_tables.dims()[1];
+  meta_data.block_size = key_cache.dims()[2];
+  meta_data.batch_size = cum_offsets.dims()[0];
+
+  switch (qkv.dtype()) {
+    case paddle::DataType::FLOAT16: {
+      return AppendAttentionKernel<paddle::DataType::FLOAT16>(
+          meta_data,
+          qkv,
+          key_cache,
+          value_cache,
+          seq_lens_encoder,
+          seq_lens_decoder,
+          seq_lens_this_time,
+          padding_offsets,
+          cum_offsets,
+          block_tables,
+          encoder_batch_ids,
+          encoder_tile_ids_per_batch,
+          encoder_num_blocks,
+          kv_batch_ids,
+          kv_tile_ids_per_batch,
+          kv_num_blocks,
+          decoder_batch_ids,
+          decoder_tile_ids_per_batch,
+          decoder_num_blocks,
+          max_enc_len_this_time,
+          max_dec_len_this_time,
+          max_len_kv,
+          rotary_embs,
+          attn_mask,
+          qkv_bias,
+          qkv_out_scales,
+          cache_k_quant_scales,
+          cache_v_quant_scales,
+          cache_k_dequant_scales,
+          cache_v_dequant_scales,
+          cache_k_zp,
+          cache_v_zp,
+          out_linear_shifts,
+          out_linear_smooths,
+          excess_blocks,
+          cache_quant_type_str,
+          use_neox_rotary_style,
+          max_input_length,
+          softmax_scale,
+          quant_max_bound,
+          quant_min_bound,
+          out_linear_in_scale,
+          speculate_max_draft_token_num,
+          causal,
+          speculate_decoder);
+    }
+    case paddle::DataType::BFLOAT16: {
+      return AppendAttentionKernel<paddle::DataType::BFLOAT16>(
+          meta_data,
+          qkv,
+          key_cache,
+          value_cache,
+          seq_lens_encoder,
+          seq_lens_decoder,
+          seq_lens_this_time,
+          padding_offsets,
+          cum_offsets,
+          block_tables,
+          encoder_batch_ids,
+          encoder_tile_ids_per_batch,
+          encoder_num_blocks,
+          kv_batch_ids,
+          kv_tile_ids_per_batch,
+          kv_num_blocks,
+          decoder_batch_ids,
+          decoder_tile_ids_per_batch,
+          decoder_num_blocks,
+          max_enc_len_this_time,
+          max_dec_len_this_time,
+          max_len_kv,
+          rotary_embs,
+          attn_mask,
+          qkv_bias,
+          qkv_out_scales,
+          cache_k_quant_scales,
+          cache_v_quant_scales,
+          cache_k_dequant_scales,
+          cache_v_dequant_scales,
+          cache_k_zp,
+          cache_v_zp,
+          out_linear_shifts,
+          out_linear_smooths,
+          excess_blocks,
+          cache_quant_type_str,
+          use_neox_rotary_style,
+          max_input_length,
+          softmax_scale,
+          quant_max_bound,
+          quant_min_bound,
+          out_linear_in_scale,
+          speculate_max_draft_token_num,
+          causal,
+          speculate_decoder);
+    }
+    case paddle::DataType::INT32: {
+      if (compute_dtype == "bf16") {
+        return AppendAttentionKernel<paddle::DataType::BFLOAT16>(
+            meta_data,
+            qkv,
+            key_cache,
+            value_cache,
+            seq_lens_encoder,
+            seq_lens_decoder,
+            seq_lens_this_time,
+            padding_offsets,
+            cum_offsets,
+            block_tables,
+            encoder_batch_ids,
+            encoder_tile_ids_per_batch,
+            encoder_num_blocks,
+            kv_batch_ids,
+            kv_tile_ids_per_batch,
+            kv_num_blocks,
+            decoder_batch_ids,
+            decoder_tile_ids_per_batch,
+            decoder_num_blocks,
+            max_enc_len_this_time,
+            max_dec_len_this_time,
+            max_len_kv,
+            rotary_embs,
+            attn_mask,
+            qkv_bias,
+            qkv_out_scales,
+            cache_k_quant_scales,
+            cache_v_quant_scales,
+            cache_k_dequant_scales,
+            cache_v_dequant_scales,
+            cache_k_zp,
+            cache_v_zp,
+            out_linear_shifts,
+            out_linear_smooths,
+            excess_blocks,
+            cache_quant_type_str,
+            use_neox_rotary_style,
+            max_input_length,
+            softmax_scale,
+            quant_max_bound,
+            quant_min_bound,
+            out_linear_in_scale,
+            speculate_max_draft_token_num,
+            causal,
+            speculate_decoder);
+      } else if (compute_dtype == "fp16") {
+        return AppendAttentionKernel<paddle::DataType::FLOAT16>(
+            meta_data,
+            qkv,
+            key_cache,
+            value_cache,
+            seq_lens_encoder,
+            seq_lens_decoder,
+            seq_lens_this_time,
+            padding_offsets,
+            cum_offsets,
+            block_tables,
+            encoder_batch_ids,
+            encoder_tile_ids_per_batch,
+            encoder_num_blocks,
+            kv_batch_ids,
+            kv_tile_ids_per_batch,
+            kv_num_blocks,
+            decoder_batch_ids,
+            decoder_tile_ids_per_batch,
+            decoder_num_blocks,
+            max_enc_len_this_time,
+            max_dec_len_this_time,
+            max_len_kv,
+            rotary_embs,
+            attn_mask,
+            qkv_bias,
+            qkv_out_scales,
+            cache_k_quant_scales,
+            cache_v_quant_scales,
+            cache_k_dequant_scales,
+            cache_v_dequant_scales,
+            cache_k_zp,
+            cache_v_zp,
+            out_linear_shifts,
+            out_linear_smooths,
+            excess_blocks,
+            cache_quant_type_str,
+            use_neox_rotary_style,
+            max_input_length,
+            softmax_scale,
+            quant_max_bound,
+            quant_min_bound,
+            out_linear_in_scale,
+            speculate_max_draft_token_num,
+            causal,
+            speculate_decoder);
+      } else {
+        PD_THROW("Only supported attr of compute_dtype in ['fp16', 'bf16'].");
+        break;
+      }
+    }
+    default: {
+      PD_THROW(
+          "NOT supported data type. "
+          "Only float16 and bfloat16 are supported. ");
+      break;
+    }
+  }
+  return {paddle::Tensor{}};
+}
+
+std::vector<std::vector<int64_t>> AppendAttentionInferShape(
+    const std::vector<int64_t>& qkv_shape,
+    const std::vector<int64_t>& key_cache_shape,
+    const std::vector<int64_t>& value_cache_shape,
+    const std::vector<int64_t>& seq_lens_encoder_shape,
+    const std::vector<int64_t>& seq_lens_decoder_shape,
+    const std::vector<int64_t>& seq_lens_this_time_shape,
+    const std::vector<int64_t>& padding_offsets_shape,
+    const std::vector<int64_t>& cum_offsets_shape,
+    const std::vector<int64_t>& block_tables_shape,
+    const std::vector<int64_t>& encoder_batch_ids_shape,
+    const std::vector<int64_t>& encoder_tile_ids_per_batch_shape,
+    const std::vector<int64_t>& encoder_num_blocks_shape,
+    const std::vector<int64_t>& kv_batch_ids_shape,
+    const std::vector<int64_t>& kv_tile_ids_per_batch_shape,
+    const std::vector<int64_t>& kv_num_blocks_shape,
+    const std::vector<int64_t>& decoder_batch_ids_shape,
+    const std::vector<int64_t>& decoder_tile_ids_per_batch_shape,
+    const std::vector<int64_t>& decoder_num_blocks_shape,
+    const std::vector<int64_t>& max_enc_len_this_time_shape,
+    const std::vector<int64_t>& max_dec_len_this_time_shape,
+    const std::vector<int64_t>& max_len_kv_shape,
+    const paddle::optional<std::vector<int64_t>>& rotary_embs_shape,
+    const paddle::optional<std::vector<int64_t>>& attn_mask_shape,
+    const paddle::optional<std::vector<int64_t>>& qkv_bias_shape,
+    const paddle::optional<std::vector<int64_t>>& qkv_out_scales_shape,
+    const paddle::optional<std::vector<int64_t>>& cache_k_quant_scales_shape,
+    const paddle::optional<std::vector<int64_t>>& cache_v_quant_scales_shape,
+    const paddle::optional<std::vector<int64_t>>& cache_k_dequant_scales_shape,
+    const paddle::optional<std::vector<int64_t>>& cache_v_dequant_scales_shape,
+    const paddle::optional<std::vector<int64_t>>& cache_k_zp_shape,
+    const paddle::optional<std::vector<int64_t>>& cache_v_zp_shape,
+    const paddle::optional<std::vector<int64_t>>& out_linear_shifts_shape,
+    const paddle::optional<std::vector<int64_t>>& out_linear_smooths_shape,
+    const paddle::optional<std::vector<int64_t>>& excess_blocks_shape) {
+  const int token_num = qkv_shape[0];
+  const int kv_num_heads = key_cache_shape[1];
+  const int head_dim_qk = key_cache_shape[3];
+  const int head_dim_v = value_cache_shape[3];
+  const int q_hidden_size =
+      qkv_shape[qkv_shape.size() - 1] - kv_num_heads * (head_dim_qk + head_dim_v);
+  const int num_heads = q_hidden_size / head_dim_qk;
+  return {{token_num, num_heads * head_dim_v}, qkv_shape};
+}
+
+std::vector<paddle::DataType> AppendAttentionInferDtype(
+    const paddle::DataType& qkv_dtype,
+    const paddle::DataType& key_cache_dtype,
+    const paddle::DataType& value_cache_dtype,
+    const paddle::DataType& seq_lens_encoder_dtype,
+    const paddle::DataType& seq_lens_decoder_dtype,
+    const paddle::DataType& seq_lens_this_time_dtype,
+    const paddle::DataType& padding_offsets_dtype,
+    const paddle::DataType& cum_offsets_dtype,
+    const paddle::DataType& block_tables_dtype,
+    const paddle::DataType& encoder_batch_ids_dtype,
+    const paddle::DataType& encoder_tile_ids_per_batch_dtype,
+    const paddle::DataType& encoder_num_blocks_dtype,
+    const paddle::DataType& kv_batch_ids_dtype,
+    const paddle::DataType& kv_tile_ids_per_batch_dtype,
+    const paddle::DataType& kv_num_blocks_dtype,
+    const paddle::DataType& decoder_batch_ids_dtype,
+    const paddle::DataType& decoder_tile_ids_per_batch_dtype,
+    const paddle::DataType& decoder_num_blocks_dtype,
+    const paddle::DataType& max_enc_len_this_time_dtype,
+    const paddle::DataType& max_dec_len_this_time_dtype,
+    const paddle::DataType& max_len_kv_dtype,
+    const paddle::optional<paddle::DataType>& rotary_embs_dtype,
+    const paddle::optional<paddle::DataType>& attn_mask_dtype,
+    const paddle::optional<paddle::DataType>& qkv_bias_dtype,
+    const paddle::optional<paddle::DataType>& qkv_out_scales_dtype,
+    const paddle::optional<paddle::DataType>& cache_k_quant_scales_dtype,
+    const paddle::optional<paddle::DataType>& cache_v_quant_scales_dtype,
+    const paddle::optional<paddle::DataType>& cache_k_dequant_scales_dtype,
+    const paddle::optional<paddle::DataType>& cache_v_dequant_scales_dtype,
+    const paddle::optional<paddle::DataType>& cache_k_zp_dtype,
+    const paddle::optional<paddle::DataType>& cache_v_zp_dtype,
+    const paddle::optional<paddle::DataType>& out_linear_shifts_dtype,
+    const paddle::optional<paddle::DataType>& out_linear_smooths_dtype,
+    const paddle::optional<paddle::DataType>& excess_blocks_dtype,
+    const std::string& compute_dtype,
+    const std::string& cache_quant_type_str,
+    const bool use_neox_rotary_style,
+    const int max_input_length,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float out_linear_in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool speculate_decoder) {
+  if (compute_dtype == "bf16") {
+    if (out_linear_in_scale > 0.0) {
+      if (fabs(quant_max_bound - 127.0f) < 0.000001) {
+        return {paddle::DataType::INT8, paddle::DataType::BFLOAT16};
+      } else if (fabs(quant_max_bound - 448.0f) < 0.000001) {
+        return {paddle::DataType::FLOAT8_E4M3FN, paddle::DataType::BFLOAT16};
+      }else{
+        PD_THROW("Only supported attr of quant_max_bound in ['127.0', '448.0'].");
+      }
+    } else {
+      return {paddle::DataType::BFLOAT16, paddle::DataType::BFLOAT16};
+    }
+  } else if (compute_dtype == "fp16") {
+    if (out_linear_in_scale > 0.0) {
+      if (fabs(quant_max_bound - 127.0f) < 0.000001) {
+        return {paddle::DataType::INT8, paddle::DataType::FLOAT16};
+      } else if (fabs(quant_max_bound - 448.0f) < 0.000001) {
+        return {paddle::DataType::FLOAT8_E4M3FN, paddle::DataType::FLOAT16};
+      }else{
+        PD_THROW("Only supported attr of quant_max_bound in ['127.0', '448.0'].");
+      }
+    } else {
+      return {paddle::DataType::FLOAT16, paddle::DataType::FLOAT16};
+    }
+  } else {
+    PD_THROW("Only supported attr of compute_dtype in ['fp16', 'bf16'].");
+  }
+}
+
+PD_BUILD_OP(append_attention)
+    .Inputs({"qkv",
+             "key_cache",
+             "value_cache",
+             "seq_lens_encoder",
+             "seq_lens_decoder",
+             "seq_lens_this_time",
+             "padding_offsets",
+             "cum_offsets",
+             "block_tables",
+             "encoder_batch_ids",
+             "encoder_tile_ids_per_batch",
+             "encoder_num_blocks",
+             "kv_batch_ids",
+             "kv_tile_ids_per_batch",
+             "kv_num_blocks",
+             "decoder_batch_ids",
+             "decoder_tile_ids_per_batch",
+             "decoder_num_blocks",
+             "max_enc_len_this_time",
+             "max_dec_len_this_time",
+             "max_len_kv",
+             paddle::Optional("rotary_embs"),
+             paddle::Optional("attn_mask"),
+             paddle::Optional("qkv_bias"),
+             paddle::Optional("qkv_out_scales"),
+             paddle::Optional("cache_k_quant_scales"),
+             paddle::Optional("cache_v_quant_scales"),
+             paddle::Optional("cache_k_dequant_scales"),
+             paddle::Optional("cache_v_dequant_scales"),
+             paddle::Optional("cache_k_zp"),
+             paddle::Optional("cache_v_zp"),
+             paddle::Optional("out_linear_shifts"),
+             paddle::Optional("out_linear_smooths"),
+             paddle::Optional("excess_blocks")})
+    .Outputs({"fmha_out", "qkv_out", "key_cache_out", "value_cache_out"})
+    .SetInplaceMap({{"key_cache", "key_cache_out"},
+                    {"value_cache", "value_cache_out"}})
+    .Attrs({"compute_type: std::string",
+            "cache_quant_type: std::string",
+            "use_neox_rotary_style: bool",
+            "max_input_length: int",
+            "softmax_scale: float",
+            "quant_max_bound: float",
+            "quant_min_bound: float",
+            "out_linear_in_scale: float",
+            "speculate_max_draft_token_num: int",
+            "causal: bool",
+            "speculate_decoder: bool"})
+    .SetKernelFn(PD_KERNEL(AppendAttention))
+    .SetInferShapeFn(PD_INFER_SHAPE(AppendAttentionInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(AppendAttentionInferDtype));

csrc/gpu/append_attn/append_attention_kernel.h

@@ -0,0 +1,302 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#pragma once
+
+#include "helper.h"
+#include "utils.cuh"
+
+template <typename T, typename OutT>
+void CascadeAppendAttentionC16Kernel(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);
+
+template <typename T, typename OutT>
+void CascadeAppendAttentionC8Kernel(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);
+
+template <typename T, typename OutT>
+void CascadeAppendAttentionC4Kernel(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);
+
+template <typename T, typename OutT>
+void CascadeAppendAttentionKernel(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const std::string& cache_quant_type_str,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out) {
+  if (cache_quant_type_str == "none") {
+    CascadeAppendAttentionC16Kernel<T, OutT>(meta_data,
+                                             qkv,
+                                             cache_k,
+                                             cache_v,
+                                             attn_mask,
+                                             cache_k_scale,
+                                             cache_v_scale,
+                                             cache_k_zp,
+                                             cache_v_zp,
+                                             shift_bias,
+                                             smooth_weight,
+                                             seq_lens_q,
+                                             seq_lens_kv,
+                                             seq_lens_encoder,
+                                             padding_offsets,
+                                             cum_offsets,
+                                             block_table,
+                                             batch_ids,
+                                             tile_ids_per_batch,
+                                             num_blocks,
+                                             block_shape_q,
+                                             max_seq_len,
+                                             max_dec_len,
+                                             softmax_scale,
+                                             quant_max_bound,
+                                             quant_min_bound,
+                                             in_scale,
+                                             speculate_max_draft_token_num,
+                                             causal,
+                                             is_decoder,
+                                             enable_prefill,
+                                             stream,
+                                             out);
+  } else if (cache_quant_type_str == "cache_int8") {
+    CascadeAppendAttentionC8Kernel<T, OutT>(meta_data,
+                                            qkv,
+                                            cache_k,
+                                            cache_v,
+                                            attn_mask,
+                                            cache_k_scale,
+                                            cache_v_scale,
+                                            cache_k_zp,
+                                            cache_v_zp,
+                                            shift_bias,
+                                            smooth_weight,
+                                            seq_lens_q,
+                                            seq_lens_kv,
+                                            seq_lens_encoder,
+                                            padding_offsets,
+                                            cum_offsets,
+                                            block_table,
+                                            batch_ids,
+                                            tile_ids_per_batch,
+                                            num_blocks,
+                                            block_shape_q,
+                                            max_seq_len,
+                                            max_dec_len,
+                                            softmax_scale,
+                                            quant_max_bound,
+                                            quant_min_bound,
+                                            in_scale,
+                                            speculate_max_draft_token_num,
+                                            causal,
+                                            is_decoder,
+                                            enable_prefill,
+                                            stream,
+                                            out);
+  } else if (cache_quant_type_str == "cache_int4_zp") {
+    CascadeAppendAttentionC4Kernel<T, OutT>(meta_data,
+                                            qkv,
+                                            cache_k,
+                                            cache_v,
+                                            attn_mask,
+                                            cache_k_scale,
+                                            cache_v_scale,
+                                            cache_k_zp,
+                                            cache_v_zp,
+                                            shift_bias,
+                                            smooth_weight,
+                                            seq_lens_q,
+                                            seq_lens_kv,
+                                            seq_lens_encoder,
+                                            padding_offsets,
+                                            cum_offsets,
+                                            block_table,
+                                            batch_ids,
+                                            tile_ids_per_batch,
+                                            num_blocks,
+                                            block_shape_q,
+                                            max_seq_len,
+                                            max_dec_len,
+                                            softmax_scale,
+                                            quant_max_bound,
+                                            quant_min_bound,
+                                            in_scale,
+                                            speculate_max_draft_token_num,
+                                            causal,
+                                            is_decoder,
+                                            enable_prefill,
+                                            stream,
+                                            out);
+  } else {
+    PD_THROW(
+        "cache_quant_type_str should be one of [none, cache_int8, "
+        "cache_int4_zp]");
+  }
+}

csrc/gpu/append_attn/decode_attention_func.cuh

@@ -0,0 +1,237 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#pragma once
+
+
+#include "multi_head_latent_attention_kernel.h"
+
+template <size_t vec_size, typename T>
+struct softmax_state_t {
+  AlignedVector<T, vec_size> o;
+  T m;
+  T d;
+  
+  __device__ __forceinline__ void init() {
+    if constexpr (std::is_same<T, half>::value) {
+#pragma unroll
+      for (int i = 0; i < vec_size / 2; ++i) {
+        *((half2*)(&o) + i) = make_half2(0, 0);
+      }
+    } else if constexpr (std::is_same<T, __nv_bfloat16>::value) {
+#pragma unroll
+      for (int i = 0; i < vec_size / 2; ++i) {
+        *((nv_bfloat162*)(&o) + i) = make_bfloat162(0, 0);
+      }
+    }
+    d = 1.f;
+    if constexpr (std::is_same<T, half>::value) {
+      m = __float2half(-5e4f);
+    } else if constexpr (std::is_same<T, nv_bfloat16>::value) {
+      m = __float2bfloat16(-3.38953e38f);
+    }
+  }
+
+  __device__ __forceinline__ softmax_state_t() {
+    init();
+  }
+
+  __device__ __forceinline__ void merge(const AlignedVector<T, vec_size>& other_o, 
+                                        T other_m,
+                                        T other_d) {
+    // using kType = typename cascade_attn_nv_type2_traits<T>::type;
+    T m_prev = m, d_prev = d;
+    m = m_prev > other_m ? m_prev : other_m;
+    T scale1 = hexp(m_prev - m), scale2 = hexp(other_m - m);
+
+    d = d_prev * scale1 + other_d * scale2;
+
+#pragma unroll
+    for (size_t i = 0; i < vec_size; ++i) {
+      o[i] = o[i] * scale1 + other_o[i] * scale2;
+    }
+  }
+
+  __device__ __forceinline__ void normalize() {
+
+#pragma unroll
+    for (size_t i = 0; i < vec_size; ++i) {
+      o[i] /= d;
+    }
+  }
+
+};
+
+template <size_t vec_size, typename T, uint32_t num_tiles = 0>
+struct softmax_state_ts {
+  uint32_t num_tiles_ = num_tiles;
+  AlignedVector<T, vec_size> o[num_tiles];
+  float m;
+  float d;
+  
+  __device__ __forceinline__ void init() {
+#pragma unroll
+    for (uint32_t tile_id = 0; tile_id < num_tiles_; ++tile_id) {
+      if constexpr (std::is_same<T, half>::value) {
+#pragma unroll
+        for (int i = 0; i < vec_size / 2; ++i) {
+          *((half2*)(&o[tile_id]) + i) = make_half2(0, 0);
+        }
+      } else if constexpr (std::is_same<T, __nv_bfloat16>::value) {
+#pragma unroll
+        for (int i = 0; i < vec_size / 2; ++i) {
+          *((nv_bfloat162*)(&o[tile_id]) + i) = make_bfloat162(0, 0);
+        }
+      }
+    }
+    d = 1.f;
+    if constexpr (std::is_same<T, half>::value) {
+      m = -5e4f;
+    } else if constexpr (std::is_same<T, nv_bfloat16>::value) {
+      m = -3.38953e38f;
+    }
+  }
+
+  __device__ __forceinline__ softmax_state_ts() {
+    init();
+  }
+
+  __device__ __forceinline__ void normalize(const uint32_t tile_id) {
+
+#pragma unroll
+    for (size_t i = 0; i < vec_size; i++) {
+      o[tile_id][i] /= d;
+    }
+  }
+
+};
+
+template <SharedMemFillMode fill_mode, uint32_t HEAD_DIM_QK, uint32_t vec_size, uint32_t NUM_VEC_PER_HEAD, uint32_t bdx, uint32_t BLOCK_SIZE, uint32_t CACHE_VEC_SIZE, typename CacheT>
+__device__ __forceinline__ void produce_kv(CacheT *smem,
+                                          CacheT *kv_base_gptr,
+                                          const int * block_table_smem,
+                                          const uint32_t seq_offset_gmem,
+                                          const uint32_t seq_offset_smem,
+                                          const uint32_t kv_head_idx,
+                                          const uint32_t kv_num_heads,
+                                          const uint32_t tidx,
+                                          const uint32_t chunk_start,
+                                          const uint32_t chunk_end) {
+  int block_id = __ldg(&block_table_smem[seq_offset_gmem / BLOCK_SIZE]);
+  if (block_id < 0) {
+    block_id = 0;
+  }
+  const uint32_t block_offset = seq_offset_gmem % BLOCK_SIZE;
+  // 8/16 T/int8 each time
+  const uint32_t k_offset_base = ((block_id * kv_num_heads + kv_head_idx) * BLOCK_SIZE + block_offset) * HEAD_DIM_QK;
+  const uint32_t smem_offset_base = seq_offset_smem * HEAD_DIM_QK;
+  for(uint32_t vid = tidx; vid < NUM_VEC_PER_HEAD; vid += bdx) {
+    pred_load<128, PrefetchMode::kPrefetch, fill_mode, CacheT>(
+      smem + smem_offset_base + vid * CACHE_VEC_SIZE,
+      kv_base_gptr + k_offset_base + vid * CACHE_VEC_SIZE,
+      seq_offset_gmem < chunk_end
+    );
+  }
+}
+
+template <uint32_t vec_size, uint32_t NUM_VEC_PER_HEAD, uint32_t bdx, uint32_t bdy, uint32_t HEAD_DIM, uint32_t DEAL_EACH_TIME, uint32_t num_tile_v, typename T, typename CacheT>
+__device__ __forceinline__ void compute_qk(const T* cu_q_smem,
+                                           const CacheT* k_smem,
+                                           const uint32_t kv_idx_base,
+                                           const uint32_t stage_idx,
+                                           const uint32_t iter_base, 
+                                           const uint32_t iter_bound,
+                                           const uint32_t tidx,
+                                           const uint32_t gid,
+                                           const float scale,
+                                           float *s,
+                                           softmax_state_ts<vec_size, T, num_tile_v>& st) {
+  const CacheT* smem;
+  AlignedVector<T, vec_size> q_vec;
+  AlignedVector<T, vec_size> k_vec;
+  float m_prev = st.m;
+  // smem = base_smem + (stage_idx * DEAL_EACH_TIME + zid * tile_size) * HEAD_DIM;
+  smem = k_smem + stage_idx * DEAL_EACH_TIME * HEAD_DIM;
+#pragma unroll
+  for (uint32_t j = 0; j < DEAL_EACH_TIME; ++j) {
+    if (iter_base + j < iter_bound) {
+      if constexpr (std::is_same<T, half>::value) {
+        s[j] = 0.f;
+      } else if constexpr (std::is_same<T, __nv_bfloat16>::value) {
+        s[j] = 0.f;
+      }
+#pragma unroll
+      for(uint32_t vid = tidx; vid < NUM_VEC_PER_HEAD; vid += bdx) {
+        Load<T, vec_size>(cu_q_smem + vid * vec_size, &q_vec);
+        Load<CacheT, vec_size>(smem + j * HEAD_DIM + vid * vec_size, &k_vec);
+        for (uint32_t i = 0; i < vec_size; ++i) {
+          s[j] += static_cast<float>(q_vec[i] * k_vec[i]);
+        }
+      }
+#pragma unroll
+      for (uint32_t offset = bdx / 2; offset > 0; offset /= 2) {
+        s[j] += __shfl_xor_sync(-1, s[j], offset, 32);
+      }    
+      __syncthreads();
+    } else {
+      if constexpr (std::is_same<T, half>::value) {
+        s[j] = -5e4f;
+      } else if constexpr (std::is_same<T, __nv_bfloat16>::value) {
+        s[j] = -3.38953e38f;
+      }
+    }
+    st.m = st.m > s[j] ? st.m : s[j];
+  }
+
+  // T o_scale = hexp(m_prev - st.m);
+  float o_scale = __expf(m_prev - st.m);
+  st.d *= o_scale;
+  
+#pragma unroll
+  for (uint32_t j = 0; j < DEAL_EACH_TIME; ++j) {
+    // s[j] = hexp(s[j] - st.m);
+    s[j] = __expf(s[j] - st.m);
+    st.d += s[j];
+  }
+#pragma unroll
+  for (uint32_t tile_id = 0; tile_id < num_tile_v; ++tile_id) {
+    for (uint32_t i = 0; i < vec_size; ++i) {
+      st.o[tile_id][i] *= o_scale;
+    }
+  }
+}
+
+template<uint32_t vec_size, uint32_t NUM_VEC_PER_HEAD, uint32_t bdx, uint32_t DEAL_EACH_TIME, uint32_t HEAD_DIM_QK, uint32_t num_tile, typename T, typename CacheT>
+__device__ __forceinline__ void compute_sv(const float *s,
+                                           const CacheT *base_v_smem,
+                                           const uint32_t stage_idx,
+                                           const uint32_t iter_base, 
+                                           const uint32_t iter_bound,
+                                           const uint32_t tidx,
+                                           softmax_state_ts<vec_size, T, num_tile>& st) {
+  const CacheT* v_smem;
+  AlignedVector<T, vec_size> v_vec;
+#pragma unroll
+  for (int j = 0; (j < DEAL_EACH_TIME) && (iter_base + j < iter_bound); ++j) {
+    v_smem = base_v_smem + stage_idx * DEAL_EACH_TIME * HEAD_DIM_QK + j * HEAD_DIM_QK;
+    for(uint32_t vid = tidx; vid < NUM_VEC_PER_HEAD; vid += bdx) {
+      Load<T, vec_size>(v_smem + vid * vec_size, &v_vec);
+      uint32_t tile_id = vid / bdx;
+#pragma unroll
+      for (int reg_id = 0; reg_id < vec_size; ++reg_id) {
+        st.o[tile_id][reg_id] += static_cast<T>(s[j]) * v_vec[reg_id];
+      }
+    }
+  }
+}
+

csrc/gpu/append_attn/decode_attention_kernel.cu

@@ -0,0 +1,562 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+// 
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// 
+//     http://www.apache.org/licenses/LICENSE-2.0
+// 
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "decode_attention_func.cuh"
+
+#define CHECK(call)                                                             \
+do                                                                              \
+{                                                                               \
+    const cudaError_t error_code = call;                                        \
+    if (error_code != cudaSuccess)                                              \
+    {                                                                           \
+        printf("CUDA Error:\n");                                                \
+        printf("     File:      %s\n", __FILE__);                               \
+        printf("     Line       %d:\n", __LINE__);                              \
+        printf("     Error code:%d\n", error_code);                             \
+        printf("     Error text:%s\n", cudaGetErrorString(error_code));         \
+        exit(1);                                                                \
+    }                                                                           \
+}while(0)
+
+template <typename T, typename OutT, int vec_size, uint32_t bdy, uint32_t HEAD_DIM>
+__global__ void merge_varlen_multi_chunks_v2_kernel(const T * __restrict__ multi_out, // [bsz, num_chunks, num_heads, head_dim]
+                                                    const T * __restrict__ multi_m, // [bsz, num_chunks, num_heads]
+                                                    const T * __restrict__ multi_d, // [bsz, num_chunks, num_heads]
+                                                    const int * __restrict__ seq_lens_q,
+                                                    const int * __restrict__ seq_lens_kv,
+                                                    const int * __restrict__ cum_offsets,
+                                                    const T * __restrict__ shift_bias, // [q_num_heads * HEAD_DIM]
+                                                    const T * __restrict__ smooth_weight, // [q_num_heads * HEAD_DIM]
+                                                    OutT * __restrict__ out, // [token_num, num_heads, head_dim]
+                                                    const float in_scale,
+                                                    const int num_chunks,
+                                                    const int chunk_size,
+                                                    const int max_seq_len,
+                                                    const int num_heads,
+                                                    const int head_dim) {
+  const int vid = threadIdx.x, ty = threadIdx.y;
+  const int qid = blockIdx.x, hid = blockIdx.y;
+  const int seq_len_q = seq_lens_q[qid];
+  if (seq_len_q == 0) return;
+  int seq_len_kv = seq_lens_kv[qid];
+  if (seq_len_kv == 0) return;
+  seq_len_kv += seq_len_q;
+  const int num_chunks_this_seq = div_up(seq_len_kv, chunk_size);
+  if (num_chunks_this_seq == 1 || ty >= num_chunks_this_seq) {
+    return;
+  }
+  __shared__ T smem[bdy * HEAD_DIM];
+  __shared__ T md_smem[bdy * 2];
+
+  const int start_token_ids = qid * max_seq_len - __ldg(&cum_offsets[qid]);
+  using LoadT = AlignedVector<T, vec_size>;
+  LoadT load_vec;
+  LoadT res_vec;
+  if constexpr (std::is_same<T, half>::value) {
+#pragma unroll
+    for (int i = 0; i < vec_size / 2; ++i) {
+      *((half2*)(&res_vec) + i) = make_half2(0, 0);
+    }
+  } else if constexpr (std::is_same<T, nv_bfloat16>::value) {
+#pragma unroll
+    for (int i = 0; i < vec_size / 2; ++i) {
+      *((nv_bfloat162*)(&res_vec) + i) = make_bfloat162(0, 0);
+    }
+  }
+  T m;
+  T d = 1.f;
+  if constexpr (std::is_same<T, half>::value) {
+    m = __float2half(-5e4f);
+  } else if constexpr (std::is_same<T, nv_bfloat16>::value) {
+    m = __float2bfloat16(-3.38953e38f);
+  }
+  // merge per ty
+#pragma unroll 2
+  for (int i = ty; i < num_chunks_this_seq; i += bdy) {
+    uint32_t offset = (qid * num_chunks + i) * num_heads + hid;
+    T m_prev = m;
+    T d_prev = d;
+    const T m_now = multi_m[offset];
+    const T d_now = multi_d[offset];
+    m = m_prev > m_now ? m_prev : m_now;
+    offset = (qid * num_chunks * num_heads + i * num_heads + hid) * head_dim + vid * vec_size;
+    Load<T, vec_size>(&multi_out[offset], &load_vec);
+    const T scale1 = hexp(m_prev - m), scale2 = hexp(m_now - m);
+    d = d * scale1 + d_now * scale2;
+#pragma once
+    for (int j = 0; j < vec_size; j++) {
+      res_vec[j] = res_vec[j] * scale1 + load_vec[j] * scale2;
+    }
+  }
+  // store ty res
+  Store<T, vec_size>(res_vec, &smem[ty * head_dim + vid * vec_size]);
+  md_smem[2 * ty] = m;
+  md_smem[2 * ty + 1] = d;
+  __syncthreads();
+
+  // merge bdy
+  softmax_state_t<vec_size, T> st{};
+  const uint32_t iter_num = min(num_chunks_this_seq, bdy);
+#pragma once
+  for (int i = 0; i < iter_num; i++) {
+    Load<T, vec_size>(&smem[i * head_dim + vid * vec_size], &load_vec);
+    const T m_tmp = md_smem[2 * i], d_tmp = md_smem[2 * i + 1];
+    st.merge(load_vec, m_tmp, d_tmp);
+  }
+  st.normalize();
+
+  AlignedVector<OutT, vec_size> out_vec;
+
+#pragma unroll
+  for (int i = 0; i < vec_size; ++i) {
+    out_vec[i] = static_cast<OutT>(st.o[i]);
+  }
+  Store<OutT, vec_size>(out_vec, &out[(start_token_ids * num_heads + hid) * head_dim + vid * vec_size]);
+}
+
+template <bool partition_kv, typename T, typename OutT, typename CacheT, uint32_t NUM_STAGES, uint32_t DEAL_EACH_TIME, uint32_t GROUP_SIZE, uint32_t HEAD_DIM_QK, uint32_t HEAD_DIM_V, 
+          uint32_t BLOCK_SIZE, uint32_t VEC_SIZE, uint32_t CACHE_VEC_SIZE, uint32_t bdx, uint32_t bdy>
+__global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [token_num, num_heads, head_dim]
+                                                    CacheT * __restrict__ cache_k, // [max_block_num, num_heads, block_size, head_dim]
+                                                    CacheT * __restrict__ cache_v,
+                                                    const T * __restrict__ shift_bias, // [q_num_heads * HEAD_DIM]
+                                                    const T * __restrict__ smooth_weight, // [q_num_heads * HEAD_DIM]
+                                                    const int * __restrict__ seq_lens_q,
+                                                    const int * __restrict__ seq_lens_kv,
+                                                    const int * __restrict__ cum_offsets,
+                                                    const int * __restrict__ block_table, // [bsz, block_num_per_seq]
+                                                    const int max_seq_len,
+                                                    const int max_dec_len,
+                                                    const int max_block_num_per_seq,
+                                                    const float scale,
+                                                    const float in_scale,
+                                                    const uint32_t chunk_size,
+                                                    T * __restrict__ tmp_workspace, // [batch_size, num_chunks, num_heads, head_dim]
+                                                    T * __restrict__ tmp_m, // [batch_size, num_chunks, num_heads]
+                                                    T * __restrict__ tmp_d, // [batch_size, num_chunks, num_heads]
+                                                    OutT * __restrict__ out) {
+  constexpr uint32_t q_block_num = GROUP_SIZE / bdy;
+  const uint32_t bidx = blockIdx.x, kv_head_idx = blockIdx.z;
+  const uint32_t bid = bidx / q_block_num, q_block_idx = bidx % q_block_num;
+  const uint32_t bidy = threadIdx.y;
+  const uint32_t gid = q_block_idx * bdy + threadIdx.y;
+  const uint32_t tidx = threadIdx.x;
+  constexpr uint32_t num_vec_per_head_qk = HEAD_DIM_QK / VEC_SIZE;
+  constexpr uint32_t num_vec_per_head_v = HEAD_DIM_V / VEC_SIZE;
+  constexpr uint32_t num_tile_v = (num_vec_per_head_v + bdx - 1) / bdx;
+
+  const uint32_t q_head_idx = kv_head_idx * GROUP_SIZE + gid;
+  const uint32_t kv_num_heads = gridDim.z;
+  const uint32_t q_num_heads = kv_num_heads * GROUP_SIZE;
+  
+  const int *block_table_now = block_table + bid * max_block_num_per_seq;
+  
+  const uint32_t num_chunks = gridDim.y;
+  const uint32_t chunk_id = blockIdx.y;
+  const uint32_t q_len = seq_lens_q[bid];
+  if (q_len <= 0) {
+    return;
+  }
+  uint32_t kv_len = seq_lens_kv[bid]; // !!!!!!!!
+  if (kv_len <= 0) {
+     return;
+  }
+  kv_len += q_len;
+  const uint32_t num_chunk_this_seq = div_up(kv_len, chunk_size);
+  const uint32_t q_start_idx = bid * max_seq_len - __ldg(&cum_offsets[bid]);
+  const uint32_t q_write_idx = bid * max_seq_len - __ldg(&cum_offsets[bid]);
+  if (chunk_id >= num_chunk_this_seq) {
+    return;
+  }
+
+  const uint32_t chunk_start = partition_kv ? chunk_id * chunk_size : 0;
+  const uint32_t chunk_end = partition_kv ? min(kv_len, chunk_start + chunk_size) : kv_len;
+  const uint32_t chunk_len = chunk_end - chunk_start;
+
+  extern __shared__ uint8_t smem[];
+  const T *q_now = q + (q_start_idx * q_num_heads + q_head_idx) * HEAD_DIM_QK;
+  T *q_smem = reinterpret_cast<T*>(smem); // [HEAD_DIM_QK * sizeof(T)]
+  T *cu_q_smem = q_smem + bidy * HEAD_DIM_QK;
+#pragma unroll
+  for(uint32_t vid = tidx; vid < num_vec_per_head_qk; vid += bdx) {
+    ((float4*)(&cu_q_smem[vid * VEC_SIZE]))[0] = ((float4*)(&q_now[vid * VEC_SIZE]))[0];
+  }
+  __syncthreads();
+  using VecT = AlignedVector<T, VEC_SIZE>;
+  VecT q_vec;
+#pragma unroll
+  for(uint32_t vid = tidx; vid < num_vec_per_head_qk; vid += bdx) {
+    Load<T, VEC_SIZE>(cu_q_smem + vid * VEC_SIZE, &q_vec);
+    for (uint32_t i = 0; i < VEC_SIZE; ++i) {
+      q_vec[i] *= scale;
+    }
+    Store<T, VEC_SIZE>(q_vec, cu_q_smem + vid * VEC_SIZE);
+  }
+
+
+  CacheT *kv_smem = reinterpret_cast<CacheT*>(smem + bdy * HEAD_DIM_QK * sizeof(CacheT)); 
+  uint32_t stage_idx = 0;  
+  constexpr int loop_times = DEAL_EACH_TIME / bdy;
+#pragma unroll
+  for (int i = 0; i < NUM_STAGES; ++i) {
+#pragma unroll
+    for (int j = 0; j < loop_times; ++j) {
+      const uint32_t k_seq_offset = i * DEAL_EACH_TIME + j * bdy + bidy;
+      const uint32_t k_seq_id = chunk_start + k_seq_offset;
+      produce_kv<SharedMemFillMode::kNoFill, HEAD_DIM_QK, VEC_SIZE, num_vec_per_head_qk, bdx, BLOCK_SIZE, CACHE_VEC_SIZE>(
+        kv_smem,
+        cache_k,
+        block_table_now,
+        k_seq_id,
+        k_seq_offset,
+        kv_head_idx,
+        kv_num_heads,
+        tidx,
+        chunk_start,
+        chunk_end
+      );
+    }
+    commit_group();
+    stage_idx = (stage_idx + 1) % NUM_STAGES;
+  }
+
+
+  softmax_state_ts<VEC_SIZE, T, num_tile_v> st;
+  float s[DEAL_EACH_TIME];
+  
+  const uint32_t num_iters = div_up(chunk_len, DEAL_EACH_TIME);
+  for (int iter = 0; iter < num_iters; ++iter) {
+    wait_group<NUM_STAGES - 1>();
+    __syncthreads();
+    // compute qk
+    compute_qk<VEC_SIZE, num_vec_per_head_qk, bdx, bdy, HEAD_DIM_QK, DEAL_EACH_TIME, num_tile_v>(
+      cu_q_smem,
+      kv_smem,
+      chunk_start + iter * DEAL_EACH_TIME,
+      stage_idx,
+      iter * DEAL_EACH_TIME,
+      chunk_len,
+      tidx,
+      gid,
+      scale,
+      s,
+      st
+    );
+    __syncthreads();
+
+    // compute sv
+    compute_sv<VEC_SIZE, num_vec_per_head_v, bdx, DEAL_EACH_TIME, HEAD_DIM_QK, num_tile_v>(
+      s,
+      kv_smem,
+      stage_idx,
+      iter * DEAL_EACH_TIME,
+      chunk_len,
+      tidx,
+      st
+    );
+    __syncthreads();
+
+#pragma unroll
+    for (int j = 0; j < loop_times; ++j) {
+      const uint32_t k_seq_offset = j * bdy + bidy;
+      produce_kv<SharedMemFillMode::kNoFill, HEAD_DIM_QK, VEC_SIZE, num_vec_per_head_qk, bdx, BLOCK_SIZE, CACHE_VEC_SIZE>(
+        kv_smem,
+        cache_k,
+        block_table_now,
+        chunk_start + k_seq_offset + (iter + NUM_STAGES) * DEAL_EACH_TIME,
+        stage_idx * DEAL_EACH_TIME + k_seq_offset,
+        kv_head_idx,
+        kv_num_heads,
+        tidx,
+        chunk_start,
+        chunk_end
+      );
+    }
+    commit_group();
+    stage_idx = (stage_idx + 1) % NUM_STAGES;
+  }
+  wait_group<0>();
+  __syncthreads();
+
+  // normize if not partition_kv
+  for(uint32_t vid = tidx; vid < num_vec_per_head_v; vid += bdx) {
+    const uint32_t tile_id = vid / bdx;
+    if (!partition_kv || num_chunk_this_seq == 1) {
+      st.normalize(tile_id);
+    }
+    if (partition_kv && num_chunk_this_seq > 1) {
+      const uint32_t head_idx = (bid * num_chunks + chunk_id) * q_num_heads + q_head_idx;
+      Store<T, VEC_SIZE>(st.o[tile_id], tmp_workspace + head_idx * HEAD_DIM_V + vid * VEC_SIZE);
+      tmp_m[head_idx] = st.m;
+      tmp_d[head_idx] = st.d;
+    } else {
+      Store<OutT, VEC_SIZE>(st.o[tile_id], out + (q_write_idx * q_num_heads + q_head_idx) * HEAD_DIM_V + vid * VEC_SIZE);
+    }
+  }
+}
+
+
+template <typename T, uint32_t GROUP_SIZE, uint32_t HEAD_DIM_QK, uint32_t HEAD_DIM_V, uint32_t BLOCK_SIZE, bool CAUSAL, uint32_t NUM_STAGE, uint32_t cache_bytes, uint32_t DEAL_EACH_TIME>
+void MultiQueryDecoderAttention(
+  const AppendAttnMetaData& meta_data,
+  cudaStream_t &stream,
+  const paddle::Tensor &q,
+  const paddle::Tensor &cache_k, // [max_block_num, num_kv_heads, block_size, head_dim]
+  const paddle::Tensor &cache_v, // [num_kv_heads, head_dim]
+  const paddle::optional<paddle::Tensor>& attn_mask,
+  const paddle::optional<paddle::Tensor>& shift_bias,
+  const paddle::optional<paddle::Tensor>& smooth_weight,
+  const paddle::Tensor &seq_lens_q,
+  const paddle::Tensor &seq_lens_kv,
+  const paddle::Tensor &padding_offsets,
+  const paddle::Tensor &cum_offsets,
+  const paddle::Tensor &block_table,
+  const int max_seq_len,
+  const int max_dec_len,
+  const float rope_scale,
+  const float rope_theta,
+  const float softmax_scale,
+  const float in_scale,
+  paddle::Tensor *out) {
+  using NV_TYPE = typename cascade_attn_type_traits<T>::type;
+
+  auto num_heads = meta_data.q_num_heads;
+  auto kv_num_heads = meta_data.kv_num_heads;
+  auto token_num = meta_data.token_nums;
+  auto bsz = meta_data.batch_size;
+  auto max_block_num_per_seq = meta_data.max_blocks_per_seq;
+  constexpr int num_stages = NUM_STAGE;
+
+  constexpr int vec_size = 16 / sizeof(T); // 8 16 32
+  constexpr int cache_vec_size = 128 / cache_bytes; // 8 16 32
+  constexpr int blockxc = HEAD_DIM_QK / cache_vec_size;
+  constexpr int num_vec_per_head = HEAD_DIM_QK / vec_size;
+  constexpr int blockx = num_vec_per_head < 32 ? num_vec_per_head : 32;
+
+  constexpr int blocky = GROUP_SIZE > 16 ? 16 : GROUP_SIZE;
+  const int gridx = bsz * (GROUP_SIZE / blocky);
+  
+  constexpr int num_threads = blockx * blocky; 
+
+  auto splitkv_kernel = multi_query_decode_attention_kernel<true, NV_TYPE, NV_TYPE, NV_TYPE, num_stages, DEAL_EACH_TIME, GROUP_SIZE, HEAD_DIM_QK, HEAD_DIM_V,
+                                                                        BLOCK_SIZE, vec_size, cache_vec_size, blockx, blocky>;
+  uint32_t cache_smem_bytes = 0;
+  
+  const T *shift_bias_ptr = shift_bias ? shift_bias.get().data<T>() : nullptr;
+  const T *smooth_weight_ptr = smooth_weight ? smooth_weight.get().data<T>() : nullptr;
+  cache_smem_bytes = num_stages * DEAL_EACH_TIME * HEAD_DIM_QK * sizeof(T);
+  
+  const uint32_t chunk_size = get_max_partition_size(bsz);
+  const int num_chunks = div_up(max_dec_len, chunk_size);
+  size_t smem_size = cache_smem_bytes + blocky * HEAD_DIM_QK * sizeof(T);
+
+  if (smem_size >= 48 * 1024) {
+    cudaFuncSetAttribute(
+      splitkv_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+  }
+  const int dev_id = 0;
+  int sm_count;
+  int act_blocks_per_sm;
+  cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev_id);
+  cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+    &act_blocks_per_sm, splitkv_kernel, num_threads, smem_size);
+  assert(act_blocks_per_sm > 1);
+  
+  const int num_blocks_per_wave = sm_count * act_blocks_per_sm;
+  const int num_blocks_need = gridx * num_chunks * kv_num_heads;
+  const int max_num_chunks = div_up(num_blocks_per_wave, num_blocks_need);
+  const float ratio = static_cast<float>(num_blocks_need) / static_cast<float>(num_blocks_per_wave);
+
+  dim3 grids(gridx, num_chunks, kv_num_heads);
+  dim3 blocks(blockx, blocky);
+  if (num_chunks <= 1) {
+    auto no_splitkv_kernel = multi_query_decode_attention_kernel<false, NV_TYPE, NV_TYPE, NV_TYPE, num_stages, DEAL_EACH_TIME, GROUP_SIZE, HEAD_DIM_QK, HEAD_DIM_V, BLOCK_SIZE, vec_size, 
+                                                                             cache_vec_size, blockx, blocky>;
+    if (smem_size >= 48 * 1024) {
+      cudaFuncSetAttribute(
+        no_splitkv_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+    }
+    no_splitkv_kernel<<<grids, blocks, smem_size, stream>>>(
+      reinterpret_cast<NV_TYPE*>(const_cast<T*>(q.data<T>())),
+      reinterpret_cast<NV_TYPE*>(const_cast<T*>(cache_k.data<T>())),
+      reinterpret_cast<NV_TYPE*>(const_cast<T*>(cache_v.data<T>())),
+      reinterpret_cast<NV_TYPE*>(const_cast<T*>(shift_bias_ptr)),
+      reinterpret_cast<NV_TYPE*>(const_cast<T*>(smooth_weight_ptr)),
+      seq_lens_q.data<int>(),
+      seq_lens_kv.data<int>(),
+      cum_offsets.data<int>(),
+      block_table.data<int>(),
+      max_seq_len,
+      max_dec_len,
+      max_block_num_per_seq,
+      softmax_scale,
+      in_scale,
+      chunk_size,
+      nullptr,
+      nullptr,
+      nullptr,
+      reinterpret_cast<NV_TYPE*>(const_cast<T*>(out->data<T>()))
+    );
+
+    // CHECK(cudaGetLastError());
+    // CHECK(cudaDeviceSynchronize());
+  } else {
+    auto *allocator = paddle::GetAllocator(q.place());
+    phi::Allocator::AllocationPtr tmp_workspace, tmp_m, tmp_d;
+    tmp_workspace = allocator->Allocate(
+        phi::SizeOf(q.dtype()) *
+        static_cast<size_t>(bsz * num_chunks * num_heads * HEAD_DIM_V));
+    tmp_m = allocator->Allocate(
+        phi::SizeOf(q.dtype()) *
+        static_cast<size_t>(bsz * num_chunks * num_heads));
+    tmp_d = allocator->Allocate(
+        phi::SizeOf(q.dtype()) *
+        static_cast<size_t>(bsz * num_chunks * num_heads));
+
+    splitkv_kernel<<<grids, blocks, smem_size, stream>>>(
+      reinterpret_cast<NV_TYPE*>(const_cast<T*>(q.data<T>())),
+      reinterpret_cast<NV_TYPE*>(const_cast<T*>(cache_k.data<T>())),
+      reinterpret_cast<NV_TYPE*>(const_cast<T*>(cache_v.data<T>())),
+      reinterpret_cast<NV_TYPE*>(const_cast<T*>(shift_bias_ptr)),
+      reinterpret_cast<NV_TYPE*>(const_cast<T*>(smooth_weight_ptr)),
+      seq_lens_q.data<int>(),
+      seq_lens_kv.data<int>(),
+      cum_offsets.data<int>(),
+      block_table.data<int>(),
+      max_seq_len,
+      max_dec_len,
+      max_block_num_per_seq,
+      softmax_scale,
+      in_scale,
+      chunk_size,
+      reinterpret_cast<NV_TYPE*>(tmp_workspace->ptr()),
+      reinterpret_cast<NV_TYPE*>(tmp_m->ptr()),
+      reinterpret_cast<NV_TYPE*>(tmp_d->ptr()),
+      reinterpret_cast<NV_TYPE*>(const_cast<T*>(out->data<T>()))
+    );
+    // CHECK(cudaGetLastError());
+    // CHECK(cudaDeviceSynchronize());
+    
+    constexpr int mblockx = HEAD_DIM_V / vec_size;
+    constexpr int bdy = 256 / mblockx;
+    dim3 grids_merge(bsz, num_heads);
+    dim3 blocks_merge(mblockx, bdy);
+    merge_varlen_multi_chunks_v2_kernel<NV_TYPE, NV_TYPE, vec_size, bdy, HEAD_DIM_V><<<grids_merge, blocks_merge, 0, stream>>>(
+      reinterpret_cast<NV_TYPE*>(tmp_workspace->ptr()),
+      reinterpret_cast<NV_TYPE*>(tmp_m->ptr()),
+      reinterpret_cast<NV_TYPE*>(tmp_d->ptr()),
+      seq_lens_q.data<int>(),
+      seq_lens_kv.data<int>(),
+      cum_offsets.data<int>(),
+      reinterpret_cast<NV_TYPE*>(const_cast<T*>(shift_bias_ptr)),
+      reinterpret_cast<NV_TYPE*>(const_cast<T*>(smooth_weight_ptr)),
+      reinterpret_cast<NV_TYPE*>(const_cast<T*>(out->data<T>())),
+      in_scale,
+      num_chunks,
+      chunk_size,
+      max_seq_len,
+      num_heads,
+      HEAD_DIM_V
+    );
+  }
+  // CHECK(cudaGetLastError());
+  // CHECK(cudaDeviceSynchronize());
+}
+
+template <typename T>
+void DecodeMLAAttentionKernel(
+  const AppendAttnMetaData& meta_data,
+  const paddle::Tensor &q, // [token_num, num_heads, head_dim]
+  const paddle::Tensor &cache_k,
+  const paddle::Tensor &cache_v,
+  const paddle::optional<paddle::Tensor>& attn_mask,
+  const paddle::optional<paddle::Tensor>& shift_bias,
+  const paddle::optional<paddle::Tensor>& smooth_weight,
+  const paddle::Tensor &seq_lens_q, // q_seq_len is 1
+  const paddle::Tensor &seq_lens_kv,
+  const paddle::Tensor &padding_offsets,
+  const paddle::Tensor &cum_offsets,
+  const paddle::Tensor &block_table,
+  int max_seq_len,
+  int max_dec_len,
+  float softmax_scale,
+  float in_scale,
+  bool causal,
+  cudaStream_t &stream,
+  paddle::Tensor *out) { 
+  const auto token_num = meta_data.token_nums;
+  const auto block_size = meta_data.block_size;
+  const auto bsz = meta_data.batch_size;
+  const auto num_heads = meta_data.q_num_heads;
+  const auto group_size = meta_data.q_num_heads / meta_data.kv_num_heads;
+  const auto head_dim_qk = meta_data.head_dims;
+  const auto head_dim_v = meta_data.head_dims_v;
+  const float rope_scale = 0.0;
+  const float rope_theta = 0.0;
+  const uint32_t deal_each_time = get_cascade_attention_deal_each_time();
+  const uint32_t num_stage = get_cascade_attention_num_stages();
+  const uint32_t num_threads = get_cascade_attention_num_threads();
+
+  DISPATCH_CAUSAL(causal, CAUSAL,
+    {DISPATCH_MLA_GROUP_SIZE(group_size, GROUP_SIZE,
+      {DISPATCH_MLA_HEAD_DIM(head_dim_qk, HEAD_DIM_QK,  
+        {DISPATCH_MLA_HEAD_DIM(head_dim_v, HEAD_DIM_V, 
+          {DISPATCH_BLOCK_SIZE(block_size, BLOCK_SIZE, 
+              {DISPATCH_DEAL_EACH_TIME(deal_each_time, DEAL_EACH_TIME,
+                  {MultiQueryDecoderAttention<T, GROUP_SIZE, HEAD_DIM_QK, HEAD_DIM_V, BLOCK_SIZE, CAUSAL, 2, 16, DEAL_EACH_TIME>(
+                  meta_data, stream, q, cache_k, cache_v, attn_mask, shift_bias, smooth_weight, seq_lens_q, seq_lens_kv, padding_offsets, cum_offsets, 
+                  block_table, max_seq_len, max_dec_len, rope_scale, rope_theta, softmax_scale, in_scale, out);})})})})})});
+}
+
+template void DecodeMLAAttentionKernel<paddle::bfloat16>(
+  const AppendAttnMetaData& meta_data,
+  const paddle::Tensor &q, // [token_num, num_heads, head_dim]
+  const paddle::Tensor &cache_k,
+  const paddle::Tensor &cache_v,
+  const paddle::optional<paddle::Tensor>& attn_mask,
+  const paddle::optional<paddle::Tensor>& shift_bias,
+  const paddle::optional<paddle::Tensor>& smooth_weight,
+  const paddle::Tensor &seq_lens_q, // q_seq_len is 1
+  const paddle::Tensor &seq_lens_kv,
+  const paddle::Tensor &padding_offsets,
+  const paddle::Tensor &cum_offsets,
+  const paddle::Tensor &block_table,
+  int max_seq_len,
+  int max_dec_len,
+  float softmax_scale,
+  float in_scale,
+  bool causal,
+  cudaStream_t &stream,
+  paddle::Tensor *out);
+
+template void DecodeMLAAttentionKernel<paddle::float16>(
+  const AppendAttnMetaData& meta_data,
+  const paddle::Tensor &q, // [token_num, num_heads, head_dim]
+  const paddle::Tensor &cache_k,
+  const paddle::Tensor &cache_v,
+  const paddle::optional<paddle::Tensor>& attn_mask,
+  const paddle::optional<paddle::Tensor>& shift_bias,
+  const paddle::optional<paddle::Tensor>& smooth_weight,
+  const paddle::Tensor &seq_lens_q, // q_seq_len is 1
+  const paddle::Tensor &seq_lens_kv,
+  const paddle::Tensor &padding_offsets,
+  const paddle::Tensor &cum_offsets,
+  const paddle::Tensor &block_table,
+  int max_seq_len,
+  int max_dec_len,
+ float softmax_scale,
+  float in_scale,
+  bool causal,
+  cudaStream_t &stream,
+  paddle::Tensor *out);

csrc/gpu/append_attn/decoder_write_cache_with_rope_kernel.cu

@@ -0,0 +1,722 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "decoder_write_cache_with_rope_kernel.h"
+#include "utils.cuh"
+
+
+template <typename T>
+void DecoderWriteCacheKV(const AppendAttnMetaData& meta_data,
+                         const paddle::Tensor& qkv,
+                         const paddle::Tensor& seq_lens,
+                         const paddle::Tensor& seq_lens_encoder,
+                         const paddle::Tensor& padding_offsets,
+                         const paddle::Tensor& cum_offsets,
+                         const paddle::Tensor& block_tables,
+                         const int max_seq_len,
+                         cudaStream_t& stream,
+                         paddle::Tensor* key_cache_out,
+                         paddle::Tensor* value_cache_out) {
+  auto max_blocks_per_seq = meta_data.max_blocks_per_seq;
+  auto bsz = meta_data.batch_size;
+  auto block_size = meta_data.block_size;
+  auto head_dim_qk = meta_data.head_dims;
+  auto head_dim_v = meta_data.head_dims_v;
+  auto num_heads = meta_data.q_num_heads;
+  auto kv_num_heads = meta_data.kv_num_heads;
+  const uint32_t elem_nums = bsz * kv_num_heads * (head_dim_qk + head_dim_v);
+
+  constexpr int PackSize = 16 / sizeof(T);
+  const int pack_num = elem_nums / PackSize;
+  const int blocksize = 128;
+  int grid_size = 1;
+  GetNumBlocks<128>(pack_num, &grid_size);
+
+  append_decode_cache_T_kernel<T, PackSize>
+      <<<grid_size, blocksize, 0, stream>>>(
+          reinterpret_cast<T*>(const_cast<T*>(qkv.data<T>())),
+          reinterpret_cast<T*>(key_cache_out->data<T>()),
+          reinterpret_cast<T*>(value_cache_out->data<T>()),
+          block_tables.data<int>(),
+          padding_offsets.data<int>(),
+          cum_offsets.data<int>(),
+          seq_lens.data<int>(),
+          seq_lens_encoder.data<int>(),
+          max_seq_len,
+          max_blocks_per_seq,
+          num_heads,
+          head_dim_qk,
+          head_dim_v,
+          block_size,
+          elem_nums,
+          kv_num_heads);
+}
+
+template <typename T, typename QKV_TYPE>
+void append_decode_cache_rope(const QKV_TYPE* qkv,
+                              T* key_cache,
+                              T* value_cache,
+                              T* qkv_out,
+                              const int* block_tables,
+                              const int* padding_offsets,
+                              const int* cum_offsets,
+                              const int* seq_lens,
+                              const int* seq_lens_encoder,
+                              const float* cos_emb,
+                              const float* sin_emb,
+                              const float* qkv_out_scales,
+                              const T* qkv_biases,
+                              const int max_seq_len,
+                              const int max_blocks_per_seq,
+                              const int num_heads,
+                              const int kv_num_heads,
+                              const int dim_head,
+                              const int block_size,
+                              const int bsz,
+                              const cudaStream_t& stream,
+                              const bool use_neox_style) {
+  const uint32_t elem_nums =
+      use_neox_style ? bsz * (num_heads + 2 * kv_num_heads) * dim_head / 2
+                     : bsz * (num_heads + 2 * kv_num_heads) * dim_head;
+
+  constexpr int PackSize = 16 / sizeof(T);
+  const int pack_num = elem_nums / PackSize;
+  const int blocksize = 128;
+  int grid_size = 1;
+  GetNumBlocks<128>(pack_num, &grid_size);
+  if (use_neox_style) {
+    if (qkv_out_scales) {
+      append_decode_cache_T_neox_rope_kernel<T, PackSize>
+          <<<grid_size, blocksize, 0, stream>>>(
+              reinterpret_cast<const int*>(qkv),
+              key_cache,
+              value_cache,
+              qkv_out,
+              block_tables,
+              padding_offsets,
+              cum_offsets,
+              seq_lens,
+              seq_lens_encoder,
+              cos_emb,
+              sin_emb,
+              qkv_out_scales,
+              qkv_biases,
+              max_seq_len,
+              max_blocks_per_seq,
+              num_heads,
+              dim_head,
+              block_size,
+              elem_nums,
+              kv_num_heads);
+    } else {
+      append_decode_cache_T_neox_rope_kernel<T, PackSize>
+          <<<grid_size, blocksize, 0, stream>>>(reinterpret_cast<const T*>(qkv),
+                                                key_cache,
+                                                value_cache,
+                                                qkv_out,
+                                                block_tables,
+                                                padding_offsets,
+                                                cum_offsets,
+                                                seq_lens,
+                                                seq_lens_encoder,
+                                                cos_emb,
+                                                sin_emb,
+                                                max_seq_len,
+                                                max_blocks_per_seq,
+                                                num_heads,
+                                                dim_head,
+                                                block_size,
+                                                elem_nums,
+                                                kv_num_heads);
+    }
+  } else {
+    if (qkv_out_scales) {
+      append_decode_cache_T_rope_kernel<T, PackSize>
+          <<<grid_size, blocksize, 0, stream>>>(
+              reinterpret_cast<const int*>(qkv),
+              key_cache,
+              value_cache,
+              qkv_out,
+              block_tables,
+              padding_offsets,
+              cum_offsets,
+              seq_lens,
+              seq_lens_encoder,
+              cos_emb,
+              sin_emb,
+              qkv_out_scales,
+              qkv_biases,
+              max_seq_len,
+              max_blocks_per_seq,
+              num_heads,
+              dim_head,
+              block_size,
+              elem_nums,
+              kv_num_heads);
+    } else {
+      append_decode_cache_T_rope_kernel<T, PackSize>
+          <<<grid_size, blocksize, 0, stream>>>(reinterpret_cast<const T*>(qkv),
+                                                key_cache,
+                                                value_cache,
+                                                qkv_out,
+                                                block_tables,
+                                                padding_offsets,
+                                                cum_offsets,
+                                                seq_lens,
+                                                seq_lens_encoder,
+                                                cos_emb,
+                                                sin_emb,
+                                                max_seq_len,
+                                                max_blocks_per_seq,
+                                                num_heads,
+                                                dim_head,
+                                                block_size,
+                                                elem_nums,
+                                                kv_num_heads);
+    }
+  }
+}
+
+template <typename T, typename QKV_TYPE>
+void append_decode_cache_int8_rope(const QKV_TYPE* qkv,
+                                   uint8_t* key_cache,
+                                   uint8_t* value_cache,
+                                   T* qkv_out,
+                                   const int* block_tables,
+                                   const int* padding_offsets,
+                                   const int* cum_offsets,
+                                   const int* seq_lens,
+                                   const int* seq_lens_encoder,
+                                   const float* cos_emb,
+                                   const float* sin_emb,
+                                   const float* qkv_out_scales,
+                                   const T* qkv_biases,
+                                   const T* cache_k_scale,
+                                   const T* cache_v_scale,
+                                   const int max_seq_len,
+                                   const int max_blocks_per_seq,
+                                   const int num_heads,
+                                   const int kv_num_heads,
+                                   const int dim_head,
+                                   const int block_size,
+                                   const int bsz,
+                                   const cudaStream_t& stream,
+                                   const bool use_neox_style) {
+  constexpr int num_warps = 4;
+  const int all_warps =
+      ((num_heads + 2 * kv_num_heads) + num_warps - 1) / num_warps * num_warps;
+  dim3 grids(bsz, all_warps / num_warps);
+  if (use_neox_style) {
+    if (qkv_out_scales) {
+      append_decode_cache_int8_neox_rope_kernel<T, 4>
+          <<<grids, num_warps * 32, 0, stream>>>(
+              reinterpret_cast<const int*>(qkv),
+              key_cache,
+              value_cache,
+              qkv_out,
+              block_tables,
+              padding_offsets,
+              cum_offsets,
+              seq_lens,
+              seq_lens_encoder,
+              cos_emb,
+              sin_emb,
+              qkv_out_scales,
+              qkv_biases,
+              cache_k_scale,
+              cache_v_scale,
+              max_seq_len,
+              max_blocks_per_seq,
+              num_heads,
+              block_size,
+              127.0f,
+              -127.0f,
+              kv_num_heads);
+    } else {
+      append_decode_cache_int8_neox_rope_kernel<T, 4>
+          <<<grids, num_warps * 32, 0, stream>>>(
+              reinterpret_cast<const T*>(qkv),
+              key_cache,
+              value_cache,
+              qkv_out,
+              block_tables,
+              padding_offsets,
+              cum_offsets,
+              seq_lens,
+              seq_lens_encoder,
+              cos_emb,
+              sin_emb,
+              cache_k_scale,
+              cache_v_scale,
+              max_seq_len,
+              max_blocks_per_seq,
+              num_heads,
+              block_size,
+              127.0f,
+              -127.0f,
+              kv_num_heads);
+    }
+  } else {
+    if (qkv_out_scales) {
+      append_decode_cache_int8_rope_kernel<T, 4>
+          <<<grids, num_warps * 32, 0, stream>>>(
+              reinterpret_cast<const int*>(qkv),
+              key_cache,
+              value_cache,
+              qkv_out,
+              block_tables,
+              padding_offsets,
+              cum_offsets,
+              seq_lens,
+              seq_lens_encoder,
+              cos_emb,
+              sin_emb,
+              qkv_out_scales,
+              qkv_biases,
+              cache_k_scale,
+              cache_v_scale,
+              max_seq_len,
+              max_blocks_per_seq,
+              num_heads,
+              block_size,
+              127.0f,
+              -127.0f,
+              kv_num_heads);
+    } else {
+      append_decode_cache_int8_rope_kernel<T, 4>
+          <<<grids, num_warps * 32, 0, stream>>>(
+              reinterpret_cast<const T*>(qkv),
+              key_cache,
+              value_cache,
+              qkv_out,
+              block_tables,
+              padding_offsets,
+              cum_offsets,
+              seq_lens,
+              seq_lens_encoder,
+              cos_emb,
+              sin_emb,
+              cache_k_scale,
+              cache_v_scale,
+              max_seq_len,
+              max_blocks_per_seq,
+              num_heads,
+              block_size,
+              127.0f,
+              -127.0f,
+              kv_num_heads);
+    }
+  }
+}
+
+template <typename T, typename QKV_TYPE>
+void append_decode_cache_int4_rope(const QKV_TYPE* qkv,
+                                   uint8_t* key_cache,
+                                   uint8_t* value_cache,
+                                   T* qkv_out,
+                                   const int* block_tables,
+                                   const int* padding_offsets,
+                                   const int* cum_offsets,
+                                   const int* seq_lens,
+                                   const int* seq_lens_encoder,
+                                   const float* cos_emb,
+                                   const float* sin_emb,
+                                   const float* qkv_out_scales,
+                                   const T* qkv_biases,
+                                   const T* cache_k_scale,
+                                   const T* cache_v_scale,
+                                   const T* cache_k_zp,
+                                   const T* cache_v_zp,
+                                   const int max_seq_len,
+                                   const int max_blocks_per_seq,
+                                   const int num_heads,
+                                   const int kv_num_heads,
+                                   const int dim_head,
+                                   const int block_size,
+                                   const int bsz,
+                                   const cudaStream_t& stream,
+                                   const bool use_neox_style) {
+  constexpr int num_warps = 4;
+  const int all_warps =
+      ((num_heads + 2 * kv_num_heads) + num_warps - 1) / num_warps * num_warps;
+  dim3 grids(bsz, all_warps / num_warps);
+  if (use_neox_style) {
+    if (qkv_out_scales) {
+      append_decode_cache_int4_neox_rope_kernel<T, 4>
+          <<<grids, num_warps * 32, 0, stream>>>(
+              reinterpret_cast<const int*>(qkv),
+              key_cache,
+              value_cache,
+              qkv_out,
+              block_tables,
+              padding_offsets,
+              cum_offsets,
+              seq_lens,
+              seq_lens_encoder,
+              cos_emb,
+              sin_emb,
+              qkv_out_scales,
+              qkv_biases,
+              cache_k_scale,
+              cache_v_scale,
+              cache_k_zp,
+              cache_v_zp,
+              max_seq_len,
+              max_blocks_per_seq,
+              num_heads,
+              block_size,
+              7.0f,
+              -8.0f,
+              kv_num_heads);
+    } else {
+      append_decode_cache_int4_neox_rope_kernel<T, 4>
+          <<<grids, num_warps * 32, 0, stream>>>(
+              reinterpret_cast<const T*>(qkv),
+              key_cache,
+              value_cache,
+              qkv_out,
+              block_tables,
+              padding_offsets,
+              cum_offsets,
+              seq_lens,
+              seq_lens_encoder,
+              cos_emb,
+              sin_emb,
+              cache_k_scale,
+              cache_v_scale,
+              cache_k_zp,
+              cache_v_zp,
+              max_seq_len,
+              max_blocks_per_seq,
+              num_heads,
+              block_size,
+              7.0f,
+              -8.0f,
+              kv_num_heads);
+    }
+  } else {
+    if (qkv_out_scales) {
+      append_decode_cache_int4_rope_kernel<T, 4>
+          <<<grids, num_warps * 32, 0, stream>>>(
+              reinterpret_cast<const int*>(qkv),
+              key_cache,
+              value_cache,
+              qkv_out,
+              block_tables,
+              padding_offsets,
+              cum_offsets,
+              seq_lens,
+              seq_lens_encoder,
+              cos_emb,
+              sin_emb,
+              qkv_out_scales,
+              qkv_biases,
+              cache_k_scale,
+              cache_v_scale,
+              cache_k_zp,
+              cache_v_zp,
+              max_seq_len,
+              max_blocks_per_seq,
+              num_heads,
+              block_size,
+              7.0f,
+              -8.0f,
+              kv_num_heads);
+    } else {
+      append_decode_cache_int4_rope_kernel<T, 4>
+          <<<grids, num_warps * 32, 0, stream>>>(
+              reinterpret_cast<const T*>(qkv),
+              key_cache,
+              value_cache,
+              qkv_out,
+              block_tables,
+              padding_offsets,
+              cum_offsets,
+              seq_lens,
+              seq_lens_encoder,
+              cos_emb,
+              sin_emb,
+              cache_k_scale,
+              cache_v_scale,
+              cache_k_zp,
+              cache_v_zp,
+              max_seq_len,
+              max_blocks_per_seq,
+              num_heads,
+              block_size,
+              7.0f,
+              -8.0f,
+              kv_num_heads);
+    }
+  }
+}
+template <typename T, typename QKV_TYPE>
+void DecoderWriteCacheWithRoPEKernel(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& qkv_biases,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const std::string& cache_quant_type_str,
+    const bool use_neox_rotary_style,
+    const int max_seq_len,
+    cudaStream_t& stream,
+    paddle::Tensor* qkv_out,
+    paddle::Tensor* key_cache_out,
+    paddle::Tensor* value_cache_out) {
+  typedef cascade_attn_type_traits<T> traits_;
+  typedef cascade_attn_type_traits<QKV_TYPE> qkt_nv_type_;
+  typedef typename traits_::type DataType_;
+  typedef typename qkt_nv_type_::type QKV_Data_TYPE;
+  const QKV_TYPE* qkv_ptr = qkv.data<QKV_TYPE>();
+
+  auto max_blocks_per_seq = meta_data.max_blocks_per_seq;
+  auto bsz = meta_data.batch_size;
+  auto block_size = meta_data.block_size;
+  auto dim_head = meta_data.head_dims;
+  auto num_heads = meta_data.q_num_heads;
+  auto kv_num_heads = meta_data.kv_num_heads;
+
+  if (rotary_embs) {
+    const float* cos_emb = rotary_embs.get().data<float>();
+    const float* sin_emb =
+        use_neox_rotary_style
+            ? rotary_embs.get().data<float>() + max_seq_len * dim_head
+            : rotary_embs.get().data<float>() + max_seq_len * dim_head / 2;
+    if (cache_quant_type_str == "none") {
+      append_decode_cache_rope(
+          reinterpret_cast<const QKV_TYPE*>(qkv_ptr),
+          reinterpret_cast<DataType_*>(key_cache_out->data<T>()),
+          reinterpret_cast<DataType_*>(value_cache_out->data<T>()),
+          reinterpret_cast<DataType_*>(qkv_out->data<T>()),
+          block_tables.data<int>(),
+          padding_offsets.data<int>(),
+          cum_offsets.data<int>(),
+          seq_lens.data<int>(),
+          seq_lens_encoder.data<int>(),
+          cos_emb,
+          sin_emb,
+          qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
+          qkv_biases ? reinterpret_cast<DataType_*>(
+                           const_cast<T*>(qkv_biases.get().data<T>()))
+                     : nullptr,
+          max_seq_len,
+          max_blocks_per_seq,
+          num_heads,
+          kv_num_heads,
+          dim_head,
+          block_size,
+          bsz,
+          stream,
+          use_neox_rotary_style);
+    } else if (cache_quant_type_str == "cache_int8") {
+      append_decode_cache_int8_rope(
+          reinterpret_cast<const QKV_TYPE*>(qkv_ptr),
+          key_cache_out->data<uint8_t>(),
+          value_cache_out->data<uint8_t>(),
+          reinterpret_cast<DataType_*>(qkv_out->data<T>()),
+          block_tables.data<int>(),
+          padding_offsets.data<int>(),
+          cum_offsets.data<int>(),
+          seq_lens.data<int>(),
+          seq_lens_encoder.data<int>(),
+          cos_emb,
+          sin_emb,
+          qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
+          qkv_biases ? reinterpret_cast<DataType_*>(
+                           const_cast<T*>(qkv_biases.get().data<T>()))
+                     : nullptr,
+          cache_k_scale ? reinterpret_cast<DataType_*>(
+                              const_cast<T*>(cache_k_scale.get().data<T>()))
+                        : nullptr,
+          cache_v_scale ? reinterpret_cast<DataType_*>(
+                              const_cast<T*>(cache_v_scale.get().data<T>()))
+                        : nullptr,
+          max_seq_len,
+          max_blocks_per_seq,
+          num_heads,
+          kv_num_heads,
+          dim_head,
+          block_size,
+          bsz,
+          stream,
+          use_neox_rotary_style);
+    } else if (cache_quant_type_str == "cache_int4_zp") {
+      append_decode_cache_int4_rope(
+          reinterpret_cast<const QKV_TYPE*>(qkv_ptr),
+          key_cache_out->data<uint8_t>(),
+          value_cache_out->data<uint8_t>(),
+          reinterpret_cast<DataType_*>(const_cast<T*>(qkv_out->data<T>())),
+          block_tables.data<int>(),
+          padding_offsets.data<int>(),
+          cum_offsets.data<int>(),
+          seq_lens.data<int>(),
+          seq_lens_encoder.data<int>(),
+          cos_emb,
+          sin_emb,
+          qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
+          qkv_biases ? reinterpret_cast<DataType_*>(
+                           const_cast<T*>(qkv_biases.get().data<T>()))
+                     : nullptr,
+          cache_k_scale ? reinterpret_cast<DataType_*>(
+                              const_cast<T*>(cache_k_scale.get().data<T>()))
+                        : nullptr,
+          cache_v_scale ? reinterpret_cast<DataType_*>(
+                              const_cast<T*>(cache_v_scale.get().data<T>()))
+                        : nullptr,
+          cache_k_zp ? reinterpret_cast<DataType_*>(
+                           const_cast<T*>(cache_k_zp.get().data<T>()))
+                     : nullptr,
+          cache_v_zp ? reinterpret_cast<DataType_*>(
+                           const_cast<T*>(cache_v_zp.get().data<T>()))
+                     : nullptr,
+          max_seq_len,
+          max_blocks_per_seq,
+          num_heads,
+          kv_num_heads,
+          dim_head,
+          block_size,
+          bsz,
+          stream,
+          use_neox_rotary_style);
+    } else {
+      PD_THROW(
+          "cache_quant_type_str should be one of [none, cache_int8, "
+          "cache_int4_zp]");
+    }
+  } else {
+    DecoderWriteCacheKV<QKV_TYPE>(meta_data,
+                                  qkv,
+                                  seq_lens,
+                                  seq_lens_encoder,
+                                  padding_offsets,
+                                  cum_offsets,
+                                  block_tables,
+                                  max_seq_len,
+                                  stream,
+                                  key_cache_out,
+                                  value_cache_out);
+  }
+}
+
+
+template void DecoderWriteCacheWithRoPEKernel<paddle::bfloat16, int>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor&
+        qkv,  // [token_num, 3, num_head, head_dim] ([token_num, num_head + 2 *
+              // kv_num_heads, head_dim] if GQA)
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& qkv_biases,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const std::string& cache_quant_type_str,
+    const bool use_neox_rotary_style,
+    const int max_seq_len,
+    cudaStream_t& stream,
+    paddle::Tensor* qkv_out,
+    paddle::Tensor* key_cache_out,
+    paddle::Tensor* value_cache_out);
+
+template void
+DecoderWriteCacheWithRoPEKernel<paddle::bfloat16, paddle::bfloat16>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor&
+        qkv,  // [token_num, 3, num_head, head_dim] ([token_num, num_head + 2 *
+              // kv_num_heads, head_dim] if GQA)
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& qkv_biases,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const std::string& cache_quant_type_str,
+    const bool use_neox_rotary_style,
+    const int max_seq_len,
+    cudaStream_t& stream,
+    paddle::Tensor* qkv_out,
+    paddle::Tensor* key_cache_out,
+    paddle::Tensor* value_cache_out);
+
+template void DecoderWriteCacheWithRoPEKernel<paddle::float16, int>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor&
+        qkv,  // [token_num, 3, num_head, head_dim] ([token_num, num_head + 2 *
+              // kv_num_heads, head_dim] if GQA)
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& qkv_biases,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const std::string& cache_quant_type_str,
+    const bool use_neox_rotary_style,
+    const int max_seq_len,
+    cudaStream_t& stream,
+    paddle::Tensor* qkv_out,
+    paddle::Tensor* key_cache_out,
+    paddle::Tensor* value_cache_out);
+
+template void DecoderWriteCacheWithRoPEKernel<paddle::float16, paddle::float16>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor&
+        qkv,  // [token_num, 3, num_head, head_dim] ([token_num, num_head + 2 *
+              // kv_num_heads, head_dim] if GQA)
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& qkv_biases,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const std::string& cache_quant_type_str,
+    const bool use_neox_rotary_style,
+    const int max_seq_len,
+    cudaStream_t& stream,
+    paddle::Tensor* qkv_out,
+    paddle::Tensor* key_cache_out,
+    paddle::Tensor* value_cache_out);

csrc/gpu/append_attn/decoder_write_cache_with_rope_kernel.h

@@ -0,0 +1,42 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#pragma once
+
+#include "decoder_write_cache_with_rope_impl.cuh"
+
+template <typename T, typename QKV_TYPE = int>
+void DecoderWriteCacheWithRoPEKernel(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor&
+        qkv,  // [token_num, 3, num_head, head_dim] ([token_num, num_head + 2 *
+              // kv_num_heads, head_dim] if GQA)
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& qkv_biases,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const std::string& cache_quant_type_str,
+    const bool use_neox_rotary_style,
+    const int max_seq_len,
+    cudaStream_t& stream,
+    paddle::Tensor* qkv_out,
+    paddle::Tensor* key_cache_out,
+    paddle::Tensor* value_cache_out);

csrc/gpu/append_attn/encoder_write_cache_with_rope_kernel.h

@@ -0,0 +1,160 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#pragma once
+
+#include "encoder_write_cache_with_rope_impl.cuh"
+
+template <typename T, typename QKV_TYPE>
+void EncoderWriteCacheWithRopeKernel(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor&
+        qkv,  // [token_num, 3, num_head, head_dim] ([token_num, num_head + 2 *
+              // kv_num_heads, head_dim] if GQA)
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& qkv_biases,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const paddle::optional<paddle::Tensor>& excess_blocks,
+    const std::string& cache_quant_type_str,
+    const int num_blocks,
+    const int max_seq_len,
+    const bool use_neox_style,
+    cudaStream_t& stream,
+    paddle::Tensor* qkv_out,
+    paddle::Tensor* key_cache_out,
+    paddle::Tensor* value_cache_out) {
+  auto token_num = meta_data.token_nums;
+  auto num_heads = meta_data.q_num_heads;
+  auto kv_num_heads = meta_data.kv_num_heads;
+  auto head_dim = meta_data.head_dims;
+  int bsz = cum_offsets.dims()[0];
+  if (rotary_embs) {
+    if (num_heads == kv_num_heads) {
+      rotary_qk_variable(
+          qkv_out->data<T>(),
+          qkv.data<QKV_TYPE>(),
+          qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
+          qkv_biases ? qkv_biases.get().data<T>() : nullptr,
+          rotary_embs.get().data<float>(),
+          padding_offsets.data<int>(),
+          seq_lens_encoder.data<int>(),
+          seq_lens_decoder.data<int>(),
+          token_num,
+          num_heads,
+          max_seq_len,
+          rotary_embs.get().dims()[2],
+          head_dim,
+          stream,
+          use_neox_style);
+    } else {
+      gqa_rotary_qk_variable(
+          qkv_out->data<T>(),
+          qkv.data<QKV_TYPE>(),
+          qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
+          qkv_biases ? qkv_biases.get().data<T>() : nullptr,
+          rotary_embs.get().data<float>(),
+          padding_offsets.data<int>(),
+          seq_lens_encoder.data<int>(),
+          seq_lens_decoder.data<int>(),
+          token_num,
+          num_heads,
+          kv_num_heads,
+          max_seq_len,
+          rotary_embs.get().dims()[2],
+          head_dim,
+          stream,
+          use_neox_style);
+    }
+  }
+  
+  const uint32_t block_size = meta_data.block_size;
+  if (cache_quant_type_str == "none") {
+    CascadeAppendWriteCacheKVQKV<T>(meta_data,
+                                    *qkv_out,
+                                    block_tables,
+                                    padding_offsets,
+                                    cum_offsets,
+                                    seq_lens_encoder,
+                                    seq_lens_decoder,
+                                    max_seq_len,
+                                    bsz,
+                                    excess_blocks,
+                                    stream,
+                                    key_cache_out,
+                                    value_cache_out);
+  } else if (cache_quant_type_str == "cache_int8") {
+    DISPATCH_GQA_HEAD_DIM(
+        head_dim, HEAD_DIM, {DISPATCH_BLOCK_SIZE(block_size, BLOCK_SIZE, {
+          CascadeAppendWriteCacheKVC8QKV<T, HEAD_DIM, BLOCK_SIZE>(
+              meta_data,
+              *key_cache_out,
+              *value_cache_out,
+              *qkv_out,
+              cache_k_scale.get(),
+              cache_v_scale.get(),
+              seq_lens_this_time,
+              seq_lens_decoder,
+              padding_offsets,
+              cum_offsets,
+              block_tables,
+              batch_ids,
+              tile_ids,
+              num_blocks,
+              max_seq_len,
+              stream,
+              key_cache_out,
+              value_cache_out);
+        })})
+  } else if (cache_quant_type_str == "cache_int4_zp") {
+    DISPATCH_GQA_HEAD_DIM(
+        head_dim, HEAD_DIM, {DISPATCH_BLOCK_SIZE(block_size, BLOCK_SIZE, {
+          CascadeAppendWriteCacheKVC4QKV<T, HEAD_DIM, BLOCK_SIZE>(
+              meta_data,
+              *key_cache_out,
+              *value_cache_out,
+              *qkv_out,
+              cache_k_scale.get(),
+              cache_v_scale.get(),
+              cache_k_zp.get(),
+              cache_v_zp.get(),
+              seq_lens_this_time,
+              seq_lens_decoder,
+              padding_offsets,
+              cum_offsets,
+              block_tables,
+              batch_ids,
+              tile_ids,
+              num_blocks,
+              max_seq_len,
+              stream,
+              key_cache_out,
+              value_cache_out);
+        })})
+  } else {
+    PD_THROW(
+        "cache_quant_type_str should be one of [none, cache_int8, "
+        "cache_int4_zp]");
+  }
+}

csrc/gpu/append_attn/get_block_shape_and_split_kv_block.cu

@@ -0,0 +1,457 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "helper.h"
+#include "paddle/extension.h"
+#include "utils.cuh"
+#include "cute/tensor.hpp"
+
+template <typename T>
+inline __device__ __host__ T div_up(T m, T n) {
+  return (m + n - 1) / n;
+}
+
+__global__ void split_q_block(const int* __restrict__ seq_lens_q,
+                              const int* __restrict__ seq_lens_encoder,
+                              int* __restrict__ batch_ids,
+                              int* __restrict__ tile_ids_per_batch,
+                              int* __restrict__ num_blocks_x,
+                              const int bsz,
+                              const int num_rows_per_block,
+                              const int group_size) {
+  if (threadIdx.x == 0) {
+    int gridx = 0;
+    int index = 0;
+    for (uint32_t bid = 0; bid < bsz; bid++) {
+      int seq_len = seq_lens_q[bid];
+      if (seq_lens_encoder && seq_lens_encoder[bid] > 0) {
+        seq_len = 0;
+      }
+      const int loop_times =
+          div_up(seq_len * group_size, num_rows_per_block);
+      for (uint32_t tile_id = 0; tile_id < loop_times; tile_id++) {
+        batch_ids[index] = bid;
+        tile_ids_per_batch[index++] = tile_id;
+      }
+      gridx += loop_times;
+    }
+    *num_blocks_x = gridx;
+  }
+}
+
+template <uint32_t config_size>
+__global__ void search_chunk_size_for_mla(
+  const int * __restrict__ seq_lens_q,
+  const int * __restrict__ seq_lens_encoder,
+  const int * __restrict__ seq_lens_decoder,
+  int * __restrict__ num_blocks_x,
+  int * __restrict__ res_chunk_size,
+  const int bsz,
+  const int set_chunk_size,
+  const int block_size,
+  const int sm_cout) {
+  const uint32_t conf_id = threadIdx.x;
+  int gridx = 0;
+  if (set_chunk_size > 0 && conf_id == 0) {
+    for (uint32_t bid = 0; bid < bsz; bid++) {
+      int seq_len = seq_lens_q[bid];
+      int seq_len_encoder = seq_lens_encoder[bid];
+      int seq_len_decoder = seq_lens_decoder[bid] + seq_len;
+      if (seq_len == 0 || seq_len_encoder > 0) continue;
+
+      int loop_times;
+      loop_times = cute::ceil_div(seq_len_decoder, set_chunk_size);
+      gridx += loop_times;
+    }
+    *num_blocks_x = gridx;
+    *res_chunk_size = set_chunk_size;
+  } else if (conf_id < config_size) {
+    __shared__ int gridx_shared[config_size];
+    // chunk_size is a multiple of 64
+    const int chunk_size = block_size << conf_id;
+    for (uint32_t bid = 0; bid < bsz; bid++) {
+      int seq_len = seq_lens_q[bid];
+      int seq_len_encoder = seq_lens_encoder[bid];
+      int seq_len_decoder = seq_lens_decoder[bid] + seq_len;
+      if (seq_len == 0 || seq_len_encoder > 0) continue;
+
+      int loop_times;
+      loop_times = cute::ceil_div(seq_len_decoder, chunk_size);
+      gridx += loop_times;
+    }
+    gridx_shared[conf_id] = gridx;
+    __syncthreads();
+    if (threadIdx.x == 0) {
+      uint32_t res_id = 0;
+      uint32_t max_last_wave_block = 0;
+      for (uint32_t i = 1; i < config_size; i++) {
+        uint32_t last_wave_block = gridx_shared[i] % sm_cout;
+        if (last_wave_block >= max_last_wave_block) {
+          res_id = i;
+          max_last_wave_block = last_wave_block;
+        }
+      }
+      *num_blocks_x = gridx_shared[res_id];
+      *res_chunk_size = block_size << res_id;
+    }
+  }
+}
+
+
+
+__global__ void split_block_for_mla(
+  const int * __restrict__ seq_lens_q,
+  const int * __restrict__ seq_lens_encoder,
+  const int * __restrict__ seq_lens_decoder,
+  int * __restrict__ batch_ids,
+  int * __restrict__ tile_ids_per_batch,
+  const int bsz,
+  const int chunk_size) {
+  if (threadIdx.x == 0) {
+    int index = 0;
+    for (uint32_t bid = 0; bid < bsz; bid++) {
+      int seq_len = seq_lens_q[bid];
+      int seq_len_encoder = seq_lens_encoder[bid];
+      int seq_len_decoder = seq_lens_decoder[bid] + seq_len;
+
+      if (seq_len == 0) continue;
+
+      int loop_times;
+      loop_times = cute::ceil_div(seq_len_decoder, chunk_size);
+      if (seq_len_encoder > 0) {
+        loop_times = 0;
+      }
+      for (uint32_t tile_id = 0; tile_id < loop_times; tile_id++) {
+        batch_ids[index] = bid;
+        tile_ids_per_batch[index++] = tile_id;
+      }
+    }
+  }
+}
+
+__global__ void split_kv_block(const int* __restrict__ seq_lens_decoder,
+                               const int* __restrict__ seq_lens_encoder,
+                               int* __restrict__ batch_ids,
+                               int* __restrict__ tile_ids_per_batch,
+                               int* __restrict__ num_blocks_x,
+                               const int bsz,
+                               const int pad_len,
+                               const int num_row_per_block) {
+  if (threadIdx.x == 0) {
+    int gridx = 0;
+    int index = 0;
+    for (uint32_t bid = 0; bid < bsz; bid++) {
+      const int start_len = seq_lens_decoder[bid];
+      int seq_len = seq_lens_encoder[bid] + start_len % pad_len;
+      if (seq_lens_encoder[bid] == 0) {
+        seq_len = 0;
+      }
+      const int loop_times = div_up(seq_len, num_row_per_block);
+      for (uint32_t tile_id = 0; tile_id < loop_times; tile_id++) {
+        batch_ids[index] = bid;
+        tile_ids_per_batch[index++] = tile_id;
+      }
+      gridx += loop_times;
+    }
+    *num_blocks_x = gridx;
+  }
+}
+
+template <int THREADBLOCK_SIZE>
+__global__ void get_max_len_kv_ernel(int* max_seq_lens_out,
+                                  const int* seq_lens_this_time,
+                                  const int* seq_lens_decoder,
+                                  const int batch_size) {
+  const int tid = threadIdx.x;
+
+  
+  typedef cub::BlockReduce<int, THREADBLOCK_SIZE> BlockReduce;
+  __shared__ typename BlockReduce::TempStorage temp_storage;
+
+  int max_len_this_thread = 0;
+  for (int i = tid; i < batch_size; i += blockDim.x) {
+    if (seq_lens_decoder[i] == 0) continue;
+    max_len_this_thread = max(seq_lens_this_time[i] + seq_lens_decoder[i], max_len_this_thread);
+  }
+  int total = BlockReduce(temp_storage).Reduce(max_len_this_thread, MaxOp<int>());
+  if (tid == 0) {
+    *max_seq_lens_out = total;
+  }
+}
+
+std::vector<paddle::Tensor> GetBlockShapeAndSplitKVBlock(
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& max_enc_len_this_time,
+    const paddle::Tensor& max_dec_len_this_time,
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& cum_offsets,
+    const int group_size,
+    const int block_size,
+    const int decoder_step_token_num) {
+  paddle::Tensor encoder_batch_ids, encoder_tile_ids_per_batch, encoder_num_blocks_x_cpu,
+    kv_batch_ids, kv_tile_ids_per_batch, kv_num_blocks_x_cpu, decoder_batch_ids,
+    decoder_tile_ids_per_batch, decoder_num_blocks_x, decoder_num_blocks_x_cpu, decoder_chunk_size_device, decoder_chunk_size_cpu;
+  auto stream = seq_lens_this_time.stream();
+  int bsz = cum_offsets.shape()[0];
+  const int encoder_block_shape_q = get_encoder_block_shape_q();
+  const int decoder_block_shape_q = get_decoder_block_shape_q();
+
+  // max_len
+  auto max_len_kv =
+      GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_decoder.place());
+  get_max_len_kv_ernel<128><<<1, 128, 0, stream>>>(
+    max_len_kv.data<int>(),
+    seq_lens_this_time.data<int>(),
+    seq_lens_decoder.data<int>(),
+    bsz
+  );
+  auto max_len_kv_cpu =
+      max_len_kv.copy_to(paddle::CPUPlace(), false);
+
+  // decoder
+  int max_dec_len_this_time_data = max_dec_len_this_time.data<int>()[0];
+  if (max_dec_len_this_time_data > 0) {
+    const bool mla_use_tensorcore = GetMlaUseTensorcore();
+    if (mla_use_tensorcore && group_size <=64) {
+      const int set_chunk_size = get_mla_dec_chunk_size(bsz);
+      decoder_chunk_size_device =
+          GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_encoder.place());
+      decoder_num_blocks_x =
+          GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_encoder.place());
+
+      int device;
+      cudaGetDevice(&device);
+      int sm_cout;
+      cudaDeviceGetAttribute(&sm_cout, cudaDevAttrMultiProcessorCount, device);
+      constexpr int config_size = 12; // search space for chunk size:[64, 128, 256, ... 131072]
+
+      search_chunk_size_for_mla<config_size><<<1, 32, 0, stream>>>(
+        seq_lens_this_time.data<int>(),
+        seq_lens_encoder.data<int>(),
+        seq_lens_decoder.data<int>(),
+        decoder_num_blocks_x.data<int>(),
+        decoder_chunk_size_device.data<int>(),
+        bsz,
+        set_chunk_size,
+        block_size,
+        sm_cout);
+
+      decoder_chunk_size_cpu =
+        decoder_chunk_size_device.copy_to(paddle::CPUPlace(), false);
+      decoder_num_blocks_x_cpu =
+        decoder_num_blocks_x.copy_to(paddle::CPUPlace(), false);
+
+      const int chunk_size = decoder_chunk_size_cpu.data<int>()[0];
+      const int num_blocks = decoder_num_blocks_x_cpu.data<int>()[0];
+      decoder_batch_ids = 
+          GetEmptyTensor({num_blocks},
+                        paddle::DataType::INT32,
+                        seq_lens_encoder.place());
+      decoder_tile_ids_per_batch =
+          GetEmptyTensor({num_blocks},
+                        paddle::DataType::INT32,
+                        seq_lens_encoder.place());
+
+      split_block_for_mla<<<1, 32, 0, stream>>>(
+        seq_lens_this_time.data<int>(),
+        seq_lens_encoder.data<int>(),
+        seq_lens_decoder.data<int>(),
+        decoder_batch_ids.data<int>(),
+        decoder_tile_ids_per_batch.data<int>(),
+        bsz,
+        chunk_size
+      );
+    } else {
+      const uint32_t decoder_max_tile_size_per_bs_q =
+          div_up((decoder_step_token_num * group_size), decoder_block_shape_q);
+      decoder_batch_ids =
+          GetEmptyTensor({bsz * decoder_max_tile_size_per_bs_q},
+                        paddle::DataType::INT32,
+                        seq_lens_encoder.place());
+      decoder_tile_ids_per_batch =
+          GetEmptyTensor({bsz * decoder_max_tile_size_per_bs_q},
+                        paddle::DataType::INT32,
+                        seq_lens_encoder.place());
+      decoder_num_blocks_x =
+          GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_encoder.place());
+      split_q_block<<<1, 32, 0, stream>>>(seq_lens_this_time.data<int>(),
+                                          seq_lens_encoder.data<int>(),
+                                          decoder_batch_ids.data<int>(),
+                                          decoder_tile_ids_per_batch.data<int>(),
+                                          decoder_num_blocks_x.data<int>(),
+                                          bsz,
+                                          decoder_block_shape_q,
+                                          group_size);
+      decoder_num_blocks_x_cpu =
+        decoder_num_blocks_x.copy_to(paddle::CPUPlace(), false);
+      decoder_chunk_size_cpu =
+        paddle::full({1}, 131072, paddle::DataType::INT32, paddle::CPUPlace());
+    }
+  } else {
+    decoder_batch_ids =
+        paddle::full({1}, -1, paddle::DataType::INT32, paddle::GPUPlace());
+    decoder_tile_ids_per_batch =
+        paddle::full({1}, -1, paddle::DataType::INT32, paddle::GPUPlace());
+    decoder_chunk_size_cpu =
+        paddle::full({1}, 131072, paddle::DataType::INT32, paddle::CPUPlace());
+    decoder_num_blocks_x = 
+        paddle::full({1}, -1, paddle::DataType::INT32, paddle::GPUPlace());
+    decoder_num_blocks_x_cpu =
+        paddle::full({1}, -1, paddle::DataType::INT32, paddle::CPUPlace());
+  }
+
+  // encoder
+  int max_enc_len_this_time_data = max_enc_len_this_time.data<int>()[0];
+  if (max_enc_len_this_time_data > 0) {
+    const uint32_t encoder_max_tile_size_per_bs_q = div_up(
+        (max_enc_len_this_time_data * group_size), encoder_block_shape_q);
+    encoder_batch_ids =
+        GetEmptyTensor({bsz * encoder_max_tile_size_per_bs_q},
+                      paddle::DataType::INT32,
+                      seq_lens_encoder.place());
+    encoder_tile_ids_per_batch =
+        GetEmptyTensor({bsz * encoder_max_tile_size_per_bs_q},
+                      paddle::DataType::INT32,
+                      seq_lens_encoder.place());
+    auto encoder_num_blocks_x =
+        GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_encoder.place());
+    split_q_block<<<1, 32, 0, stream>>>(seq_lens_encoder.data<int>(),
+                                        nullptr,
+                                        encoder_batch_ids.data<int>(),
+                                        encoder_tile_ids_per_batch.data<int>(),
+                                        encoder_num_blocks_x.data<int>(),
+                                        bsz,
+                                        encoder_block_shape_q,
+                                        group_size);
+    encoder_num_blocks_x_cpu =
+        encoder_num_blocks_x.copy_to(paddle::CPUPlace(), false);
+
+    // kv
+    const uint32_t max_tile_size_per_bs_kv =
+        div_up(max_enc_len_this_time_data, block_size);
+    kv_batch_ids = GetEmptyTensor({bsz * max_tile_size_per_bs_kv},
+                                      paddle::DataType::INT32,
+                                      seq_lens_encoder.place());
+    kv_tile_ids_per_batch = GetEmptyTensor({bsz * max_tile_size_per_bs_kv},
+                                                paddle::DataType::INT32,
+                                                seq_lens_encoder.place());
+    auto kv_num_blocks_x =
+        GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_encoder.place());
+    split_kv_block<<<1, 32, 0, stream>>>(seq_lens_decoder.data<int>(),
+                                        seq_lens_encoder.data<int>(),
+                                        kv_batch_ids.data<int>(),
+                                        kv_tile_ids_per_batch.data<int>(),
+                                        kv_num_blocks_x.data<int>(),
+                                        bsz,
+                                        block_size,
+                                        block_size);
+    kv_num_blocks_x_cpu = kv_num_blocks_x.copy_to(paddle::CPUPlace(), false);
+  } else {
+    encoder_batch_ids =
+        paddle::full({1}, -1, paddle::DataType::INT32, paddle::GPUPlace());
+    encoder_tile_ids_per_batch =
+        paddle::full({1}, -1, paddle::DataType::INT32, paddle::GPUPlace());
+    encoder_num_blocks_x_cpu =
+        paddle::full({1}, -1, paddle::DataType::INT32, paddle::CPUPlace());
+    kv_batch_ids =
+        paddle::full({1}, -1, paddle::DataType::INT32, paddle::GPUPlace());
+    kv_tile_ids_per_batch =
+        paddle::full({1}, -1, paddle::DataType::INT32, paddle::GPUPlace());
+    kv_num_blocks_x_cpu = 
+        paddle::full({1}, -1, paddle::DataType::INT32, paddle::CPUPlace());
+  }
+  return {encoder_batch_ids,
+          encoder_tile_ids_per_batch,
+          encoder_num_blocks_x_cpu, /*cpu*/
+          kv_batch_ids,
+          kv_tile_ids_per_batch,
+          kv_num_blocks_x_cpu, /*cpu*/
+          decoder_batch_ids,
+          decoder_tile_ids_per_batch,
+          decoder_num_blocks_x,
+          decoder_num_blocks_x_cpu, /*cpu*/
+          decoder_chunk_size_cpu, /*cpu*/
+          max_len_kv_cpu /*cpu*/};
+}
+
+std::vector<paddle::DataType> GetBlockShapeAndSplitKVBlockInferDtype(
+    const paddle::DataType& seq_lens_encoder_dtype,
+    const paddle::DataType& seq_lens_decoder_dtype,
+    const paddle::DataType& max_enc_len_this_time_dtype,
+    const paddle::DataType& max_dec_len_this_time_dtype,
+    const paddle::DataType& seq_lens_this_time_dtype,
+    const paddle::DataType& cum_offsets_dtype) {
+  return {paddle::DataType::INT32,
+          paddle::DataType::INT32,
+          paddle::DataType::INT32,
+          paddle::DataType::INT32,
+          paddle::DataType::INT32,
+          paddle::DataType::INT32,
+          paddle::DataType::INT32,
+          paddle::DataType::INT32,
+          paddle::DataType::INT32,
+          paddle::DataType::INT32,
+          paddle::DataType::INT32,
+          paddle::DataType::INT32};
+}
+
+std::vector<std::vector<int64_t>> GetBlockShapeAndSplitKVBlockInferShape(
+    const std::vector<int64_t>& seq_lens_encoder_shape,
+    const std::vector<int64_t>& seq_lens_decoder_shape,
+    const std::vector<int64_t>& max_enc_len_this_time_shape,
+    const std::vector<int64_t>& max_dec_len_this_time_shape,
+    const std::vector<int64_t>& seq_lens_this_time_shape,
+    const std::vector<int64_t>& cum_offsets_shape) {
+  std::vector<int64_t> dynamic_shape = {-1};
+
+  return {dynamic_shape,
+          dynamic_shape,
+          {1},
+          dynamic_shape,
+          dynamic_shape,
+          {1},
+          dynamic_shape,
+          dynamic_shape,
+          {1},
+          {1},
+          {1},
+          {1}};
+}
+
+PD_BUILD_OP(get_block_shape_and_split_kv_block)
+    .Inputs({"seq_lens_encoder",
+             "seq_lens_decoder",
+             "max_enc_len_this_time",
+             "max_dec_len_this_time",
+             "seq_lens_this_time",
+             "cum_offsets"})
+    .Outputs({"encoder_batch_ids",
+              "encoder_tile_ids_per_batch",
+              "encoder_num_blocks",
+              "kv_batch_ids",
+              "kv_tile_ids_per_batch",
+              "kv_num_blocks",
+              "decoder_batch_ids",
+              "decoder_tile_ids_per_batch",
+              "decoder_num_blocks",
+              "decoder_num_blocks_cpu",
+              "decoder_chunk_size_cpu",
+              "max_len_kv"})
+    .Attrs({"group_size: int",
+            "block_size: int",
+            "decoder_step_token_num: int"})
+    .SetKernelFn(PD_KERNEL(GetBlockShapeAndSplitKVBlock))
+    .SetInferShapeFn(PD_INFER_SHAPE(GetBlockShapeAndSplitKVBlockInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(GetBlockShapeAndSplitKVBlockInferDtype));

csrc/gpu/append_attn/mem_util.cuh

@@ -0,0 +1,318 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#pragma once
+
+#include <cuda_runtime.h>
+#include <stdint.h>
+
+enum class SharedMemFillMode { kFillZero, kNoFill };
+
+enum class PrefetchMode { kNoPrefetch, kPrefetch };
+
+template <typename T>
+__device__ __forceinline__ void ldmatrix_m8n8x4_impl(uint32_t* R, T* smem_ptr) {
+  uint32_t smem_int_ptr =
+      static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  asm volatile(
+      "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
+      : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3])
+      : "r"(smem_int_ptr));
+}
+
+template <typename T>
+__device__ __forceinline__ void ldmatrix_m8n8x4_trans_impl(uint32_t* R,
+                                                           T* smem_ptr) {
+  uint32_t smem_int_ptr =
+      static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  asm volatile(
+      "ldmatrix.sync.aligned.trans.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
+      : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3])
+      : "r"(smem_int_ptr));
+}
+
+__device__ __forceinline__ void commit_group() {
+  asm volatile("cp.async.commit_group;\n" ::);
+}
+
+template <size_t n>
+__device__ __forceinline__ void wait_group() {
+  asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
+}
+
+template <PrefetchMode prefetch_mode, typename T>
+__device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) {
+  uint32_t smem_int_ptr =
+      static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
+    asm volatile(
+        "cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(
+            smem_int_ptr),
+        "l"(gmem_ptr),
+        "n"(16),
+        "r"(16));
+  } else {
+    asm volatile(
+        "cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
+        "l"(gmem_ptr),
+        "n"(16),
+        "r"(16));
+  }
+}
+
+template <PrefetchMode prefetch_mode, SharedMemFillMode fill_mode, typename T>
+__device__ __forceinline__ void pred_load_128b(T* smem_ptr,
+                                               const T* gmem_ptr,
+                                               bool predicate) {
+  uint32_t smem_int_ptr =
+      static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
+    int src_in_bytes = predicate ? 16 : 0;
+    if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
+      asm volatile(
+          "cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(
+              smem_int_ptr),
+          "l"(gmem_ptr),
+          "n"(16),
+          "r"(src_in_bytes));
+    } else {
+      asm volatile(
+          "cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
+          "l"(gmem_ptr),
+          "n"(16),
+          "r"(src_in_bytes));
+    }
+  } else {
+    if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
+      asm volatile(
+          "{\n"
+          " .reg .pred p;\n"
+          " setp.ne.b32 p, %0, 0;\n"
+          " @p cp.async.cg.shared.global.L2::128B [%1], [%2], %3;\n"
+          "}\n" ::"r"((int)predicate),
+          "r"(smem_int_ptr),
+          "l"(gmem_ptr),
+          "n"(16));
+    } else {
+      asm volatile(
+          "{\n"
+          " .reg .pred p;\n"
+          " setp.ne.b32 p, %0, 0;\n"
+          " @p cp.async.cg.shared.global [%1], [%2], %3;\n"
+          "}\n" ::"r"((int)predicate),
+          "r"(smem_int_ptr),
+          "l"(gmem_ptr),
+          "n"(16));
+    }
+  }
+}
+
+template <PrefetchMode prefetch_mode, SharedMemFillMode fill_mode, typename T>
+__device__ __forceinline__ void pred_load_64b(T* smem_ptr,
+                                              const T* gmem_ptr,
+                                              bool predicate) {
+  uint32_t smem_int_ptr =
+      static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
+    int src_in_bytes = predicate ? 8 : 0;
+    asm volatile(
+        "cp.async.ca.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
+        "l"(gmem_ptr),
+        "n"(8),
+        "r"(src_in_bytes));
+  } else {
+    asm volatile(
+        "{\n"
+        " .reg .pred p;\n"
+        " setp.ne.b32 p, %0, 0;\n"
+        " @p cp.async.ca.shared.global [%1], [%2], %3;\n"
+        "}\n" ::"r"((int)predicate),
+        "r"(smem_int_ptr),
+        "l"(gmem_ptr),
+        "n"(8));
+  }
+}
+
+template <PrefetchMode prefetch_mode, SharedMemFillMode fill_mode, typename T>
+__device__ __forceinline__ void pred_load_32b(T* smem_ptr,
+                                              const T* gmem_ptr,
+                                              bool predicate) {
+  uint32_t smem_int_ptr =
+      static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
+    int src_in_bytes = predicate ? 4 : 0;
+    asm volatile(
+        "cp.async.ca.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
+        "l"(gmem_ptr),
+        "n"(4),
+        "r"(src_in_bytes));
+  } else {
+    asm volatile(
+        "{\n"
+        " .reg .pred p;\n"
+        " setp.ne.b32 p, %0, 0;\n"
+        " @p cp.async.ca.shared.global [%1], [%2], %3;\n"
+        "}\n" ::"r"((int)predicate),
+        "r"(smem_int_ptr),
+        "l"(gmem_ptr),
+        "n"(4));
+  }
+}
+
+template <size_t num_bits, PrefetchMode prefetch_mode, typename T>
+__device__ __forceinline__ void load(T* smem_ptr, const T* gmem_ptr) {
+  static_assert(num_bits == 128, "num_bits must be 128");
+  load_128b<prefetch_mode>(smem_ptr, gmem_ptr);
+}
+
+template <size_t num_bits,
+          PrefetchMode prefetch_mode,
+          SharedMemFillMode fill_mode,
+          typename T>
+__device__ __forceinline__ void pred_load(T* smem_ptr,
+                                          const T* gmem_ptr,
+                                          bool predicate) {
+  static_assert(num_bits == 128 || num_bits == 64 || num_bits == 32,
+                "num_bits must be 128, 64 or 32.");
+  if constexpr (num_bits == 128) {
+    pred_load_128b<prefetch_mode, fill_mode>(smem_ptr, gmem_ptr, predicate);
+  } else if constexpr (num_bits == 64) {
+    pred_load_64b<prefetch_mode, fill_mode>(smem_ptr, gmem_ptr, predicate);
+  } else if constexpr (num_bits == 32) {
+    pred_load_32b<prefetch_mode, fill_mode>(smem_ptr, gmem_ptr, predicate);
+  }
+}
+
+using b32_t = uint32_t;
+using b64_t = uint2;
+using b128_t = uint4;
+
+template <typename T>
+constexpr __host__ __device__ __forceinline__ uint32_t num_elems_per_128b() {
+  return sizeof(b128_t) / sizeof(T);
+}
+
+struct smem_t {
+  // The base pointer.
+  b128_t* base;
+  __device__ __forceinline__ smem_t() : base(nullptr) {}
+  template <typename T>
+  __device__ __forceinline__ smem_t(T* base) : base((b128_t*)base) {}
+
+
+  template <uint32_t stride, uint32_t inv_stride = 0>
+  static __device__ __forceinline__ uint32_t get_permuted_offset(uint32_t i,
+                                                                 uint32_t j) {
+    if constexpr (inv_stride <= 1) {
+      return i * stride + (j ^ (i % 8));
+    } else {
+      return i / inv_stride * 8 + ((j + (i % inv_stride) * stride)) ^
+             ((i / inv_stride) % 8);
+    }
+  }
+
+  template <uint32_t step_size, uint32_t row_stride = 8>
+  static __device__ __forceinline__ uint32_t
+  advance_offset_by_column(uint32_t offset, uint32_t step_idx) {
+    if constexpr (row_stride == 2) {
+      static_assert(step_size == 2, "Unsupported step size");
+      return offset + step_size;
+    } else if constexpr (row_stride == 4) {
+      static_assert(step_size == 2 || step_size == 4, "Unsupported step size");
+      if constexpr (step_size == 2) {
+        return (offset ^ 0x2) + (step_idx % 2 == 1) * 4;
+      } else {
+        return offset + step_size;
+      }
+    } else {
+      static_assert(step_size == 2 || step_size == 4 || step_size % 8 == 0,
+                    "Unsupported step size");
+      if constexpr (step_size == 2) {
+        return (offset ^ (0x2 + (0x4 * (step_idx % 2 == 1)))) +
+               (step_idx % 4 == 3) * 8;
+      } else if constexpr (step_size == 4) {
+        return (offset ^ 0x4) + (step_idx % 2 == 1) * 8;
+      } else {
+        // step_size % 8 == 0
+        return offset + step_size;
+      }
+    }
+  }
+
+  template <uint32_t step_size, uint32_t row_stride>
+  static __device__ __forceinline__ uint32_t
+  advance_offset_by_row(uint32_t offset) {
+    if constexpr (row_stride == 2) {
+      static_assert(step_size == 16 || step_size % 32 == 0,
+                    "Unsupported step size");
+      if constexpr (step_size == 16) {
+        return (offset ^ 0x4) + step_size * row_stride;
+      } else {
+        // step_size % 32 == 0
+        return offset + step_size * row_stride;
+      }
+    } else if constexpr (row_stride == 4) {
+      static_assert(step_size == 8 || step_size % 16 == 0,
+                    "Unsupported step size");
+      if constexpr (step_size == 8) {
+        return (offset ^ 0x4) + step_size * row_stride;
+      } else {
+        // step_size % 16 == 0
+        return offset + step_size * row_stride;
+      }
+    } else {
+      static_assert(step_size == 4 || step_size % 8 == 0,
+                    "Unsupported step size");
+      if constexpr (step_size == 4) {
+        return (offset ^ 0x4) + step_size * row_stride;
+      } else {
+        // step_size % 8 == 0
+        return offset + step_size * row_stride;
+      }
+    }
+  }
+
+  __device__ __forceinline__ void ldmatrix_m8n8x4(uint32_t offset,
+                                                  uint32_t* R) {
+    b128_t* smem_ptr = base + offset;
+    ldmatrix_m8n8x4_impl(R, smem_ptr);
+  }
+
+  __device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t offset,
+                                                        uint32_t* R) {
+    b128_t* smem_ptr = base + offset;
+    ldmatrix_m8n8x4_trans_impl(R, smem_ptr);
+  }
+
+  template <SharedMemFillMode fill_mode, typename T>
+  __device__ __forceinline__ void load_128b_async(uint32_t offset,
+                                                  const T* gptr,
+                                                  bool predicate) {
+    b128_t* smem_ptr = base + offset;
+    pred_load_128b<PrefetchMode::kPrefetch, fill_mode>(
+        smem_ptr, reinterpret_cast<const b128_t*>(gptr), predicate);
+  }
+
+  template <typename T>
+  __device__ __forceinline__ void load_128b_async(uint32_t offset,
+                                                  const T* gptr) {
+    b128_t* smem_ptr = base + offset;
+    load_128b<PrefetchMode::kPrefetch>(smem_ptr,
+                                       reinterpret_cast<const b128_t*>(gptr));
+  }
+
+  template <typename T>
+  __device__ __forceinline__ void store_128b(uint32_t offset, T* gptr) {
+    *reinterpret_cast<b128_t*>(gptr) = *(base + offset);
+  }
+};

csrc/gpu/append_attn/mla_cache_kernel.cu

@@ -0,0 +1,291 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#pragma once
+
+#include "mla_cache_kernel.cuh"
+
+template <paddle::DataType T>
+std::vector<paddle::Tensor> PrefillMLAWriteCache(
+                    const AppendAttnMetaData& meta_data,
+                    const paddle::Tensor& kv_nope,
+                    const paddle::Tensor& kv_pe,
+                    const paddle::Tensor& seq_lens,
+                    const paddle::Tensor& seq_lens_decoder,
+                    const paddle::Tensor& padding_offsets,
+                    const paddle::Tensor& cum_offsets,
+                    const paddle::Tensor& block_tables,
+                    const int max_seq_len,
+                    cudaStream_t& stream,
+                    paddle::Tensor* kv_cache) {
+  typedef PDTraits<T> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
+
+  auto max_blocks_per_seq = meta_data.max_blocks_per_seq;
+  auto num_tokens = meta_data.token_nums;
+  auto block_size = meta_data.block_size;
+  auto nope_size = meta_data.head_dims_v;
+  auto all_size = meta_data.head_dims;
+  int pe_size = all_size - nope_size;
+  auto kv_num_heads = meta_data.kv_num_heads;
+  const uint32_t elem_nums = num_tokens * kv_num_heads * all_size;
+
+  constexpr int PackSize = 16 / sizeof(DataType_);
+  const int pack_num = elem_nums / PackSize;
+  const int blocksize = 128;
+  int grid_size = 1;
+  GetNumBlocks<128>(pack_num, &grid_size);
+
+  prefill_absorb_cache_kernel<DataType_, PackSize>
+      <<<grid_size, blocksize, 0, stream>>>(
+          reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_nope.data<data_t>())),
+          reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_pe.data<data_t>())),
+          reinterpret_cast<DataType_*>(kv_cache->data<data_t>()),
+          block_tables.data<int>(),
+          padding_offsets.data<int>(),
+          cum_offsets.data<int>(),
+          seq_lens.data<int>(),
+          seq_lens_decoder.data<int>(),
+          max_seq_len,
+          max_blocks_per_seq,
+          kv_num_heads,
+          nope_size,
+          pe_size,
+          block_size,
+          elem_nums);
+  return {};
+}
+
+std::vector<paddle::Tensor> PrefillMLAWriteCacheKernel(
+    const paddle::Tensor& kv_nope,
+    const paddle::Tensor& kv_pe,
+    const paddle::Tensor& kv_cache,
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const std::string& cache_quant_type_str,
+    const int max_seq_len) {
+  cudaStream_t stream = kv_pe.stream();
+  AppendAttnMetaData meta_data;
+  const auto& kv_nope_dims = kv_nope.dims();
+  const auto& kv_pe_dims = kv_pe.dims();
+  const auto& kv_cache_dims = kv_cache.dims();
+  meta_data.kv_num_heads = kv_cache_dims[1];
+  const auto nope_size = kv_nope_dims[kv_nope_dims.size() - 1] / meta_data.kv_num_heads;
+  meta_data.token_nums = kv_nope_dims[0];
+  meta_data.head_dims = kv_cache_dims[3];
+  meta_data.head_dims_v = nope_size;
+
+  meta_data.max_blocks_per_seq = block_tables.dims()[1];
+  meta_data.block_size = kv_cache_dims[2];
+  meta_data.batch_size = cum_offsets.dims()[0];
+  switch (kv_pe.dtype()) {
+    case paddle::DataType::BFLOAT16: {
+      return PrefillMLAWriteCache<paddle::DataType::BFLOAT16>(meta_data,
+                              kv_nope,
+                              kv_pe,
+                              seq_lens,
+                              seq_lens_decoder,
+                              padding_offsets,
+                              cum_offsets,
+                              block_tables,
+                              max_seq_len,
+                              stream,
+                              const_cast<paddle::Tensor*>(&kv_cache));
+    }
+    case paddle::DataType::FLOAT16: {
+      return PrefillMLAWriteCache<paddle::DataType::FLOAT16>(meta_data,
+                              kv_nope,
+                              kv_pe,
+                              seq_lens,
+                              seq_lens_decoder,
+                              padding_offsets,
+                              cum_offsets,
+                              block_tables,
+                              max_seq_len,
+                              stream,
+                              const_cast<paddle::Tensor*>(&kv_cache));
+    }
+  }
+  return {};
+}
+
+template <paddle::DataType T>
+std::vector<paddle::Tensor> DecodeMLAWriteCache(
+                    const AppendAttnMetaData& meta_data,
+                    const paddle::Tensor& kv_nope,
+                    const paddle::Tensor& kv_pe,
+                    const paddle::Tensor& seq_lens,
+                    const paddle::Tensor& seq_lens_encoder,
+                    const paddle::Tensor& padding_offsets,
+                    const paddle::Tensor& cum_offsets,
+                    const paddle::Tensor& block_tables,
+                    const int max_seq_len,
+                    const bool speculate_decoder,
+                    cudaStream_t& stream,
+                    paddle::Tensor* kv_cache) {
+  typedef PDTraits<T> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
+  
+  auto max_blocks_per_seq = meta_data.max_blocks_per_seq;
+  auto bsz = meta_data.batch_size;
+  auto token_num = meta_data.token_nums;
+  auto block_size = meta_data.block_size;
+  auto nope_size = meta_data.head_dims_v;
+  auto all_size = meta_data.head_dims;
+  int pe_size = all_size - nope_size;
+  auto kv_num_heads = meta_data.kv_num_heads;
+  constexpr int PackSize = 16 / sizeof(DataType_);
+  const int blocksize = 128;
+  int grid_size = 1;
+
+
+  if (speculate_decoder) {
+    const uint32_t elem_nums = token_num * kv_num_heads * all_size;
+    const int pack_num = elem_nums / PackSize;
+    GetNumBlocks<128>(pack_num, &grid_size);
+    speculate_decode_absorb_cache_kernel<DataType_, PackSize>
+        <<<grid_size, blocksize, 0, stream>>>(
+            reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_nope.data<data_t>())),
+            reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_pe.data<data_t>())),
+            reinterpret_cast<DataType_*>(kv_cache->data<data_t>()),
+            block_tables.data<int>(),
+            padding_offsets.data<int>(),
+            cum_offsets.data<int>(),
+            seq_lens.data<int>(),
+            seq_lens_encoder.data<int>(),
+            max_seq_len,
+            max_blocks_per_seq,
+            kv_num_heads,
+            nope_size,
+            pe_size,
+            block_size,
+            elem_nums);
+  } else {
+    const uint32_t elem_nums = bsz * kv_num_heads * all_size;
+    const int pack_num = elem_nums / PackSize;
+    GetNumBlocks<128>(pack_num, &grid_size);
+    decode_absorb_cache_kernel<DataType_, PackSize>
+        <<<grid_size, blocksize, 0, stream>>>(
+            reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_nope.data<data_t>())),
+            reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_pe.data<data_t>())),
+            reinterpret_cast<DataType_*>(kv_cache->data<data_t>()),
+            block_tables.data<int>(),
+            cum_offsets.data<int>(),
+            seq_lens.data<int>(),
+            seq_lens_encoder.data<int>(),
+            max_seq_len,
+            max_blocks_per_seq,
+            kv_num_heads,
+            nope_size,
+            pe_size,
+            block_size,
+            elem_nums);
+  }
+  return {};
+}
+
+std::vector<paddle::Tensor> DecodeMLAWriteCacheKernel(
+    const paddle::Tensor& kv_nope,
+    const paddle::Tensor& kv_pe,
+    const paddle::Tensor& kv_cache,
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const std::string& cache_quant_type_str,
+    const int max_seq_len,
+    const bool speculate_decoder) {
+  cudaStream_t stream = kv_pe.stream();
+  AppendAttnMetaData meta_data;
+  const auto& kv_nope_dims = kv_nope.dims();
+  const auto& kv_pe_dims = kv_pe.dims();
+  const auto& kv_cache_dims = kv_cache.dims();
+  meta_data.kv_num_heads = kv_cache_dims[1];
+  const auto nope_size = kv_nope_dims[kv_nope_dims.size() - 1] / meta_data.kv_num_heads;
+  meta_data.token_nums = kv_nope_dims[0];
+  meta_data.head_dims = kv_cache_dims[3];
+  meta_data.head_dims_v = nope_size;
+
+  meta_data.max_blocks_per_seq = block_tables.dims()[1];
+  meta_data.block_size = kv_cache_dims[2];
+  meta_data.batch_size = cum_offsets.dims()[0];
+  switch (kv_pe.dtype()) {
+    case paddle::DataType::BFLOAT16: {
+      return DecodeMLAWriteCache<paddle::DataType::BFLOAT16>(meta_data,
+                              kv_nope,
+                              kv_pe,
+                              seq_lens,
+                              seq_lens_encoder,
+                              padding_offsets,
+                              cum_offsets,
+                              block_tables,
+                              max_seq_len,
+                              speculate_decoder,
+                              stream,
+                              const_cast<paddle::Tensor*>(&kv_cache));
+    }
+    case paddle::DataType::FLOAT16: {
+      return DecodeMLAWriteCache<paddle::DataType::FLOAT16>(meta_data,
+                              kv_nope,
+                              kv_pe,
+                              seq_lens,
+                              seq_lens_encoder,
+                              padding_offsets,
+                              cum_offsets,
+                              block_tables,
+                              max_seq_len,
+                              speculate_decoder,
+                              stream,
+                              const_cast<paddle::Tensor*>(&kv_cache));
+    }
+  }
+  return {};
+}
+
+
+PD_BUILD_OP(prefill_mla_write_cache)
+    .Inputs({"kv_nope",
+             "kv_pe",
+             "kv_cache",
+             "seq_lens",
+             "seq_lens_decoder",
+             "padding_offsets",
+             "cum_offsets",
+             "block_tables"})
+    .Outputs({"kv_cache_out"})
+    .SetInplaceMap({{"kv_cache", "kv_cache_out"}})
+    .Attrs({"cache_quant_type_str: std::string",
+            "max_seq_len: int"})
+    .SetKernelFn(PD_KERNEL(PrefillMLAWriteCacheKernel));
+
+PD_BUILD_OP(decode_mla_write_cache)
+    .Inputs({"kv_nope",
+             "kv_pe",
+             "kv_cache",
+             "seq_lens",
+             "seq_lens_encoder",
+             "padding_offsets",
+             "cum_offsets",
+             "block_tables"})
+    .Outputs({"kv_cache_out"})
+    .SetInplaceMap({{"kv_cache", "kv_cache_out"}})
+    .Attrs({"cache_quant_type_str: std::string",
+            "max_seq_len: int",
+            "speculate_decoder: bool"})
+    .SetKernelFn(PD_KERNEL(DecodeMLAWriteCacheKernel));

csrc/gpu/append_attn/mla_cache_kernel.cuh

@@ -0,0 +1,242 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#pragma once
+
+#include "helper.h"
+#include "mem_util.cuh"
+#include "utils.cuh"
+
+template <typename T, int VecSize = 1>
+__global__ void decode_absorb_cache_kernel(
+    const T* __restrict__ kv_nope,  // [bsz, kv_num_heads, pe_size] 512
+    const T* __restrict__ kv_pe,  // [bsz, kv_num_heads, nope_size] 64
+    T* __restrict__ kv_cache,    // [num_blocks, kv_num_heads, block_size,
+                                  // nope_size]                             
+    const int* __restrict__ block_tables,     // [bsz, max_blocks_per_seq]
+    const int* __restrict__ cum_offsets,
+    const int* __restrict__ seq_lens,          // [bsz]
+    const int* __restrict__ seq_lens_encoder,  // [bsz]
+    const int max_seq_len,
+    const int max_blocks_per_seq,
+    const int kv_num_heads,
+    const int nope_size,
+    const int pe_size,
+    const int block_size,
+    const uint32_t elem_cnt) {
+  using LoadT = AlignedVector<T, VecSize>;
+  constexpr int HalfVecSize = VecSize / 2;
+  LoadT src_vec;
+
+  int64_t global_thread_idx = blockDim.x * blockIdx.x + threadIdx.x;
+  const uint32_t nope_hidden_size = kv_num_heads * nope_size;
+  const uint32_t pe_hidden_size = kv_num_heads * pe_size;
+  const uint32_t all_size = nope_size + pe_size;
+  const int64_t hidden_size = nope_hidden_size + pe_hidden_size;
+
+  for (int32_t linear_index = global_thread_idx * VecSize,
+               step = gridDim.x * blockDim.x * VecSize;
+       linear_index < elem_cnt;
+       linear_index += step) {
+    const int ori_bi = linear_index / hidden_size;
+    const int bias = linear_index % hidden_size;
+    const int start_token_idx = ori_bi * max_seq_len - cum_offsets[ori_bi];
+    if (seq_lens_encoder[ori_bi] > 0) return;
+    const int write_seq_id = seq_lens[ori_bi];
+    
+    if (write_seq_id == 0) continue;
+
+    const int* block_table_now = nullptr;
+
+    block_table_now = block_tables + ori_bi * max_blocks_per_seq;
+    const int block_idx = block_table_now[write_seq_id / block_size];
+    const int block_offset = write_seq_id % block_size;
+
+    if (bias < nope_hidden_size) { // pe
+      const uint32_t inner_bias = bias;
+      const uint32_t hi = inner_bias / nope_size;
+      const uint32_t h_bias = inner_bias % nope_size;
+      const uint32_t tgt_idx = block_idx * kv_num_heads * block_size * all_size +
+                             hi * block_size * all_size +
+                             block_offset * all_size + h_bias;
+      const uint32_t ori_idx =
+          start_token_idx * nope_hidden_size + inner_bias;
+      Load<T, VecSize>(&kv_nope[ori_idx], &src_vec);
+      Store<T, VecSize>(src_vec, &kv_cache[tgt_idx]);
+    } else {
+      const uint32_t inner_bias = bias - nope_hidden_size;
+      const uint32_t hi = inner_bias / pe_size;
+      const uint32_t h_bias = inner_bias % pe_size;
+      const uint32_t tgt_idx = block_idx * kv_num_heads * block_size * all_size +
+                             hi * block_size * all_size +
+                             block_offset * all_size + nope_size + h_bias;
+      const uint32_t ori_idx =
+          start_token_idx * pe_hidden_size + inner_bias;
+      Load<T, VecSize>(&kv_pe[ori_idx], &src_vec);
+      Store<T, VecSize>(src_vec, &kv_cache[tgt_idx]);
+    }
+  }
+}
+
+template <typename T, int VecSize = 1>
+__global__ void speculate_decode_absorb_cache_kernel(
+    const T* __restrict__ kv_nope,  // [bsz, kv_num_heads, pe_size] 512
+    const T* __restrict__ kv_pe,  // [bsz, kv_num_heads, nope_size] 64
+    T* __restrict__ kv_cache,    // [num_blocks, kv_num_heads, block_size,
+                                  // nope_size]                             
+    const int* __restrict__ block_tables,     // [bsz, max_blocks_per_seq]
+    const int* __restrict__ padding_offsets,
+    const int* __restrict__ cum_offsets,
+    const int* __restrict__ seq_lens,          // [bsz]
+    const int* __restrict__ seq_lens_encoder,  // [bsz]
+    const int max_seq_len,
+    const int max_blocks_per_seq,
+    const int kv_num_heads,
+    const int nope_size,
+    const int pe_size,
+    const int block_size,
+    const uint32_t elem_cnt) {
+  using LoadT = AlignedVector<T, VecSize>;
+  constexpr int HalfVecSize = VecSize / 2;
+  LoadT src_vec;
+
+  int64_t global_thread_idx = blockDim.x * blockIdx.x + threadIdx.x;
+  const uint32_t nope_hidden_size = kv_num_heads * nope_size;
+  const uint32_t pe_hidden_size = kv_num_heads * pe_size;
+  const uint32_t all_size = nope_size + pe_size;
+  const int64_t hidden_size = nope_hidden_size + pe_hidden_size;
+
+  for (int32_t linear_index = global_thread_idx * VecSize,
+               step = gridDim.x * blockDim.x * VecSize;
+       linear_index < elem_cnt;
+       linear_index += step) {
+    const int token_id = linear_index / hidden_size;
+    const int ori_bi = (token_id + padding_offsets[token_id]) / max_seq_len;
+    if (seq_lens[ori_bi] == 0) continue;
+    const int bias = linear_index % hidden_size;
+    const int start_token_idx = ori_bi * max_seq_len - cum_offsets[ori_bi];
+    const int write_seq_id = 
+        seq_lens[ori_bi] + token_id - start_token_idx;
+    if (write_seq_id == 0) continue;
+
+    const int* block_table_now = nullptr;
+
+    block_table_now = block_tables + ori_bi * max_blocks_per_seq;
+    const int block_idx = block_table_now[write_seq_id / block_size];
+    const int block_offset = write_seq_id % block_size;
+    if (block_idx < 0) {
+      printf(
+          "Fatal Error!!!, block idx %d when write_seq_id is %d\n some key var "
+          "%d %d %d %d\n",
+          block_idx,
+          write_seq_id,
+          ori_bi,
+          seq_lens[ori_bi],
+          token_id,
+          cum_offsets[ori_bi]);
+    }
+    if (bias < nope_hidden_size) { // pe
+      const uint32_t inner_bias = bias;
+      const uint32_t hi = inner_bias / nope_size;
+      const uint32_t h_bias = inner_bias % nope_size;
+      const uint32_t tgt_idx = block_idx * kv_num_heads * block_size * all_size +
+                             hi * block_size * all_size +
+                             block_offset * all_size + h_bias;
+      const uint32_t ori_idx =
+          token_id * nope_hidden_size + inner_bias;
+      Load<T, VecSize>(&kv_nope[ori_idx], &src_vec);
+      Store<T, VecSize>(src_vec, &kv_cache[tgt_idx]);
+    } else {
+      const uint32_t inner_bias = bias - nope_hidden_size;
+      const uint32_t hi = inner_bias / pe_size;
+      const uint32_t h_bias = inner_bias % pe_size;
+      const uint32_t tgt_idx = block_idx * kv_num_heads * block_size * all_size +
+                             hi * block_size * all_size +
+                             block_offset * all_size + nope_size + h_bias;
+      const uint32_t ori_idx =
+          token_id * pe_hidden_size + inner_bias;
+      Load<T, VecSize>(&kv_pe[ori_idx], &src_vec);
+      Store<T, VecSize>(src_vec, &kv_cache[tgt_idx]);
+    }
+  }
+}
+
+template <typename T, int VecSize = 1>
+__global__ void prefill_absorb_cache_kernel(
+    const T* __restrict__ kv_nope,  // [bsz, kv_num_heads, pe_size] 512
+    const T* __restrict__ kv_pe,  // [bsz, kv_num_heads, nope_size] 64
+    T* __restrict__ kv_cache,    // [num_blocks, kv_num_heads, block_size,
+                                  // nope_size]                             
+    const int* __restrict__ block_tables,     // [bsz, max_blocks_per_seq]
+    const int* __restrict__ padding_offsets,
+    const int* __restrict__ cum_offsets,
+    const int* __restrict__ seq_lens,          // [bsz]
+    const int* __restrict__ seq_lens_decoder,  // [bsz]
+    const int max_seq_len,
+    const int max_blocks_per_seq,
+    const int kv_num_heads,
+    const int nope_size,
+    const int pe_size,
+    const int block_size,
+    const uint32_t elem_cnt) {
+  using LoadT = AlignedVector<T, VecSize>;
+  LoadT src_vec;
+
+  int64_t global_thread_idx = blockDim.x * blockIdx.x + threadIdx.x;
+  const uint32_t nope_hidden_size = kv_num_heads * nope_size;
+  const uint32_t pe_hidden_size = kv_num_heads * pe_size;
+  const uint32_t all_size = nope_size + pe_size;
+  const int64_t hidden_size = nope_hidden_size + pe_hidden_size;
+
+  for (int32_t linear_index = global_thread_idx * VecSize,
+               step = gridDim.x * blockDim.x * VecSize;
+       linear_index < elem_cnt;
+       linear_index += step) {
+    const uint32_t token_idx = linear_index / hidden_size;
+    const uint32_t bias = linear_index % hidden_size;
+    const uint32_t ori_token_idx = token_idx + padding_offsets[token_idx];
+    const uint32_t ori_bi = ori_token_idx / max_seq_len;
+    if (seq_lens[ori_bi] == 0) continue;
+    const uint32_t ori_seq_id =
+        ori_token_idx % max_seq_len + seq_lens_decoder[ori_bi];  
+
+    const int* block_table_now = nullptr;
+    block_table_now = block_tables + ori_bi * max_blocks_per_seq;
+    const uint32_t block_idx = block_table_now[ori_seq_id / block_size];
+    const uint32_t block_offset = ori_seq_id % block_size;
+
+    if (bias < nope_hidden_size) { // pe
+      const uint32_t inner_bias = bias;
+      const uint32_t hi = inner_bias / nope_size;
+      const uint32_t h_bias = inner_bias % nope_size;
+      const uint32_t tgt_idx = block_idx * kv_num_heads * block_size * all_size +
+                             hi * block_size * all_size +
+                             block_offset * all_size + h_bias;
+      const uint32_t ori_idx =
+          token_idx * nope_hidden_size + inner_bias;
+      Load<T, VecSize>(&kv_nope[ori_idx], &src_vec);
+      Store<T, VecSize>(src_vec, &kv_cache[tgt_idx]);
+    } else {
+      const uint32_t inner_bias = bias - nope_hidden_size;
+      const uint32_t hi = inner_bias / pe_size;
+      const uint32_t h_bias = inner_bias % pe_size;
+      const uint32_t tgt_idx = block_idx * kv_num_heads * block_size * all_size +
+                             hi * block_size * all_size +
+                             block_offset * all_size + nope_size + h_bias;
+      const uint32_t ori_idx =
+          token_idx * pe_hidden_size + inner_bias;
+      Load<T, VecSize>(&kv_pe[ori_idx], &src_vec);
+      Store<T, VecSize>(src_vec, &kv_cache[tgt_idx]);
+    }
+  }
+}

csrc/gpu/append_attn/mma_tensor_op.cuh

@@ -0,0 +1,296 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#pragma once
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+enum class MMAMode {
+  kInit = 0U,
+  kInplaceUpdate = 1U,
+};
+
+template <typename T, MMAMode mma_mode = MMAMode::kInplaceUpdate>
+__device__ __forceinline__ void mma_sync_m16n16k32_row_col_i8i8i32(
+    int* C,         // 8
+    uint32_t* A,    // 4
+    uint32_t* B) {  // 4
+  if constexpr (mma_mode == MMAMode::kInit) {
+    asm volatile(
+        "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
+        "{%0,  %1,  %2,  %3},"
+        "{%4,  %5,  %6,  %7},"
+        "{%8,  %9},"
+        "{%10, %11, %12, %13};\n"
+        : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+        : "r"(A[0]),
+          "r"(A[1]),
+          "r"(A[2]),
+          "r"(A[3]),
+          "r"(B[0]),
+          "r"(B[1]),
+          "r"(0),
+          "r"(0),
+          "r"(0),
+          "r"(0));
+    asm volatile(
+        "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
+        "{%0,  %1,  %2,  %3},"
+        "{%4,  %5,  %6,  %7},"
+        "{%8,  %9},"
+        "{%10, %11, %12, %13};\n"
+        : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
+        : "r"(A[0]),
+          "r"(A[1]),
+          "r"(A[2]),
+          "r"(A[3]),
+          "r"(B[2]),
+          "r"(B[3]),
+          "r"(0),
+          "r"(0),
+          "r"(0),
+          "r"(0));
+  } else {
+    asm volatile(
+        "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
+        "{%0,  %1,  %2,  %3},"
+        "{%4,  %5,  %6,  %7},"
+        "{%8,  %9},"
+        "{%10, %11, %12, %13};\n"
+        : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+        : "r"(A[0]),
+          "r"(A[1]),
+          "r"(A[2]),
+          "r"(A[3]),
+          "r"(B[0]),
+          "r"(B[1]),
+          "r"(C[0]),
+          "r"(C[1]),
+          "r"(C[2]),
+          "r"(C[3]));
+    asm volatile(
+        "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
+        "{%0,  %1,  %2,  %3},"
+        "{%4,  %5,  %6,  %7},"
+        "{%8,  %9},"
+        "{%10, %11, %12, %13};\n"
+        : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
+        : "r"(A[0]),
+          "r"(A[1]),
+          "r"(A[2]),
+          "r"(A[3]),
+          "r"(B[2]),
+          "r"(B[3]),
+          "r"(C[4]),
+          "r"(C[5]),
+          "r"(C[6]),
+          "r"(C[7]));
+  }
+}
+
+template <typename T, MMAMode mma_mode = MMAMode::kInplaceUpdate>
+__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(
+    float* C, uint32_t* A, uint32_t* B) {
+  if constexpr (mma_mode == MMAMode::kInit) {
+    if constexpr (std::is_same<T, half>::value) {  // fp16
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+          "{%0,  %1,  %2,  %3},"
+          "{%4,  %5,  %6,  %7},"
+          "{%8,  %9},"
+          "{%10, %11, %12, %13};\n"
+          : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+          : "r"(A[0]),
+            "r"(A[1]),
+            "r"(A[2]),
+            "r"(A[3]),
+            "r"(B[0]),
+            "r"(B[1]),
+            "f"(0.f),
+            "f"(0.f),
+            "f"(0.f),
+            "f"(0.f));
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+          "{%0,  %1,  %2,  %3},"
+          "{%4,  %5,  %6,  %7},"
+          "{%8,  %9},"
+          "{%10, %11, %12, %13};\n"
+          : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
+          : "r"(A[0]),
+            "r"(A[1]),
+            "r"(A[2]),
+            "r"(A[3]),
+            "r"(B[2]),
+            "r"(B[3]),
+            "f"(0.f),
+            "f"(0.f),
+            "f"(0.f),
+            "f"(0.f));
+    } else {  // bf16
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+          "{%0,  %1,  %2,  %3},"
+          "{%4,  %5,  %6,  %7},"
+          "{%8,  %9},"
+          "{%10, %11, %12, %13};\n"
+          : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+          : "r"(A[0]),
+            "r"(A[1]),
+            "r"(A[2]),
+            "r"(A[3]),
+            "r"(B[0]),
+            "r"(B[1]),
+            "f"(0.f),
+            "f"(0.f),
+            "f"(0.f),
+            "f"(0.f));
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+          "{%0,  %1,  %2,  %3},"
+          "{%4,  %5,  %6,  %7},"
+          "{%8,  %9},"
+          "{%10, %11, %12, %13};\n"
+          : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
+          : "r"(A[0]),
+            "r"(A[1]),
+            "r"(A[2]),
+            "r"(A[3]),
+            "r"(B[2]),
+            "r"(B[3]),
+            "f"(0.f),
+            "f"(0.f),
+            "f"(0.f),
+            "f"(0.f));
+    }
+  } else {
+    if constexpr (std::is_same<T, half>::value) {
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+          "{%0,  %1,  %2,  %3},"
+          "{%4,  %5,  %6,  %7},"
+          "{%8,  %9},"
+          "{%10, %11, %12, %13};\n"
+          : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+          : "r"(A[0]),
+            "r"(A[1]),
+            "r"(A[2]),
+            "r"(A[3]),
+            "r"(B[0]),
+            "r"(B[1]),
+            "f"(C[0]),
+            "f"(C[1]),
+            "f"(C[2]),
+            "f"(C[3]));
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+          "{%0,  %1,  %2,  %3},"
+          "{%4,  %5,  %6,  %7},"
+          "{%8,  %9},"
+          "{%10, %11, %12, %13};\n"
+          : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
+          : "r"(A[0]),
+            "r"(A[1]),
+            "r"(A[2]),
+            "r"(A[3]),
+            "r"(B[2]),
+            "r"(B[3]),
+            "f"(C[4]),
+            "f"(C[5]),
+            "f"(C[6]),
+            "f"(C[7]));
+    } else {
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+          "{%0,  %1,  %2,  %3},"
+          "{%4,  %5,  %6,  %7},"
+          "{%8,  %9},"
+          "{%10, %11, %12, %13};\n"
+          : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+          : "r"(A[0]),
+            "r"(A[1]),
+            "r"(A[2]),
+            "r"(A[3]),
+            "r"(B[0]),
+            "r"(B[1]),
+            "f"(C[0]),
+            "f"(C[1]),
+            "f"(C[2]),
+            "f"(C[3]));
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+          "{%0,  %1,  %2,  %3},"
+          "{%4,  %5,  %6,  %7},"
+          "{%8,  %9},"
+          "{%10, %11, %12, %13};\n"
+          : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
+          : "r"(A[0]),
+            "r"(A[1]),
+            "r"(A[2]),
+            "r"(A[3]),
+            "r"(B[2]),
+            "r"(B[3]),
+            "f"(C[4]),
+            "f"(C[5]),
+            "f"(C[6]),
+            "f"(C[7]));
+    }
+  }
+}
+
+template <typename DType>
+__device__ __forceinline__ void rowsum_f16f16f32(float* d, DType* s) {
+  static_assert(sizeof(DType) == 2, "DType must be 16bit floating data type");
+  uint32_t* s_u32 = (uint32_t*)(s);
+  if constexpr (std::is_same<DType, half>::value) {
+    asm volatile(
+        "{\n"
+        ".reg .f32 ph;\n"
+        "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+        "{%0,  ph,  %1,  ph},"
+        "{%2,  %3,  %4,  %5},"
+        "{%6,  %7},"
+        "{%8,  0.,  %9,  0.};\n"
+        "}\n"
+        : "=f"(d[0]), "=f"(d[1])
+        : "r"(s_u32[0]),
+          "r"(s_u32[1]),
+          "r"(s_u32[2]),
+          "r"(s_u32[3]),
+          "r"(1006648320),
+          "r"(1006648320),
+          "f"(d[0]),
+          "f"(d[1]));
+  } else {
+    asm volatile(
+        "{\n"
+        ".reg .f32 ph;\n"
+        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+        "{%0,  ph,  %1,  ph},"
+        "{%2,  %3,  %4,  %5},"
+        "{%6,  %7},"
+        "{%8,  0.,  %9,  0.};\n"
+        "}\n"
+        : "=f"(d[0]), "=f"(d[1])
+        : "r"(s_u32[0]),
+          "r"(s_u32[1]),
+          "r"(s_u32[2]),
+          "r"(s_u32[3]),
+          "r"(1065369472),
+          "r"(1065369472),
+          "f"(d[0]),
+          "f"(d[1]));
+  }
+}

csrc/gpu/append_attn/multi_head_latent_attention_kernel.h

@@ -0,0 +1,38 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#pragma once
+#include "helper.h"
+#include "utils.cuh"
+
+template <typename T>
+void DecodeMLAAttentionKernel(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor &q, // [token_num, num_heads, head_dim]
+    const paddle::Tensor &cache_k,
+    const paddle::Tensor &cache_v,
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>& shift_bias,
+    const paddle::optional<paddle::Tensor>& smooth_weight,
+    const paddle::Tensor &seq_lens_q, // q_seq_len is 1
+    const paddle::Tensor &seq_lens_kv,
+    const paddle::Tensor &padding_offsets,
+    const paddle::Tensor &cum_offsets,
+    const paddle::Tensor &block_table,
+    int max_seq_len,
+    int max_dec_len,
+    float softmax_scale,
+    float in_scale,
+    bool causal,
+    cudaStream_t &stream,
+    paddle::Tensor *out);

csrc/gpu/append_attn/speculate_write_cache_with_rope_kernel.cu

@@ -0,0 +1,608 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "speculate_write_cache_with_rope_kernel.h"
+#include "utils.cuh"
+
+template <typename T>
+void SpeculateWriteCacheKV(const AppendAttnMetaData& meta_data,
+                         const paddle::Tensor& qkv,
+                         const paddle::Tensor& seq_lens,
+                         const paddle::Tensor& padding_offsets,
+                         const paddle::Tensor& cum_offsets,
+                         const paddle::Tensor& block_tables,
+                         const int max_seq_len,
+                         cudaStream_t& stream,
+                         paddle::Tensor* key_cache_out,
+                         paddle::Tensor* value_cache_out) {
+  auto max_blocks_per_seq = meta_data.max_blocks_per_seq;
+  auto bsz = meta_data.batch_size;
+  auto block_size = meta_data.block_size;
+  auto head_dim_qk = meta_data.head_dims;
+  auto head_dim_v = meta_data.head_dims_v;
+  auto num_heads = meta_data.q_num_heads;
+  auto kv_num_heads = meta_data.kv_num_heads;
+  auto token_num = meta_data.token_nums;
+  const uint32_t elem_nums = token_num * kv_num_heads * (head_dim_qk + head_dim_v);
+
+  constexpr int PackSize = 16 / sizeof(T);
+  const int pack_num = elem_nums / PackSize;
+  const int blocksize = 128;
+  int grid_size = 1;
+  GetNumBlocks<128>(pack_num, &grid_size);
+
+  append_speculate_cache_kernel<T, PackSize>
+      <<<grid_size, blocksize, 0, stream>>>(
+          reinterpret_cast<T*>(const_cast<T*>(qkv.data<T>())),
+          reinterpret_cast<T*>(key_cache_out->data<T>()),
+          reinterpret_cast<T*>(value_cache_out->data<T>()),
+          block_tables.data<int>(),
+          padding_offsets.data<int>(),
+          cum_offsets.data<int>(),
+          seq_lens.data<int>(),
+          max_seq_len,
+          max_blocks_per_seq,
+          num_heads,
+          head_dim_qk,
+          head_dim_v,
+          block_size,
+          elem_nums,
+          kv_num_heads);
+}
+
+// rope + write
+template <typename T, typename QKV_TYPE>
+void append_speculate_cache_rope(const QKV_TYPE* qkv,
+                                 T* key_cache,
+                                 T* value_cache,
+                                 T* qkv_out,
+                                 const int* block_tables,
+                                 const int* padding_offsets,
+                                 const int* cum_offsets,
+                                 const int* seq_lens,
+                                 const int* seq_lens_encoder,
+                                 const float* cos_emb,
+                                 const float* sin_emb,
+                                 const float* qkv_out_scales,
+                                 const T* qkv_biases,
+                                 const int max_seq_len,
+                                 const int max_blocks_per_seq,
+                                 const int num_heads,
+                                 const int kv_num_heads,
+                                 const int dim_head,
+                                 const int block_size,
+                                 const int bsz,
+                                 const int token_num,
+                                 const cudaStream_t& stream,
+                                 const bool use_neox_style) {
+  int output_inner_dim = num_heads + 2 * kv_num_heads;
+
+  const uint32_t elem_nums =
+      use_neox_style ? token_num * (num_heads + 2 * kv_num_heads) * dim_head / 2
+                     : token_num * (num_heads + 2 * kv_num_heads) * dim_head;
+  constexpr int PackSize = 16 / sizeof(T);
+  const int pack_num = elem_nums / PackSize;
+
+  const int threads_per_block = 128;
+  int grid_size = 1;
+  GetNumBlocks(pack_num, &grid_size);
+  if (use_neox_style) {
+    append_speculate_cache_neox_rope_kernel<T, PackSize>
+        <<<grid_size, threads_per_block, 0, stream>>>(
+            qkv,  // [token_num, num_heads + 2 * gqa_group_size, head_size]
+            key_cache,
+            value_cache,
+            qkv_out,
+            block_tables,
+            padding_offsets,
+            cum_offsets,
+            seq_lens,
+            cos_emb,
+            sin_emb,
+            qkv_out_scales,
+            qkv_biases,  // [num_head + 2 * gqa_group_size, dim_head]
+            max_seq_len,
+            max_blocks_per_seq,
+            num_heads,
+            output_inner_dim,
+            dim_head,
+            block_size,
+            elem_nums,
+            kv_num_heads);
+  } else {
+    append_speculate_cache_rope_kernel<T, PackSize>
+        <<<grid_size, threads_per_block, 0, stream>>>(
+            qkv,  // [token_num, num_heads + 2 * gqa_group_size, head_size]
+            key_cache,
+            value_cache,
+            qkv_out,
+            block_tables,
+            padding_offsets,
+            cum_offsets,
+            seq_lens,
+            cos_emb,
+            sin_emb,
+            qkv_out_scales,
+            qkv_biases,  // [num_head + 2 * gqa_group_size, dim_head]
+            max_seq_len,
+            max_blocks_per_seq,
+            num_heads,
+            output_inner_dim,
+            dim_head,
+            block_size,
+            elem_nums,
+            kv_num_heads);
+  }
+}
+
+template <typename T, typename QKV_TYPE>
+void append_speculate_cache_int8_rope(const QKV_TYPE* qkv,
+                                      uint8_t* key_cache,
+                                      uint8_t* value_cache,
+                                      T* qkv_out,
+                                      const int* block_tables,
+                                      const int* padding_offsets,
+                                      const int* cum_offsets,
+                                      const int* seq_lens,
+                                      const int* seq_lens_encoder,
+                                      const float* cos_emb,
+                                      const float* sin_emb,
+                                      const float* qkv_out_scales,
+                                      const T* qkv_biases,
+                                      const T* cache_k_scale,
+                                      const T* cache_v_scale,
+                                      const int max_seq_len,
+                                      const int max_blocks_per_seq,
+                                      const int num_heads,
+                                      const int kv_num_heads,
+                                      const int dim_head,
+                                      const int block_size,
+                                      const int bsz,
+                                      const int token_num,
+                                      const cudaStream_t& stream,
+                                      const bool use_neox_style) {
+  constexpr int num_warps = 4;
+  const int all_warps =
+      ((num_heads + 2 * kv_num_heads) + num_warps - 1) / num_warps * num_warps;
+  dim3 grids(token_num, all_warps / num_warps);
+
+  append_clear_cache_int8_block<4>
+      <<<grids, num_warps * 32, 0, stream>>>(key_cache,
+                                             value_cache,
+                                             seq_lens,
+                                             block_tables,
+                                             padding_offsets,
+                                             cum_offsets,
+                                             seq_lens_encoder,
+                                             max_seq_len,
+                                             max_blocks_per_seq,
+                                             num_heads,
+                                             block_size,
+                                             kv_num_heads);
+  if (use_neox_style) {
+    append_speculate_cache_int8_neox_rope_kernel<T, 4>
+        <<<grids, num_warps * 32, 0, stream>>>(qkv,
+                                               key_cache,
+                                               value_cache,
+                                               qkv_out,
+                                               block_tables,
+                                               padding_offsets,
+                                               cum_offsets,
+                                               seq_lens,
+                                               seq_lens_encoder,
+                                               cos_emb,
+                                               sin_emb,
+                                               qkv_out_scales,
+                                               qkv_biases,
+                                               cache_k_scale,
+                                               cache_v_scale,
+                                               max_seq_len,
+                                               max_blocks_per_seq,
+                                               num_heads,
+                                               block_size,
+                                               127.0f,
+                                               -127.0f,
+                                               kv_num_heads);
+  } else {
+    append_speculate_cache_int8_rope_kernel<T, 4>
+        <<<grids, num_warps * 32, 0, stream>>>(qkv,
+                                               key_cache,
+                                               value_cache,
+                                               qkv_out,
+                                               block_tables,
+                                               padding_offsets,
+                                               cum_offsets,
+                                               seq_lens,
+                                               seq_lens_encoder,
+                                               cos_emb,
+                                               sin_emb,
+                                               qkv_out_scales,
+                                               qkv_biases,
+                                               cache_k_scale,
+                                               cache_v_scale,
+                                               max_seq_len,
+                                               max_blocks_per_seq,
+                                               num_heads,
+                                               block_size,
+                                               127.0f,
+                                               -127.0f,
+                                               kv_num_heads);
+  }
+}
+
+template <typename T, typename QKV_TYPE>
+void append_speculate_cache_int4_rope(const QKV_TYPE* qkv,
+                                      uint8_t* key_cache,
+                                      uint8_t* value_cache,
+                                      T* qkv_out,
+                                      const int* block_tables,
+                                      const int* padding_offsets,
+                                      const int* cum_offsets,
+                                      const int* seq_lens,
+                                      const int* seq_lens_encoder,
+                                      const float* cos_emb,
+                                      const float* sin_emb,
+                                      const float* qkv_out_scales,
+                                      const T* qkv_biases,
+                                      const T* cache_k_scale,
+                                      const T* cache_v_scale,
+                                      const T* cache_k_zp,
+                                      const T* cache_v_zp,
+                                      const int max_seq_len,
+                                      const int max_blocks_per_seq,
+                                      const int num_heads,
+                                      const int kv_num_heads,
+                                      const int dim_head,
+                                      const int block_size,
+                                      const int bsz,
+                                      const int token_num,
+                                      const cudaStream_t& stream,
+                                      const bool use_neox_style) {
+  constexpr int num_warps = 4;
+  const int all_warps =
+      ((num_heads + 2 * kv_num_heads) + num_warps - 1) / num_warps * num_warps;
+  dim3 grids(token_num, all_warps / num_warps);
+
+  append_clear_cache_int4_block<4>
+      <<<grids, num_warps * 32, 0, stream>>>(key_cache,
+                                             value_cache,
+                                             seq_lens,
+                                             block_tables,
+                                             padding_offsets,
+                                             cum_offsets,
+                                             seq_lens_encoder,
+                                             max_seq_len,
+                                             max_blocks_per_seq,
+                                             num_heads,
+                                             block_size,
+                                             kv_num_heads);
+  if (use_neox_style) {
+    append_speculate_cache_int4_neox_rope_kernel<T, 4>
+        <<<grids, num_warps * 32, 0, stream>>>(qkv,
+                                               key_cache,
+                                               value_cache,
+                                               qkv_out,
+                                               block_tables,
+                                               padding_offsets,
+                                               cum_offsets,
+                                               seq_lens,
+                                               seq_lens_encoder,
+                                               cos_emb,
+                                               sin_emb,
+                                               qkv_out_scales,
+                                               qkv_biases,
+                                               cache_k_scale,
+                                               cache_v_scale,
+                                               cache_k_zp,
+                                               cache_v_zp,
+                                               max_seq_len,
+                                               max_blocks_per_seq,
+                                               num_heads,
+                                               block_size,
+                                               7.0f,
+                                               -8.0f,
+                                               kv_num_heads);
+  } else {
+    append_speculate_cache_int4_rope_kernel<T, 4>
+        <<<grids, num_warps * 32, 0, stream>>>(qkv,
+                                               key_cache,
+                                               value_cache,
+                                               qkv_out,
+                                               block_tables,
+                                               padding_offsets,
+                                               cum_offsets,
+                                               seq_lens,
+                                               seq_lens_encoder,
+                                               cos_emb,
+                                               sin_emb,
+                                               qkv_out_scales,
+                                               qkv_biases,
+                                               cache_k_scale,
+                                               cache_v_scale,
+                                               cache_k_zp,
+                                               cache_v_zp,
+                                               max_seq_len,
+                                               max_blocks_per_seq,
+                                               num_heads,
+                                               block_size,
+                                               7.0f,
+                                               -8.0f,
+                                               kv_num_heads);
+  }
+}
+template <typename T, typename QKV_TYPE>
+void SpeculateWriteCacheWithRoPEKernel(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& qkv_biases,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const std::string& cache_quant_type_str,
+    const bool use_neox_rotary_style,
+    const int max_seq_len,
+    cudaStream_t& stream,
+    paddle::Tensor* qkv_out,
+    paddle::Tensor* key_cache_out,
+    paddle::Tensor* value_cache_out) {
+  typedef cascade_attn_type_traits<T> traits_;
+  typedef cascade_attn_type_traits<QKV_TYPE> qkt_nv_type_;
+  typedef typename traits_::type DataType_;
+  typedef typename qkt_nv_type_::type QKV_Data_TYPE;
+  const QKV_TYPE* qkv_ptr = qkv.data<QKV_TYPE>();
+
+  auto max_blocks_per_seq = meta_data.max_blocks_per_seq;
+  auto bsz = meta_data.batch_size;
+  auto token_nums = meta_data.token_nums;
+  auto block_size = meta_data.block_size;
+  auto dim_head = meta_data.head_dims;
+  auto num_heads = meta_data.q_num_heads;
+  auto kv_num_heads = meta_data.kv_num_heads;
+
+  if (rotary_embs) {
+    const float* cos_emb =
+        rotary_embs ? rotary_embs.get().data<float>() : nullptr;
+    const float* sin_emb =
+            use_neox_rotary_style
+                ? rotary_embs.get().data<float>() + max_seq_len * dim_head
+                : rotary_embs.get().data<float>() + max_seq_len * dim_head / 2;
+    
+    if (cache_quant_type_str == "none") {
+        append_speculate_cache_rope(
+            reinterpret_cast<const QKV_TYPE*>(qkv_ptr),
+            reinterpret_cast<DataType_*>(key_cache_out->data<T>()),
+            reinterpret_cast<DataType_*>(value_cache_out->data<T>()),
+            reinterpret_cast<DataType_*>(qkv_out->data<T>()),
+            block_tables.data<int>(),
+            padding_offsets.data<int>(),
+            cum_offsets.data<int>(),
+            seq_lens.data<int>(),
+            seq_lens_encoder.data<int>(),
+            cos_emb,
+            sin_emb,
+            qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
+            qkv_biases ? reinterpret_cast<DataType_*>(
+                            const_cast<T*>(qkv_biases.get().data<T>()))
+                    : nullptr,
+            max_seq_len,
+            max_blocks_per_seq,
+            num_heads,
+            kv_num_heads,
+            dim_head,
+            block_size,
+            bsz,
+            token_nums,
+            stream,
+            use_neox_rotary_style);
+    } else if (cache_quant_type_str == "cache_int8") {
+        append_speculate_cache_int8_rope(
+            reinterpret_cast<const QKV_TYPE*>(qkv_ptr),
+            key_cache_out->data<uint8_t>(),
+            value_cache_out->data<uint8_t>(),
+            reinterpret_cast<DataType_*>(qkv_out->data<T>()),
+            block_tables.data<int>(),
+            padding_offsets.data<int>(),
+            cum_offsets.data<int>(),
+            seq_lens.data<int>(),
+            seq_lens_encoder.data<int>(),
+            cos_emb,
+            sin_emb,
+            qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
+            qkv_biases ? reinterpret_cast<DataType_*>(
+                            const_cast<T*>(qkv_biases.get().data<T>()))
+                    : nullptr,
+            cache_k_scale ? reinterpret_cast<DataType_*>(
+                                const_cast<T*>(cache_k_scale.get().data<T>()))
+                        : nullptr,
+            cache_v_scale ? reinterpret_cast<DataType_*>(
+                                const_cast<T*>(cache_v_scale.get().data<T>()))
+                        : nullptr,
+            max_seq_len,
+            max_blocks_per_seq,
+            num_heads,
+            kv_num_heads,
+            dim_head,
+            block_size,
+            bsz,
+            token_nums,
+            stream,
+            use_neox_rotary_style);
+    } else if (cache_quant_type_str == "cache_int4_zp") {
+        append_speculate_cache_int4_rope(
+            reinterpret_cast<const QKV_TYPE*>(qkv_ptr),
+            key_cache_out->data<uint8_t>(),
+            value_cache_out->data<uint8_t>(),
+            reinterpret_cast<DataType_*>(const_cast<T*>(qkv_out->data<T>())),
+            block_tables.data<int>(),
+            padding_offsets.data<int>(),
+            cum_offsets.data<int>(),
+            seq_lens.data<int>(),
+            seq_lens_encoder.data<int>(),
+            cos_emb,
+            sin_emb,
+            qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
+            qkv_biases ? reinterpret_cast<DataType_*>(
+                            const_cast<T*>(qkv_biases.get().data<T>()))
+                    : nullptr,
+            cache_k_scale ? reinterpret_cast<DataType_*>(
+                                const_cast<T*>(cache_k_scale.get().data<T>()))
+                        : nullptr,
+            cache_v_scale ? reinterpret_cast<DataType_*>(
+                                const_cast<T*>(cache_v_scale.get().data<T>()))
+                        : nullptr,
+            cache_k_zp ? reinterpret_cast<DataType_*>(
+                            const_cast<T*>(cache_k_zp.get().data<T>()))
+                    : nullptr,
+            cache_v_zp ? reinterpret_cast<DataType_*>(
+                            const_cast<T*>(cache_v_zp.get().data<T>()))
+                    : nullptr,
+            max_seq_len,
+            max_blocks_per_seq,
+            num_heads,
+            kv_num_heads,
+            dim_head,
+            block_size,
+            bsz,
+            token_nums,
+            stream,
+            use_neox_rotary_style);
+    } else {
+        PD_THROW(
+            "cache_quant_type_str should be one of [none, cache_int8, "
+            "cache_int4_zp]");
+    }
+  } else {
+    SpeculateWriteCacheKV<QKV_TYPE>(meta_data,
+                                    qkv,
+                                    seq_lens,
+                                    padding_offsets,
+                                    cum_offsets,
+                                    block_tables,
+                                    max_seq_len,
+                                    stream,
+                                    key_cache_out,
+                                    value_cache_out);
+  }
+}
+
+template void SpeculateWriteCacheWithRoPEKernel<paddle::bfloat16, int>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor&
+        qkv,  // [token_num, 3, num_head, head_dim] ([token_num, num_head + 2 *
+              // gqa_group_size, head_dim] if GQA)
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& qkv_biases,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const std::string& cache_quant_type_str,
+    const bool use_neox_rotary_style,
+    const int max_seq_len,
+    cudaStream_t& stream,
+    paddle::Tensor* qkv_out,
+    paddle::Tensor* key_cache_out,
+    paddle::Tensor* value_cache_out);
+
+template void
+SpeculateWriteCacheWithRoPEKernel<paddle::bfloat16, paddle::bfloat16>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor&
+        qkv,  // [token_num, 3, num_head, head_dim] ([token_num, num_head + 2 *
+              // gqa_group_size, head_dim] if GQA)
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& qkv_biases,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const std::string& cache_quant_type_str,
+    const bool use_neox_rotary_style,
+    const int max_seq_len,
+    cudaStream_t& stream,
+    paddle::Tensor* qkv_out,
+    paddle::Tensor* key_cache_out,
+    paddle::Tensor* value_cache_out);
+
+template void SpeculateWriteCacheWithRoPEKernel<paddle::float16, int>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor&
+        qkv,  // [token_num, 3, num_head, head_dim] ([token_num, num_head + 2 *
+              // gqa_group_size, head_dim] if GQA)
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& qkv_biases,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const std::string& cache_quant_type_str,
+    const bool use_neox_rotary_style,
+    const int max_seq_len,
+    cudaStream_t& stream,
+    paddle::Tensor* qkv_out,
+    paddle::Tensor* key_cache_out,
+    paddle::Tensor* value_cache_out);
+
+
+template void
+SpeculateWriteCacheWithRoPEKernel<paddle::float16, paddle::float16>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor&
+        qkv,  // [token_num, 3, num_head, head_dim] ([token_num, num_head + 2 *
+              // gqa_group_size, head_dim] if GQA)
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& qkv_biases,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const std::string& cache_quant_type_str,
+    const bool use_neox_rotary_style,
+    const int max_seq_len,
+    cudaStream_t& stream,
+    paddle::Tensor* qkv_out,
+    paddle::Tensor* key_cache_out,
+    paddle::Tensor* value_cache_out);

csrc/gpu/append_attn/speculate_write_cache_with_rope_kernel.h

@@ -0,0 +1,42 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#pragma once
+
+#include "speculate_write_cache_with_rope_impl.cuh"
+
+template <typename T, typename QKV_TYPE = int>
+void SpeculateWriteCacheWithRoPEKernel(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor&
+        qkv,  // [token_num, 3, num_head, head_dim] ([token_num, num_head + 2 *
+              // gqa_group_size, head_dim] if GQA)
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& qkv_biases,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const std::string& cache_quant_type_str,
+    const bool use_neox_rotary_style,
+    const int max_seq_len,
+    cudaStream_t& stream,
+    paddle::Tensor* qkv_out,
+    paddle::Tensor* key_cache_out,
+    paddle::Tensor* value_cache_out);

csrc/gpu/append_attn/template_instantiation/append_attention_c16_bfloat16_bfloat16_kernel.cu

@@ -0,0 +1,58 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c16_impl.cuh"
+
+
+template void CascadeAppendAttentionC16Kernel<paddle::bfloat16, paddle::bfloat16>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/append_attention_c16_bfloat16_fp8_kernel.cu

@@ -0,0 +1,57 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c16_impl.cuh"
+
+template void CascadeAppendAttentionC16Kernel<paddle::bfloat16, paddle::float8_e4m3fn>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/append_attention_c16_bfloat16_int8_kernel.cu

@@ -0,0 +1,57 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c16_impl.cuh"
+
+template void CascadeAppendAttentionC16Kernel<paddle::bfloat16, int8_t>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/append_attention_c16_float16_float16_kernel.cu

@@ -0,0 +1,57 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c16_impl.cuh"
+
+template void CascadeAppendAttentionC16Kernel<paddle::float16, paddle::float16>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/append_attention_c16_float16_fp8_kernel.cu

@@ -0,0 +1,57 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c16_impl.cuh"
+
+template void CascadeAppendAttentionC16Kernel<paddle::float16, paddle::float8_e4m3fn>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/append_attention_c16_float16_int8_kernel.cu

@@ -0,0 +1,57 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c16_impl.cuh"
+
+template void CascadeAppendAttentionC16Kernel<paddle::float16, int8_t>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/append_attention_c4_bfloat16_bfloat16_kernel.cu

@@ -0,0 +1,57 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c4_impl.cuh"
+
+template void CascadeAppendAttentionC4Kernel<paddle::bfloat16, paddle::bfloat16>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/append_attention_c4_bfloat16_fp8_kernel.cu

@@ -0,0 +1,57 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c4_impl.cuh"
+
+template void CascadeAppendAttentionC4Kernel<paddle::bfloat16, paddle::float8_e4m3fn>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/append_attention_c4_bfloat16_int8_kernel.cu

@@ -0,0 +1,57 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c4_impl.cuh"
+
+template void CascadeAppendAttentionC4Kernel<paddle::bfloat16, int8_t>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/append_attention_c4_float16_float16_kernel.cu

@@ -0,0 +1,58 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c4_impl.cuh"
+
+
+template void CascadeAppendAttentionC4Kernel<paddle::float16, paddle::float16>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/append_attention_c4_float16_fp8_kernel.cu

@@ -0,0 +1,57 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c4_impl.cuh"
+
+template void CascadeAppendAttentionC4Kernel<paddle::float16, paddle::float8_e4m3fn>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/append_attention_c4_float16_int8_kernel.cu

@@ -0,0 +1,57 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c4_impl.cuh"
+
+template void CascadeAppendAttentionC4Kernel<paddle::float16, int8_t>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/append_attention_c8_bfloat16_bfloat16_kernel.cu

@@ -0,0 +1,59 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c8_impl.cuh"
+
+
+template void
+CascadeAppendAttentionC8Kernel<paddle::bfloat16, paddle::bfloat16>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/append_attention_c8_bfloat16_fp8_kernel.cu

@@ -0,0 +1,57 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c8_impl.cuh"
+
+template void CascadeAppendAttentionC8Kernel<paddle::bfloat16, paddle::float8_e4m3fn>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/append_attention_c8_bfloat16_int8_kernel.cu

@@ -0,0 +1,57 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c8_impl.cuh"
+
+template void CascadeAppendAttentionC8Kernel<paddle::bfloat16, int8_t>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/append_attention_c8_float16_float16_kernel.cu

@@ -0,0 +1,57 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c8_impl.cuh"
+
+template void CascadeAppendAttentionC8Kernel<paddle::float16, paddle::float16>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/append_attention_c8_float16_fp8_kerne.cu

@@ -0,0 +1,57 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c8_impl.cuh"
+
+template void CascadeAppendAttentionC8Kernel<paddle::float16, paddle::float8_e4m3fn>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/append_attention_c8_float16_int8_kerne.cu

@@ -0,0 +1,57 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../append_attention_c8_impl.cuh"
+
+template void CascadeAppendAttentionC8Kernel<paddle::float16, int8_t>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& qkv,  // [token_num, num_heads, head_dim]
+    const paddle::Tensor&
+        cache_k,  // [max_block_num, num_heads, block_size, head_dim]
+    const paddle::Tensor&
+        cache_v,  // [max_block_num, num_heads, head_dim, block_size]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_q,
+    const paddle::Tensor& seq_lens_kv,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_table,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const int num_blocks,
+    const int block_shape_q,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool is_decoder,
+    const bool enable_prefill,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/append_attn/template_instantiation/encoder_write_cache_with_rope_bfloat16_bfloat16_kernel.cu

@@ -0,0 +1,45 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../encoder_write_cache_with_rope_kernel.h"
+
+template void
+EncoderWriteCacheWithRopeKernel<paddle::bfloat16, paddle::bfloat16>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor&
+        qkv,  // [token_num, 3, num_head, head_dim] ([token_num, num_head + 2 *
+              // kv_num_heads, head_dim] if GQA)
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& qkv_biases,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const paddle::optional<paddle::Tensor>& excess_blocks,
+    const std::string& cache_quant_type_str,
+    const int num_blocks,
+    const int max_seq_len,
+    const bool use_neox_style,
+    cudaStream_t& stream,
+    paddle::Tensor* qkv_out,
+    paddle::Tensor* key_cache_out,
+    paddle::Tensor* value_cache_out);

csrc/gpu/append_attn/template_instantiation/encoder_write_cache_with_rope_bfloat16_int_kernel.cu

@@ -0,0 +1,44 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../encoder_write_cache_with_rope_kernel.h"
+
+template void EncoderWriteCacheWithRopeKernel<paddle::bfloat16, int>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor&
+        qkv,  // [token_num, 3, num_head, head_dim] ([token_num, num_head + 2 *
+              // kv_num_heads, head_dim] if GQA)
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& qkv_biases,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const paddle::optional<paddle::Tensor>& excess_blocks,
+    const std::string& cache_quant_type_str,
+    const int num_blocks,
+    const int max_seq_len,
+    const bool use_neox_style,
+    cudaStream_t& stream,
+    paddle::Tensor* qkv_out,
+    paddle::Tensor* key_cache_out,
+    paddle::Tensor* value_cache_out);

csrc/gpu/append_attn/template_instantiation/encoder_write_cache_with_rope_float16_float16_kernel.cu

@@ -0,0 +1,44 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../encoder_write_cache_with_rope_kernel.h"
+
+template void EncoderWriteCacheWithRopeKernel<paddle::float16, paddle::float16>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor&
+        qkv,  // [token_num, 3, num_head, head_dim] ([token_num, num_head + 2 *
+              // kv_num_heads, head_dim] if GQA)
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& qkv_biases,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const paddle::optional<paddle::Tensor>& excess_blocks,
+    const std::string& cache_quant_type_str,
+    const int num_blocks,
+    const int max_seq_len,
+    const bool use_neox_style,
+    cudaStream_t& stream,
+    paddle::Tensor* qkv_out,
+    paddle::Tensor* key_cache_out,
+    paddle::Tensor* value_cache_out);

csrc/gpu/append_attn/template_instantiation/encoder_write_cache_with_rope_float16_int_kernel.cu

@@ -0,0 +1,44 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "../encoder_write_cache_with_rope_kernel.h"
+
+template void EncoderWriteCacheWithRopeKernel<paddle::float16, int>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor&
+        qkv,  // [token_num, 3, num_head, head_dim] ([token_num, num_head + 2 *
+              // kv_num_heads, head_dim] if GQA)
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& qkv_biases,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const paddle::optional<paddle::Tensor>& excess_blocks,
+    const std::string& cache_quant_type_str,
+    const int num_blocks,
+    const int max_seq_len,
+    const bool use_neox_style,
+    cudaStream_t& stream,
+    paddle::Tensor* qkv_out,
+    paddle::Tensor* key_cache_out,
+    paddle::Tensor* value_cache_out);

csrc/gpu/append_attn/utils.cuh

@@ -0,0 +1,474 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#pragma once
+
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+
+#include "mem_util.cuh"
+
+struct AppendAttnMetaData {
+  int batch_size;
+  int block_size;
+  int q_num_heads;
+  int kv_num_heads;
+  int token_nums;
+  int head_dims;
+  int head_dims_v;
+  int max_blocks_per_seq;
+};
+
+__forceinline__ __host__ __device__ int div_up(int a, int b) {
+  return (a + b - 1) / b;
+}
+
+enum PosEncMode { kNonePos, kRoPE, kAliBi };
+
+enum CacheType { CacheT, CacheInt8Hw, CacheInt4CwZp };
+
+template <typename T>
+struct cascade_attn_type_traits {
+  using type = T;
+};
+
+template <>
+struct cascade_attn_type_traits<phi::dtype::bfloat16> {
+  using type = __nv_bfloat16;
+};
+
+template <>
+struct cascade_attn_type_traits<phi::dtype::float16> {
+  using type = half;
+};
+
+template <>
+struct cascade_attn_type_traits<phi::dtype::float8_e4m3fn> {
+  using type = __nv_fp8_e4m3;
+};
+
+template <typename T>
+struct cascade_attn_nv_type2_traits {
+  using type = T;
+};
+
+template <>
+struct cascade_attn_nv_type2_traits<__nv_bfloat16> {
+  using type = __nv_bfloat162;
+};
+
+template <>
+struct cascade_attn_nv_type2_traits<half> {
+  using type = half2;
+};
+
+template <CacheType cache_type>
+struct vec_traits {
+  using type = b128_t;
+};
+
+template <>
+struct vec_traits<CacheType::CacheInt8Hw> {
+  using type = b64_t;
+};
+
+template <>
+struct vec_traits<CacheType::CacheInt4CwZp> {
+  using type = b32_t;
+};
+
+template <typename T, CacheType cache_type>
+struct cache_type_traits {
+  using type = T;
+};
+
+template <typename T>
+struct cache_type_traits<T, CacheType::CacheInt8Hw> {
+  using type = uint8_t;
+};
+
+template <typename T>
+struct cache_type_traits<T, CacheType::CacheInt4CwZp> {
+  using type = uint8_t;
+};
+
+__device__ __forceinline__ uint32_t sub_if_greater_or_zero(uint32_t x,
+                                                           uint32_t y) {
+  return (x > y) ? x - y : 0U;
+}
+
+/******************************FASTER CAST*********************************/
+inline __device__ static void convert_int8(
+    __nv_bfloat16* result, const uint32_t& source) {  // 4 int8 each time
+  uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(result);
+  uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
+
+  static constexpr uint32_t fp32_base = 0x4B000000;
+  float fp32_intermediates[4];
+
+  uint32_t* fp32_intermediates_casted =
+      reinterpret_cast<uint32_t*>(fp32_intermediates);
+  fp32_intermediates_casted[0] = __byte_perm(i8s, fp32_base, 0x7650);
+  fp32_intermediates_casted[1] = __byte_perm(i8s, fp32_base, 0x7651);
+  fp32_intermediates_casted[2] = __byte_perm(i8s, fp32_base, 0x7652);
+  fp32_intermediates_casted[3] = __byte_perm(i8s, fp32_base, 0x7653);
+
+#pragma unroll
+  for (int ii = 0; ii < 4; ++ii) {
+    fp32_intermediates[ii] -= 8388736.f;  // (8388608.f + 128.f);
+  }
+
+#pragma unroll
+  for (int ii = 0; ii < 2; ++ii) {
+    bf16_result_ptr[ii] = __byte_perm(fp32_intermediates_casted[2 * ii + 0],
+                                      fp32_intermediates_casted[2 * ii + 1],
+                                      0x7632);
+  }
+}
+
+inline __device__ static void convert_int8(
+    half* result, const uint32_t& source) {  // 4 int8 each time
+  uint32_t* fp16_result_ptr = reinterpret_cast<uint32_t*>(result);
+  uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
+  static constexpr uint32_t mask_for_elt_01 = 0x5150;
+  static constexpr uint32_t mask_for_elt_23 = 0x5352;
+  static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
+
+  asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+               : "=r"(fp16_result_ptr[0])
+               : "r"(i8s), "n"(start_byte_for_fp16), "n"(mask_for_elt_01));
+  asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+               : "=r"(fp16_result_ptr[1])
+               : "r"(i8s), "n"(start_byte_for_fp16), "n"(mask_for_elt_23));
+
+  static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480;
+  asm volatile("sub.f16x2 %0, %1, %2;\n"
+               : "=r"(fp16_result_ptr[0])
+               : "r"(fp16_result_ptr[0]), "r"(I8s_TO_F16s_MAGIC_NUM));
+  asm volatile("sub.f16x2 %0, %1, %2;\n"
+               : "=r"(fp16_result_ptr[1])
+               : "r"(fp16_result_ptr[1]), "r"(I8s_TO_F16s_MAGIC_NUM));
+}
+
+inline __device__ static void convert_int4(
+    __nv_bfloat16* result, const uint32_t& source) {  // 8 int4 each time
+  uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(result);
+
+  static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;
+  static constexpr uint32_t MASK = 0x0f0f0f0f;  // 0xf -> 0b1111 select 0,4
+  static constexpr uint32_t I4s_TO_FP32s_MAGIC_NUM = 0x43434343;
+  static constexpr uint32_t mask_for_elt_01 = 0x5150;
+  static constexpr uint32_t mask_for_elt_23 = 0x5352;
+
+  uint32_t tmp1 = source & MASK;       // 0 1 2 3
+  uint32_t tmp2 = source >> 4 & MASK;  // 4 5 6 7
+
+  bf16_result_ptr[0] = __byte_perm(tmp1,
+                                   I4s_TO_FP32s_MAGIC_NUM,
+                                   mask_for_elt_01);  // 0 1
+  bf16_result_ptr[1] = __byte_perm(tmp1,
+                                   I4s_TO_FP32s_MAGIC_NUM,
+                                   mask_for_elt_23);  // 2 3
+  bf16_result_ptr[2] = __byte_perm(tmp2,
+                                   I4s_TO_FP32s_MAGIC_NUM,
+                                   mask_for_elt_01);  // 4 5
+  bf16_result_ptr[3] = __byte_perm(tmp2,
+                                   I4s_TO_FP32s_MAGIC_NUM,
+                                   mask_for_elt_23);  // 6 7
+}
+
+inline __device__ static void convert_int4(
+    half* result, const uint32_t& source) {  // 7 5 3 1 6 4 2 0
+  uint32_t* fp16_result_ptr = reinterpret_cast<uint32_t*>(result);
+
+  static constexpr uint32_t MASK =
+      0x0f0f0f0f;  // 0xf -> 0b1111 select 0,1;   7 5 3 1 6 4 2 0
+  static constexpr uint32_t I4s_TO_FP32s_MAGIC_NUM = 0x64646464;
+  static constexpr uint32_t mask_for_elt_01 = 0x5150;
+  static constexpr uint32_t mask_for_elt_23 = 0x5352;
+
+  uint32_t tmp1 = source & MASK;       // 0 1 2 3
+  uint32_t tmp2 = source >> 4 & MASK;  // 4 5 6 7
+  fp16_result_ptr[0] = __byte_perm(tmp1,
+                                   I4s_TO_FP32s_MAGIC_NUM,
+                                   mask_for_elt_01);  // 0 1
+  fp16_result_ptr[1] = __byte_perm(tmp1,
+                                   I4s_TO_FP32s_MAGIC_NUM,
+                                   mask_for_elt_23);  // 2 3
+  fp16_result_ptr[2] = __byte_perm(tmp2,
+                                   I4s_TO_FP32s_MAGIC_NUM,
+                                   mask_for_elt_01);  // 4 5
+  fp16_result_ptr[3] = __byte_perm(tmp2,
+                                   I4s_TO_FP32s_MAGIC_NUM,
+                                   mask_for_elt_23);  // 6 7
+}
+
+/******************* vec_t type cast *******************/
+
+template <typename dst_t, typename src_t, size_t vec_size>
+__forceinline__ __host__ __device__ void vec_cast(dst_t* dst,
+                                                  const src_t* src) {
+#pragma unroll
+  for (size_t i = 0; i < vec_size; ++i) {
+    dst[i] = src[i];
+  }
+}
+
+template <size_t vec_size>
+__forceinline__ __host__ __device__ void vec_cast<float, half>(
+    float* dst, const half* src) {
+#pragma unroll
+  for (size_t i = 0; i < vec_size / 2; ++i) {
+    ((float2*)dst)[i] = __half22float2(((half2*)src)[i]);
+  }
+}
+
+template <size_t vec_size>
+__forceinline__ __host__ __device__ void vec_cast<half, float>(
+    half* dst, const float* src) {
+#pragma unroll
+  for (size_t i = 0; i < vec_size / 2; ++i) {
+    ((half2*)dst)[i] = __float22half2_rn(((float2*)src)[i]);
+  }
+}
+
+template <size_t vec_size>
+__forceinline__ __host__ __device__ void vec_cast<float, nv_bfloat16>(
+    float* dst, const nv_bfloat16* src) {
+#pragma unroll
+  for (size_t i = 0; i < vec_size / 2; ++i) {
+    ((float2*)dst)[i] = __bfloat1622float2(((nv_bfloat162*)src)[i]);
+  }
+}
+
+template <size_t vec_size>
+__forceinline__ __host__ __device__ void vec_cast<nv_bfloat16, float>(
+    nv_bfloat16* dst, const float* src) {
+#pragma unroll
+  for (size_t i = 0; i < vec_size / 2; ++i) {
+    ((nv_bfloat162*)dst)[i] = __float22bfloat162_rn(((float2*)src)[i]);
+  }
+}
+
+#define CHECK_CUDA_CALL(func, ...)                                      \
+  {                                                                     \
+    cudaError_t e = (func);                                             \
+    if (e != cudaSuccess) {                                             \
+      std::cerr << "CUDA Error: " << cudaGetErrorString(e) << " (" << e \
+                << ") " << __FILE__ << ": line " << __LINE__            \
+                << " at function " << STR(func) << std::endl;           \
+      return e;                                                         \
+    }                                                                   \
+  }
+
+#define DISPATCH_GQA_HEAD_DIM(head_dim, HEAD_DIM, ...)  \
+  switch (head_dim) {                                   \
+    case 128: {                                         \
+      constexpr size_t HEAD_DIM = 128;                  \
+      __VA_ARGS__                                       \
+      break;                                            \
+    }                                                   \
+    case 192: {                                         \
+      constexpr size_t HEAD_DIM = 192;                  \
+      __VA_ARGS__                                       \
+      break;                                            \
+    }                                                   \
+    default: {                                          \
+      PD_THROW("not support the head_dim: ", head_dim); \
+    }                                                   \
+  }
+
+#define DISPATCH_MLA_HEAD_DIM(head_dim, HEAD_DIM, ...)  \
+  switch (head_dim) {                                   \
+    case 128: {                                         \
+      constexpr size_t HEAD_DIM = 128;                  \
+      __VA_ARGS__                                       \
+      break;                                            \
+    }                                                   \
+    case 192: {                                         \
+      constexpr size_t HEAD_DIM = 192;                  \
+      __VA_ARGS__                                       \
+      break;                                            \
+    }                                                   \
+    case 512: {                                         \
+      constexpr size_t HEAD_DIM = 512;                  \
+      __VA_ARGS__                                       \
+      break;                                            \
+    }                                                   \
+    case 576: {                                         \
+      constexpr size_t HEAD_DIM = 576;                  \
+      __VA_ARGS__                                       \
+      break;                                            \
+    }                                                   \
+    default: {                                          \
+      PD_THROW("not support the head_dim: ", head_dim); \
+    }                                                   \
+  }
+
+#define DISPATCH_NUM_STAGE(num_stage, NUM_STAGE, ...)   \
+  if (num_stage == 2) {                                 \
+    constexpr size_t NUM_STAGE = 2;                     \
+    __VA_ARGS__                                         \
+  } else {                                              \
+    PD_THROW("not support the num_stage: ", num_stage); \
+  }
+
+#define DISPATCH_CACHE_TYPE(cache_type, cache_type_now, cache_bytes, ...) \
+  if (cache_type == 0) {                                                  \
+    constexpr CacheType cache_type_now = CacheType::CacheT;               \
+    constexpr size_t cache_bytes = 16;                                    \
+    __VA_ARGS__                                                           \
+  } else if (cache_type == 1) {                                           \
+    constexpr CacheType cache_type_now = CacheType::CacheInt8Hw;          \
+    constexpr size_t cache_bytes = 8;                                     \
+    __VA_ARGS__                                                           \
+  } else if (cache_type == 2) {                                           \
+    constexpr CacheType cache_type_now = CacheType::CacheInt4CwZp;        \
+    constexpr size_t cache_bytes = 4;                                     \
+    __VA_ARGS__                                                           \
+  } else {                                                                \
+    PD_THROW("not support the cache_type: ", cache_type);                 \
+  }
+
+#define DISPATCH_DEAL_EACH_TIME(deal_each_time, DEAL_EACH_TIME, ...) \
+  if (deal_each_time == 16) {                                        \
+    constexpr size_t DEAL_EACH_TIME = 16;                            \
+    __VA_ARGS__                                                      \
+  } else if (deal_each_time == 32) {                                 \
+    constexpr size_t DEAL_EACH_TIME = 32;                            \
+    __VA_ARGS__                                                      \
+  } else {                                                           \
+    PD_THROW("not support the deal_each_time: ", deal_each_time);    \
+  }
+
+#define DISPATCH_GQA_GROUP_SIZE(group_size, GROUP_SIZE, ...) \
+  if (group_size == 1) {                                     \
+    constexpr size_t GROUP_SIZE = 1;                         \
+    __VA_ARGS__                                              \
+  } else if (group_size == 2) {                              \
+    constexpr size_t GROUP_SIZE = 2;                         \
+    __VA_ARGS__                                              \
+  } else if (group_size == 3) {                              \
+    constexpr size_t GROUP_SIZE = 3;                         \
+    __VA_ARGS__                                              \
+  } else if (group_size == 4) {                              \
+    constexpr size_t GROUP_SIZE = 4;                         \
+    __VA_ARGS__                                              \
+  } else if (group_size == 5) {                              \
+    constexpr size_t GROUP_SIZE = 5;                         \
+    __VA_ARGS__                                              \
+  } else if (group_size == 6) {                              \
+    constexpr size_t GROUP_SIZE = 6;                         \
+    __VA_ARGS__                                              \
+  } else if (group_size == 7) {                              \
+    constexpr size_t GROUP_SIZE = 7;                         \
+    __VA_ARGS__                                              \
+  } else if (group_size == 8) {                              \
+    constexpr size_t GROUP_SIZE = 8;                         \
+    __VA_ARGS__                                              \
+  } else if (group_size == 16) {                             \
+    constexpr size_t GROUP_SIZE = 16;                        \
+    __VA_ARGS__                                              \
+  } else {                                                   \
+    PD_THROW("not support the group_size: ", group_size);    \
+  }
+
+#define DISPATCH_MLA_GROUP_SIZE(group_size, GROUP_SIZE, ...) \
+  if (group_size == 1) {                                     \
+    constexpr size_t GROUP_SIZE = 1;                         \
+    __VA_ARGS__                                              \
+  } else if (group_size == 2) {                              \
+    constexpr size_t GROUP_SIZE = 2;                         \
+    __VA_ARGS__                                              \
+  } else if (group_size == 4) {                              \
+    constexpr size_t GROUP_SIZE = 4;                         \
+    __VA_ARGS__                                              \
+  } else if (group_size == 8) {                              \
+    constexpr size_t GROUP_SIZE = 8;                         \
+    __VA_ARGS__                                              \
+  } else if (group_size == 16) {                             \
+    constexpr size_t GROUP_SIZE = 16;                        \
+    __VA_ARGS__                                              \
+  } else if (group_size == 128) {                            \
+    constexpr size_t GROUP_SIZE = 128;                       \
+    __VA_ARGS__                                              \
+  } else {                                                   \
+    PD_THROW("not support the group_size: ", group_size);    \
+  }
+
+#define DISPATCH_BLOCKSHAPE_Q(block_shape_q, BLOCK_SHAPE_Q, NUM_WARP_Q, ...) \
+  if (block_shape_q <= 16) {                                                 \
+    constexpr size_t BLOCK_SHAPE_Q = 16;                                     \
+    constexpr size_t NUM_WARP_Q = 1;                                         \
+    __VA_ARGS__                                                              \
+  } else if (block_shape_q <= 64) {                                          \
+    constexpr size_t BLOCK_SHAPE_Q = 64;                                     \
+    constexpr size_t NUM_WARP_Q = 4;                                         \
+    __VA_ARGS__                                                              \
+  } else {                                                                   \
+    PD_THROW("not support the block_shape_q: ", block_shape_q);              \
+  }
+
+#define DISPATCH_CAUSAL(causal, CAUSAL, ...)      \
+  if (causal) {                                   \
+    constexpr bool CAUSAL = true;                 \
+    __VA_ARGS__                                   \
+  } else {                                        \
+    PD_THROW("not support the causal: ", causal); \
+  }
+
+#define DISPATCH_ENABLE_PREFILL(enable_prefill, ENABLE_PREFILL, ...) \
+  if (enable_prefill) {                                              \
+    constexpr bool ENABLE_PREFILL = 1;                               \
+    __VA_ARGS__                                                      \
+  } else {                                                           \
+    constexpr bool ENABLE_PREFILL = 0;                               \
+    __VA_ARGS__                                                      \
+  }
+
+#define DISPATCH_BLOCK_SIZE(block_size, BLOCK_SIZE, ...)  \
+  if (block_size == 64) {                                 \
+    constexpr size_t BLOCK_SIZE = 64;                     \
+    __VA_ARGS__                                           \
+  } else {                                                \
+    PD_THROW("not support the block_size: ", block_size); \
+  }
+
+#define DISPATCH_BLOCKSHAPE_Q_SYSTEM(                           \
+    block_shape_q, BLOCK_SHAPE_Q, NUM_WARP_Q, ...)              \
+  if (block_shape_q <= 16) {                                    \
+    constexpr size_t BLOCK_SHAPE_Q = 16;                        \
+    constexpr size_t NUM_WARP_Q = 1;                            \
+    __VA_ARGS__                                                 \
+  } else if (block_shape_q <= 32) {                             \
+    constexpr size_t BLOCK_SHAPE_Q = 32;                        \
+    constexpr size_t NUM_WARP_Q = 1;                            \
+    __VA_ARGS__                                                 \
+  } else {                                                      \
+    PD_THROW("not support the block_shape_q: ", block_shape_q); \
+  }
+
+template <typename T>
+inline HOSTDEVICE T roundWithTiesToEven(T x) {
+  T xLower = floor(x);
+  T xUpper = ceil(x);
+  // x is in interval [xl,xu]. Choose closest of two bounds, breaking ties to
+  // even.
+  T dLower = x - xLower;
+  T dUpper = xUpper - x;
+  return static_cast<T>(
+      (dLower == dUpper ? fmod(xLower, 2.0F) == 0.0F : dLower < dUpper)
+          ? xLower
+          : xUpper);
+}

csrc/gpu/cpp_extensions.cu

@@ -0,0 +1,365 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/extension.h"
+#include "all_reduce.h"
+
+namespace py = pybind11;
+
+std::vector<paddle::Tensor> AppendAttention(
+    const paddle::Tensor& qkv,
+    const paddle::Tensor& key_cache,
+    const paddle::Tensor& value_cache,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& encoder_batch_ids,
+    const paddle::Tensor& encoder_tile_ids_per_batch,
+    const paddle::Tensor& encoder_num_blocks,
+    const paddle::Tensor& kv_batch_ids,
+    const paddle::Tensor& kv_tile_ids_per_batch,
+    const paddle::Tensor& kv_num_blocks,
+    const paddle::Tensor& decoder_batch_ids,
+    const paddle::Tensor& decoder_tile_ids_per_batch,
+    const paddle::Tensor& decoder_num_blocks,
+    const paddle::Tensor& max_enc_len_this_time,
+    const paddle::Tensor& max_dec_len_this_time,
+    const paddle::Tensor& max_len_kv,
+    const paddle::optional<paddle::Tensor>& rotary_embs,
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>& qkv_bias,
+    const paddle::optional<paddle::Tensor>& qkv_out_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_quant_scales,
+    const paddle::optional<paddle::Tensor>& cache_v_quant_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_dequant_scales,
+    const paddle::optional<paddle::Tensor>& cache_v_dequant_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const paddle::optional<paddle::Tensor>& out_linear_shifts,
+    const paddle::optional<paddle::Tensor>& out_linear_smooths,
+    const paddle::optional<paddle::Tensor>& excess_blocks,
+    const std::string& compute_dtype,
+    const std::string& cache_quant_type_str,
+    const bool use_neox_rotary_style,
+    const int max_input_length,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float out_linear_in_scale,
+    const int speculate_max_draft_token_num,
+    const bool causal,
+    const bool speculate_decoder);
+
+void FusedRotaryPositionEncoding(
+    paddle::Tensor& query,  // [num_tokens, num_heads, head_size] or
+                            // [num_tokens, num_heads * head_size]
+    paddle::Tensor& key,
+    // [num_tokens, num_kv_heads, head_size] or [num_tokens, num_kv_heads *
+    // head_size]
+    const paddle::Tensor& position_ids,   // [num_tokens]
+    const paddle::Tensor& cos_sin_cache,  // [max_position, rot_dim]
+    int head_size,
+    bool is_neox);
+
+std::vector<paddle::Tensor> MultiHeadLatentAttention(
+    const paddle::Tensor& query,
+    const paddle::Tensor& key_cache,
+    const paddle::Tensor& value_cache,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& encoder_batch_ids,
+    const paddle::Tensor& encoder_tile_ids_per_batch,
+    const paddle::Tensor& encoder_num_blocks,
+    const paddle::Tensor& kv_batch_ids,
+    const paddle::Tensor& kv_tile_ids_per_batch,
+    const paddle::Tensor& kv_num_blocks,
+    const paddle::Tensor& decoder_batch_ids,
+    const paddle::Tensor& decoder_tile_ids_per_batch,
+    const paddle::Tensor& decoder_num_blocks,
+    const paddle::Tensor& decoder_num_blocks_cpu,
+    const paddle::Tensor& decoder_chunk_size_cpu,
+    const paddle::Tensor& max_enc_len_this_time,
+    const paddle::Tensor& max_dec_len_this_time,
+    const paddle::Tensor& max_len_kv,
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>& query_bias,
+    const paddle::optional<paddle::Tensor>& query_out_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_quant_scales,
+    const paddle::optional<paddle::Tensor>& cache_v_quant_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_dequant_scales,
+    const paddle::optional<paddle::Tensor>& cache_v_dequant_scales,
+    const paddle::optional<paddle::Tensor>& cache_k_zp,
+    const paddle::optional<paddle::Tensor>& cache_v_zp,
+    const paddle::optional<paddle::Tensor>& out_linear_shifts,
+    const paddle::optional<paddle::Tensor>& out_linear_smooths,
+    const std::string& compute_dtype,
+    const std::string& cache_quant_type_str,
+    const int nope_size,
+    const int max_input_length,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float out_linear_in_scale,
+    const int speculate_draft_total_token_num,
+    const bool causal,
+    const bool speculate_decoder);
+
+std::vector<paddle::Tensor> GetBlockShapeAndSplitKVBlock(
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& max_enc_len_this_time,
+    const paddle::Tensor& max_dec_len_this_time,
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& cum_offsets,
+    const int group_size,
+    const int block_size,
+    const int decoder_step_token_num);
+
+std::vector<paddle::Tensor> NoauxTc(paddle::Tensor& scores,
+                                    paddle::Tensor& scores_with_bias,
+                                    int n_group,
+                                    int topk_group,
+                                    int topk,
+                                    float routed_scaling_factor);
+
+std::vector<paddle::Tensor> PrefillMLAWriteCacheKernel(
+    const paddle::Tensor& kv_nope,
+    const paddle::Tensor& kv_pe,
+    const paddle::Tensor& kv_cache,
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const std::string& cache_quant_type_str,
+    const int max_seq_len);
+
+std::vector<paddle::Tensor> DecodeMLAWriteCacheKernel(
+    const paddle::Tensor& kv_nope,
+    const paddle::Tensor& kv_pe,
+    const paddle::Tensor& kv_cache,
+    const paddle::Tensor& seq_lens,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const std::string& cache_quant_type_str,
+    const int max_seq_len,
+    const bool speculate_decoder);
+
+paddle::Tensor cutlass_fp8_fp8_half_block_gemm_fused_func(
+    const paddle::Tensor& x,
+    const paddle::Tensor& y,
+    const paddle::Tensor& x_scale,
+    const paddle::Tensor& y_scale,
+    const paddle::optional<paddle::Tensor>& bias,
+    bool trans_x,
+    bool trans_y,
+    std::string output_dtype,
+    std::string activation_type);
+
+void GetPositionIdsAndMaskEncoderBatch(
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& position_ids,
+    const paddle::Tensor& mask_encoder_batch);
+
+void SetPreidsTokenPenaltyMultiScores(const paddle::Tensor& pre_ids,
+                                      const paddle::Tensor& input_ids,
+                                      const paddle::Tensor& seq_lens_encoder,
+                                      const paddle::Tensor& seq_lens_decoder,
+                                      const paddle::Tensor& step_idx,
+                                      const paddle::Tensor& stop_flags,
+                                      const paddle::Tensor& logits,
+                                      const paddle::Tensor& penalty_scores,
+                                      const paddle::Tensor& frequency_scores,
+                                      const paddle::Tensor& presence_scores,
+                                      const paddle::Tensor& temperatures,
+                                      const paddle::Tensor& bad_tokens,
+                                      const paddle::Tensor& cur_len,
+                                      const paddle::Tensor& min_len,
+                                      const paddle::Tensor& eos_token_id);
+
+void UpdateInputesV2(const paddle::Tensor& stop_flags,
+               const paddle::Tensor& step_idx,
+               const paddle::Tensor& not_need_stop, // cpu
+               const paddle::Tensor& seq_lens_this_time,
+               const paddle::Tensor& seq_lens_encoder,
+               const paddle::Tensor& seq_lens_decoder,
+               const paddle::Tensor& max_dec_len,
+               const paddle::Tensor& input_ids,
+               const paddle::Tensor& stop_nums,
+               const paddle::Tensor& next_tokens,
+               const paddle::Tensor& is_block_step,
+               const paddle::Tensor& end_ids,
+               const paddle::Tensor& kwargs_next_tokens);
+
+paddle::Tensor RebuildPaddingV2Func(const paddle::Tensor& tmp_out, // [token_num, dim_embed]
+                                    const paddle::Tensor& cum_offsets, // [bsz, 1]
+                                    const paddle::Tensor& seq_lens_decoder,
+                                    const paddle::Tensor& seq_lens_encoder,
+                                    const paddle::optional<paddle::Tensor>& output_padding_offset,
+                                    int max_input_length);
+
+std::vector<paddle::Tensor> PerTokenGroupQuant(const paddle::Tensor& x,
+                                        const int group_size,
+                                        const bool transpose_scale,
+                                        const float quant_max_bound,
+                                        const float quant_min_bound);
+
+std::vector<paddle::Tensor> PerTensorQuantFp8(const paddle::Tensor& x, const paddle::optional<paddle::Tensor>& scale);
+
+std::vector<paddle::Tensor> GetPaddingOffsetV2(const paddle::Tensor& input_ids,
+                                               const paddle::Tensor& cum_offsets,
+                                               const paddle::Tensor& token_num,
+                                               const paddle::Tensor& seq_len,
+                                               const paddle::optional<paddle::Tensor>& draft_tokens,
+                                               const paddle::optional<paddle::Tensor>& seq_lens_encoder);
+
+void SaveOutMmsg(const paddle::Tensor& x,
+                 const paddle::Tensor& not_need_stop, // cpu
+                 int64_t rank_id);
+
+void GetOutput(const paddle::Tensor& x,
+               int64_t rank_id,
+               bool wait_flag);
+
+void StepPaddle(const paddle::Tensor &stop_flags,
+                const paddle::Tensor &seq_lens_this_time,
+                const paddle::Tensor &ori_seq_lens_encoder,
+                const paddle::Tensor &seq_lens_encoder,
+                const paddle::Tensor &seq_lens_decoder,
+                const paddle::Tensor &block_tables,  // [bsz, block_num_per_seq]
+                const paddle::Tensor &encoder_block_lens,
+                const paddle::Tensor &is_block_step,
+                const paddle::Tensor &step_block_list,
+                const paddle::Tensor &step_lens,
+                const paddle::Tensor &recover_block_list,
+                const paddle::Tensor &recover_lens,
+                const paddle::Tensor &need_block_list,
+                const paddle::Tensor &need_block_len,
+                const paddle::Tensor &used_list_len,
+                const paddle::Tensor &free_list,
+                const paddle::Tensor &free_list_len,
+                const paddle::Tensor &input_ids,
+                const paddle::Tensor &pre_ids,
+                const paddle::Tensor &step_idx,
+                const paddle::Tensor &next_tokens,
+                const paddle::Tensor &first_token_ids,
+                const int block_size,
+                const int encoder_decoder_block_num);
+
+void SaveOutputDygraph(
+    const paddle::Tensor& all_token_ids,
+    const paddle::Tensor& tokens,
+    const paddle::Tensor& result_ids,
+    const paddle::Tensor& step_idx
+);
+
+PYBIND11_MODULE(paddlenlp_ops, m) {
+  /**
+   * all_reduce.cu
+   */
+  m.def("init_custom_all_reduce", &init_custom_all_reduce, "init all reduce class function");
+  m.def("all_reduce", &all_reduce, "all reduce function");
+  m.def("dispose", &dispose, "del function for python");
+  m.def("meta_size", &meta_size, "meta_size function for Signal struct");
+  m.def("register_buffer", &register_buffer, "register ipc buffer");
+  m.def("f_append_attention", &AppendAttention, "AppendAttention");
+  m.def("f_fused_rotary_position_encoding", &FusedRotaryPositionEncoding, "FusedRotaryPositionEncoding");
+  m.def("f_multi_head_latent_attention", &MultiHeadLatentAttention, "MultiHeadLatentAttention");
+  m.def("f_noaux_tc", &NoauxTc, "NoauxTc");
+  m.def("f_get_block_shape_and_split_kv_block", &GetBlockShapeAndSplitKVBlock, "GetBlockShapeAndSplitKVBlock");
+  m.def("f_prefill_mla_write_cache", &PrefillMLAWriteCacheKernel, "PrefillMLAWriteCacheKernel");
+  m.def("f_decode_mla_write_cache", &DecodeMLAWriteCacheKernel, "DecodeMLAWriteCacheKernel");
+  m.def("f_get_position_ids_and_mask_encoder_batch", &GetPositionIdsAndMaskEncoderBatch, "GetPositionIdsAndMaskEncoderBatch");
+  m.def("f_set_preids_token_penalty_multi_scores", &SetPreidsTokenPenaltyMultiScores, "SetPreidsTokenPenaltyMultiScores");
+  m.def("f_update_inputs_v2", &UpdateInputesV2, "UpdateInputesV2");
+  m.def("f_rebuild_padding_v2", &RebuildPaddingV2Func, "RebuildPaddingV2Func");
+  m.def("f_per_token_group_quant", &PerTokenGroupQuant, "PerTokenGroupQuant");
+  m.def("f_per_tensor_quant_fp8", &PerTensorQuantFp8, "PerTensorQuantFp8");
+  m.def("f_get_padding_offset_v2", &GetPaddingOffsetV2, "GetPaddingOffsetV2");
+  m.def("f_save_output", &SaveOutMmsg, "SaveOutMmsg");
+  m.def("f_get_output", &GetOutput, "GetOutput");
+  m.def("f_step_paddle", &StepPaddle, "StepPaddle");
+  m.def("f_save_output_dygraph", &SaveOutputDygraph, "SaveOutputDygraph");
+//   m.def("f_cutlass_fp8_fp8_half_block_gemm_fused", &cutlass_fp8_fp8_half_block_gemm_fused_func, "cutlass_fp8_fp8_half_block_gemm_fused_func");
+}
+
+PYBIND11_MODULE(paddlenlp_ops_80, m) {
+  /**
+   * all_reduce.cu
+   */
+  m.def("init_custom_all_reduce", &init_custom_all_reduce, "init all reduce class function");
+  m.def("all_reduce", &all_reduce, "all reduce function");
+  m.def("dispose", &dispose, "del function for python");
+  m.def("meta_size", &meta_size, "meta_size function for Signal struct");
+  m.def("register_buffer", &register_buffer, "register ipc buffer");
+  m.def("f_append_attention", &AppendAttention, "AppendAttention");
+  m.def("f_fused_rotary_position_encoding", &FusedRotaryPositionEncoding, "FusedRotaryPositionEncoding");
+  m.def("f_multi_head_latent_attention", &MultiHeadLatentAttention, "MultiHeadLatentAttention");
+  m.def("f_noaux_tc", &NoauxTc, "NoauxTc");
+  m.def("f_get_block_shape_and_split_kv_block", &GetBlockShapeAndSplitKVBlock, "GetBlockShapeAndSplitKVBlock");
+  m.def("f_prefill_mla_write_cache", &PrefillMLAWriteCacheKernel, "PrefillMLAWriteCacheKernel");
+  m.def("f_decode_mla_write_cache", &DecodeMLAWriteCacheKernel, "DecodeMLAWriteCacheKernel");
+  m.def("f_get_position_ids_and_mask_encoder_batch", &GetPositionIdsAndMaskEncoderBatch, "GetPositionIdsAndMaskEncoderBatch");
+  m.def("f_set_preids_token_penalty_multi_scores", &SetPreidsTokenPenaltyMultiScores, "SetPreidsTokenPenaltyMultiScores");
+  m.def("f_update_inputs_v2", &UpdateInputesV2, "UpdateInputesV2");
+  m.def("f_rebuild_padding_v2", &RebuildPaddingV2Func, "RebuildPaddingV2Func");
+  m.def("f_per_token_group_quant", &PerTokenGroupQuant, "PerTokenGroupQuant");
+  m.def("f_per_tensor_quant_fp8", &PerTensorQuantFp8, "PerTensorQuantFp8");
+  m.def("f_get_padding_offset_v2", &GetPaddingOffsetV2, "GetPaddingOffsetV2");
+  m.def("f_save_output", &SaveOutMmsg, "SaveOutMmsg");
+  m.def("f_get_output", &GetOutput, "GetOutput");
+  m.def("f_step_paddle", &StepPaddle, "StepPaddle");
+  m.def("f_save_output_dygraph", &SaveOutputDygraph, "SaveOutputDygraph");
+}
+
+PYBIND11_MODULE(paddlenlp_ops_90, m) {
+  /**
+   * all_reduce.cu
+   */
+  m.def("init_custom_all_reduce", &init_custom_all_reduce, "init all reduce class function");
+  m.def("all_reduce", &all_reduce, "all reduce function");
+  m.def("dispose", &dispose, "del function for python");
+  m.def("meta_size", &meta_size, "meta_size function for Signal struct");
+  m.def("register_buffer", &register_buffer, "register ipc buffer");
+  m.def("f_append_attention", &AppendAttention, "AppendAttention");
+  m.def("f_fused_rotary_position_encoding", &FusedRotaryPositionEncoding, "FusedRotaryPositionEncoding");
+  m.def("f_multi_head_latent_attention", &MultiHeadLatentAttention, "MultiHeadLatentAttention");
+  m.def("f_noaux_tc", &NoauxTc, "NoauxTc");
+  m.def("f_get_block_shape_and_split_kv_block", &GetBlockShapeAndSplitKVBlock, "GetBlockShapeAndSplitKVBlock");
+  m.def("f_prefill_mla_write_cache", &PrefillMLAWriteCacheKernel, "PrefillMLAWriteCacheKernel");
+  m.def("f_decode_mla_write_cache", &DecodeMLAWriteCacheKernel, "DecodeMLAWriteCacheKernel");
+  m.def("f_get_position_ids_and_mask_encoder_batch", &GetPositionIdsAndMaskEncoderBatch, "GetPositionIdsAndMaskEncoderBatch");
+  m.def("f_set_preids_token_penalty_multi_scores", &SetPreidsTokenPenaltyMultiScores, "SetPreidsTokenPenaltyMultiScores");
+  m.def("f_update_inputs_v2", &UpdateInputesV2, "UpdateInputesV2");
+  m.def("f_rebuild_padding_v2", &RebuildPaddingV2Func, "RebuildPaddingV2Func");
+  m.def("f_per_token_group_quant", &PerTokenGroupQuant, "PerTokenGroupQuant");
+  m.def("f_per_tensor_quant_fp8", &PerTensorQuantFp8, "PerTensorQuantFp8");
+  m.def("f_get_padding_offset_v2", &GetPaddingOffsetV2, "GetPaddingOffsetV2");
+  m.def("f_save_output", &SaveOutMmsg, "SaveOutMmsg");
+  m.def("f_get_output", &GetOutput, "GetOutput");
+  m.def("f_step_paddle", &StepPaddle, "StepPaddle");
+  m.def("f_save_output_dygraph", &SaveOutputDygraph, "SaveOutputDygraph");
+}

csrc/gpu/cutlass_kernels/cutlass_helper.h

@@ -0,0 +1,113 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+// 
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// 
+//     http://www.apache.org/licenses/LICENSE-2.0
+// 
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include <mutex>
+
+#include "helper.h"
+#include "cutlass/half.h"
+#include "cutlass/bfloat16.h"
+#include "paddle/extension.h"
+
+template <paddle::DataType D>
+class CutlassDtypeTraits;
+
+template <>
+class CutlassDtypeTraits<paddle::DataType::FLOAT32> {
+public:
+  typedef float DataType;
+  typedef float data_t;
+};
+
+template <>
+class CutlassDtypeTraits<paddle::DataType::FLOAT16> {
+public:
+  typedef cutlass::half_t DataType;
+  typedef paddle::float16 data_t;
+};
+
+template <>
+class CutlassDtypeTraits<paddle::DataType::BFLOAT16> {
+public:
+  typedef cutlass::bfloat16_t DataType;
+  typedef paddle::bfloat16 data_t;
+};
+
+class CutlassGemmConfigMannager {
+public:
+    static CutlassGemmConfigMannager& getInstance() {
+        static CutlassGemmConfigMannager instance;
+        return instance;
+    }
+
+    CutlassGemmConfigMannager(const CutlassGemmConfigMannager&) = delete;
+    CutlassGemmConfigMannager& operator=(const CutlassGemmConfigMannager&) = delete;
+
+    void up_date_configs(const nlohmann::json& j){
+      std::lock_guard<std::mutex> lock(mutex_);
+      for (auto it = j.begin(); it != j.end(); ++it) {
+          json_[it.key()] = it.value();
+      }
+    }
+
+    nlohmann::json* get_gemm_best_configs(const std::string & config_file_path) {
+      if (!load_initialized_) {
+        std::ifstream file(config_file_path);
+        if(!file.good()){
+            throw std::runtime_error("cutlass gemm_best_config can not be found, please set gemm_best_config'path as FLAGS_use_cutlass_device_best_config_path, or unset FLAGS_use_cutlass_device_best_config_path to tune gemm_best_config");
+        }
+        json_ = ReadJsonFromFile(config_file_path);
+        load_initialized_ = true;
+        save_initialized_ = false;
+      } 
+      return &json_;
+    }
+
+private:
+    void save_gemm_best_configs_(const std::string & config_file_path) {
+      std::ifstream file(config_file_path);
+      if(!file.good()){
+        std::ofstream new_file(config_file_path);
+        new_file << json_.dump(4);
+        new_file.close();
+      } else {
+        nlohmann::json old_json = ReadJsonFromFile(config_file_path);
+        for (auto it = json_.begin(); it != json_.end(); ++it) {
+            old_json[it.key()] = it.value();
+        }
+        json_ = old_json;
+        std::ofstream new_file(config_file_path, std::ios::out | std::ios::trunc);
+        new_file << json_.dump(4);
+        new_file.close();
+        file.close();
+      }
+      return;
+    }
+
+    CutlassGemmConfigMannager() : json_(nullptr), load_initialized_(false) , save_initialized_(true){}
+    ~CutlassGemmConfigMannager() {
+      std::lock_guard<std::mutex> lock(mutex_);
+      if(save_initialized_){
+        std::string config_file_path = "fp8_fuse_gemm_config.json";
+        save_gemm_best_configs_(config_file_path);
+      }
+      save_initialized_=true;
+      load_initialized_=false;
+      json_.clear();
+    }
+    mutable std::mutex mutex_;
+    nlohmann::json json_;
+    bool load_initialized_;
+    bool save_initialized_;
+};

csrc/gpu/cutlass_kernels/fp8_gemm_fused/dual_gemm/device/dual_gemm.h

@@ -0,0 +1,544 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Performs a dual gemm in one fused kernel:
+```
+D0 = epilogue0(X @ B0, C0)
+D1 = epilogue1(X @ B1, C1)
+D2 = element_wise(D0, D1)
+```
+*/
+
+#pragma once
+
+#include "cutlass/arch/arch.h"
+#include "cutlass/cutlass.h"
+#include "cutlass/device_kernel.h"
+#include "cutlass/numeric_types.h"
+
+#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
+
+#include "cutlass/epilogue/thread/linear_combination_relu.h"
+#include "cutlass/epilogue/threadblock/default_epilogue_tensor_op.h"
+#include "cutlass/gemm/device/default_gemm_configuration.h"
+#include "cutlass/gemm/threadblock/default_mma.h"
+
+#include "fp8_gemm_fused/dual_gemm/dual_gemm_common.h"
+#include "fp8_gemm_fused/dual_gemm/kernel/dual_gemm.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace device {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B0 matrix operand
+    typename LayoutB0_,
+    /// Layout type for B1 matrix operand
+    typename LayoutB1_,
+    /// Element type for C matrix operands
+    typename ElementC_,
+    /// Layout type for D matrix operands
+    typename ElementD_,
+    /// Layout type for C and D matrix operands
+    typename LayoutC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator_,
+    /// Operator class tag
+    typename OperatorClass_,
+    /// Tag indicating architecture to tune for
+    typename ArchTag_,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape_,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape_,
+    /// Instruction-level tile size (concept: GemmShape)
+    typename InstructionShape_,
+    /// Epilogue output operator
+    typename EpilogueOutputOp0_,
+    typename EpilogueOutputOp1_,
+    typename EpilogueOutputOp2_,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle_ =
+        threadblock::GemmIdentityThreadblockSwizzle<>,
+    /// Number of stages used in the pipelined mainloop
+    int Stages = DefaultGemmConfiguration<OperatorClass_,
+                                          ArchTag_,
+                                          ElementA_,
+                                          ElementB_,
+                                          ElementC_,
+                                          ElementAccumulator_>::kStages,
+    bool StoreD0 = true,
+    bool StoreD1 = true,
+    /// If true, kernel supports split-K with serial reduction
+    bool SplitKSerial = false,
+    /// Access granularity of A matrix in units of elements
+    int AlignmentA = DefaultGemmConfiguration<OperatorClass_,
+                                              ArchTag_,
+                                              ElementA_,
+                                              ElementB_,
+                                              ElementC_,
+                                              ElementAccumulator_>::kAlignmentA,
+    /// Access granularity of B matrix in units of elements
+    int AlignmentB = DefaultGemmConfiguration<OperatorClass_,
+                                              ArchTag_,
+                                              ElementA_,
+                                              ElementB_,
+                                              ElementC_,
+                                              ElementAccumulator_>::kAlignmentB,
+    /// Operation performed by GEMM
+    typename Operator_ =
+        typename DefaultGemmConfiguration<OperatorClass_,
+                                          ArchTag_,
+                                          ElementA_,
+                                          ElementB_,
+                                          ElementC_,
+                                          ElementAccumulator_>::Operator>
+class DualGemm {
+ public:
+  using ElementA = ElementA_;
+  using LayoutA = LayoutA_;
+  using TensorRefA = TensorRef<ElementA const, LayoutA>;
+  using ElementB = ElementB_;
+  using LayoutB0 = LayoutB0_;
+  using LayoutB1 = LayoutB1_;
+  using TensorRefB0 = TensorRef<ElementB const, LayoutB0>;
+  using TensorRefB1 = TensorRef<ElementB const, LayoutB1>;
+  using ElementC = ElementC_;
+  using ElementD = ElementD_;
+  using LayoutC = LayoutC_;
+  using TensorRefC = TensorRef<ElementC const, LayoutC>;
+  using TensorRefD = TensorRef<ElementD, LayoutC>;
+  using ElementAccumulator = ElementAccumulator_;
+  using OperatorClass = OperatorClass_;
+  using ArchTag = ArchTag_;
+  using ThreadblockShape = ThreadblockShape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = InstructionShape_;
+  using EpilogueOutputOp0 = EpilogueOutputOp0_;
+  using EpilogueOutputOp1 = EpilogueOutputOp1_;
+  using EpilogueOutputOp2 = EpilogueOutputOp2_;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  using Operator = Operator_;
+  static int const kStages = Stages;
+  static int const kAlignmentA = AlignmentA;
+  static int const kAlignmentB = AlignmentB;
+  static int const kAlignmentC = EpilogueOutputOp1::kCount;
+  static bool const kSplitKSerial = SplitKSerial;
+  static bool constexpr kStoreD0 = StoreD0;
+  static bool constexpr kStoreD1 = StoreD1;
+  static ComplexTransform const kTransformA = ComplexTransform::kNone;
+  static ComplexTransform const kTransformB = ComplexTransform::kNone;
+
+  using LayoutScaleBias = layout::RowMajor;
+  /// Define the kernel
+  /// Define the threadblock-scoped matrix multiply-accumulate
+  static_assert(ArchTag::kMinComputeCapability >= 80,
+                "Only multistage is implemented");
+  static_assert(kStages >= 3, "Only multistage is implemented");
+  using Mma0 =
+      typename cutlass::gemm::threadblock::DefaultMma<ElementA,
+                                                      LayoutA,
+                                                      kAlignmentA,
+                                                      ElementB,
+                                                      LayoutB0,
+                                                      kAlignmentB,
+                                                      ElementAccumulator,
+                                                      layout::RowMajor,
+                                                      arch::OpClassTensorOp,
+                                                      ArchTag,
+                                                      ThreadblockShape,
+                                                      WarpShape,
+                                                      InstructionShape,
+                                                      Stages,
+                                                      Operator>::ThreadblockMma;
+  using Mma1 =
+      typename cutlass::gemm::threadblock::DefaultMma<ElementA,
+                                                      LayoutA,
+                                                      kAlignmentA,
+                                                      ElementB,
+                                                      LayoutB1,
+                                                      kAlignmentB,
+                                                      ElementAccumulator,
+                                                      layout::RowMajor,
+                                                      arch::OpClassTensorOp,
+                                                      ArchTag,
+                                                      ThreadblockShape,
+                                                      WarpShape,
+                                                      InstructionShape,
+                                                      Stages,
+                                                      Operator>::ThreadblockMma;
+  using DualMma =
+      threadblock::DualMmaMultistage<typename Mma0::Shape,
+                                     typename Mma0::IteratorA,
+                                     typename Mma0::SmemIteratorA,
+                                     Mma0::kCacheOpA,
+                                     typename Mma0::IteratorB,
+                                     typename Mma0::SmemIteratorB,
+                                     Mma0::kCacheOpB,
+                                     typename Mma1::IteratorB,
+                                     typename Mma1::SmemIteratorB,
+                                     typename Mma0::ElementC,
+                                     typename Mma0::LayoutC,
+                                     typename Mma0::Policy,
+                                     typename Mma1::Policy,
+                                     Mma0::kStages,
+                                     SharedMemoryClearOption::kNone>;
+
+  static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK;
+
+  /// Define the epilogue
+  using Epilogue0 =
+      typename cutlass::epilogue::threadblock::DefaultEpilogueTensorOp<
+          ThreadblockShape,
+          typename DualMma::Operator0,
+          kPartitionsK,
+          EpilogueOutputOp0,
+          EpilogueOutputOp0::kCount>::Epilogue;
+  using Epilogue1 =
+      typename cutlass::epilogue::threadblock::DefaultEpilogueTensorOp<
+          ThreadblockShape,
+          typename DualMma::Operator1,
+          kPartitionsK,
+          EpilogueOutputOp1,
+          EpilogueOutputOp1::kCount>::Epilogue;
+  using Epilogue2 =
+      typename cutlass::epilogue::threadblock::DefaultEpilogueTensorOp<
+          ThreadblockShape,
+          typename DualMma::Operator1,
+          kPartitionsK,
+          EpilogueOutputOp2,
+          EpilogueOutputOp2::kCount>::Epilogue;
+
+  /// Define the kernel-level GEMM operator.
+  using DualGemmKernel = kernel::DualGemm<DualMma,
+                                          Epilogue0,
+                                          Epilogue1,
+                                          Epilogue2,
+                                          EpilogueOutputOp2,
+                                          ThreadblockSwizzle,
+                                          kSplitKSerial,
+                                          kStoreD0,
+                                          kStoreD1>;
+
+  /// Argument structure
+  struct Arguments {
+    //
+    // Data members
+    //
+
+    DualGemmMode mode;
+    GemmCoord problem_size;
+    TensorRef<ElementA const, LayoutA> ref_A0;
+    TensorRef<ElementB const, LayoutB0> ref_B0;
+    TensorRef<ElementC const, LayoutC> ref_C0;
+    TensorRef<ElementC, LayoutC> ref_D0;
+    TensorRef<ElementB const, LayoutB1> ref_B1;
+    TensorRef<ElementC const, LayoutC> ref_C1;
+    TensorRef<ElementC, LayoutC> ref_D1;
+    TensorRef<ElementD, LayoutC> ref_D2;
+    typename EpilogueOutputOp0::Params epilogue0;
+    typename EpilogueOutputOp1::Params epilogue1;
+    typename EpilogueOutputOp2::Params epilogue2;
+    int split_k_slices;
+
+    int batch_count;
+    int64_t batch_stride_A;
+    int64_t batch_stride_B0;
+    int64_t batch_stride_B1;
+    int64_t batch_stride_C;
+    int64_t batch_stride_D;
+
+    //
+    // Methods
+    //
+
+    /// Default ctor
+    CUTLASS_HOST_DEVICE
+    Arguments() : problem_size(0, 0, 0), split_k_slices(1) {}
+
+    /// Constructs an Arguments structure
+    CUTLASS_HOST_DEVICE
+    Arguments(DualGemmMode mode,
+              GemmCoord problem_size_,
+              TensorRef<ElementA const, LayoutA> ref_A0_,
+              TensorRef<ElementB const, LayoutB0> ref_B0_,
+              TensorRef<ElementC const, LayoutC> ref_C0_,
+              TensorRef<ElementC, LayoutC> ref_D0_,
+              TensorRef<ElementB const, LayoutB1> ref_B1_,
+              TensorRef<ElementC const, LayoutC> ref_C1_,
+              TensorRef<ElementC, LayoutC> ref_D1_,
+              TensorRef<ElementD, LayoutC> ref_D2_,
+              typename EpilogueOutputOp0::Params epilogue0_ =
+                  typename EpilogueOutputOp0::Params(),
+              typename EpilogueOutputOp1::Params epilogue1_ =
+                  typename EpilogueOutputOp1::Params(),
+              typename EpilogueOutputOp2::Params epilogue2_ =
+                  typename EpilogueOutputOp2::Params(),
+              int split_k_slices_ = 1,
+              int batch_count = 1,
+              int64_t batch_stride_A = 0,
+              int64_t batch_stride_B0 = 0,
+              int64_t batch_stride_B1 = 0,
+              int64_t batch_stride_C = 0,
+              int64_t batch_stride_D = 0)
+        : mode(mode),
+          problem_size(problem_size_),
+          ref_A0(ref_A0_),
+          ref_B0(ref_B0_),
+          ref_C0(ref_C0_),
+          ref_D0(ref_D0_),
+          ref_B1(ref_B1_),
+          ref_C1(ref_C1_),
+          ref_D1(ref_D1_),
+          ref_D2(ref_D2_),
+          epilogue0(epilogue0_),
+          epilogue1(epilogue1_),
+          epilogue2(epilogue2_),
+          split_k_slices(split_k_slices_),
+          batch_count(batch_count),
+          batch_stride_A(batch_stride_A),
+          batch_stride_B0(batch_stride_B0),
+          batch_stride_B1(batch_stride_B1),
+          batch_stride_C(batch_stride_C),
+          batch_stride_D(batch_stride_D) {}
+  };
+
+ private:
+  /// Kernel parameters object
+  typename DualGemmKernel::Params params_;
+
+ public:
+  /// Constructs the GEMM.
+  DualGemm() = default;
+
+  /// Determines whether the GEMM can execute the given problem.
+  static Status can_implement(Arguments const &args) {
+    if (args.mode == DualGemmMode::kBatched && kSplitKSerial) {
+      return Status::kErrorInvalidProblem;
+    }
+    if (!kSplitKSerial && args.split_k_slices > 1) {
+      return Status::kErrorInvalidProblem;
+    }
+    if (kStoreD0 != (args.ref_D0.data() != nullptr)) {
+      return Status::kErrorInternal;
+    }
+    if (kStoreD1 != (args.ref_D1.data() != nullptr)) {
+      return Status::kErrorInternal;
+    }
+
+    Status status = DualGemmKernel::can_implement(args.problem_size,
+                                                  args.ref_A0.non_const_ref(),
+                                                  args.ref_B0.non_const_ref(),
+                                                  args.ref_C0.non_const_ref(),
+                                                  args.ref_D0,
+                                                  args.ref_B1.non_const_ref(),
+                                                  args.ref_C1.non_const_ref(),
+                                                  args.ref_D1,
+                                                  args.ref_D2);
+
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    return Status::kSuccess;
+  }
+
+  /// Gets the workspace size
+  static size_t get_workspace_size(Arguments const &args) {
+    size_t bytes = 0;
+
+    if (kSplitKSerial && args.split_k_slices > 1) {
+      // Determine grid shape
+      ThreadblockSwizzle threadblock_swizzle;
+
+      cutlass::gemm::GemmCoord tiled_shape =
+          threadblock_swizzle.get_tiled_shape(args.problem_size,
+                                              {ThreadblockShape::kM,
+                                               ThreadblockShape::kN,
+                                               ThreadblockShape::kK},
+                                              args.split_k_slices);
+
+      bytes += sizeof(int) * size_t(tiled_shape.m()) * size_t(tiled_shape.n());
+    }
+
+    return bytes;
+  }
+
+  /// Initializes GEMM state from arguments.
+  Status initialize(Arguments const &args,
+                    void *workspace = nullptr,
+                    cudaStream_t stream = nullptr) {
+    // Determine grid shape
+    ThreadblockSwizzle threadblock_swizzle;
+
+    cutlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape(
+        args.problem_size,
+        {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
+        args.mode == DualGemmMode::kBatched ? args.batch_count
+                                            : args.split_k_slices);
+
+    if (kSplitKSerial) {
+      if (args.split_k_slices > 1) {
+        if (!workspace) {
+          return Status::kErrorWorkspaceNull;
+        }
+
+        size_t bytes = get_workspace_size(args);
+
+        cudaError_t result = cudaMemsetAsync(workspace, 0, bytes, stream);
+
+        if (result != cudaSuccess) {
+          return Status::kErrorInternal;
+        }
+      }
+    } else {
+      if (args.split_k_slices > 1) {
+        return Status::kErrorInvalidProblem;
+      }
+    }
+
+    // Initialize the Params structure
+    params_ = typename DualGemmKernel::Params{
+        args.mode,
+        args.problem_size,
+        grid_shape,
+        args.ref_A0.non_const_ref(),
+        args.ref_B0.non_const_ref(),
+        args.ref_C0.non_const_ref(),
+        args.ref_D0,
+        args.ref_B1.non_const_ref(),
+        args.ref_C1.non_const_ref(),
+        args.ref_D1,
+        args.ref_D2,
+        args.epilogue0,
+        args.epilogue1,
+        args.epilogue2,
+        reinterpret_cast<int *>(workspace),
+        args.batch_stride_A,
+        args.batch_stride_B0,
+        args.batch_stride_B1,
+        args.batch_stride_C,
+        args.batch_stride_D,
+    };
+
+    return Status::kSuccess;
+  }
+
+  /// Lightweight update given a subset of arguments
+  Status update(Arguments const &args, void *workspace = nullptr) {
+    if (kSplitKSerial && args.split_k_slices > 1) {
+      if (!workspace) {
+        return Status::kErrorWorkspaceNull;
+      }
+    }
+
+    params_.ref_A0.reset(args.ref_A0.non_const_ref().data());
+    params_.ref_B0.reset(args.ref_B0.non_const_ref().data());
+    params_.ref_C0.reset(args.ref_C0.non_const_ref().data());
+    params_.ref_D0.reset(args.ref_D0.data());
+    params_.ref_B1.reset(args.ref_B1.non_const_ref().data());
+    params_.ref_C1.reset(args.ref_C1.non_const_ref().data());
+    params_.ref_D1.reset(args.ref_D1.data());
+    params_.ref_D2.reset(args.ref_D2.data());
+    params_.output_op_0 = args.epilogue0;
+    params_.output_op_1 = args.epilogue1;
+    params_.output_op_2 = args.epilogue2;
+    params_.semaphore = reinterpret_cast<int *>(workspace);
+
+    return Status::kSuccess;
+  }
+
+  /// Runs the kernel using initialized state.
+  Status run(cudaStream_t stream = nullptr) {
+    ThreadblockSwizzle threadblock_swizzle;
+
+    dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape);
+    dim3 block(DualGemmKernel::kThreadCount, 1, 1);
+
+    cudaError_t result;
+
+    int smem_size =
+        int(sizeof(typename DualGemmKernel::SharedStorage));  // NOLINT
+    if (smem_size >= (48 << 10)) {
+      result = cudaFuncSetAttribute(Kernel<DualGemmKernel>,
+                                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                                    smem_size);
+
+      if (result != cudaSuccess) {
+        printf(
+            "cudaFuncSetAttribute(Kernel<DualGemmKernel>, "
+            "cudaFuncAttributeMaxDynamicSharedMemorySize, %d) returned an "
+            "error\n",
+            smem_size);
+        return Status::kErrorInternal;
+      }
+    }
+
+    cutlass::Kernel<DualGemmKernel>
+        <<<grid, block, smem_size, stream>>>(params_);
+
+    result = cudaGetLastError();
+
+    return result == cudaSuccess ? Status::kSuccess : Status::kErrorInternal;
+  }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(cudaStream_t stream = nullptr) { return run(stream); }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(Arguments const &args,
+                    void *workspace = nullptr,
+                    cudaStream_t stream = nullptr) {
+    Status status = initialize(args, workspace, stream);
+
+    if (status == Status::kSuccess) {
+      status = run(stream);
+    }
+
+    return status;
+  }
+};
+
+}  // namespace device
+}  // namespace gemm
+}  // namespace cutlass

csrc/gpu/cutlass_kernels/fp8_gemm_fused/dual_gemm/dual_gemm_common.h

@@ -0,0 +1,43 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Defines common types used for all DualGemm operators.
+*/
+#pragma once
+
+namespace cutlass {
+namespace gemm {
+
+enum class DualGemmMode { kGemm, kBatched, kInvalid };
+
+}  // namespace gemm
+}  // namespace cutlass