diff --git a/baselines_1/Llama-2-7b-hf_bi_idx_10.json b/baselines_1/Llama-2-7b-hf_bi_idx_10.json new file mode 100644 index 0000000000000000000000000000000000000000..4548c1bb35b151f690771aadba9ec984b0bfe05a --- /dev/null +++ b/baselines_1/Llama-2-7b-hf_bi_idx_10.json @@ -0,0 +1,34 @@ +[ + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 4, + 4 +] \ No newline at end of file diff --git a/baselines_1/Llama-2-7b-hf_bi_idx_5.json b/baselines_1/Llama-2-7b-hf_bi_idx_5.json new file mode 100644 index 0000000000000000000000000000000000000000..fe2595cad704583b8cb3e90f038f7be61e95904d --- /dev/null +++ b/baselines_1/Llama-2-7b-hf_bi_idx_5.json @@ -0,0 +1,34 @@ +[ + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 2, + 2, + 2, + 2, + 2, + 4, + 4, + 4 +] \ No newline at end of file diff --git a/baselines_1/Llama-2-7b-hf_zd_idx_10.json b/baselines_1/Llama-2-7b-hf_zd_idx_10.json new file mode 100644 index 0000000000000000000000000000000000000000..8c22663c27c654e791172192258b277f329cb4ec --- /dev/null +++ b/baselines_1/Llama-2-7b-hf_zd_idx_10.json @@ -0,0 +1,34 @@ +[ + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 2, + 2, + 4, + 2, + 2, + 2, + 2, + 4, + 2, + 2, + 2, + 2, + 4, + 4, + 4 +] \ No newline at end of file diff --git a/baselines_1/Llama-2-7b-hf_zd_idx_5.json b/baselines_1/Llama-2-7b-hf_zd_idx_5.json new file mode 100644 index 0000000000000000000000000000000000000000..f147c5627c10f8d78e815f510bacaf8d7c3f14f5 --- /dev/null +++ b/baselines_1/Llama-2-7b-hf_zd_idx_5.json @@ -0,0 +1,34 @@ +[ + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 2, + 2, + 4, + 2, + 4, + 2, + 2, + 4, + 4, + 4 +] \ No newline at end of file diff --git a/llm-awq/.gitignore b/llm-awq/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..6417a78a68f4d3e0b5cc6728592630d0f8a50c01 --- /dev/null +++ b/llm-awq/.gitignore @@ -0,0 +1,173 @@ +.DS_Store + +data/ +checkpoints +demo_images +serve_images +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +*.pyc +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ + +*.pt +**/*.pt +**/*.pyc +*.json +__pycache__ diff --git a/llm-awq/LICENSE b/llm-awq/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..aca327a505563c97f6ce15cbb88098e8e72f3965 --- /dev/null +++ b/llm-awq/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 MIT HAN Lab + +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. diff --git a/llm-awq/awq.egg-info/top_level.txt b/llm-awq/awq.egg-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..7928383645d42ce1f3209d78cf2e3d8345ed8273 --- /dev/null +++ b/llm-awq/awq.egg-info/top_level.txt @@ -0,0 +1,3 @@ +awq +figures +tinychat diff --git a/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention.h b/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention.h new file mode 100644 index 0000000000000000000000000000000000000000..f68d48574900cd67487174784bd1bc26b037ac8b --- /dev/null +++ b/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention.h @@ -0,0 +1,185 @@ +// Downloaded from from FasterTransformer v5.2.1 +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.2.1_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention.h +/* + * Copyright (c) 2020-2022, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cuda_bf16_wrapper.h" +#include +#include +#include +#include +#include + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#define CHECK_CUDA(call) \ + do { \ + cudaError_t status_ = call; \ + if (status_ != cudaSuccess) { \ + fprintf(stderr, "CUDA error (%s:%d): %s\n", __FILE__, __LINE__, cudaGetErrorString(status_)); \ + exit(1); \ + } \ + } while (0) + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// The structure of parameters for the masked multihead attention kernel. +// +// We use the following terminology to describe the different dimensions. +// +// B: Batch size (number of sequences), +// L: Sequence length, +// D: Hidden dimension, +// H: Number of heads, +// Dh: Hidden dimension per head - Dh = D / H. + +template +struct Multihead_attention_params_base { + + // The output buffer. Dimensions B x D. + T* out = nullptr; + + // The input Qs and the associated bias. Dimensions B x D and D, resp. + const T *q = nullptr, *q_bias = nullptr; + // The input Ks and the associated bias. Dimensions B x D and D, resp. + const T *k = nullptr, *k_bias = nullptr; + // The input Vs and the associated bias. Dimensions B x D and D, resp. + const T *v = nullptr, *v_bias = nullptr; + + // The cache for the Ks. The size must be at least B x L x D. + T* k_cache = nullptr; + // The cache for the Vs. The size must be at least B x L x D. + T* v_cache = nullptr; + // The indirections to use for cache when beam sampling. + const int* cache_indir = nullptr; + + // Stride to handle the case when KQV is a single buffer + int stride = 0; + + // The batch size. + int batch_size = 0; + // The beam width + int beam_width = 0; + // The sequence length. + int memory_max_len = 0; + // The number of heads (H). + int num_heads = 0; + // The number of heads for KV cache. + int num_kv_heads = 0; + // The hidden dimension per head (Dh). + int hidden_size_per_head = 0; + // The per-head latent space reserved for rotary embeddings. + int rotary_embedding_dim = 0; + bool neox_rotary_style = false; + float rotary_base = 0.0f; + float rotary_scale = 1.0f; + // The maximum length of input sentences. + int max_input_length = 0; + // The current timestep. TODO(bhsueh) Check that do we only this param in cross attention? + int timestep = 0; + // The current timestep of each sentences (support different timestep for different sentences) + + // The 1.f / sqrt(Dh). Computed on the host. + float inv_sqrt_dh = 0.0f; + + // Used when we have some input context like gpt + const int* total_padding_tokens = nullptr; + + const bool* masked_tokens = nullptr; + const int* prefix_prompt_lengths = nullptr; + int max_prefix_prompt_length = 0; + + const T* relative_attention_bias = nullptr; + int relative_attention_bias_stride = 0; + // The slope per head of linear position bias to attention score (H). + const float* linear_bias_slopes = nullptr; + + const T* ia3_key_weights = nullptr; + const T* ia3_value_weights = nullptr; + const int* ia3_tasks = nullptr; + + const float* qkv_scale_out = nullptr; + const float* attention_out_scale = nullptr; + int int8_mode = 0; +}; + +template +struct Multihead_attention_params: public Multihead_attention_params_base { + // output cross attentions + float* cross_attention_out = nullptr; + int max_decoder_seq_len = 0; + bool is_return_cross_attentions = false; + + // allows to exist attention eary + bool* finished = nullptr; + + // required in case of cross attention + // will need it here till if constexpr in c++17 + int* memory_length_per_sample = nullptr; + + // required in case of masked attention with different length + const int* length_per_sample = nullptr; +}; + +template +struct Multihead_attention_params: public Multihead_attention_params_base { + // output cross attentions + float* cross_attention_out = nullptr; + int max_decoder_seq_len = 0; + bool is_return_cross_attentions = false; + + // allows to exist attention eary + bool* finished = nullptr; + + // required in case of cross attention + int* memory_length_per_sample = nullptr; + + // required in case of masked attention with different length + const int* length_per_sample = nullptr; +}; + +template +using Masked_multihead_attention_params = Multihead_attention_params; + +template +using Cross_multihead_attention_params = Multihead_attention_params; + +template +struct outputCrossAttentionParam { + // max decoder output length + int max_decoder_seq_len = 0; + T* cross_attention_out = nullptr; + bool is_return_cross_attentions = false; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +void masked_multihead_attention(const Masked_multihead_attention_params& params, const cudaStream_t& stream); +void masked_multihead_attention(const Masked_multihead_attention_params& params, const cudaStream_t& stream); +#ifdef ENABLE_BF16 +void masked_multihead_attention(const Masked_multihead_attention_params<__nv_bfloat16>& params, + const cudaStream_t& stream); +#endif +void cross_multihead_attention(const Cross_multihead_attention_params& params, const cudaStream_t& stream); +void cross_multihead_attention(const Cross_multihead_attention_params& params, const cudaStream_t& stream); +#ifdef ENABLE_BF16 +void cross_multihead_attention(const Cross_multihead_attention_params<__nv_bfloat16>& params, + const cudaStream_t& stream); +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention_template.hpp b/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention_template.hpp new file mode 100644 index 0000000000000000000000000000000000000000..c336d8017e026f66d77af8e40f2bb0494ea61214 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention_template.hpp @@ -0,0 +1,1608 @@ +// Downloaded from from FasterTransformer v5.2.1 +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.2.1_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp +/* + * Copyright (c) 2020-2022, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "decoder_masked_multihead_attention.h" +#include "decoder_masked_multihead_attention_utils.h" +#include "cuda_bf16_wrapper.h" +#include "cuda_bf16_fallbacks.cuh" +#include +#include +#include + +// #define MMHA_USE_HMMA_FOR_REDUCTION + +// Below are knobs to extend FP32 accumulation for higher FP16 accuracy + +// Does not seem to affect the accuracy that much +#define MMHA_USE_FP32_ACUM_FOR_FMA + +// Seems to slightly improve the accuracy +#define MMHA_USE_FP32_ACUM_FOR_OUT + +#if 0 && defined(MMHA_USE_FP32_ACUM_FOR_OUT) + // Does not seem to improve the accuracy + //#define MMHA_USE_FP32_ACUM_FOR_LOGITS +#endif + +namespace mmha { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// +// We use the following terminology to describe the different dimensions. +// +// B: Batch size (number of sequences), +// L: Sequence length, +// D: Hidden dimension, +// H: Number of heads, +// Dh: Hidden dimension per head - Dh = D / H. +// +// The different kernels assign a threadblock for B x H pair. The grid has size (1, B, H). We use +// 64, 128 and 256 threads per block. +// +// Each threadblock loads Dh values from Q and its associated bias. The kernels run a loop to +// compute Q * K^T where K is loaded from a cache buffer -- except for the current timestep. The +// cache buffer helps with memory accesses and contains keys with bias. +// +// The layout of the cache buffer for the keys is [B, H, Dh/x, L, x] where x == 8 for FP16 and +// x == 4 for FP32 where the fastest moving dimension (contiguous data) is the rightmost one. The +// values for x are chosen to create chunks of 16 bytes. +// +// The different kernels use 1, 2 or 4 threads per key (THREADS_PER_KEY). The size of the LDGs +// depends on the number of threads per key. Each thread sums Dh / THREADS_PER_KEY elements. At +// the end of each iteration of the Q * K^T loop, we perform a reduction between lanes using an +// HMMA instruction (Tensor Core). Each Q * K^T valuey is stored in shared memory in FP32. +// +// After that loop, a parallel softmax is computed across the different Q * K^T values stored in +// shared memory. +// +// The kernel ends with a loop over the values in V. We use THREADS_PER_VALUE to control how many +// timesteps are computed by loop iteration. As with the keys, the values are read from a cache +// except for the current timestep. The layout of the cache buffer for the values is much simpler +// as it is [B, H, L, Dh]. +// + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct Qk_vec_ { +}; + +template<> +struct Qk_vec_ { + using Type = float; +}; +template<> +struct Qk_vec_ { + using Type = float2; +}; +template<> +struct Qk_vec_ { + using Type = float4; +}; +template<> +struct Qk_vec_ { + using Type = float4; +}; +template<> +struct Qk_vec_ { + using Type = uint32_t; +}; +template<> +struct Qk_vec_ { + using Type = uint32_t; +}; +template<> +struct Qk_vec_ { + using Type = uint2; +}; +template<> +struct Qk_vec_ { + using Type = uint4; +}; +#ifdef ENABLE_BF16 +template<> +struct Qk_vec_<__nv_bfloat16, 32> { + using Type = __nv_bfloat162; +}; +template<> +struct Qk_vec_<__nv_bfloat16, 64> { + using Type = __nv_bfloat162; +}; +template<> +struct Qk_vec_<__nv_bfloat16, 128> { + using Type = bf16_4_t; +}; +template<> +struct Qk_vec_<__nv_bfloat16, 256> { + using Type = bf16_8_t; +}; +#endif // ENABLE_BF16 +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct K_vec_ { +}; + +template<> +struct K_vec_ { + using Type = float; +}; +template<> +struct K_vec_ { + using Type = float2; +}; +template<> +struct K_vec_ { + using Type = float4; +}; +template<> +struct K_vec_ { + using Type = uint32_t; +}; +template<> +struct K_vec_ { + using Type = uint2; +}; +template<> +struct K_vec_ { + using Type = uint4; +}; +#ifdef ENABLE_BF16 +template<> +struct K_vec_<__nv_bfloat16, 4> { + using Type = __nv_bfloat162; +}; +template<> +struct K_vec_<__nv_bfloat16, 2> { + using Type = bf16_4_t; +}; +template<> +struct K_vec_<__nv_bfloat16, 1> { + using Type = bf16_8_t; +}; +#endif // ENABLE_BF16 +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct V_vec_ { +}; + +template<> +struct V_vec_ { + using Type = float; +}; +template<> +struct V_vec_ { + using Type = float2; +}; +template<> +struct V_vec_ { + using Type = float4; +}; +template<> +struct V_vec_ { + using Type = uint32_t; +}; +template<> +struct V_vec_ { + using Type = uint2; +}; +template<> +struct V_vec_ { + using Type = uint4; +}; +#ifdef ENABLE_BF16 +template<> +struct V_vec_<__nv_bfloat16, 2> { + using Type = __nv_bfloat162; +}; +template<> +struct V_vec_<__nv_bfloat16, 4> { + using Type = bf16_4_t; +}; +template<> +struct V_vec_<__nv_bfloat16, 8> { + using Type = bf16_8_t; +}; +#endif // ENABLE_BF16 +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef MMHA_USE_FP32_ACUM_FOR_FMA +template +struct Qk_vec_acum_fp32_ { +}; + +template<> +struct Qk_vec_acum_fp32_ { + using Type = float; +}; +template<> +struct Qk_vec_acum_fp32_ { + using Type = float2; +}; +template<> +struct Qk_vec_acum_fp32_ { + using Type = float4; +}; +// template<> struct Qk_vec_acum_fp32_ { using Type = float; }; +template<> +struct Qk_vec_acum_fp32_ { + using Type = float2; +}; +template<> +struct Qk_vec_acum_fp32_ { + using Type = Float4_; +}; +template<> +struct Qk_vec_acum_fp32_ { + using Type = Float8_; +}; +template<> +struct Qk_vec_acum_fp32_<__nv_bfloat16> { + using Type = float; +}; +template<> +struct Qk_vec_acum_fp32_<__nv_bfloat162> { + using Type = float2; +}; +template<> +struct Qk_vec_acum_fp32_ { + using Type = Float4_; +}; +template<> +struct Qk_vec_acum_fp32_ { + using Type = Float8_; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct K_vec_acum_fp32_ { +}; + +template<> +struct K_vec_acum_fp32_ { + using Type = float; +}; +template<> +struct K_vec_acum_fp32_ { + using Type = float2; +}; +template<> +struct K_vec_acum_fp32_ { + using Type = float4; +}; +template<> +struct K_vec_acum_fp32_ { + using Type = float2; +}; +template<> +struct K_vec_acum_fp32_ { + using Type = Float4_; +}; +template<> +struct K_vec_acum_fp32_ { + using Type = Float8_; +}; +template<> +struct K_vec_acum_fp32_<__nv_bfloat16> { + using Type = float; +}; +template<> +struct K_vec_acum_fp32_<__nv_bfloat162> { + using Type = float2; +}; +template<> +struct K_vec_acum_fp32_ { + using Type = Float4_; +}; +template<> +struct K_vec_acum_fp32_ { + using Type = Float8_; +}; +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef MMHA_USE_FP32_ACUM_FOR_OUT +template +struct V_vec_acum_fp32_ { +}; + +template<> +struct V_vec_acum_fp32_ { + using Type = float; +}; +template<> +struct V_vec_acum_fp32_ { + using Type = float2; +}; +template<> +struct V_vec_acum_fp32_ { + using Type = float4; +}; +template<> +struct V_vec_acum_fp32_ { + using Type = float2; +}; +template<> +struct V_vec_acum_fp32_ { + using Type = Float4_; +}; +template<> +struct V_vec_acum_fp32_ { + using Type = Float8_; +}; +#ifdef ENABLE_BF16 +template<> +struct V_vec_acum_fp32_<__nv_bfloat162> { + using Type = float2; +}; +template<> +struct V_vec_acum_fp32_ { + using Type = Float4_; +}; +template<> +struct V_vec_acum_fp32_ { + using Type = Float8_; +}; +#endif // ENABLE_BF16 +#endif +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline __device__ float qk_dot_(const K_vec (&q)[N], const K_vec (&k)[N]) +{ +#ifdef MMHA_USE_FP32_ACUM_FOR_FMA + using K_vec_acum = typename K_vec_acum_fp32_::Type; +#else + using K_vec_acum = K_vec; +#endif + // Compute the parallel products for Q*K^T (treat vector lanes separately). + K_vec_acum qk_vec = mul(q[0], k[0]); +#pragma unroll + for (int ii = 1; ii < N; ++ii) { + qk_vec = fma(q[ii], k[ii], qk_vec); + } + + // Finalize the reduction across lanes. + float qk = sum(qk_vec); +#pragma unroll + for (int mask = THREADS_PER_KEY / 2; mask >= 1; mask /= 2) { + qk += __shfl_xor_sync(uint32_t(-1), qk, mask); + } + return qk; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct Qk_dot { + template + static inline __device__ float dot(const K_vec (&q)[N], const K_vec (&k)[N]) + { + return qk_dot_(q, k); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float4 hmma_fp32(const uint2& a, uint32_t b) +{ + float4 c; + float zero = 0.f; + asm volatile("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 \n" + " {%0, %1, %2, %3}, \n" + " {%4, %5}, \n" + " {%6}, \n" + " {%7, %7, %7, %7}; \n" + + : "=f"(c.x), "=f"(c.y), "=f"(c.z), "=f"(c.w) + : "r"(a.x) "r"(a.y), "r"(b), "f"(zero)); + return c; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline __device__ float qk_hmma_dot_(const uint32_t (&q)[N], const uint32_t (&k)[N]) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 750 +#ifdef MMHA_USE_FP32_ACUM_FOR_FMA + using K_vec_acum = typename K_vec_acum_fp32_::Type; +#else + using K_vec_acum = uint32_t; +#endif + K_vec_acum qk_vec = mul(q[0], k[0]); +#pragma unroll + for (int ii = 1; ii < N; ++ii) { + qk_vec = fma(q[ii], k[ii], qk_vec); + } +#ifdef MMHA_USE_FP32_ACUM_FOR_FMA + uint32_t qk_vec_ = float2_to_half2(qk_vec); + return hmma_fp32(make_uint2(qk_vec_, 0u), 0x3c003c00u).x; +#else + return hmma_fp32(make_uint2(qk_vec, 0u), 0x3c003c00u).x; +#endif +#else + return 0.f; +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +struct Qk_dot { + template + static inline __device__ float dot(const uint32_t (&q)[N], const uint32_t (&k)[N]) + { +#if __CUDA_ARCH__ >= 750 && defined(MMHA_USE_HMMA_FOR_REDUCTION) + return qk_hmma_dot_(q, k); +#else + return qk_dot_<4>(q, k); +#endif // defined MMHA_USE_HMMA_FOR_REDUCTION + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline __device__ float block_sum(float* red_smem, float sum) +{ + + // Decompose the thread index into warp / lane. + int warp = threadIdx.x / WARP_SIZE; + int lane = threadIdx.x % WARP_SIZE; + +// Compute the sum per warp. +#pragma unroll + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { + sum += __shfl_xor_sync(uint32_t(-1), sum, mask); + } + + // Warp leaders store the data to shared memory. + if (lane == 0) { + red_smem[warp] = sum; + } + + // Make sure the data is in shared memory. + __syncthreads(); + + // The warps compute the final sums. + if (lane < WARPS_PER_BLOCK) { + sum = red_smem[lane]; + } + +// Parallel reduction inside the warp. +#pragma unroll + for (int mask = WARPS_PER_BLOCK / 2; mask >= 1; mask /= 2) { + sum += __shfl_xor_sync(uint32_t(-1), sum, mask); + } + + // Broadcast to other threads. + return __shfl_sync(uint32_t(-1), sum, 0); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(float& dst, float src) +{ + dst = src; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(uint16_t& dst, float src) +{ + dst = float_to_half(src); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(uint32_t& dst, float2 src) +{ + dst = float2_to_half2(src); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +#ifdef ENABLE_BF16 +inline __device__ void convert_from_float(__nv_bfloat16& dst, float src) +{ + dst = __float2bfloat16(src); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(__nv_bfloat162& dst, float2 src) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + dst = __float22bfloat162_rn(src); +#else + dst = __floats2bfloat162_rn(src.x, src.y); +#endif +} +#endif // ENABLE_BF16 +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(uint2& dst, Float4_ src) +{ + dst.x = float2_to_half2(src.x); + dst.y = float2_to_half2(src.y); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(uint2& dst, float4 src) +{ + convert_from_float(dst, Float4_{make_float2(src.x, src.y), make_float2(src.z, src.w)}); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(uint4& dst, Float8_ src) +{ + dst.x = float2_to_half2(src.x); + dst.y = float2_to_half2(src.y); + dst.z = float2_to_half2(src.z); + dst.w = float2_to_half2(src.w); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef ENABLE_BF16 +inline __device__ void convert_from_float(bf16_4_t& dst, Float4_ src) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + dst.x = __float22bfloat162_rn(src.x); + dst.y = __float22bfloat162_rn(src.y); +#else + dst.x = __floats2bfloat162_rn(src.x.x, src.x.y); + dst.y = __floats2bfloat162_rn(src.y.x, src.y.y); +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(bf16_4_t& dst, float4 src) +{ + convert_from_float(dst, Float4_{make_float2(src.x, src.y), make_float2(src.z, src.w)}); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(bf16_8_t& dst, Float8_ src) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + dst.x = __float22bfloat162_rn(src.x); + dst.y = __float22bfloat162_rn(src.y); + dst.z = __float22bfloat162_rn(src.z); + dst.w = __float22bfloat162_rn(src.w); +#else + dst.x = __floats2bfloat162_rn(src.x.x, src.x.y); + dst.y = __floats2bfloat162_rn(src.y.x, src.y.y); + dst.z = __floats2bfloat162_rn(src.z.x, src.z.y); + dst.w = __floats2bfloat162_rn(src.w.x, src.w.y); +#endif +} +#endif // ENABLE_BF16 + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(float2& dst, float2 src) +{ + dst = src; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ void convert_from_float(float4& dst, float4 src) +{ + dst = src; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float convert_to_float(float4 u) +{ + return u.x; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float convert_to_float(uint4 u) +{ + float2 tmp = half2_to_float2(u.x); + return tmp.x; +} + +#if defined(MMHA_USE_FP32_ACUM_FOR_LOGITS) + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float cast_to_float(float u) +{ + return u; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float2 cast_to_float(float2 u) +{ + return u; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float4 cast_to_float(float4 u) +{ + return u; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float4_ cast_to_float(Float4_ u) +{ + return u; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float8_ cast_to_float(Float8_ u) +{ + return u; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float2 cast_to_float(uint32_t u) +{ + return half2_to_float2(u); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float4_ cast_to_float(uint2 u) +{ + Float4_ tmp; + tmp.x = half2_to_float2(u.x); + tmp.y = half2_to_float2(u.y); + return tmp; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ Float8_ cast_to_float(uint4 u) +{ + Float8_ tmp; + tmp.x = half2_to_float2(u.x); + tmp.y = half2_to_float2(u.y); + tmp.z = half2_to_float2(u.z); + tmp.w = half2_to_float2(u.w); + return tmp; +} + +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float float_from_int8(int8_t u) +{ + return u; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float2 float_from_int8(int16_t u) +{ + union { + int16_t int16; + int8_t int8[2]; + }; + int16 = u; + return make_float2(int8[0], int8[1]); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ float4 float_from_int8(int32_t u) +{ + union { + int32_t int32; + int8_t int8[4]; + }; + int32 = u; + return make_float4(int8[0], int8[1], int8[2], int8[3]); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// clang-format off +inline __device__ Float8_ float_from_int8(int64_t u) +{ + union { + int64_t int64; + int16_t int16[4]; + }; + int64 = u; + return Float8_ {float_from_int8(int16[0]), + float_from_int8(int16[1]), + float_from_int8(int16[2]), + float_from_int8(int16[3])}; +} +// clang-format on + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ int8_t cast_to_int8(float val) +{ + union { + int8_t int8[2]; + int16_t int16; + }; + asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=h"(int16) : "f"(val)); + return int8[0]; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ int32_t cast_to_int8(float4 val) +{ + union { + int8_t int8[4]; + int32_t int32; + }; + int8[0] = cast_to_int8(val.x); + int8[1] = cast_to_int8(val.y); + int8[2] = cast_to_int8(val.z); + int8[3] = cast_to_int8(val.w); + return int32; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ int64_t cast_to_int8(Float8_ val) +{ + union { + int8_t int8[8]; + int64_t int64; + }; + int8[0] = cast_to_int8(val.x.x); + int8[1] = cast_to_int8(val.x.y); + int8[2] = cast_to_int8(val.y.x); + int8[3] = cast_to_int8(val.y.y); + int8[4] = cast_to_int8(val.z.x); + int8[5] = cast_to_int8(val.z.y); + int8[6] = cast_to_int8(val.w.x); + int8[7] = cast_to_int8(val.w.y); + return int64; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline __device__ __host__ T div_up(T m, T n) +{ + return (m + n - 1) / n; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline size_t smem_size_in_bytes(const Multihead_attention_params& params, + int threads_per_value, + int threads_per_block) +{ + // The amount of shared memory needed to store the Q*K^T values in float. + const int max_timesteps = min(params.timestep, params.memory_max_len); + size_t qk_sz = (DO_CROSS_ATTENTION) ? div_up(params.memory_max_len + 1, 4) * 16 : div_up(max_timesteps + 1, 4) * 16; + + // The extra memory needed if we are not using floats for the final logits. + size_t logits_sz = 0; +#ifndef MMHA_USE_FP32_ACUM_FOR_LOGITS + if (sizeof(T) != 4) { + // TDOD + logits_sz = (DO_CROSS_ATTENTION) ? div_up(params.memory_max_len + 1, 4) * 4 * sizeof(T) : + div_up(max_timesteps + 1, 4) * 4 * sizeof(T); + } +#endif + + // The total size needed during softmax. + size_t softmax_sz = qk_sz + logits_sz; + + // The number of partial rows to reduce in the final reduction. + int rows_per_red = threads_per_block / threads_per_value; + // The amount of storage needed to finalize the outputs. + size_t red_sz = rows_per_red * params.hidden_size_per_head * sizeof(T) / 2; + + size_t transpose_rotary_size = 0; + if (params.rotary_embedding_dim > 0 && params.neox_rotary_style) { + transpose_rotary_size = 2 * params.rotary_embedding_dim * sizeof(T); + } + + // The max. + return max(max(softmax_sz, red_sz), transpose_rotary_size); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline __device__ constexpr uint32_t shfl_mask(int threads) +{ + return threads == 32 ? uint32_t(-1) : (1u << threads) - 1u; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template< + // The type of the inputs. Supported types: float and half. + typename T, + // The hidden dimension per head. + int Dh, + int Dh_MAX, + // The number of threads per key. + int THREADS_PER_KEY, + // The number of threads per value. + int THREADS_PER_VALUE, + // The number of threads in a threadblock. + int THREADS_PER_BLOCK, + bool DO_CROSS_ATTENTION> +__global__ void masked_multihead_attention_kernel(Multihead_attention_params params) +{ + + // Make sure the hidden dimension per head is a multiple of the number of threads per key. + static_assert(Dh_MAX % THREADS_PER_KEY == 0, ""); + // Make sure the hidden dimension per head is a multiple of the number of threads per value. + static_assert(Dh_MAX % THREADS_PER_VALUE == 0, ""); + + // The size of a warp. + constexpr int WARP_SIZE = 32; + // The number of warps in a threadblock. + constexpr int WARPS_PER_BLOCK = THREADS_PER_BLOCK / WARP_SIZE; + + // Use smem_size_in_bytes (above) to determine the amount of shared memory. + extern __shared__ char smem_[]; + + // The shared memory for the Q*K^T values and partial logits in softmax. + float* qk_smem = reinterpret_cast(smem_); + + // The shared memory for the logits. For FP32, that's the same buffer as qk_smem. + char* logits_smem_ = smem_; +#ifndef MMHA_USE_FP32_ACUM_FOR_LOGITS + if (sizeof(T) != 4) { + // TODO - change to tlength + const int max_timesteps = min(params.timestep, params.memory_max_len); + logits_smem_ += + (DO_CROSS_ATTENTION) ? div_up(params.memory_max_len + 1, 4) * 16 : div_up(max_timesteps + 1, 4) * 16; + } + T* logits_smem = reinterpret_cast(logits_smem_); +#else + float* logits_smem = reinterpret_cast(logits_smem_); +#endif + + // The shared memory to do the final reduction for the output values. Reuse qk_smem. + T* out_smem = reinterpret_cast(smem_); + + // The shared memory buffers for the block-wide reductions. One for max, one for sum. + __shared__ float red_smem[WARPS_PER_BLOCK * 2]; + + // A vector of Q or K elements for the current timestep. + using Qk_vec = typename Qk_vec_::Type; + + // Use alignment for safely casting the shared buffers as Qk_vec. + // Shared memory to store Q inputs. + __shared__ __align__(sizeof(Qk_vec)) T q_smem[Dh_MAX]; + + // This is one of the reasons we should have a separate kernel for cross attention + __shared__ __align__(sizeof(Qk_vec)) T bias_smem[DO_CROSS_ATTENTION ? Dh_MAX : 1]; + + // A vector of Q or K elements for the current timestep. + using Qk_vec = typename Qk_vec_::Type; + // The number of elements per vector. + constexpr int QK_VEC_SIZE = sizeof(Qk_vec) / sizeof(T); + // Make sure the hidden size per head is a multiple of the vector size. + static_assert(Dh_MAX % QK_VEC_SIZE == 0, ""); + // We will use block wide reduction if needed + // static_assert(Dh_MAX / QK_VEC_SIZE <= WARP_SIZE, ""); + // The number of vectors per warp. + constexpr int QK_VECS_PER_WARP = Dh_MAX / QK_VEC_SIZE; + + // The layout of the cache is [B, H, Dh/x, L, x] with x == 4/8 for FP32/FP16. Since each thread + // owns x elements, we have to decompose the linear index into chunks of x values and the posi- + // tion of the thread in that chunk. + + // The number of elements in a chunk of 16B (that's the x in the above formula). + constexpr int QK_ELTS_IN_16B = 16 / sizeof(T); + // The number of K vectors in 16B. + constexpr int QK_VECS_IN_16B = 16 / sizeof(Qk_vec); + + // The batch/beam idx + const int bi = blockIdx.y; + if (params.finished != nullptr && params.finished[bi] == true) { + return; + } + // The beam idx + const int beami = bi % params.beam_width; + // The "beam-aware" batch idx + const int bbi = bi / params.beam_width; + // The head. + const int num_kv_heads = params.num_kv_heads; + const int kv_rep = (params.num_heads / num_kv_heads); + const int hi = blockIdx.x; + const int hi_kv = hi / kv_rep; + + // Combine the batch and the head indices. + const int bhi = bi * params.num_heads + hi; + const int bhi_kv = bi * (params.num_heads / kv_rep) + hi_kv; + // Combine the "beam-aware" batch idx and the head indices. + const int bbhi = bbi * params.beam_width * params.num_heads + hi; + const int bbhi_kv = bbi * params.beam_width * (params.num_heads / kv_rep) + hi_kv; + // The thread in the block. + const int tidx = threadIdx.x; + + const bool handle_kv = !DO_CROSS_ATTENTION || (DO_CROSS_ATTENTION && params.timestep == 0); + // Every kv_rep threads have the same kv_cache values. So only the first one writes back. + const int write_kv_cache = handle_kv && (hi % kv_rep == 0); + + // While doing the product Q*K^T for the different keys we track the max. + float qk_max = -FLT_MAX; + + float qk = 0.0F; + + // int qkv_base_offset = (params.stride == 0) ? bhi * Dh : bi * params.stride + hi * Dh; + const int q_base_offset = bi * params.stride + hi * Dh; + const int k_base_offset = bi * params.stride + hi_kv * Dh; + const int v_base_offset = k_base_offset; + + const size_t bi_seq_len_offset = bi * params.memory_max_len; + + // int tlength = (DO_CROSS_ATTENTION)? params.memory_length_per_sample[bi] - 1 : params.timestep; + int tlength = (DO_CROSS_ATTENTION) ? params.memory_length_per_sample[bi] - 1 : + (params.length_per_sample == nullptr) ? + params.timestep : + params.length_per_sample[bi] + params.max_prefix_prompt_length; + const int first_step = max(0, tlength + 1 - params.memory_max_len); + const int tlength_circ = tlength % params.memory_max_len; + + // First QK_VECS_PER_WARP load Q and K + the bias values for the current timestep. + const bool is_masked = tidx >= QK_VECS_PER_WARP; + + // The offset in the Q and K buffer also accounts for the batch. + // int qk_offset = qkv_base_offset + tidx * QK_VEC_SIZE; + int q_offset = q_base_offset + tidx * QK_VEC_SIZE; + int k_offset = k_base_offset + tidx * QK_VEC_SIZE; + int v_offset = k_offset; + + // The offset in the bias buffer. + // int qk_bias_offset = hi * Dh + tidx * QK_VEC_SIZE; + int q_bias_offset = hi * Dh + tidx * QK_VEC_SIZE; + int k_bias_offset = hi_kv * Dh + tidx * QK_VEC_SIZE; + int v_bias_offset = k_bias_offset; + + const bool do_ia3 = handle_kv && params.ia3_tasks != nullptr; + const int ia3_task_id = do_ia3 ? params.ia3_tasks[bbi] : 0; + + // Trigger the loads from the Q and K buffers. + Qk_vec q; + zero(q); + if (!is_masked && (Dh == Dh_MAX || tidx * QK_VEC_SIZE < Dh)) { + if (params.int8_mode == 2) { + using Packed_Int8_t = typename packed_type::value>::type; + using Packed_Float_t = typename packed_type::value>::type; + const auto q_scaling = params.qkv_scale_out[0]; + const auto q_quant = + *reinterpret_cast(&reinterpret_cast(params.q)[q_offset]); + + convert_from_float(q, mul(q_scaling, float_from_int8(q_quant))); + } + else { + q = *reinterpret_cast(¶ms.q[q_offset]); + } + } + + Qk_vec k; + zero(k); + if (DO_CROSS_ATTENTION) { + // The 16B chunk written by the thread. + int co = tidx / QK_VECS_IN_16B; + // The position of the thread in that 16B chunk. + int ci = tidx % QK_VECS_IN_16B * QK_VEC_SIZE; + + // Two chunks are separated by L * x elements. A thread write QK_VEC_SIZE elements. + int offset = bhi_kv * params.memory_max_len * Dh + co * params.memory_max_len * QK_ELTS_IN_16B + + // params.timestep*QK_ELTS_IN_16B + + tlength * QK_ELTS_IN_16B + ci; + k = !is_masked && (Dh == Dh_MAX || tidx * QK_VEC_SIZE < Dh) ? + *reinterpret_cast(¶ms.k_cache[offset]) : + k; + } + else { + if (!is_masked && (Dh == Dh_MAX || tidx * QK_VEC_SIZE < Dh)) { + if (params.int8_mode == 2) { + using Packed_Int8_t = typename packed_type::value>::type; + using Packed_Float_t = typename packed_type::value>::type; + const auto k_scaling = params.qkv_scale_out[1]; + const auto k_quant = + *reinterpret_cast(&reinterpret_cast(params.k)[k_offset]); + + convert_from_float(k, mul(k_scaling, float_from_int8(k_quant))); + } + else { + k = *reinterpret_cast(¶ms.k[k_offset]); + } + } + } + + // Trigger the loads from the Q and K bias buffers. + Qk_vec q_bias; + zero(q_bias); + q_bias = (!is_masked && Dh == Dh_MAX || tidx * QK_VEC_SIZE < Dh) && params.q_bias != nullptr ? + *reinterpret_cast(¶ms.q_bias[q_bias_offset]) : + q_bias; + + Qk_vec k_bias; + zero(k_bias); + if (handle_kv) { + k_bias = !is_masked && (Dh == Dh_MAX || tidx * QK_VEC_SIZE < Dh) && params.k_bias != nullptr ? + *reinterpret_cast(¶ms.k_bias[k_bias_offset]) : + k_bias; + } + + // Computes the Q/K values with bias. + q = add(q, q_bias); + if (handle_kv) { + k = add(k, k_bias); + } + if (do_ia3 && !is_masked) { + k = mul( + k, + *reinterpret_cast( + ¶ms.ia3_key_weights[(ia3_task_id * params.num_heads + hi) * Dh + tidx * QK_VEC_SIZE])); + } + + // Padded len + const int padd_len = (params.total_padding_tokens == nullptr) ? 0 : params.total_padding_tokens[bi]; + if (params.rotary_embedding_dim > 0 && !params.neox_rotary_style) { + if (handle_kv) { + apply_rotary_embedding(q, k, tidx, params.rotary_embedding_dim, tlength - padd_len, params.rotary_base, params.rotary_scale); + } + else { + apply_rotary_embedding(q, tidx, params.rotary_embedding_dim, tlength - padd_len, params.rotary_base, params.rotary_scale); + } + } + else if (params.rotary_embedding_dim > 0 && params.neox_rotary_style) { + const bool do_rotary = !is_masked && QK_VEC_SIZE * tidx < params.rotary_embedding_dim; + + T* q_smem = reinterpret_cast(smem_); + T* k_smem = q_smem + params.rotary_embedding_dim; + + const int half_rotary_dim = params.rotary_embedding_dim / 2; + const int half_idx = (tidx * QK_VEC_SIZE) / half_rotary_dim; + const int intra_half_idx = (tidx * QK_VEC_SIZE) % half_rotary_dim; + const int smem_pitch = half_rotary_dim; // TODO: adjust for bank conflicts + + assert(half_rotary_dim % QK_VEC_SIZE == 0); + + if (do_rotary) { + *reinterpret_cast(q_smem + half_idx * smem_pitch + intra_half_idx) = q; + + if (handle_kv) { + *reinterpret_cast(k_smem + half_idx * smem_pitch + intra_half_idx) = k; + } + } + + __syncthreads(); + + const int transpose_idx = half_idx * (half_rotary_dim / 2) + intra_half_idx / 2; + constexpr int tidx_factor = (QK_VEC_SIZE > 1) ? QK_VEC_SIZE / 2 : 1; + if (do_rotary) { + mmha::vec_from_smem_transpose(q, q_smem, transpose_idx, smem_pitch); + + if (handle_kv) { + mmha::vec_from_smem_transpose(k, k_smem, transpose_idx, smem_pitch); + + mmha::apply_rotary_embedding( + q, k, transpose_idx / tidx_factor, params.rotary_embedding_dim, tlength - padd_len, params.rotary_base, params.rotary_scale); + + mmha::write_smem_transpose(k, k_smem, transpose_idx, smem_pitch); + } + else { + mmha::apply_rotary_embedding( + q, transpose_idx / tidx_factor, params.rotary_embedding_dim, tlength, params.rotary_base, params.rotary_scale); + } + mmha::write_smem_transpose(q, q_smem, transpose_idx, smem_pitch); + } + + __syncthreads(); + + if (do_rotary) { + q = *reinterpret_cast(q_smem + half_idx * smem_pitch + intra_half_idx); + if (handle_kv) { + k = *reinterpret_cast(k_smem + half_idx * smem_pitch + intra_half_idx); + } + } + + __syncthreads(); + } + + if (!is_masked) { + // Store the Q values to shared memory. + *reinterpret_cast(&q_smem[tidx * QK_VEC_SIZE]) = q; + + // Store Dh values of k_bias into smem, since will need to add later + // if params.timestep == 0 + if (DO_CROSS_ATTENTION && params.timestep == 0) { + *reinterpret_cast(&bias_smem[tidx * QK_VEC_SIZE]) = k_bias; + } + + // Write the K values to the global memory cache. + // + // NOTE: The stores are uncoalesced as we have multiple chunks of 16B spread across the memory + // system. We designed it this way as it allows much better memory loads (and there are many + // more loads) + the stores are really "write and forget" since we won't need the ack before + // the end of the kernel. There's plenty of time for the transactions to complete. + + // The 16B chunk written by the thread. + int co = tidx / QK_VECS_IN_16B; + // The position of the thread in that 16B chunk. + int ci = tidx % QK_VECS_IN_16B * QK_VEC_SIZE; + + // Two chunks are separated by L * x elements. A thread write QK_VEC_SIZE elements. + int offset = bhi_kv * params.memory_max_len * Dh + co * params.memory_max_len * QK_ELTS_IN_16B + + // params.timestep*QK_ELTS_IN_16B + + tlength_circ * QK_ELTS_IN_16B + ci; + + if (write_kv_cache) { + // Trigger the stores to global memory. + if (Dh == Dh_MAX || co < Dh / QK_ELTS_IN_16B) { + *reinterpret_cast(¶ms.k_cache[offset]) = k; + } + } + + // Compute \sum_i Q[i] * K^T[i] for the current timestep. +#ifdef MMHA_USE_FP32_ACUM_FOR_FMA + using Qk_vec_acum = typename Qk_vec_acum_fp32_::Type; +#else + using Qk_vec_acum = Qk_vec; +#endif + qk = dot(q, k); + if (QK_VECS_PER_WARP <= WARP_SIZE) { +#pragma unroll + for (int mask = QK_VECS_PER_WARP / 2; mask >= 1; mask /= 2) { + qk += __shfl_xor_sync(shfl_mask(QK_VECS_PER_WARP), qk, mask); + } + } + } + + if (QK_VECS_PER_WARP > WARP_SIZE) { + constexpr int WARPS_PER_RED = (QK_VECS_PER_WARP + WARP_SIZE - 1) / WARP_SIZE; + qk = block_sum(&red_smem[WARPS_PER_RED], qk); + } + + // Store that value in shared memory. Keep the Q*K^T value in register for softmax. + if (tidx == 0) { + // Normalize qk. + qk *= params.inv_sqrt_dh; + if (params.relative_attention_bias != nullptr) { + // TODO (Haotian): check whether we should replace hi with hi_kv, + // although params.relative_attention_bias is usually not used. + qk = add(qk, + params.relative_attention_bias[hi * params.relative_attention_bias_stride + * params.relative_attention_bias_stride + + (tlength - padd_len) * params.relative_attention_bias_stride + + (tlength - padd_len)]); + } + // Add alibi positional encoding + // qk += (alibi_slope != 0) ? alibi_slope * (params.timestep - params.memory_max_len) : 0; + // We don't need to apply the linear position bias here since qi - ki = 0 yields the position bias 0. + + qk_max = qk; + qk_smem[tlength - first_step] = qk; + // qk_smem[params.timestep] = qk; + } + + // Make sure the data is in shared memory. + __syncthreads(); + + // The type of queries and keys for the math in the Q*K^T product. + using K_vec = typename K_vec_::Type; + // The number of elements per vector. + constexpr int K_VEC_SIZE = sizeof(K_vec) / sizeof(T); + // Make sure the hidden size per head is a multiple of the vector size. + static_assert(Dh_MAX % K_VEC_SIZE == 0, ""); + // The number of elements per thread. + constexpr int K_ELTS_PER_THREAD = Dh_MAX / THREADS_PER_KEY; + // The number of vectors per thread. + constexpr int K_VECS_PER_THREAD = K_ELTS_PER_THREAD / K_VEC_SIZE; + + // The position the first key loaded by each thread from the cache buffer (for this B * H). + int ko = tidx / THREADS_PER_KEY; + // The position of the thread in the chunk of keys. + int ki = tidx % THREADS_PER_KEY * K_VEC_SIZE; + + static_assert(Dh_MAX == THREADS_PER_KEY * K_VEC_SIZE * K_VECS_PER_THREAD); + + // Load the Q values from shared memory. The values are reused during the loop on K. + K_vec q_vec[K_VECS_PER_THREAD]; +#pragma unroll + for (int ii = 0; ii < K_VECS_PER_THREAD; ++ii) { + q_vec[ii] = *reinterpret_cast(&q_smem[ki + ii * THREADS_PER_KEY * K_VEC_SIZE]); + } + + K_vec k_bias_vec[DO_CROSS_ATTENTION ? K_VECS_PER_THREAD : 1]; + if (DO_CROSS_ATTENTION && params.timestep == 0) { +#pragma unroll + for (int ii = 0; ii < K_VECS_PER_THREAD; ++ii) { + k_bias_vec[ii] = *reinterpret_cast(&bias_smem[ki + ii * THREADS_PER_KEY * K_VEC_SIZE]); + } + } + + // The number of timesteps loaded per iteration. + constexpr int K_PER_ITER = THREADS_PER_BLOCK / THREADS_PER_KEY; + // The number of keys per warp. + constexpr int K_PER_WARP = WARP_SIZE / THREADS_PER_KEY; + + // The base pointer for the key in the cache buffer. + T* k_cache = ¶ms.k_cache[bhi_kv * params.memory_max_len * Dh + ki]; + // Base pointer for the beam's batch, before offsetting with indirection buffer + T* k_cache_batch = ¶ms.k_cache[bbhi_kv * params.memory_max_len * Dh + ki]; + + // Pick a number of keys to make sure all the threads of a warp enter (due to shfl_sync). + // int ti_end = div_up(params.timestep, K_PER_WARP) * K_PER_WARP; + int ti_end = div_up(tlength - first_step, K_PER_WARP) * K_PER_WARP + first_step; + + // prefix prompt length if has + const int prefix_prompt_length = (params.prefix_prompt_lengths == nullptr) ? 0 : params.prefix_prompt_lengths[bi]; + + // Iterate over the keys/timesteps to compute the various (Q*K^T)_{ti} values. + const bool has_beams = params.cache_indir != nullptr; + const int* beam_indices = has_beams ? ¶ms.cache_indir[bi_seq_len_offset] : nullptr; + + for (int ti = first_step + ko; ti < ti_end; ti += K_PER_ITER) { + const int ti_circ = ti % params.memory_max_len; + + // The keys loaded from the key cache. + K_vec k[K_VECS_PER_THREAD]; + K_vec k_vec_zero; + zero(k_vec_zero); +#pragma unroll + for (int ii = 0; ii < K_VECS_PER_THREAD; ++ii) { + int jj = ii * params.memory_max_len + ti_circ; + // if( ti < params.timestep ) { + const bool within_bounds = (Dh == Dh_MAX || jj * QK_ELTS_IN_16B < Dh * params.memory_max_len); + if (ti < tlength) { + if (!within_bounds) { + k[ii] = k_vec_zero; + } + else { + if (has_beams) { + const int beam_offset = beam_indices[ti_circ] * params.num_heads * params.memory_max_len * Dh; + k[ii] = *reinterpret_cast(&k_cache_batch[beam_offset + jj * QK_ELTS_IN_16B]); + } + else { + k[ii] = *reinterpret_cast(&k_cache_batch[jj * QK_ELTS_IN_16B]); + } + } + // add bias and update k_cache + if (DO_CROSS_ATTENTION && params.timestep == 0) { + k[ii] = add(k[ii], k_bias_vec[ii]); + + if (do_ia3) { + k[ii] = mul( + k[ii], + *reinterpret_cast( + ¶ms.ia3_key_weights[(ia3_task_id * params.num_heads + hi) * Dh + ki + + ii * THREADS_PER_KEY * K_VEC_SIZE])); + } + + if (Dh == Dh_MAX || jj * QK_ELTS_IN_16B < Dh * params.memory_max_len) { + *reinterpret_cast(&k_cache[jj * QK_ELTS_IN_16B]) = k[ii]; + } + } + } + } + + // Perform the dot product and normalize qk. + // + // WARNING: ALL THE THREADS OF A WARP MUST ENTER!!! + float qk = Qk_dot::dot(q_vec, k) * params.inv_sqrt_dh; + bool is_mask = (params.masked_tokens != nullptr) && params.masked_tokens[bi_seq_len_offset + ti]; + + // Store the product to shared memory. There's one qk value per timestep. Update the max. + // if( ti < params.timestep && tidx % THREADS_PER_KEY == 0 ) { + if (ti < tlength && tidx % THREADS_PER_KEY == 0) { + if (params.relative_attention_bias != nullptr) { + qk = add(qk, + params.relative_attention_bias[hi * params.relative_attention_bias_stride + * params.relative_attention_bias_stride + + tlength * params.relative_attention_bias_stride + ti]); + } + if (params.linear_bias_slopes != nullptr) { + // Apply the linear position bias: (ki - qi) * slope[hi]. + // The padding token locates between the input context and the generated tokens. + // We need to remove the number of padding tokens in the distance computation. + // ti : 0 1 2 3 4 5 6 7 8 9(tlength) + // token: i i i i p p p o o o where i=input, p=pad, o=output. + // e.g. ti = 2, dist = (9 - 3) - 2 = 4. + int max_context_length = params.max_prefix_prompt_length + params.max_input_length; + float dist = (ti < max_context_length ? ti + padd_len : ti) - tlength; + + qk += mul(params.linear_bias_slopes[hi], dist); + } + // Add alibi positional encoding + // qk += (alibi_slope != 0) ? alibi_slope * (params.timestep - params.memory_max_len) : 0; + qk_max = is_mask ? qk_max : fmaxf(qk_max, qk); + qk_smem[ti - first_step] = qk; + } + } + +// Perform the final reduction to compute the max inside each warp. +// +// NOTE: In a group of THREADS_PER_KEY threads, the leader already has the max value for the +// group so it's not needed to run the reduction inside the group (again). +#pragma unroll + for (int mask = WARP_SIZE / 2; mask >= THREADS_PER_KEY; mask /= 2) { + qk_max = fmaxf(qk_max, __shfl_xor_sync(uint32_t(-1), qk_max, mask)); + } + + // Decompose the thread index into warp and lane. + const int warp = tidx / WARP_SIZE; + const int lane = tidx % WARP_SIZE; + + // The warp leader writes the max to shared memory. + if (lane == 0) { + red_smem[warp] = qk_max; + } + + // Make sure the products are in shared memory. + __syncthreads(); + + // The warps finalize the reduction. + qk_max = lane < WARPS_PER_BLOCK ? red_smem[lane] : -FLT_MAX; +#pragma unroll + for (int mask = WARPS_PER_BLOCK / 2; mask >= 1; mask /= 2) { + qk_max = fmaxf(qk_max, __shfl_xor_sync(uint32_t(-1), qk_max, mask)); + } + + // Broadcast to all the threads in the warp. + qk_max = __shfl_sync(uint32_t(-1), qk_max, 0); + + // Compute the logits and start the sum. + float sum = 0.f; + // for( int ti = tidx; ti <= params.timestep; ti += THREADS_PER_BLOCK ) { + for (int ti = first_step + tidx; ti <= tlength; ti += THREADS_PER_BLOCK) { + bool is_mask = (params.masked_tokens != nullptr) && params.masked_tokens[bi_seq_len_offset + ti]; + float logit = is_mask ? 0.f : __expf(qk_smem[ti - first_step] - qk_max); + sum += logit; + qk_smem[ti - first_step] = logit; + } + + // Compute the sum. + sum = block_sum(&red_smem[WARPS_PER_BLOCK], sum); + + // Normalize the logits. + float inv_sum = __fdividef(1.f, sum + 1.e-6f); + // for( int ti = tidx; ti <= params.timestep; ti += THREADS_PER_BLOCK ) { + const size_t cross_attention_out_offset = + params.is_return_cross_attentions ? + bhi_kv * params.max_decoder_seq_len * params.memory_max_len + params.timestep * params.memory_max_len : + 0; + for (int ti = first_step + tidx; ti <= tlength; ti += THREADS_PER_BLOCK) { + float logit = qk_smem[ti - first_step] * inv_sum; + if (params.is_return_cross_attentions) { + params.cross_attention_out[cross_attention_out_offset + ti] = logit; + } + convert_from_float(logits_smem[ti - first_step], logit); + } + + // Put Values part below so we leverage __syncthreads + // from the previous step + + // The number of elements per vector. + constexpr int V_VEC_SIZE = Dh_MAX / THREADS_PER_VALUE; + // A vector of V elements for the current timestep. + using V_vec = typename V_vec_::Type; + + // The value computed by this thread. + int vo = tidx / THREADS_PER_VALUE; + // The hidden dimensions computed by this particular thread. + int vi = tidx % THREADS_PER_VALUE * V_VEC_SIZE; + + // The base pointer for the value in the cache buffer. + T* v_cache = ¶ms.v_cache[bhi_kv * params.memory_max_len * Dh + vi]; + // Base pointer for the beam's batch, before offsetting with indirection buffer + T* v_cache_batch = ¶ms.v_cache[bbhi_kv * params.memory_max_len * Dh + vi]; + + // The number of values processed per iteration of the loop. + constexpr int V_PER_ITER = THREADS_PER_BLOCK / THREADS_PER_VALUE; + + // One group of threads computes the product(s) for the current timestep. + V_vec v_bias; + zero(v_bias); + // if( vo == params.timestep % V_PER_ITER ) { + if (Dh == Dh_MAX || vi < Dh) { + if (handle_kv) { + if (vo == tlength % V_PER_ITER) { + // Trigger the loads from the V bias buffer. + if (params.v_bias != nullptr) { + v_bias = *reinterpret_cast(¶ms.v_bias[hi_kv * Dh + vi]); + } + if (DO_CROSS_ATTENTION) { + *reinterpret_cast(&bias_smem[vi]) = v_bias; + } + } + } + } + + // From previous, before values, step + // Also make sure the logits are in shared memory. + __syncthreads(); + + // Values continued +#ifdef MMHA_USE_FP32_ACUM_FOR_OUT + using V_vec_acum = typename V_vec_acum_fp32_::Type; +#else + using V_vec_acum = V_vec; +#endif + // The partial outputs computed by each thread. + V_vec_acum out; + zero(out); + + // Loop over the timesteps to compute the partial outputs. + // for( int ti = vo; ti < params.timestep; ti += V_PER_ITER ) { + if (Dh == Dh_MAX || vi < Dh) { + for (int ti = first_step + vo; ti < tlength; ti += V_PER_ITER) { + const int ti_circ = ti % params.memory_max_len; + + // Fetch offset based on cache_indir when beam sampling + const int beam_src = (params.cache_indir != nullptr) ? params.cache_indir[bi_seq_len_offset + ti_circ] : 0; + const int beam_offset = beam_src * params.num_heads * params.memory_max_len * Dh; + // Load the values from the cache. + V_vec v = *reinterpret_cast(&v_cache_batch[beam_offset + ti_circ * Dh]); + if (DO_CROSS_ATTENTION && params.timestep == 0) { + v = add(v, *reinterpret_cast(&bias_smem[vi])); + if (do_ia3) { + v = mul( + v, + *reinterpret_cast( + ¶ms.ia3_value_weights[(ia3_task_id * params.num_heads + hi) * Dh + vi])); + } + *reinterpret_cast(&v_cache[ti * Dh]) = v; + } + // Load the logits from shared memory. +#if defined(MMHA_USE_FP32_ACUM_FOR_LOGITS) + float logit = logits_smem[ti - first_step]; + out = fma(logit, cast_to_float(v), out); +#else + T logit = logits_smem[ti - first_step]; + + // Update the partial sums. + out = fma(logit, v, out); +#endif + } + } + + // One group of threads computes the product(s) for the current timestep. + // if( vo == params.timestep % V_PER_ITER ) { + if (vo == tlength % V_PER_ITER && (Dh == Dh_MAX || vi < Dh)) { + + V_vec v; + if (DO_CROSS_ATTENTION) { + v = *reinterpret_cast(&v_cache[tlength * Dh]); + } + else { + // Trigger the loads from the V buffer. + const auto v_offset = v_base_offset + vi; + if (params.int8_mode == 2) { + using Packed_Int8_t = typename packed_type::value>::type; + using Packed_Float_t = typename packed_type::value>::type; + const auto v_scaling = params.qkv_scale_out[2]; + const auto v_quant = + *reinterpret_cast(&reinterpret_cast(params.v)[v_offset]); + + convert_from_float(v, mul(v_scaling, float_from_int8(v_quant))); + } + else { + v = *reinterpret_cast(¶ms.v[v_offset]); + } + // Trigger the loads from the V bias buffer. + // V_vec v_bias = *reinterpret_cast(¶ms.v_bias[hi*Dh + vi]); + } + + // Compute the V values with bias. + v = add(v, v_bias); + if (write_kv_cache) { + + if (do_ia3) { + v = mul( + v, + *reinterpret_cast( + ¶ms.ia3_value_weights[(ia3_task_id * params.num_heads + hi) * Dh + vi])); + } + + // Store the values with bias back to global memory in the cache for V. + //*reinterpret_cast(&v_cache[params.timestep*Dh]) = v; + *reinterpret_cast(&v_cache[tlength_circ * Dh]) = v; + } + + // Initialize the output value with the current timestep. +#if defined(MMHA_USE_FP32_ACUM_FOR_LOGITS) + // out = fma(logits_smem[params.timestep], cast_to_float(v), out); + out = fma(logits_smem[tlength - first_step], cast_to_float(v), out); +#else + // out = fma(logits_smem[params.timestep], v, out); + out = fma(logits_smem[tlength - first_step], v, out); +#endif + } + + // Make sure we can start writing to shared memory. + __syncthreads(); + + // Run the final reduction amongst the different groups computing different partial outputs. + if (Dh == Dh_MAX || vi < Dh) { +#pragma unroll + for (int active_groups = V_PER_ITER; active_groups >= 2; active_groups /= 2) { + + // The midpoint in the number of active groups. + int midpoint = active_groups / 2; + + // The upper part of active threads store to shared memory. + if (vo >= midpoint && vo < active_groups && (Dh == Dh_MAX || vi < Dh)) { +#ifdef MMHA_USE_FP32_ACUM_FOR_OUT + convert_from_float(*reinterpret_cast(&out_smem[(vo - midpoint) * Dh + vi]), out); +#else + *reinterpret_cast(&out_smem[(vo - midpoint) * Dh + vi]) = out; +#endif + } + __syncthreads(); + + // The bottom warps update their values. + if (vo < midpoint && (Dh == Dh_MAX || vi < Dh)) { + out = add(*reinterpret_cast(&out_smem[vo * Dh + vi]), out); + } + __syncthreads(); + } + } + + // Output the final values. + if (vo == 0 && (Dh == Dh_MAX || vi < Dh)) { +#ifdef MMHA_USE_FP32_ACUM_FOR_OUT + if (params.int8_mode == 2) { + using Packed_Int8_t = typename packed_type::value>::type; + out = mul(*params.attention_out_scale, out); + *reinterpret_cast(&(reinterpret_cast(params.out)[bhi * Dh + vi])) = + cast_to_int8(out); + } + else { + convert_from_float(*reinterpret_cast(¶ms.out[bhi * Dh + vi]), out); + } +#else + // TODO: support int8_mode? + *reinterpret_cast(¶ms.out[bhi * Dh + vi]) = out; +#endif + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mmha + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +void mmha_launch_kernel(const KERNEL_PARAMS_TYPE& params, const cudaStream_t& stream); diff --git a/llm-awq/awq/kernels/csrc/attention/setup.py b/llm-awq/awq/kernels/csrc/attention/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..6c19b38b9cebf7a44634e31442ba93639cfbd72b --- /dev/null +++ b/llm-awq/awq/kernels/csrc/attention/setup.py @@ -0,0 +1,159 @@ +# Adapted from https://github.com/NVIDIA/apex/blob/master/setup.py +import sys +import warnings +import os +from packaging.version import parse, Version + +from setuptools import setup, find_packages +import subprocess + +import torch +from torch.utils.cpp_extension import ( + BuildExtension, + CppExtension, + CUDAExtension, + CUDA_HOME, +) + + +# ninja build does not work unless include_dirs are abs path +this_dir = os.path.dirname(os.path.abspath(__file__)) + + +def get_cuda_bare_metal_version(cuda_dir): + raw_output = subprocess.check_output( + [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True + ) + output = raw_output.split() + release_idx = output.index("release") + 1 + bare_metal_version = parse(output[release_idx].split(",")[0]) + + return raw_output, bare_metal_version + + +def check_cuda_torch_binary_vs_bare_metal(cuda_dir): + raw_output, bare_metal_version = get_cuda_bare_metal_version(cuda_dir) + torch_binary_version = parse(torch.version.cuda) + + print("\nCompiling cuda extensions with") + print(raw_output + "from " + cuda_dir + "/bin\n") + + if bare_metal_version != torch_binary_version: + raise RuntimeError( + "Cuda extensions are being compiled with a version of Cuda that does " + "not match the version used to compile Pytorch binaries. " + "Pytorch binaries were compiled with Cuda {}.\n".format(torch.version.cuda) + + "In some cases, a minor-version mismatch will not cause later errors: " + "https://github.com/NVIDIA/apex/pull/323#discussion_r287021798. " + "You can try commenting out this check (at your own risk)." + ) + + +def raise_if_cuda_home_none(global_option: str) -> None: + if CUDA_HOME is not None: + return + raise RuntimeError( + f"{global_option} was requested, but nvcc was not found. Are you sure your environment has nvcc available? " + "If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, " + "only images whose names contain 'devel' will provide nvcc." + ) + + +def append_nvcc_threads(nvcc_extra_args): + _, bare_metal_version = get_cuda_bare_metal_version(CUDA_HOME) + if bare_metal_version >= Version("11.2"): + return nvcc_extra_args + ["--threads", "4"] + return nvcc_extra_args + + +if not torch.cuda.is_available(): + # https://github.com/NVIDIA/apex/issues/486 + # Extension builds after https://github.com/pytorch/pytorch/pull/23408 attempt to query torch.cuda.get_device_capability(), + # which will fail if you are compiling in an environment without visible GPUs (e.g. during an nvidia-docker build command). + print( + "\nWarning: Torch did not find available GPUs on this system.\n", + "If your intention is to cross-compile, this is not an error.\n" + "By default, Apex will cross-compile for Pascal (compute capabilities 6.0, 6.1, 6.2),\n" + "Volta (compute capability 7.0), Turing (compute capability 7.5),\n" + "and, if the CUDA version is >= 11.0, Ampere (compute capability 8.0).\n" + "If you wish to cross-compile for a single specific architecture,\n" + 'export TORCH_CUDA_ARCH_LIST="compute capability" before running setup.py.\n', + ) + if os.environ.get("TORCH_CUDA_ARCH_LIST", None) is None and CUDA_HOME is not None: + _, bare_metal_version = get_cuda_bare_metal_version(CUDA_HOME) + if bare_metal_version >= Version("11.8"): + os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5;8.0;8.6;9.0" + elif bare_metal_version >= Version("11.1"): + os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5;8.0;8.6" + elif bare_metal_version == Version("11.0"): + os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5;8.0" + else: + os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5" + + +print("\n\ntorch.__version__ = {}\n\n".format(torch.__version__)) +TORCH_MAJOR = int(torch.__version__.split(".")[0]) +TORCH_MINOR = int(torch.__version__.split(".")[1]) + +cmdclass = {} +ext_modules = [] + +# Check, if ATen/CUDAGeneratorImpl.h is found, otherwise use ATen/cuda/CUDAGeneratorImpl.h +# See https://github.com/pytorch/pytorch/pull/70650 +generator_flag = [] +torch_dir = torch.__path__[0] +if os.path.exists(os.path.join(torch_dir, "include", "ATen", "CUDAGeneratorImpl.h")): + generator_flag = ["-DOLD_GENERATOR_PATH"] + +raise_if_cuda_home_none("--ft_attention") +# Check, if CUDA11 is installed for compute capability 8.0 +cc_flag = [] +_, bare_metal_version = get_cuda_bare_metal_version(CUDA_HOME) +if bare_metal_version < Version("11.0"): + raise RuntimeError("ft_attention is only supported on CUDA 11 and above") +cc_flag.append("-gencode") +cc_flag.append("arch=compute_70,code=sm_70") +cc_flag.append("-gencode") +cc_flag.append("arch=compute_80,code=sm_80") +if bare_metal_version >= Version("11.8"): + cc_flag.append("-gencode") + cc_flag.append("arch=compute_90,code=sm_90") + +ext_modules.append( + CUDAExtension( + name="ft_attention", + sources=[ + "ft_attention.cpp", + "decoder_masked_multihead_attention.cu", + ], + extra_compile_args={ + "cxx": ["-O3", "-DENABLE_BF16"] + generator_flag, + "nvcc": append_nvcc_threads( + [ + "-DENABLE_BF16", # TODO + "-O3", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_BFLOAT16_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + "-U__CUDA_NO_BFLOAT162_OPERATORS__", + "-U__CUDA_NO_BFLOAT162_CONVERSIONS__", + "--expt-relaxed-constexpr", + "--expt-extended-lambda", + "--use_fast_math", + ] + + generator_flag + + cc_flag + ), + }, + include_dirs=[this_dir], + ) +) + +setup( + name="ft_attention", + version="0.1", + description="Attention for single query from FasterTransformer", + ext_modules=ext_modules, + cmdclass={"build_ext": BuildExtension} if ext_modules else {}, +) diff --git a/llm-awq/awq/kernels/csrc/layernorm/layernorm.h b/llm-awq/awq/kernels/csrc/layernorm/layernorm.h new file mode 100644 index 0000000000000000000000000000000000000000..de43ccac688d65b540b8fc9838e9e3c44b1758a1 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/layernorm/layernorm.h @@ -0,0 +1,3 @@ +#include + +void layernorm_forward_cuda(torch::Tensor _input, torch::Tensor _gamma, torch::Tensor _out, float eps); diff --git a/llm-awq/awq/kernels/csrc/layernorm/reduction.cuh b/llm-awq/awq/kernels/csrc/layernorm/reduction.cuh new file mode 100644 index 0000000000000000000000000000000000000000..678160e8fdf5788757a82060bab1ca6b6f6d3baf --- /dev/null +++ b/llm-awq/awq/kernels/csrc/layernorm/reduction.cuh @@ -0,0 +1,82 @@ +/* + +Adapted from NVIDIA FasterTransformer: +https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/kernels/reduce_kernel_utils.cuh +*/ + +#pragma once +#include +#if ((__CUDACC_VER_MAJOR__ > 11) || (__CUDACC_VER_MAJOR__ == 11 && __CUDACC_VER_MINOR__ >= 0)) +#include +#else +#include +#endif +#include +#include +#include +#include + +static const float HALF_FLT_MAX = 65504.F; +#define FINAL_MASK 0xffffffff + + +template +inline __device__ T add(T a, T b) { + return a + b; +} + +template<> +inline __device__ half2 add(half2 a, half2 b) { + return __hadd2(a, b); +} + +template<> +inline __device__ half add(half a, half b) { + return __hadd(a, b); +} + +template +__inline__ __device__ T warpReduceSum(T val) +{ +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + val = add(val, __shfl_xor_sync(FINAL_MASK, val, mask, 32)); //__shfl_sync bf16 return float when sm < 80 + return val; +} + +/* Calculate the sum of all elements in a block */ +template +__inline__ __device__ T blockReduceSum(T val) +{ + static __shared__ T shared[32]; + int lane = threadIdx.x & 0x1f; + int wid = threadIdx.x >> 5; + + val = warpReduceSum(val); + + if (lane == 0) + shared[wid] = val; + + __syncthreads(); + + // Modify from blockDim.x << 5 to blockDim.x / 32. to prevent + // blockDim.x is not divided by 32 + val = (threadIdx.x < (blockDim.x / 32.f)) ? shared[lane] : (T)(0.0f); + val = warpReduceSum(val); + + return val; +} + + +template +__device__ __forceinline__ T clamp_inf_for_half(const float input) +{ + return input; +} + +template<> +__device__ __forceinline__ half clamp_inf_for_half(const float input) +{ + // clamp inf values to enable fp16 training + return input > 0.0f ? __float2half(min(input, HALF_FLT_MAX - 1000)) : __float2half(max(input, -HALF_FLT_MAX + 1000)); +} diff --git a/llm-awq/awq/kernels/csrc/position_embedding/pos_encoding_kernels.cu b/llm-awq/awq/kernels/csrc/position_embedding/pos_encoding_kernels.cu new file mode 100644 index 0000000000000000000000000000000000000000..883b59c41b74cf1947c1b1e1bacde716f0f21857 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/position_embedding/pos_encoding_kernels.cu @@ -0,0 +1,88 @@ +/* + +Adapted from the VLLM project: +https://github.com/vllm-project/vllm/blob/main/csrc/pos_encoding_kernels.cu + +*/ + +#include +#include +#include "pos_encoding.h" + +template +__global__ void rotary_embedding_neox_kernel( + const int64_t* __restrict__ positions, // [num_tokens] + scalar_t* __restrict__ query, // [num_tokens, num_heads, head_size] + scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size] + const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim // 2] + const int rot_dim, + const int stride, + const int num_heads, + const int head_size) { + // Each thread block is responsible for one token. + const int token_idx = blockIdx.x; + int64_t pos = positions[token_idx]; + const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim; + + const int embed_dim = rot_dim / 2; + const int n = num_heads * embed_dim; + for (int i = threadIdx.x; i < n; i += blockDim.x) { + const int head_idx = i / embed_dim; + const int token_head = token_idx * stride + head_idx * head_size; + + const int rot_offset = i % embed_dim; + const int x_index = rot_offset; + const int y_index = embed_dim + rot_offset; + + const int out_x = token_idx * stride + head_idx * head_size + x_index; + const int out_y = token_idx * stride + head_idx * head_size + y_index; + + const scalar_t cos = __ldg(cache_ptr + x_index); + const scalar_t sin = __ldg(cache_ptr + y_index); + + const scalar_t q_x = query[token_head + x_index]; + const scalar_t q_y = query[token_head + y_index]; + query[out_x] = q_x * cos - q_y * sin; + query[out_y] = q_y * cos + q_x * sin; + + const scalar_t k_x = key[token_head + x_index]; + const scalar_t k_y = key[token_head + y_index]; + key[out_x] = k_x * cos - k_y * sin; + key[out_y] = k_y * cos + k_x * sin; + } +} + +void rotary_embedding_neox( + torch::Tensor& positions, // [b, num_tokens] + torch::Tensor& query, // [b, num_tokens, 1, num_heads, head_size] + torch::Tensor& key, // [b, num_tokens, 1, num_heads, head_size] + int head_size, + torch::Tensor& cos_sin_cache) // [max_position, rot_dim] +{ + int num_tokens = query.size(0) * query.size(1); + int rot_dim = cos_sin_cache.size(1); + int num_heads = query.size(-2); + int stride = num_heads * head_size; + // TORCH_CHECK(stride == key.stride(0)); + + dim3 grid(num_tokens); + dim3 block(std::min(num_heads * rot_dim / 2, 512)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + query.scalar_type(), + "rotary_embedding_neox", + [&] { + rotary_embedding_neox_kernel<<>>( + positions.data_ptr(), + query.data_ptr(), + key.data_ptr(), + cos_sin_cache.data_ptr(), + rot_dim, + stride, + num_heads, + head_size); + }); +} + diff --git a/llm-awq/awq/kernels/csrc/quantization/dequantize.cuh b/llm-awq/awq/kernels/csrc/quantization/dequantize.cuh new file mode 100644 index 0000000000000000000000000000000000000000..5d333b35c148d0cd01c8cb6fffd5deb28db28d33 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization/dequantize.cuh @@ -0,0 +1,79 @@ +/* +Modified from NVIDIA FasterTransformer: https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h + +@article{lin2023awq, + title={AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration}, + author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Dang, Xingyu and Han, Song}, + journal={arXiv}, + year={2023} +} +*/ + +#pragma once + + +__device__ uint4 dequantize_s4_to_fp16x2(uint32_t const& source) +{ + uint4 result; + + uint32_t* h = reinterpret_cast(&result); + uint32_t const i4s = reinterpret_cast(source); + + // First, we extract the i4s and construct an intermediate fp16 number. + static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + static constexpr uint32_t BOTTOM_MASK = 0x000f000f; + static constexpr uint32_t TOP_MASK = 0x00f000f0; + static constexpr uint32_t I4s_TO_F16s_MAGIC_NUM = 0x64006400; + + // Note that the entire sequence only requires 1 shift instruction. This is thanks to the register packing + // format and the fact that we force our integers to be unsigned, and account for this in the fp16 subtractions. + // In addition, I exploit the fact that sub and fma have the same throughput in order to convert elt_23 and + // elt_67 to fp16 without having to shift them to the bottom bits before hand. + + // Shift right by 8 to now consider elt_45 and elt_67. Issue first to hide RAW dependency if we issue + // immediately before required. + const uint32_t top_i4s = i4s >> 8; + // Extract elt_01 - (i4s & 0x000f000f) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[0]) + : "r"(i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_23 (i4s & 0x00f000f0) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[1]) + : "r"(i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_45 (top_i4s & 0x000f000f) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[2]) + : "r"(top_i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_67 (top_i4s & 0x00f000f0) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[3]) + : "r"(top_i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + + // I use inline PTX below because I am not sure if the compiler will emit float2half instructions if I use the + // half2 ctor. In this case, I chose performance reliability over code readability. + + // This is the half2 {1032, 1032} represented as an integer. + // static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64086408; + // Haotian: subtract {1024, 1024} instead, we do not need to map to [-8, 7] + static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64006400; + // This is the half2 {1 / 16, 1 / 16} represented as an integer. + static constexpr uint32_t ONE_SIXTEENTH = 0x2c002c00; + // This is the half2 {-72, -72} represented as an integer. + // static constexpr uint32_t NEG_72 = 0xd480d480; + // Haotian: Let's use {-64, -64}. + static constexpr uint32_t NEG_64 = 0xd400d400; + + // Finally, we construct the output numbers. + // Convert elt_01 + asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(FP16_TOP_MAGIC_NUM)); + // Convert elt_23 + asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[1]) : "r"(h[1]), "r"(ONE_SIXTEENTH), "r"(NEG_64)); + // Convert elt_45 + asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[2]) : "r"(h[2]), "r"(FP16_TOP_MAGIC_NUM)); + // Convert elt_67 + asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[3]) : "r"(h[3]), "r"(ONE_SIXTEENTH), "r"(NEG_64)); + + return result; +} + diff --git a/llm-awq/awq/kernels/csrc/quantization/gemm_cuda.h b/llm-awq/awq/kernels/csrc/quantization/gemm_cuda.h new file mode 100644 index 0000000000000000000000000000000000000000..4c0846c84866b23a44dc9a640358bf54b160bec9 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization/gemm_cuda.h @@ -0,0 +1,4 @@ +#include + +torch::Tensor gemm_forward_cuda(torch::Tensor _in_feats, torch::Tensor _kernel, + torch::Tensor _scaling_factors, torch::Tensor _zeros, int group_size, int split_k_iters); diff --git a/llm-awq/awq/kernels/csrc/quantization/gemm_cuda_gen.cu b/llm-awq/awq/kernels/csrc/quantization/gemm_cuda_gen.cu new file mode 100644 index 0000000000000000000000000000000000000000..231220b203c81cd5a008b37d9e85d3eb4322e63d --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization/gemm_cuda_gen.cu @@ -0,0 +1,298 @@ +// Inspired by NVIDIA's FasterTransformer +/* + +@article{lin2023awq, + title={AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration}, + author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Dang, Xingyu and Han, Song}, + journal={arXiv}, + year={2023} +} + +*/ + +#include +#include "gemm_cuda.h" +#include +#include + + +// Pack two half values. +static inline __device__ __host__ unsigned +__pack_half2(const half x, const half y) { + unsigned v0 = *((unsigned short *)&x); + unsigned v1 = *((unsigned short *)&y); + return (v1 << 16) | v0; +} + +__device__ __forceinline__ int make_divisible(int c, int divisor){ + return (c + divisor - 1) / divisor; +} + +template +__global__ void __launch_bounds__(128) gemm_forward_4bit_cuda_m128n64k32(int split_k_iters, half* __restrict__ A, int* __restrict__ B, half* __restrict__ scaling_factors, int* zeros, int M, int IC, int OC, half* __restrict__ C) +{ + static constexpr uint32_t ZERO = 0x0; + float C_warp[64]; + __shared__ half A_shared[128 * (32 + 8)]; + __shared__ half B_shared[64 * (32 + 8)]; + + // __shared__ half scaling_factors_shared[64]; + // __shared__ half zeros_shared[64]; + + int j_factors1 = ((OC + 64 - 1) / 64); + + int blockIdx_x = 0; + int blockIdx_y = blockIdx.x % ((M + 128 - 1) / 128 * j_factors1); + int blockIdx_z = blockIdx.x / ((M + 128 - 1) / 128 * j_factors1); + + half A_shared_warp[32]; + half B_shared_warp[16]; + for (int i_0_3_init = 0; i_0_3_init < 4; ++i_0_3_init) { + for (int j_0_4_init = 0; j_0_4_init < 2; ++j_0_4_init) { + for (int i = 0; i < 8; ++i) { + C_warp[((i_0_3_init * 16) + (j_0_4_init * 8)) + i] = 0.0; + } + } + } + + static constexpr int row_stride_warp = 32 * 8 / 32; + static constexpr int row_stride_A = 4 * 32 * 8 / 32; + static constexpr int row_stride = 4 * 32 * 8 / 32; + const int make_divisible_multipler = 128 / G; + const int zeros_w = make_divisible(make_divisible(IC / G, 8), make_divisible_multipler) * make_divisible_multipler; + const int sf_w = zeros_w * 8; + + bool ld_zero_flag = (threadIdx.y * 32 + threadIdx.x) * 8 < 64; + int ld_A_row = (blockIdx_y / j_factors1 * 128 + threadIdx.y * row_stride_warp + threadIdx.x * 8 / 32); // threadIdx.y is warp_id + // bool wb_C_flag = (threadIdx.x / 4) < M; + + half* A_ptr = A + + (((int)blockIdx_y) / j_factors1 * 128 + (((int)threadIdx.y) * row_stride_warp) + ((int)threadIdx.x) / (32 / 8)) * IC + + (((int)threadIdx.x) % (32 / 8)) * 8; + + int* B_ptr = B + + ((int)threadIdx.y) * (IC / 8) * 8 + + (((int)threadIdx.x) / (32 / 8)) * (IC / 8) + + (((int)blockIdx_y) % j_factors1) * 64 * (IC / 8) + + (((int)threadIdx.x) % (32 / 8)) * 1; + +// Why * 1 in the above line? + + half* A_shared_ptr = A_shared + + ((int)threadIdx.y) * row_stride_warp * (32 + 8) + + (((int)threadIdx.x) / (32 / 8)) * (32 + 8) + + (((int)threadIdx.x) % (32 / 8) ) * 8; + + half* B_shared_ptr = B_shared + + ((int)threadIdx.y) * (row_stride / 4) * (32 + 8) + + (((int)threadIdx.x) / (32 / 8)) * (32 + 8) + + (((int)threadIdx.x) % (32 / 8)) * 8; + + + int* zeros_ptr = zeros + + ((int)threadIdx.y) * zeros_w * 8 + + (((int)threadIdx.x) / (32 / 8)) * zeros_w + + (((int)blockIdx_y) % j_factors1) * 64 * zeros_w + // this term is zero + + (((int)threadIdx.x) % (32 / 8)) / G ; + + half* scaling_factors_ptr = scaling_factors + + ((int)threadIdx.y) * sf_w * 8 + + (((int)threadIdx.x) / (32 / 8)) * sf_w + + (((int)blockIdx_y) % j_factors1) * (64) * sf_w + // this term is zero + + (((int)threadIdx.x) % (32 / 8)) * 8 / G; + + + // Haotian: TBD, check, May 29 11:46 AM PST + half* C_ptr = C + + blockIdx_z * M * OC // blockIdx_z -> split_k dim + + (((int)blockIdx_y) % j_factors1) * 64 + + (((int)threadIdx.y) / 2) * 32 + + (((int)threadIdx.x) % 4) * 2; + + // preload s.f. and zeros + int k_bound = make_divisible(IC / 32, split_k_iters); // (IC / 32 + split_k_iters - 1) / split_k_iters; + if ((k_bound - 1) * 32 + blockIdx_z >= IC) k_bound -= 1; + + // TODO (Haotian): load scales and zero points to smem + + for (int _k_0_0 = 0; _k_0_0 < k_bound; ++_k_0_0) { + int k_0_0 = _k_0_0 * split_k_iters + blockIdx_z; + __syncthreads(); + // TODO: Haotian: Here we assume M % cta_M = 0. + for (int ax0_ax1_fused_0 = 0; ax0_ax1_fused_0 < 4; ++ax0_ax1_fused_0) + { + if (ld_A_row + ax0_ax1_fused_0 * row_stride_A < M) + { + *(uint4*)(A_shared_ptr + ax0_ax1_fused_0 * row_stride_A * 40) = *(uint4*)(A_ptr + (ax0_ax1_fused_0 * row_stride_A * IC) + (k_0_0 * 32)); + } + else + { + *(uint4*)(A_shared_ptr + ax0_ax1_fused_0 * row_stride_A * 40) = make_uint4(0, 0, 0, 0); + } + } + + + int* zeros_ptr_local = zeros_ptr + k_0_0 * 32 / G / 8; + half* scaling_factors_ptr_local = scaling_factors_ptr + k_0_0 * 32 / G; + + // uint4 B_loaded_scale = make_uint4(0, 0, 0, 0); + int* B_ptr_local = B_ptr + k_0_0 * (32 / 8); + + for (int ax0_ax1_fused_0 = 0; ax0_ax1_fused_0 < 2; ++ax0_ax1_fused_0) { + + // B: 32 x 136 (128+8) float16 + // each warp: 32 x 4 + // each thr: read 32 bit -> convert to 8xFP16 (a UINT4) -> scale and minus zero -> WB UINT4 + // row stride in shared memory: (NWARPS * 32 * 8 / cta_N) + int B_loaded_current = *(B_ptr_local + ax0_ax1_fused_0 * row_stride * (IC / 8)); + int zeros_loaded = *(zeros_ptr_local + ax0_ax1_fused_0 * row_stride * zeros_w); + zeros_loaded >>= ((k_0_0 * 32 / G) % 8) * 4; + float current_zeros = (float)(zeros_loaded & 0xF); + half scaling_factors_loaded = *(scaling_factors_ptr_local + ax0_ax1_fused_0 * row_stride * sf_w); + half B_loaded_fp16[8]; + #pragma unroll + for (int ic_1 = 0; ic_1 < 8; ic_1++){ + float current_single_weight_fp = (float)(B_loaded_current & 0xF); + half dequantized_weight = __float2half(__half2float(scaling_factors_loaded) * (current_single_weight_fp - current_zeros)); + B_loaded_current = B_loaded_current >> 4; + B_loaded_fp16[ic_1] = dequantized_weight; + } + // write back + *(uint4*)(B_shared_ptr + ax0_ax1_fused_0 * row_stride * (32 + 8)) = *reinterpret_cast(B_loaded_fp16); + } + __syncthreads(); + for (int k_0_1 = 0; k_0_1 < 2; ++k_0_1) { + for (int ax0_0 = 0; ax0_0 < 4; ++ax0_0) { + { + unsigned int addr; + __asm__ __volatile__( + "{ .reg .u64 addr; cvta.to.shared.u64 addr, %1; cvt.u32.u64 %0, addr; }\n" + : "=r"(addr) + : "l"((void *)((&(A_shared[((((((int)threadIdx.y) & 1) * 2560) + (ax0_0 * 640)) + (k_0_1 * 16))])) + (((((int)threadIdx.x) & 15) * 40) + ((((int)threadIdx.x) >> 4) * 8)))) + ); + __asm__ __volatile__( + "ldmatrix.sync.aligned.m8n8.x4.shared.b16" + "{%0, %1, %2, %3}, [%4];\n" + : "=r"(((unsigned *)(A_shared_warp + (ax0_0 * 8)))[0]), "=r"(((unsigned *)(A_shared_warp + (ax0_0 * 8)))[1]), "=r"(((unsigned *)(A_shared_warp + (ax0_0 * 8)))[2]), "=r"(((unsigned *)(A_shared_warp + (ax0_0 * 8)))[3]) + : "r"(addr) + ); + } + } + + for (int ax0_0_1 = 0; ax0_0_1 < 2; ++ax0_0_1) { + { + unsigned int addr; + __asm__ __volatile__( + "{ .reg .u64 addr; cvta.to.shared.u64 addr, %1; cvt.u32.u64 %0, addr; }\n" + : "=r"(addr) + : "l"((void *)((&(B_shared[((((((int)threadIdx.y) >> 1) * 1280) + (ax0_0_1 * 640)) + (k_0_1 * 16))])) + ((((((int)threadIdx.x) >> 4) * 320) + ((((int)threadIdx.x) & 7) * 40)) + (((((int)threadIdx.x) & 15) >> 3) * 8)))) + ); + __asm__ __volatile__( + "ldmatrix.sync.aligned.m8n8.x4.shared.b16" + "{%0, %1, %2, %3}, [%4];\n" + : "=r"(((unsigned *)(B_shared_warp + (ax0_0_1 * 8)))[0]), "=r"(((unsigned *)(B_shared_warp + (ax0_0_1 * 8)))[1]), "=r"(((unsigned *)(B_shared_warp + (ax0_0_1 * 8)))[2]), "=r"(((unsigned *)(B_shared_warp + (ax0_0_1 * 8)))[3]) + : "r"(addr) + ); + } + } + + for (int i_0_3 = 0; i_0_3 < 4; ++i_0_3) { + for (int j_0_4 = 0; j_0_4 < 2; ++j_0_4) { + + { + __asm__ __volatile__( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32" + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};\n" + : "=f"(((float *)(C_warp + ((i_0_3 * 16) + (j_0_4 * 8))))[0]), "=f"(((float *)(C_warp + ((i_0_3 * 16) + (j_0_4 * 8))))[1]), "=f"(((float *)(C_warp + ((i_0_3 * 16) + (j_0_4 * 8))))[2]), "=f"(((float *)(C_warp + ((i_0_3 * 16) + (j_0_4 * 8))))[3]) + : "r"(((unsigned *)(A_shared_warp + (i_0_3 * 8)))[0]), "r"(((unsigned *)(A_shared_warp + (i_0_3 * 8)))[1]), "r"(((unsigned *)(A_shared_warp + (i_0_3 * 8)))[2]), "r"(((unsigned *)(A_shared_warp + (i_0_3 * 8)))[3]), "r"(((unsigned *)(B_shared_warp + (j_0_4 * 8)))[0]), "r"(((unsigned *)(B_shared_warp + (j_0_4 * 8)))[1]), "f"(((float *)(C_warp + ((i_0_3 * 16) + (j_0_4 * 8))))[0]), "f"(((float *)(C_warp + ((i_0_3 * 16) + (j_0_4 * 8))))[1]), "f"(((float *)(C_warp + ((i_0_3 * 16) + (j_0_4 * 8))))[2]), "f"(((float *)(C_warp + ((i_0_3 * 16) + (j_0_4 * 8))))[3])); + } + + { + __asm__ __volatile__( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32" + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};\n" + : "=f"(((float *)(C_warp + (((i_0_3 * 16) + (j_0_4 * 8)) + 4)))[0]), "=f"(((float *)(C_warp + (((i_0_3 * 16) + (j_0_4 * 8)) + 4)))[1]), "=f"(((float *)(C_warp + (((i_0_3 * 16) + (j_0_4 * 8)) + 4)))[2]), "=f"(((float *)(C_warp + (((i_0_3 * 16) + (j_0_4 * 8)) + 4)))[3]) + : "r"(((unsigned *)(A_shared_warp + (i_0_3 * 8)))[0]), "r"(((unsigned *)(A_shared_warp + (i_0_3 * 8)))[1]), "r"(((unsigned *)(A_shared_warp + (i_0_3 * 8)))[2]), "r"(((unsigned *)(A_shared_warp + (i_0_3 * 8)))[3]), "r"(((unsigned *)(B_shared_warp + ((j_0_4 * 8) + 4)))[0]), "r"(((unsigned *)(B_shared_warp + ((j_0_4 * 8) + 4)))[1]), "f"(((float *)(C_warp + (((i_0_3 * 16) + (j_0_4 * 8)) + 4)))[0]), "f"(((float *)(C_warp + (((i_0_3 * 16) + (j_0_4 * 8)) + 4)))[1]), "f"(((float *)(C_warp + (((i_0_3 * 16) + (j_0_4 * 8)) + 4)))[2]), "f"(((float *)(C_warp + (((i_0_3 * 16) + (j_0_4 * 8)) + 4)))[3])); + } + } + } + } + } + +// Haotian: Here (May 29 11:46AM PST) +// TODO: Shang: Hoist loop invariance. + for (int ax0_0_2 = 0; ax0_0_2 < 4; ++ax0_0_2) { + for (int ax1_0 = 0; ax1_0 < 2; ++ax1_0) { + for (int local_id = 0; local_id < 8; ++local_id) { + int row_offset = (((int)blockIdx_y) / j_factors1) * 128 + (threadIdx.y % 2) * 64 + ax0_0_2 * 16 + (local_id % 4) / 2 * 8 + ((int)threadIdx.x) / 4; + if (row_offset < M) + { + *(C_ptr + ax1_0 * 16 + row_offset * OC + (local_id / 4) * 8 + local_id % 2) = __float2half(C_warp[(ax0_0_2 * 16) + (ax1_0 * 8) + local_id]); + } + } + } + } +} + +// in_feats: M, IC [float16] +// kernel: IC, OC // 8 [int32] -> cast to IC, OC [uint4b] +// scaling_factors: IC // G, OC [float16] +// zeros: IC // G, OC // 8 [int32] -> cast to IC // G, OC [uint4b] +// assume that batch_size < 16 for now + +torch::Tensor gemm_forward_cuda( + torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _scaling_factors, + torch::Tensor _zeros, + int group_size, + int split_k_iters) +{ + int num_in_feats = _in_feats.size(0); + int num_in_channels = _in_feats.size(1); + const at::cuda::OptionalCUDAGuard device_guard(device_of(_in_feats)); + + auto options = torch::TensorOptions().dtype(_in_feats.dtype()).device(_in_feats.device()); + // for int4, need _kernel.size(1) * 8 + at::Tensor _out_feats = torch::empty({split_k_iters, num_in_feats, _kernel.size(0)}, options); + int num_out_feats = _out_feats.size(-2); + int num_out_channels = _out_feats.size(-1); + + auto in_feats = reinterpret_cast(_in_feats.data_ptr()); + auto kernel = reinterpret_cast(_kernel.data_ptr()); + auto out_feats = reinterpret_cast(_out_feats.data_ptr()); + auto scaling_factors = reinterpret_cast(_scaling_factors.data_ptr()); + auto zeros = reinterpret_cast(_zeros.data_ptr()); + + // blockIdx_x: i_factors[0] * j_factors[0] + // blockIdx_y: i_factors[1] * j_factors[1] + + if (num_out_channels % 64 != 0) + throw std::invalid_argument("OC is not multiple of cta_N = 64"); + if (num_out_channels % 8 != 0) + throw std::invalid_argument("OC is not multiple of pack_num = 8"); + int j_factors1 = num_out_channels / 64 / 1; + dim3 num_blocks((num_out_feats + 128 - 1) / 128 * j_factors1 * split_k_iters); + + // threadIdx.x: 32 + // threadIdx.y: i_factors[2] * j_factors[2] + dim3 threads_per_block(32, 4); + if (group_size == 128) + { + gemm_forward_4bit_cuda_m128n64k32<128><<>>( + split_k_iters, in_feats, kernel, scaling_factors, zeros, num_in_feats, num_in_channels, num_out_channels, out_feats); + } + else if (group_size == 64) + { + gemm_forward_4bit_cuda_m128n64k32<64><<>>( + split_k_iters, in_feats, kernel, scaling_factors, zeros, num_in_feats, num_in_channels, num_out_channels, out_feats); + } + else + { + throw std::invalid_argument("Group size temporarily not supported."); + } + return _out_feats.sum(0); +} + diff --git a/llm-awq/awq/kernels/csrc/quantization/gemv_cuda.h b/llm-awq/awq/kernels/csrc/quantization/gemv_cuda.h new file mode 100644 index 0000000000000000000000000000000000000000..748abc5d1bcef5cca0c056c6e1e2279c3bb262a9 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization/gemv_cuda.h @@ -0,0 +1,9 @@ +#pragma once +#include + +torch::Tensor gemv_forward_cuda( + torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _scaling_factors, + torch::Tensor _zeros, + int group_size); diff --git a/llm-awq/awq/kernels/csrc/quantization_new/dequantize.cuh b/llm-awq/awq/kernels/csrc/quantization_new/dequantize.cuh new file mode 100644 index 0000000000000000000000000000000000000000..9917ec94911ca106844858ea855f4fc15902a2a7 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization_new/dequantize.cuh @@ -0,0 +1,123 @@ +/* +Modified from NVIDIA FasterTransformer: https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h + +@article{lin2023awq, + title={AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration}, + author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Dang, Xingyu and Han, Song}, + journal={arXiv}, + year={2023} +} +*/ +#include +#include +#pragma once + +template +__inline__ __device__ void dequantize_s4_to_fp16x2(half2 const &source, uint4 *result); + +template <> +__inline__ __device__ void dequantize_s4_to_fp16x2(half2 const &source, uint4 *result) +{ + // uint4 result; + + uint32_t *h = reinterpret_cast(result); + uint32_t const i4s = reinterpret_cast(source); + + // First, we extract the i4s and construct an intermediate fp16 number. + constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + constexpr uint32_t BOTTOM_MASK = 0x000f000f; + constexpr uint32_t TOP_MASK = 0x00f000f0; + constexpr uint32_t I4s_TO_F16s_MAGIC_NUM = 0x64006400; + + // Note that the entire sequence only requires 1 shift instruction. This is thanks to the register packing + // format and the fact that we force our integers to be unsigned, and account for this in the fp16 subtractions. + // In addition, I exploit the fact that sub and fma have the same throughput in order to convert elt_23 and + // elt_67 to fp16 without having to shift them to the bottom bits before hand. + + // Shift right by 8 to now consider elt_45 and elt_67. Issue first to hide RAW dependency if we issue + // immediately before required. + const uint32_t top_i4s = i4s >> 8; + // Extract elt_01 - (i4s & 0x000f000f) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[0]) + : "r"(i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_23 (i4s & 0x00f000f0) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[1]) + : "r"(i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_45 (top_i4s & 0x000f000f) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[2]) + : "r"(top_i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_67 (top_i4s & 0x00f000f0) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[3]) + : "r"(top_i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + + // I use inline PTX below because I am not sure if the compiler will emit float2half instructions if I use the + // half2 ctor. In this case, I chose performance reliability over code readability. + + // This is the half2 {1032, 1032} represented as an integer. + // static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64086408; + // Haotian: subtract {1024, 1024} instead, we do not need to map to [-8, 7] + static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64006400; + // This is the half2 {1 / 16, 1 / 16} represented as an integer. + static constexpr uint32_t ONE_SIXTEENTH = 0x2c002c00; + // This is the half2 {-72, -72} represented as an integer. + // static constexpr uint32_t NEG_72 = 0xd480d480; + // Haotian: Let's use {-64, -64}. + static constexpr uint32_t NEG_64 = 0xd400d400; + + // Finally, we construct the output numbers. + // Convert elt_01 + asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(FP16_TOP_MAGIC_NUM)); + // Convert elt_23 + asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[1]) : "r"(h[1]), "r"(ONE_SIXTEENTH), "r"(NEG_64)); + // Convert elt_45 + asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[2]) : "r"(h[2]), "r"(FP16_TOP_MAGIC_NUM)); + // Convert elt_67 + asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[3]) : "r"(h[3]), "r"(ONE_SIXTEENTH), "r"(NEG_64)); +} + +template <> +__inline__ __device__ void dequantize_s4_to_fp16x2(half2 const &source, uint4 *result) +{ + // uint4 result; + + uint32_t *h = reinterpret_cast(result); + uint32_t const i4s = reinterpret_cast(source); + + // First, we extract the i4s and construct an intermediate bf16 number. + constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + constexpr uint32_t BOTTOM_MASK = 0x000f000f; + constexpr uint32_t I4s_TO_BF16s_MAGIC_NUM = 0x43004300; + + // Shift right by 4, 8, 12 to consider elt_23, elt_45 and elt_67. + const uint32_t i4s1 = i4s >> 4; + const uint32_t i4s2 = i4s >> 8; + const uint32_t i4s3 = i4s >> 12; + // Extract elt_01 - (i4s & 0x000f000f) | 0x43004300 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[0]) + : "r"(i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_BF16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_23 - (i4s & 0x000f000f) | 0x43004300 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[1]) + : "r"(i4s1), "n"(BOTTOM_MASK), "n"(I4s_TO_BF16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_45 - (i4s & 0x000f000f) | 0x43004300 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[2]) + : "r"(i4s2), "n"(BOTTOM_MASK), "n"(I4s_TO_BF16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_67 - (i4s & 0x000f000f) | 0x43004300 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[3]) + : "r"(i4s3), "n"(BOTTOM_MASK), "n"(I4s_TO_BF16s_MAGIC_NUM), "n"(immLut)); + + // This is the nv_bfloat162 {128, 128} represented as an integer + static constexpr uint32_t BF16_TOP_MAGIC_NUM = 0x43004300; + + reinterpret_cast<__nv_bfloat162*>(h)[0] = __hsub2(reinterpret_cast<__nv_bfloat162*>(h)[0], reinterpret_cast(BF16_TOP_MAGIC_NUM)); + reinterpret_cast<__nv_bfloat162*>(h)[1] = __hsub2(reinterpret_cast<__nv_bfloat162*>(h)[1], reinterpret_cast(BF16_TOP_MAGIC_NUM)); + reinterpret_cast<__nv_bfloat162*>(h)[2] = __hsub2(reinterpret_cast<__nv_bfloat162*>(h)[2], reinterpret_cast(BF16_TOP_MAGIC_NUM)); + reinterpret_cast<__nv_bfloat162*>(h)[3] = __hsub2(reinterpret_cast<__nv_bfloat162*>(h)[3], reinterpret_cast(BF16_TOP_MAGIC_NUM)); +} \ No newline at end of file diff --git a/llm-awq/awq/kernels/csrc/quantization_new/dispatch_utils.cuh b/llm-awq/awq/kernels/csrc/quantization_new/dispatch_utils.cuh new file mode 100644 index 0000000000000000000000000000000000000000..cb7773838cad0c767e6bb4aca1668ee2cf7e2154 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization_new/dispatch_utils.cuh @@ -0,0 +1,18 @@ +#pragma once +#include +#include +#include +#include + +#define DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(pytorch_dtype, c_type, ...) \ + if (pytorch_dtype == at::ScalarType::Half) { \ + using c_type = half; \ + __VA_ARGS__ \ + } else if (pytorch_dtype == at::ScalarType::BFloat16) { \ + using c_type = nv_bfloat16; \ + __VA_ARGS__ \ + } else { \ + std::ostringstream oss; \ + oss << __PRETTY_FUNCTION__ << " failed to dispatch data type " << pytorch_dtype; \ + TORCH_CHECK(false, oss.str()); \ + } diff --git a/llm-awq/awq/kernels/csrc/quantization_new/gemm/gemm_cuda.cu b/llm-awq/awq/kernels/csrc/quantization_new/gemm/gemm_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..d4d297da5a919d1cc7871a57b49924fdc36b39b9 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization_new/gemm/gemm_cuda.cu @@ -0,0 +1,1236 @@ +#include +#include "semaphore.h" +#include "gemm_cuda.h" +#include "../dequantize.cuh" +#include "../dispatch_utils.cuh" +#include +#include + +#define kInterleave 4 +#define OP_M 16 +#define OP_N 8 +#define OP_K 16 +#define INTRIN_M 16 +#define INTRIN_N 16 +#define INTRIN_K 16 +#define WARP_SIZE 32 +#define SMEM_PAD_A 0 +#define SMEM_PAD_B 0 +#define PACK_SIZE 8 +#if (__CUDACC_VER_MAJOR__ >= 11) && (__CUDACC_VER_MINOR__ >= 4) +#define L2_CACHEHINT(size) ".L2::" #size "B" +#else +#define L2_CACHEHINT(size) +#endif + +#define KERNEL_LAUNCH_CODE \ + int num_mn_tiles = (num_in_feats + CTA_M - 1) / CTA_M * (num_out_channels + CTA_N - 1) / CTA_N; \ + torch::Tensor _semaphores = torch::empty({num_mn_tiles}, options_int); \ + auto semaphores = reinterpret_cast(_semaphores.data_ptr()); \ + constexpr int NUM_WARPS = (CTA_M / WARP_M) * (CTA_N / WARP_N) * (CTA_K / WARP_K); \ + constexpr int SCALES_SMEM_SIZE = (G >= CTA_K) ? (CTA_N / (G / CTA_K) * STAGES * 2) : (CTA_N * (CTA_K / G) * STAGES * 2); \ + constexpr int kSmemByteSize = (CTA_M * (CTA_K + SMEM_PAD_A) + CTA_N * (CTA_K + SMEM_PAD_B) / kInterleave + SCALES_SMEM_SIZE) * STAGES * sizeof(ctype); \ + if (kSmemByteSize >= 99 * 1024) \ + { \ + printf("This kernel requires %d Bytes of shared memory, which exceeds device limit.\n", kSmemByteSize); \ + return _out_feats; \ + } \ + int j_factors1 = num_out_channels / CTA_N / 1; \ + dim3 num_blocks((num_out_feats + CTA_M - 1) / CTA_M * j_factors1 * SPLITK); \ + dim3 threads_per_block(WARP_SIZE, NUM_WARPS); \ + auto kernel_func = gemm_w4a16_T1; \ + cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemByteSize); \ + kernel_func<<>>( \ + in_feats, kernel, scales, zeros, out_feats, semaphores, num_in_feats, num_out_channels, num_in_channels); + +template +__inline__ __host__ __device__ int get_log_tile(int n) +{ + if (N >= 8 && n >= 6) + return 3; + else if (N >= 4 && n >= 3) + return 2; + else if (N >= 2 && n >= 2) + return 1; + else + return 0; +} + +__inline__ __device__ uint2 get_block_idx_mapping(int blockIdx_x, int blockIdx_y, int log_tile) +{ + return make_uint2((blockIdx_x >> log_tile), (blockIdx_y << log_tile) + ((blockIdx_x) & ((1 << (log_tile)) - 1))); +} + +template +__device__ void sync_slice(int slice_id) +{ + if constexpr (SLICES == 1) + { + __syncthreads(); + } + else + { + constexpr int SLICE_GROUP = (SLICES + 7) / 8; + constexpr uint32_t num_threads = NUM_WARPS_MN * WARP_SIZE; + const uint32_t barrier_id = slice_id / SLICE_GROUP + 1; + asm volatile("bar.sync %0, %1;" : : "r"(barrier_id), "n"(num_threads)); + } +} + +__inline__ __device__ uint32_t cast_smem_ptr_to_uint(void const *const ptr) +{ + uint32_t smem_int_ptr; + + asm("{.reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }\n" + : "=r"(smem_int_ptr) + : "l"(ptr)); + + return smem_int_ptr; +} + +template +__inline__ __device__ void ldmatrix_m8n8_x4_b16(T *shared_warp, int ax0_0, uint32_t addr) +{ + __asm__ __volatile__( + "ldmatrix.sync.aligned.m8n8.x4.shared.b16" + "{%0, %1, %2, %3}, [%4];" + : "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[0]), "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[1]), "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[2]), "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[3]) + : "r"(addr)); +} + +template +__inline__ __device__ void ldmatrix_m8n8_x4_trans_b16(T *shared_warp, int ax0_0, uint32_t addr) +{ + __asm__ __volatile__( + "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16" + "{%0, %1, %2, %3}, [%4];" + : "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[0]), "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[1]), "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[2]), "=r"(((unsigned *)(shared_warp + (ax0_0 * 8)))[3]) + : "r"(addr)); +} + +__inline__ __device__ void cp_async_cg_A(uint32_t smem_int_ptr, const uint4 *__restrict__ src, bool mask) +{ + const int cp_size = 16; + asm volatile("{" + " .reg .pred p;" + " setp.ne.b32 p, %0, 0;" + " @p cp.async.cg.shared.global" L2_CACHEHINT(128) " [%1], [%2], %3;" + "}" ::"r"((int)mask), + "r"(smem_int_ptr), + "l"(src), + "n"(cp_size)); +} + +__device__ __inline__ void mma_m16n8k16_f16f16f16(half *C_warp, half *A_shared_warp, half *B_shared_warp) +{ + __asm__ __volatile__( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16" + "{%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%8, %9};" + : "=r"(((unsigned *)C_warp)[0]), "=r"(((unsigned *)C_warp)[1]) + : "r"(((unsigned *)A_shared_warp)[0]), "r"(((unsigned *)A_shared_warp)[1]), "r"(((unsigned *)A_shared_warp)[2]), "r"(((unsigned *)A_shared_warp)[3]), "r"(((unsigned *)B_shared_warp)[0]), "r"(((unsigned *)B_shared_warp)[1]), "r"(((unsigned *)C_warp)[0]), "r"(((unsigned *)C_warp)[1])); +} + +__device__ __inline__ void mma_m16n8k16_bf16bf16f32(float *C_warp, nv_bfloat16 *A_shared_warp, nv_bfloat16 *B_shared_warp) +{ + + __asm__ __volatile__( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32" + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};" + : "=f"(C_warp[0]), "=f"(C_warp[1]), "=f"(C_warp[2]), "=f"(C_warp[3]) + : "r"(((unsigned *)A_shared_warp)[0]), "r"(((unsigned *)A_shared_warp)[1]), "r"(((unsigned *)A_shared_warp)[2]), "r"(((unsigned *)A_shared_warp)[3]), "r"(((unsigned *)B_shared_warp)[0]), "r"(((unsigned *)B_shared_warp)[1]), "f"(C_warp[0]), "f"(C_warp[1]), "f"(C_warp[2]), "f"(C_warp[3])); +} + +template +__device__ __inline__ void global_to_share_one_stage_A(T *src, T *dst, int global_nrows, int global_ncols, int cta_offset_m, int cta_offset_n, int cta_offset_k, int global_iter_k, int shared_iter_k, bool mask) +{ + constexpr int threads_needed = (CTA_M * CTA_K) / PACK_SIZE / SHARED_K_ITERS; + constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE; + constexpr int total_global_iters = (CTA_M * CTA_K) / PACK_SIZE / threads_used; + constexpr int partial_global_iters = (total_global_iters + SHARED_K_ITERS - 1) / SHARED_K_ITERS; + constexpr int cta_step_m_or_n = (threads_used * PACK_SIZE) / CTA_K; + constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int threads_per_row = CTA_K / PACK_SIZE; + constexpr int kSmemCol = CTA_K + SMEM_PAD_A; + bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used); + int ld_col = (threadIdx.x % threads_per_row); +#pragma unroll + for (int _global_iter = 0; _global_iter < partial_global_iters; ++_global_iter) + { + int global_iter = shared_iter_k * partial_global_iters + _global_iter; + int ld_row = global_iter * cta_step_m_or_n + threadIdx.y * warp_step_m_or_n + (threadIdx.x / threads_per_row); + int ld_col_swizzled = (ld_col ^ (ld_row) & 7) * PACK_SIZE; + void *dst_ptr = (void *)(dst + ld_row * kSmemCol + ld_col_swizzled); + uint4 *src_ptr = (uint4 *)(src + (ld_row + cta_offset_m) * global_ncols + ld_col * PACK_SIZE + global_iter_k * CTA_K + cta_offset_k); // cta_offset_m * global_ncols + global_iter * cta_step_m_or_n * global_ncols + threadIdx.y * warp_step_m_or_n * global_ncols + (threadIdx.x / threads_per_row) * global_ncols + global_iter_k * CTA_K + (threadIdx.x % threads_per_row) * PACK_SIZE); + if constexpr (STAGES > 1) + { + uint32_t addr = cast_smem_ptr_to_uint(dst_ptr); + cp_async_cg_A(addr, src_ptr, local_mask & (ld_row + cta_offset_m < global_nrows)); + } + else + { + if (local_mask & (ld_row + cta_offset_m < global_nrows)) + *(uint4 *)dst_ptr = *src_ptr; + } + } +} + +template +__device__ __inline__ void global_to_share_one_stage_B(T *src, T *dst, int global_ncols, int cta_offset_m, int cta_offset_n, int cta_offset_k, int global_iter_k, int shared_iter_k, bool mask) +{ + constexpr int threads_needed = (CTA_N / kInterleave * CTA_K) / PACK_SIZE / SHARED_K_ITERS; + constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE; + constexpr int total_global_iters = (CTA_N / kInterleave * CTA_K) / PACK_SIZE / threads_used; + constexpr int partial_global_iters = (total_global_iters + SHARED_K_ITERS - 1) / SHARED_K_ITERS; + constexpr int cta_step_m_or_n = (threads_used * PACK_SIZE) / CTA_K; + constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int threads_per_row = CTA_K / PACK_SIZE; + constexpr int kSmemCol = CTA_K + SMEM_PAD_B; + bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used); +#pragma unroll + for (int _global_iter = 0; _global_iter < partial_global_iters; ++_global_iter) + { + int global_iter = shared_iter_k * partial_global_iters + _global_iter; + + int ld_row = global_iter * cta_step_m_or_n + threadIdx.y * warp_step_m_or_n + (threadIdx.x / threads_per_row); + int ld_col = (threadIdx.x % threads_per_row); + int ld_col_swizzled = ld_col ^ (ld_row % 2) & 7; + void *dst_ptr = (void *)(dst + (ld_row * kSmemCol + ld_col_swizzled * PACK_SIZE)); + uint4 *src_ptr = (uint4 *)(src + global_iter_k * CTA_K + cta_offset_n / kInterleave * global_ncols + ld_row * global_ncols + ld_col * PACK_SIZE + cta_offset_k); + if constexpr (STAGES > 1) + { + uint32_t addr = cast_smem_ptr_to_uint(dst_ptr); + cp_async_cg_A(addr, src_ptr, local_mask); + } + else + { + if (local_mask) + *(uint4 *)dst_ptr = *src_ptr; + } + } +} + +template +__device__ __inline__ void global_to_share_one_stage_scales(T *src, T *dst, T *src_z, T *dst_z, int global_ncols, int cta_offset_m, int cta_offset_n, int cta_offset_k, int global_iter_k, int shared_iter_k, bool mask) +{ + constexpr int LD_AMOUNT = (G >= CTA_K) ? CTA_N : CTA_N * CTA_K / G; + constexpr int threads_needed = LD_AMOUNT / PACK_SIZE / 1; + constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE; + constexpr int total_global_iters = LD_AMOUNT / PACK_SIZE / threads_used; + constexpr int threads_per_row = CTA_N / PACK_SIZE; + constexpr int kSmemCol = CTA_N; + bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used); + int g_idx = (cta_offset_k + global_iter_k * CTA_K) / G; + + void *dst_ptr = (void *)(dst + (threadIdx.x / threads_per_row) * kSmemCol + (threadIdx.x % threads_per_row) * PACK_SIZE); + uint4 *src_ptr = (uint4 *)(src + g_idx * global_ncols + cta_offset_n + (threadIdx.x / threads_per_row) * global_ncols + (threadIdx.x % threads_per_row) * PACK_SIZE); + void *dst_ptr_z = (void *)(dst_z + (threadIdx.x / threads_per_row) * kSmemCol + (threadIdx.x % threads_per_row) * PACK_SIZE); + uint4 *src_ptr_z = (uint4 *)(src_z + g_idx * global_ncols + cta_offset_n + (threadIdx.x / threads_per_row) * global_ncols + (threadIdx.x % threads_per_row) * PACK_SIZE); + if (STAGES > 1) + { + uint32_t addr = cast_smem_ptr_to_uint(dst_ptr); + cp_async_cg_A(addr, src_ptr, local_mask); + uint32_t addr_z = cast_smem_ptr_to_uint(dst_ptr_z); + cp_async_cg_A(addr_z, src_ptr_z, local_mask); + } + else + { + if (local_mask) + { + *(uint4 *)dst_ptr = *src_ptr; + *(uint4 *)dst_ptr_z = *src_ptr_z; + } + } +} + +template +__device__ __inline__ void share_to_reg_one_stage_A(T *src, T *dst, int warp_offset_m, int warp_offset_n, int warp_offset_k, int k_0_1) +{ + constexpr int kSmemCol = CTA_K + SMEM_PAD_A; + + for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) + { + + int ld_row = warp_offset_m + shared_iter * OP_M + (threadIdx.x % 16); + int ld_col = k_0_1 * 16 + (threadIdx.x / 16) * 8 + warp_offset_k; + int ld_col_swizzled = ((ld_col / PACK_SIZE) ^ (ld_row) & 7) * PACK_SIZE; + void *addr_ptr = (void *)(src + ld_row * kSmemCol + ld_col_swizzled); + + uint32_t addr = cast_smem_ptr_to_uint(addr_ptr); + ldmatrix_m8n8_x4_b16(dst, shared_iter, addr); + } +} + +template +__device__ __inline__ void share_to_reg_one_stage_B(T *src, T *src_scales, T *src_zeros, T *dst, T *dst_fp16, int warp_offset_m, int warp_offset_n, int warp_offset_k, int k_0_1) +{ + using T2 = typename std::conditional::value, half2, nv_bfloat162>::type; + constexpr int kSmemCol = CTA_K + SMEM_PAD_B; + int r0 = ((threadIdx.x / 8 / 2) * 8 + threadIdx.x % 8); + int c0 = ((threadIdx.x / 8) % 2) * 8; + int r = r0 / 4; + int c = (r0 % 4) * 16 + c0; + int c_swizzled = ((c / PACK_SIZE) ^ (r % 2) & 7) * PACK_SIZE; + + if constexpr (ldmatrix) + { +#pragma unroll + for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) + { + void *addr_ptr = (void *)(src + warp_offset_n / kInterleave * kSmemCol + shared_iter * 16 / kInterleave * kSmemCol + k_0_1 * 16 + r * kSmemCol + c_swizzled + warp_offset_k); + uint32_t addr = cast_smem_ptr_to_uint(addr_ptr); + ldmatrix_m8n8_x4_b16(dst, shared_iter, addr); + } + } + +#pragma unroll + for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) + { + T scale = src_scales[(warp_offset_k / G) * CTA_N + warp_offset_n + 16 * shared_iter + 8 * (k_0_1 % 2) + threadIdx.x / 4]; + T zero = src_zeros[(warp_offset_k / G) * CTA_N + warp_offset_n + 16 * shared_iter + 8 * (k_0_1 % 2) + threadIdx.x / 4]; + T2 scale2, zero2; + if constexpr (std::is_same::value) + { + scale2 = __half2half2(scale); + zero2 = __half2half2(zero); + } + else + { + scale2 = __bfloat162bfloat162(scale); + zero2 = __bfloat162bfloat162(zero); + } + T2 loaded[4]; + dequantize_s4_to_fp16x2(*reinterpret_cast(dst + (k_0_1 % 2) * 4 + (k_0_1 / 2 * 2) + shared_iter * 8), reinterpret_cast(loaded)); +#pragma unroll + for (int i = 0; i < 4; i++) + { + loaded[i] = __hfma2(loaded[i], scale2, zero2); + } + *reinterpret_cast(dst_fp16 + shared_iter * 16 + 8 * (k_0_1 % 2)) = *reinterpret_cast(loaded); + } +} + +template +__global__ void gemm_w4a16_T1(T *__restrict__ A, T *__restrict__ B, T *__restrict__ scales, T *__restrict__ zeros, T *__restrict__ C, int *__restrict__ semaphores, int M, int N, int K) +{ + using DTypeAccum = typename std::conditional::value, half, float>::type; + constexpr int NUM_WARPS_MN = CTA_M / WARP_M * CTA_N / WARP_N; + constexpr int NUM_WARPS = NUM_WARPS_MN * CTA_K / WARP_K; + constexpr int CTA_SIZE = NUM_WARPS * WARP_SIZE; + constexpr int CTA_SIZE_MN = NUM_WARPS_MN * WARP_SIZE; + constexpr int SLICES = CTA_K / WARP_K; + int num_blocks_n = (N + CTA_N - 1) / CTA_N; + int num_blocks_m = (M + CTA_M - 1) / CTA_M; + int blockIdx_x = 0; + int blockIdx_y = blockIdx.x % (num_blocks_m * num_blocks_n); + int blockIdx_z = blockIdx.x / (num_blocks_m * num_blocks_n); + const int log_tile = get_log_tile<1>((N + CTA_N - 1) / CTA_N); + int blockIdx_m = blockIdx_y / (num_blocks_n >> log_tile); + int blockIdx_n = blockIdx_y % (num_blocks_n >> log_tile); + const uint2 block_idx_mapping = get_block_idx_mapping(blockIdx_m, blockIdx_n, log_tile); + blockIdx_m = block_idx_mapping.x; + blockIdx_n = block_idx_mapping.y; + + DTypeAccum C_warp[CTA_M * CTA_N / CTA_SIZE_MN]; + constexpr int kSmemPadKA = CTA_K + SMEM_PAD_A; + constexpr int kSmemPadKB = CTA_K + SMEM_PAD_B; + constexpr int kSmemSizeAPerStage = CTA_M * kSmemPadKA; + constexpr int kSmemSizeBPerStage = CTA_N / kInterleave * kSmemPadKB; + constexpr int kSmemSizeA = kSmemSizeAPerStage * STAGES; + constexpr int kSmemSizeB = kSmemSizeBPerStage * STAGES; + constexpr int scales_load_interval = G >= CTA_K ? G / CTA_K : 1; + constexpr int scales_per_load = G < CTA_K ? CTA_K / G : 1; + constexpr int kSmemSizeScales = CTA_N * STAGES / scales_load_interval * scales_per_load; + constexpr int kSmemSizeZeros = CTA_N * STAGES / scales_load_interval * scales_per_load; + extern __shared__ half mem_shared[]; + T *A_shared = (T*)mem_shared; + T *B_shared = (T*)mem_shared + kSmemSizeA; + T *scales_shared = (T*)mem_shared + kSmemSizeA + kSmemSizeB; + T *zeros_shared = (T*)mem_shared + kSmemSizeA + kSmemSizeB + kSmemSizeScales; + T *C_shared = (T*)(mem_shared); + T A_shared_warp_[2][WARP_M * INTRIN_K / + WARP_SIZE]; + T B_shared_warp_[2][WARP_N * 32 / + WARP_SIZE]; + T B_shared_warp_tmp_[2][WARP_N * 16 / + WARP_SIZE]; + int cta_offset_m = blockIdx_m * CTA_M; + int cta_offset_n = blockIdx_n * CTA_N; + int cta_offset_k = blockIdx_z * (K / SPLITK); + int warp_mn = threadIdx.y % NUM_WARPS_MN; + int slice_id = threadIdx.y / NUM_WARPS_MN; + int warp_offset_n = (warp_mn % (CTA_N / WARP_N)) * WARP_N; + int warp_offset_m = (warp_mn / (CTA_N / WARP_N)) * WARP_M; + int warp_offset_k = slice_id * WARP_K; + + for (int i = 0; i < CTA_M * CTA_N / CTA_SIZE_MN; i++) + C_warp[i] = 0.0; + + int gemm_iters = (K + CTA_K - 1) / CTA_K / SPLITK; + int k_0_0_ld = 0; + int k_0_0 = 0; + constexpr int prologue_stages = STAGES == 1 ? 1 : STAGES - 1; +#pragma unroll + for (k_0_0_ld = 0; k_0_0_ld < prologue_stages; ++k_0_0_ld) + { + global_to_share_one_stage_A(A, A_shared + k_0_0_ld * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, 0, true); + global_to_share_one_stage_B(B, B_shared + k_0_0_ld * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, 0, true); + global_to_share_one_stage_scales( + scales, scales_shared + (k_0_0_ld / scales_load_interval * scales_per_load) * CTA_N, + zeros, zeros_shared + (k_0_0_ld / scales_load_interval * scales_per_load) * CTA_N, + N, cta_offset_m, cta_offset_n, cta_offset_k, + k_0_0_ld, 0, k_0_0_ld < gemm_iters && k_0_0_ld % scales_load_interval == 0); + if constexpr (STAGES > 1) + __pipeline_commit(); + } + if constexpr (STAGES > 1) + __pipeline_wait_prior(STAGES - 2); + __syncthreads(); + + share_to_reg_one_stage_A(A_shared, A_shared_warp_[0], warp_offset_m, warp_offset_n, warp_offset_k, 0); + share_to_reg_one_stage_B(B_shared, scales_shared, zeros_shared, B_shared_warp_tmp_[0], B_shared_warp_[0], warp_offset_m, warp_offset_n, warp_offset_k, 0); + constexpr int SHARED_K_ITERS = WARP_K / INTRIN_K; + + for (; k_0_0 < gemm_iters; ++k_0_0, ++k_0_0_ld) + { + int ld_stage = k_0_0_ld % STAGES; + int compute_stage = k_0_0 % STAGES; + T *A_shared_this_compute_stage; + T *B_shared_this_compute_stage; + T *scales_shared_this_compute_stage; + T *zeros_shared_this_compute_stage; + +#pragma unroll + for (int iter_k = 0; iter_k < SHARED_K_ITERS; ++iter_k) + { + A_shared_this_compute_stage = A_shared + compute_stage * kSmemSizeAPerStage; + B_shared_this_compute_stage = B_shared + compute_stage * kSmemSizeBPerStage; + scales_shared_this_compute_stage = scales_shared + (compute_stage / scales_load_interval * scales_per_load) * CTA_N; + zeros_shared_this_compute_stage = zeros_shared + (compute_stage / scales_load_interval * scales_per_load) * CTA_N; + share_to_reg_one_stage_A(A_shared_this_compute_stage, A_shared_warp_[(iter_k + 1) % 2], warp_offset_m, warp_offset_n, warp_offset_k, (iter_k + 1) % SHARED_K_ITERS); + if ((iter_k + 1) % kInterleave == 0) + { + if (compute_stage % 2 == 1) + { + share_to_reg_one_stage_B( + B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage, + B_shared_warp_tmp_[1], B_shared_warp_[((iter_k + 1) / 2) % 2], + warp_offset_m, warp_offset_n, warp_offset_k, (iter_k + 1) % SHARED_K_ITERS); + } + else + { + share_to_reg_one_stage_B( + B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage, + B_shared_warp_tmp_[0], B_shared_warp_[((iter_k + 1) / 2) % 2], + warp_offset_m, warp_offset_n, warp_offset_k, (iter_k + 1) % SHARED_K_ITERS); + } + } + else + { + if (compute_stage % 2 == 1) + { + share_to_reg_one_stage_B( + B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage, + B_shared_warp_tmp_[1], B_shared_warp_[((iter_k + 1) / 2) % 2], + warp_offset_m, warp_offset_n, warp_offset_k, (iter_k + 1) % SHARED_K_ITERS); + } + else + { + share_to_reg_one_stage_B( + B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage, + B_shared_warp_tmp_[0], B_shared_warp_[((iter_k + 1) / 2) % 2], + warp_offset_m, warp_offset_n, warp_offset_k, (iter_k + 1) % SHARED_K_ITERS); + } + } + T *A_shared_warp = A_shared_warp_[iter_k % 2]; + T *B_shared_warp = B_shared_warp_[(iter_k / 2) % 2]; + + for (int i_0_3 = 0; i_0_3 < WARP_M / INTRIN_M; ++i_0_3) + { + for (int j_0_4 = 0; j_0_4 < WARP_N / INTRIN_N; ++j_0_4) + { + if constexpr (std::is_same::value) + { + mma_m16n8k16_f16f16f16(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4); + mma_m16n8k16_f16f16f16(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8 + 4, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4 + 8); + } + else + { + mma_m16n8k16_bf16bf16f32(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4); + mma_m16n8k16_bf16bf16f32(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8 + 4, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4 + 8); + } + } + } + + if (iter_k < WARP_K / INTRIN_K - 1) + { + if constexpr (STAGES == 1) + __syncthreads(); + global_to_share_one_stage_A(A, A_shared + ld_stage * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, iter_k, k_0_0_ld < gemm_iters); + global_to_share_one_stage_B(B, B_shared + ld_stage * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, iter_k, k_0_0_ld < gemm_iters); + } + + if (iter_k == WARP_K / INTRIN_K - 2) + { + if constexpr (STAGES == 1 && WARP_K / INTRIN_K > 2) + { + __syncthreads(); + } + global_to_share_one_stage_A(A, A_shared + ld_stage * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, iter_k + 1, k_0_0_ld < gemm_iters); + global_to_share_one_stage_B(B, B_shared + ld_stage * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, iter_k + 1, k_0_0_ld < gemm_iters); + global_to_share_one_stage_scales( + scales, scales_shared + (ld_stage / scales_load_interval * scales_per_load) * CTA_N, + zeros, zeros_shared + (ld_stage / scales_load_interval * scales_per_load) * CTA_N, + N, cta_offset_m, cta_offset_n, cta_offset_k, + k_0_0_ld, iter_k, k_0_0_ld < gemm_iters && k_0_0_ld % scales_load_interval == 0); + if constexpr (STAGES > 1) + { + __pipeline_commit(); + __pipeline_wait_prior(STAGES - 2); + } + compute_stage = (k_0_0 + 1) % STAGES; + __syncthreads(); + } + } + } + __pipeline_commit(); + __pipeline_wait_prior(0); + __syncthreads(); + + if constexpr (std::is_same::value) + { + if constexpr (SLICES > 1) + { + #pragma unroll + for (int z = 0; z < SLICES; ++z) + { + if (slice_id == z) + { + #pragma unroll + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + #pragma unroll + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + #pragma unroll + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) + { + if (z > 0) + { + C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] += C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2]; + } + C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2] = C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id]; + }; + } + } + } + __syncthreads(); + } + if (slice_id == 0) + { + #pragma unroll + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + #pragma unroll + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + #pragma unroll + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) + { + C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] = C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2]; + }; + } + } + } + } + + if (slice_id == 0) + { + Semaphore semaphore(semaphores + blockIdx_y, threadIdx.x); + + if constexpr (SPLITK > 1) + { + semaphore.fetch(); + } + + if (blockIdx_z != 0) + { + semaphore.wait(blockIdx_z); + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2) + { + int write_row = cta_offset_m + warp_offset_m + ax0_0_1 * OP_M + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)); + + if (write_row < M) + { + half2 *existing_psum_ptr = reinterpret_cast( + C + write_row * N + + cta_offset_n + warp_offset_n + ax1_0_1 * 16 + + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2); + + *existing_psum_ptr = __hadd2(*existing_psum_ptr, + *reinterpret_cast(C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 + + ax1_0_1 * 8 + local_id)); + } + }; + } + } + } + else + { + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2) + { + int write_row = cta_offset_m + warp_offset_m + ax0_0_1 * OP_M + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)); + if (write_row < M) + { + *reinterpret_cast( + C + write_row * N + + cta_offset_n + warp_offset_n + ax1_0_1 * 16 + + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2) = + *reinterpret_cast(C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 + + ax1_0_1 * 8 + local_id); + } + }; + } + } + } + + if constexpr (SPLITK > 1) + { + + int lock = 0; + if (SPLITK == blockIdx_z + 1) + { + + lock = 0; + } + else + { + lock = blockIdx_z + 1; + } + semaphore.release(lock); + } + } + } + else + { + // first convert fp32 to bf16 + nv_bfloat16 C_warp16[CTA_M * CTA_N / CTA_SIZE_MN]; +#pragma unroll + for (int i = 0; i < CTA_M * CTA_N / CTA_SIZE_MN / 2; ++i) + { + ((nv_bfloat162*)C_warp16)[i] = __float22bfloat162_rn(((float2*)C_warp)[i]); + } + + // the following is the same as fp16. Maybe there is a neat way to implement this. + if constexpr (SLICES > 1) + { +#pragma unroll + for (int z = 0; z < SLICES; ++z) + { + if (slice_id == z) + { +#pragma unroll + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { +#pragma unroll + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { +#pragma unroll + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) + { + if (z > 0) + { + C_warp16[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] += C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2]; + } + C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2] = C_warp16[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id]; + }; + } + } + } + __syncthreads(); + } + if (slice_id == 0) + { +#pragma unroll + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { +#pragma unroll + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { +#pragma unroll + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) + { + C_warp16[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] = C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2]; + }; + } + } + } + } + + if (slice_id == 0) + { + Semaphore semaphore(semaphores + blockIdx_y, threadIdx.x); + + if constexpr (SPLITK > 1) + { + semaphore.fetch(); + } + + if (blockIdx_z != 0) + { + semaphore.wait(blockIdx_z); + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2) + { + int write_row = cta_offset_m + warp_offset_m + ax0_0_1 * OP_M + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)); + + if (write_row < M) + { + nv_bfloat162 *existing_psum_ptr = reinterpret_cast( + C + write_row * N + + cta_offset_n + warp_offset_n + ax1_0_1 * 16 + + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2); + + *existing_psum_ptr = __hadd2(*existing_psum_ptr, + *reinterpret_cast(C_warp16 + ax0_0_1 * WARP_N / INTRIN_N * 8 + + ax1_0_1 * 8 + local_id)); + } + }; + } + } + } + else + { + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2) + { + int write_row = cta_offset_m + warp_offset_m + ax0_0_1 * OP_M + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)); + if (write_row < M) + { + *reinterpret_cast( + C + write_row * N + + cta_offset_n + warp_offset_n + ax1_0_1 * 16 + + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2) = + *reinterpret_cast(C_warp16 + ax0_0_1 * WARP_N / INTRIN_N * 8 + + ax1_0_1 * 8 + local_id); + } + }; + } + } + } + + if constexpr (SPLITK > 1) + { + + int lock = 0; + if (SPLITK == blockIdx_z + 1) + { + + lock = 0; + } + else + { + lock = blockIdx_z + 1; + } + semaphore.release(lock); + } + } + } +} + +template +__device__ __inline__ void global_to_share_one_stage_A_T2(T *src, T *dst, int global_nrows, int global_ncols, int cta_offset_m, int cta_offset_n, int global_iter_k, int shared_iter_k, bool mask) +{ + constexpr int threads_needed = (CTA_M * CTA_K) / PACK_SIZE / SHARED_K_ITERS; + constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE; + constexpr int total_global_iters = (CTA_M * CTA_K) / PACK_SIZE / threads_used; + constexpr int partial_global_iters = (total_global_iters + SHARED_K_ITERS - 1) / SHARED_K_ITERS; + constexpr int cta_step_m_or_n = (threads_used * PACK_SIZE) / CTA_K; + constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int threads_per_row = CTA_K / PACK_SIZE; + constexpr int kSmemCol = CTA_K + SMEM_PAD_A; + bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used); + int ld_col = (threadIdx.x % threads_per_row); +#pragma unroll + for (int _global_iter = 0; _global_iter < partial_global_iters; ++_global_iter) + { + int global_iter = shared_iter_k * partial_global_iters + _global_iter; + int ld_row = global_iter * cta_step_m_or_n + threadIdx.y * warp_step_m_or_n + (threadIdx.x / threads_per_row); + int ld_col_swizzled = (ld_col ^ (ld_row) & 7) * PACK_SIZE; + void *dst_ptr = (void *)(dst + ld_row * kSmemCol + ld_col_swizzled); + uint4 *src_ptr = (uint4 *)(src + (ld_row + cta_offset_m) * global_ncols + ld_col * PACK_SIZE + global_iter_k * CTA_K); // cta_offset_m * global_ncols + global_iter * cta_step_m_or_n * global_ncols + threadIdx.y * warp_step_m_or_n * global_ncols + (threadIdx.x / threads_per_row) * global_ncols + global_iter_k * CTA_K + (threadIdx.x % threads_per_row) * PACK_SIZE); + if constexpr (STAGES > 1) + { + uint32_t addr = cast_smem_ptr_to_uint(dst_ptr); + cp_async_cg_A(addr, src_ptr, local_mask & (ld_row + cta_offset_m < global_nrows)); + } + else + { + if (local_mask & (ld_row + cta_offset_m < global_nrows)) + *(uint4 *)dst_ptr = *src_ptr; + } + } +} + +template +__device__ __inline__ void global_to_share_one_stage_B_T2(T *src, T *dst, int global_ncols, int cta_offset_m, int cta_offset_n, int global_iter_k, int shared_iter_k, bool mask) +{ + constexpr int threads_needed = (CTA_N / kInterleave * CTA_K) / PACK_SIZE / SHARED_K_ITERS; + constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE; + constexpr int total_global_iters = (CTA_N / kInterleave * CTA_K) / PACK_SIZE / threads_used; + constexpr int partial_global_iters = (total_global_iters + SHARED_K_ITERS - 1) / SHARED_K_ITERS; + constexpr int cta_step_m_or_n = (threads_used * PACK_SIZE) / CTA_K; + constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int threads_per_row = CTA_K / PACK_SIZE; + constexpr int kSmemCol = CTA_K + SMEM_PAD_B; + bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used); +#pragma unroll + for (int _global_iter = 0; _global_iter < partial_global_iters; ++_global_iter) + { + int global_iter = shared_iter_k * partial_global_iters + _global_iter; + + int ld_row = global_iter * cta_step_m_or_n + threadIdx.y * warp_step_m_or_n + (threadIdx.x / threads_per_row); + int ld_col = (threadIdx.x % threads_per_row); + int ld_col_swizzled = ld_col ^ (ld_row % 2) & 7; + void *dst_ptr = (void *)(dst + (ld_row * kSmemCol + ld_col_swizzled * PACK_SIZE)); + uint4 *src_ptr = (uint4 *)(src + global_iter_k * CTA_K + cta_offset_n / kInterleave * global_ncols + ld_row * global_ncols + ld_col * PACK_SIZE); + if constexpr (STAGES > 1) + { + uint32_t addr = cast_smem_ptr_to_uint(dst_ptr); + cp_async_cg_A(addr, src_ptr, local_mask); + } + else + { + if (local_mask) + *(uint4 *)dst_ptr = *src_ptr; + } + } +} + +template +__device__ __inline__ void global_to_share_one_stage_scales_T2(T *src, T *dst, T *src_z, T *dst_z, int global_ncols, int cta_offset_m, int cta_offset_n, int global_iter_k, int shared_iter_k, bool mask) +{ + constexpr int threads_needed = CTA_N / PACK_SIZE / 1; + constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE; + constexpr int total_global_iters = CTA_N / PACK_SIZE / threads_used; + constexpr int threads_per_row = CTA_N / PACK_SIZE; + constexpr int kSmemCol = CTA_N; + bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used); + int g_idx = global_iter_k * CTA_K / G; + + void *dst_ptr = (void *)(dst + (threadIdx.x % threads_per_row) * PACK_SIZE); + uint4 *src_ptr = (uint4 *)(src + g_idx * global_ncols + cta_offset_n + (threadIdx.x % threads_per_row) * PACK_SIZE); + void *dst_ptr_z = (void *)(dst_z + (threadIdx.x % threads_per_row) * PACK_SIZE); + uint4 *src_ptr_z = (uint4 *)(src_z + g_idx * global_ncols + cta_offset_n + (threadIdx.x % threads_per_row) * PACK_SIZE); + if (STAGES > 1) + { + uint32_t addr = cast_smem_ptr_to_uint(dst_ptr); + cp_async_cg_A(addr, src_ptr, local_mask); + uint32_t addr_z = cast_smem_ptr_to_uint(dst_ptr_z); + cp_async_cg_A(addr_z, src_ptr_z, local_mask); + } + else + { + if (local_mask) + { + *(uint4 *)dst_ptr = *src_ptr; + *(uint4 *)dst_ptr_z = *src_ptr_z; + } + } +} + +template +__device__ __inline__ void share_to_reg_one_stage_A_T2(T *src, T *dst, int warp_offset_m, int warp_offset_n, int k_0_1) +{ + constexpr int kSmemCol = CTA_K + SMEM_PAD_A; + + for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) + { + + int ld_row = warp_offset_m + shared_iter * OP_M + (threadIdx.x % 16); + int ld_col = k_0_1 * 16 + (threadIdx.x / 16) * 8; + int ld_col_swizzled = ((ld_col / PACK_SIZE) ^ (ld_row) & 7) * PACK_SIZE; + void *addr_ptr = (void *)(src + ld_row * kSmemCol + ld_col_swizzled); + + uint32_t addr = cast_smem_ptr_to_uint(addr_ptr); + ldmatrix_m8n8_x4_b16(dst, shared_iter, addr); + } +} + +template +__device__ __inline__ void share_to_reg_one_stage_B_T2(T *src, T *src_scales, T *src_zeros, T *dst, T *dst_fp16, int warp_offset_m, int warp_offset_n, int k_0_1) +{ + using T2 = typename std::conditional::value, half2, nv_bfloat162>::type; + constexpr int kSmemCol = CTA_K + SMEM_PAD_B; + int r0 = ((threadIdx.x / 8 / 2) * 8 + threadIdx.x % 8); + int c0 = ((threadIdx.x / 8) % 2) * 8; + int r = r0 / 4; + int c = (r0 % 4) * 16 + c0; + int c_swizzled = ((c / PACK_SIZE) ^ (r % 2) & 7) * PACK_SIZE; + + if constexpr (ldmatrix) + { +#pragma unroll + for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) + { + void *addr_ptr = (void *)(src + warp_offset_n / kInterleave * kSmemCol + shared_iter * 16 / kInterleave * kSmemCol + k_0_1 * 16 + r * kSmemCol + c_swizzled); + uint32_t addr = cast_smem_ptr_to_uint(addr_ptr); + ldmatrix_m8n8_x4_b16(dst, shared_iter, addr); + } + } + +#pragma unroll + for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) + { + T scale = src_scales[warp_offset_n + 16 * shared_iter + 8 * (k_0_1 % 2) + threadIdx.x / 4]; + T zero = src_zeros[warp_offset_n + 16 * shared_iter + 8 * (k_0_1 % 2) + threadIdx.x / 4]; + T2 scale2, zero2; + if constexpr (std::is_same::value) + { + scale2 = __half2half2(scale); + zero2 = __half2half2(zero); + } + else + { + scale2 = __bfloat162bfloat162(scale); + zero2 = __bfloat162bfloat162(zero); + } + T2 loaded[4]; + dequantize_s4_to_fp16x2(*reinterpret_cast(dst + (k_0_1 % 2) * 4 + (k_0_1 / 2 * 2) + shared_iter * 8), reinterpret_cast(loaded)); +#pragma unroll + for (int i = 0; i < 4; i++) + { + loaded[i] = __hfma2(loaded[i], scale2, zero2); + } + *reinterpret_cast(dst_fp16 + shared_iter * 16 + 8 * (k_0_1 % 2)) = *reinterpret_cast(loaded); + } +} + +template +__global__ void gemm_w4a16_T2(T *__restrict__ A, T *__restrict__ B, T *__restrict__ scales, T *__restrict__ zeros, T *__restrict__ C, int M, int N, int K) +{ + using DTypeAccum = typename std::conditional::value, half, float>::type; + constexpr int NUM_WARPS = CTA_M / WARP_M * CTA_N / WARP_N; + constexpr int CTA_SIZE = NUM_WARPS * WARP_SIZE; + int num_blocks_n = (N + CTA_N - 1) / CTA_N; + int num_blocks_m = (M + CTA_M - 1) / CTA_M; + int blockIdx_x = 0; + int blockIdx_y = blockIdx.x % (num_blocks_m * num_blocks_n); + int blockIdx_z = blockIdx.x / (num_blocks_m * num_blocks_n); + const int log_tile = get_log_tile<1>((N + CTA_N - 1) / CTA_N); + int blockIdx_m = blockIdx_y / (num_blocks_n >> log_tile); + int blockIdx_n = blockIdx_y % (num_blocks_n >> log_tile); + const uint2 block_idx_mapping = get_block_idx_mapping(blockIdx_m, blockIdx_n, log_tile); + blockIdx_m = block_idx_mapping.x; + blockIdx_n = block_idx_mapping.y; + + DTypeAccum C_warp[CTA_M * CTA_N / CTA_SIZE]; + constexpr int kSmemPadKA = CTA_K + SMEM_PAD_A; + constexpr int kSmemPadKB = CTA_K + SMEM_PAD_B; + constexpr int kSmemSizeAPerStage = CTA_M * kSmemPadKA; + constexpr int kSmemSizeBPerStage = CTA_N / kInterleave * kSmemPadKB; + constexpr int kSmemSizeA = kSmemSizeAPerStage * STAGES; + constexpr int kSmemSizeB = kSmemSizeBPerStage * STAGES; + constexpr int kSmemSizeScales = CTA_N * STAGES / 2; + constexpr int kSmemSizeZeros = CTA_N * STAGES / 2; + constexpr int scales_load_interval = G / CTA_K; + extern __shared__ half mem_shared[]; + T *A_shared = (T*)mem_shared; + T *B_shared = (T*)mem_shared + kSmemSizeA; + T *scales_shared = (T*)mem_shared + kSmemSizeA + kSmemSizeB; + T *zeros_shared = (T*)mem_shared + kSmemSizeA + kSmemSizeB + kSmemSizeScales; + T A_shared_warp_[2][WARP_M * INTRIN_K / + WARP_SIZE]; + T B_shared_warp_[2][WARP_N * 32 / + WARP_SIZE]; + T B_shared_warp_tmp_[2][WARP_N * 16 / + WARP_SIZE]; + int cta_offset_m = blockIdx_m * CTA_M; + int cta_offset_n = blockIdx_n * CTA_N; + int warp_offset_m = (threadIdx.y % (CTA_M / WARP_M)) * WARP_M; + int warp_offset_n = (threadIdx.y / (CTA_M / WARP_M)) * WARP_N; + + for (int i = 0; i < CTA_M * CTA_N / CTA_SIZE; i++) + C_warp[i] = 0.0; + + int gemm_iters = (K + CTA_K - 1) / CTA_K; + int k_0_0_ld = 0; + int k_0_0 = 0; + constexpr int prologue_stages = STAGES == 1 ? 1 : STAGES - 1; +#pragma unroll + for (k_0_0_ld = 0; k_0_0_ld < prologue_stages; ++k_0_0_ld) + { + global_to_share_one_stage_A_T2(A, A_shared + k_0_0_ld * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, k_0_0_ld, 0, true); + global_to_share_one_stage_B_T2(B, B_shared + k_0_0_ld * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, k_0_0_ld, 0, true); + global_to_share_one_stage_scales_T2( + scales, scales_shared + (k_0_0_ld / scales_load_interval) * CTA_N, + zeros, zeros_shared + (k_0_0_ld / scales_load_interval) * CTA_N, + N, cta_offset_m, cta_offset_n, k_0_0_ld, 0, k_0_0_ld < gemm_iters && k_0_0_ld % scales_load_interval == 0); + if constexpr (STAGES > 1) + __pipeline_commit(); + } + if constexpr (STAGES > 1) + __pipeline_wait_prior(STAGES - 2); + __syncthreads(); + + share_to_reg_one_stage_A_T2(A_shared, A_shared_warp_[0], warp_offset_m, warp_offset_n, 0); + share_to_reg_one_stage_B_T2(B_shared, scales_shared, zeros_shared, B_shared_warp_tmp_[0], B_shared_warp_[0], warp_offset_m, warp_offset_n, 0); + constexpr int SHARED_K_ITERS = WARP_K / INTRIN_K; + + for (; k_0_0 < gemm_iters; ++k_0_0, ++k_0_0_ld) + { + int ld_stage = k_0_0_ld % STAGES; + int compute_stage = k_0_0 % STAGES; + T *A_shared_this_compute_stage; + T *B_shared_this_compute_stage; + T *scales_shared_this_compute_stage; + T *zeros_shared_this_compute_stage; + + for (int iter_k = 0; iter_k < SHARED_K_ITERS; ++iter_k) + { + A_shared_this_compute_stage = A_shared + compute_stage * kSmemSizeAPerStage; + B_shared_this_compute_stage = B_shared + compute_stage * kSmemSizeBPerStage; + scales_shared_this_compute_stage = scales_shared + (compute_stage / scales_load_interval) * CTA_N; + zeros_shared_this_compute_stage = zeros_shared + (compute_stage / scales_load_interval) * CTA_N; + share_to_reg_one_stage_A_T2(A_shared_this_compute_stage, A_shared_warp_[(iter_k + 1) % 2], warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS); + if ((iter_k + 1) % kInterleave == 0) + { + if (compute_stage % 2 == 1) + { + share_to_reg_one_stage_B_T2( + B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage, + B_shared_warp_tmp_[1], B_shared_warp_[((iter_k + 1) / 2) % 2], + warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS); + } + else + { + share_to_reg_one_stage_B_T2( + B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage, + B_shared_warp_tmp_[0], B_shared_warp_[((iter_k + 1) / 2) % 2], + warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS); + } + } + else + { + if (compute_stage % 2 == 1) + { + share_to_reg_one_stage_B_T2( + B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage, + B_shared_warp_tmp_[1], B_shared_warp_[((iter_k + 1) / 2) % 2], + warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS); + } + else + { + share_to_reg_one_stage_B_T2( + B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage, + B_shared_warp_tmp_[0], B_shared_warp_[((iter_k + 1) / 2) % 2], + warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS); + } + } + __syncthreads(); + T *A_shared_warp = A_shared_warp_[iter_k % 2]; + T *B_shared_warp = B_shared_warp_[(iter_k / 2) % 2]; + for (int i_0_3 = 0; i_0_3 < WARP_M / INTRIN_M; ++i_0_3) + { + for (int j_0_4 = 0; j_0_4 < WARP_N / INTRIN_N; ++j_0_4) + { + if constexpr (std::is_same::value) + { + mma_m16n8k16_f16f16f16(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4); + mma_m16n8k16_f16f16f16(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8 + 4, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4 + 8); + } + else + { + mma_m16n8k16_bf16bf16f32(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4); + mma_m16n8k16_bf16bf16f32(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8 + 4, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4 + 8); + } + } + } + + if (iter_k < WARP_K / INTRIN_K - 1) + { + if constexpr (STAGES == 1) + __syncthreads(); + global_to_share_one_stage_A_T2(A, A_shared + ld_stage * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, k_0_0_ld, iter_k, k_0_0_ld < gemm_iters); + global_to_share_one_stage_B_T2(B, B_shared + ld_stage * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, k_0_0_ld, iter_k, k_0_0_ld < gemm_iters); + } + + if (iter_k == WARP_K / INTRIN_K - 2) + { + if constexpr (STAGES == 1 && WARP_K / INTRIN_K > 2) + { + __syncthreads(); + } + global_to_share_one_stage_A_T2(A, A_shared + ld_stage * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, k_0_0_ld, iter_k + 1, k_0_0_ld < gemm_iters); + global_to_share_one_stage_B_T2(B, B_shared + ld_stage * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, k_0_0_ld, iter_k + 1, k_0_0_ld < gemm_iters); + global_to_share_one_stage_scales_T2( + scales, scales_shared + (ld_stage / scales_load_interval) * CTA_N, + zeros, zeros_shared + (ld_stage / scales_load_interval) * CTA_N, + N, cta_offset_m, cta_offset_n, k_0_0_ld, iter_k, k_0_0_ld < gemm_iters && k_0_0_ld % scales_load_interval == 0); + if constexpr (STAGES > 1) + { + __pipeline_commit(); + __pipeline_wait_prior(STAGES - 2); + } + compute_stage = (k_0_0 + 1) % STAGES; + __syncthreads(); + } + } + } + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2) + { + int write_row = cta_offset_m + warp_offset_m + ax0_0_1 * OP_M + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)); + if (write_row < M) + { + if constexpr (std::is_same::value) + { + *reinterpret_cast( + C + write_row * N + + cta_offset_n + warp_offset_n + ax1_0_1 * 16 + + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2) = + (*reinterpret_cast(C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 + + ax1_0_1 * 8 + local_id)); + } + else + { + *reinterpret_cast( + C + write_row * N + + cta_offset_n + warp_offset_n + ax1_0_1 * 16 + + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2) = + (__float22bfloat162_rn(*reinterpret_cast(C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 + + ax1_0_1 * 8 + local_id))); + } + } + }; + } + } +} + +torch::Tensor gemm_forward_cuda_new( + torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _scales, + torch::Tensor _zeros) +{ + std::vector output_shape = _in_feats.sizes().vec(); + output_shape.back() = _kernel.size(0) * kInterleave; + int num_in_feats = _in_feats.numel() / _in_feats.size(-1); + int num_in_channels = _in_feats.size(-1); + auto options = + torch::TensorOptions().dtype(_in_feats.dtype()).device(_in_feats.device()); + auto options_int = + torch::TensorOptions().dtype(torch::kInt32).device(_in_feats.device()); + at::Tensor _out_feats = torch::empty(output_shape, options); + int num_out_feats = _out_feats.numel() / _out_feats.size(-1); + int num_out_channels = _out_feats.size(-1); + + auto data_type = _in_feats.scalar_type(); + TORCH_CHECK(_scales.scalar_type() == data_type); + TORCH_CHECK(_zeros.scalar_type() == data_type); + + DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(data_type, ctype, { + auto in_feats = reinterpret_cast(_in_feats.data_ptr()); + auto kernel = reinterpret_cast(_kernel.data_ptr()); + auto scales = reinterpret_cast(_scales.data_ptr()); + auto zeros = reinterpret_cast(_zeros.data_ptr()); + auto out_feats = reinterpret_cast(_out_feats.data_ptr()); + + if (num_out_feats <= 32) + { + constexpr int G = 128; + constexpr int CTA_M = 16; + constexpr int CTA_N = 128; + constexpr int CTA_K = 128; + constexpr int WARP_M = 16; + constexpr int WARP_N = 32; + constexpr int WARP_K = 64; + constexpr int SPLITK = 2; + constexpr int STAGES = 4; + KERNEL_LAUNCH_CODE + } + else if (num_out_feats <= 64) + { + constexpr int G = 128; + constexpr int CTA_M = 16; + constexpr int CTA_N = 128; + constexpr int CTA_K = 128; + constexpr int WARP_M = 16; + constexpr int WARP_N = 32; + constexpr int WARP_K = 64; + constexpr int SPLITK = 1; + constexpr int STAGES = 3; + KERNEL_LAUNCH_CODE + } + else if (num_out_feats <= 128) + { + constexpr int G = 128; + constexpr int CTA_M = 32; + constexpr int CTA_N = 128; + constexpr int CTA_K = 128; + constexpr int WARP_M = 32; + constexpr int WARP_N = 32; + constexpr int WARP_K = 64; + constexpr int SPLITK = 1; + constexpr int STAGES = 4; + KERNEL_LAUNCH_CODE + } + else if (num_out_feats <= 192) + { + constexpr int G = 128; + constexpr int CTA_M = 64; + constexpr int CTA_N = 128; + constexpr int CTA_K = 64; + constexpr int WARP_M = 64; + constexpr int WARP_N = 32; + constexpr int WARP_K = 64; + constexpr int SPLITK = 1; + constexpr int STAGES = 4; + KERNEL_LAUNCH_CODE + } + else + { + constexpr int G = 128; + constexpr int CTA_M = 64; + constexpr int CTA_N = 128; + constexpr int CTA_K = 64; + constexpr int WARP_M = 64; + constexpr int WARP_N = 32; + constexpr int WARP_K = 64; + constexpr int STAGES = 4; + + constexpr int NUM_WARPS = (CTA_M / WARP_M) * (CTA_N / WARP_N); + constexpr int kSmemByteSize = (CTA_M * (CTA_K + SMEM_PAD_A) + CTA_N * (CTA_K + SMEM_PAD_B) / kInterleave + CTA_N) * STAGES * sizeof(ctype); + if (kSmemByteSize >= 99 * 1024) + { + printf("This kernel requires %d Bytes of shared memory, which exceeds device limit.\n", kSmemByteSize); + return _out_feats; + } + int j_factors1 = num_out_channels / CTA_N / 1; + dim3 num_blocks((num_out_feats + CTA_M - 1) / CTA_M * j_factors1); + dim3 threads_per_block(WARP_SIZE, NUM_WARPS); + auto kernel_func = gemm_w4a16_T2; + cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemByteSize); + kernel_func<<>>( + in_feats, kernel, scales, zeros, out_feats, num_in_feats, num_out_channels, num_in_channels); + } + }); + + return _out_feats; +} \ No newline at end of file diff --git a/llm-awq/awq/kernels/csrc/quantization_new/gemm/gemm_cuda.h b/llm-awq/awq/kernels/csrc/quantization_new/gemm/gemm_cuda.h new file mode 100644 index 0000000000000000000000000000000000000000..d5a7b158587a51e768163c5868e95f02dd8413d9 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization_new/gemm/gemm_cuda.h @@ -0,0 +1,3 @@ +#include + +torch::Tensor gemm_forward_cuda_new(torch::Tensor _in_feats, torch::Tensor _kernel, torch::Tensor _scales, torch::Tensor _zeros); diff --git a/llm-awq/awq/kernels/csrc/quantization_new/gemm/semaphore.h b/llm-awq/awq/kernels/csrc/quantization_new/gemm/semaphore.h new file mode 100644 index 0000000000000000000000000000000000000000..acc636f745c53fec08521c7ec863b5d1baf675f4 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization_new/gemm/semaphore.h @@ -0,0 +1,109 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Implementation of a CTA-wide semaphore for inter-CTA synchronization. +*/ + +#pragma once + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// namespace cutlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// CTA-wide semaphore for inter-CTA synchronization. +class Semaphore +{ +public: + int *lock; + bool wait_thread; + int state; + +public: + /// Implements a semaphore to wait for a flag to reach a given value + __host__ __device__ Semaphore(int *lock_, int thread_id) : lock(lock_), + wait_thread(thread_id < 0 || thread_id == 0), + state(-1) + { + } + + /// Permit fetching the synchronization mechanism early + __device__ void fetch() + { + if (wait_thread) + { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); +#else + asm volatile("ld.global.cg.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); +#endif + } + } + + /// Gets the internal state + __device__ int get_state() const + { + return state; + } + + /// Waits until the semaphore is equal to the given value + __device__ void wait(int status = 0) + { + while (__syncthreads_and(state != status)) + { + fetch(); + } + + __syncthreads(); + } + + /// Updates the lock with the given result + __device__ void release(int status = 0) + { + __syncthreads(); + + if (wait_thread) + { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("st.global.release.gpu.b32 [%0], %1;\n" : : "l"(lock), "r"(status)); +#else + asm volatile("st.global.cg.b32 [%0], %1;\n" : : "l"(lock), "r"(status)); +#endif + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// } // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/llm-awq/awq/kernels/csrc/quantization_new/gemv/gemv_cuda.cu b/llm-awq/awq/kernels/csrc/quantization_new/gemv/gemv_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..235293e43990a87bb4b1e6abb38cbc148c077478 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization_new/gemv/gemv_cuda.cu @@ -0,0 +1,339 @@ +/* + * Modified from NVIDIA [TRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/tree/d37b507f41a87457fe9f10f7459d08f5db235745/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv) + * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* +@article{lin2023awq, + title={AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration}, + author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Dang, Xingyu and Han, Song}, + journal={arXiv}, + year={2023} +} +*/ + +#include +#include +#include +#include "gemv_cuda.h" +#include "../dequantize.cuh" +#include "../dispatch_utils.cuh" +#define PACK_FACTOR 8 +#define WARP_SIZE 32 +#define MEM_ACCESS_SIZE 128 + +// Reduce sum within the warp using the tree reduction algorithm. +template +__device__ __forceinline__ static void warp_reduce(T* psum, float (*out_smem)[Num * 4]) +{ + // kInterleave = 4 + float fpsum[Num]; + #pragma unroll + for (int i = 0; i < Num; ++i) + { + fpsum[i] = static_cast(psum[i]); + } + + #pragma unroll + for (int i = 0; i < Num; ++i) + { + // T0 + T1 + T8 + T9 + T16 + T17 + T24 + T25 (kInterleave = 4) + fpsum[i] += __shfl_xor_sync(~0, fpsum[i], 16); + fpsum[i] += __shfl_xor_sync(~0, fpsum[i], 8); + fpsum[i] += __shfl_xor_sync(~0, fpsum[i], 1); + } + __syncthreads(); + int warp = threadIdx.x / WarpSize, lane = threadIdx.x % WarpSize; + if (lane == 0 || lane == 2 || lane == 4 || lane == 6) + { + #pragma unroll + for (int i = 0; i < Num; ++i) + { + out_smem[warp][i * 4 + lane / 2] = fpsum[i]; + } + } + __syncthreads(); +}; + +__device__ __forceinline__ int make_divisible(int c, int divisor){ + return (c + divisor - 1) / divisor; +} + +template +__global__ void gemv_kernel( + const T* inputs, const uint32_t* weight, const T* scales, const T* zeros, T* outputs, + const int IC, const int OC) +{ + const int kStride = 64; + const int kElemsPerThread = MEM_ACCESS_SIZE / 4; + const int kThreadsNumPerTile = kStride / kElemsPerThread; + // assert(MEM_ACCESS_SIZE == 128); + + using T2 = typename std::conditional< + std::is_same::value, + half2, + nv_bfloat162 + >::type; + + static constexpr int kShuffleSize = 32; + static constexpr int kShuffleBasicTile = 2; + static constexpr int kShuffleContinous = 4; + static constexpr int kShuffleStrided = 4; + + constexpr int Num = NPerBlock * Batch; + constexpr int kInterleave = 4; + + T local_inputs[kElemsPerThread]; + uint32_t local_qweights[MEM_ACCESS_SIZE / 32]; + T half_weight_buffer[kElemsPerThread]; + T dequantized_weight[kElemsPerThread * NPerBlock]; + T local_scale[NPerBlock]; + T local_scaled_zeros[NPerBlock]; + + T psum[Num]; + for (int i = 0; i < Num; ++i) + psum[i] = static_cast(0.f); + + // extern __shared__ uint8_t shmem[]; + // float(*out_smem)[Num * kInterleave] = reinterpret_cast(shmem); + __shared__ float out_smem[BlockSize / WARP_SIZE * 2][Num * kInterleave]; + + const int blk_row_offset = blockIdx.x * NPerBlock * kInterleave; + const int thd_row_offset = (threadIdx.x / kThreadsNumPerTile) % kInterleave; + const int act_k_offset = threadIdx.x / (kThreadsNumPerTile * kInterleave) * kStride + + (threadIdx.x % kThreadsNumPerTile) * kElemsPerThread; + const int group_offset = act_k_offset / GroupSize; + // TODO: use make_divisible + const uint32_t* blk_weight_ptr = weight + blk_row_offset * IC / PACK_FACTOR; + const T* scale_ptr = scales + blk_row_offset + thd_row_offset + group_offset * OC; + const T* zeros_ptr = zeros + blk_row_offset + thd_row_offset + group_offset * OC; + const T* inputs_ptr = inputs + act_k_offset; + + const int act_forward_step = BlockSize * kElemsPerThread / kInterleave; + const int scale_forward_step = act_forward_step / GroupSize * OC; + + // Main loop iteration, each block completes the outputs for several OCs + for (int kk = threadIdx.x * kElemsPerThread; kk < IC * kInterleave; kk += BlockSize * kElemsPerThread) + { + // Load qweight, scales and scaled_zeros + #pragma unroll + for (int idx = 0; idx < NPerBlock; ++idx) + { + // use float4 to load weights, each thread load 32 int4 numbers (1 x float4, 128 bit) + *((float4*)(local_qweights)) = + *((float4*)(blk_weight_ptr + (idx * kInterleave * IC + kk)/ PACK_FACTOR)); + local_scale[idx] = *(scale_ptr + idx * kInterleave); + local_scaled_zeros[idx] = *(zeros_ptr + idx * kInterleave); + + // Map int4 qweight to fp format + #pragma unroll + for (int i = 0; i < MEM_ACCESS_SIZE / 32; ++i) + { + // Converts 32 bits (8 x int4) to 8 fp16 + dequantize_s4_to_fp16x2(*reinterpret_cast(local_qweights + i), reinterpret_cast(half_weight_buffer + i * PACK_FACTOR)); + } + + // Dequantize (apply s/z) and shuffle elements to match the weight packing format + #pragma unroll + for (int i = 0; i < kShuffleContinous; ++i) + { + #pragma unroll + for (int j = 0; j < kShuffleStrided; ++j) + { + T2 w = + *reinterpret_cast( + half_weight_buffer + (i + j * kShuffleContinous)* kShuffleBasicTile + ); + if constexpr (std::is_same::value) + { + w = __hfma2(w, __half2half2(local_scale[idx]), __half2half2(local_scaled_zeros[idx])); + } + else + { + w = __hfma2(w, __bfloat162bfloat162(local_scale[idx]), __bfloat162bfloat162(local_scaled_zeros[idx])); + } + dequantized_weight[((i * kShuffleStrided + j) * kShuffleBasicTile + 0) + * NPerBlock + idx] + = w.x; + dequantized_weight[((i * kShuffleStrided + j) * kShuffleBasicTile + 1) + * NPerBlock + idx] + = w.y; + } + } + } + #pragma unroll + for (int batch_idx = 0; batch_idx < Batch; ++batch_idx) + { + const T* local_inputs_ptr = inputs_ptr + batch_idx * IC; + #pragma unroll + for (int idx = 0; idx < kElemsPerThread / 8; ++idx) + { + // load activation, 8 halves (128 bits) / step. + *((float4*)(local_inputs + idx * 8)) = *((float4*)(local_inputs_ptr + idx * 8)); + } + // Perform the MACs + #pragma unroll + for (int x = 0; x < NPerBlock / 2; ++x) + { + #pragma unroll + for (int y = 0; y < kElemsPerThread; ++y) + { + if constexpr (std::is_same::value) + { + *reinterpret_cast(psum + batch_idx * NPerBlock + x * 2) + = __hfma2(*reinterpret_cast(dequantized_weight + y * NPerBlock + x * 2), + __half2half2(local_inputs[y]), + *reinterpret_cast(psum + batch_idx * NPerBlock + x * 2)); + } + else + { + *reinterpret_cast(psum + batch_idx * NPerBlock + x * 2) + = __hfma2(*reinterpret_cast(dequantized_weight + y * NPerBlock + x * 2), + __bfloat162bfloat162(local_inputs[y]), + *reinterpret_cast(psum + batch_idx * NPerBlock + x * 2)); + } + } + } + } + inputs_ptr += act_forward_step; + scale_ptr += scale_forward_step; + zeros_ptr += scale_forward_step; + } + + warp_reduce(psum, out_smem); + + // Num * Interleave = batch * NPerBlock * Interleave -> 1 thread_block write back num + for (int i = threadIdx.x; i < Num * kInterleave; i += BlockSize) + { + int batch_idx = i / (NPerBlock * kInterleave); + int oc_idx = i % (NPerBlock * kInterleave); + float acc = 0.f; + for (int j = 0; j < BlockSize / WARP_SIZE; ++j) + { + acc += out_smem[j][i]; + } + outputs[batch_idx * OC + blk_row_offset + oc_idx] = static_cast(acc); + } +} + +/* +Computes GEMV (PyTorch interface). + +Args: + _in_feats: tensor of shape [B, IC]; + _kernel: int tensor of shape [OC, IC // 8]; + _zeros: int tensor of shape [OC, IC // G // 8]; + _scaling_factors: tensor of shape [OC, IC // G]; + blockDim_x: size of thread block, dimension x, where blockDim_x * workload_per_thread = IC; + blockDim_y: size of thread block, dimension y, where blockDim_y * gridDim_y = OC; + +Returns: + out_feats: tensor of shape [B, OC]; +*/ +torch::Tensor gemv_forward_cuda_new( + torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _scaling_factors, + torch::Tensor _zeros, + int m, + int n, + int k, + int group_size) +{ + + std::vector output_shape = _in_feats.sizes().vec(); + output_shape.back() = n; + + auto data_type = _in_feats.scalar_type(); + TORCH_CHECK(_scaling_factors.scalar_type() == data_type); + TORCH_CHECK(_zeros.scalar_type() == data_type); + + auto options = torch::TensorOptions().dtype(_in_feats.dtype()).device(_in_feats.device()); + at::Tensor _out_feats = torch::empty(output_shape, options); + + DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(data_type, ctype, { + auto in_feats = reinterpret_cast(_in_feats.data_ptr()); + auto kernel = reinterpret_cast(_kernel.data_ptr()); + auto zeros = reinterpret_cast(_zeros.data_ptr()); + auto scaling_factors = reinterpret_cast(_scaling_factors.data_ptr()); + auto out_feats = reinterpret_cast(_out_feats.data_ptr()); + + static constexpr int N_PER_BLOCK = 2; + static constexpr int K_INTERLEAVE = 4; + static constexpr int BLOCK_SIZE = 256; + + dim3 num_blocks(n / N_PER_BLOCK / K_INTERLEAVE); + dim3 num_threads(BLOCK_SIZE); + + // if (group_size == 64) + // { + // gemv_kernel_g64<<>>( + // // pointers + // in_feats, kernel, zeros, scaling_factors, out_feats, + // // constants + // num_in_channels, num_out_channels + // ); + // } + if (group_size == 128) + { + switch (m) + { + case 1: + gemv_kernel<<>>( + in_feats, kernel, scaling_factors, zeros, out_feats, k, n + ); + break; + case 2: + gemv_kernel<<>>( + in_feats, kernel, scaling_factors, zeros, out_feats, k, n + ); + break; + case 3: + gemv_kernel<<>>( + in_feats, kernel, scaling_factors, zeros, out_feats, k, n + ); + break; + case 4: + gemv_kernel<<>>( + in_feats, kernel, scaling_factors, zeros, out_feats, k, n + ); + break; + case 5: + gemv_kernel<<>>( + in_feats, kernel, scaling_factors, zeros, out_feats, k, n + ); + break; + case 6: + gemv_kernel<<>>( + in_feats, kernel, scaling_factors, zeros, out_feats, k, n + ); + break; + case 7: + gemv_kernel<<>>( + in_feats, kernel, scaling_factors, zeros, out_feats, k, n + ); + break; + default: + throw std::runtime_error("Unsupported batch size for gemv kernel.\n"); + } + } + else + { + throw std::runtime_error("Unsupported group size for gemv kernel.\n"); + } + }); + return _out_feats; +} + diff --git a/llm-awq/awq/kernels/csrc/quantization_new/gemv/gemv_cuda.h b/llm-awq/awq/kernels/csrc/quantization_new/gemv/gemv_cuda.h new file mode 100644 index 0000000000000000000000000000000000000000..181637cfc8716a9293ebea721593d4f121d2dfe7 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/quantization_new/gemv/gemv_cuda.h @@ -0,0 +1,12 @@ +#pragma once +#include + +torch::Tensor gemv_forward_cuda_new( + torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _scaling_factors, + torch::Tensor _zeros, + int m, + int n, + int k, + int group_size); diff --git a/llm-awq/awq/kernels/csrc/rope_new/fused_rope_with_pos.cu b/llm-awq/awq/kernels/csrc/rope_new/fused_rope_with_pos.cu new file mode 100644 index 0000000000000000000000000000000000000000..b81ebff6da00a691de7e001d03562a9477d6f8a6 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/rope_new/fused_rope_with_pos.cu @@ -0,0 +1,407 @@ +// Modified from https://github.com/NVIDIA/TransformerEngine +// Modified by Shang Yang. + +/************************************************************************* + * Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights + *reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include +#include + +#include "fused_rope_with_pos.h" +// #include + +// #include "../common.h" +// #include "../util/logging.h" +// #include "../utils.cuh" +#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) + +#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) + +#define THREADS_PER_WARP 32 + +template +__device__ void fused_rope_with_pos_block_forward( + const scalar_t *src, const float *freqs, scalar_t *dst, + const int offset_block, const int offset_block_dst, const int h, + const int d, const int d2, const int stride_h, const int stride_d, + const int o_stride_h, const int o_stride_d) { + int s_id = blockIdx.x; + int s = gridDim.x; + int b_id = blockIdx.y; +#pragma unroll + for (int d_id = threadIdx.x; d_id < d2; d_id += blockDim.x) { + float v_cos, v_sin; + sincosf(freqs[(b_id * s + s_id) * d2 + d_id], &v_sin, &v_cos); +#pragma unroll + for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) { + int offset_src = offset_block + h_id * stride_h + d_id * stride_d; + int offset_dst = offset_block_dst + h_id * o_stride_h + d_id * o_stride_d; + float v_src = src[offset_src]; + float v_src_rotate = + (d_id + d2 / 2 < d2) + ? -static_cast(src[offset_src + (d2 / 2) * stride_d]) + : static_cast(src[offset_src + (d2 / 2 - d2) * stride_d]); + dst[offset_dst] = v_src * v_cos + v_src_rotate * v_sin; + } + } + + // copy the rest + if (d > d2) { +#pragma unroll + for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) { + int offset_head = offset_block + h_id * stride_h; + int offset_head_dst = offset_block_dst + h_id * o_stride_h; +#pragma unroll + for (int d_id = d2 + threadIdx.x; d_id < d; d_id += blockDim.x) { + dst[offset_head_dst + d_id * o_stride_d] = + src[offset_head + d_id * stride_d]; + } + } + } +} + +// template +// __device__ void fused_rope_block_backward(const scalar_t *src, const float +// *freqs, scalar_t *dst, +// const int offset_block, const int +// offset_block_dst, const int h, +// const int d, const int d2, const +// int stride_h, const int stride_d, +// const int o_stride_h, const int +// o_stride_d) { +// int s_id = blockIdx.x; +// #pragma unroll +// for (int d_id = threadIdx.x; d_id < d2; d_id += blockDim.x) { +// float v_cos = cosf(freqs[s_id * d2 + d_id]); +// float v_sin = (d_id + d2 / 2 < d2) ? sinf(freqs[s_id * d2 + d_id + d2 / +// 2]) +// : -sinf(freqs[s_id * d2 + d_id + d2 / +// 2 - d2]); +// #pragma unroll +// for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) { +// int offset_src = offset_block + h_id * stride_h + d_id * stride_d; +// int offset_dst = offset_block_dst + h_id * o_stride_h + d_id * +// o_stride_d; float v_src = src[offset_src]; float v_src_rotate = (d_id + +// d2 / 2 < d2) ? src[offset_src + (d2 / 2) * stride_d] +// : src[offset_src + (d2 / 2 - +// d2) * stride_d]; +// dst[offset_dst] = v_src * v_cos + v_src_rotate * v_sin; +// } +// } + +// // handle the tail +// if (d > d2) { +// #pragma unroll +// for (int h_id = threadIdx.y; h_id < h; h_id += blockDim.y) { +// int offset_head = offset_block + h_id * stride_h; +// int offset_head_dst = offset_block_dst + h_id * o_stride_h; +// #pragma unroll +// for (int d_id = d2 + threadIdx.x; d_id < d; d_id += blockDim.x) { +// dst[offset_head_dst + d_id * o_stride_d] = src[offset_head + d_id * +// stride_d]; +// } +// } +// } +// } + +template +__global__ void fused_rope_with_pos_forward_kernel( + const scalar_t *src, const float *freqs, scalar_t *dst, const int h, + const int d, const int d2, const int stride_s, const int stride_b, + const int stride_h, const int stride_d, const int o_stride_s, + const int o_stride_b, const int o_stride_h, const int o_stride_d) { + int s_id = blockIdx.x, b_id = blockIdx.y; + int offset_block = s_id * stride_s + b_id * stride_b; + int offset_block_dst = s_id * o_stride_s + b_id * o_stride_b; + fused_rope_with_pos_block_forward( + src, freqs, dst, offset_block, offset_block_dst, h, d, d2, stride_h, + stride_d, o_stride_h, o_stride_d); +} + +// template +// __global__ void fused_rope_backward_kernel(const scalar_t *src, const float +// *freqs, scalar_t *dst, +// const int h, const int d, const +// int d2, const int stride_s, const +// int stride_b, const int stride_h, +// const int stride_d, const int +// o_stride_s, const int o_stride_b, +// const int o_stride_h, const int +// o_stride_d) { +// int s_id = blockIdx.x, b_id = blockIdx.y; +// int offset_block = s_id * stride_s + b_id * stride_b; +// int offset_block_dst = s_id * o_stride_s + b_id * o_stride_b; +// fused_rope_block_backward(src, freqs, dst, offset_block, +// offset_block_dst, h, d, d2, stride_h, +// stride_d, o_stride_h, o_stride_d); +// } + +template +void fused_rope_with_pos_forward_launcher( + const scalar_t *input, const float *freqs, scalar_t *output, const int s, + const int b, const int h, const int d, const int d2, const int stride_s, + const int stride_b, const int stride_h, const int stride_d, + const int o_stride_s, const int o_stride_b, const int o_stride_h, + const int o_stride_d, cudaStream_t stream) { + int warps_per_block = h < 16 ? 4 : 8; + dim3 blocks(s, b); + dim3 threads(THREADS_PER_WARP, warps_per_block); + + fused_rope_with_pos_forward_kernel<<>>( + input, freqs, output, h, d, d2, stride_s, stride_b, stride_h, stride_d, + o_stride_s, o_stride_b, o_stride_h, o_stride_d); + // NVTE_CHECK_CUDA(cudaGetLastError()); +} + +// template +// void fused_rope_backward_launcher(const scalar_t *output_grads, const float +// *freqs, +// scalar_t *input_grads, const int s, const +// int b, const int h, const int d, const int +// d2, const int stride_s, const int stride_b, +// const int stride_h, const int stride_d, +// const int o_stride_s, const int o_stride_b, +// const int o_stride_h, const int o_stride_d, +// cudaStream_t stream) { +// int warps_per_block = h < 16 ? 4 : 8; +// dim3 blocks(s, b); +// dim3 threads(THREADS_PER_WARP, warps_per_block); + +// fused_rope_backward_kernel<<>>( +// output_grads, freqs, input_grads, h, d, d2, stride_s, stride_b, +// stride_h, stride_d, o_stride_s, o_stride_b, o_stride_h, o_stride_d); +// // NVTE_CHECK_CUDA(cudaGetLastError()); +// } + +template +void fused_rope_with_pos_forward(const at::Tensor &input, + const at::Tensor &freqs, at::Tensor &output, + const int s, const int b, const int h, + const int d, const int d2, const int stride_s, + const int stride_b, const int stride_h, + const int stride_d, const int o_stride_s, + const int o_stride_b, const int o_stride_h, + const int o_stride_d, cudaStream_t stream) { + // TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + // input.data.dtype, scalar_t, + fused_rope_with_pos_forward_launcher( + reinterpret_cast(input.data_ptr()), + reinterpret_cast(freqs.data_ptr()), + reinterpret_cast(output.data_ptr()), s, b, h, d, d2, stride_s, + stride_b, stride_h, stride_d, o_stride_s, o_stride_b, o_stride_h, + o_stride_d, stream); + // ); +} + +// template +// void fused_rope_backward(const at::Tensor &output_grads, const at::Tensor +// &freqs, at::Tensor &input_grads, +// const int s, const int b, const int h, const int d, +// const int d2, const int stride_s, const int +// stride_b, const int stride_h, const int stride_d, +// const int o_stride_s, const int o_stride_b, const +// int o_stride_h, const int o_stride_d, cudaStream_t +// stream) { +// // TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( +// // output_grads.data.dtype, scalar_t, +// fused_rope_backward_launcher(reinterpret_cast(output_grads.data_ptr()), +// reinterpret_cast(freqs.data_ptr()), +// reinterpret_cast(input_grads.data_ptr()), s, b, h, d, +// d2, stride_s, stride_b, stride_h, +// stride_d, o_stride_s, o_stride_b, +// o_stride_h, o_stride_d, stream); +// // ); +// } + +template +void nvte_fused_rope_with_pos_forward( + const at::Tensor input, const at::Tensor freqs, at::Tensor output, + const int s, const int b, const int h, const int d, const int d2, + const int stride_s, const int stride_b, const int stride_h, + const int stride_d, const int o_stride_s, const int o_stride_b, + const int o_stride_h, const int o_stride_d, cudaStream_t stream) { + // NVTE_API_CALL(nvte_fused_rope_forward); + // using namespace transformer_engine; + fused_rope_with_pos_forward( + input, freqs, output, s, b, h, d, d2, stride_s, stride_b, stride_h, + stride_d, o_stride_s, o_stride_b, o_stride_h, o_stride_d, stream); +} + +// template +// void nvte_fused_rope_backward(const at::Tensor output_grads, const at::Tensor +// freqs, +// at::Tensor input_grads, const int s, const int +// b, const int h, const int d, const int d2, +// const int stride_s, const int stride_b, const +// int stride_h, const int stride_d, const int +// o_stride_s, const int o_stride_b, const int +// o_stride_h, const int o_stride_d, cudaStream_t +// stream) { +// // NVTE_API_CALL(nvte_fused_rope_backward); +// // using namespace transformer_engine; +// fused_rope_backward(output_grads, freqs, input_grads, s, b, h, d, +// d2, stride_s, stride_b, +// stride_h, stride_d, o_stride_s, o_stride_b, o_stride_h, +// o_stride_d, stream); +// } + +// Interface for Python +at::Tensor fused_rope_with_pos_forward_func( + const at::Tensor &input, const at::Tensor &freqs, + const bool transpose_output_memory) { + // using namespace transformer_engine; + // TORCH_CHECK(input.dim() == 4, "expected 4D tensor"); + // TORCH_CHECK(freqs.dim() == 4, "expected 4D tensor"); + // TORCH_CHECK(input.size(0) <= freqs.size(0), + // "expected freqs tensor has a longer sequence length than + // input"); + // TORCH_CHECK(freqs.size(1) == 1 && freqs.size(2) == 1, + // "expected the second and third dims of the freqs tensor equal + // 1"); + // TORCH_CHECK(input.size(3) >= freqs.size(3), + // "expected the last dim of the input tensor equals or is " + // "greater than the freqs tensor"); + // TORCH_CHECK(freqs.scalar_type() == at::ScalarType::Float, + // "Dtype of the freqs tensor must be float"); + + // input sizes: (s, b, h, d) + // s: sequence length + // b: batch size + // h: head num + // d: dim of each head + const int s = input.size(0); + const int b = input.size(1); + const int h = input.size(2); + const int d = input.size(3); + // input strides + const int stride_s = input.stride(0); + const int stride_b = input.stride(1); + const int stride_h = input.stride(2); + const int stride_d = input.stride(3); + // freqs' shape is always (s, 1, 1, d2), so the strides are same under + // different memory formats + // freqs' shape is now (B, S, D) + const int d2 = freqs.size(-1); + + // output + auto act_options = input.options().requires_grad(false); + at::Tensor output; + if (transpose_output_memory) { + output = torch::empty({b, s, h, d}, act_options).transpose(0, 1); + } else { + output = torch::empty({s, b, h, d}, act_options); + } + // output strides + const int o_stride_s = output.stride(0); + const int o_stride_b = output.stride(1); + const int o_stride_h = output.stride(2); + const int o_stride_d = output.stride(3); + + auto input_cu = input; + auto freqs_cu = freqs; + auto output_cu = output; + + VLLM_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "nvte_fused_rope_forward", [&] { + nvte_fused_rope_with_pos_forward( + input_cu.data(), freqs_cu.data(), output_cu.data(), s, b, h, d, d2, + stride_s, stride_b, stride_h, stride_d, o_stride_s, o_stride_b, + o_stride_h, o_stride_d, at::cuda::getCurrentCUDAStream()); + }); + + // nvte_fused_rope_forward(input_cu.data(), + // freqs_cu.data(), output_cu.data(), s, b, h, d, d2, + // stride_s, stride_b, stride_h, stride_d, o_stride_s, + // o_stride_b, o_stride_h, o_stride_d, + // at::cuda::getCurrentCUDAStream()); + + return output; +} + +// // Interface for Python +// at::Tensor fused_rope_backward_func(const at::Tensor &output_grads, const +// at::Tensor &freqs, +// const bool transpose_output_memory) { +// // using namespace transformer_engine; +// // TORCH_CHECK(output_grads.dim() == 4, "expected 4D tensor"); +// // TORCH_CHECK(freqs.dim() == 4, "expected 4D tensor"); +// // TORCH_CHECK(output_grads.size(0) <= freqs.size(0), +// // "expected freqs tensor has a longer sequence length than +// output_grads"); +// // TORCH_CHECK(freqs.size(1) == 1 && freqs.size(2) == 1, +// // "expected the second and third dims of the freqs tensor +// equal 1"); +// // TORCH_CHECK(output_grads.size(3) >= freqs.size(3), +// // "expected the last dim of the output_grads tensor equals or +// is " +// // "greater than the freqs tensor"); +// // TORCH_CHECK(freqs.scalar_type() == at::ScalarType::Float, +// // "Dtype of the freqs tensor must be float"); + +// // output_grads sizes: (s, b, h, d) +// // s: sequence length +// // b: batch size +// // h: head num +// // d: dim of each head +// const int s = output_grads.size(0); +// const int b = output_grads.size(1); +// const int h = output_grads.size(2); +// const int d = output_grads.size(3); +// // output_grads strides +// const int stride_s = output_grads.stride(0); +// const int stride_b = output_grads.stride(1); +// const int stride_h = output_grads.stride(2); +// const int stride_d = output_grads.stride(3); +// // freqs' shape is always (s, 1, 1, d2), so the strides are same under +// // different memory formats +// const int d2 = freqs.size(3); + +// auto act_options = output_grads.options().requires_grad(false); +// at::Tensor input_grads; +// if (transpose_output_memory) { +// input_grads = torch::empty({b, s, h, d}, act_options).transpose(0, 1); +// } else { +// input_grads = torch::empty({s, b, h, d}, act_options); +// } +// const int o_stride_s = input_grads.stride(0); +// const int o_stride_b = input_grads.stride(1); +// const int o_stride_h = input_grads.stride(2); +// const int o_stride_d = input_grads.stride(3); + +// auto output_grads_cu = output_grads; +// auto freqs_cu = freqs; +// auto input_grads_cu = input_grads; + +// VLLM_DISPATCH_FLOATING_TYPES( +// output_grads.scalar_type(), "nvte_fused_rope_forward", [&] { +// nvte_fused_rope_backward(output_grads_cu.data(), +// freqs_cu.data(), input_grads_cu.data(), s, b, h, +// d, d2, stride_s, stride_b, stride_h, +// stride_d, o_stride_s, o_stride_b, +// o_stride_h, o_stride_d, +// at::cuda::getCurrentCUDAStream()); +// }); + +// // nvte_fused_rope_backward(output_grads_cu.data(), freqs_cu.data(), +// input_grads_cu.data(), s, b, h, +// // d, d2, stride_s, stride_b, stride_h, stride_d, +// o_stride_s, o_stride_b, +// // o_stride_h, o_stride_d, +// at::cuda::getCurrentCUDAStream()); + +// return input_grads; +// } diff --git a/llm-awq/awq/kernels/csrc/rope_new/fused_rope_with_pos.h b/llm-awq/awq/kernels/csrc/rope_new/fused_rope_with_pos.h new file mode 100644 index 0000000000000000000000000000000000000000..6f999d186e1f2ed76f8c31878b8591b7b94d12ab --- /dev/null +++ b/llm-awq/awq/kernels/csrc/rope_new/fused_rope_with_pos.h @@ -0,0 +1,5 @@ +#include + +at::Tensor fused_rope_with_pos_forward_func(const at::Tensor &input, + const at::Tensor &freqs, + const bool transpose_output_memory); diff --git a/llm-awq/awq/kernels/csrc/w8a8/act.cu b/llm-awq/awq/kernels/csrc/w8a8/act.cu new file mode 100644 index 0000000000000000000000000000000000000000..5dce6c1c55d7ba33c94094a97a6dc09ee6aeee59 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/act.cu @@ -0,0 +1,141 @@ +#include +#include +#include + +#include "dispatch_utils.h" +#include "utils.cuh" +#include "reduction_utils.cuh" + +namespace vllm { + +template __device__ __forceinline__ T silu(const T &x) { + // x * sigmoid(x) + return (T)(((float)x) / (1.0f + expf((float)-x))); +} + +template __device__ __forceinline__ T gelu_new(const T &x) { + const half x3 = (half)(x * x * x); + const T t = (T)tanhf((T)((T)0.79788456f * (half)(x + (T)((T)0.044715f * x3)))); + return ((T)0.5) * x * (((T)1.0) + t); +} + +template +__device__ __forceinline__ T gelu_fast(const T &x) { + const half f = (half)x; + const T t = + (T)tanhf(((T)(f * (T)0.79788456f)) * (((T)1.0) + (T)((T)0.044715f * f) * x)); + return ((T)0.5) * x * (((T)1.0) + t); +} + + + +// dequant int32 input, apply silu and mul, then per token quant to int8 +template +__global__ void gelu_and_quant_kernel( + int8_t *__restrict__ out, // [..., d] + half *__restrict__ input, // [..., d] + const int d, + scale_type * scale_out, // [num_tokens] + half *__restrict__ tmp = nullptr // [num_tokens, d] +) { + const int token_idx = blockIdx.x; + const float max_value= 127.0f; + if constexpr (use_per_token_quant) { + float amax_val = 0.0f; + const half zero = 0.0001f; + + for (int idx = threadIdx.x; idx < d; idx += blockDim.x) { + const half x = + (half)__ldg(&input[token_idx * d + idx]); + half t = gelu_fast(x); + tmp[token_idx * d + idx] = t; + t = t > zero ? t : -t; + if ((float)t > amax_val) + amax_val = (float)t; + } + + __shared__ float s_amax; + const float block_amax_val = blockReduceMax(amax_val); + if (threadIdx.x == 0) { + s_amax = block_amax_val; + scale_out[token_idx] = half(block_amax_val / max_value); + } + __syncthreads(); + + float tmp_scale = max_value / s_amax; + for (int idx = threadIdx.x; idx < d; idx += blockDim.x) { + out[token_idx * d + idx] = + float_to_int8_rn((half)tmp_scale * tmp[token_idx * d + idx]); + } + } else { + for (int idx = threadIdx.x; idx < d; idx += blockDim.x) { + const float x = + (float)__ldg(&input[token_idx * d + idx]); + out[token_idx * d + idx] = float_to_int8_rn((half)gelu_fast(x) / scale_out[0]); + } + } +} +} // namespace vllm + + + +void gelu_and_quant( + torch::Tensor &out, // [..., d] + torch::Tensor &input, // [..., d] + torch::Tensor &scale_out, // [...] + torch::Tensor &tmp // [num_tokens, d] + ) { + int64_t num_tokens = input.numel() / input.size(-1); + int d = input.size(-1); + dim3 grid(num_tokens); + dim3 block(std::min(d, 128)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + vllm::gelu_and_quant_kernel<<>>( + out.data_ptr(), reinterpret_cast(input.data_ptr()), d, reinterpret_cast(scale_out.data_ptr()),reinterpret_cast(tmp.data_ptr())); +} + + + +namespace vllm { + +template +__global__ void silu_and_mul_kernel( + scalar_t* __restrict__ out, // [..., d] + const scalar_t* __restrict__ input, // [..., 2 * d] + const int d) { + + const int token_idx = blockIdx.x; + const int64_t token_idx_d = token_idx * int64_t(d); + const int64_t token_idx_2d = token_idx_d * 2; + for (int idx = threadIdx.x; idx < d; idx += blockDim.x) { + const scalar_t x = __ldg(&input[token_idx_2d + idx]); + const scalar_t y = __ldg(&input[token_idx_2d + d + idx]); + out[token_idx_d + idx] = silu(x) * y; + } +} +} // namespace vllm + + + +torch::Tensor silu_and_mul( + torch::Tensor& input) // [..., 2 * d] +{ + int64_t num_tokens = input.numel() / input.size(-1); + int d = input.size(-1) / 2; + + std::vector output_shape = input.sizes().vec(); + output_shape[output_shape.size() - 1]=d; + auto options = + torch::TensorOptions().dtype(input.dtype()).device(input.device()); + at::Tensor output = torch::empty(output_shape, options); + + + dim3 grid(num_tokens); + dim3 block(std::min(d, 256)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "silu_and_mul_kernel", [&] { + vllm::silu_and_mul_kernel<<>>( + output.data_ptr(), input.data_ptr(), d); + }); + return output; +} \ No newline at end of file diff --git a/llm-awq/awq/kernels/csrc/w8a8/act.h b/llm-awq/awq/kernels/csrc/w8a8/act.h new file mode 100644 index 0000000000000000000000000000000000000000..98c92c2cfbaf8929321b921789cf48405d10f7a0 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/act.h @@ -0,0 +1,29 @@ +// Inspired by TRT-LLM. +// Modified by Shang Yang and Haotian Tang. +// @article{lin2024awq, +// title={AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration}, +// author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Chen, Wei-Ming and Wang, Wei-Chen and Xiao, Guangxuan and Dang, Xingyu and Gan, Chuang and Han, Song}, +// journal={Proceedings of Machine Learning and Systems}, +// volume={6}, +// pages={87--100}, +// year={2024} +// } + +#include +#include +// Inspired by vLLM-SmoothQuant: https://github.com/vllm-project/vllm/pull/1112. +#include + + +void gelu_and_quant(torch::Tensor &out, // [..., d] + torch::Tensor &input, // [..., d] + torch::Tensor &scale_out, // [num_tokens] + torch::Tensor &tmp // [num_tokens, d] +); + +torch::Tensor silu_and_mul(torch::Tensor &input // [..., 2 * d] +); + + + + diff --git a/llm-awq/awq/kernels/csrc/w8a8/dispatch_utils.h b/llm-awq/awq/kernels/csrc/w8a8/dispatch_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..7c0c49d392a9806f1da29f0df61d362ed7170f59 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/dispatch_utils.h @@ -0,0 +1,14 @@ +/* + * Adapted from + * https://github.com/pytorch/pytorch/blob/v2.0.1/aten/src/ATen/Dispatch.h + */ +#include + +#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) + +#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH( \ + TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) diff --git a/llm-awq/awq/kernels/csrc/w8a8/layernorm.cu b/llm-awq/awq/kernels/csrc/w8a8/layernorm.cu new file mode 100644 index 0000000000000000000000000000000000000000..510d4e1e471fe67df5dd37656f23163a70fc2bf6 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/layernorm.cu @@ -0,0 +1,232 @@ +// Inspired by QServe https://github.com/mit-han-lab/qserve/tree/main. +// Modified by Yuming Lou. +// @article{lin2024awq, +// title={AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration}, +// author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Chen, Wei-Ming and Wang, Wei-Chen and Xiao, Guangxuan and Dang, Xingyu and Gan, Chuang and Han, Song}, +// journal={Proceedings of Machine Learning and Systems}, +// volume={6}, +// pages={87--100}, +// year={2024} +// } +#include +#include +#include "dispatch_utils.h" +#include "utils.cuh" +#include "reduction_utils.cuh" + + +namespace vllm { + +// from TRTLLM +template +__inline__ __device__ Tf compute_layernorm(Tf val, float s_mean, float s_variance, const T* gamma, const T* beta, int i) +{ + Tf ret = (val - s_mean) * s_variance * cuda_cast(gamma[i]); + if (beta != nullptr) + { + ret = ret + cuda_cast(beta[i]); + } + return ret; +} + +// from TRTLLM +/* Computes the layernorm https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html + * normed_output <- ( (input - E[input]) / Sqrt(Var[input] + eps) ) * gamma + beta + * input is [tokens, hidden_dim]. Mean and Variance are per-row (i.e. per-token) + * + * One CTA handles one row. + * + * with USE_DIFF_OF_SQUARES set to false: + * First pass (loop) computes the mean. + * Second computes the variance via Var[x] = E[(x - E[x])²]. + * Third pass computes and writes normed_output + * For better speedup, we set USE_DIFF_OF_SQUARES to true (may be faster but less accurate): + * It turns out the accuracy dosen't drop. + * First pass (loop) computes the mean and variance via Var[x] = E[x²] - E[x]² + * Second pass computes and writes normed_output + * + * + * use_shmem controls if we cache input values into shared memory + * + * Optional: with dynamic scaling, the last pass doesn't write immediately but finds the + * amax per row. A final pass scales to int8 accordingly, and writes output to + * normed_output_quant. + */ +template +__global__ void generalLayerNorm(const T* input, const T* gamma, const T* beta, T* normed_output, const float eps, + int tokens, int hidden_dim, const scale_type* scale_orig_quant_per_tensor, scale_type* scale_orig_quant_per_token, + int8_t* normed_output_quant, bool use_shmem) +{ + constexpr auto num_elems_T = num_elems::value; + using int8_packed_t = typename packed_as::type; + using float_packed_t = typename packed_as::type; + using T_scalar = typename packed_as::type; + + extern __shared__ __align__(sizeof(float)) char _shmem[]; + T* shmem = reinterpret_cast(_shmem); + __shared__ float s_mean; + __shared__ float s_variance; + + const int tidx = threadIdx.x; + const int bidx = blockIdx.x; + + float mean = 0.0f; + float variance = 0.0f; + float local_sum = 0.0f; + float local_var_sum = 0.0f; + const int n_elems = hidden_dim / num_elems_T; + for (int i = tidx; i < n_elems; i += blockDim.x) + { + const T val = input[bidx * n_elems + i]; + if (use_shmem) + { + shmem[i] = val; + } + const float_packed_t val_f = cuda_cast(val); + local_sum += cuda_sum(val_f); + if (USE_DIFF_OF_SQUARES) + { + local_var_sum += cuda_sum(val_f * val_f); + } + } + //Compute mean + if (USE_DIFF_OF_SQUARES) + { + float packed[2] = {local_sum, local_var_sum}; + blockReduceSumV2(packed); + mean = packed[0]; + variance = packed[1]; + } + else + { + mean = blockReduceSum(local_sum); + } + + if (threadIdx.x == 0) + { + mean = mean / hidden_dim; + s_mean = mean; + if (USE_DIFF_OF_SQUARES) + { + variance = (variance / hidden_dim) - (mean * mean); // Var[x] = E[x²] - E[x]² + s_variance = rsqrtf(variance + eps); + } + } + __syncthreads(); + + + if (!USE_DIFF_OF_SQUARES) + { + for (int i = tidx; i < n_elems; i += blockDim.x) + { + const T val = use_shmem ? shmem[i] : input[bidx * n_elems + i]; + float_packed_t diff = cuda_cast(val); // - s_mean; + local_var_sum += cuda_sum(diff * diff); + } + variance = blockReduceSum(local_var_sum); + + if (threadIdx.x == 0) + { + s_variance = rsqrtf(variance / hidden_dim + eps); + } + __syncthreads(); + } + + // Compute LN and Quantize + const bool with_per_token_scaling = scale_orig_quant_per_token != nullptr; + const bool with_per_tensor_scaling = scale_orig_quant_per_tensor != nullptr; + const float_packed_t scale_orig_quant + = cuda_cast(with_per_tensor_scaling ? __half2float(*scale_orig_quant_per_tensor) : 0.0f); + T_scalar amax = 1e-6f; + + for (int i = tidx; i < n_elems; i += blockDim.x) + { + const int index = bidx * n_elems + i; + const float_packed_t val_f = cuda_cast(use_shmem ? shmem[i] : input[index]); + const T val = cuda_cast(compute_layernorm(val_f, s_mean, s_variance, gamma, beta, i)); + + if (with_per_token_scaling) + { + amax = cuda_max(cuda_max(cuda_abs(val)), amax); + if (use_shmem) + { + shmem[i] = val; + } + } + else if (with_per_tensor_scaling) + { + reinterpret_cast(normed_output_quant)[index] + = cuda_cast(cuda_cast(val) * scale_orig_quant); + } + else + { + normed_output[index] = val; + } + } + + if (with_per_token_scaling) + { + float abs_max_f = blockAllReduceMax(cuda_cast(amax)); + const float dynamic_per_token_scale = 127.f / abs_max_f; + for (int i = tidx; i < n_elems; i += blockDim.x) + { + const int index = bidx * n_elems + i; + float_packed_t val_f = cuda_cast(use_shmem ? shmem[i] : input[index]); + if (!use_shmem) + { + val_f = compute_layernorm(val_f, s_mean, s_variance, gamma, beta, i); + } + + reinterpret_cast(normed_output_quant)[index] + = cuda_cast(val_f * cuda_cast(dynamic_per_token_scale)); + } + if (tidx == 0) + { + scale_orig_quant_per_token[bidx] = abs_max_f / 127.f; + } + } +} + + +} // namespace vllm + +void rms_norm_general(torch::Tensor &out, // [..., hidden_size] + torch::Tensor &input, // [..., hidden_size] + torch::Tensor &weight, // [hidden_size] + torch::Tensor &bias, // [hidden_size] + torch::Tensor &scaling, // [tokens] or [1] + float epsilon, + bool use_per_token_quant = true) { + int hidden_size = input.size(-1); + int num_tokens = input.numel() / hidden_size; + dim3 grid(num_tokens); + dim3 block(std::min(hidden_size, 128));//Reduce the idle probability of threads + block.x = 32 * ((block.x + 31) / 32); + + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "generalLayerNorm", [&] { + using T = typename FloatTypeConverter::Type; + if (use_per_token_quant) { + // per-token + vllm::generalLayerNorm<<>>( + reinterpret_cast(input.data_ptr()), + reinterpret_cast(weight.data_ptr()), + reinterpret_cast(bias.data_ptr()), + nullptr, epsilon, num_tokens, hidden_size, nullptr, scaling.data_ptr(), + out.data_ptr(), false + ); + // input, gamma, beta, normed_output, eps, tokens, hidden_dim, per_tensor_scale, per_token_scale + // normed_output_quant, use_shmem + // out.data_ptr(), input.data_ptr(), + // weight.data_ptr(), epsilon, num_tokens, hidden_size); + } else { + // per-tensor + vllm::generalLayerNorm<<>>( + reinterpret_cast(input.data_ptr()), + reinterpret_cast(weight.data_ptr()), nullptr, + nullptr, epsilon, num_tokens, hidden_size, scaling.data_ptr(), nullptr, + out.data_ptr(), false + ); + } + }); +} \ No newline at end of file diff --git a/llm-awq/awq/kernels/csrc/w8a8/layernorm.h b/llm-awq/awq/kernels/csrc/w8a8/layernorm.h new file mode 100644 index 0000000000000000000000000000000000000000..9e5d740520cf3b1f1b910e61e59ffe0dc570d810 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/layernorm.h @@ -0,0 +1,21 @@ +// Inspired by TRT-LLM. +// Modified by Shang Yang and Haotian Tang. +// @article{lin2024awq, +// title={AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration}, +// author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Chen, Wei-Ming and Wang, Wei-Chen and Xiao, Guangxuan and Dang, Xingyu and Gan, Chuang and Han, Song}, +// journal={Proceedings of Machine Learning and Systems}, +// volume={6}, +// pages={87--100}, +// year={2024} +// } + +#include +#include +void rms_norm_general(torch::Tensor &out, // [..., hidden_size] + torch::Tensor &input, // [..., hidden_size] + torch::Tensor &weight, // [hidden_size] + torch::Tensor &bias, // [hidden_size] + torch::Tensor &scaling, // [tokens] or [1] + float epsilon, + bool use_per_token_quant); + diff --git a/llm-awq/awq/kernels/csrc/w8a8/quantization.cu b/llm-awq/awq/kernels/csrc/w8a8/quantization.cu new file mode 100644 index 0000000000000000000000000000000000000000..43d02507448645e6441c7b131b69399dcad6ad3b --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/quantization.cu @@ -0,0 +1,113 @@ +// Inspired by vLLM-SmoothQuant: https://github.com/vllm-project/vllm/pull/1112 and TensorRT-LLM. +// Modified by Shang Yang and Haotian Tang. +// @article{lin2024awq, +// title={AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration}, +// author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Chen, Wei-Ming and Wang, Wei-Chen and Xiao, Guangxuan and Dang, Xingyu and Gan, Chuang and Han, Song}, +// journal={Proceedings of Machine Learning and Systems}, +// volume={6}, +// pages={87--100}, +// year={2024} +// } +#include +#include + +#include "utils.cuh" +#include +#include +#include "quantization.h" + +#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) +#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) + +template +__inline__ __device__ T warpReduceMax(T val) +{ +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + val = max(val, __shfl_xor_sync(0xffffffff, val, mask, 32)); + return val; +} + +/* Calculate the maximum of all elements in a block */ +template +__inline__ __device__ T blockReduceMax(T val) +{ + static __shared__ T shared[32]; + int lane = threadIdx.x & 0x1f; // in-warp idx + int wid = threadIdx.x >> 5; // warp idx + val = warpReduceMax(val); // get maxx in each warp + if (lane == 0) // record in-warp maxx by warp Idx + shared[wid] = val; + __syncthreads(); + // Modify from blockDim.x << 5 to blockDim.x / 32. to prevent + // blockDim.x is not divided by 32 + val = (threadIdx.x < (blockDim.x / 32.f)) ? shared[lane] : -1e20f; + val = warpReduceMax(val); + return val; +} + + + +namespace vllm { +template +__global__ void quant_kernel(const T *__restrict__ input, + int8_t *__restrict__ output, scale_type scale, + int num_tokens, int hidden_size) { + const int tid = threadIdx.x; + const int token_idx = blockIdx.x; + + if constexpr (use_per_token_quant) { + float amax_val = 0.0f; + const float zero = 0.0f; + + for (int i = tid; i < hidden_size; i += blockDim.x) { + float val = (float)input[token_idx * hidden_size + i]; + val = val > zero ? val : -val; + if (val > amax_val) + amax_val = val; + } + + __shared__ float s_amax; + const float block_amax_val = blockReduceMax(amax_val); + if (tid == 0) { + s_amax = block_amax_val; + scale[token_idx] = __float2half_rn(block_amax_val / 127.0f); + } + __syncthreads(); + + float tmp_scale = 127.0f / s_amax; + for (int i = tid; i < hidden_size; i += blockDim.x) { + output[token_idx * hidden_size + i] = + float_to_int8_rn(((float)input[token_idx * hidden_size + i]) * tmp_scale); + } + } else { + for (int i = tid; i < hidden_size; i += blockDim.x) { + output[token_idx * hidden_size + i] = + float_to_int8_rn(((float)input[token_idx * hidden_size + i]) / __half2float(scale)); + } + } +} +} + + + +void invoke_quant(torch::Tensor &out, // [..., hidden_size] + torch::Tensor &input, // [..., hidden_size] + torch::Tensor &scale) { // [num_tokens] + assert(input.is_contiguous()); + assert(out.is_contiguous()); + int hidden_size = input.size(-1); + int num_tokens = input.numel() / hidden_size; + dim3 grid(num_tokens); + dim3 block(std::min(hidden_size, 1024)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "quant_kernel", [&] { + vllm::quant_kernel<<>>( + input.data_ptr(), out.data_ptr(), + scale.data_ptr(), num_tokens, hidden_size); + }); +} + diff --git a/llm-awq/awq/kernels/csrc/w8a8/quantization.h b/llm-awq/awq/kernels/csrc/w8a8/quantization.h new file mode 100644 index 0000000000000000000000000000000000000000..a5db72d730aa30075c2548a4cf4b5ac50b7f9b22 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/quantization.h @@ -0,0 +1,4 @@ +#include +void invoke_quant(torch::Tensor &out, // [..., hidden_size] + torch::Tensor &input, // [..., hidden_size] + torch::Tensor &scale); // [num_tokens] \ No newline at end of file diff --git a/llm-awq/awq/kernels/csrc/w8a8/utils.cuh b/llm-awq/awq/kernels/csrc/w8a8/utils.cuh new file mode 100644 index 0000000000000000000000000000000000000000..cd4d45df253273bb1de6e2652836d86ecee4054b --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/utils.cuh @@ -0,0 +1,469 @@ +// Adated from FasterTransformer, https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp +// Modified by Haotian Tang +#pragma once + +#include +#include +#include +#include +#include + +template +struct FloatTypeConverter +{ + using Type = T; +}; + +template <> +struct FloatTypeConverter +{ + using Type = half; +}; + +template <> +struct FloatTypeConverter +{ + using Type = __nv_bfloat16; +}; + +template <> +struct FloatTypeConverter +{ + using Type = float; +}; + + + +template struct num_elems; +template <> struct num_elems { static constexpr int value = 1; }; +template <> struct num_elems { static constexpr int value = 2; }; +template <> struct num_elems { static constexpr int value = 4; }; +template <> struct num_elems { static constexpr int value = 1; }; +template <> struct num_elems { static constexpr int value = 2; }; +#ifdef ENABLE_BF16 +template <> struct num_elems<__nv_bfloat16> { static constexpr int value = 1; }; +template <> struct num_elems<__nv_bfloat162> { static constexpr int value = 2; }; +#endif +#ifdef ENABLE_FP8 +template <> struct num_elems<__nv_fp8_e4m3> { static constexpr int value = 1; }; +template <> struct num_elems<__nv_fp8x2_e4m3> { static constexpr int value = 2; }; +#endif + +template struct packed_as; +template struct packed_as { using type = T; }; +template<> struct packed_as { using type = half2; }; +template<> struct packed_as { using type = float2; }; +template<> struct packed_as { using type = int16_t; }; +template<> struct packed_as { using type = int2; }; +template<> struct packed_as { using type = half; }; +template<> struct packed_as { using type = float; }; +#ifdef ENABLE_BF16 +template<> struct packed_as<__nv_bfloat16, 2> { using type = __nv_bfloat162; }; +template<> struct packed_as<__nv_bfloat162, 1> { using type = __nv_bfloat16; }; +#endif +#ifdef ENABLE_FP8 +template<> struct packed_as<__nv_fp8_e4m3, 2> { using type = __nv_fp8x2_e4m3; }; +template<> struct packed_as<__nv_fp8x2_e4m3, 1> { using type = __nv_fp8_e4m3; }; +template<> struct packed_as<__nv_fp8_e5m2, 2> { using type = __nv_fp8x2_e5m2; }; +template<> struct packed_as<__nv_fp8x2_e5m2, 1> { using type = __nv_fp8_e5m2; }; +#endif + +inline __device__ float2 operator*(float2 a, float2 b) { return make_float2(a.x * b.x, a.y * b.y); } +inline __device__ float2 operator+(float2 a, float2 b) { return make_float2(a.x + b.x, a.y + b.y); } +inline __device__ float2 operator-(float2 a, float2 b) { return make_float2(a.x - b.x, a.y - b.y); } + +inline __device__ float2 operator*(float2 a, float b) { return make_float2(a.x * b, a.y * b); } +inline __device__ float2 operator+(float2 a, float b) { return make_float2(a.x + b, a.y + b); } +inline __device__ float2 operator-(float2 a, float b) { return make_float2(a.x - b, a.y - b); } + +static inline __device__ int8_t float_to_int8_rn(float x) +{ + uint32_t dst; + asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x)); + return reinterpret_cast(dst); +} + +template +inline __device__ T ldg(const T* val) { + return __ldg(val); +} + +#if ENABLE_BF16 +#define bf1622float2 __bfloat1622float2 +#define float22bf162 __float22bfloat162_rn +#define bf162bf162 __bfloat162bfloat162 +inline __device__ int16_t bf1622int16(__nv_bfloat162 val) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float2 f_val; + f_val.x = max(min(__low2float(val), 127.f), -128.f); + f_val.y = max(min(__high2float(val), 127.f), -128.f); + + union + { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = static_cast(static_cast(f_val.x)); + int8[1] = static_cast(static_cast(f_val.y)); + return int16; +#else + val = __hmin2(val, make_bfloat162(127., 127.)); + val = __hmax2(val, make_bfloat162(-128., -128.)); + + union + { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = static_cast(static_cast(val.x)); + int8[1] = static_cast(static_cast(val.y)); + return int16; +#endif +} +#endif + +#if ENABLE_BF16 +template<> +inline __device__ __nv_bfloat162 ldg(const __nv_bfloat162* val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return val[0]; +#else + return __ldg(val); +#endif +} + +template<> +inline __device__ __nv_bfloat16 ldg(const __nv_bfloat16* val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return val[0]; +#else + return __ldg(val); +#endif +} +#endif // ENABLE_BF16 + +template +__device__ inline T_OUT cuda_cast(T_IN val) +{ + return val; +} + +template <> +__device__ inline float2 cuda_cast(int2 val) +{ + return make_float2(val.x, val.y); +} + +template <> +__device__ inline float2 cuda_cast(float val) +{ + return make_float2(val, val); +} + +template <> +__device__ inline float2 cuda_cast(half2 val) +{ + return __half22float2(val); +} + +template <> +__device__ inline half2 cuda_cast(float2 val) +{ + return __float22half2_rn(val); +} + +template <> +__device__ inline half2 cuda_cast(float val) +{ + return __float2half2_rn(val); +} + +template <> +__device__ inline half2 cuda_cast(half val) +{ + return __half2half2(val); +} + +template <> +__device__ inline int8_t cuda_cast(half val) +{ + union + { + int8_t int8[2]; + int16_t int16; + }; + + union + { + half fp16; + int16_t int16_in; + }; + + fp16 = val; + asm volatile("cvt.rni.sat.s8.f16 %0, %1;" : "=h"(int16) : "h"(int16_in)); + return int8[0]; +} + +template <> +__device__ inline int16_t cuda_cast(half2 val) +{ + union + { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = cuda_cast(val.x); + int8[1] = cuda_cast(val.y); + return int16; +} + +template <> +__device__ inline int8_t cuda_cast(float val) +{ + union + { + int8_t int8[2]; + int16_t int16; + }; + + asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=h"(int16) : "f"(val)); + return int8[0]; +} + +template <> +__device__ inline int16_t cuda_cast(float2 val) +{ + union + { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = cuda_cast(val.x); + int8[1] = cuda_cast(val.y); + return int16; +} + +template <> +__device__ inline half2 cuda_cast(int16_t val) +{ + union + { + int8_t int8[2]; + int16_t int16; + }; + + int16 = val; + return make_half2(int8[0], int8[1]); +} + +template <> +__device__ inline float2 cuda_cast(int16_t val) +{ + union + { + int8_t int8[2]; + int16_t int16; + }; + + int16 = val; + return make_float2(int8[0], int8[1]); +} + +#ifdef ENABLE_BF16 +template <> +__device__ inline __nv_bfloat16 cuda_cast(int32_t val) +{ + return static_cast(val); +} + +template <> +__device__ inline __nv_bfloat16 cuda_cast(int8_t val) +{ + return static_cast(val); +} + +template <> +__device__ inline int8_t cuda_cast(__nv_bfloat16 val) +{ + return static_cast(val); +} + +template <> +__device__ inline float cuda_cast(__nv_bfloat16 val) +{ + return __bfloat162float(val); +} + +template <> +__device__ inline float2 cuda_cast(__nv_bfloat162 val) +{ + return bf1622float2(val); +} + +template <> +__device__ inline half cuda_cast(__nv_bfloat16 val) +{ + return __float2half(__bfloat162float(val)); +} + +template <> +__device__ inline int16_t cuda_cast(__nv_bfloat162 val) +{ + return bf1622int16(val); +} + +template <> +__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, float>(float val) +{ + return __float2bfloat16(val); +} + +template <> +__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, half>(half val) +{ + return __float2bfloat16(__half2float(val)); +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, __nv_bfloat16>(__nv_bfloat16 val) +{ + return bf162bf162(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float>(float val) +{ + return __float2bfloat162_rn(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float2>(float2 val) +{ + return float22bf162(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, int16_t>(int16_t val) +{ + union + { + int8_t int8[2]; + int16_t int16; + }; + + int16 = val; + __nv_bfloat162 res; + res.x = cuda_cast<__nv_bfloat16>(int8[0]); + res.y = cuda_cast<__nv_bfloat16>(int8[1]); + return res; +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, half2>(half2 val) +{ + return float22bf162(__half22float2(val)); +} + +#endif // ENABLE BF16 + +template +__device__ inline To cuda_sum(Ti val) +{ + return cuda_cast(val); +}; + +template +__device__ inline To cuda_sum(float2 val) +{ + return cuda_cast(val.x + val.y); +}; + +// Unary maximum: compute the max of a vector type +template +__device__ inline To cuda_max(Ti val) +{ + return cuda_cast(val); +}; + +template <> +__device__ inline float cuda_max(float2 val) +{ + return fmaxf(val.x, val.y); +} + +template <> +__device__ inline half cuda_max(half2 val) +{ + return __hmax(val.x, val.y); +} + +#ifdef ENABLE_BF16 +template <> +__device__ inline __nv_bfloat16 cuda_max(__nv_bfloat162 val) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)) + return __hmax(val.x, val.y); +#endif +} +#endif + +// Binary maximum: compute the max of two scalar types +template +__device__ inline T cuda_max(T val1, T val2) +{ + return (val1 > val2) ? val1 : val2; +} + +template +__device__ inline T cuda_abs(T val) +{ + assert(false); + return {}; +} + +template <> +__device__ inline float cuda_abs(float val) +{ + return fabs(val); +} + +template <> +__device__ inline float2 cuda_abs(float2 val) +{ + return make_float2(fabs(val.x), fabs(val.y)); +} + +template <> +__device__ inline half cuda_abs(half val) +{ + return __habs(val); +} + +template <> +__device__ inline half2 cuda_abs(half2 val) +{ + return __habs2(val); +} + +#ifdef ENABLE_BF16 + +#if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__) +template <> +__device__ inline __nv_bfloat16 cuda_abs(__nv_bfloat16 val) +{ + return __habs(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_abs(__nv_bfloat162 val) +{ + return __habs2(val); +} +#endif + +#endif // ENABLE_FP16 \ No newline at end of file diff --git a/llm-awq/awq/kernels/csrc/w8a8/w8a8_gemm_cuda.cu b/llm-awq/awq/kernels/csrc/w8a8/w8a8_gemm_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..96a72d7683d4d47f05d848e29ac803176ec5c2b6 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/w8a8_gemm_cuda.cu @@ -0,0 +1,953 @@ +// Inspired by QServe https://github.com/mit-han-lab/qserve/tree/main. +// Modified by Yuming Lou. +// @article{lin2024awq, +// title={AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration}, +// author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Chen, Wei-Ming and Wang, Wei-Chen and Xiao, Guangxuan and Dang, Xingyu and Gan, Chuang and Han, Song}, +// journal={Proceedings of Machine Learning and Systems}, +// volume={6}, +// pages={87--100}, +// year={2024} +// } + +#include "w8a8_gemm_cuda.h" +#include +#include +#include + +#define OP_M 16 +#define OP_N 8 +#define OP_K 32 +#define INTRIN_M 16 +#define INTRIN_N 16 +#define INTRIN_K 32 +#define WARP_SIZE 32 +#define SMEM_PAD_A 0 +#define SMEM_PAD_B 0 +#define PACK_SIZE 16 +#if (__CUDACC_VER_MAJOR__ >= 11) && (__CUDACC_VER_MINOR__ >= 4) +#define L2_CACHEHINT(size) ".L2::" #size "B" +#else +#define L2_CACHEHINT(size) +#endif +#define KERNEL_LAUNCH_CODE_FUSE_BIAS \ + constexpr int NUM_WARPS = (CTA_M / WARP_M) * (CTA_N / WARP_N) * (CTA_K / WARP_K); \ + constexpr int kSmemByteSize = \ + (CTA_M * (CTA_K + SMEM_PAD_A) + CTA_N * (CTA_K + SMEM_PAD_B)) * STAGES * \ + sizeof(int8_t) + CTA_N * sizeof(float); \ + if (kSmemByteSize >= 99 * 1024) \ + { \ + printf("This kernel requires %d Bytes of shared memory, which exceeds " \ + "device limit.\n", \ + kSmemByteSize); \ + return ; \ + } \ + int num_blocks_m = (num_out_feats + CTA_M - 1) / CTA_M; \ + int num_blocks_n = (num_out_channels+ CTA_N - 1) / CTA_N / 1; \ + const int log_tile = get_log_tile<8>((num_out_feats + CTA_M - 1) / CTA_M); \ + const int tile_shift = 1 << log_tile; \ + dim3 num_blocks(num_blocks_n *tile_shift, \ + (num_blocks_m + tile_shift - 1) / tile_shift); \ + dim3 threads_per_block(WARP_SIZE, NUM_WARPS); \ + auto kernel_func = \ + dense_kernel0_fuse_bias; \ + cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, \ + kSmemByteSize); \ + kernel_func<<>>( \ + in_feats, kernel, wscales, ascales, out_feats, bias, num_in_feats, num_out_channels, \ + num_in_channels); + + +#define KERNEL_LAUNCH_CODE \ + constexpr int NUM_WARPS = (CTA_M / WARP_M) * (CTA_N / WARP_N) * (CTA_K / WARP_K); \ + constexpr int kSmemByteSize = \ + (CTA_M * (CTA_K + SMEM_PAD_A) + CTA_N * (CTA_K + SMEM_PAD_B)) * STAGES * \ + sizeof(int8_t); \ + if (kSmemByteSize >= 99 * 1024) \ + { \ + printf("This kernel requires %d Bytes of shared memory, which exceeds " \ + "device limit.\n", \ + kSmemByteSize); \ + return ; \ + } \ + int num_blocks_m = (num_out_feats + CTA_M - 1) / CTA_M; \ + int num_blocks_n = num_out_channels / CTA_N / 1; \ + const int log_tile = get_log_tile<8>((num_out_feats + CTA_M - 1) / CTA_M); \ + const int tile_shift = 1 << log_tile; \ + dim3 num_blocks(num_blocks_n *tile_shift, \ + (num_blocks_m + tile_shift - 1) / tile_shift); \ + dim3 threads_per_block(WARP_SIZE, NUM_WARPS); \ + auto kernel_func = \ + dense_kernel0; \ + cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, \ + kSmemByteSize); \ + kernel_func<<>>( \ + in_feats, kernel, wscales, ascales, out_feats, num_in_feats, num_out_channels, \ + num_in_channels); + + + +template +__inline__ __host__ __device__ int get_log_tile(int n) +{ + if (N >= 8 && n >= 6) + return 3; + else if (N >= 4 && n >= 3) + return 2; + else if (N >= 2 && n >= 2) + return 1; + else + return 0; +} + +__inline__ __device__ uint2 get_block_idx_mapping(int blockIdx_x, + int blockIdx_y, + int log_tile) +{ + return make_uint2((blockIdx_x >> log_tile), + (blockIdx_y << log_tile) + + ((blockIdx_x) & ((1 << (log_tile)) - 1))); +} + +__inline__ __device__ uint32_t cast_smem_ptr_to_uint(void const *const ptr) +{ + uint32_t smem_int_ptr; + + asm("{.reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, " + "smem_ptr; }\n" + : "=r"(smem_int_ptr) + : "l"(ptr)); + + return smem_int_ptr; +} + +__inline__ __device__ void ldmatrix_m8n8_x4_b16(int8_t *shared_warp, int ax0_0, + uint32_t addr) +{ + __asm__ __volatile__("ldmatrix.sync.aligned.m8n8.x4.shared.b16" + "{%0, %1, %2, %3}, [%4];" + : "=r"(((unsigned *)(shared_warp + (ax0_0 * 16)))[0]), + "=r"(((unsigned *)(shared_warp + (ax0_0 * 16)))[1]), + "=r"(((unsigned *)(shared_warp + (ax0_0 * 16)))[2]), + "=r"(((unsigned *)(shared_warp + (ax0_0 * 16)))[3]) + : "r"(addr)); +} + +__inline__ __device__ void +ldmatrix_m8n8_x4_trans_b16(int8_t *shared_warp, int ax0_0, uint32_t addr) +{ + __asm__ __volatile__("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16" + "{%0, %1, %2, %3}, [%4];" + : "=r"(((unsigned *)(shared_warp + (ax0_0 * 16)))[0]), + "=r"(((unsigned *)(shared_warp + (ax0_0 * 16)))[1]), + "=r"(((unsigned *)(shared_warp + (ax0_0 * 16)))[2]), + "=r"(((unsigned *)(shared_warp + (ax0_0 * 16)))[3]) + : "r"(addr)); +} + +// function from lmdeploy +__inline__ __device__ void +cp_async_cg_A(uint32_t smem_int_ptr, const uint4 *__restrict__ src, bool mask)//256 * int8 +{ + const int cp_size = 16; + asm volatile("{" + " .reg .pred p;" + " setp.ne.b32 p, %0, 0;" + " @p cp.async.cg.shared.global" L2_CACHEHINT(128) " [%1], [%2], %3;" + "}" ::"r"((int)mask), + "r"(smem_int_ptr), + "l"(src), + "n"(cp_size)); +} + +__device__ __inline__ void mma_m16n8k32(void *C_warp, void *A_shared_warp, + void *B_shared_warp) +{ + __asm__ __volatile__( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32" + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};" + : "=r"(((int *)C_warp)[0]), "=r"(((int *)C_warp)[1]), + "=r"(((int *)C_warp)[2]), "=r"(((int *)C_warp)[3]) + : "r"(((unsigned *)A_shared_warp)[0]), + "r"(((unsigned *)A_shared_warp)[1]), + "r"(((unsigned *)A_shared_warp)[2]), + "r"(((unsigned *)A_shared_warp)[3]), + "r"(((unsigned *)B_shared_warp)[0]), + "r"(((unsigned *)B_shared_warp)[1]), "r"(((int *)C_warp)[0]), + "r"(((int *)C_warp)[1]), "r"(((int *)C_warp)[2]), + "r"(((int *)C_warp)[3])); +} + +template +__device__ __inline__ void +global_to_share_one_stage_A(int8_t *src, int8_t *dst, int global_ncols, + int cta_offset_m, int cta_offset_n, + int global_iter_k, int shared_iter_k, bool mask, + bool *preds) +{ + constexpr int total_global_iters = (CTA_M * CTA_K) / PACK_SIZE / CTA_SIZE; + constexpr int partial_global_iters = total_global_iters / SHARED_K_ITERS; + constexpr int cta_step_m_or_n = (CTA_SIZE * PACK_SIZE) / CTA_K; + constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int threads_per_row = CTA_K / PACK_SIZE; + constexpr int kSmemCol = CTA_K + SMEM_PAD_A; + int8_t *dst_hoisted = dst; + int8_t *src_hoisted = src + global_iter_k * CTA_K; + + if (mask) + { +#pragma unroll + for (int _global_iter = 0; _global_iter < partial_global_iters; + ++_global_iter) + { + int global_iter = shared_iter_k * partial_global_iters + _global_iter; + + void *dst_ptr = + (void *)(dst_hoisted + global_iter * cta_step_m_or_n * kSmemCol); + uint4 *src_ptr = + (uint4 *)(src_hoisted + global_iter * cta_step_m_or_n * global_ncols); + if constexpr (STAGES > 1) + { + uint32_t addr = cast_smem_ptr_to_uint(dst_ptr); + cp_async_cg_A(addr, src_ptr, preds[global_iter]); + } + else + { + if (preds[global_iter]) + *(uint4 *)dst_ptr = *src_ptr; + } + } + } +} + +template +__device__ __inline__ void +global_to_share_one_stage_B(int8_t *src, int8_t *dst, int global_ncols, + int cta_offset_m, int cta_offset_n, + int global_iter_k, int shared_iter_k, bool mask, bool *preds) +{ + constexpr int total_global_iters = (CTA_N * CTA_K) / PACK_SIZE / CTA_SIZE; + constexpr int partial_global_iters = total_global_iters / SHARED_K_ITERS; + constexpr int cta_step_m_or_n = (CTA_SIZE * PACK_SIZE) / CTA_K; + constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int threads_per_row = CTA_K / PACK_SIZE; + constexpr int kSmemCol = CTA_K + SMEM_PAD_B; + int8_t *dst_hoisted = dst; + int8_t *src_hoisted = src + global_iter_k * CTA_K; +#pragma unroll + for (int _global_iter = 0; _global_iter < partial_global_iters; + ++_global_iter) + { + int global_iter = shared_iter_k * partial_global_iters + _global_iter; + + void *dst_ptr = + (void *)(dst_hoisted + global_iter * cta_step_m_or_n * kSmemCol); + uint4 *src_ptr = + (uint4 *)(src_hoisted + global_iter * cta_step_m_or_n * global_ncols); + if constexpr (STAGES > 1) + { + uint32_t addr = cast_smem_ptr_to_uint(dst_ptr); + cp_async_cg_A(addr, src_ptr, preds[global_iter]); + } + else + { + if (preds[global_iter]) + *(uint4 *)dst_ptr = *src_ptr; + } + } +} + +template +__device__ __inline__ void +share_to_reg_one_stage_A(int8_t *src, int8_t *dst, int warp_offset_m, + int warp_offset_n, int k_0_1, int shared_iters) +{ + constexpr int kSmemCol = CTA_K + SMEM_PAD_A; + int ld_col = (k_0_1 * INTRIN_K + (threadIdx.x / 16) * 16) / PACK_SIZE; + + for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) + { + int ld_row = warp_offset_m + shared_iter * INTRIN_M + (threadIdx.x % 16); + int ld_col_swizzled = ld_col ^ (ld_row / 2) & 3; + void *addr_ptr = + (void *)(src + ld_row * kSmemCol + ld_col_swizzled * PACK_SIZE); + uint32_t addr = cast_smem_ptr_to_uint(addr_ptr); + ldmatrix_m8n8_x4_b16(dst, shared_iter, addr); + } +} + +template +__device__ __inline__ void +share_to_reg_one_stage_B(int8_t *src, int8_t *dst, int warp_offset_m, + int warp_offset_n, int k_0_1, int shared_iters) +{ + constexpr int kSmemCol = CTA_K + SMEM_PAD_B; + int ld_col = (k_0_1 * INTRIN_K + ((threadIdx.x / 8) % 2) * 16) / PACK_SIZE; + + for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) + { + int ld_row = warp_offset_n + shared_iter * INTRIN_N + ((threadIdx.x / 8 / 2) * 8 + threadIdx.x % 8); + int ld_col_swizzled = ld_col ^ (ld_row / 2) & 3; + void *addr_ptr = + (void *)(src + ld_row * kSmemCol + ld_col_swizzled * PACK_SIZE); + uint32_t addr = cast_smem_ptr_to_uint(addr_ptr); + ldmatrix_m8n8_x4_b16(dst, shared_iter, addr); + } +} + +template +__global__ void dense_kernel0_fuse_bias(int8_t *__restrict__ A, int8_t *__restrict__ B, + half2 *__restrict__ wscales, half *__restrict__ ascales, + half *__restrict__ C, half *__restrict__ Bias, + int M, int N, int K) +{ + constexpr int NUM_WARPS_MN = CTA_M / WARP_M * CTA_N / WARP_N; + constexpr int NUM_WARPS = NUM_WARPS_MN * CTA_K / WARP_K; + constexpr int CTA_SIZE = NUM_WARPS * WARP_SIZE; + constexpr int CTA_SIZE_MN = NUM_WARPS_MN * WARP_SIZE; + constexpr int SLICES = CTA_K / WARP_K; + int num_blocks_n = (N + CTA_N - 1) / CTA_N; + int num_blocks_m = (M + CTA_M - 1) / CTA_M; + + int blockIdx_n = blockIdx.x; + int blockIdx_m = blockIdx.y; + const int log_tile = get_log_tile<8>((M + CTA_M - 1) / CTA_M); + const uint2 block_idx_mapping = + get_block_idx_mapping(blockIdx_n, blockIdx_m, log_tile); + blockIdx_n = block_idx_mapping.x; + blockIdx_m = block_idx_mapping.y; + + int C_warp[CTA_M * CTA_N / CTA_SIZE_MN]; + constexpr int kSmemPadKA = CTA_K + SMEM_PAD_A; + constexpr int kSmemPadKB = CTA_K + SMEM_PAD_B; + constexpr int kSmemSizeAPerStage = CTA_M * kSmemPadKA; + constexpr int kSmemSizeBPerStage = CTA_N * kSmemPadKB; + constexpr int kSmemSizeA = kSmemSizeAPerStage * STAGES; + constexpr int kSmemSizeB = kSmemSizeBPerStage * STAGES; + extern __shared__ int8_t mem_shared[]; + int8_t *A_shared = mem_shared; + int8_t *B_shared = mem_shared + kSmemSizeA; + float *Bias_shared= reinterpret_cast(mem_shared + kSmemSizeA + kSmemSizeB); + int8_t A_shared_warp_[2][WARP_M * WARP_K / + WARP_SIZE]; + int8_t B_shared_warp_[2][WARP_N * WARP_K / + WARP_SIZE]; + constexpr int A_total_global_iters = (CTA_M * CTA_K) / PACK_SIZE / CTA_SIZE; + constexpr int B_total_global_iters = (CTA_N * CTA_K) / PACK_SIZE / CTA_SIZE; + constexpr int A_src_step_m = (CTA_SIZE * PACK_SIZE) / CTA_K; + constexpr int B_src_step_k = (CTA_SIZE * PACK_SIZE) / CTA_K; + constexpr int A_warp_step_m = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int B_warp_step_n = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int A_threads_per_row = CTA_K / PACK_SIZE; + constexpr int B_threads_per_row = CTA_K / PACK_SIZE; + int cta_offset_m = blockIdx_m * CTA_M; + int cta_offset_n = blockIdx_n * CTA_N; + int warp_mn = threadIdx.y % NUM_WARPS_MN; + int slice_id = threadIdx.y / NUM_WARPS_MN; // Always zero if threadIdx.z==0! + int warp_offset_m = (warp_mn % (CTA_M / WARP_M)) * WARP_M; + int warp_offset_n = (warp_mn / (CTA_M / WARP_M)) * WARP_N; + int warp_offset_k = slice_id * WARP_K; + + for (int i = 0; i < CTA_M * CTA_N / CTA_SIZE_MN; i++) + C_warp[i] = 0; + + int gemm_iters = (K + CTA_K - 1) / CTA_K; + int k_0_0_ld = 0; + int k_0_0 = 0; + constexpr int prologue_stages = STAGES == 1 ? 1 : STAGES - 1; + int A_hoisted_row = threadIdx.y * A_warp_step_m + (threadIdx.x / A_threads_per_row); + int A_hoisted_col = (threadIdx.x % A_threads_per_row); + int A_hoisted_col_swizzled = A_hoisted_col ^ (A_hoisted_row / 2) & 3; + + int B_hoisted_row = threadIdx.y * B_warp_step_n + (threadIdx.x / B_threads_per_row); + int B_hoisted_col = (threadIdx.x % B_threads_per_row); + int B_hoisted_col_swizzled = B_hoisted_col ^ (B_hoisted_row / 2) & 3; + + int8_t *A_shared_hoisted = A_shared + + A_hoisted_row * kSmemPadKA + + A_hoisted_col_swizzled * PACK_SIZE; + int8_t *B_shared_hoisted = B_shared + B_hoisted_row * kSmemPadKB + + B_hoisted_col_swizzled * PACK_SIZE; + int8_t *A_hoisted = A + cta_offset_m * K + A_hoisted_row * K + + A_hoisted_col * PACK_SIZE; + int8_t *B_hoisted = B + cta_offset_n * K + B_hoisted_row * K + + B_hoisted_col * PACK_SIZE; + bool A_g2s_preds[A_total_global_iters]; + bool B_g2s_preds[B_total_global_iters]; + //debug + // printf("A: %d ",A_total_global_iters); + // printf("B: %d ",B_total_global_iters); + // printf("prologue_stages: %d ",prologue_stages); + // __shared__ float2 Bias_shared[CTA_N]; + #pragma unroll + for (int i = 0; i < CTA_N ; i++) + { + Bias_shared[i] = __half2float(Bias[cta_offset_n+i]); + } + + +#pragma unroll + for (int i = 0; i < A_total_global_iters; i++) + { + A_g2s_preds[i] = (cta_offset_m + A_hoisted_row + i * A_src_step_m) < M; + } + #pragma unroll + for (int i = 0; i < B_total_global_iters; i++) + { + B_g2s_preds[i] = cta_offset_n + B_hoisted_row + i * B_src_step_k < N; + } + int *C_shared = reinterpret_cast(mem_shared); +#pragma unroll + for (k_0_0_ld = 0; k_0_0_ld < prologue_stages; ++k_0_0_ld) + { + global_to_share_one_stage_A( + A_hoisted, A_shared_hoisted + k_0_0_ld * kSmemSizeAPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, 0, true, A_g2s_preds); + global_to_share_one_stage_B( + B_hoisted, B_shared_hoisted + k_0_0_ld * kSmemSizeBPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, 0, true, B_g2s_preds); + if constexpr (STAGES > 1) + __pipeline_commit(); + } + if constexpr (STAGES > 1) + __pipeline_wait_prior(STAGES - 2); + __syncthreads(); + +// global_to_share_bias(Bias,Bias_shared,cta_offset_n); + + share_to_reg_one_stage_A( + A_shared + warp_offset_k, A_shared_warp_[0], warp_offset_m, warp_offset_n, 0, + WARP_M / INTRIN_M); + share_to_reg_one_stage_B( + B_shared + warp_offset_k, B_shared_warp_[0], warp_offset_m, warp_offset_n, 0, + WARP_N / INTRIN_N); + constexpr int SHARED_K_ITERS = WARP_K / INTRIN_K; + + for (; k_0_0 < gemm_iters; ++k_0_0, ++k_0_0_ld) + { + int ld_stage = k_0_0_ld % STAGES; + int compute_stage = k_0_0 % STAGES; + int8_t *A_shared_this_compute_stage; + int8_t *B_shared_this_compute_stage; + + for (int iter_k = 0; iter_k < SHARED_K_ITERS; ++iter_k) + { + A_shared_this_compute_stage = + A_shared + compute_stage * kSmemSizeAPerStage + warp_offset_k; + B_shared_this_compute_stage = + B_shared + compute_stage * kSmemSizeBPerStage + warp_offset_k; + share_to_reg_one_stage_A( + A_shared_this_compute_stage, A_shared_warp_[(iter_k + 1) % 2], + warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS, + WARP_M / INTRIN_M); + share_to_reg_one_stage_B( + B_shared_this_compute_stage, B_shared_warp_[(iter_k + 1) % 2], + warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS, + WARP_N / INTRIN_N); + int8_t *A_shared_warp = A_shared_warp_[iter_k % 2]; + int8_t *B_shared_warp = B_shared_warp_[iter_k % 2]; + for (int i_0_3 = 0; i_0_3 < WARP_M / INTRIN_M; ++i_0_3) + { + for (int j_0_4 = 0; j_0_4 < WARP_N / INTRIN_N; ++j_0_4) + { + mma_m16n8k32( + (void *)(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8), + (void *)(A_shared_warp + i_0_3 * 16), + (void *)(B_shared_warp + j_0_4 * 16)); + mma_m16n8k32( + (void *)(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8 + 4), + (void *)(A_shared_warp + i_0_3 * 16), + (void *)(B_shared_warp + j_0_4 * 16 + 8)); + } + } + + if (iter_k < SHARED_K_ITERS - 1) + { + if constexpr (STAGES == 1) + __syncthreads(); + global_to_share_one_stage_A( + A_hoisted, A_shared_hoisted + ld_stage * kSmemSizeAPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, iter_k, + k_0_0_ld < gemm_iters, A_g2s_preds); + global_to_share_one_stage_B( + B_hoisted, B_shared_hoisted + ld_stage * kSmemSizeBPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, iter_k, + k_0_0_ld < gemm_iters, B_g2s_preds); + } + + if (iter_k == SHARED_K_ITERS - 2) + { + if constexpr (STAGES == 1 && SHARED_K_ITERS > 2) + { + __syncthreads(); + } + global_to_share_one_stage_A( + A_hoisted, A_shared_hoisted + ld_stage * kSmemSizeAPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, iter_k + 1, + k_0_0_ld < gemm_iters, A_g2s_preds); + global_to_share_one_stage_B( + B_hoisted, B_shared_hoisted + ld_stage * kSmemSizeBPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, iter_k + 1, + k_0_0_ld < gemm_iters, B_g2s_preds); + if constexpr (STAGES > 1) + { + __pipeline_commit(); + __pipeline_wait_prior(STAGES - 2); + } + compute_stage = (k_0_0 + 1) % STAGES; + __syncthreads(); + } + } + } + __pipeline_commit(); + __pipeline_wait_prior(0); + __syncthreads(); + + if constexpr (SLICES > 1) + { +#pragma unroll + for (int z = 0; z < SLICES; ++z) + { + if (slice_id == z) + { +#pragma unroll + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { +#pragma unroll + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { +#pragma unroll + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) + { + if (z > 0) + { + C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] += C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2]; + } + C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2] = C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id]; + }; + } + } + } + __syncthreads(); + } + if (slice_id == 0) + { +#pragma unroll + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { +#pragma unroll + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { +#pragma unroll + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) + { + C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] = C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2]; + }; + } + } + } + } + + int row_wb_thd = cta_offset_m + warp_offset_m + (threadIdx.x / 4); + int col_wb_thd = cta_offset_n + warp_offset_n + (threadIdx.x % 4) * 2; + if (slice_id == 0) + { + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + int row_wb_1 = row_wb_thd + ax0_0_1 * OP_M; + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + int col_wb_1 = col_wb_thd + ax1_0_1 * 16; + int *C_warp_local = C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8; + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2) + { + int row_wb = row_wb_1 + (local_id % 4) / 2 * 8; + int col_wb = col_wb_1 + (local_id / 4) * 8 + (local_id % 2); + if (row_wb < M && col_wb < N ){ + float2 wscale = __half22float2(*(wscales + col_wb / 2)); + float ascale = __half2float(ascales[row_wb]); + float2 psums = make_float2(__int2float_rn(C_warp_local[local_id]), __int2float_rn(C_warp_local[local_id + 1])); + psums.x = psums.x * wscale.x * ascale + Bias_shared[col_wb % CTA_N]; + psums.y = psums.y * wscale.y * ascale + Bias_shared[col_wb % CTA_N + 1]; + *reinterpret_cast(C + row_wb * N + col_wb) = __float22half2_rn(psums); + } + }; + } + } + } +} + +void w8a8_gemm_fuse_bias_forward_cuda(torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _wscales, + torch::Tensor _ascales, + torch::Tensor _out_feats, + torch::Tensor _bias) +{ + int num_in_feats = _in_feats.size(0); + int num_in_channels = _in_feats.size(1); + auto in_feats = reinterpret_cast(_in_feats.data_ptr()); + auto kernel = reinterpret_cast(_kernel.data_ptr()); + auto wscales = reinterpret_cast(_wscales.data_ptr()); + auto ascales = reinterpret_cast(_ascales.data_ptr()); + auto bias = reinterpret_cast(_bias.data_ptr()); + // auto options = + // torch::TensorOptions().dtype(torch::kFloat16).device(_in_feats.device()); + // at::Tensor _out_feats = + // torch::empty({num_in_feats, _kernel.size(0)}, options); + int num_out_feats = _out_feats.size(-2); + int num_out_channels = _out_feats.size(-1); + + + auto out_feats = reinterpret_cast(_out_feats.data_ptr()); + + if (num_out_feats > 128) + { + constexpr int CTA_M = 128; + constexpr int CTA_N = 128; + constexpr int CTA_K = 64; + constexpr int WARP_M = 64; + constexpr int WARP_N = 32; + constexpr int WARP_K = 64; + constexpr int STAGES = 6; + KERNEL_LAUNCH_CODE_FUSE_BIAS + } + else + { + constexpr int CTA_M = 64; + constexpr int CTA_N = 64; + constexpr int CTA_K = 64; + constexpr int WARP_M = 32; + constexpr int WARP_N = 16; + constexpr int WARP_K = 64; + constexpr int STAGES = 6; + KERNEL_LAUNCH_CODE_FUSE_BIAS + } + return ; +} + +template +__global__ void dense_kernel0(int8_t *__restrict__ A, int8_t *__restrict__ B, + half2 *__restrict__ wscales, half *__restrict__ ascales, + half *__restrict__ C, int M, int N, int K) +{ + constexpr int NUM_WARPS_MN = CTA_M / WARP_M * CTA_N / WARP_N; + constexpr int NUM_WARPS = NUM_WARPS_MN * CTA_K / WARP_K; + constexpr int CTA_SIZE = NUM_WARPS * WARP_SIZE; + constexpr int CTA_SIZE_MN = NUM_WARPS_MN * WARP_SIZE; + constexpr int SLICES = CTA_K / WARP_K; + int num_blocks_n = (N + CTA_N - 1) / CTA_N; + int num_blocks_m = (M + CTA_M - 1) / CTA_M; + + int blockIdx_n = blockIdx.x; + int blockIdx_m = blockIdx.y; + const int log_tile = get_log_tile<8>((M + CTA_M - 1) / CTA_M); + const uint2 block_idx_mapping = + get_block_idx_mapping(blockIdx_n, blockIdx_m, log_tile); + blockIdx_n = block_idx_mapping.x; + blockIdx_m = block_idx_mapping.y; + + int C_warp[CTA_M * CTA_N / CTA_SIZE_MN]; + constexpr int kSmemPadKA = CTA_K + SMEM_PAD_A; + constexpr int kSmemPadKB = CTA_K + SMEM_PAD_B; + constexpr int kSmemSizeAPerStage = CTA_M * kSmemPadKA; + constexpr int kSmemSizeBPerStage = CTA_N * kSmemPadKB; + constexpr int kSmemSizeA = kSmemSizeAPerStage * STAGES; + constexpr int kSmemSizeB = kSmemSizeBPerStage * STAGES; + extern __shared__ int8_t mem_shared[]; + int8_t *A_shared = mem_shared; + int8_t *B_shared = mem_shared + kSmemSizeA; + int8_t A_shared_warp_[2][WARP_M * WARP_K / + WARP_SIZE]; + int8_t B_shared_warp_[2][WARP_N * WARP_K / + WARP_SIZE]; + constexpr int A_total_global_iters = (CTA_M * CTA_K) / PACK_SIZE / CTA_SIZE; + constexpr int B_total_global_iters = (CTA_N * CTA_K) / PACK_SIZE / CTA_SIZE; + constexpr int A_src_step_m = (CTA_SIZE * PACK_SIZE) / CTA_K; + constexpr int B_src_step_k = (CTA_SIZE * PACK_SIZE) / CTA_K; + constexpr int A_warp_step_m = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int B_warp_step_n = (WARP_SIZE * PACK_SIZE) / CTA_K; + constexpr int A_threads_per_row = CTA_K / PACK_SIZE; + constexpr int B_threads_per_row = CTA_K / PACK_SIZE; + int cta_offset_m = blockIdx_m * CTA_M; + int cta_offset_n = blockIdx_n * CTA_N; + int warp_mn = threadIdx.y % NUM_WARPS_MN; + int slice_id = threadIdx.y / NUM_WARPS_MN; + int warp_offset_m = (warp_mn % (CTA_M / WARP_M)) * WARP_M; + int warp_offset_n = (warp_mn / (CTA_M / WARP_M)) * WARP_N; + int warp_offset_k = slice_id * WARP_K; + + for (int i = 0; i < CTA_M * CTA_N / CTA_SIZE_MN; i++) + C_warp[i] = 0; + + int gemm_iters = (K + CTA_K - 1) / CTA_K; + int k_0_0_ld = 0; + int k_0_0 = 0; + constexpr int prologue_stages = STAGES == 1 ? 1 : STAGES - 1; + int A_hoisted_row = threadIdx.y * A_warp_step_m + (threadIdx.x / A_threads_per_row); + int A_hoisted_col = (threadIdx.x % A_threads_per_row); + int A_hoisted_col_swizzled = A_hoisted_col ^ (A_hoisted_row / 2) & 3; + + int B_hoisted_row = threadIdx.y * B_warp_step_n + (threadIdx.x / B_threads_per_row); + int B_hoisted_col = (threadIdx.x % B_threads_per_row); + int B_hoisted_col_swizzled = B_hoisted_col ^ (B_hoisted_row / 2) & 3; + + int8_t *A_shared_hoisted = A_shared + + A_hoisted_row * kSmemPadKA + + A_hoisted_col_swizzled * PACK_SIZE; + int8_t *B_shared_hoisted = B_shared + B_hoisted_row * kSmemPadKB + + B_hoisted_col_swizzled * PACK_SIZE; + int8_t *A_hoisted = A + cta_offset_m * K + A_hoisted_row * K + + A_hoisted_col * PACK_SIZE; + int8_t *B_hoisted = B + cta_offset_n * K + B_hoisted_row * K + + B_hoisted_col * PACK_SIZE; + bool A_g2s_preds[A_total_global_iters]; +#pragma unroll + for (int i = 0; i < A_total_global_iters; i++) + { + A_g2s_preds[i] = (cta_offset_m + A_hoisted_row + i * A_src_step_m) < M; + } + bool B_g2s_preds[B_total_global_iters]; + #pragma unroll + for (int i = 0; i < B_total_global_iters; i++) + { + B_g2s_preds[i] = (cta_offset_n + B_hoisted_col + i) < N; + } + int *C_shared = reinterpret_cast(mem_shared); +#pragma unroll + for (k_0_0_ld = 0; k_0_0_ld < prologue_stages; ++k_0_0_ld) + { + global_to_share_one_stage_A( + A_hoisted, A_shared_hoisted + k_0_0_ld * kSmemSizeAPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, 0, true, A_g2s_preds); + global_to_share_one_stage_B( + B_hoisted, B_shared_hoisted + k_0_0_ld * kSmemSizeBPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, 0, true, B_g2s_preds); + if constexpr (STAGES > 1) + __pipeline_commit(); + } + if constexpr (STAGES > 1) + __pipeline_wait_prior(STAGES - 2); + __syncthreads(); + + share_to_reg_one_stage_A( + A_shared + warp_offset_k, A_shared_warp_[0], warp_offset_m, warp_offset_n, 0, + WARP_M / INTRIN_M); + share_to_reg_one_stage_B( + B_shared + warp_offset_k, B_shared_warp_[0], warp_offset_m, warp_offset_n, 0, + WARP_N / INTRIN_N); + constexpr int SHARED_K_ITERS = WARP_K / INTRIN_K; + + for (; k_0_0 < gemm_iters; ++k_0_0, ++k_0_0_ld) + { + int ld_stage = k_0_0_ld % STAGES; + int compute_stage = k_0_0 % STAGES; + int8_t *A_shared_this_compute_stage; + int8_t *B_shared_this_compute_stage; + + for (int iter_k = 0; iter_k < SHARED_K_ITERS; ++iter_k) + { + A_shared_this_compute_stage = + A_shared + compute_stage * kSmemSizeAPerStage + warp_offset_k; + B_shared_this_compute_stage = + B_shared + compute_stage * kSmemSizeBPerStage + warp_offset_k; + share_to_reg_one_stage_A( + A_shared_this_compute_stage, A_shared_warp_[(iter_k + 1) % 2], + warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS, + WARP_M / INTRIN_M); + share_to_reg_one_stage_B( + B_shared_this_compute_stage, B_shared_warp_[(iter_k + 1) % 2], + warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS, + WARP_N / INTRIN_N); + int8_t *A_shared_warp = A_shared_warp_[iter_k % 2]; + int8_t *B_shared_warp = B_shared_warp_[iter_k % 2]; + for (int i_0_3 = 0; i_0_3 < WARP_M / INTRIN_M; ++i_0_3) + { + for (int j_0_4 = 0; j_0_4 < WARP_N / INTRIN_N; ++j_0_4) + { + mma_m16n8k32( + (void *)(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8), + (void *)(A_shared_warp + i_0_3 * 16), + (void *)(B_shared_warp + j_0_4 * 16)); + mma_m16n8k32( + (void *)(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8 + 4), + (void *)(A_shared_warp + i_0_3 * 16), + (void *)(B_shared_warp + j_0_4 * 16 + 8)); + } + } + + if (iter_k < SHARED_K_ITERS - 1) + { + if constexpr (STAGES == 1) + __syncthreads(); + global_to_share_one_stage_A( + A_hoisted, A_shared_hoisted + ld_stage * kSmemSizeAPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, iter_k, + k_0_0_ld < gemm_iters, A_g2s_preds); + global_to_share_one_stage_B( + B_hoisted, B_shared_hoisted + ld_stage * kSmemSizeBPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, iter_k, + k_0_0_ld < gemm_iters, B_g2s_preds); + } + + if (iter_k == SHARED_K_ITERS - 2) + { + if constexpr (STAGES == 1 && SHARED_K_ITERS > 2) + { + __syncthreads(); + } + global_to_share_one_stage_A( + A_hoisted, A_shared_hoisted + ld_stage * kSmemSizeAPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, iter_k + 1, + k_0_0_ld < gemm_iters, A_g2s_preds); + global_to_share_one_stage_B( + B_hoisted, B_shared_hoisted + ld_stage * kSmemSizeBPerStage, K, + cta_offset_m, cta_offset_n, k_0_0_ld, iter_k + 1, + k_0_0_ld < gemm_iters, B_g2s_preds); + if constexpr (STAGES > 1) + { + __pipeline_commit(); + __pipeline_wait_prior(STAGES - 2); + } + compute_stage = (k_0_0 + 1) % STAGES; + __syncthreads(); + } + } + } + __pipeline_commit(); + __pipeline_wait_prior(0); + __syncthreads(); + + if constexpr (SLICES > 1) + { +#pragma unroll + for (int z = 0; z < SLICES; ++z) + { + if (slice_id == z) + { +#pragma unroll + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { +#pragma unroll + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { +#pragma unroll + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) + { + if (z > 0) + { + C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] += C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2]; + } + C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2] = C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id]; + }; + } + } + } + __syncthreads(); + } + if (slice_id == 0) + { +#pragma unroll + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { +#pragma unroll + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { +#pragma unroll + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) + { + C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] = C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2]; + }; + } + } + } + } + + int row_wb_thd = cta_offset_m + warp_offset_m + (threadIdx.x / 4); + int col_wb_thd = cta_offset_n + warp_offset_n + (threadIdx.x % 4) * 2; + if (slice_id == 0) + { + for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) + { + int row_wb_1 = row_wb_thd + ax0_0_1 * OP_M; + for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) + { + int col_wb_1 = col_wb_thd + ax1_0_1 * 16; + int *C_warp_local = C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8; + for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2) + { + int row_wb = row_wb_1 + (local_id % 4) / 2 * 8; + int col_wb = col_wb_1 + (local_id / 4) * 8 + (local_id % 2); + if (row_wb < M && col_wb < N){ + int col_wb = col_wb_1 + (local_id / 4) * 8 + (local_id % 2); + float2 wscale = __half22float2(*(wscales + col_wb / 2)); + float ascale = __half2float(ascales[row_wb]); + float2 psums = make_float2(__int2float_rn(C_warp_local[local_id]), __int2float_rn(C_warp_local[local_id + 1])); + psums.x *= wscale.x * ascale; + psums.y *= wscale.y * ascale; + *reinterpret_cast(C + row_wb * N + col_wb) = __float22half2_rn(psums); + } + }; + } + } + } +} + +void w8a8_gemm_forward_cuda(torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _wscales, + torch::Tensor _ascales, + torch::Tensor _out_feats) +{ + int num_in_feats = _in_feats.size(0); + int num_in_channels = _in_feats.size(1); + auto in_feats = reinterpret_cast(_in_feats.data_ptr()); + auto kernel = reinterpret_cast(_kernel.data_ptr()); + auto wscales = reinterpret_cast(_wscales.data_ptr()); + auto ascales = reinterpret_cast(_ascales.data_ptr()); + + // auto options = + // torch::TensorOptions().dtype(torch::kFloat16).device(_in_feats.device()); + // at::Tensor _out_feats = + // torch::empty({num_in_feats, _kernel.size(0)}, options); + int num_out_feats = _out_feats.size(-2); + int num_out_channels = _out_feats.size(-1); + + + auto out_feats = reinterpret_cast(_out_feats.data_ptr()); + + if (num_out_feats > 128) + { + constexpr int CTA_M = 128; + constexpr int CTA_N = 128; + constexpr int CTA_K = 64; + constexpr int WARP_M = 128; + constexpr int WARP_N = 32; + constexpr int WARP_K = 64; + constexpr int STAGES = 3; + KERNEL_LAUNCH_CODE + } + else + { + constexpr int CTA_M = 64; + constexpr int CTA_N = 64; + constexpr int CTA_K = 64; + constexpr int WARP_M = 32; + constexpr int WARP_N = 32; + constexpr int WARP_K = 64; + constexpr int STAGES = 6; + KERNEL_LAUNCH_CODE + } + return ; +} \ No newline at end of file diff --git a/llm-awq/awq/kernels/csrc/w8a8/w8a8_gemm_cuda.h b/llm-awq/awq/kernels/csrc/w8a8/w8a8_gemm_cuda.h new file mode 100644 index 0000000000000000000000000000000000000000..6355fb99ec4d6fbd956df23c008e4236060512b8 --- /dev/null +++ b/llm-awq/awq/kernels/csrc/w8a8/w8a8_gemm_cuda.h @@ -0,0 +1,4 @@ +#include + +void w8a8_gemm_forward_cuda(torch::Tensor _in_feats, torch::Tensor _kernel, torch::Tensor _wscales, torch::Tensor _ascales, torch::Tensor _out_feats); +void w8a8_gemm_fuse_bias_forward_cuda(torch::Tensor _in_feats, torch::Tensor _kernel, torch::Tensor _wscales, torch::Tensor _ascales, torch::Tensor _out_feats, torch::Tensor _bias); \ No newline at end of file diff --git a/llm-awq/awq/quantize/__init__.py b/llm-awq/awq/quantize/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..17ddad2a142490a81b6d44ebedd7629678d321c1 --- /dev/null +++ b/llm-awq/awq/quantize/__init__.py @@ -0,0 +1,2 @@ +from .w8a8_linear import * +from .smooth import * diff --git a/llm-awq/awq/quantize/__pycache__/w8a8_linear.cpython-311.pyc b/llm-awq/awq/quantize/__pycache__/w8a8_linear.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..15f83ba6d994799a071f95f953178bee9b0425d1 Binary files /dev/null and b/llm-awq/awq/quantize/__pycache__/w8a8_linear.cpython-311.pyc differ diff --git a/llm-awq/awq/quantize/auto_clip.py b/llm-awq/awq/quantize/auto_clip.py new file mode 100644 index 0000000000000000000000000000000000000000..0714bae50bfa37c8873df2a8ae9ae6a038bd4d04 --- /dev/null +++ b/llm-awq/awq/quantize/auto_clip.py @@ -0,0 +1,98 @@ +import torch +import torch.nn as nn +from .quantizer import pseudo_quantize_tensor +import gc + +__all__ = ["auto_clip_block"] + + +# weight quantization +@torch.no_grad() +def auto_clip_layer( + w, input_feat, n_bit, q_config, n_grid=20, max_shrink=0.5, n_sample_token=512 +): + assert w.dim() == 2 + org_w_shape = w.shape + # w [co, ci] -> [co, 1, n_group, group size] + # input_feat [n_token, ci] -> [1, n_token, n_group, group size] + group_size = ( + q_config["q_group_size"] if q_config["q_group_size"] > 0 else w.shape[1] + ) + input_feat = input_feat.view(-1, input_feat.shape[-1]) + input_feat = input_feat.reshape(1, input_feat.shape[0], -1, group_size) + input_feat = input_feat[:, 0 :: input_feat.shape[1] // n_sample_token] + w = w.reshape(w.shape[0], 1, -1, group_size) + + oc_batch_size = 256 if w.shape[0] % 256 == 0 else 64 # prevent OOM + assert w.shape[0] % oc_batch_size == 0 + w_all = w + best_max_val_all = [] + + for i_b in range(w.shape[0] // oc_batch_size): + w = w_all[i_b * oc_batch_size : (i_b + 1) * oc_batch_size] + + org_max_val = w.abs().amax(dim=-1, keepdim=True) # co, 1, n_group, 1 + + best_max_val = org_max_val.clone() + min_errs = torch.ones_like(org_max_val) * 1e9 + input_feat = input_feat.to(w.device) + org_out = (input_feat * w).sum(dim=-1) # co, n_token, n_group + + for i_s in range(int(max_shrink * n_grid)): + max_val = org_max_val * (1 - i_s / n_grid) + min_val = -max_val + cur_w = torch.clamp(w, min_val, max_val) + q_w = pseudo_quantize_tensor(cur_w, n_bit=n_bit, **q_config) + cur_out = (input_feat * q_w).sum(dim=-1) + + # co, 1, n_group, 1 + err = (cur_out - org_out).pow(2).mean(dim=1).view(min_errs.shape) + del cur_w + del cur_out + cur_best_idx = err < min_errs + min_errs[cur_best_idx] = err[cur_best_idx] + best_max_val[cur_best_idx] = max_val[cur_best_idx] + best_max_val_all.append(best_max_val) + + best_max_val = torch.cat(best_max_val_all, dim=0) + + del input_feat + del org_out + gc.collect() + torch.cuda.empty_cache() + return best_max_val.squeeze(1) + + +@torch.no_grad() +def auto_clip_block(module, w_bit, q_config, input_feat): + named_linears = { + name: m for name, m in module.named_modules() if isinstance(m, nn.Linear) + } + + clip_list = [] + for name in named_linears: + # due to qk bmm, it is hard to clip precisely + if any([_ in name for _ in ["q_", "k_", "query", "key", "Wqkv"]]): + continue + named_linears[name].cuda() + max_val = auto_clip_layer( + named_linears[name].weight, input_feat[name], n_bit=w_bit, q_config=q_config + ) + clip_list.append((name, max_val)) + named_linears[name].cpu() + return clip_list + + +@torch.no_grad() +def apply_clip(module, clip_list): + from ..utils.module import get_op_by_name + + for name, max_val in clip_list: + layer = get_op_by_name(module, name) + layer.cuda() + max_val = max_val.to(layer.weight.device).to(layer.weight.dtype) + org_shape = layer.weight.shape + layer.weight.data = layer.weight.data.reshape(*max_val.shape[:2], -1) + layer.weight.data = torch.clamp(layer.weight.data, -max_val, max_val) + layer.weight.data = layer.weight.data.reshape(org_shape) + layer.cpu() diff --git a/llm-awq/awq/quantize/auto_scale.py b/llm-awq/awq/quantize/auto_scale.py new file mode 100644 index 0000000000000000000000000000000000000000..fe604908ad9aa6a97b597d6dfa43c7b427913b17 --- /dev/null +++ b/llm-awq/awq/quantize/auto_scale.py @@ -0,0 +1,480 @@ +import gc +import torch +import torch.nn as nn + +from transformers.models.bloom.modeling_bloom import BloomBlock, BloomGelu +from transformers.models.opt.modeling_opt import OPTDecoderLayer +from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaRMSNorm +from transformers.activations import GELUActivation +from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm, Qwen2DecoderLayer + +from .qmodule import ScaledActivation +from ..utils.module import get_op_by_name, get_op_name, set_op_by_name + +__all__ = ["auto_scale_block", "apply_scale"] + + +@torch.no_grad() +def get_weight_scale(weight, q_group_size=-1): + org_shape = weight.shape + if q_group_size > 0: + weight = weight.view(-1, q_group_size) + scale = weight.abs() / weight.abs().amax(dim=1, keepdim=True) + scale = scale.view(org_shape) + scale = scale.mean(0) + return scale + + +@torch.no_grad() +def get_act_scale(x): + return x.abs().view(-1, x.shape[-1]).mean(0) + + +@torch.no_grad() +def scale_ln_fcs(ln, fcs, scales): + if not isinstance(fcs, list): + fcs = [fcs] + + scales = scales.to(ln.weight.device).to(ln.weight.dtype) + + ln.weight.div_(scales) + if hasattr(ln, "bias") and ln.bias is not None: + ln.bias.div_(scales) + + for fc in fcs: + fc.weight.mul_(scales.view(1, -1)) + + for p in ln.parameters(): + assert torch.isnan(p).sum() == 0 + for fc in fcs: + for p in fc.parameters(): + assert torch.isnan(p).sum() == 0 + + +@torch.no_grad() +def scale_fc_fc(fc1, fc2, scales): + assert isinstance(fc1, nn.Linear) + assert isinstance(fc2, nn.Linear) + # assert fc1.out_features == fc2.in_features + + scales = scales.to(fc1.weight.device).to(fc1.weight.dtype) + + # fc1.weight.div_(scales.view(-1, 1)) + fc1.weight[-scales.size(0) :].div_(scales.view(-1, 1)) + if fc1.bias is not None: + fc1.bias.div_(scales.view(-1)) + + fc2.weight.mul_(scales.view(1, -1)) + + for p in fc1.parameters(): + assert torch.isnan(p).sum() == 0 + for p in fc2.parameters(): + assert torch.isnan(p).sum() == 0 + + +@torch.no_grad() +def scale_gelu_fc(gelu, fc, scales): + assert isinstance(gelu, (nn.GELU, BloomGelu, GELUActivation)) + assert isinstance(fc, nn.Linear) + + fc.weight.mul_(scales.view(1, -1).to(fc.weight.device).to(fc.weight.dtype)) + + for p in fc.parameters(): + assert torch.isnan(p).sum() == 0 + + +@torch.no_grad() +def auto_scale_block(module, module_kwargs, w_bit, q_config, input_feat): + from .quantizer import pseudo_quantize_tensor + + # firstly, get the weight quantize function + if w_bit is not None: + + def w_quantize_func(p): + return pseudo_quantize_tensor( + p, + n_bit=w_bit, + **q_config, + ).detach() + + else: + + def w_quantize_func(p): + return p + + if "use_cache" in module_kwargs: + module_kwargs.pop("use_cache") + + # find the best scale ratio + def _search_module_scale(block, linears2scale: list, x, kwargs={}): + # w: co, ci + # x: n, ci + x = x.to(next(block.parameters()).device) + with torch.no_grad(): + org_out = block(x, **kwargs) + if isinstance(org_out, tuple): + org_out = org_out[0] + + x_max = get_act_scale(x) + + best_error = float("inf") + best_ratio = -1 + best_scales = None + + n_grid = 20 + history = [] + + org_sd = {k: v.cpu() for k, v in block.state_dict().items()} + for ratio in range(n_grid): + ratio = ratio * 1 / n_grid + scales = x_max.pow(ratio).clamp(min=1e-4).view(-1) + scales = scales / (scales.max() * scales.min()).sqrt() + for fc in linears2scale: + fc.weight.mul_(scales.view(1, -1).to(fc.weight.device)) + fc.weight.data = w_quantize_func(fc.weight.data) / (scales.view(1, -1)) + out = block(x, **kwargs) + if isinstance(out, tuple): + out = out[0] + + loss = ( + (org_out - out).float().pow(2).mean().item() + ) # float prevents overflow + history.append(loss) + is_best = loss < best_error + if is_best: + best_error = loss + best_ratio = ratio + best_scales = scales + block.load_state_dict(org_sd) + if best_ratio == -1: + print(history) + raise Exception + # print(best_ratio) + best_scales = best_scales.view(-1) + + assert torch.isnan(best_scales).sum() == 0, best_scales + return best_scales.detach() + + def _auto_get_scale(prev_op, layers, inp, module2inspect=None, kwargs={}): + # module2inspect: if given, we will check the output diff of this module instead of layers + if module2inspect is None: + assert len(layers) == 1 + module2inspect = layers[0] + + scales = _search_module_scale(module2inspect, layers, inp, kwargs) + scales = scales.detach().cpu() + # prev_op_name, [layer_name], scale + return ( + get_op_name(module, prev_op), + tuple([get_op_name(module, m) for m in layers]), + scales, + ) + + scales_list = [] # return the searched scales + + if isinstance(module, OPTDecoderLayer): + # attention input + scales_list.append( + _auto_get_scale( + prev_op=module.self_attn_layer_norm, + layers=[ + module.self_attn.q_proj, + module.self_attn.k_proj, + module.self_attn.v_proj, + ], + inp=input_feat["self_attn.q_proj"], + module2inspect=module.self_attn, + kwargs=module_kwargs, + ) + ) + # attn out + scales_list.append( + _auto_get_scale( + prev_op=module.self_attn.v_proj, + layers=[module.self_attn.out_proj], + inp=input_feat["self_attn.out_proj"], + ) + ) + # fc1 + scales_list.append( + _auto_get_scale( + prev_op=module.final_layer_norm, + layers=[module.fc1], + inp=input_feat["fc1"], + ) + ) + # fc2 + scales_list.append( + _auto_get_scale( + prev_op=module.fc1, + layers=[module.fc2], + inp=input_feat["fc2"], + ) + ) + + elif isinstance(module, (LlamaDecoderLayer, Qwen2DecoderLayer)): + # attention input + scales_list.append( + _auto_get_scale( + prev_op=module.input_layernorm, + layers=[ + module.self_attn.q_proj, + module.self_attn.k_proj, + module.self_attn.v_proj, + ], + inp=input_feat["self_attn.q_proj"], + module2inspect=module.self_attn, + kwargs=module_kwargs, + ) + ) + # attn out + # Please refer to https://github.com/mit-han-lab/llm-awq/pull/67#issue-1850622696 + if module.self_attn.v_proj.weight.shape == module.self_attn.o_proj.weight.shape: + scales_list.append( + _auto_get_scale( + prev_op=module.self_attn.v_proj, + layers=[module.self_attn.o_proj], + inp=input_feat["self_attn.o_proj"], + ) + ) + # fc1 + scales_list.append( + _auto_get_scale( + prev_op=module.post_attention_layernorm, + layers=[module.mlp.gate_proj, module.mlp.up_proj], + inp=input_feat["mlp.gate_proj"], + module2inspect=module.mlp, + ) + ) + # fc2 + scales_list.append( + _auto_get_scale( + prev_op=module.mlp.up_proj, + layers=[module.mlp.down_proj], + inp=input_feat["mlp.down_proj"], + ) + ) + + elif isinstance(module, BloomBlock): + # attention input + scales_list.append( + _auto_get_scale( + prev_op=module.input_layernorm, + layers=[module.self_attention.query_key_value], + inp=input_feat["self_attention.query_key_value"], + module2inspect=module, + kwargs=module_kwargs, + ) + ) + # attn out + # Please refer to https://github.com/mit-han-lab/llm-awq/issues/2#issuecomment-1606297469 + """ + scales_list.append(_auto_get_scale( + prev_op=module.self_attention.query_key_value, + layers=[module.self_attention.dense], + inp=input_feat['self_attention.dense'], + )) + """ + # fc1 + scales_list.append( + _auto_get_scale( + prev_op=module.post_attention_layernorm, + layers=[module.mlp.dense_h_to_4h], + inp=input_feat["mlp.dense_h_to_4h"], + module2inspect=module, + kwargs=module_kwargs, + ) + ) + # fc2 + scales_list.append( + _auto_get_scale( + prev_op=module.mlp.gelu_impl, + layers=[module.mlp.dense_4h_to_h], + inp=input_feat["mlp.dense_4h_to_h"], + ) + ) + elif "mpt" in str(module.__class__).lower(): + # attention input + scales_list.append( + _auto_get_scale( + prev_op=module.norm_1, + layers=[module.attn.Wqkv], + inp=input_feat["attn.Wqkv"], + module2inspect=module.attn, + kwargs=module_kwargs, + ) + ) + + # attn out + scales_list.append( + _auto_get_scale( + prev_op=module.attn.Wqkv, + layers=[module.attn.out_proj], + inp=input_feat["attn.out_proj"], + ) + ) + # fc1 + scales_list.append( + _auto_get_scale( + prev_op=module.norm_2, + layers=[module.ffn.up_proj], + inp=input_feat["ffn.up_proj"], + module2inspect=module.ffn, + ) + ) + # fc2 + scales_list.append( + _auto_get_scale( + prev_op=module.ffn.act, + layers=[module.ffn.down_proj], + inp=input_feat["ffn.down_proj"], + ) + ) + + elif "falcon" in str(module.__class__).lower(): + # attn out + # Haotian: TBD: need to handle repeated scales for MQ + """ + scales_list.append(_auto_get_scale( + prev_op=module.self_attention.query_key_value, + layers=[module.self_attention.dense], + inp=input_feat['self_attention.dense'], + )) + """ + # fc1, as long as it is scaled, everything is screwed up + if "falcon-7b" in str(module.__class__).lower(): + scales_list.append( + _auto_get_scale( + prev_op=module.input_layernorm, + layers=[ + module.mlp.dense_h_to_4h, + module.self_attention.query_key_value, + ], + inp=input_feat["self_attention.query_key_value"], + module2inspect=module, + kwargs=module_kwargs, + ) + ) + elif "falcon-40b" in str(module.__class__).lower(): + scales_list.append( + _auto_get_scale( + prev_op=module.ln_attn, + layers=[module.self_attention.query_key_value], + inp=input_feat["self_attention.query_key_value"], + module2inspect=module, + kwargs=module_kwargs, + ) + ) + scales_list.append( + _auto_get_scale( + prev_op=module.ln_mlp, + layers=[module.mlp.dense_h_to_4h], + inp=input_feat["mlp.dense_h_to_4h"], + module2inspect=module, + kwargs=module_kwargs, + ) + ) + else: + raise NotImplementedError( + "Unknown Falcon architecture, currently only falcon-7b and falcon-40b are supported" + ) + # fc2 + scales_list.append( + _auto_get_scale( + prev_op=module.mlp.act, + layers=[module.mlp.dense_4h_to_h], + inp=input_feat["mlp.dense_4h_to_h"], + ) + ) + elif "bigcode" in str(module.__class__).lower(): + scales_list.append( + _auto_get_scale( + prev_op=module.ln_1, + layers=[module.attn.c_attn], + inp=input_feat["attn.c_attn"], + module2inspect=module.attn, + kwargs=module_kwargs, + ) + ) + # fc1 + scales_list.append( + _auto_get_scale( + prev_op=module.ln_2, + layers=[module.mlp.c_fc], + inp=input_feat["mlp.c_fc"], + module2inspect=module.mlp, + ) + ) + # fc2 + scales_list.append( + _auto_get_scale( + prev_op=module.mlp.act, + layers=[module.mlp.c_proj], + inp=input_feat["mlp.c_proj"], + ) + ) + elif "neox" in str(module.__class__).lower(): + scales_list.append( + _auto_get_scale( + prev_op=module.input_layernorm, + layers=[module.attention.query_key_value], + inp=input_feat["attention.query_key_value"], + module2inspect=module.attention, + kwargs=module_kwargs, + ) + ) + # fc1 + scales_list.append( + _auto_get_scale( + prev_op=module.post_attention_layernorm, + layers=[module.mlp.dense_h_to_4h], + inp=input_feat["mlp.dense_h_to_4h"], + module2inspect=module.mlp, + ) + ) + # fc2 + scales_list.append( + _auto_get_scale( + prev_op=module.mlp.act, + layers=[module.mlp.dense_4h_to_h], + inp=input_feat["mlp.dense_4h_to_h"], + ) + ) + else: + raise NotImplementedError(f"{type(module)} not supported yet!") + + return scales_list + + +def apply_scale(module, scales_list, input_feat_dict=None): + for prev_op_name, layer_names, scales in scales_list: + prev_op = get_op_by_name(module, prev_op_name) + layers = [get_op_by_name(module, name) for name in layer_names] + + prev_op.cuda() + for layer in layers: + layer.cuda() + scales.cuda() + + if isinstance(prev_op, nn.Linear): + assert len(layers) == 1 + scale_fc_fc(prev_op, layers[0], scales) + elif isinstance(prev_op, (nn.LayerNorm, LlamaRMSNorm, Qwen2RMSNorm)): + scale_ln_fcs(prev_op, layers, scales) + elif isinstance(prev_op, (nn.GELU, BloomGelu, GELUActivation, nn.SiLU)): + new_module = ScaledActivation(prev_op, scales) + set_op_by_name(module, prev_op_name, new_module) + scale_gelu_fc(prev_op, layers[0], scales) + else: + raise NotImplementedError(f"prev_op {type(prev_op)} not supported yet!") + + # apply the scaling to input feat if given; prepare it for clipping + if input_feat_dict is not None: + for layer_name in layer_names: + inp = input_feat_dict[layer_name] + inp.div_(scales.view(1, -1).to(inp.device).to(inp.dtype)) + + prev_op.cpu() + for layer in layers: + layer.cpu() + scales.cpu() diff --git a/llm-awq/awq/quantize/pre_quant.py b/llm-awq/awq/quantize/pre_quant.py new file mode 100644 index 0000000000000000000000000000000000000000..23a5ecbe11befb977a79535701802703be31042e --- /dev/null +++ b/llm-awq/awq/quantize/pre_quant.py @@ -0,0 +1,254 @@ +import torch +import torch.nn as nn +import tqdm +import gc +import functools +from collections import defaultdict +from typing import List + +from transformers.models.bloom.modeling_bloom import BloomForCausalLM +from transformers.models.opt.modeling_opt import OPTForCausalLM +from transformers.models.llama.modeling_llama import LlamaForCausalLM +try: + from tinychat.models import LlavaLlamaForCausalLM +except ImportError as e: + pass + +from transformers.models.qwen2.modeling_qwen2 import Qwen2ForCausalLM + +from .auto_scale import auto_scale_block, apply_scale +from .auto_clip import auto_clip_block, apply_clip + +__all__ = ["run_awq"] + + +def get_named_linears(module): + return {name: m for name, m in module.named_modules() if isinstance(m, nn.Linear)} + + +def get_blocks(model): + if model.__class__.__name__ in ("LlamaForCausalLM", "Qwen2ForCausalLM"): + layers = model.model.layers + elif model.__class__.__name__ == "InternVL3": + layers = model.language_model.model.layers + # layers = [model.language_model.model.layers, model.vision_model.encoder.layers] + elif model.__class__.__name__ == "LlavaLlamaForCausalLM": + # layers = [model.model.layers, model.model.vision_tower.vision_tower.vision_model.encoder.layers] + layers = model.model.layers + elif isinstance(model, OPTForCausalLM): + layers = model.model.decoder.layers + elif isinstance(model, BloomForCausalLM): + layers = model.transformer.h + elif "mpt" in str(model.__class__).lower(): + layers = model.transformer.blocks + elif "falcon" in str(model.__class__).lower(): + layers = model.transformer.h + elif "bigcode" in str(model.__class__).lower(): + layers = model.transformer.h + elif "neox" in str(model.__class__).lower(): + layers = model.gpt_neox.layers + elif model.__class__.__name__ == "LlavaLlamaModel": + layers = model.llm.model.layers + else: + raise NotImplementedError(type(model)) + return layers + + +def move_embed(model, device): + if isinstance(model, (LlamaForCausalLM, Qwen2ForCausalLM)): + model.model.embed_tokens = model.model.embed_tokens.to(device) + model.model.rotary_emb = model.model.rotary_emb.to(device) + elif model.__class__.__name__ == "InternVL3": + model.language_model.model.embed_tokens = ( + model.language_model.model.embed_tokens.to(device) + ) + model.language_model.model.rotary_emb = ( + model.language_model.model.rotary_emb.to(device) + ) + model.vision_model.embeddings.to(device) + elif isinstance(model, LlavaLlamaForCausalLM): + model.model.embed_tokens = model.model.embed_tokens.to(device) + model.model.vision_tower.vision_tower.vision_model.embeddings.to(device) + elif isinstance(model, OPTForCausalLM): + model.model.decoder.embed_tokens = model.model.decoder.embed_tokens.to(device) + model.model.decoder.embed_positions = model.model.decoder.embed_positions.to( + device + ) + elif isinstance(model, BloomForCausalLM): + model.transformer.word_embeddings = model.transformer.word_embeddings.to(device) + model.transformer.word_embeddings_layernorm = ( + model.transformer.word_embeddings_layernorm.to(device) + ) + elif "mpt" in str(model.__class__).lower(): + model.transformer.wte = model.transformer.wte.to(device) + model.transformer.emb_drop = model.transformer.emb_drop.to(device) + elif "falcon" in str(model.__class__).lower(): + model.transformer.word_embeddings = model.transformer.word_embeddings.to(device) + elif "bigcode" in str(model.__class__).lower(): + model.transformer.wte = model.transformer.wte.to(device) + model.transformer.wpe = model.transformer.wpe.to(device) + model.transformer.drop = model.transformer.drop.to(device) + elif "neox" in str(model.__class__).lower(): + model.gpt_neox.embed_in = model.gpt_neox.embed_in.to(device) + model.gpt_neox.emb_dropout = model.gpt_neox.emb_dropout.to(device) + model.embed_out = model.embed_out.to(device) + elif "llavallamamodel" in str(model.__class__).lower(): + model.llm.model.embed_tokens = model.llm.model.embed_tokens.to(device) + else: + raise NotImplementedError(type(model)) + + +@torch.no_grad() +def run_awq( + model, + enc, + w_bit, + q_config, + n_samples=512, + seqlen=512, + auto_scale=True, + mse_range=True, + # some configs for ablation study + calib_data="pileval", +): + from ..utils.calib_data import get_calib_dataset + from ..utils.module import append_str_prefix, get_op_name + + if "bigcode" in str(model.__class__).lower(): + # otherwise attention_mask will always be on cpu. + model.transformer.bias = model.transformer.bias.to("cuda") + + layers = get_blocks(model) + + samples = get_calib_dataset( + data=calib_data, tokenizer=enc, n_samples=n_samples, block_size=seqlen + ) + samples = torch.cat(samples, dim=0) + + inps = [] + layer_kwargs = {} + + layers[0] = layers[0].cuda() + move_embed(model, "cuda") + + # get input and kwargs to layer 0 + # with_kwargs is only supported in PyTorch 2.0 + # use this Catcher hack for now + class Catcher(nn.Module): + def __init__(self, module): + super().__init__() + self.module = module + + def forward(self, inp, **kwargs): + inps.append(inp) + layer_kwargs.update(kwargs) + raise ValueError # early exit to break later inference + + # patch layer 0 to catch input and kwargs + layers[0] = Catcher(layers[0]) + try: + if model.__class__.__name__ == "LlavaLlamaModel": + model.llm(samples.to(next(model.parameters()).device)) + elif model.__class__.__name__ == "InternVL3": + model.language_model(samples.to(next(model.parameters()).device)) + else: + model(samples.to(next(model.parameters()).device)) + except ValueError: # work with early exit + pass + del samples + layers[0] = layers[0].module # restore + inps = inps[0] + + layers[0] = layers[0].cpu() + move_embed(model, "cpu") + + gc.collect() + torch.cuda.empty_cache() + + awq_results = { + "scale": [], + "clip": [], + } + + # solve layer by layer + for i in tqdm.tqdm(range(len(layers)), desc="Running AWQ..."): + layer = layers[i] + layer = layer.cuda() + named_linears = get_named_linears(layer) + + # firstly, get input features of all linear layers + def cache_input_hook(m, x, y, name, feat_dict): + x = x[0] + x = x.detach().cpu() + feat_dict[name].append(x) + + input_feat = defaultdict(list) + handles = [] + for name in named_linears: + handles.append( + named_linears[name].register_forward_hook( + functools.partial(cache_input_hook, name=name, feat_dict=input_feat) + ) + ) + inps = inps.to(next(layer.parameters()).device) # in case multi-gpu + # get output as next layer's input + inps = layer(inps, **layer_kwargs)[0] + for h in handles: + h.remove() + # now solve for scaling and clipping + input_feat = {k: torch.cat(v, dim=0) for k, v in input_feat.items()} + + # Clear GPU memory + torch.cuda.empty_cache() + + if ( + auto_scale + ): # if it applies, we should also modify the input_feat with scales + scales_list = auto_scale_block( + layer, + layer_kwargs, + w_bit=w_bit, + q_config=q_config, + input_feat=input_feat, + ) + # apply_scale(layer, scales_list, input_feat_dict=input_feat) + apply_scale(layers[i], scales_list, input_feat_dict=input_feat) + # append prefix to make names global + awq_results["scale"] += append_str_prefix( + scales_list, get_op_name(model, layer) + "." + ) + + # Clear GPU memory + torch.cuda.empty_cache() + # for line in torch.cuda.memory_summary().splitlines(): + # if "Allocated" in line: + # print(line) + + if mse_range: + clip_list = auto_clip_block( + layer, + w_bit=w_bit, + q_config=q_config, + input_feat=input_feat, + ) + apply_clip(layer, clip_list) + # append prefix to make names global + awq_results["clip"] += append_str_prefix( + clip_list, get_op_name(model, layer) + "." + ) + + layer = layer.cpu() + # Haotian: check activation replacement + del input_feat + gc.collect() + torch.cuda.empty_cache() + # for line in torch.cuda.memory_summary().splitlines(): + # if "Allocated" in line: + # print(line) + + return awq_results + + +def apply_awq(model, awq_results): + apply_scale(model, awq_results["scale"]) + apply_clip(model, awq_results["clip"]) diff --git a/llm-awq/awq/quantize/qmodule.py b/llm-awq/awq/quantize/qmodule.py new file mode 100644 index 0000000000000000000000000000000000000000..d7b477c3b07bf66cfbe63752bb09f18374ac830e --- /dev/null +++ b/llm-awq/awq/quantize/qmodule.py @@ -0,0 +1,235 @@ +import math +import torch +import torch.nn as nn +import awq_inference_engine # with CUDA kernels + + +def make_divisible(c, divisor): + return (c + divisor - 1) // divisor + + +def calculate_zeros_width(in_features, group_size=128, pack_num=8): + if group_size >= 128: + size_multiplier = 1 + elif group_size == 64: + size_multiplier = 2 + elif group_size == 32: + size_multiplier = 4 + else: + raise NotImplementedError + + base_width = make_divisible(in_features // group_size, pack_num) + base_width = make_divisible(base_width, size_multiplier) * size_multiplier + return base_width + + +def pack_intweight(unpacked_qweight, interleave, kstride): + # unpacked_qweight: [N, K] + N = unpacked_qweight.shape[0] + K = unpacked_qweight.shape[1] + + Packed_Kernel = unpacked_qweight.cpu().numpy().reshape(N, K // 32, 32) + # np.arange(32).reshape(4, 4, 2).transpose(1, 0, 2) => [0, 1, 8, 9, 16, 17, 24, 25, ...] + Packed_Kernel = Packed_Kernel.reshape(N, K // 32, 4, 4, 2).transpose(0, 1, 3, 2, 4) + Packed_Kernel = Packed_Kernel.reshape(N, K // 32, 32) + + # reorder each 8 weights for fast dequantization + # [0, 1, 2, 3, 4, 5, 6, 7] => [0, 2, 4, 6, 1, 3, 5, 7] + Packed_Kernel = Packed_Kernel.reshape(N, K // 32, 4, 8) + Packed_Kernel = Packed_Kernel.reshape(N, K // 32, 4, 4, 2).transpose(0, 1, 2, 4, 3) + Packed_Kernel = Packed_Kernel.reshape(N, K) + + # interleaving every four rows + Packed_Kernel = Packed_Kernel.reshape( + N // interleave, interleave, K // kstride, kstride + ) + # N // 4, K // 64, 4, 64 + Packed_Kernel = Packed_Kernel.transpose(0, 2, 1, 3) + Packed_Kernel = Packed_Kernel.reshape( + N // interleave, K // kstride, kstride, interleave + ) + # Packing -> (N // 4, K // 64, 64) + Packed_Kernel = ( + Packed_Kernel[..., 0] + | (Packed_Kernel[..., 1] << 4) + | (Packed_Kernel[..., 2] << 8) + | (Packed_Kernel[..., 3] << 12) + ) + # reshape to (N // 4, K), FP16 format + Packed_Kernel = Packed_Kernel.reshape(N // interleave, K) + qweight = ( + torch.tensor(Packed_Kernel.astype("int16")) + .to(unpacked_qweight.device) + .contiguous() + ) + return qweight + + +class ScaledActivation(nn.Module): + def __init__(self, module, scales): + super().__init__() + self.act = module + self.scales = nn.Parameter(scales.data) + + def forward(self, x): + return self.act(x) / self.scales.view(1, 1, -1).to(x.device) + + +class WQLinear(nn.Module): + def __init__(self, w_bit, group_size, in_features, out_features, bias, dev, dtype=torch.float16): + super().__init__() + + if w_bit not in [4]: + raise NotImplementedError("Only 4-bit are supported for now.") + + self.in_features = in_features + self.out_features = out_features + self.w_bit = w_bit + self.group_size = group_size if group_size != -1 else in_features + self.split_k_iters = 8 + self.interleave = 4 + # quick sanity check (make sure aligment) + assert self.in_features % self.group_size == 0 + assert out_features % (32 // self.w_bit) == 0 + pack_num = 32 // self.w_bit + int16_pack_num = 16 // self.w_bit + + assert out_features % (self.interleave) == 0 + self.register_buffer( + "qweight", + torch.zeros( + ( + out_features // self.interleave, + in_features // int16_pack_num * self.interleave, + ), + dtype=torch.int16, + device=dev, + ), + ) + self.register_buffer( + "scales", + torch.zeros( + ( + calculate_zeros_width(in_features, self.group_size) * pack_num, + out_features, + ), + dtype=dtype, + device=dev, + ), + ) + self.register_buffer( + "scaled_zeros", + torch.zeros( + ( + calculate_zeros_width(in_features, self.group_size) * pack_num, + out_features, + ), + dtype=dtype, + device=dev, + ), + ) + + if bias: + self.register_buffer( + "bias", torch.zeros((out_features), dtype=dtype, device=dev) + ) + else: + self.bias = None + + @classmethod + def from_linear( + cls, linear, w_bit, group_size, init_only=False, scales=None, zeros=None + ): + awq_linear = cls( + w_bit, + group_size, + linear.in_features, + linear.out_features, + linear.bias is not None, + linear.weight.device, + dtype=linear.weight.data.dtype + ) + if init_only: # just prepare for loading sd + return awq_linear + + # need scales and zeros info for real quantization + assert scales is not None and zeros is not None + scale_zeros = zeros * scales + + dtype = scales.dtype + + pack_num = 32 // awq_linear.w_bit + qscales = torch.zeros( + ( + scales.shape[0], + calculate_zeros_width(linear.in_features, group_size) * pack_num, + ), + dtype=dtype, + device=scales.device, + ) + qscales[:, : scales.shape[1]] = scales + # awq_linear.scales = scales.clone().half() + awq_linear.scales = qscales.transpose(1, 0).contiguous() + if linear.bias is not None: + awq_linear.bias = linear.bias.clone().to(dtype) + + intweight = [] + for idx in range(awq_linear.in_features): + intweight.append( + torch.round( + (linear.weight.data[:, idx] + scale_zeros[:, idx // group_size]) + / qscales[:, idx // group_size] + ).to(torch.int)[:, None] + ) + intweight = torch.cat(intweight, dim=1) + # intweight = intweight.t().contiguous() + intweight = intweight.to(dtype=torch.int32) + awq_linear.qweight = pack_intweight( + intweight.contiguous(), interleave=4, kstride=64 + ) + + zeros = zeros.to(dtype=torch.int32) + scaled_zeros = torch.zeros_like(qscales) + # scaled_zeros[:, :scales.shape[1]] = -(qscales[:, :scales.shape[1]] * (zeros.to(torch.float32) - 8.0)).to(torch.float16) + scaled_zeros[:, : scales.shape[1]] = -( + qscales[:, : scales.shape[1]] * (zeros.to(torch.float32)) + ).to(dtype) + awq_linear.scaled_zeros = scaled_zeros.transpose(1, 0).contiguous() + + return awq_linear + + @torch.no_grad() + def forward(self, x): + # out_shape = x.shape[:-1] + (self.out_features,) + # inputs = x.reshape(-1, x.shape[-1]) + inputs = x + if inputs.numel() / inputs.shape[-1] < 8: + out = awq_inference_engine.gemv_forward_cuda_new( + inputs, + self.qweight, + self.scales, + self.scaled_zeros, + inputs.numel() // inputs.shape[-1], + self.out_features, + self.in_features, + self.group_size, + ) + else: + out = awq_inference_engine.gemm_forward_cuda_new( + inputs, self.qweight, self.scales, self.scaled_zeros + ) # - 8.0 * self.scales) + out = out + self.bias if self.bias is not None else out + # print(out) + # assert 0 + return out + + def extra_repr(self) -> str: + return ( + "in_features={}, out_features={}, bias={}, w_bit={}, group_size={}".format( + self.in_features, + self.out_features, + self.bias is not None, + self.w_bit, + self.group_size, + ) + ) diff --git a/llm-awq/awq/quantize/quantizer.py b/llm-awq/awq/quantize/quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..7d15191651545f504ace5e0f284525421e423bf1 --- /dev/null +++ b/llm-awq/awq/quantize/quantizer.py @@ -0,0 +1,165 @@ +import torch +import torch.nn as nn +from tqdm import tqdm +import gc +from .qmodule import ScaledActivation +from ..utils.module import set_op_by_name + +from transformers.models.bloom.modeling_bloom import BloomBlock + +EMBEDDING_KEYWORDS = ["embed"] +LM_HEAD_KEYWORDS = ["lm_head", "embed_out", "output"] + + +def scale_activations(module): + param = next(module.parameters()) + dtype = param.dtype + device = param.device + if isinstance(module, BloomBlock): + if isinstance(module.mlp.gelu_impl, ScaledActivation): + return + c = module.mlp.dense_h_to_4h.out_features + act = ScaledActivation( + module.mlp.gelu_impl, torch.ones(c, dtype=dtype, device=device) + ) + set_op_by_name(module, "mlp.gelu_impl", act) + elif "mptblock" in str(module.__class__.__name__).lower(): + if isinstance(module.ffn.act, ScaledActivation): + return + c = module.ffn.up_proj.out_features + act = ScaledActivation( + module.ffn.act, torch.ones(c, dtype=dtype, device=device) + ) + set_op_by_name(module, "ffn.act", act) + elif "falcon" in str(module.__class__).lower(): + if isinstance(module.mlp.act, ScaledActivation): + return + c = module.mlp.dense_h_to_4h.out_features + act = ScaledActivation( + module.mlp.act, torch.ones(c, dtype=dtype, device=device) + ) + set_op_by_name(module, "mlp.act", act) + elif "bigcode" in str(module.__class__).lower(): + if isinstance(module.mlp.act, ScaledActivation): + return + c = module.mlp.c_proj.out_features + act = ScaledActivation( + module.mlp.act, torch.ones(c, dtype=dtype, device=device) + ) + set_op_by_name(module, "mlp.act", act) + elif "neox" in str(module.__class__).lower(): + if isinstance(module.mlp.act, ScaledActivation): + return + c = module.mlp.dense_h_to_4h.out_features + act = ScaledActivation( + module.mlp.act, torch.ones(c, dtype=dtype, device=device) + ) + set_op_by_name(module, "mlp.act", act) + + +# core quantization method (simulated quantization) +def pseudo_quantize_tensor( + w, n_bit=8, zero_point=True, q_group_size=-1, inplace=False, get_scale_zp=False +): + org_w_shape = w.shape + if q_group_size > 0: + assert org_w_shape[-1] % q_group_size == 0 + w = w.reshape(-1, q_group_size) + assert w.dim() == 2 + if zero_point: + max_val = w.amax(dim=1, keepdim=True) + min_val = w.amin(dim=1, keepdim=True) + max_int = 2**n_bit - 1 + min_int = 0 + scales = (max_val - min_val).clamp(min=1e-5) / max_int + zeros = (-torch.round(min_val / scales)).clamp_(min_int, max_int) + else: # we actually never used this + assert min_val is None + max_val = w.abs().amax(dim=1, keepdim=True) + max_val = max_val.clamp(min=1e-5) + max_int = 2 ** (n_bit - 1) - 1 + min_int = -(2 ** (n_bit - 1)) + scales = max_val / max_int + zeros = 0 + + assert torch.isnan(scales).sum() == 0 + assert torch.isnan(w).sum() == 0 + + if inplace: + ( + (w.div_(scales).round_().add_(zeros)).clamp_(min_int, max_int).sub_(zeros) + ).mul_(scales) + else: + w = ( + torch.clamp(torch.round(w / scales) + zeros, min_int, max_int) - zeros + ) * scales + assert torch.isnan(w).sum() == 0 + + w = w.reshape(org_w_shape) + + if get_scale_zp: + return w, scales.view(w.shape[0], -1), zeros.view(w.shape[0], -1) + else: + return w + + +@torch.no_grad() +def pseudo_quantize_model_weight( + model, + w_bit, + q_config, +): + from .pre_quant import get_blocks, get_named_linears + + layers = get_blocks(model) + for i in tqdm(range(len(layers)), desc="pseudo weight quantization..."): + named_linears = get_named_linears(layers[i]) + for n, m in named_linears.items(): + m.cuda() + m.weight.data = pseudo_quantize_tensor( + m.weight.data, n_bit=w_bit, **q_config + ) + m.cpu() + + +@torch.no_grad() +def real_quantize_model_weight(model, w_bit, q_config, init_only=False): + from .qmodule import WQLinear + from .pre_quant import get_blocks, get_named_linears + + assert q_config["zero_point"], "We only support zero_point quantization now." + + layers = get_blocks(model) + for i in tqdm( + range(len(layers)), + desc="real weight quantization..." + ("(init only)" if init_only else ""), + ): + layer = layers[i] + named_linears = get_named_linears(layer) + scale_activations(layer) + + for name, module in named_linears.items(): + if init_only: + q_linear = WQLinear.from_linear( + module, w_bit, q_config["q_group_size"], True + ) + q_linear.to(next(layer.parameters()).device) + set_op_by_name(layer, name, q_linear) + else: + module.cuda() + module.weight.data, scales, zeros = pseudo_quantize_tensor( + module.weight.data, n_bit=w_bit, get_scale_zp=True, **q_config + ) + # scales = scales.t().contiguous() + # zeros = zeros.t().contiguous() + q_linear = WQLinear.from_linear( + module, w_bit, q_config["q_group_size"], False, scales, zeros + ) + module.cpu() + q_linear.to(next(layer.parameters()).device) + set_op_by_name(layer, name, q_linear) + torch.cuda.empty_cache() + gc.collect() + + torch.cuda.empty_cache() + gc.collect() diff --git a/llm-awq/awq/quantize/smooth.py b/llm-awq/awq/quantize/smooth.py new file mode 100644 index 0000000000000000000000000000000000000000..f4393f81edf2c0dbccb055c570c713ce3ebb3502 --- /dev/null +++ b/llm-awq/awq/quantize/smooth.py @@ -0,0 +1,246 @@ +# Adapted from SmoothQuant (https://github.com/mit-han-lab/smoothquant) and modified by Yuming Lou + + +import torch.nn as nn +try: + import llava + from llava.media import Image, Video + from llava.utils.media import extract_media + from llava.constants import DEFAULT_IMAGE_TOKEN + from llava.mm_utils import process_image, process_images +except ImportError: + print("VILA is not installed. Multimodal features will not be available. To activate, please install VILA at https://github.com/NVlabs/VILA.") + +import torch +from collections import defaultdict +from functools import partial +from tqdm import tqdm +import numpy as np +import functools + + +@torch.no_grad() +def get_act_scales(model, data): + num_samples = data.shape[0] + model.eval() + act_scales = {} + + def stat_tensor(name, tensor): + hidden_dim = tensor.shape[-1] + tensor = tensor.view(-1, hidden_dim).abs().detach() + comming_max = torch.max(tensor, dim=0)[0].float().cpu() + if name in act_scales: + act_scales[name] = torch.max(act_scales[name], comming_max) + else: + act_scales[name] = comming_max + + def stat_input_hook(m, x, y, name): + if isinstance(x, tuple): + x = x[0] + stat_tensor(name, x) + + hooks = [] + for name, m in model.named_modules(): + if isinstance(m, nn.Linear): + hooks.append( + m.register_forward_hook(functools.partial(stat_input_hook, name=name)) + ) + + for i in tqdm(range(num_samples)): + input = data[i : i + 1] + model(input) + + for h in hooks: + h.remove() + + return act_scales + + +@torch.no_grad() +def get_static_decoder_layer_scales( + model, + data, +): + num_samples = data.shape[1] + model.eval() + device = next(model.parameters()).device + + act_dict = defaultdict(dict) + + def stat_io_hook(m, x, y, name): + if isinstance(x, tuple): + x = x[0] + if name not in act_dict or "input" not in act_dict[name]: + act_dict[name]["input"] = x.detach().abs().max().item() + else: + act_dict[name]["input"] = max( + act_dict[name]["input"], x.detach().abs().max().item() + ) + if isinstance(y, tuple): + y = y[0] + if name not in act_dict or "output" not in act_dict[name]: + act_dict[name]["output"] = y.detach().abs().max().item() + else: + act_dict[name]["output"] = max( + act_dict[name]["output"], y.detach().abs().max().item() + ) + + hooks = [] + for name, m in model.named_modules(): + if isinstance(m, torch.nn.Linear): + hooks.append(m.register_forward_hook(partial(stat_io_hook, name=name))) + pbar = tqdm(range(num_samples)) + for i in pbar: + model(data[i : i + 1]) + mean_scale = np.mean([v["input"] for v in act_dict.values()]) + pbar.set_description(f"Mean input scale: {mean_scale:.2f}") + for hook in hooks: + hook.remove() + decoder_layer_scales = [] + for idx in range(model.config.num_hidden_layers): + scale_dict = {} + scale_dict["attn_input_scale"] = ( + act_dict[ + f"vision_tower.vision_model.encoder.layers.{idx}.self_attn.q_proj" + ]["input"] + / 127 + ) + scale_dict["q_output_scale"] = ( + act_dict[ + f"vision_tower.vision_model.encoder.layers.{idx}.self_attn.q_proj" + ]["output"] + / 127 + ) + scale_dict["k_output_scale"] = ( + act_dict[ + f"vision_tower.vision_model.encoder.layers.{idx}.self_attn.k_proj" + ]["output"] + / 127 + ) + scale_dict["v_output_scale"] = ( + act_dict[ + f"vision_tower.vision_model.encoder.layers.{idx}.self_attn.v_proj" + ]["output"] + / 127 + ) + scale_dict["out_input_scale"] = ( + act_dict[ + f"vision_tower.vision_model.encoder.layers.{idx}.self_attn.out_proj" + ]["input"] + / 127 + ) + scale_dict["fc1_input_scale"] = ( + act_dict[f"vision_tower.vision_model.encoder.layers.{idx}.mlp.fc1"]["input"] + / 127 + ) + scale_dict["fc2_input_scale"] = ( + act_dict[f"vision_tower.vision_model.encoder.layers.{idx}.mlp.fc2"]["input"] + / 127 + ) + decoder_layer_scales.append(scale_dict) + + return decoder_layer_scales, act_dict + + +def get_smooth_scale(model_path, media): + # Load model + model = llava.load(model_path, devices=[0]) + del model.llm + del model.mm_projector + torch.cuda.empty_cache() + model = model.cuda().eval() + prompt = [] + if media is not None: + for m in media or []: + if any(m.endswith(ext) for ext in [".jpg", ".jpeg", ".png"]): + m = Image(m) + elif any(m.endswith(ext) for ext in [".mp4", ".mkv", ".webm"]): + m = Video(m) + else: + raise ValueError(f"Unsupported media type: {m}") + prompt.append(m) + conversation = [{"from": "human", "value": prompt}] + media = extract_media(conversation, model.config) + for name in media: + if name == "image": + if ( + len(media["image"]) == 1 + and model.config.image_aspect_ratio == "dynamic" + ): + model.config.image_processor = model.vision_tower.image_processor + images = process_image( + media["image"][0], model.config, None, enable_dynamic_res=True + ).half() + conversation[0]["value"] = conversation[0]["value"].replace( + DEFAULT_IMAGE_TOKEN, f"{DEFAULT_IMAGE_TOKEN}\n" * images.shape[0] + ) + else: + images = process_images( + media["image"], model.vision_tower.image_processor, model.config + ).half() + media[name] = [image for image in images] + elif name == "video": + media[name] = [ + process_images( + images, model.vision_tower.image_processor, model.config + ).half() + for images in media[name] + ] + else: + raise ValueError(f"Unsupported media type: {name}") + images = torch.cat(media["video"], dim=1) + model.vision_tower = model.vision_tower.eval() + decoder_layer_scales = get_act_scales(model.vision_tower, images) + return decoder_layer_scales + + +@torch.no_grad() +def smooth_ln_fcs(ln, fcs, act_scales, alpha=0.5): + if not isinstance(fcs, list): + fcs = [fcs] + assert isinstance(ln, nn.LayerNorm) + for fc in fcs: + assert isinstance(fc, nn.Linear) + assert ln.weight.numel() == fc.in_features == act_scales.numel() + + device, dtype = fcs[0].weight.device, fcs[0].weight.dtype + act_scales = act_scales.to(device=device, dtype=dtype) + weight_scales = torch.cat( + [fc.weight.abs().max(dim=0, keepdim=True)[0] for fc in fcs], dim=0 + ) + weight_scales = weight_scales.max(dim=0)[0].clamp(min=1e-5) + + scales = ( + (act_scales.pow(alpha) / weight_scales.pow(1 - alpha)) + .clamp(min=1e-5) + .to(device) + .to(dtype) + ) + + ln.weight.div_(scales) + ln.bias.div_(scales) + + for fc in fcs: + fc.weight.mul_(scales.view(1, -1)) + + +@torch.no_grad() +def smooth_lm(model, scales, alpha=0.5): + if "siglip" in str(model.__class__).lower(): + num = 0 + for name, module in model.named_modules(): + if "siglipencoderlayer" in str(module.__class__).lower(): + attn_ln = module.layer_norm1 + qkv = [ + module.self_attn.q_proj, + module.self_attn.k_proj, + module.self_attn.v_proj, + ] + qkv_input_scales = scales[name + ".self_attn.q_proj"] + smooth_ln_fcs(attn_ln, qkv, qkv_input_scales, alpha) + + ffn_ln = module.layer_norm2 + fc1 = module.mlp.fc1 + fc1_input_scales = scales[name + ".mlp.fc1"] + smooth_ln_fcs(ffn_ln, fc1, fc1_input_scales, alpha) + num += 1 diff --git a/llm-awq/awq/quantize/w8a8_linear.py b/llm-awq/awq/quantize/w8a8_linear.py new file mode 100644 index 0000000000000000000000000000000000000000..b733a1d5974ef2b3c2f3650f32d4d5dd95d32e1e --- /dev/null +++ b/llm-awq/awq/quantize/w8a8_linear.py @@ -0,0 +1,276 @@ +# Adapted from qserve (https://github.com/mit-han-lab/qserve/tree/main) and modified by Yuming Lou + + +from typing import Optional, Union +from torch.nn import Parameter +import awq_inference_engine +import torch +import gc +from awq.utils.module import set_op_by_name +from tqdm import tqdm + + +class W8A8OF16LinearStaticScale(torch.nn.Module): + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = True, + scale: Union[torch.tensor, float] = 1.0, + params_dtype: Optional[torch.dtype] = None, + ): + super().__init__() + + # Keep input parameters + self.in_features = in_features + self.out_features = out_features + # size [1] or size [oc] + self.register_buffer( + "dequant_scale", torch.ones(out_features, dtype=torch.half) + ) + # Parameters. + # NOTE: torch.nn.functional.linear performs XA^T + b and as a result + # we allocate the transpose. + self.create_weights() + + if bias: + self.bias = torch.empty( + self.out_features, + device=torch.cuda.current_device(), + dtype=torch.float16, + ) + else: + self.register_parameter("bias", None) + + def create_weights(self) -> None: + self.register_buffer( + "weight", + torch.empty( + self.out_features, + self.in_features, + dtype=torch.int8, + requires_grad=False, + ), + ) + + def apply_weights( + self, + x: torch.Tensor, + bias: Optional[torch.Tensor], + ) -> torch.Tensor: + raise NotImplementedError + + def forward(self, input_): + # Matrix multiply. + output = self.apply_weights(input_, self.bias) + output_bias = self.bias + return output, output_bias + + +class W8A8OF16LinearDynamicInputScale(W8A8OF16LinearStaticScale): + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = True, + scale: Union[torch.tensor, float] = 1.0, + params_dtype: Optional[torch.dtype] = None, + ): + super().__init__( + in_features=in_features, + out_features=out_features, + bias=bias, + scale=scale, + params_dtype=params_dtype, + ) + if bias: + self.apply_weights = self.apply_weights_bias + else: + self.apply_weights = self.apply_weights_no_bias + + #W bias. Fused bias and W8A8 GEMM + def apply_weights_bias( + self, + # [batch, tokens, channels] + x: torch.Tensor, + # [batch * tokens] + input_scale: torch.Tensor, + output_buffer: torch.Tensor, + bias: torch.Tensor = None, + ): + x_shape = x.shape + if len(x.shape) > 2: + assert 0, "Not implemented" + x = x.view(-1, x_shape[-1]) + # If use awq_inference_engine.w8a8_gemm_fuse_bias_forward_cuda + awq_inference_engine.w8a8_gemm_fuse_bias_forward_cuda( + x, self.weight, self.dequant_scale, input_scale, output_buffer, bias + ) + if len(x.shape) > 2: + assert 0, "Not implemented 2" + output_buffer = output_buffer.view(*x_shape[:-1], -1) + + #W/H bias. W8A8 GEMM + def apply_weights_no_bias( + self, + # [batch, tokens, channels] + x: torch.Tensor, + # [batch * tokens] + input_scale: torch.Tensor, + output_buffer: torch.Tensor, + bias: torch.Tensor = None, + ): + x_shape = x.shape + if len(x.shape) > 2: + assert 0, "Not implemented" + x = x.view(-1, x_shape[-1]) + # If use awq_inference_engine.w8a8_gemm_forward_cuda + awq_inference_engine.w8a8_gemm_forward_cuda( + x, self.weight, self.dequant_scale, input_scale, output_buffer + ) + if len(x.shape) > 2: + assert 0, "Not implemented 2" + output_buffer = output_buffer.view(*x_shape[:-1], -1) + + def forward(self, input_, input_scale, output_buffer): + # Matrix multiply. + self.apply_weights(input_, input_scale, output_buffer, self.bias) + + @classmethod + def from_linear( + cls, + linear, + init_only=False, + s1_scale=None, + fc1=False, + ): + q_linear = cls( + linear.in_features, + linear.out_features, + linear.bias is not None, + ) + if init_only: # just prepare for loading sd + return q_linear + if s1_scale is None: + s1_scale, _ = torch.max(abs(linear.weight.data), dim=-1, keepdim=True) + s1_scale = s1_scale.clamp_(min=1e-5).div_(127) + + if linear.bias is not None: + q_linear.bias = linear.bias.clone().half().contiguous().cuda() + ## Quantize the weights + # ---- Quantize the weights to int8 ---- # + linear_weight = linear.weight.data # OC, IC + linear_weight = linear_weight.div_(s1_scale.to(linear_weight.device)) + linear_weight = linear_weight.round_().to(torch.int8) + + q_linear.weight.data[:, :] = linear_weight.half().contiguous().cuda() + + # ---- Pack the scales ---- # + q_linear.dequant_scale.data[:] = ( + s1_scale.reshape(-1).half().contiguous().cuda() + ) + return q_linear.cuda() + + @classmethod + def from_qkv( + cls, + q, + k, + v, + init_only=False, + s1_scale=None, + ): + q_linear = cls( + q.in_features, + q.out_features + k.out_features + v.out_features, + q.bias is not None, + ) + if init_only: # just prepare for loading sd + return q_linear + weight = torch.cat([q.weight.data, k.weight.data, v.weight.data], dim=0) + + if s1_scale is None: + s1_scale, _ = torch.max(abs(weight), dim=-1, keepdim=True) + s1_scale = s1_scale.clamp_(min=1e-5).div_(127) + + if q.bias is not None: + bias = torch.cat([q.bias, k.bias, v.bias], dim=0) + q_linear.bias = bias.clone().half().contiguous().cuda() + # ---- Quantize the weights to int8 ---- # + weight = weight.div_(s1_scale.to(weight.device)) + weight = weight.round_().to(torch.int8) + + q_linear.weight.data[:, :] = weight.contiguous().cuda() + + # ---- Pack the scales ---- # + q_linear.dequant_scale.data[:] = ( + s1_scale.reshape(q.out_features + k.out_features + v.out_features) + .half() + .contiguous().cuda() + ) + return q_linear.cuda() + + +class FakeW8A8Linear(torch.nn.Module): + def __init__( + self, in_features: int, out_features: int, bias: bool = True, wbit: int = 8 + ): + super().__init__() + self.weight = torch.nn.Parameter( + torch.empty(out_features, in_features, dtype=torch.half) + ) + if bias: + self.bias = torch.nn.Parameter( + torch.empty(1, out_features, dtype=torch.half) + ) + else: + self.bias = None + self.wbit = wbit + self.maxv = 2 ** (wbit - 1) - 1 + + def forward(self, input): + t_shape = input.shape + input.view(-1, t_shape[-1]) + scales = input.abs().max(dim=-1, keepdim=True)[0] + scales.clamp_(min=1e-5).div_(self.maxv) + input.div_(scales).round_().mul_(scales) + output = torch.functional.F.linear(input, self.weight, self.bias) + return output + + @classmethod + def from_linear(cls, linear: torch.nn.Linear, wbit=8): + fake_linear = cls( + linear.in_features, linear.out_features, linear.bias is not None, wbit + ) + maxv = 2 ** (wbit - 1) - 1 + scale = ( + torch.max(abs(linear.weight.data.detach()), -1, keepdim=True)[0] + .clamp_(min=1e-5) + .div_(maxv) + ) + weight = linear.weight.data / scale + weight = weight.round_() + weight = weight * scale + fake_linear.weight.copy_(weight.contiguous()) + if linear.bias is not None: + fake_linear.bias.copy_( + linear.bias.detach().half().reshape(1, linear.out_features).contiguous() + ) + else: + linear.bias = None + del linear, scale, weight + torch.cuda.empty_cache() + return fake_linear + + +def fake_quant(model, wbit=8): + for name, m in tqdm( + model.named_modules(), + desc="Fake quantizing", + total=len(list(model.named_modules())), + ): + if isinstance(m, torch.nn.Linear): + FQlinear = FakeW8A8Linear.from_linear(m, wbit) + del m + torch.cuda.empty_cache() + set_op_by_name(model, name, FQlinear) diff --git a/llm-awq/awq/utils/__init__.py b/llm-awq/awq/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/llm-awq/awq/utils/__pycache__/__init__.cpython-311.pyc b/llm-awq/awq/utils/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..466bfee723919b6b33537604eca3898fadaf74b6 Binary files /dev/null and b/llm-awq/awq/utils/__pycache__/__init__.cpython-311.pyc differ diff --git a/llm-awq/awq/utils/__pycache__/parallel.cpython-311.pyc b/llm-awq/awq/utils/__pycache__/parallel.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7616a7a63fda4388c326d62c7f6ee234c7712c69 Binary files /dev/null and b/llm-awq/awq/utils/__pycache__/parallel.cpython-311.pyc differ diff --git a/llm-awq/awq/utils/lm_eval_adaptor.py b/llm-awq/awq/utils/lm_eval_adaptor.py new file mode 100644 index 0000000000000000000000000000000000000000..8115702971b44876dcf237b92b6a34d8cd93a5f4 --- /dev/null +++ b/llm-awq/awq/utils/lm_eval_adaptor.py @@ -0,0 +1,116 @@ +import transformers +import torch +from lm_eval.base import BaseLM +import fnmatch + + +class LMEvalAdaptor(BaseLM): + def __init__(self, model_name, model, tokenizer, batch_size=1, max_length=-1): + super().__init__() + + assert isinstance(batch_size, int) + + self.model_name = model_name + self.model = model + self.model.eval() + + self.tokenizer = tokenizer + + # assert isinstance(self.tokenizer, ( + # transformers.GPT2Tokenizer, transformers.GPT2TokenizerFast, + # transformers.T5Tokenizer, transformers.T5TokenizerFast, + # )), "this tokenizer has not been checked for compatibility yet!" + + self.vocab_size = self.tokenizer.vocab_size + + self._batch_size = batch_size + + self._max_length = max_length + + @property + def eot_token_id(self): + # we use EOT because end of *text* is more accurate for what we're doing than end of *sentence* + return self.tokenizer.eos_token_id + + @property + def max_length(self): + if self._max_length != -1: + return self._max_length + if hasattr(self.model.config, "n_ctx"): + return self.model.config.n_ctx + elif hasattr(self.model.config, "max_position_embeddings"): + return self.model.config.max_position_embeddings + elif hasattr(self.model.config, "n_positions"): + return self.model.config.n_positions + elif "bloom" in self.model_name: + return 2048 + elif "llama" in self.model_name: + return 2048 # TODO: did not check this + elif "mpt" in self.model_name: + return 2048 + elif "falcon" in self.model_name: + return 2048 + else: + print(self.model.config) + raise NotImplementedError + + @property + def max_gen_toks(self): + return 256 + + @property + def batch_size(self): + return self._batch_size + + @property + def device(self): + return "cuda" + + def tok_encode(self, string: str): + return self.tokenizer.encode(string, add_special_tokens=False) + + def tok_decode(self, tokens): + return self.tokenizer.decode(tokens) + + def _model_call(self, inps): + """ + inps: a torch tensor of shape [batch, sequence] + the size of sequence may vary from call to call + + returns: a torch tensor of shape [batch, sequence, vocab] with the + logits returned from the model + """ + with torch.no_grad(): + if isinstance( + self.model, + transformers.models.t5.modeling_t5.T5ForConditionalGeneration, + ): + dec_inps = torch.cat( + [ + torch.tensor( + self.model.generation_config.decoder_start_token_id, + ) + .tile(len(inps), 1) + .to(inps), + inps, + ], + dim=1, + ) + + kwargs = { + "decoder_input_ids": dec_inps, + } + else: + kwargs = {} + out = self.model(inps, **kwargs)[0] + if ( + "opt" in self.model_name + ): # there are a few extra tokens in opt, which we should omit + return out[:, :, :50257] + else: + return out # [:, :, :self.tokenizer.vocab_size] + + def _model_generate(self, context, max_length, eos_token_id): + return self.model.generate( + context, max_length=max_length, eos_token_id=eos_token_id, do_sample=False + ) diff --git a/llm-awq/awq/utils/module.py b/llm-awq/awq/utils/module.py new file mode 100644 index 0000000000000000000000000000000000000000..3c150b59e59927ad4632947307459acdd76a9437 --- /dev/null +++ b/llm-awq/awq/utils/module.py @@ -0,0 +1,39 @@ +def get_op_by_name(module, op_name): + # get the op by its name relative to the module + for name, m in module.named_modules(): + if name == op_name: + return m + raise ValueError(f"Cannot find op {op_name} in module {module}") + + +def set_op_by_name(layer, name, new_module): + levels = name.split(".") + if len(levels) > 1: + mod_ = layer + for l_idx in range(len(levels) - 1): + if levels[l_idx].isdigit(): + mod_ = mod_[int(levels[l_idx])] + else: + mod_ = getattr(mod_, levels[l_idx]) + setattr(mod_, levels[-1], new_module) + else: + setattr(layer, name, new_module) + + +def get_op_name(module, op): + # get the name of the op relative to the module + for name, m in module.named_modules(): + if m is op: + return name + raise ValueError(f"Cannot find op {op} in module {module}") + + +def append_str_prefix(x, prefix): + if isinstance(x, str): + return prefix + x + elif isinstance(x, tuple): + return tuple([append_str_prefix(y, prefix) for y in x]) + elif isinstance(x, list): + return [append_str_prefix(y, prefix) for y in x] + else: + return x diff --git a/llm-awq/awq/utils/parallel.py b/llm-awq/awq/utils/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..9aed0f657eae43e8b33c545f628b5d0da47cb02f --- /dev/null +++ b/llm-awq/awq/utils/parallel.py @@ -0,0 +1,28 @@ +import os +import torch +import gc + + +def auto_parallel(args): + model_size = args.model_path.split("-")[-1] + if model_size.endswith("m"): + model_gb = 1 + else: + model_gb = float(model_size[:-1]) + if model_gb < 20: + n_gpu = 1 + elif model_gb < 50: + n_gpu = 4 + else: + n_gpu = 8 + args.parallel = n_gpu > 1 + cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None) + if isinstance(cuda_visible_devices, str): + cuda_visible_devices = cuda_visible_devices.split(",") + else: + cuda_visible_devices = list(range(8)) + os.environ["CUDA_VISIBLE_DEVICES"] = ",".join( + [str(dev) for dev in cuda_visible_devices[:n_gpu]] + ) + print("CUDA_VISIBLE_DEVICES: ", os.environ["CUDA_VISIBLE_DEVICES"]) + return cuda_visible_devices diff --git a/llm-awq/awq/utils/utils.py b/llm-awq/awq/utils/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..708e62edd897c7f8264c8fdc0abe883528aea4a9 --- /dev/null +++ b/llm-awq/awq/utils/utils.py @@ -0,0 +1,51 @@ +import torch +import accelerate + + +def get_module_by_name_suffix(model, module_name: str): + for name, module in model.named_modules(): + if name.endswith(module_name): + return module + + +def simple_dispatch_model(model, device_map): + from accelerate.hooks import add_hook_to_module, AlignDevicesHook + + if "" in device_map: + d = device_map[""] + model = model.to(torch.device(d)) + model.hf_device_map = device_map + return model + + tied_params = accelerate.utils.modeling.find_tied_parameters(model) + if set(device_map.values()) == {"cpu"} or set(device_map.values()) == { + "cpu", + "disk", + }: + main_device = "cpu" + else: + main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0] + + cpu_offload_group = [(n, d) for n, d in device_map.items() if d == "cpu"] + prev_hook = None + for idx, (n, d) in enumerate(cpu_offload_group): + m = get_module_by_name_suffix(model, n) + _, prev_hook = accelerate.cpu_offload_with_hook( + m, execution_device=main_device, prev_module_hook=prev_hook + ) + # set first cpu offload module's prev_module_hook to the last cpu offload module's hook + if len(cpu_offload_group) > 1: + get_module_by_name_suffix( + model, cpu_offload_group[0][0] + )._hf_hook.prev_module_hook = prev_hook + + for n, d in device_map.items(): + m = get_module_by_name_suffix(model, n) + if d != "cpu": + d = torch.device(d) + hook = AlignDevicesHook(d, io_same_device=True, place_submodules=True) + add_hook_to_module(m, hook) + accelerate.utils.modeling.retie_parameters(model, tied_params) + model.hf_device_map = device_map + + return model diff --git a/llm-awq/pyproject.toml b/llm-awq/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..7a170943c8fdf5f5cfb208aca1da1668d453a1e2 --- /dev/null +++ b/llm-awq/pyproject.toml @@ -0,0 +1,31 @@ +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "awq" +version = "0.1.0" +description = "An efficient and accurate low-bit weight quantization(INT3/4) method for LLMs." +readme = "README.md" +requires-python = ">=3.8" +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: Apache Software License", +] +dependencies = [ + "accelerate==0.34.2", "sentencepiece", "tokenizers>=0.12.1", + "torch==2.3.0", "torchvision==0.18.0", + "transformers==4.46.0", + "lm_eval==0.3.0", "texttable", + "toml", "attributedict", + "protobuf", + "gradio==3.35.2", "gradio_client==0.2.9", + "fastapi", "uvicorn", + "pydantic==1.10.19" +] + +[tool.setuptools.packages.find] +exclude = ["results*", "scripts*", "examples*"] + +[tool.wheel] +exclude = ["results*", "scripts*", "examples*"] diff --git a/llm-awq/scripts/codellama_example.sh b/llm-awq/scripts/codellama_example.sh new file mode 100644 index 0000000000000000000000000000000000000000..6c5375f2553894095b40a2292a0f1f63657333b3 --- /dev/null +++ b/llm-awq/scripts/codellama_example.sh @@ -0,0 +1,25 @@ +MODEL=CodeLlama-13b-Instruct + +# run AWQ search (optional; we provided the pre-computed results) +python -m awq.entry --model_path /dataset/codellama-hf/$MODEL \ + --w_bit 4 --q_group_size 128 \ + --run_awq --dump_awq awq_cache/$MODEL-w4-g128.pt + +# evaluate the AWQ quantize model (simulated pseudo quantization) +python -m awq.entry --model_path /dataset/codellama-hf/$MODEL \ + --tasks wikitext \ + --w_bit 4 --q_group_size 128 \ + --load_awq awq_cache/$MODEL-w4-g128.pt \ + --q_backend fake + +# generate real quantized weights (w4) +python -m awq.entry --model_path /dataset/codellama-hf/$MODEL \ + --w_bit 4 --q_group_size 128 \ + --load_awq awq_cache/$MODEL-w4-g128.pt \ + --q_backend real --dump_quant quant_cache/$MODEL-w4-g128-awq.pt + +# load and evaluate the real quantized model (smaller gpu memory usage) +python -m awq.entry --model_path /dataset/codellama-hf/$MODEL \ + --tasks wikitext \ + --w_bit 4 --q_group_size 128 \ + --load_quant quant_cache/$MODEL-w4-g128-awq.pt \ No newline at end of file diff --git a/llm-awq/scripts/llama2_example.sh b/llm-awq/scripts/llama2_example.sh new file mode 100644 index 0000000000000000000000000000000000000000..dff0a71f9ca3c733b1543dffbed12089b078d50f --- /dev/null +++ b/llm-awq/scripts/llama2_example.sh @@ -0,0 +1,25 @@ +MODEL=llama-2-7b + +# run AWQ search (optional; we provided the pre-computed results) +python -m awq.entry --model_path /dataset/llama2-hf/$MODEL \ + --w_bit 4 --q_group_size 128 \ + --run_awq --dump_awq awq_cache/$MODEL-w4-g128.pt + +# evaluate the AWQ quantize model (simulated pseudo quantization) +python -m awq.entry --model_path /dataset/llama2-hf/$MODEL \ + --tasks wikitext \ + --w_bit 4 --q_group_size 128 \ + --load_awq awq_cache/$MODEL-w4-g128.pt \ + --q_backend fake + +# generate real quantized weights (w4) +python -m awq.entry --model_path /dataset/llama2-hf/$MODEL \ + --w_bit 4 --q_group_size 128 \ + --load_awq awq_cache/$MODEL-w4-g128.pt \ + --q_backend real --dump_quant quant_cache/$MODEL-w4-g128-awq.pt + +# load and evaluate the real quantized model (smaller gpu memory usage) +python -m awq.entry --model_path /dataset/llama2-hf/$MODEL \ + --tasks wikitext \ + --w_bit 4 --q_group_size 128 \ + --load_quant quant_cache/$MODEL-w4-g128-awq.pt \ No newline at end of file diff --git a/llm-awq/scripts/llama3_example.sh b/llm-awq/scripts/llama3_example.sh new file mode 100644 index 0000000000000000000000000000000000000000..9ee886c7daf887db80ae96746d743d4fd77f1663 --- /dev/null +++ b/llm-awq/scripts/llama3_example.sh @@ -0,0 +1,25 @@ +MODEL=llama3-8b + +# run AWQ search (optional; we provided the pre-computed results) +python -m awq.entry --model_path /dataset/models/llama3/$MODEL \ + --w_bit 4 --q_group_size 128 \ + --run_awq --dump_awq awq_cache/$MODEL-w4-g128.pt + +# evaluate the AWQ quantize model (simulated pseudo quantization) +python -m awq.entry --model_path /dataset/models/llama3/$MODEL \ + --tasks wikitext \ + --w_bit 4 --q_group_size 128 \ + --load_awq awq_cache/$MODEL-w4-g128.pt \ + --q_backend fake + +# generate real quantized weights (w4) +python -m awq.entry --model_path /dataset/models/llama3/$MODEL \ + --w_bit 4 --q_group_size 128 \ + --load_awq awq_cache/$MODEL-w4-g128.pt \ + --q_backend real --dump_quant quant_cache/$MODEL-w4-g128-awq.pt + +# load and evaluate the real quantized model (smaller gpu memory usage) +python -m awq.entry --model_path /dataset/models/llama3/$MODEL \ + --tasks wikitext \ + --w_bit 4 --q_group_size 128 \ + --load_quant quant_cache/$MODEL-w4-g128-awq.pt \ No newline at end of file diff --git a/lm-evaluation-harness/lm_eval/models/__pycache__/ibm_watsonx_ai.cpython-310.pyc b/lm-evaluation-harness/lm_eval/models/__pycache__/ibm_watsonx_ai.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2c0dc97c8054cc0fdf35d926605a4e18f8999963 Binary files /dev/null and b/lm-evaluation-harness/lm_eval/models/__pycache__/ibm_watsonx_ai.cpython-310.pyc differ diff --git a/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_Social.yaml b/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_Social.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3042a419e6e3902ddd0090028fc4b875a148a213 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/aexams/aexams_Social.yaml @@ -0,0 +1,4 @@ +"dataset_name": "Social" +"description": "قم بالإجابة على مايلي في مجال العلوم الإجتماعية \n\n" +"include": "_default_template_yaml" +"task": "aexams_Social" diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_amh.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_amh.yaml new file mode 100644 index 0000000000000000000000000000000000000000..23007e3657c85b3b42ac5591180096c54740a240 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_amh.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: amh +include: afrimgsm_yaml +task: afrimgsm_amh_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_eng.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_eng.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1d5694225089b96dfeb06d331482bafa821cede8 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_eng.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: eng +include: afrimgsm_yaml +task: afrimgsm_eng_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_ewe.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_ewe.yaml new file mode 100644 index 0000000000000000000000000000000000000000..68980400de696f30ce325ff260b5f8cefd5d95dc --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_ewe.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: ewe +include: afrimgsm_yaml +task: afrimgsm_ewe_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..aef377d292006471eecc551eb414ced3c751eaa5 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_hau.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: hau +include: afrimgsm_yaml +task: afrimgsm_hau_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_lin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_lin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..be6a24ceb1cdec606a11b31865f75a8ef5188b4a --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_lin.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: lin +include: afrimgsm_yaml +task: afrimgsm_lin_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_lug.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_lug.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4bffe69f252d7edda8e146a6f61b81dc0bd550c5 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_lug.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: lug +include: afrimgsm_yaml +task: afrimgsm_lug_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_sna.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_sna.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1afa6bb3a5455f405d9735b933844a0974ec0899 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_sna.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: sna +include: afrimgsm_yaml +task: afrimgsm_sna_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_sot.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_sot.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f6cd4cdcd7e4860f73fa0bbdffe16ce02b2d2234 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_sot.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: sot +include: afrimgsm_yaml +task: afrimgsm_sot_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_twi.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_twi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a82235d7883d0fb0b23e92129ed373d2503a31e1 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_twi.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: twi +include: afrimgsm_yaml +task: afrimgsm_twi_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_vai.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_vai.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e04d28f0d33f1da1d6282431c8d4e1655a1f175b --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_vai.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: vai +include: afrimgsm_yaml +task: afrimgsm_vai_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_yor.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_yor.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ba89f9ace0bf2aff3e467e021408d6de790ddd0f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_1/afrimgsm_yor.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: yor +include: afrimgsm_yaml +task: afrimgsm_yor_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_eng.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_eng.yaml new file mode 100644 index 0000000000000000000000000000000000000000..940000376a025ef42c7352c2db416b2fe1a9b38f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_eng.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: eng +include: afrimgsm_yaml +task: afrimgsm_eng_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_ewe.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_ewe.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d0cd4926120ab7de80c13ba9b13bd327c81866bb --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_ewe.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: ewe +include: afrimgsm_yaml +task: afrimgsm_ewe_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3e7d62ab615e7776f714f717d3b0f246782c8f21 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_hau.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: hau +include: afrimgsm_yaml +task: afrimgsm_hau_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_kin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_kin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9c7b65251577ccb7bfa1fe74ddbe06114247dae5 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_kin.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: kin +include: afrimgsm_yaml +task: afrimgsm_kin_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_lin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_lin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..71594885ab6bb615f95c4a0a4f17d24ba41bbd1d --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_lin.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: lin +include: afrimgsm_yaml +task: afrimgsm_lin_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_lug.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_lug.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cf4633745e06b7af03cf890f7d2b9426179bbdbc --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_lug.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: lug +include: afrimgsm_yaml +task: afrimgsm_lug_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_orm.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_orm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e3d9d96ed900fb7583da0b676616f069ec9bf5ee --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_orm.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: orm +include: afrimgsm_yaml +task: afrimgsm_orm_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_sot.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_sot.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5ad77562229ec921e07ebcfe8a35ce9c7b072d99 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_sot.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: sot +include: afrimgsm_yaml +task: afrimgsm_sot_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_swa.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_swa.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fea74a3d1bd1bad4946491457ef689d953b8f66a --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_swa.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: swa +include: afrimgsm_yaml +task: afrimgsm_swa_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_twi.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_twi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..311639a12c103194dec5ccc6894bb175ac67cc26 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_twi.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: twi +include: afrimgsm_yaml +task: afrimgsm_twi_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_vai.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_vai.yaml new file mode 100644 index 0000000000000000000000000000000000000000..655b23dec64404bb271b558726ef5f279096f4c0 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_vai.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: vai +include: afrimgsm_yaml +task: afrimgsm_vai_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_wol.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_wol.yaml new file mode 100644 index 0000000000000000000000000000000000000000..493551623d4071cfebbf59f0a179f487a0f46866 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_wol.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: wol +include: afrimgsm_yaml +task: afrimgsm_wol_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_xho.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_xho.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1c076be5a556970703a4f1581700d1c0b2a3217d --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_xho.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: xho +include: afrimgsm_yaml +task: afrimgsm_xho_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_zul.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_zul.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0b340608c1289ef7a6257cec74ea0956039d12a6 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_2/afrimgsm_zul.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: zul +include: afrimgsm_yaml +task: afrimgsm_zul_prompt_2 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_eng.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_eng.yaml new file mode 100644 index 0000000000000000000000000000000000000000..87efd748800c67b8703318535649280952b21208 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_eng.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: eng +include: afrimgsm_yaml +task: afrimgsm_eng_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_ewe.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_ewe.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4636069a5a6e9291fd071b1221c87927f895264f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_ewe.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: ewe +include: afrimgsm_yaml +task: afrimgsm_ewe_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_hau.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_hau.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c0cfcbe44e862238c16a5c12feadbedca447a3ee --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_hau.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: hau +include: afrimgsm_yaml +task: afrimgsm_hau_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_kin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_kin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f01edaae271f79a4f3f9e6ce444e62af86a85032 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_kin.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: kin +include: afrimgsm_yaml +task: afrimgsm_kin_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_orm.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_orm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8a5ff414f680b347e5d9b7f9504be35b968b4d7d --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_orm.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: orm +include: afrimgsm_yaml +task: afrimgsm_orm_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_sna.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_sna.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4000b9dc41fa79e16f673db2f54fa31d77dded02 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_sna.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: sna +include: afrimgsm_yaml +task: afrimgsm_sna_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_twi.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_twi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..12652498abab497f610192e3dd6d377110e913d4 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_twi.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: twi +include: afrimgsm_yaml +task: afrimgsm_twi_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_wol.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_wol.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2a8fd58b055584229c5143e5cc862c250a2b3e0b --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_wol.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: wol +include: afrimgsm_yaml +task: afrimgsm_wol_prompt_3 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_yaml new file mode 100644 index 0000000000000000000000000000000000000000..8dd3f5ca74a7c6d4dcba9daab9d8c7653b9c8e6e --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_3/afrimgsm_yaml @@ -0,0 +1,34 @@ +tag: + - afrimgsm_tasks + - afrimgsm_tasks_prompt_3 +dataset_path: masakhane/afrimgsm +output_type: generate_until +test_split: test +doc_to_target: '{% if answer is not none %}{{answer[21:]}}{% else %}{{answer_number|string}}{% endif %}' +doc_to_text: "Solve the following math question \n\nQuestion: {{question}} \nAnswer: " +target_delimiter: "" +generation_kwargs: + do_sample: false + until: + - 'Question:' + - + - <|im_end|> +filter_list: + - name: remove_whitespace + filter: + - function: remove_whitespace + - function: take_first + - filter: + - function: regex + group_select: -1 + regex_pattern: (-?[$0-9.,]{2,})|(-?[0-9]+) + - function: take_first + name: flexible-extract +metric_list: + - metric: exact_match + aggregation: mean + higher_is_better: true + ignore_case: true + ignore_punctuation: true +metadata: + version: 2.0 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_amh.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_amh.yaml new file mode 100644 index 0000000000000000000000000000000000000000..00e45eb4d433718f0ee36896297a20e3aced9dc6 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_amh.yaml @@ -0,0 +1,7 @@ +# Generated by utils.py +dataset_name: amh +doc_to_text: "Answer the given question with the appropriate numerical value, ensuring\ + \ that the response is clear and without any supplementary information. \n\nQuestion:\ + \ {{question}} \nAnswer: " +include: afrimgsm_yaml +task: afrimgsm_amh_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_eng.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_eng.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0e3abef526bb0571d2a6fbac07a2abeff6acb3a8 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_eng.yaml @@ -0,0 +1,7 @@ +# Generated by utils.py +dataset_name: eng +doc_to_text: "Answer the given question with the appropriate numerical value, ensuring\ + \ that the response is clear and without any supplementary information. \n\nQuestion:\ + \ {{question}} \nAnswer: " +include: afrimgsm_yaml +task: afrimgsm_eng_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_fra.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_fra.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b7e5c155519d43922d3391bfe903ab8720d9ddbf --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_fra.yaml @@ -0,0 +1,7 @@ +# Generated by utils.py +dataset_name: fra +doc_to_text: "Answer the given question with the appropriate numerical value, ensuring\ + \ that the response is clear and without any supplementary information. \n\nQuestion:\ + \ {{question}} \nAnswer: " +include: afrimgsm_yaml +task: afrimgsm_fra_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_lin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_lin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..44d2d2f39bc83011268c1bcffaa4ea8564e3b5a4 --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_lin.yaml @@ -0,0 +1,7 @@ +# Generated by utils.py +dataset_name: lin +doc_to_text: "Answer the given question with the appropriate numerical value, ensuring\ + \ that the response is clear and without any supplementary information. \n\nQuestion:\ + \ {{question}} \nAnswer: " +include: afrimgsm_yaml +task: afrimgsm_lin_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_twi.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_twi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f3e030d0bc072ba50afa51a988ee0fa4d37f8e7f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_twi.yaml @@ -0,0 +1,7 @@ +# Generated by utils.py +dataset_name: twi +doc_to_text: "Answer the given question with the appropriate numerical value, ensuring\ + \ that the response is clear and without any supplementary information. \n\nQuestion:\ + \ {{question}} \nAnswer: " +include: afrimgsm_yaml +task: afrimgsm_twi_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_xho.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_xho.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a759549811b8e6f009ed49c7b4bd16bc46163d8f --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct/prompt_4/afrimgsm_xho.yaml @@ -0,0 +1,7 @@ +# Generated by utils.py +dataset_name: xho +doc_to_text: "Answer the given question with the appropriate numerical value, ensuring\ + \ that the response is clear and without any supplementary information. \n\nQuestion:\ + \ {{question}} \nAnswer: " +include: afrimgsm_yaml +task: afrimgsm_xho_prompt_4 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct_cot/prompt_1/afrimgsm_cot_lin.yaml b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct_cot/prompt_1/afrimgsm_cot_lin.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ecbf38d859c25896eb039d87ed9b36a4797a77de --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/direct_cot/prompt_1/afrimgsm_cot_lin.yaml @@ -0,0 +1,4 @@ +# Generated by utils.py +dataset_name: lin +include: afrimgsm_cot_yaml +task: afrimgsm_cot_lin_prompt_1 diff --git a/lm-evaluation-harness/lm_eval/tasks/afrimgsm/gen_utils.py b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/gen_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ecef389f3a4051e57b652f617b19ddd15d3c26ca --- /dev/null +++ b/lm-evaluation-harness/lm_eval/tasks/afrimgsm/gen_utils.py @@ -0,0 +1,122 @@ +import argparse +import os + +import yaml + + +class FunctionTag: + def __init__(self, value): + self.value = value + + +def prompt_func(mode, lang): + prompt_map = { + "prompt_4": "Answer the given question with the step by step solution appropriate numerical value, ensuring that the response is " + "clear and without any supplementary information. \n\nQuestion: {{question}} \nStep by step answer: ", + "prompt_5": f"For mathematical questions provided in {lang} language. Supply the accurate step by step answer to the " + "provided question. \n\nQuestion: {{question}} \nStep by step answer: ", + } + return prompt_map[mode] + + +def gen_lang_yamls(output_dir: str, overwrite: bool, mode: str) -> None: + """ + Generate a yaml file for each language. + + :param output_dir: The directory to output the files to. + :param overwrite: Whether to overwrite files if they already exist. + """ + err = [] + languages = { + "eng": "English", + "amh": "Amharic", + "ibo": "Igbo", + "fra": "French", + "sna": "chiShona", + "wol": "Wolof", + "ewe": "Ewe", + "lin": "Lingala", + "lug": "Luganda", + "xho": "isiXhosa", + "kin": "Kinyarwanda", + "twi": "Twi", + "zul": "Zulu", + "orm": "Oromo", + "yor": "Yoruba", + "hau": "Hausa", + "sot": "Sesotho", + "swa": "Swahili", + "vai": "Vai", + } + + for lang in languages.keys(): + try: + file_name = f"afrimgsm_cot_{lang}.yaml" + task_name = f"afrimgsm_cot_{lang}_{mode}" + yaml_template = "afrimgsm_cot_yaml" + if "translate" in output_dir.split("/")[-1]: + file_name = f"afrimgsm_cot_translate_{lang}.yaml" + task_name = f"afrimgsm_cot_translate_{lang}_{mode}" + yaml_template = "afrimgsm_cot_translate_yaml" + if int(mode.split("_")[-1]) > 3: + yaml_details = { + "include": yaml_template, + "task": task_name, + "dataset_name": lang, + "doc_to_text": prompt_func(mode, languages[lang]), + } + else: + yaml_details = { + "include": yaml_template, + "task": task_name, + "dataset_name": lang, + } + os.makedirs(f"{output_dir}/{mode}", exist_ok=True) + with open( + f"{output_dir}/{mode}/{file_name}", + "w" if overwrite else "x", + encoding="utf8", + ) as f: + f.write("# Generated by utils.py\n") + yaml.dump( + yaml_details, + f, + allow_unicode=True, + ) + except FileExistsError: + err.append(file_name) + + if len(err) > 0: + raise FileExistsError( + "Files were not created because they already exist (use --overwrite flag):" + f" {', '.join(err)}" + ) + + +def main() -> None: + """Parse CLI args and generate language-specific yaml files.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--overwrite", + default=True, + action="store_true", + help="Overwrite files if they already exist", + ) + parser.add_argument( + "--output-dir", + default="./translate_cot", + help="Directory to write yaml files to", + ) + parser.add_argument( + "--mode", + default="prompt_5", + choices=["prompt_1", "prompt_2", "prompt_3", "prompt_4", "prompt_5"], + help="Prompt number", + ) + args = parser.parse_args() + + gen_lang_yamls(output_dir=args.output_dir, overwrite=args.overwrite, mode=args.mode) + + +if __name__ == "__main__": + main()