diff --git a/.cache/torch/comm_lib_trace_rank_0 b/.cache/torch/comm_lib_trace_rank_0
new file mode 100644
index 0000000000000000000000000000000000000000..0d806759497252c1149322598775e6325d84e359
Binary files /dev/null and b/.cache/torch/comm_lib_trace_rank_0 differ
diff --git a/.cache/torch/comm_lib_trace_rank_1 b/.cache/torch/comm_lib_trace_rank_1
new file mode 100644
index 0000000000000000000000000000000000000000..0d806759497252c1149322598775e6325d84e359
Binary files /dev/null and b/.cache/torch/comm_lib_trace_rank_1 differ
diff --git a/.conda/aau_token b/.conda/aau_token
new file mode 100644
index 0000000000000000000000000000000000000000..3efc135b3ac2062c15b92a99ecabdd7dc545313a
--- /dev/null
+++ b/.conda/aau_token
@@ -0,0 +1 @@
+YN_RyWTyaweE0R_BuNYxb-
\ No newline at end of file
diff --git a/.conda/aau_token_host b/.conda/aau_token_host
new file mode 100644
index 0000000000000000000000000000000000000000..ab11a34065f0107321da54ad5f3240587563a182
--- /dev/null
+++ b/.conda/aau_token_host
@@ -0,0 +1 @@
+zHxE_XAQ
\ No newline at end of file
diff --git a/LSAQ_CoreCode/lsaq_quant.py b/LSAQ_CoreCode/lsaq_quant.py
new file mode 100644
index 0000000000000000000000000000000000000000..494aebc03cda552c9d68d86d1ec6e842a054f091
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/LSAQ_CoreCode/main.ipynb b/LSAQ_CoreCode/main.ipynb
new file mode 100644
index 0000000000000000000000000000000000000000..3bd73b3e4767e3b11d0286d3fabdce2274427164
--- /dev/null
+++ b/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
+}
diff --git a/README.md b/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..e1f08db8d1c83d85e73a71c4785725c8f2c5d600
--- /dev/null
+++ b/README.md
@@ -0,0 +1 @@
+# quantization
\ No newline at end of file
diff --git a/__pycache__/zscore.cpython-310.pyc b/__pycache__/zscore.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f019d7e95de416bd885b6b7da640c55b527a5d92
Binary files /dev/null and b/__pycache__/zscore.cpython-310.pyc differ
diff --git a/__pycache__/zscore.cpython-311.pyc b/__pycache__/zscore.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..12cbb90eb09cc09b4e41916ccd18e9f629fa672f
Binary files /dev/null and b/__pycache__/zscore.cpython-311.pyc differ
diff --git a/baselines/Llama-2-7b-hf_alpha_idx_10.json b/baselines/Llama-2-7b-hf_alpha_idx_10.json
new file mode 100644
index 0000000000000000000000000000000000000000..e948afb0fcf30e40ab627bbae3f1182ae0dbb07c
--- /dev/null
+++ b/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
+]
\ No newline at end of file
diff --git a/baselines/Llama-2-7b-hf_alpha_idx_5.json b/baselines/Llama-2-7b-hf_alpha_idx_5.json
new file mode 100644
index 0000000000000000000000000000000000000000..41620c390d6fd8ea9eecb80c44fbdf719c6f1b44
--- /dev/null
+++ b/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
+]
\ No newline at end of file
diff --git a/baselines/Llama-2-7b-hf_kurtosis_idx_10.json b/baselines/Llama-2-7b-hf_kurtosis_idx_10.json
new file mode 100644
index 0000000000000000000000000000000000000000..0ac0e90089a37bb5a0a3427233dceabce7bed3de
--- /dev/null
+++ b/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
+]
\ No newline at end of file
diff --git a/baselines/Llama-2-7b-hf_kurtosis_idx_5.json b/baselines/Llama-2-7b-hf_kurtosis_idx_5.json
new file mode 100644
index 0000000000000000000000000000000000000000..d21b7074b69afa377c31b1f8e836769bd03d2259
--- /dev/null
+++ b/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
+]
\ No newline at end of file
diff --git a/baselines/Llama-2-7b-hf_z_idx_10.json b/baselines/Llama-2-7b-hf_z_idx_10.json
new file mode 100644
index 0000000000000000000000000000000000000000..f9c210e0e88b56cfc45828d1f26aab3cc4341fb2
--- /dev/null
+++ b/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
+]
\ No newline at end of file
diff --git a/baselines/Llama-2-7b-hf_z_idx_5.json b/baselines/Llama-2-7b-hf_z_idx_5.json
new file mode 100644
index 0000000000000000000000000000000000000000..93f91c3050675723ff1463c79c824e2734fab004
--- /dev/null
+++ b/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
+]
\ No newline at end of file
diff --git a/eval.sh b/eval.sh
new file mode 100644
index 0000000000000000000000000000000000000000..d3092ca054b58d1c09b50256c541fa6d9d83bb38
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/eval_coherence.sh b/eval_coherence.sh
new file mode 100644
index 0000000000000000000000000000000000000000..ac39c2c8b7d08e1abdda7709f1bf98d5085f2798
--- /dev/null
+++ b/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"
+
diff --git a/eval_fg.sh b/eval_fg.sh
new file mode 100644
index 0000000000000000000000000000000000000000..912ddfb5f729e089e255fb0bc438fb1caf79ea4a
--- /dev/null
+++ b/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"
+
diff --git a/eval_hd.sh b/eval_hd.sh
new file mode 100644
index 0000000000000000000000000000000000000000..867983c961a3eada02d72f6b5c202500ffde7718
--- /dev/null
+++ b/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"
diff --git a/eval_layer_llama.sh b/eval_layer_llama.sh
new file mode 100644
index 0000000000000000000000000000000000000000..14f3d9475be49bb07e0edc10610138234f4d4cf0
--- /dev/null
+++ b/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
+
+
diff --git a/eval_layer_qwen.sh b/eval_layer_qwen.sh
new file mode 100644
index 0000000000000000000000000000000000000000..cf4d4139cc49e1fecac0471838d5b907cc157092
--- /dev/null
+++ b/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
+
+
diff --git a/eval_zd.sh b/eval_zd.sh
new file mode 100644
index 0000000000000000000000000000000000000000..cb568fb9e70363bc57951dceed1df207116f255a
--- /dev/null
+++ b/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"
diff --git a/inference.py b/inference.py
new file mode 100644
index 0000000000000000000000000000000000000000..b5b20419c0174551c3a9857dd4f136ace4ba60a5
--- /dev/null
+++ b/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']))
\ No newline at end of file
diff --git a/layerwise-awq.py b/layerwise-awq.py
new file mode 100644
index 0000000000000000000000000000000000000000..280dd4591eb63c84b3b2ff6f5c4b3f893ae589ad
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/llm-awq/README.md b/llm-awq/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..2d8fff5ef785f1697dde3744c276d90bae14e5cb
--- /dev/null
+++ b/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)
+
diff --git a/llm-awq/awq/__pycache__/entry.cpython-311.pyc b/llm-awq/awq/__pycache__/entry.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..63d499804044a63a7b825a56224f5ba0d8f90dc9
Binary files /dev/null and b/llm-awq/awq/__pycache__/entry.cpython-311.pyc differ
diff --git a/llm-awq/awq/entry.py b/llm-awq/awq/entry.py
new file mode 100644
index 0000000000000000000000000000000000000000..134fc1c02a12ad5b34c3ab3e2e3ea528ccd5b1ae
--- /dev/null
+++ b/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()
diff --git a/llm-awq/awq/kernels/csrc/attention/README.md b/llm-awq/awq/kernels/csrc/attention/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..ec0aae55e81684ebce5f0d5fd13680470ced7845
--- /dev/null
+++ b/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 .
+```
diff --git a/llm-awq/awq/kernels/csrc/attention/cuda_bf16_fallbacks.cuh b/llm-awq/awq/kernels/csrc/attention/cuda_bf16_fallbacks.cuh
new file mode 100644
index 0000000000000000000000000000000000000000..f5641f61609172090da1c8e77e43f9f4694ccca0
--- /dev/null
+++ b/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
diff --git a/llm-awq/awq/kernels/csrc/attention/cuda_bf16_wrapper.h b/llm-awq/awq/kernels/csrc/attention/cuda_bf16_wrapper.h
new file mode 100644
index 0000000000000000000000000000000000000000..efb6e798730879bc2cd16088b2091991862a6074
--- /dev/null
+++ b/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
diff --git a/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention.cu b/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention.cu
new file mode 100644
index 0000000000000000000000000000000000000000..e5a6690086489a46c39452d7c9d3d14c7edf2ddf
--- /dev/null
+++ b/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
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention_utils.h b/llm-awq/awq/kernels/csrc/attention/decoder_masked_multihead_attention_utils.h
new file mode 100644
index 0000000000000000000000000000000000000000..7f6b7d81f69204c0f16cd5c2647196486c66ae08
--- /dev/null
+++ b/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
diff --git a/llm-awq/awq/kernels/csrc/attention/ft_attention.cpp b/llm-awq/awq/kernels/csrc/attention/ft_attention.cpp
new file mode 100644
index 0000000000000000000000000000000000000000..37d25bad49963ea1797cdd2b58202aefe1bcec13
--- /dev/null
+++ b/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;
+}
\ No newline at end of file
diff --git a/llm-awq/awq/kernels/csrc/attention/ft_attention.h b/llm-awq/awq/kernels/csrc/attention/ft_attention.h
new file mode 100644
index 0000000000000000000000000000000000000000..53037116aae1d7858c63870125a146b4c2f7df87
--- /dev/null
+++ b/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);
\ No newline at end of file
diff --git a/llm-awq/awq/kernels/csrc/layernorm/layernorm.cu b/llm-awq/awq/kernels/csrc/layernorm/layernorm.cu
new file mode 100644
index 0000000000000000000000000000000000000000..8f2de9a199f5f7f2bcbeac050da7dfce01192eba
--- /dev/null
+++ b/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);
+    });
+}
diff --git a/llm-awq/awq/kernels/csrc/position_embedding/pos_encoding.h b/llm-awq/awq/kernels/csrc/position_embedding/pos_encoding.h
new file mode 100644
index 0000000000000000000000000000000000000000..04e205b238c56142a5568294978338c592a2005d
--- /dev/null
+++ b/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);
\ No newline at end of file
diff --git a/llm-awq/awq/kernels/csrc/pybind.cpp b/llm-awq/awq/kernels/csrc/pybind.cpp
new file mode 100644
index 0000000000000000000000000000000000000000..30424ee05147bcb5d66f2df2a7819470934b7051
--- /dev/null
+++ b/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");
+}
diff --git a/llm-awq/awq/kernels/csrc/quantization/gemv_cuda.cu b/llm-awq/awq/kernels/csrc/quantization/gemv_cuda.cu
new file mode 100644
index 0000000000000000000000000000000000000000..3a55e662e82d77d5c9322fef32df0f24bc2bb911
--- /dev/null
+++ b/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;
+;}
+
diff --git a/llm-awq/awq/kernels/setup.py b/llm-awq/awq/kernels/setup.py
new file mode 100644
index 0000000000000000000000000000000000000000..a44e1a5bb2d88352e6da5de690e8a2964ffbb9ac
--- /dev/null
+++ b/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"],
+)
diff --git a/llm-awq/tinychat/models/falcon.py b/llm-awq/tinychat/models/falcon.py
new file mode 100644
index 0000000000000000000000000000000000000000..91021aaab1324653c014a4fb6b23f6a520a10f6b
--- /dev/null
+++ b/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()
diff --git a/llm-awq/tinychat/models/internvl3.py b/llm-awq/tinychat/models/internvl3.py
new file mode 100644
index 0000000000000000000000000000000000000000..7bbc462f8b3841424eead65166397243a6d7aa6e
--- /dev/null
+++ b/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,
+        )
+    
diff --git a/llm-awq/tinychat/models/nvila/builder.py b/llm-awq/tinychat/models/nvila/builder.py
new file mode 100644
index 0000000000000000000000000000000000000000..6a6c094c3b6651e40ea0cd4965091e253f956aab
--- /dev/null
+++ b/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
diff --git a/llm-awq/tinychat/models/nvila/llava_arch.py b/llm-awq/tinychat/models/nvila/llava_arch.py
new file mode 100644
index 0000000000000000000000000000000000000000..e1aeb4d23b733c1968f363577762a2fa6643ffa9
--- /dev/null
+++ b/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
diff --git a/llm-awq/tinychat/models/qwen2.py b/llm-awq/tinychat/models/qwen2.py
new file mode 100644
index 0000000000000000000000000000000000000000..e7a4593f3294173ad4ed51c88b4cfae900772da0
--- /dev/null
+++ b/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)
diff --git a/llm-awq/tinychat/models/vila_llama.py b/llm-awq/tinychat/models/vila_llama.py
new file mode 100644
index 0000000000000000000000000000000000000000..a2abacbc75c581ebd11c041b94ab258d7b1b3114
--- /dev/null
+++ b/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
diff --git a/llm-awq/tinychat/modules/fused_attn.py b/llm-awq/tinychat/modules/fused_attn.py
new file mode 100644
index 0000000000000000000000000000000000000000..90a0f3db9541aa1805747cc0ad52f6bb0ceddfca
--- /dev/null
+++ b/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)
diff --git a/llm-awq/tinychat/modules/fused_internencoder.py b/llm-awq/tinychat/modules/fused_internencoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..2ba12a2ca4d5dc0718cf424548d7a6851547ec53
--- /dev/null
+++ b/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
+
+
diff --git a/llm-awq/tinychat/modules/fused_norm.py b/llm-awq/tinychat/modules/fused_norm.py
new file mode 100644
index 0000000000000000000000000000000000000000..e8e1f0d49709f510957d91f58abf829ffeaf6efd
--- /dev/null
+++ b/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)
diff --git a/llm-awq/tinychat/modules/fused_siglipdecoder.py b/llm-awq/tinychat/modules/fused_siglipdecoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..4614baa4df75f896f1a9b8f6234ff17fe263c858
--- /dev/null
+++ b/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,
+        )
diff --git a/llm-awq/tinychat/modules/fused_vision_attn.py b/llm-awq/tinychat/modules/fused_vision_attn.py
new file mode 100644
index 0000000000000000000000000000000000000000..bda32e53edc03c0c697dbffa88e25c40f9d27feb
--- /dev/null
+++ b/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)
diff --git a/llm-awq/tinychat/scripts/internvl_demo.sh b/llm-awq/tinychat/scripts/internvl_demo.sh
new file mode 100644
index 0000000000000000000000000000000000000000..de0933bc9a1865e6d7bb682e33855731d83e0770
--- /dev/null
+++ b/llm-awq/tinychat/scripts/internvl_demo.sh
@@ -0,0 +1,18 @@
+MODEL_PATH=PATH_TO_INTERNVL
+MODEL_NAME=InternVL3-8B
+
+# run AWQ search
+python -m awq.entry --model_path $MODEL_PATH \
+    --w_bit 4 --q_group_size 128 \
+    --run_awq --dump_awq awq_cache/$MODEL_NAME-w4-g128.pt
+
+# generate real quantized weights (w4)
+python -m awq.entry --model_path $MODEL_PATH \
+    --w_bit 4 --q_group_size 128 --load_awq awq_cache/$MODEL_NAME-w4-g128.pt \
+    --q_backend real --dump_quant quant_cache/$MODEL_NAME-w4-128-awq.pt
+
+# Run the TinyChat demo:
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python internvl_demo.py --model-path $MODEL_PATH \
+    --quant_path quant_cache/$MODEL_NAME-w4-128-awq-v2.pt \
+    --media ../figures/vila-logo.jpg --max_seq_len 4096 --chunk \
+    --model_type internvl3 --quant_VT --quant_llm
\ No newline at end of file
diff --git a/llm-awq/tinychat/scripts/nvila_demo.sh b/llm-awq/tinychat/scripts/nvila_demo.sh
new file mode 100644
index 0000000000000000000000000000000000000000..58f8bbc7c7150b37522d3d3e872d5061161c84ca
--- /dev/null
+++ b/llm-awq/tinychat/scripts/nvila_demo.sh
@@ -0,0 +1,22 @@
+MODEL_PATH=PATH_TO_NVILA
+MODEL_NAME=NVILA-8B
+
+# run AWQ search
+python -m awq.entry --model_path $MODEL_PATH \
+    --smooth_scale --media_path https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-VL/space_woaudio.mp4  \
+    --act_scale_path awq_cache/$MODEL_NAME-smooth-scale.pt --vila-20 \
+    --w_bit 4 --q_group_size 128 \
+    --run_awq --dump_awq awq_cache/$MODEL_NAME.pt
+
+# generate real quantized weights (w4)
+python -m awq.entry --model_path $MODEL_PATH/llm \
+    --w_bit 4 --q_group_size 128  \
+    --load_awq awq_cache/$MODEL_NAME.pt \
+    --q_backend real --dump_quant quant_cache/$MODEL_NAME-w4-g128-awq.pt --vila-20
+
+# Run the TinyChat demo:
+python nvila_demo.py --model-path $MODEL_PATH    \
+       --quant_path quant_cache/$MODEL_NAME-w4-g128-awq.pt      \
+       --media ../figures/nvila-logo.jpg     \
+       --act_scale_path awq_cache/$MODEL_NAME-smooth-scale.pt    \
+       --all --chunk --model_type nvila --vis_image
\ No newline at end of file
diff --git a/llm-awq/tinychat/serve/examples/CPR.jpg b/llm-awq/tinychat/serve/examples/CPR.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..d7793d8b3281146fe79b97a5546891159bf934d0
Binary files /dev/null and b/llm-awq/tinychat/serve/examples/CPR.jpg differ
diff --git a/llm-awq/tinychat/serve/examples/icl-logo/adobe.jpg b/llm-awq/tinychat/serve/examples/icl-logo/adobe.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..73ae44cb7cc6a596910c33f8d3005bc11f1dd9b6
Binary files /dev/null and b/llm-awq/tinychat/serve/examples/icl-logo/adobe.jpg differ
diff --git a/llm-awq/tinychat/serve/examples/icl-logo/apple.jpg b/llm-awq/tinychat/serve/examples/icl-logo/apple.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..c8259d04d1d9f292ed95ec1083e720cf3865e8d5
Binary files /dev/null and b/llm-awq/tinychat/serve/examples/icl-logo/apple.jpg differ
diff --git a/llm-awq/tinychat/serve/examples/icl-logo/google.webp b/llm-awq/tinychat/serve/examples/icl-logo/google.webp
new file mode 100644
index 0000000000000000000000000000000000000000..2fbba7a73c686782e915097791a0ec46d576e66f
Binary files /dev/null and b/llm-awq/tinychat/serve/examples/icl-logo/google.webp differ
diff --git a/llm-awq/tinychat/serve/examples/icl-logo/nvidia.png b/llm-awq/tinychat/serve/examples/icl-logo/nvidia.png
new file mode 100644
index 0000000000000000000000000000000000000000..980e6f4cfc9f6732babbad637af6d8b2e40eb042
Binary files /dev/null and b/llm-awq/tinychat/serve/examples/icl-logo/nvidia.png differ
diff --git a/llm-awq/tinychat/serve/gradio_web_server.py b/llm-awq/tinychat/serve/gradio_web_server.py
new file mode 100644
index 0000000000000000000000000000000000000000..f7de6e4a5b319521bb1123bf028298723542db21
--- /dev/null
+++ b/llm-awq/tinychat/serve/gradio_web_server.py
@@ -0,0 +1,1200 @@
+#    Modified from https://github.com/haotian-liu/LLaVA
+#    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 argparse
+import datetime
+import json
+import os
+import time
+
+import gradio as gr
+import requests
+
+LOGDIR = "."
+
+from tinychat.serve.llava_conv import (
+    default_conversation,
+    conv_templates,
+    get_conversation,
+    SeparatorStyle,
+)
+from tinychat.utils.log_utils import (
+    build_logger,
+    server_error_msg,
+    violates_moderation,
+    moderation_msg,
+)
+import hashlib
+
+IMAGE_BOX_NUM = 3
+BUTTON_LIST_LEN = 2
+
+logger = build_logger("gradio_web_server", "gradio_web_server.log")
+
+headers = {"User-Agent": "TinyChat AWQ Chatbot"}
+
+no_change_btn = gr.Button.update()
+enable_btn = gr.Button.update(interactive=True)
+disable_btn = gr.Button.update(interactive=False)
+
+from tinychat.utils.constants import (
+    LLAVA_DEFAULT_IMAGE_TOKEN,
+    LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER,
+    AUTO_FILL_IM_TOKEN_HOLDER,
+)
+
+# IMAGE_TOKEN_VIS = "**[IMAGE]**"
+IMAGE_TOKEN_VIS = "**\<image\>**"
+
+priority = {
+    "vicuna-13b": "aaaaaaa",
+    "koala-13b": "aaaaaab",
+}
+
+
+def get_conv_log_filename():
+    t = datetime.datetime.now()
+    name = os.path.join(LOGDIR, f"{t.year}-{t.month:02d}-{t.day:02d}-conv.json")
+    return name
+
+
+def get_model_list():
+    ret = requests.post(args.controller_url + "/refresh_all_workers")
+    assert ret.status_code == 200
+    ret = requests.post(args.controller_url + "/list_models")
+    models = ret.json()["models"]
+    models.sort(key=lambda x: priority.get(x, x))
+    logger.info(f"Models: {models}")
+    return models
+
+
+get_window_url_params = """
+function() {
+    const params = new URLSearchParams(window.location.search);
+    url_params = Object.fromEntries(params);
+    console.log(url_params);
+    return url_params;
+    }
+"""
+
+
+def load_demo(url_params, prompt_style_btn, request: gr.Request):
+    logger.info(f"load_demo. ip: {request.client.host}. params: {url_params}")
+
+    dropdown_update = gr.Dropdown.update(visible=True)
+    if "model" in url_params:
+        model = url_params["model"]
+        if model in models:
+            dropdown_update = gr.Dropdown.update(value=model, visible=True)
+    state = get_conversation(prompt_style_btn)
+    # state = default_conversation.copy()
+    return state, dropdown_update
+
+
+def load_demo_refresh_model_list(prompt_style_btn, request: gr.Request):
+    logger.info(f"load_demo. ip: {request.client.host}")
+    models = get_model_list()
+    state = get_conversation(prompt_style_btn)
+    # state = default_conversation.copy()
+    dropdown_update = gr.Dropdown.update(
+        choices=models, value=models[0] if len(models) > 0 else ""
+    )
+    return state, dropdown_update
+
+
+# def vote_last_response(state, vote_type, model_selector, request: gr.Request):
+#     with open(get_conv_log_filename(), "a") as fout:
+#         data = {
+#             "tstamp": round(time.time(), 4),
+#             "type": vote_type,
+#             "model": model_selector,
+#             "state": state.dict(),
+#             "ip": request.client.host,
+#         }
+#         fout.write(json.dumps(data) + "\n")
+
+
+# def upvote_last_response(state, model_selector, request: gr.Request):
+#     logger.info(f"upvote. ip: {request.client.host}")
+#     vote_last_response(state, "upvote", model_selector, request)
+#     return ("",) + (disable_btn,) * 3
+
+
+# def downvote_last_response(state, model_selector, request: gr.Request):
+#     logger.info(f"downvote. ip: {request.client.host}")
+#     vote_last_response(state, "downvote", model_selector, request)
+#     return ("",) + (disable_btn,) * 3
+
+
+# def flag_last_response(state, model_selector, request: gr.Request):
+#     logger.info(f"flag. ip: {request.client.host}")
+#     vote_last_response(state, "flag", model_selector, request)
+#     return ("",) + (disable_btn,) * 3
+
+
+def regenerate(state, image_process_mode, request: gr.Request):
+    logger.info(f"regenerate. ip: {request.client.host}")
+    state.messages[-1][-1] = None
+    prev_human_msg = state.messages[-2]
+    if type(prev_human_msg[1]) in (tuple, list):
+        prev_human_msg[1] = (*prev_human_msg[1][:2], image_process_mode)
+    state.skip_next = False
+    return (state, state.to_gradio_chatbot(), "") + (disable_btn,) * BUTTON_LIST_LEN
+
+
+def change_prompt_style(state, prompt_style_btn, request: gr.Request):
+    if state.version != prompt_style_btn:
+        state = get_conversation(prompt_style_btn)
+    return state
+
+
+def clear_history(prompt_style_btn, request: gr.Request):
+    logger.info(f"clear_history. ip: {request.client.host}")
+    state = get_conversation(prompt_style_btn)
+    return (
+        (state, state.to_gradio_chatbot(), "")
+        + (None,) * IMAGE_BOX_NUM
+        + (None,)  # Videobox
+        + (disable_btn,) * BUTTON_LIST_LEN
+    )
+
+
+def clear_text_history(state, prompt_style_btn, request: gr.Request):
+    state = get_conversation(prompt_style_btn)
+    return (state, state.to_gradio_chatbot())
+
+
+def clear_after_click_example_1_video(videobox, textbox):
+    imagebox = None
+    imagebox_2 = None
+    imagebox_3 = None
+    state = get_conversation("default")
+    prompt_style_btn = "default"
+    return (state, imagebox, imagebox_2, imagebox_3, videobox, prompt_style_btn)
+
+
+def clear_after_click_example_1_image(imagebox, textbox):
+    imagebox_2 = None
+    imagebox_3 = None
+    videobox = None
+    state = get_conversation("default")
+    prompt_style_btn = "default"
+    return (state, imagebox, imagebox_2, imagebox_3, videobox, prompt_style_btn)
+
+
+def clear_after_click_example_2_image(imagebox, imagebox_2, textbox):
+    imagebox_3 = None
+    videobox = None
+    state = get_conversation("default")
+    prompt_style_btn = "default"
+    return (state, imagebox, imagebox_2, imagebox_3, videobox, prompt_style_btn)
+
+
+def clear_after_click_example_3_image(imagebox, imagebox_2, imagebox_3, textbox):
+    videobox = None
+    state = get_conversation("default")
+    prompt_style_btn = "default"
+    return (state, imagebox, imagebox_2, imagebox_3, videobox, prompt_style_btn)
+
+
+def clear_after_click_example_3_image_icl(imagebox, imagebox_2, imagebox_3, textbox):
+    videobox = None
+    state = get_conversation("no-sys")
+    prompt_style_btn = "no-sys"
+    return (state, imagebox, imagebox_2, imagebox_3, videobox, prompt_style_btn)
+
+
+def add_images(
+    state,
+    imagebox,
+    imagebox_2,
+    imagebox_3,
+    videobox,
+    image_process_mode,
+    request: gr.Request,
+):
+    if state.image_loaded:
+        # return (state,) + (None,) * IMAGE_BOX_NUM
+        return state
+
+    def extract_frames(video_path):
+        import cv2
+        from PIL import Image
+
+        vidcap = cv2.VideoCapture(video_path)
+        fps = vidcap.get(cv2.CAP_PROP_FPS)
+        frame_count = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+        duration = frame_count / fps
+
+        frame_interval = frame_count // 8
+        print(
+            "duration:", duration, "frames:", frame_count, "intervals:", frame_interval
+        )
+        # frame_interval = 10
+
+        def get_frame(max_frames):
+            # frame_id = int(fps * stamp)
+            # vidcap.set(cv2.CAP_PROP_POS_FRAMES, frame_id)
+            # ret, frame = vidcap.read()
+            images = []
+            count = 0
+            success = True
+            while success:
+                success, frame = vidcap.read()
+                if count % frame_interval == 0:
+                    img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                    im_pil = Image.fromarray(img)
+                    images.append(im_pil)
+                    if len(images) == max_frames:
+                        return images
+
+                count += 1
+            # assert ret, "videocap.read fails!"
+            # img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            # im_pil = Image.fromarray(img)
+            # print(f"loading {stamp} success")
+            return images
+
+        # return [get_frame(0), get_frame(stamp1), get_frame(stamp2)]
+        # img = get_frame(0)
+        # img1 = get_frame(frame_interval * 1)
+        # return [img, img1, img, img1, img, img1,]
+        return get_frame(8)
+
+    frames = [
+        None,
+    ]
+    if videobox is not None:
+        frames = extract_frames(videobox)
+        # add frames as regular images
+        logger.info(f"Got videobox: {videobox}.")
+
+    logger.info(f"add_image. ip: {request.client.host}.")
+    image_list = [imagebox, imagebox_2, imagebox_3, *frames]
+    logger.info(f"image_list: {image_list}")
+
+    im_count = 0
+    for image in image_list:
+        if image is not None:
+            im_count += 1
+    for image in image_list:
+        if image is not None:
+            if args.auto_pad_image_token or im_count == 1:
+                text = (AUTO_FILL_IM_TOKEN_HOLDER, image, image_process_mode)
+            else:
+                text = ("", image, image_process_mode)
+            state.append_message(None, text)
+            state.append_message(
+                None, None
+            )  # in order to match the input-output pair for textbox outputs
+            # state.append_message(state.roles[0], text)
+            # state.append_message(state.roles[1], None)
+    # state.skip_next = False
+    logger.info(f"im_count {im_count}. ip: {request.client.host}.")
+    state.image_loaded = True
+    # return (state,) + (None,) * IMAGE_BOX_NUM
+    return state
+
+
+def add_text_only(state, text, request: gr.Request):
+    logger.info(f"add_text_only. ip: {request.client.host}. len: {len(text)}")
+
+    if args.moderate:
+        flagged = violates_moderation(text)
+        if flagged:
+            state.skip_next = True
+            return (state, moderation_msg) + (no_change_btn,) * BUTTON_LIST_LEN
+
+    # This is 1536 characters, rather than tokens
+    text = text[:1536]  # Hard cut-off
+    state.append_message(state.roles[0], text)
+    state.append_message(state.roles[1], None)
+    state.skip_next = False
+    return (state, "") + (disable_btn,) * BUTTON_LIST_LEN
+
+
+def add_text(
+    state, text, image, image_process_mode, prompt_style_btn, request: gr.Request
+):
+    logger.info(f"add_text. ip: {request.client.host}. len: {len(text)}")
+    if len(text) <= 0 and image is None:
+        state.skip_next = True
+        return (state, state.to_gradio_chatbot(), "", None) + (
+            no_change_btn,
+        ) * BUTTON_LIST_LEN
+    if args.moderate:
+        flagged = violates_moderation(text)
+        if flagged:
+            state.skip_next = True
+            return (state, state.to_gradio_chatbot(), moderation_msg, None) + (
+                no_change_btn,
+            ) * BUTTON_LIST_LEN
+
+    text = text[:1536]  # Hard cut-off
+    if image is not None:
+        text = text[:1200]  # Hard cut-off for images
+        if "<image>" not in text:
+            # text = '<Image><image></Image>' + text
+            text = text + "\n<image>"
+        text = (text, image, image_process_mode)
+        if len(state.get_images(return_pil=True)) > 0:
+            state = get_conversation(prompt_style_btn)
+            # state = default_conversation.copy()
+    state.append_message(state.roles[0], text)
+    state.append_message(state.roles[1], None)
+    state.skip_next = False
+    return (state, state.to_gradio_chatbot(), "", None) + (
+        disable_btn,
+    ) * BUTTON_LIST_LEN
+
+
+def http_bot(
+    state,
+    model_selector,
+    temperature,
+    top_p,
+    max_new_tokens,
+    prompt_style_btn,
+    request: gr.Request,
+):
+    logger.info(f"http_bot. ip: {request.client.host}")
+    start_tstamp = time.time()
+    model_name = model_selector
+
+    if state.skip_next:
+        # This generate call is skipped due to invalid inputs
+        yield (state, state.to_gradio_chatbot()) + (no_change_btn,) * BUTTON_LIST_LEN
+        return
+    if len(state.messages) == state.offset + 2:
+        # First round of conversation
+        if "llava" in model_name.lower():
+            if "llama-2" in model_name.lower():
+                template_name = "llava_llama_2"
+            elif "v1" in model_name.lower():
+                if "mmtag" in model_name.lower():
+                    template_name = "v1_mmtag"
+                elif (
+                    "plain" in model_name.lower()
+                    and "finetune" not in model_name.lower()
+                ):
+                    template_name = "v1_mmtag"
+                else:
+                    template_name = "llava_v1"
+            elif "mpt" in model_name.lower():
+                template_name = "mpt"
+            else:
+                if "mmtag" in model_name.lower():
+                    template_name = "v0_mmtag"
+                elif (
+                    "plain" in model_name.lower()
+                    and "finetune" not in model_name.lower()
+                ):
+                    template_name = "v0_mmtag"
+                else:
+                    template_name = "llava_v0"
+        elif "mpt" in model_name:
+            template_name = "mpt_text"
+        elif "llama-2" in model_name:
+            template_name = "llama_2"
+        else:
+            template_name = "vicuna_v1"
+        if prompt_style_btn == "no-sys":
+            new_state = get_conversation(prompt_style_btn)
+        else:
+            new_state = conv_templates[template_name].copy()
+        new_state.append_message(new_state.roles[0], state.messages[-2][1])
+        new_state.append_message(new_state.roles[1], None)
+        state = new_state
+
+    # Query worker address
+    controller_url = args.controller_url
+    ret = requests.post(
+        controller_url + "/get_worker_address", json={"model": model_name}
+    )
+    worker_addr = ret.json()["address"]
+    logger.info(f"model_name: {model_name}, worker_addr: {worker_addr}")
+
+    # No available worker
+    if worker_addr == "":
+        state.messages[-1][-1] = server_error_msg
+        yield (
+            state,
+            state.to_gradio_chatbot(),
+            # disable_btn,
+            # disable_btn,
+            # disable_btn,
+            enable_btn,
+            enable_btn,
+        )
+        return
+
+    # Construct prompt
+    prompt = state.get_prompt()
+
+    all_images = state.get_images(return_pil=True)
+    all_image_hash = [hashlib.md5(image.tobytes()).hexdigest() for image in all_images]
+    for image, hash in zip(all_images, all_image_hash):
+        t = datetime.datetime.now()
+        filename = os.path.join(
+            LOGDIR, "serve_images", f"{t.year}-{t.month:02d}-{t.day:02d}", f"{hash}.jpg"
+        )
+        if not os.path.isfile(filename):
+            os.makedirs(os.path.dirname(filename), exist_ok=True)
+            image.save(filename)
+
+    # Make requests
+    pload = {
+        "model": model_name,
+        "prompt": prompt,
+        "temperature": float(temperature),
+        "top_p": float(top_p),
+        "max_new_tokens": min(int(max_new_tokens), 1536),
+        "stop": (
+            state.sep
+            if state.sep_style in [SeparatorStyle.SINGLE, SeparatorStyle.MPT]
+            else state.sep2
+        ),
+        "images": f"List of {len(state.get_images())} images: {all_image_hash}",
+    }
+
+    image_num = len(state.get_images())
+    if image_num == 0:
+        state.messages[-1][
+            -1
+        ] = "**NO INPUT IMAGE RECEIVED BY THE SERVER. PLEASE CHECK YOUR INTERNET CONNECTION AND REFRESH THE PAGE.**"
+        yield (
+            state,
+            state.to_gradio_chatbot(),
+            # disable_btn,
+            # disable_btn,
+            # disable_btn,
+            enable_btn,
+            enable_btn,
+        )
+        return
+
+    count_auto_im_token = prompt.count(AUTO_FILL_IM_TOKEN_HOLDER)
+    count_manual_im_token = prompt.count(LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER)
+    if (count_auto_im_token == image_num) and (
+        count_manual_im_token == 0
+    ):  # Use default system prompt
+        prompt = prompt.replace(AUTO_FILL_IM_TOKEN_HOLDER, LLAVA_DEFAULT_IMAGE_TOKEN)
+    elif (count_auto_im_token == image_num) and (
+        count_manual_im_token == image_num
+    ):  # Use <image> token inserted by user
+        prompt = prompt.replace(AUTO_FILL_IM_TOKEN_HOLDER, "")
+        prompt = prompt.replace(
+            LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER, LLAVA_DEFAULT_IMAGE_TOKEN
+        )
+    elif (count_auto_im_token == 0) and (count_manual_im_token == image_num):
+        prompt = prompt.replace(
+            LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER, LLAVA_DEFAULT_IMAGE_TOKEN
+        )
+    else:
+        state.messages[-1][
+            -1
+        ] = "**IMAGE NUM MISMATCHES IMAGE TOKEN PLACEHOLDER. PLEASE CHECK YOUR INPUT AND REFRESH THE PAGE.**"
+        yield (
+            state,
+            state.to_gradio_chatbot(),
+            # disable_btn,
+            # disable_btn,
+            # disable_btn,
+            enable_btn,
+            enable_btn,
+        )
+        return
+
+    pload["prompt"] = prompt
+    logger.info(f"==== request ====\n{pload}")
+    pload["images"] = state.get_images()
+
+    state.messages[-1][-1] = "▌"
+    ret = state.to_gradio_chatbot()
+    ret[0][0] = ret[0][0].replace(LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER, IMAGE_TOKEN_VIS)
+    yield (state, ret) + (disable_btn,) * BUTTON_LIST_LEN
+
+    try:
+        # Stream output
+        response = requests.post(
+            worker_addr + "/worker_generate_stream",
+            headers=headers,
+            json=pload,
+            stream=True,
+            timeout=10,
+        )
+        for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
+            if chunk:
+                data = json.loads(chunk.decode())
+                if data["error_code"] == 0:
+                    output = data["text"][len(prompt) :].strip()
+                    state.messages[-1][-1] = output + "▌"
+                    ret = state.to_gradio_chatbot()
+                    ret[0][0] = ret[0][0].replace(
+                        LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER, IMAGE_TOKEN_VIS
+                    )
+                    yield (state, ret) + (disable_btn,) * BUTTON_LIST_LEN
+                else:
+                    output = data["text"] + f" (error_code: {data['error_code']})"
+                    state.messages[-1][-1] = output
+                    ret = state.to_gradio_chatbot()
+                    ret[0][0] = ret[0][0].replace(
+                        LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER, IMAGE_TOKEN_VIS
+                    )
+                    yield (state, ret) + (
+                        # disable_btn,
+                        # disable_btn,
+                        # disable_btn,
+                        enable_btn,
+                        enable_btn,
+                    )
+                    return
+                time.sleep(0.03)
+    except requests.exceptions.RequestException as e:
+        state.messages[-1][-1] = server_error_msg
+        yield (state, state.to_gradio_chatbot()) + (
+            # disable_btn,
+            # disable_btn,
+            # disable_btn,
+            enable_btn,
+            enable_btn,
+        )
+        return
+
+    state.messages[-1][-1] = state.messages[-1][-1][:-1]
+    ret = state.to_gradio_chatbot()
+    ret[0][0] = ret[0][0].replace(LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER, IMAGE_TOKEN_VIS)
+    yield (state, ret) + (enable_btn,) * BUTTON_LIST_LEN
+
+    finish_tstamp = time.time()
+    logger.info(f"{output}")
+
+    with open(get_conv_log_filename(), "a") as fout:
+        data = {
+            "tstamp": round(finish_tstamp, 4),
+            "type": "chat",
+            "model": model_name,
+            "start": round(start_tstamp, 4),
+            "finish": round(finish_tstamp, 4),
+            "state": state.dict(),
+            "images": all_image_hash,
+            "ip": request.client.host,
+        }
+        fout.write(json.dumps(data) + "\n")
+
+
+title_markdown = """
+# VILA: On Pre-training for Visual Language Models
+[\[Paper\]](https://arxiv.org/abs/2312.07533)  [\[Github\]](https://github.com/NVlabs/VILA)
+### Powered by [TinyChat](https://github.com/mit-han-lab/llm-awq/tree/main/tinychat) with 4-bit [AWQ](https://arxiv.org/abs/2306.00978).
+"""
+
+tos_markdown = """
+### Terms of Use
+By using this service, users are required to agree to the following terms:
+The service is a research preview intended for non-commercial use only. It only provides limited safety measures and may generate offensive content. It must not be used for any illegal, harmful, violent, racist, or sexual purposes. The service may collect user dialogue data for future research.
+Please click the "Flag" button if you get any inappropriate answer! We will collect those to keep improving our moderator.
+For an optimal experience, please use desktop computers for this demo, as mobile devices may compromise its quality.
+"""
+
+
+learn_more_markdown = """
+### License
+The service is a research preview intended for non-commercial use only, subject to the model [License](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) of LLaMA, [Terms of Use](https://openai.com/policies/terms-of-use) of the data generated by OpenAI, and [Privacy Practices](https://chrome.google.com/webstore/detail/sharegpt-share-your-chatg/daiacboceoaocpibfodeljbdfacokfjb) of ShareGPT. Please contact us if you find any potential violation.
+"""
+
+ack_markdown = """
+### Acknowledgement
+This demo is inspired by [LLaVA](https://github.com/haotian-liu/LLaVA). We thank LLaVA for providing an elegant way to build the Gradio Web UI.          
+"""
+
+block_css = """
+
+#buttons button {
+    min-width: min(120px,100%);
+}
+
+"""
+
+
+def build_demo(embed_mode):
+    textbox = gr.Textbox(
+        show_label=False, placeholder="Enter text and press ENTER", container=False
+    )
+    with gr.Blocks(
+        title="VILA on TinyChat", theme=gr.themes.Default(), css=block_css
+    ) as demo:
+        state = gr.State()
+
+        if not embed_mode:
+            gr.Markdown(title_markdown)
+
+        with gr.Row():
+            with gr.Column(scale=8):
+                with gr.Row():
+                    imagebox = gr.Image(type="pil")
+                    imagebox_2 = gr.Image(type="pil")
+                    imagebox_3 = gr.Image(type="pil")
+                    videobox = gr.Video(label="1 video = 8 frames")
+                image_process_mode = gr.Radio(
+                    ["Crop", "Resize", "Pad", "Default"],
+                    value="Default",
+                    label="Preprocess for non-square image",
+                    visible=False,
+                )
+                # imagebox_out = gr.Image(height=150)
+                with gr.Row():
+                    with gr.Column(scale=5):
+                        textbox.render()
+                    with gr.Column(scale=1, min_width=100):
+                        submit_btn = gr.Button(value="Send", variant="primary")
+                    with gr.Column(scale=1, min_width=100):
+                        clear_btn = gr.Button(
+                            value="🗑️  Clear", variant="primary", interactive=False
+                        )
+                    with gr.Column(scale=1, min_width=100):
+                        regenerate_btn = gr.Button(
+                            value="🔄  Retry", variant="primary", interactive=False
+                        )
+                with gr.Row():
+                    gr.Markdown(
+                        "### *** Before changing the current images, uploading new images or switching the prompt style, please click the clear button."
+                    )
+                chatbot = gr.Chatbot(
+                    elem_id="chatbot", label="TinyChat Assistant", height=550
+                )
+
+            with gr.Column(scale=4):
+                with gr.Row(equal_height=True):
+                    with gr.Column(scale=1, min_width=50):
+                        model_selector = gr.Dropdown(
+                            choices=models,
+                            value=models[0] if len(models) > 0 else "",
+                            label="Model",
+                            interactive=True,
+                            show_label=True,
+                            container=False,
+                        )
+                    with gr.Column(scale=1, min_width=50):
+                        prompt_style_btn = gr.Radio(
+                            ["default", "no-sys"],
+                            label="Prompt style",
+                            value="default",
+                            interactive=True,
+                        )
+
+                # with gr.Row():
+                # with gr.Column(scale=1, min_width=50):
+                #     im_submit_btn = gr.Button(value="Add image", variant="primary")
+                # with gr.Column(scale=1, min_width=50):
+                #     submit_btn_1 = gr.Button(value="Send", variant="primary")
+
+                cur_dir = os.path.dirname(os.path.abspath(__file__))
+                with gr.Row(equal_height=True):
+                    gr.Examples(
+                        examples=[
+                            [
+                                f"{cur_dir}/examples/video/qZDF__7LNKc.4.mp4",
+                                "Elaborate on the visual and narrative elements of the video in detail.",
+                            ],
+                        ],
+                        label="Video Example",
+                        inputs=[videobox, textbox],
+                        fn=clear_after_click_example_1_video,
+                        outputs=[
+                            state,
+                            imagebox,
+                            imagebox_2,
+                            imagebox_3,
+                            videobox,
+                            prompt_style_btn,
+                        ],
+                        run_on_click=True,
+                    )
+                with gr.Row(equal_height=True):
+                    with gr.Column(scale=1, min_width=50):
+                        gr.Examples(
+                            examples=[
+                                [
+                                    f"{cur_dir}/examples/pedestrain.png",
+                                    "<image> What is the person in the center of the image doing?",
+                                ],
+                            ],
+                            label="Image Example 1",
+                            inputs=[imagebox, textbox],
+                            fn=clear_after_click_example_1_image,
+                            outputs=[
+                                state,
+                                imagebox,
+                                imagebox_2,
+                                imagebox_3,
+                                videobox,
+                                prompt_style_btn,
+                            ],
+                            run_on_click=True,
+                        )
+                    with gr.Column(scale=1, min_width=50):
+                        gr.Examples(
+                            examples=[
+                                [
+                                    f"{cur_dir}/examples/car_repair.png",
+                                    "<image> What is the brand of the silver car in the image?",
+                                ],
+                            ],
+                            label="Image Example 2",
+                            inputs=[imagebox, textbox],
+                            fn=clear_after_click_example_1_image,
+                            outputs=[
+                                state,
+                                imagebox,
+                                imagebox_2,
+                                imagebox_3,
+                                videobox,
+                                prompt_style_btn,
+                            ],
+                            run_on_click=True,
+                        )
+                with gr.Row(equal_height=True):
+                    with gr.Column(scale=1, min_width=50):
+                        gr.Examples(
+                            examples=[
+                                [
+                                    f"{cur_dir}/examples/CPR.jpg",
+                                    "<image> What are the people doing in this image?",
+                                ],
+                            ],
+                            label="Image Example 3",
+                            inputs=[imagebox, textbox],
+                            fn=clear_after_click_example_1_image,
+                            outputs=[
+                                state,
+                                imagebox,
+                                imagebox_2,
+                                imagebox_3,
+                                videobox,
+                                prompt_style_btn,
+                            ],
+                            run_on_click=True,
+                        )
+                    with gr.Column(scale=1, min_width=50):
+                        gr.Examples(
+                            examples=[
+                                [
+                                    f"{cur_dir}/examples/Wall_fissure.png",
+                                    "<image> What are the likely service needed for this building?",
+                                ],
+                            ],
+                            label="Image Example 4",
+                            inputs=[imagebox, textbox],
+                            fn=clear_after_click_example_1_image,
+                            outputs=[
+                                state,
+                                imagebox,
+                                imagebox_2,
+                                imagebox_3,
+                                videobox,
+                                prompt_style_btn,
+                            ],
+                            run_on_click=True,
+                        )
+
+                with gr.Row(equal_height=True):
+                    with gr.Column(scale=1, min_width=50):
+                        gr.Examples(
+                            examples=[
+                                [
+                                    f"{cur_dir}/examples/animal_blocking.png",
+                                    "<image> What is unusual in this image?",
+                                ],
+                            ],
+                            label="Image Example 5",
+                            inputs=[imagebox, textbox],
+                            fn=clear_after_click_example_1_image,
+                            outputs=[
+                                state,
+                                imagebox,
+                                imagebox_2,
+                                imagebox_3,
+                                videobox,
+                                prompt_style_btn,
+                            ],
+                            run_on_click=True,
+                        )
+                    with gr.Column(scale=1, min_width=50):
+                        gr.Examples(
+                            examples=[
+                                [
+                                    f"{cur_dir}/examples/windmill_people.png",
+                                    "<image> Can you describe what is happening?",
+                                ],
+                            ],
+                            label="Image Example 6",
+                            inputs=[imagebox, textbox],
+                            fn=clear_after_click_example_1_image,
+                            outputs=[
+                                state,
+                                imagebox,
+                                imagebox_2,
+                                imagebox_3,
+                                videobox,
+                                prompt_style_btn,
+                            ],
+                            run_on_click=True,
+                        )
+
+                gr.Examples(
+                    examples=[
+                        [
+                            f"{cur_dir}/examples/climate_change/climate_change_1.png",
+                            f"{cur_dir}/examples/climate_change/climate_change_2.png",
+                            "<image> <image> What is the implication of temperature based on this image?",
+                        ],
+                    ],
+                    inputs=[imagebox, imagebox_2, textbox],
+                    label="Multi-image Example 1",
+                    fn=clear_after_click_example_2_image,
+                    outputs=[
+                        state,
+                        imagebox,
+                        imagebox_2,
+                        imagebox_3,
+                        videobox,
+                        prompt_style_btn,
+                    ],
+                    run_on_click=True,
+                )
+
+                gr.Examples(
+                    examples=[
+                        [
+                            f"{cur_dir}/examples/palms/palm1.png",
+                            f"{cur_dir}/examples/palms/palm2.png",
+                            f"{cur_dir}/examples/palms/palm3.png",
+                            "8:15am: <image> 12:45pm: <image> 16:00pm: <image> When did I have lunch and what did I eat for lunch?",
+                        ],
+                    ],
+                    inputs=[imagebox, imagebox_2, imagebox_3, textbox],
+                    label="Multi-image Example 2",
+                    fn=clear_after_click_example_3_image,
+                    outputs=[
+                        state,
+                        imagebox,
+                        imagebox_2,
+                        imagebox_3,
+                        videobox,
+                        prompt_style_btn,
+                    ],
+                    run_on_click=True,
+                )
+
+                gr.Examples(
+                    examples=[
+                        [
+                            f"{cur_dir}/examples/golf/Golfman1.png",
+                            f"{cur_dir}/examples/golf/Golfman2.png",
+                            f"{cur_dir}/examples/golf/Golfman3.png",
+                            "<image> <image> <image> What happens to the man after hitting the ball? And why does it happen?",
+                        ],
+                    ],
+                    inputs=[imagebox, imagebox_2, imagebox_3, textbox],
+                    label="Multi-image Example 3",
+                    fn=clear_after_click_example_3_image,
+                    outputs=[
+                        state,
+                        imagebox,
+                        imagebox_2,
+                        imagebox_3,
+                        videobox,
+                        prompt_style_btn,
+                    ],
+                    run_on_click=True,
+                )
+
+                gr.Examples(
+                    examples=[
+                        [
+                            f"{cur_dir}/examples/icl-logo/google.webp",
+                            f"{cur_dir}/examples/icl-logo/apple.jpg",
+                            f"{cur_dir}/examples/icl-logo/nvidia.png",
+                            "<image> is famous for its search engine. <image> is famous for Mac and iPhone. <image> ",
+                        ],
+                    ],
+                    inputs=[imagebox, imagebox_2, imagebox_3, textbox],
+                    label="In-context Learning Example 1 (Please switch the prompt style to 'no-sys')",
+                    fn=clear_after_click_example_3_image_icl,
+                    outputs=[
+                        state,
+                        imagebox,
+                        imagebox_2,
+                        imagebox_3,
+                        videobox,
+                        prompt_style_btn,
+                    ],
+                    run_on_click=True,
+                )
+
+                gr.Examples(
+                    examples=[
+                        [
+                            f"{cur_dir}/examples/icl-building/csail_building.jpeg",
+                            f"{cur_dir}/examples/icl-building/Toronto_Tower.jpeg",
+                            f"{cur_dir}/examples/icl-building/Golden_State_Bridge.jpeg",
+                            "<image> Boston. <image> Toronto. <image> ",
+                        ],
+                    ],
+                    inputs=[imagebox, imagebox_2, imagebox_3, textbox],
+                    label="In-context Learning Example 2 (Please switch the prompt style to 'no-sys')",
+                    fn=clear_after_click_example_3_image_icl,
+                    outputs=[
+                        state,
+                        imagebox,
+                        imagebox_2,
+                        imagebox_3,
+                        videobox,
+                        prompt_style_btn,
+                    ],
+                    run_on_click=True,
+                )
+
+                gr.Examples(
+                    examples=[
+                        [
+                            f"{cur_dir}/examples/arts/sunflowers.jpg",
+                            f"{cur_dir}/examples/arts/the_persistence_of_memory.png",
+                            f"{cur_dir}/examples/arts/impression_sunrise.png",
+                            "<image> Vincent Van Gogh. <image> Salvador Dalí. <image>",
+                        ],
+                    ],
+                    inputs=[imagebox, imagebox_2, imagebox_3, textbox],
+                    label="In-context Learning Example 3 (Please switch the prompt style to 'no-sys')",
+                    fn=clear_after_click_example_3_image_icl,
+                    outputs=[
+                        state,
+                        imagebox,
+                        imagebox_2,
+                        imagebox_3,
+                        videobox,
+                        prompt_style_btn,
+                    ],
+                    run_on_click=True,
+                )
+
+                with gr.Accordion("Parameters", open=False) as parameter_row:
+                    temperature = gr.Slider(
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=0.2,
+                        step=0.1,
+                        interactive=True,
+                        label="Temperature",
+                    )
+                    top_p = gr.Slider(
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=1.0,
+                        step=0.1,
+                        interactive=True,
+                        label="Top P",
+                    )
+                    max_output_tokens = gr.Slider(
+                        minimum=0,
+                        maximum=1024,
+                        value=512,
+                        step=64,
+                        interactive=True,
+                        label="Max output tokens",
+                    )
+
+                # with gr.Row(elem_id="buttons") as button_row:
+                # upvote_btn = gr.Button(value="👍  Upvote", interactive=False)
+                # downvote_btn = gr.Button(value="👎  Downvote", interactive=False)
+                # flag_btn = gr.Button(value="⚠️  Flag", interactive=False)
+                # stop_btn = gr.Button(value="⏹️  Stop Generation", interactive=False)
+                # regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
+
+        if not embed_mode:
+            gr.Markdown(tos_markdown)
+            gr.Markdown(learn_more_markdown)
+            gr.Markdown(ack_markdown)
+        url_params = gr.JSON(visible=False)
+
+        # Register listeners
+        # btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn]
+        btn_list = [regenerate_btn, clear_btn]
+
+        # upvote_btn.click(
+        #     upvote_last_response,
+        #     [state, model_selector],
+        #     [textbox, upvote_btn, downvote_btn, flag_btn],
+        #     queue=False,
+        # )
+        # downvote_btn.click(
+        #     downvote_last_response,
+        #     [state, model_selector],
+        #     [textbox, upvote_btn, downvote_btn, flag_btn],
+        #     queue=False,
+        # )
+        # flag_btn.click(
+        #     flag_last_response,
+        #     [state, model_selector],
+        #     [textbox, upvote_btn, downvote_btn, flag_btn],
+        #     queue=False,
+        # )
+
+        regenerate_btn.click(
+            regenerate,
+            [state, image_process_mode],
+            [state, chatbot, textbox] + btn_list,
+            queue=False,
+        ).then(
+            http_bot,
+            [
+                state,
+                model_selector,
+                temperature,
+                top_p,
+                max_output_tokens,
+                prompt_style_btn,
+            ],
+            [state, chatbot] + btn_list,
+        )
+
+        prompt_style_btn.change(
+            change_prompt_style, [state, prompt_style_btn], [state], queue=False
+        )
+
+        clear_btn.click(
+            clear_history,
+            [prompt_style_btn],
+            [state, chatbot, textbox, imagebox, imagebox_2, imagebox_3, videobox]
+            + btn_list,
+            queue=False,
+        )
+
+        # textbox.submit(
+        #     add_text,
+        #     [state, textbox, imagebox, image_process_mode],
+        #     [state, chatbot, textbox, imagebox] + btn_list,
+        #     queue=False
+        # ).then(
+        #     http_bot,
+        #     [state, model_selector, temperature, top_p, max_output_tokens],
+        #     [state, chatbot] + btn_list
+        # )
+
+        # im_submit_btn.click(
+        #     mirror,
+        #     inputs=[imagebox],
+        #     outputs=[imagebox_out]
+        # ).then(
+        #     add_image,
+        #     [state, imagebox, image_process_mode],
+        #     [state, imagebox] + btn_list,
+        #     queue=False
+        # )
+
+        textbox.submit(
+            clear_text_history, [state, prompt_style_btn], [state, chatbot], queue=False
+        ).then(
+            add_images,
+            [state, imagebox, imagebox_2, imagebox_3, videobox, image_process_mode],
+            [state],
+            queue=False,
+        ).then(
+            add_text_only, [state, textbox], [state, textbox] + btn_list, queue=False
+        ).then(
+            http_bot,
+            [
+                state,
+                model_selector,
+                temperature,
+                top_p,
+                max_output_tokens,
+                prompt_style_btn,
+            ],
+            [state, chatbot] + btn_list,
+        )
+
+        submit_btn.click(
+            clear_text_history, [state, prompt_style_btn], [state, chatbot], queue=False
+        ).then(
+            add_images,
+            [state, imagebox, imagebox_2, imagebox_3, videobox, image_process_mode],
+            [state],
+            queue=False,
+        ).then(
+            add_text_only, [state, textbox], [state, textbox] + btn_list, queue=False
+        ).then(
+            http_bot,
+            [
+                state,
+                model_selector,
+                temperature,
+                top_p,
+                max_output_tokens,
+                prompt_style_btn,
+            ],
+            [state, chatbot] + btn_list,
+        )
+
+        if args.model_list_mode == "once":
+            demo.load(
+                load_demo,
+                [url_params, prompt_style_btn],
+                [state, model_selector],
+                _js=get_window_url_params,
+                queue=False,
+            )
+        elif args.model_list_mode == "reload":
+            demo.load(
+                load_demo_refresh_model_list,
+                [prompt_style_btn],
+                [state, model_selector],
+                queue=False,
+            )
+        else:
+            raise ValueError(f"Unknown model list mode: {args.model_list_mode}")
+
+    return demo
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser.add_argument("--port", type=int)
+    parser.add_argument("--controller-url", type=str, default="http://localhost:21001")
+    parser.add_argument("--concurrency-count", type=int, default=10)
+    parser.add_argument(
+        "--model-list-mode", type=str, default="once", choices=["once", "reload"]
+    )
+    parser.add_argument("--share", action="store_true")
+    parser.add_argument("--moderate", action="store_true")
+    parser.add_argument("--embed", action="store_true")
+    parser.add_argument(
+        "--auto-pad-image-token",
+        action="store_true",
+        help="Automatically pad <image> token to the before of the prompt if no user inputs.",
+    )
+    # NOTE: For single image input, we still auto pad <image> token even if the --auto-pad-image-token is False
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    models = get_model_list()
+
+    logger.info(args)
+    demo = build_demo(args.embed)
+    demo.queue(concurrency_count=args.concurrency_count, api_open=False).launch(
+        server_name=args.host, server_port=args.port, share=args.share
+    )
diff --git a/llm-awq/tinychat/serve/llava_conv.py b/llm-awq/tinychat/serve/llava_conv.py
new file mode 100644
index 0000000000000000000000000000000000000000..bd5e6c1d238c0d01e572e1e0e82bebe905a6acf8
--- /dev/null
+++ b/llm-awq/tinychat/serve/llava_conv.py
@@ -0,0 +1,454 @@
+#    Modified from https://github.com/haotian-liu/LLaVA
+#    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 dataclasses
+from enum import auto, Enum
+from typing import List, Tuple
+
+from tinychat.utils.constants import AUTO_FILL_IM_TOKEN_HOLDER
+
+
+class SeparatorStyle(Enum):
+    """Different separator style."""
+
+    SINGLE = auto()
+    TWO = auto()
+    MPT = auto()
+    PLAIN = auto()
+    LLAMA_2 = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    offset: int
+    sep_style: SeparatorStyle = SeparatorStyle.SINGLE
+    sep: str = "###"
+    sep2: str = None
+    version: str = "Unknown"
+
+    skip_next: bool = False
+    image_loaded: bool = False
+
+    def get_prompt(self):
+        messages = self.messages
+        if len(messages) > 0 and type(messages[0][1]) is tuple:
+            messages = self.messages.copy()
+            init_role, init_msg = messages[0].copy()
+            # init_msg = init_msg[0].replace(AUTO_FILL_IM_TOKEN_HOLDER, "").strip()
+            init_msg = init_msg[0].replace("", "").strip()
+            if "mmtag" in self.version:
+                messages[0] = (init_role, init_msg)
+                messages.insert(0, (self.roles[0], "<Image><image></Image>"))
+                messages.insert(1, (self.roles[1], "Received."))
+            else:
+                # messages[0] = (init_role, AUTO_FILL_IM_TOKEN_HOLDER + init_msg)
+                messages[0] = (init_role, "" + init_msg)
+
+        if self.sep_style == SeparatorStyle.SINGLE:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if role:
+                    if message:
+                        if type(message) is tuple:
+                            message, _, _ = message
+                        ret += role + ": " + message + self.sep
+                    else:
+                        ret += role + ":"
+                else:
+                    if message:
+                        if type(message) is tuple:
+                            message, _, _ = message
+                        ret += message
+        elif self.sep_style == SeparatorStyle.TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(messages):
+                if role:
+                    if message:
+                        if type(message) is tuple:
+                            message, _, _ = message
+                        ret += role + ": " + message + seps[i % 2]
+                    else:
+                        ret += role + ":"
+                else:
+                    if message:
+                        if type(message) is tuple:
+                            message, _, _ = message
+                        ret += message + seps[i % 2]
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+        elif self.sep_style == SeparatorStyle.LLAMA_2:
+            wrap_sys = lambda msg: f"<<SYS>>\n{msg}\n<</SYS>>\n\n"
+            wrap_inst = lambda msg: f"[INST] {msg} [/INST]"
+            ret = ""
+
+            for i, (role, message) in enumerate(messages):
+                if i == 0:
+                    assert message, "first message should not be none"
+                    assert role == self.roles[0], "first message should come from user"
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    if i == 0:
+                        message = wrap_sys(self.system) + message
+                    if i % 2 == 0:
+                        message = wrap_inst(message)
+                        ret += self.sep + message
+                    else:
+                        ret += " " + message + " " + self.sep2
+                else:
+                    ret += ""
+            ret = ret.lstrip(self.sep)
+        elif self.sep_style == SeparatorStyle.PLAIN:
+            seps = [self.sep, self.sep2]
+            ret = self.system
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += message + seps[i % 2]
+                else:
+                    ret += ""
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+        return ret
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def get_images(self, return_pil=False):
+        images = []
+        for i, (role, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    import base64
+                    from io import BytesIO
+                    from PIL import Image
+
+                    msg, image, image_process_mode = msg
+                    if image_process_mode == "Pad":
+
+                        def expand2square(pil_img, background_color=(122, 116, 104)):
+                            width, height = pil_img.size
+                            if width == height:
+                                return pil_img
+                            elif width > height:
+                                result = Image.new(
+                                    pil_img.mode, (width, width), background_color
+                                )
+                                result.paste(pil_img, (0, (width - height) // 2))
+                                return result
+                            else:
+                                result = Image.new(
+                                    pil_img.mode, (height, height), background_color
+                                )
+                                result.paste(pil_img, ((height - width) // 2, 0))
+                                return result
+
+                        image = expand2square(image)
+                    elif image_process_mode in ["Default", "Crop"]:
+                        pass
+                    elif image_process_mode == "Resize":
+                        image = image.resize((336, 336))
+                    else:
+                        raise ValueError(
+                            f"Invalid image_process_mode: {image_process_mode}"
+                        )
+                    max_hw, min_hw = max(image.size), min(image.size)
+                    aspect_ratio = max_hw / min_hw
+                    max_len, min_len = 800, 400
+                    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+                    longest_edge = int(shortest_edge * aspect_ratio)
+                    W, H = image.size
+                    if longest_edge != max(image.size):
+                        if H > W:
+                            H, W = longest_edge, shortest_edge
+                        else:
+                            H, W = shortest_edge, longest_edge
+                        image = image.resize((W, H))
+                    if return_pil:
+                        images.append(image)
+                    else:
+                        buffered = BytesIO()
+                        image.save(buffered, format="PNG")
+                        img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+                        images.append(img_b64_str)
+        return images
+
+    def to_gradio_chatbot(self):
+        ret = []
+        # count the figures to skip visualizing them in the text box.
+        cur_num_fig = 0
+        for i, (role, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    # Skip the visualization of image in the chatbox
+                    cur_num_fig += 1
+                    continue
+                    # import base64
+                    # from io import BytesIO
+                    # msg, image, image_process_mode = msg
+                    # max_hw, min_hw = max(image.size), min(image.size)
+                    # aspect_ratio = max_hw / min_hw
+                    # max_len, min_len = 800, 400
+                    # shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+                    # longest_edge = int(shortest_edge * aspect_ratio)
+                    # W, H = image.size
+                    # if H > W:
+                    #     H, W = longest_edge, shortest_edge
+                    # else:
+                    #     H, W = shortest_edge, longest_edge
+                    # image = image.resize((W, H))
+                    # buffered = BytesIO()
+                    # image.save(buffered, format="JPEG")
+                    # img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+                    # img_str = f'<img src="data:image/png;base64,{img_b64_str}" alt="user upload image" />'
+                    # msg = img_str + msg.replace('<image>', '').strip()
+                    # ret.append([msg[0], None])
+                else:
+                    ret.append([msg, None])
+            else:
+                if cur_num_fig > 0:
+                    cur_num_fig -= 1
+                    continue
+                ret[-1][-1] = msg
+        return ret
+
+    def copy(self):
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            version=self.version,
+            image_loaded=self.image_loaded,
+        )
+
+    def dict(self):
+        if len(self.get_images()) > 0:
+            return {
+                "system": self.system,
+                "roles": self.roles,
+                "messages": [
+                    [x, y[0] if type(y) is tuple else y] for x, y in self.messages
+                ],
+                "offset": self.offset,
+                "sep": self.sep,
+                "sep2": self.sep2,
+            }
+        return {
+            "system": self.system,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+            "sep": self.sep,
+            "sep2": self.sep2,
+        }
+
+
+conv_vicuna_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+        (
+            "Human",
+            "What are the key differences between renewable and non-renewable energy sources?",
+        ),
+        (
+            "Assistant",
+            "Renewable energy sources are those that can be replenished naturally in a relatively "
+            "short amount of time, such as solar, wind, hydro, geothermal, and biomass. "
+            "Non-renewable energy sources, on the other hand, are finite and will eventually be "
+            "depleted, such as coal, oil, and natural gas. Here are some key differences between "
+            "renewable and non-renewable energy sources:\n"
+            "1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable "
+            "energy sources are finite and will eventually run out.\n"
+            "2. Environmental impact: Renewable energy sources have a much lower environmental impact "
+            "than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, "
+            "and other negative effects.\n"
+            "3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically "
+            "have lower operational costs than non-renewable sources.\n"
+            "4. Reliability: Renewable energy sources are often more reliable and can be used in more remote "
+            "locations than non-renewable sources.\n"
+            "5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different "
+            "situations and needs, while non-renewable sources are more rigid and inflexible.\n"
+            "6. Sustainability: Renewable energy sources are more sustainable over the long term, while "
+            "non-renewable sources are not, and their depletion can lead to economic and social instability.\n",
+        ),
+    ),
+    offset=2,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+empty_conv = Conversation(
+    system="",
+    roles=("", ""),
+    version="no-sys",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep="",
+    sep2="</s>",
+)
+
+conv_vicuna_v1 = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="default",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llama_2 = Conversation(
+    system="""You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
+
+If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.""",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_llava_llama_2 = Conversation(
+    system="You are a helpful language and vision assistant. "
+    "You are able to understand the visual content that the user provides, "
+    "and assist the user with a variety of tasks using natural language.",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_mpt = Conversation(
+    system="""<|im_start|>system
+A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.""",
+    roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
+    version="mpt",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MPT,
+    sep="<|im_end|>",
+)
+
+conv_llava_plain = Conversation(
+    system="",
+    roles=("", ""),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.PLAIN,
+    sep="\n",
+)
+
+conv_llava_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_llava_v0_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+    "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("Human", "Assistant"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+    version="v0_mmtag",
+)
+
+conv_llava_v1 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llava_v1_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+    "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("USER", "ASSISTANT"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+    version="v1_mmtag",
+)
+
+default_conversation = conv_vicuna_v1
+conv_templates = {
+    "no-sys": empty_conv,
+    "default": conv_vicuna_v1,
+    "v0": conv_vicuna_v0,
+    "v1": conv_vicuna_v1,
+    "vicuna_v1": conv_vicuna_v1,
+    "llama_2": conv_llama_2,
+    "plain": conv_llava_plain,
+    "v0_plain": conv_llava_plain,
+    "llava_v0": conv_llava_v0,
+    "v0_mmtag": conv_llava_v0_mmtag,
+    "llava_v1": conv_llava_v1,
+    "v1_mmtag": conv_llava_v1_mmtag,
+    "llava_llama_2": conv_llava_llama_2,
+    "mpt": conv_mpt,
+}
+
+
+def get_conversation(version: str):
+    conv = conv_templates.get(version, conv_vicuna_v1).copy()
+    return conv
+
+
+if __name__ == "__main__":
+    print(default_conversation.get_prompt())
diff --git a/llm-awq/tinychat/serve/model_worker.py b/llm-awq/tinychat/serve/model_worker.py
new file mode 100644
index 0000000000000000000000000000000000000000..cca412a3645b21e3c3e136f320884ddf67447a18
--- /dev/null
+++ b/llm-awq/tinychat/serve/model_worker.py
@@ -0,0 +1,433 @@
+#    Modified from https://github.com/haotian-liu/LLaVA
+#    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.
+
+"""
+A model worker executes the model.
+"""
+import argparse
+import asyncio
+import json
+import time
+import threading
+import uuid
+
+from fastapi import FastAPI, Request, BackgroundTasks
+from fastapi.responses import StreamingResponse
+import requests
+import torch
+import uvicorn
+from functools import partial
+from tqdm import tqdm
+
+import tinychat.utils.constants
+from tinychat.utils.constants import (
+    WORKER_HEART_BEAT_INTERVAL,
+    LLAVA_DEFAULT_IMAGE_TOKEN_IDX,
+    LLAVA_DEFAULT_IMAGE_TOKEN,
+    LLAVA_DEFAULT_IM_START_TOKEN,
+    LLAVA_DEFAULT_IM_END_TOKEN,
+)
+from tinychat.utils.log_utils import (
+    build_logger,
+    server_error_msg,
+    pretty_print_semaphore,
+)
+from tinychat.stream_generators.llava_stream_gen import tokenizer_image_token
+from tinychat.utils.llava_image_processing import process_images, load_image_from_base64
+from tinychat.models.llava_llama import LlavaLlamaForCausalLM
+from tinychat.stream_generators.llava_stream_gen import LlavaStreamGenerator
+from tinychat.utils.prompt_templates import (
+    get_prompter,
+    get_stop_token_ids,
+    get_image_token,
+)
+from tinychat.utils.conversation_utils import gen_params
+
+from transformers import AutoConfig, AutoTokenizer
+from accelerate import load_checkpoint_and_dispatch
+
+# import os
+# os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+GB = 1 << 30
+
+worker_id = str(uuid.uuid4())[:6]
+logger = build_logger("model_worker", f"model_worker_{worker_id}.log")
+global_counter = 0
+
+model_semaphore = None
+
+
+def heart_beat_worker(controller):
+    while True:
+        time.sleep(WORKER_HEART_BEAT_INTERVAL)
+        controller.send_heart_beat()
+
+
+def skip(*args, **kwargs):
+    pass
+
+
+class ModelWorker:
+    def __init__(
+        self,
+        controller_addr,
+        worker_addr,
+        worker_id,
+        no_register,
+        model_type,
+        model_path,
+        model_name,
+        quant_path,
+        precision,
+        device,
+    ):
+        self.controller_addr = controller_addr
+        self.worker_addr = worker_addr
+        self.worker_id = worker_id
+        self.model_type = model_type
+        self.model_path = model_path
+        if model_path.endswith("/"):
+            model_path = model_path[:-1]
+        if model_name is None:
+            model_paths = model_path.split("/")
+            if model_paths[-1].startswith("checkpoint-"):
+                self.model_name = model_paths[-2] + "_" + model_paths[-1]
+            else:
+                self.model_name = model_paths[-1]
+        else:
+            self.model_name = model_name
+        if precision == "W4A16":
+            self.model_name = self.model_name + "-4bit-AWQ"
+        self.device = device
+
+        # Load TinyChat model
+        logger.info(f"Loading the model {self.model_name} on worker {worker_id} ...")
+
+        setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+        setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+        torch.nn.init.kaiming_uniform_ = skip
+        torch.nn.init.kaiming_normal_ = skip
+        torch.nn.init.uniform_ = skip
+        torch.nn.init.normal_ = skip
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+        tinychat.utils.constants.LLAVA_DEFAULT_IMAGE_PATCH_TOKEN_IDX = (
+            self.tokenizer.convert_tokens_to_ids(
+                [tinychat.utils.constants.LLAVA_DEFAULT_IMAGE_PATCH_TOKEN]
+            )[0]
+        )
+        config = AutoConfig.from_pretrained(args.model_path, trust_remote_code=True)
+        config.min_max_range_path = args.model_path + "/emb_min_max.pt"
+        model = LlavaLlamaForCausalLM(config, args.device).half()
+        vision_tower = model.get_model().vision_tower
+        if not vision_tower.is_loaded:
+            vision_tower.load_model()
+        vision_tower = vision_tower.half()
+        self.image_processor = vision_tower.image_processor
+
+        if precision == "W16A16":
+            pbar = tqdm(range(1))
+            pbar.set_description("Loading checkpoint shards")
+            for i in pbar:
+                model = load_checkpoint_and_dispatch(
+                    model,
+                    model_path,
+                    no_split_module_classes=[
+                        "OPTDecoderLayer",
+                        "LlamaDecoderLayer",
+                        "BloomBlock",
+                        "MPTBlock",
+                        "DecoderLayer",
+                        "CLIPEncoderLayer",
+                    ],
+                ).to(device)
+        elif precision == "W4A16":
+            from tinychat.utils.load_quant import load_awq_model
+
+            model = load_awq_model(model, quant_path, 4, 128, device)
+            from tinychat.modules import (
+                make_quant_norm,
+                make_quant_attn,
+            )
+
+            make_quant_attn(model, device)
+            make_quant_norm(model)
+            model = model.to(device)
+        else:
+            raise NotImplementedError(f"Precision {precision} is not supported.")
+
+        self.model = model
+        self.is_multimodal = (
+            "llava" in self.model_name.lower() or "vila" in self.model_name.lower()
+        )
+
+        if not no_register:
+            self.register_to_controller()
+            self.heart_beat_thread = threading.Thread(
+                target=heart_beat_worker, args=(self,)
+            )
+            self.heart_beat_thread.start()
+
+    def register_to_controller(self):
+        logger.info("Register to controller")
+
+        url = self.controller_addr + "/register_worker"
+        data = {
+            "worker_name": self.worker_addr,
+            "check_heart_beat": True,
+            "worker_status": self.get_status(),
+        }
+        r = requests.post(url, json=data)
+        assert r.status_code == 200
+
+    def send_heart_beat(self):
+        logger.info(
+            f"Send heart beat. Models: {[self.model_name]}. "
+            f"Semaphore: {pretty_print_semaphore(model_semaphore)}. "
+            f"global_counter: {global_counter}"
+        )
+
+        url = self.controller_addr + "/receive_heart_beat"
+
+        while True:
+            try:
+                ret = requests.post(
+                    url,
+                    json={
+                        "worker_name": self.worker_addr,
+                        "queue_length": self.get_queue_length(),
+                    },
+                    timeout=5,
+                )
+                exist = ret.json()["exist"]
+                break
+            except requests.exceptions.RequestException as e:
+                logger.error(f"heart beat error: {e}")
+            time.sleep(5)
+
+        if not exist:
+            self.register_to_controller()
+
+    def get_queue_length(self):
+        if model_semaphore is None:
+            return 0
+        else:
+            return (
+                args.limit_model_concurrency
+                - model_semaphore._value
+                + (
+                    len(model_semaphore._waiters)
+                    if model_semaphore._waiters is not None
+                    else 0
+                )
+            )
+
+    def get_status(self):
+        return {
+            "model_names": [self.model_name],
+            "speed": 1,
+            "queue_length": self.get_queue_length(),
+        }
+
+    @torch.inference_mode()
+    def generate_stream(self, params):
+        tokenizer, model, image_processor = (
+            self.tokenizer,
+            self.model,
+            self.image_processor,
+        )
+
+        prompt = params["prompt"]
+        ori_prompt = prompt
+        images = params.get("images", None)
+        if images is not None and len(images) > 0 and self.is_multimodal:
+            if len(images) > 0:
+                if len(images) != prompt.count(LLAVA_DEFAULT_IMAGE_TOKEN):
+                    raise ValueError(
+                        "Number of images does not match number of <image> tokens in prompt"
+                    )
+
+                images = [load_image_from_base64(image) for image in images]
+                images = process_images(images, image_processor, model.config)
+
+                if type(images) is list:
+                    images = [
+                        image.to(model.device, dtype=torch.float16) for image in images
+                    ]
+                else:
+                    images = images.to(model.device, dtype=torch.float16)
+
+                replace_token = LLAVA_DEFAULT_IMAGE_TOKEN
+                if getattr(model.config, "mm_use_im_start_end", False):
+                    replace_token = (
+                        LLAVA_DEFAULT_IM_START_TOKEN
+                        + replace_token
+                        + LLAVA_DEFAULT_IM_END_TOKEN
+                    )
+                prompt = prompt.replace(LLAVA_DEFAULT_IMAGE_TOKEN, replace_token)
+            else:
+                images = None
+        else:
+            images = None
+
+        gen_params.temp = float(params.get("temperature", 1.0))
+        gen_params.top_p = float(params.get("top_p", 1.0))
+        gen_params.n_predict = min(int(params.get("max_new_tokens", 256)), 1024)
+
+        stream_generator = LlavaStreamGenerator
+        stop_token_ids = get_stop_token_ids(self.model_type, self.model_path)
+        image_token = get_image_token(model, self.model_path)
+        image_token_holder = (
+            tinychat.utils.constants.LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER
+        )
+        prompt = prompt.replace(image_token_holder, image_token)
+
+        # print("=" * 50)
+        # print(prompt)
+        # print('=' * 50)
+        output_stream = stream_generator(
+            model,
+            tokenizer,
+            prompt,
+            gen_params,
+            device=model.device,
+            stop_token_ids=stop_token_ids,
+            image_tensor=images,
+        )
+
+        generated_text = ori_prompt
+        pre = 0
+        for outputs in output_stream:
+            output_text = outputs["text"]
+            output_text = output_text.strip().split(" ")
+            now = len(output_text) - 1
+            if now > pre:
+                generated_text += " ".join(output_text[pre:now]) + " "
+                yield json.dumps(
+                    {"text": generated_text, "error_code": 0}
+                ).encode() + b"\0"
+                pre = now
+        generated_text += " ".join(output_text[pre:])
+        yield json.dumps({"text": generated_text, "error_code": 0}).encode() + b"\0"
+
+    def generate_stream_gate(self, params):
+        try:
+            for x in self.generate_stream(params):
+                yield x
+        except ValueError as e:
+            print("Caught ValueError:", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+        except torch.cuda.CudaError as e:
+            print("Caught torch.cuda.CudaError:", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+        except Exception as e:
+            print("Caught Unknown Error", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+
+
+app = FastAPI()
+
+
+def release_model_semaphore(fn=None):
+    model_semaphore.release()
+    if fn is not None:
+        fn()
+
+
+@app.post("/worker_generate_stream")
+async def generate_stream(request: Request):
+    global model_semaphore, global_counter
+    global_counter += 1
+    params = await request.json()
+
+    if model_semaphore is None:
+        model_semaphore = asyncio.Semaphore(args.limit_model_concurrency)
+    await model_semaphore.acquire()
+    worker.send_heart_beat()
+    generator = worker.generate_stream_gate(params)
+    background_tasks = BackgroundTasks()
+    background_tasks.add_task(
+        partial(release_model_semaphore, fn=worker.send_heart_beat)
+    )
+    return StreamingResponse(generator, background=background_tasks)
+
+
+@app.post("/worker_get_status")
+async def get_status(request: Request):
+    return worker.get_status()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=21002)
+    parser.add_argument("--worker-address", type=str, default="http://localhost:21002")
+    parser.add_argument(
+        "--controller-address", type=str, default="http://localhost:21001"
+    )
+    parser.add_argument(
+        "--model-type",
+        type=str,
+        default="LLaMa",
+        help="type of the (base) language model",
+    )
+    parser.add_argument("--model-path", type=str, default="facebook/opt-350m")
+    parser.add_argument("--model-name", type=str)
+    parser.add_argument("--quant-path", type=str, default=None)
+    parser.add_argument("--precision", type=str, default="W4A16")
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument(
+        "--multi-modal",
+        action="store_true",
+        help="Multimodal mode is automatically detected with model name, please make sure `llava` is included in the model path.",
+    )
+    parser.add_argument("--limit-model-concurrency", type=int, default=5)
+    parser.add_argument("--stream-interval", type=int, default=1)
+    parser.add_argument("--no-register", action="store_true")
+
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    if args.multi_modal:
+        logger.warning(
+            "Multimodal mode is automatically detected with model name, please make sure `llava` is included in the model path."
+        )
+
+    worker = ModelWorker(
+        args.controller_address,
+        args.worker_address,
+        worker_id,
+        args.no_register,
+        args.model_type,
+        args.model_path,
+        args.model_name,
+        args.quant_path,
+        args.precision,
+        args.device,
+    )
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")
diff --git a/llm-awq/tinychat/serve/model_worker_new.py b/llm-awq/tinychat/serve/model_worker_new.py
new file mode 100644
index 0000000000000000000000000000000000000000..debbe2ebf44eb5bdc758e80e26278ceb553e9d83
--- /dev/null
+++ b/llm-awq/tinychat/serve/model_worker_new.py
@@ -0,0 +1,445 @@
+#    Modified from https://github.com/haotian-liu/LLaVA
+#    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.
+
+"""
+A model worker executes the model.
+"""
+import os
+import argparse
+import asyncio
+import json
+import time
+import threading
+import uuid
+
+from fastapi import FastAPI, Request, BackgroundTasks
+from fastapi.responses import StreamingResponse
+import requests
+import torch
+import uvicorn
+from functools import partial
+from tqdm import tqdm
+
+import tinychat.utils.constants
+from tinychat.utils.constants import (
+    WORKER_HEART_BEAT_INTERVAL,
+    LLAVA_DEFAULT_IMAGE_TOKEN_IDX,
+    LLAVA_DEFAULT_IMAGE_TOKEN,
+    LLAVA_DEFAULT_IM_START_TOKEN,
+    LLAVA_DEFAULT_IM_END_TOKEN,
+)
+from tinychat.utils.log_utils import (
+    build_logger,
+    server_error_msg,
+    pretty_print_semaphore,
+)
+from tinychat.stream_generators.llava_stream_gen import tokenizer_image_token
+from tinychat.utils.llava_image_processing import process_images, load_image_from_base64
+
+# from tinychat.models.llava_llama import LlavaLlamaForCausalLM
+from tinychat.models.vila_llama import VilaLlamaForCausalLM
+from tinychat.stream_generators.llava_stream_gen import LlavaStreamGenerator
+from tinychat.utils.prompt_templates import (
+    get_prompter,
+    get_stop_token_ids,
+    get_image_token,
+)
+from tinychat.utils.conversation_utils import gen_params
+
+from transformers import AutoConfig, AutoTokenizer
+from accelerate import load_checkpoint_and_dispatch
+
+# import os
+# os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+GB = 1 << 30
+
+worker_id = str(uuid.uuid4())[:6]
+logger = build_logger("model_worker", f"model_worker_{worker_id}.log")
+global_counter = 0
+
+model_semaphore = None
+
+
+def heart_beat_worker(controller):
+    while True:
+        time.sleep(WORKER_HEART_BEAT_INTERVAL)
+        controller.send_heart_beat()
+
+
+def skip(*args, **kwargs):
+    pass
+
+
+class ModelWorker:
+    def __init__(
+        self,
+        controller_addr,
+        worker_addr,
+        worker_id,
+        no_register,
+        model_type,
+        model_path,
+        model_name,
+        quant_path,
+        precision,
+        device,
+    ):
+        self.controller_addr = controller_addr
+        self.worker_addr = worker_addr
+        self.worker_id = worker_id
+        self.model_type = model_type
+        self.model_path = model_path
+        if model_path.endswith("/"):
+            model_path = model_path[:-1]
+        if model_name is None:
+            model_paths = model_path.split("/")
+            if model_paths[-1].startswith("checkpoint-"):
+                self.model_name = model_paths[-2] + "_" + model_paths[-1]
+            else:
+                self.model_name = model_paths[-1]
+        else:
+            self.model_name = model_name
+        if precision == "W4A16":
+            self.model_name = self.model_name + "-4bit-AWQ"
+        self.device = device
+
+        # Load TinyChat model
+        logger.info(f"Loading the model {self.model_name} on worker {worker_id} ...")
+
+        setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+        setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+        torch.nn.init.kaiming_uniform_ = skip
+        torch.nn.init.kaiming_normal_ = skip
+        torch.nn.init.uniform_ = skip
+        torch.nn.init.normal_ = skip
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            os.path.join(args.model_path, "llm"), use_fast=False
+        )
+        tinychat.utils.constants.LLAVA_DEFAULT_IMAGE_PATCH_TOKEN_IDX = (
+            self.tokenizer.convert_tokens_to_ids(
+                [tinychat.utils.constants.LLAVA_DEFAULT_IMAGE_PATCH_TOKEN]
+            )[0]
+        )
+        config = AutoConfig.from_pretrained(args.model_path, trust_remote_code=True)
+        model = VilaLlamaForCausalLM(config).half()
+        tinychat.utils.constants.LLAVA_DEFAULT_IMAGE_PATCH_TOKEN_IDX = (
+            self.tokenizer.convert_tokens_to_ids(
+                [tinychat.utils.constants.LLAVA_DEFAULT_IMAGE_PATCH_TOKEN]
+            )[0]
+        )
+        vision_tower = model.get_vision_tower()
+        # if not vision_tower.is_loaded:
+        #     vision_tower.load_model()
+        vision_tower = vision_tower.half()
+        self.image_processor = vision_tower.image_processor
+
+        if precision == "W16A16":
+            pbar = tqdm(range(1))
+            pbar.set_description("Loading checkpoint shards")
+            for i in pbar:
+                model.llm = load_checkpoint_and_dispatch(
+                    model.llm,
+                    os.path.join(args.model_path, "llm"),
+                    no_split_module_classes=[
+                        "OPTDecoderLayer",
+                        "LlamaDecoderLayer",
+                        "BloomBlock",
+                        "MPTBlock",
+                        "DecoderLayer",
+                        "CLIPEncoderLayer",
+                    ],
+                ).to(device)
+            model = model.to(args.device)
+            model.eval()
+
+        elif precision == "W4A16":
+            from tinychat.utils.load_quant import load_awq_model
+
+            model.llm = load_awq_model(model.llm, quant_path, 4, 128, device)
+            from tinychat.modules import (
+                make_quant_norm,
+                make_quant_attn,
+            )
+
+            make_quant_attn(model.llm, device)
+            make_quant_norm(model.llm)
+            model = model.to(device)
+        else:
+            raise NotImplementedError(f"Precision {precision} is not supported.")
+
+        self.model = model
+        self.is_multimodal = (
+            "llava" in self.model_name.lower() or "vila" in self.model_name.lower()
+        )
+
+        if not no_register:
+            self.register_to_controller()
+            self.heart_beat_thread = threading.Thread(
+                target=heart_beat_worker, args=(self,)
+            )
+            self.heart_beat_thread.start()
+
+    def register_to_controller(self):
+        logger.info("Register to controller")
+
+        url = self.controller_addr + "/register_worker"
+        data = {
+            "worker_name": self.worker_addr,
+            "check_heart_beat": True,
+            "worker_status": self.get_status(),
+        }
+        r = requests.post(url, json=data)
+        assert r.status_code == 200
+
+    def send_heart_beat(self):
+        logger.info(
+            f"Send heart beat. Models: {[self.model_name]}. "
+            f"Semaphore: {pretty_print_semaphore(model_semaphore)}. "
+            f"global_counter: {global_counter}"
+        )
+
+        url = self.controller_addr + "/receive_heart_beat"
+
+        while True:
+            try:
+                ret = requests.post(
+                    url,
+                    json={
+                        "worker_name": self.worker_addr,
+                        "queue_length": self.get_queue_length(),
+                    },
+                    timeout=5,
+                )
+                exist = ret.json()["exist"]
+                break
+            except requests.exceptions.RequestException as e:
+                logger.error(f"heart beat error: {e}")
+            time.sleep(5)
+
+        if not exist:
+            self.register_to_controller()
+
+    def get_queue_length(self):
+        if model_semaphore is None:
+            return 0
+        else:
+            return (
+                args.limit_model_concurrency
+                - model_semaphore._value
+                + (
+                    len(model_semaphore._waiters)
+                    if model_semaphore._waiters is not None
+                    else 0
+                )
+            )
+
+    def get_status(self):
+        return {
+            "model_names": [self.model_name],
+            "speed": 1,
+            "queue_length": self.get_queue_length(),
+        }
+
+    @torch.inference_mode()
+    def generate_stream(self, params):
+        tokenizer, model, image_processor = (
+            self.tokenizer,
+            self.model,
+            self.image_processor,
+        )
+
+        prompt = params["prompt"]
+        ori_prompt = prompt
+        images = params.get("images", None)
+        if images is not None and len(images) > 0 and self.is_multimodal:
+            if len(images) > 0:
+                if len(images) != prompt.count(LLAVA_DEFAULT_IMAGE_TOKEN):
+                    raise ValueError(
+                        "Number of images does not match number of <image> tokens in prompt"
+                    )
+
+                images = [load_image_from_base64(image) for image in images]
+                images = process_images(images, image_processor, model.config)
+
+                if type(images) is list:
+                    images = [
+                        image.to(model.device, dtype=torch.float16) for image in images
+                    ]
+                else:
+                    images = images.to(model.device, dtype=torch.float16)
+
+                replace_token = LLAVA_DEFAULT_IMAGE_TOKEN
+                if getattr(model.config, "mm_use_im_start_end", False):
+                    replace_token = (
+                        LLAVA_DEFAULT_IM_START_TOKEN
+                        + replace_token
+                        + LLAVA_DEFAULT_IM_END_TOKEN
+                    )
+                prompt = prompt.replace(LLAVA_DEFAULT_IMAGE_TOKEN, replace_token)
+            else:
+                images = None
+        else:
+            images = None
+
+        gen_params.temp = float(params.get("temperature", 1.0))
+        gen_params.top_p = float(params.get("top_p", 1.0))
+        gen_params.n_predict = min(int(params.get("max_new_tokens", 256)), 1024)
+
+        stream_generator = LlavaStreamGenerator
+        stop_token_ids = get_stop_token_ids(self.model_type, self.model_path)
+        image_token = get_image_token(model, self.model_path)
+        image_token_holder = (
+            tinychat.utils.constants.LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER
+        )
+        prompt = prompt.replace(image_token_holder, image_token)
+
+        # print("=" * 50)
+        # print(prompt)
+        # print('=' * 50)
+        output_stream = stream_generator(
+            model,
+            tokenizer,
+            prompt,
+            gen_params,
+            device=model.device,
+            stop_token_ids=stop_token_ids,
+            image_tensor=images,
+        )
+
+        generated_text = ori_prompt
+        pre = 0
+        for outputs in output_stream:
+            output_text = outputs["text"]
+            output_text = output_text.strip().split(" ")
+            now = len(output_text) - 1
+            if now > pre:
+                generated_text += " ".join(output_text[pre:now]) + " "
+                yield json.dumps(
+                    {"text": generated_text, "error_code": 0}
+                ).encode() + b"\0"
+                pre = now
+        generated_text += " ".join(output_text[pre:])
+        yield json.dumps({"text": generated_text, "error_code": 0}).encode() + b"\0"
+
+    def generate_stream_gate(self, params):
+        try:
+            for x in self.generate_stream(params):
+                yield x
+        except ValueError as e:
+            print("Caught ValueError:", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+        except torch.cuda.CudaError as e:
+            print("Caught torch.cuda.CudaError:", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+        except Exception as e:
+            print("Caught Unknown Error", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+
+
+app = FastAPI()
+
+
+def release_model_semaphore(fn=None):
+    model_semaphore.release()
+    if fn is not None:
+        fn()
+
+
+@app.post("/worker_generate_stream")
+async def generate_stream(request: Request):
+    global model_semaphore, global_counter
+    global_counter += 1
+    params = await request.json()
+
+    if model_semaphore is None:
+        model_semaphore = asyncio.Semaphore(args.limit_model_concurrency)
+    await model_semaphore.acquire()
+    worker.send_heart_beat()
+    generator = worker.generate_stream_gate(params)
+    background_tasks = BackgroundTasks()
+    background_tasks.add_task(
+        partial(release_model_semaphore, fn=worker.send_heart_beat)
+    )
+    return StreamingResponse(generator, background=background_tasks)
+
+
+@app.post("/worker_get_status")
+async def get_status(request: Request):
+    return worker.get_status()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=21002)
+    parser.add_argument("--worker-address", type=str, default="http://localhost:21002")
+    parser.add_argument(
+        "--controller-address", type=str, default="http://localhost:21001"
+    )
+    parser.add_argument(
+        "--model-type",
+        type=str,
+        default="LLaMa",
+        help="type of the (base) language model",
+    )
+    parser.add_argument("--model-path", type=str, default="facebook/opt-350m")
+    parser.add_argument("--model-name", type=str)
+    parser.add_argument("--quant-path", type=str, default=None)
+    parser.add_argument("--precision", type=str, default="W4A16")
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument(
+        "--multi-modal",
+        action="store_true",
+        help="Multimodal mode is automatically detected with model name, please make sure `llava` is included in the model path.",
+    )
+    parser.add_argument("--limit-model-concurrency", type=int, default=5)
+    parser.add_argument("--stream-interval", type=int, default=1)
+    parser.add_argument("--no-register", action="store_true")
+
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    if args.multi_modal:
+        logger.warning(
+            "Multimodal mode is automatically detected with model name, please make sure `llava` is included in the model path."
+        )
+
+    worker = ModelWorker(
+        args.controller_address,
+        args.worker_address,
+        worker_id,
+        args.no_register,
+        args.model_type,
+        args.model_path,
+        args.model_name,
+        args.quant_path,
+        args.precision,
+        args.device,
+    )
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")
diff --git a/llm-awq/tinychat/stream_generators/internvl_stream_gen.py b/llm-awq/tinychat/stream_generators/internvl_stream_gen.py
new file mode 100644
index 0000000000000000000000000000000000000000..6f9f723a4d1025ecde8c3a3e2be2199499d7319a
--- /dev/null
+++ b/llm-awq/tinychat/stream_generators/internvl_stream_gen.py
@@ -0,0 +1,204 @@
+import torch
+import gc
+import time
+from typing import Optional
+
+from .llava_stream_gen import prepare_logits_processor
+
+context_tokens = 0
+context_time = 0.0
+total_tokens = 0
+generation_time_list = []
+
+
+@torch.inference_mode()
+def InternVLStreamGenerator(
+    model,
+    gen_params,
+    input: str,
+    media=None,
+    media_cfg=None,
+    start_pos: int = 0,
+    device: str = "cuda:0",
+    stream_interval: int = 2,
+    echo: bool = False,
+    stop_token_ids=[],
+    image_tensor: Optional[torch.FloatTensor] = None,
+    chunk_prefilling: bool = False,
+    quant_llm: bool = False,
+):
+    if chunk_prefilling and start_pos != 0:
+        input = "<|im_start|>" + input
+        
+    if media is not None and "image" in media:
+        num_patches_list = [image.size(0) for image in media["image"]]
+            
+        IMG_START_TOKEN = '<img>'
+        IMG_END_TOKEN = '</img>'
+        IMG_CONTEXT_TOKEN = '<IMG_CONTEXT>'
+        NUM_IMAGE_TOKEN = 256
+        
+        img_context_token_id = model.tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
+        model.img_context_token_id = img_context_token_id
+        for num_patches in num_patches_list:
+            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * NUM_IMAGE_TOKEN * num_patches + IMG_END_TOKEN
+            input = input.replace('<image>', image_tokens, 1)
+            
+    if media is not None and "video" in media:
+        num_patches_list = [video.size(0) for video in media["video"]]
+        
+        IMG_START_TOKEN = '<img>'
+        IMG_END_TOKEN = '</img>'
+        IMG_CONTEXT_TOKEN = '<IMG_CONTEXT>'
+        NUM_IMAGE_TOKEN = 256
+        
+        img_context_token_id = model.tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
+        model.img_context_token_id = img_context_token_id
+        for num_patches in num_patches_list:
+            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * NUM_IMAGE_TOKEN * num_patches + IMG_END_TOKEN
+            input = input.replace('<image>', image_tokens, 1)
+    
+    input_ids = model.tokenizer(input)["input_ids"]
+    output_ids = list(input_ids)
+    input_echo_len = len(output_ids)
+    len_input = len(input)
+    if gen_params.top_k <= 0:
+        top_k = gen_params.n_vocab
+    else:
+        top_k = gen_params.top_k
+    logits_processor = prepare_logits_processor(
+        gen_params.temp, gen_params.repeat_penalty, gen_params.top_p, top_k
+    )
+    past_key_values = out = None
+    stop_token_ids.append(model.tokenizer.eos_token_id)
+    max_new_tokens = gen_params.n_predict
+
+    for i in range(max_new_tokens):
+        torch.cuda.synchronize()
+        t_st = time.time()
+        if i == 0:
+            inputs = torch.as_tensor([input_ids], device=device)
+        else:
+            inputs = torch.as_tensor([[token]], device=device)
+        out, length = model.stream_gen(
+            input_ids=inputs,
+            media=media,
+            media_cfg=media_cfg,
+            start_pos=start_pos,
+            chunk_prefilling=chunk_prefilling,
+            quant_llm=quant_llm,
+        )
+        start_pos += length
+        logits = out
+        torch.cuda.synchronize()
+        t_ed = time.time()
+        media = None
+        media_cfg = None
+        if torch.sum(torch.isinf(logits)):
+            print(
+                "{a} of {b}".format(
+                    a=torch.sum(torch.isinf(logits)).item(), b=logits.numel()
+                )
+            )
+            print("{},{}".format(torch.max(logits), torch.min(logits)))
+        # Processing the logits
+        if logits_processor:
+            if gen_params.repeat_penalty > 1.0:
+                tmp_output_ids = torch.as_tensor([output_ids], device=logits.device)
+                # tmp_output_ids = output_ids[0].unsqueeze(0)
+            else:
+                tmp_output_ids = None
+            last_token_logits = logits_processor(tmp_output_ids, logits[:, -1, :])[0]
+        else:
+            last_token_logits = logits[:, -1, :]
+        if gen_params.temp < 1e-5 or gen_params.top_p < 1e-8:  # greedy
+            token = int(torch.argmax(last_token_logits))
+        else:
+            probs = torch.softmax(last_token_logits.float(), dim=-1)
+            if torch.any(torch.isinf(probs)) or torch.any(torch.isnan(probs)):
+                print(
+                    "[Error] Invalid probabilities detected (Inf/Nan exists). Saving the tensor and exiting..."
+                )
+                torch.save(last_token_logits, "last_token_logits.pt")
+                exit()
+            token = int(torch.multinomial(probs, num_samples=1))
+        output_ids.append(token)
+
+        global context_time
+        global context_tokens
+        global total_tokens
+        global generation_time_list
+        if i == 0:
+            context_time = t_ed - t_st
+            context_tokens = length
+            generation_time_list = []
+        else:
+            generation_time_list.append(t_ed - t_st)
+
+        if token in stop_token_ids:
+            stopped = True
+        else:
+            stopped = False
+
+        if i % stream_interval == 0 or i == max_new_tokens - 1 or stopped:
+            if echo:
+                tmp_output_ids = output_ids
+                rfind_start = len_input
+            else:
+                tmp_output_ids = output_ids[input_echo_len:]
+                rfind_start = 0
+
+            output = model.tokenizer.decode(
+                tmp_output_ids,
+                skip_special_tokens=True,
+                spaces_between_special_tokens=False,
+            )
+
+            partially_stopped = False
+
+            # prevent yielding partial stop sequence
+            if not partially_stopped:
+                yield {
+                    "text": output,
+                    "usage": {
+                        "prompt_tokens": input_echo_len,
+                        "completion_tokens": i,
+                        "total_tokens": input_echo_len + i,
+                    },
+                    "finish_reason": None,
+                    "timing": None,
+                }
+
+        if stopped:
+            break
+
+    # finish stream event, which contains finish reason
+    if i == max_new_tokens - 1:
+        finish_reason = "length"
+    elif stopped:
+        finish_reason = "stop"
+    else:
+        finish_reason = None
+
+    total_tokens = context_tokens + len(generation_time_list)
+    yield {
+        "text": output,
+        "usage": {
+            "prompt_tokens": input_echo_len,
+            "completion_tokens": i,
+            "total_tokens": input_echo_len + i,
+        },
+        "finish_reason": finish_reason,
+        "timing": {
+            "context_tokens": context_tokens,
+            "context_time": context_time,
+            "total_tokens": total_tokens,
+            "generation_time_list": generation_time_list,
+        },
+    }
+
+    del past_key_values, out
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    # return context_tokens, context_time, total_tokens, generation_time_list
diff --git a/llm-awq/tinychat/stream_generators/stream_gen.py b/llm-awq/tinychat/stream_generators/stream_gen.py
new file mode 100644
index 0000000000000000000000000000000000000000..9fd8e1e9203be1e18e233b2a17de5d2b185f9162
--- /dev/null
+++ b/llm-awq/tinychat/stream_generators/stream_gen.py
@@ -0,0 +1,213 @@
+import torch
+import gc
+import time
+
+from transformers.generation.logits_process import (
+    LogitsProcessorList,
+    RepetitionPenaltyLogitsProcessor,
+    TemperatureLogitsWarper,
+    TopKLogitsWarper,
+    TopPLogitsWarper,
+)
+
+context_tokens = 0
+context_time = 0.0
+total_tokens = 0
+generation_time_list = []
+
+
+def prepare_logits_processor(
+    temperature: float, repetition_penalty: float, top_p: float, top_k: int
+) -> LogitsProcessorList:
+    processor_list = LogitsProcessorList()
+    # TemperatureLogitsWarper doesn't accept 0.0, 1.0 makes it a no-op so we skip two cases.
+    if temperature >= 1e-5 and temperature != 1.0:
+        processor_list.append(TemperatureLogitsWarper(temperature))
+    if repetition_penalty > 1.0:
+        processor_list.append(RepetitionPenaltyLogitsProcessor(repetition_penalty))
+    if 1e-8 <= top_p < 1.0:
+        processor_list.append(TopPLogitsWarper(top_p))
+    if top_k > 0:
+        processor_list.append(TopKLogitsWarper(top_k))
+    return processor_list
+
+
+@torch.inference_mode()
+def StreamGenerator(
+    model,
+    tokenizer,
+    input: str,
+    start_pos: int,
+    gen_params: dict,
+    device: str = "cuda:0",
+    stream_interval: int = 2,
+    echo: bool = False,
+    stop_token_ids=[],
+    chunk_prefilling=False,
+    quant_llm=False,
+):
+    if chunk_prefilling and start_pos != 0:
+        input_ids = tokenizer(input).input_ids[
+            1:
+        ]  # tokenizer will add a <s> at the beginning, so to delete it (important for chunk_prefilling)
+    else:
+        input_ids = tokenizer(input)["input_ids"]
+    input_echo_len = len(input_ids)
+    output_ids = list(input_ids)
+    len_input = len(input)
+
+    if gen_params.top_k <= 0:
+        top_k = gen_params.n_vocab
+    else:
+        top_k = gen_params.top_k
+    logits_processor = prepare_logits_processor(
+        gen_params.temp, gen_params.repeat_penalty, gen_params.top_p, top_k
+    )
+
+    past_key_values = out = None
+    stop_token_ids.append(tokenizer.eos_token_id)
+    max_new_tokens = gen_params.n_predict
+    for i in range(max_new_tokens):
+        torch.cuda.synchronize()
+        t_st = time.time()
+
+        if i == 0:
+            inputs = torch.as_tensor([input_ids], device=device)
+        else:
+            inputs = torch.as_tensor([[token]], device=device)
+
+        if (
+            "llama" not in model.__class__.__name__.lower()
+            and "mpt" not in model.__class__.__name__.lower()
+            and "falcon" not in model.__class__.__name__.lower()
+            and "qwen" not in model.__class__.__name__.lower()
+            and "internvl" not in model.__class__.__name__.lower()
+        ):
+            if i == 0:  # Context Stage
+                # out = model(inputs, use_cache=True)
+                out = model(inputs)
+                logits = out.logits
+                past_key_values = out.past_key_values
+            else:
+                out = model(
+                    input_ids=inputs,
+                    use_cache=True,
+                    past_key_values=past_key_values,
+                )
+                logits = out.logits
+                past_key_values = out.past_key_values
+        else:
+            if (
+                "llama" in model.__class__.__name__.lower()
+                or "qwen" in model.__class__.__name__.lower()
+                or "internvl" in model.__class__.__name__.lower()
+            ) and not quant_llm:
+                out = model(
+                    inputs,
+                    start_pos=start_pos,
+                    chunk_prefilling=chunk_prefilling,
+                    quant=quant_llm,
+                )
+            else:
+                out = model(
+                    inputs, start_pos=start_pos, chunk_prefilling=chunk_prefilling
+                )
+            start_pos += inputs.shape[1]
+            logits = out
+        torch.cuda.synchronize()
+        t_ed = time.time()
+
+        # Processing the logits
+        if logits_processor:
+            if gen_params.repeat_penalty > 1.0:
+                tmp_output_ids = torch.as_tensor([output_ids], device=logits.device)
+            else:
+                tmp_output_ids = None
+            last_token_logits = logits_processor(tmp_output_ids, logits[:, -1, :])[0]
+        else:
+            last_token_logits = logits[0, -1, :]
+        if gen_params.temp < 1e-5 or gen_params.top_p < 1e-8:  # greedy
+            token = int(torch.argmax(last_token_logits))
+        else:
+            probs = torch.softmax(last_token_logits, dim=-1)
+            token = int(torch.multinomial(probs, num_samples=1))
+        output_ids.append(token)
+
+        global context_time
+        global context_tokens
+        global total_tokens
+        global generation_time_list
+        if i == 0:
+            context_time = t_ed - t_st
+            context_tokens = inputs.shape[1]
+            generation_time_list = []
+        else:
+            generation_time_list.append(t_ed - t_st)
+
+        if token in stop_token_ids:
+            stopped = True
+        else:
+            stopped = False
+
+        if i % stream_interval == 0 or i == max_new_tokens - 1 or stopped:
+            if echo:
+                tmp_output_ids = output_ids
+                rfind_start = len_input
+            else:
+                tmp_output_ids = output_ids[input_echo_len:]
+                rfind_start = 0
+
+            output = tokenizer.decode(
+                tmp_output_ids,
+                skip_special_tokens=True,
+                spaces_between_special_tokens=False,
+            )
+
+            partially_stopped = False
+
+            # prevent yielding partial stop sequence
+            if not partially_stopped:
+                yield {
+                    "text": output,
+                    "usage": {
+                        "prompt_tokens": input_echo_len,
+                        "completion_tokens": i,
+                        "total_tokens": input_echo_len + i,
+                    },
+                    "finish_reason": None,
+                    "timing": None,
+                }
+
+        if stopped:
+            break
+
+    # finish stream event, which contains finish reason
+    if i == max_new_tokens - 1:
+        finish_reason = "length"
+    elif stopped:
+        finish_reason = "stop"
+    else:
+        finish_reason = None
+
+    total_tokens = context_tokens + len(generation_time_list)
+    yield {
+        "text": output,
+        "usage": {
+            "prompt_tokens": input_echo_len,
+            "completion_tokens": i,
+            "total_tokens": input_echo_len + i,
+        },
+        "finish_reason": finish_reason,
+        "timing": {
+            "context_tokens": context_tokens,
+            "context_time": context_time,
+            "total_tokens": total_tokens,
+            "generation_time_list": generation_time_list,
+        },
+    }
+
+    del past_key_values, out
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    # return context_tokens, context_time, total_tokens, generation_time_list
diff --git a/llm-awq/tinychat/utils/__init__.py b/llm-awq/tinychat/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..fc4c2513717be2c25b103a248e00dcce64862c59
--- /dev/null
+++ b/llm-awq/tinychat/utils/__init__.py
@@ -0,0 +1,3 @@
+import tinychat.utils.constants as constants
+
+constants.init()
diff --git a/llm-awq/tinychat/utils/conversation_utils.py b/llm-awq/tinychat/utils/conversation_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..1b98c1d8d3076799f3f50ff8f527761b9cc5fc48
--- /dev/null
+++ b/llm-awq/tinychat/utils/conversation_utils.py
@@ -0,0 +1,85 @@
+from typing import Dict
+import numpy as np
+from attributedict.collections import AttributeDict
+
+gen_params = AttributeDict(
+    [
+        ("seed", -1),  # RNG seed
+        ("n_threads", 1),  # TODO: fix this
+        ("n_predict", 512),  # new tokens to predict
+        ("n_parts", -1),  # amount of model parts (-1: determine from model dimensions)
+        ("n_ctx", 512),  # context size
+        ("n_batch", 512),  # batch size for prompt processing (must be >=32 to use BLAS)
+        ("n_keep", 0),  # number of tokens to keep from initial prompt
+        ("n_vocab", 50272),  # vocabulary size
+        # sampling parameters
+        ("logit_bias", dict()),  # logit bias for specific tokens: <int, float>
+        ("top_k", 50),  # <= 0 to use vocab size
+        ("top_p", 0.95),  # 1.0 = disabled
+        ("tfs_z", 1.00),  # 1.0 = disabled
+        ("typical_p", 1.00),  # 1.0 = disabled
+        ("temp", 0.20),  # 1.0 = disabled
+        ("repeat_penalty", 1.10),  # 1.0 = disabled
+        (
+            "repeat_last_n",
+            64,
+        ),  # last n tokens to penalize (0 = disable penalty, -1 = context size)
+        ("frequency_penalty", 0.00),  # 0.0 = disabled
+        ("presence_penalty", 0.00),  # 0.0 = disabled
+        ("mirostat", 0),  # 0 = disabled, 1 = mirostat, 2 = mirostat 2.0
+        ("mirostat_tau", 5.00),  # target entropy
+        ("mirostat_eta", 0.10),  # learning rate
+    ]
+)
+
+
+class TimeStats:
+    def __init__(self):
+        self.total_tokens = 0
+        self.context_tokens = 0
+        self.context_time = 0.0
+
+        self.generation_tokens = 0
+        self.generation_time_list = []
+
+    def update(self, timing: Dict):
+        self.context_tokens = timing["context_tokens"]
+        self.context_time = timing["context_time"]
+        self.total_tokens = timing["total_tokens"]
+        self.generation_time_list = timing["generation_time_list"]
+        self.generation_tokens = len(self.generation_time_list)
+        self.average_speed = (self.context_time + np.sum(self.generation_time_list)) / (
+            self.context_tokens + self.generation_tokens
+        )
+
+    def show(self):
+        if self.total_tokens == 0:
+            # No stats to show.
+            return
+
+        print("*" * 50)
+        print(
+            f"Speed of Generation : {np.average(self.generation_time_list)*1000:.3f} ms/token"
+        )
+        print("*" * 50)
+
+
+def stream_output(output_stream, time_stats: TimeStats = None):
+    pre = 0
+    for outputs in output_stream:
+        output_text = outputs["text"]
+        output_text = output_text.strip().split(" ")
+        now = len(output_text) - 1
+        if now > pre:
+            print(" ".join(output_text[pre:now]), end=" ", flush=True)
+            pre = now
+    print(" ".join(output_text[pre:]), flush=True)
+    if "timing" in outputs and outputs["timing"] is not None:
+        timing = outputs["timing"]
+        total_tokens = timing["total_tokens"]
+        if time_stats is not None:
+            time_stats.update(timing)
+        prompt_tokens = timing["context_tokens"]
+    print("-" * 50)
+    print("TTFT: {:.3f} s for {} tokens.".format(timing["context_time"], prompt_tokens))
+    return " ".join(output_text), total_tokens
diff --git a/llm-awq/tinychat/utils/llava_image_processing.py b/llm-awq/tinychat/utils/llava_image_processing.py
new file mode 100644
index 0000000000000000000000000000000000000000..de0b6d80317f05eb76e466adc8aad667b3e33936
--- /dev/null
+++ b/llm-awq/tinychat/utils/llava_image_processing.py
@@ -0,0 +1,113 @@
+#    Modified from https://github.com/haotian-liu/LLaVA
+#    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 torch
+from PIL import Image
+from io import BytesIO
+import requests
+import os
+import base64
+
+
+def load_image_from_base64(image):
+    return Image.open(BytesIO(base64.b64decode(image)))
+
+
+def load_image(image_file):
+    if image_file.startswith("http://") or image_file.startswith("https://"):
+        response = requests.get(image_file)
+        image = Image.open(BytesIO(response.content)).convert("RGB")
+    else:
+        image = Image.open(image_file).convert("RGB")
+    return image
+
+
+def load_images(image_files):
+    out = []
+    for image_file in image_files:
+        image = load_image(image_file)
+        out.append(image)
+    return out
+
+
+def vis_images(image_files):
+    if len(image_files) == 1:
+        image = image_files[0]
+        os.system(f"termvisage --query-timeout 1 {image} -H left --height 12")
+
+    else:
+        # Concat images
+        system_inst = "convert "
+        inst_template1 = " \\( {image} -background none -resize x500 \\) "
+        inst_template2 = " \\( {image} -background none -resize x500 -splice 100x0 \\) "
+        count = 0
+        for image in image_files:
+            count += 1
+            if count == 1:
+                system_inst += inst_template1.format(image=image)
+            else:
+                system_inst += inst_template2.format(image=image)
+        system_inst += " +append .vis.jpg"
+        os.system(system_inst)
+
+        os.system(f"termvisage --query-timeout 1 .vis.jpg -H left")
+
+
+def expand2square(pil_img, background_color):
+    """
+    Copy from Llava codebase for image preprocessing.
+    """
+    width, height = pil_img.size
+    if width == height:
+        return pil_img
+    elif width > height:
+        result = Image.new(pil_img.mode, (width, width), background_color)
+        result.paste(pil_img, (0, (width - height) // 2))
+        return result
+    else:
+        result = Image.new(pil_img.mode, (height, height), background_color)
+        result.paste(pil_img, ((height - width) // 2, 0))
+        return result
+
+
+def process_images(images, image_processor, model_cfg):
+    """
+    Copy from Llava codebase for image preprocessing.
+    """
+    image_aspect_ratio = getattr(model_cfg, "image_aspect_ratio", None)
+    new_images = []
+    if image_aspect_ratio == "pad":
+        for image in images:
+            image = expand2square(
+                image, tuple(int(x * 255) for x in image_processor.image_mean)
+            )
+            image = image_processor.preprocess(image, return_tensors="pt")[
+                "pixel_values"
+            ][0]
+            if "intern" in image_processor.__class__.__name__.lower():
+                # special case
+                new_images.append(image.unsqueeze(0))
+            else:
+                new_images.append(image)
+    else:
+        ret = image_processor(images, return_tensors="pt")["pixel_values"]
+        if "intern" in image_processor.__class__.__name__.lower():
+            # special case
+            ret = [x.unsqueeze(0) for x in ret]
+        return ret
+    if all(x.shape == new_images[0].shape for x in new_images):
+        new_images = torch.stack(new_images, dim=0)
+
+    return new_images
diff --git a/llm-awq/tinychat/utils/prompt_templates.py b/llm-awq/tinychat/utils/prompt_templates.py
new file mode 100644
index 0000000000000000000000000000000000000000..2ac65909bd5b6724fb5143544b386ac6a3f637dc
--- /dev/null
+++ b/llm-awq/tinychat/utils/prompt_templates.py
@@ -0,0 +1,399 @@
+from typing import List
+from tinychat.utils.constants import (
+    LLAVA_DEFAULT_IMAGE_TOKEN,
+    LLAVA_DEFAULT_IMAGE_PATCH_TOKEN,
+)
+
+
+def get_image_token(model, model_name):
+    return LLAVA_DEFAULT_IMAGE_TOKEN + "\\n "
+    # if "llava" in model_name.lower():
+    #     return LLAVA_DEFAULT_IMAGE_TOKEN + "\\n "
+    # elif "vila" in model_name.lower():
+    #     vision_config = model.get_vision_tower().vision_tower.config
+    #     image_token_len = (vision_config.image_size // vision_config.patch_size) ** 2
+    #     if (
+    #         "downsample" in model.config.mm_projector_type
+    #         or "ds" in model.config.mm_projector_type
+    #     ):
+    #         image_token_len = image_token_len // 4
+    #     if "p32" in model_name:  # extra leading patches
+    #         image_token_len += 32
+    #     elif "se" in model.config.mm_projector_type:
+    #         image_token_len += 2
+    #     return LLAVA_DEFAULT_IMAGE_PATCH_TOKEN * image_token_len + "\\n "
+    # return ""
+
+
+class BasePrompter:
+    def __init__(
+        self,
+        system_inst,
+        role1,
+        role2,
+        sen_spliter="\n",
+        qa_spliter="\n",
+        decorator: List[str] = None,
+        colon=":",
+    ):
+        self.system_inst = system_inst  # System Instruction
+        self.role1 = role1  # The name of USER
+        self.role2 = role2  # The name of AI-Assistant
+        self.sen_spliter = sen_spliter  # How to split system/user/assistant outputs
+        self.qa_spliter = qa_spliter  # How to split Q&A rounds
+        self.decorator = decorator
+        self.colon = colon
+        if self.decorator == None:
+            self.starter = ""
+            self.stopper = ""
+        else:
+            self.starter = self.decorator[0]
+            self.stopper = self.decorator[1]
+        if self.system_inst == None:
+            self.template = (
+                self.starter
+                + self.role1
+                + self.colon
+                + " {prompt}"
+                + self.stopper
+                + self.sen_spliter
+                + self.starter
+                + self.role2
+                + self.colon
+            )
+        else:
+
+            self.template = (
+                self.starter
+                + self.system_inst
+                + self.stopper
+                + self.sen_spliter
+                + self.starter
+                + self.role1
+                + self.colon
+                + " {prompt}"
+                + self.stopper
+                + self.sen_spliter
+                + self.starter
+                + self.role2
+                + self.colon
+            )
+        self.model_input = None
+
+    def insert_prompt(self, input_prompt):
+        self.model_input = self.template.format(prompt=input_prompt)
+
+    def update_template(self, outputs, chunk_prefilling=0):
+        if chunk_prefilling:
+            self.template = (
+                self.role1
+                + self.colon
+                + "{prompt}"
+                + self.stopper
+                + self.sen_spliter  # blank space
+                + self.starter
+                + self.role2
+                + self.colon
+            )
+        else:
+            self.template = (
+                self.model_input
+                + " "
+                + outputs.strip()
+                + self.stopper
+                + self.qa_spliter
+                + self.starter
+                + self.role1
+                + self.colon
+                + "{prompt}"
+                + self.stopper
+                + self.sen_spliter
+                + self.starter
+                + self.role2
+                + self.colon
+            )
+        self.model_input = None
+
+
+class OneShotBasePrompter(BasePrompter):
+    def __init__(
+        self,
+        oneshot_example: List[str],  # User prompt + Assistant responce
+        system_inst,
+        role1,
+        role2,
+        sen_spliter="\n",
+        qa_spliter="\n",
+        decorator: List[str] = None,
+    ):
+        super().__init__(system_inst, role1, role2, sen_spliter, qa_spliter)
+        assert len(oneshot_example) == 2, "One-shot example must be a List of 2 strs."
+        self.user_example = oneshot_example[0]
+        self.assistant_example = oneshot_example[1]
+        self.insert_prompt(self.user_example)
+        self.update_template(self.assistant_example)
+
+
+class EmptyPrompter(BasePrompter):
+    def __init__(self):
+        system_inst = ""
+        role1 = ""
+        role2 = ""
+        sen_spliter = ""
+        qa_spliter = "</s>"
+        super().__init__(system_inst, role1, role2, sen_spliter, qa_spliter)
+
+
+class VicunaPrompter(BasePrompter):
+    def __init__(self):
+        system_inst = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions."
+        role1 = "USER"
+        role2 = "ASSISTANT"
+        sen_spliter = " "
+        qa_spliter = "</s>"
+        super().__init__(system_inst, role1, role2, sen_spliter, qa_spliter)
+
+
+class Llama2Prompter(OneShotBasePrompter):
+    def __init__(self, short_prompt=False):
+        system_inst = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions."
+        role1 = "### Human"
+        role2 = "### Assistant"
+        sen_spliter = "\n"
+        qa_spliter = "</s>"
+        user_example = "Got any creative ideas for a 10 year old's birthday?"
+        if short_prompt:
+            assistant_example = (
+                "Of course! Here are some creative ideas for a 10-year-old's birthday party:\n"
+                + "1. Treasure Hunt: Organize a treasure hunt in your backyard or nearby park. Create clues and riddles for the kids to solve, leading them to hidden treasures and surprises.\n"
+                + "2. Science Party: Plan a science-themed party where kids can engage in fun and interactive experiments. You can set up different stations with activities like making slime, erupting volcanoes, or creating simple chemical reactions.\n"
+                + "Remember to tailor the activities to the birthday child's interests and preferences. Have a great celebration!"
+            )
+        else:
+            assistant_example = (
+                "Of course! Here are some creative ideas for a 10-year-old's birthday party:\n"
+                + "1. Treasure Hunt: Organize a treasure hunt in your backyard or nearby park. Create clues and riddles for the kids to solve, leading them to hidden treasures and surprises.\n"
+                + "2. Science Party: Plan a science-themed party where kids can engage in fun and interactive experiments. You can set up different stations with activities like making slime, erupting volcanoes, or creating simple chemical reactions.\n"
+                + "3. Outdoor Movie Night: Set up a backyard movie night with a projector and a large screen or white sheet. Create a cozy seating area with blankets and pillows, and serve popcorn and snacks while the kids enjoy a favorite movie under the stars.\n"
+                + "4. DIY Crafts Party: Arrange a craft party where kids can unleash their creativity. Provide a variety of craft supplies like beads, paints, and fabrics, and let them create their own unique masterpieces to take home as party favors.\n"
+                + "5. Sports Olympics: Host a mini Olympics event with various sports and games. Set up different stations for activities like sack races, relay races, basketball shooting, and obstacle courses. Give out medals or certificates to the participants.\n"
+                + "6. Cooking Party: Have a cooking-themed party where the kids can prepare their own mini pizzas, cupcakes, or cookies. Provide toppings, frosting, and decorating supplies, and let them get hands-on in the kitchen.\n"
+                + "7. Superhero Training Camp: Create a superhero-themed party where the kids can engage in fun training activities. Set up an obstacle course, have them design their own superhero capes or masks, and organize superhero-themed games and challenges.\n"
+                + "8. Outdoor Adventure: Plan an outdoor adventure party at a local park or nature reserve. Arrange activities like hiking, nature scavenger hunts, or a picnic with games. Encourage exploration and appreciation for the outdoors.\n"
+                + "Remember to tailor the activities to the birthday child's interests and preferences. Have a great celebration!"
+            )
+        oneshot_example = [user_example, assistant_example]
+        super().__init__(
+            oneshot_example, system_inst, role1, role2, sen_spliter, qa_spliter
+        )
+
+
+class Llama3Prompter(BasePrompter):
+    """
+    Example:
+    <|start_header_id|>user<|end_header_id|>
+
+    Show me some attractions in Boston.<|eot_id|>
+
+    <|start_header_id|>assistant<|end_header_id|>
+
+    """
+
+    def __init__(self):
+        system_inst = ""
+        role1 = "<|start_header_id|>user<|end_header_id|>\n\n"
+        role2 = "<|start_header_id|>assistant<|end_header_id|>\n\n"
+        sen_spliter = "<|eot_id|>"
+        qa_spliter = ""
+        colon = ""
+        super().__init__(
+            system_inst, role1, role2, sen_spliter, qa_spliter, colon=colon
+        )
+
+
+class QwenPrompter(BasePrompter):
+    def __init__(self):
+        system_inst = "<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n"
+        role1 = "<|im_start|>user\n"
+        role2 = "<|im_start|>assistant\n"
+        sen_spliter = "<|im_end|>\n"
+        qa_spliter = ""
+        colon = ""
+        super().__init__(
+            system_inst, role1, role2, sen_spliter, qa_spliter, colon=colon
+        )
+
+
+class LlavaLlamaPrompter(BasePrompter):
+    def __init__(self):
+        system_inst = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions."
+        role1 = "USER"
+        role2 = "ASSISTANT"
+        sen_spliter = " "
+        qa_spliter = "</s>"
+        super().__init__(system_inst, role1, role2, sen_spliter, qa_spliter)
+
+
+class LlavaLlama3Prompter(BasePrompter):
+    """
+    Example:
+    <|start_header_id|>user<|end_header_id|>
+
+    Show me some attractions in Boston.<|eot_id|>
+
+    <|start_header_id|>assistant<|end_header_id|>
+
+    """
+
+    def __init__(self):
+        system_inst = (
+            "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nYou are a helpful language and vision assistant. "
+            + "You are able to understand the visual content that the user provides, "
+            + "and assist the user with a variety of tasks using natural language."
+        )
+        role1 = "<|start_header_id|>user<|end_header_id|>\n\n"
+        role2 = "<|start_header_id|>assistant<|end_header_id|>\n\n"
+        sen_spliter = "<|end_of_text|>"
+        qa_spliter = ""
+        colon = ""
+        super().__init__(
+            system_inst, role1, role2, sen_spliter, qa_spliter, colon=colon
+        )
+
+
+class FalconSimplePrompter(BasePrompter):
+    def __init__(self):
+        system_inst = None
+        role1 = "User"
+        role2 = "Assistant"
+        sen_spliter = "\n\n"
+        qa_spliter = "\n\n"
+        super().__init__(system_inst, role1, role2, sen_spliter, qa_spliter)
+
+
+class FalconPrompter(BasePrompter):
+    def __init__(self):
+        system_inst = (
+            "The following is a conversation between a highly knowledgeable and intelligent AI assistant, called Falcon, "
+            + "and a human user, called User. In the following interactions, User and Falcon will converse in natural language, "
+            + "and Falcon will answer User's questions. Falcon was built to be respectful, polite and inclusive. "
+            + "Falcon was built by the Technology Innovation Institute in Abu Dhabi. "
+            + "Falcon will never decline to answer a question, and always attempts to give an answer that User would be satisfied with. "
+            + "It knows a lot, and always tells the truth. The conversation begins."
+        )
+        role1 = "User"
+        role2 = "Assistant"
+        sen_spliter = "\n"
+        qa_spliter = "\n"
+        super().__init__(system_inst, role1, role2, sen_spliter, qa_spliter)
+
+
+class MPTPrompter(BasePrompter):
+    def __init__(self):
+        system_inst = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions."
+        role1 = "### Human"
+        role2 = "### Assistant"
+        sen_spliter = "\n"
+        qa_spliter = "\n"
+        super().__init__(system_inst, role1, role2, sen_spliter, qa_spliter)
+
+
+class MPTChatPrompter(BasePrompter):
+    def __init__(self):
+        system_inst = (
+            "system\n"
+            + "- You are a helpful assistant chatbot trained by MosaicML.\n"
+            + "- You answer questions.\n"
+            + "- You are excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.\n"
+            + "- You are more than just an information source, you are also able to write poetry, short stories, and make jokes."
+        )
+        role1 = "user"
+        role2 = "assistant"
+        sen_spliter = "\n"
+        qa_spliter = "\n"
+        decorator = ["<|im_start|>", "<|im_end|>"]
+        super().__init__(system_inst, role1, role2, sen_spliter, qa_spliter, decorator)
+
+
+class NVILAPrompter(BasePrompter):
+    def __init__(self):
+        system_inst = "system\n" + "You are a helpful assistant<|im_end|>\n"
+        role1 = "user"
+        role2 = "assistant"
+        sen_spliter = "\n"
+        qa_spliter = "\n"
+        decorator = ["<|im_start|>", "<|im_end|>"]
+        super().__init__(
+            system_inst, role1, role2, sen_spliter, qa_spliter, decorator, "\n"
+        )
+
+class InternVL3Prompter(BasePrompter):
+    def __init__(self):
+        system_inst = "system\n你是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。你可以理解用户提供的视觉内容，并使用自然语言帮助用户完成各种任务。"
+        role1 = "user"
+        role2 = "assistant"
+        sen_splitter = "\n"
+        qa_splitter = "\n"
+        decorator = ["<|im_start|>", "<|im_end|>"]
+        super().__init__(
+            system_inst, role1, role2, sen_splitter, qa_splitter, decorator, "\n"
+        )
+
+
+def get_prompter(model_type, model_path="", short_prompt=False, empty_prompt=False):
+    if empty_prompt:
+        return EmptyPrompter()
+    if model_type.lower() == "llama":
+        if "vicuna" in model_path.lower():
+            return VicunaPrompter()
+        elif (
+            "llama-3" in model_path.lower() or "llama3" in model_path.lower()
+        ) and "30b" not in model_path.lower():
+            if "vila" in model_path.lower():
+                # with system prompt by default
+                return LlavaLlama3Prompter()
+            else:
+                return Llama3Prompter()
+        elif "llava" in model_path.lower() or "vila" in model_path.lower():
+            return LlavaLlamaPrompter()
+        else:
+            return Llama2Prompter(short_prompt)
+    elif model_type.lower() == "falcon":
+        # return FalconPrompter()
+        return FalconSimplePrompter()
+    elif "qwen" in model_path.lower() or "qwen" in model_type.lower():
+        return QwenPrompter()
+    elif model_type.lower() == "mpt":
+        if "mpt" and "chat" in model_path.lower():
+            return MPTChatPrompter()
+        else:
+            return MPTPrompter()
+    elif model_type.lower() == "nvila":
+        return NVILAPrompter()
+    elif model_type.lower() == "internvl3":
+        return InternVL3Prompter()
+    else:
+        raise ValueError(f"model type {model_type} is not supported")
+
+
+def get_stop_token_ids(model_type, model_path=""):
+    if model_type.lower() == "llama":
+        if (
+            "llama-3" in model_path.lower() or "llama3" in model_path.lower()
+        ) and "30b" not in model_path.lower():
+            # llama3
+            return [128001, 128009]
+        return []
+    elif model_type.lower() == "qwen" or model_type.lower() == "internvl3":
+        return [151643, 151645]
+    elif model_type.lower() == "falcon":
+        return [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
+    elif model_type.lower() == "mpt":
+        if "mpt" and "chat" in model_path:
+            return [50278, 0]
+        else:
+            return []
+    elif model_type.lower() == "nvila":
+        return [151645]
+    else:
+        raise ValueError(f"model type {model_type} is not supported")
diff --git a/llm-awq/tinychat/vila10_demo.py b/llm-awq/tinychat/vila10_demo.py
new file mode 100644
index 0000000000000000000000000000000000000000..ecd1460bb703b2cc17a5653786ccbe42abe557c6
--- /dev/null
+++ b/llm-awq/tinychat/vila10_demo.py
@@ -0,0 +1,233 @@
+import argparse
+import torch
+
+from PIL import Image
+from tqdm import tqdm
+
+from transformers import AutoConfig, AutoTokenizer
+from accelerate import load_checkpoint_and_dispatch
+
+from tinychat.utils.tune import (
+    device_warmup,
+    tune_all_wqlinears,
+    tune_llava_patch_embedding,
+)
+from tinychat.utils.prompt_templates import (
+    get_prompter,
+    get_stop_token_ids,
+    get_image_token,
+)
+from tinychat.utils.llava_image_processing import (
+    process_images,
+    load_images,
+    vis_images,
+)
+import tinychat.utils.constants
+from tinychat.models.llava_llama import LlavaLlamaForCausalLM
+from tinychat.stream_generators.llava_stream_gen import LlavaStreamGenerator
+from tinychat.utils.conversation_utils import gen_params, stream_output, TimeStats
+
+import os
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+
+def image_parser(args):
+    out = args.image_file.split(args.im_sep)
+    return out
+
+
+def skip(*args, **kwargs):
+    pass
+
+
+def main(args):
+    # Accelerate model initialization
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+    torch.nn.init.kaiming_uniform_ = skip
+    torch.nn.init.kaiming_normal_ = skip
+    torch.nn.init.uniform_ = skip
+    torch.nn.init.normal_ = skip
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model_path, use_fast=False)
+    tinychat.utils.constants.LLAVA_DEFAULT_IMAGE_PATCH_TOKEN_IDX = (
+        tokenizer.convert_tokens_to_ids(
+            [tinychat.utils.constants.LLAVA_DEFAULT_IMAGE_PATCH_TOKEN]
+        )[0]
+    )
+    config = AutoConfig.from_pretrained(args.model_path, trust_remote_code=True)
+    config.min_max_range_path = args.model_path + "/emb_min_max.pt"
+    model = LlavaLlamaForCausalLM(config, args.device).half()
+    vision_tower = model.get_model().vision_tower
+    if not vision_tower.is_loaded:
+        vision_tower.load_model()
+    image_processor = vision_tower.image_processor
+    vision_tower = vision_tower.half()
+
+    if args.precision == "W16A16":
+        pbar = tqdm(range(1))
+        pbar.set_description("Loading checkpoint shards")
+        for i in pbar:
+            model = load_checkpoint_and_dispatch(
+                model,
+                args.model_path,
+                no_split_module_classes=[
+                    "OPTDecoderLayer",
+                    "LlamaDecoderLayer",
+                    "BloomBlock",
+                    "MPTBlock",
+                    "DecoderLayer",
+                    "CLIPEncoderLayer",
+                ],
+            ).to(args.device)
+
+    elif args.precision == "W4A16":
+        from tinychat.utils.load_quant import load_awq_model
+
+        model = load_awq_model(model, args.quant_path, 4, 128, args.device)
+        from tinychat.modules import (
+            make_quant_norm,
+            make_quant_attn,
+            make_fused_mlp,
+            make_fused_vision_attn,
+        )
+
+        make_quant_attn(model, args.device)
+        make_quant_norm(model)
+        # make_fused_mlp(model)
+        # make_fused_vision_attn(model,args.device)
+        model = model.to(args.device)
+
+    else:
+        raise NotImplementedError(f"Precision {args.precision} is not supported.")
+
+    image_files = image_parser(args)
+    image_num = len(image_files)
+    images = load_images(image_files)
+    if args.vis_image:
+        print("=" * 50)
+        print("Input Image:")
+        vis_images(image_files)
+    # Similar operation in model_worker.py
+    image_tensor = process_images(images, image_processor, model.config)
+    if type(image_tensor) is list:
+        image_tensor = [
+            image.to(args.device, dtype=torch.float16) for image in image_tensor
+        ]
+    else:
+        image_tensor = image_tensor.to(args.device, dtype=torch.float16)
+
+    device_warmup(args.device)
+    tune_llava_patch_embedding(vision_tower, device=args.device)
+
+    stream_generator = LlavaStreamGenerator
+
+    if args.max_seq_len <= 1024:
+        short_prompt = True
+    else:
+        short_prompt = False
+    model_prompter = get_prompter(
+        args.model_type, args.model_path, short_prompt, args.empty_prompt
+    )
+    stop_token_ids = get_stop_token_ids(args.model_type, args.model_path)
+    count = 0
+
+    if args.empty_prompt:
+        input_indicator = "Input: "
+        output_indicator = "Generated: "
+    else:
+        input_indicator = "USER: "
+        output_indicator = "ASSISTANT: "
+
+    model.eval()
+    time_stats = TimeStats()
+    while True:
+        # Get input from the user
+        print("=" * 50)
+        input_prompt = input(input_indicator)
+        print("-" * 50)
+        if input_prompt == "":
+            print("EXIT...")
+            time_stats.show()
+            break
+        if count == 0:  # Insert image here
+            image_token = get_image_token(model, args.model_path)
+            image_token_holder = (
+                tinychat.utils.constants.LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER
+            )
+            im_token_count = input_prompt.count(image_token_holder)
+            if im_token_count == 0:
+                model_prompter.insert_prompt(image_token * image_num + input_prompt)
+            else:
+                assert im_token_count == image_num
+                input_prompt = input_prompt.replace(image_token_holder, image_token)
+                model_prompter.insert_prompt(input_prompt)
+        else:
+            model_prompter.insert_prompt(input_prompt)
+        output_stream = stream_generator(
+            model,
+            tokenizer,
+            model_prompter.model_input,
+            gen_params,
+            device=args.device,
+            stop_token_ids=stop_token_ids,
+            image_tensor=image_tensor,
+        )
+        print(output_indicator, end="", flush=True)
+        if count == 0:
+            outputs = stream_output(output_stream, time_stats)
+        else:
+            outputs = stream_output(output_stream)
+        if (
+            args.single_round is not True and args.max_seq_len > 512
+        ):  # Only memorize previous conversations when kv_cache_size > 512
+            model_prompter.update_template(outputs)
+        count += 1
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_type", type=str, default="LLaMa", help="type of the model"
+    )
+    parser.add_argument(
+        "--model-path", type=str, default="/data/llm/checkpoints/llava/llava-v1.5-7b"
+    )
+    parser.add_argument(
+        "--quant-path",
+        type=str,
+        default="/data/llm/checkpoints/llava/llava-v1.5-7b-w4-g128-awq.pt",
+    )
+    parser.add_argument(
+        "--precision", type=str, default="W4A16", help="compute precision"
+    )
+    parser.add_argument(
+        "--image-file",
+        type=str,
+        default="https://llava.hliu.cc/file=/nobackup/haotian/code/LLaVA/llava/serve/examples/extreme_ironing.jpg",
+    )
+    parser.add_argument(
+        "--im-sep",
+        type=str,
+        default=",",
+    )
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--max_seq_len", type=int, default=2048)
+    parser.add_argument(
+        "--single_round",
+        action="store_true",
+        help="whether to memorize previous conversations",
+    )
+    parser.add_argument(
+        "--vis-image",
+        action="store_true",
+        help="whether to visualize the image while chatting",
+    )
+    parser.add_argument(
+        "--empty-prompt",
+        action="store_true",
+        help="whether to use empty prompt template",
+    )
+    args = parser.parse_args()
+    main(args)
diff --git a/lm-evaluation-harness/.gitignore b/lm-evaluation-harness/.gitignore
new file mode 100644
index 0000000000000000000000000000000000000000..9ae167be97686d8e332e469e3d84708879860091
--- /dev/null
+++ b/lm-evaluation-harness/.gitignore
@@ -0,0 +1,47 @@
+# macOS system files
+.DS_Store
+
+# Virtual environments
+.venv/
+venv/
+ENV/
+env/
+*.env
+
+# Python bytecode and build artifacts
+__pycache__/
+*.py[cod]
+*.so
+*.egg-info/
+build/
+dist/
+
+# IDE & editor settings
+.vscode/
+.idea/
+
+# Jupyter
+.ipynb_checkpoints/
+profile_default/
+ipython_config.py
+
+# Output and data
+output/
+data/
+temp/
+test_logs/
+
+# Caching
+lm_eval/caching/.cache
+lm_cache/
+
+# Logging
+*.log
+logs/
+
+# wandb experiment tracking
+wandb/
+examples/wandb/
+
+# PyInstaller
+*.spec
diff --git a/lm-evaluation-harness/.pre-commit-config.yaml b/lm-evaluation-harness/.pre-commit-config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..af3f9f086976f5caf5046e623e9cf2d9f2785057
--- /dev/null
+++ b/lm-evaluation-harness/.pre-commit-config.yaml
@@ -0,0 +1,60 @@
+# Ignore test linting to avoid conflicting changes to version stability.
+exclude: ^tests/testdata/
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v5.0.0
+    hooks:
+      - id: check-added-large-files
+      - id: check-ast
+      - id: check-byte-order-marker
+      - id: check-case-conflict
+      - id: check-json
+      - id: check-merge-conflict
+        args: [--assume-in-merge]
+      - id: check-symlinks
+      - id: check-yaml
+        args: ["--unsafe"]
+      - id: destroyed-symlinks
+      - id: detect-private-key
+      - id: end-of-file-fixer
+      - id: no-commit-to-branch
+        always_run: false
+      - id: requirements-txt-fixer
+      - id: trailing-whitespace
+        args: [--markdown-linebreak-ext=md]
+      - id: fix-byte-order-marker
+        exclude: docs/CNAME
+      - id: fix-encoding-pragma
+        args: [--remove]
+      - id: mixed-line-ending
+        args: [--fix=lf]
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.11.10
+    hooks:
+      # Run the linter.
+      - id: ruff
+        args:
+          - --fix
+        # Run the formatter.
+      - id: ruff-format
+  - repo: https://github.com/codespell-project/codespell
+    rev: v2.4.1
+    hooks:
+      - id: codespell
+        exclude: >
+          (?x)^(
+              .*\.json|ignore.txt|lm_eval/tasks/.*|.*yaml|.*\.ipynb
+          )$
+        args: [--check-filenames, --check-hidden, --ignore-words=ignore.txt]
+  - repo: https://github.com/jackdewinter/pymarkdown
+    rev: v0.9.29
+    hooks:
+      - id: pymarkdown
+        exclude: ^(lm_eval/tasks/.*|docs/footguns\.md)$
+        args: [fix, -r]
+#  - repo: https://github.com/pre-commit/mirrors-mypy
+#    rev: v1.5.1
+#    hooks:
+#    - id: mypy
+#      additional_dependencies: [".[sentencepiece,multilingual,promptsource,gptq]", "types-PyYAML", "types-requests"]
+#      exclude: ^tests/.*$
diff --git a/obtain_metric.py b/obtain_metric.py
new file mode 100644
index 0000000000000000000000000000000000000000..6c7119d74e62c0ddc60f65c86d815fe308580a66
--- /dev/null
+++ b/obtain_metric.py
@@ -0,0 +1,242 @@
+import json
+import numpy as np
+import os
+from zscore import kurtosis_outlier_layers
+
+def mean_and_var(arr):
+    """
+    输入: arr (np.ndarray 或能转成 np.array 的对象)
+    输出: (mean, variance)
+    """
+    arr = np.asarray(arr, dtype=float)
+    mean = np.mean(arr)
+    var = np.var(arr)   # 默认是总体方差，如果要无偏估计可用 ddof=1
+    return mean, var
+
+def select_layers_shaped(kurtosis, alpha, k=5, return_scores=False,
+                         pos_gaussians=None, neg_gaussians=None):
+    """
+    强塑形联合指标:
+      score_i = (alpha_i / kurtosis_i) * ( 1
+                                           + sum_j b_j * exp(-((kurtosis_i - c_j)^2) / s_j^2)
+                                           - sum_t d_t * exp(-((kurtosis_i - u_t)^2) / r_t^2) )
+
+    参数
+    ----
+    kurtosis : list/ndarray   κ
+    alpha    : list/ndarray   α
+    k        : 取前 k 个下标
+    pos_gaussians : [(c, s, b), ...]  正向高斯(中心c, 带宽s>0, 系数b>0)
+    neg_gaussians : [(u, r, d), ...]  负向高斯(中心u, 带宽r>0, 系数d>0)
+
+    返回
+    ----
+    idx_topk : list[int]
+    (可选) scores : ndarray
+    """
+    krt = np.asarray(kurtosis, dtype=float)
+    alp = np.asarray(alpha, dtype=float)
+    if krt.shape != alp.shape:
+        raise ValueError("kurtosis 与 alpha 形状不一致")
+
+    base = alp - krt
+
+    c, sigma = mean_and_var(krt)
+    beta = 1.5
+    weight = 1.0 + beta * np.exp(-((krt - c) ** 2) / (sigma ** 2))
+    scores = base * weight
+    return scores
+
+def joint_score(kurtosis, alpha,
+                alpha0=4.20, s_alpha=0.090, 
+                s_k=1.254, w_alpha=9.48, w_k=1.02):
+    """
+    Kurtosis-Alpha 联合指标 (Taguchi 损失型)
+    -------------------------
+    参数
+    kurtosis : list/ndarray
+        每层的 kurtosis 值
+    alpha : list/ndarray
+        每层的 alpha 值
+    alpha0 : float
+        alpha 的目标值（名义最佳点）
+    s_alpha : float
+        alpha 的尺度因子
+    s_k : float
+        kurtosis 的尺度因子
+    w_alpha, w_k : float
+        alpha 和 kurtosis 的权重
+
+    返回
+    scores : ndarray
+        每层的综合得分（越大越好）
+    """
+    k = np.array(kurtosis)
+    a = np.array(alpha)
+
+    loss_k = (k / s_k) ** 2
+    loss_a = ((a - alpha0) / s_alpha) ** 2
+
+    scores = -(w_k * loss_k + w_alpha * loss_a)
+    return scores
+
+def simple_joint_score(kurtosis, alpha, stable_rank):
+    """
+    极简联合指标:
+    S_l = -k_l * (alpha_l - Q80(alpha))^2
+    
+    参数
+    ----
+    kurtosis : list/ndarray
+        每层的 kurtosis 值
+    alpha : list/ndarray
+        每层的 alpha 值
+    
+    返回
+    ----
+    scores : ndarray
+        每层的综合得分（越大越好）
+    """
+    k = np.array(kurtosis)
+    a = np.array(alpha)
+    s = np.array(stable_rank)
+    print('alpha_hat_datas')
+    # h = np.array(alpha_hat_datas)
+    k = (k - k.min()) / (k.max() - k.min())
+    # k_exp = np.exp(k)   # 防止溢出
+    # k = k_exp / np.sum(k_exp)
+    
+    a = 1 / a
+    a = (a - a.min()) / (a.max() - a.min())
+    # a_exp = np.exp(a)   # 防止溢出
+    # a = a_exp / np.sum(a_exp)
+
+    s = 1 / s
+    s = (s - s.min()) / (s.max() - s.min())
+    # s_exp = np.exp(s)  
+    # s = s_exp / np.sum(s_exp)
+
+    # h = 1 / h
+    # # h= (h - h.min()) / (h.max() - h.min())
+    # h_exp = np.exp(h)   # 防止溢出
+    # h = h_exp / np.sum(h_exp)
+    
+    # alpha_target = np.percentile(a, 80)   # α 的 80 分位数
+    # # alpha_target = 0
+    # scores = -k * (a - alpha_target) ** 2
+    
+    # k = k * 0.5
+    # h = h * 0.5
+    # print(k)
+    # print(a)
+    # print(h)
+    # a = a + 0.3 * h
+
+
+    # a = a +1
+    # k = k +1 
+    # h = h +1
+    print('a', {i:aa for i, aa in enumerate(a)})
+    print()
+    print('k', {i:aa for i, aa in enumerate(k)})
+    print()
+    print('s', {i:aa for i, aa in enumerate(s)})
+    # k = 100 * k
+    # scores = (k / a) * (k - a)
+    k = 0.2 * k
+    s = 1.2 * s
+    scores = k + a + s
+    # scores = (k * h) * (k + h)
+    # print()
+    # print('scores', {i:aa for i, aa in enumerate(scores)})
+    return scores
+
+model_names = ['Llama-2-7b-hf',] 
+            #    'Llama-2-13b-hf', 'Qwen3-8B', 'Qwen3-4B', 'Mistral-7B-Instruct-v0.3','Llama-3.2-3B-Instruct']
+
+for model in model_names:
+    # model = "Llama-2-7b-hf"
+    print(model)
+    alpha_path = f'/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/alpha_values/alpha_values_{model}.json'
+    kurtosis_path = f'/mnt/bn/life-mllm/users/cxr/quantization/lm-quant-toolkit/kurtosis_means/kurtosis_means-{model}.json'
+    alpha_hat_path = f'/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/alpha_hat/alpha_values_{model}.json'
+    stable_rank_path = f'/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/stable_rank/stable_rank_{model}.json'
+    zd_path = f'/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/ZD/ZD_{model}.json'
+    bi_path = f'/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/BI/BI_{model}.json'
+
+    with open(stable_rank_path, 'r', encoding='utf-8') as f:
+        stable_rank_datas = json.load(f)
+
+    with open(alpha_path, 'r', encoding='utf-8') as f:
+        alpha_datas = json.load(f)
+
+    with open(kurtosis_path, 'r', encoding='utf-8') as f:
+        kurtosis_datas = json.load(f)
+
+    with open(alpha_hat_path, 'r', encoding='utf-8') as f:
+        alpha_hat_datas = json.load(f)
+
+    with open(zd_path, 'r', encoding='utf-8') as f:
+        zd_datas = json.load(f)
+
+    with open(bi_path, 'r', encoding='utf-8') as f:
+        bi_datas = json.load(f)
+
+    print(len(kurtosis_datas))
+    print('kurtosis_datas', np.argsort(kurtosis_datas))
+    print('stable_rank_datas', np.argsort([-a for a in stable_rank_datas]))
+    print('alpha_datas', np.argsort([-a for a in alpha_datas]))
+    print('alpha_hat_datas', np.argsort([-a for a in alpha_hat_datas]))
+    
+    print('zd_datas', np.argsort([-a for a in zd_datas]))
+    print('bi_datas', np.argsort([a for a in bi_datas]))
+
+    scores = simple_joint_score(kurtosis_datas, alpha_datas, stable_rank_datas)
+    idx_sorted = np.argsort(scores)
+    # print('scores:', scores)
+    print('idx_sorted:',idx_sorted)
+    print()
+
+    kurtosis_idx = np.argsort(kurtosis_datas).tolist()
+    stable_rank_idx = np.argsort([-a for a in stable_rank_datas]).tolist()
+    alpha_idx = np.argsort([-a for a in alpha_datas]).tolist()
+    bi_idx = np.argsort([a for a in bi_datas]).tolist()
+    z_idx = kurtosis_outlier_layers(kurtosis_datas)
+    zd_idx = np.argsort([-a for a in zd_datas]).tolist()
+    layrs = [5, 10]
+
+    for layr in layrs:
+    #     bits = [4 for _ in range(len(alpha_datas))]
+    #     print(kurtosis_idx[:layr])
+    #     for i in kurtosis_idx[:layr]:
+    #         bits[i] = 2
+    #     with open(os.path.join('/mnt/bn/life-mllm/users/cxr/quantization/baselines', f'{model}_kurtosis_idx_{layr}.json'), "w", encoding="utf-8") as f:
+    #         json.dump(bits, f, ensure_ascii=False, indent=4)
+
+    #     bits = [4 for _ in range(len(alpha_datas))]
+    #     for i in alpha_idx[:layr]:
+    #         bits[i] = 2
+    #     with open(os.path.join('/mnt/bn/life-mllm/users/cxr/quantization/baselines', f'{model}_alpha_idx_{layr}.json'), "w", encoding="utf-8") as f:
+    #         json.dump(bits, f, ensure_ascii=False, indent=4)
+
+    #     bits = [4 for _ in range(len(alpha_datas))]
+    #     for i in z_idx[:layr]:
+    #         bits[i] = 2
+    #     with open(os.path.join('/mnt/bn/life-mllm/users/cxr/quantization/baselines', f'{model}_z_idx_{layr}.json'), "w", encoding="utf-8") as f:
+    #         json.dump(bits, f, ensure_ascii=False, indent=4)
+
+        # bits = [4 for _ in range(len(alpha_datas))]
+        # print(bi_idx[:layr])
+        # for i in bi_idx[:layr]:
+        #     bits[i] = 2
+        # with open(os.path.join('/mnt/bn/life-mllm/users/cxr/quantization/baselines_1', f'{model}_bi_idx_{layr}.json'), "w", encoding="utf-8") as f:
+        #     json.dump(bits, f, ensure_ascii=False, indent=4)
+
+        kurtosis_datas, stable_rank_datas
+        bits = [4 for _ in range(len(alpha_datas))]
+        print(idx_sorted[:layr])
+        for i in idx_sorted[:layr]:
+            bits[i] = 2
+        with open(os.path.join('/mnt/bn/life-mllm/users/cxr/quantization/ours2', f'{model}_aks_plus_idx_sorted_{layr}.json'), "w", encoding="utf-8") as f:
+            json.dump(bits, f, ensure_ascii=False, indent=4)
+        
\ No newline at end of file
diff --git a/test.py b/test.py
new file mode 100644
index 0000000000000000000000000000000000000000..c89a5ffe659a7f7b9d6d95451eda87c8dfed62fc
--- /dev/null
+++ b/test.py
@@ -0,0 +1,24 @@
+from huggingface_hub import HfApi
+
+# 初始化 API
+api = HfApi()
+
+# 上传超大文件夹（自动分片+续传）
+api.upload_large_folder(
+    folder_path="./clean-model-files",  # 清理后的本地目录
+    repo_id="你的用户名/你的仓库名",     # 如 zhangsan/my-llm-model
+    repo_type="model",                  # 仓库类型
+    allow_patterns=[                    # 仅上传以下文件（白名单）
+        "*.safetensors",
+        "*.bin",
+        "config.json",
+        "tokenizer*.json",
+        "vocab.txt"
+    ],
+    ignore_patterns=[                   # 排除以下文件（黑名单）
+        ".cache/**",
+        "*.log",
+        "__pycache__/**"
+    ],
+    overwrite=True,                     # 覆盖已有文件
+)
\ No newline at end of file
diff --git a/test.sh b/test.sh
new file mode 100644
index 0000000000000000000000000000000000000000..b1d9ab1fcadb27239fb01db109f84369766da043
--- /dev/null
+++ b/test.sh
@@ -0,0 +1 @@
+conda activate lm-eval
\ No newline at end of file
diff --git a/test2.py b/test2.py
new file mode 100644
index 0000000000000000000000000000000000000000..6af55604589ae5b7a5f9f1f2541a588eee865aff
--- /dev/null
+++ b/test2.py
@@ -0,0 +1,51 @@
+import json
+import heapq
+def get_top_k_indices(json_file_path, k):
+    """
+    读取JSON文件中的列表，返回最大的k个元素的索引
+    
+    Args:
+        json_file_path: JSON文件路径
+        k: 需要获取的最大元素的个数
+        
+    Returns:
+        list: 按元素大小降序排列的索引列表
+        
+    Raises:
+        FileNotFoundError: 文件不存在时抛出
+        ValueError: k值无效或数据格式错误时抛出
+    """
+    # 读取JSON文件
+    try:
+        with open(json_file_path, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+    except FileNotFoundError:
+        raise FileNotFoundError(f"文件 {json_file_path} 不存在")
+    except json.JSONDecodeError:
+        raise ValueError("JSON文件格式错误")
+    
+    # 验证数据类型
+    if not isinstance(data, list):
+        raise ValueError("JSON文件内容不是列表")
+    
+    # 验证k值的有效性
+    if k <= 0 or k > len(data):
+        raise ValueError(f"k值无效，应在1到{len(data)}之间")
+    
+    # 获取元素值和索引的元组列表 [(value, index), ...]
+    value_index_pairs = [(value, idx) for idx, value in enumerate(data)]
+    
+    # 方法1：使用heapq获取最大的k个元素（效率更高，O(n log k)）
+    top_k_pairs = heapq.nlargest(k, value_index_pairs, key=lambda x: x[0])
+    
+    # 方法2：使用排序（简单直观，O(n log n)）
+    # sorted_pairs = sorted(value_index_pairs, key=lambda x: x[0], reverse=True)
+    # top_k_pairs = sorted_pairs[:k]
+    
+    # 提取索引
+    top_k_indices = [pair[1] for pair in top_k_pairs]
+    
+    return top_k_indices
+
+a = get_top_k_indices('/mnt/bn/life-mllm/users/cxr/quantization/quantization_metric/metrics/alpha/alpha_mlp_Llama-2-7b-hf.json', 10)
+print(a)
\ No newline at end of file
diff --git a/zscore.py b/zscore.py
new file mode 100644
index 0000000000000000000000000000000000000000..45bc481e0d11acdeaea0e782520b02bea5dc6e05
--- /dev/null
+++ b/zscore.py
@@ -0,0 +1,44 @@
+import numpy as np
+import json
+
+def kurtosis_outlier_layers(kurtosis_values, threshold=3.0, method="subtract"):
+    """
+    根据 Kurtosis 值列表，计算差分并用 z-score 检测异常层
+
+    参数:
+        kurtosis_values (list or np.ndarray): Kurtosis 值序列 (s1, s2, ..., sn)
+        threshold (float): z-score 阈值 (默认 3.0)
+        method (str): 差分方式
+            "subtract" -> si+1 - si
+            "divide"   -> si+1 / si
+
+    返回:
+        z_scores (np.ndarray): 差分的 z-score
+        outlier_indices (list): 被判定为异常的层索引 (对应原始 Kurtosis 序列中的层号)
+    """
+    values = np.array(kurtosis_values, dtype=float)
+
+    # Step 1: 差分
+    if method == "subtract":
+        diffs = np.diff(values)  # s2-s1, s3-s2, ...
+    elif method == "divide":
+        diffs = values[1:] / values[:-1]
+    else:
+        raise ValueError("method must be 'subtract' or 'divide'")
+
+    # Step 2: 计算 z-score
+    mean = np.mean(diffs)
+    std = np.std(diffs, ddof=1)  # 样本标准差
+    z_scores = abs((diffs - mean) / std)
+    results = [(i+1, z) for i, z in enumerate(z_scores)]
+    results.sort(key=lambda x: x[1], reverse=False)
+    print(results)
+    results = [i[0] for i in results]
+    
+    return results
+
+k_path = '/mnt/bn/life-mllm/users/cxr/quantization/lm-quant-toolkit/kurtosis_means/kurtosis_means-Llama-2-7b-hf.json'
+with open(k_path, "r", encoding="utf-8") as f:
+    kurtosis_values = json.load(f)
+
+kurtosis_outlier_layers(kurtosis_values)
\ No newline at end of file