Add files using upload-large-folder tool

.conda/aau_token

@@ -0,0 +1 @@
+YN_RyWTyaweE0R_BuNYxb-

.conda/aau_token_host

@@ -0,0 +1 @@
+zHxE_XAQ

LSAQ_CoreCode/lsaq_quant.py

@@ -0,0 +1,170 @@
+import os
+import torch
+import torch.nn as nn
+import numpy as np
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import tqdm
+import json
+import math
+import torch.nn.functional as F
+
+
+from datasets import load_dataset
+
+@torch.no_grad()
+def quantize_weight_per_channel_absmax(w, n_bits=8):
+    # w: (out_features, in_features)
+    scales = w.abs().max(dim=-1, keepdim=True)[0]
+    q_max = 2 ** (n_bits - 1) - 1
+    scales.clamp_(min=1e-5).div_(q_max)
+    w.div_(scales).round_().mul_(scales)
+    return w
+
+
+@torch.no_grad()
+def quantize_weight_per_tensor_absmax(w, n_bits=8):
+    # w: (out_features, in_features)
+    scales = w.abs().max()
+    q_max = 2 ** (n_bits - 1) - 1
+    scales.clamp_(min=1e-5).div_(q_max)
+    w.div_(scales).round_().mul_(scales)
+    return w
+
+class W8A16Linear(nn.Module):
+    def __init__(
+        self,
+        # bit_width,
+        in_features,
+        out_features,
+        bias=True,
+        quantize_output=False,
+    ):
+        super().__init__()
+        # self.bit_width = bit_width
+        self.in_features = in_features
+        self.out_features = out_features
+
+        self.register_buffer(
+            "weight",
+            torch.randn(
+                self.out_features,
+                self.in_features,
+                dtype=torch.float16,
+                requires_grad=False,
+            ),
+        )
+        if bias:
+            self.register_buffer(
+                "bias",
+                torch.zeros(
+                    (1, self.out_features), dtype=torch.float16, requires_grad=False
+                ),
+            )
+        else:
+            self.register_buffer("bias", None)
+
+    def to(self, *args, **kwargs):
+        super(W8A16Linear, self).to(*args, **kwargs)
+        self.weight = self.weight.to(*args, **kwargs)
+        if self.bias is not None:
+            self.bias = self.bias.to(*args, **kwargs)
+        return self
+
+    @torch.no_grad()
+    def forward(self, x):
+        y = torch.functional.F.linear(x, self.weight, self.bias)
+        return y
+
+    @staticmethod
+    def from_float(
+        bit, module, weight_quant="per_channel", quantize_output=False
+    ):
+        assert isinstance(module, torch.nn.Linear)
+        new_module = W8A16Linear(
+            # bit,
+            module.in_features,
+            module.out_features,
+            module.bias is not None,
+            quantize_output=quantize_output,
+        )
+        if weight_quant == "per_channel":
+            new_module.weight = quantize_weight_per_channel_absmax(module.weight, bit)
+        elif weight_quant == "per_tensor":
+            new_module.weight = quantize_weight_per_tensor_absmax(module.weight, bit)
+        else:
+            raise ValueError(f"Invalid weight_quant: {weight_quant}")
+        new_module.weight_quant_name = weight_quant
+        if module.bias is not None:
+            new_module.bias = module.bias
+        return new_module
+
+    def __repr__(self):
+        return f"W8A16Linear({self.in_features}, {self.out_features}, bias={self.bias is not None}, weight_quant={self.weight_quant_name})"
+
+def quantize_llama_like(
+    model, mlp_quant, self_attn_quant, low_bit, weight_quant="per_channel", quantize_bmm_input=False
+):
+    from transformers.models.llama.modeling_llama import (
+        LlamaAttention,
+        LlamaMLP,
+    )
+
+    for name, m in model.model.named_modules():
+        if isinstance(m, LlamaMLP):
+            if low_bit == 0:
+                continue
+            else:
+                if name in mlp_quant:
+                    bit = low_bit
+                    print(f'{name} {bit} bit quant ')
+                else:
+                    if low_bit == 4:
+                        bit = 8
+                        print(f'{name} {bit} bit quant ')
+                    elif low_bit == 8:
+                        continue
+
+            m.gate_proj = W8A16Linear.from_float(
+                bit, m.gate_proj, weight_quant=weight_quant
+            )
+            m.up_proj = W8A16Linear.from_float(
+                bit, m.up_proj, weight_quant=weight_quant
+            )
+            m.down_proj = W8A16Linear.from_float(
+                bit, m.down_proj, weight_quant=weight_quant
+            )
+        elif isinstance(m, LlamaAttention):
+            if low_bit == 0:
+                continue
+            else:
+                if name in self_attn_quant:
+                    bit = low_bit
+                else:
+                    if low_bit == 4:
+                        bit = 8
+                    elif low_bit == 8:
+                        continue
+
+            m.q_proj = W8A16Linear.from_float(
+                bit,  
+                m.q_proj,
+                weight_quant=weight_quant,
+                quantize_output=quantize_bmm_input,
+            )
+            m.k_proj = W8A16Linear.from_float(
+                bit, 
+                m.k_proj,
+                weight_quant=weight_quant,
+                quantize_output=quantize_bmm_input,
+            )
+            m.v_proj = W8A16Linear.from_float(
+                bit, 
+                m.v_proj,
+                weight_quant=weight_quant,
+                quantize_output=quantize_bmm_input,
+            )
+            m.o_proj = W8A16Linear.from_float(
+                bit, m.o_proj, weight_quant=weight_quant
+            )
+            
+    return model

LSAQ_CoreCode/main.ipynb

@@ -0,0 +1,321 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import GPUtil\n",
+    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
+    "import tqdm\n",
+    "from functools import partial"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Resource Detection"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "gpus = GPUtil.getGPUs()\n",
+    "free_memory = []\n",
+    "\n",
+    "for gpu in gpus:\n",
+    "    free_memory.append(gpu.memoryFree)\n",
+    "\n",
+    "memory_sort = sorted(range(len(free_memory)), key=lambda i: free_memory[i])\n",
+    "\n",
+    "gpu_id = memory_sort[-1]\n",
+    "gpu_memory = free_memory[memory_sort[-1]]\n",
+    "\n",
+    "print(f'gpu_id:{gpu_id}; gpu_memory:{gpu_memory}')\n",
+    "\n",
+    "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(gpu_id)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Model Selection"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_name = \"/data/LLMs/Llama-2-7b-hf\"\n",
+    "\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
+    "model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16, device_map=\"auto\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Layer Importance Detection"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def encode(tok, text, padding=True, truncation=True, max_length=None):\n",
+    "    # 将文本转换为输入 IDs\n",
+    "    input_ids = [tok.bos_id] + tok.encode(text)\n",
+    "\n",
+    "    # 生成注意力掩码\n",
+    "    attention_mask = [1] * len(input_ids)\n",
+    "\n",
+    "    # 如果进行了填充，则调整注意力掩码\n",
+    "    if padding:\n",
+    "        padding_length = max_length - len(input_ids)\n",
+    "        attention_mask = [0] * padding_length + attention_mask\n",
+    "        input_ids = [tok.eos_id] * padding_length + input_ids\n",
+    "\n",
+    "    encoded_input = {\n",
+    "        'input_ids': input_ids,\n",
+    "        'attention_mask': attention_mask\n",
+    "    }\n",
+    "    return encoded_input\n",
+    "\n",
+    "def batch_encode_plus(tok, texts, max_length=None, return_tensors=None):\n",
+    "    encoded_inputs = []\n",
+    "\n",
+    "    # 循环处理每个文本\n",
+    "    if max_length is None:\n",
+    "        max_length = -1\n",
+    "        for text in texts:\n",
+    "            # if isinstance(text, list):\n",
+    "            #     text = text[0]\n",
+    "            # print(text)\n",
+    "            len_ = len([tok.bos_id] + tok.encode(text))\n",
+    "            if len_ > max_length:\n",
+    "                max_length = len_\n",
+    "    for text in texts:\n",
+    "        # if isinstance(text, list):\n",
+    "        #     text = text[0]\n",
+    "        encoded_input = encode(tok, text, max_length = max_length)\n",
+    "        encoded_inputs.append(encoded_input)\n",
+    "\n",
+    "    # 合并结果\n",
+    "    batch_encoded = {\n",
+    "        'input_ids': [encoded_input['input_ids'] for encoded_input in encoded_inputs],\n",
+    "        'attention_mask': [encoded_input['attention_mask'] for encoded_input in encoded_inputs]\n",
+    "    }\n",
+    "\n",
+    "    batch_encoded = {key: torch.tensor(val) for key, val in batch_encoded.items()}\n",
+    "\n",
+    "    return batch_encoded"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tokenizer.bos_token = tokenizer.eos_token\n",
+    "tokenizer.bos_id = tokenizer.bos_token_id\n",
+    "tokenizer.eos_id = tokenizer.eos_token_id\n",
+    "importances = [0 for i in range(len(model.model.layers))]  # layer-wise importance scores"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import load_dataset\n",
+    "\n",
+    "dataset = load_dataset(\"wikitext\", \"wikitext-2-raw-v1\", split=\"test\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "MAX_SEQ_LEN = 1024\n",
+    "batch_size = 1\n",
+    "dataset_size = 200"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def jaccard_set(list1, list2):\n",
+    "    \"\"\"Define Jaccard Similarity function for two sets\"\"\"\n",
+    "    intersection = len(list(set(list1).intersection(list2)))\n",
+    "    union = (len(list1) + len(list2)) - intersection\n",
+    "    return float(intersection) / union"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "k = 20\n",
+    "\n",
+    "for i in tqdm.tqdm(range(0, dataset_size, batch_size), total = dataset_size / batch_size):\n",
+    "    \n",
+    "    prompts = dataset['text'][i:i + batch_size]\n",
+    "    max_seq_len = MAX_SEQ_LEN\n",
+    "    stride = 256\n",
+    "    max_gen_len = 0\n",
+    "\n",
+    "\n",
+    "    prompt_tokens = batch_encode_plus(\n",
+    "        tokenizer,\n",
+    "        prompts,\n",
+    "        return_tensors='pt'\n",
+    "    )\n",
+    "    input_ids = prompt_tokens['input_ids']\n",
+    "    attn_mask = prompt_tokens['attention_mask']\n",
+    "    max_prompt_len = max(len(t) for t in input_ids)\n",
+    "    all_jac_sim = [0 for i in range(len(model.model.layers))] \n",
+    "    E = model.get_input_embeddings().weight.detach()\n",
+    "    \n",
+    "    # authors use a sliding window of size 1024 with a shift of 256\n",
+    "    for start in range(0, max_prompt_len, stride):\n",
+    "        seq_ids = (attn_mask.sum(dim=-1) > start).nonzero().squeeze()\n",
+    "        seq_ids = seq_ids.unsqueeze(0) if seq_ids.dim() == 0 else seq_ids  # ensure 2d\n",
+    "        inputs = input_ids[seq_ids, start:start+max_seq_len]\n",
+    "        attn = attn_mask[seq_ids, start:start+max_seq_len]\n",
+    "\n",
+    "        if max_gen_len == 0:\n",
+    "            outputs = model(\n",
+    "                input_ids=inputs.to(\"cuda\"),\n",
+    "                attention_mask=attn.to(\"cuda\"),\n",
+    "                output_hidden_states=True,\n",
+    "            )\n",
+    "        else:\n",
+    "            outputs = model.generate(\n",
+    "                input_ids=inputs.to(\"cuda\"),\n",
+    "                attention_mask=attn.to(\"cuda\"),\n",
+    "                max_new_tokens=max_gen_len, \n",
+    "                output_hidden_states=True,\n",
+    "                return_dict_in_generate=True,\n",
+    "            )\n",
+    "\n",
+    "        hiddens = outputs.hidden_states\n",
+    "\n",
+    "        for i in range(len(hiddens) - 1):\n",
+    "            in_hidden = hiddens[i][:,-1,:]\n",
+    "            out_hidden = hiddens[i+1][:,-1,:]\n",
+    "\n",
+    "            in_projs = in_hidden @ E.T\n",
+    "            out_projs = out_hidden @ E.T\n",
+    "\n",
+    "            in_projs = in_projs.detach().cpu().numpy()\n",
+    "            ot_projs = out_projs.detach().cpu().numpy()\n",
+    "\n",
+    "            in_ind = np.argsort(-in_projs)\n",
+    "            ot_ind = np.argsort(-ot_projs)\n",
+    "\n",
+    "            in_topks = [tokenizer.decode(i) for i in in_ind[0][:k]]\n",
+    "            ot_topks = [tokenizer.decode(i) for i in ot_ind[0][:k]]\n",
+    "\n",
+    "            all_jac_sim[i] += jaccard_set(in_topks, ot_topks)\n",
+    "\n",
+    "      \n",
+    "    importances = [x + y for x, y in zip(importances, all_jac_sim)]\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import math\n",
+    "def normalize(lst, range_min=0, range_max=1):\n",
+    "    min_val = min(lst)\n",
+    "    max_val = max(lst)\n",
+    "    normalized = [(range_max - range_min) * (x - min_val) / (max_val - min_val) + range_min for x in lst]\n",
+    "    return normalized\n",
+    "\n",
+    "filtered_values = [0 if math.isinf(value) else value for value in importances] \n",
+    "normalized_lst = normalize(filtered_values)\n",
+    "\n",
+    "sorted_indices = sorted(range(len(normalized_lst)), key=lambda i: normalized_lst[i])\n",
+    "reversed_list = list(reversed(sorted_indices))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Quantize"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from lsaq_quant import quantize_llama_like\n",
+    "\n",
+    "num_of_layer2quant = 8\n",
+    "bit = 8\n",
+    "\n",
+    "layer_to_quant = reversed_list[0:num_of_layer2quant]\n",
+    "\n",
+    "mlp_quant = [f'layers.{item}.mlp' for item in layer_to_quant]\n",
+    "self_attn_quant = [f'layers.{item}.self_attn' for item in layer_to_quant]\n",
+    "\n",
+    "print(f'quanting ... ')\n",
+    "model_lsaq = quantize_llama_like(model, mlp_quant, self_attn_quant, bit)\n",
+    "print(f'quanted')"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "smoothquant",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.19"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

README.md

@@ -0,0 +1 @@
+# quantization

baselines/Llama-2-7b-hf_alpha_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,
+    4,
+    2,
+    2,
+    2,
+    2,
+    4
+]

baselines/Llama-2-7b-hf_alpha_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,
+    2,
+    2,
+    2,
+    4,
+    2,
+    4,
+    4,
+    4,
+    4,
+    4,
+    2,
+    4
+]

baselines/Llama-2-7b-hf_kurtosis_idx_10.json

@@ -0,0 +1,34 @@
+[
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    2,
+    2,
+    4,
+    2,
+    2,
+    2,
+    4,
+    2,
+    4,
+    2,
+    4,
+    2,
+    4,
+    2,
+    2,
+    4,
+    4,
+    4
+]

baselines/Llama-2-7b-hf_kurtosis_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,
+    2,
+    4,
+    4,
+    4,
+    4,
+    2,
+    4,
+    2,
+    4,
+    2,
+    2,
+    4,
+    4,
+    4
+]

baselines/Llama-2-7b-hf_z_idx_10.json

@@ -0,0 +1,34 @@
+[
+    4,
+    4,
+    4,
+    2,
+    2,
+    2,
+    4,
+    2,
+    4,
+    2,
+    2,
+    4,
+    2,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    2,
+    4,
+    4,
+    4,
+    2,
+    4,
+    2,
+    4,
+    4,
+    4,
+    4
+]

baselines/Llama-2-7b-hf_z_idx_5.json

@@ -0,0 +1,34 @@
+[
+    4,
+    4,
+    4,
+    4,
+    4,
+    2,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    2,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    4,
+    2,
+    4,
+    4,
+    4,
+    2,
+    4,
+    2,
+    4,
+    4,
+    4,
+    4
+]

eval.sh

@@ -0,0 +1,59 @@
+export HTTP_PROXY=http://sys-proxy-rd-relay.byted.org:8118
+export http_proxy=http://sys-proxy-rd-relay.byted.org:8118
+export https_proxy=http://sys-proxy-rd-relay.byted.org:8118
+export no_proxy="$no_proxy,.byteintl.net"
+export HF_ENDPOINT=https://hf-mirror.com
+
+
+cd quantization_metric/
+model_id=/mnt/bn/life-mllm/users/cxr/quantization/models/Qwen/Qwen2.5-14B
+model_name=$(basename "$model_id")
+cuda_id=4
+
+model=/mnt/bn/life-mllm/users/cxr/quantization/models/${model_name}-quantization-layer-mlp
+# output_dir=Alpha_values_mlp
+# bit_layers_dir=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/bit_layers
+# result_dir=/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results
+
+
+
+start=$(date +%s.%N)
+# rm -rf $model
+
+modes=("mlp")
+for mode in ${modes[@]}; do
+    for idx in {32..47}; do
+        echo $mode $idx
+        cd ../quantization_metric
+        python -u main.py --bit_layer_idx $idx  --save_dir ${model} --mode $mode --model_id $model_id --cuda_id $cuda_id
+        cd ../lm-evaluation-harness
+        bash run_scripts/eval_base_qwen2_5_14b.sh ${model} ${mode}_${idx} 
+        rm -rf ${model}
+        end=$(date +%s.%N)
+        runtime=$(awk "BEGIN {print $end - $start}")
+        echo "Execution time: $runtime seconds"
+
+    done
+done
+
+
+model=/mnt/bn/life-mllm/users/cxr/quantization/models/${model_name}-quantization-layer
+
+start=$(date +%s.%N)
+# rm -rf $model
+
+modes=("self_attn")
+for mode in ${modes[@]}; do
+    for idx in {32..47}; do
+        echo $mode $idx
+        cd ../quantization_metric
+        python -u main.py --bit_layer_idx $idx  --save_dir ${model} --mode $mode --model_id $model_id --cuda_id $cuda_id
+        cd ../lm-evaluation-harness
+        bash run_scripts/eval_base_qwen2_5_14b.sh ${model} ${mode}_${idx} 
+        rm -rf ${model}
+        end=$(date +%s.%N)
+        runtime=$(awk "BEGIN {print $end - $start}")
+        echo "Execution time: $runtime seconds"
+
+    done
+done

eval_coherence.sh

@@ -0,0 +1,29 @@
+export HTTP_PROXY=http://sys-proxy-rd-relay.byted.org:8118
+export http_proxy=http://sys-proxy-rd-relay.byted.org:8118
+export https_proxy=http://sys-proxy-rd-relay.byted.org:8118
+export no_proxy="$no_proxy,.byteintl.net"
+export HF_ENDPOINT=https://hf-mirror.com
+
+cd quantization_metric/
+model=../models/patch/Llama-2-7b-hf-quantization
+# output_dir=Alpha_values_mlp
+tasks=piqa,winogrande,arc_easy,arc_challenge,hellaswag,boolq
+# bit_layers_dir=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/bit_layers
+# result_dir=/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results
+
+
+start=$(date +%s.%N)
+# rm -rf $model
+
+file=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/metrics/coherence/coherence_self_attn_Llama-2-7b-hf.json
+echo "$file"
+cd ../quantization_metric
+configure_id=$(basename $file .json)
+python -u main_low.py --bit_layers $file  --save_dir ${model} --k 5 --reverse False
+cd ../lm-evaluation-harness
+bash run_scripts/eval.sh ${configure_id} ${model} ${tasks}
+rm -rf ${model}
+end=$(date +%s.%N)
+runtime=$(awk "BEGIN {print $end - $start}")
+echo "Execution time: $runtime seconds"
+

eval_fg.sh

@@ -0,0 +1,89 @@
+export HTTP_PROXY=http://sys-proxy-rd-relay.byted.org:8118
+export http_proxy=http://sys-proxy-rd-relay.byted.org:8118
+export https_proxy=http://sys-proxy-rd-relay.byted.org:8118
+export no_proxy="$no_proxy,.byteintl.net"
+export HF_ENDPOINT=https://hf-mirror.com
+
+
+cd quantization_metric/
+model_id=/mnt/bn/life-mllm/users/cxr/quantization/models/Qwen/Qwen2.5-7B
+model_name=$(basename "$model_id")
+
+model=/mnt/bn/life-mllm/users/cxr/quantization/models/${model_name}-quantization-fg
+# output_dir=Alpha_values_mlp
+# bit_layers_dir=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/bit_layers
+# result_dir=/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results
+
+start=$(date +%s.%N)
+# rm -rf $model
+cd ../quantization_metric
+# fg1
+# self_attn_layer_to_quant="4 1 2 8 23"
+# mlp_layer_to_quant="27 16 19 17 25"
+
+
+# save_fg=fg2
+# self_attn_layer_to_quant="23 22 25 24 26"
+# mlp_layer_to_quant="27 16 19 17 25"
+
+
+# save_fg=fg3
+# self_attn_layer_to_quant="23 22 25 24 26"
+# mlp_layer_to_quant="27 16 19"
+
+
+
+# save_fg=fg4
+# self_attn_layer_to_quant="23 22 25 24 26"
+# mlp_layer_to_quant="27"
+
+
+# save_fg=fg5
+# self_attn_layer_to_quant="27 16 19 17 25"
+# mlp_layer_to_quant="27 16 19 17 25"
+
+# save_fg=baseline_BI
+# self_attn_layer_to_quant="16 17 15 14 13"
+# mlp_layer_to_quant="16 17 15 14 13"
+
+
+save_fg=f6
+self_attn_layer_to_quant="4 1 2 8 23 22 25 5 24 7 26 6 20 12 19 17 21 11 10 9 18"
+mlp_layer_to_quant="27 16 19"
+
+python -u main_fg.py --cuda_id 6  --save_dir ${model} --model_id $model_id  --self_attn_layer_to_quant "${self_attn_layer_to_quant}" --mlp_layer_to_quant "${mlp_layer_to_quant}"
+cd ../lm-evaluation-harness
+bash run_scripts/eval_base_fg.sh ${model} ${save_fg}
+rm -rf ${model}
+end=$(date +%s.%N)
+runtime=$(awk "BEGIN {print $end - $start}")
+echo "Execution time: $runtime seconds"
+
+
+
+model_id=/mnt/bn/life-mllm/users/cxr/quantization/models/meta-llama/Llama-3.1-8B
+model_name=$(basename "$model_id")
+
+model=/mnt/bn/life-mllm/users/cxr/quantization/models/${model_name}-quantization-fg
+# output_dir=Alpha_values_mlp
+# bit_layers_dir=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/bit_layers
+# result_dir=/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results
+
+start=$(date +%s.%N)
+# rm -rf $model
+cd ../quantization_metric
+# save_fg=baseline_BI
+# self_attn_layer_to_quant="24 25 23 26 27"
+# mlp_layer_to_quant="24 25 23 26 27"
+save_fg=fg6
+self_attn_layer_to_quant="29 23 24 30 18 28 26 20 16 27 25 17 19 21"
+mlp_layer_to_quant="26 20 22"
+
+python -u main_fg.py --cuda_id 6  --save_dir ${model} --model_id $model_id  --self_attn_layer_to_quant "${self_attn_layer_to_quant}" --mlp_layer_to_quant "${mlp_layer_to_quant}"
+cd ../lm-evaluation-harness
+bash run_scripts/eval_base_fg.sh ${model} ${save_fg}
+rm -rf ${model}
+end=$(date +%s.%N)
+runtime=$(awk "BEGIN {print $end - $start}")
+echo "Execution time: $runtime seconds"
+

eval_hd.sh

@@ -0,0 +1,28 @@
+export HTTP_PROXY=http://sys-proxy-rd-relay.byted.org:8118
+export http_proxy=http://sys-proxy-rd-relay.byted.org:8118
+export https_proxy=http://sys-proxy-rd-relay.byted.org:8118
+export no_proxy="$no_proxy,.byteintl.net"
+export HF_ENDPOINT=https://hf-mirror.com
+
+cd quantization_metric/
+model=../models/patch/Llama-2-7b-hf-quantization
+# output_dir=Alpha_values_mlp
+tasks=piqa,winogrande,arc_easy,arc_challenge,hellaswag,boolq
+# bit_layers_dir=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/bit_layers
+# result_dir=/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results
+
+
+start=$(date +%s.%N)
+# rm -rf $model
+
+file=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/metrics/head_diversity/head_diversity_self_attn_Llama-2-7b-hf.json
+echo "$file"
+cd ../quantization_metric
+configure_id=$(basename $file .json)
+python -u main_low.py --bit_layers $file  --save_dir ${model} --k 5 --reverse False
+cd ../lm-evaluation-harness
+bash run_scripts/eval.sh ${configure_id} ${model} ${tasks}
+rm -rf ${model}
+end=$(date +%s.%N)
+runtime=$(awk "BEGIN {print $end - $start}")
+echo "Execution time: $runtime seconds"

eval_layer_llama.sh

@@ -0,0 +1,39 @@
+export HTTP_PROXY=http://sys-proxy-rd-relay.byted.org:8118
+export http_proxy=http://sys-proxy-rd-relay.byted.org:8118
+export https_proxy=http://sys-proxy-rd-relay.byted.org:8118
+export no_proxy="$no_proxy,.byteintl.net"
+export HF_ENDPOINT=https://hf-mirror.com
+
+
+cd quantization_metric/
+cuda_id=0
+model_id="/mnt/bn/life-mllm/users/cxr/quantization/models/meta-llama/Llama-3.1-8B"
+model_name=$(basename "$model_id")
+
+model=/mnt/bn/life-mllm/users/cxr/quantization/models/${model_name}-quantization-layer
+# output_dir=Alpha_values_mlp
+# bit_layers_dir=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/bit_layers
+# result_dir=/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results
+
+
+
+start=$(date +%s.%N)
+# rm -rf $model
+
+modes=("self_attn" "mlp")
+for mode in ${modes[@]}; do
+    for idx in {-1..31}; do
+        echo $mode $idx
+        cd ../quantization_metric
+        python -u main.py --bit_layer_idx $idx  --save_dir ${model} --mode $mode --model_id $model_id --cuda_id $cuda_id
+        cd ../lm-evaluation-harness
+        bash run_scripts/eval_base_qwen2_5_7b.sh ${model} ${mode}_${idx} $cuda_id
+        rm -rf ${model}
+        end=$(date +%s.%N)
+        runtime=$(awk "BEGIN {print $end - $start}")
+        echo "Execution time: $runtime seconds"
+
+    done
+done
+
+

eval_layer_qwen.sh

@@ -0,0 +1,39 @@
+export HTTP_PROXY=http://sys-proxy-rd-relay.byted.org:8118
+export http_proxy=http://sys-proxy-rd-relay.byted.org:8118
+export https_proxy=http://sys-proxy-rd-relay.byted.org:8118
+export no_proxy="$no_proxy,.byteintl.net"
+export HF_ENDPOINT=https://hf-mirror.com
+
+
+cd quantization_metric/
+cuda_id=1
+model_id=/mnt/bn/life-mllm/users/cxr/quantization/models/Qwen/Qwen2.5-7B
+model_name=$(basename "$model_id")
+
+model=/mnt/bn/life-mllm/users/cxr/quantization/models/${model_name}-quantization-layer
+# output_dir=Alpha_values_mlp
+# bit_layers_dir=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/bit_layers
+# result_dir=/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results
+
+
+
+start=$(date +%s.%N)
+# rm -rf $model
+
+modes=("self_attn" "mlp")
+for mode in ${modes[@]}; do
+    for idx in {-1..27}; do
+        echo $mode $idx
+        cd ../quantization_metric
+        python -u main.py --bit_layer_idx $idx  --save_dir ${model} --mode $mode --model_id $model_id --cuda_id $cuda_id
+        cd ../lm-evaluation-harness
+        bash run_scripts/eval_base_qwen2_5_7b.sh ${model} ${mode}_${idx} $cuda_id
+        rm -rf ${model}
+        end=$(date +%s.%N)
+        runtime=$(awk "BEGIN {print $end - $start}")
+        echo "Execution time: $runtime seconds"
+
+    done
+done
+
+

eval_zd.sh

@@ -0,0 +1,29 @@
+export HTTP_PROXY=http://sys-proxy-rd-relay.byted.org:8118
+export http_proxy=http://sys-proxy-rd-relay.byted.org:8118
+export https_proxy=http://sys-proxy-rd-relay.byted.org:8118
+export no_proxy="$no_proxy,.byteintl.net"
+export HF_ENDPOINT=https://hf-mirror.com
+
+cd quantization_metric/
+model=../models/patch/Llama-2-7b-hf-quantization-zd
+# output_dir=Alpha_values_mlp
+tasks=piqa,winogrande,arc_easy,arc_challenge,hellaswag,boolq
+# bit_layers_dir=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/bit_layers
+# result_dir=/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results
+
+
+start=$(date +%s.%N)
+# rm -rf $model
+
+# file=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/metrics/ZD/ZD_mlp_Llama-2-7b-hf.json
+file=/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/metrics/ZD/ZD_self_attn_Llama-2-7b-hf.json
+echo "$file"
+cd ../quantization_metric
+configure_id=$(basename $file .json)
+python -u main_low.py --bit_layers $file  --save_dir ${model} --k 5
+cd ../lm-evaluation-harness
+bash run_scripts/eval.sh ${configure_id} ${model} ${tasks}
+rm -rf ${model}
+end=$(date +%s.%N)
+runtime=$(awk "BEGIN {print $end - $start}")
+echo "Execution time: $runtime seconds"

inference.py

@@ -0,0 +1,19 @@
+
+a= {"results": {
+    "arc_easy": {
+      "alias": "arc_easy",
+      "acc,none": 0.6902356902356902,
+      "acc_stderr,none": 0.00948817285190372,
+      "acc_norm,none": 0.6422558922558923,
+      "acc_norm_stderr,none": 0.00983577275734336
+    },
+    "arc_easy": {
+      "alias": "arc_easy",
+      "acc,none": 0.6902356902356902,
+      "acc_stderr,none": 0.00948817285190372,
+      "acc_norm,none": 0.6422558922558923,
+      "acc_norm_stderr,none": 0.00983577275734336
+    }
+  }
+}
+print(len(a['results']))

layerwise-awq.py

