chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
+47
@@ -0,0 +1,47 @@
|
||||
# Extensions
|
||||
|
||||
*.a
|
||||
*.bat
|
||||
*.bin
|
||||
*.dll
|
||||
*.dot
|
||||
*.etag
|
||||
*.exe
|
||||
*.gcda
|
||||
*.gcno
|
||||
*.gcov
|
||||
*.gguf
|
||||
*.gguf.json
|
||||
*.lastModified
|
||||
*.log
|
||||
*.metallib
|
||||
*.o
|
||||
*.so
|
||||
*.tmp
|
||||
|
||||
# IDE / OS
|
||||
|
||||
.cache/
|
||||
.ccls-cache/
|
||||
.direnv/
|
||||
.DS_Store
|
||||
.envrc
|
||||
.idea/
|
||||
.swiftpm
|
||||
.vs/
|
||||
.vscode/
|
||||
nppBackup
|
||||
|
||||
# Models
|
||||
models/*
|
||||
gpu/checkpoints/*
|
||||
|
||||
# Python
|
||||
|
||||
/.venv
|
||||
__pycache__/
|
||||
*/poetry.lock
|
||||
poetry.toml
|
||||
|
||||
build/
|
||||
logs/
|
||||
@@ -0,0 +1,4 @@
|
||||
[submodule "3rdparty/llama.cpp"]
|
||||
path = 3rdparty/llama.cpp
|
||||
url = https://github.com/Eddie-Wang1120/llama.cpp.git
|
||||
branch = merge-dev
|
||||
@@ -0,0 +1,78 @@
|
||||
cmake_minimum_required(VERSION 3.14) # for add_link_options and implicit target directories.
|
||||
project("bitnet.cpp" C CXX)
|
||||
include(CheckIncludeFileCXX)
|
||||
|
||||
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
|
||||
|
||||
if (NOT XCODE AND NOT MSVC AND NOT CMAKE_BUILD_TYPE)
|
||||
set(CMAKE_BUILD_TYPE Release CACHE STRING "Build type" FORCE)
|
||||
set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "Debug" "Release" "MinSizeRel" "RelWithDebInfo")
|
||||
endif()
|
||||
|
||||
set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
|
||||
|
||||
# option list
|
||||
option(BITNET_ARM_TL1 "bitnet.cpp: use tl1 on arm platform" OFF)
|
||||
option(BITNET_X86_TL2 "bitnet.cpp: use tl2 on x86 platform" OFF)
|
||||
|
||||
|
||||
set(CMAKE_CXX_STANDARD_REQUIRED true)
|
||||
set(CMAKE_C_STANDARD 11)
|
||||
set(CMAKE_C_STANDARD_REQUIRED true)
|
||||
set(THREADS_PREFER_PTHREAD_FLAG ON)
|
||||
|
||||
# override ggml options
|
||||
set(GGML_BITNET_ARM_TL1 ${BITNET_ARM_TL1})
|
||||
set(GGML_BITNET_X86_TL2 ${BITNET_X86_TL2})
|
||||
|
||||
if (GGML_BITNET_ARM_TL1)
|
||||
add_compile_definitions(GGML_BITNET_ARM_TL1)
|
||||
endif()
|
||||
if (GGML_BITNET_X86_TL2)
|
||||
add_compile_definitions(GGML_BITNET_X86_TL2)
|
||||
endif()
|
||||
|
||||
if (CMAKE_C_COMPILER_ID STREQUAL "GNU" OR CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
|
||||
add_compile_options(-fpermissive)
|
||||
endif()
|
||||
|
||||
find_package(Threads REQUIRED)
|
||||
|
||||
add_subdirectory(src)
|
||||
set(LLAMA_BUILD_SERVER ON CACHE BOOL "Build llama.cpp server" FORCE)
|
||||
add_subdirectory(3rdparty/llama.cpp)
|
||||
|
||||
# install
|
||||
|
||||
include(GNUInstallDirs)
|
||||
include(CMakePackageConfigHelpers)
|
||||
|
||||
set(LLAMA_INCLUDE_INSTALL_DIR ${CMAKE_INSTALL_INCLUDEDIR}
|
||||
CACHE PATH "Location of header files")
|
||||
set(LLAMA_LIB_INSTALL_DIR ${CMAKE_INSTALL_LIBDIR}
|
||||
CACHE PATH "Location of library files")
|
||||
set(LLAMA_BIN_INSTALL_DIR ${CMAKE_INSTALL_BINDIR}
|
||||
CACHE PATH "Location of binary files")
|
||||
set(LLAMA_BUILD_NUMBER ${BUILD_NUMBER})
|
||||
set(LLAMA_BUILD_COMMIT ${BUILD_COMMIT})
|
||||
set(LLAMA_INSTALL_VERSION 0.0.${BUILD_NUMBER})
|
||||
|
||||
get_target_property(GGML_DIRECTORY ggml SOURCE_DIR)
|
||||
get_directory_property(GGML_DIR_DEFINES DIRECTORY ${GGML_DIRECTORY} COMPILE_DEFINITIONS)
|
||||
get_target_property(GGML_TARGET_DEFINES ggml COMPILE_DEFINITIONS)
|
||||
set(GGML_TRANSIENT_DEFINES ${GGML_TARGET_DEFINES} ${GGML_DIR_DEFINES})
|
||||
get_target_property(GGML_LINK_LIBRARIES ggml LINK_LIBRARIES)
|
||||
|
||||
get_directory_property(LLAMA_TRANSIENT_DEFINES COMPILE_DEFINITIONS)
|
||||
|
||||
write_basic_package_version_file(
|
||||
${CMAKE_CURRENT_BINARY_DIR}/LlamaConfigVersion.cmake
|
||||
VERSION ${LLAMA_INSTALL_VERSION}
|
||||
COMPATIBILITY SameMajorVersion)
|
||||
|
||||
install(FILES ${CMAKE_CURRENT_BINARY_DIR}/LlamaConfig.cmake
|
||||
${CMAKE_CURRENT_BINARY_DIR}/LlamaConfigVersion.cmake
|
||||
DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/Llama)
|
||||
|
||||
set_target_properties(llama PROPERTIES PUBLIC_HEADER ${CMAKE_CURRENT_SOURCE_DIR}/llama.h)
|
||||
install(TARGETS llama LIBRARY PUBLIC_HEADER)
|
||||
@@ -0,0 +1,9 @@
|
||||
# Microsoft Open Source Code of Conduct
|
||||
|
||||
This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/).
|
||||
|
||||
Resources:
|
||||
|
||||
- [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/)
|
||||
- [Microsoft Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/)
|
||||
- Contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with questions or concerns
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) Microsoft Corporation.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE
|
||||
@@ -0,0 +1,339 @@
|
||||
# bitnet.cpp
|
||||
[](https://opensource.org/licenses/MIT)
|
||||
|
||||
|
||||
[<img src="./assets/header_model_release.png" alt="BitNet Model on Hugging Face" width="800"/>](https://huggingface.co/microsoft/BitNet-b1.58-2B-4T)
|
||||
|
||||
Try it out via this [demo](https://demo-bitnet-h0h8hcfqeqhrf5gf.canadacentral-01.azurewebsites.net/), or build and run it on your own [CPU](https://github.com/microsoft/BitNet?tab=readme-ov-file#build-from-source) or [GPU](https://github.com/microsoft/BitNet/blob/main/gpu/README.md).
|
||||
|
||||
bitnet.cpp is the official inference framework for 1-bit LLMs (e.g., BitNet b1.58). It offers a suite of optimized kernels, that support **fast** and **lossless** inference of 1.58-bit models on CPU and GPU (NPU support will coming next).
|
||||
|
||||
The first release of bitnet.cpp is to support inference on CPUs. bitnet.cpp achieves speedups of **1.37x** to **5.07x** on ARM CPUs, with larger models experiencing greater performance gains. Additionally, it reduces energy consumption by **55.4%** to **70.0%**, further boosting overall efficiency. On x86 CPUs, speedups range from **2.37x** to **6.17x** with energy reductions between **71.9%** to **82.2%**. Furthermore, bitnet.cpp can run a 100B BitNet b1.58 model on a single CPU, achieving speeds comparable to human reading (5-7 tokens per second), significantly enhancing the potential for running LLMs on local devices. Please refer to the [technical report](https://arxiv.org/abs/2410.16144) for more details.
|
||||
|
||||
**Latest optimization** introduces parallel kernel implementations with configurable tiling and embedding quantization support, achieving **1.15x to 2.1x** additional speedup over the original implementation across different hardware platforms and workloads. For detailed technical information, see the [optimization guide](src/README.md).
|
||||
|
||||
<img src="./assets/performance.png" alt="performance_comparison" width="800"/>
|
||||
|
||||
|
||||
## Demo
|
||||
|
||||
A demo of bitnet.cpp running a BitNet b1.58 3B model on Apple M2:
|
||||
|
||||
https://github.com/user-attachments/assets/7f46b736-edec-4828-b809-4be780a3e5b1
|
||||
|
||||
## What's New:
|
||||
- 01/15/2026 [BitNet CPU Inference Optimization](https://github.com/microsoft/BitNet/blob/main/src/README.md) 
|
||||
- 05/20/2025 [BitNet Official GPU inference kernel](https://github.com/microsoft/BitNet/blob/main/gpu/README.md)
|
||||
- 04/14/2025 [BitNet Official 2B Parameter Model on Hugging Face](https://huggingface.co/microsoft/BitNet-b1.58-2B-4T)
|
||||
- 02/18/2025 [Bitnet.cpp: Efficient Edge Inference for Ternary LLMs](https://arxiv.org/abs/2502.11880)
|
||||
- 11/08/2024 [BitNet a4.8: 4-bit Activations for 1-bit LLMs](https://arxiv.org/abs/2411.04965)
|
||||
- 10/21/2024 [1-bit AI Infra: Part 1.1, Fast and Lossless BitNet b1.58 Inference on CPUs](https://arxiv.org/abs/2410.16144)
|
||||
- 10/17/2024 bitnet.cpp 1.0 released.
|
||||
- 03/21/2024 [The-Era-of-1-bit-LLMs__Training_Tips_Code_FAQ](https://github.com/microsoft/unilm/blob/master/bitnet/The-Era-of-1-bit-LLMs__Training_Tips_Code_FAQ.pdf)
|
||||
- 02/27/2024 [The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits](https://arxiv.org/abs/2402.17764)
|
||||
- 10/17/2023 [BitNet: Scaling 1-bit Transformers for Large Language Models](https://arxiv.org/abs/2310.11453)
|
||||
|
||||
## Acknowledgements
|
||||
|
||||
This project is based on the [llama.cpp](https://github.com/ggerganov/llama.cpp) framework. We would like to thank all the authors for their contributions to the open-source community. Also, bitnet.cpp's kernels are built on top of the Lookup Table methodologies pioneered in [T-MAC](https://github.com/microsoft/T-MAC/). For inference of general low-bit LLMs beyond ternary models, we recommend using T-MAC.
|
||||
## Official Models
|
||||
<table>
|
||||
</tr>
|
||||
<tr>
|
||||
<th rowspan="2">Model</th>
|
||||
<th rowspan="2">Parameters</th>
|
||||
<th rowspan="2">CPU</th>
|
||||
<th colspan="3">Kernel</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<th>I2_S</th>
|
||||
<th>TL1</th>
|
||||
<th>TL2</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td rowspan="2"><a href="https://huggingface.co/microsoft/BitNet-b1.58-2B-4T">BitNet-b1.58-2B-4T</a></td>
|
||||
<td rowspan="2">2.4B</td>
|
||||
<td>x86</td>
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
<td>✅</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>ARM</td>
|
||||
<td>✅</td>
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
## Supported Models
|
||||
❗️**We use existing 1-bit LLMs available on [Hugging Face](https://huggingface.co/) to demonstrate the inference capabilities of bitnet.cpp. We hope the release of bitnet.cpp will inspire the development of 1-bit LLMs in large-scale settings in terms of model size and training tokens.**
|
||||
|
||||
<table>
|
||||
</tr>
|
||||
<tr>
|
||||
<th rowspan="2">Model</th>
|
||||
<th rowspan="2">Parameters</th>
|
||||
<th rowspan="2">CPU</th>
|
||||
<th colspan="3">Kernel</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<th>I2_S</th>
|
||||
<th>TL1</th>
|
||||
<th>TL2</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td rowspan="2"><a href="https://huggingface.co/1bitLLM/bitnet_b1_58-large">bitnet_b1_58-large</a></td>
|
||||
<td rowspan="2">0.7B</td>
|
||||
<td>x86</td>
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
<td>✅</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>ARM</td>
|
||||
<td>✅</td>
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td rowspan="2"><a href="https://huggingface.co/1bitLLM/bitnet_b1_58-3B">bitnet_b1_58-3B</a></td>
|
||||
<td rowspan="2">3.3B</td>
|
||||
<td>x86</td>
|
||||
<td>❌</td>
|
||||
<td>❌</td>
|
||||
<td>✅</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>ARM</td>
|
||||
<td>❌</td>
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td rowspan="2"><a href="https://huggingface.co/HF1BitLLM/Llama3-8B-1.58-100B-tokens">Llama3-8B-1.58-100B-tokens</a></td>
|
||||
<td rowspan="2">8.0B</td>
|
||||
<td>x86</td>
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
<td>✅</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>ARM</td>
|
||||
<td>✅</td>
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td rowspan="2"><a href="https://huggingface.co/collections/tiiuae/falcon3-67605ae03578be86e4e87026">Falcon3 Family</a></td>
|
||||
<td rowspan="2">1B-10B</td>
|
||||
<td>x86</td>
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
<td>✅</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>ARM</td>
|
||||
<td>✅</td>
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td rowspan="2"><a href="https://huggingface.co/collections/tiiuae/falcon-edge-series-6804fd13344d6d8a8fa71130">Falcon-E Family</a></td>
|
||||
<td rowspan="2">1B-3B</td>
|
||||
<td>x86</td>
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
<td>✅</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>ARM</td>
|
||||
<td>✅</td>
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
|
||||
|
||||
## Installation
|
||||
|
||||
### Requirements
|
||||
- python>=3.9
|
||||
- cmake>=3.22
|
||||
- clang>=18
|
||||
- For Windows users, install [Visual Studio 2022](https://visualstudio.microsoft.com/downloads/). In the installer, toggle on at least the following options(this also automatically installs the required additional tools like CMake):
|
||||
- Desktop-development with C++
|
||||
- C++-CMake Tools for Windows
|
||||
- Git for Windows
|
||||
- C++-Clang Compiler for Windows
|
||||
- MS-Build Support for LLVM-Toolset (clang)
|
||||
- For Debian/Ubuntu users, you can download with [Automatic installation script](https://apt.llvm.org/)
|
||||
|
||||
`bash -c "$(wget -O - https://apt.llvm.org/llvm.sh)"`
|
||||
- conda (highly recommend)
|
||||
|
||||
### Build from source
|
||||
|
||||
> [!IMPORTANT]
|
||||
> If you are using Windows, please remember to always use a Developer Command Prompt / PowerShell for VS2022 for the following commands. Please refer to the FAQs below if you see any issues.
|
||||
|
||||
1. Clone the repo
|
||||
```bash
|
||||
git clone --recursive https://github.com/microsoft/BitNet.git
|
||||
cd BitNet
|
||||
```
|
||||
2. Install the dependencies
|
||||
```bash
|
||||
# (Recommended) Create a new conda environment
|
||||
conda create -n bitnet-cpp python=3.9
|
||||
conda activate bitnet-cpp
|
||||
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
3. Build the project
|
||||
```bash
|
||||
# Manually download the model and run with local path
|
||||
huggingface-cli download microsoft/BitNet-b1.58-2B-4T-gguf --local-dir models/BitNet-b1.58-2B-4T
|
||||
python setup_env.py -md models/BitNet-b1.58-2B-4T -q i2_s
|
||||
|
||||
```
|
||||
<pre>
|
||||
usage: setup_env.py [-h] [--hf-repo {1bitLLM/bitnet_b1_58-large,1bitLLM/bitnet_b1_58-3B,HF1BitLLM/Llama3-8B-1.58-100B-tokens,tiiuae/Falcon3-1B-Instruct-1.58bit,tiiuae/Falcon3-3B-Instruct-1.58bit,tiiuae/Falcon3-7B-Instruct-1.58bit,tiiuae/Falcon3-10B-Instruct-1.58bit}] [--model-dir MODEL_DIR] [--log-dir LOG_DIR] [--quant-type {i2_s,tl1}] [--quant-embd]
|
||||
[--use-pretuned]
|
||||
|
||||
Setup the environment for running inference
|
||||
|
||||
optional arguments:
|
||||
-h, --help show this help message and exit
|
||||
--hf-repo {1bitLLM/bitnet_b1_58-large,1bitLLM/bitnet_b1_58-3B,HF1BitLLM/Llama3-8B-1.58-100B-tokens,tiiuae/Falcon3-1B-Instruct-1.58bit,tiiuae/Falcon3-3B-Instruct-1.58bit,tiiuae/Falcon3-7B-Instruct-1.58bit,tiiuae/Falcon3-10B-Instruct-1.58bit}, -hr {1bitLLM/bitnet_b1_58-large,1bitLLM/bitnet_b1_58-3B,HF1BitLLM/Llama3-8B-1.58-100B-tokens,tiiuae/Falcon3-1B-Instruct-1.58bit,tiiuae/Falcon3-3B-Instruct-1.58bit,tiiuae/Falcon3-7B-Instruct-1.58bit,tiiuae/Falcon3-10B-Instruct-1.58bit}
|
||||
Model used for inference
|
||||
--model-dir MODEL_DIR, -md MODEL_DIR
|
||||
Directory to save/load the model
|
||||
--log-dir LOG_DIR, -ld LOG_DIR
|
||||
Directory to save the logging info
|
||||
--quant-type {i2_s,tl1}, -q {i2_s,tl1}
|
||||
Quantization type
|
||||
--quant-embd Quantize the embeddings to f16
|
||||
--use-pretuned, -p Use the pretuned kernel parameters
|
||||
</pre>
|
||||
## Usage
|
||||
### Basic usage
|
||||
```bash
|
||||
# Run inference with the quantized model
|
||||
python run_inference.py -m models/BitNet-b1.58-2B-4T/ggml-model-i2_s.gguf -p "You are a helpful assistant" -cnv
|
||||
```
|
||||
<pre>
|
||||
usage: run_inference.py [-h] [-m MODEL] [-n N_PREDICT] -p PROMPT [-t THREADS] [-c CTX_SIZE] [-temp TEMPERATURE] [-cnv]
|
||||
|
||||
Run inference
|
||||
|
||||
optional arguments:
|
||||
-h, --help show this help message and exit
|
||||
-m MODEL, --model MODEL
|
||||
Path to model file
|
||||
-n N_PREDICT, --n-predict N_PREDICT
|
||||
Number of tokens to predict when generating text
|
||||
-p PROMPT, --prompt PROMPT
|
||||
Prompt to generate text from
|
||||
-t THREADS, --threads THREADS
|
||||
Number of threads to use
|
||||
-c CTX_SIZE, --ctx-size CTX_SIZE
|
||||
Size of the prompt context
|
||||
-temp TEMPERATURE, --temperature TEMPERATURE
|
||||
Temperature, a hyperparameter that controls the randomness of the generated text
|
||||
-cnv, --conversation Whether to enable chat mode or not (for instruct models.)
|
||||
(When this option is turned on, the prompt specified by -p will be used as the system prompt.)
|
||||
</pre>
|
||||
|
||||
### Benchmark
|
||||
We provide scripts to run the inference benchmark providing a model.
|
||||
|
||||
```
|
||||
usage: e2e_benchmark.py -m MODEL [-n N_TOKEN] [-p N_PROMPT] [-t THREADS]
|
||||
|
||||
Setup the environment for running the inference
|
||||
|
||||
required arguments:
|
||||
-m MODEL, --model MODEL
|
||||
Path to the model file.
|
||||
|
||||
optional arguments:
|
||||
-h, --help
|
||||
Show this help message and exit.
|
||||
-n N_TOKEN, --n-token N_TOKEN
|
||||
Number of generated tokens.
|
||||
-p N_PROMPT, --n-prompt N_PROMPT
|
||||
Prompt to generate text from.
|
||||
-t THREADS, --threads THREADS
|
||||
Number of threads to use.
|
||||
```
|
||||
|
||||
Here's a brief explanation of each argument:
|
||||
|
||||
- `-m`, `--model`: The path to the model file. This is a required argument that must be provided when running the script.
|
||||
- `-n`, `--n-token`: The number of tokens to generate during the inference. It is an optional argument with a default value of 128.
|
||||
- `-p`, `--n-prompt`: The number of prompt tokens to use for generating text. This is an optional argument with a default value of 512.
|
||||
- `-t`, `--threads`: The number of threads to use for running the inference. It is an optional argument with a default value of 2.
|
||||
- `-h`, `--help`: Show the help message and exit. Use this argument to display usage information.
|
||||
|
||||
For example:
|
||||
|
||||
```sh
|
||||
python utils/e2e_benchmark.py -m /path/to/model -n 200 -p 256 -t 4
|
||||
```
|
||||
|
||||
This command would run the inference benchmark using the model located at `/path/to/model`, generating 200 tokens from a 256 token prompt, utilizing 4 threads.
|
||||
|
||||
For the model layout that do not supported by any public model, we provide scripts to generate a dummy model with the given model layout, and run the benchmark on your machine:
|
||||
|
||||
```bash
|
||||
python utils/generate-dummy-bitnet-model.py models/bitnet_b1_58-large --outfile models/dummy-bitnet-125m.tl1.gguf --outtype tl1 --model-size 125M
|
||||
|
||||
# Run benchmark with the generated model, use -m to specify the model path, -p to specify the prompt processed, -n to specify the number of token to generate
|
||||
python utils/e2e_benchmark.py -m models/dummy-bitnet-125m.tl1.gguf -p 512 -n 128
|
||||
```
|
||||
|
||||
### Convert from `.safetensors` Checkpoints
|
||||
|
||||
```sh
|
||||
# Prepare the .safetensors model file
|
||||
huggingface-cli download microsoft/bitnet-b1.58-2B-4T-bf16 --local-dir ./models/bitnet-b1.58-2B-4T-bf16
|
||||
|
||||
# Convert to gguf model
|
||||
python ./utils/convert-helper-bitnet.py ./models/bitnet-b1.58-2B-4T-bf16
|
||||
```
|
||||
|
||||
### FAQ (Frequently Asked Questions)📌
|
||||
|
||||
#### Q1: The build dies with errors building llama.cpp due to issues with std::chrono in log.cpp?
|
||||
|
||||
**A:**
|
||||
This is an issue introduced in recent version of llama.cpp. Please refer to this [commit](https://github.com/tinglou/llama.cpp/commit/4e3db1e3d78cc1bcd22bcb3af54bd2a4628dd323) in the [discussion](https://github.com/abetlen/llama-cpp-python/issues/1942) to fix this issue.
|
||||
|
||||
#### Q2: How to build with clang in conda environment on windows?
|
||||
|
||||
**A:**
|
||||
Before building the project, verify your clang installation and access to Visual Studio tools by running:
|
||||
```
|
||||
clang -v
|
||||
```
|
||||
|
||||
This command checks that you are using the correct version of clang and that the Visual Studio tools are available. If you see an error message such as:
|
||||
```
|
||||
'clang' is not recognized as an internal or external command, operable program or batch file.
|
||||
```
|
||||
|
||||
It indicates that your command line window is not properly initialized for Visual Studio tools.
|
||||
|
||||
• If you are using Command Prompt, run:
|
||||
```
|
||||
"C:\Program Files\Microsoft Visual Studio\2022\Professional\Common7\Tools\VsDevCmd.bat" -startdir=none -arch=x64 -host_arch=x64
|
||||
```
|
||||
|
||||
• If you are using Windows PowerShell, run the following commands:
|
||||
```
|
||||
Import-Module "C:\Program Files\Microsoft Visual Studio\2022\Professional\Common7\Tools\Microsoft.VisualStudio.DevShell.dll" Enter-VsDevShell 3f0e31ad -SkipAutomaticLocation -DevCmdArguments "-arch=x64 -host_arch=x64"
|
||||
```
|
||||
|
||||
These steps will initialize your environment and allow you to use the correct Visual Studio tools.
|
||||
@@ -0,0 +1,7 @@
|
||||
# WeHub 来源说明
|
||||
|
||||
- 原始项目：`microsoft/BitNet`
|
||||
- 原始仓库：https://github.com/microsoft/BitNet
|
||||
- 导入方式：上游默认分支的最新快照
|
||||
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
|
||||
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
<!-- BEGIN MICROSOFT SECURITY.MD V0.0.9 BLOCK -->
|
||||
|
||||
## Security
|
||||
|
||||
Microsoft takes the security of our software products and services seriously, which includes all source code repositories managed through our GitHub organizations, which include [Microsoft](https://github.com/Microsoft), [Azure](https://github.com/Azure), [DotNet](https://github.com/dotnet), [AspNet](https://github.com/aspnet) and [Xamarin](https://github.com/xamarin).
|
||||
|
||||
If you believe you have found a security vulnerability in any Microsoft-owned repository that meets [Microsoft's definition of a security vulnerability](https://aka.ms/security.md/definition), please report it to us as described below.
|
||||
|
||||
## Reporting Security Issues
|
||||
|
||||
**Please do not report security vulnerabilities through public GitHub issues.**
|
||||
|
||||
Instead, please report them to the Microsoft Security Response Center (MSRC) at [https://msrc.microsoft.com/create-report](https://aka.ms/security.md/msrc/create-report).
|
||||
|
||||
If you prefer to submit without logging in, send email to [secure@microsoft.com](mailto:secure@microsoft.com). If possible, encrypt your message with our PGP key; please download it from the [Microsoft Security Response Center PGP Key page](https://aka.ms/security.md/msrc/pgp).
|
||||
|
||||
You should receive a response within 24 hours. If for some reason you do not, please follow up via email to ensure we received your original message. Additional information can be found at [microsoft.com/msrc](https://www.microsoft.com/msrc).
|
||||
|
||||
Please include the requested information listed below (as much as you can provide) to help us better understand the nature and scope of the possible issue:
|
||||
|
||||
* Type of issue (e.g. buffer overflow, SQL injection, cross-site scripting, etc.)
|
||||
* Full paths of source file(s) related to the manifestation of the issue
|
||||
* The location of the affected source code (tag/branch/commit or direct URL)
|
||||
* Any special configuration required to reproduce the issue
|
||||
* Step-by-step instructions to reproduce the issue
|
||||
* Proof-of-concept or exploit code (if possible)
|
||||
* Impact of the issue, including how an attacker might exploit the issue
|
||||
|
||||
This information will help us triage your report more quickly.
|
||||
|
||||
If you are reporting for a bug bounty, more complete reports can contribute to a higher bounty award. Please visit our [Microsoft Bug Bounty Program](https://aka.ms/security.md/msrc/bounty) page for more details about our active programs.
|
||||
|
||||
## Preferred Languages
|
||||
|
||||
We prefer all communications to be in English.
|
||||
|
||||
## Policy
|
||||
|
||||
Microsoft follows the principle of [Coordinated Vulnerability Disclosure](https://aka.ms/security.md/cvd).
|
||||
|
||||
<!-- END MICROSOFT SECURITY.MD BLOCK -->
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 14 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 13 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 16 KiB |
@@ -0,0 +1,49 @@
|
||||
Codegen for TL1 and TL2
|
||||
------------------------
|
||||
|
||||
codegen_tl1.py and codegen_tl2.py are using params to generate kernel codes in different devices to achieve fastest performance for TL1 and TL2.
|
||||
|
||||
We cutting weight into multiple compute blocks to best utilize hardware capabilities.
|
||||
|
||||
### Example
|
||||
bitnet_b1_58-large:
|
||||
|
||||
- Make sure Matmul kernels shapes \
|
||||
For example, bitnet_b1_58-large Matmul kernel shapes are:\
|
||||
[1536, 4096]\
|
||||
[1536, 1536]\
|
||||
[4096, 1536]
|
||||
|
||||
- Make sure each BM, BK, bm for each kernel to meet the requirements below
|
||||
- Generate codes\
|
||||
For example, for bitnet_b1_58-large, we can gencode like:
|
||||
|
||||
```bash
|
||||
# For TL1
|
||||
python utils/codegen_tl1.py --model bitnet_b1_58-large --BM 256,128,256 --BK 128,64,128 --bm 32,64,32
|
||||
|
||||
# For TL2
|
||||
python utils/codegen_tl2.py --model bitnet_b1_58-large --BM 256,128,256 --BK 96,192,96 --bm 32,32,32
|
||||
```
|
||||
|
||||
### TL1:
|
||||
|
||||
|
||||
For TL1, we cut weight into M / BM weights, each weight shape is (BM, K). Then we cut weight into K / BK weights, each weight shape is (BM, BK). As for (BM, BK) weight, we cut it the same way into (bm, compute_num / bm) compute blocks, and finish computing in it.
|
||||
|
||||
Thus, we need to make sure
|
||||
- M % BM == 0
|
||||
- K % BK == 0
|
||||
- BM % bm == 0
|
||||
- bm choose in [32, 64]
|
||||
|
||||
### TL2:
|
||||
|
||||
|
||||
For TL2, things got a little more complicated. Due to TL2 needs BK % 6 == 0, we need to split K into threeK and twoK, in which compute in TL2 for (M, threeK), compute in TL1 for (M, two_K).
|
||||
|
||||
Thus, we needs to make sure
|
||||
- M % BM == 0
|
||||
- K % BK % 32 == 0
|
||||
- BM % bm == 0
|
||||
- bm choose in \[32\]
|
||||
Executable
+107
@@ -0,0 +1,107 @@
|
||||
# BitNet Inference Kernel
|
||||
|
||||
This repository provides a highly efficient GEMV kernel implementation for the BitNet model, optimized for W2A8 inference — 2-bit weights and 8-bit activations. It is tailored for use with the [BitNet-b1.58-2B-4T](https://arxiv.org/abs/2504.12285) model.
|
||||
|
||||
## Features
|
||||
|
||||
- Support for W2A8 (2-bit weight × 8-bit activation) GEMV computation
|
||||
- Custom CUDA kernels with low-latency execution
|
||||
- Optimizations for memory access, decoding, and compute throughput
|
||||
|
||||
## Usage
|
||||
|
||||
Installation and kernel performance tests:
|
||||
|
||||
```bash
|
||||
# (Recommended) Create a new conda environment
|
||||
conda create --name bitnet-gpu "python<3.13"
|
||||
conda activate bitnet-gpu
|
||||
|
||||
# Install dependencies
|
||||
pip install -r requirements.txt
|
||||
|
||||
# Build the kernel
|
||||
cd bitnet_kernels
|
||||
bash compile.sh
|
||||
cd ..
|
||||
|
||||
# Run performance tests
|
||||
python test.py
|
||||
```
|
||||
|
||||
End-to-end inference:
|
||||
|
||||
```bash
|
||||
# Download and convert the BitNet-b1.58-2B model
|
||||
mkdir checkpoints
|
||||
huggingface-cli download microsoft/bitnet-b1.58-2B-4T-bf16 --local-dir ./checkpoints/bitnet-b1.58-2B-4T-bf16
|
||||
python ./convert_safetensors.py --safetensors_file ./checkpoints/bitnet-b1.58-2B-4T-bf16/model.safetensors --output checkpoints/model_state.pt --model_name 2B
|
||||
python ./convert_checkpoint.py --input ./checkpoints/model_state.pt
|
||||
rm ./checkpoints/model_state.pt
|
||||
|
||||
# Inference
|
||||
python3 ./generate.py ./checkpoints/ --interactive --chat_format
|
||||
```
|
||||
|
||||
## Optimizations
|
||||
|
||||
### Weight Permutation
|
||||
|
||||
The weight matrix is divided into 16×32 blocks to optimize memory access patterns.
|
||||
|
||||
Within each block, values are stored contiguously in memory and permuted to facilitate efficient access and processing.
|
||||
|
||||
See `convert_checkpoint.py` for details.
|
||||
|
||||
### Fast Decoding
|
||||
|
||||
Every 16 two-bit values are packed into a single 32-bit integer using the following interleaving pattern:
|
||||
```
|
||||
[0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15]
|
||||
```
|
||||
|
||||
This layout is designed to accelerate decoding by enabling efficient extraction of 4 values at a time into `int8`.
|
||||
|
||||
### `dp4a` Instruction
|
||||
|
||||
We use the `dp4a` instruction to accelerate low-precision dot product operations.
|
||||
|
||||
This instruction performs a dot product between two 4-element vectors (each stored in a 32-bit word as 8-bit integers) and accumulates the result into a 32-bit integer.
|
||||
|
||||
It significantly improves GEMV throughput when processing quantized weights and activations.
|
||||
|
||||
|
||||
## Performance
|
||||
|
||||
### Kernel Benchmarks
|
||||
|
||||
Tested on NVIDIA A100 40GB GPU, our custom W2A8 kernel shows significant speedups over standard BF16 implementations:
|
||||
|
||||
| Shape (N×K) | W2A8 Latency (us) | BF16 Latency (us) | Speedup Ratio |
|
||||
|---------------------|-------------------|-------------------|----------------------|
|
||||
| 2560 × 2560 | 13.32 | 18.32 | 1.38 |
|
||||
| 3840 × 2560 | 14.90 | 18.87 | 1.27 |
|
||||
| 13824 × 2560 | 18.75 | 59.51 | 3.17 |
|
||||
| 2560 × 6912 | 14.49 | 37.78 | 2.61 |
|
||||
| 3200 × 3200 | 14.61 | 19.08 | 1.31 |
|
||||
| 4800 × 3200 | 13.09 | 21.84 | 1.67 |
|
||||
| 3200 × 10240 | 19.64 | 60.79 | 3.10 |
|
||||
| 20480 × 3200 | 30.99 | 112.39 | 3.63 |
|
||||
|
||||
### End-to-End Generation Latency
|
||||
|
||||
Compared to a similarly-sized BF16 model (Gemma-2-2B using vLLM), BitNet-b1.58-2B with our kernel achieves consistent speedups across workloads:
|
||||
|
||||
| Input Length | Output Length | BF16 Latency (ms) | W2A8 Latency (ms) | Speedup Ratio |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| 64 | 16 | 187.64 | 57.40 | 3.27 |
|
||||
| 64 | 32 | 353.50 | 112.22 | 3.15 |
|
||||
| 64 | 64 | 683.23 | 221.08 | 3.09 |
|
||||
| 256 | 16 | 183.14 | 61.24 | 2.99 |
|
||||
| 256 | 32 | 353.14 | 115.47 | 3.06 |
|
||||
| 256 | 64 | 684.24 | 224.16 | 3.05 |
|
||||
| 512 | 16 | 208.99 | 68.06 | 3.07 |
|
||||
| 512 | 32 | 354.33 | 122.72 | 2.89 |
|
||||
| 512 | 64 | 709.65 | 231.82 | 3.06 |
|
||||
|
||||
*Note: Comparison uses equivalent-sized models (2B parameters) on NVIDIA A100 40GB GPU.*
|
||||
@@ -0,0 +1,37 @@
|
||||
#include "bitnet_kernels.h"
|
||||
|
||||
extern "C" void bitlinear_int8xint2(int8_t* input0, int8_t* input1, __nv_bfloat16* output0, __nv_bfloat16* s, __nv_bfloat16* ws, int M, int N, int K, cudaStream_t stream){
|
||||
if (M == 1 && N == 3840 && K == 2560){
|
||||
ladder_int8xint2_kernel<1, 3840, 2560, 3, 8, 16><<<dim3(240, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if (M == 1 && N == 2560 && K == 2560){
|
||||
ladder_int8xint2_kernel<1, 2560, 2560, 1, 8, 16><<<dim3(160, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if (M == 1 && N == 13824 && K == 2560){
|
||||
ladder_int8xint2_kernel<1, 13824, 2560, 2, 8, 16><<<dim3(864, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if (M == 1 && N == 2560 && K == 6912){
|
||||
ladder_int8xint2_kernel<1, 2560, 6912, 1, 8, 16><<<dim3(160, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if(M == 1 && N == 4800 && K == 3200){
|
||||
ladder_int8xint2_kernel<1, 4800, 3200, 6, 8, 16><<<dim3(300, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if(M == 1 && N == 3200 && K == 3200){
|
||||
ladder_int8xint2_kernel<1, 3200, 3200, 1, 8, 16><<<dim3(200, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if(M == 1 && N == 20480 && K == 3200){
|
||||
ladder_int8xint2_kernel<1, 20480, 3200, 2, 8, 16><<<dim3(1280, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if(M == 1 && N == 3200 && K == 10240){
|
||||
ladder_int8xint2_kernel<1, 3200, 10240, 1, 8, 16><<<dim3(200, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if(M == 1 && N == 5120 && K == 27648){
|
||||
ladder_int8xint2_kernel<1, 5120, 27648, 1, 8, 16><<<dim3(320, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if(M == 1 && N == 55296 && K == 5120){
|
||||
ladder_int8xint2_kernel<1, 55296, 5120, 1, 8, 16><<<dim3(3456, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else{
|
||||
std::cout << "required ladder gemm kernel: M " << M << ", N " << N << ", K " << K << std::endl;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
#include <cuda_runtime.h>
|
||||
#include <math_constants.h>
|
||||
#include <math.h>
|
||||
#include <mma.h>
|
||||
#include <iostream>
|
||||
#include <cuda.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_bf16.h>
|
||||
|
||||
|
||||
#if (((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 4)) || (__CUDACC_VER_MAJOR__ > 11))
|
||||
#define TVM_ENABLE_L2_PREFETCH 1
|
||||
#else
|
||||
#define TVM_ENABLE_L2_PREFETCH 0
|
||||
#endif
|
||||
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 800
|
||||
#define TVM_ENBALE_EFFICIENT_SMEM_PTR_CAST 1
|
||||
#else
|
||||
#define TVM_ENBALE_EFFICIENT_SMEM_PTR_CAST 0
|
||||
#endif
|
||||
|
||||
template <typename T1, typename T2>
|
||||
__device__ void decode_i2s_to_i8s(T1 *_i2s, T2 *_i8s, const int N = 16)
|
||||
{
|
||||
// convert 8 int2b_t to 8 int8b_t -> 2 int32
|
||||
uint *i8s = reinterpret_cast<uint *>(_i8s);
|
||||
|
||||
// i2s = {e0, e4, e8, e12, e1, e5, e9, e13, e2, e6, e10, e14, e3, e7, e11, e15}
|
||||
uint const i2s = *_i2s;
|
||||
|
||||
static constexpr uint immLut = (0xf0 & 0xcc) | 0xaa; // 0b11101010
|
||||
static constexpr uint BOTTOM_MASK = 0x03030303; // 0xf -> 0b11 select 0,3
|
||||
static constexpr uint I4s_TO_I8s_MAGIC_NUM = 0x00000000;
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < (N / 4); i++)
|
||||
{
|
||||
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
|
||||
: "=r"(i8s[i])
|
||||
: "r"(i2s >> (2 * i)), "n"(BOTTOM_MASK), "n"(I4s_TO_I8s_MAGIC_NUM), "n"(immLut));
|
||||
i8s[i] = __vsubss4(i8s[i], 0x02020202);
|
||||
}
|
||||
}
|
||||
|
||||
template <int M, int N, int K, int ws_num, int K_block_size, int N_block_size>
|
||||
__global__ void __launch_bounds__(128) ladder_int8xint2_kernel(int8_t* __restrict__ A, int8_t* __restrict__ B, __nv_bfloat16* __restrict__ dtype_transform, __nv_bfloat16* __restrict__ s, __nv_bfloat16* __restrict__ ws) {
|
||||
constexpr int K_per_loop = 16;
|
||||
constexpr int wmma_K = 32;
|
||||
constexpr int wmma_N = 16;
|
||||
int in_thread_C_local[1];
|
||||
signed char A_local[K_per_loop];
|
||||
int B_reshape_local[1];
|
||||
signed char B_decode_local[K_per_loop];
|
||||
int red_buf0[1];
|
||||
in_thread_C_local[0] = 0;
|
||||
#pragma unroll
|
||||
for (int k_0 = 0; k_0 < K/(K_per_loop * K_block_size); ++k_0) {
|
||||
*(int4*)(A_local + 0) = *(int4*)(A + ((k_0 * K_per_loop * K_block_size) + (((int)threadIdx.x) * K_per_loop)));
|
||||
B_reshape_local[0] = *(int*)(B +
|
||||
(((int)blockIdx.x) * N_block_size * K / 4) +
|
||||
(k_0 * K_block_size * K_per_loop * wmma_N / 4) +
|
||||
((((int)threadIdx.x) >> 1) * wmma_K * wmma_N / 4) +
|
||||
((((int)threadIdx.y) >> 3) * (wmma_K * wmma_N / 2) / 4) +
|
||||
((((int)threadIdx.x) & 1) * (wmma_K * wmma_N / 4) / 4) +
|
||||
((((int)threadIdx.y) & 7) * (wmma_K / 2) / 4)
|
||||
);
|
||||
decode_i2s_to_i8s(B_reshape_local, B_decode_local, 16);
|
||||
#pragma unroll
|
||||
for (int k_2_0 = 0; k_2_0 < 4; ++k_2_0) {
|
||||
in_thread_C_local[0] = __dp4a(*(int *)&A_local[((k_2_0 * 4))],*(int *)&B_decode_local[((k_2_0 * 4))], in_thread_C_local[0]);
|
||||
}
|
||||
}
|
||||
red_buf0[0] = in_thread_C_local[0];
|
||||
#pragma unroll
|
||||
for (int offset = K_block_size/2; offset > 0; offset /= 2) {
|
||||
red_buf0[0] += __shfl_down_sync(__activemask(), red_buf0[0], offset, K_block_size);
|
||||
}
|
||||
int out_idx = ((((int)blockIdx.x) * N_block_size) + ((int)threadIdx.y));
|
||||
int ws_idx = out_idx / (N / ws_num);
|
||||
if (threadIdx.x == 0)
|
||||
dtype_transform[out_idx] = (__nv_bfloat16)(((float)red_buf0[0])/(float)s[0]*(float)ws[ws_idx]);
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
nvcc -std=c++17 -Xcudafe --diag_suppress=177 --compiler-options -fPIC -lineinfo --shared bitnet_kernels.cu -lcuda -gencode=arch=compute_80,code=compute_80 -o libbitnet.so
|
||||
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
from setuptools import setup
|
||||
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
||||
|
||||
setup(
|
||||
name='bitlinear_cpp',
|
||||
ext_modules=[
|
||||
CUDAExtension('bitlinear_cuda', [
|
||||
'bitnet_kernels.cu',
|
||||
])
|
||||
],
|
||||
cmdclass={
|
||||
'build_ext': BuildExtension
|
||||
})
|
||||
Executable
+100
@@ -0,0 +1,100 @@
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from dataclasses import dataclass
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from safetensors.torch import save_file
|
||||
import model
|
||||
from pack_weight import convert_weight_int8_to_int2
|
||||
|
||||
@torch.inference_mode()
|
||||
def convert_ts_checkpoint(
|
||||
*,
|
||||
input_path: str = "",
|
||||
) -> None:
|
||||
|
||||
config = model.ModelArgs()
|
||||
print(f"Model config {config.__dict__}")
|
||||
|
||||
def quant_weight_int8(weight):
|
||||
s = 1.0 / weight.abs().mean().clamp_(min=1e-5)
|
||||
new_weight = (weight * s).round().clamp(-1, 1).to(torch.int8)
|
||||
new_scale = (1.0 / s).to(torch.bfloat16)
|
||||
return new_weight, new_scale.reshape(1)
|
||||
|
||||
def quant_weight_fp16(weight):
|
||||
s = 1.0 / weight.abs().mean().clamp_(min=1e-5)
|
||||
new_weight = (weight * s).round().clamp(-1, 1) / s
|
||||
return new_weight
|
||||
|
||||
def convert_int8_to_int2(weight):
|
||||
return convert_weight_int8_to_int2(weight)
|
||||
|
||||
merged_result = torch.load(input_path, map_location="cpu", mmap=True, weights_only=True)
|
||||
int2_result = {}
|
||||
fp16_result = {}
|
||||
zero = torch.zeros(1).to(torch.bfloat16)
|
||||
for key, value in merged_result.items():
|
||||
if 'wqkv' in key:
|
||||
wq = value[:config.dim]
|
||||
wk = value[config.dim:config.dim // config.n_heads * config.n_kv_heads + config.dim]
|
||||
wv = value[config.dim // config.n_heads * config.n_kv_heads + config.dim:]
|
||||
wq_weight, wa_scale = quant_weight_int8(wq)
|
||||
wk_weight, wb_scale = quant_weight_int8(wk)
|
||||
wv_weight, wc_scale = quant_weight_int8(wv)
|
||||
wqkv_weight = torch.cat([wq_weight, wk_weight, wv_weight], dim=0)
|
||||
wqkv_scale = torch.cat([wa_scale, wb_scale, wc_scale, zero], dim=0)
|
||||
int2_result[key] = convert_int8_to_int2(wqkv_weight)
|
||||
int2_result[key.replace('weight', 'weight_scale')] = wqkv_scale
|
||||
|
||||
wq_weight = quant_weight_fp16(wq)
|
||||
wk_weight = quant_weight_fp16(wk)
|
||||
wv_weight = quant_weight_fp16(wv)
|
||||
wqkv_weight = torch.cat([wq_weight, wk_weight, wv_weight], dim=0)
|
||||
fp16_result[key] = wqkv_weight
|
||||
elif 'w13' in key:
|
||||
w1 = value[:config.ffn_dim]
|
||||
w3 = value[config.ffn_dim:]
|
||||
w1_weight, w1_scale = quant_weight_int8(w1)
|
||||
w3_weight, w3_scale = quant_weight_int8(w3)
|
||||
w13_weight = torch.cat([w1_weight, w3_weight], dim=0)
|
||||
w13_scale = torch.cat([w1_scale, w3_scale, zero, zero], dim=0)
|
||||
int2_result[key] = convert_int8_to_int2(w13_weight)
|
||||
int2_result[key.replace('weight', 'weight_scale')] = w13_scale
|
||||
|
||||
w1_weight = quant_weight_fp16(w1)
|
||||
w3_weight = quant_weight_fp16(w3)
|
||||
w13_weight = torch.cat([w1_weight, w3_weight], dim=0)
|
||||
fp16_result[key] = w13_weight
|
||||
elif 'w2' in key or 'wo' in key:
|
||||
weight, scale = quant_weight_int8(value)
|
||||
scale = torch.cat([scale, zero, zero, zero], dim=0)
|
||||
int2_result[key] = convert_int8_to_int2(weight)
|
||||
int2_result[key.replace('weight', 'weight_scale')] = scale
|
||||
|
||||
weight = quant_weight_fp16(value)
|
||||
fp16_result[key] = weight
|
||||
else:
|
||||
int2_result[key] = value.clone()
|
||||
fp16_result[key] = value.clone()
|
||||
|
||||
output_dir = os.path.dirname(input_path)
|
||||
print(f"Saving checkpoint to {output_dir}/model_state_int2.pt")
|
||||
torch.save(int2_result, f"{output_dir}/model_state_int2.pt")
|
||||
|
||||
print(f"Saving checkpoint to {output_dir}/model_state_fp16.pt")
|
||||
torch.save(fp16_result, f"{output_dir}/model_state_fp16.pt")
|
||||
|
||||
if __name__ == '__main__':
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser(description='Convert TorchScale checkpoint.')
|
||||
parser.add_argument('--input', type=str)
|
||||
|
||||
args = parser.parse_args()
|
||||
convert_ts_checkpoint(
|
||||
input_path=args.input,
|
||||
)
|
||||
@@ -0,0 +1,116 @@
|
||||
import re
|
||||
import torch
|
||||
from pathlib import Path
|
||||
from safetensors.torch import load_file
|
||||
from einops import rearrange
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
transformer_configs = {
|
||||
"2B": dict(n_layer=30, n_head=20, dim=2560, vocab_size=128256, n_local_heads=5, intermediate_size=6912),
|
||||
}
|
||||
|
||||
@dataclass
|
||||
class ModelArgs:
|
||||
block_size: int = 4096
|
||||
vocab_size: int = 32000
|
||||
n_layer: int = 32
|
||||
n_head: int = 32
|
||||
dim: int = 4096
|
||||
intermediate_size: int = None
|
||||
n_local_heads: int = -1
|
||||
head_dim: int = 64
|
||||
rope_base: float = 10000
|
||||
norm_eps: float = 1e-5
|
||||
|
||||
def __post_init__(self):
|
||||
if self.n_local_heads == -1:
|
||||
self.n_local_heads = self.n_head
|
||||
if self.intermediate_size is None:
|
||||
hidden_dim = 4 * self.dim
|
||||
n_hidden = int(2 * hidden_dim / 3)
|
||||
self.intermediate_size = n_hidden + (256 - n_hidden % 256) if n_hidden % 256 else n_hidden
|
||||
self.head_dim = self.dim // self.n_head
|
||||
|
||||
@classmethod
|
||||
def from_name(cls, name: str):
|
||||
if name in transformer_configs:
|
||||
return cls(**transformer_configs[name])
|
||||
config = [k for k in transformer_configs if k in name.upper() or k in name]
|
||||
assert len(config) == 1, f"Unknown model name: {name}"
|
||||
return cls(**transformer_configs[config[0]])
|
||||
|
||||
def invert_convert_q(w: torch.Tensor, config: ModelArgs) -> torch.Tensor:
|
||||
return rearrange(w, '(h l d) i -> (h d l) i', h=config.n_head, l=2)
|
||||
|
||||
def invert_convert_k(w: torch.Tensor, config: ModelArgs) -> torch.Tensor:
|
||||
return rearrange(w, '(h l d) i -> (h d l) i', h=config.n_local_heads, l=2)
|
||||
|
||||
def convert_back(
|
||||
safetensors_path: str,
|
||||
output_file: str,
|
||||
model_name: Optional[str] = None,
|
||||
):
|
||||
st_dict = load_file(safetensors_path)
|
||||
|
||||
cfg = ModelArgs.from_name(model_name)
|
||||
print(f"Using model configurations: {cfg}")
|
||||
|
||||
recovered: dict = {}
|
||||
|
||||
for layer in range(cfg.n_layer):
|
||||
base = f"model.layers.{layer}."
|
||||
|
||||
wq = st_dict[f"{base}self_attn.q_proj.weight"]
|
||||
wk = st_dict[f"{base}self_attn.k_proj.weight"]
|
||||
wv = st_dict[f"{base}self_attn.v_proj.weight"]
|
||||
|
||||
wq = invert_convert_q(wq, cfg)
|
||||
wk = invert_convert_k(wk, cfg)
|
||||
|
||||
wqkv = torch.cat([wq, wk, wv], dim=0)
|
||||
recovered[f"layers.{layer}.attention.wqkv.weight"] = wqkv
|
||||
|
||||
recovered[f"layers.{layer}.attention.wo.weight"] = st_dict[f"{base}self_attn.o_proj.weight"]
|
||||
|
||||
recovered[f"layers.{layer}.attention_norm.weight"] = st_dict[f"{base}input_layernorm.weight"]
|
||||
recovered[f"layers.{layer}.ffn_norm.weight"] = st_dict[f"{base}post_attention_layernorm.weight"]
|
||||
recovered[f"layers.{layer}.attention.attn_sub_norm.weight"] = st_dict[f"{base}self_attn.attn_sub_norm.weight"]
|
||||
recovered[f"layers.{layer}.feed_forward.ffn_sub_norm.weight"] = st_dict[f"{base}mlp.ffn_sub_norm.weight"]
|
||||
|
||||
gate = st_dict[f"{base}mlp.gate_proj.weight"]
|
||||
up = st_dict[f"{base}mlp.up_proj.weight"]
|
||||
w13 = torch.cat([gate, up], dim=0)
|
||||
recovered[f"layers.{layer}.feed_forward.w13.weight"] = w13
|
||||
|
||||
recovered[f"layers.{layer}.feed_forward.w2.weight"] = st_dict[f"{base}mlp.down_proj.weight"]
|
||||
|
||||
recovered["tok_embeddings.weight"] = st_dict["model.embed_tokens.weight"]
|
||||
recovered["output.weight"] = st_dict["model.embed_tokens.weight"]
|
||||
recovered["norm.weight"] = st_dict["model.norm.weight"]
|
||||
|
||||
print(f"Saving to {output_file}")
|
||||
torch.save(recovered, output_file)
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser(description="Convert Safetensors back to Torch .pth checkpoint")
|
||||
parser.add_argument(
|
||||
"--safetensors_file", type=str, required=True,
|
||||
help="Path to input .safetensors file"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output", type=str, default="./checkpoints/model_state.pt",
|
||||
help="Path to output .pt file"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model_name", type=str, default="2B",
|
||||
help="Model configuration name to use (e.g. 2B)"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
convert_back(
|
||||
safetensors_path=args.safetensors_file,
|
||||
output_file=args.output,
|
||||
model_name=args.model_name,
|
||||
)
|
||||
Executable
+359
@@ -0,0 +1,359 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the BSD license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import json
|
||||
import os
|
||||
import readline # type: ignore # noqa
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Iterable, Optional, Tuple, Union
|
||||
|
||||
import fire
|
||||
import model as fast
|
||||
import torch
|
||||
from stats import Stats
|
||||
from tokenizer import Tokenizer, ChatFormat
|
||||
import sample_utils
|
||||
from xformers.ops.fmha.attn_bias import (
|
||||
BlockDiagonalCausalWithOffsetPaddedKeysMask as AttnBias,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenArgs:
|
||||
gen_length: int = 32
|
||||
gen_bsz: int = 1
|
||||
prompt_length: int = 64
|
||||
|
||||
use_sampling: bool = False
|
||||
temperature: float = 0.8
|
||||
top_p: float = 0.9
|
||||
|
||||
|
||||
class FastGen:
|
||||
GRAPH_WARMUPS: int = 1
|
||||
tokenizer: Tokenizer
|
||||
|
||||
@staticmethod
|
||||
def build(
|
||||
ckpt_dir: str,
|
||||
gen_args: GenArgs,
|
||||
device: Union[torch.device, str],
|
||||
tokenizer_path: Optional[str] = None,
|
||||
num_layers: int = 13,
|
||||
use_full_vocab: bool = False,
|
||||
) -> "FastGen":
|
||||
"""
|
||||
Load a Llama or Code Llama checkpoint and return a new
|
||||
generator for this model.
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
model_args_prefill = fast.ModelArgs(use_kernel=False)
|
||||
model_args_decode = fast.ModelArgs(use_kernel=True)
|
||||
tokenizer = Tokenizer("./tokenizer.model")
|
||||
|
||||
torch.set_default_device(device)
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
|
||||
prefill_model = fast.Transformer(model_args_prefill)
|
||||
decode_model = fast.Transformer(model_args_decode)
|
||||
|
||||
fp16_ckpt_path = str(Path(ckpt_dir) / "model_state_fp16.pt")
|
||||
fp16_checkpoint = torch.load(fp16_ckpt_path, map_location="cpu", weights_only=True)
|
||||
int2_ckpt_path = str(Path(ckpt_dir) / "model_state_int2.pt")
|
||||
int2_checkpoint = torch.load(int2_ckpt_path, map_location="cpu", weights_only=True)
|
||||
prefill_model.load_state_dict(fp16_checkpoint, strict=True)
|
||||
decode_model.load_state_dict(int2_checkpoint, strict=True)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
print(f"loaded model in {time.time() - start_time:.2f} seconds")
|
||||
start_time = time.time()
|
||||
|
||||
return FastGen(gen_args, model_args_prefill, prefill_model, decode_model, tokenizer)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
args: GenArgs,
|
||||
model_args: fast.ModelArgs,
|
||||
prefill_model: fast.Transformer,
|
||||
decode_model: fast.Transformer,
|
||||
tokenizer: Tokenizer,
|
||||
):
|
||||
self.gen_args = args
|
||||
self.max_seq_length = args.prompt_length + args.gen_length
|
||||
self.model_args = model_args
|
||||
# self.model = model
|
||||
self.prefill_model = prefill_model
|
||||
self.decode_model = decode_model
|
||||
self.tokenizer = tokenizer
|
||||
self._prefill_cuda_graph, self._prefill_compile_model, self._prefill_inputs, self._prefill_logits = None, None, None, None
|
||||
self._generate_cuda_graph, self._generate_compile_model, self._generate_inputs, self._generate_logits = None, None, None, None
|
||||
self._cache = None
|
||||
start_time = time.time()
|
||||
self._prefill_compile_model = self.compile_prefill()
|
||||
self._generate_compile_model = self.compile_generate()
|
||||
print(f"compiled model in {time.time() - start_time:.2f} seconds")
|
||||
|
||||
def compile_prefill(self):
|
||||
|
||||
if self._cache is None:
|
||||
self._cache = fast.make_cache(
|
||||
args=self.model_args,
|
||||
length=self.gen_args.gen_bsz * self.max_seq_length,
|
||||
)
|
||||
|
||||
seq_lens = [self.gen_args.prompt_length for _ in range(self.gen_args.gen_bsz)]
|
||||
|
||||
bias = AttnBias.from_seqlens(
|
||||
q_seqlen=seq_lens,
|
||||
kv_seqlen=seq_lens,
|
||||
kv_padding=self.max_seq_length,
|
||||
)
|
||||
bias.q_seqinfo.to("cuda")
|
||||
bias.k_seqinfo.to("cuda")
|
||||
|
||||
tokens = torch.IntTensor([1] * self.gen_args.gen_bsz * self.gen_args.prompt_length).cuda()
|
||||
self._prefill_inputs = (tokens, bias)
|
||||
|
||||
s = torch.cuda.Stream()
|
||||
s.wait_stream(torch.cuda.current_stream())
|
||||
|
||||
with torch.cuda.stream(s):
|
||||
_ = self.prefill_model.forward_with_attn_bias(
|
||||
token_values=self._prefill_inputs[0],
|
||||
attn_bias=self._prefill_inputs[1],
|
||||
cache=self._cache,
|
||||
)
|
||||
torch.cuda.current_stream().wait_stream(s)
|
||||
|
||||
self._prefill_cuda_graph = torch.cuda.CUDAGraph()
|
||||
recording_kwargs = {}
|
||||
if "capture_error_mode" in torch.cuda.graph.__init__.__annotations__:
|
||||
# In PyTorch 2.1+ and nightlies from late Aug 2023,
|
||||
# we can do this to maybe avoid watchdog-related crashes
|
||||
recording_kwargs["capture_error_mode"] = "thread_local"
|
||||
with torch.cuda.graph(self._prefill_cuda_graph, **recording_kwargs):
|
||||
self._prefill_logits = self.prefill_model.forward_with_attn_bias(
|
||||
token_values=self._prefill_inputs[0],
|
||||
attn_bias=self._prefill_inputs[1],
|
||||
cache=self._cache,
|
||||
)
|
||||
|
||||
def replay(tokens, seq_lens=None):
|
||||
self._prefill_inputs[0].copy_(tokens)
|
||||
if seq_lens is not None:
|
||||
self._prefill_inputs[1].k_seqinfo.seqlen.copy_(seq_lens)
|
||||
|
||||
self._prefill_cuda_graph.replay()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
return self._prefill_logits
|
||||
|
||||
return replay
|
||||
|
||||
def compile_generate(self):
|
||||
|
||||
if self._cache is None:
|
||||
self._cache = fast.make_cache(
|
||||
args=self.model_args,
|
||||
length=self.gen_args.gen_bsz * self.max_seq_length,
|
||||
)
|
||||
|
||||
seq_lens = [1 for _ in range(self.gen_args.gen_bsz)]
|
||||
kv_seq_lens = [self.gen_args.prompt_length for _ in range(self.gen_args.gen_bsz)]
|
||||
|
||||
bias = AttnBias.from_seqlens(
|
||||
q_seqlen=seq_lens,
|
||||
kv_seqlen=kv_seq_lens,
|
||||
kv_padding=self.max_seq_length,
|
||||
)
|
||||
bias.q_seqinfo.to("cuda")
|
||||
bias.k_seqinfo.to("cuda")
|
||||
|
||||
tokens = torch.IntTensor([1] * self.gen_args.gen_bsz).cuda()
|
||||
self._generate_inputs = (tokens, bias)
|
||||
|
||||
s = torch.cuda.Stream()
|
||||
s.wait_stream(torch.cuda.current_stream())
|
||||
|
||||
with torch.cuda.stream(s):
|
||||
_ = self.decode_model.forward_with_attn_bias(
|
||||
token_values=self._generate_inputs[0],
|
||||
attn_bias=self._generate_inputs[1],
|
||||
cache=self._cache,
|
||||
)
|
||||
torch.cuda.current_stream().wait_stream(s)
|
||||
|
||||
self._generate_cuda_graph = torch.cuda.CUDAGraph()
|
||||
recording_kwargs = {}
|
||||
if "capture_error_mode" in torch.cuda.graph.__init__.__annotations__:
|
||||
# In PyTorch 2.1+ and nightlies from late Aug 2023,
|
||||
# we can do this to maybe avoid watchdog-related crashes
|
||||
recording_kwargs["capture_error_mode"] = "thread_local"
|
||||
with torch.cuda.graph(self._generate_cuda_graph, **recording_kwargs):
|
||||
self._generate_logits = self.decode_model.forward_with_attn_bias(
|
||||
token_values=self._generate_inputs[0],
|
||||
attn_bias=self._generate_inputs[1],
|
||||
cache=self._cache,
|
||||
)
|
||||
|
||||
def replay(tokens, seq_lens):
|
||||
self._generate_inputs[0].copy_(tokens)
|
||||
self._generate_inputs[1].k_seqinfo.seqlen.copy_(seq_lens)
|
||||
|
||||
self._generate_cuda_graph.replay()
|
||||
|
||||
return self._generate_logits
|
||||
|
||||
return replay
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate_all(
|
||||
self, prompts: list[list[int]], use_cuda_graphs: bool, use_sampling: bool
|
||||
) -> Tuple[Stats, list[list[int]]]:
|
||||
bs = len(prompts)
|
||||
prompt_lens = [len(p) for p in prompts]
|
||||
padded_prompt_lens = [self.gen_args.prompt_length] * bs
|
||||
max_prompt_length = max(prompt_lens)
|
||||
gen_length = self.gen_args.gen_length
|
||||
max_seq_length = max_prompt_length + gen_length
|
||||
print(max_prompt_length, gen_length)
|
||||
|
||||
bias = AttnBias.from_seqlens(
|
||||
q_seqlen=padded_prompt_lens,
|
||||
kv_seqlen=prompt_lens,
|
||||
kv_padding=max_seq_length,
|
||||
)
|
||||
bias.q_seqinfo.to("cuda")
|
||||
bias.k_seqinfo.to("cuda")
|
||||
|
||||
# Input tensors to the cuda graph
|
||||
kv_seqlen = bias.k_seqinfo.seqlen
|
||||
prompts = [prompt + [1] * (self.gen_args.prompt_length - len(prompt)) for prompt in prompts]
|
||||
tokens = torch.IntTensor(sum(prompts, [])).cuda()
|
||||
out_tokens = torch.zeros((max_seq_length, bs), dtype=torch.int)
|
||||
|
||||
stats = Stats()
|
||||
torch.cuda.synchronize()
|
||||
stats.phase("prefill" if use_cuda_graphs else "total")
|
||||
# stats.phase("total")
|
||||
|
||||
output = self._prefill_compile_model(tokens, None)
|
||||
|
||||
logits = output[kv_seqlen - 1, :]
|
||||
logits = logits.view(bs, self.model_args.vocab_size)
|
||||
|
||||
if use_sampling:
|
||||
temp = 0.7
|
||||
top_p = 0.95
|
||||
probs = torch.softmax(logits / temp, dim=-1)
|
||||
next_token = sample_utils.top_p(probs, top_p)
|
||||
else:
|
||||
next_token = torch.argmax(logits, dim=-1)
|
||||
|
||||
next_token = next_token.reshape(bs)
|
||||
out_tokens[0, :] = next_token
|
||||
|
||||
torch.cuda.synchronize()
|
||||
stats.phase("decode" if use_cuda_graphs else "total")
|
||||
|
||||
eos_id = self.tokenizer.eot_id
|
||||
for niter in range(1, gen_length):
|
||||
kv_seqlen.add_(kv_seqlen < max_seq_length)
|
||||
output = self._generate_compile_model(next_token, kv_seqlen)
|
||||
|
||||
logits = output.view(bs, self.model_args.vocab_size)
|
||||
|
||||
if use_sampling:
|
||||
temp = 0.7
|
||||
top_p = 0.95
|
||||
probs = torch.softmax(logits / temp, dim=-1)
|
||||
next_token = sample_utils.top_p(probs, top_p)
|
||||
else:
|
||||
next_token = torch.argmax(logits, dim=-1)
|
||||
|
||||
next_token = next_token.reshape(bs)
|
||||
out_tokens[niter, :] = next_token
|
||||
|
||||
if next_token.eq(eos_id).any():
|
||||
break
|
||||
|
||||
torch.cuda.synchronize()
|
||||
stats.end_phase(tokens=niter * bs)
|
||||
|
||||
def trim_answer(prompt_len, tokens):
|
||||
# print(prompt, tokens)
|
||||
"""Trim the answer to end it on an eos token."""
|
||||
tokens = tokens[: max_seq_length - prompt_len]
|
||||
eos_id = self.tokenizer.eot_id
|
||||
if eos_id in tokens:
|
||||
return tokens[: tokens.index(eos_id) + 1]
|
||||
else:
|
||||
return tokens
|
||||
|
||||
answers = [
|
||||
trim_answer(prompt_len, answer)
|
||||
for prompt_len, answer in zip(prompt_lens, out_tokens.t().tolist())
|
||||
]
|
||||
return stats, answers
|
||||
|
||||
|
||||
def get_prompts(interactive: bool) -> Iterable[list[str]]:
|
||||
if interactive:
|
||||
while True:
|
||||
try:
|
||||
prompts = input("enter prompt: ").split("\n")
|
||||
except EOFError:
|
||||
print("exiting")
|
||||
sys.exit(0)
|
||||
yield prompts
|
||||
else:
|
||||
yield [
|
||||
"Hello, my name is",
|
||||
]
|
||||
|
||||
|
||||
def main(ckpt_dir: str, interactive: bool = False, chat_format: bool = False, sampling: bool = False):
|
||||
|
||||
local_rank = 0
|
||||
device = f"cuda:{local_rank}"
|
||||
torch.cuda.set_device(local_rank)
|
||||
|
||||
g = FastGen.build(ckpt_dir, GenArgs(), device)
|
||||
|
||||
if chat_format:
|
||||
g.tokenizer = ChatFormat(g.tokenizer)
|
||||
|
||||
for prompts in get_prompts(interactive):
|
||||
# prompts = [f"{prompt}\n" for prompt in prompts]
|
||||
if chat_format:
|
||||
# prompts = [f'<|begin_of_text|>User: {prompt}<|eot_id|>Assistant: ' for prompt in prompts]
|
||||
tokens = [g.tokenizer.encode_dialog_prompt(dialog=[{"role": "user", "content": prompt}], completion=True) for prompt in prompts]
|
||||
else:
|
||||
tokens = [g.tokenizer.encode(x, bos=False, eos=False) for x in prompts]
|
||||
|
||||
print(tokens)
|
||||
stats, out_tokens = g.generate_all(
|
||||
tokens, use_cuda_graphs="NO_CUDA_GRAPHS" not in os.environ, use_sampling=sampling,
|
||||
)
|
||||
|
||||
for i, prompt in enumerate(prompts):
|
||||
print(f"> {prompt}")
|
||||
answer = g.tokenizer.decode(out_tokens[i])
|
||||
print(answer)
|
||||
print("---------------")
|
||||
|
||||
for phase_stats in stats.phases:
|
||||
print(phase_stats.show())
|
||||
|
||||
print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
fire.Fire(main)
|
||||
Executable
+366
@@ -0,0 +1,366 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the BSD license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from xformers.ops import RMSNorm, fmha, rope_padded
|
||||
from xformers.ops.fmha.attn_bias import (
|
||||
BlockDiagonalCausalWithOffsetPaddedKeysMask as AttnBias,
|
||||
)
|
||||
|
||||
import ctypes
|
||||
bitnet_lib = ctypes.CDLL('bitnet_kernels/libbitnet.so')
|
||||
|
||||
def bitnet_int8xint2_linear(input0, input1, s, ws):
|
||||
out_shape = list(input0.shape)
|
||||
out_shape[-1] = input1.shape[0]
|
||||
|
||||
stream = torch.cuda.current_stream()
|
||||
|
||||
M = input0.shape[0]
|
||||
if len(out_shape) == 3:
|
||||
M *= input0.shape[1]
|
||||
N = input1.shape[0]
|
||||
K = input1.shape[1] * 4
|
||||
|
||||
ret = torch.zeros(*out_shape, dtype=torch.bfloat16, device=input0.device)
|
||||
|
||||
bitnet_lib.bitlinear_int8xint2(*[ctypes.c_void_p(input0.data_ptr()), ctypes.c_void_p(input1.data_ptr()), ctypes.c_void_p(ret.data_ptr()), ctypes.c_void_p(s.data_ptr()), ctypes.c_void_p(ws.data_ptr()), ctypes.c_int(M), ctypes.c_int(N), ctypes.c_int(K), ctypes.c_void_p(stream.cuda_stream)])
|
||||
|
||||
return ret
|
||||
|
||||
@dataclass
|
||||
class ModelArgs:
|
||||
dim: int = 2560
|
||||
n_layers: int = 30
|
||||
n_heads: int = 20
|
||||
n_kv_heads: int = 5
|
||||
vocab_size: int = 128256
|
||||
ffn_dim: int = 6912
|
||||
norm_eps: float = 1e-5
|
||||
rope_theta: float = 500000.0
|
||||
use_kernel: bool = False
|
||||
|
||||
|
||||
LayerCache = Tuple[torch.Tensor, torch.Tensor]
|
||||
|
||||
class BitLinearKernel(nn.Module):
|
||||
in_features: int
|
||||
out_features: int
|
||||
weight: torch.Tensor
|
||||
weight_scale: torch.Tensor
|
||||
|
||||
def __init__(self, in_features: int, out_features: int, bias: bool = False):
|
||||
super().__init__()
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
|
||||
self.weight = torch.nn.Parameter(torch.zeros(out_features, in_features//4, dtype=torch.int8), requires_grad=False)
|
||||
self.weight_scale = torch.nn.Parameter(torch.zeros(4, dtype=torch.bfloat16), requires_grad=False)
|
||||
|
||||
@torch.compile
|
||||
def quant_input(self, input):
|
||||
s = 127 / input.abs().max(dim=-1, keepdim=True).values.clamp_(min=1e-5)
|
||||
return (input * s).round().clamp(-128, 127).to(torch.int8), s
|
||||
|
||||
def forward(self, input):
|
||||
input, s = self.quant_input(input)
|
||||
return bitnet_int8xint2_linear(input, self.weight, s, self.weight_scale)
|
||||
|
||||
class BitLinear(nn.Linear):
|
||||
@torch.compile
|
||||
def quant_input(self, input):
|
||||
s = 127 / input.abs().max(dim=-1, keepdim=True).values.clamp_(min=1e-5)
|
||||
return (input * s).round().clamp(-128, 127) / s
|
||||
|
||||
def forward(self, input):
|
||||
input = self.quant_input(input)
|
||||
return F.linear(input, self.weight)
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
head_dim: int,
|
||||
n_heads: int,
|
||||
n_kv_heads: int,
|
||||
rope_theta: float,
|
||||
norm_eps: float,
|
||||
use_kernel: bool,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.head_dim = head_dim
|
||||
self.rope_theta = rope_theta
|
||||
|
||||
self.n_local_heads = n_heads
|
||||
self.n_local_kv_heads = n_kv_heads
|
||||
|
||||
Linear = BitLinearKernel if use_kernel else BitLinear
|
||||
|
||||
self.wqkv = Linear(
|
||||
dim,
|
||||
(self.n_local_heads + 2 * self.n_local_kv_heads) * head_dim,
|
||||
bias=False,
|
||||
)
|
||||
self.wo = Linear(
|
||||
self.n_local_heads * head_dim,
|
||||
dim,
|
||||
bias=False,
|
||||
)
|
||||
|
||||
self.attn_sub_norm = RMSNorm(dim, norm_eps)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
cache: LayerCache,
|
||||
attn_bias: AttnBias,
|
||||
) -> torch.Tensor:
|
||||
|
||||
xqkv = self.wqkv(x)
|
||||
xq = xqkv[:, : (self.n_local_heads * self.head_dim)]
|
||||
xkv = xqkv[:, (self.n_local_heads * self.head_dim) :]
|
||||
xk, xv = xkv.chunk(2, 1)
|
||||
|
||||
output_shape = xq.shape
|
||||
heads_per_group = self.n_local_heads // self.n_local_kv_heads
|
||||
xq = xq.view(
|
||||
1, xq.shape[0], self.n_local_kv_heads, heads_per_group, self.head_dim
|
||||
)
|
||||
xk = xk.view(1, xk.shape[0], self.n_local_kv_heads, 1, self.head_dim)
|
||||
# xq = rearrange(xq, 'b (g h l d) -> 1 b h g (d l)', g=heads_per_group, h=self.n_local_kv_heads, d=self.head_dim // 2, l=2)
|
||||
# xk = rearrange(xk, 'b (g l d) -> 1 b g 1 (d l)', g=self.n_local_kv_heads, d=self.head_dim // 2)
|
||||
xv = xv.view(1, xv.shape[0], self.n_local_kv_heads, 1, self.head_dim)
|
||||
cache_k, cache_v = cache
|
||||
|
||||
xq = rope_padded(
|
||||
xq=xq,
|
||||
xk=xk,
|
||||
xv=xv,
|
||||
cache_k=cache_k,
|
||||
cache_v=cache_v,
|
||||
attn_bias=attn_bias,
|
||||
theta=self.rope_theta,
|
||||
)
|
||||
|
||||
output = fmha.memory_efficient_attention_forward(
|
||||
xq, cache_k, cache_v, attn_bias, op = fmha.flash.FwOp
|
||||
)
|
||||
|
||||
output = output.reshape(output_shape)
|
||||
output = self.attn_sub_norm(output)
|
||||
output = self.wo(output)
|
||||
|
||||
return output
|
||||
|
||||
@torch.compile
|
||||
def squared_relu(x: torch.Tensor) -> torch.Tensor:
|
||||
return F.relu(x) ** 2
|
||||
|
||||
class FeedForward(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
hidden_dim: int,
|
||||
norm_eps: float,
|
||||
use_kernel: bool,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
Linear = BitLinearKernel if use_kernel else BitLinear
|
||||
|
||||
self.w13 = Linear(
|
||||
dim,
|
||||
2 * hidden_dim,
|
||||
bias=False,
|
||||
)
|
||||
self.w2 = Linear(
|
||||
hidden_dim,
|
||||
dim,
|
||||
bias=False,
|
||||
)
|
||||
self.ffn_sub_norm = RMSNorm(hidden_dim, norm_eps)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x13 = self.w13(x)
|
||||
x1, x3 = x13.chunk(2, -1)
|
||||
inner = self.ffn_sub_norm(squared_relu(x1) * x3)
|
||||
output = self.w2(inner)
|
||||
return output
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
assert args.dim % args.n_heads == 0
|
||||
head_dim = args.dim // args.n_heads
|
||||
if args.n_kv_heads is not None:
|
||||
n_kv_heads = args.n_kv_heads
|
||||
else:
|
||||
n_kv_heads = args.n_heads
|
||||
|
||||
assert args.n_heads % n_kv_heads == 0
|
||||
|
||||
self.attention = Attention(
|
||||
dim=args.dim,
|
||||
head_dim=head_dim,
|
||||
n_heads=args.n_heads,
|
||||
n_kv_heads=n_kv_heads,
|
||||
rope_theta=args.rope_theta,
|
||||
norm_eps=args.norm_eps,
|
||||
use_kernel=args.use_kernel,
|
||||
)
|
||||
self.feed_forward = FeedForward(
|
||||
dim=args.dim,
|
||||
hidden_dim=args.ffn_dim,
|
||||
norm_eps=args.norm_eps,
|
||||
use_kernel=args.use_kernel,
|
||||
)
|
||||
self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
|
||||
self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
cache: LayerCache,
|
||||
attn_bias: AttnBias,
|
||||
) -> torch.Tensor:
|
||||
h = x + self.attention.forward(
|
||||
self.attention_norm(x),
|
||||
cache,
|
||||
attn_bias,
|
||||
)
|
||||
out = h + self.feed_forward(self.ffn_norm(h))
|
||||
return out
|
||||
|
||||
|
||||
class Transformer(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
assert args.vocab_size > 0
|
||||
|
||||
self.tok_embeddings = nn.Embedding(
|
||||
num_embeddings=args.vocab_size,
|
||||
embedding_dim=args.dim,
|
||||
)
|
||||
|
||||
self.layers = nn.ModuleList()
|
||||
for _ in range(args.n_layers):
|
||||
self.layers.append(TransformerBlock(args))
|
||||
|
||||
self.norm = RMSNorm(args.dim, eps=args.norm_eps)
|
||||
|
||||
self.output = nn.Linear(
|
||||
args.dim,
|
||||
args.vocab_size,
|
||||
bias=False,
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def forward_with_attn_bias(
|
||||
self,
|
||||
token_values: torch.Tensor,
|
||||
attn_bias: AttnBias,
|
||||
cache: list[LayerCache],
|
||||
) -> torch.Tensor:
|
||||
h = self.tok_embeddings(token_values)
|
||||
|
||||
for i, layer in enumerate(self.layers):
|
||||
h = layer(h, cache[i], attn_bias)
|
||||
|
||||
logits = self.output(self.norm(h))
|
||||
return logits.float()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
token_values: torch.Tensor,
|
||||
token_lengths: torch.Tensor,
|
||||
start_pos: torch.Tensor,
|
||||
cache: list[LayerCache],
|
||||
kv_padding: int,
|
||||
) -> torch.Tensor:
|
||||
attn_bias = AttnBias.from_seqlens(
|
||||
q_seqlen=token_lengths.tolist(),
|
||||
kv_seqlen=(start_pos + token_lengths).tolist(),
|
||||
kv_padding=kv_padding,
|
||||
)
|
||||
return self.forward_with_attn_bias(token_values, attn_bias, cache)
|
||||
|
||||
|
||||
def make_cache(
|
||||
args: ModelArgs,
|
||||
length: int,
|
||||
device: Optional[Union[str, torch.device]] = None,
|
||||
n_layers: Optional[int] = None,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
) -> list[LayerCache]:
|
||||
"""
|
||||
Allocate a cache to be used with the Transformer module.
|
||||
|
||||
Args:
|
||||
args (ModelArgs): the model configuration.
|
||||
length (int): per layer cache size.
|
||||
It is usually budgeted as ``max_batch * max_seq``
|
||||
device (torch.device, optional): the device on which
|
||||
the cache should be allocated.
|
||||
n_layers (int, optional): the number of layers to
|
||||
allocate a cache for (defaults to the model
|
||||
settings).
|
||||
dtype (torch.dtype, optional): the dtype to use for
|
||||
cache entries (defaults to the default dtype).
|
||||
|
||||
Returns:
|
||||
The cache object to pass to ``Tranformer.forward``.
|
||||
"""
|
||||
|
||||
head_dim = args.dim // args.n_heads
|
||||
n_kv_heads = args.n_kv_heads
|
||||
if n_kv_heads is None:
|
||||
n_kv_heads = args.n_heads
|
||||
n_local_kv_heads = n_kv_heads
|
||||
|
||||
if n_layers is None:
|
||||
n_layers = args.n_layers
|
||||
|
||||
shape = (1, length, n_local_kv_heads, 1, head_dim)
|
||||
heads_per_group = args.n_heads // n_kv_heads
|
||||
expansion = (-1, -1, -1, heads_per_group, -1)
|
||||
return [
|
||||
(
|
||||
torch.zeros(shape, device=device, dtype=dtype).expand(expansion),
|
||||
torch.zeros(shape, device=device, dtype=dtype).expand(expansion),
|
||||
)
|
||||
for _ in range(n_layers)
|
||||
]
|
||||
|
||||
|
||||
def cache_prefix(cache: list[LayerCache], length: int) -> list[LayerCache]:
|
||||
"""
|
||||
Take a prefix view of a larger cache.
|
||||
|
||||
The original cache object remains of identical size and valid
|
||||
after the shrinked alias has been used. This function is useful
|
||||
when a cache was allocated for a larger batch size than what is
|
||||
necessary.
|
||||
|
||||
Args:
|
||||
cache: the cache to take a view in.
|
||||
length (int): the desired length
|
||||
|
||||
Returns:
|
||||
A view in the input cache object.
|
||||
"""
|
||||
|
||||
if len(cache) > 0:
|
||||
assert cache[0][0].shape[1] >= length
|
||||
|
||||
return [(ck[:, :length], cv[:, :length]) for ck, cv in cache]
|
||||
Executable
+98
@@ -0,0 +1,98 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
def B_global_16x32_to_shared_load_16x32_layout(i, j):
|
||||
"""
|
||||
stride * 8 * (tx // HALF_WARP_expr)
|
||||
+ (tx % 8) * stride
|
||||
+ 16 * ((tx % HALF_WARP_expr) // 8)
|
||||
"""
|
||||
thread_id = i * 2 + j // 16
|
||||
row = (thread_id // 16) * 8 + (thread_id % 8)
|
||||
col = (j % 16) + 16 * ((thread_id % 16) // 8)
|
||||
return row, col
|
||||
|
||||
|
||||
def permutate_weight_fastest(weight):
|
||||
wmma_n = 16
|
||||
wmma_k = 32
|
||||
N = weight.shape[0]
|
||||
K = weight.shape[1]
|
||||
|
||||
# Create a lookup table for the permutation
|
||||
mapping = np.zeros((wmma_n, wmma_k, 2), dtype=int)
|
||||
for ii in range(wmma_n):
|
||||
for jj in range(wmma_k):
|
||||
mapping[ii, jj] = B_global_16x32_to_shared_load_16x32_layout(ii, jj)
|
||||
|
||||
# Reshape weight for the final format
|
||||
permutated_weight = np.zeros((N // wmma_n, K // wmma_k, wmma_n, wmma_k), dtype="int8")
|
||||
|
||||
# Use advanced indexing for the entire operation
|
||||
i_indices = np.arange(N // wmma_n)[:, np.newaxis, np.newaxis, np.newaxis]
|
||||
j_indices = np.arange(K // wmma_k)[np.newaxis, :, np.newaxis, np.newaxis]
|
||||
|
||||
# Create the source indices
|
||||
src_i = i_indices * wmma_n + mapping[:, :, 0]
|
||||
src_j = j_indices * wmma_k + mapping[:, :, 1]
|
||||
|
||||
# Extract and reshape in one go
|
||||
permutated_weight = weight[src_i, src_j]
|
||||
|
||||
return permutated_weight
|
||||
|
||||
|
||||
def compress_int2_to_int8(int2_weight):
|
||||
int8_weight = np.zeros(
|
||||
(*int2_weight.shape[:-1], int2_weight.shape[-1] // 4), dtype=np.int8
|
||||
)
|
||||
for j in range(int2_weight.shape[-1] // 4):
|
||||
for k in range(4):
|
||||
int8_weight[:, :, :, j] |= int2_weight[:, :, :, j * 4 + k] << (k * 2)
|
||||
return int8_weight
|
||||
|
||||
|
||||
def interleave_weight_int8(qweight, nbits=2):\
|
||||
# reinterpret the data type of qweight to int32
|
||||
# shift = [ 0, 8, 16, 24, 2, 10, 18, 26, 4, 12, 20, 28, 6, 14, 22, 30]
|
||||
# index: [ 0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15]
|
||||
qweight = qweight.view(np.int32)
|
||||
new_qweight = np.zeros_like(qweight)
|
||||
bits_stride = 8
|
||||
mask = (1 << nbits) - 1 # for 4bit the val is 0x0000000f
|
||||
num_groups = 32 // bits_stride # 4
|
||||
elems_per_group = bits_stride // nbits # 4
|
||||
for i in range(num_groups):
|
||||
for j in range(elems_per_group):
|
||||
offset = i * elems_per_group + j
|
||||
shift = (offset % num_groups) * bits_stride + (offset // num_groups) * nbits
|
||||
|
||||
new_qweight |= ((qweight >> (nbits * offset)) & mask) << shift
|
||||
return new_qweight.view(np.int8)
|
||||
|
||||
|
||||
|
||||
def convert_weight_int8_to_int2(weight):
|
||||
N = weight.shape[0]
|
||||
K = weight.shape[1]
|
||||
|
||||
weight = weight+2
|
||||
|
||||
weight = weight.cpu().numpy()
|
||||
|
||||
# print(weight)
|
||||
# print(torch.max(weight), torch.min(weight))
|
||||
|
||||
# permutated_weight_slow = permutate_weight(weight)
|
||||
permutated_weight = permutate_weight_fastest(weight)
|
||||
# assert np.all(permutated_weight_slow == permutated_weight)
|
||||
# print("Permutation is correct")
|
||||
compressed_weight = compress_int2_to_int8(permutated_weight)
|
||||
interleaved_weight = interleave_weight_int8(compressed_weight, 2)
|
||||
|
||||
ret = torch.from_numpy(interleaved_weight)
|
||||
|
||||
ret = torch.reshape(ret, (N, K // 4))
|
||||
|
||||
return ret
|
||||
Executable
+9
@@ -0,0 +1,9 @@
|
||||
fire
|
||||
sentencepiece
|
||||
torch>=2.2.0
|
||||
xformers>=0.0.22
|
||||
tiktoken
|
||||
blobfile
|
||||
flask
|
||||
einops
|
||||
transformers
|
||||
Executable
+31
@@ -0,0 +1,31 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the BSD license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
|
||||
@torch.compile
|
||||
def top_p(probs: torch.Tensor, p: float) -> torch.Tensor:
|
||||
"""
|
||||
Perform top-p (nucleus) sampling on a probability distribution.
|
||||
|
||||
Args:
|
||||
probs (torch.Tensor): probability distribution tensor.
|
||||
p (float): probability threshold for top-p sampling.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: sampled token indices.
|
||||
|
||||
Note:
|
||||
Top-p sampling selects the smallest set of tokens whose cumulative
|
||||
probability mass exceeds the threshold p. The distribution is
|
||||
renormalized based on the selected tokens.
|
||||
"""
|
||||
probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
|
||||
probs_sum = torch.cumsum(probs_sort, dim=-1)
|
||||
mask = probs_sum - probs_sort > p
|
||||
probs_sort[mask] = 0.0
|
||||
next_token = torch.multinomial(probs_sort, num_samples=1)
|
||||
next_token = torch.gather(probs_idx, -1, next_token)
|
||||
return next_token
|
||||
Executable
+57
@@ -0,0 +1,57 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the BSD license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class PhaseStats:
|
||||
name: str
|
||||
tokens: int
|
||||
time: float
|
||||
|
||||
def show(self) -> str:
|
||||
tps = self.tokens / self.time
|
||||
return (
|
||||
f"[{self.name}] "
|
||||
f"generated tokens: {self.tokens}"
|
||||
f" - total time: {self.time:.3f}s"
|
||||
f" - {tps:.1f} tokens per second"
|
||||
)
|
||||
|
||||
|
||||
class Stats:
|
||||
"""
|
||||
Generation stats, split by phases.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.phases = []
|
||||
self.current = None
|
||||
|
||||
def end_phase(self, tokens: int, now: Optional[float] = None):
|
||||
"""Terminate the current phase."""
|
||||
if self.current is None:
|
||||
return
|
||||
if now is None:
|
||||
now = time.time()
|
||||
cname, ctokens, ctime = self.current
|
||||
stats = PhaseStats(
|
||||
name=cname,
|
||||
tokens=tokens - ctokens,
|
||||
time=now - ctime,
|
||||
)
|
||||
self.phases.append(stats)
|
||||
|
||||
def phase(self, name: str, tokens: int = 0):
|
||||
"""
|
||||
Start a new phase, and terminate the current one,
|
||||
if one is ongoing.
|
||||
"""
|
||||
now = time.time()
|
||||
self.end_phase(tokens, now)
|
||||
self.current = (name, tokens, now)
|
||||
+99
@@ -0,0 +1,99 @@
|
||||
import torch
|
||||
from torch.utils import benchmark
|
||||
from torch import nn
|
||||
|
||||
from pack_weight import convert_weight_int8_to_int2
|
||||
from torch.profiler import profile, record_function, ProfilerActivity
|
||||
import ctypes
|
||||
import numpy as np
|
||||
# set all seed
|
||||
torch.manual_seed(42)
|
||||
np.random.seed(42)
|
||||
|
||||
bitnet_lib = ctypes.CDLL('bitnet_kernels/libbitnet.so')
|
||||
|
||||
def bitnet_int8xint2_linear(input0, input1, s, ws, ret):
|
||||
out_shape = list(input0.shape)
|
||||
out_shape[-1] = input1.shape[0]
|
||||
|
||||
stream = torch.cuda.current_stream()
|
||||
|
||||
M = input0.shape[0]
|
||||
if len(out_shape) == 3:
|
||||
M *= input0.shape[1]
|
||||
N = input1.shape[0]
|
||||
K = input1.shape[1] * 4
|
||||
|
||||
bitnet_lib.bitlinear_int8xint2(*[ctypes.c_void_p(input0.data_ptr()), ctypes.c_void_p(input1.data_ptr()), ctypes.c_void_p(ret.data_ptr()), ctypes.c_void_p(s.data_ptr()), ctypes.c_void_p(ws.data_ptr()), ctypes.c_int(M), ctypes.c_int(N), ctypes.c_int(K), ctypes.c_void_p(stream.cuda_stream)])
|
||||
|
||||
return ret
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_list = [
|
||||
(2560, 2560),
|
||||
(3840, 2560),
|
||||
(13824, 2560),
|
||||
(2560, 6912) ,
|
||||
(3200, 3200),
|
||||
(4800, 3200),
|
||||
(3200, 10240),
|
||||
(20480, 3200),
|
||||
]
|
||||
for N,K in test_list:
|
||||
weight = torch.randint(-1, 2, (N, K), dtype=torch.int8, device='cuda')
|
||||
weight_scale = torch.ones(1, dtype=torch.bfloat16, device='cuda')
|
||||
weight_compressed = convert_weight_int8_to_int2(weight).to('cuda')
|
||||
|
||||
for i in range(1):
|
||||
input0 = torch.randint(-128,127,(1, K),dtype=torch.int8, device='cuda')
|
||||
input0_bf16 = input0.to(torch.bfloat16)
|
||||
input_np = input0.cpu().to(torch.int32).numpy()
|
||||
weight_np = weight.cpu().to(torch.int32).T.numpy()
|
||||
out_np = np.matmul(input_np,weight_np)
|
||||
out_np = torch.tensor(out_np).cuda().to(torch.bfloat16)
|
||||
|
||||
s = torch.ones(1, dtype=torch.bfloat16, device='cuda')
|
||||
ws = torch.ones(6, dtype=torch.bfloat16, device='cuda')
|
||||
|
||||
ret = torch.empty((1,N), dtype=torch.bfloat16, device=input0.device)
|
||||
out = bitnet_int8xint2_linear(input0, weight_compressed, s, ws, ret)
|
||||
|
||||
print(f'custom == np {torch.all(out==out_np)}')
|
||||
|
||||
input0 = torch.randint(-128,127,(1, K),dtype=torch.int8, device='cuda')
|
||||
input0_fp16 = input0.to(torch.float16)
|
||||
input0_bf16 = input0.to(torch.bfloat16)
|
||||
weight_fp16 = weight.to(torch.float16).T
|
||||
weight_bf16 = weight.to(torch.bfloat16).T
|
||||
ret = torch.empty((1,N), dtype=torch.bfloat16, device=input0.device)
|
||||
s = torch.ones(1, dtype=torch.bfloat16, device='cuda')
|
||||
ws = torch.ones(6, dtype=torch.bfloat16, device='cuda')
|
||||
t0 = benchmark.Timer(
|
||||
stmt="bitnet_int8xint2_linear(input0, weight_compressed, s, ws, ret)",
|
||||
setup="from __main__ import input0, weight_compressed, s, ws, ret, bitnet_int8xint2_linear",
|
||||
num_threads=1,
|
||||
)
|
||||
|
||||
t1 = benchmark.Timer(
|
||||
stmt="torch.matmul(input0_bf16,weight_bf16)",
|
||||
setup="from __main__ import input0_bf16, weight_bf16",
|
||||
num_threads=1,
|
||||
)
|
||||
|
||||
time0 = t0.timeit(50)
|
||||
time1 = t1.timeit(50)
|
||||
|
||||
print(f'Shape{N,K}, W2A8: {time0.mean * 1e6:.2f}us, torch BF16: {time1.mean * 1e6:.2f}us')
|
||||
# activities = [ ProfilerActivity.CUDA,
|
||||
# # ProfilerActivity.CPU
|
||||
# ]
|
||||
# sort_by_keyword = 'cuda' + "_time_total"
|
||||
# with profile(activities=activities, record_shapes=True) as prof:
|
||||
# with record_function("model_inference1"):
|
||||
# for _ in range(10):
|
||||
# bitnet_int8xint2_linear(input0, weight_compressed, s, ws, ret)
|
||||
# torch.matmul(input0_fp16,weight_fp16)
|
||||
# torch.matmul(input0_bf16,weight_bf16)
|
||||
|
||||
# print(prof.key_averages().table(sort_by=sort_by_keyword, row_limit=15))
|
||||
|
||||
Executable
+128000
File diff suppressed because it is too large
Load Diff
Executable
+257
@@ -0,0 +1,257 @@
|
||||
import os
|
||||
from logging import getLogger
|
||||
from pathlib import Path
|
||||
from typing import (
|
||||
AbstractSet,
|
||||
cast,
|
||||
Collection,
|
||||
Dict,
|
||||
Iterator,
|
||||
List,
|
||||
Literal,
|
||||
Sequence,
|
||||
TypedDict,
|
||||
Union,
|
||||
)
|
||||
|
||||
import tiktoken
|
||||
from tiktoken.load import load_tiktoken_bpe
|
||||
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
||||
Role = Literal["system", "user", "assistant"]
|
||||
|
||||
|
||||
class Message(TypedDict):
|
||||
role: Role
|
||||
content: str
|
||||
|
||||
|
||||
Dialog = Sequence[Message]
|
||||
|
||||
|
||||
class Tokenizer:
|
||||
"""
|
||||
Tokenizing and encoding/decoding text using the Tiktoken tokenizer.
|
||||
"""
|
||||
|
||||
special_tokens: Dict[str, int]
|
||||
|
||||
num_reserved_special_tokens = 256
|
||||
|
||||
pat_str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+" # noqa: E501
|
||||
|
||||
def __init__(self, model_path: str):
|
||||
"""
|
||||
Initializes the Tokenizer with a Tiktoken model.
|
||||
|
||||
Args:
|
||||
model_path (str): The path to the Tiktoken model file.
|
||||
"""
|
||||
assert os.path.isfile(model_path), model_path
|
||||
|
||||
mergeable_ranks = load_tiktoken_bpe(model_path)
|
||||
num_base_tokens = len(mergeable_ranks)
|
||||
special_tokens = [
|
||||
"<|begin_of_text|>",
|
||||
"<|end_of_text|>",
|
||||
"<|reserved_special_token_0|>",
|
||||
"<|reserved_special_token_1|>",
|
||||
"<|reserved_special_token_2|>",
|
||||
"<|reserved_special_token_3|>",
|
||||
"<|start_header_id|>",
|
||||
"<|end_header_id|>",
|
||||
"<|reserved_special_token_4|>",
|
||||
"<|eot_id|>", # end of turn
|
||||
] + [
|
||||
f"<|reserved_special_token_{i}|>"
|
||||
for i in range(5, self.num_reserved_special_tokens - 5)
|
||||
]
|
||||
self.special_tokens = {
|
||||
token: num_base_tokens + i for i, token in enumerate(special_tokens)
|
||||
}
|
||||
self.model = tiktoken.Encoding(
|
||||
name=Path(model_path).name,
|
||||
pat_str=self.pat_str,
|
||||
mergeable_ranks=mergeable_ranks,
|
||||
special_tokens=self.special_tokens,
|
||||
)
|
||||
logger.info(f"Reloaded tiktoken model from {model_path}")
|
||||
|
||||
self.n_words: int = self.model.n_vocab
|
||||
# BOS / EOS token IDs
|
||||
self.bos_id: int = self.special_tokens["<|begin_of_text|>"]
|
||||
self.eos_id: int = self.special_tokens["<|end_of_text|>"]
|
||||
self.pad_id: int = self.n_words - 1
|
||||
self.stop_tokens = {
|
||||
self.special_tokens["<|end_of_text|>"],
|
||||
self.special_tokens["<|eot_id|>"],
|
||||
}
|
||||
logger.info(
|
||||
f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"
|
||||
)
|
||||
|
||||
def encode(
|
||||
self,
|
||||
s: str,
|
||||
*,
|
||||
bos: bool,
|
||||
eos: bool,
|
||||
allowed_special: Union[Literal["all"], AbstractSet[str]] = set(),
|
||||
disallowed_special: Union[Literal["all"], Collection[str]] = (),
|
||||
) -> List[int]:
|
||||
"""
|
||||
Encodes a string into a list of token IDs.
|
||||
|
||||
Args:
|
||||
s (str): The input string to be encoded.
|
||||
bos (bool): Whether to prepend the beginning-of-sequence token.
|
||||
eos (bool): Whether to append the end-of-sequence token.
|
||||
allowed_tokens ("all"|set[str]): allowed special tokens in string
|
||||
disallowed_tokens ("all"|set[str]): special tokens that raise an error when in string
|
||||
|
||||
Returns:
|
||||
list[int]: A list of token IDs.
|
||||
|
||||
By default, setting disallowed_special=() encodes a string by ignoring
|
||||
special tokens. Specifically:
|
||||
- Setting `disallowed_special` to () will cause all text corresponding
|
||||
to special tokens to be encoded as natural text (insteading of raising
|
||||
an error).
|
||||
- Setting `allowed_special` to "all" will treat all text corresponding
|
||||
to special tokens to be encoded as special tokens.
|
||||
"""
|
||||
assert type(s) is str
|
||||
|
||||
# The tiktoken tokenizer can handle <=400k chars without
|
||||
# pyo3_runtime.PanicException.
|
||||
TIKTOKEN_MAX_ENCODE_CHARS = 400_000
|
||||
|
||||
# https://github.com/openai/tiktoken/issues/195
|
||||
# Here we iterate over subsequences and split if we exceed the limit
|
||||
# of max consecutive non-whitespace or whitespace characters.
|
||||
MAX_NO_WHITESPACES_CHARS = 25_000
|
||||
|
||||
substrs = (
|
||||
substr
|
||||
for i in range(0, len(s), TIKTOKEN_MAX_ENCODE_CHARS)
|
||||
for substr in self._split_whitespaces_or_nonwhitespaces(
|
||||
s[i : i + TIKTOKEN_MAX_ENCODE_CHARS], MAX_NO_WHITESPACES_CHARS
|
||||
)
|
||||
)
|
||||
t: List[int] = []
|
||||
for substr in substrs:
|
||||
t.extend(
|
||||
self.model.encode(
|
||||
substr,
|
||||
allowed_special=allowed_special,
|
||||
disallowed_special=disallowed_special,
|
||||
)
|
||||
)
|
||||
if bos:
|
||||
t.insert(0, self.bos_id)
|
||||
if eos:
|
||||
t.append(self.eos_id)
|
||||
return t
|
||||
|
||||
def decode(self, t: Sequence[int]) -> str:
|
||||
"""
|
||||
Decodes a list of token IDs into a string.
|
||||
|
||||
Args:
|
||||
t (List[int]): The list of token IDs to be decoded.
|
||||
|
||||
Returns:
|
||||
str: The decoded string.
|
||||
"""
|
||||
# Typecast is safe here. Tiktoken doesn't do anything list-related with the sequence.
|
||||
return self.model.decode(cast(List[int], t))
|
||||
|
||||
@staticmethod
|
||||
def _split_whitespaces_or_nonwhitespaces(
|
||||
s: str, max_consecutive_slice_len: int
|
||||
) -> Iterator[str]:
|
||||
"""
|
||||
Splits the string `s` so that each substring contains no more than `max_consecutive_slice_len`
|
||||
consecutive whitespaces or consecutive non-whitespaces.
|
||||
"""
|
||||
current_slice_len = 0
|
||||
current_slice_is_space = s[0].isspace() if len(s) > 0 else False
|
||||
slice_start = 0
|
||||
|
||||
for i in range(len(s)):
|
||||
is_now_space = s[i].isspace()
|
||||
|
||||
if current_slice_is_space ^ is_now_space:
|
||||
current_slice_len = 1
|
||||
current_slice_is_space = is_now_space
|
||||
else:
|
||||
current_slice_len += 1
|
||||
if current_slice_len > max_consecutive_slice_len:
|
||||
yield s[slice_start:i]
|
||||
slice_start = i
|
||||
current_slice_len = 1
|
||||
yield s[slice_start:]
|
||||
|
||||
class ChatFormat:
|
||||
def __init__(self, tokenizer: Tokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
self.eot_id = tokenizer.special_tokens["<|eot_id|>"]
|
||||
|
||||
def decode(self, tokens: List[int]) -> str:
|
||||
# Decode the tokens to a string.
|
||||
decoded_str = self.tokenizer.decode(tokens)
|
||||
# Remove the special tokens from the decoded string.
|
||||
decoded_str = decoded_str.replace("<|eot_id|>", "")
|
||||
return decoded_str
|
||||
|
||||
def encode_header(self, message: Message) -> List[int]:
|
||||
tokens = []
|
||||
if message["role"] == "system":
|
||||
tokens.extend(self.tokenizer.encode("System: ", bos=False, eos=False))
|
||||
elif message["role"] == "user":
|
||||
tokens.extend(self.tokenizer.encode("User: ", bos=False, eos=False))
|
||||
elif message["role"] == "assistant":
|
||||
tokens.extend(self.tokenizer.encode("Assistant: ", bos=False, eos=False))
|
||||
else:
|
||||
raise NotImplementedError(f"Role {message['role']} not implemented.")
|
||||
# tokens.append(self.tokenizer.special_tokens["<|start_header_id|>"])
|
||||
# tokens.extend(self.tokenizer.encode(message["role"], bos=False, eos=False))
|
||||
# tokens.append(self.tokenizer.special_tokens["<|end_header_id|>"])
|
||||
# tokens.extend(self.tokenizer.encode("\n\n", bos=False, eos=False))
|
||||
return tokens
|
||||
|
||||
def encode_message(self, message: Message, return_target=False) -> List[int]:
|
||||
tokens, targets = [], []
|
||||
headers = self.encode_header(message)
|
||||
contents = self.tokenizer.encode(message["content"].strip(), bos=False, eos=False)
|
||||
contents.append(self.tokenizer.special_tokens["<|eot_id|>"])
|
||||
tokens = headers + contents
|
||||
|
||||
if message["role"] == "assistant":
|
||||
targets = [-1] * len(headers) + contents
|
||||
else:
|
||||
targets = [-1] * len(tokens)
|
||||
|
||||
if return_target:
|
||||
return tokens, targets
|
||||
|
||||
return tokens, None
|
||||
|
||||
def encode_dialog_prompt(self, dialog: Dialog, completion=False, return_target=False) -> List[int]:
|
||||
tokens = [self.tokenizer.special_tokens["<|begin_of_text|>"]]
|
||||
targets = [-1]
|
||||
for message in dialog:
|
||||
_tokens, _targets = self.encode_message(message, return_target=return_target)
|
||||
tokens.extend(_tokens)
|
||||
if _targets is not None:
|
||||
targets.extend(_targets)
|
||||
# Add the start of an assistant message for the model to complete.
|
||||
if completion:
|
||||
tokens.extend(self.encode_header({"role": "assistant", "content": ""}))
|
||||
|
||||
if return_target:
|
||||
return tokens, targets
|
||||
|
||||
return tokens
|
||||
@@ -0,0 +1,35 @@
|
||||
#define ACT_PARALLEL
|
||||
#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) || defined(__SSSE3__)
|
||||
#if defined(ACT_PARALLEL)
|
||||
#define ROW_BLOCK_SIZE 4
|
||||
#define COL_BLOCK_SIZE 128
|
||||
#define PARALLEL_SIZE 4
|
||||
#else
|
||||
#define ROW_BLOCK_SIZE 128
|
||||
#define COL_BLOCK_SIZE 32
|
||||
#define PARALLEL_SIZE 8
|
||||
#endif // ACT_PARALLEL
|
||||
#elif defined(__ARM_NEON)
|
||||
#if defined(__ARM_FEATURE_DOTPROD)
|
||||
#if defined(ACT_PARALLEL)
|
||||
#define ROW_BLOCK_SIZE 8
|
||||
#define COL_BLOCK_SIZE 256
|
||||
#define PARALLEL_SIZE 8
|
||||
#else
|
||||
#define ROW_BLOCK_SIZE 64
|
||||
#define COL_BLOCK_SIZE 16
|
||||
#define PARALLEL_SIZE 2
|
||||
#endif // ACT_PARALLEL
|
||||
#else
|
||||
#if defined(ACT_PARALLEL)
|
||||
#define ROW_BLOCK_SIZE 8
|
||||
#define COL_BLOCK_SIZE 256
|
||||
#define PARALLEL_SIZE 4
|
||||
#else
|
||||
#define ROW_BLOCK_SIZE 128
|
||||
#define COL_BLOCK_SIZE 32
|
||||
#define PARALLEL_SIZE 4
|
||||
#endif // ACT_PARALLEL
|
||||
#endif // __ARM_FEATURE_DOTPROD
|
||||
#endif // __AVX__
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
#pragma once
|
||||
|
||||
#include "ggml.h"
|
||||
#include "ggml-backend.h"
|
||||
|
||||
#ifdef __ARM_NEON
|
||||
#include <arm_neon.h>
|
||||
typedef float32_t bitnet_float_type;
|
||||
#else
|
||||
typedef float bitnet_float_type;
|
||||
#endif
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
struct bitnet_tensor_extra {
|
||||
int lut_scales_size;
|
||||
int BK;
|
||||
int n_tile_num;
|
||||
uint8_t * qweights;
|
||||
bitnet_float_type * scales;
|
||||
};
|
||||
|
||||
GGML_API void ggml_bitnet_init(void);
|
||||
GGML_API void ggml_bitnet_free(void);
|
||||
// src0->type == Q4_0/IQ2_XXS/IQ3_XXS
|
||||
// bitnet.cpp currently only supports BitNet quantization or GPTQ-like quantization (only scales, without zeros)
|
||||
// If use i-quantization gguf models, the results will be wrong
|
||||
// TODO: add customized block types Q2_0/Q3_0
|
||||
GGML_API bool ggml_bitnet_can_mul_mat(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst);
|
||||
GGML_API size_t ggml_bitnet_mul_mat_get_wsize(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst);
|
||||
GGML_API void ggml_bitnet_mul_mat_task_init(void * src1, void * qlut, void * lut_scales, void * lut_biases, int n, int k, int m, int bits);
|
||||
GGML_API void ggml_bitnet_mul_mat_task_compute(void * src0, void * scales, void * qlut, void * lut_scales, void * lut_biases, void * dst, int n, int k, int m, int bits);
|
||||
GGML_API void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor);
|
||||
GGML_API int ggml_bitnet_get_type_bits(enum ggml_type type);
|
||||
GGML_API void ggml_bitnet_set_n_threads(int n_threads);
|
||||
#if defined(GGML_BITNET_ARM_TL1)
|
||||
GGML_API void ggml_qgemm_lut(int m, int k, void* A, void* LUT, void* Scales, void* LUT_Scales, void* C);
|
||||
GGML_API void ggml_preprocessor(int m, int k, void* B, void* LUT_Scales, void* QLUT);
|
||||
#endif
|
||||
#if defined(GGML_BITNET_X86_TL2)
|
||||
GGML_API void ggml_qgemm_lut(int bs, int m, int k, int BK, void* A, void* sign, void* LUT, void* Scales, void* LUT_Scales, void* C);
|
||||
GGML_API void ggml_preprocessor(int bs, int m, int three_k, int two_k, void* B, void* LUT_Scales, void* Three_QLUT, void* Two_QLUT);
|
||||
#endif
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 36 KiB |
Binary file not shown.
@@ -0,0 +1,771 @@
|
||||
#if defined(GGML_BITNET_ARM_TL1)
|
||||
#include "ggml-bitnet.h"
|
||||
#define GGML_BITNET_MAX_NODES 8192
|
||||
static bool initialized = false;
|
||||
static bitnet_tensor_extra * bitnet_tensor_extras = nullptr;
|
||||
static size_t bitnet_tensor_extras_index = 0;
|
||||
static void * aligned_malloc(size_t size) {{
|
||||
#if defined(_WIN32)
|
||||
return _aligned_malloc(size, 64);
|
||||
#else
|
||||
void * ptr = nullptr;
|
||||
posix_memalign(&ptr, 64, size);
|
||||
return ptr;
|
||||
#endif
|
||||
}}
|
||||
static void aligned_free(void * ptr) {{
|
||||
#if defined(_WIN32)
|
||||
_aligned_free(ptr);
|
||||
#else
|
||||
free(ptr);
|
||||
#endif
|
||||
}}
|
||||
|
||||
void per_tensor_quant(int k, void* lut_scales_, void* b_) {{
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;
|
||||
bitnet_float_type* b = (bitnet_float_type*)b_;
|
||||
#ifdef __ARM_NEON
|
||||
float32x4_t temp_max = vdupq_n_f32(0);
|
||||
for (int i=0; i < k / 4; i++) {{
|
||||
float32x4_t vec_bs = vld1q_f32(b + 4 * i);
|
||||
float32x4_t abssum = vabsq_f32(vec_bs);
|
||||
temp_max = vmaxq_f32(abssum, temp_max);
|
||||
}}
|
||||
float32_t scales = 127 / vmaxvq_f32(temp_max);
|
||||
*lut_scales = scales;
|
||||
#elif defined __AVX2__
|
||||
__m256 max_vec = _mm256_set1_ps(0.f);
|
||||
const __m256 vec_sign = _mm256_set1_ps(-0.0f);
|
||||
// #pragma unroll
|
||||
for (int i = 0; i < k / 8; i++) {{
|
||||
__m256 vec_b = _mm256_loadu_ps(b + i * 8);
|
||||
__m256 vec_babs = _mm256_andnot_ps(vec_sign, vec_b);
|
||||
max_vec = _mm256_max_ps(vec_babs, max_vec);
|
||||
}}
|
||||
__m128 max1 = _mm_max_ps(_mm256_extractf128_ps(max_vec, 1), _mm256_castps256_ps128(max_vec));
|
||||
max1 = _mm_max_ps(max1, _mm_movehl_ps(max1, max1));
|
||||
max1 = _mm_max_ss(max1, _mm_movehdup_ps(max1));
|
||||
float scales = 127 / _mm_cvtss_f32(max1);
|
||||
*lut_scales = scales;
|
||||
#endif
|
||||
}}
|
||||
|
||||
void partial_max_reset(void* lut_scales_) {{
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;
|
||||
*lut_scales = 0.0;
|
||||
}}
|
||||
|
||||
#ifdef __ARM_NEON
|
||||
inline void Transpose_8_8(
|
||||
int16x8_t *v0,
|
||||
int16x8_t *v1,
|
||||
int16x8_t *v2,
|
||||
int16x8_t *v3,
|
||||
int16x8_t *v4,
|
||||
int16x8_t *v5,
|
||||
int16x8_t *v6,
|
||||
int16x8_t *v7)
|
||||
{{
|
||||
int16x8x2_t q04 = vzipq_s16(*v0, *v4);
|
||||
int16x8x2_t q15 = vzipq_s16(*v1, *v5);
|
||||
int16x8x2_t q26 = vzipq_s16(*v2, *v6);
|
||||
int16x8x2_t q37 = vzipq_s16(*v3, *v7);
|
||||
|
||||
int16x8x2_t q0246_0 = vzipq_s16(q04.val[0], q26.val[0]);
|
||||
int16x8x2_t q0246_1 = vzipq_s16(q04.val[1], q26.val[1]);
|
||||
int16x8x2_t q1357_0 = vzipq_s16(q15.val[0], q37.val[0]);
|
||||
int16x8x2_t q1357_1 = vzipq_s16(q15.val[1], q37.val[1]);
|
||||
|
||||
int16x8x2_t q_fin_0 = vzipq_s16(q0246_0.val[0], q1357_0.val[0]);
|
||||
int16x8x2_t q_fin_1 = vzipq_s16(q0246_0.val[1], q1357_0.val[1]);
|
||||
int16x8x2_t q_fin_2 = vzipq_s16(q0246_1.val[0], q1357_1.val[0]);
|
||||
int16x8x2_t q_fin_3 = vzipq_s16(q0246_1.val[1], q1357_1.val[1]);
|
||||
|
||||
*v0 = q_fin_0.val[0];
|
||||
*v1 = q_fin_0.val[1];
|
||||
*v2 = q_fin_1.val[0];
|
||||
*v3 = q_fin_1.val[1];
|
||||
*v4 = q_fin_2.val[0];
|
||||
*v5 = q_fin_2.val[1];
|
||||
*v6 = q_fin_3.val[0];
|
||||
*v7 = q_fin_3.val[1];
|
||||
}}
|
||||
#endif
|
||||
|
||||
template<int act_k>
|
||||
inline void lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_type* lut_scales) {{
|
||||
#ifdef __ARM_NEON
|
||||
int16x8_t vec_lut[16];
|
||||
float32_t scales = *lut_scales;
|
||||
uint8_t tbl_mask[16];
|
||||
tbl_mask[0] = 0;
|
||||
tbl_mask[1] = 2;
|
||||
tbl_mask[2] = 4;
|
||||
tbl_mask[3] = 6;
|
||||
tbl_mask[4] = 8;
|
||||
tbl_mask[5] = 10;
|
||||
tbl_mask[6] = 12;
|
||||
tbl_mask[7] = 14;
|
||||
tbl_mask[8] = 1;
|
||||
tbl_mask[9] = 3;
|
||||
tbl_mask[10] = 5;
|
||||
tbl_mask[11] = 7;
|
||||
tbl_mask[12] = 9;
|
||||
tbl_mask[13] = 11;
|
||||
tbl_mask[14] = 13;
|
||||
tbl_mask[15] = 15;
|
||||
uint8x16_t tbl_mask_q = vld1q_u8(tbl_mask);
|
||||
#pragma unroll
|
||||
for (int k = 0; k < act_k / 16; ++k) {{
|
||||
float32x4x2_t vec_bs_x0 = vld2q_f32(b + k * 16);
|
||||
float32x4x2_t vec_bs_x1 = vld2q_f32(b + k * 16 + 8);
|
||||
float32x4_t vec_f_0 = vmulq_n_f32(vec_bs_x0.val[0], scales);
|
||||
float32x4_t vec_f_1 = vmulq_n_f32(vec_bs_x0.val[1], scales);
|
||||
float32x4_t vec_f_2 = vmulq_n_f32(vec_bs_x1.val[0], scales);
|
||||
float32x4_t vec_f_3 = vmulq_n_f32(vec_bs_x1.val[1], scales);
|
||||
int32x4_t vec_b_0 = vcvtnq_s32_f32(vec_f_0);
|
||||
int32x4_t vec_b_1 = vcvtnq_s32_f32(vec_f_1);
|
||||
int32x4_t vec_b_2 = vcvtnq_s32_f32(vec_f_2);
|
||||
int32x4_t vec_b_3 = vcvtnq_s32_f32(vec_f_3);
|
||||
int16x4_t vec_b16_0 = vmovn_s32(vec_b_0);
|
||||
int16x4_t vec_b16_1 = vmovn_s32(vec_b_1);
|
||||
int16x4_t vec_b16_2 = vmovn_s32(vec_b_2);
|
||||
int16x4_t vec_b16_3 = vmovn_s32(vec_b_3);
|
||||
int16x8_t vec_bs_0 = vcombine_s16(vec_b16_0, vec_b16_2);
|
||||
int16x8_t vec_bs_1 = vcombine_s16(vec_b16_1, vec_b16_3);
|
||||
vec_lut[0] = vdupq_n_s16(0);
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_0;
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_1;
|
||||
vec_lut[1] = vdupq_n_s16(0);
|
||||
vec_lut[1] = vec_lut[1] - vec_bs_0;
|
||||
vec_lut[2] = vdupq_n_s16(0);
|
||||
vec_lut[2] = vec_lut[2] - vec_bs_0;
|
||||
vec_lut[2] = vec_lut[2] + vec_bs_1;
|
||||
vec_lut[3] = vdupq_n_s16(0);
|
||||
vec_lut[3] = vec_lut[3] - vec_bs_1;
|
||||
vec_lut[4] = vdupq_n_s16(0);
|
||||
vec_lut[5] = vec_bs_1;
|
||||
vec_lut[6] = vec_bs_0;
|
||||
vec_lut[6] = vec_lut[6] - vec_bs_1;
|
||||
vec_lut[7] = vec_bs_0;
|
||||
vec_lut[8] = vec_bs_0;
|
||||
vec_lut[8] = vec_lut[8] + vec_bs_1;
|
||||
Transpose_8_8(&(vec_lut[0]), &(vec_lut[1]), &(vec_lut[2]), &(vec_lut[3]),
|
||||
&(vec_lut[4]), &(vec_lut[5]), &(vec_lut[6]), &(vec_lut[7]));
|
||||
Transpose_8_8(&(vec_lut[8]), &(vec_lut[9]), &(vec_lut[10]), &(vec_lut[11]),
|
||||
&(vec_lut[12]), &(vec_lut[13]), &(vec_lut[14]), &(vec_lut[15]));
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < 8; idx++) {{
|
||||
int8x16_t q0_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx]), tbl_mask_q);
|
||||
int8x8_t q0_low = vget_low_s8(q0_s);
|
||||
int8x8_t q0_high = vget_high_s8(q0_s);
|
||||
int8x16_t q1_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx + 8]), tbl_mask_q);
|
||||
int8x8_t q1_low = vget_low_s8(q1_s);
|
||||
int8x8_t q1_high = vget_high_s8(q1_s);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2, q0_high);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 8, q1_high);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 16, q0_low);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 24, q1_low);
|
||||
}}
|
||||
}}
|
||||
#endif
|
||||
}}
|
||||
|
||||
static bool is_type_supported(enum ggml_type type) {{
|
||||
if (type == GGML_TYPE_Q4_0 ||
|
||||
type == GGML_TYPE_TL1) {{
|
||||
return true;
|
||||
}} else {{
|
||||
return false;
|
||||
}}
|
||||
}}
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM14336_4096 256
|
||||
#define BBK14336_4096 128
|
||||
inline void tbl_impl_14336_4096(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK14336_4096 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[8];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM14336_4096; i += 64) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<8; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 2; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[2] += vec_v_left_1.val[0];
|
||||
vec_c[2] += vec_v_right_1.val[0];
|
||||
vec_c[3] += vec_v_left_1.val[1];
|
||||
vec_c[3] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[4] += vec_v_left_2.val[0];
|
||||
vec_c[4] += vec_v_right_2.val[0];
|
||||
vec_c[5] += vec_v_left_2.val[1];
|
||||
vec_c[5] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[6] += vec_v_left_3.val[0];
|
||||
vec_c[6] += vec_v_right_3.val[0];
|
||||
vec_c[7] += vec_v_left_3.val[1];
|
||||
vec_c[7] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
int32x4_t vec_v_bot_low_low_4 = vmovl_s16(vget_low_s16(vec_c[4]));
|
||||
int32x4_t vec_v_bot_low_high_4 = vmovl_high_s16(vec_c[4]);
|
||||
vst1q_s32(c + i + 32, vld1q_s32(c + i + 32) + vec_v_bot_low_low_4);
|
||||
vst1q_s32(c + i + 36, vld1q_s32(c + i + 36) + vec_v_bot_low_high_4);
|
||||
int32x4_t vec_v_bot_low_low_5 = vmovl_s16(vget_low_s16(vec_c[5]));
|
||||
int32x4_t vec_v_bot_low_high_5 = vmovl_high_s16(vec_c[5]);
|
||||
vst1q_s32(c + i + 40, vld1q_s32(c + i + 40) + vec_v_bot_low_low_5);
|
||||
vst1q_s32(c + i + 44, vld1q_s32(c + i + 44) + vec_v_bot_low_high_5);
|
||||
int32x4_t vec_v_bot_low_low_6 = vmovl_s16(vget_low_s16(vec_c[6]));
|
||||
int32x4_t vec_v_bot_low_high_6 = vmovl_high_s16(vec_c[6]);
|
||||
vst1q_s32(c + i + 48, vld1q_s32(c + i + 48) + vec_v_bot_low_low_6);
|
||||
vst1q_s32(c + i + 52, vld1q_s32(c + i + 52) + vec_v_bot_low_high_6);
|
||||
int32x4_t vec_v_bot_low_low_7 = vmovl_s16(vget_low_s16(vec_c[7]));
|
||||
int32x4_t vec_v_bot_low_high_7 = vmovl_high_s16(vec_c[7]);
|
||||
vst1q_s32(c + i + 56, vld1q_s32(c + i + 56) + vec_v_bot_low_low_7);
|
||||
vst1q_s32(c + i + 60, vld1q_s32(c + i + 60) + vec_v_bot_low_high_7);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_14336_4096(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM14336_4096];
|
||||
memset(&(CBits[0]), 0, BM14336_4096 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 4096 / BBK14336_4096; ++k_outer) {
|
||||
tbl_impl_14336_4096((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK14336_4096 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK14336_4096 / 2 / 2 * BM14336_4096)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM14336_4096; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM4096_14336 256
|
||||
#define BBK4096_14336 128
|
||||
inline void tbl_impl_4096_14336(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK4096_14336 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM4096_14336; i += 32) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<4; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 4; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[0] += vec_v_left_1.val[0];
|
||||
vec_c[0] += vec_v_right_1.val[0];
|
||||
vec_c[1] += vec_v_left_1.val[1];
|
||||
vec_c[1] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[2] += vec_v_left_2.val[0];
|
||||
vec_c[2] += vec_v_right_2.val[0];
|
||||
vec_c[3] += vec_v_left_2.val[1];
|
||||
vec_c[3] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[2] += vec_v_left_3.val[0];
|
||||
vec_c[2] += vec_v_right_3.val[0];
|
||||
vec_c[3] += vec_v_left_3.val[1];
|
||||
vec_c[3] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_4096_14336(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM4096_14336];
|
||||
memset(&(CBits[0]), 0, BM4096_14336 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 14336 / BBK4096_14336; ++k_outer) {
|
||||
tbl_impl_4096_14336((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK4096_14336 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK4096_14336 / 2 / 2 * BM4096_14336)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM4096_14336; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM1024_4096 128
|
||||
#define BBK1024_4096 64
|
||||
inline void tbl_impl_1024_4096(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK1024_4096 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[8];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM1024_4096; i += 64) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<8; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 2; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[2] += vec_v_left_1.val[0];
|
||||
vec_c[2] += vec_v_right_1.val[0];
|
||||
vec_c[3] += vec_v_left_1.val[1];
|
||||
vec_c[3] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[4] += vec_v_left_2.val[0];
|
||||
vec_c[4] += vec_v_right_2.val[0];
|
||||
vec_c[5] += vec_v_left_2.val[1];
|
||||
vec_c[5] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[6] += vec_v_left_3.val[0];
|
||||
vec_c[6] += vec_v_right_3.val[0];
|
||||
vec_c[7] += vec_v_left_3.val[1];
|
||||
vec_c[7] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
int32x4_t vec_v_bot_low_low_4 = vmovl_s16(vget_low_s16(vec_c[4]));
|
||||
int32x4_t vec_v_bot_low_high_4 = vmovl_high_s16(vec_c[4]);
|
||||
vst1q_s32(c + i + 32, vld1q_s32(c + i + 32) + vec_v_bot_low_low_4);
|
||||
vst1q_s32(c + i + 36, vld1q_s32(c + i + 36) + vec_v_bot_low_high_4);
|
||||
int32x4_t vec_v_bot_low_low_5 = vmovl_s16(vget_low_s16(vec_c[5]));
|
||||
int32x4_t vec_v_bot_low_high_5 = vmovl_high_s16(vec_c[5]);
|
||||
vst1q_s32(c + i + 40, vld1q_s32(c + i + 40) + vec_v_bot_low_low_5);
|
||||
vst1q_s32(c + i + 44, vld1q_s32(c + i + 44) + vec_v_bot_low_high_5);
|
||||
int32x4_t vec_v_bot_low_low_6 = vmovl_s16(vget_low_s16(vec_c[6]));
|
||||
int32x4_t vec_v_bot_low_high_6 = vmovl_high_s16(vec_c[6]);
|
||||
vst1q_s32(c + i + 48, vld1q_s32(c + i + 48) + vec_v_bot_low_low_6);
|
||||
vst1q_s32(c + i + 52, vld1q_s32(c + i + 52) + vec_v_bot_low_high_6);
|
||||
int32x4_t vec_v_bot_low_low_7 = vmovl_s16(vget_low_s16(vec_c[7]));
|
||||
int32x4_t vec_v_bot_low_high_7 = vmovl_high_s16(vec_c[7]);
|
||||
vst1q_s32(c + i + 56, vld1q_s32(c + i + 56) + vec_v_bot_low_low_7);
|
||||
vst1q_s32(c + i + 60, vld1q_s32(c + i + 60) + vec_v_bot_low_high_7);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_1024_4096(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM1024_4096];
|
||||
memset(&(CBits[0]), 0, BM1024_4096 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 4096 / BBK1024_4096; ++k_outer) {
|
||||
tbl_impl_1024_4096((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK1024_4096 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK1024_4096 / 2 / 2 * BM1024_4096)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM1024_4096; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM4096_4096 128
|
||||
#define BBK4096_4096 64
|
||||
inline void tbl_impl_4096_4096(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK4096_4096 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM4096_4096; i += 32) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<4; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 4; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[0] += vec_v_left_1.val[0];
|
||||
vec_c[0] += vec_v_right_1.val[0];
|
||||
vec_c[1] += vec_v_left_1.val[1];
|
||||
vec_c[1] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[2] += vec_v_left_2.val[0];
|
||||
vec_c[2] += vec_v_right_2.val[0];
|
||||
vec_c[3] += vec_v_left_2.val[1];
|
||||
vec_c[3] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[2] += vec_v_left_3.val[0];
|
||||
vec_c[2] += vec_v_right_3.val[0];
|
||||
vec_c[3] += vec_v_left_3.val[1];
|
||||
vec_c[3] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_4096_4096(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM4096_4096];
|
||||
memset(&(CBits[0]), 0, BM4096_4096 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 4096 / BBK4096_4096; ++k_outer) {
|
||||
tbl_impl_4096_4096((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK4096_4096 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK4096_4096 / 2 / 2 * BM4096_4096)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM4096_4096; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
|
||||
template<int K>
|
||||
void preprocessor_k(void* B, void* LUT_Scales, void* QLUT) {{
|
||||
partial_max_reset((&(((bitnet_float_type*)LUT_Scales)[0])));
|
||||
per_tensor_quant(K, (&(((bitnet_float_type*)LUT_Scales)[0])), (&(((bitnet_float_type*)B)[0])));
|
||||
|
||||
lut_ctor<K>((&(((int8_t*)QLUT)[0])), (&(((bitnet_float_type*)B)[0])), (&(((bitnet_float_type*)LUT_Scales)[0])));
|
||||
}}
|
||||
void ggml_preprocessor(int m, int k, void* B, void* LUT_Scales, void* QLUT) {
|
||||
if (m == 14336 && k == 4096) {
|
||||
preprocessor_k<4096>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
else if (m == 4096 && k == 14336) {
|
||||
preprocessor_k<14336>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
else if (m == 1024 && k == 4096) {
|
||||
preprocessor_k<4096>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
else if (m == 4096 && k == 4096) {
|
||||
preprocessor_k<4096>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
}
|
||||
void ggml_qgemm_lut(int m, int k, void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
if (m == 14336 && k == 4096) {
|
||||
qgemm_lut_14336_4096(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
else if (m == 4096 && k == 14336) {
|
||||
qgemm_lut_4096_14336(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
else if (m == 1024 && k == 4096) {
|
||||
qgemm_lut_1024_4096(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
else if (m == 4096 && k == 4096) {
|
||||
qgemm_lut_4096_4096(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
}
|
||||
|
||||
void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor) {
|
||||
if (!(is_type_supported(tensor->type) && tensor->backend == GGML_BACKEND_TYPE_CPU && tensor->extra == nullptr)) {
|
||||
return;
|
||||
}
|
||||
|
||||
int k = tensor->ne[0];
|
||||
int m = tensor->ne[1];
|
||||
const int lut_scales_size = 1;
|
||||
const int scales_size = 1;
|
||||
int bk = 0;
|
||||
int bm = 0;
|
||||
|
||||
if (m == 14336 && k == 4096) {
|
||||
bm = BM14336_4096;
|
||||
bk = BBK14336_4096;
|
||||
}
|
||||
else if (m == 4096 && k == 14336) {
|
||||
bm = BM4096_14336;
|
||||
bk = BBK4096_14336;
|
||||
}
|
||||
else if (m == 1024 && k == 4096) {
|
||||
bm = BM1024_4096;
|
||||
bk = BBK1024_4096;
|
||||
}
|
||||
else if (m == 4096 && k == 4096) {
|
||||
bm = BM4096_4096;
|
||||
bk = BBK4096_4096;
|
||||
}
|
||||
|
||||
const int n_tile_num = m / bm;
|
||||
const int BK = bk;
|
||||
uint8_t * qweights;
|
||||
bitnet_float_type * scales;
|
||||
|
||||
scales = (bitnet_float_type *) aligned_malloc(sizeof(bitnet_float_type));
|
||||
qweights = (uint8_t *) tensor->data;
|
||||
float * i2_scales = (float * )(qweights + k * m / 4);
|
||||
scales[0] = (bitnet_float_type) i2_scales[0];
|
||||
|
||||
tensor->extra = bitnet_tensor_extras + bitnet_tensor_extras_index;
|
||||
bitnet_tensor_extras[bitnet_tensor_extras_index++] = {
|
||||
/* .lut_scales_size = */ lut_scales_size,
|
||||
/* .scales_size = */ scales_size,
|
||||
/* .n_tile_num = */ n_tile_num,
|
||||
/* .qweights = */ qweights,
|
||||
/* .scales = */ scales
|
||||
};
|
||||
}
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,28 @@
|
||||
[Kernels_0]
|
||||
m = 14336
|
||||
k = 4096
|
||||
bm = 256
|
||||
bk = 128
|
||||
bmm = 64
|
||||
|
||||
[Kernels_1]
|
||||
m = 4096
|
||||
k = 14336
|
||||
bm = 256
|
||||
bk = 128
|
||||
bmm = 32
|
||||
|
||||
[Kernels_2]
|
||||
m = 1024
|
||||
k = 4096
|
||||
bm = 128
|
||||
bk = 64
|
||||
bmm = 64
|
||||
|
||||
[Kernels_3]
|
||||
m = 4096
|
||||
k = 4096
|
||||
bm = 128
|
||||
bk = 64
|
||||
bmm = 32
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
[Kernels_0]
|
||||
m = 14336
|
||||
k = 4096
|
||||
bm = 256
|
||||
bk = 96
|
||||
bmm = 32
|
||||
|
||||
[Kernels_1]
|
||||
m = 4096
|
||||
k = 14336
|
||||
bm = 128
|
||||
bk = 96
|
||||
bmm = 32
|
||||
|
||||
[Kernels_2]
|
||||
m = 1024
|
||||
k = 4096
|
||||
bm = 256
|
||||
bk = 96
|
||||
bmm = 32
|
||||
|
||||
[Kernels_3]
|
||||
m = 4096
|
||||
k = 4096
|
||||
bm = 128
|
||||
bk = 96
|
||||
bmm = 32
|
||||
|
||||
@@ -0,0 +1,627 @@
|
||||
#if defined(GGML_BITNET_ARM_TL1)
|
||||
#include "ggml-bitnet.h"
|
||||
#define GGML_BITNET_MAX_NODES 8192
|
||||
static bool initialized = false;
|
||||
static bitnet_tensor_extra * bitnet_tensor_extras = nullptr;
|
||||
static size_t bitnet_tensor_extras_index = 0;
|
||||
static void * aligned_malloc(size_t size) {{
|
||||
#if defined(_WIN32)
|
||||
return _aligned_malloc(size, 64);
|
||||
#else
|
||||
void * ptr = nullptr;
|
||||
posix_memalign(&ptr, 64, size);
|
||||
return ptr;
|
||||
#endif
|
||||
}}
|
||||
static void aligned_free(void * ptr) {{
|
||||
#if defined(_WIN32)
|
||||
_aligned_free(ptr);
|
||||
#else
|
||||
free(ptr);
|
||||
#endif
|
||||
}}
|
||||
|
||||
void per_tensor_quant(int k, void* lut_scales_, void* b_) {{
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;
|
||||
bitnet_float_type* b = (bitnet_float_type*)b_;
|
||||
#ifdef __ARM_NEON
|
||||
float32x4_t temp_max = vdupq_n_f32(0);
|
||||
for (int i=0; i < k / 4; i++) {{
|
||||
float32x4_t vec_bs = vld1q_f32(b + 4 * i);
|
||||
float32x4_t abssum = vabsq_f32(vec_bs);
|
||||
temp_max = vmaxq_f32(abssum, temp_max);
|
||||
}}
|
||||
float32_t scales = 127 / vmaxvq_f32(temp_max);
|
||||
*lut_scales = scales;
|
||||
#elif defined __AVX2__
|
||||
__m256 max_vec = _mm256_set1_ps(0.f);
|
||||
const __m256 vec_sign = _mm256_set1_ps(-0.0f);
|
||||
// #pragma unroll
|
||||
for (int i = 0; i < k / 8; i++) {{
|
||||
__m256 vec_b = _mm256_loadu_ps(b + i * 8);
|
||||
__m256 vec_babs = _mm256_andnot_ps(vec_sign, vec_b);
|
||||
max_vec = _mm256_max_ps(vec_babs, max_vec);
|
||||
}}
|
||||
__m128 max1 = _mm_max_ps(_mm256_extractf128_ps(max_vec, 1), _mm256_castps256_ps128(max_vec));
|
||||
max1 = _mm_max_ps(max1, _mm_movehl_ps(max1, max1));
|
||||
max1 = _mm_max_ss(max1, _mm_movehdup_ps(max1));
|
||||
float scales = 127 / _mm_cvtss_f32(max1);
|
||||
*lut_scales = scales;
|
||||
#endif
|
||||
}}
|
||||
|
||||
void partial_max_reset(void* lut_scales_) {{
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;
|
||||
*lut_scales = 0.0;
|
||||
}}
|
||||
|
||||
#ifdef __ARM_NEON
|
||||
inline void Transpose_8_8(
|
||||
int16x8_t *v0,
|
||||
int16x8_t *v1,
|
||||
int16x8_t *v2,
|
||||
int16x8_t *v3,
|
||||
int16x8_t *v4,
|
||||
int16x8_t *v5,
|
||||
int16x8_t *v6,
|
||||
int16x8_t *v7)
|
||||
{{
|
||||
int16x8x2_t q04 = vzipq_s16(*v0, *v4);
|
||||
int16x8x2_t q15 = vzipq_s16(*v1, *v5);
|
||||
int16x8x2_t q26 = vzipq_s16(*v2, *v6);
|
||||
int16x8x2_t q37 = vzipq_s16(*v3, *v7);
|
||||
|
||||
int16x8x2_t q0246_0 = vzipq_s16(q04.val[0], q26.val[0]);
|
||||
int16x8x2_t q0246_1 = vzipq_s16(q04.val[1], q26.val[1]);
|
||||
int16x8x2_t q1357_0 = vzipq_s16(q15.val[0], q37.val[0]);
|
||||
int16x8x2_t q1357_1 = vzipq_s16(q15.val[1], q37.val[1]);
|
||||
|
||||
int16x8x2_t q_fin_0 = vzipq_s16(q0246_0.val[0], q1357_0.val[0]);
|
||||
int16x8x2_t q_fin_1 = vzipq_s16(q0246_0.val[1], q1357_0.val[1]);
|
||||
int16x8x2_t q_fin_2 = vzipq_s16(q0246_1.val[0], q1357_1.val[0]);
|
||||
int16x8x2_t q_fin_3 = vzipq_s16(q0246_1.val[1], q1357_1.val[1]);
|
||||
|
||||
*v0 = q_fin_0.val[0];
|
||||
*v1 = q_fin_0.val[1];
|
||||
*v2 = q_fin_1.val[0];
|
||||
*v3 = q_fin_1.val[1];
|
||||
*v4 = q_fin_2.val[0];
|
||||
*v5 = q_fin_2.val[1];
|
||||
*v6 = q_fin_3.val[0];
|
||||
*v7 = q_fin_3.val[1];
|
||||
}}
|
||||
#endif
|
||||
|
||||
template<int act_k>
|
||||
inline void lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_type* lut_scales) {{
|
||||
#ifdef __ARM_NEON
|
||||
int16x8_t vec_lut[16];
|
||||
float32_t scales = *lut_scales;
|
||||
uint8_t tbl_mask[16];
|
||||
tbl_mask[0] = 0;
|
||||
tbl_mask[1] = 2;
|
||||
tbl_mask[2] = 4;
|
||||
tbl_mask[3] = 6;
|
||||
tbl_mask[4] = 8;
|
||||
tbl_mask[5] = 10;
|
||||
tbl_mask[6] = 12;
|
||||
tbl_mask[7] = 14;
|
||||
tbl_mask[8] = 1;
|
||||
tbl_mask[9] = 3;
|
||||
tbl_mask[10] = 5;
|
||||
tbl_mask[11] = 7;
|
||||
tbl_mask[12] = 9;
|
||||
tbl_mask[13] = 11;
|
||||
tbl_mask[14] = 13;
|
||||
tbl_mask[15] = 15;
|
||||
uint8x16_t tbl_mask_q = vld1q_u8(tbl_mask);
|
||||
#pragma unroll
|
||||
for (int k = 0; k < act_k / 16; ++k) {{
|
||||
float32x4x2_t vec_bs_x0 = vld2q_f32(b + k * 16);
|
||||
float32x4x2_t vec_bs_x1 = vld2q_f32(b + k * 16 + 8);
|
||||
float32x4_t vec_f_0 = vmulq_n_f32(vec_bs_x0.val[0], scales);
|
||||
float32x4_t vec_f_1 = vmulq_n_f32(vec_bs_x0.val[1], scales);
|
||||
float32x4_t vec_f_2 = vmulq_n_f32(vec_bs_x1.val[0], scales);
|
||||
float32x4_t vec_f_3 = vmulq_n_f32(vec_bs_x1.val[1], scales);
|
||||
int32x4_t vec_b_0 = vcvtnq_s32_f32(vec_f_0);
|
||||
int32x4_t vec_b_1 = vcvtnq_s32_f32(vec_f_1);
|
||||
int32x4_t vec_b_2 = vcvtnq_s32_f32(vec_f_2);
|
||||
int32x4_t vec_b_3 = vcvtnq_s32_f32(vec_f_3);
|
||||
int16x4_t vec_b16_0 = vmovn_s32(vec_b_0);
|
||||
int16x4_t vec_b16_1 = vmovn_s32(vec_b_1);
|
||||
int16x4_t vec_b16_2 = vmovn_s32(vec_b_2);
|
||||
int16x4_t vec_b16_3 = vmovn_s32(vec_b_3);
|
||||
int16x8_t vec_bs_0 = vcombine_s16(vec_b16_0, vec_b16_2);
|
||||
int16x8_t vec_bs_1 = vcombine_s16(vec_b16_1, vec_b16_3);
|
||||
vec_lut[0] = vdupq_n_s16(0);
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_0;
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_1;
|
||||
vec_lut[1] = vdupq_n_s16(0);
|
||||
vec_lut[1] = vec_lut[1] - vec_bs_0;
|
||||
vec_lut[2] = vdupq_n_s16(0);
|
||||
vec_lut[2] = vec_lut[2] - vec_bs_0;
|
||||
vec_lut[2] = vec_lut[2] + vec_bs_1;
|
||||
vec_lut[3] = vdupq_n_s16(0);
|
||||
vec_lut[3] = vec_lut[3] - vec_bs_1;
|
||||
vec_lut[4] = vdupq_n_s16(0);
|
||||
vec_lut[5] = vec_bs_1;
|
||||
vec_lut[6] = vec_bs_0;
|
||||
vec_lut[6] = vec_lut[6] - vec_bs_1;
|
||||
vec_lut[7] = vec_bs_0;
|
||||
vec_lut[8] = vec_bs_0;
|
||||
vec_lut[8] = vec_lut[8] + vec_bs_1;
|
||||
Transpose_8_8(&(vec_lut[0]), &(vec_lut[1]), &(vec_lut[2]), &(vec_lut[3]),
|
||||
&(vec_lut[4]), &(vec_lut[5]), &(vec_lut[6]), &(vec_lut[7]));
|
||||
Transpose_8_8(&(vec_lut[8]), &(vec_lut[9]), &(vec_lut[10]), &(vec_lut[11]),
|
||||
&(vec_lut[12]), &(vec_lut[13]), &(vec_lut[14]), &(vec_lut[15]));
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < 8; idx++) {{
|
||||
int8x16_t q0_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx]), tbl_mask_q);
|
||||
int8x8_t q0_low = vget_low_s8(q0_s);
|
||||
int8x8_t q0_high = vget_high_s8(q0_s);
|
||||
int8x16_t q1_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx + 8]), tbl_mask_q);
|
||||
int8x8_t q1_low = vget_low_s8(q1_s);
|
||||
int8x8_t q1_high = vget_high_s8(q1_s);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2, q0_high);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 8, q1_high);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 16, q0_low);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 24, q1_low);
|
||||
}}
|
||||
}}
|
||||
#endif
|
||||
}}
|
||||
|
||||
static bool is_type_supported(enum ggml_type type) {{
|
||||
if (type == GGML_TYPE_Q4_0 ||
|
||||
type == GGML_TYPE_TL1) {{
|
||||
return true;
|
||||
}} else {{
|
||||
return false;
|
||||
}}
|
||||
}}
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM3200_8640 160
|
||||
#define BBK3200_8640 64
|
||||
inline void tbl_impl_3200_8640(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK3200_8640 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM3200_8640; i += 32) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<4; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 4; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[0] += vec_v_left_1.val[0];
|
||||
vec_c[0] += vec_v_right_1.val[0];
|
||||
vec_c[1] += vec_v_left_1.val[1];
|
||||
vec_c[1] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[2] += vec_v_left_2.val[0];
|
||||
vec_c[2] += vec_v_right_2.val[0];
|
||||
vec_c[3] += vec_v_left_2.val[1];
|
||||
vec_c[3] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[2] += vec_v_left_3.val[0];
|
||||
vec_c[2] += vec_v_right_3.val[0];
|
||||
vec_c[3] += vec_v_left_3.val[1];
|
||||
vec_c[3] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_3200_8640(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM3200_8640];
|
||||
memset(&(CBits[0]), 0, BM3200_8640 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 8640 / BBK3200_8640; ++k_outer) {
|
||||
tbl_impl_3200_8640((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK3200_8640 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK3200_8640 / 2 / 2 * BM3200_8640)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM3200_8640; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM3200_3200 320
|
||||
#define BBK3200_3200 128
|
||||
inline void tbl_impl_3200_3200(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK3200_3200 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[8];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM3200_3200; i += 64) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<8; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 2; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[2] += vec_v_left_1.val[0];
|
||||
vec_c[2] += vec_v_right_1.val[0];
|
||||
vec_c[3] += vec_v_left_1.val[1];
|
||||
vec_c[3] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[4] += vec_v_left_2.val[0];
|
||||
vec_c[4] += vec_v_right_2.val[0];
|
||||
vec_c[5] += vec_v_left_2.val[1];
|
||||
vec_c[5] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[6] += vec_v_left_3.val[0];
|
||||
vec_c[6] += vec_v_right_3.val[0];
|
||||
vec_c[7] += vec_v_left_3.val[1];
|
||||
vec_c[7] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
int32x4_t vec_v_bot_low_low_4 = vmovl_s16(vget_low_s16(vec_c[4]));
|
||||
int32x4_t vec_v_bot_low_high_4 = vmovl_high_s16(vec_c[4]);
|
||||
vst1q_s32(c + i + 32, vld1q_s32(c + i + 32) + vec_v_bot_low_low_4);
|
||||
vst1q_s32(c + i + 36, vld1q_s32(c + i + 36) + vec_v_bot_low_high_4);
|
||||
int32x4_t vec_v_bot_low_low_5 = vmovl_s16(vget_low_s16(vec_c[5]));
|
||||
int32x4_t vec_v_bot_low_high_5 = vmovl_high_s16(vec_c[5]);
|
||||
vst1q_s32(c + i + 40, vld1q_s32(c + i + 40) + vec_v_bot_low_low_5);
|
||||
vst1q_s32(c + i + 44, vld1q_s32(c + i + 44) + vec_v_bot_low_high_5);
|
||||
int32x4_t vec_v_bot_low_low_6 = vmovl_s16(vget_low_s16(vec_c[6]));
|
||||
int32x4_t vec_v_bot_low_high_6 = vmovl_high_s16(vec_c[6]);
|
||||
vst1q_s32(c + i + 48, vld1q_s32(c + i + 48) + vec_v_bot_low_low_6);
|
||||
vst1q_s32(c + i + 52, vld1q_s32(c + i + 52) + vec_v_bot_low_high_6);
|
||||
int32x4_t vec_v_bot_low_low_7 = vmovl_s16(vget_low_s16(vec_c[7]));
|
||||
int32x4_t vec_v_bot_low_high_7 = vmovl_high_s16(vec_c[7]);
|
||||
vst1q_s32(c + i + 56, vld1q_s32(c + i + 56) + vec_v_bot_low_low_7);
|
||||
vst1q_s32(c + i + 60, vld1q_s32(c + i + 60) + vec_v_bot_low_high_7);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_3200_3200(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM3200_3200];
|
||||
memset(&(CBits[0]), 0, BM3200_3200 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 3200 / BBK3200_3200; ++k_outer) {
|
||||
tbl_impl_3200_3200((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK3200_3200 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK3200_3200 / 2 / 2 * BM3200_3200)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM3200_3200; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM8640_3200 320
|
||||
#define BBK8640_3200 64
|
||||
inline void tbl_impl_8640_3200(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK8640_3200 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM8640_3200; i += 32) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<4; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 4; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[0] += vec_v_left_1.val[0];
|
||||
vec_c[0] += vec_v_right_1.val[0];
|
||||
vec_c[1] += vec_v_left_1.val[1];
|
||||
vec_c[1] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[2] += vec_v_left_2.val[0];
|
||||
vec_c[2] += vec_v_right_2.val[0];
|
||||
vec_c[3] += vec_v_left_2.val[1];
|
||||
vec_c[3] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[2] += vec_v_left_3.val[0];
|
||||
vec_c[2] += vec_v_right_3.val[0];
|
||||
vec_c[3] += vec_v_left_3.val[1];
|
||||
vec_c[3] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_8640_3200(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM8640_3200];
|
||||
memset(&(CBits[0]), 0, BM8640_3200 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 3200 / BBK8640_3200; ++k_outer) {
|
||||
tbl_impl_8640_3200((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK8640_3200 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK8640_3200 / 2 / 2 * BM8640_3200)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM8640_3200; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
|
||||
template<int K>
|
||||
void preprocessor_k(void* B, void* LUT_Scales, void* QLUT) {{
|
||||
partial_max_reset((&(((bitnet_float_type*)LUT_Scales)[0])));
|
||||
per_tensor_quant(K, (&(((bitnet_float_type*)LUT_Scales)[0])), (&(((bitnet_float_type*)B)[0])));
|
||||
|
||||
lut_ctor<K>((&(((int8_t*)QLUT)[0])), (&(((bitnet_float_type*)B)[0])), (&(((bitnet_float_type*)LUT_Scales)[0])));
|
||||
}}
|
||||
void ggml_preprocessor(int m, int k, void* B, void* LUT_Scales, void* QLUT) {
|
||||
if (m == 3200 && k == 8640) {
|
||||
preprocessor_k<8640>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
else if (m == 3200 && k == 3200) {
|
||||
preprocessor_k<3200>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
else if (m == 8640 && k == 3200) {
|
||||
preprocessor_k<3200>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
}
|
||||
void ggml_qgemm_lut(int m, int k, void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
if (m == 3200 && k == 8640) {
|
||||
qgemm_lut_3200_8640(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
else if (m == 3200 && k == 3200) {
|
||||
qgemm_lut_3200_3200(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
else if (m == 8640 && k == 3200) {
|
||||
qgemm_lut_8640_3200(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
}
|
||||
|
||||
void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor) {
|
||||
if (!(is_type_supported(tensor->type) && tensor->backend == GGML_BACKEND_TYPE_CPU && tensor->extra == nullptr)) {
|
||||
return;
|
||||
}
|
||||
|
||||
int k = tensor->ne[0];
|
||||
int m = tensor->ne[1];
|
||||
const int lut_scales_size = 1;
|
||||
const int scales_size = 1;
|
||||
int bk = 0;
|
||||
int bm = 0;
|
||||
|
||||
if (m == 3200 && k == 8640) {
|
||||
bm = BM3200_8640;
|
||||
bk = BBK3200_8640;
|
||||
}
|
||||
else if (m == 3200 && k == 3200) {
|
||||
bm = BM3200_3200;
|
||||
bk = BBK3200_3200;
|
||||
}
|
||||
else if (m == 8640 && k == 3200) {
|
||||
bm = BM8640_3200;
|
||||
bk = BBK8640_3200;
|
||||
}
|
||||
|
||||
const int n_tile_num = m / bm;
|
||||
const int BK = bk;
|
||||
uint8_t * qweights;
|
||||
bitnet_float_type * scales;
|
||||
|
||||
scales = (bitnet_float_type *) aligned_malloc(sizeof(bitnet_float_type));
|
||||
qweights = (uint8_t *) tensor->data;
|
||||
float * i2_scales = (float * )(qweights + k * m / 4);
|
||||
scales[0] = (bitnet_float_type) i2_scales[0];
|
||||
|
||||
tensor->extra = bitnet_tensor_extras + bitnet_tensor_extras_index;
|
||||
bitnet_tensor_extras[bitnet_tensor_extras_index++] = {
|
||||
/* .lut_scales_size = */ lut_scales_size,
|
||||
/* .scales_size = */ scales_size,
|
||||
/* .n_tile_num = */ n_tile_num,
|
||||
/* .qweights = */ qweights,
|
||||
/* .scales = */ scales
|
||||
};
|
||||
}
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,21 @@
|
||||
[Kernels_0]
|
||||
m = 3200
|
||||
k = 8640
|
||||
bm = 160
|
||||
bk = 64
|
||||
bmm = 32
|
||||
|
||||
[Kernels_1]
|
||||
m = 3200
|
||||
k = 3200
|
||||
bm = 320
|
||||
bk = 128
|
||||
bmm = 64
|
||||
|
||||
[Kernels_2]
|
||||
m = 8640
|
||||
k = 3200
|
||||
bm = 320
|
||||
bk = 64
|
||||
bmm = 32
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
[Kernels_0]
|
||||
m = 3200
|
||||
k = 8640
|
||||
bm = 160
|
||||
bk = 96
|
||||
bmm = 32
|
||||
|
||||
[Kernels_1]
|
||||
m = 3200
|
||||
k = 3200
|
||||
bm = 320
|
||||
bk = 96
|
||||
bmm = 32
|
||||
|
||||
[Kernels_2]
|
||||
m = 8640
|
||||
k = 3200
|
||||
bm = 320
|
||||
bk = 96
|
||||
bmm = 32
|
||||
|
||||
@@ -0,0 +1,627 @@
|
||||
#if defined(GGML_BITNET_ARM_TL1)
|
||||
#include "ggml-bitnet.h"
|
||||
#define GGML_BITNET_MAX_NODES 8192
|
||||
static bool initialized = false;
|
||||
static bitnet_tensor_extra * bitnet_tensor_extras = nullptr;
|
||||
static size_t bitnet_tensor_extras_index = 0;
|
||||
static void * aligned_malloc(size_t size) {{
|
||||
#if defined(_WIN32)
|
||||
return _aligned_malloc(size, 64);
|
||||
#else
|
||||
void * ptr = nullptr;
|
||||
posix_memalign(&ptr, 64, size);
|
||||
return ptr;
|
||||
#endif
|
||||
}}
|
||||
static void aligned_free(void * ptr) {{
|
||||
#if defined(_WIN32)
|
||||
_aligned_free(ptr);
|
||||
#else
|
||||
free(ptr);
|
||||
#endif
|
||||
}}
|
||||
|
||||
void per_tensor_quant(int k, void* lut_scales_, void* b_) {{
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;
|
||||
bitnet_float_type* b = (bitnet_float_type*)b_;
|
||||
#ifdef __ARM_NEON
|
||||
float32x4_t temp_max = vdupq_n_f32(0);
|
||||
for (int i=0; i < k / 4; i++) {{
|
||||
float32x4_t vec_bs = vld1q_f32(b + 4 * i);
|
||||
float32x4_t abssum = vabsq_f32(vec_bs);
|
||||
temp_max = vmaxq_f32(abssum, temp_max);
|
||||
}}
|
||||
float32_t scales = 127 / vmaxvq_f32(temp_max);
|
||||
*lut_scales = scales;
|
||||
#elif defined __AVX2__
|
||||
__m256 max_vec = _mm256_set1_ps(0.f);
|
||||
const __m256 vec_sign = _mm256_set1_ps(-0.0f);
|
||||
// #pragma unroll
|
||||
for (int i = 0; i < k / 8; i++) {{
|
||||
__m256 vec_b = _mm256_loadu_ps(b + i * 8);
|
||||
__m256 vec_babs = _mm256_andnot_ps(vec_sign, vec_b);
|
||||
max_vec = _mm256_max_ps(vec_babs, max_vec);
|
||||
}}
|
||||
__m128 max1 = _mm_max_ps(_mm256_extractf128_ps(max_vec, 1), _mm256_castps256_ps128(max_vec));
|
||||
max1 = _mm_max_ps(max1, _mm_movehl_ps(max1, max1));
|
||||
max1 = _mm_max_ss(max1, _mm_movehdup_ps(max1));
|
||||
float scales = 127 / _mm_cvtss_f32(max1);
|
||||
*lut_scales = scales;
|
||||
#endif
|
||||
}}
|
||||
|
||||
void partial_max_reset(void* lut_scales_) {{
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;
|
||||
*lut_scales = 0.0;
|
||||
}}
|
||||
|
||||
#ifdef __ARM_NEON
|
||||
inline void Transpose_8_8(
|
||||
int16x8_t *v0,
|
||||
int16x8_t *v1,
|
||||
int16x8_t *v2,
|
||||
int16x8_t *v3,
|
||||
int16x8_t *v4,
|
||||
int16x8_t *v5,
|
||||
int16x8_t *v6,
|
||||
int16x8_t *v7)
|
||||
{{
|
||||
int16x8x2_t q04 = vzipq_s16(*v0, *v4);
|
||||
int16x8x2_t q15 = vzipq_s16(*v1, *v5);
|
||||
int16x8x2_t q26 = vzipq_s16(*v2, *v6);
|
||||
int16x8x2_t q37 = vzipq_s16(*v3, *v7);
|
||||
|
||||
int16x8x2_t q0246_0 = vzipq_s16(q04.val[0], q26.val[0]);
|
||||
int16x8x2_t q0246_1 = vzipq_s16(q04.val[1], q26.val[1]);
|
||||
int16x8x2_t q1357_0 = vzipq_s16(q15.val[0], q37.val[0]);
|
||||
int16x8x2_t q1357_1 = vzipq_s16(q15.val[1], q37.val[1]);
|
||||
|
||||
int16x8x2_t q_fin_0 = vzipq_s16(q0246_0.val[0], q1357_0.val[0]);
|
||||
int16x8x2_t q_fin_1 = vzipq_s16(q0246_0.val[1], q1357_0.val[1]);
|
||||
int16x8x2_t q_fin_2 = vzipq_s16(q0246_1.val[0], q1357_1.val[0]);
|
||||
int16x8x2_t q_fin_3 = vzipq_s16(q0246_1.val[1], q1357_1.val[1]);
|
||||
|
||||
*v0 = q_fin_0.val[0];
|
||||
*v1 = q_fin_0.val[1];
|
||||
*v2 = q_fin_1.val[0];
|
||||
*v3 = q_fin_1.val[1];
|
||||
*v4 = q_fin_2.val[0];
|
||||
*v5 = q_fin_2.val[1];
|
||||
*v6 = q_fin_3.val[0];
|
||||
*v7 = q_fin_3.val[1];
|
||||
}}
|
||||
#endif
|
||||
|
||||
template<int act_k>
|
||||
inline void lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_type* lut_scales) {{
|
||||
#ifdef __ARM_NEON
|
||||
int16x8_t vec_lut[16];
|
||||
float32_t scales = *lut_scales;
|
||||
uint8_t tbl_mask[16];
|
||||
tbl_mask[0] = 0;
|
||||
tbl_mask[1] = 2;
|
||||
tbl_mask[2] = 4;
|
||||
tbl_mask[3] = 6;
|
||||
tbl_mask[4] = 8;
|
||||
tbl_mask[5] = 10;
|
||||
tbl_mask[6] = 12;
|
||||
tbl_mask[7] = 14;
|
||||
tbl_mask[8] = 1;
|
||||
tbl_mask[9] = 3;
|
||||
tbl_mask[10] = 5;
|
||||
tbl_mask[11] = 7;
|
||||
tbl_mask[12] = 9;
|
||||
tbl_mask[13] = 11;
|
||||
tbl_mask[14] = 13;
|
||||
tbl_mask[15] = 15;
|
||||
uint8x16_t tbl_mask_q = vld1q_u8(tbl_mask);
|
||||
#pragma unroll
|
||||
for (int k = 0; k < act_k / 16; ++k) {{
|
||||
float32x4x2_t vec_bs_x0 = vld2q_f32(b + k * 16);
|
||||
float32x4x2_t vec_bs_x1 = vld2q_f32(b + k * 16 + 8);
|
||||
float32x4_t vec_f_0 = vmulq_n_f32(vec_bs_x0.val[0], scales);
|
||||
float32x4_t vec_f_1 = vmulq_n_f32(vec_bs_x0.val[1], scales);
|
||||
float32x4_t vec_f_2 = vmulq_n_f32(vec_bs_x1.val[0], scales);
|
||||
float32x4_t vec_f_3 = vmulq_n_f32(vec_bs_x1.val[1], scales);
|
||||
int32x4_t vec_b_0 = vcvtnq_s32_f32(vec_f_0);
|
||||
int32x4_t vec_b_1 = vcvtnq_s32_f32(vec_f_1);
|
||||
int32x4_t vec_b_2 = vcvtnq_s32_f32(vec_f_2);
|
||||
int32x4_t vec_b_3 = vcvtnq_s32_f32(vec_f_3);
|
||||
int16x4_t vec_b16_0 = vmovn_s32(vec_b_0);
|
||||
int16x4_t vec_b16_1 = vmovn_s32(vec_b_1);
|
||||
int16x4_t vec_b16_2 = vmovn_s32(vec_b_2);
|
||||
int16x4_t vec_b16_3 = vmovn_s32(vec_b_3);
|
||||
int16x8_t vec_bs_0 = vcombine_s16(vec_b16_0, vec_b16_2);
|
||||
int16x8_t vec_bs_1 = vcombine_s16(vec_b16_1, vec_b16_3);
|
||||
vec_lut[0] = vdupq_n_s16(0);
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_0;
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_1;
|
||||
vec_lut[1] = vdupq_n_s16(0);
|
||||
vec_lut[1] = vec_lut[1] - vec_bs_0;
|
||||
vec_lut[2] = vdupq_n_s16(0);
|
||||
vec_lut[2] = vec_lut[2] - vec_bs_0;
|
||||
vec_lut[2] = vec_lut[2] + vec_bs_1;
|
||||
vec_lut[3] = vdupq_n_s16(0);
|
||||
vec_lut[3] = vec_lut[3] - vec_bs_1;
|
||||
vec_lut[4] = vdupq_n_s16(0);
|
||||
vec_lut[5] = vec_bs_1;
|
||||
vec_lut[6] = vec_bs_0;
|
||||
vec_lut[6] = vec_lut[6] - vec_bs_1;
|
||||
vec_lut[7] = vec_bs_0;
|
||||
vec_lut[8] = vec_bs_0;
|
||||
vec_lut[8] = vec_lut[8] + vec_bs_1;
|
||||
Transpose_8_8(&(vec_lut[0]), &(vec_lut[1]), &(vec_lut[2]), &(vec_lut[3]),
|
||||
&(vec_lut[4]), &(vec_lut[5]), &(vec_lut[6]), &(vec_lut[7]));
|
||||
Transpose_8_8(&(vec_lut[8]), &(vec_lut[9]), &(vec_lut[10]), &(vec_lut[11]),
|
||||
&(vec_lut[12]), &(vec_lut[13]), &(vec_lut[14]), &(vec_lut[15]));
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < 8; idx++) {{
|
||||
int8x16_t q0_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx]), tbl_mask_q);
|
||||
int8x8_t q0_low = vget_low_s8(q0_s);
|
||||
int8x8_t q0_high = vget_high_s8(q0_s);
|
||||
int8x16_t q1_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx + 8]), tbl_mask_q);
|
||||
int8x8_t q1_low = vget_low_s8(q1_s);
|
||||
int8x8_t q1_high = vget_high_s8(q1_s);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2, q0_high);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 8, q1_high);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 16, q0_low);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 24, q1_low);
|
||||
}}
|
||||
}}
|
||||
#endif
|
||||
}}
|
||||
|
||||
static bool is_type_supported(enum ggml_type type) {{
|
||||
if (type == GGML_TYPE_Q4_0 ||
|
||||
type == GGML_TYPE_TL1) {{
|
||||
return true;
|
||||
}} else {{
|
||||
return false;
|
||||
}}
|
||||
}}
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM1536_4096 256
|
||||
#define BBK1536_4096 128
|
||||
inline void tbl_impl_1536_4096(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK1536_4096 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM1536_4096; i += 32) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<4; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 4; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[0] += vec_v_left_1.val[0];
|
||||
vec_c[0] += vec_v_right_1.val[0];
|
||||
vec_c[1] += vec_v_left_1.val[1];
|
||||
vec_c[1] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[2] += vec_v_left_2.val[0];
|
||||
vec_c[2] += vec_v_right_2.val[0];
|
||||
vec_c[3] += vec_v_left_2.val[1];
|
||||
vec_c[3] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[2] += vec_v_left_3.val[0];
|
||||
vec_c[2] += vec_v_right_3.val[0];
|
||||
vec_c[3] += vec_v_left_3.val[1];
|
||||
vec_c[3] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_1536_4096(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM1536_4096];
|
||||
memset(&(CBits[0]), 0, BM1536_4096 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 4096 / BBK1536_4096; ++k_outer) {
|
||||
tbl_impl_1536_4096((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK1536_4096 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK1536_4096 / 2 / 2 * BM1536_4096)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM1536_4096; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM1536_1536 128
|
||||
#define BBK1536_1536 64
|
||||
inline void tbl_impl_1536_1536(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK1536_1536 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[8];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM1536_1536; i += 64) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<8; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 2; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[2] += vec_v_left_1.val[0];
|
||||
vec_c[2] += vec_v_right_1.val[0];
|
||||
vec_c[3] += vec_v_left_1.val[1];
|
||||
vec_c[3] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[4] += vec_v_left_2.val[0];
|
||||
vec_c[4] += vec_v_right_2.val[0];
|
||||
vec_c[5] += vec_v_left_2.val[1];
|
||||
vec_c[5] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[6] += vec_v_left_3.val[0];
|
||||
vec_c[6] += vec_v_right_3.val[0];
|
||||
vec_c[7] += vec_v_left_3.val[1];
|
||||
vec_c[7] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
int32x4_t vec_v_bot_low_low_4 = vmovl_s16(vget_low_s16(vec_c[4]));
|
||||
int32x4_t vec_v_bot_low_high_4 = vmovl_high_s16(vec_c[4]);
|
||||
vst1q_s32(c + i + 32, vld1q_s32(c + i + 32) + vec_v_bot_low_low_4);
|
||||
vst1q_s32(c + i + 36, vld1q_s32(c + i + 36) + vec_v_bot_low_high_4);
|
||||
int32x4_t vec_v_bot_low_low_5 = vmovl_s16(vget_low_s16(vec_c[5]));
|
||||
int32x4_t vec_v_bot_low_high_5 = vmovl_high_s16(vec_c[5]);
|
||||
vst1q_s32(c + i + 40, vld1q_s32(c + i + 40) + vec_v_bot_low_low_5);
|
||||
vst1q_s32(c + i + 44, vld1q_s32(c + i + 44) + vec_v_bot_low_high_5);
|
||||
int32x4_t vec_v_bot_low_low_6 = vmovl_s16(vget_low_s16(vec_c[6]));
|
||||
int32x4_t vec_v_bot_low_high_6 = vmovl_high_s16(vec_c[6]);
|
||||
vst1q_s32(c + i + 48, vld1q_s32(c + i + 48) + vec_v_bot_low_low_6);
|
||||
vst1q_s32(c + i + 52, vld1q_s32(c + i + 52) + vec_v_bot_low_high_6);
|
||||
int32x4_t vec_v_bot_low_low_7 = vmovl_s16(vget_low_s16(vec_c[7]));
|
||||
int32x4_t vec_v_bot_low_high_7 = vmovl_high_s16(vec_c[7]);
|
||||
vst1q_s32(c + i + 56, vld1q_s32(c + i + 56) + vec_v_bot_low_low_7);
|
||||
vst1q_s32(c + i + 60, vld1q_s32(c + i + 60) + vec_v_bot_low_high_7);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_1536_1536(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM1536_1536];
|
||||
memset(&(CBits[0]), 0, BM1536_1536 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 1536 / BBK1536_1536; ++k_outer) {
|
||||
tbl_impl_1536_1536((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK1536_1536 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK1536_1536 / 2 / 2 * BM1536_1536)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM1536_1536; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM4096_1536 256
|
||||
#define BBK4096_1536 128
|
||||
inline void tbl_impl_4096_1536(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK4096_1536 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM4096_1536; i += 32) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<4; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 4; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[0] += vec_v_left_1.val[0];
|
||||
vec_c[0] += vec_v_right_1.val[0];
|
||||
vec_c[1] += vec_v_left_1.val[1];
|
||||
vec_c[1] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[2] += vec_v_left_2.val[0];
|
||||
vec_c[2] += vec_v_right_2.val[0];
|
||||
vec_c[3] += vec_v_left_2.val[1];
|
||||
vec_c[3] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[2] += vec_v_left_3.val[0];
|
||||
vec_c[2] += vec_v_right_3.val[0];
|
||||
vec_c[3] += vec_v_left_3.val[1];
|
||||
vec_c[3] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_4096_1536(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM4096_1536];
|
||||
memset(&(CBits[0]), 0, BM4096_1536 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 1536 / BBK4096_1536; ++k_outer) {
|
||||
tbl_impl_4096_1536((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK4096_1536 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK4096_1536 / 2 / 2 * BM4096_1536)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM4096_1536; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
|
||||
template<int K>
|
||||
void preprocessor_k(void* B, void* LUT_Scales, void* QLUT) {{
|
||||
partial_max_reset((&(((bitnet_float_type*)LUT_Scales)[0])));
|
||||
per_tensor_quant(K, (&(((bitnet_float_type*)LUT_Scales)[0])), (&(((bitnet_float_type*)B)[0])));
|
||||
|
||||
lut_ctor<K>((&(((int8_t*)QLUT)[0])), (&(((bitnet_float_type*)B)[0])), (&(((bitnet_float_type*)LUT_Scales)[0])));
|
||||
}}
|
||||
void ggml_preprocessor(int m, int k, void* B, void* LUT_Scales, void* QLUT) {
|
||||
if (m == 1536 && k == 4096) {
|
||||
preprocessor_k<4096>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
else if (m == 1536 && k == 1536) {
|
||||
preprocessor_k<1536>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
else if (m == 4096 && k == 1536) {
|
||||
preprocessor_k<1536>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
}
|
||||
void ggml_qgemm_lut(int m, int k, void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
if (m == 1536 && k == 4096) {
|
||||
qgemm_lut_1536_4096(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
else if (m == 1536 && k == 1536) {
|
||||
qgemm_lut_1536_1536(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
else if (m == 4096 && k == 1536) {
|
||||
qgemm_lut_4096_1536(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
}
|
||||
|
||||
void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor) {
|
||||
if (!(is_type_supported(tensor->type) && tensor->backend == GGML_BACKEND_TYPE_CPU && tensor->extra == nullptr)) {
|
||||
return;
|
||||
}
|
||||
|
||||
int k = tensor->ne[0];
|
||||
int m = tensor->ne[1];
|
||||
const int lut_scales_size = 1;
|
||||
const int scales_size = 1;
|
||||
int bk = 0;
|
||||
int bm = 0;
|
||||
|
||||
if (m == 1536 && k == 4096) {
|
||||
bm = BM1536_4096;
|
||||
bk = BBK1536_4096;
|
||||
}
|
||||
else if (m == 1536 && k == 1536) {
|
||||
bm = BM1536_1536;
|
||||
bk = BBK1536_1536;
|
||||
}
|
||||
else if (m == 4096 && k == 1536) {
|
||||
bm = BM4096_1536;
|
||||
bk = BBK4096_1536;
|
||||
}
|
||||
|
||||
const int n_tile_num = m / bm;
|
||||
const int BK = bk;
|
||||
uint8_t * qweights;
|
||||
bitnet_float_type * scales;
|
||||
|
||||
scales = (bitnet_float_type *) aligned_malloc(sizeof(bitnet_float_type));
|
||||
qweights = (uint8_t *) tensor->data;
|
||||
float * i2_scales = (float * )(qweights + k * m / 4);
|
||||
scales[0] = (bitnet_float_type) i2_scales[0];
|
||||
|
||||
tensor->extra = bitnet_tensor_extras + bitnet_tensor_extras_index;
|
||||
bitnet_tensor_extras[bitnet_tensor_extras_index++] = {
|
||||
/* .lut_scales_size = */ lut_scales_size,
|
||||
/* .scales_size = */ scales_size,
|
||||
/* .n_tile_num = */ n_tile_num,
|
||||
/* .qweights = */ qweights,
|
||||
/* .scales = */ scales
|
||||
};
|
||||
}
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,21 @@
|
||||
[Kernels_0]
|
||||
m = 1536
|
||||
k = 4096
|
||||
bm = 256
|
||||
bk = 128
|
||||
bmm = 32
|
||||
|
||||
[Kernels_1]
|
||||
m = 1536
|
||||
k = 1536
|
||||
bm = 128
|
||||
bk = 64
|
||||
bmm = 64
|
||||
|
||||
[Kernels_2]
|
||||
m = 4096
|
||||
k = 1536
|
||||
bm = 256
|
||||
bk = 128
|
||||
bmm = 32
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
[Kernels_0]
|
||||
m = 1536
|
||||
k = 4096
|
||||
bm = 256
|
||||
bk = 96
|
||||
bmm = 32
|
||||
|
||||
[Kernels_1]
|
||||
m = 1536
|
||||
k = 1536
|
||||
bm = 128
|
||||
bk = 192
|
||||
bmm = 32
|
||||
|
||||
[Kernels_2]
|
||||
m = 4096
|
||||
k = 1536
|
||||
bm = 256
|
||||
bk = 96
|
||||
bmm = 64
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
# These requirements include all dependencies for all top-level python scripts
|
||||
# for llama.cpp. Avoid adding packages here directly.
|
||||
#
|
||||
# Package versions must stay compatible across all top-level python scripts.
|
||||
#
|
||||
|
||||
-r 3rdparty/llama.cpp/requirements/requirements-convert_legacy_llama.txt
|
||||
-r 3rdparty/llama.cpp/requirements/requirements-convert_hf_to_gguf.txt
|
||||
-r 3rdparty/llama.cpp/requirements/requirements-convert_hf_to_gguf_update.txt
|
||||
-r 3rdparty/llama.cpp/requirements/requirements-convert_llama_ggml_to_gguf.txt
|
||||
-r 3rdparty/llama.cpp/requirements/requirements-convert_lora_to_gguf.txt
|
||||
@@ -0,0 +1,56 @@
|
||||
import os
|
||||
import sys
|
||||
import signal
|
||||
import platform
|
||||
import argparse
|
||||
import subprocess
|
||||
|
||||
def run_command(command, shell=False):
|
||||
"""Run a system command and ensure it succeeds."""
|
||||
try:
|
||||
subprocess.run(command, shell=shell, check=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"Error occurred while running command: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
def run_inference():
|
||||
build_dir = "build"
|
||||
if platform.system() == "Windows":
|
||||
main_path = os.path.join(build_dir, "bin", "Release", "llama-cli.exe")
|
||||
if not os.path.exists(main_path):
|
||||
main_path = os.path.join(build_dir, "bin", "llama-cli")
|
||||
else:
|
||||
main_path = os.path.join(build_dir, "bin", "llama-cli")
|
||||
command = [
|
||||
f'{main_path}',
|
||||
'-m', args.model,
|
||||
'-n', str(args.n_predict),
|
||||
'-t', str(args.threads),
|
||||
'-p', args.prompt,
|
||||
'-ngl', '0',
|
||||
'-c', str(args.ctx_size),
|
||||
'--temp', str(args.temperature),
|
||||
"-b", "1",
|
||||
]
|
||||
if args.conversation:
|
||||
command.append("-cnv")
|
||||
run_command(command)
|
||||
|
||||
def signal_handler(sig, frame):
|
||||
print("Ctrl+C pressed, exiting...")
|
||||
sys.exit(0)
|
||||
|
||||
if __name__ == "__main__":
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
# Usage: python run_inference.py -p "Microsoft Corporation is an American multinational corporation and technology company headquartered in Redmond, Washington."
|
||||
parser = argparse.ArgumentParser(description='Run inference')
|
||||
parser.add_argument("-m", "--model", type=str, help="Path to model file", required=False, default="models/bitnet_b1_58-3B/ggml-model-i2_s.gguf")
|
||||
parser.add_argument("-n", "--n-predict", type=int, help="Number of tokens to predict when generating text", required=False, default=128)
|
||||
parser.add_argument("-p", "--prompt", type=str, help="Prompt to generate text from", required=True)
|
||||
parser.add_argument("-t", "--threads", type=int, help="Number of threads to use", required=False, default=2)
|
||||
parser.add_argument("-c", "--ctx-size", type=int, help="Size of the prompt context", required=False, default=2048)
|
||||
parser.add_argument("-temp", "--temperature", type=float, help="Temperature, a hyperparameter that controls the randomness of the generated text", required=False, default=0.8)
|
||||
parser.add_argument("-cnv", "--conversation", action='store_true', help="Whether to enable chat mode or not (for instruct models.)")
|
||||
|
||||
args = parser.parse_args()
|
||||
run_inference()
|
||||
@@ -0,0 +1,64 @@
|
||||
import os
|
||||
import sys
|
||||
import signal
|
||||
import platform
|
||||
import argparse
|
||||
import subprocess
|
||||
|
||||
def run_command(command, shell=False):
|
||||
"""Run a system command and ensure it succeeds."""
|
||||
try:
|
||||
subprocess.run(command, shell=shell, check=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"Error occurred while running command: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
def run_server():
|
||||
build_dir = "build"
|
||||
if platform.system() == "Windows":
|
||||
server_path = os.path.join(build_dir, "bin", "Release", "llama-server.exe")
|
||||
if not os.path.exists(server_path):
|
||||
server_path = os.path.join(build_dir, "bin", "llama-server")
|
||||
else:
|
||||
server_path = os.path.join(build_dir, "bin", "llama-server")
|
||||
|
||||
command = [
|
||||
f'{server_path}',
|
||||
'-m', args.model,
|
||||
'-c', str(args.ctx_size),
|
||||
'-t', str(args.threads),
|
||||
'-n', str(args.n_predict),
|
||||
'-ngl', '0',
|
||||
'--temp', str(args.temperature),
|
||||
'--host', args.host,
|
||||
'--port', str(args.port),
|
||||
'-cb' # Enable continuous batching
|
||||
]
|
||||
|
||||
if args.prompt:
|
||||
command.extend(['-p', args.prompt])
|
||||
|
||||
# Note: -cnv flag is removed as it's not supported by the server
|
||||
|
||||
print(f"Starting server on {args.host}:{args.port}")
|
||||
run_command(command)
|
||||
|
||||
def signal_handler(sig, frame):
|
||||
print("Ctrl+C pressed, shutting down server...")
|
||||
sys.exit(0)
|
||||
|
||||
if __name__ == "__main__":
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
|
||||
parser = argparse.ArgumentParser(description='Run llama.cpp server')
|
||||
parser.add_argument("-m", "--model", type=str, help="Path to model file", required=False, default="models/bitnet_b1_58-3B/ggml-model-i2_s.gguf")
|
||||
parser.add_argument("-p", "--prompt", type=str, help="System prompt for the model", required=False)
|
||||
parser.add_argument("-n", "--n-predict", type=int, help="Number of tokens to predict", required=False, default=4096)
|
||||
parser.add_argument("-t", "--threads", type=int, help="Number of threads to use", required=False, default=2)
|
||||
parser.add_argument("-c", "--ctx-size", type=int, help="Size of the context window", required=False, default=2048)
|
||||
parser.add_argument("--temperature", type=float, help="Temperature for sampling", required=False, default=0.8)
|
||||
parser.add_argument("--host", type=str, help="IP address to listen on", required=False, default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, help="Port to listen on", required=False, default=8080)
|
||||
|
||||
args = parser.parse_args()
|
||||
run_server()
|
||||
+244
@@ -0,0 +1,244 @@
|
||||
import subprocess
|
||||
import signal
|
||||
import sys
|
||||
import os
|
||||
import platform
|
||||
import argparse
|
||||
import logging
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger("setup_env")
|
||||
|
||||
SUPPORTED_HF_MODELS = {
|
||||
"1bitLLM/bitnet_b1_58-large": {
|
||||
"model_name": "bitnet_b1_58-large",
|
||||
},
|
||||
"1bitLLM/bitnet_b1_58-3B": {
|
||||
"model_name": "bitnet_b1_58-3B",
|
||||
},
|
||||
"HF1BitLLM/Llama3-8B-1.58-100B-tokens": {
|
||||
"model_name": "Llama3-8B-1.58-100B-tokens",
|
||||
},
|
||||
"tiiuae/Falcon3-7B-Instruct-1.58bit": {
|
||||
"model_name": "Falcon3-7B-Instruct-1.58bit",
|
||||
},
|
||||
"tiiuae/Falcon3-7B-1.58bit": {
|
||||
"model_name": "Falcon3-7B-1.58bit",
|
||||
},
|
||||
"tiiuae/Falcon3-10B-Instruct-1.58bit": {
|
||||
"model_name": "Falcon3-10B-Instruct-1.58bit",
|
||||
},
|
||||
"tiiuae/Falcon3-10B-1.58bit": {
|
||||
"model_name": "Falcon3-10B-1.58bit",
|
||||
},
|
||||
"tiiuae/Falcon3-3B-Instruct-1.58bit": {
|
||||
"model_name": "Falcon3-3B-Instruct-1.58bit",
|
||||
},
|
||||
"tiiuae/Falcon3-3B-1.58bit": {
|
||||
"model_name": "Falcon3-3B-1.58bit",
|
||||
},
|
||||
"tiiuae/Falcon3-1B-Instruct-1.58bit": {
|
||||
"model_name": "Falcon3-1B-Instruct-1.58bit",
|
||||
},
|
||||
"microsoft/BitNet-b1.58-2B-4T": {
|
||||
"model_name": "BitNet-b1.58-2B-4T",
|
||||
},
|
||||
"tiiuae/Falcon-E-3B-Instruct": {
|
||||
"model_name": "Falcon-E-3B-Instruct",
|
||||
},
|
||||
"tiiuae/Falcon-E-1B-Instruct": {
|
||||
"model_name": "Falcon-E-1B-Instruct",
|
||||
},
|
||||
"tiiuae/Falcon-E-3B-Base": {
|
||||
"model_name": "Falcon-E-3B-Base",
|
||||
},
|
||||
"tiiuae/Falcon-E-1B-Base": {
|
||||
"model_name": "Falcon-E-1B-Base",
|
||||
},
|
||||
}
|
||||
|
||||
SUPPORTED_QUANT_TYPES = {
|
||||
"arm64": ["i2_s", "tl1"],
|
||||
"x86_64": ["i2_s", "tl2"]
|
||||
}
|
||||
|
||||
COMPILER_EXTRA_ARGS = {
|
||||
"arm64": ["-DBITNET_ARM_TL1=OFF"],
|
||||
"x86_64": ["-DBITNET_X86_TL2=OFF"]
|
||||
}
|
||||
|
||||
OS_EXTRA_ARGS = {
|
||||
"Windows":["-T", "ClangCL"],
|
||||
}
|
||||
|
||||
ARCH_ALIAS = {
|
||||
"AMD64": "x86_64",
|
||||
"x86": "x86_64",
|
||||
"x86_64": "x86_64",
|
||||
"aarch64": "arm64",
|
||||
"arm64": "arm64",
|
||||
"ARM64": "arm64",
|
||||
}
|
||||
|
||||
def system_info():
|
||||
return platform.system(), ARCH_ALIAS[platform.machine()]
|
||||
|
||||
def get_model_name():
|
||||
if args.hf_repo:
|
||||
return SUPPORTED_HF_MODELS[args.hf_repo]["model_name"]
|
||||
return os.path.basename(os.path.normpath(args.model_dir))
|
||||
|
||||
def run_command(command, shell=False, log_step=None):
|
||||
"""Run a system command and ensure it succeeds."""
|
||||
if log_step:
|
||||
log_file = os.path.join(args.log_dir, log_step + ".log")
|
||||
with open(log_file, "w") as f:
|
||||
try:
|
||||
subprocess.run(command, shell=shell, check=True, stdout=f, stderr=f)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logging.error(f"Error occurred while running command: {e}, check details in {log_file}")
|
||||
sys.exit(1)
|
||||
else:
|
||||
try:
|
||||
subprocess.run(command, shell=shell, check=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logging.error(f"Error occurred while running command: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
def prepare_model():
|
||||
_, arch = system_info()
|
||||
hf_url = args.hf_repo
|
||||
model_dir = args.model_dir
|
||||
quant_type = args.quant_type
|
||||
quant_embd = args.quant_embd
|
||||
if hf_url is not None:
|
||||
# download the model
|
||||
model_dir = os.path.join(model_dir, SUPPORTED_HF_MODELS[hf_url]["model_name"])
|
||||
Path(model_dir).mkdir(parents=True, exist_ok=True)
|
||||
logging.info(f"Downloading model {hf_url} from HuggingFace to {model_dir}...")
|
||||
run_command(["huggingface-cli", "download", hf_url, "--local-dir", model_dir], log_step="download_model")
|
||||
elif not os.path.exists(model_dir):
|
||||
logging.error(f"Model directory {model_dir} does not exist.")
|
||||
sys.exit(1)
|
||||
else:
|
||||
logging.info(f"Loading model from directory {model_dir}.")
|
||||
gguf_path = os.path.join(model_dir, "ggml-model-" + quant_type + ".gguf")
|
||||
if not os.path.exists(gguf_path) or os.path.getsize(gguf_path) == 0:
|
||||
logging.info(f"Converting HF model to GGUF format...")
|
||||
if quant_type.startswith("tl"):
|
||||
run_command([sys.executable, "utils/convert-hf-to-gguf-bitnet.py", model_dir, "--outtype", quant_type, "--quant-embd"], log_step="convert_to_tl")
|
||||
else: # i2s
|
||||
# convert to f32
|
||||
run_command([sys.executable, "utils/convert-hf-to-gguf-bitnet.py", model_dir, "--outtype", "f32"], log_step="convert_to_f32_gguf")
|
||||
f32_model = os.path.join(model_dir, "ggml-model-f32.gguf")
|
||||
i2s_model = os.path.join(model_dir, "ggml-model-i2_s.gguf")
|
||||
# quantize to i2s
|
||||
if platform.system() != "Windows":
|
||||
if quant_embd:
|
||||
run_command(["./build/bin/llama-quantize", "--token-embedding-type", "f16", f32_model, i2s_model, "I2_S", "1", "1"], log_step="quantize_to_i2s")
|
||||
else:
|
||||
run_command(["./build/bin/llama-quantize", f32_model, i2s_model, "I2_S", "1"], log_step="quantize_to_i2s")
|
||||
else:
|
||||
if quant_embd:
|
||||
run_command(["./build/bin/Release/llama-quantize", "--token-embedding-type", "f16", f32_model, i2s_model, "I2_S", "1", "1"], log_step="quantize_to_i2s")
|
||||
else:
|
||||
run_command(["./build/bin/Release/llama-quantize", f32_model, i2s_model, "I2_S", "1"], log_step="quantize_to_i2s")
|
||||
|
||||
logging.info(f"GGUF model saved at {gguf_path}")
|
||||
else:
|
||||
logging.info(f"GGUF model already exists at {gguf_path}")
|
||||
|
||||
def setup_gguf():
|
||||
# Install the pip package
|
||||
run_command([sys.executable, "-m", "pip", "install", "3rdparty/llama.cpp/gguf-py"], log_step="install_gguf")
|
||||
|
||||
def gen_code():
|
||||
_, arch = system_info()
|
||||
|
||||
llama3_f3_models = set([model['model_name'] for model in SUPPORTED_HF_MODELS.values() if model['model_name'].startswith("Falcon") or model['model_name'].startswith("Llama")])
|
||||
|
||||
if arch == "arm64":
|
||||
if args.use_pretuned:
|
||||
pretuned_kernels = os.path.join("preset_kernels", get_model_name())
|
||||
if not os.path.exists(pretuned_kernels):
|
||||
logging.error(f"Pretuned kernels not found for model {args.hf_repo}")
|
||||
sys.exit(1)
|
||||
if args.quant_type == "tl1":
|
||||
shutil.copyfile(os.path.join(pretuned_kernels, "bitnet-lut-kernels-tl1.h"), "include/bitnet-lut-kernels.h")
|
||||
shutil.copyfile(os.path.join(pretuned_kernels, "kernel_config_tl1.ini"), "include/kernel_config.ini")
|
||||
elif args.quant_type == "tl2":
|
||||
shutil.copyfile(os.path.join(pretuned_kernels, "bitnet-lut-kernels-tl2.h"), "include/bitnet-lut-kernels.h")
|
||||
shutil.copyfile(os.path.join(pretuned_kernels, "kernel_config_tl2.ini"), "include/kernel_config.ini")
|
||||
if get_model_name() == "bitnet_b1_58-large":
|
||||
run_command([sys.executable, "utils/codegen_tl1.py", "--model", "bitnet_b1_58-large", "--BM", "256,128,256", "--BK", "128,64,128", "--bm", "32,64,32"], log_step="codegen")
|
||||
elif get_model_name() in llama3_f3_models:
|
||||
run_command([sys.executable, "utils/codegen_tl1.py", "--model", "Llama3-8B-1.58-100B-tokens", "--BM", "256,128,256,128", "--BK", "128,64,128,64", "--bm", "32,64,32,64"], log_step="codegen")
|
||||
elif get_model_name() == "bitnet_b1_58-3B":
|
||||
run_command([sys.executable, "utils/codegen_tl1.py", "--model", "bitnet_b1_58-3B", "--BM", "160,320,320", "--BK", "64,128,64", "--bm", "32,64,32"], log_step="codegen")
|
||||
elif get_model_name() == "BitNet-b1.58-2B-4T":
|
||||
run_command([sys.executable, "utils/codegen_tl1.py", "--model", "bitnet_b1_58-3B", "--BM", "160,320,320", "--BK", "64,128,64", "--bm", "32,64,32"], log_step="codegen")
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
else:
|
||||
if args.use_pretuned:
|
||||
# cp preset_kernels/model_name/bitnet-lut-kernels_tl1.h to include/bitnet-lut-kernels.h
|
||||
pretuned_kernels = os.path.join("preset_kernels", get_model_name())
|
||||
if not os.path.exists(pretuned_kernels):
|
||||
logging.error(f"Pretuned kernels not found for model {args.hf_repo}")
|
||||
sys.exit(1)
|
||||
shutil.copyfile(os.path.join(pretuned_kernels, "bitnet-lut-kernels-tl2.h"), "include/bitnet-lut-kernels.h")
|
||||
if get_model_name() == "bitnet_b1_58-large":
|
||||
run_command([sys.executable, "utils/codegen_tl2.py", "--model", "bitnet_b1_58-large", "--BM", "256,128,256", "--BK", "96,192,96", "--bm", "32,32,32"], log_step="codegen")
|
||||
elif get_model_name() in llama3_f3_models:
|
||||
run_command([sys.executable, "utils/codegen_tl2.py", "--model", "Llama3-8B-1.58-100B-tokens", "--BM", "256,128,256,128", "--BK", "96,96,96,96", "--bm", "32,32,32,32"], log_step="codegen")
|
||||
elif get_model_name() == "bitnet_b1_58-3B":
|
||||
run_command([sys.executable, "utils/codegen_tl2.py", "--model", "bitnet_b1_58-3B", "--BM", "160,320,320", "--BK", "96,96,96", "--bm", "32,32,32"], log_step="codegen")
|
||||
elif get_model_name() == "BitNet-b1.58-2B-4T":
|
||||
run_command([sys.executable, "utils/codegen_tl2.py", "--model", "bitnet_b1_58-3B", "--BM", "160,320,320", "--BK", "96,96,96", "--bm", "32,32,32"], log_step="codegen")
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
def compile():
|
||||
# Check if cmake is installed
|
||||
cmake_exists = subprocess.run(["cmake", "--version"], capture_output=True)
|
||||
if cmake_exists.returncode != 0:
|
||||
logging.error("Cmake is not available. Please install CMake and try again.")
|
||||
sys.exit(1)
|
||||
_, arch = system_info()
|
||||
if arch not in COMPILER_EXTRA_ARGS.keys():
|
||||
logging.error(f"Arch {arch} is not supported yet")
|
||||
exit(0)
|
||||
logging.info("Compiling the code using CMake.")
|
||||
run_command(["cmake", "-B", "build", *COMPILER_EXTRA_ARGS[arch], *OS_EXTRA_ARGS.get(platform.system(), []), "-DCMAKE_C_COMPILER=clang", "-DCMAKE_CXX_COMPILER=clang++"], log_step="generate_build_files")
|
||||
# run_command(["cmake", "--build", "build", "--target", "llama-cli", "--config", "Release"])
|
||||
run_command(["cmake", "--build", "build", "--config", "Release"], log_step="compile")
|
||||
|
||||
def main():
|
||||
setup_gguf()
|
||||
gen_code()
|
||||
compile()
|
||||
prepare_model()
|
||||
|
||||
def parse_args():
|
||||
_, arch = system_info()
|
||||
parser = argparse.ArgumentParser(description='Setup the environment for running the inference')
|
||||
parser.add_argument("--hf-repo", "-hr", type=str, help="Model used for inference", choices=SUPPORTED_HF_MODELS.keys())
|
||||
parser.add_argument("--model-dir", "-md", type=str, help="Directory to save/load the model", default="models")
|
||||
parser.add_argument("--log-dir", "-ld", type=str, help="Directory to save the logging info", default="logs")
|
||||
parser.add_argument("--quant-type", "-q", type=str, help="Quantization type", choices=SUPPORTED_QUANT_TYPES[arch], default="i2_s")
|
||||
parser.add_argument("--quant-embd", action="store_true", help="Quantize the embeddings to f16")
|
||||
parser.add_argument("--use-pretuned", "-p", action="store_true", help="Use the pretuned kernel parameters")
|
||||
return parser.parse_args()
|
||||
|
||||
def signal_handler(sig, frame):
|
||||
logging.info("Ctrl+C pressed, exiting...")
|
||||
sys.exit(0)
|
||||
|
||||
if __name__ == "__main__":
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
args = parse_args()
|
||||
Path(args.log_dir).mkdir(parents=True, exist_ok=True)
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
main()
|
||||
@@ -0,0 +1,10 @@
|
||||
set(GGML_HEADERS_BITNET ../include/ggml-bitnet.h)
|
||||
set(GGML_SOURCES_BITNET ggml-bitnet-mad.cpp)
|
||||
set(GGML_SOURCES_BITNET ggml-bitnet-lut.cpp)
|
||||
|
||||
include_directories(3rdparty/llama.cpp/ggml/include)
|
||||
|
||||
if (NOT (CMAKE_C_COMPILER_ID MATCHES "Clang" OR CMAKE_C_COMPILER_ID STREQUAL "GNU") OR
|
||||
NOT (CMAKE_CXX_COMPILER_ID MATCHES "Clang" OR CMAKE_CXX_COMPILER_ID STREQUAL "GNU"))
|
||||
message(FATAL_ERROR "Clang or GCC is required for Bitnet.cpp compilation")
|
||||
endif()
|
||||
+205
@@ -0,0 +1,205 @@
|
||||
# BitNet CPU Inference Optimization
|
||||
|
||||
This update provides significant performance improvements for BitNet inference on CPU through paralleled kernel implementations, native I2_S GEMM/GEMV support, configurable tiling block size and embedding quantization.
|
||||
|
||||
## Update
|
||||
|
||||
- **Parallel Weight & Activation Computation**
|
||||
Implemented parallel processing of weights and activations in the W2A8 vet_dot kernel, achieving improved throughput on both x86 and ARM architectures.
|
||||
|
||||
- **Native I2_S GEMM & GEMV Support**
|
||||
Integrated I2_S GEMM and GEMV operations into ggml library, making them fully compatible with the llama.cpp architecture. This enables seamless integration with existing inference pipelines.
|
||||
|
||||
- **Configurable Tiling & Parallelism**
|
||||
Introduced configurable GEMM & GEMV block sizes and parallelism levels, allowing performance fine-tuning for different CPU architectures.
|
||||
|
||||
- **Embedding Quantization**
|
||||
Added support for embedding layer quantization with Q6_K format, reducing memory footprint and improving inference speed while maintaining high accuracy.
|
||||
|
||||
## Usage
|
||||
|
||||
### Configuration Options
|
||||
|
||||
The `include/gemm-config.h` file controls kernel behavior:
|
||||
|
||||
```c
|
||||
#define ROW_BLOCK_SIZE 4
|
||||
#define COL_BLOCK_SIZE 128
|
||||
#define PARALLEL_SIZE 4
|
||||
```
|
||||
|
||||
Modify these values based on your CPU cache size and architecture for optimal performance. Users can fine-tune performance on their machine through `include/gemm-config.h`.
|
||||
|
||||
### Enabling Embedding Quantization
|
||||
|
||||
To use embedding quantization for additional speedup:
|
||||
|
||||
**Using setup_env.py:**
|
||||
```bash
|
||||
python setup_env.py --quant-embd
|
||||
```
|
||||
This automatically converts embeddings to Q6_K format.
|
||||
|
||||
**Manual conversion:**
|
||||
```bash
|
||||
build/bin/llama-quantize --token-embedding-type Q6_K models/BitNet-b1.58-2B-4T/ggml-model-f32.gguf models/BitNet-b1.58-2B-4T/ggml-model-i2_s-embed-q6_k.gguf I2_S 1 1
|
||||
```
|
||||
|
||||
## Optimizations
|
||||
|
||||
### 1. Weight & Activation Parallelism
|
||||
|
||||
The kernel implements two parallelization strategies:
|
||||
|
||||
- **Weight Parallel:** Processes multiple weight rows/columns in a single kernel call, reducing kernel launch overhead.
|
||||
|
||||
- **Activation Parallel:** Built on top of weight parallel, amortizes the I2_S weight unpacking cost across multiple activation elements.
|
||||
|
||||
**Recommendation:** For I2_S quantization format, activation parallel is recommended due to the unpack operation benefits. The current kernel defaults to activation parallel.
|
||||
|
||||
**Kernel Performance Comparison:**
|
||||
|
||||
<div align="center">
|
||||
|
||||
Test configuration: AMD EPYC 7V13 (x86), 1 threads, time in milliseconds (mean±std)
|
||||
|
||||
| Matrix Size | No Parallel | Weight Parallel | Activation Parallel |
|
||||
|:---:|:---:|:---:|:---:|
|
||||
| [1, 2048] × [2048, 2048] | 0.075±0.012 | **0.058±0.007** | 0.076±0.011 |
|
||||
| [32, 2048] × [2048, 2048] | 2.400±0.041 | 1.599±0.020 | **1.202±0.018** |
|
||||
| [128, 2048] × [2048, 2048] | 10.820±0.039 | 6.458±0.168 | **5.805±0.039** |
|
||||
| [256, 2048] × [2048, 2048] | 21.669±0.080 | 12.739±0.183 | **11.882±0.040** |
|
||||
| [512, 2048] × [2048, 2048] | 43.257±0.083 | 25.680±0.335 | **23.342±0.082** |
|
||||
| [2048, 2048] × [2048, 2048] | 173.175±0.214 | 103.112±0.552 | **93.276±0.612** |
|
||||
| [128, 2048] × [2048, 8192] | 43.345±0.090 | 25.541±0.239 | **23.528±0.052** |
|
||||
| [128, 8192] × [8192, 2048] | 38.085±0.162 | 23.866±0.096 | **22.569±0.132** |
|
||||
|
||||
</div>
|
||||
|
||||
### 2. GEMM/GEMV Integration with llama.cpp
|
||||
|
||||
Integrated I2_S quantization format into llama.cpp's compute graph:
|
||||
|
||||
- **GEMV Operations:** Optimized matrix-vector multiplication for token generation.
|
||||
- **GEMM Operations:** Efficient matrix-matrix multiplication for prompt processing.
|
||||
- **Tiling Strategy:** Configurable block sizes for optimal cache utilization.
|
||||
|
||||
### 3. Configuration Fine-tuning
|
||||
|
||||
Fine-tuning kernel parameters for optimal performance on specific hardware:
|
||||
|
||||
**Example Configuration (x86, AMD EPYC 7V13):**
|
||||
- Method: Activation Parallel
|
||||
- Threads: 8
|
||||
- Workload: 128 prompt tokens (pp128)
|
||||
|
||||
**Fine-tuning Parameters:**
|
||||
- **Parallelism Degree:** [2, 4, 8]
|
||||
- **Row Block Size:** [2, 4, 8, 16, 32]
|
||||
- **Column Block Size:** [32, 64, 128, 256, 512, 1024]
|
||||
|
||||
**Fine-tuning Results:**
|
||||
|
||||
<div align="center">
|
||||
|
||||
<img src="./assets/fine_tuning_result.png" alt="fine_tune_result" width="800"/>
|
||||
|
||||
*Shows throughput (tokens/s) for various configurations.*
|
||||
|
||||
</div>
|
||||
|
||||
**Optimal Configuration:** Under this setup (x86, 8 threads, pp128), the best performance is achieved with parallelism degree = 4, row block size = 4, and column block size = 128.
|
||||
|
||||
### 4. Embedding Quantization
|
||||
|
||||
Evaluated multiple embedding quantization formats to balance memory usage, model quality, and inference speed:
|
||||
|
||||
**Perplexity Comparison:**
|
||||
|
||||
<div align="center">
|
||||
|
||||
Test configuration: BitNet-b1.58-2B-4T, TG128
|
||||
|
||||
| Embedding Type | Wikitext | PTB | LAMBADA | IMDB | AG NEWS |
|
||||
|:---:|:---:|:---:|:---:|:---:|:---:|
|
||||
| **F32** | 17.1090±0.1278 | 33.0858±0.4886 | 43.2850±0.6363 | 29.3016±0.2890 | 36.7686±0.3920 |
|
||||
| **F16** | 17.1090±0.1278 | 33.0858±0.4886 | 43.2850±0.6363 | 29.3016±0.2890 | 36.7686±0.3920 |
|
||||
| **Q8_0** | 17.1197±0.1280 | 33.1181±0.4893 | 43.2891±0.6364 | 29.3133±0.2892 | 36.7740±0.3920 |
|
||||
| **Q6_K** | 17.1487±0.1282 | 33.2203±0.4914 | 43.3046±0.6362 | 29.3491±0.2897 | 36.7972±0.3921 |
|
||||
| **Q5_0** | 17.2379±0.1288 | 33.2439±0.4907 | 43.4631±0.6379 | 29.5481±0.2920 | 36.8539±0.3924 |
|
||||
| **Q4_0** | 17.3529±0.1300 | 33.7754±0.5001 | 44.4552±0.6559 | 30.1044±0.2978 | 37.3985±0.3997 |
|
||||
| **Q3_K** | 17.6434±0.1320 | 34.3914±0.5089 | 45.4591±0.6735 | 30.8476±0.3069 | 39.5692±0.4259 |
|
||||
| **I2_S** | N/A | N/A | N/A | N/A | N/A |
|
||||
|
||||
**N/A indicates model failure due to extreme quantization.*
|
||||
|
||||
</div>
|
||||
|
||||
**Inference Speed Comparison:**
|
||||
|
||||
<div align="center">
|
||||
|
||||
<img src="./assets/embedding_throughput.png" alt="embedding_throughput" width="800"/>
|
||||
|
||||
*Token generation throughput (tg128) for different embedding quantization types.*
|
||||
|
||||
</div>
|
||||
|
||||
**Recommendation:** Based on comprehensive evaluation of memory footprint, perplexity preservation, and inference speed, **Q6_K** is selected as the optimal embedding quantization format.
|
||||
|
||||
## Performance
|
||||
|
||||
Comparison of optimized parallel kernels vs. original implementation:
|
||||
|
||||
**Test Configuration:**
|
||||
- Model: BitNet-b1.58-2B-4T
|
||||
- Hardware: AMD EPYC 7V13
|
||||
- Threads: 1 / 2 / 4 / 8 / 12 / 16
|
||||
- Test: 128 prompt tokens (pp128) + 128 generated tokens (tg128)
|
||||
- Method: Activation Parallel
|
||||
|
||||
<div align="center">
|
||||
|
||||
<img src="./assets/performance_comparison_amd_epyc.png" alt="performance_comparison_amd_epyc" width="800"/>
|
||||
|
||||
</div>
|
||||
|
||||
**Test Configuration:**
|
||||
- Model: BitNet-b1.58-2B-4T
|
||||
- Hardware: Intel i7-13800H
|
||||
- Threads: 1 / 2 / 4 / 6
|
||||
- Test: 128 prompt tokens (pp128) + 128 generated tokens (tg128)
|
||||
- Method: Activation Parallel
|
||||
|
||||
<div align="center">
|
||||
|
||||
<img src="./assets/performance_comparison_i7-13800h.png" alt="performance_comparison_i7-13800h" width="800"/>
|
||||
|
||||
</div>
|
||||
|
||||
**Test Configuration:**
|
||||
- Model: BitNet-b1.58-2B-4T
|
||||
- Hardware: Cobalt 100
|
||||
- Threads: 1 / 2 / 4 / 8
|
||||
- Test: 128 prompt tokens (pp128) + 128 generated tokens (tg128)
|
||||
- Method: Activation Parallel
|
||||
|
||||
<div align="center">
|
||||
|
||||
<img src="./assets/performance_comparison_cobalt100_dotprod.png" alt="performance_comparison_cobalt100_dotprod" width="800"/>
|
||||
|
||||
</div>
|
||||
|
||||
## Technical Details
|
||||
|
||||
### Key Files Modified
|
||||
|
||||
- `src/ggml-bitnet-mad.cpp`: Parallel kernel implementations
|
||||
- `3rdparty/llama.cpp/ggml/src/ggml.c`: GEMM/GEMV integration
|
||||
- `include/gemm-config.h`: Configuration file
|
||||
|
||||
### Supported Architectures
|
||||
|
||||
- ✅ x86-64 with AVX2
|
||||
- ✅ ARM with NEON
|
||||
- ✅ ARM with DOTPROD extension
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 183 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 341 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 313 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 290 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 260 KiB |
@@ -0,0 +1,167 @@
|
||||
#include <vector>
|
||||
#include <type_traits>
|
||||
|
||||
#include <string.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
#include "ggml-bitnet.h"
|
||||
#include "ggml-quants.h"
|
||||
#include "bitnet-lut-kernels.h"
|
||||
|
||||
#if defined(GGML_BITNET_ARM_TL1)
|
||||
|
||||
void ggml_bitnet_init(void) {
|
||||
// LOG(INFO) << "ggml_bitnet_init";
|
||||
|
||||
if (initialized) {
|
||||
return;
|
||||
}
|
||||
initialized = true;
|
||||
|
||||
// if (wrapper == nullptr) {
|
||||
// wrapper = new BITNET::BITNETGeMMWrapper<bitnet_bitnet_float_type>();
|
||||
// }
|
||||
if (bitnet_tensor_extras == nullptr) {
|
||||
bitnet_tensor_extras = new bitnet_tensor_extra[GGML_BITNET_MAX_NODES];
|
||||
}
|
||||
bitnet_tensor_extras_index = 0;
|
||||
}
|
||||
|
||||
void ggml_bitnet_free(void) {
|
||||
// LOG(INFO) << "ggml_bitnet_free";
|
||||
|
||||
if (!initialized) {
|
||||
return;
|
||||
}
|
||||
initialized = false;
|
||||
|
||||
// delete wrapper;
|
||||
// wrapper = nullptr;
|
||||
for (size_t i = 0; i < bitnet_tensor_extras_index; i++) {
|
||||
// aligned_free(bitnet_tensor_extras[i].qweights);
|
||||
// aligned_free(bitnet_tensor_extras[i].scales);
|
||||
}
|
||||
delete[] bitnet_tensor_extras;
|
||||
bitnet_tensor_extras = nullptr;
|
||||
}
|
||||
|
||||
static bool do_permutate(enum ggml_type type) {
|
||||
if (type == GGML_TYPE_TL1) {
|
||||
// Add additional args to decide if permuted I2 or naive I2
|
||||
return false;
|
||||
} else {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
bool ggml_bitnet_can_mul_mat(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst) {
|
||||
if ((is_type_supported(src0->type)) &&
|
||||
src1->type == GGML_TYPE_F32 &&
|
||||
dst->type == GGML_TYPE_F32 &&
|
||||
src0->backend == GGML_BACKEND_TYPE_CPU) {
|
||||
if (src1->ne[1] <= 1) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t ggml_bitnet_mul_mat_get_wsize(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst) {
|
||||
const size_t ne01 = src0->ne[1];
|
||||
const size_t ne10 = src1->ne[0];
|
||||
const size_t ne11 = src1->ne[1];
|
||||
const int bits = ggml_bitnet_get_type_bits(src0->type);
|
||||
|
||||
size_t wsize = ne10 * ne11 * 15 * sizeof(int8_t) + 1 * ne11 * 2 * sizeof(bitnet_float_type);
|
||||
if (sizeof(bitnet_float_type) == 2) {
|
||||
// Need fp32 to fp16 conversion
|
||||
wsize += std::max(ne10, ne01) * ne11 * sizeof(bitnet_float_type);
|
||||
}
|
||||
wsize = ((wsize - 1) / 64 + 1) * 64;
|
||||
return wsize;
|
||||
}
|
||||
|
||||
int ggml_bitnet_get_type_bits(enum ggml_type type) {
|
||||
switch (type) {
|
||||
case GGML_TYPE_TL1:
|
||||
return 2;
|
||||
case GGML_TYPE_Q4_0:
|
||||
return 4;
|
||||
default:
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
|
||||
#endif
|
||||
#if defined(GGML_BITNET_X86_TL2)
|
||||
void ggml_bitnet_init(void) {
|
||||
// LOG(INFO) << "ggml_bitnet_init";
|
||||
|
||||
if (initialized) {
|
||||
return;
|
||||
}
|
||||
initialized = true;
|
||||
|
||||
// if (wrapper == nullptr) {
|
||||
// wrapper = new BITNET::BITNETGeMMWrapper<bitnet_bitnet_float_type>();
|
||||
// }
|
||||
if (bitnet_tensor_extras == nullptr) {
|
||||
bitnet_tensor_extras = new bitnet_tensor_extra[GGML_BITNET_MAX_NODES];
|
||||
}
|
||||
bitnet_tensor_extras_index = 0;
|
||||
}
|
||||
|
||||
void ggml_bitnet_free(void) {
|
||||
// LOG(INFO) << "ggml_bitnet_free";
|
||||
|
||||
if (!initialized) {
|
||||
return;
|
||||
}
|
||||
initialized = false;
|
||||
|
||||
// delete wrapper;
|
||||
// wrapper = nullptr;
|
||||
for (size_t i = 0; i < bitnet_tensor_extras_index; i++) {
|
||||
// aligned_free(bitnet_tensor_extras[i].qweights);
|
||||
// aligned_free(bitnet_tensor_extras[i].scales);
|
||||
}
|
||||
delete[] bitnet_tensor_extras;
|
||||
bitnet_tensor_extras = nullptr;
|
||||
}
|
||||
|
||||
bool ggml_bitnet_can_mul_mat(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst) {
|
||||
if ((is_type_supported(src0->type)) &&
|
||||
src1->type == GGML_TYPE_F32 &&
|
||||
dst->type == GGML_TYPE_F32 &&
|
||||
src0->backend == GGML_BACKEND_TYPE_CPU) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t ggml_bitnet_mul_mat_get_wsize(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst) {
|
||||
const size_t ne01 = src0->ne[1];
|
||||
const size_t ne10 = src1->ne[0];
|
||||
const size_t ne11 = src1->ne[1];
|
||||
|
||||
size_t wsize = ne10 * ne11 * 11 * sizeof(int8_t) + 2 * ne11 * 2 * sizeof(bitnet_float_type);
|
||||
if (sizeof(bitnet_float_type) == 2) {
|
||||
// Need fp32 to fp16 conversion
|
||||
wsize += std::max(ne10, ne01) * ne11 * sizeof(bitnet_float_type);
|
||||
}
|
||||
wsize = ((wsize - 1) / 64 + 1) * 64;
|
||||
return wsize;
|
||||
}
|
||||
|
||||
int ggml_bitnet_get_type_bits(enum ggml_type type) {
|
||||
switch (type) {
|
||||
case GGML_TYPE_TL2:
|
||||
return 2;
|
||||
case GGML_TYPE_Q4_0:
|
||||
return 4;
|
||||
default:
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,442 @@
|
||||
import argparse
|
||||
import os
|
||||
from configparser import ConfigParser
|
||||
|
||||
def gen_ctor_code():
|
||||
kernel_code = "\n\
|
||||
#include \"ggml-bitnet.h\"\n\
|
||||
#define GGML_BITNET_MAX_NODES 8192\n\
|
||||
static bool initialized = false;\n\
|
||||
static bitnet_tensor_extra * bitnet_tensor_extras = nullptr;\n\
|
||||
static size_t bitnet_tensor_extras_index = 0;\n\
|
||||
static void * aligned_malloc(size_t size) {{\n\
|
||||
#if defined(_WIN32)\n\
|
||||
return _aligned_malloc(size, 64);\n\
|
||||
#else\n\
|
||||
void * ptr = nullptr;\n\
|
||||
posix_memalign(&ptr, 64, size);\n\
|
||||
return ptr;\n\
|
||||
#endif\n\
|
||||
}}\n\
|
||||
static void aligned_free(void * ptr) {{\n\
|
||||
#if defined(_WIN32)\n\
|
||||
_aligned_free(ptr);\n\
|
||||
#else\n\
|
||||
free(ptr);\n\
|
||||
#endif\n\
|
||||
}}\n\
|
||||
\n\
|
||||
void per_tensor_quant(int k, void* lut_scales_, void* b_) {{\n\
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;\n\
|
||||
bitnet_float_type* b = (bitnet_float_type*)b_;\n\
|
||||
#ifdef __ARM_NEON\n\
|
||||
float32x4_t temp_max = vdupq_n_f32(0);\n\
|
||||
for (int i=0; i < k / 4; i++) {{\n\
|
||||
float32x4_t vec_bs = vld1q_f32(b + 4 * i);\n\
|
||||
float32x4_t abssum = vabsq_f32(vec_bs);\n\
|
||||
temp_max = vmaxq_f32(abssum, temp_max);\n\
|
||||
}}\n\
|
||||
float32_t scales = 127 / vmaxvq_f32(temp_max);\n\
|
||||
*lut_scales = scales;\n\
|
||||
#elif defined __AVX2__\n\
|
||||
__m256 max_vec = _mm256_set1_ps(0.f);\n\
|
||||
const __m256 vec_sign = _mm256_set1_ps(-0.0f);\n\
|
||||
// #pragma unroll\n\
|
||||
for (int i = 0; i < k / 8; i++) {{\n\
|
||||
__m256 vec_b = _mm256_loadu_ps(b + i * 8);\n\
|
||||
__m256 vec_babs = _mm256_andnot_ps(vec_sign, vec_b);\n\
|
||||
max_vec = _mm256_max_ps(vec_babs, max_vec);\n\
|
||||
}}\n\
|
||||
__m128 max1 = _mm_max_ps(_mm256_extractf128_ps(max_vec, 1), _mm256_castps256_ps128(max_vec));\n\
|
||||
max1 = _mm_max_ps(max1, _mm_movehl_ps(max1, max1));\n\
|
||||
max1 = _mm_max_ss(max1, _mm_movehdup_ps(max1));\n\
|
||||
float scales = 127 / _mm_cvtss_f32(max1);\n\
|
||||
*lut_scales = scales;\n\
|
||||
#endif\n\
|
||||
}}\n\
|
||||
\n\
|
||||
void partial_max_reset(void* lut_scales_) {{\n\
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;\n\
|
||||
*lut_scales = 0.0;\n\
|
||||
}}\n\
|
||||
\n\
|
||||
#ifdef __ARM_NEON\n\
|
||||
inline void Transpose_8_8(\n\
|
||||
int16x8_t *v0,\n\
|
||||
int16x8_t *v1,\n\
|
||||
int16x8_t *v2,\n\
|
||||
int16x8_t *v3,\n\
|
||||
int16x8_t *v4,\n\
|
||||
int16x8_t *v5,\n\
|
||||
int16x8_t *v6,\n\
|
||||
int16x8_t *v7)\n\
|
||||
{{\n\
|
||||
int16x8x2_t q04 = vzipq_s16(*v0, *v4);\n\
|
||||
int16x8x2_t q15 = vzipq_s16(*v1, *v5);\n\
|
||||
int16x8x2_t q26 = vzipq_s16(*v2, *v6);\n\
|
||||
int16x8x2_t q37 = vzipq_s16(*v3, *v7);\n\
|
||||
\n\
|
||||
int16x8x2_t q0246_0 = vzipq_s16(q04.val[0], q26.val[0]);\n\
|
||||
int16x8x2_t q0246_1 = vzipq_s16(q04.val[1], q26.val[1]);\n\
|
||||
int16x8x2_t q1357_0 = vzipq_s16(q15.val[0], q37.val[0]);\n\
|
||||
int16x8x2_t q1357_1 = vzipq_s16(q15.val[1], q37.val[1]);\n\
|
||||
\n\
|
||||
int16x8x2_t q_fin_0 = vzipq_s16(q0246_0.val[0], q1357_0.val[0]);\n\
|
||||
int16x8x2_t q_fin_1 = vzipq_s16(q0246_0.val[1], q1357_0.val[1]);\n\
|
||||
int16x8x2_t q_fin_2 = vzipq_s16(q0246_1.val[0], q1357_1.val[0]);\n\
|
||||
int16x8x2_t q_fin_3 = vzipq_s16(q0246_1.val[1], q1357_1.val[1]);\n\
|
||||
\n\
|
||||
*v0 = q_fin_0.val[0];\n\
|
||||
*v1 = q_fin_0.val[1];\n\
|
||||
*v2 = q_fin_1.val[0];\n\
|
||||
*v3 = q_fin_1.val[1];\n\
|
||||
*v4 = q_fin_2.val[0];\n\
|
||||
*v5 = q_fin_2.val[1];\n\
|
||||
*v6 = q_fin_3.val[0];\n\
|
||||
*v7 = q_fin_3.val[1];\n\
|
||||
}}\n\
|
||||
#endif\n\
|
||||
\n\
|
||||
template<int act_k>\n\
|
||||
inline void lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_type* lut_scales) {{\n\
|
||||
#ifdef __ARM_NEON\n\
|
||||
int16x8_t vec_lut[16];\n\
|
||||
float32_t scales = *lut_scales;\n\
|
||||
uint8_t tbl_mask[16];\n\
|
||||
tbl_mask[0] = 0;\n\
|
||||
tbl_mask[1] = 2;\n\
|
||||
tbl_mask[2] = 4;\n\
|
||||
tbl_mask[3] = 6;\n\
|
||||
tbl_mask[4] = 8;\n\
|
||||
tbl_mask[5] = 10;\n\
|
||||
tbl_mask[6] = 12;\n\
|
||||
tbl_mask[7] = 14;\n\
|
||||
tbl_mask[8] = 1;\n\
|
||||
tbl_mask[9] = 3;\n\
|
||||
tbl_mask[10] = 5;\n\
|
||||
tbl_mask[11] = 7;\n\
|
||||
tbl_mask[12] = 9;\n\
|
||||
tbl_mask[13] = 11;\n\
|
||||
tbl_mask[14] = 13;\n\
|
||||
tbl_mask[15] = 15;\n\
|
||||
uint8x16_t tbl_mask_q = vld1q_u8(tbl_mask);\n\
|
||||
#pragma unroll\n\
|
||||
for (int k = 0; k < act_k / 16; ++k) {{\n\
|
||||
float32x4x2_t vec_bs_x0 = vld2q_f32(b + k * 16);\n\
|
||||
float32x4x2_t vec_bs_x1 = vld2q_f32(b + k * 16 + 8);\n\
|
||||
float32x4_t vec_f_0 = vmulq_n_f32(vec_bs_x0.val[0], scales);\n\
|
||||
float32x4_t vec_f_1 = vmulq_n_f32(vec_bs_x0.val[1], scales);\n\
|
||||
float32x4_t vec_f_2 = vmulq_n_f32(vec_bs_x1.val[0], scales);\n\
|
||||
float32x4_t vec_f_3 = vmulq_n_f32(vec_bs_x1.val[1], scales);\n\
|
||||
int32x4_t vec_b_0 = vcvtnq_s32_f32(vec_f_0);\n\
|
||||
int32x4_t vec_b_1 = vcvtnq_s32_f32(vec_f_1);\n\
|
||||
int32x4_t vec_b_2 = vcvtnq_s32_f32(vec_f_2);\n\
|
||||
int32x4_t vec_b_3 = vcvtnq_s32_f32(vec_f_3);\n\
|
||||
int16x4_t vec_b16_0 = vmovn_s32(vec_b_0);\n\
|
||||
int16x4_t vec_b16_1 = vmovn_s32(vec_b_1);\n\
|
||||
int16x4_t vec_b16_2 = vmovn_s32(vec_b_2);\n\
|
||||
int16x4_t vec_b16_3 = vmovn_s32(vec_b_3);\n\
|
||||
int16x8_t vec_bs_0 = vcombine_s16(vec_b16_0, vec_b16_2);\n\
|
||||
int16x8_t vec_bs_1 = vcombine_s16(vec_b16_1, vec_b16_3);\n\
|
||||
vec_lut[0] = vdupq_n_s16(0);\n\
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_0;\n\
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_1;\n\
|
||||
vec_lut[1] = vdupq_n_s16(0);\n\
|
||||
vec_lut[1] = vec_lut[1] - vec_bs_0;\n\
|
||||
vec_lut[2] = vdupq_n_s16(0);\n\
|
||||
vec_lut[2] = vec_lut[2] - vec_bs_0;\n\
|
||||
vec_lut[2] = vec_lut[2] + vec_bs_1;\n\
|
||||
vec_lut[3] = vdupq_n_s16(0);\n\
|
||||
vec_lut[3] = vec_lut[3] - vec_bs_1;\n\
|
||||
vec_lut[4] = vdupq_n_s16(0);\n\
|
||||
vec_lut[5] = vec_bs_1;\n\
|
||||
vec_lut[6] = vec_bs_0;\n\
|
||||
vec_lut[6] = vec_lut[6] - vec_bs_1;\n\
|
||||
vec_lut[7] = vec_bs_0;\n\
|
||||
vec_lut[8] = vec_bs_0;\n\
|
||||
vec_lut[8] = vec_lut[8] + vec_bs_1;\n\
|
||||
Transpose_8_8(&(vec_lut[0]), &(vec_lut[1]), &(vec_lut[2]), &(vec_lut[3]),\n\
|
||||
&(vec_lut[4]), &(vec_lut[5]), &(vec_lut[6]), &(vec_lut[7]));\n\
|
||||
Transpose_8_8(&(vec_lut[8]), &(vec_lut[9]), &(vec_lut[10]), &(vec_lut[11]),\n\
|
||||
&(vec_lut[12]), &(vec_lut[13]), &(vec_lut[14]), &(vec_lut[15]));\n\
|
||||
#pragma unroll\n\
|
||||
for (int idx = 0; idx < 8; idx++) {{\n\
|
||||
int8x16_t q0_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx]), tbl_mask_q);\n\
|
||||
int8x8_t q0_low = vget_low_s8(q0_s);\n\
|
||||
int8x8_t q0_high = vget_high_s8(q0_s);\n\
|
||||
int8x16_t q1_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx + 8]), tbl_mask_q);\n\
|
||||
int8x8_t q1_low = vget_low_s8(q1_s);\n\
|
||||
int8x8_t q1_high = vget_high_s8(q1_s);\n\
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2, q0_high);\n\
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 8, q1_high);\n\
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 16, q0_low);\n\
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 24, q1_low);\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
#endif\n\
|
||||
}}\n\
|
||||
\n\
|
||||
static bool is_type_supported(enum ggml_type type) {{\n\
|
||||
if (type == GGML_TYPE_Q4_0 ||\n\
|
||||
type == GGML_TYPE_TL1) {{\n\
|
||||
return true;\n\
|
||||
}} else {{\n\
|
||||
return false;\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
"
|
||||
return kernel_code
|
||||
|
||||
def gen_body_core_code(bm, by):
|
||||
length = 4
|
||||
all_code = ""
|
||||
for i in range(length):
|
||||
core_code = "\n\
|
||||
uint8x16_t vec_a_{0} = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + {0} * 16);\n\
|
||||
uint8x16_t vec_a{0}_top = vshrq_n_u8(vec_a_{0}, 4);\n\
|
||||
uint8x16_t vec_a{0}_bot = vandq_u8(vec_a_{0}, vec_mask);\n\
|
||||
int8x16_t vec_v_{0}_left_tmp0 = vqtbl1q_s8(vec_lut[{1} * k + {2}], vec_a{0}_top);\n\
|
||||
int8x16_t vec_v_{0}_left_tmp1 = vqtbl1q_s8(vec_lut[{1} * k + {3}], vec_a{0}_top);\n\
|
||||
int8x16_t vec_v_{0}_right_tmp0 = vqtbl1q_s8(vec_lut[{1} * k + {4}], vec_a{0}_bot);\n\
|
||||
int8x16_t vec_v_{0}_right_tmp1 = vqtbl1q_s8(vec_lut[{1} * k + {5}], vec_a{0}_bot);\n\
|
||||
int8x16x2_t vec_v_left_{0} = vzipq_s8(vec_v_{0}_left_tmp1, vec_v_{0}_left_tmp0);\n\
|
||||
int8x16x2_t vec_v_right_{0} = vzipq_s8(vec_v_{0}_right_tmp1, vec_v_{0}_right_tmp0);\n\
|
||||
vec_c[{6}] += vec_v_left_{0}.val[0];\n\
|
||||
vec_c[{6}] += vec_v_right_{0}.val[0];\n\
|
||||
vec_c[{7}] += vec_v_left_{0}.val[1];\n\
|
||||
vec_c[{7}] += vec_v_right_{0}.val[1];\n\
|
||||
".format(i, 2 * by // 2, (4 * i) % (2 * by // 2), (4 * i + 1) % (2 * by // 2), (4 * i + 2) % (2 * by // 2), (4 * i + 3) % (2 * by // 2), (i * 2) // (by // 2) * 2 + 0, (i * 2) // (by // 2) * 2 + 1)
|
||||
|
||||
all_code = "".join([all_code, core_code])
|
||||
|
||||
all_code = "".join([all_code, "\n }\n\n"])
|
||||
|
||||
for i in range(bm // 8):
|
||||
core_code = "\
|
||||
int32x4_t vec_v_bot_low_low_{0} = vmovl_s16(vget_low_s16(vec_c[{0}]));\n\
|
||||
int32x4_t vec_v_bot_low_high_{0} = vmovl_high_s16(vec_c[{0}]);\n\
|
||||
vst1q_s32(c + i + {1}, vld1q_s32(c + i + {1}) + vec_v_bot_low_low_{0});\n\
|
||||
vst1q_s32(c + i + {2}, vld1q_s32(c + i + {2}) + vec_v_bot_low_high_{0});\n".format(i, i * 8, i * 8 + 4)
|
||||
all_code = "".join([all_code, core_code])
|
||||
|
||||
return all_code
|
||||
|
||||
def gen_tbl_impl(pre, BM, BK, bm, k):
|
||||
|
||||
kernel_code = "\
|
||||
#include <arm_neon.h>\n\
|
||||
\n\
|
||||
#define BM{0} {1}\n\
|
||||
#define BBK{0} {2}\n\
|
||||
inline void tbl_impl_{0}(int32_t* c, int8_t* lut, uint8_t* a) {{\n\
|
||||
#ifdef __ARM_NEON\n\
|
||||
const int KK = BBK{0} / 2;\n\
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);\n\
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);\n\
|
||||
int8x16_t vec_lut[2 * KK];\n\
|
||||
".format(pre, BM, BK)
|
||||
|
||||
kernel_code = "".join([kernel_code, " int16x8_t vec_c[{}];".format(bm // 8)])
|
||||
|
||||
kernel_code = "".join([kernel_code, "\n\
|
||||
#pragma unroll\n\
|
||||
for (int k = 0; k < 2 * KK; k++) {\n\
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);\n\
|
||||
}\n"])
|
||||
|
||||
pre_core_code = "\n\
|
||||
#pragma unroll\n\
|
||||
for (int i = 0; i < BM{}; i += {}) {{\n\
|
||||
#pragma unroll\n\
|
||||
for (int i=0; i<{}; i++) {{\n\
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);\n\
|
||||
}}\n".format(pre, bm, bm // 8)
|
||||
|
||||
body_core_pre_code = "\n\
|
||||
#pragma unroll\n\
|
||||
for (int k = 0; k < KK / {}; k++) {{\n\
|
||||
".format(256 // bm // 2)
|
||||
|
||||
body_core_post_code = "\n\
|
||||
}\n\
|
||||
\
|
||||
#endif\n\
|
||||
}\n"
|
||||
|
||||
kernel_code = "".join([kernel_code, pre_core_code, body_core_pre_code, gen_body_core_code(bm, 256 // bm), body_core_post_code])
|
||||
|
||||
kernel_code = "".join([kernel_code, "\n\
|
||||
int32_t qgemm_lut_{0}(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {{\n\
|
||||
alignas({1}) uint32_t CBits[BM{0}];\n\
|
||||
memset(&(CBits[0]), 0, BM{0} * sizeof(int32_t));\n\
|
||||
#pragma unroll\n\
|
||||
for (int32_t k_outer = 0; k_outer < {2} / BBK{0}; ++k_outer) {{\n\
|
||||
tbl_impl_{0}((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK{0} / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK{0} / 2 / 2 * BM{0})])));\n\
|
||||
}}\n\
|
||||
#pragma unroll\n\
|
||||
for (int i = 0; i < BM{0}; i++) {{\n\
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];\n\
|
||||
}}\n\
|
||||
return 0;\n\
|
||||
}};\n".format(pre, min(32, BK), k)])
|
||||
|
||||
return kernel_code
|
||||
|
||||
def gen_top_api(kernel_shapes):
|
||||
|
||||
kernel_code = "void ggml_preprocessor(int m, int k, void* B, void* LUT_Scales, void* QLUT) {{\n\
|
||||
if (m == {0} && k == {1}) {{\n\
|
||||
preprocessor_k<{1}>(B, LUT_Scales, QLUT);\n\
|
||||
}}\n\
|
||||
".format(kernel_shapes[0][0], kernel_shapes[0][1])
|
||||
for i in range(1, len(kernel_shapes)):
|
||||
kernel_code = "".join([kernel_code, " else if (m == {0} && k == {1}) {{\n\
|
||||
preprocessor_k<{1}>(B, LUT_Scales, QLUT);\n\
|
||||
}}\n".format(kernel_shapes[i][0], kernel_shapes[i][1])])
|
||||
kernel_code = "".join([kernel_code, "}\n"])
|
||||
kernel_code = "".join([kernel_code, "void ggml_qgemm_lut(int m, int k, void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {{\n\
|
||||
if (m == {0} && k == {1}) {{\n\
|
||||
qgemm_lut_{0}_{1}(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}\n\
|
||||
".format(kernel_shapes[0][0], kernel_shapes[0][1])])
|
||||
for i in range(1, len(kernel_shapes)):
|
||||
kernel_code = "".join([kernel_code, " else if (m == {0} && k == {1}) {{\n\
|
||||
qgemm_lut_{0}_{1}(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}\n\
|
||||
".format(kernel_shapes[i][0], kernel_shapes[i][1])])
|
||||
kernel_code = "".join([kernel_code, "}\n"])
|
||||
return kernel_code
|
||||
|
||||
def gen_preprocess_code():
|
||||
kernel_code = "\n\
|
||||
template<int K>\n\
|
||||
void preprocessor_k(void* B, void* LUT_Scales, void* QLUT) {{\n\
|
||||
partial_max_reset((&(((bitnet_float_type*)LUT_Scales)[0])));\n\
|
||||
per_tensor_quant(K, (&(((bitnet_float_type*)LUT_Scales)[0])), (&(((bitnet_float_type*)B)[0])));\n\
|
||||
\n\
|
||||
lut_ctor<K>((&(((int8_t*)QLUT)[0])), (&(((bitnet_float_type*)B)[0])), (&(((bitnet_float_type*)LUT_Scales)[0])));\n\
|
||||
}}\n"
|
||||
return kernel_code
|
||||
|
||||
def gen_transform_code(kernel_shape):
|
||||
kernel_code = "\n\
|
||||
void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor) {\n\
|
||||
if (!(is_type_supported(tensor->type) && tensor->backend == GGML_BACKEND_TYPE_CPU && tensor->extra == nullptr)) {\n\
|
||||
return;\n\
|
||||
}\n\
|
||||
\n\
|
||||
int k = tensor->ne[0];\n\
|
||||
int m = tensor->ne[1];\n\
|
||||
const int lut_scales_size = 1;\n\
|
||||
const int scales_size = 1;\n\
|
||||
int bk = 0;\n\
|
||||
int bm = 0;\n"
|
||||
|
||||
kernel_code = "".join([kernel_code, "\n\
|
||||
if (m == {0} && k == {1}) {{\n\
|
||||
bm = BM{0}_{1};\n\
|
||||
bk = BBK{0}_{1};\n\
|
||||
}}\n".format(kernel_shapes[0][0], kernel_shapes[0][1])])
|
||||
|
||||
for i in range(1, len(kernel_shapes)):
|
||||
kernel_code = "".join([kernel_code, "else if (m == {0} && k == {1}) {{\n\
|
||||
bm = BM{0}_{1};\n\
|
||||
bk = BBK{0}_{1};\n\
|
||||
}}\n".format(kernel_shapes[i][0], kernel_shapes[i][1])])
|
||||
|
||||
kernel_code = "".join([kernel_code, "\n\
|
||||
const int n_tile_num = m / bm;\n\
|
||||
const int BK = bk;\n\
|
||||
uint8_t * qweights;\n\
|
||||
bitnet_float_type * scales;\n\
|
||||
\n\
|
||||
scales = (bitnet_float_type *) aligned_malloc(sizeof(bitnet_float_type));\n\
|
||||
qweights = (uint8_t *) tensor->data;\n\
|
||||
float * i2_scales = (float * )(qweights + k * m / 4);\n\
|
||||
scales[0] = (bitnet_float_type) i2_scales[0];\n\
|
||||
\n\
|
||||
tensor->extra = bitnet_tensor_extras + bitnet_tensor_extras_index;\n\
|
||||
bitnet_tensor_extras[bitnet_tensor_extras_index++] = {\n\
|
||||
/* .lut_scales_size = */ lut_scales_size,\n\
|
||||
/* .BK = */ BK,\n\
|
||||
/* .n_tile_num = */ n_tile_num,\n\
|
||||
/* .qweights = */ qweights,\n\
|
||||
/* .scales = */ scales\n\
|
||||
};\n\
|
||||
}\n"])
|
||||
|
||||
return kernel_code
|
||||
|
||||
if __name__ == "__main__":
|
||||
ModelShapeDict = {
|
||||
"bitnet_b1_58-large" : [[1536, 4096],
|
||||
[1536, 1536],
|
||||
[4096, 1536]],
|
||||
"bitnet_b1_58-3B" : [[3200, 8640],
|
||||
[3200, 3200],
|
||||
[8640, 3200]],
|
||||
"Llama3-8B-1.58-100B-tokens" : [[14336, 4096],
|
||||
[4096, 14336],
|
||||
[1024, 4096],
|
||||
[4096, 4096]]
|
||||
}
|
||||
|
||||
parser = argparse.ArgumentParser(description='gen impl')
|
||||
parser.add_argument('--model',default="input", type=str, dest="model",
|
||||
help="choose from bitnet_b1_58-large/bitnet_b1_58-3B/Llama3-8B-1.58-100B-tokens.")
|
||||
parser.add_argument('--BM',default="input", type=str,
|
||||
help="block length when cutting one weight (M, K) into M / BM weights (BM, K).")
|
||||
parser.add_argument('--BK',default="input", type=str,
|
||||
help="block length when cutting one weight (M, K) into K / BK weights (M, BK).")
|
||||
parser.add_argument('--bm',default="input", type=str,
|
||||
help="using simd instructions to compute (bm, 256 / bm) in one block")
|
||||
args = parser.parse_args()
|
||||
|
||||
kernel_shapes = ModelShapeDict[args.model]
|
||||
|
||||
BM_list = [int(item) for item in args.BM.split(',')]
|
||||
BK_list = [int(item) for item in args.BK.split(',')]
|
||||
bm_list = [int(item) for item in args.bm.split(',')]
|
||||
|
||||
assert(len(BM_list) == len(BK_list) == len(bm_list) == len(kernel_shapes)), "number of BM / BK / bm shoud be {}".format(len(kernel_shapes))
|
||||
|
||||
for i in range(len(kernel_shapes)):
|
||||
assert kernel_shapes[i][0] % BM_list[i] == 0, "M %% BM should be 0"
|
||||
assert kernel_shapes[i][1] % BK_list[i] == 0, "K %% BK should be 0"
|
||||
assert bm_list[i] in [32, 64], "choose bm from [32, 64]"
|
||||
|
||||
tbl_impl_code = []
|
||||
|
||||
for i in range(len(kernel_shapes)):
|
||||
tbl_impl_code.append(
|
||||
gen_tbl_impl("{}_{}".format(kernel_shapes[i][0], kernel_shapes[i][1]), BM_list[i], BK_list[i], bm_list[i], kernel_shapes[i][1])
|
||||
)
|
||||
api_code = gen_top_api(kernel_shapes)
|
||||
pre_code = gen_preprocess_code()
|
||||
ctor_code = gen_ctor_code()
|
||||
trans_code = gen_transform_code(kernel_shapes)
|
||||
|
||||
output_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "include")
|
||||
|
||||
with open(''.join([output_dir, "/bitnet-lut-kernels.h"]), 'w') as f:
|
||||
f.write(''.join("#if defined(GGML_BITNET_ARM_TL1)"))
|
||||
f.write(''.join(ctor_code))
|
||||
for code in tbl_impl_code:
|
||||
f.write(''.join(code))
|
||||
f.write(''.join(pre_code))
|
||||
f.write(''.join(api_code))
|
||||
f.write(''.join(trans_code))
|
||||
f.write(''.join("#endif"))
|
||||
|
||||
config = ConfigParser()
|
||||
|
||||
for i in range(len(kernel_shapes)):
|
||||
config.add_section('Kernels_{}'.format(i))
|
||||
config.set('Kernels_{}'.format(i), 'M'.format(i), str(kernel_shapes[i][0]))
|
||||
config.set('Kernels_{}'.format(i), 'K'.format(i), str(kernel_shapes[i][1]))
|
||||
config.set('Kernels_{}'.format(i), 'BM'.format(i), str(BM_list[i]))
|
||||
config.set('Kernels_{}'.format(i), 'BK'.format(i), str(BK_list[i]))
|
||||
config.set('Kernels_{}'.format(i), 'bmm'.format(i), str(bm_list[i]))
|
||||
|
||||
with open(''.join([output_dir, "/kernel_config.ini"]), 'w') as configfile:
|
||||
config.write(configfile)
|
||||
@@ -0,0 +1,757 @@
|
||||
import argparse
|
||||
import os
|
||||
from configparser import ConfigParser
|
||||
|
||||
def gen_ctor_code():
|
||||
kernel_code = "\n\
|
||||
#include \"ggml-bitnet.h\"\n\
|
||||
#include <cstring>\n\
|
||||
#include <immintrin.h>\n\
|
||||
#define GGML_BITNET_MAX_NODES 8192\n\
|
||||
static bool initialized = false;\n\
|
||||
static bitnet_tensor_extra * bitnet_tensor_extras = nullptr;\n\
|
||||
static size_t bitnet_tensor_extras_index = 0;\n\
|
||||
static void * aligned_malloc(size_t size) {\n\
|
||||
#if defined(_WIN32)\n\
|
||||
return _aligned_malloc(size, 64);\n\
|
||||
#else\n\
|
||||
void * ptr = nullptr;\n\
|
||||
posix_memalign(&ptr, 64, size);\n\
|
||||
return ptr;\n\
|
||||
#endif\n\
|
||||
}\n\
|
||||
\n\
|
||||
static void aligned_free(void * ptr) {\n\
|
||||
#if defined(_WIN32)\n\
|
||||
_aligned_free(ptr);\n\
|
||||
#else\n\
|
||||
free(ptr);\n\
|
||||
#endif\n\
|
||||
}\n\
|
||||
#define BK2 32\n\
|
||||
#if defined __AVX2__\n\
|
||||
inline void _mm256_merge_epi32(const __m256i v0, const __m256i v1, __m256i *vl, __m256i *vh)\n\
|
||||
{\n\
|
||||
__m256i va = _mm256_permute4x64_epi64(v0, _MM_SHUFFLE(3, 1, 2, 0));\n\
|
||||
__m256i vb = _mm256_permute4x64_epi64(v1, _MM_SHUFFLE(3, 1, 2, 0));\n\
|
||||
*vl = _mm256_unpacklo_epi32(va, vb);\n\
|
||||
*vh = _mm256_unpackhi_epi32(va, vb);\n\
|
||||
}\n\
|
||||
inline void _mm256_merge_epi64(const __m256i v0, const __m256i v1, __m256i *vl, __m256i *vh)\n\
|
||||
{\n\
|
||||
__m256i va = _mm256_permute4x64_epi64(v0, _MM_SHUFFLE(3, 1, 2, 0));\n\
|
||||
__m256i vb = _mm256_permute4x64_epi64(v1, _MM_SHUFFLE(3, 1, 2, 0));\n\
|
||||
*vl = _mm256_unpacklo_epi64(va, vb);\n\
|
||||
*vh = _mm256_unpackhi_epi64(va, vb);\n\
|
||||
}\n\
|
||||
inline void _mm256_merge_si128(const __m256i v0, const __m256i v1, __m256i *vl, __m256i *vh)\n\
|
||||
{\n\
|
||||
*vl = _mm256_permute2x128_si256(v0, v1, _MM_SHUFFLE(0, 2, 0, 0));\n\
|
||||
*vh = _mm256_permute2x128_si256(v0, v1, _MM_SHUFFLE(0, 3, 0, 1));\n\
|
||||
}\n\
|
||||
inline void Transpose_8_8(\n\
|
||||
__m256i *v0,\n\
|
||||
__m256i *v1,\n\
|
||||
__m256i *v2,\n\
|
||||
__m256i *v3,\n\
|
||||
__m256i *v4,\n\
|
||||
__m256i *v5,\n\
|
||||
__m256i *v6,\n\
|
||||
__m256i *v7)\n\
|
||||
{\n\
|
||||
__m256i w0, w1, w2, w3, w4, w5, w6, w7;\n\
|
||||
__m256i x0, x1, x2, x3, x4, x5, x6, x7;\n\
|
||||
_mm256_merge_epi32(*v0, *v1, &w0, &w1);\n\
|
||||
_mm256_merge_epi32(*v2, *v3, &w2, &w3);\n\
|
||||
_mm256_merge_epi32(*v4, *v5, &w4, &w5);\n\
|
||||
_mm256_merge_epi32(*v6, *v7, &w6, &w7);\n\
|
||||
_mm256_merge_epi64(w0, w2, &x0, &x1);\n\
|
||||
_mm256_merge_epi64(w1, w3, &x2, &x3);\n\
|
||||
_mm256_merge_epi64(w4, w6, &x4, &x5);\n\
|
||||
_mm256_merge_epi64(w5, w7, &x6, &x7);\n\
|
||||
_mm256_merge_si128(x0, x4, v0, v1);\n\
|
||||
_mm256_merge_si128(x1, x5, v2, v3);\n\
|
||||
_mm256_merge_si128(x2, x6, v4, v5);\n\
|
||||
_mm256_merge_si128(x3, x7, v6, v7);\n\
|
||||
}\n\
|
||||
#endif\n\
|
||||
inline int32_t per_tensor_quant(int k, void* lut_scales_, void* b_) {\n\
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;\n\
|
||||
bitnet_float_type* b = (bitnet_float_type*)b_;\n\
|
||||
#if defined __AVX2__\n\
|
||||
__m256 max_vec = _mm256_set1_ps(0.f);\n\
|
||||
const __m256 vec_sign = _mm256_set1_ps(-0.0f);\n\
|
||||
for (int i = 0; i < k / 8; i++) {\n\
|
||||
__m256 vec_b = _mm256_loadu_ps(b + i * 8);\n\
|
||||
__m256 vec_babs = _mm256_andnot_ps(vec_sign, vec_b);\n\
|
||||
max_vec = _mm256_max_ps(vec_babs, max_vec);\n\
|
||||
}\n\
|
||||
__m128 max1 = _mm_max_ps(_mm256_extractf128_ps(max_vec, 1), _mm256_castps256_ps128(max_vec));\n\
|
||||
max1 = _mm_max_ps(max1, _mm_movehl_ps(max1, max1));\n\
|
||||
max1 = _mm_max_ss(max1, _mm_movehdup_ps(max1));\n\
|
||||
float scales = 127 / _mm_cvtss_f32(max1);\n\
|
||||
*lut_scales = scales;\n\
|
||||
#endif\n\
|
||||
return 0;\n\
|
||||
}\n\
|
||||
inline int32_t partial_max_reset(int32_t bs, void* lut_scales_) {\n\
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;\n\
|
||||
#pragma unroll\n\
|
||||
for (int i=0; i< bs; i++) {\n\
|
||||
lut_scales[i] = 0.0;\n\
|
||||
}\n\
|
||||
return 0;\n\
|
||||
}\n\
|
||||
template<int act_k>\n\
|
||||
inline int32_t three_lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_type* lut_scales) {\n\
|
||||
#if defined __AVX2__\n\
|
||||
__m256i vec_lut[16];\n\
|
||||
const __m256i vec_bi = _mm256_set_epi32(84, 72, 60, 48, 36, 24, 12, 0);\n\
|
||||
float scales = *lut_scales;\n\
|
||||
__m256i shuffle_mask = _mm256_set_epi8(\n\
|
||||
0x0f, 0x0d, 0x0b, 0x09, 0x07, 0x05, 0x03, 0x01,\n\
|
||||
0x0e, 0x0c, 0x0a, 0x08, 0x06, 0x04, 0x02, 0x00,\n\
|
||||
0x0f, 0x0d, 0x0b, 0x09, 0x07, 0x05, 0x03, 0x01,\n\
|
||||
0x0e, 0x0c, 0x0a, 0x08, 0x06, 0x04, 0x02, 0x00\n\
|
||||
);\n\
|
||||
#pragma unroll\n\
|
||||
for (int k = 0; k < act_k / 24; ++k) {\n\
|
||||
__m256 vec_b0 = _mm256_i32gather_ps(b + k * 24 + 0, vec_bi, 1);\n\
|
||||
__m256 vec_b1 = _mm256_i32gather_ps(b + k * 24 + 1, vec_bi, 1);\n\
|
||||
__m256 vec_b2 = _mm256_i32gather_ps(b + k * 24 + 2, vec_bi, 1);\n\
|
||||
\n\
|
||||
__m256i vec_b0i = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(vec_b0, _mm256_set1_ps(scales)), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));\n\
|
||||
__m256i vec_b1i = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(vec_b1, _mm256_set1_ps(scales)), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));\n\
|
||||
__m256i vec_b2i = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(vec_b2, _mm256_set1_ps(scales)), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));\n\
|
||||
\n\
|
||||
vec_lut[15] = _mm256_setzero_si256();\n\
|
||||
vec_lut[14] = _mm256_setzero_si256();\n\
|
||||
vec_lut[13] = vec_b0i;\n\
|
||||
vec_lut[13] = _mm256_add_epi32(vec_lut[13], vec_b1i);\n\
|
||||
vec_lut[13] = _mm256_add_epi32(vec_lut[13], vec_b2i);\n\
|
||||
vec_lut[12] = vec_b0i;\n\
|
||||
vec_lut[12] = _mm256_add_epi32(vec_lut[12], vec_b1i);\n\
|
||||
vec_lut[11] = vec_b0i;\n\
|
||||
vec_lut[11] = _mm256_add_epi32(vec_lut[11], vec_b1i);\n\
|
||||
vec_lut[11] = _mm256_sub_epi32(vec_lut[11], vec_b2i);\n\
|
||||
vec_lut[10] = vec_b0i;\n\
|
||||
vec_lut[10] = _mm256_add_epi32(vec_lut[10], vec_b2i);\n\
|
||||
vec_lut[9] = vec_b0i;\n\
|
||||
vec_lut[8] = vec_b0i;\n\
|
||||
vec_lut[8] = _mm256_sub_epi32(vec_lut[8], vec_b2i);\n\
|
||||
vec_lut[7] = vec_b0i;\n\
|
||||
vec_lut[7] = _mm256_sub_epi32(vec_lut[7], vec_b1i);\n\
|
||||
vec_lut[7] = _mm256_add_epi32(vec_lut[7], vec_b2i);\n\
|
||||
vec_lut[6] = vec_b0i;\n\
|
||||
vec_lut[6] = _mm256_sub_epi32(vec_lut[6], vec_b1i);\n\
|
||||
vec_lut[5] = vec_b0i;\n\
|
||||
vec_lut[5] = _mm256_sub_epi32(vec_lut[5], vec_b1i);\n\
|
||||
vec_lut[5] = _mm256_sub_epi32(vec_lut[5], vec_b2i);\n\
|
||||
vec_lut[4] = vec_b1i;\n\
|
||||
vec_lut[4] = _mm256_add_epi32(vec_lut[4], vec_b2i);\n\
|
||||
vec_lut[3] = vec_b1i;\n\
|
||||
vec_lut[2] = vec_b1i;\n\
|
||||
vec_lut[2] = _mm256_sub_epi32(vec_lut[2], vec_b2i);\n\
|
||||
vec_lut[1] = vec_b2i;\n\
|
||||
vec_lut[0] = _mm256_setzero_si256();\n\
|
||||
__m256i ix[16];\n\
|
||||
\n\
|
||||
#pragma unroll\n\
|
||||
for (int g = 0; g < 16; ++g) {\n\
|
||||
ix[g] = vec_lut[g];\n\
|
||||
}\n\
|
||||
\n\
|
||||
Transpose_8_8(&(ix[0]), &(ix[1]), &(ix[2]), &(ix[3]), &(ix[4]), &(ix[5]),&(ix[6]), &(ix[7]));\n\
|
||||
Transpose_8_8(&(ix[8]), &(ix[9]), &(ix[10]), &(ix[11]), &(ix[12]), &(ix[13]),&(ix[14]), &(ix[15]));\n\
|
||||
\n\
|
||||
#pragma unroll\n\
|
||||
for (int g = 0; g < 8; ++g) {\n\
|
||||
ix[g] = _mm256_packs_epi32(ix[g], ix[g + 8]);\n\
|
||||
ix[g] = _mm256_permute4x64_epi64(ix[g], _MM_SHUFFLE(3, 1, 2, 0));\n\
|
||||
ix[g] = _mm256_shuffle_epi8(ix[g], shuffle_mask);\n\
|
||||
ix[g] = _mm256_permute4x64_epi64(ix[g], _MM_SHUFFLE(3, 1, 2, 0));\n\
|
||||
}\n\
|
||||
int8_t* qlut_i8 = reinterpret_cast<int8_t*>(qlut);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 0 * 32 + 0), ix[0]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 1 * 32 + 0), ix[1]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 2 * 32 + 0), ix[2]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 3 * 32 + 0), ix[3]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 4 * 32 + 0), ix[4]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 5 * 32 + 0), ix[5]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 6 * 32 + 0), ix[6]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 7 * 32 + 0), ix[7]);\n\
|
||||
\n\
|
||||
}\n\
|
||||
\n\
|
||||
*lut_scales = scales;\n\
|
||||
#endif\n\
|
||||
return 0;\n\
|
||||
}\n\
|
||||
\n\
|
||||
template<int act_k>\n\
|
||||
inline int32_t two_lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_type* lut_scales) {\n\
|
||||
#if defined __AVX2__\n\
|
||||
__m256i vec_lut[16];\n\
|
||||
const __m256i vec_bi = _mm256_set_epi32(56, 48, 40, 32, 24, 16, 8, 0);\n\
|
||||
float scales = *lut_scales;\n\
|
||||
__m256i shuffle_mask = _mm256_set_epi8(\n\
|
||||
0x0f, 0x0d, 0x0b, 0x09, 0x07, 0x05, 0x03, 0x01,\n\
|
||||
0x0e, 0x0c, 0x0a, 0x08, 0x06, 0x04, 0x02, 0x00,\n\
|
||||
0x0f, 0x0d, 0x0b, 0x09, 0x07, 0x05, 0x03, 0x01,\n\
|
||||
0x0e, 0x0c, 0x0a, 0x08, 0x06, 0x04, 0x02, 0x00\n\
|
||||
);\n\
|
||||
#pragma unroll\n\
|
||||
for (int k = 0; k < act_k / 16; ++k) {\n\
|
||||
__m256 vec_b0f = _mm256_i32gather_ps(b + k * 16 + 0, vec_bi, 1);\n\
|
||||
__m256 vec_b1f = _mm256_i32gather_ps(b + k * 16 + 1, vec_bi, 1);\n\
|
||||
\n\
|
||||
__m256i vec_b0 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(vec_b0f, _mm256_set1_ps(scales)), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));\n\
|
||||
__m256i vec_b1 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(vec_b1f, _mm256_set1_ps(scales)), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));\n\
|
||||
vec_lut[15] = _mm256_setzero_si256();\n\
|
||||
vec_lut[14] = _mm256_setzero_si256();\n\
|
||||
vec_lut[13] = _mm256_setzero_si256();\n\
|
||||
vec_lut[12] = _mm256_setzero_si256();\n\
|
||||
vec_lut[11] = _mm256_setzero_si256();\n\
|
||||
vec_lut[10] = _mm256_setzero_si256();\n\
|
||||
vec_lut[9] = _mm256_setzero_si256();\n\
|
||||
vec_lut[8] = vec_b0;\n\
|
||||
vec_lut[8] = _mm256_add_epi32(vec_lut[8], vec_b1);\n\
|
||||
vec_lut[7] = vec_b0;\n\
|
||||
vec_lut[6] = vec_b0;\n\
|
||||
vec_lut[6] = _mm256_sub_epi32(vec_lut[6], vec_b1);\n\
|
||||
vec_lut[5] = vec_b1;\n\
|
||||
vec_lut[4] = _mm256_setzero_si256();\n\
|
||||
vec_lut[3] = _mm256_setzero_si256();\n\
|
||||
vec_lut[3] = _mm256_sub_epi32(vec_lut[3], vec_b1);\n\
|
||||
vec_lut[2] = _mm256_setzero_si256();\n\
|
||||
vec_lut[2] = _mm256_sub_epi32(vec_lut[2], vec_b0);\n\
|
||||
vec_lut[2] = _mm256_add_epi32(vec_lut[2], vec_b1);\n\
|
||||
vec_lut[1] = _mm256_setzero_si256();\n\
|
||||
vec_lut[1] = _mm256_sub_epi32(vec_lut[1], vec_b0);\n\
|
||||
vec_lut[0] = _mm256_setzero_si256();\n\
|
||||
vec_lut[0] = _mm256_sub_epi32(vec_lut[0], vec_b0);\n\
|
||||
vec_lut[0] = _mm256_sub_epi32(vec_lut[0], vec_b1);\n\
|
||||
\n\
|
||||
__m256i ix[16];\n\
|
||||
#pragma unroll\n\
|
||||
for (int g = 0; g < 16; ++g) {\n\
|
||||
ix[g] = vec_lut[g];\n\
|
||||
}\n\
|
||||
\n\
|
||||
Transpose_8_8(&(ix[0]), &(ix[1]), &(ix[2]), &(ix[3]), &(ix[4]), &(ix[5]),&(ix[6]), &(ix[7]));\n\
|
||||
Transpose_8_8(&(ix[8]), &(ix[9]), &(ix[10]), &(ix[11]), &(ix[12]), &(ix[13]),&(ix[14]), &(ix[15]));\n\
|
||||
\n\
|
||||
#pragma unroll\n\
|
||||
for (int g = 0; g < 8; ++g) {\n\
|
||||
ix[g] = _mm256_packs_epi32(ix[g], ix[g + 8]);\n\
|
||||
ix[g] = _mm256_permute4x64_epi64(ix[g], _MM_SHUFFLE(3, 1, 2, 0));\n\
|
||||
ix[g] = _mm256_shuffle_epi8(ix[g], shuffle_mask);\n\
|
||||
ix[g] = _mm256_permute4x64_epi64(ix[g], _MM_SHUFFLE(3, 1, 2, 0));\n\
|
||||
}\n\
|
||||
\n\
|
||||
int8_t* qlut_i8 = reinterpret_cast<int8_t*>(qlut);\n\
|
||||
\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 0 * 32 + 0), ix[0]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 1 * 32 + 0), ix[1]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 2 * 32 + 0), ix[2]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 3 * 32 + 0), ix[3]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 4 * 32 + 0), ix[4]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 5 * 32 + 0), ix[5]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 6 * 32 + 0), ix[6]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 7 * 32 + 0), ix[7]);\n\
|
||||
\n\
|
||||
}\n\
|
||||
*lut_scales = scales;\n\
|
||||
#endif\n\
|
||||
return 0;\n\
|
||||
}\n\
|
||||
static bool is_type_supported(enum ggml_type type) {\n\
|
||||
if (type == GGML_TYPE_Q4_0 ||\n\
|
||||
type == GGML_TYPE_TL2) {\n\
|
||||
return true;\n\
|
||||
} else {\n\
|
||||
return false;\n\
|
||||
}\n\
|
||||
}\n\
|
||||
"
|
||||
return kernel_code
|
||||
|
||||
def gen_tbl_impl(pre, BM, BK, bm, k_list):
|
||||
|
||||
kernel_code = "\
|
||||
#include <immintrin.h>\n\
|
||||
\n\
|
||||
#define BM{0} {1}\n\
|
||||
#define BBK{0} {2}\n\
|
||||
template<int batch_size, int K3>\n\
|
||||
inline void three_tbl_impl_{0}(int32_t* c, int8_t* lut, uint8_t* a, uint8_t* sign) {{\n\
|
||||
".format(pre, BM, BK)
|
||||
|
||||
kernel_code = "".join([kernel_code, "\
|
||||
#ifdef __AVX2__\n\
|
||||
const __m256i vec_mask = _mm256_set1_epi8(0x0f);\n\
|
||||
const __m256i vec_sign_mask = _mm256_set1_epi16(0x8000);\n\
|
||||
const __m256i vec_zero = _mm256_set1_epi8(0x00);\n\
|
||||
const __m256i vec_one = _mm256_set1_epi8(0xff);\n\
|
||||
const int KK = BBK{0} / 3;\n\
|
||||
#pragma unroll\n\
|
||||
for (int i = 0; i < BM{0}; i += 32) {{\n\
|
||||
__m256i vec_as[KK / 2];\n\
|
||||
__m256i vec_signs[KK / 8];\n\
|
||||
#pragma unroll\n\
|
||||
for (int ai = 0; ai < KK / 2; ai++) {{\n\
|
||||
vec_as[ai] = _mm256_loadu_si256(reinterpret_cast<__m256i*>(a + i * KK / 2 + ai * 32));\n\
|
||||
}}\n\
|
||||
#pragma unroll\n\
|
||||
for (int as = 0; as < KK / 8; as++) {{\n\
|
||||
vec_signs[as] = _mm256_loadu_si256(reinterpret_cast<__m256i*>(sign + i * KK / 8 + as * 32));\n\
|
||||
}}\n\
|
||||
#pragma unroll\n\
|
||||
for (int bs = 0; bs < batch_size; bs++) {{\n\
|
||||
__m256i vec_c0 = _mm256_setzero_si256();\n\
|
||||
__m256i vec_c1 = _mm256_setzero_si256();\n\
|
||||
#pragma unroll\n\
|
||||
for (int k = 0; k < KK / 8; k++) {{\n\
|
||||
__m256i vec_sign = vec_signs[k];\n\
|
||||
__m256i vec_a_0 = vec_as[k * 4 + 0];\n\
|
||||
__m128i vec_k1_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 0 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k2_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 16 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k3_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 32 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k4_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 48 + K3 / 3 * 32 * bs));\n\
|
||||
__m256i vec_sign_left_hi_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0)), 15);\n\
|
||||
__m256i vec_sign_left_lo_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0 + 1)), 15);\n\
|
||||
__m256i vec_v_top_0 = _mm256_and_si256(_mm256_srli_epi16(vec_a_0, 4), vec_mask);\n\
|
||||
__m256i vec_v_top_fir_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_0, vec_k1_0), vec_v_top_0);\n\
|
||||
__m256i vec_v_top_sec_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_0, vec_k2_0), vec_v_top_0);\n\
|
||||
__m256i vec_sign_right_hi_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0 + 2)), 15);\n\
|
||||
__m256i vec_sign_right_lo_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0 + 3)), 15);\n\
|
||||
__m256i vec_v_bot_0 = _mm256_and_si256(vec_a_0, vec_mask);\n\
|
||||
__m256i vec_v_bot_fir_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_0, vec_k3_0), vec_v_bot_0);\n\
|
||||
__m256i vec_v_bot_sec_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_0, vec_k4_0), vec_v_bot_0);\n\
|
||||
__m256i vec_v_top_lo_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_0, vec_v_top_sec_0), vec_sign_left_lo_0), vec_sign_left_lo_0);\n\
|
||||
__m256i vec_v_top_hi_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_0, vec_v_top_sec_0), vec_sign_left_hi_0), vec_sign_left_hi_0);\n\
|
||||
__m256i vec_v_bot_lo_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_0, vec_v_bot_sec_0), vec_sign_right_lo_0), vec_sign_right_lo_0);\n\
|
||||
__m256i vec_v_bot_hi_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_0, vec_v_bot_sec_0), vec_sign_right_hi_0), vec_sign_right_hi_0);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_0);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_0);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_0);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_0);\n\
|
||||
__m256i vec_a_1 = vec_as[k * 4 + 1];\n\
|
||||
__m128i vec_k1_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 0 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k2_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 16 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k3_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 32 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k4_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 48 + K3 / 3 * 32 * bs));\n\
|
||||
__m256i vec_sign_left_hi_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1)), 15);\n\
|
||||
__m256i vec_sign_left_lo_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1 + 1)), 15);\n\
|
||||
__m256i vec_v_top_1 = _mm256_and_si256(_mm256_srli_epi16(vec_a_1, 4), vec_mask);\n\
|
||||
__m256i vec_v_top_fir_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_1, vec_k1_1), vec_v_top_1);\n\
|
||||
__m256i vec_v_top_sec_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_1, vec_k2_1), vec_v_top_1);\n\
|
||||
__m256i vec_sign_right_hi_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1 + 2)), 15);\n\
|
||||
__m256i vec_sign_right_lo_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1 + 3)), 15);\n\
|
||||
__m256i vec_v_bot_1 = _mm256_and_si256(vec_a_1, vec_mask);\n\
|
||||
__m256i vec_v_bot_fir_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_1, vec_k3_1), vec_v_bot_1);\n\
|
||||
__m256i vec_v_bot_sec_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_1, vec_k4_1), vec_v_bot_1);\n\
|
||||
__m256i vec_v_top_lo_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_1, vec_v_top_sec_1), vec_sign_left_lo_1), vec_sign_left_lo_1);\n\
|
||||
__m256i vec_v_top_hi_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_1, vec_v_top_sec_1), vec_sign_left_hi_1), vec_sign_left_hi_1);\n\
|
||||
__m256i vec_v_bot_lo_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_1, vec_v_bot_sec_1), vec_sign_right_lo_1), vec_sign_right_lo_1);\n\
|
||||
__m256i vec_v_bot_hi_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_1, vec_v_bot_sec_1), vec_sign_right_hi_1), vec_sign_right_hi_1);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_1);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_1);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_1);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_1);\n\
|
||||
__m256i vec_a_2 = vec_as[k * 4 + 2];\n\
|
||||
__m128i vec_k1_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 0 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k2_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 16 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k3_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 32 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k4_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 48 + K3 / 3 * 32 * bs));\n\
|
||||
__m256i vec_sign_left_hi_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2)), 15);\n\
|
||||
__m256i vec_sign_left_lo_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2 + 1)), 15);\n\
|
||||
__m256i vec_v_top_2 = _mm256_and_si256(_mm256_srli_epi16(vec_a_2, 4), vec_mask);\n\
|
||||
__m256i vec_v_top_fir_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_2, vec_k1_2), vec_v_top_2);\n\
|
||||
__m256i vec_v_top_sec_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_2, vec_k2_2), vec_v_top_2);\n\
|
||||
__m256i vec_sign_right_hi_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2 + 2)), 15);\n\
|
||||
__m256i vec_sign_right_lo_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2 + 3)), 15);\n\
|
||||
__m256i vec_v_bot_2 = _mm256_and_si256(vec_a_2, vec_mask);\n\
|
||||
__m256i vec_v_bot_fir_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_2, vec_k3_2), vec_v_bot_2);\n\
|
||||
__m256i vec_v_bot_sec_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_2, vec_k4_2), vec_v_bot_2);\n\
|
||||
__m256i vec_v_top_lo_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_2, vec_v_top_sec_2), vec_sign_left_lo_2), vec_sign_left_lo_2);\n\
|
||||
__m256i vec_v_top_hi_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_2, vec_v_top_sec_2), vec_sign_left_hi_2), vec_sign_left_hi_2);\n\
|
||||
__m256i vec_v_bot_lo_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_2, vec_v_bot_sec_2), vec_sign_right_lo_2), vec_sign_right_lo_2);\n\
|
||||
__m256i vec_v_bot_hi_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_2, vec_v_bot_sec_2), vec_sign_right_hi_2), vec_sign_right_hi_2);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_2);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_2);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_2);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_2);\n\
|
||||
__m256i vec_a_3 = vec_as[k * 4 + 3];\n\
|
||||
__m128i vec_k1_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 0 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k2_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 16 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k3_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 32 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k4_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 48 + K3 / 3 * 32 * bs));\n\
|
||||
__m256i vec_sign_left_hi_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3)), 15);\n\
|
||||
__m256i vec_sign_left_lo_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3 + 1)), 15);\n\
|
||||
__m256i vec_v_top_3 = _mm256_and_si256(_mm256_srli_epi16(vec_a_3, 4), vec_mask);\n\
|
||||
__m256i vec_v_top_fir_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_3, vec_k1_3), vec_v_top_3);\n\
|
||||
__m256i vec_v_top_sec_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_3, vec_k2_3), vec_v_top_3);\n\
|
||||
__m256i vec_sign_right_hi_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3 + 2)), 15);\n\
|
||||
__m256i vec_sign_right_lo_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3 + 3)), 15);\n\
|
||||
__m256i vec_v_bot_3 = _mm256_and_si256(vec_a_3, vec_mask);\n\
|
||||
__m256i vec_v_bot_fir_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_3, vec_k3_3), vec_v_bot_3);\n\
|
||||
__m256i vec_v_bot_sec_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_3, vec_k4_3), vec_v_bot_3);\n\
|
||||
__m256i vec_v_top_lo_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_3, vec_v_top_sec_3), vec_sign_left_lo_3), vec_sign_left_lo_3);\n\
|
||||
__m256i vec_v_top_hi_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_3, vec_v_top_sec_3), vec_sign_left_hi_3), vec_sign_left_hi_3);\n\
|
||||
__m256i vec_v_bot_lo_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_3, vec_v_bot_sec_3), vec_sign_right_lo_3), vec_sign_right_lo_3);\n\
|
||||
__m256i vec_v_bot_hi_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_3, vec_v_bot_sec_3), vec_sign_right_hi_3), vec_sign_right_hi_3);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_3);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_3);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_3);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_3);\n\
|
||||
}}\n\
|
||||
__m256i vec_gc0 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + BM{0} * bs));\n\
|
||||
__m256i vec_gc1 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM{0} * bs));\n\
|
||||
__m256i vec_gc2 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM{0} * bs));\n\
|
||||
__m256i vec_gc3 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM{0} * bs));\n\
|
||||
vec_gc0 = _mm256_add_epi32(vec_gc0, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c0)));\n\
|
||||
vec_gc1 = _mm256_add_epi32(vec_gc1, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c0, 1)));\n\
|
||||
vec_gc2 = _mm256_add_epi32(vec_gc2, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c1)));\n\
|
||||
vec_gc3 = _mm256_add_epi32(vec_gc3, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c1, 1)));\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + BM{0} * bs), vec_gc0);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM{0} * bs), vec_gc1);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM{0} * bs), vec_gc2);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM{0} * bs), vec_gc3);\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
#endif\n\
|
||||
}}\n\
|
||||
\n\
|
||||
template<int batch_size, int K2>\n\
|
||||
inline int32_t two_tbl_impl{0}(int32_t* c, int8_t* lut, uint8_t* a) {{\n\
|
||||
#ifdef __AVX2__\n\
|
||||
const __m256i vec_mask = _mm256_set1_epi8(0x0f);\n\
|
||||
const int KK = BK2 / 2;\n\
|
||||
#pragma unroll\n\
|
||||
for (int i = 0; i < BM{0}; i += 32) {{\n\
|
||||
__m256i vec_as[KK / 2];\n\
|
||||
#pragma unroll\n\
|
||||
for (int ai = 0; ai < KK / 2; ai++) {{\n\
|
||||
vec_as[ai] = _mm256_loadu_si256(reinterpret_cast<__m256i*>(a + i * KK / 2 + ai * 32));\n\
|
||||
}}\n\
|
||||
#pragma unroll\n\
|
||||
for (int bs = 0; bs < batch_size; bs++) {{\n\
|
||||
__m256i vec_c0 = _mm256_setzero_si256();\n\
|
||||
__m256i vec_c1 = _mm256_setzero_si256();\n\
|
||||
#pragma unroll\n\
|
||||
for (int k = 0; k < KK / 8; k++) {{\n\
|
||||
#pragma unroll\n\
|
||||
for (int j = 0; j < 4; j++) {{\n\
|
||||
__m256i vec_a = vec_as[k * 4 + j];\n\
|
||||
\n\
|
||||
__m128i vec_k1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 0 + K2 / 2 * 32 * bs));\n\
|
||||
__m128i vec_k2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 16 + K2 / 2 * 32 * bs));\n\
|
||||
__m128i vec_k3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 32 + K2 / 2 * 32 * bs));\n\
|
||||
__m128i vec_k4 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 48 + K2 / 2 * 32 * bs));\n\
|
||||
\n\
|
||||
__m256i vec_v_top = _mm256_and_si256(_mm256_srli_epi16(vec_a, 4), vec_mask);\n\
|
||||
__m256i vec_v_top_fir = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1, vec_k1), vec_v_top);\n\
|
||||
__m256i vec_v_top_sec = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2, vec_k2), vec_v_top);\n\
|
||||
\n\
|
||||
__m256i vec_v_bot = _mm256_and_si256(vec_a, vec_mask);\n\
|
||||
__m256i vec_v_bot_fir = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3, vec_k3), vec_v_bot);\n\
|
||||
__m256i vec_v_bot_sec = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4, vec_k4), vec_v_bot);\n\
|
||||
\n\
|
||||
__m256i vec_v_top_lo = _mm256_unpackhi_epi8(vec_v_top_fir, vec_v_top_sec);\n\
|
||||
__m256i vec_v_top_hi = _mm256_unpacklo_epi8(vec_v_top_fir, vec_v_top_sec);\n\
|
||||
__m256i vec_v_bot_lo = _mm256_unpackhi_epi8(vec_v_bot_fir, vec_v_bot_sec);\n\
|
||||
__m256i vec_v_bot_hi = _mm256_unpacklo_epi8(vec_v_bot_fir, vec_v_bot_sec);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo); \n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
\n\
|
||||
__m256i vec_gc0 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + BM{0} * bs));\n\
|
||||
__m256i vec_gc1 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM{0} * bs));\n\
|
||||
__m256i vec_gc2 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM{0} * bs));\n\
|
||||
__m256i vec_gc3 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM{0} * bs));\n\
|
||||
\n\
|
||||
vec_gc0 = _mm256_add_epi32(vec_gc0, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c0)));\n\
|
||||
vec_gc1 = _mm256_add_epi32(vec_gc1, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c0, 1)));\n\
|
||||
vec_gc2 = _mm256_add_epi32(vec_gc2, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c1)));\n\
|
||||
vec_gc3 = _mm256_add_epi32(vec_gc3, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c1, 1)));\n\
|
||||
\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + BM{0} * bs), vec_gc0);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM{0} * bs), vec_gc1);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM{0} * bs), vec_gc2);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM{0} * bs), vec_gc3);\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
#endif\n\
|
||||
return 0;\n\
|
||||
}}\n\
|
||||
\n\
|
||||
template<int BATCH_SIZE>\n\
|
||||
int32_t three_qgemm_lut_{0}(void* A, void* sign, void* LUT, void* Scales, void* LUT_Scales, void* C) {{\n\
|
||||
alignas(32) uint32_t CBits[BATCH_SIZE * BM{0}];\n\
|
||||
memset(&(CBits[0]), 0, BATCH_SIZE * BM{0} * sizeof(int32_t));\n\
|
||||
#pragma unroll\n\
|
||||
for (int32_t k_outer = 0; k_outer < {1} / BBK{0}; ++k_outer) {{\n\
|
||||
three_tbl_impl_{0}<BATCH_SIZE, {1}>((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK{0} / 3 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK{0} / 3 / 2 * BM{0})])), (&(((uint8_t*)sign)[(k_outer * BBK{0} / 3 / 8 * BM{0})])));\n\
|
||||
}}\n\
|
||||
#pragma unroll\n\
|
||||
for (int bs = 0; bs < BATCH_SIZE; bs++) {{\n\
|
||||
#pragma unroll\n\
|
||||
for (int i = 0; i < BM{0}; i++) {{\n\
|
||||
((int32_t*)C)[i] = (int32_t)(((int32_t*)CBits)[i + bs * BM{0}]);\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
return 0;\n\
|
||||
}}\n\
|
||||
\n\
|
||||
template<int BATCH_SIZE>\n\
|
||||
int32_t two_qgemm_lut_{0}(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {{\n\
|
||||
alignas(32) uint32_t CBits[BATCH_SIZE * BM{0}];\n\
|
||||
memset(&(CBits[0]), 0, BATCH_SIZE * BM{0} * sizeof(int32_t));\n\
|
||||
#pragma unroll\n\
|
||||
for (int32_t k_outer = 0; k_outer < {2} / 32; ++k_outer) {{\n\
|
||||
two_tbl_impl{0}<BATCH_SIZE, {2}>((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BK2 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BK2 / 2 / 2 * BM{0})])));\n\
|
||||
}}\n\
|
||||
#pragma unroll\n\
|
||||
for (int bs = 0; bs < BATCH_SIZE; bs++) {{\n\
|
||||
#pragma unroll\n\
|
||||
for (int i = 0; i < BM{0}; i++) {{\n\
|
||||
((int32_t*)C)[i] += (int32_t)(((int32_t*)CBits)[i + bs * BM{0}]);\n\
|
||||
((float*)C)[i] = (float)(((int32_t*)C)[i]) / ((float*)LUT_Scales)[bs] * ((float*)Scales)[0];\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
return 0;\n\
|
||||
}}\n\
|
||||
\n\
|
||||
".format(pre, k_list[1], k_list[0])])
|
||||
return kernel_code
|
||||
|
||||
def gen_top_api(kernel_shapes, k_list):
|
||||
|
||||
kernel_code = "void ggml_preprocessor(int bs, int m, int three_k, int two_k, void* B, void* LUT_Scales, void* Three_QLUT, void* Two_QLUT) {{\n\
|
||||
partial_max_reset(bs, (&(((float*)LUT_Scales)[0])));\n\
|
||||
if (m == {0} && two_k == {1} && three_k == {2}) {{\n\
|
||||
for (int32_t b = 0; b < bs; b++) {{\n\
|
||||
per_tensor_quant(two_k + three_k, (&(((float*)LUT_Scales)[b])), (&(((float*)B)[b * (two_k + three_k)])));\n\
|
||||
three_lut_ctor<{2}>((&(((int8_t*)Three_QLUT)[b * three_k / 3 * 32])), (&(((float*)B)[b * (three_k + two_k)])), (&(((float*)LUT_Scales)[b])));\n\
|
||||
two_lut_ctor<{1}>((&(((int8_t*)Two_QLUT)[b * two_k / 2 * 32])), (&(((float*)B)[b * (three_k + two_k) + {2}])), (&(((float*)LUT_Scales)[b])));\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
".format(kernel_shapes[0][0], k_list[0][0], k_list[0][1])
|
||||
for i in range(1, len(kernel_shapes)):
|
||||
kernel_code = "".join([kernel_code, " else if (m == {0} && two_k == {1} && three_k == {2}) {{\n\
|
||||
for (int32_t b = 0; b < bs; b++) {{\n\
|
||||
per_tensor_quant(two_k + three_k, (&(((float*)LUT_Scales)[b])), (&(((float*)B)[b * (two_k + three_k)])));\n\
|
||||
three_lut_ctor<{2}>((&(((int8_t*)Three_QLUT)[b * three_k / 3 * 32])), (&(((float*)B)[b * (three_k + two_k)])), (&(((float*)LUT_Scales)[b])));\n\
|
||||
two_lut_ctor<{1}>((&(((int8_t*)Two_QLUT)[b * two_k / 2 * 32])), (&(((float*)B)[b * (three_k + two_k) + {2}])), (&(((float*)LUT_Scales)[b])));\n\
|
||||
}}\n\
|
||||
}}\n".format(kernel_shapes[i][0], k_list[i][0], k_list[i][1])])
|
||||
kernel_code = "".join([kernel_code, "}\n"])
|
||||
|
||||
|
||||
kernel_code = "".join([kernel_code, "void ggml_qgemm_lut(int bs, int m, int k, int BK, void* A, void* sign, void* LUT, void* Scales, void* LUT_Scales, void* C) {{\n\
|
||||
if (m == {0} && k == {1}) {{\n\
|
||||
if (BK == {2}) {{\n\
|
||||
if (bs == 1) {{\n\
|
||||
two_qgemm_lut_{4}<1>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 8) {{\n\
|
||||
two_qgemm_lut_{4}<8>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 32) {{\n\
|
||||
two_qgemm_lut_{4}<32>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 128) {{\n\
|
||||
two_qgemm_lut_{4}<128>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 256) {{\n\
|
||||
two_qgemm_lut_{4}<256>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 512) {{\n\
|
||||
two_qgemm_lut_{4}<512>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
else if (BK == {3}) {{\n\
|
||||
if (bs == 1) {{\n\
|
||||
three_qgemm_lut_{4}<1>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 8) {{\n\
|
||||
three_qgemm_lut_{4}<8>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 32) {{\n\
|
||||
three_qgemm_lut_{4}<32>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 128) {{\n\
|
||||
three_qgemm_lut_{4}<128>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 256) {{\n\
|
||||
three_qgemm_lut_{4}<256>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 512) {{\n\
|
||||
three_qgemm_lut_{4}<512>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
".format(kernel_shapes[0][0], kernel_shapes[0][1], k_list[0][0], k_list[0][1], "{}_{}".format(kernel_shapes[0][0], kernel_shapes[0][1]))])
|
||||
for i in range(1, len(kernel_shapes)):
|
||||
kernel_code = "".join([kernel_code, " else if (m == {0} && k == {1}) {{\n\
|
||||
if (BK == {2}) {{\n\
|
||||
if (bs == 1) {{\n\
|
||||
two_qgemm_lut_{4}<1>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 8) {{\n\
|
||||
two_qgemm_lut_{4}<8>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 32) {{\n\
|
||||
two_qgemm_lut_{4}<32>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 128) {{\n\
|
||||
two_qgemm_lut_{4}<128>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 256) {{\n\
|
||||
two_qgemm_lut_{4}<256>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 512) {{\n\
|
||||
two_qgemm_lut_{4}<512>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
else if (BK == {3}) {{\n\
|
||||
if (bs == 1) {{\n\
|
||||
three_qgemm_lut_{4}<1>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 8) {{\n\
|
||||
three_qgemm_lut_{4}<8>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 32) {{\n\
|
||||
three_qgemm_lut_{4}<32>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 128) {{\n\
|
||||
three_qgemm_lut_{4}<128>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 256) {{\n\
|
||||
three_qgemm_lut_{4}<256>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 512) {{\n\
|
||||
three_qgemm_lut_{4}<512>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
".format(kernel_shapes[i][0], kernel_shapes[i][1], k_list[i][0], k_list[i][1], "{}_{}".format(kernel_shapes[i][0], kernel_shapes[i][1]))])
|
||||
kernel_code = "".join([kernel_code, "}\n"])
|
||||
return kernel_code
|
||||
|
||||
def gen_transform_code(kernel_shapes):
|
||||
kernel_code = "\n\
|
||||
void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor) {\n\
|
||||
if (!(is_type_supported(tensor->type) && tensor->backend == GGML_BACKEND_TYPE_CPU && tensor->extra == nullptr)) {\n\
|
||||
return;\n\
|
||||
}\n\
|
||||
\n\
|
||||
int k = tensor->ne[0];\n\
|
||||
int m = tensor->ne[1];\n\
|
||||
const int lut_scales_size = 1;\n\
|
||||
int bk = 0;\n\
|
||||
int bm = 0;\n"
|
||||
|
||||
kernel_code = "".join([kernel_code, "\n\
|
||||
if (m == {0} && k == {1}) {{\n\
|
||||
bm = BM{0}_{1};\n\
|
||||
bk = BBK{0}_{1};\n\
|
||||
}}\n".format(kernel_shapes[0][0], kernel_shapes[0][1])])
|
||||
|
||||
for i in range(1, len(kernel_shapes)):
|
||||
kernel_code = "".join([kernel_code, "else if (m == {0} && k == {1}) {{\n\
|
||||
bm = BM{0}_{1};\n\
|
||||
bk = BBK{0}_{1};\n\
|
||||
}}\n".format(kernel_shapes[i][0], kernel_shapes[i][1])])
|
||||
|
||||
kernel_code = "".join([kernel_code, "\n\
|
||||
const int n_tile_num = m / bm;\n\
|
||||
const int BK = bk;\n\
|
||||
uint8_t * qweights;\n\
|
||||
bitnet_float_type * scales;\n\
|
||||
\n\
|
||||
scales = (bitnet_float_type *) aligned_malloc(sizeof(bitnet_float_type));\n\
|
||||
qweights = (uint8_t *) tensor->data;\n\
|
||||
int nbytes = (k - 256) * m / 3 * 5 / 8 + 256 * m / 2 * 4 / 8;\n\
|
||||
if (nbytes % 32 != 0) nbytes = 32 - nbytes % 32 + nbytes;\n\
|
||||
float * i2_scales = (float * )(qweights + nbytes);\n\
|
||||
scales[0] = (bitnet_float_type) i2_scales[0];\n\
|
||||
\n\
|
||||
tensor->extra = bitnet_tensor_extras + bitnet_tensor_extras_index;\n\
|
||||
bitnet_tensor_extras[bitnet_tensor_extras_index++] = {\n\
|
||||
/* .lut_scales_size = */ lut_scales_size,\n\
|
||||
/* .BK = */ BK,\n\
|
||||
/* .n_tile_num = */ n_tile_num,\n\
|
||||
/* .qweights = */ qweights,\n\
|
||||
/* .scales = */ scales\n\
|
||||
};\n\
|
||||
}\n"])
|
||||
|
||||
return kernel_code
|
||||
|
||||
def get_three_k_two_k(K, bk):
|
||||
bk_num = K // bk
|
||||
three_k = bk_num * bk
|
||||
two_k = K - three_k
|
||||
return two_k, three_k
|
||||
|
||||
if __name__ == "__main__":
|
||||
ModelShapeDict = {
|
||||
"bitnet_b1_58-large" : [[1536, 4096],
|
||||
[1536, 1536],
|
||||
[4096, 1536]],
|
||||
"bitnet_b1_58-3B" : [[3200, 8640],
|
||||
[3200, 3200],
|
||||
[8640, 3200]],
|
||||
"Llama3-8B-1.58-100B-tokens" : [[14336, 4096],
|
||||
[4096, 14336],
|
||||
[1024, 4096],
|
||||
[4096, 4096]]
|
||||
}
|
||||
|
||||
parser = argparse.ArgumentParser(description='gen impl')
|
||||
parser.add_argument('--model',default="input", type=str, dest="model",
|
||||
help="choose from bitnet_b1_58-large/bitnet_b1_58-3B/Llama3-8B-1.58-100B-tokens.")
|
||||
parser.add_argument('--BM',default="input", type=str,
|
||||
help="block length when cutting one weight (M, K) into M / BM weights (BM, K).")
|
||||
parser.add_argument('--BK',default="input", type=str,
|
||||
help="block length when cutting one weight (M, K) into K / BK weights (M, BK).")
|
||||
parser.add_argument('--bm',default="input", type=str,
|
||||
help="using simd instructions to compute (bm, 192 / bm) in one block")
|
||||
args = parser.parse_args()
|
||||
|
||||
kernel_shapes = ModelShapeDict[args.model]
|
||||
|
||||
BM_list = [int(item) for item in args.BM.split(',')]
|
||||
BK_list = [int(item) for item in args.BK.split(',')]
|
||||
bm_list = [int(item) for item in args.bm.split(',')]
|
||||
|
||||
tbl_impl_code = []
|
||||
k_list = []
|
||||
|
||||
for i in range(len(kernel_shapes)):
|
||||
k_list.append(get_three_k_two_k(kernel_shapes[i][1], BK_list[i]))
|
||||
|
||||
for i in range(len(kernel_shapes)):
|
||||
tbl_impl_code.append(
|
||||
gen_tbl_impl("{}_{}".format(kernel_shapes[i][0], kernel_shapes[i][1]), BM_list[i], BK_list[i], bm_list[i], k_list[i])
|
||||
)
|
||||
|
||||
assert(len(BM_list) == len(BK_list) == len(bm_list) == len(kernel_shapes)), "number of BM / BK / bm shoud be {}".format(len(kernel_shapes))
|
||||
|
||||
for i in range(len(kernel_shapes)):
|
||||
assert kernel_shapes[i][0] % BM_list[i] == 0, "M %% BM should be 0"
|
||||
assert (kernel_shapes[i][1] % BK_list[i]) % 32 == 0, "K %% BK %% 32 should be 0"
|
||||
assert bm_list[i] in [32], "choose bm from [32]"
|
||||
|
||||
ctor_code = gen_ctor_code()
|
||||
api_code = gen_top_api(kernel_shapes, k_list)
|
||||
trans_code = gen_transform_code(kernel_shapes)
|
||||
|
||||
output_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "include")
|
||||
|
||||
with open(''.join([output_dir, "/bitnet-lut-kernels.h"]), 'w') as f:
|
||||
f.write(''.join("#if defined(GGML_BITNET_X86_TL2)"))
|
||||
f.write(''.join(ctor_code))
|
||||
for code in tbl_impl_code:
|
||||
f.write(''.join(code))
|
||||
f.write(''.join(api_code))
|
||||
f.write(''.join(trans_code))
|
||||
f.write(''.join("#endif"))
|
||||
|
||||
config = ConfigParser()
|
||||
|
||||
for i in range(len(kernel_shapes)):
|
||||
config.add_section('Kernels_{}'.format(i))
|
||||
config.set('Kernels_{}'.format(i), 'M'.format(i), str(kernel_shapes[i][0]))
|
||||
config.set('Kernels_{}'.format(i), 'K'.format(i), str(kernel_shapes[i][1]))
|
||||
config.set('Kernels_{}'.format(i), 'BM'.format(i), str(BM_list[i]))
|
||||
config.set('Kernels_{}'.format(i), 'BK'.format(i), str(BK_list[i]))
|
||||
config.set('Kernels_{}'.format(i), 'bmm'.format(i), str(bm_list[i]))
|
||||
|
||||
with open(''.join([output_dir, "/kernel_config.ini"]), 'w') as configfile:
|
||||
config.write(configfile)
|
||||
@@ -0,0 +1,134 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import sys
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
def run_command(command_list, cwd=None, check=True):
|
||||
print(f"Executing: {' '.join(map(str, command_list))}")
|
||||
try:
|
||||
process = subprocess.run(command_list, cwd=cwd, check=check, capture_output=False, text=True)
|
||||
return process
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"Error executing command: {' '.join(map(str, e.cmd))}")
|
||||
print(f"Return code: {e.returncode}")
|
||||
raise
|
||||
|
||||
def main():
|
||||
if len(sys.argv) < 2:
|
||||
script_name = Path(sys.argv[0]).name
|
||||
print(f"Usage: python {script_name} <model-directory>")
|
||||
sys.exit(1)
|
||||
|
||||
model_dir_arg = sys.argv[1]
|
||||
model_dir = Path(model_dir_arg).resolve()
|
||||
|
||||
if not model_dir.is_dir():
|
||||
print(f"Error: Model directory '{model_dir}' not found or is not a directory.")
|
||||
sys.exit(1)
|
||||
|
||||
utils_dir = Path(__file__).parent.resolve()
|
||||
project_root_dir = utils_dir.parent
|
||||
|
||||
preprocess_script = utils_dir / "preprocess-huggingface-bitnet.py"
|
||||
convert_script = utils_dir / "convert-ms-to-gguf-bitnet.py"
|
||||
|
||||
llama_quantize_binary = project_root_dir / "build" / "bin" / "llama-quantize"
|
||||
|
||||
input_file = model_dir / "model.safetensors"
|
||||
input_backup_file = model_dir / "model.safetensors.backup"
|
||||
preprocessed_output_file = model_dir / "model.safetensors"
|
||||
|
||||
gguf_f32_output = model_dir / "ggml-model-f32-bitnet.gguf"
|
||||
gguf_i2s_output = model_dir / "ggml-model-i2s-bitnet.gguf"
|
||||
|
||||
if not preprocess_script.is_file():
|
||||
print(f"Error: Preprocess script not found at '{preprocess_script}'")
|
||||
sys.exit(1)
|
||||
if not convert_script.is_file():
|
||||
print(f"Error: Convert script not found at '{convert_script}'")
|
||||
sys.exit(1)
|
||||
if not llama_quantize_binary.is_file():
|
||||
print(f"Error: llama-quantize binary not found at '{llama_quantize_binary}'")
|
||||
sys.exit(1)
|
||||
|
||||
if not input_file.is_file():
|
||||
print(f"Error: Input safetensors file not found at '{input_file}'")
|
||||
sys.exit(1)
|
||||
|
||||
try:
|
||||
print(f"Backing up '{input_file}' to '{input_backup_file}'")
|
||||
if input_backup_file.exists():
|
||||
print(f"Warning: Removing existing backup file '{input_backup_file}'")
|
||||
input_backup_file.unlink()
|
||||
shutil.move(input_file, input_backup_file)
|
||||
|
||||
print("Preprocessing huggingface checkpoint...")
|
||||
cmd_preprocess = [
|
||||
sys.executable,
|
||||
str(preprocess_script),
|
||||
"--input", str(input_backup_file),
|
||||
"--output", str(preprocessed_output_file)
|
||||
]
|
||||
run_command(cmd_preprocess)
|
||||
|
||||
print("Converting to GGUF (f32)...")
|
||||
cmd_convert = [
|
||||
sys.executable,
|
||||
str(convert_script),
|
||||
str(model_dir),
|
||||
"--vocab-type", "bpe",
|
||||
"--outtype", "f32",
|
||||
"--concurrency", "1",
|
||||
"--outfile", str(gguf_f32_output)
|
||||
]
|
||||
run_command(cmd_convert)
|
||||
|
||||
print("Quantizing model to I2_S...")
|
||||
cmd_quantize = [
|
||||
str(llama_quantize_binary),
|
||||
str(gguf_f32_output),
|
||||
str(gguf_i2s_output),
|
||||
"I2_S",
|
||||
"1"
|
||||
]
|
||||
run_command(cmd_quantize)
|
||||
|
||||
print("Convert successfully.")
|
||||
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}")
|
||||
finally:
|
||||
print("Cleaning up intermediate files...")
|
||||
if preprocessed_output_file.exists() and preprocessed_output_file != input_backup_file:
|
||||
print(f"Removing preprocessed file: {preprocessed_output_file}")
|
||||
try:
|
||||
preprocessed_output_file.unlink()
|
||||
except OSError as e:
|
||||
print(f"Warning: Could not remove {preprocessed_output_file}: {e}")
|
||||
|
||||
# if gguf_f32_output.exists():
|
||||
# print(f"Removing f32 GGUF: {gguf_f32_output}")
|
||||
# try:
|
||||
# gguf_f32_output.unlink()
|
||||
# except OSError as e:
|
||||
# print(f"Warning: Could not remove {gguf_f32_output}: {e}")
|
||||
|
||||
if input_backup_file.exists():
|
||||
if not input_file.exists():
|
||||
print(f"Restoring original '{input_file}' from '{input_backup_file}'")
|
||||
try:
|
||||
shutil.move(input_backup_file, input_file)
|
||||
except Exception as e:
|
||||
print(f"Warning: Could not restore {input_file} from backup: {e}")
|
||||
else:
|
||||
print(f"Removing backup '{input_backup_file}' as original '{input_file}' should be present.")
|
||||
try:
|
||||
input_backup_file.unlink()
|
||||
except OSError as e:
|
||||
print(f"Warning: Could not remove backup {input_backup_file}: {e}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
+1610
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,59 @@
|
||||
import os
|
||||
import sys
|
||||
import logging
|
||||
import argparse
|
||||
import platform
|
||||
import subprocess
|
||||
|
||||
def run_command(command, shell=False, log_step=None):
|
||||
"""Run a system command and ensure it succeeds."""
|
||||
if log_step:
|
||||
log_file = os.path.join(args.log_dir, log_step + ".log")
|
||||
with open(log_file, "w") as f:
|
||||
try:
|
||||
subprocess.run(command, shell=shell, check=True, stdout=f, stderr=f)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logging.error(f"Error occurred while running command: {e}, check details in {log_file}")
|
||||
sys.exit(1)
|
||||
else:
|
||||
try:
|
||||
subprocess.run(command, shell=shell, check=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logging.error(f"Error occurred while running command: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
def run_benchmark():
|
||||
build_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "build")
|
||||
if platform.system() == "Windows":
|
||||
bench_path = os.path.join(build_dir, "bin", "Release", "llama-bench.exe")
|
||||
if not os.path.exists(bench_path):
|
||||
bench_path = os.path.join(build_dir, "bin", "llama-bench")
|
||||
else:
|
||||
bench_path = os.path.join(build_dir, "bin", "llama-bench")
|
||||
if not os.path.exists(bench_path):
|
||||
logging.error(f"Benchmark binary not found, please build first.")
|
||||
sys.exit(1)
|
||||
command = [
|
||||
f'{bench_path}',
|
||||
'-m', args.model,
|
||||
'-n', str(args.n_token),
|
||||
'-ngl', '0',
|
||||
'-b', '1',
|
||||
'-t', str(args.threads),
|
||||
'-p', str(args.n_prompt),
|
||||
'-r', '5'
|
||||
]
|
||||
run_command(command)
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description='Setup the environment for running the inference')
|
||||
parser.add_argument("-m", "--model", type=str, help="Path to model file", required=True)
|
||||
parser.add_argument("-n", "--n-token", type=int, help="Number of generated tokens", required=False, default=128)
|
||||
parser.add_argument("-p", "--n-prompt", type=int, help="Prompt to generate text from", required=False, default=512)
|
||||
parser.add_argument("-t", "--threads", type=int, help="Number of threads to use", required=False, default=2)
|
||||
return parser.parse_args()
|
||||
|
||||
if __name__ == "__main__":
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
args = parse_args()
|
||||
run_benchmark()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,50 @@
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import save_file
|
||||
import torch
|
||||
|
||||
def quant_weight_fp16(weight):
|
||||
weight = weight.to(torch.float)
|
||||
s = 1.0 / weight.abs().mean().clamp_(min=1e-5)
|
||||
new_weight = (weight * s).round().clamp(-1, 1) / s
|
||||
return new_weight
|
||||
|
||||
def quant_model(input, output):
|
||||
tensors = {}
|
||||
|
||||
with safe_open(input, framework='pt') as f:
|
||||
for name in f.keys():
|
||||
tensors[name] = f.get_tensor(name)
|
||||
|
||||
keyword_list = [
|
||||
'q_proj.weight',
|
||||
'k_proj.weight',
|
||||
'v_proj.weight',
|
||||
'o_proj.weight',
|
||||
'gate_proj.weight',
|
||||
'up_proj.weight',
|
||||
'down_proj.weight'
|
||||
]
|
||||
|
||||
if any(keyword in name for keyword in keyword_list):
|
||||
print(f'[INFO] Quantizing {name}')
|
||||
tensors[name] = quant_weight_fp16(tensors[name])
|
||||
|
||||
print(f'[INFO] Saving to {output}\nThis may take a while.')
|
||||
save_file(tensors, output)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser(description="Convert Safetensors back to Torch .pth checkpoint")
|
||||
parser.add_argument(
|
||||
"--input", type=str, required=True,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output", type=str, required=True,
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
quant_model(
|
||||
input=args.input,
|
||||
output=args.output,
|
||||
)
|
||||
@@ -0,0 +1,473 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Embedding Quantization Script
|
||||
This script converts ggml-model-f32.gguf to multiple quantized versions
|
||||
with different token embedding types.
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import os
|
||||
import argparse
|
||||
import re
|
||||
import csv
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class EmbeddingQuantizer:
|
||||
def __init__(self, input_model, output_dir, quantize_bin="../build/bin/llama-quantize",
|
||||
bench_bin="../build/bin/llama-bench", stats_dir="../stats", csv_output=None):
|
||||
self.input_model = Path(input_model)
|
||||
self.output_dir = Path(output_dir)
|
||||
self.quantize_bin = Path(quantize_bin)
|
||||
self.bench_bin = Path(bench_bin)
|
||||
self.stats_dir = Path(stats_dir)
|
||||
self.csv_output = Path(csv_output) if csv_output else None
|
||||
|
||||
# Verify input file exists
|
||||
if not self.input_model.exists():
|
||||
raise FileNotFoundError(f"Input model not found: {self.input_model}")
|
||||
|
||||
# Verify quantize tool exists
|
||||
if not self.quantize_bin.exists():
|
||||
raise FileNotFoundError(f"Quantize binary not found: {self.quantize_bin}")
|
||||
|
||||
# Verify bench tool exists
|
||||
if not self.bench_bin.exists():
|
||||
raise FileNotFoundError(f"Benchmark binary not found: {self.bench_bin}")
|
||||
|
||||
# Create output directories
|
||||
self.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.stats_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.results = []
|
||||
self.newly_created_files = set() # Track newly created files
|
||||
|
||||
def quantize(self, embedding_type, output_suffix):
|
||||
"""
|
||||
Perform single quantization
|
||||
|
||||
Args:
|
||||
embedding_type: Token embedding type (uppercase format, e.g., Q6_K)
|
||||
output_suffix: Output file suffix (lowercase format, e.g., q6_k)
|
||||
|
||||
Returns:
|
||||
bool: Whether successful
|
||||
"""
|
||||
output_file = self.output_dir / f"ggml-model-i2_s-embed-{output_suffix}.gguf"
|
||||
|
||||
# Check if file already exists
|
||||
file_already_existed = output_file.exists()
|
||||
|
||||
if file_already_existed:
|
||||
print(f"ℹ️ File already exists: {output_file}")
|
||||
print(f" Skipping quantization, will use existing file for benchmark")
|
||||
return True
|
||||
|
||||
cmd = [
|
||||
str(self.quantize_bin),
|
||||
"--token-embedding-type", embedding_type,
|
||||
str(self.input_model),
|
||||
str(output_file),
|
||||
"I2_S",
|
||||
"1",
|
||||
"1"
|
||||
]
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print(f"🔄 Quantizing with embedding type: {embedding_type}")
|
||||
print(f"📥 Input: {self.input_model}")
|
||||
print(f"📤 Output: {output_file}")
|
||||
print(f"💻 Command: {' '.join(cmd)}")
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
start_time = datetime.now()
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
cwd=os.getcwd(),
|
||||
timeout=600 # 10 minute timeout
|
||||
)
|
||||
|
||||
end_time = datetime.now()
|
||||
duration = (end_time - start_time).total_seconds()
|
||||
|
||||
if result.returncode == 0:
|
||||
# Get output file size
|
||||
file_size_mb = output_file.stat().st_size / (1024 * 1024)
|
||||
|
||||
print(f"✅ Success! Duration: {duration:.2f}s, Size: {file_size_mb:.2f} MB")
|
||||
|
||||
# Record newly created file
|
||||
if not file_already_existed:
|
||||
self.newly_created_files.add(output_file)
|
||||
|
||||
# Print part of output
|
||||
if result.stdout:
|
||||
print("\n📊 Quantization output:")
|
||||
print(result.stdout[-500:] if len(result.stdout) > 500 else result.stdout)
|
||||
|
||||
return True
|
||||
else:
|
||||
print(f"❌ Failed with return code {result.returncode}")
|
||||
print(f"Error: {result.stderr}")
|
||||
return False
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
print(f"❌ Timeout (exceeded 10 minutes)")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Exception: {e}")
|
||||
return False
|
||||
|
||||
def benchmark_model(self, output_suffix):
|
||||
"""
|
||||
Benchmark model
|
||||
|
||||
Args:
|
||||
output_suffix: Output file suffix (lowercase format, e.g., q6_k)
|
||||
|
||||
Returns:
|
||||
dict: Dictionary with benchmark results, or None if failed
|
||||
"""
|
||||
model_file = self.output_dir / f"ggml-model-i2_s-embed-{output_suffix}.gguf"
|
||||
|
||||
if not model_file.exists():
|
||||
print(f"❌ Model file not found for benchmarking: {model_file}")
|
||||
return None
|
||||
|
||||
cmd = [
|
||||
str(self.bench_bin),
|
||||
"-m", str(model_file),
|
||||
"-p", "128",
|
||||
"-n", "0",
|
||||
"-t", "1,2,4,8",
|
||||
"-ngl", "0"
|
||||
]
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print(f"🏃 Running benchmark for: {output_suffix}")
|
||||
print(f"💻 Command: {' '.join(cmd)}")
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
cwd=os.getcwd(),
|
||||
timeout=300 # 5 minute timeout
|
||||
)
|
||||
|
||||
if result.returncode == 0:
|
||||
print("✅ Benchmark completed successfully")
|
||||
print("\n📊 Benchmark output:")
|
||||
print(result.stdout)
|
||||
|
||||
# 解析输出
|
||||
bench_results = self.parse_benchmark_output(result.stdout, output_suffix)
|
||||
return bench_results
|
||||
else:
|
||||
print(f"❌ Benchmark failed with return code {result.returncode}")
|
||||
print(f"Error: {result.stderr}")
|
||||
return None
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
print(f"❌ Benchmark timeout (exceeded 5 minutes)")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Benchmark exception: {e}")
|
||||
return None
|
||||
|
||||
def parse_benchmark_output(self, output, output_suffix):
|
||||
"""
|
||||
Parse benchmark output to extract t/s data (mean±std)
|
||||
|
||||
Args:
|
||||
output: Benchmark command output
|
||||
output_suffix: Output file suffix
|
||||
|
||||
Returns:
|
||||
dict: Dictionary with parsed results
|
||||
"""
|
||||
results = {
|
||||
'embedding_type': output_suffix,
|
||||
'threads_1': None,
|
||||
'threads_2': None,
|
||||
'threads_4': None,
|
||||
'threads_8': None,
|
||||
}
|
||||
|
||||
# Parse table data
|
||||
# Find lines containing pp128 and t/s
|
||||
lines = output.strip().split('\n')
|
||||
|
||||
for line in lines:
|
||||
# Skip header and separator lines
|
||||
if '|' not in line or 'model' in line or '---' in line:
|
||||
continue
|
||||
|
||||
# Try to extract data
|
||||
# Format similar to: | bitnet-25 2B I2_S - 2 bpw ternary | 1012.28 MiB | 2.74 B | CPU | 12 | pp128 | 405.73 ± 3.69 |
|
||||
parts = [p.strip() for p in line.split('|')]
|
||||
|
||||
if len(parts) >= 8 and 'pp128' in parts[6]:
|
||||
threads_str = parts[5].strip()
|
||||
throughput_str = parts[7].strip()
|
||||
|
||||
# Extract thread count
|
||||
try:
|
||||
threads = int(threads_str)
|
||||
except:
|
||||
continue
|
||||
|
||||
# Extract t/s data (format: "405.73 ± 3.69" or "405.73")
|
||||
# Try to match "mean ± std" format
|
||||
match_with_std = re.search(r'([\d.]+)\s*±\s*([\d.]+)', throughput_str)
|
||||
if match_with_std:
|
||||
mean = float(match_with_std.group(1))
|
||||
std = float(match_with_std.group(2))
|
||||
throughput = f"{mean:.2f}±{std:.2f}"
|
||||
else:
|
||||
# Only mean, no std
|
||||
match = re.search(r'([\d.]+)', throughput_str)
|
||||
if match:
|
||||
throughput = f"{float(match.group(1)):.2f}"
|
||||
else:
|
||||
continue
|
||||
|
||||
# Store result based on thread count
|
||||
if threads == 1:
|
||||
results['threads_1'] = throughput
|
||||
elif threads == 2:
|
||||
results['threads_2'] = throughput
|
||||
elif threads == 4:
|
||||
results['threads_4'] = throughput
|
||||
elif threads == 8:
|
||||
results['threads_8'] = throughput
|
||||
|
||||
return results
|
||||
|
||||
def cleanup_model(self, output_suffix):
|
||||
"""
|
||||
Cleanup model files (only delete newly created files)
|
||||
|
||||
Args:
|
||||
output_suffix: Output file suffix
|
||||
"""
|
||||
model_file = self.output_dir / f"ggml-model-i2_s-embed-{output_suffix}.gguf"
|
||||
|
||||
if model_file in self.newly_created_files:
|
||||
try:
|
||||
model_file.unlink()
|
||||
print(f"🗑️ Deleted newly created file: {model_file}")
|
||||
self.newly_created_files.remove(model_file)
|
||||
except Exception as e:
|
||||
print(f"⚠️ Failed to delete {model_file}: {e}")
|
||||
else:
|
||||
print(f"ℹ️ Keeping existing file: {model_file}")
|
||||
|
||||
def run_all_quantizations(self, types_to_quantize):
|
||||
"""
|
||||
Run all quantizations
|
||||
|
||||
Args:
|
||||
types_to_quantize: List of quantization types, tuples of (embedding_type, output_suffix)
|
||||
"""
|
||||
print(f"\n{'='*80}")
|
||||
print(f"🚀 Starting Embedding Quantization and Benchmarking")
|
||||
print(f"{'='*80}")
|
||||
print(f"📥 Input model: {self.input_model}")
|
||||
print(f"📤 Output directory: {self.output_dir}")
|
||||
print(f"📊 Stats directory: {self.stats_dir}")
|
||||
print(f"🔢 Total quantizations: {len(types_to_quantize)}")
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
total_start = datetime.now()
|
||||
|
||||
for i, (embedding_type, output_suffix) in enumerate(types_to_quantize, 1):
|
||||
print(f"\n{'#'*80}")
|
||||
print(f"[{i}/{len(types_to_quantize)}] Processing {output_suffix} ({embedding_type})")
|
||||
print(f"{'#'*80}\n")
|
||||
|
||||
# Quantize model
|
||||
success = self.quantize(embedding_type, output_suffix)
|
||||
|
||||
if not success:
|
||||
print(f"⚠️ Skipping benchmark for {output_suffix} due to quantization failure")
|
||||
continue
|
||||
|
||||
# Run benchmark
|
||||
bench_results = self.benchmark_model(output_suffix)
|
||||
|
||||
if bench_results:
|
||||
self.results.append(bench_results)
|
||||
else:
|
||||
print(f"⚠️ Benchmark failed for {output_suffix}")
|
||||
|
||||
# Cleanup model files (only delete newly created files)
|
||||
self.cleanup_model(output_suffix)
|
||||
|
||||
print(f"\n{'#'*80}")
|
||||
print(f"✅ Completed {output_suffix}")
|
||||
print(f"{'#'*80}\n")
|
||||
|
||||
total_end = datetime.now()
|
||||
total_duration = (total_end - total_start).total_seconds()
|
||||
|
||||
# 保存结果到CSV
|
||||
self.save_results_to_csv()
|
||||
|
||||
# 打印总结
|
||||
self.print_summary(total_duration)
|
||||
|
||||
def save_results_to_csv(self):
|
||||
"""将benchmark结果保存到CSV文件"""
|
||||
if not self.results:
|
||||
print("⚠️ No results to save")
|
||||
return
|
||||
|
||||
# Use user-specified CSV path, otherwise use default path
|
||||
if self.csv_output:
|
||||
csv_file = self.csv_output
|
||||
# Ensure parent directory exists
|
||||
csv_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
else:
|
||||
csv_file = self.stats_dir / f"embedding_benchmark.csv"
|
||||
|
||||
print(f"\n💾 Saving results to: {csv_file}")
|
||||
|
||||
try:
|
||||
with open(csv_file, 'w', newline='') as f:
|
||||
fieldnames = ['embedding_type', 'threads_1', 'threads_2', 'threads_4', 'threads_8']
|
||||
writer = csv.DictWriter(f, fieldnames=fieldnames)
|
||||
|
||||
writer.writeheader()
|
||||
for result in self.results:
|
||||
writer.writerow(result)
|
||||
|
||||
print(f"✅ Results saved successfully")
|
||||
|
||||
# Also print table
|
||||
print(f"\n📊 Benchmark Results:")
|
||||
print(f"{'Type':<15} {'1 thread':<18} {'2 threads':<18} {'4 threads':<18} {'8 threads':<18}")
|
||||
print("-" * 87)
|
||||
for result in self.results:
|
||||
t1 = result['threads_1'] if result['threads_1'] else "N/A"
|
||||
t2 = result['threads_2'] if result['threads_2'] else "N/A"
|
||||
t4 = result['threads_4'] if result['threads_4'] else "N/A"
|
||||
t8 = result['threads_8'] if result['threads_8'] else "N/A"
|
||||
print(f"{result['embedding_type']:<15} {t1:<18} {t2:<18} {t4:<18} {t8:<18}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to save results: {e}")
|
||||
|
||||
def print_summary(self, total_duration):
|
||||
"""Print quantization summary"""
|
||||
print(f"\n\n{'='*80}")
|
||||
print(f"📊 QUANTIZATION AND BENCHMARK SUMMARY")
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
successful = len(self.results)
|
||||
total = len(self.results)
|
||||
|
||||
print(f"✅ Completed: {successful} benchmarks")
|
||||
print(f"⏱️ Total duration: {total_duration/60:.2f} minutes\n")
|
||||
|
||||
if self.results:
|
||||
if self.csv_output and self.csv_output.exists():
|
||||
print(f"📁 Results saved to: {self.csv_output}")
|
||||
else:
|
||||
csv_files = list(self.stats_dir.glob("embedding_benchmark*.csv"))
|
||||
if csv_files:
|
||||
latest_csv = max(csv_files, key=lambda p: p.stat().st_mtime)
|
||||
print(f"📁 Results saved to: {latest_csv}")
|
||||
|
||||
print(f"\n{'='*80}\n")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Quantize model embeddings to multiple formats')
|
||||
parser.add_argument('--input', '-i',
|
||||
default='../models/BitNet-b1.58-2B-4T/ggml-model-f32.gguf',
|
||||
help='Input model path (default: ../models/BitNet-b1.58-2B-4T/ggml-model-f32.gguf)')
|
||||
parser.add_argument('--output-dir', '-o',
|
||||
default='../models/BitNet-b1.58-2B-4T',
|
||||
help='Output directory (default: ../models/BitNet-b1.58-2B-4T)')
|
||||
parser.add_argument('--quantize-bin', '-q',
|
||||
default='../build/bin/llama-quantize',
|
||||
help='Path to llama-quantize binary (default: ../build/bin/llama-quantize)')
|
||||
parser.add_argument('--bench-bin', '-b',
|
||||
default='../build/bin/llama-bench',
|
||||
help='Path to llama-bench binary (default: ../build/bin/llama-bench)')
|
||||
parser.add_argument('--stats-dir',
|
||||
default='../stats',
|
||||
help='Directory to save benchmark results (default: ../stats)')
|
||||
parser.add_argument('--csv-output', '-c',
|
||||
help='Custom path for CSV output file (e.g., stats/my_results.csv)')
|
||||
parser.add_argument('--types', '-t',
|
||||
nargs='+',
|
||||
help='Specific types to quantize (e.g., f32 q6_k q4_0)')
|
||||
parser.add_argument('--skip-existing', '-s',
|
||||
action='store_true',
|
||||
help='Skip quantization if output file already exists (will still benchmark existing files)')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Define all supported quantization types
|
||||
# Format: (embedding_type for command line, output_suffix for filename)
|
||||
all_types = [
|
||||
('F32', 'f32'),
|
||||
('F16', 'f16'),
|
||||
('Q8_0', 'q8_0'),
|
||||
('Q6_K', 'q6_k'),
|
||||
('Q5_0', 'q5_0'),
|
||||
('Q4_0', 'q4_0'),
|
||||
('Q3_K', 'q3_k'),
|
||||
('TQ2_0', 'tq2_0'),
|
||||
]
|
||||
|
||||
# If specific types are specified, filter the list
|
||||
if args.types:
|
||||
types_lower = [t.lower() for t in args.types]
|
||||
types_to_quantize = [(et, os) for et, os in all_types if os.lower() in types_lower]
|
||||
if not types_to_quantize:
|
||||
print(f"❌ No valid types specified. Available types: {', '.join([os for _, os in all_types])}")
|
||||
return
|
||||
else:
|
||||
types_to_quantize = all_types
|
||||
|
||||
# If skip existing files is enabled, no need to filter
|
||||
# Because new logic will automatically detect and skip during quantization, but will still benchmark
|
||||
|
||||
# 创建量化器并运行
|
||||
try:
|
||||
quantizer = EmbeddingQuantizer(
|
||||
args.input,
|
||||
args.output_dir,
|
||||
args.quantize_bin,
|
||||
args.bench_bin,
|
||||
args.stats_dir,
|
||||
args.csv_output
|
||||
)
|
||||
quantizer.run_all_quantizations(types_to_quantize)
|
||||
except FileNotFoundError as e:
|
||||
print(f"❌ Error: {e}")
|
||||
return 1
|
||||
except KeyboardInterrupt:
|
||||
print("\n\n⚠️ Quantization interrupted by user")
|
||||
return 1
|
||||
except Exception as e:
|
||||
print(f"\n❌ Unexpected error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
exit(main() or 0)
|
||||
Executable
+573
@@ -0,0 +1,573 @@
|
||||
#!/bin/bash
|
||||
# Unified GEMM kernel benchmark script
|
||||
# Builds, tests, and benchmarks the GEMM kernel with configurable output
|
||||
|
||||
set -e
|
||||
|
||||
# Default values
|
||||
BUILD_DIR="../build"
|
||||
ITERATIONS=1000
|
||||
OUTPUT_CSV=""
|
||||
SKIP_BUILD=false
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
# Print usage
|
||||
print_usage() {
|
||||
cat << EOF
|
||||
Usage: $0 [options]
|
||||
|
||||
Options:
|
||||
-o, --output <path> Output CSV file path (default: ../stats/gemm_kernel_test_noparal.csv)
|
||||
-i, --iterations <num> Number of iterations per test (default: 1000)
|
||||
-s, --skip-build Skip building the benchmark binary
|
||||
-h, --help Show this help message
|
||||
|
||||
Examples:
|
||||
# Run with default settings
|
||||
$0
|
||||
|
||||
# Specify custom output file
|
||||
$0 -o /path/to/my_results.csv
|
||||
|
||||
# Quick test with fewer iterations
|
||||
$0 -i 100 -o quick_test.csv
|
||||
|
||||
# Skip build if already compiled
|
||||
$0 -s -o results.csv
|
||||
EOF
|
||||
}
|
||||
|
||||
# Parse command line arguments
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case $1 in
|
||||
-o|--output)
|
||||
OUTPUT_CSV="$2"
|
||||
shift 2
|
||||
;;
|
||||
-i|--iterations)
|
||||
ITERATIONS="$2"
|
||||
shift 2
|
||||
;;
|
||||
-s|--skip-build)
|
||||
SKIP_BUILD=true
|
||||
shift
|
||||
;;
|
||||
-h|--help)
|
||||
print_usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "Unknown option: $1"
|
||||
print_usage
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
# Set default output CSV if not specified
|
||||
if [ -z "$OUTPUT_CSV" ]; then
|
||||
OUTPUT_CSV="${SCRIPT_DIR}/../stats/gemm_kernel_test_noparal.csv"
|
||||
fi
|
||||
|
||||
# Create output directory first
|
||||
mkdir -p "$(dirname "$OUTPUT_CSV")"
|
||||
|
||||
# Convert to absolute path
|
||||
if [[ "$OUTPUT_CSV" = /* ]]; then
|
||||
# Already absolute path
|
||||
OUTPUT_CSV="$OUTPUT_CSV"
|
||||
else
|
||||
# Convert relative path to absolute
|
||||
OUTPUT_CSV="$(cd "$(dirname "$OUTPUT_CSV")" && pwd)/$(basename "$OUTPUT_CSV")"
|
||||
fi
|
||||
|
||||
echo "=========================================="
|
||||
echo "GEMM Kernel Benchmark Suite"
|
||||
echo "=========================================="
|
||||
echo "Configuration:"
|
||||
echo " Iterations: $ITERATIONS"
|
||||
echo " Output CSV: $OUTPUT_CSV"
|
||||
echo " Skip build: $SKIP_BUILD"
|
||||
echo "=========================================="
|
||||
echo ""
|
||||
|
||||
# Build the benchmark binary
|
||||
if [ "$SKIP_BUILD" = false ]; then
|
||||
echo "Step 1: Building GEMM kernel benchmark..."
|
||||
echo "------------------------------------------"
|
||||
|
||||
CXX=${CXX:-g++}
|
||||
|
||||
# Create build directory if it doesn't exist
|
||||
mkdir -p "${SCRIPT_DIR}/${BUILD_DIR}"
|
||||
|
||||
# Create temporary C++ source file
|
||||
TEMP_CPP="${SCRIPT_DIR}/${BUILD_DIR}/test_gemm_kernel_temp.cpp"
|
||||
|
||||
cat > "${TEMP_CPP}" << 'EOF'
|
||||
/**
|
||||
* Standalone benchmark for ggml_gemm_i2_i8_s kernel
|
||||
*
|
||||
* This program tests the performance of the ggml_gemm_i2_i8_s kernel
|
||||
* with configurable matrix sizes and iteration counts.
|
||||
*
|
||||
* Usage: ./test_gemm_kernel [options]
|
||||
* -n <size> : embedding dimension (must be divisible by 4, default: 2048)
|
||||
* -r <rows> : number of rows in matrix Y (default: 32)
|
||||
* -c <cols> : number of columns in matrix X (default: 128)
|
||||
* -i <iters> : number of iterations (default: 1000)
|
||||
* -w <warmup> : number of warmup iterations (default: 10)
|
||||
*/
|
||||
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
#include <time.h>
|
||||
#include <stdint.h>
|
||||
#include <math.h>
|
||||
#include <assert.h>
|
||||
|
||||
// Include necessary headers
|
||||
#include "../include/gemm-config.h"
|
||||
|
||||
// Function declarations (from ggml-quants.h)
|
||||
extern "C" void ggml_vec_dot_i2_i8_s(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc);
|
||||
|
||||
// GEMM kernel definition
|
||||
void ggml_gemm_i2_i8_s(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) {
|
||||
#if defined(ACT_PARALLEL)
|
||||
const int64_t row_block = ROW_BLOCK_SIZE;
|
||||
const int64_t col_block = COL_BLOCK_SIZE;
|
||||
|
||||
for (int64_t c0 = 0; c0 < nc; c0 += col_block) {
|
||||
int64_t cur_c = (c0 + col_block <= nc) ? col_block : (nc - c0);
|
||||
for (int64_t r0 = 0; r0 < nr; r0 += row_block) {
|
||||
int64_t cur_r = (r0 + row_block <= nr) ? row_block : (nr - r0);
|
||||
const void * vy_r = (const uint8_t *)vy + r0 * n;
|
||||
for (int64_t c = 0; c < cur_c; ++c) {
|
||||
const int64_t col = c0 + c;
|
||||
float * s_col = s + col;
|
||||
const void * vx_col = (const uint8_t *)vx + col * n / 4;
|
||||
ggml_vec_dot_i2_i8_s(n, s_col + r0 * bs, bs, vx_col, n, vy_r, n, cur_r);
|
||||
}
|
||||
}
|
||||
}
|
||||
#else
|
||||
const int64_t row_block = ROW_BLOCK_SIZE;
|
||||
const int64_t col_block = COL_BLOCK_SIZE;
|
||||
|
||||
for (int64_t r0 = 0; r0 < nr; r0 += row_block) {
|
||||
int64_t cur_r = (r0 + row_block <= nr) ? row_block : (nr - r0);
|
||||
for (int64_t c0 = 0; c0 < nc; c0 += col_block) {
|
||||
int64_t cur_c = (c0 + col_block <= nc) ? col_block : (nc - c0);
|
||||
const void * vx_c = (const uint8_t *)vx + c0 * n / 4;
|
||||
for (int64_t r = 0; r < cur_r; ++r) {
|
||||
const int64_t row = r0 + r;
|
||||
float * s_row = s + row * bs;
|
||||
const void * vy_row = (const uint8_t *)vy + row * n;
|
||||
ggml_vec_dot_i2_i8_s(n, s_row + c0, bs, vx_c, n, vy_row, n, cur_c);
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// Helper function to get current time in nanoseconds
|
||||
double get_time_ns() {
|
||||
struct timespec ts;
|
||||
clock_gettime(CLOCK_MONOTONIC, &ts);
|
||||
return ts.tv_sec * 1e9 + ts.tv_nsec;
|
||||
}
|
||||
|
||||
// Initialize matrix with random i2 values (2-bit quantized)
|
||||
void init_matrix_i2(uint8_t* data, int n, int cols) {
|
||||
// i2 format: 4 values per byte (2 bits each)
|
||||
int total_bytes = n * cols / 4;
|
||||
for (int i = 0; i < total_bytes; i++) {
|
||||
data[i] = rand() & 0xFF;
|
||||
}
|
||||
}
|
||||
|
||||
// Initialize matrix with random i8 values
|
||||
void init_matrix_i8(int8_t* data, int n, int rows) {
|
||||
int total_elements = n * rows;
|
||||
for (int i = 0; i < total_elements; i++) {
|
||||
data[i] = (int8_t)((rand() % 256) - 128);
|
||||
}
|
||||
}
|
||||
|
||||
// Benchmark configuration
|
||||
struct BenchmarkConfig {
|
||||
int n; // embedding dimension (must be divisible by 4)
|
||||
int nr; // number of rows in Y matrix
|
||||
int nc; // number of columns in X matrix
|
||||
int iterations; // number of benchmark iterations
|
||||
int warmup; // number of warmup iterations
|
||||
};
|
||||
|
||||
void print_config(const BenchmarkConfig& config) {
|
||||
printf("=" "=%.78s\n", "===============================================================================");
|
||||
printf("Benchmark Configuration:\n");
|
||||
printf("=" "=%.78s\n", "===============================================================================");
|
||||
printf(" Embedding dimension (n) : %d\n", config.n);
|
||||
printf(" Matrix Y rows (nr) : %d\n", config.nr);
|
||||
printf(" Matrix X columns (nc) : %d\n", config.nc);
|
||||
printf(" Iterations : %d\n", config.iterations);
|
||||
printf(" Warmup iterations : %d\n", config.warmup);
|
||||
printf("\nMatrix sizes:\n");
|
||||
printf(" X (i2): %d x %d (%.2f KB)\n", config.nc, config.n,
|
||||
(config.nc * config.n / 4) / 1024.0);
|
||||
printf(" Y (i8): %d x %d (%.2f KB)\n", config.nr, config.n,
|
||||
(config.nr * config.n) / 1024.0);
|
||||
printf(" S (f32): %d x %d (%.2f KB)\n", config.nr, config.nc,
|
||||
(config.nr * config.nc * sizeof(float)) / 1024.0);
|
||||
printf("\nGEMM Config:\n");
|
||||
#if defined(ACT_PARALLEL)
|
||||
printf(" ACT_PARALLEL : ON\n");
|
||||
#else
|
||||
printf(" ACT_PARALLEL : OFF\n");
|
||||
#endif
|
||||
printf(" ROW_BLOCK_SIZE : %d\n", ROW_BLOCK_SIZE);
|
||||
printf(" COL_BLOCK_SIZE : %d\n", COL_BLOCK_SIZE);
|
||||
printf(" PARALLEL_SIZE : %d\n", PARALLEL_SIZE);
|
||||
printf("=" "=%.78s\n\n", "===============================================================================");
|
||||
}
|
||||
|
||||
void run_benchmark(const BenchmarkConfig& config) {
|
||||
// Allocate matrices
|
||||
printf("Allocating matrices...\n");
|
||||
|
||||
// X matrix (i2 format): nc x n, but stored as nc x (n/4) bytes
|
||||
// Align to 64 bytes for AVX-512, which is backward compatible with AVX2 (32 bytes)
|
||||
size_t x_size = config.nc * config.n / 4;
|
||||
size_t x_size_aligned = ((x_size + 63) / 64) * 64;
|
||||
uint8_t* X = (uint8_t*)aligned_alloc(64, x_size_aligned);
|
||||
|
||||
// Y matrix (i8 format): nr x n
|
||||
size_t y_size = config.nr * config.n;
|
||||
size_t y_size_aligned = ((y_size + 63) / 64) * 64;
|
||||
int8_t* Y = (int8_t*)aligned_alloc(64, y_size_aligned);
|
||||
|
||||
// Result matrix (float32): nr x nc
|
||||
size_t s_size = config.nr * config.nc * sizeof(float);
|
||||
size_t s_size_aligned = ((s_size + 63) / 64) * 64;
|
||||
float* S = (float*)aligned_alloc(64, s_size_aligned);
|
||||
|
||||
if (!X || !Y || !S) {
|
||||
fprintf(stderr, "Failed to allocate memory\n");
|
||||
exit(1);
|
||||
}
|
||||
|
||||
// Initialize matrices with random data
|
||||
printf("Initializing matrices with random data...\n");
|
||||
srand(time(NULL));
|
||||
init_matrix_i2(X, config.n, config.nc);
|
||||
init_matrix_i8(Y, config.n, config.nr);
|
||||
memset(S, 0, config.nr * config.nc * sizeof(float));
|
||||
|
||||
// Warmup
|
||||
printf("Running %d warmup iterations...\n", config.warmup);
|
||||
for (int i = 0; i < config.warmup; i++) {
|
||||
ggml_gemm_i2_i8_s(config.n, S, config.nc, X, Y, config.nr, config.nc);
|
||||
}
|
||||
|
||||
// Benchmark
|
||||
printf("Running %d benchmark iterations...\n", config.iterations);
|
||||
double total_time = 0.0;
|
||||
double min_time = 1e20;
|
||||
double max_time = 0.0;
|
||||
|
||||
for (int i = 0; i < config.iterations; i++) {
|
||||
double start = get_time_ns();
|
||||
ggml_gemm_i2_i8_s(config.n, S, config.nc, X, Y, config.nr, config.nc);
|
||||
double end = get_time_ns();
|
||||
|
||||
double elapsed = end - start;
|
||||
total_time += elapsed;
|
||||
if (elapsed < min_time) min_time = elapsed;
|
||||
if (elapsed > max_time) max_time = elapsed;
|
||||
|
||||
if ((i + 1) % 100 == 0) {
|
||||
printf(" Progress: %d/%d iterations\n", i + 1, config.iterations);
|
||||
}
|
||||
}
|
||||
|
||||
// Calculate statistics
|
||||
double avg_time_ns = total_time / config.iterations;
|
||||
double avg_time_ms = avg_time_ns / 1e6;
|
||||
double min_time_ms = min_time / 1e6;
|
||||
double max_time_ms = max_time / 1e6;
|
||||
|
||||
// Calculate GFLOPS
|
||||
// For GEMM: nr x nc x n multiply-adds = 2 * nr * nc * n FLOPs
|
||||
double flops = 2.0 * config.nr * config.nc * config.n;
|
||||
double gflops = (flops / avg_time_ns);
|
||||
|
||||
// Calculate throughput (tokens/s assuming each column is a token)
|
||||
double throughput = (config.nc * 1e9) / avg_time_ns;
|
||||
|
||||
// Print results
|
||||
printf("\n");
|
||||
printf("=" "=%.78s\n", "===============================================================================");
|
||||
printf("Benchmark Results:\n");
|
||||
printf("=" "=%.78s\n", "===============================================================================");
|
||||
printf(" Average time : %.3f ms\n", avg_time_ms);
|
||||
printf(" Min time : %.3f ms\n", min_time_ms);
|
||||
printf(" Max time : %.3f ms\n", max_time_ms);
|
||||
printf(" Std dev : %.3f ms\n", sqrt((max_time_ms - min_time_ms) * (max_time_ms - min_time_ms) / 12));
|
||||
printf("\nPerformance:\n");
|
||||
printf(" GFLOPS : %.2f\n", gflops);
|
||||
printf(" Throughput : %.2f tokens/s\n", throughput);
|
||||
printf(" Latency/token : %.3f us\n", (avg_time_ms * 1000) / config.nc);
|
||||
printf("=" "=%.78s\n", "===============================================================================");
|
||||
|
||||
// Cleanup
|
||||
free(X);
|
||||
free(Y);
|
||||
free(S);
|
||||
}
|
||||
|
||||
void print_usage(const char* program) {
|
||||
printf("Usage: %s [options]\n", program);
|
||||
printf("Options:\n");
|
||||
printf(" -n <size> Embedding dimension (must be divisible by 4, default: 2048)\n");
|
||||
printf(" -r <rows> Number of rows in matrix Y (default: 32)\n");
|
||||
printf(" -c <cols> Number of columns in matrix X (default: 128)\n");
|
||||
printf(" -i <iters> Number of iterations (default: 1000)\n");
|
||||
printf(" -w <warmup> Number of warmup iterations (default: 10)\n");
|
||||
printf(" -h Show this help message\n");
|
||||
}
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
BenchmarkConfig config = {
|
||||
.n = 2048,
|
||||
.nr = 32,
|
||||
.nc = 128,
|
||||
.iterations = 1000,
|
||||
.warmup = 10
|
||||
};
|
||||
|
||||
// Parse command line arguments
|
||||
for (int i = 1; i < argc; i++) {
|
||||
if (strcmp(argv[i], "-n") == 0 && i + 1 < argc) {
|
||||
config.n = atoi(argv[++i]);
|
||||
} else if (strcmp(argv[i], "-r") == 0 && i + 1 < argc) {
|
||||
config.nr = atoi(argv[++i]);
|
||||
} else if (strcmp(argv[i], "-c") == 0 && i + 1 < argc) {
|
||||
config.nc = atoi(argv[++i]);
|
||||
} else if (strcmp(argv[i], "-i") == 0 && i + 1 < argc) {
|
||||
config.iterations = atoi(argv[++i]);
|
||||
} else if (strcmp(argv[i], "-w") == 0 && i + 1 < argc) {
|
||||
config.warmup = atoi(argv[++i]);
|
||||
} else if (strcmp(argv[i], "-h") == 0) {
|
||||
print_usage(argv[0]);
|
||||
return 0;
|
||||
} else {
|
||||
fprintf(stderr, "Unknown option: %s\n", argv[i]);
|
||||
print_usage(argv[0]);
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
// Validate configuration
|
||||
if (config.n % 4 != 0) {
|
||||
fprintf(stderr, "Error: Embedding dimension (-n) must be divisible by 4\n");
|
||||
return 1;
|
||||
}
|
||||
|
||||
if (config.n <= 0 || config.nr <= 0 || config.nc <= 0 || config.iterations <= 0) {
|
||||
fprintf(stderr, "Error: All size parameters must be positive\n");
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Run benchmark
|
||||
print_config(config);
|
||||
run_benchmark(config);
|
||||
|
||||
return 0;
|
||||
}
|
||||
EOF
|
||||
|
||||
# Compiler flags
|
||||
CXXFLAGS="-O3 -march=native -mtune=native -std=c++17 -fopenmp"
|
||||
CXXFLAGS+=" -I${SCRIPT_DIR}/.. -I${SCRIPT_DIR}/../include"
|
||||
CXXFLAGS+=" -I${SCRIPT_DIR}/../3rdparty/llama.cpp/ggml/include"
|
||||
CXXFLAGS+=" -I${SCRIPT_DIR}/../3rdparty/llama.cpp/ggml/src"
|
||||
CXXFLAGS+=" -I${SCRIPT_DIR}/../3rdparty/llama.cpp/include"
|
||||
CXXFLAGS+=" -DNDEBUG -ffast-math"
|
||||
|
||||
# Link flags
|
||||
LDFLAGS="-lm -lpthread"
|
||||
|
||||
# Link with pre-built libraries
|
||||
GGML_LIB_DIR="${SCRIPT_DIR}/../build/3rdparty/llama.cpp/ggml/src"
|
||||
GGML_SO="${GGML_LIB_DIR}/libggml.so"
|
||||
|
||||
if [ ! -f "${GGML_SO}" ]; then
|
||||
echo "❌ Error: Cannot find libggml.so at ${GGML_SO}"
|
||||
echo "Please build the project first with: cmake --build build"
|
||||
rm -f "${TEMP_CPP}"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
LDFLAGS+=" -L${GGML_LIB_DIR} -lggml -Wl,-rpath,${GGML_LIB_DIR}"
|
||||
|
||||
# Output binary
|
||||
BENCHMARK_BIN="${SCRIPT_DIR}/${BUILD_DIR}/test_gemm_kernel"
|
||||
|
||||
echo "Compiler: ${CXX}"
|
||||
echo "Building from embedded source..."
|
||||
echo ""
|
||||
|
||||
# Build
|
||||
${CXX} ${CXXFLAGS} "${TEMP_CPP}" -o ${BENCHMARK_BIN} ${LDFLAGS}
|
||||
|
||||
if [ $? -eq 0 ]; then
|
||||
echo "✅ Build successful!"
|
||||
rm -f "${TEMP_CPP}"
|
||||
echo ""
|
||||
else
|
||||
echo "❌ Build failed!"
|
||||
rm -f "${TEMP_CPP}"
|
||||
exit 1
|
||||
fi
|
||||
else
|
||||
echo "Step 1: Skipping build (using existing binary)"
|
||||
echo "------------------------------------------"
|
||||
BENCHMARK_BIN="${SCRIPT_DIR}/${BUILD_DIR}/test_gemm_kernel"
|
||||
|
||||
if [ ! -f "${BENCHMARK_BIN}" ]; then
|
||||
echo "❌ Error: Benchmark binary not found at ${BENCHMARK_BIN}"
|
||||
echo "Please run without -s to build it first."
|
||||
exit 1
|
||||
fi
|
||||
echo "✅ Found existing binary"
|
||||
echo ""
|
||||
fi
|
||||
|
||||
# Set LD_LIBRARY_PATH to include the GGML library directory
|
||||
GGML_LIB_DIR="${SCRIPT_DIR}/../build/3rdparty/llama.cpp/ggml/src"
|
||||
export LD_LIBRARY_PATH="${GGML_LIB_DIR}:${LD_LIBRARY_PATH}"
|
||||
|
||||
echo "Step 2: Running benchmark tests"
|
||||
echo "------------------------------------------"
|
||||
echo "Library path: ${GGML_LIB_DIR}"
|
||||
echo ""
|
||||
|
||||
# Write CSV header
|
||||
echo "test_name,n,nr,nc,time_ms,gflops,throughput_tokens_per_sec" > "$OUTPUT_CSV"
|
||||
echo "Results will be saved to: $OUTPUT_CSV"
|
||||
echo ""
|
||||
|
||||
# Function to extract metrics and append to CSV
|
||||
extract_and_save() {
|
||||
local test_name="$1"
|
||||
local output="$2"
|
||||
|
||||
# Extract values using grep and awk
|
||||
local n=$(echo "$output" | grep "Embedding dimension" | awk '{print $5}')
|
||||
local nr=$(echo "$output" | grep "Matrix Y rows" | awk '{print $6}')
|
||||
local nc=$(echo "$output" | grep "Matrix X columns" | awk '{print $6}')
|
||||
local avg_time=$(echo "$output" | grep "Average time" | awk '{print $4}')
|
||||
local min_time=$(echo "$output" | grep "Min time" | awk '{print $4}')
|
||||
local max_time=$(echo "$output" | grep "Max time" | awk '{print $4}')
|
||||
local gflops=$(echo "$output" | grep "GFLOPS" | awk '{print $3}')
|
||||
local throughput=$(echo "$output" | grep "Throughput" | awk '{print $3}')
|
||||
|
||||
# Check if values were extracted successfully
|
||||
if [ -z "$avg_time" ] || [ -z "$min_time" ] || [ -z "$max_time" ]; then
|
||||
echo "Warning: Failed to extract timing data for ${test_name}"
|
||||
echo "${test_name},${n},${nr},${nc},N/A,N/A,N/A" >> "$OUTPUT_CSV"
|
||||
return
|
||||
fi
|
||||
|
||||
# Calculate standard deviation estimate from range
|
||||
# Using awk with proper variable passing
|
||||
local std_time=$(awk -v min="$min_time" -v max="$max_time" 'BEGIN {printf "%.4f", (max - min) / 4}')
|
||||
|
||||
# Format as mean±std
|
||||
local time_formatted="${avg_time}±${std_time}"
|
||||
|
||||
# Append to CSV
|
||||
echo "${test_name},${n},${nr},${nc},${time_formatted},${gflops},${throughput}" >> "$OUTPUT_CSV"
|
||||
}
|
||||
|
||||
# Run benchmark tests
|
||||
echo "=========================================="
|
||||
echo "BitNet-2B Typical Shapes Performance Test"
|
||||
echo "=========================================="
|
||||
echo ""
|
||||
|
||||
echo "Test 1: Single Token Generation (Attention QKV projection)"
|
||||
echo " Scenario: Generating 1 token at a time"
|
||||
echo " Shape: n=2048, r=1, c=2048"
|
||||
OUTPUT=$($BENCHMARK_BIN -n 2048 -r 1 -c 2048 -i $ITERATIONS 2>&1)
|
||||
echo "$OUTPUT"
|
||||
extract_and_save "single_token_gen" "$OUTPUT"
|
||||
echo ""
|
||||
|
||||
echo "Test 2: Small Batch Prompt Processing (Attention QKV projection)"
|
||||
echo " Scenario: Processing prompt with 128 tokens, batch size 1"
|
||||
echo " Shape: n=2048, r=128, c=2048"
|
||||
OUTPUT=$($BENCHMARK_BIN -n 2048 -r 128 -c 2048 -i $ITERATIONS 2>&1)
|
||||
echo "$OUTPUT"
|
||||
extract_and_save "small_batch_prompt" "$OUTPUT"
|
||||
echo ""
|
||||
|
||||
echo "Test 3: Medium Batch Prompt Processing (Attention QKV projection)"
|
||||
echo " Scenario: Processing prompt with 256 tokens or batch of 256"
|
||||
echo " Shape: n=2048, r=256, c=2048"
|
||||
OUTPUT=$($BENCHMARK_BIN -n 2048 -r 256 -c 2048 -i $ITERATIONS 2>&1)
|
||||
echo "$OUTPUT"
|
||||
extract_and_save "medium_batch_prompt" "$OUTPUT"
|
||||
echo ""
|
||||
|
||||
echo "Test 4: Large Batch Processing (Attention QKV projection)"
|
||||
echo " Scenario: Processing 512 tokens or batch of 512"
|
||||
echo " Shape: n=2048, r=512, c=2048"
|
||||
OUTPUT=$($BENCHMARK_BIN -n 2048 -r 512 -c 2048 -i $ITERATIONS 2>&1)
|
||||
echo "$OUTPUT"
|
||||
extract_and_save "large_batch_prompt" "$OUTPUT"
|
||||
echo ""
|
||||
|
||||
echo "Test 5: FFN Up-projection (Small batch)"
|
||||
echo " Scenario: Feed-forward network expansion, 128 tokens"
|
||||
echo " Shape: n=2048, r=128, c=8192"
|
||||
OUTPUT=$($BENCHMARK_BIN -n 2048 -r 128 -c 8192 -i $ITERATIONS 2>&1)
|
||||
echo "$OUTPUT"
|
||||
extract_and_save "ffn_up_projection" "$OUTPUT"
|
||||
echo ""
|
||||
|
||||
echo "Test 6: FFN Down-projection (Small batch)"
|
||||
echo " Scenario: Feed-forward network reduction, 128 tokens"
|
||||
echo " Shape: n=8192, r=128, c=2048"
|
||||
OUTPUT=$($BENCHMARK_BIN -n 8192 -r 128 -c 2048 -i $ITERATIONS 2>&1)
|
||||
echo "$OUTPUT"
|
||||
extract_and_save "ffn_down_projection" "$OUTPUT"
|
||||
echo ""
|
||||
|
||||
echo "Test 7: Long Context Processing"
|
||||
echo " Scenario: Processing very long context (2048 tokens)"
|
||||
echo " Shape: n=2048, r=2048, c=2048"
|
||||
OUTPUT=$($BENCHMARK_BIN -n 2048 -r 2048 -c 2048 -i $ITERATIONS 2>&1)
|
||||
echo "$OUTPUT"
|
||||
extract_and_save "long_context" "$OUTPUT"
|
||||
echo ""
|
||||
|
||||
echo "Test 8: Batched Token Generation"
|
||||
echo " Scenario: Generating tokens for 32 sequences simultaneously"
|
||||
echo " Shape: n=2048, r=32, c=2048"
|
||||
OUTPUT=$($BENCHMARK_BIN -n 2048 -r 32 -c 2048 -i $ITERATIONS 2>&1)
|
||||
echo "$OUTPUT"
|
||||
extract_and_save "batched_token_gen" "$OUTPUT"
|
||||
echo ""
|
||||
|
||||
echo "=========================================="
|
||||
echo "All tests completed successfully!"
|
||||
echo "=========================================="
|
||||
echo "Results saved to: $OUTPUT_CSV"
|
||||
echo ""
|
||||
echo "Summary:"
|
||||
wc -l "$OUTPUT_CSV" | awk '{print " Total records:", $1 - 1}'
|
||||
echo " Output file: $OUTPUT_CSV"
|
||||
echo "=========================================="
|
||||
@@ -0,0 +1,608 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Perplexity Test Script
|
||||
Tests GGUF model perplexity on multiple datasets using llama-perplexity.
|
||||
"""
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
import csv
|
||||
import re
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
import argparse
|
||||
import tempfile
|
||||
import shutil
|
||||
import statistics
|
||||
|
||||
|
||||
class PerplexityTester:
|
||||
def __init__(self, model_path, llama_perplexity_bin="../build/bin/llama-perplexity",
|
||||
data_dir="../data", output_dir="perplexity_results", quick_mode=False,
|
||||
quantize_bin="../build/bin/llama-quantize", test_embeddings=False, csv_output=None):
|
||||
self.model_path = Path(model_path)
|
||||
self.llama_perplexity_bin = Path(llama_perplexity_bin)
|
||||
self.quantize_bin = Path(quantize_bin)
|
||||
self.data_dir = Path(data_dir)
|
||||
self.output_dir = Path(output_dir)
|
||||
self.quick_mode = quick_mode
|
||||
self.test_embeddings = test_embeddings
|
||||
self.csv_output = Path(csv_output) if csv_output else None
|
||||
self.results = []
|
||||
self.created_models = set() # Track newly created model files
|
||||
self.temp_files = [] # Track temporary files for cleanup
|
||||
|
||||
# Embedding types to test
|
||||
self.embedding_types = [
|
||||
('F32', 'f32'),
|
||||
('F16', 'f16'),
|
||||
('Q8_0', 'q8_0'),
|
||||
('Q6_K', 'q6_k'),
|
||||
('Q5_0', 'q5_0'),
|
||||
('Q4_0', 'q4_0'),
|
||||
('Q3_K', 'q3_k'),
|
||||
('TQ2_0', 'tq2_0'),
|
||||
]
|
||||
|
||||
# Create output directory
|
||||
self.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Verify llama-perplexity binary exists
|
||||
if not self.llama_perplexity_bin.exists():
|
||||
raise FileNotFoundError(f"llama-perplexity binary not found: {self.llama_perplexity_bin}")
|
||||
|
||||
# Verify quantize binary exists if testing embeddings
|
||||
if self.test_embeddings and not self.quantize_bin.exists():
|
||||
raise FileNotFoundError(f"llama-quantize binary not found: {self.quantize_bin}")
|
||||
|
||||
# Verify model file exists
|
||||
if not self.model_path.exists():
|
||||
raise FileNotFoundError(f"Model file not found: {self.model_path}")
|
||||
|
||||
def find_datasets(self):
|
||||
"""Find all test.txt files in dataset directories."""
|
||||
datasets = []
|
||||
|
||||
if not self.data_dir.exists():
|
||||
print(f"❌ Data directory not found: {self.data_dir}")
|
||||
return datasets
|
||||
|
||||
print(f"\n🔍 Searching for datasets in {self.data_dir}...")
|
||||
|
||||
# Look for test.txt files in subdirectories
|
||||
for dataset_dir in sorted(self.data_dir.iterdir()):
|
||||
if dataset_dir.is_dir():
|
||||
test_file = dataset_dir / "test.txt"
|
||||
if test_file.exists():
|
||||
size_mb = test_file.stat().st_size / (1024 * 1024)
|
||||
datasets.append({
|
||||
'name': dataset_dir.name,
|
||||
'path': test_file,
|
||||
'size': test_file.stat().st_size,
|
||||
'size_mb': size_mb
|
||||
})
|
||||
print(f" ✅ {dataset_dir.name:<20} ({size_mb:.2f} MB)")
|
||||
else:
|
||||
print(f" ⚠️ {dataset_dir.name:<20} (no test.txt found)")
|
||||
|
||||
return datasets
|
||||
|
||||
def create_quick_dataset(self, dataset_path, num_chars=4096):
|
||||
"""Create a temporary dataset with only the first N characters for quick testing."""
|
||||
temp_file = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt', encoding='utf-8')
|
||||
self.temp_files.append(temp_file.name)
|
||||
|
||||
try:
|
||||
with open(dataset_path, 'r', encoding='utf-8', errors='ignore') as f:
|
||||
content = f.read(num_chars)
|
||||
temp_file.write(content)
|
||||
temp_file.close()
|
||||
return Path(temp_file.name)
|
||||
except Exception as e:
|
||||
print(f"⚠️ Failed to create quick dataset: {e}")
|
||||
temp_file.close()
|
||||
return dataset_path
|
||||
|
||||
def cleanup_temp_files(self):
|
||||
"""Clean up temporary files."""
|
||||
for temp_file in self.temp_files:
|
||||
try:
|
||||
os.unlink(temp_file)
|
||||
except:
|
||||
pass
|
||||
self.temp_files = []
|
||||
|
||||
def run_perplexity_test(self, dataset_name, dataset_path, threads=16, ctx_size=512, model_override=None):
|
||||
"""Run perplexity test on a single dataset."""
|
||||
test_model = model_override if model_override else self.model_path
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print(f"📊 Testing on dataset: {dataset_name}")
|
||||
print(f" File: {dataset_path}")
|
||||
print(f" Model: {test_model.name}")
|
||||
print(f"{'='*80}")
|
||||
|
||||
cmd = [
|
||||
str(self.llama_perplexity_bin),
|
||||
"-m", str(test_model),
|
||||
"-f", str(dataset_path),
|
||||
"-t", str(threads),
|
||||
"-c", str(ctx_size),
|
||||
"-ngl", "0" # CPU only
|
||||
]
|
||||
|
||||
print(f"💻 Command: {' '.join(cmd)}")
|
||||
print(f"⏱️ Starting test...\n")
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=3600, # 1 hour timeout
|
||||
cwd=os.getcwd()
|
||||
)
|
||||
|
||||
elapsed_time = time.time() - start_time
|
||||
|
||||
if result.returncode == 0:
|
||||
# Parse perplexity from output (check both stdout and stderr)
|
||||
combined_output = result.stdout + "\n" + result.stderr
|
||||
ppl = self.parse_perplexity(combined_output)
|
||||
|
||||
if ppl is not None:
|
||||
print(f"\n✅ Perplexity: {ppl}")
|
||||
print(f"⏱️ Time: {elapsed_time:.2f}s ({elapsed_time/60:.2f} min)")
|
||||
status = "success"
|
||||
else:
|
||||
print(f"\n⚠️ Test completed but could not parse perplexity")
|
||||
print(f"Last 500 chars of stdout:")
|
||||
print(result.stdout[-500:])
|
||||
print(f"Last 500 chars of stderr:")
|
||||
print(result.stderr[-500:])
|
||||
status = "parse_error"
|
||||
ppl = None
|
||||
else:
|
||||
print(f"\n❌ Test failed with return code {result.returncode}")
|
||||
print(f"Error: {result.stderr[:500]}")
|
||||
status = "failed"
|
||||
ppl = None
|
||||
elapsed_time = time.time() - start_time
|
||||
|
||||
return {
|
||||
'dataset': dataset_name,
|
||||
'perplexity': ppl,
|
||||
'time': elapsed_time,
|
||||
'status': status,
|
||||
'stdout': result.stdout,
|
||||
'stderr': result.stderr
|
||||
}
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
elapsed_time = time.time() - start_time
|
||||
print(f"\n❌ Timeout after {elapsed_time:.2f}s")
|
||||
return {
|
||||
'dataset': dataset_name,
|
||||
'perplexity': None,
|
||||
'time': elapsed_time,
|
||||
'status': 'timeout',
|
||||
'stdout': '',
|
||||
'stderr': 'Test exceeded 1 hour timeout'
|
||||
}
|
||||
except Exception as e:
|
||||
elapsed_time = time.time() - start_time
|
||||
print(f"\n❌ Error: {e}")
|
||||
return {
|
||||
'dataset': dataset_name,
|
||||
'perplexity': None,
|
||||
'time': elapsed_time,
|
||||
'status': 'error',
|
||||
'stdout': '',
|
||||
'stderr': str(e)
|
||||
}
|
||||
|
||||
def parse_perplexity(self, output):
|
||||
"""Parse perplexity value (mean±std format) from llama-perplexity output."""
|
||||
# First try to match "PPL = mean +/- std" format
|
||||
pattern_with_std = r'PPL\s*=\s*(\d+\.?\d*)\s*\+/-\s*(\d+\.?\d*)'
|
||||
match = re.search(pattern_with_std, output, re.IGNORECASE | re.MULTILINE)
|
||||
if match:
|
||||
try:
|
||||
mean = float(match.group(1))
|
||||
std = float(match.group(2))
|
||||
return f"{mean:.4f}±{std:.4f}"
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# Fallback to patterns without std
|
||||
patterns = [
|
||||
r'Final estimate:\s*PPL\s*=\s*(\d+\.?\d*)',
|
||||
r'Final perplexity:\s*(\d+\.?\d*)',
|
||||
r'PPL\s*=\s*(\d+\.?\d*)',
|
||||
r'PPL:\s*(\d+\.?\d*)',
|
||||
r'perplexity:\s*(\d+\.?\d*)',
|
||||
r'ppl\s*=\s*(\d+\.?\d*)',
|
||||
r'Perplexity:\s*(\d+\.?\d*)',
|
||||
]
|
||||
|
||||
for pattern in patterns:
|
||||
match = re.search(pattern, output, re.IGNORECASE | re.MULTILINE)
|
||||
if match:
|
||||
try:
|
||||
return f"{float(match.group(1)):.4f}"
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
return None
|
||||
|
||||
def quantize_embedding(self, embedding_type, output_suffix):
|
||||
"""
|
||||
Quantize model with specific embedding type.
|
||||
|
||||
Args:
|
||||
embedding_type: Token embedding type (uppercase, e.g., 'Q6_K')
|
||||
output_suffix: Output file suffix (lowercase, e.g., 'q6_k')
|
||||
|
||||
Returns:
|
||||
Path to quantized model or None if failed
|
||||
"""
|
||||
# Construct output path
|
||||
model_dir = self.model_path.parent
|
||||
output_path = model_dir / f"ggml-model-i2_s-embed-{output_suffix}.gguf"
|
||||
|
||||
# Check if file already exists
|
||||
file_existed = output_path.exists()
|
||||
|
||||
if file_existed:
|
||||
print(f"ℹ️ Model already exists: {output_path.name}")
|
||||
return output_path
|
||||
|
||||
cmd = [
|
||||
str(self.quantize_bin),
|
||||
"--token-embedding-type", embedding_type,
|
||||
str(self.model_path),
|
||||
str(output_path),
|
||||
"I2_S",
|
||||
"1",
|
||||
"1"
|
||||
]
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print(f"🔄 Quantizing with embedding type: {embedding_type}")
|
||||
print(f"📥 Input: {self.model_path.name}")
|
||||
print(f"📤 Output: {output_path.name}")
|
||||
print(f"💻 Command: {' '.join(cmd)}")
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
cwd=os.getcwd(),
|
||||
timeout=600 # 10 minutes timeout
|
||||
)
|
||||
|
||||
duration = time.time() - start_time
|
||||
|
||||
if result.returncode == 0:
|
||||
file_size_mb = output_path.stat().st_size / (1024 * 1024)
|
||||
print(f"✅ Quantization successful!")
|
||||
print(f" Duration: {duration:.2f}s")
|
||||
print(f" Size: {file_size_mb:.2f} MB")
|
||||
|
||||
# Mark as newly created
|
||||
self.created_models.add(output_path)
|
||||
return output_path
|
||||
else:
|
||||
print(f"❌ Quantization failed with return code {result.returncode}")
|
||||
print(f"Error: {result.stderr[:500]}")
|
||||
return None
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
print(f"❌ Quantization timeout (exceeded 10 minutes)")
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"❌ Quantization error: {e}")
|
||||
return None
|
||||
|
||||
def cleanup_model(self, model_path):
|
||||
"""Delete model file if it was created during this session."""
|
||||
if model_path in self.created_models:
|
||||
try:
|
||||
model_path.unlink()
|
||||
print(f"🗑️ Deleted: {model_path.name}")
|
||||
self.created_models.remove(model_path)
|
||||
except Exception as e:
|
||||
print(f"⚠️ Failed to delete {model_path.name}: {e}")
|
||||
else:
|
||||
print(f"ℹ️ Keeping existing file: {model_path.name}")
|
||||
|
||||
def run_all_tests(self, threads=16, ctx_size=512):
|
||||
"""Run perplexity tests on all datasets."""
|
||||
datasets = self.find_datasets()
|
||||
|
||||
if not datasets:
|
||||
print(f"\n❌ No datasets found in {self.data_dir}")
|
||||
print(f" Make sure each dataset directory has a test.txt file")
|
||||
return
|
||||
|
||||
# Quick mode: test all datasets but only first 4096 chars with smaller context
|
||||
if self.quick_mode:
|
||||
ctx_size = min(ctx_size, 128) # Use smaller context in quick mode
|
||||
print(f"\n⚡ QUICK TEST MODE ENABLED")
|
||||
print(f" - Testing all datasets with first 4096 characters only")
|
||||
print(f" - Using reduced context size: {ctx_size}")
|
||||
|
||||
# Determine models to test
|
||||
if self.test_embeddings:
|
||||
print(f"\n{'='*80}")
|
||||
print(f"🧪 EMBEDDING QUANTIZATION TEST MODE")
|
||||
print(f"{'='*80}")
|
||||
print(f"📦 Base model: {self.model_path.name}")
|
||||
print(f"🔢 Embedding types to test: {len(self.embedding_types)}")
|
||||
print(f"📊 Datasets: {len(datasets)}")
|
||||
print(f"🧵 Threads: {threads}")
|
||||
print(f"📏 Context size: {ctx_size}")
|
||||
print(f"{'='*80}")
|
||||
|
||||
total_start = time.time()
|
||||
|
||||
# Test each embedding type
|
||||
for i, (embedding_type, output_suffix) in enumerate(self.embedding_types, 1):
|
||||
print(f"\n\n{'#'*80}")
|
||||
print(f"[{i}/{len(self.embedding_types)}] Testing embedding type: {output_suffix} ({embedding_type})")
|
||||
print(f"{'#'*80}")
|
||||
|
||||
# Quantize model
|
||||
quantized_model = self.quantize_embedding(embedding_type, output_suffix)
|
||||
|
||||
if quantized_model is None:
|
||||
print(f"⚠️ Skipping tests for {output_suffix} due to quantization failure")
|
||||
continue
|
||||
|
||||
# Test on all datasets
|
||||
for j, dataset in enumerate(datasets, 1):
|
||||
print(f"\n[{j}/{len(datasets)}] Testing {dataset['name']} with {output_suffix}...")
|
||||
|
||||
# Use quick dataset if in quick mode
|
||||
test_path = dataset['path']
|
||||
if self.quick_mode:
|
||||
test_path = self.create_quick_dataset(dataset['path'])
|
||||
|
||||
result = self.run_perplexity_test(
|
||||
f"{dataset['name']}_embed-{output_suffix}",
|
||||
test_path,
|
||||
threads,
|
||||
ctx_size,
|
||||
model_override=quantized_model
|
||||
)
|
||||
self.results.append(result)
|
||||
|
||||
# Cleanup model after testing
|
||||
print(f"\n🧹 Cleaning up {output_suffix} model...")
|
||||
self.cleanup_model(quantized_model)
|
||||
|
||||
print(f"\n{'#'*80}")
|
||||
print(f"✅ Completed {output_suffix}")
|
||||
print(f"{'#'*80}")
|
||||
|
||||
total_time = time.time() - total_start
|
||||
|
||||
else:
|
||||
# Regular single model test
|
||||
print(f"\n{'='*80}")
|
||||
print(f"🚀 PERPLEXITY TEST SESSION{' (QUICK MODE)' if self.quick_mode else ''}")
|
||||
print(f"{'='*80}")
|
||||
print(f"📦 Model: {self.model_path.name}")
|
||||
print(f"📁 Model path: {self.model_path}")
|
||||
print(f"📊 Datasets {'to test' if self.quick_mode else 'found'}: {len(datasets)}")
|
||||
print(f"🧵 Threads: {threads}")
|
||||
print(f"📏 Context size: {ctx_size}")
|
||||
print(f"{'='*80}")
|
||||
|
||||
total_start = time.time()
|
||||
|
||||
# Run tests
|
||||
for i, dataset in enumerate(datasets, 1):
|
||||
print(f"\n\n[{i}/{len(datasets)}] Processing {dataset['name']}...")
|
||||
|
||||
# Use quick dataset if in quick mode
|
||||
test_path = dataset['path']
|
||||
if self.quick_mode:
|
||||
test_path = self.create_quick_dataset(dataset['path'])
|
||||
|
||||
result = self.run_perplexity_test(
|
||||
dataset['name'],
|
||||
test_path,
|
||||
threads,
|
||||
ctx_size
|
||||
)
|
||||
self.results.append(result)
|
||||
|
||||
total_time = time.time() - total_start
|
||||
|
||||
# Clean up temporary files
|
||||
if self.quick_mode:
|
||||
print(f"\n🧹 Cleaning up temporary files...")
|
||||
self.cleanup_temp_files()
|
||||
|
||||
# Save results
|
||||
self.save_results()
|
||||
|
||||
# Print summary
|
||||
self.print_summary(total_time)
|
||||
|
||||
def save_results(self):
|
||||
"""Save results to CSV file."""
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
model_name = self.model_path.stem
|
||||
|
||||
# Use custom CSV path if provided
|
||||
if self.csv_output:
|
||||
csv_file = self.csv_output
|
||||
# Create parent directory if needed
|
||||
csv_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
else:
|
||||
csv_file = self.output_dir / f"ppl_{model_name}_{timestamp}.csv"
|
||||
|
||||
print(f"\n💾 Saving results...")
|
||||
|
||||
with open(csv_file, 'w', newline='') as f:
|
||||
writer = csv.DictWriter(f, fieldnames=['dataset', 'perplexity', 'time_seconds', 'status'])
|
||||
writer.writeheader()
|
||||
for result in self.results:
|
||||
writer.writerow({
|
||||
'dataset': result['dataset'],
|
||||
'perplexity': result['perplexity'] if result['perplexity'] is not None else 'N/A',
|
||||
'time_seconds': f"{result['time']:.2f}",
|
||||
'status': result['status']
|
||||
})
|
||||
|
||||
print(f" ✅ CSV saved: {csv_file}")
|
||||
|
||||
# Save detailed log
|
||||
log_file = self.output_dir / f"ppl_{model_name}_{timestamp}.log"
|
||||
with open(log_file, 'w') as f:
|
||||
f.write(f"Perplexity Test Results\n")
|
||||
f.write(f"{'='*80}\n")
|
||||
f.write(f"Model: {self.model_path}\n")
|
||||
f.write(f"Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
|
||||
f.write(f"{'='*80}\n\n")
|
||||
|
||||
for result in self.results:
|
||||
f.write(f"\n{'='*80}\n")
|
||||
f.write(f"Dataset: {result['dataset']}\n")
|
||||
f.write(f"Perplexity: {result['perplexity']}\n")
|
||||
f.write(f"Time: {result['time']:.2f}s\n")
|
||||
f.write(f"Status: {result['status']}\n")
|
||||
f.write(f"\nOutput:\n{result['stdout']}\n")
|
||||
if result['stderr']:
|
||||
f.write(f"\nErrors:\n{result['stderr']}\n")
|
||||
|
||||
print(f" ✅ Log saved: {log_file}")
|
||||
|
||||
def print_summary(self, total_time):
|
||||
"""Print summary of all tests."""
|
||||
print(f"\n\n{'='*80}")
|
||||
print(f"📊 TEST SUMMARY")
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
# Sort results by perplexity (lower is better)
|
||||
successful = [r for r in self.results if r['perplexity'] is not None]
|
||||
failed = [r for r in self.results if r['perplexity'] is None]
|
||||
|
||||
if successful:
|
||||
# Extract numeric value from "mean±std" format for sorting
|
||||
def get_ppl_value(result):
|
||||
ppl = result['perplexity']
|
||||
if isinstance(ppl, str) and '±' in ppl:
|
||||
return float(ppl.split('±')[0])
|
||||
elif isinstance(ppl, str):
|
||||
try:
|
||||
return float(ppl)
|
||||
except ValueError:
|
||||
return float('inf')
|
||||
return ppl
|
||||
|
||||
successful_sorted = sorted(successful, key=get_ppl_value)
|
||||
|
||||
print(f"{'Dataset':<20} {'Perplexity':>20} {'Time (s)':>12} {'Status':<15}")
|
||||
print(f"{'-'*80}")
|
||||
|
||||
for result in successful_sorted:
|
||||
ppl_str = str(result['perplexity']) if result['perplexity'] is not None else 'N/A'
|
||||
print(f"{result['dataset']:<20} {ppl_str:>20} "
|
||||
f"{result['time']:>12.2f} {result['status']:<15}")
|
||||
|
||||
best_ppl = str(successful_sorted[0]['perplexity'])
|
||||
print(f"\n🏆 Best result: {successful_sorted[0]['dataset']} "
|
||||
f"(PPL: {best_ppl})")
|
||||
|
||||
if failed:
|
||||
print(f"\n❌ Failed tests ({len(failed)}):")
|
||||
for result in failed:
|
||||
print(f" - {result['dataset']}: {result['status']}")
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print(f"✅ Completed: {len(successful)}/{len(self.results)}")
|
||||
print(f"⏱️ Total time: {total_time:.2f}s ({total_time/60:.2f} min)")
|
||||
print(f"📁 Results saved in: {self.output_dir}")
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Test model perplexity on multiple datasets')
|
||||
parser.add_argument('--model', '-m',
|
||||
required=True,
|
||||
help='Path to GGUF model file')
|
||||
parser.add_argument('--data-dir', '-d',
|
||||
default='data',
|
||||
help='Directory containing dataset folders (default: data)')
|
||||
parser.add_argument('--threads', '-t',
|
||||
type=int,
|
||||
default=16,
|
||||
help='Number of threads (default: 16)')
|
||||
parser.add_argument('--ctx-size', '-c',
|
||||
type=int,
|
||||
default=512,
|
||||
help='Context size (default: 512)')
|
||||
parser.add_argument('--output-dir', '-o',
|
||||
default='perplexity_results',
|
||||
help='Output directory for results (default: perplexity_results)')
|
||||
parser.add_argument('--llama-perplexity',
|
||||
default='./build/bin/llama-perplexity',
|
||||
help='Path to llama-perplexity binary (default: ./build/bin/llama-perplexity)')
|
||||
parser.add_argument('--quick', '-q',
|
||||
action='store_true',
|
||||
help='Quick test mode: test all datasets with first 4096 characters and reduced context size (128)')
|
||||
parser.add_argument('--test-embeddings', '-e',
|
||||
action='store_true',
|
||||
help='Test different embedding quantization types (f32, f16, q8_0, q6_k, q5_0, q4_0, q3_k, tq2_0)')
|
||||
parser.add_argument('--csv-output',
|
||||
help='Custom path for CSV output file (e.g., results/my_ppl_results.csv)')
|
||||
parser.add_argument('--quantize-bin',
|
||||
default='./build/bin/llama-quantize',
|
||||
help='Path to llama-quantize binary (default: ./build/bin/llama-quantize)')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
tester = PerplexityTester(
|
||||
model_path=args.model,
|
||||
llama_perplexity_bin=args.llama_perplexity,
|
||||
data_dir=args.data_dir,
|
||||
output_dir=args.output_dir,
|
||||
quick_mode=args.quick,
|
||||
quantize_bin=args.quantize_bin,
|
||||
test_embeddings=args.test_embeddings,
|
||||
csv_output=args.csv_output
|
||||
)
|
||||
|
||||
tester.run_all_tests(
|
||||
threads=args.threads,
|
||||
ctx_size=args.ctx_size
|
||||
)
|
||||
|
||||
except FileNotFoundError as e:
|
||||
print(f"❌ Error: {e}")
|
||||
return 1
|
||||
except KeyboardInterrupt:
|
||||
print("\n\n⚠️ Test interrupted by user")
|
||||
return 1
|
||||
except Exception as e:
|
||||
print(f"\n❌ Unexpected error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return 1
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
exit(main())
|
||||
Executable
+151
@@ -0,0 +1,151 @@
|
||||
#!/bin/bash
|
||||
# Monitor power consumption for llama-bench with different thread configurations
|
||||
# Usage: ./monitor_power.sh <model_path> <output_csv> <pp_threads> <tg_threads>
|
||||
# Example: ./monitor_power.sh models/model.gguf results.csv "1,2,4,8" "1,2,4,8"
|
||||
|
||||
set -e
|
||||
|
||||
# Parse arguments
|
||||
if [ $# -ne 4 ]; then
|
||||
echo "Usage: $0 <model_path> <output_csv> <pp_threads> <tg_threads>"
|
||||
echo "Example: $0 models/model.gguf results.csv \"1,2,4,8\" \"1,2,4,8\""
|
||||
exit 1
|
||||
fi
|
||||
|
||||
MODEL_PATH="$1"
|
||||
OUTPUT_CSV="$2"
|
||||
PP_THREADS="$3"
|
||||
TG_THREADS="$4"
|
||||
|
||||
TEMP_LOG="/tmp/power_monitor_$$.log"
|
||||
PID_FILE="/tmp/monitor_$$.pid"
|
||||
BENCH_OUTPUT="/tmp/bench_output_$$.txt"
|
||||
|
||||
# Validate model exists
|
||||
if [ ! -f "$MODEL_PATH" ]; then
|
||||
echo "Error: Model file not found: $MODEL_PATH"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Create output directory if needed
|
||||
mkdir -p "$(dirname "$OUTPUT_CSV")"
|
||||
|
||||
# Function to monitor CPU stats
|
||||
monitor_cpu() {
|
||||
local log_file="$1"
|
||||
echo "Timestamp,CPU_Usage(%),Avg_Freq(MHz)" > "$log_file"
|
||||
while [ -f "$PID_FILE" ]; do
|
||||
cpu_usage=$(top -bn1 | grep "Cpu(s)" | awk '{print 100-$8}')
|
||||
avg_freq=$(grep "cpu MHz" /proc/cpuinfo | awk '{sum+=$4; count++} END {printf "%.0f", sum/count}')
|
||||
timestamp=$(date +%s.%N)
|
||||
echo "$timestamp,$cpu_usage,$avg_freq" >> "$log_file"
|
||||
sleep 0.5
|
||||
done
|
||||
}
|
||||
|
||||
# Function to calculate average power
|
||||
calculate_power() {
|
||||
local log_file="$1"
|
||||
awk -F',' 'NR>1 {sum_cpu+=$2; count++} END {
|
||||
if (count > 0) {
|
||||
avg_cpu = sum_cpu/count
|
||||
est_power = avg_cpu * 200 / 100
|
||||
printf "%.2f", est_power
|
||||
} else {
|
||||
print "0"
|
||||
}
|
||||
}' "$log_file"
|
||||
}
|
||||
|
||||
# Function to extract throughput from llama-bench output
|
||||
extract_throughput() {
|
||||
local bench_output="$1"
|
||||
local workload="$2"
|
||||
grep "$workload" "$bench_output" | awk '{
|
||||
# Extract mean from "mean ± std" format
|
||||
for (i=1; i<=NF; i++) {
|
||||
if ($(i+1) == "±") {
|
||||
printf "%.2f", $i
|
||||
exit
|
||||
}
|
||||
}
|
||||
}'
|
||||
}
|
||||
|
||||
# Function to run single benchmark
|
||||
run_benchmark() {
|
||||
local workload="$1" # "pp" or "tg"
|
||||
local threads="$2"
|
||||
local n_flag=""
|
||||
|
||||
if [ "$workload" = "pp" ]; then
|
||||
n_flag="-n 0"
|
||||
workload_name="pp128"
|
||||
else
|
||||
n_flag="-n 128"
|
||||
workload_name="tg128"
|
||||
fi
|
||||
|
||||
# Output progress to stderr (won't be captured in CSV)
|
||||
echo "Testing $workload_name with $threads threads..." >&2
|
||||
|
||||
# Start monitoring
|
||||
touch "$PID_FILE"
|
||||
monitor_cpu "$TEMP_LOG" &
|
||||
local monitor_pid=$!
|
||||
|
||||
# Run benchmark
|
||||
./build/bin/llama-bench -m "$MODEL_PATH" -p 128 $n_flag -t "$threads" -ngl 0 > "$BENCH_OUTPUT" 2>&1
|
||||
|
||||
# Stop monitoring
|
||||
rm -f "$PID_FILE"
|
||||
wait $monitor_pid 2>/dev/null || true
|
||||
|
||||
# Extract results
|
||||
local throughput=$(extract_throughput "$BENCH_OUTPUT" "$workload_name")
|
||||
local power=$(calculate_power "$TEMP_LOG")
|
||||
|
||||
if [ -z "$throughput" ] || [ "$throughput" = "0" ]; then
|
||||
echo "Warning: Failed to extract throughput for $workload_name, threads=$threads" >&2
|
||||
throughput="0"
|
||||
fi
|
||||
|
||||
# Calculate J/t (Joules per token)
|
||||
local j_per_token=$(awk -v p="$power" -v t="$throughput" 'BEGIN {
|
||||
if (t > 0) printf "%.4f", p/t; else print "0"
|
||||
}')
|
||||
|
||||
# Output progress to stderr
|
||||
echo " Throughput: $throughput t/s, Power: $power W, Energy: $j_per_token J/t" >&2
|
||||
|
||||
# Only output CSV line to stdout (this will be captured)
|
||||
echo "$workload_name,$threads,$throughput,$power,$j_per_token"
|
||||
}
|
||||
|
||||
# Initialize CSV
|
||||
echo "Workload,Threads,Throughput(t/s),Power(W),Energy(J/t)" > "$OUTPUT_CSV"
|
||||
|
||||
# Test PP workloads
|
||||
IFS=',' read -ra PP_ARRAY <<< "$PP_THREADS"
|
||||
for threads in "${PP_ARRAY[@]}"; do
|
||||
threads=$(echo "$threads" | xargs) # trim whitespace
|
||||
result=$(run_benchmark "pp" "$threads")
|
||||
echo "$result" >> "$OUTPUT_CSV"
|
||||
done
|
||||
|
||||
# Test TG workloads
|
||||
IFS=',' read -ra TG_ARRAY <<< "$TG_THREADS"
|
||||
for threads in "${TG_ARRAY[@]}"; do
|
||||
threads=$(echo "$threads" | xargs) # trim whitespace
|
||||
result=$(run_benchmark "tg" "$threads")
|
||||
echo "$result" >> "$OUTPUT_CSV"
|
||||
done
|
||||
|
||||
# Cleanup
|
||||
rm -f "$TEMP_LOG" "$BENCH_OUTPUT" "$PID_FILE"
|
||||
|
||||
echo ""
|
||||
echo "=== Benchmark Complete ==="
|
||||
echo "Results saved to: $OUTPUT_CSV"
|
||||
echo ""
|
||||
cat "$OUTPUT_CSV"
|
||||
@@ -0,0 +1,362 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
GEMM Configuration Tuning Script
|
||||
This script automatically tunes ROW_BLOCK_SIZE, COL_BLOCK_SIZE, and PARALLEL_SIZE
|
||||
to find the optimal configuration for maximum throughput (t/s).
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import os
|
||||
import re
|
||||
import csv
|
||||
import shutil
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
import argparse
|
||||
|
||||
|
||||
class GemmTuner:
|
||||
def __init__(self, config_path, model_path, threads=16):
|
||||
self.config_path = Path(config_path)
|
||||
self.model_path = model_path
|
||||
self.threads = threads
|
||||
self.backup_path = self.config_path.parent / f"gemm-config.h.backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
|
||||
self.build_dir = Path("../build")
|
||||
self.results = []
|
||||
|
||||
def backup_config(self):
|
||||
"""Backup current configuration file"""
|
||||
print(f"📦 Backing up current config to {self.backup_path}")
|
||||
shutil.copy2(self.config_path, self.backup_path)
|
||||
|
||||
def restore_config(self):
|
||||
"""Restore original configuration file"""
|
||||
print(f"♻️ Restoring original config from {self.backup_path}")
|
||||
shutil.copy2(self.backup_path, self.config_path)
|
||||
|
||||
def generate_config(self, act_parallel, row_block_size, col_block_size, parallel_size):
|
||||
"""Generate new configuration file with simplified format"""
|
||||
content = ""
|
||||
|
||||
# Simplified configuration format
|
||||
if act_parallel:
|
||||
content += "#define ACT_PARALLEL\n"
|
||||
|
||||
content += f"#define ROW_BLOCK_SIZE {row_block_size}\n"
|
||||
content += f"#define COL_BLOCK_SIZE {col_block_size}\n"
|
||||
content += f"#define PARALLEL_SIZE {parallel_size}\n"
|
||||
|
||||
with open(self.config_path, 'w') as f:
|
||||
f.write(content)
|
||||
|
||||
def rebuild_project(self):
|
||||
"""Rebuild project"""
|
||||
print("🔨 Rebuilding project...")
|
||||
result = subprocess.run(
|
||||
["cmake", "--build", str(self.build_dir), "--target", "llama-bench"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
cwd=os.getcwd()
|
||||
)
|
||||
if result.returncode != 0:
|
||||
print(f"⚠️ Build warning/error: {result.stderr}")
|
||||
return False
|
||||
return True
|
||||
|
||||
def run_benchmark(self):
|
||||
"""Run benchmark test"""
|
||||
cmd = [
|
||||
f"{self.build_dir}/bin/llama-bench",
|
||||
"-m", self.model_path,
|
||||
"-p", "128",
|
||||
"-n", "0",
|
||||
"-t", str(self.threads),
|
||||
"-ngl", "0"
|
||||
]
|
||||
|
||||
print(f"⚡ Running benchmark: {' '.join(cmd)}")
|
||||
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
cwd=os.getcwd(),
|
||||
timeout=300 # 5分钟超时
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
print(f"❌ Benchmark failed: {result.stderr}")
|
||||
return None
|
||||
|
||||
return result.stdout
|
||||
|
||||
def parse_throughput(self, output):
|
||||
"""Parse pp128 throughput from output"""
|
||||
# 匹配 pp128: | pp128 | 501.06 ± 11.37 |
|
||||
pp_pattern = r'\|\s+pp128\s+\|\s+([\d.]+)\s+±\s+([\d.]+)\s+\|'
|
||||
pp_match = re.search(pp_pattern, output)
|
||||
|
||||
if pp_match:
|
||||
pp_throughput = float(pp_match.group(1))
|
||||
pp_std_dev = float(pp_match.group(2))
|
||||
|
||||
return {
|
||||
'pp_throughput': pp_throughput,
|
||||
'pp_std_dev': pp_std_dev
|
||||
}
|
||||
|
||||
return None
|
||||
|
||||
def test_configuration(self, act_parallel, row_block_size, col_block_size, parallel_size):
|
||||
"""Test single configuration"""
|
||||
config_name = f"ACT_{'ON' if act_parallel else 'OFF'}_R{row_block_size}_C{col_block_size}_P{parallel_size}"
|
||||
print(f"\n{'='*80}")
|
||||
print(f"🧪 Testing configuration: {config_name}")
|
||||
print(f" ACT_PARALLEL: {act_parallel}")
|
||||
print(f" ROW_BLOCK_SIZE: {row_block_size}")
|
||||
print(f" COL_BLOCK_SIZE: {col_block_size}")
|
||||
print(f" PARALLEL_SIZE: {parallel_size}")
|
||||
print(f"{'='*80}")
|
||||
|
||||
# Generate configuration
|
||||
self.generate_config(act_parallel, row_block_size, col_block_size, parallel_size)
|
||||
|
||||
# Rebuild project
|
||||
if not self.rebuild_project():
|
||||
print("⚠️ Build failed, skipping this configuration")
|
||||
return None
|
||||
|
||||
# Run benchmark test
|
||||
output = self.run_benchmark()
|
||||
if output is None:
|
||||
return None
|
||||
|
||||
# Parse results
|
||||
metrics = self.parse_throughput(output)
|
||||
|
||||
if metrics is not None:
|
||||
result = {
|
||||
'act_parallel': act_parallel,
|
||||
'row_block_size': row_block_size,
|
||||
'col_block_size': col_block_size,
|
||||
'parallel_size': parallel_size,
|
||||
'config_name': config_name,
|
||||
**metrics
|
||||
}
|
||||
self.results.append(result)
|
||||
print(f"✅ PP128: {metrics['pp_throughput']:.2f} ± {metrics['pp_std_dev']:.2f} t/s")
|
||||
return result
|
||||
else:
|
||||
print("❌ Failed to parse throughput")
|
||||
return None
|
||||
|
||||
def save_results(self, csv_path):
|
||||
"""Save results to CSV file"""
|
||||
print(f"\n💾 Saving results to {csv_path}")
|
||||
|
||||
with open(csv_path, 'w', newline='') as f:
|
||||
writer = csv.DictWriter(f, fieldnames=[
|
||||
'config_name', 'act_parallel', 'row_block_size',
|
||||
'col_block_size', 'parallel_size',
|
||||
'pp_throughput', 'pp_std_dev'
|
||||
])
|
||||
writer.writeheader()
|
||||
writer.writerows(self.results)
|
||||
|
||||
def find_best_config(self):
|
||||
"""Find the best configuration with highest throughput"""
|
||||
if not self.results:
|
||||
print("❌ No valid results found")
|
||||
return None
|
||||
|
||||
best = max(self.results, key=lambda x: x['pp_throughput'])
|
||||
return best
|
||||
|
||||
def run_tuning(self, configurations, output_csv=None):
|
||||
"""Run test for all configurations"""
|
||||
print(f"\n🚀 Starting tuning process with {len(configurations)} configurations")
|
||||
print(f"📊 Model: {self.model_path}")
|
||||
print(f"🧵 Threads: {self.threads}\n")
|
||||
|
||||
# Backup configuration
|
||||
self.backup_config()
|
||||
|
||||
try:
|
||||
# Test all configurations
|
||||
for i, config in enumerate(configurations, 1):
|
||||
print(f"\n[{i}/{len(configurations)}]")
|
||||
self.test_configuration(**config)
|
||||
|
||||
# Save results
|
||||
if output_csv is None:
|
||||
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
|
||||
csv_path = f"../stats/tuning_results_{timestamp}.csv"
|
||||
else:
|
||||
csv_path = output_csv
|
||||
|
||||
# Ensure stats directory exists
|
||||
os.makedirs(os.path.dirname(csv_path), exist_ok=True)
|
||||
self.save_results(csv_path)
|
||||
|
||||
# Find best configuration
|
||||
best = self.find_best_config()
|
||||
if best:
|
||||
print(f"\n{'='*80}")
|
||||
print(f"🏆 BEST CONFIGURATION FOUND!")
|
||||
print(f"{'='*80}")
|
||||
print(f"Configuration: {best['config_name']}")
|
||||
print(f"ACT_PARALLEL: {best['act_parallel']}")
|
||||
print(f"ROW_BLOCK_SIZE: {best['row_block_size']}")
|
||||
print(f"COL_BLOCK_SIZE: {best['col_block_size']}")
|
||||
print(f"PARALLEL_SIZE: {best['parallel_size']}")
|
||||
print(f"PP128 Throughput: {best['pp_throughput']:.2f} ± {best['pp_std_dev']:.2f} t/s")
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
# Show the configuration that will be written
|
||||
print("Configuration to be written to gemm-config.h:")
|
||||
print("-" * 80)
|
||||
if best['act_parallel']:
|
||||
print("#define ACT_PARALLEL")
|
||||
print(f"#define ROW_BLOCK_SIZE {best['row_block_size']}")
|
||||
print(f"#define COL_BLOCK_SIZE {best['col_block_size']}")
|
||||
print(f"#define PARALLEL_SIZE {best['parallel_size']}")
|
||||
print("-" * 80)
|
||||
|
||||
# Apply best configuration
|
||||
apply = input("\nDo you want to apply this configuration to gemm-config.h? (y/n): ").strip().lower()
|
||||
if apply == 'y':
|
||||
self.generate_config(
|
||||
best['act_parallel'],
|
||||
best['row_block_size'],
|
||||
best['col_block_size'],
|
||||
best['parallel_size']
|
||||
)
|
||||
self.rebuild_project()
|
||||
print("✅ Best configuration applied and project rebuilt!")
|
||||
else:
|
||||
self.restore_config()
|
||||
print("✅ Original configuration restored")
|
||||
|
||||
# Clean up backup file
|
||||
if self.backup_path.exists():
|
||||
self.backup_path.unlink()
|
||||
print(f"🗑️ Removed backup file: {self.backup_path}")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n⚠️ Tuning interrupted by user")
|
||||
self.restore_config()
|
||||
# Clean up backup file
|
||||
if self.backup_path.exists():
|
||||
self.backup_path.unlink()
|
||||
print(f"🗑️ Removed backup file: {self.backup_path}")
|
||||
except Exception as e:
|
||||
print(f"\n❌ Error during tuning: {e}")
|
||||
self.restore_config()
|
||||
# Clean up backup file
|
||||
if self.backup_path.exists():
|
||||
self.backup_path.unlink()
|
||||
print(f"🗑️ Removed backup file: {self.backup_path}")
|
||||
raise
|
||||
|
||||
|
||||
def generate_configurations():
|
||||
"""Generate list of configurations to test"""
|
||||
configurations = []
|
||||
|
||||
act_parallel_options = [True]
|
||||
|
||||
row_sizes = [2, 4, 8]#[2, 4, 8, 16, 32]
|
||||
col_sizes = [32, 64]#[32, 64, 128, 256, 512, 1024]
|
||||
parallelism_degree = [4]
|
||||
|
||||
for act_parallel in act_parallel_options:
|
||||
for row in row_sizes:
|
||||
for col in col_sizes:
|
||||
for parallel in parallelism_degree:
|
||||
# Add filtering conditions
|
||||
if act_parallel:
|
||||
# When ACT_PARALLEL=True, only calculate combinations with parallel < row
|
||||
if parallel > row:
|
||||
continue
|
||||
else:
|
||||
# When ACT_PARALLEL=False, only calculate combinations with parallel < col
|
||||
if parallel > col:
|
||||
continue
|
||||
|
||||
configurations.append({
|
||||
'act_parallel': act_parallel,
|
||||
'row_block_size': row,
|
||||
'col_block_size': col,
|
||||
'parallel_size': parallel
|
||||
})
|
||||
|
||||
return configurations
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Tune GEMM configuration for optimal performance')
|
||||
parser.add_argument('--config', default='../include/gemm-config.h',
|
||||
help='Path to gemm-config.h file')
|
||||
parser.add_argument('--model', default='../models/BitNet-b1.58-2B-4T/ggml-model-i2_s-embed-q6_k.gguf',
|
||||
help='Path to model file')
|
||||
parser.add_argument('--threads', type=int, default=8,
|
||||
help='Number of threads to use')
|
||||
parser.add_argument('--quick', action='store_true',
|
||||
help='Quick test with fewer configurations')
|
||||
parser.add_argument('--custom', action='store_true',
|
||||
help='Manually specify configurations to test')
|
||||
parser.add_argument('--output', type=str, default=None,
|
||||
help='Output CSV file path (default: stats/tuning_results_<timestamp>.csv)')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
tuner = GemmTuner(args.config, args.model, args.threads)
|
||||
|
||||
if args.custom:
|
||||
# Custom configuration mode
|
||||
print("Custom configuration mode")
|
||||
configurations = []
|
||||
while True:
|
||||
print("\nEnter configuration (or 'done' to finish):")
|
||||
act = input("ACT_PARALLEL (y/n): ").strip().lower() == 'y'
|
||||
if input == 'done':
|
||||
break
|
||||
row = int(input("ROW_BLOCK_SIZE: "))
|
||||
col = int(input("COL_BLOCK_SIZE: "))
|
||||
par = int(input("PARALLEL_SIZE: "))
|
||||
configurations.append({
|
||||
'act_parallel': act,
|
||||
'row_block_size': row,
|
||||
'col_block_size': col,
|
||||
'parallel_size': par
|
||||
})
|
||||
elif args.quick:
|
||||
# Quick test mode - test only a few key configurations
|
||||
configurations = [
|
||||
{'act_parallel': True, 'row_block_size': 4, 'col_block_size': 128, 'parallel_size': 4},
|
||||
{'act_parallel': True, 'row_block_size': 8, 'col_block_size': 128, 'parallel_size': 4},
|
||||
{'act_parallel': True, 'row_block_size': 4, 'col_block_size': 64, 'parallel_size': 4},
|
||||
{'act_parallel': False, 'row_block_size': 32, 'col_block_size': 4, 'parallel_size': 4},
|
||||
{'act_parallel': False, 'row_block_size': 16, 'col_block_size': 4, 'parallel_size': 4},
|
||||
]
|
||||
else:
|
||||
# Full test mode
|
||||
configurations = generate_configurations()
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print(f"GEMM Configuration Tuner")
|
||||
print(f"{'='*80}")
|
||||
print(f"Total configurations to test: {len(configurations)}")
|
||||
print(f"Estimated time: ~{len(configurations) * 0.5:.1f} minutes (assuming 30s per test)")
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
proceed = input("Proceed with tuning? (y/n): ").strip().lower()
|
||||
if proceed != 'y':
|
||||
print("Tuning cancelled")
|
||||
return
|
||||
|
||||
tuner.run_tuning(configurations, output_csv=args.output)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user