chore: import upstream snapshot with attribution
Book-CI / test (macos-latest) (push) Has been cancelled
Book-CI / test (ubuntu-latest) (push) Has been cancelled
Book-CI / test (windows-latest) (push) Has been cancelled
Release Fake Tag / publish (push) Has been cancelled
Deploy / deploy (macos-latest) (push) Has been cancelled
Deploy / deploy (ubuntu-latest) (push) Has been cancelled
Deploy / deploy (windows-latest) (push) Has been cancelled
Release to PyPI / Build & publish sglang-kt (push) Has been cancelled
Release to PyPI / Build kt-kernel (Python 3.11) (push) Has been cancelled
Release to PyPI / Build kt-kernel (Python 3.12) (push) Has been cancelled
Release to PyPI / Publish kt-kernel to PyPI (push) Has been cancelled
Book-CI / test (macos-latest) (push) Has been cancelled
Book-CI / test (ubuntu-latest) (push) Has been cancelled
Book-CI / test (windows-latest) (push) Has been cancelled
Release Fake Tag / publish (push) Has been cancelled
Deploy / deploy (macos-latest) (push) Has been cancelled
Deploy / deploy (ubuntu-latest) (push) Has been cancelled
Deploy / deploy (windows-latest) (push) Has been cancelled
Release to PyPI / Build & publish sglang-kt (push) Has been cancelled
Release to PyPI / Build kt-kernel (Python 3.11) (push) Has been cancelled
Release to PyPI / Build kt-kernel (Python 3.12) (push) Has been cancelled
Release to PyPI / Publish kt-kernel to PyPI (push) Has been cancelled
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,84 @@
|
||||
|
||||
# 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, 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 [INSERT CONTACT METHOD]. 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.0, available at [https://www.contributor-covenant.org/version/2/0/code_of_conduct.html][v2.0].
|
||||
|
||||
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.0]: https://www.contributor-covenant.org/version/2/0/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
|
||||
@@ -0,0 +1,139 @@
|
||||
## Before Commit!
|
||||
|
||||
Your commit message must follow Conventional Commits (https://www.conventionalcommits.org/) and your code should be formatted. The Git hooks will do most of the work automatically:
|
||||
|
||||
### Tool Requirements
|
||||
|
||||
You need a recent `clang-format` (>= 18). In a conda environment you can install:
|
||||
|
||||
```shell
|
||||
conda install -c conda-forge clang-format=18
|
||||
```
|
||||
|
||||
If you previously configured with an older version, remove the build directory and reconfigure:
|
||||
|
||||
```shell
|
||||
rm -rf kt-kernel/build
|
||||
```
|
||||
|
||||
Install `black` for Python formatting:
|
||||
|
||||
```shell
|
||||
conda install black
|
||||
```
|
||||
|
||||
### Install hook:
|
||||
```shell
|
||||
bash kt-kernel/scripts/install-git-hooks.sh
|
||||
#or just cmake the kt-kernel
|
||||
cmake -S kt-kernel -B kt-kernel/build
|
||||
```
|
||||
|
||||
There are manual commands if you need format.
|
||||
|
||||
```shell
|
||||
cmake -S kt-kernel -B kt-kernel/build
|
||||
cmake --build kt-kernel/build --target format
|
||||
```
|
||||
|
||||
## Developer Note
|
||||
|
||||
Formatting and commit message rules are enforced by Git hooks. After installing `clang-format` and `black`, just commit normally—the hooks will run formatting for you.
|
||||
|
||||
> [!NOTE]
|
||||
> If formatting modifies files, the commit is aborted after staging those changes. Review them and run `git commit` again. Repeat until no further formatting changes appear.
|
||||
|
||||
---
|
||||
|
||||
### Conventional Commit Regex (Reference)
|
||||
|
||||
The commit-msg hook enforces this pattern:
|
||||
|
||||
```text
|
||||
regex='^\[(feat|fix|docs|style|refactor|perf|test|build|ci|chore|revert|wip)\](\([^\)]+\))?(!)?: .+'
|
||||
```
|
||||
|
||||
Meaning (English):
|
||||
* `[type]` required — one of feat|fix|docs|style|refactor|perf|test|build|ci|chore|revert|wip
|
||||
* Optional scope: `(scope)` — any chars except `)`
|
||||
* Optional breaking change marker: `!` right after type or scope
|
||||
* Separator: `: ` (colon + space)
|
||||
* Subject: free text (at least one character)
|
||||
|
||||
Examples:
|
||||
```text
|
||||
[feat]: add adaptive batching
|
||||
[fix(parser)]: handle empty token list
|
||||
[docs]!: update API section for breaking rename
|
||||
```
|
||||
|
||||
You can bypass locally (not recommended) with:
|
||||
```shell
|
||||
git commit --no-verify
|
||||
```
|
||||
## 提交前提醒
|
||||
|
||||
提交信息必须满足 Conventional Commits 规范 (https://www.conventionalcommits.org/)，代码需要符合格式要求。Git 钩子已经集成了大部分工作：
|
||||
### 软件要求
|
||||
|
||||
需要较新的 `clang-format` (>= 18)，在 conda 环境中安装：
|
||||
|
||||
```shell
|
||||
conda install -c conda-forge clang-format=18
|
||||
```
|
||||
|
||||
如果之前用老版本配置过，请删除构建目录重新配置：
|
||||
|
||||
```shell
|
||||
rm -rf kt-kernel/build
|
||||
```
|
||||
|
||||
安装 `black` 以进行 Python 文件格式化：
|
||||
|
||||
```shell
|
||||
conda install black
|
||||
```
|
||||
### 安装钩子
|
||||
```shell
|
||||
bash kt-kernel/scripts/install-git-hooks.sh
|
||||
#or just cmake the kt-kernel
|
||||
cmake -S kt-kernel -B kt-kernel/build
|
||||
```
|
||||
如果你需要手动格式化：
|
||||
```shell
|
||||
cmake -S kt-kernel -B kt-kernel/build
|
||||
cmake --build kt-kernel/build --target format
|
||||
```
|
||||
|
||||
## 开发者说明
|
||||
|
||||
本仓库通过 Git hooks 自动执行代码格式化与提交信息规范检查。只需安装好 `clang-format` 与 `black` 后正常执行提交即可，钩子会自动格式化。
|
||||
|
||||
> [!NOTE]
|
||||
> 如果格式化修改了文件，钩子会终止提交并已暂存这些改动。请查看修改后再次执行 `git commit`，重复直到没有新的格式化变更。
|
||||
|
||||
### 提交信息正则（参考）
|
||||
|
||||
钩子使用如下正则检查提交信息：
|
||||
```text
|
||||
regex='^\[(feat|fix|docs|style|refactor|perf|test|build|ci|chore|revert|wip)\](\([^\)]+\))?(!)?: .+'
|
||||
```
|
||||
含义：
|
||||
* `[type]` 必填：feat|fix|docs|style|refactor|perf|test|build|ci|chore|revert|wip
|
||||
* 作用域可选：`(scope)`，不能包含右括号
|
||||
* 可选的破坏性标记：`!`
|
||||
* 分隔符：冒号+空格 `: `
|
||||
* 描述：至少一个字符
|
||||
|
||||
示例：
|
||||
```text
|
||||
[feat]: 增加自适应 batch 功能
|
||||
[fix(tokenizer)]: 修复空 token 列表处理
|
||||
[docs]!: 更新接口文档（存在破坏性修改）
|
||||
```
|
||||
|
||||
跳过钩子（不推荐，仅紧急时）：
|
||||
```shell
|
||||
git commit --no-verify
|
||||
```
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
name: "\U0001F41B Bug / Help"
|
||||
description: Create a report to help us improve the ktransformers project
|
||||
labels: ["pending"]
|
||||
body:
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: |
|
||||
Issues included in **[FAQs](https://github.com/kvcache-ai/ktransformers/issues/1608)** or those with **insufficient** information may be closed without a response.
|
||||
已经包含在 **[常见问题](https://github.com/kvcache-ai/ktransformers/issues/1608)** 内或提供信息**不完整**的 issues 可能不会被回复。
|
||||
|
||||
- type: checkboxes
|
||||
id: reminder
|
||||
attributes:
|
||||
label: Reminder
|
||||
description: |
|
||||
Please ensure you have read the above rules carefully and searched the existing issues (including FAQs).
|
||||
请确保您已经认真阅读了上述规则并且搜索过现有的 issues（包括常见问题）。
|
||||
|
||||
options:
|
||||
- label: I have read the above rules and searched the existing issues.
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: system-info
|
||||
validations:
|
||||
required: true
|
||||
attributes:
|
||||
label: System Info
|
||||
description: |
|
||||
Please share your system info with us. You can run the command **lscpu**, ** nvidia-smi ** etc. and copy-paste its output below.
|
||||
请提供您的系统信息。您可以在命令行运行 **lscpu**, **nvidia-smi** 等命令，并将其输出复制到该文本框中。
|
||||
|
||||
placeholder: ktransformers version,sglang version, platform, python version, cpu info, GPU/NPU info ...
|
||||
|
||||
- type: textarea
|
||||
id: reproduction
|
||||
validations:
|
||||
required: true
|
||||
attributes:
|
||||
label: Reproduction
|
||||
description: |
|
||||
Please provide entry arguments, error messages and stack traces that reproduces the problem.
|
||||
请提供入口参数，错误日志以及异常堆栈以便于我们复现问题。
|
||||
|
||||
value: |
|
||||
```text
|
||||
Put your message here.
|
||||
```
|
||||
|
||||
- type: textarea
|
||||
id: others
|
||||
validations:
|
||||
required: false
|
||||
attributes:
|
||||
label: Others
|
||||
@@ -0,0 +1,41 @@
|
||||
name: "\U0001F680 Feature request"
|
||||
description: Submit a request for a new feature
|
||||
labels: ["enhancement", "pending"]
|
||||
body:
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: |
|
||||
Please do not create issues that are not related to new features under this category.
|
||||
请勿在此分类下创建和新特性无关的 issues。
|
||||
|
||||
- type: checkboxes
|
||||
id: reminder
|
||||
attributes:
|
||||
label: Reminder
|
||||
description: |
|
||||
Please ensure you have read the above rules carefully and searched the existing issues.
|
||||
请确保您已经认真阅读了上述规则并且搜索过现有的 issues。
|
||||
|
||||
options:
|
||||
- label: I have read the above rules and searched the existing issues.
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: description
|
||||
validations:
|
||||
required: true
|
||||
attributes:
|
||||
label: Description
|
||||
description: |
|
||||
A clear and concise description of the feature proposal.
|
||||
请详细描述您希望加入的新功能特性。
|
||||
|
||||
- type: textarea
|
||||
id: contribution
|
||||
validations:
|
||||
required: false
|
||||
attributes:
|
||||
label: Pull Request
|
||||
description: |
|
||||
Have you already created the relevant PR and submitted the code?
|
||||
您是否已经创建了相关 PR 并提交了代码？
|
||||
@@ -0,0 +1,5 @@
|
||||
blank_issues_enabled: false
|
||||
contact_links:
|
||||
- name: 📚 FAQs | 常见问题
|
||||
url: https://github.com/kvcache-ai/ktransformers/issues/1608
|
||||
about: Reading in advance is recommended | 建议提前阅读
|
||||
@@ -0,0 +1,8 @@
|
||||
# What does this PR do?
|
||||
|
||||
Fixes # (issue)
|
||||
|
||||
## Before submitting
|
||||
|
||||
- [ ] Did you read the [contributor guideline](https://github.com/kvcache-ai/ktransformers/blob/main/.github/CONTRIBUTING.md)?
|
||||
- [ ] Did you write any new necessary tests?
|
||||
@@ -0,0 +1,7 @@
|
||||
# Reporting Security Issues
|
||||
|
||||
To report a security issue, please use the GitHub Security Advisory ["Report a Vulnerability"](https://github.com/kvcache-ai/ktransformers/security/advisories/new) tab.
|
||||
|
||||
We will send a response indicating the next steps in handling your report. After the initial reply to your report, the security team will keep you informed of the progress towards a fix and full announcement, and may ask for additional information or guidance.
|
||||
|
||||
Report security bugs in third-party modules to the person or team maintaining the module.
|
||||
@@ -0,0 +1,32 @@
|
||||
name: Book-CI
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
# - server_support
|
||||
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
# - server_support
|
||||
jobs:
|
||||
test:
|
||||
name: test
|
||||
runs-on: ${{ matrix.os }}
|
||||
strategy:
|
||||
matrix:
|
||||
os: [ubuntu-latest, macos-latest, windows-latest]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Install Rust
|
||||
run: |
|
||||
rustup set profile minimal
|
||||
rustup toolchain install stable
|
||||
rustup default stable
|
||||
- name: Setup mdBook
|
||||
uses: peaceiris/actions-mdbook@v2
|
||||
with:
|
||||
mdbook-version: "latest"
|
||||
# - name: Run tests
|
||||
# run: mdbook test
|
||||
@@ -0,0 +1,49 @@
|
||||
name: Deploy
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
# - server_support
|
||||
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
# - server_support
|
||||
|
||||
defaults:
|
||||
run:
|
||||
shell: bash
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
|
||||
jobs:
|
||||
deploy:
|
||||
runs-on: ${{ matrix.os }}
|
||||
strategy:
|
||||
matrix:
|
||||
os: [ubuntu-latest, macos-latest, windows-latest]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Install Rust
|
||||
run: |
|
||||
rustup set profile minimal
|
||||
rustup toolchain install stable
|
||||
rustup default stable
|
||||
- name: Setup mdBook
|
||||
uses: peaceiris/actions-mdbook@v2
|
||||
with:
|
||||
mdbook-version: "latest"
|
||||
- run: mdbook build
|
||||
# - name: Copy Assets
|
||||
# run: |
|
||||
# chmod +x ci/copy-assets.sh
|
||||
# ci/copy-assets.sh ${{ matrix.os }}
|
||||
- name: Deploy
|
||||
uses: peaceiris/actions-gh-pages@v3
|
||||
# or || github.ref == 'refs/heads/server_support'
|
||||
if: ${{ github.ref == 'refs/heads/main' }}
|
||||
with:
|
||||
github_token: ${{ secrets.GITHUB_TOKEN }}
|
||||
publish_dir: ./book
|
||||
@@ -0,0 +1,142 @@
|
||||
name: DockerHub CI
|
||||
|
||||
on:
|
||||
release:
|
||||
types: [published]
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
push_to_dockerhub:
|
||||
description: 'Push image to DockerHub? (true/false)'
|
||||
required: true
|
||||
default: 'false'
|
||||
type: boolean
|
||||
cuda_version:
|
||||
description: 'CUDA version (e.g., 12.8.1)'
|
||||
required: false
|
||||
default: '12.8.1'
|
||||
type: string
|
||||
push_simplified_tag:
|
||||
description: 'Also push simplified tag? (true/false)'
|
||||
required: false
|
||||
default: 'true'
|
||||
type: boolean
|
||||
ubuntu_mirror:
|
||||
description: 'Use Tsinghua Ubuntu mirror? (0/1)'
|
||||
required: false
|
||||
default: '0'
|
||||
type: string
|
||||
|
||||
# push:
|
||||
# branches:
|
||||
# - main
|
||||
env:
|
||||
DOCKERHUB_REPO: ${{ secrets.DOCKERHUB_USERNAME }}/ktransformers
|
||||
jobs:
|
||||
test:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Run tests
|
||||
run: |
|
||||
if [ -f docker-compose.test.yml ]; then
|
||||
docker-compose --file docker-compose.test.yml build
|
||||
docker-compose --file docker-compose.test.yml run sut
|
||||
else
|
||||
docker build . --file docker/Dockerfile
|
||||
fi
|
||||
|
||||
build-and-push:
|
||||
needs: test
|
||||
name: Build and Push Multi-Variant Docker Image
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Move Docker data directory
|
||||
run: |
|
||||
sudo systemctl stop docker
|
||||
sudo mkdir -p /mnt/docker
|
||||
sudo rsync -avz /var/lib/docker/ /mnt/docker
|
||||
sudo rm -rf /var/lib/docker
|
||||
sudo ln -s /mnt/docker /var/lib/docker
|
||||
sudo systemctl start docker
|
||||
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v3
|
||||
|
||||
- name: Login to Docker Hub
|
||||
uses: docker/login-action@v3
|
||||
with:
|
||||
username: ${{ secrets.DOCKERHUB_USERNAME }}
|
||||
password: ${{ secrets.DOCKERHUB_TOKEN }}
|
||||
|
||||
- name: Determine build parameters
|
||||
id: params
|
||||
run: |
|
||||
# Determine if we should push
|
||||
if [ "${{ github.event_name }}" = "release" ]; then
|
||||
echo "should_push=true" >> $GITHUB_OUTPUT
|
||||
echo "push_simplified=true" >> $GITHUB_OUTPUT
|
||||
elif [ "${{ github.event_name }}" = "workflow_dispatch" ]; then
|
||||
echo "should_push=${{ inputs.push_to_dockerhub }}" >> $GITHUB_OUTPUT
|
||||
echo "push_simplified=${{ inputs.push_simplified_tag }}" >> $GITHUB_OUTPUT
|
||||
else
|
||||
echo "should_push=false" >> $GITHUB_OUTPUT
|
||||
echo "push_simplified=false" >> $GITHUB_OUTPUT
|
||||
fi
|
||||
|
||||
# Determine CUDA version
|
||||
if [ "${{ github.event_name }}" = "workflow_dispatch" ] && [ -n "${{ inputs.cuda_version }}" ]; then
|
||||
echo "cuda_version=${{ inputs.cuda_version }}" >> $GITHUB_OUTPUT
|
||||
else
|
||||
echo "cuda_version=12.8.1" >> $GITHUB_OUTPUT
|
||||
fi
|
||||
|
||||
# Determine Ubuntu mirror setting
|
||||
if [ "${{ github.event_name }}" = "workflow_dispatch" ] && [ -n "${{ inputs.ubuntu_mirror }}" ]; then
|
||||
echo "ubuntu_mirror=${{ inputs.ubuntu_mirror }}" >> $GITHUB_OUTPUT
|
||||
else
|
||||
echo "ubuntu_mirror=0" >> $GITHUB_OUTPUT
|
||||
fi
|
||||
|
||||
- name: Build and push Docker image
|
||||
run: |
|
||||
cd docker
|
||||
|
||||
# Build command arguments
|
||||
BUILD_ARGS=(
|
||||
--cuda-version "${{ steps.params.outputs.cuda_version }}"
|
||||
--ubuntu-mirror "${{ steps.params.outputs.ubuntu_mirror }}"
|
||||
--repository "${{ env.DOCKERHUB_REPO }}"
|
||||
)
|
||||
|
||||
# Add simplified tag option if enabled
|
||||
if [ "${{ steps.params.outputs.push_simplified }}" = "true" ]; then
|
||||
BUILD_ARGS+=(--also-push-simplified)
|
||||
fi
|
||||
|
||||
# Add HTTP proxy if available
|
||||
if [ -n "${{ secrets.HTTP_PROXY }}" ]; then
|
||||
BUILD_ARGS+=(--http-proxy "${{ secrets.HTTP_PROXY }}")
|
||||
fi
|
||||
|
||||
# Add HTTPS proxy if available
|
||||
if [ -n "${{ secrets.HTTPS_PROXY }}" ]; then
|
||||
BUILD_ARGS+=(--https-proxy "${{ secrets.HTTPS_PROXY }}")
|
||||
fi
|
||||
|
||||
# Dry run if not pushing
|
||||
if [ "${{ steps.params.outputs.should_push }}" != "true" ]; then
|
||||
BUILD_ARGS+=(--dry-run)
|
||||
fi
|
||||
|
||||
# Execute build script
|
||||
./push-to-dockerhub.sh "${BUILD_ARGS[@]}"
|
||||
|
||||
- name: Display image information
|
||||
if: steps.params.outputs.should_push == 'true'
|
||||
run: |
|
||||
echo "::notice title=Docker Image::Image pushed successfully to ${{ env.DOCKERHUB_REPO }}"
|
||||
echo "Pull command: docker pull ${{ env.DOCKERHUB_REPO }}:v\$(VERSION)-cu\$(CUDA_SHORT)"
|
||||
@@ -0,0 +1,104 @@
|
||||
name: PR KT-Kernel Test
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- develop
|
||||
types: [synchronize, labeled]
|
||||
workflow_dispatch:
|
||||
|
||||
concurrency:
|
||||
group: pr-kt-kernel-test-${{ github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
# =============================================== check changes ====================================================
|
||||
check-changes:
|
||||
runs-on: ubuntu-latest
|
||||
outputs:
|
||||
kt_kernel: ${{ steps.filter.outputs.kt_kernel }}
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Fail if the PR does not have the 'run-ci' label
|
||||
if: github.event_name == 'pull_request' && !contains(github.event.pull_request.labels.*.name, 'run-ci')
|
||||
run: |
|
||||
echo "This pull request does not have the 'run-ci' label. Failing the workflow."
|
||||
exit 1
|
||||
|
||||
- name: Fail if the PR is a draft
|
||||
if: github.event_name == 'pull_request' && github.event.pull_request.draft == true
|
||||
run: |
|
||||
echo "This pull request is a draft. Failing the workflow."
|
||||
exit 1
|
||||
|
||||
- name: Detect file changes
|
||||
id: filter
|
||||
uses: dorny/paths-filter@v3
|
||||
with:
|
||||
filters: |
|
||||
kt_kernel:
|
||||
- "kt-kernel/**"
|
||||
- ".github/workflows/kt-kernel-tests.yml"
|
||||
|
||||
# =============================================== KT-Kernel tests ====================================================
|
||||
per-commit-kt-kernel-cpu:
|
||||
needs: [check-changes]
|
||||
if: always() && !failure() && !cancelled() &&
|
||||
(needs.check-changes.outputs.kt_kernel == 'true' || github.event_name == 'workflow_dispatch')
|
||||
runs-on: kt-cpu
|
||||
continue-on-error: false
|
||||
steps:
|
||||
- name: Cleanup
|
||||
run: |
|
||||
sudo rm -rf $GITHUB_WORKSPACE/* || true
|
||||
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
submodules: recursive
|
||||
|
||||
- name: Install KT-Kernel
|
||||
run: |
|
||||
cd kt-kernel
|
||||
bash install.sh build
|
||||
|
||||
- name: Run KT-Kernel CPU tests
|
||||
timeout-minutes: 60
|
||||
run: |
|
||||
cd kt-kernel/test
|
||||
python3 run_suite.py --hw cpu --suite default
|
||||
|
||||
# =============================================== finish ====================================================
|
||||
pr-test-kt-kernel-finish:
|
||||
needs: [check-changes, per-commit-kt-kernel-cpu]
|
||||
if: always()
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Check all dependent job statuses
|
||||
run: |
|
||||
# Convert the 'needs' context to a JSON string
|
||||
json_needs='${{ toJson(needs) }}'
|
||||
|
||||
# Get a list of all job names from the JSON keys
|
||||
job_names=$(echo "$json_needs" | jq -r 'keys_unsorted[]')
|
||||
|
||||
for job in $job_names; do
|
||||
# For each job, extract its result
|
||||
result=$(echo "$json_needs" | jq -r --arg j "$job" '.[$j].result')
|
||||
|
||||
# Print the job name and its result
|
||||
echo "$job: $result"
|
||||
|
||||
# Check for failure or cancellation and exit if found
|
||||
if [[ "$result" == "failure" || "$result" == "cancelled" ]]; then
|
||||
echo "The above jobs failed."
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
|
||||
# If the loop completes, all jobs were successful
|
||||
echo "All jobs completed successfully"
|
||||
exit 0
|
||||
@@ -0,0 +1,41 @@
|
||||
name: Release Fake Tag
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
paths:
|
||||
- "version.py"
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
|
||||
jobs:
|
||||
publish:
|
||||
if: github.repository == 'kvcache-ai/ktransformers'
|
||||
runs-on: ubuntu-latest
|
||||
environment: 'prod'
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
token: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
- name: Get version
|
||||
id: get_version
|
||||
run: |
|
||||
version=$(cat version.py | grep '__version__' | cut -d'"' -f2)
|
||||
echo "TAG=v$version" >> $GITHUB_OUTPUT
|
||||
|
||||
- name: Create and push tag
|
||||
run: |
|
||||
TAG=${{ steps.get_version.outputs.TAG }}
|
||||
git config user.name "ktransformers-bot"
|
||||
git config user.email "ktransformers-bot@users.noreply.github.com"
|
||||
if git ls-remote --tags --exit-code origin "refs/tags/${TAG}" > /dev/null 2>&1; then
|
||||
echo "Tag ${TAG} already exists on origin, skipping."
|
||||
exit 0
|
||||
fi
|
||||
git tag "${TAG}"
|
||||
git push origin "${TAG}"
|
||||
@@ -0,0 +1,283 @@
|
||||
name: Release to PyPI
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
paths:
|
||||
- "version.py"
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
test_pypi:
|
||||
description: 'Publish to TestPyPI instead of PyPI (for testing)'
|
||||
required: false
|
||||
default: 'false'
|
||||
type: choice
|
||||
options:
|
||||
- 'true'
|
||||
- 'false'
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
# ── sglang-kt (must be on PyPI before users can pip install kt-kernel) ──
|
||||
build-and-publish-sglang-kt:
|
||||
name: Build & publish sglang-kt
|
||||
runs-on: [self-hosted, linux, x64]
|
||||
if: github.repository == 'kvcache-ai/ktransformers' && github.ref == 'refs/heads/main'
|
||||
environment: prod
|
||||
permissions:
|
||||
id-token: write
|
||||
contents: read
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
submodules: recursive
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: '3.12'
|
||||
|
||||
- name: Install build tools
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install build wheel setuptools twine
|
||||
|
||||
- name: Build sglang-kt wheel
|
||||
working-directory: third_party/sglang/python
|
||||
run: |
|
||||
KT_VERSION=$(python3 -c "exec(open('${{ github.workspace }}/version.py').read()); print(__version__)")
|
||||
export SGLANG_KT_VERSION="$KT_VERSION"
|
||||
echo "Building sglang-kt v${KT_VERSION} wheel..."
|
||||
python -m build --wheel -v
|
||||
ls dist/ | grep -q "sglang_kt" || (echo "ERROR: Wheel name does not contain sglang_kt" && exit 1)
|
||||
|
||||
- name: Publish sglang-kt to PyPI
|
||||
if: github.event.inputs.test_pypi != 'true'
|
||||
env:
|
||||
TWINE_USERNAME: __token__
|
||||
TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
|
||||
run: |
|
||||
python -m twine upload --skip-existing --verbose third_party/sglang/python/dist/*.whl
|
||||
|
||||
- name: Publish sglang-kt to TestPyPI (if requested)
|
||||
if: github.event.inputs.test_pypi == 'true'
|
||||
env:
|
||||
TWINE_USERNAME: __token__
|
||||
TWINE_PASSWORD: ${{ secrets.TEST_PYPI_API_TOKEN }}
|
||||
run: |
|
||||
python -m twine upload --repository testpypi --skip-existing --verbose third_party/sglang/python/dist/*.whl
|
||||
|
||||
# ── kt-kernel ──
|
||||
build-kt-kernel:
|
||||
name: Build kt-kernel (Python ${{ matrix.python-version }})
|
||||
runs-on: [self-hosted, linux, x64, gpu]
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
python-version: ['3.11', '3.12']
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
submodules: recursive
|
||||
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
|
||||
- name: Verify CUDA availability
|
||||
run: |
|
||||
nvidia-smi || (echo "ERROR: GPU not available" && exit 1)
|
||||
nvcc --version || (echo "ERROR: CUDA toolkit not found" && exit 1)
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
# System packages (cmake/libhwloc-dev/pkg-config/libnuma-dev) are expected to be
|
||||
# preinstalled on self-hosted runners. Skip apt-get to avoid sudo dependency.
|
||||
for pkg in cmake pkg-config; do command -v $pkg >/dev/null || { echo "ERROR: $pkg missing on runner"; exit 1; }; done
|
||||
python -m pip install --upgrade pip
|
||||
pip install build wheel setuptools
|
||||
pip install torch --index-url https://download.pytorch.org/whl/cu128
|
||||
|
||||
- name: Build kt-kernel wheel
|
||||
working-directory: kt-kernel
|
||||
env:
|
||||
CPUINFER_BUILD_ALL_VARIANTS: '1'
|
||||
CPUINFER_ENABLE_CPPTRACE: '0'
|
||||
CPUINFER_USE_CUDA: '1'
|
||||
CPUINFER_CUDA_ARCHS: '80;86;89;90;120'
|
||||
CPUINFER_CUDA_STATIC_RUNTIME: '1'
|
||||
CPUINFER_BUILD_TYPE: 'Release'
|
||||
CPUINFER_PARALLEL: '4'
|
||||
CPUINFER_FORCE_REBUILD: '1'
|
||||
CUDA_HOME: '/usr/local/cuda-12.8'
|
||||
run: |
|
||||
echo "Building kt-kernel with:"
|
||||
echo " - CUDA support (SM 80, 86, 89, 90, 120)"
|
||||
echo " - CPU multi-variant (AMX, AVX512, AVX2)"
|
||||
python -m build --wheel -v
|
||||
|
||||
- name: Verify wheel
|
||||
working-directory: kt-kernel
|
||||
run: |
|
||||
echo "Generated wheel:"
|
||||
ls -lh dist/
|
||||
|
||||
# Install and test
|
||||
pip install dist/*.whl
|
||||
python -c "import kt_kernel; print(f'✓ Version: {kt_kernel.__version__}')"
|
||||
python -c "import kt_kernel; print(f'✓ CPU variant: {kt_kernel.__cpu_variant__}')"
|
||||
|
||||
# Verify CUDA support
|
||||
python -c "
|
||||
from kt_kernel import kt_kernel_ext
|
||||
cpu_infer = kt_kernel_ext.CPUInfer(4)
|
||||
methods = dir(cpu_infer)
|
||||
has_cuda = 'submit_with_cuda_stream' in methods
|
||||
print(f'✓ CUDA support: {has_cuda}')
|
||||
"
|
||||
|
||||
# Verify CPU multi-variant support
|
||||
echo "Checking CPU variants in wheel..."
|
||||
python -m zipfile -l dist/*.whl | grep "_kt_kernel_ext_" || echo "Warning: No variant .so files found"
|
||||
python -m zipfile -l dist/*.whl | grep "_kt_kernel_ext_amx.cpython" && echo "✓ AMX variant found" || echo "Note: AMX variant missing"
|
||||
python -m zipfile -l dist/*.whl | grep "_kt_kernel_ext_avx512" && echo "✓ AVX512 variants found" || echo "Note: AVX512 variants missing"
|
||||
python -m zipfile -l dist/*.whl | grep "_kt_kernel_ext_avx2.cpython" && echo "✓ AVX2 variant found" || echo "Note: AVX2 variant missing"
|
||||
|
||||
# Verify static linking (should NOT depend on libcudart.so).
|
||||
# Use $RUNNER_TEMP (honors TMPDIR redirect to /mnt) — /tmp is the
|
||||
# system disk on self-hosted runners and can be tight.
|
||||
CHECK_DIR="${RUNNER_TEMP:-/tmp}/check"
|
||||
rm -rf "$CHECK_DIR"
|
||||
unzip -q dist/*.whl -d "$CHECK_DIR"
|
||||
if ldd "$CHECK_DIR"/kt_kernel/*.so 2>/dev/null | grep -q "libcudart.so"; then
|
||||
echo "ERROR: Dynamic cudart found, should be statically linked"
|
||||
exit 1
|
||||
else
|
||||
echo "✓ CUDA runtime statically linked"
|
||||
fi
|
||||
|
||||
- name: Repair wheel for manylinux
|
||||
working-directory: kt-kernel
|
||||
run: |
|
||||
pip install auditwheel patchelf
|
||||
mkdir -p wheelhouse
|
||||
for wheel in dist/*.whl; do
|
||||
auditwheel repair "$wheel" --plat manylinux_2_35_x86_64 --exclude libcuda.so.1 -w wheelhouse/
|
||||
done
|
||||
rm -f dist/*.whl && cp wheelhouse/*.whl dist/
|
||||
|
||||
- name: Upload artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: kt-kernel-wheels-py${{ matrix.python-version }}
|
||||
path: kt-kernel/dist/*.whl
|
||||
retention-days: 7
|
||||
|
||||
publish-pypi:
|
||||
name: Publish kt-kernel to PyPI
|
||||
needs: [build-and-publish-sglang-kt, build-kt-kernel]
|
||||
runs-on: [self-hosted, linux, x64]
|
||||
if: github.repository == 'kvcache-ai/ktransformers' && github.ref == 'refs/heads/main'
|
||||
environment: prod
|
||||
permissions:
|
||||
id-token: write # For trusted publishing (OIDC)
|
||||
contents: read
|
||||
|
||||
steps:
|
||||
- name: Download all wheel artifacts
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
path: artifacts/
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: '3.12'
|
||||
|
||||
- name: Organize wheels into dist/
|
||||
run: |
|
||||
mkdir -p dist/
|
||||
find artifacts/ -name "*.whl" -exec cp {} dist/ \;
|
||||
echo "Wheels to publish:"
|
||||
ls -lh dist/
|
||||
|
||||
- name: Get version from wheel
|
||||
id: get_version
|
||||
run: |
|
||||
# Extract version from first wheel filename
|
||||
wheel_name=$(ls dist/*.whl | head -1 | xargs basename)
|
||||
# Extract version (format: kt_kernel-X.Y.Z-...)
|
||||
version=$(echo "$wheel_name" | sed 's/kt_kernel-\([0-9.]*\)-.*/\1/')
|
||||
echo "VERSION=$version" >> $GITHUB_OUTPUT
|
||||
echo "Publishing version: $version"
|
||||
|
||||
- name: Install twine
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install twine
|
||||
|
||||
- name: Publish to TestPyPI (if requested)
|
||||
if: github.event.inputs.test_pypi == 'true'
|
||||
env:
|
||||
TWINE_USERNAME: __token__
|
||||
TWINE_PASSWORD: ${{ secrets.TEST_PYPI_API_TOKEN }}
|
||||
run: |
|
||||
python -m twine upload \
|
||||
--repository testpypi \
|
||||
--skip-existing \
|
||||
--verbose \
|
||||
dist/*.whl
|
||||
|
||||
- name: Publish to PyPI
|
||||
if: github.event.inputs.test_pypi != 'true'
|
||||
env:
|
||||
TWINE_USERNAME: __token__
|
||||
TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
|
||||
run: |
|
||||
python -m twine upload \
|
||||
--skip-existing \
|
||||
--verbose \
|
||||
dist/*.whl
|
||||
|
||||
- name: Create release summary
|
||||
run: |
|
||||
echo "## 🎉 kt-kernel v${{ steps.get_version.outputs.VERSION }} Published to PyPI" >> $GITHUB_STEP_SUMMARY
|
||||
echo "" >> $GITHUB_STEP_SUMMARY
|
||||
echo "### Installation" >> $GITHUB_STEP_SUMMARY
|
||||
echo '```bash' >> $GITHUB_STEP_SUMMARY
|
||||
echo "pip install kt-kernel==${{ steps.get_version.outputs.VERSION }}" >> $GITHUB_STEP_SUMMARY
|
||||
echo '```' >> $GITHUB_STEP_SUMMARY
|
||||
echo "" >> $GITHUB_STEP_SUMMARY
|
||||
echo "### Published Wheels" >> $GITHUB_STEP_SUMMARY
|
||||
echo "Total: $(ls -1 dist/*.whl | wc -l) wheels (Python 3.10, 3.11, 3.12)" >> $GITHUB_STEP_SUMMARY
|
||||
echo "" >> $GITHUB_STEP_SUMMARY
|
||||
echo "### Features" >> $GITHUB_STEP_SUMMARY
|
||||
echo "" >> $GITHUB_STEP_SUMMARY
|
||||
echo "**CPU Multi-Variant Support:**" >> $GITHUB_STEP_SUMMARY
|
||||
echo "- ✅ AMX (Intel Sapphire Rapids+, 2023)" >> $GITHUB_STEP_SUMMARY
|
||||
echo "- ✅ AVX512 Base/VNNI/VBMI/BF16 (Intel Skylake-X/Ice Lake/Cascade Lake, 2017+)" >> $GITHUB_STEP_SUMMARY
|
||||
echo "- ✅ AVX2 (Maximum compatibility, 2013+)" >> $GITHUB_STEP_SUMMARY
|
||||
echo "- 🔧 Runtime CPU detection: Automatically selects optimal variant" >> $GITHUB_STEP_SUMMARY
|
||||
echo "" >> $GITHUB_STEP_SUMMARY
|
||||
echo "**CUDA Support:**" >> $GITHUB_STEP_SUMMARY
|
||||
echo "- ✅ SM 80 (Ampere: A100, RTX 3000 series)" >> $GITHUB_STEP_SUMMARY
|
||||
echo "- ✅ SM 86 (Ampere: RTX 3060-3090)" >> $GITHUB_STEP_SUMMARY
|
||||
echo "- ✅ SM 89 (Ada Lovelace: RTX 4000 series)" >> $GITHUB_STEP_SUMMARY
|
||||
echo "- ✅ SM 90 (Hopper: H100)" >> $GITHUB_STEP_SUMMARY
|
||||
echo "- 🔧 Static CUDA runtime: Compatible with CUDA 11.8+ and 12.x drivers" >> $GITHUB_STEP_SUMMARY
|
||||
echo "- 🔧 Works on CPU-only systems (CUDA features disabled gracefully)" >> $GITHUB_STEP_SUMMARY
|
||||
echo "" >> $GITHUB_STEP_SUMMARY
|
||||
echo "**Requirements:**" >> $GITHUB_STEP_SUMMARY
|
||||
echo "- Python 3.10, 3.11, or 3.12" >> $GITHUB_STEP_SUMMARY
|
||||
echo "- Linux x86-64 (manylinux_2_17 compatible)" >> $GITHUB_STEP_SUMMARY
|
||||
echo "- For CUDA features: NVIDIA driver with CUDA 11.8+ or 12.x support" >> $GITHUB_STEP_SUMMARY
|
||||
echo "" >> $GITHUB_STEP_SUMMARY
|
||||
echo "PyPI link: https://pypi.org/project/kt-kernel/${{ steps.get_version.outputs.VERSION }}/" >> $GITHUB_STEP_SUMMARY
|
||||
@@ -0,0 +1,129 @@
|
||||
name: Release sglang-kt to PyPI
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
paths:
|
||||
- "third_party/sglang"
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
test_pypi:
|
||||
description: 'Publish to TestPyPI instead of PyPI (for testing)'
|
||||
required: false
|
||||
default: 'false'
|
||||
type: choice
|
||||
options:
|
||||
- 'true'
|
||||
- 'false'
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
build-sglang-kt:
|
||||
name: Build sglang-kt wheel
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
submodules: recursive
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: '3.12'
|
||||
|
||||
- name: Install build tools
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install build wheel setuptools
|
||||
|
||||
- name: Build sglang-kt wheel
|
||||
working-directory: third_party/sglang/python
|
||||
run: |
|
||||
# Read version from ktransformers version.py
|
||||
KT_VERSION=$(python3 -c "exec(open('${{ github.workspace }}/version.py').read()); print(__version__)")
|
||||
export SGLANG_KT_VERSION="$KT_VERSION"
|
||||
echo "Building sglang-kt v${KT_VERSION} wheel..."
|
||||
python -m build --wheel -v
|
||||
|
||||
- name: Verify wheel
|
||||
working-directory: third_party/sglang/python
|
||||
run: |
|
||||
echo "Generated wheel:"
|
||||
ls -lh dist/
|
||||
# Verify the wheel has the correct package name
|
||||
ls dist/ | grep -q "sglang_kt" || (echo "ERROR: Wheel name does not contain sglang_kt" && exit 1)
|
||||
echo "Wheel name verified."
|
||||
|
||||
- name: Upload artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: sglang-kt-wheel
|
||||
path: third_party/sglang/python/dist/*.whl
|
||||
retention-days: 7
|
||||
|
||||
publish-pypi:
|
||||
name: Publish sglang-kt to PyPI
|
||||
needs: [build-sglang-kt]
|
||||
runs-on: ubuntu-latest
|
||||
if: github.repository == 'kvcache-ai/ktransformers' && github.ref == 'refs/heads/main'
|
||||
environment: prod
|
||||
permissions:
|
||||
id-token: write
|
||||
contents: read
|
||||
|
||||
steps:
|
||||
- name: Download wheel artifact
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: sglang-kt-wheel
|
||||
path: dist/
|
||||
|
||||
- name: Display wheels
|
||||
run: |
|
||||
echo "Wheels to publish:"
|
||||
ls -lh dist/
|
||||
|
||||
- name: Install twine
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install twine
|
||||
|
||||
- name: Publish to TestPyPI (if requested)
|
||||
if: github.event.inputs.test_pypi == 'true'
|
||||
env:
|
||||
TWINE_USERNAME: __token__
|
||||
TWINE_PASSWORD: ${{ secrets.TEST_PYPI_API_TOKEN }}
|
||||
run: |
|
||||
python -m twine upload \
|
||||
--repository testpypi \
|
||||
--skip-existing \
|
||||
--verbose \
|
||||
dist/*.whl
|
||||
|
||||
- name: Publish to PyPI
|
||||
if: github.event.inputs.test_pypi != 'true'
|
||||
env:
|
||||
TWINE_USERNAME: __token__
|
||||
TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
|
||||
run: |
|
||||
python -m twine upload \
|
||||
--skip-existing \
|
||||
--verbose \
|
||||
dist/*.whl
|
||||
|
||||
- name: Create release summary
|
||||
run: |
|
||||
echo "## sglang-kt Published to PyPI" >> $GITHUB_STEP_SUMMARY
|
||||
echo "" >> $GITHUB_STEP_SUMMARY
|
||||
echo "### Installation" >> $GITHUB_STEP_SUMMARY
|
||||
echo '```bash' >> $GITHUB_STEP_SUMMARY
|
||||
echo "pip install sglang-kt" >> $GITHUB_STEP_SUMMARY
|
||||
echo '```' >> $GITHUB_STEP_SUMMARY
|
||||
echo "" >> $GITHUB_STEP_SUMMARY
|
||||
echo "This is the kvcache-ai fork of SGLang with kt-kernel support." >> $GITHUB_STEP_SUMMARY
|
||||
echo "PyPI link: https://pypi.org/project/sglang-kt/" >> $GITHUB_STEP_SUMMARY
|
||||
@@ -0,0 +1,81 @@
|
||||
name: Sync sglang submodule
|
||||
|
||||
on:
|
||||
schedule:
|
||||
# Run daily at 08:00 UTC
|
||||
- cron: "0 8 * * *"
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
pull-requests: write
|
||||
|
||||
jobs:
|
||||
sync:
|
||||
name: Check for sglang-kt updates
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
submodules: true
|
||||
fetch-depth: 0
|
||||
token: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
- name: Update sglang submodule to latest main
|
||||
id: update
|
||||
run: |
|
||||
OLD_SHA=$(git -C third_party/sglang rev-parse HEAD)
|
||||
git submodule update --remote third_party/sglang
|
||||
NEW_SHA=$(git -C third_party/sglang rev-parse HEAD)
|
||||
|
||||
echo "old_sha=$OLD_SHA" >> "$GITHUB_OUTPUT"
|
||||
echo "new_sha=$NEW_SHA" >> "$GITHUB_OUTPUT"
|
||||
|
||||
if [ "$OLD_SHA" = "$NEW_SHA" ]; then
|
||||
echo "changed=false" >> "$GITHUB_OUTPUT"
|
||||
echo "sglang submodule is already up to date ($OLD_SHA)"
|
||||
else
|
||||
echo "changed=true" >> "$GITHUB_OUTPUT"
|
||||
|
||||
# Collect commit log between old and new
|
||||
COMMITS=$(git -C third_party/sglang log --oneline "$OLD_SHA..$NEW_SHA" | head -20)
|
||||
echo "commits<<EOF" >> "$GITHUB_OUTPUT"
|
||||
echo "$COMMITS" >> "$GITHUB_OUTPUT"
|
||||
echo "EOF" >> "$GITHUB_OUTPUT"
|
||||
|
||||
# sglang-kt version = ktransformers version (from version.py)
|
||||
VERSION=$(python3 -c "exec(open('version.py').read()); print(__version__)" 2>/dev/null || echo "unknown")
|
||||
echo "version=$VERSION" >> "$GITHUB_OUTPUT"
|
||||
|
||||
echo "sglang submodule updated: $OLD_SHA -> $NEW_SHA (v$VERSION)"
|
||||
fi
|
||||
|
||||
- name: Create pull request
|
||||
if: steps.update.outputs.changed == 'true'
|
||||
uses: peter-evans/create-pull-request@v6
|
||||
with:
|
||||
token: ${{ secrets.GITHUB_TOKEN }}
|
||||
commit-message: |
|
||||
[build]: sync sglang submodule to ${{ steps.update.outputs.new_sha }}
|
||||
branch: auto/sync-sglang
|
||||
delete-branch: true
|
||||
title: "[build] Sync sglang-kt submodule (v${{ steps.update.outputs.version }})"
|
||||
body: |
|
||||
Automated sync of `third_party/sglang` submodule to latest `main`.
|
||||
|
||||
**Old ref:** `${{ steps.update.outputs.old_sha }}`
|
||||
**New ref:** `${{ steps.update.outputs.new_sha }}`
|
||||
**sglang-kt version:** `${{ steps.update.outputs.version }}`
|
||||
|
||||
### Commits included
|
||||
```
|
||||
${{ steps.update.outputs.commits }}
|
||||
```
|
||||
|
||||
---
|
||||
*This PR was created automatically by the [sync-sglang-submodule](${{ github.server_url }}/${{ github.repository }}/actions/workflows/sync-sglang-submodule.yml) workflow.*
|
||||
labels: |
|
||||
dependencies
|
||||
automated
|
||||
+33
@@ -0,0 +1,33 @@
|
||||
__pycache__
|
||||
build
|
||||
.vscode
|
||||
*.so
|
||||
*.cache
|
||||
server.db
|
||||
logs
|
||||
node_modules
|
||||
*.nsys-rep
|
||||
.vs/
|
||||
*pycache*
|
||||
*build/
|
||||
.DS_Store
|
||||
compile_commands.json
|
||||
*.egg-info*
|
||||
*dist/
|
||||
ktransformers/server/local_store/
|
||||
ktransformers/server_test1.db
|
||||
*.patch
|
||||
img/
|
||||
tmp*.txt
|
||||
test.txt
|
||||
book
|
||||
ktransformers/tests/chat_txt.txt
|
||||
mmlu_result*
|
||||
ktransformers/ktransformers_ext/cuda_musa/
|
||||
test_prompt.txt
|
||||
csrc/demo
|
||||
build*
|
||||
CMakeFiles/
|
||||
kvc2/
|
||||
sched/
|
||||
*.png
|
||||
+14
@@ -0,0 +1,14 @@
|
||||
[submodule "third_party/llama.cpp"]
|
||||
path = third_party/llama.cpp
|
||||
url = https://github.com/ggerganov/llama.cpp.git
|
||||
[submodule "third_party/pybind11"]
|
||||
path = third_party/pybind11
|
||||
url = https://github.com/pybind/pybind11.git
|
||||
[submodule "third_party/custom_flashinfer"]
|
||||
path = third_party/custom_flashinfer
|
||||
url = https://github.com/kvcache-ai/custom_flashinfer.git
|
||||
branch = fix-precision-mla-merge-main
|
||||
[submodule "third_party/sglang"]
|
||||
path = third_party/sglang
|
||||
url = https://github.com/kvcache-ai/sglang.git
|
||||
branch = main
|
||||
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
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,0 +1,34 @@
|
||||
# Maintainers
|
||||
|
||||
This document lists the current maintainers and outlines their responsibilities.
|
||||
|
||||
## Current Maintainers
|
||||
|
||||
| Name | GitHub | Role | Affiliation | Email |
|
||||
|------|--------|------|-------------|-------|
|
||||
| Weiyu Xie | [@ErvinXie](https://github.com/ErvinXie) | Maintainer | [MADSys Lab](https://madsys.cs.tsinghua.edu.cn/) @ Tsinghua University | xwy21@mails.tsinghua.edu.cn |
|
||||
| Hongtao Chen | [@chenht2022](https://github.com/chenht2022) | Maintainer | [MADSys Lab](https://madsys.cs.tsinghua.edu.cn/) @ Tsinghua University | cht22@mails.tsinghua.edu.cn |
|
||||
| Jianwei Dong | [@ovowei](https://github.com/ovowei) | Maintainer | [MADSys Lab](https://madsys.cs.tsinghua.edu.cn/) @ Tsinghua University | dongjw24@mails.tsinghua.edu.cn |
|
||||
| Ziwei Yuan | [@KMSorSMS](https://github.com/KMSorSMS) | Maintainer | [Approaching.AI](http://approaching.ai/) | 2022090910005@std.uestc.edu.cn |
|
||||
| Qingliang Ou | [@ouqingliang](https://github.com/ouqingliang) | Maintainer | [MADSys Lab](https://madsys.cs.tsinghua.edu.cn/) @ Tsinghua University | oql@bupt.edu.cn |
|
||||
| Jiaqi Liao | [@SkqLiao](https://github.com/SkqLiao) | Maintainer | [Approaching.AI](http://approaching.ai/) | jiaqi.liao@bit.edu.cn |
|
||||
| Peilin Li | [@JimmyPeilinLi](https://github.com/JimmyPeilinLi) | Maintainer | [Approaching.AI](http://approaching.ai/) | lipeilin@mail.nwpu.edu.cn |
|
||||
| Xingxing Hao | [@mrhaoxx](https://github.com/mrhaoxx) | Maintainer | [Approaching.AI](http://approaching.ai/) | mr.haoxx@gmail.com |
|
||||
| Boxin Zhang | [@Atream](https://github.com/Atream) | Maintainer | [MADSys Lab](https://madsys.cs.tsinghua.edu.cn/) @ Tsinghua University | zhangbx24@mails.tsinghua.edu.cn |
|
||||
| Jingqi Tang | [@Azure-Tang](https://github.com/Azure-Tang) | Maintainer | [MADSys Lab](https://madsys.cs.tsinghua.edu.cn/) @ Tsinghua University | tangjq25@mails.tsinghua.edu.cn |
|
||||
| Jiahao Wang | [@qiyuxinlin](https://github.com/qiyuxinlin) | Maintainer | [Approaching.AI](http://approaching.ai/) | 202241050020@hdu.edu.cn |
|
||||
|
||||
## Responsibilities
|
||||
|
||||
Maintainers steward the project and keep it healthy for users and contributors.
|
||||
|
||||
- Review and approve pull requests; ensure changes meet quality, testing, and documentation standards.
|
||||
- Triage issues, keep labels organized, and respond to questions in a timely manner.
|
||||
- Uphold the project’s code of conduct and report violations when needed.
|
||||
- Maintain CI reliability and address regressions promptly.
|
||||
- Oversee releases and keep compatibility with supported dependency versions.
|
||||
- Protect project security and follow the security disclosure process.
|
||||
|
||||
## Becoming a Maintainer
|
||||
|
||||
We welcome contributors who show sustained, high-quality contributions and collaborative behavior. If you are interested, please contact an existing maintainer and share your recent contributions and areas of focus.
|
||||
@@ -0,0 +1 @@
|
||||
include version.py
|
||||
@@ -0,0 +1,163 @@
|
||||
<div align="center">
|
||||
<p align="center">
|
||||
|
||||
<picture>
|
||||
<img alt="KTransformers" src="https://github.com/user-attachments/assets/d5a2492f-a415-4456-af99-4ab102f13f8b" width=50%>
|
||||
|
||||
</picture>
|
||||
|
||||
</p>
|
||||
<h3>A Flexible Framework for Experiencing Cutting-edge LLM Inference/Fine-tune Optimizations</h3>
|
||||
<strong><a href="#-overview">🎯 Overview</a> | <a href="#-inference---high-performance-kt-kernel-serving">🚀 Inference</a> | <a href="#-sft---fine-tuning-with-llama-factory">🎓 SFT</a> | <a href="#-citation">🔥 Citation</a> | <a href="https://github.com/kvcache-ai/ktransformers/issues/1921">🚀 Roadmap(2026Q2)</a> </strong>
|
||||
</div>
|
||||
|
||||
## 🎯 Overview
|
||||
|
||||
KTransformers is a research project focused on efficient inference and fine-tuning of large language models through CPU-GPU heterogeneous computing. The project now exposes two user-facing capabilities from the kt-kernel source tree: [Inference](./kt-kernel/README.md) and [SFT](./doc/en/SFT/KTransformers-Fine-Tuning_Quick-Start.md).
|
||||
|
||||
## 🔥 Updates
|
||||
* **June 21, 2026**: MiniMax-M3 Day0 Support! ([Tutorial](./doc/en/kt-kernel/MiniMax-M3-Tutorial.md))
|
||||
* **June 17, 2026**: GLM-5.2 Day0 Support! ([Tutorial](./doc/en/kt-kernel/GLM-5.2-Tutorial.md))
|
||||
* **May 6, 2026**: KTransformers at [GOSIM Paris 2026](https://paris2026.gosim.org/zh/schedule/) — "Agentic AI on Edge" track. We'll present KT's inference performance on consumer hardware.
|
||||
* **May 02, 2026**: DeepSeek-V4-Flash Support! ([Tutorial](./doc/en/DeepSeek-V4-Flash.md))
|
||||
* **Apr 30, 2026**: KTransformers v0.6.1 refreshes kt-kernel inference and SFT docs with separate [Inference](./kt-kernel/README.md) and [SFT Quick Start](./doc/en/SFT/KTransformers-Fine-Tuning_Quick-Start.md) entry points.
|
||||
* **Mar 26, 2026**: Support AVX2-only CPU backend for KT-Kernel inference. ([Tutorial](./doc/en/kt-kernel/AVX2-Tutorial.md))
|
||||
* **Feb 13, 2026**: MiniMax-M2.5 Day0 Support! ([Tutorial](./doc/en/MiniMax-M2.5.md))
|
||||
* **Feb 12, 2026**: GLM-5 Day0 Support! ([Tutorial](./doc/en/kt-kernel/GLM-5-Tutorial.md))
|
||||
* **Jan 27, 2026**: Kimi-K2.5 Day0 Support! ([Tutorial](./doc/en/Kimi-K2.5.md)) ([SFT Tutorial](./doc/en/SFT_Installation_Guide_KimiK2.5.md))
|
||||
* **Jan 22, 2026**: Support [CPU-GPU Expert Scheduling](./doc/en/kt-kernel/experts-sched-Tutorial.md), [Native BF16 and FP8 per channel Precision](./doc/en/kt-kernel/Native-Precision-Tutorial.md) and [AutoDL unified fine-tuning and inference](./doc/zh/【云端低价训推】%20KTransformers%2BAutoDL%2BLlamaFactory：随用随租的低成本超大模型「微调%2B推理」一体化流程.pdf)
|
||||
* **Dec 24, 2025**: Support Native MiniMax-M2.1 inference. ([Tutorial](./doc/en/kt-kernel/MiniMax-M2.1-Tutorial.md))
|
||||
* **Dec 22, 2025**: Support RL-DPO fine-tuning with LLaMA-Factory. ([Tutorial](./doc/en/SFT/DPO_tutorial.md))
|
||||
* **Dec 5, 2025**: Support Native Kimi-K2-Thinking inference ([Tutorial](./doc/en/kt-kernel/Kimi-K2-Thinking-Native.md))
|
||||
* **Nov 6, 2025**: Support Kimi-K2-Thinking inference ([Tutorial](./doc/en/Kimi-K2-Thinking.md)) and fine-tune ([Tutorial](./doc/en/SFT_Installation_Guide_KimiK2.md))
|
||||
* **Nov 4, 2025**: KTransformers Fine-Tuning × LLaMA-Factory Integration. ([Tutorial](./doc/en/SFT/KTransformers-Fine-Tuning_User-Guide.md))
|
||||
* **Oct 27, 2025**: Support Ascend NPU. ([Tutorial](./doc/zh/DeepseekR1_V3_tutorial_zh_for_Ascend_NPU.md))
|
||||
* **Oct 10, 2025**: Integrating into SGLang. ([Roadmap](https://github.com/sgl-project/sglang/issues/11425), [Blog](https://lmsys.org/blog/2025-10-22-KTransformers/))
|
||||
* **Sept 11, 2025**: Support Qwen3-Next. ([Tutorial](./doc/en/Qwen3-Next.md))
|
||||
* **Sept 05, 2025**: Support Kimi-K2-0905. ([Tutorial](./doc/en/Kimi-K2.md))
|
||||
* **July 26, 2025**: Support SmallThinker and GLM4-MoE. ([Tutorial](./doc/en/SmallThinker_and_Glm4moe.md))
|
||||
* **July 11, 2025**: Support Kimi-K2. ([Tutorial](./doc/en/Kimi-K2.md))
|
||||
* **June 30, 2025**: Support 3-layer (GPU-CPU-Disk) [prefix cache](./doc/en/prefix_cache.md) reuse.
|
||||
* **May 14, 2025**: Support Intel Arc GPU ([Tutorial](./doc/en/xpu.md)).
|
||||
* **Apr 29, 2025**: Support AMX-Int8、 AMX-BF16 and Qwen3MoE ([Tutorial](./doc/en/AMX.md))
|
||||
* **Apr 9, 2025**: Experimental support for LLaMA 4 models ([Tutorial](./doc/en/llama4.md)).
|
||||
* **Apr 2, 2025**: Support Multi-concurrency. ([Tutorial](./doc/en/balance-serve.md)).
|
||||
* **Mar 15, 2025**: Support ROCm on AMD GPU ([Tutorial](./doc/en/ROCm.md)).
|
||||
* **Mar 5, 2025**: Support unsloth 1.58/2.51 bits weights and [IQ1_S/FP8 hybrid](./doc/en/fp8_kernel.md) weights. Support 139K [Longer Context](./doc/en/DeepseekR1_V3_tutorial.md#v022--v023-longer-context--fp8-kernel) for DeepSeek-V3 and R1 in 24GB VRAM.
|
||||
* **Feb 25, 2025**: Support [FP8 GPU kernel](./doc/en/fp8_kernel.md) for DeepSeek-V3 and R1; [Longer Context](./doc/en/DeepseekR1_V3_tutorial.md#v022-longer-context).
|
||||
* **Feb 15, 2025**: Longer Context (from 4K to 8K for 24GB VRAM) & Slightly Faster Speed （+15%, up to 16 Tokens/s), update [docs](./doc/en/DeepseekR1_V3_tutorial.md) and [online books](https://kvcache-ai.github.io/ktransformers/).
|
||||
* **Feb 10, 2025**: Support Deepseek-R1 and V3 on single (24GB VRAM)/multi gpu and 382G DRAM, up to 3~28x speedup. For detailed show case and reproduction tutorial, see [here](./doc/en/DeepseekR1_V3_tutorial.md).
|
||||
* **Aug 28, 2024**: Decrease DeepseekV2's required VRAM from 21G to 11G.
|
||||
* **Aug 15, 2024**: Update detailed [tutorial](doc/en/injection_tutorial.md) for injection and multi-GPU.
|
||||
* **Aug 14, 2024**: Support llamfile as linear backend.
|
||||
* **Aug 12, 2024**: Support multiple GPU; Support new model: mixtral 8\*7B and 8\*22B; Support q2k, q3k, q5k dequant on gpu.
|
||||
* **Aug 9, 2024**: Support windows native.
|
||||
|
||||
---
|
||||
|
||||
## 📦 Capabilities
|
||||
|
||||
### 🚀 [Inference](./kt-kernel/README.md) - High-Performance kt-kernel Serving
|
||||
|
||||
CPU-optimized kernel operations for heterogeneous LLM inference.
|
||||
|
||||
<img width="1049" height="593" alt="image" src="https://github.com/user-attachments/assets/68f423da-3f55-4025-bdc9-9ceaa554f00b" />
|
||||
|
||||
|
||||
**Key Features:**
|
||||
- **AMX/AVX Acceleration**: Intel AMX and AVX512/AVX2 optimized kernels for INT4/INT8 quantized inference
|
||||
- **MoE Optimization**: Efficient Mixture-of-Experts inference with NUMA-aware memory management
|
||||
- **Quantization Support**: CPU-side INT4/INT8 quantized weights, GPU-side GPTQ support
|
||||
- **Easy Integration**: Clean Python API for SGLang and other frameworks
|
||||
|
||||
**Quick Start:**
|
||||
```bash
|
||||
cd kt-kernel
|
||||
pip install .
|
||||
```
|
||||
|
||||
**Use Cases:**
|
||||
|
||||
- CPU-GPU hybrid inference for large MoE models
|
||||
- Integration with SGLang for production serving
|
||||
- Heterogeneous expert placement (hot experts on GPU, cold experts on CPU)
|
||||
|
||||
**Performance Examples:**
|
||||
| Model | Hardware Configuration | Total Throughput | Output Throughput |
|
||||
|-------|------------------------|------------------|-------------------|
|
||||
| DeepSeek-R1-0528 (FP8) | 8×L20 GPU + Xeon Gold 6454S | 227.85 tokens/s | 87.58 tokens/s (8-way concurrency) |
|
||||
|
||||
👉 **[Full Documentation →](./kt-kernel/README.md)**
|
||||
|
||||
---
|
||||
|
||||
### 🎓 [SFT](./doc/en/SFT/KTransformers-Fine-Tuning_Quick-Start.md) - Fine-Tuning with LLaMA-Factory
|
||||
|
||||
KTransformers × LLaMA-Factory integration for ultra-large MoE model fine-tuning.
|
||||
|
||||
|
||||
|
||||
**Key Features:**
|
||||
- **Multi-Backend Support**: CPU/GPU hybrid fine-tuning with INT8/INT4 quantization
|
||||
- **Ultra-Large MoE Support**: Fine-tune models like DeepSeek-V3/R1 on limited GPU memory
|
||||
- **Faster than ZeRO-Offload**: 6-12x training speedup in benchmarked MoE SFT workloads
|
||||
- **Lower CPU Memory**: About half the CPU memory of the previous KT SFT path in the benchmarked setup
|
||||
- **LLaMA-Factory Integration**: Seamless integration with popular fine-tuning framework
|
||||
|
||||
| Model | GPU Memory | Training Speed | Hardware |
|
||||
|-------|------------|----------------|----------|
|
||||
| DeepSeek-V3 | ~80GB total | 3.7 it/s | 4x RTX 4090 |
|
||||
| DeepSeek-R1 | ~80GB total | 3.7 it/s | 4x RTX 4090 |
|
||||
| Qwen3-30B-A3B | ~24GB total | 8+ it/s | 1x RTX 4090 |
|
||||
|
||||
**Quick Start:**
|
||||
```bash
|
||||
cd /path/to/LLaMA-Factory
|
||||
pip install -e .
|
||||
pip install -r requirements/ktransformers.txt
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch \
|
||||
--config_file examples/ktransformers/accelerate/fsdp2_kt_int8.yaml \
|
||||
src/train.py \
|
||||
examples/ktransformers/train_lora/qwen3_5moe_lora_sft_kt.yaml
|
||||
```
|
||||
|
||||
👉 **[Quick Start →](./doc/en/SFT/KTransformers-Fine-Tuning_Quick-Start.md)**
|
||||
👉 **[Full Documentation →](./doc/en/SFT/KTransformers-Fine-Tuning_User-Guide.md)**
|
||||
|
||||
---
|
||||
|
||||
## 🔥 Citation
|
||||
|
||||
If you use KTransformers in your research, please cite our paper:
|
||||
|
||||
```bibtex
|
||||
@inproceedings{10.1145/3731569.3764843,
|
||||
title = {KTransformers: Unleashing the Full Potential of CPU/GPU Hybrid Inference for MoE Models},
|
||||
author = {Chen, Hongtao and Xie, Weiyu and Zhang, Boxin and Tang, Jingqi and Wang, Jiahao and Dong, Jianwei and Chen, Shaoyuan and Yuan, Ziwei and Lin, Chen and Qiu, Chengyu and Zhu, Yuening and Ou, Qingliang and Liao, Jiaqi and Chen, Xianglin and Ai, Zhiyuan and Wu, Yongwei and Zhang, Mingxing},
|
||||
booktitle = {Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles},
|
||||
year = {2025}
|
||||
}
|
||||
```
|
||||
|
||||
## 👥 Contributors & Team
|
||||
|
||||
Developed and maintained by:
|
||||
- [MADSys Lab](https://madsys.cs.tsinghua.edu.cn/) @ Tsinghua University
|
||||
- [Approaching.AI](http://approaching.ai/)
|
||||
- [9#AISoft](https://github.com/aisoft9)
|
||||
- Community contributors
|
||||
|
||||
We welcome contributions! Please feel free to submit issues and pull requests.
|
||||
|
||||
## 💬 Community & Support
|
||||
|
||||
- **GitHub Issues**: [Report bugs or request features](https://github.com/kvcache-ai/ktransformers/issues)
|
||||
- **WeChat Group**: See [archive/WeChatGroup.png](./archive/WeChatGroup.png)
|
||||
|
||||
## 📦 KT original Code
|
||||
|
||||
The original integrated KTransformers framework has been archived to the [`archive/`](./archive/) directory for reference. The project now organizes the two capabilities above from the kt-kernel source tree for clearer documentation and maintenance.
|
||||
|
||||
For the original documentation with full quick-start guides and examples, see:
|
||||
- [archive/README.md](./archive/README.md) (English)
|
||||
- [archive/README_ZH.md](./archive/README_ZH.md) (中文)
|
||||
@@ -0,0 +1,7 @@
|
||||
# WeHub 来源说明
|
||||
|
||||
- 原始项目：`kvcache-ai/ktransformers`
|
||||
- 原始仓库：https://github.com/kvcache-ai/ktransformers
|
||||
- 导入方式：上游默认分支的最新快照
|
||||
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
|
||||
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
|
||||
+151
@@ -0,0 +1,151 @@
|
||||
<div align="center">
|
||||
<p align="center">
|
||||
|
||||
<picture>
|
||||
<img alt="KTransformers" src="https://github.com/user-attachments/assets/d5a2492f-a415-4456-af99-4ab102f13f8b" width=50%>
|
||||
|
||||
</picture>
|
||||
|
||||
</p>
|
||||
<h3>一个用于体验尖端 LLM 推理/微调优化的灵活框架</h3>
|
||||
<strong><a href="#-概览">🎯 概览</a> | <a href="#-推理---kt-kernel-高性能推理">🚀 推理</a> | <a href="#-sft---llama-factory-微调">🎓 SFT</a> | <a href="#-引用">🔥 引用</a> </strong>
|
||||
</div>
|
||||
|
||||
## 🎯 概览
|
||||
|
||||
KTransformers 是一个专注于通过 CPU-GPU 异构计算实现大语言模型高效推理和微调的研究项目。目前两个面向用户的能力都来自 kt-kernel 源码目录：[推理](./kt-kernel/README.md) 和 [SFT](./doc/en/SFT/KTransformers-Fine-Tuning_Quick-Start.md)。
|
||||
|
||||
## 🔥 更新
|
||||
|
||||
* **2026 年 4 月 30 日**：KTransformers v0.6.1 更新 kt-kernel 推理和 SFT 文档，提供独立的[推理](./kt-kernel/README.md)和 [SFT Quick Start](./doc/en/SFT/KTransformers-Fine-Tuning_Quick-Start.md)入口。
|
||||
* **2025 年 12 月 5 日**：支持原生 Kimi-K2-Thinking 推理（[教程](./doc/en/kt-kernel/Kimi-K2-Thinking-Native.md)）
|
||||
* **2025 年 11 月 6 日**：支持 Kimi-K2-Thinking 推理（[教程](./doc/en/Kimi-K2-Thinking.md)）和微调（[教程](./doc/en/SFT_Installation_Guide_KimiK2.md)）
|
||||
* **2025 年 11 月 4 日**：KTransformers 微调 × LLaMA-Factory 集成（[教程](./doc/en/SFT/KTransformers-Fine-Tuning_User-Guide.md)）
|
||||
* **2025 年 10 月 27 日**：支持昇腾 NPU（[教程](./doc/zh/DeepseekR1_V3_tutorial_zh_for_Ascend_NPU.md)）
|
||||
* **2025 年 10 月 10 日**：集成到 SGLang（[路线图](https://github.com/sgl-project/sglang/issues/11425)，[博客](https://lmsys.org/blog/2025-10-22-KTransformers/)）
|
||||
* **2025 年 9 月 11 日**：支持 Qwen3-Next（[教程](./doc/en/Qwen3-Next.md)）
|
||||
* **2025 年 9 月 5 日**：支持 Kimi-K2-0905（[教程](./doc/en/Kimi-K2.md)）
|
||||
* **2025 年 7 月 26 日**：支持 SmallThinker 和 GLM4-MoE（[教程](./doc/en/SmallThinker_and_Glm4moe.md)）
|
||||
* **2025 年 7 月 11 日**：支持 Kimi-K2（[教程](./doc/en/Kimi-K2.md)）
|
||||
* **2025 年 6 月 30 日**：支持 3 层（GPU-CPU-磁盘）[前缀缓存](./doc/en/prefix_cache.md)复用
|
||||
* **2025 年 5 月 14 日**：支持 Intel Arc GPU（[教程](./doc/en/xpu.md)）
|
||||
* **2025 年 4 月 29 日**：支持 AMX-Int8、AMX-BF16 和 Qwen3MoE（[教程](./doc/en/AMX.md)）
|
||||
* **2025 年 4 月 9 日**：实验性支持 LLaMA 4 模型（[教程](./doc/en/llama4.md)）
|
||||
* **2025 年 4 月 2 日**：支持多并发（[教程](./doc/en/balance-serve.md)）
|
||||
* **2025 年 3 月 15 日**：支持 AMD GPU 上的 ROCm（[教程](./doc/en/ROCm.md)）
|
||||
* **2025 年 3 月 5 日**：支持 unsloth 1.58/2.51 位权重和 [IQ1_S/FP8 混合](./doc/en/fp8_kernel.md)权重。在 24GB VRAM 中支持 DeepSeek-V3 和 R1 的 139K [更长上下文](./doc/en/DeepseekR1_V3_tutorial.md#v022--v023-longer-context--fp8-kernel)
|
||||
* **2025 年 2 月 25 日**：为 DeepSeek-V3 和 R1 支持 [FP8 GPU 内核](./doc/en/fp8_kernel.md)；[更长上下文](./doc/en/DeepseekR1_V3_tutorial.md#v022-longer-context)
|
||||
* **2025 年 2 月 15 日**：更长上下文（24GB VRAM 从 4K 到 8K）& 速度稍快（+15%，最高 16 Tokens/s），更新[文档](./doc/en/DeepseekR1_V3_tutorial.md)和[在线手册](https://kvcache-ai.github.io/ktransformers/)
|
||||
* **2025 年 2 月 10 日**：支持 Deepseek-R1 和 V3 在单 GPU（24GB VRAM）/多 GPU 和 382GB DRAM 上运行，速度提升高达 3~28 倍。详细案例展示和复现教程请参见[这里](./doc/en/DeepseekR1_V3_tutorial.md)
|
||||
* **2024 年 8 月 28 日**：将 DeepseekV2 所需的 VRAM 从 21GB 降低到 11GB
|
||||
* **2024 年 8 月 15 日**：更新了关于注入和多 GPU 的详细[教程](doc/en/injection_tutorial.md)
|
||||
* **2024 年 8 月 14 日**：支持 llamfile 作为线性后端
|
||||
* **2024 年 8 月 12 日**：支持多 GPU；支持新模型：mixtral 8\*7B 和 8\*22B；支持 GPU 上的 q2k、q3k、q5k 去量化
|
||||
* **2024 年 8 月 9 日**：支持 Windows 原生环境
|
||||
|
||||
---
|
||||
|
||||
## 📦 功能入口
|
||||
|
||||
### 🚀 [推理](./kt-kernel/README.md) - kt-kernel 高性能推理
|
||||
|
||||
用于异构 LLM 推理的 CPU 优化内核操作。
|
||||
|
||||
|
||||
|
||||
**主要特性：**
|
||||
- **AMX/AVX 加速**：Intel AMX 和 AVX512/AVX2 优化的内核，用于 INT4/INT8 量化推理
|
||||
- **MoE 优化**：高效的专家混合推理，具有 NUMA 感知内存管理
|
||||
- **量化支持**：CPU 端 INT4/INT8 量化权重，GPU 端 GPTQ 支持
|
||||
- **易于集成**：为 SGLang 和其他框架提供简洁的 Python API
|
||||
|
||||
**快速开始：**
|
||||
```bash
|
||||
cd kt-kernel
|
||||
pip install .
|
||||
```
|
||||
|
||||
**使用场景：**
|
||||
|
||||
- 大型 MoE 模型的 CPU-GPU 混合推理
|
||||
- 与 SGLang 集成用于生产服务
|
||||
- 异构专家放置（热专家在 GPU 上，冷专家在 CPU 上）
|
||||
|
||||
**性能示例：**
|
||||
| 模型 | 硬件配置 | 总吞吐量 | 输出吞吐量 |
|
||||
|-------|------------------------|------------------|-------------------|
|
||||
| DeepSeek-R1-0528 (FP8) | 8×L20 GPU + Xeon Gold 6454S | 227.85 tokens/s | 87.58 tokens/s（8 路并发）|
|
||||
|
||||
👉 **[完整文档 →](./kt-kernel/README.md)**
|
||||
|
||||
---
|
||||
|
||||
### 🎓 [SFT](./doc/en/SFT/KTransformers-Fine-Tuning_Quick-Start.md) - LLaMA-Factory 微调
|
||||
|
||||
KTransformers × LLaMA-Factory 集成，面向超大 MoE 模型微调。
|
||||
|
||||
|
||||
|
||||
**主要特性:**
|
||||
- **多后端支持**: CPU/GPU 混合微调，支持 INT8/INT4 量化
|
||||
- **超大 MoE 支持**: 在有限 GPU 内存下微调 DeepSeek-V3/R1 等模型
|
||||
- **相对 ZeRO-Offload 加速**: 在基准 MoE SFT 任务中训练速度提升 6-12 倍
|
||||
- **降低 CPU 内存**: 相比上一版 KT SFT 路径，基准配置下 CPU 内存约降至 1/2
|
||||
- **LLaMA-Factory 集成**: 与流行微调框架无缝集成
|
||||
|
||||
| 模型 | GPU 内存 | 训练速度 | 硬件 |
|
||||
|-------|------------|----------------|----------|
|
||||
| DeepSeek-V3 | ~80GB 总计 | 3.7 it/s | 4x RTX 4090 |
|
||||
| DeepSeek-R1 | ~80GB 总计 | 3.7 it/s | 4x RTX 4090 |
|
||||
| Qwen3-30B-A3B | ~24GB 总计 | 8+ it/s | 1x RTX 4090 |
|
||||
|
||||
**快速开始:**
|
||||
```bash
|
||||
cd /path/to/LLaMA-Factory
|
||||
pip install -e .
|
||||
pip install -r requirements/ktransformers.txt
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch \
|
||||
--config_file examples/ktransformers/accelerate/fsdp2_kt_int8.yaml \
|
||||
src/train.py \
|
||||
examples/ktransformers/train_lora/qwen3_5moe_lora_sft_kt.yaml
|
||||
```
|
||||
|
||||
👉 **[Quick Start →](./doc/en/SFT/KTransformers-Fine-Tuning_Quick-Start.md)**
|
||||
👉 **[完整文档 →](./doc/en/SFT/KTransformers-Fine-Tuning_User-Guide.md)**
|
||||
|
||||
---
|
||||
|
||||
## 🔥 引用
|
||||
|
||||
如果您在研究中使用了 KTransformers，请引用我们的论文：
|
||||
|
||||
```bibtex
|
||||
@inproceedings{10.1145/3731569.3764843,
|
||||
title = {KTransformers: Unleashing the Full Potential of CPU/GPU Hybrid Inference for MoE Models},
|
||||
author = {Chen, Hongtao and Xie, Weiyu and Zhang, Boxin and Tang, Jingqi and Wang, Jiahao and Dong, Jianwei and Chen, Shaoyuan and Yuan, Ziwei and Lin, Chen and Qiu, Chengyu and Zhu, Yuening and Ou, Qingliang and Liao, Jiaqi and Chen, Xianglin and Ai, Zhiyuan and Wu, Yongwei and Zhang, Mingxing},
|
||||
booktitle = {Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles},
|
||||
year = {2025}
|
||||
}
|
||||
```
|
||||
|
||||
## 👥 贡献者与团队
|
||||
|
||||
由以下团队开发和维护：
|
||||
- 清华大学 [MADSys 实验室](https://madsys.cs.tsinghua.edu.cn/)
|
||||
- [Approaching.AI](http://approaching.ai/)
|
||||
- 社区贡献者
|
||||
|
||||
我们欢迎贡献！请随时提交问题和拉取请求。
|
||||
|
||||
## 💬 社区与支持
|
||||
|
||||
- **GitHub Issues**：[报告问题或请求功能](https://github.com/kvcache-ai/ktransformers/issues)
|
||||
- **微信群**：请参见 [archive/WeChatGroup.png](./archive/WeChatGroup.png)
|
||||
|
||||
## 📦 KT原仓库
|
||||
|
||||
原始的集成 KTransformers 框架已归档到 [`archive/`](./archive/) 目录以供参考。该项目现在围绕 kt-kernel 源码树中的上述两个能力入口组织文档和维护。
|
||||
|
||||
有关原始文档以及完整的快速入门指南和示例，请参见：
|
||||
- [archive/README.md](./archive/README.md)（英文）
|
||||
- [archive/README_ZH.md](./archive/README_ZH.md)（中文）
|
||||
@@ -0,0 +1,19 @@
|
||||
FROM pytorch/pytorch:2.5.1-cuda12.1-cudnn9-devel as compile_server
|
||||
WORKDIR /workspace
|
||||
ENV CUDA_HOME /usr/local/cuda
|
||||
RUN <<EOF
|
||||
apt update -y && apt install -y --no-install-recommends \
|
||||
git \
|
||||
wget \
|
||||
vim \
|
||||
gcc \
|
||||
g++ \
|
||||
cmake &&
|
||||
rm -rf /var/lib/apt/lists/* &&
|
||||
pip install --upgrade pip &&
|
||||
pip install ninja pyproject numpy cpufeature &&
|
||||
pip install flash-attn &&
|
||||
cp /usr/lib/x86_64-linux-gnu/libstdc++.so.6 /opt/conda/lib/
|
||||
EOF
|
||||
# Set the default shell to bash
|
||||
CMD ["/bin/bash"]
|
||||
@@ -0,0 +1,34 @@
|
||||
{
|
||||
"name": "Ktrans Dev Container",
|
||||
"privileged": true,
|
||||
"build": {
|
||||
"dockerfile": "Dockerfile",
|
||||
"context": "..",
|
||||
"args": {
|
||||
"http_proxy": "${env:http_proxy}",
|
||||
"https_proxy": "${env:https_proxy}",
|
||||
}
|
||||
},
|
||||
"runArgs": [
|
||||
"--network=host",
|
||||
"--gpus",
|
||||
"all"
|
||||
// "--gpu all"
|
||||
],
|
||||
"workspaceFolder": "/workspace",
|
||||
"workspaceMount": "source=${localWorkspaceFolder},target=/workspace,type=bind,consistency=cached",
|
||||
"mounts": [
|
||||
"source=/mnt/data,target=/mnt/incontainer,type=bind,consistency=cached"
|
||||
],
|
||||
"customizations": {
|
||||
"vscode": {
|
||||
"extensions": [
|
||||
],
|
||||
"settings": {
|
||||
"terminal.integrated.shell.linux": "/bin/bash",
|
||||
"cmake.configureOnOpen": true,
|
||||
"cmake.generator": "Ninja"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
[flake8]
|
||||
max-line-length = 120
|
||||
extend-select = B950
|
||||
extend-ignore = E203,E501,E701, B001,B006,B007,B008,B009,B010,B011,B016,B028,B031,B950,E265,E266,E401,E402,E711,E712,E713,E721,E722,E731,F401,F403,F405,F541,F811,F821,F841,W391
|
||||
@@ -0,0 +1,28 @@
|
||||
[submodule "third_party/llama.cpp"]
|
||||
path = archive/third_party/llama.cpp
|
||||
url = https://github.com/ggerganov/llama.cpp.git
|
||||
[submodule "third_party/pybind11"]
|
||||
path = archive/third_party/pybind11
|
||||
url = https://github.com/pybind/pybind11.git
|
||||
[submodule "third_party/spdlog"]
|
||||
path = archive/third_party/spdlog
|
||||
url = https://github.com/gabime/spdlog.git
|
||||
[submodule "third_party/custom_flashinfer"]
|
||||
path = archive/third_party/custom_flashinfer
|
||||
url = https://github.com/kvcache-ai/custom_flashinfer.git
|
||||
branch = fix-precision-mla-merge-main
|
||||
[submodule "third_party/xxHash"]
|
||||
path = archive/third_party/xxHash
|
||||
url = https://github.com/Cyan4973/xxHash.git
|
||||
[submodule "third_party/prometheus-cpp"]
|
||||
path = archive/third_party/prometheus-cpp
|
||||
url = https://github.com/jupp0r/prometheus-cpp
|
||||
[submodule "third_party/PhotonLibOS"]
|
||||
path = archive/third_party/PhotonLibOS
|
||||
url = https://github.com/alibaba/PhotonLibOS.git
|
||||
[submodule "kt-kernel/third_party/llama.cpp"]
|
||||
path = kt-kernel/third_party/llama.cpp
|
||||
url = https://github.com/ggerganov/llama.cpp.git
|
||||
[submodule "kt-kernel/third_party/pybind11"]
|
||||
path = kt-kernel/third_party/pybind11
|
||||
url = https://github.com/pybind/pybind11.git
|
||||
@@ -0,0 +1,6 @@
|
||||
[MASTER]
|
||||
extension-pkg-whitelist=pydantic
|
||||
max-line-length=120
|
||||
|
||||
[MESSAGES CONTROL]
|
||||
disable=missing-function-docstring
|
||||
@@ -0,0 +1,64 @@
|
||||
FROM pytorch/pytorch:2.5.1-cuda12.1-cudnn9-devel as compile_server
|
||||
|
||||
|
||||
ARG CPU_INSTRUCT=NATIVE
|
||||
|
||||
# 设置工作目录和 CUDA 路径
|
||||
WORKDIR /workspace
|
||||
ENV CUDA_HOME=/usr/local/cuda
|
||||
|
||||
|
||||
|
||||
# 安装依赖
|
||||
RUN apt update -y
|
||||
RUN apt install -y --no-install-recommends \
|
||||
libtbb-dev \
|
||||
libssl-dev \
|
||||
libcurl4-openssl-dev \
|
||||
libaio1 \
|
||||
libaio-dev \
|
||||
libfmt-dev \
|
||||
libgflags-dev \
|
||||
zlib1g-dev \
|
||||
patchelf \
|
||||
git \
|
||||
wget \
|
||||
vim \
|
||||
gcc \
|
||||
g++ \
|
||||
cmake
|
||||
# 拷贝代码
|
||||
RUN git clone https://github.com/kvcache-ai/ktransformers.git
|
||||
# 清理 apt 缓存
|
||||
RUN rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# 进入项目目录
|
||||
WORKDIR /workspace/ktransformers
|
||||
# 初始化子模块
|
||||
RUN git submodule update --init --recursive
|
||||
|
||||
# 升级 pip
|
||||
RUN pip install --upgrade pip
|
||||
|
||||
# 安装构建依赖
|
||||
RUN pip install ninja pyproject numpy cpufeature aiohttp zmq openai
|
||||
|
||||
# 安装 flash-attn（提前装可以避免后续某些编译依赖出错）
|
||||
RUN pip install flash-attn
|
||||
|
||||
# 安装 ktransformers 本体（含编译）
|
||||
RUN CPU_INSTRUCT=${CPU_INSTRUCT} \
|
||||
USE_BALANCE_SERVE=1 \
|
||||
KTRANSFORMERS_FORCE_BUILD=TRUE \
|
||||
TORCH_CUDA_ARCH_LIST="8.0;8.6;8.7;8.9;9.0+PTX" \
|
||||
pip install . --no-build-isolation --verbose
|
||||
|
||||
RUN pip install third_party/custom_flashinfer/
|
||||
# 清理 pip 缓存
|
||||
RUN pip cache purge
|
||||
|
||||
# 拷贝 C++ 运行时库
|
||||
RUN cp /usr/lib/x86_64-linux-gnu/libstdc++.so.6 /opt/conda/lib/
|
||||
|
||||
# 保持容器运行（调试用）
|
||||
ENTRYPOINT ["tail", "-f", "/dev/null"]
|
||||
@@ -0,0 +1,68 @@
|
||||
# Base image
|
||||
FROM intel/oneapi-basekit:2025.0.1-0-devel-ubuntu22.04
|
||||
|
||||
ARG http_proxy
|
||||
ARG https_proxy
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
ENV CONDA_DIR=/opt/conda
|
||||
|
||||
# Install dependencies
|
||||
RUN apt-get update && apt-get install -y \
|
||||
wget \
|
||||
curl \
|
||||
bash \
|
||||
git \
|
||||
vim \
|
||||
ca-certificates \
|
||||
binutils \
|
||||
cmake \
|
||||
g++ \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install Miniforge
|
||||
RUN wget https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh -O /tmp/miniforge.sh && \
|
||||
bash /tmp/miniforge.sh -b -p $CONDA_DIR && \
|
||||
rm /tmp/miniforge.sh && \
|
||||
$CONDA_DIR/bin/conda clean -afy
|
||||
|
||||
# Add conda to PATH
|
||||
ENV PATH=$CONDA_DIR/bin:$PATH
|
||||
|
||||
RUN bash -c "\
|
||||
source /opt/conda/etc/profile.d/conda.sh && \
|
||||
conda create --name ktransformers python=3.11 -y && \
|
||||
conda activate ktransformers && \
|
||||
conda env list && \
|
||||
conda install -c conda-forge libstdcxx-ng -y && \
|
||||
strings \$(find /opt/conda/envs/ktransformers/lib -name 'libstdc++.so.6') | grep GLIBCXX | grep 3.4.32 \
|
||||
"
|
||||
|
||||
RUN bash -c "\
|
||||
source /opt/conda/etc/profile.d/conda.sh && \
|
||||
conda activate ktransformers && \
|
||||
pip install ipex-llm[xpu_2.6]==2.3.0b20250518 --extra-index-url https://download.pytorch.org/whl/xpu && \
|
||||
pip uninstall -y torch torchvision torchaudio && \
|
||||
pip install torch==2.7+xpu torchvision torchaudio --index-url https://download.pytorch.org/whl/test/xpu && \
|
||||
pip uninstall -y intel-opencl-rt dpcpp-cpp-rt && \
|
||||
pip list \
|
||||
"
|
||||
|
||||
# Clone and set up ktransformers repo
|
||||
RUN bash -c "\
|
||||
source $CONDA_DIR/etc/profile.d/conda.sh && \
|
||||
conda activate ktransformers && \
|
||||
git clone https://github.com/kvcache-ai/ktransformers.git && \
|
||||
cd ktransformers && \
|
||||
git submodule update --init && \
|
||||
sed -i 's/torch\.xpu\.is_available()/True/g' setup.py && \
|
||||
bash install.sh --dev xpu \
|
||||
"
|
||||
|
||||
# Init conda and prepare bashrc
|
||||
RUN conda init bash && \
|
||||
echo "source $CONDA_DIR/etc/profile.d/conda.sh" >> ~/.bashrc && \
|
||||
echo "conda activate ktransformers" >> ~/.bashrc
|
||||
|
||||
WORKDIR /ktransformers/
|
||||
CMD ["bash"]
|
||||
+201
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
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,0 +1,13 @@
|
||||
graft third_party
|
||||
graft ktransformers
|
||||
graft local_chat.py
|
||||
graft csrc
|
||||
include LICENSE README.md
|
||||
prune ktransformers/website
|
||||
prune ktransformers/logs
|
||||
prune ktransformers.egg-info
|
||||
prune third_party/llama.cpp/models
|
||||
graft ktransformers/website/dist
|
||||
global-exclude __pycache__
|
||||
include KTransformersOps.*.so
|
||||
include cpuinfer_ext.*.so
|
||||
@@ -0,0 +1,32 @@
|
||||
flake_find:
|
||||
cd ktransformers && flake8 | grep -Eo '[A-Z][0-9]{3}' | sort | uniq| paste -sd ',' -
|
||||
format:
|
||||
@cd ktransformers && black .
|
||||
@black setup.py
|
||||
dev_install:
|
||||
# clear build dirs
|
||||
rm -rf build
|
||||
rm -rf *.egg-info
|
||||
rm -rf ktransformers/ktransformers_ext/build
|
||||
rm -rf ktransformers/ktransformers_ext/cuda/build
|
||||
rm -rf ktransformers/ktransformers_ext/cuda/dist
|
||||
rm -rf ktransformers/ktransformers_ext/cuda/*.egg-info
|
||||
|
||||
# install ktransformers
|
||||
echo "Installing python dependencies from requirements.txt"
|
||||
pip install -r requirements-local_chat.txt
|
||||
|
||||
echo "Installing ktransformers"
|
||||
KTRANSFORMERS_FORCE_BUILD=TRUE pip install -e . -v --no-build-isolation
|
||||
echo "Installation completed successfully"
|
||||
clean:
|
||||
rm -rf build
|
||||
rm -rf *.egg-info
|
||||
rm -rf ktransformers/ktransformers_ext/build
|
||||
rm -rf ktransformers/ktransformers_ext/cuda/build
|
||||
rm -rf ktransformers/ktransformers_ext/cuda/dist
|
||||
rm -rf ktransformers/ktransformers_ext/cuda/*.egg-info
|
||||
install_numa:
|
||||
USE_NUMA=1 make dev_install
|
||||
install_no_numa:
|
||||
env -u USE_NUMA make dev_install
|
||||
@@ -0,0 +1,136 @@
|
||||
<div align="center">
|
||||
<p align="center">
|
||||
<picture>
|
||||
<img alt="KTransformers" src="https://github.com/user-attachments/assets/d5a2492f-a415-4456-af99-4ab102f13f8b" width=50%>
|
||||
</picture>
|
||||
</p>
|
||||
<h3>High-Performance CPU-GPU Hybrid Inference for Large Language Models</h3>
|
||||
</div>
|
||||
|
||||
## 🎯 Overview
|
||||
|
||||
KTransformers is a research project focused on efficient inference and fine-tuning of large language models through CPU-GPU heterogeneous computing. The project has evolved into **two core modules**: [kt-kernel](./kt-kernel/) and [kt-sft](./kt-sft/).
|
||||
|
||||
## 🔥 Updates
|
||||
|
||||
* **Nov 6, 2025**: Support Kimi-K2-Thinking inference and fine-tune
|
||||
* **Nov 4, 2025**: KTransformers Fine-Tuning × LLaMA-Factory Integration
|
||||
* **Oct 27, 2025**: Support Ascend NPU
|
||||
* **Oct 10, 2025**: Integrating into SGLang ([Roadmap](https://github.com/sgl-project/sglang/issues/11425), [Blog](https://lmsys.org/blog/2025-10-22-KTransformers/))
|
||||
* **Sept 11, 2025**: Support Qwen3-Next
|
||||
* **Sept 05, 2025**: Support Kimi-K2-0905
|
||||
* **July 26, 2025**: Support SmallThinker and GLM4-MoE
|
||||
* **June 30, 2025**: Support 3-layer (GPU-CPU-Disk) prefix cache reuse
|
||||
* **May 14, 2025**: Support Intel Arc GPU
|
||||
* **Apr 29, 2025**: Support AMX-Int8、AMX-BF16 and Qwen3MoE
|
||||
* **Apr 9, 2025**: Experimental support for LLaMA 4 models
|
||||
* **Apr 2, 2025**: Support Multi-concurrency
|
||||
* **Mar 15, 2025**: Support ROCm on AMD GPU
|
||||
* **Mar 5, 2025**: Support unsloth 1.58/2.51 bits weights and IQ1_S/FP8 hybrid weights; 139K longer context for DeepSeek-V3/R1
|
||||
* **Feb 25, 2025**: Support FP8 GPU kernel for DeepSeek-V3 and R1
|
||||
* **Feb 10, 2025**: Support Deepseek-R1 and V3, up to 3~28x speedup
|
||||
|
||||
---
|
||||
|
||||
## 📦 Core Modules
|
||||
|
||||
### 🚀 [kt-kernel](./kt-kernel/) - High-Performance Inference Kernels
|
||||
|
||||
CPU-optimized kernel operations for heterogeneous LLM inference.
|
||||
|
||||
|
||||
|
||||
**Key Features:**
|
||||
- **AMX/AVX Acceleration**: Intel AMX and AVX512/AVX2 optimized kernels for INT4/INT8 quantized inference
|
||||
- **MoE Optimization**: Efficient Mixture-of-Experts inference with NUMA-aware memory management
|
||||
- **Quantization Support**: CPU-side INT4/INT8 quantized weights, GPU-side GPTQ support
|
||||
- **Easy Integration**: Clean Python API for SGLang and other frameworks
|
||||
|
||||
**Quick Start:**
|
||||
```bash
|
||||
cd kt-kernel
|
||||
pip install .
|
||||
```
|
||||
|
||||
**Use Cases:**
|
||||
|
||||
- CPU-GPU hybrid inference for large MoE models
|
||||
- Integration with SGLang for production serving
|
||||
- Heterogeneous expert placement (hot experts on GPU, cold experts on CPU)
|
||||
|
||||
**Performance Examples:**
|
||||
| Model | Hardware Configuration | Total Throughput | Output Throughput |
|
||||
|-------|------------------------|------------------|-------------------|
|
||||
| DeepSeek-R1-0528 (FP8) | 8×L20 GPU + Xeon Gold 6454S | 227.85 tokens/s | 87.58 tokens/s (8-way concurrency) |
|
||||
|
||||
👉 **[Full Documentation →](./kt-kernel/README.md)**
|
||||
|
||||
---
|
||||
|
||||
### 🎓 [kt-sft](./kt-sft/) - Fine-Tuning Framework
|
||||
|
||||
KTransformers × LLaMA-Factory integration for ultra-large MoE model fine-tuning.
|
||||
|
||||
|
||||
|
||||
**Key Features:**
|
||||
|
||||
- **Resource Efficient**: Fine-tune 671B DeepSeek-V3 with just **70GB GPU memory** + 1.3TB RAM
|
||||
- **LoRA Support**: Full LoRA fine-tuning with heterogeneous acceleration
|
||||
- **LLaMA-Factory Integration**: Seamless integration with popular fine-tuning framework
|
||||
- **Production Ready**: Chat, batch inference, and metrics evaluation
|
||||
|
||||
**Performance Examples:**
|
||||
|
||||
| Model | Configuration | Throughput | GPU Memory |
|
||||
|-------|--------------|------------|------------|
|
||||
| DeepSeek-V3 (671B) | LoRA + AMX | ~40 tokens/s | 70GB (multi-GPU) |
|
||||
| DeepSeek-V2-Lite (14B) | LoRA + AMX | ~530 tokens/s | 6GB |
|
||||
|
||||
**Quick Start:**
|
||||
```bash
|
||||
cd kt-sft
|
||||
# Install environment following kt-sft/README.md
|
||||
USE_KT=1 llamafactory-cli train examples/train_lora/deepseek3_lora_sft_kt.yaml
|
||||
```
|
||||
|
||||
👉 **[Full Documentation →](./kt-sft/README.md)**
|
||||
|
||||
---
|
||||
|
||||
## 🔥 Citation
|
||||
|
||||
If you use KTransformers in your research, please cite our paper:
|
||||
|
||||
```bibtex
|
||||
@inproceedings{10.1145/3731569.3764843,
|
||||
title = {KTransformers: Unleashing the Full Potential of CPU/GPU Hybrid Inference for MoE Models},
|
||||
author = {Chen, Hongtao and Xie, Weiyu and Zhang, Boxin and Tang, Jingqi and Wang, Jiahao and Dong, Jianwei and Chen, Shaoyuan and Yuan, Ziwei and Lin, Chen and Qiu, Chengyu and Zhu, Yuening and Ou, Qingliang and Liao, Jiaqi and Chen, Xianglin and Ai, Zhiyuan and Wu, Yongwei and Zhang, Mingxing},
|
||||
booktitle = {Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles},
|
||||
year = {2025}
|
||||
}
|
||||
```
|
||||
|
||||
## 👥 Contributors & Team
|
||||
|
||||
Developed and maintained by:
|
||||
- [MADSys Lab](https://madsys.cs.tsinghua.edu.cn/) @ Tsinghua University
|
||||
- [Approaching.AI](http://approaching.ai/)
|
||||
- Community contributors
|
||||
|
||||
We welcome contributions! Please feel free to submit issues and pull requests.
|
||||
|
||||
## 💬 Community & Support
|
||||
|
||||
- **GitHub Issues**: [Report bugs or request features](https://github.com/kvcache-ai/ktransformers/issues)
|
||||
- **GitHub Discussions**: [Ask questions and share ideas](https://github.com/kvcache-ai/ktransformers/discussions)
|
||||
- **WeChat Group**: See [archive/WeChatGroup.png](./archive/WeChatGroup.png)
|
||||
|
||||
## 📦 Legacy Code
|
||||
|
||||
The original integrated KTransformers framework has been archived to the [`archive/`](./archive/) directory for reference. The project now focuses on the two core modules above for better modularity and maintainability.
|
||||
|
||||
For the original documentation with full quick-start guides and examples, see:
|
||||
- [archive/README_LEGACY.md](./archive/README_LEGACY.md) (English)
|
||||
- [archive/README_ZH_LEGACY.md](./archive/README_ZH_LEGACY.md) (中文)
|
||||
|
||||
@@ -0,0 +1,217 @@
|
||||
<div align="center">
|
||||
<!-- <h1>KTransformers</h1> -->
|
||||
<p align="center">
|
||||
|
||||
<picture>
|
||||
<img alt="KTransformers" src="https://github.com/user-attachments/assets/d5a2492f-a415-4456-af99-4ab102f13f8b" width=50%>
|
||||
|
||||
</picture>
|
||||
|
||||
</p>
|
||||
<h3>A Flexible Framework for Experiencing Cutting-edge LLM Inference Optimizations</h3>
|
||||
<strong><a href="#show-cases">🌟 Show Cases</a> | <a href="#quick-start">🚀 Quick Start</a> | <a href="#tutorial">📃 Tutorial</a> | <a href="#Citation">🔥 Citation </a> | <a href="https://github.com/kvcache-ai/ktransformers/discussions">💬 Discussion </a>|<a href="#FAQ"> 🙋 FAQ</a> </strong>
|
||||
</div>
|
||||
|
||||
<h2 id="intro">🎉 Introduction</h2>
|
||||
KTransformers, pronounced as Quick Transformers, is designed to enhance your 🤗 <a href="https://github.com/huggingface/transformers">Transformers</a> experience with advanced kernel optimizations and placement/parallelism strategies.
|
||||
<br/><br/>
|
||||
KTransformers is a flexible, Python-centric framework designed with extensibility at its core.
|
||||
By implementing and injecting an optimized module with a single line of code, users gain access to a Transformers-compatible
|
||||
interface, RESTful APIs compliant with OpenAI and Ollama, and even a simplified ChatGPT-like web UI.
|
||||
<br/><br/>
|
||||
Our vision for KTransformers is to serve as a flexible platform for experimenting with innovative LLM inference optimizations. Please let us know if you need any other features.
|
||||
|
||||
<h2 id="Updates">🔥 Updates</h2>
|
||||
|
||||
* **Nov 6, 2025**: Support Kimi-K2-Thinking inference ([Tutorial](./doc/en/Kimi-K2-Thinking.md)) and fine-tune ([Tutorial](./doc/en/SFT_Installation_Guide_KimiK2.md))
|
||||
* **Nov 4, 2025**: KTransformers Fine-Tuning × LLaMA-Factory Integration. ([Tutorial](./doc/en/KTransformers-Fine-Tuning_User-Guide.md))
|
||||
* **Oct 27, 2025**: Support Ascend NPU. ([Tutorial](./doc/zh/DeepseekR1_V3_tutorial_zh_for_Ascend_NPU.md))
|
||||
* **Oct 10, 2025**: Integrating into SGLang. ([Roadmap](https://github.com/sgl-project/sglang/issues/11425))
|
||||
* **Sept 11, 2025**: Support Qwen3-Next. ([Tutorial](./doc/en/Qwen3-Next.md))
|
||||
* **Sept 05, 2025**: Support Kimi-K2-0905. ([Tutorial](./doc/en/Kimi-K2.md))
|
||||
* **July 26, 2025**: Support SmallThinker and GLM4-MoE. ([Tutorial](./doc/en/SmallThinker_and_Glm4moe.md))
|
||||
* **July 11, 2025**: Support Kimi-K2. ([Tutorial](./doc/en/Kimi-K2.md))
|
||||
* **June 30, 2025**: Support 3-layer (GPU-CPU-Disk) [prefix cache](./doc/en/prefix_cache.md) reuse.
|
||||
* **May 14, 2025**: Support Intel Arc GPU ([Tutorial](./doc/en/xpu.md)).
|
||||
* **Apr 29, 2025**: Support AMX-Int8、 AMX-BF16 and Qwen3MoE ([Tutorial](./doc/en/AMX.md))
|
||||
|
||||
https://github.com/user-attachments/assets/fafe8aec-4e22-49a8-8553-59fb5c6b00a2
|
||||
|
||||
* **Apr 9, 2025**: Experimental support for LLaMA 4 models ([Tutorial](./doc/en/llama4.md)).
|
||||
* **Apr 2, 2025**: Support Multi-concurrency. ([Tutorial](./doc/en/balance-serve.md)).
|
||||
|
||||
https://github.com/user-attachments/assets/faa3bda2-928b-45a7-b44f-21e12ec84b8a
|
||||
|
||||
* **Mar 15, 2025**: Support ROCm on AMD GPU ([Tutorial](./doc/en/ROCm.md)).
|
||||
* **Mar 5, 2025**: Support unsloth 1.58/2.51 bits weights and [IQ1_S/FP8 hybrid](./doc/en/fp8_kernel.md) weights. Support 139K [Longer Context](./doc/en/DeepseekR1_V3_tutorial.md#v022--v023-longer-context--fp8-kernel) for DeepSeek-V3 and R1 in 24GB VRAM.
|
||||
* **Feb 25, 2025**: Support [FP8 GPU kernel](./doc/en/fp8_kernel.md) for DeepSeek-V3 and R1; [Longer Context](./doc/en/DeepseekR1_V3_tutorial.md#v022-longer-context).
|
||||
* **Feb 15, 2025**: Longer Context (from 4K to 8K for 24GB VRAM) & Slightly Faster Speed （+15%, up to 16 Tokens/s), update [docs](./doc/en/DeepseekR1_V3_tutorial.md) and [online books](https://kvcache-ai.github.io/ktransformers/).
|
||||
* **Feb 10, 2025**: Support Deepseek-R1 and V3 on single (24GB VRAM)/multi gpu and 382G DRAM, up to 3~28x speedup. For detailed show case and reproduction tutorial, see [here](./doc/en/DeepseekR1_V3_tutorial.md).
|
||||
* **Aug 28, 2024**: Decrease DeepseekV2's required VRAM from 21G to 11G.
|
||||
* **Aug 15, 2024**: Update detailed [tutorial](doc/en/injection_tutorial.md) for injection and multi-GPU.
|
||||
* **Aug 14, 2024**: Support llamfile as linear backend.
|
||||
* **Aug 12, 2024**: Support multiple GPU; Support new model: mixtral 8\*7B and 8\*22B; Support q2k, q3k, q5k dequant on gpu.
|
||||
* **Aug 9, 2024**: Support windows native.
|
||||
|
||||
<!-- * **Aug 28, 2024**: Support 1M context under the InternLM2.5-7B-Chat-1M model, utilizing 24GB of VRAM and 150GB of DRAM. The detailed tutorial is [here](./doc/en/long_context_tutorial.md). -->
|
||||
|
||||
<h2 id="show-cases">🌟 Show Cases</h2>
|
||||
|
||||
<div>
|
||||
<h3>GPT-4/o1-level Local VSCode Copilot on a Desktop with only 24GB VRAM</h3>
|
||||
</div>
|
||||
|
||||
https://github.com/user-attachments/assets/ebd70bfa-b2c1-4abb-ae3b-296ed38aa285
|
||||
|
||||
</p>
|
||||
|
||||
- **[NEW!!!] Local 671B DeepSeek-Coder-V3/R1:** Running its Q4_K_M version using only 14GB VRAM and 382GB DRAM([Tutorial](./doc/en/DeepseekR1_V3_tutorial.md)).
|
||||
|
||||
- Prefill Speed (tokens/s):
|
||||
- KTransformers: 54.21 (32 cores) → 74.362 (dual-socket, 2×32 cores) → 255.26 (optimized AMX-based MoE kernel, V0.3 only) → 286.55 (selectively using 6 experts, V0.3 only)
|
||||
- Compared to 10.31 tokens/s in llama.cpp with 2×32 cores, achieving up to **27.79× speedup**.
|
||||
- Decode Speed (tokens/s):
|
||||
- KTransformers: 8.73 (32 cores) → 11.26 (dual-socket, 2×32 cores) → 13.69 (selectively using 6 experts, V0.3 only)
|
||||
- Compared to 4.51 tokens/s in llama.cpp with 2×32 cores, achieving up to **3.03× speedup**.
|
||||
- Upcoming Open Source Release:
|
||||
- AMX optimizations and selective expert activation will be open-sourced in V0.3.
|
||||
- Currently available only in preview binary distribution, which can be downloaded [here](./doc/en/DeepseekR1_V3_tutorial.md).
|
||||
- **Local 236B DeepSeek-Coder-V2:** Running its Q4_K_M version using only 21GB VRAM and 136GB DRAM, attainable on a local desktop machine, which scores even better than GPT4-0613 in [BigCodeBench](https://huggingface.co/blog/leaderboard-bigcodebench).
|
||||
|
||||
<p align="center">
|
||||
<picture>
|
||||
<img alt="DeepSeek-Coder-V2 Score" src="https://github.com/user-attachments/assets/d052924e-8631-44de-aad2-97c54b965693" width=100%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
- **Faster Speed:** Achieving 126 tokens/s for 2K prompt prefill and 13.6 tokens/s for generation through MoE offloading and injecting advanced kernels from [Llamafile](https://github.com/Mozilla-Ocho/llamafile/tree/main) and [Marlin](https://github.com/IST-DASLab/marlin).
|
||||
- **VSCode Integration:** Wrapped into an OpenAI and Ollama compatible API for seamless integration as a backend for [Tabby](https://github.com/TabbyML/tabby) and various other frontends.
|
||||
|
||||
<p align="center">
|
||||
|
||||
https://github.com/user-attachments/assets/4c6a8a38-05aa-497d-8eb1-3a5b3918429c
|
||||
|
||||
</p>
|
||||
|
||||
<!-- <h3>1M Context Local Inference on a Desktop with Only 24GB VRAM</h3>
|
||||
<p align="center">
|
||||
|
||||
https://github.com/user-attachments/assets/a865e5e4-bca3-401e-94b8-af3c080e6c12
|
||||
|
||||
* **1M Context InternLM 2.5 7B**: Operates at full bf16 precision, utilizing 24GB VRAM and 150GB DRAM, which is feasible on a local desktop setup. It achieves a 92.88% success rate on the 1M "Needle In a Haystack" test and 100% on the 128K NIAH test.
|
||||
|
||||
<p align="center">
|
||||
<picture>
|
||||
<img alt="Single Needle Retrieval 128K" src="./doc/assets/needle_128K.png" width=100%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<picture>
|
||||
<img alt="Single Needle Retrieval 1000K" src="./doc/assets/needle_1M.png" width=100%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
* **Enhanced Speed**: Reaches 16.91 tokens/s for generation with a 1M context using sparse attention, powered by llamafile kernels. This method is over 10 times faster than full attention approach of llama.cpp.
|
||||
|
||||
* **Flexible Sparse Attention Framework**: Offers a flexible block sparse attention framework for CPU offloaded decoding. Compatible with SnapKV, Quest, and InfLLm. Further information is available [here](./doc/en/long_context_introduction.md).
|
||||
-->
|
||||
|
||||
<strong>More advanced features will coming soon, so stay tuned!</strong>
|
||||
|
||||
<h2 id="quick-start">🚀 Quick Start</h2>
|
||||
|
||||
Getting started with KTransformers is simple! Follow the steps below to set up and start using it.
|
||||
|
||||
we have already supported vendors:
|
||||
|
||||
- Metax
|
||||
- Sanechips (ZhuFeng V1.0)
|
||||
- Intel
|
||||
- Ascend
|
||||
- Kunpeng
|
||||
- AMD
|
||||
|
||||
### 📥 Installation
|
||||
|
||||
To install KTransformers, follow the official [Installation Guide](https://kvcache-ai.github.io/ktransformers/en/install.html).
|
||||
|
||||
<h2 id="tutorial">📃 Brief Injection Tutorial</h2>
|
||||
At the heart of KTransformers is a user-friendly, template-based injection framework.
|
||||
This allows researchers to easily replace original torch modules with optimized variants. It also simplifies the process of combining multiple optimizations, allowing the exploration of their synergistic effects.
|
||||
|
||||
</br>
|
||||
<p align="center">
|
||||
<picture>
|
||||
<img alt="Inject-Struction" src="https://github.com/user-attachments/assets/6b4c1e54-9f6d-45c5-a3fc-8fa45e7d257e" width=65%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
Given that vLLM already serves as a great framework for large-scale deployment optimizations, KTransformers is particularly focused on local deployments that are constrained by limited resources. We pay special attention to heterogeneous computing opportunities, such as GPU/CPU offloading of quantized models. For example, we support the efficient <a herf="https://github.com/Mozilla-Ocho/llamafile/tree/main">Llamafile</a> and <a herf="https://github.com/IST-DASLab/marlin">Marlin</a> kernels for CPU and GPU, respectively. More details can be found <a herf="doc/en/operators/llamafile.md">here</a>.
|
||||
|
||||
<h3>Example Usage</h3>
|
||||
To utilize the provided kernels, users only need to create a YAML-based injection template and add the call to `optimize_and_load_gguf` before using the Transformers model.
|
||||
|
||||
```python
|
||||
with torch.device("meta"):
|
||||
model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
|
||||
optimize_and_load_gguf(model, optimize_config_path, gguf_path, config)
|
||||
...
|
||||
generated = prefill_and_generate(model, tokenizer, input_tensor.cuda(), max_new_tokens=1000)
|
||||
```
|
||||
|
||||
In this example, the AutoModel is first initialized on the meta device to avoid occupying any memory resources. Then, `optimize_and_load_gguf` iterates through all sub-modules of the model, matches rules specified in your YAML rule file, and replaces them with advanced modules as specified.
|
||||
|
||||
After injection, the original `generate` interface is available, but we also provide a compatible `prefill_and_generate` method, which enables further optimizations like CUDAGraph to improve generation speed.
|
||||
|
||||
<h3>How to custom your model</h3>
|
||||
|
||||
A detailed tutorial of the injection and multi-GPU using DeepSeek-V2 as an example is given [here](doc/en/injection_tutorial.md).
|
||||
|
||||
Below is an example of a YAML template for replacing all original Linear modules with Marlin, an advanced 4-bit quantization kernel.
|
||||
|
||||
```yaml
|
||||
- match:
|
||||
name: "^model\\.layers\\..*$" # regular expression
|
||||
class: torch.nn.Linear # only match modules matching name and class simultaneously
|
||||
replace:
|
||||
class: ktransformers.operators.linear.KTransformerLinear # optimized Kernel on quantized data types
|
||||
device: "cpu" # which devices to load this module when initializing
|
||||
kwargs:
|
||||
generate_device: "cuda"
|
||||
generate_linear_type: "QuantizedLinearMarlin"
|
||||
```
|
||||
|
||||
Each rule in the YAML file has two parts: `match` and `replace`. The `match` part specifies which module should be replaced, and the `replace` part specifies the module to be injected into the model along with the initialization keywords.
|
||||
|
||||
You can find example rule templates for optimizing DeepSeek-V2 and Qwen2-57B-A14, two SOTA MoE models, in the [ktransformers/optimize/optimize_rules](ktransformers/optimize/optimize_rules) directory. These templates are used to power the `local_chat.py` demo.
|
||||
|
||||
If you are interested in our design principles and the implementation of the injection framework, please refer to the [design document](doc/en/deepseek-v2-injection.md).
|
||||
|
||||
<h2 id="Citation">🔥 Citation</h2>
|
||||
|
||||
If you use KTransformers for your research, please cite our [paper](https://madsys.cs.tsinghua.edu.cn/publication/ktransformers-unleashing-the-full-potential-of-cpu/gpu-hybrid-inference-for-moe-models/):
|
||||
|
||||
```
|
||||
@inproceedings{10.1145/3731569.3764843,
|
||||
title = {KTransformers: Unleashing the Full Potential of CPU/GPU Hybrid Inference for MoE Models},
|
||||
author = {Chen, Hongtao and Xie, Weiyu and Zhang, Boxin and Tang, Jingqi and Wang, Jiahao and Dong, Jianwei and Chen, Shaoyuan and Yuan, Ziwei and Lin, Chen and Qiu, Chengyu and Zhu, Yuening and Ou, Qingliang and Liao, Jiaqi and Chen, Xianglin and Ai, Zhiyuan and Wu, Yongwei and Zhang, Mingxing},
|
||||
booktitle = {Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles},
|
||||
year = {2025}
|
||||
}
|
||||
```
|
||||
|
||||
<h2 id="ack">Acknowledgment and Contributors</h2>
|
||||
|
||||
The development of KTransformers is based on the flexible and versatile framework provided by Transformers. We also benefit from advanced kernels such as GGUF/GGML, Llamafile, Marlin, sglang and flashinfer. We are planning to contribute back to the community by upstreaming our modifications.
|
||||
|
||||
KTransformers is actively maintained and developed by contributors from the <a href="https://madsys.cs.tsinghua.edu.cn/">MADSys group</a> at Tsinghua University and members from <a href="http://approaching.ai/">Approaching.AI</a>. We welcome new contributors to join us in making KTransformers faster and easier to use.
|
||||
|
||||
<h2 id="ack">Discussion</h2>
|
||||
|
||||
If you have any questions, feel free to open an issue. Alternatively, you can join our WeChat group for further discussion. QR Code: [WeChat Group](WeChatGroup.png)
|
||||
|
||||
<h2 id="FAQ">🙋 FAQ</h2>
|
||||
|
||||
Some common questions are answered in the [FAQ](doc/en/FAQ.md).
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
<div align="center">
|
||||
<p align="center">
|
||||
<picture>
|
||||
<img alt="KTransformers" src="https://github.com/user-attachments/assets/d5a2492f-a415-4456-af99-4ab102f13f8b" width=50%>
|
||||
</picture>
|
||||
</p>
|
||||
<h3>高性能 CPU-GPU 异构大语言模型推理</h3>
|
||||
</div>
|
||||
|
||||
## 🎯 项目概述
|
||||
|
||||
KTransformers 是一个专注于大语言模型高效推理和微调的研究项目，通过 CPU-GPU 异构计算实现资源受限环境下的模型部署。项目已演进为**两个核心模块**：[kt-kernel](./kt-kernel/) 和 [kt-sft](./kt-sft/)。
|
||||
|
||||
## 🔥 更新
|
||||
|
||||
* **2025年11月6日**：支持 Kimi-K2-Thinking 推理和微调
|
||||
* **2025年11月4日**：KTransformers 微调 × LLaMA-Factory 集成
|
||||
* **2025年10月27日**：支持 Ascend NPU
|
||||
* **2025年10月10日**：集成到 SGLang ([路线图](https://github.com/sgl-project/sglang/issues/11425), [博客](https://lmsys.org/blog/2025-10-22-KTransformers/))
|
||||
* **2025年9月11日**：支持 Qwen3-Next
|
||||
* **2025年9月5日**：支持 Kimi-K2-0905
|
||||
* **2025年7月26日**：支持 SmallThinker 和 GLM4-MoE
|
||||
* **2025年6月30日**：支持 3层（GPU-CPU-磁盘）前缀缓存复用
|
||||
* **2025年5月14日**：支持 Intel Arc GPU
|
||||
* **2025年4月29日**：支持 AMX-Int8、AMX-BF16 和 Qwen3MoE
|
||||
* **2025年4月9日**：实验性支持 LLaMA 4 模型
|
||||
* **2025年4月2日**：支持多并发
|
||||
* **2025年3月15日**：支持 AMD GPU 的 ROCm
|
||||
* **2025年3月5日**：支持 unsloth 1.58/2.51 bits 权重和 IQ1_S/FP8 混合权重；DeepSeek-V3/R1 支持 139K 长上下文
|
||||
* **2025年2月25日**：支持 DeepSeek-V3 和 R1 的 FP8 GPU 内核
|
||||
* **2025年2月10日**：支持 Deepseek-R1 和 V3，速度提升最高达 3~28 倍
|
||||
|
||||
---
|
||||
|
||||
## 📦 核心模块
|
||||
|
||||
### 🚀 [kt-kernel](./kt-kernel/) - 高性能推理内核
|
||||
|
||||
面向异构 LLM 推理的 CPU 优化内核操作库。
|
||||
|
||||
|
||||
|
||||
**核心特性：**
|
||||
- **AMX/AVX 加速**：Intel AMX 和 AVX512/AVX2 优化内核，支持 INT4/INT8 量化推理
|
||||
- **MoE 优化**：高效的专家混合推理，支持 NUMA 感知内存管理
|
||||
- **量化支持**：CPU 端 INT4/INT8 量化权重，GPU 端 GPTQ 支持
|
||||
- **易于集成**：简洁的 Python API，可集成到 SGLang 等框架
|
||||
|
||||
**快速开始：**
|
||||
```bash
|
||||
cd kt-kernel
|
||||
pip install .
|
||||
```
|
||||
|
||||
**应用场景：**
|
||||
- 大型 MoE 模型的 CPU-GPU 混合推理
|
||||
- 与 SGLang 集成用于生产服务
|
||||
- 异构专家放置（热门专家在 GPU，冷门专家在 CPU）
|
||||
|
||||
**性能示例：**
|
||||
| 模型 | 硬件配置 | 总吞吐量 | 输出吞吐量 |
|
||||
|------|---------|---------|-----------|
|
||||
| DeepSeek-R1-0528 (FP8) | 8×L20 GPU + Xeon Gold 6454S | 227.85 tokens/s | 87.58 tokens/s（8路并发）|
|
||||
|
||||
👉 **[完整文档 →](./kt-kernel/README.md)**
|
||||
|
||||
---
|
||||
|
||||
### 🎓 [kt-sft](./kt-sft/) - 微调框架
|
||||
|
||||
KTransformers × LLaMA-Factory 集成，支持超大 MoE 模型微调。
|
||||
|
||||
|
||||
|
||||
**核心特性：**
|
||||
- **资源高效**：仅需 **70GB 显存** + 1.3TB 内存即可微调 671B DeepSeek-V3
|
||||
- **LoRA 支持**：完整的 LoRA 微调与异构加速
|
||||
- **LLaMA-Factory 集成**：与流行微调框架无缝集成
|
||||
- **生产就绪**：支持对话、批量推理和指标评估
|
||||
|
||||
**性能示例：**
|
||||
| 模型 | 配置 | 吞吐量 | GPU 显存 |
|
||||
|------|------|--------|----------|
|
||||
| DeepSeek-V3 (671B) | LoRA + AMX | ~40 tokens/s | 70GB (多卡) |
|
||||
| DeepSeek-V2-Lite (14B) | LoRA + AMX | ~530 tokens/s | 6GB |
|
||||
|
||||
**快速开始：**
|
||||
```bash
|
||||
cd kt-sft
|
||||
# 按照 kt-sft/README.md 安装环境
|
||||
USE_KT=1 llamafactory-cli train examples/train_lora/deepseek3_lora_sft_kt.yaml
|
||||
```
|
||||
|
||||
👉 **[完整文档 →](./kt-sft/README.md)**
|
||||
|
||||
---
|
||||
|
||||
## 🔥 引用
|
||||
|
||||
如果您在研究中使用了 KTransformers，请引用我们的论文：
|
||||
|
||||
```bibtex
|
||||
@inproceedings{10.1145/3731569.3764843,
|
||||
title = {KTransformers: Unleashing the Full Potential of CPU/GPU Hybrid Inference for MoE Models},
|
||||
author = {Chen, Hongtao and Xie, Weiyu and Zhang, Boxin and Tang, Jingqi and Wang, Jiahao and Dong, Jianwei and Chen, Shaoyuan and Yuan, Ziwei and Lin, Chen and Qiu, Chengyu and Zhu, Yuening and Ou, Qingliang and Liao, Jiaqi and Chen, Xianglin and Ai, Zhiyuan and Wu, Yongwei and Zhang, Mingxing},
|
||||
booktitle = {Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles},
|
||||
year = {2025}
|
||||
}
|
||||
```
|
||||
|
||||
## 👥 贡献者与团队
|
||||
|
||||
由以下团队开发和维护：
|
||||
- 清华大学 [MADSys 实验室](https://madsys.cs.tsinghua.edu.cn/)
|
||||
- [Approaching.AI](http://approaching.ai/)
|
||||
- 社区贡献者
|
||||
|
||||
我们欢迎贡献！请随时提交 issues 和 pull requests。
|
||||
|
||||
## 💬 社区与支持
|
||||
|
||||
- **GitHub Issues**：[报告 bug 或请求功能](https://github.com/kvcache-ai/ktransformers/issues)
|
||||
- **GitHub Discussions**：[提问和分享想法](https://github.com/kvcache-ai/ktransformers/discussions)
|
||||
- **微信群**：查看 [archive/WeChatGroup.png](./archive/WeChatGroup.png)
|
||||
|
||||
## 📦 历史代码
|
||||
|
||||
原完整的 KTransformers 框架代码已归档至 [`archive/`](./archive/) 目录供参考。项目现专注于上述两个核心模块，以实现更好的模块化和可维护性。
|
||||
|
||||
关于原始完整文档（包含快速入门指南和示例），请查看：
|
||||
- [archive/README_LEGACY.md](./archive/README_LEGACY.md) (English)
|
||||
- [archive/README_ZH_LEGACY.md](./archive/README_ZH_LEGACY.md) (中文)
|
||||
@@ -0,0 +1,166 @@
|
||||
<div align="center">
|
||||
<!-- <h1>KTransformers</h1> -->
|
||||
<p align="center">
|
||||
|
||||
<picture>
|
||||
<img alt="KTransformers" src="https://github.com/user-attachments/assets/d5a2492f-a415-4456-af99-4ab102f13f8b" width=50%>
|
||||
|
||||
</picture>
|
||||
|
||||
</p>
|
||||
<h3>一个用于体验尖端 LLM 推理优化的灵活框架</h3>
|
||||
<strong><a href="#show-cases">🌟 案例展示</a> | <a href="#quick-start">🚀 快速入门</a> | <a href="#tutorial">📃 教程</a> | <a href="https://github.com/kvcache-ai/ktransformers/discussions">💬 讨论</a> | <a href="#FAQ">🙋 常见问题</a> </strong>
|
||||
</div>
|
||||
|
||||
<h2 id="intro">🎉 介绍</h2>
|
||||
KTransformers（发音为 Quick Transformers）旨在通过先进的内核优化和放置/并行策略来增强您对 🤗 [Transformers](https://github.com/huggingface/transformers) 的体验。
|
||||
<br/><br/>
|
||||
KTransformers 是一个以 Python 为中心的灵活框架，其核心是可扩展性。通过用一行代码实现并注入优化模块，用户可以获得与 Transformers 兼容的接口、符合 OpenAI 和 Ollama 的 RESTful API，甚至是一个简化的类似 ChatGPT 的 Web 界面。
|
||||
<br/><br/>
|
||||
我们对 KTransformers 的愿景是成为一个用于实验创新 LLM 推理优化的灵活平台。如果您需要任何其他功能，请告诉我们。
|
||||
|
||||
<h2 id="Updates">🔥 更新</h2>
|
||||
|
||||
* **2025 年 2 月 15 日**：为DeepSeek-V3/R1支持[FP8 GPU内核](./doc/en/fp8_kernel.md); 支持更长的上下文([教程](./doc/en/DeepseekR1_V3_tutorial.md#v022-longer-context)).
|
||||
* **2025 年 2 月 15 日**：长上下文(从4K到8K，24GB VRAM) & 稍快的速度(+15%)(最快 16 Tokens/s)，文档请参见 [这里](./doc/en/DeepseekR1_V3_tutorial.md) 和 [在线指南](https://kvcache-ai.github.io/ktransformers/) 。
|
||||
* **2025 年 2 月 10 日**：支持 Deepseek-R1 和 V3 在单个（24GB VRAM）/多 GPU 和 382G DRAM 上运行，速度提升高达 3~28 倍。详细教程请参见 [这里](./doc/en/DeepseekR1_V3_tutorial.md)。
|
||||
* **2024 年 8 月 28 日**：支持 InternLM2.5-7B-Chat-1M 模型下的 1M 上下文，使用 24GB 的 VRAM 和 150GB 的 DRAM。详细教程请参见 [这里](./doc/en/long_context_tutorial.md)。
|
||||
* **2024 年 8 月 28 日**：将 DeepseekV2 所需的 VRAM 从 21G 降低到 11G。
|
||||
* **2024 年 8 月 15 日**：更新了详细的 [教程](doc/en/injection_tutorial.md)，介绍注入和多 GPU 的使用。
|
||||
* **2024 年 8 月 14 日**：支持 llamfile 作为线性后端。
|
||||
* **2024 年 8 月 12 日**：支持多 GPU；支持新模型：mixtral 8\*7B 和 8\*22B；支持 q2k、q3k、q5k 在 GPU 上的去量化。
|
||||
* **2024 年 8 月 9 日**：支持 Windows。
|
||||
|
||||
<h2 id="show-cases">🌟 案例展示</h2>
|
||||
|
||||
<div>
|
||||
<h3>在仅 24GB VRAM 的桌面上运行 GPT-4/o1 级别的本地 VSCode Copilot</h3>
|
||||
</div>
|
||||
|
||||
https://github.com/user-attachments/assets/ebd70bfa-b2c1-4abb-ae3b-296ed38aa285
|
||||
|
||||
</p>
|
||||
|
||||
- **[NEW!!!] 本地 671B DeepSeek-Coder-V3/R1**：使用其 Q4_K_M 版本，仅需 14GB VRAM 和 382GB DRAM 即可运行（教程请参见 [这里](./doc/en/DeepseekR1_V3_tutorial.md)）。
|
||||
- 预填充速度（tokens/s）：
|
||||
- KTransformers：54.21（32 核）→ 74.362（双插槽，2×32 核）→ 255.26（优化的 AMX 基 MoE 内核，仅 V0.3）→ 286.55（选择性使用 6 个专家，仅 V0.3）
|
||||
- 与 llama.cpp 在 2×32 核下相比，达到 **27.79× 速度提升**。
|
||||
- 解码速度（tokens/s）：
|
||||
- KTransformers：8.73（32 核）→ 11.26（双插槽，2×32 核）→ 13.69（选择性使用 6 个专家，仅 V0.3）
|
||||
- 与 llama.cpp 在 2×32 核下相比，达到 **3.03× 速度提升**。
|
||||
- 即将开源发布：
|
||||
- AMX 优化和选择性专家激活将在 V0.3 中开源。
|
||||
- 目前仅在预览二进制分发中可用，可从 [这里](./doc/en/DeepseekR1_V3_tutorial.md) 下载。
|
||||
|
||||
- **本地 236B DeepSeek-Coder-V2**：使用其 Q4_K_M 版本，仅需 21GB VRAM 和 136GB DRAM 即可运行，甚至在 [BigCodeBench](https://huggingface.co/blog/leaderboard-bigcodebench) 中得分超过 GPT4-0613。
|
||||
|
||||
<p align="center">
|
||||
<picture>
|
||||
<img alt="DeepSeek-Coder-V2 Score" src="https://github.com/user-attachments/assets/d052924e-8631-44de-aad2-97c54b965693" width=100%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
- **更快的速度**：通过 MoE 卸载和注入来自 [Llamafile](https://github.com/Mozilla-Ocho/llamafile/tree/main) 和 [Marlin](https://github.com/IST-DASLab/marlin) 的高级内核，实现了 2K 提示预填充 126 tokens/s 和生成 13.6 tokens/s 的速度。
|
||||
- **VSCode 集成**：封装成符合 OpenAI 和 Ollama 的 API，可无缝集成到 [Tabby](https://github.com/TabbyML/tabby) 和其他前端的后端。
|
||||
|
||||
<p align="center">
|
||||
|
||||
https://github.com/user-attachments/assets/4c6a8a38-05aa-497d-8eb1-3a5b3918429c
|
||||
|
||||
</p>
|
||||
|
||||
<!-- <h3>在仅 24GB VRAM 的桌面上进行 1M 上下文本地推理</h3>
|
||||
<p align="center"> -->
|
||||
|
||||
<!-- https://github.com/user-attachments/assets/a865e5e4-bca3-401e-94b8-af3c080e6c12 -->
|
||||
<!--
|
||||
* **1M 上下文 InternLM 2.5 7B**：以全 bf16 精度运行，使用 24GB VRAM 和 150GB DRAM，可在本地桌面设置中实现。在 1M "针在干草堆中" 测试中达到 92.88% 的成功率，在 128K NIAH 测试中达到 100%。
|
||||
|
||||
<p align="center">
|
||||
<picture>
|
||||
<img alt="Single Needle Retrieval 128K" src="./doc/assets/needle_128K.png" width=100%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<picture>
|
||||
<img alt="Single Needle Retrieval 1000K" src="./doc/assets/needle_1M.png" width=100%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
* **增强的速度**：使用稀疏注意力，通过 llamafile 内核实现 1M 上下文生成 16.91 tokens/s 的速度。这种方法比 llama.cpp 的全注意力方法快 10 倍以上。
|
||||
|
||||
* **灵活的稀疏注意力框架**：提供了一个灵活的块稀疏注意力框架，用于 CPU 卸载解码。与 SnapKV、Quest 和 InfLLm 兼容。更多信息请参见 [这里](./doc/en/long_context_introduction.md)。 -->
|
||||
|
||||
<strong>更多高级功能即将推出，敬请期待！</strong>
|
||||
|
||||
<h2 id="quick-start">🚀 快速入门</h2>
|
||||
|
||||
|
||||
KTransformers 的入门非常简单！请参考我们的[安装指南]((https://kvcache-ai.github.io/ktransformers/))进行安装。
|
||||
|
||||
<h2 id="tutorial">📃 简要注入教程</h2>
|
||||
KTransformers 的核心是一个用户友好的、基于模板的注入框架。这使得研究人员可以轻松地将原始 torch 模块替换为优化的变体。它还简化了多种优化的组合过程，允许探索它们的协同效应。
|
||||
</br>
|
||||
<p align="center">
|
||||
<picture>
|
||||
<img alt="Inject-Struction" src="https://github.com/user-attachments/assets/6b4c1e54-9f6d-45c5-a3fc-8fa45e7d257e" width=65%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
鉴于 vLLM 已经是一个用于大规模部署优化的优秀框架，KTransformers 特别关注受资源限制的本地部署。我们特别关注异构计算时机，例如量化模型的 GPU/CPU 卸载。例如，我们支持高效的 <a herf="https://github.com/Mozilla-Ocho/llamafile/tree/main">Llamafile</a> 和<a herf="https://github.com/IST-DASLab/marlin">Marlin</a> 内核，分别用于 CPU 和 GPU。 更多详细信息可以在 <a herf="doc/en/operators/llamafile.md">这里</a>找到。
|
||||
|
||||
|
||||
<h3>示例用法</h3>
|
||||
要使用提供的内核，用户只需创建一个基于 YAML 的注入模板，并在使用 Transformers 模型之前添加对 `optimize_and_load_gguf` 的调用。
|
||||
|
||||
```python
|
||||
with torch.device("meta"):
|
||||
model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
|
||||
optimize_and_load_gguf(model, optimize_config_path, gguf_path, config)
|
||||
...
|
||||
generated = prefill_and_generate(model, tokenizer, input_tensor.cuda(), max_new_tokens=1000)
|
||||
```
|
||||
|
||||
在这个示例中，首先在 meta 设备上初始化 AutoModel，以避免占用任何内存资源。然后，`optimize_and_load_gguf` 遍历模型的所有子模块，匹配您的 YAML 规则文件中指定的规则，并将它们替换为指定的高级模块。
|
||||
|
||||
注入后，原始的 `generate` 接口仍然可用，但我们还提供了一个兼容的 `prefill_and_generate` 方法，这使得可以进一步优化，例如使用 CUDAGraph 提高生成速度。
|
||||
|
||||
<h3>如何自定义您的模型</h3>
|
||||
|
||||
一个详细的使用 DeepSeek-V2 作为示例的注入和 multi-GPU 教程在 [这里](doc/en/injection_tutorial.md)。
|
||||
|
||||
以下是一个将所有原始 Linear 模块替换为 Marlin 的 YAML 模板示例，Marlin 是一个高级的 4 位量化内核。
|
||||
|
||||
```yaml
|
||||
- match:
|
||||
name: "^model\\.layers\\..*$" # 正则表达式
|
||||
class: torch.nn.Linear # 仅匹配同时符合名称和类的模块
|
||||
replace:
|
||||
class: ktransformers.operators.linear.KTransformerLinear # 量化数据类型的优化内核
|
||||
device: "cpu" # 初始化时加载该模块的 device
|
||||
kwargs:
|
||||
generate_device: "cuda"
|
||||
generate_linear_type: "QuantizedLinearMarlin"
|
||||
```
|
||||
|
||||
YAML 文件中的每个规则都有两部分：`match` 和 `replace`。`match` 部分指定应替换的模块，`replace` 部分指定要注入到模型中的模块以及初始化关键字。
|
||||
|
||||
您可以在 [ktransformers/optimize/optimize_rules](ktransformers/optimize/optimize_rules) 目录中找到用于优化 DeepSeek-V2 和 Qwen2-57B-A14 的示例规则模板。这些模板用于为 `local_chat.py` 示例提供支持。
|
||||
|
||||
如果您对我们的设计原则和注入框架的实现感兴趣，请参考 [设计文档](doc/en/deepseek-v2-injection.md)。
|
||||
|
||||
<h2 id="ack">致谢和贡献者</h2>
|
||||
|
||||
KTransformers 的开发基于 Transformers 提供的灵活和多功能框架。我们还受益于 GGUF/GGML、Llamafile 、 Marlin、sglang和flashinfer 等高级内核。我们计划通过向上游贡献我们的修改来回馈社区。
|
||||
|
||||
KTransformers 由清华大学 <a href="https://madsys.cs.tsinghua.edu.cn/">MADSys group</a> 小组的成员以及 <a href="http://approaching.ai/">Approaching.AI</a> 的成员积极维护和开发。我们欢迎新的贡献者加入我们，使 KTransformers 更快、更易于使用。
|
||||
|
||||
|
||||
<h2 id="ack">讨论</h2>
|
||||
|
||||
如果您有任何问题，欢迎随时提出 issue。或者，您可以加入我们的微信群进行进一步讨论。二维码： [微信群](WeChatGroup.png)
|
||||
|
||||
<h2 id="FAQ">🙋 常见问题</h2>
|
||||
|
||||
一些常见问题的答案可以在 [FAQ](doc/en/FAQ.md) 中找到。
|
||||
@@ -0,0 +1,21 @@
|
||||
# Security Policy
|
||||
|
||||
## Supported Versions
|
||||
|
||||
Use this section to tell people about which versions of your project are
|
||||
currently being supported with security updates.
|
||||
|
||||
| Version | Supported |
|
||||
| ------- | ------------------ |
|
||||
| 5.1.x | :white_check_mark: |
|
||||
| 5.0.x | :x: |
|
||||
| 4.0.x | :white_check_mark: |
|
||||
| < 4.0 | :x: |
|
||||
|
||||
## Reporting a Vulnerability
|
||||
|
||||
Use this section to tell people how to report a vulnerability.
|
||||
|
||||
Tell them where to go, how often they can expect to get an update on a
|
||||
reported vulnerability, what to expect if the vulnerability is accepted or
|
||||
declined, etc.
|
||||
@@ -0,0 +1,18 @@
|
||||
[book]
|
||||
authors = ["kvcache-ai"]
|
||||
language = "zh-CN"
|
||||
title = "Ktransformers"
|
||||
src = "doc"
|
||||
|
||||
[output.html]
|
||||
git-repository-url = "https://github.com/kvcache-ai/ktransformers"
|
||||
edit-url-template = "https://github.com/kvcache-ai/ktransformers/edit/main/{path}"
|
||||
|
||||
[output.html.playground]
|
||||
editable = true
|
||||
copy-js = true
|
||||
# line-numbers = true
|
||||
|
||||
[output.html.fold]
|
||||
enable = true
|
||||
level = 0
|
||||
@@ -0,0 +1,106 @@
|
||||
option(KTRANSFORMERS_USE_NPU "ktransformers: use NPU" OFF)
|
||||
if(KTRANSFORMERS_USE_NPU)
|
||||
add_definitions(-DKTRANSFORMERS_USE_NPU=1)
|
||||
endif()
|
||||
|
||||
if(KTRANSFORMERS_USE_NPU)
|
||||
set(ASCEND_HOME_PATH "$ENV{ASCEND_HOME_PATH}")
|
||||
message(STATUS "ASCEND_HOME_PATH is ${ASCEND_HOME_PATH}")
|
||||
include_directories(${ASCEND_HOME_PATH}/include)
|
||||
|
||||
link_directories(${TORCH_INSTALL_PREFIX}/../torch.libs)
|
||||
# find torch_npu
|
||||
execute_process(
|
||||
COMMAND python -c "import torch; import torch_npu; print(torch_npu.__path__[0])"
|
||||
OUTPUT_VARIABLE TORCH_NPU_PATH
|
||||
OUTPUT_STRIP_TRAILING_WHITESPACE
|
||||
)
|
||||
message(STATUS "Found PTA at: ${TORCH_NPU_PATH}")
|
||||
find_library(PTA_LIBRARY torch_npu PATH "${TORCH_NPU_PATH}/lib")
|
||||
endif()
|
||||
|
||||
cmake_minimum_required(VERSION 3.21)
|
||||
find_program(GCC_COMPILER NAMES g++-13 g++-12 g++-11 g++ REQUIRED)
|
||||
set(CMAKE_CXX_COMPILER ${GCC_COMPILER})
|
||||
|
||||
# 显示选定的编译器
|
||||
message(STATUS "Using compiler: ${CMAKE_CXX_COMPILER}")
|
||||
|
||||
|
||||
project(balance_serve VERSION 0.1.0)
|
||||
|
||||
set(CMAKE_CXX_STANDARD 20)
|
||||
set(CMAKE_CXX_FLAGS "-Og -march=native -Wall -Wextra -g -fPIC")
|
||||
set(CMAKE_BUILD_TYPE "Debug")
|
||||
# set(CMAKE_CXX_FLAGS "-O3 -march=native -Wall -Wextra -fPIC")
|
||||
# set(CMAKE_BUILD_TYPE "Release")
|
||||
|
||||
|
||||
if(NOT DEFINED _GLIBCXX_USE_CXX11_ABI)
|
||||
find_package(Python3 REQUIRED COMPONENTS Interpreter)
|
||||
|
||||
execute_process(
|
||||
COMMAND ${Python3_EXECUTABLE} -c
|
||||
"import torch; print('1' if torch.compiled_with_cxx11_abi() else '0')"
|
||||
OUTPUT_VARIABLE ABI_FLAG
|
||||
OUTPUT_STRIP_TRAILING_WHITESPACE
|
||||
)
|
||||
|
||||
set(_GLIBCXX_USE_CXX11_ABI ${ABI_FLAG} CACHE STRING "C++11 ABI setting from PyTorch" FORCE)
|
||||
endif()
|
||||
|
||||
# 无论是否是自动检测，都传给编译器
|
||||
add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=${_GLIBCXX_USE_CXX11_ABI})
|
||||
|
||||
message(STATUS "_GLIBCXX_USE_CXX11_ABI=${_GLIBCXX_USE_CXX11_ABI}")
|
||||
|
||||
file(GLOB_RECURSE FMT_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/*.hpp" "${CMAKE_CURRENT_SOURCE_DIR}/*.h")
|
||||
|
||||
add_custom_target(
|
||||
format
|
||||
COMMAND clang-format
|
||||
-i
|
||||
-style=file
|
||||
${FMT_SOURCES}
|
||||
COMMENT "Running clang-format on all source files"
|
||||
)
|
||||
|
||||
set(BUILD_SHARED_LIBS ON)
|
||||
set(ENABLE_PUSH OFF)
|
||||
set(ENABLE_COMPRESSION OFF)
|
||||
|
||||
# set(CMAKE_BUILD_TYPE "Release")
|
||||
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
|
||||
|
||||
set(THIRD_PARTY_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../../third_party)
|
||||
set(THIRD_PARTY_BUILD_DIR ${CMAKE_CURRENT_BINARY_DIR}/third_party)
|
||||
add_subdirectory(${THIRD_PARTY_DIR}/prometheus-cpp ${THIRD_PARTY_BUILD_DIR}/prometheus-cpp EXCLUDE_FROM_ALL)
|
||||
add_subdirectory(${THIRD_PARTY_DIR}/xxHash/cmake_unofficial ${THIRD_PARTY_BUILD_DIR}/xxHash EXCLUDE_FROM_ALL)
|
||||
set_target_properties(xxhash PROPERTIES POSITION_INDEPENDENT_CODE ON)
|
||||
|
||||
# add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third_party/prometheus-cpp ${CMAKE_CURRENT_BINARY_DIR}/third_party/prometheus-cpp)
|
||||
set(SPDLOG_DIR ${THIRD_PARTY_DIR}/spdlog)
|
||||
set(FMT_DIR ${THIRD_PARTY_DIR}/fmt)
|
||||
|
||||
set(KVC2_INCLUDE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/kvc2/src)
|
||||
|
||||
include_directories(${THIRD_PARTY_DIR})
|
||||
|
||||
add_subdirectory(${THIRD_PARTY_DIR}/pybind11 ${THIRD_PARTY_BUILD_DIR}/pybind11)
|
||||
|
||||
execute_process(
|
||||
COMMAND python3 -c "import torch; print(torch.__path__[0])"
|
||||
OUTPUT_VARIABLE TORCH_INSTALL_PREFIX
|
||||
OUTPUT_STRIP_TRAILING_WHITESPACE
|
||||
)
|
||||
|
||||
message(STATUS "Found PyTorch at: ${TORCH_INSTALL_PREFIX}")
|
||||
|
||||
# set(TORCH_INSTALL_PREFIX "/home/xwy/.conda/envs/kvc/lib/python3.12/site-packages/torch")
|
||||
find_library(TORCH_PYTHON_LIBRARY torch_python PATH "${TORCH_INSTALL_PREFIX}/lib")
|
||||
find_package(Torch REQUIRED PATHS "${TORCH_INSTALL_PREFIX}/share/cmake/Torch" NO_DEFAULT_PATH)
|
||||
|
||||
add_subdirectory(kvc2)
|
||||
add_subdirectory(sched)
|
||||
|
||||
# add_subdirectory(test)
|
||||
@@ -0,0 +1,44 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : Azure-Tang
|
||||
* @Date : 2024-07-25 13:38:30
|
||||
* @Version : 1.0.0
|
||||
* @LastEditors : kkk1nak0
|
||||
* @LastEditTime : 2024-08-12 03:05:04
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
|
||||
#include "gptq_marlin/ops.h"
|
||||
// Python bindings
|
||||
#include <pybind11/pybind11.h>
|
||||
#include <pybind11/stl.h>
|
||||
#include <torch/extension.h>
|
||||
#include <torch/library.h>
|
||||
#include <torch/torch.h>
|
||||
// namespace py = pybind11;
|
||||
|
||||
PYBIND11_MODULE(vLLMMarlin, m) {
|
||||
|
||||
/*m.def("dequantize_q8_0", &dequantize_q8_0, "Function to dequantize q8_0
|
||||
data.", py::arg("data"), py::arg("blk_size"), py::arg("device"));
|
||||
m.def("dequantize_q6_k", &dequantize_q6_k, "Function to dequantize q6_k
|
||||
data.", py::arg("data"), py::arg("blk_size"), py::arg("device"));
|
||||
m.def("dequantize_q5_k", &dequantize_q5_k, "Function to dequantize q5_k
|
||||
data.", py::arg("data"), py::arg("blk_size"), py::arg("device"));
|
||||
m.def("dequantize_q4_k", &dequantize_q4_k, "Function to dequantize q4_k
|
||||
data.", py::arg("data"), py::arg("blk_size"), py::arg("device"));
|
||||
m.def("dequantize_q3_k", &dequantize_q3_k, "Function to dequantize q3_k
|
||||
data.", py::arg("data"), py::arg("blk_size"), py::arg("device"));
|
||||
m.def("dequantize_q2_k", &dequantize_q2_k, "Function to dequantize q2_k
|
||||
data.", py::arg("data"), py::arg("blk_size"), py::arg("device"));
|
||||
m.def("dequantize_iq4_xs", &dequantize_iq4_xs, "Function to dequantize
|
||||
iq4_xs data.", py::arg("data"), py::arg("blk_size"), py::arg("device"));*/
|
||||
m.def("gptq_marlin_gemm", &gptq_marlin_gemm,
|
||||
"Function to perform GEMM using Marlin quantization.", py::arg("a"),
|
||||
py::arg("b_q_weight"), py::arg("b_scales"), py::arg("g_idx"),
|
||||
py::arg("perm"), py::arg("workspace"), py::arg("num_bits"), py::arg("size_m_tensor"),
|
||||
py::arg("size_m"), py::arg("size_n"), py::arg("size_k"),
|
||||
py::arg("sms"), py::arg("is_k_full"));
|
||||
m.def("gptq_marlin_repack", &gptq_marlin_repack,
|
||||
"gptq_marlin repack from GPTQ");
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,76 @@
|
||||
// Adapted from
|
||||
// https://github.com/vllm-project/vllm/tree/main/csrc/quantization/gptq_marlin
|
||||
// Copyrigth 2024 The vLLM team.
|
||||
// Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
#pragma once
|
||||
|
||||
#include <torch/all.h>
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <cuda.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <iostream>
|
||||
|
||||
namespace gptq_marlin {
|
||||
|
||||
// 8 warps are a good choice since every SM has 4 schedulers and having more
|
||||
// than 1 warp per schedule allows some more latency hiding. At the same time,
|
||||
// we want relatively few warps to have many registers per warp and small tiles.
|
||||
static constexpr int default_threads = 256;
|
||||
|
||||
static constexpr int pipe_stages =
|
||||
4; // 4 pipeline stages fit into shared memory
|
||||
|
||||
static constexpr int min_thread_n = 64;
|
||||
static constexpr int min_thread_k = 64;
|
||||
|
||||
static constexpr int tile_size = 16;
|
||||
static constexpr int max_par = 16;
|
||||
|
||||
template <typename T, int n> struct Vec {
|
||||
T elems[n];
|
||||
__device__ T &operator[](int i) { return elems[i]; }
|
||||
};
|
||||
|
||||
using I4 = Vec<int, 4>;
|
||||
|
||||
constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; }
|
||||
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
|
||||
// No support for async
|
||||
#else
|
||||
|
||||
__device__ inline void cp_async4_pred(void *smem_ptr, const void *glob_ptr,
|
||||
bool pred = true) {
|
||||
const int BYTES = 16;
|
||||
uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
|
||||
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)pred),
|
||||
"r"(smem), "l"(glob_ptr), "n"(BYTES));
|
||||
}
|
||||
|
||||
__device__ inline void cp_async4(void *smem_ptr, const void *glob_ptr) {
|
||||
const int BYTES = 16;
|
||||
uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
|
||||
asm volatile("{\n"
|
||||
" cp.async.cg.shared.global [%0], [%1], %2;\n"
|
||||
"}\n" ::"r"(smem),
|
||||
"l"(glob_ptr), "n"(BYTES));
|
||||
}
|
||||
|
||||
__device__ inline void cp_async_fence() {
|
||||
asm volatile("cp.async.commit_group;\n" ::);
|
||||
}
|
||||
|
||||
template <int n> __device__ inline void cp_async_wait() {
|
||||
asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
} // namespace gptq_marlin
|
||||
@@ -0,0 +1,77 @@
|
||||
// Adapted from
|
||||
// https://github.com/vllm-project/vllm/tree/main/csrc/quantization/gptq_marlin
|
||||
// Copyrigth 2024 The vLLM team.
|
||||
// Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
#ifndef _data_types_cuh
|
||||
#define _data_types_cuh
|
||||
#include "gptq_marlin.cuh"
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
|
||||
namespace gptq_marlin {
|
||||
|
||||
template <typename scalar_t> class ScalarType {};
|
||||
|
||||
template <> class ScalarType<half> {
|
||||
public:
|
||||
using scalar_t = half;
|
||||
using scalar_t2 = half2;
|
||||
|
||||
// Matrix fragments for tensor core instructions; their precise layout is
|
||||
// documented here:
|
||||
// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type
|
||||
using FragA = Vec<half2, 4>;
|
||||
using FragB = Vec<half2, 2>;
|
||||
using FragC = Vec<float, 4>;
|
||||
using FragS = Vec<half2, 1>;
|
||||
|
||||
static __device__ float inline num2float(const half x) {
|
||||
return __half2float(x);
|
||||
}
|
||||
|
||||
static __device__ half2 inline num2num2(const half x) {
|
||||
return __half2half2(x);
|
||||
}
|
||||
|
||||
static __device__ half2 inline nums2num2(const half x1, const half x2) {
|
||||
return __halves2half2(x1, x2);
|
||||
}
|
||||
|
||||
static __host__ __device__ half inline float2num(const float x) {
|
||||
return __float2half(x);
|
||||
}
|
||||
};
|
||||
|
||||
template <> class ScalarType<nv_bfloat16> {
|
||||
public:
|
||||
using scalar_t = nv_bfloat16;
|
||||
using scalar_t2 = nv_bfloat162;
|
||||
|
||||
using FragA = Vec<nv_bfloat162, 4>;
|
||||
using FragB = Vec<nv_bfloat162, 2>;
|
||||
using FragC = Vec<float, 4>;
|
||||
using FragS = Vec<nv_bfloat162, 1>;
|
||||
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
static __device__ float inline num2float(const nv_bfloat16 x) {
|
||||
return __bfloat162float(x);
|
||||
}
|
||||
|
||||
static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) {
|
||||
return __bfloat162bfloat162(x);
|
||||
}
|
||||
|
||||
static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1,
|
||||
const nv_bfloat16 x2) {
|
||||
return __halves2bfloat162(x1, x2);
|
||||
}
|
||||
|
||||
static __host__ __device__ nv_bfloat16 inline float2num(const float x) {
|
||||
return __float2bfloat16(x);
|
||||
}
|
||||
#endif
|
||||
};
|
||||
|
||||
} // namespace gptq_marlin
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,350 @@
|
||||
#include "gptq_marlin.cuh"
|
||||
|
||||
namespace gptq_marlin {
|
||||
|
||||
static constexpr int repack_stages = 8;
|
||||
|
||||
static constexpr int repack_threads = 256;
|
||||
|
||||
static constexpr int tile_k_size = tile_size;
|
||||
static constexpr int tile_n_size = tile_k_size * 4;
|
||||
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
|
||||
|
||||
template <int const num_threads, int const num_bits, bool const has_perm>
|
||||
__global__ void marlin_repack_kernel(
|
||||
uint32_t const* __restrict__ b_q_weight_ptr,
|
||||
uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr,
|
||||
int size_k, int size_n) {}
|
||||
|
||||
} // namespace gptq_marlin
|
||||
|
||||
torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm,
|
||||
int64_t size_k, int64_t size_n,
|
||||
int64_t num_bits) {
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(
|
||||
false, "marlin_repack_from_gptq(..) requires CUDA_ARCH >= 8.0");
|
||||
return torch::empty({1, 1});
|
||||
}
|
||||
|
||||
#else
|
||||
|
||||
template <int const num_threads, int const num_bits, bool const has_perm>
|
||||
__global__ void marlin_repack_kernel(
|
||||
uint32_t const* __restrict__ b_q_weight_ptr,
|
||||
uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr,
|
||||
int size_k, int size_n) {
|
||||
constexpr int pack_factor = 32 / num_bits;
|
||||
|
||||
int k_tiles = size_k / tile_k_size;
|
||||
int n_tiles = size_n / tile_n_size;
|
||||
int block_k_tiles = div_ceil(k_tiles, gridDim.x);
|
||||
|
||||
int start_k_tile = blockIdx.x * block_k_tiles;
|
||||
if (start_k_tile >= k_tiles) {
|
||||
return;
|
||||
}
|
||||
|
||||
int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles);
|
||||
|
||||
// Wait until the next thread tile has been loaded to shared memory.
|
||||
auto wait_for_stage = [&]() {
|
||||
// We only have `stages - 2` active fetches since we are double buffering
|
||||
// and can only issue the next fetch when it is guaranteed that the previous
|
||||
// shared memory load is fully complete (as it may otherwise be
|
||||
// overwritten).
|
||||
cp_async_wait<repack_stages - 2>();
|
||||
__syncthreads();
|
||||
};
|
||||
|
||||
extern __shared__ int4 sh[];
|
||||
|
||||
constexpr int perm_size = tile_k_size / 4;
|
||||
|
||||
int4* sh_perm_ptr = sh;
|
||||
int4* sh_pipe_ptr = sh_perm_ptr;
|
||||
if constexpr (has_perm) {
|
||||
sh_pipe_ptr += perm_size;
|
||||
}
|
||||
|
||||
constexpr int tile_ints = tile_k_size / pack_factor;
|
||||
|
||||
constexpr int stage_n_threads = tile_n_size / 4;
|
||||
constexpr int stage_k_threads = has_perm ? tile_k_size : tile_ints;
|
||||
constexpr int stage_size = stage_k_threads * stage_n_threads;
|
||||
|
||||
auto load_perm_to_shared = [&](int k_tile_id) {
|
||||
int first_k_int4 = (k_tile_id * tile_k_size) / 4;
|
||||
|
||||
int4 const* perm_int4_ptr = reinterpret_cast<int4 const*>(perm_ptr);
|
||||
|
||||
if (threadIdx.x < perm_size) {
|
||||
sh_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x];
|
||||
}
|
||||
__syncthreads();
|
||||
};
|
||||
|
||||
auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) {
|
||||
if (n_tile_id >= n_tiles) {
|
||||
cp_async_fence();
|
||||
return;
|
||||
}
|
||||
|
||||
int first_n = n_tile_id * tile_n_size;
|
||||
|
||||
int4* sh_ptr = sh_pipe_ptr + stage_size * pipe;
|
||||
|
||||
if constexpr (has_perm) {
|
||||
if (threadIdx.x < stage_size) {
|
||||
int k_id = threadIdx.x / stage_n_threads;
|
||||
int n_id = threadIdx.x % stage_n_threads;
|
||||
|
||||
uint32_t const* sh_perm_int_ptr =
|
||||
reinterpret_cast<uint32_t const*>(sh_perm_ptr);
|
||||
|
||||
int src_k = sh_perm_int_ptr[k_id];
|
||||
int src_k_packed = src_k / pack_factor;
|
||||
|
||||
cp_async4(
|
||||
&sh_ptr[k_id * stage_n_threads + n_id],
|
||||
reinterpret_cast<int4 const*>(&(
|
||||
b_q_weight_ptr[src_k_packed * size_n + first_n + (n_id * 4)])));
|
||||
}
|
||||
|
||||
} else {
|
||||
if (threadIdx.x < stage_size) {
|
||||
int k_id = threadIdx.x / stage_n_threads;
|
||||
int n_id = threadIdx.x % stage_n_threads;
|
||||
|
||||
int first_k = k_tile_id * tile_k_size;
|
||||
int first_k_packed = first_k / pack_factor;
|
||||
|
||||
cp_async4(&sh_ptr[k_id * stage_n_threads + n_id],
|
||||
reinterpret_cast<int4 const*>(
|
||||
&(b_q_weight_ptr[(first_k_packed + k_id) * size_n +
|
||||
first_n + (n_id * 4)])));
|
||||
}
|
||||
}
|
||||
|
||||
cp_async_fence();
|
||||
};
|
||||
|
||||
auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) {
|
||||
if (n_tile_id >= n_tiles) {
|
||||
return;
|
||||
}
|
||||
|
||||
int warp_id = threadIdx.x / 32;
|
||||
int th_id = threadIdx.x % 32;
|
||||
|
||||
if (warp_id >= 4) {
|
||||
return;
|
||||
}
|
||||
|
||||
int tc_col = th_id / 4;
|
||||
int tc_row = (th_id % 4) * 2;
|
||||
|
||||
constexpr int tc_offsets[4] = {0, 1, 8, 9};
|
||||
|
||||
int cur_n = warp_id * 16 + tc_col;
|
||||
|
||||
constexpr int sh_stride = 64;
|
||||
constexpr uint32_t mask = (1 << num_bits) - 1;
|
||||
|
||||
int4* sh_stage_ptr = sh_pipe_ptr + stage_size * pipe;
|
||||
uint32_t* sh_stage_int_ptr = reinterpret_cast<uint32_t*>(sh_stage_ptr);
|
||||
|
||||
uint32_t* sh_perm_int_ptr = reinterpret_cast<uint32_t*>(sh_perm_ptr);
|
||||
|
||||
uint32_t vals[8];
|
||||
|
||||
if constexpr (has_perm) {
|
||||
for (int i = 0; i < 4; i++) {
|
||||
int k_idx = tc_row + tc_offsets[i];
|
||||
|
||||
uint32_t src_k = sh_perm_int_ptr[k_idx];
|
||||
uint32_t src_k_pos = src_k % pack_factor;
|
||||
|
||||
uint32_t b1_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n];
|
||||
uint32_t b1_cur_val = (b1_val >> (src_k_pos * num_bits)) & mask;
|
||||
|
||||
uint32_t b2_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n + 8];
|
||||
uint32_t b2_cur_val = (b2_val >> (src_k_pos * num_bits)) & mask;
|
||||
|
||||
vals[i] = b1_cur_val;
|
||||
vals[4 + i] = b2_cur_val;
|
||||
}
|
||||
|
||||
} else {
|
||||
uint32_t b1_vals[tile_ints];
|
||||
uint32_t b2_vals[tile_ints];
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < tile_ints; i++) {
|
||||
b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i];
|
||||
b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i];
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
int cur_elem = tc_row + tc_offsets[i];
|
||||
int cur_int = cur_elem / pack_factor;
|
||||
int cur_pos = cur_elem % pack_factor;
|
||||
|
||||
vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask;
|
||||
vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask;
|
||||
}
|
||||
}
|
||||
|
||||
constexpr int tile_size = tile_k_size * tile_n_size / pack_factor;
|
||||
int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size;
|
||||
|
||||
// Result of:
|
||||
// https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h
|
||||
if constexpr (num_bits == 4) {
|
||||
constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7};
|
||||
|
||||
uint32_t res = 0;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; i++) {
|
||||
res |= vals[pack_idx[i]] << (i * 4);
|
||||
}
|
||||
|
||||
out_ptr[out_offset + th_id * 4 + warp_id] = res;
|
||||
|
||||
} else {
|
||||
constexpr int pack_idx[4] = {0, 2, 1, 3};
|
||||
|
||||
uint32_t res1 = 0;
|
||||
uint32_t res2 = 0;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
res1 |= vals[pack_idx[i]] << (i * 8);
|
||||
res2 |= vals[4 + pack_idx[i]] << (i * 8);
|
||||
}
|
||||
|
||||
out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1;
|
||||
out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2;
|
||||
}
|
||||
};
|
||||
|
||||
auto start_pipes = [&](int k_tile_id, int n_tile_id) {
|
||||
#pragma unroll
|
||||
for (int pipe = 0; pipe < repack_stages - 1; pipe++) {
|
||||
fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe);
|
||||
}
|
||||
|
||||
wait_for_stage();
|
||||
};
|
||||
#pragma unroll
|
||||
for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) {
|
||||
int n_tile_id = 0;
|
||||
|
||||
if constexpr (has_perm) {
|
||||
load_perm_to_shared(k_tile_id);
|
||||
}
|
||||
|
||||
start_pipes(k_tile_id, n_tile_id);
|
||||
|
||||
while (n_tile_id < n_tiles) {
|
||||
#pragma unroll
|
||||
for (int pipe = 0; pipe < repack_stages; pipe++) {
|
||||
fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id,
|
||||
n_tile_id + pipe + repack_stages - 1);
|
||||
repack_tile(pipe, k_tile_id, n_tile_id + pipe);
|
||||
wait_for_stage();
|
||||
}
|
||||
n_tile_id += repack_stages;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace gptq_marlin
|
||||
|
||||
#define CALL_IF(NUM_BITS, HAS_PERM) \
|
||||
else if (num_bits == NUM_BITS && has_perm == HAS_PERM) { \
|
||||
cudaFuncSetAttribute( \
|
||||
gptq_marlin::marlin_repack_kernel<gptq_marlin::repack_threads, \
|
||||
NUM_BITS, HAS_PERM>, \
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \
|
||||
gptq_marlin::marlin_repack_kernel<gptq_marlin::repack_threads, NUM_BITS, \
|
||||
HAS_PERM> \
|
||||
<<<blocks, gptq_marlin::repack_threads, max_shared_mem, stream>>>( \
|
||||
b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \
|
||||
}
|
||||
|
||||
torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm,
|
||||
int64_t size_k, int64_t size_n,
|
||||
int64_t num_bits) {
|
||||
// Verify compatibility with marlin tile of 16x64
|
||||
TORCH_CHECK(size_k % gptq_marlin::tile_k_size == 0, "size_k = ", size_k,
|
||||
" is not divisible by tile_k_size = ", gptq_marlin::tile_k_size);
|
||||
TORCH_CHECK(size_n % gptq_marlin::tile_n_size == 0, "size_n = ", size_n,
|
||||
" is not divisible by tile_n_size = ", gptq_marlin::tile_n_size);
|
||||
|
||||
TORCH_CHECK(num_bits == 4 || num_bits == 8,
|
||||
"num_bits must be 4 or 8. Got = ", num_bits);
|
||||
int const pack_factor = 32 / num_bits;
|
||||
|
||||
// Verify B
|
||||
TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0),
|
||||
"Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0),
|
||||
", size_k = ", size_k, ", pack_factor = ", pack_factor);
|
||||
TORCH_CHECK(b_q_weight.size(1) == size_n,
|
||||
"b_q_weight.size(1) = ", b_q_weight.size(1),
|
||||
" is not size_n = ", size_n);
|
||||
|
||||
// Verify device and strides
|
||||
TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU");
|
||||
TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
|
||||
TORCH_CHECK(b_q_weight.dtype() == at::kInt, "b_q_weight type is not kInt");
|
||||
|
||||
TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU");
|
||||
TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
|
||||
TORCH_CHECK(perm.dtype() == at::kInt, "perm type is not at::kInt");
|
||||
|
||||
// Alloc buffers
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight));
|
||||
auto options = torch::TensorOptions()
|
||||
.dtype(b_q_weight.dtype())
|
||||
.device(b_q_weight.device());
|
||||
torch::Tensor out =
|
||||
torch::empty({size_k / gptq_marlin::tile_size,
|
||||
size_n * gptq_marlin::tile_size / pack_factor},
|
||||
options);
|
||||
|
||||
// Detect if there is act_order
|
||||
bool has_perm = perm.size(0) != 0;
|
||||
|
||||
// Get ptrs
|
||||
uint32_t const* b_q_weight_ptr =
|
||||
reinterpret_cast<uint32_t const*>(b_q_weight.data_ptr());
|
||||
uint32_t const* perm_ptr = reinterpret_cast<uint32_t const*>(perm.data_ptr());
|
||||
uint32_t* out_ptr = reinterpret_cast<uint32_t*>(out.data_ptr());
|
||||
|
||||
// Get dev info
|
||||
int dev = b_q_weight.get_device();
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev);
|
||||
int blocks;
|
||||
cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev);
|
||||
|
||||
int max_shared_mem = 0;
|
||||
cudaDeviceGetAttribute(&max_shared_mem,
|
||||
cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
|
||||
TORCH_CHECK(max_shared_mem > 0);
|
||||
|
||||
if (false) {
|
||||
}
|
||||
CALL_IF(4, false)
|
||||
CALL_IF(4, true)
|
||||
CALL_IF(8, false)
|
||||
CALL_IF(8, true)
|
||||
else {
|
||||
TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
|
||||
", has_perm = ", has_perm);
|
||||
}
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,24 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : Azure
|
||||
* @Date : 2024-07-22 09:27:55
|
||||
* @Version : 1.0.0
|
||||
* @LastEditors : Azure
|
||||
* @LastEditTime : 2024-07-26 08:35:00
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
#pragma once
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <torch/library.h>
|
||||
#include <torch/torch.h>
|
||||
|
||||
torch::Tensor gptq_marlin_gemm(torch::Tensor &a, torch::Tensor &b_q_weight,
|
||||
torch::Tensor &b_scales, torch::Tensor &g_idx,
|
||||
torch::Tensor &perm, torch::Tensor &workspace,
|
||||
int64_t num_bits, torch::Tensor size_m_tensor, int64_t size_m, int64_t size_n,
|
||||
int64_t size_k, int sms, bool is_k_full);
|
||||
|
||||
torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor&perm,
|
||||
int64_t size_k, int64_t size_n,
|
||||
int64_t num_bits);
|
||||
@@ -0,0 +1,25 @@
|
||||
from setuptools import setup, Extension
|
||||
from torch.utils import cpp_extension
|
||||
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
||||
setup(
|
||||
name='vLLMMarlin',
|
||||
ext_modules=[
|
||||
CUDAExtension(
|
||||
'vLLMMarlin', [
|
||||
#'custom_gguf/dequant.cu',
|
||||
'binding.cpp',
|
||||
'gptq_marlin/gptq_marlin.cu',
|
||||
'gptq_marlin/gptq_marlin_repack.cu',
|
||||
],
|
||||
extra_compile_args={
|
||||
'cxx': ['-O3'],
|
||||
'nvcc': [
|
||||
'-O3',
|
||||
'--use_fast_math',
|
||||
'-Xcompiler', '-fPIC',
|
||||
]
|
||||
},
|
||||
)
|
||||
],
|
||||
cmdclass={'build_ext': BuildExtension}
|
||||
)
|
||||
@@ -0,0 +1,335 @@
|
||||
import csv
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import vLLMMarlin
|
||||
torch.set_grad_enabled(False)
|
||||
from utils.marlin_utils import (
|
||||
MarlinWorkspace,
|
||||
marlin_quantize,
|
||||
GPTQ_MARLIN_MIN_THREAD_N,
|
||||
GPTQ_MARLIN_MIN_THREAD_K,
|
||||
GPTQ_MARLIN_MAX_PARALLEL,
|
||||
)
|
||||
|
||||
def setup_seed(seed):
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
setup_seed(20241223)
|
||||
|
||||
torch.set_grad_enabled(False)
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
global_dtype=torch.bfloat16
|
||||
global_device=torch.device("cuda",0)
|
||||
global_num_cases:int=int(50)
|
||||
torch.cuda.set_device(0)
|
||||
torch.backends.cudnn.enabled =True
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
max_batch_size = 512
|
||||
max_tp = 8
|
||||
L2_size = 73728 * 1024
|
||||
|
||||
def get_usable_mem():
|
||||
properties = torch.cuda.get_device_properties(global_device)
|
||||
#print(f"Total memory: {properties.total_memory / (1024 ** 3):.2f} GB")
|
||||
allocated_memory = torch.cuda.memory_allocated(global_device)
|
||||
#print(f"Currently allocated memory: {allocated_memory / (1024 ** 2):.2f} MB")
|
||||
reserved_memory = torch.cuda.memory_reserved(global_device)
|
||||
#print(f"Currently reserved memory: {reserved_memory / (1024 ** 2):.2f} MB")
|
||||
return properties.total_memory - 512 * 1024 ** 2 - allocated_memory# - reserved_memory
|
||||
|
||||
def exp_range(start, stop, step = 2):
|
||||
now = start
|
||||
while now <= stop:
|
||||
yield now
|
||||
now *= step
|
||||
|
||||
def timing(func, iters, epochs=100):
|
||||
#warmup
|
||||
for idx in range(iters):
|
||||
func(idx)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
cuda_graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(cuda_graph):
|
||||
for idx in range(iters):
|
||||
func(idx)
|
||||
|
||||
for _ in range(2000):
|
||||
cuda_graph.replay()
|
||||
|
||||
start_event = torch.cuda.Event(enable_timing=True)
|
||||
end_event = torch.cuda.Event(enable_timing=True)
|
||||
stream = torch.cuda.Stream()
|
||||
torch.cuda.synchronize()
|
||||
#with torch.cuda.stream(stream):
|
||||
start_event.record()
|
||||
for _ in range(10):
|
||||
cuda_graph.replay()
|
||||
end_event.record()
|
||||
torch.cuda.synchronize()
|
||||
elapsed_time_ms0 = start_event.elapsed_time(end_event)
|
||||
|
||||
start_event = torch.cuda.Event(enable_timing=True)
|
||||
end_event = torch.cuda.Event(enable_timing=True)
|
||||
torch.cuda.synchronize()
|
||||
#with torch.cuda.stream(stream):
|
||||
start_event.record()
|
||||
for _ in range(epochs+10):
|
||||
cuda_graph.replay()
|
||||
end_event.record()
|
||||
torch.cuda.synchronize()
|
||||
elapsed_time_ms = start_event.elapsed_time(end_event) - elapsed_time_ms0
|
||||
|
||||
#print(elapsed_time_ms0, elapsed_time_ms)
|
||||
return elapsed_time_ms/iters/epochs
|
||||
|
||||
class LinearMarlin(nn.Linear):
|
||||
marlin_q_w: torch.Tensor
|
||||
marlin_s: torch.Tensor
|
||||
g_idx: torch.Tensor
|
||||
sort_indices: torch.Tensor
|
||||
has_bias: bool
|
||||
def __init__(
|
||||
self,
|
||||
in_features,
|
||||
out_features,
|
||||
bias = False,
|
||||
device: str = "cuda",
|
||||
num_bits: int = 4, # 4-bit/8-bit is supported
|
||||
group_size: int = 64, # -1, 32, 64, 128
|
||||
act_order: bool = False,
|
||||
is_k_full=True,
|
||||
sms = -1, # sms in GPU
|
||||
**kwargs,
|
||||
):
|
||||
self.padding = False
|
||||
assert device.lower() != "cpu", "Marlin quantized linear only supports GPU device"
|
||||
if in_features%GPTQ_MARLIN_MIN_THREAD_K!=0 or out_features%GPTQ_MARLIN_MIN_THREAD_K!=0:
|
||||
#print(f"warning!, in_features={in_features} or out_features={out_features} is undivisible by GPTQ_MARLIN_MIN_THREAD_K={GPTQ_MARLIN_MIN_THREAD_K} and GPTQ_MARLIN_MIN_THREAD_N={GPTQ_MARLIN_MIN_THREAD_N}, padding")
|
||||
self.padding = True
|
||||
self.orin_in_features = in_features
|
||||
self.orin_out_features = out_features
|
||||
in_features = (in_features+GPTQ_MARLIN_MIN_THREAD_K-1)//GPTQ_MARLIN_MIN_THREAD_K*GPTQ_MARLIN_MIN_THREAD_K
|
||||
out_features = (out_features+GPTQ_MARLIN_MIN_THREAD_N-1)//GPTQ_MARLIN_MIN_THREAD_N*GPTQ_MARLIN_MIN_THREAD_N
|
||||
#print(f"After padding: in_features={in_features}, out_features={out_features}")
|
||||
|
||||
|
||||
super().__init__(in_features, out_features, bias, device)
|
||||
self.has_bias = bias
|
||||
self.device = device
|
||||
self.num_bits = num_bits
|
||||
self.group_size = group_size
|
||||
self.act_order = act_order
|
||||
# TODO: optimize every shape GEMM
|
||||
|
||||
blocks_k, blocks_n = in_features//128, out_features//128
|
||||
|
||||
self.sms = sms
|
||||
|
||||
self.is_k_full = is_k_full
|
||||
|
||||
self.weight.requires_grad = False
|
||||
self.weight.t_()
|
||||
# Pack Marlin linear
|
||||
#w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize(
|
||||
# self.weight, self.num_bits, self.group_size, self.act_order
|
||||
#)
|
||||
marlin_q_w = torch.randint(int(-1e9), int(1e9), (in_features//16, out_features*2), device=device, dtype=torch.int)
|
||||
marlin_s = torch.randn((in_features//64, out_features), device=device)
|
||||
self.workspace = MarlinWorkspace(
|
||||
self.out_features, GPTQ_MARLIN_MIN_THREAD_N, GPTQ_MARLIN_MAX_PARALLEL, self.device
|
||||
)
|
||||
self.marlin_q_w = marlin_q_w
|
||||
self.marlin_s = marlin_s
|
||||
self.g_idx = torch.empty((0), dtype=torch.int32, device=self.device)
|
||||
self.sort_indices = torch.empty((0), dtype=torch.int32, device=self.device)
|
||||
self.k = self.weight.shape[0]
|
||||
self.n = self.weight.shape[1]
|
||||
self.weight = None
|
||||
"""
|
||||
print(in_features, out_features)
|
||||
print(marlin_q_w.shape)
|
||||
print(marlin_q_w.dtype)
|
||||
print(marlin_s.shape)
|
||||
print(marlin_s.dtype)
|
||||
print(self.workspace.scratch.shape)
|
||||
print(self.workspace.scratch.dtype)
|
||||
print(self.g_idx.shape)
|
||||
print(self.g_idx.dtype)
|
||||
print(self.sort_indices.shape)
|
||||
print(self.sort_indices.dtype)
|
||||
#print(w_ref.shape)
|
||||
#print(w_ref.dtype)
|
||||
"""
|
||||
#w_ref = None
|
||||
|
||||
def forward(self, x: torch.Tensor, bsz_tensor: torch.Tensor) -> torch.Tensor:
|
||||
# Only support input x as BF16 and FP16
|
||||
x = x.to(self.device)
|
||||
orig_shape = list(x.shape)
|
||||
orig_dtype = x.dtype
|
||||
x = x.reshape(-1, x.shape[-1])
|
||||
if self.padding:
|
||||
padding_input=torch.empty(x.shape[0], self.in_features, device=x.device, dtype=x.dtype)
|
||||
padding_input[:,:self.orin_in_features] = x
|
||||
x = padding_input
|
||||
marlin_s = self.marlin_s.to(x.dtype)
|
||||
#print(self.sms * ((orig_shape[0]+63)//64))
|
||||
|
||||
sms = self.sms
|
||||
|
||||
x = vLLMMarlin.gptq_marlin_gemm(
|
||||
x,
|
||||
self.marlin_q_w,
|
||||
marlin_s,
|
||||
self.g_idx,
|
||||
self.sort_indices,
|
||||
self.workspace.scratch,
|
||||
self.num_bits,
|
||||
bsz_tensor,
|
||||
x.shape[0],
|
||||
self.n,
|
||||
x.shape[-1],
|
||||
sms,
|
||||
self.is_k_full,
|
||||
)
|
||||
# TODO: don't padding bias
|
||||
if self.has_bias:
|
||||
x = x + self.bias
|
||||
if self.padding:
|
||||
x = x[:,:self.orin_out_features]
|
||||
orig_shape[-1] = self.orin_out_features
|
||||
else:
|
||||
orig_shape[-1] = self.out_features
|
||||
return x.reshape(orig_shape).to(orig_dtype)
|
||||
|
||||
def benchLinearMarlin(input_dim, output_dim):#, out_file
|
||||
print("benchmarking MLP Marlin")
|
||||
print("-----------------------------------------------------------")
|
||||
headers = ["batch_size", "tp", "used_time", "bandwidth GB/s", "TFLOPS", "cases", "padding", "sms"]
|
||||
print(" | ".join(headers) + "\n")
|
||||
rows = []
|
||||
for batch_size in exp_range(1, 64):
|
||||
for tp in exp_range(1, max_tp):
|
||||
torch.cuda.empty_cache()
|
||||
if output_dim % tp != 0:
|
||||
continue
|
||||
cur_output_dim = output_dim // tp
|
||||
modules = []
|
||||
inputs = []
|
||||
data_size = int(0.53125*input_dim*cur_output_dim)
|
||||
input_size = int(2*batch_size*input_dim)
|
||||
output_size = int(2*batch_size*cur_output_dim)
|
||||
usable_mem = get_usable_mem() - 2 * input_dim * cur_output_dim
|
||||
min_cases = max(global_num_cases, (2*L2_size) // (data_size+input_size))
|
||||
cases = int(min(min_cases, (usable_mem * 0.8) // (data_size+input_size)))
|
||||
#print(usable_mem, data_size, input_size, cases)
|
||||
|
||||
bsz_tensor = torch.tensor([batch_size], device=global_device, dtype=torch.int32)
|
||||
|
||||
if cases == 0:
|
||||
row = [f"{batch_size}", "OOM", "OOM", "OOM", "0", "False"]
|
||||
rows.append(row)
|
||||
break
|
||||
for _ in range(cases):
|
||||
modules.append(LinearMarlin(input_dim, cur_output_dim, sms=56, non_equal_division=False).to(device=global_device).eval())
|
||||
inputs.append(torch.randn(batch_size, 1, input_dim, device=global_device))
|
||||
|
||||
def forward(case_id):
|
||||
modules[case_id](inputs[case_id], bsz_tensor)
|
||||
|
||||
used_time = timing(forward, iters=cases)
|
||||
bandwidth = (data_size+input_size+output_size)/used_time/1e6
|
||||
flops = 2*batch_size*input_dim*cur_output_dim
|
||||
tflops = flops/used_time/1e9
|
||||
cur_sms = modules[0].sms
|
||||
row = [f"{batch_size}", f"{tp}", f"{used_time}", f"{bandwidth}", f"{tflops}", f"{cases}", modules[0].padding, cur_sms]
|
||||
rows.append(row)
|
||||
print(f"{batch_size}", f"{tp}", f"{used_time}", f"{bandwidth}", f"{tflops}", f"{cases}", modules[0].padding, cur_sms)
|
||||
|
||||
"""
|
||||
with open(out_file, 'w', newline='') as csvfile:
|
||||
csvwriter = csv.writer(csvfile)
|
||||
csvwriter.writerow(headers)
|
||||
for row in rows:
|
||||
csvwriter.writerow(row)
|
||||
"""
|
||||
|
||||
"""
|
||||
markdown_table = " | ".join(headers) + "\n"
|
||||
markdown_table += " | ".join(["---"] * len(headers)) + "\n"
|
||||
for row in rows:
|
||||
markdown_table += " | ".join(row) + "\n"
|
||||
|
||||
print(markdown_table)
|
||||
"""
|
||||
#print("finish write file", out_file)
|
||||
#print("-------------------------------------------------------------")
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
benchLinearMarlin(5120, 3584)
|
||||
exit(0)
|
||||
|
||||
max_batch = 1
|
||||
cur_batch = 1
|
||||
|
||||
|
||||
marlin_linear = LinearMarlin(5120, 3584)
|
||||
|
||||
input_tensor = torch.randn(max_batch, 1, 5120, device="cuda", dtype=torch.bfloat16)
|
||||
bsz_tensor = torch.tensor([max_batch], device="cuda", dtype=torch.int32)
|
||||
|
||||
out_truth = marlin_linear(input_tensor, bsz_tensor)
|
||||
|
||||
print(out_truth)
|
||||
|
||||
g = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(g):
|
||||
out_buf = marlin_linear(input_tensor, bsz_tensor)
|
||||
|
||||
for i in range(10000):
|
||||
g.replay()
|
||||
|
||||
#torch.testing.assert_close(out_buf, out_truth, rtol=1e-3, atol=1e-3)
|
||||
|
||||
marlin_linear = LinearMarlin(5120, 3584)
|
||||
g = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(g):
|
||||
out_buf = marlin_linear(input_tensor, bsz_tensor)
|
||||
|
||||
new_input = torch.randn(cur_batch, 1, 5120, device="cuda", dtype=torch.bfloat16)
|
||||
bsz_tensor.copy_(torch.tensor([cur_batch], device="cuda", dtype=torch.int32))
|
||||
|
||||
new_out_truth = marlin_linear(new_input, bsz_tensor)
|
||||
input_tensor[:cur_batch].copy_(new_input)
|
||||
input_tensor[cur_batch:] = 0
|
||||
|
||||
g.replay()
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
def printMinMax(tensor):
|
||||
abs_tensor = torch.abs(tensor)
|
||||
|
||||
min_val = torch.min(abs_tensor)
|
||||
max_val = torch.max(abs_tensor)
|
||||
|
||||
min_indices = (abs_tensor == min_val).nonzero(as_tuple=True)
|
||||
max_indices = (abs_tensor == max_val).nonzero(as_tuple=True)
|
||||
|
||||
print(f"min: {min_val.item()}")
|
||||
print(f"min idx: {min_indices}")
|
||||
print(f"max: {max_val.item()}")
|
||||
print(f"max idx: {max_indices}")
|
||||
|
||||
print(out_buf[:cur_batch].shape)
|
||||
print(new_out_truth.shape)
|
||||
|
||||
|
||||
printMinMax(out_buf[:cur_batch])
|
||||
printMinMax(new_out_truth)
|
||||
|
||||
#torch.testing.assert_close(out_buf[:cur_batch, 0, :], new_out_truth[:cur_batch, 0, :], rtol=1e-3, atol=1e-3)
|
||||
@@ -0,0 +1,308 @@
|
||||
#
|
||||
# Modified by Roberto Lopez Castro (roberto.lopez.castro@udc.es).
|
||||
#
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
# This is PyTorch implementation of main part of reorder_meta()
|
||||
# function, from tools/util/include/cutlass/util/host_reorder.h file
|
||||
# of CUTLASS source tree. Furthermore, CUTLASS template for sparse
|
||||
# GEMM decides upon layout of this matrix, and at the moment for the
|
||||
# sparse GEMM executed on tensor cores, this is layout described by
|
||||
# ColumnMajorInterleaved<2> data structure, in
|
||||
# include/cutlass/layout/matrix.h of CUTLASS source tree. The
|
||||
# reordering of meta matrix into meta_reordered matrix calculated
|
||||
# according to these segments of CUTLASS code is re-implemented here.
|
||||
# Note that this calculation produces offsets for scattering metadata
|
||||
# matrix elements into reordered metadata matrix elements (or,
|
||||
# equivalently, for gathering reordered metadata matrix element back
|
||||
# into metadata matrix elements).
|
||||
def _calculate_meta_reordering_scatter_offsets(m, meta_ncols, meta_dtype,
|
||||
device):
|
||||
dst_rows = torch.arange(0, m, device=device)[:, None].repeat(1, meta_ncols)
|
||||
dst_cols = torch.arange(0, meta_ncols, device=device).repeat(m, 1)
|
||||
|
||||
# Reorder the rows, then swizzle the 2x2 blocks.
|
||||
group_x = 64
|
||||
group_y = 32 if meta_dtype.itemsize == 2 else 16
|
||||
|
||||
dst_rows = (dst_rows // group_x * group_x + (dst_rows % 2) * 2 +
|
||||
(dst_rows % 8) // 4 + ((dst_rows % group_y) % 4) // 2 * 32 +
|
||||
((dst_rows % group_x) // 8) * 4)
|
||||
|
||||
topright = ((dst_rows % 2 == 0) & (dst_cols % 2 == 1)).to(torch.int8)
|
||||
bottomleft = ((dst_rows % 2 == 1) & (dst_cols % 2 == 0)).to(torch.int8)
|
||||
dst_rows += topright - bottomleft
|
||||
dst_cols -= topright - bottomleft
|
||||
|
||||
# Assumed that meta tensor is to be stored in CUTLASS
|
||||
# InterleavedColumnMajor layout, and reverse engineered
|
||||
# corresponding code to store values into this tensor.
|
||||
interleave = 2
|
||||
cols_maj = dst_cols // interleave
|
||||
cols_min = dst_cols % interleave
|
||||
return (cols_maj * m * interleave + dst_rows * interleave +
|
||||
cols_min).view(-1)
|
||||
|
||||
|
||||
# This function converts dense matrix into sparse semi-structured
|
||||
# representation, producing "compressed" matrix, in the layout used by
|
||||
# CUTLASS backend, and corresponding metadata matrix.
|
||||
def sparse_semi_structured_from_dense_cutlass(dense):
|
||||
if dense.dim() != 2:
|
||||
raise RuntimeError(
|
||||
f"Expected 2-dimensional dense tensor, got {dense.dim()}-dimensional tensor" # noqa: E501
|
||||
)
|
||||
|
||||
m, k = dense.shape
|
||||
device = dense.device
|
||||
|
||||
meta_dtype = torch.int8
|
||||
if dense.dtype == torch.int8:
|
||||
meta_dtype = torch.int32
|
||||
elif dense.dtype in [torch.half, torch.bfloat16, torch.float, torch.int32]:
|
||||
meta_dtype = torch.int16
|
||||
else:
|
||||
raise RuntimeError(f"Invalid datatype {dense.dtype} of dense matrix")
|
||||
quadbits_per_meta_elem = meta_dtype.itemsize * 8 // 4
|
||||
if quadbits_per_meta_elem not in (4, 8):
|
||||
raise RuntimeError(
|
||||
"Invalid number of elements per meta element calculated")
|
||||
|
||||
if meta_dtype == torch.int32:
|
||||
if m % 16 != 0:
|
||||
raise RuntimeError(
|
||||
f"Number of rows of dense matrix {m} must be divisible by 16")
|
||||
else:
|
||||
if m % 32 != 0:
|
||||
raise RuntimeError(
|
||||
f"Number of rows of dense matrix {m} must be divisible by 32")
|
||||
if k % (4 * quadbits_per_meta_elem) != 0:
|
||||
raise RuntimeError(
|
||||
f"Number of columns of dense matrix {k} must be divisible by {4 * quadbits_per_meta_elem}" # noqa: E501
|
||||
)
|
||||
|
||||
if dense.dtype != torch.float:
|
||||
ksparse = 4
|
||||
dense_4 = dense.view(-1, k // ksparse, ksparse)
|
||||
m0, m1, m2, m3 = (dense_4 != 0).unbind(-1)
|
||||
else:
|
||||
ksparse = 2
|
||||
dense_2 = dense.view(-1, k // ksparse, ksparse)
|
||||
m0, m2 = m1, m3 = (dense_2 != 0).unbind(-1)
|
||||
meta_ncols = k // (ksparse * quadbits_per_meta_elem)
|
||||
|
||||
# Encoding quadruples of True/False values as follows:
|
||||
# [True, True, False, False] -> 0b0100
|
||||
# [True, False, True, False] -> 0b1000
|
||||
# [False, True, True, False] -> 0b1001
|
||||
# [True, False, False, True ] -> 0b1100
|
||||
# [False, True, False, True ] -> 0b1101
|
||||
# [False, False, True, True ] -> 0b1110
|
||||
# Thus, lower two bits in the encoding are index of the True value
|
||||
# at the lowest index in the quadruple, and the higher two bits in
|
||||
# the encoding are index of the other True value in the quadruple.
|
||||
# In case there are less than two True values, than False value or
|
||||
# values at some index or indices are considered True for the
|
||||
# encoding. In case there are more than two True values, then the
|
||||
# excess True value(s) at some indices are considered False for
|
||||
# the encoding. The exact encodings used for these cases are as
|
||||
# follows:
|
||||
# [False, False, False, False] -> 0b1110
|
||||
# [False, False, False, True ] -> 0b1110
|
||||
# [False, False, True, False] -> 0b1110
|
||||
# [False, True, False, False] -> 0b1001
|
||||
# [False, True, True, True ] -> 0b1101
|
||||
# [True, False, False, False] -> 0b1000
|
||||
# [True, False, True, True ] -> 0b1100
|
||||
# [True, True, False, True ] -> 0b0100
|
||||
# [True, True, True, False] -> 0b0100
|
||||
# [True, True, True, True ] -> 0b0100
|
||||
# These particular encodings are chosen, with the help of Espresso
|
||||
# logic minimizer software, for the purpose of minimization of
|
||||
# corresponding Boolean functions, that translate non-zero flags
|
||||
# into encoding bits. Note also possible choices for the first
|
||||
# and last of these encodings were limited only to (0b0100,
|
||||
# 0b1110), in order to produce valid encodings for 1:2 sparsity
|
||||
# case.
|
||||
|
||||
expr0 = m0 & m1
|
||||
expr1 = ~m0 & m1
|
||||
expr2 = ~m0 & ~m1
|
||||
bit0 = expr1
|
||||
bit1 = expr2
|
||||
bit2 = expr0 | expr2 | m3
|
||||
bit3 = expr1 | ~m1
|
||||
idxs0 = bit0 | (bit1.to(torch.int64) << 1)
|
||||
idxs1 = bit2 | (bit3.to(torch.int64) << 1)
|
||||
|
||||
if dense.dtype != torch.float:
|
||||
sparse0 = dense_4.gather(
|
||||
-1, idxs0.unsqueeze(-1)) # type: ignore[possibly-undefined]
|
||||
sparse1 = dense_4.gather(-1, idxs1.unsqueeze(-1))
|
||||
sparse = torch.stack((sparse0, sparse1), dim=-1).view(m, k // 2)
|
||||
else:
|
||||
sparse = dense_2.gather(-1,
|
||||
idxs0.unsqueeze(-1) // 2).view(
|
||||
m,
|
||||
k // 2) # type: ignore[possibly-undefined]
|
||||
|
||||
meta_4 = idxs0 | (idxs1 << 2)
|
||||
meta_n = meta_4.view(
|
||||
(-1, meta_ncols, quadbits_per_meta_elem)).to(meta_dtype)
|
||||
|
||||
if quadbits_per_meta_elem == 4:
|
||||
meta = (meta_n[:, :, 0]
|
||||
| (meta_n[:, :, 1] << 4)
|
||||
| (meta_n[:, :, 2] << 8)
|
||||
| (meta_n[:, :, 3] << 12))
|
||||
elif quadbits_per_meta_elem == 8:
|
||||
meta = (meta_n[:, :, 0]
|
||||
| (meta_n[:, :, 1] << 4)
|
||||
| (meta_n[:, :, 2] << 8)
|
||||
| (meta_n[:, :, 3] << 12)
|
||||
| (meta_n[:, :, 4] << 16)
|
||||
| (meta_n[:, :, 5] << 20)
|
||||
| (meta_n[:, :, 6] << 24)
|
||||
| (meta_n[:, :, 7] << 28))
|
||||
|
||||
# Reorder meta tensor elements.
|
||||
meta_reordered = meta.new_empty(
|
||||
(m * meta_ncols, )) # type: ignore[possibly-undefined]
|
||||
meta_offsets = _calculate_meta_reordering_scatter_offsets(
|
||||
m, meta_ncols, meta_dtype, device)
|
||||
meta_reordered.scatter_(0, meta_offsets, meta.view(-1))
|
||||
|
||||
return (sparse, meta_reordered.view(m, meta_ncols))
|
||||
|
||||
|
||||
# This function performs reverse of the function above - it
|
||||
# reconstructs dense matrix from a pair of "compressed" matrix, given
|
||||
# in the layout used by CUTLASS backend, and accompanying metadata
|
||||
# matrix.
|
||||
def sparse_semi_structured_to_dense_cutlass(sparse, meta_reordered):
|
||||
if sparse.dim() != 2:
|
||||
raise RuntimeError(
|
||||
f"Expected 2-dimensional sparse tensor, got {sparse.dim()}-dimensional tensor" # noqa: E501
|
||||
)
|
||||
|
||||
m, k = sparse.shape
|
||||
device = sparse.device
|
||||
|
||||
if meta_reordered.dim() != 2:
|
||||
raise RuntimeError(
|
||||
f"Expected 2-dimensional meta tensor, got {meta_reordered.dim()}-dimensional tensor" # noqa: E501
|
||||
)
|
||||
if meta_reordered.device != device:
|
||||
raise RuntimeError(
|
||||
f"Expected meta matrix to be on {device} device, got matrix on {meta_reordered.device} device" # noqa: E501
|
||||
)
|
||||
|
||||
meta_dtype = meta_reordered.dtype
|
||||
if meta_dtype not in (torch.int16, torch.int32):
|
||||
raise RuntimeError(f"Invalid datatype {meta_dtype} of meta matrix")
|
||||
quadbits_per_meta_elem = meta_dtype.itemsize * 8 // 4
|
||||
|
||||
ksparse = 4 if sparse.dtype != torch.float else 2
|
||||
|
||||
meta_nrows, meta_ncols = meta_reordered.shape
|
||||
if meta_nrows != m:
|
||||
raise RuntimeError(
|
||||
f"Number of rows of meta matrix {meta_nrows} must be equal to number of columns of spase matrix {m}" # noqa: E501
|
||||
)
|
||||
if meta_ncols * ksparse * quadbits_per_meta_elem != 2 * k:
|
||||
raise RuntimeError(
|
||||
f"Number of columns of sparse matrix {k} different from the {meta_ncols * ksparse * quadbits_per_meta_elem // 2}, " # noqa: E501
|
||||
"expected according to the number of columns of meta matrix")
|
||||
|
||||
# Undo meta tensor elements reordering.
|
||||
meta_offsets = _calculate_meta_reordering_scatter_offsets(
|
||||
m, meta_ncols, meta_dtype, device)
|
||||
meta = torch.gather(meta_reordered.view(-1), 0,
|
||||
meta_offsets).view(m, meta_ncols)
|
||||
|
||||
# Unpack sparse tensor back to original dense tensor, using
|
||||
# information provided by meta tensor. Note that torch.float
|
||||
# datatype is handled pretty much the same as
|
||||
# torch.half/torch.bfloat16, as metadata for a pair of torch.float
|
||||
# value is encoded as if underlying 8 bytes contain four
|
||||
# torch.half/torch.bfloat16 values, where either first two or last
|
||||
# two are zeros.
|
||||
meta_2 = torch.empty(
|
||||
(m, meta_ncols, 2 * quadbits_per_meta_elem),
|
||||
dtype=meta_dtype,
|
||||
device=device,
|
||||
)
|
||||
if quadbits_per_meta_elem == 4:
|
||||
meta_2[:, :, 0] = meta & 0b11
|
||||
meta_2[:, :, 1] = (meta >> 2) & 0b11
|
||||
meta_2[:, :, 2] = (meta >> 4) & 0b11
|
||||
meta_2[:, :, 3] = (meta >> 6) & 0b11
|
||||
meta_2[:, :, 4] = (meta >> 8) & 0b11
|
||||
meta_2[:, :, 5] = (meta >> 10) & 0b11
|
||||
meta_2[:, :, 6] = (meta >> 12) & 0b11
|
||||
meta_2[:, :, 7] = (meta >> 14) & 0b11
|
||||
elif quadbits_per_meta_elem == 8:
|
||||
meta_2[:, :, 0] = meta & 0b11
|
||||
meta_2[:, :, 1] = (meta >> 2) & 0b11
|
||||
meta_2[:, :, 2] = (meta >> 4) & 0b11
|
||||
meta_2[:, :, 3] = (meta >> 6) & 0b11
|
||||
meta_2[:, :, 4] = (meta >> 8) & 0b11
|
||||
meta_2[:, :, 5] = (meta >> 10) & 0b11
|
||||
meta_2[:, :, 6] = (meta >> 12) & 0b11
|
||||
meta_2[:, :, 7] = (meta >> 14) & 0b11
|
||||
meta_2[:, :, 8] = (meta >> 16) & 0b11
|
||||
meta_2[:, :, 9] = (meta >> 18) & 0b11
|
||||
meta_2[:, :, 10] = (meta >> 20) & 0b11
|
||||
meta_2[:, :, 11] = (meta >> 22) & 0b11
|
||||
meta_2[:, :, 12] = (meta >> 24) & 0b11
|
||||
meta_2[:, :, 13] = (meta >> 26) & 0b11
|
||||
meta_2[:, :, 14] = (meta >> 28) & 0b11
|
||||
meta_2[:, :, 15] = (meta >> 30) & 0b11
|
||||
|
||||
dense_offsets = meta_2.view(-1) + (
|
||||
torch.arange(0, 2 * m * k // ksparse, device=device) * 4).view(
|
||||
-1, 1).repeat(1, 2).view(-1)
|
||||
|
||||
dense = torch.zeros((m * 2 * k, ), dtype=sparse.dtype, device=device)
|
||||
if sparse.dtype != torch.float:
|
||||
# dense.scatter_(0, dense_offsets, sparse.view(-1))
|
||||
dense.scatter_(0, dense_offsets, sparse.reshape(-1))
|
||||
else:
|
||||
dense.view(torch.half).scatter_(0, dense_offsets,
|
||||
sparse.view(torch.half).view(-1))
|
||||
|
||||
return dense.view(m, 2 * k)
|
||||
|
||||
|
||||
def mask_creator(tensor):
|
||||
"""
|
||||
Class for creating N:M sparsity masks.
|
||||
Masks will be created using the N:M ratio, where for every block of
|
||||
M weights, N will be pruned based on ranked weight value. Each mask
|
||||
will correspond to the given tensor.
|
||||
|
||||
:param N: The number of weights in a group to keep
|
||||
:param M: The size of a weight group
|
||||
"""
|
||||
N = 2
|
||||
M = 4
|
||||
|
||||
mask = None
|
||||
# for i, tensor in enumerate(tensors):
|
||||
if tensor.numel() % M != 0:
|
||||
raise ValueError(
|
||||
f"Tensor of size {tensor.shape} can't be evenly divided into "
|
||||
f"{M} groups")
|
||||
|
||||
num_groups = tensor.numel() // M
|
||||
|
||||
# N:M sparsity for linear layers
|
||||
tensor_temp = tensor.detach().abs().reshape(num_groups, M)
|
||||
index = torch.argsort(tensor_temp, dim=1)[:, :int(M - N)]
|
||||
|
||||
w_b = torch.ones(tensor_temp.shape, device=tensor_temp.device)
|
||||
mask = w_b.scatter_(dim=1, index=index, value=0).reshape(tensor.shape)
|
||||
|
||||
return mask
|
||||
@@ -0,0 +1,65 @@
|
||||
'''
|
||||
Date: 2024-11-08 02:46:07
|
||||
LastEditors: djw
|
||||
LastEditTime: 2024-11-08 02:46:41
|
||||
'''
|
||||
"""This file is used for /tests and /benchmarks"""
|
||||
from typing import Dict, List
|
||||
|
||||
import numpy
|
||||
import torch
|
||||
|
||||
|
||||
# Precompute permutations for Marlin24 weight and scale shuffling # noqa: E501
|
||||
#
|
||||
# Marlin works on [16*2,64] tiles. The goal of the permutations is to reorder the weight data so that it is compatible noqa: # noqa: E501
|
||||
# with the tensor-core format that is described here:
|
||||
# https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type # noqa: E501
|
||||
#
|
||||
# As a result of this reordering, the vector loads inside the kernel will get the data as it is needed for tensor-core # noqa: E501
|
||||
# (without the need to use ldmatrix instructions) # noqa: E501
|
||||
def get_perms_24(num_bits: int):
|
||||
perm_list: List[int] = []
|
||||
for i in range(32):
|
||||
perm1: List[int] = []
|
||||
col = i // 4
|
||||
col_o = col // 2
|
||||
for block in [0, 1]:
|
||||
for row in [
|
||||
2 * (i % 4),
|
||||
2 * (i % 4) + 1,
|
||||
2 * (i % 4 + 4),
|
||||
2 * (i % 4 + 4) + 1,
|
||||
]:
|
||||
perm1.append(16 * row + col_o * 256 + 8 * (col % 2) +
|
||||
4 * block)
|
||||
for j in range(4):
|
||||
perm_list.extend([p + 1 * j for p in perm1])
|
||||
perm = numpy.array(perm_list)
|
||||
|
||||
if num_bits == 4:
|
||||
interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7])
|
||||
elif num_bits == 8:
|
||||
interleave = numpy.array([0, 2, 1, 3])
|
||||
else:
|
||||
raise ValueError("num_bits must be 4 or 8, got {}".format(num_bits))
|
||||
|
||||
perm = perm.reshape((-1, len(interleave)))[:, interleave].ravel()
|
||||
perm = torch.from_numpy(perm)
|
||||
scale_perm: List[int] = []
|
||||
for i in range(8):
|
||||
scale_perm.extend([i * 8 + j for j in [0, 4, 1, 5, 2, 6, 3, 7]])
|
||||
scale_perm_single: List[int] = []
|
||||
for i in range(8):
|
||||
scale_perm_single.extend([8 * i + j for j in [0, 1, 2, 3, 4, 5, 6, 7]])
|
||||
return perm, scale_perm, scale_perm_single
|
||||
|
||||
|
||||
marlin_24_perm: Dict[int, torch.Tensor] = {}
|
||||
marlin_24_scale_perm: Dict[int, List[int]] = {}
|
||||
marlin_24_scale_perm_single: Dict[int, List[int]] = {}
|
||||
for num_bits in [4, 8]:
|
||||
perm_24, scale_perm_24, scale_perm_single_24 = get_perms_24(num_bits)
|
||||
marlin_24_perm[num_bits] = perm_24
|
||||
marlin_24_scale_perm[num_bits] = scale_perm_24
|
||||
marlin_24_scale_perm_single[num_bits] = scale_perm_single_24
|
||||
@@ -0,0 +1,65 @@
|
||||
'''
|
||||
Date: 2024-11-08 02:46:47
|
||||
LastEditors: djw
|
||||
LastEditTime: 2024-11-08 02:46:55
|
||||
'''
|
||||
"""This file is used for /tests and /benchmarks"""
|
||||
from typing import Dict, List
|
||||
|
||||
import numpy
|
||||
import torch
|
||||
|
||||
|
||||
# Precompute permutations for Marlin weight and scale shuffling # noqa: E501
|
||||
#
|
||||
# Marlin works on [16,64] tiles. The goal of the permutations is to reorder the weight data so that it is compatible noqa: # noqa: E501
|
||||
# with the tensor-core format that is described here:
|
||||
# https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type # noqa: E501
|
||||
#
|
||||
# As a result of this reordering, the vector loads inside the kernel will get the data as it is needed for tensor-core # noqa: E501
|
||||
# (without the need to use ldmatrix instructions) # noqa: E501
|
||||
def get_perms(num_bits: int):
|
||||
perm_list: List[int] = []
|
||||
for i in range(32):
|
||||
perm1: List[int] = []
|
||||
col = i // 4
|
||||
for block in [0, 1]:
|
||||
for row in [
|
||||
2 * (i % 4),
|
||||
2 * (i % 4) + 1,
|
||||
2 * (i % 4 + 4),
|
||||
2 * (i % 4 + 4) + 1,
|
||||
]:
|
||||
perm1.append(16 * row + col + 8 * block)
|
||||
for j in range(4):
|
||||
perm_list.extend([p + 256 * j for p in perm1])
|
||||
|
||||
perm = numpy.array(perm_list)
|
||||
|
||||
if num_bits == 4:
|
||||
interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7])
|
||||
elif num_bits == 8:
|
||||
interleave = numpy.array([0, 2, 1, 3])
|
||||
else:
|
||||
raise Exception("num_bits must be 4 or 8, got {}".format(num_bits))
|
||||
|
||||
perm = perm.reshape((-1, len(interleave)))[:, interleave].ravel()
|
||||
perm = torch.from_numpy(perm)
|
||||
scale_perm: List[int] = []
|
||||
for i in range(8):
|
||||
scale_perm.extend([i + 8 * j for j in range(8)])
|
||||
scale_perm_single: List[int] = []
|
||||
for i in range(4):
|
||||
scale_perm_single.extend(
|
||||
[2 * i + j for j in [0, 1, 8, 9, 16, 17, 24, 25]])
|
||||
return perm, scale_perm, scale_perm_single
|
||||
|
||||
|
||||
marlin_perm: Dict[int, torch.Tensor] = {}
|
||||
marlin_scale_perm: Dict[int, List[int]] = {}
|
||||
marlin_scale_perm_single: Dict[int, List[int]] = {}
|
||||
for num_bits in [4, 8]:
|
||||
perm, scale_perm, scale_perm_single = get_perms(num_bits)
|
||||
marlin_perm[num_bits] = perm
|
||||
marlin_scale_perm[num_bits] = scale_perm
|
||||
marlin_scale_perm_single[num_bits] = scale_perm_single
|
||||
@@ -0,0 +1,234 @@
|
||||
"""This file is used for /tests and /benchmarks"""
|
||||
import random
|
||||
|
||||
import numpy
|
||||
import torch
|
||||
|
||||
from .format24 import (
|
||||
mask_creator, sparse_semi_structured_from_dense_cutlass)
|
||||
from .marlin_24_perms import (
|
||||
marlin_24_perm, marlin_24_scale_perm, marlin_24_scale_perm_single)
|
||||
from .marlin_perms import (
|
||||
marlin_perm, marlin_scale_perm, marlin_scale_perm_single)
|
||||
from .quant_utils import (
|
||||
get_pack_factor, quantize_weights, sort_weights, dequantize_weights)
|
||||
|
||||
|
||||
|
||||
__cuda_arch = torch.cuda.get_device_capability()
|
||||
|
||||
MARLIN_TILE = 16
|
||||
|
||||
GPTQ_MARLIN_TILE = 16
|
||||
GPTQ_MARLIN_MIN_THREAD_N = 64
|
||||
GPTQ_MARLIN_MIN_THREAD_K = 128
|
||||
GPTQ_MARLIN_MAX_PARALLEL = 16
|
||||
|
||||
GPTQ_MARLIN_SUPPORTED_NUM_BITS = [4, 8]
|
||||
GPTQ_MARLIN_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128]
|
||||
GPTQ_MARLIN_SUPPORTED_SYM = [True]
|
||||
|
||||
def is_marlin_supported():
|
||||
return __cuda_arch[0] >= 8
|
||||
|
||||
|
||||
def marlin_permute_weights(q_w, size_k, size_n, perm, tile=MARLIN_TILE):
|
||||
assert q_w.shape == (size_k, size_n)
|
||||
assert size_k % tile == 0, f"size_k = {size_k}, tile = {tile}"
|
||||
assert size_n % tile == 0, f"size_k = {size_n}, tile = {tile}"
|
||||
|
||||
# Permute weights to 16x64 marlin tiles
|
||||
q_w = q_w.reshape((size_k // tile, tile, size_n // tile, tile))
|
||||
q_w = q_w.permute((0, 2, 1, 3))
|
||||
q_w = q_w.reshape((size_k // tile, size_n * tile))
|
||||
|
||||
q_w = q_w.reshape((-1, perm.numel()))[:, perm].reshape(q_w.shape)
|
||||
|
||||
return q_w
|
||||
|
||||
|
||||
def marlin_weights(q_w, size_k, size_n, num_bits, perm):
|
||||
# Permute
|
||||
q_w = marlin_permute_weights(q_w, size_k, size_n, perm)
|
||||
|
||||
# Pack
|
||||
pack_factor = get_pack_factor(num_bits)
|
||||
orig_device = q_w.device
|
||||
|
||||
q_w = q_w.cpu().numpy().astype(numpy.uint32)
|
||||
|
||||
q_packed = numpy.zeros((q_w.shape[0], q_w.shape[1] // pack_factor),
|
||||
dtype=numpy.uint32)
|
||||
for i in range(pack_factor):
|
||||
q_packed |= q_w[:, i::pack_factor] << num_bits * i
|
||||
|
||||
q_packed = torch.from_numpy(q_packed.astype(numpy.int32)).to(orig_device)
|
||||
|
||||
return q_packed
|
||||
|
||||
|
||||
def marlin_permute_scales(s, size_k, size_n, group_size, scale_perm,
|
||||
scale_perm_single):
|
||||
if group_size < size_k and group_size != -1:
|
||||
s = s.reshape((-1, len(scale_perm)))[:, scale_perm]
|
||||
else:
|
||||
s = s.reshape((-1, len(scale_perm_single)))[:, scale_perm_single]
|
||||
s = s.reshape((-1, size_n)).contiguous()
|
||||
|
||||
return s
|
||||
|
||||
|
||||
def marlin_quantize(
|
||||
w: torch.Tensor,
|
||||
num_bits: int,
|
||||
group_size: int,
|
||||
act_order: bool,
|
||||
):
|
||||
size_k, size_n = w.shape
|
||||
|
||||
# Normalize group_size
|
||||
if group_size == -1:
|
||||
group_size = size_k
|
||||
assert group_size <= size_k
|
||||
|
||||
# Quantize (and apply act_order if provided)
|
||||
w_ref, q_w, s, g_idx, rand_perm = quantize_weights(w, num_bits, group_size,
|
||||
act_order)
|
||||
|
||||
# For act_order, sort the "weights" and "g_idx" so that group ids are
|
||||
# increasing
|
||||
sort_indices = torch.empty(0, dtype=torch.int, device=w.device)
|
||||
if act_order:
|
||||
q_w, g_idx, sort_indices = sort_weights(q_w, g_idx)
|
||||
|
||||
# Reformat to marlin
|
||||
marlin_q_w = marlin_weights(q_w, size_k, size_n, num_bits,
|
||||
marlin_perm[num_bits])
|
||||
marlin_s = marlin_permute_scales(s, size_k, size_n, group_size,
|
||||
marlin_scale_perm[num_bits],
|
||||
marlin_scale_perm_single[num_bits])
|
||||
|
||||
# Create result
|
||||
res_list = [w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, rand_perm]
|
||||
for i in range(len(res_list)):
|
||||
res_list[i] = res_list[i].to(w.device)
|
||||
|
||||
return res_list
|
||||
|
||||
|
||||
def inject_24(w, size_k, size_n):
|
||||
assert w.shape == (size_k, size_n)
|
||||
|
||||
mask = mask_creator(w.t()).t().cuda().bool()
|
||||
|
||||
return (mask * w).contiguous(), mask.contiguous()
|
||||
|
||||
|
||||
def check_24(w, num_rows_to_sample=50, _verbose=False):
|
||||
BLOCK_SIZE = 4
|
||||
MAX_NON_ZEROS = 2
|
||||
|
||||
w = w.t().contiguous()
|
||||
|
||||
print("check_24: w.shape = {}".format(w.shape))
|
||||
|
||||
num_rows, num_cols = w.shape
|
||||
sampled_row_idxs = random.choices(range(num_rows), k=num_rows_to_sample)
|
||||
if _verbose:
|
||||
print(f"Sampled row idxs = {sampled_row_idxs}")
|
||||
|
||||
total_segments = 0
|
||||
non_24_segments = 0
|
||||
for i in sampled_row_idxs:
|
||||
for j in range(0, num_cols - BLOCK_SIZE, BLOCK_SIZE):
|
||||
total_segments += 1
|
||||
block = w[i, j:j + BLOCK_SIZE]
|
||||
num_nonzero = torch.count_nonzero(block)
|
||||
if num_nonzero > MAX_NON_ZEROS:
|
||||
print("i = {} j = {} block = {}".format(i, j, block))
|
||||
non_24_segments += 1
|
||||
|
||||
print(f"{non_24_segments} / {total_segments} do not have 2:4 structure.")
|
||||
|
||||
|
||||
def compress_quantized_24_weight(q_24, size_k, size_n, num_bits):
|
||||
assert q_24.shape == (size_k, size_n)
|
||||
|
||||
# Remove zp to normalize over 0
|
||||
max_q_val = (1 << num_bits) - 1
|
||||
zp = (max_q_val + 1) // 2
|
||||
q_24_no_zp = q_24 - zp
|
||||
|
||||
# Compress
|
||||
q_24_no_zp = q_24_no_zp.t().contiguous()
|
||||
q_24_no_zp_comp, meta = sparse_semi_structured_from_dense_cutlass(
|
||||
q_24_no_zp)
|
||||
q_24_no_zp_comp = q_24_no_zp_comp.t().contiguous()
|
||||
|
||||
# Restore zp
|
||||
q_24_comp = q_24_no_zp_comp + zp
|
||||
|
||||
# Resize meta to its actual shape (without moving any data)
|
||||
meta = meta.resize_(meta.shape[1] // 2, meta.shape[0] * 2)
|
||||
|
||||
return q_24_comp, meta
|
||||
|
||||
|
||||
def marlin_24_quantize(
|
||||
w: torch.Tensor,
|
||||
num_bits: int,
|
||||
group_size: int,
|
||||
):
|
||||
size_k, size_n = w.shape
|
||||
|
||||
# Normalize group_size
|
||||
if group_size == -1:
|
||||
group_size = size_k
|
||||
assert group_size <= size_k
|
||||
|
||||
# Inject 2:4 sparsity
|
||||
w_24, mask_24 = inject_24(w, size_k, size_n)
|
||||
|
||||
# Quantize
|
||||
w_24_ref, q_w_24, s, g_idx, rand_perm = quantize_weights(w_24,
|
||||
num_bits,
|
||||
group_size,
|
||||
act_order=False)
|
||||
|
||||
# Compress quantized weight
|
||||
q_w_24_comp, meta = compress_quantized_24_weight(q_w_24, size_k, size_n,
|
||||
num_bits)
|
||||
size_k_comp = size_k // 2
|
||||
|
||||
# Reformat to marlin
|
||||
marlin_24_q_w_comp = marlin_weights(q_w_24_comp, size_k_comp, size_n,
|
||||
num_bits, marlin_24_perm[num_bits])
|
||||
marlin_24_s = marlin_permute_scales(s, size_k, size_n, group_size,
|
||||
marlin_24_scale_perm[num_bits],
|
||||
marlin_24_scale_perm_single[num_bits])
|
||||
|
||||
# Create result
|
||||
res_list = [w_24_ref, marlin_24_q_w_comp, meta, marlin_24_s]
|
||||
for i in range(len(res_list)):
|
||||
res_list[i] = res_list[i].to(w.device)
|
||||
|
||||
return res_list
|
||||
|
||||
|
||||
def compute_max_diff(output, output_ref):
|
||||
return torch.mean(torch.abs(output - output_ref)) / torch.mean(
|
||||
torch.abs(output_ref))
|
||||
|
||||
|
||||
class MarlinWorkspace:
|
||||
|
||||
def __init__(self, out_features, min_thread_n, max_parallel, device):
|
||||
assert (out_features % min_thread_n == 0), (
|
||||
"out_features = {} is undivisible by min_thread_n = {}".format(
|
||||
out_features, min_thread_n))
|
||||
|
||||
max_workspace_size = ((out_features // min_thread_n) * max_parallel)
|
||||
|
||||
self.scratch = torch.zeros(max_workspace_size,
|
||||
dtype=torch.int,
|
||||
device=device)
|
||||
@@ -0,0 +1,195 @@
|
||||
"""This file is used for /tests and /benchmarks"""
|
||||
import numpy
|
||||
import torch
|
||||
|
||||
SUPPORTED_NUM_BITS = [4, 8]
|
||||
SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128]
|
||||
|
||||
|
||||
def get_pack_factor(num_bits):
|
||||
assert num_bits in SUPPORTED_NUM_BITS, f"Unsupported num_bits = {num_bits}"
|
||||
return 32 // num_bits
|
||||
|
||||
|
||||
def permute_rows(q_w: torch.Tensor, w_ref: torch.Tensor, group_size: int):
|
||||
assert q_w.shape == w_ref.shape
|
||||
|
||||
orig_device = q_w.device
|
||||
k_size, _ = q_w.shape
|
||||
|
||||
g_idx = torch.zeros((k_size, ), dtype=torch.int32)
|
||||
for i in range(k_size):
|
||||
g_idx[i] = i // group_size
|
||||
|
||||
# Simulate act_order by doing a random permutation on K
|
||||
rand_perm = torch.randperm(k_size)
|
||||
|
||||
g_idx = g_idx[rand_perm].contiguous()
|
||||
q_w = q_w[rand_perm, :].contiguous()
|
||||
w_ref = w_ref[rand_perm, :].contiguous()
|
||||
|
||||
return (
|
||||
w_ref.to(device=orig_device),
|
||||
q_w.to(device=orig_device),
|
||||
g_idx.to(device=orig_device),
|
||||
rand_perm.to(device=orig_device),
|
||||
)
|
||||
|
||||
|
||||
# Function: Dequantize quantized weights
|
||||
def dequantize_weights(qweight, qzeros, scales, g_idx, bits=4, group_size=128, device='cuda:0'):
|
||||
# Create a tensor for bitwise right shift operation
|
||||
wf = torch.tensor(list(range(0, 32, bits)), dtype=torch.int32, device=device).unsqueeze(0)
|
||||
|
||||
# Apply bitwise right shift and convert qzeros to the appropriate type
|
||||
zeros = torch.bitwise_right_shift(torch.unsqueeze(qzeros, 2).expand(-1, -1, 32 // bits), wf.unsqueeze(0)).to(torch.int16 if bits == 8 else torch.int8)
|
||||
torch.bitwise_and(zeros, (2 ** bits) - 1, out=zeros)
|
||||
|
||||
# Reshape the zeros tensor
|
||||
zeros = zeros + 1
|
||||
zeros = zeros.reshape(-1, 1, zeros.shape[1] * zeros.shape[2])
|
||||
|
||||
# Reshape the scales tensor
|
||||
scales = scales.reshape(-1, 1, scales.shape[-1])
|
||||
|
||||
# Similar bitwise right shift operation for qweight and reshape
|
||||
weight = torch.bitwise_right_shift(torch.unsqueeze(qweight, 1).expand(-1, 32 // bits, -1), wf.unsqueeze(-1)).to(torch.int16 if bits == 8 else torch.int8)
|
||||
torch.bitwise_and(weight, (2 ** bits) - 1, out=weight)
|
||||
weight = weight.reshape(-1, group_size, weight.shape[2])
|
||||
|
||||
# Apply dequantization formula and reshape the final weight
|
||||
weight = (scales * (weight - zeros))
|
||||
weight = weight.reshape(weight.shape[0] * weight.shape[1], weight.shape[2])
|
||||
|
||||
# Return the transposed weight
|
||||
return weight.transpose(0, 1)
|
||||
|
||||
def quantize_weights(w: torch.Tensor, num_bits: int, group_size: int,
|
||||
act_order: bool):
|
||||
orig_device = w.device
|
||||
size_k, size_n = w.shape
|
||||
|
||||
assert w.is_floating_point(), "w must be float"
|
||||
assert num_bits in SUPPORTED_NUM_BITS, f"Unsupported num_bits = {num_bits}"
|
||||
assert group_size in SUPPORTED_GROUP_SIZES + [
|
||||
size_k
|
||||
], f"Unsupported groupsize = {group_size}"
|
||||
|
||||
if group_size == -1:
|
||||
group_size = size_k
|
||||
assert group_size <= size_k
|
||||
|
||||
max_q_val = 2**num_bits - 1
|
||||
half_q_val = (max_q_val + 1) // 2
|
||||
|
||||
# Reshape to [groupsize, -1]
|
||||
if group_size < size_k:
|
||||
w = w.view((-1, group_size, size_n))
|
||||
w = w.permute(1, 0, 2)
|
||||
w = w.reshape((group_size, -1))
|
||||
|
||||
# Compute scale for each group
|
||||
s = torch.max(torch.abs(w), 0, keepdim=True)[0]
|
||||
s *= 2 / max_q_val # 2 => symmetric
|
||||
|
||||
# Quantize
|
||||
q_w = torch.round(w / s).int()
|
||||
q_w += half_q_val
|
||||
q_w = torch.clamp(q_w, 0, max_q_val)
|
||||
|
||||
# Compute ref (dequantized)
|
||||
w_ref = (q_w - half_q_val).half() * s
|
||||
|
||||
# Restore original shapes
|
||||
if group_size < size_k:
|
||||
|
||||
def reshape_w(w):
|
||||
w = w.reshape((group_size, -1, size_n))
|
||||
w = w.permute(1, 0, 2)
|
||||
w = w.reshape((size_k, size_n)).contiguous()
|
||||
return w
|
||||
|
||||
q_w = reshape_w(q_w)
|
||||
w_ref = reshape_w(w_ref)
|
||||
|
||||
s = s.reshape((-1, size_n)).contiguous()
|
||||
|
||||
# Apply act_order
|
||||
g_idx = torch.empty(0, dtype=torch.int, device=w.device)
|
||||
rand_perm = torch.empty(0, dtype=torch.int, device=w.device)
|
||||
if act_order:
|
||||
assert (
|
||||
group_size < size_k
|
||||
), "For act_order, groupsize = {} must be less than size_k = {}".format(
|
||||
group_size, size_k)
|
||||
|
||||
w_ref, q_w, g_idx, rand_perm = permute_rows(q_w, w_ref, group_size)
|
||||
|
||||
return (
|
||||
w_ref.to(device=orig_device),
|
||||
q_w.to(device=orig_device),
|
||||
s.to(device=orig_device),
|
||||
g_idx.to(device=orig_device),
|
||||
rand_perm.to(device=orig_device),
|
||||
)
|
||||
|
||||
|
||||
def sort_weights(q_w: torch.Tensor, g_idx: torch.Tensor):
|
||||
orig_device = q_w.device
|
||||
|
||||
sort_indices = torch.argsort(g_idx).to(
|
||||
dtype=torch.int32) # Sort based on g_idx
|
||||
|
||||
g_idx = g_idx[sort_indices].contiguous()
|
||||
q_w = q_w[sort_indices, :].contiguous()
|
||||
|
||||
return (
|
||||
q_w.to(device=orig_device),
|
||||
g_idx.to(device=orig_device),
|
||||
sort_indices.to(device=orig_device),
|
||||
)
|
||||
|
||||
|
||||
def gptq_pack(
|
||||
q_w: torch.Tensor,
|
||||
num_bits: int,
|
||||
size_k: int,
|
||||
size_n: int,
|
||||
):
|
||||
assert q_w.shape == (size_k, size_n)
|
||||
|
||||
pack_factor = get_pack_factor(num_bits)
|
||||
assert size_k % pack_factor == 0
|
||||
|
||||
orig_device = q_w.device
|
||||
|
||||
q_w = q_w.cpu().numpy().astype(numpy.uint32)
|
||||
|
||||
q_res = numpy.zeros((size_k // pack_factor, size_n), dtype=numpy.uint32)
|
||||
|
||||
for i in range(pack_factor):
|
||||
q_res |= q_w[i::pack_factor, :] << num_bits * i
|
||||
|
||||
q_res = torch.from_numpy(q_res.astype(numpy.int32)).to(orig_device)
|
||||
return q_res
|
||||
|
||||
def gptq_unpack(
|
||||
q_res: torch.Tensor,
|
||||
num_bits: int,
|
||||
size_k: int,
|
||||
size_n: int,
|
||||
):
|
||||
pack_factor = 32 // num_bits
|
||||
assert size_k % pack_factor == 0
|
||||
|
||||
orig_device = q_res.device
|
||||
|
||||
q_res = q_res.cpu().numpy()
|
||||
|
||||
q_w = numpy.zeros((size_k, size_n), dtype=numpy.uint32)
|
||||
|
||||
for i in range(pack_factor):
|
||||
q_w[i::pack_factor, :] = (q_res >> (num_bits * i)) & ((1 << num_bits) - 1)
|
||||
|
||||
q_w = torch.from_numpy(q_w.astype(numpy.int32)).to(orig_device)
|
||||
return q_w
|
||||
@@ -0,0 +1,415 @@
|
||||
cmake_minimum_required(VERSION 3.16)
|
||||
project(cpuinfer_ext VERSION 0.1.0)
|
||||
|
||||
|
||||
set(CMAKE_CXX_STANDARD 17)
|
||||
|
||||
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -ffast-math -fopenmp")
|
||||
add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=${_GLIBCXX_USE_CXX11_ABI})
|
||||
set(CMAKE_BUILD_TYPE "Release")
|
||||
# set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -ffast-math -fopenmp")
|
||||
# set(CMAKE_BUILD_TYPE "Debug")
|
||||
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
|
||||
|
||||
|
||||
include(CheckCXXCompilerFlag)
|
||||
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
|
||||
|
||||
|
||||
option(LLAMA_NATIVE "llama: enable -march=native flag" ON)
|
||||
|
||||
# instruction set specific
|
||||
if (LLAMA_NATIVE)
|
||||
set(INS_ENB OFF)
|
||||
else()
|
||||
set(INS_ENB ON)
|
||||
endif()
|
||||
|
||||
option(LLAMA_AVX "llama: enable AVX" OFF)
|
||||
option(LLAMA_AVX2 "llama: enable AVX2" OFF)
|
||||
option(LLAMA_AVX512 "llama: enable AVX512" OFF)
|
||||
option(LLAMA_AVX512_VBMI "llama: enable AVX512-VBMI" OFF)
|
||||
option(LLAMA_AVX512_VNNI "llama: enable AVX512-VNNI" OFF)
|
||||
option(LLAMA_AVX512_BF16 "llama: enable AVX512-BF16" OFF)
|
||||
option(LLAMA_FMA "llama: enable FMA" OFF)
|
||||
# in MSVC F16C is implied with AVX2/AVX512
|
||||
if (NOT MSVC)
|
||||
option(LLAMA_F16C "llama: enable F16C" OFF)
|
||||
endif()
|
||||
option(LLAMA_AVX512_FANCY_SIMD "llama: enable AVX512-VL, AVX512-BW, AVX512-DQ, AVX512-VNNI" OFF)
|
||||
option(KTRANSFORMERS_USE_CUDA "ktransformers: use CUDA" ON)
|
||||
option(KTRANSFORMERS_USE_MUSA "ktransformers: use MUSA" OFF)
|
||||
option(KTRANSFORMERS_USE_ROCM "ktransformers: use ROCM" OFF)
|
||||
option(KTRANSFORMERS_USE_XPU "ktransformers: use XPU" OFF)
|
||||
option(KTRANSFORMERS_USE_NPU "ktransformers: use NPU" OFF)
|
||||
|
||||
if(KTRANSFORMERS_USE_NPU)
|
||||
add_definitions(-DKTRANSFORMERS_USE_NPU=1)
|
||||
endif()
|
||||
|
||||
# Architecture specific
|
||||
# TODO: probably these flags need to be tweaked on some architectures
|
||||
# feel free to update the Makefile for your architecture and send a pull request or issue
|
||||
message(STATUS "CMAKE_SYSTEM_PROCESSOR: ${CMAKE_SYSTEM_PROCESSOR}")
|
||||
if (MSVC)
|
||||
string(TOLOWER "${CMAKE_GENERATOR_PLATFORM}" CMAKE_GENERATOR_PLATFORM_LWR)
|
||||
message(STATUS "CMAKE_GENERATOR_PLATFORM: ${CMAKE_GENERATOR_PLATFORM}")
|
||||
else ()
|
||||
set(CMAKE_GENERATOR_PLATFORM_LWR "")
|
||||
endif ()
|
||||
|
||||
if (NOT MSVC)
|
||||
if (LLAMA_STATIC)
|
||||
add_link_options(-static)
|
||||
if (MINGW)
|
||||
add_link_options(-static-libgcc -static-libstdc++)
|
||||
endif()
|
||||
endif()
|
||||
if (LLAMA_GPROF)
|
||||
add_compile_options(-pg)
|
||||
endif()
|
||||
endif()
|
||||
|
||||
set(ARCH_FLAGS "")
|
||||
|
||||
if (CMAKE_OSX_ARCHITECTURES STREQUAL "arm64" OR CMAKE_GENERATOR_PLATFORM_LWR STREQUAL "arm64" OR
|
||||
(NOT CMAKE_OSX_ARCHITECTURES AND NOT CMAKE_GENERATOR_PLATFORM_LWR AND
|
||||
CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm.*|ARM64)$"))
|
||||
message(STATUS "ARM detected")
|
||||
if (MSVC)
|
||||
add_compile_definitions(__aarch64__) # MSVC defines _M_ARM64 instead
|
||||
add_compile_definitions(__ARM_NEON)
|
||||
add_compile_definitions(__ARM_FEATURE_FMA)
|
||||
|
||||
set(CMAKE_REQUIRED_FLAGS_PREV ${CMAKE_REQUIRED_FLAGS})
|
||||
string(JOIN " " CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS} "/arch:armv8.2")
|
||||
check_cxx_source_compiles("#include <arm_neon.h>\nint main() { int8x16_t _a, _b; int32x4_t _s = vdotq_s32(_s, _a, _b); return 0; }" GGML_COMPILER_SUPPORT_DOTPROD)
|
||||
if (GGML_COMPILER_SUPPORT_DOTPROD)
|
||||
add_compile_definitions(__ARM_FEATURE_DOTPROD)
|
||||
endif ()
|
||||
check_cxx_source_compiles("#include <arm_neon.h>\nint main() { float16_t _a; float16x8_t _s = vdupq_n_f16(_a); return 0; }" GGML_COMPILER_SUPPORT_FP16_VECTOR_ARITHMETIC)
|
||||
if (GGML_COMPILER_SUPPORT_FP16_VECTOR_ARITHMETIC)
|
||||
add_compile_definitions(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC)
|
||||
endif ()
|
||||
set(CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS_PREV})
|
||||
else()
|
||||
if(KTRANSFORMERS_USE_NPU)
|
||||
list(APPEND ARCH_FLAGS -march=armv8.2-a+fp16+fp16fml+dotprod -lnuma)
|
||||
endif()
|
||||
check_cxx_compiler_flag(-mfp16-format=ieee COMPILER_SUPPORTS_FP16_FORMAT_I3E)
|
||||
if (NOT "${COMPILER_SUPPORTS_FP16_FORMAT_I3E}" STREQUAL "")
|
||||
list(APPEND ARCH_FLAGS -mfp16-format=ieee)
|
||||
endif()
|
||||
if (${CMAKE_SYSTEM_PROCESSOR} MATCHES "armv6")
|
||||
# Raspberry Pi 1, Zero
|
||||
list(APPEND ARCH_FLAGS -mfpu=neon-fp-armv8 -mno-unaligned-access)
|
||||
endif()
|
||||
if (${CMAKE_SYSTEM_PROCESSOR} MATCHES "armv7")
|
||||
if ("${CMAKE_SYSTEM_NAME}" STREQUAL "Android")
|
||||
# Android armeabi-v7a
|
||||
list(APPEND ARCH_FLAGS -mfpu=neon-vfpv4 -mno-unaligned-access -funsafe-math-optimizations)
|
||||
else()
|
||||
# Raspberry Pi 2
|
||||
list(APPEND ARCH_FLAGS -mfpu=neon-fp-armv8 -mno-unaligned-access -funsafe-math-optimizations)
|
||||
endif()
|
||||
endif()
|
||||
if (${CMAKE_SYSTEM_PROCESSOR} MATCHES "armv8")
|
||||
# Android arm64-v8a
|
||||
# Raspberry Pi 3, 4, Zero 2 (32-bit)
|
||||
list(APPEND ARCH_FLAGS -mno-unaligned-access)
|
||||
endif()
|
||||
endif()
|
||||
elseif (CMAKE_OSX_ARCHITECTURES STREQUAL "x86_64" OR CMAKE_GENERATOR_PLATFORM_LWR MATCHES "^(x86_64|i686|amd64|x64|win32)$" OR
|
||||
(NOT CMAKE_OSX_ARCHITECTURES AND NOT CMAKE_GENERATOR_PLATFORM_LWR AND
|
||||
CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|i686|AMD64)$"))
|
||||
message(STATUS "x86 detected")
|
||||
if(NOT KTRANSFORMERS_USE_NPU)
|
||||
set(HOST_IS_X86 TRUE)
|
||||
set(HAS_AVX512 TRUE)
|
||||
set(__HAS_AMX__ TRUE)
|
||||
add_compile_definitions(__x86_64__)
|
||||
# check AVX512
|
||||
execute_process(
|
||||
COMMAND lscpu
|
||||
OUTPUT_VARIABLE LSCPU_OUTPUT
|
||||
OUTPUT_STRIP_TRAILING_WHITESPACE
|
||||
)
|
||||
# message(STATUS "LSCPU_OUTPUT: ${LSCPU_OUTPUT}")
|
||||
|
||||
string(FIND "${LSCPU_OUTPUT}" "avx512" COMPILER_SUPPORTS_AVX512F)
|
||||
|
||||
if (COMPILER_SUPPORTS_AVX512F GREATER -1)
|
||||
message(STATUS "Compiler and CPU support AVX512F (tested by compiling a program)")
|
||||
add_compile_definitions(__HAS_AVX512F__)
|
||||
else()
|
||||
message(STATUS "Compiler and/or CPU do NOT support AVX512F")
|
||||
set(HAS_AVX512 False)
|
||||
endif()
|
||||
|
||||
# check AMX
|
||||
string(FIND "${LSCPU_OUTPUT}" "amx" COMPILER_SUPPORTS_AMX)
|
||||
|
||||
if(COMPILER_SUPPORTS_AMX GREATER -1)
|
||||
message(STATUS "Compiler supports AMX")
|
||||
add_compile_definitions(__HAS_AMX__)
|
||||
else()
|
||||
message(STATUS "Compiler does NOT support AMX")
|
||||
endif()
|
||||
endif()
|
||||
if (MSVC)
|
||||
# instruction set detection for MSVC only
|
||||
if (LLAMA_NATIVE)
|
||||
include(cmake/FindSIMD.cmake)
|
||||
endif ()
|
||||
if (LLAMA_AVX512)
|
||||
list(APPEND ARCH_FLAGS /arch:AVX512)
|
||||
# MSVC has no compile-time flags enabling specific
|
||||
# AVX512 extensions, neither it defines the
|
||||
# macros corresponding to the extensions.
|
||||
# Do it manually.
|
||||
if (LLAMA_AVX512_VBMI)
|
||||
add_compile_definitions($<$<COMPILE_LANGUAGE:C>:__AVX512VBMI__>)
|
||||
add_compile_definitions($<$<COMPILE_LANGUAGE:CXX>:__AVX512VBMI__>)
|
||||
endif()
|
||||
if (LLAMA_AVX512_VNNI)
|
||||
add_compile_definitions($<$<COMPILE_LANGUAGE:C>:__AVX512VNNI__>)
|
||||
add_compile_definitions($<$<COMPILE_LANGUAGE:CXX>:__AVX512VNNI__>)
|
||||
endif()
|
||||
if (LLAMA_AVX512_FANCY_SIMD)
|
||||
add_compile_definitions($<$<COMPILE_LANGUAGE:C>:__AVX512VL__>)
|
||||
add_compile_definitions($<$<COMPILE_LANGUAGE:CXX>:__AVX512VL__>)
|
||||
add_compile_definitions($<$<COMPILE_LANGUAGE:C>:__AVX512BW__>)
|
||||
add_compile_definitions($<$<COMPILE_LANGUAGE:CXX>:__AVX512BW__>)
|
||||
add_compile_definitions($<$<COMPILE_LANGUAGE:C>:__AVX512DQ__>)
|
||||
add_compile_definitions($<$<COMPILE_LANGUAGE:CXX>:__AVX512DQ__>)
|
||||
add_compile_definitions($<$<COMPILE_LANGUAGE:C>:__AVX512VNNI__>)
|
||||
add_compile_definitions($<$<COMPILE_LANGUAGE:CXX>:__AVX512VNNI__>)
|
||||
endif()
|
||||
if (LLAMA_AVX512_BF16)
|
||||
add_compile_definitions($<$<COMPILE_LANGUAGE:C>:__AVX512BF16__>)
|
||||
add_compile_definitions($<$<COMPILE_LANGUAGE:CXX>:__AVX512BF16__>)
|
||||
endif()
|
||||
elseif (LLAMA_AVX2)
|
||||
list(APPEND ARCH_FLAGS /arch:AVX2)
|
||||
elseif (LLAMA_AVX)
|
||||
list(APPEND ARCH_FLAGS /arch:AVX)
|
||||
endif()
|
||||
else()
|
||||
if (LLAMA_NATIVE)
|
||||
list(APPEND ARCH_FLAGS -mfma -mavx -mavx2)
|
||||
list(APPEND ARCH_FLAGS -march=native)
|
||||
endif()
|
||||
if (LLAMA_F16C)
|
||||
list(APPEND ARCH_FLAGS -mf16c)
|
||||
endif()
|
||||
if (LLAMA_FMA)
|
||||
list(APPEND ARCH_FLAGS -mfma)
|
||||
endif()
|
||||
if (LLAMA_AVX)
|
||||
list(APPEND ARCH_FLAGS -mavx)
|
||||
endif()
|
||||
if (LLAMA_AVX2)
|
||||
list(APPEND ARCH_FLAGS -mavx2)
|
||||
endif()
|
||||
if (LLAMA_AVX512)
|
||||
list(APPEND ARCH_FLAGS -mavx512f)
|
||||
list(APPEND ARCH_FLAGS -mavx512bw)
|
||||
endif()
|
||||
if (LLAMA_AVX512_VBMI)
|
||||
list(APPEND ARCH_FLAGS -mavx512vbmi)
|
||||
endif()
|
||||
if (LLAMA_AVX512_VNNI)
|
||||
list(APPEND ARCH_FLAGS -mavx512vnni)
|
||||
endif()
|
||||
if (LLAMA_AVX512_FANCY_SIMD)
|
||||
message(STATUS "AVX512-VL, AVX512-BW, AVX512-DQ, AVX512-VNNI enabled")
|
||||
list(APPEND ARCH_FLAGS -mavx512vl)
|
||||
list(APPEND ARCH_FLAGS -mavx512bw)
|
||||
list(APPEND ARCH_FLAGS -mavx512dq)
|
||||
list(APPEND ARCH_FLAGS -mavx512vnni)
|
||||
list(APPEND ARCH_FLAGS -mavx512vpopcntdq)
|
||||
endif()
|
||||
if (LLAMA_AVX512_BF16)
|
||||
list(APPEND ARCH_FLAGS -mavx512bf16)
|
||||
endif()
|
||||
endif()
|
||||
elseif (${CMAKE_SYSTEM_PROCESSOR} MATCHES "ppc64")
|
||||
message(STATUS "PowerPC detected")
|
||||
if (${CMAKE_SYSTEM_PROCESSOR} MATCHES "ppc64le")
|
||||
list(APPEND ARCH_FLAGS -mcpu=powerpc64le)
|
||||
else()
|
||||
list(APPEND ARCH_FLAGS -mcpu=native -mtune=native)
|
||||
#TODO: Add targets for Power8/Power9 (Altivec/VSX) and Power10(MMA) and query for big endian systems (ppc64/le/be)
|
||||
endif()
|
||||
else()
|
||||
message(STATUS "Unknown architecture")
|
||||
endif()
|
||||
|
||||
# message(STATUS "CUDAToolkit_ROOT:${CUDAToolkit_ROOT}")
|
||||
# find_package(FindCUDAToolkit REQUIRED)
|
||||
# if(CUDAToolkit_FOUND)
|
||||
# message(STATUS "Found CUDA cudart lib at:${CUDAToolkit_LIBRARY_DIR}")
|
||||
# else()
|
||||
# message(STATUS "Can't found CUDA lib")
|
||||
# endif()
|
||||
|
||||
if (NOT EXISTS $ENV{ROCM_PATH})
|
||||
if (NOT EXISTS /opt/rocm)
|
||||
set(ROCM_PATH /usr)
|
||||
else()
|
||||
set(ROCM_PATH /opt/rocm)
|
||||
endif()
|
||||
else()
|
||||
set(ROCM_PATH $ENV{ROCM_PATH})
|
||||
endif()
|
||||
|
||||
list(APPEND CMAKE_PREFIX_PATH ${ROCM_PATH})
|
||||
list(APPEND CMAKE_PREFIX_PATH "${ROCM_PATH}/lib64/cmake")
|
||||
|
||||
if (NOT EXISTS $ENV{MUSA_PATH})
|
||||
if (NOT EXISTS /opt/musa)
|
||||
set(MUSA_PATH /usr/local/musa)
|
||||
else()
|
||||
set(MUSA_PATH /opt/musa)
|
||||
endif()
|
||||
else()
|
||||
set(MUSA_PATH $ENV{MUSA_PATH})
|
||||
endif()
|
||||
|
||||
list(APPEND CMAKE_MODULE_PATH "${MUSA_PATH}/cmake")
|
||||
|
||||
add_compile_options("$<$<COMPILE_LANGUAGE:CXX>:${ARCH_FLAGS}>")
|
||||
add_compile_options("$<$<COMPILE_LANGUAGE:C>:${ARCH_FLAGS}>")
|
||||
|
||||
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/pybind11 ${CMAKE_CURRENT_BINARY_DIR}/third_party/pybind11)
|
||||
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/llama.cpp ${CMAKE_CURRENT_BINARY_DIR}/third_party/llama.cpp)
|
||||
|
||||
include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../../third_party)
|
||||
if (WIN32)
|
||||
include_directories("$ENV{CUDA_PATH}/include")
|
||||
add_compile_definitions(KTRANSFORMERS_USE_CUDA=1)
|
||||
elseif (UNIX)
|
||||
if (KTRANSFORMERS_USE_ROCM)
|
||||
find_package(HIP REQUIRED)
|
||||
if(HIP_FOUND)
|
||||
include_directories("${HIP_INCLUDE_DIRS}")
|
||||
add_compile_definitions(KTRANSFORMERS_USE_ROCM=1)
|
||||
endif()
|
||||
elseif (KTRANSFORMERS_USE_MUSA)
|
||||
if (NOT EXISTS $ENV{MUSA_PATH})
|
||||
if (NOT EXISTS /opt/musa)
|
||||
set(MUSA_PATH /usr/local/musa)
|
||||
else()
|
||||
set(MUSA_PATH /opt/musa)
|
||||
endif()
|
||||
else()
|
||||
set(MUSA_PATH $ENV{MUSA_PATH})
|
||||
endif()
|
||||
|
||||
list(APPEND CMAKE_MODULE_PATH "${MUSA_PATH}/cmake")
|
||||
|
||||
find_package(MUSAToolkit)
|
||||
if (MUSAToolkit_FOUND)
|
||||
message(STATUS "MUSA Toolkit found")
|
||||
add_compile_definitions(KTRANSFORMERS_USE_MUSA=1)
|
||||
endif()
|
||||
elseif (KTRANSFORMERS_USE_XPU)
|
||||
add_compile_definitions(KTRANSFORMERS_USE_XPU=1)
|
||||
elseif (KTRANSFORMERS_USE_CUDA)
|
||||
find_package(CUDA REQUIRED)
|
||||
include_directories("${CUDA_INCLUDE_DIRS}")
|
||||
include(CheckLanguage)
|
||||
check_language(CUDA)
|
||||
if(CMAKE_CUDA_COMPILER)
|
||||
message(STATUS "CUDA detected")
|
||||
find_package(CUDAToolkit REQUIRED)
|
||||
include_directories(${CUDAToolkit_INCLUDE_DIRS})
|
||||
endif()
|
||||
message(STATUS "enabling CUDA")
|
||||
enable_language(CUDA)
|
||||
add_compile_definitions(KTRANSFORMERS_USE_CUDA=1)
|
||||
endif()
|
||||
endif()
|
||||
|
||||
aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR} SOURCE_DIR1)
|
||||
aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/cpu_backend SOURCE_DIR2)
|
||||
aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/operators/llamafile SOURCE_DIR3)
|
||||
# aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/llamafile SOURCE_DIR4)
|
||||
file(GLOB LLAMAFILE_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/llamafile/*.cpp")
|
||||
list(REMOVE_ITEM LLAMAFILE_SOURCES
|
||||
"${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/llamafile/sgemm_arm.cpp"
|
||||
"${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/llamafile/sgemm_x86.cpp"
|
||||
)
|
||||
set(SOURCE_DIR4 ${LLAMAFILE_SOURCES})
|
||||
aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/operators/kvcache SOURCE_DIR5)
|
||||
|
||||
if (HOST_IS_X86 AND HAS_AVX512 AND __HAS_AMX__)
|
||||
aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/operators/amx SOURCE_DIR6)
|
||||
endif()
|
||||
|
||||
|
||||
set(ALL_SOURCES ${SOURCE_DIR1} ${SOURCE_DIR2} ${SOURCE_DIR3} ${SOURCE_DIR4} ${SOURCE_DIR5} ${SOURCE_DIR6})
|
||||
|
||||
file(GLOB_RECURSE FMT_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/*.hpp" "${CMAKE_CURRENT_SOURCE_DIR}/*.h")
|
||||
|
||||
add_custom_target(
|
||||
format
|
||||
COMMAND clang-format
|
||||
-i
|
||||
-style=file
|
||||
${FMT_SOURCES}
|
||||
COMMENT "Running clang-format on all source files"
|
||||
)
|
||||
|
||||
|
||||
add_library(llamafile STATIC ${SOURCE_DIR4})
|
||||
|
||||
message(STATUS "CMAKE_CXX_FLAGS: ${CMAKE_CXX_FLAGS}")
|
||||
message(STATUS "ARCH_FLAGS: ${ARCH_FLAGS}")
|
||||
pybind11_add_module(${PROJECT_NAME} MODULE ${ALL_SOURCES})
|
||||
target_link_libraries(${PROJECT_NAME} PRIVATE llama)
|
||||
|
||||
|
||||
if(WIN32)
|
||||
target_link_libraries(${PROJECT_NAME} PRIVATE "$ENV{CUDA_PATH}/lib/x64/cudart.lib")#CUDA::cudart
|
||||
elseif(UNIX)
|
||||
if (KTRANSFORMERS_USE_ROCM)
|
||||
add_compile_definitions(USE_HIP=1)
|
||||
target_link_libraries(${PROJECT_NAME} PRIVATE "${ROCM_PATH}/lib/libamdhip64.so")
|
||||
message(STATUS "Building for HIP")
|
||||
elseif(KTRANSFORMERS_USE_MUSA)
|
||||
target_link_libraries(${PROJECT_NAME} PRIVATE MUSA::musart)
|
||||
elseif(KTRANSFORMERS_USE_XPU)
|
||||
elseif(KTRANSFORMERS_USE_CUDA AND NOT KTRANSFORMERS_USE_MUSA)
|
||||
target_link_libraries(${PROJECT_NAME} PRIVATE "${CUDAToolkit_LIBRARY_DIR}/libcudart.so")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# Define the USE_NUMA option
|
||||
option(USE_NUMA "Disable NUMA support" OFF)
|
||||
|
||||
# Check if the USE_NUMA environment variable is set
|
||||
if(DEFINED ENV{USE_NUMA})
|
||||
set(USE_NUMA ON)
|
||||
endif()
|
||||
|
||||
if(USE_NUMA)
|
||||
message(STATUS "NUMA support is enabled")
|
||||
else()
|
||||
message(STATUS "NUMA support is disabled")
|
||||
endif()
|
||||
|
||||
find_library(NUMA_LIBRARY NAMES numa)
|
||||
|
||||
if(NUMA_LIBRARY AND USE_NUMA)
|
||||
message(STATUS "NUMA library found: ${NUMA_LIBRARY} - enabling NUMA support")
|
||||
target_link_libraries(${PROJECT_NAME} PRIVATE ${NUMA_LIBRARY})
|
||||
target_compile_definitions(${PROJECT_NAME} PRIVATE USE_NUMA)
|
||||
else()
|
||||
if(USE_NUMA)
|
||||
message(FATAL_ERROR "NUMA library not found - maybe sudo apt install libnuma-dev")
|
||||
else()
|
||||
message(STATUS "NUMA library not found or user not set USE_NUMA - disabling NUMA support")
|
||||
endif()
|
||||
endif()
|
||||
@@ -0,0 +1,178 @@
|
||||
#!/usr/bin/env python
|
||||
# coding=utf-8
|
||||
"""
|
||||
Description :
|
||||
Author : Jianwei Dong
|
||||
Date : 2024-08-28 10:32:05
|
||||
Version : 1.0.0
|
||||
LastEditors : Jianwei Dong
|
||||
LastEditTime : 2024-08-28 10:32:05
|
||||
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
"""
|
||||
import os, sys
|
||||
import time
|
||||
|
||||
sys.path.append(os.path.dirname(__file__) + "/../build")
|
||||
import cpuinfer_ext
|
||||
import torch
|
||||
|
||||
layer_num = 10
|
||||
kv_head_num = 8
|
||||
q_head_num = 32
|
||||
head_dim = 128
|
||||
block_len = 128
|
||||
anchor_num = 1
|
||||
|
||||
anchor_type = cpuinfer_ext.kvcache.AnchorType.DYNAMIC
|
||||
kv_type = cpuinfer_ext.kvcache.ggml_type.FP16
|
||||
retrieval_type = cpuinfer_ext.kvcache.RetrievalType.LAYER
|
||||
layer_step: int = 1
|
||||
token_step: int = 1
|
||||
layer_offset: int = 0
|
||||
max_thread_num: int = 64
|
||||
max_batch_size: int = 1
|
||||
max_block_num: int = 1024
|
||||
CPUInfer = cpuinfer_ext.CPUInfer(max_thread_num)
|
||||
|
||||
warm_up_iter = 1000
|
||||
test_iter = 10000
|
||||
|
||||
|
||||
def bench_linear(cache_seqlen: int):
|
||||
with torch.inference_mode(mode=True):
|
||||
cache_seqlens = torch.tensor([cache_seqlen], dtype=torch.int32, device="cpu")
|
||||
seqlens_zero = torch.zeros((1,), dtype=torch.int32, device="cpu")
|
||||
|
||||
config = cpuinfer_ext.kvcache.KVCacheConfig(
|
||||
layer_num,
|
||||
kv_head_num,
|
||||
q_head_num,
|
||||
head_dim,
|
||||
block_len,
|
||||
anchor_num,
|
||||
anchor_type,
|
||||
kv_type,
|
||||
retrieval_type,
|
||||
layer_step,
|
||||
token_step,
|
||||
layer_offset,
|
||||
max_block_num,
|
||||
max_batch_size,
|
||||
max_thread_num,
|
||||
)
|
||||
local_kvcache = cpuinfer_ext.kvcache.KVCache(config)
|
||||
block_table = (
|
||||
torch.arange(max_block_num, dtype=torch.int32, device="cpu")
|
||||
.contiguous()
|
||||
.view(1, -1)
|
||||
)
|
||||
|
||||
for layer_idx in range(layer_num):
|
||||
k_cache = torch.randn(
|
||||
(1, cache_seqlen, kv_head_num, head_dim),
|
||||
dtype=torch.float16,
|
||||
device="cpu",
|
||||
).contiguous()
|
||||
v_cache = torch.randn(
|
||||
(1, cache_seqlen, kv_head_num, head_dim),
|
||||
dtype=torch.float16,
|
||||
device="cpu",
|
||||
).contiguous()
|
||||
|
||||
CPUInfer.submit(
|
||||
local_kvcache.update_kvcache_fp16(
|
||||
k_cache.data_ptr(),
|
||||
v_cache.data_ptr(),
|
||||
layer_idx,
|
||||
block_table.data_ptr(),
|
||||
1,
|
||||
max_block_num,
|
||||
seqlens_zero.data_ptr(),
|
||||
cache_seqlen,
|
||||
)
|
||||
)
|
||||
CPUInfer.sync()
|
||||
|
||||
input = torch.randn(
|
||||
(1, 1, q_head_num, head_dim), dtype=torch.float16, device="cpu"
|
||||
).contiguous()
|
||||
output = torch.empty(
|
||||
(1, 1, q_head_num, head_dim), dtype=torch.float16, device="cpu"
|
||||
).contiguous()
|
||||
|
||||
# attn_lse: (bsz, q_len, q_head_num)
|
||||
attn_lse = torch.empty(
|
||||
(1, 1, q_head_num), dtype=torch.float32, device="cpu"
|
||||
).contiguous()
|
||||
input = input / 100
|
||||
|
||||
# warm up
|
||||
for i in range(warm_up_iter):
|
||||
CPUInfer.submit(
|
||||
local_kvcache.attn(
|
||||
input.data_ptr(),
|
||||
output.data_ptr(),
|
||||
attn_lse.data_ptr(),
|
||||
i % layer_num,
|
||||
0,
|
||||
1,
|
||||
1,
|
||||
max_block_num,
|
||||
block_table.data_ptr(),
|
||||
cache_seqlens.data_ptr(),
|
||||
-1,
|
||||
-1,
|
||||
-1,
|
||||
)
|
||||
)
|
||||
CPUInfer.sync()
|
||||
|
||||
# test
|
||||
start = time.perf_counter()
|
||||
for i in range(test_iter):
|
||||
CPUInfer.submit(
|
||||
local_kvcache.attn(
|
||||
input.data_ptr(),
|
||||
output.data_ptr(),
|
||||
attn_lse.data_ptr(),
|
||||
i % layer_num,
|
||||
0,
|
||||
1,
|
||||
1,
|
||||
max_block_num,
|
||||
block_table.data_ptr(),
|
||||
cache_seqlens.data_ptr(),
|
||||
-1,
|
||||
-1,
|
||||
-1,
|
||||
)
|
||||
)
|
||||
CPUInfer.sync()
|
||||
end = time.perf_counter()
|
||||
total_time = end - start
|
||||
print("cache sequence length: ", cache_seqlen)
|
||||
print("Time(s): ", total_time)
|
||||
print("Iteration: ", test_iter)
|
||||
print("Time(us) per iteration: ", total_time / test_iter * 1000000)
|
||||
print(
|
||||
"Bandwidth: ",
|
||||
cache_seqlen
|
||||
* kv_head_num
|
||||
* head_dim
|
||||
* 2
|
||||
* 2
|
||||
* test_iter
|
||||
/ total_time
|
||||
/ 1000
|
||||
/ 1000
|
||||
/ 1000,
|
||||
"GB/s",
|
||||
)
|
||||
print("")
|
||||
|
||||
|
||||
bench_linear(1024)
|
||||
bench_linear(4096)
|
||||
bench_linear(16384)
|
||||
bench_linear(32768)
|
||||
bench_linear(65536)
|
||||
@@ -0,0 +1,94 @@
|
||||
#!/usr/bin/env python
|
||||
# coding=utf-8
|
||||
"""
|
||||
Description :
|
||||
Author : Jianwei Dong
|
||||
Date : 2024-08-28 10:32:05
|
||||
Version : 1.0.0
|
||||
LastEditors : Jianwei Dong
|
||||
LastEditTime : 2024-08-28 10:32:05
|
||||
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
"""
|
||||
import os, sys
|
||||
import time
|
||||
|
||||
sys.path.append(os.path.dirname(__file__) + "/../build")
|
||||
import cpuinfer_ext
|
||||
import torch
|
||||
|
||||
layer_num = 10
|
||||
kv_head_num = 8
|
||||
q_head_num = 32
|
||||
head_dim = 128
|
||||
block_len = 128
|
||||
anchor_num = 1
|
||||
warm_up_iter = 1000
|
||||
test_iter = 10000
|
||||
|
||||
|
||||
def bench_linear(cache_seqlen: int, device):
|
||||
with torch.inference_mode(mode=True):
|
||||
|
||||
kvcaches = []
|
||||
|
||||
for layer_idx in range(layer_num):
|
||||
k_cache = torch.randn(
|
||||
(1, 32, cache_seqlen, head_dim),
|
||||
dtype=torch.float16,
|
||||
device=device,
|
||||
).contiguous()
|
||||
v_cache = torch.randn(
|
||||
(1, 32, cache_seqlen, head_dim),
|
||||
dtype=torch.float16,
|
||||
device=device,
|
||||
).contiguous()
|
||||
|
||||
kvcaches.append((k_cache, v_cache))
|
||||
|
||||
input = torch.randn(
|
||||
(1, q_head_num, 1, head_dim), dtype=torch.float16, device=device
|
||||
).contiguous()
|
||||
input = input / 100
|
||||
|
||||
# warm up
|
||||
for i in range(warm_up_iter):
|
||||
k_cache = kvcaches[i % layer_num][0]
|
||||
v_cache = kvcaches[i % layer_num][1]
|
||||
torch.nn.functional.scaled_dot_product_attention(input, k_cache, v_cache)
|
||||
|
||||
# test
|
||||
start = time.perf_counter()
|
||||
for i in range(test_iter):
|
||||
k_cache = kvcaches[i % layer_num][0]
|
||||
v_cache = kvcaches[i % layer_num][1]
|
||||
torch.nn.functional.scaled_dot_product_attention(input, k_cache, v_cache)
|
||||
end = time.perf_counter()
|
||||
total_time = end - start
|
||||
print("cache sequence length: ", cache_seqlen)
|
||||
print("Time(s): ", total_time)
|
||||
print("Iteration: ", test_iter)
|
||||
print("Time(us) per iteration: ", total_time / test_iter * 1000000)
|
||||
print(
|
||||
"Bandwidth: ",
|
||||
cache_seqlen
|
||||
* q_head_num
|
||||
* head_dim
|
||||
* 2
|
||||
* 2
|
||||
* test_iter
|
||||
/ total_time
|
||||
/ 1000
|
||||
/ 1000
|
||||
/ 1000,
|
||||
"GB/s",
|
||||
)
|
||||
print("")
|
||||
|
||||
|
||||
bench_linear(1024, "cpu")
|
||||
bench_linear(4096, "cpu")
|
||||
bench_linear(1024, "cuda")
|
||||
bench_linear(4096, "cuda")
|
||||
bench_linear(16384, "cuda")
|
||||
bench_linear(32768, "cuda")
|
||||
bench_linear(65536, "cuda")
|
||||
@@ -0,0 +1,121 @@
|
||||
#!/usr/bin/env python
|
||||
# coding=utf-8
|
||||
'''
|
||||
Description :
|
||||
Author : chenht2022
|
||||
Date : 2024-07-25 10:31:59
|
||||
Version : 1.0.0
|
||||
LastEditors : chenht2022
|
||||
LastEditTime : 2024-08-06 10:35:35
|
||||
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
'''
|
||||
import os, sys
|
||||
import time
|
||||
sys.path.append(os.path.dirname(__file__) + '/../build')
|
||||
import cpuinfer_ext
|
||||
import torch
|
||||
|
||||
input_size = 16384
|
||||
output_size = 5120
|
||||
stride = 16
|
||||
group_max_len = 1024
|
||||
layer_num = 10
|
||||
qlen = 1
|
||||
CPUInfer = cpuinfer_ext.CPUInfer(64)
|
||||
warm_up_iter = 1000
|
||||
test_iter = 10000
|
||||
|
||||
def bench_linear(quant_mode: str):
|
||||
with torch.inference_mode(mode=True):
|
||||
|
||||
hidden_type = 30 # ggml_type::GGML_TYPE_BF16
|
||||
if quant_mode == "fp32":
|
||||
proj_type = 0 # ggml_type::GGML_TYPE_F32
|
||||
bytes_per_elem = 4.000000
|
||||
elif quant_mode == "fp16":
|
||||
proj_type = 1 # ggml_type::GGML_TYPE_F16
|
||||
bytes_per_elem = 2.000000
|
||||
elif quant_mode == "bf16":
|
||||
proj_type = 30 # ggml_type::GGML_TYPE_BF16
|
||||
bytes_per_elem = 2.000000
|
||||
elif quant_mode == "q8_0":
|
||||
proj_type = 8 # ggml_type::GGML_TYPE_Q8_0
|
||||
bytes_per_elem = 1.062500
|
||||
elif quant_mode == "q6_k":
|
||||
proj_type = 14 # ggml_type::GGML_TYPE_Q6_K
|
||||
bytes_per_elem = 0.820312
|
||||
elif quant_mode == "q5_k_m":
|
||||
proj_type = 13 # ggml_type::GGML_TYPE_Q5_K
|
||||
bytes_per_elem = 0.687500
|
||||
elif quant_mode == "q4_k_m":
|
||||
proj_type = 12 # ggml_type::GGML_TYPE_Q4_K
|
||||
bytes_per_elem = 0.562500
|
||||
elif quant_mode == "q3_k_m":
|
||||
proj_type = 11 # ggml_type::GGML_TYPE_Q3_K
|
||||
bytes_per_elem = 0.429688
|
||||
elif quant_mode == "q2_k":
|
||||
proj_type = 10 # ggml_type::GGML_TYPE_Q2_K
|
||||
bytes_per_elem = 0.328125
|
||||
elif quant_mode == "iq3_xs":
|
||||
proj_type = 21 # ggml_type::GGML_TYPE_IQ3_S
|
||||
bytes_per_elem = 0.429688
|
||||
elif quant_mode == "iq2_xxs":
|
||||
proj_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS
|
||||
bytes_per_elem = 0.257812
|
||||
else:
|
||||
assert(False)
|
||||
|
||||
linears = []
|
||||
projs = []
|
||||
for _ in range(layer_num):
|
||||
proj = torch.randn((output_size, input_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
config = cpuinfer_ext.linear.LinearConfig(input_size, output_size, stride, group_max_len, proj.data_ptr(), proj_type, hidden_type)
|
||||
linear = cpuinfer_ext.linear.Linear(config)
|
||||
projs.append(proj)
|
||||
linears.append(linear)
|
||||
input = torch.randn((layer_num, qlen, input_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
|
||||
output = torch.empty((layer_num, qlen, output_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
|
||||
|
||||
# warm up
|
||||
for i in range(warm_up_iter):
|
||||
CPUInfer.submit(
|
||||
linears[i % layer_num].forward(
|
||||
qlen,
|
||||
input[i % layer_num].data_ptr(),
|
||||
output[i % layer_num].data_ptr()
|
||||
)
|
||||
)
|
||||
CPUInfer.sync()
|
||||
|
||||
# test
|
||||
start = time.perf_counter()
|
||||
for i in range(test_iter):
|
||||
CPUInfer.submit(
|
||||
linears[i % layer_num].forward(
|
||||
qlen,
|
||||
input[i % layer_num].data_ptr(),
|
||||
output[i % layer_num].data_ptr()
|
||||
)
|
||||
)
|
||||
CPUInfer.sync()
|
||||
end = time.perf_counter()
|
||||
total_time = end - start
|
||||
print('Quant mode: ', quant_mode)
|
||||
print('Time(s): ', total_time)
|
||||
print('Iteration: ', test_iter)
|
||||
print('Time(us) per iteration: ', total_time / test_iter * 1000000)
|
||||
print('Bandwidth: ', input_size * output_size * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s')
|
||||
print('')
|
||||
|
||||
bench_linear("fp32")
|
||||
bench_linear("fp16")
|
||||
bench_linear("bf16")
|
||||
bench_linear("q8_0")
|
||||
bench_linear("q6_k")
|
||||
bench_linear("q5_k_m")
|
||||
bench_linear("q4_k_m")
|
||||
bench_linear("q3_k_m")
|
||||
bench_linear("q2_k")
|
||||
# Not supported on __x86_64__
|
||||
# bench_linear("iq3_xs")
|
||||
# bench_linear("iq2_xxs")
|
||||
@@ -0,0 +1,83 @@
|
||||
#!/usr/bin/env python
|
||||
# coding=utf-8
|
||||
'''
|
||||
Description :
|
||||
Author : chenht2022
|
||||
Date : 2024-07-25 10:31:59
|
||||
Version : 1.0.0
|
||||
LastEditors : chenht2022
|
||||
LastEditTime : 2024-07-25 10:32:48
|
||||
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
'''
|
||||
import os, sys
|
||||
import time
|
||||
import torch
|
||||
import torch.nn.quantized as nnq
|
||||
|
||||
scale, zero_point = 0.1, 0 # Adjust scale and zero_point based on your dataset
|
||||
|
||||
input_size = 16384
|
||||
output_size = 5120
|
||||
layer_num = 10
|
||||
qlen = 1
|
||||
warm_up_iter = 1000
|
||||
test_iter = 10000
|
||||
|
||||
def bench_linear(quant_mode: str):
|
||||
with torch.inference_mode(mode=True):
|
||||
if quant_mode == "fp32":
|
||||
proj_type = torch.float32
|
||||
bytes_per_elem = 4.000000
|
||||
elif quant_mode == "fp16":
|
||||
proj_type = torch.float16
|
||||
bytes_per_elem = 2.000000
|
||||
elif quant_mode == "bf16":
|
||||
proj_type = torch.bfloat16
|
||||
bytes_per_elem = 2.000000
|
||||
elif quant_mode == "qint8":
|
||||
proj_type = torch.qint8
|
||||
bytes_per_elem = 1.000000
|
||||
else:
|
||||
assert(False)
|
||||
|
||||
projs = []
|
||||
for _ in range(layer_num):
|
||||
proj = torch.randn((output_size, input_size), dtype = torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
if quant_mode == "qint8":
|
||||
proj_q = torch.quantize_per_tensor(proj, scale, zero_point, torch.qint8)
|
||||
quantized_layer = nnq.Linear(input_size, output_size)
|
||||
quantized_layer.set_weight_bias(proj_q, None)
|
||||
projs.append(quantized_layer)
|
||||
else:
|
||||
projs.append(proj.to(proj_type))
|
||||
input = torch.randn((layer_num, qlen, input_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
|
||||
|
||||
# warm up
|
||||
for i in range(warm_up_iter):
|
||||
if isinstance(projs[i % layer_num], nnq.Linear):
|
||||
input_q = torch.quantize_per_tensor(input[i % layer_num].to(torch.float32), scale, zero_point, torch.quint8)
|
||||
t_output = projs[i % layer_num](input_q)
|
||||
else:
|
||||
t_output = torch.mm(input[i % layer_num].to(proj_type), projs[i % layer_num].t())
|
||||
|
||||
# test
|
||||
start = time.perf_counter()
|
||||
for i in range(test_iter):
|
||||
if isinstance(projs[i % layer_num], nnq.Linear):
|
||||
input_q = torch.quantize_per_tensor(input[i % layer_num].to(torch.float32), scale, zero_point, torch.quint8)
|
||||
t_output = projs[i % layer_num](input_q)
|
||||
else:
|
||||
t_output = torch.mm(input[i % layer_num].to(proj_type), projs[i % layer_num].t())
|
||||
end = time.perf_counter()
|
||||
total_time = end - start
|
||||
print('Quant mode: ', quant_mode)
|
||||
print('Time(s): ', total_time)
|
||||
print('Iteration: ', test_iter)
|
||||
print('Time(us) per iteration: ', total_time / test_iter * 1000000)
|
||||
print('Bandwidth: ', input_size * output_size * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s')
|
||||
print('')
|
||||
|
||||
bench_linear("fp32")
|
||||
bench_linear("fp16")
|
||||
bench_linear("bf16")
|
||||
bench_linear("qint8")
|
||||
@@ -0,0 +1,150 @@
|
||||
#!/usr/bin/env python
|
||||
# coding=utf-8
|
||||
'''
|
||||
Description :
|
||||
Author : chenht2022
|
||||
Date : 2024-07-16 10:43:18
|
||||
Version : 1.0.0
|
||||
LastEditors : chenht2022
|
||||
LastEditTime : 2024-08-06 10:36:04
|
||||
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
'''
|
||||
import os, sys
|
||||
import time
|
||||
sys.path.append(os.path.dirname(__file__) + '/../build')
|
||||
import cpuinfer_ext
|
||||
import torch
|
||||
|
||||
hidden_size = 5120
|
||||
intermediate_size = 3072
|
||||
stride = 16
|
||||
group_max_len = 1024
|
||||
layer_num = 10
|
||||
qlen = 1
|
||||
CPUInfer = cpuinfer_ext.CPUInfer(64)
|
||||
warm_up_iter = 1000
|
||||
test_iter = 10000
|
||||
|
||||
def bench_mlp(quant_mode: str):
|
||||
with torch.inference_mode(mode=True):
|
||||
|
||||
hidden_type = 30 # ggml_type::GGML_TYPE_BF16
|
||||
if quant_mode == "fp32":
|
||||
gate_type = 0 # ggml_type::GGML_TYPE_F32
|
||||
up_type = 0 # ggml_type::GGML_TYPE_F32
|
||||
down_type = 0 # ggml_type::GGML_TYPE_F32
|
||||
bytes_per_elem = 4.000000
|
||||
elif quant_mode == "fp16":
|
||||
gate_type = 1 # ggml_type::GGML_TYPE_F16
|
||||
up_type = 1 # ggml_type::GGML_TYPE_F16
|
||||
down_type = 1 # ggml_type::GGML_TYPE_F16
|
||||
bytes_per_elem = 2.000000
|
||||
elif quant_mode == "bf16":
|
||||
gate_type = 30 # ggml_type::GGML_TYPE_BF16
|
||||
up_type = 30 # ggml_type::GGML_TYPE_BF16
|
||||
down_type = 30 # ggml_type::GGML_TYPE_BF16
|
||||
bytes_per_elem = 2.000000
|
||||
elif quant_mode == "q8_0":
|
||||
gate_type = 8 # ggml_type::GGML_TYPE_Q8_0
|
||||
up_type = 8 # ggml_type::GGML_TYPE_Q8_0
|
||||
down_type = 8 # ggml_type::GGML_TYPE_Q8_0
|
||||
bytes_per_elem = 1.062500
|
||||
elif quant_mode == "q6_k":
|
||||
gate_type = 14 # ggml_type::GGML_TYPE_Q6_K
|
||||
up_type = 14 # ggml_type::GGML_TYPE_Q6_K
|
||||
down_type = 14 # ggml_type::GGML_TYPE_Q6_K
|
||||
bytes_per_elem = 0.820312
|
||||
elif quant_mode == "q5_k_m":
|
||||
gate_type = 13 # ggml_type::GGML_TYPE_Q5_K
|
||||
up_type = 13 # ggml_type::GGML_TYPE_Q5_K
|
||||
down_type = 14 # ggml_type::GGML_TYPE_Q6_K
|
||||
bytes_per_elem = 0.731771
|
||||
elif quant_mode == "q4_k_m":
|
||||
gate_type = 12 # ggml_type::GGML_TYPE_Q4_K
|
||||
up_type = 12 # ggml_type::GGML_TYPE_Q4_K
|
||||
down_type = 14 # ggml_type::GGML_TYPE_Q6_K
|
||||
bytes_per_elem = 0.648437
|
||||
elif quant_mode == "q3_k_m":
|
||||
gate_type = 11 # ggml_type::GGML_TYPE_Q3_K
|
||||
up_type = 11 # ggml_type::GGML_TYPE_Q3_K
|
||||
down_type = 13 # ggml_type::GGML_TYPE_Q5_K
|
||||
bytes_per_elem = 0.515625
|
||||
elif quant_mode == "q2_k":
|
||||
gate_type = 10 # ggml_type::GGML_TYPE_Q2_K
|
||||
up_type = 10 # ggml_type::GGML_TYPE_Q2_K
|
||||
down_type = 11 # ggml_type::GGML_TYPE_Q3_K
|
||||
bytes_per_elem = 0.328125
|
||||
elif quant_mode == "iq3_xs":
|
||||
gate_type = 21 # ggml_type::GGML_TYPE_IQ3_S
|
||||
up_type = 21 # ggml_type::GGML_TYPE_IQ3_S
|
||||
down_type = 21 # ggml_type::GGML_TYPE_IQ3_S
|
||||
bytes_per_elem = 0.429688
|
||||
elif quant_mode == "iq2_xxs":
|
||||
gate_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS
|
||||
up_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS
|
||||
down_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS
|
||||
bytes_per_elem = 0.257812
|
||||
else:
|
||||
assert(False)
|
||||
|
||||
|
||||
mlps = []
|
||||
gate_projs = []
|
||||
up_projs = []
|
||||
down_projs = []
|
||||
for _ in range(layer_num):
|
||||
gate_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
up_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
down_proj = torch.randn((hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
config = cpuinfer_ext.mlp.MLPConfig(hidden_size, intermediate_size, stride, group_max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), gate_type, up_type, down_type, hidden_type)
|
||||
mlp = cpuinfer_ext.mlp.MLP(config)
|
||||
gate_projs.append(gate_proj)
|
||||
up_projs.append(up_proj)
|
||||
down_projs.append(down_proj)
|
||||
mlps.append(mlp)
|
||||
input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
|
||||
output = torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
|
||||
|
||||
# warm up
|
||||
for i in range(warm_up_iter):
|
||||
CPUInfer.submit(
|
||||
mlps[i % layer_num].forward(
|
||||
qlen,
|
||||
input[i % layer_num].data_ptr(),
|
||||
output[i % layer_num].data_ptr()
|
||||
)
|
||||
)
|
||||
CPUInfer.sync()
|
||||
|
||||
# test
|
||||
start = time.perf_counter()
|
||||
for i in range(test_iter):
|
||||
CPUInfer.submit(
|
||||
mlps[i % layer_num].forward(
|
||||
qlen,
|
||||
input[i % layer_num].data_ptr(),
|
||||
output[i % layer_num].data_ptr()
|
||||
)
|
||||
)
|
||||
CPUInfer.sync()
|
||||
end = time.perf_counter()
|
||||
total_time = end - start
|
||||
print('Quant mode: ', quant_mode)
|
||||
print('Time(s): ', total_time)
|
||||
print('Iteration: ', test_iter)
|
||||
print('Time(us) per iteration: ', total_time / test_iter * 1000000)
|
||||
print('Bandwidth: ', hidden_size * intermediate_size * 3 * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s')
|
||||
print('')
|
||||
|
||||
bench_mlp("fp32")
|
||||
bench_mlp("fp16")
|
||||
bench_mlp("bf16")
|
||||
bench_mlp("q8_0")
|
||||
bench_mlp("q6_k")
|
||||
bench_mlp("q5_k_m")
|
||||
bench_mlp("q4_k_m")
|
||||
bench_mlp("q3_k_m")
|
||||
bench_mlp("q2_k")
|
||||
# Not supported on __x86_64__
|
||||
# bench_linear("iq3_xs")
|
||||
# bench_linear("iq2_xxs")
|
||||
@@ -0,0 +1,110 @@
|
||||
#!/usr/bin/env python
|
||||
# coding=utf-8
|
||||
'''
|
||||
Description :
|
||||
Author : chenht2022
|
||||
Date : 2024-07-16 10:43:18
|
||||
Version : 1.0.0
|
||||
LastEditors : chenht2022
|
||||
LastEditTime : 2024-07-25 10:32:53
|
||||
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
'''
|
||||
import os, sys
|
||||
import time
|
||||
import torch
|
||||
import torch.nn.quantized as nnq
|
||||
|
||||
scale, zero_point = 0.1, 0 # Adjust scale and zero_point based on your dataset
|
||||
|
||||
hidden_size = 5120
|
||||
intermediate_size = 3072
|
||||
layer_num = 10
|
||||
qlen = 1
|
||||
warm_up_iter = 1000
|
||||
test_iter = 10000
|
||||
|
||||
def act_fn(x):
|
||||
return x / (1.0 + torch.exp(-x))
|
||||
|
||||
def mlp_torch(input, gate_proj, up_proj, down_proj):
|
||||
if isinstance(gate_proj, nnq.Linear):
|
||||
input_q = torch.quantize_per_tensor(input.to(torch.float32), scale, zero_point, torch.quint8)
|
||||
gate_buf = gate_proj(input_q)
|
||||
up_buf = up_proj(input_q)
|
||||
gate_buf = gate_buf.dequantize()
|
||||
up_buf = up_buf.dequantize()
|
||||
intermediate = act_fn(gate_buf) * up_buf
|
||||
intermediate_q = torch.quantize_per_tensor(intermediate, scale, zero_point, torch.quint8)
|
||||
expert_output = down_proj(intermediate_q)
|
||||
ret = expert_output.dequantize()
|
||||
else:
|
||||
gate_buf = torch.mm(input.to(gate_proj.dtype), gate_proj.t())
|
||||
up_buf = torch.mm(input.to(up_proj.dtype), up_proj.t())
|
||||
intermediate = act_fn(gate_buf) * up_buf
|
||||
ret = torch.mm(intermediate.to(down_proj.dtype), down_proj.t())
|
||||
return ret
|
||||
|
||||
def bench_mlp(quant_mode: str):
|
||||
with torch.inference_mode(mode=True):
|
||||
if quant_mode == "fp32":
|
||||
proj_type = torch.float32
|
||||
bytes_per_elem = 4.000000
|
||||
elif quant_mode == "fp16":
|
||||
proj_type = torch.float16
|
||||
bytes_per_elem = 2.000000
|
||||
elif quant_mode == "bf16":
|
||||
proj_type = torch.bfloat16
|
||||
bytes_per_elem = 2.000000
|
||||
elif quant_mode == "qint8":
|
||||
proj_type = torch.qint8
|
||||
bytes_per_elem = 1.000000
|
||||
else:
|
||||
assert(False)
|
||||
|
||||
gate_projs = []
|
||||
up_projs = []
|
||||
down_projs = []
|
||||
for _ in range(layer_num):
|
||||
gate_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
up_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
down_proj = torch.randn((hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
if quant_mode == "qint8":
|
||||
gate_proj_q = torch.quantize_per_tensor(gate_proj, scale, zero_point, torch.qint8)
|
||||
quantized_gate = nnq.Linear(hidden_size, intermediate_size)
|
||||
quantized_gate.set_weight_bias(gate_proj_q, None)
|
||||
up_proj_q = torch.quantize_per_tensor(up_proj, scale, zero_point, torch.qint8)
|
||||
quantized_up = nnq.Linear(hidden_size, intermediate_size)
|
||||
quantized_up.set_weight_bias(up_proj_q, None)
|
||||
down_proj_q = torch.quantize_per_tensor(down_proj, scale, zero_point, torch.qint8)
|
||||
quantized_down = nnq.Linear(intermediate_size, hidden_size)
|
||||
quantized_down.set_weight_bias(down_proj_q, None)
|
||||
gate_projs.append(quantized_gate)
|
||||
up_projs.append(quantized_up)
|
||||
down_projs.append(quantized_down)
|
||||
else:
|
||||
gate_projs.append(gate_proj.to(proj_type))
|
||||
up_projs.append(up_proj.to(proj_type))
|
||||
down_projs.append(down_proj.to(proj_type))
|
||||
input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
|
||||
|
||||
# warm up
|
||||
for i in range(warm_up_iter):
|
||||
mlp_torch(input[i % layer_num], gate_projs[i % layer_num], up_projs[i % layer_num], down_projs[i % layer_num])
|
||||
|
||||
# test
|
||||
start = time.perf_counter()
|
||||
for i in range(test_iter):
|
||||
mlp_torch(input[i % layer_num], gate_projs[i % layer_num], up_projs[i % layer_num], down_projs[i % layer_num])
|
||||
end = time.perf_counter()
|
||||
total_time = end - start
|
||||
print('Quant mode: ', quant_mode)
|
||||
print('Time(s): ', total_time)
|
||||
print('Iteration: ', test_iter)
|
||||
print('Time(us) per iteration: ', total_time / test_iter * 1000000)
|
||||
print('Bandwidth: ', hidden_size * intermediate_size * 3 * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s')
|
||||
print('')
|
||||
|
||||
bench_mlp("fp32")
|
||||
bench_mlp("fp16")
|
||||
bench_mlp("bf16")
|
||||
bench_mlp("qint8")
|
||||
@@ -0,0 +1,160 @@
|
||||
#!/usr/bin/env python
|
||||
# coding=utf-8
|
||||
'''
|
||||
Description :
|
||||
Author : chenht2022
|
||||
Date : 2024-07-25 10:32:05
|
||||
Version : 1.0.0
|
||||
LastEditors : chenht2022
|
||||
LastEditTime : 2024-08-06 10:41:28
|
||||
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
'''
|
||||
import os, sys
|
||||
import time
|
||||
sys.path.append(os.path.dirname(__file__) + '/../build')
|
||||
import cpuinfer_ext
|
||||
import torch
|
||||
|
||||
expert_num = 160
|
||||
hidden_size = 5120
|
||||
intermediate_size = 1536
|
||||
stride = 16
|
||||
group_min_len = 10
|
||||
group_max_len = 1024
|
||||
n_routed_experts = 6
|
||||
layer_num = 10
|
||||
qlen = 1
|
||||
CPUInfer = cpuinfer_ext.CPUInfer(64)
|
||||
warm_up_iter = 1000
|
||||
test_iter = 10000
|
||||
|
||||
def bench_moe(quant_mode: str):
|
||||
with torch.inference_mode(mode=True):
|
||||
hidden_type = 30 # ggml_type::GGML_TYPE_BF16
|
||||
if quant_mode == "fp32":
|
||||
gate_type = 0 # ggml_type::GGML_TYPE_F32
|
||||
up_type = 0 # ggml_type::GGML_TYPE_F32
|
||||
down_type = 0 # ggml_type::GGML_TYPE_F32
|
||||
bytes_per_elem = 4.000000
|
||||
elif quant_mode == "fp16":
|
||||
gate_type = 1 # ggml_type::GGML_TYPE_F16
|
||||
up_type = 1 # ggml_type::GGML_TYPE_F16
|
||||
down_type = 1 # ggml_type::GGML_TYPE_F16
|
||||
bytes_per_elem = 2.000000
|
||||
elif quant_mode == "bf16":
|
||||
gate_type = 30 # ggml_type::GGML_TYPE_BF16
|
||||
up_type = 30 # ggml_type::GGML_TYPE_BF16
|
||||
down_type = 30 # ggml_type::GGML_TYPE_BF16
|
||||
bytes_per_elem = 2.000000
|
||||
elif quant_mode == "q8_0":
|
||||
gate_type = 8 # ggml_type::GGML_TYPE_Q8_0
|
||||
up_type = 8 # ggml_type::GGML_TYPE_Q8_0
|
||||
down_type = 8 # ggml_type::GGML_TYPE_Q8_0
|
||||
bytes_per_elem = 1.062500
|
||||
elif quant_mode == "q6_k":
|
||||
gate_type = 14 # ggml_type::GGML_TYPE_Q6_K
|
||||
up_type = 14 # ggml_type::GGML_TYPE_Q6_K
|
||||
down_type = 14 # ggml_type::GGML_TYPE_Q6_K
|
||||
bytes_per_elem = 0.820312
|
||||
elif quant_mode == "q5_k_m":
|
||||
gate_type = 13 # ggml_type::GGML_TYPE_Q5_K
|
||||
up_type = 13 # ggml_type::GGML_TYPE_Q5_K
|
||||
down_type = 14 # ggml_type::GGML_TYPE_Q6_K
|
||||
bytes_per_elem = 0.731771
|
||||
elif quant_mode == "q4_k_m":
|
||||
gate_type = 12 # ggml_type::GGML_TYPE_Q4_K
|
||||
up_type = 12 # ggml_type::GGML_TYPE_Q4_K
|
||||
down_type = 14 # ggml_type::GGML_TYPE_Q6_K
|
||||
bytes_per_elem = 0.648437
|
||||
elif quant_mode == "q3_k_m":
|
||||
gate_type = 11 # ggml_type::GGML_TYPE_Q3_K
|
||||
up_type = 11 # ggml_type::GGML_TYPE_Q3_K
|
||||
down_type = 13 # ggml_type::GGML_TYPE_Q5_K
|
||||
bytes_per_elem = 0.515625
|
||||
elif quant_mode == "q2_k":
|
||||
gate_type = 10 # ggml_type::GGML_TYPE_Q2_K
|
||||
up_type = 10 # ggml_type::GGML_TYPE_Q2_K
|
||||
down_type = 11 # ggml_type::GGML_TYPE_Q3_K
|
||||
bytes_per_elem = 0.328125
|
||||
elif quant_mode == "iq3_xs":
|
||||
gate_type = 21 # ggml_type::GGML_TYPE_IQ3_S
|
||||
up_type = 21 # ggml_type::GGML_TYPE_IQ3_S
|
||||
down_type = 21 # ggml_type::GGML_TYPE_IQ3_S
|
||||
bytes_per_elem = 0.429688
|
||||
elif quant_mode == "iq2_xxs":
|
||||
gate_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS
|
||||
up_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS
|
||||
down_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS
|
||||
bytes_per_elem = 0.257812
|
||||
else:
|
||||
assert(False)
|
||||
|
||||
|
||||
moes = []
|
||||
gate_projs = []
|
||||
up_projs = []
|
||||
down_projs = []
|
||||
for _ in range(layer_num):
|
||||
gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
up_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
down_proj = torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
config = cpuinfer_ext.moe.MOEConfig(expert_num, n_routed_experts, hidden_size, intermediate_size, stride, group_min_len, group_max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), gate_type, up_type, down_type, hidden_type)
|
||||
moe = cpuinfer_ext.moe.MOE(config)
|
||||
gate_projs.append(gate_proj)
|
||||
up_projs.append(up_proj)
|
||||
down_projs.append(down_proj)
|
||||
moes.append(moe)
|
||||
expert_ids = torch.stack([torch.stack([torch.randperm(expert_num, dtype=torch.int64, device = "cuda")[:n_routed_experts] for _ in range(qlen)]) for _ in range(layer_num)]).to("cpu").contiguous()
|
||||
weights = torch.rand((layer_num, qlen, n_routed_experts), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
|
||||
output = torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
|
||||
|
||||
# warm up
|
||||
for i in range(warm_up_iter):
|
||||
CPUInfer.submit(
|
||||
moes[i % layer_num].forward(
|
||||
qlen,
|
||||
n_routed_experts,
|
||||
expert_ids[i % layer_num].data_ptr(),
|
||||
weights[i % layer_num].data_ptr(),
|
||||
input[i % layer_num].data_ptr(),
|
||||
output[i % layer_num].data_ptr()
|
||||
)
|
||||
)
|
||||
CPUInfer.sync()
|
||||
|
||||
# test
|
||||
start = time.perf_counter()
|
||||
for i in range(test_iter):
|
||||
CPUInfer.submit(
|
||||
moes[i % layer_num].forward(
|
||||
qlen,
|
||||
n_routed_experts,
|
||||
expert_ids[i % layer_num].data_ptr(),
|
||||
weights[i % layer_num].data_ptr(),
|
||||
input[i % layer_num].data_ptr(),
|
||||
output[i % layer_num].data_ptr()
|
||||
)
|
||||
)
|
||||
CPUInfer.sync()
|
||||
end = time.perf_counter()
|
||||
total_time = end - start
|
||||
print('Quant mode: ', quant_mode)
|
||||
print('Time(s): ', total_time)
|
||||
print('Iteration: ', test_iter)
|
||||
print('Time(us) per iteration: ', total_time / test_iter * 1000000)
|
||||
print('Bandwidth: ', hidden_size * intermediate_size * 3 * n_routed_experts * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s')
|
||||
print('')
|
||||
|
||||
bench_moe("fp32")
|
||||
bench_moe("fp16")
|
||||
bench_moe("bf16")
|
||||
bench_moe("q8_0")
|
||||
bench_moe("q6_k")
|
||||
bench_moe("q5_k_m")
|
||||
bench_moe("q4_k_m")
|
||||
bench_moe("q3_k_m")
|
||||
bench_moe("q2_k")
|
||||
# Not supported on __x86_64__
|
||||
# bench_linear("iq3_xs")
|
||||
# bench_linear("iq2_xxs")
|
||||
@@ -0,0 +1,107 @@
|
||||
#!/usr/bin/env python
|
||||
# coding=utf-8
|
||||
'''
|
||||
Description :
|
||||
Author : chenht2022
|
||||
Date : 2025-04-25 18:28:12
|
||||
Version : 1.0.0
|
||||
LastEditors : chenht2022
|
||||
LastEditTime : 2025-04-25 18:28:12
|
||||
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
'''
|
||||
import os, sys
|
||||
import time
|
||||
sys.path.append(os.path.dirname(__file__) + '/../build')
|
||||
import cpuinfer_ext
|
||||
import torch
|
||||
|
||||
expert_num = 8
|
||||
hidden_size = 7168
|
||||
intermediate_size = 2048
|
||||
max_len = 25600
|
||||
n_routed_experts = 8
|
||||
layer_num = 10
|
||||
qlen = 1024
|
||||
CPUInfer = cpuinfer_ext.CPUInfer(65)
|
||||
warm_up_iter = 100
|
||||
test_iter = 100
|
||||
|
||||
def bench_moe(quant_mode: str):
|
||||
with torch.inference_mode(mode=True):
|
||||
if quant_mode == "bf16":
|
||||
bytes_per_elem = 2.000000
|
||||
elif quant_mode == "int8":
|
||||
bytes_per_elem = 1.000000
|
||||
else:
|
||||
assert(False)
|
||||
|
||||
|
||||
moes = []
|
||||
gate_projs = []
|
||||
up_projs = []
|
||||
down_projs = []
|
||||
for _ in range(layer_num):
|
||||
gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
up_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
down_proj = torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
config = cpuinfer_ext.moe.AMX_MOEConfig(expert_num, n_routed_experts, hidden_size, intermediate_size, max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr())
|
||||
if quant_mode == "bf16":
|
||||
moe = cpuinfer_ext.moe.AMXBF16_MOE(config)
|
||||
CPUInfer.submit(moe.load_weights())
|
||||
CPUInfer.sync()
|
||||
elif quant_mode == "int8":
|
||||
moe = cpuinfer_ext.moe.AMXInt8_MOE(config)
|
||||
CPUInfer.submit(moe.load_weights())
|
||||
CPUInfer.sync()
|
||||
gate_projs.append(gate_proj)
|
||||
up_projs.append(up_proj)
|
||||
down_projs.append(down_proj)
|
||||
moes.append(moe)
|
||||
expert_ids = torch.stack([torch.stack([torch.randperm(expert_num, dtype=torch.int64, device = "cuda")[:n_routed_experts] for _ in range(qlen)]) for _ in range(layer_num)]).to("cpu").contiguous()
|
||||
weights = torch.rand((layer_num, qlen, n_routed_experts), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
|
||||
output = torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
|
||||
qlen_tensor = torch.tensor([qlen], dtype=torch.int32)
|
||||
|
||||
# warm up
|
||||
for i in range(warm_up_iter):
|
||||
CPUInfer.submit(
|
||||
moes[i % layer_num].forward(
|
||||
qlen,
|
||||
n_routed_experts,
|
||||
expert_ids[i % layer_num].data_ptr(),
|
||||
weights[i % layer_num].data_ptr(),
|
||||
input[i % layer_num].data_ptr(),
|
||||
output[i % layer_num].data_ptr(),
|
||||
qlen_tensor.data_ptr()
|
||||
)
|
||||
)
|
||||
CPUInfer.sync()
|
||||
|
||||
# test
|
||||
start = time.perf_counter()
|
||||
for i in range(test_iter):
|
||||
CPUInfer.submit(
|
||||
moes[i % layer_num].forward(
|
||||
qlen,
|
||||
n_routed_experts,
|
||||
expert_ids[i % layer_num].data_ptr(),
|
||||
weights[i % layer_num].data_ptr(),
|
||||
input[i % layer_num].data_ptr(),
|
||||
output[i % layer_num].data_ptr(),
|
||||
qlen_tensor.data_ptr()
|
||||
)
|
||||
)
|
||||
CPUInfer.sync()
|
||||
end = time.perf_counter()
|
||||
total_time = end - start
|
||||
print('Quant mode: ', quant_mode)
|
||||
print('Time(s): ', total_time)
|
||||
print('Iteration: ', test_iter)
|
||||
print('Time(us) per iteration: ', total_time / test_iter * 1000000)
|
||||
print('Bandwidth: ', hidden_size * intermediate_size * 3 * n_routed_experts * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s')
|
||||
print('Flops: ', hidden_size * intermediate_size * qlen * 3 * n_routed_experts * 2 * test_iter / total_time / 1000 / 1000 / 1000, 'GFLOPS')
|
||||
print('')
|
||||
|
||||
bench_moe("bf16")
|
||||
bench_moe("int8")
|
||||
@@ -0,0 +1,152 @@
|
||||
#!/usr/bin/env python
|
||||
# coding=utf-8
|
||||
'''
|
||||
Description :
|
||||
Author : chenht2022
|
||||
Date : 2024-07-25 10:32:05
|
||||
Version : 1.0.0
|
||||
LastEditors : chenht2022
|
||||
LastEditTime : 2024-07-25 10:32:57
|
||||
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
'''
|
||||
import os, sys
|
||||
import time
|
||||
import torch
|
||||
import torch.nn.quantized as nnq
|
||||
|
||||
scale, zero_point = 0.1, 0 # Adjust scale and zero_point based on your dataset
|
||||
|
||||
expert_num = 160
|
||||
hidden_size = 5120
|
||||
intermediate_size = 1536
|
||||
n_routed_experts = 6
|
||||
layer_num = 10
|
||||
qlen = 1
|
||||
warm_up_iter = 1000
|
||||
test_iter = 10000
|
||||
|
||||
def act_fn(x):
|
||||
return x / (1.0 + torch.exp(-x))
|
||||
|
||||
def mlp_torch(input, gate_proj, up_proj, down_proj):
|
||||
if isinstance(gate_proj, nnq.Linear):
|
||||
input_q = torch.quantize_per_tensor(input.to(torch.float32), scale, zero_point, torch.quint8)
|
||||
gate_buf = gate_proj(input_q)
|
||||
up_buf = up_proj(input_q)
|
||||
gate_buf = gate_buf.dequantize()
|
||||
up_buf = up_buf.dequantize()
|
||||
intermediate = act_fn(gate_buf) * up_buf
|
||||
intermediate_q = torch.quantize_per_tensor(intermediate, scale, zero_point, torch.quint8)
|
||||
expert_output = down_proj(intermediate_q)
|
||||
ret = expert_output.dequantize()
|
||||
else:
|
||||
gate_buf = torch.mm(input.to(gate_proj.dtype), gate_proj.t())
|
||||
up_buf = torch.mm(input.to(up_proj.dtype), up_proj.t())
|
||||
intermediate = act_fn(gate_buf) * up_buf
|
||||
ret = torch.mm(intermediate.to(down_proj.dtype), down_proj.t())
|
||||
return ret
|
||||
|
||||
def moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj):
|
||||
cnts = expert_ids.new_zeros((expert_ids.shape[0], expert_num))
|
||||
cnts.scatter_(1, expert_ids, 1)
|
||||
tokens_per_expert = cnts.sum(dim=0)
|
||||
idxs = expert_ids.view(-1).argsort()
|
||||
sorted_tokens = input[idxs // expert_ids.shape[1]]
|
||||
|
||||
outputs = []
|
||||
start_idx = 0
|
||||
for i, num_tokens in enumerate(tokens_per_expert):
|
||||
end_idx = start_idx + num_tokens
|
||||
if num_tokens == 0:
|
||||
continue
|
||||
tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
|
||||
expert_out = mlp_torch(tokens_for_this_expert, gate_proj[i], up_proj[i], down_proj[i])
|
||||
outputs.append(expert_out)
|
||||
start_idx = end_idx
|
||||
|
||||
outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
|
||||
|
||||
new_x = torch.empty_like(outs)
|
||||
new_x[idxs] = outs
|
||||
t_output = (
|
||||
new_x.view(*expert_ids.shape, -1)
|
||||
.type(weights.dtype)
|
||||
.mul_(weights.unsqueeze(dim=-1))
|
||||
.sum(dim=1)
|
||||
.type(new_x.dtype)
|
||||
)
|
||||
return t_output
|
||||
|
||||
def bench_moe(quant_mode: str):
|
||||
with torch.inference_mode(mode=True):
|
||||
if quant_mode == "fp32":
|
||||
proj_type = torch.float32
|
||||
bytes_per_elem = 4.000000
|
||||
elif quant_mode == "fp16":
|
||||
proj_type = torch.float16
|
||||
bytes_per_elem = 2.000000
|
||||
elif quant_mode == "bf16":
|
||||
proj_type = torch.bfloat16
|
||||
bytes_per_elem = 2.000000
|
||||
elif quant_mode == "qint8":
|
||||
proj_type = torch.qint8
|
||||
bytes_per_elem = 1.000000
|
||||
else:
|
||||
assert(False)
|
||||
|
||||
gate_projs = []
|
||||
up_projs = []
|
||||
down_projs = []
|
||||
for _ in range(layer_num):
|
||||
gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
up_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
down_proj = torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
if quant_mode == "qint8":
|
||||
quantized_gate_proj = []
|
||||
quantized_up_proj = []
|
||||
quantized_down_proj = []
|
||||
for i in range(expert_num):
|
||||
gate_proj_q = torch.quantize_per_tensor(gate_proj[i], scale, zero_point, torch.qint8)
|
||||
quantized_gate = nnq.Linear(hidden_size, intermediate_size)
|
||||
quantized_gate.set_weight_bias(gate_proj_q, None)
|
||||
quantized_gate_proj.append(quantized_gate)
|
||||
up_proj_q = torch.quantize_per_tensor(up_proj[i], scale, zero_point, torch.qint8)
|
||||
quantized_up = nnq.Linear(hidden_size, intermediate_size)
|
||||
quantized_up.set_weight_bias(up_proj_q, None)
|
||||
quantized_up_proj.append(quantized_up)
|
||||
down_proj_q = torch.quantize_per_tensor(down_proj[i], scale, zero_point, torch.qint8)
|
||||
quantized_down = nnq.Linear(intermediate_size, hidden_size)
|
||||
quantized_down.set_weight_bias(down_proj_q, None)
|
||||
quantized_down_proj.append(quantized_down)
|
||||
gate_projs.append(quantized_gate_proj)
|
||||
up_projs.append(quantized_up_proj)
|
||||
down_projs.append(quantized_down_proj)
|
||||
else:
|
||||
gate_projs.append(gate_proj.to(proj_type))
|
||||
up_projs.append(up_proj.to(proj_type))
|
||||
down_projs.append(down_proj.to(proj_type))
|
||||
expert_ids = torch.stack([torch.stack([torch.randperm(expert_num, dtype=torch.int64, device = "cuda")[:n_routed_experts] for _ in range(qlen)]) for _ in range(layer_num)]).to("cpu").contiguous()
|
||||
weights = torch.rand((layer_num, qlen, n_routed_experts), dtype=torch.float32, device = "cuda").to("cpu").contiguous()
|
||||
input = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device = "cuda").to("cpu").contiguous()
|
||||
|
||||
# warm up
|
||||
for i in range(warm_up_iter):
|
||||
moe_torch(input[i % layer_num], expert_ids[i % layer_num], weights[i % layer_num], gate_projs[i % layer_num], up_projs[i % layer_num], down_projs[i % layer_num])
|
||||
|
||||
# test
|
||||
start = time.perf_counter()
|
||||
for i in range(test_iter):
|
||||
moe_torch(input[i % layer_num], expert_ids[i % layer_num], weights[i % layer_num], gate_projs[i % layer_num], up_projs[i % layer_num], down_projs[i % layer_num])
|
||||
end = time.perf_counter()
|
||||
total_time = end - start
|
||||
print('Quant mode: ', quant_mode)
|
||||
print('Time(s): ', total_time)
|
||||
print('Iteration: ', test_iter)
|
||||
print('Time(us) per iteration: ', total_time / test_iter * 1000000)
|
||||
print('Bandwidth: ', hidden_size * intermediate_size * 3 * n_routed_experts * bytes_per_elem * test_iter / total_time / 1000 / 1000 / 1000, 'GB/s')
|
||||
print('')
|
||||
|
||||
bench_moe("fp32")
|
||||
bench_moe("fp16")
|
||||
bench_moe("bf16")
|
||||
bench_moe("qint8")
|
||||
@@ -0,0 +1,100 @@
|
||||
include(CheckCSourceRuns)
|
||||
|
||||
set(AVX_CODE "
|
||||
#include <immintrin.h>
|
||||
int main()
|
||||
{
|
||||
__m256 a;
|
||||
a = _mm256_set1_ps(0);
|
||||
return 0;
|
||||
}
|
||||
")
|
||||
|
||||
set(AVX512_CODE "
|
||||
#include <immintrin.h>
|
||||
int main()
|
||||
{
|
||||
__m512i a = _mm512_set_epi8(0, 0, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0, 0);
|
||||
__m512i b = a;
|
||||
__mmask64 equality_mask = _mm512_cmp_epi8_mask(a, b, _MM_CMPINT_EQ);
|
||||
return 0;
|
||||
}
|
||||
")
|
||||
|
||||
set(AVX2_CODE "
|
||||
#include <immintrin.h>
|
||||
int main()
|
||||
{
|
||||
__m256i a = {0};
|
||||
a = _mm256_abs_epi16(a);
|
||||
__m256i x;
|
||||
_mm256_extract_epi64(x, 0); // we rely on this in our AVX2 code
|
||||
return 0;
|
||||
}
|
||||
")
|
||||
|
||||
set(FMA_CODE "
|
||||
#include <immintrin.h>
|
||||
int main()
|
||||
{
|
||||
__m256 acc = _mm256_setzero_ps();
|
||||
const __m256 d = _mm256_setzero_ps();
|
||||
const __m256 p = _mm256_setzero_ps();
|
||||
acc = _mm256_fmadd_ps( d, p, acc );
|
||||
return 0;
|
||||
}
|
||||
")
|
||||
|
||||
macro(check_sse type flags)
|
||||
set(__FLAG_I 1)
|
||||
set(CMAKE_REQUIRED_FLAGS_SAVE ${CMAKE_REQUIRED_FLAGS})
|
||||
foreach (__FLAG ${flags})
|
||||
if (NOT ${type}_FOUND)
|
||||
set(CMAKE_REQUIRED_FLAGS ${__FLAG})
|
||||
check_c_source_runs("${${type}_CODE}" HAS_${type}_${__FLAG_I})
|
||||
if (HAS_${type}_${__FLAG_I})
|
||||
set(${type}_FOUND TRUE CACHE BOOL "${type} support")
|
||||
set(${type}_FLAGS "${__FLAG}" CACHE STRING "${type} flags")
|
||||
endif()
|
||||
math(EXPR __FLAG_I "${__FLAG_I}+1")
|
||||
endif()
|
||||
endforeach()
|
||||
set(CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS_SAVE})
|
||||
|
||||
if (NOT ${type}_FOUND)
|
||||
set(${type}_FOUND FALSE CACHE BOOL "${type} support")
|
||||
set(${type}_FLAGS "" CACHE STRING "${type} flags")
|
||||
endif()
|
||||
|
||||
mark_as_advanced(${type}_FOUND ${type}_FLAGS)
|
||||
endmacro()
|
||||
|
||||
# flags are for MSVC only!
|
||||
check_sse("AVX" " ;/arch:AVX")
|
||||
if (NOT ${AVX_FOUND})
|
||||
set(LLAMA_AVX OFF)
|
||||
else()
|
||||
set(LLAMA_AVX ON)
|
||||
endif()
|
||||
|
||||
check_sse("AVX2" " ;/arch:AVX2")
|
||||
check_sse("FMA" " ;/arch:AVX2")
|
||||
if ((NOT ${AVX2_FOUND}) OR (NOT ${FMA_FOUND}))
|
||||
set(LLAMA_AVX2 OFF)
|
||||
else()
|
||||
set(LLAMA_AVX2 ON)
|
||||
endif()
|
||||
|
||||
check_sse("AVX512" " ;/arch:AVX512")
|
||||
if (NOT ${AVX512_FOUND})
|
||||
set(LLAMA_AVX512 OFF)
|
||||
else()
|
||||
set(LLAMA_AVX512 ON)
|
||||
endif()
|
||||
@@ -0,0 +1,154 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : chenht2022
|
||||
* @Date : 2024-07-22 02:03:05
|
||||
* @Version : 1.0.0
|
||||
* @LastEditors : chenht2022
|
||||
* @LastEditTime : 2024-07-25 10:33:34
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
|
||||
#include "backend.h"
|
||||
|
||||
#ifdef USE_NUMA
|
||||
#include <numa.h>
|
||||
#include <numaif.h>
|
||||
|
||||
thread_local int Backend::numa_node = -1;
|
||||
#endif
|
||||
|
||||
thread_local int Backend::thread_local_id = -1;
|
||||
|
||||
Backend::Backend(int max_thread_num) {
|
||||
max_thread_num_ = max_thread_num;
|
||||
thread_state_.resize(max_thread_num_);
|
||||
for (int i = 0; i < max_thread_num_; i++) {
|
||||
thread_state_[i].curr = std::make_unique<std::atomic<int>>();
|
||||
thread_state_[i].status =
|
||||
std::make_unique<std::atomic<ThreadStatus>>(ThreadStatus::WAITING);
|
||||
}
|
||||
workers_.resize(max_thread_num_);
|
||||
for (int i = 1; i < max_thread_num_; i++) {
|
||||
workers_[i] = std::thread(&Backend::worker_thread, this, i);
|
||||
}
|
||||
}
|
||||
|
||||
Backend::~Backend() {
|
||||
for (int i = 0; i < max_thread_num_; i++) {
|
||||
thread_state_[i].status->store(ThreadStatus::EXIT,
|
||||
std::memory_order_release);
|
||||
}
|
||||
for (int i = 1; i < max_thread_num_; i++) {
|
||||
if (workers_[i].joinable()) {
|
||||
workers_[i].join();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int Backend::get_thread_num() { return max_thread_num_; }
|
||||
|
||||
void Backend::do_work_stealing_job(int task_num,
|
||||
std::function<void(int)> init_func,
|
||||
std::function<void(int)> compute_func,
|
||||
std::function<void(int)> finalize_func) {
|
||||
init_func_ = init_func;
|
||||
compute_func_ = compute_func;
|
||||
finalize_func_ = finalize_func;
|
||||
#ifdef USE_NUMA
|
||||
// numa node location will be calculated based on the number of threads
|
||||
thread_num_ = max_thread_num_;
|
||||
#else
|
||||
thread_num_ = std::min(max_thread_num_, task_num);
|
||||
#endif
|
||||
int base = task_num / thread_num_;
|
||||
int remain = task_num % thread_num_;
|
||||
thread_state_[0].end = base + (0 < remain);
|
||||
|
||||
// 为主线程设置 thread_local_id
|
||||
thread_local_id = 0;
|
||||
|
||||
for (int i = 1; i < thread_num_; i++) {
|
||||
thread_state_[i].curr->store(thread_state_[i - 1].end,
|
||||
std::memory_order_relaxed);
|
||||
thread_state_[i].end = thread_state_[i - 1].end + base + (i < remain);
|
||||
thread_state_[i].status->store(ThreadStatus::WORKING,
|
||||
std::memory_order_release);
|
||||
}
|
||||
thread_state_[0].curr->store(0, std::memory_order_relaxed);
|
||||
thread_state_[0].status->store(ThreadStatus::WORKING,
|
||||
std::memory_order_release);
|
||||
process_tasks(0);
|
||||
for (int i = 1; i < thread_num_; i++) {
|
||||
while (thread_state_[i].status->load(std::memory_order_acquire) ==
|
||||
ThreadStatus::WORKING) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void Backend::process_tasks(int thread_id) {
|
||||
|
||||
#ifdef USE_NUMA
|
||||
if(numa_node == -1){
|
||||
numa_node = thread_id * numa_num_configured_nodes() / thread_num_;
|
||||
struct bitmask* mask = numa_bitmask_alloc(numa_num_configured_nodes());
|
||||
numa_bitmask_setbit(mask, numa_node);
|
||||
numa_bind(mask);
|
||||
}
|
||||
#endif
|
||||
|
||||
if (init_func_ != nullptr) {
|
||||
init_func_(thread_id);
|
||||
}
|
||||
while (true) {
|
||||
int task_id = thread_state_[thread_id].curr->fetch_add(
|
||||
1, std::memory_order_acq_rel);
|
||||
if (task_id >= thread_state_[thread_id].end) {
|
||||
break;
|
||||
}
|
||||
compute_func_(task_id);
|
||||
}
|
||||
for (int t_offset = 1; t_offset < thread_num_; t_offset++) {
|
||||
int t_i = (thread_id + t_offset) % thread_num_;
|
||||
if (thread_state_[t_i].status->load(std::memory_order_acquire) !=
|
||||
ThreadStatus::WORKING) {
|
||||
continue;
|
||||
}
|
||||
while (true) {
|
||||
int task_id = thread_state_[t_i].curr->fetch_add(
|
||||
1, std::memory_order_acq_rel);
|
||||
if (task_id >= thread_state_[t_i].end) {
|
||||
break;
|
||||
}
|
||||
compute_func_(task_id);
|
||||
}
|
||||
}
|
||||
if (finalize_func_ != nullptr) {
|
||||
finalize_func_(thread_id);
|
||||
}
|
||||
thread_state_[thread_id].status->store(ThreadStatus::WAITING,
|
||||
std::memory_order_release);
|
||||
}
|
||||
|
||||
void Backend::worker_thread(int thread_id) {
|
||||
auto start = std::chrono::steady_clock::now();
|
||||
thread_local_id = thread_id; // 设置线程本地变量
|
||||
while (true) {
|
||||
ThreadStatus status =
|
||||
thread_state_[thread_id].status->load(std::memory_order_acquire);
|
||||
if (status == ThreadStatus::WORKING) {
|
||||
process_tasks(thread_id);
|
||||
start = std::chrono::steady_clock::now();
|
||||
} else if (status == ThreadStatus::WAITING) {
|
||||
auto now = std::chrono::steady_clock::now();
|
||||
auto duration =
|
||||
std::chrono::duration_cast<std::chrono::milliseconds>(now -
|
||||
start)
|
||||
.count();
|
||||
if (duration > 50) {
|
||||
std::this_thread::sleep_for(std::chrono::milliseconds(1));
|
||||
}
|
||||
} else if (status == ThreadStatus::EXIT) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : chenht2022
|
||||
* @Date : 2024-07-22 02:03:05
|
||||
* @Version : 1.0.0
|
||||
* @LastEditors : chenht2022
|
||||
* @LastEditTime : 2024-07-25 10:33:38
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
#ifndef CPUINFER_BACKEND_H
|
||||
#define CPUINFER_BACKEND_H
|
||||
|
||||
#include <atomic>
|
||||
#include <condition_variable>
|
||||
#include <cstdio>
|
||||
#include <functional>
|
||||
#include <mutex>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
|
||||
enum ThreadStatus {
|
||||
WORKING,
|
||||
WAITING,
|
||||
EXIT,
|
||||
};
|
||||
|
||||
struct ThreadState {
|
||||
std::unique_ptr<std::atomic<ThreadStatus>> status;
|
||||
std::unique_ptr<std::atomic<int>> curr;
|
||||
int end;
|
||||
};
|
||||
|
||||
class Backend {
|
||||
public:
|
||||
Backend(int);
|
||||
~Backend();
|
||||
int get_thread_num();
|
||||
void do_work_stealing_job(int, std::function<void(int)>,
|
||||
std::function<void(int)>,
|
||||
std::function<void(int)>);
|
||||
#ifdef USE_NUMA
|
||||
static thread_local int numa_node;
|
||||
#endif
|
||||
static thread_local int thread_local_id;
|
||||
|
||||
private:
|
||||
int thread_num_;
|
||||
int max_thread_num_;
|
||||
std::vector<ThreadState> thread_state_; // [thread_num]
|
||||
std::function<void(int)> init_func_;
|
||||
std::function<void(int)> compute_func_;
|
||||
std::function<void(int)> finalize_func_;
|
||||
std::vector<std::thread> workers_;
|
||||
|
||||
void process_tasks(int);
|
||||
void worker_thread(int);
|
||||
};
|
||||
#endif
|
||||
@@ -0,0 +1,98 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : chenht2022
|
||||
* @Date : 2024-07-16 10:43:18
|
||||
* @Version : 1.0.0
|
||||
* @LastEditors : chenht2022
|
||||
* @LastEditTime : 2024-08-07 09:47:43
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
#ifndef CPUINFER_CPUINFER_H
|
||||
#define CPUINFER_CPUINFER_H
|
||||
|
||||
#include <atomic>
|
||||
#include <condition_variable>
|
||||
#include <functional>
|
||||
#include <mutex>
|
||||
#include <queue>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
#include <stdexcept>
|
||||
#ifdef KTRANSFORMERS_USE_CUDA
|
||||
#include "vendors/cuda.h"
|
||||
#elif KTRANSFORMERS_USE_MUSA
|
||||
#include "vendors/musa.h"
|
||||
#elif KTRANSFORMERS_USE_ROCM
|
||||
#define __HIP_PLATFORM_AMD__
|
||||
#include "vendors/hip.h"
|
||||
#endif
|
||||
|
||||
#include "backend.h"
|
||||
#include "task_queue.h"
|
||||
#include "./vendors/vendor.h"
|
||||
|
||||
#include "llama.cpp/ggml-impl.h"
|
||||
|
||||
class CPUInfer {
|
||||
public:
|
||||
CPUInfer(int thread_num) {
|
||||
backend_ = new Backend(thread_num - 1);
|
||||
task_queue_ = new TaskQueue();
|
||||
for (int i = 0; i < (1 << 16); ++i) {
|
||||
ggml_table_f32_f16[i] = GGML_COMPUTE_FP16_TO_FP32(i);
|
||||
}
|
||||
}
|
||||
|
||||
~CPUInfer() {
|
||||
delete backend_;
|
||||
delete task_queue_;
|
||||
}
|
||||
|
||||
template <typename Func, typename Obj, typename... Args>
|
||||
void enqueue(Func f, Obj* obj, Args... args) {
|
||||
task_queue_->enqueue([=]() {
|
||||
std::invoke(f, *obj, args..., backend_);
|
||||
});
|
||||
}
|
||||
|
||||
void submit(std::pair<intptr_t, intptr_t> params) {
|
||||
void (*func)(void*) = (void (*)(void*))params.first;
|
||||
void* args = (void*)params.second;
|
||||
*((CPUInfer**)args) = this;
|
||||
func(args);
|
||||
}
|
||||
|
||||
void sync() {
|
||||
task_queue_->sync();
|
||||
}
|
||||
|
||||
void submit_with_cuda_stream(intptr_t user_cuda_stream, std::pair<intptr_t, intptr_t> params) {
|
||||
#if defined(KTRANSFORMERS_USE_CUDA) || defined(KTRANSFORMERS_USE_MUSA) || defined(KTRANSFORMERS_USE_ROCM)
|
||||
void (*func)(void*) = (void (*)(void*))params.first;
|
||||
void* args = (void*)params.second;
|
||||
*((CPUInfer**)args) = this;
|
||||
cudaLaunchHostFunc((cudaStream_t)user_cuda_stream, (cudaHostFn_t)func, args);
|
||||
#else
|
||||
throw std::runtime_error("submit_with_cuda_stream is not supported on this platforma");
|
||||
#endif
|
||||
}
|
||||
|
||||
static void sync_(void* cpu_infer_ptr) {
|
||||
CPUInfer* cpuinfer = (CPUInfer*)cpu_infer_ptr;
|
||||
cpuinfer->sync();
|
||||
}
|
||||
|
||||
void sync_with_cuda_stream(intptr_t user_cuda_stream) {
|
||||
#if defined(KTRANSFORMERS_USE_CUDA) || defined(KTRANSFORMERS_USE_MUSA) || defined(KTRANSFORMERS_USE_ROCM)
|
||||
cudaLaunchHostFunc((cudaStream_t)user_cuda_stream, (cudaHostFn_t)&sync_, (void*)this);
|
||||
#else
|
||||
throw std::runtime_error("sync_with_cuda_stream is not supported on this platforma");
|
||||
#endif
|
||||
}
|
||||
|
||||
public:
|
||||
Backend* backend_;
|
||||
TaskQueue* task_queue_;
|
||||
};
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,56 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : chenht2022
|
||||
* @Date : 2024-08-05 04:49:08
|
||||
* @Version : 1.0.0
|
||||
* @LastEditors : chenht2022
|
||||
* @LastEditTime : 2024-08-05 09:21:29
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
#include "shared_mem_buffer.h"
|
||||
#include <cstdio>
|
||||
|
||||
SharedMemBuffer::SharedMemBuffer() {
|
||||
buffer_ = nullptr;
|
||||
size_ = 0;
|
||||
}
|
||||
|
||||
SharedMemBuffer::~SharedMemBuffer() {
|
||||
if (buffer_) {
|
||||
free(buffer_);
|
||||
}
|
||||
}
|
||||
|
||||
void SharedMemBuffer::alloc(void* object, std::vector<std::pair<void**, uint64_t>> requests) {
|
||||
uint64_t size = 0;
|
||||
for (auto& request : requests) {
|
||||
size += request.second;
|
||||
}
|
||||
if (size > size_) {
|
||||
if (buffer_) {
|
||||
free(buffer_);
|
||||
}
|
||||
buffer_ = std::aligned_alloc(64, size);
|
||||
|
||||
size_ = size;
|
||||
for (auto& obj_requests : hist_requests_) {
|
||||
for (auto& requests : obj_requests.second) {
|
||||
arrange(requests);
|
||||
}
|
||||
}
|
||||
}
|
||||
arrange(requests);
|
||||
hist_requests_[object].push_back(requests);
|
||||
}
|
||||
|
||||
void SharedMemBuffer::dealloc(void* object) {
|
||||
hist_requests_.erase(object);
|
||||
}
|
||||
|
||||
void SharedMemBuffer::arrange(std::vector<std::pair<void**, uint64_t>> requests) {
|
||||
uint64_t offset = 0;
|
||||
for (auto& request : requests) {
|
||||
*(request.first) = (uint8_t*)buffer_ + offset;
|
||||
offset += request.second;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : chenht2022
|
||||
* @Date : 2024-08-05 04:49:08
|
||||
* @Version : 1.0.0
|
||||
* @LastEditors : chenht2022
|
||||
* @LastEditTime : 2024-08-05 06:36:41
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
|
||||
#ifndef CPUINFER_SHAREDMEMBUFFER_H
|
||||
#define CPUINFER_SHAREDMEMBUFFER_H
|
||||
|
||||
#include <cstdint>
|
||||
#include <cstdlib>
|
||||
#include <map>
|
||||
#include <vector>
|
||||
|
||||
class SharedMemBuffer {
|
||||
public:
|
||||
SharedMemBuffer();
|
||||
~SharedMemBuffer();
|
||||
|
||||
void alloc(void* object, std::vector<std::pair<void**, uint64_t>> requests);
|
||||
void dealloc(void* object);
|
||||
|
||||
private:
|
||||
void* buffer_;
|
||||
uint64_t size_;
|
||||
std::map<void*, std::vector<std::vector<std::pair<void**, uint64_t>>>> hist_requests_;
|
||||
|
||||
void arrange(std::vector<std::pair<void**, uint64_t>> requests);
|
||||
};
|
||||
|
||||
static SharedMemBuffer shared_mem_buffer;
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,67 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : chenht2022
|
||||
* @Date : 2024-07-17 12:25:51
|
||||
* @Version : 1.0.0
|
||||
* @LastEditors : chenht2022
|
||||
* @LastEditTime : 2024-10-09 11:08:10
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
#include "task_queue.h"
|
||||
|
||||
TaskQueue::TaskQueue() {
|
||||
worker = std::thread(&TaskQueue::processTasks, this);
|
||||
sync_flag.store(true, std::memory_order_seq_cst);
|
||||
exit_flag.store(false, std::memory_order_seq_cst);
|
||||
}
|
||||
|
||||
TaskQueue::~TaskQueue() {
|
||||
{
|
||||
mutex.lock();
|
||||
exit_flag.store(true, std::memory_order_seq_cst);
|
||||
mutex.unlock();
|
||||
}
|
||||
cv.notify_all();
|
||||
if (worker.joinable()) {
|
||||
worker.join();
|
||||
}
|
||||
}
|
||||
|
||||
void TaskQueue::enqueue(std::function<void()> task) {
|
||||
{
|
||||
mutex.lock();
|
||||
tasks.push(task);
|
||||
sync_flag.store(false, std::memory_order_seq_cst);
|
||||
mutex.unlock();
|
||||
}
|
||||
cv.notify_one();
|
||||
}
|
||||
|
||||
void TaskQueue::sync() {
|
||||
while (!sync_flag.load(std::memory_order_seq_cst))
|
||||
;
|
||||
}
|
||||
|
||||
void TaskQueue::processTasks() {
|
||||
while (true) {
|
||||
std::function<void()> task;
|
||||
{
|
||||
mutex.lock();
|
||||
cv.wait(mutex, [this]() { return !tasks.empty() || exit_flag.load(std::memory_order_seq_cst); });
|
||||
if (exit_flag.load(std::memory_order_seq_cst) && tasks.empty()) {
|
||||
return;
|
||||
}
|
||||
task = tasks.front();
|
||||
tasks.pop();
|
||||
mutex.unlock();
|
||||
}
|
||||
task();
|
||||
{
|
||||
mutex.lock();
|
||||
if (tasks.empty()) {
|
||||
sync_flag.store(true, std::memory_order_seq_cst);
|
||||
}
|
||||
mutex.unlock();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,138 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : chenht2022
|
||||
* @Date : 2024-07-16 10:43:18
|
||||
* @Version : 1.0.0
|
||||
* @LastEditors : chenht
|
||||
* @LastEditTime : 2024-10-09 11:08:07
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
#ifndef CPUINFER_TASKQUEUE_H
|
||||
#define CPUINFER_TASKQUEUE_H
|
||||
|
||||
#include <atomic>
|
||||
#include <condition_variable>
|
||||
#include <functional>
|
||||
#include <mutex>
|
||||
#include <queue>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
#ifdef _WIN32
|
||||
#include <windows.h>
|
||||
#endif
|
||||
|
||||
class custom_mutex {
|
||||
private:
|
||||
#ifdef _WIN32
|
||||
CRITICAL_SECTION cs;
|
||||
#else
|
||||
std::mutex mtx;
|
||||
#endif
|
||||
|
||||
public:
|
||||
custom_mutex() {
|
||||
#ifdef _WIN32
|
||||
InitializeCriticalSection(&cs);
|
||||
#else
|
||||
// No initialization required for std::mutex
|
||||
#endif
|
||||
}
|
||||
|
||||
~custom_mutex() {
|
||||
#ifdef _WIN32
|
||||
DeleteCriticalSection(&cs);
|
||||
#endif
|
||||
}
|
||||
|
||||
void lock() {
|
||||
#ifdef _WIN32
|
||||
EnterCriticalSection(&cs);
|
||||
#else
|
||||
mtx.lock();
|
||||
#endif
|
||||
}
|
||||
|
||||
void unlock() {
|
||||
#ifdef _WIN32
|
||||
LeaveCriticalSection(&cs);
|
||||
#else
|
||||
mtx.unlock();
|
||||
#endif
|
||||
}
|
||||
|
||||
#ifdef _WIN32
|
||||
CRITICAL_SECTION* get_handle() {
|
||||
return &cs;
|
||||
}
|
||||
#else
|
||||
std::mutex* get_handle() {
|
||||
return &mtx;
|
||||
}
|
||||
#endif
|
||||
};
|
||||
|
||||
class custom_condition_variable {
|
||||
private:
|
||||
#ifdef _WIN32
|
||||
CONDITION_VARIABLE cond_var;
|
||||
#else
|
||||
std::condition_variable cond_var;
|
||||
#endif
|
||||
|
||||
public:
|
||||
custom_condition_variable() {
|
||||
#ifdef _WIN32
|
||||
InitializeConditionVariable(&cond_var);
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename Predicate>
|
||||
void wait(custom_mutex& mutex, Predicate pred) {
|
||||
#ifdef _WIN32
|
||||
while (!pred()) {
|
||||
SleepConditionVariableCS(&cond_var, mutex.get_handle(), INFINITE);
|
||||
}
|
||||
#else
|
||||
std::unique_lock<std::mutex> lock(*mutex.get_handle(), std::adopt_lock);
|
||||
cond_var.wait(lock, pred);
|
||||
lock.release();
|
||||
#endif
|
||||
}
|
||||
|
||||
void notify_one() {
|
||||
#ifdef _WIN32
|
||||
WakeConditionVariable(&cond_var);
|
||||
#else
|
||||
cond_var.notify_one();
|
||||
#endif
|
||||
}
|
||||
|
||||
void notify_all() {
|
||||
#ifdef _WIN32
|
||||
WakeAllConditionVariable(&cond_var);
|
||||
#else
|
||||
cond_var.notify_all();
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
class TaskQueue {
|
||||
public:
|
||||
TaskQueue();
|
||||
~TaskQueue();
|
||||
|
||||
void enqueue(std::function<void()>);
|
||||
|
||||
void sync();
|
||||
|
||||
private:
|
||||
void processTasks();
|
||||
|
||||
std::queue<std::function<void()>> tasks;
|
||||
custom_mutex mutex;
|
||||
custom_condition_variable cv;
|
||||
std::thread worker;
|
||||
std::atomic<bool> sync_flag;
|
||||
std::atomic<bool> exit_flag;
|
||||
};
|
||||
#endif
|
||||
@@ -0,0 +1,3 @@
|
||||
## TODO
|
||||
|
||||
This directory can be removed after updating the version of `llama.cpp`.
|
||||
@@ -0,0 +1,15 @@
|
||||
#pragma once
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda.h>
|
||||
#include <cublas_v2.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
|
||||
#if CUDART_VERSION < 11020
|
||||
#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED CU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED
|
||||
#define CUBLAS_TF32_TENSOR_OP_MATH CUBLAS_TENSOR_OP_MATH
|
||||
#define CUBLAS_COMPUTE_16F CUDA_R_16F
|
||||
#define CUBLAS_COMPUTE_32F CUDA_R_32F
|
||||
#define cublasComputeType_t cudaDataType_t
|
||||
#endif // CUDART_VERSION < 11020
|
||||
@@ -0,0 +1,172 @@
|
||||
#pragma once
|
||||
|
||||
#define HIP_ENABLE_WARP_SYNC_BUILTINS 1
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <hipblas/hipblas.h>
|
||||
#include <hip/hip_fp16.h>
|
||||
#include <hip/hip_bfloat16.h>
|
||||
#ifdef __HIP_PLATFORM_AMD__
|
||||
// for rocblas_initialize()
|
||||
#include "rocblas/rocblas.h"
|
||||
#endif // __HIP_PLATFORM_AMD__
|
||||
|
||||
#define CUBLAS_COMPUTE_16F HIPBLAS_R_16F
|
||||
#define CUBLAS_COMPUTE_32F HIPBLAS_R_32F
|
||||
#define CUBLAS_COMPUTE_32F_FAST_16F HIPBLAS_R_32F
|
||||
#define CUBLAS_GEMM_DEFAULT HIPBLAS_GEMM_DEFAULT
|
||||
#define CUBLAS_GEMM_DEFAULT_TENSOR_OP HIPBLAS_GEMM_DEFAULT
|
||||
#define CUBLAS_OP_N HIPBLAS_OP_N
|
||||
#define CUBLAS_OP_T HIPBLAS_OP_T
|
||||
#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS
|
||||
#define CUBLAS_TF32_TENSOR_OP_MATH 0
|
||||
#define CUDA_R_16F HIPBLAS_R_16F
|
||||
#define CUDA_R_32F HIPBLAS_R_32F
|
||||
#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED hipDeviceAttributeVirtualMemoryManagementSupported
|
||||
#define CU_MEM_ALLOC_GRANULARITY_RECOMMENDED hipMemAllocationGranularityRecommended
|
||||
#define CU_MEM_ALLOCATION_TYPE_PINNED hipMemAllocationTypePinned
|
||||
#define CU_MEM_LOCATION_TYPE_DEVICE hipMemLocationTypeDevice
|
||||
#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE hipMemAccessFlagsProtReadWrite
|
||||
#define CU_CHECK(fn) {hipError_t err = fn; if(err != hipSuccess) { GGML_ABORT("HipVMM Failure: %s\n", hipGetErrorString(err)); }}
|
||||
#define __shfl_sync(mask, var, laneMask, width) __shfl(var, laneMask, width)
|
||||
#define __shfl_xor_sync(mask, var, laneMask, width) __shfl_xor(var, laneMask, width)
|
||||
#define cublasComputeType_t hipblasDatatype_t //deprecated, new hipblasComputeType_t not in 5.6
|
||||
#define cublasCreate hipblasCreate
|
||||
#define cublasDestroy hipblasDestroy
|
||||
#define cublasGemmEx hipblasGemmEx
|
||||
#define cublasGemmBatchedEx hipblasGemmBatchedEx
|
||||
#define cublasGemmStridedBatchedEx hipblasGemmStridedBatchedEx
|
||||
#define cublasHandle_t hipblasHandle_t
|
||||
#define cublasSetMathMode(handle, mode) CUBLAS_STATUS_SUCCESS
|
||||
#define cublasSetStream hipblasSetStream
|
||||
#define cublasSgemm hipblasSgemm
|
||||
#define cublasStatus_t hipblasStatus_t
|
||||
#define cublasOperation_t hipblasOperation_t
|
||||
#define cudaDataType_t hipblasDatatype_t //deprecated, new hipblasDatatype not in 5.6
|
||||
#define cudaDeviceCanAccessPeer hipDeviceCanAccessPeer
|
||||
#define cudaDeviceDisablePeerAccess hipDeviceDisablePeerAccess
|
||||
#define cudaDeviceEnablePeerAccess hipDeviceEnablePeerAccess
|
||||
#define cudaDeviceProp hipDeviceProp_t
|
||||
#define cudaDeviceSynchronize hipDeviceSynchronize
|
||||
#define cudaError_t hipError_t
|
||||
#define cudaErrorPeerAccessAlreadyEnabled hipErrorPeerAccessAlreadyEnabled
|
||||
#define cudaErrorPeerAccessNotEnabled hipErrorPeerAccessNotEnabled
|
||||
#define cudaEventCreateWithFlags hipEventCreateWithFlags
|
||||
#define cudaEventDisableTiming hipEventDisableTiming
|
||||
#define cudaEventRecord hipEventRecord
|
||||
#define cudaEventSynchronize hipEventSynchronize
|
||||
#define cudaEvent_t hipEvent_t
|
||||
#define cudaEventDestroy hipEventDestroy
|
||||
#define cudaFree hipFree
|
||||
#define cudaFreeHost hipHostFree
|
||||
#define cudaGetDevice hipGetDevice
|
||||
#define cudaGetDeviceCount hipGetDeviceCount
|
||||
#define cudaGetDeviceProperties hipGetDeviceProperties
|
||||
#define cudaGetErrorString hipGetErrorString
|
||||
#define cudaGetLastError hipGetLastError
|
||||
#define cudaHostRegister hipHostRegister
|
||||
#define cudaHostRegisterPortable hipHostRegisterPortable
|
||||
#define cudaHostRegisterReadOnly hipHostRegisterReadOnly
|
||||
#define cudaHostUnregister hipHostUnregister
|
||||
#define cudaLaunchHostFunc hipLaunchHostFunc
|
||||
#define cudaMalloc hipMalloc
|
||||
#define cudaMallocHost(ptr, size) hipHostMalloc(ptr, size, hipHostMallocDefault)
|
||||
#define cudaMemcpy hipMemcpy
|
||||
#define cudaMemcpyAsync hipMemcpyAsync
|
||||
#define cudaMemcpyPeerAsync hipMemcpyPeerAsync
|
||||
#define cudaMemcpy2DAsync hipMemcpy2DAsync
|
||||
#define cudaMemcpyDeviceToDevice hipMemcpyDeviceToDevice
|
||||
#define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost
|
||||
#define cudaMemcpyHostToDevice hipMemcpyHostToDevice
|
||||
#define cudaMemcpyKind hipMemcpyKind
|
||||
#define cudaMemset hipMemset
|
||||
#define cudaMemsetAsync hipMemsetAsync
|
||||
#define cudaMemGetInfo hipMemGetInfo
|
||||
#define cudaOccupancyMaxPotentialBlockSize hipOccupancyMaxPotentialBlockSize
|
||||
#define cudaSetDevice hipSetDevice
|
||||
#define cuDeviceGet hipDeviceGet
|
||||
#define CUdevice hipDevice_t
|
||||
#define CUdeviceptr hipDeviceptr_t
|
||||
#define cuMemUnmap hipMemUnmap
|
||||
#define CUmemAccessDesc hipMemAccessDesc
|
||||
#define cuMemAddressFree hipMemAddressFree
|
||||
#define cuMemRelease hipMemRelease
|
||||
#define CUmemGenericAllocationHandle hipMemGenericAllocationHandle_t
|
||||
#define cuMemCreate hipMemCreate
|
||||
#define cuMemAddressReserve hipMemAddressReserve
|
||||
#define cuMemMap hipMemMap
|
||||
#define cuMemSetAccess hipMemSetAccess
|
||||
#define cuMemGetAllocationGranularity hipMemGetAllocationGranularity
|
||||
#define CUmemAllocationProp hipMemAllocationProp
|
||||
#define cuDeviceGetAttribute hipDeviceGetAttribute
|
||||
#define cudaStreamCreateWithFlags hipStreamCreateWithFlags
|
||||
#define cudaStreamDestroy hipStreamDestroy
|
||||
#define cudaStreamFireAndForget hipStreamFireAndForget
|
||||
#define cudaStreamNonBlocking hipStreamNonBlocking
|
||||
#define cudaStreamPerThread hipStreamPerThread
|
||||
#define cudaStreamSynchronize hipStreamSynchronize
|
||||
#define cudaStreamWaitEvent(stream, event, flags) hipStreamWaitEvent(stream, event, flags)
|
||||
#define cudaGraphExec_t hipGraphExec_t
|
||||
#define cudaGraphNode_t hipGraphNode_t
|
||||
#define cudaKernelNodeParams hipKernelNodeParams
|
||||
#define cudaKernelNodeParams hipKernelNodeParams
|
||||
#define cudaGraphExecDestroy hipGraphExecDestroy
|
||||
#define cudaGraphLaunch hipGraphLaunch
|
||||
#define cudaErrorGraphExecUpdateFailure hipErrorGraphExecUpdateFailure
|
||||
#define cudaGraphExecUpdateResultInfo hipGraphExecUpdateResult
|
||||
#define cudaGraphNodeType hipGraphNodeType
|
||||
#define cudaGraphNodeTypeKernel hipGraphNodeTypeKernel
|
||||
#define cudaGraphInstantiate hipGraphInstantiate
|
||||
#define cudaStreamEndCapture hipStreamEndCapture
|
||||
#define cudaGraphDestroy hipGraphDestroy
|
||||
#define cudaGraphKernelNodeSetParams hipGraphKernelNodeSetParams
|
||||
#define cudaErrorInvalidDeviceFunction hipErrorInvalidDeviceFunction
|
||||
#define cudaGraphKernelNodeGetParams hipGraphKernelNodeGetParams
|
||||
#define cudaGraphNodeGetType hipGraphNodeGetType
|
||||
#define cudaGraphGetNodes hipGraphGetNodes
|
||||
#define cudaGraphExecUpdate hipGraphExecUpdate
|
||||
#define cudaStreamCaptureModeRelaxed hipStreamCaptureModeRelaxed
|
||||
#define cudaStreamBeginCapture hipStreamBeginCapture
|
||||
#define cudaGraph_t hipGraph_t
|
||||
#define cudaStream_t hipStream_t
|
||||
#define cudaSuccess hipSuccess
|
||||
#define cudaHostFn_t hipHostFn_t
|
||||
#define __trap() do { abort(); __builtin_unreachable(); } while(0)
|
||||
#define CUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS
|
||||
#define CUBLAS_STATUS_NOT_INITIALIZED HIPBLAS_STATUS_NOT_INITIALIZED
|
||||
#define CUBLAS_STATUS_ALLOC_FAILED HIPBLAS_STATUS_ALLOC_FAILED
|
||||
#define CUBLAS_STATUS_INVALID_VALUE HIPBLAS_STATUS_INVALID_VALUE
|
||||
#define CUBLAS_STATUS_ARCH_MISMATCH HIPBLAS_STATUS_ARCH_MISMATCH
|
||||
#define CUBLAS_STATUS_MAPPING_ERROR HIPBLAS_STATUS_MAPPING_ERROR
|
||||
#define CUBLAS_STATUS_EXECUTION_FAILED HIPBLAS_STATUS_EXECUTION_FAILED
|
||||
#define CUBLAS_STATUS_INTERNAL_ERROR HIPBLAS_STATUS_INTERNAL_ERROR
|
||||
#define CUBLAS_STATUS_NOT_SUPPORTED HIPBLAS_STATUS_NOT_SUPPORTED
|
||||
|
||||
#define __CUDA_ARCH__ 1300
|
||||
|
||||
#if defined(__gfx803__) || defined(__gfx900__) || defined(__gfx906__)
|
||||
#define GCN
|
||||
#endif
|
||||
|
||||
#if defined(__gfx908__) || defined(__gfx90a__) || defined(__gfx942__)
|
||||
#define CDNA
|
||||
#endif
|
||||
|
||||
#if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || \
|
||||
defined(__gfx1150__) || defined(__gfx1151__)
|
||||
#define RDNA3
|
||||
#endif
|
||||
|
||||
#if defined(__gfx1030__) || defined(__gfx1031__) || defined(__gfx1032__) || defined(__gfx1033__) || \
|
||||
defined(__gfx1034__) || defined(__gfx1035__) || defined(__gfx1036__) || defined(__gfx1037__)
|
||||
#define RDNA2
|
||||
#endif
|
||||
|
||||
#if defined(__gfx1010__) || defined(__gfx1012__)
|
||||
#define RDNA1
|
||||
#endif
|
||||
|
||||
#ifndef __has_builtin
|
||||
#define __has_builtin(x) 0
|
||||
#endif
|
||||
|
||||
typedef hip_bfloat16 nv_bfloat16;
|
||||
@@ -0,0 +1,137 @@
|
||||
#pragma once
|
||||
|
||||
#include <musa_runtime.h>
|
||||
#include <musa.h>
|
||||
#include <mublas.h>
|
||||
#include <musa_bf16.h>
|
||||
#include <musa_fp16.h>
|
||||
#define CUBLAS_COMPUTE_16F CUDA_R_16F
|
||||
#define CUBLAS_COMPUTE_32F CUDA_R_32F
|
||||
#define CUBLAS_COMPUTE_32F_FAST_16F MUBLAS_COMPUTE_32F_FAST_16F
|
||||
#define CUBLAS_GEMM_DEFAULT MUBLAS_GEMM_DEFAULT
|
||||
#define CUBLAS_GEMM_DEFAULT_TENSOR_OP MUBLAS_GEMM_DEFAULT
|
||||
#define CUBLAS_OP_N MUBLAS_OP_N
|
||||
#define CUBLAS_OP_T MUBLAS_OP_T
|
||||
#define CUBLAS_STATUS_SUCCESS MUBLAS_STATUS_SUCCESS
|
||||
#define CUBLAS_TF32_TENSOR_OP_MATH MUBLAS_MATH_MODE_DEFAULT
|
||||
#define CUDA_R_16F MUSA_R_16F
|
||||
#define CUDA_R_32F MUSA_R_32F
|
||||
#define cublasComputeType_t cudaDataType_t
|
||||
#define cublasCreate mublasCreate
|
||||
#define cublasDestroy mublasDestroy
|
||||
#define cublasGemmEx mublasGemmEx
|
||||
#define cublasGemmBatchedEx mublasGemmBatchedEx
|
||||
#define cublasGemmStridedBatchedEx mublasGemmStridedBatchedEx
|
||||
#define cublasHandle_t mublasHandle_t
|
||||
#define cublasSetMathMode mublasSetMathMode
|
||||
#define cublasSetStream mublasSetStream
|
||||
#define cublasSgemm mublasSgemm
|
||||
#define cublasStatus_t mublasStatus_t
|
||||
#define cublasOperation_t mublasOperation_t
|
||||
#define cublasGetStatusString mublasStatus_to_string
|
||||
#define cudaDataType_t musaDataType_t
|
||||
#define cudaDeviceCanAccessPeer musaDeviceCanAccessPeer
|
||||
#define cudaDeviceDisablePeerAccess musaDeviceDisablePeerAccess
|
||||
#define cudaDeviceEnablePeerAccess musaDeviceEnablePeerAccess
|
||||
#define cudaDeviceProp musaDeviceProp
|
||||
#define cudaDeviceSynchronize musaDeviceSynchronize
|
||||
#define cudaError_t musaError_t
|
||||
#define cudaErrorPeerAccessAlreadyEnabled musaErrorPeerAccessAlreadyEnabled
|
||||
#define cudaErrorPeerAccessNotEnabled musaErrorPeerAccessNotEnabled
|
||||
#define cudaEventCreateWithFlags musaEventCreateWithFlags
|
||||
#define cudaEventDisableTiming musaEventDisableTiming
|
||||
#define cudaEventRecord musaEventRecord
|
||||
#define cudaEventSynchronize musaEventSynchronize
|
||||
#define cudaEvent_t musaEvent_t
|
||||
#define cudaEventDestroy musaEventDestroy
|
||||
#define cudaFree musaFree
|
||||
#define cudaFreeHost musaFreeHost
|
||||
#define cudaGetDevice musaGetDevice
|
||||
#define cudaGetDeviceCount musaGetDeviceCount
|
||||
#define cudaGetDeviceProperties musaGetDeviceProperties
|
||||
#define cudaGetErrorString musaGetErrorString
|
||||
#define cudaGetLastError musaGetLastError
|
||||
#define cudaHostRegister musaHostRegister
|
||||
#define cudaHostRegisterPortable musaHostRegisterPortable
|
||||
#define cudaHostRegisterReadOnly musaHostRegisterReadOnly
|
||||
#define cudaHostUnregister musaHostUnregister
|
||||
#define cudaLaunchHostFunc musaLaunchHostFunc
|
||||
#define cudaMalloc musaMalloc
|
||||
#define cudaMallocHost musaMallocHost
|
||||
#define cudaMallocManaged musaMallocManaged
|
||||
#define cudaMemcpy musaMemcpy
|
||||
#define cudaMemcpyAsync musaMemcpyAsync
|
||||
#define cudaMemcpyPeerAsync musaMemcpyPeerAsync
|
||||
#define cudaMemcpy2DAsync musaMemcpy2DAsync
|
||||
#define cudaMemcpyDeviceToDevice musaMemcpyDeviceToDevice
|
||||
#define cudaMemcpyDeviceToHost musaMemcpyDeviceToHost
|
||||
#define cudaMemcpyHostToDevice musaMemcpyHostToDevice
|
||||
#define cudaMemcpyKind musaMemcpyKind
|
||||
#define cudaMemset musaMemset
|
||||
#define cudaMemsetAsync musaMemsetAsync
|
||||
#define cudaMemGetInfo musaMemGetInfo
|
||||
#define cudaOccupancyMaxPotentialBlockSize musaOccupancyMaxPotentialBlockSize
|
||||
#define cudaSetDevice musaSetDevice
|
||||
#define cudaStreamCreateWithFlags musaStreamCreateWithFlags
|
||||
#define cudaStreamDestroy musaStreamDestroy
|
||||
#define cudaStreamFireAndForget musaStreamFireAndForget
|
||||
#define cudaStreamNonBlocking musaStreamNonBlocking
|
||||
#define cudaStreamPerThread musaStreamPerThread
|
||||
#define cudaStreamSynchronize musaStreamSynchronize
|
||||
#define cudaStreamWaitEvent musaStreamWaitEvent
|
||||
#define cudaStream_t musaStream_t
|
||||
#define cudaSuccess musaSuccess
|
||||
|
||||
// Additional mappings for MUSA virtual memory pool
|
||||
#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED MU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED
|
||||
#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE MU_MEM_ACCESS_FLAGS_PROT_READWRITE
|
||||
#define CU_MEM_ALLOC_GRANULARITY_RECOMMENDED MU_MEM_ALLOC_GRANULARITY_RECOMMENDED
|
||||
#define CU_MEM_ALLOCATION_TYPE_PINNED MU_MEM_ALLOCATION_TYPE_PINNED
|
||||
#define CU_MEM_LOCATION_TYPE_DEVICE MU_MEM_LOCATION_TYPE_DEVICE
|
||||
#define CUdevice MUdevice
|
||||
#define CUdeviceptr MUdeviceptr
|
||||
#define CUmemAccessDesc MUmemAccessDesc
|
||||
#define CUmemAllocationProp MUmemAllocationProp
|
||||
#define CUmemGenericAllocationHandle MUmemGenericAllocationHandle
|
||||
#define cuDeviceGet muDeviceGet
|
||||
#define cuDeviceGetAttribute muDeviceGetAttribute
|
||||
#define cuMemAddressFree muMemAddressFree
|
||||
#define cuMemAddressReserve muMemAddressReserve
|
||||
#define cuMemCreate muMemCreate
|
||||
#define cuMemGetAllocationGranularity muMemGetAllocationGranularity
|
||||
#define cuMemMap muMemMap
|
||||
#define cuMemRelease muMemRelease
|
||||
#define cuMemSetAccess muMemSetAccess
|
||||
#define cuMemUnmap muMemUnmap
|
||||
#define cudaFuncAttributeMaxDynamicSharedMemorySize musaFuncAttributeMaxDynamicSharedMemorySize
|
||||
#define cudaFuncSetAttribute musaFuncSetAttribute
|
||||
#define cudaMemcpy3DPeerParms musaMemcpy3DPeerParms
|
||||
#define make_cudaExtent make_musaExtent
|
||||
#define make_cudaPitchedPtr make_musaPitchedPtr
|
||||
|
||||
// Additional mappings for MUSA graphs
|
||||
#define CUDA_SUCCESS MUSA_SUCCESS
|
||||
#define CUresult MUresult
|
||||
#define cuGetErrorString muGetErrorString
|
||||
#define cudaErrorGraphExecUpdateFailure musaErrorGraphExecUpdateFailure
|
||||
#define cudaErrorInvalidDeviceFunction musaErrorInvalidDeviceFunction
|
||||
#define cudaGraphDestroy musaGraphDestroy
|
||||
#define cudaGraphExecDestroy musaGraphExecDestroy
|
||||
#define cudaGraphExec_t musaGraphExec_t
|
||||
#define cudaGraphExecUpdate musaGraphExecUpdate
|
||||
#define cudaGraphExecUpdateResultInfo musaGraphExecUpdateResult
|
||||
#define cudaGraphGetNodes musaGraphGetNodes
|
||||
#define cudaGraphInstantiate musaGraphInstantiate
|
||||
#define cudaGraphKernelNodeGetParams musaGraphKernelNodeGetParams
|
||||
#define cudaGraphKernelNodeSetParams musaGraphKernelNodeSetParams
|
||||
#define cudaGraphLaunch musaGraphLaunch
|
||||
#define cudaGraphNodeGetType musaGraphNodeGetType
|
||||
#define cudaGraphNode_t musaGraphNode_t
|
||||
#define cudaGraphNodeType musaGraphNodeType
|
||||
#define cudaGraphNodeTypeKernel musaGraphNodeTypeKernel
|
||||
#define cudaGraph_t musaGraph_t
|
||||
#define cudaKernelNodeParams musaKernelNodeParams
|
||||
#define cudaStreamCaptureModeRelaxed musaStreamCaptureModeRelaxed
|
||||
#define cudaStreamEndCapture musaStreamEndCapture
|
||||
|
||||
typedef mt_bfloat16 nv_bfloat16;
|
||||
@@ -0,0 +1,13 @@
|
||||
#ifndef CPUINFER_VENDOR_VENDOR_H
|
||||
#define CPUINFER_VENDOR_VENDOR_H
|
||||
|
||||
#ifdef USE_CUDA
|
||||
#include "cuda.h"
|
||||
#elif USE_HIP
|
||||
#define __HIP_PLATFORM_AMD__
|
||||
#include "hip.h"
|
||||
#elif USE_MUSA
|
||||
#include "musa.h"
|
||||
#endif
|
||||
|
||||
#endif // CPUINFER_VENDOR_VENDOR_H
|
||||
@@ -0,0 +1,71 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : Azure-Tang, Boxin Zhang
|
||||
* @Date : 2024-07-25 13:38:30
|
||||
* @Version : 0.2.2
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
|
||||
#include "custom_gguf/ops.h"
|
||||
#ifdef KTRANSFORMERS_USE_CUDA
|
||||
#include "gptq_marlin/ops.h"
|
||||
#endif
|
||||
// Python bindings
|
||||
#include <pybind11/pybind11.h>
|
||||
#include <pybind11/stl.h>
|
||||
#include <torch/library.h>
|
||||
#include <torch/extension.h>
|
||||
#include <torch/torch.h>
|
||||
// namespace py = pybind11;
|
||||
|
||||
PYBIND11_MODULE(KTransformersOps, m) {
|
||||
|
||||
m.def("dequantize_q8_0", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) {
|
||||
torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype);
|
||||
return dequantize_q8_0((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype);
|
||||
}, "Function to dequantize q8_0 data.",
|
||||
py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype"));
|
||||
|
||||
m.def("dequantize_q6_k", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) {
|
||||
torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype);
|
||||
return dequantize_q6_k((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype);
|
||||
}, "Function to dequantize q6_k data.",
|
||||
py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype"));
|
||||
|
||||
m.def("dequantize_q5_k", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) {
|
||||
torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype);
|
||||
return dequantize_q5_k((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype);
|
||||
}, "Function to dequantize q5_k data.",
|
||||
py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype"));
|
||||
|
||||
m.def("dequantize_q4_k", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) {
|
||||
torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype);
|
||||
return dequantize_q4_k((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype);
|
||||
}, "Function to dequantize q4_k data.",
|
||||
py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype"));
|
||||
|
||||
m.def("dequantize_q3_k", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) {
|
||||
torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype);
|
||||
return dequantize_q3_k((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype);
|
||||
}, "Function to dequantize q3_k data.",
|
||||
py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype"));
|
||||
|
||||
m.def("dequantize_q2_k", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) {
|
||||
torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype);
|
||||
return dequantize_q2_k((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype);
|
||||
}, "Function to dequantize q2_k data.",
|
||||
py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype"));
|
||||
|
||||
m.def("dequantize_iq4_xs", [](const intptr_t data, int num_bytes, int blk_size, const int ele_per_blk, torch::Device device, py::object target_dtype) {
|
||||
torch::Dtype dtype = torch::python::detail::py_object_to_dtype(target_dtype);
|
||||
return dequantize_iq4_xs((int8_t*)data, num_bytes, blk_size, ele_per_blk, device, dtype);
|
||||
}, "Function to dequantize iq4_xs data.",
|
||||
py::arg("data"), py::arg("num_bytes"), py::arg("blk_size"), py::arg("ele_per_blk"), py::arg("device"), py::arg("target_dtype"));
|
||||
|
||||
#ifdef KTRANSFORMERS_USE_CUDA
|
||||
m.def("gptq_marlin_gemm", &gptq_marlin_gemm, "Function to perform GEMM using Marlin quantization.",
|
||||
py::arg("a"), py::arg("b_q_weight"), py::arg("b_scales"), py::arg("g_idx"),
|
||||
py::arg("perm"), py::arg("workspace"), py::arg("num_bits"), py::arg("size_m"),
|
||||
py::arg("size_n"), py::arg("size_k"), py::arg("is_k_full"));
|
||||
#endif
|
||||
}
|
||||
@@ -0,0 +1,882 @@
|
||||
/*
|
||||
* @Description :
|
||||
* @Author : Azure-Tang, Boxin Zhang
|
||||
* @Date : 2024-07-25 13:38:30
|
||||
* @Version : 0.2.2
|
||||
* Adapted from https://github.com/ggerganov/ggml/blob/fca1caafea7de9fbd7efc733b9818f9cf2da3050/src/ggml-quants.c
|
||||
* Copyright (c) 2023-2024 The ggml authors
|
||||
* Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
*/
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <torch/library.h>
|
||||
#include <torch/extension.h>
|
||||
#include <torch/torch.h>
|
||||
#include <cstdint>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
|
||||
#ifdef __HIP_PLATFORM_AMD__
|
||||
typedef __hip_bfloat16 nv_bfloat16;
|
||||
#endif
|
||||
|
||||
__global__ void dequantize_q8_0_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x){
|
||||
float* __restrict__ output_blk = (float*)(output + block_id * ele_per_blk);
|
||||
const int8_t* cur_block = data + block_id * blk_size;
|
||||
float scale = __half2float(*((half*)cur_block));
|
||||
cur_block += 2;
|
||||
for (int i = 0; i < ele_per_blk; i++){
|
||||
output_blk[i] = scale * cur_block[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_q8_0_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x) {
|
||||
__half* __restrict__ output_blk = (__half*)(output + block_id * ele_per_blk);
|
||||
const int8_t* cur_block = data + block_id * blk_size;
|
||||
float scale = __half2float(*((half*)cur_block));
|
||||
cur_block += 2;
|
||||
for (int i = 0; i < ele_per_blk; i++) {
|
||||
output_blk[i] = __float2half(scale * cur_block[i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_q8_0_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x) {
|
||||
nv_bfloat16* __restrict__ output_blk = (nv_bfloat16*)(output + block_id * ele_per_blk);
|
||||
const int8_t* cur_block = data + block_id * blk_size;
|
||||
float scale = __half2float(*((half*)cur_block));
|
||||
cur_block += 2;
|
||||
for (int i = 0; i < ele_per_blk; i++) {
|
||||
output_blk[i] = __float2bfloat16(scale * cur_block[i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// __device__ void get_scale_min_k4(int j, const uint8_t * __restrict__ q, uint8_t * __restrict__ d, uint8_t * __restrict__ m) {
|
||||
__device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t * __restrict__ d, uint8_t * __restrict__ m) {
|
||||
if (j < 4) {
|
||||
*d = q[j] & 63; *m = q[j + 4] & 63;
|
||||
} else {
|
||||
*d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4);
|
||||
*m = (q[j+4] >> 4) | ((q[j-0] >> 6) << 4);
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_q2_k_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id=global_idx; block_id<num_blocks; block_id+= blockDim.x * gridDim.x){
|
||||
float* __restrict__ output_blk = (float*)(output + block_id * ele_per_blk);
|
||||
|
||||
const float d = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 80)));
|
||||
const float min = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 82)));
|
||||
|
||||
const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 16);
|
||||
|
||||
int is = 0;
|
||||
float dl, ml;
|
||||
|
||||
for (int n = 0; n < 256; n += 128) {
|
||||
int shift = 0;
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
uint8_t* scales = (uint8_t*)(data + block_id * blk_size + (is++));
|
||||
uint8_t sc = *scales;
|
||||
dl = d * (sc & 0xF); ml = min * (sc >> 4);
|
||||
for (int l = 0; l < 16; ++l) *output_blk++ = dl * ((int8_t)((q[l] >> shift) & 3)) - ml;
|
||||
|
||||
scales = (uint8_t*)(data + block_id * blk_size + (is++));
|
||||
sc = *scales;
|
||||
|
||||
dl = d * (sc & 0xF); ml = min * (sc >> 4);
|
||||
for (int l = 0; l < 16; ++l) *output_blk++ = dl * ((int8_t)((q[l+16] >> shift) & 3)) - ml;
|
||||
|
||||
shift += 2;
|
||||
}
|
||||
q += 32;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_q2_k_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id=global_idx; block_id<num_blocks; block_id+= blockDim.x * gridDim.x){
|
||||
__half* __restrict__ output_blk = (__half*)(output + block_id * ele_per_blk);
|
||||
|
||||
const float d = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 80)));
|
||||
const float min = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 82)));
|
||||
|
||||
const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 16);
|
||||
|
||||
int is = 0;
|
||||
float dl, ml;
|
||||
|
||||
for (int n = 0; n < 256; n += 128) {
|
||||
int shift = 0;
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
uint8_t* scales = (uint8_t*)(data + block_id * blk_size + (is++));
|
||||
uint8_t sc = *scales;
|
||||
dl = d * (sc & 0xF); ml = min * (sc >> 4);
|
||||
for (int l = 0; l < 16; ++l) *output_blk++ = __float2half(dl * ((int8_t)((q[l] >> shift) & 3)) - ml);
|
||||
|
||||
scales = (uint8_t*)(data + block_id * blk_size + (is++));
|
||||
sc = *scales;
|
||||
|
||||
dl = d * (sc & 0xF); ml = min * (sc >> 4);
|
||||
for (int l = 0; l < 16; ++l) *output_blk++ = __float2half(dl * ((int8_t)((q[l+16] >> shift) & 3)) - ml);
|
||||
|
||||
shift += 2;
|
||||
}
|
||||
q += 32;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_q2_k_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id=global_idx; block_id<num_blocks; block_id+= blockDim.x * gridDim.x){
|
||||
nv_bfloat16* __restrict__ output_blk = (nv_bfloat16*)(output + block_id * ele_per_blk);
|
||||
|
||||
const float d = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 80)));
|
||||
const float min = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 82)));
|
||||
|
||||
const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 16);
|
||||
|
||||
int is = 0;
|
||||
float dl, ml;
|
||||
|
||||
for (int n = 0; n < 256; n += 128) {
|
||||
int shift = 0;
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
uint8_t* scales = (uint8_t*)(data + block_id * blk_size + (is++));
|
||||
uint8_t sc = *scales;
|
||||
dl = d * (sc & 0xF); ml = min * (sc >> 4);
|
||||
for (int l = 0; l < 16; ++l) *output_blk++ = __float2bfloat16(dl * ((int8_t)((q[l] >> shift) & 3)) - ml);
|
||||
|
||||
scales = (uint8_t*)(data + block_id * blk_size + (is++));
|
||||
sc = *scales;
|
||||
|
||||
dl = d * (sc & 0xF); ml = min * (sc >> 4);
|
||||
for (int l = 0; l < 16; ++l) *output_blk++ = __float2bfloat16(dl * ((int8_t)((q[l+16] >> shift) & 3)) - ml);
|
||||
|
||||
shift += 2;
|
||||
}
|
||||
q += 32;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_q3_k_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const uint32_t kmask1 = 0x03030303;
|
||||
const uint32_t kmask2 = 0x0f0f0f0f;
|
||||
for (long long block_id=global_idx; block_id<num_blocks; block_id+= blockDim.x * gridDim.x){
|
||||
float* __restrict__ output_blk = (float*)(output + block_id * ele_per_blk);
|
||||
|
||||
uint32_t aux[4];
|
||||
const int8_t * scales = (const int8_t*)aux;
|
||||
const float d_all = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 108)));
|
||||
|
||||
const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 32);
|
||||
const uint8_t * __restrict__ hm = (uint8_t*)(data + block_id * blk_size + 0);
|
||||
uint8_t m = 1;
|
||||
|
||||
|
||||
uint8_t* block_scales = (uint8_t*)(data + block_id * blk_size + 96);
|
||||
|
||||
for (int i = 0; i < 3; i++) {
|
||||
aux[i] = 0;
|
||||
for (int j = 0; j < 4; j++) {
|
||||
aux[i] |= ((uint32_t)block_scales[i * 4 + j]) << (j * 8);
|
||||
}
|
||||
}
|
||||
|
||||
uint32_t tmp = aux[2];
|
||||
aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4);
|
||||
aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4);
|
||||
aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4);
|
||||
aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4);
|
||||
|
||||
int is = 0;
|
||||
float dl;
|
||||
for (int n = 0; n < 256; n += 128) {
|
||||
int shift = 0;
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
|
||||
dl = d_all * (scales[is++] - 32);
|
||||
for (int l = 0; l < 16; ++l) {
|
||||
*output_blk++ = dl * ((int8_t)((q[l+ 0] >> shift) & 3) - ((hm[l+ 0] & m) ? 0 : 4));
|
||||
}
|
||||
|
||||
dl = d_all * (scales[is++] - 32);
|
||||
for (int l = 0; l < 16; ++l) {
|
||||
*output_blk++ = dl * ((int8_t)((q[l+16] >> shift) & 3) - ((hm[l+16] & m) ? 0 : 4));
|
||||
}
|
||||
|
||||
shift += 2;
|
||||
m <<= 1;
|
||||
}
|
||||
q += 32;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_q3_k_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const uint32_t kmask1 = 0x03030303;
|
||||
const uint32_t kmask2 = 0x0f0f0f0f;
|
||||
for (long long block_id=global_idx; block_id<num_blocks; block_id+= blockDim.x * gridDim.x){
|
||||
__half* __restrict__ output_blk = (__half*)(output + block_id * ele_per_blk);
|
||||
|
||||
uint32_t aux[4];
|
||||
const int8_t * scales = (const int8_t*)aux;
|
||||
const float d_all = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 108)));
|
||||
|
||||
const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 32);
|
||||
const uint8_t * __restrict__ hm = (uint8_t*)(data + block_id * blk_size + 0);
|
||||
uint8_t m = 1;
|
||||
|
||||
|
||||
uint8_t* block_scales = (uint8_t*)(data + block_id * blk_size + 96);
|
||||
|
||||
for (int i = 0; i < 3; i++) {
|
||||
aux[i] = 0;
|
||||
for (int j = 0; j < 4; j++) {
|
||||
aux[i] |= ((uint32_t)block_scales[i * 4 + j]) << (j * 8);
|
||||
}
|
||||
}
|
||||
|
||||
uint32_t tmp = aux[2];
|
||||
aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4);
|
||||
aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4);
|
||||
aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4);
|
||||
aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4);
|
||||
|
||||
int is = 0;
|
||||
float dl;
|
||||
for (int n = 0; n < 256; n += 128) {
|
||||
int shift = 0;
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
|
||||
dl = d_all * (scales[is++] - 32);
|
||||
for (int l = 0; l < 16; ++l) {
|
||||
*output_blk++ = __float2half(dl * ((int8_t)((q[l+ 0] >> shift) & 3) - ((hm[l+ 0] & m) ? 0 : 4)));
|
||||
}
|
||||
|
||||
dl = d_all * (scales[is++] - 32);
|
||||
for (int l = 0; l < 16; ++l) {
|
||||
*output_blk++ = __float2half(dl * ((int8_t)((q[l+16] >> shift) & 3) - ((hm[l+16] & m) ? 0 : 4)));
|
||||
}
|
||||
|
||||
shift += 2;
|
||||
m <<= 1;
|
||||
}
|
||||
q += 32;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_q3_k_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const uint32_t kmask1 = 0x03030303;
|
||||
const uint32_t kmask2 = 0x0f0f0f0f;
|
||||
for (long long block_id=global_idx; block_id<num_blocks; block_id+= blockDim.x * gridDim.x){
|
||||
nv_bfloat16* __restrict__ output_blk = (nv_bfloat16*)(output + block_id * ele_per_blk);
|
||||
|
||||
uint32_t aux[4];
|
||||
const int8_t * scales = (const int8_t*)aux;
|
||||
const float d_all = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 108)));
|
||||
|
||||
const uint8_t * __restrict__ q = (uint8_t*)(data + block_id * blk_size + 32);
|
||||
const uint8_t * __restrict__ hm = (uint8_t*)(data + block_id * blk_size + 0);
|
||||
uint8_t m = 1;
|
||||
|
||||
|
||||
uint8_t* block_scales = (uint8_t*)(data + block_id * blk_size + 96);
|
||||
|
||||
for (int i = 0; i < 3; i++) {
|
||||
aux[i] = 0;
|
||||
for (int j = 0; j < 4; j++) {
|
||||
aux[i] |= ((uint32_t)block_scales[i * 4 + j]) << (j * 8);
|
||||
}
|
||||
}
|
||||
|
||||
uint32_t tmp = aux[2];
|
||||
aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4);
|
||||
aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4);
|
||||
aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4);
|
||||
aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4);
|
||||
|
||||
int is = 0;
|
||||
float dl;
|
||||
for (int n = 0; n < 256; n += 128) {
|
||||
int shift = 0;
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
|
||||
dl = d_all * (scales[is++] - 32);
|
||||
for (int l = 0; l < 16; ++l) {
|
||||
*output_blk++ = __float2bfloat16(dl * ((int8_t)((q[l+ 0] >> shift) & 3) - ((hm[l+ 0] & m) ? 0 : 4)));
|
||||
}
|
||||
|
||||
dl = d_all * (scales[is++] - 32);
|
||||
for (int l = 0; l < 16; ++l) {
|
||||
*output_blk++ = __float2bfloat16(dl * ((int8_t)((q[l+16] >> shift) & 3) - ((hm[l+16] & m) ? 0 : 4)));
|
||||
}
|
||||
|
||||
shift += 2;
|
||||
m <<= 1;
|
||||
}
|
||||
q += 32;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
__global__ void dequantize_q4_k_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id=global_idx; block_id<num_blocks; block_id+=blockDim.x * gridDim.x){
|
||||
float* __restrict__ output_blk = (float*)(output + block_id * ele_per_blk);
|
||||
// const uint8_t * q = data[i].qs;
|
||||
const uint8_t * q = (uint8_t*)(data + block_id * 144 + 16);
|
||||
|
||||
const float d = __half2float(*(reinterpret_cast<const half*>(data + block_id * 144 + 0)));
|
||||
const float min = __half2float(*(reinterpret_cast<const half*>(data + block_id * 144 + 2)));
|
||||
int is = 0;
|
||||
uint8_t sc, m;
|
||||
for (int j = 0; j < ele_per_blk; j += 64) {
|
||||
uint8_t* scales = (uint8_t*)(data + block_id * 144 + 4);
|
||||
get_scale_min_k4(is + 0, scales, &sc, &m);
|
||||
const float d1 = d * sc; const float m1 = min * m;
|
||||
get_scale_min_k4(is + 1, scales, &sc, &m);
|
||||
const float d2 = d * sc; const float m2 = min * m;
|
||||
for (int l = 0; l < 32; ++l) *output_blk++ = d1 * (q[l] & 0xF) - m1;
|
||||
for (int l = 0; l < 32; ++l) *output_blk++ = d2 * (q[l] >> 4) - m2;
|
||||
q += 32; is += 2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_q4_k_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id=global_idx; block_id<num_blocks; block_id+=blockDim.x * gridDim.x){
|
||||
__half* __restrict__ output_blk = (__half*)(output + block_id * ele_per_blk);
|
||||
// const uint8_t * q = data[i].qs;
|
||||
const uint8_t * q = (uint8_t*)(data + block_id * 144 + 16);
|
||||
|
||||
const float d = __half2float(*(reinterpret_cast<const half*>(data + block_id * 144 + 0)));
|
||||
const float min = __half2float(*(reinterpret_cast<const half*>(data + block_id * 144 + 2)));
|
||||
int is = 0;
|
||||
uint8_t sc, m;
|
||||
for (int j = 0; j < ele_per_blk; j += 64) {
|
||||
uint8_t* scales = (uint8_t*)(data + block_id * 144 + 4);
|
||||
get_scale_min_k4(is + 0, scales, &sc, &m);
|
||||
const float d1 = d * sc; const float m1 = min * m;
|
||||
get_scale_min_k4(is + 1, scales, &sc, &m);
|
||||
const float d2 = d * sc; const float m2 = min * m;
|
||||
for (int l = 0; l < 32; ++l) *output_blk++ = __float2half(d1 * (q[l] & 0xF) - m1);
|
||||
for (int l = 0; l < 32; ++l) *output_blk++ = __float2half(d2 * (q[l] >> 4) - m2);
|
||||
q += 32; is += 2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_q4_k_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id=global_idx; block_id<num_blocks; block_id+=blockDim.x * gridDim.x){
|
||||
nv_bfloat16* __restrict__ output_blk = (nv_bfloat16*)(output + block_id * ele_per_blk);
|
||||
// const uint8_t * q = data[i].qs;
|
||||
const uint8_t * q = (uint8_t*)(data + block_id * 144 + 16);
|
||||
|
||||
const float d = __half2float(*(reinterpret_cast<const half*>(data + block_id * 144 + 0)));
|
||||
const float min = __half2float(*(reinterpret_cast<const half*>(data + block_id * 144 + 2)));
|
||||
int is = 0;
|
||||
uint8_t sc, m;
|
||||
for (int j = 0; j < ele_per_blk; j += 64) {
|
||||
uint8_t* scales = (uint8_t*)(data + block_id * 144 + 4);
|
||||
get_scale_min_k4(is + 0, scales, &sc, &m);
|
||||
const float d1 = d * sc; const float m1 = min * m;
|
||||
get_scale_min_k4(is + 1, scales, &sc, &m);
|
||||
const float d2 = d * sc; const float m2 = min * m;
|
||||
for (int l = 0; l < 32; ++l) *output_blk++ = __float2bfloat16(d1 * (q[l] & 0xF) - m1);
|
||||
for (int l = 0; l < 32; ++l) *output_blk++ = __float2bfloat16(d2 * (q[l] >> 4) - m2);
|
||||
q += 32; is += 2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_q5_k_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x){
|
||||
float* __restrict__ output_blk = (float*)(output + block_id * ele_per_blk);
|
||||
|
||||
const float d = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 0)));
|
||||
const float min = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 2)));
|
||||
|
||||
const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 16);
|
||||
const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size + 48);
|
||||
|
||||
int is = 0;
|
||||
uint8_t sc, m;
|
||||
uint8_t u1 = 1, u2 = 2;
|
||||
uint8_t* scales = (uint8_t*)(data + block_id * blk_size + 4);
|
||||
|
||||
for (int j = 0; j < 256; j += 64) {
|
||||
get_scale_min_k4(is + 0, scales, &sc, &m);
|
||||
const float d1 = d * sc; const float m1 = min * m;
|
||||
get_scale_min_k4(is + 1, scales, &sc, &m);
|
||||
const float d2 = d * sc; const float m2 = min * m;
|
||||
for (int l = 0; l < 32; ++l) *output_blk++ = d1 * ((ql[l] & 0xF) + (qh[l] & u1 ? 16 : 0)) - m1;
|
||||
for (int l = 0; l < 32; ++l) *output_blk++ = d2 * ((ql[l] >> 4) + (qh[l] & u2 ? 16 : 0)) - m2;
|
||||
ql += 32; is += 2;
|
||||
u1 <<= 2; u2 <<= 2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_q5_k_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x){
|
||||
__half* __restrict__ output_blk = (__half*)(output + block_id * ele_per_blk);
|
||||
|
||||
const float d = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 0)));
|
||||
const float min = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 2)));
|
||||
|
||||
const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 16);
|
||||
const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size + 48);
|
||||
|
||||
int is = 0;
|
||||
uint8_t sc, m;
|
||||
uint8_t u1 = 1, u2 = 2;
|
||||
uint8_t* scales = (uint8_t*)(data + block_id * blk_size + 4);
|
||||
|
||||
for (int j = 0; j < 256; j += 64) {
|
||||
get_scale_min_k4(is + 0, scales, &sc, &m);
|
||||
const float d1 = d * sc; const float m1 = min * m;
|
||||
get_scale_min_k4(is + 1, scales, &sc, &m);
|
||||
const float d2 = d * sc; const float m2 = min * m;
|
||||
for (int l = 0; l < 32; ++l) *output_blk++ = __float2half(d1 * ((ql[l] & 0xF) + (qh[l] & u1 ? 16 : 0)) - m1);
|
||||
for (int l = 0; l < 32; ++l) *output_blk++ = __float2half(d2 * ((ql[l] >> 4) + (qh[l] & u2 ? 16 : 0)) - m2);
|
||||
ql += 32; is += 2;
|
||||
u1 <<= 2; u2 <<= 2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_q5_k_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id = global_idx; block_id < num_blocks; block_id += blockDim.x * gridDim.x){
|
||||
nv_bfloat16* __restrict__ output_blk = (nv_bfloat16*)(output + block_id * ele_per_blk);
|
||||
|
||||
const float d = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 0)));
|
||||
const float min = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 2)));
|
||||
|
||||
const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 16);
|
||||
const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size + 48);
|
||||
|
||||
int is = 0;
|
||||
uint8_t sc, m;
|
||||
uint8_t u1 = 1, u2 = 2;
|
||||
uint8_t* scales = (uint8_t*)(data + block_id * blk_size + 4);
|
||||
|
||||
for (int j = 0; j < 256; j += 64) {
|
||||
get_scale_min_k4(is + 0, scales, &sc, &m);
|
||||
const float d1 = d * sc; const float m1 = min * m;
|
||||
get_scale_min_k4(is + 1, scales, &sc, &m);
|
||||
const float d2 = d * sc; const float m2 = min * m;
|
||||
for (int l = 0; l < 32; ++l) *output_blk++ = __float2bfloat16(d1 * ((ql[l] & 0xF) + (qh[l] & u1 ? 16 : 0)) - m1);
|
||||
for (int l = 0; l < 32; ++l) *output_blk++ = __float2bfloat16(d2 * ((ql[l] >> 4) + (qh[l] & u2 ? 16 : 0)) - m2);
|
||||
ql += 32; is += 2;
|
||||
u1 <<= 2; u2 <<= 2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_q6_k_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id=global_idx; block_id<num_blocks;block_id+=blockDim.x * gridDim.x){
|
||||
float* __restrict__ output_blk = (float*)(output + block_id * ele_per_blk);
|
||||
const float d = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 208)));
|
||||
|
||||
const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size);
|
||||
const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 128);
|
||||
const int8_t * __restrict__ sc = (int8_t*)(data + block_id * blk_size + 192);
|
||||
|
||||
|
||||
for (int n = 0; n < ele_per_blk; n += 128) {
|
||||
for (int l = 0; l < 32; ++l) {
|
||||
int is = l/16;
|
||||
const int8_t q1 = (int8_t)((ql[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32;
|
||||
const int8_t q2 = (int8_t)((ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32;
|
||||
const int8_t q3 = (int8_t)((ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32;
|
||||
const int8_t q4 = (int8_t)((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32;
|
||||
output_blk[l + 0] = d * sc[is + 0] * q1;
|
||||
output_blk[l + 32] = d * sc[is + 2] * q2;
|
||||
output_blk[l + 64] = d * sc[is + 4] * q3;
|
||||
output_blk[l + 96] = d * sc[is + 6] * q4;
|
||||
}
|
||||
output_blk += 128;
|
||||
ql += 64;
|
||||
qh += 32;
|
||||
sc += 8;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_q6_k_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id=global_idx; block_id<num_blocks;block_id+=blockDim.x * gridDim.x){
|
||||
__half* __restrict__ output_blk = (__half*)(output + block_id * ele_per_blk);
|
||||
const float d = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 208)));
|
||||
|
||||
const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size);
|
||||
const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 128);
|
||||
const int8_t * __restrict__ sc = (int8_t*)(data + block_id * blk_size + 192);
|
||||
|
||||
|
||||
for (int n = 0; n < ele_per_blk; n += 128) {
|
||||
for (int l = 0; l < 32; ++l) {
|
||||
int is = l/16;
|
||||
const int8_t q1 = (int8_t)((ql[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32;
|
||||
const int8_t q2 = (int8_t)((ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32;
|
||||
const int8_t q3 = (int8_t)((ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32;
|
||||
const int8_t q4 = (int8_t)((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32;
|
||||
output_blk[l + 0] = __float2half(d * sc[is + 0] * q1);
|
||||
output_blk[l + 32] = __float2half(d * sc[is + 2] * q2);
|
||||
output_blk[l + 64] = __float2half(d * sc[is + 4] * q3);
|
||||
output_blk[l + 96] = __float2half(d * sc[is + 6] * q4);
|
||||
}
|
||||
output_blk += 128;
|
||||
ql += 64;
|
||||
qh += 32;
|
||||
sc += 8;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_q6_k_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id=global_idx; block_id<num_blocks;block_id+=blockDim.x * gridDim.x){
|
||||
nv_bfloat16* __restrict__ output_blk = (nv_bfloat16*)(output + block_id * ele_per_blk);
|
||||
const float d = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size + 208)));
|
||||
|
||||
const uint8_t * __restrict__ ql = (uint8_t*)(data + block_id * blk_size);
|
||||
const uint8_t * __restrict__ qh = (uint8_t*)(data + block_id * blk_size + 128);
|
||||
const int8_t * __restrict__ sc = (int8_t*)(data + block_id * blk_size + 192);
|
||||
|
||||
|
||||
for (int n = 0; n < ele_per_blk; n += 128) {
|
||||
for (int l = 0; l < 32; ++l) {
|
||||
int is = l/16;
|
||||
const int8_t q1 = (int8_t)((ql[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32;
|
||||
const int8_t q2 = (int8_t)((ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32;
|
||||
const int8_t q3 = (int8_t)((ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32;
|
||||
const int8_t q4 = (int8_t)((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32;
|
||||
output_blk[l + 0] = __float2bfloat16(d * sc[is + 0] * q1);
|
||||
output_blk[l + 32] = __float2bfloat16(d * sc[is + 2] * q2);
|
||||
output_blk[l + 64] = __float2bfloat16(d * sc[is + 4] * q3);
|
||||
output_blk[l + 96] = __float2bfloat16(d * sc[is + 6] * q4);
|
||||
}
|
||||
output_blk += 128;
|
||||
ql += 64;
|
||||
qh += 32;
|
||||
sc += 8;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static constexpr __device__ int8_t kvalues_iq4nl[16] = {-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113};
|
||||
|
||||
__global__ void dequantize_iq4_xs_fp32_kernel(const int8_t* data, float* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id=global_idx; block_id<num_blocks; block_id+=blockDim.x * gridDim.x) {
|
||||
float* __restrict__ output_blk = (float*)(output + block_id * ele_per_blk);
|
||||
const float d = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size)));
|
||||
const uint16_t scales_h = *(reinterpret_cast<const uint16_t*>(data + block_id * blk_size + 2));
|
||||
const uint8_t* scales_l = (uint8_t*)(data + block_id * blk_size + 2 + 2);
|
||||
const uint8_t* qs = (uint8_t*)(data + block_id * blk_size + 2 + 2 + 4);
|
||||
|
||||
for (int ib = 0; ib < 8; ++ib) {
|
||||
const int ls = ((scales_l[ib / 2] >> 4 * (ib % 2)) & 0xf) | (((scales_h >> 2 * ib) & 3) << 4);
|
||||
const float dl = d * (ls - 32);
|
||||
for (int j = 0; j < 16; ++j) {
|
||||
output_blk[j + 0] = dl * kvalues_iq4nl[qs[j] & 0xf];
|
||||
output_blk[j + 16] = dl * kvalues_iq4nl[qs[j] >> 4];
|
||||
}
|
||||
output_blk += 32;
|
||||
qs += 16;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_iq4_xs_fp16_kernel(const int8_t* data, __half* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id=global_idx; block_id<num_blocks; block_id+=blockDim.x * gridDim.x) {
|
||||
__half* __restrict__ output_blk = (__half*)(output + block_id * ele_per_blk);
|
||||
const float d = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size)));
|
||||
const uint16_t scales_h = *(reinterpret_cast<const uint16_t*>(data + block_id * blk_size + 2));
|
||||
const uint8_t* scales_l = (uint8_t*)(data + block_id * blk_size + 2 + 2);
|
||||
const uint8_t* qs = (uint8_t*)(data + block_id * blk_size + 2 + 2 + 4);
|
||||
|
||||
for (int ib = 0; ib < 8; ++ib) {
|
||||
const int ls = ((scales_l[ib / 2] >> 4 * (ib % 2)) & 0xf) | (((scales_h >> 2 * ib) & 3) << 4);
|
||||
const float dl = d * (ls - 32);
|
||||
for (int j = 0; j < 16; ++j) {
|
||||
output_blk[j + 0] = __float2half(dl * kvalues_iq4nl[qs[j] & 0xf]);
|
||||
output_blk[j + 16] = __float2half(dl * kvalues_iq4nl[qs[j] >> 4]);
|
||||
}
|
||||
output_blk += 32;
|
||||
qs += 16;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void dequantize_iq4_xs_bf16_kernel(const int8_t* data, nv_bfloat16* output, const int blk_size, const int ele_per_blk, const int num_blocks) {
|
||||
long long global_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
for (long long block_id=global_idx; block_id<num_blocks; block_id+=blockDim.x * gridDim.x) {
|
||||
nv_bfloat16* __restrict__ output_blk = (nv_bfloat16*)(output + block_id * ele_per_blk);
|
||||
const float d = __half2float(*(reinterpret_cast<const half*>(data + block_id * blk_size)));
|
||||
const uint16_t scales_h = *(reinterpret_cast<const uint16_t*>(data + block_id * blk_size + 2));
|
||||
const uint8_t* scales_l = (uint8_t*)(data + block_id * blk_size + 2 + 2);
|
||||
const uint8_t* qs = (uint8_t*)(data + block_id * blk_size + 2 + 2 + 4);
|
||||
|
||||
for (int ib = 0; ib < 8; ++ib) {
|
||||
const int ls = ((scales_l[ib / 2] >> 4 * (ib % 2)) & 0xf) | (((scales_h >> 2 * ib) & 3) << 4);
|
||||
const float dl = d * (ls - 32);
|
||||
for (int j = 0; j < 16; ++j) {
|
||||
output_blk[j + 0] = __float2bfloat16(dl * kvalues_iq4nl[qs[j] & 0xf]);
|
||||
output_blk[j + 16] = __float2bfloat16(dl * kvalues_iq4nl[qs[j] >> 4]);
|
||||
}
|
||||
output_blk += 32;
|
||||
qs += 16;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor dequantize_q8_0(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) {
|
||||
int num_blocks = num_bytes / blk_size;
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device);
|
||||
|
||||
auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous);
|
||||
auto data_gpu = torch::empty({ num_bytes }, options);
|
||||
|
||||
cudaMemcpy(data_gpu.data_ptr<int8_t>(), data, num_bytes, cudaMemcpyHostToDevice);
|
||||
//data_gpu.copy_(data, false);
|
||||
|
||||
// Create output tensor
|
||||
auto output = torch::zeros({ num_blocks, 32 }, torch::dtype(target_dtype).device(device));
|
||||
|
||||
switch (target_dtype) {
|
||||
case torch::kFloat16:
|
||||
dequantize_q8_0_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
case torch::kBFloat16:
|
||||
dequantize_q8_0_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
case torch::kFloat32:
|
||||
dequantize_q8_0_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), output.data_ptr<float>(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
default:
|
||||
printf("target type not support\n");
|
||||
exit(0);
|
||||
}
|
||||
|
||||
cudaDeviceSynchronize();
|
||||
return output;
|
||||
}
|
||||
|
||||
|
||||
torch::Tensor dequantize_q6_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) {
|
||||
// data.numel%blk_size should be 0, else raise err
|
||||
int num_blocks = num_bytes / blk_size;
|
||||
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device);
|
||||
auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous);
|
||||
auto data_gpu = torch::empty({num_bytes}, options);
|
||||
|
||||
cudaMemcpy(data_gpu.data_ptr<int8_t>(), data, num_bytes, cudaMemcpyHostToDevice);
|
||||
//data_gpu.copy_(data, false);
|
||||
|
||||
// Create output tensor
|
||||
auto output = torch::zeros({num_blocks, 256}, torch::dtype(target_dtype).device(device));
|
||||
|
||||
switch (target_dtype) {
|
||||
case torch::kFloat16:
|
||||
dequantize_q6_k_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
case torch::kBFloat16:
|
||||
dequantize_q6_k_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
case torch::kFloat32:
|
||||
dequantize_q6_k_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), output.data_ptr<float>(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
default:
|
||||
printf("target type not support\n");
|
||||
exit(0);
|
||||
}
|
||||
cudaDeviceSynchronize();
|
||||
return output;
|
||||
}
|
||||
|
||||
torch::Tensor dequantize_q5_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) {
|
||||
int num_blocks = num_bytes / blk_size;
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device);
|
||||
|
||||
auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous);
|
||||
auto data_gpu = torch::empty({num_bytes}, options);
|
||||
|
||||
cudaMemcpy(data_gpu.data_ptr<int8_t>(), data, num_bytes, cudaMemcpyHostToDevice);
|
||||
//data_gpu.copy_(data, false);
|
||||
|
||||
// Create output tensor
|
||||
auto output = torch::zeros({num_blocks, 256}, torch::dtype(target_dtype).device(device));
|
||||
|
||||
switch (target_dtype) {
|
||||
case torch::kFloat16:
|
||||
dequantize_q5_k_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
case torch::kBFloat16:
|
||||
dequantize_q5_k_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
case torch::kFloat32:
|
||||
dequantize_q5_k_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), output.data_ptr<float>(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
default:
|
||||
printf("target type not support\n");
|
||||
exit(0);
|
||||
}
|
||||
cudaDeviceSynchronize();
|
||||
return output;
|
||||
}
|
||||
|
||||
torch::Tensor dequantize_q4_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) {
|
||||
// data.numel%blk_size should be 0, else raise err
|
||||
int num_blocks = num_bytes / blk_size;
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device);
|
||||
|
||||
auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous);
|
||||
auto data_gpu = torch::empty({num_bytes}, options);
|
||||
|
||||
cudaMemcpy(data_gpu.data_ptr<int8_t>(), data, num_bytes, cudaMemcpyHostToDevice);
|
||||
//data_gpu.copy_(data, false);
|
||||
|
||||
// Create output tensor
|
||||
auto output = torch::zeros({num_blocks, 256}, torch::dtype(target_dtype).device(device));
|
||||
|
||||
switch (target_dtype) {
|
||||
case torch::kFloat16:
|
||||
dequantize_q4_k_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
case torch::kBFloat16:
|
||||
dequantize_q4_k_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
case torch::kFloat32:
|
||||
dequantize_q4_k_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), output.data_ptr<float>(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
default:
|
||||
printf("target type not support\n");
|
||||
exit(0);
|
||||
}
|
||||
cudaDeviceSynchronize();
|
||||
return output;
|
||||
}
|
||||
|
||||
torch::Tensor dequantize_q3_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) {
|
||||
int num_blocks = num_bytes / blk_size;
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device);
|
||||
|
||||
auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous);
|
||||
auto data_gpu = torch::empty({num_bytes}, options);
|
||||
|
||||
cudaMemcpy(data_gpu.data_ptr<int8_t>(), data, num_bytes, cudaMemcpyHostToDevice);
|
||||
//data_gpu.copy_(data, false);
|
||||
|
||||
// Create output tensor
|
||||
auto output = torch::zeros({num_blocks, 256}, torch::dtype(target_dtype).device(device));
|
||||
|
||||
switch (target_dtype) {
|
||||
case torch::kFloat16:
|
||||
dequantize_q3_k_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
case torch::kBFloat16:
|
||||
dequantize_q3_k_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
case torch::kFloat32:
|
||||
dequantize_q3_k_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), output.data_ptr<float>(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
default:
|
||||
printf("target type not support\n");
|
||||
exit(0);
|
||||
}
|
||||
cudaDeviceSynchronize();
|
||||
return output;
|
||||
}
|
||||
|
||||
torch::Tensor dequantize_q2_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) {
|
||||
int num_blocks = num_bytes / blk_size;
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device);
|
||||
|
||||
auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous);
|
||||
auto data_gpu = torch::empty({num_bytes}, options);
|
||||
|
||||
cudaMemcpy(data_gpu.data_ptr<int8_t>(), data, num_bytes, cudaMemcpyHostToDevice);
|
||||
//data_gpu.copy_(data, false);
|
||||
|
||||
// Create output tensor
|
||||
auto output = torch::zeros({num_blocks, 256}, torch::dtype(target_dtype).device(device));
|
||||
|
||||
switch (target_dtype) {
|
||||
case torch::kFloat16:
|
||||
dequantize_q2_k_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
case torch::kBFloat16:
|
||||
dequantize_q2_k_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
case torch::kFloat32:
|
||||
dequantize_q2_k_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), output.data_ptr<float>(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
default:
|
||||
printf("target type not support\n");
|
||||
exit(0);
|
||||
}
|
||||
cudaDeviceSynchronize();
|
||||
return output;
|
||||
}
|
||||
|
||||
torch::Tensor dequantize_iq4_xs(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype) {
|
||||
int num_blocks = num_bytes / blk_size;
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device);
|
||||
|
||||
auto options = torch::TensorOptions().dtype(torch::kInt8).device(device).memory_format(torch::MemoryFormat::Contiguous);
|
||||
auto data_gpu = torch::empty({num_bytes}, options);
|
||||
|
||||
cudaMemcpy(data_gpu.data_ptr<int8_t>(), data, num_bytes, cudaMemcpyHostToDevice);
|
||||
//data_gpu.copy_(data, false);
|
||||
|
||||
// Create output tensor
|
||||
auto output = torch::zeros({num_blocks, 256}, torch::dtype(target_dtype).device(device));
|
||||
|
||||
switch (target_dtype) {
|
||||
case torch::kFloat16:
|
||||
dequantize_iq4_xs_fp16_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), (__half*)output.data_ptr(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
case torch::kBFloat16:
|
||||
dequantize_iq4_xs_bf16_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), (nv_bfloat16*)output.data_ptr(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
case torch::kFloat32:
|
||||
dequantize_iq4_xs_fp32_kernel<<<512, 256>>>(data_gpu.data_ptr<int8_t>(), output.data_ptr<float>(), blk_size, ele_per_blk, num_blocks);
|
||||
break;
|
||||
default:
|
||||
printf("target type not support\n");
|
||||
exit(0);
|
||||
}
|
||||
cudaDeviceSynchronize();
|
||||
return output;
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : Azure-Tang
|
||||
* @Date : 2024-07-22 09:27:55
|
||||
* @Version : 1.0.0
|
||||
* @LastEditors : kkk1nak0
|
||||
* @LastEditTime : 2024-08-12 03:48:46
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
#pragma once
|
||||
|
||||
#include <torch/library.h>
|
||||
#include <torch/extension.h>
|
||||
#include <torch/torch.h>
|
||||
|
||||
torch::Tensor dequantize_q8_0(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype);
|
||||
torch::Tensor dequantize_q6_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype);
|
||||
torch::Tensor dequantize_q5_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype);
|
||||
torch::Tensor dequantize_q4_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype);
|
||||
torch::Tensor dequantize_q3_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype);
|
||||
torch::Tensor dequantize_q2_k(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype);
|
||||
torch::Tensor dequantize_iq4_xs(const int8_t* data, const int num_bytes, const int blk_size, const int ele_per_blk, const torch::Device device, const torch::Dtype target_dtype);
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,80 @@
|
||||
// Adapted from
|
||||
// https://github.com/vllm-project/vllm/tree/main/csrc/quantization/gptq_marlin
|
||||
// Copyrigth 2024 The vLLM team.
|
||||
// Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
#pragma once
|
||||
|
||||
#include <torch/all.h>
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <cuda.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <iostream>
|
||||
|
||||
namespace gptq_marlin {
|
||||
|
||||
// 8 warps are a good choice since every SM has 4 schedulers and having more
|
||||
// than 1 warp per schedule allows some more latency hiding. At the same time,
|
||||
// we want relatively few warps to have many registers per warp and small tiles.
|
||||
static constexpr int default_threads = 256;
|
||||
|
||||
static constexpr int pipe_stages =
|
||||
4; // 4 pipeline stages fit into shared memory
|
||||
|
||||
static constexpr int min_thread_n = 64;
|
||||
static constexpr int min_thread_k = 64;
|
||||
|
||||
static constexpr int tile_size = 16;
|
||||
static constexpr int max_par = 16;
|
||||
|
||||
template <typename T, int n>
|
||||
struct Vec {
|
||||
T elems[n];
|
||||
__device__ T& operator[](int i) { return elems[i]; }
|
||||
};
|
||||
|
||||
using I4 = Vec<int, 4>;
|
||||
|
||||
constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; }
|
||||
|
||||
#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800) || defined (__HIP_PLATFORM_AMD__)
|
||||
// No support for async
|
||||
#else
|
||||
|
||||
__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
|
||||
bool pred = true) {
|
||||
const int BYTES = 16;
|
||||
uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
|
||||
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)pred),
|
||||
"r"(smem), "l"(glob_ptr), "n"(BYTES));
|
||||
}
|
||||
|
||||
__device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) {
|
||||
const int BYTES = 16;
|
||||
uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" cp.async.cg.shared.global [%0], [%1], %2;\n"
|
||||
"}\n" ::"r"(smem),
|
||||
"l"(glob_ptr), "n"(BYTES));
|
||||
}
|
||||
|
||||
__device__ inline void cp_async_fence() {
|
||||
asm volatile("cp.async.commit_group;\n" ::);
|
||||
}
|
||||
|
||||
template <int n>
|
||||
__device__ inline void cp_async_wait() {
|
||||
asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
} // namespace gptq_marlin
|
||||
@@ -0,0 +1,85 @@
|
||||
// Adapted from
|
||||
// https://github.com/vllm-project/vllm/tree/main/csrc/quantization/gptq_marlin
|
||||
// Copyrigth 2024 The vLLM team.
|
||||
// Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
#ifndef _data_types_cuh
|
||||
#define _data_types_cuh
|
||||
#include "gptq_marlin.cuh"
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_bf16.h>
|
||||
|
||||
#ifdef __HIP_PLATFORM_AMD__
|
||||
typedef __hip_bfloat16 nv_bfloat16;
|
||||
typedef __hip_bfloat162 nv_bfloat162;
|
||||
#endif
|
||||
|
||||
namespace gptq_marlin {
|
||||
|
||||
template <typename scalar_t>
|
||||
class ScalarType {};
|
||||
|
||||
template <>
|
||||
class ScalarType<half> {
|
||||
public:
|
||||
using scalar_t = half;
|
||||
using scalar_t2 = half2;
|
||||
|
||||
// Matrix fragments for tensor core instructions; their precise layout is
|
||||
// documented here:
|
||||
// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type
|
||||
using FragA = Vec<half2, 4>;
|
||||
using FragB = Vec<half2, 2>;
|
||||
using FragC = Vec<float, 4>;
|
||||
using FragS = Vec<half2, 1>;
|
||||
|
||||
static __device__ float inline num2float(const half x) {
|
||||
return __half2float(x);
|
||||
}
|
||||
|
||||
static __device__ half2 inline num2num2(const half x) {
|
||||
return __half2half2(x);
|
||||
}
|
||||
|
||||
static __device__ half2 inline nums2num2(const half x1, const half x2) {
|
||||
return __halves2half2(x1, x2);
|
||||
}
|
||||
|
||||
static __host__ __device__ half inline float2num(const float x) {
|
||||
return __float2half(x);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
class ScalarType<nv_bfloat16> {
|
||||
public:
|
||||
using scalar_t = nv_bfloat16;
|
||||
using scalar_t2 = nv_bfloat162;
|
||||
|
||||
using FragA = Vec<nv_bfloat162, 4>;
|
||||
using FragB = Vec<nv_bfloat162, 2>;
|
||||
using FragC = Vec<float, 4>;
|
||||
using FragS = Vec<nv_bfloat162, 1>;
|
||||
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
static __device__ float inline num2float(const nv_bfloat16 x) {
|
||||
return __bfloat162float(x);
|
||||
}
|
||||
|
||||
static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) {
|
||||
return __bfloat162bfloat162(x);
|
||||
}
|
||||
|
||||
static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1,
|
||||
const nv_bfloat16 x2) {
|
||||
return __halves2bfloat162(x1, x2);
|
||||
}
|
||||
|
||||
static __host__ __device__ nv_bfloat16 inline float2num(const float x) {
|
||||
return __float2bfloat16(x);
|
||||
}
|
||||
#endif
|
||||
};
|
||||
|
||||
} // namespace gptq_marlin
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,24 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : Azure
|
||||
* @Date : 2024-07-22 09:27:55
|
||||
* @Version : 1.0.0
|
||||
* @LastEditors : Azure
|
||||
* @LastEditTime : 2024-07-26 08:35:00
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
#pragma once
|
||||
|
||||
#include <torch/library.h>
|
||||
#include <torch/extension.h>
|
||||
#include <torch/torch.h>
|
||||
|
||||
torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
|
||||
torch::Tensor& b_scales, torch::Tensor& g_idx,
|
||||
torch::Tensor& perm, torch::Tensor& workspace,
|
||||
int64_t num_bits, int64_t size_m, int64_t size_n,
|
||||
int64_t size_k, bool is_k_full);
|
||||
|
||||
// torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm,
|
||||
// int64_t size_k, int64_t size_n,
|
||||
// int64_t num_bits);
|
||||
@@ -0,0 +1,26 @@
|
||||
|
||||
from setuptools import setup, Extension
|
||||
from torch.utils import cpp_extension
|
||||
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
||||
setup(
|
||||
name='KTransformersOps',
|
||||
ext_modules=[
|
||||
CUDAExtension(
|
||||
'KTransformersOps', [
|
||||
'custom_gguf/dequant.cu',
|
||||
'binding.cpp',
|
||||
'gptq_marlin/gptq_marlin.cu',
|
||||
# 'gptq_marlin_repack.cu',
|
||||
],
|
||||
extra_compile_args={
|
||||
'cxx': ['-O3'],
|
||||
'nvcc': [
|
||||
'-O3',
|
||||
'--use_fast_math',
|
||||
'-Xcompiler', '-fPIC',
|
||||
]
|
||||
},
|
||||
)
|
||||
],
|
||||
cmdclass={'build_ext': BuildExtension}
|
||||
)
|
||||
@@ -0,0 +1,16 @@
|
||||
import os
|
||||
import sys
|
||||
sys.path.insert(0,"/home/zbx/ktransformers")
|
||||
from ktransformers.util.custom_loader import GGUFLoader
|
||||
import torch
|
||||
|
||||
gguf_loader_1 = GGUFLoader("/mnt/data/model/DeepseekV3-q4km-gguf")
|
||||
gguf_loader_2 = GGUFLoader("/mnt/data/chenht/model/gguf_for_ktransformers/DeepSeek-V3-bf16/")
|
||||
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
|
||||
tensor_1 = gguf_loader_1.load_gguf_tensor("blk.0.attn_kv_a_mqa.weight", "cuda")
|
||||
tensor_2 = gguf_loader_2.load_gguf_tensor("blk.0.attn_kv_a_mqa.weight", "cuda")
|
||||
|
||||
print(tensor_1[0, -64:])
|
||||
print(tensor_2[0, -64:])
|
||||
@@ -0,0 +1,142 @@
|
||||
#!/usr/bin/env python
|
||||
# coding=utf-8
|
||||
"""
|
||||
Description :
|
||||
Author : Jianwei Dong
|
||||
Date : 2024-08-28 10:32:05
|
||||
Version : 1.0.0
|
||||
LastEditors : chenht2022
|
||||
LastEditTime : 2024-08-28 10:32:05
|
||||
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
"""
|
||||
import os, sys
|
||||
import time
|
||||
|
||||
sys.path.append(os.path.dirname(__file__) + "/../build")
|
||||
import cpuinfer_ext
|
||||
from flash_attn import flash_attn_with_kvcache
|
||||
import torch
|
||||
|
||||
layer_num = 10
|
||||
kv_head_num = 8
|
||||
q_head_num = 32
|
||||
head_dim = 128
|
||||
block_len = 128
|
||||
anchor_num = 1
|
||||
cache_seqlen = 8192
|
||||
cache_seqlens = torch.tensor([cache_seqlen], dtype=torch.int32, device="cpu")
|
||||
seqlens_zero = torch.zeros((1,), dtype=torch.int32, device="cpu")
|
||||
anchor_type = cpuinfer_ext.kvcache.AnchorType.DYNAMIC
|
||||
kv_type = cpuinfer_ext.kvcache.ggml_type.FP16
|
||||
retrieval_type = cpuinfer_ext.kvcache.RetrievalType.LAYER
|
||||
layer_step: int = 1
|
||||
token_step: int = 1
|
||||
layer_offset: int = 0
|
||||
max_thread_num: int = 2
|
||||
max_batch_size: int = 1
|
||||
max_block_num: int = 512
|
||||
CPUInfer = cpuinfer_ext.CPUInfer(max_thread_num)
|
||||
validation_iter = 100
|
||||
|
||||
with torch.inference_mode(mode=True):
|
||||
config = cpuinfer_ext.kvcache.KVCacheConfig(
|
||||
layer_num,
|
||||
kv_head_num,
|
||||
q_head_num,
|
||||
head_dim,
|
||||
block_len,
|
||||
anchor_num,
|
||||
anchor_type,
|
||||
kv_type,
|
||||
retrieval_type,
|
||||
layer_step,
|
||||
token_step,
|
||||
layer_offset,
|
||||
max_block_num,
|
||||
max_batch_size,
|
||||
max_thread_num,
|
||||
)
|
||||
local_kvcache = cpuinfer_ext.kvcache.KVCache(config)
|
||||
|
||||
kvcaches = []
|
||||
block_table = (
|
||||
torch.arange(max_block_num, dtype=torch.int32, device="cpu")
|
||||
.contiguous()
|
||||
.view(1, -1)
|
||||
)
|
||||
|
||||
for layer_idx in range(layer_num):
|
||||
k_cache = torch.randn(
|
||||
(1, cache_seqlen, kv_head_num, head_dim), dtype=torch.float16, device="cpu"
|
||||
).contiguous()
|
||||
v_cache = torch.randn(
|
||||
(1, cache_seqlen, kv_head_num, head_dim), dtype=torch.float16, device="cpu"
|
||||
).contiguous()
|
||||
|
||||
CPUInfer.submit(
|
||||
local_kvcache.update_kvcache_fp16(
|
||||
k_cache.data_ptr(),
|
||||
v_cache.data_ptr(),
|
||||
layer_idx,
|
||||
block_table.data_ptr(),
|
||||
1,
|
||||
max_block_num,
|
||||
seqlens_zero.data_ptr(),
|
||||
cache_seqlen,
|
||||
)
|
||||
)
|
||||
CPUInfer.sync()
|
||||
|
||||
kvcaches.append((k_cache.to("cuda"), v_cache.to("cuda")))
|
||||
|
||||
# validation
|
||||
for i in range(validation_iter):
|
||||
|
||||
k_cache = kvcaches[i % layer_num][0]
|
||||
v_cache = kvcaches[i % layer_num][1]
|
||||
input = torch.randn(
|
||||
(1, 1, q_head_num, head_dim), dtype=torch.float16, device="cpu"
|
||||
).contiguous()
|
||||
output = torch.empty(
|
||||
(1, 1, q_head_num, head_dim), dtype=torch.float16, device="cpu"
|
||||
).contiguous()
|
||||
|
||||
# attn_lse: (bsz, q_len, q_head_num)
|
||||
attn_lse = torch.empty(
|
||||
(1, 1, q_head_num), dtype=torch.float32, device="cpu"
|
||||
).contiguous()
|
||||
input = input / 100
|
||||
|
||||
CPUInfer.submit(
|
||||
local_kvcache.attn(
|
||||
input.data_ptr(),
|
||||
output.data_ptr(),
|
||||
attn_lse.data_ptr(),
|
||||
i % layer_num,
|
||||
0,
|
||||
1,
|
||||
1,
|
||||
max_block_num,
|
||||
block_table.data_ptr(),
|
||||
cache_seqlens.data_ptr(),
|
||||
-1,
|
||||
-1,
|
||||
-1,
|
||||
)
|
||||
)
|
||||
CPUInfer.sync()
|
||||
# print("cpuinfer output", output)
|
||||
|
||||
t_output = flash_attn_with_kvcache(
|
||||
q=input.to("cuda"),
|
||||
k_cache=k_cache,
|
||||
v_cache=v_cache,
|
||||
cache_seqlens=cache_seqlens.to("cuda"),
|
||||
)
|
||||
# print("torch output", t_output)
|
||||
|
||||
diff = torch.mean(torch.abs(output.to("cuda") - t_output)) / torch.mean(
|
||||
torch.abs(t_output)
|
||||
)
|
||||
print("diff = ", diff)
|
||||
assert diff < 0.001
|
||||
@@ -0,0 +1,62 @@
|
||||
#!/usr/bin/env python
|
||||
# coding=utf-8
|
||||
'''
|
||||
Description :
|
||||
Author : chenht2022
|
||||
Date : 2024-07-25 10:32:05
|
||||
Version : 1.0.0
|
||||
LastEditors : chenht2022
|
||||
LastEditTime : 2024-08-06 10:36:59
|
||||
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
'''
|
||||
import os, sys
|
||||
import time
|
||||
sys.path.append(os.path.dirname(__file__) + '/../build')
|
||||
import cpuinfer_ext
|
||||
import torch
|
||||
|
||||
input_size = 16384
|
||||
output_size = 5120
|
||||
stride = 32
|
||||
group_max_len = 1024
|
||||
proj_type = 1 # ggml_type::GGML_TYPE_F16
|
||||
hidden_type = 1 # ggml_type::GGML_TYPE_F16
|
||||
qlen = 30
|
||||
layer_num = 10
|
||||
CPUInfer = cpuinfer_ext.CPUInfer(48)
|
||||
validation_iter = 100
|
||||
|
||||
with torch.inference_mode(mode=True):
|
||||
linears = []
|
||||
projs = []
|
||||
for _ in range(layer_num):
|
||||
proj = torch.randn((output_size, input_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous()
|
||||
config = cpuinfer_ext.linear.LinearConfig(input_size, output_size, stride, group_max_len, proj.data_ptr(), proj_type, hidden_type)
|
||||
linear = cpuinfer_ext.linear.Linear(config)
|
||||
projs.append(proj)
|
||||
linears.append(linear)
|
||||
|
||||
# validation
|
||||
for i in range(validation_iter):
|
||||
linear = linears[i % layer_num]
|
||||
input = torch.randn((qlen, input_size), dtype=torch.float16).contiguous()
|
||||
output = torch.empty((qlen, output_size), dtype=torch.float16).contiguous()
|
||||
input = input / 100
|
||||
|
||||
CPUInfer.submit(
|
||||
linear.forward(
|
||||
qlen,
|
||||
input.data_ptr(),
|
||||
output.data_ptr()
|
||||
)
|
||||
)
|
||||
CPUInfer.sync()
|
||||
# print('cpuinfer output', output)
|
||||
|
||||
proj = projs[i%layer_num]
|
||||
t_output = torch.mm(input, proj.t())
|
||||
# print('torch output', t_output)
|
||||
|
||||
diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output))
|
||||
print('diff = ', diff)
|
||||
assert(diff < 0.001)
|
||||
@@ -0,0 +1,82 @@
|
||||
#!/usr/bin/env python
|
||||
# coding=utf-8
|
||||
'''
|
||||
Description :
|
||||
Author : chenht2022
|
||||
Date : 2024-07-25 10:32:05
|
||||
Version : 1.0.0
|
||||
LastEditors : chenht2022
|
||||
LastEditTime : 2024-08-06 10:37:28
|
||||
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
'''
|
||||
import os, sys
|
||||
import time
|
||||
sys.path.append(os.path.dirname(__file__) + '/../build')
|
||||
import cpuinfer_ext
|
||||
import torch
|
||||
|
||||
hidden_size = 5120
|
||||
intermediate_size = 3072
|
||||
stride = 32
|
||||
group_max_len = 1024
|
||||
gate_type = 1 # ggml_type::GGML_TYPE_F16
|
||||
up_type = 1 # ggml_type::GGML_TYPE_F16
|
||||
down_type = 1 # ggml_type::GGML_TYPE_F16
|
||||
hidden_type = 1 # ggml_type::GGML_TYPE_F16
|
||||
qlen = 30
|
||||
layer_num = 10
|
||||
CPUInfer = cpuinfer_ext.CPUInfer(48)
|
||||
validation_iter = 100
|
||||
|
||||
def act_fn(x):
|
||||
return x / (1.0 + torch.exp(-x))
|
||||
|
||||
def mlp_torch(input, gate_proj, up_proj, down_proj):
|
||||
gate_buf = torch.mm(input, gate_proj.t())
|
||||
up_buf = torch.mm(input, up_proj.t())
|
||||
intermediate = act_fn(gate_buf) * up_buf
|
||||
ret = torch.mm(intermediate, down_proj.t())
|
||||
return ret
|
||||
|
||||
with torch.inference_mode(mode=True):
|
||||
mlps = []
|
||||
gate_projs = []
|
||||
up_projs = []
|
||||
down_projs = []
|
||||
for _ in range(layer_num):
|
||||
gate_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous()
|
||||
up_proj = torch.randn((intermediate_size, hidden_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous()
|
||||
down_proj = torch.randn((hidden_size, intermediate_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous()
|
||||
config = cpuinfer_ext.mlp.MLPConfig(hidden_size, intermediate_size, stride, group_max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), gate_type, up_type, down_type, hidden_type)
|
||||
mlp = cpuinfer_ext.mlp.MLP(config)
|
||||
gate_projs.append(gate_proj)
|
||||
up_projs.append(up_proj)
|
||||
down_projs.append(down_proj)
|
||||
mlps.append(mlp)
|
||||
|
||||
# validation
|
||||
for i in range(validation_iter):
|
||||
mlp = mlps[i % layer_num]
|
||||
input = torch.randn((qlen, hidden_size), dtype=torch.float16).contiguous()
|
||||
output = torch.empty((qlen, hidden_size), dtype=torch.float16).contiguous()
|
||||
input = input / 100
|
||||
|
||||
CPUInfer.submit(
|
||||
mlp.forward(
|
||||
qlen,
|
||||
input.data_ptr(),
|
||||
output.data_ptr()
|
||||
)
|
||||
)
|
||||
CPUInfer.sync()
|
||||
# print('cpuinfer output', output)
|
||||
|
||||
gate_proj = gate_projs[i%layer_num]
|
||||
up_proj = up_projs[i%layer_num]
|
||||
down_proj = down_projs[i%layer_num]
|
||||
t_output = mlp_torch(input, gate_proj, up_proj, down_proj)
|
||||
# print('torch output', t_output)
|
||||
|
||||
diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output))
|
||||
print('diff = ', diff)
|
||||
assert(diff < 0.001)
|
||||
@@ -0,0 +1,121 @@
|
||||
#!/usr/bin/env python
|
||||
# coding=utf-8
|
||||
'''
|
||||
Description :
|
||||
Author : chenht2022
|
||||
Date : 2024-07-25 10:32:05
|
||||
Version : 1.0.0
|
||||
LastEditors : chenht2022
|
||||
LastEditTime : 2024-08-06 10:38:05
|
||||
Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
'''
|
||||
import os, sys
|
||||
import time
|
||||
sys.path.append(os.path.dirname(__file__) + '/../build')
|
||||
import cpuinfer_ext
|
||||
import torch
|
||||
|
||||
expert_num = 160
|
||||
hidden_size = 5120
|
||||
intermediate_size = 1536
|
||||
stride = 32
|
||||
group_min_len = 10
|
||||
group_max_len = 1024
|
||||
gate_type = 1 # ggml_type::GGML_TYPE_F16
|
||||
up_type = 1 # ggml_type::GGML_TYPE_F16
|
||||
down_type = 1 # ggml_type::GGML_TYPE_F16
|
||||
hidden_type = 1 # ggml_type::GGML_TYPE_F16
|
||||
n_routed_experts = 6
|
||||
qlen = 30
|
||||
layer_num = 10
|
||||
CPUInfer = cpuinfer_ext.CPUInfer(48)
|
||||
validation_iter = 100
|
||||
|
||||
def act_fn(x):
|
||||
return x / (1.0 + torch.exp(-x))
|
||||
|
||||
def mlp_torch(input, gate_proj, up_proj, down_proj):
|
||||
gate_buf = torch.mm(input, gate_proj.t())
|
||||
up_buf = torch.mm(input, up_proj.t())
|
||||
intermediate = act_fn(gate_buf) * up_buf
|
||||
ret = torch.mm(intermediate, down_proj.t())
|
||||
return ret
|
||||
|
||||
def moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj):
|
||||
cnts = expert_ids.new_zeros((expert_ids.shape[0], expert_num))
|
||||
cnts.scatter_(1, expert_ids, 1)
|
||||
tokens_per_expert = cnts.sum(dim=0)
|
||||
idxs = expert_ids.view(-1).argsort()
|
||||
sorted_tokens = input[idxs // expert_ids.shape[1]]
|
||||
|
||||
outputs = []
|
||||
start_idx = 0
|
||||
for i, num_tokens in enumerate(tokens_per_expert):
|
||||
end_idx = start_idx + num_tokens
|
||||
if num_tokens == 0:
|
||||
continue
|
||||
tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
|
||||
expert_out = mlp_torch(tokens_for_this_expert, gate_proj[i], up_proj[i], down_proj[i])
|
||||
outputs.append(expert_out)
|
||||
start_idx = end_idx
|
||||
|
||||
outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
|
||||
|
||||
new_x = torch.empty_like(outs)
|
||||
new_x[idxs] = outs
|
||||
t_output = (
|
||||
new_x.view(*expert_ids.shape, -1)
|
||||
.type(weights.dtype)
|
||||
.mul_(weights.unsqueeze(dim=-1))
|
||||
.sum(dim=1)
|
||||
.type(new_x.dtype)
|
||||
)
|
||||
return t_output
|
||||
|
||||
with torch.inference_mode(mode=True):
|
||||
moes = []
|
||||
gate_projs = []
|
||||
up_projs = []
|
||||
down_projs = []
|
||||
for _ in range(layer_num):
|
||||
gate_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous()
|
||||
up_proj = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous()
|
||||
down_proj = torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous()
|
||||
config = cpuinfer_ext.moe.MOEConfig(expert_num, n_routed_experts, hidden_size, intermediate_size, stride, group_min_len, group_max_len, gate_proj.data_ptr(), up_proj.data_ptr(), down_proj.data_ptr(), gate_type, up_type, down_type, hidden_type)
|
||||
moe = cpuinfer_ext.moe.MOE(config)
|
||||
gate_projs.append(gate_proj)
|
||||
up_projs.append(up_proj)
|
||||
down_projs.append(down_proj)
|
||||
moes.append(moe)
|
||||
|
||||
# validation
|
||||
for i in range(validation_iter):
|
||||
expert_ids = torch.stack([torch.randperm(expert_num)[:n_routed_experts] for _ in range(qlen)]).contiguous()
|
||||
weights = torch.rand((qlen, n_routed_experts), dtype=torch.float32).contiguous()
|
||||
input = torch.randn((qlen, hidden_size), dtype=torch.float16).contiguous()
|
||||
output = torch.empty((qlen, hidden_size), dtype=torch.float16).contiguous()
|
||||
input = input / 100
|
||||
|
||||
moe = moes[i % layer_num]
|
||||
CPUInfer.submit(
|
||||
moe.forward(
|
||||
qlen,
|
||||
n_routed_experts,
|
||||
expert_ids.data_ptr(),
|
||||
weights.data_ptr(),
|
||||
input.data_ptr(),
|
||||
output.data_ptr()
|
||||
)
|
||||
)
|
||||
CPUInfer.sync()
|
||||
# print('cpuinfer output', output)
|
||||
|
||||
gate_proj = gate_projs[i%layer_num]
|
||||
up_proj = up_projs[i%layer_num]
|
||||
down_proj = down_projs[i%layer_num]
|
||||
t_output = moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj)
|
||||
# print('torch output', t_output)
|
||||
|
||||
diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output))
|
||||
print('diff = ', diff)
|
||||
assert(diff < 0.001)
|
||||
@@ -0,0 +1,790 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : chenht2022, Jianwei Dong
|
||||
* @Date : 2024-07-22 02:03:22
|
||||
* @Version : 1.0.0
|
||||
* @LastEditors : Jianwei Dong
|
||||
* @LastEditTime : 2024-08-26 22:47:06
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
// Python bindings
|
||||
#include "cpu_backend/cpuinfer.h"
|
||||
#if !defined(KTRANSFORMERS_USE_ROCM) && !defined(KTRANSFORMERS_USE_XPU) && !defined(KTRANSFORMERS_USE_NPU)
|
||||
#include "device_launch_parameters.h"
|
||||
#endif
|
||||
#include "llamafile/flags.h"
|
||||
#include "operators/kvcache/kvcache.h"
|
||||
#include "operators/llamafile/linear.h"
|
||||
#include "operators/llamafile/mlp.h"
|
||||
#include "operators/llamafile/moe.h"
|
||||
|
||||
#if defined(__x86_64__) && defined(__HAS_AVX512F__) && defined(__HAS_AMX__)
|
||||
#include "operators/amx/moe.hpp"
|
||||
#endif
|
||||
|
||||
#include "pybind11/functional.h"
|
||||
#include "pybind11/operators.h"
|
||||
#include "pybind11/pybind11.h"
|
||||
#include "pybind11/stl.h"
|
||||
#include <cstdint>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
|
||||
namespace py = pybind11;
|
||||
using namespace pybind11::literals;
|
||||
|
||||
// Binding functions for the KVCache class
|
||||
class KVCacheBindings {
|
||||
public:
|
||||
class AttnBindings {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
KVCache *kv_cache;
|
||||
const ggml_fp16_t *q_in;
|
||||
ggml_fp16_t *output;
|
||||
float *attn_lse;
|
||||
int layer_idx;
|
||||
int generate_token_idx;
|
||||
int q_len;
|
||||
int batch_size;
|
||||
int max_block_num;
|
||||
int *block_table;
|
||||
int *cache_seqlens;
|
||||
int pick_block_num;
|
||||
int init_block_num;
|
||||
int local_block_num;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(
|
||||
&KVCache::attn, args_->kv_cache, args_->q_in, args_->output,
|
||||
args_->attn_lse, args_->layer_idx, args_->generate_token_idx,
|
||||
args_->q_len, args_->batch_size, args_->max_block_num,
|
||||
args_->block_table, args_->cache_seqlens, args_->pick_block_num,
|
||||
args_->init_block_num, args_->local_block_num);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t>
|
||||
cpuinfer_interface(KVCache &kv_cache, intptr_t q_in, intptr_t output,
|
||||
intptr_t attn_lse, int layer_idx,
|
||||
int generate_token_idx, int q_len, int batch_size,
|
||||
int max_block_num, intptr_t block_table,
|
||||
intptr_t cache_seqlens, int pick_block_num,
|
||||
int init_block_num, int local_block_num) {
|
||||
Args *args = new Args{nullptr,
|
||||
&kv_cache,
|
||||
(const ggml_fp16_t *)q_in,
|
||||
(ggml_fp16_t *)output,
|
||||
(float *)attn_lse,
|
||||
layer_idx,
|
||||
generate_token_idx,
|
||||
q_len,
|
||||
batch_size,
|
||||
max_block_num,
|
||||
(int *)block_table,
|
||||
(int *)cache_seqlens,
|
||||
pick_block_num,
|
||||
init_block_num,
|
||||
local_block_num};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
|
||||
class GetAllKVCacheOneLayerBindings {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
KVCache *kv_cache;
|
||||
int layer_id;
|
||||
ggml_fp16_t *k_in;
|
||||
ggml_fp16_t *v_in;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(&KVCache::get_all_kvcache_one_layer,
|
||||
args_->kv_cache, args_->layer_id,
|
||||
args_->k_in, args_->v_in);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t>
|
||||
cpuinfer_interface(KVCache &kv_cache, intptr_t k_in, intptr_t v_in,
|
||||
int layer_id) {
|
||||
Args *args = new Args{nullptr, &kv_cache, layer_id,
|
||||
(ggml_fp16_t *)k_in, (ggml_fp16_t *)v_in};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
|
||||
class GetAndUpdateKVCacheFp16Bindings {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
KVCache *kv_cache;
|
||||
ggml_fp16_t *k_in;
|
||||
ggml_fp16_t *v_in;
|
||||
int layer_id;
|
||||
int *block_table;
|
||||
int batch_size;
|
||||
int max_block_num;
|
||||
int *cache_seqlens;
|
||||
int q_len;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(&KVCache::get_and_update_kvcache_fp16,
|
||||
args_->kv_cache, args_->k_in, args_->v_in,
|
||||
args_->layer_id, args_->block_table,
|
||||
args_->batch_size, args_->max_block_num,
|
||||
args_->cache_seqlens, args_->q_len);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t>
|
||||
cpuinfer_interface(KVCache &kv_cache, intptr_t k_in, intptr_t v_in,
|
||||
int layer_id, intptr_t block_table, int batch_size,
|
||||
int max_block_num, intptr_t cache_seqlens,
|
||||
int q_len) {
|
||||
Args *args = new Args{nullptr,
|
||||
&kv_cache,
|
||||
(ggml_fp16_t *)k_in,
|
||||
(ggml_fp16_t *)v_in,
|
||||
layer_id,
|
||||
(int *)block_table,
|
||||
batch_size,
|
||||
max_block_num,
|
||||
(int *)cache_seqlens,
|
||||
q_len};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
class GetKVCacheFp16Bindings {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
KVCache *kv_cache;
|
||||
ggml_fp16_t *k_in;
|
||||
ggml_fp16_t *v_in;
|
||||
int layer_id;
|
||||
int *block_table;
|
||||
int batch_size;
|
||||
int max_block_num;
|
||||
int *cache_seqlens;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(
|
||||
&KVCache::get_kvcache_fp16, args_->kv_cache, args_->k_in,
|
||||
args_->v_in, args_->layer_id, args_->block_table,
|
||||
args_->batch_size, args_->max_block_num, args_->cache_seqlens);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t>
|
||||
cpuinfer_interface(KVCache &kv_cache, intptr_t k_in, intptr_t v_in,
|
||||
int layer_id, intptr_t block_table, int batch_size,
|
||||
int max_block_num, intptr_t cache_seqlens) {
|
||||
Args *args = new Args{nullptr,
|
||||
&kv_cache,
|
||||
(ggml_fp16_t *)k_in,
|
||||
(ggml_fp16_t *)v_in,
|
||||
layer_id,
|
||||
(int *)block_table,
|
||||
batch_size,
|
||||
max_block_num,
|
||||
(int *)cache_seqlens};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
|
||||
class UpdateKVCacheFp16Bindings {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
KVCache *kv_cache;
|
||||
ggml_fp16_t *k_in;
|
||||
ggml_fp16_t *v_in;
|
||||
int layer_id;
|
||||
int *block_table;
|
||||
int batch_size;
|
||||
int max_block_num;
|
||||
int *cache_seqlens;
|
||||
int q_len;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(&KVCache::update_kvcache_fp16,
|
||||
args_->kv_cache, args_->k_in, args_->v_in,
|
||||
args_->layer_id, args_->block_table,
|
||||
args_->batch_size, args_->max_block_num,
|
||||
args_->cache_seqlens, args_->q_len);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t>
|
||||
cpuinfer_interface(KVCache &kv_cache, intptr_t k_in, intptr_t v_in,
|
||||
int layer_id, intptr_t block_table, int batch_size,
|
||||
int max_block_num, intptr_t cache_seqlens,
|
||||
int q_len) {
|
||||
Args *args = new Args{nullptr,
|
||||
&kv_cache,
|
||||
(ggml_fp16_t *)k_in,
|
||||
(ggml_fp16_t *)v_in,
|
||||
layer_id,
|
||||
(int *)block_table,
|
||||
batch_size,
|
||||
max_block_num,
|
||||
(int *)cache_seqlens,
|
||||
q_len};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
|
||||
class UpdateImportanceBindings {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
KVCache *kv_cache;
|
||||
const ggml_fp16_t *importance;
|
||||
int layer_id;
|
||||
int *block_table;
|
||||
int batch_size;
|
||||
int max_block_num;
|
||||
int *offset;
|
||||
int width;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(
|
||||
&KVCache::update_importance, args_->kv_cache, args_->importance,
|
||||
args_->layer_id, args_->block_table, args_->batch_size,
|
||||
args_->max_block_num, args_->offset, args_->width);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t>
|
||||
cpuinfer_interface(KVCache &kv_cache, intptr_t importance, int layer_id,
|
||||
intptr_t block_table, int batch_size,
|
||||
int max_block_num, intptr_t offset, int width) {
|
||||
Args *args = new Args{nullptr,
|
||||
&kv_cache,
|
||||
(const ggml_fp16_t *)importance,
|
||||
layer_id,
|
||||
(int *)block_table,
|
||||
batch_size,
|
||||
max_block_num,
|
||||
(int *)offset,
|
||||
width};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
|
||||
class AttnWithKVCacheBindings {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
KVCache *kv_cache;
|
||||
const ggml_fp16_t *q_in;
|
||||
const ggml_fp16_t *k_in;
|
||||
const ggml_fp16_t *v_in;
|
||||
ggml_fp16_t *output;
|
||||
float *attn_lse;
|
||||
int layer_idx;
|
||||
int generate_token_idx;
|
||||
int q_len;
|
||||
int batch_size;
|
||||
int max_block_num;
|
||||
int *block_table;
|
||||
int *cache_seqlens;
|
||||
int topk;
|
||||
int local;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(
|
||||
&KVCache::attn_with_kvcache, args_->kv_cache, args_->q_in,
|
||||
args_->k_in, args_->v_in, args_->output, args_->attn_lse,
|
||||
args_->layer_idx, args_->generate_token_idx, args_->q_len,
|
||||
args_->batch_size, args_->max_block_num, args_->block_table,
|
||||
args_->cache_seqlens, args_->topk, args_->local);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t>
|
||||
cpuinfer_interface(KVCache &kv_cache, intptr_t q_in, intptr_t k_in,
|
||||
intptr_t v_in, intptr_t output, intptr_t attn_lse,
|
||||
int layer_idx, int generate_token_idx, int q_len,
|
||||
int batch_size, int max_block_num,
|
||||
intptr_t block_table, intptr_t cache_seqlens,
|
||||
int topk, int local) {
|
||||
Args *args = new Args{nullptr,
|
||||
&kv_cache,
|
||||
(const ggml_fp16_t *)q_in,
|
||||
(const ggml_fp16_t *)k_in,
|
||||
(const ggml_fp16_t *)v_in,
|
||||
(ggml_fp16_t *)output,
|
||||
(float *)attn_lse,
|
||||
layer_idx,
|
||||
generate_token_idx,
|
||||
q_len,
|
||||
batch_size,
|
||||
max_block_num,
|
||||
(int *)block_table,
|
||||
(int *)cache_seqlens,
|
||||
topk,
|
||||
local};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
|
||||
class ClearImportanceAllLayersBindings {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
KVCache *kv_cache;
|
||||
int *block_table;
|
||||
int *cache_seqlens;
|
||||
int batch_size;
|
||||
int max_block_num;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(&KVCache::clear_importance_all_layers,
|
||||
args_->kv_cache, args_->block_table,
|
||||
args_->cache_seqlens, args_->batch_size,
|
||||
args_->max_block_num);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t>
|
||||
cpuinfer_interface(KVCache &kv_cache, intptr_t block_table,
|
||||
intptr_t cache_seqlens, int batch_size,
|
||||
int max_block_num) {
|
||||
Args *args = new Args{nullptr,
|
||||
&kv_cache,
|
||||
(int *)block_table,
|
||||
(int *)cache_seqlens,
|
||||
batch_size,
|
||||
max_block_num};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
|
||||
class CalcAnchorAllLayersBindinds {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
KVCache *kv_cache;
|
||||
int *block_table;
|
||||
int *cache_seqlens;
|
||||
int batch_size;
|
||||
int max_block_num;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(&KVCache::calc_anchor_all_layers,
|
||||
args_->kv_cache, args_->block_table,
|
||||
args_->cache_seqlens, args_->batch_size,
|
||||
args_->max_block_num);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t>
|
||||
cpuinfer_interface(KVCache &kv_cache, intptr_t block_table,
|
||||
intptr_t cache_seqlens, int batch_size,
|
||||
int max_block_num) {
|
||||
Args *args = new Args{nullptr,
|
||||
&kv_cache,
|
||||
(int *)block_table,
|
||||
(int *)cache_seqlens,
|
||||
batch_size,
|
||||
max_block_num};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
|
||||
class LoadKVCacheBindings {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
KVCache *kv_cache;
|
||||
std::string tensor_file_path;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(&KVCache::load_kvcache, args_->kv_cache,
|
||||
args_->tensor_file_path);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t>
|
||||
cpuinfer_interface(KVCache &kv_cache, std::string tensor_file_path) {
|
||||
Args *args =
|
||||
new Args{nullptr, &kv_cache, (std::string)tensor_file_path};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
class DumpKVCacheBindings {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
KVCache *kv_cache;
|
||||
int *block_table;
|
||||
int cache_total_len;
|
||||
std::string tensor_file_path;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(&KVCache::dump_kvcache, args_->kv_cache,
|
||||
args_->block_table, args_->cache_total_len,
|
||||
args_->tensor_file_path);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t>
|
||||
cpuinfer_interface(KVCache &kv_cache, intptr_t block_table,
|
||||
int cache_total_len, std::string tensor_file_path) {
|
||||
Args *args =
|
||||
new Args{nullptr, &kv_cache, (int *)block_table,
|
||||
cache_total_len, (std::string)tensor_file_path};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
};
|
||||
|
||||
class LinearBindings {
|
||||
public:
|
||||
class WarmUpBindinds {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
Linear *linear;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(&Linear::warm_up, args_->linear);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t>
|
||||
cpuinfer_interface(Linear &linear) {
|
||||
Args *args = new Args{nullptr, &linear};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
class ForwardBindings {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
Linear *linear;
|
||||
int qlen;
|
||||
const void *input;
|
||||
void *output;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(&Linear::forward, args_->linear,
|
||||
args_->qlen, args_->input, args_->output);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t>
|
||||
cpuinfer_interface(Linear &linear, int qlen, intptr_t input,
|
||||
intptr_t output) {
|
||||
Args *args = new Args{nullptr, &linear, qlen, (const void *)input,
|
||||
(void *)output};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
};
|
||||
|
||||
class MLPBindings {
|
||||
public:
|
||||
class WarmUpBindinds {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
MLP *mlp;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(&MLP::warm_up, args_->mlp);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t> cpuinfer_interface(MLP &mlp) {
|
||||
Args *args = new Args{nullptr, &mlp};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
class ForwardBindings {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
MLP *mlp;
|
||||
int qlen;
|
||||
const void *input;
|
||||
void *output;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(&MLP::forward, args_->mlp, args_->qlen,
|
||||
args_->input, args_->output);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t>
|
||||
cpuinfer_interface(MLP &mlp, int qlen, intptr_t input,
|
||||
intptr_t output) {
|
||||
Args *args = new Args{nullptr, &mlp, qlen, (const void *)input,
|
||||
(void *)output};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
};
|
||||
|
||||
class MOEBindings {
|
||||
public:
|
||||
class WarmUpBindinds {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
MOE *moe;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(&MOE::warm_up, args_->moe);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t> cpuinfer_interface(MOE &moe) {
|
||||
Args *args = new Args{nullptr, &moe};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
class ForwardBindings {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
MOE *moe;
|
||||
int qlen;
|
||||
int k;
|
||||
const uint64_t *expert_ids;
|
||||
const float *weights;
|
||||
const void *input;
|
||||
void *output;
|
||||
int *batch_size_tensor;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(
|
||||
&MOE::forward, args_->moe, args_->qlen, args_->k,
|
||||
args_->expert_ids, args_->weights, args_->input, args_->output, args_->batch_size_tensor);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t>
|
||||
cpuinfer_interface(MOE &moe, int qlen, int k, intptr_t expert_ids,
|
||||
intptr_t weights, intptr_t input, intptr_t output, intptr_t batch_size_tensor) {
|
||||
Args *args = new Args{nullptr,
|
||||
&moe,
|
||||
qlen,
|
||||
k,
|
||||
(const uint64_t *)expert_ids,
|
||||
(const float *)weights,
|
||||
(const void *)input,
|
||||
(void *)output,
|
||||
(int *)batch_size_tensor};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
};
|
||||
|
||||
|
||||
#if defined(__x86_64__) && defined(__HAS_AVX512F__) && defined(__HAS_AMX__)
|
||||
template<class T>
|
||||
class AMX_MOEBindings {
|
||||
public:
|
||||
class WarmUpBindings {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
AMX_MOE<T> *moe;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(&AMX_MOE<T>::warm_up, args_->moe);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t> cpuinfer_interface(AMX_MOE<T> &moe) {
|
||||
Args *args = new Args{nullptr, &moe};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
class LoadWeightsBindings {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
AMX_MOE<T> *moe;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(&AMX_MOE<T>::load_weights, args_->moe);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t> cpuinfer_interface(AMX_MOE<T> &moe) {
|
||||
Args *args = new Args{nullptr, &moe};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
class ForwardBindings {
|
||||
public:
|
||||
struct Args {
|
||||
CPUInfer *cpuinfer;
|
||||
AMX_MOE<T> *moe;
|
||||
int qlen;
|
||||
int k;
|
||||
const uint64_t *expert_ids;
|
||||
const float *weights;
|
||||
const void *input;
|
||||
void *output;
|
||||
int *batch_size_tensor;
|
||||
};
|
||||
static void inner(void *args) {
|
||||
Args *args_ = (Args *)args;
|
||||
args_->cpuinfer->enqueue(
|
||||
&AMX_MOE<T>::forward, args_->moe, args_->qlen, args_->k,
|
||||
args_->expert_ids, args_->weights, args_->input, args_->output, args_->batch_size_tensor);
|
||||
}
|
||||
static std::pair<intptr_t, intptr_t>
|
||||
cpuinfer_interface(AMX_MOE<T> &moe, int qlen, int k, intptr_t expert_ids,
|
||||
intptr_t weights, intptr_t input, intptr_t output, intptr_t batch_size_tensor) {
|
||||
Args *args = new Args{nullptr,
|
||||
&moe,
|
||||
qlen,
|
||||
k,
|
||||
(const uint64_t *)expert_ids,
|
||||
(const float *)weights,
|
||||
(const void *)input,
|
||||
(void *)output,
|
||||
(int *)batch_size_tensor};
|
||||
return std::make_pair((intptr_t)&inner, (intptr_t)args);
|
||||
}
|
||||
};
|
||||
};
|
||||
#endif
|
||||
|
||||
PYBIND11_MODULE(cpuinfer_ext, m) {
|
||||
py::class_<CPUInfer>(m, "CPUInfer")
|
||||
.def(py::init<int>())
|
||||
.def("submit", &CPUInfer::submit)
|
||||
.def("submit_with_cuda_stream", &CPUInfer::submit_with_cuda_stream)
|
||||
.def("sync", &CPUInfer::sync)
|
||||
.def("sync_with_cuda_stream", &CPUInfer::sync_with_cuda_stream);
|
||||
|
||||
auto linear_module = m.def_submodule("linear");
|
||||
py::class_<LinearConfig>(linear_module, "LinearConfig")
|
||||
.def(py::init([](int hidden_size, int intermediate_size, int stride,
|
||||
int group_max_len, intptr_t proj, int proj_type,
|
||||
int hidden_type) {
|
||||
return LinearConfig(hidden_size, intermediate_size, stride,
|
||||
group_max_len, (void *)proj,
|
||||
(ggml_type)proj_type, (ggml_type)hidden_type);
|
||||
}));
|
||||
py::class_<Linear>(linear_module, "Linear")
|
||||
.def(py::init<LinearConfig>())
|
||||
.def("warm_up", &LinearBindings::WarmUpBindinds::cpuinfer_interface)
|
||||
.def("forward", &LinearBindings::ForwardBindings::cpuinfer_interface);
|
||||
|
||||
auto mlp_module = m.def_submodule("mlp");
|
||||
py::class_<MLPConfig>(mlp_module, "MLPConfig")
|
||||
.def(py::init([](int hidden_size, int intermediate_size, int stride,
|
||||
int group_max_len, intptr_t gate_proj,
|
||||
intptr_t up_proj, intptr_t down_proj, int gate_type,
|
||||
int up_type, int down_type, int hidden_type) {
|
||||
return MLPConfig(hidden_size, intermediate_size, stride,
|
||||
group_max_len, (void *)gate_proj, (void *)up_proj,
|
||||
(void *)down_proj, (ggml_type)gate_type,
|
||||
(ggml_type)up_type, (ggml_type)down_type,
|
||||
(ggml_type)hidden_type);
|
||||
}));
|
||||
py::class_<MLP>(mlp_module, "MLP")
|
||||
.def(py::init<MLPConfig>())
|
||||
.def("warm_up", &MLPBindings::WarmUpBindinds::cpuinfer_interface)
|
||||
.def("forward", &MLPBindings::ForwardBindings::cpuinfer_interface);
|
||||
|
||||
auto moe_module = m.def_submodule("moe");
|
||||
py::class_<MOEConfig>(moe_module, "MOEConfig")
|
||||
.def(py::init([](int expert_num, int routed_expert_num, int hidden_size,
|
||||
int intermediate_size, int stride, int group_min_len,
|
||||
int group_max_len, bool use_silu, intptr_t gate_proj,
|
||||
intptr_t up_proj, intptr_t down_proj, int gate_type,
|
||||
int up_type, int down_type, int hidden_type) {
|
||||
return MOEConfig(expert_num, routed_expert_num, hidden_size,
|
||||
intermediate_size, stride, group_min_len,
|
||||
group_max_len, use_silu, (void *)gate_proj, (void *)up_proj,
|
||||
(void *)down_proj, (ggml_type)gate_type,
|
||||
(ggml_type)up_type, (ggml_type)down_type,
|
||||
(ggml_type)hidden_type);
|
||||
}));
|
||||
py::class_<MOE>(moe_module, "MOE")
|
||||
.def(py::init<MOEConfig>())
|
||||
.def("warm_up", &MOEBindings::WarmUpBindinds::cpuinfer_interface)
|
||||
.def("forward", &MOEBindings::ForwardBindings::cpuinfer_interface);
|
||||
|
||||
|
||||
#if defined(__x86_64__) && defined(__HAS_AVX512F__) && defined(__HAS_AMX__)
|
||||
py::class_<AMX_MOEConfig>(moe_module, "AMX_MOEConfig")
|
||||
.def(py::init([](int expert_num, int routed_expert_num, int hidden_size,
|
||||
int intermediate_size,
|
||||
int max_len, bool use_silu, intptr_t gate_proj,
|
||||
intptr_t up_proj, intptr_t down_proj) {
|
||||
return AMX_MOEConfig(expert_num, routed_expert_num, hidden_size,
|
||||
intermediate_size,
|
||||
max_len, use_silu, (void *)gate_proj,
|
||||
(void *)up_proj, (void *)down_proj);
|
||||
}));
|
||||
|
||||
py::class_<AMX_MOE<amx::GemmKernel224BF>>(moe_module, "AMXBF16_MOE")
|
||||
.def(py::init<AMX_MOEConfig>())
|
||||
.def("warm_up", &AMX_MOEBindings<amx::GemmKernel224BF>::WarmUpBindings::cpuinfer_interface)
|
||||
.def("load_weights", &AMX_MOEBindings<amx::GemmKernel224BF>::LoadWeightsBindings::cpuinfer_interface)
|
||||
.def("forward", &AMX_MOEBindings<amx::GemmKernel224BF>::ForwardBindings::cpuinfer_interface);
|
||||
py::class_<AMX_MOE<amx::GemmKernel224Int8>>(moe_module, "AMXInt8_MOE")
|
||||
.def(py::init<AMX_MOEConfig>())
|
||||
.def("warm_up", &AMX_MOEBindings<amx::GemmKernel224Int8>::WarmUpBindings::cpuinfer_interface)
|
||||
.def("load_weights", &AMX_MOEBindings<amx::GemmKernel224Int8>::LoadWeightsBindings::cpuinfer_interface)
|
||||
.def("forward", &AMX_MOEBindings<amx::GemmKernel224Int8>::ForwardBindings::cpuinfer_interface);
|
||||
|
||||
#endif
|
||||
|
||||
auto kvcache_module = m.def_submodule("kvcache");
|
||||
|
||||
py::enum_<AnchorType>(kvcache_module, "AnchorType")
|
||||
.value("FIXED", AnchorType::FIXED_ANCHOR)
|
||||
.value("DYNAMIC", AnchorType::DYNAMIC)
|
||||
.value("QUEST", AnchorType::QUEST)
|
||||
.value("BLOCK_MAX", AnchorType::BLOCK_MAX)
|
||||
.value("BLOCK_MEAN", AnchorType::BLOCK_MEAN);
|
||||
py::enum_<ggml_type>(kvcache_module, "ggml_type")
|
||||
.value("FP16", ggml_type::GGML_TYPE_F16)
|
||||
.value("FP32", ggml_type::GGML_TYPE_F32)
|
||||
.value("Q4_0", ggml_type::GGML_TYPE_Q4_0)
|
||||
.value("Q8_0", ggml_type::GGML_TYPE_Q8_0);
|
||||
py::enum_<RetrievalType>(kvcache_module, "RetrievalType")
|
||||
.value("LAYER", RetrievalType::LAYER)
|
||||
.value("KVHEAD", RetrievalType::KVHEAD)
|
||||
.value("QHEAD", RetrievalType::QHEAD);
|
||||
|
||||
py::class_<KVCacheConfig>(kvcache_module, "KVCacheConfig")
|
||||
.def(py::init<int, int, int, int, int, int, AnchorType, ggml_type,
|
||||
RetrievalType, int, int, int, int, int, int>())
|
||||
.def_readwrite("layer_num", &KVCacheConfig::layer_num)
|
||||
.def_readwrite("kv_head_num", &KVCacheConfig::kv_head_num)
|
||||
.def_readwrite("q_head_num", &KVCacheConfig::q_head_num)
|
||||
.def_readwrite("head_dim", &KVCacheConfig::head_dim)
|
||||
.def_readwrite("block_len", &KVCacheConfig::block_len)
|
||||
.def_readwrite("anchor_num", &KVCacheConfig::anchor_num)
|
||||
.def_readwrite("anchor_type", &KVCacheConfig::anchor_type)
|
||||
.def_readwrite("kv_type", &KVCacheConfig::kv_type)
|
||||
.def_readwrite("retrieval_type", &KVCacheConfig::retrieval_type)
|
||||
.def_readwrite("layer_step", &KVCacheConfig::layer_step)
|
||||
.def_readwrite("token_step", &KVCacheConfig::token_step)
|
||||
.def_readwrite("layer_offset", &KVCacheConfig::layer_offset)
|
||||
.def_readwrite("max_block_num", &KVCacheConfig::max_block_num)
|
||||
.def_readwrite("max_batch_size", &KVCacheConfig::max_batch_size)
|
||||
.def_readwrite("max_thread_num", &KVCacheConfig::max_thread_num);
|
||||
py::class_<KVCache>(kvcache_module, "KVCache")
|
||||
.def(py::init<KVCacheConfig>())
|
||||
.def("get_cache_total_len", &KVCache::get_cache_total_len)
|
||||
.def("update_cache_total_len",
|
||||
[](KVCache &kvcache, int cache_total_len) {
|
||||
kvcache.update_cache_total_len(cache_total_len);
|
||||
})
|
||||
.def("attn", &KVCacheBindings::AttnBindings::cpuinfer_interface)
|
||||
.def(
|
||||
"get_all_kvcache_one_layer",
|
||||
&KVCacheBindings::GetAllKVCacheOneLayerBindings::cpuinfer_interface)
|
||||
.def("get_and_update_kvcache_fp16",
|
||||
&KVCacheBindings::GetAndUpdateKVCacheFp16Bindings::
|
||||
cpuinfer_interface)
|
||||
.def("get_kvcache_fp16",
|
||||
&KVCacheBindings::GetKVCacheFp16Bindings::cpuinfer_interface)
|
||||
.def("update_kvcache_fp16",
|
||||
&KVCacheBindings::UpdateKVCacheFp16Bindings::cpuinfer_interface)
|
||||
.def("update_importance",
|
||||
&KVCacheBindings::UpdateImportanceBindings::cpuinfer_interface)
|
||||
.def("attn_with_kvcache",
|
||||
&KVCacheBindings::AttnWithKVCacheBindings::cpuinfer_interface)
|
||||
.def("clear_importance_all_layers",
|
||||
&KVCacheBindings::ClearImportanceAllLayersBindings::
|
||||
cpuinfer_interface)
|
||||
.def("calc_anchor_all_layers",
|
||||
&KVCacheBindings::CalcAnchorAllLayersBindinds::cpuinfer_interface);
|
||||
}
|
||||
@@ -0,0 +1,974 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : chenht2022
|
||||
* @Date : 2025-04-25 18:28:12
|
||||
* @Version : 1.0.0
|
||||
* @LastEditors : chenht2022
|
||||
* @LastEditTime : 2025-04-25 18:28:12
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
|
||||
#pragma once
|
||||
#include <array>
|
||||
#include <cassert>
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
#include <immintrin.h>
|
||||
#include <iostream>
|
||||
#include <random>
|
||||
#include <stdexcept>
|
||||
#include <stdlib.h>
|
||||
#include <sys/syscall.h>
|
||||
#include <unistd.h>
|
||||
|
||||
#include "utils.hpp"
|
||||
#include <memory>
|
||||
|
||||
#if (defined(_WIN32) || defined(_WIN64))
|
||||
#define RESTRICT __restrict
|
||||
#else
|
||||
#define RESTRICT __restrict__
|
||||
#endif
|
||||
|
||||
#if (defined(_WIN32) || defined(_WIN64))
|
||||
#define ALWAYS_INLINE __forceinline
|
||||
#elif __has_attribute(always_inline) || defined(__GNUC__)
|
||||
#define ALWAYS_INLINE __attribute__((__always_inline__)) inline
|
||||
#else
|
||||
#define ALWAYS_INLINE inline
|
||||
#endif
|
||||
|
||||
namespace amx {
|
||||
|
||||
#define ARCH_GET_XCOMP_PERM 0x1022
|
||||
#define ARCH_REQ_XCOMP_PERM 0x1023
|
||||
#define XFEATURE_XTILECFG 17
|
||||
#define XFEATURE_XTILEDATA 18
|
||||
|
||||
const int TMMCount = 8;
|
||||
const int MaxTileHeight = 16;
|
||||
const int MaxTileWidth = 64;
|
||||
|
||||
const int AMX_BLK_SIZE = 32;
|
||||
|
||||
#define TMM0 0
|
||||
#define TMM1 1
|
||||
#define TMM2 2
|
||||
#define TMM3 3
|
||||
#define TMM4 4
|
||||
#define TMM5 5
|
||||
#define TMM6 6
|
||||
#define TMM7 7
|
||||
|
||||
inline bool enable_amx() {
|
||||
static thread_local bool initialized = false;
|
||||
if (initialized) {
|
||||
return true;
|
||||
}
|
||||
initialized = true;
|
||||
|
||||
if (syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM, XFEATURE_XTILEDATA)) {
|
||||
printf("\n Fail to do XFEATURE_XTILEDATA \n\n");
|
||||
return false;
|
||||
} else {
|
||||
// printf("\n TILE DATA USE SET - OK \n\n");
|
||||
return true;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
struct alignas(64) TileConfig {
|
||||
uint8_t palette;
|
||||
uint8_t start_row;
|
||||
std::array<uint8_t, 14> __0 = {};
|
||||
std::array<uint16_t, 8> colsb;
|
||||
std::array<uint8_t, 16> __1 = {};
|
||||
std::array<uint8_t, 8> rows;
|
||||
std::array<uint8_t, 8> __2 = {};
|
||||
|
||||
TileConfig() {
|
||||
palette = 1;
|
||||
start_row = 0;
|
||||
for (int i = 0; i < 8; i++) {
|
||||
set_row_col(i, 0, 0);
|
||||
}
|
||||
}
|
||||
|
||||
void set_row_col(int i, uint8_t row, uint16_t col) {
|
||||
colsb[i] = col;
|
||||
rows[i] = row;
|
||||
}
|
||||
|
||||
void set_config() { _tile_loadconfig(this); }
|
||||
|
||||
static void load_data(int to, void *from, size_t stride) {
|
||||
switch (to) {
|
||||
case 0:
|
||||
_tile_loadd(0, from, stride);
|
||||
break;
|
||||
case 1:
|
||||
_tile_loadd(1, from, stride);
|
||||
break;
|
||||
case 2:
|
||||
_tile_loadd(2, from, stride);
|
||||
break;
|
||||
case 3:
|
||||
_tile_loadd(3, from, stride);
|
||||
break;
|
||||
case 4:
|
||||
_tile_loadd(4, from, stride);
|
||||
break;
|
||||
case 5:
|
||||
_tile_loadd(5, from, stride);
|
||||
break;
|
||||
case 6:
|
||||
_tile_loadd(6, from, stride);
|
||||
break;
|
||||
case 7:
|
||||
_tile_loadd(7, from, stride);
|
||||
break;
|
||||
default:
|
||||
throw std::runtime_error("no such tile");
|
||||
}
|
||||
}
|
||||
|
||||
static void store_data(int from, void *to, size_t stride) {
|
||||
switch (from) {
|
||||
case 0:
|
||||
_tile_stored(0, to, stride);
|
||||
break;
|
||||
case 1:
|
||||
_tile_stored(1, to, stride);
|
||||
break;
|
||||
case 2:
|
||||
_tile_stored(2, to, stride);
|
||||
break;
|
||||
case 3:
|
||||
_tile_stored(3, to, stride);
|
||||
break;
|
||||
case 4:
|
||||
_tile_stored(4, to, stride);
|
||||
break;
|
||||
case 5:
|
||||
_tile_stored(5, to, stride);
|
||||
break;
|
||||
case 6:
|
||||
_tile_stored(6, to, stride);
|
||||
break;
|
||||
case 7:
|
||||
_tile_stored(7, to, stride);
|
||||
break;
|
||||
default:
|
||||
throw std::runtime_error("no such tile");
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
static_assert(sizeof(TileConfig) == 64);
|
||||
|
||||
inline void debug_tile(int t) {
|
||||
printf("Tile %d\n", t);
|
||||
uint8_t data[16][64] = {};
|
||||
TileConfig::store_data(t, data, 64);
|
||||
for (int i = 0; i < 16; i++) {
|
||||
for (int j = 0; j < 64; j++) {
|
||||
printf("%3d ", data[i][j]);
|
||||
}
|
||||
printf("\n");
|
||||
}
|
||||
printf("\n");
|
||||
}
|
||||
|
||||
inline void debug_tiles(int to = 8) {
|
||||
for (int i = 0; i < to; i++) {
|
||||
debug_tile(i);
|
||||
}
|
||||
}
|
||||
|
||||
inline void debug_m512(__m512 x) {
|
||||
float data[16];
|
||||
_mm512_storeu_ps(data, x);
|
||||
for (int i = 0; i < 16; i++) {
|
||||
printf("%f ", data[i]);
|
||||
}
|
||||
printf("\n");
|
||||
}
|
||||
|
||||
// transpose utils
|
||||
inline void transpose_16x16_32bit(__m512i *v) {
|
||||
__m512i v1[16];
|
||||
v1[0] = _mm512_unpacklo_epi32(v[0], v[1]);
|
||||
v1[1] = _mm512_unpackhi_epi32(v[0], v[1]);
|
||||
v1[2] = _mm512_unpacklo_epi32(v[2], v[3]);
|
||||
v1[3] = _mm512_unpackhi_epi32(v[2], v[3]);
|
||||
v1[4] = _mm512_unpacklo_epi32(v[4], v[5]);
|
||||
v1[5] = _mm512_unpackhi_epi32(v[4], v[5]);
|
||||
v1[6] = _mm512_unpacklo_epi32(v[6], v[7]);
|
||||
v1[7] = _mm512_unpackhi_epi32(v[6], v[7]);
|
||||
v1[8] = _mm512_unpacklo_epi32(v[8], v[9]);
|
||||
v1[9] = _mm512_unpackhi_epi32(v[8], v[9]);
|
||||
v1[10] = _mm512_unpacklo_epi32(v[10], v[11]);
|
||||
v1[11] = _mm512_unpackhi_epi32(v[10], v[11]);
|
||||
v1[12] = _mm512_unpacklo_epi32(v[12], v[13]);
|
||||
v1[13] = _mm512_unpackhi_epi32(v[12], v[13]);
|
||||
v1[14] = _mm512_unpacklo_epi32(v[14], v[15]);
|
||||
v1[15] = _mm512_unpackhi_epi32(v[14], v[15]);
|
||||
|
||||
v[0] = _mm512_unpacklo_epi64(v1[0], v1[2]);
|
||||
v[1] = _mm512_unpackhi_epi64(v1[0], v1[2]);
|
||||
v[2] = _mm512_unpacklo_epi64(v1[1], v1[3]);
|
||||
v[3] = _mm512_unpackhi_epi64(v1[1], v1[3]);
|
||||
v[4] = _mm512_unpacklo_epi64(v1[4], v1[6]);
|
||||
v[5] = _mm512_unpackhi_epi64(v1[4], v1[6]);
|
||||
v[6] = _mm512_unpacklo_epi64(v1[5], v1[7]);
|
||||
v[7] = _mm512_unpackhi_epi64(v1[5], v1[7]);
|
||||
v[8] = _mm512_unpacklo_epi64(v1[8], v1[10]);
|
||||
v[9] = _mm512_unpackhi_epi64(v1[8], v1[10]);
|
||||
v[10] = _mm512_unpacklo_epi64(v1[9], v1[11]);
|
||||
v[11] = _mm512_unpackhi_epi64(v1[9], v1[11]);
|
||||
v[12] = _mm512_unpacklo_epi64(v1[12], v1[14]);
|
||||
v[13] = _mm512_unpackhi_epi64(v1[12], v1[14]);
|
||||
v[14] = _mm512_unpacklo_epi64(v1[13], v1[15]);
|
||||
v[15] = _mm512_unpackhi_epi64(v1[13], v1[15]);
|
||||
|
||||
v1[0] = _mm512_shuffle_i32x4(v[0], v[4], 0x88);
|
||||
v1[1] = _mm512_shuffle_i32x4(v[1], v[5], 0x88);
|
||||
v1[2] = _mm512_shuffle_i32x4(v[2], v[6], 0x88);
|
||||
v1[3] = _mm512_shuffle_i32x4(v[3], v[7], 0x88);
|
||||
v1[4] = _mm512_shuffle_i32x4(v[0], v[4], 0xdd);
|
||||
v1[5] = _mm512_shuffle_i32x4(v[1], v[5], 0xdd);
|
||||
v1[6] = _mm512_shuffle_i32x4(v[2], v[6], 0xdd);
|
||||
v1[7] = _mm512_shuffle_i32x4(v[3], v[7], 0xdd);
|
||||
v1[8] = _mm512_shuffle_i32x4(v[8], v[12], 0x88);
|
||||
v1[9] = _mm512_shuffle_i32x4(v[9], v[13], 0x88);
|
||||
v1[10] = _mm512_shuffle_i32x4(v[10], v[14], 0x88);
|
||||
v1[11] = _mm512_shuffle_i32x4(v[11], v[15], 0x88);
|
||||
v1[12] = _mm512_shuffle_i32x4(v[8], v[12], 0xdd);
|
||||
v1[13] = _mm512_shuffle_i32x4(v[9], v[13], 0xdd);
|
||||
v1[14] = _mm512_shuffle_i32x4(v[10], v[14], 0xdd);
|
||||
v1[15] = _mm512_shuffle_i32x4(v[11], v[15], 0xdd);
|
||||
|
||||
v[0] = _mm512_shuffle_i32x4(v1[0], v1[8], 0x88);
|
||||
v[1] = _mm512_shuffle_i32x4(v1[1], v1[9], 0x88);
|
||||
v[2] = _mm512_shuffle_i32x4(v1[2], v1[10], 0x88);
|
||||
v[3] = _mm512_shuffle_i32x4(v1[3], v1[11], 0x88);
|
||||
v[4] = _mm512_shuffle_i32x4(v1[4], v1[12], 0x88);
|
||||
v[5] = _mm512_shuffle_i32x4(v1[5], v1[13], 0x88);
|
||||
v[6] = _mm512_shuffle_i32x4(v1[6], v1[14], 0x88);
|
||||
v[7] = _mm512_shuffle_i32x4(v1[7], v1[15], 0x88);
|
||||
v[8] = _mm512_shuffle_i32x4(v1[0], v1[8], 0xdd);
|
||||
v[9] = _mm512_shuffle_i32x4(v1[1], v1[9], 0xdd);
|
||||
v[10] = _mm512_shuffle_i32x4(v1[2], v1[10], 0xdd);
|
||||
v[11] = _mm512_shuffle_i32x4(v1[3], v1[11], 0xdd);
|
||||
v[12] = _mm512_shuffle_i32x4(v1[4], v1[12], 0xdd);
|
||||
v[13] = _mm512_shuffle_i32x4(v1[5], v1[13], 0xdd);
|
||||
v[14] = _mm512_shuffle_i32x4(v1[6], v1[14], 0xdd);
|
||||
v[15] = _mm512_shuffle_i32x4(v1[7], v1[15], 0xdd);
|
||||
}
|
||||
|
||||
/*
|
||||
Transpose 16x16 32-bit elements
|
||||
Note that v must be 64 byte aligned
|
||||
*/
|
||||
inline void transpose_16x16_32bit(__m512i *v, size_t stride) {
|
||||
assert(reinterpret_cast<intptr_t>(v) % 64 == 0 && "v must be 64 aligned");
|
||||
|
||||
auto stride_v = [=](int i) { return offset_pointer(v, i * stride); };
|
||||
__m512i v1[16];
|
||||
|
||||
v1[0] = _mm512_unpacklo_epi32(*stride_v(0), *stride_v(1));
|
||||
v1[1] = _mm512_unpackhi_epi32(*stride_v(0), *stride_v(1));
|
||||
v1[2] = _mm512_unpacklo_epi32(*stride_v(2), *stride_v(3));
|
||||
v1[3] = _mm512_unpackhi_epi32(*stride_v(2), *stride_v(3));
|
||||
v1[4] = _mm512_unpacklo_epi32(*stride_v(4), *stride_v(5));
|
||||
v1[5] = _mm512_unpackhi_epi32(*stride_v(4), *stride_v(5));
|
||||
v1[6] = _mm512_unpacklo_epi32(*stride_v(6), *stride_v(7));
|
||||
v1[7] = _mm512_unpackhi_epi32(*stride_v(6), *stride_v(7));
|
||||
v1[8] = _mm512_unpacklo_epi32(*stride_v(8), *stride_v(9));
|
||||
v1[9] = _mm512_unpackhi_epi32(*stride_v(8), *stride_v(9));
|
||||
v1[10] = _mm512_unpacklo_epi32(*stride_v(10), *stride_v(11));
|
||||
v1[11] = _mm512_unpackhi_epi32(*stride_v(10), *stride_v(11));
|
||||
v1[12] = _mm512_unpacklo_epi32(*stride_v(12), *stride_v(13));
|
||||
v1[13] = _mm512_unpackhi_epi32(*stride_v(12), *stride_v(13));
|
||||
v1[14] = _mm512_unpacklo_epi32(*stride_v(14), *stride_v(15));
|
||||
v1[15] = _mm512_unpackhi_epi32(*stride_v(14), *stride_v(15));
|
||||
|
||||
*stride_v(0) = _mm512_unpacklo_epi64(v1[0], v1[2]);
|
||||
*stride_v(1) = _mm512_unpackhi_epi64(v1[0], v1[2]);
|
||||
*stride_v(2) = _mm512_unpacklo_epi64(v1[1], v1[3]);
|
||||
*stride_v(3) = _mm512_unpackhi_epi64(v1[1], v1[3]);
|
||||
*stride_v(4) = _mm512_unpacklo_epi64(v1[4], v1[6]);
|
||||
*stride_v(5) = _mm512_unpackhi_epi64(v1[4], v1[6]);
|
||||
*stride_v(6) = _mm512_unpacklo_epi64(v1[5], v1[7]);
|
||||
*stride_v(7) = _mm512_unpackhi_epi64(v1[5], v1[7]);
|
||||
*stride_v(8) = _mm512_unpacklo_epi64(v1[8], v1[10]);
|
||||
*stride_v(9) = _mm512_unpackhi_epi64(v1[8], v1[10]);
|
||||
*stride_v(10) = _mm512_unpacklo_epi64(v1[9], v1[11]);
|
||||
*stride_v(11) = _mm512_unpackhi_epi64(v1[9], v1[11]);
|
||||
*stride_v(12) = _mm512_unpacklo_epi64(v1[12], v1[14]);
|
||||
*stride_v(13) = _mm512_unpackhi_epi64(v1[12], v1[14]);
|
||||
*stride_v(14) = _mm512_unpacklo_epi64(v1[13], v1[15]);
|
||||
*stride_v(15) = _mm512_unpackhi_epi64(v1[13], v1[15]);
|
||||
|
||||
v1[0] = _mm512_shuffle_i32x4(*stride_v(0), *stride_v(4), 0x88);
|
||||
v1[1] = _mm512_shuffle_i32x4(*stride_v(1), *stride_v(5), 0x88);
|
||||
v1[2] = _mm512_shuffle_i32x4(*stride_v(2), *stride_v(6), 0x88);
|
||||
v1[3] = _mm512_shuffle_i32x4(*stride_v(3), *stride_v(7), 0x88);
|
||||
v1[4] = _mm512_shuffle_i32x4(*stride_v(0), *stride_v(4), 0xdd);
|
||||
v1[5] = _mm512_shuffle_i32x4(*stride_v(1), *stride_v(5), 0xdd);
|
||||
v1[6] = _mm512_shuffle_i32x4(*stride_v(2), *stride_v(6), 0xdd);
|
||||
v1[7] = _mm512_shuffle_i32x4(*stride_v(3), *stride_v(7), 0xdd);
|
||||
v1[8] = _mm512_shuffle_i32x4(*stride_v(8), *stride_v(12), 0x88);
|
||||
v1[9] = _mm512_shuffle_i32x4(*stride_v(9), *stride_v(13), 0x88);
|
||||
v1[10] = _mm512_shuffle_i32x4(*stride_v(10), *stride_v(14), 0x88);
|
||||
v1[11] = _mm512_shuffle_i32x4(*stride_v(11), *stride_v(15), 0x88);
|
||||
v1[12] = _mm512_shuffle_i32x4(*stride_v(8), *stride_v(12), 0xdd);
|
||||
v1[13] = _mm512_shuffle_i32x4(*stride_v(9), *stride_v(13), 0xdd);
|
||||
v1[14] = _mm512_shuffle_i32x4(*stride_v(10), *stride_v(14), 0xdd);
|
||||
v1[15] = _mm512_shuffle_i32x4(*stride_v(11), *stride_v(15), 0xdd);
|
||||
|
||||
*stride_v(0) = _mm512_shuffle_i32x4(v1[0], v1[8], 0x88);
|
||||
*stride_v(1) = _mm512_shuffle_i32x4(v1[1], v1[9], 0x88);
|
||||
*stride_v(2) = _mm512_shuffle_i32x4(v1[2], v1[10], 0x88);
|
||||
*stride_v(3) = _mm512_shuffle_i32x4(v1[3], v1[11], 0x88);
|
||||
*stride_v(4) = _mm512_shuffle_i32x4(v1[4], v1[12], 0x88);
|
||||
*stride_v(5) = _mm512_shuffle_i32x4(v1[5], v1[13], 0x88);
|
||||
*stride_v(6) = _mm512_shuffle_i32x4(v1[6], v1[14], 0x88);
|
||||
*stride_v(7) = _mm512_shuffle_i32x4(v1[7], v1[15], 0x88);
|
||||
*stride_v(8) = _mm512_shuffle_i32x4(v1[0], v1[8], 0xdd);
|
||||
*stride_v(9) = _mm512_shuffle_i32x4(v1[1], v1[9], 0xdd);
|
||||
*stride_v(10) = _mm512_shuffle_i32x4(v1[2], v1[10], 0xdd);
|
||||
*stride_v(11) = _mm512_shuffle_i32x4(v1[3], v1[11], 0xdd);
|
||||
*stride_v(12) = _mm512_shuffle_i32x4(v1[4], v1[12], 0xdd);
|
||||
*stride_v(13) = _mm512_shuffle_i32x4(v1[5], v1[13], 0xdd);
|
||||
*stride_v(14) = _mm512_shuffle_i32x4(v1[6], v1[14], 0xdd);
|
||||
*stride_v(15) = _mm512_shuffle_i32x4(v1[7], v1[15], 0xdd);
|
||||
}
|
||||
|
||||
struct GemmKernel224BF {
|
||||
using dt = ggml_bf16_t;
|
||||
using output_t = float;
|
||||
static const int TILE_M = 16;
|
||||
static const int TILE_K = 32;
|
||||
static const int TILE_N = 16;
|
||||
static const int VNNI_BLK = 2;
|
||||
|
||||
static const int M_STEP = TILE_M * 2;
|
||||
static const int N_STEP = TILE_N * 2;
|
||||
static const int K_STEP = TILE_K;
|
||||
|
||||
static inline const int N_BLOCK = 256;
|
||||
static inline const int K_BLOCK = 1792;
|
||||
|
||||
static int recommended_nth(int n) { return (n + N_BLOCK - 1) / N_BLOCK; }
|
||||
|
||||
static std::pair<int, int> split_range_n(int n, int ith, int nth) {
|
||||
int n_start = N_BLOCK * ith;
|
||||
int n_end = std::min(n, N_BLOCK * (ith + 1));
|
||||
return {n_start, n_end};
|
||||
}
|
||||
|
||||
static void config() {
|
||||
enable_amx();
|
||||
TileConfig tile_config;
|
||||
|
||||
// size is 16 x 32
|
||||
for (int i = 0; i < 2; i++)
|
||||
tile_config.set_row_col(i, TILE_M, TILE_K * sizeof(dt));
|
||||
|
||||
// size is 16 x 32
|
||||
for (int i = 2; i < 4; i++)
|
||||
tile_config.set_row_col(i, TILE_K / VNNI_BLK, TILE_N * VNNI_BLK * sizeof(dt));
|
||||
|
||||
// size is 16 x 16
|
||||
for (int i = 4; i < 8; i++)
|
||||
tile_config.set_row_col(i, TILE_M, TILE_N * sizeof(output_t));
|
||||
|
||||
tile_config.set_config();
|
||||
}
|
||||
|
||||
static void load_a(dt *a, size_t lda) {
|
||||
_tile_loadd(0, a, lda);
|
||||
_tile_loadd(1, offset_pointer(a, lda * TILE_M), lda);
|
||||
}
|
||||
|
||||
static void load_b(dt *b, size_t ldb) {
|
||||
_tile_loadd(2, b, ldb);
|
||||
_tile_loadd(3, offset_pointer(b, ldb * TILE_N), ldb);
|
||||
}
|
||||
|
||||
static void clean_c() {
|
||||
_tile_zero(4);
|
||||
_tile_zero(5);
|
||||
_tile_zero(6);
|
||||
_tile_zero(7);
|
||||
}
|
||||
|
||||
static void load_c(output_t *c, size_t ldc) {
|
||||
_tile_loadd(4, c, ldc);
|
||||
_tile_loadd(5, offset_pointer(c, TILE_N * sizeof(output_t)), ldc);
|
||||
_tile_loadd(6, offset_pointer(c, ldc * TILE_M), ldc);
|
||||
_tile_loadd(7, offset_pointer(c, ldc * TILE_M + TILE_N * sizeof(output_t)), ldc);
|
||||
}
|
||||
|
||||
static void store_c(output_t *c, size_t ldc) {
|
||||
_tile_stored(4, c, ldc);
|
||||
_tile_stored(5, offset_pointer(c, TILE_N * sizeof(output_t)), ldc);
|
||||
_tile_stored(6, offset_pointer(c, ldc * TILE_M), ldc);
|
||||
_tile_stored(7, offset_pointer(c, ldc * TILE_M + TILE_N * sizeof(output_t)), ldc);
|
||||
}
|
||||
|
||||
static void run_tile() {
|
||||
_tile_dpbf16ps(4, 0, 2);
|
||||
_tile_dpbf16ps(5, 0, 3);
|
||||
_tile_dpbf16ps(6, 1, 2);
|
||||
_tile_dpbf16ps(7, 1, 3);
|
||||
}
|
||||
|
||||
struct BufferA {
|
||||
ggml_bf16_t *a;
|
||||
int max_m, k;
|
||||
|
||||
static size_t required_size(int max_m, int k) { return max_m * k * sizeof(ggml_bf16_t); }
|
||||
|
||||
BufferA(int max_m, int k, void *ptr) : max_m(max_m), k(k) {
|
||||
assert(reinterpret_cast<intptr_t>(ptr) % 64 == 0);
|
||||
assert(max_m % M_STEP == 0);
|
||||
assert(k % K_STEP == 0);
|
||||
a = reinterpret_cast<ggml_bf16_t *>(ptr);
|
||||
}
|
||||
|
||||
void from_mat(int m, ggml_bf16_t *src, int ith, int nth) {
|
||||
assert(m <= max_m);
|
||||
assert(ith == 0 && nth == 1);
|
||||
int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP;
|
||||
for (int m_begin = 0; m_begin < m; m_begin += M_STEP) {
|
||||
for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) {
|
||||
int k_block_size = std::min(K_BLOCK, k - k_block_begin);
|
||||
for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) {
|
||||
for (int i = 0; i < M_STEP && m_begin + i < m; i++) {
|
||||
__m512i *s = (__m512i *)(src + (m_begin + i) * k + k_block_begin + k_begin);
|
||||
__m512i *d = (__m512i *)(a + k_block_begin * m_block_size + m_begin * k_block_size + k_begin * M_STEP +
|
||||
i * K_STEP);
|
||||
avx512_copy_32xbf16(s, d);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
ggml_bf16_t *get_submat(int m, int k, int m_begin, int k_begin) {
|
||||
int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP;
|
||||
int k_block_begin = k_begin / K_BLOCK * K_BLOCK;
|
||||
k_begin -= k_block_begin;
|
||||
int k_block_size = std::min(K_BLOCK, k - k_block_begin);
|
||||
return a + k_block_begin * m_block_size + m_begin * k_block_size + k_begin * M_STEP;
|
||||
}
|
||||
};
|
||||
|
||||
struct BufferB {
|
||||
ggml_bf16_t *b;
|
||||
int n, k;
|
||||
|
||||
static size_t required_size(int n, int k) { return n * k * sizeof(ggml_bf16_t); }
|
||||
|
||||
BufferB(int n, int k, void *ptr) : n(n), k(k) {
|
||||
assert(reinterpret_cast<intptr_t>(ptr) % 64 == 0);
|
||||
assert(n % N_STEP == 0);
|
||||
assert(k % K_STEP == 0);
|
||||
b = reinterpret_cast<ggml_bf16_t *>(ptr);
|
||||
}
|
||||
|
||||
void from_mat(ggml_bf16_t *src, int ith, int nth) {
|
||||
auto [n_start, n_end] = split_range_n(n, ith, nth);
|
||||
int n_block_begin = n_start;
|
||||
int n_block_size = n_end - n_block_begin;
|
||||
for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) {
|
||||
for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) {
|
||||
int k_block_size = std::min(K_BLOCK, k - k_block_begin);
|
||||
for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) {
|
||||
for (int i = 0; i < N_STEP; i++) {
|
||||
__m512i *s = (__m512i *)(src + (n_block_begin + n_begin + i) * k + k_block_begin + k_begin);
|
||||
__m512i *d = (__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size +
|
||||
k_begin * N_STEP + i * K_STEP);
|
||||
avx512_copy_32xbf16(s, d);
|
||||
}
|
||||
transpose_16x16_32bit((__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size +
|
||||
n_begin * k_block_size + k_begin * N_STEP));
|
||||
transpose_16x16_32bit((__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size +
|
||||
n_begin * k_block_size + k_begin * N_STEP + TILE_N * K_STEP));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
ggml_bf16_t *get_submat(int n, int k, int n_begin, int k_begin) {
|
||||
int n_block_begin = n_begin / N_BLOCK * N_BLOCK;
|
||||
n_begin -= n_block_begin;
|
||||
int n_block_size = std::min(N_BLOCK, n - n_block_begin);
|
||||
int k_block_begin = k_begin / K_BLOCK * K_BLOCK;
|
||||
k_begin -= k_block_begin;
|
||||
int k_block_size = std::min(K_BLOCK, k - k_block_begin);
|
||||
return b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + k_begin * N_STEP;
|
||||
}
|
||||
};
|
||||
|
||||
struct BufferC {
|
||||
float *c;
|
||||
int max_m, n;
|
||||
|
||||
static size_t required_size(int max_m, int n) { return max_m * n * sizeof(float); }
|
||||
|
||||
BufferC(int max_m, int n, void *ptr) : max_m(max_m), n(n) {
|
||||
assert(reinterpret_cast<intptr_t>(ptr) % 64 == 0);
|
||||
assert(max_m % M_STEP == 0);
|
||||
assert(n % N_STEP == 0);
|
||||
c = reinterpret_cast<float *>(ptr);
|
||||
}
|
||||
|
||||
void to_mat(int m, ggml_bf16_t *dst, int ith, int nth) {
|
||||
assert(m <= max_m);
|
||||
auto [n_start, n_end] = split_range_n(n, ith, nth);
|
||||
int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP;
|
||||
int n_block_begin = n_start;
|
||||
int n_block_size = n_end - n_block_begin;
|
||||
for (int m_begin = 0; m_begin < m; m_begin += M_STEP) {
|
||||
for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) {
|
||||
for (int i = 0; i < M_STEP && m_begin + i < m; i++) {
|
||||
__m512 *x0 =
|
||||
(__m512 *)(c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP + i * N_STEP);
|
||||
__m512 *x1 = (__m512 *)(c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP +
|
||||
i * N_STEP + 16);
|
||||
avx512_32xfp32_to_32xbf16(x0, x1, (__m512i *)(dst + (m_begin + i) * n + n_block_begin + n_begin));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
float *get_submat(int m, int n, int m_begin, int n_begin) {
|
||||
int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP;
|
||||
int n_block_begin = n_begin / N_BLOCK * N_BLOCK;
|
||||
int n_block_size = std::min(N_BLOCK, n - n_block_begin);
|
||||
n_begin -= n_block_begin;
|
||||
return c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP;
|
||||
}
|
||||
};
|
||||
};
|
||||
|
||||
struct GemmKernel224Int8 {
|
||||
using dt = int8_t;
|
||||
using output_t = int32_t;
|
||||
static const int TILE_M = 16;
|
||||
static const int TILE_K = 64;
|
||||
static const int TILE_N = 16;
|
||||
static const int VNNI_BLK = 4;
|
||||
|
||||
static const int M_STEP = TILE_M * 2;
|
||||
static const int N_STEP = TILE_N * 2;
|
||||
static const int K_STEP = TILE_K;
|
||||
|
||||
static inline const int N_BLOCK = 256;
|
||||
static inline const int K_BLOCK = 3584;
|
||||
|
||||
static int recommended_nth(int n) { return (n + N_BLOCK - 1) / N_BLOCK; }
|
||||
|
||||
static std::pair<int, int> split_range_n(int n, int ith, int nth) {
|
||||
int n_start = N_BLOCK * ith;
|
||||
int n_end = std::min(n, N_BLOCK * (ith + 1));
|
||||
return {n_start, n_end};
|
||||
}
|
||||
|
||||
static void config() {
|
||||
enable_amx();
|
||||
TileConfig tile_config;
|
||||
|
||||
// size is 16 x 64
|
||||
for (int i = 0; i < 2; i++)
|
||||
tile_config.set_row_col(i, TILE_M, TILE_K * sizeof(dt));
|
||||
|
||||
// size is 16 x 64
|
||||
for (int i = 2; i < 4; i++)
|
||||
tile_config.set_row_col(i, TILE_K / VNNI_BLK, TILE_N * VNNI_BLK * sizeof(dt));
|
||||
|
||||
// size is 16 x 16
|
||||
for (int i = 4; i < 8; i++)
|
||||
tile_config.set_row_col(i, TILE_M, TILE_N * sizeof(output_t));
|
||||
|
||||
tile_config.set_config();
|
||||
}
|
||||
|
||||
static void load_a(dt *a, size_t lda) {
|
||||
_tile_loadd(0, a, lda);
|
||||
_tile_loadd(1, offset_pointer(a, lda * TILE_M), lda);
|
||||
}
|
||||
|
||||
static void load_b(dt *b, size_t ldb) {
|
||||
_tile_loadd(2, b, ldb);
|
||||
_tile_loadd(3, offset_pointer(b, ldb * TILE_N), ldb);
|
||||
}
|
||||
|
||||
static void clean_c() {
|
||||
_tile_zero(4);
|
||||
_tile_zero(5);
|
||||
_tile_zero(6);
|
||||
_tile_zero(7);
|
||||
}
|
||||
|
||||
static void load_c(output_t *c, size_t ldc) {
|
||||
_tile_loadd(4, c, ldc);
|
||||
_tile_loadd(5, offset_pointer(c, TILE_N * sizeof(output_t)), ldc);
|
||||
_tile_loadd(6, offset_pointer(c, ldc * TILE_M), ldc);
|
||||
_tile_loadd(7, offset_pointer(c, ldc * TILE_M + TILE_N * sizeof(output_t)), ldc);
|
||||
}
|
||||
|
||||
static void store_c(output_t *c, size_t ldc) {
|
||||
_tile_stored(4, c, ldc);
|
||||
_tile_stored(5, offset_pointer(c, TILE_N * sizeof(output_t)), ldc);
|
||||
_tile_stored(6, offset_pointer(c, ldc * TILE_M), ldc);
|
||||
_tile_stored(7, offset_pointer(c, ldc * TILE_M + TILE_N * sizeof(output_t)), ldc);
|
||||
}
|
||||
|
||||
static void run_tile() {
|
||||
_tile_dpbssd(4, 0, 2);
|
||||
_tile_dpbssd(5, 0, 3);
|
||||
_tile_dpbssd(6, 1, 2);
|
||||
_tile_dpbssd(7, 1, 3);
|
||||
}
|
||||
|
||||
struct BufferA {
|
||||
int8_t *a;
|
||||
float *d;
|
||||
int max_m, k;
|
||||
|
||||
static size_t required_size(int max_m, int k) { return max_m * k * sizeof(int8_t) + max_m * sizeof(float); }
|
||||
|
||||
BufferA(int max_m, int k, void *ptr) : max_m(max_m), k(k) {
|
||||
assert(reinterpret_cast<intptr_t>(ptr) % 64 == 0);
|
||||
assert(max_m % M_STEP == 0);
|
||||
assert(k % K_STEP == 0);
|
||||
a = reinterpret_cast<int8_t *>(ptr);
|
||||
d = reinterpret_cast<float *>(a + max_m * k);
|
||||
}
|
||||
|
||||
void from_mat(int m, ggml_bf16_t *src, int ith, int nth) {
|
||||
assert(m <= max_m);
|
||||
assert(ith == 0 && nth == 1);
|
||||
for (int m_begin = 0; m_begin < m; m_begin += M_STEP) {
|
||||
for (int i = 0; i < M_STEP && m_begin + i < m; i++) {
|
||||
float amax = 0.0f;
|
||||
for (int j = 0; j < k; j += 32) {
|
||||
__m512 f0, f1;
|
||||
avx512_32xbf16_to_32xfp32((__m512i *)(src + (m_begin + i) * k + j), &f0, &f1);
|
||||
amax = MAX(amax, _mm512_reduce_max_ps(_mm512_abs_ps(f0)));
|
||||
amax = MAX(amax, _mm512_reduce_max_ps(_mm512_abs_ps(f1)));
|
||||
}
|
||||
d[m_begin + i] = amax / ((1 << 7) - 1);
|
||||
}
|
||||
}
|
||||
int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP;
|
||||
for (int m_begin = 0; m_begin < m; m_begin += M_STEP) {
|
||||
for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) {
|
||||
int k_block_size = std::min(K_BLOCK, k - k_block_begin);
|
||||
for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) {
|
||||
for (int i = 0; i < M_STEP && m_begin + i < m; i++) {
|
||||
__m512 id = _mm512_set1_ps(d[m_begin + i] ? 1.0f / d[m_begin + i] : 0.0f);
|
||||
int8_t *dst = a + k_block_begin * m_block_size + m_begin * k_block_size + k_begin * M_STEP + i * K_STEP;
|
||||
__m512 f0, f1, f2, f3;
|
||||
avx512_32xbf16_to_32xfp32((__m512i *)(src + (m_begin + i) * k + k_block_begin + k_begin), &f0, &f1);
|
||||
avx512_32xbf16_to_32xfp32((__m512i *)(src + (m_begin + i) * k + k_block_begin + k_begin) + 1, &f2, &f3);
|
||||
__m512i i0 = _mm512_cvtps_epi32(_mm512_mul_ps(f0, id));
|
||||
__m512i i1 = _mm512_cvtps_epi32(_mm512_mul_ps(f1, id));
|
||||
__m512i i2 = _mm512_cvtps_epi32(_mm512_mul_ps(f2, id));
|
||||
__m512i i3 = _mm512_cvtps_epi32(_mm512_mul_ps(f3, id));
|
||||
__m128i s0 = _mm512_cvtsepi32_epi8(i0);
|
||||
__m128i s1 = _mm512_cvtsepi32_epi8(i1);
|
||||
__m128i s2 = _mm512_cvtsepi32_epi8(i2);
|
||||
__m128i s3 = _mm512_cvtsepi32_epi8(i3);
|
||||
_mm_storeu_si128((__m128i *)dst, s0);
|
||||
_mm_storeu_si128((__m128i *)(dst + 16), s1);
|
||||
_mm_storeu_si128((__m128i *)(dst + 32), s2);
|
||||
_mm_storeu_si128((__m128i *)(dst + 48), s3);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int8_t *get_submat(int m, int k, int m_begin, int k_begin) {
|
||||
int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP;
|
||||
int k_block_begin = k_begin / K_BLOCK * K_BLOCK;
|
||||
k_begin -= k_block_begin;
|
||||
int k_block_size = std::min(K_BLOCK, k - k_block_begin);
|
||||
return a + k_block_begin * m_block_size + m_begin * k_block_size + k_begin * M_STEP;
|
||||
}
|
||||
|
||||
float *get_scale(int m, int m_begin) { return d + m_begin; }
|
||||
};
|
||||
|
||||
struct BufferB {
|
||||
int8_t *b;
|
||||
float *d;
|
||||
int n, k;
|
||||
|
||||
static size_t required_size(int n, int k) { return n * k * sizeof(int8_t) + n * sizeof(float); }
|
||||
|
||||
BufferB(int n, int k, void *ptr) : n(n), k(k) {
|
||||
assert(reinterpret_cast<intptr_t>(ptr) % 64 == 0);
|
||||
assert(n % N_STEP == 0);
|
||||
assert(k % K_STEP == 0);
|
||||
b = reinterpret_cast<int8_t *>(ptr);
|
||||
d = reinterpret_cast<float *>(b + n * k);
|
||||
}
|
||||
|
||||
void from_mat(ggml_bf16_t *src, int ith, int nth) {
|
||||
auto [n_start, n_end] = split_range_n(n, ith, nth);
|
||||
int n_block_begin = n_start;
|
||||
int n_block_size = n_end - n_block_begin;
|
||||
for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) {
|
||||
for (int i = 0; i < N_STEP; i++) {
|
||||
float amax = 0.0f;
|
||||
for (int j = 0; j < k; j += 32) {
|
||||
__m512 f0, f1;
|
||||
avx512_32xbf16_to_32xfp32((__m512i *)(src + (n_block_begin + n_begin + i) * k + j), &f0, &f1);
|
||||
amax = MAX(amax, _mm512_reduce_max_ps(_mm512_abs_ps(f0)));
|
||||
amax = MAX(amax, _mm512_reduce_max_ps(_mm512_abs_ps(f1)));
|
||||
}
|
||||
d[n_block_begin + n_begin + i] = amax / ((1 << 7) - 1);
|
||||
}
|
||||
}
|
||||
for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) {
|
||||
for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) {
|
||||
int k_block_size = std::min(K_BLOCK, k - k_block_begin);
|
||||
for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) {
|
||||
for (int i = 0; i < N_STEP; i++) {
|
||||
__m512 id = _mm512_set1_ps(d[n_block_begin + n_begin + i] ? 1.0f / d[n_block_begin + n_begin + i] : 0.0f);
|
||||
int8_t *dst = b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size +
|
||||
k_begin * N_STEP + i * K_STEP;
|
||||
__m512 f0, f1, f2, f3;
|
||||
avx512_32xbf16_to_32xfp32((__m512i *)(src + (n_block_begin + n_begin + i) * k + k_block_begin + k_begin),
|
||||
&f0, &f1);
|
||||
avx512_32xbf16_to_32xfp32(
|
||||
(__m512i *)(src + (n_block_begin + n_begin + i) * k + k_block_begin + k_begin) + 1, &f2, &f3);
|
||||
__m512i i0 = _mm512_cvtps_epi32(_mm512_mul_ps(f0, id));
|
||||
__m512i i1 = _mm512_cvtps_epi32(_mm512_mul_ps(f1, id));
|
||||
__m512i i2 = _mm512_cvtps_epi32(_mm512_mul_ps(f2, id));
|
||||
__m512i i3 = _mm512_cvtps_epi32(_mm512_mul_ps(f3, id));
|
||||
__m128i s0 = _mm512_cvtsepi32_epi8(i0);
|
||||
__m128i s1 = _mm512_cvtsepi32_epi8(i1);
|
||||
__m128i s2 = _mm512_cvtsepi32_epi8(i2);
|
||||
__m128i s3 = _mm512_cvtsepi32_epi8(i3);
|
||||
_mm_storeu_si128((__m128i *)dst, s0);
|
||||
_mm_storeu_si128((__m128i *)(dst + 16), s1);
|
||||
_mm_storeu_si128((__m128i *)(dst + 32), s2);
|
||||
_mm_storeu_si128((__m128i *)(dst + 48), s3);
|
||||
}
|
||||
transpose_16x16_32bit((__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size +
|
||||
n_begin * k_block_size + k_begin * N_STEP));
|
||||
transpose_16x16_32bit((__m512i *)(b + n_block_begin * k + k_block_begin * n_block_size +
|
||||
n_begin * k_block_size + k_begin * N_STEP + TILE_N * K_STEP));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int8_t *get_submat(int n, int k, int n_begin, int k_begin) {
|
||||
int n_block_begin = n_begin / N_BLOCK * N_BLOCK;
|
||||
n_begin -= n_block_begin;
|
||||
int n_block_size = std::min(N_BLOCK, n - n_block_begin);
|
||||
int k_block_begin = k_begin / K_BLOCK * K_BLOCK;
|
||||
k_begin -= k_block_begin;
|
||||
int k_block_size = std::min(K_BLOCK, k - k_block_begin);
|
||||
return b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + k_begin * N_STEP;
|
||||
}
|
||||
|
||||
float *get_scale(int n, int n_begin) { return d + n_begin; }
|
||||
};
|
||||
|
||||
struct BufferC {
|
||||
float *c;
|
||||
int max_m, n;
|
||||
|
||||
static size_t required_size(int max_m, int n) { return max_m * n * sizeof(float); }
|
||||
|
||||
BufferC(int max_m, int n, void *ptr) : max_m(max_m), n(n) {
|
||||
assert(reinterpret_cast<intptr_t>(ptr) % 64 == 0);
|
||||
assert(max_m % M_STEP == 0);
|
||||
assert(n % N_STEP == 0);
|
||||
c = reinterpret_cast<float *>(ptr);
|
||||
}
|
||||
|
||||
void to_mat(int m, ggml_bf16_t *dst, int ith, int nth) {
|
||||
assert(m <= max_m);
|
||||
auto [n_start, n_end] = split_range_n(n, ith, nth);
|
||||
int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP;
|
||||
int n_block_begin = n_start;
|
||||
int n_block_size = n_end - n_block_begin;
|
||||
for (int m_begin = 0; m_begin < m; m_begin += M_STEP) {
|
||||
for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) {
|
||||
for (int i = 0; i < M_STEP && m_begin + i < m; i++) {
|
||||
__m512 *x0 =
|
||||
(__m512 *)(c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP + i * N_STEP);
|
||||
__m512 *x1 = (__m512 *)(c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP +
|
||||
i * N_STEP + 16);
|
||||
avx512_32xfp32_to_32xbf16(x0, x1, (__m512i *)(dst + (m_begin + i) * n + n_block_begin + n_begin));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
float *get_submat(int m, int n, int m_begin, int n_begin) {
|
||||
int m_block_size = (m + M_STEP - 1) / M_STEP * M_STEP;
|
||||
int n_block_begin = n_begin / N_BLOCK * N_BLOCK;
|
||||
int n_block_size = std::min(N_BLOCK, n - n_block_begin);
|
||||
n_begin -= n_block_begin;
|
||||
return c + m_block_size * n_block_begin + m_begin * n_block_size + n_begin * M_STEP;
|
||||
}
|
||||
};
|
||||
};
|
||||
|
||||
inline void mat_mul(int m, int n, int k, std::shared_ptr<GemmKernel224BF::BufferA> ba,
|
||||
std::shared_ptr<GemmKernel224BF::BufferB> bb, std::shared_ptr<GemmKernel224BF::BufferC> bc, int ith,
|
||||
int nth, bool use_amx) {
|
||||
using K = GemmKernel224BF;
|
||||
assert(n % K::N_STEP == 0);
|
||||
assert(k % K::K_STEP == 0);
|
||||
|
||||
auto [n_start, n_end] = K::split_range_n(n, ith, nth);
|
||||
|
||||
for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K::K_BLOCK) {
|
||||
for (int m_begin = 0; m_begin < m; m_begin += K::M_STEP) {
|
||||
for (int n_begin = n_start; n_begin < n_end; n_begin += K::N_STEP) {
|
||||
|
||||
float *c = bc->get_submat(m, n, m_begin, n_begin);
|
||||
if (!use_amx) {
|
||||
__m512 *c512 = (__m512 *)c;
|
||||
if (k_block_begin == 0) {
|
||||
for (int m_i = 0; m_i < m && m_i < K::M_STEP; m_i++) {
|
||||
c512[m_i * 2] = _mm512_setzero_ps();
|
||||
c512[m_i * 2 + 1] = _mm512_setzero_ps();
|
||||
}
|
||||
}
|
||||
|
||||
for (int k_begin = 0; k_begin < K::K_BLOCK && k_block_begin + k_begin < k; k_begin += K::K_STEP) {
|
||||
int32_t *a32 = (int32_t *)ba->get_submat(m, k, m_begin, k_block_begin + k_begin);
|
||||
__m512bh *b512 = (__m512bh *)bb->get_submat(n, k, n_begin, k_block_begin + k_begin);
|
||||
for (int m_i = 0; m_i < m && m_i < K::M_STEP; m_i++) {
|
||||
for (int k_i = 0; k_i < 16; k_i++) {
|
||||
__m512bh ma = (__m512bh)_mm512_set1_epi32(a32[m_i * 16 + k_i]);
|
||||
for (int n_i = 0; n_i < 2; n_i++) {
|
||||
c512[m_i * 2 + n_i] = _mm512_dpbf16_ps(c512[m_i * 2 + n_i], ma, b512[n_i * 16 + k_i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} else {
|
||||
if (k_block_begin == 0) {
|
||||
K::clean_c();
|
||||
} else {
|
||||
K::load_c(c, K::N_STEP * sizeof(float));
|
||||
}
|
||||
for (int k_begin = 0; k_begin < K::K_BLOCK && k_block_begin + k_begin < k; k_begin += K::K_STEP) {
|
||||
K::load_a(ba->get_submat(m, k, m_begin, k_block_begin + k_begin), K::K_STEP * sizeof(ggml_bf16_t));
|
||||
K::load_b(bb->get_submat(n, k, n_begin, k_block_begin + k_begin), K::K_STEP * sizeof(ggml_bf16_t));
|
||||
K::run_tile();
|
||||
}
|
||||
K::store_c(c, K::N_STEP * sizeof(float));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
inline __m512i _mm512_dpbssd_epi32(__m512i src, __m512i a, __m512i b) {
|
||||
__m256i a_lo = _mm512_extracti64x4_epi64(a, 0);
|
||||
__m256i a_hi = _mm512_extracti64x4_epi64(a, 1);
|
||||
__m256i b_lo = _mm512_extracti64x4_epi64(b, 0);
|
||||
__m256i b_hi = _mm512_extracti64x4_epi64(b, 1);
|
||||
|
||||
b_lo = _mm256_sign_epi8(b_lo, a_lo);
|
||||
b_hi = _mm256_sign_epi8(b_hi, a_hi);
|
||||
|
||||
b = _mm512_inserti64x4(b, b_lo, 0);
|
||||
b = _mm512_inserti64x4(b, b_hi, 1);
|
||||
|
||||
a = _mm512_abs_epi8(a);
|
||||
|
||||
return _mm512_dpbusd_epi32(src, a, b);
|
||||
}
|
||||
|
||||
inline void mat_mul(int m, int n, int k, std::shared_ptr<GemmKernel224Int8::BufferA> ba,
|
||||
std::shared_ptr<GemmKernel224Int8::BufferB> bb, std::shared_ptr<GemmKernel224Int8::BufferC> bc,
|
||||
int ith, int nth, bool use_amx) {
|
||||
using K = GemmKernel224Int8;
|
||||
assert(n % K::N_STEP == 0);
|
||||
assert(k % K::K_STEP == 0);
|
||||
|
||||
auto [n_start, n_end] = K::split_range_n(n, ith, nth);
|
||||
|
||||
for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K::K_BLOCK) {
|
||||
for (int m_begin = 0; m_begin < m; m_begin += K::M_STEP) {
|
||||
for (int n_begin = n_start; n_begin < n_end; n_begin += K::N_STEP) {
|
||||
float *c = bc->get_submat(m, n, m_begin, n_begin);
|
||||
|
||||
if (!use_amx) {
|
||||
__m512i *c512 = (__m512i *)c;
|
||||
if (k_block_begin == 0) {
|
||||
for (int m_i = 0; m_i < m && m_i < K::M_STEP; m_i++) {
|
||||
c512[m_i * 2] = _mm512_setzero_si512();
|
||||
c512[m_i * 2 + 1] = _mm512_setzero_si512();
|
||||
}
|
||||
}
|
||||
|
||||
for (int k_begin = 0; k_begin < K::K_BLOCK && k_block_begin + k_begin < k; k_begin += K::K_STEP) {
|
||||
static_assert(K::K_STEP * sizeof(int8_t) == sizeof(__m512i));
|
||||
static_assert(K::N_STEP / K::TILE_N == 2, "Must be lke this");
|
||||
|
||||
int32_t *a32 = (int32_t *)ba->get_submat(m, k, m_begin, k_block_begin + k_begin);
|
||||
__m512i *b512 = (__m512i *)bb->get_submat(n, k, n_begin, k_block_begin + k_begin);
|
||||
for (int m_i = 0; m_i < m && m_i < K::M_STEP; m_i++) {
|
||||
for (int k_i = 0; k_i < 16; k_i++) {
|
||||
__m512i ma = _mm512_set1_epi32(a32[m_i * 16 + k_i]);
|
||||
for (int n_i = 0; n_i < 2; n_i++) {
|
||||
c512[m_i * 2 + n_i] = _mm512_dpbssd_epi32(c512[m_i * 2 + n_i], ma, b512[n_i * 16 + k_i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if (k_block_begin == 0) {
|
||||
K::clean_c();
|
||||
} else {
|
||||
K::load_c((int32_t *)c, K::N_STEP * sizeof(int32_t));
|
||||
}
|
||||
for (int k_begin = 0; k_begin < K::K_BLOCK && k_block_begin + k_begin < k; k_begin += K::K_STEP) {
|
||||
K::load_a(ba->get_submat(m, k, m_begin, k_block_begin + k_begin), K::K_STEP * sizeof(int8_t));
|
||||
K::load_b(bb->get_submat(n, k, n_begin, k_block_begin + k_begin), K::K_STEP * sizeof(int8_t));
|
||||
K::run_tile();
|
||||
}
|
||||
K::store_c((int32_t *)c, K::N_STEP * sizeof(int32_t));
|
||||
}
|
||||
|
||||
if (k_block_begin + K::K_BLOCK >= k) {
|
||||
int to = m - m_begin;
|
||||
if (m - m_begin > K::M_STEP) {
|
||||
to = K::M_STEP;
|
||||
}
|
||||
for (int i = 0; i < to; i++) {
|
||||
__m512 as = _mm512_set1_ps(*ba->get_scale(m, m_begin + i));
|
||||
__m512 bs = _mm512_load_ps(bb->get_scale(n, n_begin));
|
||||
__m512i now = _mm512_load_si512((__m512i *)(c + i * K::N_STEP));
|
||||
__m512 result = _mm512_mul_ps(_mm512_mul_ps(as, bs), _mm512_cvtepi32_ps(now));
|
||||
_mm512_store_ps((__m512 *)(c + i * K::N_STEP), result);
|
||||
bs = _mm512_load_ps(bb->get_scale(n, n_begin) + K::TILE_N);
|
||||
now = _mm512_load_si512((__m512i *)(c + i * K::N_STEP + K::TILE_N));
|
||||
result = _mm512_mul_ps(_mm512_mul_ps(as, bs), _mm512_cvtepi32_ps(now));
|
||||
_mm512_store_ps((__m512 *)(c + i * K::N_STEP + K::TILE_N), result);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace amx
|
||||
@@ -0,0 +1,46 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : chenht2022
|
||||
* @Date : 2025-04-25 18:28:12
|
||||
* @Version : 1.0.0
|
||||
* @LastEditors : chenht2022
|
||||
* @LastEditTime : 2025-04-25 18:28:12
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
|
||||
#pragma once
|
||||
#include <cstdint>
|
||||
|
||||
|
||||
template <typename T>
|
||||
T* offset_pointer(T* ptr, std::size_t byte_offset) {
|
||||
return reinterpret_cast<T*>(reinterpret_cast<char*>(ptr) + byte_offset);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
const T* offset_pointer(const T* ptr, std::size_t byte_offset) {
|
||||
return reinterpret_cast<const T*>(reinterpret_cast<const char*>(ptr) + byte_offset);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T* offset_pointer_row_major(T* t, int row, int col, std::size_t ld) {
|
||||
return offset_pointer(t, row * ld) + col;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T* offset_pointer_col_major(T* t, int row, int col, std::size_t ld) {
|
||||
return offset_pointer(t, col * ld) + row;
|
||||
}
|
||||
|
||||
static inline void avx512_copy_32xbf16(__m512i* src, __m512i* dst) {
|
||||
_mm512_storeu_si512(dst, _mm512_loadu_si512(src));
|
||||
}
|
||||
|
||||
static inline void avx512_32xfp32_to_32xbf16(__m512* src0, __m512* src1, __m512i* dst) {
|
||||
_mm512_storeu_si512(dst, __m512i(_mm512_cvtne2ps_pbh(*src1, *src0)));
|
||||
}
|
||||
|
||||
static inline void avx512_32xbf16_to_32xfp32(__m512i* src, __m512* dst0, __m512* dst1) {
|
||||
_mm512_storeu_ps(dst0, _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256((const __m256i *)(src))), 16)));
|
||||
_mm512_storeu_ps(dst1, _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256((const __m256i *)(src) + 1)), 16)));
|
||||
}
|
||||
@@ -0,0 +1,422 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : chenht2022
|
||||
* @Date : 2025-04-25 18:28:12
|
||||
* @Version : 1.0.0
|
||||
* @LastEditors : chenht2022
|
||||
* @LastEditTime : 2025-04-25 18:28:12
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
#ifndef CPUINFER_OPERATOR_AMX_MOE_H
|
||||
#define CPUINFER_OPERATOR_AMX_MOE_H
|
||||
|
||||
#include <cmath>
|
||||
#include <cstdio>
|
||||
#include <functional>
|
||||
#include <mutex>
|
||||
#include <vector>
|
||||
|
||||
#include "../../cpu_backend/backend.h"
|
||||
#include "../../cpu_backend/shared_mem_buffer.h"
|
||||
#include "llama.cpp/ggml-impl.h"
|
||||
#include "llama.cpp/ggml-quants.h"
|
||||
#include "llama.cpp/ggml.h"
|
||||
#include "llamafile/sgemm.h"
|
||||
|
||||
#include "la/amx.hpp"
|
||||
|
||||
#ifdef USE_NUMA
|
||||
#include <numa.h>
|
||||
#include <numaif.h>
|
||||
void *numa_alloc_aligned(size_t size, int node, size_t alignment) {
|
||||
void *ptr = numa_alloc_onnode(size, node);
|
||||
assert(reinterpret_cast<intptr_t>(ptr) % 64 == 0);
|
||||
return ptr;
|
||||
}
|
||||
#endif
|
||||
|
||||
static inline __m512 exp_avx512(__m512 x) {
|
||||
const __m512 log2e = _mm512_set1_ps(1.44269504089f);
|
||||
const __m512 c1 = _mm512_set1_ps(0.69314718056f);
|
||||
|
||||
__m512 y = _mm512_mul_ps(x, log2e);
|
||||
__m512i int_part = _mm512_cvtps_epi32(y);
|
||||
__m512 frac_part = _mm512_sub_ps(y, _mm512_cvtepi32_ps(int_part));
|
||||
|
||||
const __m512 poly_1 = _mm512_set1_ps(0.9999999995f);
|
||||
const __m512 poly_2 = _mm512_set1_ps(0.6931471805f);
|
||||
const __m512 poly_3 = _mm512_set1_ps(0.2402265069f);
|
||||
const __m512 poly_4 = _mm512_set1_ps(0.0555041087f);
|
||||
const __m512 poly_5 = _mm512_set1_ps(0.0096181291f);
|
||||
const __m512 poly_6 = _mm512_set1_ps(0.0013333558f);
|
||||
|
||||
__m512 frac_exp = _mm512_fmadd_ps(
|
||||
frac_part, poly_6,
|
||||
_mm512_fmadd_ps(frac_part, poly_5,
|
||||
_mm512_fmadd_ps(frac_part, poly_4,
|
||||
_mm512_fmadd_ps(frac_part, poly_3, _mm512_fmadd_ps(frac_part, poly_2, poly_1)))));
|
||||
|
||||
__m512 two_pow_i = _mm512_scalef_ps(_mm512_set1_ps(1.0f), _mm512_cvtepi32_ps(int_part));
|
||||
return _mm512_mul_ps(two_pow_i, frac_exp);
|
||||
}
|
||||
|
||||
static inline __m512 act_fn(__m512 gate_val, __m512 up_val) {
|
||||
__m512 neg_gate_val = _mm512_sub_ps(_mm512_setzero_ps(), gate_val);
|
||||
__m512 exp_neg_gate = exp_avx512(neg_gate_val);
|
||||
__m512 denom = _mm512_add_ps(_mm512_set1_ps(1.0f), exp_neg_gate);
|
||||
__m512 act_val = _mm512_div_ps(gate_val, denom);
|
||||
|
||||
return _mm512_mul_ps(act_val, up_val);
|
||||
}
|
||||
|
||||
static inline __m512 relu_act_fn(__m512 gate_val, __m512 up_val) {
|
||||
__m512 zero_vec = _mm512_setzero_ps();
|
||||
__m512 act_val = _mm512_max_ps(zero_vec, gate_val);
|
||||
return _mm512_mul_ps(act_val, up_val);
|
||||
}
|
||||
|
||||
struct AMX_MOEConfig {
|
||||
int expert_num;
|
||||
int routed_expert_num;
|
||||
int hidden_size;
|
||||
int intermediate_size;
|
||||
int max_len;
|
||||
bool use_silu;
|
||||
void *gate_proj;
|
||||
void *up_proj;
|
||||
void *down_proj;
|
||||
|
||||
AMX_MOEConfig() {}
|
||||
|
||||
AMX_MOEConfig(int expert_num, int routed_expert_num, int hidden_size, int intermediate_size, int max_len, bool use_silu,
|
||||
void *gate_proj, void *up_proj, void *down_proj)
|
||||
: expert_num(expert_num), routed_expert_num(routed_expert_num), hidden_size(hidden_size),
|
||||
intermediate_size(intermediate_size), max_len(max_len), use_silu(use_silu), gate_proj(gate_proj), up_proj(up_proj),
|
||||
down_proj(down_proj) {}
|
||||
};
|
||||
|
||||
template <class T> class AMX_MOE {
|
||||
private:
|
||||
AMX_MOEConfig config_;
|
||||
void *gate_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)]
|
||||
void *up_proj_; // [expert_num * intermediate_size * hidden_size ( /32 if quantized)]
|
||||
void *down_proj_; // [expert_num * hidden_size * intermediate_size ( /32 if quantized)]
|
||||
|
||||
ggml_bf16_t *m_local_input_; // [routed_expert_num * max_len * hidden_size]
|
||||
ggml_bf16_t *m_local_gate_output_; // [routed_expert_num * max_len * intermediate_size]
|
||||
ggml_bf16_t *m_local_up_output_; // [routed_expert_num * max_len * intermediate_size]
|
||||
ggml_bf16_t *m_local_down_output_; // [routed_expert_num * max_len * hidden_size]
|
||||
|
||||
std::vector<std::vector<int>> m_local_pos_; // [max_len, routed_expert_num]
|
||||
std::vector<int> m_local_num_; // [expert_num]
|
||||
std::vector<int> m_expert_id_map_; // [expert_num]
|
||||
std::vector<ggml_bf16_t *> m_local_input_ptr_; // [expert_num]
|
||||
std::vector<ggml_bf16_t *> m_local_gate_output_ptr_; // [expert_num]
|
||||
std::vector<ggml_bf16_t *> m_local_up_output_ptr_; // [expert_num]
|
||||
std::vector<ggml_bf16_t *> m_local_down_output_ptr_; // [expert_num]
|
||||
|
||||
std::vector<std::shared_ptr<typename T::BufferA>> gate_up_ba_;
|
||||
std::vector<std::shared_ptr<typename T::BufferC>> gate_bc_;
|
||||
std::vector<std::shared_ptr<typename T::BufferC>> up_bc_;
|
||||
std::vector<std::shared_ptr<typename T::BufferA>> down_ba_;
|
||||
std::vector<std::shared_ptr<typename T::BufferC>> down_bc_;
|
||||
|
||||
#ifdef USE_NUMA
|
||||
std::vector<std::vector<std::shared_ptr<typename T::BufferB>>> gate_bb_numa_;
|
||||
std::vector<std::vector<std::shared_ptr<typename T::BufferB>>> up_bb_numa_;
|
||||
std::vector<std::vector<std::shared_ptr<typename T::BufferB>>> down_bb_numa_;
|
||||
#else
|
||||
std::vector<std::shared_ptr<typename T::BufferB>> gate_bb_;
|
||||
std::vector<std::shared_ptr<typename T::BufferB>> up_bb_;
|
||||
std::vector<std::shared_ptr<typename T::BufferB>> down_bb_;
|
||||
#endif
|
||||
|
||||
public:
|
||||
AMX_MOE(AMX_MOEConfig config) {
|
||||
config_ = config;
|
||||
gate_proj_ = config_.gate_proj;
|
||||
up_proj_ = config_.up_proj;
|
||||
down_proj_ = config_.down_proj;
|
||||
|
||||
std::vector<std::pair<void **, uint64_t>> m_mem_requests;
|
||||
m_mem_requests.push_back({(void **)&m_local_input_,
|
||||
sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size});
|
||||
m_mem_requests.push_back({(void **)&m_local_gate_output_, sizeof(ggml_bf16_t) * config_.routed_expert_num *
|
||||
config_.max_len * config_.intermediate_size});
|
||||
m_mem_requests.push_back({(void **)&m_local_up_output_, sizeof(ggml_bf16_t) * config_.routed_expert_num *
|
||||
config_.max_len * config_.intermediate_size});
|
||||
m_mem_requests.push_back({(void **)&m_local_down_output_,
|
||||
sizeof(ggml_bf16_t) * config_.routed_expert_num * config_.max_len * config_.hidden_size});
|
||||
std::vector<void *> gate_up_ba_ptr(config_.expert_num);
|
||||
std::vector<void *> gate_bc_ptr(config_.expert_num);
|
||||
std::vector<void *> up_bc_ptr(config_.expert_num);
|
||||
std::vector<void *> down_ba_ptr(config_.expert_num);
|
||||
std::vector<void *> down_bc_ptr(config_.expert_num);
|
||||
for (int i = 0; i < config_.expert_num; i++) {
|
||||
m_mem_requests.push_back(
|
||||
{(void **)&gate_up_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.hidden_size)});
|
||||
m_mem_requests.push_back(
|
||||
{(void **)&gate_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.intermediate_size)});
|
||||
m_mem_requests.push_back(
|
||||
{(void **)&up_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.intermediate_size)});
|
||||
m_mem_requests.push_back(
|
||||
{(void **)&down_ba_ptr[i], T::BufferA::required_size(config_.max_len, config_.intermediate_size)});
|
||||
m_mem_requests.push_back(
|
||||
{(void **)&down_bc_ptr[i], T::BufferC::required_size(config_.max_len, config_.hidden_size)});
|
||||
}
|
||||
shared_mem_buffer.alloc(this, m_mem_requests);
|
||||
|
||||
m_local_pos_.resize(config_.max_len);
|
||||
for (int i = 0; i < config_.max_len; i++) {
|
||||
m_local_pos_[i].resize(config_.routed_expert_num);
|
||||
}
|
||||
m_expert_id_map_.resize(config_.expert_num);
|
||||
m_local_num_.resize(config_.expert_num);
|
||||
m_local_input_ptr_.resize(config_.expert_num);
|
||||
m_local_gate_output_ptr_.resize(config_.expert_num);
|
||||
m_local_up_output_ptr_.resize(config_.expert_num);
|
||||
m_local_down_output_ptr_.resize(config_.expert_num);
|
||||
|
||||
for (uint64_t i = 0; i < config_.expert_num; i++) {
|
||||
gate_up_ba_.push_back(
|
||||
std::make_shared<typename T::BufferA>(config_.max_len, config_.hidden_size, gate_up_ba_ptr[i]));
|
||||
gate_bc_.push_back(
|
||||
std::make_shared<typename T::BufferC>(config_.max_len, config_.intermediate_size, gate_bc_ptr[i]));
|
||||
up_bc_.push_back(std::make_shared<typename T::BufferC>(config_.max_len, config_.intermediate_size, up_bc_ptr[i]));
|
||||
down_ba_.push_back(
|
||||
std::make_shared<typename T::BufferA>(config_.max_len, config_.intermediate_size, down_ba_ptr[i]));
|
||||
down_bc_.push_back(std::make_shared<typename T::BufferC>(config_.max_len, config_.hidden_size, down_bc_ptr[i]));
|
||||
|
||||
#ifdef USE_NUMA
|
||||
int numa_nodes = numa_num_configured_nodes();
|
||||
gate_bb_numa_.resize(numa_nodes);
|
||||
up_bb_numa_.resize(numa_nodes);
|
||||
down_bb_numa_.resize(numa_nodes);
|
||||
for (int j = 0; j < numa_nodes; j++) {
|
||||
void *gate_bb_ptr =
|
||||
numa_alloc_aligned(T::BufferB::required_size(config_.intermediate_size, config_.hidden_size), j, 64);
|
||||
gate_bb_numa_[j].push_back(
|
||||
std::make_shared<typename T::BufferB>(config_.intermediate_size, config_.hidden_size, gate_bb_ptr));
|
||||
void *up_bb_ptr =
|
||||
numa_alloc_aligned(T::BufferB::required_size(config_.intermediate_size, config_.hidden_size), j, 64);
|
||||
up_bb_numa_[j].push_back(
|
||||
std::make_shared<typename T::BufferB>(config_.intermediate_size, config_.hidden_size, up_bb_ptr));
|
||||
void *down_bb_ptr =
|
||||
numa_alloc_aligned(T::BufferB::required_size(config_.hidden_size, config_.intermediate_size), j, 64);
|
||||
down_bb_numa_[j].push_back(
|
||||
std::make_shared<typename T::BufferB>(config_.hidden_size, config_.intermediate_size, down_bb_ptr));
|
||||
}
|
||||
#else
|
||||
void *gate_bb_ptr =
|
||||
std::aligned_alloc(64, T::BufferB::required_size(config_.intermediate_size, config_.hidden_size));
|
||||
gate_bb_.push_back(
|
||||
std::make_shared<typename T::BufferB>(config_.intermediate_size, config_.hidden_size, gate_bb_ptr));
|
||||
|
||||
void *up_bb_ptr =
|
||||
std::aligned_alloc(64, T::BufferB::required_size(config_.intermediate_size, config_.hidden_size));
|
||||
up_bb_.push_back(
|
||||
std::make_shared<typename T::BufferB>(config_.intermediate_size, config_.hidden_size, up_bb_ptr));
|
||||
|
||||
void *down_bb_ptr =
|
||||
std::aligned_alloc(64, T::BufferB::required_size(config_.hidden_size, config_.intermediate_size));
|
||||
down_bb_.push_back(
|
||||
std::make_shared<typename T::BufferB>(config_.hidden_size, config_.intermediate_size, down_bb_ptr));
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
~AMX_MOE() { shared_mem_buffer.dealloc(this); }
|
||||
|
||||
void load_weights(Backend *backend) {
|
||||
int nth = T::recommended_nth(config_.intermediate_size);
|
||||
backend->do_work_stealing_job(
|
||||
nth * config_.expert_num, nullptr,
|
||||
[&](int task_id) {
|
||||
uint64_t expert_idx = task_id / nth;
|
||||
int ith = task_id % nth;
|
||||
#ifdef USE_NUMA
|
||||
int numa_nodes = numa_num_configured_nodes();
|
||||
for (int j = 0; j < numa_nodes; j++) {
|
||||
gate_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.gate_proj +
|
||||
expert_idx * config_.intermediate_size * config_.hidden_size,
|
||||
ith, nth);
|
||||
up_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.up_proj +
|
||||
expert_idx * config_.intermediate_size * config_.hidden_size,
|
||||
ith, nth);
|
||||
}
|
||||
#else
|
||||
gate_bb_[expert_idx]->from_mat((ggml_bf16_t *)config_.gate_proj +
|
||||
expert_idx * config_.intermediate_size * config_.hidden_size,
|
||||
ith, nth);
|
||||
up_bb_[expert_idx]->from_mat(
|
||||
(ggml_bf16_t *)config_.up_proj + expert_idx * config_.intermediate_size * config_.hidden_size, ith, nth);
|
||||
#endif
|
||||
},
|
||||
nullptr);
|
||||
nth = T::recommended_nth(config_.hidden_size);
|
||||
backend->do_work_stealing_job(
|
||||
nth * config_.expert_num, nullptr,
|
||||
[&](int task_id) {
|
||||
uint64_t expert_idx = task_id / nth;
|
||||
int ith = task_id % nth;
|
||||
#ifdef USE_NUMA
|
||||
int numa_nodes = numa_num_configured_nodes();
|
||||
for (int j = 0; j < numa_nodes; j++) {
|
||||
down_bb_numa_[j][expert_idx]->from_mat((ggml_bf16_t *)config_.down_proj +
|
||||
expert_idx * config_.hidden_size * config_.intermediate_size,
|
||||
ith, nth);
|
||||
}
|
||||
#else
|
||||
down_bb_[expert_idx]->from_mat((ggml_bf16_t *)config_.down_proj +
|
||||
expert_idx * config_.hidden_size * config_.intermediate_size,
|
||||
ith, nth);
|
||||
#endif
|
||||
},
|
||||
nullptr);
|
||||
}
|
||||
|
||||
void warm_up(Backend *backend) {}
|
||||
|
||||
void forward(int qlen, int k, const uint64_t *expert_ids, const float *weights, const void *input, void *output,
|
||||
int *batch_size_tensor, Backend *backend) {
|
||||
qlen = batch_size_tensor[0];
|
||||
bool use_amx = (qlen > 4 * config_.expert_num / config_.routed_expert_num);
|
||||
int activated_expert = 0;
|
||||
for (int i = 0; i < config_.expert_num; i++) {
|
||||
m_local_num_[i] = 0;
|
||||
}
|
||||
for (int i = 0; i < qlen; i++) {
|
||||
for (int j = 0; j < k; j++) {
|
||||
m_local_pos_[i][j] = m_local_num_[expert_ids[i * k + j]]++;
|
||||
}
|
||||
}
|
||||
for (int i = 0; i < config_.expert_num; i++) {
|
||||
if (m_local_num_[i] > 0) {
|
||||
m_expert_id_map_[activated_expert] = i;
|
||||
activated_expert++;
|
||||
}
|
||||
}
|
||||
uint64_t offset = 0;
|
||||
for (int i = 0; i < config_.expert_num; i++) {
|
||||
m_local_input_ptr_[i] = m_local_input_ + offset * config_.hidden_size;
|
||||
m_local_gate_output_ptr_[i] = m_local_gate_output_ + offset * config_.intermediate_size;
|
||||
m_local_up_output_ptr_[i] = m_local_up_output_ + offset * config_.intermediate_size;
|
||||
m_local_down_output_ptr_[i] = m_local_down_output_ + offset * config_.hidden_size;
|
||||
offset += m_local_num_[i];
|
||||
}
|
||||
backend->do_work_stealing_job(
|
||||
qlen, nullptr,
|
||||
[&](int i) {
|
||||
for (int j = 0; j < k; j++) {
|
||||
memcpy(m_local_input_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size,
|
||||
(ggml_bf16_t *)input + i * config_.hidden_size, sizeof(ggml_bf16_t) * config_.hidden_size);
|
||||
}
|
||||
},
|
||||
nullptr);
|
||||
backend->do_work_stealing_job(
|
||||
activated_expert, nullptr,
|
||||
[&](int task_id) {
|
||||
int expert_idx = m_expert_id_map_[task_id];
|
||||
|
||||
gate_up_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_input_ptr_[expert_idx], 0, 1);
|
||||
},
|
||||
nullptr);
|
||||
int nth = T::recommended_nth(config_.intermediate_size);
|
||||
backend->do_work_stealing_job(
|
||||
nth * activated_expert, [&](int _) { T::config(); },
|
||||
[&](int task_id) {
|
||||
int expert_idx = m_expert_id_map_[task_id / nth];
|
||||
int ith = task_id % nth;
|
||||
#ifdef USE_NUMA
|
||||
amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size,
|
||||
gate_up_ba_[expert_idx], gate_bb_numa_[Backend::numa_node][expert_idx], gate_bc_[expert_idx],
|
||||
ith, nth, use_amx);
|
||||
amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size,
|
||||
gate_up_ba_[expert_idx], up_bb_numa_[Backend::numa_node][expert_idx], up_bc_[expert_idx], ith,
|
||||
nth, use_amx);
|
||||
#else
|
||||
amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size,
|
||||
gate_up_ba_[expert_idx], gate_bb_[expert_idx], gate_bc_[expert_idx], ith, nth, use_amx);
|
||||
amx::mat_mul(m_local_num_[expert_idx], config_.intermediate_size, config_.hidden_size,
|
||||
gate_up_ba_[expert_idx], up_bb_[expert_idx], up_bc_[expert_idx], ith, nth, use_amx);
|
||||
#endif
|
||||
gate_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_gate_output_ptr_[expert_idx], ith, nth);
|
||||
up_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_up_output_ptr_[expert_idx], ith, nth);
|
||||
auto [n_start, n_end] = T::split_range_n(config_.intermediate_size, ith, nth);
|
||||
if (config_.use_silu) {
|
||||
for (int i = 0; i < m_local_num_[expert_idx]; i++) {
|
||||
ggml_bf16_t *gate_output_ptr = &m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size];
|
||||
ggml_bf16_t *up_output_ptr = &m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size];
|
||||
for (int j = n_start; j < n_end; j += 32) {
|
||||
__m512 gate_val0, gate_val1, up_val0, up_val1;
|
||||
avx512_32xbf16_to_32xfp32((__m512i *)(gate_output_ptr + j), &gate_val0, &gate_val1);
|
||||
avx512_32xbf16_to_32xfp32((__m512i *)(up_output_ptr + j), &up_val0, &up_val1);
|
||||
__m512 result0 = act_fn(gate_val0, up_val0);
|
||||
__m512 result1 = act_fn(gate_val1, up_val1);
|
||||
avx512_32xfp32_to_32xbf16(&result0, &result1, (__m512i *)(gate_output_ptr + j));
|
||||
}
|
||||
}
|
||||
}
|
||||
else {
|
||||
for (int i = 0; i < m_local_num_[expert_idx]; i++) {
|
||||
ggml_bf16_t *gate_output_ptr = &m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size];
|
||||
ggml_bf16_t *up_output_ptr = &m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size];
|
||||
for (int j = n_start; j < n_end; j += 32) {
|
||||
__m512 gate_val0, gate_val1, up_val0, up_val1;
|
||||
avx512_32xbf16_to_32xfp32((__m512i *)(gate_output_ptr + j), &gate_val0, &gate_val1);
|
||||
avx512_32xbf16_to_32xfp32((__m512i *)(up_output_ptr + j), &up_val0, &up_val1);
|
||||
__m512 result0 = relu_act_fn(gate_val0, up_val0);
|
||||
__m512 result1 = relu_act_fn(gate_val1, up_val1);
|
||||
avx512_32xfp32_to_32xbf16(&result0, &result1, (__m512i *)(gate_output_ptr + j));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
},
|
||||
nullptr);
|
||||
backend->do_work_stealing_job(
|
||||
activated_expert, nullptr,
|
||||
[&](int task_id) {
|
||||
int expert_idx = m_expert_id_map_[task_id];
|
||||
down_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_gate_output_ptr_[expert_idx], 0, 1);
|
||||
},
|
||||
nullptr);
|
||||
nth = T::recommended_nth(config_.hidden_size);
|
||||
backend->do_work_stealing_job(
|
||||
nth * activated_expert, [&](int _) { T::config(); },
|
||||
[&](int task_id) {
|
||||
int expert_idx = m_expert_id_map_[task_id / nth];
|
||||
int ith = task_id % nth;
|
||||
#ifdef USE_NUMA
|
||||
amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, down_ba_[expert_idx],
|
||||
down_bb_numa_[Backend::numa_node][expert_idx], down_bc_[expert_idx], ith, nth, use_amx);
|
||||
#else
|
||||
amx::mat_mul(m_local_num_[expert_idx], config_.hidden_size, config_.intermediate_size, down_ba_[expert_idx],
|
||||
down_bb_[expert_idx], down_bc_[expert_idx], ith, nth, use_amx);
|
||||
#endif
|
||||
down_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_down_output_ptr_[expert_idx], ith, nth);
|
||||
},
|
||||
nullptr);
|
||||
backend->do_work_stealing_job(
|
||||
qlen, nullptr,
|
||||
[&](int i) {
|
||||
for (int e = 0; e < config_.hidden_size; e += 32) {
|
||||
__m512 x0 = _mm512_setzero_ps();
|
||||
__m512 x1 = _mm512_setzero_ps();
|
||||
for (int j = 0; j < k; j++) {
|
||||
__m512 weight = _mm512_set1_ps(weights[i * k + j]);
|
||||
__m512 down_output0, down_output1;
|
||||
avx512_32xbf16_to_32xfp32((__m512i *)(m_local_down_output_ptr_[expert_ids[i * k + j]] +
|
||||
m_local_pos_[i][j] * config_.hidden_size + e),
|
||||
&down_output0, &down_output1);
|
||||
x0 = _mm512_fmadd_ps(down_output0, weight, x0);
|
||||
x1 = _mm512_fmadd_ps(down_output1, weight, x1);
|
||||
}
|
||||
avx512_32xfp32_to_32xbf16(&x0, &x1, (__m512i *)((ggml_bf16_t *)output + i * config_.hidden_size + e));
|
||||
}
|
||||
},
|
||||
nullptr);
|
||||
}
|
||||
};
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,727 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : Jianwei Dong
|
||||
* @Date : 2024-08-26 22:47:06
|
||||
* @Version : 1.0.0
|
||||
* @LastEditors : Jianwei Dong
|
||||
* @LastEditTime : 2024-08-26 22:47:06
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
|
||||
#ifndef CPUINFER_OPERATOR_KVCACHE_H
|
||||
#define CPUINFER_OPERATOR_KVCACHE_H
|
||||
|
||||
#include <algorithm>
|
||||
#include <atomic>
|
||||
#include <cassert>
|
||||
#include <condition_variable>
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
#include <cstring>
|
||||
#include <fstream>
|
||||
#include <functional>
|
||||
#include <future>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <queue>
|
||||
#include <random>
|
||||
#include <stdexcept>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
|
||||
#include "../../cpu_backend/backend.h"
|
||||
#include "llama.cpp/ggml-common.h"
|
||||
#include "llama.cpp/ggml-impl.h"
|
||||
#include "llama.cpp/ggml-quants.h"
|
||||
#include "llama.cpp/ggml.h"
|
||||
#include "llamafile/sgemm.h"
|
||||
|
||||
#define CHUNK_SIZE 32
|
||||
|
||||
/**
|
||||
* @brief Converts a ggml_type enum value to its corresponding string
|
||||
* representation.
|
||||
*
|
||||
* This function provides a human-readable string representation for a given
|
||||
* ggml_type enum value. The string can be used for logging, debugging, or
|
||||
* displaying information in a user interface.
|
||||
*
|
||||
* @param type The ggml_type enum value to convert.
|
||||
* @return A string representation of the enum value.
|
||||
*/
|
||||
std::string ggml_type_to_string(ggml_type type);
|
||||
|
||||
/**
|
||||
* @enum AnchorType
|
||||
* @brief Defines the types of anchors used in attention mechanisms.
|
||||
*
|
||||
* This enum specifies different types of anchors that can be used in attention
|
||||
* mechanisms, such as fixed anchors, dynamic anchors, or special anchors like
|
||||
* QUEST, BLOCK_MEAN, or BLOCK_MAX.
|
||||
*/
|
||||
enum AnchorType {
|
||||
FIXED_ANCHOR, /**< A fixed anchor that does not change. */
|
||||
DYNAMIC, /**< A dynamic anchor that can change over time. */
|
||||
QUEST, /**< A special anchor type used for QUEST (Query and Embedding Space
|
||||
Transformation). */
|
||||
BLOCK_MEAN, /**< An anchor based on the mean of a block of data. */
|
||||
BLOCK_MAX /**< An anchor based on the maximum value within a block of data.
|
||||
*/
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Converts an AnchorType enum value to its corresponding string
|
||||
* representation.
|
||||
*
|
||||
* This function provides a human-readable string representation for a given
|
||||
* AnchorType enum value. The string can be used for logging, debugging, or
|
||||
* displaying information in a user interface.
|
||||
*
|
||||
* @param anchor_type The AnchorType enum value to convert.
|
||||
* @return A string representation of the enum value.
|
||||
*/
|
||||
std::string AnchorTypeToString(AnchorType anchor_type);
|
||||
|
||||
/**
|
||||
* @enum RetrievalType
|
||||
* @brief Defines the types of retrieval strategies in attention mechanisms.
|
||||
*
|
||||
* This enum specifies different retrieval strategies that can be used in
|
||||
* attention mechanisms, such as layer-level retrieval, key-value head-level
|
||||
* retrieval, or query head-level retrieval.
|
||||
*/
|
||||
enum RetrievalType {
|
||||
LAYER, /**< Retrieval at the layer level. */
|
||||
KVHEAD, /**< Retrieval at the key-value head level. */
|
||||
QHEAD /**< Retrieval at the query head level. */
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Converts a RetrievalType enum value to its corresponding string
|
||||
* representation.
|
||||
*
|
||||
* This function provides a human-readable string representation for a given
|
||||
* RetrievalType enum value. The string can be used for logging, debugging, or
|
||||
* displaying information in a user interface.
|
||||
*
|
||||
* @param retrieval_type The RetrievalType enum value to convert.
|
||||
* @return A string representation of the enum value.
|
||||
*/
|
||||
std::string RetrievalTypeToString(RetrievalType retrieval_type);
|
||||
|
||||
/**
|
||||
* @struct KVCacheConfig
|
||||
* @brief Configuration structure for Key-Value (KV) Cache.
|
||||
*
|
||||
* This structure holds configuration parameters for setting up and managing
|
||||
* a Key-Value (KV) Cache used in various attention mechanisms. It includes
|
||||
* parameters such as the number of layers, the number of heads, the dimension
|
||||
* of each head, block length, anchor information, and memory-related settings.
|
||||
*/
|
||||
struct KVCacheConfig {
|
||||
int layer_num; /**< Number of layers in the model. */
|
||||
int kv_head_num; /**< Number of heads in the KV Cache. */
|
||||
int q_head_num; /**< Number of heads in the query. */
|
||||
int head_dim; /**< Dimension of each head. */
|
||||
int block_len; /**< Length of each block in the cache. */
|
||||
int anchor_num; /**< Number of anchors used in attention. */
|
||||
|
||||
ggml_type kv_type; /**< Data type of the KV Cache (e.g., fp16, q8_0). */
|
||||
|
||||
// Controls the pre-allocated memory size
|
||||
int max_block_num; /**< Maximum number of blocks that can be allocated. */
|
||||
int max_batch_size; /**< Maximum batch size that can be processed. */
|
||||
int max_thread_num; /**< Maximum number of threads that can be used. */
|
||||
|
||||
AnchorType
|
||||
anchor_type; /**< Type of anchors used in the attention mechanism. */
|
||||
RetrievalType
|
||||
retrieval_type; /**< Type of retrieval strategy used in the cache. */
|
||||
|
||||
int layer_step; /**< Step size between layers. */
|
||||
int token_step; /**< Step size between tokens. */
|
||||
int layer_offset; /**< Offset value for layers. */
|
||||
|
||||
/**
|
||||
* @brief Default constructor for KVCacheConfig.
|
||||
*
|
||||
* Initializes the configuration with default values. This constructor
|
||||
* does not initialize any member variables explicitly.
|
||||
*/
|
||||
KVCacheConfig() = default;
|
||||
|
||||
/**
|
||||
* @brief Parameterized constructor for KVCacheConfig.
|
||||
*
|
||||
* This constructor initializes the configuration with specific values
|
||||
* for all member variables.
|
||||
*
|
||||
* @param layer_num The number of layers in the model.
|
||||
* @param kv_head_num The number of heads in the KV Cache.
|
||||
* @param q_head_num The number of heads in the query.
|
||||
* @param head_dim The dimension of each head.
|
||||
* @param block_len The length of each block in the cache.
|
||||
* @param anchor_num The number of anchors used in attention.
|
||||
* @param anchor_type The type of anchors used in the attention mechanism.
|
||||
* @param kv_type The data type of the KV Cache (e.g., fp16, q8_0).
|
||||
* @param retrieval_type The type of retrieval strategy used in the cache.
|
||||
* @param layer_step The step size between layers.
|
||||
* @param token_step The step size between tokens.
|
||||
* @param layer_offset The offset value for layers.
|
||||
* @param max_block_num The maximum number of blocks that can be allocated.
|
||||
* @param max_batch_size The maximum batch size that can be processed.
|
||||
* @param max_thread_num The maximum number of threads that can be used.
|
||||
*/
|
||||
KVCacheConfig(int layer_num, int kv_head_num, int q_head_num, int head_dim,
|
||||
int block_len, int anchor_num, AnchorType anchor_type,
|
||||
ggml_type kv_type, RetrievalType retrieval_type,
|
||||
int layer_step, int token_step, int layer_offset,
|
||||
int max_block_num, int max_batch_size, int max_thread_num);
|
||||
};
|
||||
|
||||
/**
|
||||
* @class KVCache
|
||||
* @brief Manages the Key-Value (KV) Cache used in attention mechanisms.
|
||||
*
|
||||
* The KVCache class provides functionality for managing the Key-Value Cache,
|
||||
* including resizing the cache, retrieving configuration parameters, and
|
||||
* updating internal states. This class is typically used in transformer models
|
||||
* to store and manage past key and value states for efficient attention
|
||||
* computations.
|
||||
*/
|
||||
class KVCache {
|
||||
public:
|
||||
/**
|
||||
* @brief Constructs a KVCache object with the given configuration.
|
||||
*
|
||||
* Initializes the KVCache with the specified configuration parameters,
|
||||
* such as the number of layers, heads, head dimensions, and other
|
||||
* relevant settings.
|
||||
*
|
||||
* @param config The configuration object containing initialization
|
||||
* parameters.
|
||||
*/
|
||||
KVCache(KVCacheConfig config);
|
||||
|
||||
/**
|
||||
* @brief Resizes the number of threads used by the cache.
|
||||
*
|
||||
* This function adjusts the number of threads that the cache can utilize.
|
||||
* It allows dynamic reconfiguration of the parallel processing capabilities
|
||||
* based on the current workload or system resources.
|
||||
*
|
||||
* @param thread_num The new number of threads to use.
|
||||
*/
|
||||
void ThreadResize(int thread_num);
|
||||
|
||||
/**
|
||||
* @brief Resizes the batch size managed by the cache.
|
||||
*
|
||||
* This function adjusts the batch size that the cache can handle. It
|
||||
* is useful when the input batch size changes dynamically, allowing
|
||||
* the cache to be reconfigured accordingly.
|
||||
*
|
||||
* @param batch_size The new batch size.
|
||||
*/
|
||||
void BatchResize(int batch_size);
|
||||
|
||||
/**
|
||||
* @brief Resizes the number of blocks managed by the cache.
|
||||
*
|
||||
* This function adjusts the number of blocks that the cache can manage.
|
||||
* It allows dynamic reconfiguration of the block structure based on the
|
||||
* current sequence length or other factors.
|
||||
*
|
||||
* @param block_num The new number of blocks.
|
||||
*/
|
||||
void BlockResize(int block_num);
|
||||
|
||||
/**
|
||||
* @brief Gets the number of layers in the cache.
|
||||
*
|
||||
* @return The number of layers configured in the cache.
|
||||
*/
|
||||
int get_layer_num() { return config_.layer_num; }
|
||||
|
||||
/**
|
||||
* @brief Gets the number of KV heads in the cache.
|
||||
*
|
||||
* @return The number of KV heads configured in the cache.
|
||||
*/
|
||||
int get_kv_head_num() { return config_.kv_head_num; }
|
||||
|
||||
/**
|
||||
* @brief Gets the number of query heads in the cache.
|
||||
*
|
||||
* @return The number of query heads configured in the cache.
|
||||
*/
|
||||
int get_q_head_num() { return config_.q_head_num; }
|
||||
|
||||
/**
|
||||
* @brief Gets the dimension of each head in the cache.
|
||||
*
|
||||
* @return The dimension of each head.
|
||||
*/
|
||||
int get_head_dim() { return config_.head_dim; }
|
||||
|
||||
/**
|
||||
* @brief Gets the length of each block in the cache.
|
||||
*
|
||||
* @return The length of each block.
|
||||
*/
|
||||
int get_block_len() { return config_.block_len; }
|
||||
|
||||
/**
|
||||
* @brief Gets the number of blocks for a specific layer.
|
||||
*
|
||||
* @param layer_id The ID of the layer for which to retrieve the block
|
||||
* number.
|
||||
* @return The number of blocks in the specified layer.
|
||||
*/
|
||||
int get_block_num(int layer_id) { return past_block_num_[layer_id]; }
|
||||
|
||||
/**
|
||||
* @brief Gets the number of anchors in the cache.
|
||||
*
|
||||
* @return The number of anchors configured in the cache.
|
||||
*/
|
||||
int get_anchor_num() { return config_.anchor_num; }
|
||||
|
||||
/**
|
||||
* @brief Gets the total length of the cache.
|
||||
*
|
||||
* @return The total length of the cache.
|
||||
*/
|
||||
int get_cache_total_len() { return cache_total_len_; }
|
||||
|
||||
/**
|
||||
* @brief Gets the total number of blocks in the cache.
|
||||
*
|
||||
* This function computes and returns the total number of blocks in the
|
||||
* cache based on the total cache length and the block length configuration.
|
||||
*
|
||||
* @return The total number of blocks in the cache.
|
||||
*/
|
||||
int get_cache_total_block_num() {
|
||||
return (cache_total_len_ + config_.block_len - 1) / config_.block_len;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Updates the total length of the cache.
|
||||
*
|
||||
* This function sets a new total length for the cache, allowing dynamic
|
||||
* adjustment of the cache size during runtime.
|
||||
*
|
||||
* @param cache_total_len The new total length of the cache.
|
||||
*/
|
||||
void update_cache_total_len(int cache_total_len) {
|
||||
cache_total_len_ = cache_total_len;
|
||||
}
|
||||
void attn(const ggml_fp16_t *q_in, ggml_fp16_t *output, float *attn_lse,
|
||||
int layer_idx, int generate_token_idx, int q_len, int batch_size,
|
||||
int max_block_num, int *block_table, int *cache_seqlens,
|
||||
int pick_block_num, int init_block_num, int local_block_num,
|
||||
Backend *backend);
|
||||
|
||||
void update_kvcache_one_block_fp16(const ggml_fp16_t *k_in,
|
||||
const ggml_fp16_t *v_in, int layer_id,
|
||||
int block_idx, Backend *backend);
|
||||
|
||||
void get_kvcache_one_block_fp16(ggml_fp16_t *k_in, ggml_fp16_t *v_in,
|
||||
int layer_id, int block_idx,
|
||||
Backend *backend);
|
||||
|
||||
void update_importance_one_block(const ggml_fp16_t *importance,
|
||||
int layer_id, int block_idx,
|
||||
Backend *backend);
|
||||
void get_importance_one_block(ggml_fp16_t *importance, int layer_id,
|
||||
int block_idx, Backend *backend);
|
||||
|
||||
void get_anchor_one_block(ggml_fp16_t *anchor, int layer_id, int block_idx,
|
||||
Backend *backend);
|
||||
|
||||
void update_anchor_one_block(const ggml_fp16_t *anchor, int layer_id,
|
||||
int block_idx, Backend *backend);
|
||||
|
||||
void calc_anchor_all_layers(int *block_table, int *cache_seqlens,
|
||||
int batch_size, int max_block_num,
|
||||
Backend *backend);
|
||||
|
||||
void load_kvcache(std::string tensor_file_path, Backend *backend);
|
||||
void dump_kvcache(int *block_table, int cache_total_len,
|
||||
std::string tensor_file_path, Backend *backend);
|
||||
|
||||
void get_and_update_kvcache_fp16(ggml_fp16_t *k_in, ggml_fp16_t *v_in,
|
||||
int layer_id, int *block_table,
|
||||
int batch_size, int max_block_num,
|
||||
int *cache_seqlens, int q_len,
|
||||
Backend *backend);
|
||||
|
||||
void get_kvcache_fp16(ggml_fp16_t *k_in, ggml_fp16_t *v_in, int layer_id,
|
||||
int *block_table, int batch_size, int max_block_num,
|
||||
int *cache_seqlens, Backend *backend);
|
||||
|
||||
void update_kvcache_fp16(const ggml_fp16_t *k_in, const ggml_fp16_t *v_in,
|
||||
int layer_id, int *block_table, int batch_size,
|
||||
int max_block_num, int *cache_seqlens, int q_len,
|
||||
Backend *backend);
|
||||
|
||||
void update_importance(const ggml_fp16_t *importance, int layer_id,
|
||||
int *block_table, int batch_size, int max_block_num,
|
||||
int *offset, int width, Backend *backend);
|
||||
|
||||
void attn_with_kvcache(const ggml_fp16_t *q_in, const ggml_fp16_t *k_in,
|
||||
const ggml_fp16_t *v_in, ggml_fp16_t *output,
|
||||
float *attn_lse, int layer_idx,
|
||||
int generate_token_idx, int q_len, int batch_size,
|
||||
int max_block_num, int *block_table,
|
||||
int *cache_seqlens, int topk, int local,
|
||||
Backend *backend);
|
||||
|
||||
void clear_importance_all_layers(int *block_table, int *cache_seqlens,
|
||||
int batch_size, int max_block_num,
|
||||
Backend *backend);
|
||||
|
||||
void clear_kvcache_all_layers(int *block_table, int *cache_seqlens,
|
||||
int batch_size, int max_block_num,
|
||||
Backend *backend);
|
||||
|
||||
void get_sincos(ggml_fp16_t *sin, ggml_fp16_t *cos, int seqlen);
|
||||
|
||||
void get_attn_sparsity(const ggml_fp16_t *q_in, float *attn_sparsity,
|
||||
int layer_idx, int generate_token_idx, int q_len,
|
||||
int batch_size, int max_block_num, int *block_table,
|
||||
int *cache_seqlens, int *block_table_origin,
|
||||
int *cache_seqlens_origin, int max_block_num_origin,
|
||||
int topk, int local, Backend *backend);
|
||||
|
||||
void get_all_kvcache_one_layer(int layer_id, ggml_fp16_t *k_in,
|
||||
ggml_fp16_t *v_in, Backend *backend);
|
||||
|
||||
private:
|
||||
// Persistent data
|
||||
KVCacheConfig config_;
|
||||
int n_gqa_; // q_head_num / kv_head_num
|
||||
int cache_total_len_; // Number of tokens in cache
|
||||
std::vector<uint64_t> past_block_num_; // [layer_num]
|
||||
std::vector<std::vector<std::vector<std::vector<block_q4_0>>>>
|
||||
k_cache_q4; // [layer_num, kv_head_num, past_block_num, block_len *
|
||||
// (head_dim / QK_4)]
|
||||
std::vector<std::vector<std::vector<std::vector<block_q4_0>>>>
|
||||
v_cache_q4; // [layer_num, kv_head_num, past_block_num, head_dim *
|
||||
// (block_len / QK_4)]
|
||||
std::vector<std::vector<std::vector<std::vector<block_q8_0>>>>
|
||||
k_cache_q8; // [layer_num, kv_head_num, past_block_num, block_len *
|
||||
// (head_dim / QK_8)]
|
||||
std::vector<std::vector<std::vector<std::vector<block_q8_0>>>>
|
||||
v_cache_q8; // [layer_num, kv_head_num, past_block_num, head_dim *
|
||||
// (block_len / QK_8)]
|
||||
|
||||
std::vector<std::vector<std::vector<std::vector<ggml_fp16_t>>>>
|
||||
k_cache_fp16_; // [layer_num, kv_head_num, past_block_num, block_len *
|
||||
// head_dim]
|
||||
std::vector<std::vector<std::vector<std::vector<ggml_fp16_t>>>>
|
||||
v_cache_fp16_; // [layer_num, kv_head_num, past_block_num, head_dim *
|
||||
// block_len]
|
||||
|
||||
std::vector<std::vector<std::vector<std::vector<ggml_fp16_t>>>>
|
||||
importance_; // [layer_num, past_block_num, block_len,
|
||||
// attention_head_num]
|
||||
|
||||
std::vector<ggml_fp16_t>
|
||||
anchor_; // [layer_num * past_block_num * anchor_num *
|
||||
// attention_head_num * head_dim]
|
||||
|
||||
// Runtime data
|
||||
int64_t layer_id_;
|
||||
int64_t block_idx_;
|
||||
int *block_table_;
|
||||
uint64_t block_num_;
|
||||
int max_block_num_after_retrieval_;
|
||||
|
||||
// Rotary positional embeddings
|
||||
std::vector<std::vector<ggml_fp16_t>> sin_; // [seq_len, head_dim]
|
||||
std::vector<std::vector<ggml_fp16_t>> cos_; // [seq_len, head_dim]
|
||||
|
||||
// update/get
|
||||
int seq_len_;
|
||||
uint16_t *k_scales_; // q4_0
|
||||
uint8_t *k_in_; // q4_0
|
||||
uint16_t *v_scales_; // q4_0
|
||||
uint8_t *v_in_; // q4_0
|
||||
uint16_t *k_data_; // fp16
|
||||
uint16_t *v_data_; // fp16
|
||||
uint16_t *importance_data_; // fp16
|
||||
uint16_t *anchor_data_; // fp16
|
||||
|
||||
// sparsity = (sigma(block lse / lse))
|
||||
std::vector<std::vector<std::vector<float>>>
|
||||
block_lse_; // [batch_size, max_block_num, q_head_num]
|
||||
std::vector<std::vector<float>> attn_sparsity_; // [batch_size, q_head_num]
|
||||
|
||||
// attn
|
||||
std::vector<std::vector<float>>
|
||||
avg_q; // [batch_size, q_head_num * head_dim]
|
||||
|
||||
std::vector<std::vector<ggml_fp16_t>>
|
||||
avg_q_fp16; // [batch_size, q_head_num * head_dim]
|
||||
std::vector<
|
||||
std::priority_queue<std::pair<float, int>,
|
||||
std::vector<std::pair<float, int>>, std::greater<>>>
|
||||
top_similar_block_;
|
||||
|
||||
std::vector<std::vector<float>> block_similar_;
|
||||
std::vector<std::vector<std::vector<float>>> block_similar_kv_head_;
|
||||
std::vector<std::vector<std::vector<float>>> block_similar_q_head_;
|
||||
|
||||
std::vector<int> cache_seqlens_; // [batch_size]
|
||||
std::vector<int> selected_blocks_num_history_; // [layer_num // layer_step]
|
||||
|
||||
std::vector<std::vector<std::vector<int>>> selected_blocks_history_;
|
||||
// [layer_num // layer_step, batch_size, max_block_num]
|
||||
|
||||
std::vector<std::vector<std::vector<std::vector<int>>>>
|
||||
selected_blocks_history_kvhead_; // [layer_num // layer_step,
|
||||
// batch_size, max_block_num,
|
||||
// kv_head_num]
|
||||
|
||||
std::vector<std::vector<int>>
|
||||
block_table_before_retrieval_; // [batch_size, max_block_num]
|
||||
std::vector<std::vector<int>>
|
||||
block_table_after_retrieval_; // [batch_size, pick_block_num]
|
||||
|
||||
std::vector<std::vector<std::vector<int>>>
|
||||
block_table_before_retrieval_qhead_; // [batch_size, max_block_num,
|
||||
// q_head_num]
|
||||
std::vector<std::vector<std::vector<int>>>
|
||||
block_table_after_retrieval_qhead_; // [batch_size, pick_block_num,
|
||||
// q_head_num]
|
||||
|
||||
std::vector<std::vector<std::vector<int>>>
|
||||
block_table_before_retrieval_kvhead_; // [batch_size, max_block_num,
|
||||
// kv_head_num]
|
||||
std::vector<std::vector<std::vector<int>>>
|
||||
block_table_after_retrieval_kvhead_; // [batch_size, pick_block_num,
|
||||
// kv_head_num]
|
||||
|
||||
std::vector<std::vector<std::unique_ptr<std::mutex>>>
|
||||
mutex_; // [batch_size, kv_head_num]
|
||||
std::vector<std::vector<std::vector<block_q8_0>>>
|
||||
q_q8_0_; // [batch_size, kv_head_num, n_gqa * head_dim / QK8_0]
|
||||
std::vector<std::vector<std::vector<float>>>
|
||||
q_fp32_; // [batch_size, kv_head_num, n_gqa * head_dim]
|
||||
|
||||
std::vector<std::vector<std::vector<float>>>
|
||||
output_fp32_; // [batch_size, kv_head_num, n_gqa * head_dim]
|
||||
std::vector<std::vector<std::vector<float>>>
|
||||
attn_lse_; // [batch_size, kv_head_num, n_gqa]
|
||||
|
||||
std::vector<std::pair<int, int>> thread_cur_head_idx_; // [thread_num]
|
||||
|
||||
std::vector<std::vector<block_q8_0>>
|
||||
thread_local_output_q8_0_; // [thread_num, n_gqa * head_dim / QK8_0]
|
||||
std::vector<std::vector<float>>
|
||||
thread_local_attn_score_; // [thread_num, n_gqa * block_len]
|
||||
std::vector<std::vector<float>>
|
||||
thread_local_output_fp32_; // [thread_num, n_gqa * head_dim]
|
||||
std::vector<std::vector<float>>
|
||||
thread_local_attn_lse_; // [thread_num, n_gqa]
|
||||
std::vector<std::vector<float>>
|
||||
thread_local_cur_output_fp32_; // [thread_num, n_gqa * head_dim]
|
||||
std::vector<std::vector<float>>
|
||||
thread_local_cur_attn_lse_; // [thread_num, n_gqa]
|
||||
std::vector<std::vector<uint8_t>>
|
||||
thread_local_attn_mask_; // [thread_num, block_len // 8]
|
||||
std::vector<std::vector<char>>
|
||||
thread_local_draft_; // [thread_num, 2 * n_gqa * block_len + 6 * n_gqa *
|
||||
// head_dim + 2 * block_len * head_dim]
|
||||
|
||||
// tmp space
|
||||
std::vector<float> q_fp32; // [n_gqa * head_dim]
|
||||
|
||||
void quantize_q_(const uint16_t *q_in_data, int batch_size);
|
||||
void attn_initialize_layer_(int batch_size, int layer_idx, int *block_table,
|
||||
int &max_block_num, int *cache_seqlens);
|
||||
void attn_initialize_kvhead_(int batch_size, int layer_idx,
|
||||
int *block_table, int &max_block_num,
|
||||
int *cache_seqlens);
|
||||
void retrieval_kvcache_layer_(const uint16_t *q_in_data, int init_block_num,
|
||||
int local_block_num, int pick_block_num,
|
||||
int q_len, int generate_token_idx,
|
||||
int batch_size, int layer_idx,
|
||||
int *cache_seqlens, int &max_block_num,
|
||||
Backend *backend);
|
||||
void retrieval_kvcache_kvhead_(const uint16_t *q_in_data,
|
||||
int init_block_num, int local_block_num,
|
||||
int pick_block_num, int q_len,
|
||||
int generate_token_idx, int batch_size,
|
||||
int layer_idx, int *cache_seqlens,
|
||||
int &max_block_num, Backend *backend);
|
||||
|
||||
void calculate_block_similarity_layer_(
|
||||
const uint16_t *q_in_data, int batch_size, int layer_idx, int q_len,
|
||||
int max_block_num, int *cache_seqlens, int init_block_num,
|
||||
int local_block_num, int pick_block_num, Backend *backend);
|
||||
void calculate_block_similarity_kvhead_(
|
||||
const uint16_t *q_in_data, int batch_size, int layer_idx, int q_len,
|
||||
int max_block_num, int *cache_seqlens, int init_block_num,
|
||||
int local_block_num, int pick_block_num, Backend *backend);
|
||||
|
||||
void select_block_layer_(int batch_size, int layer_idx, int max_block_num,
|
||||
int init_block_num, int local_block_num,
|
||||
int pick_block_num);
|
||||
void select_block_kvhead_(int batch_size, int layer_idx, int max_block_num,
|
||||
int init_block_num, int local_block_num,
|
||||
int pick_block_num);
|
||||
|
||||
void calculate_sparsity_layer_(const uint16_t *q_in_data,
|
||||
float *attn_sparsity, int batch_size,
|
||||
int max_block_num, int *block_table,
|
||||
int *cache_seqlens, Backend *backend);
|
||||
void calculate_sparsity_kvhead_(const uint16_t *q_in_data,
|
||||
float *attn_sparsity, int batch_size,
|
||||
int max_block_num, int *block_table,
|
||||
int *cache_seqlens, Backend *backend);
|
||||
|
||||
void attention_kvhead_(const uint16_t *q_in_data, ggml_fp16_t *output,
|
||||
float *attn_lse, int batch_size, Backend *backend);
|
||||
void attention_layer_(const uint16_t *q_in_data, ggml_fp16_t *output,
|
||||
float *attn_lse, int batch_size, Backend *backend);
|
||||
|
||||
/**
|
||||
* @brief Computes attention with KV cache for one block.
|
||||
*
|
||||
* This function performs attention computation for one block using KV
|
||||
* cache. The function supports different data types for Q, K, and V caches,
|
||||
* and provides options for quantization. The function does not perform any
|
||||
* dynamic memory allocation internally, so all necessary buffers must be
|
||||
* pre-allocated externally.
|
||||
*
|
||||
* @param head_dim The dimension of the head.
|
||||
* @param bsz The batch size.
|
||||
* @param q_type The data type of Q (GGML data type). Only supports fp16 and
|
||||
* q8_0.
|
||||
* @param q Pointer to the Q tensor [bsz, head_dim]. The quantization is
|
||||
* always applied along the head_dim dimension. The size must be
|
||||
* bsz * head_dim/32 * qtype_size. If head_dim % 32 != 0, an error
|
||||
* will be raised.
|
||||
* @param past_kv_len The length of the past KV cache.
|
||||
* @param past_kv_offset The offset in the past KV cache.
|
||||
* @param is_full_attn Boolean flag indicating whether to use full attention
|
||||
* (true for full 1 mask).
|
||||
* @param attn_mask Pointer to the attention mask [bsz, past_kv_len]. If
|
||||
* is_full_attn = false, a bit matrix is passed to
|
||||
* represent the mask.
|
||||
* @param k_type The data type of K cache (GGML data type). Only supports
|
||||
* fp16, q4_0, and q8_0.
|
||||
* @param k_quant_type Quantization type for K cache. 0 for per_token, 1 for
|
||||
* per_channel. Other values will raise an error.
|
||||
* @param k_cache Pointer to the K cache tensor [seq_len, head_dim]. If
|
||||
* quant_type == 0, head_dim % 32 must be 0. If quant_type ==
|
||||
* 1, seq_len % 32 must be 0.
|
||||
* @param num_k_anchor The number of K anchors. If num_k_anchor == 0, it
|
||||
* means no anchor is present.
|
||||
* @param k_cache_anchors Pointer to the K cache anchors [num_k_anchor,
|
||||
* head_dim]. The k_anchor_type must be fp16.
|
||||
* @param k_cache_anchor_pos Pointer to the K cache anchor positions. Each
|
||||
* token is associated with the nearest previous anchor position.
|
||||
* @param v_type The data type of V cache (GGML data type).
|
||||
* @param v_quant_type Quantization type for V cache.
|
||||
* @param v_cache Pointer to the V cache tensor [head_dim, seq_len].
|
||||
* @param num_v_anchor The number of V anchors.
|
||||
* @param v_cache_anchors Pointer to the V cache anchors.
|
||||
* @param v_cache_anchor_pos Pointer to the V cache anchor positions.
|
||||
* @param attn_score Pre-allocated buffer for attention scores [bsz,
|
||||
* past_kv_len].
|
||||
* @param output Output tensor [bsz, head_dim] with the same type as q_type.
|
||||
* @param lse Pre-allocated buffer [bsz] for the log-sum-exp of the
|
||||
* attention scores.
|
||||
* @param draft Pre-allocated temporary buffer. The buffer size should be
|
||||
* enough to hold (2 * bsz * past_kv_len + 6 * bsz * head_dim + 2 *
|
||||
* past_kv_len * head_dim + past_kv_len * head_dim / 32) bytes.
|
||||
* @param rotary_angle Pointer to the rotary angle tensor.
|
||||
* @param rotary_cos Pointer to the cosine values for rotary embedding.
|
||||
* @param rotary_sin Pointer to the sine values for rotary embedding.
|
||||
*/
|
||||
void attn_with_kvcache_one_block_(
|
||||
int head_dim, int bsz,
|
||||
ggml_type q_type, // GGML data type of `Q`, only supports fp16 and q8_0
|
||||
// [bsz, head_dim]
|
||||
// Quantization is always on the head_dim dimension (per_token). If
|
||||
// head_dim % 32 != 0, an error will be raised. The size must be bsz *
|
||||
// head_dim/32 * qtype_size.
|
||||
const void *q,
|
||||
|
||||
int past_kv_len, int past_kv_offset,
|
||||
bool is_full_attn, // true indicates a full 1 mask
|
||||
// If is_full_attn = false, a bit matrix representing the mask is
|
||||
// passed. [bsz, past_kv_len]
|
||||
const uint8_t *attn_mask,
|
||||
|
||||
ggml_type k_type, // GGML data type of `K Cache`, only supports fp16,
|
||||
// q4_0, q8_0
|
||||
int k_quant_type, // 0 for per_token, 1 for per_channel, others raise an
|
||||
// error
|
||||
// [seq_len, head_dim]
|
||||
// If quant_type == 0, head_dim % 32 must be 0.
|
||||
// If quant_type == 1, seq_len % 32 must be 0.
|
||||
const void *k_cache,
|
||||
|
||||
// k_anchor_type must be fp16
|
||||
int num_k_anchor, // num_k_anchor == 0 indicates no anchor
|
||||
// [num_k_anchor, head_dim]
|
||||
const void *k_cache_anchors,
|
||||
// Each token is associated with the nearest previous position's anchor,
|
||||
// with the same distance.
|
||||
const int *k_cache_anchor_pos,
|
||||
|
||||
// v_cache similar to k_cache
|
||||
ggml_type v_type, int v_quant_type,
|
||||
// [head_dim, seq_len]
|
||||
const void *v_cache, int num_v_anchor, const void *v_cache_anchors,
|
||||
const int *v_cache_anchor_pos,
|
||||
|
||||
// Pre-allocated buffer for intermediate calculations [bsz,
|
||||
// past_kv_len]. No malloc is performed inside this function.
|
||||
float *attn_score,
|
||||
|
||||
// Output: [bsz, head_dim], with the same type as q_type
|
||||
void *output,
|
||||
// [bsz]
|
||||
float *lse,
|
||||
|
||||
// Pre-allocated temporary buffer with sufficient size:
|
||||
// (2 * bsz * past_kv_len + 6 * bsz * head_dim + 2 * past_kv_len *
|
||||
// head_dim + past_kv_len * head_dim / 32) bytes.
|
||||
void *draft,
|
||||
|
||||
// Apply rotary embedding online
|
||||
const int *rotary_angle, const void *rotary_cos, const void *rotary_sin
|
||||
// rotary_cos=None,
|
||||
// rotary_sin=None,
|
||||
// cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None,
|
||||
// cache_batch_idx: Optional[torch.Tensor] = None,
|
||||
// rotary_interleaved=True,
|
||||
|
||||
// // Not supported for now
|
||||
// window_size=(-1, -1), # -1 means infinite context window
|
||||
// alibi_slopes=None,
|
||||
);
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Scales a float32 vector by a given scalar value.
|
||||
*
|
||||
* This function multiplies each element of the input vector `y` by a scalar
|
||||
* `v`. It uses platform-specific optimizations if available, such as Apple's
|
||||
* Accelerate framework or SIMD instructions. If no specific optimization is
|
||||
* available, the function falls back to a simple scalar multiplication loop.
|
||||
*
|
||||
* @param n The number of elements in the vector `y`.
|
||||
* @param y The input vector to be scaled. The result will be stored in the same
|
||||
* vector.
|
||||
* @param v The scalar value by which to scale the vector.
|
||||
*/
|
||||
void ggml_vec_scale_f32(const int n, float *y, const float v);
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,126 @@
|
||||
/**
|
||||
* @Description :
|
||||
* @Author : Jianwei Dong
|
||||
* @Date : 2024-08-26 22:47:06
|
||||
* @Version : 1.0.0
|
||||
* @LastEditors : Jianwei Dong
|
||||
* @LastEditTime : 2024-08-26 22:47:06
|
||||
* @Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
|
||||
**/
|
||||
|
||||
#include "kvcache.h"
|
||||
|
||||
#include <chrono>
|
||||
|
||||
void KVCache::load_kvcache(std::string tensor_file_path, Backend *backend) {
|
||||
// Timer start
|
||||
auto start = std::chrono::high_resolution_clock::now();
|
||||
std::ifstream ifs_tensor(tensor_file_path, std::ios::binary);
|
||||
if (!ifs_tensor) {
|
||||
throw std::runtime_error("Failed to open tensor file");
|
||||
}
|
||||
ifs_tensor.read(reinterpret_cast<char *>(&cache_total_len_),
|
||||
sizeof(cache_total_len_));
|
||||
int past_block_num =
|
||||
(cache_total_len_ + config_.block_len - 1) / config_.block_len;
|
||||
printf("cache_total_len: %d, past_block_num: %d\n", cache_total_len_,
|
||||
past_block_num);
|
||||
for (int i = 0; i < config_.layer_num; ++i) {
|
||||
past_block_num_[i] = past_block_num;
|
||||
}
|
||||
ifs_tensor.read(reinterpret_cast<char *>(anchor_.data()),
|
||||
anchor_.size() * sizeof(ggml_fp16_t));
|
||||
for (int i = 0; i < config_.layer_num; ++i) {
|
||||
for (int j = 0; j < config_.kv_head_num; ++j) {
|
||||
for (int k = 0; k < past_block_num_[i]; ++k) {
|
||||
if (config_.kv_type == GGML_TYPE_F16) {
|
||||
ifs_tensor.read(
|
||||
reinterpret_cast<char *>(k_cache_fp16_[i][j][k].data()),
|
||||
k_cache_fp16_[i][j][k].size() * sizeof(ggml_fp16_t));
|
||||
ifs_tensor.read(
|
||||
reinterpret_cast<char *>(v_cache_fp16_[i][j][k].data()),
|
||||
v_cache_fp16_[i][j][k].size() * sizeof(ggml_fp16_t));
|
||||
} else if (config_.kv_type == GGML_TYPE_Q4_0) {
|
||||
ifs_tensor.read(
|
||||
reinterpret_cast<char *>(k_cache_q4[i][j][k].data()),
|
||||
k_cache_q4[i][j][k].size() * sizeof(block_q4_0));
|
||||
ifs_tensor.read(
|
||||
reinterpret_cast<char *>(v_cache_q4[i][j][k].data()),
|
||||
v_cache_q4[i][j][k].size() * sizeof(block_q4_0));
|
||||
}
|
||||
}
|
||||
}
|
||||
for (int k = 0; k < past_block_num_[i]; ++k) {
|
||||
for (int l = 0; l < config_.block_len; l++) {
|
||||
ifs_tensor.read(
|
||||
reinterpret_cast<char *>(importance_[i][k][l].data()),
|
||||
importance_[i][k][l].size() * sizeof(ggml_fp16_t));
|
||||
}
|
||||
}
|
||||
}
|
||||
ifs_tensor.close();
|
||||
// Timer end
|
||||
auto end = std::chrono::high_resolution_clock::now();
|
||||
std::chrono::duration<double> diff = end - start;
|
||||
printf("time of load: %f s\n", diff.count());
|
||||
}
|
||||
void KVCache::dump_kvcache(int *block_table, int cache_total_len,
|
||||
std::string tensor_file_path, Backend *backend) {
|
||||
// Timer start
|
||||
auto start = std::chrono::high_resolution_clock::now();
|
||||
std::ofstream ofs(tensor_file_path, std::ios::binary);
|
||||
printf("dump_kvcache: %s\n", tensor_file_path.c_str());
|
||||
if (!ofs.is_open()) {
|
||||
std::cerr << "Cannot open file " << tensor_file_path << std::endl;
|
||||
return;
|
||||
}
|
||||
ofs.write(reinterpret_cast<const char *>(&cache_total_len),
|
||||
sizeof(cache_total_len));
|
||||
int past_block_num =
|
||||
(cache_total_len + config_.block_len - 1) / config_.block_len;
|
||||
printf("cache_total_len: %d, past_block_num: %d\n", cache_total_len,
|
||||
past_block_num);
|
||||
ofs.write(reinterpret_cast<const char *>(anchor_.data()),
|
||||
anchor_.size() * sizeof(ggml_fp16_t));
|
||||
for (int i = 0; i < config_.layer_num; ++i) {
|
||||
for (int j = 0; j < config_.kv_head_num; ++j) {
|
||||
for (int k = 0; k < past_block_num; ++k) {
|
||||
int block_idx = block_table[k];
|
||||
if (config_.kv_type == GGML_TYPE_F16) {
|
||||
ofs.write(reinterpret_cast<const char *>(
|
||||
k_cache_fp16_[i][j][block_idx].data()),
|
||||
k_cache_fp16_[i][j][block_idx].size() *
|
||||
sizeof(ggml_fp16_t));
|
||||
ofs.write(reinterpret_cast<const char *>(
|
||||
v_cache_fp16_[i][j][block_idx].data()),
|
||||
v_cache_fp16_[i][j][block_idx].size() *
|
||||
sizeof(ggml_fp16_t));
|
||||
|
||||
} else if (config_.kv_type == GGML_TYPE_Q4_0) {
|
||||
ofs.write(reinterpret_cast<const char *>(
|
||||
k_cache_q4[i][j][block_idx].data()),
|
||||
k_cache_q4[i][j][block_idx].size() *
|
||||
sizeof(block_q4_0));
|
||||
ofs.write(reinterpret_cast<const char *>(
|
||||
v_cache_q4[i][j][block_idx].data()),
|
||||
v_cache_q4[i][j][block_idx].size() *
|
||||
sizeof(block_q4_0));
|
||||
}
|
||||
}
|
||||
}
|
||||
for (int k = 0; k < past_block_num; ++k) {
|
||||
int block_idx = block_table[k];
|
||||
for (int l = 0; l < config_.block_len; l++) {
|
||||
ofs.write(reinterpret_cast<const char *>(
|
||||
importance_[i][block_idx][l].data()),
|
||||
importance_[i][block_idx][l].size() *
|
||||
sizeof(ggml_fp16_t));
|
||||
}
|
||||
}
|
||||
}
|
||||
ofs.close();
|
||||
// Timer end
|
||||
auto end = std::chrono::high_resolution_clock::now();
|
||||
std::chrono::duration<double> diff = end - start;
|
||||
printf("time of dump: %f s\n", diff.count());
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user