Add files using upload-large-folder tool

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
+]

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
+]

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
+]

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
+]

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__

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.

llm-awq/awq.egg-info/top_level.txt

@@ -0,0 +1,3 @@
+awq
+figures
+tinychat

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 <cuda_fp16.h>
+#include <cuda_runtime_api.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#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<typename T>
+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<typename T, bool CROSS_ATTENTION>
+struct Multihead_attention_params: public Multihead_attention_params_base<T> {
+    // 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<typename T>
+struct Multihead_attention_params<T, true>: public Multihead_attention_params_base<T> {
+    // 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<class T>
+using Masked_multihead_attention_params = Multihead_attention_params<T, false>;
+
+template<class T>
+using Cross_multihead_attention_params = Multihead_attention_params<T, true>;
+
+template<typename T>
+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<float>& params, const cudaStream_t& stream);
+void masked_multihead_attention(const Masked_multihead_attention_params<uint16_t>& 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<float>& params, const cudaStream_t& stream);
+void cross_multihead_attention(const Cross_multihead_attention_params<uint16_t>& 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
+
+////////////////////////////////////////////////////////////////////////////////////////////////////

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 <assert.h>
+#include <float.h>
+#include <type_traits>
+
+// #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<typename T, int Dh>
+struct Qk_vec_ {
+};
+
+template<>
+struct Qk_vec_<float, 32> {
+    using Type = float;
+};
+template<>
+struct Qk_vec_<float, 64> {
+    using Type = float2;
+};
+template<>
+struct Qk_vec_<float, 128> {
+    using Type = float4;
+};
+template<>
+struct Qk_vec_<float, 256> {
+    using Type = float4;
+};
+template<>
+struct Qk_vec_<uint16_t, 32> {
+    using Type = uint32_t;
+};
+template<>
+struct Qk_vec_<uint16_t, 64> {
+    using Type = uint32_t;
+};
+template<>
+struct Qk_vec_<uint16_t, 128> {
+    using Type = uint2;
+};
+template<>
+struct Qk_vec_<uint16_t, 256> {
+    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<typename T, int THREADS_PER_KEY>
+struct K_vec_ {
+};
+
+template<>
+struct K_vec_<float, 4> {
+    using Type = float;
+};
+template<>
+struct K_vec_<float, 2> {
+    using Type = float2;
+};
+template<>
+struct K_vec_<float, 1> {
+    using Type = float4;
+};
+template<>
+struct K_vec_<uint16_t, 4> {
+    using Type = uint32_t;
+};
+template<>
+struct K_vec_<uint16_t, 2> {
+    using Type = uint2;
+};
+template<>
+struct K_vec_<uint16_t, 1> {
+    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<typename T, int V_VEC_SIZE>
+struct V_vec_ {
+};
+
+template<>
+struct V_vec_<float, 1> {
+    using Type = float;
+};
+template<>
+struct V_vec_<float, 2> {
+    using Type = float2;
+};
+template<>
+struct V_vec_<float, 4> {
+    using Type = float4;
+};
+template<>
+struct V_vec_<uint16_t, 2> {
+    using Type = uint32_t;
+};
+template<>
+struct V_vec_<uint16_t, 4> {
+    using Type = uint2;
+};
+template<>
+struct V_vec_<uint16_t, 8> {
+    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<typename T>
+struct Qk_vec_acum_fp32_ {
+};
+
+template<>
+struct Qk_vec_acum_fp32_<float> {
+    using Type = float;
+};
+template<>
+struct Qk_vec_acum_fp32_<float2> {
+    using Type = float2;
+};
+template<>
+struct Qk_vec_acum_fp32_<float4> {
+    using Type = float4;
+};
+// template<> struct Qk_vec_acum_fp32_<uint16_t> { using Type = float;        };
+template<>
+struct Qk_vec_acum_fp32_<uint32_t> {
+    using Type = float2;
+};
+template<>
+struct Qk_vec_acum_fp32_<uint2> {
+    using Type = Float4_;
+};
+template<>
+struct Qk_vec_acum_fp32_<uint4> {
+    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_<bf16_4_t> {
+    using Type = Float4_;
+};
+template<>
+struct Qk_vec_acum_fp32_<bf16_8_t> {
+    using Type = Float8_;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T>
+struct K_vec_acum_fp32_ {
+};
+
+template<>
+struct K_vec_acum_fp32_<float> {
+    using Type = float;
+};
+template<>
+struct K_vec_acum_fp32_<float2> {
+    using Type = float2;
+};
+template<>
+struct K_vec_acum_fp32_<float4> {
+    using Type = float4;
+};
+template<>
+struct K_vec_acum_fp32_<uint32_t> {
+    using Type = float2;
+};
+template<>
+struct K_vec_acum_fp32_<uint2> {
+    using Type = Float4_;
+};
+template<>
+struct K_vec_acum_fp32_<uint4> {
+    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_<bf16_4_t> {
+    using Type = Float4_;
+};
+template<>
+struct K_vec_acum_fp32_<bf16_8_t> {
+    using Type = Float8_;
+};
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#ifdef MMHA_USE_FP32_ACUM_FOR_OUT
+template<typename T>
+struct V_vec_acum_fp32_ {
+};
+
+template<>
+struct V_vec_acum_fp32_<float> {
+    using Type = float;
+};
+template<>
+struct V_vec_acum_fp32_<float2> {
+    using Type = float2;
+};
+template<>
+struct V_vec_acum_fp32_<float4> {
+    using Type = float4;
+};
+template<>
+struct V_vec_acum_fp32_<uint32_t> {
+    using Type = float2;
+};
+template<>
+struct V_vec_acum_fp32_<uint2> {
+    using Type = Float4_;
+};
+template<>
+struct V_vec_acum_fp32_<uint4> {
+    using Type = Float8_;
+};
+#ifdef ENABLE_BF16
+template<>
+struct V_vec_acum_fp32_<__nv_bfloat162> {
+    using Type = float2;
+};
+template<>
+struct V_vec_acum_fp32_<bf16_4_t> {
+    using Type = Float4_;
+};
+template<>
+struct V_vec_acum_fp32_<bf16_8_t> {
+    using Type = Float8_;
+};
+#endif  // ENABLE_BF16
+#endif
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<int THREADS_PER_KEY, typename K_vec, int N>
+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_<K_vec>::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<K_vec_acum, K_vec, K_vec>(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<typename T, int THREADS_PER_KEY>
+struct Qk_dot {
+    template<typename K_vec, int N>
+    static inline __device__ float dot(const K_vec (&q)[N], const K_vec (&k)[N])
+    {
+        return qk_dot_<THREADS_PER_KEY>(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<int N>
+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_<uint32_t>::Type;
+#else
+    using K_vec_acum = uint32_t;
+#endif
+    K_vec_acum qk_vec = mul<K_vec_acum, uint32_t, uint32_t>(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<uint16_t, 4> {
+    template<int N>
+    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<int WARPS_PER_BLOCK, int WARP_SIZE = 32>
+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<typename T>
+inline __device__ __host__ T div_up(T m, T n)
+{
+    return (m + n - 1) / n;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T, bool DO_CROSS_ATTENTION>
+inline size_t smem_size_in_bytes(const Multihead_attention_params<T, DO_CROSS_ATTENTION>& 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<T, DO_CROSS_ATTENTION> 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<float*>(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<T*>(logits_smem_);
+#else
+    float* logits_smem = reinterpret_cast<float*>(logits_smem_);
+#endif
+
+    // The shared memory to do the final reduction for the output values. Reuse qk_smem.
+    T* out_smem = reinterpret_cast<T*>(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_<T, Dh_MAX>::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_<T, Dh_MAX>::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<int8_t, num_elems<Qk_vec>::value>::type;
+            using Packed_Float_t = typename packed_type<float, num_elems<Qk_vec>::value>::type;
+            const auto q_scaling = params.qkv_scale_out[0];
+            const auto q_quant =
+                *reinterpret_cast<const Packed_Int8_t*>(&reinterpret_cast<const int8_t*>(params.q)[q_offset]);
+
+            convert_from_float(q, mul<Packed_Float_t, float>(q_scaling, float_from_int8(q_quant)));
+        }
+        else {
+            q = *reinterpret_cast<const Qk_vec*>(&params.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<const Qk_vec*>(&params.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<int8_t, num_elems<Qk_vec>::value>::type;
+                using Packed_Float_t = typename packed_type<float, num_elems<Qk_vec>::value>::type;
+                const auto k_scaling = params.qkv_scale_out[1];
+                const auto k_quant =
+                    *reinterpret_cast<const Packed_Int8_t*>(&reinterpret_cast<const int8_t*>(params.k)[k_offset]);
+
+                convert_from_float(k, mul<Packed_Float_t, float>(k_scaling, float_from_int8(k_quant)));
+            }
+            else {
+                k = *reinterpret_cast<const Qk_vec*>(&params.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<const Qk_vec*>(&params.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<const Qk_vec*>(&params.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<Qk_vec, Qk_vec, Qk_vec>(
+            k,
+            *reinterpret_cast<const Qk_vec*>(
+                &params.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<T*>(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<Qk_vec*>(q_smem + half_idx * smem_pitch + intra_half_idx) = q;
+
+            if (handle_kv) {
+                *reinterpret_cast<Qk_vec*>(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<Qk_vec*>(q_smem + half_idx * smem_pitch + intra_half_idx);
+            if (handle_kv) {
+                k = *reinterpret_cast<Qk_vec*>(k_smem + half_idx * smem_pitch + intra_half_idx);
+            }
+        }
+
+        __syncthreads();
+    }
+
+    if (!is_masked) {
+        // Store the Q values to shared memory.
+        *reinterpret_cast<Qk_vec*>(&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<Qk_vec*>(&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<Qk_vec*>(&params.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_<Qk_vec>::Type;
+#else
+        using Qk_vec_acum = Qk_vec;
+#endif
+        qk = dot<Qk_vec_acum, Qk_vec>(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<WARPS_PER_RED>(&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_<T, THREADS_PER_KEY>::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<const K_vec*>(&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<const K_vec*>(&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 = &params.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 = &params.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 ? &params.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<const K_vec*>(&k_cache_batch[beam_offset + jj * QK_ELTS_IN_16B]);
+                    }
+                    else {
+                        k[ii] = *reinterpret_cast<const K_vec*>(&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_vec, K_vec, K_vec>(
+                            k[ii],
+                            *reinterpret_cast<const K_vec*>(
+                                &params.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_vec*>(&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<T, THREADS_PER_KEY>::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<float, float, float>(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<WARPS_PER_BLOCK>(&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_<T, V_VEC_SIZE>::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 = &params.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 = &params.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<const V_vec*>(&params.v_bias[hi_kv * Dh + vi]);
+                }
+                if (DO_CROSS_ATTENTION) {
+                    *reinterpret_cast<V_vec*>(&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_<V_vec>::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<const V_vec*>(&v_cache_batch[beam_offset + ti_circ * Dh]);
+            if (DO_CROSS_ATTENTION && params.timestep == 0) {
+                v = add(v, *reinterpret_cast<V_vec*>(&bias_smem[vi]));
+                if (do_ia3) {
+                    v = mul<V_vec, V_vec, V_vec>(
+                        v,
+                        *reinterpret_cast<const V_vec*>(
+                            &params.ia3_value_weights[(ia3_task_id * params.num_heads + hi) * Dh + vi]));
+                }
+                *reinterpret_cast<V_vec*>(&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<const V_vec*>(&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<int8_t, num_elems<V_vec>::value>::type;
+                using Packed_Float_t = typename packed_type<float, num_elems<V_vec>::value>::type;
+                const auto v_scaling = params.qkv_scale_out[2];
+                const auto v_quant =
+                    *reinterpret_cast<const Packed_Int8_t*>(&reinterpret_cast<const int8_t*>(params.v)[v_offset]);
+
+                convert_from_float(v, mul<Packed_Float_t, float>(v_scaling, float_from_int8(v_quant)));
+            }
+            else {
+                v = *reinterpret_cast<const V_vec*>(&params.v[v_offset]);
+            }
+            // Trigger the loads from the V bias buffer.
+            // V_vec v_bias = *reinterpret_cast<const V_vec*>(&params.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_vec, V_vec, V_vec>(
+                    v,
+                    *reinterpret_cast<const V_vec*>(
+                        &params.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_vec*>(&v_cache[params.timestep*Dh]) = v;
+            *reinterpret_cast<V_vec*>(&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<V_vec*>(&out_smem[(vo - midpoint) * Dh + vi]), out);
+#else
+                *reinterpret_cast<V_vec*>(&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<const V_vec*>(&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<int8_t, num_elems<V_vec_acum>::value>::type;
+            out                 = mul<V_vec_acum, float>(*params.attention_out_scale, out);
+            *reinterpret_cast<Packed_Int8_t*>(&(reinterpret_cast<int8_t*>(params.out)[bhi * Dh + vi])) =
+                cast_to_int8(out);
+        }
+        else {
+            convert_from_float(*reinterpret_cast<V_vec*>(&params.out[bhi * Dh + vi]), out);
+        }
+#else
+        // TODO: support int8_mode?
+        *reinterpret_cast<V_vec*>(&params.out[bhi * Dh + vi]) = out;
+#endif
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace mmha
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T, int Dh, int Dh_MAX, typename KERNEL_PARAMS_TYPE>
+void mmha_launch_kernel(const KERNEL_PARAMS_TYPE& params, const cudaStream_t& stream);

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 {},
+)

llm-awq/awq/kernels/csrc/layernorm/layernorm.h

@@ -0,0 +1,3 @@
+#include <torch/extension.h>
+
+void layernorm_forward_cuda(torch::Tensor _input, torch::Tensor _gamma, torch::Tensor _out, float eps);

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 <assert.h>
+#if ((__CUDACC_VER_MAJOR__ > 11) || (__CUDACC_VER_MAJOR__ == 11 && __CUDACC_VER_MINOR__ >= 0))
+#include <cooperative_groups/reduce.h>
+#else
+#include <cooperative_groups.h>
+#endif
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <float.h>
+#include <type_traits>
+
+static const float HALF_FLT_MAX = 65504.F;
+#define FINAL_MASK 0xffffffff
+
+
+template<typename T>
+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<typename T>
+__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<typename T>
+__inline__ __device__ T blockReduceSum(T val)
+{
+    static __shared__ T shared[32];
+    int                 lane = threadIdx.x & 0x1f;
+    int                 wid  = threadIdx.x >> 5;
+
+    val = warpReduceSum<T>(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<T>(val);
+
+    return val;
+}
+
+
+template<typename T>
+__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));
+}

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 <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include "pos_encoding.h"
+
+template<typename scalar_t>
+__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<scalar_t><<<grid, block, 0, stream>>>(
+        positions.data_ptr<int64_t>(),
+        query.data_ptr<scalar_t>(),
+        key.data_ptr<scalar_t>(),
+        cos_sin_cache.data_ptr<scalar_t>(),
+        rot_dim,
+        stride,
+        num_heads,
+        head_size);
+    });
+}
+

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<uint32_t*>(&result);
+    uint32_t const i4s = reinterpret_cast<uint32_t const&>(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;
+}
+

llm-awq/awq/kernels/csrc/quantization/gemm_cuda.h

@@ -0,0 +1,4 @@
+#include <torch/extension.h>
+
+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);

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 <torch/extension.h>
+#include "gemm_cuda.h"
+#include <cuda_fp16.h>
+#include <c10/cuda/CUDAGuard.h>
+
+
+// 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 <int G>
+__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<uint4*>(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<half*>(_in_feats.data_ptr<at::Half>());
+    auto kernel = reinterpret_cast<int*>(_kernel.data_ptr<int>());
+    auto out_feats = reinterpret_cast<half*>(_out_feats.data_ptr<at::Half>());
+    auto scaling_factors = reinterpret_cast<half*>(_scaling_factors.data_ptr<at::Half>());
+    auto zeros = reinterpret_cast<int*>(_zeros.data_ptr<int>());
+
+    // 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><<<num_blocks, threads_per_block>>>(
+        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><<<num_blocks, threads_per_block>>>(
+        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);
+}
+

llm-awq/awq/kernels/csrc/quantization/gemv_cuda.h

@@ -0,0 +1,9 @@
+#pragma once
+#include <torch/extension.h>
+
+torch::Tensor gemv_forward_cuda(
+    torch::Tensor _in_feats,
+    torch::Tensor _kernel,
+    torch::Tensor _scaling_factors,
+    torch::Tensor _zeros,
+    int group_size);

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 <cuda_fp16.h>
+#include <cuda_bf16.h>
+#pragma once
+
+template <typename T = half>
+__inline__ __device__ void dequantize_s4_to_fp16x2(half2 const &source, uint4 *result);
+
+template <>
+__inline__ __device__ void dequantize_s4_to_fp16x2<half>(half2 const &source, uint4 *result)
+{
+    // uint4 result;
+
+    uint32_t *h = reinterpret_cast<uint32_t *>(result);
+    uint32_t const i4s = reinterpret_cast<uint32_t const &>(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<nv_bfloat16>(half2 const &source, uint4 *result)
+{
+    // uint4 result;
+
+    uint32_t *h = reinterpret_cast<uint32_t *>(result);
+    uint32_t const i4s = reinterpret_cast<uint32_t const &>(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<const __nv_bfloat162&>(BF16_TOP_MAGIC_NUM));
+    reinterpret_cast<__nv_bfloat162*>(h)[1] = __hsub2(reinterpret_cast<__nv_bfloat162*>(h)[1], reinterpret_cast<const __nv_bfloat162&>(BF16_TOP_MAGIC_NUM));
+    reinterpret_cast<__nv_bfloat162*>(h)[2] = __hsub2(reinterpret_cast<__nv_bfloat162*>(h)[2], reinterpret_cast<const __nv_bfloat162&>(BF16_TOP_MAGIC_NUM));
+    reinterpret_cast<__nv_bfloat162*>(h)[3] = __hsub2(reinterpret_cast<__nv_bfloat162*>(h)[3], reinterpret_cast<const __nv_bfloat162&>(BF16_TOP_MAGIC_NUM));
+}

llm-awq/awq/kernels/csrc/quantization_new/dispatch_utils.cuh

@@ -0,0 +1,18 @@
+#pragma once
+#include <torch/extension.h>
+#include <cstdint>
+#include <sstream>
+#include <stdexcept>
+
+#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());                                                      \
+  }

llm-awq/awq/kernels/csrc/quantization_new/gemm/gemm_cuda.cu

@@ -0,0 +1,1236 @@
+#include <cuda_fp16.h>
+#include "semaphore.h"
+#include "gemm_cuda.h"
+#include "../dequantize.cuh"
+#include "../dispatch_utils.cuh"
+#include <torch/extension.h>
+#include <cuda_pipeline_primitives.h>
+
+#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<int *>(_semaphores.data_ptr<int>());                                                                               \
+  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<CTA_M, CTA_N, CTA_K, WARP_M, WARP_N, WARP_K, STAGES, G, SPLITK, ctype>;                                                     \
+  cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemByteSize);                                                        \
+  kernel_func<<<num_blocks, threads_per_block, kSmemByteSize>>>(                                                                                        \
+      in_feats, kernel, scales, zeros, out_feats, semaphores, num_in_feats, num_out_channels, num_in_channels);
+
+template <int N>
+__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 <int SLICES, int NUM_WARPS_MN>
+__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 <typename T>
+__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 <typename T>
+__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 <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int SHARED_K_ITERS, int STAGES, typename T>
+__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 <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int SHARED_K_ITERS, int STAGES, typename T>
+__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 <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int STAGES, int G, typename T>
+__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 <int CTA_M, int CTA_N, int CTA_K, int STAGES, int shared_iters, typename T>
+__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 <int CTA_M, int CTA_N, int CTA_K, int STAGES, bool ldmatrix, int shared_iters, int G, typename T>
+__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<std::is_same<T, half>::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<T, half>::value)
+    {
+      scale2 = __half2half2(scale);
+      zero2 = __half2half2(zero);
+    }
+    else
+    {
+      scale2 = __bfloat162bfloat162(scale);
+      zero2 = __bfloat162bfloat162(zero);   
+    }
+    T2 loaded[4];
+    dequantize_s4_to_fp16x2<T>(*reinterpret_cast<half2 *>(dst + (k_0_1 % 2) * 4 + (k_0_1 / 2 * 2) + shared_iter * 8), reinterpret_cast<uint4 *>(loaded));
+#pragma unroll
+    for (int i = 0; i < 4; i++)
+    {
+      loaded[i] = __hfma2(loaded[i], scale2, zero2);
+    }
+    *reinterpret_cast<uint4 *>(dst_fp16 + shared_iter * 16 + 8 * (k_0_1 % 2)) = *reinterpret_cast<uint4 *>(loaded);
+  }
+}
+
+template <int CTA_M, int CTA_N, int CTA_K, int WARP_M, int WARP_N, int WARP_K, int STAGES, int G, int SPLITK, typename T>
+__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<std::is_same<T, half>::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<CTA_M, CTA_N, CTA_K, CTA_SIZE, 1, STAGES>(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<CTA_M, CTA_N, CTA_K, CTA_SIZE, 1, STAGES>(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<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES, G>(
+        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<CTA_M, CTA_N, CTA_K, STAGES, WARP_M / INTRIN_M>(A_shared, A_shared_warp_[0], warp_offset_m, warp_offset_n, warp_offset_k, 0);
+  share_to_reg_one_stage_B<CTA_M, CTA_N, CTA_K, STAGES, true, WARP_N / INTRIN_N, G>(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<CTA_M, CTA_N, CTA_K, STAGES, WARP_M / INTRIN_M>(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<CTA_M, CTA_N, CTA_K, STAGES, true, WARP_N / INTRIN_N, G>(
+              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<CTA_M, CTA_N, CTA_K, STAGES, true, WARP_N / INTRIN_N, G>(
+              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<CTA_M, CTA_N, CTA_K, STAGES, false, WARP_N / INTRIN_N, G>(
+              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<CTA_M, CTA_N, CTA_K, STAGES, false, WARP_N / INTRIN_N, G>(
+              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<T, half>::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<CTA_M, CTA_N, CTA_K, CTA_SIZE, WARP_K / INTRIN_K, STAGES>(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<CTA_M, CTA_N, CTA_K, CTA_SIZE, WARP_K / INTRIN_K, STAGES>(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<CTA_M, CTA_N, CTA_K, CTA_SIZE, WARP_K / INTRIN_K, STAGES>(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<CTA_M, CTA_N, CTA_K, CTA_SIZE, WARP_K / INTRIN_K, STAGES>(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<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES, G>(
+            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<T, half>::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<half2 *>(
+                    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<half2 *>(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<half2 *>(
+                    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<half2 *>(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<nv_bfloat162 *>(
+                    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<nv_bfloat162 *>(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<nv_bfloat162 *>(
+                    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<nv_bfloat162 *>(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 <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int SHARED_K_ITERS, int STAGES, typename T>
+__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 <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int SHARED_K_ITERS, int STAGES, typename T>
+__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 <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int STAGES, int G, typename T>
+__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 <int CTA_M, int CTA_N, int CTA_K, int STAGES, int shared_iters, typename T>
+__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 <int CTA_M, int CTA_N, int CTA_K, int STAGES, bool ldmatrix, int shared_iters, int G, typename T>
+__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<std::is_same<T, half>::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<T, half>::value)
+    {
+      scale2 = __half2half2(scale);
+      zero2 = __half2half2(zero);
+    }
+    else
+    {
+      scale2 = __bfloat162bfloat162(scale);
+      zero2 = __bfloat162bfloat162(zero);   
+    }
+    T2 loaded[4];
+    dequantize_s4_to_fp16x2<T>(*reinterpret_cast<half2 *>(dst + (k_0_1 % 2) * 4 + (k_0_1 / 2 * 2) + shared_iter * 8), reinterpret_cast<uint4 *>(loaded));
+#pragma unroll
+    for (int i = 0; i < 4; i++)
+    {
+      loaded[i] = __hfma2(loaded[i], scale2, zero2);
+    }
+    *reinterpret_cast<uint4 *>(dst_fp16 + shared_iter * 16 + 8 * (k_0_1 % 2)) = *reinterpret_cast<uint4 *>(loaded);
+  }
+}
+
+template <int CTA_M, int CTA_N, int CTA_K, int WARP_M, int WARP_N, int WARP_K, int STAGES, int G, typename T>
+__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<std::is_same<T, half>::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<CTA_M, CTA_N, CTA_K, CTA_SIZE, 1, STAGES>(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<CTA_M, CTA_N, CTA_K, CTA_SIZE, 1, STAGES>(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<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES, G>(
+        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<CTA_M, CTA_N, CTA_K, STAGES, WARP_M / INTRIN_M>(A_shared, A_shared_warp_[0], warp_offset_m, warp_offset_n, 0);
+  share_to_reg_one_stage_B_T2<CTA_M, CTA_N, CTA_K, STAGES, true, WARP_N / INTRIN_N, G>(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<CTA_M, CTA_N, CTA_K, STAGES, WARP_M / INTRIN_M>(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<CTA_M, CTA_N, CTA_K, STAGES, true, WARP_N / INTRIN_N, G>(
+              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<CTA_M, CTA_N, CTA_K, STAGES, true, WARP_N / INTRIN_N, G>(
+              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<CTA_M, CTA_N, CTA_K, STAGES, false, WARP_N / INTRIN_N, G>(
+              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<CTA_M, CTA_N, CTA_K, STAGES, false, WARP_N / INTRIN_N, G>(
+              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<T, half>::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<CTA_M, CTA_N, CTA_K, CTA_SIZE, WARP_K / INTRIN_K, STAGES>(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<CTA_M, CTA_N, CTA_K, CTA_SIZE, WARP_K / INTRIN_K, STAGES>(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<CTA_M, CTA_N, CTA_K, CTA_SIZE, WARP_K / INTRIN_K, STAGES>(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<CTA_M, CTA_N, CTA_K, CTA_SIZE, WARP_K / INTRIN_K, STAGES>(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<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES, G>(
+            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<T, half>::value)
+          {
+            *reinterpret_cast<half2 *>(
+                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<half2 *>(C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 +
+                                                              ax1_0_1 * 8 + local_id));
+          }
+          else
+          {
+            *reinterpret_cast<nv_bfloat162 *>(
+                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<float2 *>(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<int64_t> 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<ctype *>(_in_feats.data_ptr());
+    auto kernel = reinterpret_cast<ctype *>(_kernel.data_ptr());
+    auto scales = reinterpret_cast<ctype *>(_scales.data_ptr());
+    auto zeros = reinterpret_cast<ctype *>(_zeros.data_ptr());
+    auto out_feats = reinterpret_cast<ctype *>(_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<CTA_M, CTA_N, CTA_K, WARP_M, WARP_N, WARP_K, STAGES, G, ctype>;
+      cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemByteSize);
+      kernel_func<<<num_blocks, threads_per_block, kSmemByteSize>>>(
+          in_feats, kernel, scales, zeros, out_feats, num_in_feats, num_out_channels, num_in_channels);
+    }
+  });
+
+  return _out_feats;
+}

llm-awq/awq/kernels/csrc/quantization_new/gemm/gemm_cuda.h

@@ -0,0 +1,3 @@
+#include <torch/extension.h>
+
+torch::Tensor gemm_forward_cuda_new(torch::Tensor _in_feats, torch::Tensor _kernel, torch::Tensor _scales, torch::Tensor _zeros);

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
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

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 <cuda_fp16.h>
+#include <stdio.h>
+#include <torch/extension.h>
+#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 <int Num, int WarpSize, typename T>
+__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<float>(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 <int NPerBlock, int Batch, int BlockSize, int GroupSize, typename T>
+__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<T, half>::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<T>(0.f);
+    
+    // extern __shared__ uint8_t shmem[];
+    // float(*out_smem)[Num * kInterleave] = reinterpret_cast<float(*)[Num * kInterleave]>(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<T>(*reinterpret_cast<half2 *>(local_qweights + i), reinterpret_cast<uint4 *>(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<T2*>(
+                          half_weight_buffer + (i + j * kShuffleContinous)* kShuffleBasicTile
+                        );
+                    if constexpr (std::is_same<T, half>::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<T, half>::value)
+                    {                       
+                      *reinterpret_cast<half2*>(psum + batch_idx * NPerBlock + x * 2)
+                          = __hfma2(*reinterpret_cast<half2*>(dequantized_weight + y * NPerBlock + x * 2),
+                              __half2half2(local_inputs[y]),
+                              *reinterpret_cast<half2*>(psum + batch_idx * NPerBlock + x * 2));
+                    }
+                    else
+                    {
+                      *reinterpret_cast<nv_bfloat162*>(psum + batch_idx * NPerBlock + x * 2)
+                          = __hfma2(*reinterpret_cast<nv_bfloat162*>(dequantized_weight + y * NPerBlock + x * 2),
+                              __bfloat162bfloat162(local_inputs[y]),
+                              *reinterpret_cast<nv_bfloat162*>(psum + batch_idx * NPerBlock + x * 2));   
+                    }
+                }
+            }
+        }
+        inputs_ptr += act_forward_step;
+        scale_ptr += scale_forward_step;
+        zeros_ptr += scale_forward_step;
+    }
+
+    warp_reduce<Num, WARP_SIZE>(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<T>(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<int64_t> 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<ctype*>(_in_feats.data_ptr());
+      auto kernel = reinterpret_cast<uint32_t*>(_kernel.data_ptr());
+      auto zeros = reinterpret_cast<ctype*>(_zeros.data_ptr());
+      auto scaling_factors = reinterpret_cast<ctype*>(_scaling_factors.data_ptr());
+      auto out_feats = reinterpret_cast<ctype*>(_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<<<num_blocks, num_threads>>>(
+      //     // 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<N_PER_BLOCK, 1, BLOCK_SIZE, 128><<<num_blocks, num_threads>>>(
+            in_feats, kernel, scaling_factors, zeros, out_feats, k, n
+          );
+          break;
+        case 2:
+          gemv_kernel<N_PER_BLOCK, 2, BLOCK_SIZE, 128><<<num_blocks, num_threads>>>(
+            in_feats, kernel, scaling_factors, zeros, out_feats, k, n
+          );
+          break;
+        case 3:
+          gemv_kernel<N_PER_BLOCK, 3, BLOCK_SIZE, 128><<<num_blocks, num_threads>>>(
+            in_feats, kernel, scaling_factors, zeros, out_feats, k, n
+          );
+          break;
+        case 4:
+          gemv_kernel<N_PER_BLOCK, 4, BLOCK_SIZE, 128><<<num_blocks, num_threads>>>(
+            in_feats, kernel, scaling_factors, zeros, out_feats, k, n
+          );
+          break;
+        case 5:
+          gemv_kernel<N_PER_BLOCK, 5, BLOCK_SIZE, 128><<<num_blocks, num_threads>>>(
+            in_feats, kernel, scaling_factors, zeros, out_feats, k, n
+          );
+          break;
+        case 6:
+          gemv_kernel<N_PER_BLOCK, 6, BLOCK_SIZE, 128><<<num_blocks, num_threads>>>(
+            in_feats, kernel, scaling_factors, zeros, out_feats, k, n
+          );
+          break;
+        case 7:
+          gemv_kernel<N_PER_BLOCK, 7, BLOCK_SIZE, 128><<<num_blocks, num_threads>>>(
+            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;
+}
+

llm-awq/awq/kernels/csrc/quantization_new/gemv/gemv_cuda.h

@@ -0,0 +1,12 @@
+#pragma once
+#include <torch/extension.h>
+
+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);

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 <ATen/cuda/CUDAContext.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+#include "fused_rope_with_pos.h"
+// #include <transformer_engine/fused_rope.h>
+
+// #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 <typename scalar_t>
+__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<float>(src[offset_src + (d2 / 2) * stride_d])
+              : static_cast<float>(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 <typename scalar_t>
+// __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 <typename scalar_t>
+__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<scalar_t>(
+      src, freqs, dst, offset_block, offset_block_dst, h, d, d2, stride_h,
+      stride_d, o_stride_h, o_stride_d);
+}
+
+// template <typename scalar_t>
+// __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<scalar_t>(src, freqs, dst, offset_block,
+//   offset_block_dst, h, d, d2, stride_h,
+//                             stride_d, o_stride_h, o_stride_d);
+// }
+
+template <typename scalar_t>
+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<scalar_t><<<blocks, threads, 0, stream>>>(
+      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 <typename scalar_t>
+// 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<scalar_t><<<blocks, threads, 0, stream>>>(
+//       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 <typename scalar_t>
+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<scalar_t>(
+      reinterpret_cast<const scalar_t *>(input.data_ptr()),
+      reinterpret_cast<const float *>(freqs.data_ptr()),
+      reinterpret_cast<scalar_t *>(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 <typename scalar_t>
+// 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<scalar_t>(reinterpret_cast<const scalar_t
+//       *>(output_grads.data_ptr()),
+//                                    reinterpret_cast<const float
+//                                    *>(freqs.data_ptr()),
+//                                    reinterpret_cast<scalar_t
+//                                    *>(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 <typename scalar_t>
+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<scalar_t>(
+      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 <typename scalar_t>
+// 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<scalar_t>(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<scalar_t>(
+            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.scalar_type()>(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<scalar_t>(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<float>(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;
+// }

llm-awq/awq/kernels/csrc/rope_new/fused_rope_with_pos.h

@@ -0,0 +1,5 @@
+#include <torch/extension.h>
+
+at::Tensor fused_rope_with_pos_forward_func(const at::Tensor &input,
+                                            const at::Tensor &freqs,
+                                            const bool transpose_output_memory);

llm-awq/awq/kernels/csrc/w8a8/act.cu

@@ -0,0 +1,141 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+#include <cuda_fp16.h>
+
+#include "dispatch_utils.h"
+#include "utils.cuh"
+#include "reduction_utils.cuh"
+
+namespace vllm {
+
+template <typename T> __device__ __forceinline__ T silu(const T &x) {
+  // x * sigmoid(x)
+  return (T)(((float)x) / (1.0f + expf((float)-x)));
+}
+
+template <typename T> __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 <typename T>
+__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 <typename scale_type, bool use_per_token_quant>
+__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<half, true><<<grid, block, 0, stream>>>(
+      out.data_ptr<int8_t>(), reinterpret_cast<half *>(input.data_ptr<at::Half>()), d, reinterpret_cast<half *>(scale_out.data_ptr<at::Half>()),reinterpret_cast<half *>(tmp.data_ptr<at::Half>()));
+}
+
+
+
+namespace vllm {
+  
+template<typename scalar_t>
+__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<int64_t> 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<scalar_t><<<grid, block, 0, stream>>>(
+        output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);
+  });
+  return output;
+}

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 <torch/extension.h>
+#include <cuda_fp16.h>
+// Inspired by vLLM-SmoothQuant: https://github.com/vllm-project/vllm/pull/1112.
+#include <torch/extension.h>
+
+
+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]
+);
+
+
+        
+

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 <torch/extension.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__))

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 <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+#include "dispatch_utils.h"
+#include "utils.cuh"
+#include "reduction_utils.cuh"
+
+
+namespace vllm {
+
+// from TRTLLM
+template <typename Tf, typename T>
+__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<Tf>(gamma[i]);
+    if (beta != nullptr)
+    {
+        ret = ret + cuda_cast<Tf>(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 <typename T, typename scale_type, bool USE_DIFF_OF_SQUARES = true>
+__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<T>::value;
+    using int8_packed_t = typename packed_as<int8_t, num_elems_T>::type;
+    using float_packed_t = typename packed_as<float, num_elems_T>::type;
+    using T_scalar = typename packed_as<T, 1>::type;
+
+    extern __shared__ __align__(sizeof(float)) char _shmem[];
+    T* shmem = reinterpret_cast<T*>(_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<float_packed_t>(val);
+        local_sum += cuda_sum<float>(val_f);
+        if (USE_DIFF_OF_SQUARES)
+        {
+            local_var_sum += cuda_sum<float>(val_f * val_f);
+        }
+    }
+    //Compute mean
+    if (USE_DIFF_OF_SQUARES)
+    {
+        float packed[2] = {local_sum, local_var_sum};
+        blockReduceSumV2<float, 2>(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<float_packed_t>(val); // - s_mean;
+            local_var_sum += cuda_sum<float>(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<float_packed_t>(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<float_packed_t>(use_shmem ? shmem[i] : input[index]);
+        const T val = cuda_cast<T>(compute_layernorm(val_f, s_mean, s_variance, gamma, beta, i));
+
+        if (with_per_token_scaling)
+        {
+            amax = cuda_max(cuda_max<T_scalar, T>(cuda_abs(val)), amax);
+            if (use_shmem)
+            {
+                shmem[i] = val;
+            }
+        }
+        else if (with_per_tensor_scaling)
+        {
+            reinterpret_cast<int8_packed_t*>(normed_output_quant)[index]
+                = cuda_cast<int8_packed_t>(cuda_cast<float_packed_t>(val) * scale_orig_quant);
+        }
+        else
+        {
+            normed_output[index] = val;
+        }
+    }
+
+    if (with_per_token_scaling)
+    {
+        float abs_max_f = blockAllReduceMax(cuda_cast<float>(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<float_packed_t>(use_shmem ? shmem[i] : input[index]);
+            if (!use_shmem)
+            {
+                val_f = compute_layernorm(val_f, s_mean, s_variance, gamma, beta, i);
+            }
+
+            reinterpret_cast<int8_packed_t*>(normed_output_quant)[index]
+                = cuda_cast<int8_packed_t>(val_f * cuda_cast<float_packed_t>(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<scalar_t>::Type;
+    if (use_per_token_quant) {
+      // per-token
+      vllm::generalLayerNorm<T, at::Half><<<grid, block, 0, stream>>>(
+        reinterpret_cast<T*>(input.data_ptr<scalar_t>()), 
+        reinterpret_cast<T*>(weight.data_ptr<scalar_t>()), 
+        reinterpret_cast<T*>(bias.data_ptr<scalar_t>()),
+        nullptr, epsilon, num_tokens, hidden_size, nullptr, scaling.data_ptr<at::Half>(),
+        out.data_ptr<int8_t>(), false
+      );
+      // input, gamma, beta, normed_output, eps, tokens, hidden_dim, per_tensor_scale, per_token_scale
+      // normed_output_quant, use_shmem
+        // out.data_ptr<int8_t>(), input.data_ptr<scalar_t>(),
+        // weight.data_ptr<scalar_t>(), epsilon, num_tokens, hidden_size);
+    } else {
+      // per-tensor
+      vllm::generalLayerNorm<T, at::Half><<<grid, block, 0, stream>>>(
+        reinterpret_cast<T*>(input.data_ptr<scalar_t>()), 
+        reinterpret_cast<T*>(weight.data_ptr<scalar_t>()), nullptr,
+        nullptr, epsilon, num_tokens, hidden_size, scaling.data_ptr<at::Half>(), nullptr,
+        out.data_ptr<int8_t>(), false
+      );
+    }
+  });
+}

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 <torch/extension.h>
+#include <cuda_fp16.h>
+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);
+

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 <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+
+#include "utils.cuh"
+#include <cuda_fp16.h>
+#include <cassert>
+#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<typename T>
+__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<typename T>
+__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 <typename T, typename scale_type, bool use_per_token_quant>
+__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<scalar_t, at::Half *, true><<<grid, block, 0, stream>>>(
+        input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
+        scale.data_ptr<at::Half>(), num_tokens, hidden_size);
+  });
+}
+

llm-awq/awq/kernels/csrc/w8a8/quantization.h

@@ -0,0 +1,4 @@
+#include <torch/extension.h>
+void invoke_quant(torch::Tensor &out,   // [..., hidden_size]
+                  torch::Tensor &input, // [..., hidden_size]
+                  torch::Tensor &scale);  // [num_tokens]

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 <assert.h>
+#include <stdint.h>
+#include <float.h>
+#include <type_traits>
+#include <torch/extension.h>
+
+template <typename T>
+struct FloatTypeConverter
+{
+    using Type = T;
+};
+
+template <>
+struct FloatTypeConverter<at::Half>
+{
+    using Type = half;
+};
+
+template <>
+struct FloatTypeConverter<at::BFloat16>
+{
+    using Type = __nv_bfloat16;
+};
+
+template <>
+struct FloatTypeConverter<float>
+{
+    using Type = float;
+};
+
+
+
+template<typename T> struct num_elems;
+template <>          struct num_elems<float>           { static constexpr int value = 1; };
+template <>          struct num_elems<float2>          { static constexpr int value = 2; };
+template <>          struct num_elems<float4>          { static constexpr int value = 4; };
+template <>          struct num_elems<half>            { static constexpr int value = 1; };
+template <>          struct num_elems<half2>           { 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<typename T, int num> struct packed_as;
+template<typename T>          struct packed_as<T, 1>              { using type = T; };
+template<>                    struct packed_as<half,  2>          { using type = half2; };
+template<>                    struct packed_as<float,  2>         { using type = float2; };
+template<>                    struct packed_as<int8_t, 2>         { using type = int16_t; };
+template<>                    struct packed_as<int32_t, 2>        { using type = int2; };
+template<>                    struct packed_as<half2, 1>          { using type = half; };
+template<>                    struct packed_as<float2, 1>         { 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<const int8_t&>(dst);
+}
+
+template<typename T>
+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<int8_t>(static_cast<short>(f_val.x));
+    int8[1] = static_cast<int8_t>(static_cast<short>(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<int8_t>(static_cast<short>(val.x));
+    int8[1] = static_cast<int8_t>(static_cast<short>(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 <typename T_OUT, typename T_IN>
+__device__ inline T_OUT cuda_cast(T_IN val)
+{
+    return val;
+}
+
+template <>
+__device__ inline float2 cuda_cast<float2, int2>(int2 val)
+{
+    return make_float2(val.x, val.y);
+}
+
+template <>
+__device__ inline float2 cuda_cast<float2, float>(float val)
+{
+    return make_float2(val, val);
+}
+
+template <>
+__device__ inline float2 cuda_cast<float2, half2>(half2 val)
+{
+    return __half22float2(val);
+}
+
+template <>
+__device__ inline half2 cuda_cast<half2, float2>(float2 val)
+{
+    return __float22half2_rn(val);
+}
+
+template <>
+__device__ inline half2 cuda_cast<half2, float>(float val)
+{
+    return __float2half2_rn(val);
+}
+
+template <>
+__device__ inline half2 cuda_cast<half2, half>(half val)
+{
+    return __half2half2(val);
+}
+
+template <>
+__device__ inline int8_t cuda_cast<int8_t, half>(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<int16_t, half2>(half2 val)
+{
+    union
+    {
+        int8_t int8[2];
+        int16_t int16;
+    };
+
+    int8[0] = cuda_cast<int8_t>(val.x);
+    int8[1] = cuda_cast<int8_t>(val.y);
+    return int16;
+}
+
+template <>
+__device__ inline int8_t cuda_cast<int8_t, float>(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<int16_t, float2>(float2 val)
+{
+    union
+    {
+        int8_t int8[2];
+        int16_t int16;
+    };
+
+    int8[0] = cuda_cast<int8_t>(val.x);
+    int8[1] = cuda_cast<int8_t>(val.y);
+    return int16;
+}
+
+template <>
+__device__ inline half2 cuda_cast<half2, int16_t>(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<float2, int16_t>(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<float>(val);
+}
+
+template <>
+__device__ inline __nv_bfloat16 cuda_cast(int8_t val)
+{
+    return static_cast<float>(val);
+}
+
+template <>
+__device__ inline int8_t cuda_cast(__nv_bfloat16 val)
+{
+    return static_cast<float>(val);
+}
+
+template <>
+__device__ inline float cuda_cast<float, __nv_bfloat16>(__nv_bfloat16 val)
+{
+    return __bfloat162float(val);
+}
+
+template <>
+__device__ inline float2 cuda_cast<float2, __nv_bfloat162>(__nv_bfloat162 val)
+{
+    return bf1622float2(val);
+}
+
+template <>
+__device__ inline half cuda_cast<half, __nv_bfloat16>(__nv_bfloat16 val)
+{
+    return __float2half(__bfloat162float(val));
+}
+
+template <>
+__device__ inline int16_t cuda_cast<int16_t, __nv_bfloat162>(__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 <typename To, typename Ti>
+__device__ inline To cuda_sum(Ti val)
+{
+    return cuda_cast<To>(val);
+};
+
+template <typename To>
+__device__ inline To cuda_sum(float2 val)
+{
+    return cuda_cast<To>(val.x + val.y);
+};
+
+// Unary maximum: compute the max of a vector type
+template <typename To, typename Ti>
+__device__ inline To cuda_max(Ti val)
+{
+    return cuda_cast<To>(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 <typename T>
+__device__ inline T cuda_max(T val1, T val2)
+{
+    return (val1 > val2) ? val1 : val2;
+}
+
+template <typename T>
+__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

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 <cuda_fp16.h>
+#include <cuda_pipeline_primitives.h>
+#include <torch/extension.h>
+
+#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<CTA_M, CTA_N, CTA_K, WARP_M, WARP_N, WARP_K, STAGES>;        \
+  cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, \
+                       kSmemByteSize);                                           \
+  kernel_func<<<num_blocks, threads_per_block, kSmemByteSize>>>(                 \
+      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<CTA_M, CTA_N, CTA_K, WARP_M, WARP_N, WARP_K, STAGES>;        \
+  cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, \
+                       kSmemByteSize);                                           \
+  kernel_func<<<num_blocks, threads_per_block, kSmemByteSize>>>(                 \
+      in_feats, kernel, wscales, ascales, out_feats, num_in_feats, num_out_channels,               \
+      num_in_channels);
+
+
+      
+template <int N>
+__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 <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int SHARED_K_ITERS,
+          int STAGES>
+__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 <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int SHARED_K_ITERS,
+          int STAGES>
+__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 <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int STAGES>
+__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 <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int STAGES>
+__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 <int CTA_M, int CTA_N, int CTA_K, int WARP_M, int WARP_N, int WARP_K,
+          int STAGES>
+__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<float*>(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<int *>(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<CTA_M, CTA_N, CTA_K, CTA_SIZE, 1, STAGES>(
+        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<CTA_M, CTA_N, CTA_K, CTA_SIZE, 1, STAGES>(
+        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<CTA_N,CTA_SIZE>(Bias,Bias_shared,cta_offset_n);
+
+  share_to_reg_one_stage_A<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES>(
+      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<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES>(
+      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<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES>(
+          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<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES>(
+          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<CTA_M, CTA_N, CTA_K, CTA_SIZE,
+                                    WARP_K / INTRIN_K, STAGES>(
+            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<CTA_M, CTA_N, CTA_K, CTA_SIZE,
+                                    WARP_K / INTRIN_K, STAGES>(
+            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<CTA_M, CTA_N, CTA_K, CTA_SIZE,
+                                    WARP_K / INTRIN_K, STAGES>(
+            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<CTA_M, CTA_N, CTA_K, CTA_SIZE,
+                                    WARP_K / INTRIN_K, STAGES>(
+            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<half2 *>(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<int8_t *>(_in_feats.data_ptr<int8_t>());
+  auto kernel = reinterpret_cast<int8_t *>(_kernel.data_ptr<int8_t>());
+  auto wscales = reinterpret_cast<half2 *>(_wscales.data_ptr());
+  auto ascales = reinterpret_cast<half *>(_ascales.data_ptr());
+  auto bias = reinterpret_cast<half *>(_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<half *>(_out_feats.data_ptr<at::Half>());
+
+  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 <int CTA_M, int CTA_N, int CTA_K, int WARP_M, int WARP_N, int WARP_K,
+          int STAGES>
+__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<int *>(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<CTA_M, CTA_N, CTA_K, CTA_SIZE, 1, STAGES>(
+        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<CTA_M, CTA_N, CTA_K, CTA_SIZE, 1, STAGES>(
+        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<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES>(
+      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<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES>(
+      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<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES>(
+          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<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES>(
+          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<CTA_M, CTA_N, CTA_K, CTA_SIZE,
+                                    WARP_K / INTRIN_K, STAGES>(
+            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<CTA_M, CTA_N, CTA_K, CTA_SIZE,
+                                    WARP_K / INTRIN_K, STAGES>(
+            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<CTA_M, CTA_N, CTA_K, CTA_SIZE,
+                                    WARP_K / INTRIN_K, STAGES>(
+            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<CTA_M, CTA_N, CTA_K, CTA_SIZE,
+                                    WARP_K / INTRIN_K, STAGES>(
+            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<half2 *>(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<int8_t *>(_in_feats.data_ptr<int8_t>());
+  auto kernel = reinterpret_cast<int8_t *>(_kernel.data_ptr<int8_t>());
+  auto wscales = reinterpret_cast<half2 *>(_wscales.data_ptr());
+  auto ascales = reinterpret_cast<half *>(_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<half *>(_out_feats.data_ptr<at::Half>());
+
+  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 ;
+}

llm-awq/awq/kernels/csrc/w8a8/w8a8_gemm_cuda.h

@@ -0,0 +1,4 @@
+#include <torch/extension.h>
+
+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);

llm-awq/awq/quantize/__init__.py

@@ -0,0 +1,2 @@
+from .w8a8_linear import *
+from .smooth import *

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()

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()

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"])

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,
+            )
+        )

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()

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

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)

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
+        )

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