@@ -0,0 +1,322 @@
+# -*- encoding:utf-8 -*-
+@torch.no_grad()
+def run_awq(
+    model,
+    enc,
+    w_bit,
+    q_config,
+    n_samples=512,
+    seqlen=512,
+    auto_scale=True,
+    mse_range=True,
+    calib_data="pileval",  # data for calibration
+    skip_first: int = 0,  # number of initial layers to keep in full precision
+    first_n: int = 0,  # number of initial layers to apply first quant
+    w_bit_first: int | None = None,
+    w_bit_rest: int | None = None,
+    # --- mixed-precision strategy --------------------------------------------------
+    strategy: str = "layer",  # "layer" (default): original solve layer-by-layer; "auto": structured mixed-precision
+    m_auto: int | None = None,  # number of high-bit layers when strategy == "auto"; defaults to 25% of L
+    hi_bit: int = 4,
+    lo_bit: int = 2,
+    alpha: float = 1 / 3,
+    beta: float = 1 / 3,
+    gamma: float = 1 / 3,
+    k_energy: int = 32,
+    metrics_csv: str | None = None,  # optional explicit path to metrics CSV (delta_ppl,erank_diff,topk_energy_diff)
+):
+    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": [],
+    }
+
+    # ---------------------------------------------------------------------------
+    # Determine per-layer bit-widths according to the requested *strategy*
+    # ---------------------------------------------------------------------------
+
+    if strategy.lower() == "auto":
+        # -------------------------------------------------------------------
+        # Use qpRANK pre-computed diagnostics to decide per-layer precision.
+        # Users may place the JSON files (drop_layer_ppl.json, diff_erank_values.json)
+        # under the project root (default path) or supply env QPRANK_METRICS_DIR.
+        # -------------------------------------------------------------------
+
+        import json, os, math, csv
+
+        def _load_metrics_from_csv(csv_path: str):
+            """Return delta_ppl, erank_diff, topk_energy_diff lists from a csv file."""
+            delta_ppl, erank, topk = [], [], []
+            with open(csv_path, "r", encoding="utf-8") as f:
+                reader = csv.DictReader(f)
+                for row in reader:
+                    delta_ppl.append(float(row.get("delta_ppl", 0)))
+                    erank.append(abs(float(row.get("erank_diff", 0))))
+                    topk_val = row.get("topk_energy_diff")
+                    if topk_val is not None and topk_val != "":
+                        topk.append(float(topk_val))
+            # Ensure all same length
+            assert len(delta_ppl) == len(erank), "CSV length mismatch"
+            if len(topk) != len(delta_ppl):
+                topk = [0.0] * len(delta_ppl)
+            return delta_ppl, erank, topk
+
+        delta_ppl: List[float]
+        delta_r: List[float]
+        delta_e: List[float]
+
+        # Priority 1: explicit CSV path
+        if metrics_csv is not None and os.path.isfile(metrics_csv):
+            delta_ppl, delta_r, delta_e = _load_metrics_from_csv(metrics_csv)
+        else:
+            # Priority 2: auto-detect inside QPRANK directory structure
+            base_dir = os.getenv("QPRANK_METRICS_DIR", os.path.expanduser("~/qpRANK/src"))
+
+            # Derive a crude model identifier from config
+            cfg_name = getattr(model, "config", None)
+            model_id = (
+                getattr(cfg_name, "_name_or_path", "model").replace("/", "_")
+                if cfg_name is not None
+                else "model"
+            )
+
+            # Traverse to find a metrics_long.csv matching pattern
+            candidate_csv = None
+            for root, dirs, files in os.walk(base_dir):
+                if "metrics_long.csv" in files and model_id in root:
+                    candidate_csv = os.path.join(root, "metrics_long.csv")
+                    break
+
+            if candidate_csv and os.path.isfile(candidate_csv):
+                delta_ppl, delta_r, delta_e = _load_metrics_from_csv(candidate_csv)
+            else:
+                # Fallback to old JSON files (legacy)
+                metrics_dir = os.getenv("QPRANK_METRICS_DIR", os.path.expanduser("~/qpRANK"))
+                ppl_path = os.path.join(metrics_dir, "drop_layer_ppl.json")
+                erank_path = os.path.join(metrics_dir, "diff_erank_values.json")
+
+                if not (os.path.isfile(ppl_path) and os.path.isfile(erank_path)):
+                    raise FileNotFoundError(
+                        "Cannot locate per-layer metric files for auto strategy. Provide metrics_csv path or set QPRANK_METRICS_DIR appropriately."
+                    )
+
+                delta_ppl = json.load(open(ppl_path, "r"))["delta_ppl"]
+                erank_json = json.load(open(erank_path, "r"))
+
+                keys = [k for k in ("q", "k", "v") if k in erank_json]
+                delta_r = [
+                    sum(erank_json[k][i] for k in keys) / len(keys)
+                    for i in range(len(delta_ppl))
+                ]
+
+                delta_e = erank_json.get("topk_energy_diff", [0.0] * len(delta_ppl))
+        #! layer 的数量
+        L_total = len(delta_ppl)
+
+        # Normalise
+        def _norm(arr):
+            m = max(arr) if max(arr) > 0 else 1.0
+            return [x / m for x in arr]
+
+        ppl_hat = _norm(delta_ppl)
+        r_hat = _norm(delta_r)
+        e_hat = _norm(delta_e)
+
+        scores = [
+            alpha * ppl_hat[i] + beta * r_hat[i] + gamma * e_hat[i]
+            for i in range(L_total)
+        ]
+
+        #! 1/4 的 layer 
+        if m_auto is None:
+            m_auto = max(1, L_total // 4)
+
+        idx_sorted = sorted(range(L_total), key=lambda i: scores[i], reverse=True)
+        #! 前 1/4 的 layer 用 high bit, 其他的用 low bit
+        hi_set = set(idx_sorted[:m_auto])
+        
+        #! 每个 layer 的 bit 数量的分配
+        #! 我们也是在这边修改成得到我们的 layer 分配就好了
+        bits_per_layer = [hi_bit if i in hi_set else lo_bit for i in range(L_total)]
+
+        # ---- verbose print & log ----
+        try:
+            import logging
+            _logger = logging.getLogger(__name__)
+        except ImportError:
+            _logger = None
+
+        print("[AUTO] Per-layer bit-width allocation (index:bit):")
+        mapping_str = ", ".join(f"{idx}:{bits_per_layer[idx]}b" for idx in range(L_total))
+        print(mapping_str)
+
+        if _logger is not None:
+            _logger.info("AUTO bit-width allocation: " + mapping_str)
+
+        print(f"[AUTO] Layers @ {hi_bit}-bit: {sorted(list(hi_set))}")
+        print(f"[AUTO] Layers @ {lo_bit}-bit: {sorted([i for i in range(L_total) if i not in hi_set])}")
+
+        if _logger is not None:
+            _logger.info(f"Layers_{hi_bit}bit: {sorted(list(hi_set))}")
+            _logger.info(f"Layers_{lo_bit}bit: {[i for i in range(L_total) if i not in hi_set]}")
+
+    else:
+        # Fallback to original scheme (uniform or head/tail mixed precision).
+        bits_per_layer = None  # will be decided on the fly as before
+
+    # solve layer by layer
+    for i in tqdm.tqdm(range(len(layers)), desc="Running AWQ..."):
+        # print(f"Layer {i} of {len(layers)-1}")
+        layer = layers[i]
+
+        # Flag: whether to apply quantization to this layer
+        #! 他们也指定了超参数从第几层开始量化
+        quantize_this = i >= skip_first
+
+        # Determine bit-width for this layer
+        if strategy.lower() == "auto" and bits_per_layer is not None:
+            current_w_bit = bits_per_layer[i]
+            if i == 0:
+                # show a brief summary once for user awareness
+                print(
+                    f"[AUTO] Using structured mixed-precision: {sum(b == hi_bit for b in bits_per_layer)} layers @ {hi_bit}-bit, {sum(b == lo_bit for b in bits_per_layer)} layers @ {lo_bit}-bit."
+                )
+        else:
+            # original rule-based selection
+            if i < first_n:
+                current_w_bit = w_bit_first if w_bit_first is not None else w_bit
+                print(
+                    f"Layer {i} is quantizing with {current_w_bit} bits. (when this sentence isnt printed, it is quantizing with {w_bit_rest} bits)"
+                )
+            else:
+                current_w_bit = w_bit_rest if w_bit_rest is not None else w_bit
+
+        
+        #! 从这边往后就和原来的代码一样
+        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=current_w_bit, #! 改成 current_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=current_w_bit, #! 改成 current_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

llm-awq/README.md

@@ -0,0 +1,292 @@
+# AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration 
+[[Paper](https://arxiv.org/abs/2306.00978)][[Website](https://hanlab.mit.edu/projects/awq)]
+
+**Efficient and accurate** low-bit weight quantization (INT3/4) for LLMs, supporting **instruction-tuned** models and **multi-modal** LMs.
+
+![overview](figures/overview.png)
+
+The current release supports: 
+
+- AWQ search for accurate quantization. 
+- Pre-computed AWQ model zoo for LLMs (Llama-1/2/3, OPT, CodeLlama, StarCoder, Vicuna, VILA, LLaVA; load to generate quantized weights).
+- Memory-efficient 4-bit Linear in PyTorch.
+- Efficient CUDA kernel implementation for fast inference (support context and decoding stage).
+- Examples on 4-bit inference of an instruction-tuned model (Vicuna) and **multi-modal LM** (VILA).
+- Chunk prefilling for faster prefilling in multi-round Q&A setting.
+- State-of-the-art prefilling speed of LLMs/VLMs on edge devices: [TinyChat 2.0](./tinychat).
+
+**Thanks to AWQ, TinyChat can deliver more efficient responses with LLM/VLM chatbots through 4-bit inference.**
+
+* TinyChat with LLaMA-3-8b on RTX 4090 (2.7x faster than FP16):
+
+![TinyChat with LLaMA-3-8b on RTX 4090: W4A16 is 2.7x faster than FP16](./tinychat/figures/4090_example_new.gif)
+
+* TinyChat with LLaMA-3-8b on Jetson Orin (2.9x faster than FP16):
+
+![TinyChat with LLaMA-3-8b on Jetson Orin: W4A16 is 2.9x faster than FP16](./tinychat/figures/orin_example_new.gif)
+
+
+**TinyChat also supports inference with vision language models (e.g., VILA, LLaVA). In the following examples, W4A16 quantized models from VILA family are launched with TinyChat.**
+
+* TinyChat with NVILA-8B on RTX 4090 (single-image inputs):
+
+![TinyChat with NVILA on 4090 single image](./tinychat/figures/4090_nvila_single.gif)
+
+* TinyChat with NVILA-8B on RTX 4090 (multi-image inputs):
+
+![TinyChat with NVILA on 4090 multiple images](./tinychat/figures/4090_nvila_multi.gif)
+
+<!-- Check out [TinyChat](tinychat), which delievers **30 tokens/second** inference performance (**3.2x faster** than FP16) for the **Llama2** chatbot on the resource-constrained NVIDIA Jetson Orin!  -->
+
+* TinyChat with video reasoning:
+
+https://github.com/user-attachments/assets/b68a7a0d-5175-4030-985b-5ae0ae94f874
+
+**Prompt:** What might be the next step according to the video?
+
+**Answer:** The next step in the video could be to place the shaped dough onto a baking sheet and let it rise before baking.
+
+**Online demo:** https://vila.hanlab.ai
+
+Check out [TinyChat](tinychat), which offers a turn-key solution for **on-device inference** of LLMs and VLMs on **resource-constrained edge platforms**. With TinyChat, it is now possible to efficiently run **large** models on **small** and **low-power** devices even without Internet connection!
+
+
+## News
+- [2025/04] 🔥 AWQ now supports DeepSeek-R1-Distilled models. Try our example [here](https://github.com/mit-han-lab/llm-awq/blob/main/scripts/DeepSeek_R1_Distill_example.sh)!
+- [2025/02] AWQ now supports BF16 precision. See example [here](https://github.com/mit-han-lab/llm-awq/blob/main/scripts/qwen_example.sh).
+- [2024/10] 🔥⚡ Explore advancements in [TinyChat 2.0](./tinychat), the latest version with significant advancements in prefilling speed of Edge LLMs and VLMs, **1.5-1.7x** faster than the previous version of TinyChat. Please refer to the [README](./tinychat/README.md) and [blog](https://hanlab.mit.edu/blog/tinychat20) for more details. 
+- [2024/05] 🏆 AWQ receives the **Best Paper Award** at **MLSys 2024**. 🎉 
+- [2024/05] 🔥 The **VILA-1.5** model family which features **video understanding** is now supported in AWQ and TinyChat. Check out out online demo powered by TinyChat [here](https://vila.hanlab.ai). Example is [here](scripts/vila15_example.sh).
+- [2024/05] 🔥 [AMD](https://community.amd.com/t5/ai/reduce-memory-footprint-and-improve-performance-running-llms-on/ba-p/686157) adopts AWQ to improve LLM serving efficiency.
+- [2024/04] 🔥 We released AWQ and TinyChat support for The **Llama-3** model family! Check out our example [here](scripts/llama3_example.sh).
+- [2024/02] 🔥 AWQ has been accepted to **MLSys 2024**!
+- [2024/02] 🔥 We supported [VILA Vision Languague Models](https://arxiv.org/abs/2312.07533) in AWQ & TinyChat! Check our latest demos with multi-image inputs!
+- [2024/02] 🔥 We released new version of quantized GEMM/GEMV kernels in [**TinyChat**](tinychat), leading to **38 tokens/second** inference speed on NVIDIA Jetson Orin!
+- [2024/01] 🔥 AWQ has been integrated by [Google Vertex AI](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama-2-quantized)!
+- [2023/11] 🔥 AWQ has been integrated by [Amazon Sagemaker Containers](https://aws.amazon.com/blogs/machine-learning/boost-inference-performance-for-llms-with-new-amazon-sagemaker-containers/)!
+- [2023/11] 🔥 We added AWQ support and pre-computed search results for CodeLlama, StarCoder, StableCode models. Checkout our model zoo [here](https://huggingface.co/datasets/mit-han-lab/awq-model-zoo)!
+- [2023/11] 🔥 AWQ is now integrated natively in Hugging Face transformers through `from_pretrained`. You can either load quantized models from the Hub or your own HF quantized models.
+- [2023/10] AWQ is integrated into NVIDIA [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/)
+- [2023/09] AWQ is integrated into [Intel Neural Compressor](https://github.com/intel/neural-compressor), [FastChat](https://github.com/lm-sys/FastChat/blob/main/docs/awq.md), [vLLM](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/awq.py), [HuggingFace TGI](https://github.com/huggingface/text-generation-inference/pull/1054), and [LMDeploy](https://github.com/InternLM/lmdeploy). 
+- [2023/09] ⚡ Check out our latest [**TinyChat**](tinychat), which is ~2x faster than the first release on Orin!
+- [2023/09] ⚡ Check out [**AutoAWQ**](https://github.com/casper-hansen/AutoAWQ), a third-party implementation to make AWQ easier to expand to new models, improve inference speed, and integrate into Huggingface.
+- [2023/07] 🔥 We released **TinyChat**, an efficient and lightweight chatbot interface based on AWQ. TinyChat enables efficient LLM inference on both cloud and edge GPUs. Llama-2-chat models are supported! Check out our implementation [here](tinychat).
+- [2023/07] 🔥 We added AWQ support and pre-computed search results for Llama-2 models (7B & 13B). Checkout our model zoo [here](https://huggingface.co/datasets/mit-han-lab/awq-model-zoo)!
+- [2023/07] We extended the support for more LLM models including MPT, Falcon, and BLOOM. 
+
+## Contents
+
+- [AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration](#awq-activation-aware-weight-quantization-for-llm-compression-and-acceleration)
+  - [News](#news)
+  - [Contents](#contents)
+  - [Helpful Links](#helpful-links)
+  - [Install](#install)
+  - [AWQ Model Zoo](#awq-model-zoo)
+  - [Examples](#examples)
+  - [Usage](#usage)
+  - [Results on Visual Language Models](#results-on-visual-language-models)
+  - [Reference](#reference)
+  - [Related Projects](#related-projects)
+
+## Helpful Links
+
+- [VILA online demo](vila.hanlab.ai): Visual Language Models efficiently supported by AWQ & TinyChat.
+- [LLM on the Edge](https://github.com/mit-han-lab/llm-awq/tree/nv_laptop?tab=readme-ov-file#install): AWQ and TinyChat support edge GPUs such as NVIDIA Jetson Orin.
+- [VLMs on Laptop](https://github.com/mit-han-lab/llm-awq/tree/nv_laptop?tab=readme-ov-file#run-vila-on-laptop): Follow the instructions to deploy VLMs on NVIDIA Laptops with TinyChat.
+- [Gradio Server](https://github.com/mit-han-lab/llm-awq/tree/nv_laptop/tinychat/serve#gradio-demo-vila-with-tinychat): Try to build your own VLM online demo with AWQ and TinyChat!
+- [QServe](https://github.com/mit-han-lab/qserve): 🔥 **[New]** Efficient and accurate serving system for large-scale LLM inference.
+  
+## Install
+
+1. Clone this repository and navigate to AWQ folder
+```
+git clone https://github.com/mit-han-lab/llm-awq
+cd llm-awq
+```
+
+2. Install Package
+```
+conda create -n awq python=3.10 -y
+conda activate awq
+pip install --upgrade pip  # enable PEP 660 support
+pip install -e .
+```
+
+* For **edge devices** like Orin, before running the commands above, please:
+
+    1. Modify [pyproject.toml](pyproject.toml) by commenting out [this line](https://github.com/mit-han-lab/llm-awq/blob/3fce69061682fdd528824e5da3d03a8a8b545f2a/pyproject.toml#L17).
+    2. Manually install precompiled PyTorch binaries (>=2.0.0) from [NVIDIA](https://forums.developer.nvidia.com/t/pytorch-for-jetson/72048). You also need to install torchvision from this website when running NVILA.
+    3. Set the appropriate Python version for conda environment (e.g., `conda create -n awq python=3.8 -y` for JetPack 5).
+  
+3. Install efficient W4A16 (4-bit weight, 16-bit activation) CUDA kernel and optimized FP16 kernels (e.g. layernorm, positional encodings).
+```
+cd awq/kernels
+python setup.py install
+```
+
+4. Install Flash Attention
+```
+pip install flash-attn --no-build-isolation
+```
+
+We recommend starting an interactive python CLI interface and run `import flash_attn` to check whether FlashAttention-2 is installed successfully. If not, we recommend downloading pre-built wheels from [here](https://github.com/Dao-AILab/flash-attention/releases/tag/v2.5.8). Please notice:
+
+- PyTorch version needs to exactly match with the version specified in the `.whl` name;
+- Check out both `cxx11abiTRUE` and `cxx11abiFALSE` wheels if one of them does not work;
+- It's recommended to match CUDA version specified in the `.whl` filename, but minor mismatches (e.g. 12.1 vs 12.2, or even 11.8 vs 12.2) usually do not matter.
+
+
+5. [Optional] In order to run AWQ and TinyChat with NVILA model family, please install VILA:
+
+```bash
+git clone https://github.com/NVlabs/VILA.git
+cd VILA
+pip install -e .
+```
+
+## AWQ Model Zoo
+
+We provide pre-computed AWQ search results for multiple model families, including LLaMA, OPT, Vicuna, and LLaVA. To get the pre-computed AWQ search results, run:
+
+```bash
+# git lfs install  # install git lfs if not already
+git clone https://huggingface.co/datasets/mit-han-lab/awq-model-zoo awq_cache
+```
+
+The detailed support list:
+
+| Models | Sizes                       | INT4-g128 | INT3-g128 |
+| ------ | --------------------------- | --------- | --------- |
+| [DeepSeek-R1-Distill](/scripts/DeepSeek_R1_Distill_example.sh) | 1.5B/7B/8B |  ✅         |         |
+| [Qwen-2.5](/scripts/qwen_example.sh) | 7B/72B            | ✅         |          |
+| [NVILA](/scripts/nvila_example.sh) | 3B/8B |  ✅         |         |
+| [VILA-1.5](/scripts/vila15_example.sh)  | 3B/8B/13B/40B  | ✅         | ✅        |
+| [Llama3](/scripts/llama_example.sh)  | 8B/70B  | ✅         | ✅        |
+| [VILA](/scripts/vila_example.sh)    | 7B/13B                     | ✅         |           |
+| [Llama2](/scripts/llama_example.sh)  | 7B/13B/70B  | ✅         | ✅        |
+| [LLaMA](/scripts/llama2_example.sh)  | 7B/13B/30B/65B              | ✅         | ✅        |
+| [OPT](/scripts/opt_example.sh)    | 125m/1.3B/2.7B/6.7B/13B/30B | ✅         | ✅        |
+| [CodeLlama](/scripts/codellama_example.sh) | 7B/13B/34B               | ✅         | ✅        |
+| [StarCoder](/scripts/starcoder_example.sh) | 15.5B                    | ✅         | ✅        |
+| [Vicuna-v1.1](/scripts/vicuna_example.sh) | 7B/13B                 | ✅         |           |
+| [LLaVA-v0](/scripts/llava_example.sh) | 13B                       | ✅         |           |
+
+Note: We only list models that we have prepare the [AWQ searching results](https://huggingface.co/datasets/mit-han-lab/awq-model-zoo/tree/main) in the table above. AWQ also supports models such as LLaVA-v1.5 7B, and you may need to run the [AWQ search](#usage) on your own to quantize these models. For our latest VLM NVILA, quantized weights are available [here](https://huggingface.co/Efficient-Large-Model/NVILA-AWQ).
+
+## Examples
+
+AWQ can be easily applied to various LMs thanks to its good generalization, including instruction-tuned models and multi-modal LMs. It provides an easy-to-use tool to reduce the serving cost of LLMs.
+
+Here we provide two examples of AWQ application: Vicuna-7B (chatbot) and LLaVA-13B (visual reasoning) under `./examples` directory. AWQ can easily reduce the GPU memory of model serving and speed up token generation. It provides accurate quantization, providing reasoning outputs. You should be able to observe **memory savings** when running the models with 4-bit weights. 
+
+Note that we perform AWQ using only textual calibration data, depsite we are running on multi-modal input. Please refer to `./examples` for details.
+
+![overview](figures/example_vis.jpg)
+
+## Usage
+
+We provide several sample script to run AWQ (please refer to `./scripts`). We use Llama3-8B as an example.
+
+1. Perform AWQ search and save search results (we already did it for you):
+```bash
+python -m awq.entry --model_path /PATH/TO/LLAMA3/llama3-8b \
+    --w_bit 4 --q_group_size 128 \
+    --run_awq --dump_awq awq_cache/llama3-8b-w4-g128.pt
+```
+
+2. Evaluate the AWQ quantized model on WikiText-2 (simulated pseudo quantization)
+```bash
+python -m awq.entry --model_path /PATH/TO/LLAMA3/llama3-8b \
+    --tasks wikitext \
+    --w_bit 4 --q_group_size 128 \
+    --load_awq awq_cache/llama3-8b-w4-g128.pt \
+    --q_backend fake
+```
+
+3. Generate real quantized weights (INT4)
+```bash
+mkdir quant_cache
+python -m awq.entry --model_path /PATH/TO/LLAMA3/llama3-8b \
+    --w_bit 4 --q_group_size 128 \
+    --load_awq awq_cache/llama3-8b-w4-g128.pt \
+    --q_backend real --dump_quant quant_cache/llama3-8b-w4-g128-awq.pt
+```
+
+4. Load and evaluate the real quantized model (now you can see smaller gpu memory usage)
+```bash
+python -m awq.entry --model_path /PATH/TO/LLAMA3/llama3-8b \
+    --tasks wikitext \
+    --w_bit 4 --q_group_size 128 \
+    --load_quant quant_cache/llama3-8b-w4-g128-awq.pt
+```
+## Results on Visual Language Models
+
+AWQ also seamlessly supports large multi-modal models (LMMs). Please refer to [TinyChat](./tinychat/README.md) for more details.
+
+
+<!-- AWQ also seamlessly supports large multi-modal models (LMMs). We demonstrate the results on the recent [VILA-1.5](https://github.com/Efficient-Large-Model/VILA) model family. -->
+
+<!-- 
+| VILA-1.5-3B   | VQA-v2            | GQA               | VizWiz  | ScienceQA         | TextVQA           | POPE    | MME     | MMBench           | MMBench-CN    | SEED    |
+| ----------- |:-----------------:|:-----------------:|:-------:|:-----------------:|:-----------------:|:-------:|:-------:|:-----------------:|:-------------:|:-------:|
+| FP16        | 80.4  | 61.5 | 53.5   | 69.0  | 60.4  | 85.9 | 1442.4 | 63.4 | 52.7   | 60.9 |
+| AWQ-INT4    | 80.0  | 61.1 | 53.8   | 67.8  | 60.4  | 85.9 | 1437.3 | 63.3 | 51.4   | 59.8 | 
+
+| VILA-1.5-8B    | VQA-v2            | GQA               | VizWiz  | ScienceQA         | TextVQA           | POPE    | MME     | MMBench           | MMBench-CN    | SEED    |
+| ----------- |:-----------------:|:-----------------:|:-------:|:-----------------:|:-----------------:|:-------:|:-------:|:-----------------:|:-------------:|:-------:|
+| FP16        | 80.9  | 61.9 | 58.7   | 79.9  | 66.3  | 84.4 | 1577.01 | 72.3 | 66.2   | 64.2 |
+| AWQ-INT4    | 80.3  | 61.7 | 59.3   | 79.0  | 65.4  | 82.9 | 1593.65 | 71.0 | 64.9   | 64.0 |
+
+| VILA-1.5-13B    | VQA-v2            | GQA               | VizWiz  | ScienceQA         | TextVQA           | POPE    | MME     | MMBench           | MMBench-CN    | SEED    |
+| ----------- |:-----------------:|:-----------------:|:-------:|:-----------------:|:-----------------:|:-------:|:-------:|:-----------------:|:-------------:|:-------:|
+| FP16       | 82.8  | 64.3 | 62.6   | 80.1  | 65.0  | 86.3 | 1569.55 | 74.9 | 66.3   | 65.1 |
+| AWQ-INT4    | 82.7  | 64.5 | 63.3   | 79.7  | 64.7  | 86.7 | 1531.35 | 74.7 | 66.7   | 65.1 |
+
+
+| VILA-1.5-40B    | VQA-v2            | GQA               | VizWiz  | ScienceQA         | TextVQA           | POPE    | MME     | MMBench           | MMBench-CN    | SEED    |
+| ----------- |:-----------------:|:-----------------:|:-------:|:-----------------:|:-----------------:|:-------:|:-------:|:-----------------:|:-------------:|:-------:|
+| FP16      | 84.3  | 64.6 | 62.2   | 87.2  | 73.6  | 87.3 | 1726.82 | 82.4 | 80.2   | 69.1 |
+| AWQ-INT4   | 84.1  | 64.4 | 61.3   | 86.7  | 73.2  | 88.2 | 1714.79 | 83.2 | 79.6   | 68.9 | 
+
+
+## Inference speed ( Token/sec )
+
+| $~~~~~~$               | Precision | A100  | 4090  | Orin |
+| ---------------------- | --------- | ----- | ----- | ---- |
+| VILA1.5-3B           | fp16      | 104.6 | 137.6 | 25.4 |
+| VILA1.5-3B-AWQ       | int4      | 182.8 | 215.5 | 42.5 |
+| VILA1.5-3B-S2        | fp16      | 104.3 | 137.2 | 24.6 |
+| VILA1.5-3B-S2-AWQ    | int4      | 180.2 | 219.3 | 40.1 |
+| Llama-3-VILA1.5-8B     | fp16      | 74.9  | 57.4  | 10.2 |
+| Llama-3-VILA1.5-8B-AWQ | int4      | 168.9 | 150.2 | 28.7 |
+| VILA1.5-13B            | fp16      | 50.9  | OOM   | 6.1  |
+| VILA1.5-13B-AWQ        | int4      | 115.9 | 105.7 | 20.6 |
+| VILA1.5-40B            | fp16      | OOM  | OOM   | --  |
+| VILA1.5-40B-AWQ        | int4      | 57.0 | OOM | -- | -->
+
+
+## Reference
+
+If you find AWQ useful or relevant to your research, please kindly cite our paper:
+
+```
+@inproceedings{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 Chen, Wei-Ming and Wang, Wei-Chen and Xiao, Guangxuan and Dang, Xingyu and Gan, Chuang and Han, Song},
+  booktitle={MLSys},
+  year={2024}
+}
+```
+
+## Related Projects
+
+[SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models](https://github.com/mit-han-lab/smoothquant)
+
+[GPTQ: Accurate Post-training Compression for Generative Pretrained Transformers](https://arxiv.org/abs/2210.17323)
+
+[Vicuna and FastChat](https://github.com/lm-sys/FastChat#readme)
+
+[LLaVA: Large Language and Vision Assistant](https://github.com/haotian-liu/LLaVA)
+
+[VILA: On Pre-training for Visual Language Models](https://github.com/Efficient-Large-Model/VILA)
+

llm-awq/awq/entry.py

@@ -0,0 +1,357 @@
+from lm_eval import evaluator, tasks
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
+import torch
+import argparse
+import os
+import json
+from accelerate import (
+    init_empty_weights,
+    infer_auto_device_map,
+    dispatch_model,
+    load_checkpoint_in_model,
+)
+from accelerate.utils.modeling import get_balanced_memory
+from awq.utils.parallel import auto_parallel
+from awq.quantize.pre_quant import run_awq, apply_awq
+from awq.quantize.quantizer import (
+    pseudo_quantize_model_weight,
+    real_quantize_model_weight,
+)
+from awq.utils.lm_eval_adaptor import LMEvalAdaptor
+from awq.utils.utils import simple_dispatch_model
+from datasets import load_dataset
+from torch import nn
+import tqdm
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model_path", type=str, help="path of the hf model")
+parser.add_argument("--dtype", type=str, default="float16", choices=["float16", "bfloat16"])
+parser.add_argument("--batch_size", type=int, default=1, help="batch size")
+parser.add_argument("--tasks", default=None, type=str)
+parser.add_argument("--output_path", default=None, type=str)
+parser.add_argument("--num_fewshot", type=int, default=0)
+# model config
+parser.add_argument("--parallel", action="store_true", help="enable model parallelism")
+# max memory to offload larger models to CPU
+parser.add_argument(
+    "--max_memory",
+    type=str,
+    nargs="*",
+    help="List of device_id:max_memory pairs to be parsed into a dictionary; "
+    + "Example: 0:10GiB 1:10GiB cpu:30GiB; "
+    + "mode details here: "
+    + "https://huggingface.co/docs/accelerate/usage_guides/big_modeling",
+)
+parser.add_argument(
+    "--auto_parallel",
+    action="store_true",
+    help="automatically set parallel and batch_size",
+)
+# quantization config
+parser.add_argument("--w_bit", type=int, default=None)
+parser.add_argument("--q_group_size", type=int, default=-1)
+parser.add_argument("--no_zero_point", action="store_true", help="disable zero_point")
+parser.add_argument("--q_backend", type=str, default="fake", choices=["fake", "real"])
+# save/load real quantized weights
+parser.add_argument("--dump_quant", type=str, default=None, help="save quantized model")
+parser.add_argument(
+    "--dump_fake", type=str, default=None, help="save fake-quantized model"
+)
+parser.add_argument("--load_quant", type=str, default=None, help="load quantized model")
+# apply/save/load awq
+parser.add_argument("--run_awq", action="store_true", help="perform awq search process")
+parser.add_argument(
+    "--dump_awq", type=str, default=None, help="save the awq search results"
+)
+parser.add_argument(
+    "--load_awq", type=str, default=None, help="load the awq search results"
+)
+parser.add_argument(
+    "--vila-15",
+    action="store_true",
+    help="quantizing vila 1.5",
+)
+parser.add_argument(
+    "--vila-20",
+    action="store_true",
+    help="quantizing or smoothing vila 2.0 (NVILA)",
+)
+parser.add_argument(
+    "--smooth_scale",
+    action="store_true",
+    help="generate the act scale of visiontower",
+)
+parser.add_argument(
+    "--media_path",
+    type=str,
+    nargs="+",
+    help="The input video to get act scale for visiontower",
+)
+parser.add_argument(
+    "--act_scale_path",
+    type=str,
+    default=None,
+    help="Path to save act scale",
+)
+args = parser.parse_args()
+assert (
+    args.act_scale_path is not None and len(args.media_path) > 0
+) or not args.smooth_scale
+vila_10_quant_mode = (
+    ("llava" in args.model_path.lower() or "vila" in args.model_path.lower())
+    and not args.vila_15
+    and not args.vila_20
+)
+
+max_memory = [v.split(":") for v in (args.max_memory or [])]
+max_memory = {(int(k) if k.isdigit() else k): v for k, v in max_memory}
+
+if args.auto_parallel:
+    gpu_list = auto_parallel(args)
+
+# get quantization config (apart from w_bit)
+q_config = {
+    "zero_point": not args.no_zero_point,  # by default True
+    "q_group_size": args.q_group_size,  # whether to use group quantization
+}
+print("Quantization config:", q_config)
+
+# build model and tokenizer
+
+
+def build_model_and_enc(model_path, dtype):
+    torch_dtype = torch.float16 if dtype == "float16" else torch.bfloat16
+    if not os.path.exists(model_path):  # look into ssd
+        raise FileNotFoundError(f"{model_path} not found!")
+    print(f"* Building model {model_path}")
+
+    # all hf model
+    if vila_10_quant_mode:
+        from llava.model.builder import load_pretrained_model
+        from llava.mm_utils import get_model_name_from_path
+
+        enc, model, image_processor, context_len = load_pretrained_model(
+            model_path=model_path,
+            model_base=None,
+            model_name=get_model_name_from_path(model_path),
+            device="cpu",
+            **{"use_cache": False},
+        )
+    else:
+        config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+        # Note (Haotian): To avoid OOM after huggingface transformers 4.36.2
+        config.use_cache = False
+        if "mpt" in config.__class__.__name__.lower():
+            enc = AutoTokenizer.from_pretrained(
+                config.tokenizer_name, trust_remote_code=True
+            )
+        else:
+            enc = AutoTokenizer.from_pretrained(
+                model_path, use_fast=False, trust_remote_code=True
+            )
+
+    if args.load_quant:  # directly load quantized weights
+        print("Loading pre-computed quantized weights...")
+        with init_empty_weights():
+            model = AutoModelForCausalLM.from_config(
+                config=config, torch_dtype=torch_dtype, trust_remote_code=True
+            )
+        real_quantize_model_weight(
+            model, w_bit=args.w_bit, q_config=q_config, init_only=True
+        )
+
+        model.tie_weights()
+
+        # Infer device map
+        kwargs = {"max_memory": max_memory} if len(max_memory) else {}
+        device_map = infer_auto_device_map(
+            model,
+            no_split_module_classes=[
+                "OPTDecoderLayer",
+                "LlamaDecoderLayer",
+                "BloomBlock",
+                "MPTBlock",
+                "DecoderLayer",
+            ],
+            **kwargs,
+        )
+        # Load checkpoint in the model
+        load_checkpoint_in_model(
+            model,
+            checkpoint=args.load_quant,
+            device_map=device_map,
+            offload_state_dict=True,
+        )
+        # Dispatch model
+        model = simple_dispatch_model(model, device_map=device_map)
+
+        model.eval()
+    else:  # fp16 to quantized
+        args.run_awq &= not args.load_awq  # if load_awq, no need to run awq
+        # Init model on CPU:
+        kwargs = {"torch_dtype": torch_dtype, "low_cpu_mem_usage": True}
+        if not vila_10_quant_mode:
+            model = AutoModelForCausalLM.from_pretrained(
+                model_path, config=config, trust_remote_code=True, **kwargs
+            )
+
+        model.eval()
+
+        if args.run_awq:
+            assert args.dump_awq, "Please save the awq results with --dump_awq"
+
+            awq_results = run_awq(
+                model,
+                enc,
+                w_bit=args.w_bit,
+                q_config=q_config,
+                n_samples=128,
+                seqlen=512,
+            )
+            if args.dump_awq:
+                dirpath = os.path.dirname(args.dump_awq)
+                os.makedirs(dirpath, exist_ok=True)
+
+                torch.save(awq_results, args.dump_awq)
+                print("AWQ results saved at", args.dump_awq)
+
+            exit(0)
+
+        if args.load_awq:
+            print("Loading pre-computed AWQ results from", args.load_awq)
+            awq_results = torch.load(args.load_awq, map_location="cpu")
+            apply_awq(model, awq_results)
+
+        # weight quantization
+        if args.w_bit is not None:
+            if args.q_backend == "fake":
+                assert (
+                    args.dump_quant is None
+                ), "Need to use real quantization to dump quantized weights"
+                pseudo_quantize_model_weight(model, w_bit=args.w_bit, q_config=q_config)
+                if args.dump_fake:
+                    model.save_pretrained(args.dump_fake)
+                    print("Pseudo-quantized models saved at", args.dump_fake)
+            elif args.q_backend == "real":  # real quantization
+                real_quantize_model_weight(model, w_bit=args.w_bit, q_config=q_config)
+                if args.dump_quant:
+                    if not args.dump_quant.endswith("v2.pt"):
+                        print("[Info] Auto-change the dump_quant file name to *v2.pt")
+                        args.dump_quant = args.dump_quant.replace(".pt", "-v2.pt")
+                    dirpath = os.path.dirname(args.dump_quant)
+                    os.makedirs(dirpath, exist_ok=True)
+
+                    print(f"Saving the quantized model at {args.dump_quant}...")
+                    torch.save(model.cpu().state_dict(), args.dump_quant)
+                    exit(0)
+            else:
+                raise NotImplementedError
+
+        # Move the model to GPU (as much as possible) for LM evaluation
+        kwargs = {
+            "max_memory": get_balanced_memory(
+                model, max_memory if len(max_memory) > 0 else None
+            )
+        }
+        device_map = infer_auto_device_map(
+            model,
+            # TODO: can we remove this?
+            no_split_module_classes=[
+                "OPTDecoderLayer",
+                "LlamaDecoderLayer",
+                "BloomBlock",
+                "MPTBlock",
+                "DecoderLayer",
+            ],
+            **kwargs,
+        )
+        model = dispatch_model(model, device_map=device_map)
+
+    return model, enc
+
+
+def main():
+    if args.output_path is not None and os.path.exists(args.output_path):
+        # print(f"Results {args.output_path} already generated. Exit.")
+        print(f"Results {args.output_path} already generated. Overwrite.")
+        # exit()
+
+    # a hack here to auto set model group
+    if args.smooth_scale and args.vila_20:
+        if os.path.exists(args.act_scale_path):
+            print(f"Found existing Smooth Scales {args.act_scale_path}, skip.")
+        else:
+            from awq.quantize import get_smooth_scale
+
+            act_scale = get_smooth_scale(args.model_path, args.media_path)
+            os.makedirs(os.path.dirname(args.act_scale_path), exist_ok=True)
+            torch.save(act_scale, args.act_scale_path)
+            print("Save act scales at " + str(args.act_scale_path))
+            args.model_path = args.model_path + "/llm"
+        if args.dump_awq is None and args.dump_quant is None:
+            exit()
+
+    if args.dump_awq and os.path.exists(args.dump_awq):
+        print(f"Found existing AWQ results {args.dump_awq}, exit.")
+        exit()
+    model, enc = build_model_and_enc(args.model_path, args.dtype)
+
+    if args.tasks is not None:
+        # https://github.com/IST-DASLab/gptq/blob/2d65066eeb06a5c9ff5184d8cebdf33662c67faf/llama.py#L206
+        if args.tasks == "wikitext":
+            testenc = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
+            testenc = enc("\n\n".join(testenc["text"]), return_tensors="pt")
+            model.seqlen = 2048
+            testenc = testenc.input_ids.to(model.device)
+            nsamples = testenc.numel() // model.seqlen
+            model = model.eval()
+            nlls = []
+            for i in tqdm.tqdm(range(nsamples), desc="evaluating..."):
+                batch = testenc[:, (i * model.seqlen) : ((i + 1) * model.seqlen)].to(
+                    model.device
+                )
+                with torch.no_grad():
+                    lm_logits = model(batch).logits
+                shift_logits = lm_logits[:, :-1, :].contiguous().float()
+                shift_labels = testenc[
+                    :, (i * model.seqlen) : ((i + 1) * model.seqlen)
+                ][:, 1:]
+                loss_fct = nn.CrossEntropyLoss()
+                loss = loss_fct(
+                    shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
+                )
+                neg_log_likelihood = loss.float() * model.seqlen
+                nlls.append(neg_log_likelihood)
+
+            ppl = torch.exp(torch.stack(nlls).sum() / (nsamples * model.seqlen))
+            print(ppl.item())
+
+            results = {"ppl": ppl.item()}
+            if args.output_path is not None:
+                os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
+                with open(args.output_path, "w") as f:
+                    json.dump(results, f, indent=2)
+        else:
+            task_names = args.tasks.split(",")
+
+            lm_eval_model = LMEvalAdaptor(args.model_path, model, enc, args.batch_size)
+            results = evaluator.simple_evaluate(
+                model=lm_eval_model,
+                tasks=task_names,
+                batch_size=args.batch_size,
+                no_cache=True,
+                num_fewshot=args.num_fewshot,
+            )
+
+            print(evaluator.make_table(results))
+
+        if args.output_path is not None:
+            os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
+            # otherwise cannot save
+            results["config"]["model"] = args.model_path
+            with open(args.output_path, "w") as f:
+                json.dump(results, f, indent=2)
+
+
+if __name__ == "__main__":
+    main()

llm-awq/awq/kernels/csrc/attention/README.md

@@ -0,0 +1,8 @@
+# Attention kernel from FasterTransformer
+
+This CUDA extension wraps the single-query attention [kernel](https://github.com/NVIDIA/FasterTransformer/blob/release/v5.2.1_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp) from
+FasterTransformer v5.2.1 for benchmarking purpose.
+
+```sh
+cd csrc/ft_attention && pip install .
+```

llm-awq/awq/kernels/csrc/attention/cuda_bf16_fallbacks.cuh

@@ -0,0 +1,257 @@
+// Downloaded from from FasterTransformer v5.2.1
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.2.1_tag/src/fastertransformer/utils/cuda_bf16_fallbacks.cuh
+/*
+ * Copyright (c) 2019-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>
+
+namespace fastertransformer {
+
+#ifdef ENABLE_BF16
+inline __device__ float2 bf1622float2(const __nv_bfloat162 val) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    float2 f_val;
+    f_val.x = __low2float(val);
+    f_val.y = __high2float(val);
+    return f_val;
+#else
+    return __bfloat1622float2(val);
+#endif
+}
+
+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
+}
+
+inline __device__ __nv_bfloat162 float22bf162(const float2 val) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    return __floats2bfloat162_rn(val.x, val.y);
+#else
+    return __float22bfloat162_rn(val);
+#endif
+}
+
+inline __device__ __nv_bfloat162 bf162bf162(const __nv_bfloat16 val) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    __nv_bfloat162 val2;
+    val2.x = val;
+    val2.y = val;
+    return val2;
+#else
+    return __bfloat162bfloat162(val);
+#endif
+}
+
+inline __device__ __nv_bfloat162 bf16hadd2(const __nv_bfloat162 x, const __nv_bfloat162 y) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    float fxl, fxh, fyl, fyh;
+    fxl = __low2float(x);
+    fxh = __high2float(x);
+    fyl = __low2float(y);
+    fyh = __high2float(y);
+    return __floats2bfloat162_rn(fxl + fyl, fxh + fyh);
+#else
+    return __hadd2(x, y);
+#endif
+}
+
+inline __device__ __nv_bfloat16 bf16hadd(const __nv_bfloat16 x, const __nv_bfloat16 y) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    return __float2bfloat16( __bfloat162float(x) + __bfloat162float(y) );
+#else
+    return __hadd(x, y);
+#endif
+}
+
+inline __device__ __nv_bfloat162 bf16hsub2(const __nv_bfloat162 x, const __nv_bfloat162 y) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    float fxl, fxh, fyl, fyh;
+    fxl = __low2float(x);
+    fxh = __high2float(x);
+    fyl = __low2float(y);
+    fyh = __high2float(y);
+    return __floats2bfloat162_rn(fxl - fyl, fxh - fyh);
+#else
+    return __hsub2(x, y);
+#endif
+}
+
+inline __device__ __nv_bfloat16 bf16hsub(const __nv_bfloat16 x, const __nv_bfloat16 y) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    return __float2bfloat16( __bfloat162float(x) - __bfloat162float(y) );
+#else
+    return __hsub(x, y);
+#endif
+}
+
+inline __device__ __nv_bfloat162 bf16hmul2(const __nv_bfloat162 x, const __nv_bfloat162 y) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    float fxl, fxh, fyl, fyh;
+    fxl = __low2float(x);
+    fxh = __high2float(x);
+    fyl = __low2float(y);
+    fyh = __high2float(y);
+    return __floats2bfloat162_rn(fxl * fyl, fxh * fyh);
+#else
+    return __hmul2(x, y);
+#endif
+}
+
+inline __device__ __nv_bfloat16 bf16hmul(const __nv_bfloat16 x, const __nv_bfloat16 y) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    return __float2bfloat16( __bfloat162float(x) * __bfloat162float(y) );
+#else
+    return __hmul(x, y);
+#endif
+}
+
+inline __device__ __nv_bfloat162 bf16hfma2(const __nv_bfloat162 x, const __nv_bfloat162 y, const __nv_bfloat162 z) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    float fxl, fxh, fyl, fyh, fzl, fzh;
+    fxl = __low2float(x);
+    fxh = __high2float(x);
+    fyl = __low2float(y);
+    fyh = __high2float(y);
+    fzl = __low2float(z);
+    fzh = __high2float(z);
+    return __floats2bfloat162_rn(fxl * fyl + fzl, fxh * fyh + fzh);
+#else
+    return __hfma2(x, y, z);
+#endif
+}
+
+inline __device__ __nv_bfloat16 bf16hfma(const __nv_bfloat16 x, const __nv_bfloat16 y, const __nv_bfloat16 z) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    return __float2bfloat16( __bfloat162float(x) * __bfloat162float(y) + __bfloat162float(z));
+#else
+    return __hfma(x, y, z);
+#endif
+}
+
+inline __device__ __nv_bfloat162 bf16exp2(const __nv_bfloat162 x) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    float fxl, fxh;
+    fxl = __low2float(x);
+    fxh = __high2float(x);;
+    return __floats2bfloat162_rn(expf(fxl), expf(fxh));
+#else
+    return h2exp(x);
+#endif
+}
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)
+inline __device__ __nv_bfloat162 operator*(const __nv_bfloat162 x, const __nv_bfloat162 y) { return bf16hmul2(x, y); };
+inline __device__ __nv_bfloat162 operator+(const __nv_bfloat162 x, const __nv_bfloat162 y) { return bf16hadd2(x, y); };
+
+inline __device__ __nv_bfloat162 make_bfloat162(const __nv_bfloat16 x, const __nv_bfloat16 y)
+{
+    __nv_bfloat162 t; t.x = x; t.y = y; return t;
+}
+
+#endif
+
+inline __device__ __nv_bfloat16 bf16hadd(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    return __float2bfloat16(__bfloat162float(a) + __bfloat162float(b) + __bfloat162float(c));
+#else
+    return a + b + c;
+#endif
+}
+
+inline __device__ __nv_bfloat16 bf16hadd(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c, __nv_bfloat16 d) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    return __float2bfloat16(__bfloat162float(a) + __bfloat162float(b) + __bfloat162float(c) + __bfloat162float(d));
+#else
+    return (__nv_bfloat16)((float)a + (float)b + (float)c + (float)d);
+#endif
+}
+
+inline __device__ __nv_bfloat162 bf16hadd2(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    float fal, fah, fbl, fbh, fcl, fch;
+    fal = __low2float(a);
+    fah = __high2float(a);
+    fbl = __low2float(b);
+    fbh = __high2float(b);
+    fcl = __low2float(c);
+    fch = __high2float(c);
+    return __floats2bfloat162_rn(fal + fbl + fcl, fah + fbh + fch);
+#else
+    return a + b + c;
+#endif
+}
+
+inline __device__ __nv_bfloat16 bf16hmul(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    return __float2bfloat16(__bfloat162float(a) * __bfloat162float(b) * __bfloat162float(c));
+#else
+    return a * b * c;
+#endif
+}
+
+inline __device__ __nv_bfloat162 bf16hmul2(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    float fal, fah, fbl, fbh, fcl, fch;
+    fal = __low2float(a);
+    fah = __high2float(a);
+    fbl = __low2float(b);
+    fbh = __high2float(b);
+    fcl = __low2float(c);
+    fch = __high2float(c);
+    return __floats2bfloat162_rn(fal * fbl * fcl, fah * fbh * fch);
+#else
+    return a * b * c;
+#endif
+}
+
+inline __device__ __nv_bfloat162 bf16hfma2(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c, __nv_bfloat162 d) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+    float fal, fah, fbl, fbh, fcl, fch, fdl, fdh;
+    fal = __low2float(a);
+    fah = __high2float(a);
+    fbl = __low2float(b);
+    fbh = __high2float(b);
+    fcl = __low2float(c);
+    fch = __high2float(c);
+    fdl = __low2float(d);
+    fdh = __high2float(d);
+    return __floats2bfloat162_rn(fal * fbl * fcl + fdl, fah * fbh * fch + fdh);
+#else
+    return a * b * c + d;
+#endif
+}
+
+#endif // ENABLE_BF16
+
+}  // namespace fastertransformer

llm-awq/awq/kernels/csrc/attention/cuda_bf16_wrapper.h

@@ -0,0 +1,23 @@
+// Downloaded from from FasterTransformer v5.2.1
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.2.1_tag/src/fastertransformer/utils/cuda_bf16_wrapper.h
+/*
+ * Copyright (c) 2019-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
+
+#ifdef ENABLE_BF16
+#include <cuda_bf16.h>
+#endif

llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention.cu

@@ -0,0 +1,154 @@
+// Adapted 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_128.cu
+/*
+ * 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.
+ */
+
+#include "decoder_masked_multihead_attention.h"
+#include "decoder_masked_multihead_attention_utils.h"
+#include "cuda_bf16_wrapper.h"
+#include <assert.h>
+#include <float.h>
+#include <type_traits>
+
+#include "decoder_masked_multihead_attention_template.hpp"
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#define MMHA_LAUNCH_KERNEL(T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, DO_CROSS_ATTENTION, stream)    \
+    size_t smem_sz = mmha::smem_size_in_bytes<T, DO_CROSS_ATTENTION>(params, THDS_PER_VALUE, THDS_PER_BLOCK);          \
+    auto kernel = mmha::masked_multihead_attention_kernel<T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE,                 \
+                                                          THDS_PER_BLOCK, DO_CROSS_ATTENTION>;                         \
+    if (smem_sz >= 48 * 1024) {                                                                                        \
+        cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_sz);                            \
+    }                                                                                                                  \
+    dim3 grid(params.num_heads, params.batch_size);                                                                    \
+    kernel<<<grid, THDS_PER_BLOCK, smem_sz, stream>>>(params)
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// !!! Specialize the launcher for Cross attention
+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)
+{
+    constexpr int  THREADS_PER_VALUE  = Dh_MAX * sizeof(T) / 16;
+    constexpr bool DO_CROSS_ATTENTION = std::is_same<KERNEL_PARAMS_TYPE, Cross_multihead_attention_params<T>>::value;
+    int            tlength            = (DO_CROSS_ATTENTION) ? params.memory_max_len : params.timestep;
+    // printf("tlength, CROSS_ATTENTION = %d, %d\n", tlength, DO_CROSS_ATTENTION);
+    if (tlength < 32) {
+        MMHA_LAUNCH_KERNEL(T, Dh, Dh_MAX, 4, THREADS_PER_VALUE, 64, DO_CROSS_ATTENTION, stream);
+    }
+    else if (tlength < 2048) {
+        MMHA_LAUNCH_KERNEL(T, Dh, Dh_MAX, 2, THREADS_PER_VALUE, 128, DO_CROSS_ATTENTION, stream);
+    }
+    else {
+        MMHA_LAUNCH_KERNEL(T, Dh, Dh_MAX, 1, THREADS_PER_VALUE, 256, DO_CROSS_ATTENTION, stream);
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#undef MMHA_LAUNCH_KERNEL
+
+template<typename T, typename KERNEL_PARAMS_TYPE>
+void multihead_attention_(const KERNEL_PARAMS_TYPE& params, const cudaStream_t& stream)
+{
+    switch (params.hidden_size_per_head) {
+        case 32:
+            mmha_launch_kernel<T, 32, 32, KERNEL_PARAMS_TYPE>(params, stream);
+            break;
+        case 48:
+            mmha_launch_kernel<T, 48, 64, KERNEL_PARAMS_TYPE>(params, stream);
+            break;
+        case 64:
+            mmha_launch_kernel<T, 64, 64, KERNEL_PARAMS_TYPE>(params, stream);
+            break;
+        case 80:
+            mmha_launch_kernel<T, 80, 128, KERNEL_PARAMS_TYPE>(params, stream);
+            break;
+        case 96:
+            mmha_launch_kernel<T, 96, 128, KERNEL_PARAMS_TYPE>(params, stream);
+            break;
+        case 112:
+            mmha_launch_kernel<T, 112, 128, KERNEL_PARAMS_TYPE>(params, stream);
+            break;
+        case 128:
+            mmha_launch_kernel<T, 128, 128, KERNEL_PARAMS_TYPE>(params, stream);
+            break;
+        case 160:
+            mmha_launch_kernel<T, 160, 256, KERNEL_PARAMS_TYPE>(params, stream);
+            break;
+        case 192:
+            mmha_launch_kernel<T, 192, 256, KERNEL_PARAMS_TYPE>(params, stream);
+            break;
+        case 224:
+            mmha_launch_kernel<T, 224, 256, KERNEL_PARAMS_TYPE>(params, stream);
+            break;
+        case 256:
+            mmha_launch_kernel<T, 256, 256, KERNEL_PARAMS_TYPE>(params, stream);
+            break;
+        default:
+            assert(false);
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+void masked_multihead_attention(const Masked_multihead_attention_params<float>& params, const cudaStream_t& stream)
+{
+    multihead_attention_<float, Masked_multihead_attention_params<float>>(params, stream);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+void masked_multihead_attention(const Masked_multihead_attention_params<uint16_t>& params, const cudaStream_t& stream)
+{
+    multihead_attention_<uint16_t, Masked_multihead_attention_params<uint16_t>>(params, stream);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#ifdef ENABLE_BF16
+void masked_multihead_attention(const Masked_multihead_attention_params<__nv_bfloat16>& params,
+                                const cudaStream_t&                                     stream)
+{
+    multihead_attention_<__nv_bfloat16, Masked_multihead_attention_params<__nv_bfloat16>>(params, stream);
+}
+#endif
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+void cross_multihead_attention(const Cross_multihead_attention_params<float>& params, const cudaStream_t& stream)
+{
+    multihead_attention_<float, Cross_multihead_attention_params<float>>(params, stream);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+void cross_multihead_attention(const Cross_multihead_attention_params<uint16_t>& params, const cudaStream_t& stream)
+{
+    multihead_attention_<uint16_t, Cross_multihead_attention_params<uint16_t>>(params, stream);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#ifdef ENABLE_BF16
+void cross_multihead_attention(const Cross_multihead_attention_params<__nv_bfloat16>& params,
+                               const cudaStream_t&                                    stream)
+{
+    multihead_attention_<__nv_bfloat16, Cross_multihead_attention_params<__nv_bfloat16>>(params, stream);
+}
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////

llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention_utils.h

@@ -0,0 +1,1795 @@
+// 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_utils.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_bf16_fallbacks.cuh"
+#include <stdint.h>
+
+using namespace fastertransformer;
+
+namespace mmha {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Float8_ {
+    float2 x;
+    float2 y;
+    float2 z;
+    float2 w;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Float4_ {
+    float2 x;
+    float2 y;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#ifdef ENABLE_BF16
+struct bf16_4_t {
+    __nv_bfloat162 x;
+    __nv_bfloat162 y;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct bf16_8_t {
+    __nv_bfloat162 x;
+    __nv_bfloat162 y;
+    __nv_bfloat162 z;
+    __nv_bfloat162 w;
+};
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+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<Float4_> {
+    static constexpr int value = 4;
+};
+template<>
+struct num_elems<Float8_> {
+    static constexpr int value = 8;
+};
+
+template<>
+struct num_elems<uint32_t> {
+    static constexpr int value = 2;
+};
+template<>
+struct num_elems<uint2> {
+    static constexpr int value = 4;
+};
+template<>
+struct num_elems<uint4> {
+    static constexpr int value = 8;
+};
+
+#ifdef ENABLE_BF16
+template<>
+struct num_elems<__nv_bfloat162> {
+    static constexpr int value = 2;
+};
+template<>
+struct num_elems<bf16_4_t> {
+    static constexpr int value = 4;
+};
+template<>
+struct num_elems<bf16_8_t> {
+    static constexpr int value = 8;
+};
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T, int N>
+struct packed_type;
+template<typename T>
+struct packed_type<T, 1> {
+    using type = T;
+};
+template<>
+struct packed_type<int8_t, 2> {
+    using type = int16_t;
+};
+template<>
+struct packed_type<int8_t, 4> {
+    using type = int32_t;
+};
+template<>
+struct packed_type<int8_t, 8> {
+    using type = int64_t;
+};
+
+template<>
+struct packed_type<float, 2> {
+    using type = float2;
+};
+template<>
+struct packed_type<float, 4> {
+    using type = float4;
+};
+template<>
+struct packed_type<float, 8> {
+    using type = Float8_;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float add(float a, float b)
+{
+    return a + b;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float2 add(float2 a, float2 b)
+{
+    float2 c;
+    c.x = add(a.x, b.x);
+    c.y = add(a.y, b.y);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float4 add(float4 a, float4 b)
+{
+    float4 c;
+    c.x = add(a.x, b.x);
+    c.y = add(a.y, b.y);
+    c.z = add(a.z, b.z);
+    c.w = add(a.w, b.w);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#ifdef ENABLE_BF16
+inline __device__ __nv_bfloat16 add(__nv_bfloat16 a, __nv_bfloat16 b)
+{
+    return a + b;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ __nv_bfloat162 add(__nv_bfloat162 a, __nv_bfloat162 b)
+{
+    return bf16hadd2(a, b);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ bf16_4_t add(bf16_4_t a, bf16_4_t b)
+{
+    bf16_4_t c;
+    c.x = add(a.x, b.x);
+    c.y = add(a.y, b.y);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ bf16_8_t add(bf16_8_t a, bf16_8_t b)
+{
+    bf16_8_t c;
+    c.x = add(a.x, b.x);
+    c.y = add(a.y, b.y);
+    c.z = add(a.z, b.z);
+    c.w = add(a.w, b.w);
+    return c;
+}
+#endif  // ENABLE_BF16
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ uint16_t add(uint16_t a, uint16_t b)
+{
+    uint16_t c;
+    asm volatile("add.f16 %0, %1, %2;\n" : "=h"(c) : "h"(a), "h"(b));
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ uint32_t add(uint32_t a, uint32_t b)
+{
+    uint32_t c;
+    asm volatile("add.f16x2 %0, %1, %2;\n" : "=r"(c) : "r"(a), "r"(b));
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ uint2 add(uint2 a, uint2 b)
+{
+    uint2 c;
+    c.x = add(a.x, b.x);
+    c.y = add(a.y, b.y);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ uint4 add(uint4 a, uint4 b)
+{
+    uint4 c;
+    c.x = add(a.x, b.x);
+    c.y = add(a.y, b.y);
+    c.z = add(a.z, b.z);
+    c.w = add(a.w, b.w);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ uint16_t float_to_half(float f)
+{
+    union {
+        uint32_t u32;
+        uint16_t u16[2];
+    } tmp;
+#if 0 && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800  // Is it better?
+  float zero = 0.f;
+  asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;\n" : "=r"(tmp.u32) : "f"(zero), "f"(f));
+#else
+    asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[0]) : "f"(f));
+#endif
+    return tmp.u16[0];
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ uint32_t float2_to_half2(float2 f)
+{
+    union {
+        uint32_t u32;
+        uint16_t u16[2];
+    } tmp;
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+    asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;\n" : "=r"(tmp.u32) : "f"(f.y), "f"(f.x));
+#else
+    asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[0]) : "f"(f.x));
+    asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[1]) : "f"(f.y));
+#endif
+    return tmp.u32;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float half_to_float(uint16_t h)
+{
+    float f;
+    asm volatile("cvt.f32.f16 %0, %1;\n" : "=f"(f) : "h"(h));
+    return f;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float2 half2_to_float2(uint32_t v)
+{
+    uint16_t lo, hi;
+    asm volatile("mov.b32 {%0, %1}, %2;\n" : "=h"(lo), "=h"(hi) : "r"(v));
+    return make_float2(half_to_float(lo), half_to_float(hi));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float add(float a, uint16_t b)
+{
+    return a + half_to_float(b);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#ifdef ENABLE_BF16
+inline __device__ float add(float a, __nv_bfloat16 b)
+{
+    return a + __bfloat162float(b);
+}
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float2 add(uint32_t a, float2 fb)
+{
+    float2 fa = half2_to_float2(a);
+    return add(fa, fb);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ Float4_ add(uint2 a, Float4_ fb)
+{
+    Float4_ fc;
+    fc.x = add(a.x, fb.x);
+    fc.y = add(a.y, fb.y);
+    return fc;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ Float8_ add(uint4 a, Float8_ fb)
+{
+    Float8_ fc;
+    fc.x = add(a.x, fb.x);
+    fc.y = add(a.y, fb.y);
+    fc.z = add(a.z, fb.z);
+    fc.w = add(a.w, fb.w);
+    return fc;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ uint32_t h0_h0(uint16_t a)
+{
+    uint32_t b;
+    asm volatile("mov.b32 %0, {%1, %1};" : "=r"(b) : "h"(a));
+    return b;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float fma(float a, float b, float c)
+{
+    return a * b + c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float2 fma(float2 a, float2 b, float2 c)
+{
+    float2 d;
+    d.x = fma(a.x, b.x, c.x);
+    d.y = fma(a.y, b.y, c.y);
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float2 fma(float a, float2 b, float2 c)
+{
+    float2 d;
+    d.x = fma(a, b.x, c.x);
+    d.y = fma(a, b.y, c.y);
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float4 fma(float4 a, float4 b, float4 c)
+{
+    float4 d;
+    d.x = fma(a.x, b.x, c.x);
+    d.y = fma(a.y, b.y, c.y);
+    d.z = fma(a.z, b.z, c.z);
+    d.w = fma(a.w, b.w, c.w);
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float4 fma(float a, float4 b, float4 c)
+{
+    float4 d;
+    d.x = fma(a, b.x, c.x);
+    d.y = fma(a, b.y, c.y);
+    d.z = fma(a, b.z, c.z);
+    d.w = fma(a, b.w, c.w);
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ Float4_ fma(float a, Float4_ b, Float4_ c)
+{
+    Float4_ d;
+    d.x = fma(a, b.x, c.x);
+    d.y = fma(a, b.y, c.y);
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ Float8_ fma(float a, Float8_ b, Float8_ c)
+{
+    Float8_ d;
+    d.x = fma(a, b.x, c.x);
+    d.y = fma(a, b.y, c.y);
+    d.z = fma(a, b.z, c.z);
+    d.w = fma(a, b.w, c.w);
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#ifdef ENABLE_BF16
+inline __device__ float2 add(__nv_bfloat162 a, float2 fb)
+{
+    float2 fa = bf1622float2(a);
+    return add(fa, fb);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ Float4_ add(bf16_4_t a, Float4_ fb)
+{
+    Float4_ fc;
+    fc.x = add(a.x, fb.x);
+    fc.y = add(a.y, fb.y);
+    return fc;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ Float8_ add(bf16_8_t a, Float8_ fb)
+{
+    Float8_ fc;
+    fc.x = add(a.x, fb.x);
+    fc.y = add(a.y, fb.y);
+    fc.z = add(a.z, fb.z);
+    fc.w = add(a.w, fb.w);
+    return fc;
+}
+#endif  // ENABLE_BF16
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ uint32_t fma(uint32_t a, uint32_t b, uint32_t c)
+{
+    uint32_t d;
+    asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(d) : "r"(a), "r"(b), "r"(c));
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ uint32_t fma(uint16_t a, uint32_t b, uint32_t c)
+{
+    return fma(h0_h0(a), b, c);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ uint2 fma(uint2 a, uint2 b, uint2 c)
+{
+    uint2 d;
+    d.x = fma(a.x, b.x, c.x);
+    d.y = fma(a.y, b.y, c.y);
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ uint2 fma(uint16_t a, uint2 b, uint2 c)
+{
+    uint32_t s = h0_h0(a);
+    uint2    d;
+    d.x = fma(s, b.x, c.x);
+    d.y = fma(s, b.y, c.y);
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ uint4 fma(uint4 a, uint4 b, uint4 c)
+{
+    uint4 d;
+    d.x = fma(a.x, b.x, c.x);
+    d.y = fma(a.y, b.y, c.y);
+    d.z = fma(a.z, b.z, c.z);
+    d.w = fma(a.w, b.w, c.w);
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ uint4 fma(uint16_t a, uint4 b, uint4 c)
+{
+    uint32_t s = h0_h0(a);
+    uint4    d;
+    d.x = fma(s, b.x, c.x);
+    d.y = fma(s, b.y, c.y);
+    d.z = fma(s, b.z, c.z);
+    d.w = fma(s, b.w, c.w);
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float fma(uint16_t a, uint16_t b, float fc)
+{
+    float fa = half_to_float(a);
+    float fb = half_to_float(b);
+    return fa * fb + fc;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float2 fma(uint32_t a, uint32_t b, float2 fc)
+{
+    float2 fa = half2_to_float2(a);
+    float2 fb = half2_to_float2(b);
+    return fma(fa, fb, fc);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float2 fma(uint16_t a, uint32_t b, float2 fc)
+{
+    return fma(h0_h0(a), b, fc);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ Float4_ fma(uint2 a, uint2 b, Float4_ fc)
+{
+    Float4_ fd;
+    fd.x = fma(a.x, b.x, fc.x);
+    fd.y = fma(a.y, b.y, fc.y);
+    return fd;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ Float4_ fma(uint16_t a, uint2 b, Float4_ fc)
+{
+    uint32_t s = h0_h0(a);
+    Float4_  fd;
+    fd.x = fma(s, b.x, fc.x);
+    fd.y = fma(s, b.y, fc.y);
+    return fd;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ Float8_ fma(uint4 a, uint4 b, Float8_ fc)
+{
+    Float8_ fd;
+    fd.x = fma(a.x, b.x, fc.x);
+    fd.y = fma(a.y, b.y, fc.y);
+    fd.z = fma(a.z, b.z, fc.z);
+    fd.w = fma(a.w, b.w, fc.w);
+    return fd;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ Float8_ fma(uint16_t a, uint4 b, Float8_ fc)
+{
+    uint32_t s = h0_h0(a);
+    Float8_  fd;
+    fd.x = fma(s, b.x, fc.x);
+    fd.y = fma(s, b.y, fc.y);
+    fd.z = fma(s, b.z, fc.z);
+    fd.w = fma(s, b.w, fc.w);
+    return fd;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+#ifdef ENABLE_BF16
+inline __device__ __nv_bfloat162 fma(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c)
+{
+    return bf16hfma2(a, b, c);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ __nv_bfloat162 fma(__nv_bfloat16 a, __nv_bfloat162 b, __nv_bfloat162 c)
+{
+    return bf16hfma2(bf162bf162(a), b, c);
+}
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ bf16_4_t fma(bf16_4_t a, bf16_4_t b, bf16_4_t c)
+{
+    bf16_4_t d;
+    d.x = fma(a.x, b.x, c.x);
+    d.y = fma(a.y, b.y, c.y);
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ bf16_4_t fma(__nv_bfloat16 a, bf16_4_t b, bf16_4_t c)
+{
+    __nv_bfloat162 s = bf162bf162(a);
+    bf16_4_t       d;
+    d.x = fma(s, b.x, c.x);
+    d.y = fma(s, b.y, c.y);
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ bf16_8_t fma(bf16_8_t a, bf16_8_t b, bf16_8_t c)
+{
+    bf16_8_t d;
+    d.x = fma(a.x, b.x, c.x);
+    d.y = fma(a.y, b.y, c.y);
+    d.z = fma(a.z, b.z, c.z);
+    d.w = fma(a.w, b.w, c.w);
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ bf16_8_t fma(__nv_bfloat16 a, bf16_8_t b, bf16_8_t c)
+{
+    __nv_bfloat162 s = bf162bf162(a);
+    bf16_8_t       d;
+    d.x = fma(s, b.x, c.x);
+    d.y = fma(s, b.y, c.y);
+    d.z = fma(s, b.z, c.z);
+    d.w = fma(s, b.w, c.w);
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float fma(__nv_bfloat16 a, __nv_bfloat16 b, float fc)
+{
+    return __bfloat162float(a) * __bfloat162float(b) + fc;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float2 fma(__nv_bfloat162 a, __nv_bfloat162 b, float2 fc)
+{
+    float2 fa = bf1622float2(a);
+    float2 fb = bf1622float2(b);
+    return fma(fa, fb, fc);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float2 fma(__nv_bfloat16 a, __nv_bfloat162 b, float2 fc)
+{
+    return fma(bf162bf162(a), b, fc);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ Float4_ fma(bf16_4_t a, bf16_4_t b, Float4_ fc)
+{
+    Float4_ fd;
+    fd.x = fma(a.x, b.x, fc.x);
+    fd.y = fma(a.y, b.y, fc.y);
+    return fd;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ Float4_ fma(__nv_bfloat16 a, bf16_4_t b, Float4_ fc)
+{
+    __nv_bfloat162 s = bf162bf162(a);
+    Float4_        fd;
+    fd.x = fma(s, b.x, fc.x);
+    fd.y = fma(s, b.y, fc.y);
+    return fd;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ Float8_ fma(bf16_8_t a, bf16_8_t b, Float8_ fc)
+{
+    Float8_ fd;
+    fd.x = fma(a.x, b.x, fc.x);
+    fd.y = fma(a.y, b.y, fc.y);
+    fd.z = fma(a.z, b.z, fc.z);
+    fd.w = fma(a.w, b.w, fc.w);
+    return fd;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ Float8_ fma(__nv_bfloat16 a, bf16_8_t b, Float8_ fc)
+{
+    __nv_bfloat162 s = bf162bf162(a);
+    Float8_        fd;
+    fd.x = fma(s, b.x, fc.x);
+    fd.y = fma(s, b.y, fc.y);
+    fd.z = fma(s, b.z, fc.z);
+    fd.w = fma(s, b.w, fc.w);
+    return fd;
+}
+#endif  // ENABLE_BF16
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename Acc, typename A, typename B>
+inline __device__ Acc mul(A a, B b)
+{
+    return a * b;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ float mul<float, float>(float a, float b)
+{
+    return a * b;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ float2 mul(float2 a, float2 b)
+{
+    float2 c;
+    c.x = a.x * b.x;
+    c.y = a.y * b.y;
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ float2 mul(float a, float2 b)
+{
+    float2 c;
+    c.x = a * b.x;
+    c.y = a * b.y;
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ float4 mul(float4 a, float4 b)
+{
+    float4 c;
+    c.x = a.x * b.x;
+    c.y = a.y * b.y;
+    c.z = a.z * b.z;
+    c.w = a.w * b.w;
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ float4 mul(float a, float4 b)
+{
+    float4 c;
+    c.x = a * b.x;
+    c.y = a * b.y;
+    c.z = a * b.z;
+    c.w = a * b.w;
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ Float8_ mul(float a, Float8_ b)
+{
+    Float8_ c;
+    c.x = make_float2(a * b.x.x, a * b.x.y);
+    c.y = make_float2(a * b.y.x, a * b.y.y);
+    c.z = make_float2(a * b.z.x, a * b.z.y);
+    c.w = make_float2(a * b.w.x, a * b.w.y);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ uint16_t mul(uint16_t a, uint16_t b)
+{
+    uint16_t c;
+    asm volatile("mul.f16 %0, %1, %2;\n" : "=h"(c) : "h"(a), "h"(b));
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ uint32_t mul(uint32_t a, uint32_t b)
+{
+    uint32_t c;
+    asm volatile("mul.f16x2 %0, %1, %2;\n" : "=r"(c) : "r"(a), "r"(b));
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ uint32_t mul(uint16_t a, uint32_t b)
+{
+    return mul<uint32_t, uint32_t, uint32_t>(h0_h0(a), b);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ uint2 mul(uint2 a, uint2 b)
+{
+    uint2 c;
+    c.x = mul<uint32_t, uint32_t, uint32_t>(a.x, b.x);
+    c.y = mul<uint32_t, uint32_t, uint32_t>(a.y, b.y);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ uint2 mul(uint16_t a, uint2 b)
+{
+    uint32_t s = h0_h0(a);
+    uint2    c;
+    c.x = mul<uint32_t, uint32_t, uint32_t>(s, b.x);
+    c.y = mul<uint32_t, uint32_t, uint32_t>(s, b.y);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ uint4 mul(uint4 a, uint4 b)
+{
+    uint4 c;
+    c.x = mul<uint32_t, uint32_t, uint32_t>(a.x, b.x);
+    c.y = mul<uint32_t, uint32_t, uint32_t>(a.y, b.y);
+    c.z = mul<uint32_t, uint32_t, uint32_t>(a.z, b.z);
+    c.w = mul<uint32_t, uint32_t, uint32_t>(a.w, b.w);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ uint4 mul(uint16_t a, uint4 b)
+{
+    uint32_t s = h0_h0(a);
+    uint4    c;
+    c.x = mul<uint32_t, uint32_t, uint32_t>(s, b.x);
+    c.y = mul<uint32_t, uint32_t, uint32_t>(s, b.y);
+    c.z = mul<uint32_t, uint32_t, uint32_t>(s, b.z);
+    c.w = mul<uint32_t, uint32_t, uint32_t>(s, b.w);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ float mul(uint16_t a, uint16_t b)
+{
+    float fa = half_to_float(a);
+    float fb = half_to_float(b);
+    return fa * fb;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ float mul(uint16_t a, float b)
+{
+    return half_to_float(a) * b;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ float2 mul(uint32_t a, uint32_t b)
+{
+    float2 fa = half2_to_float2(a);
+    float2 fb = half2_to_float2(b);
+    return mul<float2, float2, float2>(fa, fb);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ float2 mul(uint16_t a, uint32_t b)
+{
+    return mul<float2, uint32_t, uint32_t>(h0_h0(a), b);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ Float4_ mul(uint2 a, uint2 b)
+{
+    Float4_ fc;
+    fc.x = mul<float2, uint32_t, uint32_t>(a.x, b.x);
+    fc.y = mul<float2, uint32_t, uint32_t>(a.y, b.y);
+    return fc;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ Float4_ mul(uint16_t a, uint2 b)
+{
+    uint32_t s = h0_h0(a);
+    Float4_  fc;
+    fc.x = mul<float2, uint32_t, uint32_t>(s, b.x);
+    fc.y = mul<float2, uint32_t, uint32_t>(s, b.y);
+    return fc;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ Float8_ mul(uint4 a, uint4 b)
+{
+    Float8_ fc;
+    fc.x = mul<float2, uint32_t, uint32_t>(a.x, b.x);
+    fc.y = mul<float2, uint32_t, uint32_t>(a.y, b.y);
+    fc.z = mul<float2, uint32_t, uint32_t>(a.z, b.z);
+    fc.w = mul<float2, uint32_t, uint32_t>(a.w, b.w);
+    return fc;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ Float8_ mul(uint16_t a, uint4 b)
+{
+    uint32_t s = h0_h0(a);
+    Float8_  fc;
+    fc.x = mul<float2, uint32_t, uint32_t>(s, b.x);
+    fc.y = mul<float2, uint32_t, uint32_t>(s, b.y);
+    fc.z = mul<float2, uint32_t, uint32_t>(s, b.z);
+    fc.w = mul<float2, uint32_t, uint32_t>(s, b.w);
+    return fc;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#ifdef ENABLE_BF16
+template<>
+inline __device__ __nv_bfloat16 mul(__nv_bfloat16 a, __nv_bfloat16 b)
+{
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+    return __hmul(a, b);
+#else
+    return bf16hmul(a, b);
+#endif
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ __nv_bfloat162 mul(__nv_bfloat162 a, __nv_bfloat162 b)
+{
+    return bf16hmul2(a, b);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ __nv_bfloat162 mul(__nv_bfloat16 a, __nv_bfloat162 b)
+{
+    return mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(bf162bf162(a), b);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ bf16_4_t mul(bf16_4_t a, bf16_4_t b)
+{
+    bf16_4_t c;
+    c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.x, b.x);
+    c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.y, b.y);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ bf16_4_t mul(__nv_bfloat16 a, bf16_4_t b)
+{
+    __nv_bfloat162 s = bf162bf162(a);
+    bf16_4_t       c;
+    c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.x);
+    c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.y);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ bf16_8_t mul(bf16_8_t a, bf16_8_t b)
+{
+    bf16_8_t c;
+    c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.x, b.x);
+    c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.y, b.y);
+    c.z = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.z, b.z);
+    c.w = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.w, b.w);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ bf16_8_t mul(__nv_bfloat16 a, bf16_8_t b)
+{
+    __nv_bfloat162 s = bf162bf162(a);
+    bf16_8_t       c;
+    c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.x);
+    c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.y);
+    c.z = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.z);
+    c.w = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.w);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ float mul(__nv_bfloat16 a, __nv_bfloat16 b)
+{
+    float fa = (float)a;
+    float fb = (float)b;
+    return fa * fb;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ float mul(__nv_bfloat16 a, float b)
+{
+    return __bfloat162float(a) * b;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ float2 mul(__nv_bfloat162 a, __nv_bfloat162 b)
+{
+    float2 fa = bf1622float2(a);
+    float2 fb = bf1622float2(b);
+    return mul<float2, float2, float2>(fa, fb);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ float2 mul(__nv_bfloat16 a, __nv_bfloat162 b)
+{
+    return mul<float2, __nv_bfloat162, __nv_bfloat162>(bf162bf162(a), b);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ Float4_ mul(bf16_4_t a, bf16_4_t b)
+{
+    Float4_ fc;
+    fc.x = mul<float2, __nv_bfloat162, __nv_bfloat162>(a.x, b.x);
+    fc.y = mul<float2, __nv_bfloat162, __nv_bfloat162>(a.y, b.y);
+    return fc;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ Float4_ mul(__nv_bfloat16 a, bf16_4_t b)
+{
+    __nv_bfloat162 s = bf162bf162(a);
+    Float4_        fc;
+    fc.x = mul<float2, __nv_bfloat162, __nv_bfloat162>(s, b.x);
+    fc.y = mul<float2, __nv_bfloat162, __nv_bfloat162>(s, b.y);
+    return fc;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ Float8_ mul(bf16_8_t a, bf16_8_t b)
+{
+    Float8_ fc;
+    fc.x = mul<float2, __nv_bfloat162, __nv_bfloat162>(a.x, b.x);
+    fc.y = mul<float2, __nv_bfloat162, __nv_bfloat162>(a.y, b.y);
+    fc.z = mul<float2, __nv_bfloat162, __nv_bfloat162>(a.z, b.z);
+    fc.w = mul<float2, __nv_bfloat162, __nv_bfloat162>(a.w, b.w);
+    return fc;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+inline __device__ Float8_ mul(__nv_bfloat16 a, bf16_8_t b)
+{
+    __nv_bfloat162 s = bf162bf162(a);
+    Float8_        fc;
+    fc.x = mul<float2, __nv_bfloat162, __nv_bfloat162>(s, b.x);
+    fc.y = mul<float2, __nv_bfloat162, __nv_bfloat162>(s, b.y);
+    fc.z = mul<float2, __nv_bfloat162, __nv_bfloat162>(s, b.z);
+    fc.w = mul<float2, __nv_bfloat162, __nv_bfloat162>(s, b.w);
+    return fc;
+}
+#endif  // ENABLE_BF16
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float sum(float v)
+{
+    return v;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float sum(float2 v)
+{
+    return v.x + v.y;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float sum(float4 v)
+{
+    return v.x + v.y + v.z + v.w;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#ifdef ENABLE_BF16
+inline __device__ float sum(__nv_bfloat162 v)
+{
+    float2 vf = bf1622float2(v);
+    return vf.x + vf.y;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float sum(bf16_4_t v)
+{
+    return sum(v.x) + sum(v.y);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float sum(bf16_8_t v)
+{
+    return sum(v.x) + sum(v.y) + sum(v.z) + sum(v.w);
+}
+#endif  // ENABLE_BF16
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float sum(uint16_t v)
+{
+    return half_to_float(v);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float sum(uint32_t v)
+{
+    float2 tmp = half2_to_float2(v);
+    return tmp.x + tmp.y;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float sum(uint2 v)
+{
+    uint32_t c = add(v.x, v.y);
+    return sum(c);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float sum(uint4 v)
+{
+#if 1
+    uint32_t c = add(v.x, v.y);
+    c          = add(c, v.z);
+    c          = add(c, v.w);
+#else
+    uint32_t c = add(v.x, v.y);
+    uint32_t d = add(v.z, v.w);
+    c          = add(c, d);
+#endif
+    return sum(c);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float sum(Float4_ v)
+{
+    return v.x.x + v.x.y + v.y.x + v.y.y;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float sum(Float8_ v)
+{
+    return v.x.x + v.x.y + v.y.x + v.y.y + v.z.x + v.z.y + v.w.x + v.w.y;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T>
+inline __device__ float dot(T a, T b)
+{
+    return sum(mul<T, T, T>(a, b));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename A, typename T>
+inline __device__ float dot(T a, T b)
+{
+    return sum(mul<A, T, T>(a, b));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void zero(uint16_t& dst)
+{
+    dst = uint16_t(0);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T>
+inline __device__ void zero(T& dst)
+{
+    constexpr int WORDS = sizeof(T) / 4;
+    union {
+        T        raw;
+        uint32_t words[WORDS];
+    } tmp;
+#pragma unroll
+    for (int ii = 0; ii < WORDS; ++ii) {
+        tmp.words[ii] = 0u;
+    }
+    dst = tmp.raw;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// inline __device__ float2 rotary_embedding_coefficient(const int zid, const int rot_embed_dim, const float t_step, const float base)
+// {
+//     const float inv_freq = t_step / pow(base, zid / (float)rot_embed_dim);
+//     return {cos(inv_freq), sin(inv_freq)};
+// }
+
+// with scale
+inline __device__ float2 rotary_embedding_coefficient(
+        const int zid, const int rot_embed_dim, const float t_step, const float base, const float scale)
+{
+    const float inv_freq = (t_step * scale) / pow(base, zid / (float)rot_embed_dim);
+    return {cos(inv_freq), sin(inv_freq)};
+}
+
+
+inline __device__ float2 rotary_embedding_transform(const float2 v, const float2 coef)
+{
+    float2 rot_v;
+    rot_v.x = coef.x * v.x - coef.y * v.y;
+    rot_v.y = coef.x * v.y + coef.y * v.x;
+    return rot_v;
+}
+
+inline __device__ uint32_t rotary_embedding_transform(const uint32_t v, const float2 coef)
+{
+    float2 fv     = half2_to_float2(v);
+    float2 rot_fv = rotary_embedding_transform(fv, coef);
+    return float2_to_half2(rot_fv);
+}
+
+#ifdef ENABLE_BF16
+inline __device__ __nv_bfloat162 rotary_embedding_transform(const __nv_bfloat162 v, const float2 coef)
+{
+    float2 fv     = bf1622float2(v);
+    float2 rot_fv = rotary_embedding_transform(fv, coef);
+    return __floats2bfloat162_rn(rot_fv.x, rot_fv.y);
+}
+#endif
+
+inline __device__ void apply_rotary_embedding(float& q, int zid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    return;
+}
+
+inline __device__ void apply_rotary_embedding(float& q, float& k, int zid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    return;
+}
+
+inline __device__ void apply_rotary_embedding(float2& q, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    if (2 * tid >= rot_embed_dim) {
+        return;
+    }
+    const auto coef = rotary_embedding_coefficient(2 * tid, rot_embed_dim, t_step, base, scale);
+    q               = rotary_embedding_transform(q, coef);
+}
+
+inline __device__ void apply_rotary_embedding(float2& q, float2& k, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    if (2 * tid >= rot_embed_dim) {
+        return;
+    }
+    const auto coef = rotary_embedding_coefficient(2 * tid, rot_embed_dim, t_step, base, scale);
+    q               = rotary_embedding_transform(q, coef);
+    k               = rotary_embedding_transform(k, coef);
+}
+
+inline __device__ void apply_rotary_embedding(float4& q, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    if (4 * tid >= rot_embed_dim) {
+        return;
+    }
+
+    Float4_&   q_    = *reinterpret_cast<Float4_*>(&q);
+    const auto coef0 = rotary_embedding_coefficient(4 * tid, rot_embed_dim, t_step, base, scale);
+    q_.x             = rotary_embedding_transform(q_.x, coef0);
+    const auto coef1 = rotary_embedding_coefficient(4 * tid + 2, rot_embed_dim, t_step, base, scale);
+    q_.y             = rotary_embedding_transform(q_.y, coef1);
+}
+
+inline __device__ void apply_rotary_embedding(float4& q, float4& k, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    if (4 * tid >= rot_embed_dim) {
+        return;
+    }
+
+    Float4_&   q_    = *reinterpret_cast<Float4_*>(&q);
+    Float4_&   k_    = *reinterpret_cast<Float4_*>(&k);
+    const auto coef0 = rotary_embedding_coefficient(4 * tid, rot_embed_dim, t_step, base, scale);
+    q_.x             = rotary_embedding_transform(q_.x, coef0);
+    k_.x             = rotary_embedding_transform(k_.x, coef0);
+    const auto coef1 = rotary_embedding_coefficient(4 * tid + 2, rot_embed_dim, t_step, base, scale);
+    q_.y             = rotary_embedding_transform(q_.y, coef1);
+    k_.y             = rotary_embedding_transform(k_.y, coef1);
+}
+
+inline __device__ void apply_rotary_embedding(uint32_t& q, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    if (2 * tid >= rot_embed_dim) {
+        return;
+    }
+    const auto coef = rotary_embedding_coefficient(2 * tid, rot_embed_dim, t_step, base, scale);
+    q               = rotary_embedding_transform(q, coef);
+}
+
+inline __device__ void apply_rotary_embedding(uint32_t& q, uint32_t& k, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    if (2 * tid >= rot_embed_dim) {
+        return;
+    }
+    const auto coef = rotary_embedding_coefficient(2 * tid, rot_embed_dim, t_step, base, scale);
+    q               = rotary_embedding_transform(q, coef);
+    k               = rotary_embedding_transform(k, coef);
+}
+
+inline __device__ void apply_rotary_embedding(uint2& q, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    if (4 * tid >= rot_embed_dim) {
+        return;
+    }
+    const auto coef0 = rotary_embedding_coefficient(4 * tid, rot_embed_dim, t_step, base, scale);
+    q.x              = rotary_embedding_transform(q.x, coef0);
+    const auto coef1 = rotary_embedding_coefficient(4 * tid + 2, rot_embed_dim, t_step, base, scale);
+    q.y              = rotary_embedding_transform(q.y, coef1);
+}
+
+inline __device__ void apply_rotary_embedding(uint2& q, uint2& k, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    if (4 * tid >= rot_embed_dim) {
+        return;
+    }
+    const auto coef0 = rotary_embedding_coefficient(4 * tid, rot_embed_dim, t_step, base, scale);
+    q.x              = rotary_embedding_transform(q.x, coef0);
+    k.x              = rotary_embedding_transform(k.x, coef0);
+    const auto coef1 = rotary_embedding_coefficient(4 * tid + 2, rot_embed_dim, t_step, base, scale);
+    q.y              = rotary_embedding_transform(q.y, coef1);
+    k.y              = rotary_embedding_transform(k.y, coef1);
+}
+
+inline __device__ void apply_rotary_embedding(uint4& q, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    if (8 * tid >= rot_embed_dim) {
+        return;
+    }
+    const auto coef0 = rotary_embedding_coefficient(8 * tid, rot_embed_dim, t_step, base, scale);
+    q.x              = rotary_embedding_transform(q.x, coef0);
+    const auto coef1 = rotary_embedding_coefficient(8 * tid + 2, rot_embed_dim, t_step, base, scale);
+    q.y              = rotary_embedding_transform(q.y, coef1);
+    const auto coef2 = rotary_embedding_coefficient(8 * tid + 4, rot_embed_dim, t_step, base, scale);
+    q.z              = rotary_embedding_transform(q.z, coef2);
+    const auto coef3 = rotary_embedding_coefficient(8 * tid + 6, rot_embed_dim, t_step, base, scale);
+    q.w              = rotary_embedding_transform(q.w, coef3);
+}
+
+inline __device__ void apply_rotary_embedding(uint4& q, uint4& k, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    if (8 * tid >= rot_embed_dim) {
+        return;
+    }
+    const auto coef0 = rotary_embedding_coefficient(8 * tid, rot_embed_dim, t_step, base, scale);
+    q.x              = rotary_embedding_transform(q.x, coef0);
+    k.x              = rotary_embedding_transform(k.x, coef0);
+    const auto coef1 = rotary_embedding_coefficient(8 * tid + 2, rot_embed_dim, t_step, base, scale);
+    q.y              = rotary_embedding_transform(q.y, coef1);
+    k.y              = rotary_embedding_transform(k.y, coef1);
+    const auto coef2 = rotary_embedding_coefficient(8 * tid + 4, rot_embed_dim, t_step, base, scale);
+    q.z              = rotary_embedding_transform(q.z, coef2);
+    k.z              = rotary_embedding_transform(k.z, coef2);
+    const auto coef3 = rotary_embedding_coefficient(8 * tid + 6, rot_embed_dim, t_step, base, scale);
+    q.w              = rotary_embedding_transform(q.w, coef3);
+    k.w              = rotary_embedding_transform(k.w, coef3);
+}
+
+#ifdef ENABLE_BF16
+inline __device__ void apply_rotary_embedding(__nv_bfloat162& q, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    if (2 * tid >= rot_embed_dim) {
+        return;
+    }
+    const auto coef = rotary_embedding_coefficient(2 * tid, rot_embed_dim, t_step, base, scale);
+    q               = rotary_embedding_transform(q, coef);
+}
+
+inline __device__ void
+apply_rotary_embedding(__nv_bfloat162& q, __nv_bfloat162& k, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    if (2 * tid >= rot_embed_dim) {
+        return;
+    }
+    const auto coef = rotary_embedding_coefficient(2 * tid, rot_embed_dim, t_step, base, scale);
+    q               = rotary_embedding_transform(q, coef);
+    k               = rotary_embedding_transform(k, coef);
+}
+
+inline __device__ void apply_rotary_embedding(bf16_4_t& q, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    if (4 * tid >= rot_embed_dim) {
+        return;
+    }
+    const auto coef0 = rotary_embedding_coefficient(4 * tid, rot_embed_dim, t_step, base, scale);
+    q.x              = rotary_embedding_transform(q.x, coef0);
+    const auto coef1 = rotary_embedding_coefficient(4 * tid + 2, rot_embed_dim, t_step, base, scale);
+    q.y              = rotary_embedding_transform(q.y, coef1);
+}
+
+inline __device__ void apply_rotary_embedding(bf16_4_t& q, bf16_4_t& k, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    if (4 * tid >= rot_embed_dim) {
+        return;
+    }
+    const auto coef0 = rotary_embedding_coefficient(4 * tid, rot_embed_dim, t_step, base, scale);
+    q.x              = rotary_embedding_transform(q.x, coef0);
+    k.x              = rotary_embedding_transform(k.x, coef0);
+    const auto coef1 = rotary_embedding_coefficient(4 * tid + 2, rot_embed_dim, t_step, base, scale);
+    q.y              = rotary_embedding_transform(q.y, coef1);
+    k.y              = rotary_embedding_transform(k.y, coef1);
+}
+
+inline __device__ void apply_rotary_embedding(bf16_8_t& q, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    if (8 * tid >= rot_embed_dim) {
+        return;
+    }
+    const auto coef0 = rotary_embedding_coefficient(8 * tid, rot_embed_dim, t_step, base, scale);
+    q.x              = rotary_embedding_transform(q.x, coef0);
+    const auto coef1 = rotary_embedding_coefficient(8 * tid + 2, rot_embed_dim, t_step, base, scale);
+    q.y              = rotary_embedding_transform(q.y, coef1);
+    const auto coef2 = rotary_embedding_coefficient(8 * tid + 4, rot_embed_dim, t_step, base, scale);
+    q.z              = rotary_embedding_transform(q.z, coef2);
+    const auto coef3 = rotary_embedding_coefficient(8 * tid + 6, rot_embed_dim, t_step, base, scale);
+    q.w              = rotary_embedding_transform(q.w, coef3);
+}
+
+inline __device__ void apply_rotary_embedding(bf16_8_t& q, bf16_8_t& k, int tid, int rot_embed_dim, int t_step, const float base=10000.0f, const float scale=1.0f)
+{
+    if (8 * tid >= rot_embed_dim) {
+        return;
+    }
+    const auto coef0 = rotary_embedding_coefficient(8 * tid, rot_embed_dim, t_step, base, scale);
+    q.x              = rotary_embedding_transform(q.x, coef0);
+    k.x              = rotary_embedding_transform(k.x, coef0);
+    const auto coef1 = rotary_embedding_coefficient(8 * tid + 2, rot_embed_dim, t_step, base, scale);
+    q.y              = rotary_embedding_transform(q.y, coef1);
+    k.y              = rotary_embedding_transform(k.y, coef1);
+    const auto coef2 = rotary_embedding_coefficient(8 * tid + 4, rot_embed_dim, t_step, base, scale);
+    q.z              = rotary_embedding_transform(q.z, coef2);
+    k.z              = rotary_embedding_transform(k.z, coef2);
+    const auto coef3 = rotary_embedding_coefficient(8 * tid + 6, rot_embed_dim, t_step, base, scale);
+    q.w              = rotary_embedding_transform(q.w, coef3);
+    k.w              = rotary_embedding_transform(k.w, coef3);
+}
+#endif  // ENABLE_BF16
+
+template<typename Vec_T, typename T>
+__device__ __inline__ void vec_from_smem_transpose(Vec_T& vec, T* smem, int transpose_idx, int smem_pitch);
+
+template<>
+__device__ __inline__ void vec_from_smem_transpose(float& vec, float* smem, int transpose_idx, int smem_pitch)
+{
+    return;
+}
+
+template<>
+__device__ __inline__ void vec_from_smem_transpose(uint32_t& vec, uint16_t* smem, int transpose_idx, int smem_pitch)
+{
+    union {
+        uint32_t u32;
+        uint16_t u16[2];
+    } tmp;
+    tmp.u16[0] = smem[transpose_idx];
+    tmp.u16[1] = smem[smem_pitch + transpose_idx];
+
+    vec = tmp.u32;
+}
+
+template<>
+__device__ __inline__ void vec_from_smem_transpose(uint2& vec, uint16_t* smem, int transpose_idx, int smem_pitch)
+{
+    union {
+        uint32_t u32;
+        uint16_t u16[2];
+    } tmp_1, tmp_2;
+    tmp_1.u32 = *reinterpret_cast<uint32_t*>(&smem[transpose_idx]);
+    tmp_2.u32 = *reinterpret_cast<uint32_t*>(&smem[smem_pitch + transpose_idx]);
+
+    union {
+        uint2    u32x2;
+        uint16_t u16[4];
+    } tmp_3;
+    tmp_3.u16[0] = tmp_1.u16[0];
+    tmp_3.u16[1] = tmp_2.u16[0];
+    tmp_3.u16[2] = tmp_1.u16[1];
+    tmp_3.u16[3] = tmp_2.u16[1];
+
+    vec = tmp_3.u32x2;
+}
+
+template<>
+__device__ __inline__ void vec_from_smem_transpose(uint4& vec, uint16_t* smem, int transpose_idx, int smem_pitch)
+{
+    union {
+        uint64_t u64;
+        uint16_t u16[4];
+    } tmp_1, tmp_2;
+    tmp_1.u64 = *reinterpret_cast<uint64_t*>(&smem[transpose_idx]);
+    tmp_2.u64 = *reinterpret_cast<uint64_t*>(&smem[smem_pitch + transpose_idx]);
+
+    union {
+        uint4    u32x4;
+        uint16_t u16[8];
+    } tmp_3;
+    tmp_3.u16[0] = tmp_1.u16[0];
+    tmp_3.u16[1] = tmp_2.u16[0];
+    tmp_3.u16[2] = tmp_1.u16[1];
+    tmp_3.u16[3] = tmp_2.u16[1];
+    tmp_3.u16[4] = tmp_1.u16[2];
+    tmp_3.u16[5] = tmp_2.u16[2];
+    tmp_3.u16[6] = tmp_1.u16[3];
+    tmp_3.u16[7] = tmp_2.u16[3];
+
+    vec = tmp_3.u32x4;
+}
+
+#ifdef ENABLE_BF16
+template<>
+__device__ __inline__ void
+vec_from_smem_transpose(bf16_4_t& vec, __nv_bfloat16* smem, int transpose_idx, int smem_pitch)
+{
+    union {
+        uint32_t      u32;
+        __nv_bfloat16 bf16[2];
+    } tmp_1, tmp_2;
+    tmp_1.u32 = *reinterpret_cast<uint32_t*>(&smem[transpose_idx]);
+    tmp_2.u32 = *reinterpret_cast<uint32_t*>(&smem[smem_pitch + transpose_idx]);
+
+    vec.x = __nv_bfloat162{tmp_1.bf16[0], tmp_2.bf16[0]};
+    vec.y = __nv_bfloat162{tmp_1.bf16[1], tmp_2.bf16[1]};
+}
+
+template<>
+__device__ __inline__ void
+vec_from_smem_transpose(bf16_8_t& vec, __nv_bfloat16* smem, int transpose_idx, int smem_pitch)
+{
+    union {
+        uint64_t      u64;
+        __nv_bfloat16 bf16[4];
+    } tmp_1, tmp_2;
+    tmp_1.u64 = *reinterpret_cast<uint64_t*>(&smem[transpose_idx]);
+    tmp_2.u64 = *reinterpret_cast<uint64_t*>(&smem[smem_pitch + transpose_idx]);
+
+    vec.x = __nv_bfloat162{tmp_1.bf16[0], tmp_2.bf16[0]};
+    vec.y = __nv_bfloat162{tmp_1.bf16[1], tmp_2.bf16[1]};
+    vec.z = __nv_bfloat162{tmp_1.bf16[2], tmp_2.bf16[2]};
+    vec.w = __nv_bfloat162{tmp_1.bf16[3], tmp_2.bf16[3]};
+}
+#endif  // ENABLE_BF16
+
+template<>
+__device__ __inline__ void vec_from_smem_transpose(float4& vec, float* smem, int transpose_idx, int smem_pitch)
+{
+    vec.x = smem[transpose_idx];
+    vec.z = smem[transpose_idx + 1];
+    vec.y = smem[smem_pitch + transpose_idx];
+    vec.w = smem[smem_pitch + transpose_idx + 1];
+}
+
+template<>
+__device__ __inline__ void vec_from_smem_transpose(uint32_t& vec, half* smem, int transpose_idx, int smem_pitch)
+{
+    union {
+        uint32_t u32;
+        half     u16[2];
+    } tmp;
+    tmp.u16[0] = smem[transpose_idx];
+    tmp.u16[1] = smem[smem_pitch + transpose_idx];
+
+    vec = tmp.u32;
+}
+
+#ifdef ENABLE_BF16
+template<>
+__device__ __inline__ void
+vec_from_smem_transpose(__nv_bfloat162& vec, __nv_bfloat16* smem, int transpose_idx, int smem_pitch)
+{
+    vec.x = smem[transpose_idx];
+    vec.y = smem[smem_pitch + transpose_idx];
+}
+#endif
+
+template<>
+__device__ __inline__ void vec_from_smem_transpose(float2& vec, float* smem, int transpose_idx, int smem_pitch)
+{
+    vec.x = smem[transpose_idx];
+    vec.y = smem[smem_pitch + transpose_idx];
+}
+
+template<typename Vec_T, typename T>
+__device__ __inline__ void write_smem_transpose(const Vec_T& vec, T* smem, int transpose_idx, int smem_pitch);
+
+template<>
+__device__ __inline__ void write_smem_transpose(const float& vec, float* smem, int transpose_idx, int smem_pitch)
+{
+    return;
+}
+
+template<>
+__device__ __inline__ void write_smem_transpose(const uint4& vec, uint16_t* smem, int transpose_idx, int smem_pitch)
+{
+    union {
+        uint64_t u64;
+        uint16_t u16[4];
+    } tmp_1, tmp_2;
+
+    union {
+        uint4    u32x4;
+        uint16_t u16[8];
+    } tmp_3;
+    tmp_3.u32x4  = vec;
+    tmp_1.u16[0] = tmp_3.u16[0];
+    tmp_2.u16[0] = tmp_3.u16[1];
+    tmp_1.u16[1] = tmp_3.u16[2];
+    tmp_2.u16[1] = tmp_3.u16[3];
+    tmp_1.u16[2] = tmp_3.u16[4];
+    tmp_2.u16[2] = tmp_3.u16[5];
+    tmp_1.u16[3] = tmp_3.u16[6];
+    tmp_2.u16[3] = tmp_3.u16[7];
+
+    *reinterpret_cast<uint64_t*>(&smem[transpose_idx])              = tmp_1.u64;
+    *reinterpret_cast<uint64_t*>(&smem[smem_pitch + transpose_idx]) = tmp_2.u64;
+}
+
+template<>
+__device__ __inline__ void write_smem_transpose(const uint2& vec, uint16_t* smem, int transpose_idx, int smem_pitch)
+{
+    union {
+        uint32_t u32;
+        uint16_t u16[2];
+    } tmp_1, tmp_2;
+
+    union {
+        uint2    u32x2;
+        uint16_t u16[4];
+    } tmp_3;
+    tmp_3.u32x2  = vec;
+    tmp_1.u16[0] = tmp_3.u16[0];
+    tmp_2.u16[0] = tmp_3.u16[1];
+    tmp_1.u16[1] = tmp_3.u16[2];
+    tmp_2.u16[1] = tmp_3.u16[3];
+
+    *reinterpret_cast<uint32_t*>(&smem[transpose_idx])              = tmp_1.u32;
+    *reinterpret_cast<uint32_t*>(&smem[smem_pitch + transpose_idx]) = tmp_2.u32;
+}
+
+template<>
+__device__ __inline__ void write_smem_transpose(const uint32_t& vec, uint16_t* smem, int transpose_idx, int smem_pitch)
+{
+    union {
+        uint32_t u32;
+        uint16_t u16[2];
+    } tmp;
+    tmp.u32 = vec;
+
+    smem[transpose_idx]              = tmp.u16[0];
+    smem[smem_pitch + transpose_idx] = tmp.u16[1];
+}
+
+template<>
+__device__ __inline__ void write_smem_transpose(const float4& vec, float* smem, int transpose_idx, int smem_pitch)
+{
+    smem[transpose_idx]                  = vec.x;
+    smem[transpose_idx + 1]              = vec.z;
+    smem[smem_pitch + transpose_idx]     = vec.y;
+    smem[smem_pitch + transpose_idx + 1] = vec.w;
+}
+
+template<>
+__device__ __inline__ void write_smem_transpose(const uint32_t& vec, half* smem, int transpose_idx, int smem_pitch)
+{
+    union {
+        uint32_t u32;
+        half     u16[2];
+    } tmp;
+
+    tmp.u32                          = vec;
+    smem[transpose_idx]              = tmp.u16[0];
+    smem[smem_pitch + transpose_idx] = tmp.u16[1];
+}
+
+#ifdef ENABLE_BF16
+template<>
+__device__ __inline__ void
+write_smem_transpose(const __nv_bfloat162& vec, __nv_bfloat16* smem, int transpose_idx, int smem_pitch)
+{
+    smem[transpose_idx]              = vec.x;
+    smem[smem_pitch + transpose_idx] = vec.y;
+}
+
+template<>
+__device__ __inline__ void
+write_smem_transpose(const bf16_4_t& vec, __nv_bfloat16* smem, int transpose_idx, int smem_pitch)
+{
+    write_smem_transpose(reinterpret_cast<const uint2&>(vec), reinterpret_cast<uint16_t*>(smem), transpose_idx, smem_pitch);
+}
+
+template<>
+__device__ __inline__ void
+write_smem_transpose(const bf16_8_t& vec, __nv_bfloat16* smem, int transpose_idx, int smem_pitch)
+{
+    write_smem_transpose(reinterpret_cast<const uint4&>(vec), reinterpret_cast<uint16_t*>(smem), transpose_idx, smem_pitch);
+}
+#endif
+
+template<>
+__device__ __inline__ void write_smem_transpose(const float2& vec, float* smem, int transpose_idx, int smem_pitch)
+{
+    smem[transpose_idx]              = vec.x;
+    smem[smem_pitch + transpose_idx] = vec.y;
+}
+
+}  // namespace mmha

llm-awq/awq/kernels/csrc/attention/ft_attention.cpp

@@ -0,0 +1,185 @@
+// Adapted from NVIDIA/FasterTransformer and FlashAttention
+
+#include <torch/extension.h>
+#include "ATen/cuda/CUDAContext.h"
+#include <c10/cuda/CUDAGuard.h>
+
+#include "ft_attention.h"
+#include "decoder_masked_multihead_attention.h"
+
+#define CHECK_DEVICE(x) TORCH_CHECK(x.device().type() == torch::kCUDA, #x " must be on CUDA")
+#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+
+#define DISPATCH_FLOAT_AND_HALF_AND_BF16(TYPE, NAME, ...)                  \
+  if (TYPE == at::ScalarType::Half) {                                      \
+    using scalar_t = at::Half;                                             \
+    __VA_ARGS__();                                                         \
+  } else if (TYPE == at::ScalarType::BFloat16) {                           \
+    using scalar_t = at::BFloat16;                                         \
+    __VA_ARGS__();                                                         \
+  } else if (TYPE == at::ScalarType::Float)  {                             \
+    using scalar_t = float;                                                \
+    __VA_ARGS__();                                                         \
+  } else {                                                                 \
+    AT_ERROR(#NAME, " not implemented for type '", toString(TYPE), "'"); \
+  }
+
+template<typename T>
+void masked_multihead_attention(const Masked_multihead_attention_params<T>& params,
+                                const cudaStream_t& stream);
+
+template<typename T>
+void cross_multihead_attention(const Masked_multihead_attention_params<T>& params,
+                               const cudaStream_t& stream);
+
+template<typename T>
+struct SATypeConverter {
+    using Type = T;
+};
+
+template<>
+struct SATypeConverter<at::Half> {
+    using Type = uint16_t;
+};
+
+template<>
+struct SATypeConverter<at::BFloat16> {
+    using Type = __nv_bfloat16;
+};
+
+template <typename T>
+void set_params(Masked_multihead_attention_params<T> &params,
+                const size_t batch_size,
+                const size_t nheads,
+                const size_t nheads_kv,
+                const size_t memory_max_seqlen,
+                const size_t headdim,
+                const int timestep,
+                const int rotary_embedding_dim,
+                const float rotary_base,
+                const float rotary_scale,
+                const bool neox_rotary_style,
+                const int qkv_batch_stride,
+                T *q_ptr,
+                T *k_ptr,
+                T *v_ptr,
+                T *k_cache_ptr,
+                T *v_cache_ptr,
+                int *length_per_sample,
+                float *alibi_slopes_ptr,
+                T *out_ptr) {
+    // Reset the parameters
+    memset(&params, 0, sizeof(params));
+    params.q = q_ptr;
+    params.k = k_ptr;
+    params.v = v_ptr;
+    params.q_bias = nullptr;
+    params.k_bias = nullptr;
+    params.v_bias = nullptr;
+    params.k_cache = k_cache_ptr;
+    params.v_cache = v_cache_ptr;
+    params.linear_bias_slopes = alibi_slopes_ptr;
+    params.out = out_ptr;
+    params.cache_indir = nullptr;
+    params.stride = qkv_batch_stride;
+    params.batch_size = batch_size;
+    params.beam_width = 1;
+    params.memory_max_len = memory_max_seqlen;
+    params.num_heads = nheads;
+    params.num_kv_heads = nheads_kv;
+    params.hidden_size_per_head = headdim;
+    params.rotary_embedding_dim = rotary_embedding_dim;
+    params.rotary_base = rotary_base;
+    params.rotary_scale = rotary_scale;
+    params.neox_rotary_style = neox_rotary_style;
+    params.timestep = timestep;
+    params.inv_sqrt_dh = 1.f / sqrt(float(headdim));
+    params.total_padding_tokens = nullptr;
+    params.masked_tokens = nullptr;
+    params.prefix_prompt_lengths = nullptr;
+    params.max_prefix_prompt_length = 0;
+    params.relative_attention_bias = nullptr;
+    params.relative_attention_bias_stride = 0;
+    params.cross_attention_out = nullptr;
+    params.max_decoder_seq_len = 0;
+    params.is_return_cross_attentions = false;
+    params.finished = nullptr;
+    params.memory_length_per_sample = nullptr;
+    params.length_per_sample = length_per_sample;
+}
+
+torch::Tensor single_query_attention(const torch::Tensor q,
+                                     const torch::Tensor k,
+                                     const torch::Tensor v,
+                                     torch::Tensor k_cache,
+                                     torch::Tensor v_cache,
+                                     c10::optional<const torch::Tensor> length_per_sample_,
+                                     c10::optional<const torch::Tensor> alibi_slopes_,
+                                     const int timestep,
+                                     const int rotary_embedding_dim,
+                                     const float rotary_base,
+                                     const float rotary_scale,
+                                     // neox_rotary_style = not interleaved
+                                     const bool neox_rotary_style) {
+    CHECK_DEVICE(q); CHECK_DEVICE(k); CHECK_DEVICE(v); CHECK_DEVICE(k_cache); CHECK_DEVICE(v_cache);
+    int batch_size = v_cache.size(0);
+    int nheads = q.size(1);
+    int nheads_kv = v_cache.size(1);
+    int memory_max_seqlen = v_cache.size(2);
+    int headdim = v_cache.size(3);
+    CHECK_SHAPE(q, batch_size, nheads, headdim);
+    CHECK_SHAPE(k, batch_size, nheads_kv, headdim);
+    CHECK_SHAPE(v, batch_size, nheads_kv, headdim);
+    CHECK_SHAPE(v_cache, batch_size, nheads_kv, memory_max_seqlen, headdim);
+    // k_cache shape: [B, H, Dh/x, L, x] where x=8 for fp16 and x=4 for fp32
+    int packsize = k_cache.dtype() == torch::kFloat32 ? 4 : 8;
+    CHECK_SHAPE(k_cache, batch_size, nheads_kv, headdim / packsize, memory_max_seqlen, packsize);
+    TORCH_CHECK(q.stride(2) == 1 && q.stride(1) == headdim);
+    TORCH_CHECK(k.stride(2) == 1 && k.stride(1) == headdim);
+    TORCH_CHECK(v.stride(2) == 1 && v.stride(1) == headdim);
+    // TORCH_CHECK(q.stride(0) == k.stride(0) && q.stride(0) == v.stride(0));
+    CHECK_CONTIGUOUS(v_cache); CHECK_CONTIGUOUS(k_cache);
+
+    if (length_per_sample_.has_value()) {
+        auto length_per_sample = length_per_sample_.value();
+        CHECK_DEVICE(length_per_sample);
+        CHECK_SHAPE(length_per_sample, batch_size);
+        CHECK_CONTIGUOUS(length_per_sample);
+        TORCH_CHECK(length_per_sample.dtype() == torch::kInt32);
+    }
+
+    if (alibi_slopes_.has_value()) {
+      auto alibi_slopes = alibi_slopes_.value();
+      CHECK_DEVICE(alibi_slopes);
+      CHECK_SHAPE(alibi_slopes, nheads);
+      CHECK_CONTIGUOUS(alibi_slopes); 
+      TORCH_CHECK(alibi_slopes.dtype() == torch::kFloat32);
+    }
+
+    // Otherwise the kernel will be launched from cuda:0 device
+    // Cast to char to avoid compiler warning about narrowing
+    at::cuda::CUDAGuard device_guard{(char)q.get_device()};
+
+    torch::Tensor out = torch::empty_like(q);
+
+    DISPATCH_FLOAT_AND_HALF_AND_BF16(q.scalar_type(), "single_query_attention", [&] {
+        using DataType = typename SATypeConverter<scalar_t>::Type;
+        Masked_multihead_attention_params<DataType> params;
+        set_params(params, batch_size, nheads, nheads_kv, memory_max_seqlen, headdim, 
+                   timestep, rotary_embedding_dim, rotary_base, rotary_scale, neox_rotary_style, q.stride(0),
+                   reinterpret_cast<DataType*>(q.data_ptr()),
+                   reinterpret_cast<DataType*>(k.data_ptr()),
+                   reinterpret_cast<DataType*>(v.data_ptr()),
+                   reinterpret_cast<DataType*>(k_cache.data_ptr()),
+                   reinterpret_cast<DataType*>(v_cache.data_ptr()),
+                   length_per_sample_.has_value()
+                       ? length_per_sample_.value().data_ptr<int>() : nullptr,
+                   alibi_slopes_.has_value() 
+                       ? alibi_slopes_.value().data_ptr<float>(): nullptr,
+                   reinterpret_cast<DataType*>(out.data_ptr()));
+        auto stream = at::cuda::getCurrentCUDAStream();
+        masked_multihead_attention(params, stream);
+    });
+    return out;
+}

llm-awq/awq/kernels/csrc/attention/ft_attention.h

@@ -0,0 +1,16 @@
+#pragma once
+#include <torch/extension.h>
+
+
+torch::Tensor single_query_attention(const torch::Tensor q,
+                                     const torch::Tensor k,
+                                     const torch::Tensor v,
+                                     torch::Tensor k_cache,
+                                     torch::Tensor v_cache,
+                                     c10::optional<const torch::Tensor> length_per_sample_,
+                                     c10::optional<const torch::Tensor> alibi_slopes_,
+                                     const int timestep,
+                                     const int rotary_embedding_dim = 0,
+                                     const float rotary_base = 10000.0f,
+                                     const float rotary_scale = 1.0f,
+                                     const bool neox_rotary_style=true);

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

@@ -0,0 +1,131 @@
+/*
+
+Adapted from NVIDIA FasterTransformer:
+https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/kernels/layernorm_kernels.cu
+
+*/
+
+#include <torch/extension.h>
+#include <cuda_fp16.h>
+#include "reduction.cuh"
+#include "layernorm.h"
+#include <cuda_runtime.h>
+#include <c10/cuda/CUDAGuard.h>
+
+#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());                                                      \
+  }
+
+static inline __device__ float to_float(half src)
+{
+    return __half2float(src);
+}
+
+static inline __device__ float to_float(float src)
+{
+    return src;
+}
+
+template<typename T>
+__global__ void generalT5LayerNorm(
+    const T* __restrict input, const T* __restrict gamma, T* output, const float layernorm_eps, int m, int n)
+{
+    // layernorm module in the T5 style No bias and no subtraction of mean.
+    const int tid = threadIdx.x;
+
+    __shared__ float s_variance;
+    float            variance = 0.0f;
+
+    float local_var_sum = 0.0f;
+    for (int i = tid; i < n; i += blockDim.x) {
+        float diff = to_float(__ldg(&input[blockIdx.x * n + i]));
+        local_var_sum += diff * diff;
+    }
+    variance = blockReduceSum(local_var_sum);
+
+    if (threadIdx.x == 0) {
+        s_variance = rsqrtf(variance / (float)n + layernorm_eps);
+    }
+    __syncthreads();
+
+    for (int i = tid; i < n; i += blockDim.x) {
+        output[blockIdx.x * n + i] =
+            clamp_inf_for_half<T>((to_float(input[blockIdx.x * n + i]) * s_variance) * to_float(__ldg(&gamma[i])));
+    }
+}
+
+
+template<typename T>
+void invokeGeneralT5LayerNorm(T*           out,
+                              const T*     input,
+                              const T*     gamma,
+                              // const T*     beta,
+                              const float  layernorm_eps,
+                              const int    m,
+                              const int    n)
+{
+    dim3 grid(m);
+    dim3 block(min(n, 1024));
+
+    /* For general cases, n is equal to hidden_units, e.g., 512/1024.
+        Since we have warp shuffle inside the code, block.x % 32 should be 0.
+    */
+    if (n % 32 != 0) {
+        block.x = 1024;
+    }
+
+    block.x = block.x / (4 / sizeof(T));  // if using half, only need half of block.x
+
+    /* should pay attention to the rsqrt precision*/
+    generalT5LayerNorm<T><<<grid, block>>>(input, gamma, out, layernorm_eps, m, n);  // For gpt-3
+}
+
+template void invokeGeneralT5LayerNorm(half*           out,
+                              const half*     input,
+                              const half*     gamma,
+                              // const half*     beta,
+                              const float  layernorm_eps,
+                              const int    m,
+                              const int    n);
+
+template void invokeGeneralT5LayerNorm(float*           out,
+                              const float*     input,
+                              const float*     gamma,
+                              // const half*     beta,
+                              const float  layernorm_eps,
+                              const int    m,
+                              const int    n);
+
+
+
+// input b, n, c
+void layernorm_forward_cuda(
+    torch::Tensor _input,
+    torch::Tensor _gamma,
+    torch::Tensor _out,
+    float eps)
+{
+    int m = _input.size(0) * _input.size(1);
+    int n = _input.size(2);
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(_input));
+
+    auto data_type = _input.scalar_type();
+    TORCH_CHECK(_gamma.scalar_type() == data_type);
+    TORCH_CHECK(_out.scalar_type() == data_type);
+
+    DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(data_type, ctype, {
+        auto input = reinterpret_cast<ctype*>(_input.data_ptr());
+        auto gamma = reinterpret_cast<ctype*>(_gamma.data_ptr());
+        auto out = reinterpret_cast<ctype*>(_out.data_ptr());
+        invokeGeneralT5LayerNorm(out, input, gamma, eps, m, n);
+    });
+}

llm-awq/awq/kernels/csrc/position_embedding/pos_encoding.h

@@ -0,0 +1,9 @@
+#pragma once
+#include <torch/extension.h>
+
+void rotary_embedding_neox(
+  torch::Tensor& positions,
+  torch::Tensor& query,
+  torch::Tensor& key,
+  int head_size,
+  torch::Tensor& cos_sin_cache);

llm-awq/awq/kernels/csrc/pybind.cpp

@@ -0,0 +1,38 @@
+#include <pybind11/pybind11.h>
+#include <torch/extension.h>
+#include "attention/ft_attention.h"
+#include "layernorm/layernorm.h"
+#include "quantization/gemm_cuda.h"
+#include "quantization/gemv_cuda.h"
+#include "quantization_new/gemm/gemm_cuda.h"
+#include "quantization_new/gemv/gemv_cuda.h"
+#include "position_embedding/pos_encoding.h"
+#include "rope_new/fused_rope_with_pos.h"
+#include "w8a8/w8a8_gemm_cuda.h"
+#include "w8a8/quantization.h"
+#include "w8a8/layernorm.h"
+#include "w8a8/act.h"
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("layernorm_forward_cuda", &layernorm_forward_cuda, "FasterTransformer layernorm kernel");
+    m.def("gemm_forward_cuda", &gemm_forward_cuda, "Quantized GEMM kernel.");
+    m.def("gemv_forward_cuda", &gemv_forward_cuda, "Quantized GEMV kernel.");
+    m.def("gemm_forward_cuda_new", &gemm_forward_cuda_new, "New quantized GEMM kernel.");
+    m.def("gemv_forward_cuda_new", &gemv_forward_cuda_new, "New quantized GEMV kernel.");
+    m.def("rotary_embedding_neox", &rotary_embedding_neox, "Apply GPT-NeoX style rotary embedding to query and key");
+    m.def("single_query_attention", &single_query_attention, "Attention with a single query",
+          py::arg("q"), py::arg("k"), py::arg("v"), py::arg("k_cache"), py::arg("v_cache"),
+          py::arg("length_per_sample_"), py::arg("alibi_slopes_"), py::arg("timestep"), py::arg("rotary_embedding_dim")=0,
+          py::arg("rotary_base")=10000.0f, py::arg("rotary_scale")=1.0f, py::arg("neox_rotary_style")=true);
+    m.def("fused_rope_with_pos_forward_func", &fused_rope_with_pos_forward_func,"Fused rope forward function with B,S,D embedding");
+    m.def("w8a8_gemm_forward_cuda", &w8a8_gemm_forward_cuda, "our w8a8 gemm kernel");
+    m.def("w8a8_gemm_fuse_bias_forward_cuda", &w8a8_gemm_fuse_bias_forward_cuda, "our w8a8 gemm fused bias kernel");
+    m.def("invoke_quant", &invoke_quant, "fp16->int8 quantization");
+    m.def("rms_norm_general", &rms_norm_general, py::arg("out"), py::arg("input"),
+        py::arg("weight"), py::arg("bias"),py::arg("scaling"), py::arg("epsilon"), py::arg("use_per_token_quant") = true,
+        "Apply Root Mean Square (RMS) Normalization to the input tensor (TRTLLM kernel).");
+    m.def("silu_and_mul", &silu_and_mul, "Activation function.");
+    m.def("gelu_and_quant",&gelu_and_quant, "Apply gelu act and quant output");
+}

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

@@ -0,0 +1,247 @@
+// Inspired by https://github.com/ankan-ban/llama_cu_awq
+/*
+
+@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"
+#define VECTORIZE_FACTOR 8
+#define Q_VECTORIZE_FACTOR 8
+#define PACK_FACTOR 8
+#define WARP_SIZE 32
+
+
+// Reduce sum within the warp using the tree reduction algorithm.
+__device__ __forceinline__ float warp_reduce_sum(float sum) {
+  #pragma unroll
+  for(int i = 4; i >= 0; i--){
+    sum += __shfl_down_sync(0xffffffff, sum, 1<<i);
+  }
+  /*
+  // Equivalent to the following tree reduction implementation:
+  sum += __shfl_down_sync(0xffffffff, sum, 16);
+  sum += __shfl_down_sync(0xffffffff, sum, 8);
+  sum += __shfl_down_sync(0xffffffff, sum, 4);
+  sum += __shfl_down_sync(0xffffffff, sum, 2);
+  sum += __shfl_down_sync(0xffffffff, sum, 1);
+  */
+  return sum;
+}
+
+__device__ __forceinline__ int make_divisible(int c, int divisor){
+  return (c + divisor - 1) / divisor;
+}
+
+
+/*
+Computes GEMV (group_size = 64).
+
+Args:
+  inputs: vector of shape [batch_size, IC];
+  weight: matrix of shape [OC, IC / 8];
+  output: vector of shape [OC];
+  zeros: matrix of shape [OC, IC / group_size / 8];
+  scaling_factors: matrix of shape [OC, IC / group_size];
+
+Notes:
+  One cannot infer group_size from the shape of scaling factors.
+  the second dimension is rounded up to a multiple of PACK_FACTOR.
+*/
+__global__ void gemv_kernel_g64(
+  const float4* _inputs, const uint32_t* weight, const uint32_t* zeros, const half* scaling_factors, half* _outputs, 
+  const int IC, const int OC){
+    const int group_size = 64;
+    float psum = 0;
+    const int batch_idx = blockIdx.z;
+    const int oc_idx = blockIdx.y * blockDim.y + threadIdx.y; 
+    const float4* inputs = _inputs + batch_idx * IC / PACK_FACTOR;
+    half* outputs = _outputs + batch_idx * OC;
+    // This is essentially zeros_w.
+    const int num_groups_packed = make_divisible(make_divisible(IC / group_size, PACK_FACTOR), 2) * 2;
+    const int weight_w = IC / PACK_FACTOR;
+    // TODO (Haotian): zeros_w is incorrect, after fixing we got misaligned address
+    const int zeros_w = make_divisible(make_divisible(IC / group_size, PACK_FACTOR), 2) * 2;
+    // consistent with input shape
+    const int sf_w = make_divisible(make_divisible(IC / group_size, PACK_FACTOR), 2) * 2 * PACK_FACTOR;
+    // if(blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 0 && threadIdx.y == 0) printf("%d %d %d %d %d\n", IC, group_size, PACK_FACTOR, zeros_w, sf_w);
+    // tile size: 4 OC x 1024 IC per iter
+    for(int packed_group_idx = 0; packed_group_idx < num_groups_packed / 2; packed_group_idx++){
+      // 1024 numbers in one iteration across warp. Need 1024 / group_size zeros.
+      uint64_t packed_zeros = *reinterpret_cast<const uint64_t*>(zeros + oc_idx * zeros_w + packed_group_idx * 2);
+      uint32_t packed_weights[4];
+      // use float4 to load weights, each thread load 32 int4 numbers (1 x float4)
+      *((float4*)(packed_weights)) = *((float4*)(weight + oc_idx * weight_w + packed_group_idx * (WARP_SIZE * 4) + threadIdx.x * 4));
+      // load scaling factors
+      // g64: two threads -> 64 numbers -> 1 group; 1 warp = 16 groups.
+      float scaling_factor = __half2float(scaling_factors[oc_idx * sf_w + packed_group_idx * 16 + (threadIdx.x / 2)]);
+      float current_zeros = (float)((packed_zeros >> (threadIdx.x / 2 * 4)) & 0xF);
+      int inputs_ptr_delta = packed_group_idx * WARP_SIZE * 4 + threadIdx.x * 4; 
+      const float4* inputs_ptr = inputs + inputs_ptr_delta;
+      // multiply 32 weights with 32 inputs
+      #pragma unroll
+      for (int ic_0 = 0; ic_0 < 4; ic_0++){
+        // iterate over different uint32_t packed_weights in this loop
+        uint32_t current_packed_weight = packed_weights[ic_0];
+        half packed_inputs[PACK_FACTOR];
+        // each thread load 8 inputs, starting index is packed_group_idx * 128 * 8 (because each iter loads 128*8)
+        if (inputs_ptr_delta + ic_0 < IC / PACK_FACTOR) {
+          *((float4*)packed_inputs) = *(inputs_ptr + ic_0);
+          #pragma unroll
+          for (int ic_1 = 0; ic_1 < PACK_FACTOR; ic_1++){
+            // iterate over 8 numbers packed within each uint32_t number
+            float current_single_weight_fp = (float)(current_packed_weight & 0xF);
+            float dequantized_weight = scaling_factor * (current_single_weight_fp - current_zeros);
+            //if(blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 0 && threadIdx.y == 0 && ic_0 == 0 && ic_1 == 0 && packed_group_idx == 0) printf("%f %f %f %f %X %X\n", dequantized_weight, current_single_weight_fp, scaling_factor, current_zeros, current_packed_weight, packed_zeros);
+            psum += dequantized_weight * __half2float(packed_inputs[ic_1]);
+            current_packed_weight = current_packed_weight >> 4;
+          }
+        }
+      }
+    }
+    psum = warp_reduce_sum(psum);
+    if (threadIdx.x == 0) {
+     outputs[oc_idx] = __float2half(psum); 
+    }
+}
+
+
+/*
+Computes GEMV (group_size = 128).
+
+Args:
+  inputs: vector of shape [batch_size, IC];
+  weight: matrix of shape [OC, IC / 8];
+  output: vector of shape [OC];
+  zeros: matrix of shape [OC, IC / group_size / 8];
+  scaling_factors: matrix of shape [OC, IC / group_size];
+
+Notes:
+  One cannot infer group_size from the shape of scaling factors.
+  the second dimension is rounded up to a multiple of PACK_FACTOR.
+*/
+__global__ void gemv_kernel_g128(
+  const float4* _inputs, const uint32_t* weight, const uint32_t* zeros, const half* scaling_factors, half* _outputs, 
+  const int IC, const int OC){
+    const int group_size = 128;
+    float psum = 0;
+    const int batch_idx = blockIdx.z;
+    const int oc_idx = blockIdx.y * blockDim.y + threadIdx.y; 
+    const float4* inputs = _inputs + batch_idx * IC / PACK_FACTOR;
+    half* outputs = _outputs + batch_idx * OC;
+    const int num_groups_packed = make_divisible(IC / group_size, PACK_FACTOR);
+    const int weight_w = IC / PACK_FACTOR;
+    // TODO (Haotian): zeros_w is incorrect, after fixing we got misaligned address
+    const int zeros_w = make_divisible(IC / group_size, PACK_FACTOR);
+    // consistent with input shape
+    const int sf_w = make_divisible(IC / group_size, PACK_FACTOR) * PACK_FACTOR;
+    //if(blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 0 && threadIdx.y == 0) printf("%d %d %d %d\n", IC, group_size, PACK_FACTOR, zeros_w);
+    // tile size: 4 OC x 1024 IC per iter
+    for(int packed_group_idx = 0; packed_group_idx < num_groups_packed; packed_group_idx++){
+      // 1024 numbers in one iteration across warp. Need 1024 / group_size zeros.
+      uint32_t packed_zeros = *(zeros + oc_idx * zeros_w + packed_group_idx);
+      uint32_t packed_weights[4];
+      // use float4 to load weights, each thread load 32 int4 numbers (1 x float4)
+      *((float4*)(packed_weights)) = *((float4*)(weight + oc_idx * weight_w + packed_group_idx * (WARP_SIZE * 4) + threadIdx.x * 4));
+      // load scaling factors
+      // g128: four threads -> 128 numbers -> 1 group; 1 warp = 8 groups.
+      float scaling_factor = __half2float(scaling_factors[oc_idx * sf_w + packed_group_idx * 8 + (threadIdx.x / 4)]);
+      float current_zeros = (float)((packed_zeros >> (threadIdx.x / 4 * 4)) & 0xF);
+      int inputs_ptr_delta = packed_group_idx * WARP_SIZE * 4 + threadIdx.x * 4; 
+      const float4* inputs_ptr = inputs + inputs_ptr_delta;
+      // multiply 32 weights with 32 inputs
+      #pragma unroll
+      for (int ic_0 = 0; ic_0 < 4; ic_0++){
+        // iterate over different uint32_t packed_weights in this loop
+        uint32_t current_packed_weight = packed_weights[ic_0];
+        half packed_inputs[PACK_FACTOR];
+        // each thread load 8 inputs, starting index is packed_group_idx * 128 * 8 (because each iter loads 128*8)
+        if (inputs_ptr_delta + ic_0 < IC / PACK_FACTOR) {
+          *((float4*)packed_inputs) = *(inputs_ptr + ic_0);
+          #pragma unroll
+          for (int ic_1 = 0; ic_1 < PACK_FACTOR; ic_1++){
+            // iterate over 8 numbers packed within each uint32_t number
+            float current_single_weight_fp = (float)(current_packed_weight & 0xF);
+            float dequantized_weight = scaling_factor * (current_single_weight_fp - current_zeros);
+            //if(blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 0 && threadIdx.y == 0 && ic_0 == 0 && ic_1 == 0 && packed_group_idx == 0) printf("%f %f %f %f %X %X\n", dequantized_weight, current_single_weight_fp, scaling_factor, current_zeros, current_packed_weight, packed_zeros);
+            psum += dequantized_weight * __half2float(packed_inputs[ic_1]);
+            current_packed_weight = current_packed_weight >> 4;
+          }
+        }
+      }
+    }
+    psum = warp_reduce_sum(psum);
+    if (threadIdx.x == 0) {
+     outputs[oc_idx] = __float2half(psum); 
+    }
+}
+
+
+/*
+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(
+    torch::Tensor _in_feats,
+    torch::Tensor _kernel,
+    torch::Tensor _scaling_factors,
+    torch::Tensor _zeros,
+    int group_size)
+{
+    int num_in_feats = _in_feats.size(0);
+    int num_in_channels = _in_feats.size(1);
+    // int kernel_volume = _out_in_map.size(1);
+    auto in_feats = reinterpret_cast<float4*>(_in_feats.data_ptr<at::Half>());
+    auto kernel = reinterpret_cast<uint32_t*>(_kernel.data_ptr<int>());
+    auto zeros = reinterpret_cast<uint32_t*>(_zeros.data_ptr<int>());
+    auto scaling_factors = reinterpret_cast<half*>(_scaling_factors.data_ptr<at::Half>());
+    // auto out_in_map = _out_in_map.data_ptr<int>();
+    auto options =
+    torch::TensorOptions().dtype(_in_feats.dtype()).device(_in_feats.device());
+    // kernel is [OC, IC]
+    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>());
+    int blockDim_z = num_out_feats;
+    dim3 num_blocks(1, num_out_channels / 4, num_out_feats);
+    dim3 num_threads(32, 4);
+    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
+      );
+    }
+    else if (group_size == 128)
+    {
+      gemv_kernel_g128<<<num_blocks, num_threads>>>(
+        // pointers
+        in_feats, kernel, zeros, scaling_factors, out_feats,
+        // constants
+        num_in_channels, num_out_channels
+      );
+    }
+    return _out_feats;
+;}
+

llm-awq/awq/kernels/setup.py

@@ -0,0 +1,51 @@
+from setuptools import find_packages, setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension
+
+
+extra_compile_args = {
+    "cxx": ["-g", "-O3", "-fopenmp", "-lgomp", "-std=c++17", "-DENABLE_BF16"],
+    "nvcc": [
+        "-O3",
+        "-std=c++17",
+        "-DENABLE_BF16",  # TODO
+        "-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",
+        "--threads=8",
+    ],
+}
+
+setup(
+    name="awq_inference_engine",
+    packages=find_packages(),
+    ext_modules=[
+        CUDAExtension(
+            name="awq_inference_engine",
+            sources=[
+                "csrc/pybind.cpp",
+                "csrc/quantization/gemm_cuda_gen.cu",
+                "csrc/quantization/gemv_cuda.cu",
+                "csrc/quantization_new/gemv/gemv_cuda.cu",
+                "csrc/quantization_new/gemm/gemm_cuda.cu",
+                "csrc/layernorm/layernorm.cu",
+                "csrc/position_embedding/pos_encoding_kernels.cu",
+                "csrc/attention/ft_attention.cpp",
+                "csrc/attention/decoder_masked_multihead_attention.cu",
+                "csrc/rope_new/fused_rope_with_pos.cu",
+                "csrc/w8a8/w8a8_gemm_cuda.cu",
+                "csrc/w8a8/quantization.cu",
+                "csrc/w8a8/act.cu",
+                "csrc/w8a8/layernorm.cu"
+            ],
+            extra_compile_args=extra_compile_args,
+        ),
+    ],
+    cmdclass={"build_ext": BuildExtension},
+    install_requires=["torch"],
+)

llm-awq/tinychat/models/falcon.py

@@ -0,0 +1,304 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This software may be used and distributed according to the terms of the GNU General Public License version 3.
+
+from typing import Optional, Tuple
+from dataclasses import dataclass
+import math
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import awq_inference_engine
+
+import tinychat.utils.constants
+
+max_batch_size = tinychat.utils.constants.max_batch_size
+max_seq_len = tinychat.utils.constants.max_seq_len
+
+
+# rotary pos emb helpers (torch.jit.script does not seem to support staticmethod...)
+def rotate_half(x):
+    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
+    return torch.cat(
+        (-x2, x1), dim=x1.ndim - 1
+    )  # dim=-1 triggers a bug in torch < 1.8.0
+
+
+class RotaryEmbedding(nn.Module):
+    """Implementation of RotaryEmbedding from GPT-NeoX.
+    This implementation is design to operate on queries and keys that are compatible with
+    [batch_size, n_heads_per_partition, seq_len, head_dim] (e.g. MinGPTAttention format).
+    """
+
+    def __init__(
+        self,
+        head_dim: int,
+        base=10000,
+    ):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.head_dim = head_dim
+        self.seq_len_cached = None
+        self.batch_size_cached = None
+        self.cos_cached: torch.Tensor | None = None
+        self.sin_cached: torch.Tensor | None = None
+
+    def cos_sin(
+        self,
+        seq_len: int,
+        device="cuda",
+        dtype=torch.bfloat16,
+    ) -> torch.Tensor:
+        if seq_len != self.seq_len_cached:
+            self.seq_len_cached = seq_len
+            t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
+            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1).to(device)
+
+            if dtype in [torch.float16, torch.bfloat16]:
+                emb = emb.float()
+
+            self.cos_cached = emb.cos()[None, :, :]
+            self.sin_cached = emb.sin()[None, :, :]
+
+            self.cos_cached = self.cos_cached.type(dtype)
+            self.sin_cached = self.sin_cached.type(dtype)
+
+        return self.cos_cached, self.sin_cached
+
+    def forward(self, _q, _k):
+        batch, seq_len, num_heads, head_dim = _q.shape
+        q = _q.permute(0, 2, 1, 3).contiguous().reshape(-1, seq_len, head_dim)
+        k = _k.permute(0, 2, 1, 3).contiguous().reshape(-1, seq_len, head_dim)
+        cos, sin = self.cos_sin(seq_len, q.device, q.dtype)
+        return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
+
+
+class FalconAttentionFused(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.n_local_heads = args.n_head
+        self.head_dim = args.hidden_size // args.n_head
+
+        self.query_key_value = nn.Linear(
+            args.hidden_size,
+            args.n_head * self.head_dim + 2 * self.head_dim,
+            bias=False,
+        )
+
+        self.dense = nn.Linear(
+            args.n_head * self.head_dim,
+            args.hidden_size,
+            bias=False,
+        )
+
+        # following fastertransformer definition
+
+        self.cache_v = (
+            torch.zeros(
+                (
+                    max_batch_size,
+                    1,
+                    max_seq_len,
+                    self.head_dim,
+                )
+            )
+            .cuda()
+            .half()
+        )  # added to half
+        # 8: pack 8 fp16 in FT, if fp32 then use 4
+        self.cache_k = (
+            torch.zeros(
+                (
+                    max_batch_size,
+                    1,
+                    self.head_dim // 8,
+                    max_seq_len,
+                    8,
+                )
+            )
+            .cuda()
+            .half()
+        )  # added to half
+
+        self.rotary_emb = RotaryEmbedding(self.head_dim)
+        self.rope_theta = args.rope_theta
+        self.rope_scaling = args.rope_scaling
+        if self.rope_scaling is None:
+            self.rope_scaling = 1.0
+        else:
+            self.rope_scaling = 1.0 / self.rope_scaling["factor"]
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        mask: Optional[torch.Tensor],
+    ):
+        bsz, seqlen, _ = x.shape
+
+        xqkv = self.query_key_value(x)
+        xqkv = xqkv.view(bsz, seqlen, self.n_local_heads + 2, self.head_dim)
+        xq = xqkv[:, :, :-2]
+        xk = xqkv[:, :, [-2]]
+        xv = xqkv[:, :, [-1]]
+
+        if seqlen > 1:
+            xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+            xk = xk.view(bsz, seqlen, 1, self.head_dim)
+            xv = xv.view(bsz, seqlen, 1, self.head_dim)
+
+            xq, xk = self.rotary_emb(xq, xk)
+            xq = (
+                xq.reshape(bsz, self.n_local_heads, seqlen, self.head_dim)
+                .permute(0, 2, 1, 3)
+                .contiguous()
+            )
+            xk = (
+                xk.reshape(bsz, 1, seqlen, self.head_dim)
+                .permute(0, 2, 1, 3)
+                .contiguous()
+            )
+
+            self.cache_k = self.cache_k.to(xq)
+            self.cache_v = self.cache_v.to(xq)
+
+            values_store = xv.transpose(2, 1)
+            keys_store = (
+                xk.reshape(bsz, seqlen, 1, self.head_dim // 8, 8)
+                .permute(0, 2, 3, 1, 4)
+                .contiguous()
+            )
+
+            self.cache_v[:bsz, :, start_pos : start_pos + seqlen, :] = values_store
+            self.cache_k[:bsz, :, :, start_pos : start_pos + seqlen, :] = keys_store
+
+            keys = xk
+            values = xv
+
+            xq = xq.transpose(1, 2)
+            keys = keys.transpose(1, 2)
+            values = values.transpose(1, 2)
+            scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
+            if mask is not None:
+                scores = scores + mask  # (bs, n_local_heads, slen, cache_len + slen)
+            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
+            output = torch.matmul(scores, values)  # (bs, n_local_heads, slen, head_dim)
+            output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
+        else:
+            # xq = xq[:, 0, :, :]
+            # xk = xk[:, 0, :, :]
+            # xv = xv[:, 0, :, :]
+            xq = xq.view(bsz, self.n_local_heads, self.head_dim)
+            xk = xk.view(bsz, 1, self.head_dim)
+            xv = xv.view(bsz, 1, self.head_dim)
+
+            output = awq_inference_engine.single_query_attention(
+                xq,
+                xk,
+                xv,
+                self.cache_k,
+                self.cache_v,
+                None,
+                # alibi position encodings
+                None,
+                start_pos,
+                self.head_dim,
+                self.rope_theta,
+                self.rope_scaling,
+                True,
+            )
+            output = output.reshape(bsz, 1, -1)
+
+        return self.dense(output)
+
+
+class FalconMLP(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+    ):
+        super().__init__()
+        self.dense_h_to_4h = nn.Linear(dim, 4 * dim, bias=False)
+        self.act = nn.GELU()
+        self.dense_4h_to_h = nn.Linear(4 * dim, dim, bias=False)
+
+    def forward(self, x):
+        x = self.act(self.dense_h_to_4h(x))
+        x = self.dense_4h_to_h(x)
+        return x
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, layer_id: int, args):
+        super().__init__()
+        self.n_heads = args.n_head
+        self.dim = args.hidden_size
+        self.head_dim = args.hidden_size // args.n_head
+        self.self_attention = FalconAttentionFused(args)
+        self.mlp = FalconMLP(dim=args.hidden_size)
+        self.layer_id = layer_id
+        self.input_layernorm = nn.LayerNorm(
+            args.hidden_size, eps=args.layer_norm_epsilon
+        )
+        # self.post_attention_layernorm = nn.LayerNorm(args.dim, eps=args.norm_eps)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        mask: Optional[torch.Tensor],
+    ):
+        layernorm_output = self.input_layernorm(x)
+        h_attn = x + self.self_attention.forward(layernorm_output, start_pos, mask)
+        h_mlp = self.mlp(layernorm_output)
+        out = h_attn + h_mlp
+        return out
+
+
+class Transformer(nn.Module):
+    def __init__(self, params):
+        super().__init__()
+        self.params = params
+        self.vocab_size = params.vocab_size
+        self.n_layers = params.n_layer
+
+        self.word_embeddings = nn.Embedding(params.vocab_size, params.hidden_size)
+
+        self.h = torch.nn.ModuleList()
+        for layer_id in range(params.n_layer):
+            self.h.append(TransformerBlock(layer_id, params))
+
+        self.ln_f = nn.LayerNorm(params.hidden_size, eps=params.layer_norm_epsilon)
+
+    @torch.inference_mode()
+    def forward(self, tokens: torch.Tensor, start_pos: int):
+        _bsz, seqlen = tokens.shape
+        h = self.word_embeddings(tokens)
+
+        mask = None
+        if seqlen > 1:
+            mask = torch.full(
+                (1, 1, seqlen, seqlen), float("-inf"), device=tokens.device
+            )
+            mask = torch.triu(mask, diagonal=start_pos + 1).type_as(h)
+        for layer in self.h:
+            h = layer(h, start_pos, mask)
+        h = self.ln_f(h)
+        return h
+
+
+class FalconForCausalLM(nn.Module):
+    def __init__(self, params):
+        super().__init__()
+        self.config = params
+        self.transformer = Transformer(params)
+        self.lm_head = nn.Linear(params.hidden_size, params.vocab_size, bias=False)
+
+    @torch.inference_mode()
+    def forward(self, tokens: torch.Tensor, start_pos: int):
+        h = self.transformer(tokens, start_pos)
+        output = self.lm_head(h)  # only compute last logits
+        return output.float()

llm-awq/tinychat/models/internvl3.py

@@ -0,0 +1,383 @@
+import os
+from collections import defaultdict, deque
+from typing import Dict, List, Optional, Tuple, Union, Any
+import warnings
+from time import time
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+import transformers
+from transformers import (AutoConfig, 
+                          AutoModel, 
+                          AutoTokenizer,
+                          GenerationConfig,
+                          PretrainedConfig, 
+                          PreTrainedModel)
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import ModelOutput, logging
+from transformers import modeling_utils
+
+from .internvl.configuration_internvl import InternVisionConfig, InternVLChatConfig
+from .internvl.internvit import InternVisionModel
+from .internvl.conversation import get_conv_template
+from .internvl.media import load_image, load_video
+
+from llava.media import Image, Video
+
+from .qwen2 import Qwen2ForCausalLM
+from .llama import LlamaForCausalLM
+
+try:
+    import flash_attn
+    has_flash_attn = True
+except ImportError:
+    print('FlashAttention2 is not installed.')
+    has_flash_attn = False
+
+def skip(*args, **kwargs):
+    pass
+
+torch.nn.init.kaiming_uniform_ = skip
+torch.nn.init.kaiming_normal_ = skip
+torch.nn.init.uniform_ = skip
+torch.nn.init.normal_ = skip
+
+modeling_utils._init_weights = False
+
+
+logger = logging.get_logger(__name__)
+
+
+class InternVL3(PreTrainedModel):
+    config_class = InternVLChatConfig
+    main_input_name = 'pixel_values'
+    base_model_prefix = 'language_model'
+    _supports_flash_attn_2 = True
+    supports_gradient_checkpointing = True
+    _no_split_modules = ['InternVisionModel', 'LlamaDecoderLayer', 'Qwen2DecoderLayer']
+        
+    def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None, use_flash_attn=True):
+        super().__init__(config)
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(config.name_or_path, trust_remote_code=True, use_fast=False)
+
+        image_size = config.force_image_size or config.vision_config.image_size
+        patch_size = config.vision_config.patch_size
+        self.patch_size = patch_size
+        self.select_layer = config.select_layer
+        self.template = config.template
+        self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2))
+        self.downsample_ratio = config.downsample_ratio
+        self.ps_version = config.ps_version
+        use_flash_attn = use_flash_attn if has_flash_attn else False
+        config.vision_config.use_flash_attn = True if use_flash_attn else False
+        config.llm_config._attn_implementation = 'flash_attention_2' if use_flash_attn else 'eager'
+
+        logger.info(f'num_image_token: {self.num_image_token}')
+        logger.info(f'ps_version: {self.ps_version}')
+        if vision_model is not None:
+            self.vision_model = vision_model
+        else:
+            self.vision_model = InternVisionModel(config.vision_config)
+        if language_model is not None:
+            self.language_model = language_model
+        else:
+            if config.llm_config.architectures[0] == 'LlamaForCausalLM':
+                self.language_model = LlamaForCausalLM(config.llm_config)
+            elif config.llm_config.architectures[0] == 'Qwen2ForCausalLM':
+                self.language_model = Qwen2ForCausalLM(config.llm_config)
+            else:
+                raise NotImplementedError(f'{config.llm_config.architectures[0]} is not implemented.')
+
+        vit_hidden_size = config.vision_config.hidden_size
+        llm_hidden_size = config.llm_config.hidden_size
+
+        self.mlp1 = nn.Sequential(
+            nn.LayerNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2),
+            nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size),
+            nn.GELU(),
+            nn.Linear(llm_hidden_size, llm_hidden_size)
+        )
+
+        self.img_context_token_id = None
+        self.conv_template = get_conv_template(self.template)
+        self.system_message = self.conv_template.system_message
+    
+    def freezed_module_patch(self):
+        self.vision_model.eval()
+        self.language_model.eval()
+        self.mlp1.eval()
+    
+    def pixel_shuffle(self, x, scale_factor=0.5):
+        n, w, h, c = x.size()
+        # N, W, H, C --> N, W, H * scale, C // scale
+        x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))
+        # N, W, H * scale, C // scale --> N, H * scale, W, C // scale
+        x = x.permute(0, 2, 1, 3).contiguous()
+        # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)
+        x = x.view(n, int(h * scale_factor), int(w * scale_factor),
+                   int(c / (scale_factor * scale_factor)))
+        if self.ps_version == 'v1':
+            warnings.warn("In ps_version 'v1', the height and width have not been swapped back, "
+                          'which results in a transposed image.')
+        else:
+            x = x.permute(0, 2, 1, 3).contiguous()
+        return x
+    
+    @torch.inference_mode()
+    def prepare_media(self, conversation):
+        prompt = conversation[0]["value"]
+        media = {"image": [], "video": []}
+        for item in prompt:
+            if isinstance(item, Image):
+                media["image"].append(load_image(item.path))
+            if isinstance(item, Video):
+                pixel_values, num_patches_list = load_video(item.path)
+                media["video"].extend(pixel_values)
+        
+        return media, num_patches_list if media["video"] else None
+        
+    @torch.inference_mode()
+    def extract_features(self, pixel_values):
+        if self.select_layer == -1:
+            vit_embeds = self.vision_model(
+                pixel_values=pixel_values,
+                output_hidden_states=False,
+                return_dict=True).last_hidden_state
+        else:
+            vit_embeds = self.vision_model(
+                pixel_values=pixel_values,
+                output_hidden_states=True,
+                return_dict=True).hidden_states[self.select_layer]
+        vit_embeds = vit_embeds[:, 1:, :]
+
+        h = w = int(vit_embeds.shape[1] ** 0.5)
+        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
+        vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)
+        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
+        vit_embeds = self.mlp1(vit_embeds)
+        return vit_embeds
+        
+    @torch.inference_mode()
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        media: Dict[str, List[torch.Tensor]],
+        media_config: Dict[str, Dict[str, Any]],
+        labels: Optional[torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+    ):
+        attention_mask = (
+            attention_mask
+            if attention_mask is not None
+            else torch.ones_like(input_ids, dtype=torch.bool)
+        )
+        
+        if media["image"]:
+            pixel_values = torch.cat(media["image"], dim=0).half().cuda()
+        elif media["video"]:
+            pixel_values = torch.cat(media["video"], dim=0).half().cuda()
+            
+        vit_embeds = self.extract_features(pixel_values)
+        
+        input_embeds = self.language_model.get_input_embeddings()(input_ids)
+        B, N, C = input_embeds.shape
+        input_embeds = input_embeds.reshape(B * N, C)
+        
+        input_ids = input_ids.reshape(B * N)
+        selected = (input_ids == self.img_context_token_id)
+        
+        input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)
+        
+        input_embeds = input_embeds.reshape(B, N, C)
+        
+        return input_embeds, None, attention_mask
+    
+    @torch.inference_mode()
+    def benchmark(self, prompt: Union[str, List], quant_llm) -> None:
+        media = {"image": [], "video": []}
+        question = ""
+        for item in prompt:
+            if isinstance(item, str):
+                question += item
+            if isinstance(item, Image):
+                media["image"].append(load_image(item.path))
+            if isinstance(item, Video):
+                pixel_values, num_patches_list = load_video(item.path)
+                media["video"].extend(pixel_values)
+        
+        if media["image"]:
+            num_patches_list = [image.size(0) for image in media["image"]]
+        
+        if media["image"] and '<image>' not in question:
+            question = '<image>\n' + question
+        
+        if media["video"] and '<image>' not in question:
+            video_prefix = ''.join([f'Frame{i+1}: <image>\n' for i in range(len(num_patches_list))])
+            question = video_prefix + question
+        
+        template = get_conv_template(self.template)
+        template.system_message = self.system_message
+        eos_token_id = self.tokenizer.convert_tokens_to_ids(template.sep.strip())
+        
+        template.append_message(template.roles[0], question)
+        template.append_message(template.roles[1], None)
+        query = template.get_prompt()
+        
+        IMG_START_TOKEN = '<img>'
+        IMG_END_TOKEN = '</img>'
+        IMG_CONTEXT_TOKEN = '<IMG_CONTEXT>'
+        
+        img_context_token_id = self.tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
+        self.img_context_token_id = img_context_token_id
+        
+        for num_patches in num_patches_list:
+            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
+            query = query.replace('<image>', image_tokens, 1)
+            
+        model_inputs = self.tokenizer(query, return_tensors='pt')
+        input_ids = model_inputs['input_ids'].to(self.device)
+        attention_mask = model_inputs['attention_mask'].to(self.device)
+        
+        for i in range(10):
+            torch.cuda.synchronize()
+            t_st = time()
+            inputs_embeds, _, attention_mask = self._embed(
+                input_ids=input_ids, 
+                media=media,
+                media_config=None,
+                labels=None,
+                attention_mask=attention_mask
+            )
+            torch.cuda.synchronize()
+            t_ed = time()
+            torch.cuda.empty_cache()
+            
+        if media["image"]:
+            print(
+                "Time of vision tower and others is {:.5f} s for {} images ({} x {} x {})".format(
+                    t_ed - t_st, sum(num_patches_list), media["image"][0].shape[1], media["image"][0].shape[2], media["image"][0].shape[3]
+                )
+            )
+        elif media["video"]:
+            print(
+                "Time of vision tower and others is {:.5f} s for {} video frames ({} x {} x {})".format(
+                    t_ed - t_st, sum(num_patches_list),  media["video"][0].shape[1], media["video"][0].shape[2], media["video"][0].shape[3]
+                )
+            )
+        output = self.language_model.benchmark(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            quant_llm=quant_llm
+        )
+        response = self.tokenizer.decode(output[0], skip_special_tokens=True).strip()
+        
+        return response
+
+    @torch.inference_mode()
+    def stream_gen(
+        self,
+        input_ids,
+        media,
+        media_cfg,
+        start_pos,
+        chunk_prefilling,
+        quant_llm,
+        attention_mask=None,
+    ) -> str:
+        if media is None:
+            inputs_embeds = self.language_model.get_input_embeddings()(input_ids ).clone()
+        else:
+            inputs_embeds, _, _ = self._embed(input_ids, media, None, None, attention_mask)
+    
+        length = inputs_embeds.shape[1]
+        if quant_llm:
+            out = self.language_model(None, start_pos, inputs_embeds, chunk_prefilling)
+        else:
+            out = self.language_model.forwardfp16(None, start_pos, inputs_embeds, chunk_prefilling)
+        return out, length
+
+    @torch.inference_mode()
+    def forward(
+            self,
+            pixel_values: torch.FloatTensor,
+            input_ids: torch.LongTensor = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            image_flags: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        image_flags = image_flags.squeeze(-1)
+        input_embeds = self.language_model.get_input_embeddings()(input_ids).clone()
+
+        vit_embeds = self.extract_feature(pixel_values)
+        vit_embeds = vit_embeds[image_flags == 1]
+        vit_batch_size = pixel_values.shape[0]
+
+        B, N, C = input_embeds.shape
+        input_embeds = input_embeds.reshape(B * N, C)
+
+        if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:
+            print(f'dynamic ViT batch size: {vit_batch_size}, images per sample: {vit_batch_size / B}, dynamic token length: {N}')
+
+        input_ids = input_ids.reshape(B * N)
+        selected = (input_ids == self.img_context_token_id)
+        try:
+            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)
+        except Exception as e:
+            vit_embeds = vit_embeds.reshape(-1, C)
+            print(f'warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, '
+                  f'vit_embeds.shape={vit_embeds.shape}')
+            n_token = min(selected.sum(), vit_embeds.size(0))
+            input_embeds[selected][:n_token] = input_embeds[selected][:n_token] * 0.0 + vit_embeds[:n_token]
+
+        input_embeds = input_embeds.reshape(B, N, C)
+
+        outputs = self.language_model(
+            inputs_embeds=input_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        logits = outputs.logits
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+    

llm-awq/tinychat/models/nvila/builder.py

@@ -0,0 +1,291 @@
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# 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.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+import os
+import os.path as osp
+import warnings
+from dataclasses import asdict
+from typing import Tuple
+
+import torch
+from huggingface_hub import file_exists, repo_exists
+from huggingface_hub.utils import HFValidationError
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    PretrainedConfig,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+)
+
+from llava.constants import MEDIA_TOKENS
+from llava.model.utils import packing
+from llava.utils.logging import logger
+from llava.utils.tokenizer import infer_stop_tokens
+
+
+def has_tokenizer(repo_id_or_path: str) -> bool:
+    # Check if the tokenizer is in a local directory
+    if osp.exists(osp.join(repo_id_or_path, "tokenizer_config.json")):
+        return True
+
+    # Check if the tokenizer is in a Hugging Face Hub repo
+    try:
+        return repo_exists(repo_id_or_path) and file_exists(
+            repo_id_or_path, "tokenizer_config.json"
+        )
+    except HFValidationError:
+        return False
+
+
+def context_length_extension(config):
+    orig_ctx_len = getattr(config, "max_position_embeddings", None)
+    model_max_length = getattr(config, "model_max_length", None)
+    if orig_ctx_len and model_max_length > orig_ctx_len:
+        print(f"Scaling RoPE from {orig_ctx_len} to {model_max_length}")
+        scaling_factor = float(math.ceil(model_max_length / orig_ctx_len))
+        config.rope_scaling = {"type": "linear", "factor": scaling_factor}
+    return config
+
+
+def build_llm_and_tokenizer(
+    model_name_or_path: str,
+    config: PretrainedConfig,
+    attn_implementation=None,
+    model_max_length=None,
+    *args,
+    **kwargs,
+) -> Tuple[PreTrainedModel, PreTrainedTokenizer]:
+    # print(model_name_or_path)
+    llm_cfg = AutoConfig.from_pretrained(model_name_or_path)
+    llm_cfg._attn_implementation = attn_implementation
+    llm_cfg.model_max_length = model_max_length
+    if model_max_length is not None:
+        context_length_extension(llm_cfg)
+
+    # Quantization related
+    quantization_restore_from_checkpoint = False
+    if kwargs.get("quantize_model_class") is not None:
+        assert kwargs.get("model_args") is not None
+        quantize_model_class = kwargs.pop("quantize_model_class", None)
+        model_args = kwargs.pop("model_args", None)
+
+        if (
+            quantize_model_class == "QLlamaForCausalLM"
+        ):  # TODO: Also change the name of this class
+            from .qllama import QLlamaConfig
+
+            llm_cfg.architectures = "QLlamaForCausalLM"
+            _attn_implementation = llm_cfg._attn_implementation
+            llm_cfg = QLlamaConfig(**llm_cfg.to_dict())
+            llm_cfg._attn_implementation = _attn_implementation
+        elif (
+            quantize_model_class == "QMemLlamaForCausalLM"
+        ):  # TODO: Also change the name of this class
+            from .qmemllama import QMemLlamaConfig
+
+            llm_cfg.architectures = "QMemLlamaForCausalLM"
+            llm_cfg = QMemLlamaConfig(**llm_cfg.to_dict())
+        elif quantize_model_class == "FP8LinearQwen2ForCausalLM":
+            from .configuration_quantize import QuantizationConfig
+            from .fp8linearqwen2 import FP8LinearQwen2Config
+
+            llm_cfg.architectures = "FP8LinearQwen2ForCausalLM"
+            coat_fp8_args = QuantizationConfig(**asdict(model_args))
+
+            # Remove the quantization args from llm_cfg and make it a independent config
+            model_args_dict = asdict(model_args)
+            for key in asdict(coat_fp8_args).keys():
+                model_args_dict.pop(key, None)
+
+            llm_cfg.coat_fp8_args = asdict(coat_fp8_args)
+            _attn_implementation = llm_cfg._attn_implementation
+
+            llm_cfg = FP8LinearQwen2Config(**llm_cfg.to_dict())
+            llm_cfg._attn_implementation = _attn_implementation
+
+        elif quantize_model_class == "FP8ActivationQwen2ForCausalLM":
+            from ..coat.activation.models._fp8_quantization_config import (
+                QuantizationConfig,
+            )
+            from .fp8activationqwen2 import FP8ActivationQwen2Config
+
+            quantization_restore_from_checkpoint = True
+
+            llm_cfg.architectures = "FP8ActivationQwen2ForCausalLM"
+            coat_fp8_args = QuantizationConfig(**asdict(model_args))
+
+            # Remove the quantization args from llm_cfg and make it a independent config
+            model_args_dict = asdict(model_args)
+            for key in asdict(coat_fp8_args).keys():
+                model_args_dict.pop(key, None)
+
+            llm_cfg.coat_fp8_args = asdict(coat_fp8_args)
+            _attn_implementation = llm_cfg._attn_implementation
+
+            llm_cfg = FP8ActivationQwen2Config(**llm_cfg.to_dict())
+            llm_cfg._attn_implementation = _attn_implementation
+
+        elif quantize_model_class == "FP8ActivationResidualQwen2ForCausalLM":
+            from ..coat.activation.models._fp8_quantization_config import (
+                QuantizationConfig,
+            )
+            from .fp8activationresidualqwen2 import FP8ActivationResidualQwen2Config
+
+            quantization_restore_from_checkpoint = True
+
+            llm_cfg.architectures = "FP8ActivationResidualQwen2ForCausalLM"
+            coat_fp8_args = QuantizationConfig(**asdict(model_args))
+
+            # Remove the quantization args from llm_cfg and make it a independent config
+            model_args_dict = asdict(model_args)
+            for key in asdict(coat_fp8_args).keys():
+                model_args_dict.pop(key, None)
+
+            llm_cfg.coat_fp8_args = asdict(coat_fp8_args)
+            _attn_implementation = llm_cfg._attn_implementation
+
+            llm_cfg = FP8ActivationResidualQwen2Config(**llm_cfg.to_dict())
+            llm_cfg._attn_implementation = _attn_implementation
+        else:
+            raise ValueError(
+                f"{quantize_model_class} is not supported quantize_model_class."
+            )
+
+        kwargs.pop("quantize_model_class", None)
+
+        if quantize_model_class in [
+            "FP8LinearQwen2ForCausalLM",
+            "FP8ActivationQwen2ForCausalLM",
+            "FP8ActivationResidualQwen2ForCausalLM",
+        ]:  # Remove the quantization args from llm_cfg and make it a independent config
+            llm_cfg.update(model_args_dict)
+        else:
+            llm_cfg.update(asdict(model_args))
+        # print(model_args)
+
+    if quantization_restore_from_checkpoint:
+        fp8_model_name_or_path = kwargs.pop("fp8_llm_cfg", None)
+
+        llm = AutoModelForCausalLM.from_pretrained(
+            fp8_model_name_or_path,
+            config=llm_cfg,
+            torch_dtype=eval(config.model_dtype),
+            *args,
+            **kwargs,
+        )
+
+    else:
+        llm = AutoModelForCausalLM.from_pretrained(
+            model_name_or_path,
+            config=llm_cfg,
+            torch_dtype=eval(config.model_dtype),
+            *args,
+            **kwargs,
+        )
+    packing.patch(llm)
+
+    # Locate the tokenizer.
+    llm_path = model_name_or_path
+    if not has_tokenizer(llm_path):
+        llm_path = osp.join(llm_path, "llm")
+    if not has_tokenizer(llm_path):
+        raise ValueError(f"Cannot find tokenizer in {llm_path}.")
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        llm_path, padding_side="right", use_fast=False, legacy=False
+    )
+    if model_max_length is not None:
+        tokenizer.model_max_length = model_max_length
+
+    # Load chat template if specified.
+    if getattr(config, "chat_template", None) is not None:
+        logger.info(f"Using chat template: {config.chat_template}")
+        fpath = os.path.join(
+            os.path.dirname(__file__), "chat_templates", f"{config.chat_template}.jinja"
+        )
+        with open(fpath) as fd:
+            chat_template = fd.read()
+        tokenizer.chat_template = chat_template.replace("    ", "").replace("\n", "")
+
+    # Set stop tokens for the tokenizer
+    tokenizer.stop_tokens = infer_stop_tokens(tokenizer)
+    tokenizer.stop_token_ids = tokenizer.convert_tokens_to_ids(tokenizer.stop_tokens)
+
+    # Add media tokens to the tokenizer
+    tokenizer.media_tokens = MEDIA_TOKENS
+    tokenizer.media_token_ids = {}
+    for name, token in MEDIA_TOKENS.items():
+        tokenizer.add_tokens([token], special_tokens=True)
+        tokenizer.media_token_ids[name] = tokenizer.convert_tokens_to_ids(token)
+
+    # TODO(ligeng): is this necessary for llava?
+    config.hidden_size = llm.config.hidden_size
+    return llm, tokenizer
+
+
+def build_tokenizer(
+    model_name_or_path: str,
+    config: PretrainedConfig,
+    attn_implementation=None,
+    model_max_length=None,
+    *args,
+    **kwargs,
+) -> Tuple[PreTrainedModel, PreTrainedTokenizer]:
+    # print(model_name_or_path)
+    llm_cfg = AutoConfig.from_pretrained(model_name_or_path)
+    llm_cfg._attn_implementation = attn_implementation
+    llm_cfg.model_max_length = model_max_length
+    if model_max_length is not None:
+        context_length_extension(llm_cfg)
+
+    # Locate the tokenizer.
+    llm_path = model_name_or_path
+    if not has_tokenizer(llm_path):
+        llm_path = osp.join(llm_path, "llm")
+    if not has_tokenizer(llm_path):
+        raise ValueError(f"Cannot find tokenizer in {llm_path}.")
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        llm_path, padding_side="right", use_fast=False, legacy=False
+    )
+    if model_max_length is not None:
+        tokenizer.model_max_length = model_max_length
+
+    # Load chat template if specified.
+    if getattr(config, "chat_template", None) is not None:
+        logger.info(f"Using chat template: {config.chat_template}")
+        fpath = os.path.join(
+            os.path.dirname(__file__), "chat_templates", f"{config.chat_template}.jinja"
+        )
+        with open(fpath) as fd:
+            chat_template = fd.read()
+        tokenizer.chat_template = chat_template.replace("    ", "").replace("\n", "")
+
+    # Set stop tokens for the tokenizer
+    tokenizer.stop_tokens = infer_stop_tokens(tokenizer)
+    tokenizer.stop_token_ids = tokenizer.convert_tokens_to_ids(tokenizer.stop_tokens)
+
+    # Add media tokens to the tokenizer
+    tokenizer.media_tokens = MEDIA_TOKENS
+    tokenizer.media_token_ids = {}
+    for name, token in MEDIA_TOKENS.items():
+        tokenizer.add_tokens([token], special_tokens=True)
+        tokenizer.media_token_ids[name] = tokenizer.convert_tokens_to_ids(token)
+
+    return tokenizer

llm-awq/tinychat/models/nvila/llava_arch.py

@@ -0,0 +1,909 @@
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+import copy
+import json
+import logging
+import os
+import os.path as osp
+import warnings
+from abc import ABC
+from collections import OrderedDict, defaultdict, deque
+from itertools import chain
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from einops import rearrange
+from hydra.utils import instantiate
+from transformers import AutoConfig, GenerationConfig, PreTrainedModel
+from transformers.modeling_utils import ContextManagers, no_init_weights
+from time import time
+from llava.constants import DEFAULT_IMAGE_TOKEN, IGNORE_INDEX
+from llava.mm_utils import process_image, process_images
+from llava.model.configuration_llava import LlavaConfig
+from llava.model.language_model.builder import build_llm_and_tokenizer
+from llava.model.multimodal_encoder.builder import build_vision_tower
+from llava.model.multimodal_projector.builder import build_mm_projector
+from llava.model.utils import get_model_config
+
+# from llava.train.sequence_parallel import get_pg_manager
+from llava.utils import distributed as dist
+from llava.utils.media import extract_media
+from llava.utils.tokenizer import tokenize_conversation
+from .builder import build_tokenizer
+
+
+class LlavaMetaModel(ABC):
+    def init_vlm(self, config, *args, **kwargs):
+        # TODO(ligeng): figure out how from_config and from_pretrained works in HF implementation.
+        if (
+            hasattr(self, "llm")
+            or hasattr(self, "vision_tower")
+            or hasattr(self, "mm_projector")
+        ):
+            # already initialized, skipped
+            return
+
+        model_dtype = getattr(config, "model_dtype", "torch.float16")
+        if not hasattr(config, "model_dtype"):
+            warnings.warn(
+                "model_dtype not found in config, defaulting to torch.float16."
+            )
+            config.model_dtype = model_dtype
+
+        cfgs = get_model_config(config)
+        if len(cfgs) == 3:
+            self.llm_cfg, vision_tower_cfg, mm_projector_cfg = cfgs
+        else:
+            raise ValueError(
+                "`llm_cfg` `mm_projector_cfg` `vision_tower_cfg` not found in the config."
+            )
+        self.tokenizer = build_tokenizer(self.llm_cfg, config, *args, **kwargs)
+        self.vision_tower = build_vision_tower(vision_tower_cfg, config)
+        self.mm_projector = build_mm_projector(mm_projector_cfg, config)
+
+        self.encoders = {}
+        for name in ["image", "video"]:
+            config = getattr(self.config, f"{name}_encoder")
+            if isinstance(config, str):
+                config = json.loads(config)
+            self.encoders[name] = instantiate(config, parent=self)
+
+        self.post_config()
+        self.is_loaded = True
+
+        assert (
+            self.vision_tower is not None or self.mm_projector is not None
+        ), "At least one of the components must be instantiated."
+
+    @classmethod
+    def load_from_config(cls, model_path_or_config, *args, **kwargs):
+        pass
+
+    ## FIXME we will use this function to load model in the future
+    @classmethod
+    def load_pretrained(cls, model_path_or_config, *args, **kwargs):
+        kwargs.pop("config", None)
+
+        if isinstance(model_path_or_config, str):
+            config = AutoConfig.from_pretrained(model_path_or_config)
+        elif isinstance(model_path_or_config, LlavaConfig):
+            config = model_path_or_config
+        else:
+            raise NotImplementedError(
+                f"wrong type, {type(model_path_or_config)} \
+                                      {isinstance(model_path_or_config, LlavaConfig)}"
+            )
+
+        model_dtype = getattr(config, "model_dtype", "torch.float16")
+        if not hasattr(config, "model_dtype"):
+            warnings.warn(
+                "model_dtype not found in config, defaulting to torch.float16."
+            )
+            config.model_dtype = model_dtype
+
+        cfgs = get_model_config(config)
+        if len(cfgs) == 3:
+            llm_cfg, vision_tower_cfg, mm_projector_cfg = cfgs
+        else:
+            raise ValueError(
+                "`llm_cfg` `mm_projector_cfg` `vision_tower_cfg` not found in the config."
+            )
+
+        # print(llm_cfg, vision_tower_cfg, mm_projector_cfg); input("DEBUG load_pretrained")
+        init_context = [
+            no_init_weights(_enable=True),
+        ]
+        # print("Before Init Context")
+        # if hasattr(config, "deepspeed") and "mics" in config.deepspeed:
+        #     print("Using MiCS_Init")
+        #     import deepspeed
+        #     init_context.append(deepspeed.zero.MiCS_Init(config_dict_or_path=config.deepspeed))
+        with ContextManagers(init_context):
+            vlm = cls(config, *args, **kwargs)
+        # print(llm_cfg, vision_tower_cfg, mm_projector_cfg); input("DEBUG load_pretrained finish")
+
+        if (
+            hasattr(vlm, "llm")
+            or hasattr(vlm, "vision_tower")
+            or hasattr(vlm, "mm_projector")
+        ):
+            if vlm.is_loaded:
+                return vlm
+
+        vlm.llm, vlm.tokenizer = build_llm_and_tokenizer(
+            llm_cfg, config, *args, **kwargs
+        )
+        vlm.vision_tower = build_vision_tower(vision_tower_cfg, config)
+        vlm.mm_projector = build_mm_projector(mm_projector_cfg, config)
+
+        self.post_config()
+        self.is_loaded = True
+
+        # FIXME(ligeng, yunhao): llm should never be none here.
+        assert (
+            vlm.llm is not None
+            or vlm.vision_tower is not None
+            or vlm.mm_projector is not None
+        ), "At least one of the components must be instantiated."
+        return vlm
+
+    ## FIXME we will use this function to save the model in the future
+    def save_pretrained(self, output_dir, state_dict=None):
+        if state_dict is None:
+            # other wise fetch from deepspeed
+            # state_dict = accelerator.get_state_dict(is_deepspeed_enabled)
+            state_dict = self.state_dict()
+
+        if getattr(self, "tokenizer", None):
+            self.tokenizer.save_pretrained(osp.join(output_dir, "llm"))
+
+        if self.get_llm():
+            print(f"saving llm to {osp.join(output_dir, 'llm')}")
+            self.llm.config._name_or_path = osp.join(output_dir, "llm")
+            llm_state_dict = OrderedDict(
+                {k.split("llm.")[-1]: v for k, v in state_dict.items() if "llm" in k}
+            )
+            self.llm.save_pretrained(
+                os.path.join(output_dir, "llm"), state_dict=llm_state_dict
+            )
+            self.config.llm_cfg = self.llm.config
+
+        if self.get_vision_tower():
+            print(f"saving vision_tower to {osp.join(output_dir, 'vision_tower')}")
+            self.vision_tower.config._name_or_path = osp.join(
+                output_dir, "vision_tower"
+            )
+            vision_tower_state_dict = OrderedDict(
+                {
+                    k.split("vision_tower.vision_tower.")[-1]: v
+                    for k, v in state_dict.items()
+                    if "vision_tower" in k
+                }
+            )
+            self.vision_tower.vision_tower.save_pretrained(
+                os.path.join(output_dir, "vision_tower"),
+                state_dict=vision_tower_state_dict,
+            )
+            self.vision_tower.image_processor.save_pretrained(
+                os.path.join(output_dir, "vision_tower")
+            )
+            self.config.vision_tower_cfg = self.vision_tower.config
+            if hasattr(self.config.vision_tower_cfg, "auto_map"):
+                if "radio" not in self.get_vision_tower().__class__.__name__.lower():
+                    delattr(self.config.vision_tower_cfg, "auto_map")
+
+        if self.get_mm_projector():
+            print(f"saving mm_projector to {osp.join(output_dir, 'mm_projector')}")
+            self.mm_projector.config._name_or_path = osp.join(
+                output_dir, "mm_projector"
+            )
+            mm_projector_state_dict = OrderedDict(
+                {
+                    k.split("mm_projector.")[-1]: v
+                    for k, v in state_dict.items()
+                    if "mm_projector" in k
+                }
+            )
+            self.mm_projector.save_pretrained(
+                os.path.join(output_dir, "mm_projector"),
+                state_dict=mm_projector_state_dict,
+            )
+            self.config.mm_projector_cfg = self.mm_projector.config
+        ## update and save top-level config
+        self.config._name_or_path = output_dir
+        self.config.architectures = [self.__class__.__name__]
+        self.config.save_pretrained(output_dir)
+
+    def get_llm(self):
+        llm = getattr(self, "llm", None)
+        if type(llm) is list:
+            llm = llm[0]
+        return llm
+
+    def get_lm_head(self):
+        lm_head = getattr(self.get_llm(), "lm_head", None)
+        return lm_head
+
+    def get_vision_tower(self):
+        vision_tower = getattr(self, "vision_tower", None)
+        if type(vision_tower) is list:
+            vision_tower = vision_tower[0]
+        return vision_tower
+
+    def get_mm_projector(self):
+        mm_projector = getattr(self, "mm_projector", None)
+        if type(mm_projector) is list:
+            mm_projector = mm_projector[0]
+        return mm_projector
+
+    def post_config(self):
+
+        if getattr(self.config, "vision_tower_cfg", None) is None:
+            self.config.vision_tower_cfg = self.vision_tower.config
+        if getattr(self.config, "mm_projector_cfg", None) is None:
+            self.config.mm_projector_cfg = self.mm_projector.config
+
+    @staticmethod
+    def merge_chessboard(x, num_split_h, num_split_w):
+        """
+        x: b * n * c or b * h * w * c
+        out: b * c * h * w
+        Assuming x contains num_split**2 sub-squares concatenated along batch dimension, merge the sub-squares back to the original whole square.
+        """
+        B = x.shape[0]
+        if x.dim() == 3:
+            N = x.shape[1]
+            x = rearrange(x, "b (h w) c -> b c h w", h=int(N**0.5), w=int(N**0.5))
+
+        assert B % (num_split_h * num_split_w) == 0
+        b = B // (num_split_h * num_split_w)
+
+        x_merge = torch.cat(
+            [
+                torch.cat(
+                    [
+                        x[(i * num_split_w + j) * b : (i * num_split_w + j + 1) * b]
+                        for j in range(num_split_w)
+                    ],
+                    dim=-1,
+                )
+                for i in range(num_split_h)
+            ],
+            dim=-2,
+        )
+
+        return x_merge
+
+    @staticmethod
+    def split_chessboard(x, num_split_h, num_split_w):
+        """
+        x: b * c * h * w
+        out: b * c * h * w
+        Deividing x into num_split**2 sub-squares, and concatenate all the sub-squares on the batch dimension
+        """
+        B, C, H, W = x.shape
+        assert H % num_split_h == 0 and W % num_split_w == 0
+        h, w = H // num_split_h, W // num_split_w
+        x_split = torch.cat(
+            [
+                x[:, :, i * h : (i + 1) * h, j * w : (j + 1) * w]
+                for i in range(num_split_h)
+                for j in range(num_split_w)
+            ],
+            dim=0,
+        )
+        return x_split
+
+    def merge_features_for_dynamic_s2(self, image_features, block_sizes):
+        scales = self.get_vision_tower().scales
+        resize_output_to_scale_idx = self.get_vision_tower().resize_output_to_scale_idx
+
+        image_features_each_image = []
+        new_block_sizes = []
+        block_cnt = 0
+        for block_size_each_image in block_sizes:
+            if block_size_each_image is None:
+                cur_features = image_features[block_cnt : block_cnt + 1]
+                cur_features = rearrange(
+                    cur_features,
+                    "1 (h w) c -> 1 c h w",
+                    h=int(cur_features.shape[1] ** 0.5),
+                )
+                cur_features = cur_features.repeat(1, len(scales), 1, 1)
+                image_features_each_image.append(cur_features)
+                new_block_sizes.append((1, 1))
+                block_cnt += 1
+            else:
+                cur_features_each_scale = []
+                for scale in scales[:-1]:
+                    num_blocks_this_scale = (scale // scales[0]) ** 2
+                    cur_features_each_scale.append(
+                        self.merge_chessboard(
+                            image_features[
+                                block_cnt : block_cnt + num_blocks_this_scale
+                            ],
+                            num_split_h=scale // scales[0],
+                            num_split_w=scale // scales[0],
+                        )
+                    )  # 1 * C * H * W
+                    block_cnt += num_blocks_this_scale
+                num_blocks_last_scale = (
+                    block_size_each_image[0] * block_size_each_image[1]
+                )
+                cur_features_each_scale.append(
+                    self.merge_chessboard(
+                        image_features[block_cnt : block_cnt + num_blocks_last_scale],
+                        num_split_h=block_size_each_image[0],
+                        num_split_w=block_size_each_image[1],
+                    )
+                )  # 1 * C * H * W
+                block_cnt += num_blocks_last_scale
+
+                # resize and concat features from different scales
+                output_size = cur_features_each_scale[resize_output_to_scale_idx].shape[
+                    -2:
+                ]
+                cur_features = torch.cat(
+                    [
+                        F.interpolate(
+                            cur_features_each_scale[i].to(torch.float32),
+                            size=output_size,
+                            mode="area",
+                        ).to(cur_features_each_scale[i].dtype)
+                        for i in range(len(cur_features_each_scale))
+                    ],
+                    dim=1,
+                )
+                # cur_features = rearrange(cur_features, "1 c h w -> (h w) c")
+
+                image_features_each_image.append(cur_features)
+
+                if (
+                    resize_output_to_scale_idx == len(scales) - 1
+                    or resize_output_to_scale_idx == -1
+                ):
+                    new_block_sizes.append(block_size_each_image)
+                else:
+                    new_block_sizes.append(
+                        (
+                            scales[resize_output_to_scale_idx] // scales[0],
+                            scales[resize_output_to_scale_idx] // scales[0],
+                        )
+                    )
+
+        assert block_cnt == len(image_features)
+
+        return image_features_each_image, new_block_sizes
+
+    def encode_images(
+        self, images, block_sizes: Optional[Optional[Tuple[int, ...]]] = None
+    ):
+        if block_sizes is None:
+            block_sizes = [None] * len(images)
+        if getattr(self.config, "dynamic_s2", False):
+            image_features = self.get_vision_tower()(images)
+            image_features, new_block_sizes = self.merge_features_for_dynamic_s2(
+                image_features, block_sizes
+            )
+
+            image_features = [
+                self.split_chessboard(x, block_size[0], block_size[1])
+                for x, block_size in zip(image_features, new_block_sizes)
+            ]  # list of B * C * H * W tensors
+            image_features = torch.cat(
+                [rearrange(x, "b c h w -> b (h w) c") for x in image_features], dim=0
+            )  # B * N * C
+            image_features = self.get_mm_projector()(image_features)
+            image_features = list(
+                image_features.split(
+                    [block_size[0] * block_size[1] for block_size in new_block_sizes],
+                    dim=0,
+                )
+            )
+            image_features = [
+                self.merge_chessboard(x, block_size[0], block_size[1])
+                for x, block_size in zip(image_features, new_block_sizes)
+            ]  # list of 1 * C * H * W tensors
+            image_features = [
+                rearrange(x, "1 c h w -> (h w) c") for x in image_features
+            ]  # list of N * C tensors
+            image_features = torch.stack(image_features, dim=0)
+        else:
+            image_features = self.get_vision_tower()(images)
+            image_features = self.get_mm_projector()(image_features)
+        return image_features
+
+    ## @yunhao: is there a better way to handle function call and attributes for llm?
+    ## support beam search
+    def _temporary_reorder_cache(self, past_key_values, sorted_idx):
+        return self.get_llm()._temporary_reorder_cache(past_key_values, sorted_idx)
+
+    def get_input_embeddings(self):
+        return self.get_llm().get_input_embeddings()
+
+    def get_output_embeddings(self):
+        return self.get_llm().get_output_embeddings()
+
+    def resize_token_embeddings(self, embed_size):
+        self.get_llm().resize_token_embeddings(embed_size)
+
+
+class LlavaMetaForCausalLM(ABC):
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        media: Dict[str, List[torch.Tensor]],
+        media_config: Dict[str, Dict[str, Any]],
+        labels: Optional[torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        labels = (
+            labels if labels is not None else torch.full_like(input_ids, IGNORE_INDEX)
+        )
+        attention_mask = (
+            attention_mask
+            if attention_mask is not None
+            else torch.ones_like(input_ids, dtype=torch.bool)
+        )
+
+        # Extract text and media embeddings
+        text_embeds = self.llm.model.embed_tokens(input_ids)
+        media_embeds = self.__embed_media_tokens(media, media_config)
+
+        # This is a workaround to make sure the dummy embeddings are consumed
+        while media_embeds.get("dummy"):
+            dummy_embed = media_embeds["dummy"].popleft()
+            text_embeds += torch.sum(dummy_embed) * 0
+        # Remove padding
+        batch_size = labels.shape[0]
+        text_embeds = [text_embeds[k][attention_mask[k]] for k in range(batch_size)]
+        labels = [labels[k][attention_mask[k]] for k in range(batch_size)]
+
+        # Build inverse mapping from token ID to media name
+        media_tokens = {}
+        for name, token_id in self.tokenizer.media_token_ids.items():
+            media_tokens[token_id] = name
+
+        # Fuse text and media embeddings
+        inputs_m, labels_m = [], []
+        for k in range(batch_size):
+            inputs_mk, labels_mk = [], []
+            pos = 0
+            while pos < len(labels[k]):
+                if input_ids[k][pos].item() in media_tokens:
+                    end = pos + 1
+                    name = media_tokens[input_ids[k][pos].item()]
+                    input = media_embeds[name].popleft()
+                    label = torch.full(
+                        [input.shape[0]],
+                        IGNORE_INDEX,
+                        device=labels[k].device,
+                        dtype=labels[k].dtype,
+                    )
+                else:
+                    end = pos
+                    while (
+                        end < len(labels[k])
+                        and input_ids[k][end].item() not in media_tokens
+                    ):
+                        end += 1
+                    input = text_embeds[k][pos:end]
+                    label = labels[k][pos:end]
+                inputs_mk.append(input)
+                labels_mk.append(label)
+                pos = end
+            inputs_m.append(torch.cat(inputs_mk, dim=0))
+            labels_m.append(torch.cat(labels_mk, dim=0))
+        inputs, labels = inputs_m, labels_m
+
+        # Check if all media embeddings are consumed
+        for name in media_embeds:
+            if media_embeds[name]:
+                raise ValueError(f"Not all {name} embeddings are consumed!")
+
+        # Truncate sequences to `model_max_length` as media embeddings are inserted
+        inputs, labels = self.__truncate_sequence(inputs, labels)
+
+        # Pad sequences to the longest one in the batch
+        return self.__batchify_sequence(inputs, labels)
+
+    def __embed_media_tokens(
+        self,
+        media: Dict[str, List[torch.Tensor]],
+        media_config: Dict[str, Dict[str, Any]],
+    ) -> Dict[str, List[torch.Tensor]]:
+        embeds = defaultdict(deque)
+        for name in media:
+            embeds[name] = deque(self.encoders[name](media[name], media_config[name]))
+        return embeds
+
+    def __truncate_sequence(
+        self, inputs: List[torch.Tensor], labels: List[torch.Tensor]
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if any(len(input) > self.tokenizer.model_max_length for input in inputs):
+            warnings.warn(
+                f"Truncating sequences to `model_max_length` ({self.tokenizer.model_max_length})."
+            )
+            inputs = [input[: self.tokenizer.model_max_length] for input in inputs]
+            labels = [label[: self.tokenizer.model_max_length] for label in labels]
+        return inputs, labels
+
+    def __batchify_sequence(
+        self, inputs: List[torch.Tensor], labels: List[torch.Tensor]
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        batch_size = len(inputs)
+        device = inputs[0].device
+        hidden_size = inputs[0].shape[1]
+        max_length = max(inputs[k].shape[0] for k in range(batch_size))
+        attention_mask = torch.ones(
+            (batch_size, max_length), dtype=torch.bool, device=device
+        )
+
+        inputs_p, labels_p = [], []
+        for k in range(batch_size):
+            size_pk = max_length - inputs[k].shape[0]
+            inputs_pk = torch.zeros(
+                (size_pk, hidden_size), dtype=inputs[k].dtype, device=device
+            )
+            labels_pk = torch.full(
+                (size_pk,), IGNORE_INDEX, dtype=labels[k].dtype, device=device
+            )
+            if self.tokenizer.padding_side == "right":
+                attention_mask[k, inputs[k].shape[0] :] = False
+                inputs_pk = torch.cat([inputs[k], inputs_pk], dim=0)
+                labels_pk = torch.cat([labels[k], labels_pk], dim=0)
+            else:
+                attention_mask[k, : -inputs[k].shape[0]] = False
+                inputs_pk = torch.cat([inputs_pk, inputs[k]], dim=0)
+                labels_pk = torch.cat([labels_pk, labels[k]], dim=0)
+            inputs_p.append(inputs_pk)
+            labels_p.append(labels_pk)
+
+        inputs = torch.stack(inputs_p, dim=0)
+        labels = torch.stack(labels_p, dim=0)
+        return inputs, labels, attention_mask
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        input_ids: Optional[torch.FloatTensor] = None,
+        media: Optional[Dict[str, List[torch.Tensor]]] = None,
+        media_config: Dict[str, Dict[str, Any]] = None,
+        attention_mask: Optional[torch.LongTensor] = None,
+        quant_llm: Optional[bool] = True,
+        **generation_kwargs,
+    ):
+        inputs_embeds, _, attention_mask = self._embed(
+            input_ids, media, media_config, None, attention_mask
+        )
+        return self.llm.generate(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            quant_llm=quant_llm,
+            **generation_kwargs,
+        )
+
+    @torch.inference_mode()
+    def generate_content(
+        self,
+        prompt: Union[str, List],
+        generation_config: Optional[GenerationConfig] = None,
+        quant_llm: Optional[bool] = True,
+    ) -> str:
+        # TODO(zhijianl): Support directly taking conversation as input
+        conversation = [{"from": "human", "value": prompt}]
+
+        # Extract media from the conversation
+
+        # TODO (extract and preprocess should be done together, as the preprocess of image and video can be different, i.e. when dynamic res is used)
+        media = extract_media(conversation, self.config)
+
+        # Process media
+        media_config = defaultdict(dict)
+        for name in media:
+            if name == "image":
+                if len(media["image"]) == 1 and self.config.image_aspect_ratio in [
+                    "dynamic",
+                    "dynamic_s2",
+                ]:
+                    self.config.image_processor = self.vision_tower.image_processor
+                    if self.config.image_aspect_ratio == "dynamic":
+                        images = process_image(
+                            media["image"][0],
+                            self.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:
+                        if type(self.config.s2_scales) is str:
+                            self.config.s2_scales = list(
+                                map(int, self.config.s2_scales.split(","))
+                            )
+                        images, block_sizes = process_image(
+                            media["image"][0], self.config, None, enable_dynamic_s2=True
+                        )
+                        images = images.half()
+                        media_config[name]["block_sizes"] = [block_sizes]
+                else:
+                    images = process_images(
+                        media["image"], self.vision_tower.image_processor, self.config
+                    ).half()
+                media[name] = [image for image in images]
+            elif name == "video":
+                media[name] = [
+                    process_images(
+                        images, self.vision_tower.image_processor, self.config
+                    ).half()
+                    for images in media[name]
+                ]
+            else:
+                raise ValueError(f"Unsupported media type: {name}")
+
+        # Tokenize the conversation
+        input_ids = (
+            tokenize_conversation(
+                conversation, self.tokenizer, add_generation_prompt=True
+            )
+            .cuda()
+            .unsqueeze(0)
+        )
+
+        # Set up the generation config
+        generation_config = generation_config or self.default_generation_config
+        # Generate the response
+        try:
+            output_ids = self.generate(
+                input_ids=input_ids,
+                media=media,
+                media_config=media_config,
+                generation_config=generation_config,
+                quant_llm=quant_llm,
+            )
+        except ValueError:
+            if not generation_config.do_sample:
+                raise
+            # FIXME(zhijianl): This is a temporary workaround for the sampling issue
+            logging.warning(
+                "Generation failed with sampling, retrying with greedy decoding."
+            )
+            generation_config.do_sample = False
+            output_ids = self.generate(
+                input_ids=input_ids,
+                media=media,
+                media_config=media_config,
+                generation_config=generation_config,
+            )
+
+        # Decode the response
+        response = self.tokenizer.decode(
+            output_ids[0], skip_special_tokens=True
+        ).strip()
+        return response
+
+    @torch.inference_mode()
+    def benchmark(self, prompt: Union[str, List], quant_llm) -> None:
+        # TODO(zhijianl): Support directly taking conversation as input
+        conversation = [{"from": "human", "value": prompt}]
+
+        # Extract media from the conversation
+
+        # TODO (extract and preprocess should be done together, as the preprocess of image and video can be different, i.e. when dynamic res is used)
+        media = extract_media(conversation, self.config)
+
+        # Process media
+        media_config = defaultdict(dict)
+        image_num = 0
+        for name in media:
+            if name == "image":
+                if len(media["image"]) == 1 and self.config.image_aspect_ratio in [
+                    "dynamic",
+                    "dynamic_s2",
+                ]:
+                    self.config.image_processor = self.vision_tower.image_processor
+                    if self.config.image_aspect_ratio == "dynamic":
+                        images = process_image(
+                            media["image"][0],
+                            self.config,
+                            None,
+                            enable_dynamic_res=True,
+                        ).half()
+                        if len(images.shape) == 3:
+                            images = images.reshape(1, *images.shape)
+                        image_num += images.shape[0]
+                        size = images.shape[1:]
+                        conversation[0]["value"] = conversation[0]["value"].replace(
+                            DEFAULT_IMAGE_TOKEN,
+                            f"{DEFAULT_IMAGE_TOKEN}\n" * images.shape[0],
+                        )
+                    else:
+                        if type(self.config.s2_scales) is str:
+                            self.config.s2_scales = list(
+                                map(int, self.config.s2_scales.split(","))
+                            )
+                        images, block_sizes = process_image(
+                            media["image"][0], self.config, None, enable_dynamic_s2=True
+                        )
+                        images = images.half()
+                        if len(images.shape) == 3:
+                            images = images.reshape(1, *images.shape)
+                        image_num += images.shape[0]
+                        size = images.shape[1:]
+                        media_config[name]["block_sizes"] = [block_sizes]
+                else:
+                    images = process_images(
+                        media["image"], self.vision_tower.image_processor, self.config
+                    ).half()
+                    image_num += images.shape[0]
+                    size = images.shape[1:]
+                media[name] = [image for image in images]
+            elif name == "video":
+                media[name] = [
+                    process_images(
+                        images, self.vision_tower.image_processor, self.config
+                    ).half()
+                    for images in media[name]
+                ]
+                for images in media[name]:
+                    image_num += images.shape[0]
+                    size = images.shape[1:]
+            else:
+                raise ValueError(f"Unsupported media type: {name}")
+
+        # Tokenize the conversation
+        input_ids = (
+            tokenize_conversation(
+                conversation, self.tokenizer, add_generation_prompt=True
+            )
+            .cuda()
+            .unsqueeze(0)
+        )
+
+        # Set up the generation config
+        for i in range(10):
+            torch.cuda.synchronize()
+            t_st = time()
+            inputs_embeds, _, attention_mask = self._embed(
+                input_ids, media, media_config, None, None
+            )
+            torch.cuda.synchronize()
+            t_ed = time()
+            torch.cuda.empty_cache()
+        print(
+            "Time of vision tower and others is {:.5f} s for {} images ({} x {} x {})".format(
+                t_ed - t_st, image_num, size[0], size[1], size[2]
+            )
+        )
+        output = self.llm.benchmark(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            quant_llm=quant_llm,
+        )
+        # response = self.tokenizer.decode(output, skip_special_tokens=True).strip()
+        return
+
+    @property
+    def default_generation_config(self) -> GenerationConfig:
+        generation_config = copy.deepcopy(self.generation_config or GenerationConfig())
+        if self.tokenizer.eos_token_id is None:
+            raise ValueError("Tokenizer must have an EOS token")
+        if generation_config.max_length == GenerationConfig().max_length:
+            generation_config.max_length = self.tokenizer.model_max_length
+        if generation_config.pad_token_id is None:
+            generation_config.pad_token_id = (
+                self.tokenizer.pad_token_id or self.tokenizer.eos_token_id
+            )
+        if generation_config.bos_token_id is None:
+            generation_config.bos_token_id = (
+                self.tokenizer.bos_token_id or self.tokenizer.eos_token_id
+            )
+        if generation_config.eos_token_id is None:
+            generation_config.eos_token_id = self.tokenizer.stop_token_ids
+        return generation_config
+
+    # Prepare media
+
+    # Process media
+    @torch.inference_mode()
+    def prepare_media(self, conversation):
+        media = extract_media(conversation, self.config)
+
+        # Process media
+        media_config = defaultdict(dict)
+        for name in media:
+            if name == "image":
+                if len(media["image"]) == 1 and self.config.image_aspect_ratio in [
+                    "dynamic",
+                    "dynamic_s2",
+                ]:
+                    self.config.image_processor = self.vision_tower.image_processor
+                    if self.config.image_aspect_ratio == "dynamic":
+                        images = process_image(
+                            media["image"][0],
+                            self.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:
+                        if type(self.config.s2_scales) is str:
+                            self.config.s2_scales = list(
+                                map(int, self.config.s2_scales.split(","))
+                            )
+                        images, block_sizes = process_image(
+                            media["image"][0], self.config, None, enable_dynamic_s2=True
+                        )
+                        images = images.half()
+                        media_config[name]["block_sizes"] = [block_sizes]
+                else:
+                    images = process_images(
+                        media["image"], self.vision_tower.image_processor, self.config
+                    ).half()
+                media[name] = [image for image in images]
+            elif name == "video":
+                media[name] = [
+                    process_images(
+                        images, self.vision_tower.image_processor, self.config
+                    ).half()
+                    for images in media[name]
+                ]
+            else:
+                raise ValueError(f"Unsupported media type: {name}")
+        return media, media_config
+
+    @torch.inference_mode()
+    def stream_gen(
+        self,
+        input_ids,
+        media,
+        media_cfg,
+        start_pos,
+        chunk_prefilling,
+        quant_llm,
+        attention_mask=None,
+    ) -> str:
+        if media is None:
+            inputs_embeds = self.llm.model.embed_tokens(input_ids)
+        else:
+            image_num = torch.sum(input_ids == 151649)
+            if image_num == 1 and self.config.image_aspect_ratio == "dynamic":
+                patch_num = len(media["image"])
+                new_input_ids = []
+                for i, id in enumerate(input_ids[0]):
+                    if id == 151649:
+                        new_input_ids.extend(input_ids[0, 0:i])
+                        new_input_ids.extend([198, 151649, 198] * patch_num)
+                        new_input_ids.extend(input_ids[0, i + 1 :])
+                        break
+                input_ids = torch.tensor(
+                    [new_input_ids], dtype=torch.int, device="cuda"
+                )
+            inputs_embeds, _, _ = self._embed(
+                input_ids, media, media_cfg, None, attention_mask=None
+            )
+        length = inputs_embeds.shape[1]
+        if quant_llm:
+            out = self.llm(None, start_pos, inputs_embeds, chunk_prefilling)
+        else:
+            out = self.llm.forwardfp16(None, start_pos, inputs_embeds, chunk_prefilling)
+        return out, length

llm-awq/tinychat/models/qwen2.py

@@ -0,0 +1,511 @@
+# Modified from https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2/modeling_qwen2.py
+"""PyTorch Qwen2 model."""
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+import awq_inference_engine
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+from transformers.activations import ACT2FN
+import tinychat
+import torch.nn.functional as F
+import time
+from tqdm import tqdm
+from transformers import GenerationMixin
+from transformers.models.qwen2 import Qwen2ForCausalLM
+from flash_attn import flash_attn_func
+
+max_batch_size = tinychat.utils.constants.max_batch_size
+max_seq_len = tinychat.utils.constants.max_seq_len
+
+
+class Qwen2RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = torch.empty_like(x)
+        awq_inference_engine.layernorm_forward_cuda(x, self.weight, output, self.eps)
+        return output
+
+
+def precompute_freqs_cis(
+    dim: int, end: int, theta: float = 10000.0, scale: float = 1.0
+):
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)  # type: ignore
+    freqs = torch.outer(t * scale, freqs).float()  # type: ignore
+
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
+    return freqs_cis
+
+
+def precompute_freqs(
+    dim: int, end: int, theta: float = 10000.0, scale: float = 1.0, device=None
+):
+    inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float().to(device) / dim))
+    seq = torch.arange(end, dtype=inv_freq.dtype, device=device)
+    freqs = torch.einsum("i , j -> i j", seq, inv_freq)
+    freqs = freqs.reshape(freqs.shape[0], 1, 1, -1)
+    return torch.cat((freqs, freqs), dim=-1)
+
+
+def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
+    ndim = x.ndim
+    assert 0 <= 1 < ndim
+    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
+    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+    return freqs_cis.view(*shape)
+
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cis: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    # xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    # k_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    xq_ = torch.view_as_complex(
+        xq.float().reshape(*xq.shape[:-1], 2, -1).transpose(-2, -1).contiguous()
+    )
+    xk_ = torch.view_as_complex(
+        xk.float().reshape(*xk.shape[:-1], 2, -1).transpose(-2, -1).contiguous()
+    )
+    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
+    xq_out = torch.view_as_real(xq_ * freqs_cis).transpose(-2, -1).flatten(3)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).transpose(-2, -1).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+class Qwen2MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, hidden_state):
+        return self.down_proj(
+            self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)
+        )
+
+
+# Copied from transformers.models.llama.modeling_llama.repeat_kv
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = x.shape
+    if n_rep == 1:
+        return x
+    x = x[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return x.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+class Qwen2AttentionFused(nn.Module):
+    """
+    Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer
+    and "Generating Long Sequences with Sparse Transformers".
+    """
+
+    def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.args = config
+        self.layer_idx = layer_idx
+        if layer_idx is None:
+            print(
+                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.is_causal = True
+        self.attention_dropout = config.attention_dropout
+        self.rope_scaling = config.rope_scaling
+        if self.rope_scaling is None:
+            self.rope_scaling = 1.0
+        elif isinstance(self.rope_scaling, dict):
+            self.rope_scaling = self.rope_scaling.get("factor", 1.0)
+
+        if (self.head_dim * self.num_heads) != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {self.num_heads})."
+            )
+        self.q_proj = nn.Linear(
+            self.hidden_size, self.num_heads * self.head_dim, bias=True
+        )
+        self.k_proj = nn.Linear(
+            self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, self.hidden_size, bias=False
+        )
+        self.kv_max_seq_len = min(max_seq_len, self.max_position_embeddings)
+        # following fastertransformer definition
+        self.cache_v = (
+            torch.zeros(
+                (
+                    max_batch_size,
+                    self.num_key_value_heads,
+                    # args.max_position_embeddings,
+                    self.kv_max_seq_len,
+                    self.head_dim,
+                )
+            )
+            .cuda()
+            .half()
+        )  # added to half
+        # 8: pack 8 fp16 in FT, if fp32 then use 4
+        self.cache_k = (
+            torch.zeros(
+                (
+                    max_batch_size,
+                    self.num_key_value_heads,
+                    self.head_dim // 8,
+                    # args.max_position_embeddings,
+                    self.kv_max_seq_len,
+                    8,
+                )
+            )
+            .cuda()
+            .half()
+        )  # added to half
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        freqs: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        chunk_prefilling: bool = False,
+    ):
+        bsz, seqlen, _ = x.shape
+
+        query_states = self.q_proj(x)
+        key_states = self.k_proj(x)
+        value_states = self.v_proj(x)
+
+        if seqlen > 1:
+            xq = query_states.view(bsz, seqlen, self.num_heads, self.head_dim)
+            xk = key_states.view(bsz, seqlen, self.num_key_value_heads, self.head_dim)
+            xv = value_states.view(bsz, seqlen, self.num_key_value_heads, self.head_dim)
+
+            xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs)
+
+            self.cache_k = self.cache_k.to(xq)
+            self.cache_v = self.cache_v.to(xq)
+
+            values_store = xv.transpose(2, 1)
+
+            keys_store = (
+                xk.reshape(bsz, seqlen, self.num_key_value_heads, self.head_dim // 8, 8)
+                .permute(0, 2, 3, 1, 4)
+                .contiguous()
+            )
+
+            self.cache_v[:bsz, :, start_pos : start_pos + seqlen, :] = values_store
+            self.cache_k[:bsz, :, :, start_pos : start_pos + seqlen, :] = keys_store
+            if chunk_prefilling:
+                keys = self.cache_k[:, :, :, 0 : start_pos + seqlen, :]
+                keys = (
+                    keys.permute(0, 3, 1, 2, 4)
+                    .reshape(
+                        bsz, start_pos + seqlen, self.num_key_value_heads, self.head_dim
+                    )
+                    .contiguous()
+                )
+                values = self.cache_v[:, :, 0 : start_pos + seqlen, :]
+                values = (
+                    values.transpose(2, 1)
+                    .reshape(
+                        bsz, start_pos + seqlen, self.num_key_value_heads, self.head_dim
+                    )
+                    .contiguous()
+                )
+            else:
+                keys = xk
+                values = xv
+            output = flash_attn_func(
+                q=xq,
+                k=keys,
+                v=values,
+                causal=True,
+            )
+            output = output.contiguous().view(bsz, seqlen, -1)
+        else:
+            xq = query_states.view(bsz, self.num_heads, self.head_dim)
+            xk = key_states.view(bsz, self.num_key_value_heads, self.head_dim)
+            xv = value_states.view(bsz, self.num_key_value_heads, self.head_dim)
+            
+            output = awq_inference_engine.single_query_attention(
+                xq,
+                xk,
+                xv,
+                self.cache_k,
+                self.cache_v,
+                None,
+                # alibi position encodings
+                None,
+                start_pos,
+                self.head_dim,
+                self.rope_theta,
+                self.rope_scaling,
+                True,
+            )
+            output = output.reshape(bsz, 1, -1)
+
+        return self.o_proj(output)
+
+
+class Qwen2DecoderLayer(nn.Module):
+    def __init__(self, config: Qwen2Config, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = Qwen2AttentionFused(config, layer_idx)
+
+        self.mlp = Qwen2MLP(config)
+        self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen2RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        freqs: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        chunk_prefilling: bool = False,
+    ):
+        residual = x
+        x = self.input_layernorm(x)
+
+        # Self Attention
+        x = self.self_attn(
+            x=x,
+            start_pos=start_pos,
+            freqs=freqs,
+            mask=mask,
+            chunk_prefilling=chunk_prefilling,
+        )
+        x = residual + x
+
+        # Fully Connected
+        residual = x
+        x = self.post_attention_layernorm(x)
+        x = self.mlp(x)
+        x = residual + x
+        return x
+
+
+class Qwen2Model(nn.Module):
+    def __init__(self, config: Qwen2Config):
+        super().__init__()
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [
+                Qwen2DecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        # Note (Haotian): rope_theta has to be defined here, otherwise context stage is wrong.
+        rope_scale = config.rope_scaling
+        if rope_scale is None:
+            rope_scale = 1.0
+        else:
+            rope_scale = 1.0 / rope_scale["factor"]
+        self.freqs = precompute_freqs(
+            config.hidden_size // config.num_attention_heads,
+            config.max_position_embeddings * 2,
+            config.rope_theta,
+            rope_scale,
+        )
+        self.freqs_cis = precompute_freqs_cis(
+            config.hidden_size // config.num_attention_heads,
+            config.max_position_embeddings * 2,
+            config.rope_theta,
+            rope_scale,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        start_pos: Optional[int] = 0,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        chunk_prefilling: bool = False,
+    ):
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        seqlen = inputs_embeds.shape[1]
+
+        self.freqs = self.freqs.to(inputs_embeds.device)
+        freqs = self.freqs[start_pos : start_pos + seqlen]
+
+        mask = None
+        if seqlen > 1:
+            mask = torch.full(
+                (1, 1, seqlen, seqlen), float("-inf"), device=inputs_embeds.device
+            )
+            mask = torch.triu(mask, diagonal=1).type_as(inputs_embeds)
+            if chunk_prefilling:
+                mask_history = torch.zeros(
+                    (1, 1, seqlen, start_pos),
+                    dtype=torch.float16,
+                    device=inputs_embeds.device,
+                ).type_as(inputs_embeds)
+                mask = torch.cat((mask_history, mask), dim=-1)
+        x = inputs_embeds
+
+        for decoder_layer in self.layers:
+            x = decoder_layer(x, start_pos, freqs, mask, chunk_prefilling)
+        x = x[:, -1:, :]
+        x = self.norm(x)
+
+        return x
+
+    def forwardfp16(
+        self,
+        input_ids: torch.LongTensor = None,
+        start_pos: Optional[int] = 0,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        chunk_prefilling: bool = False,
+    ):
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        seqlen = inputs_embeds.shape[1]
+
+        self.freqs_cis = self.freqs_cis.to(inputs_embeds.device)
+        freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]
+
+        mask = None
+        if seqlen > 1:
+            mask = torch.full(
+                (1, 1, seqlen, seqlen), float("-inf"), device=inputs_embeds.device
+            )
+            mask = torch.triu(mask, diagonal=1).type_as(inputs_embeds)
+            if chunk_prefilling:
+                mask_history = torch.zeros(
+                    (1, 1, seqlen, start_pos),
+                    dtype=torch.float16,
+                    device=inputs_embeds.device,
+                ).type_as(inputs_embeds)
+                mask = torch.cat((mask_history, mask), dim=-1)
+        x = inputs_embeds
+
+        for decoder_layer in self.layers:
+            x = decoder_layer(x, start_pos, freqs_cis, mask, chunk_prefilling)
+        x = x[:, -1:, :]
+        x = self.norm(x)
+
+        return x
+
+
+class Qwen2ForCausalLM(Qwen2ForCausalLM):
+    def __init__(self, config):
+
+        def skip(*args, **kwargs):
+            pass
+
+        torch.nn.init.kaiming_uniform_ = skip
+        torch.nn.init.kaiming_normal_ = skip
+        torch.nn.init.uniform_ = skip
+        torch.nn.init.normal_ = skip
+        from transformers import modeling_utils
+
+        modeling_utils._init_weights = False
+
+        super().__init__(config)
+        self.model = Qwen2Model(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.config = config
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        start_pos: int = 0,
+        inputs_embeds: torch.Tensor = None,
+        chunk_prefilling: bool = False,
+        quant=True,
+    ):
+        if quant:
+            outputs = self.model(
+                input_ids=input_ids,
+                inputs_embeds=inputs_embeds,
+                start_pos=start_pos,
+                chunk_prefilling=chunk_prefilling,
+            )
+        else:
+            outputs = self.model.forwardfp16(
+                input_ids=input_ids,
+                inputs_embeds=inputs_embeds,
+                start_pos=start_pos,
+                chunk_prefilling=chunk_prefilling,
+            )
+        logits = self.lm_head(outputs)
+        return logits
+
+    def benchmark(self, inputs_embeds, attention_mask, max_output=128, quant_llm=True):
+        output_list = []
+        start_pos = 0
+        for i in range(10):
+            torch.cuda.synchronize()
+            tst = time.time()
+            token = self.forward(None, start_pos, inputs_embeds, quant=quant_llm)
+            torch.cuda.synchronize()
+            ted = time.time()
+        print(
+            "LLM TTFT: {:.6f} s for {} tokens".format(
+                (ted - tst), inputs_embeds.shape[1]
+            )
+        )
+        start_pos = inputs_embeds.shape[1]
+        token = torch.argmax(token, keepdim=True)[0]
+        output_list.append(token)
+
+        torch.cuda.synchronize()
+        tst = time.time()
+        for _ in range(max_output):
+            token = self.forward(token, start_pos)
+            token = torch.argmax(token, keepdim=True)[
+                0
+            ]  # Only fixed-length eager decoding is supported now
+            output_list.append(token)
+            start_pos += 1
+        torch.cuda.synchronize()
+        ted = time.time()
+        print("Decoding througput: {:.6f} tokens/s".format(max_output / (ted - tst)))
+
+        return torch.cat(output_list, dim=1)

llm-awq/tinychat/models/vila_llama.py

@@ -0,0 +1,109 @@
+import os
+import warnings
+import shutil
+import torch
+import torch.nn as nn
+from typing import List, Optional, Tuple, Union
+import time
+
+from transformers import AutoConfig, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from llava.model.utils import get_model_config
+from llava.model.language_model.builder import build_llm_and_tokenizer
+from llava.model.multimodal_encoder.builder import build_vision_tower
+from llava.model.multimodal_projector.builder import build_mm_projector
+from llava.model.llava_arch import LlavaMetaModel, LlavaMetaForCausalLM
+from .llama import LlamaForCausalLM, Transformer
+
+
+class VilaLlamaForCausalLM(LlavaMetaModel, LlavaMetaForCausalLM, PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.init_vlm(config)
+
+    def init_vlm(self, config=None, *args, **kwargs):
+        if (
+            hasattr(self, "llm")
+            or hasattr(self, "vision_tower")
+            or hasattr(self, "mm_projector")
+        ):
+            # already initialized, skipped
+            return
+
+        model_dtype = getattr(config, "model_dtype", "torch.float16")
+        if not hasattr(config, "model_dtype"):
+            warnings.warn(
+                "model_dtype not found in config, defaulting to torch.float16."
+            )
+            config.model_dtype = model_dtype
+
+        # print("init_vlm(): config", config); input("DEBUG init_vlm")
+        cfgs = get_model_config(config)
+        if len(cfgs) == 3:
+            llm_cfg, vision_tower_cfg, mm_projector_cfg = cfgs
+        else:
+            raise ValueError(
+                "`llm_cfg` `mm_projector_cfg` `vision_tower_cfg` not found in the config."
+            )
+        # print("init_vlm():", cfgs); input("DEBUG init_vlm")
+        llm_cfg = AutoConfig.from_pretrained(llm_cfg)
+
+        # self.llm, self.tokenizer = build_llm_and_tokenizer(llm_cfg, config, *args, **kwargs)
+        self.llm = LlamaForCausalLM(llm_cfg)
+        self.vision_tower = build_vision_tower(vision_tower_cfg, config)
+        self.mm_projector = build_mm_projector(mm_projector_cfg, config)
+
+        self.post_config()
+        self.is_loaded = True
+
+        assert (
+            self.llm is not None
+            or self.vision_tower is not None
+            or self.mm_projector is not None
+        ), "At least one of the components must be instantiated."
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        start_pos: int = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        images: Optional[torch.FloatTensor] = None,
+        return_dict: Optional[bool] = None,
+        special_token: bool = False,
+        chunk_prefilling: bool = False,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        self.freezed_module_patch()
+        if inputs_embeds is None:
+            (
+                _,
+                _,
+                _,
+                _,
+                inputs_embeds,
+                _,
+            ) = self.prepare_inputs_labels_for_multimodal(
+                input_ids, position_ids, attention_mask, past_key_values, labels, images
+            )
+        if inputs_embeds is not None:
+            outputs = self.llm.forward(
+                tokens=None,
+                start_pos=start_pos,
+                inputs_embeds=inputs_embeds,
+                chunk_prefilling=chunk_prefilling,
+            )
+        else:  # tokens
+            outputs = self.llm.forward(
+                tokens=input_ids,
+                start_pos=start_pos,
+                inputs_embeds=None,
+                chunk_prefilling=chunk_prefilling,
+            )
+        return outputs

llm-awq/tinychat/modules/fused_attn.py

@@ -0,0 +1,634 @@
+import math
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from transformers.models.llama.modeling_llama import (
+    LlamaAttention,
+    LlamaRotaryEmbedding,
+    apply_rotary_pos_emb,
+)
+from typing import Optional
+from awq.quantize.qmodule import WQLinear
+import awq_inference_engine
+from tinychat.models.llama import apply_rotary_emb
+import gc
+
+import tinychat.utils.constants
+from flash_attn import flash_attn_func
+from tinychat.models.llama import LlamaAttentionFused
+from tinychat.models.qwen2 import Qwen2AttentionFused
+
+max_batch_size = tinychat.utils.constants.max_batch_size
+max_seq_len = tinychat.utils.constants.max_seq_len
+
+
+class QuantLlamaRotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
+        super().__init__()
+
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        inv_freq = 1.0 / (
+            self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)
+        )
+        self.register_buffer("inv_freq", inv_freq)
+        # Build here to make `torch.jit.trace` work.
+        self._set_cos_sin_cache(
+            seq_len=max_position_embeddings,
+            device=self.inv_freq.device,
+            dtype=torch.get_default_dtype(),
+        )
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(
+            self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype
+        )
+
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        # emb = torch.cat((freqs, freqs), dim=-1)
+
+        cos = freqs.cos()
+        sin = freqs.sin()
+        cache = torch.cat((cos, sin), dim=-1)
+
+        # self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
+        # self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)
+        self.register_buffer("cos_sin_cache", cache.half(), persistent=False)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        positions: torch.Tensor,
+    ):
+        # Apply rotary embedding to the query and key before passing them
+        # to the attention op.
+        # print(positions.shape, query.shape, key.shape, self.cos_sin_cache.shape)
+        query = query.contiguous()
+        key = key.contiguous()
+        awq_inference_engine.rotary_embedding_neox(
+            positions,
+            query,
+            key,
+            self.dim,
+            self.cos_sin_cache,
+        )
+        return query, key
+
+
+class QuantLlamaAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, hidden_size, num_heads, qkv_proj, o_proj, dev):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+
+        if (self.head_dim * num_heads) != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {num_heads})."
+            )
+        self.qkv_proj = qkv_proj
+        self.o_proj = o_proj
+        self.rotary_emb = QuantLlamaRotaryEmbedding(
+            self.head_dim, max_position_embeddings=2048, device=dev
+        )
+
+    def forward(
+        self,
+        hidden_states,
+        past_key_value=None,
+        attention_mask=None,
+        position_ids=None,
+        output_attentions=False,
+        use_cache=False,
+    ):
+        """Input shape: Batch x Time x Channel"""
+
+        bsz, q_len, _ = hidden_states.size()
+
+        qkv_states = self.qkv_proj(hidden_states)
+        qkv_states = qkv_states.view(bsz, q_len, 3, self.num_heads, self.head_dim)
+
+        # This updates the query and key states in-place, saving VRAM.
+        query_states, key_states, value_states = torch.split(qkv_states, 1, dim=2)
+        query_states, key_states = self.rotary_emb(
+            query_states, key_states, position_ids
+        )
+
+        del qkv_states
+        query_states = query_states.view(
+            bsz, q_len, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        key_states = key_states.view(
+            bsz, q_len, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        value_states = value_states.view(
+            bsz, q_len, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+
+        is_causal = past_key_value is None
+
+        kv_seq_len = q_len
+        if past_key_value is not None:
+            kv_seq_len += past_key_value[0].shape[-2]
+
+        value_states = value_states.to("cuda:0")
+
+        if past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+        if use_cache:
+            # Since qkv_proj is fused, query_states etc will hold a reference to the original qkv_states tensor
+            # which can cause excessive memory usage by the cache. `contiguous` is a convenient way to workaround this.
+            key_states = key_states.contiguous()
+            value_states = value_states.contiguous()
+            query_states = query_states.contiguous()
+
+        past_key_value = (key_states, value_states) if use_cache else None
+
+        # with torch.backends.cuda.sdp_kernel(enable_math=False):
+        attn_output = F.scaled_dot_product_attention(
+            query_states, key_states, value_states, is_causal=is_causal
+        )
+        del query_states, key_states, value_states
+
+        attn_output = attn_output.transpose(1, 2).reshape(bsz, q_len, self.hidden_size)
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, None, past_key_value
+
+
+class QuantLlamaAttentionFused(nn.Module):
+    def __init__(
+        self, hidden_size, num_heads, kv_max_seq_len, qkv_layer, o_proj, dev, args
+    ):
+        super().__init__()
+
+        self.args = args
+        self.n_local_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.num_heads = args.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+
+        self.num_key_value_heads = args.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = args.max_position_embeddings
+        self.rope_theta = args.rope_theta
+        self.rope_scaling = args.rope_scaling
+        if self.rope_scaling is None:
+            self.rope_scaling = 1.0
+        if isinstance(self.rope_scaling, dict):
+            self.rope_scaling = self.rope_scaling.get("factor", 1.0)
+
+        self.qkv_proj = qkv_layer
+        self.o_proj = o_proj
+
+        self.kv_max_seq_len = kv_max_seq_len
+
+        # following fastertransformer definition
+        self.cache_v = (
+            torch.zeros(
+                (
+                    max_batch_size,
+                    self.num_key_value_heads,
+                    # args.max_position_embeddings,
+                    self.kv_max_seq_len,
+                    self.head_dim,
+                )
+            )
+            .to(dev)
+            .half()
+        )  # added to half
+        # 8: pack 8 fp16 in FT, if fp32 then use 4
+        self.cache_k = (
+            torch.zeros(
+                (
+                    max_batch_size,
+                    self.num_key_value_heads,
+                    self.head_dim // 8,
+                    # args.max_position_embeddings,
+                    self.kv_max_seq_len,
+                    8,
+                )
+            )
+            .to(dev)
+            .half()
+        )  # added to half
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        freqs: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        chunk_prefilling: bool = False,
+    ):
+        bsz, seqlen, _ = x.shape
+        xqkv = self.qkv_proj(x)
+        xqkv = xqkv.view(
+            bsz,
+            seqlen,
+            self.n_local_heads + self.num_key_value_heads * 2,
+            self.head_dim,
+        )
+        xq = xqkv[:, :, 0 : self.n_local_heads]
+        xk = xqkv[
+            :, :, self.n_local_heads : (self.n_local_heads + self.num_key_value_heads)
+        ]
+        xv = xqkv[:, :, -self.num_key_value_heads :]
+
+        if seqlen > 1:
+            xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+            xk = xk.view(bsz, seqlen, self.num_key_value_heads, self.head_dim)
+            xv = xv.view(bsz, seqlen, self.num_key_value_heads, self.head_dim)
+
+            xq = awq_inference_engine.fused_rope_with_pos_forward_func(xq, freqs, True)
+            xk = awq_inference_engine.fused_rope_with_pos_forward_func(xk, freqs, True)
+
+            self.cache_k = self.cache_k.to(xq)
+            self.cache_v = self.cache_v.to(xq)
+
+            values_store = xv.transpose(2, 1)
+            keys_store = (
+                xk.reshape(bsz, seqlen, self.num_key_value_heads, self.head_dim // 8, 8)
+                .permute(0, 2, 3, 1, 4)
+                .contiguous()
+            )
+
+            self.cache_v[:bsz, :, start_pos : start_pos + seqlen, :] = values_store
+            self.cache_k[:bsz, :, :, start_pos : start_pos + seqlen, :] = keys_store
+            if chunk_prefilling:
+                keys = self.cache_k[:, :, :, 0:start_pos, :]
+                keys = (
+                    keys.permute(0, 3, 1, 2, 4)
+                    .reshape(bsz, start_pos, self.num_key_value_heads, self.head_dim)
+                    .contiguous()
+                )
+                keys = torch.cat((keys, xk), dim=1)
+                values = self.cache_v[:, :, 0:start_pos, :]
+                values = (
+                    values.transpose(2, 1)
+                    .reshape(bsz, start_pos, self.num_key_value_heads, self.head_dim)
+                    .contiguous()
+                )
+                values = torch.cat((values, xv), dim=1)
+            else:
+                keys = xk
+                values = xv
+
+            keys = torch.repeat_interleave(
+                keys, dim=2, repeats=self.num_key_value_groups
+            )
+            values = torch.repeat_interleave(
+                values, dim=2, repeats=self.num_key_value_groups
+            )
+
+            xq = xq.transpose(1, 2)
+            keys = keys.transpose(1, 2)
+            values = values.transpose(1, 2)
+            scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
+            if mask is not None:
+                scores = scores + mask  # (bs, n_local_heads, slen, cache_len + slen)
+            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
+            output = torch.matmul(scores, values)  # (bs, n_local_heads, slen, head_dim)
+            output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
+        else:
+            xq = xq.view(bsz, self.n_local_heads, self.head_dim)
+            xk = xk.view(bsz, self.num_key_value_heads, self.head_dim)
+            xv = xv.view(bsz, self.num_key_value_heads, self.head_dim)
+
+            output = awq_inference_engine.single_query_attention(
+                xq,
+                xk,
+                xv,
+                self.cache_k,
+                self.cache_v,
+                None,
+                None,
+                start_pos,
+                self.head_dim,
+                self.rope_theta,
+                self.rope_scaling,
+                True,
+            )
+            output = output.reshape(bsz, 1, -1)
+
+        return self.o_proj(output)
+
+
+class QuantLlamaAttentionFusedFlash(nn.Module):
+    """Flash_attn_func from 'Flash{A}ttention-2: Faster Attention with Better Parallelism and Work Partitioning' paper"""
+
+    """This function is faster than the varlen one but only supports single-batch inference"""
+
+    def __init__(
+        self, hidden_size, num_heads, kv_max_seq_len, qkv_layer, o_proj, dev, args
+    ):
+        super().__init__()
+
+        self.args = args
+        self.n_local_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.num_heads = args.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+
+        self.num_key_value_heads = args.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = args.max_position_embeddings
+        self.rope_theta = args.rope_theta
+        self.rope_scaling = args.rope_scaling
+        if self.rope_scaling is None:
+            self.rope_scaling = 1.0
+        elif isinstance(self.rope_scaling, dict):
+            self.rope_scaling = self.rope_scaling.get("factor", 1.0)
+
+        self.qkv_proj = qkv_layer
+        self.o_proj = o_proj
+
+        self.kv_max_seq_len = kv_max_seq_len
+        # following fastertransformer definition
+        # For short seqlence, we use fused kernel to accelerate decoding.
+        if self.kv_max_seq_len <= 8192:
+            self.cache_v = (
+                torch.zeros(
+                    (
+                        max_batch_size,
+                        self.num_key_value_heads,
+                        # args.max_position_embeddings,
+                        self.kv_max_seq_len,
+                        self.head_dim,
+                    )
+                )
+                .to(dev)
+                .half()
+            )  # added to half
+            # 8: pack 8 fp16 in FT, if fp32 then use 4
+            self.cache_k = (
+                torch.zeros(
+                    (
+                        max_batch_size,
+                        self.num_key_value_heads,
+                        self.head_dim // 8,
+                        # args.max_position_embeddings,
+                        kv_max_seq_len,
+                        8,
+                    )
+                )
+                .to(dev)
+                .half()
+            )  # added to half
+            self.forward = self.short_forward
+        # For long sequence, we use flash attantion for both prefilling and decoding to avoid OOM.
+        else:
+            self.cache_v = (
+                torch.zeros(
+                    (
+                        max_batch_size,
+                        self.kv_max_seq_len,
+                        self.num_key_value_heads,
+                        self.head_dim,
+                    )
+                )
+                .to(dev)
+                .half()
+            )  # added to half
+            self.cache_k = (
+                torch.zeros(
+                    (
+                        max_batch_size,
+                        self.kv_max_seq_len,
+                        self.num_key_value_heads,
+                        self.head_dim,
+                    )
+                )
+                .to(dev)
+                .half()
+            )  # added to half
+            self.forward = self.long_forward
+
+    def short_forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        freqs: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        chunk_prefilling: bool = False,
+    ):
+        bsz, seqlen, _ = x.shape
+        xqkv = self.qkv_proj(x)
+        xqkv = xqkv.view(
+            bsz,
+            seqlen,
+            self.n_local_heads + self.num_key_value_heads * 2,
+            self.head_dim,
+        )
+        xq = xqkv[:, :, 0 : self.n_local_heads]
+        xk = xqkv[
+            :, :, self.n_local_heads : (self.n_local_heads + self.num_key_value_heads)
+        ]
+        xv = xqkv[:, :, -self.num_key_value_heads :]
+
+        if seqlen > 1:
+            xq = awq_inference_engine.fused_rope_with_pos_forward_func(xq, freqs, True)
+            xk = awq_inference_engine.fused_rope_with_pos_forward_func(xk, freqs, True)
+
+            self.cache_k = self.cache_k.to(xq)
+            self.cache_v = self.cache_v.to(xq)
+
+            values_store = xv.transpose(2, 1)
+            keys_store = (
+                xk.reshape(bsz, seqlen, self.num_key_value_heads, self.head_dim // 8, 8)
+                .permute(0, 2, 3, 1, 4)
+                .contiguous()
+            )
+
+            self.cache_v[:bsz, :, start_pos : start_pos + seqlen, :] = values_store
+            self.cache_k[:bsz, :, :, start_pos : start_pos + seqlen, :] = keys_store
+
+            if chunk_prefilling:
+                keys = self.cache_k[:, :, :, 0 : start_pos + seqlen, :]
+                keys = (
+                    keys.permute(0, 3, 1, 2, 4)
+                    .reshape(
+                        bsz, start_pos + seqlen, self.num_key_value_heads, self.head_dim
+                    )
+                    .contiguous()
+                )
+                values = self.cache_v[:, :, 0 : start_pos + seqlen, :]
+                values = (
+                    values.transpose(2, 1)
+                    .reshape(
+                        bsz, start_pos + seqlen, self.num_key_value_heads, self.head_dim
+                    )
+                    .contiguous()
+                )
+            else:
+                keys = xk
+                values = xv
+
+            output = flash_attn_func(
+                q=xq,
+                k=keys,
+                v=values,
+                causal=True,
+            )
+            output = output.contiguous().view(bsz, seqlen, -1)
+        else:
+            xq = xq.view(bsz, self.n_local_heads, self.head_dim)
+            xk = xk.view(bsz, self.num_key_value_heads, self.head_dim)
+            xv = xv.view(bsz, self.num_key_value_heads, self.head_dim)
+            output = awq_inference_engine.single_query_attention(
+                xq,
+                xk,
+                xv,
+                self.cache_k,
+                self.cache_v,
+                None,
+                None,
+                start_pos,
+                self.head_dim,
+                self.rope_theta,
+                self.rope_scaling,
+                True,
+            )
+            output = output.reshape(bsz, 1, -1)
+        return self.o_proj(output)
+
+    def long_forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int,
+        freqs: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        chunk_prefilling: bool = False,
+    ):
+        bsz, seqlen, _ = x.shape
+        xqkv = self.qkv_proj(x)
+        xqkv = xqkv.view(
+            bsz,
+            seqlen,
+            self.n_local_heads + self.num_key_value_heads * 2,
+            self.head_dim,
+        )
+        xq = xqkv[:, :, 0 : self.n_local_heads]
+        xk = xqkv[
+            :, :, self.n_local_heads : (self.n_local_heads + self.num_key_value_heads)
+        ]
+        xv = xqkv[:, :, -self.num_key_value_heads :]
+
+        xq = awq_inference_engine.fused_rope_with_pos_forward_func(xq, freqs, True)
+        xk = awq_inference_engine.fused_rope_with_pos_forward_func(xk, freqs, True)
+
+        self.cache_k = self.cache_k.to(xq)
+        self.cache_v = self.cache_v.to(xq)
+
+        self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv
+        self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
+
+        keys = self.cache_k[:, 0 : start_pos + seqlen]
+        values = self.cache_v[:, 0 : start_pos + seqlen]
+
+        output = flash_attn_func(
+            q=xq,
+            k=keys,
+            v=values,
+            causal=True,
+        )
+        output = output.view(bsz, seqlen, -1)
+        return self.o_proj(output)
+
+
+def make_quant_attn(model, dev, flash_attn=True):
+    """
+    Replace all LlamaAttention modules with QuantLlamaAttention modules, fusing the q, k, v projections.
+    """
+    model = model.cpu()
+    for name, m in model.named_modules():
+        if not m.__class__.__name__ in [
+            "LlamaAttention",
+            "LlamaAttentionFused",
+            "Qwen2AttentionFused",
+        ]:
+            continue
+
+        q_proj = m.q_proj
+        k_proj = m.k_proj
+        v_proj = m.v_proj
+
+        qweights = torch.cat([q_proj.qweight, k_proj.qweight, v_proj.qweight], dim=0)
+        scaled_zeros = torch.cat(
+            [q_proj.scaled_zeros, k_proj.scaled_zeros, v_proj.scaled_zeros], dim=1
+        ).contiguous()
+        scales = torch.cat(
+            [q_proj.scales, k_proj.scales, v_proj.scales], dim=1
+        ).contiguous()
+        # g_idx = torch.cat([q_proj.g_idx, k_proj.g_idx, v_proj.g_idx], dim=0)
+        g_idx = None
+        bias = (
+            torch.cat([q_proj.bias, k_proj.bias, v_proj.bias], dim=0)
+            if q_proj.bias is not None
+            else None
+        )
+
+        qkv_layer = WQLinear(
+            q_proj.w_bit,
+            q_proj.group_size,
+            q_proj.in_features,
+            q_proj.out_features + k_proj.out_features + v_proj.out_features,
+            q_proj.bias is not None,
+            q_proj.qweight.device,
+        )
+        qkv_layer.qweight = qweights
+        qkv_layer.scaled_zeros = scaled_zeros
+        qkv_layer.scales = scales
+
+        qkv_layer.bias = bias
+        qkv_layer.split_k_iters = q_proj.split_k_iters
+        # We're dropping the rotary embedding layer m.rotary_emb here. We don't need it in the triton branch.
+        if isinstance(m, LlamaAttention):
+            attn = QuantLlamaAttention(
+                m.hidden_size, m.num_heads, qkv_layer, m.o_proj, dev
+            )
+        else:
+            if flash_attn:
+                attn = QuantLlamaAttentionFusedFlash(
+                    m.args.hidden_size,
+                    m.args.num_attention_heads,
+                    m.kv_max_seq_len,
+                    qkv_layer,
+                    m.o_proj,
+                    dev,
+                    m.args,
+                )
+            else:
+                attn = QuantLlamaAttentionFused(
+                    m.args.hidden_size,
+                    m.args.num_attention_heads,
+                    m.kv_max_seq_len,
+                    qkv_layer,
+                    m.o_proj,
+                    dev,
+                    m.args,
+                )
+        if "." in name:
+            parent_name = name.rsplit(".", 1)[0]
+            child_name = name[len(parent_name) + 1 :]
+            parent = model.get_submodule(parent_name)
+        else:
+            parent_name = ""
+            parent = model
+            child_name = name
+
+        # print(f"Replacing {name} with quant_attn; parent: {parent_name}, child's name: {child_name}")
+        setattr(parent, child_name, attn)
+        gc.collect()
+        torch.cuda.empty_cache()
+    model = model.to(dev)

llm-awq/tinychat/modules/fused_internencoder.py

@@ -0,0 +1,237 @@
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from einops import rearrange
+from timm.layers import DropPath
+from torch import nn
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (BaseModelOutput,
+                                           BaseModelOutputWithPooling)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+from awq.quantize import W8A8OF16LinearDynamicInputScale
+import awq_inference_engine
+
+from tinychat.models.internvl.internvit import (FlashAttention,
+                                                InternRMSNorm,
+                                                InternVisionEmbeddings,
+                                                InternAttention,
+                                                InternMLP,
+                                                InternVisionEncoderLayer,
+                                                InternVisionEncoder)
+from tinychat.models.internvl.configuration_internvl import InternVisionConfig
+
+try:
+    from flash_attn.bert_padding import pad_input, unpad_input
+    from flash_attn.flash_attn_interface import \
+        flash_attn_varlen_qkvpacked_func
+    has_flash_attn = True
+except:
+    print('FlashAttention2 is not installed.')
+    has_flash_attn = False
+
+logger = logging.get_logger(__name__)
+
+
+class QuantInternVisionEncoder(nn.Module):
+    """
+    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+    [`InternEncoderLayer`].
+
+    Args:
+        config (`InternConfig`):
+            The corresponding vision configuration for the `InternEncoder`.
+    """
+
+    def __init__(self, module: InternVisionEncoder, bsz=64, seqlen=1024):
+        super().__init__()
+        self.config = module.config
+        # stochastic depth decay rule
+        self.layers = nn.ModuleList([QuantInternVisionEncoderLayer(layer, self.config) for layer in module.layers])
+        self.gradient_checkpointing = True
+        self.bsz = bsz
+        self.seqlen = seqlen
+
+    def forward(
+        self,
+        inputs_embeds,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        r"""
+        Args:
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+                Embedded representation of the inputs. Should be float, not int tokens.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
+                for more detail.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        """
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        encoder_states = () if output_hidden_states else None
+        hidden_states = inputs_embeds
+
+        for idx, encoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_states,)
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    encoder_layer,
+                    hidden_states)
+            else:
+                layer_outputs = encoder_layer(
+                    hidden_states,
+                )
+            hidden_states = layer_outputs
+
+        if output_hidden_states:
+            encoder_states = encoder_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, encoder_states] if v is not None)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states, hidden_states=encoder_states
+        )
+
+class QuantInternRMSNorm(nn.Module):
+    def __init__(self, module: nn.Module, use_per_token_quant=True):
+        super().__init__()
+        self.weight = nn.Parameter(module.weight.data, requires_grad=False)
+        self.bias = nn.Parameter(module.bias.data, requires_grad=False)
+        self.variance_epsilon = module.eps
+        self.use_per_token_quant = use_per_token_quant
+
+    def forward(self, hidden_states):
+        bsz, seqlen, hidden_size = hidden_states.shape
+        output = torch.empty((bsz * seqlen), hidden_size, device=hidden_states.device, dtype=torch.int8)
+        scale = torch.empty((bsz * seqlen), device=hidden_states.device, dtype=hidden_states.dtype)
+        awq_inference_engine.rms_norm_general(
+            output,
+            hidden_states,
+            self.weight,
+            self.bias,
+            scale,
+            self.variance_epsilon,
+            self.use_per_token_quant,
+        )
+        return output, scale
+
+class QuantInternAttention(nn.Module):
+    def __init__(self, module: InternAttention, config: InternVisionConfig, init_only=False):
+        super().__init__()
+        self.config = config
+        self.embed_dim = module.embed_dim
+        self.num_heads = module.num_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        self.scale = module.scale
+        self.use_flash_attn = config.use_flash_attn
+
+        self.qkv = W8A8OF16LinearDynamicInputScale.from_linear(module.qkv, init_only=init_only)
+        self.proj = W8A8OF16LinearDynamicInputScale.from_linear(module.proj, init_only=init_only)
+
+        self.qk_normalization = module.qk_normalization
+        if self.qk_normalization:
+            self.q_norm = QuantInternRMSNorm(module.q_norm)
+            self.k_norm = QuantInternRMSNorm(module.k_norm)
+
+        if self.use_flash_attn:
+            from tinychat.models.internvl.internvit import FlashAttention
+            self.inner_attn = FlashAttention(attention_dropout=config.attention_dropout)
+
+    def forward(self, hidden_states: torch.Tensor, scale_in: torch.Tensor):
+        bsz, seqlen, hidden_size = hidden_states.shape
+
+        qkv_out = torch.empty(bsz * seqlen, 3 * hidden_size, dtype=torch.float16, device=hidden_states.device)
+        self.qkv(hidden_states.reshape(-1, hidden_size), scale_in, qkv_out)
+
+        qkv = rearrange(qkv_out.view(bsz, seqlen, -1), 'b s (three h d) -> b s three h d', three=3, h=self.num_heads)
+
+        if self.qk_normalization:
+            q, k, v = qkv.unbind(2)
+            q, _ = self.q_norm(q.flatten(-2, -1)); q = q.view_as(q)
+            k, _ = self.k_norm(k.flatten(-2, -1)); k = k.view_as(k)
+            qkv = torch.stack([q, k, v], dim=2)
+
+        attn_out, _ = self.inner_attn(qkv, need_weights=False, causal=False)
+        attn_out = rearrange(attn_out, 'b s h d -> (b s) (h d)')
+
+        quant_out = torch.empty_like(attn_out, dtype=torch.int8)
+        scale_proj_in = torch.empty(bsz * seqlen, device=hidden_states.device, dtype=torch.float16)
+        awq_inference_engine.invoke_quant(quant_out, attn_out, scale_proj_in)
+
+        proj_out = torch.empty_like(attn_out)
+        self.proj(quant_out, scale_proj_in, proj_out)
+
+        return proj_out
+
+class QuantInternMLP(nn.Module):
+    def __init__(self, module: InternMLP, config: InternVisionConfig):
+        super().__init__()
+        self.config = config
+        self.act = module.act
+        self.fc1 = W8A8OF16LinearDynamicInputScale.from_linear(module.fc1)
+        self.fc2 = W8A8OF16LinearDynamicInputScale.from_linear(module.fc2)
+
+    def forward(self, hidden_states: torch.Tensor, scale_in: torch.Tensor):
+        bsz, seqlen, hidden_size = hidden_states.shape
+        device = hidden_states.device
+        
+        fc1_out = torch.empty((bsz * seqlen), self.config.intermediate_size, dtype=torch.float16, device=device)
+        self.fc1(hidden_states.reshape(-1, hidden_size), scale_in, fc1_out)
+        
+        tmp = torch.empty(
+            ((bsz * seqlen) * self.config.intermediate_size),
+            device=device,
+            dtype=torch.float16,
+        )
+        act_out = torch.empty_like(fc1_out, dtype=torch.int8)
+        scale_act = torch.empty(bsz * seqlen, device=device, dtype=torch.float16)
+        awq_inference_engine.gelu_and_quant(act_out, fc1_out, scale_act, tmp)
+
+        fc2_out = torch.empty((bsz * seqlen), hidden_size, dtype=torch.float16, device=device)
+        self.fc2(act_out, scale_act, fc2_out)
+
+        return fc2_out
+
+class QuantInternVisionEncoderLayer(nn.Module):
+    def __init__(self, module: InternVisionEncoderLayer, config: InternVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+
+        self.attn = QuantInternAttention(module.attn, config)
+        self.mlp = QuantInternMLP(module.mlp, config)
+
+        self.norm1 = QuantInternRMSNorm(module.norm1)
+        self.norm2 = QuantInternRMSNorm(module.norm2)
+
+        self.ls1 = module.ls1
+        self.ls2 = module.ls2
+
+    def forward(self, hidden_states: torch.Tensor):
+        bsz, seqlen, hidden_size = hidden_states.shape
+
+        residual = hidden_states
+        norm1_out, scale1 = self.norm1(hidden_states)
+        attn_out = self.attn(norm1_out.view(bsz, seqlen, hidden_size), scale1)
+        hidden_states = residual + attn_out.view(bsz, seqlen, hidden_size) * self.ls1
+
+        residual = hidden_states
+        norm2_out, scale2 = self.norm2(hidden_states)
+        mlp_out = self.mlp(norm2_out.view(bsz, seqlen, hidden_size), scale2)
+        hidden_states = residual + mlp_out.view(bsz, seqlen, hidden_size) * self.ls2
+
+        return hidden_states
+
+

llm-awq/tinychat/modules/fused_norm.py

@@ -0,0 +1,46 @@
+import torch
+from torch import nn
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
+import awq_inference_engine
+
+
+class FTLlamaRMSNorm(nn.Module):
+    def __init__(self, weight, eps=1e-6):
+        """
+        LlamaRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = weight
+        self.variance_epsilon = eps
+
+    def forward(self, x):
+        output = torch.empty_like(x)
+        awq_inference_engine.layernorm_forward_cuda(
+            x, self.weight, output, self.variance_epsilon
+        )
+        return output
+
+
+def make_quant_norm(model):
+    """
+    Replace all LlamaRMSNorm modules with FTLlamaRMSNorm modules
+    """
+
+    for name, m in model.named_modules():
+        if not isinstance(m, LlamaRMSNorm):
+            continue
+
+        norm = FTLlamaRMSNorm(m.weight, m.variance_epsilon)
+
+        if "." in name:
+            parent_name = name.rsplit(".", 1)[0]
+            child_name = name[len(parent_name) + 1 :]
+            parent = model.get_submodule(parent_name)
+        else:
+            parent_name = ""
+            parent = model
+            child_name = name
+
+        # print(f"Replacing {name} with quant_attn; parent: {parent_name}, child's name: {child_name}")
+
+        setattr(parent, child_name, norm)

llm-awq/tinychat/modules/fused_siglipdecoder.py

@@ -0,0 +1,282 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from awq.quantize import W8A8OF16LinearDynamicInputScale
+from llava.model.multimodal_encoder.siglip.modeling_siglip import (
+    SiglipMLP,
+    SiglipEncoder,
+    SiglipAttention,
+    SiglipEncoderLayer,
+)
+from tinychat.utils.input_metadata import ActivationBuffer
+from transformers.modeling_outputs import BaseModelOutput
+from typing import Optional, Tuple, Union
+from flash_attn import flash_attn_func
+import time
+
+CLIP_RANGE = 5
+
+
+import awq_inference_engine
+
+
+class QuantSiglipEncoder(nn.Module):
+    def __init__(self, module: SiglipEncoder, bsz=64, seqlen=1024):
+        super().__init__()
+        self.config = module.config
+        self.layers = [QuantSiglipEncoderLayer(layer) for layer in module.layers]
+        self.buffer = ActivationBuffer(module)
+        self.bsz = bsz
+        self.seqlen = seqlen
+        self.buffer.allocate_activation_buffer(self.bsz * self.seqlen)
+
+    # Ignore copy
+    def forward(
+        self,
+        inputs_embeds,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,  # dummy
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        # TODO Find why this code is necessary
+        # torch.sum(inputs_embeds!=inputs_embeds)
+        bsz, seqlen, _ = inputs_embeds.shape
+        if self.bsz != bsz or self.seqlen != seqlen:
+            self.buffer.allocate_activation_buffer(bsz * seqlen)
+            self.bsz = bsz
+            self.seqlen = seqlen
+
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        encoder_states = () if output_hidden_states else None
+
+        hidden_states = inputs_embeds
+        for i, encoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                encoder_states = encoder_states + (
+                    hidden_states.reshape(bsz, seqlen, -1),
+                )
+            hidden_states = encoder_layer(
+                hidden_states, self.buffer, attention_mask, bsz, seqlen
+            )
+
+        if output_hidden_states:
+            encoder_states = encoder_states + (hidden_states.reshape(bsz, seqlen, -1),)
+        if not return_dict:
+            return tuple(v for v in [hidden_states, encoder_states] if v is not None)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states.reshape(bsz, seqlen, -1),
+            hidden_states=encoder_states,
+            attentions=None,
+        )
+
+
+class QuantSiglipMLP(nn.Module):
+    def __init__(self, siglipmlp, init_only=False):
+        super().__init__()
+        self.config = siglipmlp.config
+        self.activation_fn = siglipmlp.activation_fn
+        self.fc1 = W8A8OF16LinearDynamicInputScale.from_linear(
+            siglipmlp.fc1, init_only=init_only, fc1=False
+        )
+        self.fc2 = W8A8OF16LinearDynamicInputScale.from_linear(
+            siglipmlp.fc2, init_only=init_only
+        )
+        self.invoke_quant = self.invoke_quant_mlp
+
+    def invoke_quant_mlp(self, buffer, actfn_output):
+        awq_inference_engine.invoke_quant(
+            buffer.quantized_mlp_act_buffer,
+            actfn_output,
+            buffer.quantized_scale_buffer,
+        )
+
+    def forward(self, buffer: ActivationBuffer) -> torch.Tensor:
+        # INT8 in, FP16 out
+        self.fc1(
+            buffer.quantized_hidden_states_buffer,
+            buffer.quantized_scale_buffer,
+            buffer.fc1_buffer,
+        )
+        # Act & quantization
+        awq_inference_engine.gelu_and_quant(
+            buffer.quantized_mlp_act_buffer,
+            buffer.fc1_buffer,
+            buffer.quantized_scale_buffer,
+            buffer.tmp,
+        )
+        # INT8 in, FP16 out
+        self.fc2(
+            buffer.quantized_mlp_act_buffer,
+            buffer.quantized_scale_buffer,
+            buffer.in_out_fc2_act_buffer,
+        )
+
+
+class QuantSiglipFlashAttention2(nn.Module):
+    def __init__(
+        self,
+        module: SiglipAttention,
+        init_only=False,
+    ):
+        super().__init__()
+        self.config = module.config
+        self.embed_dim = module.embed_dim
+        self.num_heads = module.num_heads
+        self.head_dim = self.embed_dim // self.num_heads
+
+        self.qkv_proj = W8A8OF16LinearDynamicInputScale.from_qkv(
+            module.q_proj, module.k_proj, module.v_proj, init_only=init_only
+        )
+        self.out_proj = W8A8OF16LinearDynamicInputScale.from_linear(
+            module.out_proj, init_only=init_only
+        )
+        self.invoke_quant = self.invoke_quant_wo
+
+    def invoke_quant_wo(self, buffer, attn_output):
+        awq_inference_engine.invoke_quant(
+            buffer.quantized_hidden_states_buffer,
+            attn_output,
+            buffer.quantized_scale_buffer,
+        )
+
+    # Adapted from transformers.models.llama.modeling_llama.LlamaFlashAttention2.forward
+    def forward(
+        self, buffer: ActivationBuffer, bsz=64, seqlen=1024
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        # qkv
+        self.qkv_proj(
+            buffer.quantized_hidden_states_buffer,
+            buffer.quantized_scale_buffer,
+            buffer.qkv_proj_act_buffer,
+        )
+        q, k, v = buffer.qkv_proj_act_buffer.split(
+            [self.embed_dim, self.embed_dim, self.embed_dim], dim=-1
+        )
+        q = q.reshape(bsz, seqlen, self.num_heads, self.head_dim)
+        k = k.reshape(bsz, seqlen, self.num_heads, self.head_dim)
+        v = v.reshape(bsz, seqlen, self.num_heads, self.head_dim)
+        attn_output = flash_attn_func(q, k, v, softmax_scale=None, causal=False)
+        attn_output = attn_output.reshape(bsz * seqlen, -1)
+        # FP16 -> int8
+        self.invoke_quant(buffer, attn_output)
+        # INT8 in, FP16 out
+        self.out_proj(
+            buffer.quantized_hidden_states_buffer,
+            buffer.quantized_scale_buffer,
+            buffer.in_out_fc2_act_buffer,
+        )
+
+
+class QuantSiglipEncoderLayer(nn.Module):
+    def __init__(self, module: SiglipEncoderLayer):
+        super().__init__()
+        self.embed_dim = module.embed_dim
+        self.self_attn = QuantSiglipFlashAttention2(module.self_attn)
+        self.layer_norm1 = RMSNormGeneral(
+            module.layer_norm1.weight.data,
+            module.layer_norm1.bias.data,
+            module.layer_norm1.eps,
+            True,
+        ).cuda()
+        self.mlp = QuantSiglipMLP(module.mlp)
+        self.layer_norm2 = RMSNormGeneral(
+            module.layer_norm2.weight.data,
+            module.layer_norm2.bias.data,
+            module.layer_norm2.eps,
+            True,
+        ).cuda()
+        self.quant = self.invoke_quant_norm
+
+    def invoke_quant_norm(self, buffer, normfn_output):
+        awq_inference_engine.invoke_quant(
+            buffer.quantized_hidden_states_buffer,
+            normfn_output,
+            buffer.quantized_scale_buffer,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        buffer: ActivationBuffer,
+        attention_mask,
+        bsz,
+        seqlen,
+    ) -> Tuple[torch.FloatTensor]:
+        # Attention block
+        # FP16 in int8 out, layernorm & quantization
+        residual = hidden_states
+        self.layer_norm1(
+            hidden_states.reshape(-1, self.embed_dim),
+            buffer.quantized_hidden_states_buffer,
+            buffer.quantized_scale_buffer,
+        )
+
+        # INT8 -> FP16
+        self.self_attn(buffer, bsz, seqlen)
+        hidden_states = (
+            residual.reshape(-1, self.embed_dim) + buffer.in_out_fc2_act_buffer
+        )
+        # Fully Connected
+        residual = hidden_states
+        # FP16 in int8 out, layernorm & quantization
+        self.layer_norm2(
+            hidden_states.reshape(-1, self.embed_dim),
+            buffer.quantized_hidden_states_buffer,
+            buffer.quantized_scale_buffer,
+        )
+
+        # INT8 -> FP16
+        self.mlp(buffer)
+        hidden_states = (
+            residual.reshape(-1, self.embed_dim) + buffer.in_out_fc2_act_buffer
+        )
+        return hidden_states
+
+
+class RMSNormGeneral(nn.Module):
+    """Root mean square normalization (w/ per-token or per-tensor quant).
+
+    Computes x -> w * x / sqrt(E[x^2] + eps) where w is the learned weight.
+    Refer to https://arxiv.org/abs/1910.07467
+    """
+
+    def __init__(
+        self,
+        weight: torch.tensor,
+        bias: torch.tensor,
+        eps: float = 1e-6,
+        use_per_token_quant: bool = True,
+    ) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(weight, requires_grad=False)
+        self.bias = nn.Parameter(bias, requires_grad=False)
+        self.variance_epsilon = eps
+        self.use_per_token_quant = use_per_token_quant
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        quantized_hidden_states_buffer: torch.Tensor,
+        quantized_scale_buffer: torch.Tensor,
+        quantized_sum_buffer: torch.Tensor = None,
+    ) -> torch.Tensor:
+        # quantized_sum_buffer is not used, only to keep the consistency of the interface
+        awq_inference_engine.rms_norm_general(
+            quantized_hidden_states_buffer,
+            x,
+            self.weight.data,
+            self.bias.data,
+            quantized_scale_buffer,
+            self.variance_epsilon,
+            self.use_per_token_quant,
+        )

llm-awq/tinychat/modules/fused_vision_attn.py

@@ -0,0 +1,272 @@
+import math
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from typing import Optional, Tuple
+
+# from awq.quantize.qmodule import WQLinear
+# import awq_inference_engine
+# from tinychat.models.llama import apply_rotary_emb
+import gc
+
+import tinychat.utils.constants
+
+max_batch_size = tinychat.utils.constants.max_batch_size
+max_seq_len = tinychat.utils.constants.max_seq_len
+
+from transformers.activations import ACT2FN
+from transformers.models.clip.configuration_clip import (
+    CLIPConfig,
+    CLIPTextConfig,
+    CLIPVisionConfig,
+)
+from transformers.models.clip.modeling_clip import CLIPAttention
+
+
+class CLIPAttentionFused(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(
+        self, hidden_size, num_heads, qkv_proj, out_proj, dev, attention_dropout=0.0
+    ):
+        super().__init__()
+        self.embed_dim = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.scale = self.head_dim**-0.5
+        self.dropout = attention_dropout
+
+        if (self.head_dim * num_heads) != self.embed_dim:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {num_heads})."
+            )
+        self.qkv_proj = qkv_proj
+        self.out_proj = out_proj
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return (
+            tensor.view(bsz, seq_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+            .contiguous()
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        causal_attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        """Input shape: Batch x Time x Channel"""
+
+        bsz, tgt_len, embed_dim = hidden_states.size()
+
+        qkv_states = self.qkv_proj(hidden_states)
+        qkv_states = qkv_states.view(bsz, tgt_len, 3, self.num_heads, self.head_dim)
+
+        # This updates the query and key states in-place, saving VRAM.
+        query_states, key_states, value_states = torch.split(qkv_states, 1, dim=2)
+        del qkv_states
+
+        proj_shape = (bsz * self.num_heads, -1, self.head_dim)
+
+        query_states = (
+            query_states.view(bsz, tgt_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+            .view(*proj_shape)
+            * self.scale
+        )
+        key_states = (
+            key_states.view(bsz, tgt_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+            .view(*proj_shape)
+        )
+        value_states = (
+            value_states.view(bsz, tgt_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+            .view(*proj_shape)
+        )
+
+        src_len = key_states.size(1)
+        attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
+
+        if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        # apply the causal_attention_mask first
+        if causal_attention_mask is not None:
+            if causal_attention_mask.size() != (bsz, 1, tgt_len, src_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is"
+                    f" {causal_attention_mask.size()}"
+                )
+            attn_weights = (
+                attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+                + causal_attention_mask
+            )
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, tgt_len, src_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = (
+                attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+                + attention_mask
+            )
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+
+        if output_attentions:
+            # this operation is a bit akward, but it's required to
+            # make sure that attn_weights keeps its gradient.
+            # In order to do so, attn_weights have to reshaped
+            # twice and have to be reused in the following
+            attn_weights_reshaped = attn_weights.view(
+                bsz, self.num_heads, tgt_len, src_len
+            )
+            attn_weights = attn_weights_reshaped.view(
+                bsz * self.num_heads, tgt_len, src_len
+            )
+        else:
+            attn_weights_reshaped = None
+
+        attn_probs = nn.functional.dropout(
+            attn_weights, p=self.dropout, training=self.training
+        )
+
+        attn_output = torch.bmm(attn_probs, value_states)
+
+        if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
+        attn_output = attn_output.transpose(1, 2)
+        attn_output = attn_output.reshape(bsz, tgt_len, embed_dim)
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, attn_weights_reshaped
+
+
+class CLIPMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.activation_fn = ACT2FN[config.hidden_act]
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+
+class CLIPEncoderLayer(nn.Module):
+    def __init__(self, config: CLIPConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.self_attn = CLIPAttention(config)
+        self.layer_norm1 = nn.LayerNorm(self.embed_dim)
+        self.mlp = CLIPMLP(config)
+        self.layer_norm2 = nn.LayerNorm(self.embed_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        causal_attention_mask: torch.Tensor,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.FloatTensor]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`): attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+                `(config.encoder_attention_heads,)`.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+        """
+        residual = hidden_states
+
+        hidden_states = self.layer_norm1(hidden_states)
+        hidden_states, attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            causal_attention_mask=causal_attention_mask,
+            output_attentions=output_attentions,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+
+def make_fused_vision_attn(model, dev):
+    """
+    Replace all LlamaAttention modules with QuantLlamaAttention modules, fusing the q, k, v projections.
+    """
+    model = model.cpu()
+    for name, m in model.named_modules():
+        if not m.__class__.__name__ in ["CLIPAttention", "CLIPAttentionFused"]:
+            continue
+
+        q_proj = m.q_proj
+        k_proj = m.k_proj
+        v_proj = m.v_proj
+
+        weights = torch.cat([q_proj.weight, k_proj.weight, v_proj.weight], dim=0)
+        bias = (
+            torch.cat([q_proj.bias, k_proj.bias, v_proj.bias], dim=0)
+            if q_proj.bias is not None
+            else None
+        )
+
+        qkv_layer = nn.Linear(
+            q_proj.in_features,
+            q_proj.out_features + k_proj.out_features + v_proj.out_features,
+            q_proj.bias is not None,
+            q_proj.weight.device,
+        )
+        qkv_layer.weight.data = weights
+
+        qkv_layer.bias.data = bias
+        if isinstance(m, CLIPAttention):
+            attn = CLIPAttentionFused(
+                m.embed_dim, m.num_heads, qkv_layer, m.out_proj, dev
+            )
+        if "." in name:
+            parent_name = name.rsplit(".", 1)[0]
+            child_name = name[len(parent_name) + 1 :]
+            parent = model.get_submodule(parent_name)
+        else:
+            parent_name = ""
+            parent = model
+            child_name = name
+
+        # print(f"Replacing {name} with quant_attn; parent: {parent_name}, child's name: {child_name}")
+        setattr(parent, child_name, attn)
+        gc.collect()
+        torch.cuda.empty_cache()
+    model = model.to(dev)